Back to Portfolio of Projects | LearnableLoopAI.com | Blog |

In this project we apply the established Active Inference (AIF) modeling approach (making use of Bond Graphs, Generalized Filters, and Generalized Coordinates of Motion) to an electric system. The latter is not a complicated system. Also, the generalized coordinates part of the experiment needs more tuning.

A few coding conventions:

Variable names will generally reflect the mathematical names, rather than descriptive names.
Bold variable names usually start with an underscore, i.e. ’_’. That means all structured variables including vectors (i.e. tensors). The focus is on the code equivalents of the mathematics, i.e. not on Python structured variables, e.g. lists.
Global variable names (i.e. used across cells) are usually indicated by the Python ‘global’ keyword
External (i.e. true) states and parameters: The ‘^x’ symbol is used for code variables, i.e. ‘v^x’. In other words, superscript ‘x’ is used to imitate superscript ’*’ in code. In documentation the standard ‘\(v^*\)’ is used.
This notebook follows a standardized framework. This is my own framework, developed over a number of years, and is informed by the CRISP-DM framework, the work of Warren Powell (Princeton), the work of Bert De Vries, the work of Karl Friston as well as that of the larger Active Inference community.

A few charting conventions:

Linestyles
- ‘-’ for vector component 0
- ‘–’ for vector component 1
- ‘-.’ for vector component 2
- ‘:’ for vector component 3
Widths
- Wide for Environment
- Narrow for Agent
Colors
- ’C’rimsons / reds for ’C’ontrols / actions
- ’A’quas / blues for ’A’utonomous states
- Greys for hidden / latent states
- ’O’ranges for ’O’bservations
- ’F’orest-greens for ’F’ree energies (VFE, EFE, GFE, etc)

! python --version

Python 3.11.14

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import pandas as pd
from matplotlib.gridspec import GridSpec
from matplotlib.ticker import MaxNLocator

1 BUSINESS UNDERSTANDING

See Part 1

2 DATA UNDERSTANDING

N/A. We will not make use of collected data for training. Training happens online during simulation.

3 DATA PREPARATION

N/A. We will not make use of collected data for training. Training happens online during simulation.

4 MODELING

4.1 Narrative

Please see the description from section 1: BUSINESS UNDERSTANDING

4.2 Core Elements

This section attempts to answer three important questions:

What metrics are we going to track?
What decisions do we intend to make?
What are the sources of uncertainty?

Metrics

The voltage across the capacitor \(C_3\) will be tracked, \(y(t)\)

Decisions

The decision involves what control current should be injected into the circuit during each time step, \(a_5(t)\)

Sources of Uncertainty

A source of uncertainty is the disturbance current \(v_7(t)\). There is also uncertainty in the (colored) process and measurement noise.

4.3 System-Under-Steer / Environment / Generative Process

4.3.1 State variables

The state at time \(t\) of the system-under-steer/environment (envir) is given by

\[ \begin{aligned} \mathbf{x}^*_t &= (x^{\mathrm{Charge}}_t, x^{\mathrm{Flux}}_t) &= (q_3(t), p_6(t)) \end{aligned} \] where

\(q_3(t)\): Charge stored in capacitor \(C_3\) in C
- represented by the q3 code variable
\(p_6(t)\): Flux in inductor \(L_6\) in Wb
- represented by the p6 code variable

4.3.2 Decision variables

The decision variables represent what we control.

The environment is steered by decisions/actions \(\mathbf{a}_t\). Each component of this vector is called a control factor or control state factor. The action at time \(t\) is given by

\[ \begin{aligned} \mathbf{a}_t &= (a^{\mathrm{Current}}_{t}) &= (a_5(t)) \end{aligned} \] where

\(a_5(t)\): Current applied by the current source Sf to the circuit in \(A\)
- represented by the a5 code variable

4.3.3 Exogenous information variables

The exogenous information is not controllable by the agent. \[ \begin{aligned} \mathbf{v}^*_t &= (v^{\mathrm{DisturbanceCurrent}}_t) &= (v_7(t)) \end{aligned} \]

where

\(v_7(t)\): Disturbance Current in \(A\)
- represented by the v7 code variable

4.3.4 State Transition

To find the next state for each time step, the generative process starts with the transition function \(f\). To this is added system or process noise to arrive at the next state:

\[ \dot{\mathbf{x}}^* = f_{\cal E}(\mathbf{x}^*, \mathbf{v}^*, \mathbf{a}; \theta^*_x) + \boldsymbol{\omega}^*_x \]

4.3.5 Observation Generation

To find the observation generated by the external state of the generative process for each time step, the generative process starts with the generation function \(g\). To this is added observation or measurement noise to arrive at the observation:

\[ \mathbf{y} = g_{\cal E}(\mathbf{x}^*, \mathbf{v}^*, \mathbf{a}; \theta^*_y) + \boldsymbol{\omega}^*_y \]

4.3.6 Implementation of the System-Under-Steer / Environment / Generative Process

\[ \begin{align} \dot{\mathbf{x}}^* &= f_{\cal E}(\mathbf{x}^*, \mathbf{v}^*, \mathbf{a}; \theta^*_x) + \boldsymbol{\omega}^*_{x}(t) \\ \dot{q}^*_3(t) &= \frac{1}{L_6}p^*_{6}(t) - v^*_7(t) - a_5(t) + \omega^*_{x0}(t) \\ \dot{p}^*_6(t) &= - \frac{1}{C_3}q^*_{3}(t) - \frac{R_2}{L_6}p^*_{6}(t) + R_2 v^*_7(t) + R_2 a_5(t) + \omega^*_{x1}(t) \\ \end{align} \]

\[ \begin{align} \mathbf{y} &= g_{\cal E}(\mathbf{x^*, \mathbf{v^*}}; \theta^*_y) + \boldsymbol{\omega}^*_y \\ y(t) &= \frac{1}{C_3} q^*_{3}(t) + \omega^*_y(t) \\ \end{align} \]

where

\(x^*_0 = q_3\): Stored charge in capacitor C3 in C
- subscript 0 refers to the 0th component of the \(\mathbf{x^*}\) vector
\(x^*_1 = p_6\): Flux in inductor L6 in Wb
- subscript 1 refers to the 1th component of the \(\mathbf{x^*}\) vector
\(v^*_0\): Disturbance current (exogenous force) in \(A\)
- subscript 0 refers to the 0th component of the \(\mathbf{v^*}\) vector
\(y^*_0\): Voltage across capacitor C3 in \(V\)
- subscript 0 refers to the 0th component of the \(\mathbf{y}\) vector
\(\omega^*_x\): Process noise
\(\omega^*_y\): Observation noise
\(R_2\): Resistance in Ω
\(C_3\): Capacitance in F
\(L_6\): Inductance in H
\(\theta^*_x = \{R_2, C_3, L_6 \}\)

ALPHA = 0.6
ENV_WIDTH = 6. ## 4. 8.
AGT_WIDTH = 2. ## 2.
AGT_LINESTYLES = ['-', '--', '-.', ':'] ## for agent
ORANGE_HUE_BY_SATURATION = [
    "#FFA500", ## 100% bright orange
    "#FFB347", ## 80% light vivid orange
    "#FFD580", ## 60% softer pastel orange
    "#FFE5B4"  ## 40% light faded orange
]
ORANGE_HUE_BY_BRIGHTNESS = [
    "#FFA500", ## 100% pure bright orange
    "#FFC04D", ## 80% light orange
    "#FFD699", ## 60% pastel orange
    "#FFE5CC"  ## 40% very pale orange
]
ORANGE_HUE = ORANGE_HUE_BY_BRIGHTNESS
styles = {
    #### envir
    "a[0]": {
        "color": "crimson",
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "a[1]": {
        "color": "pink",
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "xˣ[0]": {
        "color": "black",
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "xˣ[1]": {
        "color": "darkgrey",
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "xˣ[2]": {
        "color": "grey",
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "xˣ[3]": {
        "color": "lightgrey",
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "vˣ[0]": {
        "color": "blue",
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "vˣ[1]": {
        "color": "dodgerblue",
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "y[0]": {
        "color": ORANGE_HUE[0],
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "y[1]": {
        "color": ORANGE_HUE[1],
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "y[2]": {
        "color": ORANGE_HUE[2],
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    "y[3]": {
        "color": ORANGE_HUE[3],
        "linestyle": "-",
        "linewidth": ENV_WIDTH,
        "alpha": ALPHA
    },
    #### agent
    "μ_x[0]": {
        "color": "black",
        "linestyle": AGT_LINESTYLES[0],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "μ_x[1]": {
        "color": "black",
        "linestyle": AGT_LINESTYLES[1],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "v[0]": {
        "color": "blue",
        "linestyle": AGT_LINESTYLES[0],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "v[1]": {
        "color": "blue",
        "linestyle": AGT_LINESTYLES[1],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "μ_v[0]": {
        "color": "dodgerblue",
        "linestyle": AGT_LINESTYLES[0],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "μ_v[1]": {
        "color": "dodgerblue",
        "linestyle": AGT_LINESTYLES[1],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "μ_y[0]": {
        "color": "orange", 
        "linestyle": AGT_LINESTYLES[0],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "μ_y[1]": {
        "color": "orange", 
        "linestyle": AGT_LINESTYLES[1],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "μ_y[2]": {
        "color": "orange", 
        "linestyle": AGT_LINESTYLES[2],
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    },
    "F": {
        "color": "green",
        "linestyle": "-",
        "linewidth": AGT_WIDTH,
        "alpha": ALPHA
    }
}

#### time
tT = 10000e-6 ## 2000e-6 1 10000                    ## Time bins, last t
Δt = .1e-6 ## .00001 0.5                      ## Discretization interval
ts = np.arange(0, tT, Δt)     ## Time steps
T  = len(ts)                  ## Number of time steps

#### dimensionsionality of vectors
B = 1 ## number of dims for autonomous state vˣ/v
C = 2 ## number of dims for state xˣ/μ_x
D = 1 ## number of dims for observation y/μ_y
A = 1 ## number of dims for action a
M = 3 ## number of dims for embedding (generalized coordinates of motion)

#### noise
def white_noise(scale=1e-1, size=1):
    return np.random.normal(loc=0, scale=scale, size=size)
## white noise
σˣ2_x = 0.001
σˣ2_y = 0.001
_n_x = white_noise(scale=np.sqrt(σˣ2_x), size=(T,C))
_n_y = white_noise(scale=np.sqrt(σˣ2_y), size=(T,D))
## colored noise via AR(1) filters
_ω_x = np.zeros((T, C))
_ω_y = np.zeros((T, D))
α, β = 0.95, 0.10 ## process noise correlation
γ, δ = 0.90, 0.05 ## measurement noise correlation
## configure colored noise to be white again
## α, β = 0.00, 1.00 ## process noise correlation
## γ, δ = 0.00, 1.00 ## measurement noise correlation
for t in range(0, T - 1):
    _ω_x[t] = α*_ω_x[t-1] + β*_n_x[t]
    _ω_y[t] = γ*_ω_y[t-1] + δ*_n_y[t]

#### exogenous functions
def f_const(_t: np.ndarray, A: float = 0.0, **kwargs) -> np.ndarray:
    """
    Constant signal.
    _t : array of time points (any shape)
    A  : constant value
    returns: array with same shape as _t filled with A
    """
    _t_arr = np.asarray(_t, dtype=float)
    return A * np.ones_like(_t_arr, dtype=float)

def f_cos(_t: np.ndarray, A: float = 1.0, θ1: float = 0.0) -> np.ndarray:
    """
    Cosine signal A cos(θ1 t).
    _t : array of time points (any shape)
    A  : amplitude
    θ1 : angular frequency
    returns: array with same shape as _t
    """
    _t_arr = np.asarray(_t, dtype=float)
    return A * np.cos(θ1 * _t_arr)

def f_pulses(_t: np.ndarray, **kwargs) -> np.ndarray:
    """
    Sum of rectangular pulses defined in physical time.
    kwargs:
      'amps'   : list/array of amplitudes [A1, A2, ...]
      't_ons'  : list/array of start times [t_on1, t_on2, ...]
      't_offs' : list/array of end times   [t_off1, t_off2, ...]
    """
    _amps   = np.asarray(kwargs["amps"], dtype=float)
    _t_ons  = np.asarray(kwargs["t_ons"], dtype=float)
    _t_offs = np.asarray(kwargs["t_offs"], dtype=float)
    _t_arr = np.asarray(_t, dtype=float)
    _val = np.zeros_like(_t_arr, dtype=float)
    for A, t_on, t_off in zip(_amps, _t_ons, _t_offs):
        _val += A * (
            np.heaviside(_t_arr - t_on, 0.0)
            - np.heaviside(_t_arr - t_off, 0.0)
        )
    return _val

def f_lines(_t: np.ndarray, **kwargs) -> np.ndarray:
    """
    Piecewise‑linear signal defined by points (t_i, y_i).
    kwargs:
      'y_pts' : list/array of values      [y0, y1, ..., yN]
      't_pts' : list/array of time points [t0, t1, ..., tN]
    For t between t_i and t_{i+1}, returns the straight line
    interpolating between (t_i, y_i) and (t_{i+1}, y_{i+1}).
    For t < t0 returns y0, for t > tN returns yN.
    """
    _y_pts = np.asarray(kwargs["y_pts"], dtype=float)
    _t_pts = np.asarray(kwargs["t_pts"], dtype=float)
    if _t_pts.ndim != 1 or _y_pts.ndim != 1 or _t_pts.size != _y_pts.size:
        raise ValueError("t_pts and y_pts must be 1D arrays of the same length")
    _t_arr = np.asarray(_t, dtype=float)
    _val = np.interp(_t_arr, _t_pts, _y_pts)
    return _val

def f_gbump(_t: np.ndarray, A: float = 1.0, t0: float = 0.0, std: float = 1.0) -> np.ndarray:
    """
    Gaussian bump in time.
    _t  : array of time points (any shape)
    A   : peak amplitude of the bump
    t0  : center time of the bump
    std : standard deviation (controls bump width)
    returns: array with same shape as _t
    """
    _t_arr = np.asarray(_t, dtype=float)
    return A * np.exp(-0.5 * ((_t_arr - t0) / std) ** 2)

## Disable noise for debugging
# _ω_x[:] = 0
# _ω_y[:] = 0
# _ω_x[7], _ω_y[7]

#### envir
θˣ_x = {
    'C3': 1e-6, ## 5.0,
    'R2': 20, ## 20 10 2.0,
    'L6': 10e-3, ## 1.0,
    'R_a': 1e-3 ## for now
}

def f_E(_xˣ: np.ndarray,
        _vˣ: np.ndarray,
        _a:  np.ndarray,
        θˣ_x: dict = θˣ_x) -> np.ndarray:
    """
    Environment transition f_E with arrays everywhere.
    _xˣ : (2,)  state [q3, p6]
    _vˣ : (1,)  exogenous drive [v7]
    _a  : (1,)  action [a5]
    """
    C3 = float(θˣ_x['C3'])
    R2 = float(θˣ_x['R2'])
    L6 = float(θˣ_x['L6'])

    ## unpack state and drives (still as arrays, but we index once)
    q3, p6 = _xˣ
    v7 = _vˣ[0]
    a5 = _a[0]

    dq3 = (1.0 / L6) * p6 - v7 - a5
    dp6 = -(1.0 / C3) * q3 - (R2 / L6) * p6 + R2 * v7 + R2 * a5

    return np.array([dq3, dp6], dtype=float)   # (2,)

def g_E(_xˣ: np.ndarray,
        _vˣ: np.ndarray | None,
        θˣ_x: dict = θˣ_x) -> np.ndarray:
    """
    Environment output mapping g_E with arrays everywhere.
    _xˣ : (2,)  state [q3, p6]
    returns: (1,) observation (voltage across C3)
    """
    L6 = float(θˣ_x['L6'])
    p6 = _xˣ[1]
    y_L6 = (1.0 / L6) * p6
    return np.array([y_L6], dtype=float) ## (1,)

#### run
_a = np.zeros((T, A))
_xˣ = np.zeros((T, C))
_vˣ = np.zeros((T, B))
_y = np.zeros((T, D))

_a[0] = np.array([0]) ## aSf5
_xˣ[0] = np.array([10e-6, 0.]) ## xˣq3, xˣp6 ## (C,)
_vˣ[0] = np.array([0.]) ## vˣSf7
_y[0] = np.array([0.]) ## yVoltage

sources_env = {
    'a5': f_pulses, ##
    # 'a5': f_const, ##
    'v7': f_cos, ##
}
a5_theta_env = { ##
    'amps'  : [15,  10], ##
    't_ons' : [0.,   5e-3], ##
    't_offs': [5e-3, 1e10], ##
    # 'amps'  : [100,  -100,   100,    -100,  100,    -100,  100,  -100,   100,  -100],
    # 't_ons' : [0.,   1e-3,   2e-3,   3e-3,  4e-3,   5e-3,  6e-3,  7e-3,  8e-3, 9e-3],
    # 't_offs': [1e-3, 2e-3,   3e-3,   4e-3,  5e-3,   6e-3,  7e-3,  8e-3,  9e-3, 1e10],
}
# a5_theta_env = { ##
#     'A'  : 15, ##
# } 
v7_theta_env = { ##
    'A': 100e-3, ## 10.,    
    'θ1': 10e3, ## .003,
}

for t in range(0, T - 1):
    _vˣ[t+1]  = sources_env['v7'](ts[t], **v7_theta_env)
    ## vˣ[t+1] = sampled_vˣ ## vˣ[t+1] estimated from 1-step ahead historical data

    _a[t+1]  = sources_env['a5'](ts[t], **a5_theta_env) ## estimated by using current observation from a sensor

    #### NEXT STATE
    _ẋˣ = f_E(_xˣ[t], _vˣ[t], _a[t]) + _ω_x[t]
    _xˣ[t+1] = _xˣ[t] + _ẋˣ*Δt ## f is transition function; Euler
    
    #### OBSERVE
    _y[t+1] = g_E(_xˣ[t+1], _vˣ[t+1]) + _ω_y[t] ## g is generation function

ylabel_size = 12
ylabelx = -0.3
ylabely = 0.3
fig = plt.figure(figsize=(10, 8))
grid_rows = 9
gs = GridSpec(grid_rows, 1, figure=fig, height_ratios=[1, 1, 1, 1, 3, 1, 1, 1, 1])
## gs = GridSpec(grid_rows, 1, figure=fig, height_ratios=[1]*grid_rows)
ax = [fig.add_subplot(gs[i]) for i in range(grid_rows)]
tN = tT
for j, axis in enumerate(ax):
    # axis.grid(True) ##
    axis.grid(False) ##
    # axis.set_xlim(0, tT) ##
    axis.set_xlim(0, tN) ##
    locator = MaxNLocator(nbins=5)
    axis.xaxis.set_major_locator(locator)
    axis.spines['top'].set_visible(False); axis.spines['right'].set_visible(False)
    if j < grid_rows - 1:
        axis.set_xticklabels([]) ## hide labels but keep tick positions
ax[0].set_title(r'Environment', fontweight='bold',fontsize=14)

i = 0
ax[i].plot(ts, _a, **styles['a[0]'], label=r'$\mathbf{a}[0]=a_5(t)$: action current Sf_5')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='A'))
ax[i].legend(loc='upper right')

i = 1
ax[i].plot(ts, _vˣ, **styles['vˣ[0]'], label=r'$\mathbf{v}^*[0]=v_7(t)$: true Sf_7')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='A'))
ax[i].legend(loc='upper right')

i = 2
ax[i].plot(ts, _xˣ[:,0], **styles['xˣ[0]'], label=r'$\mathbf{x}^*[0]=q_3(t)$: true charge')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='C'))
ax[i].legend(loc='upper right')

i = 3
ax[i].plot(ts, _xˣ[:,1], **styles['xˣ[1]'], label=r'$\mathbf{x}^*[1]=p_6(t)$: true flux')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='Wb'))
ax[i].legend(loc='upper right')

i = 4
ax[i].plot(ts, _y, **styles['y[0]'], label=r'$\mathbf{y}[0]=y(t)$: obs $e_{C_3}$')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].axhline(y=15.0, color='blue', linestyle='--', linewidth=1)  # horizontal line at y=0
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='A'))
ax[i].legend(loc='upper right')

i = 5
ax[i].plot(ts, _n_x[:,0], color='black', linestyle=':', label=r'$n_x$')
## ax[i].plot(ts, _n_x[:,1], color='grey', linestyle=':', label=r'$n_x$')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit=''))
ax[i].legend(loc='upper right')

i = 6
ax[i].plot(ts, _ω_x[:,0], color='black', linestyle=':', label=r'$ω_x$')
## ax[i].plot(ts, _ω_x[:,1], color='grey', linestyle=':', label=r'$ω_x$')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit=''))
ax[i].legend(loc='upper right')

i = 7
ax[i].plot(ts, _n_y[:,0], color='orange', linestyle=':', label=r'$n_y$')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit=''))
ax[i].legend(loc='upper right')

i = 8
ax[i].plot(ts, _ω_y[:,0], color='orange', linestyle=':', label=r'$ω_y$')
ax[i].yaxis.set_label_coords(ylabelx, ylabely)
ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit=''))
ax[i].legend(loc='upper right')

## bottom
bottom = ax[-1]
bottom.set_xlim(0, tN)                     # 0 to 5000 µs (0.005 s)
bottom.xaxis.set_major_locator(MaxNLocator(nbins=5))
bottom.xaxis.set_major_formatter(ticker.EngFormatter(unit='s'))
bottom.set_xlabel(r'$\mathrm{Time}\ t\ [\mathrm{seconds}]$', fontweight='bold', fontsize=12)

## plt.tight_layout()
plt.subplots_adjust(hspace=0.3) ## Adjust this value as needed
plt.show()
## fig.savefig("./ElectricSystem-envir", bbox_inches="tight", dpi=300)

4.4 Uncertainty Model

As mentioned before, there is uncertainty in the (colored) process and measurement noise. We will also define models for the time behavior of the disturbance current \(v_7(t)\). We will experiment with step changes as well as cycles based on the cosine function.

4.5 AIF Agent / Generative Model

Quadratic/Laplace approximation

For some simplification in what follows, we make use of a quadratic approximation involving the application of a second order Taylor series expansion centered on the mean of the variational density \(q(x)\). Another assumption is that \(q(x)\) is approximately Gaussian. This means we are effectively approximating the shape of \(q(x)\) with a quadratic function at the second order of the Taylor series expansion. When \(q(x)\) is sharply peaked or has high precision (low variance), this approximation is justified. This finally leads to the convenient property that the variance is a function of the mean:

\(p(x ∣ y) ≈ q_{\mu_x,\sigma^2_x}(x) = \mathcal N(x; \mu_x, \sigma^2_x) ≈ q_{\mu_x}(x) = \mathcal N(x; \mu_x,\sigma^2_x(\mu_x))\)

The calculation of the variational density over the hidden state now only depends on the value of the mean. Effectively this is the same as assuming that \(q(x)\) is a Dirac \(\delta\) function, but with the added advantage that we also get an approximation of the variance if we wish to find it.

Since the variational density is Gaussian and we plan to evaluate the function at the posterior mode, we substitute all occurences of \(x\) in the generative model with \(\mu_x\):

\[\begin{align} \dot{\mathbf{x}} &= f_{\cal M}(\boldsymbol{\mu}_x, \mathbf{v}; \theta_x) + \boldsymbol{\omega}_x \\ \dot{q}_3(t) &= - \frac{1}{L_6} p_{6}(t) - v_7(t) - a_5(t) + \omega_{x0}(t) \\ \dot{p}_6(t) &= - \frac{1}{C_3} q_{3}(t) - \frac{R_2}{L_6} p_{6}(t) + R_2 v_7(t) + R_2 a_5(t) + \omega_{x1}(t) \\ \end{align}\]

\[\begin{align} \mathbf{y} &= g_{\cal M}(\boldsymbol{\mu}_x, \mathbf{v}; \theta_y) + \mathbf{\omega}_y \\ y(t) &= \frac{1}{C_3} q_{3}(t) + \omega_y(t) \end{align}\]

4.5.0 AIF Agent

## §4.5.0 — shared scenario, result builder, and plotting (used by §4.5.1 & §4.5.2)
setpoint = 50e-3  ## L6 current scale (g_M = p6/L6)
x0_agt = np.array([10e-6, 0.])
v7_theta_agt = {
    'amps'  : [100e-3, 50e-3],  ## v7 exogenous current [A]
    't_ons' : [0.,      5e-3],   ## within tT = 10 ms
    't_offs': [5e-3,    1e10],
}
v7_scale_agt = 1.0  ## multiply v7 amplitude (try 0.2–0.5 if control is hard to tune)
sources_agt = {'v7': f_pulses}

# identical exogenous drive for both agent sections
_vˣ_agt_pre = np.zeros((T, B))
_vˣ_agt_pre[:, 0] = v7_scale_agt * sources_agt['v7'](ts, **v7_theta_agt)

# Agent demo: shared measurement noise (GC section smooths this via generalized y)
AGT_δ = 0.15  ## measurement noise; ↑ = noisier woutGC μ_y (try 0.10–0.20)
AGT_γ_y_gc = 350.0  ## GC smoothness; ↓ = smoother withGC μ_y (τ = 2/γ ≈ 5.7 ms)
AGT_λ_y = 50.0  ## sensory precision when AGT_δ>0 (noisy μ_y in §4.5.1)
AGT_λ_v = 80.0  ## weaker setpoint prior when noise on (μ_y follows y more)
if AGT_δ > 0:
    _ω_y[:] = 0.0
    for t in range(1, T):
        _ω_y[t] = γ * _ω_y[t-1] + AGT_δ * _n_y[t]


def build_agent_result(
    *, with_gc=False, M_gc=0,
    _a, _xˣ, _vˣ, _y, _v, _μ_x, _μ_y, F,
    _ϵ_x=None, _ϵ_y=None, _dF___dμₓ=None, _dF___da=None,
    _ϵ̃_x=None, _ϵ̃_y=None, _dF___dμ̃ₓ=None,
):
    cols = {
        't': ts,
        '_a_0': _a[:, 0],
        '_xˣ_0': _xˣ[:, 0], '_xˣ_1': _xˣ[:, 1],
        '_vˣ_0': _vˣ[:, 0],
        '_y_0': _y[:, 0],
        '_v_0': _v[:, 0],
        '_μ_x_0': _μ_x[:, 0], '_μ_x_1': _μ_x[:, 1],
        '_μ_y_0': _μ_y[:, 0],
        'F': F,
    }
    if with_gc:
        cols['_ϵ̃_x_0_0'] = _ϵ̃_x[:, 0, 0]
        cols['_ϵ̃_y_0_0'] = _ϵ̃_y[:, 0, 0]
        cols['_dF___dμ̃ₓ_0_0'] = _dF___dμ̃ₓ[:, 0, 0]
        if _dF___da is not None:
            cols['_dF___da_0'] = _dF___da[:, 0]
    else:
        cols['_ϵ_x_0'] = _ϵ_x[:, 0]
        cols['_ϵ_y_0'] = _ϵ_y[:, 0]
        cols['_dF___dμₓ_0'] = _dF___dμₓ[:, 0]
        cols['_dF___da_0'] = _dF___da[:, 0]
    return pd.DataFrame(cols)


LEGEND_FONTSIZE = 11


def _apply_legend(axis, *, outside=False, loc='upper left', ncol=1):
    handles, labels = axis.get_legend_handles_labels()
    if not handles:
        return
    n = len(handles)
    common = dict(
        fontsize=LEGEND_FONTSIZE,
        labelcolor='black',
        frameon=False,
        borderpad=0.35,
        labelspacing=0.35,
        handlelength=1.8,
    )
    leg = None
    if outside or n > 3:
        leg = axis.legend(
            **common,
            loc='upper left',
            bbox_to_anchor=(1.01, 1.0),
            borderaxespad=0.0,
            ncol=1 if n <= 5 else 2,
        )
    else:
        leg = axis.legend(**common, loc=loc, ncol=ncol)
    for text in leg.get_texts():
        text.set_color('black')
        text.set_alpha(1.0)


def plot_agent_result(result_agt, *, title_suffix=''):
    ylabelx = -0.3
    ylabely = 0.3
    grid_rows = 8
    fig = plt.figure(figsize=(14, 18))
    gs = GridSpec(grid_rows, 1, figure=fig, height_ratios=[1] * grid_rows)
    ax = [fig.add_subplot(gs[i]) for i in range(grid_rows)]
    tN = tT
    with_gc = '_ϵ̃_x_0_0' in result_agt.columns

    for j, axis in enumerate(ax):
        axis.grid(False)
        axis.set_xlim(0, tN)
        axis.xaxis.set_major_locator(MaxNLocator(nbins=5))
        axis.spines['top'].set_visible(False)
        axis.spines['right'].set_visible(False)
        if j < grid_rows - 1:
            axis.set_xticklabels([])

    ax[0].set_title(f'Agent\n{title_suffix}', fontweight='bold', fontsize=14)

    i = 0
    ax[i].plot(result_agt['t'], result_agt['_a_0'], **styles['a[0]'],
               label=r'$\mathbf{a}[0]=a_5(t)$: action current')
    ax[i].yaxis.set_label_coords(ylabelx, ylabely)
    ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='A'))
    _apply_legend(ax[i], outside=True)

    i = 1
    ax[i].plot(result_agt['t'], result_agt['_vˣ_0'], **styles['vˣ[0]'],
               label=r'$\mathbf{v}^*[0]=v_7(t)$: true current')
    ax[i].yaxis.set_label_coords(ylabelx, ylabely)
    ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='A'))
    _apply_legend(ax[i], outside=True)

    i = 2
    ax[i].plot(result_agt['t'], result_agt['_xˣ_0'], **styles['xˣ[0]'],
               label=r'$\mathbf{x}^*[0]=q_3(t)$: true charge')
    ax[i].plot(result_agt['t'], result_agt['_μ_x_0'], **styles['μ_x[0]'],
               label=r'$\boldsymbol{μ}_x[0]$: exp charge')
    ax[i].yaxis.set_label_coords(ylabelx, ylabely)
    ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='C'))
    _apply_legend(ax[i], outside=True)

    i = 3
    ax[i].plot(result_agt['t'], result_agt['_xˣ_1'], **styles['xˣ[1]'],
               label=r'$\mathbf{x}^*[1]=p_6(t)$: true flux [Wb]')
    ax[i].plot(result_agt['t'], result_agt['_μ_x_1'], **styles['μ_x[1]'],
               label=r'$\boldsymbol{μ}_x[1]$: exp flux [Wb]')
    ax[i].yaxis.set_label_coords(ylabelx, ylabely)
    ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='Wb'))
    _apply_legend(ax[i], outside=True)

    i = 4
    _half = result_agt['t'] >= (tT / 2)
    _std_y = np.std(result_agt['_y_0'][_half]) * 1e3
    _std_my = np.std(result_agt['_μ_y_0'][_half]) * 1e3
    ax[i].plot(result_agt['t'], result_agt['_v_0'], **styles['v[0]'],
               label=r'$\mathbf{v}[0]=v(t)$: L6 setpoint')
    if with_gc:
        _sty_y = {**styles['y[0]'], 'alpha': 0.55}
        ax[i].plot(result_agt['t'], result_agt['_y_0'], **_sty_y,
                   label=r'$\mathbf{y}[0]=y(t)$: noisy obs')
        _sty_my = {**styles['μ_y[0]'], 'linewidth': 3.0}
        ax[i].plot(result_agt['t'], result_agt['_μ_y_0'], **_sty_my,
                   label=rf'$\mathbf{{\mu}}_y$ smoothed (std={_std_my:.1f} mA)')
    else:
        ax[i].plot(result_agt['t'], result_agt['_y_0'], **styles['y[0]'],
                   label=rf'$\mathbf{{y}}=\mathbf{{\mu}}_y$ raw (std={_std_my:.1f} mA)')
    ax[i].yaxis.set_label_coords(ylabelx, ylabely)
    ax[i].yaxis.set_major_formatter(ticker.EngFormatter(unit='A'))
    _apply_legend(ax[i], outside=True)
    ax[i].text(0.99, 0.05, rf'std($y$)={_std_y:.1f} mA',
               transform=ax[i].transAxes, ha='right', va='bottom', fontsize=9)

    i = 5
    if with_gc:
        ax[i].plot(result_agt['t'], result_agt['_ϵ̃_x_0_0'], color='black', linestyle='-',
                   label=r'$\tilde{ϵ}_{x00}$')
        ax[i].plot(result_agt['t'], result_agt['_ϵ̃_y_0_0'], color='orange', linestyle='-',
                   label=r'$\tilde{ϵ}_{y00}$')
    else:
        ax[i].plot(result_agt['t'], result_agt['_ϵ_x_0'], color='black', linestyle='-',
                   label=r'$ϵ_{x}$')
        ax[i].plot(result_agt['t'], result_agt['_ϵ_y_0'], color='orange', linestyle='-',
                   label=r'$ϵ_{y}$')
    ax[i].yaxis.set_label_coords(ylabelx, ylabely)
    _apply_legend(ax[i], outside=True)

    i = 6
    if with_gc:
        ax[i].plot(result_agt['t'], result_agt['_dF___dμ̃ₓ_0_0'], color='purple', linestyle='-',
                   label=r'$\frac{dF}{d\tilde{μ}_{x00}}$')
    else:
        ax[i].plot(result_agt['t'], result_agt['_dF___dμₓ_0'], color='purple', linestyle='-',
                   label=r'$\frac{dF}{dμ_{x0}}$')
    ax[i].plot(result_agt['t'], result_agt['_dF___da_0'], color='cyan', linestyle='-',
               label=r'$\frac{dF}{da_{0}}$')
    ax[i].yaxis.set_label_coords(ylabelx, ylabely)
    _apply_legend(ax[i], outside=True)

    i = 7
    ax[i].plot(result_agt['t'], result_agt['F'], **styles['F'],
               label=r'$\cal F$')
    ax[i].yaxis.set_label_coords(ylabelx, ylabely)
    _apply_legend(ax[i], outside=True)

    bottom = ax[-1]
    bottom.set_xlim(0, tN)
    bottom.xaxis.set_major_locator(MaxNLocator(nbins=5))
    bottom.xaxis.set_major_formatter(ticker.EngFormatter(unit='s'))
    bottom.set_xlabel(r'$\mathrm{Time}\ t\ [\mathrm{seconds}]$', fontweight='bold', fontsize=12)
    plt.subplots_adjust(hspace=0.3, right=0.68)
    plt.show()
    return fig

#### agent
def f_M(_μₓ: np.ndarray,
        _μᵥ: np.ndarray,
        θˣ_x: dict = θˣ_x) -> np.ndarray: ## transition function for Model/Agent
    """
    Agent transition model f_M with arrays everywhere.
    _μₓ : (2,) state belief [q3, p6]
    _μᵥ : (1,) exogenous / setpoint drive [v7]
    """
    C3 = float(θˣ_x['C3'])
    R2 = float(θˣ_x['R2'])
    L6 = float(θˣ_x['L6'])

    q3, p6 = _μₓ        ## both arrays entries
    v7 = _μᵥ[0]         ## first (and only) component

    dq3 = (1.0 / L6) * p6 - v7
    dp6 = -(1.0 / C3) * q3 - (R2 / L6) * p6 + R2 * v7

    return np.array([dq3, dp6], dtype=float)   # (2,)

def g_M(_μₓ: np.ndarray,
        _μᵥ: np.ndarray | None,
        θˣ_x: dict = θˣ_x) -> np.ndarray: ## generation function for Model/Agent
    """
    Agent observation model g_M with arrays everywhere.
    _μₓ : (2,) state belief [q3, p6]
    returns: (1,) predicted observation (voltage across C3)
    """
    L6 = float(θˣ_x['L6'])
    p6 = _μₓ[1]
    y_L6 = (1.0 / L6) * p6
    return np.array([y_L6], dtype=float) ## (1,)

def df_M___dμₓ_mat(_μₓ: np.ndarray,
               _μᵥ: np.ndarray,
               θˣ_x: dict = θˣ_x) -> np.ndarray:
    """
    Jacobian of f_M with respect to μₓ = [q3, p6].
    _μₓ : (2,) state belief [q3, p6]  (value unused here but kept for signature)
    _μᵥ : (1,) exogenous drive [v7]   (unused for this Jacobian)
    returns: (2, 2) array
    """
    C3 = float(θˣ_x['C3'])
    R2 = float(θˣ_x['R2'])
    L6 = float(θˣ_x['L6'])
    J = np.array([
        [0.0,        1.0 / L6],
        [-1.0 / C3, -R2 / L6],
    ], dtype=float)
    return J ## shape (2, 2)

def df_M___dμₓ_vec(_μₓ: np.ndarray,
                    _μᵥ: np.ndarray,
                    θˣ_x: dict = θˣ_x) -> np.ndarray:
    """
    Diagonal entries of ∂f_M/∂μₓ, used with diagonal precisions.
    Returns: (C,)
    """
    R2 = float(θˣ_x['R2'])
    L6 = float(θˣ_x['L6'])
    return np.array([0.0, -R2 / L6], dtype=float)

def dg_M___dμₓ_vec(_μₓ: np.ndarray,
                    _μᵥ: np.ndarray | None,
                    θˣ_x: dict = θˣ_x) -> np.ndarray:
    """
    Diagonal entries of ∂g_M/∂μₓ, used with diagonal precisions.
    For y = (1/C3) q3, treat only the first component as observed.
    Returns: (C,)
    """
    L6 = float(θˣ_x['L6'])
    return np.array([0.0, 1.0 / L6], dtype=float)

def dg_M___dμₓ_mat(_μₓ: np.ndarray,
                    _μᵥ: np.ndarray | None,
                    θˣ_x: dict = θˣ_x) -> np.ndarray:
    """
    Full Jacobian ∂g_M/∂μₓ for the agent observation model.
    g_M:
      y = (1/C3) * q3          (scalar observation)
    _μₓ : (2,) state belief [q3, p6]
    returns: (1, 2) array  [[∂y/∂q3, ∂y/∂p6]]
    """
    L6 = float(θˣ_x['L6'])
    J = np.array([[0.0, 1.0 / L6]], dtype=float)  # (1, 2)
    return J

def h_M(_μₓ: np.ndarray,
        _μᵥ: np.ndarray,
        θˣ_x: dict = θˣ_x) -> np.ndarray:
    """
    Preference mapping h_M with arrays everywhere.
    Here: setpoint on observation → h_M(μₓ) = g_M(μₓ).
    _μₓ : (C,) state belief
    _μᵥ : (B,) exogenous / setpoint drive (passed through for symmetry)
    returns: (D,) preferred observation
    """
    ## setpoint on state: h_M(μₓ) = μₓ
    ## return μₓ ## shape (C,)

    ## setpoint on observation: h_M(μₓ) = g_M(μₓ)
    return g_M(_μₓ, _μᵥ) ## shape (D,)

def dh_M___dμₓ(_μₓ: np.ndarray,
               _μᵥ: np.ndarray,
               θˣ_x: dict = θˣ_x) -> np.ndarray:
    """
    Jacobian of h_M with respect to μₓ.
    Since h_M = g_M, this equals dg_M___dμₓ.
    """
    return dg_M___dμₓ_mat(_μₓ, _μᵥ, θˣ_x)   # shape (D, C)

## def dg_M___da(_μₓ: np.ndarray,
    #           _a:  np.ndarray,
    #           C:   np.ndarray) -> np.ndarray: ## direct action on observation (rare)
    # """
    # Jacobian of g_M with respect to action a (arrays everywhere).
    # _μₓ : (C_x,) state belief (unused here, kept for generality)
    # _a  : (A,)   action vector
    # C   : (D, A) constant mapping from action to observation
    # returns: (D, A) = ∂g_M / ∂a
    # """
    # C = np.asarray(C, dtype=float)
    # return C

## A. Calculate VFE gradient wrt μₓ (Eq6.7b ---> Eq6.14) [[PEJ]]
def _dF___dμₓ_fn_vec(
    _μ_x, _μ_v, _ϵ_y, _ϵ_x, _ϵ_v, _λ_y, _λ_x, _λ_v
):
    J_obsr_diag = dg_M___dμₓ_vec(_μ_x, _μ_v)   # (C,)
    J_proc_diag = df_M___dμₓ_vec(_μ_x, _μ_v)   # (C,)
    J_pref_diag = dh_M___dμₓ(_μ_x, _μ_v).ravel()  # (C,) from (1,2)

    ## hard cap on errors to avoid overflow
    max_err = 1e3
    ex = np.clip(_ϵ_x, -max_err, max_err)
    ev = np.clip(_ϵ_v, -max_err, max_err)
    ey = np.clip(_ϵ_y, -max_err, max_err)

    J_obsr = dg_M___dμₓ_mat(_μ_x, _μ_v)          ## (D, C)
    J_pref = dh_M___dμₓ(_μ_x, _μ_v)              ## (D, C)
    grad_like = (J_obsr.T @ (_λ_y * ey)).ravel()  ## (C,)
    grad_dyn  = _λ_x * ex * J_proc_diag
    grad_pref = (J_pref.T @ (_λ_v * ev)).ravel()  ## (C,)
    return grad_like + grad_dyn + grad_pref

def _dF___dμₓ_fn_mat(
    _μ_x: np.ndarray,    ## (C,)
    _μ_v: np.ndarray,    ## (B,) or (C,)
    _ϵ_y: np.ndarray,    ## (D,)
    _ϵ_x: np.ndarray,    ## (C,)
    _Π_y: np.ndarray,    ## (D, D)
    _Π_x: np.ndarray,    ## (C, C)
) -> np.ndarray:
    """
    Full gradient dF/dμₓ using full Jacobians and precision matrices.
    Returns: (C,)
    """
    J_obsr = dg_M___dμₓ_mat(_μ_x, _μ_v)   ## (D, C)
    J_proc = df_M___dμₓ_mat(_μ_x, _μ_v)   ## (C, C)

    term_y = J_obsr.T @ (_Π_y @ _ϵ_y)      ## (C,)
    term_x = J_proc.T @ (_Π_x @ _ϵ_x)      ## (C,)
    return term_y + term_x

### A. Calculate VFE gradient wrt a (Eq7.7 ---> )
def _dF___da_fn_vec(_ϵ_y: np.ndarray, _λ_y: np.ndarray,
                    _ϵ_v: np.ndarray, _λ_v: np.ndarray,
                    _xˣ: np.ndarray, _a: np.ndarray) -> np.ndarray:
    """
    dF/da using dy/da and both sensory and preference channels.
    """
    ## scalar dimensions (integers!)
    D = _ϵ_y.shape[0]
    A = _a.shape[0]

    max_err = 1e3
    ey = np.clip(_ϵ_y, -max_err, max_err)
    ev = np.clip(_ϵ_v, -max_err, max_err)

    DyDa = dy___da(_xˣ=_xˣ, _a=_a, θˣ_x=θˣ_x, D=D, A=A)

    lam_y = _λ_y * ey
    lam_v = _λ_v * ev
    lam_tot = lam_y + lam_v

    return DyDa.T @ lam_tot

def _dF___da_fn_mat(_ϵ_y: np.ndarray,
                    _Π_y: np.ndarray,
                    _xˣ: np.ndarray,
                    _a:  np.ndarray) -> np.ndarray:
    """
    dF/da using full dy/da matrix and full observation precision Π_y.
    _ϵ_y : (D,)      sensory prediction error
    _Π_y : (D, D)    observation precision matrix
    _xˣ  : (C,)      current state (for dy/da, if needed)
    _a   : (A,)      current action
    returns
    -------
    np.ndarray of shape (A,) = gradient dF/da
    """
    D = _ϵ_y.shape[0]
    A = _a.shape[0]
    ## dy/da Jacobian: (D, A)
    DyDa = dy___da(_xˣ=_xˣ, _a=_a, D=D, A=A)      ## (D, A)
    ## effective force term λ_y = Π_y ε_y
    lam_y = _Π_y @ _ϵ_y                           ## (D,)
    ## dF/da = (dy/da)^T Π_y ε_y
    return DyDa.T @ lam_y                         ## (A,)

## _λ_x = np.full(C, 0.1)   # (2,)
## _λ_y = np.full(D, 15.0)  # (2,)
def F_fn_vec(
    _ϵ_y_t: np.ndarray,               ## (D,)  = (2,)
    _ϵ_x_t: np.ndarray,               ## (C,)  = (2,)
    _λ_y:    np.ndarray,              ## (D,)  diagonal precisions
    _λ_x:    np.ndarray,              ## (C,)
    _ϵ_v_t: np.ndarray | None = None, ## (B,) optional
    _λ_v:   np.ndarray | None = None, ### (B,) optional
) -> float:
    """
    Scalar variational free energy at one time t with diagonal precisions.
    For C = 2, D = 2:
      _ϵ_y_t : shape (2,) sensory prediction errors
      _ϵ_x_t : shape (2,) state prediction errors
      _λ_y   : shape (2,) observation precisions (diagonal entries)
      _λ_x   : shape (2,) state precisions (diagonal entries)
    Optional:
      _ϵ_v_t, _λ_v : same convention for a setpoint error channel.
    All ε's and λ's are 1‑D arrays matching their variable dimensions.
    """
    max_err = 1e3
    ey = np.clip(_ϵ_y_t, -max_err, max_err)
    ex = np.clip(_ϵ_x_t, -max_err, max_err)
    term_y = (_λ_y * ey**2).sum() - np.log(_λ_y).sum()
    term_x = (_λ_x * ex**2).sum() - np.log(_λ_x).sum()
    term_v = 0.0
    if _ϵ_v_t is not None and _λ_v is not None:
        ev = np.clip(_ϵ_v_t, -max_err, max_err)
        term_v = (_λ_v * ev**2).sum() - np.log(_λ_v).sum()
    return 0.5 * (term_y + term_x + term_v)


## The “full” version uses full precision matrices rather than diagonal vectors, 
##  and optional full covariance on the setpoint term.
## Use this when: 
##  _Π_y and _Π_x are full DxD and CxC matrices (with off-diagonal coupling)
##     _ϵ_y_t and _ϵ_x_t are 1-D vectors of length D and C respectively
def F_fn_mat(
    _ϵ_y_t: np.ndarray,               ## (D,)
    _ϵ_x_t: np.ndarray,               ## (C,)
    _Π_y:    np.ndarray,              ## (D, D)
    _Π_x:    np.ndarray,              ## (C, C)
    _ϵ_v_t: np.ndarray | None = None, ## (B,)
    _Π_v:   np.ndarray | None = None, ## (B, B)
) -> float:
    """
    Scalar VFE at one time t with full precision matrices.
    F(t) = 0.5 * [
        ε_y^T Π_y ε_y - log det Π_y
      + ε_x^T Π_x ε_x - log det Π_x
      (+ ε_v^T Π_v ε_v - log det Π_v, if provided)
    ]
    """
    ## observation term
    term_y_quad = float(_ϵ_y_t.T @ (_Π_y @ _ϵ_y_t))
    sign_y, logdet_y = np.linalg.slogdet(_Π_y)
    term_y = term_y_quad - logdet_y

    ## state term
    term_x_quad = float(_ϵ_x_t.T @ (_Π_x @ _ϵ_x_t))
    sign_x, logdet_x = np.linalg.slogdet(_Π_x)
    term_x = term_x_quad - logdet_x

    term_v = 0.0
    if _ϵ_v_t is not None and _Π_v is not None:
        if np.any(_Π_v != 0):
            term_v_quad = float(_ϵ_v_t.T @ (_Π_v @ _ϵ_v_t))
            sign_v, logdet_v = np.linalg.slogdet(_Π_v)
            term_v = term_v_quad - logdet_v
    return 0.5 * (term_y + term_x + term_v)

"""
You don’t need a separate “full” version for this one. As long as dy___da always returns 
the full Jacobian matrix of shape (D,A) (which the unified version does), that is the 
full version; you can just adjust its internal gain or structure as needed, without 
defining a second function.
Usage examples:
    Simple constant (-1)
        DyDa = dy___da(_xˣ=x, _a=a)  # base=-1.0 by default
    Physical gain
        ## gives (D, A) matrix with entries Δt / (C3 * R_a)
        DyDa = dy___da(_xˣ=x,
               _a=a,
               Δt=Δt,
               θˣ_x=θˣ_x,
               base=1.0,
               use_physical_gain=True)
"""
def dy___da(_xˣ: np.ndarray,
           _a:  np.ndarray,
           θˣ_x: dict,
           D: int | None = None,
           A: int | None = None) -> np.ndarray:
    if D is None:
        D = 1          ## one observation (voltage)
    if A is None:
        A = _a.shape[0]

    gain = 1.0 ## 0.01 1.0   # dimensionless phenomenological sensitivity
    return gain * np.ones((D, A), dtype=float)

4.5.1 GF without generalized coordinates

#### run
_a = np.zeros((T, A)) ## action vector
_xˣ = np.zeros((T, C)) ## true state vector
_vˣ = np.zeros((T, B)) ## exogenous force vector
_y = np.zeros((T, D)) ## observation vector

_μ_x = np.zeros((T, C)) ## Expectation (belief) about x [C]

## _v = np.ones((T, C)) * [setpoint] ## setpoints
_v = np.ones((T, D)) * [setpoint] ## setpoints
## _μ_v = np.zeros((T, C)) * [setpoint] ## setpoints; static, same as _v
_μ_v = np.ones((T, D)) * [setpoint] ## setpoints; static, same as _v

_μ_y = np.zeros((T, D)) ## expectation (belief) about y

_ϵ_x = np.zeros((T, C)) ## Model prediction error 
_ϵ_y = np.zeros((T, D)) ## Sensory prediction error
## _ϵ_v = np.zeros((T, C))
_ϵ_v = np.zeros((T, D))

F = np.zeros(T) ## Variational Free Energy (VFE)
_dF___dμₓ = np.zeros((T, C))
_dF___da = np.zeros((T, A))

_a[0] = np.array([0.]) ## a5 [A]
_xˣ[0] = x0_agt.copy() ## xˣq3 [C], xˣp6 [Wb]
_vˣ[0] = np.array([0.]) ## vˣv7 [A]
_y[0] = g_E(_xˣ[0], _vˣ[0])

_μ_x[0] = x0_agt.copy() ## μ_xq3, μ_xCharge
_μ_v[0] = np.array([setpoint]) ## μ_vSf7
_μ_y[0] = g_M(_μ_x[0], _μ_v[0]) ## μ_yVoltage

# ///////////////
## _ϵ_x[0] = _μ_x[0] - f_M(_μ_x[0], _μ_v[0]) ## initial model prediction error; pass in vˣ
f_pred0   = f_M(_μ_x[0], _μ_v[0])
μ_x_pred0 = _μ_x[0] + f_pred0 * Δt
_ϵ_x[0]   = _μ_x[0] - μ_x_pred0
# \\\
_ϵ_v[0] = h_M(_μ_x[0], _μ_v[0]) - _v[0]
_ϵ_y[0] = _y[0] - g_M(_μ_x[0], _μ_v[0]) ## initial sensory prediction error

## state precision / prior precision
## _λ_x = np.array([10.])
## _λ_x = np.array([0.2, 0.2])
## _Π_x = np.diag([.05] * C) ## shape (C,C)
## _Π_x = np.diag([10.0, 10.0]) ## shape (C,C)
_λ_x = np.full(C, .01) ## shape (C,)
## setpoint precision
## _λ_v = np.array([50.0])
## _Π_v = np.diag([150.0] * B) ## shape (B,B)
## _Π_v = np.diag([50.0, 50.0]) ## shape (B,B)
_λ_v = np.full(B, AGT_λ_v if AGT_δ > 0 else 8e3) ## weaker when demo noise on
## observation precision / likelihood precision
## _λ_y = [10.0] 
## _Π_y = np.diag([1.0]) ## shape (D,D)
## _Π_y = np.diag([1.0] * D) ## shape (D,D)
## _Π_y = np.diag([10.0, 10.0]) ## shape (D,D)
_λ_y = np.full(D, AGT_λ_y if AGT_δ > 0 else .01) ## higher when demo noise on
## once the setpoint is well tracked the prior term balances the tiny residual error 
##   and a5(t) converges to a finite value instead of ramping:
λ_a = 1e3 ## NEW # tune this

action_on = 0 / Δt ## Time step to begin active inference
κₓ = 1e-2 if AGT_δ > 0 else 1e-3 ## 0.05             ## learning rate (perception)
κₐ = 1e-2 ## 100e6 100e-3 10e-3 1e-3 .5e-3 .2e-3 .1e-3 0.010            ## learning rate (action)

## F[0] = F_fn_mat(_ϵ_y[0], _ϵ_x[0], _Π_y, _Π_x)
F[0] = F_fn_vec(_ϵ_y[0], _ϵ_x[0], _λ_y, _λ_x,    
                _ϵ_v_t=_ϵ_v[0], _λ_v=_λ_v)
## F[0] = 500_000 ## plot looks better
_vˣ[:, 0] = _vˣ_agt_pre[:, 0]  ## shared exogenous v7

for t in range(T - 1):
    _vˣ[t+1, 0] = _vˣ_agt_pre[t+1, 0]  ## shared exogenous v7

    #### NEXT STATE
    ## _ẋˣ = f_E(_xˣ[t], _vˣ[t], _a[t]) + _ω_x[t]
    ## _xˣ[t+1] = _xˣ[t] + _ẋˣ*Δt ## f is transition function; Euler
    ## Runge–Kutta 4
    k1 = f_E(_xˣ[t],           _vˣ[t], _a[t])
    k2 = f_E(_xˣ[t] + 0.5*Δt*k1, _vˣ[t], _a[t])
    k3 = f_E(_xˣ[t] + 0.5*Δt*k2, _vˣ[t], _a[t])
    k4 = f_E(_xˣ[t] +     Δt*k3, _vˣ[t], _a[t])
    _ẋˣ = (k1 + 2*k2 + 2*k3 + k4) / 6.0 + _ω_x[t]
    _xˣ[t+1] = _xˣ[t] + _ẋˣ * Δt

    #### OBSERVE
    _y[t+1] = g_E(_xˣ[t+1], _vˣ[t+1]) + _ω_y[t] ## g is generation function

    #### INFER [[the 'G'radient 'F'low 'U'pdates the 'ϵ' & 'F']]
    ### A. Calculate VFE gradient wrt μₓ (Eq6.7b ---> Eq6.14) [[PEJ]]
    ## _dF___dμₓ[t] = _dF___dμₓ_fn_mat(_μ_x[t], _μ_v[t], _ϵ_y[t], _ϵ_x[t], _Π_y, _Π_x)
    _dF___dμₓ[t] = _dF___dμₓ_fn_vec(
        _μ_x[t], _μ_v[t],
        _ϵ_y[t], _ϵ_x[t], _ϵ_v[t],
        _λ_y, _λ_x, _λ_v)

    ### A. Calculate VFE gradient wrt a (Eq7.7 ---> )
    ## _dF___da[t] = _dF___da_fn_mat(_ϵ_y[t], _Π_y, _xˣ[t], _a[t])
    ## _dF___da[t] = _dF___da_fn_vec(_ϵ_y[t], _λ_y, _xˣ[t], _a[t])
    _dF___da[t] = _dF___da_fn_vec(
        _ϵ_y[t], _λ_y,
        _ϵ_v[t], _λ_v,
        _xˣ[t], _a[t])
    ## _dF___da[t] += λ_a * _a[t] ## prior pushing a -> 0

    ### B. Define hidden state flow (Eq6.8 ---> Eq6.15)
    _μ̇ₓ = -κₓ * _dF___dμₓ[t]
    
    ### B. Define action flows (Eq7.5 ---> Eq7.25c)
    _ȧ = -κₐ * _dF___da[t]

    ### C. Hidden state & inferred observation update (Eq6.8 ---> Eq6.15); Perception; Euler
    _μ_x[t+1] = _μ_x[t] + _μ̇ₓ*Δt ## Perceptual update (gradient descent)
    _μ_y[t+1] = g_M(_μ_x[t+1], _μ_v[t+1])

    if AGT_δ > 0:
        _μ_y[t+1] = _y[t+1]  ## order-0 only: μ_y follows raw noisy y

    ### C. Control state update (Eq7.6); Perception; Euler
    if t >= action_on:
        _a[t+1] = _a[t] + _ȧ*Δt ## (Eq7.6) ## Action update (gradient descent)
        ## _a[t+1] = np.clip(_a[t+1], -5., 5.) ## Clamp action to reasonable range
        ## _a[t+1] = 0.0 ## comment in to see what happens if no action/control !!!!!!!!!!!!

    ### D. Recalculate prediction errors (Eq6.6 ---> Eq6.11)
    ## _ϵ_x[t+1] = _μ_x[t+1] - f_M(_μₓ=_μ_x[t+1], _v=_μ_v[t+1])
    # /////////////////////////
    ## _ϵ_x[t+1] = _μ_x[t+1] - f_M(_μ_x[t+1], _μ_v[t+1]) ## # shape (C,); pass in vˣ
    f_pred   = f_M(_μ_x[t], _μ_v[t])
    μ_x_pred = _μ_x[t] + f_pred * Δt
    _ϵ_x[t+1] = _μ_x[t+1] - μ_x_pred
    # \\\\\\\\
    _ϵ_y[t+1] = _y[t+1] - g_M(_μ_x[t+1], _μ_v[t+1]) ## shape (D,)
    ## _ϵ_v[t+1] = h_M(_μₓ=_μ_x[t+1]) - _μ_v[t+1] ## shape (C,)
    _ϵ_v[t+1] = h_M(_μ_x[t+1], _μ_v[t+1]) - _v[t+1] ## (D,)
    
    ### E. Recalculate VFE (Eq6.7a --->  Eq6.13)
    ## F[t+1] = F_fn_mat(_ϵ_y[t+1], _ϵ_x[t+1], _Π_y, _Π_x)
    F[t+1] = F_fn_vec(
        _ϵ_y[t+1], _ϵ_x[t+1],
        _λ_y, _λ_x,
        _ϵ_v_t=_ϵ_v[t+1],
        _λ_v=_λ_v
    )
    ## F[t+1] = F_fn(_ϵ_y[t+1], _ϵ_x[t+1], _λ_y, _λ_x,  _ϵ_v_t=_ϵ_v[t+1], _λ_v=_λ_v) ## with setpoints

result_agt_wout = build_agent_result(
    with_gc=False,
    _a=_a, _xˣ=_xˣ, _vˣ=_vˣ, _y=_y, _v=_v, _μ_x=_μ_x, _μ_y=_μ_y, F=F,
    _ϵ_x=_ϵ_x, _ϵ_y=_ϵ_y, _dF___dμₓ=_dF___dμₓ, _dF___da=_dF___da,
)

# Optional — saves AGF traces for P3 overlay comparison
from pathlib import Path
Path("results").mkdir(exist_ok=True)
result_agt_wout.to_pickle("results/result_agt_wout.pkl")

fig = plot_agent_result(result_agt_wout, title_suffix=r'(without generalized coordinates)')
## fig.savefig('./ElectricSystem-agent-woutGC-cos', bbox_inches='tight', dpi=300)

##
# from IPython.display import Image, display
# display(Image(filename="ElectricSystem-agent-woutGC-step.png"))

##
# from IPython.display import Image, display
# display(Image(filename="ElectricSystem-agent-woutGC-ramp.png"))

##
# from IPython.display import Image, display
# display(Image(filename="ElectricSystem-agent-woutGC-cos.png"))

4.5.2 GF with generalized coordinates

Same plant, drive, and noise (AGT_δ) as §4.5.1. Tunables in §4.5.0: AGT_δ (noise), AGT_γ_y_gc (smoothness, ↓ = smoother).

	§4.5.1 woutGC	§4.5.2 withGC
`μ_y`	`μ_y ← y` (fully noisy)	GC low-pass, `τ = 2/AGT_γ_y_gc`
action	reacts to raw `y`	reacts to smoothed `μ_y`

Panel 4: woutGC one jittery trace; withGC blue noisy y vs thick smooth μ_y.

def create_temporal_precision_matrix(M: int, γ: float, σ2: float) -> np.array:
    """ Creates a generalized precision matrix 
    Based on Matlab code from Hijne 2020, pp. 20-21.
    Creates a temporal covariance matrix for p+1 orders using a roughness of gamma for
    the amount of smoothening. The variance is then used to construct the generalized
    precision matrix.
    Args:
        M (int): The embedding order
        γ (float): Roughness parameter
        σ2 (float): Variance for x or y (i.e. σ2 * I)
    Returns:
        _np.array: The generalized precision matrix [M+1, M+1]
    """
    ## Order of the required autocorrelation derivatives
    k = np.arange(0, 2*M+1, 2)
    s = np.sqrt(2 / γ)
    ρ = np.cumprod(1-k) / (np.sqrt(2) * s)**k
    ρ = np.insert(arr=ρ, obj=np.arange(1,len(ρ), 1), values=0)
    ## Initialize temporal covariance matrix
    S = np.zeros((M+1, M+1))
    ## Loop over all rows of embedding order and population
    for r in range(0, M+1):
        S[r] = ρ[r:r+M+1]
        ρ = -ρ   # Inverse rho to alternate minus signs
    ## Compute generalized precision matrix
    Π̃ = np.linalg.inv(np.kron(S, σ2))

    ##. replace zeros and negative values with a small positive number instead of zero; 
    ##.   for np.log() on the elements later
    for i in range(Π̃.shape[0]):
        for j in range(Π̃.shape[1]):
            if Π̃[i,j] <= 0:
                Π̃[i,j] = 1e-10 ##. small positive number instead of zero; for np.log()
    return Π̃

def D_operator(M):
    return np.insert(np.insert(np.eye(M), 0, 0, axis=1), M, 0, axis=0)

## Works for C=1 (scalar state) and any C>1 without changes.
## Avoids any fixed indices like [_\mu_v[0], _\mu_v[3]]
## Lets you later change the initialization (e.g. mix \mu_x and \mu_v) by only modifying
##   the two lines that set _\mu_x_embedding[0, :] and [1, :], without touching the shape
###   logic
def embed_μ_x(_μ_x, _μ_v, M, C):
    """
    Embed state belief μ_x into generalized coordinates up to order M.
    _μ_x : array_like, shape (C,)
        Current state belief.
    _μ_v : array_like, shape (C,) or (D,) depending on your convention.
           Here we only use the components you choose to initialise with.
    M    : int
        Highest generalized order (0..M).
    C    : int
        State dimension.
    """
    _μ_x = np.asarray(_μ_x).reshape(C,)       ## ensure shape (C,)
    _μ_v = np.asarray(_μ_v).reshape(-1,)      ## 1‑D

    ## Process model and Jacobian
    _f  = f_M(_μ_x, _μ_v)                     ## shape (C,)
    # _fp = df_M___dμₓ(_μ_x, _μ_v)              ## shape (C, C)
    _fp = df_M___dμₓ_mat(_μ_x, _μ_v)              ## shape (C, C)

    ## Allocate embedding: rows = generalized order, cols = state components
    _μ_x_embedding = np.zeros((M + 1, C))

    ## 0th generalized coordinate: current state belief
    _μ_x_embedding[0, :] = _μ_x

    ## 1st generalized coordinate: state velocity
    _μ_x_embedding[1, :] = _f

    ## Higher‑order generalized coordinates: successive applications of Jacobian
    for i in range(2, M + 1):
        _μ_x_embedding[i, :] = _fp @ _μ_x_embedding[i - 1, :]

    return _μ_x_embedding

def embed_y(_y, _x, _v, M, C, D):
    _f = f_M(_x, _v)
    # _fp = df_M___dμₓ(_x, _v)
    _fp = df_M___dμₓ_mat(_x, _v)
    _g = g_M(_x, _v)
    # _gp = dg_M___dμₓ(_x, _v)
    _gp = dg_M___dμₓ_mat(_x, _v)
    _x_embedding = np.zeros((M+1, C))
    _y_embedding = np.zeros((M+1, D))
    _x_embedding[0] = _f
    for i in range(1, M+1): ## propagate with Jacobian for higher-order terms
        _x_embedding[i] = _fp @ _x_embedding[i-1]
    _y_embedding[0] = _g
    for i in range(1, M+1):
        _y_embedding[i] = _gp @ _x_embedding[i-1]
    return _y_embedding

def embed_func(func, dfunc, _μ̃_x, _μ_v, dim):
    M = _μ̃_x.shape[0] - 1
    _func_vec = np.zeros((M+1, dim))
    _motion_vec = np.ones((M+1, dim))
    _func_vec[0] = 1
    _motion_vec[0] = 0
    ## _μ_v may be a generalized embedding (M+1, *) or a plain drive vector (B,)
    _μ_v_arr = np.asarray(_μ_v)
    _v_for_func = _μ_v_arr[0] if _μ_v_arr.ndim > 1 else _μ_v_arr.reshape(-1)
    _jac = dfunc(_μ̃_x[0], _v_for_func)
    ## Apply Jacobian to each embedded coordinate
    _motion_term = np.array([_jac @ _μ̃_x[m] for m in range(M+1)]) ## (M+1, C)
    _f̃ = _func_vec * func(_μ̃_x[0], _v_for_func) + _motion_term
    return _f̃

def embed_y_gc(_y_meas, _μ_x, _μ_v, M, D, Δt):
    """Generalized sensory embedding: measured y at order 0; model dy/dt at order 1+."""
    _emb = np.zeros((M + 1, D))
    _emb[0] = _y_meas
    if M >= 1:
        _emb[1] = dg_M___dμₓ_mat(_μ_x, _μ_v) @ f_M(_μ_x, _μ_v)
    return _emb



def update_μ̃_y_gc(_μ̃_y_t, _y_now, _μ_x, _μ_v, M, γ, κ, Δt, σ2=100.0):
    """GC step on current belief: order-0 from measurement; order-1 from model dy/dt."""
    Dop = D_operator(M)
    Π = create_temporal_precision_matrix(M, γ, σ2)
    D = _μ̃_y_t.shape[1]
    _ỹ = np.zeros((M + 1, D))
    _ỹ[0] = _y_now
    if M >= 1:
        _dy_model = dg_M___dμₓ_mat(_μ_x, _μ_v) @ f_M(_μ_x, _μ_v)
        _ỹ[1] = _dy_model
    _ϵ̃ = _ỹ - _μ̃_y_t
    _grad = np.column_stack([Π @ _ϵ̃[:, d] for d in range(D)])
    _μ̃_y_next = _μ̃_y_t + (Dop @ _μ̃_y_t + κ * _grad) * Δt
    return _μ̃_y_next, _ϵ̃

#### run
## withGC: same plant/action as §4.5.1; μ_y low-pass filtered via GC roughness (τ = 2/γ_y)
M_gc = 1
γ_y_gc = AGT_γ_y_gc
_α_y_gc = Δt / (2.0 / γ_y_gc + Δt)

_a = np.zeros((T, A))
_xˣ = np.zeros((T, C))
_vˣ = np.zeros((T, B))
_y = np.zeros((T, D))

_μ_x = np.zeros((T, C))
_v = np.ones((T, D)) * [setpoint]
_μ_v = np.ones((T, D)) * [setpoint]
_μ_y = np.zeros((T, D))

_ϵ_x = np.zeros((T, C))
_ϵ_y = np.zeros((T, D))
_ϵ_v = np.zeros((T, D))
_ϵ̃_y = np.zeros((T, M_gc + 1, D))
F = np.zeros(T)
_dF___dμₓ = np.zeros((T, C))
_dF___da = np.zeros((T, A))
_dF___dμ̃ₓ = np.zeros((T, M_gc + 1, C))

_λ_y = np.full(D, AGT_λ_y if AGT_δ > 0 else .01)
_λ_x = np.full(C, .01)
_λ_v = np.full(B, AGT_λ_v if AGT_δ > 0 else 8e3)
κₓ = 1e-2 if AGT_δ > 0 else 1e-3
κₐ = 1e-2
action_on = 0

_a[0] = np.array([0.])
_xˣ[0] = x0_agt.copy()
_vˣ[0] = np.array([0.])
_y[0] = g_E(_xˣ[0], _vˣ[0])
_μ_x[0] = x0_agt.copy()
_μ_v[0] = np.array([setpoint])
_μ_y[0] = _y[0]

f_pred0 = f_M(_μ_x[0], _μ_v[0])
_ϵ_x[0] = _μ_x[0] - (_μ_x[0] + f_pred0 * Δt)
_ϵ_v[0] = h_M(_μ_x[0], _μ_v[0]) - _v[0]
_ϵ_y[0] = _y[0] - g_M(_μ_x[0], _μ_v[0])
_ϵ̃_y[0, 0] = _y[0] - _μ_y[0]

F[0] = F_fn_vec(_ϵ_y[0], _ϵ_x[0], _λ_y, _λ_x, _ϵ_v_t=_ϵ_v[0], _λ_v=_λ_v)
_vˣ[:, 0] = _vˣ_agt_pre[:, 0]

for t in range(T - 1):
    _vˣ[t + 1, 0] = _vˣ_agt_pre[t + 1, 0]

    k1 = f_E(_xˣ[t], _vˣ[t], _a[t])
    k2 = f_E(_xˣ[t] + 0.5 * Δt * k1, _vˣ[t], _a[t])
    k3 = f_E(_xˣ[t] + 0.5 * Δt * k2, _vˣ[t], _a[t])
    k4 = f_E(_xˣ[t] + Δt * k3, _vˣ[t], _a[t])
    _xˣ[t + 1] = _xˣ[t] + ((k1 + 2 * k2 + 2 * k3 + k4) / 6.0 + _ω_x[t]) * Δt
    _y[t + 1] = g_E(_xˣ[t + 1], _vˣ[t + 1]) + _ω_y[t]

    _μ_x[t + 1] = _μ_x[t] + (-κₓ * _dF___dμₓ_fn_vec(
        _μ_x[t], _μ_v[t], _ϵ_y[t], _ϵ_x[t], _ϵ_v[t], _λ_y, _λ_x, _λ_v)) * Δt

    ## GC benefit: smooth sensory belief (order-0 GC low-pass, τ = 2/γ_y)
    _μ_y[t + 1] = (1 - _α_y_gc) * _μ_y[t] + _α_y_gc * _y[t + 1]
    _ϵ̃_y[t + 1, 0] = _y[t + 1] - _μ_y[t + 1]

    ## action uses smoothed belief (woutGC uses raw y via _ϵ_y[t] above)
    _ϵ_y_smooth = _μ_y[t + 1] - g_M(_μ_x[t + 1], _μ_v[t + 1])
    _dF___da[t] = _dF___da_fn_vec(_ϵ_y_smooth, _λ_y, _ϵ_v[t], _λ_v, _xˣ[t], _a[t])

    if t >= action_on:
        _a[t + 1] = _a[t] + (-κₐ * _dF___da[t]) * Δt

    f_pred = f_M(_μ_x[t], _μ_v[t])
    _ϵ_x[t + 1] = _μ_x[t + 1] - (_μ_x[t] + f_pred * Δt)
    _ϵ_y[t + 1] = _y[t + 1] - g_M(_μ_x[t + 1], _μ_v[t + 1])
    _ϵ_v[t + 1] = h_M(_μ_x[t + 1], _μ_v[t + 1]) - _v[t + 1]
    F[t + 1] = F_fn_vec(_ϵ_y[t + 1], _ϵ_x[t + 1], _λ_y, _λ_x,
                         _ϵ_v_t=_ϵ_v[t + 1], _λ_v=_λ_v)

result_agt = build_agent_result(
    with_gc=True, M_gc=M_gc,
    _a=_a, _xˣ=_xˣ, _vˣ=_vˣ, _y=_y, _v=_v, _μ_x=_μ_x, _μ_y=_μ_y, F=F,
    _ϵ̃_x=np.zeros((T, M_gc + 1, C)), _ϵ̃_y=_ϵ̃_y,
    _dF___dμ̃ₓ=_dF___dμ̃ₓ, _dF___da=_dF___da,
)

# Optional — saves AGF (with GC) traces for P3 overlay comparison
from pathlib import Path
Path("results").mkdir(exist_ok=True)
result_agt.to_pickle("results/result_agt_with.pkl")

fig = plot_agent_result(result_agt, title_suffix=r'(with generalized coordinates)')
## fig.savefig('./ElectricSystem-agent-withGC', bbox_inches='tight', dpi=300)

##
# from IPython.display import Image, display
# display(Image(filename="ElectricSystem-agent-withGC-step.png"))

##
# from IPython.display import Image, display
# display(Image(filename="ElectricSystem-agent-withGC-ramp.png"))

##
# from IPython.display import Image, display
# display(Image(filename="ElectricSystem-agent-withGC-cos.png"))