Part 4: Data-driven model of battery with DMDc¶

In this notebook, we demonstrate how we can create data-driven models with Dynamic Mode Decomposition (DMD) for a battery and use them in Jaxonomy.

Dataset¶

We utilise the same dataset as in the previous tutorial. We also resample the dataset at a frequency of 10Hz.

Data driven modelling by Dynamic Mode Decomposition with control (DMDc)¶

Here, our interest is to find a discrete-time battery model which takes $i[k]$ at a time-index $k$ and generates outputs $v_t[k]$ and the ampere-hours discharged $d[k]$. We know that the SoC ($s$) is related to the measurement of Ampere-hours discharged $d$ through the algebraic relation: $s = \frac{Q-d}{Q}$. Thus, for learning the model we work with direct measurement $d$ instead of $s$.

Considering $v_t[k]$ and $d[k]$ as the state variables and $i[k]$ as the control input, we seek an approximate model for the battery dynamics with DMDc of the following form

\begin{align} \begin{pmatrix} v_t[k+1] \\ d[k+1] \end{pmatrix} = A \begin{pmatrix} v_t[k] \\ d[k] \end{pmatrix} + B \left( i[k] \right), \tag{1} \end{align}

where $A$ is a $2\times 2$ matrix and $B$ is $1\times 1$ matrix.

Library note: This notebook uses Jaxonomy's built-in reduced-order-modelling API — the dmdc fitting routine and the DMDForecaster discrete-time predictor block from jaxonomy.library.rom (both re-exported from jaxonomy.library) — rather than a hand-rolled inline DMDc. These were promoted from this tutorial's original inline code into the maintained library; the math and workflow below are unchanged, only the plumbing now calls the library.

In [1]:

%matplotlib inline
import matplotlib.pyplot as plt

import numpy as np

from math import ceil

from jax import config

config.update("jax_enable_x64", True)

import jax
import jax.numpy as jnp

from scipy.io import loadmat

train_file_name = "dataset_25deg/10-28-18_17.44 551_UDDS_25degC_LGHG2.mat"
val_file_name = "dataset_25deg/10-29-18_14.41 551_Mixed1_25degC_LGHG2.mat"

tain_data = loadmat(train_file_name)
val_data = loadmat(val_file_name)

Q = 3.0  # battery capacity is known to be 3Ah (LG HG2 cell)


def extract_features_from_matfile(filename, Q=3.0, dt=0.1):
    data = loadmat(filename)
    t = data["meas"][0][0][0]
    vt = data["meas"][0][0][2]
    curr = -data["meas"][0][0][3]
    D = -data["meas"][0][0][4]

    # Resample
    T_end = t[-1, 0]
    t_resampled = np.linspace(0.0, T_end, ceil(T_end / dt))
    vt_resampled = np.interp(t_resampled, t[:, 0], vt[:, 0])
    curr_resampled = np.interp(t_resampled, t[:, 0], curr[:, 0])
    D_resampled = np.interp(t_resampled, t[:, 0], D[:, 0])

    return (t_resampled, vt_resampled, curr_resampled, D_resampled)


t_train, vt_train, curr_train, d_train = extract_features_from_matfile(train_file_name)
t_val, vt_val, curr_val, d_val = extract_features_from_matfile(val_file_name)

%matplotlib inline import matplotlib.pyplot as plt import numpy as np from math import ceil from jax import config config.update("jax_enable_x64", True) import jax import jax.numpy as jnp from scipy.io import loadmat train_file_name = "dataset_25deg/10-28-18_17.44 551_UDDS_25degC_LGHG2.mat" val_file_name = "dataset_25deg/10-29-18_14.41 551_Mixed1_25degC_LGHG2.mat" tain_data = loadmat(train_file_name) val_data = loadmat(val_file_name) Q = 3.0 # battery capacity is known to be 3Ah (LG HG2 cell) def extract_features_from_matfile(filename, Q=3.0, dt=0.1): data = loadmat(filename) t = data["meas"][0][0][0] vt = data["meas"][0][0][2] curr = -data["meas"][0][0][3] D = -data["meas"][0][0][4] # Resample T_end = t[-1, 0] t_resampled = np.linspace(0.0, T_end, ceil(T_end / dt)) vt_resampled = np.interp(t_resampled, t[:, 0], vt[:, 0]) curr_resampled = np.interp(t_resampled, t[:, 0], curr[:, 0]) D_resampled = np.interp(t_resampled, t[:, 0], D[:, 0]) return (t_resampled, vt_resampled, curr_resampled, D_resampled) t_train, vt_train, curr_train, d_train = extract_features_from_matfile(train_file_name) t_val, vt_val, curr_val, d_val = extract_features_from_matfile(val_file_name)

To compute the $A$ and $B$ matrices in equation (1), we use dmdc from jaxonomy.library.rom (re-exported from jaxonomy.library). Given the snapshot matrices $X$ (states at $k$), $X'$ (states at $k+1$) and $U$ (control at $k$), it forms the augmented matrix $\Omega = \begin{pmatrix} X \\ U \end{pmatrix}$ and solves the least-squares problem $\begin{pmatrix} A & B \end{pmatrix} = X'\,\Omega^{+}$ through the SVD — exactly the DMDc computation for equation (1). It returns a DMDcResult exposing the full operators .A and .B (as well as the reduced .A_tilde, .B_tilde, the POD .basis, and the .eigenvalues).

In [2]:

# DMDc fitting (`dmdc`) and the discrete-time predictor block (`DMDForecaster`)
# both live in jaxonomy's reduced-order-modelling library.
from jaxonomy.library import dmdc, DMDForecaster

# DMDc fitting (`dmdc`) and the discrete-time predictor block (`DMDForecaster`) # both live in jaxonomy's reduced-order-modelling library. from jaxonomy.library import dmdc, DMDForecaster

In [3]:

# Prepare the data matrices

state_names = ["$v_t$", "$d$"]
state_data = np.vstack([vt_train, d_train]) # all state data
control_data = curr_train.reshape((1, curr_train.size)) # all control data

X = state_data[:, :-1]  # State at time k
X_prime = state_data[:, 1:]  # State at time k+1
U = control_data[:, :-1]  # Control input at time k

# Compute A and B matrices with the library DMDc
res = dmdc(X, X_prime, U)
Aest, Best = res.A, res.B

print("Estimated A and B matrices:")
print("A matrix:")
print(Aest)

print("B matrix:")
print(Best)

# Prepare the data matrices state_names = ["$v_t$", "$d$"] state_data = np.vstack([vt_train, d_train]) # all state data control_data = curr_train.reshape((1, curr_train.size)) # all control data X = state_data[:, :-1] # State at time k X_prime = state_data[:, 1:] # State at time k+1 U = control_data[:, :-1] # Control input at time k # Compute A and B matrices with the library DMDc res = dmdc(X, X_prime, U) Aest, Best = res.A, res.B print("Estimated A and B matrices:") print("A matrix:") print(Aest) print("B matrix:") print(Best)

Estimated A and B matrices:
A matrix:
[[1.00000680e+00 7.95046837e-06]
 [2.45928067e-08 9.99999984e-01]]
B matrix:
[[-7.42636033e-05]
 [ 2.76371923e-05]]

Note: one may also use dedicated libraries to achieve this, for example pykoopman or pysindy. For example, with the pykoopman package, the same matrices could be computed as follows:

import pykoopman as pk

state_data = np.vstack([vt_train, d_train]).T
control_data = curr_train.reshape((curr_train.size, 1))

dmdc = pk.regression.DMDc(svd_rank=None)
model = pk.Koopman(regressor=dmdc)
model.fit(state_data, u=control_data)

Aest = model.ur @ model.A @ model.ur.T
Best = model.ur @ model.B

print("Estimated A and B matrices:")
print(f"{Aest=}")
print(f"{Best=}")

and the trained model can be simulated with model.simulate.

We can now simulate the model, initialising the model at the initial values of the state measurements, and see how its predictions compare with the experiment data. We write the following simulate_dmdc function for this purporse.

In [4]:

def simulate_dmdc(A, B, initial_state, U):
    """
    Simulate the system dynamics using the estimated A and B matrices from DMDc.
    
    Given an initial state and a sequence of control inputs, this function simulates the system 
    over time and returns the state trajectory.
    
    Parameters
    ----------
    A : numpy.ndarray
        The estimated state transition matrix A, of shape (n_states, n_states).
    B : numpy.ndarray
        The estimated control input matrix B, of shape (n_states, n_controls).
    initial_state : numpy.ndarray
        The initial state vector, of shape (n_states,).
    U : numpy.ndarray
        Control input matrix, of shape (n_controls, n_samples). Each column represents the control 
        input vector at a time step.
    
    Returns
    -------
    X_sim : numpy.ndarray
        The simulated state trajectory, of shape (n_states, n_samples+1). Each column represents the state 
        vector at a time step. The first column is the initial state.
    
    Example
    -------
    # Assuming A and B are already computed, and initial_state and U are defined
    initial_state = np.array([vt_train[0], d_train[0]])
    X_sim = simulate_dmdc(A, B, initial_state, U)
    
    print("Simulated state trajectory:")
    print(X_sim)
    """
    n_states = initial_state.shape[0]
    n_samples = U.shape[1]
    
    # Initialize the state trajectory matrix
    X_sim = np.zeros((n_states, n_samples + 1))
    X_sim[:, 0] = initial_state
    
    # Simulate the system dynamics
    for k in range(n_samples):
        X_sim[:, k + 1] = A @ X_sim[:, k] + B @ U[:, k]
    
    return X_sim

def simulate_dmdc(A, B, initial_state, U): """ Simulate the system dynamics using the estimated A and B matrices from DMDc. Given an initial state and a sequence of control inputs, this function simulates the system over time and returns the state trajectory. Parameters ---------- A : numpy.ndarray The estimated state transition matrix A, of shape (n_states, n_states). B : numpy.ndarray The estimated control input matrix B, of shape (n_states, n_controls). initial_state : numpy.ndarray The initial state vector, of shape (n_states,). U : numpy.ndarray Control input matrix, of shape (n_controls, n_samples). Each column represents the control input vector at a time step. Returns ------- X_sim : numpy.ndarray The simulated state trajectory, of shape (n_states, n_samples+1). Each column represents the state vector at a time step. The first column is the initial state. Example ------- # Assuming A and B are already computed, and initial_state and U are defined initial_state = np.array([vt_train[0], d_train[0]]) X_sim = simulate_dmdc(A, B, initial_state, U) print("Simulated state trajectory:") print(X_sim) """ n_states = initial_state.shape[0] n_samples = U.shape[1] # Initialize the state trajectory matrix X_sim = np.zeros((n_states, n_samples + 1)) X_sim[:, 0] = initial_state # Simulate the system dynamics for k in range(n_samples): X_sim[:, k + 1] = A @ X_sim[:, k] + B @ U[:, k] return X_sim

In [5]:

initial_state = state_data[:, 0]

pred_state = simulate_dmdc(Aest, Best, initial_state, U)

lw = 0.5
fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 5))
axs[0].plot(t_train, control_data[0,:], label="discharge current: control", lw=lw)
axs[0].legend()
for i, ax in enumerate(axs[1:]):
    ax.plot(t_train, state_data[i, :], label=state_names[i] + ": exp", lw=lw)
    ax.plot(t_train, pred_state[i, :], label=state_names[i] + ": DMDc", lw=lw)
    ax.legend(loc="best")
fig.suptitle("Training with original state")
plt.tight_layout()
plt.show()

initial_state = state_data[:, 0] pred_state = simulate_dmdc(Aest, Best, initial_state, U) lw = 0.5 fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 5)) axs[0].plot(t_train, control_data[0,:], label="discharge current: control", lw=lw) axs[0].legend() for i, ax in enumerate(axs[1:]): ax.plot(t_train, state_data[i, :], label=state_names[i] + ": exp", lw=lw) ax.plot(t_train, pred_state[i, :], label=state_names[i] + ": DMDc", lw=lw) ax.legend(loc="best") fig.suptitle("Training with original state") plt.tight_layout() plt.show()

We can also compute aggregate error quantities as follows:

In [6]:

def compute_prediction_error(pred_state, state_data):
    rms_err = []
    for y, x in zip(pred_state, state_data):
        rms_err.append(np.sqrt(np.average((y - x) ** 2)))
    return rms_err


print("RMS: error:", compute_prediction_error(pred_state, state_data))

def compute_prediction_error(pred_state, state_data): rms_err = [] for y, x in zip(pred_state, state_data): rms_err.append(np.sqrt(np.average((y - x) ** 2))) return rms_err print("RMS: error:", compute_prediction_error(pred_state, state_data))

RMS: error: [np.float64(0.635225983464464), np.float64(0.0010199622057858219)]

Augmenting the DMDc state¶

In the above results. we note that while the model predicts $d$ quite reliably, the predictions of $v_t$ have a very large error. To fix, this, one can augment the state vector with features that are functions of the state vector. By augmenting the state with additional features, we are already performing extended DMD. However, we will cover extended DMD in more detail in the next tutorial.

In this article, the authors suggest using additional features of $\phi_1 = 1/v_t^2$ and $\phi_2 = \exp (-\frac{1}{2 v_t^2})$. Let's try augmenting the state with these additional features and re-estimate the following DMDc system:

\begin{align} \begin{pmatrix} v_t[k+1] \\ d[k+1] \\ \phi_1[k+1] \\ \phi_2[k+1] \end{pmatrix} = A \begin{pmatrix} v_t[k] \\ d[k] \\ \phi_1[k] \\ \phi_2[k] \end{pmatrix} + B \left( i[k] \right), \end{align},

where $A\in \mathbb{R}^{4 \times 4}$ and $B\in \mathbb{R}^{1 \times 1}$.

Note that there are two ways to work with the above system once $A$ and $B$ have been estimated. In the first method, one can the initial values of $v_t[0]$ and $d_t[0]$ to compute the initial values of $\phi_1[0]$ and $\phi_2[0]$. Subsequently, we can let the state evolve with the above equations without any changes at each timestep. In the second method, at the end of each time-step, we can take the dynamics output $v_t[k+1]$ and $d[k+1]$ and recompute $\phi_1[k+1]$ and $\phi_2[k+1]$ from $v_t[k+1]$ and $d[k+1]$ based on our definitions of $\phi_1$ and $\phi_2$. Thus, after every timestep, the evolved $\phi_1[k+1]$ and $\phi_2[k+1]$ are discarded and replaced with their new values based on evolved $v_t[k+1]$ and $d[k+1]$. In what follows, we demonstrate the first method. The second method is demonstrated in the next tutorial on extended DMDc.

In [7]:

# create augmented state
state_names = ["$v_t$", "$d$", "$\phi_1$", "$\phi_2$"]
f_phi1 = lambda x: 1 / x**2
f_phi2 = lambda x: np.exp(1.0 / 2 / x**2)
state_data = np.vstack([vt_train, d_train, f_phi1(vt_train), f_phi2(vt_train)])
control_data = curr_train.reshape((1, curr_train.size))


X = state_data[:, :-1]  # State at time k
X_prime = state_data[:, 1:]  # State at time k+1
U = control_data[:, :-1]  # Control input at time k

# Compute A and B matrices with the library DMDc
res = dmdc(X, X_prime, U)
Aest, Best = res.A, res.B

print("Estimated A and B matrices:")
print("A matrix:")
print(Aest)

print("B matrix:")
print(Best)

# create augmented state state_names = ["$v_t$", "$d$", "$\phi_1$", "$\phi_2$"] f_phi1 = lambda x: 1 / x**2 f_phi2 = lambda x: np.exp(1.0 / 2 / x**2) state_data = np.vstack([vt_train, d_train, f_phi1(vt_train), f_phi2(vt_train)]) control_data = curr_train.reshape((1, curr_train.size)) X = state_data[:, :-1] # State at time k X_prime = state_data[:, 1:] # State at time k+1 U = control_data[:, :-1] # Control input at time k # Compute A and B matrices with the library DMDc res = dmdc(X, X_prime, U) Aest, Best = res.A, res.B print("Estimated A and B matrices:") print("A matrix:") print(Aest) print("B matrix:") print(Best)

Estimated A and B matrices:
A matrix:
[[ 9.74288727e-01 -4.68933614e-03 -3.63239828e-01  1.24116965e-01]
 [ 2.05157947e-05  1.00000373e+00  2.89099104e-04 -9.88749889e-05]
 [ 1.03526993e-03  1.90127493e-04  1.01453464e+00 -4.99292652e-03]
 [ 5.37788563e-04  9.87835817e-05  7.54895958e-03  9.97406406e-01]]
B matrix:
[[-3.94937765e-04]
 [ 2.78929873e-05]
 [ 1.61247146e-05]
 [ 8.37949564e-06]]

In [8]:

initial_state = state_data[:, 0]

pred_state = simulate_dmdc(Aest, Best, initial_state, U)

print("RMS: error:", compute_prediction_error(pred_state, state_data))

fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 7))
axs[0].plot(t_train, control_data[0,:], label="discharge current: control", lw=lw)
axs[0].legend()
for i, ax in enumerate(axs[1:]):
    ax.plot(t_train, state_data[i, :], label=state_names[i] + ": exp", lw=lw)
    ax.plot(t_train, pred_state[i, :], label=state_names[i] + ": DMDc", lw=lw)
    ax.legend(loc="best")
fig.suptitle("Training with augmented state")
plt.tight_layout()
plt.show()

initial_state = state_data[:, 0] pred_state = simulate_dmdc(Aest, Best, initial_state, U) print("RMS: error:", compute_prediction_error(pred_state, state_data)) fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 7)) axs[0].plot(t_train, control_data[0,:], label="discharge current: control", lw=lw) axs[0].legend() for i, ax in enumerate(axs[1:]): ax.plot(t_train, state_data[i, :], label=state_names[i] + ": exp", lw=lw) ax.plot(t_train, pred_state[i, :], label=state_names[i] + ": DMDc", lw=lw) ax.legend(loc="best") fig.suptitle("Training with augmented state") plt.tight_layout() plt.show()

RMS: error: [np.float64(0.039511563006929845), np.float64(0.0004265231584104484), np.float64(0.0026201385090203703), np.float64(0.0013785619786720058)]

The learnt data-driven model is not perfect, but indeed much better than the previous version. We can test the models performance on unseen data for validation

In [9]:

state_data = np.vstack([vt_val, d_val, f_phi1(vt_val), f_phi2(vt_val)])
control_data = curr_val.reshape((1, curr_val.size))

initial_state = state_data[:, 0]
pred_state = simulate_dmdc(Aest, Best, initial_state, control_data[:, :-1])

print("RMS: error:", compute_prediction_error(pred_state, state_data))

fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 7))
axs[0].plot(t_val, control_data[0, :], label="discharge current: control", lw=lw)
axs[0].legend()
for i, ax in enumerate(axs[1:]):
    ax.plot(t_val, state_data[i, :], label=state_names[i] + ": exp", lw=lw)
    ax.plot(t_val, pred_state[i, :], label=state_names[i] + ": DMDc", lw=lw)
    ax.legend(loc="best")
fig.suptitle("Validation with augmented state")
plt.tight_layout()
plt.show()

state_data = np.vstack([vt_val, d_val, f_phi1(vt_val), f_phi2(vt_val)]) control_data = curr_val.reshape((1, curr_val.size)) initial_state = state_data[:, 0] pred_state = simulate_dmdc(Aest, Best, initial_state, control_data[:, :-1]) print("RMS: error:", compute_prediction_error(pred_state, state_data)) fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 7)) axs[0].plot(t_val, control_data[0, :], label="discharge current: control", lw=lw) axs[0].legend() for i, ax in enumerate(axs[1:]): ax.plot(t_val, state_data[i, :], label=state_names[i] + ": exp", lw=lw) ax.plot(t_val, pred_state[i, :], label=state_names[i] + ": DMDc", lw=lw) ax.legend(loc="best") fig.suptitle("Validation with augmented state") plt.tight_layout() plt.show()

RMS: error: [np.float64(0.05133994375746349), np.float64(0.004327815331528889), np.float64(0.003288549980968619), np.float64(0.0017280425337799243)]

One may try to improve the data-driven model further by using additional features, or consider alternative algorithms such as the extended DMDc and/or SINDy. These alternative algorithms are showcased in subsequent tutorials.

Including the DMDc model in Jaxonomy¶

Similarly to the previous tutorials in this series, ultimately our goal is to be able to include the battery model in Jaxonomy, so that it can be utilised in a wider system model.

Jaxonomy provides the DMDForecaster block for exactly this: it wraps a fitted operator as a jax-traceable discrete-time predictor that simulates the discrete dynamics $\mathbf{x}[k+1] = A\, \mathbf{x}[k] + B \; \mathbf{u}[k]$ inside jaxonomy.simulate. We pair it with a small DiscreteSource LeafSystem that replays the discrete control-input sequence. For the latter, one may also utilise the pre-built DataSource in Jaxonomy (see Jaxonomy documentation and source: /library/data_source.py).

In [10]:

import jaxonomy
from jaxonomy.framework import LeafSystem
from jaxonomy.library import Constant

from jaxonomy.simulation import SimulatorOptions, ODESolverOptions

from typing import NamedTuple


class DiscreteSource(LeafSystem):
    class DiscreteStateType(NamedTuple):
        index: jnp.int64
        source_val: jnp.float64

    def __init__(self, val_array, dt, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.val_array = val_array
        self.dt = dt

        self.declare_discrete_state(
            default_value=self.DiscreteStateType(index=0, source_val=val_array[0]),
            as_array=False,
        )

        self.declare_periodic_update(
            self._update,
            period=self.dt,
            offset=0,
        )

        self.declare_output_port(
            self._compute_output,
            default_value=self.val_array[0],
            period=self.dt,
            offset=0.0,
            requires_inputs = False,
        )

    def _compute_output(self, time, state, **params):
        return state.discrete_state.source_val

    def _update(self, time, state, *inputs, **params):
        index = state.discrete_state.index
        index = index + 1
        source_val = self.val_array[index]
        return self.DiscreteStateType(index=index, source_val=source_val)

import jaxonomy from jaxonomy.framework import LeafSystem from jaxonomy.library import Constant from jaxonomy.simulation import SimulatorOptions, ODESolverOptions from typing import NamedTuple class DiscreteSource(LeafSystem): class DiscreteStateType(NamedTuple): index: jnp.int64 source_val: jnp.float64 def __init__(self, val_array, dt, *args, **kwargs): super().__init__(*args, **kwargs) self.val_array = val_array self.dt = dt self.declare_discrete_state( default_value=self.DiscreteStateType(index=0, source_val=val_array[0]), as_array=False, ) self.declare_periodic_update( self._update, period=self.dt, offset=0, ) self.declare_output_port( self._compute_output, default_value=self.val_array[0], period=self.dt, offset=0.0, requires_inputs = False, ) def _compute_output(self, time, state, **params): return state.discrete_state.source_val def _update(self, time, state, *inputs, **params): index = state.discrete_state.index index = index + 1 source_val = self.val_array[index] return self.DiscreteStateType(index=index, source_val=source_val)

We can now include the learnt data-driven system in Jaxonomy and simulate.

In [11]:

builder = jaxonomy.DiagramBuilder()
initial_state = jnp.array(
    [vt_train[0], d_train[0], f_phi1(vt_train[0]), f_phi2(vt_train[0])]
)
dt = 0.1

dls = builder.add(
    DMDForecaster(
        A=Aest, B=Best, dt=dt, initial_state=initial_state, name="dls"
    )
)
control = builder.add(DiscreteSource(jnp.array(curr_train), dt=dt, name="control"))

builder.connect(control.output_ports[0], dls.input_ports[0])

diagram = builder.build()
context = diagram.create_context()

recorded_signals = {
    "state": diagram["dls"].output_ports[0],
    "control": diagram["control"].output_ports[0],
}

options = SimulatorOptions(max_major_steps=ceil(t_train[-1] / dt))
sol = jaxonomy.simulate(
    diagram,
    context,
    (0.0, t_train[-1]),
    options=options,
    recorded_signals=recorded_signals,
)

builder = jaxonomy.DiagramBuilder() initial_state = jnp.array( [vt_train[0], d_train[0], f_phi1(vt_train[0]), f_phi2(vt_train[0])] ) dt = 0.1 dls = builder.add( DMDForecaster( A=Aest, B=Best, dt=dt, initial_state=initial_state, name="dls" ) ) control = builder.add(DiscreteSource(jnp.array(curr_train), dt=dt, name="control")) builder.connect(control.output_ports[0], dls.input_ports[0]) diagram = builder.build() context = diagram.create_context() recorded_signals = { "state": diagram["dls"].output_ports[0], "control": diagram["control"].output_ports[0], } options = SimulatorOptions(max_major_steps=ceil(t_train[-1] / dt)) sol = jaxonomy.simulate( diagram, context, (0.0, t_train[-1]), options=options, recorded_signals=recorded_signals, )

22:38:26.789 - [jaxonomy][INFO]: Simulator ready to start: SimulatorOptions(math_backend=jax, enable_tracing=True, max_major_step_length=None, max_major_steps=159656, ode_solver_method=auto, rtol=1e-06, atol=1e-08, min_minor_step_size=None, max_minor_step_size=None, zc_bisection_loop_count=40, save_time_series=True, recorded_signals=2, return_context=True, validate=True), Dopri5Solver(system=Diagram(root, 2 nodes), rtol=1e-06, atol=1e-08, max_step_size=None, min_step_size=None, method='auto', enable_autodiff=False, max_checkpoints=16, supports_mass_matrix=False)

/Users/reda/Dev/jaxy/jaxonomy/.claude/worktrees/surrogate-modeling-rom-support-353830/jaxonomy/simulation/errors.py:157: UserWarning: jaxonomy.simulate: the recording buffer (buffer_length=159656) filled; the trajectory was recorded at reduced resolution (79829 of 159657 samples, keeping every 2th). The recorded time-series still starts at t0 and covers the whole trajectory (the head is no longer dropped; the last kept sample may precede tf by up to 2 steps). Set SimulatorOptions(buffer_length=159658) or larger to capture every sample, loosen rtol/atol, or reduce the number of recorded signals.
  return func(*args, **kwargs)

In [12]:

state_data = np.vstack(
    [vt_train, d_train, f_phi1(vt_train), f_phi2(vt_train)]
)  # exp data for plotting

fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 7))
axs[0].plot(sol.time, sol.outputs["control"], label="discharge current: control", lw=lw)
axs[0].legend()
for i, ax in enumerate(axs[1:]):
    ax.plot(t_train, state_data[i, :], label=state_names[i] + ": exp", lw=lw)
    ax.plot(
        sol.time, sol.outputs["state"][:, i], label=state_names[i] + ": DMDc", lw=lw
    )
    ax.legend(loc="best")
fig.suptitle("Jaxonomy")
plt.tight_layout()
plt.show()

state_data = np.vstack( [vt_train, d_train, f_phi1(vt_train), f_phi2(vt_train)] ) # exp data for plotting fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 7)) axs[0].plot(sol.time, sol.outputs["control"], label="discharge current: control", lw=lw) axs[0].legend() for i, ax in enumerate(axs[1:]): ax.plot(t_train, state_data[i, :], label=state_names[i] + ": exp", lw=lw) ax.plot( sol.time, sol.outputs["state"][:, i], label=state_names[i] + ": DMDc", lw=lw ) ax.legend(loc="best") fig.suptitle("Jaxonomy") plt.tight_layout() plt.show()

Training the DMDc model as part of the Jaxonomy block¶

We could go a step further and include the DMDc training within the Jaxonomy block. This is helpful where you want to incorporate all preprocessing within the Jaxonomy block and create your own library of reusable blocks. We show this below with a DiscreteDMDc block that fits the DMDc operators from the training data (again via the library dmdc) inside its constructor and then simulates the resulting discrete-time system.

In [13]:

class DiscreteDMDc(LeafSystem):
    """
    Compute DMDc matrices A and B from the training data.
    Subsequently implement the following discrete-time system
    x[k+1] = A x[k] + B u[k]
    """

    def __init__(self, state_data, control_data, dt, initial_state=None, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.state_data = state_data
        self.control_data = control_data
        self.dt = dt
        
        self._compute_dmdc()
        
        self.declare_input_port()

        if initial_state is None:
            initial_state = jnp.zeros(nx)

        self.declare_discrete_state(default_value=initial_state)

        self.declare_periodic_update(
            self._update_state,
            period=self.dt,
            offset=0.0,
        )

        self.declare_output_port(
            self._compute_output,
            default_value=initial_state,
            period=self.dt,
            offset=0.0,
            requires_inputs=False
        )

    def _compute_output(self, time, state, **params):
        return state.discrete_state

    def _update_state(self, time, state, u, **params):
        x = state.discrete_state
        # Dynamics: y = Ax + Bu
        y = jnp.matmul(self.A, x) + jnp.matmul(self.B, jnp.atleast_1d(u))
        return y

    def _compute_dmdc(self):
        X = self.state_data[:, :-1]  # State at time k
        X_prime = self.state_data[:, 1:]  # State at time k+1
        U = self.control_data[:, :-1]  # Control input at time k

        # Delegate the DMDc regression to the library
        res = dmdc(X, X_prime, U)

        self.A = res.A
        self.B = res.B
        self.nx, self.nu = self.B.shape

class DiscreteDMDc(LeafSystem): """ Compute DMDc matrices A and B from the training data. Subsequently implement the following discrete-time system x[k+1] = A x[k] + B u[k] """ def __init__(self, state_data, control_data, dt, initial_state=None, *args, **kwargs): super().__init__(*args, **kwargs) self.state_data = state_data self.control_data = control_data self.dt = dt self._compute_dmdc() self.declare_input_port() if initial_state is None: initial_state = jnp.zeros(nx) self.declare_discrete_state(default_value=initial_state) self.declare_periodic_update( self._update_state, period=self.dt, offset=0.0, ) self.declare_output_port( self._compute_output, default_value=initial_state, period=self.dt, offset=0.0, requires_inputs=False ) def _compute_output(self, time, state, **params): return state.discrete_state def _update_state(self, time, state, u, **params): x = state.discrete_state # Dynamics: y = Ax + Bu y = jnp.matmul(self.A, x) + jnp.matmul(self.B, jnp.atleast_1d(u)) return y def _compute_dmdc(self): X = self.state_data[:, :-1] # State at time k X_prime = self.state_data[:, 1:] # State at time k+1 U = self.control_data[:, :-1] # Control input at time k # Delegate the DMDc regression to the library res = dmdc(X, X_prime, U) self.A = res.A self.B = res.B self.nx, self.nu = self.B.shape

In [14]:

state_data = np.vstack([vt_train, d_train, f_phi1(vt_train), f_phi2(vt_train)])
control_data = curr_train.reshape((1, curr_train.size))

builder = jaxonomy.DiagramBuilder()

dmdc = builder.add(
    DiscreteDMDc(
        state_data, control_data, dt=dt, initial_state=initial_state, name="dmdc"
    )
)

control = builder.add(DiscreteSource(jnp.array(curr_train), dt=dt, name="control"))

builder.connect(control.output_ports[0], dmdc.input_ports[0])

diagram = builder.build()
context = diagram.create_context()

recorded_signals = {
    "state": diagram["dmdc"].output_ports[0],
    "control": diagram["control"].output_ports[0],
}

options = SimulatorOptions(max_major_steps=ceil(t_train[-1] / dt))
sol = jaxonomy.simulate(
    diagram,
    context,
    (0.0, t_train[-1]),
    options=options,
    recorded_signals=recorded_signals,
)

state_data = np.vstack([vt_train, d_train, f_phi1(vt_train), f_phi2(vt_train)]) control_data = curr_train.reshape((1, curr_train.size)) builder = jaxonomy.DiagramBuilder() dmdc = builder.add( DiscreteDMDc( state_data, control_data, dt=dt, initial_state=initial_state, name="dmdc" ) ) control = builder.add(DiscreteSource(jnp.array(curr_train), dt=dt, name="control")) builder.connect(control.output_ports[0], dmdc.input_ports[0]) diagram = builder.build() context = diagram.create_context() recorded_signals = { "state": diagram["dmdc"].output_ports[0], "control": diagram["control"].output_ports[0], } options = SimulatorOptions(max_major_steps=ceil(t_train[-1] / dt)) sol = jaxonomy.simulate( diagram, context, (0.0, t_train[-1]), options=options, recorded_signals=recorded_signals, )

22:38:28.269 - [jaxonomy][INFO]: Simulator ready to start: SimulatorOptions(math_backend=jax, enable_tracing=True, max_major_step_length=None, max_major_steps=159656, ode_solver_method=auto, rtol=1e-06, atol=1e-08, min_minor_step_size=None, max_minor_step_size=None, zc_bisection_loop_count=40, save_time_series=True, recorded_signals=2, return_context=True, validate=True), Dopri5Solver(system=Diagram(root, 2 nodes), rtol=1e-06, atol=1e-08, max_step_size=None, min_step_size=None, method='auto', enable_autodiff=False, max_checkpoints=16, supports_mass_matrix=False)

In [15]:

fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 7))
axs[0].plot(sol.time, sol.outputs["control"], label="discharge current: control", lw=lw)
axs[0].legend()
for i, ax in enumerate(axs[1:]):
    ax.plot(t_train, state_data[i, :], label=state_names[i] + ": exp", lw=lw)
    ax.plot(
        sol.time, sol.outputs["state"][:, i], label=state_names[i] + ": DMDc", lw=lw
    )
    ax.legend(loc="best")
fig.suptitle("Jaxonomy")
plt.tight_layout()
plt.show()

fig, axs = plt.subplots(len(state_names) + 1, 1, figsize=(11, 7)) axs[0].plot(sol.time, sol.outputs["control"], label="discharge current: control", lw=lw) axs[0].legend() for i, ax in enumerate(axs[1:]): ax.plot(t_train, state_data[i, :], label=state_names[i] + ": exp", lw=lw) ax.plot( sol.time, sol.outputs["state"][:, i], label=state_names[i] + ": DMDc", lw=lw ) ax.legend(loc="best") fig.suptitle("Jaxonomy") plt.tight_layout() plt.show()

In [ ]: