ganesh 0.27.0

Minimization and sampling in Rust, simplified
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

use ganesh::{
    algorithms::mcmc::{
        aies::AIESInit, ess::ESSInit, AIESConfig, ChainStorageMode, ESSConfig, ESSMove, AIES, ESS,
    },
    core::MaxSteps,
    test_functions::rosenbrock::Rosenbrock,
    traits::Algorithm,
    DVector, Float,
};
use std::{env, process};

fn parse_usize_arg(args: &[String], index: usize, default: usize, name: &str) -> usize {
    args.get(index)
        .map(|value| {
            value.parse::<usize>().unwrap_or_else(|_| {
                eprintln!("invalid {}: {}", name, value);
                process::exit(2);
            })
        })
        .unwrap_or(default)
}

fn parse_optional_usize_arg(args: &[String], index: usize, name: &str) -> Option<usize> {
    args.get(index).map(|value| {
        value.parse::<usize>().unwrap_or_else(|_| {
            eprintln!("invalid {}: {}", name, value);
            process::exit(2);
        })
    })
}

fn walkers(n_dim: usize, n_walkers: usize) -> Vec<DVector<Float>> {
    (0..n_walkers)
        .map(|i| {
            DVector::from_fn(n_dim, |j, _| {
                let center = if (i + j) % 2 == 0 { -1.5 } else { 1.5 };
                center + 0.02 * i as Float + 0.01 * j as Float
            })
        })
        .collect()
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let sampler = args.get(1).map(String::as_str).unwrap_or("aies");
    let n_dim = parse_usize_arg(&args, 2, 4, "dimension");
    let n_walkers = parse_usize_arg(&args, 3, 32, "walker count");
    let n_steps = parse_usize_arg(&args, 4, 2000, "step count");
    let chain_storage = match args.get(5).map(String::as_str) {
        None => ChainStorageMode::Full,
        Some("rolling") => ChainStorageMode::Rolling {
            window: parse_usize_arg(&args, 6, 64, "rolling window"),
        },
        Some("sampled") => ChainStorageMode::Sampled {
            keep_every: parse_usize_arg(&args, 6, 10, "sample keep_every"),
            max_samples: match args.get(7).map(String::as_str) {
                None | Some("none") => None,
                Some(_) => parse_optional_usize_arg(&args, 7, "sample max_samples"),
            },
        },
        Some(value) => ChainStorageMode::Rolling {
            window: value.parse::<usize>().unwrap_or_else(|_| {
                eprintln!(
                    "invalid chain storage mode: {} (expected integer window, 'rolling', or 'sampled')",
                    value
                );
                process::exit(2);
            }),
        },
    };

    let problem = Rosenbrock { n: n_dim };
    let x0 = walkers(n_dim, n_walkers);

    match sampler {
        "aies" => {
            let mut solver = AIES::default();
            let init = AIESInit::new(x0).unwrap();
            let config = AIESConfig::default().with_chain_storage(chain_storage);
            let summary = solver
                .process(
                    &problem,
                    &(),
                    init,
                    config,
                    AIES::default_callbacks().with_terminator(MaxSteps(n_steps)),
                )
                .unwrap();
            println!(
                "sampler=aies walkers={} steps={} dim={} evals={} chain_storage={:?}",
                n_walkers, n_steps, n_dim, summary.n_f_evals, summary.chain_storage
            );
        }
        "ess" => {
            let mut solver = ESS::default();
            let init = ESSInit::new(x0).unwrap();
            let config = ESSConfig::default()
                .with_moves([ESSMove::gaussian(0.2), ESSMove::differential(0.8)])
                .unwrap()
                .with_n_adaptive(10)
                .with_max_steps(128)
                .with_chain_storage(chain_storage);
            let summary = solver
                .process(
                    &problem,
                    &(),
                    init,
                    config,
                    ESS::default_callbacks().with_terminator(MaxSteps(n_steps)),
                )
                .unwrap();
            println!(
                "sampler=ess walkers={} steps={} dim={} evals={} chain_storage={:?}",
                n_walkers, n_steps, n_dim, summary.n_f_evals, summary.chain_storage
            );
        }
        other => {
            eprintln!("unknown sampler: {} (expected 'aies' or 'ess')", other);
            process::exit(2);
        }
    }
}