Skip to main content

pmcore/algorithms/nonparametric/
npod.rs

1use crate::{
2    algorithms::{NonParametricRunner, Status, StopReason},
3    estimation::nonparametric::{
4        calculate_psi, ipm::burke, qr, CycleLog, NPCycle, NonParametricResult, Psi, Theta, Weights,
5    },
6};
7use pharmsol::ParameterOptimizer;
8
9use anyhow::bail;
10use anyhow::Result;
11use pharmsol::prelude::{data::Data, simulator::Equation};
12use pharmsol::{prelude::AssayErrorModel, AssayErrorModels};
13
14use ndarray::Array1;
15use rayon::prelude::{IntoParallelRefMutIterator, ParallelIterator};
16use serde::{Deserialize, Serialize};
17
18use super::error_optim::{optimize_error_models, ErrorOptimConfig};
19
20const THETA_F: f64 = 1e-2;
21const THETA_D: f64 = 1e-4;
22
23/// Configuration options for the Non-Parametric Optimal Design (NPOD) algorithm.
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct NpodConfig {
26    /// Maximum number of cycles to run the algorithm for.
27    pub max_cycles: usize,
28    /// Configuration for the error-model factor (gamma/lambda) optimization.
29    pub error_optim: ErrorOptimConfig,
30    /// Whether to print progress information during the first cycle.
31    pub progress: bool,
32}
33
34impl NpodConfig {
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    pub fn max_cycles(mut self, cycles: usize) -> Self {
40        self.max_cycles = cycles;
41        self
42    }
43
44    pub fn error_optim(mut self, config: ErrorOptimConfig) -> Self {
45        self.error_optim = config;
46        self
47    }
48
49    pub fn progress(mut self, progress: bool) -> Self {
50        self.progress = progress;
51        self
52    }
53}
54
55impl Default for NpodConfig {
56    fn default() -> Self {
57        Self {
58            max_cycles: 100,
59            error_optim: ErrorOptimConfig::default(),
60            progress: true,
61        }
62    }
63}
64
65#[derive(Debug)]
66pub struct NPOD<E: Equation + Send + 'static> {
67    equation: E,
68    psi: Psi,
69    prior: Theta,
70    theta: Theta,
71    lambda: Weights,
72    w: Weights,
73    last_objf: f64,
74    objf: f64,
75    cycle: usize,
76    gamma_delta: Vec<f64>,
77    error_models: AssayErrorModels,
78    converged: bool,
79    status: Status,
80    cycle_log: CycleLog,
81    data: Data,
82    config: NpodConfig,
83}
84
85impl<E: Equation + Send + 'static> NPOD<E> {
86    pub(crate) fn from_parts(
87        equation: E,
88        data: Data,
89        error_models: AssayErrorModels,
90        theta: Theta,
91        config: NpodConfig,
92    ) -> Result<Self> {
93        let gamma_delta = vec![config.error_optim.step; error_models.len()];
94
95        Ok(Self {
96            equation,
97            psi: Psi::new(),
98            prior: theta.clone(),
99            theta,
100            lambda: Weights::default(),
101            w: Weights::default(),
102            last_objf: -1e30,
103            objf: f64::NEG_INFINITY,
104            cycle: 0,
105            gamma_delta,
106            error_models,
107            converged: false,
108            status: Status::Continue,
109            cycle_log: CycleLog::new(),
110            data,
111            config,
112        })
113    }
114}
115
116impl<E: Equation + Send + 'static> NonParametricRunner<E> for NPOD<E> {
117    fn into_result(&self) -> Result<NonParametricResult<E>> {
118        NonParametricResult::new(
119            self.equation.clone(),
120            self.data.clone(),
121            self.error_models.clone(),
122            self.prior.clone(),
123            self.theta.clone(),
124            self.psi.clone(),
125            self.w.clone(),
126            -2. * self.objf,
127            self.cycle,
128            self.status.clone(),
129            self.cycle_log.clone(),
130        )
131    }
132
133    fn equation(&self) -> &E {
134        &self.equation
135    }
136
137    fn error_models(&self) -> &AssayErrorModels {
138        &self.error_models
139    }
140
141    fn data(&self) -> &Data {
142        &self.data
143    }
144
145    fn increment_cycle(&mut self) -> usize {
146        self.cycle += 1;
147        self.cycle
148    }
149
150    fn cycle(&self) -> usize {
151        self.cycle
152    }
153
154    fn set_theta(&mut self, theta: Theta) {
155        self.theta = theta;
156    }
157
158    fn theta(&self) -> &Theta {
159        &self.theta
160    }
161
162    fn psi(&self) -> &Psi {
163        &self.psi
164    }
165
166    fn likelihood(&self) -> f64 {
167        self.objf
168    }
169
170    fn set_status(&mut self, status: Status) {
171        self.status = status;
172    }
173
174    fn status(&self) -> &Status {
175        &self.status
176    }
177
178    fn log_cycle_state(&mut self) {
179        let state = NPCycle::new(
180            self.cycle,
181            -2. * self.objf,
182            self.error_models.clone(),
183            self.theta.clone(),
184            self.w.clone(),
185            self.theta.nspp(),
186            (self.last_objf - self.objf).abs(),
187            self.status.clone(),
188        );
189        self.cycle_log.push(state);
190        self.last_objf = self.objf;
191    }
192
193    fn evaluation(&mut self) -> Result<Status> {
194        tracing::info!("Objective function = {:.4}", -2.0 * self.objf);
195        tracing::debug!("Support points: {}", self.theta.nspp());
196        self.error_models.iter().for_each(|(outeq, em)| {
197            if AssayErrorModel::None == *em {
198                return;
199            }
200            tracing::debug!(
201                "Error model for outeq {}: {:.16}",
202                outeq,
203                em.factor().unwrap_or_default()
204            );
205        });
206        if self.last_objf > self.objf + 1e-4 {
207            tracing::warn!(
208                "Objective function decreased from {:.4} to {:.4} (delta = {})",
209                -2.0 * self.last_objf,
210                -2.0 * self.objf,
211                -2.0 * self.last_objf - -2.0 * self.objf
212            );
213        }
214
215        if (self.last_objf - self.objf).abs() <= THETA_F {
216            tracing::info!("Objective function convergence reached");
217            self.converged = true;
218            self.set_status(Status::Stop(StopReason::Converged));
219            self.log_cycle_state();
220            return Ok(self.status.clone());
221        }
222
223        if self.cycle >= self.config.max_cycles {
224            tracing::warn!("Maximum number of cycles reached");
225            self.converged = true;
226            self.set_status(Status::Stop(StopReason::MaxCycles));
227            self.log_cycle_state();
228            return Ok(self.status.clone());
229        }
230
231        if std::path::Path::new("stop").exists() {
232            tracing::warn!("Stopfile detected - breaking");
233            self.converged = true;
234            self.set_status(Status::Stop(StopReason::StopFile));
235            self.log_cycle_state();
236            return Ok(self.status.clone());
237        }
238
239        self.status = Status::Continue;
240        self.log_cycle_state();
241        Ok(self.status.clone())
242    }
243
244    fn estimation(&mut self) -> Result<()> {
245        let error_model: AssayErrorModels = self.error_models.clone();
246
247        self.psi = calculate_psi(
248            &self.equation,
249            &self.data,
250            &self.theta,
251            &error_model,
252            self.cycle == 1 && self.config.progress,
253        )?;
254
255        if let Err(err) = self.check_zero_probability_subjects() {
256            bail!(err);
257        }
258
259        (self.lambda, _) = match burke(&self.psi) {
260            Ok((lambda, objf)) => (lambda, objf),
261            Err(err) => {
262                bail!(err);
263            }
264        };
265        Ok(())
266    }
267
268    fn condensation(&mut self) -> Result<()> {
269        let max_lambda = self
270            .lambda
271            .iter()
272            .fold(f64::NEG_INFINITY, |acc, x| x.max(acc));
273
274        let mut keep = Vec::<usize>::new();
275        for (index, lam) in self.lambda.iter().enumerate() {
276            if lam > max_lambda / 1000_f64 {
277                keep.push(index);
278            }
279        }
280        if self.psi.matrix().ncols() != keep.len() {
281            tracing::debug!(
282                "Lambda (max/1000) dropped {} support point(s)",
283                self.psi.matrix().ncols() - keep.len(),
284            );
285        }
286
287        self.theta.filter_indices(keep.as_slice());
288        self.psi.filter_column_indices(keep.as_slice());
289
290        let (r, perm) = qr::qrd(&self.psi)?;
291
292        let mut keep = Vec::<usize>::new();
293        let keep_n = self.psi.matrix().ncols().min(self.psi.matrix().nrows());
294        for i in 0..keep_n {
295            let test = r.col(i).norm_l2();
296            let r_diag_val = r.get(i, i);
297            let ratio = r_diag_val / test;
298            if ratio.abs() >= 1e-8 {
299                keep.push(*perm.get(i).unwrap());
300            }
301        }
302
303        if self.psi.matrix().ncols() != keep.len() {
304            tracing::debug!(
305                "QR decomposition dropped {} support point(s)",
306                self.psi.matrix().ncols() - keep.len(),
307            );
308        }
309
310        self.theta.filter_indices(keep.as_slice());
311        self.psi.filter_column_indices(keep.as_slice());
312
313        (self.lambda, self.objf) = match burke(&self.psi) {
314            Ok((lambda, objf)) => (lambda, objf),
315            Err(err) => {
316                return Err(anyhow::anyhow!("Error in IPM: {:?}", err));
317            }
318        };
319        self.w = self.lambda.clone();
320        Ok(())
321    }
322
323    fn optimizations(&mut self) -> Result<()> {
324        optimize_error_models(
325            &self.equation,
326            &self.data,
327            &self.theta,
328            &mut self.error_models,
329            &mut self.gamma_delta,
330            &mut self.objf,
331            &mut self.lambda,
332            &mut self.psi,
333            &self.config.error_optim,
334        )
335    }
336
337    fn expansion(&mut self) -> Result<()> {
338        let pyl_col = self.psi().matrix().as_ref() * self.w.weights().as_ref();
339        let pyl: Array1<f64> = pyl_col.iter().copied().collect();
340
341        let error_model: AssayErrorModels = self.error_models.clone();
342
343        let mut candididate_points: Vec<Array1<f64>> = Vec::default();
344        for spp in self.theta.matrix().row_iter() {
345            let candidate: Vec<f64> = spp.iter().cloned().collect();
346            let spp = Array1::from(candidate);
347            candididate_points.push(spp.to_owned());
348        }
349        candididate_points.par_iter_mut().for_each(|spp| {
350            let optimizer = ParameterOptimizer::new(&self.equation, &self.data, &error_model, &pyl);
351            let candidate_point = optimizer.optimize_point(spp.to_owned()).unwrap();
352            *spp = candidate_point;
353        });
354        for cp in candididate_points {
355            self.theta.suggest_point(cp.to_vec().as_slice(), THETA_D)?;
356        }
357        Ok(())
358    }
359}