Optimization

jaxonomy.optimization

AutoTuner

PID autotuning (without a measurement filter) with constraints in the frequency domain.

Supports only SISO systems.

Supports only continuous-time plants (TODO: extend to discrete-time systems)

Parameters:

Name	Type	Description	Default
plant		LeafSystem or a Diagram. If plant is not an LTISystem, operating points x_op and u_op must be provided for linearization.	required
n		int, optional Filter coefficient for the continuous-time PID controller	100
sim_time		float, optional Simulation time for computation of the error metric	2.0
metric		str, optional Error metric to be minimized. Options are "IAE" and "IE" "IAE": Integral of the absolute error "IE": Integral of the error	'IAE'
x_op		np.ndarray, optional Operating point of state vector for linearization	None
u_op		np.ndarray, optional Operating point of control vector for linearization	None
pid_gains_0		list or Array, optional Initial guess for PID gains [kp, ki, kd]	[1.0, 10.0, 0.1]
pid_gains_upper_bounds		list or Array, optional Upper bounds for PID gains [kp, ki, kd]. Lower bounds are set to 0	None
Ms		float, optional Maximum sensitivity	100.0
Mt		float, optional Maximum complementary sensitivity	100.0
add_filter		bool, optional Add measurement filter (currently not implemented)	False
method		str, optional The method for optimization. Available options are: - "scipy-slsqp" - "scipy-cobyla" - "scipy-trust-constr" - "ipopt" - "nlopt-slsqp" - "nlopt-cobyla" - "nlopt-ld_mma" - "nlopt-isres" - "nlopt-ags" - "nlopt-direct"	'scipy-slsqp'

Notes:

The utilities plot_freq_response, plot_time_response, and plot_freq_and_time_responses can be used to visualize the frequency and time responses of the closed-loop system.

Post initialization the tune method should be called to obtain the optimal PID gains. See notebooks/opt_framework/pid_autotuning.ipynb for an example.

Source code in jaxonomy/optimization/pid_autotuning.py

class AutoTuner:
    """
    PID autotuning (without a measurement filter) with constraints in the frequency
    domain.

    Supports only SISO systems.

    Supports only continuous-time plants (TODO: extend to discrete-time systems)

    Parameters:
        plant: LeafSystem or a Diagram.
            If plant is not an LTISystem, operating points x_op and u_op must be
            provided for linearization.
        n: int, optional
            Filter coefficient for the continuous-time PID controller
        sim_time: float, optional
            Simulation time for computation of the error metric
        metric: str, optional
            Error metric to be minimized. Options are "IAE" and "IE"
                "IAE": Integral of the absolute error
                "IE": Integral of the error
        x_op: np.ndarray, optional
            Operating point of state vector for linearization
        u_op: np.ndarray, optional
            Operating point of control vector for linearization
        pid_gains_0: list or Array, optional
            Initial guess for PID gains [kp, ki, kd]
        pid_gains_upper_bounds: list or Array, optional
            Upper bounds for PID gains [kp, ki, kd]. Lower bounds are set to 0
        Ms: float, optional
            Maximum sensitivity
        Mt: float, optional
            Maximum complementary sensitivity
        add_filter: bool, optional
            Add measurement filter (currently not implemented)
        method: str, optional
            The method for optimization. Available options are:
                - "scipy-slsqp"
                - "scipy-cobyla"
                - "scipy-trust-constr"
                - "ipopt"
                - "nlopt-slsqp"
                - "nlopt-cobyla"
                - "nlopt-ld_mma"
                - "nlopt-isres"
                - "nlopt-ags"
                - "nlopt-direct"

    Notes:

    The utilities `plot_freq_response`, `plot_time_response`, and
    `plot_freq_and_time_responses` can be used to visualize the frequency and time
    responses of the closed-loop system.

    Post initialization the `tune` method should be called to obtain the optimal PID
    gains. See `notebooks/opt_framework/pid_autotuning.ipynb` for an example.

    """

    def __init__(
        self,
        plant,
        n=100,
        sim_time=2.0,
        metric="IAE",
        x_op=None,
        u_op=None,
        pid_gains_0=[1.0, 10.0, 0.1],
        pid_gains_upper_bounds=None,
        Ms=100.0,
        Mt=100.0,
        add_filter=False,  # NOTE: add measurement filter (currently not implemented)
        method="scipy-slsqp",
    ):
        if isinstance(plant, LTISystem):  # LTISystem includes TransferFunction
            linear_plant = plant
            linear_plant.create_context()
        else:
            if x_op is None or u_op is None:
                raise ValueError("Operating point x_op and u_op must be provided")

            _, linear_plant = linearize_plant(plant, x_op, u_op)

        if linear_plant.B.shape[1] != 1 or linear_plant.C.shape[0] != 1:
            raise ValueError("Plant must be SISO")

        self.A, self.B, self.C, self.D = (
            linear_plant.A,
            linear_plant.B,
            linear_plant.C,
            linear_plant.D,
        )

        self.pid, self.integrator, self.diagram = make_closed_loop_pid_system(
            linear_plant,
            metric,
            n=n,
        )

        self.lb = [0.0] * 3
        if pid_gains_upper_bounds is None:
            self.ub = [jnp.inf] * 3
        else:
            self.ub = pid_gains_upper_bounds

        self.base_context = self.diagram.create_context()

        self.n = n
        self.sim_time = sim_time
        self.metric = metric
        self.pid_gains_0 = pid_gains_0
        self.Ms = Ms
        self.Mt = Mt
        self.add_filter = add_filter
        self.method = method

        self.omega_grid = 10.0 ** jnp.linspace(-2, 2, 1000)
        # self.omega_grid = 10.0 ** jnp.linspace(-1, 2, 150)
        self.options = SimulatorOptions(
            enable_autodiff=True,
            max_major_step_length=0.01,  # rtol=1e-08, atol=1e-10
        )

        self.circle_constraint_vectorized = jax.vmap(
            self.circle_constraint_, in_axes=(None, None, None, 0, None, None)
        )  # Deprecated

        self.Ps_vectorized = jax.vmap(self.Ps, in_axes=0)
        self.Cs_vectorized = jax.vmap(self.Cs, in_axes=(None, None, None, 0))
        self.vec_absolute = jax.vmap(jnp.absolute)

    @partial(jax.jit, static_argnums=(0,))
    def objective(self, pid_params):
        kp, ki, kd = pid_params
        pid_subcontext = self.base_context[self.pid.system_id].with_parameters(
            {"kp": kp, "ki": ki, "kd": kd}
        )
        context = self.base_context.with_subcontext(self.pid.system_id, pid_subcontext)
        sol = jaxonomy.simulate(
            self.diagram, context, (0.0, self.sim_time), options=self.options
        )
        return self.integrator.output_ports[0].eval(sol.context) / self.sim_time

    @partial(jax.jit, static_argnums=(0,))
    def Ps(self, s):
        P = (
            self.C @ jnp.linalg.inv(s * jnp.eye(self.A.shape[0]) - self.A) @ self.B
            + self.D
        )
        return P[0, 0]

    @partial(jax.jit, static_argnums=(0,))
    def Cs(self, kp, ki, kd, s):
        return kp + ki / s + kd * s

    @partial(jax.jit, static_argnums=(0,))
    def circle_constraint_(self, kp, ki, kd, omega, c, r):
        """Deprecated: this is needed for `self.constraints_` which is deprecated
        and replaced by `self.constraints`.
        """
        s = omega * 1.0j
        L = self.Ps(s) * self.Cs(kp, ki, kd, s)
        return jnp.absolute(L - c) - r

    @partial(jax.jit, static_argnums=(0,))
    def constraints_(self, pid_params):
        """Deprecated: replaced by `self.constraints`"""
        kp, ki, kd = pid_params
        Ms, Mt = self.Ms, self.Mt
        g_Ms = self.circle_constraint_vectorized(
            kp, ki, kd, self.omega_grid, -1.0, 1.0 / Ms
        )
        g_Mt = self.circle_constraint_vectorized(
            kp, ki, kd, self.omega_grid, -(Mt**2) / (Mt**2 - 1.0), Mt / (Mt**2 - 1.0)
        )
        return jnp.array([jnp.min(g_Ms), jnp.min(g_Mt)])

    @partial(jax.jit, static_argnums=(0,))
    def constraints(self, pid_params):
        kp, ki, kd = pid_params
        S_grid = 1.0 / (
            1.0
            + self.Ps_vectorized(self.omega_grid * 1.0j)
            * self.Cs_vectorized(kp, ki, kd, self.omega_grid * 1.0j)
        )
        T_grid = 1.0 - S_grid

        S_grid = self.vec_absolute(S_grid)
        T_grid = self.vec_absolute(T_grid)

        return jnp.array([self.Ms - jnp.max(S_grid), self.Mt - jnp.max(T_grid)])

    def tune(self):
        x0 = jnp.array(self.pid_gains_0)
        bounds = list(zip(self.lb, self.ub))

        obj = jax.jit(self.objective)
        cons = jax.jit(self.constraints)

        obj_grad = jax.grad(self.objective)
        obj_hess = jax.jit(jax.hessian(self.objective))

        cons_jac = jax.jit(jax.jacfwd(self.constraints))

        print(f"Tuning with {self.method}")
        if self.method in SCIPY_METHODS:
            constraints_scipy = NonlinearConstraint(cons, 0.0, jnp.inf, jac=cons_jac)

            res = minimize(
                obj,
                x0,
                jac=obj_grad,
                method=SCIPY_METHODS[self.method],
                bounds=bounds,
                constraints=constraints_scipy,
                options={"maxiter": 100},
            )

        elif self.method == "ipopt":
            cons_hess = jax.hessian(self.constraints)
            cons_hess_vp = jax.jit(
                lambda x, v: jnp.sum(
                    # pylint: disable-next=not-callable
                    jnp.multiply(v[:, jnp.newaxis, jnp.newaxis], cons_hess(x)),
                    axis=0,
                )
            )

            constraints_ipopt = [
                {"type": "ineq", "fun": cons, "jac": cons_jac, "hess": cons_hess_vp}
            ]

            res = cyipopt.minimize_ipopt(
                obj,
                x0=x0,
                jac=obj_grad,
                hess=obj_hess,
                constraints=constraints_ipopt,
                bounds=bounds,
                options={
                    "max_iter": 500,
                    "disp": 5,
                },
            )

        elif self.method in NLOPT_METHODS:
            if self.method in NLOPT_METHODS_GLOBAL and any(
                ub == jnp.inf for ub in self.ub
            ):
                raise ValueError(
                    f"Method {self.method} requires finite upper bounds for all "
                    "parameters. Please specify `pid_gains_upper_bounds`."
                )

            # Define the objective function for nlopt
            def nlopt_obj(x, grad):
                if grad.size > 0:
                    grad[:] = obj_grad(jnp.array(x))
                # pylint: disable-next=not-callable
                return float(obj(jnp.array(x)))

            # Define the objective function for nlopt
            def nlopt_cons(result, x, grad):
                if grad.size > 0:
                    # pylint: disable-next=not-callable
                    grad[:, :] = -cons_jac(jnp.array(x))
                # pylint: disable-next=not-callable
                result[:] = -cons(jnp.array(x))

            # Initialize nlopt optimizer
            method = NLOPT_METHODS[self.method]()
            opt = nlopt.opt(method, len(x0))

            # Set the objective function
            opt.set_min_objective(nlopt_obj)

            # Set the constraints
            opt.add_inequality_mconstraint(nlopt_cons, [1e-6, 1e-06])

            # Set the bounds
            lower_bounds, upper_bounds = zip(*bounds)
            opt.set_lower_bounds(lower_bounds)
            opt.set_upper_bounds(upper_bounds)

            # Set stopping criteria
            opt.set_maxeval(500)
            opt.set_ftol_rel(1e-5)
            opt.set_xtol_rel(1e-6)
            opt.set_maxtime(30.0)

            # Run the optimization
            x_opt = opt.optimize(x0)
            print(f"{x_opt=}")
            minf = opt.last_optimum_value()

            nlopt_success_codes = {
                nlopt.SUCCESS: "SUCCESS",
                nlopt.STOPVAL_REACHED: "STOPVAL_REACHED",
                nlopt.FTOL_REACHED: "FTOL_REACHED",
                nlopt.XTOL_REACHED: "XTOL_REACHED",
                nlopt.MAXEVAL_REACHED: "MAXEVAL_REACHED",
                nlopt.MAXTIME_REACHED: "MAXTIME_REACHED",
            }

            nlopt_error_codes = {
                nlopt.FAILURE: "FAILURE",
                nlopt.INVALID_ARGS: "INVALID_ARGS",
                nlopt.OUT_OF_MEMORY: "OUT_OF_MEMORY",
                nlopt.ROUNDOFF_LIMITED: "ROUNDOFF_LIMITED",
                nlopt.FORCED_STOP: "FORCED_STOP",
            }

            nlopt_status_codes = {**nlopt_success_codes, **nlopt_error_codes}

            res = OptResults(
                x=x_opt,
                fun=minf,
                success=opt.last_optimize_result() in nlopt_success_codes,
                message=nlopt_status_codes[opt.last_optimize_result()],
            )

        else:
            raise ValueError("Invalid method")
        return res.x, res

    def plot_freq_response(
        self, pid_params, plant_tf_num, plant_tf_den, Ms=None, Mt=None
    ):
        if Ms is None:
            Ms = self.Ms

        if Mt is None:
            Mt = self.Mt

        kp, ki, kd = pid_params
        Cs = ct.TransferFunction([kd, kp, ki], [1, 0], name="PID")
        Ps = ct.TransferFunction(plant_tf_num, plant_tf_den, name="Plant")

        # Plot Gang of Four transfer functions
        fig1 = plt.figure()
        ct.gangof4_plot(Ps, Cs, omega=self.omega_grid)

        axs = fig1.get_axes()

        axs[3].set_title(r"$T = \dfrac{PC}{1+PC}$")
        axs[1].set_title(r"$PS = \dfrac{P}{1+PC}$")
        axs[2].set_title(r"$CS = \dfrac{C}{1+PC}$")
        axs[0].set_title(r"$S = \dfrac{1}{1+PC}$")

        if Ms is not None:
            axs[0].hlines(
                Ms,
                self.omega_grid.min(),
                self.omega_grid.max(),
                colors="r",
                linestyles="--",
            )

        if Mt is not None:
            axs[3].hlines(
                Mt,
                self.omega_grid.min(),
                self.omega_grid.max(),
                colors="b",
                linestyles="--",
            )

        # Set x-axis labels for the bottom plots
        axs[2].set_xlabel("Frequency (rad/sec)")
        axs[3].set_xlabel("Frequency (rad/sec)")

        fig1.tight_layout()
        fig1.suptitle("Frequency domain response")

        # Plot Nyquist plot
        fig2 = plt.figure()
        ct.nyquist_plot(
            Ps * Cs, omega=self.omega_grid, warn_nyquist=False, warn_encirclements=False
        )
        axs = fig2.get_axes()
        ax = axs[0]

        def gen_circle_points(c, r):
            t = jnp.linspace(-jnp.pi / 2, jnp.pi / 2, 100)
            return jnp.array([c + r * jnp.cos(t), r * jnp.sin(t)])

        if Ms is not None:
            c1, r1 = -1.0, 1.0 / Ms
            xc1, yc1 = gen_circle_points(c1, r1)
            ax.plot(xc1, yc1, "r--")

        if Mt is not None:
            c2, r2 = -(Mt**2) / (Mt**2 - 1.0), Mt / (Mt**2 - 1.0)
            xc2, yc2 = gen_circle_points(c2, r2)
            ax.plot(xc2, yc2, "b--")

        ax.set_title("Nyquist Plot")
        fig2.tight_layout()

        return fig1, fig2

    def plot_time_response(self, pid_params):
        kp, ki, kd = pid_params
        pid_subcontext = self.base_context[self.pid.system_id].with_parameters(
            {"kp": kp, "ki": ki, "kd": kd}
        )
        context = self.base_context.with_subcontext(self.pid.system_id, pid_subcontext)

        recorded_signals = {
            "objective": self.diagram["integrator"].output_ports[0],
            "ref": self.diagram["ref"].output_ports[0],
            "plant": self.diagram.output_ports[0],
            "pid": self.diagram["pid"].output_ports[0],
        }

        sol = jaxonomy.simulate(
            self.diagram,
            context,
            (0.0, self.sim_time),
            recorded_signals=recorded_signals,
        )

        fig, (ax1, ax2, ax3) = plt.subplots(3, 1)

        ax1.plot(sol.time, sol.outputs["plant"], label=r"plant: $y$")
        ax1.plot(sol.time, sol.outputs["ref"], label=r"reference: $y_r$")
        ax2.plot(
            sol.time,
            sol.outputs["objective"] / self.sim_time,
            label=f"objective: {self.metric}",
        )
        ax3.plot(sol.time, sol.outputs["pid"], label=r"pid-control: $u$")

        ax3.set_xlabel("Time (s)")
        for ax in (ax1, ax2, ax3):
            ax.legend()

        fig.suptitle("Time domain response")
        fig.tight_layout()

        print(
            f"objective = "
            f"{self.integrator.output_ports[0].eval(sol.context)/self.sim_time}"
        )
        return fig

    def plot_freq_and_time_responses(
        self, pid_params, plant_tf_num, plant_tf_den, Ms=None, Mt=None
    ):
        fig1, fig2 = self.plot_freq_response(
            pid_params, plant_tf_num, plant_tf_den, Ms, Mt
        )
        fig3 = self.plot_time_response(pid_params)
        return fig1, fig2, fig3

circle_constraint_(kp, ki, kd, omega, c, r)

Deprecated: this is needed for self.constraints_ which is deprecated and replaced by self.constraints.

Source code in jaxonomy/optimization/pid_autotuning.py

@partial(jax.jit, static_argnums=(0,))
def circle_constraint_(self, kp, ki, kd, omega, c, r):
    """Deprecated: this is needed for `self.constraints_` which is deprecated
    and replaced by `self.constraints`.
    """
    s = omega * 1.0j
    L = self.Ps(s) * self.Cs(kp, ki, kd, s)
    return jnp.absolute(L - c) - r

constraints_(pid_params)

Deprecated: replaced by self.constraints

Source code in jaxonomy/optimization/pid_autotuning.py

@partial(jax.jit, static_argnums=(0,))
def constraints_(self, pid_params):
    """Deprecated: replaced by `self.constraints`"""
    kp, ki, kd = pid_params
    Ms, Mt = self.Ms, self.Mt
    g_Ms = self.circle_constraint_vectorized(
        kp, ki, kd, self.omega_grid, -1.0, 1.0 / Ms
    )
    g_Mt = self.circle_constraint_vectorized(
        kp, ki, kd, self.omega_grid, -(Mt**2) / (Mt**2 - 1.0), Mt / (Mt**2 - 1.0)
    )
    return jnp.array([jnp.min(g_Ms), jnp.min(g_Mt)])

CompositeTransform

Bases: Transform

A composite transformation that applies a list of transformations in sequence.

Source code in jaxonomy/optimization/framework/base/transformations.py

class CompositeTransform(Transform):
    """
    A composite transformation that applies a list of transformations in sequence.
    """

    def __init__(self, transformations):
        self.transformations = transformations

    def transform(self, params: dict):
        for transformation in self.transformations:
            params = transformation.transform(params)
        return params

    def inverse_transform(self, params: dict):
        for transformation in reversed(self.transformations):
            params = transformation.inverse_transform(params)
        return params

ConfidenceIntervalResult dataclass

Confidence intervals and covariance matrix from the Laplace approximation.

All matrix/array attributes are plain numpy.ndarray for easy inspection and serialisation.

Attributes

param_names : list[str] Flat parameter names (array params expanded to "theta[0]", etc.). opt_params : dict Optimised parameter values in the original (un-transformed) space. covariance : ndarray, shape (n, n) Estimated parameter covariance matrix. correlation : ndarray, shape (n, n) Correlation matrix (covariance normalised by marginal standard deviations). standard_errors : ndarray, shape (n,) Marginal standard deviations sqrt(diag(covariance)). confidence_intervals : dict[str, tuple[float, float]] Per-parameter (lower, upper) bounds in the original space. Keys match param_names. confidence_level : float Nominal confidence level (e.g. 0.95 for 95 %). z_score : float Standard-normal quantile corresponding to confidence_level. hessian : ndarray, shape (n, n) Hessian of the objective evaluated at the optimum, in the (possibly transformed) optimisation space. hessian_eigenvalues : ndarray, shape (n,) Eigenvalues of the Hessian (ascending). hessian_condition_number : float Ratio max|λ| / min|λ|. Large values (> 1 000) signal near-collinear parameters or an ill-conditioned problem. is_positive_definite : bool True when the Hessian was positive definite at the supplied point (necessary condition for a true local minimum). residual_variance : float or None Residual variance σ² used to scale the covariance. None when n_data was not provided (pure MLE / default). n_data : int or None Number of observations used (for least-squares scaling). objective_value : float Loss at the optimum. hessian_method : str How the Hessian was computed: "AD" (automatic differentiation), "FD" (finite differences), "provided", or "failed". message : str Any warnings raised during computation (empty when all is well).

Source code in jaxonomy/optimization/confidence.py

@dataclass
class ConfidenceIntervalResult:
    """
    Confidence intervals and covariance matrix from the Laplace approximation.

    All matrix/array attributes are plain ``numpy.ndarray`` for easy inspection
    and serialisation.

    Attributes
    ----------
    param_names : list[str]
        Flat parameter names (array params expanded to ``"theta[0]"``, etc.).
    opt_params : dict
        Optimised parameter values in the **original** (un-transformed) space.
    covariance : ndarray, shape (n, n)
        Estimated parameter covariance matrix.
    correlation : ndarray, shape (n, n)
        Correlation matrix (covariance normalised by marginal standard deviations).
    standard_errors : ndarray, shape (n,)
        Marginal standard deviations ``sqrt(diag(covariance))``.
    confidence_intervals : dict[str, tuple[float, float]]
        Per-parameter ``(lower, upper)`` bounds in the **original** space.
        Keys match ``param_names``.
    confidence_level : float
        Nominal confidence level (e.g. ``0.95`` for 95 %).
    z_score : float
        Standard-normal quantile corresponding to ``confidence_level``.
    hessian : ndarray, shape (n, n)
        Hessian of the objective evaluated at the optimum, in the (possibly
        transformed) optimisation space.
    hessian_eigenvalues : ndarray, shape (n,)
        Eigenvalues of the Hessian (ascending).
    hessian_condition_number : float
        Ratio max|λ| / min|λ|.  Large values (> 1 000) signal near-collinear
        parameters or an ill-conditioned problem.
    is_positive_definite : bool
        ``True`` when the Hessian was positive definite at the supplied point
        (necessary condition for a true local minimum).
    residual_variance : float or None
        Residual variance ``σ²`` used to scale the covariance.  ``None`` when
        ``n_data`` was not provided (pure MLE / default).
    n_data : int or None
        Number of observations used (for least-squares scaling).
    objective_value : float
        Loss at the optimum.
    hessian_method : str
        How the Hessian was computed: ``"AD"`` (automatic differentiation),
        ``"FD"`` (finite differences), ``"provided"``, or ``"failed"``.
    message : str
        Any warnings raised during computation (empty when all is well).
    """

    param_names: list[str]
    opt_params: dict[str, Any]
    covariance: np.ndarray
    correlation: np.ndarray
    standard_errors: np.ndarray
    confidence_intervals: dict[str, tuple[float, float]]
    confidence_level: float
    z_score: float
    hessian: np.ndarray
    hessian_eigenvalues: np.ndarray
    hessian_condition_number: float
    is_positive_definite: bool
    residual_variance: float | None
    n_data: int | None
    objective_value: float
    hessian_method: str
    message: str

    # ------------------------------------------------------------------
    # Convenience accessors
    # ------------------------------------------------------------------

    def interval(self, param_name: str) -> tuple[float, float]:
        """Return ``(lower, upper)`` for a single parameter by name.

        Raises ``KeyError`` if the name is not found.  For array parameters
        use the expanded name, e.g. ``ci.interval("theta[0]")``.
        """
        return self.confidence_intervals[param_name]

    def contains(self, param_name: str, value: float) -> bool:
        """Return ``True`` when *value* lies within the CI for *param_name*."""
        lo, hi = self.interval(param_name)
        return lo <= value <= hi

    # ------------------------------------------------------------------
    # Display
    # ------------------------------------------------------------------

    def summary(self) -> str:
        """Return a formatted human-readable summary table."""
        level_pct = self.confidence_level * 100
        lines = [
            f"=== Parameter Confidence Intervals ({level_pct:.1f}%) "
            f"[Laplace approximation] ===",
            f"Objective at optimum       : {self.objective_value:.6g}",
            f"Hessian computation method : {self.hessian_method}",
            f"Hessian positive definite  : {'yes' if self.is_positive_definite else 'NO ← not at a true minimum'}",
            f"Hessian condition number   : {self.hessian_condition_number:.3g}",
        ]
        if self.residual_variance is not None:
            lines.append(f"Residual variance σ²       : {self.residual_variance:.4g}  "
                         f"(n_data={self.n_data})")
        if self.message:
            lines.append(f"⚠  {self.message}")
        lines += [
            "",
            f"{'Parameter':<24} {'Opt. value':>14} {'Std. error':>13} "
            f"  {level_pct:.1f}% CI",
            "-" * 72,
        ]
        for name in self.param_names:
            lo, hi = self.confidence_intervals[name]
            # Retrieve the optimal value (may be multi-element param)
            se_idx = self.param_names.index(name)
            se = self.standard_errors[se_idx]
            # Optimal value in original space
            opt_val = self._opt_val_for(name)
            lines.append(
                f"{name:<24} {opt_val:>14.6g} {se:>13.4e}  "
                f"[{lo:>12.6g}, {hi:>12.6g}]"
            )
        return "\n".join(lines)

    def _opt_val_for(self, flat_name: str) -> float:
        """Extract the scalar optimal value for a flat parameter name."""
        # scalar param
        if flat_name in self.opt_params:
            return float(np.asarray(self.opt_params[flat_name]).ravel()[0])
        # vector param: name looks like "theta[0]" or "theta[0,1]"
        bracket = flat_name.find("[")
        if bracket != -1:
            key = flat_name[:bracket]
            idx_str = flat_name[bracket + 1 : flat_name.find("]")]
            idx = tuple(int(s) for s in idx_str.split(","))
            val = np.asarray(self.opt_params.get(key, np.nan))
            try:
                return float(val[idx])
            except (IndexError, TypeError):
                pass
        return float("nan")

    def __repr__(self) -> str:
        return self.summary()

contains(param_name, value)

Return True when value lies within the CI for param_name.

Source code in jaxonomy/optimization/confidence.py

def contains(self, param_name: str, value: float) -> bool:
    """Return ``True`` when *value* lies within the CI for *param_name*."""
    lo, hi = self.interval(param_name)
    return lo <= value <= hi

interval(param_name)

Return (lower, upper) for a single parameter by name.

Raises KeyError if the name is not found. For array parameters use the expanded name, e.g. ci.interval("theta[0]").

Source code in jaxonomy/optimization/confidence.py

def interval(self, param_name: str) -> tuple[float, float]:
    """Return ``(lower, upper)`` for a single parameter by name.

    Raises ``KeyError`` if the name is not found.  For array parameters
    use the expanded name, e.g. ``ci.interval("theta[0]")``.
    """
    return self.confidence_intervals[param_name]

summary()

Return a formatted human-readable summary table.

Source code in jaxonomy/optimization/confidence.py

def summary(self) -> str:
    """Return a formatted human-readable summary table."""
    level_pct = self.confidence_level * 100
    lines = [
        f"=== Parameter Confidence Intervals ({level_pct:.1f}%) "
        f"[Laplace approximation] ===",
        f"Objective at optimum       : {self.objective_value:.6g}",
        f"Hessian computation method : {self.hessian_method}",
        f"Hessian positive definite  : {'yes' if self.is_positive_definite else 'NO ← not at a true minimum'}",
        f"Hessian condition number   : {self.hessian_condition_number:.3g}",
    ]
    if self.residual_variance is not None:
        lines.append(f"Residual variance σ²       : {self.residual_variance:.4g}  "
                     f"(n_data={self.n_data})")
    if self.message:
        lines.append(f"⚠  {self.message}")
    lines += [
        "",
        f"{'Parameter':<24} {'Opt. value':>14} {'Std. error':>13} "
        f"  {level_pct:.1f}% CI",
        "-" * 72,
    ]
    for name in self.param_names:
        lo, hi = self.confidence_intervals[name]
        # Retrieve the optimal value (may be multi-element param)
        se_idx = self.param_names.index(name)
        se = self.standard_errors[se_idx]
        # Optimal value in original space
        opt_val = self._opt_val_for(name)
        lines.append(
            f"{name:<24} {opt_val:>14.6g} {se:>13.4e}  "
            f"[{lo:>12.6g}, {hi:>12.6g}]"
        )
    return "\n".join(lines)

DistributionConfig dataclass

Structure of attributes for specifying distributions for stochastic variables

Source code in jaxonomy/optimization/framework/base/optimizable.py

@dataclass
class DistributionConfig:
    """
    Structure of attributes for specifying distributions for stochastic variables
    """

    names: list[str]
    shapes: list[tuple]
    distributions: list[str]
    distributions_configs: list[dict]

Evosax

Bases: Optimizer

Population based global optimizers from Evosax.

Parameters:

Name	Type	Description	Default
optimizable	Optimizable	The optimizable object.	required
opt_method	str	The optimization method to use. See evosax.Strategies for available methods.	'CMA_ES'
opt_method_config	dict	Configuration for the optimization method.	None
pop_size	int	The population size.	10
num_generations	int	The number of generations.	100
print_every	int	Print progress every print_every generations.	1
metrics_writer	MetricsWriter | None	Optional CSV file to write metrics to.	None
seed	int	The random seed.	None

Source code in jaxonomy/optimization/framework/optimizers_evosax.py

class Evosax(Optimizer):
    """
    Population based global optimizers from Evosax.

    Parameters:
        optimizable (Optimizable):
            The optimizable object.
        opt_method (str):
            The optimization method to use. See `evosax.Strategies` for
            available methods.
        opt_method_config (dict):
            Configuration for the optimization method.
        pop_size (int):
            The population size.
        num_generations (int):
            The number of generations.
        print_every (int):
            Print progress every `print_every` generations.
        metrics_writer (MetricsWriter|None):
            Optional CSV file to write metrics to.
        seed (int):
            The random seed.
    """

    def __init__(
        self,
        optimizable: Optimizable,
        opt_method="CMA_ES",
        opt_method_config=None,
        pop_size=10,
        num_generations=100,
        print_every=1,
        seed=None,
        metrics_writer: MetricsWriter = None,
    ):
        self.optimizable = optimizable
        self.opt_method = opt_method
        self.pop_size = pop_size
        self.num_generations = num_generations
        self.print_every = print_every
        self.metrics_writer = metrics_writer
        self.optimal_params = None
        self._fitness_history: list[float] = []

        self.num_dims = optimizable.params_0_flat.size

        if self.optimizable.has_constraints:
            raise ValueError(
                f"Optimization method evosax:{self.opt_method} "
                "does not support constraints."
            )

        if opt_method not in evosax.Strategies:
            raise ValueError(f"Unknown optimization method: {opt_method}")

        if opt_method_config is None:
            opt_method_config = {}

        self.strategy = evosax.Strategies[opt_method](self.pop_size, self.num_dims)
        self.es_params = self.strategy.default_params.replace(**opt_method_config)

        # Create bounds
        if optimizable.bounds_flat is not None:
            lower_bounds, upper_bounds = zip(*optimizable.bounds_flat)
            lb = jnp.array(lower_bounds)
            ub = jnp.array(upper_bounds)
            self.es_params = self.es_params.replace(clip_min=lb, clip_max=ub)

        # Create initialization bounds
        if optimizable.init_min_max_flat is not None:
            init_min, init_max = zip(*optimizable.init_min_max_flat)
            imin = jnp.array(init_min)
            imax = jnp.array(init_max)
            self.es_params = self.es_params.replace(init_min=imin, init_max=imax)

        else:
            # if bounds are specified, unless they are infinity, we can use
            # them for initialization. If infinity, we initialize in [-0.1,0.1]
            if optimizable.bounds_flat is not None:
                bounds = [
                    (
                        -0.1 if b[0] == -jnp.inf else b[0],
                        0.1 if b[1] == jnp.inf else b[1],
                    )
                    for b in optimizable.bounds_flat
                ]
                lower_bounds, upper_bounds = zip(*bounds)
                lb = jnp.array(lower_bounds)
                ub = jnp.array(upper_bounds)
                self.es_params = self.es_params.replace(init_min=lb, init_max=ub)

            # if strategy defaults are not zero, they are likely set to sensible values,
            # so we use them, otherwise we scale the initial params by a factor of 10
            elif self.es_params.init_min == 0 and self.es_params.init_max == 0:
                factor = 10.0
                imin = jnp.full(self.num_dims, self.optimizable.params_0_flat / factor)
                imax = jnp.full(self.num_dims, self.optimizable.params_0_flat * factor)
                self.es_params = self.es_params.replace(init_min=imin, init_max=imax)

        self.key = jr.PRNGKey(
            np.random.randint(0, 2**32, dtype=np.int64) if seed is None else seed
        )

    def optimize(self):
        """Run optimization"""
        fitness_func = jax.jit(self.optimizable.batched_objective_flat)

        state = self.strategy.initialize(self.key, self.es_params)

        # https://github.com/RobertTLange/evosax/issues/45
        state = state.replace(best_fitness=jnp.finfo(jnp.float64).max)

        for gen in range(self.num_generations):
            self.key, subkey = jr.split(self.key)
            x, state = self.strategy.ask(subkey, state, self.es_params)
            fitness = fitness_func(x)
            state = self.strategy.tell(x, fitness, state, self.es_params)
            self._fitness_history.append(float(state.best_fitness))

            if self.print_every is not None and (gen + 1) % self.print_every == 0:
                logger.info(
                    "# Gen: %3d|Fitness: %.6f|Params: %s",
                    gen + 1,
                    state.best_fitness,
                    state.best_member,
                )
            if self.metrics_writer is not None:
                self.metrics_writer.write_metrics(best_fitness=state.best_fitness)

        params = state.best_member
        self.optimal_params = self.optimizable.unflatten_params(params)
        if self.optimizable.transformation is not None:
            self.optimal_params = self.optimizable.transformation.inverse_transform(
                self.optimal_params
            )
        return OptimizationResult(
            params=self.optimal_params,
            success=True,
            nit=self.num_generations,
            nfev=self.num_generations * self.pop_size,
            message=f"Completed {self.num_generations} generations.",
            final_loss=float(state.best_fitness),
            loss_history=list(self._fitness_history),
        )

optimize()

Run optimization

Source code in jaxonomy/optimization/framework/optimizers_evosax.py

def optimize(self):
    """Run optimization"""
    fitness_func = jax.jit(self.optimizable.batched_objective_flat)

    state = self.strategy.initialize(self.key, self.es_params)

    # https://github.com/RobertTLange/evosax/issues/45
    state = state.replace(best_fitness=jnp.finfo(jnp.float64).max)

    for gen in range(self.num_generations):
        self.key, subkey = jr.split(self.key)
        x, state = self.strategy.ask(subkey, state, self.es_params)
        fitness = fitness_func(x)
        state = self.strategy.tell(x, fitness, state, self.es_params)
        self._fitness_history.append(float(state.best_fitness))

        if self.print_every is not None and (gen + 1) % self.print_every == 0:
            logger.info(
                "# Gen: %3d|Fitness: %.6f|Params: %s",
                gen + 1,
                state.best_fitness,
                state.best_member,
            )
        if self.metrics_writer is not None:
            self.metrics_writer.write_metrics(best_fitness=state.best_fitness)

    params = state.best_member
    self.optimal_params = self.optimizable.unflatten_params(params)
    if self.optimizable.transformation is not None:
        self.optimal_params = self.optimizable.transformation.inverse_transform(
            self.optimal_params
        )
    return OptimizationResult(
        params=self.optimal_params,
        success=True,
        nit=self.num_generations,
        nfev=self.num_generations * self.pop_size,
        message=f"Completed {self.num_generations} generations.",
        final_loss=float(state.best_fitness),
        loss_history=list(self._fitness_history),
    )

IPOPT

Bases: Optimizer

Interior Point Optimizer (IPOPT) for optimization of the objective function with optional constraints and bounds.

Parameters:

Name	Type	Description	Default
optimizable	Optimizable	The optimizable object.	required
options	dict	Options forwarded to cyipopt.minimize_ipopt. See https://coin-or.github.io/Ipopt/OPTIONS.html for the full list. Commonly used keys: maxiter (int, default 3000) Maximum number of IPOPT iterations. disp (int, default 5) Verbosity level (0 = silent). tol (float) Convergence tolerance on the NLP optimality conditions. acceptable_tol (float) Looser acceptable-solution tolerance.	{'disp': 5}

Source code in jaxonomy/optimization/framework/optimizers_ipopt.py

class IPOPT(Optimizer):
    """
    Interior Point Optimizer (IPOPT) for optimization of the objective function
    with optional constraints and bounds.

    Parameters:
        optimizable (Optimizable):
            The optimizable object.
        options (dict):
            Options forwarded to ``cyipopt.minimize_ipopt``.
            See https://coin-or.github.io/Ipopt/OPTIONS.html for the full
            list.  Commonly used keys:

            ``maxiter`` (int, default 3000)
                Maximum number of IPOPT iterations.
            ``disp`` (int, default 5)
                Verbosity level (0 = silent).
            ``tol`` (float)
                Convergence tolerance on the NLP optimality conditions.
            ``acceptable_tol`` (float)
                Looser acceptable-solution tolerance.
    """

    def __init__(self, optimizable: Optimizable, options: dict = {"disp": 5}):
        self.optimizable = optimizable
        self.options = options
        self.optimal_params = None

    def optimize(self) -> OptimizationResult:
        """Run optimisation and return an :class:`~jaxonomy.optimization.OptimizationResult`.

        Gradients of the objective and constraint Jacobians are computed with
        JAX automatic differentiation (``jax.grad`` / ``jax.jacrev``).

        **Hessian strategy** — JAX's ``jax.hessian`` requires forward-mode
        automatic differentiation (``jacfwd``) through the gradient, but the
        jaxonomy ODE solver uses ``custom_vjp`` which only supports
        reverse-mode.  Attempting to compute ``jax.hessian`` of a simulation
        objective therefore raises a runtime error.  IPOPT is instead
        configured with ``hessian_approximation = "limited-memory"`` (L-BFGS
        approximation) which only requires first-order gradient information and
        converges super-linearly.  For problems where you *know* the objective
        is twice-differentiable and do not use the jaxonomy ODE integrator you
        can override this by passing ``options={"hessian_approximation":
        "exact", ...}`` and providing ``hess`` via the ``_hess_fn`` constructor
        argument.
        """
        params = self.optimizable.params_0_flat

        # ── objective ─────────────────────────────────────────────────────────
        objective = jax.jit(self.optimizable.objective_flat)
        gradient = jax.jit(jax.grad(objective))

        # ── constraints (only when present) ───────────────────────────────────
        if self.optimizable.has_constraints:
            # Ensure the constraint function always returns a 1-D array so
            # cyipopt can reliably probe its Jacobian sparsity structure.
            constraints = jax.jit(_atleast_1d_output(self.optimizable.constraints_flat))
            constraints_jac = jax.jit(jax.jacrev(constraints))

            constraints_ipopt = [
                {
                    "type": "ineq",
                    "fun": constraints,
                    "jac": constraints_jac,
                    # No "hess" key: IPOPT uses L-BFGS for constraint Hessians
                    # when hessian_approximation="limited-memory"
                }
            ]
        else:
            constraints_ipopt = []

        # ── bounds ─────────────────────────────────────────────────────────────
        bounds = self.optimizable.bounds_flat

        # Jobs from the UI may put (-inf, inf) as default bounds.  The user
        # may also specify bounds this way.  cyipopt expects ``None`` to mean
        # unbounded.
        if bounds is not None:
            bounds = [
                (
                    None if b[0] == -jnp.inf else b[0],
                    None if b[1] == jnp.inf else b[1],
                )
                for b in bounds
            ]

            # If every bound is None the problem is effectively unbounded.
            flattened_bounds = [element for tup in bounds for element in tup]
            all_none = all(element is None for element in flattened_bounds)
            bounds = None if all_none else bounds

        # ── merge options — inject limited-memory unless caller overrides ──────
        effective_options = {"hessian_approximation": "limited-memory"}
        effective_options.update(self.options)

        # ── call IPOPT ─────────────────────────────────────────────────────────
        res = cyipopt.minimize_ipopt(
            objective,
            x0=params,
            jac=gradient,
            hess=None,
            constraints=constraints_ipopt,
            bounds=bounds,
            options=effective_options,
        )

        logger.info("IPOPT result:\n%s", res)
        if not getattr(res, "success", False):
            logger.warning("IPOPT did not converge: %s", getattr(res, "message", ""))

        # ── unpack result ──────────────────────────────────────────────────────
        solved_params = res.x
        self.optimal_params = self.optimizable.unflatten_params(solved_params)
        if self.optimizable.transformation is not None:
            self.optimal_params = self.optimizable.transformation.inverse_transform(
                self.optimal_params
            )

        return OptimizationResult(
            params=self.optimal_params,
            success=bool(getattr(res, "success", False)),
            nit=int(getattr(res, "nit", 0)),
            nfev=int(getattr(res, "nfev", 0)),
            message=str(getattr(res, "message", "")),
            final_loss=float(getattr(res, "fun", float("nan"))),
        )

optimize()

Run optimisation and return an :class:~jaxonomy.optimization.OptimizationResult.

Gradients of the objective and constraint Jacobians are computed with JAX automatic differentiation (jax.grad / jax.jacrev).

Hessian strategy — JAX's jax.hessian requires forward-mode automatic differentiation (jacfwd) through the gradient, but the jaxonomy ODE solver uses custom_vjp which only supports reverse-mode. Attempting to compute jax.hessian of a simulation objective therefore raises a runtime error. IPOPT is instead configured with hessian_approximation = "limited-memory" (L-BFGS approximation) which only requires first-order gradient information and converges super-linearly. For problems where you know the objective is twice-differentiable and do not use the jaxonomy ODE integrator you can override this by passing options={"hessian_approximation": "exact", ...} and providing hess via the _hess_fn constructor argument.

Source code in jaxonomy/optimization/framework/optimizers_ipopt.py

def optimize(self) -> OptimizationResult:
    """Run optimisation and return an :class:`~jaxonomy.optimization.OptimizationResult`.

    Gradients of the objective and constraint Jacobians are computed with
    JAX automatic differentiation (``jax.grad`` / ``jax.jacrev``).

    **Hessian strategy** — JAX's ``jax.hessian`` requires forward-mode
    automatic differentiation (``jacfwd``) through the gradient, but the
    jaxonomy ODE solver uses ``custom_vjp`` which only supports
    reverse-mode.  Attempting to compute ``jax.hessian`` of a simulation
    objective therefore raises a runtime error.  IPOPT is instead
    configured with ``hessian_approximation = "limited-memory"`` (L-BFGS
    approximation) which only requires first-order gradient information and
    converges super-linearly.  For problems where you *know* the objective
    is twice-differentiable and do not use the jaxonomy ODE integrator you
    can override this by passing ``options={"hessian_approximation":
    "exact", ...}`` and providing ``hess`` via the ``_hess_fn`` constructor
    argument.
    """
    params = self.optimizable.params_0_flat

    # ── objective ─────────────────────────────────────────────────────────
    objective = jax.jit(self.optimizable.objective_flat)
    gradient = jax.jit(jax.grad(objective))

    # ── constraints (only when present) ───────────────────────────────────
    if self.optimizable.has_constraints:
        # Ensure the constraint function always returns a 1-D array so
        # cyipopt can reliably probe its Jacobian sparsity structure.
        constraints = jax.jit(_atleast_1d_output(self.optimizable.constraints_flat))
        constraints_jac = jax.jit(jax.jacrev(constraints))

        constraints_ipopt = [
            {
                "type": "ineq",
                "fun": constraints,
                "jac": constraints_jac,
                # No "hess" key: IPOPT uses L-BFGS for constraint Hessians
                # when hessian_approximation="limited-memory"
            }
        ]
    else:
        constraints_ipopt = []

    # ── bounds ─────────────────────────────────────────────────────────────
    bounds = self.optimizable.bounds_flat

    # Jobs from the UI may put (-inf, inf) as default bounds.  The user
    # may also specify bounds this way.  cyipopt expects ``None`` to mean
    # unbounded.
    if bounds is not None:
        bounds = [
            (
                None if b[0] == -jnp.inf else b[0],
                None if b[1] == jnp.inf else b[1],
            )
            for b in bounds
        ]

        # If every bound is None the problem is effectively unbounded.
        flattened_bounds = [element for tup in bounds for element in tup]
        all_none = all(element is None for element in flattened_bounds)
        bounds = None if all_none else bounds

    # ── merge options — inject limited-memory unless caller overrides ──────
    effective_options = {"hessian_approximation": "limited-memory"}
    effective_options.update(self.options)

    # ── call IPOPT ─────────────────────────────────────────────────────────
    res = cyipopt.minimize_ipopt(
        objective,
        x0=params,
        jac=gradient,
        hess=None,
        constraints=constraints_ipopt,
        bounds=bounds,
        options=effective_options,
    )

    logger.info("IPOPT result:\n%s", res)
    if not getattr(res, "success", False):
        logger.warning("IPOPT did not converge: %s", getattr(res, "message", ""))

    # ── unpack result ──────────────────────────────────────────────────────
    solved_params = res.x
    self.optimal_params = self.optimizable.unflatten_params(solved_params)
    if self.optimizable.transformation is not None:
        self.optimal_params = self.optimizable.transformation.inverse_transform(
            self.optimal_params
        )

    return OptimizationResult(
        params=self.optimal_params,
        success=bool(getattr(res, "success", False)),
        nit=int(getattr(res, "nit", 0)),
        nfev=int(getattr(res, "nfev", 0)),
        message=str(getattr(res, "message", "")),
        final_loss=float(getattr(res, "fun", float("nan"))),
    )

IdentityTransform

Bases: Transform

A transformation that does nothing: y = x.

Source code in jaxonomy/optimization/framework/base/transformations.py

class IdentityTransform(Transform):
    """
    A transformation that does nothing: ``` y = x ```.
    """

    def transform(self, params: dict):
        return params

    def inverse_transform(self, params: dict):
        return params

LogTransform

Bases: Transform

A transformation that applies the natural logarithm to the values of the parameters. y = log(x).

Source code in jaxonomy/optimization/framework/base/transformations.py

class LogTransform(Transform):
    """
    A transformation that applies the natural logarithm to the values of the parameters.
    ``` y = log(x) ```.
    """

    def transform(self, params: dict):
        return {k: jnp.log(v) for k, v in params.items()}

    def inverse_transform(self, params: dict):
        return {k: jnp.exp(v) for k, v in params.items()}

LogitTransform

Bases: Transform

The logit transformation, defined as: y = log(x / (1 - x))

Source code in jaxonomy/optimization/framework/base/transformations.py

class LogitTransform(Transform):
    """
    The logit transformation, defined as:
    ``` y = log(x / (1 - x)) ```
    """

    def transform(self, params: dict):
        return {k: jnp.log(v / (1.0 - v)) for k, v in params.items()}

    def inverse_transform(self, params: dict):
        return {k: 1.0 / (1.0 + jnp.exp(-v)) for k, v in params.items()}

MultiStart

Multi-start wrapper for any jaxonomy optimizer.

Runs n_starts optimizations from different initial points and returns all results as well as the best one (lowest final_loss).

Parameters

optimizable : Optimizable The problem to optimize. Must be a jaxonomy Optimizable instance. optimizer_factory : Callable[[Optimizable], Optimizer] A factory function that takes an Optimizable (potentially with different initial parameters) and returns a ready-to-run optimizer. Example::

    factory = lambda opt: Scipy(opt, "L-BFGS-B",
                                opt_method_config={"maxiter": 40},
                                use_autodiff_grad=True)
    ms = MultiStart(optimizable, factory, n_starts=8, seed=0)
    result = ms.run()

int

Number of random restarts (default 10).

init_sampler : Callable or None Custom sampling function with signature (n_starts: int, params_0_flat: np.ndarray) -> np.ndarray returning an array of shape (n_starts, n_params). Row 0 is always replaced with the original params_0_flat when include_initial=True. If None (default), uniform sampling around params_0 is used. sample_scale : float Scale factor for the default uniform sampler. The search window for each parameter is [p0 ± sample_scale * max(|p0|, 1)] (default 1.0). seed : int or None Random seed for reproducibility. include_initial : bool When True (default), the first start always uses the original params_0, regardless of the sampler output.

Source code in jaxonomy/optimization/multi_start.py

class MultiStart:
    """
    Multi-start wrapper for any jaxonomy optimizer.

    Runs ``n_starts`` optimizations from different initial points and returns
    all results as well as the best one (lowest ``final_loss``).

    Parameters
    ----------
    optimizable : Optimizable
        The problem to optimize.  Must be a jaxonomy ``Optimizable`` instance.
    optimizer_factory : Callable[[Optimizable], Optimizer]
        A *factory* function that takes an ``Optimizable`` (potentially with
        different initial parameters) and returns a ready-to-run optimizer.
        Example::

            factory = lambda opt: Scipy(opt, "L-BFGS-B",
                                        opt_method_config={"maxiter": 40},
                                        use_autodiff_grad=True)
            ms = MultiStart(optimizable, factory, n_starts=8, seed=0)
            result = ms.run()

    n_starts : int
        Number of random restarts (default 10).
    init_sampler : Callable or None
        Custom sampling function with signature
        ``(n_starts: int, params_0_flat: np.ndarray) -> np.ndarray``
        returning an array of shape ``(n_starts, n_params)``.
        Row 0 is always replaced with the original ``params_0_flat`` when
        ``include_initial=True``.  If ``None`` (default), uniform sampling
        around ``params_0`` is used.
    sample_scale : float
        Scale factor for the default uniform sampler.  The search window
        for each parameter is
        ``[p0 ± sample_scale * max(|p0|, 1)]`` (default 1.0).
    seed : int or None
        Random seed for reproducibility.
    include_initial : bool
        When ``True`` (default), the first start always uses the original
        ``params_0``, regardless of the sampler output.
    """

    def __init__(
        self,
        optimizable: Optimizable,
        optimizer_factory: Callable,
        n_starts: int = 10,
        init_sampler: Callable | None = None,
        sample_scale: float = 1.0,
        seed: int | None = None,
        include_initial: bool = True,
    ):
        self.optimizable = optimizable
        self.optimizer_factory = optimizer_factory
        self.n_starts = n_starts
        self.init_sampler = init_sampler
        self.sample_scale = sample_scale
        self.include_initial = include_initial
        self.rng = np.random.default_rng(seed)
        self._results: list[OptimizationResult] = []

    def _generate_starts(self) -> np.ndarray:
        """Return (n_starts, n_params) array of initial parameter vectors."""
        params_0 = np.array(self.optimizable.params_0_flat)

        if self.init_sampler is not None:
            starts = np.array(self.init_sampler(self.n_starts, params_0))
        else:
            starts = _uniform_sampler(
                self.n_starts,
                params_0,
                self.sample_scale,
                self.optimizable.bounds_flat,
                self.rng,
            )

        if self.include_initial:
            starts[0] = params_0

        return starts

    def run(self) -> MultiStartResult:
        """
        Execute all starts sequentially and return a :class:`MultiStartResult`.

        Each start clones the optimizable with a new ``params_0_flat``, calls
        ``optimizer_factory(clone)`` to get a fresh optimizer, and runs
        ``optimizer.optimize()``.  Failed starts (exceptions) are recorded as
        unsuccessful ``OptimizationResult`` entries with ``success=False``.

        Returns
        -------
        MultiStartResult
        """
        starts = self._generate_starts()
        results: list[OptimizationResult] = []

        for i, p0 in enumerate(starts):
            # Shallow-copy the optimizable and override its initial params.
            opt_clone = copy.copy(self.optimizable)
            opt_clone.params_0_flat = jnp.array(p0)

            optimizer = self.optimizer_factory(opt_clone)

            try:
                result = optimizer.optimize()
                if not isinstance(result, OptimizationResult):
                    # Wrap legacy plain-dict return for compatibility
                    result = OptimizationResult(params=dict(result))
            except Exception as exc:  # noqa: BLE001
                nan_params = {
                    k: float("nan")
                    for k in self.optimizable.unflatten_params(
                        jnp.array(p0)
                    )
                }
                result = OptimizationResult(
                    params=nan_params,
                    success=False,
                    message=f"Start {i} raised {type(exc).__name__}: {exc}",
                    final_loss=float("inf"),
                )

            results.append(result)

        self._results = results

        # Pick best: lowest finite final_loss among successful runs.
        successful = [
            (i, r)
            for i, r in enumerate(results)
            if r.success
            and r.final_loss is not None
            and np.isfinite(float(r.final_loss))
        ]

        if successful:
            best_idx, best = min(successful, key=lambda ir: float(ir[1].final_loss))
        else:
            # All failed — fall back to last result
            best_idx = len(results) - 1
            best = results[best_idx]

        n_successful = sum(1 for r in results if r.success)

        return MultiStartResult(
            results=results,
            best_result=best,
            best_start_index=best_idx,
            n_starts=self.n_starts,
            n_successful=n_successful,
        )

    @property
    def results(self) -> list[OptimizationResult]:
        """Results from the last :meth:`run` call (empty before first run)."""
        return self._results

results property

Results from the last :meth:run call (empty before first run).

run()

Execute all starts sequentially and return a :class:MultiStartResult.

Each start clones the optimizable with a new params_0_flat, calls optimizer_factory(clone) to get a fresh optimizer, and runs optimizer.optimize(). Failed starts (exceptions) are recorded as unsuccessful OptimizationResult entries with success=False.

Returns

MultiStartResult

Source code in jaxonomy/optimization/multi_start.py

def run(self) -> MultiStartResult:
    """
    Execute all starts sequentially and return a :class:`MultiStartResult`.

    Each start clones the optimizable with a new ``params_0_flat``, calls
    ``optimizer_factory(clone)`` to get a fresh optimizer, and runs
    ``optimizer.optimize()``.  Failed starts (exceptions) are recorded as
    unsuccessful ``OptimizationResult`` entries with ``success=False``.

    Returns
    -------
    MultiStartResult
    """
    starts = self._generate_starts()
    results: list[OptimizationResult] = []

    for i, p0 in enumerate(starts):
        # Shallow-copy the optimizable and override its initial params.
        opt_clone = copy.copy(self.optimizable)
        opt_clone.params_0_flat = jnp.array(p0)

        optimizer = self.optimizer_factory(opt_clone)

        try:
            result = optimizer.optimize()
            if not isinstance(result, OptimizationResult):
                # Wrap legacy plain-dict return for compatibility
                result = OptimizationResult(params=dict(result))
        except Exception as exc:  # noqa: BLE001
            nan_params = {
                k: float("nan")
                for k in self.optimizable.unflatten_params(
                    jnp.array(p0)
                )
            }
            result = OptimizationResult(
                params=nan_params,
                success=False,
                message=f"Start {i} raised {type(exc).__name__}: {exc}",
                final_loss=float("inf"),
            )

        results.append(result)

    self._results = results

    # Pick best: lowest finite final_loss among successful runs.
    successful = [
        (i, r)
        for i, r in enumerate(results)
        if r.success
        and r.final_loss is not None
        and np.isfinite(float(r.final_loss))
    ]

    if successful:
        best_idx, best = min(successful, key=lambda ir: float(ir[1].final_loss))
    else:
        # All failed — fall back to last result
        best_idx = len(results) - 1
        best = results[best_idx]

    n_successful = sum(1 for r in results if r.success)

    return MultiStartResult(
        results=results,
        best_result=best,
        best_start_index=best_idx,
        n_starts=self.n_starts,
        n_successful=n_successful,
    )

MultiStartResult dataclass

Results from a multi-start optimization run.

Attributes:

Name	Type	Description
results	list[OptimizationResult]	All OptimizationResult objects — one per start.
best_result	OptimizationResult	The result with the lowest final_loss among successful runs.
best_start_index	int	Index into results of the best run.
n_starts	int	Total number of starts attempted.
n_successful	int	Number of starts that reported success=True.

Source code in jaxonomy/optimization/multi_start.py

@dataclass
class MultiStartResult:
    """
    Results from a multi-start optimization run.

    Attributes:
        results: All ``OptimizationResult`` objects — one per start.
        best_result: The result with the lowest ``final_loss`` among
            successful runs.
        best_start_index: Index into ``results`` of the best run.
        n_starts: Total number of starts attempted.
        n_successful: Number of starts that reported ``success=True``.
    """

    results: list[OptimizationResult]
    best_result: OptimizationResult
    best_start_index: int
    n_starts: int
    n_successful: int

    def summary(self) -> str:
        lines = [
            f"MultiStartResult: {self.n_successful}/{self.n_starts} starts converged.",
            f"Best start: #{self.best_start_index}  "
            f"final_loss={self.best_result.final_loss:.6g}  "
            f"params={self.best_result.params}",
        ]
        for i, r in enumerate(self.results):
            marker = " ← best" if i == self.best_start_index else ""
            loss_str = f"{r.final_loss:.6g}" if r.final_loss is not None else "N/A"
            lines.append(
                f"  [{i}] success={r.success}  loss={loss_str}  "
                f"nit={r.nit}{marker}"
            )
        return "\n".join(lines)

    def __repr__(self) -> str:
        return self.summary()

NLopt

Bases: Optimizer

Optimizers using the NLopt library.

Parameters:

Name	Type	Description	Default
optimizable	Optimizable	The optimizable object.	required
opt_method	str	The optimization method to use.	required
ftol_rel	float	Relative tolerance on function value.	1e-06
ftol_abs	float	Absolute tolerance on function value.	1e-06
xtol_rel	float	Relative tolerance on optimization parameters.	1e-06
xtol_abs	float	Absolute tolerance on optimization parameters.	1e-06
cons_tol	float	Tolerance on constraints.	1e-06
maxeval	int	Maximum number of function evaluations.	500
maxtime	float	Maximum time in seconds.	0

Source code in jaxonomy/optimization/framework/optimizers_nlopt.py

class NLopt(Optimizer):
    """
    Optimizers using the NLopt library.

    Parameters:
        optimizable (Optimizable):
            The optimizable object.
        opt_method (str):
            The optimization method to use.
        ftol_rel (float):
            Relative tolerance on function value.
        ftol_abs (float):
            Absolute tolerance on function value.
        xtol_rel (float):
            Relative tolerance on optimization parameters.
        xtol_abs (float):
            Absolute tolerance on optimization parameters.
        cons_tol (float):
            Tolerance on constraints.
        maxeval (int):
            Maximum number of function evaluations.
        maxtime (float):
            Maximum time in seconds.
    """

    def __init__(
        self,
        optimizable: Optimizable,
        opt_method: str,
        ftol_rel=1e-06,
        ftol_abs=1e-06,
        xtol_rel=1e-06,
        xtol_abs=1e-06,
        cons_tol=1e-06,
        maxeval=500,
        maxtime=0,
    ):
        self.optimizable = optimizable
        self.opt_method = opt_method
        self.ftol_rel = ftol_rel
        self.ftol_abs = ftol_abs
        self.xtol_rel = xtol_rel
        self.xtol_abs = xtol_abs
        self.cons_tol = cons_tol
        self.maxeval = maxeval
        self.maxtime = maxtime
        self.optimal_params = None

    def optimize(self):
        """Run optimization"""
        params = self.optimizable.params_0_flat
        objective = jax.jit(self.optimizable.objective_flat)
        gradient = jax.jit(jax.grad(objective))

        constraints = jax.jit(self.optimizable.constraints_flat)
        constraints_jac = jax.jit(jax.jacrev(constraints))

        def nlopt_obj(x, grad):
            if grad.size > 0:
                grad[:] = gradient(jnp.array(x))
            return float(objective(jnp.array(x)))

        def nlopt_cons(result, x, grad):
            if grad.size > 0:
                grad[:, :] = -constraints_jac(jnp.array(x))
            result[:] = -constraints(jnp.array(x))

        if (
            self.optimizable.bounds_flat is not None
            and self.opt_method not in SUPPORTS_BOUNDS
        ):
            raise ValueError(
                f"Optimization method nlopt:{self.opt_method} does not support bounds."
            )

        if (
            self.optimizable.has_constraints
            and self.opt_method not in SUPPORTS_CONSTRAINTS
        ):
            raise ValueError(
                f"Optimization method nlopt:{self.opt_method} "
                "does not support constraints."
            )

        if self.opt_method not in ALL_METHODS:
            raise ValueError(
                f"Optimization method nlopt:{self.opt_method} is not supported."
            )

        # Initialize nlopt optimizer
        opt_method = ALL_METHODS[self.opt_method]()
        opt = nlopt.opt(opt_method, len(params))

        # Set the objective function
        opt.set_min_objective(nlopt_obj)

        # Set the constraints
        if self.optimizable.has_constraints:
            num_constraints = self.optimizable.constraints_flat(jnp.array(params)).size
            opt.add_inequality_mconstraint(
                nlopt_cons, [self.cons_tol] * num_constraints
            )

        # Set the bounds
        if self.optimizable.bounds_flat is not None:
            lower_bounds, upper_bounds = zip(*self.optimizable.bounds_flat)
            opt.set_lower_bounds(lower_bounds)
            opt.set_upper_bounds(upper_bounds)

        # Set stopping criteria
        opt.set_ftol_rel(self.ftol_rel)
        opt.set_ftol_abs(self.ftol_abs)
        opt.set_xtol_rel(self.xtol_rel)
        opt.set_xtol_abs(self.xtol_abs)
        opt.set_maxeval(self.maxeval)
        opt.set_maxtime(self.maxtime)

        # Run the optimization
        params = opt.optimize(params)

        result_code = opt.last_optimize_result()
        _NLOPT_MESSAGES = {
            1: "Success",
            2: "Stopval reached",
            3: "Function tolerance reached",
            4: "X tolerance reached",
            5: "Maximum evaluations reached",
            6: "Maximum time reached",
            -1: "Failure",
            -2: "Invalid arguments",
            -3: "Out of memory",
            -4: "Roundoff limited",
            -5: "Forced stop",
        }
        _nlopt_nfev = opt.get_numevals()
        try:
            _nlopt_final_loss = float(opt.last_optimum_value())
        except Exception:
            _nlopt_final_loss = None

        self.optimal_params = self.optimizable.unflatten_params(params)
        if self.optimizable.transformation is not None:
            self.optimal_params = self.optimizable.transformation.inverse_transform(
                self.optimal_params
            )
        return OptimizationResult(
            params=self.optimal_params,
            success=result_code > 0,
            nit=_nlopt_nfev,
            nfev=_nlopt_nfev,
            message=_NLOPT_MESSAGES.get(result_code, f"Unknown code {result_code}"),
            final_loss=_nlopt_final_loss,
        )

optimize()

Run optimization

Source code in jaxonomy/optimization/framework/optimizers_nlopt.py

def optimize(self):
    """Run optimization"""
    params = self.optimizable.params_0_flat
    objective = jax.jit(self.optimizable.objective_flat)
    gradient = jax.jit(jax.grad(objective))

    constraints = jax.jit(self.optimizable.constraints_flat)
    constraints_jac = jax.jit(jax.jacrev(constraints))

    def nlopt_obj(x, grad):
        if grad.size > 0:
            grad[:] = gradient(jnp.array(x))
        return float(objective(jnp.array(x)))

    def nlopt_cons(result, x, grad):
        if grad.size > 0:
            grad[:, :] = -constraints_jac(jnp.array(x))
        result[:] = -constraints(jnp.array(x))

    if (
        self.optimizable.bounds_flat is not None
        and self.opt_method not in SUPPORTS_BOUNDS
    ):
        raise ValueError(
            f"Optimization method nlopt:{self.opt_method} does not support bounds."
        )

    if (
        self.optimizable.has_constraints
        and self.opt_method not in SUPPORTS_CONSTRAINTS
    ):
        raise ValueError(
            f"Optimization method nlopt:{self.opt_method} "
            "does not support constraints."
        )

    if self.opt_method not in ALL_METHODS:
        raise ValueError(
            f"Optimization method nlopt:{self.opt_method} is not supported."
        )

    # Initialize nlopt optimizer
    opt_method = ALL_METHODS[self.opt_method]()
    opt = nlopt.opt(opt_method, len(params))

    # Set the objective function
    opt.set_min_objective(nlopt_obj)

    # Set the constraints
    if self.optimizable.has_constraints:
        num_constraints = self.optimizable.constraints_flat(jnp.array(params)).size
        opt.add_inequality_mconstraint(
            nlopt_cons, [self.cons_tol] * num_constraints
        )

    # Set the bounds
    if self.optimizable.bounds_flat is not None:
        lower_bounds, upper_bounds = zip(*self.optimizable.bounds_flat)
        opt.set_lower_bounds(lower_bounds)
        opt.set_upper_bounds(upper_bounds)

    # Set stopping criteria
    opt.set_ftol_rel(self.ftol_rel)
    opt.set_ftol_abs(self.ftol_abs)
    opt.set_xtol_rel(self.xtol_rel)
    opt.set_xtol_abs(self.xtol_abs)
    opt.set_maxeval(self.maxeval)
    opt.set_maxtime(self.maxtime)

    # Run the optimization
    params = opt.optimize(params)

    result_code = opt.last_optimize_result()
    _NLOPT_MESSAGES = {
        1: "Success",
        2: "Stopval reached",
        3: "Function tolerance reached",
        4: "X tolerance reached",
        5: "Maximum evaluations reached",
        6: "Maximum time reached",
        -1: "Failure",
        -2: "Invalid arguments",
        -3: "Out of memory",
        -4: "Roundoff limited",
        -5: "Forced stop",
    }
    _nlopt_nfev = opt.get_numevals()
    try:
        _nlopt_final_loss = float(opt.last_optimum_value())
    except Exception:
        _nlopt_final_loss = None

    self.optimal_params = self.optimizable.unflatten_params(params)
    if self.optimizable.transformation is not None:
        self.optimal_params = self.optimizable.transformation.inverse_transform(
            self.optimal_params
        )
    return OptimizationResult(
        params=self.optimal_params,
        success=result_code > 0,
        nit=_nlopt_nfev,
        nfev=_nlopt_nfev,
        message=_NLOPT_MESSAGES.get(result_code, f"Unknown code {result_code}"),
        final_loss=_nlopt_final_loss,
    )

NegativeNegativeLogTransform

Bases: Transform

A transformation that applies the negative of the natural logarithm of the negative of the values of the parameters. y = -log(-x)

Source code in jaxonomy/optimization/framework/base/transformations.py

class NegativeNegativeLogTransform(Transform):
    """
    A transformation that applies the negative of the natural logarithm of the negative
    of the values of the parameters.
    ``` y = -log(-x) ```
    """

    def transform(self, params: dict):
        return {k: -jnp.log(-v) for k, v in params.items()}

    def inverse_transform(self, params: dict):
        return {k: -jnp.exp(-v) for k, v in params.items()}

NormalizeTransform

Bases: Transform

A transformation that normalizes the values of the parameters to the range [0, 1]. y = (x - min) / (max - min) Paramteters: - params_min: dict with the minimum values for each parameter. - params_max: dict with the maximum values for each parameter.

Source code in jaxonomy/optimization/framework/base/transformations.py

class NormalizeTransform(Transform):
    """
    A transformation that normalizes the values of the parameters to the range [0, 1].
    ``` y = (x - min) / (max - min) ```
    Paramteters:
        - params_min: dict with the minimum values for each parameter.
        - params_max: dict with the maximum values for each parameter.
    """

    def __init__(self, params_min: dict, params_max: dict):
        self.params_min = params_min
        self.params_max = params_max

    def transform(self, params: dict):
        return {
            k: (v - self.params_min[k]) / (self.params_max[k] - self.params_min[k])
            for k, v in params.items()
        }

    def inverse_transform(self, params: dict):
        return {
            k: v * (self.params_max[k] - self.params_min[k]) + self.params_min[k]
            for k, v in params.items()
        }

Optax

Bases: Optimizer

Optax optimizer without support for stochastic variables.

Paramters

optimizable (Optimizable): The optimizable object. opt_method (str): The optimization method to use. learning_rate (float): The learning rate. opt_method_config (dict): Configuration for the optimization method. num_epochs (int): The number of epochs. clip_range (tuple): The range to clip the gradients. print_every (int): Print progress every print_every epochs. metrics_writer (MetricsWriter|None): Optional CSV file to write metrics to.

Source code in jaxonomy/optimization/framework/optimizers_optax.py

class Optax(Optimizer):
    """
    Optax optimizer without support for stochastic variables.

    Paramters:
        optimizable (Optimizable):
            The optimizable object.
        opt_method (str):
            The optimization method to use.
        learning_rate (float):
            The learning rate.
        opt_method_config (dict):
            Configuration for the optimization method.
        num_epochs (int):
            The number of epochs.
        clip_range (tuple):
            The range to clip the gradients.
        print_every (int):
            Print progress every `print_every` epochs.
        metrics_writer (MetricsWriter|None):
            Optional CSV file to write metrics to.
    """

    def __init__(
        self,
        optimizable: Optimizable,
        opt_method,
        learning_rate,
        opt_method_config,
        num_epochs=100,
        clip_range=None,
        print_every=None,
        metrics_writer: MetricsWriter = None,
    ):
        self.optimizable = optimizable
        self.opt_method = opt_method
        self.num_epochs = num_epochs
        self.clip_range = clip_range
        self.print_every = print_every
        self.metrics_writer = metrics_writer
        self.optimal_params = None
        self.losses = []

        if self.optimizable.bounds_flat is not None:
            # Jobs from UI would put (-jnp.inf, jnp.inf) as defualt bounds. The user
            # may also have specified bounds this way.
            bounds = [
                (
                    None if b[0] == -jnp.inf else b[0],
                    None if b[1] == jnp.inf else b[1],
                )
                for b in self.optimizable.bounds_flat
            ]

            # Check if all bounds are None, i.e. no bounds at all, and hence Optax
            # algorithms which don't natively support bounds can be used
            flattened_bounds = [element for tup in bounds for element in tup]
            all_none = all(element is None for element in flattened_bounds)

            if not all_none:
                raise ValueError(
                    f"Optimization method {opt_method} does not support bounds."
                )

        if self.optimizable.has_constraints:
            raise ValueError(
                f"Optimization method optax:{self.opt_method} "
                "does not support constraints."
            )

        opt_func = getattr(optax, opt_method)

        # Instantiate the optimizer with validated config
        valid_opts = _remap_and_filter_valid_params(opt_func, opt_method_config)
        self.optimizer = opt_func(learning_rate, **valid_opts)

    @partial(jax.jit, static_argnums=(0,))
    def step(self, params, opt_state):
        """Take a single optimization step"""
        loss, grads = jax.value_and_grad(self.optimizable.objective_flat)(params)
        grads = jnp.clip(grads, *self.clip_range) if self.clip_range else grads
        updates, opt_state = self.optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, loss

    def optimize(self) -> dict[str, ArrayLike]:
        """Run optimization"""
        params = self.optimizable.params_0_flat
        opt_state = self.optimizer.init(params)

        for epoch in range(self.num_epochs):
            params, opt_state, loss = self.step(params, opt_state)
            mean_loss = jnp.mean(loss)
            self.losses.append(mean_loss)
            if self.print_every and epoch % self.print_every == 0:
                p: dict = self.optimizable.unflatten_params(params)
                if self.optimizable.transformation is not None:
                    p = self.optimizable.transformation.inverse_transform(p)
                p = {k: v.tolist() for k, v in p.items()}
                logger.info("Epoch %s, loss: %s", epoch, mean_loss, **logdata(params=p))
            if self.metrics_writer:
                self.metrics_writer.write_metrics(loss=mean_loss)

        self.optimal_params = self.optimizable.unflatten_params(params)
        if self.optimizable.transformation is not None:
            self.optimal_params = self.optimizable.transformation.inverse_transform(
                self.optimal_params
            )
        return OptimizationResult(
            params=self.optimal_params,
            success=True,
            nit=self.num_epochs,
            nfev=self.num_epochs,
            message=f"Completed {self.num_epochs} epochs.",
            final_loss=float(self.losses[-1]) if self.losses else None,
            loss_history=[float(x) for x in self.losses],
        )

    @property
    def metrics(self):
        return {"loss": self.losses}

optimize()

Run optimization

Source code in jaxonomy/optimization/framework/optimizers_optax.py

def optimize(self) -> dict[str, ArrayLike]:
    """Run optimization"""
    params = self.optimizable.params_0_flat
    opt_state = self.optimizer.init(params)

    for epoch in range(self.num_epochs):
        params, opt_state, loss = self.step(params, opt_state)
        mean_loss = jnp.mean(loss)
        self.losses.append(mean_loss)
        if self.print_every and epoch % self.print_every == 0:
            p: dict = self.optimizable.unflatten_params(params)
            if self.optimizable.transformation is not None:
                p = self.optimizable.transformation.inverse_transform(p)
            p = {k: v.tolist() for k, v in p.items()}
            logger.info("Epoch %s, loss: %s", epoch, mean_loss, **logdata(params=p))
        if self.metrics_writer:
            self.metrics_writer.write_metrics(loss=mean_loss)

    self.optimal_params = self.optimizable.unflatten_params(params)
    if self.optimizable.transformation is not None:
        self.optimal_params = self.optimizable.transformation.inverse_transform(
            self.optimal_params
        )
    return OptimizationResult(
        params=self.optimal_params,
        success=True,
        nit=self.num_epochs,
        nfev=self.num_epochs,
        message=f"Completed {self.num_epochs} epochs.",
        final_loss=float(self.losses[-1]) if self.losses else None,
        loss_history=[float(x) for x in self.losses],
    )

step(params, opt_state)

Take a single optimization step

Source code in jaxonomy/optimization/framework/optimizers_optax.py

@partial(jax.jit, static_argnums=(0,))
def step(self, params, opt_state):
    """Take a single optimization step"""
    loss, grads = jax.value_and_grad(self.optimizable.objective_flat)(params)
    grads = jnp.clip(grads, *self.clip_range) if self.clip_range else grads
    updates, opt_state = self.optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

OptaxWithStochasticVars

Bases: Optimizer

Optax optimizer with support for stochastic variables.

Parameters:

Name	Type	Description	Default
optimizable	OptimizableWithStochasticVars	The optimizable object.	required
opt_method	str	The optimization method to use.	required
learning_rate	float	The learning rate.	required
opt_method_config	dict	Configuration for the optimization method.	required
num_epochs	int	The number of epochs.	100
batch_size	int	The batch size.	1
num_batches	int	The number of batches.	1
clip_range	tuple	The range to clip the gradients.	None
print_every	int	Print progress every print_every epochs.	None
metrics_writer	MetricsWriter | None	Optional CSV file to write metrics to.	None

Source code in jaxonomy/optimization/framework/optimizers_optax.py

class OptaxWithStochasticVars(Optimizer):
    """
    Optax optimizer with support for stochastic variables.

    Parameters:
        optimizable (OptimizableWithStochasticVars):
            The optimizable object.
        opt_method (str):
            The optimization method to use.
        learning_rate (float):
            The learning rate.
        opt_method_config (dict):
            Configuration for the optimization method.
        num_epochs (int):
            The number of epochs.
        batch_size (int):
            The batch size.
        num_batches (int):
            The number of batches.
        clip_range (tuple):
            The range to clip the gradients.
        print_every (int):
            Print progress every `print_every` epochs.
        metrics_writer (MetricsWriter|None):
            Optional CSV file to write metrics to.
    """

    def __init__(
        self,
        optimizable: OptimizableWithStochasticVars,
        opt_method: str,
        learning_rate,
        opt_method_config,
        num_epochs=100,
        batch_size=1,
        num_batches=1,
        clip_range=None,
        print_every=None,
        metrics_writer: MetricsWriter = None,
    ):
        self.optimizable = optimizable
        self.opt_method = opt_method
        self.num_epochs = num_epochs
        self.batch_size = batch_size
        self.num_batches = num_batches
        self.clip_range = clip_range
        self.print_every = print_every
        self.metrics_writer = metrics_writer
        self.optimal_params = None
        self.losses = []

        if optimizable.bounds_flat is not None:
            # Jobs from UI would put (-jnp.inf, jnp.inf) as defualt bounds. The user
            # may also have specified bounds this way.
            bounds = [
                (
                    None if b[0] == -jnp.inf else b[0],
                    None if b[1] == jnp.inf else b[1],
                )
                for b in self.optimizable.bounds_flat
            ]

            # Check if all bounds are None, i.e. no bounds at all, and hence Optax
            # algorithms which don't natively support bounds can be used
            flattened_bounds = [element for tup in bounds for element in tup]
            all_none = all(element is None for element in flattened_bounds)

            if not all_none:
                raise ValueError(
                    f"Optimization method {opt_method} does not support bounds."
                )

        opt_func = getattr(optax, opt_method, None)
        if opt_func is None:
            raise ValueError(f"Unknown optax optimizer: {opt_method}")

        # Instantiate the optimizer with validated config
        valid_opts = _remap_and_filter_valid_params(opt_func, opt_method_config)
        self.optimizer = opt_func(learning_rate, **valid_opts)

    def batched_objective_flat(self, params, stochastic_vars_batch_flat):
        """Mean of the objective function over a batch"""
        return jnp.mean(
            self.optimizable.batched_objective_flat(params, stochastic_vars_batch_flat)
        )

    @partial(jax.jit, static_argnums=(0,))
    def step(self, params, opt_state, stochastic_vars_batch):
        """Take a single optimization step over one batch"""
        batch_loss, grads = jax.value_and_grad(self.batched_objective_flat)(
            params, stochastic_vars_batch
        )

        grads = jnp.clip(grads, *self.clip_range) if self.clip_range else grads

        updates, opt_state = self.optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, batch_loss

    def optimize(self):
        """Run optimization"""
        params = self.optimizable.params_0_flat
        opt_state = self.optimizer.init(params)

        if self.num_batches * self.batch_size == 1:
            # don't randomize over stochastic variables; use a single
            # batch of size 1 with initial stochastic variables
            data_flat, _ = ravel_pytree(self.optimizable.vars_0)
            stochastic_vars_training_data_flat = data_flat[None, None, :]
        else:
            _, stochastic_vars_training_data_flat = self.optimizable.sample_random_vars(
                self.num_batches * self.batch_size
            )

        @jax.jit
        def _scan_fun(carry, stochastic_vars_batch):
            params, opt_state = carry
            params, opt_state, batch_loss = self.step(
                params, opt_state, stochastic_vars_batch
            )
            return (params, opt_state), batch_loss

        for epoch in range(self.num_epochs):
            if self.num_batches * self.batch_size == 1:
                stochastic_vars_batches = stochastic_vars_training_data_flat
            else:
                stochastic_vars_batches = self.optimizable.generate_batches(
                    stochastic_vars_training_data_flat,
                    self.num_batches,
                    self.batch_size,
                )
            (params, opt_state), batch_losses = lax.scan(
                _scan_fun, (params, opt_state), stochastic_vars_batches
            )

            self.losses.append(jnp.mean(batch_losses))
            if self.print_every and epoch % self.print_every == 0:
                p: dict = self.optimizable.unflatten_params(params)
                if self.optimizable.transformation is not None:
                    p = self.optimizable.transformation.inverse_transform(p)
                p = {k: v.tolist() for k, v in p.items()}
                logger.info(
                    "Epoch %s, average batch loss: %s",
                    epoch,
                    jnp.mean(batch_losses),
                    **logdata(params=p),
                )
            if self.metrics_writer is not None:
                self.metrics_writer.write_metrics(loss=self.losses[-1])

        self.optimal_params = self.optimizable.unflatten_params(params)
        if self.optimizable.transformation is not None:
            self.optimal_params = self.optimizable.transformation.inverse_transform(
                self.optimal_params
            )
        return OptimizationResult(
            params=self.optimal_params,
            success=True,
            nit=self.num_epochs,
            nfev=self.num_epochs * self.num_batches,
            message=f"Completed {self.num_epochs} epochs.",
            final_loss=float(self.losses[-1]) if self.losses else None,
            loss_history=[float(x) for x in self.losses],
        )

    @property
    def metrics(self):
        return {"loss": self.losses}

batched_objective_flat(params, stochastic_vars_batch_flat)

Mean of the objective function over a batch

Source code in jaxonomy/optimization/framework/optimizers_optax.py

def batched_objective_flat(self, params, stochastic_vars_batch_flat):
    """Mean of the objective function over a batch"""
    return jnp.mean(
        self.optimizable.batched_objective_flat(params, stochastic_vars_batch_flat)
    )

optimize()

Run optimization

Source code in jaxonomy/optimization/framework/optimizers_optax.py

def optimize(self):
    """Run optimization"""
    params = self.optimizable.params_0_flat
    opt_state = self.optimizer.init(params)

    if self.num_batches * self.batch_size == 1:
        # don't randomize over stochastic variables; use a single
        # batch of size 1 with initial stochastic variables
        data_flat, _ = ravel_pytree(self.optimizable.vars_0)
        stochastic_vars_training_data_flat = data_flat[None, None, :]
    else:
        _, stochastic_vars_training_data_flat = self.optimizable.sample_random_vars(
            self.num_batches * self.batch_size
        )

    @jax.jit
    def _scan_fun(carry, stochastic_vars_batch):
        params, opt_state = carry
        params, opt_state, batch_loss = self.step(
            params, opt_state, stochastic_vars_batch
        )
        return (params, opt_state), batch_loss

    for epoch in range(self.num_epochs):
        if self.num_batches * self.batch_size == 1:
            stochastic_vars_batches = stochastic_vars_training_data_flat
        else:
            stochastic_vars_batches = self.optimizable.generate_batches(
                stochastic_vars_training_data_flat,
                self.num_batches,
                self.batch_size,
            )
        (params, opt_state), batch_losses = lax.scan(
            _scan_fun, (params, opt_state), stochastic_vars_batches
        )

        self.losses.append(jnp.mean(batch_losses))
        if self.print_every and epoch % self.print_every == 0:
            p: dict = self.optimizable.unflatten_params(params)
            if self.optimizable.transformation is not None:
                p = self.optimizable.transformation.inverse_transform(p)
            p = {k: v.tolist() for k, v in p.items()}
            logger.info(
                "Epoch %s, average batch loss: %s",
                epoch,
                jnp.mean(batch_losses),
                **logdata(params=p),
            )
        if self.metrics_writer is not None:
            self.metrics_writer.write_metrics(loss=self.losses[-1])

    self.optimal_params = self.optimizable.unflatten_params(params)
    if self.optimizable.transformation is not None:
        self.optimal_params = self.optimizable.transformation.inverse_transform(
            self.optimal_params
        )
    return OptimizationResult(
        params=self.optimal_params,
        success=True,
        nit=self.num_epochs,
        nfev=self.num_epochs * self.num_batches,
        message=f"Completed {self.num_epochs} epochs.",
        final_loss=float(self.losses[-1]) if self.losses else None,
        loss_history=[float(x) for x in self.losses],
    )

step(params, opt_state, stochastic_vars_batch)

Take a single optimization step over one batch

Source code in jaxonomy/optimization/framework/optimizers_optax.py

@partial(jax.jit, static_argnums=(0,))
def step(self, params, opt_state, stochastic_vars_batch):
    """Take a single optimization step over one batch"""
    batch_loss, grads = jax.value_and_grad(self.batched_objective_flat)(
        params, stochastic_vars_batch
    )

    grads = jnp.clip(grads, *self.clip_range) if self.clip_range else grads

    updates, opt_state = self.optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, batch_loss

Optimizable

Bases: OptimizableBase

Base class for all optimizables with no stochastic variables.

For parameters, see OptimizableBase.

The abstract method prepare_context should update the context to incorporate the optimization parameters.

This classs creates methods for evaluation of the objective and constraints from the concrete implementation of the abstract methods. This class also creates methods for batched evaluation of the objective and constraints, which are useful for optimizers that can work with batches (eg. Optax), and population-based optimizers.

Source code in jaxonomy/optimization/framework/base/optimizable.py

class Optimizable(OptimizableBase):
    """
    Base class for all optimizables with no stochastic variables.

    For parameters, see `OptimizableBase`.

    The abstract method `prepare_context` should update the context to incorporate the
    optimization parameters.

    This classs creates methods for evaluation of the objective and constraints from the
    concrete implementation of the abstract methods. This class also creates methods for
    batched evaluation of the objective and constraints, which are useful for optimizers
    that can work with batches (eg. Optax), and population-based optimizers.
    """

    def __init__(
        self,
        diagram,
        base_context,
        sim_t_span=(0.0, 1.0),
        params_0=None,
        bounds=None,
        transformation=None,
        init_min_max=None,
        seed=None,
        sim_options=None,
    ):
        super().__init__(
            diagram,
            base_context,
            sim_t_span,
            params_0,
            bounds,
            transformation,
            init_min_max,
            seed,
            sim_options,
        )
        self.batched_objective = jax.jit(jax.vmap(self.objective, in_axes=(0,)))
        self.batched_objective_flat = jax.jit(
            jax.vmap(self.objective_flat, in_axes=(0,))
        )

        self.batched_constraints = jax.jit(jax.vmap(self.constraints, in_axes=(0,)))
        self.batched_constraints_flat = jax.jit(
            jax.vmap(self.constraints_flat, in_axes=(0,))
        )

    @abstractmethod
    def prepare_context(self, context, params: dict):
        """
        Model-specific updates to incorporate the sample data and parameters.
        Return the updated context.
        """
        pass

    def run_simulation(self, params: dict):
        """
        Run simulation and return final results context.
        """
        context = self.base_context.with_time(self.start_time)
        context = self.prepare_context(context, params)
        results = self.simulator.advance_to(self.stop_time, context)
        return results.context

    def objective_flat(self, params: Array):
        """Objective function for optimization with flattened parameters input"""
        return self.objective(self.unflatten_params(jnp.atleast_1d(params)))

    def objective(self, params: dict):
        """Objective function for optimization with dict parameters input"""
        if self.transformation is not None:
            params = self.transformation.inverse_transform(params)
        results_context = self.run_simulation(params)
        return self.objective_from_context(results_context)

    def constraints_flat(self, params: Array):
        """Constraints function for optimization with flattened parameters input"""
        return self.constraints(self.unflatten_params(jnp.atleast_1d(params)))

    def constraints(self, params: dict):
        """Constraints function for optimization with dict parameters input"""
        if self.transformation is not None:
            params = self.transformation.inverse_transform(params)
        results_context = self.run_simulation(params)
        return self.constraints_from_context(results_context)

constraints(params)

Constraints function for optimization with dict parameters input

Source code in jaxonomy/optimization/framework/base/optimizable.py

def constraints(self, params: dict):
    """Constraints function for optimization with dict parameters input"""
    if self.transformation is not None:
        params = self.transformation.inverse_transform(params)
    results_context = self.run_simulation(params)
    return self.constraints_from_context(results_context)

constraints_flat(params)

Constraints function for optimization with flattened parameters input

Source code in jaxonomy/optimization/framework/base/optimizable.py

def constraints_flat(self, params: Array):
    """Constraints function for optimization with flattened parameters input"""
    return self.constraints(self.unflatten_params(jnp.atleast_1d(params)))

objective(params)

Objective function for optimization with dict parameters input

Source code in jaxonomy/optimization/framework/base/optimizable.py

def objective(self, params: dict):
    """Objective function for optimization with dict parameters input"""
    if self.transformation is not None:
        params = self.transformation.inverse_transform(params)
    results_context = self.run_simulation(params)
    return self.objective_from_context(results_context)

objective_flat(params)

Objective function for optimization with flattened parameters input

Source code in jaxonomy/optimization/framework/base/optimizable.py

def objective_flat(self, params: Array):
    """Objective function for optimization with flattened parameters input"""
    return self.objective(self.unflatten_params(jnp.atleast_1d(params)))

prepare_context(context, params) abstractmethod

Model-specific updates to incorporate the sample data and parameters. Return the updated context.

Source code in jaxonomy/optimization/framework/base/optimizable.py

@abstractmethod
def prepare_context(self, context, params: dict):
    """
    Model-specific updates to incorporate the sample data and parameters.
    Return the updated context.
    """
    pass

run_simulation(params)

Run simulation and return final results context.

Source code in jaxonomy/optimization/framework/base/optimizable.py

def run_simulation(self, params: dict):
    """
    Run simulation and return final results context.
    """
    context = self.base_context.with_time(self.start_time)
    context = self.prepare_context(context, params)
    results = self.simulator.advance_to(self.stop_time, context)
    return results.context

OptimizableWithStochasticVars

Bases: OptimizableBase

Base class for all optimizables with stochastic variables. This is designed only for Optax optimizers and without constraints. Other optimizers are unlikely to work well with stochastic variables.

This class is similar to Optimizable with the key difference that both params and vars (stochastic variables) need to be updated as opposed to params alone

Parameters:

Name	Type	Description	Default
vars_0		dict Initial stochastic variable values. If not provided, the stochastic_vars method will be used to extract these from the base context.	None
distribution_config_vars		DistributionConfig Configuration for stochastic variables. If not provided, standard normal distribution is used.	None

Source code in jaxonomy/optimization/framework/base/optimizable.py

class OptimizableWithStochasticVars(OptimizableBase):
    """
    Base class for all optimizables with stochastic variables. This is designed
    only for Optax optimizers and without constraints. Other optimizers are unlikely to
    work well with stochastic variables.

    This class is similar to `Optimizable` with the key difference that both `params`
    and `vars` (stochastic variables) need to be updated as opposed to `params` alone

    Parameters:
        vars_0: dict
            Initial stochastic variable values. If not provided, the
            `stochastic_vars` method will be used to extract these from the
            base context.
        distribution_config_vars: DistributionConfig
            Configuration for stochastic variables. If not provided, standard normal
            distribution is used.
    """

    def __init__(
        self,
        diagram,
        base_context,
        sim_t_span=(0.0, 1.0),
        params_0=None,
        vars_0=None,
        distribution_config_vars=None,
        bounds=None,
        transformation=None,
        seed=None,
        sim_options=None,
    ):
        super().__init__(
            diagram,
            base_context,
            sim_t_span,
            params_0,
            bounds,
            transformation,
            init_min_max=None,
            seed=seed,
            sim_options=sim_options,
        )

        if vars_0 is None:
            self.vars_0 = self.stochastic_vars(base_context)
        else:
            self.vars_0 = vars_0

        self.vars_0_flat, self.unflatten_vars = ravel_pytree(self.vars_0)
        self.num_stochastic_vars = self.vars_0_flat.size

        self.batched_objective = jax.jit(jax.vmap(self.objective, in_axes=(None, 0)))
        self.batched_objective_flat = jax.jit(
            jax.vmap(self.objective_flat, in_axes=(None, 0))
        )

        if distribution_config_vars is None:
            logger.warning(
                "`distribution_config_vars` is not specified. Using standard normal "
                "as the default distribution"
            )
            self.distribution_config_vars = DistributionConfig(
                names=list(self.vars_0.keys()),
                shapes=[jnp.shape(x) for x in self.vars_0.values()],
                distributions=["normal"] * len(self.vars_0),
                distributions_configs=[{}] * len(self.vars_0),
            )
        else:
            self.distribution_config_vars = distribution_config_vars

    @abstractmethod
    def prepare_context(self, context, params: dict, vars: dict):
        """
        Model-specific updates to incorporate the parameters and stochastic vars.
        Return the updated context.
        """
        pass

    @abstractmethod
    def stochastic_vars(self, context) -> dict:
        """
        Extract stochastic `vars` from the context.
        These should be in the form of a dict of Pytrees.
        """
        pass

    def run_simulation(self, params: dict, vars: dict):
        """Run simulation and return final results context."""
        context = self.base_context.with_time(self.start_time)
        context = self.prepare_context(context, params, vars)
        results = self.simulator.advance_to(self.stop_time, context)
        return results.context

    def objective_flat(self, params: Array, vars: Array):
        """Objective function for optimization with flattened parameters and vars
        input"""
        return self.objective(
            self.unflatten_params(jnp.atleast_1d(params)),
            self.unflatten_vars(jnp.atleast_1d(vars)),
        )

    def objective(self, params: dict, vars: dict):
        """Objective function for optimization with dict parameters and vars input"""
        if self.transformation is not None:
            params = self.transformation.inverse_transform(params)
        results_context = self.run_simulation(params, vars)
        return self.objective_from_context(results_context)

    def sample_random_vars(self, num_samples):
        """Generate random samples of the stochastic variables"""
        names = self.distribution_config_vars.names
        shapes = self.distribution_config_vars.shapes
        distributions = self.distribution_config_vars.distributions
        distributions_configs = self.distribution_config_vars.distributions_configs
        data, flat_data = self._generate_random_data(
            names,
            shapes,
            distributions,
            distributions_configs,
            num_samples,
        )
        return data, flat_data

    def generate_batches(
        self,
        data,
        num_batches,
        batch_size,
    ):
        """
        Given all samples `data`, generate `num_batches` random batches of size
        `batch_size` each
        """
        num_samples = data.shape[0]
        self.key, subkey = jr.split(self.key)
        batch_indices = jax.random.choice(
            subkey, num_samples, (num_batches, batch_size), replace=True
        )
        batches = data[batch_indices]
        return batches

    @staticmethod
    def _distribution(name: str, key, shape, options: dict):
        # remap options from json names to jax.random names
        # NOTE: code users should be able to pass default argnames (i.e. those used
        # by jax.random), for example, "minval" and "maxval" for uniform distribution
        if name == "normal":
            mean = options.get("mean", 0.0)
            std_dev = options.get("std_dev", 1.0)
            return jr.normal(key, shape) * std_dev + mean

        if name == "lognormal":
            mean = options.get("mean", 0.0)
            sigma = options.get("std_dev", 1.0)
            return jr.lognormal(key, sigma=sigma, shape=shape) + mean

        if name == "uniform":
            minval = options.get("min", 0.0)
            maxval = options.get("max", 1.0)
            return jr.uniform(key, shape=shape, minval=minval, maxval=maxval)

        warnings.warn(f"Unknown distribution: {name}.")
        sample_func = getattr(jr, name)
        return sample_func(key, shape, **options)

    def _generate_random_data(
        self,
        names,
        shapes,
        distributions,
        distributions_configs,
        num_samples,
    ):
        data = {}
        self.key, *subkeys = jr.split(self.key, len(names) + 1)
        for key, name, shape, distribution, distribution_config in zip(
            subkeys, names, shapes, distributions, distributions_configs
        ):
            data[name] = self._distribution(
                distribution,
                key,
                (num_samples, *shape),
                distribution_config,
            )

        def _flatten(x):
            x_flat, _ = ravel_pytree(x)
            return x_flat

        return data, _flatten(data)

generate_batches(data, num_batches, batch_size)

Given all samples data, generate num_batches random batches of size batch_size each

Source code in jaxonomy/optimization/framework/base/optimizable.py

def generate_batches(
    self,
    data,
    num_batches,
    batch_size,
):
    """
    Given all samples `data`, generate `num_batches` random batches of size
    `batch_size` each
    """
    num_samples = data.shape[0]
    self.key, subkey = jr.split(self.key)
    batch_indices = jax.random.choice(
        subkey, num_samples, (num_batches, batch_size), replace=True
    )
    batches = data[batch_indices]
    return batches

objective(params, vars)

Objective function for optimization with dict parameters and vars input

Source code in jaxonomy/optimization/framework/base/optimizable.py

def objective(self, params: dict, vars: dict):
    """Objective function for optimization with dict parameters and vars input"""
    if self.transformation is not None:
        params = self.transformation.inverse_transform(params)
    results_context = self.run_simulation(params, vars)
    return self.objective_from_context(results_context)

objective_flat(params, vars)

Objective function for optimization with flattened parameters and vars input

Source code in jaxonomy/optimization/framework/base/optimizable.py

def objective_flat(self, params: Array, vars: Array):
    """Objective function for optimization with flattened parameters and vars
    input"""
    return self.objective(
        self.unflatten_params(jnp.atleast_1d(params)),
        self.unflatten_vars(jnp.atleast_1d(vars)),
    )

prepare_context(context, params, vars) abstractmethod

Model-specific updates to incorporate the parameters and stochastic vars. Return the updated context.

Source code in jaxonomy/optimization/framework/base/optimizable.py

@abstractmethod
def prepare_context(self, context, params: dict, vars: dict):
    """
    Model-specific updates to incorporate the parameters and stochastic vars.
    Return the updated context.
    """
    pass

run_simulation(params, vars)

Run simulation and return final results context.

Source code in jaxonomy/optimization/framework/base/optimizable.py

def run_simulation(self, params: dict, vars: dict):
    """Run simulation and return final results context."""
    context = self.base_context.with_time(self.start_time)
    context = self.prepare_context(context, params, vars)
    results = self.simulator.advance_to(self.stop_time, context)
    return results.context

sample_random_vars(num_samples)

Generate random samples of the stochastic variables

Source code in jaxonomy/optimization/framework/base/optimizable.py

def sample_random_vars(self, num_samples):
    """Generate random samples of the stochastic variables"""
    names = self.distribution_config_vars.names
    shapes = self.distribution_config_vars.shapes
    distributions = self.distribution_config_vars.distributions
    distributions_configs = self.distribution_config_vars.distributions_configs
    data, flat_data = self._generate_random_data(
        names,
        shapes,
        distributions,
        distributions_configs,
        num_samples,
    )
    return data, flat_data

stochastic_vars(context) abstractmethod

Extract stochastic vars from the context. These should be in the form of a dict of Pytrees.

Source code in jaxonomy/optimization/framework/base/optimizable.py

@abstractmethod
def stochastic_vars(self, context) -> dict:
    """
    Extract stochastic `vars` from the context.
    These should be in the form of a dict of Pytrees.
    """
    pass

OptimizationResult dataclass

Unified result returned by all jaxonomy optimizers.

Supports dict-like access (result["param"]) for backward compatibility with code that treated the old return value as a plain parameter dict.

Attributes:

Name	Type	Description
params	dict[str, Any]	dict mapping parameter name → optimized value (same as the dict that optimizers used to return directly).
success	bool	True if the optimizer reported convergence.
nit	int	Number of iterations (or epochs / generations).
nfev	int	Number of objective-function evaluations.
message	str	Human-readable status message from the optimizer.
final_loss	float | None	Objective value at the optimum. None when not available (e.g. population-based methods that track fitness separately).
loss_history	list[float]	Sequence of objective values recorded during optimization (one per epoch / generation).

Source code in jaxonomy/optimization/framework/base/optimizer.py

@dataclass
class OptimizationResult:
    """
    Unified result returned by all jaxonomy optimizers.

    Supports dict-like access (``result["param"]``) for backward compatibility
    with code that treated the old return value as a plain parameter dict.

    Attributes:
        params: dict mapping parameter name → optimized value (same as the
            dict that optimizers used to return directly).
        success: ``True`` if the optimizer reported convergence.
        nit: Number of iterations (or epochs / generations).
        nfev: Number of objective-function evaluations.
        message: Human-readable status message from the optimizer.
        final_loss: Objective value at the optimum.  ``None`` when not
            available (e.g. population-based methods that track fitness
            separately).
        loss_history: Sequence of objective values recorded during
            optimization (one per epoch / generation).
    """

    params: dict[str, Any]
    success: bool = True
    nit: int = 0
    nfev: int = 0
    message: str = ""
    final_loss: float | None = None
    loss_history: list[float] = field(default_factory=list)

    # ------------------------------------------------------------------
    # Backward-compatible dict-like interface
    # ------------------------------------------------------------------
    def __getitem__(self, key: str) -> Any:
        return self.params[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self.params[key] = value

    def __contains__(self, key: object) -> bool:
        return key in self.params

    def __iter__(self):
        return iter(self.params)

    def __len__(self) -> int:
        return len(self.params)

    def items(self):
        return self.params.items()

    def keys(self):
        return self.params.keys()

    def values(self):
        return self.params.values()

    def get(self, key: str, default: Any = None) -> Any:
        return self.params.get(key, default)

    def __repr__(self) -> str:
        return (
            f"OptimizationResult("
            f"params={self.params}, "
            f"success={self.success}, "
            f"nit={self.nit}, "
            f"nfev={self.nfev}, "
            f"final_loss={self.final_loss}, "
            f"message={self.message!r}"
            f")"
        )

RLEnv

Bases: Env

Base class for reinforcement learning environments in Jaxonomy.

Source code in jaxonomy/optimization/rl_env.py

class RLEnv(BraxEnv):
    """Base class for reinforcement learning environments in Jaxonomy."""

    def __init__(self, plant: SystemBase, act_size: int, dt: float):
        if len(plant.input_ports) != 1:
            raise ValueError("Plant must have exactly one input port.")

        if len(plant.output_ports) != 1:
            raise ValueError("Plant must have exactly one output port.")

        self._act_size = act_size
        self._plant = plant
        self._plant_id = plant.system_id
        self.dt = dt

        # Embed the plant within a simple wrapper diagram that has a constant input
        input_name = "const_in"
        self.diagram = _wrapper_diagram(plant, act_size, input_name=input_name)
        input_id = self.diagram[input_name].system_id

        self._forward = partial(_forward, self.diagram, dt, input_id)
        self._get_obs = partial(_get_obs, plant, input_id)

        self._static_context = self.diagram.create_context()
        obs = self._get_obs(self._static_context, jnp.zeros(self._act_size))
        self._obs_size = obs.size

    @partial(jax.jit, static_argnums=0)
    def reset(self, rng: jax.Array) -> RLState:
        pipeline_state = self._static_context

        # Randomize the plant context as defined by the user-provided randomize function
        pipeline_state = self.randomize(pipeline_state, rng)

        obs = self._get_obs(pipeline_state, jnp.zeros(self._act_size))
        reward, done = jnp.zeros(2)
        metrics = {}
        return RLState(pipeline_state, obs, reward, done, metrics)  # pylint: disable=too-many-function-args

    @partial(jax.jit, static_argnums=0)
    def step(self, state: RLState, action: jax.Array) -> RLState:
        pipeline_state = self._forward(state.pipeline_state, action)
        obs = self._get_obs(pipeline_state, action)
        reward = self.get_reward(pipeline_state, obs, action)
        done = self.get_done(pipeline_state, obs)
        return state.replace(
            pipeline_state=pipeline_state, obs=obs, reward=reward, done=done
        )

    @property
    def action_size(self) -> int:
        return self._act_size

    @property
    def observation_size(self) -> int:
        return self._obs_size

    @property
    def backend(self) -> str:
        return "jaxonomy"

    #
    # To be overridden by subclasses
    #
    @abc.abstractmethod
    def get_reward(
        self, pipeline_state: ContextBase, obs: jax.Array, act: jax.Array
    ) -> jax.Array:
        """Return the reward for the current state and observation."""
        pass

    def get_done(self, pipeline_state: ContextBase, obs: jax.Array) -> jax.Array:
        """Return a boolean indicating whether the episode is done."""
        return 0.0

    def randomize(self, pipeline_state: ContextBase, rng: jax.Array) -> ContextBase:
        """Randomize the initial states, parameters, etc."""
        return pipeline_state

    def render(
        self,
        trajectory: list[RLState],
        height: int = 240,
        width: int = 320,
        camera: str = None,
    ) -> list[np.ndarray]:
        """Render the trajectory"""
        raise NotImplementedError(
            "Rendering is not yet supported for Jaxonomy environments."
        )

get_done(pipeline_state, obs)

Return a boolean indicating whether the episode is done.

Source code in jaxonomy/optimization/rl_env.py

def get_done(self, pipeline_state: ContextBase, obs: jax.Array) -> jax.Array:
    """Return a boolean indicating whether the episode is done."""
    return 0.0

get_reward(pipeline_state, obs, act) abstractmethod

Return the reward for the current state and observation.

Source code in jaxonomy/optimization/rl_env.py

@abc.abstractmethod
def get_reward(
    self, pipeline_state: ContextBase, obs: jax.Array, act: jax.Array
) -> jax.Array:
    """Return the reward for the current state and observation."""
    pass

randomize(pipeline_state, rng)

Randomize the initial states, parameters, etc.

Source code in jaxonomy/optimization/rl_env.py

def randomize(self, pipeline_state: ContextBase, rng: jax.Array) -> ContextBase:
    """Randomize the initial states, parameters, etc."""
    return pipeline_state

render(trajectory, height=240, width=320, camera=None)

Render the trajectory

Source code in jaxonomy/optimization/rl_env.py

def render(
    self,
    trajectory: list[RLState],
    height: int = 240,
    width: int = 320,
    camera: str = None,
) -> list[np.ndarray]:
    """Render the trajectory"""
    raise NotImplementedError(
        "Rendering is not yet supported for Jaxonomy environments."
    )

Scipy

Bases: Optimizer

Scipy/JAX-scipy optimizers.

Parameters:

Name	Type	Description	Default
optimizable	Optimizable	The optimizable object.	required
opt_method	str	The optimization method to use.	required
tol	float	Tolerance for termination. For detailed control, use opt_method_config.	None
opt_method_config	dict	Configuration for the optimization method.	None
use_autodiff_grad	bool	Whether to use autodiff for gradient computation.	True
use_jax_scipy	bool	Whether to use JAX's version of optimize.minimize.	False

Source code in jaxonomy/optimization/framework/optimizers_scipy.py

class Scipy(Optimizer):
    """
    Scipy/JAX-scipy optimizers.

    Parameters:
        optimizable (Optimizable):
            The optimizable object.
        opt_method (str):
            The optimization method to use.
        tol (float):
            Tolerance for termination. For detailed control, use `opt_method_config`.
        opt_method_config (dict):
            Configuration for the optimization method.
        use_autodiff_grad (bool):
            Whether to use autodiff for gradient computation.
        use_jax_scipy (bool):
            Whether to use JAX's version of `optimize.minimize`.
    """

    def __init__(
        self,
        optimizable: Optimizable,
        opt_method,
        tol=None,
        opt_method_config=None,
        use_autodiff_grad=True,
        use_jax_scipy=False,
        metrics_writer: MetricsWriter = None,
    ):
        self.optimizable = optimizable
        self.opt_method = opt_method
        self.tol = tol
        self.opt_method_config = opt_method_config or {}
        self.use_autodiff_grad = use_autodiff_grad
        self.use_jax_scipy = use_jax_scipy
        self.optimal_params = None
        self.metrics_writer = metrics_writer
        self._loss_history: list[float] = []

    def optimize(self):
        """Run optimization"""
        params = self.optimizable.params_0_flat
        objective = jax.jit(self.optimizable.objective_flat)

        _success = True
        _nit = 0
        _nfev = 0
        _message = ""
        _final_loss = float("nan")

        if self.use_jax_scipy:
            warnings.warn(
                "`use_jax_scipy` is True. JAX's version of optimize.minimize will be "
                "used. Consequently, `opt_method` will be set of `BFGS` and autodiff "
                "will be used for gradient computation. Constraints and bounds will "
                "be ignored. If you want to use scipy's version of minimize, set "
                " `use_jax_scipy` to False."
            )
            opt_res = jax_scipy_opt.minimize(
                objective,
                params,
                method="BFGS",
                tol=self.tol,
                options=self.opt_method_config,
            )
            params = opt_res.x
            _success = bool(getattr(opt_res, 'success', True))
            _nit = int(getattr(opt_res, 'nit', 0))
            _nfev = int(getattr(opt_res, 'nfev', 0))
            _final_loss = float(getattr(opt_res, 'fun', float('nan')))

        else:
            use_jac = False
            if self.opt_method in ACCEPTS_GRAD and self.use_autodiff_grad:
                jac = jax.jit(jax.grad(objective))
                use_jac = True

            # Handle bounds
            bounds = self.optimizable.bounds_flat

            # Jobs from UI would put (-jnp.inf, jnp.inf) as defualt bounds. The user
            # may also have specified bounds this way. Scipy expects `None` to imply
            # unboundedness.
            if bounds is not None:
                bounds = [
                    (
                        None if b[0] == -jnp.inf else b[0],
                        None if b[1] == jnp.inf else b[1],
                    )
                    for b in bounds
                ]

                # Check if all bounds are None, i.e. no bounds at all, and hence
                # algorithms that do not support bounds can be used.
                flattened_bounds = [element for tup in bounds for element in tup]
                all_none = all(element is None for element in flattened_bounds)
                bounds = None if all_none else bounds

            if bounds is not None and self.opt_method not in SUPPORTS_BOUNDS:
                raise ValueError(
                    f"Optimization method scipy:{self.opt_method} "
                    "does not support bounds."
                )

            # Handle constraints
            if (
                self.optimizable.has_constraints
                and self.opt_method not in SUPPORTS_CONSTRAINTS
            ):
                raise ValueError(
                    f"Optimization method scipy:{self.opt_method} "
                    "does not support constraints."
                )

            if self.optimizable.has_constraints:
                constraints = jax.jit(self.optimizable.constraints_flat)
                constraints_jac = jax.jit(jax.jacrev(constraints))
                constraints = sciopt.NonlinearConstraint(
                    constraints, 0.0, jnp.inf, jac=constraints_jac
                )
            else:
                constraints = None

            if self.metrics_writer is not None:
                cb = (
                    self._scipy_callback_new
                    if self.opt_method in MINIMIZE_METHODS_NEW_CB
                    else partial(self._scipy_callback_legacy, objective)
                )
            else:
                cb = None

            opt_res: "sciopt.OptimizeResult" = sciopt.minimize(
                objective,
                params,
                method=self.opt_method,
                jac=jac if use_jac else None,
                bounds=bounds,
                constraints=constraints,
                tol=self.tol,
                options=self.opt_method_config,
                callback=cb,
            )

            params = opt_res.x

            # Show the raw information from scipy. This can help with debugging.
            logger.info("Optimization result:\n%s", opt_res)

            if not opt_res.success:
                logger.warning("Optimization did not converge: %s", opt_res.message)

            _nit = int(getattr(opt_res, 'nit', 0))
            _nfev = int(getattr(opt_res, 'nfev', 0))
            _success = bool(getattr(opt_res, 'success', False))
            _message = str(getattr(opt_res, 'message', ''))
            _final_loss = float(getattr(opt_res, 'fun', float('nan')))

        self.optimal_params = self.optimizable.unflatten_params(params)
        if self.optimizable.transformation is not None:
            self.optimal_params = self.optimizable.transformation.inverse_transform(
                self.optimal_params
            )
        return OptimizationResult(
            params=self.optimal_params,
            success=_success,
            nit=_nit,
            nfev=_nfev,
            message=_message,
            final_loss=_final_loss,
            loss_history=list(self._loss_history),
        )

    # NOTE: if this turns out to be too expensive, we can throttle writes in the
    # MetricsWriter and only compute metrics when we need them.
    def _write_metrics(self, fun, x):
        metrics = {}
        if fun is not None:
            metrics["fun"] = fun
            self._loss_history.append(float(fun))
        if x is not None:
            params: dict = self.optimizable.unflatten_params(x)
            for k, v in params.items():
                if np.asarray(v).shape == ():
                    metrics[k] = v
        if len(metrics) > 0:
            self.metrics_writer.write_metrics(**metrics)

    def _scipy_callback_new(self, intermediate_result: "sciopt.OptimizeResult"):
        fun = intermediate_result.get("fun")
        if fun is not None:
            self._loss_history.append(float(fun))
        self._write_metrics(
            intermediate_result.get("fun"), intermediate_result.get("x")
        )

    def _scipy_callback_legacy(self, objective, intermediate_results: np.ndarray):
        fun = objective(intermediate_results)
        self._loss_history.append(float(fun))
        self._write_metrics(fun, intermediate_results)

optimize()

Run optimization

Source code in jaxonomy/optimization/framework/optimizers_scipy.py

def optimize(self):
    """Run optimization"""
    params = self.optimizable.params_0_flat
    objective = jax.jit(self.optimizable.objective_flat)

    _success = True
    _nit = 0
    _nfev = 0
    _message = ""
    _final_loss = float("nan")

    if self.use_jax_scipy:
        warnings.warn(
            "`use_jax_scipy` is True. JAX's version of optimize.minimize will be "
            "used. Consequently, `opt_method` will be set of `BFGS` and autodiff "
            "will be used for gradient computation. Constraints and bounds will "
            "be ignored. If you want to use scipy's version of minimize, set "
            " `use_jax_scipy` to False."
        )
        opt_res = jax_scipy_opt.minimize(
            objective,
            params,
            method="BFGS",
            tol=self.tol,
            options=self.opt_method_config,
        )
        params = opt_res.x
        _success = bool(getattr(opt_res, 'success', True))
        _nit = int(getattr(opt_res, 'nit', 0))
        _nfev = int(getattr(opt_res, 'nfev', 0))
        _final_loss = float(getattr(opt_res, 'fun', float('nan')))

    else:
        use_jac = False
        if self.opt_method in ACCEPTS_GRAD and self.use_autodiff_grad:
            jac = jax.jit(jax.grad(objective))
            use_jac = True

        # Handle bounds
        bounds = self.optimizable.bounds_flat

        # Jobs from UI would put (-jnp.inf, jnp.inf) as defualt bounds. The user
        # may also have specified bounds this way. Scipy expects `None` to imply
        # unboundedness.
        if bounds is not None:
            bounds = [
                (
                    None if b[0] == -jnp.inf else b[0],
                    None if b[1] == jnp.inf else b[1],
                )
                for b in bounds
            ]

            # Check if all bounds are None, i.e. no bounds at all, and hence
            # algorithms that do not support bounds can be used.
            flattened_bounds = [element for tup in bounds for element in tup]
            all_none = all(element is None for element in flattened_bounds)
            bounds = None if all_none else bounds

        if bounds is not None and self.opt_method not in SUPPORTS_BOUNDS:
            raise ValueError(
                f"Optimization method scipy:{self.opt_method} "
                "does not support bounds."
            )

        # Handle constraints
        if (
            self.optimizable.has_constraints
            and self.opt_method not in SUPPORTS_CONSTRAINTS
        ):
            raise ValueError(
                f"Optimization method scipy:{self.opt_method} "
                "does not support constraints."
            )

        if self.optimizable.has_constraints:
            constraints = jax.jit(self.optimizable.constraints_flat)
            constraints_jac = jax.jit(jax.jacrev(constraints))
            constraints = sciopt.NonlinearConstraint(
                constraints, 0.0, jnp.inf, jac=constraints_jac
            )
        else:
            constraints = None

        if self.metrics_writer is not None:
            cb = (
                self._scipy_callback_new
                if self.opt_method in MINIMIZE_METHODS_NEW_CB
                else partial(self._scipy_callback_legacy, objective)
            )
        else:
            cb = None

        opt_res: "sciopt.OptimizeResult" = sciopt.minimize(
            objective,
            params,
            method=self.opt_method,
            jac=jac if use_jac else None,
            bounds=bounds,
            constraints=constraints,
            tol=self.tol,
            options=self.opt_method_config,
            callback=cb,
        )

        params = opt_res.x

        # Show the raw information from scipy. This can help with debugging.
        logger.info("Optimization result:\n%s", opt_res)

        if not opt_res.success:
            logger.warning("Optimization did not converge: %s", opt_res.message)

        _nit = int(getattr(opt_res, 'nit', 0))
        _nfev = int(getattr(opt_res, 'nfev', 0))
        _success = bool(getattr(opt_res, 'success', False))
        _message = str(getattr(opt_res, 'message', ''))
        _final_loss = float(getattr(opt_res, 'fun', float('nan')))

    self.optimal_params = self.optimizable.unflatten_params(params)
    if self.optimizable.transformation is not None:
        self.optimal_params = self.optimizable.transformation.inverse_transform(
            self.optimal_params
        )
    return OptimizationResult(
        params=self.optimal_params,
        success=_success,
        nit=_nit,
        nfev=_nfev,
        message=_message,
        final_loss=_final_loss,
        loss_history=list(self._loss_history),
    )

SensitivityResult dataclass

Result of a parameter sensitivity / identifiability analysis.

All arrays are plain numpy.ndarray for easy inspection.

Attributes

param_names : list[str] Parameter names, in the same order as the flat parameter vector. params_0 : dict[str, Any] The parameter values at which the analysis was performed. objective_value : float Objective value at params_0. gradients : ndarray, shape (n_params,) Gradient of the objective w.r.t. each parameter. normalized_sensitivity : ndarray, shape (n_params,) |p_i * ∂L/∂p_i| — relative sensitivity. Dimensionless and comparable across parameters with different scales. hessian : ndarray, shape (n_params, n_params) Hessian of the objective (FIM approximation). NaN-filled when compute_hessian=False. hessian_diagonal : ndarray, shape (n_params,) Diagonal of the Hessian. eigenvalues : ndarray, shape (n_params,) Eigenvalues of the Hessian (ascending). condition_number : float Ratio of largest to smallest non-negligible eigenvalue. Large values (> 1e6) indicate near-collinear parameters. unidentifiable_params : list[str] Parameter names whose normalised sensitivity is below sensitivity_threshold * max_sensitivity. sensitivity_threshold : float Relative threshold used to flag unidentifiable parameters.

Source code in jaxonomy/optimization/sensitivity.py

@dataclass
class SensitivityResult:
    """
    Result of a parameter sensitivity / identifiability analysis.

    All arrays are plain ``numpy.ndarray`` for easy inspection.

    Attributes
    ----------
    param_names : list[str]
        Parameter names, in the same order as the flat parameter vector.
    params_0 : dict[str, Any]
        The parameter values at which the analysis was performed.
    objective_value : float
        Objective value at ``params_0``.
    gradients : ndarray, shape (n_params,)
        Gradient of the objective w.r.t. each parameter.
    normalized_sensitivity : ndarray, shape (n_params,)
        ``|p_i * ∂L/∂p_i|`` — relative sensitivity.  Dimensionless and
        comparable across parameters with different scales.
    hessian : ndarray, shape (n_params, n_params)
        Hessian of the objective (FIM approximation).  ``NaN``-filled when
        ``compute_hessian=False``.
    hessian_diagonal : ndarray, shape (n_params,)
        Diagonal of the Hessian.
    eigenvalues : ndarray, shape (n_params,)
        Eigenvalues of the Hessian (ascending).
    condition_number : float
        Ratio of largest to smallest non-negligible eigenvalue.  Large values
        (> 1e6) indicate near-collinear parameters.
    unidentifiable_params : list[str]
        Parameter names whose normalised sensitivity is below
        ``sensitivity_threshold * max_sensitivity``.
    sensitivity_threshold : float
        Relative threshold used to flag unidentifiable parameters.
    """

    param_names: list[str]
    params_0: dict[str, Any]
    objective_value: float
    gradients: np.ndarray
    normalized_sensitivity: np.ndarray
    hessian: np.ndarray
    hessian_diagonal: np.ndarray
    eigenvalues: np.ndarray
    condition_number: float
    unidentifiable_params: list[str]
    sensitivity_threshold: float = 1e-3

    def summary(self) -> str:
        """Return a human-readable summary table."""
        lines = [
            "=== Parameter Sensitivity / Identifiability Analysis ===",
            f"Objective at params_0: {self.objective_value:.6g}",
            f"Hessian condition number: {self.condition_number:.3g}",
            "",
            f"{'Parameter':<22} {'Gradient':>14} {'Norm. Sensitivity':>20} "
            f"{'Status':>12}",
            "-" * 70,
        ]
        max_s = float(np.max(self.normalized_sensitivity)) if len(self.normalized_sensitivity) else 1.0
        for name, g, s in zip(
            self.param_names, self.gradients, self.normalized_sensitivity
        ):
            relative = s / max(max_s, 1e-30)
            status = "✓ ok" if name not in self.unidentifiable_params else "✗ LOW"
            lines.append(
                f"{name:<22} {g:>14.4e} {s:>20.4e} {status:>12}"
            )
        if self.unidentifiable_params:
            lines.append(
                f"\n⚠  Low-sensitivity parameters "
                f"(threshold={self.sensitivity_threshold:.1e}): "
                f"{self.unidentifiable_params}"
            )
        if not np.any(np.isnan(self.eigenvalues)):
            lines.append(f"\nHessian eigenvalues: {self.eigenvalues}")
        return "\n".join(lines)

    def __repr__(self) -> str:
        return self.summary()

summary()

Return a human-readable summary table.

Source code in jaxonomy/optimization/sensitivity.py

def summary(self) -> str:
    """Return a human-readable summary table."""
    lines = [
        "=== Parameter Sensitivity / Identifiability Analysis ===",
        f"Objective at params_0: {self.objective_value:.6g}",
        f"Hessian condition number: {self.condition_number:.3g}",
        "",
        f"{'Parameter':<22} {'Gradient':>14} {'Norm. Sensitivity':>20} "
        f"{'Status':>12}",
        "-" * 70,
    ]
    max_s = float(np.max(self.normalized_sensitivity)) if len(self.normalized_sensitivity) else 1.0
    for name, g, s in zip(
        self.param_names, self.gradients, self.normalized_sensitivity
    ):
        relative = s / max(max_s, 1e-30)
        status = "✓ ok" if name not in self.unidentifiable_params else "✗ LOW"
        lines.append(
            f"{name:<22} {g:>14.4e} {s:>20.4e} {status:>12}"
        )
    if self.unidentifiable_params:
        lines.append(
            f"\n⚠  Low-sensitivity parameters "
            f"(threshold={self.sensitivity_threshold:.1e}): "
            f"{self.unidentifiable_params}"
        )
    if not np.any(np.isnan(self.eigenvalues)):
        lines.append(f"\nHessian eigenvalues: {self.eigenvalues}")
    return "\n".join(lines)

Trainer

Base class for optimizing model parameters via simulation.

Should probably get a more descriptive name once we're doing other kinds of training...

Source code in jaxonomy/optimization/training.py

class Trainer:
    """Base class for optimizing model parameters via simulation.

    Should probably get a more descriptive name once we're doing other kinds
    of training...
    """

    def __init__(
        self,
        simulator: Simulator,
        context,
        optimizer="adamw",
        lr=1e-3,
        print_every=10,
        clip_range=(-10.0, 10.0),
        **opt_kwargs,
    ):
        self.simulator = simulator
        self.context = context
        self.opt_state = None

        # See https://optax.readthedocs.io/en/latest/api.html for supported optimizers
        self.optimizer = getattr(optax, optimizer)(lr, **opt_kwargs)

        self.print_every = print_every
        self.clip_range = clip_range

    @abc.abstractmethod
    def optimizable_parameters(self, context):
        """Extract optimizable model-specific parameters from the context.

        These should be in the form of a PyTree (e.g. tuple, dict, array, etc)
        and should be the first arguments to `prepare_context`.
        """
        pass

    @abc.abstractmethod
    def prepare_context(self, context, *data, key=None):
        """Model-specific updates to incorporate the sample data and parameters.

        `data` should be the combination of the output of `optimizable_parameters`
        along with all the per-simulation "training data".  Parameters will
        update once per epoch, and training data will update once per sample.
        """
        pass

    @abc.abstractmethod
    def evaluate_cost(self, context):
        """Model-specific cost function, evaluated on final context"""
        pass

    def make_forward(self, start_time, stop_time):
        """Create a generic forward pass through the simulation, returning loss"""

        # Take all the data and model parameters, run a simulation, return loss.
        def _simulate(key, *data):
            context = self.context.with_time(start_time)
            context = self.prepare_context(context, *data, key=key)
            results = self.simulator.advance_to(stop_time, context)
            return self.evaluate_cost(results.context)

        return _simulate

    def make_loss_fn(self, forward, params):
        """Create a loss function based on a forward pass of the simulation

        `params` here can be any PyTree - it will get flattened to a single array
        """
        # Flatten all optimizable parameters into a single array
        p0, unflatten = ravel_pytree(params)

        # Define the loss as the mean cost function over the data set
        def _loss(p, key, *batch_data):
            # Map the forward pass over all the data points and return the loss
            loss = batch_scan(partial(forward, key, unflatten(p)), *batch_data)
            return loss

        # JIT compile the loss function and return the initial parameter
        # array and unflatten function
        return jax.jit(_loss), p0, unflatten

    def train(
        self,
        training_data,
        sim_start_time,
        sim_stop_time,
        epochs=100,
        key=None,
        params=None,
        opt_state=None,
    ):
        """Run the optimization loop over the training data"""

        if (
            self.simulator.max_major_steps is None
            or self.simulator.max_major_steps <= 0
        ):
            self.simulator.max_major_steps = estimate_max_major_steps(
                self.simulator.system,
                (sim_start_time, sim_stop_time),
                self.simulator.max_major_step_length,
            )

        if key is None:
            key = jax.random.PRNGKey(np.random.randint(0, 2**32, dtype=np.int64))

        # Create a function to evaluate the forward pass through the simulation
        forward = self.make_forward(sim_start_time, sim_stop_time)

        # Pull out the optimizable parameters from the context
        if params is None:
            params = self.optimizable_parameters(self.context)

        # Initialize the optimizer and create the loss function
        loss, p, unflatten = self.make_loss_fn(forward, params)

        if opt_state is None:
            opt_state = self.optimizer.init(p)

        self.opt_state = opt_state

        @jax.jit
        def opt_step(p, opt_state, key, batch_data):
            key, subkey = jax.random.split(key)
            if batch_data:
                loss_value, grads = jax.value_and_grad(loss)(p, subkey, *batch_data)
            else:
                loss_value, grads = jax.value_and_grad(loss)(p, subkey)

            grads = jnp.clip(grads, *self.clip_range)

            updates, opt_state = self.optimizer.update(grads, opt_state, p)
            p = optax.apply_updates(p, updates)
            return p, opt_state, key, loss_value

        def _scan_fun(carry, batch_data):
            p, opt_state, key, loss_value = opt_step(*carry, batch_data)
            return (p, opt_state, key), loss_value

        # Run the optimization loop
        for epoch in range(epochs):
            (p, self.opt_state, key), batch_loss = jax.lax.scan(
                _scan_fun, (p, self.opt_state, key), training_data
            )

            if epoch % self.print_every == 0:
                logger.info("Epoch %s, loss: %s", epoch, jnp.mean(batch_loss))

        # Return the optimized parameters
        return unflatten(p)

evaluate_cost(context) abstractmethod

Model-specific cost function, evaluated on final context

Source code in jaxonomy/optimization/training.py

@abc.abstractmethod
def evaluate_cost(self, context):
    """Model-specific cost function, evaluated on final context"""
    pass

make_forward(start_time, stop_time)

Create a generic forward pass through the simulation, returning loss

Source code in jaxonomy/optimization/training.py

def make_forward(self, start_time, stop_time):
    """Create a generic forward pass through the simulation, returning loss"""

    # Take all the data and model parameters, run a simulation, return loss.
    def _simulate(key, *data):
        context = self.context.with_time(start_time)
        context = self.prepare_context(context, *data, key=key)
        results = self.simulator.advance_to(stop_time, context)
        return self.evaluate_cost(results.context)

    return _simulate

make_loss_fn(forward, params)

Create a loss function based on a forward pass of the simulation

params here can be any PyTree - it will get flattened to a single array

Source code in jaxonomy/optimization/training.py

def make_loss_fn(self, forward, params):
    """Create a loss function based on a forward pass of the simulation

    `params` here can be any PyTree - it will get flattened to a single array
    """
    # Flatten all optimizable parameters into a single array
    p0, unflatten = ravel_pytree(params)

    # Define the loss as the mean cost function over the data set
    def _loss(p, key, *batch_data):
        # Map the forward pass over all the data points and return the loss
        loss = batch_scan(partial(forward, key, unflatten(p)), *batch_data)
        return loss

    # JIT compile the loss function and return the initial parameter
    # array and unflatten function
    return jax.jit(_loss), p0, unflatten

optimizable_parameters(context) abstractmethod

Extract optimizable model-specific parameters from the context.

These should be in the form of a PyTree (e.g. tuple, dict, array, etc) and should be the first arguments to prepare_context.

Source code in jaxonomy/optimization/training.py

@abc.abstractmethod
def optimizable_parameters(self, context):
    """Extract optimizable model-specific parameters from the context.

    These should be in the form of a PyTree (e.g. tuple, dict, array, etc)
    and should be the first arguments to `prepare_context`.
    """
    pass

prepare_context(context, *data, key=None) abstractmethod

Model-specific updates to incorporate the sample data and parameters.

data should be the combination of the output of optimizable_parameters along with all the per-simulation "training data". Parameters will update once per epoch, and training data will update once per sample.

Source code in jaxonomy/optimization/training.py

@abc.abstractmethod
def prepare_context(self, context, *data, key=None):
    """Model-specific updates to incorporate the sample data and parameters.

    `data` should be the combination of the output of `optimizable_parameters`
    along with all the per-simulation "training data".  Parameters will
    update once per epoch, and training data will update once per sample.
    """
    pass

train(training_data, sim_start_time, sim_stop_time, epochs=100, key=None, params=None, opt_state=None)

Run the optimization loop over the training data

Source code in jaxonomy/optimization/training.py

def train(
    self,
    training_data,
    sim_start_time,
    sim_stop_time,
    epochs=100,
    key=None,
    params=None,
    opt_state=None,
):
    """Run the optimization loop over the training data"""

    if (
        self.simulator.max_major_steps is None
        or self.simulator.max_major_steps <= 0
    ):
        self.simulator.max_major_steps = estimate_max_major_steps(
            self.simulator.system,
            (sim_start_time, sim_stop_time),
            self.simulator.max_major_step_length,
        )

    if key is None:
        key = jax.random.PRNGKey(np.random.randint(0, 2**32, dtype=np.int64))

    # Create a function to evaluate the forward pass through the simulation
    forward = self.make_forward(sim_start_time, sim_stop_time)

    # Pull out the optimizable parameters from the context
    if params is None:
        params = self.optimizable_parameters(self.context)

    # Initialize the optimizer and create the loss function
    loss, p, unflatten = self.make_loss_fn(forward, params)

    if opt_state is None:
        opt_state = self.optimizer.init(p)

    self.opt_state = opt_state

    @jax.jit
    def opt_step(p, opt_state, key, batch_data):
        key, subkey = jax.random.split(key)
        if batch_data:
            loss_value, grads = jax.value_and_grad(loss)(p, subkey, *batch_data)
        else:
            loss_value, grads = jax.value_and_grad(loss)(p, subkey)

        grads = jnp.clip(grads, *self.clip_range)

        updates, opt_state = self.optimizer.update(grads, opt_state, p)
        p = optax.apply_updates(p, updates)
        return p, opt_state, key, loss_value

    def _scan_fun(carry, batch_data):
        p, opt_state, key, loss_value = opt_step(*carry, batch_data)
        return (p, opt_state, key), loss_value

    # Run the optimization loop
    for epoch in range(epochs):
        (p, self.opt_state, key), batch_loss = jax.lax.scan(
            _scan_fun, (p, self.opt_state, key), training_data
        )

        if epoch % self.print_every == 0:
            logger.info("Epoch %s, loss: %s", epoch, jnp.mean(batch_loss))

    # Return the optimized parameters
    return unflatten(p)

Transform

Bases: ABC

Base class for transformations.

Source code in jaxonomy/optimization/framework/base/transformations.py

class Transform(ABC):
    """Base class for transformations."""

    @abstractmethod
    def transform(self, params: dict) -> dict:
        """
        Take original parameters dict {key:value} and output a dict with identical keys
        but transformed `values`.
        """
        pass

    @abstractmethod
    def inverse_transform(self, params: dict) -> dict:
        """
        Take transformed parameters dict {key:value} and output a dict with identical
        keys but inverse-transformed `values`.
        """
        pass

inverse_transform(params) abstractmethod

Take transformed parameters dict {key:value} and output a dict with identical keys but inverse-transformed values.

Source code in jaxonomy/optimization/framework/base/transformations.py

@abstractmethod
def inverse_transform(self, params: dict) -> dict:
    """
    Take transformed parameters dict {key:value} and output a dict with identical
    keys but inverse-transformed `values`.
    """
    pass

transform(params) abstractmethod

Take original parameters dict {key:value} and output a dict with identical keys but transformed values.

Source code in jaxonomy/optimization/framework/base/transformations.py

@abstractmethod
def transform(self, params: dict) -> dict:
    """
    Take original parameters dict {key:value} and output a dict with identical keys
    but transformed `values`.
    """
    pass

TuningResult dataclass

Result of a tune_parameters call.

Attributes:

Name	Type	Description
params	Dict[str, Array]	Optimal parameter values as a dict {name: jax.Array}.
objective	float	Final objective value (scalar).
history	list	List of (iteration, objective) tuples if tracking enabled.
success	bool	True if the optimizer reported successful convergence.
message	str	Human-readable status from the underlying optimizer.
raw	Optional[OptimizationResult]	Underlying OptimizationResult from the optimizer framework (for inspection of optimizer-specific fields).

Source code in jaxonomy/optimization/parameter_tuning.py

@dataclass
class TuningResult:
    """Result of a `tune_parameters` call.

    Attributes:
        params: Optimal parameter values as a dict {name: jax.Array}.
        objective: Final objective value (scalar).
        history: List of (iteration, objective) tuples if tracking enabled.
        success: True if the optimizer reported successful convergence.
        message: Human-readable status from the underlying optimizer.
        raw: Underlying `OptimizationResult` from the optimizer framework
            (for inspection of optimizer-specific fields).
    """

    params: Dict[str, jax.Array]
    objective: float
    history: list = field(default_factory=list)
    success: bool = True
    message: str = ""
    raw: Optional[OptimizationResult] = None

compute_confidence_intervals(optimizable, opt_params, confidence_level=0.95, n_data=None, hessian=None, eps_fd=0.0001, regularize=True)

Compute Wald-type confidence intervals for optimised parameters.

Uses the Laplace approximation: the parameter posterior is approximated as a Gaussian centred at the optimum θ* with covariance H⁻¹, where H = ∇²L(θ*) is the Hessian of the loss.

Parameters

optimizable : Optimizable The jaxonomy optimizable whose objective_flat is used. opt_params : OptimizationResult | dict | array-like Optimised parameters at which to evaluate the Hessian. Can be:

* An :class:`~jaxonomy.optimization.OptimizationResult` returned by
  any jaxonomy optimizer — the ``params`` dict is extracted and
  flattened automatically.
* A plain ``dict`` mapping parameter names to values.
* A flat 1-D array matching ``optimizable.params_0_flat``.

confidence_level : float Nominal confidence level (default 0.95 for 95 % CIs). n_data : int or None Number of observations. When provided, the covariance is scaled by the residual variance estimate

    ``σ² = 2 · L(θ*) / max(n_data − n_params, 1)``

This is appropriate for **sum-of-squares objectives**
``L = ½ Σ rᵢ²``.  For maximum-likelihood objectives leave ``None``.

hessian : ndarray or None Pre-computed Hessian matrix (e.g. from :func:compute_sensitivity). When None (default) the Hessian is computed automatically using JAX AD (with a finite-difference fallback for ODE-based objectives). eps_fd : float Step size used for the finite-difference Hessian fallback (default 1e-4). Ignored when hessian is provided or when AD succeeds. regularize : bool When True (default), negative eigenvalues of the Hessian are clipped to a small positive value before inversion. This makes the covariance well-defined even when the supplied point is not a true local minimum. A warning is recorded in result.message.

Returns

ConfidenceIntervalResult Dataclass containing the covariance matrix, standard errors, and per-parameter confidence intervals in the original (physical) parameter space.

Notes

Parameter transformations: if the Optimizable uses a transformation (e.g. :class:LogTransform), the Hessian is computed in the transformed space and the resulting CI bounds are mapped back to the original space via transformation.inverse_transform.

Validity: the Laplace approximation requires the objective to be smooth and the optimum to be a true interior local minimum (positive- definite Hessian). If is_positive_definite is False in the result, the CIs are computed but should be treated with caution.

Profile likelihood: the Laplace approximation is a first-order Gaussian approximation. For strongly nonlinear models or highly non-Gaussian posteriors, profile likelihood confidence intervals are more accurate but require repeated re-optimisation.

Examples

from jaxonomy.optimization import Scipy, compute_confidence_intervals opt = Scipy(my_opt, method="L-BFGS-B", use_autodiff_grad=True) result = opt.optimize() ci = compute_confidence_intervals(my_opt, result, confidence_level=0.95) print(ci.summary()) lo, hi = ci.interval("c")

Source code in jaxonomy/optimization/confidence.py

def compute_confidence_intervals(
    optimizable: Optimizable,
    opt_params: "OptimizationResult | dict | np.ndarray",
    confidence_level: float = 0.95,
    n_data: int | None = None,
    hessian: "np.ndarray | None" = None,
    eps_fd: float = 1e-4,
    regularize: bool = True,
) -> ConfidenceIntervalResult:
    """
    Compute Wald-type confidence intervals for optimised parameters.

    Uses the **Laplace approximation**: the parameter posterior is approximated
    as a Gaussian centred at the optimum θ* with covariance ``H⁻¹``, where
    ``H = ∇²L(θ*)`` is the Hessian of the loss.

    Parameters
    ----------
    optimizable : Optimizable
        The jaxonomy optimizable whose ``objective_flat`` is used.
    opt_params : OptimizationResult | dict | array-like
        Optimised parameters at which to evaluate the Hessian.  Can be:

        * An :class:`~jaxonomy.optimization.OptimizationResult` returned by
          any jaxonomy optimizer — the ``params`` dict is extracted and
          flattened automatically.
        * A plain ``dict`` mapping parameter names to values.
        * A flat 1-D array matching ``optimizable.params_0_flat``.
    confidence_level : float
        Nominal confidence level (default ``0.95`` for 95 % CIs).
    n_data : int or None
        Number of observations.  When provided, the covariance is scaled by
        the residual variance estimate

            ``σ² = 2 · L(θ*) / max(n_data − n_params, 1)``

        This is appropriate for **sum-of-squares objectives**
        ``L = ½ Σ rᵢ²``.  For maximum-likelihood objectives leave ``None``.
    hessian : ndarray or None
        Pre-computed Hessian matrix (e.g. from :func:`compute_sensitivity`).
        When ``None`` (default) the Hessian is computed automatically using
        JAX AD (with a finite-difference fallback for ODE-based objectives).
    eps_fd : float
        Step size used for the finite-difference Hessian fallback (default
        ``1e-4``).  Ignored when ``hessian`` is provided or when AD succeeds.
    regularize : bool
        When ``True`` (default), negative eigenvalues of the Hessian are
        clipped to a small positive value before inversion.  This makes the
        covariance well-defined even when the supplied point is not a true
        local minimum.  A warning is recorded in ``result.message``.

    Returns
    -------
    ConfidenceIntervalResult
        Dataclass containing the covariance matrix, standard errors, and
        per-parameter confidence intervals in the **original** (physical)
        parameter space.

    Notes
    -----
    **Parameter transformations**: if the ``Optimizable`` uses a
    ``transformation`` (e.g. :class:`LogTransform`), the Hessian is computed
    in the *transformed* space and the resulting CI bounds are mapped back to
    the original space via ``transformation.inverse_transform``.

    **Validity**: the Laplace approximation requires the objective to be
    smooth and the optimum to be a true interior local minimum (positive-
    definite Hessian).  If ``is_positive_definite`` is ``False`` in the
    result, the CIs are computed but should be treated with caution.

    **Profile likelihood**: the Laplace approximation is a first-order
    Gaussian approximation.  For strongly nonlinear models or highly
    non-Gaussian posteriors, profile likelihood confidence intervals are
    more accurate but require repeated re-optimisation.

    Examples
    --------
    >>> from jaxonomy.optimization import Scipy, compute_confidence_intervals
    >>> opt = Scipy(my_opt, method="L-BFGS-B", use_autodiff_grad=True)
    >>> result = opt.optimize()
    >>> ci = compute_confidence_intervals(my_opt, result, confidence_level=0.95)
    >>> print(ci.summary())
    >>> lo, hi = ci.interval("c")
    """
    # ------------------------------------------------------------------
    # 1. Resolve flat parameter vector at the optimum
    # ------------------------------------------------------------------
    opt_flat = _to_flat_array(opt_params, optimizable.unflatten_params)
    n_params = len(opt_flat)

    # ------------------------------------------------------------------
    # 2. Objective value at the optimum
    # ------------------------------------------------------------------
    obj_fn = jax.jit(optimizable.objective_flat)
    try:
        obj_val = float(obj_fn(jnp.array(opt_flat)))
    except Exception:
        obj_val = float("nan")

    # ------------------------------------------------------------------
    # 3. Hessian at the optimum
    # ------------------------------------------------------------------
    messages: list[str] = []

    if hessian is not None:
        H_raw = np.asarray(hessian, dtype=float)
        hess_method = "provided"
    else:
        H_raw, hess_method = _compute_hessian(obj_fn, opt_flat, eps=eps_fd)

    if np.any(np.isnan(H_raw)):
        messages.append(
            f"Hessian computation {hess_method!r} produced NaN — "
            "covariance and CIs are unreliable."
        )

    # ------------------------------------------------------------------
    # 4. Eigendecompose and check positive-definiteness
    # ------------------------------------------------------------------
    H_sym = 0.5 * (H_raw + H_raw.T)  # enforce symmetry

    try:
        eigvals_raw = np.linalg.eigvalsh(H_sym)
        is_pd = bool(np.all(eigvals_raw > 0))
        pos = np.abs(eigvals_raw[np.abs(eigvals_raw) > 1e-16])
        hess_cond = float(pos.max() / pos.min()) if len(pos) >= 2 else (
            1.0 if len(pos) == 1 else float("inf")
        )
    except np.linalg.LinAlgError:
        eigvals_raw = np.full(n_params, np.nan)
        is_pd = False
        hess_cond = float("inf")

    if not is_pd and not np.any(np.isnan(H_raw)):
        messages.append(
            "Hessian is not positive definite — supplied point may not be "
            "a true local minimum. "
            + ("CIs computed with regularised Hessian." if regularize else
               "CIs may be unreliable.")
        )

    # ------------------------------------------------------------------
    # 5. Regularise (clip negative eigenvalues) if requested
    # ------------------------------------------------------------------
    if regularize:
        H_for_inv, _ = _nearest_positive_definite(H_sym)
    else:
        H_for_inv = H_sym

    # ------------------------------------------------------------------
    # 6. Invert Hessian → raw covariance
    # ------------------------------------------------------------------
    H_inv, inv_msg = _safe_invert(H_for_inv)
    if inv_msg:
        messages.append(inv_msg)

    # ------------------------------------------------------------------
    # 7. Residual-variance scaling (least-squares mode)
    # ------------------------------------------------------------------
    resid_var: float | None = None
    if n_data is not None:
        dof = max(n_data - n_params, 1)
        if not np.isnan(obj_val):
            resid_var = 2.0 * obj_val / dof
        else:
            resid_var = float("nan")
            messages.append(
                "Could not compute residual variance: objective value is NaN."
            )
        if resid_var is not None and not np.isnan(resid_var):
            H_inv = H_inv * resid_var

    cov = H_inv

    # ------------------------------------------------------------------
    # 8. Correlation matrix and standard errors
    # ------------------------------------------------------------------
    diag = np.diag(cov)
    # Clip negative diagonal entries (can arise from regularisation of badly
    # conditioned problems) to avoid imaginary standard errors.
    diag_safe = np.maximum(diag, 0.0)
    if np.any(diag < 0):
        messages.append(
            "Covariance matrix has negative diagonal entries; "
            "standard errors for affected parameters are set to NaN."
        )
    std_errors = np.where(diag_safe > 0, np.sqrt(diag_safe), np.nan)

    outer_std = np.outer(std_errors, std_errors)
    with np.errstate(invalid="ignore", divide="ignore"):
        corr = np.where(outer_std > 0, cov / outer_std, np.eye(n_params))

    # ------------------------------------------------------------------
    # 9. z-score and raw (transformed-space) CIs
    # ------------------------------------------------------------------
    z = _z_quantile(confidence_level)
    lo_flat = opt_flat - z * std_errors
    hi_flat = opt_flat + z * std_errors

    # ------------------------------------------------------------------
    # 10. Back-transform CIs to original parameter space
    # ------------------------------------------------------------------
    lo_dict_t = optimizable.unflatten_params(jnp.array(lo_flat))
    hi_dict_t = optimizable.unflatten_params(jnp.array(hi_flat))
    opt_dict_t = optimizable.unflatten_params(jnp.array(opt_flat))

    if getattr(optimizable, "transformation", None) is not None:
        tf = optimizable.transformation
        lo_dict_orig = tf.inverse_transform(lo_dict_t)
        hi_dict_orig = tf.inverse_transform(hi_dict_t)
        opt_dict_orig = tf.inverse_transform(opt_dict_t)
    else:
        lo_dict_orig = lo_dict_t
        hi_dict_orig = hi_dict_t
        opt_dict_orig = opt_dict_t

    # Convert to plain Python dicts (jax arrays → numpy scalars/arrays)
    opt_dict_orig = {k: np.asarray(v) for k, v in opt_dict_orig.items()}
    lo_dict_orig  = {k: np.asarray(v) for k, v in lo_dict_orig.items()}
    hi_dict_orig  = {k: np.asarray(v) for k, v in hi_dict_orig.items()}

    # ------------------------------------------------------------------
    # 11. Build the per-parameter CI dict using expanded names
    # ------------------------------------------------------------------
    flat_names = _expand_param_names(opt_dict_orig)

    def _dict_to_flat(d: dict) -> np.ndarray:
        """Flatten a dict of arrays in the same order as flat_names."""
        vals = []
        for key, val in d.items():
            val_np = np.asarray(val)
            if val_np.ndim == 0 or val_np.size == 1:
                vals.append(float(val_np.ravel()[0]))
            else:
                vals.extend(float(x) for x in val_np.ravel())
        return np.array(vals)

    lo_orig_flat = _dict_to_flat(lo_dict_orig)
    hi_orig_flat = _dict_to_flat(hi_dict_orig)
    opt_orig_flat = _dict_to_flat(opt_dict_orig)

    # When a transform is monotone *decreasing* (unlikely but possible),
    # lo and hi may swap.  Always store as (min, max).
    ci_dict: dict[str, tuple[float, float]] = {}
    for name, lo_val, hi_val in zip(flat_names, lo_orig_flat, hi_orig_flat):
        ci_dict[name] = (min(float(lo_val), float(hi_val)),
                         max(float(lo_val), float(hi_val)))

    # ------------------------------------------------------------------
    # 12. Assemble result
    # ------------------------------------------------------------------
    return ConfidenceIntervalResult(
        param_names=flat_names,
        opt_params=opt_dict_orig,
        covariance=cov,
        correlation=corr,
        standard_errors=std_errors,
        confidence_intervals=ci_dict,
        confidence_level=confidence_level,
        z_score=z,
        hessian=H_raw,
        hessian_eigenvalues=eigvals_raw,
        hessian_condition_number=hess_cond,
        is_positive_definite=is_pd,
        residual_variance=resid_var,
        n_data=n_data,
        objective_value=obj_val,
        hessian_method=hess_method,
        message="  |  ".join(messages),
    )

compute_sensitivity(optimizable, params_0_flat=None, sensitivity_threshold=0.001, compute_hessian=True)

Compute gradient-based parameter sensitivity at a given operating point.

Uses JAX automatic differentiation — no finite differences, no extra simulations beyond two JIT-compiled evaluations (gradient + optional Hessian).

Parameters

optimizable : Optimizable The jaxonomy optimizable whose objective_flat is differentiated. params_0_flat : array-like or None Flat parameter vector to evaluate at. Defaults to optimizable.params_0_flat. sensitivity_threshold : float Relative threshold (0–1) for flagging parameters as low-sensitivity. A parameter is flagged when its normalised sensitivity is less than sensitivity_threshold × max(all normalised sensitivities). Default 1e-3. compute_hessian : bool Whether to compute the full Hessian / FIM. Can be expensive for many parameters (O(n²) simulations). Default True.

Returns

SensitivityResult

Source code in jaxonomy/optimization/sensitivity.py

def compute_sensitivity(
    optimizable: Optimizable,
    params_0_flat: "jnp.ndarray | None" = None,
    sensitivity_threshold: float = 1e-3,
    compute_hessian: bool = True,
) -> SensitivityResult:
    """
    Compute gradient-based parameter sensitivity at a given operating point.

    Uses JAX automatic differentiation — no finite differences, no extra
    simulations beyond two JIT-compiled evaluations (gradient + optional
    Hessian).

    Parameters
    ----------
    optimizable : Optimizable
        The jaxonomy optimizable whose ``objective_flat`` is differentiated.
    params_0_flat : array-like or None
        Flat parameter vector to evaluate at.  Defaults to
        ``optimizable.params_0_flat``.
    sensitivity_threshold : float
        Relative threshold (0–1) for flagging parameters as low-sensitivity.
        A parameter is flagged when its normalised sensitivity is less than
        ``sensitivity_threshold × max(all normalised sensitivities)``.
        Default ``1e-3``.
    compute_hessian : bool
        Whether to compute the full Hessian / FIM.  Can be expensive for
        many parameters (O(n²) simulations).  Default ``True``.

    Returns
    -------
    SensitivityResult
    """
    if params_0_flat is None:
        params_0_flat = optimizable.params_0_flat
    params_0_flat = jnp.array(params_0_flat, dtype=float)

    obj_fn = jax.jit(optimizable.objective_flat)
    grad_fn = jax.jit(jax.grad(obj_fn))

    # --- objective value and gradient ---
    obj_val = float(obj_fn(params_0_flat))
    grads = np.array(grad_fn(params_0_flat), dtype=float)

    # --- normalised sensitivity: |p_i * ∂L/∂p_i| ---
    p0_np = np.array(params_0_flat, dtype=float)
    # Use max(|p|, 1) so that near-zero parameters still get a meaningful scale
    scale = np.where(np.abs(p0_np) > 1e-10, np.abs(p0_np), 1.0)
    norm_sensitivity = np.abs(grads * scale)

    # --- Hessian (FIM approximation) ---
    n = len(params_0_flat)
    if compute_hessian:
        try:
            hess_fn = jax.jit(jax.hessian(obj_fn))
            hessian = np.array(hess_fn(params_0_flat), dtype=float)
        except Exception:
            # Second-order AD may fail when the objective uses ODE solvers with
            # custom_vjp (which only supports first-order differentiation).
            # Fall back to a finite-difference approximation of the Hessian.
            try:
                eps = 1e-4
                p0_np_h = np.array(params_0_flat, dtype=float)
                hessian = np.zeros((n, n), dtype=float)
                for i in range(n):
                    ei = np.zeros(n, dtype=float)
                    ei[i] = eps
                    for j in range(n):
                        ej = np.zeros(n, dtype=float)
                        ej[j] = eps
                        f_pp = float(obj_fn(jnp.array(p0_np_h + ei + ej)))
                        f_pm = float(obj_fn(jnp.array(p0_np_h + ei - ej)))
                        f_mp = float(obj_fn(jnp.array(p0_np_h - ei + ej)))
                        f_mm = float(obj_fn(jnp.array(p0_np_h - ei - ej)))
                        hessian[i, j] = (f_pp - f_pm - f_mp + f_mm) / (4 * eps * eps)
            except Exception:
                hessian = np.full((n, n), np.nan)
    else:
        hessian = np.full((n, n), np.nan)

    hess_diag = np.diag(hessian)

    # --- eigenvalues and condition number ---
    if compute_hessian and not np.any(np.isnan(hessian)):
        try:
            eigvals = np.linalg.eigvalsh(hessian)
            positive = np.abs(eigvals[np.abs(eigvals) > 1e-16])
            if len(positive) >= 2:
                condition_number = float(positive.max() / positive.min())
            elif len(positive) == 1:
                condition_number = 1.0
            else:
                condition_number = float("inf")
        except np.linalg.LinAlgError:
            eigvals = np.full(n, np.nan)
            condition_number = float("inf")
    else:
        eigvals = np.full(n, np.nan)
        condition_number = float("inf")

    # --- unidentifiable parameter detection ---
    max_sensitivity = float(np.max(norm_sensitivity)) if n > 0 else 1.0
    threshold_abs = sensitivity_threshold * max(max_sensitivity, 1e-30)
    param_dict = optimizable.unflatten_params(params_0_flat)
    # Expand vector-valued parameters to one name per flat element so the
    # names line up 1:1 with the per-element gradients / sensitivities
    # (a bare list(keys()) is per-parameter and would silently truncate the
    # zip below for any array param). Mirrors confidence._expand_param_names.
    param_names = _expand_param_names(param_dict)
    unidentifiable = [
        name
        for name, s in zip(param_names, norm_sensitivity)
        if s < threshold_abs
    ]

    return SensitivityResult(
        param_names=param_names,
        params_0=param_dict,
        objective_value=obj_val,
        gradients=grads,
        normalized_sensitivity=norm_sensitivity,
        hessian=hessian,
        hessian_diagonal=hess_diag,
        eigenvalues=eigvals,
        condition_number=condition_number,
        unidentifiable_params=unidentifiable,
        sensitivity_threshold=sensitivity_threshold,
    )

implicit_solver(solver, residual, linear_solve=None)

Make an iterative solver reverse-mode differentiable via the IFT.

Parameters:

Name	Type	Description	Default
solver	Callable	solver(theta) -> x_star — any (jit-compatible) function returning a solution as a flat array (shape (n,) or scalar). Internal control flow is unrestricted; it is never differentiated. theta may be any pytree.	required
residual	Callable	residual(x, theta) -> r with r the same shape as x and residual(solver(theta), theta) ≈ 0. Must be JAX-differentiable — this is the equation the solution satisfies, used to construct both sides of the IFT.	required
linear_solve	Optional[Callable]	optional linear_solve(A, b) -> w used for the adjoint system (∂g/∂x)ᵀ w = b. Defaults to the dense :func:jnp.linalg.solve — right for the small systems that appear inside dynamics callbacks (constraint dimensions of tens). Supply a matrix-free solver (e.g. CG) for large n.	None

Returns:

Type	Description
	A function wrapped(theta) -> x_star that is byte-equivalent
	to solver in the forward pass and reverse-differentiable.

Example — an implicit velocity law solved by Newton iteration::

def solve_v(theta):                       # while_loop inside
    def newton(v):
        g = v + jnp.tanh(theta * v) - 1.0
        dg = 1.0 + theta / jnp.cosh(theta * v) ** 2
        return v - g / dg
    def cond(carry):
        v, i = carry
        return (jnp.abs(v + jnp.tanh(theta * v) - 1.0) > 1e-12) & (i < 50)
    def body(carry):
        v, i = carry
        return newton(v), i + 1
    v, _ = jax.lax.while_loop(cond, body, (jnp.asarray(0.5), 0))
    return v

def residual(v, theta):
    return v + jnp.tanh(theta * v) - 1.0

solve_v_diff = implicit_solver(solve_v, residual)
jax.grad(solve_v_diff)(0.3)               # works; matches FD

Source code in jaxonomy/optimization/implicit.py

def implicit_solver(
    solver: Callable,
    residual: Callable,
    linear_solve: Optional[Callable] = None,
):
    """Make an iterative solver reverse-mode differentiable via the IFT.

    Args:
        solver: ``solver(theta) -> x_star`` — any (jit-compatible)
            function returning a solution as a **flat array** (shape
            ``(n,)`` or scalar). Internal control flow is unrestricted;
            it is never differentiated. ``theta`` may be any pytree.
        residual: ``residual(x, theta) -> r`` with ``r`` the same shape
            as ``x`` and ``residual(solver(theta), theta) ≈ 0``. Must be
            JAX-differentiable — this is the equation the solution
            satisfies, used to construct both sides of the IFT.
        linear_solve: optional ``linear_solve(A, b) -> w`` used for the
            adjoint system ``(∂g/∂x)ᵀ w = b``. Defaults to the dense
            :func:`jnp.linalg.solve` — right for the small systems that
            appear inside dynamics callbacks (constraint dimensions of
            tens). Supply a matrix-free solver (e.g. CG) for large ``n``.

    Returns:
        A function ``wrapped(theta) -> x_star`` that is byte-equivalent
        to ``solver`` in the forward pass and reverse-differentiable.

    Example — an implicit velocity law solved by Newton iteration::

        def solve_v(theta):                       # while_loop inside
            def newton(v):
                g = v + jnp.tanh(theta * v) - 1.0
                dg = 1.0 + theta / jnp.cosh(theta * v) ** 2
                return v - g / dg
            def cond(carry):
                v, i = carry
                return (jnp.abs(v + jnp.tanh(theta * v) - 1.0) > 1e-12) & (i < 50)
            def body(carry):
                v, i = carry
                return newton(v), i + 1
            v, _ = jax.lax.while_loop(cond, body, (jnp.asarray(0.5), 0))
            return v

        def residual(v, theta):
            return v + jnp.tanh(theta * v) - 1.0

        solve_v_diff = implicit_solver(solve_v, residual)
        jax.grad(solve_v_diff)(0.3)               # works; matches FD
    """
    if linear_solve is None:
        def linear_solve(A, b):  # noqa: ANN001 - simple default
            return jnp.linalg.solve(A, b)

    @jax.custom_vjp
    def wrapped(theta):
        return solver(theta)

    def fwd(theta):
        x_star = solver(theta)
        return x_star, (x_star, theta)

    def bwd(saved, x_bar):
        x_star, theta = saved
        x_arr = jnp.asarray(x_star)
        scalar_x = x_arr.ndim == 0
        x_flat = jnp.atleast_1d(x_arr)
        xbar_flat = jnp.atleast_1d(jnp.asarray(x_bar)).astype(x_flat.dtype)

        def g_of_x(x):
            r = residual(x[0] if scalar_x else x, theta)
            return jnp.atleast_1d(jnp.asarray(r))

        # Adjoint linear system: (∂g/∂x)ᵀ w = -x̄. Dense Jacobian by
        # default — the systems this wraps are small (constraint dims of
        # tens); pass linear_solve= for matrix-free treatment.
        J_x = jax.jacobian(g_of_x)(x_flat)
        w = linear_solve(jnp.transpose(J_x), -xbar_flat)

        # θ̄ = (∂g/∂θ)ᵀ w via a VJP of the residual in its second slot.
        def g_of_theta(th):
            r = residual(x_star, th)
            return jnp.atleast_1d(jnp.asarray(r))

        _, vjp_theta = jax.vjp(g_of_theta, theta)
        (theta_bar,) = vjp_theta(w)
        return (theta_bar,)

    wrapped.defvjp(fwd, bwd)
    return wrapped

ise_objective(builder, signal_port, reference_port=None, weight=1.0, initial_cost=0.0, name='ise')

Add blocks to compute the Integral of Squared Error (ISE).

.. math::

J = \int_0^T w \, \| \text{signal}(t) - \text{reference}(t) \|^2 \, dt

When reference_port is None the reference is implicitly zero, so the objective is :math:\int_0^T w \, \|\text{signal}(t)\|^2 \, dt.

The function adds the following blocks to builder:

• (optional) :class:~jaxonomy.library.Adder computing signal − reference
• :class:~jaxonomy.library.Power (2.0)
• :class:~jaxonomy.library.SumOfElements (handles both scalar and vector signals transparently)
• (optional) :class:~jaxonomy.library.Gain if weight ≠ 1
• :class:~jaxonomy.library.Integrator accumulating the cost

Parameters

builder: The :class:~jaxonomy.DiagramBuilder to add blocks to. signal_port: Output port of the signal to penalise. reference_port: Output port of the reference signal. None → reference is 0. weight: Scalar multiplier applied to the squared norm before integration. For per-component or matrix weighting use :func:lqr_objective. initial_cost: Initial value of the accumulating integrator (default 0.0). name: Name prefix for the added blocks.

Returns

OutputPort Scalar port whose value at the end of simulation equals J.

Examples

Minimise oscillation energy of a spring-mass system::

obj = ise_objective(b, x.output_ports[0])  # ∫ x² dt
# later:  return obj.eval(ctx)

Multi-signal ISE with a shared reference of zero::

cost_x = ise_objective(b, x.output_ports[0], name="ise_x")
cost_v = ise_objective(b, v.output_ports[0], name="ise_v")
total  = weighted_sum(b, [cost_x, cost_v], weights=[1.0, 0.5])

Source code in jaxonomy/optimization/objectives.py

def ise_objective(
    builder,
    signal_port,
    reference_port=None,
    weight: float = 1.0,
    initial_cost: float = 0.0,
    name: str = "ise",
):
    r"""Add blocks to compute the **Integral of Squared Error** (ISE).

    .. math::

        J = \int_0^T w \, \| \text{signal}(t) - \text{reference}(t) \|^2 \, dt

    When ``reference_port`` is ``None`` the reference is implicitly zero, so
    the objective is :math:`\int_0^T w \, \|\text{signal}(t)\|^2 \, dt`.

    The function adds the following blocks to *builder*:

    * (optional) :class:`~jaxonomy.library.Adder` computing
      ``signal − reference``
    * :class:`~jaxonomy.library.Power` ``(2.0)``
    * :class:`~jaxonomy.library.SumOfElements` (handles both scalar and
      vector signals transparently)
    * (optional) :class:`~jaxonomy.library.Gain` if ``weight ≠ 1``
    * :class:`~jaxonomy.library.Integrator` accumulating the cost

    Parameters
    ----------
    builder:
        The :class:`~jaxonomy.DiagramBuilder` to add blocks to.
    signal_port:
        Output port of the signal to penalise.
    reference_port:
        Output port of the reference signal.  ``None`` → reference is 0.
    weight:
        Scalar multiplier applied to the squared norm before integration.
        For per-component or matrix weighting use :func:`lqr_objective`.
    initial_cost:
        Initial value of the accumulating integrator (default ``0.0``).
    name:
        Name prefix for the added blocks.

    Returns
    -------
    OutputPort
        Scalar port whose value at the end of simulation equals *J*.

    Examples
    --------
    Minimise oscillation energy of a spring-mass system::

        obj = ise_objective(b, x.output_ports[0])  # ∫ x² dt
        # later:  return obj.eval(ctx)

    Multi-signal ISE with a shared reference of zero::

        cost_x = ise_objective(b, x.output_ports[0], name="ise_x")
        cost_v = ise_objective(b, v.output_ports[0], name="ise_v")
        total  = weighted_sum(b, [cost_x, cost_v], weights=[1.0, 0.5])
    """
    from jaxonomy.library import Adder

    # ── error ──────────────────────────────────────────────────────────────
    if reference_port is not None:
        err = builder.add(Adder(2, operators="+-", name=f"{name}_err"))
        builder.connect(signal_port, err.input_ports[0])
        builder.connect(reference_port, err.input_ports[1])
        err_port = err.output_ports[0]
    else:
        err_port = signal_port

    # ── ‖e‖² (scalar) ──────────────────────────────────────────────────────
    sq_port = _sq_sum_port(builder, err_port, name)

    # ── optional scalar weight ──────────────────────────────────────────────
    weighted_port = _scale_port(builder, sq_port, weight, name)

    # ── integrate ───────────────────────────────────────────────────────────
    return _integrate_port(builder, weighted_port, initial_cost, name)

lqr_objective(builder, state_port, Q, control_port=None, R=None, initial_cost=0.0, name='lqr')

Add blocks to compute an LQR-style quadratic cost.

.. math::

J = \int_0^T \bigl( x(t)^\top Q\, x(t) \;+\; u(t)^\top R\, u(t) \bigr)\, dt

When control_port or R is None only the state cost :math:\int x^\top Q x\, dt is computed.

The function adds a single-input :class:~jaxonomy.library.ReduceBlock for :math:x^\top Q x (and optionally one for :math:u^\top R u), an optional :class:~jaxonomy.library.Adder, and an :class:~jaxonomy.library.Integrator.

Parameters

builder: The :class:~jaxonomy.DiagramBuilder to add blocks to. state_port: Output port of the state vector :math:x. Q: Positive semi-definite state weight matrix, shape (nx, nx). control_port: Output port of the control vector :math:u. None → no control penalty. R: Positive definite control weight matrix, shape (nu, nu). Required when control_port is provided. initial_cost: Initial value of the accumulating integrator (default 0.0). name: Name prefix for the added blocks.

Returns

OutputPort Scalar port whose value at the end of simulation equals J.

Examples

Pendulum regulation::

# ∫ θ²·Q[0,0] + ω²·Q[1,1] dt  (diagonal Q)
Q = jnp.diag(jnp.array([10.0, 1.0]))
R = jnp.array([[0.1]])
cost = lqr_objective(b, x.output_ports[0], Q, u.output_ports[0], R)

Source code in jaxonomy/optimization/objectives.py

def lqr_objective(
    builder,
    state_port,
    Q,
    control_port=None,
    R=None,
    initial_cost: float = 0.0,
    name: str = "lqr",
):
    r"""Add blocks to compute an **LQR-style quadratic cost**.

    .. math::

        J = \int_0^T \bigl( x(t)^\top Q\, x(t) \;+\; u(t)^\top R\, u(t) \bigr)\, dt

    When ``control_port`` or ``R`` is ``None`` only the state cost
    :math:`\int x^\top Q x\, dt` is computed.

    The function adds a single-input :class:`~jaxonomy.library.ReduceBlock`
    for :math:`x^\top Q x` (and optionally one for :math:`u^\top R u`), an
    optional :class:`~jaxonomy.library.Adder`, and an
    :class:`~jaxonomy.library.Integrator`.

    Parameters
    ----------
    builder:
        The :class:`~jaxonomy.DiagramBuilder` to add blocks to.
    state_port:
        Output port of the state vector :math:`x`.
    Q:
        Positive semi-definite state weight matrix, shape ``(nx, nx)``.
    control_port:
        Output port of the control vector :math:`u`.  ``None`` → no control
        penalty.
    R:
        Positive definite control weight matrix, shape ``(nu, nu)``.
        Required when *control_port* is provided.
    initial_cost:
        Initial value of the accumulating integrator (default ``0.0``).
    name:
        Name prefix for the added blocks.

    Returns
    -------
    OutputPort
        Scalar port whose value at the end of simulation equals *J*.

    Examples
    --------
    Pendulum regulation::

        # ∫ θ²·Q[0,0] + ω²·Q[1,1] dt  (diagonal Q)
        Q = jnp.diag(jnp.array([10.0, 1.0]))
        R = jnp.array([[0.1]])
        cost = lqr_objective(b, x.output_ports[0], Q, u.output_ports[0], R)
    """
    from jaxonomy.library import Adder

    Q = np.asarray(Q, dtype=float)

    # ── state cost: x^T Q x ────────────────────────────────────────────────
    x_cost_block = builder.add(_QuadraticIntegrand(Q, name=f"{name}_Qcost"))
    builder.connect(state_port, x_cost_block.input_ports[0])
    cost_port = x_cost_block.output_ports[0]

    # ── optional control cost: u^T R u ─────────────────────────────────────
    if control_port is not None and R is not None:
        R = np.asarray(R, dtype=float)
        u_cost_block = builder.add(_QuadraticIntegrand(R, name=f"{name}_Rcost"))
        builder.connect(control_port, u_cost_block.input_ports[0])

        total = builder.add(Adder(2, operators="++", name=f"{name}_total"))
        builder.connect(cost_port, total.input_ports[0])
        builder.connect(u_cost_block.output_ports[0], total.input_ports[1])
        cost_port = total.output_ports[0]

    # ── integrate ───────────────────────────────────────────────────────────
    return _integrate_port(builder, cost_port, initial_cost, name)

tracking_mse(builder, signal_port, t_data, y_data, weight=1.0, interpolation='linear', initial_cost=0.0, name='tracking_mse')

Add blocks to compute the dataset tracking MSE.

Computes

.. math::

J = \int_0^T w \, \| \text{signal}(t) - y_{\text{ref}}(t) \|^2 \, dt

where :math:y_{\text{ref}}(t) is the reference signal interpolated from the dataset (t_data, y_data) at every simulation time step.

The function wires:

1. :class:~jaxonomy.library.Clock → current simulation time
2. :class:~jaxonomy.library.LookupTable1d → interpolated reference
3. :func:ise_objective → squared error integrator

Parameters

builder: The :class:~jaxonomy.DiagramBuilder to add blocks to. signal_port: Output port of the simulated signal to compare against the data. t_data: 1-D array of reference time points (must be strictly increasing). y_data: Array of reference values. Shape (N,) for scalar signals or (N, ny) for vector signals. Extrapolation clamps to the nearest endpoint value. weight: Scalar multiplier applied before integration. interpolation: Interpolation method passed to :class:~jaxonomy.library.LookupTable1d: "linear" (default), "nearest", or "flat". initial_cost: Initial value of the integrator. name: Name prefix for the added blocks.

Returns

OutputPort Scalar port equal to J at the end of simulation.

Examples

Fit a model to measured step-response data::

import numpy as np
t_meas = np.linspace(0, 5, 50)
y_meas = 1 - np.exp(-t_meas)   # first-order step response

cost = tracking_mse(b, plant.output_ports[0], t_meas, y_meas)

Source code in jaxonomy/optimization/objectives.py

def tracking_mse(
    builder,
    signal_port,
    t_data,
    y_data,
    weight: float = 1.0,
    interpolation: str = "linear",
    initial_cost: float = 0.0,
    name: str = "tracking_mse",
):
    r"""Add blocks to compute the **dataset tracking MSE**.

    Computes

    .. math::

        J = \int_0^T w \, \| \text{signal}(t) - y_{\text{ref}}(t) \|^2 \, dt

    where :math:`y_{\text{ref}}(t)` is the reference signal *interpolated*
    from the dataset ``(t_data, y_data)`` at every simulation time step.

    The function wires:

    1. :class:`~jaxonomy.library.Clock` → current simulation time
    2. :class:`~jaxonomy.library.LookupTable1d` → interpolated reference
    3. :func:`ise_objective` → squared error integrator

    Parameters
    ----------
    builder:
        The :class:`~jaxonomy.DiagramBuilder` to add blocks to.
    signal_port:
        Output port of the simulated signal to compare against the data.
    t_data:
        1-D array of reference time points (must be strictly increasing).
    y_data:
        Array of reference values.  Shape ``(N,)`` for scalar signals or
        ``(N, ny)`` for vector signals.  Extrapolation clamps to the
        nearest endpoint value.
    weight:
        Scalar multiplier applied before integration.
    interpolation:
        Interpolation method passed to :class:`~jaxonomy.library.LookupTable1d`:
        ``"linear"`` (default), ``"nearest"``, or ``"flat"``.
    initial_cost:
        Initial value of the integrator.
    name:
        Name prefix for the added blocks.

    Returns
    -------
    OutputPort
        Scalar port equal to *J* at the end of simulation.

    Examples
    --------
    Fit a model to measured step-response data::

        import numpy as np
        t_meas = np.linspace(0, 5, 50)
        y_meas = 1 - np.exp(-t_meas)   # first-order step response

        cost = tracking_mse(b, plant.output_ports[0], t_meas, y_meas)
    """
    from jaxonomy.library import Clock, LookupTable1d

    t_data = np.asarray(t_data, dtype=float)
    y_data = np.asarray(y_data, dtype=float)

    clock = builder.add(Clock(name=f"{name}_clock"))
    ref = builder.add(
        LookupTable1d(t_data, y_data, interpolation, name=f"{name}_ref")
    )
    builder.connect(clock.output_ports[0], ref.input_ports[0])

    return ise_objective(
        builder,
        signal_port,
        reference_port=ref.output_ports[0],
        weight=weight,
        initial_cost=initial_cost,
        name=name,
    )

tune_parameters(diagram, base_context, sim_t_span, params_0, set_params, objective_fn, bounds=None, optimizer='scipy-lbfgs', n_iter=100, learning_rate=0.05, sim_options=None, verbose=True)

Tune scalar parameters of a jaxonomy diagram to minimize an objective.

The simulator is differentiated through using JAX autodiff; the gradient of objective_fn with respect to each entry of params_0 is computed automatically, and an optimizer minimizes the objective.

Parameters:

Name	Type	Description	Default
diagram		A built jaxonomy diagram.	required
base_context		A Context created from the diagram. The optimizer calls set_params(base_context, params) each iteration to inject the current parameter values, then advances the simulator over sim_t_span, then evaluates objective_fn on the final context.	required
sim_t_span	Tuple[float, float]	(t0, tf) simulation interval.	required
params_0	Dict[str, Any]	Dict of initial parameter values. Each value must be a scalar or a JAX-compatible array.	required
set_params	Callable[[Any, Dict[str, Array]], Any]	Callback (context, params_dict) -> updated_context. Typical implementations write parameter values into LeafSystem parameters via context.with_parameter(...), or modify initial states.	required
objective_fn	Callable[[Any], Array]	Callback (results_context) -> scalar. The optimizer minimizes this. Must be differentiable through JAX.	required
bounds	Optional[Dict[str, Tuple[float, float]]]	Optional dict {param_name: (lb, ub)} for box constraints. Only honoured by box-constrained optimizers (l-bfgs-b, slsqp, trust-constr); ignored otherwise.	None
optimizer	str	Optimizer alias or scipy method name. Defaults to "scipy-lbfgs". See _OPTIMIZER_ALIASES for shortcuts.	'scipy-lbfgs'
n_iter	int	Maximum number of optimizer iterations.	100
learning_rate	float	Learning rate for optax optimizers (ignored by scipy).	0.05
sim_options	Optional[SimulatorOptions]	Optional SimulatorOptions. If None, a default with autodiff enabled is used.	None
verbose	bool	If True, log progress to the jaxonomy logger.	True

Returns:

Type	Description
TuningResult	A TuningResult with optimal parameter values, the final objective,
TuningResult	and a reference to the raw optimizer result.

Notes

• Discrete parameters (e.g., a horizon length N, a state-machine guard threshold) are not differentiable through the simulator and should be left as fixed hyperparameters. If you need to sweep them, wrap tune_parameters in an outer loop or grid search.
• Saturation regions (jnp.clip, jax.lax.cond with hard switches) have zero gradient. Tuning parameters whose value determines a saturation region may be impossible from a starting point already saturated; consider warm-starting away from the saturation boundary.
• Bounds enforcement: with scipy-lbfgs / slsqp / trust-constr, bounds are honoured by the solver. With optax optimizers, bounds are not enforced; clip parameters yourself in set_params if you need them.

See also

jaxonomy.optimization.Optimizable — the lower-level interface this function wraps. Use it directly if you need stochastic variables, constraints, or batched evaluations.

Source code in jaxonomy/optimization/parameter_tuning.py

def tune_parameters(
    diagram,
    base_context,
    sim_t_span: Tuple[float, float],
    params_0: Dict[str, Any],
    set_params: Callable[[Any, Dict[str, jax.Array]], Any],
    objective_fn: Callable[[Any], jax.Array],
    bounds: Optional[Dict[str, Tuple[float, float]]] = None,
    optimizer: str = "scipy-lbfgs",
    n_iter: int = 100,
    learning_rate: float = 0.05,
    sim_options: Optional[SimulatorOptions] = None,
    verbose: bool = True,
) -> TuningResult:
    """Tune scalar parameters of a jaxonomy diagram to minimize an objective.

    The simulator is differentiated through using JAX autodiff; the gradient
    of `objective_fn` with respect to each entry of `params_0` is computed
    automatically, and an optimizer minimizes the objective.

    Args:
        diagram: A built jaxonomy diagram.
        base_context: A `Context` created from the diagram. The optimizer
            calls `set_params(base_context, params)` each iteration to inject
            the current parameter values, then advances the simulator over
            `sim_t_span`, then evaluates `objective_fn` on the final context.
        sim_t_span: `(t0, tf)` simulation interval.
        params_0: Dict of initial parameter values. Each value must be a
            scalar or a JAX-compatible array.
        set_params: Callback `(context, params_dict) -> updated_context`.
            Typical implementations write parameter values into LeafSystem
            parameters via `context.with_parameter(...)`, or modify initial
            states.
        objective_fn: Callback `(results_context) -> scalar`. The optimizer
            minimizes this. Must be differentiable through JAX.
        bounds: Optional dict `{param_name: (lb, ub)}` for box constraints.
            Only honoured by box-constrained optimizers (l-bfgs-b, slsqp,
            trust-constr); ignored otherwise.
        optimizer: Optimizer alias or scipy method name. Defaults to
            "scipy-lbfgs". See `_OPTIMIZER_ALIASES` for shortcuts.
        n_iter: Maximum number of optimizer iterations.
        learning_rate: Learning rate for optax optimizers (ignored by scipy).
        sim_options: Optional `SimulatorOptions`. If None, a default with
            autodiff enabled is used.
        verbose: If True, log progress to the jaxonomy logger.

    Returns:
        A `TuningResult` with optimal parameter values, the final objective,
        and a reference to the raw optimizer result.

    Notes:
        - **Discrete parameters** (e.g., a horizon length `N`, a state-machine
          guard threshold) are not differentiable through the simulator and
          should be left as fixed hyperparameters. If you need to sweep them,
          wrap `tune_parameters` in an outer loop or grid search.
        - **Saturation regions** (`jnp.clip`, `jax.lax.cond` with hard
          switches) have zero gradient. Tuning parameters whose value
          determines a saturation region may be impossible from a starting
          point already saturated; consider warm-starting away from the
          saturation boundary.
        - **Bounds enforcement**: with `scipy-lbfgs` / `slsqp` /
          `trust-constr`, bounds are honoured by the solver. With `optax`
          optimizers, bounds are not enforced; clip parameters yourself in
          `set_params` if you need them.

    See also:
        `jaxonomy.optimization.Optimizable` — the lower-level interface this
        function wraps. Use it directly if you need stochastic variables,
        constraints, or batched evaluations.
    """
    if sim_options is None:
        sim_options = SimulatorOptions(enable_autodiff=True)

    # Default max_major_steps if user didn't provide it
    if sim_options.max_major_steps is None:
        sim_options = dataclasses.replace(
            sim_options,
            max_major_steps=estimate_max_major_steps(diagram, sim_t_span),
        )

    # Normalize bounds: the Optimizable framework expects a dict of (lb, ub)
    # tuples; missing keys mean unbounded. We pass through as-is.
    optimizable = _CallableOptimizable(
        diagram=diagram,
        base_context=base_context,
        sim_t_span=sim_t_span,
        params_0=params_0,
        bounds=bounds,
        set_params_fn=set_params,
        objective_fn=objective_fn,
        sim_options=sim_options,
    )

    # Resolve optimizer choice
    opt_kind, opt_method = _resolve_optimizer(optimizer)

    if verbose:
        logger.info(
            f"tune_parameters: optimizer={opt_kind}/{opt_method}, "
            f"params={list(params_0.keys())}, n_iter={n_iter}"
        )

    raw_result: OptimizationResult
    if opt_kind == "scipy":
        opt = Scipy(
            optimizable=optimizable,
            opt_method=opt_method,
            opt_method_config={"maxiter": int(n_iter), "disp": bool(verbose)},
        )
        raw_result = opt.optimize()
    elif opt_kind == "optax":
        # The Optax runner in jaxonomy is wired for stochastic-variable
        # optimization. Tuning deterministic diagram parameters via optax is
        # not yet exposed at this top-level API. Users who need it should
        # build an OptimizableWithStochasticVars directly.
        raise NotImplementedError(
            "Optax-based tuning for deterministic diagrams is not yet supported "
            "by `tune_parameters`. Use optimizer='scipy-lbfgs' (default) for now, "
            "or drop down to OptimizableWithStochasticVars + Optax directly."
        )
    else:
        raise ValueError(f"Unknown optimizer kind: {opt_kind!r}")

    # Reassemble result
    optimal_params = raw_result.params
    final_objective = (
        float(raw_result.final_loss)
        if raw_result.final_loss is not None
        else float("nan")
    )
    success = bool(getattr(raw_result, "success", True))
    message = str(getattr(raw_result, "message", ""))

    if verbose:
        logger.info(
            f"tune_parameters: done. final objective={final_objective:.6g}, "
            f"success={success}"
        )

    return TuningResult(
        params=optimal_params,
        objective=final_objective,
        history=list(getattr(raw_result, "loss_history", []) or []),
        success=success,
        message=message,
        raw=raw_result,
    )

weighted_sum(builder, objectives, weights=None, name='total_cost')

Combine multiple objective ports into a weighted sum.

.. math::

J_{\text{total}} = \sum_{i} w_i \, J_i

Parameters

builder: The :class:~jaxonomy.DiagramBuilder to add blocks to. objectives: Sequence of scalar output ports, one per term. weights: Scalar weights :math:w_i. None → uniform weight 1.0. Must have the same length as objectives when provided. name: Name of the final :class:~jaxonomy.library.Adder block (and prefix for :class:~jaxonomy.library.Gain blocks when weights differ from 1).

Returns

OutputPort Scalar port equal to :math:J_{\text{total}}.

Raises

ValueError If objectives is empty or weights has a different length.

Examples

Combine two ISE objectives with different priorities::

cost_pos = ise_objective(b, x.output_ports[0], name="ise_x")
cost_vel = ise_objective(b, v.output_ports[0], name="ise_v")
total    = weighted_sum(b, [cost_pos, cost_vel], weights=[10.0, 1.0])

Source code in jaxonomy/optimization/objectives.py

def weighted_sum(
    builder,
    objectives: Sequence,
    weights: Sequence[float] | None = None,
    name: str = "total_cost",
):
    r"""Combine multiple objective ports into a **weighted sum**.

    .. math::

        J_{\text{total}} = \sum_{i} w_i \, J_i

    Parameters
    ----------
    builder:
        The :class:`~jaxonomy.DiagramBuilder` to add blocks to.
    objectives:
        Sequence of scalar output ports, one per term.
    weights:
        Scalar weights :math:`w_i`.  ``None`` → uniform weight ``1.0``.
        Must have the same length as *objectives* when provided.
    name:
        Name of the final :class:`~jaxonomy.library.Adder` block (and prefix
        for :class:`~jaxonomy.library.Gain` blocks when weights differ from 1).

    Returns
    -------
    OutputPort
        Scalar port equal to :math:`J_{\text{total}}`.

    Raises
    ------
    ValueError
        If *objectives* is empty or *weights* has a different length.

    Examples
    --------
    Combine two ISE objectives with different priorities::

        cost_pos = ise_objective(b, x.output_ports[0], name="ise_x")
        cost_vel = ise_objective(b, v.output_ports[0], name="ise_v")
        total    = weighted_sum(b, [cost_pos, cost_vel], weights=[10.0, 1.0])
    """
    from jaxonomy.library import Adder, Gain

    objectives = list(objectives)
    n = len(objectives)

    if n == 0:
        raise ValueError("weighted_sum: 'objectives' must not be empty.")

    if weights is None:
        weights = [1.0] * n
    else:
        weights = [float(w) for w in weights]
        if len(weights) != n:
            raise ValueError(
                f"weighted_sum: 'objectives' has {n} elements but "
                f"'weights' has {len(weights)} elements."
            )

    # Apply individual weights via Gain blocks where w ≠ 1
    scaled = []
    for i, (port, w) in enumerate(zip(objectives, weights)):
        if w != 1.0:
            g = builder.add(Gain(w, name=f"{name}_w{i}"))
            builder.connect(port, g.input_ports[0])
            scaled.append(g.output_ports[0])
        else:
            scaled.append(port)

    # Short-circuit for single objective
    if n == 1:
        return scaled[0]

    # Sum all scaled objectives
    ops = "+" * n
    adder = builder.add(Adder(n, operators=ops, name=name))
    for i, p in enumerate(scaled):
        builder.connect(p, adder.input_ports[i])
    return adder.output_ports[0]