Search code examples
derivativeopenmdao

In OpenMDAO can I use set_check_partial_options to prevent certain derivatives from being checked?


I have an om.ExplicitComponent in which some of the derivatives are exact (analytic) and can be checked with the complex step method (cs), some of the derivatives are analytic but cannot be checked with cs, and some which can only be evaluated with finite differences.

import openmdao.api as om

from scipy.special import ellipe, ellipk
from scipy.special import hyp2f1

class MWE(om.ExplicitComponent):
    def setup(self):
        self.add_input("a")
        self.add_input("b")
        self.add_input("c")
        self.add_output("x")
        self.add_output("y")
        self.add_output("z")

    def compute(self, inputs, outputs):
        a = inputs["a"]
        b = inputs["b"]
        c = inputs["c"]
        outputs["x"] = a**2
        outputs["y"] = ellipe(b)
        outputs["z"] = hyp2f1(1 / 10, a, 1 / 2, c)

    def setup_partials(self):
        self.declare_partials("x", ["a"], method="exact")
        self.declare_partials("y", ["b"], method="exact")
        self.declare_partials("z", ["a"], method="fd")
        self.declare_partials("z", ["c"], method="exact")

    def compute_partials(self, inputs, J):
        a = inputs["a"]
        b = inputs["b"]
        c = inputs["c"]
        J["x", "a"] = 2 * a

        J["y", "b"] = (ellipe(b) - ellipk(b)) / (2 * b)

        J["z", "c"] = (a / 5) * hyp2f1(11 / 10, 1 + a, 3 / 2, c)

Is there a way I can use one or more calls to set_check_partial_options and check_partials in order to

  • Check J["x", "a"] using the cs method, since it's the most exact or demanding,
  • Check J["y", "b"] using the fd method, since the elliptic functions can't handle complexes,
  • Check J["z", "c"] using the fd method for the same reason.
  • Don't bother checking J["z", "a"] at all since it needs to be evaluated using fd anyway as there is no analytic formulation possible (at least within scipy). ?

Solution

  • You can use the set_check_partials method to get fine grained control of how the checks are performed. This methods is set inside the setup or setup_partialsof a specific component and lets you control the settings for each input separately.

    If you want to set the whole component to CS or FD you can use "*" for the wrt argument.

    Just be aware that these component level settings get over-written by the method argument you pass into the check_partials method itself. So just leave that blank and configure each component with its own local settings.

    Unfortunately, there seems to be a bug in V3.15 related to set_check_partials_options. The following is the correct syntax for the model, but an error is thrown if you uncomment the cs option. For V3.15 and less, you'll have to check all with CS or none. It should be fixed by V3.16

    import openmdao.api as om
    
    from scipy.special import ellipe, ellipk
    from scipy.special import hyp2f1
    
    class MWE(om.ExplicitComponent):
        def setup(self):
            self.add_input("a")
            self.add_input("b")
            self.add_input("c")
            self.add_output("x")
            self.add_output("y")
            self.add_output("z")
    
        def compute(self, inputs, outputs):
            a = inputs["a"]
            b = inputs["b"]
            c = inputs["c"]
            outputs["x"] = a**2
            outputs["y"] = ellipe(b)
            outputs["z"] = hyp2f1(1 / 10, a, 1 / 2, c)
    
        def setup_partials(self):
            self.declare_partials("x", ["a"], method="exact")
            self.declare_partials("y", ["b"], method="exact")
            self.declare_partials("z", ["a"], method="fd")
            self.declare_partials("z", ["c"], method="exact")
    
    
            # self.set_check_partial_options(wrt="a", method="cs") # bug in V3.15
            self.set_check_partial_options(wrt="a", form="backward") 
            self.set_check_partial_options(wrt="b", form="central")
    
        def compute_partials(self, inputs, J):
            a = inputs["a"]
            b = inputs["b"]
            c = inputs["c"]
            J["x", "a"] = 2 * a
    
            J["y", "b"] = (ellipe(b) - ellipk(b)) / (2 * b)
    
            J["z", "c"] = (a / 5) * hyp2f1(11 / 10, 1 + a, 3 / 2, c)
    
    
    
    if __name__ == "__main__":
    
        p = om.Problem()
        p.model = MWE()
    
        p.setup(force_alloc_complex=True)
    
        # set some workable initial values to avoid inf
        p['a'] = 2. 
        p['b'] = -2. 
        p['c'] = -2. 
    
        p.run_model()
        p.model.list_outputs()
    
        p.check_partials()