zsfm_flowstate/infer/
mod.rs

1use std::collections::HashMap;
2use std::io::BufReader;
3use std::sync::Mutex;
4use std::path::{Path, PathBuf};
5
6use anyhow::{Context, Result};
7use candle_core::{DType, Device, Tensor};
8use candle_core::quantized::gguf_file;
9use candle_nn::ops;
10use simdeez::prelude::*;
11
12use crate::config::FlowStateConfig;
13
14simd_runtime_generate!(
15    fn ssm_scan_step(
16        state_r: &mut [f32], state_i: &mut [f32],
17        a_r: &[f32], a_i: &[f32],
18        bu_r: &[f32], bu_i: &[f32],
19    ) {
20        let mut sr = &mut state_r[..]; let mut si = &mut state_i[..];
21        let mut ar = &a_r[..];         let mut ai = &a_i[..];
22        let mut br = &bu_r[..];        let mut bi = &bu_i[..];
23
24        while sr.len() >= S::Vf32::WIDTH {
25            let sr_v = S::Vf32::load_from_slice(sr);
26            let si_v = S::Vf32::load_from_slice(si);
27            let ar_v = S::Vf32::load_from_slice(ar);
28            let ai_v = S::Vf32::load_from_slice(ai);
29            let br_v = S::Vf32::load_from_slice(br);
30            let bi_v = S::Vf32::load_from_slice(bi);
31
32            // new_r = ar*sr - ai*si + br  (neg_mul_add(a,b,c) = c - a*b)
33            let new_r = ar_v.mul_add(sr_v, ai_v.neg_mul_add(si_v, br_v));
34            // new_i = ar*si + ai*sr + bi
35            let new_i = ar_v.mul_add(si_v, ai_v.mul_add(sr_v, bi_v));
36
37            new_r.copy_to_slice(sr);
38            new_i.copy_to_slice(si);
39
40            sr = &mut sr[S::Vf32::WIDTH..]; si = &mut si[S::Vf32::WIDTH..];
41            ar = &ar[S::Vf32::WIDTH..];     ai = &ai[S::Vf32::WIDTH..];
42            br = &br[S::Vf32::WIDTH..];     bi = &bi[S::Vf32::WIDTH..];
43        }
44
45        for j in 0..sr.len() {
46            let nr = ar[j] * sr[j] - ai[j] * si[j] + br[j];
47            let ni = ar[j] * si[j] + ai[j] * sr[j] + bi[j];
48            sr[j] = nr;
49            si[j] = ni;
50        }
51    }
52);
53
54// ---------------------------------------------------------------------------
55// Config
56// ---------------------------------------------------------------------------
57
58#[derive(Debug, Clone)]
59pub struct InferConfig {
60    pub num_layers: usize,
61    pub embed_dim: usize,
62    pub state_dim: usize,
63    pub n_inputs: usize,        // 2 for with_missing (value + mask)
64    pub decoder_dim: usize,
65    pub decoder_patch_len: usize,
66    pub quantiles: Vec<f32>,
67    pub basis_range: [f32; 2],  // e.g. [-1.0, 0.95] for "legs"
68    pub context_length: usize,
69    pub eps: f32,
70}
71
72/// `FlowStateConfig` already resolved every default (via serde) or failed to parse if a
73/// genuinely required field was missing, so this mapping is infallible and needs no further
74/// defaulting of its own.
75impl From<&FlowStateConfig> for InferConfig {
76    fn from(c: &FlowStateConfig) -> Self {
77        InferConfig {
78            num_layers:        c.encoder_num_layers as usize,
79            embed_dim:         c.embedding_feature_dim as usize,
80            state_dim:         c.encoder_state_dim as usize,
81            n_inputs:          c.n_inputs() as usize,
82            decoder_dim:       c.decoder_dim as usize,
83            decoder_patch_len: c.decoder_patch_len as usize,
84            quantiles:         c.quantiles.clone(),
85            basis_range:       c.basis_range(),
86            context_length:    c.context_length as usize,
87            eps:               1e-5,
88        }
89    }
90}
91
92impl InferConfig {
93    pub fn quantiles(&self) -> &[f32] { &self.quantiles }
94    pub fn median_index(&self) -> usize {
95        self.quantiles
96            .iter()
97            .position(|&q| (q - 0.5).abs() < 1e-6)
98            .unwrap_or(self.quantiles.len() / 2)
99    }
100}
101
102// ---------------------------------------------------------------------------
103// Builder
104// ---------------------------------------------------------------------------
105
106/// Fluent constructor for [`FlowStateModel`]: point it at a GGUF file and a config (from a
107/// parsed `config.json` via [`config_from`](FlowStateModelBuilder::config_from)), then call
108/// [`build`](FlowStateModelBuilder::build).
109///
110/// ```no_run
111/// use zsfm_flowstate::{FlowStateConfig, FlowStateModel};
112///
113/// # fn main() -> anyhow::Result<()> {
114/// let cfg = FlowStateConfig::from_json(&std::fs::read_to_string("config.json")?)?;
115/// let model = FlowStateModel::builder("flowstate.gguf").config_from(&cfg).build()?;
116/// # Ok(()) }
117/// ```
118pub struct FlowStateModelBuilder {
119    gguf_path: PathBuf,
120    config: Option<InferConfig>,
121}
122
123impl FlowStateModelBuilder {
124    fn new(gguf_path: impl Into<PathBuf>) -> Self {
125        Self { gguf_path: gguf_path.into(), config: None }
126    }
127
128    pub fn config(mut self, config: InferConfig) -> Self {
129        self.config = Some(config);
130        self
131    }
132
133    pub fn config_from(mut self, c: &FlowStateConfig) -> Self {
134        self.config = Some(InferConfig::from(c));
135        self
136    }
137
138    pub fn build(self) -> Result<FlowStateModel> {
139        let config = self
140            .config
141            .context("FlowStateModelBuilder: no config set — call .config(...) or .config_from(...)")?;
142        FlowStateModel::load(&self.gguf_path, config)
143    }
144}
145
146// ---------------------------------------------------------------------------
147// Weight structs
148// ---------------------------------------------------------------------------
149
150struct S5Weights {
151    log_lambda_real: Vec<f32>,  // [state_dim]
152    lambda_imag:     Vec<f32>,  // [state_dim]
153    b_r: Tensor,                // [state_dim, embed_dim]
154    b_i: Tensor,                // [state_dim, embed_dim]
155    c_r: Tensor,                // [embed_dim, state_dim]
156    c_i: Tensor,                // [embed_dim, state_dim]
157    d:   Tensor,                // [embed_dim]
158    log_delta: Vec<f32>,        // [state_dim]
159}
160
161struct BlockWeights {
162    ssm:        S5Weights,
163    out_weight: Tensor,   // [embed_dim, embed_dim]
164    out_bias:   Tensor,   // [embed_dim]
165    norm_weight: Tensor,  // [embed_dim]
166    norm_bias:   Tensor,  // [embed_dim]
167    // cached per scale_factor: (A_bar_real, A_bar_imag, B_bar_real_t, B_bar_imag_t)
168    disc_cache: Mutex<HashMap<u32, (Vec<f32>, Vec<f32>, Tensor, Tensor)>>,
169}
170
171pub struct FlowStateModel {
172    device:     Device,
173    pub config: InferConfig,
174    embed_w:    Tensor,   // [embed_dim, n_inputs]
175    embed_b:    Tensor,   // [embed_dim]
176    blocks:     Vec<BlockWeights>,
177    decoder_w:  Tensor,   // [n_quantiles * decoder_dim, embed_dim]
178    decoder_b:  Tensor,   // [n_quantiles * decoder_dim]
179    legendre_cache: Mutex<HashMap<usize, Tensor>>,
180}
181
182// ---------------------------------------------------------------------------
183// Loading
184// ---------------------------------------------------------------------------
185
186fn load_f32_vec(
187    content: &gguf_file::Content,
188    reader: &mut (impl std::io::Read + std::io::Seek),
189    name: &str,
190    device: &Device,
191) -> anyhow::Result<Vec<f32>> {
192    zsfm_nn::load_vec(content, reader, name, device)
193}
194
195fn load_matrix(
196    content: &gguf_file::Content,
197    reader: &mut (impl std::io::Read + std::io::Seek),
198    name: &str,
199    device: &Device,
200) -> anyhow::Result<Tensor> {
201    zsfm_nn::load_tensor(content, reader, name, device, DType::F32)
202}
203
204impl FlowStateModel {
205    pub fn load(gguf_path: &Path, config: InferConfig) -> anyhow::Result<Self> {
206        let device = Device::Cpu;
207        let file = std::fs::File::open(gguf_path)
208            .with_context(|| format!("open {}", gguf_path.display()))?;
209        let mut file = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
210        let content = gguf_file::Content::read(&mut file).context("read GGUF header")?;
211
212        let embed_w = load_matrix(&content, &mut file, "embed.weight", &device)?;
213        let embed_b = load_matrix(&content, &mut file, "embed.bias", &device)?;
214
215        let mut blocks = Vec::with_capacity(config.num_layers);
216        for n in 0..config.num_layers {
217            let ssm = S5Weights {
218                log_lambda_real: load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.log_lambda_real"), &device)?,
219                lambda_imag:     load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.lambda_imag"),     &device)?,
220                b_r: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.b_r"), &device)?,
221                b_i: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.b_i"), &device)?,
222                c_r: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.c_r"), &device)?,
223                c_i: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.c_i"), &device)?,
224                d:          load_matrix(&content, &mut file, &format!("blk.{n}.ssm.d"),         &device)?,
225                log_delta:  load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.log_delta"), &device)?,
226            };
227            blocks.push(BlockWeights {
228                ssm,
229                out_weight:  load_matrix(&content, &mut file, &format!("blk.{n}.out.weight"),  &device)?,
230                out_bias:    load_matrix(&content, &mut file, &format!("blk.{n}.out.bias"),    &device)?,
231                norm_weight: load_matrix(&content, &mut file, &format!("blk.{n}.norm.weight"), &device)?,
232                norm_bias:   load_matrix(&content, &mut file, &format!("blk.{n}.norm.bias"),   &device)?,
233                disc_cache:  Mutex::new(HashMap::new()),
234            });
235        }
236
237        let decoder_w = load_matrix(&content, &mut file, "decoder.weight", &device)?;
238        let decoder_b = load_matrix(&content, &mut file, "decoder.bias", &device)?;
239
240        Ok(Self {
241            device,
242            config,
243            embed_w,
244            embed_b,
245            blocks,
246            decoder_w,
247            decoder_b,
248            legendre_cache: Mutex::new(HashMap::new()),
249        })
250    }
251
252    /// Start building a [`FlowStateModel`] — see [`FlowStateModelBuilder`].
253    pub fn builder(gguf_path: impl Into<PathBuf>) -> FlowStateModelBuilder {
254        FlowStateModelBuilder::new(gguf_path)
255    }
256
257    // -----------------------------------------------------------------------
258    // Public inference entry point
259    // -----------------------------------------------------------------------
260
261    /// Forecast `prediction_length` steps from a univariate context series.
262    /// Returns `Vec<Vec<f32>>` of shape `[n_quantiles][prediction_length]`.
263    pub fn forecast(&self, context: &[f32], prediction_length: usize) -> anyhow::Result<Vec<Vec<f32>>> {
264        let cfg = &self.config;
265
266        // 1. Pad or trim context to match model requirements
267        let ctx_len = context.len().min(cfg.context_length);
268        let start = context.len().saturating_sub(ctx_len);
269        let context = &context[start..];
270        let seq_len = context.len();
271
272        // 2. Causal RevIN: compute prefix statistics
273        let (normed_values, final_mean, final_std) = causal_revin_norm(context, cfg.eps);
274        if std::env::var("FLOWSTATE_DEBUG").is_ok() {
275            eprintln!("RevIN final_mean={:.8} final_std={:.8}", final_mean, final_std);
276            eprintln!("normed[0..4]: {:.6} {:.6} {:.6} {:.6}",
277                      normed_values[0], normed_values[1], normed_values[2], normed_values[3]);
278            eprintln!("normed[252..256]: {:.6} {:.6} {:.6} {:.6}",
279                      normed_values[252], normed_values[253], normed_values[254], normed_values[255]);
280        }
281
282        // 3. Build input tensor [seq_len, n_inputs] with mask channel = 0 (no missing)
283        let mut input_data = vec![0.0f32; seq_len * cfg.n_inputs];
284        for t in 0..seq_len {
285            input_data[t * cfg.n_inputs] = normed_values[t];
286            if cfg.n_inputs > 1 {
287                input_data[t * cfg.n_inputs + 1] = 0.0; // mask = 0 means known
288            }
289        }
290
291        // 4. Embedding: [seq_len, n_inputs] × embed_w^T + embed_b → [seq_len, embed_dim]
292        let input_t = Tensor::from_vec(input_data, (seq_len, cfg.n_inputs), &self.device)?;
293        let mut hidden = linear(&input_t, &self.embed_w, &self.embed_b)?;
294
295        // 5. Scale factor for discretization: decoder_patch_len / prediction_length
296        let scale_factor = cfg.decoder_patch_len as f32 / prediction_length as f32;
297
298        // 6. Encoder: S5 layers
299        // Pre-allocate scratch buffers once and reuse across all blocks.
300        // This avoids repeated large Vec allocations (seq_len * state_dim floats each)
301        // that would otherwise be zero-initialised and discarded per non-last block.
302        let scratch_len = seq_len * cfg.state_dim;
303        let mut scan_r = vec![0.0f32; scratch_len];
304        let mut scan_i = vec![0.0f32; scratch_len];
305        let mut state_r = vec![0.0f32; cfg.state_dim];
306        let mut state_i = vec![0.0f32; cfg.state_dim];
307        let num_layers = self.blocks.len();
308        for (i, block) in self.blocks.iter().enumerate() {
309            let is_last = i == num_layers - 1;
310            hidden = self.apply_s5_layer(
311                hidden, block, scale_factor, is_last,
312                &mut scan_r, &mut scan_i, &mut state_r, &mut state_i,
313            )?;
314        }
315        // After last layer: hidden is [1, embed_dim]
316
317        // 7. Decoder: linear → [n_q, decoder_dim]
318        let n_q = cfg.quantiles.len();
319        let coeffs = linear(&hidden, &self.decoder_w, &self.decoder_b)?
320            .reshape((n_q, cfg.decoder_dim))?;
321
322        if std::env::var("FLOWSTATE_DEBUG").is_ok() {
323            let coeffs_data: Vec<f32> = coeffs.flatten_all()?.to_vec1()?;
324            for qi in 0..n_q {
325                for d in 0..cfg.decoder_dim {
326                    eprintln!("COEFF,{qi},{d},{:.8}", coeffs_data[qi * cfg.decoder_dim + d]);
327                }
328            }
329        }
330
331        // 8. Legendre basis [prediction_length, decoder_dim] — cached per prediction_length
332        let basis = {
333            let mut cache = self.legendre_cache.lock().unwrap();
334            if !cache.contains_key(&prediction_length) {
335                let raw = legendre_basis(prediction_length, cfg.decoder_dim, cfg.basis_range,
336                                        scale_factor, cfg.decoder_patch_len);
337                let flat: Vec<f32> = raw.into_iter().flatten().collect();
338                cache.insert(prediction_length,
339                    Tensor::from_vec(flat, (prediction_length, cfg.decoder_dim), &self.device)?);
340            }
341            cache[&prediction_length].clone()
342        };
343
344        // 9. [n_q, decoder_dim] @ [decoder_dim, prediction_length] → [n_q, prediction_length]
345        let out_t = coeffs.matmul(&basis.t()?)?;
346
347        // 10. Denormalize raw decoder channels.
348        let out_raw: Vec<f32> = out_t.flatten_all()?.to_vec1()?;
349        let mut denormed = vec![vec![0.0f32; prediction_length]; n_q];
350        for q in 0..n_q {
351            for p in 0..prediction_length {
352                denormed[q][p] = out_raw[q * prediction_length + p] * final_std + final_mean;
353            }
354        }
355
356        // 11. Quantile recalibration: FlowStateForPrediction.forward() does not use the
357        // n_q raw decoder channels directly as quantile predictions. It treats them as
358        // n_q empirical samples and re-derives quantile estimates at the configured
359        // probability levels via linear-interpolation order statistics
360        // (`torch.quantile(model_output.last_hidden_state, quantiles, dim=1)` in
361        // modeling_flowstate.py). Skipping this step caused outer quantiles (q0.1, q0.9)
362        // to diverge from the Python reference by up to ~3.8 while q0.5 stayed near-exact
363        // (interpolation index for p=0.5 lands exactly on the middle sorted sample).
364        let mut output = vec![vec![0.0f32; prediction_length]; n_q];
365        for p in 0..prediction_length {
366            let mut sorted: Vec<f32> = (0..n_q).map(|q| denormed[q][p]).collect();
367            sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
368            for (qi, &prob) in cfg.quantiles.iter().enumerate() {
369                let idx = (n_q - 1) as f32 * prob;
370                let lower = idx.floor() as usize;
371                let upper = idx.ceil() as usize;
372                let weight = idx - lower as f32;
373                output[qi][p] = sorted[lower] * (1.0 - weight) + sorted[upper] * weight;
374            }
375        }
376
377        Ok(output)
378    }
379
380    // -----------------------------------------------------------------------
381    // S5 layer (one encoder block)
382    // -----------------------------------------------------------------------
383
384    /// `scan_r` / `scan_i`: caller-owned scratch buffers of at least `seq_len * state_dim`
385    /// elements, reused across blocks to avoid repeated large heap allocations.
386    /// `state_r` / `state_i`: caller-owned scratch of at least `state_dim` elements.
387    #[allow(clippy::too_many_arguments)]
388    fn apply_s5_layer(
389        &self,
390        x: Tensor,          // [seq_len, embed_dim]
391        block: &BlockWeights,
392        scale_factor: f32,
393        is_last: bool,
394        scan_r: &mut Vec<f32>,  // scratch: seq_len * state_dim (reused across blocks)
395        scan_i: &mut Vec<f32>,
396        state_r: &mut Vec<f32>, // scratch: state_dim (running real state)
397        state_i: &mut Vec<f32>, // scratch: state_dim (running imag state)
398    ) -> anyhow::Result<Tensor> {
399        let cfg = &self.config;
400        let state_dim = cfg.state_dim;
401        let embed_dim = cfg.embed_dim;
402
403        let seq_len = x.dim(0)?;
404
405        // Save skip connection (trimmed for last layer)
406        let skip = if is_last {
407            x.narrow(0, seq_len - 1, 1)?  // [1, embed_dim]
408        } else {
409            x.clone()
410        };
411
412        // ---- Get or compute discretized SSM matrices (cached per scale_factor) ----
413        let (a_bar_r, a_bar_i, b_bar_r_t, b_bar_i_t) = {
414            let key = scale_factor.to_bits();
415            let mut cache = block.disc_cache.lock().unwrap();
416            if !cache.contains_key(&key) {
417                let result = discretize(&block.ssm, scale_factor, state_dim, embed_dim, &self.device)?;
418                cache.insert(key, result);
419            }
420            let (ar, ai, brt, bit) = &cache[&key];
421            (ar.clone(), ai.clone(), brt.clone(), bit.clone())
422        };
423
424        // B @ x for all timesteps at once: x [seq_len, embed_dim] × B^T [embed_dim, state_dim]
425        let bu_r = x.matmul(&b_bar_r_t.t()?)?;  // [seq_len, state_dim]
426        let bu_i = x.matmul(&b_bar_i_t.t()?)?;
427        let bu_r_data: Vec<f32> = bu_r.flatten_all()?.to_vec1()?;
428        let bu_i_data: Vec<f32> = bu_i.flatten_all()?.to_vec1()?;
429
430        // Sequential SSM scan: h[t] = A_bar * h[t-1] + B_bar * u[t].
431        // A_bar is diagonal complex, so each state dimension s is independent (NEON-vectorisable
432        // inner loop). The outer loop over t is the inherent sequential dependency.
433        // We reuse caller-provided scratch buffers to avoid repeated large heap allocations.
434        //
435        // Reset running state to zero at the start of each block.
436        state_r[..state_dim].fill(0.0);
437        state_i[..state_dim].fill(0.0);
438
439        let (h_r_t, h_i_t) = if is_last {
440            // Last block: only the final hidden state is consumed downstream.
441            for t in 0..seq_len {
442                let bu_r_t = &bu_r_data[t * state_dim..(t + 1) * state_dim];
443                let bu_i_t = &bu_i_data[t * state_dim..(t + 1) * state_dim];
444                ssm_scan_step(
445                    &mut state_r[..state_dim], &mut state_i[..state_dim],
446                    &a_bar_r, &a_bar_i, bu_r_t, bu_i_t,
447                );
448            }
449            let hr = Tensor::from_vec(state_r[..state_dim].to_vec(), (1, state_dim), &self.device)?;
450            let hi = Tensor::from_vec(state_i[..state_dim].to_vec(), (1, state_dim), &self.device)?;
451            (hr, hi)
452        } else {
453            // Non-last blocks: full hidden state history required for the C projection.
454            // Ensure scratch buffers are large enough (they are since caller sized them for
455            // the maximum seq_len * state_dim of the first call in this forecast).
456            let needed = seq_len * state_dim;
457            if scan_r.len() < needed { scan_r.resize(needed, 0.0); }
458            if scan_i.len() < needed { scan_i.resize(needed, 0.0); }
459            for t in 0..seq_len {
460                let bu_r_t = &bu_r_data[t * state_dim..(t + 1) * state_dim];
461                let bu_i_t = &bu_i_data[t * state_dim..(t + 1) * state_dim];
462                ssm_scan_step(
463                    &mut state_r[..state_dim], &mut state_i[..state_dim],
464                    &a_bar_r, &a_bar_i, bu_r_t, bu_i_t,
465                );
466                let row = t * state_dim;
467                scan_r[row..row + state_dim].copy_from_slice(&state_r[..state_dim]);
468                scan_i[row..row + state_dim].copy_from_slice(&state_i[..state_dim]);
469            }
470            let hr = Tensor::from_vec(scan_r[..needed].to_vec(), (seq_len, state_dim), &self.device)?;
471            let hi = Tensor::from_vec(scan_i[..needed].to_vec(), (seq_len, state_dim), &self.device)?;
472            (hr, hi)
473        };
474
475        // C @ h: y_real = C_r @ h_r - C_i @ h_i  → [out_seq_len, embed_dim]
476        let y_from_cr = h_r_t.matmul(&block.ssm.c_r.t()?)?;
477        let y_from_ci = h_i_t.matmul(&block.ssm.c_i.t()?)?;
478        let y_raw = (y_from_cr - y_from_ci)?;
479
480        // D skip: y += D * x_at_positions  (skip is the correct slice in both cases)
481        let y_t = (y_raw + skip.broadcast_mul(&block.ssm.d)?)?;
482
483        // ---- MLP: selu(y) * sigmoid(out_linear(selu(y))) ----
484        let y_selu = selu_tensor(&y_t)?;
485        let gate_pre = linear(&y_selu, &block.out_weight, &block.out_bias)?;
486        let gate = sigmoid_tensor(&gate_pre)?;
487        let y_gated = y_selu.mul(&gate)?;
488
489        // ---- LayerNorm ----
490        let y_normed = layer_norm(&y_gated, &block.norm_weight, &block.norm_bias, self.config.eps)?;
491
492        // ---- Residual ----
493        Ok((y_normed + skip)?)
494    }
495}
496
497// ---------------------------------------------------------------------------
498// Causal RevIN
499// ---------------------------------------------------------------------------
500
501/// Returns (normalized_values[seq_len], final_mean, final_std).
502/// Each position t is normalized by the cumulative mean/std of x[0..=t].
503fn causal_revin_norm(x: &[f32], eps: f32) -> (Vec<f32>, f32, f32) {
504    let n = x.len();
505    let mut normed = vec![0.0f32; n];
506    let mut cum_sum = 0.0f32;
507    let mut cum_sq_diff = 0.0f32;
508    let mut final_mean = 0.0f32;
509    let mut final_std = 1.0f32;
510
511    for t in 0..n {
512        let count = (t + 1) as f32;
513        cum_sum += x[t];
514        let mean_t = cum_sum / count;
515
516        cum_sq_diff += (x[t] - mean_t) * (x[t] - mean_t);
517        let var_t = (cum_sq_diff / count).max(0.0);
518        let std_t = (var_t + eps).sqrt();
519
520        normed[t] = (x[t] - mean_t) / std_t;
521
522        if t == n - 1 {
523            final_mean = mean_t;
524            final_std = std_t;
525        }
526    }
527
528    (normed, final_mean, final_std)
529}
530
531// ---------------------------------------------------------------------------
532// SSM discretization
533// ---------------------------------------------------------------------------
534
535/// Returns (A_bar_real, A_bar_imag, B_bar_real_tensor, B_bar_imag_tensor).
536/// B_bar tensors have shape [state_dim, embed_dim].
537fn discretize(
538    ssm: &S5Weights,
539    scale_factor: f32,
540    state_dim: usize,
541    embed_dim: usize,
542    device: &Device,
543) -> anyhow::Result<(Vec<f32>, Vec<f32>, Tensor, Tensor)> {
544    let mut a_r = vec![0.0f32; state_dim];
545    let mut a_i = vec![0.0f32; state_dim];
546    let mut coeff_r = vec![0.0f32; state_dim];
547    let mut coeff_i = vec![0.0f32; state_dim];
548
549    for s in 0..state_dim {
550        let lam_r = -ssm.log_lambda_real[s].exp();
551        let lam_i = ssm.lambda_imag[s];
552        // Delta_eff = scale_factor * exp(log_Delta), matching modeling_flowstate.py's
553        // `log_Lambda_bar = scale_factor * lambda_ * exp(log_Delta)`. NOT
554        // exp(scale_factor * log_Delta) — the two coincide only at scale_factor == 1.0
555        // (i.e. when horizon == decoder_patch_len), which masked this bug for the
556        // common single-patch case.
557        let delta = scale_factor * ssm.log_delta[s].exp();
558
559        let exp_r = lam_r * delta;
560        let exp_i = lam_i * delta;
561        let mag = exp_r.exp();
562        a_r[s] = mag * exp_i.cos();
563        a_i[s] = mag * exp_i.sin();
564
565        let num_r = a_r[s] - 1.0;
566        let num_i = a_i[s];
567        let denom_sq = lam_r * lam_r + lam_i * lam_i;
568        if denom_sq > 1e-20 {
569            coeff_r[s] = (num_r * lam_r + num_i * lam_i) / denom_sq;
570            coeff_i[s] = (num_i * lam_r - num_r * lam_i) / denom_sq;
571        } else {
572            coeff_r[s] = delta;
573            coeff_i[s] = 0.0;
574        }
575    }
576
577    let b_r_data = get_tensor_data_row_major(&ssm.b_r, state_dim, embed_dim);
578    let b_i_data = get_tensor_data_row_major(&ssm.b_i, state_dim, embed_dim);
579
580    let mut b_bar_r = vec![0.0f32; state_dim * embed_dim];
581    let mut b_bar_i = vec![0.0f32; state_dim * embed_dim];
582
583    for s in 0..state_dim {
584        for e in 0..embed_dim {
585            let br = b_r_data[s * embed_dim + e];
586            let bi = b_i_data[s * embed_dim + e];
587            b_bar_r[s * embed_dim + e] = coeff_r[s] * br - coeff_i[s] * bi;
588            b_bar_i[s * embed_dim + e] = coeff_r[s] * bi + coeff_i[s] * br;
589        }
590    }
591
592    let b_bar_r_t = Tensor::from_vec(b_bar_r, (state_dim, embed_dim), device)?;
593    let b_bar_i_t = Tensor::from_vec(b_bar_i, (state_dim, embed_dim), device)?;
594
595    Ok((a_r, a_i, b_bar_r_t, b_bar_i_t))
596}
597
598/// Extract tensor data in row-major order as Vec<f32>.
599fn get_tensor_data_row_major(t: &Tensor, rows: usize, cols: usize) -> Vec<f32> {
600    t.to_dtype(DType::F32)
601        .and_then(|t| t.reshape((rows, cols)))
602        .and_then(|t| t.flatten_all())
603        .and_then(|t| t.to_vec1())
604        .unwrap_or_else(|_| vec![0.0f32; rows * cols])
605}
606
607// ---------------------------------------------------------------------------
608// Legendre basis (FlowStateLegendreBasis equivalent)
609// ---------------------------------------------------------------------------
610
611/// Public wrapper for diagnostics.
612pub fn dump_legendre_basis(n_points: usize, degree: usize, range: [f32; 2],
613                           scale: f32, pred_dist: usize) -> Vec<Vec<f32>> {
614    legendre_basis(n_points, degree, range, scale, pred_dist)
615}
616
617/// Compute Legendre polynomial basis matrix.
618/// Returns [n_points][degree+1] scaled by 1/4 (as in get_kernel).
619fn legendre_basis(n_points: usize, degree: usize, range: [f32; 2],
620                  scale: f32, pred_dist: usize) -> Vec<Vec<f32>> {
621    let dt = scale * (range[1] - range[0]) / pred_dist as f32;
622    let t: Vec<f32> = (1..=n_points).map(|i| range[0] + i as f32 * dt).collect();
623
624    // Compute degree Legendre polynomials (P0..P_{degree-1}) per point.
625    // degree+1 scratch columns needed during recurrence, then truncated to degree.
626    let mut basis = vec![vec![0.0f32; degree + 1]; n_points];
627    for (p, &x) in t.iter().enumerate() {
628        basis[p][0] = 1.0;
629        if degree >= 1 {
630            basis[p][1] = x;
631        }
632        for k in 1..degree {
633            let kf = k as f32;
634            basis[p][k + 1] =
635                ((2.0 * kf + 1.0) * x * basis[p][k] - kf * basis[p][k - 1]) / (kf + 1.0);
636        }
637        for d in 0..degree {
638            basis[p][d] /= 4.0;
639        }
640        basis[p].truncate(degree);
641    }
642
643    basis
644}
645
646// ---------------------------------------------------------------------------
647// Neural network primitives
648// ---------------------------------------------------------------------------
649
650/// y = x @ w^T + b  (w: [out, in], b: [out])
651fn linear(x: &Tensor, w: &Tensor, b: &Tensor) -> anyhow::Result<Tensor> {
652    zsfm_nn::linear_bias(x, w, b)
653}
654
655/// SELU activation using candle ops (no Vec roundtrip).
656fn selu_tensor(x: &Tensor) -> anyhow::Result<Tensor> {
657    const SCALE: f64 = 1.0507009873554804934193349852946;
658    const ALPHA: f64 = 1.6732632423543772848170429916717;
659    const ALPHA_SCALE: f64 = SCALE * ALPHA;
660    let pos = x.relu()?;
661    let neg = (x - &pos)?;            // min(x, 0)
662    let selu_pos = (pos * SCALE)?;
663    let selu_neg = ((neg.exp()? - 1.0)? * ALPHA_SCALE)?;
664    Ok((selu_pos + selu_neg)?)
665}
666
667fn sigmoid_tensor(x: &Tensor) -> anyhow::Result<Tensor> {
668    Ok(ops::sigmoid(x)?)
669}
670
671/// LayerNorm using candle ops (no Vec roundtrip).
672fn layer_norm(x: &Tensor, weight: &Tensor, bias: &Tensor, eps: f32) -> anyhow::Result<Tensor> {
673    zsfm_nn::layer_norm(x, weight, bias, eps as f64)
674}
675
676// ---------------------------------------------------------------------------
677// zsfm-core::Forecaster
678// ---------------------------------------------------------------------------
679
680impl zsfm_core::Forecaster for FlowStateModel {
681    type Config = InferConfig;
682
683    fn load(gguf_path: &Path, config: InferConfig) -> Result<Self> {
684        FlowStateModel::load(gguf_path, config)
685    }
686
687    /// FlowState's `forecast()` is univariate-only; `context`/`mask` must carry exactly one
688    /// variate (mask is currently unused — the missing-value channel is always fed as
689    /// "known" since callers don't currently thread the mask through).
690    fn forecast(
691        &self,
692        context: &[Vec<f32>],
693        _mask: &[Vec<bool>],
694        horizon: usize,
695    ) -> Result<zsfm_core::QuantileMatrix> {
696        anyhow::ensure!(context.len() == 1, "FlowStateModel only supports univariate forecasting (1 variate)");
697        let qmat = FlowStateModel::forecast(self, &context[0], horizon)?; // [n_q][horizon]
698        Ok(qmat.into_iter().map(|row| vec![row]).collect())
699    }
700}
zsfm_flowstate/infer/mod.rs

zsfm_flowstate/infer/
mod.rs
