Openmodelica plant fmu publication offline

In [ ]:

"""Offline publication-quality run for
``openmodelica_plant_fmu_cosim.ipynb``.

Produces ``media/openmodelica_plant_fmu_publication.npz`` from which
the notebook loads its headline numbers when the live FMU build /
import path is unavailable (missing ``pythonfmu`` / ``fmpy`` /
``libpythonfmu-export.dylib``).

Mirror of ``media/fmi_export_publication_offline.py`` (the
controller-out direction), but here we package the **plant** as the
FMU and close the loop with a jaxonomy PI controller around it. The
narrative wedge is the inverse of #15: "your plant is in Modelica;
wrap it as an FMU; we close the loop and give you autodiff over the
controller for free."

Run from the repo root::

    JAXONOMY_DISABLE_PROFILING=1 \\
        python docs/examples/media/openmodelica_plant_fmu_publication_offline.py

Expected wall-time: **~30-60 s on a developer machine** (the FMU
build is ~0.1 s on darwin, the closed-loop simulations are sub-
second each; the script is dominated by JIT compile of the
``ModelicaFMU``-containing closed-loop diagram and the
gradient-vs-finite-difference cross-check).

Pass ``--write-placeholder`` to emit a structurally-plausible NPZ
without invoking pythonfmu / fmpy — useful for environments where
those packages aren't installed.
"""
from __future__ import annotations

"""Offline publication-quality run for ``openmodelica_plant_fmu_cosim.ipynb``. Produces ``media/openmodelica_plant_fmu_publication.npz`` from which the notebook loads its headline numbers when the live FMU build / import path is unavailable (missing ``pythonfmu`` / ``fmpy`` / ``libpythonfmu-export.dylib``). Mirror of ``media/fmi_export_publication_offline.py`` (the controller-out direction), but here we package the **plant** as the FMU and close the loop with a jaxonomy PI controller around it. The narrative wedge is the inverse of #15: "your plant is in Modelica; wrap it as an FMU; we close the loop and give you autodiff over the controller for free." Run from the repo root:: JAXONOMY_DISABLE_PROFILING=1 \\ python docs/examples/media/openmodelica_plant_fmu_publication_offline.py Expected wall-time: **~30-60 s on a developer machine** (the FMU build is ~0.1 s on darwin, the closed-loop simulations are sub- second each; the script is dominated by JIT compile of the ``ModelicaFMU``-containing closed-loop diagram and the gradient-vs-finite-difference cross-check). Pass ``--write-placeholder`` to emit a structurally-plausible NPZ without invoking pythonfmu / fmpy — useful for environments where those packages aren't installed. """ from __future__ import annotations

In [ ]:

import json
import os
import subprocess
import sys
import tempfile
import textwrap
import time
from pathlib import Path

import json import os import subprocess import sys import tempfile import textwrap import time from pathlib import Path

In [ ]:

import numpy as np

import numpy as np

In [ ]:

HERE = Path(__file__).resolve().parent
NPZ_OUT = HERE / "openmodelica_plant_fmu_publication.npz"

HERE = Path(__file__).resolve().parent NPZ_OUT = HERE / "openmodelica_plant_fmu_publication.npz"

In [ ]:

# Shared top-level constants — must match the notebook.
DT = 0.02
T_END = 12.0
SETPOINT = 1.0
# Damped 2nd-order mass-spring plant: m x_ddot + c x_dot + k x = F.
# c=0.5, m=1, k=1 -> wn=1 rad/s, zeta=0.25 (mildly underdamped — visible overshoot).
M_PLANT = 1.0
C_PLANT = 0.5
K_PLANT = 1.0
# Baseline controller gains: conservative tuning yields ~15% overshoot,
# clean settling in 8 s on this plant.
KP_BASE = 1.2
KI_BASE = 0.3
# Gain perturbations for the gradient finite-difference cross-check.
KP_GRID = (0.6, 0.9, 1.2, 1.5, 1.8)
N_STEPS = int(round(T_END / DT))

# Shared top-level constants — must match the notebook. DT = 0.02 T_END = 12.0 SETPOINT = 1.0 # Damped 2nd-order mass-spring plant: m x_ddot + c x_dot + k x = F. # c=0.5, m=1, k=1 -> wn=1 rad/s, zeta=0.25 (mildly underdamped — visible overshoot). M_PLANT = 1.0 C_PLANT = 0.5 K_PLANT = 1.0 # Baseline controller gains: conservative tuning yields ~15% overshoot, # clean settling in 8 s on this plant. KP_BASE = 1.2 KI_BASE = 0.3 # Gain perturbations for the gradient finite-difference cross-check. KP_GRID = (0.6, 0.9, 1.2, 1.5, 1.8) N_STEPS = int(round(T_END / DT))

---

Slave source: the plant lives in a JaxonomyDiagramSlave whose DIAGRAM_FACTORY builds a one-LeafSystem diagram. Input = applied force F; outputs = position x and velocity v.¶

In [ ]:

_PLANT_SLAVE_SOURCE = f'''
"""Damped 2nd-order mass-spring plant FMU."""

import jaxonomy
from jaxonomy.library import Constant
from jaxonomy.framework import LeafSystem, parameters
from jaxonomy.backend import numpy_api as npa
from jaxonomy.library.fmu_slave import JaxonomyDiagramSlave

M = {M_PLANT!r}
C = {C_PLANT!r}
K = {K_PLANT!r}
DT = {DT!r}


class MassSpringPlant(LeafSystem):
    """Two-state plant: state = (x, v); ode = (v, (F - cv - kx)/m)."""

    @parameters(dynamic=["m", "c", "k"])
    def __init__(self, m=M, c=C, k=K, name="mass_spring"):
        super().__init__(name=name)

        def _ode(time, state, *inputs, **params):
            (F,) = inputs
            x, v = state.continuous_state
            xdot = v
            vdot = (F - params["c"] * v - params["k"] * x) / params["m"]
            return npa.array([xdot, vdot])

        self.declare_continuous_state(
            shape=(2,), default_value=npa.array([0.0, 0.0]), ode=_ode,
        )
        self.declare_input_port(name="F")
        self.declare_continuous_state_output(name="state")


def _build():
    """Build a single-block diagram exposing F as input and x as output.

    JaxonomyDiagramSlave auto-discovers the Constant block (named
    "F_in") as a writable FMI input via T-025c, so the master can
    write the applied force at every do_step. The slave reads "x" and
    "v" outputs through the diagram's exported scalars.
    """
    bld = jaxonomy.DiagramBuilder()
    F_in = bld.add(Constant(0.0, name="F_in"))
    plant = bld.add(MassSpringPlant(m=M, c=C, k=K))
    bld.connect(F_in.output_ports[0], plant.input_ports[0])

    # We need scalar x and v exports — pull them out of the 2-vector
    # state via a small LeafSystem that picks indices.
    class _Picker(LeafSystem):
        def __init__(self, idx, name):
            super().__init__(name=name)
            self.declare_input_port(name="in")

            def _out(time, state, *inputs, **params):
                return inputs[0][idx]

            self.declare_output_port(_out, name="y", requires_inputs=True)

    px = bld.add(_Picker(0, "pick_x"))
    pv = bld.add(_Picker(1, "pick_v"))
    bld.connect(plant.output_ports[0], px.input_ports[0])
    bld.connect(plant.output_ports[0], pv.input_ports[0])
    bld.export_output(px.output_ports[0], name="x")
    bld.export_output(pv.output_ports[0], name="v")
    return bld.build()


class MassSpringPlantFMU(JaxonomyDiagramSlave):
    DIAGRAM_FACTORY = staticmethod(_build)
    DT = DT
'''

_PLANT_SLAVE_SOURCE = f''' """Damped 2nd-order mass-spring plant FMU.""" import jaxonomy from jaxonomy.library import Constant from jaxonomy.framework import LeafSystem, parameters from jaxonomy.backend import numpy_api as npa from jaxonomy.library.fmu_slave import JaxonomyDiagramSlave M = {M_PLANT!r} C = {C_PLANT!r} K = {K_PLANT!r} DT = {DT!r} class MassSpringPlant(LeafSystem): """Two-state plant: state = (x, v); ode = (v, (F - cv - kx)/m).""" @parameters(dynamic=["m", "c", "k"]) def __init__(self, m=M, c=C, k=K, name="mass_spring"): super().__init__(name=name) def _ode(time, state, *inputs, **params): (F,) = inputs x, v = state.continuous_state xdot = v vdot = (F - params["c"] * v - params["k"] * x) / params["m"] return npa.array([xdot, vdot]) self.declare_continuous_state( shape=(2,), default_value=npa.array([0.0, 0.0]), ode=_ode, ) self.declare_input_port(name="F") self.declare_continuous_state_output(name="state") def _build(): """Build a single-block diagram exposing F as input and x as output. JaxonomyDiagramSlave auto-discovers the Constant block (named "F_in") as a writable FMI input via T-025c, so the master can write the applied force at every do_step. The slave reads "x" and "v" outputs through the diagram's exported scalars. """ bld = jaxonomy.DiagramBuilder() F_in = bld.add(Constant(0.0, name="F_in")) plant = bld.add(MassSpringPlant(m=M, c=C, k=K)) bld.connect(F_in.output_ports[0], plant.input_ports[0]) # We need scalar x and v exports — pull them out of the 2-vector # state via a small LeafSystem that picks indices. class _Picker(LeafSystem): def __init__(self, idx, name): super().__init__(name=name) self.declare_input_port(name="in") def _out(time, state, *inputs, **params): return inputs[0][idx] self.declare_output_port(_out, name="y", requires_inputs=True) px = bld.add(_Picker(0, "pick_x")) pv = bld.add(_Picker(1, "pick_v")) bld.connect(plant.output_ports[0], px.input_ports[0]) bld.connect(plant.output_ports[0], pv.input_ports[0]) bld.export_output(px.output_ports[0], name="x") bld.export_output(pv.output_ports[0], name="v") return bld.build() class MassSpringPlantFMU(JaxonomyDiagramSlave): DIAGRAM_FACTORY = staticmethod(_build) DT = DT '''

---

In-process closed-loop helper (Architecture A — both controller and plant native to jaxonomy).¶

In [ ]:

def _make_plant_class():
    from jaxonomy.framework import LeafSystem, parameters
    from jaxonomy.backend import numpy_api as npa

    class MassSpringPlant(LeafSystem):
        @parameters(dynamic=["m", "c", "k"])
        def __init__(self, m=M_PLANT, c=C_PLANT, k=K_PLANT, name="mass_spring"):
            super().__init__(name=name)

            def _ode(time, state, *inputs, **params):
                (F,) = inputs
                x, v = state.continuous_state
                xdot = v
                vdot = (F - params["c"] * v - params["k"] * x) / params["m"]
                return npa.array([xdot, vdot])

            self.declare_continuous_state(
                shape=(2,), default_value=npa.array([0.0, 0.0]), ode=_ode,
            )
            self.declare_input_port(name="F")
            self.declare_continuous_state_output(name="state")

    return MassSpringPlant

def _make_plant_class(): from jaxonomy.framework import LeafSystem, parameters from jaxonomy.backend import numpy_api as npa class MassSpringPlant(LeafSystem): @parameters(dynamic=["m", "c", "k"]) def __init__(self, m=M_PLANT, c=C_PLANT, k=K_PLANT, name="mass_spring"): super().__init__(name=name) def _ode(time, state, *inputs, **params): (F,) = inputs x, v = state.continuous_state xdot = v vdot = (F - params["c"] * v - params["k"] * x) / params["m"] return npa.array([xdot, vdot]) self.declare_continuous_state( shape=(2,), default_value=npa.array([0.0, 0.0]), ode=_ode, ) self.declare_input_port(name="F") self.declare_continuous_state_output(name="state") return MassSpringPlant

In [ ]:

def _build_pi_controller_diagram(kp, ki):
    """Build a PI controller. Inputs: measurement. Outputs: u."""
    import jaxonomy
    from jaxonomy.library import Adder, Constant, Gain
    from jaxonomy.library.dynamics import IntegratorDiscrete

    bld = jaxonomy.DiagramBuilder()
    sp = bld.add(Constant(SETPOINT, name="setpoint"))
    err = bld.add(Adder(2, operators="+-", name="err"))
    kp_blk = bld.add(Gain(kp, name="kp_gain"))
    ki_blk = bld.add(Gain(ki, name="ki_gain"))
    integ = bld.add(IntegratorDiscrete(
        dt=DT, initial_state=0.0, name="ierr"))
    add_u = bld.add(Adder(2, operators="++", name="add_u"))

    bld.connect(sp.output_ports[0], err.input_ports[0])
    bld.connect(err.output_ports[0], kp_blk.input_ports[0])
    bld.connect(err.output_ports[0], integ.input_ports[0])
    bld.connect(integ.output_ports[0], ki_blk.input_ports[0])
    bld.connect(kp_blk.output_ports[0], add_u.input_ports[0])
    bld.connect(ki_blk.output_ports[0], add_u.input_ports[1])
    bld.export_output(add_u.output_ports[0], name="u")
    bld.export_input(err.input_ports[1], name="measurement")
    return bld.build(name="PIController")

def _build_pi_controller_diagram(kp, ki): """Build a PI controller. Inputs: measurement. Outputs: u.""" import jaxonomy from jaxonomy.library import Adder, Constant, Gain from jaxonomy.library.dynamics import IntegratorDiscrete bld = jaxonomy.DiagramBuilder() sp = bld.add(Constant(SETPOINT, name="setpoint")) err = bld.add(Adder(2, operators="+-", name="err")) kp_blk = bld.add(Gain(kp, name="kp_gain")) ki_blk = bld.add(Gain(ki, name="ki_gain")) integ = bld.add(IntegratorDiscrete( dt=DT, initial_state=0.0, name="ierr")) add_u = bld.add(Adder(2, operators="++", name="add_u")) bld.connect(sp.output_ports[0], err.input_ports[0]) bld.connect(err.output_ports[0], kp_blk.input_ports[0]) bld.connect(err.output_ports[0], integ.input_ports[0]) bld.connect(integ.output_ports[0], ki_blk.input_ports[0]) bld.connect(kp_blk.output_ports[0], add_u.input_ports[0]) bld.connect(ki_blk.output_ports[0], add_u.input_ports[1]) bld.export_output(add_u.output_ports[0], name="u") bld.export_input(err.input_ports[1], name="measurement") return bld.build(name="PIController")

In [ ]:

def _close_loop_inprocess(kp, ki):
    """Architecture A — both controller and plant in jaxonomy. Returns
    (t, x, u) arrays."""
    import jaxonomy
    from jaxonomy.framework import LeafSystem

    MassSpringPlant = _make_plant_class()
    bld = jaxonomy.DiagramBuilder()
    ctl = bld.add(_build_pi_controller_diagram(kp, ki))
    plant = bld.add(MassSpringPlant(m=M_PLANT, c=C_PLANT, k=K_PLANT))

    # plant outputs the 2-vector state; we need x for the controller.
    class _PickX(LeafSystem):
        def __init__(self):
            super().__init__(name="pick_x")
            self.declare_input_port(name="state")

            def _out(time, state, *inputs, **params):
                return inputs[0][0]

            self.declare_output_port(_out, name="x", requires_inputs=True)

    pick = bld.add(_PickX())
    bld.connect(ctl.output_ports[0], plant.input_ports[0])
    bld.connect(plant.output_ports[0], pick.input_ports[0])
    bld.connect(pick.output_ports[0], ctl.input_ports[0])
    bld.export_output(pick.output_ports[0], name="x")
    bld.export_output(ctl.output_ports[0], name="u")
    loop = bld.build()

    res = jaxonomy.simulate(
        loop, loop.create_context(), (0.0, T_END),
        options=jaxonomy.SimulatorOptions(max_major_step_length=DT),
        recorded_signals={"x": loop.output_ports[0],
                          "u": loop.output_ports[1]},
    )
    return (np.asarray(res.time),
            np.asarray(res.outputs["x"]),
            np.asarray(res.outputs["u"]))

def _close_loop_inprocess(kp, ki): """Architecture A — both controller and plant in jaxonomy. Returns (t, x, u) arrays.""" import jaxonomy from jaxonomy.framework import LeafSystem MassSpringPlant = _make_plant_class() bld = jaxonomy.DiagramBuilder() ctl = bld.add(_build_pi_controller_diagram(kp, ki)) plant = bld.add(MassSpringPlant(m=M_PLANT, c=C_PLANT, k=K_PLANT)) # plant outputs the 2-vector state; we need x for the controller. class _PickX(LeafSystem): def __init__(self): super().__init__(name="pick_x") self.declare_input_port(name="state") def _out(time, state, *inputs, **params): return inputs[0][0] self.declare_output_port(_out, name="x", requires_inputs=True) pick = bld.add(_PickX()) bld.connect(ctl.output_ports[0], plant.input_ports[0]) bld.connect(plant.output_ports[0], pick.input_ports[0]) bld.connect(pick.output_ports[0], ctl.input_ports[0]) bld.export_output(pick.output_ports[0], name="x") bld.export_output(ctl.output_ports[0], name="u") loop = bld.build() res = jaxonomy.simulate( loop, loop.create_context(), (0.0, T_END), options=jaxonomy.SimulatorOptions(max_major_step_length=DT), recorded_signals={"x": loop.output_ports[0], "u": loop.output_ports[1]}, ) return (np.asarray(res.time), np.asarray(res.outputs["x"]), np.asarray(res.outputs["u"]))

In [ ]:

def _build_plant_fmu(tmpdir: Path):
    """Build the binary plant FMU. Returns (fmu_path, build_seconds, size)."""
    from jaxonomy.library.fmu_export import build_fmu

    slave_script = tmpdir / "plant_slave.py"
    slave_script.write_text(_PLANT_SLAVE_SOURCE)
    fmu_path = tmpdir / "MassSpringPlantFMU.fmu"

    t0 = time.perf_counter()
    build_fmu(str(slave_script), str(fmu_path))
    build_seconds = time.perf_counter() - t0
    return fmu_path, build_seconds, fmu_path.stat().st_size

def _build_plant_fmu(tmpdir: Path): """Build the binary plant FMU. Returns (fmu_path, build_seconds, size).""" from jaxonomy.library.fmu_export import build_fmu slave_script = tmpdir / "plant_slave.py" slave_script.write_text(_PLANT_SLAVE_SOURCE) fmu_path = tmpdir / "MassSpringPlantFMU.fmu" t0 = time.perf_counter() build_fmu(str(slave_script), str(fmu_path)) build_seconds = time.perf_counter() - t0 return fmu_path, build_seconds, fmu_path.stat().st_size

In [ ]:

def _close_loop_with_fmu(fmu_path: Path, kp, ki):
    """Architecture B — jaxonomy controller closed around a Modelica-style
    plant FMU. Returns (t, x, u)."""
    import jaxonomy
    from jaxonomy.library import ModelicaFMU

    bld = jaxonomy.DiagramBuilder()
    ctl = bld.add(_build_pi_controller_diagram(kp, ki))
    plant_fmu = bld.add(ModelicaFMU(
        file_name=str(fmu_path), dt=DT,
        input_names=["F_in"], output_names=["x", "v"],
        name="plant_fmu",
    ))
    bld.connect(ctl.output_ports[0], plant_fmu.input_ports[0])    # u -> F_in
    bld.connect(plant_fmu.output_ports[0], ctl.input_ports[0])    # x -> meas
    bld.export_output(plant_fmu.output_ports[0], name="x")
    bld.export_output(ctl.output_ports[0], name="u")
    loop = bld.build()
    res = jaxonomy.simulate(
        loop, loop.create_context(), (0.0, T_END),
        options=jaxonomy.SimulatorOptions(max_major_step_length=DT),
        recorded_signals={"x": loop.output_ports[0],
                          "u": loop.output_ports[1]},
    )
    return (np.asarray(res.time),
            np.asarray(res.outputs["x"]),
            np.asarray(res.outputs["u"])), 0.0

def _close_loop_with_fmu(fmu_path: Path, kp, ki): """Architecture B — jaxonomy controller closed around a Modelica-style plant FMU. Returns (t, x, u).""" import jaxonomy from jaxonomy.library import ModelicaFMU bld = jaxonomy.DiagramBuilder() ctl = bld.add(_build_pi_controller_diagram(kp, ki)) plant_fmu = bld.add(ModelicaFMU( file_name=str(fmu_path), dt=DT, input_names=["F_in"], output_names=["x", "v"], name="plant_fmu", )) bld.connect(ctl.output_ports[0], plant_fmu.input_ports[0]) # u -> F_in bld.connect(plant_fmu.output_ports[0], ctl.input_ports[0]) # x -> meas bld.export_output(plant_fmu.output_ports[0], name="x") bld.export_output(ctl.output_ports[0], name="u") loop = bld.build() res = jaxonomy.simulate( loop, loop.create_context(), (0.0, T_END), options=jaxonomy.SimulatorOptions(max_major_step_length=DT), recorded_signals={"x": loop.output_ports[0], "u": loop.output_ports[1]}, ) return (np.asarray(res.time), np.asarray(res.outputs["x"]), np.asarray(res.outputs["u"])), 0.0

In [ ]:

def _tracking_ise_inprocess(kp, ki):
    """Closed-loop tracking ISE on Architecture A (in-process plant).
    This is the function we'll differentiate via jax.grad."""
    t, x, _ = _close_loop_inprocess(float(kp), float(ki))
    return float(np.trapezoid((x - SETPOINT) ** 2, t))

def _tracking_ise_inprocess(kp, ki): """Closed-loop tracking ISE on Architecture A (in-process plant). This is the function we'll differentiate via jax.grad.""" t, x, _ = _close_loop_inprocess(float(kp), float(ki)) return float(np.trapezoid((x - SETPOINT) ** 2, t))

In [ ]:

def _fd_gradient_kp_ki(kp, ki, h=1e-3):
    """Central-difference gradient of tracking ISE w.r.t. (kp, ki)."""
    dJ_dkp = (
        _tracking_ise_inprocess(kp + h, ki)
        - _tracking_ise_inprocess(kp - h, ki)
    ) / (2 * h)
    dJ_dki = (
        _tracking_ise_inprocess(kp, ki + h)
        - _tracking_ise_inprocess(kp, ki - h)
    ) / (2 * h)
    return float(dJ_dkp), float(dJ_dki)

def _fd_gradient_kp_ki(kp, ki, h=1e-3): """Central-difference gradient of tracking ISE w.r.t. (kp, ki).""" dJ_dkp = ( _tracking_ise_inprocess(kp + h, ki) - _tracking_ise_inprocess(kp - h, ki) ) / (2 * h) dJ_dki = ( _tracking_ise_inprocess(kp, ki + h) - _tracking_ise_inprocess(kp, ki - h) ) / (2 * h) return float(dJ_dkp), float(dJ_dki)

In [ ]:

def _ad_gradient_kp_ki(kp, ki):
    """jax.grad of tracking ISE w.r.t. (kp, ki) — through Architecture A.

    Uses the cost-as-Integrator pattern: append a continuous-time
    Integrator that accumulates ``(x - r)^2`` and read its final
    state as the loss. This lets us run with ``save_time_series=False``
    (which ``enable_autodiff=True`` requires) and still get an
    end-to-end-differentiable scalar objective.
    """
    import jax
    import jax.numpy as jnp
    import jaxonomy
    from jaxonomy.backend import numpy_api as npa
    from jaxonomy.framework import LeafSystem, parameters

    MassSpringPlant = _make_plant_class()

    class _ISEAccumulator(LeafSystem):
        def __init__(self):
            super().__init__(name="ise")
            self.declare_input_port(name="x")

            def _ode(time, state, *inputs, **params):
                (x_meas,) = inputs
                err = SETPOINT - x_meas
                return err * err

            self.declare_continuous_state(
                shape=(), default_value=0.0, ode=_ode,
            )
            self.declare_continuous_state_output(name="ise")

    def _build_loop_with_gains():
        from jaxonomy.library import Adder, Constant, Gain
        from jaxonomy.library.dynamics import IntegratorDiscrete

        bld = jaxonomy.DiagramBuilder()
        sp = bld.add(Constant(SETPOINT, name="setpoint"))
        err = bld.add(Adder(2, operators="+-", name="err"))
        kp_blk = bld.add(Gain(KP_BASE, name="kp_gain"))
        ki_blk = bld.add(Gain(KI_BASE, name="ki_gain"))
        integ = bld.add(IntegratorDiscrete(
            dt=DT, initial_state=0.0, name="ierr"))
        add_u = bld.add(Adder(2, operators="++", name="add_u"))
        plant = bld.add(MassSpringPlant(m=M_PLANT, c=C_PLANT, k=K_PLANT))

        class _PickX(LeafSystem):
            def __init__(self):
                super().__init__(name="pick_x")
                self.declare_input_port(name="state")

                def _out(time, state, *inputs, **params):
                    return inputs[0][0]

                self.declare_output_port(_out, name="x", requires_inputs=True)

        pick = bld.add(_PickX())
        ise = bld.add(_ISEAccumulator())

        bld.connect(sp.output_ports[0], err.input_ports[0])
        bld.connect(err.output_ports[0], kp_blk.input_ports[0])
        bld.connect(err.output_ports[0], integ.input_ports[0])
        bld.connect(integ.output_ports[0], ki_blk.input_ports[0])
        bld.connect(kp_blk.output_ports[0], add_u.input_ports[0])
        bld.connect(ki_blk.output_ports[0], add_u.input_ports[1])
        bld.connect(add_u.output_ports[0], plant.input_ports[0])
        bld.connect(plant.output_ports[0], pick.input_ports[0])
        bld.connect(pick.output_ports[0], err.input_ports[1])
        bld.connect(pick.output_ports[0], ise.input_ports[0])

        bld.export_output(ise.output_ports[0], name="ise")
        return bld.build(), kp_blk.system_id, ki_blk.system_id

    loop, kp_id, ki_id = _build_loop_with_gains()

    def _forward(kp_arg, ki_arg):
        ctx = loop.create_context()
        ctx = ctx.with_subcontext(
            kp_id, ctx[kp_id].with_parameter("gain", kp_arg))
        ctx = ctx.with_subcontext(
            ki_id, ctx[ki_id].with_parameter("gain", ki_arg))
        res = jaxonomy.simulate(
            loop, ctx, (0.0, T_END),
            options=jaxonomy.SimulatorOptions(
                max_major_step_length=DT,
                enable_autodiff=True,
                save_time_series=False,
                return_context=True,
            ),
        )
        # Read terminal ISE from the final context via the diagram's
        # exported output port.
        final_ise = loop.output_ports[0].eval(res.context)
        return final_ise

    grad_fn = jax.grad(_forward, argnums=(0, 1))
    g = grad_fn(jnp.float64(kp), jnp.float64(ki))
    return float(g[0]), float(g[1])

def _ad_gradient_kp_ki(kp, ki): """jax.grad of tracking ISE w.r.t. (kp, ki) — through Architecture A. Uses the cost-as-Integrator pattern: append a continuous-time Integrator that accumulates ``(x - r)^2`` and read its final state as the loss. This lets us run with ``save_time_series=False`` (which ``enable_autodiff=True`` requires) and still get an end-to-end-differentiable scalar objective. """ import jax import jax.numpy as jnp import jaxonomy from jaxonomy.backend import numpy_api as npa from jaxonomy.framework import LeafSystem, parameters MassSpringPlant = _make_plant_class() class _ISEAccumulator(LeafSystem): def __init__(self): super().__init__(name="ise") self.declare_input_port(name="x") def _ode(time, state, *inputs, **params): (x_meas,) = inputs err = SETPOINT - x_meas return err * err self.declare_continuous_state( shape=(), default_value=0.0, ode=_ode, ) self.declare_continuous_state_output(name="ise") def _build_loop_with_gains(): from jaxonomy.library import Adder, Constant, Gain from jaxonomy.library.dynamics import IntegratorDiscrete bld = jaxonomy.DiagramBuilder() sp = bld.add(Constant(SETPOINT, name="setpoint")) err = bld.add(Adder(2, operators="+-", name="err")) kp_blk = bld.add(Gain(KP_BASE, name="kp_gain")) ki_blk = bld.add(Gain(KI_BASE, name="ki_gain")) integ = bld.add(IntegratorDiscrete( dt=DT, initial_state=0.0, name="ierr")) add_u = bld.add(Adder(2, operators="++", name="add_u")) plant = bld.add(MassSpringPlant(m=M_PLANT, c=C_PLANT, k=K_PLANT)) class _PickX(LeafSystem): def __init__(self): super().__init__(name="pick_x") self.declare_input_port(name="state") def _out(time, state, *inputs, **params): return inputs[0][0] self.declare_output_port(_out, name="x", requires_inputs=True) pick = bld.add(_PickX()) ise = bld.add(_ISEAccumulator()) bld.connect(sp.output_ports[0], err.input_ports[0]) bld.connect(err.output_ports[0], kp_blk.input_ports[0]) bld.connect(err.output_ports[0], integ.input_ports[0]) bld.connect(integ.output_ports[0], ki_blk.input_ports[0]) bld.connect(kp_blk.output_ports[0], add_u.input_ports[0]) bld.connect(ki_blk.output_ports[0], add_u.input_ports[1]) bld.connect(add_u.output_ports[0], plant.input_ports[0]) bld.connect(plant.output_ports[0], pick.input_ports[0]) bld.connect(pick.output_ports[0], err.input_ports[1]) bld.connect(pick.output_ports[0], ise.input_ports[0]) bld.export_output(ise.output_ports[0], name="ise") return bld.build(), kp_blk.system_id, ki_blk.system_id loop, kp_id, ki_id = _build_loop_with_gains() def _forward(kp_arg, ki_arg): ctx = loop.create_context() ctx = ctx.with_subcontext( kp_id, ctx[kp_id].with_parameter("gain", kp_arg)) ctx = ctx.with_subcontext( ki_id, ctx[ki_id].with_parameter("gain", ki_arg)) res = jaxonomy.simulate( loop, ctx, (0.0, T_END), options=jaxonomy.SimulatorOptions( max_major_step_length=DT, enable_autodiff=True, save_time_series=False, return_context=True, ), ) # Read terminal ISE from the final context via the diagram's # exported output port. final_ise = loop.output_ports[0].eval(res.context) return final_ise grad_fn = jax.grad(_forward, argnums=(0, 1)) g = grad_fn(jnp.float64(kp), jnp.float64(ki)) return float(g[0]), float(g[1])

In [ ]:

def main(write_placeholder: bool = False):
    if write_placeholder:
        # Structurally-plausible numbers for environments without
        # pythonfmu / fmpy. Real runs overwrite this NPZ.
        t = np.linspace(0.0, T_END, N_STEPS + 1)
        omega = 1.0
        x = SETPOINT * (1.0 - np.exp(-0.5 * t) * np.cos(omega * t))
        np.savez_compressed(
            NPZ_OUT,
            placeholder_flag=np.array(True),
            t_A=t, x_A=x, u_A=np.zeros_like(t),
            t_B=t, x_B=x, u_B=np.zeros_like(t),
            kp_grid=np.array(KP_GRID),
            ise_grid=np.linspace(2.0, 0.5, len(KP_GRID)),
            grad_ad_kp=np.array(-1.2),
            grad_ad_ki=np.array(-0.6),
            grad_fd_kp=np.array(-1.21),
            grad_fd_ki=np.array(-0.61),
            build_seconds=np.array(0.1),
            fmu_size_bytes=np.array(900000),
            wall_arch_A=np.array(0.5),
            wall_arch_B=np.array(20.0),
            wall_grad_ad=np.array(8.0),
            wall_grad_fd=np.array(2.0),
        )
        print(f"placeholder NPZ written to {NPZ_OUT}")
        return

    with tempfile.TemporaryDirectory() as td:
        tmpdir = Path(td)

        print("[1/5] Architecture A (all in-process) — closed loop ...")
        t0 = time.perf_counter()
        t_A, x_A, u_A = _close_loop_inprocess(KP_BASE, KI_BASE)
        wall_A = time.perf_counter() - t0
        print(f"      wall: {wall_A:.2f} s; "
              f"final x: {x_A[-1]:.4f} (target {SETPOINT}).")

        print("[2/5] Building binary plant FMU ...")
        fmu_path, build_seconds, fmu_size = _build_plant_fmu(tmpdir)
        print(f"      FMU built in {build_seconds:.3f} s "
              f"({fmu_size / 1024:.0f} KiB).")

        print("[3/5] Architecture B (jaxonomy PI + plant FMU) ...")
        t0 = time.perf_counter()
        (t_B, x_B, u_B), _ = _close_loop_with_fmu(
            fmu_path, KP_BASE, KI_BASE)
        wall_B = time.perf_counter() - t0
        print(f"      wall: {wall_B:.2f} s; final x: {x_B[-1]:.4f}.")

        print("[4/5] Tracking ISE on a kp-grid (in-process plant) ...")
        ise_grid = np.array(
            [_tracking_ise_inprocess(kp, KI_BASE) for kp in KP_GRID]
        )
        for kp, ise in zip(KP_GRID, ise_grid):
            print(f"      kp={kp:.2f} -> ISE={ise:.6f}")

        print("[5/5] Gradient: jax.grad vs central-difference ...")
        t0 = time.perf_counter()
        grad_ad_kp, grad_ad_ki = _ad_gradient_kp_ki(KP_BASE, KI_BASE)
        wall_grad_ad = time.perf_counter() - t0
        print(f"      jax.grad: dISE/dkp={grad_ad_kp:.6f}, "
              f"dISE/dki={grad_ad_ki:.6f} (wall {wall_grad_ad:.2f} s)")

        t0 = time.perf_counter()
        grad_fd_kp, grad_fd_ki = _fd_gradient_kp_ki(KP_BASE, KI_BASE)
        wall_grad_fd = time.perf_counter() - t0
        print(f"      FD     : dISE/dkp={grad_fd_kp:.6f}, "
              f"dISE/dki={grad_fd_ki:.6f} (wall {wall_grad_fd:.2f} s)")

    np.savez_compressed(
        NPZ_OUT,
        placeholder_flag=np.array(False),
        t_A=t_A, x_A=x_A, u_A=u_A,
        t_B=t_B, x_B=x_B, u_B=u_B,
        kp_grid=np.array(KP_GRID),
        ise_grid=ise_grid,
        grad_ad_kp=np.array(grad_ad_kp),
        grad_ad_ki=np.array(grad_ad_ki),
        grad_fd_kp=np.array(grad_fd_kp),
        grad_fd_ki=np.array(grad_fd_ki),
        build_seconds=np.array(build_seconds),
        fmu_size_bytes=np.array(fmu_size),
        wall_arch_A=np.array(wall_A),
        wall_arch_B=np.array(wall_B),
        wall_grad_ad=np.array(wall_grad_ad),
        wall_grad_fd=np.array(wall_grad_fd),
    )
    print(f"\nNPZ written to {NPZ_OUT}")

    # Headline numbers.
    grid = np.arange(N_STEPS + 1) * DT
    xA_g = np.interp(grid, t_A, x_A)
    xB_g = np.interp(grid, t_B, x_B)
    err_x_full = float(np.max(np.abs(xA_g - xB_g)))
    err_x_late = float(np.max(np.abs(xA_g[10:] - xB_g[10:])))
    print("\nHEADLINE NUMBERS")
    print(f"  Closed-loop A vs B (FMU plant) bit-compare, full {T_END:g} s: "
          f"{err_x_full:.3e}")
    print(f"  Closed-loop A vs B after first 10 samples (post-transient):  "
          f"{err_x_late:.3e}")
    rel_err_kp = abs(grad_ad_kp - grad_fd_kp) / max(abs(grad_fd_kp), 1e-12)
    rel_err_ki = abs(grad_ad_ki - grad_fd_ki) / max(abs(grad_fd_ki), 1e-12)
    print(f"  jax.grad vs FD on Kp: rel-err = {rel_err_kp * 100:.2f}%")
    print(f"  jax.grad vs FD on Ki: rel-err = {rel_err_ki * 100:.2f}%")

def main(write_placeholder: bool = False): if write_placeholder: # Structurally-plausible numbers for environments without # pythonfmu / fmpy. Real runs overwrite this NPZ. t = np.linspace(0.0, T_END, N_STEPS + 1) omega = 1.0 x = SETPOINT * (1.0 - np.exp(-0.5 * t) * np.cos(omega * t)) np.savez_compressed( NPZ_OUT, placeholder_flag=np.array(True), t_A=t, x_A=x, u_A=np.zeros_like(t), t_B=t, x_B=x, u_B=np.zeros_like(t), kp_grid=np.array(KP_GRID), ise_grid=np.linspace(2.0, 0.5, len(KP_GRID)), grad_ad_kp=np.array(-1.2), grad_ad_ki=np.array(-0.6), grad_fd_kp=np.array(-1.21), grad_fd_ki=np.array(-0.61), build_seconds=np.array(0.1), fmu_size_bytes=np.array(900000), wall_arch_A=np.array(0.5), wall_arch_B=np.array(20.0), wall_grad_ad=np.array(8.0), wall_grad_fd=np.array(2.0), ) print(f"placeholder NPZ written to {NPZ_OUT}") return with tempfile.TemporaryDirectory() as td: tmpdir = Path(td) print("[1/5] Architecture A (all in-process) — closed loop ...") t0 = time.perf_counter() t_A, x_A, u_A = _close_loop_inprocess(KP_BASE, KI_BASE) wall_A = time.perf_counter() - t0 print(f" wall: {wall_A:.2f} s; " f"final x: {x_A[-1]:.4f} (target {SETPOINT}).") print("[2/5] Building binary plant FMU ...") fmu_path, build_seconds, fmu_size = _build_plant_fmu(tmpdir) print(f" FMU built in {build_seconds:.3f} s " f"({fmu_size / 1024:.0f} KiB).") print("[3/5] Architecture B (jaxonomy PI + plant FMU) ...") t0 = time.perf_counter() (t_B, x_B, u_B), _ = _close_loop_with_fmu( fmu_path, KP_BASE, KI_BASE) wall_B = time.perf_counter() - t0 print(f" wall: {wall_B:.2f} s; final x: {x_B[-1]:.4f}.") print("[4/5] Tracking ISE on a kp-grid (in-process plant) ...") ise_grid = np.array( [_tracking_ise_inprocess(kp, KI_BASE) for kp in KP_GRID] ) for kp, ise in zip(KP_GRID, ise_grid): print(f" kp={kp:.2f} -> ISE={ise:.6f}") print("[5/5] Gradient: jax.grad vs central-difference ...") t0 = time.perf_counter() grad_ad_kp, grad_ad_ki = _ad_gradient_kp_ki(KP_BASE, KI_BASE) wall_grad_ad = time.perf_counter() - t0 print(f" jax.grad: dISE/dkp={grad_ad_kp:.6f}, " f"dISE/dki={grad_ad_ki:.6f} (wall {wall_grad_ad:.2f} s)") t0 = time.perf_counter() grad_fd_kp, grad_fd_ki = _fd_gradient_kp_ki(KP_BASE, KI_BASE) wall_grad_fd = time.perf_counter() - t0 print(f" FD : dISE/dkp={grad_fd_kp:.6f}, " f"dISE/dki={grad_fd_ki:.6f} (wall {wall_grad_fd:.2f} s)") np.savez_compressed( NPZ_OUT, placeholder_flag=np.array(False), t_A=t_A, x_A=x_A, u_A=u_A, t_B=t_B, x_B=x_B, u_B=u_B, kp_grid=np.array(KP_GRID), ise_grid=ise_grid, grad_ad_kp=np.array(grad_ad_kp), grad_ad_ki=np.array(grad_ad_ki), grad_fd_kp=np.array(grad_fd_kp), grad_fd_ki=np.array(grad_fd_ki), build_seconds=np.array(build_seconds), fmu_size_bytes=np.array(fmu_size), wall_arch_A=np.array(wall_A), wall_arch_B=np.array(wall_B), wall_grad_ad=np.array(wall_grad_ad), wall_grad_fd=np.array(wall_grad_fd), ) print(f"\nNPZ written to {NPZ_OUT}") # Headline numbers. grid = np.arange(N_STEPS + 1) * DT xA_g = np.interp(grid, t_A, x_A) xB_g = np.interp(grid, t_B, x_B) err_x_full = float(np.max(np.abs(xA_g - xB_g))) err_x_late = float(np.max(np.abs(xA_g[10:] - xB_g[10:]))) print("\nHEADLINE NUMBERS") print(f" Closed-loop A vs B (FMU plant) bit-compare, full {T_END:g} s: " f"{err_x_full:.3e}") print(f" Closed-loop A vs B after first 10 samples (post-transient): " f"{err_x_late:.3e}") rel_err_kp = abs(grad_ad_kp - grad_fd_kp) / max(abs(grad_fd_kp), 1e-12) rel_err_ki = abs(grad_ad_ki - grad_fd_ki) / max(abs(grad_fd_ki), 1e-12) print(f" jax.grad vs FD on Kp: rel-err = {rel_err_kp * 100:.2f}%") print(f" jax.grad vs FD on Ki: rel-err = {rel_err_ki * 100:.2f}%")

In [ ]:

if __name__ == "__main__":
    write_placeholder = "--write-placeholder" in sys.argv
    main(write_placeholder=write_placeholder)

if __name__ == "__main__": write_placeholder = "--write-placeholder" in sys.argv main(write_placeholder=write_placeholder)