flowmatch 0.1.6

Flow matching primitives (ndarray-first; backend-agnostic) with semidiscrete FM and RFM experiments.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

//! Verify that Euler achieves O(h) and Heun achieves O(h^2) convergence.
//!
//! Uses dx/dt = -x with exact solution x(T) = x(0) * exp(-T).
//! Runs at multiple step counts and checks the error decay rate.

use flowmatch::ode::{integrate_fixed, OdeMethod};
use ndarray::Array1;

/// Solve dx/dt = -x from t=0 to t=1, return absolute error vs exact.
fn error_at_steps(method: OdeMethod, steps: usize) -> f64 {
    let x0 = Array1::from_vec(vec![1.0f32]);
    let exact = (-1.0f32).exp();
    let dt = 1.0f32 / (steps as f32);

    let result = integrate_fixed(method, &x0, 0.0, dt, steps, |x, _t| {
        Ok(Array1::from_vec(vec![-x[0]]))
    })
    .unwrap();

    (result[0] - exact).abs() as f64
}

/// Estimate convergence order from errors at doubling step counts.
///
/// Returns the average of log2(err[i] / err[i+1]) across consecutive pairs.
fn estimate_order(method: OdeMethod, step_counts: &[usize]) -> f64 {
    let errors: Vec<f64> = step_counts
        .iter()
        .map(|&s| error_at_steps(method, s))
        .collect();

    let mut ratios = Vec::new();
    for i in 0..errors.len() - 1 {
        if errors[i + 1] > 1e-12 {
            // When steps double, h halves. Order p means err(h/2) / err(h) ~ 2^{-p}.
            let ratio = (errors[i] / errors[i + 1]).ln() / 2.0f64.ln();
            ratios.push(ratio);
        }
    }

    ratios.iter().sum::<f64>() / ratios.len() as f64
}

#[test]
fn euler_achieves_first_order_convergence() {
    // Doubling step counts: 10, 20, 40, 80, 160.
    let steps = vec![10, 20, 40, 80, 160];
    let order = estimate_order(OdeMethod::Euler, &steps);

    assert!(
        (0.85..1.15).contains(&order),
        "Euler should be ~O(h^1), estimated order = {order:.3}"
    );
}

#[test]
fn heun_achieves_second_order_convergence() {
    // Doubling step counts: 10, 20, 40, 80, 160.
    let steps = vec![10, 20, 40, 80, 160];
    let order = estimate_order(OdeMethod::Heun, &steps);

    assert!(
        (1.85..2.15).contains(&order),
        "Heun should be ~O(h^2), estimated order = {order:.3}"
    );
}

/// Same test on a second ODE to confirm rates are not problem-specific.
///
/// dx/dt = cos(t), x(0) = 0. Exact: x(T) = sin(T).
fn error_cosine_at_steps(method: OdeMethod, steps: usize) -> f64 {
    let x0 = Array1::from_vec(vec![0.0f32]);
    let t_final = 1.0f32;
    let exact = t_final.sin();
    let dt = t_final / (steps as f32);

    let result = integrate_fixed(method, &x0, 0.0, dt, steps, |_x, t| {
        Ok(Array1::from_vec(vec![t.cos()]))
    })
    .unwrap();

    (result[0] - exact).abs() as f64
}

#[test]
fn euler_first_order_on_cosine_ode() {
    let steps = [20, 40, 80, 160, 320];
    let errors: Vec<f64> = steps
        .iter()
        .map(|&s| error_cosine_at_steps(OdeMethod::Euler, s))
        .collect();

    let mut ratios = Vec::new();
    for i in 0..errors.len() - 1 {
        if errors[i + 1] > 1e-12 {
            ratios.push((errors[i] / errors[i + 1]).ln() / 2.0f64.ln());
        }
    }
    let order = ratios.iter().sum::<f64>() / ratios.len() as f64;

    assert!(
        (0.85..1.15).contains(&order),
        "Euler should be ~O(h^1) on cosine ODE, estimated order = {order:.3}"
    );
}

#[test]
fn heun_second_order_on_cosine_ode() {
    // Keep step counts low to stay in the truncation-dominated regime (f32).
    let steps = [5, 10, 20, 40, 80];
    let errors: Vec<f64> = steps
        .iter()
        .map(|&s| error_cosine_at_steps(OdeMethod::Heun, s))
        .collect();

    let mut ratios = Vec::new();
    for i in 0..errors.len() - 1 {
        if errors[i + 1] > 1e-12 {
            ratios.push((errors[i] / errors[i + 1]).ln() / 2.0f64.ln());
        }
    }
    let order = ratios.iter().sum::<f64>() / ratios.len() as f64;

    assert!(
        (1.85..2.15).contains(&order),
        "Heun should be ~O(h^2) on cosine ODE, estimated order = {order:.3}"
    );
}