<img src="http://cse.lab.imtlucca.it/~bemporad/jax-sysid/images/jax-sysid-logo.png" alt="jax-sysid" width=40%/>
A Python package based on <a href="https://jax.readthedocs.io"> JAX </a> for linear and nonlinear system identification of state-space models, recurrent neural network (RNN) training, and nonlinear regression/classification.
# Contents
* [Package description](#description)
* [Installation](#install)
* [Basic usage](#basic-usage)
* [Linear state-space models](#linear)
* [Nonlinear system identification and RNNs](#nonlinear)
* [Static models and nonlinear regression/classification] (#static)
* [Contributors](#contributors)
* [Citing jax-sysid](#bibliography)
* [License](#license)
<a name="description"></a>
## Package description
**jax-sysid** is a Python package based on <a href="https://jax.readthedocs.io"> JAX </a> for linear and nonlinear system identification of state-space models, recurrent neural network (RNN) training, and nonlinear regression/classification. The algorithm can handle L1-regularization and group-Lasso regularization and relies on L-BFGS optimization for accurate modeling, fast convergence, and good sparsification of model coefficients.
The package implements the approach described in the following paper:
<a name="cite-Bem24"><a>
> [1] A. Bemporad, "[Linear and nonlinear system identification under $\ell_1$- and group-Lasso regularization via L-BFGS-B](
http://arxiv.org/abs/2403.03827)," submitted for publication. Available on arXiv at <a href="http://arxiv.org/abs/2403.03827">
http://arxiv.org/abs/2403.03827</a>, 2024. [[bib entry](#ref1)]
<a name="install"></a>
## Installation
~~~python
pip install jax-sysid
~~~
<a name="basic-usage"></a>
## Basic usage
<a name="linear"></a>
### Linear state-space models
Given input/output training data $(u_0,y_0)$, $\ldots$, $(u_{N-1},y_{N-1})$, $u_k\in R^{n_u}$, $y_k\in R^{n_y}$, we want to identify a state-space model in the following form
$$ x_{k+1}=Ax_k+Bu_k$$
$$ \hat y_k=Cx_k+Du_k $$
where $k$ denotes the sample instant, $x_k\in R^{n_x}$ is the vector of hidden states, and
$A,B,C,D$ are matrices of appropriate dimensions to be learned.
The training problem to solve is
$$\min_{z}r(z)+\frac{1}{N}\sum_{k=0}^{N-1} \|y_{k}-Cx_k-Du_k\|_2^2$$
$$\mbox{s.t.}\ x_{k+1}=Ax_k+Bu_k, \ k=0,\ldots,N-2$$
where $z=(\theta,x_0)$ and $\theta$ collecting the entries of $A,B,C,D$.
The regularization term $r(z)$ includes the following components:
$$\frac{1}{2} \rho_{\theta} \|\theta\|_2^2 $$
$$\rho_{x_0} \|x_0\|_2^2$$
$$\tau \left\|z\right\|_1$$
$$\tau_g\sum_{i=1}^{n_u} \|I_iz\|_2$$
with $\rho_\theta>0$, $\rho_{x_0}>0$, $\tau\geq 0$, $\tau_g\geq 0$. See examples below.
Let's start training a discrete-time linear model $(A,B,C,D)$ on a sequence of inputs $U=[u_0\ \ldots\ u_{N-1}]'$ and output $Y=[y_0\ \ldots\ y_{N-1}]'$, with regularization $\rho_\theta=10^{-2}$, $\rho_{x_0}=10^{-3}$, running the L-BFGS solver for at most 1000 function evaluations:
~~~python
from jax_sysid.models import LinearModel
model = LinearModel(nx, ny, nu)
model.loss(rho_x0=1.e-3, rho_th=1.e-2)
model.optimization(lbfgs_epochs=1000)
model.fit(Y,U)
Yhat, Xhat = model.predict(model.x0, U)
~~~
After identifying the model, to retrieve the resulting state-space realization you can use the following:
~~~python
A,B,C,D = model.ssdata()
~~~
Given a new test sequence of inputs and outputs, an initial state that is compatible with the identified model can be reconstructed by running an extended Kalman filter and Rauch–Tung–Striebel smoothing (cf.
[[1]](#cite-Bem24)) and used to simulate the model:
~~~python
x0_test = model.learn_x0(U_test, Y_test)
Yhat_test, Xhat_test = model.predict(x0_test, U_test)
~~~
R2-scores on training and test data can be computed as follows:
~~~python
from jax_sysid.utils import compute_scores
R2_train, R2_test, msg = compute_scores(Y, Yhat, Y_test, Yhat_test, fit='R2')
print(msg)
~~~
It is good practice to scale the input and output signals. To identify a model on scaled signals, you can use the following:
~~~python
from jax_sysid.utils import standard_scale, unscale
Ys, ymean, ygain = standard_scale(Y)
Us, umean, ugain = standard_scale(U)
model.fit(Ys, Us)
Yshat, Xhat = model.predict(model.x0, Us)
Yhat = unscale(Yshat, ymean, ygain)
~~~
Let us now retrain the model using L1-regularization
and check the sparsity of the resulting model:
~~~python
model.loss(rho_x0=1.e-3, rho_th=1.e-2, tau_th=0.03)
model.fit(Ys, Us)
print(model.sparsity_analysis())
~~~
To reduce the number of states in the model, you can use group-Lasso regularization as follows:
~~~python
model.loss(rho_x0=1.e-3, rho_th=1.e-2, tau_g=0.1)
model.group_lasso_x()
model.fit(Ys, Us)
~~~
Groups in this case are entries in $A,B,C,x_0$ related to the same state.
Group-Lasso can be also used to try to reduce the number of inputs that are relevant in the model. You can do this as follows:
~~~python
model.loss(rho_x0=1.e-3, rho_th=1.e-2, tau_g=0.15)
model.group_lasso_u()
model.fit(Ys, Us)
~~~
Groups in this case are entries in $B,D$ related to the same input.
**jax-sysid** also supports multiple training experiments. In this case, the sequences of training inputs and outputs are passed as a list of arrays. For example, if three experiments are available for training, use the following command:
~~~python
model.fit([Ys1, Ys2, Ys3], [Us1, Us2, Us3])
~~~
In case the initial state $x_0$ is trainable, one initial state per experiment is optimized. To avoid training the initial state, add `train_x0=False` when calling `model.loss`.
<a name="nonlinear"></a>
### Nonlinear system identification and RNNs
Given input/output training data $(u_0,y_0)$, $\ldots$, $(u_{N-1},y_{N-1})$, $u_k\in R^{n_u}$, $y_k\in R^{n_y}$, we want to identify a nonlinear parametric state-space model in the following form
$$ x_{k+1}=f(x_k,u_k,\theta)$$
$$ \hat y_k=g(x_k,u_k,\theta)$$
where $k$ denotes the sample instant, $x_k\in R^{n_x}$ is the vector of hidden states, and $\theta$ collects the trainable parameters of the model.
As for the linear case, the training problem to solve is
$$ \min_{z}r(z)+\frac{1}{N}\sum_{k=0}^{N-1} \|y_{k}-g(x_k,u_k,\theta)\|_2^2$$
$$\mbox{s.t.}\ x_{k+1}=f(x_k,u_k,\theta),\ k=0,\ldots,N-2$$
where $z=(\theta,x_0)$. The regularization term $r(z)$ is the same as in the linear case.
For example, let us consider the following residual RNN model without input/output feedthrough:
$$ x_{k+1}=Ax_k+Bu_k+f_x(x_k,u_k,\theta_x)$$
$$ \hat y_k=Cx_k+f_y(x_k,\theta_y)$$
where $f_x$, $f_y$ are feedforward shallow neural networks, and let $z$ collects the coefficients in $A,B,C,D,\theta_x,\theta_y$. We want to train $z$ by running 1000 Adam iterations followed by at most 1000 L-BFGS function evaluations:
~~~python
from jax_sysid.models import Model
Ys, ymean, ygain = standard_scale(Y)
Us, umean, ugain = standard_scale(U)
def sigmoid(x):
return 1. / (1. + jnp.exp(-x))
@jax.jit
def state_fcn(x,u,params):
A,B,C,W1,W2,W3,b1,b2,W4,W5,b3,b4=params
return A@x+B@u+W3@sigmoid(W1@x+W2@u+b1)+b2
@jax.jit
def output_fcn(x,u,params):
A,B,C,W1,W2,W3,b1,b2,W4,W5,b3,b4=params
return C@x+W5@sigmoid(W4@x+b3)+b4
model = Model(nx, ny, nu, state_fcn=state_fcn, output_fcn=output_fcn)
nnx = 5 # number of hidden neurons in state-update function
nny = 5 # number of hidden neurons in output function
# Parameter initialization:
A = 0.5*np.eye(nx)
B = 0.1*np.random.randn(nx,nu)
C = 0.1*np.random.randn(ny,nx)
W1 = 0.1*np.random.randn(nnx,nx)
W2 = 0.5*np.random.randn(nnx,nu)
W3 = 0.5*np.random.randn(nx,nnx)
b1 = np.zeros(nnx)
b2 = np.zeros(nx)
W4 = 0.5*np.random.randn(nny,nx)
W5 = 0.5*np.random.randn(ny,nny)
b3 = np.zeros(nny)
b4 = np.zeros(ny)
model.init(params=[A,B,C,W1,W2,W3,b1,b2,W4,W5,b3,b4])
model.loss(rho_x0=1.e-4, rho_th=1.e-4)
model.optimization(adam_epochs=1000, lbfgs_epochs=1000)
model.fit(Ys, Us)
Yshat, Xshat = model.predict(model.x0, Us)
Yhat = unscale(Yshat, ymean, ygain)
~~~
**jax-sysid** also supports recurrent neural networks defined via the **flax.linen** library:
~~~python
from jax_sysid.models import RNN
# state-update function
class FX(nn.Module):
@nn.compact
def __call__(self, x):
x = nn.Dense(features=5)(x)
x = nn.swish(x)
x = nn.Dense(features=5)(x)
x = nn.swish(x)
x = nn.Dense(features=nx)(x)
return x
# output function
class FY(nn.Module):
@nn.compact
def __call__(self, x):
x = nn.Dense(features=5)(x)
x = nn.tanh(x)
x = nn.Dense(features=ny)(x)
return x
model = RNN(nx, ny, nu, FX=FX, FY=FY, x_scaling=0.1)
model.loss(rho_x0=1.e-4, rho_th=1.e-4, tau_th=0.0001)
model.optimization(adam_epochs=0, lbfgs_epochs=2000)
model.fit(Ys, Us)
~~~
where the extra parameter `x_scaling` is used to scale down (when $0\leq$ `x_scaling` $<1$) the default initialization of the network weights instantiated by **flax**.
**jax-sysid** also supports custom loss functions penalizing the deviations of $\hat y$ from $y$. For example, to identify a system with a binary output, we can use the (modified) cross-entropy loss
$$
{\mathcal L}(\hat Y,Y)=\frac{1}{N}\sum_{k=0}^{N-1}
-y_k\log(\epsilon+\hat y_k)-(1-y_k)\log(\epsilon+1-\hat y_k)
$$
where $\hat Y=(\hat y_0,\ldots,\hat y_{N-1})$ and $Y=(y_0,\ldots, y_{N-1})$ are the sequences of predicted and measured outputs, respectively, and $\epsilon>0$ is a tolerance used to prevent numerical issues in case $\hat y_k\approx 0$ or $\hat y_k\approx 1$:
~~~python
epsil=1.e-4
@jax.jit
def cross_entropy_loss(Yhat,Y):
loss=jnp.sum(-Y*jnp.log(epsil+Yhat)-(1.-Y)*jnp.log(epsil+1.-Yhat))/Y.shape[0]
return loss
model.loss(rho_x0=0.01, rho_th=0.001, output_loss=cross_entropy_loss)
~~~
By default, **jax-sysid** minimizes the classical mean squared error
$$
{\mathcal L}(\hat Y,Y)=\frac{1}{N}\sum_{k=0}^{N-1}
\|y_k-\hat y_k\|_2^2
$$
**jax-sysid** also supports custom regularization terms $r_c(z)$, where $z=(\theta,x_0)$. You can specify such a custom regularization function when defining the overall loss. For example, say for some reason you want to impose $\|\theta\|_2^2\leq 1$ as a soft constraint, you can penalize
$$\frac{1}{2} \rho_{\theta} \|\theta\|_2^2 + \rho_{x_0} \|x_0\|_2^2 + \rho_c\max\{\|\theta\|_2^2-1,0\}^2$$
with $\rho_c\gg\rho_\theta$, $\rho_c\gg\rho_{x_0}$, for instance $\rho_c=1000$, $\rho_\theta=0.001$, $\rho_{x0}=0.01$. In Python:
~~~python
@jax.jit
def custom_reg_fcn(th,x0):
return 1000.*jnp.maximum(jnp.sum(th**2)-1.,0.)**2
model.loss(rho_x0=0.01, rho_th=0.001, custom_regularization= custom_reg_fcn)
~~~
To include lower and upper bounds on the parameters of the model and/or the initial state, use the following additional arguments when specifying the optimization problem:
~~~python
model.optimization(params_min=lb, params_max=ub, x0_min=xmin, x0_max=xmax, ...)
~~~
where `lb` and `ub` are lists of arrays with the same structure as `model.params`, while `xmin` and `xmax` are arrays of the same dimension `model.nx` of the state vector. By default, each value is set equal to `None`, i.e., the corresponding constraint is not enforced. See `example_linear_positive.py` for examples of how to use nonnegative constraints to fit a positive linear system.
<a name="static"></a>
### Static models and nonlinear regression / classification
The same optimization algorithms used to train dynamical models can be used to train static models, i.e., to solve the nonlinear regression problem:
$$ \min_{z}r(z)+\frac{1}{N}\sum_{k=0}^{N-1} \|y_{k}-f(u_k,\theta)\|_2^2$$
where $z=\theta$ is the vector of model parameters to train and $r(z)$ admits the same
regularization terms as in the case of dynamical models.
For example, if the model is a shallow neural network you can use the following code:
~~~python
from jax_sysid.models import StaticModel
from jax_sysid.utils import standard_scale, unscale
@jax.jit
def output_fcn(u, params):
W1,b1,W2,b2=params
y = W1@u.T+b1
y = W2@jnp.arctan(y)+b2
return y.T
model = StaticModel(ny, nu, output_fcn)
nn=10 # number of neurons
model.init(params=[np.random.randn(nn,nu), np.random.randn(nn,1), np.random.randn(1,nn), np.random.randn(1,1)])
model.loss(rho_th=1.e-4, tau_th=tau_th)
model.optimization(lbfgs_epochs=500)
model.fit(Ys, Us)
~~~
**jax-sysid** also supports feedforward neural networks defined via the **flax.linen** library:
~~~python
from jax_sysid.models import FNN
from flax import linen as nn
# output function
class FY(nn.Module):
@nn.compact
def __call__(self, x):
x = nn.Dense(features=20)(x)
x = nn.tanh(x)
x = nn.Dense(features=20)(x)
x = nn.tanh(x)
x = nn.Dense(features=ny)(x)
return x
model = FNN(ny, nu, FY)
model.loss(rho_th=1.e-4, tau_th=tau_th)
model.optimization(lbfgs_epochs=500)
model.fit(Ys, Us)
~~~
To include lower and upper bounds on the parameters of the model, use the following additional arguments when specifying the optimization problem:
~~~python
model.optimization(lbfgs_epochs=500, params_min=lb, params_max=ub)
~~~
where `lb` and `ub` are lists of arrays with the same structure as `model.params`. See `example_static_convex.py` for examples of how to use nonnegative constraints to fit input-convex neural networks.
To solve classification problems, you need to define a custom loss function to change the default Mean-Squared-Error loss. For example, to train a classifier for a multi-category classification problem with $K$ classes, you can specify a neural network with a linear output layer generating output predictions $\hat y\in R^K$ and define the associated cross-entropy $\ell(\hat y,y) = -\sum_{k=1}^Ky_k\log\left(\frac{e^{\hat y_k}}{\sum_{j=1}^Ke^{\hat y_j}}\right)$ function as follows:
~~~python
def cross_entropy(Yhat,Y):
return -jax.numpy.sum(jax.nn.log_softmax(Yhat, axis=1)*Y)/Y.shape[0]
model.loss(rho_th=1.e-4, output_loss=cross_entropy)
~~~
See `example_static_fashion_mist.py` for an example using **Keras** with JAX backend to define the neural network model.
<a name="contributors"><a>
## Contributors
This package was coded by Alberto Bemporad.
This software is distributed without any warranty. Please cite the paper below if you use this software.
<a name="bibliography"><a>
## Citing jax-sysid
<a name="ref1"></a>
```
@article{Bem24,
author={A. Bemporad},
title={Linear and nonlinear system identification under $\ell_1$- and group-{Lasso} regularization via {L-BFGS-B}},
note = {submitted for publication. Also available on arXiv
at \url{http://arxiv.org/abs/2403.03827}},
year=2024
}
```
<a name="license"><a>
## License
Apache 2.0
(C) 2024 A. Bemporad
Raw data
{
"_id": null,
"home_page": null,
"name": "jax-sysid",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "system identification, subspace identification, nonlinear regression",
"author": null,
"author_email": "Alberto Bemporad <alberto.bemporad@imtlucca.it>",
"download_url": "https://files.pythonhosted.org/packages/7c/7c/58f0194f35eb69f842e9250ea3e3c37d1f1acc35e5f3bb1d173480fd2e3b/jax_sysid-0.4.2.tar.gz",
"platform": null,
"description": "<img src=\"http://cse.lab.imtlucca.it/~bemporad/jax-sysid/images/jax-sysid-logo.png\" alt=\"jax-sysid\" width=40%/>\n\nA Python package based on <a href=\"https://jax.readthedocs.io\"> JAX </a> for linear and nonlinear system identification of state-space models, recurrent neural network (RNN) training, and nonlinear regression/classification.\n \n# Contents\n\n* [Package description](#description)\n\n* [Installation](#install)\n\n* [Basic usage](#basic-usage)\n * [Linear state-space models](#linear)\n * [Nonlinear system identification and RNNs](#nonlinear)\n * [Static models and nonlinear regression/classification] (#static)\n\n* [Contributors](#contributors)\n\n* [Citing jax-sysid](#bibliography)\n\n* [License](#license)\n\n\n<a name=\"description\"></a>\n## Package description \n\n**jax-sysid** is a Python package based on <a href=\"https://jax.readthedocs.io\"> JAX </a> for linear and nonlinear system identification of state-space models, recurrent neural network (RNN) training, and nonlinear regression/classification. The algorithm can handle L1-regularization and group-Lasso regularization and relies on L-BFGS optimization for accurate modeling, fast convergence, and good sparsification of model coefficients.\n\nThe package implements the approach described in the following paper:\n\n<a name=\"cite-Bem24\"><a>\n> [1] A. Bemporad, \"[Linear and nonlinear system identification under $\\ell_1$- and group-Lasso regularization via L-BFGS-B](\nhttp://arxiv.org/abs/2403.03827),\" submitted for publication. Available on arXiv at <a href=\"http://arxiv.org/abs/2403.03827\">\nhttp://arxiv.org/abs/2403.03827</a>, 2024. [[bib entry](#ref1)]\n\n\n<a name=\"install\"></a>\n## Installation\n\n~~~python\npip install jax-sysid\n~~~\n\n<a name=\"basic-usage\"></a>\n## Basic usage\n\n<a name=\"linear\"></a>\n### Linear state-space models\n\nGiven input/output training data $(u_0,y_0)$, $\\ldots$, $(u_{N-1},y_{N-1})$, $u_k\\in R^{n_u}$, $y_k\\in R^{n_y}$, we want to identify a state-space model in the following form\n\n$$ x_{k+1}=Ax_k+Bu_k$$\n\n$$ \\hat y_k=Cx_k+Du_k $$\n\nwhere $k$ denotes the sample instant, $x_k\\in R^{n_x}$ is the vector of hidden states, and\n$A,B,C,D$ are matrices of appropriate dimensions to be learned.\n\nThe training problem to solve is\n\n$$\\min_{z}r(z)+\\frac{1}{N}\\sum_{k=0}^{N-1} \\|y_{k}-Cx_k-Du_k\\|_2^2$$\n\n$$\\mbox{s.t.}\\ x_{k+1}=Ax_k+Bu_k, \\ k=0,\\ldots,N-2$$\n\nwhere $z=(\\theta,x_0)$ and $\\theta$ collecting the entries of $A,B,C,D$.\n\nThe regularization term $r(z)$ includes the following components:\n\n$$\\frac{1}{2} \\rho_{\\theta} \\|\\theta\\|_2^2 $$\n\n$$\\rho_{x_0} \\|x_0\\|_2^2$$\n\n$$\\tau \\left\\|z\\right\\|_1$$\n\n$$\\tau_g\\sum_{i=1}^{n_u} \\|I_iz\\|_2$$\n\nwith $\\rho_\\theta>0$, $\\rho_{x_0}>0$, $\\tau\\geq 0$, $\\tau_g\\geq 0$. See examples below.\n\nLet's start training a discrete-time linear model $(A,B,C,D)$ on a sequence of inputs $U=[u_0\\ \\ldots\\ u_{N-1}]'$ and output $Y=[y_0\\ \\ldots\\ y_{N-1}]'$, with regularization $\\rho_\\theta=10^{-2}$, $\\rho_{x_0}=10^{-3}$, running the L-BFGS solver for at most 1000 function evaluations:\n\n~~~python\nfrom jax_sysid.models import LinearModel\n\nmodel = LinearModel(nx, ny, nu)\nmodel.loss(rho_x0=1.e-3, rho_th=1.e-2) \nmodel.optimization(lbfgs_epochs=1000) \nmodel.fit(Y,U)\nYhat, Xhat = model.predict(model.x0, U)\n~~~\n\nAfter identifying the model, to retrieve the resulting state-space realization you can use the following:\n\n~~~python\nA,B,C,D = model.ssdata()\n~~~\n\nGiven a new test sequence of inputs and outputs, an initial state that is compatible with the identified model can be reconstructed by running an extended Kalman filter and Rauch\u2013Tung\u2013Striebel smoothing (cf. \n[[1]](#cite-Bem24)) and used to simulate the model:\n\n~~~python\nx0_test = model.learn_x0(U_test, Y_test)\nYhat_test, Xhat_test = model.predict(x0_test, U_test)\n~~~\n\nR2-scores on training and test data can be computed as follows:\n\n~~~python\nfrom jax_sysid.utils import compute_scores\n\nR2_train, R2_test, msg = compute_scores(Y, Yhat, Y_test, Yhat_test, fit='R2')\nprint(msg)\n~~~\n\nIt is good practice to scale the input and output signals. To identify a model on scaled signals, you can use the following:\n\n~~~python\nfrom jax_sysid.utils import standard_scale, unscale\n\nYs, ymean, ygain = standard_scale(Y)\nUs, umean, ugain = standard_scale(U)\nmodel.fit(Ys, Us)\nYshat, Xhat = model.predict(model.x0, Us)\nYhat = unscale(Yshat, ymean, ygain)\n~~~\n\nLet us now retrain the model using L1-regularization\nand check the sparsity of the resulting model:\n\n~~~python\nmodel.loss(rho_x0=1.e-3, rho_th=1.e-2, tau_th=0.03) \nmodel.fit(Ys, Us)\nprint(model.sparsity_analysis())\n~~~\n \nTo reduce the number of states in the model, you can use group-Lasso regularization as follows:\n\n~~~python\nmodel.loss(rho_x0=1.e-3, rho_th=1.e-2, tau_g=0.1) \nmodel.group_lasso_x()\nmodel.fit(Ys, Us)\n~~~\nGroups in this case are entries in $A,B,C,x_0$ related to the same state.\n\nGroup-Lasso can be also used to try to reduce the number of inputs that are relevant in the model. You can do this as follows:\n\n~~~python\nmodel.loss(rho_x0=1.e-3, rho_th=1.e-2, tau_g=0.15) \nmodel.group_lasso_u()\nmodel.fit(Ys, Us)\n~~~\nGroups in this case are entries in $B,D$ related to the same input.\n\n**jax-sysid** also supports multiple training experiments. In this case, the sequences of training inputs and outputs are passed as a list of arrays. For example, if three experiments are available for training, use the following command:\n\n~~~python\nmodel.fit([Ys1, Ys2, Ys3], [Us1, Us2, Us3])\n~~~\n\nIn case the initial state $x_0$ is trainable, one initial state per experiment is optimized. To avoid training the initial state, add `train_x0=False` when calling `model.loss`.\n\n<a name=\"nonlinear\"></a>\n### Nonlinear system identification and RNNs\nGiven input/output training data $(u_0,y_0)$, $\\ldots$, $(u_{N-1},y_{N-1})$, $u_k\\in R^{n_u}$, $y_k\\in R^{n_y}$, we want to identify a nonlinear parametric state-space model in the following form\n\n$$ x_{k+1}=f(x_k,u_k,\\theta)$$\n\n$$ \\hat y_k=g(x_k,u_k,\\theta)$$\n\nwhere $k$ denotes the sample instant, $x_k\\in R^{n_x}$ is the vector of hidden states, and $\\theta$ collects the trainable parameters of the model.\n\nAs for the linear case, the training problem to solve is\n\n$$ \\min_{z}r(z)+\\frac{1}{N}\\sum_{k=0}^{N-1} \\|y_{k}-g(x_k,u_k,\\theta)\\|_2^2$$\n\n$$\\mbox{s.t.}\\ x_{k+1}=f(x_k,u_k,\\theta),\\ k=0,\\ldots,N-2$$\n\nwhere $z=(\\theta,x_0)$. The regularization term $r(z)$ is the same as in the linear case.\n\nFor example, let us consider the following residual RNN model without input/output feedthrough:\n\n$$ x_{k+1}=Ax_k+Bu_k+f_x(x_k,u_k,\\theta_x)$$ \n\n$$ \\hat y_k=Cx_k+f_y(x_k,\\theta_y)$$\n\nwhere $f_x$, $f_y$ are feedforward shallow neural networks, and let $z$ collects the coefficients in $A,B,C,D,\\theta_x,\\theta_y$. We want to train $z$ by running 1000 Adam iterations followed by at most 1000 L-BFGS function evaluations:\n\n~~~python\nfrom jax_sysid.models import Model\n\nYs, ymean, ygain = standard_scale(Y)\nUs, umean, ugain = standard_scale(U)\n\ndef sigmoid(x):\n return 1. / (1. + jnp.exp(-x)) \n@jax.jit\ndef state_fcn(x,u,params):\n A,B,C,W1,W2,W3,b1,b2,W4,W5,b3,b4=params\n return A@x+B@u+W3@sigmoid(W1@x+W2@u+b1)+b2 \n@jax.jit\ndef output_fcn(x,u,params):\n A,B,C,W1,W2,W3,b1,b2,W4,W5,b3,b4=params\n return C@x+W5@sigmoid(W4@x+b3)+b4\n\nmodel = Model(nx, ny, nu, state_fcn=state_fcn, output_fcn=output_fcn)\nnnx = 5 # number of hidden neurons in state-update function\nnny = 5 # number of hidden neurons in output function\n\n# Parameter initialization:\nA = 0.5*np.eye(nx)\nB = 0.1*np.random.randn(nx,nu)\nC = 0.1*np.random.randn(ny,nx)\nW1 = 0.1*np.random.randn(nnx,nx)\nW2 = 0.5*np.random.randn(nnx,nu)\nW3 = 0.5*np.random.randn(nx,nnx)\nb1 = np.zeros(nnx)\nb2 = np.zeros(nx)\nW4 = 0.5*np.random.randn(nny,nx)\nW5 = 0.5*np.random.randn(ny,nny)\nb3 = np.zeros(nny)\nb4 = np.zeros(ny)\nmodel.init(params=[A,B,C,W1,W2,W3,b1,b2,W4,W5,b3,b4]) \n\nmodel.loss(rho_x0=1.e-4, rho_th=1.e-4)\nmodel.optimization(adam_epochs=1000, lbfgs_epochs=1000) \nmodel.fit(Ys, Us)\nYshat, Xshat = model.predict(model.x0, Us)\nYhat = unscale(Yshat, ymean, ygain)\n~~~\n\n**jax-sysid** also supports recurrent neural networks defined via the **flax.linen** library:\n\n\n~~~python\nfrom jax_sysid.models import RNN\n\n# state-update function\nclass FX(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(features=5)(x)\n x = nn.swish(x)\n x = nn.Dense(features=5)(x)\n x = nn.swish(x)\n x = nn.Dense(features=nx)(x)\n return x\n\n# output function\nclass FY(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(features=5)(x)\n x = nn.tanh(x)\n x = nn.Dense(features=ny)(x)\n return x\n \nmodel = RNN(nx, ny, nu, FX=FX, FY=FY, x_scaling=0.1)\nmodel.loss(rho_x0=1.e-4, rho_th=1.e-4, tau_th=0.0001)\nmodel.optimization(adam_epochs=0, lbfgs_epochs=2000) \nmodel.fit(Ys, Us)\n~~~\nwhere the extra parameter `x_scaling` is used to scale down (when $0\\leq$ `x_scaling` $<1$) the default initialization of the network weights instantiated by **flax**.\n\n**jax-sysid** also supports custom loss functions penalizing the deviations of $\\hat y$ from $y$. For example, to identify a system with a binary output, we can use the (modified) cross-entropy loss\n\n$$\n\t{\\mathcal L}(\\hat Y,Y)=\\frac{1}{N}\\sum_{k=0}^{N-1}\n\t-y_k\\log(\\epsilon+\\hat y_k)-(1-y_k)\\log(\\epsilon+1-\\hat y_k)\n$$\n\nwhere $\\hat Y=(\\hat y_0,\\ldots,\\hat y_{N-1})$ and $Y=(y_0,\\ldots, y_{N-1})$ are the sequences of predicted and measured outputs, respectively, and $\\epsilon>0$ is a tolerance used to prevent numerical issues in case $\\hat y_k\\approx 0$ or $\\hat y_k\\approx 1$:\n\n~~~python\nepsil=1.e-4\n@jax.jit\ndef cross_entropy_loss(Yhat,Y):\n loss=jnp.sum(-Y*jnp.log(epsil+Yhat)-(1.-Y)*jnp.log(epsil+1.-Yhat))/Y.shape[0]\n return loss\nmodel.loss(rho_x0=0.01, rho_th=0.001, output_loss=cross_entropy_loss)\n~~~\n\nBy default, **jax-sysid** minimizes the classical mean squared error\n\n$$\n\t{\\mathcal L}(\\hat Y,Y)=\\frac{1}{N}\\sum_{k=0}^{N-1}\n\t\\|y_k-\\hat y_k\\|_2^2\n$$\n\n**jax-sysid** also supports custom regularization terms $r_c(z)$, where $z=(\\theta,x_0)$. You can specify such a custom regularization function when defining the overall loss. For example, say for some reason you want to impose $\\|\\theta\\|_2^2\\leq 1$ as a soft constraint, you can penalize\n\n$$\\frac{1}{2} \\rho_{\\theta} \\|\\theta\\|_2^2 + \\rho_{x_0} \\|x_0\\|_2^2 + \\rho_c\\max\\{\\|\\theta\\|_2^2-1,0\\}^2$$\n\nwith $\\rho_c\\gg\\rho_\\theta$, $\\rho_c\\gg\\rho_{x_0}$, for instance $\\rho_c=1000$, $\\rho_\\theta=0.001$, $\\rho_{x0}=0.01$. In Python:\n\n~~~python\n@jax.jit\ndef custom_reg_fcn(th,x0):\n return 1000.*jnp.maximum(jnp.sum(th**2)-1.,0.)**2\nmodel.loss(rho_x0=0.01, rho_th=0.001, custom_regularization= custom_reg_fcn)\n~~~\n\nTo include lower and upper bounds on the parameters of the model and/or the initial state, use the following additional arguments when specifying the optimization problem:\n\n~~~python\nmodel.optimization(params_min=lb, params_max=ub, x0_min=xmin, x0_max=xmax, ...)\n~~~\n\nwhere `lb` and `ub` are lists of arrays with the same structure as `model.params`, while `xmin` and `xmax` are arrays of the same dimension `model.nx` of the state vector. By default, each value is set equal to `None`, i.e., the corresponding constraint is not enforced. See `example_linear_positive.py` for examples of how to use nonnegative constraints to fit a positive linear system.\n\n<a name=\"static\"></a>\n### Static models and nonlinear regression / classification\nThe same optimization algorithms used to train dynamical models can be used to train static models, i.e., to solve the nonlinear regression problem:\n\n$$ \\min_{z}r(z)+\\frac{1}{N}\\sum_{k=0}^{N-1} \\|y_{k}-f(u_k,\\theta)\\|_2^2$$\n\nwhere $z=\\theta$ is the vector of model parameters to train and $r(z)$ admits the same\nregularization terms as in the case of dynamical models.\n\nFor example, if the model is a shallow neural network you can use the following code:\n\n~~~python\nfrom jax_sysid.models import StaticModel\nfrom jax_sysid.utils import standard_scale, unscale\n\n@jax.jit\ndef output_fcn(u, params):\n W1,b1,W2,b2=params\n y = W1@u.T+b1\n y = W2@jnp.arctan(y)+b2\n return y.T\nmodel = StaticModel(ny, nu, output_fcn)\nnn=10 # number of neurons\nmodel.init(params=[np.random.randn(nn,nu), np.random.randn(nn,1), np.random.randn(1,nn), np.random.randn(1,1)])\nmodel.loss(rho_th=1.e-4, tau_th=tau_th) \nmodel.optimization(lbfgs_epochs=500) \nmodel.fit(Ys, Us)\n~~~\n\n**jax-sysid** also supports feedforward neural networks defined via the **flax.linen** library:\n\n~~~python\nfrom jax_sysid.models import FNN\nfrom flax import linen as nn \n\n# output function\nclass FY(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(features=20)(x)\n x = nn.tanh(x)\n x = nn.Dense(features=20)(x)\n x = nn.tanh(x)\n x = nn.Dense(features=ny)(x)\n return x\n \nmodel = FNN(ny, nu, FY)\nmodel.loss(rho_th=1.e-4, tau_th=tau_th)\nmodel.optimization(lbfgs_epochs=500)\nmodel.fit(Ys, Us)\n~~~\n\nTo include lower and upper bounds on the parameters of the model, use the following additional arguments when specifying the optimization problem:\n\n~~~python\nmodel.optimization(lbfgs_epochs=500, params_min=lb, params_max=ub)\n~~~\n\nwhere `lb` and `ub` are lists of arrays with the same structure as `model.params`. See `example_static_convex.py` for examples of how to use nonnegative constraints to fit input-convex neural networks.\n\nTo solve classification problems, you need to define a custom loss function to change the default Mean-Squared-Error loss. For example, to train a classifier for a multi-category classification problem with $K$ classes, you can specify a neural network with a linear output layer generating output predictions $\\hat y\\in R^K$ and define the associated cross-entropy $\\ell(\\hat y,y) = -\\sum_{k=1}^Ky_k\\log\\left(\\frac{e^{\\hat y_k}}{\\sum_{j=1}^Ke^{\\hat y_j}}\\right)$ function as follows: \n\n~~~python\ndef cross_entropy(Yhat,Y):\n return -jax.numpy.sum(jax.nn.log_softmax(Yhat, axis=1)*Y)/Y.shape[0] \nmodel.loss(rho_th=1.e-4, output_loss=cross_entropy)\n~~~\n\nSee `example_static_fashion_mist.py` for an example using **Keras** with JAX backend to define the neural network model.\n \n<a name=\"contributors\"><a>\n## Contributors\n\nThis package was coded by Alberto Bemporad.\n\n\nThis software is distributed without any warranty. Please cite the paper below if you use this software.\n\n<a name=\"bibliography\"><a>\n## Citing jax-sysid\n\n<a name=\"ref1\"></a>\n\n```\n@article{Bem24,\n author={A. Bemporad},\n title={Linear and nonlinear system identification under $\\ell_1$- and group-{Lasso} regularization via {L-BFGS-B}},\n note = {submitted for publication. Also available on arXiv\n at \\url{http://arxiv.org/abs/2403.03827}},\n year=2024\n}\n```\n\n<a name=\"license\"><a>\n## License\n\nApache 2.0\n\n(C) 2024 A. Bemporad\n",
"bugtrack_url": null,
"license": null,
"summary": "jax-sysid - A Python package for linear and nonlinear system identification and nonlinear regression using Jax.",
"version": "0.4.2",
"project_urls": {
"Homepage": "http://cse.lab.imtlucca.it/~bemporad/jax-sysid"
},
"split_keywords": [
"system identification",
" subspace identification",
" nonlinear regression"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "98931354c96a5a897ce914f89046f614cc89aac163551cfc4959521a4d946ddd",
"md5": "7b0068477bc1dcccd344da35dd0e3490",
"sha256": "4b79ee86cc57c4dd659da2ded23d108d0f2c5e24152938aced12a9c6b27b4ff5"
},
"downloads": -1,
"filename": "jax_sysid-0.4.2-py3-none-any.whl",
"has_sig": false,
"md5_digest": "7b0068477bc1dcccd344da35dd0e3490",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 28692,
"upload_time": "2024-05-07T20:28:32",
"upload_time_iso_8601": "2024-05-07T20:28:32.053917Z",
"url": "https://files.pythonhosted.org/packages/98/93/1354c96a5a897ce914f89046f614cc89aac163551cfc4959521a4d946ddd/jax_sysid-0.4.2-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "7c7c58f0194f35eb69f842e9250ea3e3c37d1f1acc35e5f3bb1d173480fd2e3b",
"md5": "b6a34cee39c2896263dc3be21d226929",
"sha256": "a07abba69bc08fb989ec75c7b71aa26db273f0dacfa6c595c2c68436d8529973"
},
"downloads": -1,
"filename": "jax_sysid-0.4.2.tar.gz",
"has_sig": false,
"md5_digest": "b6a34cee39c2896263dc3be21d226929",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 34150,
"upload_time": "2024-05-07T20:28:34",
"upload_time_iso_8601": "2024-05-07T20:28:34.299246Z",
"url": "https://files.pythonhosted.org/packages/7c/7c/58f0194f35eb69f842e9250ea3e3c37d1f1acc35e5f3bb1d173480fd2e3b/jax_sysid-0.4.2.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-05-07 20:28:34",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "jax-sysid"
}
JSON