Propulsion

propulsion

Propulsion discipline for the Sobieski's SSBJ use case.

Classes

JAXSobieskiPropulsion

JAXSobieskiPropulsion()

Bases: BaseJAXSobieskiDiscipline

Propulsion discipline for the Sobieski's SSBJ use case.

Initialize the JAXDiscipline.

Source code in src/gemseo_jax/problems/sobieski/propulsion.py

def __init__(self) -> None:  # noqa: D107
    super().__init__()
    self.__s_initial = array([
        self.mach_initial,
        self.h_initial,
        self.throttle_initial,
    ])
    self.__flag_temp = array([[0.95, 1.0, 1.1], [1.05, 1.0, 0.9], [0.95, 1.0, 1.1]])
    self.__bound_temp = array([0.25, 0.25, 0.25])
    self.__throttle_coeff = 16168.6
    self.__sfc_coeff = array(
        [
            1.13238425638512,
            1.53436586044561,
            -0.00003295564466,
            -0.00016378694115,
            -0.31623315541888,
            0.00000410691343,
            -0.00005248000590,
            -0.00000000008574,
            0.00000000190214,
            0.00000001059951,
        ],
    )
    self.__thua_coeff = array(
        [
            11483.7822254806,
            10856.2163466548,
            -0.5080237941,
            3200.157926969,
            -0.1466251679,
            0.0000068572,
        ],
    )
    self.default_input_data["c_3"] = np_array([self.constants[3]])
    self.__a0_g31 = 1.0
    self.__ai_g31 = array([0.3, -0.3, 0.3])
    self.__aij_g31 = array(
        [
            [0.2, 0.0794, 0.16304],
            [0.0794, -0.2, -0.12714],
            [0.16304, -0.12714, 0.2],
        ],
    )

Attributes

jax_out_func instance-attribute

jax_out_func: Callable[[DataType], DataType] = function

The JAX function to compute the outputs from the inputs.

Classes

DifferentiationMethod

Bases: StrEnum

The method to compute the Jacobian.

Functions

add_differentiated_inputs

add_differentiated_inputs(
    input_names: Iterable[str] = (),
) -> None

Add the inputs with respect to which to differentiate the outputs.

The inputs that do not represent continuous numbers are filtered out.

Parameters:

• input_names (Iterable[str], default: () ) –

  The input variables with respect to which to differentiate the outputs. If empty, use all the inputs.

Raises:

• ValueError –

  When an input name is not the name of a discipline input.

Notes

The Jacobian is also filtered to view non-differentiated static.

Source code in src/gemseo_jax/jax_discipline.py

def add_differentiated_inputs(
    self,
    input_names: Iterable[str] = (),
) -> None:
    """
    Notes:
        The Jacobian is also filtered to view non-differentiated static.
    """  # noqa: D205, D212, D415
    old_differentiated_inputs = self._differentiated_input_names.copy()
    super().add_differentiated_inputs(input_names=input_names)
    refilter = any(
        input_name not in old_differentiated_inputs
        for input_name in self._differentiated_input_names
    )
    if refilter:
        self._filter_jacobian()

add_differentiated_outputs

add_differentiated_outputs(
    output_names: Iterable[str] = (),
) -> None

Add the outputs to be differentiated.

The outputs that do not represent continuous numbers are filtered out.

Parameters:

• output_names (Iterable[str], default: () ) –

  The outputs to be differentiated. If empty, use all the outputs.

Raises:

• ValueError –

  When an output name is not the name of a discipline output.

Notes

The Jacobian is also filtered to view non-differentiated static.

Source code in src/gemseo_jax/jax_discipline.py

def add_differentiated_outputs(
    self,
    output_names: Iterable[str] = (),
) -> None:
    """
    Notes:
        The Jacobian is also filtered to view non-differentiated static.
    """  # noqa: D205, D212, D415
    old_differentiated_outputs = self._differentiated_output_names.copy()
    super().add_differentiated_outputs(output_names=output_names)
    refilter = any(
        output_name not in old_differentiated_outputs
        for output_name in self._differentiated_output_names
    )
    if refilter:
        self._filter_jacobian()

compile_jit

compile_jit(pre_run: bool = True) -> None

Apply jit compilation over function and jacobian.

Parameters:

• pre_run (bool, default: True ) –

  Whether to call jitted callables once to trigger compilation and log times.

Source code in src/gemseo_jax/jax_discipline.py

def compile_jit(
    self,
    pre_run: bool = True,
) -> None:
    """Apply jit compilation over function and jacobian.

    Args:
        pre_run: Whether to call jitted callables once to trigger compilation and
            log times.
    """
    self.jax_out_func = jit(self.jax_out_func)
    self.__jax_jac_func = jit(self.__jax_jac_func)
    if pre_run:
        self._jit_pre_run()