freq-statespace 0.1.2

Frequency-domain (non)linear state-space identification using JAX.

freq-statespace

A flexible JAX-based package for nonlinear state-space identification using frequency-domain optimization techniques, focusing on the nonlinear Linear Fractional Representation (NL-LFR) model structure, which combines an LTI system with a static feedback nonlinearity. This internal feedback formulation allows for capturing many complex real-world dynamics, with other popular block-oriented structures such as Wiener, Hammerstein, and Wiener-Hammerstein models arising as special cases. Identification of standard linear state-space models is also supported.

Basic usage

The package works with (multiple periods and realizations of) input–output data sequences $u(n)$ and $y(n)$ for $n = 0, \ldots, N-1$, assuming periodic excitation and an integer number of steady-state output periods. The specific NL-LFR structure is defined as:

  \begin{align*}
    x(n+1) &= A x(n) + B_u u(n) + B_w w(n),\\
    y(n) &= C_y x(n) + D_{yu} u(n) + D_{yw} w(n),\\
    z(n) &= C_z x(n) + D_{zu} u(n),\\ 
    w(n) &= f\big(z(n)\big),
  \end{align*}

consisting of linear state-space matrices and a static nonlinear function approximator $f(\cdot)$.

A typical step-wise identification procedure is as follows:

1. Best Linear Approximation (BLA) parametrization. Initializes the matrices $A$, $B_u$, $C_y$ and $D_{yu}$ using the frequency-domain subspace method, and refines these estimates through iterative optimization. If you're only interested in linear state-space models, you can stop the identification process here.
2. NL-LFR initialization. Applies the frequency-domain inference and learning method to efficiently initialize the remaining model parameters while keeping the BLA parameters fixed. This step requires that $f(\cdot)$ is linear in the parameters, i.e., $f(\cdot)=\beta^\top\phi(\cdot)$, with $\beta$ the parameter vector and $\phi(\cdot)$ the nonlinear feature mapping (e.g., polynomial features).
3. NL-LFR optimization. Performs iterative refinement of all model parameters using time-domain simulations. This is the most computationally demanding step, mainly due to the sequential nature of the forward simulations. Fortunately, the previous steps should have provided an initialization that is already close to a good local minimum.

It is also possible to skip the inference and learning step and go straight to nonlinear optimization. An advantage of this approach is that it puts no restriction on the structure of $f(\cdot)$, i.e., it does not require a model that is linear in its parameters.

Features

• Provides a user-friendly interface for identifying linear state-space models using frequency-domain subspace estimation based on the nonparametric BLA.
• Offers two workflows for identifying nonlinear LFR state-space models by primarily exploiting a frequency-domain formulation that enables inherent parallelism.
• Uses JAX for automatic differentiation, JIT compilation, and GPU/TPU acceleration.
• Supports Optimistix solvers (Levenberg-Marquardt, BFGS, ...) for typical system identification problems.
• Supports Optax optimizers (Adam, SGD, ...) for large-scale optimization.

Installation

Requires Python 3.10 or newer:

pip install freq-statespace

If JAX isn't already installed in your environment, the above command will install the CPU-only version. For GPU/TPU support (strongly recommended, often many times faster for mid-size to large problems), follow the JAX installation guide.

Quick example

We show an exemplary training pipeline on the Silverbox benchmark dataset, containing input-output measurements from an electronic circuit that mimics a mass-spring-damper system with a cubic spring nonlinearity.

We first estimate the BLA:

import freq_statespace as fss

data = fss.load_and_preprocess_silverbox_data()  # 8192 x 6 samples

# Step 1: BLA estimation
nx = 2  # state dimension
bla = fss.lin.subspace_id(data, nx)  # NRMSE 18.36%, non-iterative
bla = fss.lin.optimize(bla, data)  # NRMSE 13.17%, 6 iters, 1.32ms/iter

Next, we proceed with inference and learning, followed by full nonlinear optimization:

# Step 2: Inference and learning
nw, nz = 1, 1  # internal signal dimensions
phi = fss.static.basis.Polynomial(nz, degree=3)
nllfr = fss.nonlin.inference_and_learning(bla, data, phi, nw)  # NRMSE 1.11%, 45 iters, 18.4ms/iter

# Step 3: Nonlinear optimization
nllfr = fss.nonlin.optimize(nllfr, data)  # NRMSE 0.44%, 100 iters, 387ms/iter

Alternatively, we could skip inference and learning and jump straight to nonlinear optimization. In this example we use a neural network:

import jax

# Step 2: Nonlinear optimization
nw, nz = 1, 1  # internal signal dimensions
neural_net = fss.static.NeuralNetwork(nw, nz, layers=1, neurons_per_layer=10, activation=jax.nn.relu)
nllfr = fss.nonlin.connect(bla, neural_net)
nllfr = fss.nonlin.optimize(nllfr, data)  # NRMSE 0.55%, 100 iters, 354ms/iter

Note: Iteration timings were measured on an NVIDIA T600 Laptop GPU.

Serialization of models can be achieved like so:

path = "models/nllfr.zip"
fss.save_model(nllfr, path)
nllfr_loaded = fss.load_model(path)

The examples/ folder also provides Jupyter notebooks for more challenging benchmark systems, with additional notes on hyperparameter tuning and solver configurations.

Preparing your data

Every identification problem starts by casting the time-domain input-output data into the required format and supplying minimal frequency metadata (the excited frequencies and the sampling frequency); the helper function fss.create_data_object(...) is provided for this purpose. Once the data object is instantiated, the workflow proceeds exactly as in the quick example above.

Citation

If you use this code in your work, please cite it as (arXiv link):

@article{floren2025inference,
  title={Inference and Learning of Nonlinear LFR State-Space Models},
  author={Floren, Merijn and No{\"e}l, Jean-Philippe and Swevers, Jan},
  journal={IEEE Control Systems Letters},
  year={2025},
  publisher={IEEE}
}