Block library

Notes on imported neural models (ONNX / ONNXJax / PyTorch / TensorFlow):

• x64 at import. import jaxonomy enables JAX 64-bit mode (jax_enable_x64) process-wide, so float32 artifacts see float64 inputs and silently compute in different arithmetic than they were trained and validated in. Cast explicitly at block boundaries: cast_outputs_to_dtype="float32" on the block, x.astype(jnp.float32) on upstream signals.
• Discrete-time policies need a ZeroOrderHold. A sample-and-hold controller exported from a discrete-time training loop (torch/NEUROMANCER-style) is otherwise re-evaluated at every ODE solver stage, silently destroying step-for-step parity with the exporting framework. Follow the policy block with ZeroOrderHold(dt=ts) and pin the step grid with SimulatorOptions(max_major_step_length=ts, max_minor_step_size=ts) — with that, closed-loop parity is ~4e-8 over 400 steps on the two-tank benchmark.

jaxonomy.library

Saturate = _inject_saturate_limit_kwarg(Saturate) module-attribute

Clip the input signal to a specified range.

Given an input signal u and upper and lower limits ulim and llim, the output signal is:

    y = max(llim, min(ulim, u))

where max and min are the element-wise maximum and minimum functions. This is equivalent to y = clip(u, llim, ulim).

Optionally, the block can also be configured with "dynamic" limits, which will add input ports for time-varying upper and lower limits.

Input ports

(0) The input signal. (1) The upper limit, if dynamic limits are enabled. (2) The lower limit, if dynamic limits are enabled. (Will be indexed as 1 if dynamic upper limits are not enabled.)

Output ports

(0) The clipped output signal.

Parameters:

Name	Type	Description	Default
upper_limit		The upper limit of the input signal. Default is np.inf.	required
enable_dynamic_upper_limit		If True, then the upper limit can be set by an external signal. Default is False.	required
lower_limit		The lower limit of the input signal. Default is -np.inf.	required
enable_dynamic_lower_limit		If True, then the lower limit can be set by an external signal. Default is False.	required
limit		T-115-followup-saturate-symmetric-kwarg: shorthand for the symmetric case. Saturate(limit=L) expands to Saturate(upper_limit=+L, lower_limit=-L). Mutually exclusive with explicit upper_limit / lower_limit and with the dynamic-limit flags. L must be a positive finite scalar.	required

Events

The block will trigger an event when the input signal crosses either the upper or lower limit. For example, if the block is configured with static upper and lower limits and the input signal crosses the upper limit, then a zero-crossing event will be triggered.

T-115-followup-mode-flag

The mode kwarg unifies the smooth (differentiable) variant previously exposed as :class:SoftSaturate. mode="hard" (default) is byte-equivalent to the legacy behavior, including zero-crossing event declaration. mode="smooth" dispatches to :func:soft_saturate and does not declare zero-crossing events (the smooth output has no discontinuity for the solver to catch). The smooth path requires finite upper_limit / lower_limit and a sharpness > 0 (defaults to 10.0).

Abs

Bases: FeedthroughBlock

Output the absolute value of the input signal.

Input ports

None

Output ports

(0) The absolute value of the input signal.

Events

An event is triggered when the output changes from positive to negative or vice versa.

Source code in jaxonomy/library/math_ops.py

class Abs(FeedthroughBlock):
    """Output the absolute value of the input signal.

    Input ports:
        None

    Output ports:
        (0) The absolute value of the input signal.

    Events:
        An event is triggered when the output changes from positive to negative
        or vice versa.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(npa.abs, *args, **kwargs)

    def _zero_crossing(self, _time, _state, u):
        return u

    def initialize_static_data(self, context):
        # Add a zero-crossing event so ODE solvers can't try to integrate
        # through a discontinuity. For efficiency, only do this if the output is
        # fed to an ODE.
        if not self.has_zero_crossing_events and is_discontinuity(self.output_ports[0]):
            self.declare_zero_crossing(self._zero_crossing, direction="crosses_zero")

        return super().initialize_static_data(context)

Adder

Bases: ReduceBlock

Computes the sum/difference of the input.

The add/subtract operation can be switched by setting the operators parameter. For example, a 3-input block specified as Adder(3, operators="+-+") would add the first and third inputs and subtract the second input.

Input ports

(0..n_in-1) The input signals to add/subtract.

Output ports

(0) The sum/difference of the input signals.

Source code in jaxonomy/library/math_ops.py

class Adder(ReduceBlock):
    """Computes the sum/difference of the input.

    The add/subtract operation can be switched by setting the `operators` parameter.
    For example, a 3-input block specified as `Adder(3, operators="+-+")` would add
    the first and third inputs and subtract the second input.

    Input ports:
        (0..n_in-1) The input signals to add/subtract.

    Output ports:
        (0) The sum/difference of the input signals.
    """

    @parameters(static=["operators"])
    def __init__(self, n_in, *args, operators=None, dtype=None, **kwargs):
        # T-038a-followup-other-blocks: per-block dtype override; stored
        # outside the @parameters list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.  Lazy
        # import avoids a circular dep — ``precision.py`` does not
        # import block code.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(n_in, None, *args, **kwargs)

    def initialize(self, operators):
        if operators is not None and any(char not in {"+", "-"} for char in operators):
            raise BlockParameterError(
                message=f"Adder block {self.name} has invalid operators {operators}. Can only contain '+' and '-'",
                system=self,
                parameter_name="operators",
            )

        if operators is None:
            _func = sum
        else:
            signs = [1 if op == "+" else -1 for op in operators]

            def _func(inputs):
                signed_inputs = [s * u for (s, u) in zip(signs, inputs)]
                return sum(signed_inputs)

        if self._dtype is not None:
            # T-038a-followup-other-blocks: wrap the reducer so the final
            # sum is cast to the per-block dtype.
            _inner = _func
            _dtype = self._dtype

            def _func(inputs):
                return npa.asarray(_inner(inputs)).astype(_dtype)

        self.replace_op(_func)

Arithmetic

Bases: ReduceBlock

Performs addition, subtraction, multiplication, and division on the input.

The arithmetic operation is determined by setting the operators parameter. For example, a 4-input block specified as Arithmetic(4, operators="+-*/") would: - Add the first input, - Subtract the second input, - Multiply the third input, - Divide by the fourth input.

Input ports

(0..n_in-1) The input signals for the specified arithmetic operations.

Output ports

(0) The result of the specified arithmetic operations on the input signals.

Source code in jaxonomy/library/math_ops.py

class Arithmetic(ReduceBlock):
    """Performs addition, subtraction, multiplication, and division on the input.

    The arithmetic operation is determined by setting the `operators` parameter.
    For example, a 4-input block specified as `Arithmetic(4, operators="+-*/")` would:
        - Add the first input,
        - Subtract the second input,
        - Multiply the third input,
        - Divide by the fourth input.

    Input ports:
        (0..n_in-1) The input signals for the specified arithmetic operations.

    Output ports:
        (0) The result of the specified arithmetic operations on the input signals.

    """

    @parameters(static=["operators"])
    def __init__(self, n_in, *args, operators=None, **kwargs):
        super().__init__(n_in, None, *args, **kwargs)

    def initialize(self, operators):
        if operators is not None and any(
            char not in {"+", "-", "*", "/"} for char in operators
        ):
            raise BlockParameterError(
                message=f"Arithmetic block {self.name} has invalid operators {operators}. Can only contain '+', '-', '*', '/'.",
                system=self,
                parameter_name="operators",
            )

        ops = {
            "+": npa.add,
            "-": npa.subtract,
            "/": npa.divide,
            "*": npa.multiply,
        }

        def evaluate_expression(operands, operators):
            operands = operands[:]
            operators = operators[:]

            # Handle multiplication and division
            while "*" in operators or "/" in operators:
                for op in ("*", "/"):
                    if op in operators:
                        index = operators.index(op)
                        result = ops[op](operands[index], operands[index + 1])
                        operands = operands[:index] + [result] + operands[index + 2 :]
                        operators = operators[:index] + operators[index + 1 :]

            # Handle addition and subtraction
            while "+" in operators or "-" in operators:
                for op in ("-", "+"):
                    if op in operators:
                        index = operators.index(op)
                        result = ops[op](operands[index], operands[index + 1])
                        operands = operands[:index] + [result] + operands[index + 2 :]
                        operators = operators[:index] + operators[index + 1 :]

            return operands[0]

        def _func(inputs):
            inputs = list(inputs)
            if operators[0] == "/":
                inputs[0] = 1.0 / inputs[0]
            if operators[0] == "-":
                inputs[0] = -inputs[0]
            ops = operators[1:]
            return evaluate_expression(inputs, ops)

        self.replace_op(_func)

AugmentedStateEKF

Bases: LeafSystem

Extended Kalman Filter with augmented state for joint state and parameter estimation.

The block estimates both the plant state x and unknown parameters θ online by augmenting the state vector:

.. code-block:: text

z = [x; θ]   (shape: nx + n_params)

with augmented dynamics and observation:

.. code-block:: text

z[n+1] = f_aug(z[n], u[n]) + noise
       = [f(x[n], u[n], θ[n]); θ[n]]   + [G_x w_x; w_θ]
y[n]   = h(x[n], u[n], θ[n]) + v[n]

E(w_x)  = E(w_θ) = E(v) = 0
Cov(w_x) = Q_x,  Cov(w_θ) = Q_θ,  Cov(v) = R

Parameters θ follow a random-walk model (θ[n+1] = θ[n] + w_θ). Setting Q_theta small makes parameters quasi-constant; increasing it allows tracking of slowly time-varying parameters.

All Jacobians are computed automatically via jax.jacfwd.

                +--------------------+
--- u[n] ------>|                    |----> x_hat[n]
                |  AugmentedStateEKF |
--- y[n] ------>|                    |----> theta_hat[n]
                +--------------------+

Input ports

(0) u : control vector at timestep n, shape (nu,) or scalar (1) y : measurement vector at timestep n, shape (ny,) or scalar

Output ports

(0) x_hat : state estimate, shape (nx,) (1) theta_hat : parameter estimate, shape (n_params,)

Parameters:

Name	Type	Description	Default
dt		float Sampling period.	required
nx		int Dimension of the plant state x.	required
n_params		int Dimension of the parameter vector θ.	required
forward		Callable Discrete-time state transition: f(x, u, theta) -> x_next. Must be JAX-traceable.	required
observation		Callable Observation function: h(x, u, theta) -> y. Must be JAX-traceable.	required
G_x_func		Callable Process-noise input matrix for states: G_x(t) -> (nx, nw) array. Pass lambda t: jnp.eye(nx) for isotropic process noise.	required
Q_x_func		Callable Process-noise covariance for states: Q_x(t, x, u, theta) -> (nw, nw).	required
Q_theta		array_like Constant parameter diffusion covariance matrix (n_params, n_params). Small values → slow/no parameter drift.	required
R_func		Callable Measurement noise covariance: R(t) -> (ny, ny).	required
x_hat_0		array_like Initial state estimate, shape (nx,).	required
P_hat_0_x		array_like Initial state covariance, shape (nx, nx).	required
theta_hat_0		array_like Initial parameter estimate, shape (n_params,).	required
P_hat_0_theta		array_like Initial parameter covariance, shape (n_params, n_params).	required

Example::

import jax.numpy as jnp
from jaxonomy import library

# Simple first-order system:  x[n+1] = a*x[n] + b*u[n]
# where 'a' (decay rate) is unknown and must be estimated.

def forward(x, u, theta):
    a = theta[0]
    return jnp.array([a * x[0] + u[0]])

def observation(x, u, theta):
    return jnp.array([x[0]])

aekf = library.AugmentedStateEKF(
    dt=0.1,
    nx=1,
    n_params=1,
    forward=forward,
    observation=observation,
    G_x_func=lambda t: jnp.eye(1),
    Q_x_func=lambda t, x, u, th: jnp.array([[0.01]]),
    Q_theta=jnp.array([[1e-4]]),
    R_func=lambda t: jnp.array([[0.1]]),
    x_hat_0=jnp.zeros(1),
    P_hat_0_x=jnp.eye(1),
    theta_hat_0=jnp.zeros(1),
    P_hat_0_theta=jnp.eye(1),
)

Source code in jaxonomy/library/state_estimators/augmented_ekf.py

class AugmentedStateEKF(LeafSystem):
    """
    Extended Kalman Filter with augmented state for **joint** state and parameter
    estimation.

    The block estimates both the plant state *x* and unknown parameters *θ* online
    by augmenting the state vector:

    .. code-block:: text

        z = [x; θ]   (shape: nx + n_params)

    with augmented dynamics and observation:

    .. code-block:: text

        z[n+1] = f_aug(z[n], u[n]) + noise
               = [f(x[n], u[n], θ[n]); θ[n]]   + [G_x w_x; w_θ]
        y[n]   = h(x[n], u[n], θ[n]) + v[n]

        E(w_x)  = E(w_θ) = E(v) = 0
        Cov(w_x) = Q_x,  Cov(w_θ) = Q_θ,  Cov(v) = R

    Parameters *θ* follow a **random-walk** model (``θ[n+1] = θ[n] + w_θ``).
    Setting ``Q_theta`` small makes parameters quasi-constant; increasing it allows
    tracking of slowly time-varying parameters.

    All Jacobians are computed automatically via ``jax.jacfwd``.

    ```
                    +--------------------+
    --- u[n] ------>|                    |----> x_hat[n]
                    |  AugmentedStateEKF |
    --- y[n] ------>|                    |----> theta_hat[n]
                    +--------------------+
    ```

    Input ports:
        (0) u : control vector at timestep n, shape ``(nu,)`` or scalar
        (1) y : measurement vector at timestep n, shape ``(ny,)`` or scalar

    Output ports:
        (0) x_hat     : state estimate, shape ``(nx,)``
        (1) theta_hat : parameter estimate, shape ``(n_params,)``

    Parameters:
        dt : float
            Sampling period.
        nx : int
            Dimension of the plant state *x*.
        n_params : int
            Dimension of the parameter vector *θ*.
        forward : Callable
            Discrete-time state transition: ``f(x, u, theta) -> x_next``.
            Must be JAX-traceable.
        observation : Callable
            Observation function: ``h(x, u, theta) -> y``.
            Must be JAX-traceable.
        G_x_func : Callable
            Process-noise input matrix for states: ``G_x(t) -> (nx, nw)`` array.
            Pass ``lambda t: jnp.eye(nx)`` for isotropic process noise.
        Q_x_func : Callable
            Process-noise covariance for states: ``Q_x(t, x, u, theta) -> (nw, nw)``.
        Q_theta : array_like
            Constant parameter diffusion covariance matrix ``(n_params, n_params)``.
            Small values → slow/no parameter drift.
        R_func : Callable
            Measurement noise covariance: ``R(t) -> (ny, ny)``.
        x_hat_0 : array_like
            Initial state estimate, shape ``(nx,)``.
        P_hat_0_x : array_like
            Initial state covariance, shape ``(nx, nx)``.
        theta_hat_0 : array_like
            Initial parameter estimate, shape ``(n_params,)``.
        P_hat_0_theta : array_like
            Initial parameter covariance, shape ``(n_params, n_params)``.

    Example::

        import jax.numpy as jnp
        from jaxonomy import library

        # Simple first-order system:  x[n+1] = a*x[n] + b*u[n]
        # where 'a' (decay rate) is unknown and must be estimated.

        def forward(x, u, theta):
            a = theta[0]
            return jnp.array([a * x[0] + u[0]])

        def observation(x, u, theta):
            return jnp.array([x[0]])

        aekf = library.AugmentedStateEKF(
            dt=0.1,
            nx=1,
            n_params=1,
            forward=forward,
            observation=observation,
            G_x_func=lambda t: jnp.eye(1),
            Q_x_func=lambda t, x, u, th: jnp.array([[0.01]]),
            Q_theta=jnp.array([[1e-4]]),
            R_func=lambda t: jnp.array([[0.1]]),
            x_hat_0=jnp.zeros(1),
            P_hat_0_x=jnp.eye(1),
            theta_hat_0=jnp.zeros(1),
            P_hat_0_theta=jnp.eye(1),
        )
    """

    class DiscreteStateType(NamedTuple):
        """Internal filter state (both minus=predicted and plus=corrected estimates)."""

        z_hat_minus: npa.ndarray  # predicted   augmented state [nx+n_params]
        P_hat_minus: npa.ndarray  # predicted   augmented covariance
        z_hat_plus: npa.ndarray   # corrected   augmented state [nx+n_params]
        P_hat_plus: npa.ndarray   # corrected   augmented covariance

    @parameters(
        static=[
            "dt",
            "nx",
            "n_params",
            "forward",
            "observation",
            "G_x_func",
            "Q_x_func",
            "Q_theta",
            "R_func",
            "x_hat_0",
            "P_hat_0_x",
            "theta_hat_0",
            "P_hat_0_theta",
        ],
    )
    def __init__(
        self,
        dt,
        nx,
        n_params,
        forward,
        observation,
        G_x_func,
        Q_x_func,
        Q_theta,
        R_func,
        x_hat_0,
        P_hat_0_x,
        theta_hat_0,
        P_hat_0_theta,
        name=None,
        **kwargs,
    ):
        super().__init__(name=name, **kwargs)

        x_hat_0 = jnp.asarray(x_hat_0, dtype=float)
        P_hat_0_x = jnp.asarray(P_hat_0_x, dtype=float)
        theta_hat_0 = jnp.asarray(theta_hat_0, dtype=float)
        P_hat_0_theta = jnp.asarray(P_hat_0_theta, dtype=float)

        # Augmented initial state and covariance
        z_hat_0 = jnp.concatenate([x_hat_0, theta_hat_0])
        P_z_0 = jax.scipy.linalg.block_diag(P_hat_0_x, P_hat_0_theta)

        # Input ports
        self.u_in_index = self.declare_input_port(name="u")
        self.y_in_index = self.declare_input_port(name="y")

        # Internal discrete state
        self.declare_discrete_state(
            default_value=self.DiscreteStateType(
                z_hat_minus=z_hat_0,
                P_hat_minus=P_z_0,
                z_hat_plus=z_hat_0,
                P_hat_plus=P_z_0,
            ),
            as_array=False,
        )

        # Periodic update at each timestep
        self.declare_periodic_update(
            self._update,
            period=dt,
            offset=0.0,
        )

        # Build dependency tickets for feedthrough outputs
        u_ticket = self.input_ports[self.u_in_index].ticket
        y_ticket = self.input_ports[self.y_in_index].ticket
        prereqs = [DependencyTicket.xd, u_ticket, y_ticket]
        required_inputs = [self.u_in_index, self.y_in_index]

        # Output port 0: x_hat  (state estimate)
        self.declare_output_port(
            self._output_x_hat,
            period=dt,
            offset=0.0,
            default_value=x_hat_0,
            name="x_hat",
            requires_inputs=required_inputs,
            prerequisites_of_calc=prereqs,
        )

        # Output port 1: theta_hat  (parameter estimate)
        self.declare_output_port(
            self._output_theta_hat,
            period=dt,
            offset=0.0,
            default_value=theta_hat_0,
            name="theta_hat",
            requires_inputs=required_inputs,
            prerequisites_of_calc=prereqs,
        )

    def initialize(
        self,
        dt,
        nx,
        n_params,
        forward,
        observation,
        G_x_func,
        Q_x_func,
        Q_theta,
        R_func,
        x_hat_0,
        P_hat_0_x,
        theta_hat_0,
        P_hat_0_theta,
    ):
        """Called at context-creation time to store resolved parameters."""
        self.nx = nx
        self.np = n_params
        self.nz = nx + n_params

        self.forward = forward
        self.observation = observation
        self.G_x_func = G_x_func
        self.Q_x_func = Q_x_func
        self.Q_theta = jnp.asarray(Q_theta, dtype=float)
        self.R_func = R_func

        # Determine ny from R matrix shape
        self.ny = self.R_func(0.0).shape[0]

        # Build augmented dynamics and observation (closures over nx, n_params)
        _nx = nx

        def f_aug(z, u):
            x, theta = z[:_nx], z[_nx:]
            x_next = forward(x, u, theta)
            return jnp.concatenate([x_next, theta])

        def h_aug(z, u):
            x, theta = z[:_nx], z[_nx:]
            return jnp.atleast_1d(observation(x, u, theta))

        self.f_aug = f_aug
        self.h_aug = h_aug

        # Jacobian functions for the augmented system
        self.jac_f_aug = jax.jacfwd(f_aug)   # ∂f_aug/∂z  [nz × nz]
        self.jac_h_aug = jax.jacfwd(h_aug)   # ∂h_aug/∂z  [ny × nz]

        self.eye_z = jnp.eye(self.nz)

    # ──────────────────────────────────────────────────────────────────────────
    # EKF correct step
    # ──────────────────────────────────────────────────────────────────────────

    def _correct(self, time, z_hat_minus, P_hat_minus, u, y):
        """Update estimate using current measurement."""
        y = jnp.atleast_1d(jnp.asarray(y, dtype=float))
        u = jnp.atleast_1d(jnp.asarray(u, dtype=float))

        C = self.jac_h_aug(z_hat_minus, u).reshape((self.ny, self.nz))
        R = self.R_func(time)

        # Kalman gain
        S = C @ P_hat_minus @ C.T + R
        K = P_hat_minus @ C.T @ jnp.linalg.inv(S)

        # State update
        innovation = y - self.h_aug(z_hat_minus, u)
        z_hat_plus = z_hat_minus + K @ innovation

        # Covariance update
        P_hat_plus = (self.eye_z - K @ C) @ P_hat_minus

        return z_hat_plus, P_hat_plus

    # ──────────────────────────────────────────────────────────────────────────
    # EKF propagate step
    # ──────────────────────────────────────────────────────────────────────────

    def _propagate(self, time, z_hat_plus, P_hat_plus, u):
        """Predict next state from corrected estimate."""
        u = jnp.atleast_1d(jnp.asarray(u, dtype=float))

        # Augmented state Jacobian
        A = self.jac_f_aug(z_hat_plus, u).reshape((self.nz, self.nz))

        # Augmented process noise covariance
        #   G_aug = block_diag(G_x, I_theta)
        #   Q_aug = block_diag(G_x Q_x G_x^T, Q_theta)
        x_plus, theta_plus = z_hat_plus[:self.nx], z_hat_plus[self.nx:]
        G_x = self.G_x_func(time)
        Q_x = self.Q_x_func(time, x_plus, u, theta_plus)
        GQG_x = G_x @ Q_x @ G_x.T
        GQG_aug = jax.scipy.linalg.block_diag(GQG_x, self.Q_theta)

        # Propagate
        z_hat_minus = self.f_aug(z_hat_plus, u)
        P_hat_minus = A @ P_hat_plus @ A.T + GQG_aug

        return z_hat_minus, P_hat_minus

    # ──────────────────────────────────────────────────────────────────────────
    # Periodic state update
    # ──────────────────────────────────────────────────────────────────────────

    def _update(self, time, state, *inputs, **params):
        u, y = inputs
        z_hat_minus = state.discrete_state.z_hat_minus
        P_hat_minus = state.discrete_state.P_hat_minus

        z_hat_plus, P_hat_plus = self._correct(time, z_hat_minus, P_hat_minus, u, y)
        z_hat_minus_next, P_hat_minus_next = self._propagate(
            time, z_hat_plus, P_hat_plus, u
        )

        return self.DiscreteStateType(
            z_hat_minus=z_hat_minus_next,
            P_hat_minus=P_hat_minus_next,
            z_hat_plus=z_hat_plus,
            P_hat_plus=P_hat_plus,
        )

    # ──────────────────────────────────────────────────────────────────────────
    # Output callbacks  (feedthrough: recompute correction with current inputs)
    # ──────────────────────────────────────────────────────────────────────────

    def _output_x_hat(self, time, state, *inputs, **params):
        u, y = inputs
        z_hat_minus = state.discrete_state.z_hat_minus
        P_hat_minus = state.discrete_state.P_hat_minus
        z_hat_plus, _ = self._correct(time, z_hat_minus, P_hat_minus, u, y)
        return z_hat_plus[:self.nx]

    def _output_theta_hat(self, time, state, *inputs, **params):
        u, y = inputs
        z_hat_minus = state.discrete_state.z_hat_minus
        P_hat_minus = state.discrete_state.P_hat_minus
        z_hat_plus, _ = self._correct(time, z_hat_minus, P_hat_minus, u, y)
        return z_hat_plus[self.nx:]

DiscreteStateType

Bases: NamedTuple

Internal filter state (both minus=predicted and plus=corrected estimates).

Source code in jaxonomy/library/state_estimators/augmented_ekf.py

class DiscreteStateType(NamedTuple):
    """Internal filter state (both minus=predicted and plus=corrected estimates)."""

    z_hat_minus: npa.ndarray  # predicted   augmented state [nx+n_params]
    P_hat_minus: npa.ndarray  # predicted   augmented covariance
    z_hat_plus: npa.ndarray   # corrected   augmented state [nx+n_params]
    P_hat_plus: npa.ndarray   # corrected   augmented covariance

initialize(dt, nx, n_params, forward, observation, G_x_func, Q_x_func, Q_theta, R_func, x_hat_0, P_hat_0_x, theta_hat_0, P_hat_0_theta)

Called at context-creation time to store resolved parameters.

Source code in jaxonomy/library/state_estimators/augmented_ekf.py

def initialize(
    self,
    dt,
    nx,
    n_params,
    forward,
    observation,
    G_x_func,
    Q_x_func,
    Q_theta,
    R_func,
    x_hat_0,
    P_hat_0_x,
    theta_hat_0,
    P_hat_0_theta,
):
    """Called at context-creation time to store resolved parameters."""
    self.nx = nx
    self.np = n_params
    self.nz = nx + n_params

    self.forward = forward
    self.observation = observation
    self.G_x_func = G_x_func
    self.Q_x_func = Q_x_func
    self.Q_theta = jnp.asarray(Q_theta, dtype=float)
    self.R_func = R_func

    # Determine ny from R matrix shape
    self.ny = self.R_func(0.0).shape[0]

    # Build augmented dynamics and observation (closures over nx, n_params)
    _nx = nx

    def f_aug(z, u):
        x, theta = z[:_nx], z[_nx:]
        x_next = forward(x, u, theta)
        return jnp.concatenate([x_next, theta])

    def h_aug(z, u):
        x, theta = z[:_nx], z[_nx:]
        return jnp.atleast_1d(observation(x, u, theta))

    self.f_aug = f_aug
    self.h_aug = h_aug

    # Jacobian functions for the augmented system
    self.jac_f_aug = jax.jacfwd(f_aug)   # ∂f_aug/∂z  [nz × nz]
    self.jac_h_aug = jax.jacfwd(h_aug)   # ∂h_aug/∂z  [ny × nz]

    self.eye_z = jnp.eye(self.nz)

Backlash

Bases: LeafSystem

Hysteretic nonlinearity modelling mechanical slack / backlash.

A Backlash(width) block has a single discrete state last_output that tracks the most recent output value. At each sample tick the update rule is::

delta = u - last_output
if   delta >  width/2:  new_output = u - width/2
elif delta < -width/2:  new_output = u + width/2
else:                   new_output = last_output

Equivalently, the output "follows" the input only after the input has moved by more than width/2 from the last output; within that band the output sticks. This is the standard model for gearbox slack / actuator hysteresis: the input must take up the slack before the output moves.

Differentiability: the per-step update is expressed via npa.where on delta; the output is a continuous (piecewise-linear) function of both u and width, so jax.grad w.r.t. width is finite. The non-smooth "knee" at |delta| = width/2 has subgradient 0 (inside the band) or -sign(delta)/2 (outside) w.r.t. width -- both finite, as required.

Input ports

(0) The driving input signal.

Output ports

(0) The hysteretic output signal.

Parameters:

Name	Type	Description	Default
width		Positive scalar; total hysteresis band width. width=0 recovers y = u (no hysteresis).	1.0
dt		Periodic update sample time. Required: this is a discrete block.	0.01
initial_output		Initial value of the output / discrete state. Default 0.0.	0.0

Notes

For an exact match against a canonical discrete-time "Backlash" block the discrete sample time must match. For a continuous hysteresis approximation, choose dt much smaller than the dominant input timescale; the block then tracks the input modulo the width/2 slack with single-step latency.

Source code in jaxonomy/library/nonlinearities.py

class Backlash(LeafSystem):
    """Hysteretic nonlinearity modelling mechanical slack / backlash.

    A ``Backlash(width)`` block has a single discrete state ``last_output``
    that tracks the most recent output value. At each sample tick the
    update rule is::

        delta = u - last_output
        if   delta >  width/2:  new_output = u - width/2
        elif delta < -width/2:  new_output = u + width/2
        else:                   new_output = last_output

    Equivalently, the output "follows" the input only after the input has
    moved by more than ``width/2`` from the last output; within that band
    the output sticks. This is the standard model for gearbox slack /
    actuator hysteresis: the input must take up the slack before the
    output moves.

    Differentiability: the per-step update is expressed via ``npa.where``
    on ``delta``; the output is a continuous (piecewise-linear) function
    of both ``u`` and ``width``, so ``jax.grad`` w.r.t. ``width`` is
    finite. The non-smooth "knee" at ``|delta| = width/2`` has subgradient
    ``0`` (inside the band) or ``-sign(delta)/2`` (outside) w.r.t.
    ``width`` -- both finite, as required.

    Input ports:
        (0) The driving input signal.

    Output ports:
        (0) The hysteretic output signal.

    Parameters:
        width: Positive scalar; total hysteresis band width. ``width=0``
            recovers ``y = u`` (no hysteresis).
        dt: Periodic update sample time. Required: this is a discrete
            block.
        initial_output: Initial value of the output / discrete state.
            Default ``0.0``.

    Notes:
        For an exact match against a canonical discrete-time "Backlash" block
        the discrete sample time must match. For a *continuous* hysteresis
        approximation, choose ``dt`` much smaller than the dominant input
        timescale; the block then tracks the input modulo the
        ``width/2`` slack with single-step latency.
    """

    @parameters(static=["dt"], dynamic=["width", "initial_output"])
    def __init__(self, width=1.0, dt=0.01, initial_output=0.0, **kwargs):
        super().__init__(**kwargs)
        if width < 0:
            raise BlockParameterError(
                message=(
                    f"Backlash block {self.name}: width must be >= 0, "
                    f"got {width}."
                ),
                system=self,
                parameter_name="width",
            )
        if dt <= 0:
            raise BlockParameterError(
                message=(
                    f"Backlash block {self.name}: dt must be > 0, got {dt}."
                ),
                system=self,
                parameter_name="dt",
            )
        self.dt = dt
        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, width=1.0, initial_output=0.0, dt=None):
        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=self.dt,
        )

        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=0.0,
            requires_inputs=False,
            prerequisites_of_calc=[DependencyTicket.xd],
            default_value=initial_output,
        )

    def reset_default_values(self, width=1.0, initial_output=0.0):
        self.declare_discrete_state(default_value=initial_output)
        self.configure_output_port_default_value(
            self._output_port_idx, initial_output
        )

    @staticmethod
    def _apply(last_output, u, width):
        """Pure backlash kernel.

        Splitting this out makes it directly testable / gradient-friendly
        outside the LeafSystem update path (see the T-115-followup tests).
        """
        half = width / 2.0
        delta = u - last_output
        # Output snaps to the edge of the band the input has crossed, or
        # stays put if the input is still inside the band.
        upper_edge = u - half  # active when delta > +half
        lower_edge = u + half  # active when delta < -half
        # First branch: above band -> follow at upper edge.
        # Second branch: below band -> follow at lower edge.
        # Else: still inside band -> hold previous output.
        return npa.where(
            delta > half,
            upper_edge,
            npa.where(delta < -half, lower_edge, last_output),
        )

    def _update(self, _time, state, u, **params):
        return self._apply(state.discrete_state, u, params["width"])

    def _output(self, _time, state, **_params):
        return state.discrete_state

BatteryCell

Bases: LeafSystem

Dynamic electro-checmical Li-ion cell model.

Based on Tremblay and Dessaint (2009).

By using appropriate parameters, the cell model can be used to model a battery pack with the assumption that the cells of the pack behave as a single unit.

Parameters E0, K, A, below are abstract parameters used in the model presented in the reference paper. As described in the reference paper, these parameters can be extracted from typical cell manufacturer datasheets; see section 3. Section 3 also provides a table of example values for these parameters.

Input ports

(0) The current (A) flowing through the cell. Positive is discharge.

Output ports

(0) The voltage across the cell terminals (V) (1) The state of charge of the cell (normalized between 0 and 1)

Parameters:

Name	Type	Description	Default
E0	float	described as "battery constant voltage (V)" by the reference paper.	3.366
K	float	described as "polarization constant (V/Ah)" by the reference paper.	0.0076
Q	float	battery capacity in Ah	2.3
R	float	internal resistance (Ohms)	0.01
A	float	described as "exponential zone amplitude (V)" by the reference paper.	0.26422
B	float	described as "exponential zone time constant inverse (1/Ah)" by the reference paper.	26.5487
initial_SOC	float	initial state of charge, normalized between 0 and 1.	1.0

Source code in jaxonomy/library/battery_cell.py

class BatteryCell(LeafSystem):
    """Dynamic electro-checmical Li-ion cell model.

    Based on [Tremblay and Dessaint (2009)](https://doi.org/10.3390/wevj3020289).

    By using appropriate parameters, the cell model can be used to model a battery pack
    with the assumption that the cells of the pack behave as a single unit.

    Parameters E0, K, A, below are abstract parameters used in the model presented in
    the reference paper. As described in the reference paper, these parameters can be
    extracted from typical cell manufacturer datasheets; see section 3. Section 3 also
    provides a table of example values for these parameters.

    Input ports:
        (0) The current (A) flowing through the cell. Positive is discharge.

    Output ports:
        (0) The voltage across the cell terminals (V)
        (1) The state of charge of the cell (normalized between 0 and 1)

    Parameters:
        E0: described as "battery constant voltage (V)" by the reference paper.
        K: described as "polarization constant (V/Ah)" by the reference paper.
        Q: battery capacity in Ah
        R: internal resistance (Ohms)
        A: described as "exponential zone amplitude (V)" by the reference paper.
        B:
            described as "exponential zone time constant inverse (1/Ah)" by the
            reference paper.
        initial_SOC: initial state of charge, normalized between 0 and 1.
    """

    class BatteryStateType(NamedTuple):
        soc: float
        i_star: float
        i_lb: float

    class FirstOrderFilter(NamedTuple):
        A: float
        B: float
        C: float

    @parameters(dynamic=["E0", "K", "Q", "R", "tau", "A", "B"], static=["initial_SOC"])
    def __init__(
        self,
        E0: float = 3.366,
        K: float = 0.0076,
        Q: float = 2.3,
        R: float = 0.01,
        tau: float = 30.0,
        A: float = 0.26422,
        B: float = 26.5487,
        initial_SOC: float = 1.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.declare_input_port()  # Current flowing through the cell

        self.declare_output_port(
            self._voltage_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            name="voltage",
        )

        self.declare_output_port(
            self._soc_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            name="soc",
        )

    def initialize(self, E0, K, Q, R, tau, A, B, initial_SOC):
        # Filter for input current
        self.current_filter = self.FirstOrderFilter(-0.05, 1.0, 0.05)

        # Filter for loop-breaker
        self.lb_filter = self.FirstOrderFilter(-10.0, 1.0, 10.0)

        initial_state = self.BatteryStateType(
            soc=initial_SOC,
            i_star=0.0,  # Filtered input current
            i_lb=0.0,  # Filtered current for loop-breaking
        )

        self.declare_continuous_state(
            default_value=initial_state,
            as_array=False,
            ode=self._ode,
        )

    def _ode(self, _time, state, *inputs, **parameters) -> BatteryStateType:
        xc = state.continuous_state
        Q = parameters["Q"]

        (u,) = inputs

        soc_der_unsat = -u / (Q * Ah_to_As)

        # SoC must be between 0 and 1
        llim_violation = (xc.soc <= 0.0) & (soc_der_unsat < 0.0)
        ulim_violation = (xc.soc >= 1.0) & (soc_der_unsat > 0.0)

        # Saturated time derivative
        soc_der = npa.where(llim_violation | ulim_violation, 0.0, soc_der_unsat)

        # Derivative of istar, the filtered current signal
        i_star_der = self.current_filter.A * xc.i_star + self.current_filter.B * u

        # Derivative of ilb, the filtered current signal for loop-breaking
        i_lb_der = self.lb_filter.A * xc.i_lb + self.lb_filter.B * u

        return self.BatteryStateType(
            soc=soc_der,
            i_star=i_star_der,
            i_lb=i_lb_der,
        )

    def _voltage_output(self, _time, state, *_inputs, **parameters) -> Array:
        E0 = parameters["E0"]
        Q = parameters["Q"]
        K = parameters["K"]
        A = parameters["A"]
        B = parameters["B"]
        R = parameters["R"]
        xc = state.continuous_state

        # Filtered input current
        i_star = self.current_filter.C * xc.i_star

        # Loop-breaking current
        i_lb = self.lb_filter.C * xc.i_lb

        # Apply limits to state of charge
        soc = npa.clip(xc.soc, 0.0, 1.0)

        # Undo normalization by Q - this is ∫i*dt, the integral of current
        i_int = Q * (1 - soc)

        chg_mode_Q_gain = 0.1
        vdyn_den = npa.where(i_star >= 0, Q - i_int, i_int + chg_mode_Q_gain * Q)
        vdyn = i_star * K * Q / vdyn_den

        vbatt_ulim = 2 * E0  # Reasonable upper limit on battery voltage
        vbatt_presat = (
            E0 - R * i_lb - i_int * K * Q / (Q - i_int) + A * npa.exp(-B * i_int) - vdyn
        )
        return npa.clip(vbatt_presat, 0.0, vbatt_ulim)

    def _soc_output(self, _time, state, *_inputs, **_parameters) -> Array:
        return state.continuous_state.soc

BusCreator

Bases: LeafSystem

Pack n = len(field_names) signals into a single named-bus output.

The output is a collections.namedtuple (named "Bus") whose fields are exactly field_names in declaration order. NamedTuples are first-class JAX pytrees, so the bus signal flows through jax.jit, vmap, and grad without any extra registration.

Pair with :class:BusSelector to pull individual fields back out downstream. Use this when signals share a logical group identity (e.g. ("position", "velocity", "acceleration") for a vehicle state bus) and you would rather refer to them by name than by positional Mux/Demux index.

Parameters:

Name	Type	Description	Default
field_names		Tuple/list of strings — one name per input port, in declaration order. Must be unique, valid Python identifiers (NamedTuple constraint).	required
field_units		Optional mapping from field name to :class:Unit (T-117-followup-bus-units). When supplied, each input port is declared with the corresponding units= and the output port carries a :class:BusUnit so downstream :class:BusSelector blocks can recover per-field units at connect time. When None (the default), input/output ports carry no unit metadata — byte-equivalent to the T-117-fu-bus-namedtuple shipping behaviour.	None
field_shapes		Optional mapping from field name to a JAX-style shape tuple (T-117-followup-bus-array). When supplied, the named field carries an array of the declared shape rather than a scalar — useful for grouping e.g. an 8-element thermocouple readout under a single "sensors" slot without manually muxing. Fields not listed default to scalar shape (). Default None is byte-equivalent to the T-117-fu-bus-namedtuple all-scalar behaviour.	None

Input ports

(0..n-1) — one port per field name; values are packed into the corresponding slot of the output NamedTuple.

Output ports

(0) — the bus, a NamedTuple with fields field_names.

Source code in jaxonomy/library/routing.py

class BusCreator(LeafSystem):
    """Pack ``n = len(field_names)`` signals into a single named-bus output.

    The output is a ``collections.namedtuple`` (named ``"Bus"``) whose
    fields are exactly ``field_names`` in declaration order. NamedTuples
    are first-class JAX pytrees, so the bus signal flows through
    ``jax.jit``, ``vmap``, and ``grad`` without any extra registration.

    Pair with :class:`BusSelector` to pull individual fields back out
    downstream. Use this when signals share a logical group identity
    (e.g. ``("position", "velocity", "acceleration")`` for a vehicle
    state bus) and you would rather refer to them by name than by
    positional ``Mux``/``Demux`` index.

    Parameters:
        field_names: Tuple/list of strings — one name per input port,
            in declaration order. Must be unique, valid Python
            identifiers (NamedTuple constraint).
        field_units: Optional mapping from field name to :class:`Unit`
            (T-117-followup-bus-units). When supplied, each input port
            is declared with the corresponding ``units=`` and the
            output port carries a :class:`BusUnit` so downstream
            :class:`BusSelector` blocks can recover per-field units at
            connect time. When ``None`` (the default), input/output
            ports carry no unit metadata — byte-equivalent to the
            T-117-fu-bus-namedtuple shipping behaviour.
        field_shapes: Optional mapping from field name to a JAX-style
            shape tuple (T-117-followup-bus-array). When supplied, the
            named field carries an array of the declared shape rather
            than a scalar — useful for grouping e.g. an 8-element
            thermocouple readout under a single ``"sensors"`` slot
            without manually muxing. Fields not listed default to
            scalar shape ``()``. Default ``None`` is byte-equivalent
            to the T-117-fu-bus-namedtuple all-scalar behaviour.

    Input ports:
        ``(0..n-1)`` — one port per field name; values are packed into
        the corresponding slot of the output NamedTuple.

    Output ports:
        ``(0)`` — the bus, a NamedTuple with fields ``field_names``.
    """

    def __init__(
        self,
        field_names,
        *args,
        field_units=None,
        field_shapes=None,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

        # Validate up front so the failure mode is a clear ValueError on
        # construction, not a cryptic NamedTuple "Type names and field
        # names must be valid identifiers" error from deep inside the
        # collections module.
        field_names = tuple(field_names)
        if len(field_names) == 0:
            raise ValueError(
                "BusCreator requires at least one field name; "
                "got an empty tuple."
            )
        if len(set(field_names)) != len(field_names):
            raise ValueError(
                f"BusCreator field_names must be unique; got {field_names!r}."
            )
        for fname in field_names:
            if not isinstance(fname, str) or not fname.isidentifier():
                raise ValueError(
                    f"BusCreator field name {fname!r} is not a valid "
                    "Python identifier (required for NamedTuple fields)."
                )

        # T-117-followup-bus-units: optional per-field unit propagation.
        # ``field_units=None`` keeps the historic default-off path
        # byte-equivalent (no BusUnit on the output port, no units on
        # input ports). When provided, we validate the keys match
        # ``field_names`` exactly so the user gets a clear error before
        # any port is declared.
        self._bus_unit = None
        per_input_units: dict[str, object] = {fname: None for fname in field_names}
        if field_units is not None:
            from ..framework.units import BusUnit as _BusUnit, Unit as _Unit

            field_units_dict = dict(field_units)
            extra = set(field_units_dict) - set(field_names)
            missing = set(field_names) - set(field_units_dict)
            if extra or missing:
                raise ValueError(
                    "BusCreator field_units keys must match field_names "
                    f"exactly; got extra={sorted(extra)!r}, "
                    f"missing={sorted(missing)!r}."
                )
            for fname, u in field_units_dict.items():
                if not isinstance(u, _Unit):
                    raise TypeError(
                        f"BusCreator field_units[{fname!r}] must be a Unit "
                        f"instance, got {type(u).__name__}: {u!r}."
                    )
                per_input_units[fname] = u
            self._bus_unit = _BusUnit(fields=field_units_dict)

        # T-117-followup-bus-array: optional per-field array shapes.
        # ``field_shapes=None`` keeps every field scalar (``()``), which
        # is byte-equivalent to the T-117-fu-bus-namedtuple all-scalar
        # behaviour. When supplied, fields not listed in the mapping
        # default to scalar so users only spell out the array fields.
        per_field_shapes: dict[str, tuple] = {fname: () for fname in field_names}
        if field_shapes is not None:
            field_shapes_dict = dict(field_shapes)
            extra = set(field_shapes_dict) - set(field_names)
            if extra:
                raise ValueError(
                    "BusCreator field_shapes keys must be a subset of "
                    f"field_names; got unknown keys={sorted(extra)!r}."
                )
            for fname, shape in field_shapes_dict.items():
                shape_t = tuple(shape)
                for dim in shape_t:
                    if not isinstance(dim, int) or dim < 0:
                        raise ValueError(
                            f"BusCreator field_shapes[{fname!r}] must be a "
                            f"tuple of non-negative ints, got {shape!r}."
                        )
                per_field_shapes[fname] = shape_t

        self._field_names = field_names
        self._field_shapes = per_field_shapes
        self._bus_type = namedtuple("Bus", field_names)

        input_tickets = []
        for fname in field_names:
            idx = self.declare_input_port(
                name=fname, units=per_input_units[fname]
            )
            input_tickets.append(self.input_ports[idx].ticket)

        def _compute_bus(_time, _state, *inputs, **_params):
            # ``inputs`` is exactly the tuple of upstream values in port
            # order, matching ``field_names`` by construction.
            return self._bus_type(*inputs)

        # NOTE: we deliberately do NOT pass ``default_value=`` to
        # ``declare_output_port`` here — see the leaf_system.py code path
        # at line 931, which calls ``npa.array(default_value)`` and
        # would flatten our NamedTuple into a plain 1-D array (losing
        # the bus type). Letting the framework lazily compute the
        # default by calling ``_compute_bus`` on a dummy context yields
        # the correct NamedTuple-typed default. T-005 default-float64
        # is preserved transitively via the upstream ``Constant``
        # blocks' float64 zeros.
        self.declare_output_port(
            _compute_bus,
            name="bus",
            prerequisites_of_calc=input_tickets,
            requires_inputs=True,
            units=self._bus_unit,
        )

    @property
    def field_names(self) -> tuple[str, ...]:
        """The tuple of field names in declaration / port order."""
        return self._field_names

    @property
    def bus_type(self) -> type:
        """The underlying NamedTuple class for this bus."""
        return self._bus_type

    @property
    def bus_unit(self):
        """The compound :class:`BusUnit` for this bus, or ``None`` if
        ``field_units`` was not supplied at construction time."""
        return self._bus_unit

    @property
    def field_shapes(self) -> dict:
        """Mapping from field name to declared array shape tuple
        (T-117-followup-bus-array). Fields default to scalar ``()``
        when ``field_shapes`` is omitted at construction time."""
        return dict(self._field_shapes)

bus_type property

The underlying NamedTuple class for this bus.

bus_unit property

The compound :class:BusUnit for this bus, or None if field_units was not supplied at construction time.

field_names property

The tuple of field names in declaration / port order.

field_shapes property

Mapping from field name to declared array shape tuple (T-117-followup-bus-array). Fields default to scalar () when field_shapes is omitted at construction time.

BusMerge

Bases: LeafSystem

Merge two bus signals by union of fields (LeafSystem wrapper).

Wraps :func:merge_buses as a block for use inside a Diagram. The two upstream bus signals are read from input ports 0 and 1; the merged bus is produced on output port 0.

The merged-bus schema is fixed at construction time from bus_spec_a and bus_spec_b so the output port can be declared with a known NamedTuple type (and so the framework's pytree handling sees a consistent type across context-build and trace time, matching the T-117-fu-bus-namedtuple design for :class:BusCreator).

Parameters:

Name	Type	Description	Default
bus_spec_a		Schema of the first bus. Accepts a :class:BusCreator instance, a NamedTuple class, or a tuple/list of field-name strings.	required
bus_spec_b		Schema of the second bus. Same forms as bus_spec_a.	required
on_collision	str	Policy for fields that appear in both schemas: "error" (default) — raise :class:ValueError at construction time. "prefer_a" — read the colliding leaf from input port 0. "prefer_b" — read the colliding leaf from input port 1.	'error'

Input ports

(0) — bus_a (NamedTuple-shaped). (1) — bus_b (NamedTuple-shaped).

Output ports

(0) — the merged bus, a NamedTuple whose fields are the union of the two input schemas.

Source code in jaxonomy/library/routing.py

class BusMerge(LeafSystem):
    """Merge two bus signals by union of fields (LeafSystem wrapper).

    Wraps :func:`merge_buses` as a block for use inside a Diagram. The
    two upstream bus signals are read from input ports 0 and 1; the
    merged bus is produced on output port 0.

    The merged-bus schema is fixed at construction time from
    ``bus_spec_a`` and ``bus_spec_b`` so the output port can be declared
    with a known NamedTuple type (and so the framework's pytree handling
    sees a consistent type across context-build and trace time, matching
    the T-117-fu-bus-namedtuple design for :class:`BusCreator`).

    Parameters:
        bus_spec_a: Schema of the first bus. Accepts a
            :class:`BusCreator` instance, a NamedTuple class, or a
            tuple/list of field-name strings.
        bus_spec_b: Schema of the second bus. Same forms as
            ``bus_spec_a``.
        on_collision: Policy for fields that appear in both schemas:

            * ``"error"`` (default) — raise :class:`ValueError` at
              construction time.
            * ``"prefer_a"`` — read the colliding leaf from input
              port 0.
            * ``"prefer_b"`` — read the colliding leaf from input
              port 1.

    Input ports:
        ``(0)`` — bus_a (NamedTuple-shaped).
        ``(1)`` — bus_b (NamedTuple-shaped).

    Output ports:
        ``(0)`` — the merged bus, a NamedTuple whose fields are the
        union of the two input schemas.
    """

    def __init__(
        self,
        bus_spec_a,
        bus_spec_b,
        *args,
        on_collision: str = "error",
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

        fields_a = _bus_field_names(bus_spec_a)
        fields_b = _bus_field_names(bus_spec_b)
        if len(fields_a) == 0:
            raise ValueError(
                "BusMerge bus_spec_a must have at least one field; "
                "got an empty schema."
            )
        if len(fields_b) == 0:
            raise ValueError(
                "BusMerge bus_spec_b must have at least one field; "
                "got an empty schema."
            )
        merged_order, collisions = _merged_field_order(
            fields_a, fields_b, on_collision
        )

        self._fields_a = fields_a
        self._fields_b = fields_b
        self._on_collision = on_collision
        self._collisions = collisions
        self._merged_field_names = merged_order
        self._bus_type = namedtuple("MergedBus", merged_order)

        # Declare the two bus input ports. We do not pass units here —
        # the merged-bus unit/schema surface is a deeper followup (the
        # T-117-fu-bus-units machinery is per-BusCreator). Default-off
        # parity: callers who do not opt into units see no behavioural
        # change.
        idx_a = self.declare_input_port(name="bus_a")
        idx_b = self.declare_input_port(name="bus_b")
        ticket_a = self.input_ports[idx_a].ticket
        ticket_b = self.input_ports[idx_b].ticket

        # Pre-compute the per-output-field source map so the runtime
        # closure stays tight: a tuple of ``(source_index, field_name)``
        # pairs where ``source_index`` is 0 for bus_a and 1 for bus_b.
        set_a = set(fields_a)
        set_b = set(fields_b)
        source_map = []
        for name in merged_order:
            in_a = name in set_a
            in_b = name in set_b
            if in_a and in_b:
                # Collision — policy dictates the source. ``error`` is
                # rejected above in ``_merged_field_order`` so only the
                # two prefer-* cases survive here.
                source_map.append((0 if on_collision == "prefer_a" else 1, name))
            elif in_a:
                source_map.append((0, name))
            else:
                source_map.append((1, name))
        source_map_local = tuple(source_map)
        bus_type_local = self._bus_type

        def _compute_merged(_time, _state, *inputs, **_params):
            # ``inputs`` is ``(bus_a, bus_b)`` in declaration order.
            bus_a_val, bus_b_val = inputs
            buses = (bus_a_val, bus_b_val)
            leaves = tuple(
                getattr(buses[src], name) for (src, name) in source_map_local
            )
            return bus_type_local(*leaves)

        # As with BusCreator: do NOT pass ``default_value=`` — the
        # framework would call ``npa.array`` on the NamedTuple and lose
        # the typed-bus shape. The lazy default-value path computes the
        # correct NamedTuple-typed default from upstream defaults.
        self.declare_output_port(
            _compute_merged,
            name="bus",
            prerequisites_of_calc=[ticket_a, ticket_b],
            requires_inputs=True,
        )

    @property
    def field_names(self) -> tuple[str, ...]:
        """The tuple of merged field names, in output / declaration order."""
        return self._merged_field_names

    @property
    def bus_type(self) -> type:
        """The underlying NamedTuple class for the merged bus."""
        return self._bus_type

    @property
    def on_collision(self) -> str:
        """The collision-resolution policy in effect for this block."""
        return self._on_collision

    @property
    def collisions(self) -> tuple[str, ...]:
        """The tuple of field names that collided between the two
        input schemas. Empty unless ``on_collision`` is ``"prefer_a"``
        or ``"prefer_b"``."""
        return self._collisions

bus_type property

The underlying NamedTuple class for the merged bus.

collisions property

The tuple of field names that collided between the two input schemas. Empty unless on_collision is "prefer_a" or "prefer_b".

field_names property

The tuple of merged field names, in output / declaration order.

on_collision property

The collision-resolution policy in effect for this block.

BusPassthrough

Bases: LeafSystem

Identity copy of a bus signal — single input port, single output port.

The output is the input value, returned as-is. Pass-through semantics: NamedTuple-shaped buses (as produced by :class:BusCreator / :class:BusMerge) flow through unchanged, scalar/array signals likewise. This block exists to give Diagrams an explicit "junction" node for rewiring, debugging, or scheduler-boundary purposes; it has no parameters and does no computation beyond forwarding.

Differentiable: jax.grad flows from the output back to the input leaf-by-leaf (identity has Jacobian = identity), and the block is JIT-traceable since the underlying op is a no-op closure return.

Parameters:

Name	Type	Description	Default
bus_unit		Optional :class:BusUnit describing the per-field units of the bus signal. When supplied, both input and output ports are tagged with this BusUnit so connect- time unit checks see a matched pair. None (the default) preserves the unit-less behaviour byte-for-byte.	None

Input ports

(0) — the bus (or any) signal to forward.

Output ports

(0) — the same value, returned as-is.

Source code in jaxonomy/library/routing.py

class BusPassthrough(LeafSystem):
    """Identity copy of a bus signal — single input port, single output port.

    The output is the input value, returned as-is. Pass-through semantics:
    NamedTuple-shaped buses (as produced by :class:`BusCreator` /
    :class:`BusMerge`) flow through unchanged, scalar/array signals
    likewise. This block exists to give Diagrams an explicit "junction"
    node for rewiring, debugging, or scheduler-boundary purposes; it
    has no parameters and does no computation beyond forwarding.

    Differentiable: ``jax.grad`` flows from the output back to the input
    leaf-by-leaf (identity has Jacobian = identity), and the block is
    JIT-traceable since the underlying op is a no-op closure return.

    Parameters:
        bus_unit: Optional :class:`BusUnit` describing the per-field
            units of the bus signal. When supplied, both input and
            output ports are tagged with this ``BusUnit`` so connect-
            time unit checks see a matched pair. ``None`` (the default)
            preserves the unit-less behaviour byte-for-byte.

    Input ports:
        ``(0)`` — the bus (or any) signal to forward.

    Output ports:
        ``(0)`` — the same value, returned as-is.
    """

    def __init__(self, *args, bus_unit=None, **kwargs):
        super().__init__(*args, **kwargs)

        # Optional bus_unit propagation mirrors BusCreator/BusSelector:
        # when supplied, we tag both ports so the connect-time check
        # sees compatible BusUnit metadata on both sides. When None,
        # the default-off path is byte-equivalent to a no-op LeafSystem
        # with one input/one output and no unit metadata.
        if bus_unit is not None:
            from ..framework.units import BusUnit as _BusUnit

            if not isinstance(bus_unit, _BusUnit):
                raise TypeError(
                    f"BusPassthrough bus_unit must be a BusUnit instance, "
                    f"got {type(bus_unit).__name__}: {bus_unit!r}."
                )

        self._bus_unit = bus_unit
        self.declare_input_port(name="in", units=bus_unit)
        in_ticket = self.input_ports[0].ticket

        def _passthrough(_time, _state, *inputs, **_params):
            # ``inputs`` is a 1-tuple containing the upstream value;
            # we return it unchanged so NamedTuple-typed buses survive
            # without being flattened by ``npa.array``.
            (value,) = inputs
            return value

        # As with BusCreator / BusMerge: do NOT pass ``default_value=`` —
        # the leaf-system path would call ``npa.array`` on a NamedTuple-
        # shaped default and flatten it into a 1-D array, losing the
        # bus type. The lazy default-value computation re-runs the
        # closure on a dummy context and yields the correct typed value.
        self.declare_output_port(
            _passthrough,
            name="out",
            prerequisites_of_calc=[in_ticket],
            requires_inputs=True,
            units=bus_unit,
        )

    @property
    def bus_unit(self):
        """The :class:`BusUnit` propagated through this passthrough, or
        ``None`` if no unit metadata was supplied at construction."""
        return self._bus_unit

bus_unit property

The :class:BusUnit propagated through this passthrough, or None if no unit metadata was supplied at construction.

BusSelector

Bases: LeafSystem

Pull one named field out of a bus signal.

The bus input is expected to be a NamedTuple-shaped value (typically produced by :class:BusCreator). The selected field is read with plain getattr, so this block is the inverse of BusCreator when wired correctly: BusSelector("a")(BusCreator(["a","b","c"]) (a, b, c)) == a.

Both getattr and the NamedTuple constructor are transparent to JAX autodiff, so gradients flow cleanly from the selector output back to the upstream input that filled the corresponding bus slot.

T-117-followup-bus-dot-path: field_name may contain dots to descend into nested bus signals — e.g. BusSelector("chassis.suspension.spring_force") extracts the leaf in one block instead of three cascaded selectors. The path is resolved via :func:operator.attrgetter, so each segment must name a valid NamedTuple field at its level. Each segment is validated as a Python identifier at construction time.

Parameters:

Name	Type	Description	Default
field_name		Name (or dot-separated path) of the bus field to extract. Raises AttributeError at execution time if the upstream bus does not have a field at any segment of the path.	required
bus_unit		Optional :class:BusUnit describing the per-field units of the upstream bus (T-117-followup-bus-units). When supplied, the selector's input port is tagged with this BusUnit (so the connect-time check verifies that the upstream :class:BusCreator produced a compatible bus) and its output port is tagged with bus_unit.fields[field_name] so further downstream blocks see the right scalar unit. Default None preserves the unit-less behaviour byte-for-byte. When field_name contains dots, the leaf unit cannot be propagated because :class:BusUnit is flat (one Unit per top-level field); the input bus is still tagged but the output unit is None. Pass a single-segment field_name if you need leaf-unit propagation.	None

Input ports

(0) — the bus signal (NamedTuple-shaped).

Output ports

(0) — the value of bus.<field_name>, optionally sliced at slice_idx.

Source code in jaxonomy/library/routing.py

class BusSelector(LeafSystem):
    """Pull one named field out of a bus signal.

    The bus input is expected to be a NamedTuple-shaped value (typically
    produced by :class:`BusCreator`). The selected field is read with
    plain ``getattr``, so this block is the inverse of ``BusCreator``
    when wired correctly: ``BusSelector("a")(BusCreator(["a","b","c"])
    (a, b, c)) == a``.

    Both ``getattr`` and the NamedTuple constructor are transparent to
    JAX autodiff, so gradients flow cleanly from the selector output
    back to the upstream input that filled the corresponding bus slot.

    T-117-followup-bus-dot-path: ``field_name`` may contain dots to
    descend into nested bus signals — e.g.
    ``BusSelector("chassis.suspension.spring_force")`` extracts the
    leaf in one block instead of three cascaded selectors. The path is
    resolved via :func:`operator.attrgetter`, so each segment must
    name a valid NamedTuple field at its level. Each segment is
    validated as a Python identifier at construction time.

    Parameters:
        field_name: Name (or dot-separated path) of the bus field to
            extract. Raises ``AttributeError`` at execution time if
            the upstream bus does not have a field at any segment of
            the path.
        bus_unit: Optional :class:`BusUnit` describing the per-field
            units of the upstream bus (T-117-followup-bus-units). When
            supplied, the selector's *input* port is tagged with this
            ``BusUnit`` (so the connect-time check verifies that the
            upstream :class:`BusCreator` produced a compatible bus)
            and its *output* port is tagged with
            ``bus_unit.fields[field_name]`` so further downstream
            blocks see the right scalar unit. Default ``None``
            preserves the unit-less behaviour byte-for-byte.

            When ``field_name`` contains dots, the leaf unit cannot be
            propagated because :class:`BusUnit` is flat (one Unit per
            top-level field); the input bus is still tagged but the
            output unit is ``None``. Pass a single-segment
            ``field_name`` if you need leaf-unit propagation.

    Input ports:
        ``(0)`` — the bus signal (NamedTuple-shaped).

    Output ports:
        ``(0)`` — the value of ``bus.<field_name>``, optionally sliced
        at ``slice_idx``.
    """

    def __init__(
        self, field_name, *args, bus_unit=None, slice_idx=None, **kwargs
    ):
        super().__init__(*args, **kwargs)

        if not isinstance(field_name, str):
            raise ValueError(
                f"BusSelector field_name must be a str, got "
                f"{type(field_name).__name__}: {field_name!r}."
            )
        # T-117-followup-bus-dot-path: validate each segment of the
        # (possibly dotted) path individually; an all-empty / pure-dot
        # path is rejected up front.
        path_segments = field_name.split(".") if field_name else []
        if not path_segments:
            raise ValueError(
                "BusSelector field_name must be non-empty."
            )
        for seg in path_segments:
            if not seg.isidentifier():
                # Phrase the error so it matches the legacy
                # "not a valid Python identifier" wording for the
                # single-segment case, plus an extra hint for the
                # dot-path case.
                detail = (
                    f"BusSelector field_name {field_name!r} is not a "
                    f"valid Python identifier (NamedTuple field "
                    f"constraint)."
                )
                if len(path_segments) > 1:
                    detail = (
                        f"BusSelector field_name {field_name!r} "
                        f"contains segment {seg!r} that is not a valid "
                        f"Python identifier (each dot-separated segment "
                        f"must be a valid Python identifier — "
                        f"NamedTuple field constraint)."
                    )
                raise ValueError(detail)
        is_dotted = len(path_segments) > 1

        # T-117-followup-bus-array: ``slice_idx`` must be a plain
        # non-negative int when supplied (bools are technically ints so
        # we filter them out explicitly). Validating up front means an
        # obvious mistake fires at construction time rather than
        # halfway through a traced JAX path.
        if slice_idx is not None:
            if isinstance(slice_idx, bool) or not isinstance(slice_idx, int):
                raise TypeError(
                    f"BusSelector slice_idx must be an int, got "
                    f"{type(slice_idx).__name__}: {slice_idx!r}."
                )
            if slice_idx < 0:
                raise ValueError(
                    f"BusSelector slice_idx must be non-negative, "
                    f"got {slice_idx!r}."
                )

        # T-117-followup-bus-units: when a BusUnit is supplied, the
        # selector's output unit is the per-field unit from the bus.
        # We validate the field exists in the supplied schema so the
        # error fires here rather than at connect time.
        #
        # T-117-followup-bus-dot-path: dotted paths defeat the
        # flat-BusUnit lookup, so the leaf unit is silently dropped
        # (the input bus is still tagged for the connect-time check).
        output_unit = None
        if bus_unit is not None:
            from ..framework.units import BusUnit as _BusUnit

            if not isinstance(bus_unit, _BusUnit):
                raise TypeError(
                    f"BusSelector bus_unit must be a BusUnit instance, "
                    f"got {type(bus_unit).__name__}: {bus_unit!r}."
                )
            if not is_dotted:
                if field_name not in bus_unit.fields:
                    raise ValueError(
                        f"BusSelector field_name {field_name!r} is not "
                        f"present in bus_unit fields "
                        f"{sorted(bus_unit.fields)!r}."
                    )
                output_unit = bus_unit.fields[field_name]
            else:
                # Top-segment must still be present so the connect-time
                # check can verify the upstream bus has the right
                # outer shape. We don't recurse — BusUnit is flat — but
                # at least pin the outer field name.
                top = path_segments[0]
                if top not in bus_unit.fields:
                    raise ValueError(
                        f"BusSelector field_name {field_name!r}: top-"
                        f"level segment {top!r} is not present in "
                        f"bus_unit fields {sorted(bus_unit.fields)!r}."
                    )
                # output_unit stays None for nested paths.

        self._field_name = field_name
        self._bus_unit = bus_unit
        self._slice_idx = slice_idx
        self.declare_input_port(name="bus", units=bus_unit)

        # T-117-followup-bus-dot-path: ``operator.attrgetter`` resolves
        # dot-paths in a single call; for the single-segment case it
        # is exactly equivalent to ``getattr(bus, name)``. Bind it in
        # the closure so the traced computation has a fixed callable.
        # T-117-followup-bus-array: when ``slice_idx`` is supplied we
        # bind it inside the closure so the traced computation indexes
        # the field array at the static slot. ``attrgetter`` + integer
        # indexing are both transparent to ``jax.grad``/``jit``.
        import operator as _operator
        _getter = _operator.attrgetter(field_name)
        if slice_idx is None:

            def _compute_field(_time, _state, *inputs, **_params):
                (bus,) = inputs
                return _getter(bus)
        else:
            _idx = slice_idx

            def _compute_field(_time, _state, *inputs, **_params):
                (bus,) = inputs
                return _getter(bus)[_idx]

        # Use only the leaf segment as the output-port name to keep
        # port names valid Python identifiers / NamedTuple-friendly.
        leaf_name = path_segments[-1]
        self.declare_output_port(
            _compute_field,
            name=leaf_name,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
            units=output_unit,
        )

    @property
    def field_name(self) -> str:
        """The name of the bus field this block selects."""
        return self._field_name

    @property
    def bus_unit(self):
        """The :class:`BusUnit` describing the upstream bus, or
        ``None`` if no unit metadata was supplied."""
        return self._bus_unit

    @property
    def slice_idx(self):
        """The integer index into an array-valued bus field, or
        ``None`` if the selector returns the field value as-is
        (T-117-followup-bus-array)."""
        return self._slice_idx

bus_unit property

The :class:BusUnit describing the upstream bus, or None if no unit metadata was supplied.

field_name property

The name of the bus field this block selects.

slice_idx property

The integer index into an array-valued bus field, or None if the selector returns the field value as-is (T-117-followup-bus-array).

BusUnit dataclass

Compound unit carrying one :class:Unit per named bus field.

Attached to the output port of a :class:BusCreator (and the matching input port of a :class:BusSelector) so that the connect-time consistency check can verify each field's unit individually.

Attributes:

Name	Type	Description
fields	Mapping[str, Unit]	Mapping from bus field name to its :class:Unit. Stored as a plain dict (insertion order preserved).

Source code in jaxonomy/framework/units.py

@dataclass(frozen=True, eq=False)
class BusUnit:
    """Compound unit carrying one :class:`Unit` per named bus field.

    Attached to the output port of a :class:`BusCreator` (and the
    matching input port of a :class:`BusSelector`) so that the
    connect-time consistency check can verify each field's unit
    individually.

    Attributes:
        fields: Mapping from bus field name to its :class:`Unit`.
            Stored as a plain ``dict`` (insertion order preserved).
    """

    fields: Mapping[str, Unit] = field(default_factory=dict)

    def __post_init__(self):
        # Normalise to a plain dict so the class is hashable and the
        # ordering is deterministic. Validate each entry is a Unit.
        normalised: dict[str, Unit] = {}
        for k, v in dict(self.fields).items():
            if not isinstance(k, str):
                raise TypeError(
                    f"BusUnit field name must be a str, got {type(k).__name__}: {k!r}"
                )
            if not isinstance(v, Unit):
                raise TypeError(
                    f"BusUnit field {k!r} must map to a Unit instance, "
                    f"got {type(v).__name__}: {v!r}"
                )
            normalised[k] = v
        object.__setattr__(self, "fields", normalised)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, BusUnit):
            return NotImplemented
        if set(self.fields.keys()) != set(other.fields.keys()):
            return False
        return all(self.fields[k] == other.fields[k] for k in self.fields)

    def __hash__(self) -> int:
        return hash(tuple(sorted(self.fields.items(), key=lambda kv: kv[0])))

    def __repr__(self) -> str:
        body = ", ".join(f"{k}={v!r}" for k, v in self.fields.items())
        return f"BusUnit({{{body}}})"

    def field_unit(self, name: str) -> Unit | None:
        """Return the :class:`Unit` for ``name``, or ``None`` if absent.

        Used by :class:`BusSelector` to look up its output-port unit
        when wired downstream of a unit-tagged bus.
        """
        return self.fields.get(name)

field_unit(name)

Return the :class:Unit for name, or None if absent.

Used by :class:BusSelector to look up its output-port unit when wired downstream of a unit-tagged bus.

Source code in jaxonomy/framework/units.py

def field_unit(self, name: str) -> Unit | None:
    """Return the :class:`Unit` for ``name``, or ``None`` if absent.

    Used by :class:`BusSelector` to look up its output-port unit
    when wired downstream of a unit-tagged bus.
    """
    return self.fields.get(name)

BusUpdate

Bases: LeafSystem

Replace one field of a bus signal with a new value.

Two-input / one-output LeafSystem. Input port 0 carries the upstream bus (a NamedTuple-shaped value, typically produced by :class:BusCreator). Input port 1 carries the new value for the field named field_name. The output port is a fresh bus identical to the input except that the field_name slot is replaced by the new value. Field order is preserved exactly; all other fields are forwarded unchanged.

Use this instead of the BusSelector-modify-BusCreator triplet when you only need to edit one field of a wide bus -- the block makes the intent explicit and avoids manually wiring N-1 passthrough edges.

Differentiable: the underlying op is NamedTuple construction over getattr lookups on the input bus plus the new_value leaf, all of which are transparent to jax.grad and jax.jit (same as :class:BusCreator and :func:merge_buses).

Parameters:

Name	Type	Description	Default
bus_spec		Schema of the bus. Accepts a :class:BusCreator instance, a NamedTuple class, or a tuple/list of field-name strings (same forms as :class:BusMerge). Used both to validate field_name at construction time and to type the output port's NamedTuple class so the framework's pytree handling sees a consistent type across context-build and trace time.	required
field_name		The name of the field to replace. Must be one of the names declared in bus_spec; otherwise a clear :class:ValueError is raised at construction time.	required
bus_unit		Optional :class:BusUnit describing the per-field units of the upstream bus (T-117-followup-bus-update-units- prop). When supplied, the block's bus_in and bus_out ports advertise this BusUnit so the connect-time check verifies that the upstream :class:BusCreator produced a compatible bus and that downstream consumers see the right schema. The new_value input port is independently tagged with the per-field unit bus_unit.fields[field_name] so the connect-time check enforces unit compatibility on the replacement value. Composes with T-104 Phase 2 behaviour on Sum / Product / Integrator. Default None preserves the unit-less behaviour byte-for-byte.	None

Input ports

(0) — bus_in, the upstream bus signal (NamedTuple-shaped). (1) — new_value, the value to put into the field_name slot of the output bus.

Output ports

(0) — bus_out, a NamedTuple of the same shape as bus_in with field_name replaced by new_value.

Source code in jaxonomy/library/routing.py

class BusUpdate(LeafSystem):
    """Replace one field of a bus signal with a new value.

    Two-input / one-output LeafSystem. Input port 0 carries the upstream
    bus (a NamedTuple-shaped value, typically produced by
    :class:`BusCreator`). Input port 1 carries the new value for the
    field named ``field_name``. The output port is a fresh bus identical
    to the input except that the ``field_name`` slot is replaced by the
    new value. Field order is preserved exactly; all other fields are
    forwarded unchanged.

    Use this instead of the BusSelector-modify-BusCreator triplet when
    you only need to edit one field of a wide bus -- the block makes the
    intent explicit and avoids manually wiring N-1 passthrough edges.

    Differentiable: the underlying op is NamedTuple construction over
    ``getattr`` lookups on the input bus plus the ``new_value`` leaf, all
    of which are transparent to ``jax.grad`` and ``jax.jit`` (same as
    :class:`BusCreator` and :func:`merge_buses`).

    Parameters:
        bus_spec: Schema of the bus. Accepts a :class:`BusCreator`
            instance, a NamedTuple class, or a tuple/list of field-name
            strings (same forms as :class:`BusMerge`). Used both to
            validate ``field_name`` at construction time and to type the
            output port's NamedTuple class so the framework's pytree
            handling sees a consistent type across context-build and
            trace time.
        field_name: The name of the field to replace. Must be one of the
            names declared in ``bus_spec``; otherwise a clear
            :class:`ValueError` is raised at construction time.
        bus_unit: Optional :class:`BusUnit` describing the per-field
            units of the upstream bus (T-117-followup-bus-update-units-
            prop). When supplied, the block's ``bus_in`` and ``bus_out``
            ports advertise this ``BusUnit`` so the connect-time check
            verifies that the upstream :class:`BusCreator` produced a
            compatible bus and that downstream consumers see the right
            schema. The ``new_value`` input port is independently tagged
            with the per-field unit ``bus_unit.fields[field_name]`` so
            the connect-time check enforces unit compatibility on the
            replacement value. Composes with T-104 Phase 2 behaviour on
            ``Sum`` / ``Product`` / ``Integrator``. Default ``None``
            preserves the unit-less behaviour byte-for-byte.

    Input ports:
        ``(0)`` — ``bus_in``, the upstream bus signal (NamedTuple-shaped).
        ``(1)`` — ``new_value``, the value to put into the ``field_name``
        slot of the output bus.

    Output ports:
        ``(0)`` — ``bus_out``, a NamedTuple of the same shape as
        ``bus_in`` with ``field_name`` replaced by ``new_value``.
    """

    def __init__(
        self,
        bus_spec,
        field_name,
        *args,
        bus_unit=None,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

        # Reuse the bus_spec normalisation helper from the BusMerge
        # section: accepts a BusCreator instance, a NamedTuple class, or
        # a tuple/list of strings, with valid-identifier validation.
        field_names = _bus_field_names(bus_spec)
        if len(field_names) == 0:
            raise ValueError(
                "BusUpdate bus_spec must have at least one field; "
                "got an empty schema."
            )
        if not isinstance(field_name, str):
            raise TypeError(
                "BusUpdate field_name must be a string; got "
                f"{type(field_name).__name__}: {field_name!r}."
            )
        if field_name not in field_names:
            raise ValueError(
                f"BusUpdate field_name {field_name!r} is not in the "
                f"bus schema {field_names!r}."
            )

        # T-117-followup-bus-update-units-prop: validate bus_unit
        # against the schema at construction time so an obvious mistake
        # fires here rather than at the connect-time check downstream.
        new_value_unit = None
        if bus_unit is not None:
            from ..framework.units import BusUnit as _BusUnit

            if not isinstance(bus_unit, _BusUnit):
                raise TypeError(
                    f"BusUpdate bus_unit must be a BusUnit instance, "
                    f"got {type(bus_unit).__name__}: {bus_unit!r}."
                )
            # The BusUnit's field set must match the bus schema exactly;
            # a missing or stray name signals a mismatched declaration.
            if set(bus_unit.fields.keys()) != set(field_names):
                raise ValueError(
                    f"BusUpdate bus_unit fields "
                    f"{sorted(bus_unit.fields.keys())!r} do not match "
                    f"the bus schema {sorted(field_names)!r}."
                )
            new_value_unit = bus_unit.fields[field_name]

        self._field_names = field_names
        self._field_name = field_name
        self._bus_unit = bus_unit
        # The output bus is the same shape as the input -- we name the
        # NamedTuple class "Bus" to match the T-117-fu-bus-namedtuple
        # convention used by BusCreator. The class identity differs from
        # the upstream BusCreator's class (one fresh class per BusUpdate
        # instance) but the field tuple matches exactly, so JAX's pytree
        # handling treats them as structurally equivalent.
        self._bus_type = namedtuple("Bus", field_names)

        # Declare input ports: bus_in then new_value, in that order.
        # T-117-followup-bus-update-units-prop: when ``bus_unit`` is
        # supplied, the ``bus_in`` and ``bus_out`` ports advertise the
        # full BusUnit (so downstream consumers see the schema and the
        # upstream connect-time check verifies compatibility), and the
        # ``new_value`` port advertises the per-field leaf unit.
        idx_bus = self.declare_input_port(name="bus_in", units=bus_unit)
        idx_new = self.declare_input_port(
            name="new_value", units=new_value_unit,
        )
        ticket_bus = self.input_ports[idx_bus].ticket
        ticket_new = self.input_ports[idx_new].ticket

        # Pre-compute the per-output-slot source map so the runtime
        # closure stays tight: tuple of ``(use_new_value, name)`` pairs
        # in declaration order. ``use_new_value`` is True for exactly the
        # ``field_name`` slot; False for every other field (which is
        # forwarded unchanged from ``bus_in``).
        source_map_local = tuple(
            (name == field_name, name) for name in field_names
        )
        bus_type_local = self._bus_type

        def _compute_updated(_time, _state, *inputs, **_params):
            # ``inputs`` is ``(bus_in, new_value)`` in declaration order.
            bus_in_val, new_value = inputs
            leaves = tuple(
                new_value if use_new else getattr(bus_in_val, name)
                for (use_new, name) in source_map_local
            )
            return bus_type_local(*leaves)

        # As with BusCreator / BusMerge / BusPassthrough: do NOT pass
        # ``default_value=`` -- the framework would call ``npa.array``
        # on the NamedTuple and flatten it. The lazy default-value
        # path computes the correct NamedTuple-typed default by running
        # ``_compute_updated`` on a dummy context.
        # T-117-followup-bus-update-units-prop: tag the output port with
        # the same BusUnit as the input bus so downstream blocks see
        # the schema preserved through the update.
        self.declare_output_port(
            _compute_updated,
            name="bus_out",
            prerequisites_of_calc=[ticket_bus, ticket_new],
            requires_inputs=True,
            units=bus_unit,
        )

    @property
    def field_names(self) -> tuple[str, ...]:
        """The tuple of bus field names, in declaration / output order."""
        return self._field_names

    @property
    def field_name(self) -> str:
        """The name of the field this block replaces on each tick."""
        return self._field_name

    @property
    def bus_type(self) -> type:
        """The underlying NamedTuple class for the output bus."""
        return self._bus_type

    @property
    def bus_unit(self):
        """The :class:`BusUnit` propagated through this update, or
        ``None`` if no unit metadata was supplied at construction
        (T-117-followup-bus-update-units-prop)."""
        return self._bus_unit

bus_type property

The underlying NamedTuple class for the output bus.

bus_unit property

The :class:BusUnit propagated through this update, or None if no unit metadata was supplied at construction (T-117-followup-bus-update-units-prop).

field_name property

The name of the field this block replaces on each tick.

field_names property

The tuple of bus field names, in declaration / output order.

Chirp

Bases: SourceBlock

Produces a linear chirp signal — matches :func:scipy.signal.chirp.

The output signal is cos(2π·f(t)·t + phi) with the linearly swept frequency f(t) = f0 + (f1 − f0)·t/(2·stop_time). At t=0 the instantaneous frequency is f0 Hz; at t=stop_time it is f1 Hz.

See https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.chirp.html

Parameters:

Name	Type	Description	Default
f0	float	Frequency (Hz) at time t=0.	required
f1	float	Frequency (Hz) at time t=stop_time.	required
stop_time	float	Time to end the signal (seconds).	required
phi	float	Phase offset (radians).	0.0
units	str | None	T-122-followup-chirp-hz-convention — frequency-unit convention for f0 / f1. Defaults to "hz" (matches the docstring and :func:scipy.signal.chirp). Legacy diagrams that depended on the pre-2026-05 behaviour (where f0 / f1 were silently interpreted in rad/s) can opt into the old semantics by passing units="rad/s"; doing so emits a :class:DeprecationWarning because the legacy path will be removed in a future release.	'hz'

Input ports

None

Output ports

(0) The chirp signal.

Source code in jaxonomy/library/sources.py

class Chirp(SourceBlock):
    """Produces a linear chirp signal — matches :func:`scipy.signal.chirp`.

    The output signal is ``cos(2π·f(t)·t + phi)`` with the linearly
    swept frequency ``f(t) = f0 + (f1 − f0)·t/(2·stop_time)``. At
    ``t=0`` the instantaneous frequency is ``f0`` Hz; at
    ``t=stop_time`` it is ``f1`` Hz.

    See
    https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.chirp.html

    Parameters:
        f0 (float): Frequency (Hz) at time t=0.
        f1 (float): Frequency (Hz) at time t=stop_time.
        stop_time (float): Time to end the signal (seconds).
        phi (float): Phase offset (radians).
        units (str | None, optional): T-122-followup-chirp-hz-convention
            — frequency-unit convention for ``f0`` / ``f1``. Defaults to
            ``"hz"`` (matches the docstring and :func:`scipy.signal.chirp`).
            Legacy diagrams that depended on the pre-2026-05 behaviour
            (where ``f0`` / ``f1`` were silently interpreted in rad/s)
            can opt into the old semantics by passing ``units="rad/s"``;
            doing so emits a :class:`DeprecationWarning` because the
            legacy path will be removed in a future release.

    Input ports:
        None

    Output ports:
        (0) The chirp signal.
    """

    @parameters(static=["units"], dynamic=["f0", "f1", "stop_time", "phi"])
    def __init__(self, f0, f1, stop_time, phi=0.0, units="hz", **kwargs):
        # T-122-followup-chirp-hz-convention: ``units`` is a static
        # opt-in to the legacy ``rad/s`` semantics; not JAX-traced.
        if units not in ("hz", "rad/s"):
            raise BlockParameterError(
                message=(
                    f"Chirp: units must be 'hz' (default) or 'rad/s' "
                    f"(legacy), got {units!r}."
                ),
                parameter_name="units",
            )
        if units == "rad/s":
            warnings.warn(
                "Chirp(..., units='rad/s') is the pre-2026-05 legacy "
                "convention that contradicts the documented Hz "
                "semantics and the scipy.signal.chirp parity claim. It "
                "is preserved here only for backwards compatibility "
                "and will be removed in a future release. Migrate to "
                "the default (units='hz') and divide your existing "
                "f0/f1 by 2π if you actually meant angular frequency.",
                DeprecationWarning,
                stacklevel=2,
            )
        self._chirp_units = units
        super().__init__(None, **kwargs)

    def initialize(self, f0, f1, stop_time, phi, units="hz"):
        # T-122-followup-chirp-hz-convention: the Hz path multiplies by
        # 2π so the instantaneous frequency at time ``t`` matches the
        # docstring's stated f0 + (f1−f0)·t/stop_time Hz schedule. The
        # legacy ``rad/s`` path preserves the pre-fix expression
        # exactly for byte-equivalent reproducibility on existing
        # diagrams. Reads ``self._chirp_units`` so the JSON round-trip
        # path (which goes through ``@parameters(static=...)`` and
        # injects ``units`` as a re-init kwarg) and the manual path
        # both wind up with the same effective setting.
        del units  # parity arg for the @parameters decorator
        if self._chirp_units == "hz":
            two_pi = 2 * npa.pi

            def _func(time, stop_time, f0, f1, phi):
                f = f0 + (f1 - f0) * time / (2 * stop_time)
                return npa.cos(two_pi * f * time + phi)

        else:  # "rad/s" — legacy path

            def _func(time, stop_time, f0, f1, phi):
                f = f0 + (f1 - f0) * time / (2 * stop_time)
                return npa.cos(f * time + phi)

        self.replace_op(_func)

Clock

Bases: SourceBlock

Source block returning simulation time.

Input ports

None

Output ports

(0) The simulation time.

Parameters:

Name	Type	Description	Default
dtype		The data type of the output signal. The default is "None", which will default to the current default floating point precision	None

Source code in jaxonomy/library/sources.py

class Clock(SourceBlock):
    """Source block returning simulation time.

    Input ports:
        None

    Output ports:
        (0) The simulation time.

    Parameters:
        dtype:
            The data type of the output signal.  The default is "None", which will
            default to the current default floating point precision
    """

    def __init__(self, dtype=None, **kwargs):
        super().__init__(lambda t: npa.array(t, dtype=dtype), **kwargs)

Comparator

Bases: LeafSystem

Compare two signals using typical relational operators.

When using == and != operators, the block uses tolerances to determine if the expression is true or false.

Parameters:

Name	Type	Description	Default
operator		one of ("==", "!=", ">=", ">", ">=", "<")	None
atol		the absolute tolerance value used with "==" or "!="	1e-05
rtol		the relative tolerance value used with "==" or "!="	1e-08

Input Ports

(0) The left side operand (1) The right side operand

Output Ports

(0) The result of the comparison (boolean signal)

Events

An event is triggered when the output changes from true to false or vice versa.

Source code in jaxonomy/library/logic.py

class Comparator(LeafSystem):
    """Compare two signals using typical relational operators.

    When using == and != operators, the block uses tolerances to determine if the
    expression is true or false.

    Parameters:
        operator: one of ("==", "!=", ">=", ">", ">=", "<")
        atol: the absolute tolerance value used with "==" or "!="
        rtol: the relative tolerance value used with "==" or "!="

    Input Ports:
        (0) The left side operand
        (1) The right side operand

    Output Ports:
        (0) The result of the comparison (boolean signal)

    Events:
        An event is triggered when the output changes from true to false or vice versa.
    """

    @parameters(static=["operator", "atol", "rtol"])
    def __init__(self, atol=1e-5, rtol=1e-8, operator=None, **kwargs):
        super().__init__(**kwargs)
        self.declare_input_port()
        self.declare_input_port()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, atol, rtol, operator):
        func_lookup = {
            ">": npa.greater,
            ">=": npa.greater_equal,
            "<": npa.less,
            "<=": npa.less_equal,
            "==": self._equal,
            "!=": self._ne,
        }

        if operator not in func_lookup:
            message = (
                f"Comparator block '{self.name}' has invalid selection "
                + f"'{operator}' for parameter 'operator'. Valid options: "
                + ",".join([k for k in func_lookup.keys()])
            )
            raise BlockParameterError(
                message=message, system=self, parameter_name="operator"
            )

        self.rtol = rtol
        self.atol = atol

        compare = func_lookup[operator]

        def _compute_output(_time, _state, *inputs, **_params):
            return compare(*inputs)

        self.configure_output_port(
            self._output_port_idx,
            _compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )
        self.evt_direction = self._process_operator(operator)

    def _equal(self, x, y):
        if npa.issubdtype(x.dtype, npa.floating):
            return npa.isclose(x, y, self.rtol, self.atol)
        return x == y

    def _ne(self, x, y):
        if npa.issubdtype(x.dtype, npa.floating):
            return npa.logical_not(npa.isclose(x, y, self.rtol, self.atol))
        return x != y

    def _zero_crossing(self, _time, _state, *inputs, **_params):
        return inputs[0] - inputs[1]

    def _process_operator(self, operator):
        if operator in ["<", "<="]:
            return "positive_then_non_positive"
        if operator in [">", ">="]:
            return "negative_then_non_negative"
        return "crosses_zero"

    def initialize_static_data(self, context):
        # Add a zero-crossing event so ODE solvers can't try to integrate
        # through a discontinuity. For efficiency, only do this if the output is
        # fed to an ODE.
        if not self.has_zero_crossing_events and is_discontinuity(self.output_ports[0]):
            self.declare_zero_crossing(
                self._zero_crossing, direction=self.evt_direction
            )

        return super().initialize_static_data(context)

Conditional

Bases: LeafSystem

Container block that enables/disables a submodel.

Parameters:

Name	Type	Description	Default
submodel	Callable	Callable taking *inputs and returning a single output array. For a subdiagram, use jaxonomy.submodel_function to build a compatible callable, then wrap with a context-capturing lambda: lambda *u: f(context, *u).	required
n_inputs	int	Number of non-enable inputs the submodel takes. Input port 0 is always the enable signal; ports 1..n_inputs carry the submodel's inputs in order.	1
when_disabled	str	"reset", "hold", or "passthrough".	RESET
initial_value		Output value when disabled (reset or hold mode's initial state). Also used to infer output shape/dtype when the submodel has not been evaluated yet.	0.0
name		Optional block name.	required

Source code in jaxonomy/library/conditional.py

class Conditional(LeafSystem):
    """Container block that enables/disables a submodel.

    Args:
        submodel: Callable taking ``*inputs`` and returning a single
            output array.  For a subdiagram, use
            ``jaxonomy.submodel_function`` to build a compatible
            callable, then wrap with a context-capturing lambda:
            ``lambda *u: f(context, *u)``.
        n_inputs: Number of non-enable inputs the submodel takes.
            Input port 0 is always the enable signal; ports 1..n_inputs
            carry the submodel's inputs in order.
        when_disabled: ``"reset"``, ``"hold"``, or ``"passthrough"``.
        initial_value: Output value when disabled (reset or hold mode's
            initial state).  Also used to infer output shape/dtype when
            the submodel has not been evaluated yet.
        name: Optional block name.
    """

    def __init__(
        self,
        submodel: Callable,
        n_inputs: int = 1,
        when_disabled: str = WhenDisabled.RESET,
        initial_value=0.0,
        hold_period: float | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if when_disabled not in WhenDisabled.valid():
            raise ValueError(
                f"Conditional: when_disabled must be one of "
                f"{WhenDisabled.valid()!r}, got {when_disabled!r}"
            )
        if when_disabled == WhenDisabled.HOLD and not hold_period:
            raise ValueError(
                "Conditional(when_disabled='hold') requires a positive "
                "hold_period to determine the snapshot sample rate."
            )

        self._submodel = submodel
        self._when_disabled = when_disabled
        self._initial = jnp.asarray(initial_value)

        # Port 0 is always enable; remaining ports are submodel inputs.
        self.declare_input_port(name="enable")
        for i in range(n_inputs):
            self.declare_input_port(name=f"u_{i}")

        if when_disabled == WhenDisabled.HOLD:
            # Discrete state holding the last snapshot of the submodel
            # output while the block was enabled.  Updated at the
            # user-supplied ``hold_period``; between snapshots the output
            # reads from the last store.
            self.declare_discrete_state(default_value=self._initial)
            self.declare_periodic_update(
                self._hold_update, period=float(hold_period), offset=0.0,
            )
        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    # ── callbacks ─────────────────────────────────────────────────────────

    def _submodel_output(self, inputs):
        """Evaluate the wrapped submodel on the forwarded inputs."""
        user_inputs = inputs[1:]  # skip enable
        return jnp.asarray(self._submodel(*user_inputs))

    def _hold_update(self, time, state, *inputs, **params):
        enable = inputs[0]
        y_sub = self._submodel_output(inputs)
        # On disabled steps, keep the previous held value.
        return jnp.where(
            jnp.asarray(enable).astype(bool),
            y_sub,
            state.discrete_state,
        )

    def _compute_output(self, time, state, *inputs, **params):
        enable = jnp.asarray(inputs[0]).astype(bool)
        y_sub = self._submodel_output(inputs)

        if self._when_disabled == WhenDisabled.RESET:
            return jnp.where(enable, y_sub, self._initial)

        if self._when_disabled == WhenDisabled.HOLD:
            held = state.discrete_state
            return jnp.where(enable, y_sub, held)

        # passthrough: output = first user input when disabled.  This
        # requires the submodel output and the first user input to have
        # compatible shapes.
        passthrough = jnp.asarray(inputs[1]) if len(inputs) > 1 else self._initial
        return jnp.where(enable, y_sub, passthrough)

Constant

Bases: LeafSystem

A source block that emits a constant value.

Parameters:

Name	Type	Description	Default
value		The constant value of the block.	required
dtype	optional, T-038a-followup-other-blocks	If set, the constant value is cast to this dtype on output. See LookupTable1d for the per-block dtype contract.	None
units	(optional, T - 104 - followup - units - on - source - blocks)	If set, the output port advertises this :class:Unit. The connect-time consistency check (T-104) then enforces downstream ports declare a compatible unit. Default None keeps the legacy "no-units" behaviour (byte-equivalent to pre-T-104 diagrams).	None

Input ports

None

Output ports

(0) The constant value.

Source code in jaxonomy/library/sources.py

class Constant(LeafSystem):
    """A source block that emits a constant value.

    Parameters:
        value: The constant value of the block.
        dtype (optional, T-038a-followup-other-blocks):
            If set, the constant value is cast to this dtype on output.
            See ``LookupTable1d`` for the per-block dtype contract.
        units (optional, T-104-followup-units-on-source-blocks):
            If set, the output port advertises this :class:`Unit`. The
            connect-time consistency check (T-104) then enforces
            downstream ports declare a compatible unit. Default
            ``None`` keeps the legacy "no-units" behaviour
            (byte-equivalent to pre-T-104 diagrams).

    Input ports:
        None

    Output ports:
        (0) The constant value.
    """

    @parameters(dynamic=["value"])
    def __init__(self, value, *args, dtype=None, units=None, **kwargs):
        # T-038a-followup-other-blocks: dtype is stored outside the
        # @parameters dynamic list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(**kwargs)
        # T-104-followup-units-on-source-blocks: forward the optional
        # ``units=`` kwarg to the output-port declaration so the source
        # block can advertise its own output unit (rather than relying
        # on the downstream port to do so).
        self._output_port_idx = self.declare_output_port(
            name="out_0", units=units
        )

    def initialize(self, value):
        if self._dtype is None:

            def _func(time, state, *inputs, **parameters):
                return parameters["value"]

        else:
            _dtype = self._dtype

            def _func(time, state, *inputs, **parameters):
                return npa.asarray(parameters["value"]).astype(_dtype)

        self.configure_output_port(
            self._output_port_idx,
            _func,
            prerequisites_of_calc=[DependencyTicket.nothing],
            requires_inputs=False,
        )

ContinuousTimeInfiniteHorizonKalmanFilter

Bases: LeafSystem

Continuous-time Infinite Horizon Kalman Filter for the following system:

dot_x =  A x + B u + G w
y   = C x + D u + v

E(w) = E(v) = 0
E(ww') = Q
E(vv') = R
E(wv') = N = 0

Input ports

(0) u : continuous-time control vector (1) y : continuous-time measurement vector

Output ports

(1) x_hat : continuous-time state vector estimate

Parameters:

Name	Type	Description	Default
A		ndarray State transition matrix	required
B		ndarray Input matrix	required
C		ndarray Output matrix	required
D		ndarray Feedthrough matrix	required
G		ndarray Process noise matrix	required
Q		ndarray Process noise covariance matrix	required
R		ndarray Measurement noise covariance matrix	required
x_hat_0		ndarray Initial state estimate	required

Source code in jaxonomy/library/state_estimators/continuous_time_infinite_horizon_kalman_filter.py

class ContinuousTimeInfiniteHorizonKalmanFilter(LeafSystem):
    """
    Continuous-time Infinite Horizon Kalman Filter for the following system:

    ```
    dot_x =  A x + B u + G w
    y   = C x + D u + v

    E(w) = E(v) = 0
    E(ww') = Q
    E(vv') = R
    E(wv') = N = 0
    ```

    Input ports:
        (0) u : continuous-time control vector
        (1) y : continuous-time measurement vector

    Output ports:
        (1) x_hat : continuous-time state vector estimate

    Parameters:
        A: ndarray
            State transition matrix
        B: ndarray
            Input matrix
        C: ndarray
            Output matrix
        D: ndarray
            Feedthrough matrix
        G: ndarray
            Process noise matrix
        Q: ndarray
            Process noise covariance matrix
        R: ndarray
            Measurement noise covariance matrix
        x_hat_0: ndarray
            Initial state estimate
    """

    def __init__(self, A, B, C, D, G, Q, R, x_hat_0, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.A = A
        self.B = B
        self.C = C
        self.D = D
        self.G = G
        self.Q = Q
        self.R = R

        self.nx, self.nu = B.shape
        self.ny = C.shape[0]

        L, P, E = control.lqe(A, G, C, Q, R)

        self.A_minus_LC = A - npa.matmul(L, C)
        self.B_minus_LD = B - npa.matmul(L, D)
        self.L = L

        self.declare_input_port()  # u
        self.declare_input_port()  # y

        self.declare_continuous_state(
            ode=self._ode, shape=x_hat_0.shape, default_value=x_hat_0, as_array=True
        )  # continuous state for x_hat

        self.declare_continuous_state_output()

    def _ode(self, time, state, *inputs, **params):
        x_hat = state.continuous_state

        u, y = inputs

        u = npa.atleast_1d(u)
        y = npa.atleast_1d(y)

        dot_x_hat = (
            npa.dot(self.A_minus_LC, x_hat)
            + npa.dot(self.B_minus_LD, u)
            + npa.dot(self.L, y)
        )

        return dot_x_hat

    #######################################
    # Make filter for a continuous plant  #
    #######################################
    @staticmethod
    @with_resolved_parameters
    def for_continuous_plant(
        plant,
        x_eq,
        u_eq,
        Q,
        R,
        G=None,
        x_hat_bar_0=None,
        name=None,
    ):
        """
        Obtain a continuous-time Infinite Horizon Kalman Filter system for a
        continuous-time plant after linearization at equilibrium point (x_eq, u_eq)

        The input plant contains the deterministic forms of the forward and observation
        operators:

        ```
            dx/dt = f(x,u)
            y = g(x,u)
        ```

        Note: Only plants with one vector-valued input and one vector-valued output
        are currently supported. Furthermore, the plant LeafSystem/Diagram should have
        only one vector-valued integrator.

        A plant with disturbances of the following form is then considered
        following form:

        ```
            dx/dt = f(x,u) + G w
            y = g(x,u) +  v
        ```

        where:

            `w` represents the process noise,
            `v` represents the measurement noise,

        and

        ```
            E(w) = E(v) = 0
            E(ww') = Q
            E(vv') = R
            E(wv') = N = 0
        ```

        This plant with disturbances is linearized (only `f` and `q`) around the
        equilibrium point to obtain:

        ```
            d/dt (x_bar) = A x_bar + B u_bar + G w    --- (C1)
            y_bar = C x_bar + D u_bar + v             --- (C2)
        ```

        where,

        ```
            x_bar = x - x_eq
            u_bar = u - u_eq
            y_bar = y - y_bar
            y_eq = g(x_eq, u_eq)
        ```

        A continuous-time Kalman Filter estimator for the system of equations (C1) and
        (C2) is returned. This filter is in the `x_bar`, `u_bar`, and `y_bar`
        states.

        The returned system will have

        Input ports:
            (0) u_bar : continuous-time control vector relative to equilibrium point
            (1) y_bar : continuous-time measurement vector relative to equilibrium point

        Output ports:
            (1) x_hat_bar : continuous-time state vector estimate relative to
                            equilibrium point

        Parameters:
            plant : a `Plant` object which can be a LeafSystem or a Diagram.
            x_eq: ndarray
                Equilibrium state vector for discretization
            u_eq: ndarray
                Equilibrium control vector for discretization
            Q: ndarray
                Process noise covariance matrix.
            R: ndarray
                Measurement noise covariance matrix.
            G: ndarray
                Process noise matrix. If `None`, `G=B` is assumed making disrurbances
                additive to control vector `u`, i.e. `u_disturbed = u_orig + w`.
            x_hat_bar_0: ndarray
                Initial state estimate relative to equilibrium point.
                If None, an identity matrix is assumed.
        """

        y_eq, linear_plant = linearize_plant(plant, x_eq, u_eq)

        # LTISystem populates its A/B/C/D attributes in initialize(), which runs
        # at context creation — trigger it before reading them (matches the
        # pattern in state_estimators.utils).
        linear_plant.create_context()

        A, B, C, D = linear_plant.A, linear_plant.B, linear_plant.C, linear_plant.D

        nx, nu = B.shape
        ny, _ = D.shape

        if G is None:
            G = B

        if x_hat_bar_0 is None:
            x_hat_bar_0 = npa.zeros(nx)

        # Instantiate a Kalman Filter instance for the linearized plant
        kf = ContinuousTimeInfiniteHorizonKalmanFilter(
            A,
            B,
            C,
            D,
            G,
            Q,
            R,
            x_hat_bar_0,
            name=name,
        )

        return y_eq, kf

for_continuous_plant(plant, x_eq, u_eq, Q, R, G=None, x_hat_bar_0=None, name=None) staticmethod

Obtain a continuous-time Infinite Horizon Kalman Filter system for a continuous-time plant after linearization at equilibrium point (x_eq, u_eq)

The input plant contains the deterministic forms of the forward and observation operators:

    dx/dt = f(x,u)
    y = g(x,u)

Note: Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator.

A plant with disturbances of the following form is then considered following form:

    dx/dt = f(x,u) + G w
    y = g(x,u) +  v

where:

`w` represents the process noise,
`v` represents the measurement noise,

and

    E(w) = E(v) = 0
    E(ww') = Q
    E(vv') = R
    E(wv') = N = 0

This plant with disturbances is linearized (only f and q) around the equilibrium point to obtain:

    d/dt (x_bar) = A x_bar + B u_bar + G w    --- (C1)
    y_bar = C x_bar + D u_bar + v             --- (C2)

where,

    x_bar = x - x_eq
    u_bar = u - u_eq
    y_bar = y - y_bar
    y_eq = g(x_eq, u_eq)

A continuous-time Kalman Filter estimator for the system of equations (C1) and (C2) is returned. This filter is in the x_bar, u_bar, and y_bar states.

The returned system will have

Input ports

(0) u_bar : continuous-time control vector relative to equilibrium point (1) y_bar : continuous-time measurement vector relative to equilibrium point

Output ports

(1) x_hat_bar : continuous-time state vector estimate relative to equilibrium point

Parameters:

Name	Type	Description	Default
plant		a Plant object which can be a LeafSystem or a Diagram.	required
x_eq		ndarray Equilibrium state vector for discretization	required
u_eq		ndarray Equilibrium control vector for discretization	required
Q		ndarray Process noise covariance matrix.	required
R		ndarray Measurement noise covariance matrix.	required
G		ndarray Process noise matrix. If None, G=B is assumed making disrurbances additive to control vector u, i.e. u_disturbed = u_orig + w.	None
x_hat_bar_0		ndarray Initial state estimate relative to equilibrium point. If None, an identity matrix is assumed.	None

Source code in jaxonomy/library/state_estimators/continuous_time_infinite_horizon_kalman_filter.py

@staticmethod
@with_resolved_parameters
def for_continuous_plant(
    plant,
    x_eq,
    u_eq,
    Q,
    R,
    G=None,
    x_hat_bar_0=None,
    name=None,
):
    """
    Obtain a continuous-time Infinite Horizon Kalman Filter system for a
    continuous-time plant after linearization at equilibrium point (x_eq, u_eq)

    The input plant contains the deterministic forms of the forward and observation
    operators:

    ```
        dx/dt = f(x,u)
        y = g(x,u)
    ```

    Note: Only plants with one vector-valued input and one vector-valued output
    are currently supported. Furthermore, the plant LeafSystem/Diagram should have
    only one vector-valued integrator.

    A plant with disturbances of the following form is then considered
    following form:

    ```
        dx/dt = f(x,u) + G w
        y = g(x,u) +  v
    ```

    where:

        `w` represents the process noise,
        `v` represents the measurement noise,

    and

    ```
        E(w) = E(v) = 0
        E(ww') = Q
        E(vv') = R
        E(wv') = N = 0
    ```

    This plant with disturbances is linearized (only `f` and `q`) around the
    equilibrium point to obtain:

    ```
        d/dt (x_bar) = A x_bar + B u_bar + G w    --- (C1)
        y_bar = C x_bar + D u_bar + v             --- (C2)
    ```

    where,

    ```
        x_bar = x - x_eq
        u_bar = u - u_eq
        y_bar = y - y_bar
        y_eq = g(x_eq, u_eq)
    ```

    A continuous-time Kalman Filter estimator for the system of equations (C1) and
    (C2) is returned. This filter is in the `x_bar`, `u_bar`, and `y_bar`
    states.

    The returned system will have

    Input ports:
        (0) u_bar : continuous-time control vector relative to equilibrium point
        (1) y_bar : continuous-time measurement vector relative to equilibrium point

    Output ports:
        (1) x_hat_bar : continuous-time state vector estimate relative to
                        equilibrium point

    Parameters:
        plant : a `Plant` object which can be a LeafSystem or a Diagram.
        x_eq: ndarray
            Equilibrium state vector for discretization
        u_eq: ndarray
            Equilibrium control vector for discretization
        Q: ndarray
            Process noise covariance matrix.
        R: ndarray
            Measurement noise covariance matrix.
        G: ndarray
            Process noise matrix. If `None`, `G=B` is assumed making disrurbances
            additive to control vector `u`, i.e. `u_disturbed = u_orig + w`.
        x_hat_bar_0: ndarray
            Initial state estimate relative to equilibrium point.
            If None, an identity matrix is assumed.
    """

    y_eq, linear_plant = linearize_plant(plant, x_eq, u_eq)

    # LTISystem populates its A/B/C/D attributes in initialize(), which runs
    # at context creation — trigger it before reading them (matches the
    # pattern in state_estimators.utils).
    linear_plant.create_context()

    A, B, C, D = linear_plant.A, linear_plant.B, linear_plant.C, linear_plant.D

    nx, nu = B.shape
    ny, _ = D.shape

    if G is None:
        G = B

    if x_hat_bar_0 is None:
        x_hat_bar_0 = npa.zeros(nx)

    # Instantiate a Kalman Filter instance for the linearized plant
    kf = ContinuousTimeInfiniteHorizonKalmanFilter(
        A,
        B,
        C,
        D,
        G,
        Q,
        R,
        x_hat_bar_0,
        name=name,
    )

    return y_eq, kf

CoordinateRotation

Bases: LeafSystem

Computes the rotation of a 3D vector between coordinate systems.

Given sufficient information to construct a rotation matrix C_AB from orthogonal coordinate system B to orthogonal coordinate system A, along with an input vector x_B expressed in B-axes, this block will compute the matrix-vector product x_A = C_AB @ x_B.

Note that depending on the type of rotation representation, this matrix may not be explicitly computed. The types of rotations supported are Quaternion, Euler Angles, and Direction Cosine Matrix (DCM).

By default, the rotations have the following convention:

• Quaternion: The rotation is represented by a 4-component quaternion q. The rotation is carried out by the product p_A = q⁻¹ * p_B * q, where q⁻¹ is the quaternion inverse of q, * is the quaternion product, and p_A and p_B are the quaternion extensions of the vectors x_A and x_B, i.e. p_A = [0, x_A] and p_B = [0, x_B].

• Roll-Pitch-Yaw (Euler Angles): The rotation is represented by the set of Euler angles ϕ (roll), θ (pitch), and ψ (yaw), in the "1-2-3" convention for intrinsic rotations. The resulting rotation matrix C_AB(ϕ, θ, ψ) is the same as the product of the three single-axis rotation matrices C_AB = Cz(ψ) * Cy(θ) * Cx(ϕ).

  For example, if B represents a fixed "world" frame with axes xyz and A is a body-fixed frame with axes XYZ, then C_AB represents a rotation from the world frame to the body frame, in the following sequence:

  1. Right-hand rotation about the world frame x-axis by ϕ (roll), resulting in the intermediate frame x'y'z' with x' = x.
  2. Right-hand rotation about the intermediate frame y'-axis by θ (pitch), resulting in the intermediate frame x''y''z'' with y'' = y'.
  3. Right-hand rotation about the intermediate frame z''-axis by ψ (yaw), resulting in the body frame XYZ with z = z''.

• Direction Cosine Matrix: The rotation is directly represented as a 3x3 matrix C_AB. The rotation is carried out by the matrix-vector product x_A = C_AB @ x_B.

Input ports

(0): The input vector x_B expressed in the B-axes.

(1): (if enable_external_rotation_definition=True) The rotation representation (quaternion, Euler angles, or cosine matrix) that defines the rotation from B to A (or A to B if inverse=True).

Output ports

(0): The output vector x_A expressed in the A-axes.

Parameters:

Name	Type	Description	Default
rotation_type	str	The type of rotation representation to use. Must be one of ("quaternion", "roll_pitch_yaw", "dcm").	required
enable_external_rotation_definition		If True, the block will have one input port for the rotation representation (quaternion, Euler angles, or cosine matrix). Otherwise the rotation must be provided as a block parameter.	True
inverse		If True, the block will compute the inverse transformation, i.e. if the matrix representation of the rotation is C_AB from frame B to frame A, the block will compute the inverse transformation C_BA = C_AB⁻¹ = C_AB.T	False
quaternion	Array	The quaternion representation of the rotation if enable_external_rotation_definition=False.	None
roll_pitch_yaw	Array	The Euler angles representation of the rotation if enable_external_rotation_definition=False.	None
direction_cosine_matrix	Array	The direction cosine matrix representation of the rotation if enable_external_rotation_definition=False.	None

Source code in jaxonomy/library/rotations.py

class CoordinateRotation(LeafSystem):
    """Computes the rotation of a 3D vector between coordinate systems.

    Given sufficient information to construct a rotation matrix `C_AB` from orthogonal
    coordinate system `B` to orthogonal coordinate system `A`, along with an input
    vector `x_B` expressed in `B`-axes, this block will compute the matrix-vector
    product `x_A = C_AB @ x_B`.

    Note that depending on the type of rotation representation, this matrix may not be
    explicitly computed.  The types of rotations supported are Quaternion, Euler
    Angles, and Direction Cosine Matrix (DCM).

    By default, the rotations have the following convention:

    - __Quaternion:__ The rotation is represented by a 4-component quaternion `q`.
        The rotation is carried out by the product `p_A = q⁻¹ * p_B * q`, where
        `q⁻¹` is the quaternion inverse of `q`, `*` is the quaternion product, and
        `p_A` and `p_B` are the quaternion extensions of the vectors `x_A` and `x_B`,
        i.e. `p_A = [0, x_A]` and `p_B = [0, x_B]`.

    - __Roll-Pitch-Yaw (Euler Angles):__ The rotation is represented by the set of Euler angles
        ϕ (roll), θ (pitch), and ψ (yaw), in the "1-2-3" convention for intrinsic
        rotations. The resulting rotation matrix `C_AB(ϕ, θ, ψ)` is the same as the product of
        the three  single-axis rotation matrices `C_AB = Cz(ψ) * Cy(θ) * Cx(ϕ)`.

        For example, if `B` represents a fixed "world" frame with axes `xyz` and `A`
        is a body-fixed frame with axes `XYZ`, then `C_AB` represents a rotation from
        the world frame to the body frame, in the following sequence:

        1. Right-hand rotation about the world frame `x`-axis by `ϕ` (roll), resulting
            in the intermediate frame `x'y'z'` with `x' = x`.
        2. Right-hand rotation about the intermediate frame `y'`-axis by `θ` (pitch),
            resulting in the intermediate frame `x''y''z''` with `y'' = y'`.
        3. Right-hand rotation about the intermediate frame `z''`-axis by `ψ` (yaw),
            resulting in the body frame `XYZ` with `z = z''`.

    - __Direction Cosine Matrix:__ The rotation is directly represented as a
        3x3 matrix `C_AB`. The rotation is carried out by the matrix-vector product
        `x_A = C_AB @ x_B`.

    Input ports:
        (0): The input vector `x_B` expressed in the `B`-axes.

        (1): (if `enable_external_rotation_definition=True`) The rotation
            representation (quaternion, Euler angles, or cosine matrix) that defines
            the rotation from `B` to `A` (or `A` to `B` if `inverse=True`).

    Output ports:
        (0): The output vector `x_A` expressed in the `A`-axes.

    Parameters:
        rotation_type (str): The type of rotation representation to use. Must be one of
            ("quaternion", "roll_pitch_yaw", "dcm").
        enable_external_rotation_definition: If `True`, the block will have one
            input port for the rotation representation (quaternion, Euler angles, or
            cosine matrix).  Otherwise the rotation must be provided as a block
            parameter.
        inverse: If `True`, the block will compute the inverse transformation, i.e.
            if the matrix representation of the rotation is `C_AB` from frame `B` to
            frame `A`, the block will compute the inverse transformation
            `C_BA = C_AB⁻¹ = C_AB.T`
        quaternion (Array, optional): The quaternion representation of the rotation
            if `enable_external_rotation_definition=False`.
        roll_pitch_yaw (Array, optional): The Euler angles representation of the
            rotation if `enable_external_rotation_definition=False`.
        direction_cosine_matrix (Array, optional): The direction cosine matrix
            representation of the rotation if `enable_external_rotation_definition=False`.
    """

    @parameters(
        static=[
            "quaternion",
            "roll_pitch_yaw",
            "direction_cosine_matrix",
            "rotation_type",
            "enable_external_rotation_definition",
            "inverse",
        ]
    )
    def __init__(
        self,
        rotation_type,
        enable_external_rotation_definition=True,
        quaternion=None,
        roll_pitch_yaw=None,
        direction_cosine_matrix=None,
        inverse=False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.external_rotation = enable_external_rotation_definition
        self.rotation_type = rotation_type
        self.inverse = inverse

        self.vector_input_index = self.declare_input_port()

        # Note: all of the possible rotation specifications are passed as parameters
        # to make the serialization work, but only one is valid at a time. This makes
        # sense from the UI, but is a bit strange when working directly with the code.
        # In any case, the typical use case is to have the external rotation port
        # enabled, so all of these should usually be None.  If more than one is
        # provided (which can happen for instance via hidden parameters in the JSON)
        # then only the rotation corresponding to the `rotation_type` will be used, and
        # the rest will be ignored.
        rotation = self._check_config(
            rotation_type,
            quaternion,
            roll_pitch_yaw,
            direction_cosine_matrix,
        )

        if enable_external_rotation_definition:
            self.rotation_input_index = self.declare_input_port()

        else:
            # Store the static rotation as a parameter (will be None if external
            # rotation is enabled)
            self.declare_dynamic_parameter("rotation", rotation)

        self._output_port_idx = self.declare_output_port(
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    def initialize(
        self,
        rotation_type,
        enable_external_rotation_definition,
        quaternion,
        roll_pitch_yaw,
        direction_cosine_matrix,
        inverse,
        rotation=None,
    ):
        if enable_external_rotation_definition != self.external_rotation:
            raise ValueError("Cannot change external rotation definition.")

        self.rotation_type = rotation_type
        self.inverse = inverse
        if not self.external_rotation:
            rotation = self._check_config(
                rotation_type,
                quaternion,
                roll_pitch_yaw,
                direction_cosine_matrix,
            )

            def _output_func(_time, _state, *inputs, **parameters):
                vector = inputs[self.vector_input_index]
                return self._apply(rotation, vector)

        else:

            def _output_func(_time, _state, *inputs, **parameters):
                vector = inputs[self.vector_input_index]
                rotation = inputs[self.rotation_input_index]
                return self._apply(rotation, vector)

        self.configure_output_port(
            self._output_port_idx,
            _output_func,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    def _check_config(
        self, rotation_type, quaternion, roll_pitch_yaw, direction_cosine_matrix
    ):
        if rotation_type not in ("quaternion", "roll_pitch_yaw", "DCM"):
            message = f"Invalid rotation type: {rotation_type}."
            raise BlockParameterError(
                message=message, system=self, parameter_name="rotation_type"
            )

        if self.external_rotation:
            # Input type checking will be done by `check_types`
            return

        if rotation_type == "quaternion":
            if quaternion is None:
                message = (
                    "A static quaternion must be provided if external rotation "
                    + "definition is disabled."
                )
                raise BlockParameterError(
                    message=message, system=self, parameter_name="quaternion"
                )
            rotation = npa.asarray(quaternion)
            if rotation.shape != (4,):
                message = (
                    "The quaternion must have shape (4,), but has shape "
                    + f"{rotation.shape}."
                )
                raise BlockParameterError(
                    message=message, system=self, parameter_name="quaternion"
                )

        elif rotation_type == "roll_pitch_yaw":
            if roll_pitch_yaw is None:
                message = (
                    "A static roll-pitch-yaw sequence must be provided if external "
                    + "rotation definition is disabled."
                )
                raise BlockParameterError(
                    message=message, system=self, parameter_name="roll_pitch_yaw"
                )
            rotation = npa.asarray(roll_pitch_yaw)
            if rotation.shape != (3,):
                message = (
                    "The Euler angles must have shape (3,), but has shape "
                    + f"{rotation.shape}."
                )
                raise BlockParameterError(
                    message=message, system=self, parameter_name="roll_pitch_yaw"
                )

        elif rotation_type == "DCM":
            if direction_cosine_matrix is None:
                message = (
                    "A static direction cosine matrix must be provided if external "
                    + "rotation definition is disabled."
                )
                raise BlockParameterError(
                    message=message,
                    system=self,
                    parameter_name="direction_cosine_matrix",
                )
            rotation = npa.asarray(direction_cosine_matrix)
            if rotation.shape != (3, 3):
                message = (
                    "The direction cosine matrix must have shape (3, 3), but has shape "
                    + f"{rotation.shape}."
                )
                raise BlockParameterError(
                    message=message,
                    system=self,
                    parameter_name="direction_cosine_matrix",
                )

        return rotation

    def _apply(self, rotation: Rotation, vector: Array) -> Array:
        rot = {
            "quaternion": Rotation.from_quat,
            "roll_pitch_yaw": partial(Rotation.from_euler, EULER_SEQ),
            "DCM": Rotation.from_matrix,
        }[self.rotation_type](rotation)

        if self.inverse:
            rot = rot.inv()

        return rot.apply(vector)

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        vec = self.input_ports[self.vector_input_index].eval(context)

        with ErrorCollector.context(error_collector):
            if vec.shape != (3,):
                raise ShapeMismatchError(
                    system=self,
                    expected_shape=(3,),
                    actual_shape=vec.shape,
                )

        if self.external_rotation:
            rot = self.input_ports[self.rotation_input_index].eval(context)

            with ErrorCollector.context(error_collector):
                if self.rotation_type == "quaternion" and rot.shape != (4,):
                    raise ShapeMismatchError(
                        system=self,
                        expected_shape=(4,),
                        actual_shape=rot.shape,
                    )
                elif self.rotation_type == "roll_pitch_yaw" and rot.shape != (3,):
                    raise ShapeMismatchError(
                        system=self,
                        expected_shape=(3,),
                        actual_shape=rot.shape,
                    )
                elif self.rotation_type == "DCM" and rot.shape != (3, 3):
                    raise ShapeMismatchError(
                        system=self,
                        expected_shape=(3, 3),
                        actual_shape=rot.shape,
                    )

CoordinateRotationConversion

Bases: LeafSystem

Converts between different representations of rotations.

See CoordinateRotation block documentation for descriptions of the different rotation representations supported. This block supports conversion between quaternion, roll-pitch-yaw (Euler angles), and direction cosine matrix (DCM).

Note that conversions are reversible in terms of the abstract rotation, although creating a quaternion from a direction cosine matrix (and therefore creating a quaternion from roll-pitch-yaw sequence) results in an arbitrary sign assignment.

Input ports

(0): The input rotation representation.

Output ports

(1): The output rotation representation.

Parameters:

Name	Type	Description	Default
conversion_type	str	The type of rotation conversion to perform. Must be one of ("quaternion_to_euler", "quaternion_to_dcm", "euler_to_quaternion", "euler_to_dcm", "dcm_to_quaternion", "dcm_to_euler")	required

Source code in jaxonomy/library/rotations.py

class CoordinateRotationConversion(LeafSystem):
    """Converts between different representations of rotations.

    See CoordinateRotation block documentation for descriptions of the different
    rotation representations supported. This block supports conversion between
    quaternion, roll-pitch-yaw (Euler angles), and direction cosine matrix (DCM).

    Note that conversions are reversible in terms of the abstract rotation, although
    creating a quaternion from a direction cosine matrix (and therefore creating a
    quaternion from roll-pitch-yaw sequence) results in an arbitrary sign assignment.

    Input ports:
        (0): The input rotation representation.

    Output ports:
        (1): The output rotation representation.

    Parameters:
        conversion_type (str): The type of rotation conversion to perform.
            Must be one of ("quaternion_to_euler", "quaternion_to_dcm",
            "euler_to_quaternion", "euler_to_dcm", "dcm_to_quaternion", "dcm_to_euler")
    """

    @parameters(static=["conversion_type"])
    def __init__(self, conversion_type, **kwargs):
        super().__init__(**kwargs)
        self.declare_input_port()
        self._output_port_idx = self.declare_output_port(requires_inputs=True)

    def initialize(self, conversion_type):
        if conversion_type not in (
            "quaternion_to_RPY",
            "quaternion_to_DCM",
            "RPY_to_quaternion",
            "RPY_to_DCM",
            "DCM_to_quaternion",
            "DCM_to_RPY",
        ):
            message = f"Invalid rotation conversion type: {conversion_type}."
            raise BlockParameterError(
                message=message, system=self, parameter_name="conversion_type"
            )

        _func = {
            "quaternion_to_RPY": quat_to_euler,
            "quaternion_to_DCM": quat_to_dcm,
            "RPY_to_quaternion": euler_to_quat,
            "RPY_to_DCM": euler_to_dcm,
            "DCM_to_quaternion": dcm_to_quat,
            "DCM_to_RPY": dcm_to_euler,
        }[conversion_type]

        def _output(_time, _state, *inputs, **_parameters):
            (u,) = inputs
            return _func(u)

        self.configure_output_port(
            self._output_port_idx,
            _output,
            requires_inputs=True,
        )

        # Serialization
        self.conversion_type = conversion_type

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        rot = self.input_ports[0].eval(context)

        with ErrorCollector.context(error_collector):
            if self.conversion_type in (
                "quaternion_to_RPY",
                "quaternion_to_DCM",
            ) and rot.shape != (4,):
                raise ShapeMismatchError(
                    system=self,
                    expected_shape=(4,),
                    actual_shape=rot.shape,
                )
            elif self.conversion_type in (
                "RPY_to_quaternion",
                "RPY_to_DCM",
            ) and rot.shape != (3,):
                raise ShapeMismatchError(
                    system=self,
                    expected_shape=(3,),
                    actual_shape=rot.shape,
                )
            elif self.conversion_type in (
                "DCM_to_quaternion",
                "DCM_to_RPY",
            ) and rot.shape != (3, 3):
                raise ShapeMismatchError(
                    system=self,
                    expected_shape=(3, 3),
                    actual_shape=rot.shape,
                )

Counter

Bases: LeafSystem

Discrete counter that increments on rising edges of its trigger input.

The block samples a boolean/binary trigger signal every dt seconds. On each rising edge (prev_trigger == False and current_trigger == True) the internal count advances by increment. When max_count is set, the counter either saturates (clamps at max_count) or wraps to 0 after the increment that hits / exceeds max_count, according to reset_on_max.

Input ports

(0) Trigger signal — boolean / binary-valued.

Output ports

(0) Current count (integer, stored as int32).

Parameters:

Name	Type	Description	Default
initial_count		Starting count value at t = 0. Default 0.	0
dt		Sample period (seconds) of the discrete update.	required
increment		Amount to add to the count on each rising edge. Default 1.	1
max_count		Optional cap on the count. If None the counter is unbounded. Default None.	None
reset_on_max		When max_count is set, controls behaviour at saturation. True wraps the count back to 0 once it reaches / exceeds max_count. False clamps the count at max_count. Default False.	False

Notes

Edge detection is the same simple "previous-sample-was-False, current-sample-is-True" rule used by :class:EdgeDetection — adequate for boolean triggers driven at the block's own sample rate. For sub-sample-period precision use ZeroCrossingTriggeredSubsystem together with this block.

The output is integer-typed and therefore non-differentiable in the strict sense; gradient-flow tests should not expect gradients with respect to the count itself.

Source code in jaxonomy/library/sources.py

class Counter(LeafSystem):
    """Discrete counter that increments on rising edges of its trigger input.

    The block samples a boolean/binary trigger signal every ``dt`` seconds.
    On each rising edge (``prev_trigger == False`` and ``current_trigger ==
    True``) the internal count advances by ``increment``. When ``max_count``
    is set, the counter either *saturates* (clamps at ``max_count``) or
    *wraps* to ``0`` after the increment that hits / exceeds ``max_count``,
    according to ``reset_on_max``.

    Input ports:
        (0) Trigger signal — boolean / binary-valued.

    Output ports:
        (0) Current count (integer, stored as ``int32``).

    Parameters:
        initial_count:
            Starting count value at ``t = 0``. Default ``0``.
        dt:
            Sample period (seconds) of the discrete update.
        increment:
            Amount to add to the count on each rising edge. Default ``1``.
        max_count:
            Optional cap on the count. If ``None`` the counter is
            unbounded. Default ``None``.
        reset_on_max:
            When ``max_count`` is set, controls behaviour at saturation.
            ``True`` wraps the count back to ``0`` once it reaches /
            exceeds ``max_count``. ``False`` clamps the count at
            ``max_count``. Default ``False``.

    Notes:
        Edge detection is the same simple "previous-sample-was-False,
        current-sample-is-True" rule used by :class:`EdgeDetection` —
        adequate for boolean triggers driven at the block's own sample
        rate. For sub-sample-period precision use
        ``ZeroCrossingTriggeredSubsystem`` together with this block.

        The output is integer-typed and therefore non-differentiable in
        the strict sense; gradient-flow tests should not expect
        gradients with respect to the count itself.
    """

    class _DiscreteStateType(NamedTuple):
        prev_trigger: Array
        count: Array

    @parameters(
        dynamic=["initial_count"],
        static=["dt", "increment", "max_count", "reset_on_max"],
    )
    def __init__(
        self,
        dt,
        initial_count=0,
        increment=1,
        max_count=None,
        reset_on_max=False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.dt = dt
        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(
        self,
        initial_count,
        dt=None,
        increment=1,
        max_count=None,
        reset_on_max=False,
    ):
        if increment is None:
            raise BlockParameterError(
                message=f"Counter block {self.name} requires non-None increment.",
            )
        if max_count is not None and int(max_count) <= 0:
            raise BlockParameterError(
                message=(
                    f"Counter block {self.name} requires max_count > 0 "
                    f"(got {max_count})."
                ),
            )

        # Store static config as plain Python ints — they're declared as
        # static @parameters so they will not be JAX-traced.
        self._increment = int(increment)
        self._max_count = None if max_count is None else int(max_count)
        self._reset_on_max = bool(reset_on_max)

        count0 = npa.asarray(int(initial_count), dtype=npa.int32)
        prev0 = npa.asarray(False, dtype=npa.bool_)
        self.declare_discrete_state(
            default_value=self._DiscreteStateType(
                prev_trigger=prev0,
                count=count0,
            ),
            as_array=False,
        )
        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=0.0,
        )
        self.configure_output_port(
            self._output_port_idx,
            self._output,
            prerequisites_of_calc=[DependencyTicket.xd],
            requires_inputs=False,
            default_value=count0,
        )

    def reset_default_values(
        self,
        initial_count,
        dt=None,
        increment=1,
        max_count=None,
        reset_on_max=False,
    ):
        count0 = npa.asarray(int(initial_count), dtype=npa.int32)
        prev0 = npa.asarray(False, dtype=npa.bool_)
        self.configure_discrete_state_default_value(
            default_value=self._DiscreteStateType(
                prev_trigger=prev0,
                count=count0,
            ),
            as_array=False,
        )
        self.configure_output_port_default_value(self._output_port_idx, count0)

    def _update(self, _time, state, *inputs, **_params):
        (trigger,) = inputs
        # Cast trigger to bool so float/int/bool sources all yield the
        # same rising-edge semantics; mirrors EdgeDetection.
        trig = npa.asarray(trigger, dtype=npa.bool_)
        prev = state.discrete_state.prev_trigger
        count = state.discrete_state.count
        # Rising edge: previous sample was False, current sample is True.
        rising = npa.logical_and(npa.logical_not(prev), trig)

        # Tentative incremented count if a rising edge fired.
        inc = npa.asarray(self._increment, dtype=count.dtype)
        next_count = count + inc

        if self._max_count is not None:
            cap = npa.asarray(self._max_count, dtype=count.dtype)
            zero = npa.asarray(0, dtype=count.dtype)
            if self._reset_on_max:
                # Wrap: once the post-increment value would reach or
                # exceed the cap, wrap back to 0.
                next_count = npa.where(next_count >= cap, zero, next_count)
            else:
                # Saturate: clamp the post-increment value at the cap.
                next_count = npa.where(next_count > cap, cap, next_count)

        new_count = npa.where(rising, next_count, count)
        return self._DiscreteStateType(
            prev_trigger=trig,
            count=new_count,
        )

    def _output(self, _time, state, *_inputs, **_params):
        return state.discrete_state.count

CrossProduct

Bases: ReduceBlock

Compute the cross product between the inputs.

See NumPy docs for details: https://numpy.org/doc/stable/reference/generated/numpy.cross.html

Input ports

(0) The first input vector. (1) The second input vector.

Output ports

(0) The cross product of the inputs.

Source code in jaxonomy/library/math_ops.py

class CrossProduct(ReduceBlock):
    """Compute the cross product between the inputs.

    See NumPy docs for details:
    https://numpy.org/doc/stable/reference/generated/numpy.cross.html

    Input ports:
        (0) The first input vector.
        (1) The second input vector.

    Output ports:
        (0) The cross product of the inputs.
    """

    def __init__(self, *args, **kwargs):
        def _cross(inputs):
            return npa.cross(*inputs)

        super().__init__(2, _cross, *args, **kwargs)

CustomJaxBlock

Bases: LeafSystem

JAX implementation of the PythonScript block.

A few important notes and changes/limitations to this JAX implementation: - For this block all code must be written using the JAX-supported subset of Python: * Numerical operations should use jax.numpy = jnp instead of numpy = np * Standard control flow is not supported (if/else, for, while, etc.). Instead use lax.cond, lax.fori_loop, lax.while_loop, etc. https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#structured-control-flow-primitives Where possible, NumPy-style operations like jnp.where or jnp.select should be preferred to lax control flow primitives. * Functions must be pure and arrays treated as immutable. https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#in-place-updates Provided these assumptions hold, the code can be JIT compiled, differentiated, run on GPU, etc. - Variable scoping: the init_code and step_code are executed in the same scope, so variables declared in the init_code will be available in the step_code and can be modified in that scope. Internally, everything declared in init_code is treated as a single state-like cache entry. However, variables declared in the step_code will NOT persist between evaluations. Users should think of step_code as a normal Python function where locally declared variables will disappear on leaving the scope. - Persistent variables (outputs and anything declared in init_code) must have static shapes and dtypes. This means that you cannot declare x = 0.0 in init_code and then later assign x = jnp.zeros(4) in step_code.

These changes mean that many older PythonScript blocks may not be backwards compatible.

Input ports

Variable number of input ports, one for each input variable declared in inputs. The order of the input ports is the same as the order of the input variables.

Output ports

Variable number of output ports, one for each output variable declared in outputs. The order of the output ports is the same as the order of the output variables.

Parameters:

Name	Type	Description	Default
dt	float	The discrete time step of the block, or None if the block is in agnostic time mode.	None
init_script	str	A string containing Python code that will be executed once when the block is initialized. This code can be used to declare persistent variables that will be available in the step_code.	''
user_statements	str	A string containing Python code that will be executed once per time step (or per output port evaluation, in agnostic mode). This code can use the persistent variables declared in init_script and the block inputs.	''
finalize_script	str	A string containing Python code that will be executed once when the simulation completes (or is otherwise torn down). This code can use the persistent variables declared in init_script. Supported for :class:CustomPythonBlock only; raises :class:PythonScriptError for :class:CustomJaxBlock.	''
accelerate_with_jax	bool	If True, the block will be JIT compiled. If False, the block will be executed in pure Python. This parameter exists for compatibility with UI options; when creating pure Python blocks from code (e.g. for testing), explicitly create the CustomPythonBlock class.	True
time_mode	str	One of "discrete" or "agnostic". If "discrete", the block step code will be evaluated at peridodic intervals specified by "dt". If "agnostic", the block step code will be evaluated once per output port evaluation, and the block will not have a discrete time step.	'discrete'
inputs	List[str]	A list of input variable names. The order of the input ports is the same as the order of the input variables.	None
outputs	Mapping[str, Tuple[DTypeLike, ShapeLike]]	A dictionary mapping output variable names to a tuple of dtype and shape. The order of the output ports is the same as the order of the output variables.	None
static_parameters	Mapping[str, Array]	A dictionary mapping parameter names to values. Parameters are treated as immutable and cannot be modified in the step code. Static parameters can't be used in ensemble simulations or optimization workflows.	None
dynamic_parameters	Mapping[str, Array]	A dictionary mapping parameter names to values. Parameters are treated as immutable and cannot be modified in the step code. Dynamic parameters can be arrays or scalars, but must have static shapes and dtypes in order to support JIT compilation.	None

Source code in jaxonomy/library/custom.py

class CustomJaxBlock(LeafSystem):
    """JAX implementation of the PythonScript block.

    A few important notes and changes/limitations to this JAX implementation:
    - For this block all code must be written using the JAX-supported subset of Python:
        * Numerical operations should use `jax.numpy = jnp` instead of `numpy = np`
        * Standard control flow is not supported (if/else, for, while, etc.). Instead
            use `lax.cond`, `lax.fori_loop`, `lax.while_loop`, etc.
            https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#structured-control-flow-primitives
            Where possible, NumPy-style operations like `jnp.where` or `jnp.select` should
            be preferred to lax control flow primitives.
        * Functions must be pure and arrays treated as immutable.
            https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#in-place-updates
        Provided these assumptions hold, the code can be JIT compiled, differentiated,
        run on GPU, etc.
    - Variable scoping: the `init_code` and `step_code` are executed in the same scope,
        so variables declared in the `init_code` will be available in the `step_code`
        and can be modified in that scope. Internally, everything declared in
        `init_code` is treated as a single state-like cache entry.
        However, variables declared in the `step_code` will NOT persist between
        evaluations. Users should think of `step_code` as a normal Python function
        where locally declared variables will disappear on leaving the scope.
    - Persistent variables (outputs and anything declared in `init_code`) must have
        static shapes and dtypes. This means that you cannot declare `x = 0.0` in
        `init_code` and then later assign `x = jnp.zeros(4)` in `step_code`.

    These changes mean that many older PythonScript blocks may not be backwards compatible.

    Input ports:
        Variable number of input ports, one for each input variable declared in `inputs`.
        The order of the input ports is the same as the order of the input variables.

    Output ports:
        Variable number of output ports, one for each output variable declared in `outputs`.
        The order of the output ports is the same as the order of the output variables.

    Parameters:
        dt (float): The discrete time step of the block, or None if the block is
            in agnostic time mode.
        init_script (str): A string containing Python code that will be executed
            once when the block is initialized. This code can be used to declare
            persistent variables that will be available in the `step_code`.
        user_statements (str): A string containing Python code that will be executed
            once per time step (or per output port evaluation, in agnostic mode).
            This code can use the persistent variables declared in `init_script` and
            the block inputs.
        finalize_script (str): A string containing Python code that will be executed
            once when the simulation completes (or is otherwise torn down). This code
            can use the persistent variables declared in `init_script`. Supported for
            :class:`CustomPythonBlock` only; raises :class:`PythonScriptError` for
            :class:`CustomJaxBlock`.
        accelerate_with_jax (bool): If True, the block will be JIT compiled. If False,
            the block will be executed in pure Python.  This parameter exists for
            compatibility with UI options; when creating pure Python blocks from code
            (e.g. for testing), explicitly create the CustomPythonBlock class.
        time_mode (str): One of "discrete" or "agnostic". If "discrete", the block
            step code will be evaluated at peridodic intervals specified by "dt".
            If "agnostic", the block step code will be evaluated once per output
            port evaluation, and the block will not have a discrete time step.
        inputs (List[str]): A list of input variable names. The order of the input
            ports is the same as the order of the input variables.
        outputs (Mapping[str, Tuple[DTypeLike, ShapeLike]]): A dictionary mapping
            output variable names to a tuple of dtype and shape. The order of the
            output ports is the same as the order of the output variables.
        static_parameters (Mapping[str, Array]): A dictionary mapping parameter names to
            values. Parameters are treated as immutable and cannot be modified in
            the step code. Static parameters can't be used in ensemble simulations or
            optimization workflows.
        dynamic_parameters (Mapping[str, Array]): A dictionary mapping parameter names to
            values. Parameters are treated as immutable and cannot be modified in
            the step code. Dynamic parameters can be arrays or scalars, but must have static
            shapes and dtypes in order to support JIT compilation.
    """

    @declare_parameters(
        static=[
            "dt",
            "init_script",
            "user_statements",
            "finalize_script",
            "accelerate_with_jax",
            "time_mode",
        ]
    )
    def __init__(
        self,
        dt: float = None,
        init_script: str = "",
        user_statements: str = "",
        finalize_script: str = "",
        accelerate_with_jax: bool = True,
        time_mode: str = "discrete",  # [discrete, agnostic]
        inputs: List[str] = None,  # [name]
        outputs: List[str] = None,
        dynamic_parameters: Mapping[str, Array] = None,
        static_parameters: Mapping[str, Array] = None,
        strict: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        dynamic_parameters = dynamic_parameters if dynamic_parameters else {}
        static_parameters = static_parameters if static_parameters else {}

        if time_mode not in ["discrete", "agnostic"]:
            raise BlockInitializationError(
                f"Invalid time mode '{time_mode}' for PythonScript block", system=self
            )

        if time_mode == "discrete" and dt is None:
            raise BlockInitializationError(
                "When in discrete time mode, dt is required for block", system=self
            )

        self.time_mode = time_mode

        if inputs is None:
            inputs = []
        if outputs is None:
            outputs = []

        # T-036e: validate inputs / outputs port-name lists eagerly so opaque
        # downstream errors (duplicate-port AssertionError, AttributeError on
        # exec'd statement, etc.) are replaced by a clear named-block error.
        _validate_custom_block_io_names(self, inputs, outputs)

        # T-036e (deeper): when strict=True, AST-walk the init / step /
        # finalize scripts to catch typo'd symbol references at construction
        # time rather than at first eval.  Default-off for backwards-compat.
        if strict:
            param_names = list(
                (dynamic_parameters or {}).keys()
            ) + list((static_parameters or {}).keys())
            # CustomJaxBlock binds ``time`` in the exec env; CustomPythonBlock
            # historically did not (see WC-98 comment at exec_step) — but
            # both subclasses funnel through here, so we allow ``time`` and
            # let the t-vs-time hint catch the legacy mistake.
            _validate_custom_block_signature(
                self,
                init_script=init_script,
                user_statements=user_statements,
                finalize_script=finalize_script,
                inputs=inputs,
                outputs=outputs,
                parameter_names=param_names,
                has_time_binding=True,
            )

        self.dt = dt

        # Note: 'optimize' level could be lowered in debug mode
        try:
            self.init_code = compile(
                init_script, filename="<init>", mode="exec", optimize=2
            )
        except BaseException as e:
            raise PythonScriptError(
                f"Syntax error in init_script for PythonScript block '{self.name_path_str}': {e}",
                system=self,
            ) from e

        try:
            self.step_code = compile(
                user_statements, filename="<step>", mode="exec", optimize=2
            )
        except BaseException as e:
            raise PythonScriptError(
                f"Syntax error in user_statements for PythonScript block '{self.name_path_str}': {e}",
                system=self,
            ) from e

        self._has_finalize_script = bool(finalize_script.strip())
        try:
            self.finalize_code = compile(
                finalize_script, filename="<finalize>", mode="exec", optimize=2
            )
        except BaseException as e:
            raise PythonScriptError(
                f"Syntax error in finalize_script for PythonScript block '{self.name_path_str}': {e}",
                system=self,
            ) from e

        if finalize_script != "" and not isinstance(self, CustomPythonBlock):
            raise PythonScriptError(
                f"PythonScript block '{self.name_path_str}' has a finalize_script "
                "but this is only supported for CustomPythonBlock (non-JAX) blocks.",
                system=self,
                parameter_name="finalize_script",
            )

        # Declare parameters
        for param_name, value in dynamic_parameters.items():
            if isinstance(value, list):
                value = npa.asarray(value)
            as_array = isinstance(value, npa.ndarray) or npa.isscalar(value)
            self.declare_dynamic_parameter(param_name, value, as_array=as_array)

        for param_name, value in static_parameters.items():
            self.declare_static_parameter(param_name, value)

        # Run the init_script
        persistent_env = self.exec_init()

        # Declare an input port for each of the input variables
        self.input_names = inputs
        for name in inputs:
            self.declare_input_port(name)

        # Declare a cache component for each of the output variables
        self._create_cache_type(outputs)

        if time_mode == "discrete":
            self._configure_discrete(dt, outputs, persistent_env)
        else:
            self._configure_agnostic(outputs, persistent_env)

    def initialize(
        self,
        dt: float = None,
        init_script: str = "",
        user_statements: str = "",
        finalize_script: str = "",
        accelerate_with_jax: bool = True,
        time_mode: str = "discrete",  # [discrete, agnostic]
        **parameters,
    ):
        pass

    def _initialize_outputs(self, outputs, persistent_env):
        default_outputs = {name: None for name in outputs}

        for name in outputs:
            # If the initial value is set explicitly in the init script,
            # override the default value.  We don't need to do this for
            # agnostic configuration since the outputs will be calculated
            # every evaluation anyway.
            if name in persistent_env:
                value = npa.asarray(persistent_env[name])
                default_outputs[name] = value

                # Also update the persistent environment so that the data types
                # are consistent with the state.
                persistent_env[name] = value

            # Otherwise throw an error, since we don't know what the initial values
            # should be, or even what shape/dtype they should have.
            else:
                msg = (
                    f"Output variable '{name}' not explicitly initialized in "
                    "init_script for PythonScript block in 'Discrete' time mode. "
                    "Either initialize the variable as an array with the correct "
                    "shape and dtype, or make the block time mode 'Agnostic'."
                )
                raise PythonScriptError(message=msg, system=self)

        return self.CacheType(
            persistent_env=persistent_env,
            **default_outputs,
        )

    def _configure_discrete(self, dt, outputs, persistent_env):
        default_values = self._initialize_outputs(outputs, persistent_env)

        # The step function acts as a periodic update that will update all components
        # of the discrete state.
        self.step_callback_index = self.declare_cache(
            self.exec_step,
            period=dt,
            offset=dt,
            requires_inputs=True,
            default_value=default_values,
        )

        cache = self.callbacks[self.step_callback_index]

        # Get the index into the state cache (different in general from the index
        # into the callback list, since not all callbacks are cached).
        self.step_cache_index = cache.cache_index

        def _make_callback(o_port_name):
            def _output(time, state, *inputs, **parameters):
                return getattr(state.cache[self.step_cache_index], o_port_name)

            return _output

        # Declare output ports for each state variable
        for o_port_name in outputs:
            self.declare_output_port(
                _make_callback(o_port_name),
                name=o_port_name,
                prerequisites_of_calc=[cache.ticket],
                requires_inputs=False,
                period=dt,
                offset=0.0,
            )

    def _configure_agnostic(self, outputs, persistent_env):
        # Create a callback to evaluate the step code and extract the
        # output. Note that this is inefficient since the step code will
        # be evaluated once _for each output port_, but it's the only way
        # to do this unless (until) we implement some variety of block
        # or function pre-ordering.
        def _make_callback(o_port_name):
            def _output(time, state, *inputs, **parameters):
                xd = self.exec_step(time, state, *inputs, **parameters)
                return getattr(xd, o_port_name)

            return _output

        # Declare output ports for each state variable
        for o_port_name in outputs:
            self.declare_output_port(
                jit(_make_callback(o_port_name)),
                name=o_port_name,
                requires_inputs=True,
            )

        # This callback doesn't need to do anything since it's never
        # actually called - the cache here just stores the initial environment
        # and the output ports are evaluated directly.  This should be changed
        # to avoid re-evaluation with multiple output ports once we can do full
        # function ordering.
        def _cache_callback(time, state, *inputs, **parameters):
            return state.cache[self.step_cache_index]

        # Since this is the return type for `exec_step` we have to declare all
        # the output ports as entries in the namedtuple, even though those values
        # won't actually be cached in "agnostic" time mode.  This is just so that
        # both "discrete" and "agnostic" modes can share the same code.
        default_values = self.CacheType(
            persistent_env=persistent_env,
            **{o_port_name: None for o_port_name in outputs},
        )
        self.step_callback_index = self.declare_cache(
            _cache_callback,
            default_value=default_values,
            requires_inputs=False,
            prerequisites_of_calc=[inport.ticket for inport in self.input_ports],
        )

        cache = self.callbacks[self.step_callback_index]
        self.step_cache_index = cache.cache_index

    def _create_cache_type(self, outputs):
        # Store the output ports as a name for type inference and casting
        self.output_names = outputs

        # Also store the dictionary of local environment variables as a cache entry
        # This is the only persistent state of the system (besides outputs) - anything
        # declared in the "step" function will be forgotten at the end of the step

        self.CacheType = namedtuple("CacheType", self.output_names + ["persistent_env"])

    @property
    def local_env_base(self):
        # Define a starting point for the local code execution environment.
        # we have to inclide __main__ so that the code behaves like a module.
        # this allows for code like this:
        #   imports ...
        #   a = 1
        #   def f(b):
        #       return a+b
        #   out_0 = f(2)
        #
        # without getting a 'a not defined' error.
        return {
            "__main__": {},
        }

    def exec_init(self) -> dict[str, Array]:
        # Before executing the step code, we have to build up the local environment.
        # This includes specified modules, python block user defined parameters.

        default_parameters = {
            name: param.get() for name, param in self.dynamic_parameters.items()
        }

        local_env = {
            **self.local_env_base,
            **default_parameters,
        }

        # similar to above where we included __main__ so the code behaves as a module,
        # here we have to pass the local_env with __main__ as 1] globals, since that
        # is what allow the code to be executed as a module. 2] local since that is where
        # the new bindings will be written, that we need to retain since the code in step_code
        # may depend on these bindings.
        try:
            _default_exec(
                self.init_code,
                local_env,
                logger_=logger,
                system=self,
                code_name="init",
            )

        except BaseException as e:
            logger.error(
                "PythonScript block '%s' init script failed",
                self.name_path_str,
                **logdata(block=self),
            )
            raise PythonScriptError(system=self) from e

        # persistent_env contains bindings for parameters and for values from init_script
        persistent_env, static_env = _filter_non_traceable(local_env)

        # Since this is called during block initialization and not any JIT-compiled code,
        # we can safely store any untraceable variables as block attributes.  For example,
        # this may contain custom functions, classes, etc.
        self.static_env = static_env

        return persistent_env

    def exec_step(self, time: float, state: LeafState, *inputs, **parameters):
        # Before executing the step code, we have to build up the local environment.
        # This includes the persistent variables (anything declared in `init_code`),
        # time, block inputs, user-defined parameters, and specified modules.

        # Retrieve the variables declared in `init_code` from the discrete state
        full_env = state.cache[self.step_cache_index]
        persistent_env = full_env.persistent_env

        # Inputs are in order of port declaration, so they match `self.input_names`
        input_env = dict(zip(self.input_names, inputs))

        # Create a dictionary of all the information that the step function will need
        base_copy = self.local_env_base.copy()
        local_env = {
            **self.static_env,
            **base_copy,
            **persistent_env,
            **input_env,
            **parameters,
        }

        # Execute the step code in the local environment
        try:
            _default_exec(
                self.step_code,
                local_env,
                logger_=logger,
                inputs=input_env,
                system=self,
                code_name="step",
            )

        except PythonScriptError:
            raise
        except BaseException as e:
            logger.error(
                "PythonScript block '%s' step failed.",
                self.name_path_str,
                **logdata(block=self),
            )
            raise PythonScriptError(system=self) from e

        # Updated state variables are stored in the local environment
        xd = {name: local_env[name] for name in self.output_names}

        # Store the persistent variables in the corresponding discrete state
        xd["persistent_env"] = {key: local_env[key] for key in persistent_env}

        # Make sure the results have a consistent data type
        for name in self.output_names:
            xd[name] = npa.asarray(local_env[name])

            # Also make sure the value stored in the persistent environment
            # has the same data type
            if name in persistent_env:
                xd["persistent_env"][name] = xd[name]

        return self.CacheType(**xd)

    def check_types(
        self,
        context: ContextBase,
        error_collector: ErrorCollector = None,
    ):
        """Test-compile the init and step code to check for errors."""
        try:
            # Note that exec_step doesn't use parameters or time
            inputs = self.collect_inputs(context)
            jit(self.exec_step)(None, context[self.system_id].state, *inputs)
        except BaseException as exc:
            with ErrorCollector.context(error_collector):
                name_error = _caused_by_nameerror(exc)
                if name_error and name_error.name == "time":
                    raise PythonScriptTimeNotSupportedError(system=self) from exc
                if isinstance(exc, PythonScriptError):
                    raise
                raise PythonScriptError(system=self) from exc

check_types(context, error_collector=None)

Test-compile the init and step code to check for errors.

Source code in jaxonomy/library/custom.py

def check_types(
    self,
    context: ContextBase,
    error_collector: ErrorCollector = None,
):
    """Test-compile the init and step code to check for errors."""
    try:
        # Note that exec_step doesn't use parameters or time
        inputs = self.collect_inputs(context)
        jit(self.exec_step)(None, context[self.system_id].state, *inputs)
    except BaseException as exc:
        with ErrorCollector.context(error_collector):
            name_error = _caused_by_nameerror(exc)
            if name_error and name_error.name == "time":
                raise PythonScriptTimeNotSupportedError(system=self) from exc
            if isinstance(exc, PythonScriptError):
                raise
            raise PythonScriptError(system=self) from exc

CustomPythonBlock

Bases: CustomJaxBlock

Container for arbitrary user-defined Python code.

Implemented to support legacy PythonScript blocks.

Not traceable (no JIT compilation or autodiff). The internal implementation and behavior of this block differs vastly from the JAX-compatible block as this block stores state directly within the Python instance. Objects and modules can be kept as discrete state.

Note that in "agnostic" mode, the step code will be evaluated once per output port evaluation. Because locally defined environment variables (in the init script) are preserved between evaluations, any mutation of these variables will be preserved. This can lead to unexpected behavior and should be avoided. Stateful behavior should be implemented using discrete state variables instead.

Warning: The finalize_script parameter is currently accepted but not executed. This is a known limitation. Do not rely on finalize_script for cleanup operations.

Source code in jaxonomy/library/custom.py

class CustomPythonBlock(CustomJaxBlock):
    """Container for arbitrary user-defined Python code.

    Implemented to support legacy PythonScript blocks.

    Not traceable (no JIT compilation or autodiff). The internal implementation
    and behavior of this block differs vastly from the JAX-compatible block as
    this block stores state directly within the Python instance. Objects
    and modules can be kept as discrete state.

    Note that in "agnostic" mode, the step code will be evaluated _once per
    output port evaluation_. Because locally defined environment variables
    (in the init script) are preserved between evaluations, any mutation of
    these variables will be preserved. This can lead to unexpected behavior
    and should be avoided. Stateful behavior should be implemented using
    discrete state variables instead.

    Warning: The finalize_script parameter is currently accepted but not executed. 
    This is a known limitation. Do not rely on finalize_script for cleanup operations.
    """

    __exec_fn = _default_exec

    def __init__(
        self,
        dt: float = None,
        init_script: str = "",
        user_statements: str = "",
        finalize_script: str = "",
        inputs: List[str] = None,  # [name]
        outputs: List[str] = None,
        accelerate_with_jax: bool = False,
        time_mode: str = "discrete",
        static_parameters: Mapping[str, Array] = None,
        strict: bool = False,
        **kwargs,
    ):
        self._static_data_initialized = False
        self._parameters = static_parameters or {}
        self._persistent_env = {}

        # Per-block mutable module state (e.g. numpy error policy, random seed).
        # Populated after exec_init; applied + restored around every exec_step.
        # This provides isolation between multiple CustomPythonBlock instances
        # that share the same module objects from sys.modules.
        self._block_module_state: dict = {}

        # Will populate return type information during static initialization
        self.result_shape_dtypes = None
        self.return_dtypes = None

        super().__init__(
            dt=dt,
            init_script=init_script,
            user_statements=user_statements,
            finalize_script=finalize_script,
            inputs=inputs,
            outputs=outputs,
            accelerate_with_jax=accelerate_with_jax,
            time_mode=time_mode,
            static_parameters=self._parameters,
            strict=strict,
            **kwargs,
        )

        if time_mode == "agnostic" and npa.active_backend == "jax":
            logger.warning(
                "System %s is in agnostic time mode but is not traced with JAX. Be "
                "advised that the step code will be evaluated once per output port "
                "evaluation. Any mutation of the local environment should be strictly "
                "avoided as it will likely lead to unexpected behavior.",
                self.name_path_str,
            )

    def initialize(self, **kwargs):
        pass

    @property
    def has_feedthrough_side_effects(self) -> bool:
        # See explanation in `SystemBase.has_ode_side_effects`.
        return self.time_mode == "agnostic"

    @staticmethod
    def set_exec_fn(exec_fn: callable):
        CustomPythonBlock.__exec_fn = exec_fn

    @property
    def local_env_base(self):
        # Define a starting point for the local code execution environment.
        return {
            "__main__": {},
            "true": True,
            "false": False,
        }

    def exec_init(self) -> None:
        default_parameters = {
            name: param.get() for name, param in self.dynamic_parameters.items()
        }

        local_env = {
            **self.local_env_base,
            **self._parameters,
            **default_parameters,
        }

        exec_fn = functools.partial(
            CustomPythonBlock.__exec_fn,
            code=self.init_code,
            env=local_env,
            logger_=logger,
            system=self,
            code_name="init",
        )

        # Snapshot global module state BEFORE running init_script so we can
        # detect which changes the init_script introduces (e.g. numpy.seterr).
        pre_init_global = _save_module_state(local_env)

        try:
            io_callback(exec_fn, None)
        except KeyboardInterrupt as e:
            logger.error(
                "Python block '%s' init script execution was interrupted.",
                self.name,
                **logdata(block=self),
            )
            raise PythonScriptError(
                message="Python block init script execution was interrupted.",
                system=self,
            ) from e
        except PythonScriptError as e:
            logger.error("%s: init script failed.", self.name, **logdata(block=self))
            raise e
        except BaseException as e:
            logger.error("%s: init script failed.", self.name, **logdata(block=self))
            raise PythonScriptError(system=self) from e

        # Capture the module state after init_script ran.  This becomes the
        # block's "initial" isolated state for subsequent exec_step calls.
        self._block_module_state = _save_module_state(local_env)

        # Restore the global module state so that this block's init_script
        # does not contaminate other blocks' initialisation.
        _restore_module_state(pre_init_global)

        self._persistent_env = local_env

        return None

    def exec_step(self, time, state, *inputs, **parameters):
        if not self._static_data_initialized:
            # return_dtypes is inferred in initialize_static_data()
            raise PythonScriptError(
                "Trying to execute step code before static data has been initialized",
                system=self,
            )
        logger.debug(
            "Executing step for %s with state=%s, inputs=%s",
            self.name,
            state,
            inputs,
        )

        # Inputs are in order of port declaration, so they match `self.input_names`
        input_env = dict(zip(self.input_names, inputs))

        base_copy = self.local_env_base.copy()
        local_env = {
            **base_copy,
            **self._persistent_env,
            **parameters,
        }

        exec_fn = functools.partial(
            CustomPythonBlock.__exec_fn,
            code=self.step_code,
            env=local_env,
            logger_=logger,
            return_vars=self.output_names,
            return_dtypes=self.return_dtypes,
            system=self,
            code_name="step",
        )

        def wrapped_exec_fn(inputs):
            # --- module isolation: checkpoint / apply / restore ---
            # 1. Save the current global state (may have been mutated by another block)
            global_state_before = _save_module_state(local_env)
            # 2. Apply this block's remembered module state
            _restore_module_state(self._block_module_state)
            try:
                result = exec_fn(inputs=inputs)
            except KeyboardInterrupt:
                logger.error(
                    "Python block '%s' step script execution was interrupted.",
                    self.name,
                    **logdata(block=self),
                )
                raise
            except NameError as e:
                err_msg = (
                    f"Python block '{self.name}' step script execution failed with a NameError on"
                    + f" missing variable '{e.name}'."
                    + " All names used in this script should be declared in the init script."
                    + f" The execution environment contains the following names: {', '.join(list(local_env.keys()))}"
                )
                logger.error(err_msg)
                logger.error("NameError: %s", e, **logdata(block=self))
                raise PythonScriptError(system=self) from e
            except PythonScriptError as e:
                logger.error("%s: exec_step failed.", self.name, **logdata(block=self))
                raise e
            except BaseException as e:
                logger.error("%s: exec_step failed.", self.name, **logdata(block=self))
                raise PythonScriptError(system=self) from e
            else:
                # 3. Capture any module state changes made by this block's step
                self._block_module_state = _save_module_state(local_env)
                return result
            finally:
                # 4. Always restore the global state that was in place before
                #    this block ran so that other blocks are unaffected.
                _restore_module_state(global_state_before)

        return_vars = io_callback(
            wrapped_exec_fn,
            self.result_shape_dtypes,
            inputs=input_env,
        )

        # Keep local env for next step but only if defined in init_script
        # NOTE: If this restriction turns out to be counterproductive, we can
        # remove it and remove the NameError handling above as well. The thinking
        # here is that this could help avoiding stuff like `if time == 0: x = 0`
        # See https://jaxonomy.atlassian.net/browse/WC-98
        self._persistent_env = {
            key: local_env[key] for key in self._persistent_env if key in local_env
        }

        # Updated state variables are stored in the local environment
        xd = {name: return_vars[i] for i, name in enumerate(self.output_names)}

        return self.CacheType(persistent_env=None, **xd)

    def _initialize_outputs(self, outputs, _persistent_env):
        # Override the base implemenetation since `persistent_env` will be None
        # in this case. Instead, pass the class attribute where the environment
        # is actually maintained.
        default_outputs = {name: None for name in outputs}
        default_values = self.CacheType(
            persistent_env=self._persistent_env,
            **default_outputs,
        )
        default_values = super()._initialize_outputs(outputs, self._persistent_env)
        default_outputs = default_values._asdict()
        self._persistent_env = default_outputs.pop("persistent_env")

        # Determine return data types
        self._initialize_result_shape_dtypes(
            [default_outputs[output] for output in outputs]
        )

        return self.CacheType(
            persistent_env=None,
            **default_outputs,
        )

    def _initialize_result_shape_dtypes(self, outputs):
        self.result_shape_dtypes = []
        self.return_dtypes = []
        for value in outputs:
            self.result_shape_dtypes.append(
                jax.ShapeDtypeStruct(value.shape, value.dtype)
            )
            self.return_dtypes.append(value.dtype)

    def initialize_static_data(self, context):
        # If in agnostic mode, call the step function once to determine the
        # data types and then store those in result_shape_dtype and return_dtypes.
        context = LeafSystem.initialize_static_data(self, context)

        if self.result_shape_dtypes is not None:
            # These data types are already known (block is in discrete mode)
            self._static_data_initialized = True
            return context

        inputs = self.collect_inputs(context)
        input_env = dict(zip(self.input_names, inputs))

        base_copy = self.local_env_base.copy()
        local_env = {
            **base_copy,
            **self._persistent_env,
        }

        # Will not do any type conversion
        return_dtypes = [None for _ in self.output_names]

        exec_fn = functools.partial(
            CustomPythonBlock.__exec_fn,
            self.step_code,
            local_env,
            logger_=logger,
            return_vars=self.output_names,
            return_dtypes=return_dtypes,
            system=self,
            code_name="step",
        )

        return_vars = exec_fn(inputs=input_env)

        self._initialize_result_shape_dtypes(return_vars)

        self._static_data_initialized = True

        return context

    def exec_finalize(self) -> None:
        """Execute the finalize_script using the current persistent environment.

        Called once at the end of the simulation via :meth:`post_simulation_finalize`.
        The script runs in the same environment that was maintained throughout the
        simulation, so all variables declared in ``init_script`` (and updated by
        ``user_statements``) are available.

        Has no effect if ``finalize_script`` is empty.
        """
        if not self._has_finalize_script:
            return

        local_env = {
            **self.local_env_base,
            **self._persistent_env,
        }

        exec_fn = functools.partial(
            CustomPythonBlock.__exec_fn,
            code=self.finalize_code,
            env=local_env,
            logger_=logger,
            system=self,
            code_name="finalize",
        )

        try:
            io_callback(exec_fn, None)
        except KeyboardInterrupt as e:
            logger.error(
                "Python block '%s' finalize script execution was interrupted.",
                self.name,
                **logdata(block=self),
            )
            raise PythonScriptError(
                message="Python block finalize script execution was interrupted.",
                system=self,
            ) from e
        except PythonScriptError as e:
            logger.error("%s: finalize script failed.", self.name, **logdata(block=self))
            raise e
        except BaseException as e:
            logger.error("%s: finalize script failed.", self.name, **logdata(block=self))
            raise PythonScriptError(system=self) from e

    def post_simulation_finalize(self) -> None:
        """Run ``finalize_script`` and then call the base-class hook."""
        self.exec_finalize()
        return super().post_simulation_finalize()

    def check_types(
        self,
        context: ContextBase,
        error_collector=None,
    ):
        pass

exec_finalize()

Execute the finalize_script using the current persistent environment.

Called once at the end of the simulation via :meth:post_simulation_finalize. The script runs in the same environment that was maintained throughout the simulation, so all variables declared in init_script (and updated by user_statements) are available.

Has no effect if finalize_script is empty.

Source code in jaxonomy/library/custom.py

def exec_finalize(self) -> None:
    """Execute the finalize_script using the current persistent environment.

    Called once at the end of the simulation via :meth:`post_simulation_finalize`.
    The script runs in the same environment that was maintained throughout the
    simulation, so all variables declared in ``init_script`` (and updated by
    ``user_statements``) are available.

    Has no effect if ``finalize_script`` is empty.
    """
    if not self._has_finalize_script:
        return

    local_env = {
        **self.local_env_base,
        **self._persistent_env,
    }

    exec_fn = functools.partial(
        CustomPythonBlock.__exec_fn,
        code=self.finalize_code,
        env=local_env,
        logger_=logger,
        system=self,
        code_name="finalize",
    )

    try:
        io_callback(exec_fn, None)
    except KeyboardInterrupt as e:
        logger.error(
            "Python block '%s' finalize script execution was interrupted.",
            self.name,
            **logdata(block=self),
        )
        raise PythonScriptError(
            message="Python block finalize script execution was interrupted.",
            system=self,
        ) from e
    except PythonScriptError as e:
        logger.error("%s: finalize script failed.", self.name, **logdata(block=self))
        raise e
    except BaseException as e:
        logger.error("%s: finalize script failed.", self.name, **logdata(block=self))
        raise PythonScriptError(system=self) from e

post_simulation_finalize()

Run finalize_script and then call the base-class hook.

Source code in jaxonomy/library/custom.py

def post_simulation_finalize(self) -> None:
    """Run ``finalize_script`` and then call the base-class hook."""
    self.exec_finalize()
    return super().post_simulation_finalize()

DMDForecaster

Bases: LeafSystem

Discrete-time predictor for a fitted (reduced) linear operator.

Propagates x[k+1] = A x[k] (+ B u[k]) and outputs y[k] = C x[k] (C defaults to the identity, so the state itself is the output). This is the online counterpart of :func:dmd / :func:dmdc: fit A (and B) from snapshots offline, then drop the operator into a Jaxonomy diagram as a jax-traceable discrete block that runs inside :func:jaxonomy.simulate.

An input port (and use of B) is created only when B is provided.

Input ports

(0) u[k]: control input, present iff B is given.

Output ports

(0) y[k] = C x[k].

Parameters:

Name	Type	Description	Default
A		State operator (n, n) — a dynamic parameter.	required
B		Optional input operator (n, m) — a dynamic parameter when given.	None
C		Optional output operator (p, n) — a dynamic parameter when given; defaults to identity.	None
dt		Sampling period of the discrete update.	1.0
initial_state		Initial state x[0] of size n (default: zeros).	None

Source code in jaxonomy/library/rom/dmd.py

class DMDForecaster(LeafSystem):
    """Discrete-time predictor for a fitted (reduced) linear operator.

    Propagates ``x[k+1] = A x[k] (+ B u[k])`` and outputs ``y[k] = C x[k]``
    (``C`` defaults to the identity, so the state itself is the output). This is
    the online counterpart of :func:`dmd` / :func:`dmdc`: fit ``A`` (and ``B``)
    from snapshots offline, then drop the operator into a Jaxonomy diagram as a
    jax-traceable discrete block that runs inside :func:`jaxonomy.simulate`.

    An input port (and use of ``B``) is created only when ``B`` is provided.

    Input ports:
        (0) u[k]: control input, present iff ``B`` is given.

    Output ports:
        (0) y[k] = C x[k].

    Parameters:
        A: State operator ``(n, n)`` — a ``dynamic`` parameter.
        B: Optional input operator ``(n, m)`` — a ``dynamic`` parameter when given.
        C: Optional output operator ``(p, n)`` — a ``dynamic`` parameter when
            given; defaults to identity.
        dt: Sampling period of the discrete update.
        initial_state: Initial state ``x[0]`` of size ``n`` (default: zeros).
    """

    @parameters(dynamic=["A", "B", "C"], static=["dt", "initial_state"])
    def __init__(self, A, B=None, C=None, dt=1.0, initial_state=None, name=None,
                 **kwargs):
        super().__init__(name=name, **kwargs)

        A = np.asarray(A, dtype=float)
        if A.ndim == 0:
            A = A.reshape(1, 1)
        self.n = A.shape[0]
        self.has_input = B is not None
        if C is not None:
            C = np.asarray(C, dtype=float)
            if C.ndim == 1:
                C = C.reshape(1, -1)
            self.p = C.shape[0]
        else:
            self.p = self.n

        self.dt = dt
        if initial_state is None:
            initial_state = np.zeros(self.n)
        self._x0 = np.asarray(initial_state, dtype=float).reshape(-1)

        if self.has_input:
            self.declare_input_port(name="u")

        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port(name="out_0")

    def initialize(self, A, B=None, C=None, dt=1.0, initial_state=None, **kwargs):
        if initial_state is None:
            x0 = npa.array(self._x0)
        else:
            x0 = npa.reshape(npa.array(initial_state, dtype=npa.float64), (-1,))

        self.declare_discrete_state(default_value=x0)
        self.configure_periodic_update(
            self._periodic_update_idx, self._update, period=self.dt, offset=0.0
        )
        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=0.0,
            default_value=npa.zeros(self.p) if self.p > 1 else 0.0,
            requires_inputs=False,
            prerequisites_of_calc=[DependencyTicket.xd],
        )

    def _update(self, _time, state, *inputs, **params):
        x = state.discrete_state
        x_next = params["A"] @ x
        if self.has_input:
            u = jnp.atleast_1d(inputs[0])
            x_next = x_next + params["B"] @ u
        return x_next

    def _output(self, _time, state, *_inputs, **params):
        x = state.discrete_state
        C = params.get("C")
        y = x if C is None else C @ x
        if self.p == 1:
            y = jnp.atleast_1d(y)[0]
        return y

DMDResult dataclass

Exact-DMD spectral decomposition (Tu et al. 2014).

Attributes:

Name	Type	Description
modes	Any	DMD modes Φ (columns), shape (n, r), generally complex.
eigenvalues	Any	Discrete-time DMD eigenvalues λ, shape (r,). The growth/decay and oscillation of the identified linear dynamics; a mode is stable iff |λ| < 1.
amplitudes	Any	Mode amplitudes b fitting the first snapshot, shape (r,).
A_tilde	Any	Reduced r×r operator in the POD-projected coordinates.

Source code in jaxonomy/library/rom/dmd.py

@dataclass
class DMDResult:
    """Exact-DMD spectral decomposition (Tu et al. 2014).

    Attributes:
        modes: DMD modes ``Φ`` (columns), shape ``(n, r)``, generally complex.
        eigenvalues: Discrete-time DMD eigenvalues ``λ``, shape ``(r,)``. The
            growth/decay and oscillation of the identified linear dynamics; a
            mode is stable iff ``|λ| < 1``.
        amplitudes: Mode amplitudes ``b`` fitting the first snapshot, shape ``(r,)``.
        A_tilde: Reduced ``r×r`` operator in the POD-projected coordinates.
    """

    modes: Any
    eigenvalues: Any
    amplitudes: Any
    A_tilde: Any

DMDcResult dataclass

DMD-with-control operators (Proctor, Brunton & Kutz 2016).

Attributes:

Name	Type	Description
A	Any	Full n×n state operator.
B	Any	Full n×m input operator.
A_tilde	Any	Reduced r×r state operator (POD-projected).
B_tilde	Any	Reduced r×m input operator.
basis	Any	POD basis Û (columns), shape (n, r), mapping reduced ↔ full.
eigenvalues	Any	Eigenvalues of A_tilde, shape (r,).

Source code in jaxonomy/library/rom/dmd.py

@dataclass
class DMDcResult:
    """DMD-with-control operators (Proctor, Brunton & Kutz 2016).

    Attributes:
        A: Full ``n×n`` state operator.
        B: Full ``n×m`` input operator.
        A_tilde: Reduced ``r×r`` state operator (POD-projected).
        B_tilde: Reduced ``r×m`` input operator.
        basis: POD basis ``Û`` (columns), shape ``(n, r)``, mapping reduced ↔ full.
        eigenvalues: Eigenvalues of ``A_tilde``, shape ``(r,)``.
    """

    A: Any
    B: Any
    A_tilde: Any
    B_tilde: Any
    basis: Any
    eigenvalues: Any

DataSource

Bases: SourceBlock

Produces outputs from an imported data file (.csv, .npy, .npz).

CSV files are read with pandas when installed; otherwise NumPy is used.

Parameters:

Name	Type	Description	Default
file_name	str	Path to .csv, .npy, or .npz.	required
column	Optional[str]	Optional. When set, selects the signal column(s) by name or index string and overrides data_columns for CSV loading. When None, data_columns is used (default index "1" is the second column, i.e. first column is often time at index 0).	None
time_column	str	For CSV with time_samples_as_column=True, column name (e.g. "t") or index string (e.g. "0"). If the name is missing but the file has a header row, the first column is used as time.	'0'
data_columns	str	Column index, name, slice (e.g. 3:8), or list string for CSV.	'1'

Source code in jaxonomy/library/data_source.py

class DataSource(SourceBlock):
    """Produces outputs from an imported data file (.csv, .npy, .npz).

    CSV files are read with pandas when installed; otherwise NumPy is used.

    Parameters:
        file_name: Path to ``.csv``, ``.npy``, or ``.npz``.
        column: Optional. When set, selects the signal column(s) by name or index
            string and overrides ``data_columns`` for CSV loading. When ``None``,
            ``data_columns`` is used (default index ``"1"`` is the second column,
            i.e. first column is often time at index ``0``).
        time_column: For CSV with ``time_samples_as_column=True``, column name (e.g.
            ``"t"``) or index string (e.g. ``"0"``). If the name is missing but the
            file has a header row, the first column is used as time.
        data_columns: Column index, name, slice (e.g. ``3:8``), or list string for CSV.
        See module docstring for ``.npy`` / ``.npz`` layout.
    """

    @parameters(
        static=[
            "file_name",
            "data_columns",
            "column",
            "extrapolation",
            "header_as_first_row",
            "interpolation",
            "sampling_interval",
            "time_column",
            "time_samples_as_column",
        ]
    )
    def __init__(
        self,
        file_name: str,
        data_columns: str = "1",
        column: Optional[str] = None,
        extrapolation: str = "hold",
        header_as_first_row: bool = False,
        interpolation: str = "zero_order_hold",
        sampling_interval: float = 1.0,
        time_column: str = "0",
        time_samples_as_column: bool = False,
        **kwargs,
    ):
        kwargs.pop("data_integration_id", None)

        super().__init__(self._callback, **kwargs)

        effective_columns = str(column) if column is not None else str(data_columns)

        times, data = load_data_source_file(
            str(file_name),
            effective_columns,
            bool(header_as_first_row),
            float(sampling_interval),
            str(time_column),
            bool(time_samples_as_column),
        )

        times = npa.array(times)
        data = npa.array(data)

        if data.size == 0:
            raise ValueError(
                f"DataSource {self.name_path_strme} could not get the requested data columns."
            )

        max_i_zoh = len(times) - 1
        max_i_interp = max(len(times) - 2, 0)
        output_dim = data.shape[1]
        self._scalar_output = output_dim == 1

        def get_below_row_idx(time, max_i):
            time_clipped = npa.clip(time, times[0], times[-1])
            index = npa.searchsorted(times[: max_i + 1], time_clipped, side="right")
            return index - 1, time_clipped

        def _func_zoh(time):
            i, _ = get_below_row_idx(time, max_i_zoh)
            if extrapolation != "zero":
                return data[i, :]
            return npa.where(time > times[-1], npa.zeros(output_dim), data[i, :])

        def _func_interp(time):
            if len(times) < 2:
                return data[0, :]
            i, time_clipped = get_below_row_idx(time, max_i_interp)
            ap1 = data[i, :]
            ap2 = data[i + 1, :]
            if extrapolation != "zero":
                return (ap2 - ap1) / (times[i + 1] - times[i]) * (
                    time_clipped - times[i]
                ) + ap1

            return npa.where(
                time > times[-1],
                npa.zeros(output_dim),
                (ap2 - ap1) / (times[i + 1] - times[i]) * (time_clipped - times[i])
                + ap1,
            )

        def _wrap_func(_func):
            def _ds_wrapped_func(time):
                output = _func(time)
                return output[0]

            return _ds_wrapped_func

        if interpolation == "zero_order_hold":
            _func = _func_zoh
        else:
            _func = _func_interp

        if self._scalar_output:
            _func = _wrap_func(_func)

        self._func = npa.jit(_func)

    def _callback(self, time):
        return self._func(time)

DeadZone

Bases: FeedthroughBlock

Generates zero output within a specified range.

Applies the following function:

         [ input,       input < -half_range
output = | 0,           -half_range <= input <= half_range
         [ input        input > half_range

Parameters:

Name	Type	Description	Default
half_range		The range of the dead zone. Must be > 0.	1.0
mode		"hard" (default, byte-equivalent to legacy behavior) or "smooth". Smooth mode replaces the discontinuous gate with a sigmoid-blended kernel so gradients flow through the dead zone region; the output then has no discontinuity at |x| = half_range and no zero-crossing events are declared.	'hard'
sharpness		Positive scalar; default 10.0. Only used in smooth mode. Larger values give a tighter approximation to the hard dead zone (with smaller gradients inside the band).	10.0
output_shifted		False (default, byte-equivalent to legacy behavior) or True. When True, the hard-mode output outside the band is shifted by half_range * sign(input) so the output is continuous across the band boundary (slope 1 outside, value 0 at the boundary). The False form keeps the legacy "Coulomb friction" semantics where the output jumps at the band boundary. The smooth mode already produces a continuous output and is unaffected by this flag.	False

Input ports

(0) The input signal.

Output ports

(0) The input signal modified by the dead zone.

Events

An event is triggered when the signal enters or exits the dead zone in either direction (hard mode only).

T-115-followup-deadzone-backlash

The mode kwarg unifies a smooth (differentiable) variant. mode="hard" (default) is byte-equivalent to the legacy behavior, including zero-crossing event declaration. mode="smooth" dispatches to a sigmoid-blended formula (see :func:soft_dead_zone) and does not declare zero-crossing events.

T-115-followup-deadzone-bilinear

The output_shifted kwarg toggles between the legacy Coulomb-style hard dead-zone (default, output jumps at the band boundary) and the shifted-output variant (continuous across the band boundary). Default False keeps the block byte-equivalent to phase 1.

Source code in jaxonomy/library/nonlinearities.py

class DeadZone(FeedthroughBlock):
    """Generates zero output within a specified range.

    Applies the following function:
    ```
             [ input,       input < -half_range
    output = | 0,           -half_range <= input <= half_range
             [ input        input > half_range
    ```

    Parameters:
        half_range: The range of the dead zone.  Must be > 0.
        mode: ``"hard"`` (default, byte-equivalent to legacy behavior) or
            ``"smooth"``. Smooth mode replaces the discontinuous gate with
            a sigmoid-blended kernel so gradients flow through the dead
            zone region; the output then has no discontinuity at
            ``|x| = half_range`` and no zero-crossing events are declared.
        sharpness: Positive scalar; default ``10.0``. Only used in smooth
            mode. Larger values give a tighter approximation to the hard
            dead zone (with smaller gradients inside the band).
        output_shifted: ``False`` (default, byte-equivalent to legacy
            behavior) or ``True``. When ``True``, the hard-mode output
            outside the band is shifted by ``half_range * sign(input)`` so
            the output is continuous across the band boundary (slope 1
            outside, value ``0`` at the boundary). The ``False`` form keeps
            the legacy "Coulomb friction" semantics where the output jumps
            at the band boundary. The ``smooth`` mode already produces a
            continuous output and is unaffected by this flag.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The input signal modified by the dead zone.

    Events:
        An event is triggered when the signal enters or exits the dead zone
        in either direction (hard mode only).

    T-115-followup-deadzone-backlash:
        The ``mode`` kwarg unifies a smooth (differentiable) variant.
        ``mode="hard"`` (default) is byte-equivalent to the legacy
        behavior, including zero-crossing event declaration.
        ``mode="smooth"`` dispatches to a sigmoid-blended formula
        (see :func:`soft_dead_zone`) and does *not* declare zero-crossing
        events.

    T-115-followup-deadzone-bilinear:
        The ``output_shifted`` kwarg toggles between the legacy
        Coulomb-style hard dead-zone (default, output jumps at the band
        boundary) and the shifted-output variant
        (continuous across the band boundary). Default ``False`` keeps
        the block byte-equivalent to phase 1.
    """

    @parameters(dynamic=["half_range", "sharpness"], static=["mode", "output_shifted"])
    def __init__(
        self,
        half_range=1.0,
        mode="hard",
        sharpness=10.0,
        output_shifted=False,
        **kwargs,
    ):
        if mode not in ("hard", "smooth"):
            raise BlockParameterError(
                message=(
                    f"DeadZone block: mode must be 'hard' or 'smooth', "
                    f"got {mode!r}."
                ),
                parameter_name="mode",
            )
        super().__init__(self._dead_zone, **kwargs)
        if half_range <= 0:
            raise BlockParameterError(
                message=f"DeadZone block {self.name} has invalid half_range {half_range}. Must be > 0.",
                system=self,
                parameter_name="half_range",
            )
        if mode == "smooth" and sharpness <= 0:
            raise BlockParameterError(
                message=(
                    f"DeadZone block {self.name}: mode='smooth' requires "
                    f"sharpness > 0, got {sharpness}."
                ),
                system=self,
                parameter_name="sharpness",
            )
        if not isinstance(output_shifted, bool):
            raise BlockParameterError(
                message=(
                    f"DeadZone block {self.name}: output_shifted must be a "
                    f"bool, got {output_shifted!r}."
                ),
                system=self,
                parameter_name="output_shifted",
            )
        self.mode = mode
        self.output_shifted = output_shifted

    def initialize(
        self, half_range, mode="hard", sharpness=10.0, output_shifted=False
    ):
        if mode != self.mode:
            raise ValueError(
                "DeadZone: mode cannot be changed after initialization"
            )
        if output_shifted != self.output_shifted:
            raise ValueError(
                "DeadZone: output_shifted cannot be changed after initialization"
            )

    def _dead_zone(self, x, **params):
        if self.mode == "smooth":
            # T-115-followup-deadzone-backlash: differentiable variant.
            return soft_dead_zone(x, params["half_range"], params["sharpness"])
        hr = params["half_range"]
        if self.output_shifted:
            # T-115-followup-deadzone-bilinear: shifted-output variant.
            # Outside the band the output is ``x - hr*sign(x)``;
            # this yields slope 1 with value 0 at ``|x| = hr`` so the
            # output is continuous across the band boundary.
            return npa.where(abs(x) < hr, x * 0, x - hr * npa.sign(x))
        # Legacy Coulomb-style: output jumps at the band boundary.
        return npa.where(abs(x) < hr, x * 0, x)

    def _lower_limit_event_value(self, _time, _state, *inputs, **params):
        (u,) = inputs
        return u + params["half_range"]

    def _upper_limit_event_value(self, _time, _state, *inputs, **params):
        (u,) = inputs
        return u - params["half_range"]

    def initialize_static_data(self, context):
        # Add zero-crossing events so ODE solvers can't try to integrate
        # through a discontinuity.
        #
        # T-115-followup-deadzone-backlash: smooth mode has no
        # discontinuity, so we never declare zero-crossing events for it.
        if (
            self.mode == "hard"
            and not self.has_zero_crossing_events
            and (self.output_ports[0])
        ):
            self.declare_zero_crossing(
                self._lower_limit_event_value, direction="crosses_zero"
            )
            self.declare_zero_crossing(
                self._upper_limit_event_value, direction="crosses_zero"
            )

        return super().initialize_static_data(context)

Decimator

Bases: LeafSystem

Fast-to-slow rate transition: subsample-and-hold at output_dt.

Implements a discrete-time decimator that samples its input on a periodic clock at output_dt (the slow rate) and holds the value until the next slow tick. The input is assumed to be running at input_dt (the fast rate); the block is agnostic to the actual upstream sampling, but the rate-mismatch detector uses this declared pair to recognise the bridge.

Difference equation, with input u and output y::

x[k+1] = u[k * (output_dt / input_dt)]
y(t)   = x[k],   t in (t_k, t_k + output_dt)

Equivalent to a "Rate Transition (fast to slow)" block in its default ZOH-with-decimation mode.

Input ports

(0) The fast-rate input signal.

Output ports

(0) The slow-rate held signal.

Parameters:

Name	Type	Description	Default
input_dt		Sample period of the upstream (fast) source. Used for documentation / the rate-mismatch detector; the block does not actually read the upstream clock.	required
output_dt		Sample period of this block's output (slow rate). Must satisfy output_dt > input_dt for "fast to slow" semantics; the constructor warns if not.	required
initial_state		Initial output value held until the first slow tick fires. Default 0.0.	0.0
mode		How to combine the input samples within each output_dt window before emitting at the slow tick. One of: "pick_last" (default) — emit the most recent input sample at the slow tick (the standard "Rate Transition fast → slow" default). Byte-equivalent to T-123 phase 1. "mean" — emit the arithmetic mean of every input sample observed during the window. Standard anti-aliasing decimation for continuous signals; differentiable through the input (linear). "peak" — emit the input sample with the largest absolute value over the window. Preserves peak excursions for envelope tracking / detection. The selector itself is non-differentiable but gradients flow through the selected sample's value (a np.where branch picks the max-|u| candidate).	'pick_last'

Notes (mode="mean" / mode="peak"): The block declares a second periodic update at input_dt that accumulates samples into a running buffer. At each slow tick the emit-and-reset update fires first (declared before the accumulator in __init__), reads the buffer, computes the window result, and zeroes the buffer for the next window. At simultaneous slow+fast ticks the emit therefore sees the full window from the previous interval; the fast tick at the same t then starts the next window with the current input as its first sample.

Source code in jaxonomy/library/dynamics.py

class Decimator(LeafSystem):
    """Fast-to-slow rate transition: subsample-and-hold at ``output_dt``.

    Implements a discrete-time decimator that samples its input on a
    periodic clock at ``output_dt`` (the slow rate) and holds the value
    until the next slow tick.  The input is assumed to be running at
    ``input_dt`` (the fast rate); the block is agnostic to the actual
    upstream sampling, but the rate-mismatch detector uses this declared
    pair to recognise the bridge.

    Difference equation, with input ``u`` and output ``y``::

        x[k+1] = u[k * (output_dt / input_dt)]
        y(t)   = x[k],   t in (t_k, t_k + output_dt)

    Equivalent to a "Rate Transition (fast to slow)" block in
    its default ZOH-with-decimation mode.

    Input ports:
        (0) The fast-rate input signal.

    Output ports:
        (0) The slow-rate held signal.

    Parameters:
        input_dt: Sample period of the upstream (fast) source. Used for
            documentation / the rate-mismatch detector; the block does
            not actually read the upstream clock.
        output_dt: Sample period of this block's output (slow rate).
            Must satisfy ``output_dt > input_dt`` for "fast to slow"
            semantics; the constructor warns if not.
        initial_state: Initial output value held until the first slow
            tick fires.  Default 0.0.
        mode: How to combine the input samples within each ``output_dt``
            window before emitting at the slow tick.  One of:

            * ``"pick_last"`` (default) — emit the most recent input
              sample at the slow tick (the standard "Rate Transition fast →
              slow" default).  Byte-equivalent to T-123 phase 1.
            * ``"mean"`` — emit the arithmetic mean of every input
              sample observed during the window.  Standard
              anti-aliasing decimation for continuous signals;
              differentiable through the input (linear).
            * ``"peak"`` — emit the input sample with the largest
              absolute value over the window.  Preserves peak
              excursions for envelope tracking / detection.  The
              selector itself is non-differentiable but gradients flow
              through the selected sample's value (a ``np.where``
              branch picks the max-|u| candidate).

    Notes (``mode="mean"`` / ``mode="peak"``):
        The block declares a second periodic update at ``input_dt`` that
        accumulates samples into a running buffer.  At each slow tick
        the emit-and-reset update fires first (declared before the
        accumulator in ``__init__``), reads the buffer, computes the
        window result, and zeroes the buffer for the next window.  At
        simultaneous slow+fast ticks the emit therefore sees the full
        window from the previous interval; the fast tick at the same
        ``t`` then starts the next window with the current input as its
        first sample.
    """

    # T-123: marker attribute used by
    # ``jaxonomy.simulation.rate_groups.detect_rate_mismatches`` to
    # recognise this block as a rate-bridge and skip the mismatch check
    # on adjacent connections.
    _jaxonomy_rate_transition = True

    @parameters(static=["input_dt", "output_dt", "mode"], dynamic=["initial_state"])
    def __init__(
        self,
        input_dt,
        output_dt,
        initial_state=0.0,
        *args,
        mode="pick_last",
        dtype=None,
        **kwargs,
    ):
        # Mirror the per-block dtype / precision-policy plumbing used by
        # ``UnitDelay`` and ``ZeroOrderHold``.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        if mode not in _DECIMATOR_MODES:
            raise ValueError(
                f"Decimator: mode={mode!r} is not one of "
                f"{_DECIMATOR_MODES!r}."
            )
        self._mode = mode
        super().__init__(*args, **kwargs)
        self.input_dt = input_dt
        self.output_dt = output_dt

        if not (output_dt > input_dt):
            warnings.warn(
                f"Decimator block '{self.name}' got output_dt={output_dt} "
                f"<= input_dt={input_dt}; expected output_dt > input_dt "
                f"for fast-to-slow rate transitions. "
                f"Consider RateTransition(input_dt, output_dt) which "
                f"auto-picks the right block.",
                UserWarning,
                stacklevel=3,
            )

        self.declare_input_port()
        # Declaration order matters for the two-phase event scheduler:
        # at simultaneous slow+fast ticks the events fire sequentially
        # in declaration order against the *accumulated* context, so
        # registering the slow emit-and-reset first lets it see the
        # accumulated window before the fast tick begins the next one.
        self._periodic_update_idx = self.declare_periodic_update()
        if self._mode != "pick_last":
            self._fast_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    # ------------------------------------------------------------------
    # initialize / reset for pick_last (legacy single-state path) and
    # for the windowed modes (NamedTuple state + dual periodic update).
    # ------------------------------------------------------------------

    def initialize(self, initial_state, input_dt=None, output_dt=None, mode=None):
        if self._dtype is not None:
            initial_state = npa.asarray(initial_state).astype(self._dtype)
        if self._mode == "pick_last":
            # Legacy phase-1 path: single periodic update at output_dt,
            # scalar discrete state.  Byte-equivalent to the pre-followup
            # block.
            self.configure_periodic_update(
                self._periodic_update_idx,
                self._update_pick_last,
                period=self.output_dt,
                offset=self.output_dt,
            )
            self.configure_output_port(
                self._output_port_idx,
                self._output_pick_last,
                period=self.output_dt,
                offset=0.0,
                requires_inputs=False,
                prerequisites_of_calc=[DependencyTicket.xd],
                default_value=initial_state,
            )
            return

        # Windowed modes need ``initial_state.dtype`` below to keep the
        # NamedTuple state's accumulator / count / peak fields all on
        # the same dtype.  Outside any ``precision_policy`` context the
        # pick_last branch above leaves ``initial_state`` as the user's
        # raw value (often a Python float); normalise it here so the
        # ``.dtype`` access is safe under T-005 default-float64.
        initial_state = npa.asarray(initial_state)
        # Windowed modes: two periodic updates + NamedTuple state.
        if self._mode == "mean":
            init_state = _DecimatorMeanState(
                output=initial_state,
                accumulator=npa.zeros_like(initial_state),
                count=npa.asarray(0.0, dtype=initial_state.dtype),
            )
            slow_cb, fast_cb = (
                self._update_mean_emit,
                self._update_mean_accumulate,
            )
        else:  # "peak"
            init_state = _DecimatorPeakState(
                output=initial_state,
                peak_abs=npa.full_like(
                    initial_state, npa.asarray(-npa.inf, dtype=initial_state.dtype)
                ),
                peak_value=npa.zeros_like(initial_state),
            )
            slow_cb, fast_cb = (
                self._update_peak_emit,
                self._update_peak_accumulate,
            )
        self.configure_periodic_update(
            self._periodic_update_idx,
            slow_cb,
            period=self.output_dt,
            offset=self.output_dt,
        )
        self.configure_periodic_update(
            self._fast_update_idx,
            fast_cb,
            period=self.input_dt,
            offset=0.0,
        )
        self.configure_output_port(
            self._output_port_idx,
            self._output_windowed,
            period=self.output_dt,
            offset=0.0,
            requires_inputs=False,
            prerequisites_of_calc=[DependencyTicket.xd],
            default_value=initial_state,
        )

    def reset_default_values(
        self, initial_state, input_dt=None, output_dt=None, mode=None,
    ):
        if self._dtype is not None:
            initial_state = npa.asarray(initial_state).astype(self._dtype)
        if self._mode == "pick_last":
            self.declare_discrete_state(default_value=initial_state)
            self.configure_output_port_default_value(
                self._output_port_idx, initial_state
            )
            return

        initial_state = npa.asarray(initial_state)
        if self._mode == "mean":
            new_state = _DecimatorMeanState(
                output=initial_state,
                accumulator=npa.zeros_like(initial_state),
                count=npa.asarray(0.0, dtype=initial_state.dtype),
            )
        else:  # "peak"
            new_state = _DecimatorPeakState(
                output=initial_state,
                peak_abs=npa.full_like(
                    initial_state, npa.asarray(-npa.inf, dtype=initial_state.dtype)
                ),
                peak_value=npa.zeros_like(initial_state),
            )
        self.declare_discrete_state(default_value=new_state, as_array=False)
        self.configure_output_port_default_value(
            self._output_port_idx, initial_state
        )

    # ------------------------------------------------------------------
    # pick_last callbacks (T-123 phase 1, unchanged).
    # ------------------------------------------------------------------

    def _update_pick_last(self, _time, _state, u, **_params):
        # Subsample: at every slow tick, latch the current fast input.
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        return u

    def _output_pick_last(self, _time, state, **_parameters):
        return state.discrete_state

    # Legacy aliases — preserved for any external caller that imported
    # the original ``Decimator._update`` / ``Decimator._output`` names
    # (e.g. the phase-1 differentiability test that calls them directly).
    _update = _update_pick_last
    _output = _output_pick_last

    # ------------------------------------------------------------------
    # mean-mode callbacks.
    # ------------------------------------------------------------------

    def _update_mean_emit(self, _time, state, _u, **_params):
        # Slow tick: read the accumulated window, emit its mean, reset.
        xd = state.discrete_state
        # Guard against the degenerate ``count == 0`` case — the
        # ``npa.where`` keeps gradients finite.  In practice the slow
        # tick has offset=output_dt, so count is always ratio>=1 when
        # the slow update fires.
        safe_count = npa.where(
            xd.count > 0, xd.count, npa.asarray(1.0, dtype=xd.count.dtype)
        )
        mean = xd.accumulator / safe_count
        return _DecimatorMeanState(
            output=mean,
            accumulator=npa.zeros_like(xd.accumulator),
            count=npa.asarray(0.0, dtype=xd.count.dtype),
        )

    def _update_mean_accumulate(self, _time, state, u, **_params):
        # Fast tick: add the current input sample to the running sum.
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        xd = state.discrete_state
        return _DecimatorMeanState(
            output=xd.output,
            accumulator=xd.accumulator + u,
            count=xd.count + npa.asarray(1.0, dtype=xd.count.dtype),
        )

    # ------------------------------------------------------------------
    # peak-mode callbacks (max-absolute-value sample within window).
    # ------------------------------------------------------------------

    def _update_peak_emit(self, _time, state, _u, **_params):
        xd = state.discrete_state
        neg_inf = npa.full_like(
            xd.peak_abs, npa.asarray(-npa.inf, dtype=xd.peak_abs.dtype)
        )
        return _DecimatorPeakState(
            output=xd.peak_value,
            peak_abs=neg_inf,
            peak_value=npa.zeros_like(xd.peak_value),
        )

    def _update_peak_accumulate(self, _time, state, u, **_params):
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        xd = state.discrete_state
        abs_u = npa.abs(u)
        is_new_peak = abs_u > xd.peak_abs
        new_peak_abs = npa.where(is_new_peak, abs_u, xd.peak_abs)
        new_peak_value = npa.where(is_new_peak, u, xd.peak_value)
        return _DecimatorPeakState(
            output=xd.output,
            peak_abs=new_peak_abs,
            peak_value=new_peak_value,
        )

    # ------------------------------------------------------------------
    # Shared output for the windowed modes — both store the held value
    # on ``state.discrete_state.output``.
    # ------------------------------------------------------------------

    def _output_windowed(self, _time, state, **_parameters):
        return state.discrete_state.output

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        inp_data = self.eval_input(context)
        xd = context[self.system_id].discrete_state
        if self._mode != "pick_last":
            # Type-check against the ``output`` field — that's the
            # signal flowing out of the block.  The accumulator/count
            # (or peak_abs/peak_value) fields are private bookkeeping.
            xd = xd.output
        check_state_type(
            self,
            inp_data=inp_data,
            state_data=xd,
            error_collector=error_collector,
        )

Demultiplexer

Bases: LeafSystem

Split a vector signal into its components.

Input ports

(0) The vector signal to split.

Output ports

(0..n_out-1) The components of the input signal.

Source code in jaxonomy/library/routing.py

class Demultiplexer(LeafSystem):
    """Split a vector signal into its components.

    Input ports:
        (0) The vector signal to split.

    Output ports:
        (0..n_out-1) The components of the input signal.
    """

    def __init__(self, n_out, **kwargs):
        super().__init__(**kwargs)

        self.declare_input_port()

        # Need a helper function so that the lambda captures the correct value of i
        # and doesn't use something that ends up fixed in scope.
        def _declare_output(i):
            def _compute_output(_time, _state, *inputs, **_params):
                (input_vec,) = inputs
                return input_vec[i]

            self.declare_output_port(
                _compute_output,
                prerequisites_of_calc=[self.input_ports[0].ticket],
            )

        for i in npa.arange(n_out):
            _declare_output(i)

Demux

Bases: LeafSystem

Unstack a single array input into n_outputs separate signals.

This is the standard Demux block and the inverse of :class:Mux: given a 1-D input [a, b, c] it produces three scalar outputs a, b, c; given a 2-D input of shape (n_outputs, k) it produces n_outputs outputs of shape (k,).

Internally this is index-based slicing along axis 0, which is fully differentiable through every output port (each output picks one slice of the input vector).

Input ports

(0) The vector or array signal to split. Its leading axis must have length n_outputs.

Output ports

(0..n_outputs-1) The components of the input signal.

Source code in jaxonomy/library/routing.py

class Demux(LeafSystem):
    """Unstack a single array input into ``n_outputs`` separate signals.

    This is the standard ``Demux`` block and the inverse of :class:`Mux`:
    given a 1-D input ``[a, b, c]`` it produces three scalar outputs
    ``a``, ``b``, ``c``; given a 2-D input of shape ``(n_outputs, k)``
    it produces ``n_outputs`` outputs of shape ``(k,)``.

    Internally this is index-based slicing along axis 0, which is fully
    differentiable through every output port (each output picks one
    slice of the input vector).

    Input ports:
        (0) The vector or array signal to split. Its leading axis must
            have length ``n_outputs``.

    Output ports:
        (0..n_outputs-1) The components of the input signal.
    """

    def __init__(self, n_outputs, **kwargs):
        super().__init__(**kwargs)

        self.declare_input_port()

        # Helper closure so each output port captures its own ``i``.
        def _declare_output(i):
            def _compute_output(_time, _state, *inputs, **_params):
                (input_vec,) = inputs
                return input_vec[i]

            self.declare_output_port(
                _compute_output,
                prerequisites_of_calc=[self.input_ports[0].ticket],
            )

        for i in range(int(n_outputs)):
            _declare_output(i)

Derivative

Bases: LTISystem

Causal estimate of the derivative of a signal in continuous time.

This is implemented as a state-space system with matrices (A, B, C, D), which are then used to create a (first-order) LTISystem. Note that this only supports single-input, single-output derivative blocks.

The derivative is implemented as a filter with a filter coefficient of N, which is used to construct the following proper transfer function:

    H(s) = Ns / (s + N)

As N -> ∞, the transfer function approaches a pure differentiator. However, this system becomes increasingly stiff and difficult to integrate, so it is recommended to select a value of N based on the time scales of the system.

From the transfer function, scipy.signal.tf2ss is used to convert to state-space form and create an LTISystem.

Input ports

(0) u: Input (scalar)

Output ports

(0) y: Output (scalar), estimating the time derivative du/dt

Source code in jaxonomy/library/linear_system.py

class Derivative(LTISystem):
    """Causal estimate of the derivative of a signal in continuous time.

    This is implemented as a state-space system with matrices (A, B, C, D),
    which are then used to create a (first-order) LTISystem.  Note that this
    only supports single-input, single-output derivative blocks.

    The derivative is implemented as a filter with a filter coefficient of `N`,
    which is used to construct the following proper transfer function:
    ```
        H(s) = Ns / (s + N)
    ```
    As N -> ∞, the transfer function approaches a pure differentiator.  However,
    this system becomes increasingly stiff and difficult to integrate, so it is
    recommended to select a value of N based on the time scales of the system.

    From the transfer function, `scipy.signal.tf2ss` is used to convert to
    state-space form and create an LTISystem.

    Input ports:
        (0) u: Input (scalar)

    Output ports:
        (0) y: Output (scalar), estimating the time derivative du/dt
    """

    # tf2ss is not implemented in jax.scipy.signal so filter_coefficient can't be
    # a dynamic parameter.
    @parameters(static=["filter_coefficient"])
    def __init__(self, filter_coefficient=100, *args, **kwargs):
        N = filter_coefficient
        num = [N, 0]
        den = [1, N]
        A, B, C, D = signal.tf2ss(num, den)
        super().__init__(A, B, C, D, *args, **kwargs)

    def _eval_output(self, time, state, *inputs, **params):
        return self._eval_output_base(self.C, self.D, state, *inputs)

    def ode(self, time, state, u, **params):
        return super().ode(time, state, u, A=self.A, B=self.B)

    def initialize(self, filter_coefficient, **kwargs):
        N = filter_coefficient
        num = [N, 0]
        den = [1, N]

        A, B, C, D = signal.tf2ss(num, den)
        self._init_state(A, B, C, D)

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        inputs = self.collect_inputs(context)
        (u,) = inputs

        if not npa.ndim(u) == 0:
            with ErrorCollector.context(error_collector):
                raise StaticError(
                    message="Derivative must have scalar input.",
                    system=self,
                )

DerivativeDiscrete

Bases: LeafSystem

Discrete approximation to the derivative of the input signal w.r.t. time.'

By default the block uses a simple backward difference approximation:

y[k] = (u[k] - u[k-1]) / dt

However, the block can also be configured to use a recursive filter for a better approximation. In this case the filter coefficients are determined by the filter_type and filter_coefficient parameters. The filter is a pair of two-element arrays a and b and the filter equation is:

a0*y[k] + a1*y[k-1] = b0*u[k] + b1*u[k-1]

Denoting the filter_coefficient parameter by N, the following filters are available: - "none": The default, a simple finite difference approximation. - "forward": A filtered forward Euler discretization. The filter is: a = [1, (N*dt - 1)] and b = [N, -N]. - "backward": A filtered backward Euler discretization. The filter is: a = [(1 + N*dt), -1] and b = [N, -N]. - "bilinear": A filtered bilinear transform discretization. The filter is: a = [(2 + N*dt), (-2 + N*dt)] and b = [2*N, -2*N].

Input ports

(0) The input signal.

Output ports

(0) The approximate derivative of the input signal.

Parameters:

Name	Type	Description	Default
dt		The time step of the discrete approximation.	required
filter_type		One of "none", "forward", "backward", or "bilinear". This determines the type of filter used to approximate the derivative. The default is "none", corresponding to a simple backward difference approximation.	'none'
filter_coefficient		The coefficient in the filter (N in the equations above). This is only used if filter_type is not "none". The default is 1.0.	1.0

Source code in jaxonomy/library/dynamics.py

class DerivativeDiscrete(LeafSystem):
    """Discrete approximation to the derivative of the input signal w.r.t. time.'

    By default the block uses a simple backward difference approximation:
    ```
    y[k] = (u[k] - u[k-1]) / dt
    ```
    However, the block can also be configured to use a recursive filter for a
    better approximation. In this case the filter coefficients are determined
    by the `filter_type` and `filter_coefficient` parameters. The filter is
    a pair of two-element arrays `a` and `b` and the filter equation is:
    ```
    a0*y[k] + a1*y[k-1] = b0*u[k] + b1*u[k-1]
    ```

    Denoting the `filter_coefficient` parameter by `N`, the following filters are
    available:
    - "none": The default, a simple finite difference approximation.
    - "forward": A filtered forward Euler discretization. The filter is:
        `a = [1, (N*dt - 1)]` and `b = [N, -N]`.
    - "backward": A filtered backward Euler discretization. The filter is:
        `a = [(1 + N*dt), -1]` and `b = [N, -N]`.
    - "bilinear": A filtered bilinear transform discretization. The filter is:
        `a = [(2 + N*dt), (-2 + N*dt)]` and `b = [2*N, -2*N]`.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The approximate derivative of the input signal.

    Parameters:
        dt:
            The time step of the discrete approximation.
        filter_type:
            One of "none", "forward", "backward", or "bilinear". This determines the
            type of filter used to approximate the derivative. The default is "none",
            corresponding to a simple backward difference approximation.
        filter_coefficient:
            The coefficient in the filter (`N` in the equations above). This is only
            used if `filter_type` is not "none". The default is 1.0.
    """

    @parameters(static=["dt", "filter_type", "filter_coefficient"])
    def __init__(self, dt, filter_type="none", filter_coefficient=1.0, dtype=None, **kwargs):
        # T-038a-followup-other-blocks: per-block dtype override; stored
        # outside the @parameters list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(**kwargs)
        self.dt = dt
        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self.deriv_output = self.declare_output_port(
            period=dt,
            offset=0.0,
            prerequisites_of_calc=[self.input_ports[0].ticket],
        )

    def initialize(self, filter_type="none", filter_coefficient=1.0, dt=None):
        # Determine the coefficients of the filter, if applicable
        # The filter is a pair of two-element array and the filter
        # equation is:
        # a0*y[k] + a1*y[k-1] = b0*u[k] + b1*u[k-1]
        b, a = derivative_filter(
            N=filter_coefficient, dt=self.dt, filter_type=filter_type
        )
        if self._dtype is not None:
            # T-038a-followup-other-blocks: cast filter coefficients to
            # the per-block dtype so the output arithmetic runs at this
            # precision regardless of upstream/global default.
            b = npa.asarray(b).astype(self._dtype)
            a = npa.asarray(a).astype(self._dtype)
        self.filter = (b, a)

        self.declare_discrete_state(default_value=None, as_array=False)

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=0.0,
        )

        # At t=0 we have no prior information, so the output will
        # be held from its initial value (zero). At t=dt, we have
        # a previous sample, so there is enough information to estimate
        # the derivative.
        self.configure_output_port(
            self.deriv_output,
            self._output,
            period=self.dt,
            offset=self.dt,
            prerequisites_of_calc=[self.input_ports[0].ticket],
        )

    def _output(self, _time, state, *inputs, **_params):
        # Compute the filtered derivative estimate
        (u,) = inputs
        b, a = self.filter
        y_prev = state.cache[self.deriv_output]
        u_prev = state.discrete_state
        y = (b[0] * u + b[1] * u_prev - a[1] * y_prev) / a[0]
        # T-038a-followup-other-blocks: cast the output to the per-block
        # dtype so cross-dtype upstream connections promote down to the
        # requested precision (best-effort; see ``LookupTable1d`` doc).
        if self._dtype is not None:
            y = npa.asarray(y).astype(self._dtype)
        return y

    def _update(self, time, state, u, **params):
        # Every dt seconds, update the state to the current values
        # T-038a-followup-other-blocks: cast u so the saved state lands
        # the same dtype on every step.
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        return u

    def initialize_static_data(self, context):
        """Infer the size and dtype of the internal states"""
        # If building as part of a subsystem, this may not be fully connected yet.
        # That's fine, as long as it is connected by root context creation time.
        # This probably isn't a good long-term solution:
        #   see https://jaxonomy.atlassian.net/browse/WC-51
        try:
            u = self.eval_input(context)
            self._default_discrete_state = u
            local_context = context[self.system_id].with_discrete_state(u)
            self._default_cache[self.deriv_output] = 0 * u
            local_context = local_context.with_cached_value(self.deriv_output, 0 * u)
            context = context.with_subcontext(self.system_id, local_context)

        except UpstreamEvalError:
            logger.debug(
                "DerivativeDiscrete.initialize_static_data: UpstreamEvalError. "
                "Continuing without default value initialization."
            )
        return super().initialize_static_data(context)

initialize_static_data(context)

Infer the size and dtype of the internal states

Source code in jaxonomy/library/dynamics.py

def initialize_static_data(self, context):
    """Infer the size and dtype of the internal states"""
    # If building as part of a subsystem, this may not be fully connected yet.
    # That's fine, as long as it is connected by root context creation time.
    # This probably isn't a good long-term solution:
    #   see https://jaxonomy.atlassian.net/browse/WC-51
    try:
        u = self.eval_input(context)
        self._default_discrete_state = u
        local_context = context[self.system_id].with_discrete_state(u)
        self._default_cache[self.deriv_output] = 0 * u
        local_context = local_context.with_cached_value(self.deriv_output, 0 * u)
        context = context.with_subcontext(self.system_id, local_context)

    except UpstreamEvalError:
        logger.debug(
            "DerivativeDiscrete.initialize_static_data: UpstreamEvalError. "
            "Continuing without default value initialization."
        )
    return super().initialize_static_data(context)

DirectShootingNMPC

Bases: NonlinearMPCIpopt

Implementation of nonlinear MPC with a direct shooting transcription and IPOPT as the NLP solver.

Input ports

(0) x_0 : current state vector. (1) x_ref : reference state trajectory for the nonlinear MPC. (2) u_ref : reference input trajectory for the nonlinear MPC.

Output ports

(1) u_opt : the optimal control input to be applied at the current time step as determined by the nonlinear MPC.

Parameters:

Name	Type	Description	Default
plant		LeafSystem or Diagram The plant to be controlled.	required
Q		Array State weighting matrix in the cost function.	required
QN		Array Terminal state weighting matrix in the cost function.	required
R		Array Control input weighting matrix in the cost function.	required
N		int The prediction horizon, an integer specifying the number of steps to predict. Note: prediction and control horizons are identical for now.	required
nh		int Number of minor steps to take within an RK4 major step.	required
dt		float: Major time step, a scalar indicating the increment in time for each step in the prediction and control horizons.	required
lb_u		Array Lower bound on the control input vector.	None
ub_u		Array Upper bound on the control input vector.	None
u_optvars_0		Array Initial guess for the control vector optimization variables in the NLP.	None

Source code in jaxonomy/library/nmpc/direct_shooting_ipopt_nmpc.py

class DirectShootingNMPC(NonlinearMPCIpopt):
    """
    Implementation of nonlinear MPC with a direct shooting transcription and IPOPT as
    the NLP solver.

    Input ports:
        (0) x_0 : current state vector.
        (1) x_ref : reference state trajectory for the nonlinear MPC.
        (2) u_ref : reference input trajectory for the nonlinear MPC.

    Output ports:
        (1) u_opt : the optimal control input to be applied at the current time step
                    as determined by the nonlinear MPC.

    Parameters:
        plant: LeafSystem or Diagram
            The plant to be controlled.

        Q: Array
            State weighting matrix in the cost function.

        QN: Array
            Terminal state weighting matrix in the cost function.

        R: Array
            Control input weighting matrix in the cost function.

        N: int
            The prediction horizon, an integer specifying the number of steps to
            predict. Note: prediction and control horizons are identical for now.

        nh: int
            Number of minor steps to take within an RK4 major step.

        dt: float:
            Major time step, a scalar indicating the increment in time for each step in
            the prediction and control horizons.

        lb_u: Array
            Lower bound on the control input vector.

        ub_u: Array
            Upper bound on the control input vector.

        u_optvars_0: Array
            Initial guess for the control vector optimization variables in the NLP.
    """

    def __init__(
        self,
        plant,
        Q,
        QN,
        R,
        N,
        nh,
        dt,
        lb_u=None,
        ub_u=None,
        u_optvars_0=None,
        name=None,
        warm_start=False,
    ):
        self.plant = plant

        self.Q = Q
        self.QN = QN
        self.R = R

        self.N = N
        self.nh = nh
        self.dt = dt

        self.lb_u = lb_u
        self.ub_u = ub_u

        self.nx = Q.shape[0]
        self.nu = R.shape[0]

        if lb_u is None:
            self.lb_u = -1e20 * jnp.ones(self.nu)

        if ub_u is None:
            self.ub_u = 1e20 * jnp.ones(self.nu)

        # Currently guesses are not taken into account
        self.u_optvars_0 = u_optvars_0  # Currently does nothing
        if u_optvars_0 is None:
            u_optvars_0 = jnp.zeros((N, self.nu))

        self.ode_rhs = make_ode_rhs(plant, self.nu)

        nlp_structure_ipopt = NMPCProblemStructure(
            self.num_optvars,
            self._objective,
        )

        super().__init__(
            dt,
            self.nu,
            self.num_optvars,
            nlp_structure_ipopt,
            name=name,
            warm_start=warm_start,
        )

    @property
    def num_optvars(self):
        return self.N * self.nu

    @property
    def num_constraints(self):
        return 0

    @property
    def bounds_optvars(self):
        lb = jnp.tile(self.lb_u, self.N)
        ub = jnp.tile(self.ub_u, self.N)
        return (lb, ub)

    @property
    def bounds_constraints(self):
        c_lb = []
        c_ub = []
        return (c_lb, c_ub)

    @partial(jax.jit, static_argnames=("self",))
    def _objective(self, optvars, t0, x0, x_ref, u_ref):
        u_flat = optvars
        u = jnp.array(u_flat.reshape((self.N, self.nu)))

        x = jnp.zeros((self.N + 1, x0.size))
        x = x.at[0].set(x0)

        def _update_function(idx, x):
            t_major_start = t0 + self.dt * idx
            x_current = x[idx]
            u_current = u[idx]
            x_next = rk4_major_step_constant_u(
                t_major_start,
                x_current,
                u_current,
                self.dt,
                self.nh,
                self.ode_rhs,
            )
            return x.at[idx + 1].set(x_next)

        x = jax.lax.fori_loop(0, self.N, _update_function, x)

        xdiff = x - x_ref
        udiff = u - u_ref

        # compute sum of quadratic products for x_0 to x_{N-1}
        A = jnp.dot(xdiff[:-1], self.Q)
        qp_x_sum = jnp.sum(xdiff[:-1] * A, axis=None)

        # Compute quadratic product for the x_N
        xN = xdiff[-1]
        qp_x_N = jnp.dot(xN, jnp.dot(self.QN, xN))

        # compute sum of quadratic products for u_0 to u_{N-1}
        B = jnp.dot(udiff, self.R)
        qp_u_sum = jnp.sum(udiff * B, axis=None)

        # Sum the quadratic products
        total_sum = qp_x_sum + qp_x_N + qp_u_sum
        return total_sum

DirectTranscriptionNMPC

Bases: NonlinearMPCIpopt

Implementation of nonlinear MPC with direct transcription and IPOPT as the NLP solver.

Input ports

(0) x_0 : current state vector. (1) x_ref : reference state trajectory for the nonlinear MPC. (2) u_ref : reference input trajectory for the nonlinear MPC.

Output ports

(1) u_opt : the optimal control input to be applied at the current time step as determined by the nonlinear MPC.

Parameters:

Name	Type	Description	Default
plant		LeafSystem or Diagram The plant to be controlled.	required
Q		Array State weighting matrix in the cost function.	required
QN		Array Terminal state weighting matrix in the cost function.	required
R		Array Control input weighting matrix in the cost function.	required
N		int The prediction horizon, an integer specifying the number of steps to predict. Note: prediction and control horizons are identical for now.	required
nh		int Number of minor steps to take within an RK4 major step.	required
dt		float: Major time step, a scalar indicating the increment in time for each step in the prediction and control horizons.	required
lb_x		Array Lower bound on the state vector.	None
ub_x		Array Upper bound on the state vector.	None
lb_u		Array Lower bound on the control input vector.	None
ub_u		Array Upper bound on the control input vector.	None
x_optvars_0		Array Initial guess for the state vector optimization variables in the NLP.	None
u_optvars_0		Array Initial guess for the control vector optimization variables in the NLP.	None

Source code in jaxonomy/library/nmpc/direct_transcription_ipopt_nmpc.py

class DirectTranscriptionNMPC(NonlinearMPCIpopt):
    """
    Implementation of nonlinear MPC with direct transcription and IPOPT as the NLP
    solver.

    Input ports:
        (0) x_0 : current state vector.
        (1) x_ref : reference state trajectory for the nonlinear MPC.
        (2) u_ref : reference input trajectory for the nonlinear MPC.

    Output ports:
        (1) u_opt : the optimal control input to be applied at the current time step
                    as determined by the nonlinear MPC.

    Parameters:
        plant: LeafSystem or Diagram
            The plant to be controlled.

        Q: Array
            State weighting matrix in the cost function.

        QN: Array
            Terminal state weighting matrix in the cost function.

        R: Array
            Control input weighting matrix in the cost function.

        N: int
            The prediction horizon, an integer specifying the number of steps to
            predict. Note: prediction and control horizons are identical for now.

        nh: int
            Number of minor steps to take within an RK4 major step.

        dt: float:
            Major time step, a scalar indicating the increment in time for each step in
            the prediction and control horizons.

        lb_x: Array
            Lower bound on the state vector.

        ub_x: Array
            Upper bound on the state vector.

        lb_u: Array
            Lower bound on the control input vector.

        ub_u: Array
            Upper bound on the control input vector.

        x_optvars_0: Array
            Initial guess for the state vector optimization variables in the NLP.

        u_optvars_0: Array
            Initial guess for the control vector optimization variables in the NLP.
    """

    def __init__(
        self,
        plant,
        Q,
        QN,
        R,
        N,
        nh,
        dt,
        lb_x=None,
        ub_x=None,
        lb_u=None,
        ub_u=None,
        x_optvars_0=None,
        u_optvars_0=None,
        name=None,
        warm_start=False,
    ):
        self.plant = plant

        self.Q = Q
        self.QN = QN
        self.R = R

        self.N = N
        self.nh = nh
        self.dt = dt

        self.lb_x = lb_x
        self.ub_x = ub_x
        self.lb_u = lb_u
        self.ub_u = ub_u

        self.nx = Q.shape[0]
        self.nu = R.shape[0]

        if lb_x is None:
            self.lb_x = -1e20 * jnp.ones(self.nx)

        if ub_x is None:
            self.ub_x = 1e20 * jnp.ones(self.nx)

        if lb_u is None:
            self.lb_u = -1e20 * jnp.ones(self.nu)

        if ub_u is None:
            self.ub_u = 1e20 * jnp.ones(self.nu)

        # Currently guesses are not taken into account
        self.x_optvars_0 = x_optvars_0  # Currently does nothing
        self.u_optvars_0 = u_optvars_0  # Currently does nothing
        if x_optvars_0 is None:
            x_optvars_0 = jnp.zeros((N + 1, self.nx))
        if u_optvars_0 is None:
            u_optvars_0 = jnp.zeros((N, self.nu))

        self.ode_rhs = make_ode_rhs(plant, self.nu)

        nlp_structure_ipopt = NMPCProblemStructure(
            self.num_optvars,
            self._objective,
            self._constraints,
        )

        super().__init__(
            dt,
            self.nu,
            self.num_optvars,
            nlp_structure_ipopt,
            name=name,
            warm_start=warm_start,
        )

    @property
    def num_optvars(self):
        return (self.N + 1) * self.nx + self.N * self.nu

    @property
    def num_constraints(self):
        return (self.N + 1) * self.nx

    @property
    def bounds_optvars(self):
        lb = jnp.hstack([jnp.tile(self.lb_u, self.N), jnp.tile(self.lb_x, self.N + 1)])
        ub = jnp.hstack([jnp.tile(self.ub_u, self.N), jnp.tile(self.ub_x, self.N + 1)])
        return (lb, ub)

    @property
    def bounds_constraints(self):
        c_lb = jnp.zeros(self.num_constraints)
        c_ub = jnp.zeros(self.num_constraints)
        return (c_lb, c_ub)

    @partial(jax.jit, static_argnames=("self",))
    def _objective(self, optvars, t0, x0, x_ref, u_ref):
        u_and_x_flat = optvars

        u = u_and_x_flat[: self.nu * self.N].reshape((self.N, self.nu))
        x = u_and_x_flat[self.nu * self.N :].reshape((self.N + 1, self.nx))

        xdiff = x - x_ref
        udiff = u - u_ref

        # compute sum of quadratic products for x_0 to x_{N-1}
        A = jnp.dot(xdiff[:-1], self.Q)
        qp_x_sum = jnp.sum(xdiff[:-1] * A, axis=None)

        # Compute quadratic product for the x_N
        xN = xdiff[-1]
        qp_x_N = jnp.dot(xN, jnp.dot(self.QN, xN))

        # compute sum of quadratic products for u_0 to u_{N-1}
        B = jnp.dot(udiff, self.R)
        qp_u_sum = jnp.sum(udiff * B, axis=None)

        # Sum the quadratic products
        total_sum = qp_x_sum + qp_x_N + qp_u_sum
        return total_sum

    @partial(jax.jit, static_argnames=("self",))
    def _constraints(self, optvars, t0, x0, x_ref, u_ref):
        u_and_x_flat = optvars
        u = u_and_x_flat[: self.nu * self.N].reshape((self.N, self.nu))
        x = u_and_x_flat[self.nu * self.N :].reshape((self.N + 1, self.nx))

        x_sim = jnp.zeros((self.N, x0.size))

        def _update_function(idx, x_sim_l):
            t_major_start = t0 + self.dt * idx
            x_current = x[idx]
            u_current = u[idx]
            x_next = rk4_major_step_constant_u(
                t_major_start,
                x_current,
                u_current,
                self.dt,
                self.nh,
                self.ode_rhs,
            )
            return x_sim_l.at[idx].set(x_next)

        x_sim = jax.lax.fori_loop(0, self.N, _update_function, x_sim)  # x1, x2, ..., xN

        c0 = x0 - x[0]
        c_others = x[1:] - x_sim

        c_all = jnp.hstack([c0.ravel(), c_others.ravel()])

        return c_all

DiscreteClock

Bases: LeafSystem

Source block that produces the time sampled at a fixed rate.

The block maintains the most recently sampled time as a discrete state, provided to the output port during the following interval. Graphically, a discrete clock sampled at 100 Hz would have the following time series:

  x(t)                  ●━
    |                   ┆
.03 |              ●━━━━○
    |              ┆
.02 |         ●━━━━○
    |         ┆
.01 |    ●━━━━○
    |    ┆
  0 ●━━━━○----+----+----+-- t
    0   .01  .02  .03  .04

The recorded states are the closed circles, which should be interpreted at index n as the value seen by all other blocks on the interval (t[n], t[n+1]).

Input ports

None

Output ports

(0) The sampled time.

Parameters:

Name	Type	Description	Default
dt		The sampling period of the clock.	required
start_time		The simulation time at which the clock starts. Defaults to 0.	0

Source code in jaxonomy/library/sources.py

class DiscreteClock(LeafSystem):
    """Source block that produces the time sampled at a fixed rate.

    The block maintains the most recently sampled time as a discrete state, provided
    to the output port during the following interval. Graphically, a discrete clock
    sampled at 100 Hz would have the following time series:

    ```
      x(t)                  ●━
        |                   ┆
    .03 |              ●━━━━○
        |              ┆
    .02 |         ●━━━━○
        |         ┆
    .01 |    ●━━━━○
        |    ┆
      0 ●━━━━○----+----+----+-- t
        0   .01  .02  .03  .04
    ```

    The recorded states are the closed circles, which should be interpreted at index
    `n` as the value seen by all other blocks on the interval `(t[n], t[n+1])`.

    Input ports:
        None

    Output ports:
        (0) The sampled time.

    Parameters:
        dt:
            The sampling period of the clock.
        start_time:
            The simulation time at which the clock starts. Defaults to 0.
    """

    @parameters(static=["dt"])
    def __init__(self, dt, dtype=None, start_time=0, **kwargs):
        super().__init__(**kwargs)
        self.dtype = dtype or float
        start_time = npa.array(start_time, dtype=self.dtype)

        self.declare_output_port(
            self._output,
            period=dt,
            offset=0.0,
            requires_inputs=False,
            default_value=start_time,
            prerequisites_of_calc=[DependencyTicket.time],
        )

    def _output(self, time, _state, *_inputs, **_params):
        return npa.array(time, dtype=self.dtype)

DiscreteInitializer

Bases: LeafSystem

Discrete Initializer.

Outputs True for first discrete step, then outputs False there after. Or, outputs False for first discrete step, then outputs True there after. Practical for cases where it is necessary to have some signal fed initially by some initialization, but then after from else in the model.

Input ports

None

Output ports

(0) The dot product of the inputs.

Source code in jaxonomy/library/dynamics.py

class DiscreteInitializer(LeafSystem):
    """Discrete Initializer.

    Outputs True for first discrete step, then outputs False there after.
    Or, outputs False for first discrete step, then outputs True there after.
    Practical for cases where it is necessary to have some signal fed initially
    by some initialization, but then after from else in the model.

    Input ports:
        None

    Output ports:
        (0) The dot product of the inputs.
    """

    @parameters(static=["dt"], dynamic=["initial_state"])
    def __init__(self, dt, initial_state=True, **kwargs):
        super().__init__(**kwargs)
        self.dt = dt
        self.declare_output_port(self._output)
        self._periodic_update_idx = self.declare_periodic_update()

    def initialize(self, initial_state, dt=None):
        self.declare_discrete_state(default_value=initial_state, dtype=npa.bool_)
        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=npa.inf,
            offset=self.dt,
        )

    def reset_default_values(self, initial_state, dt=None):
        self.configure_discrete_state_default_value(default_value=initial_state)

    def _update(self, time, state, *_inputs, **_params):
        return npa.logical_not(state.discrete_state)

    def _output(self, _time, state, *_inputs, **_params):
        return state.discrete_state

DiscreteTimeLinearQuadraticRegulator

Bases: LeafSystem

Linear Quadratic Regulator (LQR) for a discrete-time system: x[k+1] = A x[k] + B u[k]. Computes the optimal control input: u[k] = -K x[k], where u minimises the cost function over [0, ∞)]: J = ∑(x[k].T Q x[k] + u[k].T R u[k]).

Input ports

(0) x[k]: state vector of the system.

Output ports

(0) u[k]: optimal control vector.

Parameters:

Name	Type	Description	Default
A		Array State matrix of the system.	required
B		Array Input matrix of the system.	required
Q		Array State cost matrix.	required
R		Array Input cost matrix.	required
dt		float Sampling period of the system.	required

Source code in jaxonomy/library/lqr.py

class DiscreteTimeLinearQuadraticRegulator(LeafSystem):
    """
    Linear Quadratic Regulator (LQR) for a discrete-time system:
            x[k+1] = A x[k] + B u[k].
    Computes the optimal control input:
            u[k] = -K x[k],
    where u minimises the cost function over [0, ∞)]:
            J = ∑(x[k].T Q x[k] + u[k].T R u[k]).

    Input ports:
        (0) x[k]: state vector of the system.

    Output ports:
        (0) u[k]: optimal control vector.

    Parameters:
        A: Array
            State matrix of the system.
        B: Array
            Input matrix of the system.
        Q: Array
            State cost matrix.
        R: Array
            Input cost matrix.
        dt: float
            Sampling period of the system.
    """

    def __init__(self, A, B, Q, R, dt, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.K, S, E = control.dlqr(A, B, Q, R)

        self.declare_input_port()  # for state x

        self.declare_output_port(
            self._get_opt_u,
            requires_inputs=True,
            period=dt,
            offset=0.0,
            default_value=jnp.zeros(B.shape[1]),
        )

    def _get_opt_u(self, time, state, x, **params):
        return jnp.matmul(-self.K, x)

DotProduct

Bases: ReduceBlock

Compute the dot product between the inputs.

This block dispatches to jax.numpy.dot, so the semantics, broadcasting rules, etc. are the same. See the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.dot.html

Input ports

(0) The first input vector. (1) The second input vector.

Output ports

(0) The dot product of the inputs.

Source code in jaxonomy/library/math_ops.py

class DotProduct(ReduceBlock):
    """Compute the dot product between the inputs.

    This block dispatches to `jax.numpy.dot`, so the semantics, broadcasting rules,
    etc. are the same.  See the JAX docs for details:
        https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.dot.html

    Input ports:
        (0) The first input vector.
        (1) The second input vector.

    Output ports:
        (0) The dot product of the inputs.
    """

    def __init__(self, **kwargs):
        super().__init__(2, self._compute_output, **kwargs)

    def _compute_output(self, inputs):
        return npa.dot(inputs[0], inputs[1])

EDMDResult dataclass

Fitted Koopman model (Williams, Kevrekidis & Rowley 2015).

Attributes:

Name	Type	Description
K	Any	Koopman operator on lifted observables, shape (L, L).
B	Any	Input operator on the lifted space, shape (L, m) (eDMDc) or None.
C	Any	De-lift matrix mapping lifted → physical state, shape (n, L).
dictionary	Callable	The observable dictionary g used for lifting.

Source code in jaxonomy/library/rom/koopman.py

@dataclass
class EDMDResult:
    """Fitted Koopman model (Williams, Kevrekidis & Rowley 2015).

    Attributes:
        K: Koopman operator on lifted observables, shape ``(L, L)``.
        B: Input operator on the lifted space, shape ``(L, m)`` (eDMDc) or ``None``.
        C: De-lift matrix mapping lifted → physical state, shape ``(n, L)``.
        dictionary: The observable dictionary ``g`` used for lifting.
    """

    K: Any
    B: Any
    C: Any
    dictionary: Callable

ERAResult dataclass

Minimal state-space realization from Markov parameters (Juang & Pappa 1985).

Attributes:

Name	Type	Description
A	Any	Realized r×r state matrix.
B	Any	Realized r×n_inputs input matrix.
C	Any	Realized n_outputs×r output matrix.
D	Any	Feedthrough n_outputs×n_inputs (the zeroth Markov parameter).
singular_values	Any	Hankel singular values (from the block-Hankel SVD).

Source code in jaxonomy/library/rom/dmd.py

@dataclass
class ERAResult:
    """Minimal state-space realization from Markov parameters (Juang & Pappa 1985).

    Attributes:
        A: Realized ``r×r`` state matrix.
        B: Realized ``r×n_inputs`` input matrix.
        C: Realized ``n_outputs×r`` output matrix.
        D: Feedthrough ``n_outputs×n_inputs`` (the zeroth Markov parameter).
        singular_values: Hankel singular values (from the block-Hankel SVD).
    """

    A: Any
    B: Any
    C: Any
    D: Any
    singular_values: Any

EdgeDetection

Bases: LeafSystem

Output is true only when the input signal changes in a specified way.

The block updates at a discrete rate, checking the boolean- or binary-valued input signal for changes. Available edge detection modes are: - "rising": Output is true when the input changes from False (0) to True (1). - "falling": Output is true when the input changes from True (1) to False (0). - "either": Output is true when the input changes in either direction

Input ports

(0) The input signal. Must be boolean or binary-valued.

Output ports

(0) The edge detection output signal. Boolean-valued.

Parameters:

Name	Type	Description	Default
dt		The sampling period of the block.	required
edge_detection		One of "rising", "falling", or "either". Determines the type of edge detection performed by the block.	required
initial_state		The initial value of the output signal.	False

Source code in jaxonomy/library/dynamics.py

class EdgeDetection(LeafSystem):
    """Output is true only when the input signal changes in a specified way.

    The block updates at a discrete rate, checking the boolean- or binary-valued input
    signal for changes.  Available edge detection modes are:
        - "rising": Output is true when the input changes from False (0) to True (1).
        - "falling": Output is true when the input changes from True (1) to False (0).
        - "either": Output is true when the input changes in either direction

    Input ports:
        (0) The input signal. Must be boolean or binary-valued.

    Output ports:
        (0) The edge detection output signal. Boolean-valued.

    Parameters:
        dt:
            The sampling period of the block.
        edge_detection:
            One of "rising", "falling", or "either". Determines the type of edge
            detection performed by the block.
        initial_state:
            The initial value of the output signal.
    """

    class DiscreteStateType(NamedTuple):
        prev_input: Array
        output: bool

    @parameters(dynamic=["initial_state"], static=["dt", "edge_detection"])
    def __init__(self, dt, edge_detection, initial_state=False, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dt = dt
        self.declare_input_port()

        # Declare the periodic update
        self._periodic_update_idx = self.declare_periodic_update()

        # Declare the output port
        self._output_port_idx = self.declare_output_port(
            self._output,
            prerequisites_of_calc=[DependencyTicket.xd, self.input_ports[0].ticket],
            requires_inputs=False,
        )

    def initialize(self, edge_detection, initial_state, dt=None):
        # Determine the type of edge detection
        _detection_funcs = {
            "rising": self._detect_rising,
            "falling": self._detect_falling,
            "either": self._detect_either,
        }
        if edge_detection not in _detection_funcs:
            raise ValueError(
                f"EdgeDetection block {self.name} has invalid selection "
                f"{edge_detection} for 'edge_detection'"
            )
        self._detect_edge = _detection_funcs[edge_detection]

        # The discrete state will contain the previous input value and the output.
        # T-037b: cast `prev_input` to `bool_` so the JSON round-trip can't change
        # its dtype — EdgeDetection is documented as bool-input/bool-output, and
        # without the cast a parameter-typed `initial_state` (e.g. Python `False`)
        # may downcast to float64 across JSON, while runtime updates carry the
        # input port dtype, breaking lax.cond branches in the reset map.
        self.declare_discrete_state(
            default_value=self.DiscreteStateType(
                prev_input=npa.asarray(initial_state, dtype=npa.bool_),
                output=npa.asarray(False, dtype=npa.bool_),
            ),
            as_array=False,
        )
        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=0.0,
        )

        # Declare the output port
        self.configure_output_port(
            self._output_port_idx,
            self._output,
            prerequisites_of_calc=[DependencyTicket.xd, self.input_ports[0].ticket],
            requires_inputs=False,
        )

    def reset_default_values(self, initial_state, dt=None):
        # The discrete state will contain the previous input value and the output
        self.configure_discrete_state_default_value(
            default_value=self.DiscreteStateType(
                prev_input=npa.asarray(initial_state, dtype=npa.bool_),
                output=npa.asarray(False, dtype=npa.bool_),
            ),
            as_array=False,
        )

    def _update(self, time, state, *inputs, **params):
        # Update the stored previous state
        # and the output as the result of the edge detection function.
        # T-037b: enforce the bool_ contract on every update so the reset-map
        # NamedTuple has a stable dtype regardless of how the upstream port
        # types its signal (e.g. Step emits 0/1 floats by default).
        (e,) = inputs
        return self.DiscreteStateType(
            prev_input=npa.asarray(e, dtype=npa.bool_),
            output=npa.asarray(
                self._detect_edge(time, state, e, **params), dtype=npa.bool_
            ),
        )

    def _output(self, _time, state, *_inputs, **_params):
        return state.discrete_state.output

    def _detect_rising(self, _time, state, *inputs, **_params):
        (e,) = inputs
        e_prev = state.discrete_state.prev_input
        e_prev = npa.array(e_prev)
        e = npa.array(e)
        not_e_prev = npa.logical_not(e_prev)
        return npa.logical_and(not_e_prev, e)

    def _detect_falling(self, _time, state, *inputs, **_params):
        (e,) = inputs
        e_prev = state.discrete_state.prev_input
        e_prev = npa.array(e_prev)
        e = npa.array(e)
        not_e = npa.logical_not(e)
        return npa.logical_and(e_prev, not_e)

    def _detect_either(self, _time, state, *inputs, **_params):
        (e,) = inputs
        e_prev = state.discrete_state.prev_input
        e_prev = npa.array(e_prev)
        e = npa.array(e)
        not_e_prev = npa.logical_not(e_prev)
        not_e = npa.logical_not(e)
        rising = npa.logical_and(not_e_prev, e)
        falling = npa.logical_and(e_prev, not_e)
        return npa.logical_or(rising, falling)

EnabledMode

Allowed string values for EnabledSubsystem.mode.

Source code in jaxonomy/framework/containers.py

class EnabledMode:
    """Allowed string values for ``EnabledSubsystem.mode``."""

    RESET = "reset"
    HOLD = "hold"
    PASSTHROUGH = "passthrough"

    @classmethod
    def valid(cls) -> tuple[str, ...]:
        return (cls.RESET, cls.HOLD, cls.PASSTHROUGH)

EnabledStateMode

Allowed string values for EnabledSubsystem.state_mode.

Controls how the continuous state (declared via state_dynamics) evolves while the enable signal is false:

• HOLD (default): freeze the state at its current value (xdot = 0 while disabled). Resumes integration on re-enable.
• RESET: snap the state back to initial_state on every disable→enable transition (so each enable window starts from the configured initial value). While disabled, the state is held.
• FREE: the state evolves according to state_dynamics regardless of enable. Only the output is masked per mode=.

Source code in jaxonomy/framework/containers.py

class EnabledStateMode:
    """Allowed string values for ``EnabledSubsystem.state_mode``.

    Controls how the *continuous state* (declared via ``state_dynamics``)
    evolves while the enable signal is false:

    - ``HOLD`` (default): freeze the state at its current value
      (``xdot = 0`` while disabled). Resumes integration on re-enable.
    - ``RESET``: snap the state back to ``initial_state`` on every
      disable→enable transition (so each enable window starts from the
      configured initial value). While disabled, the state is held.
    - ``FREE``: the state evolves according to ``state_dynamics``
      regardless of enable. Only the *output* is masked per ``mode=``.
    """

    HOLD = "hold"
    RESET = "reset"
    FREE = "free"

    @classmethod
    def valid(cls) -> tuple[str, ...]:
        return (cls.HOLD, cls.RESET, cls.FREE)

EnabledSubsystem

Bases: LeafSystem

Container block: run a submodel only while an enable signal is true.

This is the subsystem-framing wrapper around the existing :class:jaxonomy.library.Conditional primitive (T-009). It exists as a separate class so that:

• The block-diagram-vocabulary name EnabledSubsystem is discoverable next to the rest of the container family.
• We can later extend the mode="hold" path with subsystem-state semantics (per-block discrete-state binding) without disturbing the lighter Conditional primitive.

Parameters:

Name	Type	Description	Default
submodel	Callable	Callable f(*inputs) -> output (single output per phase 1). Must be JAX-traceable.	required
n_inputs	int	Number of submodel inputs (does NOT include the enable port). Input port 0 is always the enable signal; ports 1..n_inputs carry the submodel inputs.	1
mode	Literal['reset', 'passthrough', 'hold']	One of "reset" / "passthrough" / "hold". reset: when disabled, output = initial_value. passthrough: when disabled, output = first user input (input port 1). Submodel and passthrough output must broadcast-compatibly. hold: when disabled, output holds the most recent snapshot taken at hold_period. Requires a positive hold_period.	RESET
initial_value		Output value when disabled in reset mode, and the seed for the held discrete state in hold mode. Used to infer output shape/dtype.	0.0
hold_period	float | None	Sample period (seconds) for the held snapshot in hold mode. Required iff mode == "hold".	None
state_mode	Literal['hold', 'reset', 'free']	One of "hold" / "reset" / "free". Controls the continuous-state behaviour while disabled (independent of mode= which gates only the output). See :class:EnabledStateMode. Default "hold". Only has an effect when state_dynamics is provided; for the stateless submodel default this kwarg is validated but otherwise a no-op (so the default-off path is byte-equivalent to phase 1).	HOLD
state_dynamics	Callable | None	Optional callable f(t, x, *user_inputs) -> xdot defining a continuous state for the EnabledSubsystem itself. When provided, the block declares a continuous state seeded by initial_state (or initial_value if initial_state is None) and applies state_mode semantics around it. When omitted, the block has no continuous state and behaves exactly as in T-120 phase 1.	None
initial_state		Initial value of the continuous state. Required when state_dynamics is provided.	None
name		Optional block name.	required

Source code in jaxonomy/framework/containers.py

class EnabledSubsystem(LeafSystem):
    """Container block: run a submodel only while an enable signal is true.

    This is the subsystem-framing wrapper around the existing
    :class:`jaxonomy.library.Conditional` primitive (T-009). It exists
    as a separate class so that:

    - The block-diagram-vocabulary name ``EnabledSubsystem`` is discoverable
      next to the rest of the container family.
    - We can later extend the ``mode="hold"`` path with subsystem-state
      semantics (per-block discrete-state binding) without disturbing
      the lighter ``Conditional`` primitive.

    Args:
        submodel: Callable ``f(*inputs) -> output`` (single output per
            phase 1). Must be JAX-traceable.
        n_inputs: Number of submodel inputs (does NOT include the
            enable port). Input port 0 is always the enable signal;
            ports 1..n_inputs carry the submodel inputs.
        mode: One of ``"reset"`` / ``"passthrough"`` / ``"hold"``.

            - ``reset``: when disabled, output = ``initial_value``.
            - ``passthrough``: when disabled, output = first user input
              (input port 1). Submodel and passthrough output must
              broadcast-compatibly.
            - ``hold``: when disabled, output holds the most recent
              snapshot taken at ``hold_period``. Requires a positive
              ``hold_period``.
        initial_value: Output value when disabled in reset mode, and
            the seed for the held discrete state in hold mode. Used to
            infer output shape/dtype.
        hold_period: Sample period (seconds) for the held snapshot in
            hold mode. Required iff ``mode == "hold"``.
        state_mode: One of ``"hold"`` / ``"reset"`` / ``"free"``.
            Controls the *continuous-state* behaviour while disabled
            (independent of ``mode=`` which gates only the output).
            See :class:`EnabledStateMode`. Default ``"hold"``. Only has
            an effect when ``state_dynamics`` is provided; for the
            stateless submodel default this kwarg is validated but
            otherwise a no-op (so the default-off path is byte-equivalent
            to phase 1).
        state_dynamics: Optional callable
            ``f(t, x, *user_inputs) -> xdot`` defining a continuous
            state for the EnabledSubsystem itself. When provided, the
            block declares a continuous state seeded by ``initial_state``
            (or ``initial_value`` if ``initial_state`` is None) and
            applies ``state_mode`` semantics around it. When omitted,
            the block has no continuous state and behaves exactly as in
            T-120 phase 1.
        initial_state: Initial value of the continuous state. Required
            when ``state_dynamics`` is provided.
        name: Optional block name.
    """

    def __init__(
        self,
        submodel: Callable,
        n_inputs: int = 1,
        mode: Literal["reset", "passthrough", "hold"] = EnabledMode.RESET,
        initial_value=0.0,
        hold_period: float | None = None,
        state_mode: Literal["hold", "reset", "free"] = EnabledStateMode.HOLD,
        state_dynamics: Callable | None = None,
        initial_state=None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if mode not in EnabledMode.valid():
            raise ValueError(
                f"EnabledSubsystem: mode must be one of "
                f"{EnabledMode.valid()!r}, got {mode!r}"
            )
        if state_mode not in EnabledStateMode.valid():
            raise ValueError(
                f"EnabledSubsystem: state_mode must be one of "
                f"{EnabledStateMode.valid()!r}, got {state_mode!r}"
            )
        if mode == EnabledMode.HOLD and not hold_period:
            raise ValueError(
                "EnabledSubsystem(mode='hold') requires a positive "
                "hold_period to determine the snapshot sample rate."
            )
        if n_inputs < 0:
            raise ValueError(
                f"EnabledSubsystem: n_inputs must be >= 0, got {n_inputs}"
            )
        if mode == EnabledMode.PASSTHROUGH and n_inputs < 1:
            raise ValueError(
                "EnabledSubsystem(mode='passthrough') requires at least "
                "one user input to use as the bypass signal."
            )
        if state_dynamics is not None and initial_state is None:
            raise ValueError(
                "EnabledSubsystem: state_dynamics= requires an "
                "initial_state= value (the seed for the continuous "
                "state)."
            )

        self._submodel = submodel
        self._mode = mode
        self._initial = jnp.asarray(initial_value)
        self._state_mode = state_mode
        self._state_dynamics = state_dynamics
        self._initial_state = (
            jnp.asarray(initial_state) if initial_state is not None else None
        )

        # Port 0 is always enable; remaining ports are submodel inputs.
        self.declare_input_port(name="enable")
        for i in range(n_inputs):
            self.declare_input_port(name=f"u_{i}")

        if mode == EnabledMode.HOLD:
            self.declare_discrete_state(default_value=self._initial)
            self.declare_periodic_update(
                self._hold_update,
                period=float(hold_period),
                offset=0.0,
            )

        # T-120-followup-enabled-cont-state: declare a continuous state on
        # this block when the user supplied a state_dynamics callable. The
        # state_mode kwarg controls how the state evolves vs. the enable
        # signal. The default-off path (state_dynamics=None) bypasses this
        # block entirely → byte-equivalent to phase 1.
        if self._state_dynamics is not None:
            self.declare_continuous_state(
                default_value=self._initial_state,
                ode=self._wrapped_ode,
            )

            if self._state_mode == EnabledStateMode.RESET:
                # Snap the continuous state back to its initial value on
                # every enable transition (rising or falling), so each
                # disable window leaves the state at the seed and each
                # re-enable starts from the same point. Pair this with
                # the "hold" ode-zeroing during the disabled window so
                # the state actually stays at the seed instead of drifting.
                #
                # Treat the enable signal as a centred continuous guard
                # (``enable - 0.5``): zero-crossings of this expression
                # match enable transitions in either direction. The
                # framework's continuous detector preserves the float
                # carry-type for the guard signal — using
                # ``direction="edge_detection"`` with a boolean guard
                # would conflict with the simulator's float-typed
                # internal carry, so we go through the continuous path.
                self.declare_zero_crossing(
                    guard=self._enable_guard,
                    reset_map=self._reset_continuous_state,
                    direction="crosses_zero",
                    name="enable_reset",
                )

        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    # ── callbacks ─────────────────────────────────────────────────────────

    def _submodel_output(self, inputs):
        user_inputs = inputs[1:]  # skip enable
        return jnp.asarray(self._submodel(*user_inputs))

    def _hold_update(self, time, state, *inputs, **params):
        enable = jnp.asarray(inputs[0]).astype(bool)
        y_sub = self._submodel_output(inputs)
        return jnp.where(enable, y_sub, state.discrete_state)

    def _compute_output(self, time, state, *inputs, **params):
        enable = jnp.asarray(inputs[0]).astype(bool)
        y_sub = self._submodel_output(inputs)

        if self._mode == EnabledMode.RESET:
            return jnp.where(enable, y_sub, self._initial)

        if self._mode == EnabledMode.HOLD:
            return jnp.where(enable, y_sub, state.discrete_state)

        # passthrough
        bypass = jnp.asarray(inputs[1])
        return jnp.where(enable, y_sub, bypass)

    # ── continuous-state callbacks (T-120-followup-enabled-cont-state) ───

    def _wrapped_ode(self, time, state, *inputs, **params):
        """Apply ``state_mode`` semantics to the user's ``state_dynamics``.

        - ``hold``: multiply the user's xdot by the enable flag. While
          disabled the ode returns zero, so the integrator preserves the
          state across the disabled window.
        - ``reset``: same as ``hold`` for the per-step ode (state stays
          at the seed during the disabled window); the reset to the
          initial value is enforced by a zero-crossing reset_map on the
          enable-edge.
        - ``free``: pass the user's xdot through untouched. The state
          evolves regardless of enable (only the output port is masked).
        """
        enable = jnp.asarray(inputs[0]).astype(bool)
        user_inputs = inputs[1:]
        xc = state.continuous_state
        xdot = self._state_dynamics(time, xc, *user_inputs)
        xdot = jnp.asarray(xdot)

        if self._state_mode == EnabledStateMode.FREE:
            return xdot

        # HOLD and RESET both gate the per-step derivative. RESET layers
        # the additional snap-to-initial behaviour via the zero-crossing
        # event registered in __init__.
        scale = jnp.asarray(enable, dtype=xdot.dtype)
        return xdot * scale

    def _enable_guard(self, time, state, *inputs, **params):
        """Continuous guard for enable transitions.

        Returns ``enable - 0.5`` so any 0↔1 transition crosses zero. Used
        with ``direction="crosses_zero"`` to fire the continuous-state
        reset on either a rising or falling enable edge.
        """
        return jnp.asarray(inputs[0]) - 0.5

    def _reset_continuous_state(self, time, state, *inputs, **params):
        """Snap the continuous state back to its initial value."""
        return state.with_continuous_state(
            jnp.asarray(self._initial_state, dtype=state.continuous_state.dtype)
        )

Exponent

Bases: FeedthroughBlock

Compute the exponential of the input signal.

Input ports

(0) The input signal.

Output ports

(0) The exponential of the input signal.

Parameters:

Name	Type	Description	Default
base		One of "exp" or "2". Determines the base of the exponential function.	required

Source code in jaxonomy/library/math_ops.py

class Exponent(FeedthroughBlock):
    """Compute the exponential of the input signal.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The exponential of the input signal.

    Parameters:
        base:
            One of "exp" or "2". Determines the base of the exponential function.
    """

    @parameters(static=["base"])
    def __init__(self, base, **kwargs):
        super().__init__(None, **kwargs)

    def initialize(self, base):
        func_lookup = {"exp": npa.exp, "2": npa.exp2}
        if base not in func_lookup:
            raise BlockParameterError(
                message=f"Exponent block {self.name} has invalid selection {base} for 'base'. Valid selections: "
                + ", ".join([k for k in func_lookup.keys()]),
                parameter_name="base",
            )
        self.replace_op(func_lookup[base])

ExtendedKalmanFilter

Bases: KalmanFilterBase

Extended Kalman Filter (EKF) for the following system:

```
x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
y[n]   = g(x[n], u[n]) + v[n]

E(w[n]) = E(v[n]) = 0
E(w[n]w'[n]) = Q(t[n], x[n], u[n])
E(v[n]v'[n] = R(t[n])
E(w[n]v'[n] = N(t[n]) = 0
```

f and g are discrete-time functions of state x[n] and control u[n], while RandGare discrete-time functions of timet[n].Qis a discrete-time function oft[n], x[n], u[n]`. This last aspect is included for zero-order-hold discretization of a continuous-time system

Input ports

(0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n

Output ports

(1) x_hat[n] : state vector estimate at timestep n

Parameters:

Name	Type	Description	Default
dt		float Time step of the discrete-time system	required
forward		Callable A function with signature f(x[n], u[n]) -> x[n+1] that represents f in the above equations.	required
observation		Callable A function with signature g(x[n], u[n]) -> y[n] that represents g in the above equations.	required
G_func		Callable A function with signature G(t[n]) -> G[n] that represents G in the above equations.	required
Q_func		Callable A function with signature Q(t[n], x[n], u[n]) -> Q[n] that represents Q in the above equations.	required
R_func		Callable A function with signature R(t[n]) -> R[n] that represents R in the above equations.	required
x_hat_0		ndarray Initial state estimate	required
P_hat_0		ndarray Initial state covariance matrix estimate	required

Source code in jaxonomy/library/state_estimators/extended_kalman_filter.py

class ExtendedKalmanFilter(KalmanFilterBase):
    """
    Extended Kalman Filter (EKF) for the following system:

        ```
        x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
        y[n]   = g(x[n], u[n]) + v[n]

        E(w[n]) = E(v[n]) = 0
        E(w[n]w'[n]) = Q(t[n], x[n], u[n])
        E(v[n]v'[n] = R(t[n])
        E(w[n]v'[n] = N(t[n]) = 0
        ```

    `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`,
    while R` and `G` are discrete-time functions of time `t[n]`. `Q` is a discrete-time
    function of `t[n], x[n], u[n]`. This last aspect is included for zero-order-hold
    discretization of a continuous-time system

    Input ports:
        (0) u[n] : control vector at timestep n
        (1) y[n] : measurement vector at timestep n

    Output ports:
        (1) x_hat[n] : state vector estimate at timestep n

    Parameters:
        dt: float
            Time step of the discrete-time system
        forward: Callable
            A function with signature f(x[n], u[n]) -> x[n+1] that represents `f` in
            the above equations.
        observation: Callable
            A function with signature g(x[n], u[n]) -> y[n] that represents `g` in
            the above equations.
        G_func: Callable
            A function with signature G(t[n]) -> G[n] that represents `G` in
            the above equations.
        Q_func: Callable
            A function with signature Q(t[n], x[n], u[n]) -> Q[n] that represents `Q`
            in the above equations.
        R_func: Callable
            A function with signature R(t[n]) -> R[n] that represents `R` in
            the above equations.
        x_hat_0: ndarray
            Initial state estimate
        P_hat_0: ndarray
            Initial state covariance matrix estimate
    """

    @parameters(
        static=[
            "dt",
            "forward",
            "observation",
            "G_func",
            "Q_func",
            "R_func",
            "x_hat_0",
            "P_hat_0",
        ],
    )
    def __init__(
        self,
        dt,
        forward,
        observation,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        P_hat_0,
        is_feedthrough=True,  # TODO: determine automatically?
        name=None,
        **kwargs,
    ):
        super().__init__(dt, x_hat_0, P_hat_0, is_feedthrough, name, **kwargs)

    def initialize(
        self,
        dt,
        forward,
        observation,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        P_hat_0,
    ):
        self.G_func = G_func
        self.Q_func = Q_func
        self.R_func = R_func

        self.nx = x_hat_0.size
        self.ny = self.R_func(0.0).shape[0]

        self.forward = forward
        self.observation = observation

        self.jac_forward = jax.jacfwd(forward)
        self.jac_observation = jax.jacfwd(observation)

        self.eye_x = jnp.eye(self.nx)

    def _correct(self, time, x_hat_minus, P_hat_minus, *inputs):
        u, y = inputs
        y = jnp.atleast_1d(y)

        C = self.jac_observation(x_hat_minus, u).reshape((self.ny, self.nx))

        R = self.R_func(time)

        # Kalman gain via a linear solve instead of an explicit inverse of the
        # innovation covariance S — more numerically stable (matches the base
        # KalmanFilter._correct idiom). K = P Cᵀ S⁻¹ solves K S = P Cᵀ.
        S = C @ P_hat_minus @ C.T + R
        K = jnp.linalg.solve(S.T, (P_hat_minus @ C.T).T).T

        x_hat_plus = x_hat_minus + jnp.dot(
            K, y - self.observation(x_hat_minus, u)
        )  # n|n

        P_hat_plus = jnp.matmul(self.eye_x - jnp.matmul(K, C), P_hat_minus)  # n|n

        return x_hat_plus, P_hat_plus

    def _propagate(self, time, x_hat_plus, P_hat_plus, *inputs):
        # Predict -- x_hat_plus of current step is propagated to be the
        # x_hat_minus of the next step
        # k+1|k in current step is n|n-1 for next step

        u, y = inputs
        u = jnp.atleast_1d(u)

        A = self.jac_forward(x_hat_plus, u).reshape((self.nx, self.nx))

        G = self.G_func(time)
        Q = self.Q_func(time, x_hat_plus, u)
        GQGT = G @ Q @ G.T

        x_hat_minus = self.forward(x_hat_plus, u)  # n+1|n
        P_hat_minus = A @ P_hat_plus @ A.T + GQGT  # n+1|n

        return x_hat_minus, P_hat_minus

    #######################################
    # Make filter for a continuous plant  #
    #######################################

    @staticmethod
    @with_resolved_parameters
    def for_continuous_plant(
        plant,
        dt,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        P_hat_0,
        discretization_method="euler",
        discretized_noise=False,
        name=None,
        ui_id=None,
    ):
        """
        Extended Kalman Filter system for a continuous-time plant.

        The input plant contains the deterministic forms of the forward and observation
        operators:

        ```
            dx/dt = f(x,u)
            y = g(x,u)
        ```

        Note: (i) Only plants with one vector-valued input and one vector-valued output
        are currently supported. Furthermore, the plant LeafSystem/Diagram should have
        only one vector-valued integrator; (ii) the user may pass a plant with
        disturbances (not recommended) as the input plant. In this case, the forward
        and observation evaluations will be corrupted by noise.

        A plant with disturbances of the following form is then considered:

        ```
            dx/dt = f(x,u) + G(t) w         -- (C1)
            y = g(x,u) +  v                 -- (C2)
        ```

        where:

            `w` represents the process noise,
            `v` represents the measurement noise,

        and

        ```
            E(w) = E(v) = 0
            E(ww') = Q(t)
            E(vv') = R(t)
            E(wv') = N(t) = 0
        ```

        This plant is discretized to obtain the following form:

        ```
            x[n+1] = fd(x[n], u[n]) + Gd w[n]  -- (D1)
            y[n]   = gd(x[n], u[n]) + v[n]     -- (D2)

            E(w[n]) = E(v[n]) = 0
            E(w[n]w'[n]) = Qd
            E(v[n]v'[n] = Rd
            E(w[n]v'[n] = Nd = 0
        ```

        The above discretization is performed either via the `euler` or the `zoh`
        method, and an Extended Kalman Filter estimator for the system of equations
        (D1) and (D2) is returned.

        Note: If `discretized_noise` is True, then it is assumed that the user is
        directly providing Gd, Qd and Rd. If False, then Qd and Rd are computed from
        continuous-time Q, R, and G, and Gd is set to an Identity matrix.

        The returned system will have:

        Input ports:
            (0) u[n] : control vector at timestep n
            (1) y[n] : measurement vector at timestep n

        Output ports:
            (1) x_hat[n] : state vector estimate at timestep n

        Parameters:
            plant : a `Plant` object which can be a LeafSystem or a Diagram.
            dt: float
                Time step for the discretization.
            G_func: Callable
                A function with signature G(t) -> G that represents `G` in
                the continuous-time equations (C1) and (C2).
            Q_func: Callable
                A function with signature Q(t) -> Q that represents `Q` in
                the continuous-time equations (C1) and (C2).
            R_func: Callable
                A function with signature R(t) -> R that represents `R` in
                the continuous-time equations (C1) and (C2).
            x_hat_0: ndarray
                Initial state estimate
            P_hat_0: ndarray
                Initial state covariance matrix estimate. If `None`, an Identity
                matrix is assumed.
            discretization_method: str ("euler" or "zoh")
                Method to discretize the continuous-time plant. Default is "euler".
            discretized_noise: bool
                Whether the user is directly providing Gd, Qd and Rd. Default is False.
                If True, `G_func`, `Q_func`, and `R_func` provide Gd(t), Qd(t), and
                Rd(t), respectively.
        """

        (
            forward,
            observation,
            Gd_func,
            Qd_func,
            Rd_func,
        ) = prepare_continuous_plant_for_nonlinear_kalman_filter(
            plant,
            dt,
            G_func,
            Q_func,
            R_func,
            x_hat_0,
            discretization_method,
            discretized_noise,
        )

        nx = x_hat_0.size
        if P_hat_0 is None:
            P_hat_0 = jnp.eye(nx)

        # TODO: If Gd_func is None, compute Gd automatically with u = u + w

        ekf = ExtendedKalmanFilter(
            dt,
            forward,
            observation,
            Gd_func,
            Qd_func,
            Rd_func,
            x_hat_0,
            P_hat_0,
            name=name,
            ui_id=ui_id,
        )

        return ekf

    ###################################################################################
    # Make filter from direct specification of forward/observaton operators and noise #
    ###################################################################################

    @staticmethod
    @with_resolved_parameters
    def from_operators(
        dt,
        forward,
        observation,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        P_hat_0,
        name=None,
        ui_id=None,
    ):
        """
        Extended Kalman Filter (UKF) for the following system:

        ```
            x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
            y[n]   = g(x[n], u[n]) + v[n]

            E(w[n]) = E(v[n]) = 0
            E(w[n]w'[n]) = Q(t[n], x[n], u[n])
            E(v[n]v'[n] = R(t[n])
            E(w[n]v'[n] = N(t[n]) = 0
        ```

        `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`,
        while `Q` and `R` and `G` are discrete-time functions of time `t[n]`.

        Input ports:
            (0) u[n] : control vector at timestep n
            (1) y[n] : measurement vector at timestep n

        Output ports:
            (1) x_hat[n] : state vector estimate at timestep n

        Parameters:
            dt: float
                Time step of the discrete-time system
            forward: Callable
                A function with signature f(x[n], u[n]) -> x[n+1] that represents `f`
                in the above equations.
            observation: Callable
                A function with signature g(x[n], u[n]) -> y[n] that represents `g` in
                the above equations.
            G_func: Callable
                A function with signature G(t[n]) -> G[n] that represents `G` in
                the above equations.
            Q_func: Callable
                A function with signature Q(t[n]) -> Q[n] that represents
                `Q` in the above equations.
            R_func: Callable
                A function with signature R(t[n]) -> R[n] that represents `R` in
                the above equations.
            x_hat_0: ndarray
                Initial state estimate
            P_hat_0: ndarray
                Initial state covariance matrix estimate
        """

        def Q_func_aug(t, x_k, u_k):
            return Q_func(t)

        ekf = ExtendedKalmanFilter(
            dt,
            forward,
            observation,
            G_func,
            Q_func_aug,
            R_func,
            x_hat_0,
            P_hat_0,
            name=name,
            ui_id=ui_id,
        )

        return ekf

        return ekf

for_continuous_plant(plant, dt, G_func, Q_func, R_func, x_hat_0, P_hat_0, discretization_method='euler', discretized_noise=False, name=None, ui_id=None) staticmethod

Extended Kalman Filter system for a continuous-time plant.

The input plant contains the deterministic forms of the forward and observation operators:

    dx/dt = f(x,u)
    y = g(x,u)

Note: (i) Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator; (ii) the user may pass a plant with disturbances (not recommended) as the input plant. In this case, the forward and observation evaluations will be corrupted by noise.

A plant with disturbances of the following form is then considered:

    dx/dt = f(x,u) + G(t) w         -- (C1)
    y = g(x,u) +  v                 -- (C2)

where:

`w` represents the process noise,
`v` represents the measurement noise,

and

    E(w) = E(v) = 0
    E(ww') = Q(t)
    E(vv') = R(t)
    E(wv') = N(t) = 0

This plant is discretized to obtain the following form:

    x[n+1] = fd(x[n], u[n]) + Gd w[n]  -- (D1)
    y[n]   = gd(x[n], u[n]) + v[n]     -- (D2)

    E(w[n]) = E(v[n]) = 0
    E(w[n]w'[n]) = Qd
    E(v[n]v'[n] = Rd
    E(w[n]v'[n] = Nd = 0

The above discretization is performed either via the euler or the zoh method, and an Extended Kalman Filter estimator for the system of equations (D1) and (D2) is returned.

Note: If discretized_noise is True, then it is assumed that the user is directly providing Gd, Qd and Rd. If False, then Qd and Rd are computed from continuous-time Q, R, and G, and Gd is set to an Identity matrix.

The returned system will have:

Input ports

(0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n

Output ports

(1) x_hat[n] : state vector estimate at timestep n

Parameters:

Name	Type	Description	Default
plant		a Plant object which can be a LeafSystem or a Diagram.	required
dt		float Time step for the discretization.	required
G_func		Callable A function with signature G(t) -> G that represents G in the continuous-time equations (C1) and (C2).	required
Q_func		Callable A function with signature Q(t) -> Q that represents Q in the continuous-time equations (C1) and (C2).	required
R_func		Callable A function with signature R(t) -> R that represents R in the continuous-time equations (C1) and (C2).	required
x_hat_0		ndarray Initial state estimate	required
P_hat_0		ndarray Initial state covariance matrix estimate. If None, an Identity matrix is assumed.	required
discretization_method		str ("euler" or "zoh") Method to discretize the continuous-time plant. Default is "euler".	'euler'
discretized_noise		bool Whether the user is directly providing Gd, Qd and Rd. Default is False. If True, G_func, Q_func, and R_func provide Gd(t), Qd(t), and Rd(t), respectively.	False

Source code in jaxonomy/library/state_estimators/extended_kalman_filter.py

@staticmethod
@with_resolved_parameters
def for_continuous_plant(
    plant,
    dt,
    G_func,
    Q_func,
    R_func,
    x_hat_0,
    P_hat_0,
    discretization_method="euler",
    discretized_noise=False,
    name=None,
    ui_id=None,
):
    """
    Extended Kalman Filter system for a continuous-time plant.

    The input plant contains the deterministic forms of the forward and observation
    operators:

    ```
        dx/dt = f(x,u)
        y = g(x,u)
    ```

    Note: (i) Only plants with one vector-valued input and one vector-valued output
    are currently supported. Furthermore, the plant LeafSystem/Diagram should have
    only one vector-valued integrator; (ii) the user may pass a plant with
    disturbances (not recommended) as the input plant. In this case, the forward
    and observation evaluations will be corrupted by noise.

    A plant with disturbances of the following form is then considered:

    ```
        dx/dt = f(x,u) + G(t) w         -- (C1)
        y = g(x,u) +  v                 -- (C2)
    ```

    where:

        `w` represents the process noise,
        `v` represents the measurement noise,

    and

    ```
        E(w) = E(v) = 0
        E(ww') = Q(t)
        E(vv') = R(t)
        E(wv') = N(t) = 0
    ```

    This plant is discretized to obtain the following form:

    ```
        x[n+1] = fd(x[n], u[n]) + Gd w[n]  -- (D1)
        y[n]   = gd(x[n], u[n]) + v[n]     -- (D2)

        E(w[n]) = E(v[n]) = 0
        E(w[n]w'[n]) = Qd
        E(v[n]v'[n] = Rd
        E(w[n]v'[n] = Nd = 0
    ```

    The above discretization is performed either via the `euler` or the `zoh`
    method, and an Extended Kalman Filter estimator for the system of equations
    (D1) and (D2) is returned.

    Note: If `discretized_noise` is True, then it is assumed that the user is
    directly providing Gd, Qd and Rd. If False, then Qd and Rd are computed from
    continuous-time Q, R, and G, and Gd is set to an Identity matrix.

    The returned system will have:

    Input ports:
        (0) u[n] : control vector at timestep n
        (1) y[n] : measurement vector at timestep n

    Output ports:
        (1) x_hat[n] : state vector estimate at timestep n

    Parameters:
        plant : a `Plant` object which can be a LeafSystem or a Diagram.
        dt: float
            Time step for the discretization.
        G_func: Callable
            A function with signature G(t) -> G that represents `G` in
            the continuous-time equations (C1) and (C2).
        Q_func: Callable
            A function with signature Q(t) -> Q that represents `Q` in
            the continuous-time equations (C1) and (C2).
        R_func: Callable
            A function with signature R(t) -> R that represents `R` in
            the continuous-time equations (C1) and (C2).
        x_hat_0: ndarray
            Initial state estimate
        P_hat_0: ndarray
            Initial state covariance matrix estimate. If `None`, an Identity
            matrix is assumed.
        discretization_method: str ("euler" or "zoh")
            Method to discretize the continuous-time plant. Default is "euler".
        discretized_noise: bool
            Whether the user is directly providing Gd, Qd and Rd. Default is False.
            If True, `G_func`, `Q_func`, and `R_func` provide Gd(t), Qd(t), and
            Rd(t), respectively.
    """

    (
        forward,
        observation,
        Gd_func,
        Qd_func,
        Rd_func,
    ) = prepare_continuous_plant_for_nonlinear_kalman_filter(
        plant,
        dt,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        discretization_method,
        discretized_noise,
    )

    nx = x_hat_0.size
    if P_hat_0 is None:
        P_hat_0 = jnp.eye(nx)

    # TODO: If Gd_func is None, compute Gd automatically with u = u + w

    ekf = ExtendedKalmanFilter(
        dt,
        forward,
        observation,
        Gd_func,
        Qd_func,
        Rd_func,
        x_hat_0,
        P_hat_0,
        name=name,
        ui_id=ui_id,
    )

    return ekf

from_operators(dt, forward, observation, G_func, Q_func, R_func, x_hat_0, P_hat_0, name=None, ui_id=None) staticmethod

Extended Kalman Filter (UKF) for the following system:

    x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
    y[n]   = g(x[n], u[n]) + v[n]

    E(w[n]) = E(v[n]) = 0
    E(w[n]w'[n]) = Q(t[n], x[n], u[n])
    E(v[n]v'[n] = R(t[n])
    E(w[n]v'[n] = N(t[n]) = 0

f and g are discrete-time functions of state x[n] and control u[n], while Q and R and G are discrete-time functions of time t[n].

Input ports

(0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n

Output ports

(1) x_hat[n] : state vector estimate at timestep n

Parameters:

Name	Type	Description	Default
dt		float Time step of the discrete-time system	required
forward		Callable A function with signature f(x[n], u[n]) -> x[n+1] that represents f in the above equations.	required
observation		Callable A function with signature g(x[n], u[n]) -> y[n] that represents g in the above equations.	required
G_func		Callable A function with signature G(t[n]) -> G[n] that represents G in the above equations.	required
Q_func		Callable A function with signature Q(t[n]) -> Q[n] that represents Q in the above equations.	required
R_func		Callable A function with signature R(t[n]) -> R[n] that represents R in the above equations.	required
x_hat_0		ndarray Initial state estimate	required
P_hat_0		ndarray Initial state covariance matrix estimate	required

Source code in jaxonomy/library/state_estimators/extended_kalman_filter.py

@staticmethod
@with_resolved_parameters
def from_operators(
    dt,
    forward,
    observation,
    G_func,
    Q_func,
    R_func,
    x_hat_0,
    P_hat_0,
    name=None,
    ui_id=None,
):
    """
    Extended Kalman Filter (UKF) for the following system:

    ```
        x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
        y[n]   = g(x[n], u[n]) + v[n]

        E(w[n]) = E(v[n]) = 0
        E(w[n]w'[n]) = Q(t[n], x[n], u[n])
        E(v[n]v'[n] = R(t[n])
        E(w[n]v'[n] = N(t[n]) = 0
    ```

    `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`,
    while `Q` and `R` and `G` are discrete-time functions of time `t[n]`.

    Input ports:
        (0) u[n] : control vector at timestep n
        (1) y[n] : measurement vector at timestep n

    Output ports:
        (1) x_hat[n] : state vector estimate at timestep n

    Parameters:
        dt: float
            Time step of the discrete-time system
        forward: Callable
            A function with signature f(x[n], u[n]) -> x[n+1] that represents `f`
            in the above equations.
        observation: Callable
            A function with signature g(x[n], u[n]) -> y[n] that represents `g` in
            the above equations.
        G_func: Callable
            A function with signature G(t[n]) -> G[n] that represents `G` in
            the above equations.
        Q_func: Callable
            A function with signature Q(t[n]) -> Q[n] that represents
            `Q` in the above equations.
        R_func: Callable
            A function with signature R(t[n]) -> R[n] that represents `R` in
            the above equations.
        x_hat_0: ndarray
            Initial state estimate
        P_hat_0: ndarray
            Initial state covariance matrix estimate
    """

    def Q_func_aug(t, x_k, u_k):
        return Q_func(t)

    ekf = ExtendedKalmanFilter(
        dt,
        forward,
        observation,
        G_func,
        Q_func_aug,
        R_func,
        x_hat_0,
        P_hat_0,
        name=name,
        ui_id=ui_id,
    )

    return ekf

    return ekf

FeedthroughBlock

Bases: LeafSystem

Simple feedthrough blocks with a function of a single input

Source code in jaxonomy/library/generic.py

class FeedthroughBlock(LeafSystem):
    """Simple feedthrough blocks with a function of a single input"""

    def __init__(self, func, parameters={}, **kwargs):
        super().__init__(**kwargs)
        self.declare_input_port()
        self._output_port_idx = self.declare_output_port(
            None,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )
        self.replace_op(func)

    def replace_op(self, func):
        def _callback(time, state, *inputs, **parameters):
            return func(*inputs, **parameters)

        self.configure_output_port(
            self._output_port_idx,
            _callback,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )

FilterDiscrete

Bases: LeafSystem

Finite Impulse Response (FIR) filter.

Similar to https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.lfilter.html Note: does not implement the IIR filter.

Input ports

(0) The input signal.

Output ports

(0) The filtered signal.

Parameters:

Name	Type	Description	Default
b_coefficients		Array of filter coefficients.	required

Source code in jaxonomy/library/dynamics.py

class FilterDiscrete(LeafSystem):
    """Finite Impulse Response (FIR) filter.

    Similar to https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.lfilter.html
    Note: does not implement the IIR filter.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The filtered signal.

    Parameters:
        b_coefficients:
            Array of filter coefficients.
    """

    @parameters(static=["dt", "b_coefficients"])
    def __init__(
        self,
        dt,
        b_coefficients,
        *args,
        dtype=None,
        **kwargs,
    ):
        # T-038a-followup-other-blocks: per-block dtype override; stored
        # outside the @parameters list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(*args, **kwargs)
        self.dt = dt
        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, b_coefficients, dt=None):
        # T-037a: establish a block-level dtype contract. The discrete-state
        # delay-line and the runtime input share their dtype with
        # `b_coefficients`, so JSON round-trip (which can downcast list-typed
        # coefficients) cannot make the reset-map and the declared default
        # disagree. Without this, fresh-built blocks could end up with a
        # float32 default and a float64 update (or vice versa), tripping
        # JAX's lax.cond dtype check on the reloaded diagram.
        # T-038a-followup-other-blocks: an explicit per-block ``dtype=``
        # overrides the inferred coefficient dtype so the entire delay
        # line and feed-forward sum run at the requested precision.
        if self._dtype is not None:
            b_arr = npa.asarray(b_coefficients).astype(self._dtype)
            self._state_dtype = self._dtype
        else:
            b_arr = npa.asarray(b_coefficients)
            self._state_dtype = npa.result_type(b_arr)
        initial_state = npa.zeros(len(b_coefficients) - 1, dtype=self._state_dtype)
        self.declare_discrete_state(default_value=initial_state)

        self.is_feedthrough = bool(b_coefficients[0] != 0)
        self.b_coefficients = b_arr
        prerequisites_of_calc = []
        if self.is_feedthrough:
            prerequisites_of_calc.append(self.input_ports[0].ticket)

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=self.dt,
        )

        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=self.dt,
            requires_inputs=self.is_feedthrough,
            prerequisites_of_calc=prerequisites_of_calc,
        )

    def _update(self, _time, state, u, **_parameters):
        xd = state.discrete_state
        # T-037a: cast u to the canonical state dtype so the FIFO push lands
        # the same dtype on every step, fresh-built or round-tripped.
        u = npa.asarray(u, dtype=self._state_dtype)
        return npa.concatenate([npa.atleast_1d(u), xd[:-1]])

    def _output(self, time, state, *inputs, **parameters):
        xd = state.discrete_state

        y = npa.sum(npa.dot(self.b_coefficients[1:], xd))

        if self.is_feedthrough:
            (u,) = inputs
            y += u * self.b_coefficients[0]

        return y

FiniteHorizonLinearQuadraticRegulator

Bases: LeafSystem

Finite Horizon Linear Quadratic Regulator (LQR) for a continuous-time system. Solves the Riccati Differential Equation (RDE) to compute the optimal control for the following finitie horizon cost function over [t0, tf]:

Minimise cost J:

J = [x(tf) - xd(tf)].T Qf [x(tf) - xd(tf)]
    + ∫[(x(t) - xd(t)].T Q [(x(t) - xd(t)] dt
    + ∫[(u(t) - ud(t)].T R [(u(t) - ud(t)] dt
    + 2 ∫[(x(t) - xd(t)].T N [(u(t) - ud(t)] dt

subject to the constraints:

dx(t)/dt - dx0(t)/dt = A [x(t)-x0(t)] + B [u(t)-u0(t)] - c(t),

where, x(t) is the state vector, u(t) is the control vector, xd(t) is the desired state vector, ud(t) is the desired control vector, x0(t) is the nominal state vector, u0(t) is the nominal control vector, Q, R, and N are the state, input, and cross cost matrices, Qf is the final state cost matrix,

and A, B, and c are computed from linearisation of the plant df/dx = f(x, u) around the nominal trajectory (x0(t), u0(t)).

A = df/dx(x0(t), u0(t), t)
B = df/du(x0(t), u0(t), t)
c = f(x0(t), u0(t), t) - dx0(t)/dt

The optimal control u obtained by the solution of the above problem is output.

See Section 8.5.1 of https://underactuated.csail.mit.edu/lqr.html#finite_horizon

Parameters:

Name	Type	Description	Default
t0		float Initial time of the finite horizon.	required
tf		float Final time of the finite horizon.	required
plant		a Plant object which can be a LeafSystem or a Diagram. The plant to be controlled. This represents df/dx = f(x, u).	required
Qf		Array Final state cost matrix.	required
func_Q		Callable A function that returns the state cost matrix Q at time t: func_Q(t)->Q	required
func_R		Callable A function that returns the input cost matrix R at time t: func_R(t)->R	required
func_N		Callable A function that returns the cross cost matrix N at time t. func_N(t)->N	required
func_x_0		Callable A function that returns the nominal state vector x0 at time t. func_x_0(t)->x0	required
func_u_0		Callable A function that returns the nominal control vector u0 at time t. func_u_0(t)->u0	required
func_x_d		Callable A function that returns the desired state vector xd at time t. func_x_d(t)->xd. If None, assumed to be the same as the nominal trajectory.	required
func_u_d		Callable A function that returns the desired control vector ud at time t. func_u_d(t)->ud. If None, assumed to be the same as the nominal trajectory.	required

Source code in jaxonomy/library/lqr.py

class FiniteHorizonLinearQuadraticRegulator(LeafSystem):
    """
    Finite Horizon Linear Quadratic Regulator (LQR) for a continuous-time system.
    Solves the Riccati Differential Equation (RDE) to compute the optimal control
    for the following finitie horizon cost function over [t0, tf]:

    Minimise cost J:

        J = [x(tf) - xd(tf)].T Qf [x(tf) - xd(tf)]
            + ∫[(x(t) - xd(t)].T Q [(x(t) - xd(t)] dt
            + ∫[(u(t) - ud(t)].T R [(u(t) - ud(t)] dt
            + 2 ∫[(x(t) - xd(t)].T N [(u(t) - ud(t)] dt

    subject to the constraints:

    dx(t)/dt - dx0(t)/dt = A [x(t)-x0(t)] + B [u(t)-u0(t)] - c(t),

    where,
        x(t) is the state vector,
        u(t) is the control vector,
        xd(t) is the desired state vector,
        ud(t) is the desired control vector,
        x0(t) is the nominal state vector,
        u0(t) is the nominal control vector,
        Q, R, and N are the state, input, and cross cost matrices,
        Qf is the final state cost matrix,

    and A, B, and c are computed from linearisation of the plant `df/dx = f(x, u)`
    around the nominal trajectory (x0(t), u0(t)).

        A = df/dx(x0(t), u0(t), t)
        B = df/du(x0(t), u0(t), t)
        c = f(x0(t), u0(t), t) - dx0(t)/dt

    The optimal control `u` obtained by the solution of the above problem is output.

    See Section 8.5.1 of https://underactuated.csail.mit.edu/lqr.html#finite_horizon

    Parameters:
        t0 : float
            Initial time of the finite horizon.
        tf : float
            Final time of the finite horizon.
        plant : a `Plant` object which can be a LeafSystem or a Diagram.
            The plant to be controlled. This represents `df/dx = f(x, u)`.
        Qf : Array
            Final state cost matrix.
        func_Q : Callable
            A function that returns the state cost matrix Q at time `t`: `func_Q(t)->Q`
        func_R : Callable
            A function that returns the input cost matrix R at time `t`: `func_R(t)->R`
        func_N : Callable
            A function that returns the cross cost matrix N at time `t`. `func_N(t)->N`
        func_x_0 : Callable
            A function that returns the nominal state vector `x0` at time `t`.
            func_x_0(t)->x0
        func_u_0 : Callable
            A function that returns the nominal control vector `u0` at time `t`.
            func_u_0(t)->u0
        func_x_d : Callable
            A function that returns the desired state vector `xd` at time `t`.
            func_x_d(t)->xd.  If None, assumed to be the same as the nominal trajectory.
        func_u_d : Callable
            A function that returns the desired control vector `ud` at time `t`.
            func_u_d(t)->ud.  If None, assumed to be the same as the nominal trajectory.
    """

    def __init__(
        self,
        t0,
        tf,
        plant,
        Qf,
        func_Q,
        func_R,
        func_N,
        func_x_0,
        func_u_0,
        func_x_d=None,
        func_u_d=None,
        name=None,
    ):
        super().__init__(name=name)

        self.t0 = t0
        self.tf = tf

        if func_x_d is None:
            func_x_d = func_x_0

        if func_u_d is None:
            func_u_d = func_u_0

        self.func_R = func_R
        self.func_N = func_N
        self.func_x_0 = func_x_0
        self.func_u_0 = func_u_0
        self.func_x_d = func_x_d
        self.func_u_d = func_u_d

        func_dot_x_0 = jax.jacfwd(func_x_0)
        nu = func_R(t0).shape[0]

        ode_rhs = make_ode_rhs(plant, nu)
        get_A = jax.jacfwd(ode_rhs, argnums=0)
        self.get_B = jax.jacfwd(ode_rhs, argnums=1)

        @jax.jit
        def rde(t, rde_state, args):
            t = -t
            Sxx, sx = rde_state

            Sxx = (Sxx + Sxx.T) / 2.0

            # Get nominal trajectories, desired trajectories, and cost matrices
            x_0 = func_x_0(t)
            u_0 = func_u_0(t)

            x_d = func_x_d(t)
            u_d = func_u_d(t)

            Q = func_Q(t)
            R = func_R(t)
            N = func_N(t)

            # Calculate dynamics mismatch due to nominal traj not satisfying dynamics
            dot_x_0 = func_dot_x_0(t)
            dot_x_0_eval = ode_rhs(x_0, u_0, t)
            c = dot_x_0_eval - dot_x_0

            #  Get linearisation around x_0, u_0
            A = get_A(x_0, u_0, t)
            B = self.get_B(x_0, u_0, t)

            #  Desired trajectories relative to nominal
            x_d_0 = x_d - x_0
            u_d_0 = u_d - u_0

            #  Compute RHS of RDE
            qx = -jnp.dot(Q, x_d_0) - jnp.dot(N, u_d_0)
            ru = -jnp.dot(R, u_d_0) - jnp.dot(N.T, x_d_0)

            N_plus_Sxx_B = N + jnp.matmul(Sxx, B)

            Rinv = jnp.linalg.inv(R)
            Sxx_A = jnp.matmul(Sxx, A)

            dot_Sxx = (
                Q
                - jnp.matmul(N_plus_Sxx_B, jnp.matmul(Rinv, N_plus_Sxx_B.T))
                + Sxx_A
                + Sxx_A.T
            )

            dot_sx = (
                qx
                - jnp.dot(N_plus_Sxx_B, jnp.dot(Rinv, ru + jnp.dot(B.T, sx)))
                + jnp.dot(A.T, sx)
                + jnp.dot(Sxx, c)
            )

            return (dot_Sxx, dot_sx)

        term = diffrax.ODETerm(rde)
        solver = diffrax.Tsit5()
        stepsize_controller = diffrax.PIDController(rtol=1e-5, atol=1e-5, dtmax=0.1)
        saveat = diffrax.SaveAt(dense=True)

        # TODO: Use utilities in ../simulation/ for reduced reliance on diffrax
        self.sol_rde = diffrax.diffeqsolve(
            term,
            solver,
            -tf,
            -t0,
            y0=(Qf, -jnp.dot(Qf, func_x_d(tf) - func_x_0(tf))),
            dt0=0.0001,
            saveat=saveat,
            stepsize_controller=stepsize_controller,
        )

        # Input: current state (x)
        self.declare_input_port()

        # Output port: Optimal finite horizon LQR control
        self.declare_output_port(self._eval_output, default_value=jnp.zeros(nu))

    def _eval_output(self, time, state, x, **params):
        rde_time = jnp.clip(time, self.t0, self.tf)
        rde_time = -rde_time

        Sxx, sx = self.sol_rde.evaluate(rde_time)

        x_d = self.func_x_d(time)
        u_d = self.func_u_d(time)

        x_0 = self.func_x_0(time)
        u_0 = self.func_u_0(time)

        x_d_0 = x_d - x_0
        u_d_0 = u_d - u_0

        B = self.get_B(x_0, u_0, time)

        R = self.func_R(time)
        N = self.func_N(time)
        Rinv = jnp.linalg.inv(R)

        ru = -jnp.dot(R, u_d_0) - jnp.dot(N.T, x_d_0)

        Rinv = jnp.linalg.inv(R)
        N_plus_Sxx_B = N + jnp.matmul(Sxx, B)

        u = (
            u_0
            - jnp.dot(Rinv, jnp.dot(N_plus_Sxx_B.T, (x - x_0)))
            - jnp.dot(Rinv, ru + jnp.dot(B.T, sx))
        )

        return u

FrequencyResponse dataclass

Result of a frequency-response evaluation.

Attributes:

Name	Type	Description
omegas	Any	Angular frequency vector ω in rad/s, shape (K,).
response	Any	Complex frequency response G(jω), shape (K, n_outputs, n_inputs).
magnitudes	Any	|G(jω)|, shape (K, n_outputs, n_inputs).
phases	Any	Phase arg G(jω) in radians, shape (K, n_outputs, n_inputs).

Source code in jaxonomy/library/linearization_workflow.py

@dataclass
class FrequencyResponse:
    """Result of a frequency-response evaluation.

    Attributes:
        omegas: Angular frequency vector ``ω`` in rad/s, shape ``(K,)``.
        response: Complex frequency response ``G(jω)``, shape
            ``(K, n_outputs, n_inputs)``.
        magnitudes: ``|G(jω)|``, shape ``(K, n_outputs, n_inputs)``.
        phases: Phase ``arg G(jω)`` in radians, shape
            ``(K, n_outputs, n_inputs)``.
    """

    omegas: Any  # jax.Array, shape (K,)
    response: Any  # complex jax.Array, shape (K, p, m)
    magnitudes: Any  # real jax.Array, shape (K, p, m)
    phases: Any  # real jax.Array, shape (K, p, m)

GPModel

Fitted Gaussian-process regressor (kriging).

Stores the training inputs, the pre-solved weight vector alpha = K^-1 y, and the Cholesky factor of the (noisy) covariance matrix for variance prediction. See Rasmussen & Williams 2006, Algorithm 2.1.

Source code in jaxonomy/library/rom/surrogates.py

class GPModel:
    """Fitted Gaussian-process regressor (kriging).

    Stores the training inputs, the pre-solved weight vector ``alpha = K^-1 y``,
    and the Cholesky factor of the (noisy) covariance matrix for variance
    prediction. See Rasmussen & Williams 2006, Algorithm 2.1.
    """

    def __init__(self, X, y, alpha, L, kernel, length_scale, signal_var,
                 noise, matern_nu):
        self.X_train = X
        self.y_train = y
        self.alpha = alpha
        self.L = L
        self.kernel = kernel
        self.length_scale = float(length_scale)
        self.signal_var = float(signal_var)
        self.noise = float(noise)
        self.matern_nu = float(matern_nu)

    def predict(self, Xstar):
        """Posterior ``(mean, variance)`` at query points ``Xstar``.

        jax-traceable. ``Xstar`` may be 1-D (single point / single feature) or
        2-D ``(m, d)``. Returns arrays of shape ``(m,)``.
        """
        Xstar = _as2d(Xstar)
        Ks = _gp_kernel(Xstar, self.X_train, self.kernel, self.length_scale,
                        self.signal_var, self.matern_nu)  # (m, n)
        mean = Ks @ self.alpha
        v = jsla.solve_triangular(self.L, Ks.T, lower=True)  # (n, m)
        var = self.signal_var - jnp.sum(v * v, axis=0)
        var = jnp.clip(var, 0.0, None)
        return mean, var

    def log_marginal_likelihood(self):
        n = self.X_train.shape[0]
        data_fit = -0.5 * jnp.dot(self.y_train, self.alpha)
        complexity = -jnp.sum(jnp.log(jnp.diag(self.L)))
        return data_fit + complexity - 0.5 * n * math.log(2.0 * math.pi)

predict(Xstar)

Posterior (mean, variance) at query points Xstar.

jax-traceable. Xstar may be 1-D (single point / single feature) or 2-D (m, d). Returns arrays of shape (m,).

Source code in jaxonomy/library/rom/surrogates.py

def predict(self, Xstar):
    """Posterior ``(mean, variance)`` at query points ``Xstar``.

    jax-traceable. ``Xstar`` may be 1-D (single point / single feature) or
    2-D ``(m, d)``. Returns arrays of shape ``(m,)``.
    """
    Xstar = _as2d(Xstar)
    Ks = _gp_kernel(Xstar, self.X_train, self.kernel, self.length_scale,
                    self.signal_var, self.matern_nu)  # (m, n)
    mean = Ks @ self.alpha
    v = jsla.solve_triangular(self.L, Ks.T, lower=True)  # (n, m)
    var = self.signal_var - jnp.sum(v * v, axis=0)
    var = jnp.clip(var, 0.0, None)
    return mean, var

Gain

Bases: FeedthroughBlock

Multiply the input signal by a constant value.

Input ports

(0) The input signal.

Output ports

(0) The input signal multiplied by the gain: y = gain * u.

Parameters:

Name	Type	Description	Default
gain		The value to scale the input signal by.	required
dtype	optional, T-038a-followup-other-blocks	If set, the block's output is cast to this dtype. See LookupTable1d for the per-block dtype contract.	None

Source code in jaxonomy/library/math_ops.py

class Gain(FeedthroughBlock):
    """Multiply the input signal by a constant value.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The input signal multiplied by the gain: `y = gain * u`.

    Parameters:
        gain:
            The value to scale the input signal by.
        dtype (optional, T-038a-followup-other-blocks):
            If set, the block's output is cast to this dtype.  See
            ``LookupTable1d`` for the per-block dtype contract.
    """

    @parameters(dynamic=["gain"])
    def __init__(self, gain, *args, dtype=None, **kwargs):
        # T-038a-followup-other-blocks: dtype is stored outside the
        # @parameters dynamic list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        if dtype is None:
            super().__init__(lambda x, gain: gain * x, *args, **kwargs)
        else:
            _dtype = dtype

            def _gain_op(x, gain):
                return npa.asarray(gain * x).astype(_dtype)

            super().__init__(_gain_op, *args, **kwargs)

    def initialize(self, gain):
        pass

GaussianProcess

Bases: LeafSystem

Gaussian-process (kriging) surrogate as a feedthrough block.

Input port 0 is the feature vector u; output port 0 is the predictive mean and output port 1 the predictive variance. Training data and the Cholesky factor are stored statically on the block; the weight vector alpha and the kernel hyperparameters are dynamic parameters so the surrogate is differentiable in them (Rasmussen & Williams 2006).

Source code in jaxonomy/library/rom/surrogates.py

class GaussianProcess(LeafSystem):
    """Gaussian-process (kriging) surrogate as a feedthrough block.

    Input port 0 is the feature vector ``u``; output port 0 is the predictive
    mean and output port 1 the predictive variance. Training data and the
    Cholesky factor are stored statically on the block; the weight vector
    ``alpha`` and the kernel hyperparameters are dynamic parameters so the
    surrogate is differentiable in them (Rasmussen & Williams 2006).
    """

    def __init__(self, model: GPModel, name=None, **kwargs):
        super().__init__(name=name, **kwargs)
        self.model = model
        self.declare_input_port()
        self.declare_dynamic_parameter("alpha", jnp.asarray(model.alpha))
        self.declare_dynamic_parameter(
            "length_scale", jnp.asarray(model.length_scale, dtype=jnp.float64))
        self.declare_dynamic_parameter(
            "signal_var", jnp.asarray(model.signal_var, dtype=jnp.float64))

        self._mean_port_idx = self.declare_output_port(
            self._mean, name="mean",
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )
        self._var_port_idx = self.declare_output_port(
            self._variance, name="variance",
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )

    def _mean(self, time, state, *inputs, **params):
        Xstar = _row(inputs[0])
        Ks = _gp_kernel(Xstar, self.model.X_train, self.model.kernel,
                        params["length_scale"], params["signal_var"],
                        self.model.matern_nu)
        return (Ks @ params["alpha"])[0]

    def _variance(self, time, state, *inputs, **params):
        Xstar = _row(inputs[0])
        Ks = _gp_kernel(Xstar, self.model.X_train, self.model.kernel,
                        params["length_scale"], params["signal_var"],
                        self.model.matern_nu)
        v = jsla.solve_triangular(self.model.L, Ks.T, lower=True)
        var = params["signal_var"] - jnp.sum(v * v, axis=0)
        return jnp.clip(var, 0.0, None)[0]

HermiteSimpsonNMPC

Bases: NonlinearMPCIpopt

Implementation of nonlinear MPC with Hermite-Simpson collocation and IPOPT as the NLP solver.

Input ports

(0) x_0 : current state vector. (1) x_ref : reference state trajectory for the nonlinear MPC. (2) u_ref : reference input trajectory for the nonlinear MPC.

Output ports

(1) u_opt : the optimal control input to be applied at the current time step as determined by the nonlinear MPC.

Parameters:

Name	Type	Description	Default
plant		LeafSystem or Diagram The plant to be controlled.	required
Q		Array State weighting matrix in the cost function.	required
QN		Array Terminal state weighting matrix in the cost function.	required
R		Array Control input weighting matrix in the cost function.	required
N		int The prediction horizon, an integer specifying the number of steps to predict. Note: prediction and control horizons are identical for now.	required
dt		float: Major time step, a scalar indicating the increment in time for each step in the prediction and control horizons.	required
lb_x		Array Lower bound on the state vector.	None
ub_x		Array Upper bound on the state vector.	None
lb_u		Array Lower bound on the control input vector.	None
ub_u		Array Upper bound on the control input vector.	None
include_terminal_x_as_constraint		bool If True, the terminal state is included as a constraint in the NLP.	False
include_terminal_u_as_constraint		bool If True, the terminal control input is included as a constraint in the NLP.	False
x_optvars_0		Array Initial guess for the state vector optimization variables in the NLP.	None
u_optvars_0		Array Initial guess for the control vector optimization variables in the NLP.	None

Source code in jaxonomy/library/nmpc/hermite_simpson_ipopt_nmpc.py

class HermiteSimpsonNMPC(NonlinearMPCIpopt):
    """
    Implementation of nonlinear MPC with Hermite-Simpson collocation and IPOPT as the
    NLP solver.

    Input ports:
        (0) x_0 : current state vector.
        (1) x_ref : reference state trajectory for the nonlinear MPC.
        (2) u_ref : reference input trajectory for the nonlinear MPC.

    Output ports:
        (1) u_opt : the optimal control input to be applied at the current time step
                    as determined by the nonlinear MPC.

    Parameters:
        plant: LeafSystem or Diagram
            The plant to be controlled.

        Q: Array
            State weighting matrix in the cost function.

        QN: Array
            Terminal state weighting matrix in the cost function.

        R: Array
            Control input weighting matrix in the cost function.

        N: int
            The prediction horizon, an integer specifying the number of steps to
            predict. Note: prediction and control horizons are identical for now.

        dt: float:
            Major time step, a scalar indicating the increment in time for each step in
            the prediction and control horizons.

        lb_x: Array
            Lower bound on the state vector.

        ub_x: Array
            Upper bound on the state vector.

        lb_u: Array
            Lower bound on the control input vector.

        ub_u: Array
            Upper bound on the control input vector.

        include_terminal_x_as_constraint: bool
            If True, the terminal state is included as a constraint in the NLP.

        include_terminal_u_as_constraint: bool
            If True, the terminal control input is included as a constraint in the NLP.

        x_optvars_0: Array
            Initial guess for the state vector optimization variables in the NLP.

        u_optvars_0: Array
            Initial guess for the control vector optimization variables in the NLP.
    """

    def __init__(
        self,
        plant,
        Q,
        QN,
        R,
        N,
        dt,
        lb_x=None,
        ub_x=None,
        lb_u=None,
        ub_u=None,
        include_terminal_x_as_constraint=False,
        include_terminal_u_as_constraint=False,
        x_optvars_0=None,
        u_optvars_0=None,
        name=None,
        warm_start=False,
    ):
        self.Q = Q
        self.QN = QN
        self.R = R

        self.N = N
        self.dt = dt

        self.lb_x = lb_x
        self.ub_x = ub_x
        self.lb_u = lb_u
        self.ub_u = ub_u

        self.include_terminal_x_as_constraint = include_terminal_x_as_constraint
        self.include_terminal_u_as_constraint = include_terminal_u_as_constraint

        self.nx = Q.shape[0]
        self.nu = R.shape[0]

        if lb_x is None:
            self.lb_x = -1e20 * jnp.ones(self.nx)

        if ub_x is None:
            self.ub_x = 1e20 * jnp.ones(self.nx)

        if lb_u is None:
            self.lb_u = -1e20 * jnp.ones(self.nu)

        if ub_u is None:
            self.ub_u = 1e20 * jnp.ones(self.nu)

        # Currently guesses are not taken into account
        self.x_optvars_0 = x_optvars_0  # Currently does nothing
        self.u_optvars_0 = u_optvars_0  # Currently does nothing
        if x_optvars_0 is None:
            x_optvars_0 = jnp.zeros((N + 1, self.nx))
        if u_optvars_0 is None:
            u_optvars_0 = jnp.zeros((N + 1, self.nu))

        self.ode_rhs = make_ode_rhs(plant, self.nu)

        nlp_structure_ipopt = NMPCProblemStructure(
            self.num_optvars,
            self._objective,
            self._constraints,
        )

        super().__init__(
            dt,
            self.nu,
            self.num_optvars,
            nlp_structure_ipopt,
            name=name,
            warm_start=warm_start,
        )

    @property
    def num_optvars(self):
        return (self.N + 1) * (self.nx + self.nu)

    @property
    def num_constraints(self):
        # max size regardless of terminal constraints (for jit compilation)
        num_contraints = (self.N + 2) * self.nx + self.nu
        return num_contraints

    @property
    def bounds_optvars(self):
        lb = jnp.hstack(
            [jnp.tile(self.lb_u, self.N + 1), jnp.tile(self.lb_x, self.N + 1)]
        )
        ub = jnp.hstack(
            [jnp.tile(self.ub_u, self.N + 1), jnp.tile(self.ub_x, self.N + 1)]
        )
        return (lb, ub)

    @property
    def bounds_constraints(self):
        c_lb = jnp.zeros(self.num_constraints)
        c_ub = jnp.zeros(self.num_constraints)
        return (c_lb, c_ub)

    @partial(jax.jit, static_argnames=("self",))
    def _objective(self, optvars, t0, x0, x_ref, u_ref):
        u_and_x_flat = optvars

        u = u_and_x_flat[: self.nu * (self.N + 1)].reshape((self.N + 1, self.nu))
        x = u_and_x_flat[self.nu * (self.N + 1) :].reshape((self.N + 1, self.nx))

        xdiff = x - x_ref
        udiff = u - u_ref

        # compute sum of quadratic products for x_0 to x_{n-1}
        A = jnp.dot(xdiff[:-1], self.Q)
        qp_x_sum = jnp.sum(xdiff[:-1] * A, axis=None)

        # Compute quadratic product for the x_N
        xN = xdiff[-1]
        qp_x_N = jnp.dot(xN, jnp.dot(self.QN, xN))

        # compute sum of quadratic products for u_0 to u_{n-1}
        B = jnp.dot(udiff, self.R)
        qp_u_sum = jnp.sum(udiff * B, axis=None)

        # Sum the quadratic products
        total_sum = qp_x_sum + qp_x_N + qp_u_sum
        return total_sum

    @partial(jax.jit, static_argnames=("self",))
    def _constraints(self, optvars, t0, x0, x_ref, u_ref):
        u_and_x_flat = optvars

        u = u_and_x_flat[: self.nu * (self.N + 1)].reshape((self.N + 1, self.nu))
        x = u_and_x_flat[self.nu * (self.N + 1) :].reshape((self.N + 1, self.nx))

        h = self.dt
        t = t0 + h * jnp.arange(self.N + 1)

        dot_x = jnp.zeros((self.N + 1, self.nx))

        def loop_body_break(idx, dot_x):
            rhs = self.ode_rhs(x[idx], u[idx], t[idx])
            dot_x = dot_x.at[idx].set(rhs)
            return dot_x

        dot_x = jax.lax.fori_loop(0, self.N + 1, loop_body_break, dot_x)

        t = t0 + self.dt * jnp.arange(self.N + 1)
        t_c = 0.5 * (t[:-1] + t[1:])
        u_c = 0.5 * (u[:-1] + u[1:])
        x_c = 0.5 * (x[:-1] + x[1:]) + (h / 8.0) * (dot_x[:-1] - dot_x[1:])

        dot_x_c = (-3.0 / 2.0 / h) * (x[:-1] - x[1:]) - (1.0 / 4.0) * (
            dot_x[:-1] + dot_x[1:]
        )

        c0 = x0 - x[0]

        c_others = jnp.zeros((self.N, self.nx))

        def loop_body_colloc(idx, c_others):
            c_colocation = self.ode_rhs(x_c[idx], u_c[idx], t_c[idx]) - dot_x_c[idx]
            c_others = c_others.at[idx].set(c_colocation)
            return c_others

        c_others = jax.lax.fori_loop(0, self.N, loop_body_colloc, c_others)
        c_all = jnp.hstack([c0.ravel(), c_others.ravel()])

        c_terminal_x = x_ref[self.N] - x[self.N]
        c_terminal_u = u_ref[self.N] - u[self.N]

        c_all = cond(
            self.include_terminal_x_as_constraint,
            lambda c_all, c_terminal_x: jnp.hstack([c_all, c_terminal_x.ravel()]),
            lambda c_all, c_terminal_x: jnp.hstack([c_all, jnp.zeros(self.nx)]),
            c_all,
            c_terminal_x,
        )

        c_all = cond(
            self.include_terminal_u_as_constraint,
            lambda c_all, c_terminal_u: jnp.hstack([c_all, c_terminal_u.ravel()]),
            lambda c_all, c_terminal_x: jnp.hstack([c_all, jnp.zeros(self.nu)]),
            c_all,
            c_terminal_u,
        )

        return c_all

IOPort

Bases: FeedthroughBlock

Simple class for organizing input/output ports for groups/submodels.

Since these are treated as standalone blocks in the UI rather than specific input/output ports exported to the parent model, it is more straightforward to represent them that way here as well.

This class represents a simple one-input, one-output feedthrough block where the feedthrough function is an identity. The input (resp. output) port can then be exported to the parent model to create an Inport (resp. Outport).

Source code in jaxonomy/library/routing.py

class IOPort(FeedthroughBlock):
    """Simple class for organizing input/output ports for groups/submodels.

    Since these are treated as standalone blocks in the UI rather than specific
    input/output ports exported to the parent model, it is more straightforward
    to represent them that way here as well.

    This class represents a simple one-input, one-output feedthrough block where
    the feedthrough function is an identity.  The input (resp. output) port can then
    be exported to the parent model to create an Inport (resp. Outport).
    """

    def __init__(self, *args, **kwargs):
        super().__init__(lambda x: x, *args, **kwargs)

IfThenElse

Bases: LeafSystem

Applies a conditional expression to the input signals.

Given inputs pred, true_val, and false_val, the block computes:

y = true_val if pred else false_val

The true and false values may be any arrays, but must have the same shape and dtype.

Input ports

(0) The boolean predicate. (1) The true value. (2) The false value.

Output ports

(0) The result of the conditional expression. Shape and dtype will match the true and false values.

Events

An event is triggered when the output changes from true to false or vice versa.

Source code in jaxonomy/library/logic.py

class IfThenElse(LeafSystem):
    """Applies a conditional expression to the input signals.

    Given inputs `pred`, `true_val`, and `false_val`, the block computes:
    ```
    y = true_val if pred else false_val
    ```

    The true and false values may be any arrays, but must have the same
    shape and dtype.

    Input ports:
        (0) The boolean predicate.
        (1) The true value.
        (2) The false value.

    Output ports:
        (0) The result of the conditional expression. Shape and dtype will match
            the true and false values.

    Events:
        An event is triggered when the output changes from true to false or vice versa.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.declare_input_port()  # pred
        self.declare_input_port()  # true_val
        self.declare_input_port()  # false_val

        def _compute_output(_time, _state, *inputs, **_params):
            return npa.where(inputs[0], inputs[1], inputs[2])

        self.declare_output_port(
            _compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    def _edge_detection(self, _time, _state, *inputs, **_params):
        return npa.where(inputs[0], 1.0, -1.0)

    def initialize_static_data(self, context):
        # Add a zero-crossing event so ODE solvers can't try to integrate
        # through a discontinuity. For efficiency, only do this if the output is
        # fed to an ODE.
        if not self.has_zero_crossing_events and is_discontinuity(self.output_ports[0]):
            self.declare_zero_crossing(self._edge_detection, direction="crosses_zero")

        return super().initialize_static_data(context)

InfiniteHorizonKalmanFilter

Bases: KalmanFilterBase

Infinite Horizon Kalman Filter for the following system:

x[n+1] = A x[n] + B u[n] + G w[n]
y[n]   = C x[n] + D u[n] + v[n]

E(w[n]) = E(v[n]) = 0
E(w[n]w'[n]) = Q
E(v[n]v'[n]) = R
E(w[n]v'[n]) = N = 0

Input ports

(0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n

Output ports

(1) x_hat[n] : state vector estimate at timestep n

Parameters:

Name	Type	Description	Default
dt		float Time step of the discrete-time system	required
A		ndarray State transition matrix	required
B		ndarray Input matrix	required
C		ndarray Output matrix. If None, full state output is assumed.	None
D		ndarray Feedthrough matrix. If None, no feedthrough is assumed.	None
G		ndarray Process noise matrix. If None, G=B is assumed.	None
Q		ndarray Process noise covariance matrix. If None, Identity matrix of size compatible with G and A is assumed.	None
R		ndarray Measurement noise covariance matrix. If None, Identity matrix of size compatible with C and A is assumed.	None
x_hat_0		ndarray Initial state estimate. If None, an array of zeros is assumed.	None

Source code in jaxonomy/library/state_estimators/infinite_horizon_kalman_filter.py

class InfiniteHorizonKalmanFilter(KalmanFilterBase):
    """
    Infinite Horizon Kalman Filter for the following system:

    ```
    x[n+1] = A x[n] + B u[n] + G w[n]
    y[n]   = C x[n] + D u[n] + v[n]

    E(w[n]) = E(v[n]) = 0
    E(w[n]w'[n]) = Q
    E(v[n]v'[n]) = R
    E(w[n]v'[n]) = N = 0
    ```

    Input ports:
        (0) u[n] : control vector at timestep n
        (1) y[n] : measurement vector at timestep n

    Output ports:
        (1) x_hat[n] : state vector estimate at timestep n

    Parameters:
        dt: float
            Time step of the discrete-time system
        A: ndarray
            State transition matrix
        B: ndarray
            Input matrix
        C: ndarray
            Output matrix. If `None`, full state output is assumed.
        D: ndarray
            Feedthrough matrix. If `None`, no feedthrough is assumed.
        G: ndarray
            Process noise matrix. If `None`, `G=B` is assumed.
        Q: ndarray
            Process noise covariance matrix. If `None`, Identity matrix of size
            compatible with `G` and `A` is assumed.
        R: ndarray
            Measurement noise covariance matrix. If `None`, Identity matrix of size
            compatible with `C` and `A` is assumed.
        x_hat_0: ndarray
            Initial state estimate. If `None`, an array of zeros is assumed.
    """

    @parameters(
        static=["dt", "A", "B", "C", "D", "G", "Q", "R", "x_hat_0"],
    )
    def __init__(
        self,
        dt,
        A,
        B,
        C=None,
        D=None,
        G=None,
        Q=None,
        R=None,
        x_hat_0=None,
        name=None,
        **kwargs,
    ):
        self.nx = 0
        self.nu = 0
        self.ny = 0
        self.nd = 0
        self.A = None
        self.B = None
        self.C = None
        self.D = None
        self.G = None
        self.Q = None
        self.R = None
        self.K = None
        self.A_minus_LC = None
        self.B_minus_LD = None
        self.L = None

        # Note: This class inherits from KalmanFilterBase. Since the infinite horizon
        # kalman filter does not need P_hat and track it, a dummy_P_hat_0 is set as
        # Identity matrix of size 1, and used wherever KalmanFilterBase demands
        # P_hat-like matrices
        self.dummy_P_hat_0 = npa.eye(1)
        is_feedthrough = False if D is None else bool(not npa.allclose(D, 0.0))
        super().__init__(
            dt, x_hat_0, self.dummy_P_hat_0, is_feedthrough, name, **kwargs
        )

    def initialize(
        self, dt, A, B, C=None, D=None, G=None, Q=None, R=None, x_hat_0=None
    ):
        self.nx, self.nu = B.shape

        if C is None:
            C = npa.eye(self.nx)
            self.ny = self.nx
        else:
            self.ny = C.shape[0]

        if D is None:
            D = npa.zeros((self.ny, self.nu))
        self.is_feedthrough = bool(not npa.allclose(D, 0.0))

        if G is None:
            G = B

        _, self.nd = G.shape

        if Q is None:
            Q = npa.eye(self.nd)

        if R is None:
            R = npa.eye(self.ny)

        if x_hat_0 is None:
            x_hat_0 = npa.zeros(self.nx)

        check_shape_compatibilities(A, B, C, D, G, Q, R)

        self.A = A
        self.B = B
        self.C = C
        self.D = D
        self.G = G
        self.Q = Q
        self.R = R

        L, P, E = control.dlqe(A, G, C, Q, R)

        self.K = np.linalg.solve(A, L)

        self.A_minus_LC = A - np.matmul(L, C)
        self.B_minus_LD = B - np.matmul(L, D)
        self.L = L

    def _correct(self, time, x_hat_minus, P_hat_minus, *inputs):
        u, y = inputs
        y = npa.atleast_1d(y)

        C, D = self.C, self.D

        x_hat_plus = x_hat_minus + npa.dot(self.K, y - npa.dot(C, x_hat_minus))  # n|n

        if self.is_feedthrough:
            u = npa.atleast_1d(u)
            x_hat_plus = x_hat_plus - npa.dot(self.K, npa.dot(D, u))

        return x_hat_plus, self.dummy_P_hat_0

    def _propagate(self, time, x_hat_plus, P_hat_plus, *inputs):
        u, y = inputs
        u = npa.atleast_1d(u)

        x_hat_minus = (
            npa.dot(self.A_minus_LC, x_hat_plus)
            + npa.dot(self.B_minus_LD, u)
            + npa.dot(self.L, y)
        )  # n+1|n

        return x_hat_minus, self.dummy_P_hat_0

    #######################################
    # Make filter for a continuous plant  #
    #######################################

    @staticmethod
    @with_resolved_parameters
    def for_continuous_plant(
        plant,
        x_eq,
        u_eq,
        dt,
        Q=None,
        R=None,
        G=None,
        x_hat_bar_0=None,
        discretization_method="zoh",
        discretized_noise=False,
        name=None,
        ui_id=None,
    ):
        """
        Obtain an Infinite Horizon Kalman Filter system for a continuous-time plant
        after linearization at equilibrium point (x_eq, u_eq)

        The input plant contains the deterministic forms of the forward and observation
        operators:

        ```
            dx/dt = f(x,u)
            y = g(x,u)
        ```

        Note: (i) Only plants with one vector-valued input and one vector-valued output
        are currently supported. Furthermore, the plant LeafSystem/Diagram should have
        only one vector-valued integrator. (ii) the user may pass a plant with
        disturbances as the input plant. However, computation of `y_eq` will be fraught
        with disturbances.

        A plant with disturbances of the following form is then considered
        following form:

        ```
            dx/dt = f(x,u) + G w                        --- (C1)
            y = g(x,u) +  v                             --- (C2)
        ```

        where:

            `w` represents the process noise,
            `v` represents the measurement noise,

        and

        ```
            E(w) = E(v) = 0
            E(ww') = Q
            E(vv') = R
            E(wv') = N = 0
        ```

        This plant with disturbances is linearized (only `f` and `g`) around the
        equilibrium point to obtain:

        ```
            d/dt (x_bar) = A x_bar + B u_bar + G w
            y_bar = C x_bar + D u_bar + v
        ```

        where,

        ```
            x_bar = x - x_eq
            u_bar = u - u_eq
            y_bar = y - y_bar
            y_eq = g(x_eq, u_eq)
        ```

        The linearized plant is then discretized via `euler` or `zoh` method to obtain:

        ```
            x_bar[n] = Ad x_bar[n] + Bd u_bar[n] + Gd w[n]           --- (L1)
            y_bar[n] = Cd x_bar[n] + Dd u_bar[n] + v[n]              --- (L2)

            E(w[n]) = E(v[n]) = 0
            E(w[n]w'[n]) = Qd
            E(v[n]v'[n]) = Rd
            E(w[n]v'[n]) = Nd = 0
        ```

        Note: If `discretized_noise` is True, then it is assumed that the user is
        providing Gd, Qd and Rd. If False, then Qd and Rd are computed from
        continuous-time Q, R, and G, and Gd is set to Identity matrix.

        An Infinite Horizon Kalman Filter estimator for the system of equations (L1)
        and (L2) is returned. This filter is in the `x_bar`, `u_bar`, and `y_bar`
        states.

        This returned system will have

        Input ports:
            (0) u_bar[n] : control vector at timestep n, relative to equilibrium
            (1) y_bar[n] : measurement vector at timestep n, relative to equilibrium

        Output ports:
            (1) x_hat_bar[n] : state vector estimate at timestep n, relative to
                               equilibrium

        Parameters:
            plant : a `Plant` object which can be a LeafSystem or a Diagram.
            x_eq: ndarray
                Equilibrium state vector for discretization
            u_eq: ndarray
                Equilibrium control vector for discretization
            dt: float
                Time step for the discretization.
            Q: ndarray
                Process noise covariance matrix. If `None`, Identity matrix of size
                compatible with `G` and and linearized system's `A` is assumed.
            R: ndarray
                Measurement noise covariance matrix. If `None`, Identity matrix of size
                compatible with linearized system's `C` and `A` is assumed.
            G: ndarray
                Process noise matrix. If `None`, `G=B` is assumed making disrurbances
                additive to control vector `u`, i.e. `u_disturbed = u_orig + w`.
            x_hat_bar_0: ndarray
                Initial state estimate relative to equilibrium.
                If None, an identity matrix is assumed.
            discretization_method: str ("euler" or "zoh")
                Method to discretize the continuous-time plant. Default is "euler".
            discretized_noise: bool
                Whether the user is directly providing Gd, Qd and Rd. Default is False.
                If True, `G`, `Q`, and `R` are assumed to be Gd, Qd, and Rd,
                respectively.
        """
        (
            y_eq,
            Ad,
            Bd,
            Cd,
            Dd,
            Gd,
            Qd,
            Rd,
        ) = linearize_and_discretize_continuous_plant(
            plant, x_eq, u_eq, dt, Q, R, G, discretization_method, discretized_noise
        )

        check_shape_compatibilities(Ad, Bd, Cd, Dd, Gd, Qd, Rd)

        nx = x_eq.size

        if x_hat_bar_0 is None:
            x_hat_bar_0 = npa.zeros(nx)

        # Instantiate an Infinite Horizon Kalman Filter for the linearized plant
        kf = InfiniteHorizonKalmanFilter(
            dt,
            Ad,
            Bd,
            Cd,
            Dd,
            Gd,
            Qd,
            Rd,
            x_hat_bar_0,
            name=name,
            ui_id=ui_id,
        )

        return y_eq, kf

    ##############################################
    # Make global filter for a continuous plant  #
    ##############################################

    @staticmethod
    @with_resolved_parameters
    def global_filter_for_continuous_plant(
        plant,
        x_eq,
        u_eq,
        dt,
        Q=None,
        R=None,
        G=None,
        x_hat_0=None,
        discretization_method="euler",
        discretized_noise=False,
        name=None,
        ui_id=None,
    ):
        """
        See docs for `for_continuous_plant`, which returns the local infinite horizon
        Kalman Filter. This method additionally converts the local Kalman Filter to a
        global estimator. See docs for `make_global_estimator_from_local` for details.
        """
        (
            y_eq,
            Ad,
            Bd,
            Cd,
            Dd,
            Gd,
            Qd,
            Rd,
        ) = linearize_and_discretize_continuous_plant(
            plant, x_eq, u_eq, dt, Q, R, G, discretization_method, discretized_noise
        )

        check_shape_compatibilities(Ad, Bd, Cd, Dd, Gd, Qd, Rd)

        nx = x_eq.size

        if x_hat_0 is None:
            x_hat_bar_0 = npa.zeros(nx)
        else:
            x_hat_bar_0 = x_hat_0 - x_eq

        # Instantiate an Infinite Horizon Kalman Filter for the linearized plant
        local_kf = InfiniteHorizonKalmanFilter(
            dt,
            Ad,
            Bd,
            Cd,
            Dd,
            Gd,
            Qd,
            Rd,
            x_hat_bar_0,
            name=name,
            ui_id=ui_id,
        )

        global_kf = make_global_estimator_from_local(
            local_kf,
            x_eq,
            u_eq,
            y_eq,
            name=name,
            ui_id=ui_id,
        )

        return global_kf

for_continuous_plant(plant, x_eq, u_eq, dt, Q=None, R=None, G=None, x_hat_bar_0=None, discretization_method='zoh', discretized_noise=False, name=None, ui_id=None) staticmethod

Obtain an Infinite Horizon Kalman Filter system for a continuous-time plant after linearization at equilibrium point (x_eq, u_eq)

The input plant contains the deterministic forms of the forward and observation operators:

    dx/dt = f(x,u)
    y = g(x,u)

Note: (i) Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator. (ii) the user may pass a plant with disturbances as the input plant. However, computation of y_eq will be fraught with disturbances.

A plant with disturbances of the following form is then considered following form:

    dx/dt = f(x,u) + G w                        --- (C1)
    y = g(x,u) +  v                             --- (C2)

where:

`w` represents the process noise,
`v` represents the measurement noise,

and

    E(w) = E(v) = 0
    E(ww') = Q
    E(vv') = R
    E(wv') = N = 0

This plant with disturbances is linearized (only f and g) around the equilibrium point to obtain:

    d/dt (x_bar) = A x_bar + B u_bar + G w
    y_bar = C x_bar + D u_bar + v

where,

    x_bar = x - x_eq
    u_bar = u - u_eq
    y_bar = y - y_bar
    y_eq = g(x_eq, u_eq)

The linearized plant is then discretized via euler or zoh method to obtain:

    x_bar[n] = Ad x_bar[n] + Bd u_bar[n] + Gd w[n]           --- (L1)
    y_bar[n] = Cd x_bar[n] + Dd u_bar[n] + v[n]              --- (L2)

    E(w[n]) = E(v[n]) = 0
    E(w[n]w'[n]) = Qd
    E(v[n]v'[n]) = Rd
    E(w[n]v'[n]) = Nd = 0

Note: If discretized_noise is True, then it is assumed that the user is providing Gd, Qd and Rd. If False, then Qd and Rd are computed from continuous-time Q, R, and G, and Gd is set to Identity matrix.

An Infinite Horizon Kalman Filter estimator for the system of equations (L1) and (L2) is returned. This filter is in the x_bar, u_bar, and y_bar states.

This returned system will have

Input ports

(0) u_bar[n] : control vector at timestep n, relative to equilibrium (1) y_bar[n] : measurement vector at timestep n, relative to equilibrium

Output ports

(1) x_hat_bar[n] : state vector estimate at timestep n, relative to equilibrium

Parameters:

Name	Type	Description	Default
plant		a Plant object which can be a LeafSystem or a Diagram.	required
x_eq		ndarray Equilibrium state vector for discretization	required
u_eq		ndarray Equilibrium control vector for discretization	required
dt		float Time step for the discretization.	required
Q		ndarray Process noise covariance matrix. If None, Identity matrix of size compatible with G and and linearized system's A is assumed.	None
R		ndarray Measurement noise covariance matrix. If None, Identity matrix of size compatible with linearized system's C and A is assumed.	None
G		ndarray Process noise matrix. If None, G=B is assumed making disrurbances additive to control vector u, i.e. u_disturbed = u_orig + w.	None
x_hat_bar_0		ndarray Initial state estimate relative to equilibrium. If None, an identity matrix is assumed.	None
discretization_method		str ("euler" or "zoh") Method to discretize the continuous-time plant. Default is "euler".	'zoh'
discretized_noise		bool Whether the user is directly providing Gd, Qd and Rd. Default is False. If True, G, Q, and R are assumed to be Gd, Qd, and Rd, respectively.	False

Source code in jaxonomy/library/state_estimators/infinite_horizon_kalman_filter.py

@staticmethod
@with_resolved_parameters
def for_continuous_plant(
    plant,
    x_eq,
    u_eq,
    dt,
    Q=None,
    R=None,
    G=None,
    x_hat_bar_0=None,
    discretization_method="zoh",
    discretized_noise=False,
    name=None,
    ui_id=None,
):
    """
    Obtain an Infinite Horizon Kalman Filter system for a continuous-time plant
    after linearization at equilibrium point (x_eq, u_eq)

    The input plant contains the deterministic forms of the forward and observation
    operators:

    ```
        dx/dt = f(x,u)
        y = g(x,u)
    ```

    Note: (i) Only plants with one vector-valued input and one vector-valued output
    are currently supported. Furthermore, the plant LeafSystem/Diagram should have
    only one vector-valued integrator. (ii) the user may pass a plant with
    disturbances as the input plant. However, computation of `y_eq` will be fraught
    with disturbances.

    A plant with disturbances of the following form is then considered
    following form:

    ```
        dx/dt = f(x,u) + G w                        --- (C1)
        y = g(x,u) +  v                             --- (C2)
    ```

    where:

        `w` represents the process noise,
        `v` represents the measurement noise,

    and

    ```
        E(w) = E(v) = 0
        E(ww') = Q
        E(vv') = R
        E(wv') = N = 0
    ```

    This plant with disturbances is linearized (only `f` and `g`) around the
    equilibrium point to obtain:

    ```
        d/dt (x_bar) = A x_bar + B u_bar + G w
        y_bar = C x_bar + D u_bar + v
    ```

    where,

    ```
        x_bar = x - x_eq
        u_bar = u - u_eq
        y_bar = y - y_bar
        y_eq = g(x_eq, u_eq)
    ```

    The linearized plant is then discretized via `euler` or `zoh` method to obtain:

    ```
        x_bar[n] = Ad x_bar[n] + Bd u_bar[n] + Gd w[n]           --- (L1)
        y_bar[n] = Cd x_bar[n] + Dd u_bar[n] + v[n]              --- (L2)

        E(w[n]) = E(v[n]) = 0
        E(w[n]w'[n]) = Qd
        E(v[n]v'[n]) = Rd
        E(w[n]v'[n]) = Nd = 0
    ```

    Note: If `discretized_noise` is True, then it is assumed that the user is
    providing Gd, Qd and Rd. If False, then Qd and Rd are computed from
    continuous-time Q, R, and G, and Gd is set to Identity matrix.

    An Infinite Horizon Kalman Filter estimator for the system of equations (L1)
    and (L2) is returned. This filter is in the `x_bar`, `u_bar`, and `y_bar`
    states.

    This returned system will have

    Input ports:
        (0) u_bar[n] : control vector at timestep n, relative to equilibrium
        (1) y_bar[n] : measurement vector at timestep n, relative to equilibrium

    Output ports:
        (1) x_hat_bar[n] : state vector estimate at timestep n, relative to
                           equilibrium

    Parameters:
        plant : a `Plant` object which can be a LeafSystem or a Diagram.
        x_eq: ndarray
            Equilibrium state vector for discretization
        u_eq: ndarray
            Equilibrium control vector for discretization
        dt: float
            Time step for the discretization.
        Q: ndarray
            Process noise covariance matrix. If `None`, Identity matrix of size
            compatible with `G` and and linearized system's `A` is assumed.
        R: ndarray
            Measurement noise covariance matrix. If `None`, Identity matrix of size
            compatible with linearized system's `C` and `A` is assumed.
        G: ndarray
            Process noise matrix. If `None`, `G=B` is assumed making disrurbances
            additive to control vector `u`, i.e. `u_disturbed = u_orig + w`.
        x_hat_bar_0: ndarray
            Initial state estimate relative to equilibrium.
            If None, an identity matrix is assumed.
        discretization_method: str ("euler" or "zoh")
            Method to discretize the continuous-time plant. Default is "euler".
        discretized_noise: bool
            Whether the user is directly providing Gd, Qd and Rd. Default is False.
            If True, `G`, `Q`, and `R` are assumed to be Gd, Qd, and Rd,
            respectively.
    """
    (
        y_eq,
        Ad,
        Bd,
        Cd,
        Dd,
        Gd,
        Qd,
        Rd,
    ) = linearize_and_discretize_continuous_plant(
        plant, x_eq, u_eq, dt, Q, R, G, discretization_method, discretized_noise
    )

    check_shape_compatibilities(Ad, Bd, Cd, Dd, Gd, Qd, Rd)

    nx = x_eq.size

    if x_hat_bar_0 is None:
        x_hat_bar_0 = npa.zeros(nx)

    # Instantiate an Infinite Horizon Kalman Filter for the linearized plant
    kf = InfiniteHorizonKalmanFilter(
        dt,
        Ad,
        Bd,
        Cd,
        Dd,
        Gd,
        Qd,
        Rd,
        x_hat_bar_0,
        name=name,
        ui_id=ui_id,
    )

    return y_eq, kf

global_filter_for_continuous_plant(plant, x_eq, u_eq, dt, Q=None, R=None, G=None, x_hat_0=None, discretization_method='euler', discretized_noise=False, name=None, ui_id=None) staticmethod

See docs for for_continuous_plant, which returns the local infinite horizon Kalman Filter. This method additionally converts the local Kalman Filter to a global estimator. See docs for make_global_estimator_from_local for details.

Source code in jaxonomy/library/state_estimators/infinite_horizon_kalman_filter.py

@staticmethod
@with_resolved_parameters
def global_filter_for_continuous_plant(
    plant,
    x_eq,
    u_eq,
    dt,
    Q=None,
    R=None,
    G=None,
    x_hat_0=None,
    discretization_method="euler",
    discretized_noise=False,
    name=None,
    ui_id=None,
):
    """
    See docs for `for_continuous_plant`, which returns the local infinite horizon
    Kalman Filter. This method additionally converts the local Kalman Filter to a
    global estimator. See docs for `make_global_estimator_from_local` for details.
    """
    (
        y_eq,
        Ad,
        Bd,
        Cd,
        Dd,
        Gd,
        Qd,
        Rd,
    ) = linearize_and_discretize_continuous_plant(
        plant, x_eq, u_eq, dt, Q, R, G, discretization_method, discretized_noise
    )

    check_shape_compatibilities(Ad, Bd, Cd, Dd, Gd, Qd, Rd)

    nx = x_eq.size

    if x_hat_0 is None:
        x_hat_bar_0 = npa.zeros(nx)
    else:
        x_hat_bar_0 = x_hat_0 - x_eq

    # Instantiate an Infinite Horizon Kalman Filter for the linearized plant
    local_kf = InfiniteHorizonKalmanFilter(
        dt,
        Ad,
        Bd,
        Cd,
        Dd,
        Gd,
        Qd,
        Rd,
        x_hat_bar_0,
        name=name,
        ui_id=ui_id,
    )

    global_kf = make_global_estimator_from_local(
        local_kf,
        x_eq,
        u_eq,
        y_eq,
        name=name,
        ui_id=ui_id,
    )

    return global_kf

Integrator

Bases: LeafSystem

Integrate the input signal in time.

The Integrator block is the main primitive for building continuous-time models. It is a first-order integrator, implementing the following linear time-invariant ordinary differential equation for input values u and output values y:

    ẋ = u
    y = x

where x is the state of the integrator. The integrator is initialized with the value of the initial_state parameter.

Options

Reset: the integrator can be configured to reset its state on an input trigger. The reset value can be either the initial state of the integrator or an external value provided by an input port. Limits: the integrator can be configured such that the output and state are constrained by upper and lower limits. Hold: the integrator can be configured to hold integration based on an input trigger.

The Integrator block is also designed to detect "Zeno" behavior, where the reset events happen asymptotically closer together. This is a pathological case that can cause numerical issues in the simulation and should typically be avoided by introducing some physically realistic hysteresis into the model. However, in the event that Zeno behavior is unavoidable, the integrator will enter a "Zeno" state where the output is held constant until the trigger changes value to False. See the "bouncing ball" demo for a Zeno example.

Input ports

(0) The input signal. Must match the shape and dtype of the initial continuous state. (1) The reset trigger. Optional, only if enable_reset is True. (2) The reset value. Optional, only if enable_external_reset is True. (3) The hold trigger. Optional, only if 'enable_hold' is True.

Output ports

(0) The continuous state of the integrator.

Parameters:

Name	Type	Description	Default
initial_state		The initial value of the integrator state. Can be any array, or even a nested structure of arrays, but the data type should be floating-point.	required
enable_reset		If True, the integrator will reset its state to the initial value when the reset trigger is True. Adds an additional input port for the reset trigger. This signal should be boolean- or binary-valued.	False
enable_external_reset		If True, the integrator will reset its state to the value provided by the reset value input port when the reset trigger is True. Otherwise, the integrator will reset to the initial value. Adds an additional input port for the reset value. This signal should match the shape and dtype of the initial continuous state.	False
enable_limits		If True, the integrator will constrain its state and output to within the upper and lower limits. Either limit may be disbale by setting its value to None.	False
enable_hold		If True, the integrator will hold integration when the hold trigger is True.	False
reset_on_enter_zeno		If True, the integrator will reset its state to the initial value when the integrator enters the Zeno state. This option is ignored unless enable_reset is True.	False
zeno_tolerance		The tolerance used to determine if the integrator is in the Zeno state. If the time between events is less than this tolerance, then the integrator is in the Zeno state. This option is ignored unless enable_reset is True.	1e-06

Events

An event is triggered when the "reset" port changes.

An event is triggered when the state hit one of the limits.

An event is triggered when the "hold" port changes.

Another guard is conditionally active when the integrator is in the Zeno state, and is triggered when the "reset" port changes from True to False. This event is used to exit the Zeno state and resume normal integration.

Source code in jaxonomy/library/dynamics.py

class Integrator(LeafSystem):
    """Integrate the input signal in time.

    The Integrator block is the main primitive for building continuous-time
    models.  It is a first-order integrator, implementing the following linear
    time-invariant ordinary differential equation for input values `u` and output
    values `y`:
    ```
        ẋ = u
        y = x
    ```
    where `x` is the state of the integrator.  The integrator is initialized
    with the value of the `initial_state` parameter.

    Options:
        Reset: the integrator can be configured to reset its state on an input
            trigger.  The reset value can be either the initial state of the
            integrator or an external value provided by an input port.
        Limits: the integrator can be configured such that the output and state
            are constrained by upper and lower limits.
        Hold: the integrator can be configured to hold integration based on an
            input trigger.

    The Integrator block is also designed to detect "Zeno" behavior, where the
    reset events happen asymptotically closer together.  This is a pathological
    case that can cause numerical issues in the simulation and should typically be
    avoided by introducing some physically realistic hysteresis into the model.
    However, in the event that Zeno behavior is unavoidable, the integrator will
    enter a "Zeno" state where the output is held constant until the trigger
    changes value to False.  See the "bouncing ball" demo for a Zeno example.

    Input ports:
        (0) The input signal.  Must match the shape and dtype of the initial
            continuous state.
        (1) The reset trigger.  Optional, only if `enable_reset` is True.
        (2) The reset value.  Optional, only if `enable_external_reset` is True.
        (3) The hold trigger. Optional, only if 'enable_hold' is True.

    Output ports:
        (0) The continuous state of the integrator.

    Parameters:
        initial_state:
            The initial value of the integrator state.  Can be any array, or even
            a nested structure of arrays, but the data type should be floating-point.
        enable_reset:
            If True, the integrator will reset its state to the initial value
            when the reset trigger is True.  Adds an additional input port for
            the reset trigger.  This signal should be boolean- or binary-valued.
        enable_external_reset:
            If True, the integrator will reset its state to the value provided
            by the reset value input port when the reset trigger is True. Otherwise,
            the integrator will reset to the initial value.  Adds an additional
            input port for the reset value.  This signal should match the shape
            and dtype of the initial continuous state.
        enable_limits:
            If True, the integrator will constrain its state and output to within
            the upper and lower limits. Either limit may be disbale by setting its
            value to None.
        enable_hold:
            If True, the integrator will hold integration when the hold trigger is
            True.
        reset_on_enter_zeno:
            If True, the integrator will reset its state to the initial value
            when the integrator enters the Zeno state.  This option is ignored unless
            `enable_reset` is True.
        zeno_tolerance:
            The tolerance used to determine if the integrator is in the Zeno state.
            If the time between events is less than this tolerance, then the
            integrator is in the Zeno state.  This option is ignored unless
            `enable_reset` is True.


    Events:
        An event is triggered when the "reset" port changes.

        An event is triggered when the state hit one of the limits.

        An event is triggered when the "hold" port changes.

        Another guard is conditionally active when the integrator is in the Zeno
        state, and is triggered when the "reset" port changes from True to False.
        This event is used to exit the Zeno state and resume normal integration.
    """

    @parameters(
        static=[
            "enable_reset",
            "enable_external_reset",
            "enable_limits",
            "enable_hold",
            "reset_on_enter_zeno",
        ],
        dynamic=["zeno_tolerance", "lower_limit", "upper_limit", "initial_state"],
    )
    def __init__(
        self,
        initial_state,
        enable_reset=False,
        enable_limits=False,
        lower_limit=None,
        upper_limit=None,
        enable_hold=False,
        enable_external_reset=False,
        zeno_tolerance=1e-6,
        reset_on_enter_zeno=False,
        dtype=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.dtype = dtype
        self.enable_reset = enable_reset
        self.enable_external_reset = enable_external_reset
        self.enable_hold = enable_hold
        self.discrete_state_type = namedtuple(
            "IntegratorDiscreteState", ["zeno", "counter", "tprev"]
        )

        self.xdot_index = self.declare_input_port(name="in_0")

        x0 = npa.array(initial_state, dtype=self.dtype)
        self.dtype = self.dtype if self.dtype is not None else x0.dtype
        self._continuous_state_idx = self.declare_continuous_state(
            default_value=x0,
            ode=self._ode,
            prerequisites_of_calc=[self.input_ports[self.xdot_index].ticket],
        )

        if enable_reset:
            # Boolean input for triggering reset
            self.reset_trigger_index = self.declare_input_port(name="reset_trigger")
            # prerequisites_of_calc.append(
            #     self.input_ports[self.reset_trigger_index].ticket
            # )

            # Declare a custom discrete state to track Zeno behavior
            self.declare_discrete_state(
                default_value=self.discrete_state_type(
                    zeno=False, counter=0, tprev=0.0
                ),
                as_array=False,
            )

            #
            # Declare reset event
            #
            # when reset is triggered, execute the reset map.
            self.declare_zero_crossing(
                guard=self._reset_guard,
                reset_map=self._reset,
                name="reset_on",
                direction="negative_then_non_negative",
            )
            # when reset is deasserted, do not change the state.
            self.declare_zero_crossing(
                guard=self._reset_guard,
                name="reset_off",
                direction="positive_then_non_positive",
            )

            self.declare_zero_crossing(
                guard=self._exit_zeno_guard,
                reset_map=self._exit_zeno,
                name="exit_zeno",
                direction="positive_then_non_positive",
            )

            # Optional: reset value defined by external signal
            if enable_external_reset:
                self.reset_value_index = self.declare_input_port(name="reset_value")
                # prerequisites_of_calc.append(
                #     self.input_ports[self.reset_value_index].ticket
                # )

        if enable_hold:
            # Boolean input for triggering hold assert/deassert
            self.hold_trigger_index = self.declare_input_port(name="hold_trigger")

            def _hold_guard(_time, _state, *inputs, **_params):
                trigger = inputs[self.hold_trigger_index]
                return npa.where(trigger, 1.0, -1.0)

            self.declare_zero_crossing(
                guard=_hold_guard,
                name="hold",
                direction="crosses_zero",
            )

        self._output_port_idx = self.declare_output_port(name="out_0")

    def initialize(
        self,
        initial_state,
        enable_reset=False,
        enable_limits=False,
        lower_limit=None,
        upper_limit=None,
        enable_hold=False,
        enable_external_reset=False,
        zeno_tolerance=1e-6,
        reset_on_enter_zeno=False,
    ):
        if self.enable_reset != enable_reset:
            raise ValueError("enable_reset cannot be changed after initialization")
        if self.enable_external_reset != enable_external_reset:
            raise ValueError(
                "enable_external_reset cannot be changed after initialization"
            )
        if self.enable_hold != enable_hold:
            raise ValueError("enable_hold cannot be changed after initialization")

        # Default initial condition unless modified in context
        x0 = npa.array(initial_state, dtype=self.dtype)
        self.dtype = self.dtype if self.dtype is not None else x0.dtype

        self.configure_continuous_state(
            self._continuous_state_idx,
            default_value=x0,
            ode=self._ode,
            prerequisites_of_calc=[self.input_ports[self.xdot_index].ticket],
        )

        self.reset_on_enter_zeno = reset_on_enter_zeno

        self.enable_limits = enable_limits
        self.has_lower_limit = lower_limit is not None
        self.has_upper_limit = upper_limit is not None

        self.configure_output_port(
            self._output_port_idx,
            self._output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
        )

        if enable_limits:
            if lower_limit is not None:

                def _lower_limit_guard(_time, state, *_inputs, **params):
                    return state.continuous_state - params["lower_limit"]

                self.declare_zero_crossing(
                    guard=_lower_limit_guard,
                    name="lower_limit",
                    direction="positive_then_non_positive",
                )

            if upper_limit is not None:

                def _upper_limit_guard(_time, state, *_inputs, **params):
                    return state.continuous_state - params["upper_limit"]

                self.declare_zero_crossing(
                    guard=_upper_limit_guard,
                    name="upper_limit",
                    direction="negative_then_non_negative",
                )

    def reset_default_values(self, **dynamic_parameters):
        x0 = npa.array(dynamic_parameters["initial_state"], dtype=self.dtype)
        self.configure_continuous_state_default_value(
            self._continuous_state_idx,
            default_value=x0,
        )

    def _ode(self, _time, state, *inputs, **params):
        # Normally, just integrate the input signal
        xdot = inputs[self.xdot_index]

        # However, if the reset trigger is high or the integrator is in the Zeno state,
        # then the integrator should hold
        if self.enable_reset:
            trigger = inputs[self.reset_trigger_index]
            in_zeno_state = state.discrete_state.zeno
            xdot = npa.where((trigger | in_zeno_state), npa.zeros_like(xdot), xdot)

        # Additionally, if the limits are enabled, the derivative is set to zero if
        # either limit is presnetly violated.
        if self.enable_limits:
            xc = state.continuous_state

            if self.has_lower_limit:
                llim_violation = npa.logical_and(
                    xdot < 0.0, xc <= params["lower_limit"]
                )
            else:
                llim_violation = False

            if self.has_upper_limit:
                ulim_violation = npa.logical_and(
                    xdot > 0.0, xc >= params["upper_limit"]
                )
            else:
                ulim_violation = False

            xdot = npa.where(
                (llim_violation | ulim_violation), npa.zeros_like(xdot), xdot
            )

        if self.enable_hold:
            hold = inputs[self.hold_trigger_index]
            xdot = npa.where(hold, npa.zeros_like(xdot), xdot)

        return xdot

    def _output(self, _time, state, *_inputs, **params):
        xc = state.continuous_state
        if self.enable_limits:
            lower_limit = params["lower_limit"] if self.has_lower_limit else -np.inf
            upper_limit = params["upper_limit"] if self.has_upper_limit else np.inf
            return npa.clip(xc, lower_limit, upper_limit)

        return xc

    def _reset_guard(self, _time, _state, *inputs, **_params):
        trigger = inputs[self.reset_trigger_index]
        return npa.where(trigger, 1.0, -1.0)

    def _reset(self, time, state, *inputs, **params):
        # If the distance between events is less than the tolerance, then enter the Zeno state.
        dt = time - state.discrete_state.tprev
        zeno = (dt - params["zeno_tolerance"]) <= 0
        tprev = time

        # Handle the reset event as usual
        if self.enable_external_reset:
            xc = inputs[self.reset_value_index]
        else:
            xc = npa.array(params["initial_state"], dtype=self.dtype)

        # Don't reset if entering Zeno state
        new_continuous_state = npa.where(
            zeno & (not self.reset_on_enter_zeno),
            state.continuous_state,
            xc,
        )
        state = state.with_continuous_state(new_continuous_state)

        # Count number of resets (for debugging)
        counter = state.discrete_state.counter + 1

        # Update the discrete state
        xd_plus = self.discrete_state_type(zeno=zeno, counter=counter, tprev=tprev)
        state = state.with_discrete_state(xd_plus)

        logger.debug("Resetting to %s", state)
        return state

    def _exit_zeno_guard(self, _time, _state, *inputs, **_params):
        # This will only be active when in the Zeno state.  It monitors the boolean trigger input
        # and will go from 1.0 (when trigger=True) to 0.0 (when trigger=False)
        trigger = inputs[self.reset_trigger_index]
        return npa.array(trigger, dtype=self.dtype)

    def _exit_zeno(self, _time, state, *_inputs, **_params):
        xd = state.discrete_state._replace(zeno=False)
        return state.with_discrete_state(xd)

    def determine_active_guards(self, root_context):
        # TODO: Update this to use the new zero crossing event system
        # defined in LeafSystem.
        zero_crossing_events = self.zero_crossing_events.mark_all_active()

        if not self.enable_reset:
            return zero_crossing_events

        def _get_reset(events: LeafEventCollection):
            return events.events[0]

        context = root_context[self.system_id]
        in_zeno_state = context.discrete_state.zeno

        reset = cond(
            in_zeno_state,
            lambda e: e.mark_inactive(),
            lambda e: e.mark_active(),
            _get_reset(zero_crossing_events),
        )

        def _get_exit_zeno(events: LeafEventCollection):
            return events.events[1]

        exit_zeno: ZeroCrossingEvent = cond(
            in_zeno_state,
            lambda e: e.mark_active(),
            lambda e: e.mark_inactive(),
            _get_exit_zeno(zero_crossing_events),
        )

        zero_crossing_events = eqx.tree_at(_get_reset, zero_crossing_events, reset)
        zero_crossing_events = eqx.tree_at(
            _get_exit_zeno, zero_crossing_events, exit_zeno
        )

        return zero_crossing_events

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        u = self.eval_input(context)
        xc = context[self.system_id].continuous_state
        check_state_type(
            self,
            inp_data=u,
            state_data=xc,
            error_collector=error_collector,
        )

IntegratorDiscrete

Bases: LeafSystem

Discrete first-order integrator.

This block is a discrete-time approximation to the behavior of the Integrator block. It implements the following linear time-invariant difference equation for input values u and output values y:

    x[k+1] = x[k] + dt * u[k]
    y[k] = x[k]

where x is the state of the integrator. The integrator is initialized with the value of the initial_state parameter.

Options

Reset: the integrator can be configured to reset its state on an input trigger. The reset value can be either the initial state of the integrator or an external value provided by an input port. Limits: the integrator can be configured such that the output and state are constrained by upper and lower limits. Hold: the integrator can be configured to hold integration based on an input trigger.

Unlike the continuous-time integrator, the discrete integrator does not detect Zeno behavior, since this is not a concern in discrete-time systems.

Input ports

(0) The input signal. Must match the shape and dtype of the initial state. (1) The reset trigger. Optional, only if enable_reset is True. (2) The reset value. Optional, only if enable_external_reset is True. (3) The hold trigger. Optional, only if 'enable_hold' is True.

Output ports

(0) The current state of the integrator.

Parameters:

Name	Type	Description	Default
initial_state		The initial value of the integrator state. Can be any array, or even a nested structure of arrays, but the data type should be floating-point.	required
enable_reset		If True, the integrator will reset its state to the initial value when the reset trigger is True. Adds an additional input port for the reset trigger. This signal should be boolean- or binary-valued.	False
enable_external_reset		If True, the integrator will reset its state to the value provided by the reset value input port when the reset trigger is True. Otherwise, the integrator will reset to the initial value. Adds an additional input port for the reset value. This signal should match the shape and dtype of the initial continuous state.	False
enable_limits		If True, the integrator will constrain its state and output to within the upper and lower limits. Either limit may be disbale by setting its value to None.	False
enable_hold		If True, the integrator will hold integration when the hold trigger is True.	False

Source code in jaxonomy/library/dynamics.py

class IntegratorDiscrete(LeafSystem):
    """Discrete first-order integrator.

    This block is a discrete-time approximation to the behavior of the Integrator
    block.  It implements the following linear time-invariant difference equation
    for input values `u` and output values `y`:
    ```
        x[k+1] = x[k] + dt * u[k]
        y[k] = x[k]
    ```
    where `x` is the state of the integrator.  The integrator is initialized with
    the value of the `initial_state` parameter.

    Options:
        Reset: the integrator can be configured to reset its state on an input
            trigger.  The reset value can be either the initial state of the
            integrator or an external value provided by an input port.
        Limits: the integrator can be configured such that the output and state
            are constrained by upper and lower limits.
        Hold: the integrator can be configured to hold integration based on an
            input trigger.

    Unlike the continuous-time integrator, the discrete integrator does not detect
    Zeno behavior, since this is not a concern in discrete-time systems.

    Input ports:
        (0) The input signal.  Must match the shape and dtype of the initial
            state.
        (1) The reset trigger.  Optional, only if `enable_reset` is True.
        (2) The reset value.  Optional, only if `enable_external_reset` is True.
        (3) The hold trigger. Optional, only if 'enable_hold' is True.

    Output ports:
        (0) The current state of the integrator.

    Parameters:
        initial_state:
            The initial value of the integrator state.  Can be any array, or even
            a nested structure of arrays, but the data type should be floating-point.
        enable_reset:
            If True, the integrator will reset its state to the initial value
            when the reset trigger is True.  Adds an additional input port for
            the reset trigger.  This signal should be boolean- or binary-valued.
        enable_external_reset:
            If True, the integrator will reset its state to the value provided
            by the reset value input port when the reset trigger is True. Otherwise,
            the integrator will reset to the initial value.  Adds an additional
            input port for the reset value.  This signal should match the shape
            and dtype of the initial continuous state.
        enable_limits:
            If True, the integrator will constrain its state and output to within
            the upper and lower limits. Either limit may be disbale by setting its
            value to None.
        enable_hold:
            If True, the integrator will hold integration when the hold trigger is
            True.
    """

    @parameters(
        static=[
            "dt",
            "enable_reset",
            "enable_external_reset",
            "enable_limits",
            "enable_hold",
        ],
        dynamic=["lower_limit", "upper_limit", "initial_state"],
    )
    def __init__(
        self,
        dt,
        initial_state,
        enable_reset=False,
        enable_hold=False,
        enable_limits=False,
        lower_limit=None,
        upper_limit=None,
        enable_external_reset=False,
        dtype=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.dt = dt
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self.dtype = dtype

        self.enable_reset = enable_reset
        self.enable_external_reset = enable_external_reset

        self.xdot_index = self.declare_input_port(
            name="in_0"
        )  # One vector-valued input

        self._periodic_update_idx = self.declare_periodic_update()

        if enable_reset:
            self.reset_trigger_index = self.declare_input_port(
                name="reset_trigger"
            )  # Boolean input for triggering reset

            if enable_external_reset:
                self.reset_value_index = self.declare_input_port(
                    name="reset_value"
                )  # Optional reset value

        self.enable_hold = enable_hold
        if enable_hold:
            self.hold_trigger_index = self.declare_input_port(
                name="hold_trigger"
            )  # Boolean input for triggering hold

        self.state_output_index = self.declare_output_port(name="out_0")

    def initialize(
        self,
        initial_state,
        enable_reset=False,
        enable_hold=False,
        enable_limits=False,
        lower_limit=None,
        upper_limit=None,
        enable_external_reset=False,
        dt=None,
    ):
        if self.enable_reset != enable_reset:
            raise ValueError("enable_reset cannot be changed after initialization")
        if self.enable_external_reset != enable_external_reset:
            raise ValueError(
                "enable_external_reset cannot be changed after initialization"
            )
        if self.enable_hold != enable_hold:
            raise ValueError("enable_hold cannot be changed after initialization")

        # Default initial condition unless modified in context
        x0 = npa.array(initial_state, dtype=self.dtype)
        self.dtype = self.dtype if self.dtype is not None else x0.dtype
        self.declare_discrete_state(default_value=x0)
        self.configure_periodic_update(
            self._periodic_update_idx, self._update, period=self.dt, offset=0.0
        )

        # Since the reset is applied to the output port, having this
        # active makes the block feedthrough with respect to related
        # input ports.
        self.is_feedthrough = enable_reset

        self.enable_limits = enable_limits
        self.has_lower_limit = lower_limit is not None
        self.has_upper_limit = upper_limit is not None

        prereqs = [DependencyTicket.xd]
        if enable_reset:
            prereqs.append(self.input_ports[self.reset_trigger_index].ticket)
            if enable_external_reset:
                prereqs.append(self.input_ports[self.reset_value_index].ticket)

        self.configure_output_port(
            self.state_output_index,
            self._output,
            period=self.dt,
            offset=0.0,
            default_value=x0,
            prerequisites_of_calc=prereqs,
        )

    def reset_default_values(self, **dynamic_parameters):
        x0 = npa.array(dynamic_parameters["initial_state"], dtype=self.dtype)
        self.configure_discrete_state_default_value(default_value=x0)
        self.configure_output_port_default_value(self.state_output_index, x0)

    def _reset(self, *inputs, **params):
        if self.enable_external_reset:
            return inputs[self.reset_value_index]
        return npa.array(params["initial_state"], dtype=self.dtype)

    def _apply_reset_and_limits(self, x_new, *inputs, **params):
        # Reset and limits are applied to both the update and outputs
        # so that they respond to the discontinuities simultaneously.

        if self.enable_reset:
            # If the reset is high, then return the reset value
            trigger = inputs[self.reset_trigger_index]
            x_new = npa.where(trigger, self._reset(*inputs, **params), x_new)

        if self.enable_limits:
            lower_limit = params["lower_limit"] if self.has_lower_limit else -npa.inf
            upper_limit = params["upper_limit"] if self.has_upper_limit else npa.inf
            x_new = npa.clip(x_new, lower_limit, upper_limit)

        return x_new

    def _apply_hold(self, x, x_new, *inputs, **_params):
        # Hold is only applied to the update, but not the output

        if self.enable_hold:
            # If the reset is high, then return the reset value
            trigger = inputs[self.hold_trigger_index]
            x_new = npa.where(trigger, x, x_new)

        return x_new

    def _update(self, _time, state, *inputs, **params):
        x = state.discrete_state
        xdot = inputs[self.xdot_index]
        x_new = x + self.dt * xdot
        x_new = self._apply_hold(x, x_new, *inputs, **params)
        x_new = self._apply_reset_and_limits(x_new, *inputs, **params)
        return x_new.astype(x.dtype)

    def _output(self, _time, state, *inputs, **params):
        x = state.discrete_state
        # To ensure that the discontinuities happen simultaneously with
        # the input signal, also apply the reset and limits to the outputs.
        # this makes the block feedthrough.
        y = self._apply_reset_and_limits(x, *inputs, **params)
        return y

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        u = self.eval_input(context)
        xd = context[self.system_id].discrete_state
        check_state_type(
            self,
            inp_data=u,
            state_data=xd,
            error_collector=error_collector,
        )

InterpolationUsingPrelookup

Bases: LeafSystem

Interpolate a static output_array using a precomputed (index, fraction) tuple from :class:Prelookup.

This is the downstream half of the standard Prelookup/InterpolationUsingPrelookup pair. Plug as many of these as you like into a single :class:Prelookup's output port -- each one interpolates its OWN output_array against the shared (index, fraction) signal, avoiding the redundant binary searches you would do with N independent :class:LookupTable1d blocks.

Input ports

(0) -- the (index, fraction) NamedTuple produced by an upstream :class:Prelookup block.

Output ports

(0) -- the linearly-interpolated value (1 - alpha) * output_array[i] + alpha * output_array[i + 1].

Parameters:

Name	Type	Description	Default
output_array		1-D table values, shape (N,) where N == len(prelookup input_array). output_array[i] corresponds to the i-th breakpoint in the upstream :class:Prelookup's grid.	required
dtype	optional	If set (e.g. jnp.float32), the table values are cast to this dtype on construction. Mirrors the per-block dtype contract of :class:LookupTable1d.	None
extrapolation	(optional, T - 114 - fu - prelookup - extrap)	Out-of-range policy. Must match the upstream :class:Prelookup's extrapolation kwarg (the math is all done in the producer's alpha computation; this block just consumes alpha). Exists on this side of the API for symmetry, IDE discoverability and so model JSON / user code make intent explicit. One of "clip" (default, byte-equivalent to T-114-fu-prelookup), "linear", or "nan".	'clip'

Notes

Differentiable through output_array (every time-step the interpolation is a convex combination of two of its entries; the gradient w.r.t. the picked entries is exact). The fraction-side gradient flows through the upstream :class:Prelookup's alpha computation; the discrete index is piecewise constant (expected -- same as every searchsorted-based block in the library).

Today only linear interpolation is supported -- PCHIP/Akima downstream interpolation is a deeper followup (T-114-followup-prelookup-cubic).

Source code in jaxonomy/library/tables.py

class InterpolationUsingPrelookup(LeafSystem):
    """Interpolate a static ``output_array`` using a precomputed
    (index, fraction) tuple from :class:`Prelookup`.

    This is the downstream half of the standard
    ``Prelookup``/``InterpolationUsingPrelookup`` pair.  Plug as many of
    these as you like into a single :class:`Prelookup`'s output port --
    each one interpolates its OWN ``output_array`` against the shared
    (index, fraction) signal, avoiding the redundant binary searches you
    would do with N independent :class:`LookupTable1d` blocks.

    Input ports:
        ``(0)`` -- the (index, fraction) NamedTuple produced by an
        upstream :class:`Prelookup` block.

    Output ports:
        ``(0)`` -- the linearly-interpolated value
        ``(1 - alpha) * output_array[i] + alpha * output_array[i + 1]``.

    Parameters:
        output_array:
            1-D table values, shape ``(N,)`` where ``N == len(prelookup
            input_array)``.  ``output_array[i]`` corresponds to the
            i-th breakpoint in the upstream :class:`Prelookup`'s grid.
        dtype (optional):
            If set (e.g. ``jnp.float32``), the table values are cast to
            this dtype on construction.  Mirrors the per-block dtype
            contract of :class:`LookupTable1d`.
        extrapolation (optional, T-114-fu-prelookup-extrap):
            Out-of-range policy.  Must match the upstream
            :class:`Prelookup`'s ``extrapolation`` kwarg (the math is
            all done in the producer's ``alpha`` computation; this
            block just consumes ``alpha``).  Exists on this side of the
            API for symmetry, IDE discoverability and so model JSON /
            user code make intent explicit.  One of ``"clip"`` (default,
            byte-equivalent to T-114-fu-prelookup), ``"linear"``, or
            ``"nan"``.

    Notes:
        Differentiable through ``output_array`` (every time-step the
        interpolation is a convex combination of two of its entries; the
        gradient w.r.t. the picked entries is exact).  The fraction-side
        gradient flows through the upstream :class:`Prelookup`'s
        ``alpha`` computation; the discrete ``index`` is piecewise
        constant (expected -- same as every ``searchsorted``-based block
        in the library).

        Today only linear interpolation is
        supported -- PCHIP/Akima downstream interpolation is a deeper
        followup (``T-114-followup-prelookup-cubic``).
    """

    def __init__(self, output_array, dtype=None, extrapolation="clip", **kwargs):
        # Per-block dtype + active precision policy fallback, matching
        # the LookupTable1d / Prelookup contract.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype

        # T-114-fu-prelookup-extrap -- validate the kwarg eagerly.  The
        # producer ``Prelookup`` does all the alpha math; this block
        # just records the user's declared intent for API symmetry.
        if extrapolation not in ("clip", "linear", "nan"):
            raise ValueError(
                f"InterpolationUsingPrelookup: extrapolation must be one "
                f"of ('clip','linear','nan'), got {extrapolation!r}"
            )
        self._extrapolation = extrapolation

        _out_np = np.asarray(output_array)
        if _out_np.ndim != 1:
            raise ValueError(
                f"InterpolationUsingPrelookup: output_array must be 1-D, "
                f"got shape {_out_np.shape}"
            )
        if _out_np.size < 2:
            raise ValueError(
                f"InterpolationUsingPrelookup: output_array must have at "
                f"least 2 entries, got shape {_out_np.shape}"
            )

        if self._dtype is not None:
            self._output_array = npa.asarray(_out_np).astype(self._dtype)
        else:
            self._output_array = npa.array(_out_np)

        super().__init__(**kwargs)
        self.declare_input_port()

        # Capture the table in a local so the closure does not pull
        # ``self`` into the JAX trace.
        yp_local = self._output_array

        def _compute(_time, _state, *inputs, **_params):
            (prelookup_result,) = inputs
            # Unpack the NamedTuple.  ``index`` and ``fraction`` are
            # plain JAX arrays produced by the upstream Prelookup
            # closure; both are pytree-friendly via NamedTuple
            # registration.
            i = prelookup_result.index
            alpha = prelookup_result.fraction
            return (1.0 - alpha) * yp_local[i] + alpha * yp_local[i + 1]

        self.declare_output_port(
            _compute,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )

    @property
    def output_array(self):
        """The 1-D table being interpolated."""
        return self._output_array

    @property
    def extrapolation(self):
        """The declared OOB policy (must match the upstream Prelookup)."""
        return self._extrapolation

extrapolation property

The declared OOB policy (must match the upstream Prelookup).

output_array property

The 1-D table being interpolated.

KalmanFilter

Bases: KalmanFilterBase

Kalman Filter for the following system:

x[n+1] = A x[n] + B u[n] + G w[n]
y[n]   = C x[n] + D u[n] + v[n]

E(w[n]) = E(v[n]) = 0
E(w[n]w'[n]) = Q
E(v[n]v'[n] = R
E(w[n]v'[n] = N = 0

Input ports

(0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n

Output ports

(1) x_hat[n] : state vector estimate at timestep n

Parameters:

Name	Type	Description	Default
dt		float Time step of the discrete-time system	required
A		ndarray State transition matrix	required
B		ndarray Input matrix	required
C		ndarray Output matrix. If None, full state output is assumed.	None
D		ndarray Feedthrough matrix. If None, no feedthrough is assumed.	None
G		ndarray Process noise matrix. If None, G=B is assumed.	None
Q		ndarray Process noise covariance matrix. If None, Identity matrix of size compatible with G and A is assumed.	None
R		ndarray Measurement noise covariance matrix. If None, Identity matrix of size compatible with C and A is assumed.	None
x_hat_0		ndarray Initial state estimate. If None, an array of zeros is assumed.	None
P_hat_0		ndarray Initial state covariance matrix estimate. If None, Identity matrix of size identical to A is assumed.	None

Source code in jaxonomy/library/state_estimators/kalman_filter.py

class KalmanFilter(KalmanFilterBase):
    """
    Kalman Filter for the following system:

    ```
    x[n+1] = A x[n] + B u[n] + G w[n]
    y[n]   = C x[n] + D u[n] + v[n]

    E(w[n]) = E(v[n]) = 0
    E(w[n]w'[n]) = Q
    E(v[n]v'[n] = R
    E(w[n]v'[n] = N = 0
    ```

    Input ports:
        (0) u[n] : control vector at timestep n
        (1) y[n] : measurement vector at timestep n

    Output ports:
        (1) x_hat[n] : state vector estimate at timestep n

    Parameters:
        dt: float
            Time step of the discrete-time system
        A: ndarray
            State transition matrix
        B: ndarray
            Input matrix
        C: ndarray
            Output matrix. If `None`, full state output is assumed.
        D: ndarray
            Feedthrough matrix. If `None`, no feedthrough is assumed.
        G: ndarray
            Process noise matrix. If `None`, `G=B` is assumed.
        Q: ndarray
            Process noise covariance matrix. If `None`, Identity matrix of size
            compatible with `G` and `A` is assumed.
        R: ndarray
            Measurement noise covariance matrix. If `None`, Identity matrix of size
            compatible with `C` and `A` is assumed.
        x_hat_0: ndarray
            Initial state estimate. If `None`, an array of zeros is assumed.
        P_hat_0: ndarray
            Initial state covariance matrix estimate. If `None`, Identity matrix of size
            identical to `A` is assumed.
    """

    @parameters(
        static=["dt", "A", "B", "C", "D", "G", "Q", "R", "x_hat_0", "P_hat_0"],
    )
    def __init__(
        self,
        dt,
        A,
        B,
        C=None,
        D=None,
        G=None,
        Q=None,
        R=None,
        x_hat_0=None,
        P_hat_0=None,
        name=None,
        **kwargs,
    ):
        is_feedthrough = False if D is None else bool(not npa.allclose(D, 0.0))
        super().__init__(dt, x_hat_0, P_hat_0, is_feedthrough, name, **kwargs)

    def initialize(
        self,
        dt,
        A,
        B,
        C=None,
        D=None,
        G=None,
        Q=None,
        R=None,
        x_hat_0=None,
        P_hat_0=None,
    ):
        self.nx, self.nu = B.shape

        if C is None:
            C = jnp.eye(self.nx)
            self.ny = self.nx
        else:
            self.ny = C.shape[0]

        if D is None:
            D = jnp.zeros((self.ny, self.nu))
        self.is_feedthrough = bool(not npa.allclose(D, 0.0))

        if G is None:
            G = B

        _, self.nd = G.shape

        if Q is None:
            Q = jnp.eye(self.nd)

        if R is None:
            R = jnp.eye(self.ny)

        if x_hat_0 is None:
            x_hat_0 = jnp.zeros(self.nx)

        if P_hat_0 is None:
            P_hat_0 = jnp.eye(self.nx)

        check_shape_compatibilities(A, B, C, D, G, Q, R)

        self.A = A
        self.B = B
        self.C = C
        self.D = D
        self.G = G
        self.Q = Q
        self.R = R

        self.eye_x = jnp.eye(self.nx)
        self.GQGT = G @ Q @ G.T

    def _correct(self, time, x_hat_minus, P_hat_minus, *inputs):
        u, y = inputs
        y = jnp.atleast_1d(y)

        C, D = self.C, self.D

        # Compute Kalman gain K = P C^T S^{-1} via a linear solve instead of
        # explicitly inverting S = C P C^T + R.  This is more numerically stable.
        S = C @ P_hat_minus @ C.T + self.R
        K = jnp.linalg.solve(S.T, (P_hat_minus @ C.T).T).T

        x_hat_plus = x_hat_minus + jnp.dot(K, y - jnp.dot(C, x_hat_minus))  # n|n

        if self.is_feedthrough:
            u = npa.atleast_1d(u)
            x_hat_plus = x_hat_plus - npa.dot(K, npa.dot(D, u))

        P_hat_plus = jnp.matmul(self.eye_x - jnp.matmul(K, C), P_hat_minus)  # n|n

        return x_hat_plus, P_hat_plus

    def _propagate(self, time, x_hat_plus, P_hat_plus, *inputs):
        # Predict -- x_hat_plus of current step is propagated to be the
        # x_hat_minus of the next step
        # n+1|n in current step is n|n-1 for next step

        u, y = inputs
        u = jnp.atleast_1d(u)

        A, B = self.A, self.B

        x_hat_minus = jnp.dot(A, x_hat_plus) + jnp.dot(B, u)  # n+1|n
        P_hat_minus = A @ P_hat_plus @ A.T + self.GQGT  # n+1|n

        return x_hat_minus, P_hat_minus

    #######################################
    # Make filter for a continuous plant  #
    #######################################

    @staticmethod
    @with_resolved_parameters
    def for_continuous_plant(
        plant,
        x_eq,
        u_eq,
        dt,
        Q=None,
        R=None,
        G=None,
        x_hat_bar_0=None,
        P_hat_bar_0=None,
        discretization_method="euler",
        discretized_noise=False,
        name=None,
        ui_id=None,
    ):
        """
        Obtain a Kalman Filter system for a continuous-time plant after linearization
        at equilibrium point (x_eq, u_eq)

        The input plant contains the deterministic forms of the forward and observation
        operators:

        ```
            dx/dt = f(x,u)
            y = g(x,u)
        ```

        Note: (i) Only plants with one vector-valued input and one vector-valued output
        are currently supported. Furthermore, the plant LeafSystem/Diagram should have
        only one vector-valued integrator.

        A plant with disturbances of the following form is then considered
        following form:

        ```
            dx/dt = f(x,u) + G w                        --- (C1)
            y = g(x,u) +  v                             --- (C2)
        ```

        where:

            `w` represents the process noise,
            `v` represents the measurement noise,

        and

        ```
            E(w) = E(v) = 0
            E(ww') = Q
            E(vv') = R
            E(wv') = N = 0
        ```

        This plant with disturbances is linearized (only `f` and `g`) around the
        equilibrium point to obtain:

        ```
            d/dt (x_bar) = A x_bar + B u_bar + G w
            y_bar = C x_bar + D u_bar + v
        ```

        where,

        ```
            x_bar = x - x_eq
            u_bar = u - u_eq
            y_bar = y - y_bar
            y_eq = g(x_eq, u_eq)
        ```

        The linearized plant is then discretized via `euler` or `zoh` method to obtain:

        ```
            x_bar[n] = Ad x_bar[n] + Bd u_bar[n] + Gd w[n]           --- (L1)
            y_bar[n] = Cd x_bar[n] + Dd u_bar[n] + v[n]              --- (L2)

            E(w[n]) = E(v[n]) = 0
            E(w[n]w'[n]) = Qd
            E(v[n]v'[n]) = Rd
            E(w[n]v'[n]) = Nd = 0
        ```

        Note: If `discretized_noise` is True, then it is assumed that the user is
        providing Gd, Qd and Rd. If False, then Qd and Rd are computed from
        continuous-time Q, R, and G, and Gd is set to Identity matrix.

        A Kalman Filter estimator for the system of equations (L1) and (L2) is
        created and returned. This filter is in the `x_bar`, `u_bar`, and `y_bar`
        states.

        This returned system will have

        Input ports:
            (0) u_bar[n] : control vector at timestep n, relative to equilibrium
            (1) y_bar[n] : measurement vector at timestep n, relative to equilibrium

        Output ports:
            (1) x_hat_bar[n] : state vector estimate at timestep n, relative to
                               equilibrium

        Parameters:
            plant : a `Plant` object which can be a LeafSystem or a Diagram.
            x_eq: ndarray
                Equilibrium state vector for discretization
            u_eq: ndarray
                Equilibrium control vector for discretization
            dt: float
                Time step for the discretization.
            Q: ndarray
                Process noise covariance matrix. If `None`, Identity matrix of size
                compatible with `G` and and linearized system's `A` is assumed.
            R: ndarray
                Measurement noise covariance matrix. If `None`, Identity matrix of size
                compatible with linearized system's `C` and `A` is assumed.
            G: ndarray
                Process noise matrix. If `None`, `G=B` is assumed making disrurbances
                additive to control vector `u`, i.e. `u_disturbed = u_orig + w`.
            x_hat_bar_0: ndarray
                Initial state estimate, relative to equilirium.
                If None, an identity matrix is assumed.
            P_hat_bar_0: ndarray
                Initial covariance matrix estimate for state, relative to equilibrium.
                If `None`, an Identity matrix is assumed.
            discretization_method: str ("euler" or "zoh")
                Method to discretize the continuous-time plant. Default is "euler".
            discretized_noise: bool
                Whether the user is directly providing Gd, Qd and Rd. Default is False.
                If True, `G`, `Q`, and `R` are assumed to be Gd, Qd, and Rd,
                respectively.
        """
        (
            y_eq,
            Ad,
            Bd,
            Cd,
            Dd,
            Gd,
            Qd,
            Rd,
        ) = linearize_and_discretize_continuous_plant(
            plant, x_eq, u_eq, dt, Q, R, G, discretization_method, discretized_noise
        )

        check_shape_compatibilities(Ad, Bd, Cd, Dd, Gd, Qd, Rd)

        nx = x_eq.size

        if x_hat_bar_0 is None:
            x_hat_bar_0 = jnp.zeros(nx)

        if P_hat_bar_0 is None:
            P_hat_bar_0 = jnp.eye(nx)

        # Instantiate a Kalman Filter for the linearized plant
        kf = KalmanFilter(
            dt,
            Ad,
            Bd,
            Cd,
            Dd,
            Gd,
            Qd,
            Rd,
            x_hat_bar_0,
            P_hat_bar_0,
            name=name,
            ui_id=ui_id,
        )

        return y_eq, kf

    ##############################################
    # Make global filter for a continuous plant  #
    ##############################################

    @staticmethod
    @with_resolved_parameters
    def global_filter_for_continuous_plant(
        plant,
        x_eq,
        u_eq,
        dt,
        Q=None,
        R=None,
        G=None,
        x_hat_0=None,
        P_hat_0=None,
        discretization_method="euler",
        discretized_noise=False,
        name=None,
        ui_id=None,
    ):
        """
        See docs for `for_continuous_plant`, which returns the local Kalman
        Filter. This method additionally converts the local Kalman Filter to a
        global estimator. See docs for `make_global_estimator_from_local` for details.
        """
        (
            y_eq,
            Ad,
            Bd,
            Cd,
            Dd,
            Gd,
            Qd,
            Rd,
        ) = linearize_and_discretize_continuous_plant(
            plant, x_eq, u_eq, dt, Q, R, G, discretization_method, discretized_noise
        )

        check_shape_compatibilities(Ad, Bd, Cd, Dd, Gd, Qd, Rd)

        nx = x_eq.size

        if x_hat_0 is None:
            x_hat_bar_0 = jnp.zeros(nx)
        else:
            x_hat_bar_0 = x_hat_0 - x_eq

        if P_hat_0 is None:
            P_hat_bar_0 = jnp.eye(nx)
        else:
            P_hat_bar_0 = P_hat_0

        # Instantiate a Kalman Filter for the linearized plant
        local_kf = KalmanFilter(
            dt,
            Ad,
            Bd,
            Cd,
            Dd,
            Gd,
            Qd,
            Rd,
            x_hat_bar_0,
            P_hat_bar_0,
            name=name + "_local" if name is not None else None,
        )

        global_kf = make_global_estimator_from_local(
            local_kf,
            x_eq,
            u_eq,
            y_eq,
            name=name,
            ui_id=ui_id,
        )

        return global_kf

for_continuous_plant(plant, x_eq, u_eq, dt, Q=None, R=None, G=None, x_hat_bar_0=None, P_hat_bar_0=None, discretization_method='euler', discretized_noise=False, name=None, ui_id=None) staticmethod

Obtain a Kalman Filter system for a continuous-time plant after linearization at equilibrium point (x_eq, u_eq)

The input plant contains the deterministic forms of the forward and observation operators:

    dx/dt = f(x,u)
    y = g(x,u)

Note: (i) Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator.

A plant with disturbances of the following form is then considered following form:

    dx/dt = f(x,u) + G w                        --- (C1)
    y = g(x,u) +  v                             --- (C2)

where:

`w` represents the process noise,
`v` represents the measurement noise,

and

    E(w) = E(v) = 0
    E(ww') = Q
    E(vv') = R
    E(wv') = N = 0

This plant with disturbances is linearized (only f and g) around the equilibrium point to obtain:

    d/dt (x_bar) = A x_bar + B u_bar + G w
    y_bar = C x_bar + D u_bar + v

where,

    x_bar = x - x_eq
    u_bar = u - u_eq
    y_bar = y - y_bar
    y_eq = g(x_eq, u_eq)

The linearized plant is then discretized via euler or zoh method to obtain:

    x_bar[n] = Ad x_bar[n] + Bd u_bar[n] + Gd w[n]           --- (L1)
    y_bar[n] = Cd x_bar[n] + Dd u_bar[n] + v[n]              --- (L2)

    E(w[n]) = E(v[n]) = 0
    E(w[n]w'[n]) = Qd
    E(v[n]v'[n]) = Rd
    E(w[n]v'[n]) = Nd = 0

Note: If discretized_noise is True, then it is assumed that the user is providing Gd, Qd and Rd. If False, then Qd and Rd are computed from continuous-time Q, R, and G, and Gd is set to Identity matrix.

A Kalman Filter estimator for the system of equations (L1) and (L2) is created and returned. This filter is in the x_bar, u_bar, and y_bar states.

This returned system will have

Input ports

(0) u_bar[n] : control vector at timestep n, relative to equilibrium (1) y_bar[n] : measurement vector at timestep n, relative to equilibrium

Output ports

(1) x_hat_bar[n] : state vector estimate at timestep n, relative to equilibrium

Parameters:

Name	Type	Description	Default
plant		a Plant object which can be a LeafSystem or a Diagram.	required
x_eq		ndarray Equilibrium state vector for discretization	required
u_eq		ndarray Equilibrium control vector for discretization	required
dt		float Time step for the discretization.	required
Q		ndarray Process noise covariance matrix. If None, Identity matrix of size compatible with G and and linearized system's A is assumed.	None
R		ndarray Measurement noise covariance matrix. If None, Identity matrix of size compatible with linearized system's C and A is assumed.	None
G		ndarray Process noise matrix. If None, G=B is assumed making disrurbances additive to control vector u, i.e. u_disturbed = u_orig + w.	None
x_hat_bar_0		ndarray Initial state estimate, relative to equilirium. If None, an identity matrix is assumed.	None
P_hat_bar_0		ndarray Initial covariance matrix estimate for state, relative to equilibrium. If None, an Identity matrix is assumed.	None
discretization_method		str ("euler" or "zoh") Method to discretize the continuous-time plant. Default is "euler".	'euler'
discretized_noise		bool Whether the user is directly providing Gd, Qd and Rd. Default is False. If True, G, Q, and R are assumed to be Gd, Qd, and Rd, respectively.	False

Source code in jaxonomy/library/state_estimators/kalman_filter.py

@staticmethod
@with_resolved_parameters
def for_continuous_plant(
    plant,
    x_eq,
    u_eq,
    dt,
    Q=None,
    R=None,
    G=None,
    x_hat_bar_0=None,
    P_hat_bar_0=None,
    discretization_method="euler",
    discretized_noise=False,
    name=None,
    ui_id=None,
):
    """
    Obtain a Kalman Filter system for a continuous-time plant after linearization
    at equilibrium point (x_eq, u_eq)

    The input plant contains the deterministic forms of the forward and observation
    operators:

    ```
        dx/dt = f(x,u)
        y = g(x,u)
    ```

    Note: (i) Only plants with one vector-valued input and one vector-valued output
    are currently supported. Furthermore, the plant LeafSystem/Diagram should have
    only one vector-valued integrator.

    A plant with disturbances of the following form is then considered
    following form:

    ```
        dx/dt = f(x,u) + G w                        --- (C1)
        y = g(x,u) +  v                             --- (C2)
    ```

    where:

        `w` represents the process noise,
        `v` represents the measurement noise,

    and

    ```
        E(w) = E(v) = 0
        E(ww') = Q
        E(vv') = R
        E(wv') = N = 0
    ```

    This plant with disturbances is linearized (only `f` and `g`) around the
    equilibrium point to obtain:

    ```
        d/dt (x_bar) = A x_bar + B u_bar + G w
        y_bar = C x_bar + D u_bar + v
    ```

    where,

    ```
        x_bar = x - x_eq
        u_bar = u - u_eq
        y_bar = y - y_bar
        y_eq = g(x_eq, u_eq)
    ```

    The linearized plant is then discretized via `euler` or `zoh` method to obtain:

    ```
        x_bar[n] = Ad x_bar[n] + Bd u_bar[n] + Gd w[n]           --- (L1)
        y_bar[n] = Cd x_bar[n] + Dd u_bar[n] + v[n]              --- (L2)

        E(w[n]) = E(v[n]) = 0
        E(w[n]w'[n]) = Qd
        E(v[n]v'[n]) = Rd
        E(w[n]v'[n]) = Nd = 0
    ```

    Note: If `discretized_noise` is True, then it is assumed that the user is
    providing Gd, Qd and Rd. If False, then Qd and Rd are computed from
    continuous-time Q, R, and G, and Gd is set to Identity matrix.

    A Kalman Filter estimator for the system of equations (L1) and (L2) is
    created and returned. This filter is in the `x_bar`, `u_bar`, and `y_bar`
    states.

    This returned system will have

    Input ports:
        (0) u_bar[n] : control vector at timestep n, relative to equilibrium
        (1) y_bar[n] : measurement vector at timestep n, relative to equilibrium

    Output ports:
        (1) x_hat_bar[n] : state vector estimate at timestep n, relative to
                           equilibrium

    Parameters:
        plant : a `Plant` object which can be a LeafSystem or a Diagram.
        x_eq: ndarray
            Equilibrium state vector for discretization
        u_eq: ndarray
            Equilibrium control vector for discretization
        dt: float
            Time step for the discretization.
        Q: ndarray
            Process noise covariance matrix. If `None`, Identity matrix of size
            compatible with `G` and and linearized system's `A` is assumed.
        R: ndarray
            Measurement noise covariance matrix. If `None`, Identity matrix of size
            compatible with linearized system's `C` and `A` is assumed.
        G: ndarray
            Process noise matrix. If `None`, `G=B` is assumed making disrurbances
            additive to control vector `u`, i.e. `u_disturbed = u_orig + w`.
        x_hat_bar_0: ndarray
            Initial state estimate, relative to equilirium.
            If None, an identity matrix is assumed.
        P_hat_bar_0: ndarray
            Initial covariance matrix estimate for state, relative to equilibrium.
            If `None`, an Identity matrix is assumed.
        discretization_method: str ("euler" or "zoh")
            Method to discretize the continuous-time plant. Default is "euler".
        discretized_noise: bool
            Whether the user is directly providing Gd, Qd and Rd. Default is False.
            If True, `G`, `Q`, and `R` are assumed to be Gd, Qd, and Rd,
            respectively.
    """
    (
        y_eq,
        Ad,
        Bd,
        Cd,
        Dd,
        Gd,
        Qd,
        Rd,
    ) = linearize_and_discretize_continuous_plant(
        plant, x_eq, u_eq, dt, Q, R, G, discretization_method, discretized_noise
    )

    check_shape_compatibilities(Ad, Bd, Cd, Dd, Gd, Qd, Rd)

    nx = x_eq.size

    if x_hat_bar_0 is None:
        x_hat_bar_0 = jnp.zeros(nx)

    if P_hat_bar_0 is None:
        P_hat_bar_0 = jnp.eye(nx)

    # Instantiate a Kalman Filter for the linearized plant
    kf = KalmanFilter(
        dt,
        Ad,
        Bd,
        Cd,
        Dd,
        Gd,
        Qd,
        Rd,
        x_hat_bar_0,
        P_hat_bar_0,
        name=name,
        ui_id=ui_id,
    )

    return y_eq, kf

global_filter_for_continuous_plant(plant, x_eq, u_eq, dt, Q=None, R=None, G=None, x_hat_0=None, P_hat_0=None, discretization_method='euler', discretized_noise=False, name=None, ui_id=None) staticmethod

See docs for for_continuous_plant, which returns the local Kalman Filter. This method additionally converts the local Kalman Filter to a global estimator. See docs for make_global_estimator_from_local for details.

Source code in jaxonomy/library/state_estimators/kalman_filter.py

@staticmethod
@with_resolved_parameters
def global_filter_for_continuous_plant(
    plant,
    x_eq,
    u_eq,
    dt,
    Q=None,
    R=None,
    G=None,
    x_hat_0=None,
    P_hat_0=None,
    discretization_method="euler",
    discretized_noise=False,
    name=None,
    ui_id=None,
):
    """
    See docs for `for_continuous_plant`, which returns the local Kalman
    Filter. This method additionally converts the local Kalman Filter to a
    global estimator. See docs for `make_global_estimator_from_local` for details.
    """
    (
        y_eq,
        Ad,
        Bd,
        Cd,
        Dd,
        Gd,
        Qd,
        Rd,
    ) = linearize_and_discretize_continuous_plant(
        plant, x_eq, u_eq, dt, Q, R, G, discretization_method, discretized_noise
    )

    check_shape_compatibilities(Ad, Bd, Cd, Dd, Gd, Qd, Rd)

    nx = x_eq.size

    if x_hat_0 is None:
        x_hat_bar_0 = jnp.zeros(nx)
    else:
        x_hat_bar_0 = x_hat_0 - x_eq

    if P_hat_0 is None:
        P_hat_bar_0 = jnp.eye(nx)
    else:
        P_hat_bar_0 = P_hat_0

    # Instantiate a Kalman Filter for the linearized plant
    local_kf = KalmanFilter(
        dt,
        Ad,
        Bd,
        Cd,
        Dd,
        Gd,
        Qd,
        Rd,
        x_hat_bar_0,
        P_hat_bar_0,
        name=name + "_local" if name is not None else None,
    )

    global_kf = make_global_estimator_from_local(
        local_kf,
        x_eq,
        u_eq,
        y_eq,
        name=name,
        ui_id=ui_id,
    )

    return global_kf

KoopmanPredictor

Bases: LeafSystem

Discrete-time Koopman predictor for a nonlinear system.

Each step: lift the stored physical state z = g(x), advance the lifted linear dynamics z[k+1] = K z[k] (+ B u[k]), then de-lift x[k+1] = C z[k+1]. The physical state is the block output.

Because the lifted model is linear in z, (K, B) is a discrete linear model you can use for linear MPC / LQR-style control — design in lifted coordinates and de-lift with C (the Koopman→linear-MPC framing; Korda & Mezić 2018). Note a lifted model generally has no state that both satisfies z = g(x) and a hard terminal-equality constraint, so a terminal-cost MPC is the right fit; the current :class:~jaxonomy.library.mpc.LinearDiscreteTimeMPC block (hard terminal equality, continuous-time model input) does not compose directly — see the rom_dmdc_koopman_mpc example, which uses a compact terminal-cost MPC.

An input port (and use of B) is created only when B is provided.

Input ports

(0) u[k]: control input, present iff B is given.

Output ports

(0) x[k]: de-lifted physical state.

Parameters:

Name	Type	Description	Default
K		Koopman operator (L, L) — a dynamic parameter.	required
C		De-lift matrix (n, L) — a dynamic parameter.	required
dictionary		Observable dictionary g used for lifting (identity first).	required
B		Optional lifted input operator (L, m) — a dynamic parameter when given.	None
dt		Sampling period of the discrete update.	1.0
initial_state		Initial physical state x[0] of size n.	None

Source code in jaxonomy/library/rom/koopman.py

class KoopmanPredictor(LeafSystem):
    """Discrete-time Koopman predictor for a nonlinear system.

    Each step: lift the stored physical state ``z = g(x)``, advance the lifted
    *linear* dynamics ``z[k+1] = K z[k] (+ B u[k])``, then de-lift
    ``x[k+1] = C z[k+1]``. The physical state is the block output.

    Because the lifted model is *linear* in ``z``, ``(K, B)`` is a discrete
    linear model you can use for linear MPC / LQR-style control — design in
    lifted coordinates and de-lift with ``C`` (the Koopman→linear-MPC framing;
    Korda & Mezić 2018). Note a lifted model generally has no state that both
    satisfies ``z = g(x)`` and a hard terminal-equality constraint, so a
    *terminal-cost* MPC is the right fit; the current
    :class:`~jaxonomy.library.mpc.LinearDiscreteTimeMPC` block (hard terminal
    equality, continuous-time model input) does not compose directly — see the
    ``rom_dmdc_koopman_mpc`` example, which uses a compact terminal-cost MPC.

    An input port (and use of ``B``) is created only when ``B`` is provided.

    Input ports:
        (0) u[k]: control input, present iff ``B`` is given.

    Output ports:
        (0) x[k]: de-lifted physical state.

    Parameters:
        K: Koopman operator ``(L, L)`` — a ``dynamic`` parameter.
        C: De-lift matrix ``(n, L)`` — a ``dynamic`` parameter.
        dictionary: Observable dictionary ``g`` used for lifting (identity first).
        B: Optional lifted input operator ``(L, m)`` — a ``dynamic`` parameter when given.
        dt: Sampling period of the discrete update.
        initial_state: Initial physical state ``x[0]`` of size ``n``.
    """

    @parameters(dynamic=["K", "C", "B"], static=["dt", "initial_state"])
    def __init__(self, K, C, dictionary, B=None, dt=1.0, initial_state=None,
                 name=None, **kwargs):
        super().__init__(name=name, **kwargs)

        C = np.asarray(C, dtype=float)
        if C.ndim == 1:
            C = C.reshape(1, -1)
        self.n = C.shape[0]
        self.dictionary = dictionary
        self.has_input = B is not None
        self.dt = dt

        if initial_state is None:
            initial_state = np.zeros(self.n)
        self._x0 = np.asarray(initial_state, dtype=float).reshape(-1)

        if self.has_input:
            self.declare_input_port(name="u")

        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port(name="out_0")

    def initialize(self, K, C, dictionary=None, B=None, dt=1.0, initial_state=None,
                   **kwargs):
        if initial_state is None:
            x0 = npa.array(self._x0)
        else:
            x0 = npa.reshape(npa.array(initial_state, dtype=npa.float64), (-1,))

        self.declare_discrete_state(default_value=x0)
        self.configure_periodic_update(
            self._periodic_update_idx, self._update, period=self.dt, offset=0.0
        )
        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=0.0,
            default_value=npa.zeros(self.n) if self.n > 1 else 0.0,
            requires_inputs=False,
            prerequisites_of_calc=[DependencyTicket.xd],
        )

    def _update(self, _time, state, *inputs, **params):
        x = state.discrete_state
        z = self.dictionary(x)
        z_next = params["K"] @ z
        if self.has_input:
            u = jnp.atleast_1d(inputs[0])
            z_next = z_next + params["B"] @ u
        x_next = params["C"] @ z_next
        return x_next

    def _output(self, _time, state, *_inputs, **_params):
        x = state.discrete_state
        if self.n == 1:
            return jnp.atleast_1d(x)[0]
        return x

LTISystem

Bases: LTISystemBase

Continuous-time linear time-invariant system.

Implements the following system of ODEs:

    ẋ = Ax + Bu
    y = Cx + Du

Input ports

(0) u: Input vector of size m

Output ports

(0) y: Output vector of size p. Note that this is feedthrough from the input port if and only if D is nonzero.

Parameters:

Name	Type	Description	Default
A		State matrix of size n x n	required
B		Input matrix of size n x m	required
C		Output matrix of size p x n	required
D		Feedthrough matrix of size p x m	required
initialize_states		Initial state vector of size n (default: 0)	None

Source code in jaxonomy/library/linear_system.py

class LTISystem(LTISystemBase):
    """Continuous-time linear time-invariant system.

    Implements the following system of ODEs:
    ```
        ẋ = Ax + Bu
        y = Cx + Du
    ```

    Input ports:
        (0) u: Input vector of size m

    Output ports:
        (0) y: Output vector of size p.  Note that this is feedthrough from the input
            port if and only if D is nonzero.

    Parameters:
        A: State matrix of size n x n
        B: Input matrix of size n x m
        C: Output matrix of size p x n
        D: Feedthrough matrix of size p x m
        initialize_states: Initial state vector of size n (default: 0)
    """

    @parameters(dynamic=["A", "B", "C", "D"], static=["initialize_states"])
    def __init__(self, A, B, C, D, initialize_states=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._output_port_idx = self.declare_output_port(
            self._eval_output
        )  # Single output port (y)
        self._continuous_state_id = (
            self.declare_continuous_state()
        )  # Single continuous state (x)
        # Expose the state-space matrices immediately. They are otherwise only
        # populated in initialize() (at context-creation time), which made
        # ``linearize(...).to_lti().A`` raise AttributeError before simulation;
        # initialize() re-derives them from the resolved parameters.
        (self.A, self.B, self.C, self.D, self.n, self.m, self.p) = _reshape(A, B, C, D)

    def _init_state(self, A, B, C, D, initialize_states=None):
        super()._init_state(A, B, C, D, initialize_states)
        self.configure_output_port(
            self._output_port_idx,
            self._eval_output,
            default_value=npa.zeros(self.p) if self.p > 1 else 0.0,
            requires_inputs=self.is_feedthrough,
        )
        self.configure_continuous_state(
            self._continuous_state_id,
            ode=self.ode,
            default_value=self.initialize_states,
        )

    def initialize(self, A, B, C, D, initialize_states=None, **kwargs):
        self._init_state(A, B, C, D, initialize_states)
        self.parameters["A"].set(self.A)
        self.parameters["B"].set(self.B)
        self.parameters["C"].set(self.C)
        self.parameters["D"].set(self.D)

    def _eval_output(self, time, state, *inputs, **params):
        return self._eval_output_base(params["C"], params["D"], state, *inputs)

    def _eval_output_base(self, C, D, state, *inputs):
        x = state.continuous_state
        y = npa.matmul(C, npa.atleast_1d(x))

        if self.is_feedthrough:
            (u,) = inputs
            y += npa.matmul(D, npa.atleast_1d(u))

        # Handle the special case of scalar output
        if self.scalar_output:
            y = npa.atleast_1d(y)[0]

        return y

    def ode(self, time, state, u, **params):
        x = state.continuous_state
        A, B = params["A"], params["B"]
        Ax = npa.matmul(A, npa.atleast_1d(x))
        Bu = npa.matmul(B, npa.atleast_1d(u))
        return Ax + Bu

    @property
    def ss(self):
        """State-space representation of the system."""
        return control.ss(self.A, self.B, self.C, self.D)

ss property

State-space representation of the system.

LTISystemDiscrete

Bases: LTISystemBase

Discrete-time linear time-invariant system.

Implements the following system of ODEs:

    x[k+1] = A x[k] + B u[k]
    y[k] = C x[k] + D u[k]

Input ports

(0) u[k]: Input vector of size m

Output ports

(0) y[k]: Output vector of size p. Note that this is feedthrough from the input port if and only if D is nonzero.

Parameters:

Name	Type	Description	Default
A		State matrix of size n x n	required
B		Input matrix of size n x m	required
C		Output matrix of size p x n	required
D		Feedthrough matrix of size p x m	required
dt		Sampling period	required
initialize_states		Initial state vector of size n (default: 0)	None

Source code in jaxonomy/library/linear_system.py

class LTISystemDiscrete(LTISystemBase):
    """Discrete-time linear time-invariant system.

    Implements the following system of ODEs:
    ```
        x[k+1] = A x[k] + B u[k]
        y[k] = C x[k] + D u[k]
    ```

    Input ports:
        (0) u[k]: Input vector of size m

    Output ports:
        (0) y[k]: Output vector of size p.  Note that this is feedthrough from the
                  input port if and only if D is nonzero.

    Parameters:
        A: State matrix of size n x n
        B: Input matrix of size n x m
        C: Output matrix of size p x n
        D: Feedthrough matrix of size p x m
        dt: Sampling period
        initialize_states: Initial state vector of size n (default: 0)
    """

    @parameters(dynamic=["A", "B", "C", "D"], static=["initialize_states"])
    def __init__(self, A, B, C, D, dt, initialize_states=None, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.dt = dt
        self.declare_periodic_update(
            self._update,
            period=dt,
            offset=0.0,
        )

        self._output_port_idx = self.declare_output_port(
            self._eval_output
        )  # Single output port (y)

    def _init_state(self, A, B, C, D, initialize_states=None):
        super()._init_state(A, B, C, D, initialize_states)
        self.declare_discrete_state(
            default_value=self.initialize_states,
        )  # Single discrete state (x)
        self.configure_output_port(
            self._output_port_idx,
            self._eval_output,
            period=self.dt,
            offset=0.0,
            default_value=npa.zeros(self.p) if self.p > 1 else 0.0,
            requires_inputs=self.is_feedthrough,
        )

    def initialize(self, A, B, C, D, initialize_states=None, **kwargs):
        self._init_state(A, B, C, D, initialize_states)
        self.parameters["A"].set(self.A)
        self.parameters["B"].set(self.B)
        self.parameters["C"].set(self.C)
        self.parameters["D"].set(self.D)

    def _eval_output(self, time, state, *inputs, **params):
        x = state.discrete_state
        self.C, self.D = params["C"], params["D"]
        y = npa.matmul(self.C, npa.atleast_1d(x))

        if self.is_feedthrough:
            (u,) = inputs
            y += npa.matmul(self.D, npa.atleast_1d(u))

        # Handle the special case of scalar output
        if self.scalar_output:
            y = y[0]

        return y

    def _update(self, time, state, u, **params):
        x = state.discrete_state
        self.A, self.B = params["A"], params["B"]
        Ax = npa.matmul(self.A, npa.atleast_1d(x))
        Bu = npa.matmul(self.B, npa.atleast_1d(u))
        return Ax + Bu

    @property
    def ss(self):
        """State-space representation of the system."""
        return control.ss(self.A, self.B, self.C, self.D, self.dt)

ss property

State-space representation of the system.

LeadLag

Bases: LeafSystem

Discrete first-order lead-lag compensator.

Discretisation of the continuous compensator G(s) = K * (1 + T_lead * s) / (1 + T_lag * s) via the Tustin (bilinear) transform with s = (2/dt)*(z-1)/(z+1).

The resulting difference equation is::

c   = 2 / dt
den = 1 + T_lag * c
b0  = K * (1 + T_lead * c) / den
b1  = K * (1 - T_lead * c) / den
a1  = (1 - T_lag * c) / den
y[k] = b0 * x[k] + b1 * x[k-1] - a1 * y[k-1]

With T_lead = T_lag the s-domain pole and zero cancel and the block reduces to a pure gain K (used as an identity check in the corpus).

Input ports

(0) The input signal x.

Output ports

(0) The compensated signal y.

Parameters:

Name	Type	Description	Default
dt		Sampling period of the block (s).	required
K		Compensator gain. Differentiable.	1.0
T_lead		Lead time constant (s). Differentiable.	1.0
T_lag		Lag time constant (s). Must be > 0. Differentiable.	1.0
initial_state		Initial value of y[-1]. Default 0.0.	0.0

Notes

Differentiability: K, T_lead and T_lag flow into the biquad coefficients via smooth arithmetic, so jax.grad is finite through them.

The block is feedthrough on its input port (b0 != 0 whenever K != 0), so y[k] depends on x[k] directly.

Source code in jaxonomy/library/dynamics.py

class LeadLag(LeafSystem):
    """Discrete first-order lead-lag compensator.

    Discretisation of the continuous compensator
    ``G(s) = K * (1 + T_lead * s) / (1 + T_lag * s)``
    via the Tustin (bilinear) transform with ``s = (2/dt)*(z-1)/(z+1)``.

    The resulting difference equation is::

        c   = 2 / dt
        den = 1 + T_lag * c
        b0  = K * (1 + T_lead * c) / den
        b1  = K * (1 - T_lead * c) / den
        a1  = (1 - T_lag * c) / den
        y[k] = b0 * x[k] + b1 * x[k-1] - a1 * y[k-1]

    With ``T_lead = T_lag`` the s-domain pole and zero cancel and the
    block reduces to a pure gain ``K`` (used as an identity check in
    the corpus).

    Input ports:
        (0) The input signal ``x``.

    Output ports:
        (0) The compensated signal ``y``.

    Parameters:
        dt:
            Sampling period of the block (s).
        K:
            Compensator gain.  Differentiable.
        T_lead:
            Lead time constant (s).  Differentiable.
        T_lag:
            Lag time constant (s).  Must be > 0.  Differentiable.
        initial_state:
            Initial value of ``y[-1]``.  Default 0.0.

    Notes:
        Differentiability: ``K``, ``T_lead`` and ``T_lag`` flow into the
        biquad coefficients via smooth arithmetic, so ``jax.grad`` is
        finite through them.

        The block is feedthrough on its input port (``b0 != 0`` whenever
        ``K != 0``), so ``y[k]`` depends on ``x[k]`` directly.
    """

    class DiscreteStateType(NamedTuple):
        x_prev: Array
        y_prev: Array

    @parameters(
        static=["dt"],
        dynamic=["K", "T_lead", "T_lag", "initial_state"],
    )
    def __init__(
        self,
        dt,
        K=1.0,
        T_lead=1.0,
        T_lag=1.0,
        initial_state=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.dt = dt
        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, K, T_lead, T_lag, initial_state, dt=None):
        y0 = npa.asarray(initial_state)
        x0 = npa.zeros_like(y0)
        self.declare_discrete_state(
            default_value=self.DiscreteStateType(x_prev=x0, y_prev=y0),
            as_array=False,
        )

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=self.dt,
        )

        # Feedthrough: y[k] depends on x[k] through b0.
        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=self.dt,
            default_value=y0,
            requires_inputs=True,
            prerequisites_of_calc=[
                DependencyTicket.xd,
                self.input_ports[0].ticket,
            ],
        )

    def reset_default_values(self, **dynamic_parameters):
        y0 = npa.asarray(dynamic_parameters["initial_state"])
        x0 = npa.zeros_like(y0)
        self.configure_discrete_state_default_value(
            self.DiscreteStateType(x_prev=x0, y_prev=y0),
            as_array=False,
        )
        self.configure_output_port_default_value(self._output_port_idx, y0)

    def _coeffs(self, K, T_lead, T_lag):
        # Bilinear-transform biquad coefficients.  Pure arithmetic on
        # the dynamic parameters → JAX-traceable and differentiable.
        c = 2.0 / self.dt
        den = 1.0 + T_lag * c
        b0 = K * (1.0 + T_lead * c) / den
        b1 = K * (1.0 - T_lead * c) / den
        a1 = (1.0 - T_lag * c) / den
        return b0, b1, a1

    def _update(self, _time, state, *inputs, **params):
        x = inputs[0]
        b0, b1, a1 = self._coeffs(params["K"], params["T_lead"], params["T_lag"])
        y_prev = state.discrete_state.y_prev
        x_prev = state.discrete_state.x_prev
        y_new = b0 * x + b1 * x_prev - a1 * y_prev
        return self.DiscreteStateType(x_prev=x, y_prev=y_new)

    def _output(self, _time, state, *inputs, **params):
        # Feedthrough output: recompute y[k] from x[k] and the stored
        # x[k-1], y[k-1] so the readout is consistent with the update.
        x = inputs[0]
        b0, b1, a1 = self._coeffs(params["K"], params["T_lead"], params["T_lag"])
        y_prev = state.discrete_state.y_prev
        x_prev = state.discrete_state.x_prev
        return b0 * x + b1 * x_prev - a1 * y_prev

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        u = self.eval_input(context)
        xd = context[self.system_id].discrete_state.y_prev
        check_state_type(
            self,
            inp_data=u,
            state_data=xd,
            error_collector=error_collector,
        )

LinearDiscreteTimeMPC

Bases: LeafSystem

Model predictive control for a linear discrete-time system.

Solves a constrained quadratic program at each time step using OSQP via jax.pure_callback, making it compatible with JAX's JIT compiler.

Notes

This block is feedthrough: the QP solver runs every time the output port is evaluated. Pair with a zero-order hold so the solver is invoked only once per MPC step.

Parameters:

Name	Type	Description	Default
lin_sys		Linearized system (continuous-time A, B matrices).	required
Q		State cost matrix (n×n).	required
R		Input cost matrix (m×m).	required
N		Prediction horizon (number of steps).	required
dt		Sampling period for Euler discretization.	required
x_ref		Terminal state reference (length-n array).	required
lbu		Lower bound on control input (scalar or length-m array).	-inf
ubu		Upper bound on control input (scalar or length-m array).	inf
warm_start		Whether to warm-start the OSQP solver between solves.	False

Source code in jaxonomy/library/mpc.py

class LinearDiscreteTimeMPC(LeafSystem):
    """Model predictive control for a linear discrete-time system.

    Solves a constrained quadratic program at each time step using OSQP
    via ``jax.pure_callback``, making it compatible with JAX's JIT compiler.

    Notes:
        This block is *feedthrough*: the QP solver runs every time the output
        port is evaluated.  Pair with a zero-order hold so the solver is
        invoked only once per MPC step.

    Args:
        lin_sys: Linearized system (continuous-time A, B matrices).
        Q: State cost matrix (n×n).
        R: Input cost matrix (m×m).
        N: Prediction horizon (number of steps).
        dt: Sampling period for Euler discretization.
        x_ref: Terminal state reference (length-n array).
        lbu: Lower bound on control input (scalar or length-m array).
        ubu: Upper bound on control input (scalar or length-m array).
        warm_start: Whether to warm-start the OSQP solver between solves.
    """

    def __init__(
        self,
        lin_sys,
        Q,
        R,
        N,
        dt,
        x_ref,
        lbu=-np.inf,
        ubu=np.inf,
        name=None,
        warm_start=False,
    ):
        super().__init__(name=name)
        lin_sys.create_context()
        self.n = lin_sys.A.shape[0]
        self.m = lin_sys.B.shape[1]
        self.N = N
        self.warm_start = warm_start

        # Euler discretization
        A = jnp.eye(self.n) + dt * lin_sys.A
        B = dt * lin_sys.B

        self.declare_input_port()

        # Shape of the full primal solution vector x = [x_0, u_0, ..., x_{N-1}, u_{N-1}]
        self._result_template = jnp.zeros((self.n + self.m) * self.N)

        self._make_solver(A, B, Q, R, lbu, ubu, N, x_ref)

        # Wrap the non-JAX OSQP solve in a pure_callback so it is JIT-compatible
        self._jax_solve = partial(
            jax.pure_callback, self._np_solve, self._result_template
        )

        self.declare_output_port(
            self._output,
            requires_inputs=True,
            period=dt,
            offset=0.0,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _make_solver(self, A, B, Q, R, lbu, ubu, N, xf):
        from scipy import sparse

        n, m = self.n, self.m
        I_A = jnp.eye(n)
        I_B = jnp.eye(m)

        def e(k):
            return jnp.zeros(N).at[k].set(1.0)

        # Block-diagonal cost matrix P = diag(Q, R, Q, R, ...) of size (n+m)*N
        P_dense = linalg.block_diag(*([Q, R] * N))

        # Equality / dynamics constraints
        L0 = jnp.eye(n, N * (n + m))
        L_defect = jnp.vstack(
            [
                jnp.kron(e(k), jnp.hstack([A, B]))
                + jnp.kron(e(k + 1), jnp.hstack([-I_A, 0 * B]))
                for k in range(N - 1)
            ]
        )
        Lf = jnp.kron(e(N - 1), jnp.hstack([I_A, 0 * B]))
        L_input = jnp.vstack(
            [jnp.kron(e(k), jnp.hstack([0 * B.T, I_B])) for k in range(N)]
        )
        L_dense = jnp.vstack([L0, L_defect, Lf, L_input])

        self._L_defect_rows = L_defect.shape[0]
        self._xf = xf
        self._lbu = lbu
        self._ubu = ubu

        # Precompute JIT-compiled bounds helper (used inside pure_callback)
        def _get_bounds(x0):
            lb = jnp.hstack(
                [x0, jnp.zeros(L_defect.shape[0]), xf, jnp.full(N * m, lbu)]
            )
            ub = jnp.hstack(
                [x0, jnp.zeros(L_defect.shape[0]), xf, jnp.full(N * m, ubu)]
            )
            return lb, ub

        self._get_bounds = jax.jit(_get_bounds)

        # Set up OSQP with dummy initial bounds; updated before each solve
        lb0, ub0 = _get_bounds(jnp.zeros(n))
        self.solver = osqp.OSQP()
        self.solver.setup(
            P=sparse.csc_matrix(np.array(P_dense)),
            q=np.zeros(N * (n + m)),
            A=sparse.csc_matrix(np.array(L_dense)),
            l=np.array(lb0),
            u=np.array(ub0),
            verbose=False,
            **{_osqp_warm_start_kwarg(): self.warm_start},
        )

    def _np_solve(self, time, state, x0):
        """Non-JAX solve called via pure_callback. Inputs are concrete numpy arrays."""
        lb, ub = self._get_bounds(x0)
        self.solver.update(l=np.array(lb), u=np.array(ub))
        sol = self.solver.solve()
        return sol.x

    def _dummy_solve(self, _time, _state, *_inputs, **_params):
        """Return inf when time is inf (minor ODE steps guarding against OSQP errors)."""
        return jnp.full(self._result_template.shape, jnp.inf)

    def _output(self, time, state, *inputs):
        args = (time, state, *inputs)
        xu_flat = cond(jnp.isinf(time), self._dummy_solve, self._jax_solve, *args)
        xu_traj = xu_flat.reshape((self.n + self.m, self.N), order="F")
        return xu_traj[self.n:, 0]

LinearDiscreteTimeMPC_OSQP

Bases: LinearDiscreteTimeMPC

Deprecated alias for :class:LinearDiscreteTimeMPC.

Both classes now use OSQP via jax.pure_callback. Use :class:LinearDiscreteTimeMPC directly.

Source code in jaxonomy/library/mpc.py

class LinearDiscreteTimeMPC_OSQP(LinearDiscreteTimeMPC):
    """Deprecated alias for :class:`LinearDiscreteTimeMPC`.

    Both classes now use OSQP via ``jax.pure_callback``.
    Use :class:`LinearDiscreteTimeMPC` directly.
    """

    def __init__(self, *args, **kwargs):
        warnings.warn(
            "LinearDiscreteTimeMPC_OSQP is deprecated and will be removed in a "
            "future release. Use LinearDiscreteTimeMPC instead — both classes now "
            "use OSQP via jax.pure_callback.",
            DeprecationWarning,
            stacklevel=2,
        )
        super().__init__(*args, **kwargs)

LinearQuadraticGaussian

Bases: LeafSystem

Continuous-time infinite-horizon LQG controller (separation principle).

The observer uses the algebraic Riccati solution for the Kalman gain L given (A, G, C, Qn, Rn); the regulator uses the algebraic Riccati solution for the feedback gain K given (A, B, Qc, Rc). Both use the control library's lqe / lqr helpers.

Input / output shapes follow the plant: y is (ny,), u is (nu,), and the observer's internal state is (nx,).

Source code in jaxonomy/library/lqg.py

class LinearQuadraticGaussian(LeafSystem):
    """Continuous-time infinite-horizon LQG controller (separation principle).

    The observer uses the algebraic Riccati solution for the Kalman gain
    ``L`` given ``(A, G, C, Qn, Rn)``; the regulator uses the algebraic
    Riccati solution for the feedback gain ``K`` given ``(A, B, Qc, Rc)``.
    Both use the ``control`` library's ``lqe`` / ``lqr`` helpers.

    Input / output shapes follow the plant: ``y`` is ``(ny,)``, ``u`` is
    ``(nu,)``, and the observer's internal state is ``(nx,)``.
    """

    def __init__(
        self,
        A: np.ndarray,
        B: np.ndarray,
        C: np.ndarray,
        D: np.ndarray,
        Qn: np.ndarray,
        Rn: np.ndarray,
        Qc: np.ndarray,
        Rc: np.ndarray,
        G: Optional[np.ndarray] = None,
        x_hat_0: Optional[np.ndarray] = None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

        A = np.asarray(A)
        B = np.asarray(B)
        C = np.asarray(C)
        D = np.asarray(D)
        nx, nu = B.shape
        ny = C.shape[0]
        if G is None:
            G = np.eye(nx)
        G = np.asarray(G)

        # Observer gain from LQE (Kalman filter at steady state).
        L, _P_obs, _E_obs = control.lqe(A, G, C, Qn, Rn)

        # Regulator gain from LQR.
        K, _P_reg, _E_reg = control.lqr(A, B, Qc, Rc)

        self.A = A
        self.B = B
        self.C = C
        self.D = D
        self.L = np.asarray(L)
        self.K = np.asarray(K)

        # Pre-compute observer closed-loop matrices for cheaper ODE eval:
        #     dx̂/dt = (A − B·K − L·C) x̂ + L·y     (u = −K·x̂)
        # Rearranging lets us eval u without re-forming it.
        self.A_obs = A - B @ self.K - self.L @ C
        self.B_obs = self.L  # multiplies y

        self.nx = nx
        self.nu = nu
        self.ny = ny

        if x_hat_0 is None:
            x_hat_0 = np.zeros(nx)
        x_hat_0 = np.asarray(x_hat_0)

        # I/O: one input (y), one output (u).  The observer state is the
        # block's continuous state.
        self.declare_input_port(name="y")
        self.declare_continuous_state(
            ode=self._ode, shape=x_hat_0.shape, default_value=x_hat_0, as_array=True,
        )
        self.declare_output_port(
            self._compute_u,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            name="u",
        )

    def _ode(self, time, state, *inputs, **params):
        y = npa.atleast_1d(inputs[0])
        x_hat = state.continuous_state
        return npa.dot(self.A_obs, x_hat) + npa.dot(self.B_obs, y)

    def _compute_u(self, time, state, *inputs, **params):
        x_hat = state.continuous_state
        return -npa.dot(self.K, x_hat)

LinearQuadraticRegulator

Bases: FeedthroughBlock

Linear Quadratic Regulator (LQR) for a continuous-time system: dx/dt = A x + B u. Computes the optimal control input: u = -K x, where u minimises the cost function over [0, ∞)]: J = ∫(x.T Q x + u.T R u) dt.

Input ports

(0) x: state vector of the system.

Output ports

(0) u: optimal control vector.

Parameters:

Name	Type	Description	Default
A		Array State matrix of the system.	required
B		Array Input matrix of the system.	required
Q		Array State cost matrix.	required
R		Array Input cost matrix.	required

Source code in jaxonomy/library/lqr.py

class LinearQuadraticRegulator(FeedthroughBlock):
    """
    Linear Quadratic Regulator (LQR) for a continuous-time system:
            dx/dt = A x + B u.
    Computes the optimal control input:
            u = -K x,
    where u minimises the cost function over [0, ∞)]:
            J = ∫(x.T Q x + u.T R u) dt.

    Input ports:
        (0) x: state vector of the system.

    Output ports:
        (0) u: optimal control vector.

    Parameters:
        A: Array
            State matrix of the system.
        B: Array
            Input matrix of the system.
        Q: Array
            State cost matrix.
        R: Array
            Input cost matrix.
    """

    def __init__(self, A, B, Q, R, *args, **kwargs):
        self.K, S, E = control.lqr(A, B, Q, R)
        super().__init__(lambda x: jnp.matmul(-self.K, x), *args, **kwargs)

LinearizedSystem dataclass

State-space linearization result.

For a continuous-time linsys (dt is None):

dx/dt = Ax + Bu,    y = Cx + Du

For a discrete-time linsys (dt is a positive float — produced by :func:jaxonomy.library.linearization_workflow.discretize):

x[k+1] = Ax[k] + Bu[k],    y[k] = Cx[k] + Du[k]

Attributes:

Name	Type	Description
A	Any	State matrix (n_states, n_states)
B	Any	Input matrix (n_states, n_inputs)
C	Any	Output matrix (n_outputs, n_states)
D	Any	Feedthrough matrix (n_outputs, n_inputs)
operating_point	dict	dict with state and input values used for linearization
dt	Optional[float]	Sampling period in seconds when the linsys is discrete-time; None for continuous-time. Default None so existing constructor calls remain byte-equivalent.

Source code in jaxonomy/library/linear_system.py

@dataclass
class LinearizedSystem:
    """
    State-space linearization result.

    For a continuous-time linsys (``dt is None``):

        dx/dt = Ax + Bu,    y = Cx + Du

    For a discrete-time linsys (``dt`` is a positive float — produced by
    :func:`jaxonomy.library.linearization_workflow.discretize`):

        x[k+1] = Ax[k] + Bu[k],    y[k] = Cx[k] + Du[k]

    Attributes:
        A: State matrix (n_states, n_states)
        B: Input matrix (n_states, n_inputs)
        C: Output matrix (n_outputs, n_states)
        D: Feedthrough matrix (n_outputs, n_inputs)
        operating_point: dict with state and input values
                         used for linearization
        dt: Sampling period in seconds when the linsys is discrete-time;
            ``None`` for continuous-time. Default ``None`` so existing
            constructor calls remain byte-equivalent.
    """
    A: Any   # jax.Array
    B: Any
    C: Any
    D: Any
    operating_point: dict
    dt: Optional[float] = None

    def is_discrete(self) -> bool:
        """True if this LinearizedSystem carries a sampling period."""
        return self.dt is not None

    def to_lti(self) -> "LTISystem":
        """Convert to Jaxonomy LTISystem block."""
        return LTISystem(
            A=self.A, B=self.B, C=self.C, D=self.D
        )

    def eigenvalues(self):
        """
        Compute eigenvalues of A matrix.
        Returns complex array of shape (n_states,).
        """
        import jax.numpy as jnp
        return jnp.linalg.eigvals(self.A)

    def is_stable(self) -> bool:
        """
        True if the system is stable.  For continuous-time linsys
        (``dt is None``) the criterion is ``Re(eig(A)) < 0``; for a
        discrete-time linsys (``dt is not None``) it is
        ``|eig(A)| < 1``.
        """
        import jax.numpy as jnp
        eigs = self.eigenvalues()
        if self.dt is None:
            return bool(jnp.all(jnp.real(eigs) < 0))
        return bool(jnp.all(jnp.abs(eigs) < 1.0))

    def to_scipy_lti(self):
        """Convert to :class:`scipy.signal.StateSpace` for frequency-domain analysis.

        Returns a ``scipy.signal.StateSpace`` object, which supports MIMO systems
        (multiple inputs and/or multiple outputs) as well as SISO systems.  Use this
        for Bode plots, Nyquist diagrams, step/impulse responses, etc.

        Note:
            ``scipy.signal.lti`` is SISO-only and is **not** used here.
            ``scipy.signal.StateSpace`` (a subclass of ``lti``) is the correct target
            for state-space (A, B, C, D) representations with any number of I/Os.

        Requires scipy to be installed.
        """
        try:
            from scipy import signal
            import numpy as np
            return signal.StateSpace(
                np.array(self.A), np.array(self.B),
                np.array(self.C), np.array(self.D),
            )
        except ImportError:
            raise ImportError(
                "scipy is required for to_scipy_lti(). "
                "Install with: pip install scipy"
            )

eigenvalues()

Compute eigenvalues of A matrix. Returns complex array of shape (n_states,).

Source code in jaxonomy/library/linear_system.py

def eigenvalues(self):
    """
    Compute eigenvalues of A matrix.
    Returns complex array of shape (n_states,).
    """
    import jax.numpy as jnp
    return jnp.linalg.eigvals(self.A)

is_discrete()

True if this LinearizedSystem carries a sampling period.

Source code in jaxonomy/library/linear_system.py

def is_discrete(self) -> bool:
    """True if this LinearizedSystem carries a sampling period."""
    return self.dt is not None

is_stable()

True if the system is stable. For continuous-time linsys (dt is None) the criterion is Re(eig(A)) < 0; for a discrete-time linsys (dt is not None) it is |eig(A)| < 1.

Source code in jaxonomy/library/linear_system.py

def is_stable(self) -> bool:
    """
    True if the system is stable.  For continuous-time linsys
    (``dt is None``) the criterion is ``Re(eig(A)) < 0``; for a
    discrete-time linsys (``dt is not None``) it is
    ``|eig(A)| < 1``.
    """
    import jax.numpy as jnp
    eigs = self.eigenvalues()
    if self.dt is None:
        return bool(jnp.all(jnp.real(eigs) < 0))
    return bool(jnp.all(jnp.abs(eigs) < 1.0))

to_lti()

Convert to Jaxonomy LTISystem block.

Source code in jaxonomy/library/linear_system.py

def to_lti(self) -> "LTISystem":
    """Convert to Jaxonomy LTISystem block."""
    return LTISystem(
        A=self.A, B=self.B, C=self.C, D=self.D
    )

to_scipy_lti()

Convert to :class:scipy.signal.StateSpace for frequency-domain analysis.

Returns a scipy.signal.StateSpace object, which supports MIMO systems (multiple inputs and/or multiple outputs) as well as SISO systems. Use this for Bode plots, Nyquist diagrams, step/impulse responses, etc.

Note

scipy.signal.lti is SISO-only and is not used here. scipy.signal.StateSpace (a subclass of lti) is the correct target for state-space (A, B, C, D) representations with any number of I/Os.

Requires scipy to be installed.

Source code in jaxonomy/library/linear_system.py

def to_scipy_lti(self):
    """Convert to :class:`scipy.signal.StateSpace` for frequency-domain analysis.

    Returns a ``scipy.signal.StateSpace`` object, which supports MIMO systems
    (multiple inputs and/or multiple outputs) as well as SISO systems.  Use this
    for Bode plots, Nyquist diagrams, step/impulse responses, etc.

    Note:
        ``scipy.signal.lti`` is SISO-only and is **not** used here.
        ``scipy.signal.StateSpace`` (a subclass of ``lti``) is the correct target
        for state-space (A, B, C, D) representations with any number of I/Os.

    Requires scipy to be installed.
    """
    try:
        from scipy import signal
        import numpy as np
        return signal.StateSpace(
            np.array(self.A), np.array(self.B),
            np.array(self.C), np.array(self.D),
        )
    except ImportError:
        raise ImportError(
            "scipy is required for to_scipy_lti(). "
            "Install with: pip install scipy"
        )

Logarithm

Bases: FeedthroughBlock

Compute the logarithm of the input signal.

This block dispatches to jax.numpy.log, jax.numpy.log2, or jax.numpy.log10, so the semantics, broadcasting rules, etc. are the same. See the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log.html https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log2.html https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log10.html

Input ports

(0) The input signal.

Output ports

(0) The logarithm of the input signal.

Parameters:

Name	Type	Description	Default
base		One of "natural", "2", or "10". Determines the base of the logarithm. The default is "natural".	'natural'

Source code in jaxonomy/library/math_ops.py

class Logarithm(FeedthroughBlock):
    """Compute the logarithm of the input signal.

    This block dispatches to `jax.numpy.log`, `jax.numpy.log2`, or `jax.numpy.log10`,
    so the semantics, broadcasting rules, etc. are the same.  See the JAX docs for
    details:
        https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log.html
        https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log2.html
        https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log10.html

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The logarithm of the input signal.

    Parameters:
        base:
            One of "natural", "2", or "10". Determines the base of the logarithm.
            The default is "natural".
    """

    @parameters(static=["base"])
    def __init__(self, base="natural", **kwargs):
        super().__init__(None, **kwargs)

    def initialize(self, base="natural"):
        func_lookup = {
            "10": npa.log10,
            "2": npa.log2,
            "natural": npa.log,
        }
        if base not in func_lookup:
            # cannot pass system=self because this error must be raised BEFORE calling super.__init__()
            # in the case of inheritting from FeedthroughBlock.
            # if we call super.__init__() first, we get missing key error for func_lookup[base].
            raise BlockParameterError(
                message=f"Logarithm block {self.name} has invalid selection {base} for 'base'. Valid selections: "
                + ", ".join([k for k in func_lookup.keys()]),
                parameter_name="base",
            )
        self.replace_op(func_lookup[base])

LogicalOperator

Bases: LeafSystem

Apply a boolean function elementwise to the input signals.

This block implements the following boolean functions

• "or": same as np.logical_or
• "and": same as np.logical_and
• "not": same as np.logical_not
• "nor": equivalent to np.logical_not(np.logical_or(in_0,in_1))
• "nand": equivalent to np.logical_not(np.logical_and(in_0,in_1))
• "xor": same as np.logical_xor

Input ports

(0,1) The input signals. If numeric, they are interpreted as boolean types (so 0 is False and any other value is True).

Output ports

(0) The result of the logical operation, a boolean-valued signal.

Parameters:

Name	Type	Description	Default
function		The boolean function to apply. One of "or", "and", "not", "nor", "nand", or "xor".	required

Events

An event is triggered when the output changes from True to False or vice versa.

Source code in jaxonomy/library/logic.py

class LogicalOperator(LeafSystem):
    """Apply a boolean function elementwise to the input signals.

    This block implements the following boolean functions:
        - "or": same as np.logical_or
        - "and": same as np.logical_and
        - "not": same as np.logical_not
        - "nor": equivalent to np.logical_not(np.logical_or(in_0,in_1))
        - "nand": equivalent to np.logical_not(np.logical_and(in_0,in_1))
        - "xor": same as np.logical_xor

    Input ports:
        (0,1) The input signals.  If numeric, they are interpreted as boolean
            types (so 0 is False and any other value is True).

    Output ports:
        (0) The result of the logical operation, a boolean-valued signal.

    Parameters:
        function:
            The boolean function to apply. One of "or", "and", "not", "nor", "nand",
            or "xor".

    Events:
        An event is triggered when the output changes from True to False or vice versa.
    """

    @parameters(static=["function"])
    def __init__(self, function, **kwargs):
        super().__init__(**kwargs)
        self.declare_input_port()
        if not function == "not":
            self.declare_input_port()
        self._output_port_idx = self.declare_output_port(
            None,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
            requires_inputs=True,
        )

    def initialize(self, function):
        self.function = function
        func_lookup = {
            "or": self._or,
            "and": self._and,
            "not": self._not,
            "xor": self._xor,
            "nor": self._nor,
            "nand": self._nand,
        }
        if function not in func_lookup:
            raise BlockParameterError(
                message=f"LogicalOperator block {self.name} has invalid selection {function} for 'function'. Valid options: "
                + ", ".join([f for f in func_lookup.keys()]),
                system=self,
            )

        if function != "not" and len(self.input_ports) < 2:
            raise BlockParameterError(
                message=f"Can't change logical operator from 'not' to {function} for block {self.name}",
                system=self,
            )

        if function == "not" and len(self.input_ports) > 1:
            raise BlockParameterError(
                message=f"Can't change logical operator from {function} to 'not' for block {self.name}",
                system=self,
            )

        self._func = func_lookup[function]

        self.configure_output_port(
            self._output_port_idx,
            self._func,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
            requires_inputs=True,
        )

    def _edge_detection(self, time, state, *inputs, **params):
        outp = self._func(time, state, *inputs, **params)
        return npa.where(outp, 1.0, -1.0)

    def _or(self, time, state, *inputs, **parameters):
        return npa.logical_or(npa.array(inputs[0]), npa.array(inputs[1]))

    def _and(self, time, state, *inputs, **parameters):
        return npa.logical_and(npa.array(inputs[0]), npa.array(inputs[1]))

    def _not(self, time, state, *inputs, **parameters):
        (x,) = inputs
        return npa.logical_not(npa.array(x))

    def _xor(self, time, state, *inputs, **parameters):
        return npa.logical_xor(npa.array(inputs[0]), npa.array(inputs[1]))

    def _nor(self, time, state, *inputs, **parameters):
        return npa.logical_not(
            npa.logical_or(npa.array(inputs[0]), npa.array(inputs[1]))
        )

    def _nand(self, time, state, *inputs, **parameters):
        return npa.logical_not(
            npa.logical_and(npa.array(inputs[0]), npa.array(inputs[1]))
        )

    def initialize_static_data(self, context):
        # Add a zero-crossing event so ODE solvers can't try to integrate
        # through a discontinuity.  For efficiency, only do this if the output
        # is fed to an ODE block
        if not self.has_zero_crossing_events and is_discontinuity(self.output_ports[0]):
            self.declare_zero_crossing(self._edge_detection, direction="crosses_zero")

        return super().initialize_static_data(context)

LogicalReduce

Bases: FeedthroughBlock

Apply a boolean reduce function to the elements of the input signal.

This block implements the following boolean functions

• "any": Output is True if any input element is True.
• "all": Output is True if all input elements are True.

Input ports

(0) The input signal. If numeric, they are interpreted as boolean types (so 0 is False and any other value is True).

Output ports

(0) The result of the logical operation, a boolean-valued signal.

Parameters:

Name	Type	Description	Default
function		The boolean function to apply. One of "any", "all".	required
axis		Axis or axes along which a logical OR/AND reduction is performed.	None

Events

An event is triggered when the output changes from True to False or vice versa.

Source code in jaxonomy/library/logic.py

class LogicalReduce(FeedthroughBlock):
    """Apply a boolean reduce function to the elements of the input signal.

    This block implements the following boolean functions:
        - "any": Output is True if any input element is True.
        - "all": Output is True if all input elements are True.

    Input ports:
        (0) The input signal.  If numeric, they are interpreted as boolean
            types (so 0 is False and any other value is True).

    Output ports:
        (0) The result of the logical operation, a boolean-valued signal.

    Parameters:
        function:
            The boolean function to apply. One of "any", "all".
        axis:
            Axis or axes along which a logical OR/AND reduction is performed.

    Events:
        An event is triggered when the output changes from True to False or vice versa.
    """

    @parameters(static=["function", "axis"])
    def __init__(self, function, axis=None, **kwargs):
        super().__init__(None, **kwargs)

    def initialize(self, function, axis=None):
        self.function = function
        self.axis = int(axis) if axis is not None else None
        func_lookup = {
            "any": self._any,
            "all": self._all,
        }
        if function not in func_lookup:
            raise BlockParameterError(
                message=f"LogicalReduce block {self.name} has invalid selection {function} for 'function'. Valid options: "
                + ", ".join([f for f in func_lookup.keys()])
            )

        self._func = func_lookup[function]
        self.replace_op(self._func)

    def _edge_detection(self, _time, _state, *inputs, **_params):
        outp = self._func(inputs)
        return npa.where(outp, 1.0, -1.0)

    def _any(self, inputs):
        return npa.any(npa.array(inputs), axis=self.axis)

    def _all(self, inputs):
        return npa.all(npa.array(inputs), axis=self.axis)

    def initialize_static_data(self, context):
        # Add a zero-crossing event so ODE solvers can't try to integrate
        # through a discontinuity.  For efficiency, only do this if the output
        # is fed to an ODE block
        if not self.has_zero_crossing_events and is_discontinuity(self.output_ports[0]):
            self.declare_zero_crossing(self._edge_detection, direction="crosses_zero")

        return super().initialize_static_data(context)

LookupTable1d

Bases: FeedthroughBlock

Interpolate the input signal into a static lookup table.

If a function y = f(x) is sampled at a set of points (x_i, y_i), then this block will interpolate the input signal x to compute the output signal y. The behavior is modeled after scipy.interpolate.interp1d but is implemented in JAX. Available interpolation modes are: - "linear": Linear interpolation using jax.interp. - "pchip": Monotone cubic Hermite (Hyman/Fritsch-Carlson). T-106 phase 1 — smooth gradients everywhere, monotone on monotone data. - "akima": Akima 1970 cubic spline (T-114 phase 2). Smoother than PCHIP on non-monotone data, less prone to overshoot than natural cubic splines. Matches scipy.interpolate.Akima1DInterpolator. - "cubic": Natural cubic spline (T-114-followup-natural-cubic- spline). C^2-continuous, second derivative zero at the boundaries — the smoothest possible C^2 interpolant. Requires at least 4 breakpoints. Matches scipy.interpolate.CubicSpline(bc_type='natural'). - "nearest": Nearest-neighbor interpolation. - "flat": Flat interpolation.

Input ports

(0) The input signal, which is used as the interpolation coordinate.

Output ports

(0) The interpolated output signal.

Parameters:

Name	Type	Description	Default
input_array		The array of input values at which the output values are provided.	required
output_array		The array of output values.	required
interpolation		One of "linear", "pchip", "nearest", or "flat". Determines the type of interpolation performed by the block.	required
extrapolation	optional, T-114 phase 1	One of "clip" (default — standard lookup-table behaviour, holds the boundary value), "linear" (extends the boundary slope past the breakpoints), or "nan" (returns NaN outside [input_array[0], input_array[-1]]). Stored outside @parameters, so this kwarg is not round-tripped through model JSON — pre-existing models reload byte-equivalently.	'clip'
dtype	optional, T-038a	If set (e.g. jnp.float32), the block's input_array and output_array are cast to this dtype on construction, so the interpolation arithmetic — and therefore the output signal — runs at this precision regardless of the global x64 setting. The default (None) preserves the pre-T-038a behavior: the arrays are stored verbatim and float64 is used under the default x64-enabled install. T-038a-followup-mixed-precision-cascade: when dtype is None and a :func:jaxonomy.precision_policy context manager is active, the block falls back to the context's dtype. Explicit dtype= always wins (explicit-over-implicit). Best-effort: this enforces dtype on the block's internal arrays and on the output of jnp.interp / jnp.argmin lookups, but downstream operations (e.g. connecting to a default-dtype block) are subject to JAX's standard promotion rules — the result of the wider arithmetic may be promoted to float64.	None

Notes

Currently restricted to 1D input and output data. This may be expanded to support multi-dimensional output arrays in the future.

Source code in jaxonomy/library/tables.py

class LookupTable1d(FeedthroughBlock):
    """Interpolate the input signal into a static lookup table.

    If a function `y = f(x)` is sampled at a set of points `(x_i, y_i)`, then this
    block will interpolate the input signal `x` to compute the output signal `y`.
    The behavior is modeled after `scipy.interpolate.interp1d` but is implemented
    in JAX.  Available interpolation modes are:
        - "linear": Linear interpolation using `jax.interp`.
        - "pchip": Monotone cubic Hermite (Hyman/Fritsch-Carlson). T-106
          phase 1 — smooth gradients everywhere, monotone on monotone data.
        - "akima": Akima 1970 cubic spline (T-114 phase 2). Smoother than
          PCHIP on non-monotone data, less prone to overshoot than
          natural cubic splines. Matches
          ``scipy.interpolate.Akima1DInterpolator``.
        - "cubic": Natural cubic spline (T-114-followup-natural-cubic-
          spline). C^2-continuous, second derivative zero at the
          boundaries — the smoothest possible C^2 interpolant. Requires
          at least 4 breakpoints. Matches
          ``scipy.interpolate.CubicSpline(bc_type='natural')``.
        - "nearest": Nearest-neighbor interpolation.
        - "flat": Flat interpolation.

    Input ports:
        (0) The input signal, which is used as the interpolation coordinate.

    Output ports:
        (0) The interpolated output signal.

    Parameters:
        input_array:
            The array of input values at which the output values are provided.
        output_array:
            The array of output values.
        interpolation:
            One of "linear", "pchip", "nearest", or "flat". Determines the type
            of interpolation performed by the block.
        extrapolation (optional, T-114 phase 1):
            One of "clip" (default — standard lookup-table behaviour, holds
            the boundary value), "linear" (extends the boundary slope past the breakpoints),
            or "nan" (returns NaN outside ``[input_array[0], input_array[-1]]``).
            Stored outside ``@parameters``, so this kwarg is *not* round-tripped
            through model JSON — pre-existing models reload byte-equivalently.
        dtype (optional, T-038a):
            If set (e.g. ``jnp.float32``), the block's ``input_array`` and
            ``output_array`` are cast to this dtype on construction, so the
            interpolation arithmetic — and therefore the output signal — runs at
            this precision regardless of the global x64 setting.  The default
            (``None``) preserves the pre-T-038a behavior: the arrays are stored
            verbatim and float64 is used under the default x64-enabled install.

            T-038a-followup-mixed-precision-cascade: when ``dtype is None``
            and a :func:`jaxonomy.precision_policy` context manager is
            active, the block falls back to the context's dtype.  Explicit
            ``dtype=`` always wins (explicit-over-implicit).

            Best-effort: this enforces dtype on the block's internal arrays
            and on the output of ``jnp.interp`` / ``jnp.argmin`` lookups, but
            downstream operations (e.g. connecting to a default-dtype block)
            are subject to JAX's standard promotion rules — the result of the
            wider arithmetic may be promoted to ``float64``.

    Notes:
        Currently restricted to 1D input and output data.  This may be expanded to
        support multi-dimensional output arrays in the future.
    """

    @parameters(static=["input_array", "output_array", "interpolation"])
    def __init__(
        self,
        input_array,
        output_array,
        interpolation,
        dtype=None,
        extrapolation="clip",
        **kwargs,
    ):
        # T-038a-followup-mixed-precision-cascade — when no explicit
        # ``dtype=`` kwarg was passed (``dtype is None``), fall back to
        # the active ``precision_policy`` context manager's dtype, if
        # any.  Explicit per-block dtype always wins (explicit-over-
        # implicit).  Outside any active context the policy resolver
        # returns ``None`` and the block's default-dtype path runs
        # unchanged — byte-equivalent to pre-follow-up behavior.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        # T-038a — remember the per-block dtype override so ``initialize``
        # (called by InitializeParameterResolver after parameter resolution)
        # can apply it to the resolved array values.  ``dtype`` is *not* a
        # @parameters-tracked field: it is a build-time block-shape decision
        # that doesn't round-trip through model JSON or get JAX-traced.
        self._dtype = dtype
        # T-114 phase 1 — extrapolation policy.  Stored outside the
        # @parameters list so the kwarg does not round-trip through model
        # JSON.  The default ``"clip"`` matches the historical behavior of
        # ``jnp.interp`` so existing pipelines stay byte-equivalent.
        if extrapolation not in ("clip", "linear", "nan"):
            raise ValueError(
                f"LookupTable1d: extrapolation must be one of "
                f"('clip','linear','nan'), got {extrapolation!r}"
            )
        self._extrapolation = extrapolation
        super().__init__(None, **kwargs)
        # T-002: reject non-monotonic input_array up front so silent
        # interpolation garbage doesn't propagate downstream.
        _input_np = np.asarray(input_array)
        if _input_np.ndim == 1 and _input_np.size >= 2 and not np.all(
            np.diff(_input_np) > 0
        ):
            raise ValueError(
                f"LookupTable1d '{self.name}': input_array must be strictly "
                f"monotonically increasing, got {list(_input_np)}"
            )

    def initialize(self, input_array, output_array, interpolation):
        if self._dtype is not None:
            # T-038a — cast both lookup arrays to the per-block dtype.  Done
            # after npa.array() so the dtype override survives the backend's
            # default-float promotion.
            self.input_array = npa.asarray(input_array).astype(self._dtype)
            self.output_array = npa.asarray(output_array).astype(self._dtype)
        else:
            self.input_array = npa.array(input_array)
            self.output_array = npa.array(output_array)
        if len(self.input_array.shape) != 1:
            raise ValueError(
                f"LookupTable1d block {self.name} input_array must be 1D, got shape "
                f"{self.input_array.shape}"
            )
        if len(self.output_array.shape) != 1:
            raise ValueError(
                f"LookupTable1d block {self.name} output_array must be 1D, got shape "
                f"{self.output_array.shape}"
            )
        self.max_i = len(self.input_array) - 1

        # T-114 phase 1 — fast path: when the user requested the
        # historical default (``extrapolation="clip"`` and one of the
        # original three methods), keep the original implementations
        # untouched so this default path stays byte-equivalent with the
        # pre-T-114 code, including for the numpy backend.  Only when
        # the user opts in to ``"pchip"`` or non-clip extrapolation do we
        # route through the JAX-only ``interp_1d`` backend.
        legacy_methods = {
            "linear": self._lookup_linear,
            "nearest": self._lookup_nearest,
            "flat": self._lookup_flat,
        }
        if self._extrapolation == "clip" and interpolation in legacy_methods:
            self.replace_op(legacy_methods[interpolation])
            return

        if interpolation not in (
            "linear", "pchip", "akima", "cubic", "nearest", "flat"
        ):
            raise ValueError(
                f"LookupTable1d block {self.name} has invalid selection {interpolation} "
                "for 'interpolation'"
            )

        from .lookup_table import interp_1d as _interp_1d

        extrapolation = self._extrapolation

        def _op(x):
            return _interp_1d(
                x,
                self.input_array,
                self.output_array,
                method=interpolation,
                extrapolation=extrapolation,
            )

        self.replace_op(_op)

    def _lookup_linear(self, x):
        return npa.interp(x, self.input_array, self.output_array)

    def _lookup_nearest(self, x):
        i = npa.argmin(npa.abs(self.input_array - x))
        i = npa.clip(i, 0, self.max_i)
        return self.output_array[i]

    def _lookup_flat(self, x):
        i = npa.where(
            x < self.input_array[1],
            0,
            npa.argmin(x >= self.input_array) - 1,
        )
        return self.output_array[i]

    @classmethod
    def fit_from_data(
        cls,
        xp,
        x_data,
        y_data,
        *,
        weights=None,
        smoothness: float = 0.0,
        **block_kwargs,
    ):
        """Build a ``LookupTable1d`` whose output values are fitted by
        least squares to ``(x_data, y_data)`` at the fixed grid ``xp``.

        Ergonomic wrapper around :func:`jaxonomy.library.fit_lookup_table_1d`
        so the fitting entry point is discoverable from the block class
        itself.

        Args:
            xp: Fixed grid of breakpoints (1-D, strictly increasing).
            x_data: Measured input cloud, shape ``(K,)``.
            y_data: Measured output cloud, shape ``(K,)``.
            weights: Optional per-sample weights for weighted least
                squares.  ``None`` = OLS.
            smoothness: Non-negative discrete first-difference penalty.
                Use small values (1e-3 .. 1.0) on noisy / sparse data.
            **block_kwargs: Forwarded to
                :func:`jaxonomy.library.fit_lookup_table_1d` (e.g.
                ``interpolation=``, ``extrapolation=``, ``name=``,
                ``dtype=``).

        Returns:
            A ``LookupTable1d`` instance with ``input_array=xp`` and
            ``output_array`` set to the LS-fit table values.
        """
        # Lazy import — ``lookup_table_fitting`` imports ``LookupTable1d``
        # from this module, so doing the import at function-call time
        # breaks any circular import risk while keeping the public
        # ``LookupTable1d(...)`` constructor path untouched (existing
        # call sites stay byte-equivalent).
        from .lookup_table_fitting import fit_lookup_table_1d

        return fit_lookup_table_1d(
            xp,
            x_data,
            y_data,
            weights=weights,
            smoothness=smoothness,
            **block_kwargs,
        )

fit_from_data(xp, x_data, y_data, *, weights=None, smoothness=0.0, **block_kwargs) classmethod

Build a LookupTable1d whose output values are fitted by least squares to (x_data, y_data) at the fixed grid xp.

Ergonomic wrapper around :func:jaxonomy.library.fit_lookup_table_1d so the fitting entry point is discoverable from the block class itself.

Parameters:

Name	Type	Description	Default
xp		Fixed grid of breakpoints (1-D, strictly increasing).	required
x_data		Measured input cloud, shape (K,).	required
y_data		Measured output cloud, shape (K,).	required
weights		Optional per-sample weights for weighted least squares. None = OLS.	None
smoothness	float	Non-negative discrete first-difference penalty. Use small values (1e-3 .. 1.0) on noisy / sparse data.	0.0
**block_kwargs		Forwarded to :func:jaxonomy.library.fit_lookup_table_1d (e.g. interpolation=, extrapolation=, name=, dtype=).	{}

Returns:

Type	Description
	A LookupTable1d instance with input_array=xp and
	output_array set to the LS-fit table values.

Source code in jaxonomy/library/tables.py

@classmethod
def fit_from_data(
    cls,
    xp,
    x_data,
    y_data,
    *,
    weights=None,
    smoothness: float = 0.0,
    **block_kwargs,
):
    """Build a ``LookupTable1d`` whose output values are fitted by
    least squares to ``(x_data, y_data)`` at the fixed grid ``xp``.

    Ergonomic wrapper around :func:`jaxonomy.library.fit_lookup_table_1d`
    so the fitting entry point is discoverable from the block class
    itself.

    Args:
        xp: Fixed grid of breakpoints (1-D, strictly increasing).
        x_data: Measured input cloud, shape ``(K,)``.
        y_data: Measured output cloud, shape ``(K,)``.
        weights: Optional per-sample weights for weighted least
            squares.  ``None`` = OLS.
        smoothness: Non-negative discrete first-difference penalty.
            Use small values (1e-3 .. 1.0) on noisy / sparse data.
        **block_kwargs: Forwarded to
            :func:`jaxonomy.library.fit_lookup_table_1d` (e.g.
            ``interpolation=``, ``extrapolation=``, ``name=``,
            ``dtype=``).

    Returns:
        A ``LookupTable1d`` instance with ``input_array=xp`` and
        ``output_array`` set to the LS-fit table values.
    """
    # Lazy import — ``lookup_table_fitting`` imports ``LookupTable1d``
    # from this module, so doing the import at function-call time
    # breaks any circular import risk while keeping the public
    # ``LookupTable1d(...)`` constructor path untouched (existing
    # call sites stay byte-equivalent).
    from .lookup_table_fitting import fit_lookup_table_1d

    return fit_lookup_table_1d(
        xp,
        x_data,
        y_data,
        weights=weights,
        smoothness=smoothness,
        **block_kwargs,
    )

LookupTable2d

Bases: LeafSystem

Interpolate the input signals into a static lookup table.

The behavior is modeled on scipy.interpolate.interp2d but is implemented in JAX. "linear" (bilinear, default) and "bicubic" (Catmull-Rom) interpolation are supported. The input arrays must be 1D and the output array must be 2D.

Input ports

(0) The first input signal, used as the first interpolation coordinate. (1) The second input signal, used as the second interpolation coordinate.

Output ports

(0) The interpolated output signal.

Parameters:

Name	Type	Description	Default
input_x_array		The array of input values at which the output values are provided, corresponding to the first input signal. Must be 1D	required
input_y_array		The array of input values at which the output values are provided, corresponding to the second input signal. Must be 1D	required
output_table_array		The array of output values. Must be 2D with shape (m, n), where m = len(input_x_array) and n = len(input_y_array).	required
interpolation		"linear" (default, standard bilinear behaviour) or "bicubic" (T-114-followup-2d-bicubic — Catmull-Rom cubic-convolution kernel; C^1-continuous, exact at grid corners, smoother than bilinear off-grid). "bicubic" requires at least 4 breakpoints per axis.	'linear'
extrapolation	optional, T-114 phase 2	One of "clip" (default — standard lookup-table behaviour, holds the boundary value), "linear" (bilinear extension past the grid via edge-slope continuation), or "nan" (returns NaN outside the grid). The default "clip" matches the legacy npa.interp2d behaviour byte-equivalently. Stored outside @parameters so this kwarg is not round-tripped through model JSON; pre-existing models reload byte- equivalently.	'clip'

Source code in jaxonomy/library/tables.py

class LookupTable2d(LeafSystem):
    """Interpolate the input signals into a static lookup table.

    The behavior is modeled on `scipy.interpolate.interp2d` but is implemented
    in JAX.  ``"linear"`` (bilinear, default) and ``"bicubic"`` (Catmull-Rom)
    interpolation are supported.  The input arrays must be 1D and the output
    array must be 2D.

    Input ports:
        (0) The first input signal, used as the first interpolation coordinate.
        (1) The second input signal, used as the second interpolation coordinate.

    Output ports:
        (0) The interpolated output signal.

    Parameters:
        input_x_array:
            The array of input values at which the output values are provided,
            corresponding to the first input signal. Must be 1D
        input_y_array:
            The array of input values at which the output values are provided,
            corresponding to the second input signal. Must be 1D
        output_table_array:
            The array of output values. Must be 2D with shape `(m, n)`, where
            `m = len(input_x_array)` and `n = len(input_y_array)`.
        interpolation:
            ``"linear"`` (default, standard bilinear behaviour) or
            ``"bicubic"`` (T-114-followup-2d-bicubic — Catmull-Rom
            cubic-convolution kernel; C^1-continuous, exact at grid
            corners, smoother than bilinear off-grid).  ``"bicubic"``
            requires at least 4 breakpoints per axis.
        extrapolation (optional, T-114 phase 2):
            One of "clip" (default — standard lookup-table behaviour, holds
            the boundary value), "linear" (bilinear extension past the grid
            via edge-slope continuation), or "nan" (returns NaN outside
            the grid).  The default ``"clip"`` matches the legacy
            ``npa.interp2d`` behaviour byte-equivalently.  Stored
            outside ``@parameters`` so this kwarg is not round-tripped
            through model JSON; pre-existing models reload byte-
            equivalently.
    """

    @parameters(
        static=["input_x_array", "input_y_array", "output_table_array", "interpolation"]
    )
    def __init__(
        self,
        input_x_array,
        input_y_array,
        output_table_array,
        interpolation="linear",
        dtype=None,
        extrapolation="clip",
        **kwargs,
    ):
        # T-038a-followup-other-blocks: per-block dtype override; stored
        # outside the @parameters list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        # T-114 phase 2 — extrapolation policy.  Stored outside the
        # @parameters list so the kwarg does not round-trip through model
        # JSON.  The default ``"clip"`` matches the historical behavior of
        # ``npa.interp2d`` so existing pipelines stay byte-equivalent.
        if extrapolation not in ("clip", "linear", "nan"):
            raise ValueError(
                f"LookupTable2d: extrapolation must be one of "
                f"('clip','linear','nan'), got {extrapolation!r}"
            )
        self._extrapolation = extrapolation
        super().__init__(**kwargs)
        self.declare_input_port()
        self.declare_input_port()
        self._output_port_idx = self.declare_output_port(
            None,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
            requires_inputs=True,
        )

    # T-114-followup-lookup-table-fitted-getter — expose the stored static
    # parameters as plain read-only attributes so the fitted table can be
    # inspected (e.g. plotted) without re-running ``fit_table_2d``.
    @property
    def output_table_array(self):
        return self._static_parameters["output_table_array"].get()

    @property
    def input_x_array(self):
        return self._static_parameters["input_x_array"].get()

    @property
    def input_y_array(self):
        return self._static_parameters["input_y_array"].get()

    def initialize(
        self, input_x_array, input_y_array, output_table_array, interpolation
    ):
        if self._dtype is not None:
            # T-038a-followup-other-blocks: cast the lookup arrays to the
            # per-block dtype so the interp arithmetic runs at this
            # precision regardless of the global x64 setting.
            xp = npa.asarray(input_x_array).astype(self._dtype)
            yp = npa.asarray(input_y_array).astype(self._dtype)
            zp = npa.asarray(output_table_array).astype(self._dtype)
        else:
            xp = npa.array(input_x_array)
            yp = npa.array(input_y_array)
            zp = npa.array(output_table_array)

        if len(xp.shape) != 1:
            raise ValueError(
                f"LookupTable2d block {self.name} input_x_array must be 1D, got "
                f"shape {xp.shape}"
            )

        if len(yp.shape) != 1:
            raise ValueError(
                f"LookupTable2d block {self.name} input_y_array must be 1D, got "
                f"shape {yp.shape}"
            )

        if len(zp.shape) != 2:
            raise ValueError(
                f"LookupTable2d block {self.name} output_table_array must be 2D, "
                f"got shape {zp.shape}"
            )

        if zp.shape != (len(xp), len(yp)):
            raise ValueError(
                f"LookupTable2d block {self.name} output_table_array must have "
                f"shape (len(input_x_array), len(input_y_array)), got shape {zp.shape}"
            )

        if interpolation not in ("linear", "bicubic"):
            raise NotImplementedError(
                f"LookupTable2d block {self.name} only supports "
                f"'linear' or 'bicubic' interpolation, got {interpolation!r}."
            )

        if interpolation == "bicubic" and (len(xp) < 4 or len(yp) < 4):
            # Catmull-Rom needs 4 breakpoints per axis for a well-defined
            # stencil — fail loud rather than silently degrade.
            raise ValueError(
                f"LookupTable2d block {self.name}: interpolation='bicubic' "
                f"requires at least 4 breakpoints per axis, got "
                f"len(input_x_array)={len(xp)} and len(input_y_array)={len(yp)}"
            )

        # T-114 phase 2 — fast path: when the user requested the
        # historical default (``interpolation="linear"`` +
        # ``extrapolation="clip"``), keep the legacy ``npa.interp2d``
        # call so this default path stays byte-equivalent with the
        # pre-T-114-phase-2 code, including for the numpy backend.
        # Anything else (non-clip extrapolation OR bicubic interp)
        # routes through the JAX-only ``interp_2d`` backend.
        if interpolation == "linear" and self._extrapolation == "clip":
            self._compute_output = partial(npa.interp2d, xp, yp, zp)
        else:
            from .lookup_table import interp_2d as _interp_2d

            extrapolation = self._extrapolation
            method = interpolation

            def _op(x, y):
                return _interp_2d(
                    x, y, xp, yp, zp,
                    method=method,
                    extrapolation=extrapolation,
                )

            self._compute_output = _op

        self.configure_output_port(
            self._output_port_idx,
            self._output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
            requires_inputs=True,
        )

    def _output(self, _time, _state, *inputs, **params):
        (x, y) = inputs
        return self._compute_output(x, y)

    @classmethod
    def fit_from_data(
        cls,
        xp,
        yp,
        x_data,
        y_data,
        z_data,
        *,
        weights=None,
        smoothness: float = 0.0,
        **block_kwargs,
    ):
        """Build a ``LookupTable2d`` whose table values are fitted by
        bilinear least squares to ``(x_data, y_data, z_data)`` at the
        fixed grid ``(xp, yp)``.

        Ergonomic classmethod mirror of
        :func:`jaxonomy.library.fit_lookup_table_2d` so the 2-D fitting
        entry point is discoverable from the block class itself.

        Args:
            xp: Fixed grid along the first axis (1-D, strictly
                increasing).
            yp: Fixed grid along the second axis (1-D, strictly
                increasing).
            x_data, y_data, z_data: Measurement cloud, all shape
                ``(K,)``.
            weights: Optional per-sample weights for weighted least
                squares.  ``None`` = OLS.
            smoothness: Non-negative 5-point Laplacian penalty on the
                fitted table.  ``0.0`` (default) is pure data-fit.
            **block_kwargs: Forwarded to
                :func:`jaxonomy.library.fit_lookup_table_2d` (e.g.
                ``interpolation=``, ``extrapolation=``, ``name=``,
                ``dtype=``).

        Returns:
            A ``LookupTable2d`` instance with ``input_x_array=xp``,
            ``input_y_array=yp``, and ``output_table_array`` set to the
            LS-fit table of shape ``(len(xp), len(yp))``.
        """
        # Lazy import — ``lookup_table_fitting`` imports ``LookupTable2d``
        # from this module, so doing the import at function-call time
        # avoids the circular import while keeping the public
        # ``LookupTable2d(...)`` constructor path untouched (existing
        # call sites stay byte-equivalent).
        from .lookup_table_fitting import fit_lookup_table_2d

        return fit_lookup_table_2d(
            xp,
            yp,
            x_data,
            y_data,
            z_data,
            weights=weights,
            smoothness=smoothness,
            **block_kwargs,
        )

fit_from_data(xp, yp, x_data, y_data, z_data, *, weights=None, smoothness=0.0, **block_kwargs) classmethod

Build a LookupTable2d whose table values are fitted by bilinear least squares to (x_data, y_data, z_data) at the fixed grid (xp, yp).

Ergonomic classmethod mirror of :func:jaxonomy.library.fit_lookup_table_2d so the 2-D fitting entry point is discoverable from the block class itself.

Parameters:

Name	Type	Description	Default
xp		Fixed grid along the first axis (1-D, strictly increasing).	required
yp		Fixed grid along the second axis (1-D, strictly increasing).	required
x_data, y_data, z_data		Measurement cloud, all shape (K,).	required
weights		Optional per-sample weights for weighted least squares. None = OLS.	None
smoothness	float	Non-negative 5-point Laplacian penalty on the fitted table. 0.0 (default) is pure data-fit.	0.0
**block_kwargs		Forwarded to :func:jaxonomy.library.fit_lookup_table_2d (e.g. interpolation=, extrapolation=, name=, dtype=).	{}

Returns:

Type	Description
	A LookupTable2d instance with input_x_array=xp,
	input_y_array=yp, and output_table_array set to the
	LS-fit table of shape (len(xp), len(yp)).

Source code in jaxonomy/library/tables.py

@classmethod
def fit_from_data(
    cls,
    xp,
    yp,
    x_data,
    y_data,
    z_data,
    *,
    weights=None,
    smoothness: float = 0.0,
    **block_kwargs,
):
    """Build a ``LookupTable2d`` whose table values are fitted by
    bilinear least squares to ``(x_data, y_data, z_data)`` at the
    fixed grid ``(xp, yp)``.

    Ergonomic classmethod mirror of
    :func:`jaxonomy.library.fit_lookup_table_2d` so the 2-D fitting
    entry point is discoverable from the block class itself.

    Args:
        xp: Fixed grid along the first axis (1-D, strictly
            increasing).
        yp: Fixed grid along the second axis (1-D, strictly
            increasing).
        x_data, y_data, z_data: Measurement cloud, all shape
            ``(K,)``.
        weights: Optional per-sample weights for weighted least
            squares.  ``None`` = OLS.
        smoothness: Non-negative 5-point Laplacian penalty on the
            fitted table.  ``0.0`` (default) is pure data-fit.
        **block_kwargs: Forwarded to
            :func:`jaxonomy.library.fit_lookup_table_2d` (e.g.
            ``interpolation=``, ``extrapolation=``, ``name=``,
            ``dtype=``).

    Returns:
        A ``LookupTable2d`` instance with ``input_x_array=xp``,
        ``input_y_array=yp``, and ``output_table_array`` set to the
        LS-fit table of shape ``(len(xp), len(yp))``.
    """
    # Lazy import — ``lookup_table_fitting`` imports ``LookupTable2d``
    # from this module, so doing the import at function-call time
    # avoids the circular import while keeping the public
    # ``LookupTable2d(...)`` constructor path untouched (existing
    # call sites stay byte-equivalent).
    from .lookup_table_fitting import fit_lookup_table_2d

    return fit_lookup_table_2d(
        xp,
        yp,
        x_data,
        y_data,
        z_data,
        weights=weights,
        smoothness=smoothness,
        **block_kwargs,
    )

LookupTableND

Bases: LeafSystem

Interpolate the input signal into a static N-D lookup table.

Generalises :class:LookupTable1d and :class:LookupTable2d to an arbitrary number of axes. The block takes a single input port whose value is a length-N query vector [q_1, ..., q_N] and returns the multilinearly interpolated table value at that point.

Implementation: delegates to :func:jaxonomy.library.lookup_table.interp_nd, which performs N successive 1-D linear interpolations along each axis (no jnp.interpn exists today — see the deeper-followup note in that function's docstring).

Input ports

(0) — the query vector, shape (N,) where N is the number of grid axes.

Output ports

(0) — the multilinearly interpolated table value.

Parameters:

Name	Type	Description	Default
grid_axes		Tuple of N strictly-increasing 1-D breakpoint arrays. The i-th array has length B_i and corresponds to axis i of output_array.	required
output_array		Sample values, shape (B_1, B_2, ..., B_N). output_array[i_1, ..., i_N] = f(grid_axes[0][i_1], ..., grid_axes[N-1][i_N]).	required
interpolation		Currently only "linear" (multilinear). Reserved for future N-D smooth methods (filed under T-114-followup-phase4-nd-cubic).	'linear'
extrapolation	optional	One of "clip" (default — clips each coordinate to its axis range, the standard lookup-table default), "linear" (multilinear continuation past the grid), or "nan" (returns NaN whenever any coordinate is outside its axis range).	'clip'
dtype	optional	If set (e.g. jnp.float32), the block's grid arrays and output array are cast to this dtype on construction. Mirrors the per-block dtype contract of LookupTable1d.	None

Notes

Differentiable through both the query vector and the table values (modulo the discrete bucket-index searchsorted, whose gradient is piecewise constant — within a cell the gradient is exact).

Source code in jaxonomy/library/tables.py

class LookupTableND(LeafSystem):
    """Interpolate the input signal into a static N-D lookup table.

    Generalises :class:`LookupTable1d` and :class:`LookupTable2d` to an
    arbitrary number of axes.  The block takes a single input port whose
    value is a length-``N`` query vector ``[q_1, ..., q_N]`` and returns
    the multilinearly interpolated table value at that point.

    Implementation: delegates to
    :func:`jaxonomy.library.lookup_table.interp_nd`, which performs
    ``N`` successive 1-D linear interpolations along each axis (no
    ``jnp.interpn`` exists today — see the deeper-followup note in
    that function's docstring).

    Input ports:
        ``(0)`` — the query vector, shape ``(N,)`` where ``N`` is the
        number of grid axes.

    Output ports:
        ``(0)`` — the multilinearly interpolated table value.

    Parameters:
        grid_axes:
            Tuple of ``N`` strictly-increasing 1-D breakpoint arrays.
            The ``i``-th array has length ``B_i`` and corresponds to
            axis ``i`` of ``output_array``.
        output_array:
            Sample values, shape ``(B_1, B_2, ..., B_N)``.
            ``output_array[i_1, ..., i_N] = f(grid_axes[0][i_1], ...,
            grid_axes[N-1][i_N])``.
        interpolation:
            Currently only ``"linear"`` (multilinear).  Reserved for
            future N-D smooth methods (filed under
            ``T-114-followup-phase4-nd-cubic``).
        extrapolation (optional):
            One of ``"clip"`` (default — clips each coordinate to its
            axis range, the standard lookup-table default), ``"linear"`` (multilinear
            continuation past the grid), or ``"nan"`` (returns NaN
            whenever any coordinate is outside its axis range).
        dtype (optional):
            If set (e.g. ``jnp.float32``), the block's grid arrays and
            output array are cast to this dtype on construction.
            Mirrors the per-block dtype contract of ``LookupTable1d``.

    Notes:
        Differentiable through both the query vector and the table
        values (modulo the discrete bucket-index ``searchsorted``,
        whose gradient is piecewise constant — within a cell the
        gradient is exact).
    """

    def __init__(
        self,
        grid_axes,
        output_array,
        interpolation="linear",
        dtype=None,
        extrapolation="clip",
        **kwargs,
    ):
        # Per-block dtype + active precision policy fallback, matching
        # the LookupTable1d / LookupTable2d contract.  Stored outside
        # the @parameters list so the kwarg does not round-trip through
        # model JSON or get JAX-traced.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype

        if extrapolation not in ("clip", "linear", "nan"):
            raise ValueError(
                f"LookupTableND: extrapolation must be one of "
                f"('clip','linear','nan'), got {extrapolation!r}"
            )
        self._extrapolation = extrapolation

        if interpolation != "linear":
            raise NotImplementedError(
                f"LookupTableND only supports interpolation='linear' "
                f"(multilinear) today; got {interpolation!r}.  N-D smooth "
                f"methods are filed as T-114-followup-phase4-nd-cubic."
            )
        self._interpolation = interpolation

        # Eagerly validate grid + table shapes so misconfigured blocks
        # fail at construction rather than during context build.
        if not isinstance(grid_axes, (tuple, list)) or len(grid_axes) == 0:
            raise ValueError(
                f"LookupTableND: grid_axes must be a non-empty tuple/list "
                f"of 1-D breakpoint arrays, got {type(grid_axes).__name__}"
            )
        n_axes = len(grid_axes)
        for i, axis in enumerate(grid_axes):
            arr = np.asarray(axis)
            if arr.ndim != 1:
                raise ValueError(
                    f"LookupTableND: grid_axes[{i}] must be 1-D, got shape "
                    f"{arr.shape}"
                )
            if arr.size >= 2 and not np.all(np.diff(arr) > 0):
                raise ValueError(
                    f"LookupTableND: grid_axes[{i}] must be strictly "
                    f"monotonically increasing"
                )
        out_arr = np.asarray(output_array)
        if out_arr.ndim != n_axes:
            raise ValueError(
                f"LookupTableND: output_array.ndim ({out_arr.ndim}) must "
                f"equal len(grid_axes) ({n_axes})"
            )
        expected_shape = tuple(int(np.asarray(g).shape[0]) for g in grid_axes)
        if out_arr.shape != expected_shape:
            raise ValueError(
                f"LookupTableND: output_array.shape {out_arr.shape} must "
                f"match the grid shape {expected_shape}"
            )

        # Cast (or copy) into the backend's array type with the requested
        # dtype (or the default).  Stored on ``self`` for use in the
        # output computation closure.
        if self._dtype is not None:
            self._grid_axes = tuple(
                npa.asarray(np.asarray(g)).astype(self._dtype)
                for g in grid_axes
            )
            self._output_array = npa.asarray(out_arr).astype(self._dtype)
        else:
            self._grid_axes = tuple(npa.array(np.asarray(g)) for g in grid_axes)
            self._output_array = npa.array(out_arr)

        super().__init__(**kwargs)
        self.declare_input_port()

        from .lookup_table import interp_nd as _interp_nd

        extrapolation_local = self._extrapolation
        grid_local = self._grid_axes
        values_local = self._output_array

        def _compute(_time, _state, *inputs, **_params):
            (query,) = inputs
            return _interp_nd(
                grid_local,
                values_local,
                query,
                method="linear",
                extrapolation=extrapolation_local,
            )

        self.declare_output_port(
            _compute,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )

LowPassDiscrete

Bases: LeafSystem

Discrete first-order (single-pole RC) low-pass filter.

Implements the difference equation::

tau   = 1 / (2*pi*cutoff_hz)
alpha = dt / (dt + tau)
y[k]  = alpha * x[k] + (1 - alpha) * y[k-1]

The continuous-time analogue is H(s) = 1 / (1 + tau*s), with -3 dB crossover at f = cutoff_hz in the small-dt limit.

Input ports

(0) The input signal x.

Output ports

(0) The filtered signal y.

Parameters:

Name	Type	Description	Default
dt		Sampling period of the block (s).	required
cutoff_hz		Design cutoff frequency (Hz). Differentiable.	1.0
initial_state		Initial value of y[-1]. Default 0.0.	0.0

Notes

Differentiability: cutoff_hz enters the recursive update via smooth arithmetic (alpha = dt/(dt + 1/(2*pi*cutoff_hz))), so jax.grad is finite through it.

Source code in jaxonomy/library/dynamics.py

class LowPassDiscrete(LeafSystem):
    """Discrete first-order (single-pole RC) low-pass filter.

    Implements the difference equation::

        tau   = 1 / (2*pi*cutoff_hz)
        alpha = dt / (dt + tau)
        y[k]  = alpha * x[k] + (1 - alpha) * y[k-1]

    The continuous-time analogue is ``H(s) = 1 / (1 + tau*s)``, with -3 dB
    crossover at ``f = cutoff_hz`` in the small-``dt`` limit.

    Input ports:
        (0) The input signal ``x``.

    Output ports:
        (0) The filtered signal ``y``.

    Parameters:
        dt:
            Sampling period of the block (s).
        cutoff_hz:
            Design cutoff frequency (Hz).  Differentiable.
        initial_state:
            Initial value of ``y[-1]``.  Default 0.0.

    Notes:
        Differentiability: ``cutoff_hz`` enters the recursive update via
        smooth arithmetic (``alpha = dt/(dt + 1/(2*pi*cutoff_hz))``), so
        ``jax.grad`` is finite through it.
    """

    @parameters(
        static=["dt"],
        dynamic=["cutoff_hz", "initial_state"],
    )
    def __init__(
        self,
        dt,
        cutoff_hz=1.0,
        initial_state=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.dt = dt
        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, cutoff_hz, initial_state, dt=None):
        # Discrete-state seed: previous output ``y[k-1]``.  Use a plain
        # Python float so the T-005 default (float64 when x64 is on) is
        # honoured at trace time.
        y0 = npa.asarray(initial_state)
        self.declare_discrete_state(default_value=y0)

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=self.dt,
        )

        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=self.dt,
            default_value=y0,
            prerequisites_of_calc=[DependencyTicket.xd],
        )

    def reset_default_values(self, **dynamic_parameters):
        y0 = npa.asarray(dynamic_parameters["initial_state"])
        self.configure_discrete_state_default_value(default_value=y0)
        self.configure_output_port_default_value(self._output_port_idx, y0)

    @staticmethod
    def _alpha(dt, cutoff_hz):
        # ``tau = 1/(2*pi*fc)``, ``alpha = dt / (dt + tau)``.  Smooth in
        # ``cutoff_hz`` so jax.grad is finite.
        two_pi = 2.0 * npa.pi
        tau = 1.0 / (two_pi * cutoff_hz)
        return dt / (dt + tau)

    def _update(self, _time, state, *inputs, **params):
        x = inputs[0]
        y_prev = state.discrete_state
        alpha = self._alpha(self.dt, params["cutoff_hz"])
        return alpha * x + (1.0 - alpha) * y_prev

    def _output(self, _time, state, *_inputs, **_params):
        return state.discrete_state

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        u = self.eval_input(context)
        xd = context[self.system_id].discrete_state
        check_state_type(
            self,
            inp_data=u,
            state_data=xd,
            error_collector=error_collector,
        )

Luenberger

Bases: LeafSystem

Discrete-time Luenberger observer with user-supplied gain L.

State-update equation:

.. code-block:: text

x_hat[k+1] = A·x_hat[k] + B·u[k] + L·(y[k] - C·x_hat[k] - D·u[k])

Parameters:

Name	Type	Description	Default
dt		Discrete sample period (seconds). Must match the plant's sample period (or, for continuous plants, the discretisation period chosen for the design).	required
A, B, C, D		Plant state-space matrices (typically the output of jaxonomy.discretize(linearize(plant, op), dt)). D defaults to a zero matrix of compatible shape.	required
L		Observer gain matrix, shape (n_states, n_outputs). The caller computes this — e.g. via scipy.signal.place for pole placement, or via a steady-state Kalman gain.	required
x_hat_0		Initial state estimate. Defaults to zeros.	None

Input ports

(0) u: control input vector, shape (n_inputs,). (1) y: noisy measurement vector, shape (n_outputs,).

Output ports

(0) x_hat: state estimate vector, shape (n_states,).

Notes

This block is the simpler half of the Kalman pair — the design cost (computing L) is paid offline, leaving only the cheap runtime update. If you want online Riccati-based gain updates instead, use :class:KalmanFilter. If you have a continuous plant and want the steady-state infinite-horizon Kalman gain, use :class:InfiniteHorizonKalmanFilter.

Source code in jaxonomy/library/state_estimators/luenberger.py

class Luenberger(LeafSystem):
    """Discrete-time Luenberger observer with user-supplied gain ``L``.

    State-update equation:

    .. code-block:: text

        x_hat[k+1] = A·x_hat[k] + B·u[k] + L·(y[k] - C·x_hat[k] - D·u[k])

    Args:
        dt: Discrete sample period (seconds). Must match the plant's
            sample period (or, for continuous plants, the discretisation
            period chosen for the design).
        A, B, C, D: Plant state-space matrices (typically the output of
            ``jaxonomy.discretize(linearize(plant, op), dt)``). ``D``
            defaults to a zero matrix of compatible shape.
        L: Observer gain matrix, shape ``(n_states, n_outputs)``. The
            caller computes this — e.g. via ``scipy.signal.place`` for
            pole placement, or via a steady-state Kalman gain.
        x_hat_0: Initial state estimate. Defaults to zeros.

    Input ports:
        (0) ``u``: control input vector, shape ``(n_inputs,)``.
        (1) ``y``: noisy measurement vector, shape ``(n_outputs,)``.

    Output ports:
        (0) ``x_hat``: state estimate vector, shape ``(n_states,)``.

    Notes:
        This block is the simpler half of the Kalman pair — the design
        cost (computing ``L``) is paid offline, leaving only the cheap
        runtime update. If you want online Riccati-based gain updates
        instead, use :class:`KalmanFilter`. If you have a continuous
        plant and want the steady-state infinite-horizon Kalman gain,
        use :class:`InfiniteHorizonKalmanFilter`.
    """

    @parameters(static=["dt", "A", "B", "C", "D", "L", "x_hat_0"])
    def __init__(
        self,
        dt,
        A,
        B,
        C,
        L,
        D=None,
        x_hat_0=None,
        *,
        name=None,
        **kwargs,
    ):
        super().__init__(name=name, **kwargs)

        A_arr = jnp.asarray(A)
        B_arr = jnp.asarray(B)
        C_arr = jnp.asarray(C)
        L_arr = jnp.asarray(L)
        if A_arr.ndim != 2 or A_arr.shape[0] != A_arr.shape[1]:
            raise ValueError(
                f"Luenberger {self.name!r}: A must be square; got shape "
                f"{tuple(A_arr.shape)}."
            )
        n = A_arr.shape[0]
        m = B_arr.shape[1] if B_arr.ndim == 2 else 1
        p = C_arr.shape[0] if C_arr.ndim == 2 else 1
        if L_arr.shape != (n, p):
            raise ValueError(
                f"Luenberger {self.name!r}: L must have shape "
                f"(n_states, n_outputs) = ({n}, {p}); got "
                f"{tuple(L_arr.shape)}."
            )
        if D is None:
            D = jnp.zeros((p, m))

        self._n = n
        self._m = m
        self._p = p

        # Two input ports: u (control), y (measurement).
        self.declare_input_port()  # u
        self.declare_input_port()  # y
        self._update_idx = self.declare_periodic_update()
        self._output_idx = self.declare_output_port()

    def initialize(self, dt, A, B, C, D=None, L=None, x_hat_0=None):
        A = jnp.asarray(A)
        B = jnp.asarray(B)
        C = jnp.asarray(C)
        L = jnp.asarray(L)
        if D is None:
            D = jnp.zeros((self._p, self._m))
        else:
            D = jnp.asarray(D)
        if x_hat_0 is None:
            x_hat_0 = jnp.zeros((self._n,))
        else:
            x_hat_0 = jnp.asarray(x_hat_0)

        self.declare_discrete_state(default_value=x_hat_0, as_array=True)

        dt_f = float(dt)

        def _update(time, state, *inputs, **_params):
            u = jnp.atleast_1d(inputs[0])
            y = jnp.atleast_1d(inputs[1])
            x = state.discrete_state
            y_pred = C @ x + D @ u
            innovation = y - y_pred
            return A @ x + B @ u + L @ innovation

        # ``offset=dt`` so the first update fires at t=dt (matches the
        # UnitDelay convention: x_hat[0] = x_hat_0 visible from t=0 to dt).
        self.configure_periodic_update(
            self._update_idx,
            _update,
            period=dt_f,
            offset=dt_f,
        )

        def _output(time, state, *inputs, **_params):
            return state.discrete_state

        # Output is sampled at the same discrete cadence — period/offset
        # match the canonical discrete-output pattern used elsewhere
        # (UnitDelay, KalmanFilter etc.).
        self.configure_output_port(
            self._output_idx,
            _output,
            period=dt_f,
            offset=0.0,
            default_value=x_hat_0,
        )

MJX

Bases: MuJoCoBase

A system that wraps a MuJoCo model and provides a continuous-time ODE LeafSystem. Currently only supports a single body system.

Input ports

(0) The control input vector control.

Output ports

(0) The generalized position coordinates qpos. (1) The generalized velocity coordinates qvel. (2) The actuator coordinates act. (3) The sensor data sensor_data (if enabled). (4) The video output video as RGB frames of shape (H,W,3) (if enabled). (5) A fake output port, present only if vHIL=True and outputs Array(0.0). (6+) Custom output ports, defined with user-specified python scripts.

Parameters:

Name	Type	Description	Default
file_name	str	The path to the MuJoCo XML model file.	required
dt	float	If None, jaxonomy's internal solver will be used and this block can be considered as a continuous block. If set, the model will be run in a discrete mode with the specified timestep, using MJX's solver, more like Co-Simulation. In that case, it might be favorable to set use_mjx=False.	None
key_frame_0	int | str	The keyframe to initialize the model from.	None
qpos_0	Array	The initial generalized position coordinates.	None
qvel_0	Array	The initial generalized velocity coordinates.	None
act_0	Array	The initial actuator coordinates.	None
enable_sensor_data	bool	Whether to output the sensor data to an optional port named 'sensor_data'.	False
enable_video_output	bool	Whether to output the rendered video frames to an optional port named 'video'.	False
video_size	tuple[int, int]	The size of the video output frames as a (H,W) tuple.	None
enable_mocap_pos	bool	Whether to enable the mocap_pos input port for motion capture tracking.	False
vHIL	bool	Whether to run in virtual hardware-in-the-loop mode.	False
vHIL_dt	float	The timestep for the virtual hardware-in-the-loop mode.	0.01

Notes: (i) _model and _data refer to MuJoCo's mjModel and mjData objects respectively. model and data are the corresponding MJX objects. (ii) While sensordata output is supported as a pure callback to MuJoCo since MJX has not yet implemented this aspect. This can be expensive.

Source code in jaxonomy/library/mujoco.py

class MJX(MuJoCoBase):
    """
    A system that wraps a MuJoCo model and provides a continuous-time ODE LeafSystem.
    Currently only supports a single body system.

    Input ports:
        (0) The control input vector `control`.

    Output ports:
        (0) The generalized position coordinates `qpos`.
        (1) The generalized velocity coordinates `qvel`.
        (2) The actuator coordinates `act`.
        (3) The sensor data `sensor_data` (if enabled).
        (4) The video output `video` as RGB frames of shape (H,W,3) (if enabled).
        (5) A fake output port, present only if `vHIL=True` and outputs Array(0.0).
        (6+) Custom output ports, defined with user-specified python scripts.

    Parameters:
        file_name (str):
            The path to the MuJoCo XML model file.

        dt (float, optional):
            If None, jaxonomy's internal solver will be used and this block can
            be considered as a continuous block. If set, the model will be run in
            a discrete mode with the specified timestep, using MJX's solver, more like
            Co-Simulation. In that case, it might be favorable to set `use_mjx=False`.

        key_frame_0 (int|str, optional):
            The keyframe to initialize the model from.

        qpos_0 (Array, optional):
            The initial generalized position coordinates.

        qvel_0 (Array, optional):
            The initial generalized velocity coordinates.

        act_0 (Array, optional):
            The initial actuator coordinates.

        enable_sensor_data (bool, optional):
            Whether to output the sensor data to an optional port named 'sensor_data'.

        enable_video_output (bool, optional):
            Whether to output the rendered video frames to an optional port named 'video'.

        video_size (tuple[int, int], optional):
            The size of the video output frames as a (H,W) tuple.

        enable_mocap_pos (bool, optional):
            Whether to enable the mocap_pos input port for motion capture tracking.

        vHIL (bool, optional):
            Whether to run in virtual hardware-in-the-loop mode.

        vHIL_dt (float, optional):
            The timestep for the virtual hardware-in-the-loop mode.

    Notes:
    (i) `_model` and `_data` refer to MuJoCo's `mjModel` and `mjData` objects
    respectively. `model` and `data` are the corresponding MJX objects.
    (ii) While `sensordata` output is supported as a pure callback to MuJoCo since MJX
    has not yet implemented this aspect. This can be expensive.
    """

    @parameters(
        static=[
            "file_name",
            "dt",
            "key_frame_0",
            "qpos_0",
            "qvel_0",
            "act_0",
            "enable_sensor_data",
            "enable_video_output",
            "video_size",
            "enable_mocap_pos",
            "vHIL",
            "vHIL_dt",
        ]
    )
    def __init__(
        self,
        file_name: str,
        dt: float = None,
        key_frame_0: int | str = None,
        qpos_0: Array = None,
        qvel_0: Array = None,
        act_0: Array = None,
        enable_sensor_data=False,
        enable_video_output=False,
        video_size: tuple[int, int] = None,
        enable_mocap_pos=False,
        custom_output_scripts: dict[str, str] = None,
        vHIL=False,
        vHIL_dt=0.01,
        **kwargs,
    ):
        super().__init__(
            use_mjx=True,
            file_name=file_name,
            dt=dt,
            key_frame_0=key_frame_0,
            qpos_0=qpos_0,
            qvel_0=qvel_0,
            act_0=act_0,
            enable_sensor_data=enable_sensor_data,
            enable_video_output=enable_video_output,
            video_size=video_size,
            enable_mocap_pos=enable_mocap_pos,
            custom_output_scripts=custom_output_scripts,
            vHIL=vHIL,
            vHIL_dt=vHIL_dt,
            **kwargs,
        )

        try:
            self.model = mjx.put_model(self._model)
            self.data = mjx.put_data(self._model, self._data)
        except NotImplementedError as e:
            logger.error(
                "This robot model uses features not implemented in MJX. "
                "Please try the MuJoCo block instead (toggle use_mjx to false), "
                "or modify the MJCF file.",
                **logdata(block=self, exception=f"{type(e).__name__}: {str(e)}"),
            )
            raise e

        if dt is None or dt == 0:
            self.dt = None
            logger.info(
                "MuJoCo MJX block is running in continuous mode and will use "
                "Jaxonomy's solver.",
                **logdata(block=self),
            )

            state_0 = jnp.concatenate([self.qpos_0, self.qvel_0, self.act_0])
            self.declare_continuous_state(ode=self._ode, default_value=state_0)

        else:
            self.dt = dt
            known_solver_names = {
                mujoco.mjtSolver.mjSOL_CG: "Conjugate Gradient",
                mujoco.mjtSolver.mjSOL_NEWTON: "Newton",
            }
            logger.info(
                "MuJoCo MJX block is running in discrete mode with dt=%s, "
                "this will use MJX's solver '%s' and not Jaxonomy's solver.",
                dt,
                known_solver_names.get(
                    self.model.opt.solver,
                    str(self.model.opt.solver),
                ),
                **logdata(block=self),
            )

            # T-019b-followup-batched: stash the dtype-pytree of the
            # default mjx.Data so ``_step_cache_cb`` can cast its
            # ``mjx.step`` output to the same per-leaf dtypes.  Without
            # this, ``mjx.put_data`` (numpy) returns int32 for
            # ``contact.geom1/geom2/geom`` while ``mjx.step`` (under
            # jit + jax_enable_x64) returns int64, and the
            # ``DiscreteUpdateEvent.handle`` ``lax.cond`` rejects the
            # int32-vs-int64 branch mismatch with a TypeError.  Cast at
            # the callback boundary keeps it surgical and lets
            # ``simulate_batch`` work over batched MJX rollouts.
            self._mjx_data_dtype_tree = jax.tree.map(
                lambda x: jnp.asarray(x).dtype, self.data
            )
            callback_index = self.declare_cache(
                self._step_cache_cb,
                default_value=self.data,
                period=dt,
                offset=0.0,
                requires_inputs=True,
            )
            self.mjx_data_cache_index = self.callbacks[callback_index].cache_index

        self.declare_output_port(
            self._output_qpos,
            default_value=self.qpos_0,
            requires_inputs=False,
            name="qpos",
        )

        self.declare_output_port(
            self._output_qvel,
            default_value=self.qvel_0,
            requires_inputs=False,
            name="qvel",
        )

        self.declare_output_port(
            self._output_act,
            default_value=self.act_0,
            requires_inputs=False,
            name="act",
        )

        if enable_sensor_data:
            logger.warning(
                "Sensor data output with MJX might be very slow. Consider switching "
                "to the non-MJX MuJoCo block for better performance.",
                **logdata(block=self),
            )
            self._declare_sensor_data_port(self.dt)

        if enable_video_output:
            logger.warning(
                "Video output with MJX might be very slow. Consider switching "
                "to the non-MJX MuJoCo block for better performance.",
                **logdata(block=self),
            )
            self._declare_video_output_port(video_size)

        if vHIL:
            self._declare_vhil_fake_output_port(vHIL_dt)

        self._declare_custom_output_ports(custom_output_scripts, self.dt)

    def _cached_data(self, state: LeafState) -> mjxData:
        return state.cache[self.mjx_data_cache_index]

    def _qpos(self, state: LeafState):
        if self.dt is not None:
            return self._cached_data(state).qpos

        return state.continuous_state[self.qpos_start : self.qpos_end]

    def _qvel(self, state: LeafState):
        if self.dt is not None:
            return self._cached_data(state).qvel

        return state.continuous_state[self.qvel_start : self.qvel_end]

    def _act(self, state: LeafState):
        if self.dt is not None:
            return self._cached_data(state).act

        return state.continuous_state[self.act_start : self.act_end]

    def _ode(self, time, state, *inputs, **parameters):
        # Implementation of the ODE when running model in continuous mode, with
        # jaxonomy's internal solver.

        qpos = self._qpos(state)
        qvel = self._qvel(state)
        act = self._act(state)

        ctrl = inputs[0]

        # Normalize quaternion components of qpos to prevent numerical drift that
        # accumulates over long simulations.  MuJoCo stores quaternions interleaved
        # with positions in qpos; normalize_qpos_quat handles only the quat slots.
        qpos = self.normalize_qpos_quat(qpos)

        model, data = self.model, self.data
        data = data.replace(time=time, qpos=qpos, qvel=qvel, act=act, ctrl=ctrl)

        data = mjx.forward(model, data)

        qvel_dot = data.qacc
        qpos_dot = position_derivatives(model.jnt_type, qpos, qvel)
        act_dot = data.act_dot

        state_dot = jnp.concatenate([qpos_dot, qvel_dot, act_dot])

        return state_dot

    def _step_cache_cb(self, time, state: LeafState, *inputs, **parameters):
        # Implementation of the ODE when running model in discrete mode, with
        # MJX's solver. This is like the non-MJX variant or an FMU.

        # TODO: try a version wrapped with io_callback to compare
        # compilation times. Splitting the compute graph between jaxonomy
        # and mjx could bring improvements, but quite obviously at the cost
        # of any usefulness of mjx over mujoco (autodiff, vmap, ...).

        ctrl = inputs[0]

        data = self._cached_data(state)
        data = data.replace(time=time, ctrl=ctrl)
        data = mjx.step(self.model, data)

        # T-019b-followup-batched: cast every leaf back to the dtype of
        # the default ``mjx.Data`` so the ``DiscreteUpdateEvent.handle``
        # ``lax.cond`` true/false branches agree.  Empirically only the
        # ``contact.geom*`` integer indices flip int32→int64 under
        # jit+x64, but the tree-map covers any future divergence MJX
        # introduces without a hand-maintained field whitelist.
        data = jax.tree.map(
            lambda x, d: x.astype(d), data, self._mjx_data_dtype_tree
        )

        return data

    def _output_qpos(self, time, state, *inputs, **parameters):
        qpos = self._qpos(state)
        qpos_normalized_quats = self.normalize_qpos_quat(qpos)
        return qpos_normalized_quats

    def _output_qvel(self, time, state, *inputs, **parameters):
        return self._qvel(state)

    def _output_act(self, time, state, *inputs, **parameters):
        return self._act(state)

    def _mj_forward(self, time, qpos, qvel, act, ctrl=None):
        data = self._data
        data.time = time
        data.qpos[:] = qpos
        data.qvel[:] = qvel
        data.act[:] = act
        if ctrl is not None:
            data.ctrl[:] = ctrl
        mujoco.mj_forward(self._model, data)
        return data

    def _pure_callback_sensordata(self, time, qpos, qvel, act, ctrl):
        data = self._mj_forward(time, qpos, qvel, act, ctrl)
        return data.sensordata

    def _output_sensor_data(self, time, state, *inputs, **parameters):
        qpos = self._qpos(state)
        qvel = self._qvel(state)
        act = self._act(state)
        ctrl = inputs[0] if inputs else None

        qpos = self.normalize_qpos_quat(qpos)

        return jax.pure_callback(
            self._pure_callback_sensordata,
            self.pure_callback_sensordata_result_type,
            time,
            qpos,
            qvel,
            act,
            ctrl,
        )

    def normalize_qpos_quat(self, qpos):
        """
        Normalize the quaternion components of the generalized position coordinates.
        """
        qpos_normalized, qi = [], 0

        for jnt_typ in self.model.jnt_type:
            if jnt_typ == mjx_types.JointType.FREE:
                trans = qpos[qi : qi + 3]
                quat = qpos[qi + 3 : qi + 7]
                norm_quat = mjx_math.normalize(quat)
                qpos_normalized.append(jnp.concatenate([trans, norm_quat]))
                qi = qi + 7
            elif jnt_typ == mjx_types.JointType.BALL:
                quat = qpos[qi : qi + 4]
                norm_quat = mjx_math.normalize(quat)
                qpos_normalized.append(norm_quat)
                qi = qi + 4
            elif jnt_typ in (mjx_types.JointType.HINGE, mjx_types.JointType.SLIDE):
                trans = qpos[qi]
                qpos_normalized.append(trans[None])
                qi = qi + 1
            else:
                raise RuntimeError(f"unrecognized joint type: {jnt_typ}")

        return jnp.concatenate(qpos_normalized) if qpos_normalized else jnp.empty((0,))

    def _update_viewer(self, time, state, *inputs, **parameters):
        ctrl = inputs[0]
        return jax.pure_callback(
            self._pure_callback_update_viewer,
            self.pure_callback_update_result_type,
            ctrl,
        )

    def _pure_callback_update_viewer(self, ctrl):
        if self.viewer.is_running:
            self.viewer.sync()
            self.data_vhil.ctrl[:] = ctrl
            mujoco.mj_step(self.model_vhil, self.data_vhil)
        return jnp.array(0.0)

    def _output_video_discrete(self, time, state, *inputs, **parameters):
        def _discrete_cb(state):
            mjx_data = self._cached_data(state)
            data = mjx.get_data(self._model, mjx_data)
            if self.enable_mocap_pos:
                data.mocap_pos[:] = inputs[1]
            return self.render(data)

        return io_callback(_discrete_cb, self._video_default, time, state)

    def _output_video(self, time, state, *inputs, **parameters):
        if self.dt is not None:
            return self._output_video_discrete(time, state, *inputs, **parameters)

        def _continuous_cb(time, qpos, qvel, act, inputs):
            data = self._mj_forward(time, qpos, qvel, act)
            if self.enable_mocap_pos:
                data.mocap_pos[:] = inputs[1]
            return self.render(data)

        qpos = self._qpos(state)
        qvel = self._qvel(state)
        act = self._act(state)

        qpos = self.normalize_qpos_quat(qpos)

        return io_callback(
            _continuous_cb, self._video_default, time, qpos, qvel, act, inputs
        )

normalize_qpos_quat(qpos)

Normalize the quaternion components of the generalized position coordinates.

Source code in jaxonomy/library/mujoco.py

def normalize_qpos_quat(self, qpos):
    """
    Normalize the quaternion components of the generalized position coordinates.
    """
    qpos_normalized, qi = [], 0

    for jnt_typ in self.model.jnt_type:
        if jnt_typ == mjx_types.JointType.FREE:
            trans = qpos[qi : qi + 3]
            quat = qpos[qi + 3 : qi + 7]
            norm_quat = mjx_math.normalize(quat)
            qpos_normalized.append(jnp.concatenate([trans, norm_quat]))
            qi = qi + 7
        elif jnt_typ == mjx_types.JointType.BALL:
            quat = qpos[qi : qi + 4]
            norm_quat = mjx_math.normalize(quat)
            qpos_normalized.append(norm_quat)
            qi = qi + 4
        elif jnt_typ in (mjx_types.JointType.HINGE, mjx_types.JointType.SLIDE):
            trans = qpos[qi]
            qpos_normalized.append(trans[None])
            qi = qi + 1
        else:
            raise RuntimeError(f"unrecognized joint type: {jnt_typ}")

    return jnp.concatenate(qpos_normalized) if qpos_normalized else jnp.empty((0,))

MLP

Bases: FeedthroughBlock

A feedforward neural network block representing an Equinox multi-layer perceptron (MLP). The output y of the MLP is computed as

    y = MLP(x, theta)

where theta are the parameters of the MLP, and x is the input to the MLP. This block is differentialble w.r.t. the MLP parameters theta. Note that theta, does not include the hyperparameters representing the architecture of the MLP.

Input ports

(0) The input to the MLP.

Output ports

(0) The output of the MLP.

Parameters:

Name	Type	Description	Default
in_size	int	The dimension of the input to the MLP.	None
out_size	int	The dimension of the output of the MLP.	None
width_size	int	The width of every hidden layers of the MLP.	None
depth	int	The depth of the MLP. This represents the number of hidden layers, including the output layer.	None
seed	int	The seed for the random number generator for initialization of the MLP parameters (weights and biases of every layer). If None, a random 32-bit seed will be generated.	None
activation_str	str	The activation function to use after each internal layer of the MLP. Possible values are "relu", "sigmoid", "tanh", "elu", "swish", "gelu", "leaky_relu", "rbf", and "identity". Default is "relu".	'relu'
final_activation_str	str	The activation function to use for the output layer of the MLP. Same choices as activation_str. Default is "identity".	'identity'
use_bias	bool	Whether to add a bias to the internal layers of the MLP. Default is True.	True
use_final_bias	bool	Wheter to add a bias to the output layer of the MLP. Default is True.	True
file_name	str	Optional file name containing the serialized parameters of the MLP. If provided, the parameters are loaded from the file, and set as the parameters of the MLP. Default is None.	None

Source code in jaxonomy/library/nn.py

class MLP(FeedthroughBlock):
    """
    A feedforward neural network block representing an Equinox multi-layer
    perceptron (MLP). The output `y` of the MLP is computed as

    ```
        y = MLP(x, theta)
    ```

    where `theta` are the parameters of the MLP, and `x` is the input to the MLP.
    This block is differentialble w.r.t. the MLP parameters `theta`. Note that `theta`,
    does not include the hyperparameters representing the architecture of the MLP.

    Input ports:
        (0) The input to the MLP.

    Output ports:
        (0) The output of the MLP.

    Parameters:
        in_size (int):
            The dimension of the input to the MLP.
        out_size (int):
            The dimension of the output of the MLP.
        width_size (int):
            The width of every hidden layers of the MLP.
        depth (int):
            The depth of the MLP. This represents the number of hidden layers,
            including the output layer.
        seed (int):
            The seed for the random number generator for initialization of the
            MLP parameters (weights and biases of every layer).
            If None, a random 32-bit seed will be generated.
        activation_str (str):
            The activation function to use after each internal layer of the MLP.
            Possible values are ``"relu"``, ``"sigmoid"``, ``"tanh"``, ``"elu"``,
            ``"swish"``, ``"gelu"``, ``"leaky_relu"``, ``"rbf"``, and
            ``"identity"``. Default is ``"relu"``.
        final_activation_str (str):
            The activation function to use for the output layer of the MLP.
            Same choices as ``activation_str``. Default is ``"identity"``.
        use_bias (bool):
            Whether to add a bias to the internal layers of the MLP.
            Default is True.
        use_final_bias (bool):
            Wheter to add a bias to the output layer of the MLP.
            Default is True.
        file_name (str):
            Optional file name containing the serialized parameters of the MLP.
            If provided, the parameters are loaded from the file, and set as the
            parameters of the MLP. Default is None.
    """

    @parameters(
        static=[
            "in_size",
            "out_size",
            "width_size",
            "depth",
            "seed",
            "activation_str",
            "final_activation_str",
            "use_bias",
            "use_final_bias",
            "file_name",
        ],
    )
    def __init__(
        self,
        in_size=None,
        out_size=None,
        width_size=None,
        depth=None,
        seed=None,
        activation_str="relu",
        final_activation_str="identity",
        use_bias=True,
        use_final_bias=True,
        file_name=None,
        **kwargs,
    ):
        """
        see https://docs.kidger.site/equinox/examples/serialisation/ for rationale
        of implementation here. We can't serialize the activation function, so we
        serialize a string representing a selection for activation function amongst
        a finite set of options.
        """
        super().__init__(None, **kwargs)
        # The Equinox MLP object is built in ``initialize()``, which runs when
        # ``create_context()`` is first called. Seed a sentinel so ``self.mlp``
        # access before then raises a clear error rather than AttributeError
        # (T-B4-followup-mlp-pre-context).
        self._mlp = None

    @property
    def mlp(self):
        """The underlying Equinox ``eqx.nn.MLP`` object.

        Built lazily in :meth:`initialize`, which runs the first time
        ``create_context()`` is called on a diagram containing this block.
        Accessing it before then raises a clear error pointing at
        ``create_context()`` (T-B4-followup-mlp-pre-context).
        """
        if self._mlp is None:
            raise AttributeError(
                f"MLP block {self.name!r}: the underlying Equinox network is "
                f"not built until the block is initialized. Call "
                f"`diagram.create_context()` (or `block.create_context()` for a "
                f"standalone block) first, then access `.mlp`. The architecture "
                f"hyperparameters (in_size / out_size / width_size / depth) are "
                f"available immediately; only the parameterised network object "
                f"is deferred."
            )
        return self._mlp

    def initialize(
        self,
        in_size=None,
        out_size=None,
        width_size=None,
        depth=None,
        seed=None,
        activation_str="relu",
        final_activation_str="identity",
        use_bias=True,
        use_final_bias=True,
        file_name=None,
        mlp_params=None,
    ):
        # mlp_params is stored as a dynamic parameter.  The guard below
        # (`if "mlp_params" in self.dynamic_parameters`) ensures it is updated
        # rather than overwritten on subsequent calls (e.g. after optimization),
        # so optimized weights survive re-initialization of the block.

        if in_size is None or out_size is None or width_size is None or depth is None:
            raise ValueError("Must specify in_size, out_size, width_size, and depth.")
        else:
            # Cast to int for safety
            in_size = int(in_size)
            out_size = int(out_size)
            width_size = int(width_size)
            depth = int(depth)

        # file_name may come as an empty string through json parsing
        if file_name == "":
            file_name = None

        # Mapping from activation string to callable.
        # Extend here to add new activations; see https://jax.readthedocs.io/en/latest/jax.nn.html
        def _match_activation(activation_str):
            activation_mapping = {
                "relu": jax.nn.relu,
                "sigmoid": jax.nn.sigmoid,
                "tanh": jnp.tanh,
                "elu": jax.nn.elu,
                "swish": jax.nn.silu,
                "gelu": jax.nn.gelu,
                "leaky_relu": jax.nn.leaky_relu,
                "rbf": lambda x: jnp.exp(-(x**2)),
                "identity": lambda x: x,
            }
            if activation_str not in activation_mapping:
                warnings.warn(
                    f"Provided activation function {activation_str} not recognized. "
                    "Using Identity function as activation."
                )
            return activation_mapping.get(activation_str, lambda x: x)

        seed = (
            np.random.randint(0, 2**32, dtype=np.int64) if seed is None else int(seed)
        )
        self.key = random.PRNGKey(seed)

        self._mlp = eqx.nn.MLP(
            in_size,
            out_size,
            width_size,
            depth,
            key=self.key,
            activation=_match_activation(activation_str),
            final_activation=_match_activation(final_activation_str),
            use_bias=use_bias,
            use_final_bias=use_final_bias,
        )

        if file_name is not None:
            with open(file_name, "rb") as fp:
                self._mlp = eqx.tree_deserialise_leaves(fp, self._mlp)

        # partition into a pytree of params and static components
        mlp_params, self.mlp_static = eqx.partition(self._mlp, eqx.is_array)

        if "mlp_params" in self.dynamic_parameters:
            self.dynamic_parameters["mlp_params"].set(mlp_params)
        else:
            self.declare_dynamic_parameter("mlp_params", mlp_params, as_array=False)

        def _eval_MLP(inputs, **parameters):
            mlp_params = parameters["mlp_params"]
            mlp = eqx.combine(mlp_params, self.mlp_static)
            return mlp(inputs)

        self.replace_op(_eval_MLP)

    def serialize(self, file_name, mlp_params=None):
        """
        Serialize only the parameters of the MLP. Note that the hyperparameters
        representing the architecture of the MLP are not serialized. This is because
        of the following use-cases imagined:
        (i) The user may train the Equinox MLP outside of Jaxonomy. In this case,
        it seems unnecessary to force the user to serialize the hyperparameters of the
        MLP in the strict form chosen by Jaxonomy. It would seem much easier
        for the user to just input these hyperparameters when creating the MLP block
        in Jaxonomy UI, and upload the naturally produced serialized parameters file
        by Equinox.
        (ii) The user may want to train the Equinox MLP within Jaxonomy in a notebook,
        and then use the block within Colimator UI. In this case, while serialization of
        the hyperparameters of the MLP would be a litte more convenient compared
        to manually inputting the hyperparameters in the UI, it seems like a small
        convenience relative to disadvantages of (i). Ideally the user should be
        able to use the API to push the learnt parameters.
        (iii) When we support training in the UI, the hyperparameters are naturally
        serialzed with `declare_configuraton_parameters`, and thus, in this case too,
        only serializatio of the MLP parameters is necessary.

        The choice of an optional `mlp_params` is to enable training of the
        models in a notebook and easily seralizing them for use in the UI.
        """

        if self._mlp is None:
            # The Equinox network is built lazily in ``initialize()``, which
            # normally runs on the first ``create_context()`` call. Build it
            # here from the declared parameters (mirroring
            # ``LeafContextFactory.create_node_context``) so that a
            # freshly-constructed block can be serialized directly.
            self.initialize(**self.parameters)

        if mlp_params is None:
            mlp = self._mlp
        else:
            mlp = eqx.combine(mlp_params, self.mlp_static)
        with open(file_name, "wb") as f:
            eqx.tree_serialise_leaves(f, mlp)

mlp property

The underlying Equinox eqx.nn.MLP object.

Built lazily in :meth:initialize, which runs the first time create_context() is called on a diagram containing this block. Accessing it before then raises a clear error pointing at create_context() (T-B4-followup-mlp-pre-context).

__init__(in_size=None, out_size=None, width_size=None, depth=None, seed=None, activation_str='relu', final_activation_str='identity', use_bias=True, use_final_bias=True, file_name=None, **kwargs)

see https://docs.kidger.site/equinox/examples/serialisation/ for rationale of implementation here. We can't serialize the activation function, so we serialize a string representing a selection for activation function amongst a finite set of options.

Source code in jaxonomy/library/nn.py

@parameters(
    static=[
        "in_size",
        "out_size",
        "width_size",
        "depth",
        "seed",
        "activation_str",
        "final_activation_str",
        "use_bias",
        "use_final_bias",
        "file_name",
    ],
)
def __init__(
    self,
    in_size=None,
    out_size=None,
    width_size=None,
    depth=None,
    seed=None,
    activation_str="relu",
    final_activation_str="identity",
    use_bias=True,
    use_final_bias=True,
    file_name=None,
    **kwargs,
):
    """
    see https://docs.kidger.site/equinox/examples/serialisation/ for rationale
    of implementation here. We can't serialize the activation function, so we
    serialize a string representing a selection for activation function amongst
    a finite set of options.
    """
    super().__init__(None, **kwargs)
    # The Equinox MLP object is built in ``initialize()``, which runs when
    # ``create_context()`` is first called. Seed a sentinel so ``self.mlp``
    # access before then raises a clear error rather than AttributeError
    # (T-B4-followup-mlp-pre-context).
    self._mlp = None

serialize(file_name, mlp_params=None)

Serialize only the parameters of the MLP. Note that the hyperparameters representing the architecture of the MLP are not serialized. This is because of the following use-cases imagined: (i) The user may train the Equinox MLP outside of Jaxonomy. In this case, it seems unnecessary to force the user to serialize the hyperparameters of the MLP in the strict form chosen by Jaxonomy. It would seem much easier for the user to just input these hyperparameters when creating the MLP block in Jaxonomy UI, and upload the naturally produced serialized parameters file by Equinox. (ii) The user may want to train the Equinox MLP within Jaxonomy in a notebook, and then use the block within Colimator UI. In this case, while serialization of the hyperparameters of the MLP would be a litte more convenient compared to manually inputting the hyperparameters in the UI, it seems like a small convenience relative to disadvantages of (i). Ideally the user should be able to use the API to push the learnt parameters. (iii) When we support training in the UI, the hyperparameters are naturally serialzed with declare_configuraton_parameters, and thus, in this case too, only serializatio of the MLP parameters is necessary.

The choice of an optional mlp_params is to enable training of the models in a notebook and easily seralizing them for use in the UI.

Source code in jaxonomy/library/nn.py

def serialize(self, file_name, mlp_params=None):
    """
    Serialize only the parameters of the MLP. Note that the hyperparameters
    representing the architecture of the MLP are not serialized. This is because
    of the following use-cases imagined:
    (i) The user may train the Equinox MLP outside of Jaxonomy. In this case,
    it seems unnecessary to force the user to serialize the hyperparameters of the
    MLP in the strict form chosen by Jaxonomy. It would seem much easier
    for the user to just input these hyperparameters when creating the MLP block
    in Jaxonomy UI, and upload the naturally produced serialized parameters file
    by Equinox.
    (ii) The user may want to train the Equinox MLP within Jaxonomy in a notebook,
    and then use the block within Colimator UI. In this case, while serialization of
    the hyperparameters of the MLP would be a litte more convenient compared
    to manually inputting the hyperparameters in the UI, it seems like a small
    convenience relative to disadvantages of (i). Ideally the user should be
    able to use the API to push the learnt parameters.
    (iii) When we support training in the UI, the hyperparameters are naturally
    serialzed with `declare_configuraton_parameters`, and thus, in this case too,
    only serializatio of the MLP parameters is necessary.

    The choice of an optional `mlp_params` is to enable training of the
    models in a notebook and easily seralizing them for use in the UI.
    """

    if self._mlp is None:
        # The Equinox network is built lazily in ``initialize()``, which
        # normally runs on the first ``create_context()`` call. Build it
        # here from the declared parameters (mirroring
        # ``LeafContextFactory.create_node_context``) so that a
        # freshly-constructed block can be serialized directly.
        self.initialize(**self.parameters)

    if mlp_params is None:
        mlp = self._mlp
    else:
        mlp = eqx.combine(mlp_params, self.mlp_static)
    with open(file_name, "wb") as f:
        eqx.tree_serialise_leaves(f, mlp)

MaskedDelayBuffer

Bases: LeafSystem

Delay buffer where the delay length can be set at runtime (up to max_steps).

Like ShiftRegister, but allows the delay to be specified as an input signal. Uses masking (not dynamic indexing) for JAX compatibility.

Parameters:

Name	Type	Description	Default
max_steps	int	Maximum possible delay. STATIC.	required
signal_shape	tuple	Shape of each signal frame.	()
dt	float	Discrete update interval.	0.01

Ports

Input[0] "u": signal to delay Input[1] "delay_steps": integer scalar, 0 < delay_steps <= max_steps Output[0] "y": delayed signal

Source code in jaxonomy/library/delay.py

class MaskedDelayBuffer(LeafSystem):
    """
    Delay buffer where the delay length can be set at 
    runtime (up to max_steps).

    Like ShiftRegister, but allows the delay to be specified 
    as an input signal. Uses masking (not dynamic indexing) 
    for JAX compatibility.

    Parameters:
        max_steps (int): Maximum possible delay. STATIC.
        signal_shape (tuple): Shape of each signal frame.
        dt (float): Discrete update interval.

    Ports:
        Input[0] "u": signal to delay
        Input[1] "delay_steps": integer scalar, 
            0 < delay_steps <= max_steps
        Output[0] "y": delayed signal
    """

    @parameters(static=["max_steps", "signal_shape"])
    def __init__(
        self, 
        max_steps: int, 
        signal_shape: tuple = (), 
        dt: float = 0.01, 
        **kwargs
    ):
        super().__init__(**kwargs)
        self.dt = dt
        self.max_steps = max_steps
        self.signal_shape = signal_shape

        self.input_u_idx = self.declare_input_port()
        self.input_delay_idx = self.declare_input_port()

        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, max_steps, signal_shape):
        initial_value = npa.zeros(signal_shape)
        buffer = npa.broadcast_to(initial_value, (max_steps, *signal_shape))
        self.declare_discrete_state(default_value=buffer)

        self.configure_periodic_update(
            self._periodic_update_idx, 
            self._update, 
            period=self.dt, 
            offset=self.dt
        )

        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=0.0,
            requires_inputs=True,
            prerequisites_of_calc=[
                DependencyTicket.xd, 
                self.input_ports[self.input_delay_idx].ticket
            ],
        )

    def _update(self, _time, state, *inputs, **_params):
        u = inputs[self.input_u_idx]
        buffer = npa.roll(state.discrete_state, shift=1, axis=0)
        buffer = buffer.at[0].set(u)
        return buffer

    def _output(self, _time, state, *inputs, **_params):
        delay_steps = npa.clip(inputs[self.input_delay_idx], 1, self.max_steps)
        buffer = state.discrete_state

        mask = npa.arange(self.max_steps) == (delay_steps - 1)
        axes = tuple(range(1, 1 + len(self.signal_shape)))
        if axes:
            mask = npa.expand_dims(mask, axis=axes)

        return npa.sum(buffer * mask, axis=0)

MatrixConcatenation

Bases: ReduceBlock

Concatenate two matrices along a given axis.

Dispatches to jax.numpy.concatenate, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.concatenate.html

Parameters:

Name	Type	Description	Default
axis		The axis along which the matrices are concatenated. 0 for vertical and 1 for horizontal. Default is 0.	0

Input ports

(0, 1) The input matrices A and B

Output ports

(0) The concatenation input matrices: e.g. [A,B].

Source code in jaxonomy/library/math_ops.py

class MatrixConcatenation(ReduceBlock):
    """Concatenate two matrices along a given axis.

    Dispatches to `jax.numpy.concatenate`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.concatenate.html

    Args:
        axis: The axis along which the matrices are concatenated. 0 for vertical
            and 1 for horizontal. Default is 0.

    Input ports:
        (0, 1) The input matrices `A` and `B`

    Output ports:
        (0) The concatenation input matrices: e.g. `[A,B]`.
    """

    @parameters(static=["axis"])
    def __init__(self, n_in=2, axis=0, **kwargs):
        if n_in != 2:
            raise ValueError(
                "MatrixConcatenation block only supports two input matrices."
            )
        super().__init__(2, None, **kwargs)

    def initialize(self, axis):
        def _func(inputs):
            return npa.concatenate((inputs[0], inputs[1]), axis=int(axis))

        self.replace_op(_func)

MatrixInversion

Bases: FeedthroughBlock

Compute the matrix inverse of the input signal.

Dispatches to jax.numpy.inv, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.inv.html

Input ports

(0) The input matrix.

Output ports

(0) The inverse of the input matrix.

Source code in jaxonomy/library/math_ops.py

class MatrixInversion(FeedthroughBlock):
    """Compute the matrix inverse of the input signal.

    Dispatches to `jax.numpy.inv`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.inv.html

    Input ports:
        (0) The input matrix.

    Output ports:
        (0) The inverse of the input matrix.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(npa.linalg.inv, *args, **kwargs)

MatrixMultiplication

Bases: ReduceBlock

Compute the matrix product of the input signals.

Dispatches to jax.numpy.matmul, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.matmul.html

Input ports

(0, 1) The input matrices A and B

Output ports

(0) The matrix product of the input matrices: A @ B.

Source code in jaxonomy/library/math_ops.py

class MatrixMultiplication(ReduceBlock):
    """Compute the matrix product of the input signals.

    Dispatches to `jax.numpy.matmul`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.matmul.html

    Input ports:
        (0, 1) The input matrices `A` and `B`

    Output ports:
        (0) The matrix product of the input matrices: `A @ B`.
    """

    def __init__(
        self,
        n_in=2,
        **kwargs,
    ):
        if n_in != 2:
            raise ValueError(
                "MatrixMultiplication block only supports two input signals."
            )

        def _func(inputs):
            return npa.matmul(inputs[0], inputs[1])

        super().__init__(2, _func, **kwargs)

MatrixTransposition

Bases: FeedthroughBlock

Compute the matrix transpose of the input signal.

Dispatches to jax.numpy.transpose, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.transpose.html

Input ports

(0) The input matrix.

Output ports

(0) The transpose of the input matrix.

Source code in jaxonomy/library/math_ops.py

class MatrixTransposition(FeedthroughBlock):
    """Compute the matrix transpose of the input signal.

    Dispatches to `jax.numpy.transpose`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.transpose.html

    Input ports:
        (0) The input matrix.

    Output ports:
        (0) The transpose of the input matrix.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(npa.transpose, *args, **kwargs)

MinMax

Bases: ReduceBlock

Return the extremum of the input signals.

Input ports

(0..n_in-1) The input signals.

Output ports

(0) The minimum or maximum of the input signals.

Parameters:

Name	Type	Description	Default
operator		One of "min" or "max". Determines whether the block returns the minimum or maximum of the input signals.	required

Events

An event is triggered when the extreme input signal changes. For example, if the block is configured as a "max" block with two inputs and the second signal becomes greater than the first, a zero-crossing event will be triggered.

Source code in jaxonomy/library/math_ops.py

class MinMax(ReduceBlock):
    """Return the extremum of the input signals.

    Input ports:
        (0..n_in-1) The input signals.

    Output ports:
        (0) The minimum or maximum of the input signals.

    Parameters:
        operator:
            One of "min" or "max". Determines whether the block returns the minimum
            or maximum of the input signals.

    Events:
        An event is triggered when the extreme input signal changes.  For example,
        if the block is configured as a "max" block with two inputs and the second
        signal becomes greater than the first, a zero-crossing event will be triggered.
    """

    @parameters(static=["operator"])
    def __init__(self, n_in, operator, **kwargs):
        super().__init__(n_in, None, **kwargs)

    def initialize(self, operator):
        func_lookup = {
            "max": self._max,
            "min": self._min,
        }
        if operator not in func_lookup:
            # cannot pass system=self because this error must be raised BEFORE calling super.__init__()
            # in the case of inheritting from FeedthroughBlock.
            # if we call super.__init__() first, we get missing key error for func_lookup[base].
            raise BlockParameterError(
                message=f"MinMax block {self.name} has invalid selection {operator} for 'operator'. Valid options: "
                + ", ".join([f for f in func_lookup.keys()]),
                parameter_name="operator",
            )

        self.operator = operator

        self.replace_op(func_lookup[operator])

        guard_lookup = {
            "max": self._max_guard,
            "min": self._min_guard,
        }

        self._guard = guard_lookup[operator]

    def _min(self, inputs):
        return npa.min(npa.array(inputs))

    def _max(self, inputs):
        return npa.max(npa.array(inputs))

    def _min_guard(self, _time, _state, *inputs, **_params):
        return npa.argmin(npa.array(inputs)).astype(float)

    def _max_guard(self, _time, _state, *inputs, **_params):
        return npa.argmax(npa.array(inputs)).astype(float)

    def initialize_static_data(self, context):
        # Add a zero-crossing event so ODE solvers can't try to integrate
        # through a discontinuity. For efficiency, only do this if the output
        # is fed to an ODE block
        if not self.has_zero_crossing_events and (self.output_ports[0]):
            self.declare_zero_crossing(self._guard, direction="edge_detection")

        return super().initialize_static_data(context)

ModelicaFMU

Bases: LeafSystem

Source code in jaxonomy/library/fmu_import.py

class ModelicaFMU(LeafSystem):
    # Should we pass parameter overrides via kwargs? Sounds like it could conflict
    # in some rare cases (eg. dt, name...). The corresponding definition in
    # block_interface.py is pretty fragile in this regard.
    def __init__(
        self,
        file_name,
        dt,
        name=None,
        input_names: list[str] = None,
        output_names: list[str] = None,
        parameters: dict = None,
        start_time: float = 0.0,
        first_step_at_zero: bool = False,
        **kwargs,
    ):
        """Load and execute an FMU for Co-Simulation.

        .. warning:: **One instance per FMU per process** for FMUs built
            with pythonfmu (e.g. via :func:`jaxonomy.library.build_fmu`).
            The embedded-Python wrapper holds a process-wide
            ``Py_Initialize`` singleton, so instantiating the same
            ``.fmu`` dylib twice in one Python process fails. For
            multi-start or batched co-simulation, isolate each instance
            in its own process (``multiprocessing`` /
            ``concurrent.futures.ProcessPoolExecutor`` with the *spawn*
            start method, or a ``subprocess`` running a small driver
            script) and aggregate results afterwards. This is an
            upstream pythonfmu limitation, not a Jaxonomy one; FMUs from
            other exporters (OpenModelica, Dymola, Reference-FMUs) do
            not carry it.

        Args:
            file_name (str): path to FMU file
            dt (float): stepsize for FMU simulation
            name (str, optional): name of block
            input_names (list[str], optional): if set, only expose these inputs
            output_names (list[str], optional): if set, only expose these outputs
            parameters (dict, optional): dictionary of parameter overrides
            start_time (float, optional): FMU experiment start time.
            first_step_at_zero (bool, optional): Default ``False``. The
                first FMU step fires at ``t=dt`` (Modelica clocked-block
                convention: the periodic update is offset by one sample),
                so the block's outputs at ``t=0`` reflect the FMU's
                initial state (post-``setupExperiment``, pre-step) and
                only switch to "post-first-step" values from ``t=dt``
                onward. This introduces a one-sample phase lag versus an
                FMU exported with ``offset=0`` semantics, which weakens
                FMU round-trip byte-equivalence. Pass
                ``first_step_at_zero=True`` to fire the first step at
                ``t=0`` so the outputs leave their initial value as soon
                as the simulation begins. Surfaced as the FMU offset
                asymmetry in a follow-up finding.
            kwargs: ignored
        """
        try:
            super().__init__(name=name)
            self._init(
                file_name,
                dt,
                name=name or f"fmu_{self.system_id}",
                input_names=input_names,
                output_names=output_names,
                parameters=parameters,
                start_time=start_time,
                first_step_at_zero=first_step_at_zero,
            )
        except Exception as e:
            logger.error(
                "Failed to initialize FMU block %s (%s): %s", name, self.system_id, e
            )
            raise BlockInitializationError(str(e), system=self)

    @parameters(static=["file_name"])
    def _init(
        self,
        file_name,
        dt,
        name: str,
        input_names: list[str] = None,
        output_names: list[str] = None,
        parameters: dict = None,
        start_time: float = 0.0,
        first_step_at_zero: bool = False,
    ):
        self.dt = dt

        # read the model description
        model_description = fmpy.read_model_description(file_name)

        # extract the FMU
        unzipdir = fmpy.extract(file_name)

        # T-026: dispatch on FMI version.  FMI 3.0 has a different slave
        # class and uses type-specific getters (getFloat64, getInt32, ...)
        # in place of FMI 2.0's untyped getReal / getInteger.
        self._fmi_version = "3.0" if _is_fmi3(model_description) else "2.0"
        if self._fmi_version == "3.0":
            cs = model_description.coSimulation
            if cs is None:
                # Reference-FMUs Clocks.fmu, e.g., is scheduledExecution-only —
                # a different fmpy class (FMU3ScheduledExecution) and a
                # different stepping protocol. Out of scope for the
                # co-simulation block.
                raise BlockInitializationError(
                    f"FMU {file_name} has no co-simulation interface "
                    f"(only modelExchange / scheduledExecution). "
                    f"ModelicaFMU only supports co-simulation FMUs.",
                    system=self,
                )
            self.fmu = fmu = fmi3.FMU3Slave(
                guid=model_description.guid,
                unzipDirectory=unzipdir,
                modelIdentifier=cs.modelIdentifier,
                instanceName=name,
            )
            fmu.instantiate()
            # FMI 3.0 collapses setupExperiment into enterInitializationMode
            # via keyword args.
            fmu.enterInitializationMode(startTime=start_time)
        else:
            self.fmu = fmu = fmi2.FMU2Slave(
                guid=model_description.guid,
                unzipDirectory=unzipdir,
                modelIdentifier=model_description.coSimulation.modelIdentifier,
                instanceName=name,
            )
            fmu.instantiate()
            # setup and set startTime before entering initialization mode
            # per FMI 2.0.4 section 2.1.6.
            fmu.setupExperiment(startTime=start_time)
            # enter initialization mode before get/set params per FMI 2.0.4 section 4.2.4.
            fmu.enterInitializationMode()

        # collect the value references
        self.fmu_inputs: list[ValueReference] = []
        self.fmu_outputs: list[ValueReference] = []
        # T-026a: parallel ScalarVariable arrays so exec_step can dispatch
        # each port to the right typed getter/setter (mixed-type FMUs)
        # and reshape array I/O.
        self.fmu_input_vars: list = []
        self.fmu_output_vars: list = []

        inputs_by_name: dict[str, ScalarVariable] = {}
        outputs_by_name: dict[str, ScalarVariable] = {}
        variable_by_id: dict[int, ScalarVariable] = {}

        # FIXME: we rely on the XML file here, but jaxonomy uses a similar
        # JSON file with altered variable names.
        # TODO: implement support for parsing that file and mapping from
        # jaxonomy json name to/from xml name properly.
        def _compatible_param_name(name):
            return name.replace(".", "_")

        for variable in model_description.modelVariables:
            if variable.causality == "input":
                variable_by_id[variable.valueReference] = variable
                inputs_by_name[variable.name] = variable
            elif variable.causality == "output":
                variable_by_id[variable.valueReference] = variable
                outputs_by_name[variable.name] = variable
            elif variable.causality == "parameter" and parameters is not None:
                compat_name = _compatible_param_name(variable.name)
                parameter_value = parameters.get(compat_name, None)
                if parameter_value is None:
                    continue

                logger.debug(
                    "Setting parameter #%d '%s' <%s>: %s %s",
                    variable.valueReference,
                    variable.name,
                    variable.type,
                    parameter_value,
                    type(parameter_value),
                )

                # Values at this point have been wrapped into np.ndarray of
                # shape () via jaxonomy's JSON parsing. Enumerations are ints.
                # T-026: dispatch to v3 setter names where applicable.
                self._set_value(fmu, variable, parameter_value, name)

        # If input_names or output_names are set, we filter out the variables
        # exposed as I/O ports to match those. This so that the ports in model.json
        # actually match those in the FMU.
        # NOTE: Maybe this is unnecessarily complicated.
        # T-026a: types we can't represent inside the JAX-traced state.
        # Exclude them from the *default* port set; users who actually want
        # them must opt in via input_names/output_names and handle the
        # object dtype themselves.
        _NON_JAX_TYPES = {"String", "Binary"}

        def _accept_default(variable, role):
            if variable.type in _NON_JAX_TYPES:
                logger.warning(
                    "FMU %s: skipping %s port %r (type %s — not "
                    "representable as a JAX array; pass via "
                    "%s_names to expose it explicitly)",
                    name, role, variable.name, variable.type, role,
                )
                return False
            return True

        if input_names is not None:
            for in_name in input_names:
                if in_name not in inputs_by_name:
                    raise BlockInitializationError(
                        f"Input port {in_name} found on the block { name} "
                        + f"but not found in FMU {file_name}",
                        system=self,
                    )
                variable = inputs_by_name[in_name]
                self.fmu_inputs.append(variable.valueReference)
                self.fmu_input_vars.append(variable)
                self.declare_input_port(name=variable.name)
        else:
            for in_name, variable in inputs_by_name.items():
                if not _accept_default(variable, "input"):
                    continue
                self.fmu_inputs.append(variable.valueReference)
                self.fmu_input_vars.append(variable)
                self.declare_input_port(name=in_name)

        if output_names is not None:
            for out_name in output_names:
                if out_name not in outputs_by_name:
                    raise BlockInitializationError(
                        f"Input port {out_name} found on the block { name} "
                        + f"but not found in FMU {file_name}",
                        system=self,
                    )
                variable = outputs_by_name[out_name]
                self.fmu_outputs.append(variable.valueReference)
                self.fmu_output_vars.append(variable)
        else:
            for out_name, variable in outputs_by_name.items():
                if not _accept_default(variable, "output"):
                    continue
                self.fmu_outputs.append(variable.valueReference)
                self.fmu_output_vars.append(variable)

        # T-026a: pre-compute per-type read/write groupings so the
        # io_callback at every step can dispatch with one batched call
        # per type instead of one per port.
        self._output_groups = self._build_groups(self.fmu_output_vars, "get")
        self._input_groups = self._build_groups(self.fmu_input_vars, "set")

        # exit initialization mode after get/set params per FMI 2.0.4 section 4.2.4.
        fmu.exitInitializationMode()

        # Declare a discrete state component for each of the output
        # variables. T-026a fix: index by position into fmu_output_vars,
        # not by valueReference — FMI 3 alias variables (e.g. BouncingBall's
        # ``h`` and ``h_ft`` sharing vr=1) collapse a dict-by-vr lookup.
        self.state_names = [v.name for v in self.fmu_output_vars]
        self.DiscreteStateType = namedtuple("DiscreteState", self.state_names)

        # Create the default discrete state values
        default_values = {}
        for variable in self.fmu_output_vars:
            start_value = self._get_value(fmu, variable)
            default_values[variable.name] = start_value

        # Map the default values to array-like types so that they have shape and dtype
        default_state = jax.tree_util.tree_map(
            npa.asarray, self.DiscreteStateType(**default_values)
        )
        self.declare_discrete_state(default_value=default_state, as_array=False)

        # Declare an output port for each of the output variables
        def _make_output_callback(o_port_name):
            def _output(time, state, *inputs, **parameters):
                return getattr(state.discrete_state, o_port_name)

            return _output

        for o_port_name in default_values:
            self.declare_output_port(
                _make_output_callback(o_port_name),
                name=o_port_name,
                prerequisites_of_calc=[DependencyTicket.xd],
                requires_inputs=False,
            )

        # The step function acts as a periodic update that will update all components
        # of the discrete state.
        #
        # A5 (jax.grad-through-FMU): an FMU co-simulation step is an opaque
        # external call routed through ``io_callback`` — JAX has no derivative
        # rule for it, and a naive ``jax.grad`` otherwise dies with the generic
        # "IO callbacks do not support JVP". Wrap the step in a ``custom_jvp``
        # whose JVP rule raises a clear, FMU-specific error with concrete
        # workarounds, so the failure names the cause instead of leaving the
        # user to decode a backend message. The forward (primal) path is
        # unchanged — ``custom_jvp`` only intercepts differentiation.
        block_name = self.name

        @jax.custom_jvp
        def _fmu_step(time, state, inputs_tuple):
            return io_callback(
                self.exec_step, default_state, time, state, *inputs_tuple
            )

        @_fmu_step.defjvp
        def _fmu_step_jvp(primals, tangents):
            raise BlockRuntimeError(
                "jax.grad / jax.jvp through a ModelicaFMU block "
                f"({block_name!r}) is not supported: an FMU co-simulation step "
                "is an opaque external call (via io_callback / fmpy) with no "
                "analytic derivative. To obtain sensitivities, either (a) use "
                "finite differences over the FMU inputs/parameters (perturb in "
                "plain numpy outside jax.grad, or use jaxonomy.uq Monte Carlo / "
                "Sobol), or (b) replace the FMU with a native jaxonomy model "
                "for the part of the system you need to differentiate. The "
                "forward simulation (no jax.grad) works fine.",
                system=self,
            )

        def _step(time, state, *inputs):
            # Use the io_callback (wrapped for a clear grad-time error) so that
            # we can call the untraceable FMU object.
            return _fmu_step(time, state, tuple(inputs))

        # ``offset=dt`` (default) honors the Modelica clocked-block
        # convention so the block's outputs at ``t=0`` reflect the FMU's
        # ``setupExperiment`` state. ``first_step_at_zero=True`` fires
        # the first step at ``t=0``, eliminating the one-sample phase
        # lag for users who exported the FMU with that semantics. See
        # the constructor docstring.
        self.declare_periodic_update(
            _step,
            period=dt,
            offset=0.0 if first_step_at_zero else dt,
        )

    # T-026 / T-026a: type-dispatch helpers covering both FMI 2 and FMI 3.
    def _set_value(self, fmu, variable, value, block_name):
        """Set one variable on the FMU using the right typed setter.

        Supports scalar and array variables. ``value`` may be a scalar,
        a numpy/jax array, or any iterable; it is flattened to length
        ``prod(variable.shape)`` before the C call.
        """
        table = _accessor_table(self._fmi_version)
        vt = variable.type
        if vt not in table or table[vt][1] is None:
            raise BlockInitializationError(
                f"Unsupported FMI {self._fmi_version} variable type "
                f"{vt!r} for parameter {variable.name} in FMU block "
                f"{block_name}",
                system=self,
            )
        _, setter_name, dtype = table[vt]
        ref = [variable.valueReference]
        n = _variable_n_values(variable)
        try:
            if dtype is object:  # String / Binary — pass through as-is
                if n == 1:
                    fmu_values = [value if isinstance(value, (str, bytes))
                                  else str(value)]
                else:
                    fmu_values = list(value)
            elif dtype is np.bool_:
                arr = np.asarray(value).reshape(-1).astype(np.bool_)
                fmu_values = [bool(v) for v in arr]
            else:
                arr = np.asarray(value).reshape(-1).astype(dtype, copy=False)
                fmu_values = arr.tolist()
            if len(fmu_values) != n:
                raise ValueError(
                    f"variable {variable.name} expects {n} values "
                    f"(shape={_variable_shape(variable)}), got {len(fmu_values)}"
                )
            getattr(fmu, setter_name)(ref, fmu_values)
        except Exception as e:
            raise BlockInitializationError(
                f"Failed to set parameter {variable.name}: {e}", system=self,
            ) from e

    def _get_value(self, fmu, variable):
        """Read one variable from the FMU using the version-correct getter.

        Returns a Python scalar for shape-() variables and a numpy
        ndarray of the right dtype/shape for array variables.
        """
        table = _accessor_table(self._fmi_version)
        vt = variable.type
        if vt not in table or table[vt][0] is None:
            raise NotImplementedError(
                f"Unsupported FMI {self._fmi_version} variable type {vt!r} for "
                f"output port {variable.name}"
            )
        getter_name, _, dtype = table[vt]
        ref = [variable.valueReference]
        shape = _variable_shape(variable)
        n = _variable_n_values(variable)
        # FMI 3 typed getters take nValues for arrays. FMI 2 has no array
        # type at the C level, so n is always 1 there.
        if self._fmi_version == "3.0" and n != 1:
            raw = getattr(fmu, getter_name)(ref, n)
        else:
            raw = getattr(fmu, getter_name)(ref)
        if not shape:
            return raw[0]
        if dtype is object:
            return np.asarray(list(raw), dtype=object).reshape(shape)
        return np.asarray(raw, dtype=dtype).reshape(shape)

    def _build_groups(self, variables, mode):
        """T-026a: bucket variables by FMI type for one batched call per
        type. Returns a list of (accessor_name, dtype, port_indices,
        value_refs, n_values_per_ref, shapes).

        ``mode`` selects the column from the type table — ``"get"`` or
        ``"set"``. Per-port reshapes happen in exec_step using the
        recorded shapes; cumulative offsets are recomputed from
        n_values_per_ref to keep the structure light.
        """
        col = 0 if mode == "get" else 1
        table = _accessor_table(self._fmi_version)
        by_type: dict[str, dict] = {}
        for idx, var in enumerate(variables):
            vt = var.type
            if vt not in table or table[vt][col] is None:
                raise BlockInitializationError(
                    f"Unsupported FMI {self._fmi_version} {mode}-port type "
                    f"{vt!r} on variable {var.name}", system=self,
                )
            accessor_name, _setter, dtype = table[vt]
            if mode == "set":
                accessor_name = table[vt][1]
            bucket = by_type.setdefault(vt, {
                "accessor": accessor_name, "dtype": dtype,
                "indices": [], "refs": [], "nvals": [], "shapes": [],
            })
            bucket["indices"].append(idx)
            bucket["refs"].append(var.valueReference)
            bucket["nvals"].append(_variable_n_values(var))
            bucket["shapes"].append(_variable_shape(var))
        return list(by_type.values())

    def _create_discrete_state_type(self, fmu, fmu_outputs, variables):
        self.state_names = [variables[output_ref].name for output_ref in fmu_outputs]
        self.DiscreteStateType = namedtuple("DiscreteState", self.state_names)

    def exec_step(self, time, state, *inputs, **parameters):
        # NOTE: We should get the fmu from the context in order to build a pure
        # function but it is very unlikely this would ever work with FMUs since
        # they have their own internal hidden state. More context here:
        # https://github.com/machinavitalis/jaxonomy/pull/5330/files#r1419062533
        # Also look at that PR to see the previous implementation (it worked with
        # a single I/O port).

        try:
            fmu = self.fmu

            # Note: although it may appear that the order of operations below is
            # backwards, e.g. 1] get_outputs, 2] set_inputs, 3] step, this is
            # actually intentional.
            # Explanation by example assuming 1sec update intervals.
            # The reason get_outputs happens before set_inputs and 'step, is that
            # at t=0, the fmu outputs are already at t=0, so we can just read them.
            # Then, the fmu should get inputs at t=0, and use those to take a step
            # to t=1. The step operation, using inputs at t=0, puts the fmu in a
            # state where it outputs are now at t=1. This we cannot read them until
            # next update interval at t=1.

            # T-026a: read every output type group, then write every input
            # type group. One C call per type. Arrays are reshaped on read
            # and flattened on write.
            xd: dict = {}
            is_v3 = self._fmi_version == "3.0"
            for grp in self._output_groups:
                refs = grp["refs"]
                total_n = sum(grp["nvals"])
                if is_v3 and total_n != len(refs):
                    raw = getattr(fmu, grp["accessor"])(refs, total_n)
                else:
                    raw = getattr(fmu, grp["accessor"])(refs)
                # Walk the flat result and slice per port.
                offset = 0
                for port_idx, n, shape in zip(grp["indices"], grp["nvals"], grp["shapes"]):
                    chunk = raw[offset:offset + n]
                    offset += n
                    name = self.state_names[port_idx]
                    if not shape:
                        xd[name] = chunk[0]
                    elif grp["dtype"] is object:
                        xd[name] = np.asarray(list(chunk), dtype=object).reshape(shape)
                    else:
                        xd[name] = np.asarray(chunk, dtype=grp["dtype"]).reshape(shape)

            # Group inputs by type and flatten each port's value to that
            # type's batched ``set...`` call.
            for grp in self._input_groups:
                values_flat = []
                for port_idx, n, shape in zip(grp["indices"], grp["nvals"], grp["shapes"]):
                    val = inputs[port_idx]
                    if grp["dtype"] is object:
                        if n == 1:
                            values_flat.append(
                                val if isinstance(val, (str, bytes)) else str(val)
                            )
                        else:
                            values_flat.extend(list(val))
                    elif grp["dtype"] is np.bool_:
                        arr = np.asarray(val).reshape(-1).astype(np.bool_)
                        values_flat.extend(bool(v) for v in arr)
                    else:
                        arr = np.asarray(val).reshape(-1).astype(grp["dtype"], copy=False)
                        values_flat.extend(arr.tolist())
                getattr(fmu, grp["accessor"])(grp["refs"], values_flat)

            # Advance the FMU in time. The periodic update fires at t=dt,
            # 2dt, ..., but doStep expects currentCommunicationPoint to be
            # the *start* of the step interval — i.e. the FMU's current
            # internal time, which is one period earlier. Strict FMUs
            # (e.g. Reference-FMUs/BouncingBall) error otherwise.
            fmu.doStep(
                currentCommunicationPoint=float(time) - self.dt,
                communicationStepSize=self.dt,
            )

        except (*_fmi_call_exception(),) as e:
            logger.error(
                "Failed to run FMU block %s (%s): %s", self.name, self.system_id, e
            )
            raise BlockRuntimeError(str(e), system=self) from e

        xd = jax.tree_util.tree_map(npa.asarray, xd)

        return self.DiscreteStateType(**xd)

__init__(file_name, dt, name=None, input_names=None, output_names=None, parameters=None, start_time=0.0, first_step_at_zero=False, **kwargs)

Load and execute an FMU for Co-Simulation.

.. warning:: One instance per FMU per process for FMUs built with pythonfmu (e.g. via :func:jaxonomy.library.build_fmu). The embedded-Python wrapper holds a process-wide Py_Initialize singleton, so instantiating the same .fmu dylib twice in one Python process fails. For multi-start or batched co-simulation, isolate each instance in its own process (multiprocessing / concurrent.futures.ProcessPoolExecutor with the spawn start method, or a subprocess running a small driver script) and aggregate results afterwards. This is an upstream pythonfmu limitation, not a Jaxonomy one; FMUs from other exporters (OpenModelica, Dymola, Reference-FMUs) do not carry it.

Parameters:

Name	Type	Description	Default
file_name	str	path to FMU file	required
dt	float	stepsize for FMU simulation	required
name	str	name of block	None
input_names	list[str]	if set, only expose these inputs	None
output_names	list[str]	if set, only expose these outputs	None
parameters	dict	dictionary of parameter overrides	None
start_time	float	FMU experiment start time.	0.0
first_step_at_zero	bool	Default False. The first FMU step fires at t=dt (Modelica clocked-block convention: the periodic update is offset by one sample), so the block's outputs at t=0 reflect the FMU's initial state (post-setupExperiment, pre-step) and only switch to "post-first-step" values from t=dt onward. This introduces a one-sample phase lag versus an FMU exported with offset=0 semantics, which weakens FMU round-trip byte-equivalence. Pass first_step_at_zero=True to fire the first step at t=0 so the outputs leave their initial value as soon as the simulation begins. Surfaced as the FMU offset asymmetry in a follow-up finding.	False
kwargs		ignored	{}

Source code in jaxonomy/library/fmu_import.py

def __init__(
    self,
    file_name,
    dt,
    name=None,
    input_names: list[str] = None,
    output_names: list[str] = None,
    parameters: dict = None,
    start_time: float = 0.0,
    first_step_at_zero: bool = False,
    **kwargs,
):
    """Load and execute an FMU for Co-Simulation.

    .. warning:: **One instance per FMU per process** for FMUs built
        with pythonfmu (e.g. via :func:`jaxonomy.library.build_fmu`).
        The embedded-Python wrapper holds a process-wide
        ``Py_Initialize`` singleton, so instantiating the same
        ``.fmu`` dylib twice in one Python process fails. For
        multi-start or batched co-simulation, isolate each instance
        in its own process (``multiprocessing`` /
        ``concurrent.futures.ProcessPoolExecutor`` with the *spawn*
        start method, or a ``subprocess`` running a small driver
        script) and aggregate results afterwards. This is an
        upstream pythonfmu limitation, not a Jaxonomy one; FMUs from
        other exporters (OpenModelica, Dymola, Reference-FMUs) do
        not carry it.

    Args:
        file_name (str): path to FMU file
        dt (float): stepsize for FMU simulation
        name (str, optional): name of block
        input_names (list[str], optional): if set, only expose these inputs
        output_names (list[str], optional): if set, only expose these outputs
        parameters (dict, optional): dictionary of parameter overrides
        start_time (float, optional): FMU experiment start time.
        first_step_at_zero (bool, optional): Default ``False``. The
            first FMU step fires at ``t=dt`` (Modelica clocked-block
            convention: the periodic update is offset by one sample),
            so the block's outputs at ``t=0`` reflect the FMU's
            initial state (post-``setupExperiment``, pre-step) and
            only switch to "post-first-step" values from ``t=dt``
            onward. This introduces a one-sample phase lag versus an
            FMU exported with ``offset=0`` semantics, which weakens
            FMU round-trip byte-equivalence. Pass
            ``first_step_at_zero=True`` to fire the first step at
            ``t=0`` so the outputs leave their initial value as soon
            as the simulation begins. Surfaced as the FMU offset
            asymmetry in a follow-up finding.
        kwargs: ignored
    """
    try:
        super().__init__(name=name)
        self._init(
            file_name,
            dt,
            name=name or f"fmu_{self.system_id}",
            input_names=input_names,
            output_names=output_names,
            parameters=parameters,
            start_time=start_time,
            first_step_at_zero=first_step_at_zero,
        )
    except Exception as e:
        logger.error(
            "Failed to initialize FMU block %s (%s): %s", name, self.system_id, e
        )
        raise BlockInitializationError(str(e), system=self)

MuJoCo

Bases: MuJoCoBase

MuJoCo implementation without MJX.

Refer to MJX for the main docs.

Unlike the MJX variant of the block, this version uses the solver provided by mujoco itself and the physics are fully handled by mujoco. This behaves like a Co-Simulation environment.

This variant may be used to speed up compilation times or in situations where full JAX is not available or practical.

Source code in jaxonomy/library/mujoco.py

class MuJoCo(MuJoCoBase):
    """MuJoCo implementation without MJX.

    Refer to MJX for the main docs.

    Unlike the MJX variant of the block, this version uses the solver provided
    by mujoco itself and the physics are fully handled by mujoco. This behaves
    like a Co-Simulation environment.

    This variant may be used to speed up compilation times or in situations where
    full JAX is not available or practical.
    """

    def __init__(
        self,
        file_name: str,
        dt: float = 0.01,
        key_frame_0: int | str = None,
        qpos_0: Array = None,
        qvel_0: Array = None,
        act_0: Array = None,
        enable_sensor_data=False,
        enable_video_output=False,
        video_size: tuple[int, int] = None,
        enable_mocap_pos=False,
        custom_output_scripts: dict[str, str] = None,
        vHIL=False,
        vHIL_dt=0.01,
        **kwargs,
    ):
        super().__init__(
            use_mjx=False,
            file_name=file_name,
            dt=dt,
            key_frame_0=key_frame_0,
            qpos_0=qpos_0,
            qvel_0=qvel_0,
            act_0=act_0,
            enable_sensor_data=enable_sensor_data,
            enable_video_output=enable_video_output,
            video_size=video_size,
            enable_mocap_pos=enable_mocap_pos,
            custom_output_scripts=custom_output_scripts,
            vHIL=vHIL,
            vHIL_dt=vHIL_dt,
            **kwargs,
        )

        # This output cb implements the call to _step and is the reference callback
        # that all other outputs will depend on.
        def _qpos_cb(time, state, *inputs, **parameters):
            def cb(inputs):
                self._data.ctrl = inputs[0]
                if enable_mocap_pos:
                    self._data.mocap_pos[:] = inputs[1]
                mujoco.mj_step(self._model, self._data)
                qpos_normalized_quats = self.normalize_qpos_quat(self._qpos())
                return qpos_normalized_quats

            return io_callback(cb, self.qpos_0, inputs)

        self._step_cache_index = self.declare_output_port(
            _qpos_cb,
            default_value=self.qpos_0,
            requires_inputs=True,
            offset=dt,
            period=dt,
            name="qpos",
        )

        def _qvel_cb(time, state, *inputs, **parameters):
            return io_callback(self._qvel, self.qvel_0)

        self.declare_output_port(
            _qvel_cb,
            default_value=self.qvel_0,
            requires_inputs=True,
            offset=dt,
            period=dt,
            name="qvel",
            prerequisites_of_calc=[self._step_cache_index],
        )

        def _act_cb(time, state, *inputs, **parameters):
            return io_callback(self._act, self.act_0)

        self.declare_output_port(
            _act_cb,
            default_value=self.act_0,
            requires_inputs=True,
            offset=dt,
            period=dt,
            name="act",
            prerequisites_of_calc=[self._step_cache_index],
        )

        if enable_sensor_data:
            self._declare_sensor_data_port(dt)
        if enable_video_output:
            self._declare_video_output_port(video_size)
        if vHIL:
            self._declare_vhil_fake_output_port(vHIL_dt)
        self._declare_custom_output_ports(
            custom_output_scripts, dt, requires_inputs=False
        )

    # def post_simulation_finalize(self) -> None:
    #     # FIXME this should not be here but I had a "too many files opened" error
    #     self._model = None
    #     self._data = None
    #     return super().post_simulation_finalize()

    def _qpos(self, state=None):
        return self._data.qpos

    def _qvel(self, state=None):
        return self._data.qvel

    def _act(self, state=None):
        return self._data.act

    def _sensordata(self):
        return self._data.sensordata

    def normalize_qpos_quat(self, qpos):
        mujoco.mj_normalizeQuat(self._model, qpos)
        return qpos

    def render(self, data=None):
        if data is None:
            data = self._data
        return super().render(data)

    def _output_video(self, time, state, *inputs, **parameters):
        def cb(time):
            return self.render(self._data)

        return io_callback(
            cb,
            self._video_default,
            time,
        )

MultiPortSwitch

Bases: LeafSystem

Route one of N data signals based on an integer selector input.

Inputs are (selector, data_0, data_1, ..., data_{n-1}) and the output is data_{clip(round(selector), 0, n-1)}.

Implementation strategy: stack the data inputs along a new leading axis and pick out the selected slice with integer indexing. This is fully differentiable through the selected data input (zero gradient on the others), which is the standard documented semantics for this block. The selector is rounded and clipped to [0, n-1] so floating-point inputs are tolerated; the selector itself is non-differentiable (round/clip zero out the gradient).

All data inputs must share the same shape and dtype (the stack requires it). Mixed shapes are rejected by npa.stack at trace time.

Parameters:

Name	Type	Description	Default
n_data_inputs		number of data input ports. Must be >= 1.	required
choice_names		optional tuple of unique non-empty string labels, one per data input, used for self-documenting diagrams. When supplied its length MUST equal n_data_inputs and entries must be distinct. Resolve a friendly name to its integer selector via index_of(name) when wiring the diagram — see "Named choices" below. Default None (integer-only, byte-equivalent to phase 1).	None

Input ports

(0) selector — scalar integer-valued signal in [0, n-1]. Floating values are rounded and clipped. (1..n_data_inputs) data inputs.

Output ports

(0) The data input at index selector.

Notes

The original T-118 spec includes indexing="one-based" and mode="smooth" (softmax-blend across data inputs). Both are deferred — see T-118-followup-modes. Zero-based indexing is the only mode supported in phase 1, matching Python conventions.

Named choices (T-118-followup-multi-port-string-keys, 2026-05-13): choice_names=("low", "medium", "high") lets a diagram document which port means what. The selector port itself still expects an integer at runtime — JAX cannot trace strings, so this is the build-time-only interpretation the followup spec calls out. Look up the integer for a name on the Python side:

.. code-block:: python

    mps = MultiPortSwitch(3, choice_names=("low", "med", "high"))
    sel = library.Constant(mps.index_of("med"))   # → 1

Passing a string to ``index_of`` returns the matching integer;
passing an int returns it unchanged after a range check, so
callers can mix the two without branching. Unknown strings and
out-of-range ints raise ``BlockParameterError`` at construction
(build) time, not at trace time. The runtime _compute_output
path is unchanged when ``choice_names`` is ``None`` — the
default-off byte-equivalence guarantee.

Source code in jaxonomy/library/logic.py

class MultiPortSwitch(LeafSystem):
    """Route one of N data signals based on an integer selector input.

    Inputs are ``(selector, data_0, data_1, ..., data_{n-1})`` and the
    output is ``data_{clip(round(selector), 0, n-1)}``.

    Implementation strategy: stack the data inputs along a new leading
    axis and pick out the selected slice with integer indexing. This is
    fully differentiable through the *selected* data input (zero
    gradient on the others), which is the standard documented
    semantics for this block. The ``selector`` is rounded and clipped to ``[0, n-1]``
    so floating-point inputs are tolerated; the selector itself is
    non-differentiable (``round``/``clip`` zero out the gradient).

    All data inputs must share the same shape and dtype (the stack
    requires it). Mixed shapes are rejected by ``npa.stack`` at trace
    time.

    Parameters:
        n_data_inputs: number of data input ports. Must be ``>= 1``.
        choice_names: optional tuple of unique non-empty string labels,
            one per data input, used for self-documenting diagrams.
            When supplied its length MUST equal ``n_data_inputs`` and
            entries must be distinct. Resolve a friendly name to its
            integer selector via ``index_of(name)`` when wiring the
            diagram — see "Named choices" below. Default ``None``
            (integer-only, byte-equivalent to phase 1).

    Input ports:
        (0) selector — scalar integer-valued signal in ``[0, n-1]``.
            Floating values are rounded and clipped.
        (1..n_data_inputs) data inputs.

    Output ports:
        (0) The data input at index ``selector``.

    Notes:
        The original T-118 spec includes ``indexing="one-based"`` and
        ``mode="smooth"`` (softmax-blend across data inputs). Both are
        deferred — see T-118-followup-modes. Zero-based indexing is the
        only mode supported in phase 1, matching Python conventions.

    Named choices (T-118-followup-multi-port-string-keys, 2026-05-13):
        ``choice_names=("low", "medium", "high")`` lets a diagram
        document which port means what. The selector port itself still
        expects an integer at runtime — JAX cannot trace strings, so
        this is the build-time-only interpretation the followup spec
        calls out. Look up the integer for a name on the Python side:

        .. code-block:: python

            mps = MultiPortSwitch(3, choice_names=("low", "med", "high"))
            sel = library.Constant(mps.index_of("med"))   # → 1

        Passing a string to ``index_of`` returns the matching integer;
        passing an int returns it unchanged after a range check, so
        callers can mix the two without branching. Unknown strings and
        out-of-range ints raise ``BlockParameterError`` at construction
        (build) time, not at trace time. The runtime _compute_output
        path is unchanged when ``choice_names`` is ``None`` — the
        default-off byte-equivalence guarantee.
    """

    def __init__(self, n_data_inputs, choice_names=None, **kwargs):
        super().__init__(**kwargs)

        n = int(n_data_inputs)
        if n < 1:
            raise BlockParameterError(
                message=(
                    f"MultiPortSwitch block '{self.name}' requires "
                    f"n_data_inputs >= 1; got {n_data_inputs}."
                ),
                system=self,
                parameter_name="n_data_inputs",
            )
        self._n_data_inputs = n

        # Validate + store choice_names. ``None`` preserves the phase-1
        # path byte-for-byte: no extra branches in _compute_output, no
        # extra state visible to traces. Build-time-only by design —
        # see class docstring.
        if choice_names is not None:
            names = tuple(choice_names)
            if len(names) != n:
                raise BlockParameterError(
                    message=(
                        f"MultiPortSwitch block '{self.name}': "
                        f"choice_names has {len(names)} entries but "
                        f"n_data_inputs={n}."
                    ),
                    system=self,
                    parameter_name="choice_names",
                )
            for nm in names:
                if not isinstance(nm, str) or not nm:
                    raise BlockParameterError(
                        message=(
                            f"MultiPortSwitch block '{self.name}': "
                            f"choice_names entries must be non-empty "
                            f"strings; got {nm!r}."
                        ),
                        system=self,
                        parameter_name="choice_names",
                    )
            if len(set(names)) != len(names):
                raise BlockParameterError(
                    message=(
                        f"MultiPortSwitch block '{self.name}': "
                        f"choice_names entries must be unique; got "
                        f"{names!r}."
                    ),
                    system=self,
                    parameter_name="choice_names",
                )
            self._choice_names = names
            self._choice_index = {name: i for i, name in enumerate(names)}
        else:
            self._choice_names = None
            self._choice_index = None

        self.declare_input_port()  # selector
        for _ in range(n):
            self.declare_input_port()  # data_i

        def _compute_output(_time, _state, *inputs, **_params):
            selector = inputs[0]
            data = inputs[1:]
            stacked = npa.stack(data, axis=0)
            idx = npa.clip(npa.round(selector).astype(npa.int32), 0, n - 1)
            return stacked[idx]

        self.declare_output_port(
            _compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    @property
    def choice_names(self):
        """Tuple of channel labels, or ``None`` if unlabeled."""
        return self._choice_names

    def index_of(self, selector):
        """Resolve a string or int selector to its integer index.

        ``selector`` may be a string (looked up in ``choice_names``)
        or any object convertible via ``int()``. Strings only resolve
        when ``choice_names`` was supplied at construction.
        Out-of-range ints and unknown strings raise
        ``BlockParameterError`` at *build time*; runtime selectors
        flowing through the input port are still clipped silently by
        ``_compute_output`` (no change to the runtime path).
        """
        if isinstance(selector, str):
            if self._choice_index is None:
                raise BlockParameterError(
                    message=(
                        f"MultiPortSwitch block '{self.name}': string "
                        f"selector {selector!r} requires choice_names "
                        f"at construction."
                    ),
                    system=self,
                    parameter_name="choice_names",
                )
            if selector not in self._choice_index:
                raise BlockParameterError(
                    message=(
                        f"MultiPortSwitch block '{self.name}': unknown "
                        f"choice {selector!r}; valid options: "
                        + ",".join(self._choice_names)
                        + "."
                    ),
                    system=self,
                    parameter_name="choice_names",
                )
            return self._choice_index[selector]
        idx = int(selector)
        if idx < 0 or idx >= self._n_data_inputs:
            raise BlockParameterError(
                message=(
                    f"MultiPortSwitch block '{self.name}': integer "
                    f"selector {idx} out of range [0, "
                    f"{self._n_data_inputs - 1}]."
                ),
                system=self,
                parameter_name="choice_names",
            )
        return idx

choice_names property

Tuple of channel labels, or None if unlabeled.

index_of(selector)

Resolve a string or int selector to its integer index.

selector may be a string (looked up in choice_names) or any object convertible via int(). Strings only resolve when choice_names was supplied at construction. Out-of-range ints and unknown strings raise BlockParameterError at build time; runtime selectors flowing through the input port are still clipped silently by _compute_output (no change to the runtime path).

Source code in jaxonomy/library/logic.py

def index_of(self, selector):
    """Resolve a string or int selector to its integer index.

    ``selector`` may be a string (looked up in ``choice_names``)
    or any object convertible via ``int()``. Strings only resolve
    when ``choice_names`` was supplied at construction.
    Out-of-range ints and unknown strings raise
    ``BlockParameterError`` at *build time*; runtime selectors
    flowing through the input port are still clipped silently by
    ``_compute_output`` (no change to the runtime path).
    """
    if isinstance(selector, str):
        if self._choice_index is None:
            raise BlockParameterError(
                message=(
                    f"MultiPortSwitch block '{self.name}': string "
                    f"selector {selector!r} requires choice_names "
                    f"at construction."
                ),
                system=self,
                parameter_name="choice_names",
            )
        if selector not in self._choice_index:
            raise BlockParameterError(
                message=(
                    f"MultiPortSwitch block '{self.name}': unknown "
                    f"choice {selector!r}; valid options: "
                    + ",".join(self._choice_names)
                    + "."
                ),
                system=self,
                parameter_name="choice_names",
            )
        return self._choice_index[selector]
    idx = int(selector)
    if idx < 0 or idx >= self._n_data_inputs:
        raise BlockParameterError(
            message=(
                f"MultiPortSwitch block '{self.name}': integer "
                f"selector {idx} out of range [0, "
                f"{self._n_data_inputs - 1}]."
            ),
            system=self,
            parameter_name="choice_names",
        )
    return idx

Multiplexer

Bases: ReduceBlock

Stack the input signals into a single output signal.

Dispatches to jax.numpy.hstack, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.hstack.html

Input ports

(0..n_in-1) The input signals.

Output ports

(0) The stacked output signal.

Source code in jaxonomy/library/routing.py

class Multiplexer(ReduceBlock):
    """Stack the input signals into a single output signal.

    Dispatches to `jax.numpy.hstack`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.hstack.html

    Input ports:
        (0..n_in-1) The input signals.

    Output ports:
        (0) The stacked output signal.
    """

    def __init__(self, n_in, *args, **kwargs):
        super().__init__(n_in, npa.hstack, *args, **kwargs)

Mux

Bases: ReduceBlock

Stack n_inputs homogeneous signals into a single output signal.

This is the standard Mux block. It dispatches to npa.stack along axis 0, so:

• Mux(3)([1.0, 2.0, 3.0]) -> array([1.0, 2.0, 3.0]) (shape (3,)).
• Mux(2)([(1.0, 2.0), (3.0, 4.0)]) -> array([[1.0, 2.0], [3.0, 4.0]]) (shape (2, 2)).

All inputs must be the same shape and dtype; this matches the conventional Mux semantics and npa.stack's broadcasting rules.

For the older flatten-by-concatenation behavior (hstack), use :class:Multiplexer instead.

Input ports

(0..n_inputs-1) The input signals (must share shape and dtype).

Output ports

(0) The stacked output signal, with one extra leading axis.

Source code in jaxonomy/library/routing.py

class Mux(ReduceBlock):
    """Stack ``n_inputs`` homogeneous signals into a single output signal.

    This is the standard ``Mux`` block. It dispatches to ``npa.stack``
    along axis 0, so:

    * ``Mux(3)([1.0, 2.0, 3.0]) -> array([1.0, 2.0, 3.0])`` (shape ``(3,)``).
    * ``Mux(2)([(1.0, 2.0), (3.0, 4.0)]) -> array([[1.0, 2.0],
      [3.0, 4.0]])`` (shape ``(2, 2)``).

    All inputs must be the same shape and dtype; this matches the
    conventional ``Mux`` semantics and ``npa.stack``'s broadcasting rules.

    For the older flatten-by-concatenation behavior (``hstack``), use
    :class:`Multiplexer` instead.

    Input ports:
        (0..n_inputs-1) The input signals (must share shape and dtype).

    Output ports:
        (0) The stacked output signal, with one extra leading axis.
    """

    def __init__(self, n_inputs, *args, **kwargs):
        super().__init__(n_inputs, npa.stack, *args, **kwargs)

Notch

Bases: LeafSystem

Discrete biquad band-stop ("notch") filter.

Implements the textbook biquad notch::

omega0 = 2*pi * frequency_hz * dt
r      = 1 - pi * bandwidth_hz * dt          (pole radius)
rho2   = r*r + depth * (1 - r*r)             (zero radius**2)
b0, b1, b2 = 1, -2*sqrt(rho2)*cos(omega0), rho2
a1, a2     = -2*r*cos(omega0), r*r
y[k] = b0*x[k] + b1*x[k-1] + b2*x[k-2]
       - a1*y[k-1] - a2*y[k-2]

With depth = 1 the zeros sit on the unit circle and the notch is infinitely deep. With depth = 0 the numerator collapses to the denominator and the block becomes a unit-gain pass-through (useful as an identity check).

Input ports

(0) The input signal x.

Output ports

(0) The filtered signal y.

Parameters:

Name	Type	Description	Default
dt		Sampling period of the block (s).	required
frequency_hz		Notch centre frequency (Hz). Differentiable.	1.0
bandwidth_hz		Approximate -3 dB bandwidth of the notch (Hz). Differentiable. Must satisfy pi * bandwidth_hz * dt < 1 so the pole radius r remains in (0, 1).	0.1
depth		Notch depth in [0, 1). Default 0.99. Larger ⇒ deeper notch; depth = 1 puts the zeros on the unit circle for an infinitely deep notch (allowed but numerically borderline).	0.99
initial_state		Initial value of y[-1] (and, implicitly, y[-2], x[-1], x[-2]). Default 0.0.	0.0

Notes

Differentiability: frequency_hz, bandwidth_hz, and depth enter the biquad coefficients via smooth arithmetic (cos, sqrt), so jax.grad is finite through them.

The block is feedthrough on its input port (b0 = 1), so y[k] depends on x[k] directly.

Source code in jaxonomy/library/dynamics.py

class Notch(LeafSystem):
    """Discrete biquad band-stop ("notch") filter.

    Implements the textbook biquad notch::

        omega0 = 2*pi * frequency_hz * dt
        r      = 1 - pi * bandwidth_hz * dt          (pole radius)
        rho2   = r*r + depth * (1 - r*r)             (zero radius**2)
        b0, b1, b2 = 1, -2*sqrt(rho2)*cos(omega0), rho2
        a1, a2     = -2*r*cos(omega0), r*r
        y[k] = b0*x[k] + b1*x[k-1] + b2*x[k-2]
               - a1*y[k-1] - a2*y[k-2]

    With ``depth = 1`` the zeros sit on the unit circle and the notch is
    infinitely deep.  With ``depth = 0`` the numerator collapses to the
    denominator and the block becomes a unit-gain pass-through (useful
    as an identity check).

    Input ports:
        (0) The input signal ``x``.

    Output ports:
        (0) The filtered signal ``y``.

    Parameters:
        dt:
            Sampling period of the block (s).
        frequency_hz:
            Notch centre frequency (Hz).  Differentiable.
        bandwidth_hz:
            Approximate -3 dB bandwidth of the notch (Hz).  Differentiable.
            Must satisfy ``pi * bandwidth_hz * dt < 1`` so the pole
            radius ``r`` remains in ``(0, 1)``.
        depth:
            Notch depth in ``[0, 1)``.  Default 0.99.  Larger ⇒ deeper
            notch; ``depth = 1`` puts the zeros on the unit circle for an
            infinitely deep notch (allowed but numerically borderline).
        initial_state:
            Initial value of ``y[-1]`` (and, implicitly, ``y[-2]``,
            ``x[-1]``, ``x[-2]``).  Default 0.0.

    Notes:
        Differentiability: ``frequency_hz``, ``bandwidth_hz``, and
        ``depth`` enter the biquad coefficients via smooth arithmetic
        (``cos``, ``sqrt``), so ``jax.grad`` is finite through them.

        The block is feedthrough on its input port (``b0 = 1``), so
        ``y[k]`` depends on ``x[k]`` directly.
    """

    class DiscreteStateType(NamedTuple):
        x_prev1: Array  # x[k-1]
        x_prev2: Array  # x[k-2]
        y_prev1: Array  # y[k-1]
        y_prev2: Array  # y[k-2]

    @parameters(
        static=["dt"],
        dynamic=["frequency_hz", "bandwidth_hz", "depth", "initial_state"],
    )
    def __init__(
        self,
        dt,
        frequency_hz=1.0,
        bandwidth_hz=0.1,
        depth=0.99,
        initial_state=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.dt = dt
        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(
        self,
        frequency_hz,
        bandwidth_hz,
        depth,
        initial_state,
        dt=None,
    ):
        y0 = npa.asarray(initial_state)
        x0 = npa.zeros_like(y0)
        self.declare_discrete_state(
            default_value=self.DiscreteStateType(
                x_prev1=x0, x_prev2=x0, y_prev1=y0, y_prev2=y0
            ),
            as_array=False,
        )

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=self.dt,
        )

        # Feedthrough: y[k] depends on x[k] through b0 = 1.
        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=self.dt,
            default_value=y0,
            requires_inputs=True,
            prerequisites_of_calc=[
                DependencyTicket.xd,
                self.input_ports[0].ticket,
            ],
        )

    def reset_default_values(self, **dynamic_parameters):
        y0 = npa.asarray(dynamic_parameters["initial_state"])
        x0 = npa.zeros_like(y0)
        self.configure_discrete_state_default_value(
            self.DiscreteStateType(
                x_prev1=x0, x_prev2=x0, y_prev1=y0, y_prev2=y0
            ),
            as_array=False,
        )
        self.configure_output_port_default_value(self._output_port_idx, y0)

    def _coeffs(self, frequency_hz, bandwidth_hz, depth):
        # ``omega0 = 2*pi*f*dt`` (digital angular frequency at the notch).
        # ``r = 1 - pi*bw*dt`` is the pole radius; the standard
        # approximation that gives an *approximate* -3 dB bandwidth of
        # ``bandwidth_hz`` (cf. Steiglitz, "A Digital Signal Processing
        # Primer", §9.4).
        two_pi = 2.0 * npa.pi
        omega0 = two_pi * frequency_hz * self.dt
        r = 1.0 - npa.pi * bandwidth_hz * self.dt

        cos_w0 = npa.cos(omega0)
        r2 = r * r

        # depth in [0, 1]: 0 ⇒ rho2 = r2 (numerator = denominator ⇒
        # pass-through), 1 ⇒ rho2 = 1 (zeros on the unit circle ⇒
        # infinitely deep notch).
        rho2 = r2 + depth * (1.0 - r2)
        rho = npa.sqrt(rho2)

        b0 = 1.0
        b1 = -2.0 * rho * cos_w0
        b2 = rho2
        a1 = -2.0 * r * cos_w0
        a2 = r2

        # Normalise so DC gain ``H(1) = (b0+b1+b2)/(1+a1+a2)`` equals 1.
        # Without this, the textbook biquad has a slight (~1 %) DC bump.
        # Scaling the numerator zeros uniformly preserves the on-notch
        # attenuation.
        dc_gain = (b0 + b1 + b2) / (1.0 + a1 + a2)
        b0 = b0 / dc_gain
        b1 = b1 / dc_gain
        b2 = b2 / dc_gain
        return b0, b1, b2, a1, a2

    def _update(self, _time, state, *inputs, **params):
        x = inputs[0]
        b0, b1, b2, a1, a2 = self._coeffs(
            params["frequency_hz"],
            params["bandwidth_hz"],
            params["depth"],
        )
        xd = state.discrete_state
        y_new = (
            b0 * x
            + b1 * xd.x_prev1
            + b2 * xd.x_prev2
            - a1 * xd.y_prev1
            - a2 * xd.y_prev2
        )
        return self.DiscreteStateType(
            x_prev1=x,
            x_prev2=xd.x_prev1,
            y_prev1=y_new,
            y_prev2=xd.y_prev1,
        )

    def _output(self, _time, state, *inputs, **params):
        # Feedthrough output: recompute y[k] from x[k] and the stored
        # delay-line so the readout matches the update.
        x = inputs[0]
        b0, b1, b2, a1, a2 = self._coeffs(
            params["frequency_hz"],
            params["bandwidth_hz"],
            params["depth"],
        )
        xd = state.discrete_state
        return (
            b0 * x
            + b1 * xd.x_prev1
            + b2 * xd.x_prev2
            - a1 * xd.y_prev1
            - a2 * xd.y_prev2
        )

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        u = self.eval_input(context)
        xd = context[self.system_id].discrete_state.y_prev1
        check_state_type(
            self,
            inp_data=u,
            state_data=xd,
            error_collector=error_collector,
        )

ONNX

Bases: LeafSystem

ONNX inference block.

Parameters:

Name	Type	Description	Default
file_name	str	Path to the .onnx model file.	required
num_inputs	int	Number of input tensors the model expects.	1
num_outputs	int	Number of output tensors the model produces.	1
cast_outputs_to_dtype		Optional jnp dtype name ("float32", "float64", etc.) to cast every output to. If None, the output dtype matches the model's output spec.	None
providers		onnxruntime execution providers; defaults to CPU. Pass e.g. ("CUDAExecutionProvider", "CPUExecutionProvider") on a GPU host.	('CPUExecutionProvider',)
name		Optional block name.	required

Notes

Differentiability is best-effort — jax.pure_callback does not define a VJP, so reverse-mode autodiff through the block raises. For end-to-end gradients, look at the T-023a follow-up on a JAX-traceable conversion (onnx2jax or similar).

.. note:: float32 artifacts under jaxonomy's global x64. import jaxonomy enables jax_enable_x64 process-wide, so a float32 ONNX model receives float64 inputs unless you cast at the block boundary — a silent arithmetic change relative to the framework the model was exported and validated in. One-line idiom: pass cast_outputs_to_dtype="float32" and feed the block x.astype(jnp.float32) inputs.

Source code in jaxonomy/library/onnx_block.py

class ONNX(LeafSystem):
    """ONNX inference block.

    Args:
        file_name: Path to the ``.onnx`` model file.
        num_inputs: Number of input tensors the model expects.
        num_outputs: Number of output tensors the model produces.
        cast_outputs_to_dtype: Optional ``jnp`` dtype name (``"float32"``,
            ``"float64"``, etc.) to cast every output to.  If None, the
            output dtype matches the model's output spec.
        providers: ``onnxruntime`` execution providers; defaults to CPU.
            Pass e.g. ``("CUDAExecutionProvider", "CPUExecutionProvider")``
            on a GPU host.
        name: Optional block name.

    Notes:
        Differentiability is **best-effort** — ``jax.pure_callback`` does
        not define a VJP, so reverse-mode autodiff through the block
        raises.  For end-to-end gradients, look at the T-023a follow-up
        on a JAX-traceable conversion (``onnx2jax`` or similar).

    .. note:: **float32 artifacts under jaxonomy's global x64.**
       ``import jaxonomy`` enables ``jax_enable_x64`` process-wide, so a
       float32 ONNX model receives float64 inputs unless you cast at the
       block boundary — a silent arithmetic change relative to the
       framework the model was exported and validated in.  One-line
       idiom: pass ``cast_outputs_to_dtype="float32"`` and feed the
       block ``x.astype(jnp.float32)`` inputs.
    """

    @parameters(
        static=[
            "file_name",
            "num_inputs",
            "num_outputs",
            "cast_outputs_to_dtype",
            "providers",
        ]
    )
    def __init__(
        self,
        file_name: str,
        num_inputs: int = 1,
        num_outputs: int = 1,
        cast_outputs_to_dtype=None,
        providers=("CPUExecutionProvider",),
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self._num_inputs = int(num_inputs)
        self._num_outputs = int(num_outputs)

        for _ in range(self._num_inputs):
            self.declare_input_port()

        def _make_output_callback(idx):
            def _cb(time, state, *inputs, **params):
                outs = self._evaluate(time, state, *inputs, **params)
                return outs[idx]
            return _cb

        for i in range(self._num_outputs):
            self.declare_output_port(
                _make_output_callback(i), requires_inputs=True,
            )

    def initialize(
        self,
        file_name: str,
        num_inputs: int = 1,
        num_outputs: int = 1,
        cast_outputs_to_dtype=None,
        providers=("CPUExecutionProvider",),
    ):
        if num_inputs != self._num_inputs:
            raise ValueError(
                f"ONNX: num_inputs cannot be changed after construction "
                f"({self._num_inputs} → {num_inputs})."
            )
        if num_outputs != self._num_outputs:
            raise ValueError(
                f"ONNX: num_outputs cannot be changed after construction "
                f"({self._num_outputs} → {num_outputs})."
            )

        self._dtype_output = (
            getattr(jnp, cast_outputs_to_dtype)
            if cast_outputs_to_dtype is not None
            else None
        )

        self._session = ort.InferenceSession(
            file_name, providers=list(providers),
        )
        self._input_names = [i.name for i in self._session.get_inputs()]
        self._output_names = [o.name for o in self._session.get_outputs()]

        if len(self._input_names) != self._num_inputs:
            raise ValueError(
                f"ONNX: model has {len(self._input_names)} inputs but "
                f"the block declared num_inputs={self._num_inputs}.  "
                f"Model input names: {self._input_names}"
            )
        if len(self._output_names) != self._num_outputs:
            raise ValueError(
                f"ONNX: model has {len(self._output_names)} outputs but "
                f"the block declared num_outputs={self._num_outputs}.  "
                f"Model output names: {self._output_names}"
            )

    # ── shape inference and runtime ───────────────────────────────────────

    def initialize_static_data(self, context):
        try:
            inputs = self.collect_inputs(context)
            outs = self._pure_callback(*inputs)
            self._result_type = [
                jax.ShapeDtypeStruct(o.shape, o.dtype) for o in outs
            ]
        except UpstreamEvalError:
            logger.debug(
                "ONNX.initialize_static_data: UpstreamEvalError, deferring "
                "shape inference until root context is built."
            )
        return super().initialize_static_data(context)

    def _evaluate(self, time, state, *inputs, **params):
        return jax.pure_callback(
            self._pure_callback,
            self._result_type,
            *inputs,
        )

    def _pure_callback(self, *inputs):
        feed = {
            name: np.asarray(value)
            for name, value in zip(self._input_names, inputs)
        }
        outputs_np = self._session.run(self._output_names, feed)
        if self._dtype_output is not None:
            outputs_jax = [jnp.array(o, self._dtype_output) for o in outputs_np]
        else:
            outputs_jax = [jnp.array(o) for o in outputs_np]
        return outputs_jax

ONNXJax

Bases: LeafSystem

JAX-traceable ONNX inference (T-023a).

Parameters:

Name	Type	Description	Default
file_name	str	Path to the .onnx model file.	required
num_inputs	int	Number of input tensors the model expects.	1
num_outputs	int	Number of output tensors the model produces.	1
cast_outputs_to_dtype		Optional jnp dtype name to cast every output to ("float32" / "float64" / ...).	None
name		Optional block name.	required

Differentiability: end-to-end via jaxonnxruntime's JAX primitive implementations. Op coverage failure shows up at initialize() time as a clear RuntimeError from jaxonnxruntime.

For models that use ops outside jaxonnxruntime's coverage, fall back to :class:ONNX — same constructor signature, runs via onnxruntime host callback (no autodiff).

.. note:: float32 artifacts under jaxonomy's global x64. import jaxonomy enables jax_enable_x64 process-wide, so a float32 model here computes against float64 inputs unless you cast at the block boundary. One-line idiom: pass cast_outputs_to_dtype="float32" and feed the block x.astype(jnp.float32) inputs.

.. note:: Imported discrete-time policies need a ZeroOrderHold. A sample-and-hold controller exported from a discrete-time training loop (torch / NEUROMANCER-style: compute :math:u_k once per sample, hold for ts) is re-evaluated at every ODE solver stage (e.g. all four RK4 stages) when wired directly into a continuous plant — continuous-feedback semantics. Both loops "work", but step-for-step parity with the exporting framework is silently destroyed. Follow the block with ZeroOrderHold(dt=ts) and pin the step grid with SimulatorOptions(max_major_step_length=ts, max_minor_step_size=ts); with that, closed-loop parity is ~4e-8 over 400 steps on the two-tank benchmark.

Source code in jaxonomy/library/onnx_jax_block.py

class ONNXJax(LeafSystem):
    """JAX-traceable ONNX inference (T-023a).

    Args:
        file_name: Path to the ``.onnx`` model file.
        num_inputs: Number of input tensors the model expects.
        num_outputs: Number of output tensors the model produces.
        cast_outputs_to_dtype: Optional ``jnp`` dtype name to cast every
            output to (``"float32"`` / ``"float64"`` / ...).
        name: Optional block name.

    Differentiability: end-to-end via ``jaxonnxruntime``'s JAX
    primitive implementations.  Op coverage failure shows up at
    initialize() time as a clear ``RuntimeError`` from
    ``jaxonnxruntime``.

    For models that use ops outside ``jaxonnxruntime``'s coverage,
    fall back to :class:`ONNX` — same constructor signature, runs via
    ``onnxruntime`` host callback (no autodiff).

    .. note:: **float32 artifacts under jaxonomy's global x64.**
       ``import jaxonomy`` enables ``jax_enable_x64`` process-wide, so a
       float32 model here computes against float64 inputs unless you
       cast at the block boundary.  One-line idiom: pass
       ``cast_outputs_to_dtype="float32"`` and feed the block
       ``x.astype(jnp.float32)`` inputs.

    .. note:: **Imported discrete-time policies need a ZeroOrderHold.**
       A sample-and-hold controller exported from a discrete-time
       training loop (torch / NEUROMANCER-style: compute :math:`u_k`
       once per sample, hold for ``ts``) is re-evaluated at every ODE
       solver stage (e.g. all four RK4 stages) when wired directly into
       a continuous plant — continuous-feedback semantics.  Both loops
       "work", but step-for-step parity with the exporting framework is
       silently destroyed.  Follow the block with
       ``ZeroOrderHold(dt=ts)`` and pin the step grid with
       ``SimulatorOptions(max_major_step_length=ts,
       max_minor_step_size=ts)``; with that, closed-loop parity is
       ~4e-8 over 400 steps on the two-tank benchmark.
    """

    @parameters(
        static=[
            "file_name", "num_inputs", "num_outputs", "cast_outputs_to_dtype",
        ]
    )
    def __init__(
        self,
        file_name: str,
        num_inputs: int = 1,
        num_outputs: int = 1,
        cast_outputs_to_dtype=None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self._num_inputs = int(num_inputs)
        self._num_outputs = int(num_outputs)

        for _ in range(self._num_inputs):
            self.declare_input_port()

        def _make_output_callback(idx):
            def _cb(time, state, *inputs, **params):
                outs = self._evaluate(time, state, *inputs, **params)
                return outs[idx]
            return _cb

        for i in range(self._num_outputs):
            self.declare_output_port(
                _make_output_callback(i), requires_inputs=True,
            )

    def initialize(
        self,
        file_name: str,
        num_inputs: int = 1,
        num_outputs: int = 1,
        cast_outputs_to_dtype=None,
    ):
        if num_inputs != self._num_inputs:
            raise ValueError(
                f"ONNXJax: num_inputs cannot be changed after construction "
                f"({self._num_inputs} → {num_inputs})."
            )
        if num_outputs != self._num_outputs:
            raise ValueError(
                f"ONNXJax: num_outputs cannot be changed after construction "
                f"({self._num_outputs} → {num_outputs})."
            )

        self._dtype_output = (
            getattr(jnp, cast_outputs_to_dtype)
            if cast_outputs_to_dtype is not None
            else None
        )

        # Load model + prepare a JAX-callable executor.
        model = onnx.load(file_name)
        try:
            self._rep = jort_backend.Backend.prepare(model)
        except Exception as e:
            raise RuntimeError(
                f"ONNXJax.initialize: jaxonnxruntime failed to prepare "
                f"model {file_name!r}: {e}.  This usually means the "
                "model uses ops that jaxonnxruntime hasn't implemented "
                "yet.  Fall back to ONNX (host-callback) if you don't "
                "need gradients."
            ) from e

        # Cache input/output names for ordering.
        self._input_names = [i.name for i in model.graph.input]
        # Filter to real inputs only (initializer-shadowed inputs are not user-supplied)
        initializer_names = {init.name for init in model.graph.initializer}
        self._input_names = [n for n in self._input_names if n not in initializer_names]
        self._output_names = [o.name for o in model.graph.output]

        if len(self._input_names) != self._num_inputs:
            raise ValueError(
                f"ONNXJax: model has {len(self._input_names)} inputs "
                f"but the block declared num_inputs={self._num_inputs}.  "
                f"Model inputs: {self._input_names}"
            )
        if len(self._output_names) != self._num_outputs:
            raise ValueError(
                f"ONNXJax: model has {len(self._output_names)} outputs "
                f"but the block declared num_outputs={self._num_outputs}.  "
                f"Model outputs: {self._output_names}"
            )

    # ── shape inference & runtime ─────────────────────────────────────────

    def initialize_static_data(self, context):
        try:
            inputs = self.collect_inputs(context)
            outs = self._run(inputs)
            self._result_type = [
                jax.ShapeDtypeStruct(o.shape, o.dtype) for o in outs
            ]
        except UpstreamEvalError:
            logger.debug(
                "ONNXJax.initialize_static_data: UpstreamEvalError, "
                "deferring shape inference until root context is built."
            )
        return super().initialize_static_data(context)

    def _run(self, inputs):
        outs = self._rep.run(list(inputs))
        if self._dtype_output is not None:
            outs = [jnp.asarray(o, self._dtype_output) for o in outs]
        else:
            outs = [jnp.asarray(o) for o in outs]
        return outs

    def _evaluate(self, time, state, *inputs, **params):
        # Run inline (not via pure_callback) — jaxonnxruntime emits
        # JAX primitives directly.
        return self._run(inputs)

Offset

Bases: FeedthroughBlock

Add a constant offset or bias to the input signal.

Given an input signal u and offset value b, this will return y = u + b.

Input ports

(0) The input signal.

Output ports

(0) The input signal plus the offset.

Parameters:

Name	Type	Description	Default
offset		The constant offset to add to the input signal.	required

Source code in jaxonomy/library/math_ops.py

class Offset(FeedthroughBlock):
    """Add a constant offset or bias to the input signal.

    Given an input signal `u` and offset value `b`, this will return `y = u + b`.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The input signal plus the offset.

    Parameters:
        offset:
            The constant offset to add to the input signal.
    """

    @parameters(dynamic=["offset"])
    def __init__(self, offset, *args, **kwargs):
        super().__init__(lambda x, offset: x + offset, *args, **kwargs)

    def initialize(self, offset):
        pass

OperatingPoint dataclass

Result of :func:findop.

Attributes:

Name	Type	Description
x	Any	Equilibrium continuous state.
u	Any	Input value held fixed during the search (taken from base_context at call time).
residual_norm	float	Final ‖ẋ(x*, u)‖_∞ after Newton iterations.
converged	bool	True if residual_norm met tol within max_iter steps.
iterations	int	Number of Newton iterations actually executed.

Source code in jaxonomy/library/linearization_workflow.py

@dataclass
class OperatingPoint:
    """Result of :func:`findop`.

    Attributes:
        x: Equilibrium continuous state.
        u: Input value held fixed during the search (taken from
            ``base_context`` at call time).
        residual_norm: Final ``‖ẋ(x*, u)‖_∞`` after Newton iterations.
        converged: True if ``residual_norm`` met ``tol`` within
            ``max_iter`` steps.
        iterations: Number of Newton iterations actually executed.
    """

    x: Any
    u: Any
    residual_norm: float
    converged: bool
    iterations: int

PCEModel

Fitted polynomial-chaos expansion y = sum_k c_k Psi_k(xi).

The orthonormal basis (Askey scheme) gives closed-form statistics: the mean is the constant coefficient, the variance is the sum of squared non-constant coefficients, and Sobol indices follow from partitioning that sum by which inputs each basis term depends on (Xiu & Karniadakis 2002; Sudret, Reliab. Eng. Syst. Saf. 93(7):964--979, 2008).

Source code in jaxonomy/library/rom/surrogates.py

class PCEModel:
    """Fitted polynomial-chaos expansion ``y = sum_k c_k Psi_k(xi)``.

    The orthonormal basis (Askey scheme) gives closed-form statistics: the mean
    is the constant coefficient, the variance is the sum of squared non-constant
    coefficients, and Sobol indices follow from partitioning that sum by which
    inputs each basis term depends on (Xiu & Karniadakis 2002; Sudret,
    *Reliab. Eng. Syst. Saf.* 93(7):964--979, 2008).
    """

    def __init__(self, coeffs, multi_indices, types, loc, scale, order):
        self.coeffs = coeffs
        self.multi_indices = multi_indices  # np.ndarray (K, dim)
        self.types = types
        self.loc = loc
        self.scale = scale
        self.order = int(order)
        self.dim = multi_indices.shape[1]
        # index of the all-zero (constant) multi-index
        self._const_idx = int(np.argmin(multi_indices.sum(axis=1)))

    def predict(self, Xstar):
        """Surrogate response at ``Xstar`` (jax-traceable)."""
        Xstar = _as2d(Xstar)
        Xi = _pce_standardize(Xstar, self.loc, self.scale)
        Psi = _pce_design(Xi, self.multi_indices, self.types, self.order)
        return Psi @ self.coeffs

    def mean(self):
        """Analytic mean = constant-term coefficient."""
        return self.coeffs[self._const_idx]

    def variance(self):
        """Analytic variance = sum of squared non-constant coefficients."""
        mask = np.ones(self.coeffs.shape[0], dtype=bool)
        mask[self._const_idx] = False
        return jnp.sum(jnp.asarray(self.coeffs)[jnp.asarray(mask)] ** 2)

    def sobol_indices(self):
        """Main-effect (first-order) and total Sobol indices per input.

        Returns a dict ``{"first_order": (dim,), "total": (dim,)}``.
        """
        c2 = np.asarray(self.coeffs) ** 2
        idx = self.multi_indices
        nonconst = idx.sum(axis=1) > 0
        total_var = c2[nonconst].sum()

        first = np.zeros(self.dim)
        total = np.zeros(self.dim)
        for i in range(self.dim):
            involves_i = idx[:, i] > 0
            others_zero = (idx.sum(axis=1) == idx[:, i])
            main_mask = involves_i & others_zero
            first[i] = c2[main_mask].sum() / total_var
            total[i] = c2[involves_i].sum() / total_var
        return {"first_order": jnp.asarray(first), "total": jnp.asarray(total)}

mean()

Analytic mean = constant-term coefficient.

Source code in jaxonomy/library/rom/surrogates.py

def mean(self):
    """Analytic mean = constant-term coefficient."""
    return self.coeffs[self._const_idx]

predict(Xstar)

Surrogate response at Xstar (jax-traceable).

Source code in jaxonomy/library/rom/surrogates.py

def predict(self, Xstar):
    """Surrogate response at ``Xstar`` (jax-traceable)."""
    Xstar = _as2d(Xstar)
    Xi = _pce_standardize(Xstar, self.loc, self.scale)
    Psi = _pce_design(Xi, self.multi_indices, self.types, self.order)
    return Psi @ self.coeffs

sobol_indices()

Main-effect (first-order) and total Sobol indices per input.

Returns a dict {"first_order": (dim,), "total": (dim,)}.

Source code in jaxonomy/library/rom/surrogates.py

def sobol_indices(self):
    """Main-effect (first-order) and total Sobol indices per input.

    Returns a dict ``{"first_order": (dim,), "total": (dim,)}``.
    """
    c2 = np.asarray(self.coeffs) ** 2
    idx = self.multi_indices
    nonconst = idx.sum(axis=1) > 0
    total_var = c2[nonconst].sum()

    first = np.zeros(self.dim)
    total = np.zeros(self.dim)
    for i in range(self.dim):
        involves_i = idx[:, i] > 0
        others_zero = (idx.sum(axis=1) == idx[:, i])
        main_mask = involves_i & others_zero
        first[i] = c2[main_mask].sum() / total_var
        total[i] = c2[involves_i].sum() / total_var
    return {"first_order": jnp.asarray(first), "total": jnp.asarray(total)}

variance()

Analytic variance = sum of squared non-constant coefficients.

Source code in jaxonomy/library/rom/surrogates.py

def variance(self):
    """Analytic variance = sum of squared non-constant coefficients."""
    mask = np.ones(self.coeffs.shape[0], dtype=bool)
    mask[self._const_idx] = False
    return jnp.sum(jnp.asarray(self.coeffs)[jnp.asarray(mask)] ** 2)

PID

Bases: LTISystem

Continuous-time PID controller.

The PID controller is implemented as a state-space system with matrices (A, B, C, D), which are then used to create a (second-order) LTISystem. Note that this only supports single-input, single-output PID controllers.

The PID controller implements the following control law:

    u = kp * e + ki * ∫e + kd * ė

where e is the error signal, and ∫e and ė are the integral and derivative of the error signal, respectively.

With a filter coefficient of n (to make the transfer function proper), the state-space form of the system is:

A = [[0, 1], [0, -n]]
B = [[0], [1]]
C = [[ki * n, (kp * n + ki) - (kp + kd * n) * n]]
D = [[kp + kd * n]]

Since D is nonzero, the block is feedthrough.

Input ports

(0) e: Error signal (scalar)

Output ports

(0) u: Control signal (scalar)

Parameters:

Name	Type	Description	Default
kp		Proportional gain	required
ki		Integral gain	required
kd		Derivative gain	required
n		Derivative filter coefficient	required
initial_state		Initial state of the integral term (default: 0)	0.0

Source code in jaxonomy/library/linear_system.py

class PID(LTISystem):
    """Continuous-time PID controller.

    The PID controller is implemented as a state-space system with matrices
    (A, B, C, D), which are then used to create a (second-order) LTISystem.
    Note that this only supports single-input, single-output PID controllers.

    The PID controller implements the following control law:
    ```
        u = kp * e + ki * ∫e + kd * ė
    ```
    where e is the error signal, and ∫e and ė are the integral and derivative
    of the error signal, respectively.

    With a filter coefficient of `n` (to make the transfer function proper), the
    state-space form of the system is:
    ```
    A = [[0, 1], [0, -n]]
    B = [[0], [1]]
    C = [[ki * n, (kp * n + ki) - (kp + kd * n) * n]]
    D = [[kp + kd * n]]
    ```

    Since D is nonzero, the block is feedthrough.

    Input ports:
        (0) e: Error signal (scalar)

    Output ports:
        (0) u: Control signal (scalar)

    Parameters:
        kp: Proportional gain
        ki: Integral gain
        kd: Derivative gain
        n: Derivative filter coefficient
        initial_state: Initial state of the integral term (default: 0)
    """

    @parameters(
        dynamic=["kp", "ki", "kd", "n"],
        static=["initial_state"],
    )
    def __init__(
        self,
        kp,
        ki,
        kd,
        n,
        initial_state=0.0,
        enable_external_initial_state=False,
        **kwargs,
    ):
        if enable_external_initial_state:
            raise NotImplementedError(
                "External initial state not yet implemented for PID"
            )

        A, B, C, D = self._get_abcd(kp, ki, kd, n)
        initialize_states = npa.array([initial_state, 0.0])
        super().__init__(A, B, C, D, initialize_states=initialize_states, **kwargs)

    def _get_abcd(self, kp, ki, kd, n):
        A = npa.array([[0.0, 1.0], [0.0, -n]])
        B = npa.array([[0.0], [1.0]])
        C = npa.array([(ki * n), ((kp * n + ki) - (kp + kd * n) * n)])
        D = npa.array([(kp + kd * n)])
        return A, B, C, D

    def _eval_output(self, time, state, *inputs, **params):
        kp, ki, kd, n = params["kp"], params["ki"], params["kd"], params["n"]

        A, B, C, D = self._get_abcd(kp, ki, kd, n)
        A, B, C, D, _, _, _ = _reshape(A, B, C, D)

        return self._eval_output_base(C, D, state, *inputs)

    def ode(self, time, state, u, **params):
        kp, ki, kd, n = params["kp"], params["ki"], params["kd"], params["n"]

        A, B, C, D = self._get_abcd(kp, ki, kd, n)
        A, B, C, D, _, _, _ = _reshape(A, B, C, D)

        return super().ode(time, state, u, A=A, B=B)

    def initialize(self, kp, ki, kd, n, initial_state, **kwargs):
        A, B, C, D = self._get_abcd(kp, ki, kd, n)
        initialize_states = npa.array([initial_state, 0.0])
        self._init_state(A, B, C, D, initialize_states)

PIDController2DOF

Bases: LeafSystem

Two-degree-of-freedom discrete-time PID controller.

Implements the standard 2-DOF PID control law::

u = Kp * (b*r - y) + Ki * integral(r - y) + Kd * d/dt(c*r - y)

where r is the setpoint, y is the measurement, and b and c are setpoint weights in [0, 1] for the proportional and derivative paths respectively. With b = c = 1 the block is numerically equivalent to the existing :class:PIDDiscrete block on the error signal e = r - y; with b = c = 0 it becomes an "I-PD" controller (only the integral term reacts to setpoint changes).

The integral term uses a forward-Euler approximation::

e_int[k+1] = e_int[k] + (r[k] - y[k]) * dt

and the derivative term is computed exactly as for :class:DerivativeDiscrete / :class:PIDDiscrete, including the optional first-order filter (filter_type, filter_coefficient).

Input ports

(0) Setpoint signal r. (1) Measurement signal y. Dynamic-port appendices, declared in this deterministic order when the corresponding *_dynamic flag is True:

(a) b if b_dynamic=True (T-127-followup-external-weights) (b) c if c_dynamic=True (c) kp if kp_dynamic=True (T-127-followup-gain-scheduling) (d) ki if ki_dynamic=True (e) kd if kd_dynamic=True (f) kff if kff_dynamic=True (g) u_ext if tracking_enabled=True (T-127-followup-tracking-mode) (h) mode_flag if tracking_enabled_dynamic=True (T-127-followup-bumpless-mode-switch). Scalar input cast to a {0, 1} gate that selects whether the tracking-pull branch runs this tick (0 = OFF / AUTO, non-zero = ON / MANUAL / TRACKING). Requires tracking_enabled=True so the u_ext port exists.

Port indices skip any flag set to False, so e.g. with only kp_dynamic=True the kp port is at index 2; with b_dynamic=True + kp_dynamic=True the b port is at index 2 and kp is at index 3. The instance attributes self.b_index / self.c_index / self.kp_index / self.ki_index / self.kd_index / self.kff_index / self.u_ext_index / self.mode_flag_index expose the resolved positions.

Output ports

(0) The control signal u computed by the 2-DOF PID law.

Parameters:

Name	Type	Description	Default
kp		Proportional gain (scalar).	1.0
ki		Integral gain (scalar).	1.0
kd		Derivative gain (scalar).	1.0
b		Setpoint weight for the proportional term, in [0, 1]. Default 1.0 (matches 1-DOF PID). Ignored when b_dynamic=True (the runtime port value is used instead).	1.0
c		Setpoint weight for the derivative term, in [0, 1]. Default 1.0 (matches 1-DOF PID). Ignored when c_dynamic=True. Recommendation for real-world controllers: prefer c = 0 (a.k.a. "derivative on measurement only"). With c = 1 (textbook 2-DOF / 1-DOF PID), a step change in the setpoint produces a one-tick spike of magnitude Kd / dt in d/dt(c*r - y) — "derivative kick" — that propagates into a brief, often saturation-inducing control transient. Setting c = 0 routes the derivative through the measurement only (-d/dt(y)), so setpoint steps no longer kick the derivative term and integral / proportional action alone drive the response. See the :meth:with_derivative_on_measurement factory and the derivative_on_measurement_only=True convenience kwarg below for the standard recipe.	1.0
derivative_on_measurement_only		T-127-followup-derivative-on-measurement. Convenience flag equivalent to c=0 + c_dynamic=False. When True the derivative term sees only the measurement (-d/dt(y)) and is immune to setpoint-step kick; the user-supplied c / c_dynamic kwargs are rejected with a ValueError to keep the contract unambiguous. Default False (c=1) preserves byte-equivalence with phase 1. See :meth:with_derivative_on_measurement for the matching factory function.	required
b_dynamic		If True, the proportional setpoint weight b is read from an additional input port (index 2) rather than from the static b parameter. The static b value is then ignored at runtime; the user MUST connect a signal to the new port. Mirrors the enable_dynamic_* pattern used by :class:Saturate and :class:RateLimiter. Default False (byte-equivalent to phase 1).	False
c_dynamic		If True, the derivative setpoint weight c is read from an additional input port instead of the static c parameter. The new port lives at index 2 (when only c_dynamic is set) or index 3 (when both b_dynamic and c_dynamic are set). Default False.	False
dt		Sampling period of the block.	required
initial_state		Initial value of the integral. Default 0.0.	0.0
filter_type		One of "none", "forward", "backward", or "bilinear" — derivative-filter mode. Default "none".	'none'
filter_coefficient		Filter coefficient N for the derivative filter (the conventional "filter coefficient" PID-tuning parameter). Default 1.0.	1.0
output_min		Lower saturation limit on the control output (T-127-followup- anti-windup). None (default) disables the lower clip. When either output_min or output_max is set the saturated control value is published on port (0); the unsaturated value also feeds the anti-windup correction.	None
output_max		Upper saturation limit on the control output. None (default) disables the upper clip.	None
anti_windup_method		One of "none", "back_calc", or "clamping" (T-127-followup-anti-windup). Default "none". "none" — no anti-windup; integrator update unchanged. "back_calc" — back-calculation: subtract (u_unsat - u_sat) / anti_windup_gain * dt from the integrator each tick. Smooth, fully differentiable. "clamping" — integrator-tracking: only update the integral when the controller is not pushing further into saturation (u_unsat == u_sat OR the error sign points away from the saturated direction). Anti-windup is a no-op unless output_min or output_max is also set.	'none'
anti_windup_gain		Tracking time constant Tt used by "back_calc". Smaller values pull the integrator back faster. Default 1.0. Differentiable through jax.grad.	1.0
integrator_method		One of "forward_euler" (default), "backward_euler", or "trapezoidal" — selects the discretisation used to advance the integral term (T-127-followup-discrete-integrator- derivative). Backward-Euler is more stable for stiff loops; trapezoidal is more accurate. Defaults to byte-equivalence with phase 1. Non-default values add an extra e_i_prev delay cell to the discrete state.	'forward_euler'
derivative_method		One of "forward_diff" (default), "backward_diff", or "centered_diff" — selects the unfiltered derivative kernel. "backward_diff" uses past samples only (typical for real-time control, introduces a one-tick delay relative to "forward_diff"). "centered_diff" is less noisy but adds one extra delay cell (e_d_prev_prev). Only applies when filter_type='none' — combining a non-default derivative_method with a recursive filter raises a ValueError at construction.	'forward_diff'
kff		Feedforward gain on the setpoint (T-127-followup-feedforward). Adds kff * r to the PID output before saturation / anti-windup, so the feedforward term participates in the output_min / output_max clip and the anti-windup comparison between unsaturated and saturated control values. For a plant whose steady-state transfer function from u to y is G(0), choosing kff = 1/G(0) makes the controller track step changes in r without integrator action — fastest possible step response (pair with PID for disturbance rejection). Differentiable through jax.grad. Default 0.0 → byte-equivalent to phase 1 (no feedforward).	0.0
kp_dynamic		T-127-followup-gain-scheduling. If True, the proportional gain kp is read from a runtime input port instead of the static kp parameter. The new port is appended after r, y, and any active b / c ports (see "Input ports" above). The static kp value is then ignored at runtime; the user MUST connect a signal to the new port. Default False (byte-equivalent to phase 1).	False
ki_dynamic		Same as kp_dynamic but for the integral gain ki. Default False.	False
kd_dynamic		Same as kp_dynamic but for the derivative gain kd. Default False.	False
kff_dynamic		Same as kp_dynamic but for the feedforward gain kff. Default False.	False
error_deadband		T-127-followup-deadband-error. Non-negative scalar; when positive, the raw error signal e_raw (the unweighted r - y for the integral path and the weighted b*r - y / c*r - y for the P / D paths) is gated through a deadband before it feeds each PID term. In hard mode the gate is e = e_raw for |e_raw| > error_deadband and e = 0 otherwise; in smooth mode the gate is :func:soft_dead_zone(e_raw, error_deadband, error_deadband_sharpness). Default 0.0 disables the deadband entirely (byte-equivalent to phase 1). Differentiable through error_deadband in smooth mode; hard mode is kinked at the boundary.	0.0
error_deadband_mode		"hard" (default) or "smooth". Hard mode uses an npa.where gate; smooth mode uses :func:soft_dead_zone for a sigmoid-blended kernel with finite gradient through the band. Only relevant when error_deadband > 0.	'hard'
error_deadband_sharpness		Positive scalar controlling the steepness of the smooth deadband transition (passed straight to :func:soft_dead_zone). Default 10.0. Ignored when error_deadband_mode='hard'.	10.0
tracking_enabled		T-127-followup-tracking-mode. If True, declares an extra input port u_ext (appended after every other dynamic port) and folds a tracking-error term into the integrator update:: e_track = u_ext - u_unsat I[k+1] += (e_track / tracking_gain) * dt This is the standard "tracking mode" / "manual mode" mechanism: while another controller (or an operator) drives u_ext, the PID's integrator is pulled toward the value that would produce u_ext so the handoff back to PID- driven control is bumpless. Implementation-wise it is back- calculation on the external signal — the same kernel as anti_windup_method="back_calc" but using u_ext instead of u_sat as the target. Both mechanisms compose: their corrections sum into the integrator each tick. Default False (byte-equivalent to phase 1).	False
tracking_gain		Tracking time constant Tt for the tracking-mode back- calculation kernel (T-127-followup-tracking-mode). Smaller values pull the integrator toward u_ext faster. Differentiable through jax.grad. Default 1.0.	1.0
integrate_tracking_error		T-127-followup-i-on-error-only. Selects whether the tracking-error term (u_ext - u_unsat)/Tt * dt is folded into the integrator each tick. When True (default) the integrator update is:: I[k+1] = I[k] + Ki*(r-y)*dt + (u_ext - u_unsat)/Tt * dt preserving the T-127-followup-tracking-mode kernel exactly. When False the integrator only accumulates the regulation error:: I[k+1] = I[k] + Ki*(r-y)*dt and u_ext does NOT pull the integrator at all (the tracking signal still flows through any parallel path the user has wired up, e.g. a feedforward addition outside the block). Only meaningful when tracking_enabled=True; the flag is silently irrelevant otherwise but still round-trips through :meth:to_dict / :meth:from_dict. Default True is byte-equivalent to T-127-followup-tracking-mode.	True
tracking_enabled_dynamic		T-127-followup-bumpless-mode-switch. When True, promotes the tracking-mode flag from a static construction- time choice to a runtime SCALAR INPUT port appended after u_ext. The port value is treated as a boolean (0 = OFF / AUTO, non-zero = ON / MANUAL/TRACKING) — multiplying the per-tick tracking-pull correction by that gate. This lets a model toggle between PID-driven control and external override mid-simulation while keeping the integrator loaded with the value that would produce u_ext (so handoffs in either direction remain bumpless). Requires tracking_enabled=True; tracking_enabled=False plus tracking_enabled_dynamic=True raises ValueError at construction (without the u_ext port the runtime gate has nothing to multiply). Composes with integrate_tracking_error (the gate multiplies the correction term that flag exposes). Default False is byte-equivalent to T-127-followup-tracking-mode.	False

Gain-scheduling recipe

Each of kp_dynamic, ki_dynamic, kd_dynamic, and kff_dynamic is the natural plug for a lookup-table-driven gain. The standard wiring uses one :class:LookupTable1d (or :class:LookupTable2d for two scheduling variables) per scheduled gain and the same scheduling-variable signal source for all of them::

import jaxonomy
from jaxonomy.library import (
    Constant, LookupTable1d, PIDController2DOF,
)

builder = jaxonomy.DiagramBuilder()
r = builder.add(Constant(1.0, name="r"))
y = builder.add(Constant(0.0, name="y"))
# Scheduling variable -- e.g. engine speed, Mach number,
# tank level.  Replace with whatever source you have.
sched = builder.add(Constant(0.5, name="sched"))
# Schedule kp as a function of the scheduling variable.
kp_tbl = builder.add(
    LookupTable1d(
        input_array=[0.0, 0.5, 1.0],
        output_array=[1.0, 2.0, 4.0],
        interpolation="linear",
        name="kp_schedule",
    )
)
pid = builder.add(
    PIDController2DOF(
        dt=0.01, kp_dynamic=True, name="pid"
    )
)
builder.connect(r.output_ports[0], pid.input_ports[0])
builder.connect(y.output_ports[0], pid.input_ports[1])
builder.connect(sched.output_ports[0], kp_tbl.input_ports[0])
# kp port is at index 2 when only kp_dynamic is set.
builder.connect(kp_tbl.output_ports[0], pid.input_ports[2])

Because every dynamic port is a regular signal port, gradients flow through the lookup-table parameters (breakpoints / values) AND through the scheduling-variable signal — the standard T-114 guarantee. Multiple gains can be scheduled simultaneously; flagging ki_dynamic and kd_dynamic simply adds two more ports for the integral / derivative tables.

Source code in jaxonomy/library/dynamics.py

class PIDController2DOF(LeafSystem):
    """Two-degree-of-freedom discrete-time PID controller.

    Implements the standard 2-DOF PID control law::

        u = Kp * (b*r - y) + Ki * integral(r - y) + Kd * d/dt(c*r - y)

    where ``r`` is the setpoint, ``y`` is the measurement, and ``b`` and
    ``c`` are setpoint weights in ``[0, 1]`` for the proportional and
    derivative paths respectively.  With ``b = c = 1`` the block is
    numerically equivalent to the existing :class:`PIDDiscrete` block on
    the error signal ``e = r - y``; with ``b = c = 0`` it becomes an
    "I-PD" controller (only the integral term reacts to setpoint
    changes).

    The integral term uses a forward-Euler approximation::

        e_int[k+1] = e_int[k] + (r[k] - y[k]) * dt

    and the derivative term is computed exactly as for
    :class:`DerivativeDiscrete` / :class:`PIDDiscrete`, including the
    optional first-order filter (``filter_type``, ``filter_coefficient``).

    Input ports:
        (0) Setpoint signal ``r``.
        (1) Measurement signal ``y``.
        Dynamic-port appendices, declared in this deterministic order
        when the corresponding ``*_dynamic`` flag is True:

        (a) ``b`` if ``b_dynamic=True`` (T-127-followup-external-weights)
        (b) ``c`` if ``c_dynamic=True``
        (c) ``kp`` if ``kp_dynamic=True``
            (T-127-followup-gain-scheduling)
        (d) ``ki`` if ``ki_dynamic=True``
        (e) ``kd`` if ``kd_dynamic=True``
        (f) ``kff`` if ``kff_dynamic=True``
        (g) ``u_ext`` if ``tracking_enabled=True``
            (T-127-followup-tracking-mode)
        (h) ``mode_flag`` if ``tracking_enabled_dynamic=True``
            (T-127-followup-bumpless-mode-switch).  Scalar input cast to
            a {0, 1} gate that selects whether the tracking-pull branch
            runs this tick (0 = OFF / AUTO, non-zero = ON / MANUAL /
            TRACKING).  Requires ``tracking_enabled=True`` so the
            ``u_ext`` port exists.

        Port indices skip any flag set to False, so e.g. with only
        ``kp_dynamic=True`` the ``kp`` port is at index 2; with
        ``b_dynamic=True`` + ``kp_dynamic=True`` the ``b`` port is at
        index 2 and ``kp`` is at index 3.  The instance attributes
        ``self.b_index`` / ``self.c_index`` / ``self.kp_index`` /
        ``self.ki_index`` / ``self.kd_index`` / ``self.kff_index`` /
        ``self.u_ext_index`` / ``self.mode_flag_index`` expose the
        resolved positions.

    Output ports:
        (0) The control signal ``u`` computed by the 2-DOF PID law.

    Parameters:
        kp:
            Proportional gain (scalar).
        ki:
            Integral gain (scalar).
        kd:
            Derivative gain (scalar).
        b:
            Setpoint weight for the proportional term, in ``[0, 1]``.
            Default 1.0 (matches 1-DOF PID).  Ignored when
            ``b_dynamic=True`` (the runtime port value is used instead).
        c:
            Setpoint weight for the derivative term, in ``[0, 1]``.
            Default 1.0 (matches 1-DOF PID).  Ignored when
            ``c_dynamic=True``.

            **Recommendation for real-world controllers:** prefer
            ``c = 0`` (a.k.a. "derivative on measurement only").  With
            ``c = 1`` (textbook 2-DOF / 1-DOF PID), a step change in the
            setpoint produces a one-tick spike of magnitude ``Kd / dt``
            in ``d/dt(c*r - y)`` — "derivative kick" — that propagates
            into a brief, often saturation-inducing control transient.
            Setting ``c = 0`` routes the derivative through the
            measurement only (``-d/dt(y)``), so setpoint steps no
            longer kick the derivative term and integral / proportional
            action alone drive the response.  See the
            :meth:`with_derivative_on_measurement` factory and the
            ``derivative_on_measurement_only=True`` convenience kwarg
            below for the standard recipe.
        derivative_on_measurement_only:
            T-127-followup-derivative-on-measurement.  Convenience flag
            equivalent to ``c=0`` + ``c_dynamic=False``.  When True the
            derivative term sees only the measurement (``-d/dt(y)``)
            and is immune to setpoint-step kick; the user-supplied
            ``c`` / ``c_dynamic`` kwargs are rejected with a
            ``ValueError`` to keep the contract unambiguous.  Default
            False (``c=1``) preserves byte-equivalence with phase 1.
            See :meth:`with_derivative_on_measurement` for the matching
            factory function.
        b_dynamic:
            If True, the proportional setpoint weight ``b`` is read from
            an additional input port (index 2) rather than from the
            static ``b`` parameter.  The static ``b`` value is then
            ignored at runtime; the user MUST connect a signal to the
            new port.  Mirrors the ``enable_dynamic_*`` pattern used by
            :class:`Saturate` and :class:`RateLimiter`.  Default False
            (byte-equivalent to phase 1).
        c_dynamic:
            If True, the derivative setpoint weight ``c`` is read from
            an additional input port instead of the static ``c``
            parameter.  The new port lives at index 2 (when only
            ``c_dynamic`` is set) or index 3 (when both ``b_dynamic``
            and ``c_dynamic`` are set).  Default False.
        dt:
            Sampling period of the block.
        initial_state:
            Initial value of the integral.  Default 0.0.
        filter_type:
            One of ``"none"``, ``"forward"``, ``"backward"``, or
            ``"bilinear"`` — derivative-filter mode.  Default ``"none"``.
        filter_coefficient:
            Filter coefficient ``N`` for the derivative filter (the
            conventional "filter coefficient" PID-tuning parameter).  Default 1.0.
        output_min:
            Lower saturation limit on the control output (T-127-followup-
            anti-windup).  ``None`` (default) disables the lower clip.
            When either ``output_min`` or ``output_max`` is set the
            *saturated* control value is published on port (0); the
            unsaturated value also feeds the anti-windup correction.
        output_max:
            Upper saturation limit on the control output.  ``None``
            (default) disables the upper clip.
        anti_windup_method:
            One of ``"none"``, ``"back_calc"``, or ``"clamping"``
            (T-127-followup-anti-windup).  Default ``"none"``.

            * ``"none"`` — no anti-windup; integrator update unchanged.
            * ``"back_calc"`` — back-calculation: subtract
              ``(u_unsat - u_sat) / anti_windup_gain * dt`` from the
              integrator each tick.  Smooth, fully differentiable.
            * ``"clamping"`` — integrator-tracking: only update the
              integral when the controller is *not* pushing further into
              saturation (``u_unsat == u_sat`` OR the error sign points
              away from the saturated direction).

            Anti-windup is a no-op unless ``output_min`` or ``output_max``
            is also set.
        anti_windup_gain:
            Tracking time constant ``Tt`` used by ``"back_calc"``.  Smaller
            values pull the integrator back faster.  Default 1.0.
            Differentiable through ``jax.grad``.
        integrator_method:
            One of ``"forward_euler"`` (default), ``"backward_euler"``, or
            ``"trapezoidal"`` — selects the discretisation used to advance
            the integral term (T-127-followup-discrete-integrator-
            derivative).  Backward-Euler is more stable for stiff loops;
            trapezoidal is more accurate.  Defaults to byte-equivalence
            with phase 1.  Non-default values add an extra ``e_i_prev``
            delay cell to the discrete state.
        derivative_method:
            One of ``"forward_diff"`` (default), ``"backward_diff"``, or
            ``"centered_diff"`` — selects the unfiltered derivative
            kernel.  ``"backward_diff"`` uses past samples only (typical
            for real-time control, introduces a one-tick delay relative
            to ``"forward_diff"``).  ``"centered_diff"`` is less noisy
            but adds one extra delay cell (``e_d_prev_prev``).  Only
            applies when ``filter_type='none'`` — combining a non-default
            ``derivative_method`` with a recursive filter raises a
            ``ValueError`` at construction.
        kff:
            Feedforward gain on the setpoint (T-127-followup-feedforward).
            Adds ``kff * r`` to the PID output *before* saturation /
            anti-windup, so the feedforward term participates in the
            ``output_min`` / ``output_max`` clip and the anti-windup
            comparison between unsaturated and saturated control values.
            For a plant whose steady-state transfer function from ``u``
            to ``y`` is ``G(0)``, choosing ``kff = 1/G(0)`` makes the
            controller track step changes in ``r`` without integrator
            action — fastest possible step response (pair with PID for
            disturbance rejection).  Differentiable through ``jax.grad``.
            Default ``0.0`` → byte-equivalent to phase 1 (no
            feedforward).
        kp_dynamic:
            T-127-followup-gain-scheduling.  If True, the proportional
            gain ``kp`` is read from a runtime input port instead of the
            static ``kp`` parameter.  The new port is appended after
            ``r``, ``y``, and any active ``b`` / ``c`` ports (see "Input
            ports" above).  The static ``kp`` value is then ignored at
            runtime; the user MUST connect a signal to the new port.
            Default False (byte-equivalent to phase 1).
        ki_dynamic:
            Same as ``kp_dynamic`` but for the integral gain ``ki``.
            Default False.
        kd_dynamic:
            Same as ``kp_dynamic`` but for the derivative gain ``kd``.
            Default False.
        kff_dynamic:
            Same as ``kp_dynamic`` but for the feedforward gain ``kff``.
            Default False.
        error_deadband:
            T-127-followup-deadband-error.  Non-negative scalar; when
            positive, the raw error signal ``e_raw`` (the unweighted
            ``r - y`` for the integral path and the weighted ``b*r - y``
            / ``c*r - y`` for the P / D paths) is gated through a
            deadband before it feeds each PID term.  In hard mode the
            gate is ``e = e_raw`` for ``|e_raw| > error_deadband`` and
            ``e = 0`` otherwise; in smooth mode the gate is
            :func:`soft_dead_zone(e_raw, error_deadband,
            error_deadband_sharpness)`.  Default ``0.0`` disables the
            deadband entirely (byte-equivalent to phase 1).
            Differentiable through ``error_deadband`` in smooth mode;
            hard mode is kinked at the boundary.
        error_deadband_mode:
            ``"hard"`` (default) or ``"smooth"``.  Hard mode uses an
            ``npa.where`` gate; smooth mode uses :func:`soft_dead_zone`
            for a sigmoid-blended kernel with finite gradient through
            the band.  Only relevant when ``error_deadband > 0``.
        error_deadband_sharpness:
            Positive scalar controlling the steepness of the smooth
            deadband transition (passed straight to
            :func:`soft_dead_zone`).  Default ``10.0``.  Ignored when
            ``error_deadband_mode='hard'``.
        tracking_enabled:
            T-127-followup-tracking-mode.  If True, declares an extra
            input port ``u_ext`` (appended after every other dynamic
            port) and folds a *tracking-error* term into the integrator
            update::

                e_track = u_ext - u_unsat
                I[k+1] += (e_track / tracking_gain) * dt

            This is the standard "tracking mode" / "manual mode"
            mechanism: while another controller (or an operator) drives
            ``u_ext``, the PID's integrator is pulled toward the value
            that would produce ``u_ext`` so the handoff back to PID-
            driven control is bumpless.  Implementation-wise it is back-
            calculation on the *external* signal — the same kernel as
            ``anti_windup_method="back_calc"`` but using ``u_ext``
            instead of ``u_sat`` as the target.  Both mechanisms compose:
            their corrections sum into the integrator each tick.
            Default False (byte-equivalent to phase 1).
        tracking_gain:
            Tracking time constant ``Tt`` for the tracking-mode back-
            calculation kernel (T-127-followup-tracking-mode).  Smaller
            values pull the integrator toward ``u_ext`` faster.
            Differentiable through ``jax.grad``.  Default 1.0.
        integrate_tracking_error:
            T-127-followup-i-on-error-only.  Selects whether the
            tracking-error term ``(u_ext - u_unsat)/Tt * dt`` is folded
            into the integrator each tick.  When ``True`` (default) the
            integrator update is::

                I[k+1] = I[k] + Ki*(r-y)*dt + (u_ext - u_unsat)/Tt * dt

            preserving the T-127-followup-tracking-mode kernel exactly.
            When ``False`` the integrator only accumulates the
            regulation error::

                I[k+1] = I[k] + Ki*(r-y)*dt

            and ``u_ext`` does NOT pull the integrator at all (the
            tracking signal still flows through any parallel path the
            user has wired up, e.g. a feedforward addition outside the
            block).  Only meaningful when ``tracking_enabled=True``; the
            flag is silently irrelevant otherwise but still round-trips
            through :meth:`to_dict` / :meth:`from_dict`.  Default
            ``True`` is byte-equivalent to T-127-followup-tracking-mode.
        tracking_enabled_dynamic:
            T-127-followup-bumpless-mode-switch.  When ``True``,
            promotes the tracking-mode flag from a static construction-
            time choice to a runtime SCALAR INPUT port appended after
            ``u_ext``.  The port value is treated as a boolean (``0`` =
            OFF / AUTO, non-zero = ON / MANUAL/TRACKING) — multiplying
            the per-tick tracking-pull correction by that gate.  This
            lets a model toggle between PID-driven control and external
            override mid-simulation while keeping the integrator loaded
            with the value that would produce ``u_ext`` (so handoffs in
            either direction remain bumpless).  Requires
            ``tracking_enabled=True``; ``tracking_enabled=False`` plus
            ``tracking_enabled_dynamic=True`` raises ``ValueError`` at
            construction (without the ``u_ext`` port the runtime gate
            has nothing to multiply).  Composes with
            ``integrate_tracking_error`` (the gate multiplies the
            correction term that flag exposes).  Default ``False`` is
            byte-equivalent to T-127-followup-tracking-mode.

    Gain-scheduling recipe:
        Each of ``kp_dynamic``, ``ki_dynamic``, ``kd_dynamic``, and
        ``kff_dynamic`` is the natural plug for a lookup-table-driven
        gain.  The standard wiring uses one :class:`LookupTable1d` (or
        :class:`LookupTable2d` for two scheduling variables) per
        scheduled gain and the same scheduling-variable signal source
        for all of them::

            import jaxonomy
            from jaxonomy.library import (
                Constant, LookupTable1d, PIDController2DOF,
            )

            builder = jaxonomy.DiagramBuilder()
            r = builder.add(Constant(1.0, name="r"))
            y = builder.add(Constant(0.0, name="y"))
            # Scheduling variable -- e.g. engine speed, Mach number,
            # tank level.  Replace with whatever source you have.
            sched = builder.add(Constant(0.5, name="sched"))
            # Schedule kp as a function of the scheduling variable.
            kp_tbl = builder.add(
                LookupTable1d(
                    input_array=[0.0, 0.5, 1.0],
                    output_array=[1.0, 2.0, 4.0],
                    interpolation="linear",
                    name="kp_schedule",
                )
            )
            pid = builder.add(
                PIDController2DOF(
                    dt=0.01, kp_dynamic=True, name="pid"
                )
            )
            builder.connect(r.output_ports[0], pid.input_ports[0])
            builder.connect(y.output_ports[0], pid.input_ports[1])
            builder.connect(sched.output_ports[0], kp_tbl.input_ports[0])
            # kp port is at index 2 when only kp_dynamic is set.
            builder.connect(kp_tbl.output_ports[0], pid.input_ports[2])

        Because every dynamic port is a regular signal port, gradients
        flow through the lookup-table parameters (breakpoints / values)
        AND through the scheduling-variable signal — the standard T-114
        guarantee.  Multiple gains can be scheduled simultaneously;
        flagging ``ki_dynamic`` and ``kd_dynamic`` simply adds two more
        ports for the integral / derivative tables.
    """

    class DiscreteStateType(NamedTuple):
        integral: Array
        # Recursive filter memory for the derivative estimate.  We keep the
        # *weighted* derivative error ``e_d = c*r - y`` (and the previous
        # filtered derivative) so the same recursive-filter formulas as
        # ``PIDDiscrete`` apply.
        e_d_prev: Array
        e_dot_prev: Array

    # T-127-followup-anti-windup — supported anti-windup methods.
    _ANTI_WINDUP_METHODS = ("none", "back_calc", "clamping")

    # T-127-followup-discrete-integrator-derivative — pluggable kernels.
    _INTEGRATOR_METHODS = ("forward_euler", "backward_euler", "trapezoidal")
    _DERIVATIVE_METHODS = ("forward_diff", "backward_diff", "centered_diff")

    # T-127-followup-deadband-error — supported deadband gate modes.
    _ERROR_DEADBAND_MODES = ("hard", "smooth")

    @parameters(
        static=[
            "dt",
            "filter_type",
            "filter_coefficient",
            "anti_windup_method",
            "b_dynamic",
            "c_dynamic",
            "integrator_method",
            "derivative_method",
            "kp_dynamic",
            "ki_dynamic",
            "kd_dynamic",
            "kff_dynamic",
            "error_deadband_mode",
            "tracking_enabled",
            "integrate_tracking_error",
            "tracking_enabled_dynamic",
        ],
        dynamic=[
            "kp",
            "ki",
            "kd",
            "b",
            "c",
            "initial_state",
            "output_min",
            "output_max",
            "anti_windup_gain",
            "kff",
            "error_deadband",
            "error_deadband_sharpness",
            "tracking_gain",
        ],
    )
    def __init__(
        self,
        dt,
        kp=1.0,
        ki=1.0,
        kd=1.0,
        b=1.0,
        c=1.0,
        initial_state=0.0,
        filter_type="none",
        filter_coefficient=1.0,
        output_min=None,
        output_max=None,
        anti_windup_method="none",
        anti_windup_gain=1.0,
        b_dynamic=False,
        c_dynamic=False,
        integrator_method="forward_euler",
        derivative_method="forward_diff",
        kff=0.0,
        kp_dynamic=False,
        ki_dynamic=False,
        kd_dynamic=False,
        kff_dynamic=False,
        error_deadband=0.0,
        error_deadband_mode="hard",
        error_deadband_sharpness=10.0,
        tracking_enabled=False,
        tracking_gain=1.0,
        integrate_tracking_error=True,
        tracking_enabled_dynamic=False,
        dtype=None,
        **kwargs,
    ):
        # T-127-followup-derivative-on-measurement note: the
        # ``derivative_on_measurement_only`` convenience kwarg is
        # intercepted by an outer wrapper installed after the class
        # body (see ``_pid2dof_derivative_on_measurement_wrapper``)
        # so that the rewrite to ``c=0`` / ``c_dynamic=False`` lands
        # in the original ``kwargs`` BEFORE the ``@parameters``
        # decorator captures them.  The wrapper validates conflicting
        # ``c`` / ``c_dynamic`` overrides; by the time this body runs
        # the values are already canonical.
        # Per-block dtype override — same plumbing as PIDDiscrete (T-038a-
        # followup-other-blocks / T-038a-followup-mixed-precision-cascade).
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype

        # T-127-followup-anti-windup — validate and cache static config.
        if anti_windup_method not in self._ANTI_WINDUP_METHODS:
            raise ValueError(
                f"anti_windup_method must be one of "
                f"{self._ANTI_WINDUP_METHODS!r}; got {anti_windup_method!r}"
            )
        self._anti_windup_method = anti_windup_method
        # Anti-windup is "active" only when at least one saturation limit
        # is set.  When both are None the block is byte-equivalent to
        # phase 1 (no clip on the output, no integrator correction),
        # regardless of method string — matching the documented default-
        # off contract.
        self._anti_windup_active = (
            output_min is not None or output_max is not None
        )

        # T-127-followup-discrete-integrator-derivative — validate and
        # cache the integrator/derivative kernel selection.  Defaults
        # ("forward_euler"/"forward_diff") leave the phase 1 update path
        # bit-identical (no extra state cells, same numerical formulas).
        if integrator_method not in self._INTEGRATOR_METHODS:
            raise ValueError(
                f"integrator_method must be one of "
                f"{self._INTEGRATOR_METHODS!r}; got {integrator_method!r}"
            )
        if derivative_method not in self._DERIVATIVE_METHODS:
            raise ValueError(
                f"derivative_method must be one of "
                f"{self._DERIVATIVE_METHODS!r}; got {derivative_method!r}"
            )
        # The derivative-method kwarg only governs the unfiltered finite-
        # difference path; when ``filter_type != "none"`` the recursive
        # filter coefficients (forward/backward/bilinear Euler) own the
        # discretisation.  Reject ambiguous combinations early.
        if derivative_method != "forward_diff" and filter_type != "none":
            raise ValueError(
                "derivative_method only applies when filter_type='none'; "
                f"got derivative_method={derivative_method!r} with "
                f"filter_type={filter_type!r}"
            )
        self._integrator_method = integrator_method
        self._derivative_method = derivative_method

        # T-127-followup-deadband-error — validate the deadband mode at
        # construction so we can dispatch cheaply inside ``_apply_deadband``
        # without re-parsing the string each tick.  Mirrors the
        # ``DeadZone`` block's mode-flag pattern (see T-115-followup-
        # deadzone-backlash).  The half-width and sharpness flow through
        # dynamic parameters so ``jax.grad`` w.r.t. them is finite in
        # smooth mode.
        if error_deadband_mode not in self._ERROR_DEADBAND_MODES:
            raise ValueError(
                f"error_deadband_mode must be one of "
                f"{self._ERROR_DEADBAND_MODES!r}; got {error_deadband_mode!r}"
            )
        self._error_deadband_mode = error_deadband_mode
        # Track whether the deadband is "active" purely for the
        # byte-equivalence fast path: when ``error_deadband == 0.0`` we
        # bypass the gate entirely so phase 1 / earlier-followup tests
        # remain bit-identical.  The flag mirrors ``_anti_windup_active``.
        try:
            _eb = float(error_deadband)
        except (TypeError, ValueError):
            # Tracer / array-like deadband — assume the gate runs.
            _eb = 1.0
        self._error_deadband_active = _eb != 0.0
        # Build a per-instance state tuple: extend the base 3-field layout
        # with optional delay cells when the chosen kernels need them.
        # Defaults keep the state shape identical to phase 1.
        state_fields = ["integral", "e_d_prev", "e_dot_prev"]
        if integrator_method != "forward_euler":
            # backward_euler / trapezoidal both consume the previous
            # integral-error sample, so we have to remember it.
            state_fields.append("e_i_prev")
        if derivative_method == "centered_diff":
            # centered_diff = (e[k+1] - e[k-1]) / (2*dt) needs one extra
            # delay beyond the phase 1 ``e_d_prev``.
            state_fields.append("e_d_prev_prev")
        self._state_fields = tuple(state_fields)
        # Per-instance NamedTuple (overrides the class-level default for
        # this instance).  ``self.DiscreteStateType`` is what
        # initialize() / _update() / reset_default_values() construct.
        from collections import namedtuple as _namedtuple

        self.DiscreteStateType = _namedtuple(
            "PIDController2DOFState", self._state_fields
        )

        super().__init__(**kwargs)
        self.dt = dt
        self.setpoint_index = self.declare_input_port()  # r
        self.measurement_index = self.declare_input_port()  # y

        # T-127-followup-external-weights — optional runtime input ports
        # for the setpoint weights ``b`` (proportional) and ``c``
        # (derivative).  Order matters: when both flags are True the
        # ports are appended in the (b, c) order so the indices are
        # deterministic and can be documented up front.
        self.b_dynamic = bool(b_dynamic)
        self.c_dynamic = bool(c_dynamic)
        if self.b_dynamic:
            self.b_index = self.declare_input_port()
        if self.c_dynamic:
            self.c_index = self.declare_input_port()

        # T-127-followup-gain-scheduling — optional runtime input ports
        # for the four scalar gains (Kp, Ki, Kd, Kff).  Appended after
        # the b/c ports so the indexing remains backward-compatible with
        # T-127-followup-external-weights (existing models that only set
        # b_dynamic / c_dynamic don't see their port indices shift).  The
        # documented order (kp, ki, kd, kff) is the same order in which
        # users typically schedule them.
        self.kp_dynamic = bool(kp_dynamic)
        self.ki_dynamic = bool(ki_dynamic)
        self.kd_dynamic = bool(kd_dynamic)
        self.kff_dynamic = bool(kff_dynamic)
        if self.kp_dynamic:
            self.kp_index = self.declare_input_port()
        if self.ki_dynamic:
            self.ki_index = self.declare_input_port()
        if self.kd_dynamic:
            self.kd_index = self.declare_input_port()
        if self.kff_dynamic:
            self.kff_index = self.declare_input_port()

        # T-127-followup-tracking-mode — optional ``u_ext`` input port for
        # bumpless-transfer / manual-override.  Appended last so older
        # consumers keep their port indices; the user MUST wire the
        # port when ``tracking_enabled=True``.
        self.tracking_enabled = bool(tracking_enabled)
        if self.tracking_enabled:
            self.u_ext_index = self.declare_input_port()

        # T-127-followup-bumpless-mode-switch — optional runtime
        # ``mode_flag`` input port that gates the tracking-pull branch on
        # / off each tick.  Requires ``tracking_enabled=True`` so the
        # ``u_ext`` port exists; otherwise the runtime gate has nothing
        # to multiply.  Appended LAST (after every other dynamic port
        # including ``u_ext``) so older consumers keep their indices.
        self.tracking_enabled_dynamic = bool(tracking_enabled_dynamic)
        if self.tracking_enabled_dynamic and not self.tracking_enabled:
            raise ValueError(
                "PIDController2DOF: tracking_enabled_dynamic=True requires "
                "tracking_enabled=True (the u_ext port must exist for the "
                "runtime mode flag to gate)."
            )
        if self.tracking_enabled_dynamic:
            self.mode_flag_index = self.declare_input_port()

        # T-127-followup-i-on-error-only — controls whether the
        # tracking-error term ``(u_ext - u_unsat)/Tt * dt`` is folded
        # into the integrator update.  When False, ``u_ext`` reaches the
        # integrator ONLY via the regulation error ``r - y`` (i.e. the
        # tracking signal pulls through the parallel back-calculation
        # path is suppressed).  Default True preserves the
        # T-127-followup-tracking-mode behavior, hence byte-equivalence
        # with that followup.  Stored unconditionally so the flag round-
        # trips even when ``tracking_enabled=False`` (in which case it
        # is moot but still serializable).
        self.integrate_tracking_error = bool(integrate_tracking_error)

        # Declare the periodic update.
        self._periodic_update_idx = self.declare_periodic_update()

        # Declare an output port for the control signal.
        self.control_output = self.declare_output_port()

    # T-127-followup-discrete-integrator-derivative -----------------------
    def _make_state(self, *, integral, e_d_prev, e_dot_prev,
                    e_i_prev=None, e_d_prev_prev=None):
        """Build a ``DiscreteStateType`` tuple, supplying optional fields
        only when the configured kernels need them.

        Defaults (forward_euler / forward_diff) skip the optional
        fields, so the returned tuple is shape-identical to phase 1.
        """
        kw = dict(
            integral=integral,
            e_d_prev=e_d_prev,
            e_dot_prev=e_dot_prev,
        )
        if "e_i_prev" in self._state_fields:
            kw["e_i_prev"] = (
                e_i_prev if e_i_prev is not None else npa.zeros_like(integral)
            )
        if "e_d_prev_prev" in self._state_fields:
            kw["e_d_prev_prev"] = (
                e_d_prev_prev if e_d_prev_prev is not None
                else npa.zeros_like(integral)
            )
        return self.DiscreteStateType(**kw)

    def initialize(
        self,
        kp,
        ki,
        kd,
        b,
        c,
        initial_state,
        filter_type,
        filter_coefficient,
        dt=None,
        output_min=None,
        output_max=None,
        anti_windup_method="none",
        anti_windup_gain=1.0,
        b_dynamic=False,
        c_dynamic=False,
        integrator_method="forward_euler",
        derivative_method="forward_diff",
        kff=0.0,
        kp_dynamic=False,
        ki_dynamic=False,
        kd_dynamic=False,
        kff_dynamic=False,
        error_deadband=0.0,
        error_deadband_mode="hard",
        error_deadband_sharpness=10.0,
        tracking_enabled=False,
        tracking_gain=1.0,
        integrate_tracking_error=True,
        tracking_enabled_dynamic=False,
    ):
        # T-127-followup-external-weights — port topology is decided at
        # construction time (mirrors the RateLimiter/SoftRateLimiter
        # contract); reject mid-life flips that would silently shift
        # port indices.
        if bool(b_dynamic) != self.b_dynamic:
            raise ValueError(
                "PIDController2DOF: b_dynamic cannot be changed after "
                "initialization"
            )
        if bool(c_dynamic) != self.c_dynamic:
            raise ValueError(
                "PIDController2DOF: c_dynamic cannot be changed after "
                "initialization"
            )
        # T-127-followup-gain-scheduling — same port-topology lock for
        # the runtime-gain flags.
        if bool(kp_dynamic) != self.kp_dynamic:
            raise ValueError(
                "PIDController2DOF: kp_dynamic cannot be changed after "
                "initialization"
            )
        if bool(ki_dynamic) != self.ki_dynamic:
            raise ValueError(
                "PIDController2DOF: ki_dynamic cannot be changed after "
                "initialization"
            )
        if bool(kd_dynamic) != self.kd_dynamic:
            raise ValueError(
                "PIDController2DOF: kd_dynamic cannot be changed after "
                "initialization"
            )
        if bool(kff_dynamic) != self.kff_dynamic:
            raise ValueError(
                "PIDController2DOF: kff_dynamic cannot be changed after "
                "initialization"
            )
        # T-127-followup-discrete-integrator-derivative — kernel choice
        # is part of the static port topology; reject mid-life flips that
        # would silently change the state-tuple shape.
        if integrator_method != self._integrator_method:
            raise ValueError(
                "PIDController2DOF: integrator_method cannot be changed "
                "after initialization"
            )
        if derivative_method != self._derivative_method:
            raise ValueError(
                "PIDController2DOF: derivative_method cannot be changed "
                "after initialization"
            )
        # T-127-followup-deadband-error — mode is static so the
        # dispatch inside ``_apply_deadband`` can specialise without
        # re-parsing each tick.
        if error_deadband_mode != self._error_deadband_mode:
            raise ValueError(
                "PIDController2DOF: error_deadband_mode cannot be "
                "changed after initialization"
            )
        # T-127-followup-tracking-mode — port-topology lock for the
        # tracking-mode flag (same contract as the other ``*_dynamic`` /
        # ``*_enabled`` flags above).
        if bool(tracking_enabled) != self.tracking_enabled:
            raise ValueError(
                "PIDController2DOF: tracking_enabled cannot be changed "
                "after initialization"
            )
        # T-127-followup-i-on-error-only — gate the tracking-integrator
        # contribution at construction time so the update path can
        # specialise without re-parsing each tick.
        if bool(integrate_tracking_error) != self.integrate_tracking_error:
            raise ValueError(
                "PIDController2DOF: integrate_tracking_error cannot be "
                "changed after initialization"
            )
        # T-127-followup-bumpless-mode-switch — port-topology lock for
        # the runtime ``mode_flag`` port; same contract as the other
        # ``*_dynamic`` / ``*_enabled`` flags above.
        if bool(tracking_enabled_dynamic) != self.tracking_enabled_dynamic:
            raise ValueError(
                "PIDController2DOF: tracking_enabled_dynamic cannot be "
                "changed after initialization"
            )
        # Cast initial state and zero seeds to the per-block dtype if set.
        _zero = 0.0
        if self._dtype is not None:
            initial_state = npa.asarray(initial_state).astype(self._dtype)
            _zero = npa.asarray(0.0).astype(self._dtype)

        self.declare_discrete_state(
            default_value=self._make_state(
                integral=initial_state,
                e_d_prev=_zero,
                e_dot_prev=_zero,
                e_i_prev=_zero,
                e_d_prev_prev=_zero,
            ),
            as_array=False,
        )

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=0.0,
        )

        # Derivative-filter coefficients (b0,b1) / (a0,a1) — same helper as
        # PIDDiscrete / DerivativeDiscrete.
        self.filter_type = filter_type
        b_coef, a_coef = derivative_filter(
            N=filter_coefficient, dt=self.dt, filter_type=filter_type
        )
        if self._dtype is not None:
            b_coef = npa.asarray(b_coef).astype(self._dtype)
            a_coef = npa.asarray(a_coef).astype(self._dtype)
        self.filter = (b_coef, a_coef)

        # T-127-followup-external-weights / T-127-followup-gain-
        # scheduling — include the optional dynamic-input ports in the
        # output-prerequisite set so the scheduler eagerly evaluates
        # them before ``_output``.  When a flag is False the matching
        # port simply does not exist.
        prereqs = [
            DependencyTicket.xd,
            self.input_ports[0].ticket,
            self.input_ports[1].ticket,
        ]
        if self.b_dynamic:
            prereqs.append(self.input_ports[self.b_index].ticket)
        if self.c_dynamic:
            prereqs.append(self.input_ports[self.c_index].ticket)
        if self.kp_dynamic:
            prereqs.append(self.input_ports[self.kp_index].ticket)
        if self.ki_dynamic:
            prereqs.append(self.input_ports[self.ki_index].ticket)
        if self.kd_dynamic:
            prereqs.append(self.input_ports[self.kd_index].ticket)
        if self.kff_dynamic:
            prereqs.append(self.input_ports[self.kff_index].ticket)
        # T-127-followup-tracking-mode — ``u_ext`` only feeds the
        # integrator update (``_update``), not the output value, so we
        # only need it in the update path's prerequisite set.  The
        # scheduler resolves prerequisites of the periodic update via
        # its own input-port tickets — we list it here for parity with
        # the other dynamic ports so a tracer always sees the same
        # signature.
        if self.tracking_enabled:
            prereqs.append(self.input_ports[self.u_ext_index].ticket)
        # T-127-followup-bumpless-mode-switch — runtime mode-flag port
        # only feeds ``_update`` (gates the tracking-pull correction),
        # but list it in the prereq set for tracer parity with the
        # other dynamic ports.
        if self.tracking_enabled_dynamic:
            prereqs.append(self.input_ports[self.mode_flag_index].ticket)

        self.configure_output_port(
            self.control_output,
            self._output,
            period=self.dt,
            offset=0.0,
            default_value=initial_state,
            prerequisites_of_calc=prereqs,
        )

    def reset_default_values(self, **dynamic_parameters):
        initial_state = dynamic_parameters["initial_state"]
        _zero = 0.0
        if self._dtype is not None:
            initial_state = npa.asarray(initial_state).astype(self._dtype)
            _zero = npa.asarray(0.0).astype(self._dtype)
        self.configure_discrete_state_default_value(
            self._make_state(
                integral=initial_state,
                e_d_prev=_zero,
                e_dot_prev=_zero,
                e_i_prev=_zero,
                e_d_prev_prev=_zero,
            ),
            as_array=False,
        )
        self.configure_output_port_default_value(
            self.control_output, initial_state
        )

    # T-127-followup-external-weights ------------------------------------
    def _resolve_b(self, inputs, params):
        """Return ``b`` from the runtime port when ``b_dynamic`` is set,
        otherwise from the static dynamic-parameter ``b``."""
        if self.b_dynamic:
            return inputs[self.b_index]
        return params["b"]

    def _resolve_c(self, inputs, params):
        """Return ``c`` from the runtime port when ``c_dynamic`` is set,
        otherwise from the static dynamic-parameter ``c``."""
        if self.c_dynamic:
            return inputs[self.c_index]
        return params["c"]

    # T-127-followup-gain-scheduling -------------------------------------
    def _resolve_kp(self, inputs, params):
        """Return ``kp`` from the runtime port when ``kp_dynamic`` is
        set, otherwise from the static dynamic-parameter ``kp``."""
        if self.kp_dynamic:
            return inputs[self.kp_index]
        return params["kp"]

    def _resolve_ki(self, inputs, params):
        """Return ``ki`` from the runtime port when ``ki_dynamic`` is
        set, otherwise from the static dynamic-parameter ``ki``."""
        if self.ki_dynamic:
            return inputs[self.ki_index]
        return params["ki"]

    def _resolve_kd(self, inputs, params):
        """Return ``kd`` from the runtime port when ``kd_dynamic`` is
        set, otherwise from the static dynamic-parameter ``kd``."""
        if self.kd_dynamic:
            return inputs[self.kd_index]
        return params["kd"]

    def _resolve_kff(self, inputs, params):
        """Return ``kff`` from the runtime port when ``kff_dynamic`` is
        set, otherwise from the static dynamic-parameter ``kff`` (which
        defaults to 0.0 when feedforward is disabled)."""
        if self.kff_dynamic:
            return inputs[self.kff_index]
        return params.get("kff", 0.0)

    # T-127-followup-tracking-mode --------------------------------------
    def _resolve_u_ext(self, inputs):
        """Return the external tracking signal ``u_ext`` (the value the
        integrator is nudged toward when ``tracking_enabled=True``).

        Only callable when ``self.tracking_enabled`` is True; callers
        gate on the flag before invoking.
        """
        return inputs[self.u_ext_index]

    # T-127-followup-deadband-error -------------------------------------
    def _apply_deadband(self, e_raw, params):
        """Gate ``e_raw`` through the configured deadband.

        Returns ``e_raw`` unchanged when the deadband is inactive
        (``error_deadband == 0.0`` at construction time → byte-equivalent
        to phase 1).  In ``"hard"`` mode the gate is the standard
        ``where(|e_raw| > half_range, e_raw, 0)``; in ``"smooth"`` mode
        the gate uses :func:`soft_dead_zone`, which is differentiable
        through the band (this is the only place the deadband-followup
        differs from the inline phase 1 arithmetic).
        """
        if not self._error_deadband_active:
            return e_raw
        half_range = params.get("error_deadband", 0.0)
        if self._error_deadband_mode == "smooth":
            sharpness = params.get("error_deadband_sharpness", 10.0)
            return soft_dead_zone(e_raw, half_range, sharpness)
        # Hard gate — npa.where is kinked but the per-branch gradient is
        # finite, matching the DeadZone(mode="hard") semantics.
        return npa.where(npa.abs(e_raw) > half_range, e_raw, 0.0)

    def _eval_derivative(self, _time, state, *inputs, **params):
        # Filtered estimate of the time derivative of the *weighted*
        # derivative error e_d = c*r - y.
        c = self._resolve_c(inputs, params)
        r = inputs[self.setpoint_index]
        y = inputs[self.measurement_index]
        # T-127-followup-deadband-error — gate the raw derivative error
        # before it feeds the recursive filter / finite-difference
        # kernel.  ``e_d_prev`` was written by the previous tick's
        # ``_update``, which already applied the same gate, so the
        # filter state is consistent across ticks.
        e_d = self._apply_deadband(c * r - y, params)
        e_d_prev = state.discrete_state.e_d_prev
        b_coef, a_coef = self.filter

        if self.filter_type != "none":
            e_dot_prev = state.discrete_state.e_dot_prev
            e_dot = (
                b_coef[0] * e_d + b_coef[1] * e_d_prev - a_coef[1] * e_dot_prev
            ) / a_coef[0]
        else:
            # T-127-followup-discrete-integrator-derivative — kernel
            # dispatch.  ``forward_diff`` reproduces phase 1 exactly via
            # the (b=[1,-1], a=[dt,0]) coefficients.  The other kernels
            # ignore the filter coefficients and use explicit finite-
            # difference formulas over the (e_d, e_d_prev, e_d_prev_prev)
            # delay line.
            if self._derivative_method == "forward_diff":
                # Phase 1 path — uses the precomputed filter coefficients.
                e_dot = (b_coef[0] * e_d + b_coef[1] * e_d_prev) / a_coef[0]
            elif self._derivative_method == "backward_diff":
                # D[k] = (e[k] - e[k-1]) / dt — uses past data only, so
                # the *output* lags by one tick: read the previous tick's
                # finite difference from ``e_dot_prev`` (set during the
                # last _update).  This is what gives the documented
                # one-tick transient delay relative to forward_diff.
                e_dot = state.discrete_state.e_dot_prev
            else:  # "centered_diff"
                # D[k] = (e[k+1] - e[k-1]) / (2*dt) — uses the extra
                # delay cell stored in ``e_d_prev_prev``.
                e_d_prev_prev = state.discrete_state.e_d_prev_prev
                e_dot = (e_d - e_d_prev_prev) / (2.0 * self.dt)

        return e_dot

    # T-127-followup-anti-windup ------------------------------------------
    def _saturate(self, u, **params):
        """Clip ``u`` to ``[output_min, output_max]`` when configured.

        Returns ``u`` unchanged when no saturation limits were declared
        (preserves byte-equivalence with the phase 1 default path).
        """
        if not self._anti_windup_active:
            return u
        u_sat = u
        # Dynamic params are only present when declared at construction
        # time (None-valued ones get skipped by the @parameters decorator).
        u_max = params.get("output_max", None)
        u_min = params.get("output_min", None)
        if u_max is not None:
            u_sat = npa.minimum(u_sat, u_max)
        if u_min is not None:
            u_sat = npa.maximum(u_sat, u_min)
        return u_sat

    def _update(self, time, state, *inputs, **params):
        b = self._resolve_b(inputs, params)
        c = self._resolve_c(inputs, params)
        r = inputs[self.setpoint_index]
        y = inputs[self.measurement_index]
        # T-127-followup-deadband-error — gate every error signal that
        # feeds a downstream PID term.  Default error_deadband=0.0
        # leaves all three identities unchanged (byte-equivalent to
        # phase 1).  Each branch (P / I / D) goes through its own gate
        # so the weighted-setpoint semantics of ``b`` / ``c`` are
        # preserved (the gate is applied AFTER the weighting, on the
        # same composite signal the PID terms actually see).
        e_p = self._apply_deadband(b * r - y, params)
        e_i = self._apply_deadband(r - y, params)
        e_d = self._apply_deadband(c * r - y, params)

        e_int = state.discrete_state.integral

        # T-127-followup-discrete-integrator-derivative — when using the
        # ``backward_diff`` derivative kernel, the *next* stored
        # ``e_dot_prev`` is the finite difference computed from the
        # current samples (so the next tick's _output reads it as the
        # one-tick-delayed derivative).  For forward_diff/centered_diff
        # the field is left at zero (forward_diff doesn't read it,
        # centered_diff reads ``e_d_prev_prev`` instead).
        if self.filter_type != "none":
            # Recursive filters need e_dot_prev for the IIR update;
            # compute it via the standard filtered-derivative path.
            e_dot_next = self._eval_derivative(
                time, state, *inputs, **params
            )
        elif self._derivative_method == "backward_diff":
            # Store the just-computed (e_d - e_d_prev)/dt so the next
            # tick's _output sees it as the lagged derivative.
            e_d_prev = state.discrete_state.e_d_prev
            e_dot_next = (e_d - e_d_prev) / self.dt
        else:
            # forward_diff / centered_diff: e_dot_prev is unused on the
            # output path, keep it at the previous value (matches phase 1
            # placeholder semantics).
            e_dot_next = state.discrete_state.e_dot_prev

        # Integrator kernel dispatch.
        # forward_euler (phase 1): I[k+1] = I[k] + e[k] * dt — uses the
        # most-recent error sample.  backward_euler / trapezoidal pull
        # in the previously stored sample to differentiate them by the
        # documented one-step shift on a ramp.
        if self._integrator_method == "forward_euler":
            integral_next = e_int + e_i * self.dt
        else:
            e_i_prev = state.discrete_state.e_i_prev
            if self._integrator_method == "backward_euler":
                # I[k+1] = I[k] + e[k+1] * dt — labelled per the spec; the
                # available "next" sample at tick k is the *next-tick*
                # update's input.  We approximate by integrating the
                # previous sample, which produces the documented one-tick
                # lag relative to forward_euler.
                integral_next = e_int + e_i_prev * self.dt
            else:  # "trapezoidal"
                # I[k+1] = I[k] + (e[k] + e[k+1]) / 2 * dt — average of
                # current and previous error samples.
                integral_next = e_int + (e_i + e_i_prev) * 0.5 * self.dt

        # The integral consumed by ``_eval_control`` for the anti-windup
        # path below must match what the *next* output tick sees.  That
        # is the freshly computed ``integral_next`` (phase 1 used
        # ``e_int`` here; for byte-equivalence in the default config we
        # keep that behaviour by gating on _anti_windup_active).
        # T-127-followup-anti-windup / T-127-followup-tracking-mode -----
        # Both corrections need ``u_unsat`` — the value the controller
        # *would* publish before saturation.  Compute it once when either
        # mechanism is active and feed both branches.  When neither is
        # active this whole block is skipped → byte-equivalent to phase 1.
        aw_on = (
            self._anti_windup_active and self._anti_windup_method != "none"
        )
        tr_on = self.tracking_enabled
        if aw_on or tr_on:
            # Use the e_dot the next _output would read, for consistency
            # with the saturated-output path above.
            if self.filter_type != "none":
                e_dot_for_aw = e_dot_next
            elif self._derivative_method == "backward_diff":
                # _output reads state.e_dot_prev (the *previous* tick's
                # value); preserve that semantic here too.
                e_dot_for_aw = state.discrete_state.e_dot_prev
            elif self._derivative_method == "centered_diff":
                e_d_prev_prev = state.discrete_state.e_d_prev_prev
                e_dot_for_aw = (e_d - e_d_prev_prev) / (2.0 * self.dt)
            else:
                # forward_diff: re-use the phase-1 finite-difference
                # formula via the precomputed filter coefficients.
                b_coef, a_coef = self.filter
                e_d_prev = state.discrete_state.e_d_prev
                e_dot_for_aw = (
                    b_coef[0] * e_d + b_coef[1] * e_d_prev
                ) / a_coef[0]
            # T-127-followup-feedforward — feedforward is part of the
            # unsaturated control value the anti-windup logic compares
            # against ``u_sat``; pass ``r`` so ``_eval_control`` folds it
            # in.  Default kff=0.0 leaves the comparison unchanged.
            # T-127-followup-gain-scheduling — resolve runtime-port gains
            # here so the anti-windup comparison uses the same scheduled
            # values as the main ``_output`` path.
            kp_aw = self._resolve_kp(inputs, params)
            ki_aw = self._resolve_ki(inputs, params)
            kd_aw = self._resolve_kd(inputs, params)
            kff_aw = self._resolve_kff(inputs, params)
            u_unsat = self._eval_control(
                e_p, e_int, e_dot_for_aw,
                kp_aw, ki_aw, kd_aw, kff_aw, r=r,
            )
            if aw_on:
                u_sat = self._saturate(u_unsat, **params)
                if self._anti_windup_method == "back_calc":
                    # Back-calculation: pull the integrator toward the
                    # value that would have produced u_sat, with time
                    # constant Tt = anti_windup_gain.
                    tt = params["anti_windup_gain"]
                    integral_next = integral_next - (u_unsat - u_sat) / tt * self.dt
                elif self._anti_windup_method == "clamping":
                    # Integrator-tracking: only update when the controller
                    # is not pushing further into saturation.  ``saturating``
                    # is True when we are clamped AND the error sign matches
                    # the direction of saturation (positive sat → positive
                    # e_i pushes harder; negative sat → negative e_i pushes
                    # harder).  Use ``where`` for differentiability: gradient
                    # is zero on the clamped branch.
                    sat_excess = u_unsat - u_sat  # >0 if hit upper, <0 if lower
                    pushing_further = npa.logical_and(
                        sat_excess != 0,
                        npa.sign(e_i) == npa.sign(sat_excess),
                    )
                    integral_next = npa.where(
                        pushing_further, e_int, integral_next
                    )
            # T-127-followup-tracking-mode -----------------------------
            # Bumpless transfer: nudge the integrator toward the value
            # that would have produced ``u_ext``.  Implementation is
            # back-calculation with the *external* tracking signal in
            # place of ``u_sat``.  When anti-windup is also active the
            # two corrections sum (each is a small per-tick perturbation
            # of the integrator, so superposition is correct to first
            # order).
            # T-127-followup-i-on-error-only — when
            # ``integrate_tracking_error=False`` the tracking signal
            # ``u_ext`` MUST NOT touch the integrator.  We still declare
            # the ``u_ext`` port (callers may use it for downstream
            # pull-through paths) but the per-tick correction term is
            # suppressed entirely.  The default (``True``) preserves the
            # T-127-followup-tracking-mode kernel byte-for-byte.
            if tr_on and self.integrate_tracking_error:
                u_ext = self._resolve_u_ext(inputs)
                tt_tr = params.get("tracking_gain", 1.0)
                e_track = u_ext - u_unsat
                correction = (e_track / tt_tr) * self.dt
                # T-127-followup-bumpless-mode-switch — when the runtime
                # mode-flag port is declared, multiply the correction by
                # the gate so the user can flip auto/manual mid-sim.
                # The flag is cast to the integrator dtype so it remains
                # differentiable through ``tracking_gain`` (the gate
                # itself is a non-differentiable boolean — gradient
                # through the flag is zero).  Default
                # ``tracking_enabled_dynamic=False`` skips this
                # branch entirely, preserving the
                # T-127-followup-tracking-mode kernel byte-for-byte.
                if self.tracking_enabled_dynamic:
                    mode_flag = inputs[self.mode_flag_index]
                    # Explicit boolean coercion: any non-zero value
                    # turns tracking ON.  Multiplying by an array gate
                    # keeps the operation jit/grad-friendly.
                    gate = npa.where(
                        npa.asarray(mode_flag) != 0,
                        npa.asarray(1.0, dtype=correction.dtype),
                        npa.asarray(0.0, dtype=correction.dtype),
                    )
                    correction = correction * gate
                integral_next = integral_next + correction

        # Build the new state tuple, populating optional delay cells
        # only when the configured kernels need them.
        return self._make_state(
            integral=integral_next,
            e_d_prev=e_d,
            e_dot_prev=e_dot_next,
            e_i_prev=e_i,  # only stored when integrator_method != forward_euler
            e_d_prev_prev=state.discrete_state.e_d_prev,  # shift for centered_diff
        )

    def _eval_control(self, e_p, e_int, e_dot, kp, ki, kd, kff=0.0, r=None):
        # T-127-followup-feedforward — fold ``kff * r`` into the
        # unsaturated control sum.  ``kff`` defaults to 0.0 so the phase
        # 1 path (kff=0.0) reduces to the original
        # ``kp*e_p + ki*e_int + kd*e_dot`` arithmetic identity and
        # remains byte-equivalent.  ``r`` is only passed explicitly by
        # the saturation/anti-windup branches that need the unsaturated
        # value; older call sites without ``r`` get no feedforward
        # contribution.
        #
        # T-127-followup-gain-scheduling — the four gains are always
        # supplied by the caller (resolved from ``_resolve_k{p,i,d,ff}``),
        # so this helper neither consults ``params`` nor takes ``**params``.
        # That keeps the signature collision-free with the `**params`
        # unpacking used by ``_output`` / ``_update``.
        u = kp * e_p + ki * e_int + kd * e_dot
        if r is not None:
            u = u + kff * r
        return u

    def _output(self, time, state, *inputs, **params):
        b = self._resolve_b(inputs, params)
        r = inputs[self.setpoint_index]
        y = inputs[self.measurement_index]
        # T-127-followup-deadband-error — gate the proportional error
        # the same way ``_update`` does so the published ``u`` reflects
        # the deadband on every tick.  ``e_int`` already reflects the
        # gated integrator-error trajectory because ``_update`` writes
        # the gated ``e_i`` into the integral; ``e_dot`` is gated inside
        # ``_eval_derivative``.  Default error_deadband=0.0 is
        # byte-equivalent to phase 1.
        e_p = self._apply_deadband(b * r - y, params)
        e_int = state.discrete_state.integral
        e_dot = self._eval_derivative(time, state, *inputs, **params)
        # T-127-followup-feedforward — ``kff * r`` is added inside
        # ``_eval_control`` so it participates in saturation and the
        # anti-windup u_unsat/u_sat comparison.
        # T-127-followup-gain-scheduling — resolve each scalar gain from
        # its runtime port (if the matching ``*_dynamic`` flag is set)
        # or from the static parameter.  Default ``*_dynamic=False`` is
        # byte-equivalent to the phase 1 path.
        kp = self._resolve_kp(inputs, params)
        ki = self._resolve_ki(inputs, params)
        kd = self._resolve_kd(inputs, params)
        kff = self._resolve_kff(inputs, params)
        u = self._eval_control(e_p, e_int, e_dot, kp, ki, kd, kff, r=r)
        # T-127-followup-anti-windup — publish the saturated control
        # value when limits are configured.  Default-off path returns
        # ``u`` unchanged, matching phase 1 byte-for-byte.
        u = self._saturate(u, **params)
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        return u

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        # Use the setpoint port as the canonical input shape/dtype, matching
        # PIDDiscrete's single-input check_types pattern.
        u = self.eval_input(context, self.setpoint_index)
        xd = context[self.system_id].discrete_state.integral
        check_state_type(
            self,
            inp_data=u,
            state_data=xd,
            error_collector=error_collector,
        )

    # -----------------------------------------------------------------
    # T-127-followup-config-roundtrip — JSON-friendly config serialization.
    #
    # ``to_dict()`` captures every construction-time field that affects
    # behavior (the @parameters static + dynamic lists, plus the four
    # mode strings — anti_windup_method / integrator_method /
    # derivative_method / error_deadband_mode — and the ``*_dynamic``
    # port-topology flags).  ``from_dict()`` reconstructs an equivalent
    # block; round-tripping through ``json.dumps`` / ``json.loads`` is
    # supported because every encoded value is a Python primitive
    # (float / int / bool / str) or ``None``.
    #
    # Caveat (honest fallback documented in T-127-followup-config-
    # roundtrip): when any ``*_dynamic`` flag is True the corresponding
    # static scalar is still captured (it is the declared default that
    # would be used if the user later switched the flag off and
    # reconstructed without re-wiring); the dangling input-port
    # connection has to be re-built by the caller after ``from_dict``
    # since topology / wiring lives in the diagram, not the block.
    # -----------------------------------------------------------------

    # Static (non-dynamic-parameter) construction-time fields owned by
    # this block.  These are stored as attributes / static parameters
    # and never come from a runtime port.
    _CONFIG_STATIC_FIELDS = (
        "dt",
        "filter_type",
        "filter_coefficient",
        "anti_windup_method",
        "b_dynamic",
        "c_dynamic",
        "integrator_method",
        "derivative_method",
        "kp_dynamic",
        "ki_dynamic",
        "kd_dynamic",
        "kff_dynamic",
        "error_deadband_mode",
        "tracking_enabled",
        "integrate_tracking_error",
        "tracking_enabled_dynamic",
    )

    # Dynamic-parameter scalars.  These are declared via
    # ``declare_dynamic_parameter`` (possibly skipped when ``None`` —
    # see ``output_min`` / ``output_max``).
    _CONFIG_DYNAMIC_FIELDS = (
        "kp",
        "ki",
        "kd",
        "b",
        "c",
        "initial_state",
        "output_min",
        "output_max",
        "anti_windup_gain",
        "kff",
        "error_deadband",
        "error_deadband_sharpness",
        "tracking_gain",
    )

    @staticmethod
    def _encode_scalar(value):
        """Convert a parameter value (Python scalar, numpy / jax array)
        into a JSON-friendly primitive.

        - ``None`` stays ``None`` (used for unset saturation limits).
        - Scalar numbers / 0-D arrays become Python ``float``.
        - Booleans stay ``bool``.
        - Strings stay ``str``.
        Other types raise ``TypeError`` — config round-trip is only
        intended for the scalar-config subset.
        """
        if value is None:
            return None
        if isinstance(value, bool):
            return bool(value)
        if isinstance(value, (int, float)):
            return float(value)
        if isinstance(value, str):
            return value
        # numpy / jax array fallback: only 0-D scalars are supported.
        arr = np.asarray(value)
        if arr.shape == ():
            return float(arr)
        raise TypeError(
            f"PIDController2DOF.to_dict(): cannot serialize value "
            f"{value!r} of type {type(value).__name__}; only scalar "
            f"configuration values are supported."
        )

    def to_dict(self):
        """Return a JSON-serializable dict describing this controller.

        The dict captures every construction-time field that controls
        the block's behavior — all ``@parameters``-registered fields
        plus the mode strings and ``*_dynamic`` port-topology flags
        that live outside ``@parameters``.  Round-tripping through
        :meth:`from_dict` (optionally via ``json.dumps`` /
        ``json.loads``) produces a block with identical step-response
        behavior on a fixed input.

        Note:
            Diagram wiring (which signal feeds which input port) is
            not part of the block config and must be re-established by
            the caller after :meth:`from_dict`.  This is especially
            relevant when any ``*_dynamic`` flag is True — the
            reconstructed block still declares the runtime port, but
            it has no upstream connection until the caller wires it
            up.

        Returns:
            dict mapping each of :attr:`_CONFIG_STATIC_FIELDS` and
            :attr:`_CONFIG_DYNAMIC_FIELDS` to a JSON primitive.
        """
        data = {}
        # Static fields: read directly off the static-parameter dict
        # (where ``@parameters`` stashed them) or fall back to the
        # cached instance attribute when the kwarg lives there too.
        for name in self._CONFIG_STATIC_FIELDS:
            if name in self._static_parameters:
                value = Parameter.unwrap(self._static_parameters[name])
            else:
                value = getattr(self, name, None)
            data[name] = self._encode_scalar(value)
        # Dynamic fields: only declared when non-None (see
        # ``parameters`` decorator).  Encode missing entries as None
        # — that lets the constructor's own default reactivate on
        # ``from_dict`` and keeps ``output_min`` / ``output_max``
        # round-trippable.
        for name in self._CONFIG_DYNAMIC_FIELDS:
            if name in self._dynamic_parameters:
                value = Parameter.unwrap(self._dynamic_parameters[name])
            else:
                value = None
            data[name] = self._encode_scalar(value)
        return data

    @classmethod
    def from_dict(cls, data, **block_kwargs):
        """Reconstruct a :class:`PIDController2DOF` from a config dict.

        Extra keyword arguments (``name=``, ``system_id=``, ...) are
        forwarded to the constructor so a deserialized block can pick
        up a fresh name in its target diagram.

        Args:
            data: dict produced by :meth:`to_dict` (or any equivalent
                mapping with the same key set).  ``dt`` is the only
                required field; every other field falls back to the
                constructor default when absent.
            **block_kwargs: forwarded to ``PIDController2DOF.__init__``
                (typically ``name=...``).

        Raises:
            ValueError: if ``data`` lacks the required ``dt`` field.

        Returns:
            A new :class:`PIDController2DOF` whose configuration
            matches ``data``.
        """
        if "dt" not in data or data["dt"] is None:
            raise ValueError(
                "PIDController2DOF.from_dict(): missing required field "
                "'dt' (the controller sample period)."
            )
        # Build the constructor kwargs.  Drop missing / None-valued
        # entries for fields whose constructor default is *not* None
        # so the constructor picks them up; preserve None for
        # ``output_min`` / ``output_max`` (their default IS None and
        # round-trip needs them to stay None).
        ctor_kwargs = {}
        for name in cls._CONFIG_STATIC_FIELDS + cls._CONFIG_DYNAMIC_FIELDS:
            if name not in data:
                continue
            value = data[name]
            # output_min / output_max naturally accept None.
            if value is None and name not in ("output_min", "output_max"):
                continue
            ctor_kwargs[name] = value
        ctor_kwargs.update(block_kwargs)
        return cls(**ctor_kwargs)

    # ------------------------------------------------------------------
    # T-127-followup-derivative-on-measurement — convenience factories
    # for the two common 2-DOF PID configurations.  ``standard`` keeps
    # the textbook ``b = c = 1`` defaults (equivalent to a 1-DOF PID on
    # the error signal), while ``with_derivative_on_measurement`` ships
    # ``b = 1, c = 0`` — the standard "no derivative kick" recipe
    # recommended for real-world controllers (a step change in the
    # setpoint no longer produces a ``Kd / dt`` spike through the
    # derivative term).
    # ------------------------------------------------------------------
    @classmethod
    def standard(cls, kp, ki, kd, dt, **kwargs):
        """Construct a textbook PID with ``b = c = 1``.

        Convenience factory equivalent to ``PIDController2DOF(dt, kp,
        ki, kd)``: the proportional and derivative paths both see the
        setpoint with weight 1, so the block reduces to a 1-DOF PID on
        the error signal ``e = r - y``.  Useful as the explicit
        counterpart to :meth:`with_derivative_on_measurement` —
        callers self-document which 2-DOF configuration they want.

        Args:
            kp: Proportional gain.
            ki: Integral gain.
            kd: Derivative gain.
            dt: Sampling period.
            **kwargs: Forwarded to :class:`PIDController2DOF`.  Setting
                ``b`` or ``c`` here is allowed but discouraged (use the
                main constructor for non-standard weights).

        Returns:
            A :class:`PIDController2DOF` with ``b = c = 1``.
        """
        kwargs.setdefault("b", 1.0)
        kwargs.setdefault("c", 1.0)
        return cls(dt=dt, kp=kp, ki=ki, kd=kd, **kwargs)

    @classmethod
    def with_derivative_on_measurement(cls, kp, ki, kd, dt, **kwargs):
        """Construct a PID with derivative-on-measurement-only (``b=1, c=0``).

        Convenience factory for the standard "no derivative kick"
        recipe used by most real-world controllers.  With ``c = 0`` the
        derivative term sees only the measurement (``-d/dt(y)``), so a
        step change in the setpoint does NOT inject a ``Kd / dt`` spike
        through the derivative path — only integral and proportional
        action drive the transient.

        Args:
            kp: Proportional gain.
            ki: Integral gain.
            kd: Derivative gain.
            dt: Sampling period.
            **kwargs: Forwarded to :class:`PIDController2DOF`.  Passing
                ``c`` or ``c_dynamic`` here raises ``ValueError`` via
                the constructor's ``derivative_on_measurement_only``
                contract — by construction this factory pins ``c=0``.

        Returns:
            A :class:`PIDController2DOF` with ``b = 1`` and ``c = 0``.
        """
        kwargs.setdefault("b", 1.0)
        return cls(
            dt=dt,
            kp=kp,
            ki=ki,
            kd=kd,
            derivative_on_measurement_only=True,
            **kwargs,
        )

    # ------------------------------------------------------------------
    # T-127-followup-pid-tuning-helpers — classical tuning-rule
    # classmethods.  Each helper computes ``(Kp, Ki, Kd)`` from the
    # plant-characterisation inputs and returns a configured
    # :class:`PIDController2DOF`.  The helpers are pure factories: they
    # add no behavioural surface to the existing PID class and the gains
    # they emit are forwarded to the standard constructor unchanged.
    #
    # Three rules are shipped (the classical control-engineering set):
    #
    # 1. ``ziegler_nichols`` — closed-loop ultimate-cycle method.  Given
    #    the ultimate gain ``Ku`` (the proportional gain at which the
    #    closed loop just oscillates with sustained period ``Tu``), the
    #    Z-N table maps to PID gains for P / PI / PID modes.  The
    #    coefficients are the canonical Ziegler & Nichols (1942) values
    #    reproduced in every textbook (e.g. Astrom & Hagglund).
    #
    # 2. ``cohen_coon`` — open-loop process-reaction-curve method for a
    #    first-order-plus-dead-time (FOPDT) plant
    #    ``G(s) = K * exp(-theta*s) / (tau*s + 1)``.  Slightly more
    #    aggressive than Z-N on dead-time-dominated plants; widely used
    #    for chemical-process control.
    #
    # 3. ``tyreus_luyben`` — Z-N alternative with a much longer integral
    #    time (PI: ``Kp = Ku/3.2, Ti = 2.2*Tu``).  Trades response speed
    #    for robustness; commonly used when the Z-N gains produce too
    #    much overshoot or sensitivity to model error.
    # ------------------------------------------------------------------
    _ZIEGLER_NICHOLS_MODES = ("P", "PI", "PID")
    _COHEN_COON_MODES = ("P", "PI", "PID")

    @classmethod
    def ziegler_nichols(cls, Ku, Tu, dt, mode="PID", **kwargs):
        """Construct a PID tuned by the Ziegler-Nichols ultimate-cycle rule.

        Given the ultimate gain ``Ku`` (the proportional-only gain at
        which the closed loop just sustains oscillation) and the
        corresponding ultimate period ``Tu``, the Z-N table maps to
        controller gains as::

            P:    Kp = 0.5  * Ku,                 Ki = 0,            Kd = 0
            PI:   Kp = 0.45 * Ku, Ti = Tu / 1.2 → Ki = 0.54*Ku/Tu,   Kd = 0
            PID:  Kp = 0.6  * Ku, Ti = Tu / 2.0 → Ki = 1.2 *Ku/Tu,
                  Td = Tu / 8.0                 → Kd = 0.075*Ku*Tu

        The coefficients are the canonical Ziegler & Nichols (1942)
        values; see e.g. Astrom & Hagglund, *PID Controllers: Theory,
        Design, and Tuning* (1995), Table 4.1.

        Args:
            Ku: Ultimate gain (proportional-only gain at sustained
                oscillation). Must be positive.
            Tu: Ultimate period (period of the sustained oscillation
                in seconds). Must be positive.
            dt: Sampling period for the discrete PID.
            mode: One of ``"P"``, ``"PI"``, ``"PID"`` (default
                ``"PID"``). Selects which gains are non-zero.
            **kwargs: Forwarded to :class:`PIDController2DOF`.

        Returns:
            A :class:`PIDController2DOF` whose ``(Kp, Ki, Kd)`` match
            the Z-N table for the requested mode.

        Raises:
            ValueError: If ``mode`` is not one of ``"P"``, ``"PI"``,
                ``"PID"``, or if ``Ku`` / ``Tu`` are non-positive.
        """
        if mode not in cls._ZIEGLER_NICHOLS_MODES:
            raise ValueError(
                f"ziegler_nichols: mode must be one of "
                f"{cls._ZIEGLER_NICHOLS_MODES!r}; got {mode!r}"
            )
        Ku_f = float(Ku)
        Tu_f = float(Tu)
        if Ku_f <= 0.0:
            raise ValueError(
                f"ziegler_nichols: Ku must be positive; got {Ku_f}"
            )
        if Tu_f <= 0.0:
            raise ValueError(
                f"ziegler_nichols: Tu must be positive; got {Tu_f}"
            )
        if mode == "P":
            kp = 0.5 * Ku_f
            ki = 0.0
            kd = 0.0
        elif mode == "PI":
            kp = 0.45 * Ku_f
            ki = 0.54 * Ku_f / Tu_f
            kd = 0.0
        else:  # PID
            kp = 0.6 * Ku_f
            ki = 1.2 * Ku_f / Tu_f
            kd = 0.075 * Ku_f * Tu_f
        return cls(dt=dt, kp=kp, ki=ki, kd=kd, **kwargs)

    @classmethod
    def cohen_coon(cls, K, tau, theta, dt, mode="PID", **kwargs):
        """Construct a PID tuned by the Cohen-Coon rule for a FOPDT plant.

        For a first-order-plus-dead-time plant
        ``G(s) = K * exp(-theta*s) / (tau*s + 1)`` (process gain ``K``,
        time constant ``tau``, dead time ``theta``), the Cohen-Coon
        (1953) formulas are::

            r = theta / tau

            P:    Kp = (1/K) * (1/r) * (1 + r/3)

            PI:   Kp = (1/K) * (1/r) * (9/10 + r/12)
                  Ti = theta * (30 + 3*r) / (9 + 20*r)

            PID:  Kp = (1/K) * (1/r) * (4/3 + r/4)
                  Ti = theta * (32 + 6*r) / (13 + 8*r)
                  Td = theta * 4 / (11 + 2*r)

        with ``Ki = Kp / Ti`` and ``Kd = Kp * Td``.

        Cohen-Coon is more aggressive than Ziegler-Nichols on plants
        where ``theta / tau`` is large (dead-time-dominated processes);
        it is widely used in chemical-process control.

        Args:
            K: Process (steady-state) gain. Must be non-zero.
            tau: First-order time constant in seconds. Must be positive.
            theta: Dead time in seconds. Must be positive.
            dt: Sampling period for the discrete PID.
            mode: One of ``"P"``, ``"PI"``, ``"PID"`` (default
                ``"PID"``). Selects which gains are non-zero.
            **kwargs: Forwarded to :class:`PIDController2DOF`.

        Returns:
            A :class:`PIDController2DOF` whose ``(Kp, Ki, Kd)`` match
            the Cohen-Coon formulas for the requested mode.

        Raises:
            ValueError: If ``mode`` is not one of ``"P"``, ``"PI"``,
                ``"PID"``, or if ``K`` is zero, or if ``tau`` /
                ``theta`` are non-positive.
        """
        if mode not in cls._COHEN_COON_MODES:
            raise ValueError(
                f"cohen_coon: mode must be one of "
                f"{cls._COHEN_COON_MODES!r}; got {mode!r}"
            )
        K_f = float(K)
        tau_f = float(tau)
        theta_f = float(theta)
        if K_f == 0.0:
            raise ValueError("cohen_coon: K must be non-zero")
        if tau_f <= 0.0:
            raise ValueError(
                f"cohen_coon: tau must be positive; got {tau_f}"
            )
        if theta_f <= 0.0:
            raise ValueError(
                f"cohen_coon: theta must be positive; got {theta_f}"
            )
        # Use npa for the arithmetic so the helper composes with the
        # backend selector even when callers pass JAX scalars.  npa is
        # already imported at module scope.
        r = npa.divide(theta_f, tau_f)
        inv_K = npa.divide(1.0, K_f)
        inv_r = npa.divide(1.0, r)
        if mode == "P":
            kp = float(inv_K * inv_r * (1.0 + r / 3.0))
            ki = 0.0
            kd = 0.0
        elif mode == "PI":
            kp = float(inv_K * inv_r * (9.0 / 10.0 + r / 12.0))
            Ti = float(theta_f * (30.0 + 3.0 * r) / (9.0 + 20.0 * r))
            ki = kp / Ti
            kd = 0.0
        else:  # PID
            kp = float(inv_K * inv_r * (4.0 / 3.0 + r / 4.0))
            Ti = float(theta_f * (32.0 + 6.0 * r) / (13.0 + 8.0 * r))
            Td = float(theta_f * 4.0 / (11.0 + 2.0 * r))
            ki = kp / Ti
            kd = kp * Td
        return cls(dt=dt, kp=kp, ki=ki, kd=kd, **kwargs)

    @classmethod
    def tyreus_luyben(cls, Ku, Tu, dt, **kwargs):
        """Construct a PI controller tuned by the Tyreus-Luyben rule.

        A gentler Ziegler-Nichols alternative that trades response
        speed for robustness.  Tyreus & Luyben (1992) recommend the
        PI form for most chemical-process applications because the
        derivative term tends to amplify measurement noise::

            Kp = Ku / 3.2,  Ti = 2.2 * Tu  →  Ki = Kp / Ti

        and ``Kd = 0`` (no derivative action).  Compared with Z-N,
        Tyreus-Luyben gives roughly 1/3 the proportional gain and a
        ~4x longer integral time, producing a much less aggressive
        loop with substantially better robustness to model error.

        Args:
            Ku: Ultimate gain (proportional-only gain at sustained
                oscillation). Must be positive.
            Tu: Ultimate period (period of the sustained oscillation
                in seconds). Must be positive.
            dt: Sampling period for the discrete PID.
            **kwargs: Forwarded to :class:`PIDController2DOF`.

        Returns:
            A :class:`PIDController2DOF` configured as a PI
            controller (``Kd = 0``) with the Tyreus-Luyben gains.

        Raises:
            ValueError: If ``Ku`` / ``Tu`` are non-positive.
        """
        Ku_f = float(Ku)
        Tu_f = float(Tu)
        if Ku_f <= 0.0:
            raise ValueError(
                f"tyreus_luyben: Ku must be positive; got {Ku_f}"
            )
        if Tu_f <= 0.0:
            raise ValueError(
                f"tyreus_luyben: Tu must be positive; got {Tu_f}"
            )
        kp = Ku_f / 3.2
        Ti = 2.2 * Tu_f
        ki = kp / Ti
        kd = 0.0
        return cls(dt=dt, kp=kp, ki=ki, kd=kd, **kwargs)

cohen_coon(K, tau, theta, dt, mode='PID', **kwargs) classmethod

Construct a PID tuned by the Cohen-Coon rule for a FOPDT plant.

For a first-order-plus-dead-time plant G(s) = K * exp(-theta*s) / (tau*s + 1) (process gain K, time constant tau, dead time theta), the Cohen-Coon (1953) formulas are::

r = theta / tau

P:    Kp = (1/K) * (1/r) * (1 + r/3)

PI:   Kp = (1/K) * (1/r) * (9/10 + r/12)
      Ti = theta * (30 + 3*r) / (9 + 20*r)

PID:  Kp = (1/K) * (1/r) * (4/3 + r/4)
      Ti = theta * (32 + 6*r) / (13 + 8*r)
      Td = theta * 4 / (11 + 2*r)

with Ki = Kp / Ti and Kd = Kp * Td.

Cohen-Coon is more aggressive than Ziegler-Nichols on plants where theta / tau is large (dead-time-dominated processes); it is widely used in chemical-process control.

Parameters:

Name	Type	Description	Default
K		Process (steady-state) gain. Must be non-zero.	required
tau		First-order time constant in seconds. Must be positive.	required
theta		Dead time in seconds. Must be positive.	required
dt		Sampling period for the discrete PID.	required
mode		One of "P", "PI", "PID" (default "PID"). Selects which gains are non-zero.	'PID'
**kwargs		Forwarded to :class:PIDController2DOF.	{}

Returns:

Name	Type	Description
A		class:PIDController2DOF whose (Kp, Ki, Kd) match
		the Cohen-Coon formulas for the requested mode.

Raises:

Type	Description
ValueError	If mode is not one of "P", "PI", "PID", or if K is zero, or if tau / theta are non-positive.

Source code in jaxonomy/library/dynamics.py

@classmethod
def cohen_coon(cls, K, tau, theta, dt, mode="PID", **kwargs):
    """Construct a PID tuned by the Cohen-Coon rule for a FOPDT plant.

    For a first-order-plus-dead-time plant
    ``G(s) = K * exp(-theta*s) / (tau*s + 1)`` (process gain ``K``,
    time constant ``tau``, dead time ``theta``), the Cohen-Coon
    (1953) formulas are::

        r = theta / tau

        P:    Kp = (1/K) * (1/r) * (1 + r/3)

        PI:   Kp = (1/K) * (1/r) * (9/10 + r/12)
              Ti = theta * (30 + 3*r) / (9 + 20*r)

        PID:  Kp = (1/K) * (1/r) * (4/3 + r/4)
              Ti = theta * (32 + 6*r) / (13 + 8*r)
              Td = theta * 4 / (11 + 2*r)

    with ``Ki = Kp / Ti`` and ``Kd = Kp * Td``.

    Cohen-Coon is more aggressive than Ziegler-Nichols on plants
    where ``theta / tau`` is large (dead-time-dominated processes);
    it is widely used in chemical-process control.

    Args:
        K: Process (steady-state) gain. Must be non-zero.
        tau: First-order time constant in seconds. Must be positive.
        theta: Dead time in seconds. Must be positive.
        dt: Sampling period for the discrete PID.
        mode: One of ``"P"``, ``"PI"``, ``"PID"`` (default
            ``"PID"``). Selects which gains are non-zero.
        **kwargs: Forwarded to :class:`PIDController2DOF`.

    Returns:
        A :class:`PIDController2DOF` whose ``(Kp, Ki, Kd)`` match
        the Cohen-Coon formulas for the requested mode.

    Raises:
        ValueError: If ``mode`` is not one of ``"P"``, ``"PI"``,
            ``"PID"``, or if ``K`` is zero, or if ``tau`` /
            ``theta`` are non-positive.
    """
    if mode not in cls._COHEN_COON_MODES:
        raise ValueError(
            f"cohen_coon: mode must be one of "
            f"{cls._COHEN_COON_MODES!r}; got {mode!r}"
        )
    K_f = float(K)
    tau_f = float(tau)
    theta_f = float(theta)
    if K_f == 0.0:
        raise ValueError("cohen_coon: K must be non-zero")
    if tau_f <= 0.0:
        raise ValueError(
            f"cohen_coon: tau must be positive; got {tau_f}"
        )
    if theta_f <= 0.0:
        raise ValueError(
            f"cohen_coon: theta must be positive; got {theta_f}"
        )
    # Use npa for the arithmetic so the helper composes with the
    # backend selector even when callers pass JAX scalars.  npa is
    # already imported at module scope.
    r = npa.divide(theta_f, tau_f)
    inv_K = npa.divide(1.0, K_f)
    inv_r = npa.divide(1.0, r)
    if mode == "P":
        kp = float(inv_K * inv_r * (1.0 + r / 3.0))
        ki = 0.0
        kd = 0.0
    elif mode == "PI":
        kp = float(inv_K * inv_r * (9.0 / 10.0 + r / 12.0))
        Ti = float(theta_f * (30.0 + 3.0 * r) / (9.0 + 20.0 * r))
        ki = kp / Ti
        kd = 0.0
    else:  # PID
        kp = float(inv_K * inv_r * (4.0 / 3.0 + r / 4.0))
        Ti = float(theta_f * (32.0 + 6.0 * r) / (13.0 + 8.0 * r))
        Td = float(theta_f * 4.0 / (11.0 + 2.0 * r))
        ki = kp / Ti
        kd = kp * Td
    return cls(dt=dt, kp=kp, ki=ki, kd=kd, **kwargs)

from_dict(data, **block_kwargs) classmethod

Reconstruct a :class:PIDController2DOF from a config dict.

Extra keyword arguments (name=, system_id=, ...) are forwarded to the constructor so a deserialized block can pick up a fresh name in its target diagram.

Parameters:

Name	Type	Description	Default
data		dict produced by :meth:to_dict (or any equivalent mapping with the same key set). dt is the only required field; every other field falls back to the constructor default when absent.	required
**block_kwargs		forwarded to PIDController2DOF.__init__ (typically name=...).	{}

Raises:

Type	Description
ValueError	if data lacks the required dt field.

Returns:

Type	Description
	A new :class:PIDController2DOF whose configuration
	matches data.

Source code in jaxonomy/library/dynamics.py

@classmethod
def from_dict(cls, data, **block_kwargs):
    """Reconstruct a :class:`PIDController2DOF` from a config dict.

    Extra keyword arguments (``name=``, ``system_id=``, ...) are
    forwarded to the constructor so a deserialized block can pick
    up a fresh name in its target diagram.

    Args:
        data: dict produced by :meth:`to_dict` (or any equivalent
            mapping with the same key set).  ``dt`` is the only
            required field; every other field falls back to the
            constructor default when absent.
        **block_kwargs: forwarded to ``PIDController2DOF.__init__``
            (typically ``name=...``).

    Raises:
        ValueError: if ``data`` lacks the required ``dt`` field.

    Returns:
        A new :class:`PIDController2DOF` whose configuration
        matches ``data``.
    """
    if "dt" not in data or data["dt"] is None:
        raise ValueError(
            "PIDController2DOF.from_dict(): missing required field "
            "'dt' (the controller sample period)."
        )
    # Build the constructor kwargs.  Drop missing / None-valued
    # entries for fields whose constructor default is *not* None
    # so the constructor picks them up; preserve None for
    # ``output_min`` / ``output_max`` (their default IS None and
    # round-trip needs them to stay None).
    ctor_kwargs = {}
    for name in cls._CONFIG_STATIC_FIELDS + cls._CONFIG_DYNAMIC_FIELDS:
        if name not in data:
            continue
        value = data[name]
        # output_min / output_max naturally accept None.
        if value is None and name not in ("output_min", "output_max"):
            continue
        ctor_kwargs[name] = value
    ctor_kwargs.update(block_kwargs)
    return cls(**ctor_kwargs)

standard(kp, ki, kd, dt, **kwargs) classmethod

Construct a textbook PID with b = c = 1.

Convenience factory equivalent to PIDController2DOF(dt, kp, ki, kd): the proportional and derivative paths both see the setpoint with weight 1, so the block reduces to a 1-DOF PID on the error signal e = r - y. Useful as the explicit counterpart to :meth:with_derivative_on_measurement — callers self-document which 2-DOF configuration they want.

Parameters:

Name	Type	Description	Default
kp		Proportional gain.	required
ki		Integral gain.	required
kd		Derivative gain.	required
dt		Sampling period.	required
**kwargs		Forwarded to :class:PIDController2DOF. Setting b or c here is allowed but discouraged (use the main constructor for non-standard weights).	{}

Returns:

Name	Type	Description
A		class:PIDController2DOF with b = c = 1.

Source code in jaxonomy/library/dynamics.py

@classmethod
def standard(cls, kp, ki, kd, dt, **kwargs):
    """Construct a textbook PID with ``b = c = 1``.

    Convenience factory equivalent to ``PIDController2DOF(dt, kp,
    ki, kd)``: the proportional and derivative paths both see the
    setpoint with weight 1, so the block reduces to a 1-DOF PID on
    the error signal ``e = r - y``.  Useful as the explicit
    counterpart to :meth:`with_derivative_on_measurement` —
    callers self-document which 2-DOF configuration they want.

    Args:
        kp: Proportional gain.
        ki: Integral gain.
        kd: Derivative gain.
        dt: Sampling period.
        **kwargs: Forwarded to :class:`PIDController2DOF`.  Setting
            ``b`` or ``c`` here is allowed but discouraged (use the
            main constructor for non-standard weights).

    Returns:
        A :class:`PIDController2DOF` with ``b = c = 1``.
    """
    kwargs.setdefault("b", 1.0)
    kwargs.setdefault("c", 1.0)
    return cls(dt=dt, kp=kp, ki=ki, kd=kd, **kwargs)

to_dict()

Return a JSON-serializable dict describing this controller.

The dict captures every construction-time field that controls the block's behavior — all @parameters-registered fields plus the mode strings and *_dynamic port-topology flags that live outside @parameters. Round-tripping through :meth:from_dict (optionally via json.dumps / json.loads) produces a block with identical step-response behavior on a fixed input.

Note

Diagram wiring (which signal feeds which input port) is not part of the block config and must be re-established by the caller after :meth:from_dict. This is especially relevant when any *_dynamic flag is True — the reconstructed block still declares the runtime port, but it has no upstream connection until the caller wires it up.

Returns:

Type	Description
	dict mapping each of :attr:_CONFIG_STATIC_FIELDS and
	attr:_CONFIG_DYNAMIC_FIELDS to a JSON primitive.

Source code in jaxonomy/library/dynamics.py

def to_dict(self):
    """Return a JSON-serializable dict describing this controller.

    The dict captures every construction-time field that controls
    the block's behavior — all ``@parameters``-registered fields
    plus the mode strings and ``*_dynamic`` port-topology flags
    that live outside ``@parameters``.  Round-tripping through
    :meth:`from_dict` (optionally via ``json.dumps`` /
    ``json.loads``) produces a block with identical step-response
    behavior on a fixed input.

    Note:
        Diagram wiring (which signal feeds which input port) is
        not part of the block config and must be re-established by
        the caller after :meth:`from_dict`.  This is especially
        relevant when any ``*_dynamic`` flag is True — the
        reconstructed block still declares the runtime port, but
        it has no upstream connection until the caller wires it
        up.

    Returns:
        dict mapping each of :attr:`_CONFIG_STATIC_FIELDS` and
        :attr:`_CONFIG_DYNAMIC_FIELDS` to a JSON primitive.
    """
    data = {}
    # Static fields: read directly off the static-parameter dict
    # (where ``@parameters`` stashed them) or fall back to the
    # cached instance attribute when the kwarg lives there too.
    for name in self._CONFIG_STATIC_FIELDS:
        if name in self._static_parameters:
            value = Parameter.unwrap(self._static_parameters[name])
        else:
            value = getattr(self, name, None)
        data[name] = self._encode_scalar(value)
    # Dynamic fields: only declared when non-None (see
    # ``parameters`` decorator).  Encode missing entries as None
    # — that lets the constructor's own default reactivate on
    # ``from_dict`` and keeps ``output_min`` / ``output_max``
    # round-trippable.
    for name in self._CONFIG_DYNAMIC_FIELDS:
        if name in self._dynamic_parameters:
            value = Parameter.unwrap(self._dynamic_parameters[name])
        else:
            value = None
        data[name] = self._encode_scalar(value)
    return data

tyreus_luyben(Ku, Tu, dt, **kwargs) classmethod

Construct a PI controller tuned by the Tyreus-Luyben rule.

A gentler Ziegler-Nichols alternative that trades response speed for robustness. Tyreus & Luyben (1992) recommend the PI form for most chemical-process applications because the derivative term tends to amplify measurement noise::

Kp = Ku / 3.2,  Ti = 2.2 * Tu  →  Ki = Kp / Ti

and Kd = 0 (no derivative action). Compared with Z-N, Tyreus-Luyben gives roughly 1/3 the proportional gain and a ~4x longer integral time, producing a much less aggressive loop with substantially better robustness to model error.

Parameters:

Name	Type	Description	Default
Ku		Ultimate gain (proportional-only gain at sustained oscillation). Must be positive.	required
Tu		Ultimate period (period of the sustained oscillation in seconds). Must be positive.	required
dt		Sampling period for the discrete PID.	required
**kwargs		Forwarded to :class:PIDController2DOF.	{}

Returns:

Name	Type	Description
A		class:PIDController2DOF configured as a PI
		controller (Kd = 0) with the Tyreus-Luyben gains.

Raises:

Type	Description
ValueError	If Ku / Tu are non-positive.

Source code in jaxonomy/library/dynamics.py

@classmethod
def tyreus_luyben(cls, Ku, Tu, dt, **kwargs):
    """Construct a PI controller tuned by the Tyreus-Luyben rule.

    A gentler Ziegler-Nichols alternative that trades response
    speed for robustness.  Tyreus & Luyben (1992) recommend the
    PI form for most chemical-process applications because the
    derivative term tends to amplify measurement noise::

        Kp = Ku / 3.2,  Ti = 2.2 * Tu  →  Ki = Kp / Ti

    and ``Kd = 0`` (no derivative action).  Compared with Z-N,
    Tyreus-Luyben gives roughly 1/3 the proportional gain and a
    ~4x longer integral time, producing a much less aggressive
    loop with substantially better robustness to model error.

    Args:
        Ku: Ultimate gain (proportional-only gain at sustained
            oscillation). Must be positive.
        Tu: Ultimate period (period of the sustained oscillation
            in seconds). Must be positive.
        dt: Sampling period for the discrete PID.
        **kwargs: Forwarded to :class:`PIDController2DOF`.

    Returns:
        A :class:`PIDController2DOF` configured as a PI
        controller (``Kd = 0``) with the Tyreus-Luyben gains.

    Raises:
        ValueError: If ``Ku`` / ``Tu`` are non-positive.
    """
    Ku_f = float(Ku)
    Tu_f = float(Tu)
    if Ku_f <= 0.0:
        raise ValueError(
            f"tyreus_luyben: Ku must be positive; got {Ku_f}"
        )
    if Tu_f <= 0.0:
        raise ValueError(
            f"tyreus_luyben: Tu must be positive; got {Tu_f}"
        )
    kp = Ku_f / 3.2
    Ti = 2.2 * Tu_f
    ki = kp / Ti
    kd = 0.0
    return cls(dt=dt, kp=kp, ki=ki, kd=kd, **kwargs)

with_derivative_on_measurement(kp, ki, kd, dt, **kwargs) classmethod

Construct a PID with derivative-on-measurement-only (b=1, c=0).

Convenience factory for the standard "no derivative kick" recipe used by most real-world controllers. With c = 0 the derivative term sees only the measurement (-d/dt(y)), so a step change in the setpoint does NOT inject a Kd / dt spike through the derivative path — only integral and proportional action drive the transient.

Parameters:

Name	Type	Description	Default
kp		Proportional gain.	required
ki		Integral gain.	required
kd		Derivative gain.	required
dt		Sampling period.	required
**kwargs		Forwarded to :class:PIDController2DOF. Passing c or c_dynamic here raises ValueError via the constructor's derivative_on_measurement_only contract — by construction this factory pins c=0.	{}

Returns:

Name	Type	Description
A		class:PIDController2DOF with b = 1 and c = 0.

Source code in jaxonomy/library/dynamics.py

@classmethod
def with_derivative_on_measurement(cls, kp, ki, kd, dt, **kwargs):
    """Construct a PID with derivative-on-measurement-only (``b=1, c=0``).

    Convenience factory for the standard "no derivative kick"
    recipe used by most real-world controllers.  With ``c = 0`` the
    derivative term sees only the measurement (``-d/dt(y)``), so a
    step change in the setpoint does NOT inject a ``Kd / dt`` spike
    through the derivative path — only integral and proportional
    action drive the transient.

    Args:
        kp: Proportional gain.
        ki: Integral gain.
        kd: Derivative gain.
        dt: Sampling period.
        **kwargs: Forwarded to :class:`PIDController2DOF`.  Passing
            ``c`` or ``c_dynamic`` here raises ``ValueError`` via
            the constructor's ``derivative_on_measurement_only``
            contract — by construction this factory pins ``c=0``.

    Returns:
        A :class:`PIDController2DOF` with ``b = 1`` and ``c = 0``.
    """
    kwargs.setdefault("b", 1.0)
    return cls(
        dt=dt,
        kp=kp,
        ki=ki,
        kd=kd,
        derivative_on_measurement_only=True,
        **kwargs,
    )

ziegler_nichols(Ku, Tu, dt, mode='PID', **kwargs) classmethod

Construct a PID tuned by the Ziegler-Nichols ultimate-cycle rule.

Given the ultimate gain Ku (the proportional-only gain at which the closed loop just sustains oscillation) and the corresponding ultimate period Tu, the Z-N table maps to controller gains as::

P:    Kp = 0.5  * Ku,                 Ki = 0,            Kd = 0
PI:   Kp = 0.45 * Ku, Ti = Tu / 1.2 → Ki = 0.54*Ku/Tu,   Kd = 0
PID:  Kp = 0.6  * Ku, Ti = Tu / 2.0 → Ki = 1.2 *Ku/Tu,
      Td = Tu / 8.0                 → Kd = 0.075*Ku*Tu

The coefficients are the canonical Ziegler & Nichols (1942) values; see e.g. Astrom & Hagglund, PID Controllers: Theory, Design, and Tuning (1995), Table 4.1.

Parameters:

Name	Type	Description	Default
Ku		Ultimate gain (proportional-only gain at sustained oscillation). Must be positive.	required
Tu		Ultimate period (period of the sustained oscillation in seconds). Must be positive.	required
dt		Sampling period for the discrete PID.	required
mode		One of "P", "PI", "PID" (default "PID"). Selects which gains are non-zero.	'PID'
**kwargs		Forwarded to :class:PIDController2DOF.	{}

Returns:

Name	Type	Description
A		class:PIDController2DOF whose (Kp, Ki, Kd) match
		the Z-N table for the requested mode.

Raises:

Type	Description
ValueError	If mode is not one of "P", "PI", "PID", or if Ku / Tu are non-positive.

Source code in jaxonomy/library/dynamics.py

@classmethod
def ziegler_nichols(cls, Ku, Tu, dt, mode="PID", **kwargs):
    """Construct a PID tuned by the Ziegler-Nichols ultimate-cycle rule.

    Given the ultimate gain ``Ku`` (the proportional-only gain at
    which the closed loop just sustains oscillation) and the
    corresponding ultimate period ``Tu``, the Z-N table maps to
    controller gains as::

        P:    Kp = 0.5  * Ku,                 Ki = 0,            Kd = 0
        PI:   Kp = 0.45 * Ku, Ti = Tu / 1.2 → Ki = 0.54*Ku/Tu,   Kd = 0
        PID:  Kp = 0.6  * Ku, Ti = Tu / 2.0 → Ki = 1.2 *Ku/Tu,
              Td = Tu / 8.0                 → Kd = 0.075*Ku*Tu

    The coefficients are the canonical Ziegler & Nichols (1942)
    values; see e.g. Astrom & Hagglund, *PID Controllers: Theory,
    Design, and Tuning* (1995), Table 4.1.

    Args:
        Ku: Ultimate gain (proportional-only gain at sustained
            oscillation). Must be positive.
        Tu: Ultimate period (period of the sustained oscillation
            in seconds). Must be positive.
        dt: Sampling period for the discrete PID.
        mode: One of ``"P"``, ``"PI"``, ``"PID"`` (default
            ``"PID"``). Selects which gains are non-zero.
        **kwargs: Forwarded to :class:`PIDController2DOF`.

    Returns:
        A :class:`PIDController2DOF` whose ``(Kp, Ki, Kd)`` match
        the Z-N table for the requested mode.

    Raises:
        ValueError: If ``mode`` is not one of ``"P"``, ``"PI"``,
            ``"PID"``, or if ``Ku`` / ``Tu`` are non-positive.
    """
    if mode not in cls._ZIEGLER_NICHOLS_MODES:
        raise ValueError(
            f"ziegler_nichols: mode must be one of "
            f"{cls._ZIEGLER_NICHOLS_MODES!r}; got {mode!r}"
        )
    Ku_f = float(Ku)
    Tu_f = float(Tu)
    if Ku_f <= 0.0:
        raise ValueError(
            f"ziegler_nichols: Ku must be positive; got {Ku_f}"
        )
    if Tu_f <= 0.0:
        raise ValueError(
            f"ziegler_nichols: Tu must be positive; got {Tu_f}"
        )
    if mode == "P":
        kp = 0.5 * Ku_f
        ki = 0.0
        kd = 0.0
    elif mode == "PI":
        kp = 0.45 * Ku_f
        ki = 0.54 * Ku_f / Tu_f
        kd = 0.0
    else:  # PID
        kp = 0.6 * Ku_f
        ki = 1.2 * Ku_f / Tu_f
        kd = 0.075 * Ku_f * Tu_f
    return cls(dt=dt, kp=kp, ki=ki, kd=kd, **kwargs)

PIDDiscrete

Bases: LeafSystem

Discrete-time PID controller.

This block implements a discrete-time PID controller with a first-order approximation to the integrated error and an optional derivative filter. The integrated error term is computed as:

    e_int[k+1] = e_int[k] + e[k] * dt

where e is the error signal and dt is the sampling period. The derivative term is computed in the same way as for the DerivativeDiscrete block, including filter options described there. With the running error integral e_int and current estimate of the time derivative of the error e_dot, the output is:

    u[k] = kp * e[k] + ki * e_int[k] + kd * e_dot[k]

Input ports

(0) The error signal.

Output ports

(0) The control signal computed by the PID algorithm.

Parameters:

Name	Type	Description	Default
kp		The proportional gain (scalar)	1.0
ki		The integral gain (scalar)	1.0
kd		The derivative gain (scalar)	1.0
dt		The sampling period of the block.	required
initial_state		The initial value of the running error integral. Default is 0.	0.0
enable_external_initial_state		Source for the value used for the integrator initial state. True=from inport, False=from the initial_state parameter.	False
filter_type		One of "none", "forward", "backward", or "bilinear". Determines the type of filter used to estimate the derivative of the error signal. Default is "none". See DerivativeDiscrete documentation for details.	'none'
filter_coefficient		The filter coefficient for the derivative filter. Default is 1.0. See DerivativeDiscrete documentation for details.	1.0

Source code in jaxonomy/library/dynamics.py

class PIDDiscrete(LeafSystem):
    """Discrete-time PID controller.

    This block implements a discrete-time PID controller with a first-order
    approximation to the integrated error and an optional derivative filter.
    The integrated error term is computed as:
    ```
        e_int[k+1] = e_int[k] + e[k] * dt
    ```
    where `e` is the error signal and `dt` is the sampling period.  The derivative
    term is computed in the same way as for the DerivativeDiscrete block, including
    filter options described there.  With the running error integral `e_int` and
    current estimate of the time derivative of the error `e_dot`, the output is:
    ```
        u[k] = kp * e[k] + ki * e_int[k] + kd * e_dot[k]
    ```

    Input ports:
        (0) The error signal.

    Output ports:
        (0) The control signal computed by the PID algorithm.

    Parameters:
        kp:
            The proportional gain (scalar)
        ki:
            The integral gain (scalar)
        kd:
            The derivative gain (scalar)
        dt:
            The sampling period of the block.
        initial_state:
            The initial value of the running error integral.  Default is 0.
        enable_external_initial_state:
            Source for the value used for the integrator initial state. True=from inport,
            False=from the initial_state parameter.
        filter_type:
            One of "none", "forward", "backward", or "bilinear".  Determines the type of
            filter used to estimate the derivative of the error signal.  Default is
            "none".  See DerivativeDiscrete documentation for details.
        filter_coefficient:
            The filter coefficient for the derivative filter.  Default is 1.0.  See
            DerivativeDiscrete documentation for details.
    """

    class DiscreteStateType(NamedTuple):
        integral: Array
        # Recursive filter memory for the derivative estimate
        e_prev: Array
        e_dot_prev: Array

    @parameters(
        static=["dt", "filter_type", "filter_coefficient"],
        dynamic=["kp", "ki", "kd", "initial_state"],
    )
    def __init__(
        self,
        dt,
        kp=1.0,
        ki=1.0,
        kd=1.0,
        initial_state=0.0,
        enable_external_initial_state=False,
        filter_type="none",
        filter_coefficient=1.0,
        dtype=None,
        **kwargs,
    ):
        # T-038a-followup-other-blocks: per-block dtype override; stored
        # outside the @parameters list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(**kwargs)
        self.dt = dt
        self.input_index = self.declare_input_port()

        self.enable_external_initial_state = enable_external_initial_state
        self.initial_state_index = None
        if enable_external_initial_state:
            self.initial_state_index = self.declare_input_port()

        # Declare the periodic update
        self._periodic_update_idx = self.declare_periodic_update()

        # Declare an output port for the control signal
        self.control_output = self.declare_output_port()

        # NOTE:
        # An extra output port for the derivative value is not strictly necessary,
        # but the filtered estimate could be resused elsewhere.  Also, having the
        # previous value saved in the discrete output component of state would allows
        # it to be reused in the recursive filter without recomputing it as part of
        # the update step, a minor efficiency gain.  The tradeoff is an extra event
        # that has to be handled.  This implementation uses one output event and
        # re-does the derivative calculation when a recursive filter is used, but
        # we could always do it the other way in the future.

    def initialize(
        self,
        kp,
        ki,
        kd,
        initial_state,
        filter_type,
        filter_coefficient,
        dt=None,
    ):
        # T-038a-followup-other-blocks: when an explicit per-block dtype
        # is set, cast the discrete-state seed values (integral / e_prev /
        # e_dot_prev) to that dtype so the strict ``check_types`` pass
        # does not see a mismatch between the f32 input signal and an
        # implicit-f64 ``0.0`` default.
        _zero = 0.0
        if self._dtype is not None:
            initial_state = npa.asarray(initial_state).astype(self._dtype)
            _zero = npa.asarray(0.0).astype(self._dtype)

        # Declare an internal discrete state
        self.declare_discrete_state(
            default_value=self.DiscreteStateType(
                integral=initial_state,
                e_prev=_zero,
                e_dot_prev=_zero,
            ),
            as_array=False,
        )

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=0.0,
        )

        # Determine the coefficients of the filter, if applicable
        # The filter is a pair of two-element array and the filter
        # equation is:
        # a0*y[k] + a1*y[k-1] = b0*u[k] + b1*u[k-1]
        self.filter_type = filter_type
        b, a = derivative_filter(
            N=filter_coefficient, dt=self.dt, filter_type=filter_type
        )
        if self._dtype is not None:
            # T-038a-followup-other-blocks: cast filter coefficients to
            # the per-block dtype so the derivative arithmetic runs at
            # this precision regardless of upstream/global default.
            b = npa.asarray(b).astype(self._dtype)
            a = npa.asarray(a).astype(self._dtype)
        self.filter = (b, a)

        # T-127-followup-pid-discrete-feedthrough: the output port is
        # already sample-and-hold (``period=self.dt`` causes
        # ``configure_output_port`` to register a periodic update event
        # that writes the cache, and the actual output callback just
        # reads ``state.cache[cache_index]``). Listing the input ticket
        # in ``prerequisites_of_calc`` only serves to flag the output as
        # feedthrough to the algebraic-loop detector — a spurious
        # designation, since between sample boundaries the cache is
        # constant. Dropping the input prereq matches ``UnitDelay`` and
        # un-breaks the canonical
        # ``plant → err → PIDDiscrete → Saturate → plant`` closed-loop
        # pattern (it no longer needs a hand-inserted ``UnitDelay`` to
        # silence ``AlgebraicLoopError``). Same-tick reads still work:
        # the discrete update event collects inputs because
        # ``requires_inputs`` is the default ``True``.
        self.configure_output_port(
            self.control_output,
            self._output,
            period=self.dt,
            offset=0.0,
            default_value=initial_state,
            prerequisites_of_calc=[DependencyTicket.xd],
        )

    def reset_default_values(self, **dynamic_parameters):
        # T-038a-followup-other-blocks: keep the per-block dtype contract
        # consistent across reset_default_values; otherwise the discrete
        # state and output port end up with mixed f32/f64 components.
        initial_state = dynamic_parameters["initial_state"]
        _zero = 0.0
        if self._dtype is not None:
            initial_state = npa.asarray(initial_state).astype(self._dtype)
            _zero = npa.asarray(0.0).astype(self._dtype)
        self.configure_discrete_state_default_value(
            self.DiscreteStateType(
                integral=initial_state,
                e_prev=_zero,
                e_dot_prev=_zero,
            ),
            as_array=False,
        )
        self.configure_output_port_default_value(
            self.control_output, initial_state
        )

    def _eval_derivative(self, _time, state, *inputs, **_params):
        # Filtered derivative estimate

        e = inputs[self.input_index]  # Error signal from upstream
        e_prev = state.discrete_state.e_prev
        b, a = self.filter  # IIR filter coefficients

        # If the filter is recursive we need to reuse the previous derivative
        # estimate.
        if self.filter_type != "none":
            # Filtered estimate of the time derivative
            e_dot_prev = state.discrete_state.e_dot_prev

            # New estimate of the time derivative of the error signal
            e_dot = (b[0] * e + b[1] * e_prev - a[1] * e_dot_prev) / a[0]

        else:
            # Standard finite difference approximation - no recursion
            e_dot = (b[0] * e + b[1] * e_prev) / a[0]

        return e_dot

    def _update(self, time, state, *inputs, **params):
        e = inputs[self.input_index]  # Error signal from upstream

        # Integrated error signal
        e_int = state.discrete_state.integral

        # Update the derivative estimate if needed for a recursive filter.
        if self.filter_type != "none":
            e_dot = self._eval_derivative(time, state, *inputs, **params)
        else:
            # This state entry isn't used for the finite difference estimator.
            # Can just keep the original value as a placeholder.
            e_dot = state.discrete_state.e_dot_prev

        # Update the internal state
        return self.DiscreteStateType(
            integral=e_int + e * self.dt, e_prev=e, e_dot_prev=e_dot
        )

    def _eval_control(self, e, e_int, e_dot, **params):
        # Calculate the control signal for the PID control law
        kp, ki, kd = params["kp"], params["ki"], params["kd"]
        u = kp * e + ki * e_int + kd * e_dot
        return u

    def _output(self, time, state, *inputs, **params):
        e = inputs[self.input_index]  # Error signal from upstream
        e_int = state.discrete_state.integral
        e_dot = self._eval_derivative(time, state, *inputs, **params)
        u = self._eval_control(e, e_int, e_dot, **params)
        # T-038a-followup-other-blocks: cast the control signal to the
        # per-block dtype so cross-dtype upstream connections promote
        # down to the requested precision (best-effort).
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        return u

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        u = self.eval_input(context)
        xd = context[self.system_id].discrete_state.integral
        check_state_type(
            self,
            inp_data=u,
            state_data=xd,
            error_collector=error_collector,
        )

    def initialize_static_data(self, context):
        """Set the initial state from the input port, if specified via config"""
        if self.initial_state_index is not None:
            try:
                initial_state = self.eval_input(context, self.initial_state_index)
                default_value = self.DiscreteStateType(
                    integral=initial_state,
                    e_prev=0.0,
                    e_dot_prev=0.0,
                )
                self._default_discrete_state = default_value
                local_context = context[self.system_id].with_discrete_state(
                    default_value
                )
                context = context.with_subcontext(self.system_id, local_context)

            except UpstreamEvalError:
                # The diagram has only been partially created.  Defer the
                # inference of the initial state until the upstream block has been
                # connected.
                logger.debug(
                    "PID_Discrete.initialize_static_data: UpstreamEvalError. "
                    "Continuing without default value initialization."
                )
        return super().initialize_static_data(context)

initialize_static_data(context)

Set the initial state from the input port, if specified via config

Source code in jaxonomy/library/dynamics.py

def initialize_static_data(self, context):
    """Set the initial state from the input port, if specified via config"""
    if self.initial_state_index is not None:
        try:
            initial_state = self.eval_input(context, self.initial_state_index)
            default_value = self.DiscreteStateType(
                integral=initial_state,
                e_prev=0.0,
                e_dot_prev=0.0,
            )
            self._default_discrete_state = default_value
            local_context = context[self.system_id].with_discrete_state(
                default_value
            )
            context = context.with_subcontext(self.system_id, local_context)

        except UpstreamEvalError:
            # The diagram has only been partially created.  Defer the
            # inference of the initial state until the upstream block has been
            # connected.
            logger.debug(
                "PID_Discrete.initialize_static_data: UpstreamEvalError. "
                "Continuing without default value initialization."
            )
    return super().initialize_static_data(context)

PRBS

Bases: LeafSystem

Pseudo-Random Binary Sequence (PRBS) source.

Emits +amplitude or -amplitude at each sample_time tick, drawn from a fair Bernoulli(0.5) under jax.random.bernoulli and remapped via 2*b - 1. Useful as a broad-band excitation signal for system identification.

The amplitude parameter is differentiable (scaling the binary selector); the +1 / -1 selector itself is wrapped in lax.stop_gradient.

Input ports

None.

Output ports

(0) The most recent ±amplitude sample.

Parameters:

Name	Type	Description	Default
sample_time	float	Period (s) at which a fresh bit is drawn.	required
amplitude	float	Magnitude of the binary output (differentiable).	1.0
seed	int	Integer seed for the PRNG key. If None, a 32-bit random seed is drawn from numpy.random.	None

Notes

Phase 1 uses Bernoulli(0.5) sampling rather than a true maximal-length LFSR. Period-faithful PRBS-N (n_bits register size) is deferred — see T-122-followup-lfsr.

Per-vmap-batch independence: pass fold_in_batch_index=True (T-122-followup-vmap-fold-in) to derive a per-replica independent PRNG stream via jax.lax.axis_index("batch") inside simulate_batch(use_vmap=True) / simulate_distributed. Default False preserves bit-identical phase 1 behaviour.

Source code in jaxonomy/library/sources.py

class PRBS(LeafSystem):
    """Pseudo-Random Binary Sequence (PRBS) source.

    Emits ``+amplitude`` or ``-amplitude`` at each ``sample_time``
    tick, drawn from a fair Bernoulli(0.5) under
    ``jax.random.bernoulli`` and remapped via ``2*b - 1``. Useful as
    a broad-band excitation signal for system identification.

    The ``amplitude`` parameter is differentiable (scaling the binary
    selector); the ``+1 / -1`` selector itself is wrapped in
    ``lax.stop_gradient``.

    Input ports:
        None.

    Output ports:
        (0) The most recent ``±amplitude`` sample.

    Parameters:
        sample_time: Period (s) at which a fresh bit is drawn.
        amplitude: Magnitude of the binary output (differentiable).
        seed: Integer seed for the PRNG key. If ``None``, a 32-bit
            random seed is drawn from ``numpy.random``.

    Notes:
        Phase 1 uses Bernoulli(0.5) sampling rather than a true
        maximal-length LFSR. Period-faithful PRBS-N (n_bits register
        size) is deferred — see ``T-122-followup-lfsr``.

        Per-vmap-batch independence: pass ``fold_in_batch_index=True``
        (T-122-followup-vmap-fold-in) to derive a per-replica
        independent PRNG stream via ``jax.lax.axis_index("batch")``
        inside ``simulate_batch(use_vmap=True)`` / ``simulate_distributed``.
        Default ``False`` preserves bit-identical phase 1 behaviour.
    """

    @parameters(static=["seed", "fold_in_batch_index"], dynamic=["amplitude"])
    def __init__(
        self,
        sample_time: float,
        amplitude: float = 1.0,
        seed: int = None,
        fold_in_batch_index: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self._sample_time = float(sample_time)

        self.declare_output_port(
            self._output,
            period=sample_time,
            offset=0.0,
        )
        self.declare_periodic_update(
            self._update,
            period=sample_time,
            offset=0.0,
        )

    def initialize(
        self,
        amplitude: float = 1.0,
        seed: int = None,
        fold_in_batch_index: bool = False,
    ):
        from jax import random as _jrandom
        from jax import lax as _jlax
        import jax.numpy as _jnp

        self._jrandom = _jrandom
        self._jlax = _jlax
        self._jnp = _jnp
        self._fold_in_batch_index = bool(fold_in_batch_index)

        if seed is None:
            seed = int(np.random.randint(0, 2**31 - 1, dtype=np.int64))
        key = _jrandom.PRNGKey(int(seed))
        key, subkey = _jrandom.split(key)
        bit0 = _jrandom.bernoulli(subkey, p=0.5)
        # Map {False, True} -> {-1.0, +1.0} as float.
        sel0 = _jnp.where(bit0, 1.0, -1.0)
        val0 = float(amplitude) * sel0
        default_state = _PRNGState(key=key, val=val0)
        self.declare_discrete_state(default_value=default_state, as_array=False)

    def _output(self, _time, state, *_inputs, **_parameters):
        return state.discrete_state.val

    def _update(self, _time, state, *_inputs, **parameters):
        key, subkey = self._jrandom.split(state.discrete_state.key)
        # T-122-followup-vmap-fold-in: fold per-replica batch index into
        # subkey when running under vmap(axis_name="batch") and opted-in.
        subkey = _maybe_fold_in_batch_axis(
            self._jrandom, subkey, self._fold_in_batch_index
        )
        bit = self._jrandom.bernoulli(subkey, p=0.5)
        sel = self._jlax.stop_gradient(self._jnp.where(bit, 1.0, -1.0))
        amplitude = parameters["amplitude"]
        val = amplitude * sel
        return _PRNGState(key=key, val=val)

PRBSLFSR

Bases: LeafSystem

True maximal-length PRBS-N source built on a binary LFSR.

Emits +amplitude or -amplitude at each sample_time tick, drawn from a Linear-Feedback Shift Register (LFSR) of length register_length configured with the standard primitive feedback polynomial. The output sequence has period exactly 2^N - 1 and a flat power spectrum below 1/(2N) of the sample rate, making it the canonical "white" excitation for system identification.

Reproducibility: same seed -> bit-identical sequence. Distinct seeds traverse the same cyclic orbit at different starting phases.

Differentiability: amplitude is a dynamic parameter and flows through gradients linearly; the binary selector itself is wrapped in lax.stop_gradient (the LSB extraction is non-differentiable).

Input ports

None.

Output ports

(0) The most recent ±amplitude sample.

Parameters:

Name	Type	Description	Default
sample_time	float	Period (s) at which the LFSR advances by one step.	required
amplitude	float	Magnitude of the binary output (differentiable).	1.0
register_length	int	One of {7, 9, 11, 15, 17, 23, 31}. Determines the period (2^N - 1) and the tap polynomial.	15
seed	int	Non-zero integer seeding the LFSR register state. 0 is silently promoted to 1 since the all-zero register is a fixed point of any LFSR (would emit a constant zero).	1

Notes

Per-vmap-batch independence: pass fold_in_batch_index=True (T-122-followup-vmap-fold-in) to derive a per-replica independent starting phase on the same maximal-length cycle. Inside simulate_batch(use_vmap=True) / simulate_distributed (which wrap their vmap with axis_name="batch"), the LFSR register is XOR-perturbed by a per-replica non-zero salt derived from jax.lax.axis_index("batch") on the very first update step (tracked via a phase_advanced flag in the discrete state). The salt is masked to the low N bits of the register and promoted from 0 to 1 to avoid the all-zero fixed point. Default False preserves bit-identical behaviour with the original LFSR follow-up.

Source code in jaxonomy/library/sources.py

class PRBSLFSR(LeafSystem):
    """True maximal-length PRBS-N source built on a binary LFSR.

    Emits ``+amplitude`` or ``-amplitude`` at each ``sample_time`` tick,
    drawn from a Linear-Feedback Shift Register (LFSR) of length
    ``register_length`` configured with the standard primitive
    feedback polynomial. The output sequence has period exactly
    ``2^N - 1`` and a flat power spectrum below ``1/(2N)`` of the
    sample rate, making it the canonical "white" excitation for system
    identification.

    Reproducibility: same ``seed`` -> bit-identical sequence. Distinct
    seeds traverse the same cyclic orbit at different starting phases.

    Differentiability: ``amplitude`` is a dynamic parameter and flows
    through gradients linearly; the binary selector itself is wrapped
    in ``lax.stop_gradient`` (the LSB extraction is non-differentiable).

    Input ports:
        None.

    Output ports:
        (0) The most recent ``±amplitude`` sample.

    Parameters:
        sample_time: Period (s) at which the LFSR advances by one step.
        amplitude: Magnitude of the binary output (differentiable).
        register_length: One of ``{7, 9, 11, 15, 17, 23, 31}``.
            Determines the period (``2^N - 1``) and the tap polynomial.
        seed: Non-zero integer seeding the LFSR register state. ``0``
            is silently promoted to ``1`` since the all-zero register
            is a fixed point of any LFSR (would emit a constant zero).

    Notes:
        Per-vmap-batch independence: pass ``fold_in_batch_index=True``
        (T-122-followup-vmap-fold-in) to derive a per-replica
        independent starting phase on the same maximal-length cycle.
        Inside ``simulate_batch(use_vmap=True)`` /
        ``simulate_distributed`` (which wrap their vmap with
        ``axis_name="batch"``), the LFSR register is XOR-perturbed by a
        per-replica non-zero salt derived from
        ``jax.lax.axis_index("batch")`` on the very first update step
        (tracked via a ``phase_advanced`` flag in the discrete state).
        The salt is masked to the low N bits of the register and
        promoted from 0 to 1 to avoid the all-zero fixed point.  Default
        ``False`` preserves bit-identical behaviour with the original
        LFSR follow-up.
    """

    # Standard primitive feedback polynomials (1-indexed bit positions
    # from the LSB end). Each tuple is the set of tap positions whose
    # bits are XORed to form the feedback bit.
    _TAPS = {
        7: (7, 6),
        9: (9, 5),
        11: (11, 9),
        15: (15, 14),
        17: (17, 14),
        23: (23, 18),
        31: (31, 28),
    }

    @parameters(
        static=["seed", "register_length", "fold_in_batch_index"],
        dynamic=["amplitude"],
    )
    def __init__(
        self,
        sample_time: float,
        amplitude: float = 1.0,
        register_length: int = 15,
        seed: int = 1,
        fold_in_batch_index: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if register_length not in self._TAPS:
            valid = ", ".join(str(k) for k in sorted(self._TAPS))
            raise ValueError(
                f"PRBSLFSR: unsupported register_length={register_length!r}. "
                f"Supported lengths: {valid}."
            )

        self._sample_time = float(sample_time)
        self._register_length = int(register_length)
        self._taps = self._TAPS[register_length]
        # Mask of the low N bits, used to keep the register inside
        # ``[0, 2^N)`` after every shift.
        self._reg_mask = (1 << self._register_length) - 1

        self.declare_output_port(
            self._output,
            period=sample_time,
            offset=0.0,
        )
        self.declare_periodic_update(
            self._update,
            period=sample_time,
            offset=0.0,
        )

    def initialize(
        self,
        amplitude: float = 1.0,
        register_length: int = 15,
        seed: int = 1,
        fold_in_batch_index: bool = False,
    ):
        # Lazy JAX import (consistent with the T-122 phase 1 sources).
        from jax import lax as _jlax
        from jax import random as _jrandom
        import jax.numpy as _jnp

        self._jlax = _jlax
        self._jrandom = _jrandom
        self._jnp = _jnp
        self._fold_in_batch_index = bool(fold_in_batch_index)

        # All-zero register is a fixed point of any LFSR; promote 0 -> 1
        # rather than silently emit a constant zero output.
        seed_int = int(seed) if seed is not None else 1
        seed_int = seed_int & self._reg_mask
        if seed_int == 0:
            seed_int = 1
        self._seed_int = seed_int

        reg0 = _jnp.asarray(seed_int, dtype=_jnp.uint32)
        bit0 = reg0 & _jnp.asarray(1, dtype=_jnp.uint32)
        sel0 = _jnp.where(bit0 == 1, 1.0, -1.0)
        val0 = float(amplitude) * sel0
        # ``phase_advanced=0`` means the per-replica LFSR phase shift
        # (T-122-followup-vmap-fold-in) has not yet been applied; it
        # gets flipped to ``1`` on the first ``_update`` call when
        # ``fold_in_batch_index=True``.  Always present in the state
        # tuple so the discrete-state pytree shape is independent of
        # the kwarg (keeps JIT cache keys stable).
        flag0 = _jnp.asarray(0, dtype=_jnp.uint32)
        default_state = _LFSRState(reg=reg0, val=val0, phase_advanced=flag0)
        self.declare_discrete_state(default_value=default_state, as_array=False)

    def _maybe_perturb_lfsr_register(self, reg, phase_advanced):
        """One-shot per-replica derivation of a fresh LFSR register.

        Returns ``(new_reg, new_phase_advanced)``.  No-op when
        ``fold_in_batch_index`` is ``False``, when ``phase_advanced``
        is already ``1`` (perturbation already applied on a previous
        step), or when not running under ``vmap(axis_name="batch")``
        (the unbound-axis ``NameError`` is caught at trace time).

        Applying the perturbation on the FIRST step only (rather than
        every step) preserves the LFSR's maximal-length property: each
        replica simply starts at a distinct register state on the same
        ``2^N - 1`` cycle and then evolves under the unperturbed
        feedback polynomial.

        The fresh per-replica register is derived by folding the batch
        index into a ``PRNGKey(seed)`` then taking the low N bits of a
        ``jax.random.bits`` draw -- this gives well-distributed
        register states across replicas and avoids the pathological
        degeneracies of a naive ``seed XOR axis_index`` (which can hit
        the all-zero fixed point or collide across replicas when
        ``seed`` and ``idx`` are small).  The all-zero candidate is
        promoted to ``1`` for the same reason as in ``initialize``.
        """
        if not self._fold_in_batch_index:
            return reg, phase_advanced
        import jax as _jax
        try:
            idx = _jax.lax.axis_index("batch")
        except NameError:
            return reg, phase_advanced
        # Derive a fresh per-replica register from a PRNG seeded with
        # ``seed`` and folded with the batch index.  This avoids the
        # naive-XOR pathologies (replica-0 collisions, all-zero fixed
        # point) and gives well-distributed starting states across
        # replicas.
        base_key = self._jrandom.PRNGKey(int(self._seed_int))
        per_rep_key = self._jrandom.fold_in(base_key, idx)
        # Draw 32 random bits and mask to N bits.
        rand_u32 = self._jrandom.bits(per_rep_key, shape=(), dtype=self._jnp.uint32)
        candidate = rand_u32 & self._jnp.asarray(
            self._reg_mask, dtype=self._jnp.uint32
        )
        # Promote the all-zero candidate to ``1`` -- same pattern as
        # ``initialize`` -- to avoid the LFSR's only fixed point.
        zero = self._jnp.asarray(0, dtype=self._jnp.uint32)
        one = self._jnp.asarray(1, dtype=self._jnp.uint32)
        candidate = self._jnp.where(candidate == zero, one, candidate)
        # Apply only on the first call (phase_advanced == 0).
        already = phase_advanced != zero
        new_reg = self._jnp.where(already, reg, candidate)
        new_flag = one
        return new_reg, new_flag

    def _output(self, _time, state, *_inputs, **_parameters):
        return state.discrete_state.val

    def _update(self, _time, state, *_inputs, **parameters):
        reg = state.discrete_state.reg
        phase_advanced = state.discrete_state.phase_advanced
        # T-122-followup-vmap-fold-in: one-shot per-replica phase shift
        # on the very first update step.  Bit-identical to the original
        # LFSR follow-up when ``fold_in_batch_index=False`` or when not
        # running under ``vmap(axis_name="batch")``.
        reg, phase_advanced = self._maybe_perturb_lfsr_register(
            reg, phase_advanced
        )
        one = self._jnp.asarray(1, dtype=self._jnp.uint32)
        # Unroll the XOR-over-taps chain in Python at trace time so the
        # JAX graph is a flat sequence of bit ops with no Python loop.
        feedback = (reg >> (self._taps[0] - 1)) & one
        for tap in self._taps[1:]:
            feedback = feedback ^ ((reg >> (tap - 1)) & one)
        mask = self._jnp.asarray(self._reg_mask, dtype=self._jnp.uint32)
        new_reg = ((reg << 1) | feedback) & mask
        new_bit = new_reg & one
        sel = self._jlax.stop_gradient(
            self._jnp.where(new_bit == 1, 1.0, -1.0)
        )
        amplitude = parameters["amplitude"]
        new_val = amplitude * sel
        return _LFSRState(
            reg=new_reg, val=new_val, phase_advanced=phase_advanced
        )

PolynomialChaos

Bases: LeafSystem

Polynomial-chaos surrogate y = sum_k c_k Psi_k(u) as a feedthrough block. Input port 0 is the feature vector u; the coefficients are a dynamic parameter (Xiu & Karniadakis 2002).

Source code in jaxonomy/library/rom/surrogates.py

class PolynomialChaos(LeafSystem):
    """Polynomial-chaos surrogate ``y = sum_k c_k Psi_k(u)`` as a feedthrough
    block. Input port 0 is the feature vector ``u``; the coefficients are a
    dynamic parameter (Xiu & Karniadakis 2002)."""

    def __init__(self, model: PCEModel, name=None, **kwargs):
        super().__init__(name=name, **kwargs)
        self.model = model
        self.declare_input_port()
        self.declare_dynamic_parameter("coeffs", jnp.asarray(model.coeffs))
        self._output_port_idx = self.declare_output_port(
            self._eval_output, name="y",
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )

    def _eval_output(self, time, state, *inputs, **params):
        Xstar = _row(inputs[0])
        Xi = _pce_standardize(Xstar, self.model.loc, self.model.scale)
        Psi = _pce_design(Xi, self.model.multi_indices, self.model.types,
                          self.model.order)
        return (Psi @ params["coeffs"])[0]

Power

Bases: FeedthroughBlock

Raise the input signal to a constant power.

Dispatches to jax.numpy.power, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.power.html

For input signal u with exponent p, the output will be y = u ** p.

Input ports

(0) The input signal.

Output ports

(0) The input signal raised to the power of the exponent.

Parameters:

Name	Type	Description	Default
exponent		The exponent to which the input signal is raised.	required

Source code in jaxonomy/library/math_ops.py

class Power(FeedthroughBlock):
    """Raise the input signal to a constant power.

    Dispatches to `jax.numpy.power`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.power.html

    For input signal `u` with exponent `p`, the output will be `y = u ** p`.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The input signal raised to the power of the exponent.

    Parameters:
        exponent:
            The exponent to which the input signal is raised.
    """

    @parameters(static=["exponent"])
    def __init__(self, exponent, **kwargs):
        super().__init__(self._func, **kwargs)

        # Note that the exponent here is declared as a configuration
        # parameter and not a context parameter, making it non-differentiable.
        # This is because the derivative rule for the exponent includes a log
        # of the primal input signal, which can cause NaN values during backprop
        # if the input signal is non-positive. Specifically, for `y = u ** p`, the
        # linearization with respect to `p` is `dy = y * log(u) * dp`. If we
        # eventually want to support backprop through this block, we will need
        # to handle the log of the input signal in a way that avoids NaN values.
        # (e.g. with gradient clipping). Tracked in WC-306
        self.exponent = exponent

    def initialize(self, exponent):
        self.exponent = exponent

    def _func(self, *inputs, **parameters):
        (u,) = inputs
        return u**self.exponent

Prelookup

Bases: LeafSystem

Compute the (bucket_index, fraction) pair for a query against a precomputed grid.

This is the upstream half of the standard Prelookup/InterpolationUsingPrelookup pair. Pair with one or more :class:InterpolationUsingPrelookup blocks downstream -- each can interpolate a DIFFERENT output table that shares the same grid axis without re-running the bucket search.

The marketing wedge: when N downstream tables share one query axis (e.g. 10 lookup maps on a common engine-RPM input), this saves N - 1 binary searches per evaluation.

Input ports

(0) -- the query coordinate (scalar).

Output ports

(0) -- a NamedTuple with fields (index, fraction) ready to plug into one or more :class:InterpolationUsingPrelookup blocks.

Parameters:

Name	Type	Description	Default
input_array		1-D, strictly-increasing grid of breakpoints. Stored verbatim for the bucket search.	required
dtype	optional	If set (e.g. jnp.float32), the grid array is cast to this dtype on construction. Mirrors the per-block dtype contract of :class:LookupTable1d.	None
extrapolation	(optional, T - 114 - fu - prelookup - extrap)	Out-of-range policy for the alpha blend weight. One of "clip" (default; alpha clamped to [0, 1] so OOB queries map to the nearest endpoint), "linear" (alpha left raw -- the downstream blend extends the boundary slope linearly), or "nan" (alpha set to NaN on OOB queries -- the downstream blend propagates NaN). All three modes match the corresponding :class:LookupTable1d extrapolation policies byte-for-byte. Any paired :class:InterpolationUsingPrelookup should declare the SAME mode for API agreement (the math is owned by Prelookup's alpha computation).	'clip'

Notes

Differentiable through the query coordinate via fraction (the discrete index is piecewise-constant).

Source code in jaxonomy/library/tables.py

class Prelookup(LeafSystem):
    """Compute the (bucket_index, fraction) pair for a query against a
    precomputed grid.

    This is the upstream half of the standard
    ``Prelookup``/``InterpolationUsingPrelookup`` pair.  Pair with one or
    more :class:`InterpolationUsingPrelookup` blocks downstream -- each
    can interpolate a DIFFERENT output table that shares the same grid
    axis without re-running the bucket search.

    The marketing wedge: when N downstream tables share one query axis
    (e.g. 10 lookup maps on a common engine-RPM input), this saves
    ``N - 1`` binary searches per evaluation.

    Input ports:
        ``(0)`` -- the query coordinate (scalar).

    Output ports:
        ``(0)`` -- a NamedTuple with fields ``(index, fraction)`` ready
        to plug into one or more :class:`InterpolationUsingPrelookup`
        blocks.

    Parameters:
        input_array:
            1-D, strictly-increasing grid of breakpoints.  Stored
            verbatim for the bucket search.
        dtype (optional):
            If set (e.g. ``jnp.float32``), the grid array is cast to this
            dtype on construction.  Mirrors the per-block dtype contract
            of :class:`LookupTable1d`.
        extrapolation (optional, T-114-fu-prelookup-extrap):
            Out-of-range policy for the ``alpha`` blend weight.  One of
            ``"clip"`` (default; alpha clamped to ``[0, 1]`` so OOB
            queries map to the nearest endpoint), ``"linear"`` (alpha
            left raw -- the downstream blend extends the boundary slope
            linearly), or ``"nan"`` (alpha set to NaN on OOB queries --
            the downstream blend propagates NaN).  All three modes
            match the corresponding :class:`LookupTable1d`
            extrapolation policies byte-for-byte.  Any paired
            :class:`InterpolationUsingPrelookup` should declare the
            SAME mode for API agreement (the math is owned by
            ``Prelookup``'s alpha computation).

    Notes:
        Differentiable through the query coordinate via ``fraction``
        (the discrete ``index`` is piecewise-constant).
    """

    def __init__(self, input_array, dtype=None, extrapolation="clip", **kwargs):
        # Per-block dtype + active precision policy fallback, matching
        # the LookupTable1d contract.  Stored outside the @parameters
        # list so the kwarg does not round-trip through model JSON or
        # get JAX-traced.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype

        # T-114-fu-prelookup-extrap -- validate the extrapolation kwarg
        # eagerly so a bad value raises ValueError on construction
        # rather than at trace time.  Stored outside the @parameters
        # list (Python-string kwarg, never JAX-traced).
        if extrapolation not in ("clip", "linear", "nan"):
            raise ValueError(
                f"Prelookup: extrapolation must be one of "
                f"('clip','linear','nan'), got {extrapolation!r}"
            )
        self._extrapolation = extrapolation

        # Eagerly validate the grid up front so the failure mode is a
        # clear ValueError on construction, not a cryptic shape error
        # at trace time.
        _input_np = np.asarray(input_array)
        if _input_np.ndim != 1:
            raise ValueError(
                f"Prelookup: input_array must be 1-D, got shape "
                f"{_input_np.shape}"
            )
        if _input_np.size < 2:
            raise ValueError(
                f"Prelookup: input_array must have at least 2 entries, "
                f"got shape {_input_np.shape}"
            )
        if not np.all(np.diff(_input_np) > 0):
            raise ValueError(
                f"Prelookup: input_array must be strictly monotonically "
                f"increasing, got {list(_input_np)}"
            )

        if self._dtype is not None:
            self._input_array = npa.asarray(_input_np).astype(self._dtype)
        else:
            self._input_array = npa.array(_input_np)

        super().__init__(**kwargs)
        self.declare_input_port()

        # Capture the grid in a local so the closure does not pull
        # ``self`` into the JAX trace.
        xp_local = self._input_array
        n_local = int(self._input_array.shape[0])
        extrap_local = self._extrapolation

        def _compute(_time, _state, *inputs, **_params):
            (x_query,) = inputs
            # Clip bucket index to [0, n - 2] so the i + 1 neighbour
            # downstream stays in range.  Same convention as the
            # ``interp_1d`` / ``interp_nd`` backends.
            i = npa.clip(
                npa.searchsorted(xp_local, x_query, side="right") - 1,
                0,
                n_local - 2,
            )
            x0 = xp_local[i]
            x1 = xp_local[i + 1]
            alpha = (x_query - x0) / (x1 - x0)
            # T-114-fu-prelookup-extrap -- apply the OOB policy to
            # alpha.  The downstream blend (1 - alpha) * yp[i] + alpha
            # * yp[i+1] then naturally produces:
            #   * "clip"   -- nearest endpoint (alpha in [0, 1]);
            #   * "linear" -- boundary-slope extension (alpha raw);
            #   * "nan"    -- NaN propagation (alpha is NaN past edges).
            if extrap_local == "clip":
                alpha = npa.clip(alpha, 0.0, 1.0)
            elif extrap_local == "nan":
                # Out-of-range mask uses the ORIGINAL grid endpoints
                # ``xp[0]``/``xp[-1]`` (not the clipped bucket), so a
                # query exactly on a breakpoint is in-range.  ``npa``
                # falls back to ``jnp`` for ``where``/``isnan`` under a
                # JAX trace, so this is differentiable around the
                # finite (non-NaN) branch.
                oob = (x_query < xp_local[0]) | (x_query > xp_local[-1])
                nan_val = npa.asarray(npa.nan, dtype=alpha.dtype)
                alpha = npa.where(oob, nan_val, alpha)
            # else "linear": alpha left raw -- with i clamped to
            # [0, n-2], an OOB query gets alpha < 0 (left) or alpha > 1
            # (right), and the downstream blend extends the boundary
            # slope linearly.  No-op on this branch.
            return _PrelookupResult(index=i, fraction=alpha)

        # NOTE: deliberately NOT passing ``default_value=`` here -- the
        # framework's ``declare_output_port`` would call
        # ``npa.array(default_value)`` and flatten our NamedTuple into
        # a plain 1-D array, losing the tuple type.  Letting the
        # framework lazily compute the default by invoking ``_compute``
        # on a dummy context yields the correct NamedTuple-typed
        # default.  Same trick as :class:`BusCreator`.
        self.declare_output_port(
            _compute,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )

    @property
    def input_array(self):
        """The breakpoint array used for the bucket search."""
        return self._input_array

    @property
    def extrapolation(self):
        """The OOB policy applied to ``alpha`` (``"clip"``/``"linear"``/``"nan"``)."""
        return self._extrapolation

extrapolation property

The OOB policy applied to alpha ("clip"/"linear"/"nan").

input_array property

The breakpoint array used for the bucket search.

PrelookupInverse

Bases: LeafSystem

Compute the (bucket_index, fraction) pair for an INVERSE-direction lookup against a strictly-monotonic value array.

Forward :class:Prelookup answers: given x, find (i, alpha) s.t. xp[i] + alpha * (xp[i+1] - xp[i]) ≈ x.

Inverse :class:PrelookupInverse answers: given y, find (i, alpha) s.t. yp[i] + alpha * (yp[i+1] - yp[i]) ≈ y.

The output is the same NamedTuple-typed :class:_PrelookupResult produced by :class:Prelookup, so it plugs straight into one or more :class:InterpolationUsingPrelookup blocks connected to OTHER tables -- typically the inverse table that maps i back to the recovered x (for example the breakpoints of the forward table).

Marketing wedge: implicit equations y = f(x) where f is a monotonic 1-D table -- gain scheduling, sensor calibration, etc.

Input ports

(0) -- the query coordinate in the OUTPUT space (y).

Output ports

(0) -- a :class:_PrelookupResult NamedTuple (index, fraction) ready to plug into an :class:InterpolationUsingPrelookup block.

Parameters:

Name	Type	Description	Default
output_array		1-D, strictly-monotonic (increasing OR decreasing) array of values to invert. Stored verbatim for the bucket search; decreasing tables are handled by reversing the search direction.	required
dtype	optional	If set, the value array is cast to this dtype on construction. Mirrors the :class:Prelookup / :class:LookupTable1d dtype contract.	None
extrapolation	optional	Only "clip" is supported in this followup -- queries outside the monotone range collapse to the nearest endpoint. "linear"/"nan" are filed as T-114-followup-prelookup-inverse-extrap because the OOB definition on a value-axis interacts with the direction-flipping logic non-trivially.	'clip'

Notes

Differentiable through the query coordinate and through output_array. Non-monotonic output_array raises ValueError at construction time.

Source code in jaxonomy/library/tables.py

class PrelookupInverse(LeafSystem):
    """Compute the (bucket_index, fraction) pair for an INVERSE-direction
    lookup against a strictly-monotonic value array.

    Forward :class:`Prelookup` answers: given ``x``, find ``(i, alpha)``
    s.t. ``xp[i] + alpha * (xp[i+1] - xp[i]) ≈ x``.

    Inverse :class:`PrelookupInverse` answers: given ``y``, find
    ``(i, alpha)`` s.t. ``yp[i] + alpha * (yp[i+1] - yp[i]) ≈ y``.

    The output is the same NamedTuple-typed
    :class:`_PrelookupResult` produced by :class:`Prelookup`, so it plugs
    straight into one or more :class:`InterpolationUsingPrelookup` blocks
    connected to OTHER tables -- typically the inverse table that maps
    ``i`` back to the recovered ``x`` (for example the breakpoints of the
    forward table).

    Marketing wedge: implicit equations ``y = f(x)`` where ``f`` is a
    monotonic 1-D table -- gain scheduling, sensor calibration, etc.

    Input ports:
        ``(0)`` -- the query coordinate in the OUTPUT space (``y``).

    Output ports:
        ``(0)`` -- a :class:`_PrelookupResult` NamedTuple
        ``(index, fraction)`` ready to plug into an
        :class:`InterpolationUsingPrelookup` block.

    Parameters:
        output_array:
            1-D, strictly-monotonic (increasing OR decreasing) array of
            values to invert.  Stored verbatim for the bucket search;
            decreasing tables are handled by reversing the search
            direction.
        dtype (optional):
            If set, the value array is cast to this dtype on
            construction.  Mirrors the :class:`Prelookup` /
            :class:`LookupTable1d` dtype contract.
        extrapolation (optional):
            Only ``"clip"`` is supported in this followup -- queries
            outside the monotone range collapse to the nearest endpoint.
            ``"linear"``/``"nan"`` are filed as
            ``T-114-followup-prelookup-inverse-extrap`` because the OOB
            definition on a value-axis interacts with the
            direction-flipping logic non-trivially.

    Notes:
        Differentiable through the query coordinate and through
        ``output_array``.  Non-monotonic ``output_array`` raises
        ``ValueError`` at construction time.
    """

    def __init__(
        self, output_array, dtype=None, extrapolation="clip", **kwargs
    ):
        # Per-block dtype + active precision policy fallback, mirroring
        # Prelookup / LookupTable1d.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype

        # Honest fallback: only "clip" ships in this followup.  Reject
        # "linear"/"nan" with a clear error pointing at the deeper
        # followup ticket so callers know it is on the roadmap.
        if extrapolation != "clip":
            if extrapolation in ("linear", "nan"):
                raise NotImplementedError(
                    f"PrelookupInverse: extrapolation={extrapolation!r} is "
                    f"a deeper followup (T-114-followup-prelookup-inverse-"
                    f"extrap); only 'clip' is shipped today."
                )
            raise ValueError(
                f"PrelookupInverse: extrapolation must be 'clip' "
                f"(other modes are deferred), got {extrapolation!r}"
            )
        self._extrapolation = extrapolation

        _out_np = np.asarray(output_array)
        if _out_np.ndim != 1:
            raise ValueError(
                f"PrelookupInverse: output_array must be 1-D, got shape "
                f"{_out_np.shape}"
            )
        if _out_np.size < 2:
            raise ValueError(
                f"PrelookupInverse: output_array must have at least 2 "
                f"entries, got shape {_out_np.shape}"
            )

        # Monotonicity check.  Strictly increasing OR strictly decreasing
        # is fine; anything else is ambiguous to invert.
        diffs = np.diff(_out_np)
        if np.all(diffs > 0):
            self._direction = "increasing"
        elif np.all(diffs < 0):
            self._direction = "decreasing"
        else:
            raise ValueError(
                f"PrelookupInverse: output_array must be strictly "
                f"monotonic (increasing or decreasing) to invert; got "
                f"{list(_out_np)}"
            )

        if self._dtype is not None:
            self._output_array = npa.asarray(_out_np).astype(self._dtype)
        else:
            self._output_array = npa.array(_out_np)

        super().__init__(**kwargs)
        self.declare_input_port()

        # Capture the table + direction in locals so the closure does
        # not pull ``self`` into the JAX trace.
        yp_local = self._output_array
        n_local = int(self._output_array.shape[0])
        direction_local = self._direction

        def _compute(_time, _state, *inputs, **_params):
            (y_query,) = inputs
            if direction_local == "increasing":
                # Standard searchsorted on the value-axis.
                i = npa.clip(
                    npa.searchsorted(yp_local, y_query, side="right") - 1,
                    0,
                    n_local - 2,
                )
                y0 = yp_local[i]
                y1 = yp_local[i + 1]
            else:
                # Decreasing case: searchsorted needs an ascending array
                # to work, so search the reversed array and re-map the
                # index back into the original orientation.  After
                # remap, ``i`` still indexes into yp_local with the
                # invariant that yp_local[i] >= y_query >= yp_local[i+1]
                # (modulo the boundary clip).
                yp_rev = yp_local[::-1]
                j = npa.clip(
                    npa.searchsorted(yp_rev, y_query, side="right") - 1,
                    0,
                    n_local - 2,
                )
                # Reversed bucket j corresponds to original bucket
                # (n-2) - j.
                i = (n_local - 2) - j
                y0 = yp_local[i]
                y1 = yp_local[i + 1]
            alpha = (y_query - y0) / (y1 - y0)
            # Only "clip" ships in this followup.  Collapses OOB queries
            # to the nearest endpoint.  ``alpha`` outside [0, 1] can
            # still arise when the query is exactly at a boundary
            # due to floating-point; the clip pins it.
            alpha = npa.clip(alpha, 0.0, 1.0)
            return _PrelookupResult(index=i, fraction=alpha)

        # Same NamedTuple-output-port trick as Prelookup: do NOT pass
        # ``default_value=`` (the framework would flatten the
        # NamedTuple); let it lazily compute via ``_compute``.
        self.declare_output_port(
            _compute,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )

    @property
    def output_array(self):
        """The 1-D monotonic value array being inverted."""
        return self._output_array

    @property
    def direction(self):
        """``"increasing"`` or ``"decreasing"`` -- monotonicity sense."""
        return self._direction

    @property
    def extrapolation(self):
        """The OOB policy (always ``"clip"`` in this followup)."""
        return self._extrapolation

direction property

"increasing" or "decreasing" -- monotonicity sense.

extrapolation property

The OOB policy (always "clip" in this followup).

output_array property

The 1-D monotonic value array being inverted.

Product

Bases: ReduceBlock

Compute the product and/or quotient of the input signals.

The block will multiply or divide the input signals, depending on the specified operators. For example, if the block has three inputs u1, u2, and u3 and is configured with operators="**/", then the output signal will be y = u1 * u2 / u3. By default, the block will multiply all of the input signals.

Input ports

(0..n_in-1) The input signals.

Output ports

(0) The product and/or quotient of the input signals.

Parameters:

Name	Type	Description	Default
n_in		The number of input ports.	required
operators		A string of length n_in specifying the operators to apply to each of the input signals. Each character in the string must be either "" or "/". The default is "".	None
denominator_limit		Currently unsupported	None
divide_by_zero_behavior		Currently unsupported	None

Source code in jaxonomy/library/math_ops.py

class Product(ReduceBlock):
    """Compute the product and/or quotient of the input signals.

    The block will multiply or divide the input signals, depending on the specified
    operators.  For example, if the block has three inputs `u1`, `u2`, and `u3` and
    is configured with operators="**/", then the output signal will be
    `y = u1 * u2 / u3`.  By default, the block will multiply all of the input signals.

    Input ports:
        (0..n_in-1) The input signals.

    Output ports:
        (0) The product and/or quotient of the input signals.

    Parameters:
        n_in:
            The number of input ports.
        operators:
            A string of length `n_in` specifying the operators to apply to each of
            the input signals.  Each character in the string must be either "*" or "/".
            The default is "*".
        denominator_limit:
            Currently unsupported
        divide_by_zero_behavior:
            Currently unsupported
    """

    @parameters(static=["operators", "denominator_limit", "divide_by_zero_behavior"])
    def __init__(
        self,
        n_in,
        operators=None,  # Expect "**/*", etc
        denominator_limit=None,
        divide_by_zero_behavior=None,
        **kwargs,
    ):
        super().__init__(n_in, None, **kwargs)

    def initialize(
        self,
        operators=None,  # Expect "**/*", etc
        denominator_limit=None,
        divide_by_zero_behavior=None,
    ):
        if operators is not None and any(char not in {"*", "/"} for char in operators):
            raise BlockParameterError(
                message=f"Product block {self.name} has invalid operators {operators}. Can only contain '*' and '/'",
                system=self,
                parameter_name="operators",
            )

        if operators is not None and "/" in operators:
            num_indices = npa.array(
                [idx for idx, op in enumerate(operators) if op == "*"]
            )
            den_indices = npa.array(
                [idx for idx, op in enumerate(operators) if op == "/"]
            )

            def _func(inputs):
                ain = npa.array(inputs)
                num = npa.take(ain, num_indices, axis=0)
                den = npa.take(ain, den_indices, axis=0)
                return npa.prod(num, axis=0) / npa.prod(den, axis=0)

        else:

            def _func(inputs):
                return npa.prod(npa.array(inputs), axis=0)

        self.replace_op(_func)

ProductOfElements

Bases: FeedthroughBlock

Compute the product of the elements of the input signal.

Dispatches to jax.numpy.prod, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.prod.html

Input ports

(0) The input signal.

Output ports

(0) The product of the elements of the input signal.

Source code in jaxonomy/library/math_ops.py

class ProductOfElements(FeedthroughBlock):
    """Compute the product of the elements of the input signal.

    Dispatches to `jax.numpy.prod`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.prod.html

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The product of the elements of the input signal.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(npa.prod, *args, **kwargs)

Pulse

Bases: SourceBlock

A periodic pulse signal.

Given amplitude a, pulse width w, and period p, the output signal is:

    y(t) = a if t % p < w else 0

where % is the modulo operator.

Input ports

None

Output ports

(0) The pulse signal.

Parameters:

Name	Type	Description	Default
amplitude		The amplitude of the pulse signal.	1.0
pulse_width		The fraction of the period during which the pulse is "high".	0.5
period		The period of the pulse signal.	1.0
phase_delay		Currently unsupported.	0.0

Source code in jaxonomy/library/sources.py

class Pulse(SourceBlock):
    """A periodic pulse signal.

    Given amplitude `a`, pulse width `w`, and period `p`, the output signal is:
    ```
        y(t) = a if t % p < w else 0
    ```
    where `%` is the modulo operator.

    Input ports:
        None

    Output ports:
        (0) The pulse signal.

    Parameters:
        amplitude:
            The amplitude of the pulse signal.
        pulse_width:
            The fraction of the period during which the pulse is "high".
        period:
            The period of the pulse signal.
        phase_delay:
            Currently unsupported.
    """

    @parameters(dynamic=["amplitude", "pulse_width", "period", "phase_delay"])
    def __init__(
        self, amplitude=1.0, pulse_width=0.5, period=1.0, phase_delay=0.0, **kwargs
    ):
        super().__init__(self._func, **kwargs)

        # Initialize the floating-point tolerance.  This will be machine epsilon
        # for the floating point type of the time variable (determined in the
        # static initialization step).
        self.eps = 0.0

        if abs(phase_delay) > 1e-9:
            warnings.warn("Warning. Pulse block phase_delay not implemented.")

        # Add a dummy event so that the ODE solver doesn't try to integrate through
        # the discontinuity.
        # ad 2 events, one for the up jump, and one the down jump
        self.declare_discrete_state(default_value=False)
        self._dummy_periodic_update_idx = self.declare_periodic_update()
        self._periodic_update_idx = self.declare_periodic_update()

    def initialize(self, amplitude, pulse_width, period, phase_delay):
        if abs(phase_delay) > 1e-9:
            warnings.warn("Warning. Pulse block phase_delay not implemented.")

        self.configure_periodic_update(
            self._dummy_periodic_update_idx,
            lambda *args, **kwargs: True,
            period=period,
            offset=period,
        )

        self.configure_periodic_update(
            self._periodic_update_idx,
            lambda *args, **kwargs: True,
            period=period,
            offset=period + period * pulse_width,
        )

    def _func(self, time, **parameters):
        # Add a floating-point tolerance to the modulo operation to avoid
        # accuracy issues when the time is an "exact" multiple of the period.
        period_fraction = (
            npa.remainder(time + self.eps, parameters["period"]) / parameters["period"]
        )
        return npa.where(
            period_fraction >= parameters["pulse_width"],
            0.0,
            parameters["amplitude"],
        )

    def initialize_static_data(self, context):
        # Determine machine epsilon for the type of the time variable
        self.eps = 2 * npa.finfo(npa.result_type(context.time)).eps
        return super().initialize_static_data(context)

PyTorch

Bases: LeafSystem

Block to perform inference with a pre-trained PyTorch model saved as TorchScript.

The input to the block should be of compatible type and shape expected by the TorchScript. For example, if the TorchScript model expects a torch.float32 tensor of shape (3, 224, 224), the input to the block should be a jax.numpy array of shape (3, 224, 224) of dtype jnp.float32.

For output types, if no casting is specified through the cast_outputs_to_dtype parameter, the output of the block will have the same dtype as the TorchScript model output, but expressed as jax.numpy types. For example. if the TorchScript model outputs a torch.float32 tensor, the output of the block will be a jax.numpy array of dtype jnp.float32.

If casting is specified through cast_outputs_to_dtype parameter, all the outputs, of the block will be casted to this specific jax.numpy dtype.

.. note:: float32 models under jaxonomy's global x64. import jaxonomy enables JAX 64-bit mode (jax_enable_x64) for the whole process, so upstream signals are float64 by default. A TorchScript model traced in torch.float32 therefore receives float64 inputs (a dtype error or a silent arithmetic change) unless you cast at the block boundary. One-line idiom: pass cast_outputs_to_dtype="float32" and feed the block x.astype(jnp.float32) inputs.

Input ports

(i) The ith input to the model.

Output ports

(j) The jth output of the model.

Parameters:

Name	Type	Description	Default
file_name	str	Path to the model Torchscript .pt file.	required
num_inputs	int	The number of inputs to the model. Only required for TorchScript models.	1
num_outputs	int	The number of outputs of the model.	1
cast_outputs_to_dtype	str	The dtype to cast all the outputs of the block to. Must correspond to a jax.numpy datatype. For example, "float32", "float64", "int32", "int64".	None
add_batch_dim_to_inputs	bool	Whether to add a new first dimension to the inputs before evaluating the TorchScript or TensorFlow model. This is useful when the model expects a batch dimension.	False

Source code in jaxonomy/library/predictor.py

class PyTorch(LeafSystem):
    """
    Block to perform inference with a pre-trained PyTorch model saved as TorchScript.

    The input to the block should be of compatible type and shape expected by
    the TorchScript. For example, if the TorchScript model expects
    a `torch.float32` tensor of shape `(3, 224, 224)`, the input to the block should be
    a `jax.numpy` array of shape (3, 224, 224) of dtype `jnp.float32`.

    For output types, if no casting is specified through the `cast_outputs_to_dtype`
    parameter, the output of the block will have the same dtype as the TorchScript
    model output, but expressed as `jax.numpy` types. For example. if the
    TorchScript model outputs a `torch.float32` tensor, the output of the block will be
    a `jax.numpy` array of dtype `jnp.float32`.

    If casting is specified through `cast_outputs_to_dtype` parameter, all the outputs,
    of the block will be casted to this specific `jax.numpy` dtype.

    .. note:: **float32 models under jaxonomy's global x64.**
       ``import jaxonomy`` enables JAX 64-bit mode (``jax_enable_x64``)
       for the whole process, so upstream signals are float64 by
       default.  A TorchScript model traced in ``torch.float32``
       therefore receives float64 inputs (a dtype error or a silent
       arithmetic change) unless you cast at the block boundary.
       One-line idiom: pass ``cast_outputs_to_dtype="float32"`` and
       feed the block ``x.astype(jnp.float32)`` inputs.

    Input ports:
        (i) The ith input to the model.

    Output ports:
        (j) The jth output of the model.

    Parameters:
        file_name (str):
            Path to the model Torchscript `.pt` file.

        num_inputs (int):
            The number of inputs to the model. Only required for TorchScript models.

        num_outputs (int):
            The number of outputs of the model.

        cast_outputs_to_dtype (str):
            The dtype to cast all the outputs of the block to. Must correspond to a
            `jax.numpy` datatype. For example, "float32", "float64", "int32", "int64".

        add_batch_dim_to_inputs (bool):
            Whether to add a new first dimension to the inputs before evaluating the
            TorchScript or TensorFlow model. This is useful when the model expects a
            batch dimension.
    """

    EXTRA_FILES = {}

    @parameters(
        static=[
            "file_name",
            "num_inputs",
            "num_outputs",
            "cast_outputs_to_dtype",
            "add_batch_dim_to_inputs",
        ]
    )
    def __init__(
        self,
        file_name,
        num_inputs=1,
        num_outputs=1,
        cast_outputs_to_dtype=None,
        add_batch_dim_to_inputs=False,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

        self._num_inputs = num_inputs
        self._num_outputs = num_outputs

        for _ in range(num_inputs):
            self.declare_input_port()

        def _make_output_callback(output_index):
            def _output_callback(time, state, *inputs, **params):
                outputs = self._evaluate_output(time, state, *inputs, **params)
                return outputs[output_index]

            return _output_callback

        for output_index in range(num_outputs):
            self.declare_output_port(
                _make_output_callback(output_index),
                requires_inputs=True,
            )

    def initialize(
        self,
        file_name,
        num_inputs=1,
        num_outputs=1,
        cast_outputs_to_dtype=None,
        add_batch_dim_to_inputs=False,
    ):
        if num_inputs != self._num_inputs:
            raise ValueError("num_inputs can't be changed after initialization")
        if num_outputs != self._num_outputs:
            raise ValueError("num_outputs can't be changed after initialization")

        self.dtype_output = (
            getattr(jnp, cast_outputs_to_dtype)
            if cast_outputs_to_dtype is not None
            else None
        )

        self.add_batch_dim_to_inputs = add_batch_dim_to_inputs
        self.model = torch.jit.load(file_name)
        self.model.eval()

    def initialize_static_data(self, context):
        """Infer the output shapes and dtypes of the ML model."""
        # If building as part of a subsystem, this may not be fully connected yet.
        # That's fine, as long as it is connected by root context creation time.
        # This probably isn't a good long-term solution:
        #   see https://jaxonomy.atlassian.net/browse/WC-51
        try:
            inputs = self.collect_inputs(context)
            outputs_jax = self._pure_callback(*inputs)

            self.pure_callback_result_type = [
                jax.ShapeDtypeStruct(x.shape, x.dtype) for x in outputs_jax
            ]
        except UpstreamEvalError:
            logger.debug(
                "PyTorch.initialize_static_data: UpstreamEvalError. "
                "Continuing without default value initialization."
            )
        return super().initialize_static_data(context)

    def _evaluate_output(self, time, state, *inputs, **params):
        return jax.pure_callback(
            self._pure_callback,
            self.pure_callback_result_type,
            *inputs,
        )

    def _pure_callback(self, *inputs):
        inputs_casted = [torch.tensor(np.array(item)) for item in inputs]

        if self.add_batch_dim_to_inputs:
            inputs_casted = [x.unsqueeze(0) for x in inputs_casted]
        outputs = self.model(*inputs_casted)

        if not isinstance(outputs, tuple):
            outputs = (outputs,)

        outputs_jax = (
            [jnp.array(x, self.dtype_output) for x in outputs]
            if self.dtype_output is not None
            else [jnp.array(x) for x in outputs]
        )
        return outputs_jax

initialize_static_data(context)

Infer the output shapes and dtypes of the ML model.

Source code in jaxonomy/library/predictor.py

def initialize_static_data(self, context):
    """Infer the output shapes and dtypes of the ML model."""
    # If building as part of a subsystem, this may not be fully connected yet.
    # That's fine, as long as it is connected by root context creation time.
    # This probably isn't a good long-term solution:
    #   see https://jaxonomy.atlassian.net/browse/WC-51
    try:
        inputs = self.collect_inputs(context)
        outputs_jax = self._pure_callback(*inputs)

        self.pure_callback_result_type = [
            jax.ShapeDtypeStruct(x.shape, x.dtype) for x in outputs_jax
        ]
    except UpstreamEvalError:
        logger.debug(
            "PyTorch.initialize_static_data: UpstreamEvalError. "
            "Continuing without default value initialization."
        )
    return super().initialize_static_data(context)

QuadraticCost

Bases: ReduceBlock

LQR-type quadratic cost function for a state and input.

Computes the cost as x'Qx + u'Ru, where Q and R are the cost matrices. In order to compute a running cost, combine this with an Integrator or IntegratorDiscrete block.

Source code in jaxonomy/library/costs_and_losses.py

class QuadraticCost(ReduceBlock):
    """LQR-type quadratic cost function for a state and input.

    Computes the cost as x'Qx + u'Ru, where Q and R are the cost matrices.
    In order to compute a running cost, combine this with an `Integrator`
    or `IntegratorDiscrete` block.
    """

    def __init__(self, Q, R, name=None):
        super().__init__(2, self._cost, name=name)
        self.Q = Q
        self.R = R

    def _cost(self, inputs):
        x, u = inputs
        J = jnp.dot(x, jnp.dot(self.Q, x)) + jnp.dot(u, jnp.dot(self.R, u))
        return J.squeeze()

QuanserHAL

Bases: LeafSystem

Hardware Abstraction Layer for Quanser hardware.

This block provides an interface to virtual or physical Quanser hardware. It requires that the Quanser hardware or QLabs simulator be properly configured and that the Quanser python library is available on the system path. See the Quanser documentation for more information.

To use an idealized model of the Qube Servo hardware, see the jaxonomy.library.QubeServoModel block, which may be run without hardware or in the cloud-based simulation UI.

Input ports

(0) Control signal to the motor in volts

Output ports

(0) The observed rotor and pendulum angles in radians

Parameters:

Name	Type	Description	Default
dt		The time step of the simulation.	required
version		The version of the Qube hardware (2 or 3). By default, version 2 is used with hardware=False, or version 3 is used with hardware=True.	2
hardware		If True, connect to the physical hardware. If False, connect to the QLabs simulator.	False
name		The name of the system in the Jaxonomy model.	'QuanserHAL'

Source code in jaxonomy/library/quanser.py

class QuanserHAL(LeafSystem):
    """Hardware Abstraction Layer for Quanser hardware.

    This block provides an interface to virtual or physical Quanser hardware.
    It requires that the Quanser hardware or QLabs simulator be properly configured
    and that the Quanser python library is available on the system path.  See the
    Quanser documentation for more information.

    To use an idealized model of the Qube Servo hardware, see the
    `jaxonomy.library.QubeServoModel` block, which may be run without hardware
    or in the cloud-based simulation UI.

    Input ports:
        (0) Control signal to the motor in volts

    Output ports:
        (0) The observed rotor and pendulum angles in radians

    Parameters:
        dt: The time step of the simulation.
        version: The version of the Qube hardware (2 or 3).  By default, version 2 is
            used with `hardware=False`, or version 3 is used with `hardware=True`.
        hardware:
            If True, connect to the physical hardware. If False, connect to the
            QLabs simulator.
        name: The name of the system in the Jaxonomy model.

    """

    def __init__(self, dt, version=2, hardware=False, name="QuanserHAL", ui_id=None):
        super().__init__(name=name, ui_id=ui_id)

        if version is None:
            version = 2 if not hardware else 3

        # Init QLabs
        # HACK for macOS - Quanser's code only works on Windows
        # Insert asyncio.windows_events fake module
        if sys.platform != "win32":
            _windows_events = types.ModuleType("windows_events")
            _windows_events.INFINITE = np.iinfo(np.uint32).max
            sys.modules["asyncio.windows_events"] = _windows_events

        try:
            from pal.products.qube import QubeServo2, QubeServo3
        except Exception:
            raise ImportError(
                "Could not import QubeServo2 or QubeServo3 from pal.products.qube. "
                "Check that the hardware drivers are available on the system path."
            )

        if version not in [2, 3]:
            raise ValueError("version must be 2 or 3")

        if version == 2:
            QubeClass = QubeServo2
        else:
            QubeClass = QubeServo3

        self.qube = QubeClass(hardware=hardware, pendulum=1, frequency=1 / dt)
        self._setup_siginthandler()

        print("Initialized Qube")
        if self.qube.card is None:
            raise RuntimeError(
                "Could not find hardware. Try power-cycling and check connections."
            )
        self.qube.write_led(color=[0, 1, 0])

        self.declare_input_port()  # Inputs are the control signals to the motor

        # Periodically send the control signals to the motor
        self.declare_periodic_update(
            self.step,
            period=dt,
            offset=0.0,
        )

        # Periodically read the sensor outputs
        self.declare_output_port(
            self.output,
            period=dt,
            offset=0.0,
            requires_inputs=False,
        )

    def __exit__(self, exc_type, exc_value, traceback):
        self.terminate()
        super().__exit__(exc_type, exc_value, traceback)

    def _setup_siginthandler(self):
        self.prev_sigint_handler = signal.signal(signal.SIGINT, self._interrupt)

    def _restore_siginthandler(self):
        print("Restoring sigint handler")
        if self.prev_sigint_handler is not None:
            signal.signal(signal.SIGINT, self.prev_sigint_handler)
            self.prev_sigint_handler = None

    # This custom handler makes it possible to Interrupt the jupyter kernel
    # and still connect again to the Qube environment.
    def _interrupt(self, signum, frame):
        prev_handler = self.prev_sigint_handler
        self.terminate()
        if prev_handler is not None:
            prev_handler(signum, frame)

    def terminate(self):
        self._restore_siginthandler()
        if self.qube is not None:
            self.qube.write_led(color=[1, 1, 0])
            self.qube.terminate()
            self.qube = None

    def _impure_step(self, voltage):
        # Write the voltage to the Qube
        self.qube.write_voltage(voltage)

    def step(self, time, state, *inputs, **parameters):
        return io_callback(self._impure_step, None, *inputs)

    def _impure_output(self):
        # Read the sensor outputs
        self.qube.read_outputs()
        theta, alpha = self.qube.motorPosition, self.qube.pendulumPosition
        return jnp.array([theta, alpha])

    def output(self, time, state, *inputs, **parameters):
        return io_callback(self._impure_output, jnp.zeros(2))

Quantizer

Bases: FeedthroughBlock

Discritize the input signal into a set of discrete values.

Given an input signal u and a resolution interval, this block quantizes the input signal onto the integer multiples of interval. The output signal is y = interval * f(u / interval) where f is selected by mode:

• "round" (default): round-half-to-even (banker's rounding, IEEE-754 default). Byte-equivalent with the phase-1 implementation (which used npa.round unconditionally).
• "floor": round toward -inf (truncation in many DSP impls).
• "ceil": round toward +inf.
• "trunc": round toward zero (chops the fractional part regardless of sign).

Quantization is non-differentiable: the output is piecewise-constant with measure-zero jumps. The block wraps the rounded result in :func:jax.lax.stop_gradient (JAX backend only) so JAX always sees a zero gradient through the block. This both matches the underlying mathematical reality and prevents spurious gradient leakage if any backend ever provides a smoothed surrogate for round/floor/ ceil/trunc. Under the numpy backend the helper is the identity, preserving dtype/value byte-equivalence.

Input ports

(0) The continuous input signal. In most cases, should be scaled to the range [0, interval].

Output ports

(0) The quantized output signal, on the same scale as the input signal.

Parameters:

Name	Type	Description	Default
interval		The quantization step size — output values are integer multiples of interval.	required
mode		One of "round", "floor", "ceil", "trunc". Default "round".	'round'

Source code in jaxonomy/library/nonlinearities.py

class Quantizer(FeedthroughBlock):
    """Discritize the input signal into a set of discrete values.

    Given an input signal ``u`` and a resolution ``interval``, this block
    quantizes the input signal onto the integer multiples of ``interval``.
    The output signal is ``y = interval * f(u / interval)`` where ``f`` is
    selected by ``mode``:

    - ``"round"`` (default): round-half-to-even (banker's rounding,
      IEEE-754 default). Byte-equivalent with the phase-1 implementation
      (which used ``npa.round`` unconditionally).
    - ``"floor"``: round toward -inf (truncation in many DSP impls).
    - ``"ceil"``: round toward +inf.
    - ``"trunc"``: round toward zero (chops the fractional part regardless
      of sign).

    Quantization is non-differentiable: the output is piecewise-constant
    with measure-zero jumps. The block wraps the rounded result in
    :func:`jax.lax.stop_gradient` (JAX backend only) so JAX always sees a
    zero gradient through the block. This both matches the underlying
    mathematical reality and prevents spurious gradient leakage if any
    backend ever provides a smoothed surrogate for ``round``/``floor``/
    ``ceil``/``trunc``. Under the numpy backend the helper is the
    identity, preserving dtype/value byte-equivalence.

    Input ports:
        (0) The continuous input signal. In most cases, should be scaled to the range
            ``[0, interval]``.

    Output ports:
        (0) The quantized output signal, on the same scale as the input signal.

    Parameters:
        interval:
            The quantization step size — output values are integer
            multiples of ``interval``.
        mode:
            One of ``"round"``, ``"floor"``, ``"ceil"``, ``"trunc"``.
            Default ``"round"``.
    """

    @parameters(dynamic=["interval"])
    def __init__(self, interval, mode="round", *args, **kwargs):
        if mode not in _QUANTIZER_MODES:
            raise BlockParameterError(
                message=(
                    f"Quantizer mode must be one of {_QUANTIZER_MODES}, "
                    f"got {mode!r}."
                ),
            )
        self._mode = mode
        if mode == "round":
            _round_fn = npa.round
        elif mode == "floor":
            _round_fn = npa.floor
        elif mode == "ceil":
            _round_fn = npa.ceil
        else:  # mode == "trunc"
            _round_fn = npa.trunc

        def _op(x, interval):
            return _stop_gradient(interval * _round_fn(x / interval))

        super().__init__(_op, *args, **kwargs)

    def initialize(self, interval):
        pass

QubeServoModel

Bases: LeafSystem

Plant model for the Quanser Qube Servo Furuta Pendulum.

The Quanser Qube Servo is a pendulum controlled by a rotary arm. The rotary arm is actuated by a DC motor. The pendulum is free to rotate about the rotary arm.

The state of the system is given by the rotor angle (theta), pendulum angle (alpha), rotor angular velocity, and pendulum angular velocity. The input to the system is the voltage applied to the motor, which is converted to torque by a simple linear model.

Input ports

(0) The motor voltage signal

Output ports

(0) If full_state_output is False, the rotor angle and pendulum angle. Otherwise, will return the entire continuous state vector.

Parameters:

Name	Type	Description	Default
x0		Initial state of the system [theta, alpha, theta_dot, alpha_dot]	[0.0, 0.0, 0.0, 0.0]
Rm		Motor resistance (Ohms)	8.4
km		Back-emf constant (V-s/rad)	0.042
mr		Rotary arm mass (kg)	0.095
Lr		Rotor arm length (m)	0.085
br		Rotor arm damping coefficient (N-m-s/rad)	0.0005
mp		Pendulum mass (kg)	0.024
Lp		Pendulum arm length (m)	0.129
bp		Pendulum damping coefficient (N-m-s/rad)	2.5e-05
g		Gravitational constant (m/s^2)	9.81
kr		Feedback control to send the rotor back to zero	0.0
full_state_output		If True, output the full state vector. Otherwise, only output the rotor and pendulum angles.	False

Source code in jaxonomy/library/quanser.py

class QubeServoModel(LeafSystem):
    """Plant model for the Quanser Qube Servo Furuta Pendulum.

    The Quanser Qube Servo is a pendulum controlled by a rotary arm. The
    rotary arm is actuated by a DC motor. The pendulum is free to rotate about
    the rotary arm.

    The state of the system is given by the rotor angle (theta), pendulum angle
    (alpha), rotor angular velocity, and pendulum angular velocity. The input to
    the system is the voltage applied to the motor, which is converted to torque
    by a simple linear model.

    Input ports:
        (0) The motor voltage signal

    Output ports:
        (0) If `full_state_output` is False, the rotor angle and pendulum angle.
            Otherwise, will return the entire continuous state vector.

    Parameters:
        x0: Initial state of the system [theta, alpha, theta_dot, alpha_dot]
        Rm: Motor resistance (Ohms)
        km: Back-emf constant (V-s/rad)
        mr: Rotary arm mass (kg)
        Lr: Rotor arm length (m)
        br: Rotor arm damping coefficient (N-m-s/rad)
        mp: Pendulum mass (kg)
        Lp: Pendulum arm length (m)
        bp: Pendulum damping coefficient (N-m-s/rad)
        g: Gravitational constant (m/s^2)
        kr: Feedback control to send the rotor back to zero
        full_state_output: If True, output the full state vector. Otherwise,
            only output the rotor and pendulum angles.

    """

    def __init__(
        self,
        x0=[0.0, 0.0, 0.0, 0.0],
        Rm=8.4,
        km=0.042,
        mr=0.095,
        Lr=0.085,
        br=5e-4,
        mp=0.024,
        Lp=0.129,
        bp=2.5e-5,
        g=9.81,
        kr=0.0,
        full_state_output=False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.declare_dynamic_parameter("Rm", Rm)
        self.declare_dynamic_parameter("km", km)
        self.declare_dynamic_parameter("mr", mr)
        self.declare_dynamic_parameter("Lr", Lr)
        self.declare_dynamic_parameter("br", br)
        self.declare_dynamic_parameter("mp", mp)
        self.declare_dynamic_parameter("Lp", Lp)
        self.declare_dynamic_parameter("bp", bp)
        self.declare_dynamic_parameter("g", g)
        self.declare_dynamic_parameter("kr", kr)

        # One input: motor voltage
        self.declare_input_port()

        self.declare_continuous_state(default_value=x0, ode=self._ode)

        if full_state_output:
            self.declare_continuous_state_output()

        else:
            # Only measure (alpha, theta)
            def _obs_callback(time, state, *inputs, **parameters):
                return state.continuous_state[:2]

            self.declare_output_port(_obs_callback, requires_inputs=False)

    def _ode(self, time, state, *inputs, **parameters):
        # Unpack state
        q, dq = state.continuous_state[:2], state.continuous_state[2:]
        theta, alpha = q  # Rotor angle, pendulum angle
        theta_dot, alpha_dot = dq

        # Unpack parameters
        Rm = parameters["Rm"]
        km = parameters["km"]
        mr = parameters["mr"]
        Lr = parameters["Lr"]
        br = parameters["br"]
        mp = parameters["mp"]
        Lp = parameters["Lp"]
        bp = parameters["bp"]
        kr = parameters["kr"]
        g = parameters["g"]

        lp = Lp / 2  # Pendulum center of mass

        # Moment of inertia of the rotor arm about the motor
        Jr = mr * Lr**2 / 3

        # Moment of inertia of the pendulum about the pivot point
        Jp = mp * Lp**2 / 3

        # Unpack inputs
        (u,) = inputs
        u = npa.atleast_1d(u)

        # Feedback control to send the rotor back to zero
        u -= kr * theta

        # Mass matrix
        M = npa.array(
            [
                [Jr + Jp * npa.sin(alpha) ** 2, -mp * lp * Lr * npa.cos(alpha)],
                [-mp * lp * Lr * npa.cos(alpha), Jp],
            ]
        )

        # Coriolis matrix
        C = npa.array(
            [
                [
                    Jp * npa.sin(2 * alpha) * alpha_dot + br + 0 * km**2 / Rm,
                    mp * lp * Lr * npa.sin(alpha) * alpha_dot,
                ],
                [-0.5 * Jp * npa.sin(2 * alpha) * theta_dot, bp],
            ]
        )

        # Gravity vector
        tau_g = npa.array([0, mp * g * lp * npa.sin(alpha)])

        # Input matrix
        B = npa.array(
            [
                [km / Rm],
                [0],
            ]
        )

        # State space representation
        ddq = npa.linalg.solve(M, B @ u - (C @ dq + tau_g))
        return npa.concatenate([dq, ddq])

RBFModel

Fitted radial-basis-function interpolant with optional polynomial tail.

s(x) = sum_i w_i phi(||x - c_i||) + sum_k c_k p_k(x) (Hardy 1971; Wendland 2005).

Source code in jaxonomy/library/rom/surrogates.py

class RBFModel:
    """Fitted radial-basis-function interpolant with optional polynomial tail.

    ``s(x) = sum_i w_i phi(||x - c_i||) + sum_k c_k p_k(x)`` (Hardy 1971;
    Wendland 2005).
    """

    def __init__(self, centers, weights, poly_coeffs, poly_indices, kernel,
                 epsilon):
        self.centers = centers
        self.weights = weights
        self.poly_coeffs = poly_coeffs      # None if no tail
        self.poly_indices = poly_indices    # None if no tail
        self.kernel = kernel
        self.epsilon = float(epsilon)

    def predict(self, Xstar):
        """Interpolant value at ``Xstar`` (jax-traceable)."""
        Xstar = _as2d(Xstar)
        d2 = _sqdist(Xstar, self.centers)
        y = _rbf_phi(d2, self.kernel, self.epsilon) @ self.weights
        if self.poly_indices is not None:
            y = y + _rbf_monomials(Xstar, self.poly_indices) @ self.poly_coeffs
        return y

predict(Xstar)

Interpolant value at Xstar (jax-traceable).

Source code in jaxonomy/library/rom/surrogates.py

def predict(self, Xstar):
    """Interpolant value at ``Xstar`` (jax-traceable)."""
    Xstar = _as2d(Xstar)
    d2 = _sqdist(Xstar, self.centers)
    y = _rbf_phi(d2, self.kernel, self.epsilon) @ self.weights
    if self.poly_indices is not None:
        y = y + _rbf_monomials(Xstar, self.poly_indices) @ self.poly_coeffs
    return y

RadialBasisSurrogate

Bases: LeafSystem

RBF surrogate y = sum_i w_i phi(||u - c_i||) (+ poly tail) as a feedthrough block. Input port 0 is the feature vector u; the RBF weights (and polynomial-tail coefficients) are dynamic parameters (Hardy 1971).

Source code in jaxonomy/library/rom/surrogates.py

class RadialBasisSurrogate(LeafSystem):
    """RBF surrogate ``y = sum_i w_i phi(||u - c_i||) (+ poly tail)`` as a
    feedthrough block. Input port 0 is the feature vector ``u``; the RBF weights
    (and polynomial-tail coefficients) are dynamic parameters (Hardy 1971)."""

    def __init__(self, model: RBFModel, name=None, **kwargs):
        super().__init__(name=name, **kwargs)
        self.model = model
        self.declare_input_port()
        self.declare_dynamic_parameter("weights", jnp.asarray(model.weights))
        if model.poly_indices is not None:
            self.declare_dynamic_parameter(
                "poly_coeffs", jnp.asarray(model.poly_coeffs))
        self._output_port_idx = self.declare_output_port(
            self._eval_output, name="y",
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
        )

    def _eval_output(self, time, state, *inputs, **params):
        Xstar = _row(inputs[0])
        d2 = _sqdist(Xstar, self.model.centers)
        y = _rbf_phi(d2, self.model.kernel, self.model.epsilon) @ params["weights"]
        if self.model.poly_indices is not None:
            y = y + (_rbf_monomials(Xstar, self.model.poly_indices)
                     @ params["poly_coeffs"])
        return y[0]

Ramp

Bases: SourceBlock

Output a linear ramp signal in time.

Given a slope m, a start value y0, and a start time t0, the output signal is:

    y(t) = m * (t - t0) + y0 if t >= t0 else y0

where t is the current simulation time.

Input ports

None

Output ports

(0) The ramp signal.

Parameters:

Name	Type	Description	Default
start_value		The value of the output signal at the start time.	0.0
slope		The slope of the ramp signal.	1.0
start_time		The time at which the ramp signal begins.	1.0
units	(optional, T - 104 - followup - units - on - source - blocks)	If set, the output port advertises this :class:Unit. The connect-time consistency check (T-104) then enforces downstream ports declare a compatible unit. Default None keeps the legacy "no-units" behaviour (byte-equivalent to pre-T-104 diagrams).	None

Source code in jaxonomy/library/sources.py

class Ramp(SourceBlock):
    """Output a linear ramp signal in time.

    Given a slope `m`, a start value `y0`, and a start time `t0`, the output signal is:
    ```
        y(t) = m * (t - t0) + y0 if t >= t0 else y0
    ```
    where `t` is the current simulation time.

    Input ports:
        None

    Output ports:
        (0) The ramp signal.

    Parameters:
        start_value:
            The value of the output signal at the start time.
        slope:
            The slope of the ramp signal.
        start_time:
            The time at which the ramp signal begins.
        units (optional, T-104-followup-units-on-source-blocks):
            If set, the output port advertises this :class:`Unit`. The
            connect-time consistency check (T-104) then enforces
            downstream ports declare a compatible unit. Default
            ``None`` keeps the legacy "no-units" behaviour
            (byte-equivalent to pre-T-104 diagrams).
    """

    @parameters(dynamic=["start_value", "slope", "start_time"])
    def __init__(
        self,
        start_value=0.0,
        slope=1.0,
        start_time=1.0,
        units=None,
        **kwargs,
    ):
        super().__init__(self._func, **kwargs)
        # T-104-followup-units-on-source-blocks: see Sine for rationale.
        self.output_ports[self._output_port_idx].units = units

    def initialize(self, start_value, slope, start_time):
        pass

    def _func(self, time, **parameters):
        m = parameters["slope"]
        t0 = parameters["start_time"]
        y0 = parameters["start_value"]
        return npa.where(time >= t0, m * (time - t0) + y0, y0)

RandomNumber

Bases: LeafSystem

Discrete-time random number generator.

Generates independent, identically distributed random numbers at each time step. Dispatches to jax.random for the actual random number generation.

Supported distributions include "ball", "cauchy", "choice", "dirichlet", "exponential", "gamma", "lognormal", "maxwell", "normal", "orthogonal", "poisson", "randint", "truncated_normal", and "uniform".

See https://jax.readthedocs.io/en/latest/jax.random.html#random-samplers for a full list of available distributions and associated parameters.

Although the JAX random number generator is a deterministic function of the key, this block maintains the key as part of the discrete state, making it a stateful RNG. The block can be seeded for reproducibility by passing an integer seed; if None, a random seed will be generated using numpy.random.

Note that this block should typically not be used as a source of randomness for continuous-time systems, as it generates a discrete-time signal. For continuous systems, use a continuous-time noise source, such as WhiteNoise.

Input ports

None

Output ports

(0) The most recently generated random number.

Parameters:

Name	Type	Description	Default
dt	float	The rate at which random numbers are generated.	required
distribution	str	The name of the random distribution to sample from.	'normal'
seed	int	An integer seed for the random number generator. If None, a random 32-bit seed will be generated.	None
dtype	DTypeLike	data type of the random number. If None, the default data type for the specified distribution will be used. Not all distributions support all data types; check the JAX documentation for details.	None
distribution_parameters		A dictionary of additional parameters to pass to the distribution function.	{}

Source code in jaxonomy/library/random.py

class RandomNumber(LeafSystem):
    """Discrete-time random number generator.

    Generates independent, identically distributed random numbers at each time step.
    Dispatches to `jax.random` for the actual random number generation.

    Supported distributions include "ball", "cauchy", "choice", "dirichlet",
    "exponential", "gamma", "lognormal", "maxwell", "normal", "orthogonal",
    "poisson", "randint", "truncated_normal", and "uniform".

    See https://jax.readthedocs.io/en/latest/jax.random.html#random-samplers
    for a full list of available distributions and associated parameters.

    Although the JAX random number generator is a deterministic function of the
    key, this block maintains the key as part of the discrete state, making it a
    stateful RNG.  The block can be seeded for reproducibility by passing an integer
    seed; if None, a random seed will be generated using numpy.random.

    Note that this block should typically not be used as a source of randomness for
    continuous-time systems, as it generates a discrete-time signal. For continuous
    systems, use a continuous-time noise source, such as `WhiteNoise`.

    Input ports:
        None

    Output ports:
        (0) The most recently generated random number.

    Parameters:
        dt: The rate at which random numbers are generated.
        distribution: The name of the random distribution to sample from.
        seed: An integer seed for the random number generator. If None, a random 32-bit
            seed will be generated.
        dtype: data type of the random number.  If None, the default data type for the
            specified distribution will be used.  Not all distributions support all
            data types; check the JAX documentation for details.
        distribution_parameters: A dictionary of additional parameters to pass to the
            distribution function.
    """

    class RNGState(NamedTuple):
        key: Array
        val: Array

    @classmethod
    def with_key(cls, key: "jax.Array", **kwargs) -> "RandomNumber":
        """
        Construct RandomNumber with an explicit JAX PRNGKey.

        Use this when you need independent noise streams in 
        batched (jax.vmap) simulations.

        Example:
            keys = jax.random.split(jax.random.PRNGKey(0), 16)

            # Each diagram gets a different key
            diagrams = [
                build_diagram_with(
                    RandomNumber.with_key(keys[i], ...)
                )
                for i in range(16)
            ]

            # OR with with_parameters (preferred):
            diagram.with_parameters({"noise.key": keys[i]})

        Args:
            key: JAX PRNGKey array (shape (2,) for default RNG)
            **kwargs: other constructor arguments
        """
        instance = cls(**kwargs)
        instance._explicit_key = key
        return instance

    @parameters(static=["distribution", "seed", "shape"])
    def __init__(
        self,
        dt: float,
        distribution: str = "normal",  # UI only exposes 'normal' for now
        seed: int = None,
        dtype: DTypeLike = None,
        shape: ShapeLike = (),
        name: str = None,
        ui_id: str = None,
        **distribution_parameters,
    ):
        super().__init__(name=name, ui_id=ui_id)

        # Declare config parameters for serialization
        self.declare_static_parameters(**distribution_parameters)

        # Add to the data type if specified.  Since not all distributions
        # support this parameter (though most do), we don't want to do this
        # unconditionally.
        if dtype is not None:
            distribution_parameters["dtype"] = dtype

        self.declare_output_port(
            self._output,
            period=dt,
            offset=0.0,
        )

        self.declare_periodic_update(
            self._update,
            period=dt,
            offset=0.0,
        )

    def initialize(
        self,
        distribution: str = "normal",  # UI only exposes 'normal' for now
        seed: int = None,
        shape: ShapeLike = (),
        **distribution_parameters,
    ):
        # Supposedly all distributions support the shape parameter
        if shape is not None and shape != ():
            distribution_parameters["shape"] = shape

        self.rng = partial(getattr(random, distribution), **distribution_parameters)

        if hasattr(self, "_explicit_key"):
            key = self._explicit_key
        else:
            key = random.PRNGKey(
                np.random.randint(0, 2**32, dtype=np.int64) if seed is None else seed
            )

        # The discrete state is a tuple of (key, val) pairs.  Because of the way that
        # JAX maintains RNG state, we need to keep track of the key as well as the
        # most recently generated value.
        key, subkey = random.split(key)
        default_state = self.RNGState(
            key=key,
            val=self.rng(subkey),  # Initial random number with the right data type
        )
        self.declare_discrete_state(default_value=default_state, as_array=False)

    def _output(self, _time, state, *_inputs, **_parameters):
        return state.discrete_state.val

    def _update(self, _time, state, *_inputs, **_parameters):
        key, subkey = random.split(state.discrete_state.key)
        return self.RNGState(
            key=key,
            val=self.rng(subkey),
        )

with_key(key, **kwargs) classmethod

Construct RandomNumber with an explicit JAX PRNGKey.

Use this when you need independent noise streams in batched (jax.vmap) simulations.

Example

keys = jax.random.split(jax.random.PRNGKey(0), 16)

Each diagram gets a different key

diagrams = [ build_diagram_with( RandomNumber.with_key(keys[i], ...) ) for i in range(16) ]

OR with with_parameters (preferred):

diagram.with_parameters({"noise.key": keys[i]})

Parameters:

Name	Type	Description	Default
key	'jax.Array'	JAX PRNGKey array (shape (2,) for default RNG)	required
**kwargs		other constructor arguments	{}

Source code in jaxonomy/library/random.py

@classmethod
def with_key(cls, key: "jax.Array", **kwargs) -> "RandomNumber":
    """
    Construct RandomNumber with an explicit JAX PRNGKey.

    Use this when you need independent noise streams in 
    batched (jax.vmap) simulations.

    Example:
        keys = jax.random.split(jax.random.PRNGKey(0), 16)

        # Each diagram gets a different key
        diagrams = [
            build_diagram_with(
                RandomNumber.with_key(keys[i], ...)
            )
            for i in range(16)
        ]

        # OR with with_parameters (preferred):
        diagram.with_parameters({"noise.key": keys[i]})

    Args:
        key: JAX PRNGKey array (shape (2,) for default RNG)
        **kwargs: other constructor arguments
    """
    instance = cls(**kwargs)
    instance._explicit_key = key
    return instance

RandomSource

Bases: LeafSystem

Multi-distribution discrete-time random source.

Unified rebuild of the single-distribution UniformRandomNumber pattern from T-122 phase 1: one block, four distributions, selected at construction by a string flag plus a params dict.

Supported distributions::

distribution="uniform"     params={"low":  ..., "high": ...}
distribution="normal"      params={"mean": ..., "std":  ...}
distribution="lognormal"   params={"mu":   ..., "sigma": ...}
distribution="triangular"  params={"low":  ..., "peak": ...,
                                   "high": ...}
distribution="exponential" params={"rate": ...}
distribution="poisson"     params={"rate": ...}   # integer-typed output
distribution="bernoulli"   params={"p":    ...}   # integer-typed 0/1 output
distribution="beta"        params={"alpha": ..., "beta":  ...}
distribution="gamma"       params={"shape": ..., "scale": ...}
distribution="weibull"     params={"shape": ..., "scale": ...}
distribution="pareto"      params={"scale": ..., "alpha": ...}

"exponential" is differentiable through rate via the standard inverse-CDF reparameterisation x = -log(1-u) / rate; "poisson" is the discrete count distribution and is not differentiable through rate w.r.t. its samples (the per-sample grad is zero by construction — the sampler is wrapped in stop_gradient). See T-122-followup-poisson.

Same seed -> bit-identical sequence (determinism contract). Under simulate_batch(use_vmap=True) / simulate_distributed, pass fold_in_batch_index=True (T-122-followup-vmap-fold-in) to derive a per-replica independent stream from the same master seed via jax.lax.axis_index("batch"). Default False preserves bit-identical behaviour with the original distributions follow-up.

All named params flow through smooth, differentiable reparameterisations of an underlying Uniform[0,1) or N(0,1) draw; the random draw itself is wrapped in lax.stop_gradient so gradients of downstream losses flow cleanly through the distribution parameters but never attempt to differentiate the PRNG key.

Input ports

None.

Output ports

(0) The most recent sample.

Parameters:

Name	Type	Description	Default
sample_time	float	Period (s) at which a fresh sample is drawn.	required
distribution	str	One of "uniform", "normal", "lognormal", "triangular", "exponential", "poisson".	'uniform'
params	dict	Dict of distribution parameters (see above for keys). Each value is registered as a dynamic parameter and is differentiable / vmap-mappable.	None
seed	int	Integer seed for the PRNG key. If None, a 32-bit random seed is drawn from numpy.random.	None
shape		Output shape. Default () (scalar).	()

Notes

Honest-fallback note: the spec contemplated a lax.switch over distributions to share one block class. Because each distribution has different params keys (and triangular has three), a runtime switch would require padding/aligning the param tuples per distribution, which is clunky and defeats the whole point of named params. Instead we dispatch at Python time on the static distribution flag -- different distributions trace into different compute graphs, which is exactly the JAX-idiomatic path for static-flag polymorphism.

Source code in jaxonomy/library/sources.py

class RandomSource(LeafSystem):
    """Multi-distribution discrete-time random source.

    Unified rebuild of the single-distribution ``UniformRandomNumber``
    pattern from T-122 phase 1: one block, four distributions, selected
    at construction by a string flag plus a ``params`` dict.

    Supported distributions::

        distribution="uniform"     params={"low":  ..., "high": ...}
        distribution="normal"      params={"mean": ..., "std":  ...}
        distribution="lognormal"   params={"mu":   ..., "sigma": ...}
        distribution="triangular"  params={"low":  ..., "peak": ...,
                                           "high": ...}
        distribution="exponential" params={"rate": ...}
        distribution="poisson"     params={"rate": ...}   # integer-typed output
        distribution="bernoulli"   params={"p":    ...}   # integer-typed 0/1 output
        distribution="beta"        params={"alpha": ..., "beta":  ...}
        distribution="gamma"       params={"shape": ..., "scale": ...}
        distribution="weibull"     params={"shape": ..., "scale": ...}
        distribution="pareto"      params={"scale": ..., "alpha": ...}

    ``"exponential"`` is differentiable through ``rate`` via the
    standard inverse-CDF reparameterisation ``x = -log(1-u) / rate``;
    ``"poisson"`` is the discrete count distribution and is *not*
    differentiable through ``rate`` w.r.t. its samples (the per-sample
    grad is zero by construction — the sampler is wrapped in
    ``stop_gradient``).  See T-122-followup-poisson.

    Same seed -> bit-identical sequence (determinism contract).  Under
    ``simulate_batch(use_vmap=True)`` / ``simulate_distributed``, pass
    ``fold_in_batch_index=True`` (T-122-followup-vmap-fold-in) to derive
    a per-replica independent stream from the same master seed via
    ``jax.lax.axis_index("batch")``.  Default ``False`` preserves
    bit-identical behaviour with the original distributions follow-up.

    All named ``params`` flow through smooth, differentiable
    reparameterisations of an underlying ``Uniform[0,1)`` or ``N(0,1)``
    draw; the random draw itself is wrapped in ``lax.stop_gradient`` so
    gradients of downstream losses flow cleanly through the
    distribution parameters but never attempt to differentiate the PRNG
    key.

    Input ports:
        None.

    Output ports:
        (0) The most recent sample.

    Parameters:
        sample_time: Period (s) at which a fresh sample is drawn.
        distribution: One of ``"uniform"``, ``"normal"``, ``"lognormal"``,
            ``"triangular"``, ``"exponential"``, ``"poisson"``.
        params: Dict of distribution parameters (see above for keys).
            Each value is registered as a *dynamic* parameter and is
            differentiable / vmap-mappable.
        seed: Integer seed for the PRNG key.  If ``None``, a 32-bit
            random seed is drawn from ``numpy.random``.
        shape: Output shape.  Default ``()`` (scalar).

    Notes:
        Honest-fallback note: the spec contemplated a ``lax.switch``
        over distributions to share one block class.  Because each
        distribution has different ``params`` keys (and ``triangular``
        has three), a runtime switch would require padding/aligning
        the param tuples per distribution, which is clunky and
        defeats the whole point of named params.  Instead we dispatch
        at *Python time* on the static ``distribution`` flag --
        different distributions trace into different compute graphs,
        which is exactly the JAX-idiomatic path for static-flag
        polymorphism.
    """

    @parameters(static=["distribution", "seed", "shape", "fold_in_batch_index"])
    def __init__(
        self,
        sample_time: float,
        distribution: str = "uniform",
        params: dict = None,
        seed: int = None,
        shape=(),
        fold_in_batch_index: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if params is None:
            params = {}
        # Fail fast at construction so the user gets a clear error
        # before the simulator tries to compile a broken graph.
        _validate_random_source_distribution(distribution, params)

        self._sample_time = float(sample_time)
        self._distribution = distribution
        # Capture for ``initialize`` (which receives the static
        # parameters but not the dynamic ``params`` dict directly).
        self._param_keys = tuple(_RANDOM_SOURCE_DISTRIBUTIONS[distribution])

        # Build the per-key defaults.  Scalar-typed params get coerced
        # to float for parity with the T-122 phase 1 single-distribution
        # blocks; array-typed params (categorical ``values`` / ``probs``)
        # are converted to numpy arrays so they survive the ``initialize``
        # round-trip with stable shape / dtype.
        defaults = {}
        for key in self._param_keys:
            if (distribution, key) in _RANDOM_SOURCE_ARRAY_PARAMS:
                defaults[key] = np.asarray(params[key])
            else:
                defaults[key] = float(params[key])
        # T-122-followup-categorical: normalise ``probs`` to sum to 1 so
        # the static-table / dynamic-param pair agrees with the
        # ``Categorical`` distribution semantics in jaxonomy.uq.
        if distribution == "categorical":
            probs = np.asarray(defaults["probs"], dtype=np.float64)
            if probs.ndim != 1:
                raise ValueError(
                    f"RandomSource(categorical): probs must be 1D; got shape {probs.shape}."
                )
            values_arr = np.asarray(defaults["values"])
            if values_arr.shape[0] != probs.shape[0]:
                raise ValueError(
                    f"RandomSource(categorical): values length {values_arr.shape[0]} "
                    f"must match probs length {probs.shape[0]}."
                )
            if np.any(probs < 0.0) or probs.sum() <= 0.0:
                raise ValueError(
                    f"RandomSource(categorical): probs must be non-negative with "
                    f"positive sum; got {probs}."
                )
            defaults["probs"] = probs / probs.sum()
            defaults["values"] = values_arr
        # T-122-followup-bernoulli: validate ``p`` is a probability in
        # [0, 1].  The downstream sampler is the categorical sampler with
        # values=[0, 1], so a value of ``p`` outside [0, 1] would silently
        # produce a malformed distribution.
        if distribution == "bernoulli":
            p_val = defaults["p"]
            if not (0.0 <= p_val <= 1.0):
                raise ValueError(
                    f"RandomSource(bernoulli): p ({p_val}) must be in [0, 1]."
                )
        self._param_defaults = defaults

        # Register each distribution parameter as a *dynamic* parameter
        # so that gradients/vmap flow through them (cf. the T-122 phase 1
        # ``UniformRandomNumber.low/high`` pattern), and so they live in
        # the simulation context rather than baked into the closure.
        # Exception: keys flagged in ``_RANDOM_SOURCE_STATIC_PARAMS`` are
        # captured at Python time (e.g. categorical ``values``) — the
        # gather-by-index into a heterogeneous-dtype table cannot
        # safely flow as a single dynamic-parameter array.
        for key in self._param_keys:
            if (distribution, key) in _RANDOM_SOURCE_STATIC_PARAMS:
                continue
            internal_name = _internal_param_name(distribution, key)
            self.declare_dynamic_parameter(internal_name, self._param_defaults[key])

        self.declare_output_port(
            self._output,
            period=sample_time,
            offset=0.0,
        )
        self.declare_periodic_update(
            self._update,
            period=sample_time,
            offset=0.0,
        )

    def initialize(
        self,
        distribution: str = "uniform",
        seed: int = None,
        shape=(),
        fold_in_batch_index: bool = False,
        **_dynamic_params,
    ):
        # ``_dynamic_params`` swallows the per-distribution keys (low,
        # high, mean, std, mu, sigma, peak) that the framework passes
        # to ``initialize`` because they were declared via
        # ``declare_dynamic_parameter``.  The defaults captured at
        # construction (``self._param_defaults``) drive the initial
        # discrete-state sample, so we don't need to read them here.

        # Lazy JAX import — mirrors the UniformRandomNumber / PRBS /
        # BandLimitedNoise convention so the framework's numpy-only
        # backend stays importable.
        from jax import random as _jrandom
        from jax import lax as _jlax
        import jax.numpy as _jnp

        self._jrandom = _jrandom
        self._jlax = _jlax
        self._jnp = _jnp
        self._shape = tuple(int(s) for s in shape) if shape else ()
        self._fold_in_batch_index = bool(fold_in_batch_index)

        # T-122-followup-categorical: convert the static ``values`` table
        # to a JAX array once here so per-update gathers are zero-cost.
        # ``values`` is captured at construction (Python-time table); it
        # is not a dynamic context parameter.
        if self._distribution == "categorical":
            self._values_table = _jnp.asarray(self._param_defaults["values"])
        else:
            self._values_table = None

        if seed is None:
            seed = int(np.random.randint(0, 2**31 - 1, dtype=np.int64))
        key = _jrandom.PRNGKey(int(seed))

        # Build the initial sample using the *default* parameter values
        # captured at construction.  This keeps the discrete-state
        # pytree shape stable across periodic updates regardless of
        # later context-time parameter overrides.
        key, subkey = _jrandom.split(key)
        val0 = self._sample_initial(subkey)
        default_state = _PRNGState(key=key, val=val0)
        self.declare_discrete_state(default_value=default_state, as_array=False)

    # ------------------------------------------------------------------ #
    # Per-distribution sampling                                          #
    # ------------------------------------------------------------------ #

    def _draw_uniform(self, subkey):
        """Stop-gradient unit-uniform draw of the configured shape."""
        return self._jlax.stop_gradient(
            self._jrandom.uniform(subkey, self._shape)
        )

    def _draw_normal(self, subkey):
        """Stop-gradient standard-normal draw of the configured shape."""
        return self._jlax.stop_gradient(
            self._jrandom.normal(subkey, self._shape)
        )

    def _sample_uniform(self, subkey, params):
        u = self._draw_uniform(subkey)
        return params["low"] + (params["high"] - params["low"]) * u

    def _sample_normal(self, subkey, params):
        z = self._draw_normal(subkey)
        return params["mean"] + params["std"] * z

    def _sample_lognormal(self, subkey, params):
        z = self._draw_normal(subkey)
        return self._jnp.exp(params["mu"] + params["sigma"] * z)

    def _sample_triangular(self, subkey, params):
        # Inverse-CDF ("quantile") transform of u ~ U[0,1) for the
        # triangular distribution on [low, high] with mode ``peak``:
        #     F^{-1}(u) =
        #       low  + sqrt(u  * (high-low) * (peak-low))     if u <= c
        #       high - sqrt((1-u) * (high-low) * (high-peak)) otherwise
        # where c = (peak - low) / (high - low) is the CDF at peak.
        # Smooth in low/peak/high (away from the degenerate
        # peak == low or peak == high boundaries), so jax.grad flows
        # cleanly through all three.
        u = self._draw_uniform(subkey)
        low = params["low"]
        peak = params["peak"]
        high = params["high"]
        width = high - low
        # ``c`` is the cumulative probability at the peak; the where
        # branches on a constant-shape boolean so jit/vmap are fine.
        c = (peak - low) / width
        left = low + self._jnp.sqrt(u * width * (peak - low))
        right = high - self._jnp.sqrt((1.0 - u) * width * (high - peak))
        return self._jnp.where(u <= c, left, right)

    # T-122-followup-poisson — Exponential / Poisson sampling.

    def _sample_exponential(self, subkey, params):
        """Reparameterised exponential draw: ``-log(1 - u) / rate``.

        Smooth in ``rate``, so jax.grad flows cleanly through the
        ``rate`` parameter when ``u`` is drawn under stop_gradient.
        Equivalently ``jax.random.exponential(key) / rate`` — we use
        the inverse-CDF form so the reparameterisation is the same one
        documented in the T-122 phase 1 architecture comment.
        """
        u = self._draw_uniform(subkey)
        # ``log1p(-u)`` is numerically stable near ``u -> 0``.
        return -self._jnp.log1p(-u) / params["rate"]

    def _sample_poisson(self, subkey, params):
        """Discrete Poisson count sampler.

        Output is integer-typed (``jax.random.poisson`` returns int32/
        int64), and the sample is wrapped in ``stop_gradient`` so JAX
        never tries to differentiate the discrete count through
        ``rate``.  This makes ``jax.grad(loss, rate)`` return zero from
        this block — which is correct given Poisson is non-
        differentiable w.r.t. its rate via the sample path.
        """
        # ``jax.random.poisson(key, lam, shape)`` is the canonical entry
        # point.  Cast the (potentially traced) ``rate`` to a JAX scalar
        # so the call works under jit/grad even when ``rate`` arrives
        # as a Python float.
        return self._jlax.stop_gradient(
            self._jrandom.poisson(subkey, params["rate"], shape=self._shape)
        )

    # T-122-followup-categorical — Categorical / discrete-choice sampling.

    def _sample_categorical(self, subkey, params):
        """Discrete categorical draw from a static ``values`` table.

        Picks an index ``i`` with probability ``probs[i]`` (normalised
        at construction) and returns ``values[i]`` (possibly a vector
        for vector-typed ``values``).  The selected index is wrapped in
        ``stop_gradient`` to make the non-differentiability of the
        hard categorical sample explicit — gradients through ``probs``
        and ``values`` from this sample path are zero.  Use the
        ``Categorical.differentiable_sample`` helper in
        :mod:`jaxonomy.uq.distributions` for a Gumbel-softmax relaxation
        that *is* differentiable through ``probs``.

        ``params["probs"]`` may be a tracer (it flows through the
        simulation context as a dynamic-vector parameter); we re-
        normalise here so the sampler is robust to upstream parameter
        edits that would otherwise break the ``sum == 1`` invariant.
        """
        probs = params["probs"]
        probs = probs / self._jnp.sum(probs)
        n_cat = self._values_table.shape[0]
        idx = self._jrandom.choice(subkey, n_cat, shape=self._shape, p=probs)
        idx = self._jlax.stop_gradient(idx)
        return self._values_table[idx]

    # T-122-followup-bernoulli — Binary 0/1 sampling.

    def _sample_bernoulli(self, subkey, params):
        """Bernoulli(p) draw — returns 0 with prob ``1 - p``, 1 with prob ``p``.

        Implemented via ``jax.random.bernoulli`` (which under the hood
        compares a unit-uniform draw to ``p``) cast to int32, then
        wrapped in ``stop_gradient`` so the discrete sample never tries
        to backpropagate through ``p``.  This makes
        ``jax.grad(loss, p)`` return zero from this block via the
        sample path — the corresponding differentiable channel is the
        Gumbel-softmax helper on ``Bernoulli`` /
        ``Categorical`` in :mod:`jaxonomy.uq.distributions`.

        ``p`` may arrive as a tracer (it flows through the simulation
        context as a dynamic scalar parameter); we clip to ``[0, 1]`` to
        be robust to small numerical drift, mirroring the categorical
        re-normalisation guard.
        """
        p = self._jnp.clip(params["p"], 0.0, 1.0)
        # ``jax.random.bernoulli`` returns a boolean array; cast to
        # int32 so downstream consumers see 0/1 integers (matching the
        # ``Categorical([0, 1], [1-p, p])`` semantics this delegates to).
        sample = self._jrandom.bernoulli(subkey, p, shape=self._shape)
        return self._jlax.stop_gradient(sample.astype(self._jnp.int32))

    # T-122-followup-beta-gamma — Beta and Gamma sampling.

    def _sample_beta(self, subkey, params):
        """Beta(alpha, beta) draw on the open ``(0, 1)`` interval.

        Routes through ``jax.random.beta``, which uses an implicit-reparam
        sampler — gradients flow through ``alpha`` / ``beta`` but with
        higher variance than the inverse-CDF reparams used by the other
        continuous distributions.  See the ``Beta`` docstring in
        :mod:`jaxonomy.uq.distributions` for the gradient-variance caveat.
        """
        return self._jrandom.beta(
            subkey, params["alpha"], params["beta"], shape=self._shape
        )

    def _sample_gamma(self, subkey, params):
        """Gamma(shape, scale) draw on ``[0, inf)``.

        Routes through ``jax.random.gamma`` (Marsaglia–Tsang reparam for
        shape >= 1; boost trick for shape < 1) scaled by ``scale``.
        Gradients flow cleanly through ``scale`` via the multiplicative
        rescaling and through ``shape`` via JAX's implicit-reparam
        machinery inside ``jax.random.gamma``.
        """
        z = self._jrandom.gamma(subkey, params["shape"], shape=self._shape)
        return z * params["scale"]

    # T-122-followup-weibull — Weibull sampling via closed-form inverse CDF.

    def _sample_weibull(self, subkey, params):
        """Weibull(shape, scale) draw on ``[0, inf)``.

        Closed-form inverse-CDF reparameterisation::

            x = scale * (-log(1 - u))**(1/shape),  u ~ U[0, 1)

        Smooth in *both* ``shape`` and ``scale`` -- gradients flow
        cleanly through both parameters analytically (no implicit-
        reparam machinery needed).  ``u`` is drawn under
        ``stop_gradient`` so JAX never tries to differentiate the
        random sequence w.r.t. the key.
        """
        u = self._draw_uniform(subkey)
        # ``log1p(-u)`` is numerically stable near ``u -> 0``; clamp the
        # open right boundary at ``1 - eps`` so the log stays finite.
        one_minus_eps = 1.0 - 1e-12
        u_safe = self._jnp.minimum(u, one_minus_eps)
        return params["scale"] * self._jnp.power(
            -self._jnp.log1p(-u_safe), 1.0 / params["shape"]
        )

    # T-122-followup-pareto — Pareto sampling via closed-form inverse CDF.

    def _sample_pareto(self, subkey, params):
        """Pareto(scale, alpha) draw on ``[scale, inf)``.

        Closed-form inverse-CDF reparameterisation::

            x = scale * (1 - u)**(-1 / alpha),  u ~ U[0, 1)

        Smooth in *both* ``scale`` and ``alpha`` -- gradients flow
        cleanly through both parameters analytically (no implicit-
        reparam machinery needed).  ``u`` is drawn under
        ``stop_gradient`` so JAX never tries to differentiate the
        random sequence w.r.t. the key.

        Computed via ``exp(-log1p(-u) / alpha)`` for numerical
        stability near ``u -> 0`` (where ``1 - u`` is close to 1 and
        the naive ``(1 - u)**(-1/alpha)`` form loses precision in
        the log).
        """
        u = self._draw_uniform(subkey)
        # Clamp the open right boundary so the ``-1/alpha`` exponent
        # stays finite at ``u -> 1``.
        one_minus_eps = 1.0 - 1e-12
        u_safe = self._jnp.minimum(u, one_minus_eps)
        return params["scale"] * self._jnp.exp(
            -self._jnp.log1p(-u_safe) / params["alpha"]
        )

    def _sample(self, subkey, params):
        """Dispatch on the static distribution flag."""
        if self._distribution == "uniform":
            return self._sample_uniform(subkey, params)
        if self._distribution == "normal":
            return self._sample_normal(subkey, params)
        if self._distribution == "lognormal":
            return self._sample_lognormal(subkey, params)
        if self._distribution == "triangular":
            return self._sample_triangular(subkey, params)
        if self._distribution == "exponential":
            return self._sample_exponential(subkey, params)
        if self._distribution == "poisson":
            return self._sample_poisson(subkey, params)
        if self._distribution == "categorical":
            return self._sample_categorical(subkey, params)
        if self._distribution == "bernoulli":
            return self._sample_bernoulli(subkey, params)
        if self._distribution == "beta":
            return self._sample_beta(subkey, params)
        if self._distribution == "gamma":
            return self._sample_gamma(subkey, params)
        if self._distribution == "weibull":
            return self._sample_weibull(subkey, params)
        if self._distribution == "pareto":
            return self._sample_pareto(subkey, params)
        # Unreachable: validated at __init__.  Defensive.
        raise ValueError(
            f"RandomSource: unknown distribution {self._distribution!r}"
        )

    def _sample_initial(self, subkey):
        """Draw the initial discrete-state sample using captured defaults."""
        return self._sample(subkey, self._param_defaults)

    # ------------------------------------------------------------------ #
    # LeafSystem callbacks                                               #
    # ------------------------------------------------------------------ #

    def _output(self, _time, state, *_inputs, **_parameters):
        return state.discrete_state.val

    def _update(self, _time, state, *_inputs, **parameters):
        key, subkey = self._jrandom.split(state.discrete_state.key)
        # T-122-followup-vmap-fold-in: fold per-replica batch index into
        # subkey when running under vmap(axis_name="batch") and opted-in.
        subkey = _maybe_fold_in_batch_axis(
            self._jrandom, subkey, self._fold_in_batch_index
        )
        # ``parameters`` is the *dynamic* parameter dict at simulation
        # time; pull just the *dynamic* keys this distribution uses so
        # we don't accidentally depend on stale context entries.  Static
        # keys (e.g. categorical ``values``) come from
        # ``self._param_defaults`` because they were captured at
        # Python time rather than registered as dynamic parameters.
        params = {}
        for k in self._param_keys:
            if (self._distribution, k) in _RANDOM_SOURCE_STATIC_PARAMS:
                params[k] = self._param_defaults[k]
            else:
                # T-122-followup-beta-gamma: dynamic parameters may be
                # registered under a renamed internal name (see
                # ``_internal_param_name``) when the user-facing key
                # collides with a RandomSource static parameter
                # (notably ``gamma``'s ``shape`` vs the output-shape
                # parameter).  Look up by the internal name so the
                # rename is transparent to ``_sample``.
                internal = _internal_param_name(self._distribution, k)
                params[k] = parameters[internal]
        val = self._sample(subkey, params)
        return _PRNGState(key=key, val=val)

RateLimiter

Bases: LeafSystem

Limit the time derivative of the block output.

Given an input signal u computes the derivative of the output signal as:

    y_rate = (u(t) - y(Tprev))/(t - Tprev)

Where Tprev is the last time the block was called for output update.

When y_rate is greater than the upper_limit, the output is:

    y(t) = (t - Tprev)*upper_limit + y(Tprev)

When y_rate is less than the lower_limit, the output is:

    y(t) = (t - Tprev)*lower_limit + y(Tprev)

If the lower_limit is greater than the upper_limit, and both are being violated, the upper_limit takes precedence.

Optionally, the block can also be configured with "dynamic" limits, which will add input ports for time-varying upper and lower limits.

Presently, the block is constrainted to periodic updates.

Input ports

(0) The input signal. (1) The upper limit, if dynamic limits are enabled. (2) The lower limit, if dynamic limits are enabled. (Will be indexed as 1 if dynamic upper limits are not enabled.)

Output ports

(0) The rate limited output signal.

Parameters:

Name	Type	Description	Default
upper_limit		The upper limit of the input signal. Default is np.inf.	inf
enable_dynamic_upper_limit		If True, then the upper limit can be set by an external signal. Default is False.	False
lower_limit		The lower limit of the input signal. Default is -np.inf.	-inf
enable_dynamic_lower_limit		If True, then the lower limit can be set by an external signal. Default is False.	False

T-115-followup-mode-flag

The mode kwarg unifies the smooth (differentiable) variant previously exposed as :class:SoftRateLimiter. mode="hard" (default) is byte-equivalent to the legacy behavior. mode="smooth" replaces the inner per-step delta clip with :func:soft_saturate so gradients flow through active rate limiting. Smooth mode requires finite (static) upper_limit / lower_limit and sharpness > 0 (defaults to 10.0).

Source code in jaxonomy/library/nonlinearities.py

class RateLimiter(LeafSystem):
    """Limit the time derivative of the block output.

    Given an input signal `u` computes the derivative of the output signal as:
    ```
        y_rate = (u(t) - y(Tprev))/(t - Tprev)
    ```
    Where Tprev is the last time the block was called for output update.

    When y_rate is greater than the upper_limit, the output is:
    ```
        y(t) = (t - Tprev)*upper_limit + y(Tprev)
    ```

    When y_rate is less than the lower_limit, the output is:
    ```
        y(t) = (t - Tprev)*lower_limit + y(Tprev)
    ```

    If the lower_limit is greater than the upper_limit, and both
    are being violated, the upper_limit takes precedence.

    Optionally, the block can also be configured with "dynamic" limits, which will
    add input ports for time-varying upper and lower limits.

    Presently, the block is constrainted to periodic updates.

    Input ports:
        (0) The input signal.
        (1) The upper limit, if dynamic limits are enabled.
        (2) The lower limit, if dynamic limits are enabled. (Will be indexed as 1 if
            dynamic upper limits are not enabled.)

    Output ports:
        (0) The rate limited output signal.

    Parameters:
        upper_limit:
            The upper limit of the input signal.  Default is `np.inf`.
        enable_dynamic_upper_limit:
            If True, then the upper limit can be set by an external signal. Default
            is False.
        lower_limit:
            The lower limit of the input signal.  Default is `-np.inf`.
        enable_dynamic_lower_limit:
            If True, then the lower limit can be set by an external signal. Default
            is False.

    T-115-followup-mode-flag:
        The ``mode`` kwarg unifies the smooth (differentiable) variant
        previously exposed as :class:`SoftRateLimiter`. ``mode="hard"``
        (default) is byte-equivalent to the legacy behavior.
        ``mode="smooth"`` replaces the inner per-step delta clip with
        :func:`soft_saturate` so gradients flow through active rate
        limiting. Smooth mode requires finite (static) ``upper_limit`` /
        ``lower_limit`` and ``sharpness > 0`` (defaults to ``10.0``).
    """

    class DiscreteStateType(NamedTuple):
        y_prev: Array
        t_prev: float

    @parameters(
        static=[
            "dt",
            "enable_dynamic_upper_limit",
            "enable_dynamic_lower_limit",
            "mode",
        ],
        dynamic=["upper_limit", "lower_limit", "sharpness"],
    )
    def __init__(
        self,
        dt,
        upper_limit=np.inf,
        enable_dynamic_upper_limit=False,
        lower_limit=-np.inf,
        enable_dynamic_lower_limit=False,
        mode="hard",
        sharpness=10.0,
        **kwargs,
    ):
        if mode not in ("hard", "smooth"):
            raise BlockParameterError(
                message=(
                    f"RateLimiter block: mode must be 'hard' or 'smooth', "
                    f"got {mode!r}."
                ),
                parameter_name="mode",
            )
        super().__init__(**kwargs)
        self.primary_input_index = self.declare_input_port()
        self.enable_dynamic_upper_limit = enable_dynamic_upper_limit
        self.enable_dynamic_lower_limit = enable_dynamic_lower_limit
        self.dt = dt
        self.mode = mode

        if enable_dynamic_upper_limit:
            # If dynamic limit, simply ignore the static limit
            self.upper_limit_index = self.declare_input_port()

        if enable_dynamic_lower_limit:
            # If dynamic limit, simply ignore the static limit
            self.lower_limit_index = self.declare_input_port()

        # Smooth-mode validation: needs finite static limits and positive
        # sharpness (matches SoftRateLimiter contract).
        if mode == "smooth":
            if (
                not enable_dynamic_upper_limit
                and not np.isfinite(upper_limit)
            ):
                raise BlockParameterError(
                    message=(
                        f"RateLimiter block {self.name}: mode='smooth' requires "
                        f"finite upper_limit, got {upper_limit}."
                    ),
                    system=self,
                    parameter_name="upper_limit",
                )
            if (
                not enable_dynamic_lower_limit
                and not np.isfinite(lower_limit)
            ):
                raise BlockParameterError(
                    message=(
                        f"RateLimiter block {self.name}: mode='smooth' requires "
                        f"finite lower_limit, got {lower_limit}."
                    ),
                    system=self,
                    parameter_name="lower_limit",
                )
            if sharpness <= 0:
                raise BlockParameterError(
                    message=(
                        f"RateLimiter block {self.name}: mode='smooth' requires "
                        f"sharpness > 0, got {sharpness}."
                    ),
                    system=self,
                    parameter_name="sharpness",
                )

        self.output_index = self.declare_output_port(
            self._output,
            period=dt,
            offset=0.0,
        )

    def initialize(
        self,
        upper_limit=np.inf,
        enable_dynamic_upper_limit=False,
        lower_limit=-np.inf,
        enable_dynamic_lower_limit=False,
        mode="hard",
        sharpness=10.0,
        dt=None,
    ):
        if enable_dynamic_upper_limit != self.enable_dynamic_upper_limit:
            raise ValueError(
                "RateLimiter: enable_dynamic_upper_limit cannot be changed after initialization"
            )
        if enable_dynamic_lower_limit != self.enable_dynamic_lower_limit:
            raise ValueError(
                "RateLimiter: enable_dynamic_lower_limit cannot be changed after initialization"
            )
        if mode != self.mode:
            raise ValueError(
                "RateLimiter: mode cannot be changed after initialization"
            )

    def _output(self, time, state, *inputs, **params):
        y_prev = state.cache[self.output_index]

        u = inputs[self.primary_input_index]

        t_diff = self.dt

        ulim = (
            inputs[self.upper_limit_index]
            if self.enable_dynamic_upper_limit
            else params["upper_limit"]
        )
        llim = (
            inputs[self.lower_limit_index]
            if self.enable_dynamic_lower_limit
            else params["lower_limit"]
        )

        if self.mode == "smooth":
            # T-115-followup-mode-flag: smooth per-step delta clip via
            # soft_saturate (matches SoftRateLimiter behavior).
            delta = u - y_prev
            delta_lo = t_diff * llim
            delta_hi = t_diff * ulim
            return y_prev + soft_saturate(
                delta, delta_lo, delta_hi, params["sharpness"]
            )

        y_rate = (u - y_prev) / t_diff

        y_ulim = t_diff * ulim + y_prev
        y_llim = t_diff * llim + y_prev
        y_tmp = npa.where(y_rate < llim, y_llim, u)
        y = npa.where(y_rate > ulim, y_ulim, y_tmp)

        return y

    def initialize_static_data(self, context):
        """Infer the size and dtype of the internal states"""
        # If building as part of a subsystem, this may not be fully connected yet.
        # That's fine, as long as it is connected by root context creation time.
        # This probably isn't a good long-term solution:
        #   see https://jaxonomy.atlassian.net/browse/WC-51
        try:
            u = self.eval_input(context)
            self._default_cache[self.output_index] = u
            local_context = context[self.system_id].with_discrete_state(u)
            local_context = local_context.with_cached_value(self.output_index, u)
            context = context.with_subcontext(self.system_id, local_context)

        except UpstreamEvalError:
            logger.debug(
                "RateLimiter.initialize_static_data: UpstreamEvalError. "
                "Continuing without default value initialization."
            )
        return context

initialize_static_data(context)

Infer the size and dtype of the internal states

Source code in jaxonomy/library/nonlinearities.py

def initialize_static_data(self, context):
    """Infer the size and dtype of the internal states"""
    # If building as part of a subsystem, this may not be fully connected yet.
    # That's fine, as long as it is connected by root context creation time.
    # This probably isn't a good long-term solution:
    #   see https://jaxonomy.atlassian.net/browse/WC-51
    try:
        u = self.eval_input(context)
        self._default_cache[self.output_index] = u
        local_context = context[self.system_id].with_discrete_state(u)
        local_context = local_context.with_cached_value(self.output_index, u)
        context = context.with_subcontext(self.system_id, local_context)

    except UpstreamEvalError:
        logger.debug(
            "RateLimiter.initialize_static_data: UpstreamEvalError. "
            "Continuing without default value initialization."
        )
    return context

Reciprocal

Bases: FeedthroughBlock

Compute the reciprocal of the input signal.

Input ports

(0) The input signal.

Output ports

(0) The reciprocal of the input signal: y = 1 / u.

Source code in jaxonomy/library/math_ops.py

class Reciprocal(FeedthroughBlock):
    """Compute the reciprocal of the input signal.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The reciprocal of the input signal: `y = 1 / u`.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(lambda x: 1 / x, *args, **kwargs)

RecursiveLeastSquares

Bases: LeafSystem

Recursive Least Squares (RLS) estimator for online parameter identification in linear-in-parameters models:

``y[k] = φ[k]ᵀ θ + noise``

where

• y[k] is a scalar (or vector) measurement at timestep k,
• φ[k] is a regressor vector of size n_params,
• θ is the unknown parameter vector to be estimated.

The RLS update equations with forgetting factor λ are:

.. code-block:: text

e[k]    = y[k]  − φ[k]ᵀ θ̂[k−1]               (prediction error)
K[k]    = P[k−1] φ[k] / (λ + φ[k]ᵀ P[k−1] φ[k])  (Kalman gain)
θ̂[k]   = θ̂[k−1] + K[k] e[k]                  (parameter update)
P[k]    = (P[k−1] − K[k] φ[k]ᵀ P[k−1]) / λ    (covariance update)

A forgetting factor λ < 1 down-weights older measurements, making the estimator track slowly time-varying parameters. λ = 1 (default) is the classic batch RLS equivalent.

The block is fully JAX-traceable and compatible with JIT/autodiff.

                +--------------------+
--- phi[k] ---->|                    |----> theta_hat[k]
                |  Recursive Least   |----> P[k]
--- y[k] ------>|      Squares       |----> prediction_error[k]
                +--------------------+

Input ports

(0) phi : regressor vector at timestep k, shape (n_params,) (1) y : scalar measurement at timestep k

Output ports

(0) theta_hat : parameter estimate, shape (n_params,) (1) P : parameter covariance matrix, shape (n_params, n_params) (2) prediction_error : scalar prediction residual e = y − φᵀ θ̂

Parameters:

Name	Type	Description	Default
dt		float Sampling period.	required
n_params		int Number of parameters to estimate (dimension of θ).	required
theta_0		array_like, optional Initial parameter estimate, shape (n_params,). Defaults to the zero vector.	required
P_0		array_like, optional Initial covariance matrix, shape (n_params, n_params). Defaults to 1e4 * I, which encodes high initial uncertainty.	required
forgetting_factor		float, optional Forgetting factor λ ∈ (0, 1]. Default 1.0 (no forgetting).	required

Example::

import numpy as np
import jaxonomy
from jaxonomy import library, DiagramBuilder, SimulatorOptions

# True parameters: y = 2*phi_0 + 3*phi_1
TRUE_THETA = np.array([2.0, 3.0])
DT = 0.1

rls = library.RecursiveLeastSquares(
    dt=DT, n_params=2,
    forgetting_factor=1.0,
)

Source code in jaxonomy/library/state_estimators/rls.py

class RecursiveLeastSquares(LeafSystem):
    """
    Recursive Least Squares (RLS) estimator for online parameter identification
    in linear-in-parameters models:

        ``y[k] = φ[k]ᵀ θ + noise``

    where:
      - ``y[k]``   is a scalar (or vector) measurement at timestep k,
      - ``φ[k]``   is a regressor vector of size ``n_params``,
      - ``θ``      is the unknown parameter vector to be estimated.

    The RLS update equations with forgetting factor λ are:

    .. code-block:: text

        e[k]    = y[k]  − φ[k]ᵀ θ̂[k−1]               (prediction error)
        K[k]    = P[k−1] φ[k] / (λ + φ[k]ᵀ P[k−1] φ[k])  (Kalman gain)
        θ̂[k]   = θ̂[k−1] + K[k] e[k]                  (parameter update)
        P[k]    = (P[k−1] − K[k] φ[k]ᵀ P[k−1]) / λ    (covariance update)

    A forgetting factor ``λ < 1`` down-weights older measurements, making the
    estimator track slowly time-varying parameters.  ``λ = 1`` (default) is the
    classic batch RLS equivalent.

    The block is fully JAX-traceable and compatible with JIT/autodiff.

    ```
                    +--------------------+
    --- phi[k] ---->|                    |----> theta_hat[k]
                    |  Recursive Least   |----> P[k]
    --- y[k] ------>|      Squares       |----> prediction_error[k]
                    +--------------------+
    ```

    Input ports:
        (0) phi : regressor vector at timestep k, shape ``(n_params,)``
        (1) y   : scalar measurement at timestep k

    Output ports:
        (0) theta_hat        : parameter estimate, shape ``(n_params,)``
        (1) P                : parameter covariance matrix, shape ``(n_params, n_params)``
        (2) prediction_error : scalar prediction residual ``e = y − φᵀ θ̂``

    Parameters:
        dt : float
            Sampling period.
        n_params : int
            Number of parameters to estimate (dimension of θ).
        theta_0 : array_like, optional
            Initial parameter estimate, shape ``(n_params,)``.
            Defaults to the zero vector.
        P_0 : array_like, optional
            Initial covariance matrix, shape ``(n_params, n_params)``.
            Defaults to ``1e4 * I``, which encodes high initial uncertainty.
        forgetting_factor : float, optional
            Forgetting factor λ ∈ (0, 1].  Default ``1.0`` (no forgetting).

    Example::

        import numpy as np
        import jaxonomy
        from jaxonomy import library, DiagramBuilder, SimulatorOptions

        # True parameters: y = 2*phi_0 + 3*phi_1
        TRUE_THETA = np.array([2.0, 3.0])
        DT = 0.1

        rls = library.RecursiveLeastSquares(
            dt=DT, n_params=2,
            forgetting_factor=1.0,
        )
    """

    class DiscreteStateType(NamedTuple):
        """Internal state: current parameter estimate and covariance."""

        theta_hat: npa.ndarray  # shape (n_params,)
        P: npa.ndarray          # shape (n_params, n_params)

    @parameters(
        static=["dt", "n_params", "theta_0", "P_0", "forgetting_factor"],
    )
    def __init__(
        self,
        dt,
        n_params,
        theta_0=None,
        P_0=None,
        forgetting_factor=1.0,
        name=None,
        **kwargs,
    ):
        super().__init__(name=name, **kwargs)

        # Resolve defaults for array arguments so that ports can be declared
        # with appropriate shapes at construction time.
        if theta_0 is None:
            theta_0 = jnp.zeros(n_params)
        if P_0 is None:
            P_0 = jnp.eye(n_params) * 1e4

        theta_0 = jnp.asarray(theta_0, dtype=float)
        P_0 = jnp.asarray(P_0, dtype=float)

        # Input ports
        self.phi_in_index = self.declare_input_port(name="phi")
        self.y_in_index = self.declare_input_port(name="y")

        # Internal discrete state: (theta_hat, P)
        self.declare_discrete_state(
            default_value=self.DiscreteStateType(theta_hat=theta_0, P=P_0),
            as_array=False,
        )

        # Periodic update – runs at each timestep
        self.declare_periodic_update(
            self._update,
            period=dt,
            offset=0.0,
        )

        # Dependency tickets for feedthrough outputs
        phi_ticket = self.input_ports[self.phi_in_index].ticket
        y_ticket = self.input_ports[self.y_in_index].ticket
        prereqs = [DependencyTicket.xd, phi_ticket, y_ticket]
        required_inputs = [self.phi_in_index, self.y_in_index]

        # Output port 0: theta_hat  (feedthrough on phi, y)
        self.declare_output_port(
            self._output_theta_hat,
            period=dt,
            offset=0.0,
            default_value=theta_0,
            name="theta_hat",
            requires_inputs=required_inputs,
            prerequisites_of_calc=prereqs,
        )

        # Output port 1: P  (feedthrough on phi, y)
        self.declare_output_port(
            self._output_P,
            period=dt,
            offset=0.0,
            default_value=P_0,
            name="P",
            requires_inputs=required_inputs,
            prerequisites_of_calc=prereqs,
        )

        # Output port 2: prediction_error  (feedthrough on phi, y)
        self.declare_output_port(
            self._output_prediction_error,
            period=dt,
            offset=0.0,
            default_value=jnp.zeros(()),
            name="prediction_error",
            requires_inputs=required_inputs,
            prerequisites_of_calc=prereqs,
        )

    def initialize(
        self,
        dt,
        n_params,
        theta_0=None,
        P_0=None,
        forgetting_factor=1.0,
    ):
        """Called at context-creation time to store resolved parameters."""
        self.n = n_params
        self.lam = float(forgetting_factor)

    # ──────────────────────────────────────────────────────────────────────────
    # Core RLS computation (shared between update and outputs)
    # ──────────────────────────────────────────────────────────────────────────

    @staticmethod
    def _rls_step(theta_hat, P, phi, y, lam):
        """One RLS correction step.  Returns (theta_new, P_new, error)."""
        phi = jnp.atleast_1d(jnp.asarray(phi, dtype=float)).ravel()
        y_scalar = jnp.asarray(y, dtype=float).reshape(())

        # Prediction error
        e = y_scalar - jnp.dot(phi, theta_hat)

        # Kalman gain
        Pphi = jnp.dot(P, phi)
        denom = lam + jnp.dot(phi, Pphi)
        K = Pphi / denom

        # Parameter update
        theta_new = theta_hat + K * e

        # Covariance update  (Joseph form is more numerically robust but
        # the standard form is sufficient here and cheaper to compute)
        P_new = (P - jnp.outer(K, phi) @ P) / lam

        return theta_new, P_new, e

    # ──────────────────────────────────────────────────────────────────────────
    # Periodic state update
    # ──────────────────────────────────────────────────────────────────────────

    def _update(self, time, state, *inputs, **params):
        phi, y = inputs
        theta_hat = state.discrete_state.theta_hat
        P = state.discrete_state.P

        theta_new, P_new, _ = self._rls_step(theta_hat, P, phi, y, self.lam)

        return self.DiscreteStateType(theta_hat=theta_new, P=P_new)

    # ──────────────────────────────────────────────────────────────────────────
    # Output callbacks  (feedthrough: recompute correction with current inputs)
    # ──────────────────────────────────────────────────────────────────────────

    def _output_theta_hat(self, time, state, *inputs, **params):
        phi, y = inputs
        theta_hat = state.discrete_state.theta_hat
        P = state.discrete_state.P
        theta_new, _, _ = self._rls_step(theta_hat, P, phi, y, self.lam)
        return theta_new

    def _output_P(self, time, state, *inputs, **params):
        phi, y = inputs
        theta_hat = state.discrete_state.theta_hat
        P = state.discrete_state.P
        _, P_new, _ = self._rls_step(theta_hat, P, phi, y, self.lam)
        return P_new

    def _output_prediction_error(self, time, state, *inputs, **params):
        phi, y = inputs
        theta_hat = state.discrete_state.theta_hat
        phi = jnp.atleast_1d(jnp.asarray(phi, dtype=float)).ravel()
        y_scalar = jnp.asarray(y, dtype=float).reshape(())
        return y_scalar - jnp.dot(phi, theta_hat)

DiscreteStateType

Bases: NamedTuple

Internal state: current parameter estimate and covariance.

Source code in jaxonomy/library/state_estimators/rls.py

class DiscreteStateType(NamedTuple):
    """Internal state: current parameter estimate and covariance."""

    theta_hat: npa.ndarray  # shape (n_params,)
    P: npa.ndarray          # shape (n_params, n_params)

initialize(dt, n_params, theta_0=None, P_0=None, forgetting_factor=1.0)

Called at context-creation time to store resolved parameters.

Source code in jaxonomy/library/state_estimators/rls.py

def initialize(
    self,
    dt,
    n_params,
    theta_0=None,
    P_0=None,
    forgetting_factor=1.0,
):
    """Called at context-creation time to store resolved parameters."""
    self.n = n_params
    self.lam = float(forgetting_factor)

ReducedOrderModel dataclass

A reduced model plus its provenance.

Attributes:

Name	Type	Description
system	Any	The reduced, simulatable Jaxonomy block — an LTISystem / LTISystemDiscrete for linear MOR, or a discrete-time predictor LeafSystem for data-driven methods. Drop it straight into a diagram or jaxonomy.simulate.
method	str	The reduction method that produced it.
full_order	Optional[int]	State dimension of the source model (when known).
reduced_order	Optional[int]	State dimension of system.
info	dict	Method-specific extras — e.g. error_bound and hsv for balanced truncation, eigenvalues / basis for DMD, and the raw result object from the underlying routine.

Source code in jaxonomy/library/rom/framework.py

@dataclass
class ReducedOrderModel:
    """A reduced model plus its provenance.

    Attributes:
        system: The reduced, simulatable Jaxonomy block — an ``LTISystem`` /
            ``LTISystemDiscrete`` for linear MOR, or a discrete-time predictor
            ``LeafSystem`` for data-driven methods. Drop it straight into a
            diagram or ``jaxonomy.simulate``.
        method: The reduction method that produced it.
        full_order: State dimension of the source model (when known).
        reduced_order: State dimension of ``system``.
        info: Method-specific extras — e.g. ``error_bound`` and ``hsv`` for
            balanced truncation, ``eigenvalues`` / ``basis`` for DMD, and the
            raw ``result`` object from the underlying routine.
    """

    system: Any
    method: str
    full_order: Optional[int] = None
    reduced_order: Optional[int] = None
    info: dict = field(default_factory=dict)

    def to_block(self):
        """Return the reduced Jaxonomy block (alias for ``.system``)."""
        return self.system

    def __repr__(self):
        return (
            f"ReducedOrderModel(method={self.method!r}, "
            f"full_order={self.full_order}, reduced_order={self.reduced_order})"
        )

to_block()

Return the reduced Jaxonomy block (alias for .system).

Source code in jaxonomy/library/rom/framework.py

def to_block(self):
    """Return the reduced Jaxonomy block (alias for ``.system``)."""
    return self.system

ReferenceSubdiagram

Registry for reusable diagram templates ("reference subdiagrams").

A reference subdiagram is a parameterized diagram factory. It is registered once via :meth:register and can then be instantiated multiple times with different parameter values via :meth:create_diagram.

Example::

def my_submodel(instance_name, parameters):
    builder = DiagramBuilder()
    gain = parameters["gain"].get()
    ...
    return builder.build(instance_name)

ref_id = ReferenceSubdiagram.register(
    my_submodel,
    default_parameters=[Parameter("gain", 1.0)],
)
diagram = ReferenceSubdiagram.create_diagram(ref_id, "my_instance")

Source code in jaxonomy/library/reference_subdiagram.py

class ReferenceSubdiagram:
    """Registry for reusable diagram templates ("reference subdiagrams").

    A reference subdiagram is a parameterized diagram factory.  It is registered
    once via :meth:`register` and can then be instantiated multiple times with
    different parameter values via :meth:`create_diagram`.

    Example::

        def my_submodel(instance_name, parameters):
            builder = DiagramBuilder()
            gain = parameters["gain"].get()
            ...
            return builder.build(instance_name)

        ref_id = ReferenceSubdiagram.register(
            my_submodel,
            default_parameters=[Parameter("gain", 1.0)],
        )
        diagram = ReferenceSubdiagram.create_diagram(ref_id, "my_instance")
    """

    _registry: dict[str, Callable[[Any], "Diagram"]] = {}
    _default_parameters: dict[str, list[Parameter]] = {}  # noqa: F821

    @classmethod
    def create_diagram(
        cls,
        ref_id: str,
        instance_name: str,
        *args,
        instance_parameters: dict[str, Any] = None,
        **kwargs,
    ) -> "Diagram":
        """Create a diagram based on the given reference ID and parameters.

        Note that for submodels we evaluate all parameters, there is no
        "pure" string parameters.

        Args:
            ref_id (str): The reference ID of the diagram.
            instance_name (str): Name for this specific instance.
            *args: Variable length arguments passed to the constructor.
            instance_parameters (dict[str, Any], optional): Per-instance parameter
                overrides.  Keys must match names declared at registration time.
                Example: ``{"gain": 3.0}``
            **kwargs: Keyword arguments passed to the constructor.

        Returns:
            Diagram: The created diagram.

        Raises:
            ValueError: If the reference subdiagram with the given ref_id is not found.
            ValueError: If an instance_parameter key does not match any registered
                parameter.
        """
        if ref_id not in ReferenceSubdiagram._registry:
            raise ValueError(f"ReferenceSubdiagram with ref_id {ref_id} not found.")

        params_def = ReferenceSubdiagram.get_default_parameters(ref_id)

        default_params = {p.name: p for p in params_def}

        # override the default values with any 'modified' values.
        new_instance_parameters = {}
        if instance_parameters:
            for param_name, param in instance_parameters.items():
                if param_name not in default_params:
                    raise ValueError(
                        f"Parameter {param_name} not found in parameter definitions."
                    )
                new_instance_parameters[param_name] = Parameter(
                    name=param_name, value=param
                )

        all_params = {**default_params, **new_instance_parameters}

        diagram = ReferenceSubdiagram._registry[ref_id](
            *args,
            instance_name=instance_name,
            parameters=all_params,
            **kwargs,
        )

        diagram.ref_id = ref_id
        diagram.instance_parameters = set(new_instance_parameters.keys())

        for param in params_def:
            if param.name in new_instance_parameters:
                diagram.declare_dynamic_parameter(
                    param.name, new_instance_parameters[param.name]
                )
            else:
                diagram.declare_dynamic_parameter(param.name, param)

        return diagram

    @staticmethod
    def register(
        constructor: ReferenceSubdiagramProtocol,
        default_parameters: list[Parameter] = None,  # noqa: F821
        ref_id: str | None = None,
        # Deprecated alias – use default_parameters instead
        parameter_definitions: list[Parameter] = None,  # noqa: F821
    ) -> str:
        """Register a diagram constructor as a reusable reference subdiagram.

        Args:
            constructor: A callable that builds a :class:`Diagram` given
                ``instance_name`` and ``parameters``.
            default_parameters: Default :class:`Parameter` values for this
                subdiagram.  Instances can override individual parameters at
                creation time via :meth:`create_diagram`.
            ref_id: Optional stable identifier.  A UUID is generated if omitted.
            parameter_definitions: **Deprecated** – use ``default_parameters``.

        Returns:
            str: The ``ref_id`` that can be passed to :meth:`create_diagram`.
        """
        import warnings

        if parameter_definitions is not None:
            warnings.warn(
                "The 'parameter_definitions' argument is deprecated; "
                "use 'default_parameters' instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            if default_parameters is None:
                default_parameters = parameter_definitions

        if ref_id is None:
            ref_id = str(uuid4())
        if default_parameters is None:
            default_parameters = []

        logger.debug("Registering ReferenceSubdiagram with ref_id %s", ref_id)
        if ref_id in ReferenceSubdiagram._registry:
            logger.debug(
                "ReferenceSubdiagram with ref_id %s already registered.",
                ref_id,
            )

        ReferenceSubdiagram._registry[ref_id] = constructor
        ReferenceSubdiagram._default_parameters[ref_id] = default_parameters

        return ref_id

    @staticmethod
    def get_default_parameters(
        ref_id: str,
    ) -> list[Parameter]:  # noqa: F821
        """Return the default parameters for the given reference subdiagram."""
        if ref_id not in ReferenceSubdiagram._default_parameters:
            return []
        return ReferenceSubdiagram._default_parameters[ref_id]

    @staticmethod
    def get_parameter_definitions(
        ref_id: str,
    ) -> list[Parameter]:  # noqa: F821
        """Return the default parameters for the given reference subdiagram.

        .. deprecated::
            Use :meth:`get_default_parameters` instead.
        """
        return ReferenceSubdiagram.get_default_parameters(ref_id)

create_diagram(ref_id, instance_name, *args, instance_parameters=None, **kwargs) classmethod

Create a diagram based on the given reference ID and parameters.

Note that for submodels we evaluate all parameters, there is no "pure" string parameters.

Parameters:

Name	Type	Description	Default
ref_id	str	The reference ID of the diagram.	required
instance_name	str	Name for this specific instance.	required
*args		Variable length arguments passed to the constructor.	()
instance_parameters	dict[str, Any]	Per-instance parameter overrides. Keys must match names declared at registration time. Example: {"gain": 3.0}	None
**kwargs		Keyword arguments passed to the constructor.	{}

Returns:

Name	Type	Description
Diagram	Diagram	The created diagram.

Raises:

Type	Description
ValueError	If the reference subdiagram with the given ref_id is not found.
ValueError	If an instance_parameter key does not match any registered parameter.

Source code in jaxonomy/library/reference_subdiagram.py

@classmethod
def create_diagram(
    cls,
    ref_id: str,
    instance_name: str,
    *args,
    instance_parameters: dict[str, Any] = None,
    **kwargs,
) -> "Diagram":
    """Create a diagram based on the given reference ID and parameters.

    Note that for submodels we evaluate all parameters, there is no
    "pure" string parameters.

    Args:
        ref_id (str): The reference ID of the diagram.
        instance_name (str): Name for this specific instance.
        *args: Variable length arguments passed to the constructor.
        instance_parameters (dict[str, Any], optional): Per-instance parameter
            overrides.  Keys must match names declared at registration time.
            Example: ``{"gain": 3.0}``
        **kwargs: Keyword arguments passed to the constructor.

    Returns:
        Diagram: The created diagram.

    Raises:
        ValueError: If the reference subdiagram with the given ref_id is not found.
        ValueError: If an instance_parameter key does not match any registered
            parameter.
    """
    if ref_id not in ReferenceSubdiagram._registry:
        raise ValueError(f"ReferenceSubdiagram with ref_id {ref_id} not found.")

    params_def = ReferenceSubdiagram.get_default_parameters(ref_id)

    default_params = {p.name: p for p in params_def}

    # override the default values with any 'modified' values.
    new_instance_parameters = {}
    if instance_parameters:
        for param_name, param in instance_parameters.items():
            if param_name not in default_params:
                raise ValueError(
                    f"Parameter {param_name} not found in parameter definitions."
                )
            new_instance_parameters[param_name] = Parameter(
                name=param_name, value=param
            )

    all_params = {**default_params, **new_instance_parameters}

    diagram = ReferenceSubdiagram._registry[ref_id](
        *args,
        instance_name=instance_name,
        parameters=all_params,
        **kwargs,
    )

    diagram.ref_id = ref_id
    diagram.instance_parameters = set(new_instance_parameters.keys())

    for param in params_def:
        if param.name in new_instance_parameters:
            diagram.declare_dynamic_parameter(
                param.name, new_instance_parameters[param.name]
            )
        else:
            diagram.declare_dynamic_parameter(param.name, param)

    return diagram

get_default_parameters(ref_id) staticmethod

Return the default parameters for the given reference subdiagram.

Source code in jaxonomy/library/reference_subdiagram.py

@staticmethod
def get_default_parameters(
    ref_id: str,
) -> list[Parameter]:  # noqa: F821
    """Return the default parameters for the given reference subdiagram."""
    if ref_id not in ReferenceSubdiagram._default_parameters:
        return []
    return ReferenceSubdiagram._default_parameters[ref_id]

get_parameter_definitions(ref_id) staticmethod

Return the default parameters for the given reference subdiagram.

.. deprecated:: Use :meth:get_default_parameters instead.

Source code in jaxonomy/library/reference_subdiagram.py

@staticmethod
def get_parameter_definitions(
    ref_id: str,
) -> list[Parameter]:  # noqa: F821
    """Return the default parameters for the given reference subdiagram.

    .. deprecated::
        Use :meth:`get_default_parameters` instead.
    """
    return ReferenceSubdiagram.get_default_parameters(ref_id)

register(constructor, default_parameters=None, ref_id=None, parameter_definitions=None) staticmethod

Register a diagram constructor as a reusable reference subdiagram.

Parameters:

Name	Type	Description	Default
constructor	ReferenceSubdiagramProtocol	A callable that builds a :class:Diagram given instance_name and parameters.	required
default_parameters	list[Parameter]	Default :class:Parameter values for this subdiagram. Instances can override individual parameters at creation time via :meth:create_diagram.	None
ref_id	str | None	Optional stable identifier. A UUID is generated if omitted.	None
parameter_definitions	list[Parameter]	Deprecated – use default_parameters.	None

Returns:

Name	Type	Description
str	str	The ref_id that can be passed to :meth:create_diagram.

Source code in jaxonomy/library/reference_subdiagram.py

@staticmethod
def register(
    constructor: ReferenceSubdiagramProtocol,
    default_parameters: list[Parameter] = None,  # noqa: F821
    ref_id: str | None = None,
    # Deprecated alias – use default_parameters instead
    parameter_definitions: list[Parameter] = None,  # noqa: F821
) -> str:
    """Register a diagram constructor as a reusable reference subdiagram.

    Args:
        constructor: A callable that builds a :class:`Diagram` given
            ``instance_name`` and ``parameters``.
        default_parameters: Default :class:`Parameter` values for this
            subdiagram.  Instances can override individual parameters at
            creation time via :meth:`create_diagram`.
        ref_id: Optional stable identifier.  A UUID is generated if omitted.
        parameter_definitions: **Deprecated** – use ``default_parameters``.

    Returns:
        str: The ``ref_id`` that can be passed to :meth:`create_diagram`.
    """
    import warnings

    if parameter_definitions is not None:
        warnings.warn(
            "The 'parameter_definitions' argument is deprecated; "
            "use 'default_parameters' instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        if default_parameters is None:
            default_parameters = parameter_definitions

    if ref_id is None:
        ref_id = str(uuid4())
    if default_parameters is None:
        default_parameters = []

    logger.debug("Registering ReferenceSubdiagram with ref_id %s", ref_id)
    if ref_id in ReferenceSubdiagram._registry:
        logger.debug(
            "ReferenceSubdiagram with ref_id %s already registered.",
            ref_id,
        )

    ReferenceSubdiagram._registry[ref_id] = constructor
    ReferenceSubdiagram._default_parameters[ref_id] = default_parameters

    return ref_id

Relay

Bases: LeafSystem

Simple state machine implementing hysteresis behavior.

The input-output map is as follows:

        output
          |
on_value  |          -------<------<---------------------
          |          |                    |
          |          ⌄                    ^
          |          |                    |
off_value |----------|-------->----->-----|
          |
          |---------------------------------------------- input
                     | off_threshold      | on_threshold

Note that the "time mode" behavior of this block will follow the input signal. That is, if the input signal varies continuously in time, then the zero-crossing event from OFF->ON or vice versa will be localized in time. On the other hand, if the input signal varies only as a result of periodic updates to the discrete state, the relay will only change state at those instants. If the input signal is continuous, the block can be "forced" to this discrete-time periodic behavior by adding a ZeroOrderHold block before the input.

The exception to this is the case where there are no blocks in the system containing either discrete or continuous state. In this case the state changes will only be localized to the resolution of the major step.

Input ports

(0) The input signal.

Output ports

(0) The relay output signal, which is equal to either the on_value or the off_value, depending on the internal state of the relay.

Parameters:

Name	Type	Description	Default
on_threshold		When input rises above this value, the internal state transitions to ON.	required
off_threshold		When input falls below this value, the internal state transitions to OFF.	required
on_value		Value of the output signal when state is ON.	required
off_value		Value of the output signal when state is OFF	required
initial_state		If equal to on_value, the block will be initialized in the ON state. Otherwise, it will be initialized to the OFF state.	required

Events

There are two zero-crossing events: one to transition from OFF->ON and one for the opposite transition from ON->OFF.

Source code in jaxonomy/library/logic.py

class Relay(LeafSystem):
    """Simple state machine implementing hysteresis behavior.

    The input-output map is as follows:

    ```
            output
              |
    on_value  |          -------<------<---------------------
              |          |                    |
              |          ⌄                    ^
              |          |                    |
    off_value |----------|-------->----->-----|
              |
              |---------------------------------------------- input
                         | off_threshold      | on_threshold
    ```

    Note that the "time mode" behavior of this block will follow the input
    signal.  That is, if the input signal varies continuously in time, then
    the zero-crossing event from OFF->ON or vice versa will be localized in
    time.  On the other hand, if the input signal varies only as a result
    of periodic updates to the discrete state, the relay will only change state
    at those instants.  If the input signal is continuous, the block can
    be "forced" to this discrete-time periodic behavior by adding a ZeroOrderHold
    block before the input.

    The exception to this is the case where there are no blocks in the system
    containing either discrete or continuous state.  In this case the state changes
    will only be localized to the resolution of the major step.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The relay output signal, which is equal to either the on_value or
            the off_value, depending on the internal state of the relay.

    Parameters:
        on_threshold:
            When input rises above this value, the internal state transitions to ON.
        off_threshold:
            When input falls below this value, the internal state transitions to OFF.
        on_value:
            Value of the output signal when state is ON.
        off_value:
            Value of the output signal when state is OFF
        initial_state:
            If equal to on_value, the block will be initialized in the ON state.
            Otherwise, it will be initialized to the OFF state.

    Events:
        There are two zero-crossing events: one to transition from OFF->ON and one
        for the opposite transition from ON->OFF.
    """

    class State(IntEnum):
        OFF = 0
        ON = 1

    @parameters(
        dynamic=[
            "on_threshold",
            "off_threshold",
            "initial_state",
            "on_value",
            "off_value",
        ],
    )
    def __init__(
        self, on_threshold, off_threshold, on_value, off_value, initial_state, **kwargs
    ):
        super().__init__(**kwargs)

        self.declare_default_mode(
            self.State.ON if initial_state == on_value else self.State.OFF
        )

        self.declare_input_port()
        self.declare_output_port(
            self._output,
            requires_inputs=False,
            prerequisites_of_calc=[DependencyTicket.mode],
        )

        # transition to ON event
        def _on_guard(_time, _state, u, **parameters):
            return u - parameters["on_threshold"]

        self.declare_zero_crossing(
            guard=_on_guard,
            direction="negative_then_non_negative",
            start_mode=self.State.OFF,
            end_mode=self.State.ON,
        )

        # transition to OFF event
        def _off_guard(_time, _state, u, **parameters):
            return u - parameters["off_threshold"]

        self.declare_zero_crossing(
            guard=_off_guard,
            direction="positive_then_non_positive",
            start_mode=self.State.ON,
            end_mode=self.State.OFF,
        )

    def initialize(
        self, on_threshold, off_threshold, on_value, off_value, initial_state
    ):
        self.configure_default_mode(
            self.State.ON if initial_state == on_value else self.State.OFF
        )

    def reset_default_values(self, **dynamic_parameters):
        self.configure_default_mode(
            self.State.ON
            if dynamic_parameters["initial_state"] == dynamic_parameters["on_value"]
            else self.State.OFF
        )

    def _output(self, _time, state, **parameters):
        return npa.where(
            state.mode == self.State.ON,
            parameters["on_value"],
            parameters["off_value"],
        )

ReplicatedFunction

Bases: LeafSystem

Container block: evaluate a submodel N times in parallel via vmap.

Parameters:

Name	Type	Description	Default
submodel	Callable	Callable f(*inputs) -> output. Must be JAX-traceable so vmap can transform it.	required
n	int	Number of replicas.	required
n_inputs	int	Number of input ports the block should declare (and the number of positional inputs the submodel takes).	1
in_axes	Sequence[int | None] | None	Tuple of length n_inputs, matching the vmap in_axes convention: 0 means the corresponding input is already batched along axis 0; None means broadcast the single-instance input to all N replicas. Default is (0,) * n_inputs (all inputs batched).	None
name		Optional block name.	required

Source code in jaxonomy/library/replicated.py

class ReplicatedFunction(LeafSystem):
    """Container block: evaluate a submodel N times in parallel via vmap.

    Args:
        submodel: Callable ``f(*inputs) -> output``.  Must be
            JAX-traceable so ``vmap`` can transform it.
        n: Number of replicas.
        n_inputs: Number of input ports the block should declare (and
            the number of positional inputs the submodel takes).
        in_axes: Tuple of length ``n_inputs``, matching the ``vmap``
            ``in_axes`` convention: ``0`` means the corresponding input
            is already batched along axis 0; ``None`` means broadcast
            the single-instance input to all N replicas.  Default is
            ``(0,) * n_inputs`` (all inputs batched).
        name: Optional block name.
    """

    def __init__(
        self,
        submodel: Callable,
        n: int,
        n_inputs: int = 1,
        in_axes: Sequence[int | None] | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        if n < 1:
            raise ValueError(f"ReplicatedFunction: n must be >= 1, got {n}")
        if n_inputs < 1:
            raise ValueError(
                f"ReplicatedFunction: n_inputs must be >= 1, got {n_inputs}"
            )
        if in_axes is None:
            in_axes = (0,) * n_inputs
        if len(in_axes) != n_inputs:
            raise ValueError(
                f"ReplicatedFunction: len(in_axes) must equal n_inputs "
                f"({n_inputs}), got {len(in_axes)}"
            )
        for ax in in_axes:
            if ax is not None and ax != 0:
                raise ValueError(
                    "ReplicatedFunction: in_axes entries must be 0 or None "
                    f"(got {ax}).  If you need a non-zero axis, transpose "
                    "the input upstream."
                )

        self._n = int(n)
        self._in_axes = tuple(in_axes)
        # axis_size is needed when every in_axes entry is None (all broadcast):
        # JAX vmap cannot infer N from the inputs in that case.
        self._vmapped = jax.vmap(
            submodel, in_axes=self._in_axes, axis_size=self._n,
        )

        for i in range(n_inputs):
            self.declare_input_port(name=f"u_{i}")

        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    def _compute_output(self, time, state, *inputs, **params):
        # ``inputs`` are in port-declaration order; shape rules:
        #   in_axes[i] == 0    → inputs[i] must have leading dim N
        #   in_axes[i] is None → inputs[i] is broadcast as a single value
        # We let JAX's vmap rule validate shapes.
        return self._vmapped(*inputs)

RigidBody

Bases: LeafSystem

Implements dynamics of a single three-dimensional body.

The block models both translational and rotational degrees of freedom, for a total of 6 degrees of freedom. With second-order equations, the block has 12 state variables, 6 for the position/orientation and 6 for the velocities/rates.

Currently only a roll-pitch-yaw (Euler angle) representation is supported for the orientation.

The full 12-dof state vector is x = [p_i, Φ, vᵇ, ωᵇ], where pⁱ is the position in an inertial "world" frame i, Φ is the (roll, pitch, and yaw) Euler angle sequence defining the rotation from the inertial "world" frame to the body frame, vᵇ is the translational velocity with respect to body-fixed axes b, and ωᵇ is the angular velocity about the body-fixed axes.

The mass and inertia properties of the block can independently be defined statically as parameters, or dynamically as inputs to the block.

Input ports

(0) force_vector: 3D force vector, defined in the body-fixed coordinate frame. For example, if gravity is acting on the body, the gravity vector should be pre-rotated using CoordinateRotation.

(1) torque_vector: 3D torque vector, be defined in the body-fixed coordinate frame.

(2) inertia: If enable_external_inertia_matrix=True, this input provides the time-varying body-fixed inertia matrix.

Output ports

(0): The position in the inertial "world" frame pⁱ.

(1): The orientation of the body, represented as a roll-pitch-yaw Euler angle sequence.

(2): The translational velocity with respect to body-fixed axes vᵇ.

(3): The angular velocity about the body-fixed axes ωᵇ.

(4): (if enable_output_state_derivatives=True) The time derivatives of the position variables in the world frame ṗⁱ. Not generally equal to the state vᵇ, defining time derivatives in the body frame.

(5): (if enable_output_state_derivatives=True) The "Euler rates" Φ̇, which are the time derivatives of the Euler angles. Not generally equal to the angular velocity ωᵇ.

(6): (if enable_output_state_derivatives=True) The body-fixed acceleration vector aᵇ.

(7): (if enable_output_state_derivatives=True) The angular acceleration in body-fixed axes ω̇ᵇ.

Parameters:

Name	Type	Description	Default
initial_position	Array	The initial position in the inertial frame.	required
initial_orientation	Array	The initial orientation of the body, represented as a roll-pitch-yaw Euler angle sequence.	required
initial_velocity	Array	The initial translational velocity with respect to body-fixed axes.	required
initial_angular_velocity	Array	The initial angular velocity about the body-fixed axes.	required
enable_external_mass	bool	If True, the block will have one input port for the mass. Otherwise the mass must be provided as a block parameter.	False
mass	float	The constant value for the body mass when enable_external_mass=False. If None, will default to 1.0.	1.0
enable_external_inertia_matrix	bool	If True, the block will have one input port for a (3x3) inertia matrix. Otherwise the inertia matrix must be provided as a block parameter.	False
inertia_matrix		The constant value for the body inertia matrix when enable_external_inertia_matrix=False. If None, will default to the 3x3 identity matrix.	eye(3)
enable_output_state_derivatives	bool	If True, the block will output the time derivatives of the state variables.	False
gravity_vector	Array	The constant gravitational acceleration vector acting on the body, defined in the inertial frame. If None, will default to the zero vector.	zeros(3)

Notes

Assumes that the inertia matrix is computed at the center of mass.

Assumes that the mass and inertia matrix are quasi-steady. This means that if one or both is specified as "dynamic" inputs their time derivative is neglected in the dynamics. For instance, for pure translation (w_b=0) the approximation to Newton's law is F_net = (d/dt)(m * v) ≈ m * (dv/dt).

Source code in jaxonomy/library/rotations.py

class RigidBody(LeafSystem):
    """Implements dynamics of a single three-dimensional body.

    The block models both translational and rotational degrees of freedom, for a
    total of 6 degrees of freedom.  With second-order equations, the block has
    12 state variables, 6 for the position/orientation and 6 for the velocities/rates.

    Currently only a roll-pitch-yaw (Euler angle) representation is supported for
    the orientation.

    The full 12-dof state vector is `x = [p_i, Φ, vᵇ, ωᵇ]`, where `pⁱ` is the
    position in an inertial "world" frame `i`, `Φ` is the (roll, pitch, and yaw)
    Euler angle sequence defining the rotation from the inertial "world" frame to
    the body frame, `vᵇ` is the translational velocity with respect to body-fixed
    axes `b`, and `ωᵇ` is the angular velocity about the body-fixed axes.

    The mass and inertia properties of the block can independently be defined
    statically as parameters, or dynamically as inputs to the block.

    Input ports:
        (0) force_vector: 3D force vector, defined in the _body-fixed_ coordinate
        frame.  For example, if gravity is acting on the body, the gravity vector
        should be pre-rotated using CoordinateRotation.

        (1) torque_vector: 3D torque vector, be defined in the _body-fixed_
        coordinate frame.

        (2) inertia: If `enable_external_inertia_matrix=True`, this input provides
        the time-varying body-fixed inertia matrix.

    Output ports:
        (0): The position in the inertial "world" frame `pⁱ`.

        (1): The orientation of the body, represented as a roll-pitch-yaw Euler
        angle sequence.

        (2): The translational velocity with respect to body-fixed axes `vᵇ`.

        (3): The angular velocity about the body-fixed axes `ωᵇ`.

        (4): (if `enable_output_state_derivatives=True`) The time derivatives of the
        position variables in the world frame `ṗⁱ`. Not generally equal to the state
        `vᵇ`, defining time derivatives in the body frame.

        (5): (if `enable_output_state_derivatives=True`) The "Euler rates" `Φ̇`,
        which are the time derivatives of the Euler angles. Not generally equal to
        the angular velocity `ωᵇ`.

        (6): (if `enable_output_state_derivatives=True`) The body-fixed acceleration
        vector `aᵇ`.

        (7): (if `enable_output_state_derivatives=True`) The angular acceleration in
        body-fixed axes `ω̇ᵇ`.

    Parameters:
        initial_position (Array): The initial position in the inertial frame.

        initial_orientation (Array): The initial orientation of the body, represented
            as a roll-pitch-yaw Euler angle sequence.

        initial_velocity (Array): The initial translational velocity with respect to
            body-fixed axes.

        initial_angular_velocity (Array): The initial angular velocity about the
            body-fixed axes.

        enable_external_mass (bool, optional): If `True`, the block will have one
            input port for the mass. Otherwise the mass must be provided as a block
            parameter.

        mass (float, optional): The constant value for the body mass when
            `enable_external_mass=False`. If `None`, will default to 1.0.

        enable_external_inertia_matrix (bool, optional):  If `True`, the block will
            have one input port for a (3x3) inertia matrix. Otherwise the inertia
            matrix must be provided as a block parameter.

        inertia_matrix: The constant value for the body inertia matrix when
            `enable_external_inertia_matrix=False`. If `None`, will default to
            the 3x3 identity matrix.

        enable_output_state_derivatives (bool, optional): If `True`, the block will
            output the time derivatives of the state variables.

        gravity_vector (Array, optional): The constant gravitational acceleration vector
            acting on the body, defined in the _inertial_ frame. If `None`, will default
            to the zero vector.

    Notes:
        Assumes that the inertia matrix is computed at the center of mass.

        Assumes that the mass and inertia matrix are quasi-steady.  This means that
        if one or both is specified as "dynamic" inputs their time derivative is
        neglected in the dynamics.  For instance, for pure translation (`w_b=0`) the
        approximation to Newton's law is `F_net = (d/dt)(m * v) ≈ m * (dv/dt)`.
    """

    class RigidBodyState(NamedTuple):
        position: Array
        orientation: Array
        velocity: Array
        angular_velocity: Array

        def asarray(self):
            return npa.concatenate(
                [self.position, self.orientation, self.velocity, self.angular_velocity]
            )

    @parameters(
        static=[
            "initial_position",
            "initial_orientation",
            "initial_velocity",
            "initial_angular_velocity",
            "enable_external_mass",
            "enable_external_inertia_matrix",
            "enable_output_state_derivatives",
        ],
        dynamic=["mass", "inertia_matrix", "gravity_vector"],
    )
    def __init__(
        self,
        initial_position,
        initial_orientation,
        initial_velocity,
        initial_angular_velocity,
        enable_external_mass=False,
        mass=1.0,
        enable_external_inertia_matrix=False,
        inertia_matrix=npa.eye(3),
        enable_output_state_derivatives=False,
        gravity_vector=npa.zeros(3),
        **kwargs,
    ):
        super().__init__(**kwargs)

        self._enable_external_mass = enable_external_mass
        self._enable_external_inertia_matrix = enable_external_inertia_matrix
        self._enable_output_state_derivatives = enable_output_state_derivatives

        initial_state = self._make_initial_state(
            initial_position,
            initial_orientation,
            initial_velocity,
            initial_angular_velocity,
        )

        self._continuous_state_idx = self.declare_continuous_state(
            default_value=initial_state,
            as_array=False,
            ode=self._state_derivative,
        )

        self._configure_ports(
            initial_state,
            enable_external_mass,
            enable_external_inertia_matrix,
            enable_output_state_derivatives,
        )

    def initialize(
        self,
        initial_position,
        initial_orientation,
        initial_velocity,
        initial_angular_velocity,
        enable_external_mass,
        enable_external_inertia_matrix,
        enable_output_state_derivatives,
        mass,
        inertia_matrix,
        gravity_vector,
    ):
        if enable_external_mass != self._enable_external_mass:
            raise ValueError("Cannot change external mass definition.")
        if enable_external_inertia_matrix != self._enable_external_inertia_matrix:
            raise ValueError("Cannot change external inertia matrix definition.")
        if enable_output_state_derivatives != self._enable_output_state_derivatives:
            raise ValueError("Cannot change output state derivatives definition.")

        gravity_vector = npa.asarray(gravity_vector)
        if gravity_vector.shape != (3,):
            message = (
                "Gravity vector must have shape (3,), but has shape "
                + f"{gravity_vector.shape}."
            )
            raise BlockParameterError(
                message=message, system=self, parameter_name="gravity_vector"
            )

        initial_state = self._make_initial_state(
            initial_position,
            initial_orientation,
            initial_velocity,
            initial_angular_velocity,
        )

        self.configure_continuous_state(
            self._continuous_state_idx,
            default_value=initial_state,
            as_array=False,
            ode=self._state_derivative,
        )

        self.configure_output_port(
            self.pos_output_index,
            self._pos_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            default_value=initial_state.position,
        )

        self.configure_output_port(
            self.orientation_output_index,
            self._orientation_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            default_value=initial_state.orientation,
        )

        self.configure_output_port(
            self.vel_output_index,
            self._vel_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            default_value=initial_state.velocity,
        )

        self.configure_output_port(
            self.ang_vel_output_index,
            self._ang_vel_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            default_value=initial_state.angular_velocity,
        )

    def _make_initial_state(
        self,
        initial_position,
        initial_orientation,
        initial_velocity,
        initial_angular_velocity,
    ):
        # Validate initial state arrays and create named tuple for initial state.
        initial_position = npa.asarray(initial_position)
        if initial_position.shape != (3,):
            message = (
                "Initial position must have shape (3,), but has shape "
                + f"{initial_position.shape}."
            )
            raise BlockParameterError(
                message=message, system=self, parameter_name="initial_position"
            )

        initial_orientation = npa.asarray(initial_orientation)
        if initial_orientation.shape != (3,):
            message = (
                "Initial orientation must have shape (3,), but has shape "
                + f"{initial_orientation.shape}."
            )
            raise BlockParameterError(
                message=message, system=self, parameter_name="initial_orientation"
            )

        initial_velocity = npa.asarray(initial_velocity)
        if initial_velocity.shape != (3,):
            message = (
                "Initial velocity must have shape (3,), but has shape "
                + f"{initial_velocity.shape}."
            )
            raise BlockParameterError(
                message=message, system=self, parameter_name="initial_velocity"
            )

        initial_angular_velocity = npa.asarray(initial_angular_velocity)
        if initial_angular_velocity.shape != (3,):
            message = (
                "Initial angular velocity must have shape (3,), but has shape "
                + f"{initial_angular_velocity.shape}."
            )
            raise BlockParameterError(
                message=message, system=self, parameter_name="initial_angular_velocity"
            )

        return self.RigidBodyState(
            position=initial_position,
            orientation=initial_orientation,
            velocity=initial_velocity,
            angular_velocity=initial_angular_velocity,
        )

    @property
    def force_input(self):
        return self.input_ports[self.force_index]

    @property
    def torque_input(self):
        return self.input_ports[self.torque_index]

    @property
    def mass_input(self):
        if self.mass_index is None:
            return None
        return self.input_ports[self.mass_index]

    @property
    def inertia_input(self):
        if self.inertia_index is None:
            return None
        return self.input_ports[self.inertia_index]

    @property
    def position_output(self):
        return self.output_ports[self.pos_output_index]

    @property
    def orientation_output(self):
        return self.output_ports[self.orientation_output_index]

    @property
    def velocity_output(self):
        return self.output_ports[self.vel_output_index]

    @property
    def angular_velocity_output(self):
        return self.output_ports[self.ang_vel_output_index]

    def _configure_ports(
        self,
        initial_state,
        enable_external_mass,
        enable_external_inertia_matrix,
        enable_output_state_derivatives,
    ):
        # External force vector input
        self.force_index = self.declare_input_port(name="force_vector")

        # External torque vector input
        self.torque_index = self.declare_input_port(name="torque_vector")

        # External mass input
        self.mass_index = None
        if enable_external_mass:
            self.mass_index = self.declare_input_port(name="mass")

        # External inertia matrix input
        self.inertia_index = None
        if enable_external_inertia_matrix:
            self.inertia_index = self.declare_input_port(name="inertia_matrix")

        # Position output
        self.pos_output_index = self.declare_output_port(
            self._pos_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            default_value=initial_state.position,
            name=f"{self.name}:position",
        )

        # Orientation output
        self.orientation_output_index = self.declare_output_port(
            self._orientation_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            default_value=initial_state.orientation,
            name=f"{self.name}:orientation",
        )

        # Velocity output
        self.vel_output_index = self.declare_output_port(
            self._vel_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            default_value=initial_state.velocity,
            name=f"{self.name}:velocity",
        )

        # Angular velocity output
        self.ang_vel_output_index = self.declare_output_port(
            self._ang_vel_output,
            prerequisites_of_calc=[DependencyTicket.xc],
            requires_inputs=False,
            default_value=initial_state.angular_velocity,
            name=f"{self.name}:angular_velocity",
        )

        if enable_output_state_derivatives:
            self.pos_deriv_output_index = self.declare_output_port(
                self._pos_derivative,
                prerequisites_of_calc=[DependencyTicket.xc],
                requires_inputs=False,
                default_value=npa.zeros(3),
                name=f"{self.name}:position_dot",
            )

            self.orientation_deriv_output_index = self.declare_output_port(
                self._orientation_derivative,
                prerequisites_of_calc=[DependencyTicket.xc],
                requires_inputs=False,
                default_value=npa.zeros(3),
                name=f"{self.name}:orientation_dot",
            )

            force_ticket = self.input_ports[self.force_index].ticket
            self.vel_deriv_output_index = self.declare_output_port(
                self._vel_derivative,
                prerequisites_of_calc=[force_ticket, DependencyTicket.xc],
                requires_inputs=True,
                default_value=npa.zeros(3),
                name=f"{self.name}:velocity_dot",
            )

            torque_ticket = self.input_ports[self.torque_index].ticket
            self.ang_vel_deriv_output_index = self.declare_output_port(
                self._ang_vel_derivative,
                prerequisites_of_calc=[torque_ticket, DependencyTicket.xc],
                requires_inputs=True,
                default_value=npa.zeros(3),
                name=f"{self.name}:angular_velocity_dot",
            )

    def _pos_output(self, time, state, *inputs, **parameters):
        xc = state.continuous_state
        return xc.position

    def _orientation_output(self, time, state, *inputs, **parameters):
        xc = state.continuous_state
        return xc.orientation

    def _vel_output(self, time, state, *inputs, **parameters):
        xc = state.continuous_state
        return xc.velocity

    def _ang_vel_output(self, time, state, *inputs, **parameters):
        xc = state.continuous_state
        return xc.angular_velocity

    def _pos_derivative(self, time, state, *inputs, **parameters):
        # This function produces the inertial -> body rotation.  What we
        # want is to rotate the body-fixed velocity into the inertial frame,
        # so we need the transpose of this rotation matrix.
        xc = state.continuous_state
        C_BI = euler_to_dcm(xc.orientation)
        return C_BI.T @ xc.velocity

    def _orientation_derivative(self, time, state, *inputs, **parameters):
        # Matrix mapping angular velocity in the body-fixed frame to Euler rates
        xc = state.continuous_state
        H = euler_kinematics(xc.orientation)
        return H @ xc.angular_velocity

    def _vel_derivative(self, time, state, *inputs, **parameters):
        xc = state.continuous_state

        if self.mass_index is not None:
            m = inputs[self.mass_index]
        else:
            m = parameters["mass"]

        # Gravity vector in the inertial frame
        g_I = parameters["gravity_vector"]

        # Acceleration in body-fixed frame
        F_B = inputs[self.force_index]
        C_BI = euler_to_dcm(xc.orientation)
        a_B = F_B / m + C_BI @ g_I

        # Body-fixed acceleration is the inertial plus Coriolis terms
        return a_B - npa.cross(xc.angular_velocity, xc.velocity)

    def _ang_vel_derivative(self, time, state, *inputs, **parameters):
        xc = state.continuous_state

        if self.inertia_index is not None:
            J_B = inputs[self.inertia_index]
        else:
            J_B = parameters["inertia_matrix"]

        # Torque in body-fixed frame
        tau_B = inputs[self.torque_index]

        wJw = npa.cross(xc.angular_velocity, J_B @ xc.angular_velocity)
        return npa.linalg.solve(J_B, tau_B - wJw)

    def _state_derivative(self, time, state, *inputs, **parameters):
        # See Eq. (1.7-18) in Lewis, Johnson, Stevens
        args = (time, state, *inputs)
        return self.RigidBodyState(
            position=self._pos_derivative(*args, **parameters),
            orientation=self._orientation_derivative(*args, **parameters),
            velocity=self._vel_derivative(*args, **parameters),
            angular_velocity=self._ang_vel_derivative(*args, **parameters),
        )

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        force = self.input_ports[self.force_index].eval(context)
        torque = self.input_ports[self.torque_index].eval(context)

        with ErrorCollector.context(error_collector):
            if force.shape != (3,):
                raise ShapeMismatchError(
                    system=self,
                    expected_shape=(3,),
                    actual_shape=force.shape,
                )

            if torque.shape != (3,):
                raise ShapeMismatchError(
                    system=self,
                    expected_shape=(3,),
                    actual_shape=torque.shape,
                )

        if self.mass_index is not None:
            mass = self.input_ports[self.mass_index].eval(context)

            with ErrorCollector.context(error_collector):
                if mass.shape != ():
                    raise ShapeMismatchError(
                        system=self,
                        expected_shape=(),
                        actual_shape=mass.shape,
                    )

        if self.inertia_index is not None:
            inertia = self.input_ports[self.inertia_index].eval(context)

            with ErrorCollector.context(error_collector):
                if inertia.shape != (3, 3):
                    raise ShapeMismatchError(
                        system=self,
                        expected_shape=(3, 3),
                        actual_shape=inertia.shape,
                    )

Ros2Publisher

Bases: LeafSystem

Ros2Publisher block can emit signals to a ROS2 topic, based on input signal data.

Source code in jaxonomy/library/ros2.py

class Ros2Publisher(LeafSystem):
    """
    Ros2Publisher block can emit signals to a ROS2 topic, based on input signal data.
    """

    @parameters(static=["topic", "msg_type", "fields"])
    def __init__(
        self,
        dt: float,
        topic: str,
        msg_type: type,
        fields: dict[str, type],
        **kwargs,
    ):
        """
        Publish messages to a ROS2 topic.

        Args:
            dt: Period of the system, in both sim and real (ros2) time.
            topic: ROS2 topic to publish to. Eg. `/turtle1/cmd_vel`.
            msg_type: ROS2 message type, e.g. `Twist` from `geometry_msgs.msg`.
                      Unlike the corresponding UI parameter, this must be a Python
                      type object.
            fields: Ordered dictionary of default values to extract from the
                    received message. The keys are the full attribute path
                    (with dots) to the value in the message, and the values are
                    the default values. This is used to create the output ports
                    with valid data types. Use Python or Numpy data types, not JAX.

                    For instance, for a `geometry_msgs.msg.Twist` message, the
                    `fields` could be `{"linear.x": float, "angular.z": float}`.
        """

        super().__init__(**kwargs)
        self.logger = logger.getChild("Ros2Publisher:" + self.name_path_str)

        self.node: Node = None
        self.publisher: Publisher = None

        self.dt = dt
        self.topic = topic
        self.msg_type = msg_type
        self.fields = {field: _fixup_dtype(dtype) for field, dtype in fields.items()}

        self.declare_periodic_update(self._update, period=dt, offset=0.0)

        # Extract type & full attribute path from fields. Note that this
        # relies on the fact that the Python (3.7+) dict is ordered; The
        # order must match that of the input ports. Works well with JSON
        # because our I/O ports are ordered arrays.
        # This could likely be simplified / replaced with a Bus signal type.
        self.input_types = []  # [float, float]
        self.input_attr_path = []  # ["linear.x", "angular.z"]
        for msg_field_name, msg_field_type in self.fields.items():
            input_name = _attr2name(msg_field_name)
            self.declare_input_port(name=input_name)
            self.input_types.append(msg_field_type)
            self.input_attr_path.append(msg_field_name)

        self.pre_simulation_initialize()

    def __del__(self):
        self.post_simulation_finalize()

    def pre_simulation_initialize(self):
        if not _ros2_init():
            raise RuntimeError("ROS2 init failed")

        node_name = _NODE_NAME_REGEX.sub("_", self.name_path_str)
        rnd = np.random.randint(0, 1000)
        self.node = rclpy.create_node(f"jaxonomy_{rnd}_" + node_name)
        self.publisher = self.node.create_publisher(
            self.msg_type, self.topic, qos_profile=10
        )

        self.logger.debug(
            "ROS2 publisher %s initialized with node: %s and publisher: %s",
            self.name_path_str,
            self.node,
            self.publisher,
        )

    def post_simulation_finalize(self) -> None:
        if self.node:
            self.logger.debug("ROS2 publisher %s clean up", self.name_path_str)
            self.node.destroy_publisher(self.publisher)
            self.publisher = None
            self.node.destroy_node()
            self.node = None
            _ros2_shutdown()

    def _update(self, time, state, *inputs, **params):
        return io_callback(self._publish_message, None, *inputs)

    def _publish_message(self, *inputs):
        msg = self.msg_type()

        for i, input_value in enumerate(inputs):
            value = self.input_types[i](input_value)
            _setattr_path(msg, self.input_attr_path[i], value)

        self.logger.debug("Publishing message to topic %s: %s", self.topic, msg)
        self.publisher.publish(msg)

        # Spin rclpy loop to ensure the message is sent. Also, sync the clocks
        # using dt. This is a bit of a hack for now until we have proper clock
        # synchronization.
        rclpy.spin_once(self.node, timeout_sec=self.dt)

__init__(dt, topic, msg_type, fields, **kwargs)

Publish messages to a ROS2 topic.

Parameters:

Name	Type	Description	Default
dt	float	Period of the system, in both sim and real (ros2) time.	required
topic	str	ROS2 topic to publish to. Eg. /turtle1/cmd_vel.	required
msg_type	type	ROS2 message type, e.g. Twist from geometry_msgs.msg. Unlike the corresponding UI parameter, this must be a Python type object.	required
fields	dict[str, type]	Ordered dictionary of default values to extract from the received message. The keys are the full attribute path (with dots) to the value in the message, and the values are the default values. This is used to create the output ports with valid data types. Use Python or Numpy data types, not JAX. For instance, for a `geometry_msgs.msg.Twist` message, the `fields` could be `{"linear.x": float, "angular.z": float}`.	required

Source code in jaxonomy/library/ros2.py

@parameters(static=["topic", "msg_type", "fields"])
def __init__(
    self,
    dt: float,
    topic: str,
    msg_type: type,
    fields: dict[str, type],
    **kwargs,
):
    """
    Publish messages to a ROS2 topic.

    Args:
        dt: Period of the system, in both sim and real (ros2) time.
        topic: ROS2 topic to publish to. Eg. `/turtle1/cmd_vel`.
        msg_type: ROS2 message type, e.g. `Twist` from `geometry_msgs.msg`.
                  Unlike the corresponding UI parameter, this must be a Python
                  type object.
        fields: Ordered dictionary of default values to extract from the
                received message. The keys are the full attribute path
                (with dots) to the value in the message, and the values are
                the default values. This is used to create the output ports
                with valid data types. Use Python or Numpy data types, not JAX.

                For instance, for a `geometry_msgs.msg.Twist` message, the
                `fields` could be `{"linear.x": float, "angular.z": float}`.
    """

    super().__init__(**kwargs)
    self.logger = logger.getChild("Ros2Publisher:" + self.name_path_str)

    self.node: Node = None
    self.publisher: Publisher = None

    self.dt = dt
    self.topic = topic
    self.msg_type = msg_type
    self.fields = {field: _fixup_dtype(dtype) for field, dtype in fields.items()}

    self.declare_periodic_update(self._update, period=dt, offset=0.0)

    # Extract type & full attribute path from fields. Note that this
    # relies on the fact that the Python (3.7+) dict is ordered; The
    # order must match that of the input ports. Works well with JSON
    # because our I/O ports are ordered arrays.
    # This could likely be simplified / replaced with a Bus signal type.
    self.input_types = []  # [float, float]
    self.input_attr_path = []  # ["linear.x", "angular.z"]
    for msg_field_name, msg_field_type in self.fields.items():
        input_name = _attr2name(msg_field_name)
        self.declare_input_port(name=input_name)
        self.input_types.append(msg_field_type)
        self.input_attr_path.append(msg_field_name)

    self.pre_simulation_initialize()

Ros2Subscriber

Bases: LeafSystem

Ros2Subscriber block listens to messages over a ROS2 topic and outputs them as signals in jaxonomy.

Source code in jaxonomy/library/ros2.py

class Ros2Subscriber(LeafSystem):
    """
    Ros2Subscriber block listens to messages over a ROS2 topic and outputs them as
    signals in jaxonomy.
    """

    @parameters(static=["topic", "msg_type", "fields", "read_before_start"])
    def __init__(
        self,
        dt,
        topic: str,
        msg_type: type,
        fields: dict[str, type],
        read_before_start=True,
        **kwargs,
    ):
        """Subscribe to a ROS2 topic and extract message values to output ports.

        Args:
            dt: Period of the system, in both sim and real (ros2) time.
            topic: ROS2 topic to subscribe to. Eg. `/turtle1/pose`.
            msg_type: ROS2 message type, e.g. `Pose` from `turtlesim.msg`.
                      Unlike the corresponding UI parameter, this must be a Python
                      type object.
            fields: Ordered dictionary of default values to extract from the
                    received message. The keys are the full attribute path
                    (with dots) to the value in the message, and the values are
                    the default values. This is used to create the output ports
                    with valid data types. Use Python or Numpy data types, not JAX.

                    For instance, for a `geometry_msgs.msg.Twist` message, the
                    `fields` could be `{"linear.x": float, "angular.z": float}`.
            read_before_start: If True, the subscriber will read the first message
                    before the simulation starts. Otherwise, the initial outputs will
                    be 0.
        """

        super().__init__(**kwargs)
        self.logger = logger.getChild("Ros2Subscriber:" + self.name_path_str)

        self.node: Node = None
        self.subscription: Subscription = None
        self._last_msg = None

        if not _ros2_init():
            raise RuntimeError("ROS2 init failed")

        self.dt = dt
        self.msg_type = msg_type
        self.topic = topic
        self.fields = {field: _fixup_dtype(dtype) for field, dtype in fields.items()}
        self.read_before_start = read_before_start

        # Note: Not 100% sure this is absolutely valid, but it worked with JAX.
        # If somehow we aren't getting updates, we may need to create a cache index,
        # see custom.py. See _callback().
        self.declare_periodic_update(self._update, period=dt, offset=0.0)
        self.default_values = {field: dtype() for field, dtype in self.fields.items()}

        def _make_output_cb(field_name: str, dtype: type):
            def _output():
                last_msg = self._last_msg or self.default_values
                value = _getattr_path(last_msg, field_name)
                return dtype(value)

            def _io_cb(time, state, *inputs, **params):
                return io_callback(_output, npa.asarray(_output()))

            return _io_cb

        for field, dtype in self.fields.items():
            self.declare_output_port(
                callback=_make_output_cb(field, dtype),
                name=_attr2name(field),
                prerequisites_of_calc=[],
                requires_inputs=False,
                period=dt,
                offset=0.0,
                default_value=self.default_values[field],
            )

        self.pre_simulation_initialize()

    def __del__(self):
        self.post_simulation_finalize()

    def pre_simulation_initialize(self):
        if not _ros2_init():
            raise RuntimeError("ROS2 init failed")

        node_name = _NODE_NAME_REGEX.sub("_", self.name_path_str)
        rnd = np.random.randint(0, 1000)
        self.node = rclpy.create_node(f"jaxonomy_{rnd}_" + node_name)
        self.subscription = self.node.create_subscription(
            self.msg_type, self.topic, self._callback, qos_profile=10
        )
        self.logger.debug(
            "ROS2 subscriber %s initialized, listening on topic %s msg_type=%s",
            self.name_path_str,
            self.topic,
            self.msg_type,
        )

        if self.read_before_start:
            self._update_cb()

    def post_simulation_finalize(self) -> None:
        if self.node:
            self.logger.debug("ROS2 subscriber %s clean up", self.name_path_str)
            self.node.destroy_subscription(self.subscription)
            self.subscription = None
            self.node.destroy_node()
            self.node = None
            _ros2_shutdown()

    def _update(self, time, state, *inputs, **params):
        return io_callback(self._update_cb, None)

    def _update_cb(self):
        # This timeout does not seem to block the call
        rclpy.spin_once(self.node, timeout_sec=2.0)

    def _callback(self, msg):
        self.logger.debug("Received message on topic %s: %s", self.topic, msg)

        # This may be wrong because we're not cleanly using the cache
        # like in custom.py. But it works.
        self._last_msg = msg

__init__(dt, topic, msg_type, fields, read_before_start=True, **kwargs)

Subscribe to a ROS2 topic and extract message values to output ports.

Parameters:

Name	Type	Description	Default
dt		Period of the system, in both sim and real (ros2) time.	required
topic	str	ROS2 topic to subscribe to. Eg. /turtle1/pose.	required
msg_type	type	ROS2 message type, e.g. Pose from turtlesim.msg. Unlike the corresponding UI parameter, this must be a Python type object.	required
fields	dict[str, type]	Ordered dictionary of default values to extract from the received message. The keys are the full attribute path (with dots) to the value in the message, and the values are the default values. This is used to create the output ports with valid data types. Use Python or Numpy data types, not JAX. For instance, for a `geometry_msgs.msg.Twist` message, the `fields` could be `{"linear.x": float, "angular.z": float}`.	required
read_before_start		If True, the subscriber will read the first message before the simulation starts. Otherwise, the initial outputs will be 0.	True

Source code in jaxonomy/library/ros2.py

@parameters(static=["topic", "msg_type", "fields", "read_before_start"])
def __init__(
    self,
    dt,
    topic: str,
    msg_type: type,
    fields: dict[str, type],
    read_before_start=True,
    **kwargs,
):
    """Subscribe to a ROS2 topic and extract message values to output ports.

    Args:
        dt: Period of the system, in both sim and real (ros2) time.
        topic: ROS2 topic to subscribe to. Eg. `/turtle1/pose`.
        msg_type: ROS2 message type, e.g. `Pose` from `turtlesim.msg`.
                  Unlike the corresponding UI parameter, this must be a Python
                  type object.
        fields: Ordered dictionary of default values to extract from the
                received message. The keys are the full attribute path
                (with dots) to the value in the message, and the values are
                the default values. This is used to create the output ports
                with valid data types. Use Python or Numpy data types, not JAX.

                For instance, for a `geometry_msgs.msg.Twist` message, the
                `fields` could be `{"linear.x": float, "angular.z": float}`.
        read_before_start: If True, the subscriber will read the first message
                before the simulation starts. Otherwise, the initial outputs will
                be 0.
    """

    super().__init__(**kwargs)
    self.logger = logger.getChild("Ros2Subscriber:" + self.name_path_str)

    self.node: Node = None
    self.subscription: Subscription = None
    self._last_msg = None

    if not _ros2_init():
        raise RuntimeError("ROS2 init failed")

    self.dt = dt
    self.msg_type = msg_type
    self.topic = topic
    self.fields = {field: _fixup_dtype(dtype) for field, dtype in fields.items()}
    self.read_before_start = read_before_start

    # Note: Not 100% sure this is absolutely valid, but it worked with JAX.
    # If somehow we aren't getting updates, we may need to create a cache index,
    # see custom.py. See _callback().
    self.declare_periodic_update(self._update, period=dt, offset=0.0)
    self.default_values = {field: dtype() for field, dtype in self.fields.items()}

    def _make_output_cb(field_name: str, dtype: type):
        def _output():
            last_msg = self._last_msg or self.default_values
            value = _getattr_path(last_msg, field_name)
            return dtype(value)

        def _io_cb(time, state, *inputs, **params):
            return io_callback(_output, npa.asarray(_output()))

        return _io_cb

    for field, dtype in self.fields.items():
        self.declare_output_port(
            callback=_make_output_cb(field, dtype),
            name=_attr2name(field),
            prerequisites_of_calc=[],
            requires_inputs=False,
            period=dt,
            offset=0.0,
            default_value=self.default_values[field],
        )

    self.pre_simulation_initialize()

Sawtooth

Bases: SourceBlock

Produces a modulated linear sawtooth signal.

The signal is similar to: https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.sawtooth.html

Given amplitude a, period p, and phase delay phi, the output signal is:

    y(t) = a * ((t - phi) % p)

where % is the modulo operator.

Input ports

None

Output ports

(0) The sawtooth signal.

Source code in jaxonomy/library/sources.py

class Sawtooth(SourceBlock):
    """Produces a modulated linear sawtooth signal.

    The signal is similar to:
    https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.sawtooth.html

    Given amplitude `a`, period `p`, and phase delay `phi`, the output signal is:
    ```
        y(t) = a * ((t - phi) % p)
    ```
    where `%` is the modulo operator.

    Input ports:
        None

    Output ports:
        (0) The sawtooth signal.
    """

    # `frequency` is set as a static parameter because it reconfigures the periodic
    # update when initialize() is called which would break optimization and
    # ensemble because they don't re-create the context and therefore won't call
    # initialize() if `frequency` is updated.
    @parameters(dynamic=["amplitude", "phase_delay"], static=["frequency"])
    def __init__(self, amplitude=1.0, frequency=0.5, phase_delay=1.0, **kwargs):
        super().__init__(self._func, **kwargs)

        # Initialize the floating-point tolerance.  This will be machine epsilon
        # for the floating point type of the time variable (determined in the
        # static initialization step).
        self.eps = 0.0
        self._periodic_update_idx = self.declare_periodic_update()

    def initialize(self, amplitude, frequency, phase_delay):
        # Add a dummy event so that the ODE solver doesn't try to integrate through
        # the discontinuity.
        self.declare_discrete_state(default_value=False)

        self.period = 1 / frequency
        self.configure_periodic_update(
            self._periodic_update_idx,
            lambda *args, **kwargs: True,
            period=self.period,
            offset=phase_delay,
        )

    def _func(self, time, **parameters):
        # np.mod((t - phase_delay), (1.0 / frequency)) * amplitude
        period_fraction = npa.mod(
            time - parameters["phase_delay"] + self.eps, self.period
        )
        return period_fraction * parameters["amplitude"]

    def initialize_static_data(self, context):
        # Determine machine epsilon for the type of the time variable
        self.eps = 2 * npa.finfo(npa.result_type(context.time)).eps
        return super().initialize_static_data(context)

ScalarBroadcast

Bases: FeedthroughBlock

Broadcast a scalar to a vector or matrix.

Given a scalar input u and dimensions m and n, this block will return a vector or matrix of shape (m, n) with all elements equal to u.

Input ports

(0) The scalar input signal.

Output ports

(0) The broadcasted output signal.

Parameters:

Name	Type	Description	Default
m		The number of rows in the output matrix. If m is None, then the output will be a vector with shape (n,). To get a row vector of size (1,n), set m=1 expliclty.	required
n		The number of columns in the output matrix. If n is None, then the output will be a vector with shape (m,). To get a column vector of size (m,1), set n=1 expliclty.	required

Source code in jaxonomy/library/math_ops.py

class ScalarBroadcast(FeedthroughBlock):
    """Broadcast a scalar to a vector or matrix.

    Given a scalar input `u` and dimensions `m` and `n`, this block will return
    a vector or matrix of shape `(m, n)` with all elements equal to `u`.

    Input ports:
        (0) The scalar input signal.

    Output ports:
        (0) The broadcasted output signal.

    Parameters:
        m:
            The number of rows in the output matrix.  If `m` is None, then the output
            will be a vector with shape `(n,)`. To get a row vector of size `(1,n)`,
            set `m=1` expliclty.
        n:
            The number of columns in the output matrix.  If `n` is None, then the
            output will be a vector with shape `(m,)`. To get a column vector of size
            `(m,1)`, set `n=1` expliclty.
    """

    @parameters(static=["m", "n"])
    def __init__(self, m, n, **kwargs):
        super().__init__(None, **kwargs)

    def initialize(self, m, n):
        if m is not None:
            m = int(m)
        else:
            m = 0
        if n is not None:
            n = int(n)
        else:
            n = 0

        if m > 0 and n > 0:
            ones_ = npa.ones((m, n))
        elif m > 0:
            ones_ = npa.ones((m,))
        elif n > 0:
            ones_ = npa.ones((n,))
        else:
            raise BlockParameterError(
                message=f"ScalarBroadcast block {self.name} at least m or n must not be None or Zero"
            )
        self.replace_op(lambda x: ones_ * x)

ShiftRegister

Bases: LeafSystem

Fixed-length shift register delay line.

Delays an input signal by exactly n_steps discrete timesteps. Output at time t is the input value from n_steps timesteps ago.

Parameters:

Name	Type	Description	Default
n_steps	int	Number of steps to delay. STATIC — set at construction, cannot be changed at runtime. Must be >= 1.	required
signal_shape	tuple	Shape of each signal frame. Use () for scalar, (3,) for 3-vector, etc.	()
initial_value	array - like	Value to fill the buffer with before any input has been received. Default: zeros.	None
dt	float	Discrete update interval in seconds.	0.01

Ports

Input[0] "u": signal to delay, shape=signal_shape Output[0] "y": delayed signal, shape=signal_shape

Source code in jaxonomy/library/delay.py

class ShiftRegister(LeafSystem):
    """
    Fixed-length shift register delay line.

    Delays an input signal by exactly n_steps discrete 
    timesteps. Output at time t is the input value from 
    n_steps timesteps ago.

    Parameters:
        n_steps (int): Number of steps to delay. 
            STATIC — set at construction, cannot be changed 
            at runtime. Must be >= 1.
        signal_shape (tuple): Shape of each signal frame.
            Use () for scalar, (3,) for 3-vector, etc.
        initial_value (array-like): Value to fill the buffer 
            with before any input has been received.
            Default: zeros.
        dt (float): Discrete update interval in seconds.

    Ports:
        Input[0] "u": signal to delay, shape=signal_shape
        Output[0] "y": delayed signal, shape=signal_shape
    """

    @parameters(static=["n_steps", "signal_shape"])
    def __init__(
        self, 
        n_steps: int, 
        signal_shape: tuple = (), 
        initial_value=None, 
        dt: float = 0.01, 
        **kwargs
    ):
        super().__init__(**kwargs)
        self.dt = dt
        self.n_steps = n_steps
        self.signal_shape = signal_shape

        if initial_value is None:
            self.initial_value = npa.zeros(signal_shape)
        else:
            self.initial_value = npa.array(initial_value)

        self.input_idx = self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, n_steps, signal_shape):
        buffer = npa.broadcast_to(self.initial_value, (n_steps, *signal_shape))
        self.declare_discrete_state(default_value=buffer)

        self.configure_periodic_update(
            self._periodic_update_idx, 
            self._update, 
            period=self.dt, 
            offset=self.dt
        )

        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=0.0,
            requires_inputs=False,
            prerequisites_of_calc=[DependencyTicket.xd],
            default_value=buffer[-1],
        )

    def _update(self, _time, state, *inputs, **_params):
        u = inputs[self.input_idx]
        buffer = npa.roll(state.discrete_state, shift=1, axis=0)
        buffer = buffer.at[0].set(u)
        return buffer

    def _output(self, _time, state, **_params):
        return state.discrete_state[-1]

SignalDatatypeConversion

Bases: FeedthroughBlock

Convert the input signal to a different data type. Input ports: (0) The input signal. Output ports: (0) The input signal converted to the specified data type. Parameters: dtype: The data type to which the input signal is converted. Must be a valid NumPy data type, e.g. "float32", "int64", etc.

Source code in jaxonomy/library/routing.py

class SignalDatatypeConversion(FeedthroughBlock):
    """Convert the input signal to a different data type.
    Input ports:
        (0) The input signal.
    Output ports:
        (0) The input signal converted to the specified data type.
    Parameters:
        dtype:
            The data type to which the input signal is converted.  Must be a valid
            NumPy data type, e.g. "float32", "int64", etc.
    """

    def _op(self, dtype, x):
        # This check makes the numpy backend strict like jax
        if npa.active_backend == "numpy" and isinstance(x, (list, tuple)):
            raise ValueError(
                "SignalDatatypeConversion block does not support list or tuple inputs."
            )

        return cond(
            isinstance(x, npa.ndarray),
            lambda x: npa.astype(x, dtype),
            lambda x: npa.array(x, dtype),
            x,
        )

    @parameters(static=["convert_to_type"])
    def __init__(self, convert_to_type, *args, **kwargs):
        super().__init__(partial(self._op, np.dtype(convert_to_type)), *args, **kwargs)

    def initialize(self, convert_to_type):
        self.dtype = np.dtype(convert_to_type)
        self.replace_op(partial(self._op, np.dtype(convert_to_type)))

SimulationResultsSource

Bases: LeafSystem

Replays one recorded trajectory from :class:~jaxonomy.simulation.types.SimulationResults.

Output port y is the signal value at the current simulation time, using linear interpolation or zero-order hold. Values clamp to the first/last sample outside the recorded time range (jnp.interp semantics for linear mode).

Source code in jaxonomy/library/data_source.py

class SimulationResultsSource(LeafSystem):
    """Replays one recorded trajectory from :class:`~jaxonomy.simulation.types.SimulationResults`.

    Output port ``y`` is the signal value at the current simulation time, using linear
    interpolation or zero-order hold. Values clamp to the first/last sample outside
    the recorded time range (``jnp.interp`` semantics for linear mode).
    """

    def __init__(
        self,
        results: "SimulationResults",
        signal_name: str,
        interpolation: str = "linear",
        **kwargs,
    ):
        from jaxonomy.simulation.types import SimulationResults as SR

        if not isinstance(results, SR):
            raise TypeError(f"results must be SimulationResults, got {type(results)}")
        if results.outputs is None:
            raise ValueError("SimulationResults.outputs is None; cannot replay.")
        if signal_name not in results.outputs:
            raise KeyError(f"signal_name {signal_name!r} not in results.outputs")

        super().__init__(**kwargs)

        if interpolation not in ("linear", "zero_order_hold"):
            raise ValueError(
                f"interpolation must be 'linear' or 'zero_order_hold', got {interpolation!r}"
            )

        t = np.asarray(results.time, dtype=np.float64).reshape(-1)
        y = np.asarray(results.outputs[signal_name], dtype=np.float64)
        if y.ndim > 1 and y.shape[-1] == 1:
            y = y.reshape(-1)
        if y.ndim != 1:
            raise ValueError(
                f"Replayed signal {signal_name!r} must be 1-D per time step; got shape {y.shape}"
            )
        if t.shape[0] != y.shape[0]:
            raise ValueError(
                f"time length {t.shape[0]} != signal length {y.shape[0]} for {signal_name!r}"
            )

        self._t = npa.array(t, dtype=npa.float64)
        self._y = npa.array(y, dtype=npa.float64)

        def linear(time):
            return npa.interp(time, self._t, self._y)

        def zoh(time):
            if self._t.shape[0] == 0:
                return npa.array(0.0, dtype=self._y.dtype)
            if self._t.shape[0] == 1:
                return self._y[0]
            tc = npa.clip(time, self._t[0], self._t[-1])
            idx = npa.searchsorted(self._t, tc, side="right") - 1
            idx = npa.maximum(idx, 0)
            return self._y[idx]

        self._interp = linear if interpolation == "linear" else zoh
        self._jit_interp = npa.jit(self._interp)

        def _out_cb(time, state, *inputs, **parameters):
            return self._jit_interp(time)

        self.declare_output_port(
            _out_cb,
            name="y",
            prerequisites_of_calc=[DependencyTicket.time],
            requires_inputs=False,
        )

Sindy

Bases: LeafSystem

This class implements System Identification (SINDy) algorithm with or without control inputs for contiuous-time and discrete-time systems.

The learned continuous-time dynamical system model will be of the form:

dx/dt = f(x, u)

where x is the state vector and u is the optional control input vector. The block will output the full state vector x of the system.

The learned discrete-time dynamical system model will be of the form:

x_{k+1} = f(x_k, u_k)

where x_k is the state vector at time step k and u_k is the optional control vector. The block will update the output to x_k at an interval provided by the parameter discrete_time_update_interval.

Input ports

(0) u: control vector for the system. This port is only available if the Sindy model is trained with control inputs, i.e. control_input_columns is not None during training.

Output ports

(0) x: full state output of the system.

Parameters:

Name	Type	Description	Default
file_name	str	Path to the CSV file containing training data.	None
header_as_first_row	bool	If True, the first row of the CSV file is treated as the header.	False
state_columns	int | str | list[int] | list[str]	For training, either one of the following for CSV columns representing state variables x: - a string or integer (for a single column) - a list of strings or integers (for multiple columns) - a string representing a slice of columns, e.g. '0:3'	1
control_input_columns	int | str | list[int] | list[str]	For training, either one of the following for CSV columns representing control inputs u: - a string or integer (for a single column) - a list of strings or integers (for multiple columns) - a string representing a slice of columns, e.g. '0:3' If None, then the SINDy model will be trained without control inputs.	None
dt	float	Fixed value of dt if rows of the CSV file represent equidistant time steps.	None
time_column	(str, int)	Column name (str) for column index (int) for time data t. If time_column is provided, then fixed dt above will be ignored. If neither dt nor time_column is provided, then the SINDy model will use a fixed detault time step of dt=1.	None
state_derivatives_columns	int | str | list[int] | list[str]	For training, either one of the following for csv columns representing state derivatives x_dot: - a string or integer (for a single column) - a list of strings or integers (for multiple columns) - a string representing a slice of columns, e.g. '0:3' This field is optional. If provided, the SINDy model will estimate directly use these state derivatives for training. If not provided, the SINDy model will approximate the state derivatives dot_x = dx/dt from x by using the specified differentiation_method.	None
discrete_time	bool	If True, the SINDy model will be trained for discrete-time systems. In this case, the dynamical system is treated as a map. Rather than predicting derivatives, the right hand side functions step the system forward by one time step. If False, dynamical system is assumed to be a flow (right-hand side functions predict continuous time derivatives). See documentation for pysindy.	False
differentiation_method	str	Method to use for differentiating the state data x to obtain state derivatives dot_x = dx/dt. Available options are: 'centered difference' (default)	'centered difference'
threshold	float	Threshold for the Sequentially thresholded least squares (STLSQ) algorithm used for training SINDy model.	0.1
alpha	float	Regularization strength for the STLSQ algorithm.	0.05
max_iter	int	Maximum number of iterations for the STLSQ algorithm.	20
normalize_columns	bool	If True, normalize the columns of the data matrix before regression.	False
poly_order	int	Degree of polynomial features. Set to None to omit this library.	2
fourier_n_frequencies	int	Number of Fourier frequencies. Set to None to omit this library.	None
custom_basis_functions	list of functions	A list of custom basis functions to use for training the SINDy model. For example to include f(x) = 1/x and g(x) = exp(-x), provide [lambda x: 1.0/(x.0 + 1e-06), lamda x: jnp.exp(-x)] Currently only supported for jaxonomy interface. Calls from UI and pretrained model loading does not support custom basis functions.	None
pretrained	bool	If True, use a pretrained model specified by the pretrained_file_path argument.	False
pretrained_file_path	str	Path to the pretrained model file.	None
initial_state	ndarray	Initial state of the system for propagating the continuous-time or discrete-time system forward duiring simulation.	None
discrete_time_update_interval	float	Interval at which the discrete-time model should be updated. Default is 1.0.	1.0
equations	list of strings	(For internal UI use only) The identified system equations.	None
base_feature_names	list of strings	(For internal UI use only) Features x_i and u_i.	None
feature_names	list of strings	(For internal UI use only) Composed features with basis libraries.	None
coefficients	ndarray	(For internal UI use only) Coefficients of the identified model.	None
has_control_input	bool	(For internal UI use only) If True, the model was trained with control. For standard training from CSV file, this is inferred from the parameter control_input_columns.	True

Source code in jaxonomy/library/sindy.py

class Sindy(LeafSystem):
    """
    This class implements System Identification (SINDy) algorithm with or without
    control inputs for contiuous-time and discrete-time systems.

    The learned continuous-time dynamical system model will be of the form:

    ```
    dx/dt = f(x, u)
    ```
    where `x` is the state vector and `u` is the optional control input vector. The
    block will output the full state vector `x` of the system.

    The learned discrete-time dynamical system model will be of the form:

    ```
    x_{k+1} = f(x_k, u_k)
    ```

    where `x_k` is the state vector at time step `k` and `u_k` is the optional control
    vector. The block will update the output to `x_k` at an interval provided by the
    parameter `discrete_time_update_interval`.

    Input ports:
        (0) u: control vector for the system. This port is only available if the Sindy
            model is trained with control inputs, i.e. `control_input_columns` is not
            `None` during training.

    Output ports:
        (0) x: full state output of the system.

    Parameters:
        file_name (str):
            Path to the CSV file containing training data.

        header_as_first_row (bool):
            If True, the first row of the CSV file is treated as the header.

        state_columns (int | str | list[int] | list[str]):
            For training, either one of the following for CSV columns representing
            state variables `x`:
                - a string or integer (for a single column)
                - a list of strings or integers (for multiple columns)
                - a string representing a slice of columns, e.g. '0:3'


        control_input_columns (int | str | list[int] | list[str]):
            For training, either one of the following for CSV columns representing
            control inputs `u`:
                - a string or integer (for a single column)
                - a list of strings or integers (for multiple columns)
                - a string representing a slice of columns, e.g. '0:3'
            If None, then the SINDy model will be trained without control inputs.

        dt (float):
            Fixed value of dt if rows of the CSV file represent equidistant time steps.

        time_column (str, int):
            Column name (str) for column index (int) for time data `t`.
            If `time_column` is provided, then fixed `dt` above will be ignored.
            If neither `dt` nor `time_column` is provided, then the SINDy model will
            use a fixed detault time step of `dt=1`.

        state_derivatives_columns (int | str | list[int] | list[str]):
            For training, either one of the following for csv columns representing
            state derivatives `x_dot`:
                - a string or integer (for a single column)
                - a list of strings or integers (for multiple columns)
                - a string representing a slice of columns, e.g. '0:3'
            This field is optional. If provided, the SINDy model will estimate directly
            use these state derivatives for training. If not provided, the SINDy model
            will approximate the state derivatives `dot_x = dx/dt` from `x` by using
            the specified `differentiation_method`.

        discrete_time (bool):
            If True, the SINDy model will be trained for discrete-time systems. In
            this case, the dynamical system is treated as a map. Rather than
            predicting derivatives, the right hand side functions step the system
            forward by one time step. If False, dynamical system is assumed to be a
            flow (right-hand side functions predict continuous time derivatives).
            See documentation for `pysindy`.

        differentiation_method (str):
            Method to use for differentiating the state data `x` to obtain state
            derivatives `dot_x = dx/dt`. Available options are:
                'centered difference' (default)

        threshold (float):
            Threshold for the Sequentially thresholded least squares (STLSQ) algorithm
            used for training SINDy model.

        alpha (float):
            Regularization strength for the STLSQ algorithm.

        max_iter (int):
            Maximum number of iterations for the STLSQ algorithm.

        normalize_columns (bool):
            If True, normalize the columns of the data matrix before regression.

        poly_order (int):
            Degree of polynomial features. Set to `None` to omit this library.

        fourier_n_frequencies (int):
            Number of Fourier frequencies. Set to `None` to omit this library.

        custom_basis_functions (list of functions):
            A list of custom basis functions to use for training the SINDy model.
            For example to include `f(x) = 1/x` and `g(x) = exp(-x)`,
            provide [lambda x: 1.0/(x.0 + 1e-06), lamda x: jnp.exp(-x)]

            Currently only supported for jaxonomy interface. Calls from UI and
            pretrained model loading does not support custom basis functions.

        pretrained (bool):
            If True, use a pretrained model specified by the `pretrained_file_path`
            argument.

        pretrained_file_path (str, optional): Path to the pretrained model file.

        initial_state (ndarray):
                Initial state of the system for propagating the continuous-time
                or discrete-time system forward duiring simulation.

        discrete_time_update_interval (float):
            Interval at which the discrete-time model should be updated. Default
            is 1.0.

        equations (list of strings):
            (For internal UI use only) The identified system equations.

        base_feature_names (list of strings):
            (For internal UI use only) Features x_i and u_i.

        feature_names (list of strings):
            (For internal UI use only) Composed features with basis libraries.

        coefficients (ndarray):
            (For internal UI use only) Coefficients of the identified model.

        has_control_input (bool):
            (For internal UI use only) If True, the model was trained with control.
            For standard training from CSV file, this is inferred from the
            parameter `control_input_columns`.
    """

    @parameters(
        static=[
            "file_name",
            "header_as_first_row",
            "state_columns",
            "control_input_columns",
            "discrete_time",
            "dt",
            "time_column",
            "state_derivatives_columns",
            "differentiation_method",
            "threshold",
            "alpha",
            "max_iter",
            "normalize_columns",
            "poly_order",
            "fourier_n_frequencies",
            "pretrained",
            "equations",
            "discrete_time_update_interval",
            "pretrained_file_path",
            "coefficients",
            "base_feature_names",
            "feature_names",
            "has_control_input",
            "initial_state",
        ],
    )
    def __init__(
        self,
        file_name=None,
        header_as_first_row=False,
        state_columns=1,
        control_input_columns=None,
        dt=None,
        time_column=None,
        state_derivatives_columns=None,
        discrete_time=False,
        differentiation_method="centered difference",
        # optimizer parameters
        threshold=0.1,
        alpha=0.05,
        max_iter=20,
        normalize_columns=False,
        # Library parameters
        poly_order=2,
        fourier_n_frequencies=None,
        custom_basis_functions=None,
        pretrained=False,
        pretrained_file_path=None,
        # for parameters obtained from UI training
        equations=None,
        base_feature_names=None,
        feature_names=None,
        coefficients=None,
        has_control_input=True,
        # Simulation parameters
        initial_state=None,
        discrete_time_update_interval=1.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        _validate_leafsystem_inputs(
            pretrained,
            pretrained_file_path,
            dt,
            time_column,
            poly_order,
            fourier_n_frequencies,
        )

        ui_is_providing_pretrained_data = _validate_ui_pretrained_data(
            coefficients, feature_names, base_feature_names, self.name
        )

        if ui_is_providing_pretrained_data:
            self.equations = equations
            self.base_feature_names = base_feature_names.tolist()
            self.feature_names = feature_names.tolist()
            self.coefficients = npa.array(coefficients)
            self.has_control_input = has_control_input
            self.custom_basis_functions = None

        elif pretrained:
            with open(pretrained_file_path, "r") as f:
                deserialized_model = json.load(f)

            self.equations = deserialized_model["equations"]
            self.base_feature_names = deserialized_model["base_feature_names"]
            self.feature_names = deserialized_model["feature_names"]
            self.coefficients = npa.array(deserialized_model["coefficients"])
            self.has_control_input = deserialized_model["has_control_input"]
            self.custom_basis_functions = None

        else:
            (
                self.equations,
                self.base_feature_names,
                self.feature_names,
                self.coefficients,
                self.has_control_input,
            ) = train_from_csv(
                file_name,
                header_as_first_row=header_as_first_row,
                state_columns=state_columns,
                control_input_columns=control_input_columns,
                dt=dt,
                time_column=time_column,
                state_derivatives_columns=state_derivatives_columns,
                discrete_time=discrete_time,
                differentiation_method=differentiation_method,
                threshold=threshold,
                alpha=alpha,
                max_iter=max_iter,
                normalize_columns=normalize_columns,
                poly_order=poly_order,
                custom_basis_functions=custom_basis_functions,
                fourier_n_frequencies=fourier_n_frequencies,
            )
            self.custom_basis_functions = custom_basis_functions

        self.nx, _ = self.coefficients.shape

        if npa.all(self.coefficients == 0):
            warnings.warn(
                "No features were selected for the SINDy model. "
                "Please check the training data and the feature selection "
                "parameters."
            )

        if initial_state is not None:
            if len(initial_state) != self.nx:
                raise ValueError(
                    f"Provided initial state has {len(initial_state)} elements. "
                    f"Expected {self.nx} elements."
                )
        else:
            initial_state = npa.zeros(self.nx)

        if self.has_control_input:
            self.declare_input_port()  # one vector valued input port for u

        if discrete_time:
            self.declare_discrete_state(
                shape=(self.nx,),
                default_value=initial_state,
                as_array=True,
            )
            self.declare_periodic_update(
                (
                    self._discrete_update
                    if self.has_control_input
                    else lambda time, state, **params: self._discrete_update(
                        time, state, (), **params
                    )
                ),
                period=discrete_time_update_interval,
                offset=0.0,
            )
            self.declare_output_port(
                self._full_discrete_state_output,
                period=discrete_time_update_interval,
                offset=0.0,
                default_value=initial_state,
                requires_inputs=False,
            )

        else:
            self.declare_continuous_state(
                ode=(
                    self._ode
                    if self.has_control_input
                    else lambda time, state, **params: self._ode(
                        time, state, (), **params
                    )
                ),
                shape=(self.nx,),
                default_value=npa.array(initial_state),
            )
            self.declare_continuous_state_output()  # output of the state in ODE

        # SymPy parsing to compute $f(x,u)$
        # For continuous-time systems $\dot{x} = f(x,u)$
        # For discrete-time systems $x_{k+1} = f(x_k, u_k)$
        sympy_base_features = sp.symbols(self.base_feature_names)

        # Convert feature names to sympy expressions
        sympy_feature_expressions = []
        for name in self.feature_names:
            # Replace spaces with multiplication
            name = name.replace(" ", "*")
            expr = sp.sympify(name)
            sympy_feature_expressions.append(expr)

        x_and_u_vec = sp.Matrix(sympy_base_features)
        custom_functions_dict = (
            {f"f{idx}": func for idx, func in enumerate(self.custom_basis_functions)}
            if self.custom_basis_functions
            else None
        )
        self.features_func = sp.lambdify(
            (x_and_u_vec,),
            sympy_feature_expressions,
            modules=[custom_functions_dict, "jax"] if custom_functions_dict else "jax",
        )  # feature functions

    def _ode(self, _time, state, inputs, **_params):
        """
        The ODE system RHS. The RHS is given by `coefficients @ features`
        """
        x = state.continuous_state
        u = inputs
        x_and_u = npa.hstack([x, u])
        features_evaluated = self.features_func(x_and_u)
        x_dot = npa.matmul(self.coefficients, npa.atleast_1d(features_evaluated))
        return x_dot

    def _discrete_update(self, _time, state, inputs, **_params):
        """
        Update map is given by `coefficients @ features`
        """
        x = state.discrete_state
        u = inputs
        x_and_u = npa.hstack([x, u])
        features_evaluated = self.features_func(x_and_u)
        x_plus = npa.matmul(self.coefficients, npa.atleast_1d(features_evaluated))
        return x_plus

    def _full_discrete_state_output(self, _time, state, *_inputs, **_params):
        return state.discrete_state

    def serialize(self, filename):
        """
        Save the relevant class attributes post training
        so that model state can be restored
        """
        sindy_data = {
            "equations": self.equations,
            "base_feature_names": self.base_feature_names,
            "feature_names": self.feature_names,
            "coefficients": self.coefficients.tolist(),  # Can't serialize numpy arrays
            "has_control_input": self.has_control_input,
        }
        with open(filename, "w") as f:
            json.dump(sindy_data, f)

    @staticmethod
    def serialize_trained_pysindy_model(model, filename):
        """
        Serialize a PySindy model trained outside of Jaxonomy.
        The saved file can be used as a pretrained model in Jaxonomy.
        """

        feature_names, coefficients = _reduce(
            model.feature_names, model.get_feature_names(), model.coefficients()
        )

        has_control_input = model.model.n_input_features_ > 0

        sindy_data = {
            "equations": model.equations,
            "base_feature_names": model.feature_names,
            "feature_names": feature_names,
            "coefficients": coefficients.tolist(),  # Can't serialize numpy arrays
            "has_control_input": has_control_input,
        }

        with open(filename, "w") as f:
            json.dump(sindy_data, f)

serialize(filename)

Save the relevant class attributes post training so that model state can be restored

Source code in jaxonomy/library/sindy.py

def serialize(self, filename):
    """
    Save the relevant class attributes post training
    so that model state can be restored
    """
    sindy_data = {
        "equations": self.equations,
        "base_feature_names": self.base_feature_names,
        "feature_names": self.feature_names,
        "coefficients": self.coefficients.tolist(),  # Can't serialize numpy arrays
        "has_control_input": self.has_control_input,
    }
    with open(filename, "w") as f:
        json.dump(sindy_data, f)

serialize_trained_pysindy_model(model, filename) staticmethod

Serialize a PySindy model trained outside of Jaxonomy. The saved file can be used as a pretrained model in Jaxonomy.

Source code in jaxonomy/library/sindy.py

@staticmethod
def serialize_trained_pysindy_model(model, filename):
    """
    Serialize a PySindy model trained outside of Jaxonomy.
    The saved file can be used as a pretrained model in Jaxonomy.
    """

    feature_names, coefficients = _reduce(
        model.feature_names, model.get_feature_names(), model.coefficients()
    )

    has_control_input = model.model.n_input_features_ > 0

    sindy_data = {
        "equations": model.equations,
        "base_feature_names": model.feature_names,
        "feature_names": feature_names,
        "coefficients": coefficients.tolist(),  # Can't serialize numpy arrays
        "has_control_input": has_control_input,
    }

    with open(filename, "w") as f:
        json.dump(sindy_data, f)

Sine

Bases: SourceBlock

Generates a sinusoidal signal.

Given amplitude a, frequency f, phase phi, and bias b, the output signal is:

    y(t) = a * sin(f * t + phi) + b

Input ports

None

Output ports

(0) The sinusoidal signal.

Parameters:

Name	Type	Description	Default
amplitude		The amplitude of the sinusoidal signal.	1.0
frequency		The frequency of the sinusoidal signal.	1.0
phase		The phase of the sinusoidal signal.	0.0
bias		The bias of the sinusoidal signal.	0.0
units	(optional, T - 104 - followup - units - on - source - blocks)	If set, the output port advertises this :class:Unit. The connect-time consistency check (T-104) then enforces downstream ports declare a compatible unit. Default None keeps the legacy "no-units" behaviour (byte-equivalent to pre-T-104 diagrams).	None

Source code in jaxonomy/library/sources.py

class Sine(SourceBlock):
    """Generates a sinusoidal signal.

    Given amplitude `a`, frequency `f`, phase `phi`, and bias `b`, the output signal is:
    ```
        y(t) = a * sin(f * t + phi) + b
    ```

    Input ports:
        None

    Output ports:
        (0) The sinusoidal signal.

    Parameters:
        amplitude:
            The amplitude of the sinusoidal signal.
        frequency:
            The frequency of the sinusoidal signal.
        phase:
            The phase of the sinusoidal signal.
        bias:
            The bias of the sinusoidal signal.
        units (optional, T-104-followup-units-on-source-blocks):
            If set, the output port advertises this :class:`Unit`. The
            connect-time consistency check (T-104) then enforces
            downstream ports declare a compatible unit. Default
            ``None`` keeps the legacy "no-units" behaviour
            (byte-equivalent to pre-T-104 diagrams).
    """

    @parameters(dynamic=["amplitude", "frequency", "phase", "bias"])
    def __init__(
        self,
        amplitude=1.0,
        frequency=1.0,
        phase=0.0,
        bias=0.0,
        units=None,
        **kwargs,
    ):
        super().__init__(self._eval, **kwargs)
        # T-104-followup-units-on-source-blocks: the parent
        # ``SourceBlock.__init__`` already declared the output port; we
        # tag it with the requested unit here so the source advertises
        # its own output unit (rather than relying on the downstream
        # port to do so). Stored as a plain attribute on the OutputPort
        # to match the convention established by T-104 phase 1 in
        # ``framework/system_base.py``.
        self.output_ports[self._output_port_idx].units = units

    def initialize(self, amplitude=1.0, frequency=1.0, phase=0.0, bias=0.0):
        pass

    def _eval(self, t, **parameters):
        a = parameters["amplitude"]
        f = parameters["frequency"]
        phi = parameters["phase"]
        b = parameters["bias"]
        return a * npa.sin(f * t + phi) + b

Slice

Bases: FeedthroughBlock

Slice the input signal using Python indexing rules.

Input ports

(0) The input signal.

Output ports

(0) The sliced output signal.

Parameters:

Name	Type	Description	Default
slice_		The slice operator to apply to the input signal. Must be specified as a string input, e.g. the output u[1:3] would be created with the block Slice("1:3").	required

Notes

Currently only up to 3-dimensional slices are supported.

Source code in jaxonomy/library/routing.py

class Slice(FeedthroughBlock):
    """Slice the input signal using Python indexing rules.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The sliced output signal.

    Parameters:
        slice_:
            The slice operator to apply to the input signal.  Must be specified as a
            string input, e.g. the output `u[1:3]` would be created with the block
            `Slice("1:3")`.

    Notes:
        Currently only up to 3-dimensional slices are supported.
    """

    @parameters(static=["slice_"])
    def __init__(self, slice_, *args, **kwargs):
        super().__init__(None, *args, **kwargs)

    def initialize(self, slice_):
        # if slice was provided as numpy slice object, remove this before validating.
        if slice_.startswith("np.s_"):
            slice_ = slice_[len("np.s_") :]
        # if slice is wrapped in [], remove them temporarily.
        if slice_[0] == "[":
            slice_ = slice_[1:]
        if slice_[-1] == "]":
            slice_ = slice_[:-1]

        # validate slice_ and ensure no nefarious code.
        pattern = re.compile(r"^[0-9,:]+$")
        if not pattern.match(slice_):
            raise BlockParameterError(
                message=f"Slice block {self.name} detected invalid slice operator {slice_}. [] are optional. Valid examples: '1:3,4', '[:,4:10]'",
                parameter_name="slice_",
            )

        # replace the [] and eval to numpy slcie object
        slice_ = "np.s_[" + slice_ + "]"
        np_slice = eval(slice_)

        def _func(inp):
            return npa.array(inp)[np_slice]

        self.replace_op(_func)

SnapshotData dataclass

Container for a column-wise snapshot matrix.

Attributes:

Name	Type	Description
X	ndarray	State/output snapshots, shape (n_features, n_samples).
time	Optional[ndarray]	Optional sample times, shape (n_samples,).
inputs	Optional[ndarray]	Optional input snapshots U, shape (n_inputs, n_samples).
Xdot	Optional[ndarray]	Optional time-derivative snapshots, shape (n_features, n_samples).

Source code in jaxonomy/library/rom/snapshots.py

@dataclass
class SnapshotData:
    """Container for a column-wise snapshot matrix.

    Attributes:
        X: State/output snapshots, shape ``(n_features, n_samples)``.
        time: Optional sample times, shape ``(n_samples,)``.
        inputs: Optional input snapshots ``U``, shape ``(n_inputs, n_samples)``.
        Xdot: Optional time-derivative snapshots, shape ``(n_features, n_samples)``.
    """

    X: np.ndarray
    time: Optional[np.ndarray] = None
    inputs: Optional[np.ndarray] = None
    Xdot: Optional[np.ndarray] = None

    @property
    def n_features(self) -> int:
        return int(self.X.shape[0])

    @property
    def n_samples(self) -> int:
        return int(self.X.shape[1])

SoftRateLimiter

Bases: LeafSystem

Smooth (differentiable) rate limiter.

Drop-in differentiable variant of :class:RateLimiter. Identical discrete update semantics, except the inner hard clip on the desired step (u - y_prev) is replaced by a smooth saturation so gradients flow through the limiter even when it is actively limiting.

The smooth clip is implemented via :func:soft_saturate and recovers the hard rate limiter as sharpness -> inf.

Parameters mirror :class:RateLimiter plus: sharpness: Scalar > 0 controlling how sharply the smooth saturation transitions at the rate limits. Larger sharpness -> closer to the hard rate limiter. Default 10.0.

See :class:RateLimiter for the (non-smoothed) reference behavior.

Source code in jaxonomy/library/nonlinearities.py

class SoftRateLimiter(LeafSystem):
    """Smooth (differentiable) rate limiter.

    Drop-in differentiable variant of :class:`RateLimiter`. Identical
    discrete update semantics, except the inner hard ``clip`` on the
    desired step ``(u - y_prev)`` is replaced by a smooth saturation so
    gradients flow through the limiter even when it is actively limiting.

    The smooth clip is implemented via :func:`soft_saturate` and recovers
    the hard rate limiter as ``sharpness -> inf``.

    Parameters mirror :class:`RateLimiter` plus:
        sharpness:
            Scalar > 0 controlling how sharply the smooth saturation
            transitions at the rate limits. Larger ``sharpness`` -> closer
            to the hard rate limiter. Default ``10.0``.

    See :class:`RateLimiter` for the (non-smoothed) reference behavior.
    """

    class DiscreteStateType(NamedTuple):
        y_prev: Array
        t_prev: float

    @parameters(
        static=["dt", "enable_dynamic_upper_limit", "enable_dynamic_lower_limit"],
        dynamic=["upper_limit", "lower_limit", "sharpness"],
    )
    def __init__(
        self,
        dt,
        upper_limit=np.inf,
        enable_dynamic_upper_limit=False,
        lower_limit=-np.inf,
        enable_dynamic_lower_limit=False,
        sharpness=10.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.primary_input_index = self.declare_input_port()
        self.enable_dynamic_upper_limit = enable_dynamic_upper_limit
        self.enable_dynamic_lower_limit = enable_dynamic_lower_limit
        self.dt = dt

        if enable_dynamic_upper_limit:
            self.upper_limit_index = self.declare_input_port()

        if enable_dynamic_lower_limit:
            self.lower_limit_index = self.declare_input_port()

        self.output_index = self.declare_output_port(
            self._output,
            period=dt,
            offset=0.0,
        )

    def initialize(
        self,
        upper_limit=np.inf,
        enable_dynamic_upper_limit=False,
        lower_limit=-np.inf,
        enable_dynamic_lower_limit=False,
        sharpness=10.0,
        dt=None,
    ):
        if enable_dynamic_upper_limit != self.enable_dynamic_upper_limit:
            raise ValueError(
                "SoftRateLimiter: enable_dynamic_upper_limit cannot be changed after initialization"
            )
        if enable_dynamic_lower_limit != self.enable_dynamic_lower_limit:
            raise ValueError(
                "SoftRateLimiter: enable_dynamic_lower_limit cannot be changed after initialization"
            )

    def _output(self, time, state, *inputs, **params):
        y_prev = state.cache[self.output_index]
        u = inputs[self.primary_input_index]
        t_diff = self.dt

        ulim = (
            inputs[self.upper_limit_index]
            if self.enable_dynamic_upper_limit
            else params["upper_limit"]
        )
        llim = (
            inputs[self.lower_limit_index]
            if self.enable_dynamic_lower_limit
            else params["lower_limit"]
        )
        k = params["sharpness"]

        # Smoothly clip the per-step delta in y.
        delta = u - y_prev
        delta_lo = t_diff * llim
        delta_hi = t_diff * ulim
        delta_clipped = soft_saturate(delta, delta_lo, delta_hi, k)
        return y_prev + delta_clipped

    def initialize_static_data(self, context):
        try:
            u = self.eval_input(context)
            self._default_cache[self.output_index] = u
            local_context = context[self.system_id].with_discrete_state(u)
            local_context = local_context.with_cached_value(self.output_index, u)
            context = context.with_subcontext(self.system_id, local_context)
        except UpstreamEvalError:
            logger.debug(
                "SoftRateLimiter.initialize_static_data: UpstreamEvalError. "
                "Continuing without default value initialization."
            )
        return context

SoftSaturate

Bases: FeedthroughBlock

Smooth (differentiable) saturation block.

Drop-in differentiable variant of :class:Saturate that uses :func:soft_saturate instead of npa.clip. The original hard :class:Saturate block is unchanged.

Why a separate block: the standard :class:Saturate block returns npa.clip(u, lo, hi), whose gradient is exactly zero outside the bounds. That kills gradient signal in any optimization that drives the input past the limits. SoftSaturate keeps gradient flow alive so e.g. trajectory optimization through actuator limits actually converges.

See :func:soft_saturate for the formula. Unlike :class:Saturate, this block does not declare zero-crossing events (it has no discontinuity to catch).

Parameters:

Name	Type	Description	Default
upper_limit		Upper limit; default 1.0. Must be finite.	1.0
lower_limit		Lower limit; default 0.0. Must be finite and strictly less than upper_limit.	0.0
sharpness		Smoothing knob, > 0; default 10.0. As sharpness -> inf this approaches the hard :class:Saturate block.	10.0

Input ports

(0) The input signal.

Output ports

(0) The smoothly-saturated output signal.

Source code in jaxonomy/library/nonlinearities.py

class SoftSaturate(FeedthroughBlock):
    """Smooth (differentiable) saturation block.

    Drop-in differentiable variant of :class:`Saturate` that uses
    :func:`soft_saturate` instead of ``npa.clip``. The original hard
    :class:`Saturate` block is unchanged.

    Why a separate block: the standard :class:`Saturate` block returns
    ``npa.clip(u, lo, hi)``, whose gradient is exactly zero outside the
    bounds. That kills gradient signal in any optimization that drives
    the input past the limits. ``SoftSaturate`` keeps gradient flow alive
    so e.g. trajectory optimization through actuator limits actually
    converges.

    See :func:`soft_saturate` for the formula. Unlike :class:`Saturate`,
    this block does *not* declare zero-crossing events (it has no
    discontinuity to catch).

    Parameters:
        upper_limit: Upper limit; default ``1.0``. Must be finite.
        lower_limit: Lower limit; default ``0.0``. Must be finite and
            strictly less than ``upper_limit``.
        sharpness: Smoothing knob, > 0; default ``10.0``. As
            ``sharpness -> inf`` this approaches the hard
            :class:`Saturate` block.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The smoothly-saturated output signal.
    """

    @parameters(dynamic=["upper_limit", "lower_limit", "sharpness"])
    def __init__(
        self,
        lower_limit=0.0,
        upper_limit=1.0,
        sharpness=10.0,
        **kwargs,
    ):
        super().__init__(self._soft_saturate, **kwargs)
        if not np.isfinite(lower_limit) or not np.isfinite(upper_limit):
            raise BlockParameterError(
                message=(
                    f"SoftSaturate block {self.name} requires finite "
                    f"lower_limit/upper_limit, got lower={lower_limit}, "
                    f"upper={upper_limit}. Use the hard Saturate block "
                    "for unbounded sides."
                ),
                system=self,
                parameter_name="lower_limit",
            )
        if upper_limit <= lower_limit:
            raise BlockParameterError(
                message=(
                    f"SoftSaturate block {self.name}: upper_limit "
                    f"({upper_limit}) must be > lower_limit ({lower_limit})."
                ),
                system=self,
                parameter_name="upper_limit",
            )
        if sharpness <= 0:
            raise BlockParameterError(
                message=(
                    f"SoftSaturate block {self.name}: sharpness must be "
                    f"> 0, got {sharpness}."
                ),
                system=self,
                parameter_name="sharpness",
            )

    def initialize(self, lower_limit=0.0, upper_limit=1.0, sharpness=10.0):
        pass

    def _soft_saturate(self, u, **params):
        return soft_saturate(
            u,
            params["lower_limit"],
            params["upper_limit"],
            params["sharpness"],
        )

SourceBlock

Bases: LeafSystem

Simple blocks with a single time-dependent output

Source code in jaxonomy/library/generic.py

class SourceBlock(LeafSystem):
    """Simple blocks with a single time-dependent output"""

    def __init__(self, func: Callable, **kwargs):
        """Create a source block with a time-dependent output.

        Args:
            func (Callable):
                A function of time and parameters that returns a single value.
                Signature should be `func(time, **parameters) -> Array`.
        """
        super().__init__(**kwargs)
        self._output_port_idx = self.declare_output_port(
            None,
            name="out_0",
            prerequisites_of_calc=[DependencyTicket.time],
            requires_inputs=False,
        )
        self.replace_op(func)

    def replace_op(self, func):
        def _callback(time, state, *inputs, **parameters):
            return func(time, **parameters)

        self.configure_output_port(
            self._output_port_idx,
            _callback,
            prerequisites_of_calc=[DependencyTicket.time],
            requires_inputs=False,
        )

__init__(func, **kwargs)

Create a source block with a time-dependent output.

Parameters:

Name	Type	Description	Default
func	Callable	A function of time and parameters that returns a single value. Signature should be func(time, **parameters) -> Array.	required

Source code in jaxonomy/library/generic.py

def __init__(self, func: Callable, **kwargs):
    """Create a source block with a time-dependent output.

    Args:
        func (Callable):
            A function of time and parameters that returns a single value.
            Signature should be `func(time, **parameters) -> Array`.
    """
    super().__init__(**kwargs)
    self._output_port_idx = self.declare_output_port(
        None,
        name="out_0",
        prerequisites_of_calc=[DependencyTicket.time],
        requires_inputs=False,
    )
    self.replace_op(func)

SquareRoot

Bases: FeedthroughBlock

Compute the square root of the input signal.

Dispatches to jax.numpy.sqrt, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.sqrt.html

Input ports

(0) The input signal.

Output ports

(0) The square root of the input signal.

Source code in jaxonomy/library/math_ops.py

class SquareRoot(FeedthroughBlock):
    """Compute the square root of the input signal.

    Dispatches to `jax.numpy.sqrt`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.sqrt.html

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The square root of the input signal.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(npa.sqrt, *args, **kwargs)

Stack

Bases: ReduceBlock

Stack the input signals into a single output signal along a new axis.

Dispatches to jax.numpy.stack, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.stack.html

Input ports

(0..n_in-1) The input signals.

Output ports

(0) The stacked output signal.

Parameters:

Name	Type	Description	Default
axis		The axis along which the input signals are stacked. Default is 0.	0

Source code in jaxonomy/library/math_ops.py

class Stack(ReduceBlock):
    """Stack the input signals into a single output signal along a new axis.

    Dispatches to `jax.numpy.stack`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.stack.html

    Input ports:
        (0..n_in-1) The input signals.

    Output ports:
        (0) The stacked output signal.

    Parameters:
        axis:
            The axis along which the input signals are stacked.  Default is 0.
    """

    @parameters(static=["axis"])
    def __init__(self, n_in, axis=0, **kwargs):
        super().__init__(n_in, None, **kwargs)

    def initialize(self, axis):
        self.replace_op(partial(npa.stack, axis=int(axis)))

StateMachine

Bases: LeafSystem

Finite State Machine similar to Mealy Machine. https://en.wikipedia.org/wiki/Mealy_machine

The state machine can be executed either periodically or by zero_crossings.

Each state as 0 or more exit transitions. These are prioritized such that when 2 exits are simultaneously valid, the higher priority is executed. It is not allowed for a state to have more than one exit transition with no guard. Guardless exits only make sense in the periodic case.

Each transitions may have 0 or more actions. Each action is a python statement that modifies the value of an output. When a transitions is executed (i.e. it's guard evaluates to true), its actions are then processed.

If 'time' is needed for guards or actions, pass 'time' in from clock block.

Whether executed periodically or by zero_crossings, the states are constant between transitions executions. In the zero_crossing case, all guards for transitions exiting the current state are continuously checked, and if any 'triggers', then the earlist point in time that any guard becomes true is determined, the actions of the earliest (and highest priority if multiple trigger simultaneously) guard are executed at that time, and the simulation continues afterwards.

Input ports

User specified.

Output ports

User specified.

Parameters:

Name	Type	Description	Default
dt		Either Float or None. When not None, state machine is executed periodically. When None, the transitions are monitored by zero_crossing events.	None
accelerate_with_jax	bool	Bool. When True, the actions and guards are JIT-compiled with JAX. Default is False.	False

Source code in jaxonomy/library/state_machine.py

@parameters(static=["accelerate_with_jax"])
class StateMachine(LeafSystem):
    """Finite State Machine similar to Mealy Machine.
    https://en.wikipedia.org/wiki/Mealy_machine

    The state machine can be executed either periodically or by zero_crossings.

    Each state as 0 or more exit transitions. These are prioritized such that
    when 2 exits are simultaneously valid, the higher priority is executed.
    It is not allowed for a state to have more than one exit transition with no
    guard. Guardless exits only make sense in the periodic case.

    Each transitions may have 0 or more actions. Each action is a python
    statement that modifies the value of an output. When a transitions is executed
    (i.e. it's guard evaluates to true), its actions are then processed.

    If 'time' is needed for guards or actions, pass 'time' in from clock block.

    Whether executed periodically or by zero_crossings, the states are constant
    between transitions executions.
    In the zero_crossing case, all guards for transitions exiting the current
    state are continuously checked, and if any 'triggers', then the earlist point
    in time that any guard becomes true is determined, the actions of the earliest
    (and highest priority if multiple trigger simultaneously) guard are executed at
    that time, and the simulation continues afterwards.

    Input ports:
        User specified.

    Output ports:
        User specified.

    Parameters:
        dt:
            Either Float or None.
            When not None, state machine is executed periodically.
            When None, the transitions are monitored by zero_crossing
            events.
        accelerate_with_jax:
            Bool. When True, the actions and guards are JIT-compiled with JAX.
            Default is False.
    """

    def __init__(
        self,
        sm_data: StateMachineData,
        inputs: List[str] = None,  # [name]
        outputs: List[str] = None,  # [name]
        dt=None,
        time_mode: str = "agnostic",
        name: str = None,
        ui_id: str = None,
        accelerate_with_jax: bool = False,
        **kwargs,
    ):
        super().__init__(name=name, ui_id=ui_id)

        if time_mode not in ["discrete", "agnostic"]:
            raise BlockInitializationError(
                f"Invalid time mode '{time_mode}' for PythonScript block", system=self
            )

        if time_mode == "discrete" and dt is None:
            raise BlockInitializationError(
                "When in discrete time mode, dt is required for block", system=self
            )

        if npa.active_backend == "numpy" and accelerate_with_jax:
            raise BlockInitializationError(
                "Must use JAX numerical backend when accelerate_with_jax=True",
                system=self,
            )

        try:
            sm_data = _validate_sm_data(sm_data)
        except ValueError as e:
            raise StaticError(message=str(e), system=self) from e

        self._accelerate_with_jax = accelerate_with_jax

        if accelerate_with_jax:
            # inputs to many jax functions are expected to be jnp.arrays of
            # same shape, so we pad the arrays to the same shapes.
            self._sm = sm_data.to_padded_arrays()

            self._guards = npa.array(
                [
                    [t.guard_id for t in self._sm.states[idx].transitions]
                    for idx in self._sm.states.keys()
                ]
            )

            self._dst = npa.array(
                [
                    [t.dst for t in self._sm.states[idx].transitions]
                    for idx in self._sm.states.keys()
                ]
            )

            self._actions = npa.array(
                [
                    [t.action_ids for t in self._sm.states[idx].transitions]
                    for idx in self._sm.states.keys()
                ]
            )
        else:
            self._sm = sm_data

        self.time_mode = time_mode
        _is_periodic = time_mode == "discrete"

        if inputs is None:
            inputs = []
        if outputs is None:
            outputs = []
        elif isinstance(outputs, dict):
            outputs = list(outputs.keys())

        # delcare inputs
        self._input_names = inputs
        for name in inputs:
            self.declare_input_port(name)

        self._output_names = outputs

        # Create the default discrete state values
        self._create_discrete_state_type(include_state_idx=_is_periodic)
        default_values = self._create_initial_discrete_state(
            include_state_idx=_is_periodic
        )
        self.declare_discrete_state(default_value=default_values, as_array=False)

        # Declare output ports for each state variable
        def _make_output_callback(o_port_name):
            def _output(time, state, *inputs, **parameters):
                return getattr(state.discrete_state, o_port_name)

            return _output

        for o_port_name in outputs:
            self.declare_output_port(
                _make_output_callback(o_port_name),
                name=o_port_name,
                prerequisites_of_calc=[DependencyTicket.xd],
                requires_inputs=False,
            )

        if _is_periodic:
            # delcare the periodic update event
            self.declare_periodic_update(
                self._discrete_update,
                period=dt,
                offset=dt,
            )
        else:
            # T-033: in agnostic (zero-crossing) mode, transitions are
            # documented to fire in *priority order* — the lowest-index
            # transition out of a given source state wins when multiple
            # guards are simultaneously True. The discrete-update path
            # honours this by walking the transition list and breaking
            # at the first True guard. The zero-crossing path declares
            # one event per transition, so without explicit gating they
            # all fire and order is implementation-defined.
            #
            # Fix: each guard's effective value is AND-ed with the
            # negation of every higher-priority guard from the same
            # source state. When two guards become True at the same
            # instant, only the lowest-index one's zero-crossing
            # actually flips False -> True; the others' effective
            # truth stays False because a higher-priority guard is
            # already True, so no event fires.
            #
            # wrap the callback generation so that they do not get overwritten
            # in subsequent calls to declare_zero_crossing()
            def _eval_guard(guard_id, inputs, outputs):
                return self._sm.registry.guards[guard_id](**inputs, **outputs)

            def _make_guard_callback(t, higher_priority_guard_ids):
                def _guard(_time, state, *inputs, **parameters):
                    # Inputs are in order of port declaration, so they match `self._input_names`
                    inputs = dict(zip(self._input_names, inputs))
                    # get the values of the outputs as they are presently.
                    outputs = state.discrete_state._asdict()
                    my_truth = _eval_guard(t.guard_id, inputs, outputs)
                    # AND-NOT each higher-priority sibling guard so that
                    # the lowest-index simultaneous truth wins.
                    blocked = False
                    for hg in higher_priority_guard_ids:
                        blocked = npa.logical_or(
                            blocked, _eval_guard(hg, inputs, outputs)
                        )
                    effective = npa.logical_and(
                        my_truth, npa.logical_not(blocked)
                    )
                    # we do this so that when a guard goes False-True,
                    # it creates a zero-crossing that can be localized in time.
                    g = npa.where(effective, 1.0, -1.0)
                    return g

                return _guard

            def _make_reset_callback(t):
                def _reset(_time, state, *inputs, **p):
                    # Inputs are in order of port declaration, so they match `self._input_names`
                    inputs = dict(zip(self._input_names, inputs))
                    # get the values of the outputs as they are presently.
                    outputs = state.discrete_state._asdict()
                    if self._accelerate_with_jax:
                        updated_outputs = self._exec_actions_jax(
                            t.action_ids, inputs, outputs
                        )
                    else:
                        updated_outputs = self._exec_actions(
                            t.action_ids, inputs, outputs
                        )
                    # Actions only assign a subset of the outputs; carry the
                    # untouched ones over from the current discrete state
                    # (mirrors the merge in the discrete-update callbacks).
                    merged_outputs = {
                        k: updated_outputs[k] if k in updated_outputs else outputs[k]
                        for k in self._output_names
                    }
                    return state.with_discrete_state(
                        value=self.DiscreteStateType(**merged_outputs)
                    )

                return _reset

            # T-NEW-sm-smooth-guard: a *smooth* guard residual (``lhs - rhs``
            # for a simple comparison guard) the event-time (saltation) gradient
            # can differentiate, since the boolean trigger guard above has zero
            # gradient.  Used only by the reverse-mode saltation paths, never for
            # triggering.  ``None`` when the guard is not a simple comparison.
            grad_guards = self._sm.registry.grad_guards

            def _make_grad_guard_callback(t):
                gg = (
                    grad_guards[t.guard_id]
                    if 0 <= t.guard_id < len(grad_guards)
                    else None
                )
                if gg is None:
                    return None

                def _grad_guard(_time, state, *inputs, **parameters):
                    inputs = dict(zip(self._input_names, inputs))
                    outputs = state.discrete_state._asdict()
                    return gg(**inputs, **outputs)

                return _grad_guard

            # declare zero-crossing driven events and mode
            self.declare_default_mode(self._sm.initial_state)
            self.declare_mode_output()
            for st_idx, st in self._sm.states.items():
                # T-033: walk transitions in priority order; collect
                # each preceding guard's id so the current callback
                # can AND-NOT against it.
                higher_priority: list = []
                for t in st.transitions:
                    self.declare_zero_crossing(
                        guard=_make_guard_callback(
                            t, list(higher_priority)
                        ),
                        reset_map=_make_reset_callback(t),
                        direction="negative_then_non_negative",  # we only care when the guard transitions False->True
                        start_mode=st_idx,
                        end_mode=t.dst,
                        grad_guard=_make_grad_guard_callback(t),
                    )
                    higher_priority.append(t.guard_id)

    def _create_discrete_state_type(self, include_state_idx=True):
        if include_state_idx:
            # unique identifier for the state machine state variable.
            # State names must not use a leading underscore due to namedtuple requirements.
            st_name = "active_state_index"

            if st_name in self._input_names or st_name in self._output_names:
                msg = f"StateMachine {self.name} has port with same name as state {st_name}, this is not allowed."
                raise StaticError(message=msg, system=self)

            self._st_name = st_name

            attribs = [st_name] + self._output_names
        else:
            attribs = self._output_names
        # declare the discrete_state as a namedtuple
        self.DiscreteStateType = namedtuple("DiscreteStateType", attribs)

    def _create_initial_discrete_state(self, include_state_idx=True):
        # execute the entry point actions
        # Inputs at initialization are undefined, so we populate them with `None`.
        # This aligns with the Jaxonomy execution model; attempts to use undefined inputs will fail out downstream.
        inputs = {n: None for n in self._input_names}
        outputs = {n: None for n in self._output_names}

        initial_outputs = self._exec_actions(self._sm.initial_actions, inputs, outputs)

        # check if any initial_outputs is NaN
        for k, v in initial_outputs.items():
            if np.any(np.isnan(v)):
                msg = (
                    "StateMachine has NaN values in the initial outputs. "
                    "Inputs can't be used in initial actions."
                )
                raise BlockInitializationError(message=msg, system=self)

        # enforce that all outputs have been initialized
        initialized_output_names = set(initial_outputs.keys())
        all_output_names = set(self._output_names)
        uninitialized_output_names = all_output_names.difference(
            initialized_output_names
        )
        if uninitialized_output_names:
            msg = f"StateMachine does not initialize the following output values in the entry point actions: {uninitialized_output_names}"
            raise BlockInitializationError(message=msg, system=self)

        # get and save the output dtype,shape for use in creating the jax.pure_callback
        self.output_port_params = {
            o_port_name: {"dtype": jnp.array(val).dtype, "shape": jnp.array(val).shape}
            for o_port_name, val in initial_outputs.items()
        }

        # prepare the initial state
        if include_state_idx:
            return self.DiscreteStateType(
                active_state_index=self._sm.initial_state,
                **initial_outputs,
            )

        return self.DiscreteStateType(**initial_outputs)

    def _filter_locals(self, local_env):
        # remove any bindings from locals that are not outputs.
        filtered_locals = {}
        for key, value in local_env.items():
            if key in self._output_names:
                filtered_locals[key] = value
        return filtered_locals

    def _exec_actions(self, action_ids, inputs, outputs):
        # execute actions, in context with inputs values, when done
        # all actions, filter out any variable bindings that do
        # not correspond to outputs, then repack as dict of jnp.arrays
        updated_outputs = {}
        for action_id in action_ids:
            if action_id == -1:  # padded actions are -1
                continue
            input_args = [inputs[k] for k in self._input_names]
            output_args = [outputs[k] for k in self._output_names]
            output = self._sm.registry.actions[action_id](*input_args, *output_args)
            updated_outputs.update(output)

        updated_outputs = self._filter_locals(updated_outputs)
        updated_outputs = {k: jnp.array(v) for k, v in updated_outputs.items()}
        return updated_outputs

    def _exec_actions_jax(self, action_ids, inputs, outputs):
        """Execute actions in a JAX-compatible way."""

        def _exec_action(action_id):
            input_args = [inputs[k] for k in self._input_names]
            output_args = [outputs[k] for k in self._output_names]
            return npa.cond(
                action_id == -1,  # padded actions are -1
                lambda: ({k: v for k, v in outputs.items()}, True),
                lambda: (
                    npa.switch(
                        action_id,
                        self._sm.registry.actions,
                        *input_args,
                        *output_args,
                    ),
                    False,
                ),
            )

        # TODO: npa.vmap (implement numpy version)
        action_outputs = jax.vmap(_exec_action)(action_ids)

        def _accumulate_outputs(carry, outputs):
            output, is_pad = outputs
            update = npa.cond(
                is_pad, lambda: carry, lambda: {k: v for k, v in output.items()}
            )
            carry.update(update)
            return carry, carry

        init = {**outputs}
        updated_outputs, _ = npa.scan(_accumulate_outputs, init, action_outputs)

        updated_outputs = self._filter_locals(updated_outputs)

        return updated_outputs

    def _numpy_callback(self, present_state_index, inputs, outputs):
        """
        The concept here is to evaluate all possible exit transitions from
        the active state, and then just return the updated (state,output values)
        for the successful transition. In the case no transitions are successful,
        we just return the present state and presen_outputs. Since we have ordered
        the possible transitions in order of priority, executing the lowest index
        successful trasition is 'correct' behavior.

        jax-compatible version of this function is `_jax_callback`.
        """
        # get the active state index, and the possible exit transitions
        present_state_index = int(present_state_index)
        actv_trns = self._sm.states[present_state_index].transitions

        # evaluate the guard for each possible exit transition.
        evaluated_guards = [
            self._sm.registry.guards[transition.guard_id](**inputs, **outputs)
            for transition in actv_trns
        ]

        if np.any(evaluated_guards):
            actv_trn = actv_trns[evaluated_guards.index(True)]
            new_state = actv_trn.dst
            updated_outputs = self._exec_actions(actv_trn.action_ids, inputs, outputs)
            new_outputs = []
            for k in self._output_names:
                output = updated_outputs[k] if k in updated_outputs else outputs[k]
                new_outputs.append(np.array(output))
            retval = [np.array(new_state), new_outputs]
        else:
            outputs = [np.array(outputs[k]) for k in self._output_names]
            retval = [np.array(present_state_index), outputs]

        return retval

    def _jax_callback(self, present_state_index, inputs, outputs):
        """
        The concept here is to evaluate all possible exit transitions from
        the active state, and then just return the updated (state,output values)
        for the successful transition. In the case no transitions are successful,
        we just return the present state and present_outputs. Since we have ordered
        the possible transitions in order of priority, executing the lowest index
        successful trasition is 'correct' behavior.
        """

        active_guards = _choose(present_state_index, self._guards)

        input_args = [inputs[k] for k in self._input_names]
        output_args = [outputs[k] for k in self._output_names]
        # evaluate the guard for each possible exit transition.
        evaluated_guards = npa.array(
            [
                npa.switch(
                    guard_id,
                    self._sm.registry.guards,
                    *input_args,
                    *output_args,
                )
                for guard_id in active_guards
            ]
        )

        def on_true():
            # Find the first active transition where the guard is True
            active_dst = _choose(present_state_index, self._dst)
            active_actions = _choose(present_state_index, self._actions)

            if np.size(evaluated_guards) == 0:
                # no guards are True, so we return the present state and outputs
                # note that adding jnp.size(evaluated_guards) > 0 to the npa.cond
                # still evaluates the true branch, so we need to check the size
                # here.
                new_outputs = [jnp.array(outputs[k]) for k in self._output_names]
                return npa.array(present_state_index), new_outputs

            idx = npa.argmax(
                evaluated_guards
            )  # TODO: check that it returns the first index (lowest priority)

            new_state = _choose(idx, active_dst)
            action_ids = _choose(idx, active_actions)

            updated_outputs = self._exec_actions_jax(action_ids, inputs, outputs)
            new_outputs = []
            for k in self._output_names:
                output = updated_outputs[k] if k in updated_outputs else outputs[k]
                new_outputs.append(npa.array(output))
            return new_state.squeeze(), new_outputs

        def on_false():
            new_outputs = [jnp.array(outputs[k]) for k in self._output_names]
            return npa.array(present_state_index), new_outputs

        return npa.cond(
            npa.any(evaluated_guards),
            on_true,
            on_false,
        )

    def _discrete_update(self, _time, state: LeafState, *inputs, **params):
        # persent state index
        actv_state = state.discrete_state.active_state_index

        # Inputs are in order of port declaration, so they match `self._input_names`
        inputs = dict(zip(self._input_names, inputs))

        # get the values of the outputs as they are presently.
        outputs = {
            key: value
            for key, value in state.discrete_state._asdict().items()
            if key not in {self._st_name}
        }

        if self._accelerate_with_jax:
            new_state, new_outputs = self._jax_callback(actv_state, inputs, outputs)
        else:
            # build jax.pure_callback result_shape_dtypes
            # its a nested list like: [actv_st, [outp0, outp1, ... outpN]]
            result_shape_dtypes = [jax.ShapeDtypeStruct((), jnp.int64)]  # actv_state
            result_shape_dtypes_outps = []
            for var in self._output_names:
                port = self.output_port_params[var]
                result_shape_dtypes_outps.append(
                    jax.ShapeDtypeStruct(port["shape"], np.dtype(port["dtype"]))
                )
            result_shape_dtypes.append(result_shape_dtypes_outps)

            if npa.active_backend == "numpy":
                new_state, new_outputs = self._numpy_callback(
                    actv_state,
                    inputs,
                    outputs,
                )
            else:
                # TODO: implement jax.custom_jvp to raise useful error when trying
                # to differentiate when accelerate_with_jax is False
                new_state, new_outputs = jax.pure_callback(
                    self._numpy_callback,
                    result_shape_dtypes,
                    actv_state,
                    inputs,
                    outputs,
                )

        outputs_dict = {k: v for k, v in zip(self._output_names, new_outputs)}

        return self.DiscreteStateType(
            active_state_index=new_state,
            **outputs_dict,
        )

Step

Bases: SourceBlock

A step signal.

Given start value y0, end value y1, and step time t0, the output signal is:

    y(t) = y0 if t < t0 else y1

Input ports

None

Output ports

(0) The step signal.

Parameters:

Name	Type	Description	Default
start_value		The value of the output signal before the step time.	0.0
end_value		The value of the output signal after the step time.	1.0
step_time		The time at which the step occurs.	1.0
units	(optional, T - 104 - followup - units - on - source - blocks)	If set, the output port advertises this :class:Unit. The connect-time consistency check (T-104) then enforces downstream ports declare a compatible unit. Default None keeps the legacy "no-units" behaviour (byte-equivalent to pre-T-104 diagrams).	None

Source code in jaxonomy/library/sources.py

class Step(SourceBlock):
    """A step signal.

    Given start value `y0`, end value `y1`, and step time `t0`, the
    output signal is:
    ```
        y(t) = y0 if t < t0 else y1
    ```

    Input ports:
        None

    Output ports:
        (0) The step signal.

    Parameters:
        start_value:
            The value of the output signal before the step time.
        end_value:
            The value of the output signal after the step time.
        step_time:
            The time at which the step occurs.
        units (optional, T-104-followup-units-on-source-blocks):
            If set, the output port advertises this :class:`Unit`. The
            connect-time consistency check (T-104) then enforces
            downstream ports declare a compatible unit. Default
            ``None`` keeps the legacy "no-units" behaviour
            (byte-equivalent to pre-T-104 diagrams).
    """

    @parameters(dynamic=["start_value", "end_value"], static=["step_time"])
    def __init__(
        self,
        start_value=0.0,
        end_value=1.0,
        step_time=1.0,
        units=None,
        **kwargs,
    ):
        super().__init__(self._func, **kwargs)
        # T-104-followup-units-on-source-blocks: see Sine for rationale.
        self.output_ports[self._output_port_idx].units = units
        self._periodic_update_idx = self.declare_periodic_update()

    def initialize(self, start_value, end_value, step_time):
        # Add a dummy event so that the ODE solver doesn't try to integrate through
        # the discontinuity.
        self._step_time = step_time
        self.declare_discrete_state(default_value=False)
        self.configure_periodic_update(
            self._periodic_update_idx,
            lambda *args, **kwargs: True,
            period=np.inf,
            offset=step_time,
        )

    def _func(self, time, **parameters):
        return npa.where(
            time >= self._step_time,
            parameters["end_value"],
            parameters["start_value"],
        )

Stop

Bases: LeafSystem

Stop the simulation early as soon as the input signal becomes True.

If the input signal changes as a result of a discrete update, the simulation will terminate the major step early (before advancing continuous time).

Input ports

(0): the boolean- or binary-valued termination signal

Output ports

None

Source code in jaxonomy/library/nonlinearities.py

class Stop(LeafSystem):
    """Stop the simulation early as soon as the input signal becomes True.

    If the input signal changes as a result of a discrete update, the simulation
    will terminate the major step early (before advancing continuous time).

    Input ports:
        (0): the boolean- or binary-valued termination signal

    Output ports:
        None
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.declare_input_port()

        self.declare_zero_crossing(
            guard=self._guard,
            direction="negative_then_non_negative",
            terminal=True,
        )

    def _guard(self, time, state, u, **p):
        return npa.where(u, 1.0, -1.0)

SumOfElements

Bases: FeedthroughBlock

Compute the sum of the elements of the input signal.

Dispatches to jax.numpy.sum, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.sum.html

Input ports

(0) The input signal.

Output ports

(0) The sum of the elements of the input signal.

Source code in jaxonomy/library/math_ops.py

class SumOfElements(FeedthroughBlock):
    """Compute the sum of the elements of the input signal.

    Dispatches to `jax.numpy.sum`, so see the JAX docs for details:
    https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.sum.html

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The sum of the elements of the input signal.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(npa.sum, *args, **kwargs)

Switch

Bases: LeafSystem

Route one of two data signals based on a thresholded control signal.

Three inputs (data_a, control, data_b) and one output:

.. code-block:: python

y = data_a if criteria(control, threshold) else data_b

Default mode="where" is implemented via npa.where, so JAX gradients flow through both data branches simultaneously (the selector branch is treated as non-differentiable, which is the only well-defined choice for a hard threshold).

mode="smooth" replaces the hard where with a sigmoid blend

.. code-block:: python

alpha  = sigmoid(sharpness * sign * (control - threshold))
y      = alpha * data_a + (1 - alpha) * data_b

where sign is +1 for the >=/> criteria and -1 for the <=/< criteria, so the smooth output approaches the hard answer in the strict-active region as sharpness -> inf. This mode lets gradients flow through the threshold itself, which the hard where zeroes out — the killer feature for trajectory optimization where the threshold is a tunable parameter.

data_a and data_b must be broadcast-compatible (same as npa.where's requirements). The block does not enforce that control is a scalar — element-wise selection is supported when control and the data inputs broadcast together.

Parameters:

Name	Type	Description	Default
threshold		scalar threshold against which control is compared.	0.0
criteria		one of ">=", ">", "<=", "<", "==", "!=". Default ">=". The equality criteria ("=="/"!=") are not supported in mode="smooth" (no sigmoid approximation makes sense for them).	'>='
mode		one of "where" (default), "smooth", or "hard". "where" is byte-equivalent to T-118 phase 1 and the right pick for simulation. "smooth" is the right pick for gradient-based optimization through the threshold. "hard" dispatches to jax.lax.cond so only the active branch is evaluated — useful when one branch is much more expensive than the other or when branches have incompatible side effects. Switch(mode='hard') is incompatible with vmap; use mode='where' or mode='smooth' for batched use (e.g. under simulate_batch). On older JAX (<0.4) a batched predicate raises TracerBoolConversionError; on modern JAX the cond is silently rewritten to evaluate both branches with a select, which is numerically correct but defeats the entire point of picking mode='hard' over mode='where'.	'where'
sharpness		positive scalar controlling sigmoid steepness in mode="smooth". Default 10.0. Larger values give a tighter approximation to the hard switch but smaller (and faster vanishing) gradients in the strict-active region. Ignored when mode="where".	10.0

Input ports

(0) data_a — output when criteria(control, threshold) is True. (1) control — the selector signal compared to threshold. (2) data_b — output when criteria(control, threshold) is False.

Output ports

(0) The selected data signal, with shape determined by broadcasting between data_a and data_b.

Source code in jaxonomy/library/logic.py

class Switch(LeafSystem):
    """Route one of two data signals based on a thresholded control signal.

    Three inputs ``(data_a, control, data_b)`` and one output:

    .. code-block:: python

        y = data_a if criteria(control, threshold) else data_b

    Default ``mode="where"`` is implemented via ``npa.where``, so JAX
    gradients flow through *both* data branches simultaneously (the
    selector branch is treated as non-differentiable, which is the
    only well-defined choice for a hard threshold).

    ``mode="smooth"`` replaces the hard ``where`` with a sigmoid blend

    .. code-block:: python

        alpha  = sigmoid(sharpness * sign * (control - threshold))
        y      = alpha * data_a + (1 - alpha) * data_b

    where ``sign`` is +1 for the ``>=``/``>`` criteria and -1 for the
    ``<=``/``<`` criteria, so the smooth output approaches the hard
    answer in the strict-active region as ``sharpness -> inf``. This
    mode lets gradients flow through the *threshold itself*, which the
    hard ``where`` zeroes out — the killer feature for trajectory
    optimization where the threshold is a tunable parameter.

    ``data_a`` and ``data_b`` must be broadcast-compatible (same as
    ``npa.where``'s requirements). The block does not enforce that
    ``control`` is a scalar — element-wise selection is supported when
    ``control`` and the data inputs broadcast together.

    Parameters:
        threshold: scalar threshold against which ``control`` is compared.
        criteria: one of ``">="``, ``">"``, ``"<="``, ``"<"``, ``"=="``,
            ``"!="``. Default ``">="``. The equality criteria
            (``"=="``/``"!="``) are not supported in ``mode="smooth"``
            (no sigmoid approximation makes sense for them).
        mode: one of ``"where"`` (default), ``"smooth"``, or ``"hard"``.
            ``"where"`` is byte-equivalent to T-118 phase 1 and the
            right pick for simulation. ``"smooth"`` is the right pick
            for gradient-based optimization through the threshold.
            ``"hard"`` dispatches to ``jax.lax.cond`` so only the
            active branch is evaluated — useful when one branch is
            much more expensive than the other or when branches have
            incompatible side effects. **``Switch(mode='hard')`` is
            incompatible with vmap; use mode='where' or mode='smooth'
            for batched use** (e.g. under ``simulate_batch``). On
            older JAX (<0.4) a batched predicate raises
            ``TracerBoolConversionError``; on modern JAX the cond is
            silently rewritten to evaluate both branches with a
            select, which is numerically correct but defeats the
            entire point of picking ``mode='hard'`` over ``mode='where'``.
        sharpness: positive scalar controlling sigmoid steepness in
            ``mode="smooth"``. Default ``10.0``. Larger values give a
            tighter approximation to the hard switch but smaller (and
            faster vanishing) gradients in the strict-active region.
            Ignored when ``mode="where"``.

    Input ports:
        (0) data_a — output when criteria(control, threshold) is True.
        (1) control — the selector signal compared to ``threshold``.
        (2) data_b — output when criteria(control, threshold) is False.

    Output ports:
        (0) The selected data signal, with shape determined by
            broadcasting between data_a and data_b.
    """

    @parameters(static=["threshold", "criteria", "mode"], dynamic=["sharpness"])
    def __init__(
        self,
        threshold=0.0,
        criteria=">=",
        mode="where",
        sharpness=10.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if criteria not in _SWITCH_CRITERIA:
            raise BlockParameterError(
                message=(
                    f"Switch block '{self.name}' has invalid selection "
                    f"'{criteria}' for parameter 'criteria'. Valid options: "
                    + ",".join(_SWITCH_CRITERIA.keys())
                ),
                system=self,
                parameter_name="criteria",
            )

        if mode not in _SWITCH_VALID_MODES:
            raise BlockParameterError(
                message=(
                    f"Switch block '{self.name}' has invalid selection "
                    f"'{mode}' for parameter 'mode'. Valid options: "
                    + ",".join(_SWITCH_VALID_MODES)
                    + "."
                ),
                system=self,
                parameter_name="mode",
            )

        if mode == "smooth":
            if criteria not in _SWITCH_SMOOTH_SIGN:
                raise BlockParameterError(
                    message=(
                        f"Switch block '{self.name}': mode='smooth' is not "
                        f"compatible with criteria='{criteria}' (no sigmoid "
                        "approximation defined for equality). Use one of: "
                        + ",".join(_SWITCH_SMOOTH_SIGN.keys())
                        + "."
                    ),
                    system=self,
                    parameter_name="criteria",
                )
            if not np.isfinite(sharpness) or sharpness <= 0:
                raise BlockParameterError(
                    message=(
                        f"Switch block '{self.name}': sharpness must be "
                        f"finite and > 0, got {sharpness}."
                    ),
                    system=self,
                    parameter_name="sharpness",
                )

        self.declare_input_port()  # data_a
        self.declare_input_port()  # control
        self.declare_input_port()  # data_b
        self._output_port_idx = self.declare_output_port()

    def initialize(self, threshold, criteria, mode, sharpness=10.0):
        if mode == "where":
            compare = _SWITCH_CRITERIA[criteria]

            def _compute_output(_time, _state, *inputs, **_params):
                data_a, control, data_b = inputs
                return npa.where(compare(control, threshold), data_a, data_b)
        elif mode == "hard":
            # Local import keeps the JAX dependency lazy for non-JAX
            # backends that load primitives.py for class definitions.
            # ``lax.cond`` evaluates only the active branch — the whole
            # point of hard mode — but requires a scalar predicate, so
            # this path is incompatible with vmap (and therefore with
            # ``simulate_batch``). Documented at the class level.
            from jax import lax as _jlax

            compare = _SWITCH_CRITERIA[criteria]

            def _compute_output(_time, _state, *inputs, **_params):
                data_a, control, data_b = inputs
                pred = compare(control, threshold)
                # operand-passing form keeps the closure pure and lets
                # XLA elide unused captures; both branches must return
                # the same shape/dtype, which is broadcast-compatible
                # by the same rule as npa.where in the where path.
                return _jlax.cond(
                    pred,
                    lambda ops: ops[0],
                    lambda ops: ops[1],
                    (data_a, data_b),
                )
        else:
            # mode == "smooth"; validated in __init__ so the lookup is safe.
            sign = _SWITCH_SMOOTH_SIGN[criteria]

            def _compute_output(_time, _state, *inputs, **params):
                data_a, control, data_b = inputs
                # sharpness flows through dynamic params so optimizers can
                # anneal it; threshold is static so it appears in the
                # closure as a Python value (still differentiable via
                # jax.grad on the underlying op — the sigmoid is smooth
                # in `threshold`).
                k = params.get("sharpness", sharpness)
                alpha = _switch_sigmoid(k * sign * (control - threshold))
                return alpha * data_a + (1.0 - alpha) * data_b

        self.configure_output_port(
            self._output_port_idx,
            _compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

TableSearch

Bases: LeafSystem

Search a monotonic table for the bucket containing a query value.

Given a strictly-increasing 1-D grid xp of length n and a scalar query x, returns the bucket index i (as a float) such that xp[i] <= x < xp[i+1]. Out-of-range queries clamp to the nearest endpoint: x < xp[0] returns 0; x >= xp[-1] returns n - 1.

The standard "Direct Lookup" pattern. Different from :class:Prelookup in that the output is just the bucket index -- no fractional alpha is computed. Useful for binning, threshold detection, and inverse-table indexing.

Input ports

(0) -- scalar query coordinate x.

Output ports

(0) -- scalar bucket index, returned as a float (so it composes with the float-defaulting numeric pipeline). The output is wrapped in jax.lax.stop_gradient -- gradient is zero almost everywhere by construction (step function), so we make that non-differentiability explicit to avoid spurious grad-flow surprises.

Parameters:

Name	Type	Description	Default
xp		1-D, strictly-monotonically-increasing grid of breakpoints (length >= 2). Stored verbatim for the bucket search.	required
mode		"binary" (default) uses jnp.searchsorted (O(log n)); "linear" uses a linear scan via a sum over a comparison mask (O(n)). Both modes return byte-identical results on valid (strictly-monotonic) grids.	'binary'
dtype	optional	If set (e.g. jnp.float32), the grid array is cast to this dtype on construction. Mirrors the per-block dtype contract of :class:LookupTable1d / :class:Prelookup.	None

Notes

Index is wrapped in jax.lax.stop_gradient -- the gradient through the query coordinate is zero, by construction. Callers who need a differentiable index-like quantity should use :class:Prelookup (which exposes the fractional alpha) or :class:LookupTable1d directly.

Source code in jaxonomy/library/tables.py

class TableSearch(LeafSystem):
    """Search a monotonic table for the bucket containing a query value.

    Given a strictly-increasing 1-D grid ``xp`` of length ``n`` and a
    scalar query ``x``, returns the bucket index ``i`` (as a float) such
    that ``xp[i] <= x < xp[i+1]``.  Out-of-range queries clamp to the
    nearest endpoint: ``x < xp[0]`` returns ``0``; ``x >= xp[-1]``
    returns ``n - 1``.

    The standard "Direct Lookup" pattern.  Different from
    :class:`Prelookup` in that the output is just the bucket index --
    no fractional ``alpha`` is computed.  Useful for binning,
    threshold detection, and inverse-table indexing.

    Input ports:
        ``(0)`` -- scalar query coordinate ``x``.

    Output ports:
        ``(0)`` -- scalar bucket index, returned as a float (so it
        composes with the float-defaulting numeric pipeline).  The
        output is wrapped in ``jax.lax.stop_gradient`` -- gradient is
        zero almost everywhere by construction (step function), so we
        make that non-differentiability explicit to avoid spurious
        grad-flow surprises.

    Parameters:
        xp:
            1-D, strictly-monotonically-increasing grid of breakpoints
            (length >= 2).  Stored verbatim for the bucket search.
        mode:
            ``"binary"`` (default) uses ``jnp.searchsorted`` (O(log n));
            ``"linear"`` uses a linear scan via a sum over a comparison
            mask (O(n)).  Both modes return byte-identical results on
            valid (strictly-monotonic) grids.
        dtype (optional):
            If set (e.g. ``jnp.float32``), the grid array is cast to
            this dtype on construction.  Mirrors the per-block dtype
            contract of :class:`LookupTable1d` / :class:`Prelookup`.

    Notes:
        Index is wrapped in ``jax.lax.stop_gradient`` -- the gradient
        through the query coordinate is zero, by construction.  Callers
        who need a differentiable index-like quantity should use
        :class:`Prelookup` (which exposes the fractional ``alpha``) or
        :class:`LookupTable1d` directly.
    """

    def __init__(self, xp, mode="binary", dtype=None, **kwargs):
        # Per-block dtype + active precision policy fallback, matching
        # the LookupTable1d / Prelookup contract.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype

        if mode not in ("binary", "linear"):
            raise ValueError(
                f"TableSearch: mode must be one of ('binary','linear'), "
                f"got {mode!r}"
            )
        self._mode = mode

        # Eagerly validate the grid up front so the failure mode is a
        # clear ValueError on construction, not a cryptic shape error
        # at trace time.
        _xp_np = np.asarray(xp)
        if _xp_np.ndim != 1:
            raise ValueError(
                f"TableSearch: xp must be 1-D, got shape {_xp_np.shape}"
            )
        if _xp_np.size < 2:
            raise ValueError(
                f"TableSearch: xp must have at least 2 entries, got "
                f"shape {_xp_np.shape}"
            )
        if not np.all(np.diff(_xp_np) > 0):
            raise ValueError(
                f"TableSearch: xp must be strictly monotonically "
                f"increasing, got {list(_xp_np)}"
            )

        if self._dtype is not None:
            self._xp = npa.asarray(_xp_np).astype(self._dtype)
        else:
            self._xp = npa.array(_xp_np)

        super().__init__(**kwargs)
        self.declare_input_port()

        # Capture locals so the closure does not pull ``self`` into the
        # JAX trace.
        xp_local = self._xp
        n_local = int(self._xp.shape[0])
        mode_local = self._mode

        # Lazy-import jax.lax for the stop_gradient wrap.  Keeping this
        # at module-import time would defeat the lazy-loader pattern
        # used by the rest of this file for ``equinox``.
        from jax import lax as _jlax

        def _compute(_time, _state, *inputs, **_params):
            (x_query,) = inputs
            if mode_local == "binary":
                # ``side="right"`` then subtract 1 gives the bucket index
                # ``i`` s.t. ``xp[i] <= x < xp[i+1]``.  Clip to
                # ``[0, n - 1]`` for OOB clamping (left -> 0, right ->
                # n - 1; matches the "clip" policy used elsewhere in
                # T-114).
                i = npa.clip(
                    npa.searchsorted(xp_local, x_query, side="right") - 1,
                    0,
                    n_local - 1,
                )
            else:
                # ``"linear"`` mode -- count the breakpoints strictly
                # less-or-equal to ``x_query``, then subtract 1 to land
                # on the bucket index.  Same OOB clamp semantics as the
                # binary path.  Implemented as a sum over a comparison
                # mask so it stays jit/vmap-safe.
                mask = xp_local <= x_query
                count = npa.sum(mask.astype(xp_local.dtype))
                i = npa.clip(count - 1, 0, n_local - 1)
            # Return as a float so the output composes with the
            # library's float-defaulting numeric pipeline (T-005).
            # Wrap in ``stop_gradient`` to make the
            # piecewise-constant-gradient non-differentiability
            # explicit.
            i_float = npa.asarray(i).astype(xp_local.dtype)
            return _jlax.stop_gradient(i_float)

        self.declare_output_port(
            _compute,
            prerequisites_of_calc=[self.input_ports[0].ticket],
            requires_inputs=True,
            default_value=npa.asarray(0.0, dtype=xp_local.dtype),
        )

    @property
    def xp(self):
        """The 1-D strictly-increasing breakpoint array."""
        return self._xp

    @property
    def mode(self):
        """Search mode (``"binary"`` or ``"linear"``)."""
        return self._mode

mode property

Search mode ("binary" or "linear").

xp property

The 1-D strictly-increasing breakpoint array.

TensorFlow

Bases: LeafSystem

Block to perform inference with a pre-trained TensorFlow SavedModel.

The input to the block should be of compatible type and shape expected by the TensorFlow model. For example, if the TensorFlow SavedModel model expects a tf.float32 tensor of shape (3, 224, 224), the input to the block should be a jax.numpy array of shape (3, 224, 224) of dtype jnp.float32.

For output types, if no casting is specified through the cast_outputs_to_dtype parameter, the output of the block will have the same dtype as the TensorFlow model output, but expressed as jax.numpy types. For example. if the TensorFlow model outputs a tf.float32 tensor, the output of the block will be a jax.numpy array of dtype jnp.float32.

If casting is specified through cast_outputs_to_dtype parameter, all the outputs, of the block will be casted to this specific jax.numpy dtype.

.. note:: float32 models under jaxonomy's global x64. import jaxonomy enables JAX 64-bit mode (jax_enable_x64) for the whole process, so upstream signals are float64 by default. A SavedModel with tf.float32 signatures therefore receives float64 inputs (a dtype error or a silent arithmetic change) unless you cast at the block boundary. One-line idiom: pass cast_outputs_to_dtype="float32" and feed the block x.astype(jnp.float32) inputs.

Input ports

(i) The ith input to the model.

Output ports

(j) The jth output of the model.

Parameters:

Name	Type	Description	Default
file_name	str	Path to the model file. This should be a .zip containing the SavedModel.	required
cast_outputs_to_dtype	str	The dtype to cast all the outputs of the block to. Must correspond to a jax.numpy datatype. For example, "float32", "float64", "int32", "int64".	None
add_batch_dim_to_inputs	bool	Whether to add a new first dimension to the inputs before evaluating the TorchScript or TensorFlow model. This is useful when the model expects a batch dimension.	False

Source code in jaxonomy/library/predictor.py

class TensorFlow(LeafSystem):
    """
    Block to perform inference with a pre-trained TensorFlow SavedModel.

    The input to the block should be of compatible type and shape expected by
    the TensorFlow model. For example,  if the TensorFlow SavedModel model expects a
    `tf.float32` tensor of shape `(3, 224, 224)`, the input to the block should be a
    `jax.numpy` array of shape (3, 224, 224) of dtype `jnp.float32`.

    For output types, if no casting is specified through the `cast_outputs_to_dtype`
    parameter, the output of the block will have the same dtype as the
    TensorFlow model output, but expressed as `jax.numpy` types. For example. if the
    TensorFlow model outputs a `tf.float32` tensor, the output of the block will be
    a `jax.numpy` array of dtype `jnp.float32`.

    If casting is specified through `cast_outputs_to_dtype` parameter, all the outputs,
    of the block will be casted to this specific `jax.numpy` dtype.

    .. note:: **float32 models under jaxonomy's global x64.**
       ``import jaxonomy`` enables JAX 64-bit mode (``jax_enable_x64``)
       for the whole process, so upstream signals are float64 by
       default.  A SavedModel with ``tf.float32`` signatures therefore
       receives float64 inputs (a dtype error or a silent arithmetic
       change) unless you cast at the block boundary.  One-line idiom:
       pass ``cast_outputs_to_dtype="float32"`` and feed the block
       ``x.astype(jnp.float32)`` inputs.

    Input ports:
        (i) The ith input to the model.

    Output ports:
        (j) The jth output of the model.

    Parameters:
        file_name (str):
            Path to the model file. This should be a `.zip` containing the SavedModel.

        cast_outputs_to_dtype (str):
            The dtype to cast all the outputs of the block to. Must correspond to a
            `jax.numpy` datatype. For example, "float32", "float64", "int32", "int64".

        add_batch_dim_to_inputs (bool):
            Whether to add a new first dimension to the inputs before evaluating the
            TorchScript or TensorFlow model. This is useful when the model expects a
            batch dimension.
    """

    @parameters(
        static=[
            "file_name",
            "cast_outputs_to_dtype",
            "add_batch_dim_to_inputs",
        ]
    )
    def __init__(
        self,
        file_name,
        cast_outputs_to_dtype=None,
        add_batch_dim_to_inputs=False,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        model, num_inputs, num_outputs, num_args, kwargs_signature = self._load_model(
            file_name
        )

        self.num_inputs = num_inputs
        self.num_outputs = num_outputs

        for _ in range(self.num_inputs):
            self.declare_input_port()

        def _make_output_callback(output_index):
            def _output_callback(time, state, *inputs, **params):
                outputs = self._evaluate_output(time, state, *inputs, **params)
                return outputs[output_index]

            return _output_callback

        for output_index in range(self.num_outputs):
            self.declare_output_port(
                _make_output_callback(output_index),
                requires_inputs=True,
            )

    def _load_model(self, file_name):
        _, ext = os.path.splitext(file_name)

        if ext == ".zip":
            with tempfile.TemporaryDirectory() as model_dir:
                with zipfile.ZipFile(file_name, "r") as zip_ref:
                    zip_ref.extractall(model_dir)

                model = tf.saved_model.load(model_dir)

            model = model.signatures["serving_default"]

            num_args = len(model.structured_input_signature[0])
            kwargs_signature = model.structured_input_signature[1]
            # Assure deterministic dict order
            kwargs_signature = {k: kwargs_signature[k] for k in sorted(kwargs_signature.keys())}
            num_kwargs = len(kwargs_signature)

            num_inputs = num_args + num_kwargs
            num_outputs = len(model.structured_outputs)
        else:
            raise ValueError(f"Expected extension of file is `.zip`, but found {ext}")

        return model, num_inputs, num_outputs, num_args, kwargs_signature

    def initialize(
        self,
        file_name,
        cast_outputs_to_dtype=None,
        add_batch_dim_to_inputs=False,
    ):
        self.dtype_output = (
            getattr(jnp, cast_outputs_to_dtype)
            if cast_outputs_to_dtype is not None
            else None
        )

        self.add_batch_dim_to_inputs = add_batch_dim_to_inputs

        model, num_inputs, num_outputs, num_args, kwargs_signature = self._load_model(
            file_name
        )

        if self.num_inputs != num_inputs:
            raise ValueError("num_inputs can't be changed after initialization")
        if self.num_outputs != num_outputs:
            raise ValueError("num_outputs can't be changed after initialization")

        self.model = model
        self.num_args = num_args
        self.kwargs_signature = kwargs_signature
        self.num_kwargs = len(self.kwargs_signature)

    def initialize_static_data(self, context):
        """Infer the output shapes and dtypes of the ML model."""
        # If building as part of a subsystem, this may not be fully connected yet.
        # That's fine, as long as it is connected by root context creation time.
        # This probably isn't a good long-term solution:
        #   see https://jaxonomy.atlassian.net/browse/WC-51
        try:
            inputs = self.collect_inputs(context)

            outputs_jax = self._pure_callback(*inputs)

            self.pure_callback_result_type = [
                jax.ShapeDtypeStruct(x.shape, x.dtype) for x in outputs_jax
            ]
        except UpstreamEvalError:
            logger.debug(
                "Predictor.initialize_static_data: UpstreamEvalError. "
                "Continuing without default value initialization."
            )
        return super().initialize_static_data(context)

    def _evaluate_output(self, time, state, *inputs, **params):
        return jax.pure_callback(
            self._pure_callback,
            self.pure_callback_result_type,
            *inputs,
        )

    def _pure_callback(self, *inputs):
        inputs_casted = [
            tf.convert_to_tensor(np.array(item), dtype=sig.dtype)
            for item, sig in zip(inputs, self.kwargs_signature.values())
        ]
        args_casted = inputs_casted[: self.num_args]

        # Map sorted kwargs directly to inputs
        kwargs_casted = dict(
            zip(self.kwargs_signature.keys(), inputs_casted[self.num_args :])
        )

        if self.add_batch_dim_to_inputs:
            args_casted = [tf.expand_dims(x, axis=0) for x in args_casted]
            kwargs_casted = {
                key: tf.expand_dims(value, axis=0)
                for key, value in kwargs_casted.items()
            }

        if self.num_args == 0:
            outputs_dict = self.model(**kwargs_casted)
        else:
            outputs_dict = self.model(*args_casted, **kwargs_casted)

        outputs_jax = (
            [jnp.array(outputs_dict[k], self.dtype_output) for k in sorted(outputs_dict.keys())]
            if self.dtype_output is not None
            else [jnp.array(outputs_dict[k]) for k in sorted(outputs_dict.keys())]
        )
        return outputs_jax

initialize_static_data(context)

Infer the output shapes and dtypes of the ML model.

Source code in jaxonomy/library/predictor.py

def initialize_static_data(self, context):
    """Infer the output shapes and dtypes of the ML model."""
    # If building as part of a subsystem, this may not be fully connected yet.
    # That's fine, as long as it is connected by root context creation time.
    # This probably isn't a good long-term solution:
    #   see https://jaxonomy.atlassian.net/browse/WC-51
    try:
        inputs = self.collect_inputs(context)

        outputs_jax = self._pure_callback(*inputs)

        self.pure_callback_result_type = [
            jax.ShapeDtypeStruct(x.shape, x.dtype) for x in outputs_jax
        ]
    except UpstreamEvalError:
        logger.debug(
            "Predictor.initialize_static_data: UpstreamEvalError. "
            "Continuing without default value initialization."
        )
    return super().initialize_static_data(context)

TransferFunction

Bases: LTISystem

Continuous-time LTI system specified as a transfer function.

The transfer function is converted to state-space form using scipy.signal.tf2ss. https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.tf2ss.html

The resulting system will be in canonical controller form with matrices (A, B, C, D), which are then used to create an LTISystem. Note that this only supports single-input, single-output systems.

Input ports

(0) u: Input vector (scalar)

Output ports

(0) y: Output vector (scalar). Note that this is feedthrough from the input port iff D is nonzero.

Parameters:

Name	Type	Description	Default
num		Numerator polynomial coefficients, in descending powers of s	required
den		Denominator polynomial coefficients, in descending powers of s	required

Source code in jaxonomy/library/linear_system.py

class TransferFunction(LTISystem):
    """Continuous-time LTI system specified as a transfer function.

    The transfer function is converted to state-space form using `scipy.signal.tf2ss`.
    https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.tf2ss.html

    The resulting system will be in canonical controller form with matrices
    (A, B, C, D), which are then used to create an LTISystem.  Note that this only
    supports single-input, single-output systems.

    Input ports:
        (0) u: Input vector (scalar)

    Output ports:
        (0) y: Output vector (scalar).  Note that this is feedthrough from the input
            port iff D is nonzero.

    Parameters:
        num: Numerator polynomial coefficients, in descending powers of s
        den: Denominator polynomial coefficients, in descending powers of s
    """

    # tf2ss is not implemented in jax.scipy.signal so num and den can't be
    # dynamic parameters.
    @parameters(static=["num", "den"])
    def __init__(self, num, den, *args, **kwargs):
        A, B, C, D = signal.tf2ss(num, den)
        self._num = num
        self._den = den
        super().__init__(A, B, C, D, *args, **kwargs)

    def initialize(self, num, den, **kwargs):
        A, B, C, D = signal.tf2ss(num, den)
        self._num = num
        self._den = den
        self._init_state(A, B, C, D)
        self.parameters["A"].set(self.A)
        self.parameters["B"].set(self.B)
        self.parameters["C"].set(self.C)
        self.parameters["D"].set(self.D)

TransferFunctionDiscrete

Bases: LTISystemDiscrete

Implements a Discrete Time Transfer Function.

https://en.wikipedia.org/wiki/Z-transform#Transfer_function

The resulting system will be in canonical controller form with matrices (A, B, C, D), which are then used to create an LTISystem. Note that this only supports single-input, single-output systems.

Input ports

(0) u[k]: Input vector (scalar)

Output ports

(0) y[k]: Output vector (scalar). Note that this is feedthrough from the input port if and only if D is nonzero.

Parameters:

Name	Type	Description	Default
dt		Sampling period of the discrete system.	required
num		Numerator polynomial coefficients, in descending powers of z	required
den		Denominator polynomial coefficients, in descending powers of z	required
initialize_states		Initial state vector (default: 0)	None

Source code in jaxonomy/library/linear_system.py

class TransferFunctionDiscrete(LTISystemDiscrete):
    """Implements a Discrete Time Transfer Function.

    https://en.wikipedia.org/wiki/Z-transform#Transfer_function

    The resulting system will be in canonical controller form with matrices
    (A, B, C, D), which are then used to create an LTISystem.  Note that this only
    supports single-input, single-output systems.

    Input ports:
        (0) u[k]: Input vector (scalar)

    Output ports:
        (0) y[k]: Output vector (scalar). Note that this is feedthrough from the input
            port if and only if D is nonzero.

    Parameters:
        dt:
            Sampling period of the discrete system.
        num:
            Numerator polynomial coefficients, in descending powers of z
        den:
            Denominator polynomial coefficients, in descending powers of z
        initialize_states:
            Initial state vector (default: 0)
    """

    # tf2ss is not implemented in jax.scipy.signal so num and den can't be
    # dynamic parameters.
    @parameters(static=["num", "den"])
    def __init__(self, dt, num, den, initialize_states=None, *args, **kwargs):
        A, B, C, D = signal.tf2ss(num, den)
        super().__init__(A, B, C, D, dt, initialize_states, *args, **kwargs)

    def _eval_output(self, time, state, *inputs, **params):
        return super()._eval_output(
            time, state, *inputs, A=self.A, B=self.B, C=self.C, D=self.D
        )

    def _update(self, time, state, u, **params):
        return super()._update(time, state, u, A=self.A, B=self.B)

    def initialize(self, num, den, **kwargs):
        A, B, C, D = signal.tf2ss(num, den)
        self._init_state(A, B, C, D)

TransportDelay

Bases: LeafSystem

Continuous-time fixed transport (pure) delay.

Implements y(t) = u(t - delay_seconds) for t >= delay_seconds; for t < delay_seconds the output is initial_output (the standard "Initial output" semantics).

The block samples its input on a periodic clock with period dt and stores the most recent history_length (time, value) pairs in a discrete-state ring buffer. Output evaluation at any continuous time t is a linear interpolation over the buffered (time, value) samples at t - delay_seconds.

The delay is differentiable via the input signal (gradient flows through npa.interp over the values buffer). Differentiability w.r.t. the delay value itself is well-defined wherever the buffer interpolant is differentiable; the linear interpolant has a kink at sample boundaries — pass method="pchip" on :class:VariableTransportDelay for a C¹-smooth alternative.

The buffer is sized statically as history_length samples. To cover a delay of delay_seconds at sample period dt, you need at least ceil(delay_seconds / dt) + 1 slots; we recommend a small safety margin. history_length defaults to max(8, ceil(delay_seconds / dt) + 4) which is sufficient for the default constant delay.

Input ports

(0) The input signal u(t). Scalar or array.

Output ports

(0) The delayed signal y(t) = u(t - delay_seconds) (or initial_output while t < delay_seconds).

Parameters:

Name	Type	Description	Default
dt		Sampling period for the history buffer. Smaller dt ⇒ finer interpolation but a larger ring buffer to cover the same physical delay.	required
delay_seconds		Fixed delay τ in seconds. Dynamic parameter (may be tuned via with_parameters); see notes on differentiability above.	required
initial_output		Output value while t < delay_seconds. Default is 0.0.	0.0
history_length		Number of (time, value) pairs stored. Static (compile-time) — required for vmap/JIT-safe buffer sizing. If None, defaults to max(8, ceil(delay_seconds / dt) + 4).	None

Notes

• For arbitrary array-shaped signals the interpolation is applied elementwise via jax.vmap over the trailing axes.
• Buffer overflow (delay larger than history_length * dt) is not raised; npa.interp clamps to the boundary, which means the oldest stored sample is repeated. This is a documented T-107 follow-up; for now, size history_length generously.
• VariableTransportDelay (signal-driven τ) is the natural phase-2 extension; it reuses the same ring-buffer machinery with the delay sourced from an input port.

Source code in jaxonomy/library/dynamics.py

class TransportDelay(LeafSystem):
    """Continuous-time fixed transport (pure) delay.

    Implements ``y(t) = u(t - delay_seconds)`` for ``t >= delay_seconds``;
    for ``t < delay_seconds`` the output is ``initial_output`` (the
    standard "Initial output" semantics).

    The block samples its input on a periodic clock with period ``dt`` and
    stores the most recent ``history_length`` ``(time, value)`` pairs in a
    discrete-state ring buffer. Output evaluation at any continuous time
    ``t`` is a linear interpolation over the buffered ``(time, value)``
    samples at ``t - delay_seconds``.

    The delay is differentiable via the input signal (gradient flows
    through ``npa.interp`` over the values buffer). Differentiability
    w.r.t. the delay value itself is well-defined wherever the buffer
    interpolant is differentiable; the linear interpolant has a kink at
    sample boundaries — pass ``method="pchip"`` on
    :class:`VariableTransportDelay` for a C¹-smooth alternative.

    The buffer is sized statically as ``history_length`` samples. To cover
    a delay of ``delay_seconds`` at sample period ``dt``, you need at
    least ``ceil(delay_seconds / dt) + 1`` slots; we recommend a small
    safety margin. ``history_length`` defaults to
    ``max(8, ceil(delay_seconds / dt) + 4)`` which is sufficient for the
    default constant delay.

    Input ports:
        (0) The input signal ``u(t)``. Scalar or array.

    Output ports:
        (0) The delayed signal ``y(t) = u(t - delay_seconds)`` (or
            ``initial_output`` while ``t < delay_seconds``).

    Parameters:
        dt: Sampling period for the history buffer. Smaller ``dt`` ⇒
            finer interpolation but a larger ring buffer to cover the
            same physical delay.
        delay_seconds: Fixed delay τ in seconds. Dynamic parameter (may
            be tuned via ``with_parameters``); see notes on
            differentiability above.
        initial_output: Output value while ``t < delay_seconds``. Default
            is 0.0.
        history_length: Number of ``(time, value)`` pairs stored. Static
            (compile-time) — required for vmap/JIT-safe buffer sizing.
            If None, defaults to ``max(8, ceil(delay_seconds / dt) + 4)``.

    Notes:
        - For arbitrary array-shaped signals the interpolation is applied
          elementwise via ``jax.vmap`` over the trailing axes.
        - Buffer overflow (delay larger than ``history_length * dt``) is
          not raised; ``npa.interp`` clamps to the boundary, which means
          the oldest stored sample is repeated. This is a documented
          T-107 follow-up; for now, size ``history_length`` generously.
        - ``VariableTransportDelay`` (signal-driven τ) is the natural
          phase-2 extension; it reuses the same ring-buffer machinery
          with the delay sourced from an input port.
    """

    class _BufferState(NamedTuple):
        # Newest sample at index 0; oldest at index -1. Reversed buffers
        # form a monotonically increasing time axis for ``npa.interp``.
        times: "Array"
        values: "Array"

    @parameters(
        static=["dt", "history_length"],
        dynamic=["delay_seconds", "initial_output"],
    )
    def __init__(
        self,
        dt,
        delay_seconds,
        initial_output=0.0,
        history_length=None,
        *args,
        dtype=None,
        **kwargs,
    ):
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(*args, **kwargs)

        if dt is None or float(dt) <= 0.0:
            raise BlockParameterError(
                message=(
                    f"TransportDelay block {self.name!r} requires a positive "
                    f"sample period dt; got {dt!r}."
                ),
                parameter_name="dt",
            )
        try:
            delay_hint = float(delay_seconds)
        except (TypeError, ValueError):
            delay_hint = 0.0
        if history_length is None:
            history_length = max(8, int(np.ceil(max(delay_hint, 0.0) / dt)) + 4)
        if int(history_length) < 2:
            raise BlockParameterError(
                message=(
                    f"TransportDelay block {self.name!r} requires "
                    f"history_length >= 2; got {history_length!r}."
                ),
                parameter_name="history_length",
            )

        self.dt = float(dt)
        self.history_length = int(history_length)

        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, dt, delay_seconds, initial_output, history_length=None):
        # ``history_length`` is a static parameter resolved at __init__
        # time; the framework still passes it here for symmetry, but we
        # rely on ``self.history_length`` to size buffers.
        del history_length  # noqa: F841 — silence unused-arg lint

        initial_value = npa.asarray(initial_output)
        if self._dtype is not None:
            initial_value = initial_value.astype(self._dtype)
        self._signal_shape = tuple(initial_value.shape)

        # Pre-fill the times buffer with strictly increasing sentinels
        # below t=0 so that ``npa.interp(t - delay, times[::-1], ...)``
        # clamps to the oldest sample (== ``initial_output``) for any
        # query time before the first real sample has been written. The
        # spacing matches ``dt`` so the reversed time axis stays
        # monotonically increasing.
        sentinel_t0 = -self.dt * (self.history_length + 1) - 1.0
        times = sentinel_t0 + self.dt * np.arange(self.history_length, dtype=np.float64)
        # Newest first: reverse so position 0 is the largest sentinel.
        times = times[::-1].copy()
        if self._dtype is not None:
            times = times.astype(self._dtype)

        values = npa.broadcast_to(
            initial_value, (self.history_length, *self._signal_shape)
        )

        default_state = self._BufferState(
            times=npa.asarray(times), values=npa.asarray(values)
        )
        self.declare_discrete_state(default_value=default_state, as_array=False)

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=0.0,
        )

        # Output is continuous-time (depends on ``time`` and on the
        # discrete-state buffer) — no period; ``requires_inputs=False``
        # because the lookup reads only state + time.
        self.configure_output_port(
            self._output_port_idx,
            self._output,
            prerequisites_of_calc=[DependencyTicket.xd, DependencyTicket.time],
            requires_inputs=False,
            default_value=initial_value,
        )

    def reset_default_values(
        self, dt=None, delay_seconds=None, initial_output=None, history_length=None
    ):
        # Mirror UnitDelay's pattern: rebuild defaults if the dynamic
        # ``initial_output`` changes between calls.
        del dt, delay_seconds, history_length  # noqa: F841

        if initial_output is None:
            return
        initial_value = npa.asarray(initial_output)
        if self._dtype is not None:
            initial_value = initial_value.astype(self._dtype)
        self._signal_shape = tuple(initial_value.shape)

        sentinel_t0 = -self.dt * (self.history_length + 1) - 1.0
        times = sentinel_t0 + self.dt * np.arange(self.history_length, dtype=np.float64)
        times = times[::-1].copy()
        if self._dtype is not None:
            times = times.astype(self._dtype)

        values = npa.broadcast_to(
            initial_value, (self.history_length, *self._signal_shape)
        )
        default_state = self._BufferState(
            times=npa.asarray(times), values=npa.asarray(values)
        )
        self.configure_discrete_state_default_value(
            default_value=default_state, as_array=False
        )
        self.configure_output_port_default_value(
            self._output_port_idx, initial_value
        )

    def _update(self, time, state, *inputs, **_params):
        u = inputs[0]
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        buf = state.discrete_state
        new_times = npa.roll(buf.times, shift=1, axis=0).at[0].set(time)
        new_values = npa.roll(buf.values, shift=1, axis=0).at[0].set(u)
        return self._BufferState(times=new_times, values=new_values)

    def _output(self, time, state, *_inputs, **params):
        buf = state.discrete_state
        # Reverse so the time axis is monotonically increasing for
        # ``npa.interp``: index 0 is oldest, index -1 is newest.
        xp = buf.times[::-1]
        fp = buf.values[::-1]
        delay = params["delay_seconds"]
        initial_output = params["initial_output"]

        query_t = time - delay
        if len(self._signal_shape) == 0:
            y = npa.interp(query_t, xp, fp)
        else:
            # ``npa.interp`` only handles 1-D ``fp``; vmap over the
            # trailing axes of ``values``. ``fp`` has shape
            # ``(history_length, *signal_shape)``; reshape/iterate via
            # ``jax.numpy.apply_along_axis``-style by flattening the
            # trailing dims.
            flat_fp = fp.reshape((self.history_length, -1))
            # Loop over trailing dim count statically (it's a static
            # shape) — JIT-friendly because the loop unrolls.
            ys = [npa.interp(query_t, xp, flat_fp[:, i]) for i in range(flat_fp.shape[1])]
            y = npa.stack(ys).reshape(self._signal_shape)

        # Hold the initial output before the first physical sample is
        # available; ``npa.interp`` would otherwise return the boundary
        # of the (sentinel-filled) buffer, which already equals
        # ``initial_output`` — but explicitly gating on ``time`` keeps
        # the semantics robust to dtype/shape pre-fill quirks.
        y = npa.where(time < delay, npa.asarray(initial_output), y)
        if self._dtype is not None:
            y = npa.asarray(y).astype(self._dtype)
        return y

TriggerEdge

Allowed string values for TriggeredSubsystem.edge.

Source code in jaxonomy/framework/containers.py

class TriggerEdge:
    """Allowed string values for ``TriggeredSubsystem.edge``."""

    RISING = "rising"
    FALLING = "falling"
    EITHER = "either"

    @classmethod
    def valid(cls) -> tuple[str, ...]:
        return (cls.RISING, cls.FALLING, cls.EITHER)

TriggeredSubsystem

Bases: LeafSystem

Container block: latch the submodel output on edge transitions (the child still RUNS every step — only the output is gated).

Important: this does not skip execution of the submodel on non-triggered steps. The submodel is evaluated on every step so its inputs participate in the JAX trace; the trigger only controls whether a fresh result is latched into the held output. If you need to actually skip computation between triggers, gate it yourself with jax.lax.cond at the application level.

Phase-1 implementation runs the submodel on every step (so the inputs participate in the trace) but only latches a new output on an edge transition of the trigger signal. Between transitions the output holds the most recently latched value.

The trigger signal is sampled at sample_period. Edges are detected by comparing the current trigger sample against the previously-stored sample held in discrete state.

This is not the eventual zero-crossing-driven TriggeredSubsystem described in the T-120 architecture notes (that requires hooking into the continuous-time event detector); but it is functionally correct for any sample-rate use case and matches the behaviour documented in the test fixtures.

Parameters:

Name	Type	Description	Default
submodel	Callable	Callable f(*inputs) -> output taking the non-trigger user inputs. Must be JAX-traceable.	required
n_inputs	int	Number of user inputs (NOT counting the trigger).	1
edge	Literal['rising', 'falling', 'either']	"rising" (low→high), "falling" (high→low) or "either".	RISING
sample_period	float	Period (seconds) at which the trigger signal is sampled and the latch is updated. Must be positive.	0.0
initial_value		Latched output value before any edge has been detected. Defines output shape/dtype.	0.0
name		Optional block name.	required

Limitations (phase 1): - Trigger detection runs on the periodic sample grid, not on continuous-time zero crossings. Trigger pulses shorter than sample_period may be missed. - The latch is a single discrete state; the submodel must produce a single output array. - The submodel runs on every output evaluation; only the output is gated. Users who need to skip computation on non-triggered steps should use jax.lax.cond at the application level.

Source code in jaxonomy/framework/containers.py

class TriggeredSubsystem(LeafSystem):
    """Container block: latch the submodel output on edge transitions
    (the child still RUNS every step — only the *output* is gated).

    Important: this does **not** skip execution of the submodel on
    non-triggered steps. The submodel is evaluated on every step so its
    inputs participate in the JAX trace; the trigger only controls
    whether a fresh result is *latched* into the held output. If you
    need to actually skip computation between triggers, gate it yourself
    with ``jax.lax.cond`` at the application level.

    Phase-1 implementation runs the submodel on every step (so the
    inputs participate in the trace) but only *latches* a new output
    on an edge transition of the trigger signal. Between transitions
    the output holds the most recently latched value.

    The trigger signal is sampled at ``sample_period``. Edges are
    detected by comparing the current trigger sample against the
    previously-stored sample held in discrete state.

    This is *not* the eventual zero-crossing-driven ``TriggeredSubsystem``
    described in the T-120 architecture notes (that requires hooking
    into the continuous-time event detector); but it is functionally
    correct for any sample-rate use case and matches the behaviour
    documented in the test fixtures.

    Args:
        submodel: Callable ``f(*inputs) -> output`` taking the
            non-trigger user inputs. Must be JAX-traceable.
        n_inputs: Number of user inputs (NOT counting the trigger).
        edge: ``"rising"`` (low→high), ``"falling"`` (high→low) or
            ``"either"``.
        sample_period: Period (seconds) at which the trigger signal is
            sampled and the latch is updated. Must be positive.
        initial_value: Latched output value before any edge has been
            detected. Defines output shape/dtype.
        name: Optional block name.

    Limitations (phase 1):
        - Trigger detection runs on the periodic sample grid, not on
          continuous-time zero crossings. Trigger pulses shorter than
          ``sample_period`` may be missed.
        - The latch is a single discrete state; the submodel must
          produce a single output array.
        - The submodel runs on every output evaluation; only the
          *output* is gated. Users who need to skip computation on
          non-triggered steps should use ``jax.lax.cond`` at the
          application level.
    """

    def __init__(
        self,
        submodel: Callable,
        n_inputs: int = 1,
        edge: Literal["rising", "falling", "either"] = TriggerEdge.RISING,
        sample_period: float = 0.0,
        initial_value=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if edge not in TriggerEdge.valid():
            raise ValueError(
                f"TriggeredSubsystem: edge must be one of "
                f"{TriggerEdge.valid()!r}, got {edge!r}"
            )
        if sample_period is None or float(sample_period) <= 0.0:
            raise ValueError(
                "TriggeredSubsystem requires a positive sample_period "
                "(seconds) for trigger sampling."
            )
        if n_inputs < 0:
            raise ValueError(
                f"TriggeredSubsystem: n_inputs must be >= 0, got {n_inputs}"
            )

        self._submodel = submodel
        self._edge = edge
        self._sample_period = float(sample_period)
        self._initial = jnp.asarray(initial_value)

        # Port 0 is the trigger signal; ports 1..n_inputs are the user
        # inputs forwarded to the submodel.
        self.declare_input_port(name="trigger")
        for i in range(n_inputs):
            self.declare_input_port(name=f"u_{i}")

        # Discrete state pair: (latched_output, previous_trigger).
        # Pack as a flat 1-D array so the existing scalar-friendly
        # discrete-state machinery handles them uniformly. The two
        # pieces have different shapes in general, so use a tuple-like
        # tree via two separate periodic updates? Simpler: pack as a
        # dict. ``LeafSystem.declare_discrete_state`` only accepts a
        # single default_value, so we encode (latch, prev_trigger) as
        # a flat concatenation when both are scalar. For phase 1 we
        # require a scalar trigger so this packing is safe.
        #
        # Layout: discrete_state[..., 0] holds the previous trigger
        # sample; discrete_state[..., 1:] holds the latched output
        # (flattened). For a scalar latch this collapses to length 2.
        flat_init = jnp.concatenate(
            [
                jnp.zeros((1,), dtype=self._initial.dtype),
                jnp.atleast_1d(self._initial).reshape(-1),
            ]
        )
        self._latch_shape = self._initial.shape
        self._latch_size = int(jnp.atleast_1d(self._initial).reshape(-1).size)
        self.declare_discrete_state(default_value=flat_init)
        self.declare_periodic_update(
            self._latch_update,
            period=self._sample_period,
            offset=0.0,
        )
        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    # ── helpers ───────────────────────────────────────────────────────────

    def _unpack(self, ds):
        prev_trig = ds[0]
        latch_flat = ds[1 : 1 + self._latch_size]
        latch = latch_flat.reshape(self._latch_shape)
        return prev_trig, latch

    def _pack(self, prev_trig, latch):
        return jnp.concatenate(
            [
                jnp.atleast_1d(prev_trig).reshape(-1)[:1],
                jnp.atleast_1d(latch).reshape(-1),
            ]
        )

    def _edge_detected(self, prev_trig, cur_trig):
        prev_b = jnp.asarray(prev_trig).astype(bool)
        cur_b = jnp.asarray(cur_trig).astype(bool)
        if self._edge == TriggerEdge.RISING:
            return jnp.logical_and(jnp.logical_not(prev_b), cur_b)
        if self._edge == TriggerEdge.FALLING:
            return jnp.logical_and(prev_b, jnp.logical_not(cur_b))
        # either
        return jnp.not_equal(prev_b, cur_b)

    # ── callbacks ─────────────────────────────────────────────────────────

    def _submodel_output(self, inputs):
        user_inputs = inputs[1:]  # skip trigger
        return jnp.asarray(self._submodel(*user_inputs))

    def _latch_update(self, time, state, *inputs, **params):
        ds = state.discrete_state
        prev_trig, latch = self._unpack(ds)
        cur_trig = jnp.asarray(inputs[0])
        edge = self._edge_detected(prev_trig, cur_trig)
        y_sub = self._submodel_output(inputs)
        new_latch = jnp.where(edge, y_sub, latch)
        # Store the current trigger sample (cast to the same dtype as
        # the rest of the discrete-state vector).
        new_prev = jnp.asarray(cur_trig).astype(ds.dtype).reshape(())
        return self._pack(new_prev, new_latch)

    def _compute_output(self, time, state, *inputs, **params):
        # Output reads the latched value. The submodel is *also* invoked
        # via the latch_update path on the periodic sample grid; here we
        # additionally consult the current trigger so a same-step rising
        # edge surfaces immediately rather than one sample later.
        ds = state.discrete_state
        prev_trig, latch = self._unpack(ds)
        cur_trig = jnp.asarray(inputs[0])
        edge = self._edge_detected(prev_trig, cur_trig)
        y_sub = self._submodel_output(inputs)
        return jnp.where(edge, y_sub, latch)

Trigonometric

Bases: FeedthroughBlock

Apply a trigonometric function to the input signal.

Available functions are

sin, cos, tan, asin, acos, atan, sinh, cosh, tanh, asinh, acosh, atanh

Dispatches to jax.numpy.sin, jax.numpy.cos, etc, so see the JAX docs for details.

Input ports

(0) The input signal.

Output ports

(0) The trigonometric function applied to the input signal.

Parameters:

Name	Type	Description	Default
function		The trigonometric function to apply to the input signal. Must be one of "sin", "cos", "tan", "asin", "acos", "atan", "sinh", "cosh", "tanh", "asinh", "acosh", "atanh".	required

Source code in jaxonomy/library/math_ops.py

class Trigonometric(FeedthroughBlock):
    """Apply a trigonometric function to the input signal.

    Available functions are:
        sin, cos, tan, asin, acos, atan, sinh, cosh, tanh, asinh, acosh, atanh

    Dispatches to `jax.numpy.sin`, `jax.numpy.cos`, etc, so see the JAX docs for details.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The trigonometric function applied to the input signal.

    Parameters:
        function:
            The trigonometric function to apply to the input signal.  Must be one of
            "sin", "cos", "tan", "asin", "acos", "atan", "sinh", "cosh", "tanh",
            "asinh", "acosh", "atanh".
    """

    @parameters(static=["function"])
    def __init__(self, function, **kwargs):
        super().__init__(None, **kwargs)

    def initialize(self, function):
        func_lookup = {
            "sin": npa.sin,
            "cos": npa.cos,
            "tan": npa.tan,
            "asin": npa.arcsin,
            "acos": npa.arccos,
            "atan": npa.arctan,
            "sinh": npa.sinh,
            "cosh": npa.cosh,
            "tanh": npa.tanh,
            "asinh": npa.arcsinh,
            "acosh": npa.arccosh,
            "atanh": npa.arctanh,
        }
        if function not in func_lookup:
            raise BlockParameterError(
                message=f"Trigonometric block {self.name} has invalid selection {function} for 'function'. Valid options: "
                + ", ".join([f for f in func_lookup.keys()]),
                parameter_name="function",
            )
        self.replace_op(func_lookup[function])

TruthTable

Bases: LeafSystem

Evaluate a fixed truth table over boolean-castable inputs.

Given a list of (input_pattern, output) rows, this block compares its inputs against each pattern and emits the output of the first matching row (or default_output if none match). Patterns are tuples of bool values or the string "X" as a wildcard.

Example — a 2-input AND gate:

.. code-block:: python

tt = TruthTable(
    rows=[
        ((True, True), 1.0),
        ((True, False), 0.0),
        ((False, True), 0.0),
        ((False, False), 0.0),
    ],
    n_inputs=2,
    default_output=0.0,
)

Wildcard example — ignore the first input:

.. code-block:: python

tt = TruthTable(
    rows=[(("X", True), 1.0), (("X", False), 0.0)],
    n_inputs=2,
    default_output=0.0,
)

Callable output example — row output depends on raw input values (T-119-followup-numeric-output):

.. code-block:: python

tt = TruthTable(
    rows=[
        ((True, True), lambda a, b: a + b),
        ((True, False), lambda a, b: a - b),
        ((False, "X"), 0.0),  # constant fallback row
    ],
    n_inputs=2,
    default_output=0.0,
)

Parameters:

Name	Type	Description	Default
rows		list of (pattern, output) tuples. pattern is a length-n_inputs tuple whose entries are bool (matched literally) or the string "X" (wildcard, matches anything). output may be a scalar/array (constant for that row) or a callable f(*inputs) -> scalar_or_vector invoked with the RAW inputs (every row callable is evaluated at every step under JAX's branchless where semantics; the result is selected only when the row matches). All row outputs (and the callable return values) must broadcast against default_output.	required
n_inputs		number of input ports.	required
default_output		value emitted when no row matches. May be a scalar/array (constant fallback) or a callable f(*inputs) -> scalar_or_vector invoked with the RAW inputs (T-119-followup-default-callable). For a callable default, the output shape/dtype is determined at trace time by the callable's return value; for a constant default, it is determined statically from the value.	required

Input ports

(0..n_inputs-1) Boolean-castable scalars. Non-boolean inputs are coerced to bool before pattern matching (any non-zero is True).

Output ports

(0) The output of the first matching row, or default_output.

Notes

Earlier rows take precedence: if multiple patterns would match the same input combination, the one listed first in rows wins. The static-completeness/ambiguity checker is deferred (see T-119-followup-completeness-checker); JSON serialization of the rows table is deferred (see T-119-followup-serialization).

Source code in jaxonomy/library/logic.py

class TruthTable(LeafSystem):
    """Evaluate a fixed truth table over boolean-castable inputs.

    Given a list of ``(input_pattern, output)`` rows, this block compares
    its inputs against each pattern and emits the output of the first
    matching row (or ``default_output`` if none match). Patterns are tuples
    of ``bool`` values or the string ``"X"`` as a wildcard.

    Example — a 2-input AND gate:

    .. code-block:: python

        tt = TruthTable(
            rows=[
                ((True, True), 1.0),
                ((True, False), 0.0),
                ((False, True), 0.0),
                ((False, False), 0.0),
            ],
            n_inputs=2,
            default_output=0.0,
        )

    Wildcard example — ignore the first input:

    .. code-block:: python

        tt = TruthTable(
            rows=[(("X", True), 1.0), (("X", False), 0.0)],
            n_inputs=2,
            default_output=0.0,
        )

    Callable output example — row output depends on raw input values
    (T-119-followup-numeric-output):

    .. code-block:: python

        tt = TruthTable(
            rows=[
                ((True, True), lambda a, b: a + b),
                ((True, False), lambda a, b: a - b),
                ((False, "X"), 0.0),  # constant fallback row
            ],
            n_inputs=2,
            default_output=0.0,
        )

    Parameters:
        rows: list of ``(pattern, output)`` tuples. ``pattern`` is a
            length-``n_inputs`` tuple whose entries are ``bool`` (matched
            literally) or the string ``"X"`` (wildcard, matches anything).
            ``output`` may be a scalar/array (constant for that row) or
            a callable ``f(*inputs) -> scalar_or_vector`` invoked with
            the RAW inputs (every row callable is evaluated at every
            step under JAX's branchless ``where`` semantics; the result
            is *selected* only when the row matches). All row outputs
            (and the callable return values) must broadcast against
            ``default_output``.
        n_inputs: number of input ports.
        default_output: value emitted when no row matches. May be a
            scalar/array (constant fallback) or a callable
            ``f(*inputs) -> scalar_or_vector`` invoked with the RAW
            inputs (T-119-followup-default-callable). For a callable
            default, the output shape/dtype is determined at trace
            time by the callable's return value; for a constant
            default, it is determined statically from the value.

    Input ports:
        (0..n_inputs-1) Boolean-castable scalars. Non-boolean inputs are
        coerced to bool before pattern matching (any non-zero is True).

    Output ports:
        (0) The output of the first matching row, or ``default_output``.

    Notes:
        Earlier rows take precedence: if multiple patterns would match
        the same input combination, the one listed first in ``rows`` wins.
        The static-completeness/ambiguity checker is deferred
        (see ``T-119-followup-completeness-checker``); JSON serialization
        of the rows table is deferred (see
        ``T-119-followup-serialization``).
    """

    def __init__(self, rows, n_inputs, default_output, input_names=None, **kwargs):
        super().__init__(**kwargs)

        n = int(n_inputs)
        if n < 1:
            raise BlockParameterError(
                message=(
                    f"TruthTable block '{self.name}' requires n_inputs >= 1; "
                    f"got {n_inputs}."
                ),
                system=self,
                parameter_name="n_inputs",
            )

        # T-119-followup-truth-table-named-ports — labels forwarded from
        # ``TruthTableBuilder(input_names=...)`` (or supplied directly) become
        # the names of the declared input ports. Without this, the labels
        # survive only as ``.row(...)`` keyword targets and the block's
        # ``input_ports`` show up as anonymous ``in_0`` / ``in_1`` slots in
        # ``print_schedule`` and model JSON.
        if input_names is None:
            resolved_input_names = None
        else:
            resolved_input_names = tuple(input_names)
            if len(resolved_input_names) != n:
                raise BlockParameterError(
                    message=(
                        f"TruthTable block '{self.name}' input_names length "
                        f"({len(resolved_input_names)}) does not match "
                        f"n_inputs={n}."
                    ),
                    system=self,
                    parameter_name="input_names",
                )

        # Validate row shapes up front so misconfiguration fails at __init__,
        # not deep inside a JAX trace.
        # Each entry is ``(pattern, output_or_callable, is_callable)`` —
        # constant outputs are pre-coerced via ``npa.asarray``; callables
        # are stored as-is and invoked inside ``_compute_output``.
        # (T-119-followup-numeric-output)
        validated_rows: list[tuple[tuple, object, bool]] = []
        for row_idx, row in enumerate(rows):
            if not (isinstance(row, tuple) and len(row) == 2):
                raise BlockParameterError(
                    message=(
                        f"TruthTable block '{self.name}' row {row_idx} must "
                        f"be a (pattern, output) tuple; got {row!r}."
                    ),
                    system=self,
                    parameter_name="rows",
                )
            pattern, output = row
            if not (isinstance(pattern, tuple) and len(pattern) == n):
                raise BlockParameterError(
                    message=(
                        f"TruthTable block '{self.name}' row {row_idx} pattern "
                        f"must be a length-{n} tuple; got {pattern!r}."
                    ),
                    system=self,
                    parameter_name="rows",
                )
            for p in pattern:
                if p == "X":
                    continue
                if not isinstance(p, (bool, np.bool_)):
                    raise BlockParameterError(
                        message=(
                            f"TruthTable block '{self.name}' row {row_idx} "
                            f"pattern entries must be bool or 'X'; got {p!r}."
                        ),
                        system=self,
                        parameter_name="rows",
                    )
            if callable(output):
                # Defer evaluation to runtime; stored as-is so the
                # closure can call ``output(*inputs)`` with raw values.
                validated_rows.append((pattern, output, True))
            else:
                validated_rows.append((pattern, npa.asarray(output), False))

        # Stored as plain Python attributes — see module-header comment.
        # ``_default_output_is_callable`` mirrors the per-row
        # ``is_callable`` flag and selects the runtime branch in
        # ``_compute_output``. (T-119-followup-default-callable)
        self._rows = validated_rows
        self._n_inputs = n
        self._default_output_is_callable = callable(default_output)
        if self._default_output_is_callable:
            # Defer evaluation to runtime — the callable is invoked with
            # the raw inputs inside ``_compute_output``. We keep the
            # callable as-is so ``to_dict`` / ``to_csv`` can detect and
            # reject it.
            self._default_output = default_output
        else:
            self._default_output = npa.asarray(default_output)

        for i in range(n):
            if resolved_input_names is None:
                self.declare_input_port()
            else:
                self.declare_input_port(name=resolved_input_names[i])

        # Capture by closure to keep the compute function pure for JAX trace.
        rows_local = validated_rows
        default_local = self._default_output
        default_is_callable = self._default_output_is_callable

        # T-119-followup-truth-table-true-vectorise — when every row's
        # output is constant AND the default is constant (the common case
        # for combinational logic tables), pre-stack the outputs and the
        # pattern matrix so ``_compute_output`` compiles to one
        # ``argmax`` + ``take`` selection instead of N sequential
        # ``where`` operations. Earlier rows win on conflict, matching the
        # row-by-row loop semantics.
        all_constant = (
            not default_is_callable
            and all(not is_callable for _, _, is_callable in rows_local)
        )
        if all_constant and rows_local:
            # ``pattern_codes[r, i]`` is -1 for "don't care" (X), 0 for
            # False, 1 for True. Used at runtime to build a row-match
            # mask without re-walking the Python pattern tuples.
            pattern_codes = np.full((len(rows_local), n), -1, dtype=np.int8)
            for r, (pattern, _, _) in enumerate(rows_local):
                for i, p in enumerate(pattern):
                    if p == "X":
                        continue
                    pattern_codes[r, i] = int(bool(p))
            pattern_codes_local = pattern_codes
            outputs_stack_local = npa.stack(
                [output for _, output, _ in rows_local], axis=0
            )

            def _compute_output(_time, _state, *inputs, **_params):
                bool_inputs = npa.stack(
                    [npa.asarray(x).astype(npa.int8) for x in inputs], axis=0
                )
                # ``per_input_ok[r, i] = (pattern_codes[r, i] == -1) | (pattern_codes[r, i] == bool_inputs[i])``
                pc = npa.asarray(pattern_codes_local)
                per_input_ok = (pc == -1) | (pc == bool_inputs)
                match_vec = npa.all(per_input_ok, axis=1)
                any_match = npa.any(match_vec)
                # ``argmax`` returns the FIRST True, matching the
                # row-precedence semantics.
                first_idx = npa.argmax(match_vec)
                selected = outputs_stack_local[first_idx]
                return npa.where(any_match, selected, default_local)
        else:
            def _compute_output(_time, _state, *inputs, **_params):
                bool_inputs = tuple(npa.asarray(x).astype(bool) for x in inputs)
                if default_is_callable:
                    # Same semantics as a callable row output: evaluated
                    # unconditionally with raw inputs, selected by the
                    # ``where`` chain when no row matches.
                    # (T-119-followup-default-callable)
                    result = npa.asarray(default_local(*inputs))
                else:
                    result = default_local
                # Evaluate rows in reverse so that earlier rows take precedence
                # (each ``where`` overwrites the running result on a match).
                for pattern, output, is_callable in reversed(rows_local):
                    match = npa.array(True)
                    for i, p in enumerate(pattern):
                        if p == "X":
                            continue
                        match = match & (bool_inputs[i] == bool(p))
                    if is_callable:
                        # Pass RAW inputs (not bool-cast) so callable outputs
                        # see the actual numerical values. The callable is
                        # evaluated unconditionally — ``npa.where`` selects
                        # the active row, but both branches are traced by
                        # JAX. (T-119-followup-numeric-output)
                        row_value = npa.asarray(output(*inputs))
                    else:
                        row_value = output
                    result = npa.where(match, row_value, result)
                return result

        # ``default_value`` is only known statically when the default is
        # a constant array — for a callable default the shape/dtype is
        # discovered at trace time, so pass ``None`` and let the port
        # framework infer it from the first call.
        port_default = None if self._default_output_is_callable else self._default_output
        self.declare_output_port(
            _compute_output,
            default_value=port_default,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    @classmethod
    def builder(cls, n_inputs, default_output, input_names=None, **block_kwargs):
        """Construct a fluent builder for this truth table.

        See :class:`TruthTableBuilder` for usage. Equivalent to
        ``TruthTableBuilder(n_inputs, default_output, input_names, **block_kwargs)``.
        """
        return TruthTableBuilder(
            n_inputs=n_inputs,
            default_output=default_output,
            input_names=input_names,
            **block_kwargs,
        )

    # -----------------------------------------------------------------
    # T-119-followup-completeness-checker — static analysis on the truth table.
    #
    # The conventional TruthTable static analysis performs two construction-time checks:
    #   1. Completeness — every one of the 2^N input combinations is
    #      matched by at least one row's pattern (treating ``"X"`` as a
    #      wildcard). If not, those combinations would silently fall
    #      through to ``default_output``, which is a frequent source of
    #      logic bugs.
    #   2. Disjointness — no two rows match the same input combination.
    #      Jaxonomy's runtime semantics resolve overlaps by earlier-row-
    #      wins (see ``_compute_output``), so overlap is *not* a runtime
    #      error here, but flagging it surfaces unintended row shadowing.
    #
    # Default-off: ``validate(strict_completeness=False, strict_disjointness=
    # False)`` returns the report dict and never raises. Pass either flag
    # as ``True`` to escalate to ``BlockParameterError`` on the relevant
    # finding. Enumeration is pure-Python over ``itertools.product`` —
    # this runs at construction time, never inside a JAX trace.
    # -----------------------------------------------------------------

    def validate(self, strict_completeness=False, strict_disjointness=False):
        """Static analysis of the truth-table rows.

        Enumerates all ``2**n_inputs`` boolean input vectors and checks:

        * **Completeness** — each vector is matched by at least one row's
          pattern (with ``"X"`` as wildcard). Vectors that no row matches
          are reported as ``missing_patterns``; without coverage they
          silently hit ``default_output`` at runtime.
        * **Disjointness** — no two rows match the same input vector.
          Overlaps are reported as ``(earlier_idx, later_idx)`` pairs.
          Jaxonomy's runtime resolves overlaps by earlier-row-wins, so
          this is informational unless ``strict_disjointness=True``.

        Args:
            strict_completeness: if True, raise :class:`BlockParameterError`
                when any input combination is uncovered. Default False.
            strict_disjointness: if True, raise :class:`BlockParameterError`
                when any two rows match the same input combination.
                Default False.

        Returns:
            dict with keys:

            * ``covered_combinations`` (int) — number of distinct input
              vectors matched by at least one row.
            * ``total_combinations`` (int) — ``2 ** n_inputs``.
            * ``missing_patterns`` (list[tuple[bool, ...]]) — input
              vectors not matched by any row.
            * ``overlapping_pairs`` (list[tuple[int, int]]) — sorted
              ``(i, j)`` row-index pairs (``i < j``) where row ``i`` and
              row ``j`` both match at least one common input vector.

        Notes:
            For ``n_inputs > 10`` (i.e. > 1024 enumerated combinations)
            the check emits a :class:`UserWarning` since cost grows as
            ``2 ** n_inputs * len(rows)``.
        """
        import itertools

        n = self._n_inputs
        rows = self._rows

        if n > 10:
            warnings.warn(
                f"TruthTable.validate(): enumerating 2**{n} = {2 ** n} "
                f"input combinations across {len(rows)} rows may be slow. "
                f"Consider whether the static check is worth the cost for "
                f"this many inputs.",
                UserWarning,
                stacklevel=2,
            )

        # Pre-extract patterns once; we iterate them per combination.
        # Rows are 3-tuples ``(pattern, output, is_callable)`` after
        # T-119-followup-numeric-output; only the pattern matters here.
        patterns = [row[0] for row in rows]

        def _row_matches(pattern, vec):
            for p, v in zip(pattern, vec):
                if p == "X":
                    continue
                if bool(p) != bool(v):
                    return False
            return True

        total = 1 << n  # 2 ** n
        covered = 0
        missing: list[tuple[bool, ...]] = []
        # Track row-pairs that overlap on at least one vector. Use a set
        # to dedupe across enumerated vectors, then sort for stability.
        overlapping: set[tuple[int, int]] = set()

        for vec in itertools.product((False, True), repeat=n):
            matching_indices = [
                idx for idx, pat in enumerate(patterns)
                if _row_matches(pat, vec)
            ]
            if matching_indices:
                covered += 1
                if len(matching_indices) > 1:
                    for a in range(len(matching_indices)):
                        for b in range(a + 1, len(matching_indices)):
                            overlapping.add(
                                (matching_indices[a], matching_indices[b])
                            )
            else:
                missing.append(vec)

        overlapping_pairs = sorted(overlapping)

        report = {
            "covered_combinations": covered,
            "total_combinations": total,
            "missing_patterns": missing,
            "overlapping_pairs": overlapping_pairs,
        }

        if strict_completeness and missing:
            raise BlockParameterError(
                message=(
                    f"TruthTable block '{self.name}' is incomplete: "
                    f"{len(missing)} of {total} input combination(s) "
                    f"have no matching row and would fall through to "
                    f"default_output. Missing: {missing!r}."
                ),
                system=self,
                parameter_name="rows",
            )
        if strict_disjointness and overlapping_pairs:
            raise BlockParameterError(
                message=(
                    f"TruthTable block '{self.name}' has overlapping rows: "
                    f"row pair(s) {overlapping_pairs!r} match a common "
                    f"input combination. Earlier-row-wins resolves this at "
                    f"runtime, but the overlap is likely unintentional."
                ),
                system=self,
                parameter_name="rows",
            )

        return report

    # -----------------------------------------------------------------
    # T-119-followup-serialization — explicit dict/JSON round-trip.
    #
    # The runtime ``rows`` form — ``[(tuple of bool|"X", scalar|ndarray)]``
    # — is not directly JSON-friendly: tuples become lists, ``"X"``
    # mixes with bools, and ndarrays must be flattened with their
    # shape preserved. The pair below normalizes that shape into a
    # dict-of-primitives that survives ``json.dumps`` / ``json.loads``
    # untouched, then rebuilds the runtime form on the way back in
    # via the existing ``TruthTable(rows=..., n_inputs=..., ...)``
    # constructor (so the validation path is reused, not duplicated).
    #
    # Pattern encoding: bools become ``"1"``/``"0"`` and the wildcard
    # stays as ``"X"`` so the per-row pattern is a length-``n_inputs``
    # string. This keeps JSON output compact and human-readable while
    # avoiding any ambiguity between ``False`` and ``"X"``.
    #
    # Output encoding: scalars survive as Python ``float``; arrays are
    # encoded as ``{"shape": [...], "data": [...]}`` (flattened to a
    # plain list) so dtype is reconstructed via ``np.asarray`` —
    # consistent with the constructor's existing ``npa.asarray(output)``
    # contract and with T-005 default-float64 (no explicit cast).
    #
    # The ``@parameters(static=...)`` route was rejected because
    # ``declare_static_parameters`` coerces list-typed values to
    # ``np.array`` (see system_base.py), which would clobber the
    # ``[(pattern, output)]`` nested structure on first declaration.
    # Wiring this into the dashboard model serializer is a deeper
    # followup; the dict round-trip here lets callers persist /
    # reload TruthTable rows on their own.
    # -----------------------------------------------------------------

    @staticmethod
    def _encode_pattern(pattern):
        """Encode a runtime pattern tuple to a compact string.

        ``True`` -> ``"1"``, ``False`` -> ``"0"``, ``"X"`` stays ``"X"``.
        """
        chars = []
        for p in pattern:
            if p == "X":
                chars.append("X")
            else:
                chars.append("1" if bool(p) else "0")
        return "".join(chars)

    @staticmethod
    def _decode_pattern(pattern_str, n_inputs):
        """Inverse of :meth:`_encode_pattern`."""
        if not isinstance(pattern_str, str) or len(pattern_str) != n_inputs:
            raise ValueError(
                f"TruthTable pattern string must be length {n_inputs}; "
                f"got {pattern_str!r}."
            )
        decoded = []
        for ch in pattern_str:
            if ch == "X":
                decoded.append("X")
            elif ch == "1":
                decoded.append(True)
            elif ch == "0":
                decoded.append(False)
            else:
                raise ValueError(
                    f"TruthTable pattern character must be '0', '1' or "
                    f"'X'; got {ch!r}."
                )
        return tuple(decoded)

    @staticmethod
    def _encode_output(output):
        """Encode a scalar/ndarray output as a JSON-friendly value.

        Scalars (0-D arrays or plain floats) become ``float``.
        Higher-rank arrays become ``{"shape": [...], "data": [...]}``.
        """
        arr = np.asarray(output)
        if arr.shape == ():
            return float(arr)
        return {
            "shape": list(arr.shape),
            "data": arr.flatten().tolist(),
        }

    @staticmethod
    def _decode_output(encoded):
        """Inverse of :meth:`_encode_output`.

        Returns either a Python ``float`` (for scalars) or a numpy
        ``ndarray`` (for vector/matrix outputs). The constructor will
        ``npa.asarray`` either form, preserving T-005 default-float64.
        """
        if isinstance(encoded, dict):
            data = encoded["data"]
            shape = tuple(encoded["shape"])
            return np.asarray(data).reshape(shape)
        # Scalar fallback — accept int/float/bool/numpy scalar.
        return float(encoded)

    def to_dict(self):
        """Return a JSON-serializable dict describing this TruthTable.

        The dict round-trips through :meth:`from_dict` to a TruthTable
        with identical behaviour for every input combination. Pattern
        wildcards (``"X"``) and vector outputs are preserved.

        Returns:
            dict with keys:

            * ``n_inputs`` (int)
            * ``default_output`` (float | ``{"shape", "data"}``)
            * ``rows`` (list of ``{"pattern": str, "output": ...}``)
        """
        # Callable row outputs (T-119-followup-numeric-output) and
        # callable default_output (T-119-followup-default-callable)
        # cannot be round-tripped through JSON; surface a clear error
        # rather than silently dropping the arithmetic on serialize/load.
        if self._default_output_is_callable:
            raise ValueError(
                f"TruthTable.to_dict(): default_output is a callable "
                f"(got {self._default_output!r}); callable default "
                f"outputs are not JSON-serializable. Replace with a "
                f"constant scalar/array before persisting, or "
                f"reconstruct the TruthTable in code."
            )
        encoded_rows = []
        for row_idx, (pattern, output, is_callable) in enumerate(self._rows):
            if is_callable:
                raise ValueError(
                    f"TruthTable.to_dict(): row {row_idx} has a callable "
                    f"output (got {output!r}); callable row outputs are "
                    f"not JSON-serializable. Replace callable rows with "
                    f"constant outputs before persisting, or reconstruct "
                    f"the TruthTable in code."
                )
            encoded_rows.append({
                "pattern": self._encode_pattern(pattern),
                "output": self._encode_output(output),
            })
        return {
            "n_inputs": int(self._n_inputs),
            "default_output": self._encode_output(self._default_output),
            "rows": encoded_rows,
        }

    @classmethod
    def from_dict(cls, data, **block_kwargs):
        """Reconstruct a TruthTable from the dict produced by :meth:`to_dict`.

        Extra keyword arguments (``name=``, ``system_id=``, ...) are
        forwarded to the underlying ``TruthTable`` constructor, so a
        deserialized block can pick up a fresh name in its target diagram.

        Args:
            data: dict with the keys documented on :meth:`to_dict`.
            **block_kwargs: forwarded to ``TruthTable.__init__`` (e.g.
                ``name``, ``system_id``).

        Returns:
            A new :class:`TruthTable` whose ``rows`` and
            ``default_output`` match the serialized form.
        """
        n_inputs = int(data["n_inputs"])
        default_output = cls._decode_output(data["default_output"])
        rows = [
            (
                cls._decode_pattern(entry["pattern"], n_inputs),
                cls._decode_output(entry["output"]),
            )
            for entry in data["rows"]
        ]
        return cls(
            rows=rows,
            n_inputs=n_inputs,
            default_output=default_output,
            **block_kwargs,
        )

    # -----------------------------------------------------------------
    # T-119-followup-import-from-csv — load a TruthTable from a CSV file.
    #
    # CSV layout (header row required):
    #     in1,in2,in3,output
    #     T,T,T,1.0
    #     T,T,F,0.5
    #     T,F,X,0.25
    #     F,X,X,0.0
    #
    # Input cells accept ``T``/``True``/``1`` for True,
    # ``F``/``False``/``0`` for False, and ``X``/``-``/``*`` (or empty)
    # for the wildcard. Comparison is case-insensitive and surrounding
    # whitespace is stripped. The OUTPUT column (literally named
    # ``output``, case-insensitive) carries a single float per row;
    # if multiple ``output*`` columns are present (e.g. ``output_x``,
    # ``output_y``) they are stacked into a 1-D vector output per row,
    # in the column order they appear in the header. ``default_output``
    # defaults to a zero matching the row-output shape; pass it
    # explicitly via ``**block_kwargs`` to override.
    #
    # Pure-stdlib ``csv`` parsing — no pandas / numpy.loadtxt dep. The
    # existing constructor handles validation (length / dtype) so any
    # malformed row still surfaces a clear ``BlockParameterError`` once
    # the parsed rows reach ``TruthTable.__init__``; the parser itself
    # raises ``ValueError`` for header-level / cell-level mistakes.
    #
    # T-005 default-float64 is preserved: outputs are parsed with
    # ``float(...)`` and forwarded through the constructor's
    # ``npa.asarray`` (no explicit dtype cast).
    # -----------------------------------------------------------------

    # Accepted tokens for input cells (case-insensitive, whitespace-stripped).
    _CSV_TRUE_TOKENS = frozenset({"t", "true", "1"})
    _CSV_FALSE_TOKENS = frozenset({"f", "false", "0"})
    _CSV_WILDCARD_TOKENS = frozenset({"x", "-", "*", ""})

    @staticmethod
    def _parse_csv_input_cell(cell, row_idx, col_name):
        """Parse one input-column cell into ``True``, ``False`` or ``"X"``.

        Raises ``ValueError`` on unrecognised tokens, with row + column
        context so the caller can pinpoint the bad cell.
        """
        token = str(cell).strip().lower()
        if token in TruthTable._CSV_TRUE_TOKENS:
            return True
        if token in TruthTable._CSV_FALSE_TOKENS:
            return False
        if token in TruthTable._CSV_WILDCARD_TOKENS:
            return "X"
        raise ValueError(
            f"TruthTable.from_csv: row {row_idx} column {col_name!r} has "
            f"unrecognised input token {cell!r}; expected one of "
            f"T/True/1, F/False/0, X/-/* (case-insensitive)."
        )

    @classmethod
    def from_csv(cls, path, **block_kwargs):
        """Load a TruthTable from a CSV file.

        The CSV must have a header row whose last column(s) are named
        ``output`` (single scalar output) or any sequence of columns
        whose names start with ``output`` (e.g. ``output_x, output_y``)
        which are stacked into a 1-D vector output per row. All columns
        preceding the first ``output*`` column are treated as input
        columns, in order.

        Input cells accept ``T``/``True``/``1`` (True), ``F``/``False``/
        ``0`` (False), and ``X``/``-``/``*`` or an empty cell
        (wildcard). Matching is case-insensitive and whitespace is
        stripped. Output cells must parse as ``float``.

        Example CSV::

            in1,in2,in3,output
            T,T,T,1.0
            T,T,F,0.5
            T,F,X,0.25
            F,X,X,0.0

        Args:
            path: filesystem path (``str`` or ``os.PathLike``) to a
                readable CSV file.
            **block_kwargs: forwarded to ``TruthTable.__init__`` —
                typically ``name`` / ``system_id``. May also include
                ``default_output`` to override the zero-default that
                this loader picks (a scalar ``0.0`` for single-output
                CSVs, a zeros vector of the right shape for multi-output
                CSVs).

        Returns:
            A :class:`TruthTable` whose ``rows`` mirror the CSV.

        Raises:
            ValueError: if the file is empty, has no header, has no
                ``output`` column, or any row has the wrong number of
                cells / an unparseable input or output cell.
        """
        import csv

        with open(path, "r", newline="") as fh:
            reader = csv.reader(fh)
            try:
                header = next(reader)
            except StopIteration:
                raise ValueError(
                    f"TruthTable.from_csv: file {path!r} is empty; expected "
                    f"a header row followed by data rows."
                )
            # Strip surrounding whitespace from each header cell so callers
            # can author the CSV with comfortable spacing.
            header = [h.strip() for h in header]
            if not header:
                raise ValueError(
                    f"TruthTable.from_csv: file {path!r} has an empty "
                    f"header row."
                )

            # Identify the (contiguous) trailing output column(s). A column
            # is an "output column" if its name (lowercased) starts with
            # ``output``. All input columns must precede the first output
            # column; interleaving is rejected to keep the CSV layout
            # unambiguous.
            output_indices = [
                i for i, name in enumerate(header)
                if name.lower().startswith("output")
            ]
            if not output_indices:
                raise ValueError(
                    f"TruthTable.from_csv: file {path!r} has no 'output' "
                    f"column in header {header!r}; expected at least one "
                    f"column whose name starts with 'output'."
                )
            # Outputs must be contiguous and trail the inputs.
            first_out = output_indices[0]
            expected = list(range(first_out, first_out + len(output_indices)))
            if output_indices != expected:
                raise ValueError(
                    f"TruthTable.from_csv: file {path!r} has non-contiguous "
                    f"output columns at indices {output_indices!r}; all "
                    f"'output*' columns must be the trailing columns."
                )

            input_names = header[:first_out]
            output_names = header[first_out:]
            if not input_names:
                raise ValueError(
                    f"TruthTable.from_csv: file {path!r} has no input "
                    f"columns; at least one input column is required "
                    f"before the 'output' column(s)."
                )
            n_inputs = len(input_names)
            is_vector_output = len(output_names) > 1

            rows: list[tuple[tuple, object]] = []
            for row_idx, raw_row in enumerate(reader):
                # csv.reader yields empty lists for blank lines; skip them
                # so trailing newlines in the file don't blow up parsing.
                if not raw_row or all(c.strip() == "" for c in raw_row):
                    continue
                if len(raw_row) != len(header):
                    raise ValueError(
                        f"TruthTable.from_csv: row {row_idx} has "
                        f"{len(raw_row)} cell(s); expected {len(header)} "
                        f"(header: {header!r}, row: {raw_row!r})."
                    )
                pattern = tuple(
                    cls._parse_csv_input_cell(
                        raw_row[i], row_idx, input_names[i]
                    )
                    for i in range(n_inputs)
                )
                try:
                    output_cells = [
                        float(raw_row[i].strip()) for i in output_indices
                    ]
                except ValueError as exc:
                    raise ValueError(
                        f"TruthTable.from_csv: row {row_idx} output cell(s) "
                        f"{[raw_row[i] for i in output_indices]!r} did not "
                        f"parse as float: {exc}."
                    ) from None
                if is_vector_output:
                    output = np.asarray(output_cells)
                else:
                    output = output_cells[0]
                rows.append((pattern, output))

        if not rows:
            raise ValueError(
                f"TruthTable.from_csv: file {path!r} has a header but no "
                f"data rows."
            )

        # Pick a default_output matching the row-output shape unless the
        # caller passed one explicitly via block_kwargs.
        if "default_output" not in block_kwargs:
            if is_vector_output:
                block_kwargs["default_output"] = np.zeros(len(output_names))
            else:
                block_kwargs["default_output"] = 0.0

        return cls(rows=rows, n_inputs=n_inputs, **block_kwargs)

    # -----------------------------------------------------------------
    # T-119-followup-export-to-csv — write a TruthTable to a CSV file.
    #
    # Inverse of ``from_csv``: emits the same header layout
    # ``in1,in2,...,output`` (or ``output_0,output_1,...`` for vector
    # outputs), with input cells written as ``T`` / ``F`` / ``X`` and
    # output cells written as ``float(...)``. Round-tripping
    # ``TruthTable.from_csv(t.to_csv(path))`` reproduces the same rows,
    # ``n_inputs`` and ``default_output``.
    #
    # Callable row outputs (T-119-followup-numeric-output) cannot be
    # serialised — there is no portable representation of a Python
    # closure. ``to_csv`` raises ``ValueError`` in that case.
    #
    # Pure-stdlib ``csv`` writer; T-005 default-float64 is preserved by
    # writing values via ``float(...)``.
    # -----------------------------------------------------------------

    def to_csv(self, path, **csv_kwargs):
        """Write this TruthTable to a CSV file (inverse of ``from_csv``).

        Emits a header row of ``in1,in2,...,output`` (single-output) or
        ``in1,...,output_0,output_1,...`` (vector output), followed by
        one data row per ``rows`` entry. Input cells are written as
        ``T`` / ``F`` / ``X``; output cells are written as
        ``float(...)``.

        Args:
            path: filesystem path (``str`` or ``os.PathLike``) for the
                CSV file to (over)write.
            **csv_kwargs: forwarded to ``csv.writer`` (e.g. ``delimiter``,
                ``quoting``).

        Returns:
            ``path`` (so the caller can chain
            ``TruthTable.from_csv(t.to_csv(p))``).

        Raises:
            ValueError: if any row's output is a callable
                (T-119-followup-numeric-output) or ``default_output``
                is a callable (T-119-followup-default-callable) —
                callables are not representable in CSV.
        """
        import csv

        # Reject callable outputs up front so the file is never created
        # in a half-written state. (Both callable rows and a callable
        # default_output — T-119-followup-default-callable.)
        if self._default_output_is_callable:
            raise ValueError(
                f"TruthTable.to_csv: default_output is a callable; "
                f"callable default outputs cannot be serialised to CSV. "
                f"Replace with a constant scalar/array, or serialise "
                f"via a different format."
            )
        for row_idx, (_pattern, _output, is_callable) in enumerate(self._rows):
            if is_callable:
                raise ValueError(
                    f"TruthTable.to_csv: row {row_idx} has a callable "
                    f"output; callable row outputs cannot be serialised "
                    f"to CSV. Replace with a constant scalar/array, or "
                    f"serialise via a different format."
                )

        # Determine output column layout from the default_output shape
        # (which matches all row outputs by construction — the constructor
        # broadcasts row values against default_output). A 0-d / scalar
        # default produces a single ``output`` column; a 1-D vector
        # produces ``output_0, output_1, ...`` columns. This matches the
        # ``from_csv`` round-trip: a single ``output`` column yields a
        # scalar default, while ``output_*`` columns yield a vector default.
        default = np.asarray(self._default_output)
        if default.ndim == 0:
            output_names = ["output"]
            is_vector_output = False
        elif default.ndim == 1:
            output_names = [f"output_{i}" for i in range(default.shape[0])]
            is_vector_output = True
        else:
            raise ValueError(
                f"TruthTable.to_csv: default_output has shape "
                f"{default.shape!r}; only scalar (0-d) and 1-D vector "
                f"outputs are supported for CSV export."
            )

        input_names = [f"in{i + 1}" for i in range(self._n_inputs)]
        header = input_names + output_names

        with open(path, "w", newline="") as fh:
            writer = csv.writer(fh, **csv_kwargs)
            writer.writerow(header)
            for pattern, output, _is_callable in self._rows:
                pattern_cells = []
                for p in pattern:
                    if p == "X":
                        pattern_cells.append("X")
                    elif bool(p):
                        pattern_cells.append("T")
                    else:
                        pattern_cells.append("F")
                out_arr = np.asarray(output)
                if is_vector_output:
                    # Broadcast scalars (rare — constructor stores
                    # ``npa.asarray(output)``) up to the vector width.
                    out_vec = np.broadcast_to(out_arr, (len(output_names),))
                    output_cells = [float(v) for v in out_vec]
                else:
                    # Scalar output column — ``out_arr`` should be 0-d.
                    output_cells = [float(out_arr)]
                writer.writerow(pattern_cells + output_cells)

        return path

builder(n_inputs, default_output, input_names=None, **block_kwargs) classmethod

Construct a fluent builder for this truth table.

See :class:TruthTableBuilder for usage. Equivalent to TruthTableBuilder(n_inputs, default_output, input_names, **block_kwargs).

Source code in jaxonomy/library/logic.py

@classmethod
def builder(cls, n_inputs, default_output, input_names=None, **block_kwargs):
    """Construct a fluent builder for this truth table.

    See :class:`TruthTableBuilder` for usage. Equivalent to
    ``TruthTableBuilder(n_inputs, default_output, input_names, **block_kwargs)``.
    """
    return TruthTableBuilder(
        n_inputs=n_inputs,
        default_output=default_output,
        input_names=input_names,
        **block_kwargs,
    )

from_csv(path, **block_kwargs) classmethod

Load a TruthTable from a CSV file.

The CSV must have a header row whose last column(s) are named output (single scalar output) or any sequence of columns whose names start with output (e.g. output_x, output_y) which are stacked into a 1-D vector output per row. All columns preceding the first output* column are treated as input columns, in order.

Input cells accept T/True/1 (True), F/False/ 0 (False), and X/-/* or an empty cell (wildcard). Matching is case-insensitive and whitespace is stripped. Output cells must parse as float.

Example CSV::

in1,in2,in3,output
T,T,T,1.0
T,T,F,0.5
T,F,X,0.25
F,X,X,0.0

Parameters:

Name	Type	Description	Default
path		filesystem path (str or os.PathLike) to a readable CSV file.	required
**block_kwargs		forwarded to TruthTable.__init__ — typically name / system_id. May also include default_output to override the zero-default that this loader picks (a scalar 0.0 for single-output CSVs, a zeros vector of the right shape for multi-output CSVs).	{}

Returns:

Name	Type	Description
A		class:TruthTable whose rows mirror the CSV.

Raises:

Type	Description
ValueError	if the file is empty, has no header, has no output column, or any row has the wrong number of cells / an unparseable input or output cell.

Source code in jaxonomy/library/logic.py

@classmethod
def from_csv(cls, path, **block_kwargs):
    """Load a TruthTable from a CSV file.

    The CSV must have a header row whose last column(s) are named
    ``output`` (single scalar output) or any sequence of columns
    whose names start with ``output`` (e.g. ``output_x, output_y``)
    which are stacked into a 1-D vector output per row. All columns
    preceding the first ``output*`` column are treated as input
    columns, in order.

    Input cells accept ``T``/``True``/``1`` (True), ``F``/``False``/
    ``0`` (False), and ``X``/``-``/``*`` or an empty cell
    (wildcard). Matching is case-insensitive and whitespace is
    stripped. Output cells must parse as ``float``.

    Example CSV::

        in1,in2,in3,output
        T,T,T,1.0
        T,T,F,0.5
        T,F,X,0.25
        F,X,X,0.0

    Args:
        path: filesystem path (``str`` or ``os.PathLike``) to a
            readable CSV file.
        **block_kwargs: forwarded to ``TruthTable.__init__`` —
            typically ``name`` / ``system_id``. May also include
            ``default_output`` to override the zero-default that
            this loader picks (a scalar ``0.0`` for single-output
            CSVs, a zeros vector of the right shape for multi-output
            CSVs).

    Returns:
        A :class:`TruthTable` whose ``rows`` mirror the CSV.

    Raises:
        ValueError: if the file is empty, has no header, has no
            ``output`` column, or any row has the wrong number of
            cells / an unparseable input or output cell.
    """
    import csv

    with open(path, "r", newline="") as fh:
        reader = csv.reader(fh)
        try:
            header = next(reader)
        except StopIteration:
            raise ValueError(
                f"TruthTable.from_csv: file {path!r} is empty; expected "
                f"a header row followed by data rows."
            )
        # Strip surrounding whitespace from each header cell so callers
        # can author the CSV with comfortable spacing.
        header = [h.strip() for h in header]
        if not header:
            raise ValueError(
                f"TruthTable.from_csv: file {path!r} has an empty "
                f"header row."
            )

        # Identify the (contiguous) trailing output column(s). A column
        # is an "output column" if its name (lowercased) starts with
        # ``output``. All input columns must precede the first output
        # column; interleaving is rejected to keep the CSV layout
        # unambiguous.
        output_indices = [
            i for i, name in enumerate(header)
            if name.lower().startswith("output")
        ]
        if not output_indices:
            raise ValueError(
                f"TruthTable.from_csv: file {path!r} has no 'output' "
                f"column in header {header!r}; expected at least one "
                f"column whose name starts with 'output'."
            )
        # Outputs must be contiguous and trail the inputs.
        first_out = output_indices[0]
        expected = list(range(first_out, first_out + len(output_indices)))
        if output_indices != expected:
            raise ValueError(
                f"TruthTable.from_csv: file {path!r} has non-contiguous "
                f"output columns at indices {output_indices!r}; all "
                f"'output*' columns must be the trailing columns."
            )

        input_names = header[:first_out]
        output_names = header[first_out:]
        if not input_names:
            raise ValueError(
                f"TruthTable.from_csv: file {path!r} has no input "
                f"columns; at least one input column is required "
                f"before the 'output' column(s)."
            )
        n_inputs = len(input_names)
        is_vector_output = len(output_names) > 1

        rows: list[tuple[tuple, object]] = []
        for row_idx, raw_row in enumerate(reader):
            # csv.reader yields empty lists for blank lines; skip them
            # so trailing newlines in the file don't blow up parsing.
            if not raw_row or all(c.strip() == "" for c in raw_row):
                continue
            if len(raw_row) != len(header):
                raise ValueError(
                    f"TruthTable.from_csv: row {row_idx} has "
                    f"{len(raw_row)} cell(s); expected {len(header)} "
                    f"(header: {header!r}, row: {raw_row!r})."
                )
            pattern = tuple(
                cls._parse_csv_input_cell(
                    raw_row[i], row_idx, input_names[i]
                )
                for i in range(n_inputs)
            )
            try:
                output_cells = [
                    float(raw_row[i].strip()) for i in output_indices
                ]
            except ValueError as exc:
                raise ValueError(
                    f"TruthTable.from_csv: row {row_idx} output cell(s) "
                    f"{[raw_row[i] for i in output_indices]!r} did not "
                    f"parse as float: {exc}."
                ) from None
            if is_vector_output:
                output = np.asarray(output_cells)
            else:
                output = output_cells[0]
            rows.append((pattern, output))

    if not rows:
        raise ValueError(
            f"TruthTable.from_csv: file {path!r} has a header but no "
            f"data rows."
        )

    # Pick a default_output matching the row-output shape unless the
    # caller passed one explicitly via block_kwargs.
    if "default_output" not in block_kwargs:
        if is_vector_output:
            block_kwargs["default_output"] = np.zeros(len(output_names))
        else:
            block_kwargs["default_output"] = 0.0

    return cls(rows=rows, n_inputs=n_inputs, **block_kwargs)

from_dict(data, **block_kwargs) classmethod

Reconstruct a TruthTable from the dict produced by :meth:to_dict.

Extra keyword arguments (name=, system_id=, ...) are forwarded to the underlying TruthTable constructor, so a deserialized block can pick up a fresh name in its target diagram.

Parameters:

Name	Type	Description	Default
data		dict with the keys documented on :meth:to_dict.	required
**block_kwargs		forwarded to TruthTable.__init__ (e.g. name, system_id).	{}

Returns:

Type	Description
	A new :class:TruthTable whose rows and
	default_output match the serialized form.

Source code in jaxonomy/library/logic.py

@classmethod
def from_dict(cls, data, **block_kwargs):
    """Reconstruct a TruthTable from the dict produced by :meth:`to_dict`.

    Extra keyword arguments (``name=``, ``system_id=``, ...) are
    forwarded to the underlying ``TruthTable`` constructor, so a
    deserialized block can pick up a fresh name in its target diagram.

    Args:
        data: dict with the keys documented on :meth:`to_dict`.
        **block_kwargs: forwarded to ``TruthTable.__init__`` (e.g.
            ``name``, ``system_id``).

    Returns:
        A new :class:`TruthTable` whose ``rows`` and
        ``default_output`` match the serialized form.
    """
    n_inputs = int(data["n_inputs"])
    default_output = cls._decode_output(data["default_output"])
    rows = [
        (
            cls._decode_pattern(entry["pattern"], n_inputs),
            cls._decode_output(entry["output"]),
        )
        for entry in data["rows"]
    ]
    return cls(
        rows=rows,
        n_inputs=n_inputs,
        default_output=default_output,
        **block_kwargs,
    )

to_csv(path, **csv_kwargs)

Write this TruthTable to a CSV file (inverse of from_csv).

Emits a header row of in1,in2,...,output (single-output) or in1,...,output_0,output_1,... (vector output), followed by one data row per rows entry. Input cells are written as T / F / X; output cells are written as float(...).

Parameters:

Name	Type	Description	Default
path		filesystem path (str or os.PathLike) for the CSV file to (over)write.	required
**csv_kwargs		forwarded to csv.writer (e.g. delimiter, quoting).	{}

Returns:

Type	Description
	path (so the caller can chain
	TruthTable.from_csv(t.to_csv(p))).

Raises:

Type	Description
ValueError	if any row's output is a callable (T-119-followup-numeric-output) or default_output is a callable (T-119-followup-default-callable) — callables are not representable in CSV.

Source code in jaxonomy/library/logic.py

def to_csv(self, path, **csv_kwargs):
    """Write this TruthTable to a CSV file (inverse of ``from_csv``).

    Emits a header row of ``in1,in2,...,output`` (single-output) or
    ``in1,...,output_0,output_1,...`` (vector output), followed by
    one data row per ``rows`` entry. Input cells are written as
    ``T`` / ``F`` / ``X``; output cells are written as
    ``float(...)``.

    Args:
        path: filesystem path (``str`` or ``os.PathLike``) for the
            CSV file to (over)write.
        **csv_kwargs: forwarded to ``csv.writer`` (e.g. ``delimiter``,
            ``quoting``).

    Returns:
        ``path`` (so the caller can chain
        ``TruthTable.from_csv(t.to_csv(p))``).

    Raises:
        ValueError: if any row's output is a callable
            (T-119-followup-numeric-output) or ``default_output``
            is a callable (T-119-followup-default-callable) —
            callables are not representable in CSV.
    """
    import csv

    # Reject callable outputs up front so the file is never created
    # in a half-written state. (Both callable rows and a callable
    # default_output — T-119-followup-default-callable.)
    if self._default_output_is_callable:
        raise ValueError(
            f"TruthTable.to_csv: default_output is a callable; "
            f"callable default outputs cannot be serialised to CSV. "
            f"Replace with a constant scalar/array, or serialise "
            f"via a different format."
        )
    for row_idx, (_pattern, _output, is_callable) in enumerate(self._rows):
        if is_callable:
            raise ValueError(
                f"TruthTable.to_csv: row {row_idx} has a callable "
                f"output; callable row outputs cannot be serialised "
                f"to CSV. Replace with a constant scalar/array, or "
                f"serialise via a different format."
            )

    # Determine output column layout from the default_output shape
    # (which matches all row outputs by construction — the constructor
    # broadcasts row values against default_output). A 0-d / scalar
    # default produces a single ``output`` column; a 1-D vector
    # produces ``output_0, output_1, ...`` columns. This matches the
    # ``from_csv`` round-trip: a single ``output`` column yields a
    # scalar default, while ``output_*`` columns yield a vector default.
    default = np.asarray(self._default_output)
    if default.ndim == 0:
        output_names = ["output"]
        is_vector_output = False
    elif default.ndim == 1:
        output_names = [f"output_{i}" for i in range(default.shape[0])]
        is_vector_output = True
    else:
        raise ValueError(
            f"TruthTable.to_csv: default_output has shape "
            f"{default.shape!r}; only scalar (0-d) and 1-D vector "
            f"outputs are supported for CSV export."
        )

    input_names = [f"in{i + 1}" for i in range(self._n_inputs)]
    header = input_names + output_names

    with open(path, "w", newline="") as fh:
        writer = csv.writer(fh, **csv_kwargs)
        writer.writerow(header)
        for pattern, output, _is_callable in self._rows:
            pattern_cells = []
            for p in pattern:
                if p == "X":
                    pattern_cells.append("X")
                elif bool(p):
                    pattern_cells.append("T")
                else:
                    pattern_cells.append("F")
            out_arr = np.asarray(output)
            if is_vector_output:
                # Broadcast scalars (rare — constructor stores
                # ``npa.asarray(output)``) up to the vector width.
                out_vec = np.broadcast_to(out_arr, (len(output_names),))
                output_cells = [float(v) for v in out_vec]
            else:
                # Scalar output column — ``out_arr`` should be 0-d.
                output_cells = [float(out_arr)]
            writer.writerow(pattern_cells + output_cells)

    return path

to_dict()

Return a JSON-serializable dict describing this TruthTable.

The dict round-trips through :meth:from_dict to a TruthTable with identical behaviour for every input combination. Pattern wildcards ("X") and vector outputs are preserved.

Returns:

Type	Description
	dict with keys:
	n_inputs (int)
	default_output (float | {"shape", "data"})
	rows (list of {"pattern": str, "output": ...})

Source code in jaxonomy/library/logic.py

def to_dict(self):
    """Return a JSON-serializable dict describing this TruthTable.

    The dict round-trips through :meth:`from_dict` to a TruthTable
    with identical behaviour for every input combination. Pattern
    wildcards (``"X"``) and vector outputs are preserved.

    Returns:
        dict with keys:

        * ``n_inputs`` (int)
        * ``default_output`` (float | ``{"shape", "data"}``)
        * ``rows`` (list of ``{"pattern": str, "output": ...}``)
    """
    # Callable row outputs (T-119-followup-numeric-output) and
    # callable default_output (T-119-followup-default-callable)
    # cannot be round-tripped through JSON; surface a clear error
    # rather than silently dropping the arithmetic on serialize/load.
    if self._default_output_is_callable:
        raise ValueError(
            f"TruthTable.to_dict(): default_output is a callable "
            f"(got {self._default_output!r}); callable default "
            f"outputs are not JSON-serializable. Replace with a "
            f"constant scalar/array before persisting, or "
            f"reconstruct the TruthTable in code."
        )
    encoded_rows = []
    for row_idx, (pattern, output, is_callable) in enumerate(self._rows):
        if is_callable:
            raise ValueError(
                f"TruthTable.to_dict(): row {row_idx} has a callable "
                f"output (got {output!r}); callable row outputs are "
                f"not JSON-serializable. Replace callable rows with "
                f"constant outputs before persisting, or reconstruct "
                f"the TruthTable in code."
            )
        encoded_rows.append({
            "pattern": self._encode_pattern(pattern),
            "output": self._encode_output(output),
        })
    return {
        "n_inputs": int(self._n_inputs),
        "default_output": self._encode_output(self._default_output),
        "rows": encoded_rows,
    }

validate(strict_completeness=False, strict_disjointness=False)

Static analysis of the truth-table rows.

Enumerates all 2**n_inputs boolean input vectors and checks:

• Completeness — each vector is matched by at least one row's pattern (with "X" as wildcard). Vectors that no row matches are reported as missing_patterns; without coverage they silently hit default_output at runtime.
• Disjointness — no two rows match the same input vector. Overlaps are reported as (earlier_idx, later_idx) pairs. Jaxonomy's runtime resolves overlaps by earlier-row-wins, so this is informational unless strict_disjointness=True.

Parameters:

Name	Type	Description	Default
strict_completeness		if True, raise :class:BlockParameterError when any input combination is uncovered. Default False.	False
strict_disjointness		if True, raise :class:BlockParameterError when any two rows match the same input combination. Default False.	False

Returns:

Type	Description
	dict with keys:
	covered_combinations (int) — number of distinct input vectors matched by at least one row.
	total_combinations (int) — 2 ** n_inputs.
	missing_patterns (list[tuple[bool, ...]]) — input vectors not matched by any row.
	overlapping_pairs (list[tuple[int, int]]) — sorted (i, j) row-index pairs (i < j) where row i and row j both match at least one common input vector.

Notes

For n_inputs > 10 (i.e. > 1024 enumerated combinations) the check emits a :class:UserWarning since cost grows as 2 ** n_inputs * len(rows).

Source code in jaxonomy/library/logic.py

def validate(self, strict_completeness=False, strict_disjointness=False):
    """Static analysis of the truth-table rows.

    Enumerates all ``2**n_inputs`` boolean input vectors and checks:

    * **Completeness** — each vector is matched by at least one row's
      pattern (with ``"X"`` as wildcard). Vectors that no row matches
      are reported as ``missing_patterns``; without coverage they
      silently hit ``default_output`` at runtime.
    * **Disjointness** — no two rows match the same input vector.
      Overlaps are reported as ``(earlier_idx, later_idx)`` pairs.
      Jaxonomy's runtime resolves overlaps by earlier-row-wins, so
      this is informational unless ``strict_disjointness=True``.

    Args:
        strict_completeness: if True, raise :class:`BlockParameterError`
            when any input combination is uncovered. Default False.
        strict_disjointness: if True, raise :class:`BlockParameterError`
            when any two rows match the same input combination.
            Default False.

    Returns:
        dict with keys:

        * ``covered_combinations`` (int) — number of distinct input
          vectors matched by at least one row.
        * ``total_combinations`` (int) — ``2 ** n_inputs``.
        * ``missing_patterns`` (list[tuple[bool, ...]]) — input
          vectors not matched by any row.
        * ``overlapping_pairs`` (list[tuple[int, int]]) — sorted
          ``(i, j)`` row-index pairs (``i < j``) where row ``i`` and
          row ``j`` both match at least one common input vector.

    Notes:
        For ``n_inputs > 10`` (i.e. > 1024 enumerated combinations)
        the check emits a :class:`UserWarning` since cost grows as
        ``2 ** n_inputs * len(rows)``.
    """
    import itertools

    n = self._n_inputs
    rows = self._rows

    if n > 10:
        warnings.warn(
            f"TruthTable.validate(): enumerating 2**{n} = {2 ** n} "
            f"input combinations across {len(rows)} rows may be slow. "
            f"Consider whether the static check is worth the cost for "
            f"this many inputs.",
            UserWarning,
            stacklevel=2,
        )

    # Pre-extract patterns once; we iterate them per combination.
    # Rows are 3-tuples ``(pattern, output, is_callable)`` after
    # T-119-followup-numeric-output; only the pattern matters here.
    patterns = [row[0] for row in rows]

    def _row_matches(pattern, vec):
        for p, v in zip(pattern, vec):
            if p == "X":
                continue
            if bool(p) != bool(v):
                return False
        return True

    total = 1 << n  # 2 ** n
    covered = 0
    missing: list[tuple[bool, ...]] = []
    # Track row-pairs that overlap on at least one vector. Use a set
    # to dedupe across enumerated vectors, then sort for stability.
    overlapping: set[tuple[int, int]] = set()

    for vec in itertools.product((False, True), repeat=n):
        matching_indices = [
            idx for idx, pat in enumerate(patterns)
            if _row_matches(pat, vec)
        ]
        if matching_indices:
            covered += 1
            if len(matching_indices) > 1:
                for a in range(len(matching_indices)):
                    for b in range(a + 1, len(matching_indices)):
                        overlapping.add(
                            (matching_indices[a], matching_indices[b])
                        )
        else:
            missing.append(vec)

    overlapping_pairs = sorted(overlapping)

    report = {
        "covered_combinations": covered,
        "total_combinations": total,
        "missing_patterns": missing,
        "overlapping_pairs": overlapping_pairs,
    }

    if strict_completeness and missing:
        raise BlockParameterError(
            message=(
                f"TruthTable block '{self.name}' is incomplete: "
                f"{len(missing)} of {total} input combination(s) "
                f"have no matching row and would fall through to "
                f"default_output. Missing: {missing!r}."
            ),
            system=self,
            parameter_name="rows",
        )
    if strict_disjointness and overlapping_pairs:
        raise BlockParameterError(
            message=(
                f"TruthTable block '{self.name}' has overlapping rows: "
                f"row pair(s) {overlapping_pairs!r} match a common "
                f"input combination. Earlier-row-wins resolves this at "
                f"runtime, but the overlap is likely unintentional."
            ),
            system=self,
            parameter_name="rows",
        )

    return report

TruthTableBuilder

Fluent builder for :class:TruthTable rows by named-input keywords.

Example — a 2-input AND gate:

.. code-block:: python

tt = (
    TruthTable.builder(n_inputs=2, default_output=0.0)
    .row(in1=True, in2=True, output=1.0)
    .row(in1=True, in2=False, output=0.0)
    .row(in1=False, in2="X", output=0.0)
    .build()
)

Inputs omitted from a .row(...) call default to the wildcard "X", so partial decision tables are concise. Custom input names may be supplied via the input_names= constructor argument; the default names are in1, in2, ..., inN.

Source code in jaxonomy/library/logic.py

class TruthTableBuilder:
    """Fluent builder for :class:`TruthTable` rows by named-input keywords.

    Example — a 2-input AND gate:

    .. code-block:: python

        tt = (
            TruthTable.builder(n_inputs=2, default_output=0.0)
            .row(in1=True, in2=True, output=1.0)
            .row(in1=True, in2=False, output=0.0)
            .row(in1=False, in2="X", output=0.0)
            .build()
        )

    Inputs omitted from a ``.row(...)`` call default to the wildcard
    ``"X"``, so partial decision tables are concise. Custom input names
    may be supplied via the ``input_names=`` constructor argument; the
    default names are ``in1, in2, ..., inN``.
    """

    def __init__(
        self,
        n_inputs,
        default_output,
        input_names=None,
        **block_kwargs,
    ):
        n = int(n_inputs)
        if n < 1:
            raise ValueError(
                f"TruthTableBuilder requires n_inputs >= 1; got {n_inputs}."
            )
        if input_names is None:
            self._input_names = tuple(f"in{i + 1}" for i in range(n))
        else:
            names = tuple(input_names)
            if len(names) != n:
                raise ValueError(
                    f"TruthTableBuilder input_names must have length n_inputs="
                    f"{n}; got {len(names)} ({names!r})."
                )
            if len(set(names)) != len(names):
                raise ValueError(
                    f"TruthTableBuilder input_names must be unique; got {names!r}."
                )
            self._input_names = names
        self._n_inputs = n
        self._default_output = default_output
        self._block_kwargs = block_kwargs
        self._rows: list[tuple[tuple, object]] = []

    def row(self, output, **input_assignments):
        """Append a row, named by input keyword.

        ``output`` is the value emitted when the row matches. Each
        keyword in ``input_assignments`` must be one of the configured
        input names; omitted inputs default to the wildcard ``"X"``.
        Returns ``self`` for fluent chaining.
        """
        unknown = set(input_assignments) - set(self._input_names)
        if unknown:
            raise ValueError(
                f"TruthTableBuilder.row(...) got unknown input name(s) "
                f"{sorted(unknown)!r}; expected one of {list(self._input_names)!r}."
            )
        pattern = tuple(
            input_assignments.get(name, "X") for name in self._input_names
        )
        self._rows.append((pattern, output))
        return self

    def build(self):
        """Materialize the accumulated rows into a :class:`TruthTable`."""
        # Forward ``input_names`` so the labels survive on the built block's
        # ``input_ports`` (visible in print_schedule / model JSON / error
        # messages) — T-119-followup-truth-table-named-ports. Skip the
        # default placeholder names (``in1``/``in2``/...) so a user who
        # never set ``input_names=`` keeps the existing anonymous-port
        # behaviour.
        default_input_names = tuple(f"in{i + 1}" for i in range(self._n_inputs))
        if self._input_names == default_input_names:
            input_names = None
        else:
            input_names = self._input_names
        return TruthTable(
            rows=list(self._rows),
            n_inputs=self._n_inputs,
            default_output=self._default_output,
            input_names=input_names,
            **self._block_kwargs,
        )

build()

Materialize the accumulated rows into a :class:TruthTable.

Source code in jaxonomy/library/logic.py

def build(self):
    """Materialize the accumulated rows into a :class:`TruthTable`."""
    # Forward ``input_names`` so the labels survive on the built block's
    # ``input_ports`` (visible in print_schedule / model JSON / error
    # messages) — T-119-followup-truth-table-named-ports. Skip the
    # default placeholder names (``in1``/``in2``/...) so a user who
    # never set ``input_names=`` keeps the existing anonymous-port
    # behaviour.
    default_input_names = tuple(f"in{i + 1}" for i in range(self._n_inputs))
    if self._input_names == default_input_names:
        input_names = None
    else:
        input_names = self._input_names
    return TruthTable(
        rows=list(self._rows),
        n_inputs=self._n_inputs,
        default_output=self._default_output,
        input_names=input_names,
        **self._block_kwargs,
    )

row(output, **input_assignments)

Append a row, named by input keyword.

output is the value emitted when the row matches. Each keyword in input_assignments must be one of the configured input names; omitted inputs default to the wildcard "X". Returns self for fluent chaining.

Source code in jaxonomy/library/logic.py

def row(self, output, **input_assignments):
    """Append a row, named by input keyword.

    ``output`` is the value emitted when the row matches. Each
    keyword in ``input_assignments`` must be one of the configured
    input names; omitted inputs default to the wildcard ``"X"``.
    Returns ``self`` for fluent chaining.
    """
    unknown = set(input_assignments) - set(self._input_names)
    if unknown:
        raise ValueError(
            f"TruthTableBuilder.row(...) got unknown input name(s) "
            f"{sorted(unknown)!r}; expected one of {list(self._input_names)!r}."
        )
    pattern = tuple(
        input_assignments.get(name, "X") for name in self._input_names
    )
    self._rows.append((pattern, output))
    return self

UniformRandomNumber

Bases: LeafSystem

Discrete-time uniform random number generator.

Emits a fresh Uniform[low, high] sample every sample_time seconds, using jax.random.uniform with a key carried in the block's discrete state. Reproducible: same seed and same diagram → bit-identical sequence.

The sample is computed as low + (high - low) * u where u ~ Uniform[0, 1), so gradients of downstream losses flow cleanly through low and high via the reparameterization trick. The u draw is wrapped in lax.stop_gradient so JAX never tries to differentiate the random sequence w.r.t. the key.

Input ports

None.

Output ports

(0) The most recent uniform sample.

Parameters:

Name	Type	Description	Default
sample_time	float	Period (s) at which a fresh sample is drawn.	required
low	float	Lower bound of the uniform interval (differentiable).	0.0
high	float	Upper bound of the uniform interval (differentiable).	1.0
seed	int	Integer seed for the PRNG key. If None, a 32-bit random seed is drawn from numpy.random.	None
shape		Output shape. Default () (scalar).	()

Notes

Per-vmap-batch independence: pass fold_in_batch_index=True (T-122-followup-vmap-fold-in) to derive a per-replica independent PRNG stream via jax.lax.axis_index("batch") inside simulate_batch(use_vmap=True) / simulate_distributed. Outside any vmap context the kwarg is a no-op (the unbound-axis NameError is caught gracefully and the plain seed-derived key is used). The default False preserves bit-identical behaviour with T-122 phase 1.

Source code in jaxonomy/library/sources.py

class UniformRandomNumber(LeafSystem):
    """Discrete-time uniform random number generator.

    Emits a fresh Uniform[low, high] sample every ``sample_time``
    seconds, using ``jax.random.uniform`` with a key carried in the
    block's discrete state. Reproducible: same ``seed`` and same
    diagram → bit-identical sequence.

    The sample is computed as ``low + (high - low) * u`` where
    ``u ~ Uniform[0, 1)``, so gradients of downstream losses flow
    cleanly through ``low`` and ``high`` via the reparameterization
    trick. The ``u`` draw is wrapped in ``lax.stop_gradient`` so JAX
    never tries to differentiate the random sequence w.r.t. the key.

    Input ports:
        None.

    Output ports:
        (0) The most recent uniform sample.

    Parameters:
        sample_time: Period (s) at which a fresh sample is drawn.
        low: Lower bound of the uniform interval (differentiable).
        high: Upper bound of the uniform interval (differentiable).
        seed: Integer seed for the PRNG key. If ``None``, a 32-bit
            random seed is drawn from ``numpy.random``.
        shape: Output shape. Default ``()`` (scalar).

    Notes:
        Per-vmap-batch independence: pass ``fold_in_batch_index=True``
        (T-122-followup-vmap-fold-in) to derive a per-replica
        independent PRNG stream via ``jax.lax.axis_index("batch")``
        inside ``simulate_batch(use_vmap=True)`` / ``simulate_distributed``.
        Outside any vmap context the kwarg is a no-op (the unbound-axis
        ``NameError`` is caught gracefully and the plain seed-derived
        key is used).  The default ``False`` preserves bit-identical
        behaviour with T-122 phase 1.
    """

    @parameters(
        static=["seed", "shape", "fold_in_batch_index"],
        dynamic=["low", "high"],
    )
    def __init__(
        self,
        sample_time: float,
        low: float = 0.0,
        high: float = 1.0,
        seed: int = None,
        shape=(),
        fold_in_batch_index: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self._sample_time = float(sample_time)

        self.declare_output_port(
            self._output,
            period=sample_time,
            offset=0.0,
        )
        self.declare_periodic_update(
            self._update,
            period=sample_time,
            offset=0.0,
        )

    def initialize(
        self,
        low: float = 0.0,
        high: float = 1.0,
        seed: int = None,
        shape=(),
        fold_in_batch_index: bool = False,
    ):
        # Lazy JAX import — the framework supports a numpy-only backend,
        # but the stochastic sources require jax.random just like
        # ``RandomNumber`` and ``WhiteNoise`` already do (see
        # ``library/random.py`` module header).
        from jax import random as _jrandom
        from jax import lax as _jlax

        self._jrandom = _jrandom
        self._jlax = _jlax
        self._shape = tuple(int(s) for s in shape) if shape else ()
        self._fold_in_batch_index = bool(fold_in_batch_index)

        if seed is None:
            seed = int(np.random.randint(0, 2**31 - 1, dtype=np.int64))
        key = _jrandom.PRNGKey(int(seed))
        key, subkey = _jrandom.split(key)
        # Build initial sample with the same shape/dtype the update
        # function will produce, so the discrete-state pytree is
        # stable across periodic updates.
        u0 = _jrandom.uniform(subkey, self._shape)
        val0 = float(low) + (float(high) - float(low)) * u0
        default_state = _PRNGState(key=key, val=val0)
        self.declare_discrete_state(default_value=default_state, as_array=False)

    def _output(self, _time, state, *_inputs, **_parameters):
        return state.discrete_state.val

    def _update(self, _time, state, *_inputs, **parameters):
        key, subkey = self._jrandom.split(state.discrete_state.key)
        # T-122-followup-vmap-fold-in: when running under
        # vmap(axis_name="batch") and ``fold_in_batch_index=True``, fold
        # the per-replica batch index into the freshly-split subkey so
        # each replica draws an independent stream from the same seed.
        subkey = _maybe_fold_in_batch_axis(
            self._jrandom, subkey, self._fold_in_batch_index
        )
        # ``stop_gradient`` makes the non-differentiability of the
        # random draw explicit — gradients still flow through ``low``
        # / ``high`` via the reparameterization below.
        u = self._jlax.stop_gradient(
            self._jrandom.uniform(subkey, self._shape)
        )
        low = parameters["low"]
        high = parameters["high"]
        val = low + (high - low) * u
        return _PRNGState(key=key, val=val)

UnitDelay

Bases: LeafSystem

Hold and delay the input signal by one time step.

This block implements a "unit delay" with the following difference equation for internal state x, input signal u, and output signal y:

    x[k+1] = u[k]
    y[k] = x[k]

Or, in a hybrid context, the discrete update advances the internal state from the "pre" or "minus" value x⁻ to the "post" or "plus" value x⁺ at time tₖ = t0 + k * dt. According to the discrete update rules, this calculation happens using the input values computed during the update step (i.e. by computing upstream outputs before evaluating the inputs to this block). That is, the update rule can be written x⁺(tₖ) = f(tₖ, x⁻(tₖ), u(tₖ)). The values of u are not distinguished as "pre" or "post" because there is only one value at the update time. In the difference equation notation, x⁺(tₖ) ≡ x[k+1],x⁻(tₖ) ≡ x[k], and u(tₖ) ≡ u[k]. The hybrid update rule is then:

    x⁺(tₖ) = u(tₖ)
    y(t) = x⁻(tₖ),       between tₖ⁺ and (tₖ+dt)⁻

The output signal "seen" by all other blocks on the time interval (tₖ, tₖ+dt) is then the value of the input signal u(tₖ) at the previous update. Therefore, all downstream discrete-time blocks updating at the same time tₖ will still see the value of x⁻(tₖ), the value of the internal state prior to the update.

Input ports

(0) The input signal.

Output ports

(0) The input signal delayed by one time step

Parameters:

Name	Type	Description	Default
dt		The time step of the discrete update.	required
initial_state		The initial state of the block. Default is 0.0.	required

Note

For a multi-step / fixed transport latency, do not chain N UnitDelay blocks — use a single :class:TransportDelay (delay_seconds = N * dt), which buffers the history in one block and is differentiable through the signal. UnitDelay is the exact one-sample z⁻¹ primitive; :class:TransportDelay is the parameterized N-sample delay line.

Source code in jaxonomy/library/dynamics.py

class UnitDelay(LeafSystem):
    """Hold and delay the input signal by one time step.

    This block implements a "unit delay" with the following difference equation
    for internal state `x`, input signal `u`, and output signal `y`:
    ```
        x[k+1] = u[k]
        y[k] = x[k]
    ```
    Or, in a hybrid context, the discrete update advances the internal state from
    the "pre" or "minus" value x⁻ to the "post" or "plus" value x⁺ at time
    `tₖ = t0 + k * dt`.  According to the discrete update rules, this calculation
    happens using the input values computed during the update step (i.e. by computing
    upstream outputs before evaluating the inputs to this block). That is, the update
    rule can be written `x⁺(tₖ) = f(tₖ, x⁻(tₖ), u(tₖ))`.  The values of `u` are not
    distinguished as "pre" or "post" because there is only one value at the update
    time.  In the difference equation notation, x⁺(tₖ) ≡ x[k+1]`, `x⁻(tₖ) ≡ x[k],
    and u(tₖ) ≡ u[k].  The hybrid update rule is then:
    ```
        x⁺(tₖ) = u(tₖ)
        y(t) = x⁻(tₖ),       between tₖ⁺ and (tₖ+dt)⁻
    ```

    The output signal "seen" by all other blocks on the time interval (tₖ, tₖ+dt)
    is then the value of the input signal u(tₖ) at the previous update. Therefore, all
    downstream discrete-time blocks updating at the same time tₖ will still see the
    value of x⁻(tₖ), the value of the internal state prior to the update.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The input signal delayed by one time step

    Parameters:
        dt:
            The time step of the discrete update.
        initial_state:
            The initial state of the block.  Default is 0.0.

    Note:
        For a *multi-step* / fixed transport latency, do not chain N
        ``UnitDelay`` blocks — use a single :class:`TransportDelay`
        (``delay_seconds = N * dt``), which buffers the history in one block
        and is differentiable through the signal. ``UnitDelay`` is the exact
        one-sample ``z⁻¹`` primitive; :class:`TransportDelay` is the
        parameterized N-sample delay line.
    """

    @parameters(static=["dt"], dynamic=["initial_state"])
    def __init__(self, dt, initial_state, *args, dtype=None, **kwargs):
        # T-038a-followup-other-blocks: per-block dtype override; stored
        # outside the @parameters list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(*args, **kwargs)
        self.dt = dt
        self.declare_input_port()
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(self, initial_state, dt=None):
        if self._dtype is not None:
            initial_state = npa.asarray(initial_state).astype(self._dtype)
        self.configure_periodic_update(
            self._periodic_update_idx, self._update, period=self.dt, offset=self.dt
        )

        self.configure_output_port(
            self._output_port_idx,
            self._output,
            period=self.dt,
            offset=0.0,
            requires_inputs=False,
            prerequisites_of_calc=[DependencyTicket.xd],
            default_value=initial_state,
        )

    def reset_default_values(self, initial_state, dt=None):
        if self._dtype is not None:
            initial_state = npa.asarray(initial_state).astype(self._dtype)
        self.declare_discrete_state(default_value=initial_state)
        self.configure_output_port_default_value(self._output_port_idx, initial_state)

    def _update(self, _time, _state, u, **_params):
        # Every dt seconds, update the state to the current input value
        # T-038a-followup-other-blocks: when a per-block dtype is set,
        # cast u so the stored discrete state lands the same dtype on
        # every step, regardless of upstream promotion.
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        return u

    def _output(self, _time, state, **parameters):
        return state.discrete_state

    def check_types(
        self,
        context,
        error_collector: ErrorCollector = None,
    ):
        inp_data = self.eval_input(context)
        xd = context[self.system_id].discrete_state
        check_state_type(
            self,
            inp_data=inp_data,
            state_data=xd,
            error_collector=error_collector,
        )

UnscentedKalmanFilter

Bases: KalmanFilterBase

Unscented Kalman Filter (UKF) for the following system:

```
x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
y[n]   = g(x[n], u[n]) + v[n]

E(w[n]) = E(v[n]) = 0
E(w[n]w'[n]) = Q(t[n], x[n], u[n])
E(v[n]v'[n] = R(t[n])
E(w[n]v'[n] = N(t[n]) = 0
```

f and g are discrete-time functions of state x[n] and control u[n], while RandGare discrete-time functions of timet[n].Qis a discrete-time function oft[n], x[n], u[n]`. This last aspect is included for zero-order-hold discretization of a continuous-time system

Input ports

(0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n

Output ports

(1) x_hat[n] : state vector estimate at timestep n

Parameters:

Name	Type	Description	Default
dt		float Time step of the discrete-time system	required
forward		Callable A function with signature f(x[n], u[n]) -> x[n+1] that represents f in the above equations.	required
observation		Callable A function with signature g(x[n], u[n]) -> y[n] that represents g in the above equations.	required
G_func		Callable A function with signature G(t[n]) -> G[n] that represents G in the above equations.	required
Q_func		Callable A function with signature Q(t[n], x[n], u[n]) -> Q[n] that represents Q in the above equations.	required
R_func		Callable A function with signature R(t[n]) -> R[n] that represents R in the above equations.	required
x_hat_0		ndarray Initial state estimate	required
P_hat_0		ndarray Initial state covariance matrix estimate	required
alpha		float Sigma point spread to control the amount of nonlinearities taken into account. Usually set to a value (1e-04<= alpha <= 1.0). Default is 1.0.	1.0
beta		float Scaling constant to include prior information about the distribution of the state. Default is 0.0.	0.0
kappa		float Relatively non-critical parameter to control the kurtosis of sigma point distribution. Default is 0.0.	0.0

Source code in jaxonomy/library/state_estimators/unscented_kalman_filter.py

class UnscentedKalmanFilter(KalmanFilterBase):
    """
    Unscented Kalman Filter (UKF) for the following system:

        ```
        x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
        y[n]   = g(x[n], u[n]) + v[n]

        E(w[n]) = E(v[n]) = 0
        E(w[n]w'[n]) = Q(t[n], x[n], u[n])
        E(v[n]v'[n] = R(t[n])
        E(w[n]v'[n] = N(t[n]) = 0
        ```

    `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`,
    while R` and `G` are discrete-time functions of time `t[n]`. `Q` is a discrete-time
    function of `t[n], x[n], u[n]`. This last aspect is included for zero-order-hold
    discretization of a continuous-time system

    Input ports:
        (0) u[n] : control vector at timestep n
        (1) y[n] : measurement vector at timestep n

    Output ports:
        (1) x_hat[n] : state vector estimate at timestep n

    Parameters:
        dt: float
            Time step of the discrete-time system
        forward: Callable
            A function with signature f(x[n], u[n]) -> x[n+1] that represents `f` in
            the above equations.
        observation: Callable
            A function with signature g(x[n], u[n]) -> y[n] that represents `g` in
            the above equations.
        G_func: Callable
            A function with signature G(t[n]) -> G[n] that represents `G` in
            the above equations.
        Q_func: Callable
            A function with signature Q(t[n], x[n], u[n]) -> Q[n] that represents `Q`
            in the above equations.
        R_func: Callable
            A function with signature R(t[n]) -> R[n] that represents `R` in
            the above equations.
        x_hat_0: ndarray
            Initial state estimate
        P_hat_0: ndarray
            Initial state covariance matrix estimate
        alpha: float
            Sigma point spread to control the amount of nonlinearities taken into
            account. Usually set to a value (1e-04<= alpha <= 1.0). Default is 1.0.
        beta: float
            Scaling constant to include prior information about the distribution of
            the state. Default is 0.0.
        kappa: float
            Relatively non-critical parameter to control the kurtosis of sigma point
            distribution. Default is 0.0.
    """

    @parameters(
        static=[
            "dt",
            "forward",
            "observation",
            "G_func",
            "Q_func",
            "R_func",
            "x_hat_0",
            "P_hat_0",
            "alpha",
            "beta",
            "kappa",
        ],
    )
    def __init__(
        self,
        dt,
        forward,
        observation,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        P_hat_0,
        alpha=1.0,
        beta=0.0,
        kappa=0.0,
        is_feedthrough=True,  # TODO: determine automatically?
        name=None,
        **kwargs,
    ):
        super().__init__(dt, x_hat_0, P_hat_0, is_feedthrough, name, **kwargs)

    def initialize(
        self,
        dt,
        forward,
        observation,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        P_hat_0,
        alpha=1.0,
        beta=0.0,
        kappa=0.0,
    ):
        self.G_func = G_func
        self.Q_func = Q_func
        self.R_func = R_func

        self.nx = x_hat_0.size
        self.ny = self.R_func(0.0).shape[0]

        self.alpha = alpha
        self.beta = beta
        self.kappa = kappa

        self.forward = forward
        self.observation = observation

        self.forward_sigma_points = jax.vmap(forward, in_axes=(0, None))
        self.observation_sigma_points = jax.vmap(observation, in_axes=(0, None))

        self.num_sigma_points = 2 * self.nx + 1
        self.lamb = (self.alpha**2.0) * (self.nx + self.kappa) - self.nx
        self.lamb_plus_nx = self.lamb + self.nx

        self.weights_mean = jnp.full(2 * self.nx + 1, 0.5 / (self.lamb + self.nx))
        self.weights_mean = self.weights_mean.at[0].set(
            self.lamb / (self.lamb + self.nx)
        )

        self.weights_cov = jnp.full(2 * self.nx + 1, 0.5 / (self.lamb + self.nx))
        self.weights_cov = self.weights_cov.at[0].set(
            self.lamb / (self.lamb + self.nx) + (1.0 - alpha**2 + beta)
        )

    def _gen_sigma_points(self, mean, cov):
        chol_cov = jsp.linalg.cholesky(
            self.lamb_plus_nx * cov
        )  # upper triangular Cholesky fact.

        sigma_points_plus = mean + chol_cov
        sigma_points_minus = mean - chol_cov

        sigma_points = jnp.vstack([mean, sigma_points_plus, sigma_points_minus])

        return sigma_points

    def _get_weighted_mean_and_cov_from_sigma_points(self, sigma_points):
        mean = jnp.dot(self.weights_mean, sigma_points)
        delta_sigma_points = sigma_points - mean
        cov = delta_sigma_points.T @ jnp.diag(self.weights_cov) @ delta_sigma_points

        return mean, cov

    def _get_weighted_cross_covariance_from_sigma_points(
        self, sigma_points_x, sigma_points_y
    ):
        mean_x = jnp.dot(self.weights_mean, sigma_points_x)
        delta_sigma_points_x = sigma_points_x - mean_x

        mean_y = jnp.dot(self.weights_mean, sigma_points_y)
        delta_sigma_points_y = sigma_points_y - mean_y

        cov_xy = (
            delta_sigma_points_x.T @ jnp.diag(self.weights_cov) @ delta_sigma_points_y
        )

        return cov_xy

    def _correct(self, time, x_hat_minus, P_hat_minus, *inputs):
        u, y = inputs
        u = jnp.atleast_1d(u)
        y = jnp.atleast_1d(y)

        sigma_points_x_minus = self._gen_sigma_points(x_hat_minus, P_hat_minus).reshape(
            (self.num_sigma_points, self.nx)
        )

        sigma_points_y_minus = self.observation_sigma_points(
            sigma_points_x_minus, u
        ).reshape((self.num_sigma_points, self.ny))

        y_mean, Py = self._get_weighted_mean_and_cov_from_sigma_points(
            sigma_points_y_minus
        )

        Pxy = self._get_weighted_cross_covariance_from_sigma_points(
            sigma_points_x_minus,
            sigma_points_y_minus,
        )

        R = self.R_func(time)
        S = Py + R

        # Kalman gain via a linear solve instead of an explicit inverse of the
        # innovation covariance S — more numerically stable. K = Pxy S⁻¹ solves
        # K S = Pxy.
        K = jnp.linalg.solve(S.T, Pxy.T).T

        x_hat_plus = x_hat_minus + jnp.dot(K, y - y_mean)  # n|n
        P_hat_plus = P_hat_minus - K @ S @ K.T  # n|n

        return x_hat_plus, P_hat_plus

    def _propagate(self, time, x_hat_plus, P_hat_plus, *inputs):
        # Predict -- x_hat_plus of current step is propagated to be the
        # x_hat_minus of the next step
        # k+1|k in current step is k|k-1 for next step

        u, y = inputs
        u = jnp.atleast_1d(u)

        G = self.G_func(time)
        Q = self.Q_func(time, x_hat_plus, u)
        GQGT = G @ Q @ G.T

        sigma_points_x_plus = self._gen_sigma_points(x_hat_plus, P_hat_plus).reshape(
            (self.num_sigma_points, self.nx)
        )

        sigma_points_x_minus = self.forward_sigma_points(
            sigma_points_x_plus, u
        ).reshape((self.num_sigma_points, self.nx))

        x_hat_minus, Px = self._get_weighted_mean_and_cov_from_sigma_points(
            sigma_points_x_minus
        )  # n+1|n

        P_hat_minus = Px + GQGT  # n+1|n

        return x_hat_minus, P_hat_minus

    #######################################
    # Make filter for a continuous plant  #
    #######################################

    @staticmethod
    @with_resolved_parameters
    def for_continuous_plant(
        plant,
        dt,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        P_hat_0,
        discretization_method="euler",
        discretized_noise=False,
        alpha=1.0,
        beta=0.0,
        kappa=0.0,
        name=None,
        ui_id=None,
    ):
        """
        Unscented Kalman Filter system for a continuous-time plant.

        The input plant contains the deterministic forms of the forward and observation
        operators:

        ```
            dx/dt = f(x,u)
            y = g(x,u)
        ```

        Note: (i) Only plants with one vector-valued input and one vector-valued output
        are currently supported. Furthermore, the plant LeafSystem/Diagram should have
        only one vector-valued integrator; (ii) the user may pass a plant with
        disturbances (not recommended) as the input plant. In this case, the forward
        and observation evaluations will be corrupted by noise.

        A plant with disturbances of the following form is then considered:

        ```
            dx/dt = f(x,u) + G(t) w         -- (C1)
            y = g(x,u) +  v                 -- (C2)
        ```

        where:

            `w` represents the process noise,
            `v` represents the measurement noise,

        and

        ```
            E(w) = E(v) = 0
            E(ww') = Q(t)
            E(vv') = R(t)
            E(wv') = N(t) = 0
        ```

        This plant is discretized to obtain the following form:

        ```
            x[n+1] = fd(x[n], u[n]) + Gd w[n]  -- (D1)
            y[n]   = gd(x[n], u[n]) + v[n]     -- (D2)

            E(w[n]) = E(v[n]) = 0
            E(w[n]w'[n]) = Qd
            E(v[n]v'[n] = Rd
            E(w[n]v'[n] = Nd = 0
        ```

        The above discretization is performed either via the `euler` or the `zoh`
        method, and an Unscented Kalman Filter estimator for the system of equations
        (D1) and (D2) is returned.

        Note: If `discretized_noise` is True, then it is assumed that the user is
        directly providing Gd, Qd and Rd. If False, then Qd and Rd are computed from
        continuous-time Q, R, and G, and Gd is set to an Identity matrix.

        The returned system will have:

        Input ports:
            (0) u[n] : control vector at timestep n
            (1) y[n] : measurement vector at timestep n

        Output ports:
            (1) x_hat[n] : state vector estimate at timestep n

        Parameters:
            plant : a `Plant` object which can be a LeafSystem or a Diagram.
            dt: float
                Time step for the discretization.
            G_func: Callable
                A function with signature G(t) -> G that represents `G` in
                the continuous-time equations (C1) and (C2).
            Q_func: Callable
                A function with signature Q(t) -> Q that represents `Q` in
                the continuous-time equations (C1) and (C2).
            R_func: Callable
                A function with signature R(t) -> R that represents `R` in
                the continuous-time equations (C1) and (C2).
            x_hat_0: ndarray
                Initial state estimate
            P_hat_0: ndarray
                Initial state covariance matrix estimate. If `None`, an Identity
                matrix is assumed.
            discretization_method: str ("euler" or "zoh")
                Method to discretize the continuous-time plant. Default is "euler".
            discretized_noise: bool
                Whether the user is directly providing Gd, Qd and Rd. Default is False.
                If True, `G_func`, `Q_func`, and `R_func` provide Gd(t), Qd(t), and
                Rd(t), respectively.
            alpha: float
                Sigma point spread to control the amount of nonlinearities taken into
                account. Usually set to a value (1e-04<= alpha <= 1.0). Default is 1.0.
            beta: float
                Scaling constant to include prior information about the distribution of
                the state. Default is 0.0.
            kappa: float
                Relatively non-critical parameter to control the kurtosis of sigma
                point distribution. Default is 0.0.
        """

        (
            forward,
            observation,
            Gd_func,
            Qd_func,
            Rd_func,
        ) = prepare_continuous_plant_for_nonlinear_kalman_filter(
            plant,
            dt,
            G_func,
            Q_func,
            R_func,
            x_hat_0,
            discretization_method,
            discretized_noise,
        )

        nx = x_hat_0.size
        if P_hat_0 is None:
            P_hat_0 = jnp.eye(nx)

        # TODO: If Gd_func is None, compute Gd automatically with u = u + w

        ukf = UnscentedKalmanFilter(
            dt,
            forward,
            observation,
            Gd_func,
            Qd_func,
            Rd_func,
            x_hat_0,
            P_hat_0,
            alpha=alpha,
            beta=beta,
            kappa=kappa,
            name=name,
            ui_id=ui_id,
        )

        return ukf

    ###################################################################################
    # Make filter from direct specification of forward/observaton operators and noise #
    ###################################################################################

    @staticmethod
    @with_resolved_parameters
    def from_operators(
        dt,
        forward,
        observation,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        P_hat_0,
        alpha=1.0,
        beta=0.0,
        kappa=0.0,
        name=None,
        ui_id=None,
    ):
        """
        Unscented Kalman Filter (UKF) for the following system:

        ```
            x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
            y[n]   = g(x[n], u[n]) + v[n]

            E(w[n]) = E(v[n]) = 0
            E(w[n]w'[n]) = Q(t[n], x[n], u[n])
            E(v[n]v'[n] = R(t[n])
            E(w[n]v'[n] = N(t[n]) = 0
        ```

        `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`,
        while `Q` and `R` and `G` are discrete-time functions of time `t[n]`.

        Input ports:
            (0) u[n] : control vector at timestep n
            (1) y[n] : measurement vector at timestep n

        Output ports:
            (1) x_hat[n] : state vector estimate at timestep n

        Parameters:
            dt: float
                Time step of the discrete-time system
            forward: Callable
                A function with signature f(x[n], u[n]) -> x[n+1] that represents `f`
                in the above equations.
            observation: Callable
                A function with signature g(x[n], u[n]) -> y[n] that represents `g` in
                the above equations.
            G_func: Callable
                A function with signature G(t[n]) -> G[n] that represents `G` in
                the above equations.
            Q_func: Callable
                A function with signature Q(t[n]) -> Q[n] that represents
                `Q` in the above equations.
            R_func: Callable
                A function with signature R(t[n]) -> R[n] that represents `R` in
                the above equations.
            x_hat_0: ndarray
                Initial state estimate
            P_hat_0: ndarray
                Initial state covariance matrix estimate
            alpha: float
                Sigma point spread to control the amount of nonlinearities taken into
                account. Usually set to a value (1e-04<= alpha <= 1.0). Default is 1.0.
            beta: float
                Scaling constant to include prior information about the distribution of
                the state. Default is 0.0.
            kappa: float
                Relatively non-critical parameter to control the kurtosis of sigma
                point distribution. Default is 0.0.
        """

        def Q_func_aug(t, x_k, u_k):
            return Q_func(t)

        ukf = UnscentedKalmanFilter(
            dt,
            forward,
            observation,
            G_func,
            Q_func_aug,
            R_func,
            x_hat_0,
            P_hat_0,
            alpha=alpha,
            beta=beta,
            kappa=kappa,
            name=name,
            ui_id=ui_id,
        )

        return ukf

        return ukf

for_continuous_plant(plant, dt, G_func, Q_func, R_func, x_hat_0, P_hat_0, discretization_method='euler', discretized_noise=False, alpha=1.0, beta=0.0, kappa=0.0, name=None, ui_id=None) staticmethod

Unscented Kalman Filter system for a continuous-time plant.

The input plant contains the deterministic forms of the forward and observation operators:

    dx/dt = f(x,u)
    y = g(x,u)

Note: (i) Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator; (ii) the user may pass a plant with disturbances (not recommended) as the input plant. In this case, the forward and observation evaluations will be corrupted by noise.

A plant with disturbances of the following form is then considered:

    dx/dt = f(x,u) + G(t) w         -- (C1)
    y = g(x,u) +  v                 -- (C2)

where:

`w` represents the process noise,
`v` represents the measurement noise,

and

    E(w) = E(v) = 0
    E(ww') = Q(t)
    E(vv') = R(t)
    E(wv') = N(t) = 0

This plant is discretized to obtain the following form:

    x[n+1] = fd(x[n], u[n]) + Gd w[n]  -- (D1)
    y[n]   = gd(x[n], u[n]) + v[n]     -- (D2)

    E(w[n]) = E(v[n]) = 0
    E(w[n]w'[n]) = Qd
    E(v[n]v'[n] = Rd
    E(w[n]v'[n] = Nd = 0

The above discretization is performed either via the euler or the zoh method, and an Unscented Kalman Filter estimator for the system of equations (D1) and (D2) is returned.

Note: If discretized_noise is True, then it is assumed that the user is directly providing Gd, Qd and Rd. If False, then Qd and Rd are computed from continuous-time Q, R, and G, and Gd is set to an Identity matrix.

The returned system will have:

Input ports

(0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n

Output ports

(1) x_hat[n] : state vector estimate at timestep n

Parameters:

Name	Type	Description	Default
plant		a Plant object which can be a LeafSystem or a Diagram.	required
dt		float Time step for the discretization.	required
G_func		Callable A function with signature G(t) -> G that represents G in the continuous-time equations (C1) and (C2).	required
Q_func		Callable A function with signature Q(t) -> Q that represents Q in the continuous-time equations (C1) and (C2).	required
R_func		Callable A function with signature R(t) -> R that represents R in the continuous-time equations (C1) and (C2).	required
x_hat_0		ndarray Initial state estimate	required
P_hat_0		ndarray Initial state covariance matrix estimate. If None, an Identity matrix is assumed.	required
discretization_method		str ("euler" or "zoh") Method to discretize the continuous-time plant. Default is "euler".	'euler'
discretized_noise		bool Whether the user is directly providing Gd, Qd and Rd. Default is False. If True, G_func, Q_func, and R_func provide Gd(t), Qd(t), and Rd(t), respectively.	False
alpha		float Sigma point spread to control the amount of nonlinearities taken into account. Usually set to a value (1e-04<= alpha <= 1.0). Default is 1.0.	1.0
beta		float Scaling constant to include prior information about the distribution of the state. Default is 0.0.	0.0
kappa		float Relatively non-critical parameter to control the kurtosis of sigma point distribution. Default is 0.0.	0.0

Source code in jaxonomy/library/state_estimators/unscented_kalman_filter.py

@staticmethod
@with_resolved_parameters
def for_continuous_plant(
    plant,
    dt,
    G_func,
    Q_func,
    R_func,
    x_hat_0,
    P_hat_0,
    discretization_method="euler",
    discretized_noise=False,
    alpha=1.0,
    beta=0.0,
    kappa=0.0,
    name=None,
    ui_id=None,
):
    """
    Unscented Kalman Filter system for a continuous-time plant.

    The input plant contains the deterministic forms of the forward and observation
    operators:

    ```
        dx/dt = f(x,u)
        y = g(x,u)
    ```

    Note: (i) Only plants with one vector-valued input and one vector-valued output
    are currently supported. Furthermore, the plant LeafSystem/Diagram should have
    only one vector-valued integrator; (ii) the user may pass a plant with
    disturbances (not recommended) as the input plant. In this case, the forward
    and observation evaluations will be corrupted by noise.

    A plant with disturbances of the following form is then considered:

    ```
        dx/dt = f(x,u) + G(t) w         -- (C1)
        y = g(x,u) +  v                 -- (C2)
    ```

    where:

        `w` represents the process noise,
        `v` represents the measurement noise,

    and

    ```
        E(w) = E(v) = 0
        E(ww') = Q(t)
        E(vv') = R(t)
        E(wv') = N(t) = 0
    ```

    This plant is discretized to obtain the following form:

    ```
        x[n+1] = fd(x[n], u[n]) + Gd w[n]  -- (D1)
        y[n]   = gd(x[n], u[n]) + v[n]     -- (D2)

        E(w[n]) = E(v[n]) = 0
        E(w[n]w'[n]) = Qd
        E(v[n]v'[n] = Rd
        E(w[n]v'[n] = Nd = 0
    ```

    The above discretization is performed either via the `euler` or the `zoh`
    method, and an Unscented Kalman Filter estimator for the system of equations
    (D1) and (D2) is returned.

    Note: If `discretized_noise` is True, then it is assumed that the user is
    directly providing Gd, Qd and Rd. If False, then Qd and Rd are computed from
    continuous-time Q, R, and G, and Gd is set to an Identity matrix.

    The returned system will have:

    Input ports:
        (0) u[n] : control vector at timestep n
        (1) y[n] : measurement vector at timestep n

    Output ports:
        (1) x_hat[n] : state vector estimate at timestep n

    Parameters:
        plant : a `Plant` object which can be a LeafSystem or a Diagram.
        dt: float
            Time step for the discretization.
        G_func: Callable
            A function with signature G(t) -> G that represents `G` in
            the continuous-time equations (C1) and (C2).
        Q_func: Callable
            A function with signature Q(t) -> Q that represents `Q` in
            the continuous-time equations (C1) and (C2).
        R_func: Callable
            A function with signature R(t) -> R that represents `R` in
            the continuous-time equations (C1) and (C2).
        x_hat_0: ndarray
            Initial state estimate
        P_hat_0: ndarray
            Initial state covariance matrix estimate. If `None`, an Identity
            matrix is assumed.
        discretization_method: str ("euler" or "zoh")
            Method to discretize the continuous-time plant. Default is "euler".
        discretized_noise: bool
            Whether the user is directly providing Gd, Qd and Rd. Default is False.
            If True, `G_func`, `Q_func`, and `R_func` provide Gd(t), Qd(t), and
            Rd(t), respectively.
        alpha: float
            Sigma point spread to control the amount of nonlinearities taken into
            account. Usually set to a value (1e-04<= alpha <= 1.0). Default is 1.0.
        beta: float
            Scaling constant to include prior information about the distribution of
            the state. Default is 0.0.
        kappa: float
            Relatively non-critical parameter to control the kurtosis of sigma
            point distribution. Default is 0.0.
    """

    (
        forward,
        observation,
        Gd_func,
        Qd_func,
        Rd_func,
    ) = prepare_continuous_plant_for_nonlinear_kalman_filter(
        plant,
        dt,
        G_func,
        Q_func,
        R_func,
        x_hat_0,
        discretization_method,
        discretized_noise,
    )

    nx = x_hat_0.size
    if P_hat_0 is None:
        P_hat_0 = jnp.eye(nx)

    # TODO: If Gd_func is None, compute Gd automatically with u = u + w

    ukf = UnscentedKalmanFilter(
        dt,
        forward,
        observation,
        Gd_func,
        Qd_func,
        Rd_func,
        x_hat_0,
        P_hat_0,
        alpha=alpha,
        beta=beta,
        kappa=kappa,
        name=name,
        ui_id=ui_id,
    )

    return ukf

from_operators(dt, forward, observation, G_func, Q_func, R_func, x_hat_0, P_hat_0, alpha=1.0, beta=0.0, kappa=0.0, name=None, ui_id=None) staticmethod

Unscented Kalman Filter (UKF) for the following system:

    x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
    y[n]   = g(x[n], u[n]) + v[n]

    E(w[n]) = E(v[n]) = 0
    E(w[n]w'[n]) = Q(t[n], x[n], u[n])
    E(v[n]v'[n] = R(t[n])
    E(w[n]v'[n] = N(t[n]) = 0

f and g are discrete-time functions of state x[n] and control u[n], while Q and R and G are discrete-time functions of time t[n].

Input ports

(0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n

Output ports

(1) x_hat[n] : state vector estimate at timestep n

Parameters:

Name	Type	Description	Default
dt		float Time step of the discrete-time system	required
forward		Callable A function with signature f(x[n], u[n]) -> x[n+1] that represents f in the above equations.	required
observation		Callable A function with signature g(x[n], u[n]) -> y[n] that represents g in the above equations.	required
G_func		Callable A function with signature G(t[n]) -> G[n] that represents G in the above equations.	required
Q_func		Callable A function with signature Q(t[n]) -> Q[n] that represents Q in the above equations.	required
R_func		Callable A function with signature R(t[n]) -> R[n] that represents R in the above equations.	required
x_hat_0		ndarray Initial state estimate	required
P_hat_0		ndarray Initial state covariance matrix estimate	required
alpha		float Sigma point spread to control the amount of nonlinearities taken into account. Usually set to a value (1e-04<= alpha <= 1.0). Default is 1.0.	1.0
beta		float Scaling constant to include prior information about the distribution of the state. Default is 0.0.	0.0
kappa		float Relatively non-critical parameter to control the kurtosis of sigma point distribution. Default is 0.0.	0.0

Source code in jaxonomy/library/state_estimators/unscented_kalman_filter.py

@staticmethod
@with_resolved_parameters
def from_operators(
    dt,
    forward,
    observation,
    G_func,
    Q_func,
    R_func,
    x_hat_0,
    P_hat_0,
    alpha=1.0,
    beta=0.0,
    kappa=0.0,
    name=None,
    ui_id=None,
):
    """
    Unscented Kalman Filter (UKF) for the following system:

    ```
        x[n+1] = f(x[n], u[n]) + G(t[n]) w[n]
        y[n]   = g(x[n], u[n]) + v[n]

        E(w[n]) = E(v[n]) = 0
        E(w[n]w'[n]) = Q(t[n], x[n], u[n])
        E(v[n]v'[n] = R(t[n])
        E(w[n]v'[n] = N(t[n]) = 0
    ```

    `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`,
    while `Q` and `R` and `G` are discrete-time functions of time `t[n]`.

    Input ports:
        (0) u[n] : control vector at timestep n
        (1) y[n] : measurement vector at timestep n

    Output ports:
        (1) x_hat[n] : state vector estimate at timestep n

    Parameters:
        dt: float
            Time step of the discrete-time system
        forward: Callable
            A function with signature f(x[n], u[n]) -> x[n+1] that represents `f`
            in the above equations.
        observation: Callable
            A function with signature g(x[n], u[n]) -> y[n] that represents `g` in
            the above equations.
        G_func: Callable
            A function with signature G(t[n]) -> G[n] that represents `G` in
            the above equations.
        Q_func: Callable
            A function with signature Q(t[n]) -> Q[n] that represents
            `Q` in the above equations.
        R_func: Callable
            A function with signature R(t[n]) -> R[n] that represents `R` in
            the above equations.
        x_hat_0: ndarray
            Initial state estimate
        P_hat_0: ndarray
            Initial state covariance matrix estimate
        alpha: float
            Sigma point spread to control the amount of nonlinearities taken into
            account. Usually set to a value (1e-04<= alpha <= 1.0). Default is 1.0.
        beta: float
            Scaling constant to include prior information about the distribution of
            the state. Default is 0.0.
        kappa: float
            Relatively non-critical parameter to control the kurtosis of sigma
            point distribution. Default is 0.0.
    """

    def Q_func_aug(t, x_k, u_k):
        return Q_func(t)

    ukf = UnscentedKalmanFilter(
        dt,
        forward,
        observation,
        G_func,
        Q_func_aug,
        R_func,
        x_hat_0,
        P_hat_0,
        alpha=alpha,
        beta=beta,
        kappa=kappa,
        name=name,
        ui_id=ui_id,
    )

    return ukf

    return ukf

VariableTransportDelay

Bases: LeafSystem

Continuous-time variable transport (pure) delay.

Implements y(t) = u(t - tau(t)) where the delay tau is supplied as a runtime input signal (second input port) rather than as a static parameter. This is the T-107-followup-variable-tau extension to the fixed-delay :class:TransportDelay block (T-107 phase 1).

Mechanism: identical to :class:TransportDelay — a periodic clock at period dt writes the most recent history_length (time, u) pairs into a discrete-state ring buffer, and the (continuous-time) output port performs a linear interpolation over the buffer at t - clip(tau, 0, max_delay_seconds). The clip guards the interpolation against transient out-of-range delay values from upstream blocks; out-of-band tau is clamped (not raised) so that the block remains differentiable everywhere.

Differentiability:

• w.r.t. the data input u: via npa.interp over values, same as :class:TransportDelay.
• w.r.t. the delay input tau: under method="linear" (default, phase 3) via npa.interp's gradient w.r.t. its query coordinate — the standard linear-interp Jacobian, well defined except at sample boundaries where the gradient has a jump discontinuity. Pass method="pchip" (T-107 phase 4) to route through the T-106 backend's monotone cubic Hermite interpolant instead: smooth (C^1) gradient w.r.t. tau across every sample boundary, at the cost of one extra slope-array compute per output evaluation.

The buffer is sized statically from max_delay_seconds: at sample period dt you need at least ceil(max_delay_seconds / dt) + 1 slots; history_length defaults to max(8, ceil(max_delay_seconds / dt) + 4).

Input ports

(0) The input signal u(t). Scalar or array. (1) The delay signal tau(t) in seconds. Runtime scalar in [0, max_delay_seconds]; values outside that range are silently clamped.

Output ports

(0) The delayed signal y(t) = u(t - tau(t)) (or initial_output while t < tau(t)).

Parameters:

Name	Type	Description	Default
dt		Sampling period for the history buffer. Smaller dt ⇒ finer interpolation but a larger ring buffer to cover the same physical delay.	required
max_delay_seconds		Upper bound on the runtime delay value. Used to size the ring buffer and to clip out-of-range tau inputs. Static (compile-time).	required
initial_output		Output value while t < tau(t). Default is 0.0.	0.0
history_length		Number of (time, value) pairs stored. Static (compile-time) — required for vmap/JIT-safe buffer sizing. If None, defaults to max(8, ceil(max_delay_seconds / dt) + 4).	None

Notes

• Default-off / non-touched-block path is byte-equivalent: the existing :class:TransportDelay is untouched.
• Buffer overflow (tau > max_delay_seconds) is clamped to max_delay_seconds rather than raised; this keeps the block differentiable but means the user is responsible for choosing a sufficiently large max_delay_seconds.
• The variable-tau interpolation runs once per output evaluation (continuous-time semantics). For workloads where the delay changes only at major-step granularity, sampling tau at the periodic update would be cheaper — deferred until profiling demands it.
• For arbitrary array-shaped data signals the interpolation is applied elementwise via a static loop over the trailing axes (mirrors :class:TransportDelay).

Source code in jaxonomy/library/dynamics.py

class VariableTransportDelay(LeafSystem):
    """Continuous-time variable transport (pure) delay.

    Implements ``y(t) = u(t - tau(t))`` where the delay ``tau`` is supplied
    as a runtime input signal (second input port) rather than as a static
    parameter. This is the T-107-followup-variable-tau extension to the
    fixed-delay :class:`TransportDelay` block (T-107 phase 1).

    Mechanism: identical to :class:`TransportDelay` — a periodic clock at
    period ``dt`` writes the most recent ``history_length`` ``(time, u)``
    pairs into a discrete-state ring buffer, and the (continuous-time)
    output port performs a linear interpolation over the buffer at
    ``t - clip(tau, 0, max_delay_seconds)``. The clip guards the
    interpolation against transient out-of-range delay values from
    upstream blocks; out-of-band ``tau`` is clamped (not raised) so that
    the block remains differentiable everywhere.

    Differentiability:

    * w.r.t. the data input ``u``: via ``npa.interp`` over ``values``,
      same as :class:`TransportDelay`.
    * w.r.t. the delay input ``tau``: under ``method="linear"`` (default,
      phase 3) via ``npa.interp``'s gradient w.r.t. its query
      coordinate — the standard linear-interp Jacobian, well defined
      except at sample boundaries where the gradient has a jump
      discontinuity. Pass ``method="pchip"`` (T-107 phase 4) to route
      through the T-106 backend's monotone cubic Hermite interpolant
      instead: smooth (C^1) gradient w.r.t. tau across every sample
      boundary, at the cost of one extra slope-array compute per
      output evaluation.

    The buffer is sized statically from ``max_delay_seconds``: at sample
    period ``dt`` you need at least ``ceil(max_delay_seconds / dt) + 1``
    slots; ``history_length`` defaults to
    ``max(8, ceil(max_delay_seconds / dt) + 4)``.

    Input ports:
        (0) The input signal ``u(t)``. Scalar or array.
        (1) The delay signal ``tau(t)`` in seconds. Runtime scalar in
            ``[0, max_delay_seconds]``; values outside that range are
            silently clamped.

    Output ports:
        (0) The delayed signal ``y(t) = u(t - tau(t))`` (or
            ``initial_output`` while ``t < tau(t)``).

    Parameters:
        dt: Sampling period for the history buffer. Smaller ``dt`` ⇒
            finer interpolation but a larger ring buffer to cover the
            same physical delay.
        max_delay_seconds: Upper bound on the runtime delay value. Used
            to size the ring buffer and to clip out-of-range ``tau``
            inputs. Static (compile-time).
        initial_output: Output value while ``t < tau(t)``. Default is
            0.0.
        history_length: Number of ``(time, value)`` pairs stored. Static
            (compile-time) — required for vmap/JIT-safe buffer sizing.
            If None, defaults to
            ``max(8, ceil(max_delay_seconds / dt) + 4)``.

    Notes:
        - Default-off / non-touched-block path is byte-equivalent: the
          existing :class:`TransportDelay` is untouched.
        - Buffer overflow (``tau > max_delay_seconds``) is clamped to
          ``max_delay_seconds`` rather than raised; this keeps the block
          differentiable but means the user is responsible for choosing
          a sufficiently large ``max_delay_seconds``.
        - The variable-tau interpolation runs once per output evaluation
          (continuous-time semantics). For workloads where the delay
          changes only at major-step granularity, sampling ``tau`` at the
          periodic update would be cheaper — deferred until profiling
          demands it.
        - For arbitrary array-shaped data signals the interpolation is
          applied elementwise via a static loop over the trailing axes
          (mirrors :class:`TransportDelay`).
    """

    class _BufferState(NamedTuple):
        # Newest sample at index 0; oldest at index -1. Reversing the
        # buffer gives a monotonically increasing time axis suitable for
        # ``npa.interp``.
        times: "Array"
        values: "Array"

    @parameters(
        static=["dt", "max_delay_seconds", "history_length", "method"],
        dynamic=["initial_output"],
    )
    def __init__(
        self,
        dt,
        max_delay_seconds,
        initial_output=0.0,
        history_length=None,
        method="linear",
        *args,
        dtype=None,
        **kwargs,
    ):
        # T-005 default-float64 / T-038a-followup-mixed-precision-cascade:
        # honor the active precision policy when no explicit dtype was
        # supplied.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(*args, **kwargs)

        if dt is None or float(dt) <= 0.0:
            raise BlockParameterError(
                message=(
                    f"VariableTransportDelay block {self.name!r} requires a "
                    f"positive sample period dt; got {dt!r}."
                ),
                parameter_name="dt",
            )
        # T-107 phase 4: interpolation method over the ring buffer.
        # ``"linear"`` is the phase-1 / phase-3 default (byte-equivalent);
        # ``"pchip"`` routes through the T-106 backend (monotone cubic
        # Hermite) for smooth gradients w.r.t. tau across sample
        # boundaries.
        if method not in ("linear", "pchip"):
            raise BlockParameterError(
                message=(
                    f"VariableTransportDelay block {self.name!r}: method "
                    f"must be 'linear' or 'pchip'; got {method!r}."
                ),
                parameter_name="method",
            )
        self.method = method
        try:
            max_delay_hint = float(max_delay_seconds)
        except (TypeError, ValueError):
            max_delay_hint = 0.0
        if max_delay_hint < 0.0:
            raise BlockParameterError(
                message=(
                    f"VariableTransportDelay block {self.name!r} requires "
                    f"max_delay_seconds >= 0; got {max_delay_seconds!r}."
                ),
                parameter_name="max_delay_seconds",
            )
        if history_length is None:
            history_length = max(8, int(np.ceil(max_delay_hint / dt)) + 4)
        if int(history_length) < 2:
            raise BlockParameterError(
                message=(
                    f"VariableTransportDelay block {self.name!r} requires "
                    f"history_length >= 2; got {history_length!r}."
                ),
                parameter_name="history_length",
            )

        self.dt = float(dt)
        self.max_delay_seconds = float(max_delay_hint)
        self.history_length = int(history_length)

        # Two input ports: (0) data ``u``, (1) delay ``tau``.
        self.declare_input_port()  # u
        self.declare_input_port()  # tau (variable delay)
        self._periodic_update_idx = self.declare_periodic_update()
        self._output_port_idx = self.declare_output_port()

    def initialize(
        self,
        dt,
        max_delay_seconds,
        initial_output,
        history_length=None,
        method=None,
    ):
        # ``history_length`` and ``method`` are static parameters resolved
        # at __init__ time; the framework still passes them for symmetry —
        # drop them here.
        del history_length, method  # noqa: F841

        initial_value = npa.asarray(initial_output)
        if self._dtype is not None:
            initial_value = initial_value.astype(self._dtype)
        self._signal_shape = tuple(initial_value.shape)

        # Pre-fill the times buffer with strictly increasing sentinels
        # below t=0 so that ``npa.interp(t - tau, times[::-1], ...)``
        # clamps to the oldest sample (== ``initial_output``) for any
        # query time before the first real sample has been written.
        sentinel_t0 = -self.dt * (self.history_length + 1) - 1.0
        times = sentinel_t0 + self.dt * np.arange(self.history_length, dtype=np.float64)
        # Newest first: reverse so position 0 is the largest sentinel.
        times = times[::-1].copy()
        if self._dtype is not None:
            times = times.astype(self._dtype)

        values = npa.broadcast_to(
            initial_value, (self.history_length, *self._signal_shape)
        )

        default_state = self._BufferState(
            times=npa.asarray(times), values=npa.asarray(values)
        )
        self.declare_discrete_state(default_value=default_state, as_array=False)

        self.configure_periodic_update(
            self._periodic_update_idx,
            self._update,
            period=self.dt,
            offset=0.0,
        )

        # Output reads time + discrete-state buffer + the ``tau`` input
        # port; mark ``requires_inputs=True`` so the framework wires up
        # both input ports for the lookup.
        self.configure_output_port(
            self._output_port_idx,
            self._output,
            prerequisites_of_calc=[
                DependencyTicket.xd,
                DependencyTicket.time,
                self.input_ports[1].ticket,
            ],
            requires_inputs=True,
            default_value=initial_value,
        )

    def reset_default_values(
        self,
        dt=None,
        max_delay_seconds=None,
        initial_output=None,
        history_length=None,
        method=None,
    ):
        # Mirror TransportDelay's pattern: rebuild defaults if the
        # dynamic ``initial_output`` changes between calls.
        del dt, max_delay_seconds, history_length, method  # noqa: F841

        if initial_output is None:
            return
        initial_value = npa.asarray(initial_output)
        if self._dtype is not None:
            initial_value = initial_value.astype(self._dtype)
        self._signal_shape = tuple(initial_value.shape)

        sentinel_t0 = -self.dt * (self.history_length + 1) - 1.0
        times = sentinel_t0 + self.dt * np.arange(self.history_length, dtype=np.float64)
        times = times[::-1].copy()
        if self._dtype is not None:
            times = times.astype(self._dtype)

        values = npa.broadcast_to(
            initial_value, (self.history_length, *self._signal_shape)
        )
        default_state = self._BufferState(
            times=npa.asarray(times), values=npa.asarray(values)
        )
        self.configure_discrete_state_default_value(
            default_value=default_state, as_array=False
        )
        self.configure_output_port_default_value(
            self._output_port_idx, initial_value
        )

    def _update(self, time, state, *inputs, **_params):
        # Only the data input (port 0) is written into the ring buffer;
        # the delay input (port 1) is consumed at output evaluation time.
        u = inputs[0]
        if self._dtype is not None:
            u = npa.asarray(u).astype(self._dtype)
        buf = state.discrete_state
        new_times = npa.roll(buf.times, shift=1, axis=0).at[0].set(time)
        new_values = npa.roll(buf.values, shift=1, axis=0).at[0].set(u)
        return self._BufferState(times=new_times, values=new_values)

    def _output(self, time, state, *inputs, **params):
        buf = state.discrete_state
        # Reverse so the time axis is monotonically increasing for
        # ``npa.interp``: index 0 is oldest, index -1 is newest.
        xp = buf.times[::-1]
        fp = buf.values[::-1]
        # Variable-tau: read the delay from the second input port and
        # clamp it into ``[0, max_delay_seconds]``. Out-of-band values
        # are silently clipped (rather than raised) so the block stays
        # differentiable everywhere.
        tau_raw = inputs[1]
        tau = npa.clip(tau_raw, 0.0, self.max_delay_seconds)
        initial_output = params["initial_output"]

        query_t = time - tau

        # T-107 phase 4: dispatch on interpolation method. Linear stays
        # on ``npa.interp`` for byte-equivalence with phase 3 (the
        # established default); PCHIP routes through the T-106 backend
        # for smooth gradients w.r.t. tau across sample boundaries.
        if self.method == "pchip":
            from .lookup_table import interp_1d as _interp_1d

            def _scalar_interp(values_1d):
                # ``interp_1d`` returns a JAX array; the surrounding
                # npa context handles eager-numpy callers.
                return _interp_1d(query_t, xp, values_1d, method="pchip")
        else:
            def _scalar_interp(values_1d):
                return npa.interp(query_t, xp, values_1d)

        if len(self._signal_shape) == 0:
            y = _scalar_interp(fp)
        else:
            # The 1-D interpolator only handles 1-D ``fp``; statically
            # loop over the trailing axes (shape known at trace time).
            flat_fp = fp.reshape((self.history_length, -1))
            ys = [_scalar_interp(flat_fp[:, i]) for i in range(flat_fp.shape[1])]
            y = npa.stack(ys).reshape(self._signal_shape)

        # Hold the initial output before the first physical sample is
        # available; explicitly gate on ``time < tau`` to keep semantics
        # robust to dtype/shape pre-fill quirks.
        y = npa.where(time < tau, npa.asarray(initial_output), y)
        if self._dtype is not None:
            y = npa.asarray(y).astype(self._dtype)
        return y

VideoSink

Bases: LeafSystem

Records RGB frames to a video file.

Parameters:

Name	Type	Description	Default
dt	float	Interval at which to record frames.	required
file_name	str	Name of the video file to write to (optional).	required

Source code in jaxonomy/library/video.py

class VideoSink(LeafSystem):
    """Records RGB frames to a video file.

    Parameters:
        dt: Interval at which to record frames.
        file_name: Name of the video file to write to (optional).
    """

    @parameters(static=["dt", "file_name"])
    def __init__(self, dt: float, file_name: str, **kwargs):
        super().__init__(**kwargs)

        self.dt = dt
        self.fps = 1 / dt
        file_name = str(file_name)
        ext = ".mp4" if not file_name.endswith(".mp4") else ""
        self.file_name = file_name + ext
        self.writer: "VideoWriter" = None
        self.frame_id = 0

        self.declare_input_port("frame")

        def _io_cb(time, state, *inputs, **parameters) -> Array:
            return io_callback(self._video_cb, npa.intx(0), time, inputs[0])

        self.declare_output_port(
            _io_cb,
            name="frame_id",
            requires_inputs=True,
            period=dt,
            offset=dt,
        )

    def _init_video(self, frame: Array):
        if len(frame.shape) != 3 or frame.shape[2] != 3:
            raise StaticError(
                f"Input frame must be an RGB image, got invalid shape: {frame.shape}",
                system=self,
            )

        # A note on codecs:
        # vp9 (vp09) works in browsers, but it's a bit slow to encode
        # MPEG-4 (mp4v) is faster, but not supported in browsers
        # H264 (avc1) is supported but plagued with patents
        # av1 (AV01) broke my computer

        os.makedirs(os.path.dirname(self.file_name), exist_ok=True)

        h, w, _ = frame.shape
        self.writer = cv2.VideoWriter(
            self.file_name,
            cv2.VideoWriter_fourcc(*"vp09"),
            self.fps,
            (w, h),
        )
        if not self.writer.isOpened():
            raise StaticError(
                f"Failed to open video file {self.file_name}",
                system=self,
            )

        logger.info("Writing video of size %sx%s to file: %s", w, h, self.file_name)

    def post_simulation_finalize(self) -> None:
        if self.writer is not None:
            self.writer.release()
        return super().post_simulation_finalize()

    def _video_cb(self, time: Array, frame: Array) -> Array:
        image = np.array(frame)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        if self.writer is None:
            self._init_video(image)
        self.writer.write(image)

        # jax.debug.print(
        #     "Wrote frame {frame_id} to video file at time {time}",
        #     frame_id=self.frame_id,
        #     time=time,
        # )

        frame_id = self.frame_id
        self.frame_id += 1
        return npa.intx(frame_id)

VideoSource

Bases: LeafSystem

Reads frames from a video file.

Parameters:

Name	Type	Description	Default
file_name	str	Name of the video file to read from.	required
no_repeat		Whether to stop at the end of the video or loop back to the beginning.	False

Source code in jaxonomy/library/video.py

class VideoSource(LeafSystem):
    """Reads frames from a video file.

    Parameters:
        file_name: Name of the video file to read from.
        no_repeat: Whether to stop at the end of the video or loop back to the beginning.
    """

    @parameters(static=["file_name", "no_repeat"])
    def __init__(self, file_name: str, no_repeat=False, **kwargs):
        super().__init__(**kwargs)

        self.repeat = not no_repeat
        self.file_name = str(file_name)
        self.frame_id: np.intx = 0
        self.reached_end = False

        self.reader = cv2.VideoCapture(self.file_name)
        if not self.reader.isOpened():
            raise BlockInitializationError(
                f"Could not open video file '{self.file_name}'", system=self
            )

        self.width = int(self.reader.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.depth = 1 if bool(self.reader.get(cv2.CAP_PROP_MONOCHROME)) else 3
        self.fps = self.reader.get(cv2.CAP_PROP_FPS) or 30
        self.video_length = int(self.reader.get(cv2.CAP_PROP_FRAME_COUNT))

        logger.info(
            "Opened video file '%s' with size %sx%s, %s frames, %s fps",
            self.file_name,
            self.width,
            self.height,
            self.video_length,
            self.fps,
            **logdata(block=self),
        )

        self.last_frame = np.zeros(
            (self.height, self.width, self.depth), dtype=np.uint8
        )

        def _frame_cb(time, state, *inputs, **parameters) -> Array:
            def cb(time) -> Array:
                return self._source_cb(time)

            return io_callback(cb, self.last_frame, time)

        dt = 1 / self.fps
        self.declare_output_port(
            _frame_cb,
            name="frame",
            period=dt,
            offset=dt,
            requires_inputs=False,
        )

        def _frame_id_cb(time, state, *inputs, **parameters) -> Array:
            return io_callback(self._frame_id_cb, npa.intx(0))

        self.declare_output_port(
            _frame_id_cb,
            name="frame_id",
            period=dt,
            offset=dt,
            default_value=npa.intx(0),
            requires_inputs=False,
        )

        if not self.repeat:

            def _stopped_cb(time, state, *inputs, **parameters) -> Array:
                return io_callback(self._stopped_cb, npa.bool_(False))

            self.declare_output_port(
                _stopped_cb,
                name="stopped",
                period=dt,
                offset=dt,
                default_value=npa.bool_(False),
                requires_inputs=False,
            )

    def post_simulation_finalize(self) -> None:
        if self.reader is not None:
            self.reader.release()
        return super().post_simulation_finalize()

    def _source_cb(self, time: float) -> Array:
        if self.reached_end:
            return self.last_frame

        if not self.repeat and int(time * self.fps + FPS_EPSILON) >= self.video_length:
            self.reached_end = True
            return self.last_frame

        # jax.debug.print(
        #     "Reading frame {frame_id} to video file at time {time}",
        #     frame_id=self.reader.get(cv2.CAP_PROP_POS_FRAMES),
        #     time=time,
        # )

        self.frame_id = int(time * self.fps + FPS_EPSILON) % self.video_length
        self.reader.set(cv2.CAP_PROP_POS_FRAMES, self.frame_id)

        ret, frame = self.reader.read()
        if not ret:
            if self.repeat:
                self.reader.set(cv2.CAP_PROP_POS_FRAMES, 0)
                ret, frame = self.reader.read()
            else:
                self.reached_end = True
                self.reader.release()
                self.reader = None
                return self.last_frame

        if not ret:
            raise BlockRuntimeError("Failed to read frame from video file", system=self)

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        self.last_frame = frame
        return frame

    def _frame_id_cb(self) -> Array:
        return npa.intx(self.frame_id)

    def _stopped_cb(self) -> Array:
        return self.reached_end

WhenDisabled

Allowed string values for the when_disabled kwarg.

Source code in jaxonomy/library/conditional.py

class WhenDisabled:
    """Allowed string values for the ``when_disabled`` kwarg."""

    RESET = "reset"
    HOLD = "hold"
    PASSTHROUGH = "passthrough"

    @classmethod
    def valid(cls) -> tuple[str, ...]:
        return (cls.RESET, cls.HOLD, cls.PASSTHROUGH)

WhiteNoise

Bases: LeafSystem

Continuous-time white noise generator.

Generates a band-limited white noise signal using a sinc-interpolated random number generator. The output signal is a continuous-time signal, but the underlying random number generator is discrete-time. As a result, the signal is not truly white, but is band-limited by the sample rate. The resulting signal has the following approximate power spectral density:

S(f) = A * fs if |f| < fs else 0,

where A is the noise power and fs = 1/dt is the sample rate.

See Ch. 10.4 in Baraniuk, "Signal Processing and Modeling" for details: https://shorturl.at/floRZ

The output signal will have variance A, zero mean, and will decorrelate at the sample rate.

Input ports

None

Output ports

(0) The band-limited white noise signal with variance noise_power, zero mean, and correlation time dt.

Parameters:

Name	Type	Description	Default
correlation_time		The correlation time of the output signal and the inverse of the bandwidth. It is the sample frequency of the underlying random number generator.	required
noise_power	float	The variance of the white noise signal. Also scales the amplitude of the power spectral density.	1.0
num_samples	int	The number of samples to use for sinc interpolation. More samples will result in a more accurate approximation of the ideal power spectrum, but will also increase the computational cost. The default of 10 is sufficient for most applications.	10
seed	int	An integer seed for the random number generator. If None, a random 32-bit seed will be generated.	None
dtype	DTypeLike	data type of the random number. If None, defaults to float.	None
shape	ShapeLike	The shape of the output signal. If empty, the output will be a scalar.	()

Source code in jaxonomy/library/random.py

class WhiteNoise(LeafSystem):
    """Continuous-time white noise generator.

    Generates a band-limited white noise signal using a sinc-interpolated random
    number generator.  The output signal is a continuous-time signal, but the
    underlying random number generator is discrete-time.  As a result, the signal
    is not truly white, but is band-limited by the sample rate.  The resulting signal
    has the following approximate power spectral density:
    ```
    S(f) = A * fs if |f| < fs else 0,
    ```
    where `A` is the noise power and `fs = 1/dt` is the sample rate.

    See Ch. 10.4 in Baraniuk, "Signal Processing and Modeling" for details:
        https://shorturl.at/floRZ

    The output signal will have variance `A`, zero mean, and will decorrelate at
    the sample rate.

    Input ports:
        None

    Output ports:
        (0) The band-limited white noise signal with variance `noise_power`, zero
            mean, and correlation time `dt`.

    Parameters:
        correlation_time: The correlation time of the output signal and the inverse of
            the bandwidth. It is the sample frequency of the underlying random number
            generator.
        noise_power: The variance of the white noise signal. Also scales the amplitude
            of the power spectral density.
        num_samples: The number of samples to use for sinc interpolation.  More samples
            will result in a more accurate approximation of the ideal power spectrum,
            but will also increase the computational cost.  The default of 10 is
            sufficient for most applications.
        seed: An integer seed for the random number generator. If None, a random 32-bit
            seed will be generated.
        dtype: data type of the random number.  If None, defaults to float.
        shape: The shape of the output signal.  If empty, the output will be a scalar.
    """

    class RNGState(NamedTuple):
        key: Array
        samples: Array
        t_last: float = 0.0

    @classmethod
    def with_key(cls, key: "jax.Array", **kwargs) -> "WhiteNoise":
        """
        Construct WhiteNoise with an explicit JAX PRNGKey.

        Use this when you need independent noise streams in 
        batched (jax.vmap) simulations.

        Example:
            keys = jax.random.split(jax.random.PRNGKey(0), 16)

            # Each diagram gets a different key
            diagrams = [
                build_diagram_with(
                    WhiteNoise.with_key(keys[i], ...)
                )
                for i in range(16)
            ]

            # OR with with_parameters (preferred):
            diagram.with_parameters({"noise.key": keys[i]})

        Args:
            key: JAX PRNGKey array (shape (2,) for default RNG)
            **kwargs: other constructor arguments
        """
        instance = cls(**kwargs)
        instance._explicit_key = key
        return instance

    @parameters(
        static=["num_samples", "shape", "seed", "noise_power"],
        dynamic=["correlation_time"],
    )
    def __init__(
        self,
        correlation_time,
        noise_power: float = 1.0,
        num_samples: int = 10,
        seed: int = None,
        dtype: DTypeLike = None,
        shape: ShapeLike = (),
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.dtype = dtype

        self.declare_output_port(self._output)
        self.declare_periodic_update(
            self._update,
            period=correlation_time,
            offset=0.0,
        )

    def initialize(
        self,
        correlation_time,
        noise_power: float = 1.0,
        num_samples: int = 10,
        seed: int = None,
        shape: ShapeLike = (),
    ):
        self.shape = tuple(map(int, shape))
        self.N = num_samples

        self.noise_power = noise_power
        self.shift = np.arange(self.N) - (self.N - 1) / 2
        self.rng = partial(random.normal, dtype=self.dtype)

        # The default state is a tuple of (key, samples) pairs.  The continuous-time
        # output signal is reconstructed from the samples using a sinc interpolation.
        if hasattr(self, "_explicit_key"):
            key = self._explicit_key
        else:
            seed = (
                np.random.randint(0, 2**32, dtype=np.int64) if seed is None else int(seed)
            )
            key = random.PRNGKey(int(seed))
        key, subkey = random.split(key)
        default_state = self.RNGState(
            key=key,
            samples=self.sample(subkey, shape=(self.N, *self.shape)),
        )
        self.declare_discrete_state(default_value=default_state, as_array=False)

    def sample(self, key, shape):
        return jnp.sqrt(self.noise_power) * self.rng(key, shape)

    def _output(self, time, state, *_inputs, **parameters):
        t_last = state.discrete_state.t_last

        # Time relative to the last discrete sample, in units of
        # samples.  This is the argument to the sinc function.
        w = (time - t_last) / parameters["correlation_time"] - self.shift

        # Clip the time values to limit discontinuities resulting
        # from sample updates.
        w = jnp.clip(w, -self.N // 2, self.N // 2)

        # Shift the axes so that the last axis is the sample index.
        # This is the index that will be contracted over
        samples = jnp.moveaxis(state.discrete_state.samples, 0, -1)

        return jnp.sum(samples * jnp.sinc(w), axis=-1)

    def _update(self, time, state, *_inputs, **_parameters):
        key, subkey = random.split(state.discrete_state.key)

        new_sample = self.sample(subkey, (1, *self.shape))
        samples = jnp.concatenate((state.discrete_state.samples[1:], new_sample))

        return self.RNGState(
            key=key,
            samples=samples,
            t_last=time,
        )

with_key(key, **kwargs) classmethod

Construct WhiteNoise with an explicit JAX PRNGKey.

Use this when you need independent noise streams in batched (jax.vmap) simulations.

Example

keys = jax.random.split(jax.random.PRNGKey(0), 16)

Each diagram gets a different key

diagrams = [ build_diagram_with( WhiteNoise.with_key(keys[i], ...) ) for i in range(16) ]

OR with with_parameters (preferred):

diagram.with_parameters({"noise.key": keys[i]})

Parameters:

Name	Type	Description	Default
key	'jax.Array'	JAX PRNGKey array (shape (2,) for default RNG)	required
**kwargs		other constructor arguments	{}

Source code in jaxonomy/library/random.py

@classmethod
def with_key(cls, key: "jax.Array", **kwargs) -> "WhiteNoise":
    """
    Construct WhiteNoise with an explicit JAX PRNGKey.

    Use this when you need independent noise streams in 
    batched (jax.vmap) simulations.

    Example:
        keys = jax.random.split(jax.random.PRNGKey(0), 16)

        # Each diagram gets a different key
        diagrams = [
            build_diagram_with(
                WhiteNoise.with_key(keys[i], ...)
            )
            for i in range(16)
        ]

        # OR with with_parameters (preferred):
        diagram.with_parameters({"noise.key": keys[i]})

    Args:
        key: JAX PRNGKey array (shape (2,) for default RNG)
        **kwargs: other constructor arguments
    """
    instance = cls(**kwargs)
    instance._explicit_key = key
    return instance

ZeroOrderHold

Bases: LeafSystem

Implements a "zero-order hold" A/D conversion.

https://en.wikipedia.org/wiki/Zero-order_hold

The block implements a "zero-order hold" with the following difference equation for input signal u and output signal y:

    y[k] = u[k]

The block does not maintain an internal state, but simply holds the value of the input signal at the previous update time. As a result, the block is "feedthrough" from its inputs to outputs and cannot be used to break an algebraic loop. The data type of this hold value is inferred from upstream blocks.

Input ports

(0) The input signal.

Output ports

(0) The "hold" value of the input signal. If the input signal is continuous, then the output will be the value of the input signal at the previous update time. If the input signal is discrete and synchonous with the block, the output will be the value of the input signal at the current time (i.e. identical to the input signal).

Parameters:

Name	Type	Description	Default
dt		The time step of the discrete update.	required

Source code in jaxonomy/library/dynamics.py

class ZeroOrderHold(LeafSystem):
    """Implements a "zero-order hold" A/D conversion.

    https://en.wikipedia.org/wiki/Zero-order_hold

    The block implements a "zero-order hold" with the following difference equation
    for input signal `u` and output signal `y`:
    ```
        y[k] = u[k]
    ```

    The block does not maintain an internal state, but simply holds the value of the
    input signal at the previous update time.  As a result, the block is "feedthrough"
    from its inputs to outputs and cannot be used to break an algebraic loop. The data
    type of this hold value is inferred from upstream blocks.

    Input ports:
        (0) The input signal.

    Output ports:
        (0) The "hold" value of the input signal.  If the input signal is continuous,
            then the output will be the value of the input signal at the previous
            update time.  If the input signal is discrete and synchonous with the
            block, the output will be the value of the input signal at the current
            time (i.e. identical to the input signal).

    Parameters:
        dt:
            The time step of the discrete update.
    """

    @parameters(static=["dt"])
    def __init__(self, dt, *args, dtype=None, **kwargs):
        # T-038a-followup-other-blocks: per-block dtype override; stored
        # outside the @parameters list so it does not round-trip through
        # model JSON or get JAX-traced.
        # T-038a-followup-mixed-precision-cascade: when no explicit
        # ``dtype=`` kwarg was passed, fall back to the active
        # ``precision_policy`` context manager's dtype, if any.
        if dtype is None:
            from ..precision import active_precision_policy

            dtype = active_precision_policy()
        self._dtype = dtype
        super().__init__(*args, **kwargs)
        self.dt = dt

        self.declare_input_port()
        self.declare_output_port(
            self._output,
            period=dt,
            offset=0.0,
            prerequisites_of_calc=[self.input_ports[0].ticket, DependencyTicket.xd],
        )

    def _output(self, _time, _state, u, **_params):
        # Every dt seconds, update the state to the current input value
        if self._dtype is not None:
            # T-038a-followup-other-blocks: cast the held value to the
            # per-block dtype.
            u = npa.asarray(u).astype(self._dtype)
        return u

ForEach(submodel, n, n_inputs=1, in_axes=None, name=None)

Container block: evaluate a submodel n times in parallel.

ForEach is a block-diagram-vocabulary alias for the existing :class:jaxonomy.library.ReplicatedFunction (T-010). It exists so that users familiar with the ForEach block name can find it without paying a duplication tax: the implementation is exactly :class:ReplicatedFunction under the hood.

Parameters:

Name	Type	Description	Default
submodel	Callable	Callable f(*inputs) -> output. Must be JAX-traceable.	required
n	int	Number of replicas (the iteration count).	required
n_inputs	int	Number of input ports the block declares.	1
in_axes		As in :func:jax.vmap / ReplicatedFunction: a length-n_inputs tuple of 0 (input is batched along the leading axis) or None (input is broadcast). Default is all-batched.	None
name	str | None	Optional block name.	None

Returns:

Type	Description
	A configured :class:ReplicatedFunction instance, ready to be
	wired into a :class:DiagramBuilder.

Source code in jaxonomy/framework/containers.py

def ForEach(
    submodel: Callable,
    n: int,
    n_inputs: int = 1,
    in_axes=None,
    name: str | None = None,
):
    """Container block: evaluate a submodel ``n`` times in parallel.

    ``ForEach`` is a block-diagram-vocabulary alias for the existing
    :class:`jaxonomy.library.ReplicatedFunction` (T-010). It exists so
    that users familiar with the ``ForEach`` block name can find it
    without paying a duplication tax: the implementation
    is exactly :class:`ReplicatedFunction` under the hood.

    Args:
        submodel: Callable ``f(*inputs) -> output``. Must be
            JAX-traceable.
        n: Number of replicas (the iteration count).
        n_inputs: Number of input ports the block declares.
        in_axes: As in :func:`jax.vmap` / ``ReplicatedFunction``: a
            length-``n_inputs`` tuple of ``0`` (input is batched along
            the leading axis) or ``None`` (input is broadcast). Default
            is all-batched.
        name: Optional block name.

    Returns:
        A configured :class:`ReplicatedFunction` instance, ready to be
        wired into a :class:`DiagramBuilder`.
    """
    # Lazy import: ReplicatedFunction lives in jaxonomy.library, which
    # imports the framework. Importing it eagerly here would create a
    # cycle. The lazy import is exercised only when a user actually
    # constructs a ForEach block.
    from ..library.replicated import ReplicatedFunction

    kwargs: dict = {}
    if name is not None:
        kwargs["name"] = name
    return ReplicatedFunction(
        submodel=submodel,
        n=n,
        n_inputs=n_inputs,
        in_axes=in_axes,
        **kwargs,
    )

RateTransition(input_dt, output_dt, initial_state=0.0, *, name=None, dtype=None, **kwargs)

Auto-pick the right rate-bridging block based on input_dt vs output_dt.

• input_dt > output_dt (slow source → fast destination): :class:ZeroOrderHold at output_dt (the fast rate). The held value is whatever the upstream slow block last produced; the ZOH re-samples on every fast tick.
• input_dt < output_dt (fast source → slow destination): :class:Decimator at output_dt (the slow rate).
• input_dt == output_dt (same rate): :class:UnitDelay at input_dt — a one-step delay so adjacent same-rate blocks can still break feedthrough loops.

Both ZOH and Decimator paths are tagged with the _jaxonomy_rate_transition marker so :func:jaxonomy.simulation.rate_groups.detect_rate_mismatches silences the rate-mismatch warning across the connection. The same-rate (UnitDelay) path does not need the marker because it cannot itself be a rate mismatch.

Parameters:

Name	Type	Description	Default
input_dt		Sample period of the upstream block.	required
output_dt		Sample period of the downstream block.	required
initial_state		Initial output value (only meaningful for the same-rate UnitDelay and the fast→slow Decimator branches; ZeroOrderHold ignores it in Phase 1).	0.0
name		Optional block name.	None
dtype		Optional per-block dtype (forwarded to the underlying block). See T-038a-followup-other-blocks.	None
**kwargs		Forwarded to the underlying block constructor.	{}

Returns:

Name	Type	Description
A		class:LeafSystem instance: ZeroOrderHold,
		class:Decimator, or :class:`UnitDelay``.

Source code in jaxonomy/library/dynamics.py

def RateTransition(
    input_dt,
    output_dt,
    initial_state=0.0,
    *,
    name=None,
    dtype=None,
    **kwargs,
):
    """Auto-pick the right rate-bridging block based on ``input_dt`` vs ``output_dt``.

    * ``input_dt > output_dt`` (slow source → fast destination):
      :class:`ZeroOrderHold` at ``output_dt`` (the fast rate).  The held
      value is whatever the upstream slow block last produced; the ZOH
      re-samples on every fast tick.
    * ``input_dt < output_dt`` (fast source → slow destination):
      :class:`Decimator` at ``output_dt`` (the slow rate).
    * ``input_dt == output_dt`` (same rate): :class:`UnitDelay` at
      ``input_dt`` — a one-step delay so adjacent same-rate blocks can
      still break feedthrough loops.

    Both ZOH and Decimator paths are tagged with the
    ``_jaxonomy_rate_transition`` marker so
    :func:`jaxonomy.simulation.rate_groups.detect_rate_mismatches`
    silences the rate-mismatch warning across the connection.  The
    same-rate (``UnitDelay``) path does not need the marker because it
    cannot itself be a rate mismatch.

    Args:
        input_dt: Sample period of the upstream block.
        output_dt: Sample period of the downstream block.
        initial_state: Initial output value (only meaningful for the
            same-rate ``UnitDelay`` and the fast→slow ``Decimator``
            branches; ``ZeroOrderHold`` ignores it in Phase 1).
        name: Optional block name.
        dtype: Optional per-block dtype (forwarded to the underlying
            block).  See ``T-038a-followup-other-blocks``.
        **kwargs: Forwarded to the underlying block constructor.

    Returns:
        A :class:`LeafSystem` instance: ``ZeroOrderHold``,
        :class:`Decimator`, or :class:`UnitDelay``.
    """
    if input_dt > output_dt:
        # Slow → fast: ZOH at the fast rate.  Tag the instance with the
        # rate-transition marker so ``detect_rate_mismatches`` recognises
        # it as a bridge.
        block = ZeroOrderHold(
            dt=output_dt, name=name, dtype=dtype, **kwargs
        )
        # Instance-level attribute override: the base ``ZeroOrderHold``
        # class is *not* always a rate transition (most users place a
        # ZOH at a single rate, not as a bridge), so we tag the
        # individual instance returned by this factory.
        block._jaxonomy_rate_transition = True
        return block
    if input_dt < output_dt:
        # Fast → slow: explicit Decimator (which sets the marker on the
        # class).
        return Decimator(
            input_dt=input_dt,
            output_dt=output_dt,
            initial_state=initial_state,
            name=name,
            dtype=dtype,
            **kwargs,
        )
    # Same rate: a one-step UnitDelay so users can still break a
    # feedthrough loop at a same-rate boundary.  No bridge marker
    # needed — the rate is identical on both sides.
    return UnitDelay(
        dt=input_dt,
        initial_state=initial_state,
        name=name,
        dtype=dtype,
        **kwargs,
    )

balanced_realization(sys)

Internally-balanced realization of sys.

Returns (balanced_system, hsv) where balanced_system is an equivalent :class:LinearizedSystem whose controllability and observability Gramians are equal and diagonal, with the Hankel singular values hsv on the diagonal (Moore 1981; square-root algorithm of Laub, Heath, Paige & Ward 1987).

Requires a stable sys. A non-minimal or stiff system (Gramians only numerically semidefinite) is handled — see :func:_psd_sqrt — but its ~zero Hankel-value states are ill-defined in the full balanced form; use :func:balanced_truncation or :func:minimal_realization to remove them.

Source code in jaxonomy/library/rom/linear_mor.py

def balanced_realization(sys):
    r"""Internally-balanced realization of ``sys``.

    Returns ``(balanced_system, hsv)`` where ``balanced_system`` is an
    equivalent :class:`LinearizedSystem` whose controllability and
    observability Gramians are equal and diagonal, with the Hankel singular
    values ``hsv`` on the diagonal (Moore 1981; square-root algorithm of
    Laub, Heath, Paige & Ward 1987).

    Requires a stable ``sys``. A non-minimal or stiff system (Gramians only
    numerically semidefinite) is handled — see :func:`_psd_sqrt` — but its
    ~zero Hankel-value states are ill-defined in the *full* balanced form;
    use :func:`balanced_truncation` or :func:`minimal_realization` to remove
    them.
    """
    A, B, C, D = _abcd(sys)
    T, Tinv, hsv = _balancing_transform(A, B, C, sys.dt)
    Ab = Tinv @ A @ T
    Bb = Tinv @ B
    Cb = C @ T
    return _make_reduced(sys, Ab, Bb, Cb, D), hsv

balanced_truncation(sys, order=None, tol=None)

Balanced truncation (Moore 1981).

Balances sys and keeps the states associated with the largest Hankel singular values.

Order selection:

• order given — keep exactly that many states.
• tol given (and order is None) — keep the fewest states whose retained "energy" Σσ_kept² / Σσ² is at least 1 - tol; i.e. tol is the fraction of Gramian energy allowed to be discarded.
• neither given — no truncation (returns the balanced realization).

The returned :class:LinearizedSystem additionally exposes:

• .hsv — the full Hankel-singular-value spectrum,
• .reduced_order — the retained state count r,
• .error_bound — the a priori :math:H_\infty error bound :math:\lVert G - G_r\rVert_\infty \le 2\sum_{i>r}\sigma_i (Glover 1984 / Enns 1984).

Source code in jaxonomy/library/rom/linear_mor.py

def balanced_truncation(sys, order=None, tol=None):
    r"""Balanced truncation (Moore 1981).

    Balances ``sys`` and keeps the states associated with the largest Hankel
    singular values.

    Order selection:

    * ``order`` given — keep exactly that many states.
    * ``tol`` given (and ``order`` is ``None``) — keep the fewest states whose
      retained "energy" ``Σσ_kept² / Σσ²`` is at least ``1 - tol``; i.e. ``tol``
      is the fraction of Gramian energy allowed to be discarded.
    * neither given — no truncation (returns the balanced realization).

    The returned :class:`LinearizedSystem` additionally exposes:

    * ``.hsv`` — the full Hankel-singular-value spectrum,
    * ``.reduced_order`` — the retained state count ``r``,
    * ``.error_bound`` — the a priori :math:`H_\infty` error bound
      :math:`\lVert G - G_r\rVert_\infty \le 2\sum_{i>r}\sigma_i`
      (Glover 1984 / Enns 1984).
    """
    A, B, C, D = _abcd(sys)
    T, Tinv, hsv = _balancing_transform(A, B, C, sys.dt)
    n = A.shape[0]

    if order is not None:
        r = int(order)
    elif tol is not None:
        energy = np.cumsum(hsv**2)
        total = energy[-1]
        r = int(np.searchsorted(energy, (1.0 - tol) * total) + 1)
    else:
        r = n
    r = max(1, min(r, n))

    Ab = Tinv @ A @ T
    Bb = Tinv @ B
    Cb = C @ T

    reduced = _make_reduced(sys, Ab[:r, :r], Bb[:r, :], Cb[:, :r], D)
    reduced.hsv = hsv
    reduced.reduced_order = r
    reduced.error_bound = float(2.0 * np.sum(hsv[r:]))
    return reduced

bode_data(linsys, omegas)

Return matplotlib-ready Bode arrays for linsys.

Handles MIMO systems out of the box: when the underlying :func:frequency_response returns shape (K, p, m) with p > 1 or m > 1, the returned magnitude_db / phase_deg arrays keep the same (K, p, m) shape — one Bode pair per (output, input) channel pair. The phase is unwrapped along the frequency axis (axis=0) independently for each channel, which is the standard convention for MIMO Bode plots.

Parameters:

Name	Type	Description	Default
linsys	LinearizedSystem	A :class:LinearizedSystem (p outputs, m inputs).	required
omegas		1-D array-like of angular frequencies ω (rad/s).	required

Returns:

Type	Description
	Dictionary with keys: "omega" — angular frequencies (rad/s), shape (K,); "freq_hz" — ω / (2π) for log-Hz plotting, shape (K,); "magnitude_db" — 20 log₁₀ |G(jω)| (dB); shape (K,) for SISO (squeezed for backward-compatibility with phase 1), shape (K, p, m) for MIMO; "phase_deg" — phase in degrees, unwrapped along the frequency axis; same shape as magnitude_db.

Source code in jaxonomy/library/linearization_workflow.py

def bode_data(linsys: LinearizedSystem, omegas):
    """Return matplotlib-ready Bode arrays for ``linsys``.

    Handles MIMO systems out of the box: when the underlying
    :func:`frequency_response` returns shape ``(K, p, m)`` with ``p > 1``
    or ``m > 1``, the returned ``magnitude_db`` / ``phase_deg`` arrays
    keep the same ``(K, p, m)`` shape — one Bode pair per
    ``(output, input)`` channel pair.  The phase is unwrapped along the
    frequency axis (``axis=0``) independently for each channel, which is
    the standard convention for MIMO Bode plots.

    Args:
        linsys: A :class:`LinearizedSystem` (``p`` outputs, ``m`` inputs).
        omegas: 1-D array-like of angular frequencies ``ω`` (rad/s).

    Returns:
        Dictionary with keys:
            ``"omega"`` — angular frequencies (rad/s), shape ``(K,)``;
            ``"freq_hz"`` — ``ω / (2π)`` for log-Hz plotting, shape ``(K,)``;
            ``"magnitude_db"`` — ``20 log₁₀ |G(jω)|`` (dB); shape ``(K,)``
                for SISO (squeezed for backward-compatibility with phase 1),
                shape ``(K, p, m)`` for MIMO;
            ``"phase_deg"`` — phase in degrees, unwrapped along the
                frequency axis; same shape as ``magnitude_db``.
    """
    fr = frequency_response(linsys, omegas)
    mag = fr.magnitudes
    phase = fr.phases
    # Detect SISO from the underlying (K, p, m) shape; for SISO squeeze
    # to 1-D (K,) for backward-compatibility with the T-109 phase-1 API.
    is_siso = (mag.ndim == 3 and mag.shape[-1] == 1 and mag.shape[-2] == 1)
    if is_siso:
        mag = mag[..., 0, 0]
        phase = phase[..., 0, 0]
        # axis=-1 == axis=0 here; either works on a 1-D array.
        phase_deg = jnp.unwrap(phase) * (180.0 / jnp.pi)
    else:
        # MIMO: unwrap along the frequency axis (axis=0), independently per
        # (output, input) channel pair.  Using axis=-1 (the default) would
        # unwrap along the input axis and produce a meaningless mixture
        # across channels.
        phase_deg = jnp.unwrap(phase, axis=0) * (180.0 / jnp.pi)
    mag_db = 20.0 * jnp.log10(jnp.maximum(mag, 1e-300))
    return {
        "omega": fr.omegas,
        "freq_hz": fr.omegas / (2.0 * jnp.pi),
        "magnitude_db": mag_db,
        "phase_deg": phase_deg,
    }

collect_snapshots(results, signals=None)

Assemble a snapshot matrix from a jaxonomy SimulationResults.

Selected recorded signals (results.outputs) are stacked column-wise into X of shape (n_features, n_samples) where n_features is the total width of the selected signals and n_samples == len(results.time).

Parameters:

Name	Type	Description	Default
results		A SimulationResults with .time and an .outputs dict mapping signal name -> array of shape (n_samples,) or (n_samples, dim).	required
signals	Optional[Sequence[str]]	Names to include (in order). None selects every output.	None

Returns:

Name	Type	Description
A	SnapshotData	class:SnapshotData with X and time populated.

Source code in jaxonomy/library/rom/snapshots.py

def collect_snapshots(
    results,
    signals: Optional[Sequence[str]] = None,
) -> SnapshotData:
    """Assemble a snapshot matrix from a jaxonomy ``SimulationResults``.

    Selected recorded signals (``results.outputs``) are stacked column-wise
    into ``X`` of shape ``(n_features, n_samples)`` where ``n_features`` is the
    total width of the selected signals and ``n_samples == len(results.time)``.

    Args:
        results: A ``SimulationResults`` with ``.time`` and an ``.outputs`` dict
            mapping signal name -> array of shape ``(n_samples,)`` or
            ``(n_samples, dim)``.
        signals: Names to include (in order). ``None`` selects every output.

    Returns:
        A :class:`SnapshotData` with ``X`` and ``time`` populated.
    """
    if results.outputs is None:
        raise ValueError(
            "results.outputs is None; run simulate(...) with recorded_signals."
        )
    if signals is None:
        signals = list(results.outputs.keys())

    blocks = []
    for name in signals:
        if name not in results.outputs:
            raise KeyError(
                f"signal {name!r} not in recorded outputs "
                f"{list(results.outputs.keys())}"
            )
        blocks.append(_as_columns(results.outputs[name]))

    X = np.vstack(blocks)
    time = None if results.time is None else np.asarray(results.time)
    return SnapshotData(X=X, time=time)

controllability_gramian(A, B, dt=None)

Controllability Gramian :math:W_c.

Continuous time (dt is None) solves the Lyapunov equation

.. math:: A W_c + W_c A^\mathsf{T} = -B B^\mathsf{T}

Discrete time (dt given) solves the Stein equation

.. math:: A W_c A^\mathsf{T} - W_c + B B^\mathsf{T} = 0

Both require A stable (continuous: Re(eig) < 0; discrete: |eig| < 1) for a positive-semidefinite solution.

Source code in jaxonomy/library/rom/linear_mor.py

def controllability_gramian(A, B, dt=None):
    r"""Controllability Gramian :math:`W_c`.

    Continuous time (``dt is None``) solves the Lyapunov equation

    .. math:: A W_c + W_c A^\mathsf{T} = -B B^\mathsf{T}

    Discrete time (``dt`` given) solves the Stein equation

    .. math:: A W_c A^\mathsf{T} - W_c + B B^\mathsf{T} = 0

    Both require ``A`` stable (continuous: ``Re(eig) < 0``; discrete:
    ``|eig| < 1``) for a positive-semidefinite solution.
    """
    A = _np(A)
    B = _np(B)
    if A.ndim <= 1:
        A = A.reshape(1, 1)
    B = B.reshape(A.shape[0], -1)
    if dt is None:
        return sla.solve_continuous_lyapunov(A, -(B @ B.T))
    return sla.solve_discrete_lyapunov(A, B @ B.T)

deim(nonlinear_snapshots, rank=None, energy=None)

Greedy DEIM point selection (Chaturantabut & Sorensen 2010).

Takes an SVD basis U of the nonlinear-term snapshots and greedily selects m interpolation indices, then forms the oblique DEIM projector U (Pᵀ U)⁻¹ (P selects the chosen rows).

Parameters:

Name	Type	Description	Default
nonlinear_snapshots		Snapshots of the nonlinear term, shape (n_features, n_samples).	required
rank	Optional[int]	Number of DEIM modes/points m to keep.	None
energy	Optional[float]	Cumulative-energy threshold used when rank is None.	None

Returns:

Type	Description
ndarray	(indices, projector) — indices are m distinct row indices
ndarray	(np.ndarray of int), projector has shape (n_features, m).

Source code in jaxonomy/library/rom/pod.py

def deim(
    nonlinear_snapshots,
    rank: Optional[int] = None,
    energy: Optional[float] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """Greedy DEIM point selection (Chaturantabut & Sorensen 2010).

    Takes an SVD basis ``U`` of the nonlinear-term snapshots and greedily
    selects ``m`` interpolation indices, then forms the oblique DEIM projector
    ``U (Pᵀ U)⁻¹`` (``P`` selects the chosen rows).

    Args:
        nonlinear_snapshots: Snapshots of the nonlinear term, shape
            ``(n_features, n_samples)``.
        rank: Number of DEIM modes/points ``m`` to keep.
        energy: Cumulative-energy threshold used when ``rank`` is ``None``.

    Returns:
        ``(indices, projector)`` — ``indices`` are ``m`` distinct row indices
        (``np.ndarray`` of int), ``projector`` has shape ``(n_features, m)``.
    """
    F = np.asarray(nonlinear_snapshots, dtype=float)
    U_full, sigma, _ = np.linalg.svd(F, full_matrices=False)
    m = _select_rank(sigma, rank, energy)
    U = U_full[:, :m]

    indices = np.empty(m, dtype=int)
    indices[0] = int(np.argmax(np.abs(U[:, 0])))
    for j in range(1, m):
        Uj = U[:, :j]                       # (n, j)
        P = indices[:j]                     # selected rows
        # Solve (Pᵀ Uj) c = Pᵀ U[:, j] for interpolation coefficients.
        c = np.linalg.solve(Uj[P, :], U[P, j])
        residual = U[:, j] - Uj @ c
        indices[j] = int(np.argmax(np.abs(residual)))

    projector = U @ np.linalg.inv(U[indices, :])   # U (Pᵀ U)⁻¹, shape (n, m)
    return indices, projector

deim_galerkin_reduce(linear_rhs_fn, nonlinear_fn, basis, deim_result, x_ref=None, input_size=0, name=None)

Build a DEIM hyper-reduced POD-Galerkin ROM.

The full-order dynamics are split as ẋ = f_lin(t, x, u) + g(x) with a (affine-)linear part f_lin and an elementwise nonlinearity g. The linear operator is reduced offline to dense r×r / r×m operators, and the nonlinearity is approximated by DEIM so it is evaluated only at the selected points (Chaturantabut & Sorensen 2010).

Parameters:

Name	Type	Description	Default
linear_rhs_fn	Callable	Affine-linear part. Called linear_rhs_fn(t, x, u) if input_size > 0 else linear_rhs_fn(t, x); jax-traceable, returns (n_features,). Probed offline at t=0 to extract its reduced operators, so it must be affine in (x, u).	required
nonlinear_fn	Callable	Elementwise nonlinearity g; called with a (m,) vector of states at the DEIM points and returns (m,).	required
basis		POD trial basis Φ, shape (n_features, r).	required
deim_result	Tuple[ndarray, ndarray]	The (indices, projector) pair from :func:deim.	required
x_ref		Reference/offset state (default zeros).	None
input_size	int	Width of the single input port; 0 for autonomous.	0
name	Optional[str]	Optional block name.	None

Returns:

Type	Description
_DEIMGalerkinROM	A jaxonomy LeafSystem with r reduced continuous states whose
_DEIMGalerkinROM	per-step cost is independent of the full dimension n.

Source code in jaxonomy/library/rom/pod.py

def deim_galerkin_reduce(
    linear_rhs_fn: Callable,
    nonlinear_fn: Callable,
    basis,
    deim_result: Tuple[np.ndarray, np.ndarray],
    x_ref=None,
    input_size: int = 0,
    name: Optional[str] = None,
) -> _DEIMGalerkinROM:
    """Build a DEIM hyper-reduced POD-Galerkin ROM.

    The full-order dynamics are split as ``ẋ = f_lin(t, x, u) + g(x)`` with a
    (affine-)linear part ``f_lin`` and an elementwise nonlinearity ``g``. The
    linear operator is reduced offline to dense ``r×r`` / ``r×m`` operators, and
    the nonlinearity is approximated by DEIM so it is evaluated only at the
    selected points (Chaturantabut & Sorensen 2010).

    Args:
        linear_rhs_fn: Affine-linear part. Called ``linear_rhs_fn(t, x, u)`` if
            ``input_size > 0`` else ``linear_rhs_fn(t, x)``; jax-traceable,
            returns ``(n_features,)``. Probed offline at ``t=0`` to extract its
            reduced operators, so it must be affine in ``(x, u)``.
        nonlinear_fn: Elementwise nonlinearity ``g``; called with a
            ``(m,)`` vector of states at the DEIM points and returns ``(m,)``.
        basis: POD trial basis ``Φ``, shape ``(n_features, r)``.
        deim_result: The ``(indices, projector)`` pair from :func:`deim`.
        x_ref: Reference/offset state (default zeros).
        input_size: Width of the single input port; ``0`` for autonomous.
        name: Optional block name.

    Returns:
        A jaxonomy ``LeafSystem`` with ``r`` reduced continuous states whose
        per-step cost is independent of the full dimension ``n``.
    """
    Phi = np.asarray(basis, dtype=float)
    n, r = Phi.shape
    x_ref_arr = np.zeros(n) if x_ref is None else np.asarray(x_ref, dtype=float)

    indices, projector = deim_result
    indices = np.asarray(indices, dtype=int)
    projector = np.asarray(projector, dtype=float)   # (n, m)

    # Offline reduction of the affine-linear operator by probing f_lin.
    zeros_n = jnp.zeros(n)
    if input_size > 0:
        zeros_u = jnp.zeros(input_size)
        const = np.asarray(linear_rhs_fn(0.0, zeros_n, zeros_u), dtype=float)
        # A·Φ_k = f_lin(0, Φ_k, 0) − const  (linearity of f_lin in x)
        AtimesPhi = np.stack(
            [np.asarray(linear_rhs_fn(0.0, jnp.asarray(Phi[:, k]), zeros_u),
                        dtype=float) - const
             for k in range(r)],
            axis=1,
        )  # (n, r)
        # B·e_j = f_lin(0, 0, e_j) − const
        B = np.stack(
            [np.asarray(linear_rhs_fn(0.0, zeros_n,
                                      jnp.asarray(np.eye(input_size)[:, j])),
                        dtype=float) - const
             for j in range(input_size)],
            axis=1,
        )  # (n, input_size)
        Br = Phi.T @ B
    else:
        const = np.asarray(linear_rhs_fn(0.0, zeros_n), dtype=float)
        AtimesPhi = np.stack(
            [np.asarray(linear_rhs_fn(0.0, jnp.asarray(Phi[:, k])),
                        dtype=float) - const
             for k in range(r)],
            axis=1,
        )
        Br = np.zeros((r, 0))

    Ar = Phi.T @ AtimesPhi                              # (r, r)
    # Affine offset: Φᵀ (A x_ref + const) = Φᵀ f_lin(0, x_ref, 0).
    if input_size > 0:
        f_at_ref = np.asarray(
            linear_rhs_fn(0.0, jnp.asarray(x_ref_arr), zeros_u), dtype=float)
    else:
        f_at_ref = np.asarray(
            linear_rhs_fn(0.0, jnp.asarray(x_ref_arr)), dtype=float)
    b_off = Phi.T @ f_at_ref                            # (r,)

    DEIM_reduced = Phi.T @ projector                    # (r, m)
    Phi_P = Phi[indices, :]                             # (m, r)
    x_ref_P = x_ref_arr[indices]                        # (m,)

    return _DEIMGalerkinROM(
        Ar=jnp.asarray(Ar),
        b_off=jnp.asarray(b_off),
        Br=jnp.asarray(Br),
        DEIM_reduced=jnp.asarray(DEIM_reduced),
        Phi_P=jnp.asarray(Phi_P),
        x_ref_P=jnp.asarray(x_ref_P),
        Phi=jnp.asarray(Phi),
        x_ref=jnp.asarray(x_ref_arr),
        nonlinear_fn=nonlinear_fn,
        input_size=input_size,
        name=name,
    )

discretize(linsys, dt, *, method='zoh', base_context=None, input_port=None, output_port=None)

Discretize a continuous-time linear system (T-109 phase 4).

Two call patterns, dispatched on the type of linsys:

1. discretize(linsys: LinearizedSystem, dt, *, method) — the LTI-level path (shipped first as the T-109 phase-4 sub-piece). Wraps the matrix-level helpers in :mod:jaxonomy.library.state_estimators.utils.
2. discretize(system: SystemBase, dt, *, method, base_context, input_port, output_port) — the diagram-level lift (T-109 phase 4 completion). Linearizes system about base_context (via :func:linearize) then routes the result through path 1. Equivalent to discretize(linearize(system, base_context, ...), dt, method=method); provided so controller-design workflows can write ddiagram = jaxonomy.discretize(diagram, dt, base_context=ctx) in one call.

Converts dx/dt = Ax + Bu into x[k+1] = A_d x[k] + B_d u[k] while keeping C, D, and the operating point untouched (the output map is unaffected by discretization). The returned :class:LinearizedSystem carries dt so downstream consumers (e.g. :meth:LinearizedSystem.is_stable) interpret it as discrete-time.

Parameters:

Name	Type	Description	Default
linsys		Either a continuous-time :class:LinearizedSystem (dt must be None) or a :class:SystemBase / :class:Diagram. In the diagram case, base_context is required.	required
dt	float	Sampling period in seconds. Must be positive.	required
method	str	Discretization rule. "zoh" (default) — exact zero-order-hold: A_d = expm(A·dt), B_d = A⁻¹ (A_d − I) B (with a first-order Taylor fallback when A is near-singular, so integrator dynamics A = 0 work cleanly). "euler" — first-order forward-Euler: A_d = I + A·dt, B_d = B·dt. Cheap and JAX-clean but biased; use "zoh" unless you specifically need the Euler shape for hardware-in-the-loop parity.	'zoh'
base_context		Required when linsys is a SystemBase / Diagram; ignored when it's already a LinearizedSystem. The operating point about which to linearize.	None
input_port		Optional input port for :func:linearize (diagram path only). Defaults to the diagram's single input.	None
output_port		Optional output port for :func:linearize (diagram path only). Defaults to the diagram's single output.	None

Returns:

Type	Description
LinearizedSystem	A new :class:LinearizedSystem with discrete matrices and
LinearizedSystem	dt set. The output map (C, D) and
LinearizedSystem	operating_point are forwarded unchanged.

Raises:

Type	Description
ValueError	If dt <= 0, method is not "zoh" or "euler", linsys is already discrete, or the diagram path is invoked without base_context.

Notes

Differentiable through A, B, C, D, and dt via the JAX-traceable matrix exponential and linear solve. The diagram path is differentiable through whatever :func:linearize is itself differentiable through.

See also

:func:linearize — the continuous-time linearization step. :func:jaxonomy.library.state_estimators.utils.discretize_forward_zoh and :func:discretize_forward_euler — the matrix-level primitives the LTI path wraps.

Source code in jaxonomy/library/linearization_workflow.py

def discretize(
    linsys,
    dt: float,
    *,
    method: str = "zoh",
    base_context=None,
    input_port=None,
    output_port=None,
) -> LinearizedSystem:
    """Discretize a continuous-time linear system (T-109 phase 4).

    Two call patterns, dispatched on the type of ``linsys``:

    1. ``discretize(linsys: LinearizedSystem, dt, *, method)``  — the
       LTI-level path (shipped first as the T-109 phase-4 sub-piece).
       Wraps the matrix-level helpers in
       :mod:`jaxonomy.library.state_estimators.utils`.
    2. ``discretize(system: SystemBase, dt, *, method, base_context, input_port, output_port)``
       — the **diagram-level lift** (T-109 phase 4 completion).
       Linearizes ``system`` about ``base_context`` (via :func:`linearize`)
       then routes the result through path 1. Equivalent to
       ``discretize(linearize(system, base_context, ...), dt, method=method)``;
       provided so controller-design workflows can write
       ``ddiagram = jaxonomy.discretize(diagram, dt, base_context=ctx)``
       in one call.

    Converts ``dx/dt = Ax + Bu`` into ``x[k+1] = A_d x[k] + B_d u[k]``
    while keeping ``C``, ``D``, and the operating point untouched (the
    output map is unaffected by discretization). The returned
    :class:`LinearizedSystem` carries ``dt`` so downstream consumers
    (e.g. :meth:`LinearizedSystem.is_stable`) interpret it as
    discrete-time.

    Args:
        linsys: Either a continuous-time :class:`LinearizedSystem`
            (``dt`` must be ``None``) or a :class:`SystemBase` /
            :class:`Diagram`. In the diagram case, ``base_context`` is
            required.
        dt: Sampling period in seconds. Must be positive.
        method: Discretization rule.
            ``"zoh"`` (default) — exact zero-order-hold:
            ``A_d = expm(A·dt)``,
            ``B_d = A⁻¹ (A_d − I) B`` (with a first-order Taylor
            fallback when ``A`` is near-singular, so integrator dynamics
            ``A = 0`` work cleanly).
            ``"euler"`` — first-order forward-Euler:
            ``A_d = I + A·dt``, ``B_d = B·dt``. Cheap and JAX-clean but
            biased; use ``"zoh"`` unless you specifically need the
            Euler shape for hardware-in-the-loop parity.
        base_context: Required when ``linsys`` is a SystemBase /
            Diagram; ignored when it's already a LinearizedSystem.
            The operating point about which to linearize.
        input_port: Optional input port for :func:`linearize` (diagram
            path only). Defaults to the diagram's single input.
        output_port: Optional output port for :func:`linearize`
            (diagram path only). Defaults to the diagram's single
            output.

    Returns:
        A new :class:`LinearizedSystem` with discrete matrices and
        ``dt`` set. The output map (``C``, ``D``) and
        ``operating_point`` are forwarded unchanged.

    Raises:
        ValueError: If ``dt <= 0``, ``method`` is not ``"zoh"`` or
            ``"euler"``, ``linsys`` is already discrete, or the
            diagram path is invoked without ``base_context``.

    Notes:
        Differentiable through ``A``, ``B``, ``C``, ``D``, and ``dt``
        via the JAX-traceable matrix exponential and linear solve.
        The diagram path is differentiable through whatever
        :func:`linearize` is itself differentiable through.

    See also:
        :func:`linearize` — the continuous-time linearization step.
        :func:`jaxonomy.library.state_estimators.utils.discretize_forward_zoh`
            and :func:`discretize_forward_euler` — the matrix-level
            primitives the LTI path wraps.
    """
    # Diagram-level dispatch: linearize first, then route through the
    # LTI path. Detect by the absence of LinearizedSystem-y attributes
    # rather than isinstance so we accept any structurally-equivalent
    # wrapper.
    if not isinstance(linsys, LinearizedSystem):
        if base_context is None:
            raise ValueError(
                "discretize: when the first argument is a SystemBase / "
                "Diagram, base_context= is required (it's the operating "
                "point to linearize about)."
            )
        # Lazy import to avoid the framework→library cycle at module load.
        from .linear_system import linearize

        kwargs = {}
        if input_port is not None:
            kwargs["input_port"] = input_port
        if output_port is not None:
            kwargs["output_port"] = output_port
        linsys = linearize(linsys, base_context, **kwargs)

    from .state_estimators.utils import (
        discretize_forward_euler,
        discretize_forward_zoh,
    )

    if linsys.dt is not None:
        raise ValueError(
            f"discretize: linsys already carries dt={linsys.dt!r}; "
            f"discretizing an already-discrete LinearizedSystem is not "
            f"well-defined without re-continuization first."
        )
    if method not in ("zoh", "euler"):
        raise ValueError(
            f"discretize: unknown method {method!r}; expected "
            f"'zoh' or 'euler'."
        )
    # Concrete-only positivity check — under jax.grad / jit ``dt`` is a
    # traced array and we cannot coerce it to a Python float. Skip the
    # eager validation in that case; XLA will surface any value-domain
    # issues at runtime via the underlying linear-algebra ops.
    try:
        dt_concrete = float(dt)
    except (TypeError, jax.errors.ConcretizationTypeError):
        dt_concrete = None
    if dt_concrete is not None and not (dt_concrete > 0):
        raise ValueError(f"discretize: dt must be positive; got {dt!r}.")

    A = jnp.asarray(linsys.A)
    if A.ndim == 0:
        A = A.reshape((1, 1))
    elif A.ndim == 1:
        n = A.size
        A = A.reshape((n, n))
    n = A.shape[0]

    B = _ensure_2d(linsys.B, n, max(jnp.asarray(linsys.B).size // n, 1))

    if method == "zoh":
        Ad, Bd = discretize_forward_zoh(A, B, dt)
    else:  # "euler"
        Ad, Bd = discretize_forward_euler(A, B, dt)

    # Stamp dt as a Python float when concrete (so the dataclass repr is
    # readable and is_discrete() is cheap); pass through the traced array
    # unchanged when called under jit / grad.
    dt_field = dt_concrete if dt_concrete is not None else dt
    return LinearizedSystem(
        A=Ad,
        B=Bd,
        C=linsys.C,
        D=linsys.D,
        operating_point=linsys.operating_point,
        dt=dt_field,
    )

dmdc(X, Xp, U, rank=None, B_known=None)

Dynamic Mode Decomposition with control (Proctor, Brunton & Kutz 2016).

Fits x[k+1] ≈ A x[k] + B u[k] from snapshot pairs and control inputs.

Two cases are handled:

• Unknown B (default): regress on the augmented snapshot Ω = [X; U] so [A B] = Xp Ω⁺.
• Known B (pass B_known): subtract the known control effect first, A = (Xp − B U) X⁺.

Parameters:

Name	Type	Description	Default
X		State snapshots x[k], shape (n, k).	required
Xp		Advanced snapshots x[k+1], shape (n, k).	required
U		Control inputs u[k], shape (m, k).	required
rank		Optional POD rank r for the reduced operators (defaults full).	None
B_known		Optional known input matrix (n, m) for the known-B case.	None

Returns:

Type	Description
	class:DMDcResult with full A, B and reduced A_tilde, B_tilde.

Source code in jaxonomy/library/rom/dmd.py

def dmdc(X, Xp, U, rank=None, B_known=None):
    """Dynamic Mode Decomposition with control (Proctor, Brunton & Kutz 2016).

    Fits ``x[k+1] ≈ A x[k] + B u[k]`` from snapshot pairs and control inputs.

    Two cases are handled:

    - **Unknown ``B``** (default): regress on the augmented snapshot
      ``Ω = [X; U]`` so ``[A  B] = Xp Ω⁺``.
    - **Known ``B``** (pass ``B_known``): subtract the known control effect first,
      ``A = (Xp − B U) X⁺``.

    Args:
        X: State snapshots ``x[k]``, shape ``(n, k)``.
        Xp: Advanced snapshots ``x[k+1]``, shape ``(n, k)``.
        U: Control inputs ``u[k]``, shape ``(m, k)``.
        rank: Optional POD rank ``r`` for the reduced operators (defaults full).
        B_known: Optional known input matrix ``(n, m)`` for the known-``B`` case.

    Returns:
        :class:`DMDcResult` with full ``A, B`` and reduced ``A_tilde, B_tilde``.
    """
    X = np.asarray(X)
    Xp = np.asarray(Xp)
    U = np.atleast_2d(np.asarray(U))
    if U.shape[1] != X.shape[1]:
        U = U.T  # accept (k, m) as well
    n = X.shape[0]

    if B_known is not None:
        B = np.asarray(B_known)
        if B.ndim == 1:
            B = B.reshape(n, -1)
        A = (Xp - B @ U) @ np.linalg.pinv(X)
    else:
        Omega = np.vstack([X, U])
        G = Xp @ np.linalg.pinv(Omega)
        A, B = G[:, :n], G[:, n:]

    # Reduced operators via the leading POD modes of the advanced snapshots.
    Uhat, _, _ = _svd_truncate(Xp, rank)
    A_tilde = Uhat.conj().T @ A @ Uhat
    B_tilde = Uhat.conj().T @ B
    eigenvalues = np.linalg.eigvals(A_tilde)

    return DMDcResult(
        A=A,
        B=B,
        A_tilde=A_tilde,
        B_tilde=B_tilde,
        basis=Uhat,
        eigenvalues=eigenvalues,
    )

edmd(X, Xp, dictionary, U=None)

Extended DMD — approximate the Koopman operator on lifted snapshots.

Lifts the snapshot pair through dictionary and least-squares fits the lifted linear dynamics z[k+1] ≈ K z[k] (+ B u[k]).

Parameters:

Name	Type	Description	Default
X		State snapshots x[k], shape (n, k).	required
Xp		Advanced snapshots x[k+1], shape (n, k).	required
dictionary		Callable g(x) -> lifted vector (identity observables first).	required
U		Optional control inputs (m, k) for eDMDc.	None

Returns:

Type	Description
	class:EDMDResult with the Koopman operator K, the input operator
	B (or None), and the de-lift matrix C.

Source code in jaxonomy/library/rom/koopman.py

def edmd(X, Xp, dictionary, U=None):
    """Extended DMD — approximate the Koopman operator on lifted snapshots.

    Lifts the snapshot pair through ``dictionary`` and least-squares fits the
    lifted linear dynamics ``z[k+1] ≈ K z[k] (+ B u[k])``.

    Args:
        X: State snapshots ``x[k]``, shape ``(n, k)``.
        Xp: Advanced snapshots ``x[k+1]``, shape ``(n, k)``.
        dictionary: Callable ``g(x) -> lifted vector`` (identity observables first).
        U: Optional control inputs ``(m, k)`` for eDMDc.

    Returns:
        :class:`EDMDResult` with the Koopman operator ``K``, the input operator
        ``B`` (or ``None``), and the de-lift matrix ``C``.
    """
    X = np.asarray(X, dtype=float)
    Xp = np.asarray(Xp, dtype=float)

    Z1 = _lift_columns(X, dictionary)
    Z2 = _lift_columns(Xp, dictionary)

    if U is None:
        # K solves Z2 ≈ K Z1  (least squares).
        K = Z2 @ np.linalg.pinv(Z1)
        B = None
    else:
        U = np.atleast_2d(np.asarray(U, dtype=float))
        if U.shape[1] != Z1.shape[1]:
            U = U.T
        n_state = Z1.shape[0]
        Omega = np.vstack([Z1, U])
        G = Z2 @ np.linalg.pinv(Omega)
        K, B = G[:, :n_state], G[:, n_state:]

    # De-lift matrix C: physical state ≈ C · lifted.  Solved by least squares,
    # so it reduces to a row-selection when the identity observables are present.
    C = X @ np.linalg.pinv(Z1)

    return EDMDResult(K=K, B=B, C=C, dictionary=dictionary)

era(markov, n_inputs, n_outputs, num_rows=None, num_cols=None, rank=None)

Eigensystem Realization Algorithm (Juang & Pappa 1985).

Builds a minimal discrete-time state-space realization (A, B, C, D) from a sequence of impulse-response Markov parameters Y_0 = D, Y_1 = C B, Y_2 = C A B ...

Parameters:

Name	Type	Description	Default
markov		Markov parameters. Either an array of shape (L+1, n_outputs, n_inputs) or a length-L+1 sequence of such blocks; SISO impulse responses may be passed as a 1-D array.	required
n_inputs		Number of inputs m.	required
n_outputs		Number of outputs p.	required
num_rows		Block rows α of the Hankel matrix (default ~half the data).	None
num_cols		Block cols β of the Hankel matrix (default ~half the data).	None
rank		Optional model order r (SVD truncation of the Hankel matrix).	None

Returns:

Type	Description
	class:ERAResult with the realized (A, B, C, D) and Hankel
	singular values.

Source code in jaxonomy/library/rom/dmd.py

def era(markov, n_inputs, n_outputs, num_rows=None, num_cols=None, rank=None):
    """Eigensystem Realization Algorithm (Juang & Pappa 1985).

    Builds a minimal discrete-time state-space realization ``(A, B, C, D)`` from a
    sequence of impulse-response Markov parameters
    ``Y_0 = D``, ``Y_1 = C B``, ``Y_2 = C A B`` ...

    Args:
        markov: Markov parameters. Either an array of shape
            ``(L+1, n_outputs, n_inputs)`` or a length-``L+1`` sequence of such
            blocks; SISO impulse responses may be passed as a 1-D array.
        n_inputs: Number of inputs ``m``.
        n_outputs: Number of outputs ``p``.
        num_rows: Block rows ``α`` of the Hankel matrix (default ~half the data).
        num_cols: Block cols ``β`` of the Hankel matrix (default ~half the data).
        rank: Optional model order ``r`` (SVD truncation of the Hankel matrix).

    Returns:
        :class:`ERAResult` with the realized ``(A, B, C, D)`` and Hankel
        singular values.
    """
    Y = np.asarray(markov, dtype=float)
    L = Y.shape[0] - 1  # number of pulse-response blocks after D
    Y = Y.reshape(L + 1, n_outputs, n_inputs)

    D = Y[0]
    H = Y[1:]  # pulse response Y_1 .. Y_L (used to build the Hankel matrix)

    if num_rows is None:
        num_rows = L // 2
    if num_cols is None:
        num_cols = L - num_rows
    alpha, beta = int(num_rows), int(num_cols)
    if alpha + beta > L:
        raise ValueError(
            f"era: need num_rows + num_cols <= {L} Markov blocks, "
            f"got {alpha} + {beta}."
        )

    def _hankel(shift):
        blocks = [
            [H[i + j + shift] for j in range(beta)] for i in range(alpha)
        ]
        return np.block(blocks)

    H0 = _hankel(0)  # blocks Y_{i+j+1}
    H1 = _hankel(1)  # blocks Y_{i+j+2}

    Ur, s, Vh = np.linalg.svd(H0, full_matrices=False)
    if rank is not None:
        r = min(int(rank), s.shape[0])
    else:
        tol = max(H0.shape) * np.finfo(float).eps * (s[0] if s.size else 0.0)
        r = int(np.sum(s > tol))
        r = max(r, 1)
    Ur, s, Vh = Ur[:, :r], s[:r], Vh[:r, :]
    Vr = Vh.conj().T

    s_sqrt = np.sqrt(s)
    s_inv_sqrt = 1.0 / s_sqrt
    Obs = Ur * s_sqrt          # observability factor  R Σ^{1/2}
    Cc = (s_sqrt[:, None]) * Vr.conj().T  # controllability factor Σ^{1/2} S*

    A = (s_inv_sqrt[:, None]) * (Ur.conj().T @ H1 @ Vr) * (s_inv_sqrt[None, :])
    B = Cc[:, :n_inputs]
    C = Obs[:n_outputs, :]

    return ERAResult(A=A, B=B, C=C, D=D, singular_values=s)

estimate_frequency_response(diagram, ctx, t_span, input_port, output_port, freq_grid, *, options=None, recorded_signals_extra=None, window=True, coherence_floor=1e-12, n_segments=8, segment_overlap=0.5)

Empirically estimate the SISO transfer function of diagram.

Drives diagram with whatever signal is already wired to input_port (typically a :class:jaxonomy.library.Chirp, :class:PRBS, or :class:BandLimitedNoise source connected upstream of input_port) and records the input/output trajectories. The empirical transfer function is computed as the cross-spectral ratio G(f) = Sxy(f) / Sxx(f) where Sxx and Sxy are the (Hann- windowed) auto- and cross-spectral densities of input/output. Results are interpolated onto the user-supplied freq_grid (Hz).

This is the practical alternative to analytic :func:linearize / :func:frequency_response when:

• the system contains hard nonlinearities (lookup tables, saturation, contact dynamics) that make symbolic linearization fragile, or
• you want an empirical sanity-check against the linearized model.

Parameters:

Name	Type	Description	Default
diagram		A built diagram (typically with a chirp or PRBS source wired to the block-under-test's input port).	required
ctx		Initial simulation context.	required
t_span		(t0, tf) simulation horizon. Make this comfortably longer than the slowest period of interest in freq_grid.	required
input_port		OutputPort whose recorded trajectory provides the excitation samples u(t). This is typically the upstream source's output port (the same signal that drives the block-under-test).	required
output_port		OutputPort whose recorded trajectory provides the measured response y(t).	required
freq_grid		1-D array of frequencies (Hz) at which the empirical response should be evaluated. Frequencies outside the simulation's resolved band [1/T, fs/2] are clamped — the caller should keep freq_grid inside that band.	required
options		Optional :class:SimulatorOptions. recorded_signals is overridden internally; everything else (rtol, atol, solver, etc.) is honoured.	None
recorded_signals_extra		Optional dict[str, OutputPort] of additional signals to record alongside the input/output (useful for debugging / plotting). Not used by the estimator itself.	None
window	bool	If True (default) apply a Hann window before the FFT to suppress spectral leakage. If False (rectangular) the transfer-function ratio is more sensitive to leakage but faithful to the raw FFT.	True
coherence_floor	float	Minimum |U(f)|² in the auto-spectrum below which the ratio is set to zero — guards against division by zero at frequencies the excitation never visited.	1e-12
n_segments	int	Number of overlapping segments to average (Welch's method). More segments → less variance, lower frequency resolution. n_segments=1 falls back to a single-window FFT estimate. Default 8 is a reasonable trade-off for most chirp/PRBS excitations.	8
segment_overlap	float	Fractional overlap between consecutive segments (Welch's method), in [0, 1). Default 0.5 (50%).	0.5

Returns:

Type	Description
FrequencyResponse	class:FrequencyResponse with omegas = 2π·freq_grid, complex
FrequencyResponse	response of shape (K, 1, 1), and corresponding
FrequencyResponse	magnitudes and phases. Drop-in compatible with
FrequencyResponse	func:bode_data.

Notes

• The implementation is intentionally pure-NumPy on the post-simulation arrays; it does not need to be JAX-traceable (callers can JIT downstream code that consumes the returned response array).
• For best results pick an excitation that covers the band of interest densely: a linear :class:Chirp from f0 ≪ freq_min to f1 ≳ freq_max over a horizon of several seconds, or a :class:PRBS with sample time ≪ 1/(2·freq_max).
• The returned response is a NumPy complex array (consumers calling :func:bode_data will see jnp.asarray promotion); this is fine because :class:FrequencyResponse fields are typed Any.

Source code in jaxonomy/library/linearization_workflow.py

def estimate_frequency_response(
    diagram,
    ctx,
    t_span,
    input_port,
    output_port,
    freq_grid,
    *,
    options=None,
    recorded_signals_extra=None,
    window: bool = True,
    coherence_floor: float = 1e-12,
    n_segments: int = 8,
    segment_overlap: float = 0.5,
) -> FrequencyResponse:
    """Empirically estimate the SISO transfer function of ``diagram``.

    Drives ``diagram`` with whatever signal is already wired to ``input_port``
    (typically a :class:`jaxonomy.library.Chirp`, :class:`PRBS`, or
    :class:`BandLimitedNoise` source connected upstream of ``input_port``)
    and records the input/output trajectories.  The empirical transfer
    function is computed as the cross-spectral ratio
    ``G(f) = Sxy(f) / Sxx(f)`` where ``Sxx`` and ``Sxy`` are the (Hann-
    windowed) auto- and cross-spectral densities of input/output.  Results
    are interpolated onto the user-supplied ``freq_grid`` (Hz).

    This is the practical alternative to analytic :func:`linearize` /
    :func:`frequency_response` when:

    * the system contains hard nonlinearities (lookup tables, saturation,
      contact dynamics) that make symbolic linearization fragile, or
    * you want an empirical sanity-check against the linearized model.

    Args:
        diagram: A built diagram (typically with a chirp or PRBS source
            wired to the block-under-test's input port).
        ctx: Initial simulation context.
        t_span: ``(t0, tf)`` simulation horizon.  Make this comfortably
            longer than the slowest period of interest in ``freq_grid``.
        input_port: ``OutputPort`` whose recorded trajectory provides the
            excitation samples ``u(t)``.  This is typically the upstream
            source's output port (the same signal that drives the
            block-under-test).
        output_port: ``OutputPort`` whose recorded trajectory provides the
            measured response ``y(t)``.
        freq_grid: 1-D array of frequencies (Hz) at which the empirical
            response should be evaluated.  Frequencies outside the
            simulation's resolved band ``[1/T, fs/2]`` are clamped — the
            caller should keep ``freq_grid`` inside that band.
        options: Optional :class:`SimulatorOptions`.  ``recorded_signals``
            is overridden internally; everything else (rtol, atol, solver,
            etc.) is honoured.
        recorded_signals_extra: Optional ``dict[str, OutputPort]`` of
            additional signals to record alongside the input/output
            (useful for debugging / plotting).  Not used by the estimator
            itself.
        window: If True (default) apply a Hann window before the FFT to
            suppress spectral leakage.  If False (rectangular) the
            transfer-function ratio is more sensitive to leakage but
            faithful to the raw FFT.
        coherence_floor: Minimum ``|U(f)|²`` in the auto-spectrum below
            which the ratio is set to zero — guards against division by
            zero at frequencies the excitation never visited.
        n_segments: Number of overlapping segments to average (Welch's
            method).  More segments → less variance, lower frequency
            resolution.  ``n_segments=1`` falls back to a single-window
            FFT estimate.  Default 8 is a reasonable trade-off for most
            chirp/PRBS excitations.
        segment_overlap: Fractional overlap between consecutive segments
            (Welch's method), in ``[0, 1)``.  Default 0.5 (50%).

    Returns:
        :class:`FrequencyResponse` with ``omegas = 2π·freq_grid``, complex
        ``response`` of shape ``(K, 1, 1)``, and corresponding
        ``magnitudes`` and ``phases``.  Drop-in compatible with
        :func:`bode_data`.

    Notes:
        * The implementation is intentionally pure-NumPy on the
          post-simulation arrays; it does not need to be JAX-traceable
          (callers can JIT downstream code that consumes the returned
          ``response`` array).
        * For best results pick an excitation that covers the band of
          interest densely: a linear :class:`Chirp` from ``f0 ≪ freq_min``
          to ``f1 ≳ freq_max`` over a horizon of several seconds, or a
          :class:`PRBS` with sample time ``≪ 1/(2·freq_max)``.
        * The returned ``response`` is a NumPy complex array (consumers
          calling :func:`bode_data` will see ``jnp.asarray`` promotion);
          this is fine because :class:`FrequencyResponse` fields are typed
          ``Any``.
    """
    # Lazy import to avoid a top-level circular dependency between
    # ``jaxonomy.library`` and ``jaxonomy.simulation``.
    from jaxonomy.simulation import simulate
    from jaxonomy.simulation.types import SimulatorOptions

    recorded = {"_fre_in": input_port, "_fre_out": output_port}
    if recorded_signals_extra:
        # Rename collisions are caller-error; just merge.
        recorded.update(recorded_signals_extra)

    # Auto-bump ``buffer_length`` when the user requests a fine
    # ``max_major_step_length`` but supplies a fixed buffer too small
    # for the recording — otherwise the recorded time series gets
    # truncated to the buffer tail and the FFT sees almost no data.
    # ``buffer_length=None`` (post-T-002b auto-size default) is honoured
    # by ``_check_options`` downstream, so no bump is needed in that case.
    if options is not None and options.max_major_step_length is not None:
        t0, tf = float(t_span[0]), float(t_span[1])
        needed = int(np.ceil((tf - t0) / float(options.max_major_step_length))) + 8
        current = options.buffer_length
        if current is not None and current < needed:
            # Replace the dataclass with an updated copy so we don't
            # mutate the caller's options object.
            import dataclasses as _dc
            options = _dc.replace(options, buffer_length=needed)

    results = simulate(
        diagram,
        ctx,
        t_span=t_span,
        options=options,
        recorded_signals=recorded,
    )

    t = np.asarray(results.time)
    u_raw = np.asarray(results.outputs["_fre_in"])
    y_raw = np.asarray(results.outputs["_fre_out"])

    # Coerce vector-valued single-channel outputs to scalar-per-sample.
    if u_raw.ndim > 1:
        u_raw = u_raw.reshape(u_raw.shape[0], -1)
        if u_raw.shape[1] != 1:
            raise ValueError(
                "estimate_frequency_response(): input_port must be SISO "
                f"(got width {u_raw.shape[1]}).  Use Demux to select a "
                "single channel."
            )
        u_raw = u_raw[:, 0]
    if y_raw.ndim > 1:
        y_raw = y_raw.reshape(y_raw.shape[0], -1)
        if y_raw.shape[1] != 1:
            raise ValueError(
                "estimate_frequency_response(): output_port must be SISO "
                f"(got width {y_raw.shape[1]}).  Use Demux to select a "
                "single channel."
            )
        y_raw = y_raw[:, 0]

    if t.size < 4:
        raise ValueError(
            "estimate_frequency_response(): need at least 4 samples; got "
            f"{t.size}.  Lengthen ``t_span`` or relax ``max_major_step_length``."
        )

    # Resample to a uniform grid so np.fft.rfft is correct.
    n = int(t.size)
    t_u, u, dt = _resample_to_uniform(t, u_raw, n_samples=n)
    _, y, _ = _resample_to_uniform(t, y_raw, n_samples=n)

    # Linear-detrend (a DC offset or a slow drift — e.g. from an integrator
    # reacting to the small DC component of a chirp — would dominate the
    # low-frequency bins and bleed into higher bins via spectral leakage).
    def _linear_detrend(arr):
        idx = np.arange(arr.size, dtype=np.float64)
        # Least-squares fit y = a*idx + b in closed form.
        n_arr = arr.size
        sum_x = idx.sum()
        sum_y = arr.sum()
        sum_xx = (idx * idx).sum()
        sum_xy = (idx * arr).sum()
        denom = n_arr * sum_xx - sum_x * sum_x
        if denom == 0.0:
            return arr - arr.mean()
        slope = (n_arr * sum_xy - sum_x * sum_y) / denom
        intercept = (sum_y - slope * sum_x) / n_arr
        return arr - (slope * idx + intercept)

    u = _linear_detrend(u)
    y = _linear_detrend(y)

    # ---- Welch-style segmented average of cross-/auto-spectra. ----
    if n_segments < 1:
        raise ValueError("n_segments must be >= 1")
    if not (0.0 <= segment_overlap < 1.0):
        raise ValueError("segment_overlap must be in [0, 1)")

    if n_segments == 1:
        seg_len = n
        starts = [0]
    else:
        # Choose a segment length so ``n_segments`` overlapping windows
        # fit inside ``n``.  ``hop = seg_len * (1 - overlap)`` →
        # ``n >= seg_len + (n_segments - 1) * hop``  →
        # ``seg_len <= n / (1 + (n_segments - 1)*(1 - overlap))``.
        denom = 1.0 + (n_segments - 1) * (1.0 - segment_overlap)
        seg_len = max(int(n / denom), 8)
        if seg_len >= n:
            seg_len = n
            starts = [0]
        else:
            hop = max(int(seg_len * (1.0 - segment_overlap)), 1)
            starts = [k * hop for k in range(n_segments) if k * hop + seg_len <= n]
            if not starts:
                starts = [0]
                seg_len = n

    if window:
        w = _hann_window(seg_len)
    else:
        w = np.ones(seg_len)

    n_freqs = seg_len // 2 + 1
    Sxx_acc = np.zeros(n_freqs, dtype=np.float64)
    Sxy_acc = np.zeros(n_freqs, dtype=np.complex128)

    for s in starts:
        u_seg = u[s : s + seg_len] * w
        y_seg = y[s : s + seg_len] * w
        U = np.fft.rfft(u_seg)
        Y = np.fft.rfft(y_seg)
        Sxx_acc += (U.conj() * U).real
        Sxy_acc += U.conj() * Y

    Sxx_acc /= len(starts)
    Sxy_acc /= len(starts)
    freqs = np.fft.rfftfreq(seg_len, d=dt)

    # ---- Smoothed transfer-function estimate. ----
    # The bin-by-bin ratio ``Sxy/Sxx`` is high-variance for excitations
    # like a chirp (which dwell only briefly at any instantaneous freq).
    # The standard fix is to compute the ratio of *smoothed* spectra:
    # smooth Sxx and Sxy with a small box (or Hann) kernel along
    # frequency, then take the ratio.  This is the H1 estimator from
    # the system-ID literature.
    freq_grid = np.asarray(freq_grid, dtype=np.float64).ravel()

    # Choose smoothing kernel half-width.  We want the smoothing window
    # to be much narrower than the spacing between user-requested target
    # frequencies (so we don't mush them together).  A safe heuristic:
    # 1/4 of the smallest target-frequency gap, or 3 bins minimum.
    df = freqs[1] - freqs[0] if freqs.size > 1 else 1.0
    if freq_grid.size >= 2:
        min_gap = float(np.min(np.diff(np.sort(freq_grid))))
        kernel_half = max(3, int(min_gap / (4.0 * df)))
    else:
        kernel_half = 3
    # Cap at 1% of the spectrum to avoid degenerate cases.
    kernel_half = min(kernel_half, max(3, Sxx_acc.size // 100))
    kernel_len = 2 * kernel_half + 1
    kernel = np.ones(kernel_len) / kernel_len  # box smoother

    # Same convolution length via 'same' mode.
    Sxx_smooth = np.convolve(Sxx_acc, kernel, mode="same")
    # Smooth real and imaginary of Sxy independently to keep it complex.
    Sxy_smooth = (
        np.convolve(Sxy_acc.real, kernel, mode="same")
        + 1j * np.convolve(Sxy_acc.imag, kernel, mode="same")
    )

    safe = Sxx_smooth > coherence_floor
    G_emp = np.zeros_like(Sxy_smooth, dtype=np.complex128)
    G_emp[safe] = Sxy_smooth[safe] / Sxx_smooth[safe]

    if not np.any(safe):
        G_grid = np.zeros(freq_grid.size, dtype=np.complex128)
    else:
        # Linear interp of magnitude and unwrapped phase — std Bode plot.
        mag = np.abs(G_emp)
        phase = np.unwrap(np.angle(G_emp))
        mag_grid = np.interp(freq_grid, freqs, mag)
        phase_grid = np.interp(freq_grid, freqs, phase)
        G_grid = mag_grid * np.exp(1j * phase_grid)

    # Reshape to ``(K, 1, 1)`` to match :func:`frequency_response`.
    K = freq_grid.size
    response = G_grid.reshape((K, 1, 1))
    magnitudes = np.abs(response)
    phases = np.angle(response)
    omegas = 2.0 * np.pi * freq_grid

    # Promote via npa so consumers that JIT downstream see jax arrays.
    return FrequencyResponse(
        omegas=npa.asarray(omegas),
        response=npa.asarray(response),
        magnitudes=npa.asarray(magnitudes),
        phases=npa.asarray(phases),
    )

findop(system, base_context, *, initial_guess=None, input_port=None, tol=1e-08, max_iter=50, damping=1e-10, axis_mask=None, residual_fn=None, residual_scaling=None, scaling_eps=1e-08)

Find a continuous-state operating point x* such that ẋ(x*, u₀) ≈ 0.

Performs damped Newton iteration on the residual r(x) = ẋ(x, u₀) where u₀ is read from base_context (and held fixed for the duration of the search). The Jacobian is computed with jax.jacrev and the linear update is solved with jnp.linalg.solve plus a small Levenberg regularisation so singular Jacobians degrade to a least-squares step instead of NaN.

Parameters:

Name	Type	Description	Default
system		The system whose equilibrium is sought.	required
base_context		A context that supplies the initial state, parameter values, and (via input_port.eval) the held-fixed input.	required
initial_guess		Optional initial state. Defaults to base_context.continuous_state.	None
input_port		Input port to read u₀ from. Defaults to system.input_ports[0] (errors if system has multiple inputs and none is specified).	None
tol	float	Stop when max(|scaled residual|) < tol over the solved-for components (equals max(|ẋ|) when residual_scaling is off).	1e-08
max_iter	int	Hard cap on Newton iterations.	50
damping	float	Tikhonov damping added to JᵀJ for ill-conditioned solves.	1e-10
axis_mask		Optional selector for which state components the Newton iteration drives to zero. Either a boolean array (length = number of flat state components, True = solve this component) or a sequence of integer indices. Components not selected are held at initial_guess and excluded from both the residual and the unknowns — use this to trim systems with passive states whose equilibrium derivative is intrinsically nonzero (a cornering vehicle's heading ψ̇ = r ≠ 0, a free integrator), which would otherwise dominate the full-state Newton step and prevent convergence. None (default) solves the full state.	None
residual_fn		Optional (x) -> residual_vector overriding the default ẋ(x, u₀) — e.g. to add a custom equilibrium condition or drop terms. Receives the full state x; its output is masked / scaled like the default residual.	None
residual_scaling		Optional per-component residual weighting to put disparate units on a common footing (cf. MATLAB findop's XScaling / YScaling) — e.g. a chassis residual in m/s² ~10 alongside a wheel residual in rad/s² ~100. One of: None (no scaling), "auto" (component i scaled by 1/max(|rᵢ(x₀)|, scaling_eps) so each starts at order 1), or an explicit array (full-state length, or solved-subset length when axis_mask is given). Applied to the Newton step and the convergence test; residual_norm is then reported in scaled units.	None
scaling_eps	float	Floor for the "auto" scaling denominator.	1e-08

Returns:

Type	Description
OperatingPoint	class:OperatingPoint carrying the equilibrium state and convergence
OperatingPoint	metadata. x always has the shape of initial_guess (held
OperatingPoint	components carry their initial values when axis_mask is used).

Notes

The returned x is a JAX array, so the residual function used here is differentiable: jax.grad(lambda x0: jnp.sum(residual(x0)**2)) works. Composing :func:findop itself under jax.grad requires an implicit-differentiation wrapper which is deferred to a follow-up.

Robust fallback. Newton operating-point search can stall on stiff, strongly-coupled, or badly-scaled systems even with axis_mask and residual_scaling. The most robust equilibrium finder is simply to integrate to steady state: simulate the system from a reasonable initial condition over a horizon long relative to its slowest mode and take the final state (optionally asserting max(|ẋ|) is small there). Use that when findop reports converged=False.

Source code in jaxonomy/library/linearization_workflow.py

def findop(
    system,
    base_context,
    *,
    initial_guess=None,
    input_port=None,
    tol: float = 1e-8,
    max_iter: int = 50,
    damping: float = 1e-10,
    axis_mask=None,
    residual_fn=None,
    residual_scaling=None,
    scaling_eps: float = 1e-8,
) -> OperatingPoint:
    """Find a continuous-state operating point ``x*`` such that ``ẋ(x*, u₀) ≈ 0``.

    Performs damped Newton iteration on the residual ``r(x) = ẋ(x, u₀)`` where
    ``u₀`` is read from ``base_context`` (and held fixed for the duration of
    the search).  The Jacobian is computed with ``jax.jacrev`` and the linear
    update is solved with ``jnp.linalg.solve`` plus a small Levenberg
    regularisation so singular Jacobians degrade to a least-squares step
    instead of NaN.

    Args:
        system: The system whose equilibrium is sought.
        base_context: A context that supplies the initial state, parameter
            values, and (via ``input_port.eval``) the held-fixed input.
        initial_guess: Optional initial state.  Defaults to
            ``base_context.continuous_state``.
        input_port: Input port to read ``u₀`` from.  Defaults to
            ``system.input_ports[0]`` (errors if ``system`` has multiple
            inputs and none is specified).
        tol: Stop when ``max(|scaled residual|) < tol`` over the solved-for
            components (equals ``max(|ẋ|)`` when ``residual_scaling`` is off).
        max_iter: Hard cap on Newton iterations.
        damping: Tikhonov damping added to ``JᵀJ`` for ill-conditioned solves.
        axis_mask: Optional selector for which state components the Newton
            iteration drives to zero.  Either a boolean array (length =
            number of flat state components, ``True`` = solve this component)
            or a sequence of integer indices.  Components *not* selected are
            held at ``initial_guess`` and excluded from both the residual and
            the unknowns — use this to trim systems with *passive* states whose
            equilibrium derivative is intrinsically nonzero (a cornering
            vehicle's heading ``ψ̇ = r ≠ 0``, a free integrator), which would
            otherwise dominate the full-state Newton step and prevent
            convergence.  ``None`` (default) solves the full state.
        residual_fn: Optional ``(x) -> residual_vector`` overriding the default
            ``ẋ(x, u₀)`` — e.g. to add a custom equilibrium condition or drop
            terms.  Receives the full state ``x``; its output is masked /
            scaled like the default residual.
        residual_scaling: Optional per-component residual weighting to put
            disparate units on a common footing (cf. MATLAB ``findop``'s
            ``XScaling`` / ``YScaling``) — e.g. a chassis residual in ``m/s²``
            ~10 alongside a wheel residual in ``rad/s²`` ~100.  One of:
            ``None`` (no scaling), ``"auto"`` (component ``i`` scaled by
            ``1/max(|rᵢ(x₀)|, scaling_eps)`` so each starts at order 1), or an
            explicit array (full-state length, or solved-subset length when
            ``axis_mask`` is given).  Applied to the Newton step *and* the
            convergence test; ``residual_norm`` is then reported in scaled
            units.
        scaling_eps: Floor for the ``"auto"`` scaling denominator.

    Returns:
        :class:`OperatingPoint` carrying the equilibrium state and convergence
        metadata.  ``x`` always has the shape of ``initial_guess`` (held
        components carry their initial values when ``axis_mask`` is used).

    Notes:
        The returned ``x`` is a JAX array, so the residual function used here
        is differentiable: ``jax.grad(lambda x0: jnp.sum(residual(x0)**2))``
        works.  Composing :func:`findop` itself under ``jax.grad`` requires an
        implicit-differentiation wrapper which is deferred to a follow-up.

        **Robust fallback.**  Newton operating-point search can stall on stiff,
        strongly-coupled, or badly-scaled systems even with ``axis_mask`` and
        ``residual_scaling``.  The most robust equilibrium finder is simply to
        *integrate to steady state*: ``simulate`` the system from a reasonable
        initial condition over a horizon long relative to its slowest mode and
        take the final state (optionally asserting ``max(|ẋ|)`` is small
        there).  Use that when ``findop`` reports ``converged=False``.
    """
    if input_port is None:
        if len(system.input_ports) != 1:
            raise ValueError(
                "findop(): system has multiple input ports — pass "
                "input_port=... explicitly."
            )
        input_port = system.input_ports[0]

    default_residual, u0 = _residual_fn(system, base_context, input_port)
    residual = residual_fn if residual_fn is not None else default_residual

    if initial_guess is None:
        initial_guess = base_context.continuous_state
    x = jnp.asarray(initial_guess)
    x_shape = x.shape
    x0_flat = x.ravel()
    n_total = x0_flat.shape[0]

    # Resolve axis_mask -> integer indices of the components the Newton
    # iteration solves for.  ``None`` -> all components (legacy full solve).
    if axis_mask is None:
        free_idx = np.arange(n_total)
    else:
        m = np.asarray(axis_mask)
        if m.dtype == bool:
            if m.ravel().shape[0] != n_total:
                raise ValueError(
                    f"findop(): boolean axis_mask has {m.size} entries but the "
                    f"state has {n_total} components."
                )
            free_idx = np.where(m.ravel())[0]
        else:
            free_idx = m.ravel().astype(int)
        if free_idx.size == 0:
            raise ValueError("findop(): axis_mask selects no state components.")

    free_idx_j = jnp.asarray(free_idx)

    def _full_from_z(z):
        return x0_flat.at[free_idx_j].set(z).reshape(x_shape)

    def _masked_residual(z):
        r = jnp.atleast_1d(jnp.ravel(residual(_full_from_z(z))))
        return r[free_idx_j]

    z = x0_flat[free_idx_j]

    # Per-component residual scaling (applied to the solved-for subset).
    r0 = _masked_residual(z)
    if residual_scaling is None:
        scale = jnp.ones_like(r0)
    elif isinstance(residual_scaling, str) and residual_scaling == "auto":
        scale = 1.0 / jnp.maximum(jnp.abs(r0), scaling_eps)
    else:
        scale = jnp.asarray(residual_scaling).ravel()
        if scale.shape[0] == n_total and free_idx.size != n_total:
            scale = scale[free_idx_j]
        if scale.shape[0] != r0.shape[0]:
            raise ValueError(
                f"findop(): residual_scaling has {scale.shape[0]} entries but "
                f"the solved residual has {r0.shape[0]} components."
            )

    def _scaled_residual(z):
        return _masked_residual(z) * scale

    # Restore the previous fixed value (if any) at the end so we don't
    # accidentally mutate caller state.
    restore_fixed_val = bool(getattr(input_port, "is_fixed", False))

    jac_fn = jax.jit(jax.jacrev(_scaled_residual))
    res_fn = jax.jit(_scaled_residual)

    converged = False
    iterations = 0
    res_val = res_fn(z)
    res_norm = float(jnp.max(jnp.abs(res_val)))

    for k in range(max_iter):
        iterations = k
        if res_norm < tol:
            converged = True
            break
        J = jnp.atleast_2d(jac_fn(z))
        r = jnp.atleast_1d(res_val)
        # Damped normal-equation step: (JᵀJ + λI) Δ = Jᵀ r
        n = J.shape[1]
        JtJ = J.T @ J + damping * jnp.eye(n, dtype=J.dtype)
        rhs = J.T @ r
        delta = npa.linalg.solve(JtJ, rhs)
        z = z - delta.reshape(z.shape)
        res_val = res_fn(z)
        res_norm = float(jnp.max(jnp.abs(res_val)))
    else:
        # Loop exhausted without break: record the post-loop iteration count.
        iterations = max_iter
        if res_norm < tol:
            converged = True

    x = _full_from_z(z)

    if restore_fixed_val:
        input_port.fix_value(u0)

    return OperatingPoint(
        x=x,
        u=u0,
        residual_norm=res_norm,
        converged=converged,
        iterations=iterations,
    )

fit_gp(X, y, kernel='rbf', length_scale=1.0, signal_var=1.0, noise=1e-08, optimize=False, n_restarts=0, lr=0.05, n_steps=200, matern_nu=2.5)

Fit a Gaussian-process (kriging) surrogate.

Parameters:

Name	Type	Description	Default
X		training inputs, shape (n,) or (n, d).	required
y		training targets, shape (n,).	required
kernel		"rbf" / "squared_exponential" or "matern" / "matern32" / "matern52".	'rbf'
length_scale, signal_var, noise		kernel hyperparameters (initial values when optimize=True).	required
optimize		if True, maximize the marginal log-likelihood over (length_scale, signal_var, noise) by gradient ascent in log-space (Rasmussen & Williams 2006, Eq. 5.9).	False

Returns:

Name	Type	Description
A		class:GPModel.

Source code in jaxonomy/library/rom/surrogates.py

def fit_gp(X, y, kernel="rbf", length_scale=1.0, signal_var=1.0, noise=1e-8,
           optimize=False, n_restarts=0, lr=0.05, n_steps=200, matern_nu=2.5):
    """Fit a Gaussian-process (kriging) surrogate.

    Args:
        X: training inputs, shape ``(n,)`` or ``(n, d)``.
        y: training targets, shape ``(n,)``.
        kernel: ``"rbf"`` / ``"squared_exponential"`` or ``"matern"`` /
            ``"matern32"`` / ``"matern52"``.
        length_scale, signal_var, noise: kernel hyperparameters (initial values
            when ``optimize=True``).
        optimize: if True, maximize the marginal log-likelihood over
            ``(length_scale, signal_var, noise)`` by gradient ascent in
            log-space (Rasmussen & Williams 2006, Eq. 5.9).

    Returns:
        A :class:`GPModel`.
    """
    X = _as2d(X)
    y = jnp.asarray(y, dtype=jnp.float64).reshape(-1)

    ls = float(length_scale)
    sv = float(signal_var)
    nz = float(noise)

    if optimize:
        # Optimize in log-space to keep hyperparameters positive.
        theta0 = jnp.log(jnp.array([ls, sv, nz], dtype=jnp.float64))

        def neg_lml(theta):
            ell, s, z = jnp.exp(theta)
            alpha, L = _gp_solve(X, y, kernel, ell, s, z, matern_nu)
            n = X.shape[0]
            val = (0.5 * jnp.dot(y, alpha)
                   + jnp.sum(jnp.log(jnp.diag(L)))
                   + 0.5 * n * math.log(2.0 * math.pi))
            return val

        grad_fn = jax.jit(jax.grad(neg_lml))
        theta = theta0
        for _ in range(int(n_steps)):
            g = grad_fn(theta)
            theta = theta - lr * g
        ls, sv, nz = (float(v) for v in jnp.exp(theta))

    alpha, L = _gp_solve(X, y, kernel, ls, sv, nz, matern_nu)
    return GPModel(X, y, alpha, L, kernel, ls, sv, nz, matern_nu)

fit_lookup_table_1d(xp, x_data, y_data, *, interpolation='linear', extrapolation='clip', weights=None, smoothness=0.0, name=None, **block_kwargs)

Fit a 1-D lookup table to data and return a LookupTable1d block.

Parameters:

Name	Type	Description	Default
xp		Fixed grid of breakpoints (1-D, strictly increasing).	required
x_data		Measured input cloud, shape (K,).	required
y_data		Measured output cloud, shape (K,).	required
interpolation	str	Interpolation rule for the runtime block ("linear" / "pchip" / "nearest" / "flat"). The fit itself is always linear-LS — see the module docstring of :mod:jaxonomy.library.lookup_table for why.	'linear'
extrapolation	str	Out-of-range policy for the runtime block; see :class:jaxonomy.library.LookupTable1d.	'clip'
weights		Optional per-sample weights for weighted least- squares. None = OLS.	None
smoothness	float	Non-negative discrete first-difference penalty. Use small values (1e-3 .. 1.0) on noisy / sparse data.	0.0
name	str | None	Optional block name, forwarded to LookupTable1d.	None
**block_kwargs		Additional kwargs forwarded to the LookupTable1d constructor (e.g. dtype=).	{}

Returns:

Type	Description
	A LookupTable1d instance with input_array=xp and
	output_array set to the LS-fit table values.

Source code in jaxonomy/library/lookup_table_fitting.py

def fit_lookup_table_1d(
    xp,
    x_data,
    y_data,
    *,
    interpolation: str = "linear",
    extrapolation: str = "clip",
    weights=None,
    smoothness: float = 0.0,
    name: str | None = None,
    **block_kwargs,
):
    """Fit a 1-D lookup table to data and return a ``LookupTable1d`` block.

    Args:
        xp: Fixed grid of breakpoints (1-D, strictly increasing).
        x_data: Measured input cloud, shape ``(K,)``.
        y_data: Measured output cloud, shape ``(K,)``.
        interpolation: Interpolation rule for the *runtime* block
            (``"linear"`` / ``"pchip"`` / ``"nearest"`` / ``"flat"``).
            The fit itself is always linear-LS — see the module
            docstring of :mod:`jaxonomy.library.lookup_table` for why.
        extrapolation: Out-of-range policy for the runtime block; see
            :class:`jaxonomy.library.LookupTable1d`.
        weights: Optional per-sample weights for weighted least-
            squares.  ``None`` = OLS.
        smoothness: Non-negative discrete first-difference penalty.
            Use small values (1e-3 .. 1.0) on noisy / sparse data.
        name: Optional block name, forwarded to ``LookupTable1d``.
        **block_kwargs: Additional kwargs forwarded to the
            ``LookupTable1d`` constructor (e.g. ``dtype=``).

    Returns:
        A ``LookupTable1d`` instance with ``input_array=xp`` and
        ``output_array`` set to the LS-fit table values.
    """
    yp = fit_table_1d(
        xp,
        x_data,
        y_data,
        weights=weights,
        smoothness=smoothness,
    )
    # Lazy import: the block layer pulls in the rest of the library and
    # we don't want a fitting helper to drag that in at module-load
    # time.  This also keeps the import graph clean — the math module
    # has zero block-layer dependencies.
    from .primitives import LookupTable1d

    kwargs = dict(block_kwargs)
    if name is not None:
        kwargs["name"] = name
    return LookupTable1d(
        input_array=xp,
        output_array=yp,
        interpolation=interpolation,
        extrapolation=extrapolation,
        **kwargs,
    )

fit_lookup_table_2d(xp, yp, x_data, y_data, z_data, *, interpolation='linear', extrapolation='clip', weights=None, smoothness=0.0, name=None, **block_kwargs)

Fit a 2-D lookup table to data and return a LookupTable2d block.

Parameters:

Name	Type	Description	Default
xp		Fixed grid of breakpoints along the first axis (1-D, strictly increasing).	required
yp		Fixed grid of breakpoints along the second axis (1-D, strictly increasing).	required
x_data, y_data, z_data		Measurement cloud, all shape (K,).	required
interpolation	str	Interpolation rule for the runtime block (currently only "linear" / bilinear). The fit itself is always bilinear-LS.	'linear'
extrapolation	str	Out-of-range policy for the runtime block.	'clip'
weights		Optional per-sample weights for weighted least-squares.	None
smoothness	float	Non-negative 5-point-Laplacian smoothness penalty.	0.0
name	str | None	Optional block name.	None
**block_kwargs		Additional kwargs forwarded to the LookupTable2d constructor (e.g. dtype=).	{}

Returns:

Type	Description
	A LookupTable2d instance with input_x_array=xp,
	input_y_array=yp, and output_table_array set to the
	LS-fit table values of shape (len(xp), len(yp)).

Source code in jaxonomy/library/lookup_table_fitting.py

def fit_lookup_table_2d(
    xp,
    yp,
    x_data,
    y_data,
    z_data,
    *,
    interpolation: str = "linear",
    extrapolation: str = "clip",
    weights=None,
    smoothness: float = 0.0,
    name: str | None = None,
    **block_kwargs,
):
    """Fit a 2-D lookup table to data and return a ``LookupTable2d`` block.

    Args:
        xp: Fixed grid of breakpoints along the first axis (1-D,
            strictly increasing).
        yp: Fixed grid of breakpoints along the second axis (1-D,
            strictly increasing).
        x_data, y_data, z_data: Measurement cloud, all shape ``(K,)``.
        interpolation: Interpolation rule for the *runtime* block
            (currently only ``"linear"`` / bilinear).  The fit itself is
            always bilinear-LS.
        extrapolation: Out-of-range policy for the runtime block.
        weights: Optional per-sample weights for weighted least-squares.
        smoothness: Non-negative 5-point-Laplacian smoothness penalty.
        name: Optional block name.
        **block_kwargs: Additional kwargs forwarded to the
            ``LookupTable2d`` constructor (e.g. ``dtype=``).

    Returns:
        A ``LookupTable2d`` instance with ``input_x_array=xp``,
        ``input_y_array=yp``, and ``output_table_array`` set to the
        LS-fit table values of shape ``(len(xp), len(yp))``.
    """
    zp = fit_table_2d(
        xp,
        yp,
        x_data,
        y_data,
        z_data,
        weights=weights,
        smoothness=smoothness,
    )
    # Lazy import — same pattern as the 1-D wrapper.
    from .primitives import LookupTable2d

    kwargs = dict(block_kwargs)
    if name is not None:
        kwargs["name"] = name
    return LookupTable2d(
        input_x_array=xp,
        input_y_array=yp,
        output_table_array=zp,
        interpolation=interpolation,
        extrapolation=extrapolation,
        **kwargs,
    )

fit_lookup_table_nd(grid_axes, x_data, y_data, *, interpolation='linear', extrapolation='clip', weights=None, smoothness=0.0, name=None, **block_kwargs)

Fit an N-D lookup table to data and return a LookupTableND block.

The public N-D counterpart to :func:fit_lookup_table_1d and :func:fit_lookup_table_2d. Returns a fully-built block whose output_array is the LS-fit table.

Parameters:

Name	Type	Description	Default
grid_axes		Tuple of N strictly-increasing 1-D breakpoint arrays.	required
x_data, y_data		Measurement cloud — x_data shape (K, N), y_data shape (K,).	required
interpolation	str	Interpolation rule for the runtime block. Only "linear" (multilinear) is supported today; the fit itself is always multilinear-LS.	'linear'
extrapolation	str	Out-of-range policy for the runtime block; see :class:jaxonomy.library.LookupTableND.	'clip'
weights		Optional per-sample weights for weighted least-squares.	None
smoothness	float	Non-negative coefficient on the N-D Laplacian smoothness penalty.	0.0
name	str | None	Optional block name.	None
**block_kwargs		Additional kwargs forwarded to the LookupTableND constructor (e.g. dtype=).	{}

Returns:

Name	Type	Description
A		class:LookupTableND instance with the supplied grid_axes
		and output_array set to the LS-fit table of shape
		(len(grid_axes[0]), ..., len(grid_axes[N-1])).

Source code in jaxonomy/library/lookup_table_fitting.py

def fit_lookup_table_nd(
    grid_axes,
    x_data,
    y_data,
    *,
    interpolation: str = "linear",
    extrapolation: str = "clip",
    weights=None,
    smoothness: float = 0.0,
    name: str | None = None,
    **block_kwargs,
):
    """Fit an N-D lookup table to data and return a ``LookupTableND`` block.

    The public N-D counterpart to :func:`fit_lookup_table_1d` and
    :func:`fit_lookup_table_2d`.  Returns a fully-built block whose
    ``output_array`` is the LS-fit table.

    Args:
        grid_axes: Tuple of ``N`` strictly-increasing 1-D breakpoint
            arrays.
        x_data, y_data: Measurement cloud — ``x_data`` shape ``(K, N)``,
            ``y_data`` shape ``(K,)``.
        interpolation: Interpolation rule for the *runtime* block. Only
            ``"linear"`` (multilinear) is supported today; the fit
            itself is always multilinear-LS.
        extrapolation: Out-of-range policy for the runtime block; see
            :class:`jaxonomy.library.LookupTableND`.
        weights: Optional per-sample weights for weighted least-squares.
        smoothness: Non-negative coefficient on the N-D Laplacian
            smoothness penalty.
        name: Optional block name.
        **block_kwargs: Additional kwargs forwarded to the
            ``LookupTableND`` constructor (e.g. ``dtype=``).

    Returns:
        A :class:`LookupTableND` instance with the supplied ``grid_axes``
        and ``output_array`` set to the LS-fit table of shape
        ``(len(grid_axes[0]), ..., len(grid_axes[N-1]))``.
    """
    zp = fit_table_nd(
        grid_axes,
        x_data,
        y_data,
        weights=weights,
        smoothness=smoothness,
    )
    from .primitives import LookupTableND

    kwargs = dict(block_kwargs)
    if name is not None:
        kwargs["name"] = name
    return LookupTableND(
        grid_axes=tuple(grid_axes),
        output_array=zp,
        interpolation=interpolation,
        extrapolation=extrapolation,
        **kwargs,
    )

fit_pce(X, y, distributions, order)

Fit a polynomial-chaos expansion by least-squares regression.

Parameters:

Name	Type	Description	Default
X		training inputs, shape (n,) or (n, d).	required
y		training targets, shape (n,).	required
distributions	Sequence	per-dimension germ, e.g. [("normal", mu, sigma), ("uniform", a, b)]. Hermite basis for normal, Legendre for uniform (Wiener--Askey scheme, Xiu & Karniadakis 2002).	required
order	int	total-degree truncation.	required

Returns:

Name	Type	Description
A		class:PCEModel.

Source code in jaxonomy/library/rom/surrogates.py

def fit_pce(X, y, distributions: Sequence, order: int):
    """Fit a polynomial-chaos expansion by least-squares regression.

    Args:
        X: training inputs, shape ``(n,)`` or ``(n, d)``.
        y: training targets, shape ``(n,)``.
        distributions: per-dimension germ, e.g. ``[("normal", mu, sigma),
            ("uniform", a, b)]``. Hermite basis for normal, Legendre for uniform
            (Wiener--Askey scheme, Xiu & Karniadakis 2002).
        order: total-degree truncation.

    Returns:
        A :class:`PCEModel`.
    """
    X = _as2d(X)
    y = jnp.asarray(y, dtype=jnp.float64).reshape(-1)
    dim = X.shape[1]
    if len(distributions) != dim:
        raise ValueError(
            f"distributions has {len(distributions)} entries but X has {dim} "
            "feature dimensions")

    types, loc, scale = _parse_distributions(distributions)
    multi_indices = _total_degree_indices(dim, int(order))
    Xi = _pce_standardize(X, loc, scale)
    Psi = _pce_design(Xi, multi_indices, types, int(order))
    coeffs, *_ = jnp.linalg.lstsq(Psi, y, rcond=None)
    return PCEModel(coeffs, multi_indices, types, loc, scale, order)

fit_rbf(X, y, kernel='multiquadric', epsilon=1.0, smoothing=0.0, poly_degree=None)

Fit a radial-basis-function surrogate.

Parameters:

Name	Type	Description	Default
X		training inputs, shape (n,) or (n, d).	required
y		training targets, shape (n,).	required
kernel		"multiquadric", "inverse_multiquadric", "gaussian", or "thin_plate_spline".	'multiquadric'
epsilon		shape parameter (ignored by the thin-plate spline).	1.0
smoothing		ridge regularization added to the kernel diagonal; 0 gives exact interpolation.	0.0
poly_degree		if set, augment with a total-degree polynomial tail and solve the bordered saddle-point system (Wendland 2005, Ch. 8).	None

Returns:

Name	Type	Description
An		class:RBFModel.

Source code in jaxonomy/library/rom/surrogates.py

def fit_rbf(X, y, kernel="multiquadric", epsilon=1.0, smoothing=0.0,
            poly_degree=None):
    """Fit a radial-basis-function surrogate.

    Args:
        X: training inputs, shape ``(n,)`` or ``(n, d)``.
        y: training targets, shape ``(n,)``.
        kernel: ``"multiquadric"``, ``"inverse_multiquadric"``, ``"gaussian"``,
            or ``"thin_plate_spline"``.
        epsilon: shape parameter (ignored by the thin-plate spline).
        smoothing: ridge regularization added to the kernel diagonal; ``0`` gives
            exact interpolation.
        poly_degree: if set, augment with a total-degree polynomial tail and
            solve the bordered saddle-point system (Wendland 2005, Ch. 8).

    Returns:
        An :class:`RBFModel`.
    """
    X = _as2d(X)
    y = jnp.asarray(y, dtype=jnp.float64).reshape(-1)
    n = X.shape[0]

    d2 = _sqdist(X, X)
    A = _rbf_phi(d2, kernel, epsilon) + smoothing * jnp.eye(n)

    if poly_degree is None:
        weights = jnp.linalg.solve(A, y)
        return RBFModel(X, weights, None, None, kernel, epsilon)

    poly_indices = _total_degree_indices(X.shape[1], int(poly_degree))
    P = _rbf_monomials(X, poly_indices)  # (n, m)
    m = P.shape[1]
    top = jnp.concatenate([A, P], axis=1)
    bot = jnp.concatenate([P.T, jnp.zeros((m, m), dtype=jnp.float64)], axis=1)
    M = jnp.concatenate([top, bot], axis=0)
    rhs = jnp.concatenate([y, jnp.zeros(m, dtype=jnp.float64)])
    sol = jnp.linalg.solve(M, rhs)
    weights = sol[:n]
    poly_coeffs = sol[n:]
    return RBFModel(X, weights, poly_coeffs, poly_indices, kernel, epsilon)

fit_table_1d_with_grid(n_grid_points, x_data, y_data, x_lo=None, x_hi=None, init_xp=None, *, smoothness=0.0, optimizer='gd', max_iter=200, learning_rate=0.001, auto_normalize=True)

Jointly optimise the grid xp AND the table values yp.

This is the T-124-followup-grid-optimization deliverable. Phase 1's :func:fit_table_1d fits yp at a fixed user-supplied xp; here we ALSO move the breakpoints to better resolve regions where the data has strong features (sharp peaks, kinks).

The math: for any candidate grid xp, the inner problem is still a linear least-squares solve for yp (closed form). The outer loop minimises the resulting data residual w.r.t. xp, with monotonicity enforced via a smooth cumsum(softplus(deltas)) parametrisation rather than projection.

Parameters:

Name	Type	Description	Default
n_grid_points	int	Number of breakpoints to place (must be ≥ 2).	required
x_data		Measured input cloud, shape (K,).	required
y_data		Measured output cloud, shape (K,).	required
x_lo	float | None	Lower endpoint of the grid. None (default) = min(x_data). The endpoint is pinned — the optimiser only moves the interior breakpoints.	None
x_hi	float | None	Upper endpoint of the grid. None (default) = max(x_data).	None
init_xp		Optional initial grid (1-D, strictly increasing, spanning [x_lo, x_hi]). None (default) starts from a uniform grid.	None
smoothness	float	Forwarded to the inner LS solve as a discrete first-difference penalty on yp. 0.0 (default) is pure data-residual.	0.0
optimizer	str	"gd" (default) — hand-rolled fixed-step gradient descent on the unconstrained deltas. Differentiable end-to-end, jit-friendly, reliable. "lbfgs" — delegates the outer loop to :func:jax.scipy.optimize.minimize (BFGS). Faster on well-conditioned problems but does NOT support differentiation through itself (jax.grad of the joint fit w.r.t. y_data will fail with this option — use "gd" if you need the gradient). See T-124-followup-grid-optimization-lbfgs for the proper differentiable L-BFGS implementation.	'gd'
max_iter	int	Outer-loop iteration budget. For optimizer="gd" each iter is one gradient step; for "lbfgs" it is the BFGS maxiter.	200
learning_rate	float	Step size for optimizer="gd". Default is 1e-3 — the residual landscape in deltas-space has steep cliffs near sharp data features and aggressive step sizes overshoot. Ignored by "lbfgs". When auto_normalize=True (the default), the learning rate is applied in the normalised [-1, +1] x-space and [-1, +1] y-space rather than in the user's natural units, so a single sensible default works across orders-of-magnitude data scales.	0.001
auto_normalize	bool	When True (default, T-124-followup-grid-fit-auto-normalize), the optimiser internally rescales x_data and y_data to roughly [-1, +1] (zero-mean, unit-half-range affine transform) so the learning_rate=1e-3 default works on wide-but-smooth features (e.g. an engine-map slice with rpm ∈ [80, 650] and torque ~250 N·m) without blowing up to NaN. The optimised grid and table are transformed back to the user's natural units on return — results are byte-equivalent to the pre-normalisation path on data that was already centred near unit scale. Set to False to disable (e.g. for byte-equivalent reproduction of pre-normalisation runs, or when the learning_rate was tuned in natural units).	True

Returns:

Type	Description
	(xp_opt, yp_opt) — the optimised grid (shape
	(n_grid_points,)) and the corresponding optimal table
	values (shape (n_grid_points,)). Both differentiable
	through y_data (and through x_data modulo the discrete
	bucket index in the design matrix) when
	optimizer="gd".

Honest fallback note: this ships the gradient-descent path as the primary solver (rather than full L-BFGS) per the task spec — it's slower but more robust on the inner-outer formulation. The proper differentiable L-BFGS via implicit-function-theorem unrolling is filed as T-124-followup-grid-optimization-lbfgs.

Source code in jaxonomy/library/lookup_table_fitting.py

def fit_table_1d_with_grid(
    n_grid_points: int,
    x_data,
    y_data,
    x_lo: float | None = None,
    x_hi: float | None = None,
    init_xp=None,
    *,
    smoothness: float = 0.0,
    optimizer: str = "gd",
    max_iter: int = 200,
    learning_rate: float = 1e-3,
    auto_normalize: bool = True,
):
    """Jointly optimise the grid ``xp`` AND the table values ``yp``.

    This is the T-124-followup-grid-optimization deliverable.  Phase 1's
    :func:`fit_table_1d` fits ``yp`` at a fixed user-supplied ``xp``;
    here we ALSO move the breakpoints to better resolve regions where
    the data has strong features (sharp peaks, kinks).

    The math: for any candidate grid ``xp``, the inner problem is still
    a linear least-squares solve for ``yp`` (closed form).  The outer
    loop minimises the resulting data residual w.r.t. ``xp``, with
    monotonicity enforced via a smooth ``cumsum(softplus(deltas))``
    parametrisation rather than projection.

    Args:
        n_grid_points: Number of breakpoints to place (must be ≥ 2).
        x_data: Measured input cloud, shape ``(K,)``.
        y_data: Measured output cloud, shape ``(K,)``.
        x_lo: Lower endpoint of the grid.  ``None`` (default) = ``min(x_data)``.
            The endpoint is *pinned* — the optimiser only moves the
            interior breakpoints.
        x_hi: Upper endpoint of the grid.  ``None`` (default) = ``max(x_data)``.
        init_xp: Optional initial grid (1-D, strictly increasing,
            spanning ``[x_lo, x_hi]``).  ``None`` (default) starts from
            a uniform grid.
        smoothness: Forwarded to the inner LS solve as a discrete
            first-difference penalty on ``yp``.  ``0.0`` (default) is
            pure data-residual.
        optimizer: ``"gd"`` (default) — hand-rolled fixed-step gradient
            descent on the unconstrained ``deltas``.  Differentiable
            end-to-end, jit-friendly, reliable.  ``"lbfgs"`` —
            delegates the outer loop to
            :func:`jax.scipy.optimize.minimize` (BFGS).  Faster on
            well-conditioned problems but does NOT support
            differentiation through itself (jax.grad of the joint fit
            w.r.t. y_data will fail with this option — use ``"gd"`` if
            you need the gradient).  See
            ``T-124-followup-grid-optimization-lbfgs`` for the proper
            differentiable L-BFGS implementation.
        max_iter: Outer-loop iteration budget.  For ``optimizer="gd"``
            each iter is one gradient step; for ``"lbfgs"`` it is the
            BFGS maxiter.
        learning_rate: Step size for ``optimizer="gd"``.  Default is
            ``1e-3`` — the residual landscape in ``deltas``-space has
            steep cliffs near sharp data features and aggressive step
            sizes overshoot.  Ignored by ``"lbfgs"``.  When
            ``auto_normalize=True`` (the default), the learning rate is
            applied in the normalised ``[-1, +1]`` x-space and ``[-1, +1]``
            y-space rather than in the user's natural units, so a single
            sensible default works across orders-of-magnitude data
            scales.
        auto_normalize: When ``True`` (default,
            T-124-followup-grid-fit-auto-normalize), the optimiser
            internally rescales ``x_data`` and ``y_data`` to roughly
            ``[-1, +1]`` (zero-mean, unit-half-range affine transform)
            so the ``learning_rate=1e-3`` default works on
            wide-but-smooth features (e.g. an engine-map slice with
            ``rpm ∈ [80, 650]`` and ``torque ~250 N·m``) without
            blowing up to NaN.  The optimised grid and table are
            transformed back to the user's natural units on return —
            results are byte-equivalent to the pre-normalisation path
            on data that was already centred near unit scale.  Set to
            ``False`` to disable (e.g. for byte-equivalent
            reproduction of pre-normalisation runs, or when the
            ``learning_rate`` was tuned in natural units).

    Returns:
        ``(xp_opt, yp_opt)`` — the optimised grid (shape
        ``(n_grid_points,)``) and the corresponding optimal table
        values (shape ``(n_grid_points,)``).  Both differentiable
        through ``y_data`` (and through ``x_data`` modulo the discrete
        bucket index in the design matrix) when
        ``optimizer="gd"``.

    Honest fallback note: this ships the gradient-descent path as the
    primary solver (rather than full L-BFGS) per the task spec — it's
    slower but more robust on the inner-outer formulation.  The proper
    differentiable L-BFGS via implicit-function-theorem unrolling is
    filed as ``T-124-followup-grid-optimization-lbfgs``.
    """
    if n_grid_points < 2:
        raise ValueError(
            f"fit_table_1d_with_grid: n_grid_points must be >= 2, got "
            f"{n_grid_points}"
        )
    if optimizer not in ("gd", "lbfgs"):
        raise ValueError(
            f"fit_table_1d_with_grid: unknown optimizer {optimizer!r}; expected "
            f"one of ('gd', 'lbfgs')"
        )
    x_data = jnp.asarray(x_data)
    y_data = jnp.asarray(y_data)
    if x_data.shape != y_data.shape:
        raise ValueError(
            f"fit_table_1d_with_grid: x_data shape {x_data.shape} must match "
            f"y_data shape {y_data.shape}"
        )
    if x_data.ndim != 1:
        raise ValueError(
            f"fit_table_1d_with_grid: x_data must be 1-D, got shape "
            f"{x_data.shape}"
        )
    if smoothness < 0:
        raise ValueError(
            f"fit_table_1d_with_grid: smoothness must be >= 0, got {smoothness}"
        )

    # Resolve endpoints from data when not supplied.  Cast to the data
    # dtype so npa.float64 propagates through (T-005 default-float64).
    dtype = jnp.result_type(x_data, y_data)
    if x_lo is None:
        x_lo_v = jnp.min(x_data).astype(dtype)
    else:
        x_lo_v = jnp.asarray(x_lo, dtype=dtype)
    if x_hi is None:
        x_hi_v = jnp.max(x_data).astype(dtype)
    else:
        x_hi_v = jnp.asarray(x_hi, dtype=dtype)

    # T-124-followup-grid-fit-auto-normalize — affine-transform x and y
    # to roughly ``[-1, +1]`` so the default ``learning_rate=1e-3``
    # works across orders-of-magnitude data scales. Without this, a
    # wide-but-smooth feature (e.g. ``x ∈ [80, 650]``, ``y ~ 250``)
    # makes the residual landscape so steep in ``deltas``-space that
    # any non-trivial learning rate produces NaN gradients while safe
    # learning rates barely move the breakpoints.
    if auto_normalize:
        x_center = 0.5 * (x_lo_v + x_hi_v)
        x_half_range = 0.5 * (x_hi_v - x_lo_v)
        # Guard against degenerate ranges (all x_data identical); fall
        # back to no scaling rather than dividing by zero.
        x_half_range = jnp.where(x_half_range < 1e-30, jnp.asarray(1.0, dtype=dtype), x_half_range)
        y_center = jnp.mean(y_data).astype(dtype)
        y_half_range = jnp.maximum(
            jnp.max(jnp.abs(y_data - y_center)),
            jnp.asarray(1e-30, dtype=dtype),
        ).astype(dtype)
        x_data_norm = (x_data - x_center) / x_half_range
        y_data_norm = (y_data - y_center) / y_half_range
        x_lo_norm = (x_lo_v - x_center) / x_half_range
        x_hi_norm = (x_hi_v - x_center) / x_half_range
        # The smoothness penalty is on first differences of yp; under
        # the y-rescale by ``y_half_range``, the residual scales as
        # ``y_half_range^2`` and a fair smoothness penalty has to
        # rescale to match (otherwise the regulariser strength shifts
        # with the data scale). The user passes ``smoothness`` in
        # natural units; convert to the normalised space the inner
        # solve sees.
        smoothness_norm = smoothness  # cancels: both data residual
        # and penalty matrix-D scale the same way in normalised space.
    else:
        x_data_norm = x_data
        y_data_norm = y_data
        x_lo_norm = x_lo_v
        x_hi_norm = x_hi_v
        smoothness_norm = smoothness

    # Seed the outer-loop parameters (deltas).  A uniform grid
    # corresponds to all-equal deltas; the precise value is inverted
    # from the desired uniform spacing. ``init_xp`` is supplied in
    # natural units; convert to the normalised space when needed.
    if init_xp is None:
        init_xp_norm = jnp.linspace(x_lo_norm, x_hi_norm, n_grid_points)
    else:
        init_xp_v = jnp.asarray(init_xp, dtype=dtype)
        if init_xp_v.shape != (n_grid_points,):
            raise ValueError(
                f"fit_table_1d_with_grid: init_xp shape {init_xp_v.shape} must "
                f"be ({n_grid_points},)"
            )
        if auto_normalize:
            init_xp_norm = (init_xp_v - x_center) / x_half_range
        else:
            init_xp_norm = init_xp_v

    deltas0 = _deltas_from_xp(init_xp_norm, x_lo_norm, x_hi_norm).astype(dtype)

    def loss_fn(deltas):
        xp = _xp_from_deltas(deltas, x_lo_norm, x_hi_norm)
        return _residual_at_xp(xp, x_data_norm, y_data_norm, smoothness_norm)

    if optimizer == "gd":
        # Fixed-step gradient descent.  ``jax.lax.scan`` unrolls
        # cleanly under jit and supports the implicit-derivatives
        # path through ``jax.grad`` on the outer loss.
        grad_fn = jax.grad(loss_fn)

        def step(carry, _):
            d = carry
            g = grad_fn(d)
            d_new = d - learning_rate * g
            return d_new, None

        deltas_opt, _ = jax.lax.scan(step, deltas0, xs=None, length=max_iter)
    else:
        # optimizer == "lbfgs": forward-only path via jax.scipy BFGS.
        # Note: jax.scipy.optimize.minimize does not support
        # differentiation through itself (per its docstring), so this
        # branch breaks the grad-through-fit story.  Use "gd" if you
        # need that gradient.
        from jax.scipy.optimize import minimize as _jmin

        result = _jmin(
            loss_fn,
            deltas0,
            method="BFGS",
            options={"maxiter": int(max_iter)},
        )
        deltas_opt = result.x

    xp_opt_norm = _xp_from_deltas(deltas_opt, x_lo_norm, x_hi_norm)
    if auto_normalize:
        # Undo the affine x-transform so the returned grid lives in
        # natural units. The matching yp solve below uses the natural-
        # unit data, so the returned table values are in natural units
        # too (no extra y-transform needed on the output).
        xp_opt = xp_opt_norm * x_half_range + x_center
    else:
        xp_opt = xp_opt_norm
    # Final inner solve at the optimised grid — gives the matching yp.
    yp_opt = fit_table_1d(xp_opt, x_data, y_data, smoothness=smoothness)
    return xp_opt, yp_opt

fit_table_2d(xp, yp, x_data, y_data, z_data, weights=None, smoothness=0.0, rcond=None)

Fit a 2-D lookup table zp at the fixed grid (xp, yp) to (x_data, y_data, z_data).

Solves the bilinear least-squares problem

min_zp  Σ_k w_k * (z_data[k] - bilinear_interp(x_data[k], y_data[k]; xp, yp, zp))²
        + smoothness * Σ_{i,j} (zp[i,j] - mean(neighbours))²

via :func:jnp.linalg.lstsq on the bilinear design matrix. Linear bilinear only — the 2-D analogue of the fit_table_1d linear-only restriction. See T-124-followup-2d-pchip-fit for non-linear extensions.

Parameters:

Name	Type	Description	Default
xp		1-D, strictly increasing grid along the first axis (Nx).	required
yp		1-D, strictly increasing grid along the second axis (Ny).	required
x_data, y_data, z_data		Measurement cloud, all shape (K,).	required
weights		Optional per-sample weights, shape (K,). None means uniform weighting.	None
smoothness	float	Non-negative coefficient for the 5-point Laplacian penalty. 0.0 (default) is pure data-fit; small values (1e-3 .. 1.0) regularise on noisy / sparse data.	0.0
rcond	float | None	Forwarded to :func:jnp.linalg.lstsq.	None

Returns:

Type	Description
	zp of shape (len(xp), len(yp)) — the optimal table
	values. Differentiable through z_data (and through
	x_data / y_data modulo the discrete bucket indices).

Source code in jaxonomy/library/lookup_table_fitting.py

def fit_table_2d(
    xp,
    yp,
    x_data,
    y_data,
    z_data,
    weights=None,
    smoothness: float = 0.0,
    rcond: float | None = None,
):
    """Fit a 2-D lookup table ``zp`` at the fixed grid ``(xp, yp)`` to
    ``(x_data, y_data, z_data)``.

    Solves the bilinear least-squares problem

        min_zp  Σ_k w_k * (z_data[k] - bilinear_interp(x_data[k], y_data[k]; xp, yp, zp))²
                + smoothness * Σ_{i,j} (zp[i,j] - mean(neighbours))²

    via :func:`jnp.linalg.lstsq` on the bilinear design matrix.  Linear
    bilinear only — the 2-D analogue of the ``fit_table_1d`` linear-only
    restriction.  See ``T-124-followup-2d-pchip-fit`` for non-linear
    extensions.

    Args:
        xp: 1-D, strictly increasing grid along the first axis (``Nx``).
        yp: 1-D, strictly increasing grid along the second axis (``Ny``).
        x_data, y_data, z_data: Measurement cloud, all shape ``(K,)``.
        weights: Optional per-sample weights, shape ``(K,)``.  ``None``
            means uniform weighting.
        smoothness: Non-negative coefficient for the 5-point Laplacian
            penalty.  ``0.0`` (default) is pure data-fit; small values
            (1e-3 .. 1.0) regularise on noisy / sparse data.
        rcond: Forwarded to :func:`jnp.linalg.lstsq`.

    Returns:
        ``zp`` of shape ``(len(xp), len(yp))`` — the optimal table
        values.  Differentiable through ``z_data`` (and through
        ``x_data`` / ``y_data`` modulo the discrete bucket indices).
    """
    xp = jnp.asarray(xp)
    yp = jnp.asarray(yp)
    x_data = jnp.asarray(x_data)
    y_data = jnp.asarray(y_data)
    z_data = jnp.asarray(z_data)
    if z_data.shape != x_data.shape:
        raise ValueError(
            f"fit_table_2d: z_data shape {z_data.shape} must match x_data "
            f"shape {x_data.shape}"
        )
    if smoothness < 0:
        raise ValueError(
            f"fit_table_2d: smoothness must be >= 0, got {smoothness}"
        )

    A = _build_bilinear_design(xp, yp, x_data, y_data)
    b = z_data

    if weights is not None:
        w = jnp.asarray(weights)
        if w.shape != x_data.shape:
            raise ValueError(
                f"fit_table_2d: weights shape {w.shape} must match x_data "
                f"shape {x_data.shape}"
            )
        sqrt_w = jnp.sqrt(w)
        A = A * sqrt_w[:, None]
        b = b * sqrt_w

    nx = xp.shape[0]
    ny = yp.shape[0]
    if smoothness > 0:
        L = _build_laplacian_2d(nx, ny, A.dtype)
        A = jnp.concatenate([A, jnp.sqrt(smoothness) * L], axis=0)
        b = jnp.concatenate([b, jnp.zeros(nx * ny, dtype=b.dtype)], axis=0)

    z_flat, *_ = jnp.linalg.lstsq(A, b, rcond=rcond)
    return z_flat.reshape(nx, ny)

fit_table_nd(grid_axes, x_data, y_data, *, weights=None, smoothness=0.0, rcond=None)

Fit an N-D lookup table at fixed grid breakpoints.

Solves the multilinear least-squares problem

min_zp  Σ_k w_k * (y_data[k] - multilinear_interp(x_data[k]; grid_axes, zp))²
        + smoothness * Σ_cells (zp[cell] - mean(in-bounds neighbours))²

via :func:jnp.linalg.lstsq on the multilinear design matrix. Generalises :func:fit_table_2d (and :func:fit_table_1d for N=1) to an arbitrary number of grid axes.

Parameters:

Name	Type	Description	Default
grid_axes		Tuple of N strictly-increasing 1-D breakpoint arrays. Axis d has length B_d.	required
x_data		Query points, shape (K, N). Column d is the d-th coordinate.	required
y_data		Sample values at the query points, shape (K,).	required
weights		Optional per-sample weights, shape (K,). None means uniform weighting.	None
smoothness	float	Non-negative coefficient on the N-D-Laplacian penalty. 0.0 (default) is a pure data fit; small values (1e-3 .. 1.0) regularise on noisy / sparse data. The Laplacian is the canonical smoother on a regular grid — far better than diagonal Tikhonov, especially for cells without nearby measurements.	0.0
rcond	float | None	Forwarded to :func:jnp.linalg.lstsq.	None

Returns:

Type	Description
	zp of shape (B_1, ..., B_N) — the optimal table values.
	Layout matches :class:LookupTableND.output_array exactly, so
	the result can be passed straight through.

Memory note: builds a dense (K + prod(B_i), prod(B_i)) design matrix. For N=5 with B_i = 10 that's 10^5 columns — fine on CPU up to a few thousand measurements. For larger tables or higher-D problems, switch to a sparse solver (filed under T-104-followup-fit-table-nd-sparse).

Source code in jaxonomy/library/lookup_table_fitting.py

def fit_table_nd(
    grid_axes,
    x_data,
    y_data,
    *,
    weights=None,
    smoothness: float = 0.0,
    rcond: float | None = None,
):
    """Fit an N-D lookup table at fixed grid breakpoints.

    Solves the multilinear least-squares problem

        min_zp  Σ_k w_k * (y_data[k] - multilinear_interp(x_data[k]; grid_axes, zp))²
                + smoothness * Σ_cells (zp[cell] - mean(in-bounds neighbours))²

    via :func:`jnp.linalg.lstsq` on the multilinear design matrix.
    Generalises :func:`fit_table_2d` (and :func:`fit_table_1d` for
    ``N=1``) to an arbitrary number of grid axes.

    Args:
        grid_axes: Tuple of ``N`` strictly-increasing 1-D breakpoint
            arrays. Axis ``d`` has length ``B_d``.
        x_data: Query points, shape ``(K, N)``. Column ``d`` is the
            ``d``-th coordinate.
        y_data: Sample values at the query points, shape ``(K,)``.
        weights: Optional per-sample weights, shape ``(K,)``. ``None``
            means uniform weighting.
        smoothness: Non-negative coefficient on the N-D-Laplacian
            penalty. ``0.0`` (default) is a pure data fit; small values
            (1e-3 .. 1.0) regularise on noisy / sparse data. The
            Laplacian is the canonical smoother on a regular grid — far
            better than diagonal Tikhonov, especially for cells without
            nearby measurements.
        rcond: Forwarded to :func:`jnp.linalg.lstsq`.

    Returns:
        ``zp`` of shape ``(B_1, ..., B_N)`` — the optimal table values.
        Layout matches :class:`LookupTableND.output_array` exactly, so
        the result can be passed straight through.

    Memory note: builds a dense ``(K + prod(B_i), prod(B_i))`` design
    matrix. For ``N=5`` with ``B_i = 10`` that's ``10^5`` columns —
    fine on CPU up to a few thousand measurements. For larger tables
    or higher-D problems, switch to a sparse solver (filed under
    ``T-104-followup-fit-table-nd-sparse``).
    """
    x_data = jnp.asarray(x_data)
    y_data = jnp.asarray(y_data)
    if y_data.shape != (x_data.shape[0],):
        raise ValueError(
            f"fit_table_nd: y_data shape {y_data.shape} must equal "
            f"(K,) = ({x_data.shape[0]},)."
        )
    if smoothness < 0:
        raise ValueError(
            f"fit_table_nd: smoothness must be >= 0, got {smoothness}"
        )

    A, B, _strides = _build_multilinear_design(grid_axes, x_data)
    b = y_data

    if weights is not None:
        w = jnp.asarray(weights)
        if w.shape != (x_data.shape[0],):
            raise ValueError(
                f"fit_table_nd: weights shape {w.shape} must equal "
                f"(K,) = ({x_data.shape[0]},)."
            )
        sqrt_w = jnp.sqrt(w)
        A = A * sqrt_w[:, None]
        b = b * sqrt_w

    if smoothness > 0:
        L = _build_laplacian_nd(B, A.dtype)
        A = jnp.concatenate([A, jnp.sqrt(smoothness) * L], axis=0)
        b = jnp.concatenate([b, jnp.zeros(L.shape[0], dtype=b.dtype)], axis=0)

    z_flat, *_ = jnp.linalg.lstsq(A, b, rcond=rcond)
    return z_flat.reshape(B)

frequency_response(linsys, omegas)

Compute the frequency response of a linearized state-space system.

For a continuous-time LTI ẋ = Ax + Bu, y = Cx + Du the transfer function evaluated at s = jω is the (p, m) transfer-function matrix G(s) = C (sI − A)⁻¹ B + D. This helper vectorises that evaluation across an omegas array and naturally handles MIMO systems (m > 1 inputs and/or p > 1 outputs) — the returned array shape is always (K, p, m). SISO is the special case p = m = 1 which produces shape (K, 1, 1).

Parameters:

Name	Type	Description	Default
linsys	LinearizedSystem	A :class:LinearizedSystem (typically produced by :func:linearize). A is (n, n), B is (n, m), C is (p, n), D is (p, m).	required
omegas		1-D array-like of angular frequencies ω (rad/s).	required

Returns:

Type	Description
FrequencyResponse	class:FrequencyResponse with omegas (shape (K,)), complex
FrequencyResponse	response (shape (K, p, m)), and corresponding magnitudes
FrequencyResponse	and phases (radians). For MIMO response[k, i, j] is the
FrequencyResponse	transfer function from input j to output i evaluated at
FrequencyResponse	ω = omegas[k].

Notes

The implementation is fully JAX-traceable and differentiable through A, B, C, D and omegas so it composes with jax.grad and jax.vmap. jnp.linalg.solve solves the matrix RHS B in one shot per frequency so the per-omega cost is one LU decomposition plus m triangular back-substitutions — substantially cheaper than looping over input channels.

Source code in jaxonomy/library/linearization_workflow.py

def frequency_response(linsys: LinearizedSystem, omegas) -> FrequencyResponse:
    """Compute the frequency response of a linearized state-space system.

    For a continuous-time LTI ``ẋ = Ax + Bu, y = Cx + Du`` the transfer function
    evaluated at ``s = jω`` is the ``(p, m)`` transfer-function matrix
    ``G(s) = C (sI − A)⁻¹ B + D``.  This helper vectorises that evaluation
    across an ``omegas`` array and naturally handles MIMO systems
    (``m > 1`` inputs and/or ``p > 1`` outputs) — the returned array shape
    is always ``(K, p, m)``.  SISO is the special case ``p = m = 1`` which
    produces shape ``(K, 1, 1)``.

    Args:
        linsys: A :class:`LinearizedSystem` (typically produced by
            :func:`linearize`).  ``A`` is ``(n, n)``, ``B`` is ``(n, m)``,
            ``C`` is ``(p, n)``, ``D`` is ``(p, m)``.
        omegas: 1-D array-like of angular frequencies ``ω`` (rad/s).

    Returns:
        :class:`FrequencyResponse` with ``omegas`` (shape ``(K,)``), complex
        ``response`` (shape ``(K, p, m)``), and corresponding ``magnitudes``
        and ``phases`` (radians).  For MIMO ``response[k, i, j]`` is the
        transfer function from input ``j`` to output ``i`` evaluated at
        ``ω = omegas[k]``.

    Notes:
        The implementation is fully JAX-traceable and differentiable through
        ``A, B, C, D`` and ``omegas`` so it composes with ``jax.grad`` and
        ``jax.vmap``.  ``jnp.linalg.solve`` solves the matrix RHS ``B`` in
        one shot per frequency so the per-omega cost is one LU decomposition
        plus ``m`` triangular back-substitutions — substantially cheaper than
        looping over input channels.
    """
    omegas = jnp.asarray(omegas)
    if omegas.ndim == 0:
        omegas = omegas.reshape((1,))

    A = jnp.asarray(linsys.A)
    if A.ndim == 0:
        A = A.reshape((1, 1))
    elif A.ndim == 1:
        n = A.size
        A = A.reshape((n, n))
    n = A.shape[0]

    B = _ensure_2d(linsys.B, n, max(jnp.asarray(linsys.B).size // n, 1))
    m = B.shape[1]

    C_arr = jnp.asarray(linsys.C)
    if C_arr.ndim <= 1:
        # Single-output row vector or scalar
        C_arr = C_arr.reshape((max(C_arr.size // n, 1), n))
    p = C_arr.shape[0]

    D = _ensure_2d(linsys.D, p, m)

    # Promote to complex once.
    A_c = A.astype(jnp.complex64) if A.dtype == jnp.float32 else A.astype(jnp.complex128)
    B_c = B.astype(A_c.dtype)
    C_c = C_arr.astype(A_c.dtype)
    D_c = D.astype(A_c.dtype)
    eye = jnp.eye(n, dtype=A_c.dtype)

    # Discrete-time systems evaluate the transfer function at z = e^{jωΔt}
    # rather than at s = jω; the matrices (A, B, C, D) are already the
    # discrete realization. ``is_discrete`` is a Python bool, so the branch
    # is resolved statically under jit/vmap.
    is_discrete = linsys.dt is not None
    dt = linsys.dt

    def one(omega):
        if is_discrete:
            s = jnp.exp(1j * omega * dt).astype(A_c.dtype)
        else:
            s = (1j * omega).astype(A_c.dtype)
        # Solve (sI − A) X = B   →   X = (sI − A)⁻¹ B
        # (for discrete systems s is the z-variable e^{jωΔt}).
        X = jnp.linalg.solve(s * eye - A_c, B_c)
        return C_c @ X + D_c

    response = jax.vmap(one)(omegas)  # shape (K, p, m)
    magnitudes = jnp.abs(response)
    phases = jnp.angle(response)
    return FrequencyResponse(
        omegas=omegas,
        response=response,
        magnitudes=magnitudes,
        phases=phases,
    )

galerkin_reduce(rhs_fn, basis, x_ref=None, output_fn=None, input_size=0, test_basis=None, name=None)

Project a full-order RHS onto a reduced basis (POD-Galerkin / LSPG).

Parameters:

Name	Type	Description	Default
rhs_fn	Callable	Full-order dynamics. Called as rhs_fn(t, x_full, u) when input_size > 0 else rhs_fn(t, x_full); must be jax-traceable and return dx_full of shape (n_features,).	required
basis		Trial basis Φ, shape (n_features, r).	required
x_ref		Reference/offset state added on reconstruction (default zeros).	None
output_fn	Optional[Callable]	Optional map applied to the reconstructed full state for the output port.	None
input_size	int	Width of the single input port; 0 for an autonomous block (no input port).	0
test_basis		Optional test basis Ψ (shape (n_features, r)) for a Petrov-Galerkin/LSPG projection W = (Ψᵀ Φ)⁻¹ Ψᵀ. When None, Galerkin W = Φᵀ.	None
name	Optional[str]	Optional block name.	None

Returns:

Type	Description
_ProjectionROM	A jaxonomy LeafSystem with r reduced continuous states.

Source code in jaxonomy/library/rom/pod.py

def galerkin_reduce(
    rhs_fn: Callable,
    basis,
    x_ref=None,
    output_fn: Optional[Callable] = None,
    input_size: int = 0,
    test_basis=None,
    name: Optional[str] = None,
) -> _ProjectionROM:
    """Project a full-order RHS onto a reduced basis (POD-Galerkin / LSPG).

    Args:
        rhs_fn: Full-order dynamics. Called as ``rhs_fn(t, x_full, u)`` when
            ``input_size > 0`` else ``rhs_fn(t, x_full)``; must be
            jax-traceable and return ``dx_full`` of shape ``(n_features,)``.
        basis: Trial basis ``Φ``, shape ``(n_features, r)``.
        x_ref: Reference/offset state added on reconstruction (default zeros).
        output_fn: Optional map applied to the reconstructed full state for the
            output port.
        input_size: Width of the single input port; ``0`` for an autonomous
            block (no input port).
        test_basis: Optional test basis ``Ψ`` (shape ``(n_features, r)``) for a
            Petrov-Galerkin/LSPG projection ``W = (Ψᵀ Φ)⁻¹ Ψᵀ``. When ``None``,
            Galerkin ``W = Φᵀ``.
        name: Optional block name.

    Returns:
        A jaxonomy ``LeafSystem`` with ``r`` reduced continuous states.
    """
    Phi = np.asarray(basis, dtype=float)
    n, r = Phi.shape
    x_ref_arr = np.zeros(n) if x_ref is None else np.asarray(x_ref, dtype=float)

    if test_basis is None:
        W = Phi.T
    else:
        Psi = np.asarray(test_basis, dtype=float)
        W = np.linalg.solve(Psi.T @ Phi, Psi.T)

    return _ProjectionROM(
        Phi=jnp.asarray(Phi),
        W=jnp.asarray(W),
        x_ref=jnp.asarray(x_ref_arr),
        rhs_fn=rhs_fn,
        output_fn=output_fn,
        input_size=input_size,
        name=name,
    )

hankel_singular_values(sys)

Hankel singular values, sorted descending.

:math:\sigma_i = \sqrt{\lambda_i(W_c W_o)} for the controllability and observability Gramians of sys (Moore 1981).

Source code in jaxonomy/library/rom/linear_mor.py

def hankel_singular_values(sys):
    r"""Hankel singular values, sorted descending.

    :math:`\sigma_i = \sqrt{\lambda_i(W_c W_o)}` for the controllability and
    observability Gramians of ``sys`` (Moore 1981).
    """
    A, B, C, _ = _abcd(sys)
    Wc = controllability_gramian(A, B, dt=sys.dt)
    Wo = observability_gramian(A, C, dt=sys.dt)
    eig = np.linalg.eigvals(Wc @ Wo)
    hsv = np.sqrt(np.clip(np.real(eig), 0.0, None))
    return np.sort(hsv)[::-1]

identity_dictionary()

Trivial dictionary g(x) = x.

eDMD with this dictionary reduces to plain (linear) DMD — a useful baseline.

Source code in jaxonomy/library/rom/koopman.py

def identity_dictionary() -> Callable:
    """Trivial dictionary ``g(x) = x``.

    eDMD with this dictionary reduces to plain (linear) DMD — a useful baseline.
    """

    def g(x):
        return jnp.atleast_1d(x)

    return g

impulse_response(linsys, t_grid)

Closed-form impulse response of a continuous-time LTI system.

For zero initial state the (finite part of the) impulse response is

.. code-block:: text

y(t) = C · expm(A·t) · B            for t > 0

The Dirac component D · δ(t) is omitted from the returned samples since it is not representable on a numeric grid; consumers that need it can add D to the t = 0 sample explicitly.

Parameters:

Name	Type	Description	Default
linsys	LinearizedSystem	A :class:LinearizedSystem.	required
t_grid		Scalar or 1-D array of evaluation times.	required

Returns:

Type	Description
	Array of shape (K, p, m) for vector t_grid, or (p, m)
	for scalar t_grid. K = len(t_grid), p = n_outputs,
	m = n_inputs.

Notes

Fully differentiable through A, B, C, D.

Source code in jaxonomy/library/linearization_workflow.py

def impulse_response(linsys: LinearizedSystem, t_grid):
    """Closed-form impulse response of a continuous-time LTI system.

    For zero initial state the (finite part of the) impulse response is

    .. code-block:: text

        y(t) = C · expm(A·t) · B            for t > 0

    The Dirac component ``D · δ(t)`` is omitted from the returned
    samples since it is not representable on a numeric grid; consumers
    that need it can add ``D`` to the ``t = 0`` sample explicitly.

    Args:
        linsys: A :class:`LinearizedSystem`.
        t_grid: Scalar or 1-D array of evaluation times.

    Returns:
        Array of shape ``(K, p, m)`` for vector ``t_grid``, or ``(p, m)``
        for scalar ``t_grid``.  ``K = len(t_grid)``, ``p = n_outputs``,
        ``m = n_inputs``.

    Notes:
        Fully differentiable through ``A, B, C, D``.
    """
    from jax.scipy.linalg import expm

    if linsys.is_discrete():
        raise ValueError(
            "impulse_response is defined for continuous-time "
            "LinearizedSystems only, but got a discrete-time system "
            f"(dt={linsys.dt}). The closed-form continuous matrix-"
            "exponential formula does not apply to a discrete recurrence "
            "x[k+1] = A x[k] + B u[k]. Simulate the discrete diagram "
            "directly to obtain its impulse response (see KNOWN_GAPS.md)."
        )

    A, B, C, D, n, m, p = _coerce_state_space(linsys)

    t_arr = jnp.asarray(t_grid)
    scalar_input = (t_arr.ndim == 0)
    if scalar_input:
        t_arr = t_arr.reshape((1,))

    def one(t):
        return C @ expm(A * t) @ B           # (p, m)

    out = jax.vmap(one)(t_arr)               # (K, p, m)
    if scalar_input:
        return out[0]
    return out

linearize(system, base_context, name=None, output_index=None, input_port=None, output_port=None)

Linearize the system about an operating point specified by the base context.

Note: Deprecated return type. Previously returned LTISystem directly, now returns a LinearizedSystem object. Use .to_lti() on the result if you need an LTISystem block.

Source code in jaxonomy/library/linear_system.py

def linearize(
    system,
    base_context,
    name: str = None,
    output_index: int = None,
    input_port = None,
    output_port = None,
) -> "LinearizedSystem":
    """Linearize the system about an operating point specified by the base context.

    Note: Deprecated return type. Previously returned LTISystem directly, now returns 
    a LinearizedSystem object. Use `.to_lti()` on the result if you need an LTISystem block.
    """
    if input_port is None:
        assert len(system.input_ports) == 1, (
            "Linearization requires specifying input_port for systems with multiple inputs. "
            f"System {system.name} has {len(system.input_ports)} input ports."
        )
        input_port = system.input_ports[0]

    if output_port is None:
        if len(system.output_ports) > 1:
            if output_index is None:
                logger.warning(
                    "Multiple output ports detected when linearizing system %s, "
                    "using first port as output",
                    system.name,
                )
        if output_index is None:
            output_index = 0
        output_port = system.output_ports[output_index]

    xc0 = base_context.continuous_state
    u0 = input_port.eval(base_context)

    # --- validate operating point ---
    from jax.flatten_util import ravel_pytree
    import warnings
    import jax.core

    xc0_flat, _ = ravel_pytree(xc0)

    # T-109-followup-linearize-traceable: gate the host-side guards
    # behind a non-tracer check so ``jax.jit(linearize)`` /
    # ``jax.grad(linearize)`` can trace cleanly. Under any trace
    # (``jit``'s ``DynamicJaxprTracer``, ``grad``'s ``LinearizeTracer``,
    # ``vmap``'s ``BatchTracer``) the ``isfinite`` check and the
    # equilibrium-residual warning are no-ops: they exist purely to
    # give a human a clearer error/warning at eager call time, not to
    # enforce a runtime invariant. The traced path returns the same
    # A/B/C/D Jacobians (which are the actual mathematical output);
    # the user's outer driver is free to run an in-trace
    # ``jnp.isfinite`` check downstream if they want one.
    #
    # Note: we use ``isinstance(x, jax.core.Tracer)`` rather than
    # ``jax.core.is_concrete(x)`` because the latter returns ``True``
    # for ``LinearizeTracer`` (the grad pass), where ``float(x)`` still
    # blows up with ``ConcretizationTypeError``. The ``Tracer`` base
    # class is the reliable cross-trace marker.
    def _is_traced(x):
        return isinstance(x, jax.core.Tracer)

    _concrete = not _is_traced(xc0_flat)

    if _concrete and not bool(jnp.all(jnp.isfinite(xc0_flat))):
        raise ValueError(
            f"linearize(): operating-point state x0 contains non-finite values "
            f"({xc0_flat}). Ensure the base_context has a valid continuous state."
        )

    restore_fixed_val = input_port.is_fixed

    # Map from (state, inputs) to (state derivatives, outputs)
    @jax.jit
    def f(xc, u):
        context = base_context.with_continuous_state(xc)
        with input_port.fixed(u):
            xdot = system.eval_time_derivatives(context)
            y = output_port.eval(context)
        return xdot, y

    # --- equilibrium check: warn if ẋ(x0, u0) is not near zero ---
    xdot0, _ = f(xc0, u0)
    xdot0_flat, _ = ravel_pytree(xdot0)
    if not _is_traced(xdot0_flat):
        residual = float(jnp.max(jnp.abs(xdot0_flat)))
        if residual > 1e-4:
            warnings.warn(
                f"linearize(): the operating point does not appear to be an equilibrium — "
                f"max|ẋ(x₀, u₀)| = {residual:.3e}. "
                f"The returned A, B, C, D are exact Jacobians at this point, but a "
                f"physically meaningful linearization requires ẋ(x₀, u₀) ≈ 0. "
                f"Consider finding an equilibrium first (e.g. by simulating to steady state).",
                UserWarning,
                stacklevel=2,
            )

    @jax.jit
    def jac(xc, u):
        primals, tangents = jax.jvp(f, (xc0, u0), (xc, u))
        return tangents

    A, B, C, D = _jvp_to_ss(jac, xc0, u0)
    operating_point = {"x": xc0, "u": u0}

    if restore_fixed_val:
        input_port.fix_value(u0)

    return LinearizedSystem(A=A, B=B, C=C, D=D, operating_point=operating_point)

linearize_to_lti(system, base_context, input_port=None, output_port=None, name=None)

Linearize system at base_context and return an LTISystem.

Parameters:

Name	Type	Description	Default
system	'SystemBase'	A LeafSystem or Diagram to linearize. If it has multiple input or output ports, input_port and output_port must be supplied explicitly.	required
base_context	'ContextBase'	The operating-point context. State, inputs, and parameters read from this context define the point about which the linearization is performed.	required
input_port		Input port to linearize against. Required when system has more than one input port.	None
output_port		Output port to linearize against. Required when system has more than one output port.	None
name	Optional[str]	Optional name for the returned LTISystem block.	None

Returns:

Type	Description
'LTISystem'	An LTISystem block with the derived (A, B, C, D)
'LTISystem'	matrices. Drop this into a DiagramBuilder wherever the
'LTISystem'	original subdiagram would go; downstream blocks should be
'LTISystem'	wired to lti.output_ports[0] and upstream blocks to
'LTISystem'	lti.input_ports[0].

Source code in jaxonomy/library/linearize_container.py

def linearize_to_lti(
    system: "SystemBase",
    base_context: "ContextBase",
    input_port=None,
    output_port=None,
    name: Optional[str] = None,
) -> "LTISystem":
    """Linearize ``system`` at ``base_context`` and return an ``LTISystem``.

    Args:
        system: A ``LeafSystem`` or ``Diagram`` to linearize.  If it has
            multiple input or output ports, ``input_port`` and
            ``output_port`` must be supplied explicitly.
        base_context: The operating-point context.  State, inputs, and
            parameters read from this context define the point about
            which the linearization is performed.
        input_port: Input port to linearize against.  Required when
            ``system`` has more than one input port.
        output_port: Output port to linearize against.  Required when
            ``system`` has more than one output port.
        name: Optional name for the returned ``LTISystem`` block.

    Returns:
        An ``LTISystem`` block with the derived ``(A, B, C, D)``
        matrices.  Drop this into a ``DiagramBuilder`` wherever the
        original subdiagram would go; downstream blocks should be
        wired to ``lti.output_ports[0]`` and upstream blocks to
        ``lti.input_ports[0]``.
    """
    lin = linearize(
        system,
        base_context,
        input_port=input_port,
        output_port=output_port,
        name=name,
    )
    return lin.to_lti()

merge_buses(bus_a, bus_b, *, on_collision='error')

Merge two NamedTuple-shaped bus signals by union of fields.

The merged bus is a fresh NamedTuple (named "MergedBus") whose fields are bus_a._fields followed by the fields of bus_b._fields not already in bus_a (de-duplicated while preserving declaration order). The result is a JAX-pytree-friendly value identical in shape to what :class:BusCreator would produce for the merged schema.

Differentiability: gradients flow from each merged-bus leaf back to whichever input bus contributed the leaf — the underlying op is NamedTuple construction over getattr lookups, both of which are transparent to jax.grad / jax.jit.

Parameters:

Name	Type	Description	Default
bus_a		First bus signal. Must be a NamedTuple-shaped value (isinstance(bus_a, tuple) and hasattr(bus_a, "_fields")).	required
bus_b		Second bus signal. Same contract as bus_a.	required
on_collision	str	Policy for fields that appear in both inputs. "error" (default) — raise :class:ValueError. "prefer_a" — keep the value from bus_a. "prefer_b" — keep the value from bus_b. The merged-bus schema (field order) is independent of the policy: collisions only change which leaf value lands in the colliding slot.	'error'

Returns:

Type	Description
	A NamedTuple instance whose fields are the union of
	bus_a._fields and bus_b._fields.

Raises:

Type	Description
TypeError	If either input is not a NamedTuple-shaped value.
ValueError	If on_collision is not one of the supported policies, or if collisions exist and on_collision == "error".

Source code in jaxonomy/library/routing.py

def merge_buses(bus_a, bus_b, *, on_collision: str = "error"):
    """Merge two NamedTuple-shaped bus signals by union of fields.

    The merged bus is a fresh NamedTuple (named ``"MergedBus"``) whose
    fields are ``bus_a._fields`` followed by the fields of
    ``bus_b._fields`` not already in ``bus_a`` (de-duplicated while
    preserving declaration order). The result is a JAX-pytree-friendly
    value identical in shape to what :class:`BusCreator` would produce
    for the merged schema.

    Differentiability: gradients flow from each merged-bus leaf back to
    whichever input bus contributed the leaf — the underlying op is
    NamedTuple construction over ``getattr`` lookups, both of which are
    transparent to ``jax.grad`` / ``jax.jit``.

    Parameters:
        bus_a: First bus signal. Must be a NamedTuple-shaped value
            (``isinstance(bus_a, tuple) and hasattr(bus_a, "_fields")``).
        bus_b: Second bus signal. Same contract as ``bus_a``.
        on_collision: Policy for fields that appear in both inputs.

            * ``"error"`` (default) — raise :class:`ValueError`.
            * ``"prefer_a"`` — keep the value from ``bus_a``.
            * ``"prefer_b"`` — keep the value from ``bus_b``.

            The merged-bus *schema* (field order) is independent of the
            policy: collisions only change which leaf value lands in
            the colliding slot.

    Returns:
        A NamedTuple instance whose fields are the union of
        ``bus_a._fields`` and ``bus_b._fields``.

    Raises:
        TypeError: If either input is not a NamedTuple-shaped value.
        ValueError: If ``on_collision`` is not one of the supported
            policies, or if collisions exist and ``on_collision ==
            "error"``.
    """
    if not _is_bus_signal(bus_a):
        raise TypeError(
            "merge_buses expects bus_a to be a NamedTuple-shaped bus "
            f"signal; got {type(bus_a).__name__}: {bus_a!r}."
        )
    if not _is_bus_signal(bus_b):
        raise TypeError(
            "merge_buses expects bus_b to be a NamedTuple-shaped bus "
            f"signal; got {type(bus_b).__name__}: {bus_b!r}."
        )

    fields_a = tuple(bus_a._fields)
    fields_b = tuple(bus_b._fields)
    merged_order, _ = _merged_field_order(fields_a, fields_b, on_collision)

    # Build the leaf values in merged-order. For each name, look in the
    # input dictated by the collision policy (or the unique input if no
    # collision). Using ``getattr`` keeps the path transparent to JAX.
    set_a = set(fields_a)
    set_b = set(fields_b)
    values = []
    for name in merged_order:
        in_a = name in set_a
        in_b = name in set_b
        if in_a and in_b:
            # Collision — policy dictates the source (error case is
            # already raised by _merged_field_order above).
            if on_collision == "prefer_a":
                values.append(getattr(bus_a, name))
            else:  # on_collision == "prefer_b"
                values.append(getattr(bus_b, name))
        elif in_a:
            values.append(getattr(bus_a, name))
        else:
            values.append(getattr(bus_b, name))

    MergedBus = namedtuple("MergedBus", merged_order)
    return MergedBus(*values)

minimal_realization(sys, tol=1e-08)

Minimal realization of sys (Kalman decomposition).

Removes uncontrollable and unobservable modes by projecting onto the controllable subspace (range of the controllability matrix) and then onto the observable subspace (range of the observability matrix transposed). Ranks are decided from singular values with the relative threshold tol. The input/output transfer function is preserved.

Source code in jaxonomy/library/rom/linear_mor.py

def minimal_realization(sys, tol=1e-8):
    r"""Minimal realization of ``sys`` (Kalman decomposition).

    Removes uncontrollable and unobservable modes by projecting onto the
    controllable subspace (range of the controllability matrix) and then onto
    the observable subspace (range of the observability matrix transposed).
    Ranks are decided from singular values with the relative threshold ``tol``.
    The input/output transfer function is preserved.
    """
    A, B, C, D = _abcd(sys)

    # 1) restrict to the controllable subspace
    Vc = _range_basis(_controllability_matrix(A, B), tol)
    if Vc.shape[1] == 0:
        # nothing controllable → static system
        return _make_reduced(sys, np.zeros((0, 0)), np.zeros((0, B.shape[1])),
                             np.zeros((C.shape[0], 0)), D)
    Ac = Vc.T @ A @ Vc
    Bc = Vc.T @ B
    Cc = C @ Vc

    # 2) restrict to the observable subspace (dual of step 1)
    Vo = _range_basis(_controllability_matrix(Ac.T, Cc.T), tol)
    if Vo.shape[1] == 0:
        return _make_reduced(sys, np.zeros((0, 0)), np.zeros((0, B.shape[1])),
                             np.zeros((C.shape[0], 0)), D)
    Am = Vo.T @ Ac @ Vo
    Bm = Vo.T @ Bc
    Cm = Cc @ Vo
    return _make_reduced(sys, Am, Bm, Cm, D)

modal_truncation(sys, order=None, keep=None)

Modal truncation.

Transforms sys to a real block-diagonal modal realization and keeps the dominant (slowest) modes, discarding the rest. Complex-conjugate pairs are always kept or dropped together, so the reduced model stays real; the retained poles are exactly the retained eigenvalues.

• order — target number of retained states (a straddling conjugate pair may push the actual count to order + 1).
• keep — explicit iterable of modal-state indices to retain (expanded to whole blocks).
• neither — no truncation (returns the modal-form equivalent).

Because coupling to the discarded modes is dropped outright, the DC gain generally shifts; use :func:residualize to preserve it.

Source code in jaxonomy/library/rom/linear_mor.py

def modal_truncation(sys, order=None, keep=None):
    r"""Modal truncation.

    Transforms ``sys`` to a real block-diagonal modal realization and keeps
    the dominant (slowest) modes, discarding the rest.  Complex-conjugate
    pairs are always kept or dropped together, so the reduced model stays
    real; the retained poles are exactly the retained eigenvalues.

    * ``order`` — target number of retained states (a straddling conjugate
      pair may push the actual count to ``order + 1``).
    * ``keep`` — explicit iterable of modal-state indices to retain (expanded
      to whole blocks).
    * neither — no truncation (returns the modal-form equivalent).

    Because coupling to the discarded modes is dropped outright, the DC gain
    generally shifts; use :func:`residualize` to preserve it.
    """
    A, B, C, D = _abcd(sys)
    mf = _real_modal_form(A, B, C)
    n = A.shape[0]
    mask = _select_blocks(mf.blocks, sys.dt, order, keep, n)

    Ar = mf.A[np.ix_(mask, mask)]
    Br = mf.B[mask, :]
    Cr = mf.C[:, mask]
    return _make_reduced(sys, Ar, Br, Cr, D)

model_description_xml(diagram, *, model_name, guid=None, description='Exported by jaxonomy.library.fmu_export', generation_tool='jaxonomy')

Build the FMI 2.0 modelDescription XML as a string.

Parameters:

Name	Type	Description	Default
diagram	'Diagram'	A :class:~jaxonomy.framework.diagram.Diagram whose input and output ports define the FMU's I/O surface.	required
model_name	str	Human-readable model name. Also used as the modelIdentifier (with non-identifier characters stripped).	required
guid	str | None	Optional FMU GUID; auto-generated if None.	None
description	str	Free-form description string.	'Exported by jaxonomy.library.fmu_export'
generation_tool	str	Stored in the FMU metadata.	'jaxonomy'

Returns:

Type	Description
str	UTF-8 XML string ending with a trailing newline.

Source code in jaxonomy/library/fmu_export.py

def model_description_xml(
    diagram: "Diagram",
    *,
    model_name: str,
    guid: str | None = None,
    description: str = "Exported by jaxonomy.library.fmu_export",
    generation_tool: str = "jaxonomy",
) -> str:
    """Build the FMI 2.0 modelDescription XML as a string.

    Args:
        diagram: A :class:`~jaxonomy.framework.diagram.Diagram` whose
            input and output ports define the FMU's I/O surface.
        model_name: Human-readable model name.  Also used as the
            modelIdentifier (with non-identifier characters stripped).
        guid: Optional FMU GUID; auto-generated if None.
        description: Free-form description string.
        generation_tool: Stored in the FMU metadata.

    Returns:
        UTF-8 XML string ending with a trailing newline.
    """
    if guid is None:
        guid = _gen_guid()

    model_identifier = "".join(
        c if c.isalnum() or c == "_" else "_" for c in model_name
    )
    if not model_identifier or not model_identifier[0].isalpha():
        model_identifier = "M" + model_identifier

    root = ET.Element("fmiModelDescription", attrib={
        "fmiVersion": "2.0",
        "modelName": model_name,
        "guid": guid,
        "description": description,
        "generationTool": generation_tool,
        "generationDateAndTime": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "variableNamingConvention": "structured",
    })

    # CoSimulation element.
    ET.SubElement(root, "CoSimulation", attrib={
        "modelIdentifier": model_identifier,
        "canHandleVariableCommunicationStepSize": "true",
        "canInterpolateInputs": "false",
        "maxOutputDerivativeOrder": "0",
    })

    variables = ET.SubElement(root, "ModelVariables")

    # Number variables sequentially starting at 1 (FMI convention).
    next_value_ref = 1
    output_indices: list[int] = []

    # Inputs first.
    for port in diagram.input_ports:
        for varname, _shape in _flatten_port_name_shape(port):
            sv = ET.SubElement(variables, "ScalarVariable", attrib={
                "name": varname,
                "valueReference": str(next_value_ref),
                "causality": "input",
                "variability": "continuous",
                "initial": "exact",
            })
            ET.SubElement(sv, "Real", attrib={"start": "0.0"})
            next_value_ref += 1

    # Outputs.
    for port in diagram.output_ports:
        for varname, _shape in _flatten_port_name_shape(port):
            output_indices.append(next_value_ref)
            sv = ET.SubElement(variables, "ScalarVariable", attrib={
                "name": varname,
                "valueReference": str(next_value_ref),
                "causality": "output",
                "variability": "continuous",
                "initial": "calculated",
            })
            ET.SubElement(sv, "Real", attrib={})
            next_value_ref += 1

    # ModelStructure / Outputs.
    structure = ET.SubElement(root, "ModelStructure")
    if output_indices:
        outputs_el = ET.SubElement(structure, "Outputs")
        for i, _ref in enumerate(output_indices, start=len(diagram.input_ports) + 1):
            ET.SubElement(outputs_el, "Unknown", attrib={"index": str(i)})

    ET.indent(root, space="  ")
    return ET.tostring(root, encoding="utf-8", xml_declaration=True).decode("utf-8") + "\n"

nyquist_data(linsys, omegas)

Return Nyquist-contour arrays for linsys.

The Nyquist plot traces G(jω) through the complex plane. This helper returns the real and imaginary parts of G(jω) over the supplied positive angular frequencies and additionally the reflected negative-frequency arrays (since G(-jω) = conj(G(jω)) for a real-coefficient LTI, the reflection is given exactly by (Re, -Im)). Consumers can concatenate the negative and positive arrays to obtain the full closed contour used for encirclement counting; for stability margins computed only from the positive sweep, real and imag are sufficient.

For MIMO systems the returned real / imag arrays preserve the (K, p, m) channel structure from :func:frequency_response; for SISO they are squeezed to (K,), matching the convention of :func:bode_data.

Parameters:

Name	Type	Description	Default
linsys	LinearizedSystem	A :class:LinearizedSystem.	required
omegas		1-D array-like of positive angular frequencies ω (rad/s). Negative or zero entries are not rejected — G(0) is the DC gain and is returned unchanged, but duplicating with the mirror is not meaningful at ω = 0.	required

Returns:

Type	Description
dict	Dictionary with keys: "omega" — positive angular frequencies (rad/s), shape (K,); "real" — Re G(jω); shape (K,) for SISO, (K, p, m) for MIMO; "imag" — Im G(jω); same shape as real; "real_neg" — Re G(-jω) = Re G(jω); same shape as real (mirror; provided for convenience when plotting the full closed contour); "imag_neg" — Im G(-jω) = -Im G(jω); same shape as imag.

Notes

Fully differentiable through A, B, C, D and omegas via the underlying :func:frequency_response.

Source code in jaxonomy/library/linearization_workflow.py

def nyquist_data(linsys: LinearizedSystem, omegas) -> dict:
    """Return Nyquist-contour arrays for ``linsys``.

    The Nyquist plot traces ``G(jω)`` through the complex plane.  This
    helper returns the real and imaginary parts of ``G(jω)`` over the
    supplied positive angular frequencies and additionally the reflected
    negative-frequency arrays (since ``G(-jω) = conj(G(jω))`` for a
    real-coefficient LTI, the reflection is given exactly by
    ``(Re, -Im)``).  Consumers can concatenate the negative and positive
    arrays to obtain the full closed contour used for encirclement
    counting; for stability margins computed only from the positive
    sweep, ``real`` and ``imag`` are sufficient.

    For MIMO systems the returned ``real`` / ``imag`` arrays preserve
    the ``(K, p, m)`` channel structure from :func:`frequency_response`;
    for SISO they are squeezed to ``(K,)``, matching the convention of
    :func:`bode_data`.

    Args:
        linsys: A :class:`LinearizedSystem`.
        omegas: 1-D array-like of *positive* angular frequencies
            ``ω`` (rad/s).  Negative or zero entries are not rejected —
            ``G(0)`` is the DC gain and is returned unchanged, but
            duplicating with the mirror is not meaningful at ω = 0.

    Returns:
        Dictionary with keys:
            ``"omega"`` — positive angular frequencies (rad/s), shape
                ``(K,)``;
            ``"real"`` — ``Re G(jω)``; shape ``(K,)`` for SISO,
                ``(K, p, m)`` for MIMO;
            ``"imag"`` — ``Im G(jω)``; same shape as ``real``;
            ``"real_neg"`` — ``Re G(-jω) = Re G(jω)``; same shape as
                ``real`` (mirror; provided for convenience when plotting
                the full closed contour);
            ``"imag_neg"`` — ``Im G(-jω) = -Im G(jω)``; same shape as
                ``imag``.

    Notes:
        Fully differentiable through ``A, B, C, D`` and ``omegas`` via
        the underlying :func:`frequency_response`.
    """
    fr = frequency_response(linsys, omegas)
    response = fr.response
    # SISO squeeze, matching bode_data convention.
    is_siso = (
        response.ndim == 3 and response.shape[-1] == 1 and response.shape[-2] == 1
    )
    if is_siso:
        response = response[..., 0, 0]
    real = jnp.real(response)
    imag = jnp.imag(response)
    return {
        "omega": fr.omegas,
        "real": real,
        "imag": imag,
        "real_neg": real,
        "imag_neg": -imag,
    }

observability_gramian(A, C, dt=None)

Observability Gramian :math:W_o.

Continuous time (dt is None) solves

.. math:: A^\mathsf{T} W_o + W_o A = -C^\mathsf{T} C

Discrete time (dt given) solves

.. math:: A^\mathsf{T} W_o A - W_o + C^\mathsf{T} C = 0

Source code in jaxonomy/library/rom/linear_mor.py

def observability_gramian(A, C, dt=None):
    r"""Observability Gramian :math:`W_o`.

    Continuous time (``dt is None``) solves

    .. math:: A^\mathsf{T} W_o + W_o A = -C^\mathsf{T} C

    Discrete time (``dt`` given) solves

    .. math:: A^\mathsf{T} W_o A - W_o + C^\mathsf{T} C = 0
    """
    A = _np(A)
    C = _np(C)
    if A.ndim <= 1:
        A = A.reshape(1, 1)
    C = C.reshape(-1, A.shape[0])
    if dt is None:
        return sla.solve_continuous_lyapunov(A.T, -(C.T @ C))
    return sla.solve_discrete_lyapunov(A.T, C.T @ C)

pod_basis(X, rank=None, energy=None)

Proper-orthogonal-decomposition basis of a snapshot matrix.

Computes the (host-side) thin SVD X = U Σ Vᵀ and truncates to r left singular vectors, which are the energetically optimal orthonormal modes (Sirovich 1987).

Parameters:

Name	Type	Description	Default
X		Snapshot matrix, shape (n_features, n_samples).	required
rank	Optional[int]	Explicit number of modes to keep.	None
energy	Optional[float]	Cumulative-energy threshold in (0, 1] (used when rank is None), e.g. 0.99 keeps 99% of the snapshot energy.	None

Returns:

Type	Description
ndarray	(Phi, sigma, r) where Phi has shape (n_features, r) with
ndarray	orthonormal columns, sigma is the full singular-value vector, and
int	r is the retained rank.

Source code in jaxonomy/library/rom/pod.py

def pod_basis(
    X,
    rank: Optional[int] = None,
    energy: Optional[float] = None,
) -> Tuple[jnp.ndarray, jnp.ndarray, int]:
    """Proper-orthogonal-decomposition basis of a snapshot matrix.

    Computes the (host-side) thin SVD ``X = U Σ Vᵀ`` and truncates to ``r``
    left singular vectors, which are the energetically optimal orthonormal
    modes (Sirovich 1987).

    Args:
        X: Snapshot matrix, shape ``(n_features, n_samples)``.
        rank: Explicit number of modes to keep.
        energy: Cumulative-energy threshold in ``(0, 1]`` (used when ``rank``
            is ``None``), e.g. ``0.99`` keeps 99% of the snapshot energy.

    Returns:
        ``(Phi, sigma, r)`` where ``Phi`` has shape ``(n_features, r)`` with
        orthonormal columns, ``sigma`` is the full singular-value vector, and
        ``r`` is the retained rank.
    """
    X_np = np.asarray(X, dtype=float)
    U, sigma, _ = np.linalg.svd(X_np, full_matrices=False)
    r = _select_rank(sigma, rank, energy)
    Phi = U[:, :r]
    return jnp.asarray(Phi), jnp.asarray(sigma), int(r)

pole_zero_map(linsys)

Compute poles and zeros of a :class:LinearizedSystem.

Poles are the eigenvalues of A. Zeros are the (transmission) zeros of the SISO transfer function G(s) = C (sI − A)⁻¹ B + D, computed as the finite generalised eigenvalues of the Rosenbrock system pencil

.. code-block:: text

P(s) = [[ sI − A, −B ],
        [   C   ,  D ]]

by solving the generalised eigenproblem λ E v = M v with

.. code-block:: text

E = [[ I, 0 ],   M = [[ A, B ],
     [ 0, 0 ]]        [ C, D ]]

Finite eigenvalues (those with non-zero E-weight) of this pencil are the invariant zeros of the system; for SISO they coincide with the numerator roots of the transfer function. The high-frequency gain is reported as D[0, 0] (the asymptotic value of G(s) → D for |s| → ∞); for a strictly-proper system (D = 0) the leading-coefficient gain is harder to define unambiguously without polynomial fitting and is left to a deeper follow-up.

Parameters:

Name	Type	Description	Default
linsys	LinearizedSystem	A :class:LinearizedSystem. Phase 1 ships SISO support only — for MIMO systems the first input/output channel is used.	required

Returns:

Type	Description
dict	Dictionary with keys: "poles" — complex 1-D array of eigenvalues of A, "zeros" — complex 1-D array of (finite) invariant zeros, "gain" — feedthrough scalar D[0, 0].

Notes

Pole computation is differentiable through A (via jnp.linalg.eigvals). Zero computation uses :func:scipy.linalg.eig on the generalised problem and is therefore not currently traceable by JAX — call it eagerly, outside jit. A differentiable variant is a deeper follow-up.

Source code in jaxonomy/library/linearization_workflow.py

def pole_zero_map(linsys: LinearizedSystem) -> dict:
    """Compute poles and zeros of a :class:`LinearizedSystem`.

    Poles are the eigenvalues of ``A``.  Zeros are the (transmission)
    zeros of the SISO transfer function ``G(s) = C (sI − A)⁻¹ B + D``,
    computed as the finite generalised eigenvalues of the Rosenbrock
    system pencil

    .. code-block:: text

        P(s) = [[ sI − A, −B ],
                [   C   ,  D ]]

    by solving the generalised eigenproblem ``λ E v = M v`` with

    .. code-block:: text

        E = [[ I, 0 ],   M = [[ A, B ],
             [ 0, 0 ]]        [ C, D ]]

    Finite eigenvalues (those with non-zero ``E``-weight) of this pencil
    are the invariant zeros of the system; for SISO they coincide with
    the numerator roots of the transfer function.  The high-frequency
    gain is reported as ``D[0, 0]`` (the asymptotic value of
    ``G(s) → D`` for ``|s| → ∞``); for a strictly-proper system
    (``D = 0``) the leading-coefficient gain is harder to define
    unambiguously without polynomial fitting and is left to a deeper
    follow-up.

    Args:
        linsys: A :class:`LinearizedSystem`.  Phase 1 ships SISO support
            only — for MIMO systems the first input/output channel is
            used.

    Returns:
        Dictionary with keys:
            ``"poles"`` — complex 1-D array of eigenvalues of ``A``,
            ``"zeros"`` — complex 1-D array of (finite) invariant zeros,
            ``"gain"`` — feedthrough scalar ``D[0, 0]``.

    Notes:
        Pole computation is differentiable through ``A`` (via
        ``jnp.linalg.eigvals``).  Zero computation uses
        :func:`scipy.linalg.eig` on the generalised problem and is
        therefore not currently traceable by JAX — call it eagerly,
        outside ``jit``.  A differentiable variant is a deeper follow-up.
    """
    A, B, C, D, n, m, p = _coerce_state_space(linsys)

    # Poles: eigenvalues of A.
    poles = jnp.linalg.eigvals(A)

    # Zeros: finite generalised eigenvalues of the Rosenbrock pencil.
    # SISO assumption — restrict to first input/output for phase 1.
    B0 = B[:, :1]  # (n, 1)
    C0 = C[:1, :]  # (1, n)
    D0 = D[:1, :1]  # (1, 1)

    # Block-assemble M and E (size n+1 x n+1).
    M_top = jnp.concatenate([A, B0], axis=1)            # (n, n+1)
    M_bot = jnp.concatenate([C0, D0], axis=1)           # (1, n+1)
    M = jnp.concatenate([M_top, M_bot], axis=0)         # (n+1, n+1)

    E_top = jnp.concatenate(
        [jnp.eye(n, dtype=A.dtype), jnp.zeros((n, 1), dtype=A.dtype)],
        axis=1,
    )
    E_bot = jnp.zeros((1, n + 1), dtype=A.dtype)
    E = jnp.concatenate([E_top, E_bot], axis=0)

    # Generalised eigenproblem — defer to scipy.linalg.eig because
    # jax.numpy does not yet expose a generalised eigensolver.  Pure
    # NumPy fallback would also work; scipy gives us a stable QZ.
    try:
        from scipy.linalg import eig as _scipy_eig
        eigvals, _ = _scipy_eig(np.asarray(M), np.asarray(E))
        eigvals = np.asarray(eigvals)
        # Filter infinities — those correspond to dynamic modes the
        # pencil decoupled into the singular part of E.
        finite = np.isfinite(eigvals.real) & np.isfinite(eigvals.imag)
        zeros = jnp.asarray(eigvals[finite])
    except Exception:
        # If scipy unavailable, fall back to the strictly-proper
        # invertible-A formula: zeros are eigenvalues of (A − B*C/D) when
        # D ≠ 0.  Otherwise return an empty zeros vector.
        d_scalar = float(np.asarray(D0).reshape(-1)[0])
        if abs(d_scalar) > 1e-300:
            zeros = jnp.linalg.eigvals(A - (B0 @ C0) / d_scalar)
        else:
            zeros = jnp.zeros((0,), dtype=jnp.complex128)

    gain = jnp.asarray(D0).reshape(())

    return {
        "poles": poles,
        "zeros": zeros,
        "gain": float(np.asarray(gain)),
    }

polynomial_dictionary(degree, include_constant=True)

Monomial dictionary up to degree.

Layout: [x_1..x_n, (1), (degree-2..degree monomials)] — identity first so the state is recoverable. The constant term is included by default (it makes affine dynamics representable in the lifted space).

Source code in jaxonomy/library/rom/koopman.py

def polynomial_dictionary(degree: int, include_constant: bool = True) -> Callable:
    """Monomial dictionary up to ``degree``.

    Layout: ``[x_1..x_n, (1), (degree-2..degree monomials)]`` — identity first so
    the state is recoverable. The constant term is included by default (it makes
    affine dynamics representable in the lifted space).
    """

    def g(x):
        x = jnp.atleast_1d(x)
        n = x.shape[0]
        terms = [x]  # identity observables first
        if include_constant:
            terms.append(jnp.ones((1,), dtype=x.dtype))
        for d in range(2, int(degree) + 1):
            for combo in itertools.combinations_with_replacement(range(n), d):
                term = jnp.prod(jnp.stack([x[i] for i in combo]))
                terms.append(jnp.reshape(term, (1,)))
        return jnp.concatenate(terms)

    return g

projection_error(X, basis)

Relative projection error ‖X − ΦΦᵀX‖ / ‖X‖ of X onto basis.

basis (Φ) is assumed to have orthonormal columns.

Source code in jaxonomy/library/rom/snapshots.py

def projection_error(X, basis) -> float:
    """Relative projection error ``‖X − ΦΦᵀX‖ / ‖X‖`` of ``X`` onto ``basis``.

    ``basis`` (``Φ``) is assumed to have orthonormal columns.
    """
    X = np.asarray(X)
    Phi = np.asarray(basis)
    X_proj = Phi @ (Phi.T @ X)
    denom = np.linalg.norm(X)
    if denom == 0.0:
        return float(np.linalg.norm(X - X_proj))
    return float(np.linalg.norm(X - X_proj) / denom)

rbf_dictionary(centers, epsilon=1.0)

Gaussian radial-basis dictionary.

Lifts to [x, exp(-epsilon ||x - c_j||²) for each center c_j] — identity observables first, followed by one RBF feature per center.

Parameters:

Name	Type	Description	Default
centers		Array (n_centers, n) of RBF centers.	required
epsilon	float	Shape parameter of the Gaussian kernel.	1.0

Source code in jaxonomy/library/rom/koopman.py

def rbf_dictionary(centers, epsilon: float = 1.0) -> Callable:
    """Gaussian radial-basis dictionary.

    Lifts to ``[x, exp(-epsilon ||x - c_j||²) for each center c_j]`` — identity
    observables first, followed by one RBF feature per center.

    Args:
        centers: Array ``(n_centers, n)`` of RBF centers.
        epsilon: Shape parameter of the Gaussian kernel.
    """
    centers = np.asarray(centers, dtype=float)
    if centers.ndim == 1:
        centers = centers.reshape(-1, 1)
    centers_j = jnp.asarray(centers)
    eps = float(epsilon)

    def g(x):
        x = jnp.atleast_1d(x)
        diffs = centers_j - x[None, :]
        rbf = jnp.exp(-eps * jnp.sum(diffs * diffs, axis=1))
        return jnp.concatenate([x, rbf])

    return g

reduce(target, method='balred', *, order=None, tol=None, dt=1.0, **kwargs)

Reduce target by method and return a :class:ReducedOrderModel.

Parameters:

Name	Type	Description	Default
target		An LTI model (linear MOR) or snapshot data (data-driven).	required
method		See the module docstring for the supported names.	'balred'
order		Target reduced order, where the method takes one.	None
tol		Energy/tolerance selector for balanced truncation / minreal.	None
dt		Sampling period for the data-driven predictor blocks.	1.0
**kwargs		Forwarded to the underlying routine (e.g. keep= for modal methods, dictionary= and U= for eDMD, initial_state=).	{}

Returns:

Name	Type	Description
A		class:ReducedOrderModel whose .system is ready to simulate.

Source code in jaxonomy/library/rom/framework.py

def reduce(target, method="balred", *, order=None, tol=None, dt=1.0, **kwargs):
    """Reduce ``target`` by ``method`` and return a :class:`ReducedOrderModel`.

    Args:
        target: An LTI model (linear MOR) or snapshot data (data-driven).
        method: See the module docstring for the supported names.
        order: Target reduced order, where the method takes one.
        tol: Energy/tolerance selector for balanced truncation / ``minreal``.
        dt: Sampling period for the data-driven predictor blocks.
        **kwargs: Forwarded to the underlying routine (e.g. ``keep=`` for modal
            methods, ``dictionary=`` and ``U=`` for eDMD, ``initial_state=``).

    Returns:
        A :class:`ReducedOrderModel` whose ``.system`` is ready to simulate.
    """
    key = method.lower()

    if key in _LINEAR_METHODS:
        canonical = _LINEAR_METHODS[key]
        lin = _as_linearized(target)
        if canonical == "balred":
            red = _lmor.balred(lin, order=order, tol=tol)
        elif canonical == "minreal":
            red = _lmor.minreal(lin, **({"tol": tol} if tol is not None else {}), **kwargs)
        elif canonical == "modal":
            red = _lmor.modal_truncation(lin, order=order, keep=kwargs.get("keep"))
        else:  # residualize
            red = _lmor.residualize(lin, order=order, keep=kwargs.get("keep"))
        info = {"result": red}
        for attr in ("hsv", "error_bound", "reduced_order"):
            if hasattr(red, attr):
                info[attr] = getattr(red, attr)
        return ReducedOrderModel(
            system=_reduced_lti(red),
            method=canonical,
            full_order=int(np.asarray(lin.A).shape[0]),
            reduced_order=int(np.asarray(red.A).shape[0]),
            info=info,
        )

    if key in _DATA_METHODS:
        X, _, U, x0 = _snapshot_matrix(target)
        x0 = kwargs.pop("initial_state", x0)
        if key == "dmd":
            res = _dmd(X, rank=order)
            # Real full one-step operator from the (conjugate-symmetric) DMD
            # spectrum: A = Re(Φ diag(λ) Φ⁺). Tu et al. 2014.
            A_full = np.real(
                res.modes @ np.diag(res.eigenvalues) @ np.linalg.pinv(res.modes)
            )
            system = DMDForecaster(A=A_full, dt=dt, initial_state=np.asarray(x0, float))
            return ReducedOrderModel(
                system=system, method="dmd",
                full_order=A_full.shape[0], reduced_order=len(res.eigenvalues),
                info={"result": res, "eigenvalues": res.eigenvalues},
            )
        if key == "dmdc":
            Xp = kwargs.pop("Xp", None)
            U = kwargs.pop("U", U)
            if U is None:
                raise ValueError("dmdc needs control inputs U (kwarg or SnapshotData.inputs).")
            if Xp is None:
                X, Xp = X[:, :-1], X[:, 1:]
                U = np.asarray(U)[:, : Xp.shape[1]]
            res = _dmdc(X, Xp, U, rank=order)
            system = DMDForecaster(A=res.A, B=res.B, dt=dt, initial_state=np.asarray(x0, float))
            return ReducedOrderModel(
                system=system, method="dmdc",
                full_order=np.asarray(res.A).shape[0], reduced_order=np.asarray(res.A).shape[0],
                info={"result": res, "eigenvalues": res.eigenvalues},
            )
        # edmd
        dictionary = kwargs.pop("dictionary", None)
        if dictionary is None:
            raise ValueError("edmd needs a `dictionary` observable callable.")
        Xp = kwargs.pop("Xp", None)
        U = kwargs.pop("U", U)
        if Xp is None:
            X, Xp = X[:, :-1], X[:, 1:]
            if U is not None:
                U = np.asarray(U)[:, : Xp.shape[1]]
        res = _edmd(X, Xp, dictionary, U=U)
        system = KoopmanPredictor(
            K=res.K, C=res.C, dictionary=res.dictionary, B=res.B, dt=dt,
            initial_state=np.asarray(x0, float),
        )
        return ReducedOrderModel(
            system=system, method="edmd",
            full_order=np.asarray(res.C).shape[0], reduced_order=np.asarray(res.K).shape[0],
            info={"result": res},
        )

    if key in ("pod", "galerkin", "petrov_galerkin", "lspg", "deim"):
        raise ValueError(
            f"Projection ROM ({method!r}) needs the full-order RHS callable — "
            "use galerkin_reduce(rhs_fn, basis, ...) or "
            "deim_galerkin_reduce(...) directly."
        )

    raise ValueError(f"Unknown reduction method {method!r}.")

relative_error(x_true, x_approx)

Relative L2 (Frobenius) error ‖x_true − x_approx‖ / ‖x_true‖.

Works for a single trajectory column or a full snapshot matrix.

Source code in jaxonomy/library/rom/snapshots.py

def relative_error(x_true, x_approx) -> float:
    """Relative L2 (Frobenius) error ``‖x_true − x_approx‖ / ‖x_true‖``.

    Works for a single trajectory column or a full snapshot matrix.
    """
    x_true = np.asarray(x_true)
    x_approx = np.asarray(x_approx)
    denom = np.linalg.norm(x_true)
    if denom == 0.0:
        return float(np.linalg.norm(x_true - x_approx))
    return float(np.linalg.norm(x_true - x_approx) / denom)

residualize(sys, order=None, keep=None)

Singular-perturbation (residualization) reduction.

Like :func:modal_truncation, but instead of deleting the fast modes it sets their derivative (continuous) or their increment (discrete) to zero and solves for their quasi-steady value, folding it back into the retained model. This matches the DC gain of the discarded modes (Kokotović, Khalil & O'Reilly 1986).

Partitioning the modal realization into retained (1) and discarded (2) states, continuous time gives

.. math::

A_r &= A_{11} - A_{12} A_{22}^{-1} A_{21}, &
B_r &= B_1 - A_{12} A_{22}^{-1} B_2, \\
C_r &= C_1 - C_2 A_{22}^{-1} A_{21}, &
D_r &= D - C_2 A_{22}^{-1} B_2,

and discrete time replaces A_{22}^{-1} by -(I - A_{22})^{-1}.

Selection arguments match :func:modal_truncation.

Source code in jaxonomy/library/rom/linear_mor.py

def residualize(sys, order=None, keep=None):
    r"""Singular-perturbation (residualization) reduction.

    Like :func:`modal_truncation`, but instead of deleting the fast modes it
    sets their derivative (continuous) or their increment (discrete) to zero
    and solves for their quasi-steady value, folding it back into the
    retained model.  This matches the DC gain of the discarded modes
    (Kokotović, Khalil & O'Reilly 1986).

    Partitioning the modal realization into retained ``(1)`` and discarded
    ``(2)`` states, continuous time gives

    .. math::

        A_r &= A_{11} - A_{12} A_{22}^{-1} A_{21}, &
        B_r &= B_1 - A_{12} A_{22}^{-1} B_2, \\
        C_r &= C_1 - C_2 A_{22}^{-1} A_{21}, &
        D_r &= D - C_2 A_{22}^{-1} B_2,

    and discrete time replaces ``A_{22}^{-1}`` by ``-(I - A_{22})^{-1}``.

    Selection arguments match :func:`modal_truncation`.
    """
    A, B, C, D = _abcd(sys)
    mf = _real_modal_form(A, B, C)
    n = A.shape[0]
    keep_mask = _select_blocks(mf.blocks, sys.dt, order, keep, n)
    drop_mask = ~keep_mask

    Am, Bm, Cm = mf.A, mf.B, mf.C
    if not drop_mask.any():
        return _make_reduced(sys, Am, Bm, Cm, D)

    A11 = Am[np.ix_(keep_mask, keep_mask)]
    A12 = Am[np.ix_(keep_mask, drop_mask)]
    A21 = Am[np.ix_(drop_mask, keep_mask)]
    A22 = Am[np.ix_(drop_mask, drop_mask)]
    B1 = Bm[keep_mask, :]
    B2 = Bm[drop_mask, :]
    C1 = Cm[:, keep_mask]
    C2 = Cm[:, drop_mask]

    if sys.dt is None:
        M = np.linalg.solve(A22, np.hstack([A21, B2]))  # A22^{-1} [A21  B2]
    else:
        ImA22 = np.eye(A22.shape[0]) - A22
        # discrete steady state uses -(I - A22)^{-1}
        M = -np.linalg.solve(ImA22, np.hstack([A21, B2]))
    k = A21.shape[1]
    MA = M[:, :k]
    MB = M[:, k:]

    Ar = A11 - A12 @ MA
    Br = B1 - A12 @ MB
    Cr = C1 - C2 @ MA
    Dr = D - C2 @ MB
    return _make_reduced(sys, Ar, Br, Cr, Dr)

retained_energy(singular_values, r)

Fraction of total energy captured by the first r POD modes.

Energy is measured in squared singular values, Σ_{i<r} σ_i² / Σ_i σ_i² — monotonically non-decreasing in r.

Source code in jaxonomy/library/rom/snapshots.py

def retained_energy(singular_values, r: int) -> float:
    """Fraction of total energy captured by the first ``r`` POD modes.

    Energy is measured in squared singular values,
    ``Σ_{i<r} σ_i² / Σ_i σ_i²`` — monotonically non-decreasing in ``r``.
    """
    s = np.asarray(singular_values, dtype=float)
    total = float(np.sum(s**2))
    if total == 0.0:
        return 0.0
    r = int(max(0, min(r, s.shape[0])))
    return float(np.sum(s[:r] ** 2) / total)

soft_dead_zone(u, half_range, sharpness=10.0)

Smooth (differentiable) dead-zone gate.

Approximates the hard dead-zone where(|u| < half_range, 0, u) used by :class:DeadZone(mode="hard") with a sigmoid-blended kernel so gradients flow through the band. The blend factor is sigmoid(sharpness * (|u| - half_range)):

::

gate = 0.5 * (1.0 + tanh(sharpness * (|u| - half_range) / half_range))
y    = u * gate

Properties

• y(0) = 0 exactly.
• gate -> 1 outside the band, so y -> u for |u| >> half_range.
• gate -> 0 inside the band, so y -> 0 for |u| << half_range.
• Continuous everywhere; finite gradient even inside the band (this is the whole reason it exists).
• As sharpness -> inf the function converges to the hard gate.

Parameters:

Name	Type	Description	Default
u		Input array.	required
half_range		Positive scalar; band half-width.	required
sharpness		Positive scalar; default 10.0. Larger values give a tighter approximation to the hard dead zone.	10.0

Returns:

Type	Description
	Smoothly gated array, same shape as u.

Source code in jaxonomy/library/nonlinearities.py

def soft_dead_zone(u, half_range, sharpness=10.0):
    """Smooth (differentiable) dead-zone gate.

    Approximates the hard dead-zone ``where(|u| < half_range, 0, u)``
    used by :class:`DeadZone(mode="hard")` with a sigmoid-blended kernel
    so gradients flow through the band.  The blend factor is
    ``sigmoid(sharpness * (|u| - half_range))``:

    ::

        gate = 0.5 * (1.0 + tanh(sharpness * (|u| - half_range) / half_range))
        y    = u * gate

    Properties:
      * ``y(0) = 0`` exactly.
      * ``gate -> 1`` outside the band, so ``y -> u`` for ``|u| >> half_range``.
      * ``gate -> 0`` inside the band, so ``y -> 0`` for ``|u| << half_range``.
      * Continuous everywhere; finite gradient even inside the band
        (this is the whole reason it exists).
      * As ``sharpness -> inf`` the function converges to the hard gate.

    Args:
        u: Input array.
        half_range: Positive scalar; band half-width.
        sharpness: Positive scalar; default ``10.0``. Larger values give
            a tighter approximation to the hard dead zone.

    Returns:
        Smoothly gated array, same shape as ``u``.
    """
    # Normalise by ``half_range`` so the blend transition lives near
    # ``|u| = half_range`` regardless of band width; this matches the
    # convention used by :func:`soft_saturate`.
    gate = 0.5 * (1.0 + npa.tanh(sharpness * (npa.abs(u) - half_range) / half_range))
    return u * gate

soft_saturate(u, lower, upper, sharpness=10.0)

Smooth (differentiable) saturation between lower and upper.

Approximates npa.clip(u, lower, upper) with a tanh-based smooth clamp (per the T-115 spec):

::

mid  = (lower + upper) / 2
span = upper - lower
y    = mid + (span / 2) * tanh(sharpness * (u - mid) / span)

With sharpness = 10.0 and a unit span this gives ~99% saturation by |u - mid| = span, which is the design default.

Properties

• As sharpness -> inf the function converges to a hard clip.
• Strictly monotonically increasing in u (analytically; in finite precision the slope underflows to zero very far from mid because tanh saturates exponentially).
• Has a positive derivative across and inside the bound region -- gradients flow through saturation, which is the whole reason this exists.
• Requires finite lower / upper; for unbounded sides use the hard :class:Saturate block.

Parameters:

Name	Type	Description	Default
u		Input array.	required
lower		Lower limit (scalar or broadcastable array).	required
upper		Upper limit (scalar or broadcastable array).	required
sharpness		Positive scalar; default 10.0. Larger values give a tighter approximation to clip but smaller (and faster vanishing) gradients outside the bounds.	10.0

Returns:

Type	Description
	Smoothly saturated array, same shape as u.

Source code in jaxonomy/library/nonlinearities.py

def soft_saturate(u, lower, upper, sharpness=10.0):
    """Smooth (differentiable) saturation between ``lower`` and ``upper``.

    Approximates ``npa.clip(u, lower, upper)`` with a tanh-based smooth
    clamp (per the T-115 spec):

    ::

        mid  = (lower + upper) / 2
        span = upper - lower
        y    = mid + (span / 2) * tanh(sharpness * (u - mid) / span)

    With ``sharpness = 10.0`` and a unit span this gives ~99% saturation
    by ``|u - mid| = span``, which is the design default.

    Properties:
      * As ``sharpness -> inf`` the function converges to a hard ``clip``.
      * Strictly monotonically increasing in ``u`` (analytically; in
        finite precision the slope underflows to zero very far from
        ``mid`` because tanh saturates exponentially).
      * Has a positive derivative across and inside the bound region --
        gradients flow through saturation, which is the whole reason
        this exists.
      * Requires *finite* ``lower`` / ``upper``; for unbounded sides use
        the hard :class:`Saturate` block.

    Args:
        u: Input array.
        lower: Lower limit (scalar or broadcastable array).
        upper: Upper limit (scalar or broadcastable array).
        sharpness: Positive scalar; default ``10.0``. Larger values give
            a tighter approximation to ``clip`` but smaller (and faster
            vanishing) gradients outside the bounds.

    Returns:
        Smoothly saturated array, same shape as ``u``.
    """
    mid = (lower + upper) / 2.0
    span = upper - lower
    return mid + (span / 2.0) * npa.tanh(sharpness * (u - mid) / span)

step_response(linsys, t_grid)

Closed-form step response of a continuous-time LTI system.

For zero initial state and unit step input u(t) = 1 (for t ≥ 0) the response is

.. code-block:: text

y(t) = C · ∫₀ᵗ expm(A·s) ds · B  +  D·1

The matrix integral is computed via the augmented-matrix expm trick (see :func:_augmented_step_block) so the routine works correctly for non-invertible A (e.g. integrators). When A is invertible the same value equals C·A⁻¹·(expm(A·t) − I)·B + D.

Parameters:

Name	Type	Description	Default
linsys	LinearizedSystem	A :class:LinearizedSystem.	required
t_grid		Scalar or 1-D array of evaluation times. Negative times are evaluated formally (the closed-form result is still well defined; physically the step starts at t = 0).	required

Returns:

Type	Description
	Array of shape (K, p, m) — step response at each t for
	each (output, input) pair, where K = len(t_grid),
	p = n_outputs, m = n_inputs. If t_grid is a scalar
	the returned shape is (p, m). For SISO systems with
	scalar t_grid the result squeezes naturally to a scalar.

Notes

Fully differentiable through A, B, C, D via :func:jax.scipy.linalg.expm. For very large state dimensions (n > 50) the augmented expm may be slow — the honest fall- back is to simulate the diagram with a :class:Step source (deferred to a deeper follow-up).

Source code in jaxonomy/library/linearization_workflow.py

def step_response(linsys: LinearizedSystem, t_grid):
    """Closed-form step response of a continuous-time LTI system.

    For zero initial state and unit step input ``u(t) = 1`` (for ``t ≥ 0``)
    the response is

    .. code-block:: text

        y(t) = C · ∫₀ᵗ expm(A·s) ds · B  +  D·1

    The matrix integral is computed via the augmented-matrix expm trick
    (see :func:`_augmented_step_block`) so the routine works correctly
    for non-invertible ``A`` (e.g. integrators).  When ``A`` is
    invertible the same value equals ``C·A⁻¹·(expm(A·t) − I)·B + D``.

    Args:
        linsys: A :class:`LinearizedSystem`.
        t_grid: Scalar or 1-D array of evaluation times.  Negative times
            are evaluated formally (the closed-form result is still well
            defined; physically the step starts at ``t = 0``).

    Returns:
        Array of shape ``(K, p, m)`` — step response at each ``t`` for
        each ``(output, input)`` pair, where ``K = len(t_grid)``,
        ``p = n_outputs``, ``m = n_inputs``.  If ``t_grid`` is a scalar
        the returned shape is ``(p, m)``.  For SISO systems with
        scalar ``t_grid`` the result squeezes naturally to a scalar.

    Notes:
        Fully differentiable through ``A, B, C, D`` via
        :func:`jax.scipy.linalg.expm`.  For very large state dimensions
        (``n > 50``) the augmented expm may be slow — the honest fall-
        back is to simulate the diagram with a :class:`Step` source
        (deferred to a deeper follow-up).
    """
    if linsys.is_discrete():
        raise ValueError(
            "step_response is defined for continuous-time LinearizedSystems "
            f"only, but got a discrete-time system (dt={linsys.dt}). The "
            "closed-form continuous matrix-exponential formula does not "
            "apply to a discrete recurrence x[k+1] = A x[k] + B u[k]. "
            "Simulate the discrete diagram directly to obtain its step "
            "response (see KNOWN_GAPS.md)."
        )

    A, B, C, D, n, m, p = _coerce_state_space(linsys)

    t_arr = jnp.asarray(t_grid)
    scalar_input = (t_arr.ndim == 0)
    if scalar_input:
        t_arr = t_arr.reshape((1,))

    def one(t):
        block = _augmented_step_block(A, B, t)  # (n, m)
        return C @ block + D                    # (p, m)

    out = jax.vmap(one)(t_arr)  # (K, p, m)
    if scalar_input:
        return out[0]
    return out

with_observer(plant, observer, *, plant_u_port=0, plant_y_port=0, name='plant_with_observer')

Build a new diagram with observer attached to plant.

Wires the plant's control input through to both the plant and the observer; wires the plant's measurement output through to the observer; exports the observer's x_hat estimate as a top-level output of the augmented diagram. The plant's other output ports are not re-exported automatically — call sites that need them can wire them up in a parent diagram.

Parameters:

Name	Type	Description	Default
plant		A built :class:Diagram (or :class:SystemBase) representing the open-loop plant. Must expose at least one input port (for control u) and one output port (for measurement y).	required
observer		A built observer block (typically :class:jaxonomy.library.Luenberger or any block that accepts (u, y) inputs and produces x_hat as its first output port).	required
plant_u_port	int	Index of the plant input port carrying the control signal u. Defaults to 0.	0
plant_y_port	int	Index of the plant output port carrying the measurement y. Defaults to 0.	0
name	str	Name for the resulting wrapper diagram.	'plant_with_observer'

Returns:

Type	Description
	A new :class:Diagram containing both plant and
	observer wired as described, with:
	a single exported input port u (the control signal, fed to both the plant and the observer's u input),
	a single exported output port x_hat (the observer's state estimate).

Notes

The result is a "passive" instrumentation pattern — the observer reads (u, y) and emits an estimate, but the estimate is not fed back into the plant. To close a loop around the estimate (state-feedback control with observed state), build a controller subdiagram and wire its output back to plant.u in a parent diagram.

Source code in jaxonomy/library/linearization_workflow.py

def with_observer(
    plant,
    observer,
    *,
    plant_u_port: int = 0,
    plant_y_port: int = 0,
    name: str = "plant_with_observer",
):
    """Build a new diagram with ``observer`` attached to ``plant``.

    Wires the plant's control input through to both the plant and the
    observer; wires the plant's measurement output through to the
    observer; exports the observer's ``x_hat`` estimate as a top-level
    output of the augmented diagram. The plant's other output ports
    are not re-exported automatically — call sites that need them can
    wire them up in a parent diagram.

    Args:
        plant: A built :class:`Diagram` (or :class:`SystemBase`)
            representing the open-loop plant. Must expose at least one
            input port (for control ``u``) and one output port (for
            measurement ``y``).
        observer: A built observer block (typically
            :class:`jaxonomy.library.Luenberger` or any block that
            accepts ``(u, y)`` inputs and produces ``x_hat`` as its
            first output port).
        plant_u_port: Index of the plant input port carrying the
            control signal ``u``. Defaults to 0.
        plant_y_port: Index of the plant output port carrying the
            measurement ``y``. Defaults to 0.
        name: Name for the resulting wrapper diagram.

    Returns:
        A new :class:`Diagram` containing both ``plant`` and
        ``observer`` wired as described, with:

        * a single exported input port ``u`` (the control signal,
          fed to both the plant and the observer's ``u`` input),
        * a single exported output port ``x_hat`` (the observer's
          state estimate).

    Notes:
        The result is a "passive" instrumentation pattern — the
        observer reads ``(u, y)`` and emits an estimate, but the
        estimate is not fed back into the plant. To close a loop
        around the estimate (state-feedback control with observed
        state), build a controller subdiagram and wire its output
        back to ``plant.u`` in a parent diagram.
    """
    # Lazy import to avoid the framework→library cycle.
    from ..framework import DiagramBuilder
    from .routing import IOPort

    builder = DiagramBuilder()
    p = builder.add(plant)
    o = builder.add(observer)

    # The control signal needs to fan out to both the plant's u-input
    # and the observer's u-input. ``DiagramBuilder.export_input`` only
    # maps an exported port to one destination, so we use an IOPort
    # passthrough as the fan-out node: external "u" → IOPort → both.
    u_router = builder.add(IOPort(name="u_router"))
    builder.connect(u_router.output_ports[0], p.input_ports[plant_u_port])
    builder.connect(u_router.output_ports[0], o.input_ports[0])
    builder.connect(p.output_ports[plant_y_port], o.input_ports[1])

    builder.export_input(u_router.input_ports[0], name="u")
    builder.export_output(o.output_ports[0], name="x_hat")
    return builder.build(name=name)

write_model_description(diagram, path, *, model_name=None, guid=None, description=None)

Write a Jaxonomy diagram's FMI 2.0 modelDescription.xml to disk.

Parameters:

Name	Type	Description	Default
diagram	'Diagram'	Diagram to export.	required
path	str	Output file path.	required
model_name	str | None	Defaults to diagram.name.	None
guid	str | None	Optional GUID.	None
description	str | None	Optional free-form description.	None

Returns:

Type	Description
str	The same path argument (for chaining convenience).

Source code in jaxonomy/library/fmu_export.py

def write_model_description(
    diagram: "Diagram",
    path: str,
    *,
    model_name: str | None = None,
    guid: str | None = None,
    description: str | None = None,
) -> str:
    """Write a Jaxonomy diagram's FMI 2.0 modelDescription.xml to disk.

    Args:
        diagram: Diagram to export.
        path: Output file path.
        model_name: Defaults to ``diagram.name``.
        guid: Optional GUID.
        description: Optional free-form description.

    Returns:
        The same ``path`` argument (for chaining convenience).
    """
    if model_name is None:
        model_name = getattr(diagram, "name", "JaxonomyModel")
    xml = model_description_xml(
        diagram,
        model_name=model_name,
        guid=guid,
        description=description or f"FMI 2.0 export of {model_name}",
    )
    os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(xml)
    return path