Skip to main content

pmcore/algorithms/parametric/
saem_config.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Deserialize, Clone, Serialize)]
4#[serde(deny_unknown_fields, default)]
5pub struct SaemConfig {
6    pub k1_iterations: usize,
7    pub k2_iterations: usize,
8    pub burn_in: usize,
9    pub sa_iterations: usize,
10    pub sa_cooling_factor: f64,
11    pub mcmc_step_size: f64,
12    pub rw_init: f64,
13    pub n_chains: usize,
14    pub mcmc_iterations: usize,
15    pub omega_min_variance: f64,
16    pub use_gibbs: bool,
17    pub n_kernels: usize,
18    pub transform_par: Vec<u8>,
19    pub compute_map: bool,
20    pub compute_fim: bool,
21    pub compute_ll_is: bool,
22    pub compute_ll_gq: bool,
23    pub n_mc_is: usize,
24    pub nu_is: usize,
25    pub n_nodes_gq: usize,
26    pub n_sd_gq: f64,
27    pub display_progress: usize,
28    pub seed: u64,
29    pub fix_seed: bool,
30}
31
32impl Default for SaemConfig {
33    fn default() -> Self {
34        Self {
35            k1_iterations: 300,
36            k2_iterations: 100,
37            burn_in: 5,
38            sa_iterations: 0,
39            sa_cooling_factor: 0.97,
40            mcmc_step_size: 0.4,
41            rw_init: 0.5,
42            n_chains: 1,
43            mcmc_iterations: 1,
44            omega_min_variance: 1e-6,
45            use_gibbs: false,
46            n_kernels: 4,
47            transform_par: vec![],
48            compute_map: true,
49            compute_fim: true,
50            compute_ll_is: true,
51            compute_ll_gq: false,
52            n_mc_is: 5000,
53            nu_is: 4,
54            n_nodes_gq: 12,
55            n_sd_gq: 4.0,
56            display_progress: 10,
57            seed: 123456,
58            fix_seed: true,
59        }
60    }
61}
62
63impl SaemConfig {
64    /// Creates a new `SaemConfig` with default values.
65    pub fn new() -> Self {
66        Self::default()
67    }
68
69    /// Number of exploration-phase (K1) iterations.
70    pub fn k1_iterations(mut self, iterations: usize) -> Self {
71        self.k1_iterations = iterations;
72        self
73    }
74
75    /// Number of smoothing-phase (K2) iterations.
76    pub fn k2_iterations(mut self, iterations: usize) -> Self {
77        self.k2_iterations = iterations;
78        self
79    }
80
81    /// Number of burn-in iterations.
82    pub fn burn_in(mut self, burn_in: usize) -> Self {
83        self.burn_in = burn_in;
84        self
85    }
86
87    /// Number of MCMC chains.
88    pub fn n_chains(mut self, n_chains: usize) -> Self {
89        self.n_chains = n_chains;
90        self
91    }
92
93    /// MCMC step size.
94    pub fn mcmc_step_size(mut self, step_size: f64) -> Self {
95        self.mcmc_step_size = step_size;
96        self
97    }
98
99    /// Random-number-generator seed.
100    pub fn seed(mut self, seed: u64) -> Self {
101        self.seed = seed;
102        self
103    }
104
105    pub fn total_iterations(&self) -> usize {
106        self.k1_iterations + self.k2_iterations
107    }
108
109    pub fn is_exploration_phase(&self, iteration: usize) -> bool {
110        iteration <= self.k1_iterations
111    }
112
113    pub fn is_smoothing_phase(&self, iteration: usize) -> bool {
114        iteration > self.k1_iterations
115    }
116
117    pub fn is_sa_active(&self, iteration: usize) -> bool {
118        self.sa_iterations > 0 && iteration <= self.sa_iterations
119    }
120
121    pub fn sa_temperature(&self, iteration: usize) -> f64 {
122        if self.is_sa_active(iteration) {
123            self.sa_cooling_factor.powi(iteration as i32)
124        } else {
125            1.0
126        }
127    }
128
129    pub fn step_size(&self, iteration: usize) -> f64 {
130        if iteration <= self.k1_iterations {
131            1.0
132        } else {
133            let k_smooth = iteration - self.k1_iterations;
134            1.0 / (k_smooth as f64 + 1.0)
135        }
136    }
137
138    pub fn get_transform(&self, param_idx: usize) -> u8 {
139        self.transform_par.get(param_idx).copied().unwrap_or(1)
140    }
141
142    pub fn get_transforms(&self, n_params: usize) -> Vec<u8> {
143        let mut transforms = self.transform_par.clone();
144        while transforms.len() < n_params {
145            transforms.push(1);
146        }
147        transforms.truncate(n_params);
148        transforms
149    }
150
151    pub fn infer_transforms_from_ranges(&mut self, ranges: &[(f64, f64)]) {
152        self.transform_par = ranges
153            .iter()
154            .map(|(lower, upper)| {
155                if *lower >= 0.0 && *upper > 0.0 && lower.is_finite() && upper.is_finite() {
156                    1
157                } else if (*lower - 0.0).abs() < 1e-10 && (*upper - 1.0).abs() < 1e-10 {
158                    3
159                } else {
160                    0
161                }
162            })
163            .collect();
164    }
165}