ferrotherm 0.39.0

// DTM at the published flagship scale: L=70 grid, G12 connectivity, T=8 chained EBMs, trained on
// binarized Fashion-MNIST. This is the configuration the architecture's own paper reports
// (arXiv:2510.23972): 4,900 nodes per layer, visible sites drawn uniformly at random.
//
// Reports the measured cost of a training step at that scale, then trains within a wall-clock
// budget and reports how the model's visible statistics approach the data's.
//
// usage: cargo run --release --example dtm_scale -- <fmnist-images> [seconds] [L] [T]
use ferrotherm::dtm::{forward_step, gamma_coupling, pattern_grid, Ebm, G12};
use ferrotherm::host::Timing;
use ferrotherm::rng::Pcg;
use std::time::{Duration, Instant};

fn load_images(path: &str, limit: usize) -> (Vec<Vec<i8>>, usize) {
    let raw = std::fs::read(path).expect("read images");
    let n = u32::from_be_bytes([raw[4], raw[5], raw[6], raw[7]]) as usize;
    let rows = u32::from_be_bytes([raw[8], raw[9], raw[10], raw[11]]) as usize;
    let cols = u32::from_be_bytes([raw[12], raw[13], raw[14], raw[15]]) as usize;
    let px = rows * cols;
    let take = n.min(limit);
    let mut out = Vec::with_capacity(take);
    for i in 0..take {
        let off = 16 + i * px;
        out.push(raw[off..off + px].iter().map(|&v| if v > 127 { 1i8 } else { -1 }).collect());
    }
    (out, px)
}

fn main() {
    let a: Vec<String> = std::env::args().skip(1).collect();
    // A missing argument is a usage error, not a crash.
    //
    // This was `.expect(..)`, so running the example with no arguments printed "thread 'main'
    // panicked" and a backtrace note -- which reads as the example being broken rather than as the
    // caller having forgotten a path. Every other example in this directory runs with no arguments;
    // this one needs a dataset, and the difference should look deliberate.
    let Some(path) = a.first() else {
        eprintln!("usage: dtm_scale <fmnist-images> [steps] [L] [T]");
        eprintln!();
        eprintln!("  <fmnist-images>  an idx3-ubyte file, e.g. train-images-idx3-ubyte");
        eprintln!("  [steps]          TRAINING STEPS, not seconds          (default 2000)");
        eprintln!("  [L]              pattern-grid side                   (default 70)");
        eprintln!("  [T]              denoising steps                     (default 8)");
        std::process::exit(2);
    };
    // STEPS, NOT SECONDS, AND THE DIFFERENCE IS THE WHOLE POINT.
    //
    // This trained inside `while start.elapsed() < budget` with a default of 120 seconds, so the
    // quality it reported -- the per-pixel MAE that WORKLOADS.md publishes as the flagship
    // EBM-training result -- was a function of how fast the machine was and what else was running
    // on it. A faster box took more gradient steps and got a better number for the same command.
    // That is a division by wall-clock time reported as a property of the method, which is exactly
    // what this crate's `host` and `ledger` documentation warns against everywhere else.
    //
    // A step count is reproducible on any fabric. The wall clock survives only as a SAFETY STOP, so
    // an accidental five-million-step run on a slow machine can still be interrupted -- and it says
    // loudly when it fires, because a truncated run reporting a quality figure is the original
    // defect wearing a different hat.
    let steps_wanted: usize = a.get(1).and_then(|s| s.parse().ok()).unwrap_or(2000);
    let wall_cap = Duration::from_secs(
        std::env::var("FT_DTM_WALL_CAP").ok().and_then(|s| s.parse().ok()).unwrap_or(3600),
    );
    let l: usize = a.get(2).and_then(|s| s.parse().ok()).unwrap_or(70);
    let t_steps: usize = a.get(3).and_then(|s| s.parse().ok()).unwrap_or(8);

    let (data, px) = load_images(path, 60_000);
    println!("data: {} binarized images, {px} pixels each", data.len());

    // ---- the fabric, at the published scale ----
    let edges32 = pattern_grid(l, &G12);
    let n = l * l;
    let edges: Vec<(u16, u16)> = edges32.iter().map(|&(a, b)| (a as u16, b as u16)).collect();
    println!("grid: {l}x{l} = {n} nodes, G12, {} edges (mean degree {:.1})",
             edges.len(), 2.0 * edges.len() as f64 / n as f64);

    // visible sites drawn uniformly at random, as the paper specifies
    let mut rng = Pcg::new(0xD7A1, 1);
    let mut perm: Vec<u32> = (0..n as u32).collect();
    for i in (1..perm.len()).rev() {
        let j = (rng.f64() * (i + 1) as f64) as usize;
        perm.swap(i, j);
    }
    let visible: Vec<usize> = perm[..px].iter().map(|&v| v as usize).collect();
    println!("visible sites: {} of {n} (uniformly random), latent: {}", visible.len(), n - px);

    let mut layers: Vec<Ebm> = (0..t_steps).map(|_| Ebm::new(n, edges.clone())).collect();
    println!("layers: {t_steps}; bipartite: {}; parameters: {}",
             layers[0].is_bipartite(), t_steps * (edges.len() + n));

    // forward-process schedule
    let gamma_x = 1.0f64;
    let times: Vec<f64> = (0..=t_steps).map(|i| i as f64 * 0.35).collect();

    // ---- measured cost of one training step at this scale ----
    let k_sweeps = 25usize;
    let batch = 8usize;
    let (_, sweep_t) = Timing::around(|| {
        let mut s = vec![1i8; n];
        let extra = vec![0.0f64; n];
        layers[0].gibbs_chromatic(&mut s, &extra, 10, &mut rng);
    });
    // A per-sweep cost is a rate, and the training-step figure below is that rate multiplied by
    // four more factors -- so contention on this machine is amplified, not diluted, by the
    // projection. The budgeted loop further down is unaffected: a time budget is a stop condition,
    // and it reports the steps it actually completed.
    match sweep_t.as_measurement() {
        Some(secs) => {
            let per_sweep = secs / 10.0;
            println!("\nmeasured: {:.2} ms per chromatic sweep of {n} nodes ({:.2e} node updates/s)",
                     per_sweep * 1e3, n as f64 / per_sweep);
            let step_cost = per_sweep * k_sweeps as f64 * 2.0 * batch as f64 * t_steps as f64;
            println!("one training step (batch {batch}, K={k_sweeps}, both phases, {t_steps} layers): {:.2} s",
                     step_cost);
        }
        None => {
            println!("\nper-sweep cost: {sweep_t}");
            if let Some(c) = sweep_t.caveat() {
                println!("  {c}");
            }
            println!("the training-step projection is that rate x K x 2 x batch x layers, so it is \
                      omitted rather than multiplied by {}x.", k_sweeps * 2 * batch * t_steps);
        }
    }

    // ---- train within the budget ----
    println!("\ntraining for {steps_wanted} steps (wall cap {} s, a safety stop only)...", wall_cap.as_secs());
    let lr = 0.02;
    let start = Instant::now();
    let mut steps = 0usize;
    let mut pos_ss = vec![0.0f64; edges.len()];
    let mut pos_si = vec![0.0f64; n];
    let mut neg_ss = vec![0.0f64; edges.len()];
    let mut neg_si = vec![0.0f64; n];
    let all: Vec<usize> = (0..n).collect();
    let latents: Vec<usize> = (0..n).filter(|i| !visible.contains(i)).collect();

    let mut truncated = false;
    while steps < steps_wanted {
        if start.elapsed() >= wall_cap {
            truncated = true;
            break;
        }
        let t = steps % t_steps;
        let dt = times[t + 1] - times[t];
        let gam = gamma_coupling(gamma_x, dt, 2);
        pos_ss.iter_mut().for_each(|v| *v = 0.0);
        pos_si.iter_mut().for_each(|v| *v = 0.0);
        neg_ss.iter_mut().for_each(|v| *v = 0.0);
        neg_si.iter_mut().for_each(|v| *v = 0.0);

        for _ in 0..batch {
            // a data image, noised to time t and t+1
            let img = &data[(rng.f64() * data.len() as f64) as usize % data.len()];
            let mut x_prev: Vec<i8> = img.clone();
            if t > 0 {
                forward_step(&mut x_prev, gamma_x, times[t], &mut rng);
            }
            let mut x_next = x_prev.clone();
            forward_step(&mut x_next, gamma_x, dt, &mut rng);

            // the clamp field from x^{t+1} on the visible sites
            let mut extra = vec![0.0f64; n];
            for (k, &v) in visible.iter().enumerate() {
                extra[v] = 0.5 * gam * x_next[k] as f64;
            }
            // positive phase: visibles clamped to x^t, Gibbs over latents
            let mut s = vec![0i8; n];
            for i in 0..n {
                s[i] = if rng.f64() < 0.5 { 1 } else { -1 };
            }
            for (k, &v) in visible.iter().enumerate() {
                s[v] = x_prev[k];
            }
            layers[t].gibbs(&mut s, &latents, &extra, k_sweeps, &mut rng);
            layers[t].accumulate(&s, &mut pos_ss, &mut pos_si, 1.0);

            // negative phase: everything free under the same clamp field
            let mut sn = vec![0i8; n];
            for i in 0..n {
                sn[i] = if rng.f64() < 0.5 { 1 } else { -1 };
            }
            layers[t].gibbs(&mut sn, &all, &extra, k_sweeps, &mut rng);
            layers[t].accumulate(&sn, &mut neg_ss, &mut neg_si, 1.0);
        }

        let m = batch as f64;
        // The total-correlation penalty from the architecture's own paper. Without it the
        // negative phase must mix fully or the model's correlations are underestimated and the
        // couplings grow without bound: |J| climbing monotonically with no settling is that
        // failure, not learning. The gradient reuses the negative-phase statistics, so it is free.
        let lambda = 0.35;
        let lay = &mut layers[t];
        for k in 0..lay.edges.len() {
            let (a, b) = lay.edges[k];
            let (ma, mb) = (neg_si[a as usize] / m, neg_si[b as usize] / m);
            let tc = ma * mb - neg_ss[k] / m;
            lay.j[k] += lr * ((pos_ss[k] - neg_ss[k]) / m - lambda * tc);
        }
        for i in 0..n {
            lay.h[i] += lr * (pos_si[i] - neg_si[i]) / m;
        }
        steps += 1;
        if steps.is_multiple_of(40) {
            println!("  step {steps:5}  layer {t}  |J| mean {:.4}  elapsed {:.0} s",
                     lay.j.iter().map(|v| v.abs()).sum::<f64>() / lay.j.len() as f64,
                     start.elapsed().as_secs_f64());
        }
    }
    println!("completed {steps} steps in {:.0} s", start.elapsed().as_secs_f64());

    // ---- does the model reproduce the data's visible statistics? ----
    let mut data_m = vec![0.0f64; px];
    for img in data.iter().take(4000) {
        for k in 0..px {
            data_m[k] += img[k] as f64;
        }
    }
    data_m.iter_mut().for_each(|v| *v /= 4000.0);

    // reverse chain: start from noise, denoise through the layers
    let n_samples = 32usize;
    let mut model_m = vec![0.0f64; px];
    for _ in 0..n_samples {
        let mut x: Vec<i8> = (0..px).map(|_| if rng.f64() < 0.5 { 1 } else { -1 }).collect();
        for t in (0..t_steps).rev() {
            let dt = times[t + 1] - times[t];
            let gam = gamma_coupling(gamma_x, dt, 2);
            let mut extra = vec![0.0f64; n];
            for (k, &v) in visible.iter().enumerate() {
                extra[v] = 0.5 * gam * x[k] as f64;
            }
            let mut s = vec![0i8; n];
            for i in 0..n {
                s[i] = if rng.f64() < 0.5 { 1 } else { -1 };
            }
            layers[t].gibbs(&mut s, &all, &extra, k_sweeps, &mut rng);
            for (k, &v) in visible.iter().enumerate() {
                x[k] = s[v];
            }
        }
        for k in 0..px {
            model_m[k] += x[k] as f64;
        }
    }
    model_m.iter_mut().for_each(|v| *v /= n_samples as f64);

    let mae: f64 = (0..px).map(|k| (model_m[k] - data_m[k]).abs()).sum::<f64>() / px as f64;
    let data_mean: f64 = data_m.iter().sum::<f64>() / px as f64;
    let model_mean: f64 = model_m.iter().sum::<f64>() / px as f64;
    // an untrained chain outputs unbiased noise, so its per-pixel mean is ~0 and the MAE is
    // just the data's own |mean| profile: that is the number to beat.
    let baseline: f64 = data_m.iter().map(|v| v.abs()).sum::<f64>() / px as f64;
    println!("\nper-pixel mean activation:  data {data_mean:+.3}   model {model_mean:+.3}");
    // WHAT PRODUCED THIS NUMBER, printed beside it. A quality figure without its step count is not
    // reproducible on another machine, and that is how the published one came to be a function of
    // wall-clock speed rather than of the method.
    println!("\nrun: {steps} steps, L={l}, T={t_steps}, {} images, lr {lr}", data.len());
    println!("per-pixel MAE vs data:      model {mae:.3}   untrained-noise baseline {baseline:.3}");
    if truncated {
        println!(
            "\n*** TRUNCATED: the wall-clock safety stop fired at {:.0} s after {steps} of the\n\
             *** {steps_wanted} steps asked for. THIS NUMBER IS NOT THE ONE THIS COMMAND NAMES.\n\
             *** Quote it only with the step count above, or raise FT_DTM_WALL_CAP and re-run.",
            start.elapsed().as_secs_f64()
        );
    }
    println!("verdict: {}", if mae < baseline {
        format!("the trained chain is closer to the data than noise by {:.1}%", 100.0 * (1.0 - mae / baseline))
    } else {
        format!("no better than noise after {steps} steps")
    });
}