neural-ssm 0.1.7__tar.gz

neural_ssm-0.1.7/PKG-INFO

@@ -0,0 +1,75 @@
+Metadata-Version: 2.4
+Name: neural-ssm
+Version: 0.1.7
+Summary: Neural state space models and LRU variants in PyTorch
+Author: Leonardo Massai
+License: MIT
+Project-URL: Homepage, https://github.com/LeoMassai
+Project-URL: Repository, https://github.com/LeoMassai/neural-ssm
+Project-URL: Issues, https://github.com/LeoMassai/neural-ssm/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.2
+Requires-Dist: numpy>=1.23
+Requires-Dist: scipy>=1.10
+Requires-Dist: h5py>=3.9
+Requires-Dist: tqdm>=4.66
+Requires-Dist: matplotlib>=3.7
+Requires-Dist: einops>=0.6
+Requires-Dist: typing-extensions>=4.5
+Requires-Dist: jax>=0.4.26
+Requires-Dist: jaxlib>=0.4.26
+Requires-Dist: deel-torchlip
+Provides-Extra: dev
+Requires-Dist: black>=24.3.0; extra == "dev"
+Requires-Dist: ruff>=0.3.0; extra == "dev"
+Requires-Dist: mypy>=1.6.0; extra == "dev"
+Requires-Dist: pytest>=7.4; extra == "dev"
+Requires-Dist: pytest-cov>=4.1; extra == "dev"
+Requires-Dist: ipykernel>=6.29; extra == "dev"
+Requires-Dist: jupyter>=1.0; extra == "dev"
+Provides-Extra: examples
+Requires-Dist: seaborn>=0.13; extra == "examples"
+Requires-Dist: pandas>=2.1; extra == "examples"
+Provides-Extra: docs
+Requires-Dist: mkdocs>=1.5; extra == "docs"
+Requires-Dist: mkdocs-material>=9.5; extra == "docs"
+
+# Pytorch L2RU Architecture: LRU with l2 stability guarantees and prescribed bound
+
+A PyTorch implementation of the L2RU architecture introduced in the paper Free Parametrization of L2-bounded State Space Models. https://arxiv.org/abs/2503.23818. Application in System Identification is included as an example.
+
+## L2RU block
+The L2RU block is a discrete-time linear time-invariant system implemented in state-space form as:
+```math
+\begin{align}
+x_{k+1} = Ax_{x} + B u_k\\
+y_k = C x_k + D u_k,
+\end{align}
+```
+A parametrization is provided for the matrices ```(A, B, C, D)```, guaranteeing a prescribed l2 bound for the overall SSM.
+Moreover, the use of [parallel scan algorithms](https://en.wikipedia.org/wiki/Prefix_sum) makes execution extremely fast on modern hardware in non-core-bound scenarios.
+
+## Deep L2RU Architecture
+
+L2RU units are typically organized in a deep LRU architecture like:
+
+<div align="center">
+  <img src="architecture/L2RU.png" alt="Description of image" width="800">
+</div>
+
+
+
+
+
+

neural_ssm-0.1.7/README.md

@@ -0,0 +1,28 @@
+# Pytorch L2RU Architecture: LRU with l2 stability guarantees and prescribed bound
+
+A PyTorch implementation of the L2RU architecture introduced in the paper Free Parametrization of L2-bounded State Space Models. https://arxiv.org/abs/2503.23818. Application in System Identification is included as an example.
+
+## L2RU block
+The L2RU block is a discrete-time linear time-invariant system implemented in state-space form as:
+```math
+\begin{align}
+x_{k+1} = Ax_{x} + B u_k\\
+y_k = C x_k + D u_k,
+\end{align}
+```
+A parametrization is provided for the matrices ```(A, B, C, D)```, guaranteeing a prescribed l2 bound for the overall SSM.
+Moreover, the use of [parallel scan algorithms](https://en.wikipedia.org/wiki/Prefix_sum) makes execution extremely fast on modern hardware in non-core-bound scenarios.
+
+## Deep L2RU Architecture
+
+L2RU units are typically organized in a deep LRU architecture like:
+
+<div align="center">
+  <img src="architecture/L2RU.png" alt="Description of image" width="800">
+</div>
+
+
+
+
+
+

neural_ssm-0.1.7/pyproject.toml

@@ -0,0 +1,107 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "neural-ssm"
+version = "0.1.7"
+description = "Neural state space models and LRU variants in PyTorch"
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [
+  { name = "Leonardo Massai" }
+]
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+  "Intended Audience :: Science/Research",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+  # Core
+  "torch>=2.2",
+  "numpy>=1.23",
+  # SciPy for linear algebra and scipy.io.loadmat (MAT files in Test_files)
+  "scipy>=1.10",
+  # Optional MAT v7.3 support via HDF5 (some .mat files may require it)
+  "h5py>=3.9",
+  # Utilities commonly used in training/eval and plotting
+  "tqdm>=4.66",
+  "matplotlib>=3.7",
+  # Tensor/array rearrangements if used in layers
+  "einops>=0.6",
+  # Backport typing if running older Python or for PyTorch type hints
+  "typing-extensions>=4.5",
+      # JAX (CPU by default; GPU wheels are user specific)
+  "jax>=0.4.26",
+  "jaxlib>=0.4.26",
+  "deel-torchlip",
+]
+
+[project.optional-dependencies]
+dev = [
+  "black>=24.3.0",
+  "ruff>=0.3.0",
+  "mypy>=1.6.0",
+  "pytest>=7.4",
+  "pytest-cov>=4.1",
+  "ipykernel>=6.29",
+  "jupyter>=1.0",
+]
+examples = [
+  "seaborn>=0.13",
+  "pandas>=2.1",
+]
+docs = [
+  "mkdocs>=1.5",
+  "mkdocs-material>=9.5",
+]
+
+[project.urls]
+Homepage = "https://github.com/LeoMassai"
+Repository = "https://github.com/LeoMassai/neural-ssm"
+Issues = "https://github.com/LeoMassai/neural-ssm/issues"
+
+[tool.setuptools]
+include-package-data = true
+license-files = ["LICEN[CS]E*"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["neural_ssm*"]
+
+[tool.pytest.ini_options]
+addopts = "-q"
+testpaths = ["tests", "Test_files"]
+
+[tool.black]
+line-length = 100
+target-version = ["py39", "py310", "py311", "py312"]
+include = "\\.pyi?$"
+
+[tool.ruff]
+line-length = 100
+target-version = "py39"
+src = ["src", "tests", "Test_files"]
+extend-exclude = ["build", "dist", ".venv", ".mypy_cache", ".ruff_cache"]
+select = ["E", "F", "I", "UP", "B"]
+ignore = []
+
+[tool.mypy]
+python_version = "3.9"
+ignore_missing_imports = true
+warn_unused_ignores = true
+warn_return_any = true
+warn_unused_configs = true
+no_implicit_optional = true
+strict_equality = true
+show_error_codes = true
+pretty = true
+mypy_path = ["src"]

neural_ssm-0.1.7/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

neural_ssm-0.1.7/src/neural_ssm/__init__.py

@@ -0,0 +1,42 @@
+# python
+# file: src/neural_ssm/__init__.py
+from importlib import import_module as _imp
+
+# Re-export subpackages for discoverability
+from . import ssm as ssm
+from . import rens as rens
+from . import static_layers as layers  # public alias
+
+# Top-level classes and configs
+from .ssm.lru import LRU, L2RU, lruz, SSMConfig, SSL, DeepSSM, PureLRUR
+from .rens.ren import REN
+
+# Common layers exposed at top-level for convenience
+try:
+    from .static_layers.generic_layers import LayerConfig, GLU, MLP, TLIP
+except Exception:
+    pass
+try:
+    from .static_layers.lipschitz_mlps import LMLP
+except Exception:
+    pass
+
+__all__ = [
+    "LRU", "L2RU", "lruz", "SSMConfig", "SSL", "DeepSSM", "PureLRUR",
+    "REN",
+    "layers", "ssm", "rens",
+    "LayerConfig", "GLU", "MLP", "LMLP", "TLIP",
+]
+
+__version__ = "0.1.0"
+
+def __getattr__(name):
+    # Optional lazy/compat shims; keep internals movable
+    redirects = {
+        "layers": "neural_ssm.static_layers",
+    }
+    if name in redirects:
+        return _imp(redirects[name])
+    raise AttributeError(f"module neural_ssm has no attribute {name!r}")
+
+del _imp

neural_ssm-0.1.7/src/neural_ssm/rens/__init__.py

@@ -0,0 +1,5 @@
+# python
+# file: src/neural_ssm/rens/__init__.py
+from .ren import REN
+
+__all__ = ["REN"]

neural_ssm-0.1.7/src/neural_ssm/rens/ren.py

@@ -0,0 +1,210 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+
+# Robust REN implementation in the acyclic version
+class REN(nn.Module):
+    # ## Implementation of REN model, modified from "Recurrent Equilibrium Networks: Flexible Dynamic Models with
+    # Guaranteed Stability and Robustness" by Max Revay et al.
+    def __init__(self, dim_in: int, dim_out: int, dim_internal: int,
+                 dim_nl: int, initialization_std: float = 0.5, internal_state_init=None, gammat=None, mode="l2stable",
+                 Q=None, R=None, S=None
+                 , posdef_tol: float = 0.001):
+        super().__init__()
+
+        # set dimensions
+        self.dim_in = dim_in  # input dimension m
+        self.dim_internal = dim_internal  # state dimension n
+        self.dim_nl = dim_nl  # dimension of v(t) and w(t) l
+        self.dim_out = dim_out  # output dimension p
+
+        self.mode = mode
+        self.gammat = gammat
+        self.epsilon = posdef_tol
+        # # # # # # # # # IQC specification # # # # # # # # #
+        self.Q = Q
+        self.R = R
+        self.S = S
+        # # # # # # # # # Training parameters # # # # # # # # #
+        # Sparse training matrix parameters
+        # define matrices shapes
+        self.X_shape = (2 * dim_internal + dim_nl, 2 * dim_internal + dim_nl)
+        self.Y_shape = (dim_internal, dim_internal)
+        # nn state dynamics
+        self.B2_shape = (dim_internal, dim_in)
+        # nn output
+        self.C2_shape = (dim_out, dim_internal)
+        #self.D21_shape = (self.dim_out, self.dim_l)
+        self.D22_shape = (dim_out, dim_in)
+        # v signal
+        self.D12_shape = (dim_nl, dim_in)
+        self.Z3_shape = (abs(dim_out - dim_in), min(dim_out, dim_in))
+        self.X3_shape = (min(dim_out, dim_in), min(dim_out, dim_in))
+        self.Y3_shape = (min(dim_out, dim_in), min(dim_out, dim_in))
+        self.gamma_shape = (1, 1)
+
+        self.training_param_names = ['X', 'Y', 'B2', 'C2', 'Z3', 'X3', 'Y3', 'D12']
+
+        # Optionally define a trainable gamma
+        if self.gammat is None:
+            self.training_param_names.append('gamma')
+        else:
+            self.gamma = gammat
+
+        # define trainable params
+
+        self._init_trainable_params(initialization_std)
+
+        # # # # # # # # # Non-trainable parameters and constant tensors # # # # # # # # #
+        # masks
+        self.register_buffer('eye_mask_min', torch.eye(min(dim_in, dim_out)))
+        self.register_buffer('eye_mask_dim_in', torch.eye(dim_in))
+        self.register_buffer('eye_mask_dim_out', torch.eye(dim_out))
+        self.register_buffer('eye_mask_dim_state', torch.eye(dim_internal))
+        self.register_buffer('eye_mask_H', torch.eye(2 * dim_internal + dim_nl))
+        self.register_buffer('zeros_mask_S', torch.zeros(dim_in, dim_out))
+        self.register_buffer('zeros_mask_Q', torch.zeros(dim_out, dim_out))
+        self.register_buffer('zeros_mask_R', torch.zeros(dim_in, dim_in))
+        self.register_buffer('zeros_mask_so', torch.zeros(dim_internal, dim_out))
+        self.register_buffer('eye_mask_w', torch.eye(dim_nl))
+        self.register_buffer('D21', torch.zeros(dim_out, dim_nl))
+
+        # initialize internal state
+        if internal_state_init is None:
+            self.x = torch.zeros(1, 1, self.dim_internal, device="cuda")
+        else:
+            assert isinstance(internal_state_init, torch.Tensor)
+            self.x = internal_state_init.reshape(1, 1, self.dim_internal)
+        self.register_buffer('init_x', self.x.detach().clone())
+
+        # Auxiliary elements
+        self.set_param()
+
+    def set_param(self, gamman=None):
+        if gamman is not None:
+            self.gamma = gamman
+        gamma = torch.abs(self.gamma)
+        dim_internal, dim_nl, dim_in, dim_out = self.dim_internal, self.dim_nl, self.dim_in, self.dim_out
+
+        # Updating of Q,S,R with variable gamma if needed
+        self.Q, self.R, self.S = self._set_mode(self.mode, gamma, self.Q, self.R, self.S)
+        M = F.linear(self.X3.T, self.X3.T) + self.Y3 - self.Y3.T + F.linear(self.Z3.T,
+                                                                            self.Z3.T) + self.epsilon * self.eye_mask_min
+        if dim_out >= dim_in:
+            N = torch.vstack((F.linear(self.eye_mask_dim_in - M,
+                                       torch.inverse(self.eye_mask_dim_in + M).T),
+                              -2 * F.linear(self.Z3, torch.inverse(self.eye_mask_dim_in + M).T)))
+        else:
+            N = torch.hstack((F.linear(torch.inverse(self.eye_mask_dim_out + M),
+                                       (self.eye_mask_dim_out - M).T),
+                              -2 * F.linear(torch.inverse(self.eye_mask_dim_out + M), self.Z3)))
+
+        Lq = torch.linalg.cholesky(-self.Q).T
+        Lr = torch.linalg.cholesky(self.R - torch.matmul(self.S, torch.matmul(torch.inverse(self.Q), self.S.T))).T
+        self.D22 = -torch.matmul(torch.inverse(self.Q), self.S.T) + torch.matmul(torch.inverse(Lq),
+                                                                                 torch.matmul(N, Lr))
+        # Calculate psi_r:
+        R_cal = self.R + torch.matmul(self.S, self.D22) + torch.matmul(self.S, self.D22).T + torch.matmul(self.D22.T,
+                                                                                                          torch.matmul(
+                                                                                                              self.Q,
+                                                                                                              self.D22))
+        R_cal_inv = torch.inverse(R_cal)
+        C2_cal = torch.matmul(torch.matmul(self.D22.T, self.Q) + self.S, self.C2)
+        D21_cal = torch.matmul(torch.matmul(self.D22.T, self.Q) + self.S, self.D21) - self.D12.T
+        vec_r = torch.cat((C2_cal.T, D21_cal.T, self.B2), dim=0)
+        psi_r = torch.matmul(vec_r, torch.matmul(R_cal_inv, vec_r.T))
+        # Calculate psi_q:
+        vec_q = torch.cat((self.C2.T, self.D21.T, self.zeros_mask_so), dim=0)
+        psi_q = torch.matmul(vec_q, torch.matmul(self.Q, vec_q.T))
+        # Create H matrix:
+        H = torch.matmul(self.X.T, self.X) + self.epsilon * self.eye_mask_H + psi_r - psi_q
+        h1, h2, h3 = torch.split(H, [dim_internal, dim_nl, dim_internal], dim=0)
+        H11, H12, H13 = torch.split(h1, [dim_internal, dim_nl, dim_internal], dim=1)
+        H21, H22, _ = torch.split(h2, [dim_internal, dim_nl, dim_internal], dim=1)
+        H31, H32, H33 = torch.split(h3, [dim_internal, dim_nl, dim_internal], dim=1)
+        self.P_cal = H33
+        # NN state dynamics:
+        self.F = H31
+        self.B1 = H32
+        # NN output:
+        self.E = 0.5 * (H11 + self.P_cal + self.Y - self.Y.T)
+        # v signal:  [Change the following 2 lines if we don't want a strictly acyclic REN!]
+        self.Lambda = 0.5 * torch.diag(H22)
+        self.D11 = -torch.tril(H22, diagonal=-1)
+        self.C1 = -H21
+        # Matrix P
+        #self.P = torch.matmul(self.E.T, torch.matmul(torch.inverse(self.P_cal), self.E))
+
+    def forward(self, u):
+        decay_rate = 0.95
+        batch_size = u.shape[0]
+        w = torch.zeros(batch_size, 1, self.dim_nl, device=u.device)
+        # update each row of w using Eq. (8) with a lower triangular D11
+        for i in range(self.dim_nl):
+            #  v is element i of v with dim (batch_size, 1)
+            v = F.linear(self.x, self.C1[i, :]) + F.linear(w, self.D11[i, :]) + F.linear(u, self.D12[i, :])
+            w = w + (self.eye_mask_w[i, :] * torch.tanh(v / self.Lambda[i])).reshape(batch_size, 1, self.dim_nl)
+
+        # compute next state using Eq. 18
+        self.x = F.linear(
+            F.linear(self.x, self.F) + F.linear(w, self.B1) + F.linear(u, self.B2),
+            self.E.inverse())
+
+        # compute output
+        y = F.linear(self.x, self.C2) + F.linear(w, self.D21) + F.linear(u, self.D22)
+
+        return y
+
+    def _set_mode(self, mode, gamma, Q, R, S, eps: float = 1e-4):
+        # We set Q to be negative definite. If Q is nsd we set: Q - \epsilon I.
+        # I.e. The Q we define here is denoted as \matcal{Q} in REN paper.
+        if mode == "l2stable":
+            Q = -(1. / gamma) * self.eye_mask_dim_out
+            R = gamma * self.eye_mask_dim_in
+            S = self.zeros_mask_S
+        elif mode == "input_p":
+            if self.p != self.m:
+                raise NameError("Dimensions of u(t) and y(t) need to be the same for enforcing input passivity.")
+            Q = self.zeros_mask_Q - eps * self.eye_mask_dim_out
+            R = -2. * gamma * self.eye_mask_dim_state
+            S = self.eye_mask_dim_out
+        elif mode == "output_p":
+            if self.p != self.m:
+                raise NameError("Dimensions of u(t) and y(t) need to be the same for enforcing output passivity.")
+            Q = -2. * gamma * self.eye_mask_dim_out
+            R = self.zeros_mask_R
+            S = self.eye_mask_dim_state
+        else:
+            print("Using matrices R,Q,S given by user.")
+            # Check dimensions:
+            if not (len(R.shape) == 2 and R.shape[0] == R.shape[1] and R.shape[0] == self.m):
+                raise NameError("The matrix R is not valid. It must be a square matrix of %ix%i." % (self.m, self.m))
+            if not (len(Q.shape) == 2 and Q.shape[0] == Q.shape[1] and Q.shape[0] == self.p):
+                raise NameError("The matrix Q is not valid. It must be a square matrix of %ix%i." % (self.p, self.p))
+            if not (len(S.shape) == 2 and S.shape[0] == self.m and S.shape[1] == self.p):
+                raise NameError("The matrix S is not valid. It must be a matrix of %ix%i." % (self.m, self.p))
+            # Check R=R':
+            if not (R == R.T).prod():
+                raise NameError("The matrix R is not valid. It must be symmetric.")
+            # Check Q is nsd:
+            eigs, _ = torch.linalg.eig(Q)
+            if not (eigs.real <= 0).prod():
+                print('oh!')
+                raise NameError("The matrix Q is not valid. It must be negative semidefinite.")
+            if not (eigs.real < 0).prod():
+                # We make Q negative definite: (\mathcal{Q} in the REN paper)
+                Q = Q - eps * self.eye_mask_dim_out
+        return Q, R, S
+
+        # init trainable params
+
+    def _init_trainable_params(self, initialization_std):
+        for training_param_name in self.training_param_names:  # name of one of the training params, e.g., X
+            # read the defined shapes of the selected training param, e.g., X_shape
+            shape = getattr(self, training_param_name + '_shape')
+            # define the selected param (e.g., self.X) as nn.Parameter
+            if training_param_name == 'gamma':
+                initialization_std = 3
+            setattr(self, training_param_name, nn.Parameter((torch.randn(*shape) * initialization_std)))

neural_ssm-0.1.7/src/neural_ssm/ssm/__init__.py

@@ -0,0 +1,5 @@
+# python
+# file: src/neural_ssm/ssm/__init__.py
+from .lru import LRU, L2RU, lruz, SSMConfig, SSL, DeepSSM, PureLRUR
+
+__all__ = ["LRU", "L2RU", "lruz", "SSMConfig", "SSL", "DeepSSM", "PureLRUR"]