zilla-muf 0.1.0

Shared structured-matrix and numerical primitives for sparse attention and state space models (SSMs).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

use num_traits::Float;

/// ZOH discretization for a scalar (diagonal) state-space system.
/// Given continuous A, B and step size delta, returns (A_bar, B_bar).
pub fn zoh_discretize<T: Float>(a: T, b: T, delta: T) -> (T, T) {
	let a_bar = (delta * a).exp();
	let b_bar = if a.abs() > T::epsilon() {
		(a_bar - T::one()) / a * b
	} else {
		delta * b
	};
	(a_bar, b_bar)
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn handles_zero_a_with_linear_approximation() {
		let (a_bar, b_bar) = zoh_discretize(0.0_f64, 2.0, 0.5);
		assert!((a_bar - 1.0).abs() < 1e-12);
		assert!((b_bar - 1.0).abs() < 1e-12); // delta * b = 0.5 * 2.0
	}

	#[test]
	fn matches_euler_integration_for_nonzero_a() {
		// Independent check: integrate dh/dt = a*h + b with many tiny
		// Euler steps from h(0) = 0, and compare to a single ZOH step.
		// This validates the formula physically, rather than re-deriving
		// the same exponential expression the implementation uses.
		let a = -1.0_f64;
		let b = 1.0_f64;
		let delta = 0.1_f64;

		let steps = 200_000;
		let dt = delta / steps as f64;
		let mut h = 0.0_f64;
		for _ in 0..steps {
			h += (a * h + b) * dt;
		}

		let (a_bar, b_bar) = zoh_discretize(a, b, delta);
		let h_zoh = a_bar * 0.0 + b_bar; // h0 = 0

		assert!((h - h_zoh).abs() < 1e-6, "euler={h}, zoh={h_zoh}");
	}
}