Online estimation example

In [ ]:

# SPDX-License-Identifier: MIT
"""
Online / real-time parameter estimation examples.

Two examples demonstrating the new online estimation blocks:

  A. RecursiveLeastSquares
     Identifies the coefficients θ of a linear model y = φᵀθ one measurement
     at a time, with optional forgetting factor for time-varying systems.

  B. AugmentedStateEKF
     Jointly estimates the state and an unknown model parameter of a discrete-
     time dynamical system using an Extended Kalman Filter on the augmented
     state z = [x; θ].
"""

# SPDX-License-Identifier: MIT """ Online / real-time parameter estimation examples. Two examples demonstrating the new online estimation blocks: A. RecursiveLeastSquares Identifies the coefficients θ of a linear model y = φᵀθ one measurement at a time, with optional forgetting factor for time-varying systems. B. AugmentedStateEKF Jointly estimates the state and an unknown model parameter of a discrete- time dynamical system using an Extended Kalman Filter on the augmented state z = [x; θ]. """

In [ ]:

import numpy as np
import jax.numpy as jnp

import numpy as np import jax.numpy as jnp

In [ ]:

import jaxonomy
from jaxonomy import DiagramBuilder, SimulatorOptions
from jaxonomy.library import Constant, RecursiveLeastSquares, AugmentedStateEKF

import jaxonomy from jaxonomy import DiagramBuilder, SimulatorOptions from jaxonomy.library import Constant, RecursiveLeastSquares, AugmentedStateEKF

In [ ]:

# ─────────────────────────────────────────────────────────────────────────────
# A.  Recursive Least Squares
# ─────────────────────────────────────────────────────────────────────────────
print("=" * 60)
print("A. Recursive Least Squares parameter identification")
print("=" * 60)

# ───────────────────────────────────────────────────────────────────────────── # A. Recursive Least Squares # ───────────────────────────────────────────────────────────────────────────── print("=" * 60) print("A. Recursive Least Squares parameter identification") print("=" * 60)

In [ ]:

# True model:  y = 2·φ₀ − 1·φ₁ + 3·φ₂
TRUE_THETA = np.array([2.0, -1.0, 3.0])
N_PARAMS = 3
DT = 0.1
T_END = 10.0  # simulate 10 s  (100 steps)

# True model: y = 2·φ₀ − 1·φ₁ + 3·φ₂ TRUE_THETA = np.array([2.0, -1.0, 3.0]) N_PARAMS = 3 DT = 0.1 T_END = 10.0 # simulate 10 s (100 steps)

In [ ]:

# Constant regressor (the RLS will project onto this direction)
PHI = np.array([1.0, 2.0, 1.0])
Y_TRUE = float(PHI @ TRUE_THETA)  # = 2 + (-2) + 3 = 3 → actually 2 -2 +3 = 3

# Constant regressor (the RLS will project onto this direction) PHI = np.array([1.0, 2.0, 1.0]) Y_TRUE = float(PHI @ TRUE_THETA) # = 2 + (-2) + 3 = 3 → actually 2 -2 +3 = 3

In [ ]:

print(f"True θ = {TRUE_THETA}")
print(f"φ      = {PHI}")
print(f"y      = {Y_TRUE:.4f}")

print(f"True θ = {TRUE_THETA}") print(f"φ = {PHI}") print(f"y = {Y_TRUE:.4f}")

In [ ]:

b = DiagramBuilder()
rls = b.add(
    RecursiveLeastSquares(
        dt=DT,
        n_params=N_PARAMS,
        forgetting_factor=1.0,   # no forgetting — batch equivalent
        name="rls",
    )
)
phi_src = b.add(Constant(jnp.array(PHI), name="phi"))
y_src   = b.add(Constant(jnp.array(Y_TRUE), name="y"))
b.connect(phi_src.output_ports[0], rls.input_ports[0])
b.connect(y_src.output_ports[0],   rls.input_ports[1])
diagram = b.build()

b = DiagramBuilder() rls = b.add( RecursiveLeastSquares( dt=DT, n_params=N_PARAMS, forgetting_factor=1.0, # no forgetting — batch equivalent name="rls", ) ) phi_src = b.add(Constant(jnp.array(PHI), name="phi")) y_src = b.add(Constant(jnp.array(Y_TRUE), name="y")) b.connect(phi_src.output_ports[0], rls.input_ports[0]) b.connect(y_src.output_ports[0], rls.input_ports[1]) diagram = b.build()

In [ ]:

ctx = diagram.create_context()
nseg = int(T_END / DT)
options = SimulatorOptions(max_major_steps=10 * nseg, max_major_step_length=DT)
sol = jaxonomy.simulate(
    diagram, ctx, (0.0, T_END), options=options,
    recorded_signals={
        "theta_hat": rls.output_ports[0],
        "P_diag":    rls.output_ports[1],   # full P matrix per step
        "error":     rls.output_ports[2],
    },
)

ctx = diagram.create_context() nseg = int(T_END / DT) options = SimulatorOptions(max_major_steps=10 * nseg, max_major_step_length=DT) sol = jaxonomy.simulate( diagram, ctx, (0.0, T_END), options=options, recorded_signals={ "theta_hat": rls.output_ports[0], "P_diag": rls.output_ports[1], # full P matrix per step "error": rls.output_ports[2], }, )

In [ ]:

theta_final = np.array(sol.outputs["theta_hat"][-1])
error_final = float(sol.outputs["error"][-1])
print(f"\nAfter {nseg} steps:")
print(f"  θ̂     = {theta_final}")
print(f"  error = {error_final:.4e}  (prediction residual)")
print("  (note: with constant φ only the direction of φ is identified)")

theta_final = np.array(sol.outputs["theta_hat"][-1]) error_final = float(sol.outputs["error"][-1]) print(f"\nAfter {nseg} steps:") print(f" θ̂ = {theta_final}") print(f" error = {error_final:.4e} (prediction residual)") print(" (note: with constant φ only the direction of φ is identified)")

In [ ]:

# ─────────────────────────────────────────────────────────────────────────────
# B.  Augmented-State EKF  (joint state + parameter estimation)
# ─────────────────────────────────────────────────────────────────────────────
print()
print("=" * 60)
print("B. Augmented-State EKF: joint state + parameter estimation")
print("=" * 60)

# ───────────────────────────────────────────────────────────────────────────── # B. Augmented-State EKF (joint state + parameter estimation) # ───────────────────────────────────────────────────────────────────────────── print() print("=" * 60) print("B. Augmented-State EKF: joint state + parameter estimation") print("=" * 60)

In [ ]:

# True system:  x[n+1] = a·x[n],   y[n] = x[n]
# Unknown parameter: a  (true value 0.9)
TRUE_A = 0.9
DT_EKF = 0.1
T_EKF  = 5.0
THETA_INIT = 0.5   # initial guess for a, far from truth

# True system: x[n+1] = a·x[n], y[n] = x[n] # Unknown parameter: a (true value 0.9) TRUE_A = 0.9 DT_EKF = 0.1 T_EKF = 5.0 THETA_INIT = 0.5 # initial guess for a, far from truth

In [ ]:

print(f"True decay coefficient a = {TRUE_A}")
print(f"Initial guess           a₀ = {THETA_INIT}")

print(f"True decay coefficient a = {TRUE_A}") print(f"Initial guess a₀ = {THETA_INIT}")

In [ ]:

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

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

In [ ]:

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

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

In [ ]:

b2 = DiagramBuilder()
aekf = b2.add(
    AugmentedStateEKF(
        dt=DT_EKF,
        nx=1,
        n_params=1,
        forward=forward,
        observation=observation,
        G_x_func=lambda t: jnp.eye(1),
        Q_x_func=lambda t, x, u, th: jnp.eye(1) * 1e-4,   # small state noise
        Q_theta=jnp.eye(1) * 1e-3,                          # allow parameter drift
        R_func=lambda t: jnp.eye(1) * 1e-2,                # low measurement noise
        x_hat_0=jnp.array([1.0]),
        P_hat_0_x=jnp.eye(1) * 0.1,
        theta_hat_0=jnp.array([THETA_INIT]),
        P_hat_0_theta=jnp.eye(1) * 1.0,                    # high initial uncertainty
        name="aekf",
    )
)
u_src = b2.add(Constant(jnp.zeros(1), name="u"))
# Observation: constant y = 1.0 (system at x=1, decaying to 0 not simulated here)
y_src = b2.add(Constant(jnp.ones(1), name="y"))
b2.connect(u_src.output_ports[0], aekf.input_ports[0])
b2.connect(y_src.output_ports[0], aekf.input_ports[1])
diagram2 = b2.build()

b2 = DiagramBuilder() aekf = b2.add( AugmentedStateEKF( dt=DT_EKF, nx=1, n_params=1, forward=forward, observation=observation, G_x_func=lambda t: jnp.eye(1), Q_x_func=lambda t, x, u, th: jnp.eye(1) * 1e-4, # small state noise Q_theta=jnp.eye(1) * 1e-3, # allow parameter drift R_func=lambda t: jnp.eye(1) * 1e-2, # low measurement noise x_hat_0=jnp.array([1.0]), P_hat_0_x=jnp.eye(1) * 0.1, theta_hat_0=jnp.array([THETA_INIT]), P_hat_0_theta=jnp.eye(1) * 1.0, # high initial uncertainty name="aekf", ) ) u_src = b2.add(Constant(jnp.zeros(1), name="u")) # Observation: constant y = 1.0 (system at x=1, decaying to 0 not simulated here) y_src = b2.add(Constant(jnp.ones(1), name="y")) b2.connect(u_src.output_ports[0], aekf.input_ports[0]) b2.connect(y_src.output_ports[0], aekf.input_ports[1]) diagram2 = b2.build()

In [ ]:

ctx2 = diagram2.create_context()
nseg2 = int(T_EKF / DT_EKF)
options2 = SimulatorOptions(max_major_steps=10 * nseg2, max_major_step_length=DT_EKF)
sol2 = jaxonomy.simulate(
    diagram2, ctx2, (0.0, T_EKF), options=options2,
    recorded_signals={
        "x_hat":     aekf.output_ports[0],
        "theta_hat": aekf.output_ports[1],
    },
)

ctx2 = diagram2.create_context() nseg2 = int(T_EKF / DT_EKF) options2 = SimulatorOptions(max_major_steps=10 * nseg2, max_major_step_length=DT_EKF) sol2 = jaxonomy.simulate( diagram2, ctx2, (0.0, T_EKF), options=options2, recorded_signals={ "x_hat": aekf.output_ports[0], "theta_hat": aekf.output_ports[1], }, )

In [ ]:

x_hat_final     = float(sol2.outputs["x_hat"][-1, 0])
theta_hat_final = float(sol2.outputs["theta_hat"][-1, 0])
print(f"\nAfter {nseg2} steps:")
print(f"  x̂        = {x_hat_final:.4f}")
print(f"  θ̂ (est a) = {theta_hat_final:.4f}  (true a = {TRUE_A})")
print(f"  |θ̂ - a|   = {abs(theta_hat_final - TRUE_A):.4f}")

x_hat_final = float(sol2.outputs["x_hat"][-1, 0]) theta_hat_final = float(sol2.outputs["theta_hat"][-1, 0]) print(f"\nAfter {nseg2} steps:") print(f" x̂ = {x_hat_final:.4f}") print(f" θ̂ (est a) = {theta_hat_final:.4f} (true a = {TRUE_A})") print(f" |θ̂ - a| = {abs(theta_hat_final - TRUE_A):.4f}")

In [ ]:

print("\nDone.")

print("\nDone.")