Skip to main content

pmcore/algorithms/nonparametric/
npag.rs

1use crate::algorithms::{NonParametricRunner, Status, StopReason};
2use crate::estimation::nonparametric::{
3    calculate_psi, CycleLog, NPCycle, NonParametricResult, Psi, Theta, Weights,
4};
5
6pub(crate) use crate::estimation::nonparametric::ipm::burke;
7pub(crate) use crate::estimation::nonparametric::qr;
8
9use anyhow::bail;
10use anyhow::Result;
11use pharmsol::prelude::{
12    data::{AssayErrorModels, Data},
13    simulator::Equation,
14};
15
16use pharmsol::prelude::AssayErrorModel;
17
18use crate::estimation::nonparametric::adaptative_grid;
19
20use super::error_optim::{optimize_error_models, ErrorOptimConfig};
21
22use serde::{Deserialize, Serialize};
23
24/// Configuration options for the Non-Parametric Adaptive Grid (NPAG) algorithm.
25#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
26pub struct NpagConfig {
27    pub eps: f64,
28    pub min_eps: f64,
29    pub objective_tolerance: f64,
30    pub pyl_tolerance: f64,
31    pub prune_threshold: f64,
32    pub qr_tolerance: f64,
33    pub grid_tolerance: f64,
34    pub error_optim: ErrorOptimConfig,
35    pub max_cycles: usize,
36    pub progress: bool,
37}
38
39impl Default for NpagConfig {
40    fn default() -> Self {
41        Self {
42            eps: 0.2,
43            min_eps: 1e-4,
44            objective_tolerance: 1e-4,
45            pyl_tolerance: 1e-2,
46            prune_threshold: 1e-3,
47            qr_tolerance: 1e-8,
48            grid_tolerance: 1e-4,
49            error_optim: ErrorOptimConfig::default(),
50            max_cycles: 1000,
51            progress: true,
52        }
53    }
54}
55
56impl NpagConfig {
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    pub fn eps(mut self, eps: f64) -> Self {
62        self.eps = eps;
63        self
64    }
65
66    pub fn min_eps(mut self, min_eps: f64) -> Self {
67        self.min_eps = min_eps;
68        self
69    }
70
71    pub fn objective_tolerance(mut self, tolerance: f64) -> Self {
72        self.objective_tolerance = tolerance;
73        self
74    }
75
76    pub fn pyl_tolerance(mut self, tolerance: f64) -> Self {
77        self.pyl_tolerance = tolerance;
78        self
79    }
80
81    pub fn prune_threshold(mut self, threshold: f64) -> Self {
82        self.prune_threshold = threshold;
83        self
84    }
85
86    pub fn qr_tolerance(mut self, tolerance: f64) -> Self {
87        self.qr_tolerance = tolerance;
88        self
89    }
90
91    pub fn grid_tolerance(mut self, tolerance: f64) -> Self {
92        self.grid_tolerance = tolerance;
93        self
94    }
95
96    pub fn error_optim(mut self, config: ErrorOptimConfig) -> Self {
97        self.error_optim = config;
98        self
99    }
100
101    pub fn max_cycles(mut self, cycles: usize) -> Self {
102        self.max_cycles = cycles;
103        self
104    }
105
106    pub fn progress(mut self, progress: bool) -> Self {
107        self.progress = progress;
108        self
109    }
110}
111
112#[derive(Debug)]
113pub struct NPAG<E: Equation + Send + 'static> {
114    equation: E,
115    ranges: Vec<(f64, f64)>,
116    psi: Psi,
117    prior: Theta,
118    theta: Theta,
119    lambda: Weights,
120    w: Weights,
121    eps: f64,
122    last_objf: f64,
123    objf: f64,
124    f0: f64,
125    f1: f64,
126    cycle: usize,
127    gamma_delta: Vec<f64>,
128    error_models: AssayErrorModels,
129    status: Status,
130    cycle_log: CycleLog,
131    data: Data,
132    config: NpagConfig,
133}
134
135impl<E: Equation + Send + 'static> NPAG<E> {
136    /// Construct an `NPAG` instance from explicit parts.
137    ///
138    /// The `parameter_space` is used solely to derive the finite bounds for the
139    /// adaptive grid. Initial support points can be supplied separately via
140    /// [`NonParametricRunner::set_theta`].
141    pub(crate) fn from_parts(
142        equation: E,
143        data: Data,
144        error_models: AssayErrorModels,
145        theta: Theta,
146        config: NpagConfig,
147    ) -> Result<Self> {
148        let ranges = theta.parameters().finite_ranges();
149        let gamma_delta = vec![config.error_optim.step; error_models.len()];
150        let eps = config.eps;
151
152        Ok(Self {
153            equation,
154            ranges,
155            psi: Psi::new(),
156            prior: theta.clone(),
157            theta,
158            lambda: Weights::default(),
159            w: Weights::default(),
160            eps,
161            last_objf: -1e30,
162            objf: f64::NEG_INFINITY,
163            f0: -1e30,
164            f1: f64::default(),
165            cycle: 0,
166            gamma_delta,
167            error_models,
168            status: Status::Continue,
169            cycle_log: CycleLog::new(),
170            data,
171            config,
172        })
173    }
174}
175
176impl<E: Equation + Send + 'static> NonParametricRunner<E> for NPAG<E> {
177    fn equation(&self) -> &E {
178        &self.equation
179    }
180
181    fn into_result(&self) -> Result<NonParametricResult<E>> {
182        NonParametricResult::new(
183            self.equation.clone(),
184            self.data.clone(),
185            self.error_models.clone(),
186            self.prior.clone(),
187            self.theta.clone(),
188            self.psi.clone(),
189            self.w.clone(),
190            -2. * self.objf,
191            self.cycle,
192            self.status.clone(),
193            self.cycle_log.clone(),
194        )
195    }
196
197    fn error_models(&self) -> &AssayErrorModels {
198        &self.error_models
199    }
200
201    fn data(&self) -> &Data {
202        &self.data
203    }
204
205    fn likelihood(&self) -> f64 {
206        self.objf
207    }
208
209    fn increment_cycle(&mut self) -> usize {
210        self.cycle += 1;
211        self.cycle
212    }
213
214    fn cycle(&self) -> usize {
215        self.cycle
216    }
217
218    fn set_theta(&mut self, theta: Theta) {
219        self.theta = theta;
220    }
221
222    fn theta(&self) -> &Theta {
223        &self.theta
224    }
225
226    fn psi(&self) -> &Psi {
227        &self.psi
228    }
229
230    fn evaluation(&mut self) -> Result<Status> {
231        tracing::info!("Objective function = {:.4}", -2.0 * self.objf);
232        tracing::debug!("Support points: {}", self.theta.nspp());
233
234        self.error_models.iter().for_each(|(outeq, em)| {
235            if AssayErrorModel::None == *em {
236                return;
237            }
238            tracing::debug!(
239                "Error model for outeq {}: {:.2}",
240                outeq,
241                em.factor().unwrap_or_default()
242            );
243        });
244
245        tracing::debug!("EPS = {:.4}", self.eps);
246        // Increasing objf signals instability or model misspecification.
247        if self.last_objf > self.objf + 1e-4 {
248            tracing::warn!(
249                "Objective function decreased from {:.4} to {:.4} (delta = {})",
250                -2.0 * self.last_objf,
251                -2.0 * self.objf,
252                -2.0 * self.last_objf - -2.0 * self.objf
253            );
254        }
255
256        let psi = self.psi.matrix();
257        let w = &self.w;
258        if (self.last_objf - self.objf).abs() <= self.config.objective_tolerance
259            && self.eps > self.config.min_eps
260        {
261            self.eps /= 2.;
262            if self.eps <= self.config.min_eps {
263                let pyl = psi * w.weights();
264                self.f1 = pyl.iter().map(|x| x.ln()).sum();
265                if (self.f1 - self.f0).abs() <= self.config.pyl_tolerance {
266                    tracing::info!("The model converged after {} cycles", self.cycle,);
267                    self.set_status(Status::Stop(StopReason::Converged));
268                    self.log_cycle_state();
269                    return Ok(self.status().clone());
270                } else {
271                    self.f0 = self.f1;
272                    self.eps = self.config.eps;
273                }
274            }
275        }
276
277        // Stop if we have reached maximum number of cycles
278        if self.cycle >= self.config.max_cycles {
279            tracing::warn!("Maximum number of cycles reached");
280            self.set_status(Status::Stop(StopReason::MaxCycles));
281            self.log_cycle_state();
282            return Ok(self.status().clone());
283        }
284
285        // Stop if stopfile exists
286        if std::path::Path::new("stop").exists() {
287            tracing::warn!("Stopfile detected - breaking");
288            self.set_status(Status::Stop(StopReason::StopFile));
289            self.log_cycle_state();
290            return Ok(self.status().clone());
291        }
292
293        // Continue with normal operation
294        self.set_status(Status::Continue);
295        self.log_cycle_state();
296        Ok(self.status().clone())
297    }
298
299    fn estimation(&mut self) -> Result<()> {
300        self.psi = calculate_psi(
301            &self.equation,
302            &self.data,
303            &self.theta,
304            &self.error_models,
305            self.cycle == 1 && self.config.progress,
306        )?;
307
308        if let Err(err) = self.check_zero_probability_subjects() {
309            bail!(err);
310        }
311
312        (self.lambda, _) = match burke(&self.psi) {
313            Ok((lambda, objf)) => (lambda, objf),
314            Err(err) => {
315                bail!("Error in IPM during estimation: {:?}", err);
316            }
317        };
318        Ok(())
319    }
320
321    fn condensation(&mut self) -> Result<()> {
322        // Filter out the support points with lambda < max(lambda)/1000
323
324        let max_lambda = self
325            .lambda
326            .iter()
327            .fold(f64::NEG_INFINITY, |acc, x| x.max(acc));
328
329        let mut keep = Vec::<usize>::new();
330        for (index, lam) in self.lambda.iter().enumerate() {
331            if lam > max_lambda * self.config.prune_threshold {
332                keep.push(index);
333            }
334        }
335        if self.psi.matrix().ncols() != keep.len() {
336            tracing::debug!(
337                "Lambda (max/1000) dropped {} support point(s)",
338                self.psi.matrix().ncols() - keep.len(),
339            );
340        }
341
342        self.theta.filter_indices(keep.as_slice());
343        self.psi.filter_column_indices(keep.as_slice());
344
345        //Rank-Revealing Factorization
346        let (r, perm) = qr::qrd(&self.psi)?;
347
348        let mut keep = Vec::<usize>::new();
349
350        // The minimum between the number of subjects and the actual number of support points
351        let keep_n = self.psi.matrix().ncols().min(self.psi.matrix().nrows());
352        for i in 0..keep_n {
353            let test = r.col(i).norm_l2();
354            let r_diag_val = r.get(i, i);
355            let ratio = r_diag_val / test;
356            if ratio.abs() >= self.config.qr_tolerance {
357                keep.push(*perm.get(i).unwrap());
358            }
359        }
360
361        // If a support point is dropped, log it as a debug message
362        if self.psi.matrix().ncols() != keep.len() {
363            tracing::debug!(
364                "QR decomposition dropped {} support point(s)",
365                self.psi.matrix().ncols() - keep.len(),
366            );
367        }
368
369        // Filter to keep only the support points (rows) that are in the `keep` vector
370        self.theta.filter_indices(keep.as_slice());
371        // Filter to keep only the support points (columns) that are in the `keep` vector
372        self.psi.filter_column_indices(keep.as_slice());
373
374        self.check_zero_probability_subjects()?;
375        (self.lambda, self.objf) = match burke(&self.psi) {
376            Ok((lambda, objf)) => (lambda, objf),
377            Err(err) => {
378                return Err(anyhow::anyhow!(
379                    "Error in IPM during condensation: {:?}",
380                    err
381                ));
382            }
383        };
384        self.w = self.lambda.clone();
385        Ok(())
386    }
387
388    fn optimizations(&mut self) -> Result<()> {
389        optimize_error_models(
390            &self.equation,
391            &self.data,
392            &self.theta,
393            &mut self.error_models,
394            &mut self.gamma_delta,
395            &mut self.objf,
396            &mut self.lambda,
397            &mut self.psi,
398            &self.config.error_optim,
399        )
400    }
401
402    fn expansion(&mut self) -> Result<()> {
403        adaptative_grid(
404            &mut self.theta,
405            self.eps,
406            &self.ranges,
407            self.config.grid_tolerance,
408        )?;
409        Ok(())
410    }
411
412    fn set_status(&mut self, status: Status) {
413        self.status = status;
414    }
415
416    fn status(&self) -> &Status {
417        &self.status
418    }
419
420    fn log_cycle_state(&mut self) {
421        let state = NPCycle::new(
422            self.cycle,
423            -2. * self.objf,
424            self.error_models.clone(),
425            self.theta.clone(),
426            self.w.clone(),
427            self.theta.nspp(),
428            (self.last_objf - self.objf).abs(),
429            self.status.clone(),
430        );
431        self.cycle_log.push(state);
432        self.last_objf = self.objf;
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use crate::prelude::*;
439
440    use pharmsol::{fa, fetch_params, lag, Subject, SubjectBuilderExt};
441
442    fn simple_equation() -> pharmsol::equation::ODE {
443        pharmsol::equation::ODE::new(
444            |x, p, _t, dx, b, _rateiv, _cov| {
445                fetch_params!(p, ke);
446                dx[0] = -ke * x[0] + b[0];
447            },
448            |_p, _t, _cov| lag! {},
449            |_p, _t, _cov| fa! {},
450            |_p, _t, _cov, _x| {},
451            |x, p, _t, _cov, y| {
452                fetch_params!(p, v);
453                y[0] = x[0] / v;
454            },
455        )
456        .with_nstates(1)
457        .with_ndrugs(1)
458        .with_nout(1)
459        .with_metadata(
460            pharmsol::equation::metadata::new("npag_settings_test")
461                .parameters(["ke", "v"])
462                .states(["central"])
463                .outputs(["0"])
464                .route(pharmsol::equation::Route::bolus("0").to_state("central")),
465        )
466        .expect("metadata attachment should validate")
467    }
468
469    fn simple_data() -> Data {
470        let subject = Subject::builder("1")
471            .bolus(0.0, 100.0, 0)
472            .observation(1.0, 10.0, 0)
473            .observation(2.0, 8.0, 0)
474            .build();
475
476        Data::new(vec![subject])
477    }
478
479    #[test]
480    fn npag_runs_without_error() {
481        let parameters = ParameterSpace::bounded()
482            .add("ke", 0.001, 3.0)
483            .add("v", 25.0, 250.0);
484        let prior = Theta::sobol_default(&parameters).expect("Failed to build prior");
485        let error_models = AssayErrorModels::new()
486            .add(
487                "0",
488                AssayErrorModel::additive(ErrorPoly::new(0.0, 0.5, 0.0, 0.0), 0.0),
489            )
490            .expect("Failed to build error models");
491        let problem =
492            EstimationProblem::nonparametric(simple_equation(), simple_data(), prior, error_models)
493                .expect("Failed to build problem");
494
495        let result = problem.fit_with(NonParametricAlgorithm::npag());
496
497        assert!(
498            result.is_ok(),
499            "NPAG algorithm should run without error, but got: {:?}",
500            result.err()
501        );
502    }
503}