lumamba 0.0.1

// lumamba-rs — LuMamba EEG foundation-model inference on the RLX runtime.
// Copyright (C) 2026 Nataliya Kosmyna.
// SPDX-License-Identifier: GPL-3.0-only

//! Model and runtime configuration for LuMamba inference.
//!
//! `ModelConfig` mirrors the Python `LuMamba` hyperparameters
//! (`config/model/LuMamba_tiny.yaml` in BioFoundation). Field names match the
//! `"model"` sub-object of the JSON config shipped alongside the weights.

// ── ModelConfig ───────────────────────────────────────────────────────────────

/// LuMamba model hyperparameters (the `"model"` sub-object of the JSON config).
#[derive(Debug, Clone, serde::Deserialize)]
pub struct ModelConfig {
    // ── LUNA front-end parameters ──────────────────────────────────────────
    /// Patch size in time-samples (default 40).
    #[serde(default = "default_patch_size")]
    pub patch_size: usize,

    /// Number of learned cross-attention queries `Q` (default 6 for tiny).
    #[serde(default = "default_num_queries")]
    pub num_queries: usize,

    /// Per-query / per-channel embedding dimension `E` (default 64).
    #[serde(default = "default_embed_dim")]
    pub embed_dim: usize,

    /// Attention heads in the cross-attention front-end (default 2).
    #[serde(default = "default_num_heads")]
    pub num_heads: usize,

    /// MLP expansion ratio inside the cross-attention FFN (default 4.0).
    #[serde(default = "default_mlp_ratio")]
    pub mlp_ratio: f64,

    /// Number of output classes. `0` = reconstruction (pre-training) decoder.
    #[serde(default)]
    pub num_classes: usize,

    /// LayerNorm epsilon (default 1e-5).
    #[serde(default = "default_norm_eps")]
    pub norm_eps: f64,

    // ── FEMBA bidirectional-Mamba encoder parameters ───────────────────────
    /// Inner-dimension expansion factor inside each Mamba block (default 2).
    #[serde(default = "default_expand")]
    pub expand: usize,

    /// Number of (bidirectional) Mamba blocks (default 2 for tiny).
    #[serde(default = "default_num_blocks")]
    pub num_blocks: usize,

    /// SSM state size `N` per channel (Mamba default 16).
    #[serde(default = "default_d_state")]
    pub d_state: usize,

    /// Depthwise causal conv kernel width (Mamba default 4).
    #[serde(default = "default_d_conv")]
    pub d_conv: usize,

    /// Bidirectional combine strategy: `"add"` (default) or `"ew_multiply"`.
    #[serde(default = "default_bidir_strategy")]
    pub bidirectional_strategy: String,
}

fn default_patch_size()    -> usize  { 40 }
fn default_num_queries()   -> usize  { 6 }
fn default_embed_dim()     -> usize  { 64 }
fn default_num_heads()     -> usize  { 2 }
fn default_mlp_ratio()     -> f64    { 4.0 }
fn default_norm_eps()      -> f64    { 1e-5 }
fn default_expand()        -> usize  { 2 }
fn default_num_blocks()    -> usize  { 2 }
fn default_d_state()       -> usize  { 16 }
fn default_d_conv()        -> usize  { 4 }
fn default_bidir_strategy()-> String { "add".to_string() }

impl Default for ModelConfig {
    fn default() -> Self {
        Self {
            patch_size:  default_patch_size(),
            num_queries: default_num_queries(),
            embed_dim:   default_embed_dim(),
            num_heads:   default_num_heads(),
            mlp_ratio:   default_mlp_ratio(),
            num_classes: 0,
            norm_eps:    default_norm_eps(),
            expand:      default_expand(),
            num_blocks:  default_num_blocks(),
            d_state:     default_d_state(),
            d_conv:      default_d_conv(),
            bidirectional_strategy: default_bidir_strategy(),
        }
    }
}

impl ModelConfig {
    /// Mamba `d_model` = effective hidden dim after query concatenation: `Q · E`.
    pub fn hidden_dim(&self) -> usize {
        self.embed_dim * self.num_queries
    }

    /// Mamba inner dimension `d_inner = expand · d_model`.
    pub fn d_inner(&self) -> usize {
        self.hidden_dim() * self.expand
    }

    /// Mamba `dt_rank = ceil(d_model / 16)` (the `mamba_ssm` "auto" default).
    pub fn dt_rank(&self) -> usize {
        self.hidden_dim().div_ceil(16)
    }

    /// FFN hidden dimension inside the cross-attention front-end (`E · mlp_ratio`).
    pub fn ffn_cross_dim(&self) -> usize {
        (self.embed_dim as f64 * self.mlp_ratio) as usize
    }

    /// Cross-attention head dimension (`E / num_heads`).
    pub fn cross_head_dim(&self) -> usize {
        self.embed_dim / self.num_heads
    }

    /// Whether the bidirectional streams combine by elementwise multiply
    /// (`"ew_multiply"`) rather than addition (`"add"`, the default).
    pub fn bidir_multiply(&self) -> bool {
        self.bidirectional_strategy == "ew_multiply"
    }
}

// ── DataConfig ────────────────────────────────────────────────────────────────

/// Preprocessing / epoching parameters.
#[derive(Debug, Clone)]
pub struct DataConfig {
    /// Sampling rate after resampling (Hz).
    pub sample_rate: f32,
    /// Epoch duration in seconds.
    pub epoch_dur: f32,
    /// Lower corner of the channel-position normalisation box (metres).
    pub xyz_min: [f32; 3],
    /// Upper corner of the channel-position normalisation box (metres).
    pub xyz_max: [f32; 3],
}

impl Default for DataConfig {
    fn default() -> Self {
        Self {
            sample_rate: 256.0,
            epoch_dur:   5.0,
            xyz_min: [-0.12, -0.12, -0.12],
            xyz_max: [ 0.12,  0.12,  0.12],
        }
    }
}

impl DataConfig {
    /// Number of time samples per epoch.
    pub fn epoch_samples(&self) -> usize {
        (self.sample_rate * self.epoch_dur) as usize
    }
}