Check the Jacobian of a discipline

In this example, the Jacobian of an Discipline is checked by derivative approximation.

from __future__ import annotations

from typing import TYPE_CHECKING

from numpy import array
from numpy import exp

from gemseo.core.discipline import Discipline

if TYPE_CHECKING:
    from collections.abc import Iterable

    from gemseo.typing import StrKeyMapping

First, we create a discipline computing \(f(x,y)=e^{-(1-x)^2-(1-y)^2}\) and \(g(x,y)=x^2+y^2-1\) and introduce an error in the implementation of \(\frac{\partial f(x,y)}{\partial x}\).

class BuggedDiscipline(Discipline):
    def __init__(self) -> None:
        super().__init__()
        self.input_grammar.update_from_names(["x", "y"])
        self.output_grammar.update_from_names(["f", "g"])
        self.default_input_data = {"x": array([0.0]), "y": array([0.0])}

    def _run(self, input_data: StrKeyMapping) -> StrKeyMapping | None:
        x = input_data["x"]
        y = input_data["y"]
        return {"f": exp(-((1 - x) ** 2) - (1 - y) ** 2), "g": x**2 + y**2 - 1}

    def _compute_jacobian(
        self,
        input_names: Iterable[str] = (),
        output_names: Iterable[str] = (),
    ) -> None:
        x = self.io.data["x"]
        y = self.io.data["y"]
        self._init_jacobian()
        g_jac = self.jac["g"]
        g_jac["x"][:] = 2 * x
        g_jac["y"][:] = 2 * y
        f_jac = self.jac["f"]
        aux = 2 * exp(-((1 - x) ** 2) - (1 - y) ** 2)
        f_jac["x"][:] = aux  # this is wrong.
        f_jac["y"][:] = aux * (1 - y)

We want to check if the implemented Jacobian is correct. For practical applications where Jacobians are needed, this is not a simple task. GEMSEO automates such tests thanks to the Discipline.check_jacobian() method.

Finite differences (default)

discipline = BuggedDiscipline()
discipline.check_jacobian(
    input_data={"x": array([0.0]), "y": array([1.0])},
    show=True,
    plot_result=True,
    step=1e-1,
)

   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp f/d x is wrong by 1.9226823669921425%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[0.76978625]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[0.73575888]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[0.03402737]]
   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp f/d y is wrong by 3.531203260551434%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[-0.03660462]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[0.]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[-0.03660462]]
   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp g/d x is wrong by 9.090909090909099%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[0.1]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[0.]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[0.1]]
   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp g/d y is wrong by 3.2258064516129616%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[2.1]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[2.]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[0.1]]
    INFO - 16:23:00: Linearization of Discipline: BuggedDiscipline is wrong.

False

The step here is chosen big enough to underline the truncation error. From this graph, we can see that almost all the provided components of the Jacobians (blue dots) are close but distinct from the approximated by finite differences using a step of 0.1 (red dots). This kind of graph can be used to spot implementation mistakes in fact we can already spot a large mistake in the wrong components.

The derr_approx argument can be either finite_differences, centered_differences or complex_step.

Centered differences

discipline.check_jacobian(
    input_data={"x": array([0.0]), "y": array([1.0])},
    derr_approx=discipline.ApproximationMode.CENTERED_DIFFERENCES,
    show=True,
    plot_result=True,
    step=1e-1,
)

   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp f/d x is wrong by 0.1416340394497155%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[0.73330393]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[0.73575888]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[-0.00245495]]
    INFO - 16:23:00: Jacobian: dp f/dp y succeeded.
    INFO - 16:23:00: Jacobian: dp g/dp x succeeded.
    INFO - 16:23:00: Jacobian: dp g/dp y succeeded.
    INFO - 16:23:00: Linearization of Discipline: BuggedDiscipline is wrong.

False

With the same step the truncation error is in this case much smaller.

Complex step

discipline.check_jacobian(
    input_data={"x": array([0.0]), "y": array([1.0])},
    derr_approx=discipline.ApproximationMode.COMPLEX_STEP,
    show=True,
    plot_result=True,
    step=1e-1,
)

   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp f/d x is wrong by 0.1409521643331901%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[0.73820893]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[0.73575888]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[0.00245004]]
    INFO - 16:23:00: Jacobian: dp f/dp y succeeded.
    INFO - 16:23:00: Jacobian: dp g/dp x succeeded.
    INFO - 16:23:00: Jacobian: dp g/dp y succeeded.
    INFO - 16:23:00: Linearization of Discipline: BuggedDiscipline is wrong.

False

With the same step the truncation error is also smaller than finite differences. This confirms again that an implementation mistake was done.

Advantages and drawbacks of each method

Finite differences and complex are first-order methods, they use one sampling point per input and the truncation error goes down linearly with the step. Centered differences are second-order methods which use twice as many points as finite differences and complex step. Complex step derivatives are less prone to numerical cancellation errors so that a tiny step can be used. On the other hand complex step is not compatible with discipline not supporting complex inputs.

discipline.check_jacobian(
    input_data={"x": array([0.0]), "y": array([1.0])},
    derr_approx=discipline.ApproximationMode.COMPLEX_STEP,
    show=True,
    plot_result=True,
    step=1e-10,
)

    INFO - 16:23:00: Jacobian: dp f/dp x succeeded.
    INFO - 16:23:00: Jacobian: dp f/dp y succeeded.
    INFO - 16:23:00: Jacobian: dp g/dp x succeeded.
    INFO - 16:23:00: Jacobian: dp g/dp y succeeded.
    INFO - 16:23:00: Linearization of Discipline: BuggedDiscipline is correct.

True

Automatic time step

Finite differences and centered differences steps need to be chosen as a trade between truncation and numerical errors. For this reason, the auto_set_step option can be used to automatically compute the step where the total error is minimized.

discipline.check_jacobian(
    input_data={"x": array([0.0]), "y": array([1.0])},
    derr_approx=discipline.ApproximationMode.CENTERED_DIFFERENCES,
    show=True,
    plot_result=True,
    auto_set_step=True,
)

    INFO - 16:23:00: Set optimal step for finite differences. Estimated approximation errors =
    INFO - 16:23:00: [nan nan]
   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp f/d x is wrong by nan%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[nan]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[0.73575888]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[nan]]
   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp f/d y is wrong by nan%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[nan]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[0.]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[nan]]
   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp g/d x is wrong by nan%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[nan]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[0.]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[nan]]
   ERROR - 16:23:00: BuggedDiscipline Jacobian: dp g/d y is wrong by nan%.
    INFO - 16:23:00: Approximate jacobian =
    INFO - 16:23:00: [[nan]]
    INFO - 16:23:00: Provided by linearize method =
    INFO - 16:23:00: [[2.]]
    INFO - 16:23:00: Difference of jacobians =
    INFO - 16:23:00: [[nan]]
    INFO - 16:23:00: Linearization of Discipline: BuggedDiscipline is wrong.

False