compreSSM 1.0.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: compreSSM
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: Library built off of the compreSSM paper linked in the README
|
|
5
|
+
Author: Christian Rafferty
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/cfr081709/compreSSM_library
|
|
8
|
+
Project-URL: Repository, https://github.com/cfr081709/compreSSM_library
|
|
9
|
+
Project-URL: Issues, https://github.com/cfr081709/compreSSM_library/issues
|
|
10
|
+
Requires-Python: >=3.10
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
License-File: LICENSE
|
|
13
|
+
Requires-Dist: numpy>=1.24
|
|
14
|
+
Requires-Dist: pydantic>=2.0
|
|
15
|
+
Requires-Dist: scipy>=1.10
|
|
16
|
+
Requires-Dist: torch>=2.0
|
|
17
|
+
Provides-Extra: test
|
|
18
|
+
Requires-Dist: pytest>=7.0; extra == "test"
|
|
19
|
+
Dynamic: license-file
|
|
20
|
+
|
|
21
|
+
# CompreSSM Library
|
|
22
|
+
|
|
23
|
+
CompreSSM provides balanced-truncation tools for compressing stable discrete-time
|
|
24
|
+
state-space models, including diagonal systems used by Linear Recurrent Units
|
|
25
|
+
(LRUs). It computes controllability and observability Gramians, Hankel singular
|
|
26
|
+
values, reduced systems, and an error bound for each reduction.
|
|
27
|
+
|
|
28
|
+
## Installation
|
|
29
|
+
|
|
30
|
+
```bash
|
|
31
|
+
python -m pip install compreSSM-library
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
The package requires Python 3.10 or later. PyTorch is a runtime dependency
|
|
35
|
+
because the module also provides helpers for extracting and updating PyTorch
|
|
36
|
+
modules.
|
|
37
|
+
|
|
38
|
+
## Quick start
|
|
39
|
+
|
|
40
|
+
The published module is named `core`:
|
|
41
|
+
|
|
42
|
+
```python
|
|
43
|
+
import numpy as np
|
|
44
|
+
|
|
45
|
+
from core import (
|
|
46
|
+
CompressionConfig,
|
|
47
|
+
DiagonalDiscreteSSM,
|
|
48
|
+
compress_diagonal_numpy,
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
system = DiagonalDiscreteSSM.from_numpy(
|
|
52
|
+
lam=np.array([0.5, 0.8], dtype=complex),
|
|
53
|
+
B=np.array([[1.0], [0.01]], dtype=complex),
|
|
54
|
+
C=np.array([[1.0, 0.01]], dtype=complex),
|
|
55
|
+
D=np.array([0.0]),
|
|
56
|
+
)
|
|
57
|
+
|
|
58
|
+
reduced = compress_diagonal_numpy(
|
|
59
|
+
system,
|
|
60
|
+
config=CompressionConfig(tol=0.05, max_ratio=0.95),
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
if reduced is None:
|
|
64
|
+
print("The requested tolerance does not justify reducing this system.")
|
|
65
|
+
else:
|
|
66
|
+
print(f"Reduced state dimension: {reduced.r} (from {system.n})")
|
|
67
|
+
print(f"Balanced-truncation error bound: {reduced.error_bound}")
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
`tol` is the fraction of total Hankel singular value mass that may be discarded.
|
|
71
|
+
`max_ratio` prevents accepting a reduction when the resulting state count is not
|
|
72
|
+
small enough compared with the original system.
|
|
73
|
+
|
|
74
|
+
## Main API
|
|
75
|
+
|
|
76
|
+
- `DiscreteSSM` and `compress_numpy` handle general stable discrete-time systems.
|
|
77
|
+
- `DiagonalDiscreteSSM` and `compress_diagonal_numpy` handle diagonal systems,
|
|
78
|
+
including LRU parameterizations.
|
|
79
|
+
- `CompressionConfig` controls the sample time, retained Hankel singular value
|
|
80
|
+
mass, and minimum worthwhile reduction.
|
|
81
|
+
- `compress_module` and `compress_lru_module` adapt compatible PyTorch modules.
|
|
82
|
+
- `run_scheduled_reduction` applies reduction to blocks on a training schedule.
|
|
83
|
+
|
|
84
|
+
The `D` (direct-feedthrough) matrix is retained unchanged during reduction. The
|
|
85
|
+
reported balanced-truncation error bound assumes a stable, minimal system.
|
|
86
|
+
|
|
87
|
+
For details on every class and function, see the [full API documentation](docs.md).
|
|
88
|
+
|
|
89
|
+
## Development
|
|
90
|
+
|
|
91
|
+
Install the package and test dependencies, then run the test suite:
|
|
92
|
+
|
|
93
|
+
```bash
|
|
94
|
+
python -m pip install ".[test]"
|
|
95
|
+
python -m pytest
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
## Citation
|
|
99
|
+
|
|
100
|
+
This library is based on:
|
|
101
|
+
|
|
102
|
+
> Makram Chahine, Philipp Nazari, Daniela Rus, and T. Konstantin Rusch.
|
|
103
|
+
> “The Curious Case of In-Training Compression of State Space Models.” 2026.
|
|
104
|
+
> [arXiv:2510.02823](https://arxiv.org/abs/2510.02823).
|
|
105
|
+
|
|
106
|
+
```bibtex
|
|
107
|
+
@misc{chahine2026curiouscaseintrainingcompression,
|
|
108
|
+
title={The Curious Case of In-Training Compression of State Space Models},
|
|
109
|
+
author={Makram Chahine and Philipp Nazari and Daniela Rus and T. Konstantin Rusch},
|
|
110
|
+
year={2026},
|
|
111
|
+
eprint={2510.02823},
|
|
112
|
+
archivePrefix={arXiv},
|
|
113
|
+
primaryClass={cs.LG},
|
|
114
|
+
url={https://arxiv.org/abs/2510.02823}
|
|
115
|
+
}
|
|
116
|
+
```
|
|
@@ -0,0 +1,6 @@
|
|
|
1
|
+
core.py,sha256=1gse_gARzt9wgaEOqpDCXKukhXdQCzv8e3miMiSnafM,15127
|
|
2
|
+
compressm-1.0.0.dist-info/licenses/LICENSE,sha256=gYmCMmN-UsbtPaCxDWbUgpFREb5ydCJtE7bwJFO373o,1065
|
|
3
|
+
compressm-1.0.0.dist-info/METADATA,sha256=rD9kixU_nWwqqKpdleOtQsghW4r279t64l6reH0_tQI,3672
|
|
4
|
+
compressm-1.0.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
|
|
5
|
+
compressm-1.0.0.dist-info/top_level.txt,sha256=0ebp8k35dOTHVwbY3UcokRsjC7cWLN_iaJ70tIcmODU,5
|
|
6
|
+
compressm-1.0.0.dist-info/RECORD,,
|
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
Copyright 2026 Christian Rafferty
|
|
2
|
+
|
|
3
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
|
|
4
|
+
|
|
5
|
+
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
|
6
|
+
|
|
7
|
+
THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
core
|
core.py
ADDED
|
@@ -0,0 +1,396 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import torch
|
|
3
|
+
from collections.abc import Callable, Iterable
|
|
4
|
+
from dataclasses import dataclass
|
|
5
|
+
from torch import nn
|
|
6
|
+
from pydantic import BaseModel, Field
|
|
7
|
+
from scipy.linalg import solve_discrete_lyapunov, cholesky, svd
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class CompressionConfig(BaseModel):
|
|
11
|
+
dt: float = Field(default=1.0, gt=0.0)
|
|
12
|
+
tol: float = Field(
|
|
13
|
+
default=0.01,
|
|
14
|
+
ge=0.0,
|
|
15
|
+
le=1.0,
|
|
16
|
+
description="Fraction of the sum of Hankel singular values to discard",
|
|
17
|
+
)
|
|
18
|
+
max_ratio: float = Field(default=0.95, gt=0.0, le=1.0)
|
|
19
|
+
|
|
20
|
+
@property
|
|
21
|
+
def threshold(self) -> float:
|
|
22
|
+
"""Fraction of the sum of Hankel singular values to retain."""
|
|
23
|
+
return 1.0 - self.tol
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class ReductionSchedule(BaseModel):
|
|
27
|
+
red_start: int = Field(ge=0)
|
|
28
|
+
red_end: int = Field(ge=0)
|
|
29
|
+
red_interval: int = Field(gt=0)
|
|
30
|
+
|
|
31
|
+
def fires_at(self, step: int) -> bool:
|
|
32
|
+
return (
|
|
33
|
+
self.red_start <= step < self.red_end
|
|
34
|
+
and (step + 1) % self.red_interval == 0
|
|
35
|
+
)
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
@dataclass(frozen=True)
|
|
39
|
+
class DiscreteSSM:
|
|
40
|
+
A: np.ndarray
|
|
41
|
+
B: np.ndarray
|
|
42
|
+
C: np.ndarray
|
|
43
|
+
D: np.ndarray
|
|
44
|
+
|
|
45
|
+
@property
|
|
46
|
+
def n(self) -> int:
|
|
47
|
+
return self.A.shape[0]
|
|
48
|
+
|
|
49
|
+
@classmethod
|
|
50
|
+
def from_numpy(cls, *, A: np.ndarray, B: np.ndarray, C: np.ndarray, D: np.ndarray):
|
|
51
|
+
A, B, C, D = (np.asarray(matrix) for matrix in (A, B, C, D))
|
|
52
|
+
if A.ndim != 2 or A.shape[0] != A.shape[1] or A.shape[0] == 0:
|
|
53
|
+
raise ValueError("A must be a non-empty square matrix")
|
|
54
|
+
if B.ndim != 2 or B.shape[0] != A.shape[0]:
|
|
55
|
+
raise ValueError("B must be a matrix with the same number of rows as A")
|
|
56
|
+
if C.ndim != 2 or C.shape[1] != A.shape[0]:
|
|
57
|
+
raise ValueError("C must be a matrix with the same number of columns as A")
|
|
58
|
+
if not all(np.all(np.isfinite(matrix)) for matrix in (A, B, C, D)):
|
|
59
|
+
raise ValueError("State-space matrices must contain only finite values")
|
|
60
|
+
if np.max(np.abs(np.linalg.eigvals(A))) >= 1:
|
|
61
|
+
raise ValueError("A is not stable (spectral radius >= 1)")
|
|
62
|
+
return cls(A=A, B=B, C=C, D=D)
|
|
63
|
+
|
|
64
|
+
@classmethod
|
|
65
|
+
def from_continuous_diagonal_numpy(cls, *, a: np.ndarray, B: np.ndarray, C: np.ndarray, D: np.ndarray, dt: float):
|
|
66
|
+
"""Discretize a stable diagonal continuous-time system with ZOH input."""
|
|
67
|
+
a = np.asarray(a)
|
|
68
|
+
if a.ndim != 1 or a.size == 0:
|
|
69
|
+
raise ValueError("a must be a non-empty one-dimensional array")
|
|
70
|
+
if not np.isfinite(dt) or dt <= 0:
|
|
71
|
+
raise ValueError("dt must be finite and greater than zero")
|
|
72
|
+
if not np.all(np.isfinite(a)) or np.any(np.real(a) >= 0):
|
|
73
|
+
raise ValueError("Continuous-time diagonal rates must be finite and stable")
|
|
74
|
+
|
|
75
|
+
a_bar = np.exp(dt * a)
|
|
76
|
+
integration_factor = np.full(a.shape, dt, dtype=np.result_type(a, dt))
|
|
77
|
+
np.divide(np.expm1(dt * a), a, out=integration_factor, where=a != 0)
|
|
78
|
+
B_bar = integration_factor[:, None] * B
|
|
79
|
+
return cls.from_numpy(A=np.diag(a_bar), B=B_bar, C=C, D=D)
|
|
80
|
+
|
|
81
|
+
@classmethod
|
|
82
|
+
def from_tensors(cls, *, A_log: torch.Tensor, B: torch.Tensor, C: torch.Tensor, D: torch.Tensor, dt: float):
|
|
83
|
+
a = -np.exp(tensor_to_numpy(A_log))
|
|
84
|
+
n = a.shape[0]
|
|
85
|
+
return cls.from_continuous_diagonal_numpy(
|
|
86
|
+
a=a,
|
|
87
|
+
B=tensor_to_numpy(B).reshape(n, -1),
|
|
88
|
+
C=tensor_to_numpy(C).reshape(-1, n),
|
|
89
|
+
D=tensor_to_numpy(D),
|
|
90
|
+
dt=dt,
|
|
91
|
+
)
|
|
92
|
+
|
|
93
|
+
@classmethod
|
|
94
|
+
def from_module(cls, module: nn.Module, *, dt: float):
|
|
95
|
+
sd = module.state_dict()
|
|
96
|
+
return cls.from_tensors(
|
|
97
|
+
A_log=sd['ssm.A_log'], B=sd['ssm.B'], C=sd['ssm.C'], D=sd['ssm.D'], dt=dt
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
@dataclass(frozen=True)
|
|
102
|
+
class DiagonalDiscreteSSM:
|
|
103
|
+
lam: np.ndarray
|
|
104
|
+
B: np.ndarray
|
|
105
|
+
C: np.ndarray
|
|
106
|
+
D: np.ndarray
|
|
107
|
+
|
|
108
|
+
@property
|
|
109
|
+
def n(self) -> int:
|
|
110
|
+
return self.lam.shape[0]
|
|
111
|
+
|
|
112
|
+
@classmethod
|
|
113
|
+
def from_numpy(cls, *, lam: np.ndarray, B: np.ndarray, C: np.ndarray, D: np.ndarray):
|
|
114
|
+
lam, B, C, D = (np.asarray(matrix) for matrix in (lam, B, C, D))
|
|
115
|
+
if lam.ndim != 1 or lam.size == 0:
|
|
116
|
+
raise ValueError("lam must be a non-empty one-dimensional array")
|
|
117
|
+
if B.ndim != 2 or B.shape[0] != lam.shape[0]:
|
|
118
|
+
raise ValueError("B must be a matrix with one row per state")
|
|
119
|
+
if C.ndim != 2 or C.shape[1] != lam.shape[0]:
|
|
120
|
+
raise ValueError("C must be a matrix with one column per state")
|
|
121
|
+
if not all(np.all(np.isfinite(matrix)) for matrix in (lam, B, C, D)):
|
|
122
|
+
raise ValueError("State-space matrices must contain only finite values")
|
|
123
|
+
if np.max(np.abs(lam)) >= 1:
|
|
124
|
+
raise ValueError("lam is not stable (spectral radius >= 1)")
|
|
125
|
+
return cls(lam=lam, B=B, C=C, D=D)
|
|
126
|
+
|
|
127
|
+
@classmethod
|
|
128
|
+
def from_lru_numpy(cls, *, nu_log, theta_log, B_re, B_im, gamma_log, C_re, C_im, D):
|
|
129
|
+
lam = np.exp(-np.exp(nu_log) + 1j * np.exp(theta_log))
|
|
130
|
+
B = (B_re + 1j * B_im) * np.exp(gamma_log)[:, None]
|
|
131
|
+
C = C_re + 1j * C_im
|
|
132
|
+
return cls.from_numpy(lam=lam, B=B, C=C, D=D)
|
|
133
|
+
|
|
134
|
+
@classmethod
|
|
135
|
+
def from_lru_module(cls, module: nn.Module):
|
|
136
|
+
p = {k: tensor_to_numpy(v) for k, v in module.named_parameters()}
|
|
137
|
+
return cls.from_lru_numpy(
|
|
138
|
+
nu_log=p['nu_log'], theta_log=p['theta_log'],
|
|
139
|
+
B_re=p['B_re'], B_im=p['B_im'], gamma_log=p['gamma_log'],
|
|
140
|
+
C_re=p['C_re'], C_im=p['C_im'], D=p['D'],
|
|
141
|
+
)
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
@dataclass(frozen=True)
|
|
145
|
+
class ReducedSSM:
|
|
146
|
+
A: np.ndarray
|
|
147
|
+
B: np.ndarray
|
|
148
|
+
C: np.ndarray
|
|
149
|
+
D: np.ndarray
|
|
150
|
+
r: int
|
|
151
|
+
error_bound: float
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
@dataclass(frozen=True)
|
|
155
|
+
class ReducedDiagonalSSM:
|
|
156
|
+
lam: np.ndarray
|
|
157
|
+
B: np.ndarray
|
|
158
|
+
C: np.ndarray
|
|
159
|
+
D: np.ndarray
|
|
160
|
+
r: int
|
|
161
|
+
error_bound: float
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
@dataclass(frozen=True)
|
|
165
|
+
class Balancing:
|
|
166
|
+
Lp: np.ndarray
|
|
167
|
+
Lq: np.ndarray
|
|
168
|
+
U: np.ndarray
|
|
169
|
+
sigma: np.ndarray
|
|
170
|
+
Vh: np.ndarray
|
|
171
|
+
|
|
172
|
+
|
|
173
|
+
def tensor_to_numpy(x: torch.Tensor) -> np.ndarray:
|
|
174
|
+
return x.detach().cpu().numpy().copy()
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
class CompreSSM:
|
|
178
|
+
@staticmethod
|
|
179
|
+
def calculate_gramians_diagonal(lam: np.ndarray, *, B: np.ndarray, C: np.ndarray):
|
|
180
|
+
denom = 1.0 - lam[:, None] * lam.conj()[None, :]
|
|
181
|
+
P = (B @ B.conj().T) / denom
|
|
182
|
+
Q = (C.conj().T @ C) / denom.conj()
|
|
183
|
+
return (P + P.conj().T) / 2, (Q + Q.conj().T) / 2
|
|
184
|
+
|
|
185
|
+
@staticmethod
|
|
186
|
+
def calculate_gramians(A: np.ndarray, *, B: np.ndarray, C: np.ndarray):
|
|
187
|
+
P = solve_discrete_lyapunov(A, B @ B.conj().T)
|
|
188
|
+
Q = solve_discrete_lyapunov(A.conj().T, C.conj().T @ C)
|
|
189
|
+
return (P + P.conj().T) / 2, (Q + Q.conj().T) / 2
|
|
190
|
+
|
|
191
|
+
@staticmethod
|
|
192
|
+
def diagonalize_reduced(Ar: np.ndarray, Br: np.ndarray, Cr: np.ndarray):
|
|
193
|
+
lam_r, V = np.linalg.eig(Ar)
|
|
194
|
+
return lam_r, np.linalg.solve(V, Br), Cr @ V
|
|
195
|
+
|
|
196
|
+
@staticmethod
|
|
197
|
+
def cholesky_with_jitter(M: np.ndarray, *, name: str, jitter: float = 1e-12, max_tries: int = 8):
|
|
198
|
+
n = M.shape[0]
|
|
199
|
+
scale = float(np.real(np.trace(M))) / n
|
|
200
|
+
for k in range(max_tries):
|
|
201
|
+
eps = 0.0 if k == 0 else jitter * scale * 10.0 ** (k - 1)
|
|
202
|
+
try:
|
|
203
|
+
return cholesky(M + eps * np.eye(n), lower=True)
|
|
204
|
+
except np.linalg.LinAlgError:
|
|
205
|
+
continue
|
|
206
|
+
raise np.linalg.LinAlgError(f"Gramian not positive definite on matrix {name}, even with jitter")
|
|
207
|
+
|
|
208
|
+
@staticmethod
|
|
209
|
+
def calculate_sqrt(P: np.ndarray, *, Q: np.ndarray):
|
|
210
|
+
Lp = CompreSSM.cholesky_with_jitter(P, name="P (controllability Gramian)")
|
|
211
|
+
Lq = CompreSSM.cholesky_with_jitter(Q, name="Q (observability Gramian)")
|
|
212
|
+
return Lp, Lq
|
|
213
|
+
|
|
214
|
+
@staticmethod
|
|
215
|
+
def HSVs(Lp: np.ndarray, *, Lq: np.ndarray):
|
|
216
|
+
return svd(Lq.conj().T @ Lp)
|
|
217
|
+
|
|
218
|
+
@staticmethod
|
|
219
|
+
def balance(P: np.ndarray, *, Q: np.ndarray) -> Balancing:
|
|
220
|
+
Lp, Lq = CompreSSM.calculate_sqrt(P, Q=Q)
|
|
221
|
+
U, sigma, Vh = CompreSSM.HSVs(Lp, Lq=Lq)
|
|
222
|
+
return Balancing(Lp=Lp, Lq=Lq, U=U, sigma=sigma, Vh=Vh)
|
|
223
|
+
|
|
224
|
+
@staticmethod
|
|
225
|
+
def rank_selection(sigma: np.ndarray, *, threshold: float) -> int:
|
|
226
|
+
"""
|
|
227
|
+
Finds the smallest rank retaining the requested fraction of the HSV sum.
|
|
228
|
+
|
|
229
|
+
Args:
|
|
230
|
+
sigma: 1D array of Hankel singular values sorted in descending order.
|
|
231
|
+
threshold: Fraction of the total HSV sum to retain (1.0 - tol).
|
|
232
|
+
"""
|
|
233
|
+
sigma = np.asarray(sigma)
|
|
234
|
+
if sigma.ndim != 1 or not np.all(np.isfinite(sigma)) or np.any(sigma < 0):
|
|
235
|
+
raise ValueError("sigma must be a one-dimensional array of finite non-negative values")
|
|
236
|
+
if not np.isfinite(threshold) or not 0 <= threshold <= 1:
|
|
237
|
+
raise ValueError("threshold must be between zero and one")
|
|
238
|
+
if sigma.size == 0:
|
|
239
|
+
return 0
|
|
240
|
+
|
|
241
|
+
total = np.sum(sigma)
|
|
242
|
+
if total <= 0:
|
|
243
|
+
return len(sigma)
|
|
244
|
+
cumulative_sum = np.cumsum(sigma)
|
|
245
|
+
ranks = np.flatnonzero(cumulative_sum >= threshold * total)
|
|
246
|
+
return int(ranks[0] + 1) if ranks.size else len(sigma)
|
|
247
|
+
|
|
248
|
+
@staticmethod
|
|
249
|
+
def should_compress(r: int, n: int, *, max_ratio: float):
|
|
250
|
+
return 0 < r < max_ratio * n
|
|
251
|
+
|
|
252
|
+
@staticmethod
|
|
253
|
+
def calculate_transform_matrix(Lp: np.ndarray, *, r: int, sigma: np.ndarray, Vh: np.ndarray):
|
|
254
|
+
if np.any(sigma[:r] <= 0):
|
|
255
|
+
raise np.linalg.LinAlgError("Retained Hankel singular values must be positive")
|
|
256
|
+
Vr = Vh[:r].conj().T
|
|
257
|
+
diag_entries = sigma[:r] ** -0.5
|
|
258
|
+
return Lp @ Vr @ np.diag(diag_entries), diag_entries
|
|
259
|
+
|
|
260
|
+
@staticmethod
|
|
261
|
+
def calculate_projection_matrix(Lq: np.ndarray, *, r: int, diag_entries: np.ndarray, U: np.ndarray):
|
|
262
|
+
return np.diag(diag_entries) @ U[:, :r].conj().T @ Lq.conj().T
|
|
263
|
+
|
|
264
|
+
@staticmethod
|
|
265
|
+
def truncate_A(A: np.ndarray, *, S: np.ndarray, T: np.ndarray):
|
|
266
|
+
return S @ A @ T
|
|
267
|
+
|
|
268
|
+
@staticmethod
|
|
269
|
+
def truncate_B(S: np.ndarray, *, B: np.ndarray):
|
|
270
|
+
return S @ B
|
|
271
|
+
|
|
272
|
+
@staticmethod
|
|
273
|
+
def truncate_C(C: np.ndarray, *, T: np.ndarray):
|
|
274
|
+
return C @ T
|
|
275
|
+
|
|
276
|
+
@staticmethod
|
|
277
|
+
def truncate(A: np.ndarray, B: np.ndarray, C: np.ndarray, *, balancing: Balancing, r: int):
|
|
278
|
+
T, diag_entries = CompreSSM.calculate_transform_matrix(
|
|
279
|
+
balancing.Lp, r=r, sigma=balancing.sigma, Vh=balancing.Vh
|
|
280
|
+
)
|
|
281
|
+
S = CompreSSM.calculate_projection_matrix(
|
|
282
|
+
balancing.Lq, r=r, diag_entries=diag_entries, U=balancing.U
|
|
283
|
+
)
|
|
284
|
+
return (
|
|
285
|
+
CompreSSM.truncate_A(A, S=S, T=T),
|
|
286
|
+
CompreSSM.truncate_B(S, B=B),
|
|
287
|
+
CompreSSM.truncate_C(C, T=T),
|
|
288
|
+
)
|
|
289
|
+
|
|
290
|
+
|
|
291
|
+
def compress_from_gramians(A: np.ndarray, B: np.ndarray, C: np.ndarray, *, P: np.ndarray, Q: np.ndarray, config: CompressionConfig):
|
|
292
|
+
balancing = CompreSSM.balance(P, Q=Q)
|
|
293
|
+
r = CompreSSM.rank_selection(balancing.sigma, threshold=config.threshold)
|
|
294
|
+
|
|
295
|
+
if not CompreSSM.should_compress(r, A.shape[0], max_ratio=config.max_ratio):
|
|
296
|
+
return None
|
|
297
|
+
|
|
298
|
+
Ar, Br, Cr = CompreSSM.truncate(A, B, C, balancing=balancing, r=r)
|
|
299
|
+
error_bound = 2 * float(np.real(np.sum(balancing.sigma[r:])))
|
|
300
|
+
return Ar, Br, Cr, r, error_bound
|
|
301
|
+
|
|
302
|
+
|
|
303
|
+
def compress_numpy(system: DiscreteSSM, *, config: CompressionConfig) -> ReducedSSM | None:
|
|
304
|
+
"""Apply discrete-time balanced truncation to one stable, minimal LTI system.
|
|
305
|
+
|
|
306
|
+
The returned error bound is the standard H-infinity bound for balanced
|
|
307
|
+
truncation; it assumes the LTI system is stable and minimal.
|
|
308
|
+
"""
|
|
309
|
+
P, Q = CompreSSM.calculate_gramians(system.A, B=system.B, C=system.C)
|
|
310
|
+
result = compress_from_gramians(system.A, system.B, system.C, P=P, Q=Q, config=config)
|
|
311
|
+
if result is None:
|
|
312
|
+
return None
|
|
313
|
+
Ar, Br, Cr, r, error_bound = result
|
|
314
|
+
return ReducedSSM(A=Ar, B=Br, C=Cr, D=system.D, r=r, error_bound=error_bound)
|
|
315
|
+
|
|
316
|
+
|
|
317
|
+
def compress_diagonal_numpy(system: DiagonalDiscreteSSM, *, config: CompressionConfig) -> ReducedDiagonalSSM | None:
|
|
318
|
+
"""Balanced truncation of a diagonal system, returned in diagonal form."""
|
|
319
|
+
P, Q = CompreSSM.calculate_gramians_diagonal(system.lam, B=system.B, C=system.C)
|
|
320
|
+
result = compress_from_gramians(np.diag(system.lam), system.B, system.C, P=P, Q=Q, config=config)
|
|
321
|
+
if result is None:
|
|
322
|
+
return None
|
|
323
|
+
Ar, Br, Cr, r, error_bound = result
|
|
324
|
+
lam_r, Br_diag, Cr_diag = CompreSSM.diagonalize_reduced(Ar, Br, Cr)
|
|
325
|
+
return ReducedDiagonalSSM(lam=lam_r, B=Br_diag, C=Cr_diag, D=system.D, r=r, error_bound=error_bound)
|
|
326
|
+
|
|
327
|
+
|
|
328
|
+
def compress_module(module: nn.Module, *, config: CompressionConfig) -> ReducedSSM | None:
|
|
329
|
+
"""Extract and compress the expected fixed-step diagonal SSM; does not mutate the module."""
|
|
330
|
+
system = DiscreteSSM.from_module(module, dt=config.dt)
|
|
331
|
+
return compress_numpy(system, config=config)
|
|
332
|
+
|
|
333
|
+
|
|
334
|
+
def compress_lru_module(module: nn.Module, *, config: CompressionConfig) -> ReducedDiagonalSSM | None:
|
|
335
|
+
"""Extract and compress one LRU block; does not mutate the module."""
|
|
336
|
+
system = DiagonalDiscreteSSM.from_lru_module(module)
|
|
337
|
+
return compress_diagonal_numpy(system, config=config)
|
|
338
|
+
|
|
339
|
+
|
|
340
|
+
def lru_arrays_from_diagonal(reduced: ReducedDiagonalSSM) -> dict[str, np.ndarray]:
|
|
341
|
+
modulus = np.clip(np.abs(reduced.lam), 1e-12, 1.0 - 1e-12)
|
|
342
|
+
theta = np.maximum(np.angle(reduced.lam) % (2 * np.pi), 1e-8)
|
|
343
|
+
gamma = np.sqrt(1.0 - modulus ** 2)
|
|
344
|
+
B_stored = reduced.B / gamma[:, None]
|
|
345
|
+
return {
|
|
346
|
+
'nu_log': np.log(-np.log(modulus)),
|
|
347
|
+
'theta_log': np.log(theta),
|
|
348
|
+
'gamma_log': np.log(gamma),
|
|
349
|
+
'B_re': B_stored.real,
|
|
350
|
+
'B_im': B_stored.imag,
|
|
351
|
+
'C_re': reduced.C.real,
|
|
352
|
+
'C_im': reduced.C.imag,
|
|
353
|
+
}
|
|
354
|
+
|
|
355
|
+
|
|
356
|
+
def replace_parameter(module: nn.Module, name: str, array: np.ndarray) -> None:
|
|
357
|
+
old = getattr(module, name)
|
|
358
|
+
new = torch.tensor(np.ascontiguousarray(array), dtype=old.dtype, device=old.device)
|
|
359
|
+
setattr(module, name, nn.Parameter(new, requires_grad=old.requires_grad))
|
|
360
|
+
|
|
361
|
+
|
|
362
|
+
def write_back_lru(module: nn.Module, *, reduced: ReducedDiagonalSSM) -> None:
|
|
363
|
+
"""Replace the LRU block's parameters with the reduced ones. D is unchanged."""
|
|
364
|
+
for name, array in lru_arrays_from_diagonal(reduced).items():
|
|
365
|
+
replace_parameter(module, name, array)
|
|
366
|
+
|
|
367
|
+
|
|
368
|
+
def find_blocks(model: nn.Module, *, block_cls: type) -> list[nn.Module]:
|
|
369
|
+
return [m for m in model.modules() if isinstance(m, block_cls)]
|
|
370
|
+
|
|
371
|
+
|
|
372
|
+
def reduce_lru_blocks(blocks: Iterable[nn.Module], *, config: CompressionConfig) -> list[ReducedDiagonalSSM | None]:
|
|
373
|
+
results = []
|
|
374
|
+
for block in blocks:
|
|
375
|
+
reduced = compress_lru_module(block, config=config)
|
|
376
|
+
if reduced is not None:
|
|
377
|
+
write_back_lru(block, reduced=reduced)
|
|
378
|
+
results.append(reduced)
|
|
379
|
+
return results
|
|
380
|
+
|
|
381
|
+
|
|
382
|
+
def run_scheduled_reduction(
|
|
383
|
+
*,
|
|
384
|
+
step: int,
|
|
385
|
+
schedule: ReductionSchedule,
|
|
386
|
+
blocks: Iterable[nn.Module],
|
|
387
|
+
config: CompressionConfig,
|
|
388
|
+
make_optimizer: Callable[[], torch.optim.Optimizer],
|
|
389
|
+
):
|
|
390
|
+
"""Returns a freshly built optimizer if any block was reduced, otherwise None."""
|
|
391
|
+
if not schedule.fires_at(step):
|
|
392
|
+
return None
|
|
393
|
+
results = reduce_lru_blocks(blocks, config=config)
|
|
394
|
+
if all(result is None for result in results):
|
|
395
|
+
return None
|
|
396
|
+
return make_optimizer()
|