Sensitivity example

In [ ]:

# SPDX-License-Identifier: MIT
"""
Parameter sensitivity / identifiability analysis example
=========================================================

Before fitting parameters, check whether the objective is actually *sensitive*
to each parameter.  Near-zero gradient → the simulation output barely changes
with that parameter → the parameter is unidentifiable from this data.

This example uses ``compute_sensitivity`` to:

1. Report the gradient and normalised sensitivity for each parameter.
2. Compute the Hessian (Fisher Information Matrix approximation).
3. Flag parameters with near-zero sensitivity.

Usage::

    python docs/examples/optimization/sensitivity_example.py
"""

# SPDX-License-Identifier: MIT """ Parameter sensitivity / identifiability analysis example ========================================================= Before fitting parameters, check whether the objective is actually *sensitive* to each parameter. Near-zero gradient → the simulation output barely changes with that parameter → the parameter is unidentifiable from this data. This example uses ``compute_sensitivity`` to: 1. Report the gradient and normalised sensitivity for each parameter. 2. Compute the Hessian (Fisher Information Matrix approximation). 3. Flag parameters with near-zero sensitivity. Usage:: python docs/examples/optimization/sensitivity_example.py """

In [ ]:

import numpy as np

import numpy as np

In [ ]:

from jaxonomy import DiagramBuilder, Parameter, SimulatorOptions
from jaxonomy.library import Adder, Gain, Integrator, Power
from jaxonomy.optimization import Optimizable
from jaxonomy.optimization.sensitivity import compute_sensitivity

from jaxonomy import DiagramBuilder, Parameter, SimulatorOptions from jaxonomy.library import Adder, Gain, Integrator, Power from jaxonomy.optimization import Optimizable from jaxonomy.optimization.sensitivity import compute_sensitivity

── build diagram ─────────────────────────────────────────────────────────────

In [ ]:

params = {
    "c": Parameter(np.array(0.5)),   # damping — identifiable
    "k": Parameter(np.array(1.0)),   # stiffness — identifiable
}
b = DiagramBuilder()
k_x = b.add(Gain(params["k"], name="k_x"))
c_v = b.add(Gain(params["c"], name="c_v"))
add = b.add(Adder(2, operators="--", name="adder"))
inv = b.add(Gain(1.0,         name="inv_m"))
v   = b.add(Integrator(0.1,   name="v"))
x   = b.add(Integrator(1.0,   name="x"))
b.connect(k_x.output_ports[0], add.input_ports[0])
b.connect(c_v.output_ports[0], add.input_ports[1])
b.connect(add.output_ports[0], inv.input_ports[0])
b.connect(inv.output_ports[0], v.input_ports[0])
b.connect(v.output_ports[0],   x.input_ports[0])
b.connect(v.output_ports[0],   c_v.input_ports[0])
b.connect(x.output_ports[0],   k_x.input_ports[0])
sq_v = b.add(Power(2.0, name="sq_v"))
sq_x = b.add(Power(2.0, name="sq_x"))
cv   = b.add(Integrator(0.0, name="cv"))
cx   = b.add(Integrator(0.0, name="cx"))
obj  = b.add(Adder(2, operators="++", name="obj"))
b.connect(v.output_ports[0],   sq_v.input_ports[0])
b.connect(x.output_ports[0],   sq_x.input_ports[0])
b.connect(sq_v.output_ports[0], cv.input_ports[0])
b.connect(sq_x.output_ports[0], cx.input_ports[0])
b.connect(cv.output_ports[0],  obj.input_ports[0])
b.connect(cx.output_ports[0],  obj.input_ports[1])
diagram = b.build(parameters=params)

params = { "c": Parameter(np.array(0.5)), # damping — identifiable "k": Parameter(np.array(1.0)), # stiffness — identifiable } b = DiagramBuilder() k_x = b.add(Gain(params["k"], name="k_x")) c_v = b.add(Gain(params["c"], name="c_v")) add = b.add(Adder(2, operators="--", name="adder")) inv = b.add(Gain(1.0, name="inv_m")) v = b.add(Integrator(0.1, name="v")) x = b.add(Integrator(1.0, name="x")) b.connect(k_x.output_ports[0], add.input_ports[0]) b.connect(c_v.output_ports[0], add.input_ports[1]) b.connect(add.output_ports[0], inv.input_ports[0]) b.connect(inv.output_ports[0], v.input_ports[0]) b.connect(v.output_ports[0], x.input_ports[0]) b.connect(v.output_ports[0], c_v.input_ports[0]) b.connect(x.output_ports[0], k_x.input_ports[0]) sq_v = b.add(Power(2.0, name="sq_v")) sq_x = b.add(Power(2.0, name="sq_x")) cv = b.add(Integrator(0.0, name="cv")) cx = b.add(Integrator(0.0, name="cx")) obj = b.add(Adder(2, operators="++", name="obj")) b.connect(v.output_ports[0], sq_v.input_ports[0]) b.connect(x.output_ports[0], sq_x.input_ports[0]) b.connect(sq_v.output_ports[0], cv.input_ports[0]) b.connect(sq_x.output_ports[0], cx.input_ports[0]) b.connect(cv.output_ports[0], obj.input_ports[0]) b.connect(cx.output_ports[0], obj.input_ports[1]) diagram = b.build(parameters=params)

── define optimizable ────────────────────────────────────────────────────────

In [ ]:

class SpringOptimizable(Optimizable):
    def __init__(self, *a, **kw):
        super().__init__(*a, **kw)
        self._obj = diagram["obj"].output_ports[0]

    def optimizable_params(self, ctx):
        return {"c": ctx.parameters["c"], "k": ctx.parameters["k"]}

    def objective_from_context(self, ctx):
        return self._obj.eval(ctx)

    def prepare_context(self, ctx, p):
        return ctx.with_parameters(p)

class SpringOptimizable(Optimizable): def __init__(self, *a, **kw): super().__init__(*a, **kw) self._obj = diagram["obj"].output_ports[0] def optimizable_params(self, ctx): return {"c": ctx.parameters["c"], "k": ctx.parameters["k"]} def objective_from_context(self, ctx): return self._obj.eval(ctx) def prepare_context(self, ctx, p): return ctx.with_parameters(p)

In [ ]:

model = SpringOptimizable(
    diagram, diagram.create_context(),
    params_0={"c": 0.5, "k": 1.0},
    sim_t_span=(0.0, 2.0),
    sim_options=SimulatorOptions(max_major_steps=1),
)

model = SpringOptimizable( diagram, diagram.create_context(), params_0={"c": 0.5, "k": 1.0}, sim_t_span=(0.0, 2.0), sim_options=SimulatorOptions(max_major_steps=1), )

── sensitivity analysis ──────────────────────────────────────────────────────

In [ ]:

print("Computing sensitivity at initial parameters (c=0.5, k=1.0) ...")
sens = compute_sensitivity(model, sensitivity_threshold=1e-3)
print(sens.summary())

print("Computing sensitivity at initial parameters (c=0.5, k=1.0) ...") sens = compute_sensitivity(model, sensitivity_threshold=1e-3) print(sens.summary())

In [ ]:

print("\n--- Raw values ---")
print(f"Gradients:             {sens.gradients}")
print(f"Normalised sensitivity:{sens.normalized_sensitivity}")
print(f"Hessian diagonal:      {sens.hessian_diagonal}")
print(f"Eigenvalues:           {sens.eigenvalues}")
print(f"Condition number:      {sens.condition_number:.3g}")

print("\n--- Raw values ---") print(f"Gradients: {sens.gradients}") print(f"Normalised sensitivity:{sens.normalized_sensitivity}") print(f"Hessian diagonal: {sens.hessian_diagonal}") print(f"Eigenvalues: {sens.eigenvalues}") print(f"Condition number: {sens.condition_number:.3g}")

In [ ]:

if sens.unidentifiable_params:
    print(f"\n⚠  The following parameters have low sensitivity and may be "
          f"unidentifiable:\n  {sens.unidentifiable_params}")
else:
    print("\n✓  All parameters appear identifiable at this operating point.")

if sens.unidentifiable_params: print(f"\n⚠ The following parameters have low sensitivity and may be " f"unidentifiable:\n {sens.unidentifiable_params}") else: print("\n✓ All parameters appear identifiable at this operating point.")