zilla-muf 0.1.1

Shared structured-matrix and numerical primitives for sparse attention and state space models (SSMs).
Documentation
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) {
	// A_bar = e^{Δ·a}: the exact zero-order-hold update for the
	// homogeneous (state-decay) part over one step of size Δ.
	let a_bar = (delta * a).exp();
	// B_bar = (e^{Δ·a} - 1)/a · b is the exact input contribution, but it
	// is the indeterminate 0/0 as a -> 0. Fall back to its limit, Δ·b,
	// whenever |a| is within machine epsilon of zero so the result stays
	// finite and accurate instead of dividing by ~0.
	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}");
	}
}