Skip to main content

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}