irithyll_core/ssm/mod.rs
1//! State Space Models for streaming temporal feature extraction.
2//!
3//! This module implements diagonal and selective (Mamba-style) state space models
4//! for processing sequential data in a streaming fashion. SSMs maintain hidden
5//! state that evolves with each input timestep, capturing temporal dependencies
6//! without storing past observations.
7//!
8//! # Architecture
9//!
10//! The continuous-time SSM is defined by:
11//!
12//! ```text
13//! h'(t) = A * h(t) + B * x(t) (state evolution)
14//! y(t) = C * h(t) + D * x(t) (output equation)
15//! ```
16//!
17//! For discrete-time processing, we discretize via Zero-Order Hold (ZOH) or
18//! bilinear transform. The **selective** variant (Mamba) makes B, C, and the
19//! discretization step Delta input-dependent, enabling content-aware filtering.
20//!
21//! # Modules
22//!
23//! - [`diagonal`] -- Non-selective diagonal SSM with fixed parameters
24//! - [`selective`] -- Mamba-style selective SSM with input-dependent projections
25//! - [`discretize`] -- ZOH and bilinear discretization methods
26//! - [`init`] -- A-matrix initialization strategies (Mamba, S4D)
27//! - [`projection`] -- Linear algebra helpers and PRNG for weight initialization
28
29pub mod diagonal;
30pub mod discretize;
31pub mod init;
32pub mod projection;
33pub mod selective;
34
35pub use diagonal::DiagonalSSM;
36pub use selective::SelectiveSSM;
37
38use alloc::vec::Vec;
39
40/// Trait for SSM layers that process sequential data one timestep at a time.
41///
42/// Implementors maintain internal hidden state that evolves with each call to
43/// [`forward`](SSMLayer::forward). The hidden state captures temporal patterns
44/// from the input sequence without requiring storage of past observations.
45///
46/// # Thread Safety
47///
48/// All SSM layers are `Send + Sync`, enabling use in async pipelines and
49/// parallel prediction contexts.
50pub trait SSMLayer: Send + Sync {
51 /// Process one input timestep and return the output vector.
52 ///
53 /// This advances the internal hidden state by one step. The output
54 /// dimension equals the input dimension for selective SSMs, or 1 for
55 /// scalar diagonal SSMs.
56 ///
57 /// # Arguments
58 ///
59 /// * `input` -- feature vector for this timestep
60 fn forward(&mut self, input: &[f64]) -> Vec<f64>;
61
62 /// Get a reference to the current hidden state.
63 fn state(&self) -> &[f64];
64
65 /// Output dimension of this SSM layer.
66 fn output_dim(&self) -> usize;
67
68 /// Reset hidden state to zeros, as if no data has been seen.
69 fn reset(&mut self);
70}