pmcore/algorithms/nonparametric/
ncnpag.rs

1use crate::{
2    algorithms::{
3        nonparametric::{npag::NPAG, NpagConfig},
4        NonParametricRunner, Status, StopReason,
5    },
6    estimation::nonparametric::{
7        calculate_psi, CycleLog, NPCycle, NonParametricResult, Psi, Theta, Weights,
8    },
9};
10
11use anyhow::Result;
12use faer::Mat;
13use pharmsol::prelude::{
14    data::{AssayErrorModels, Data},
15    simulator::Equation,
16};
17
18use serde::{Deserialize, Serialize};
19
20/// Configuration options for the non-collapsing NPAG (NCNPAG) algorithm.
21#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22pub struct NcnpagConfig {
23    /// Number of NPAG cycles used to refine each surviving support point.
24    ///
25    /// `0` disables refinement, leaving a pure Bayesian reweighting of the input
26    /// grid. The default (`500`) matches the historical NPAGFULL behavior.
27    pub cycles: usize,
28    /// Whether to show NPAG progress output during refinement.
29    pub progress: bool,
30}
31
32impl Default for NcnpagConfig {
33    fn default() -> Self {
34        Self {
35            cycles: 500,
36            progress: false,
37        }
38    }
39}
40
41impl NcnpagConfig {
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    /// Set the number of NPAG refinement cycles per support point (`0` to skip).
47    pub fn cycles(mut self, cycles: usize) -> Self {
48        self.cycles = cycles;
49        self
50    }
51
52    /// Enable or disable NPAG progress output during refinement.
53    pub fn progress(mut self, progress: bool) -> Self {
54        self.progress = progress;
55        self
56    }
57}
58
59/// Non-collapsing NPAG (NCNPAG) algorithm.
60///
61/// Individualizes a set of prior support points to a subject's data in two
62/// steps, without collapsing (merging) the points:
63///
64/// 1. **Bayesian filtering.** Evaluate the likelihood `P(data | θⱼ)` for each
65///    prior support point, apply a flat (uniform) prior so the posterior weight
66///    is proportional to the likelihood (`postⱼ ∝ ∏ᵢ P(dataᵢ | θⱼ)`), drop
67///    points whose normalized weight falls below `1e-100 × max`, and renormalize
68///    the survivors.
69/// 2. **Per-point NPAG refinement.** When `cycles > 0`, seed a full NPAG run
70///    from each surviving point and replace it with the resulting daughter
71///    point, preserving its filter weight. Points whose refinement fails or
72///    yields nothing are kept at their original location.
73///
74/// The result is returned as a standard [`NonParametricResult`], so the
75/// `(theta, weights)` can be consumed exactly like any other fit.
76pub struct NCNPAG<E: Equation + Send + 'static> {
77    equation: E,
78    psi: Psi,
79    theta: Theta,
80    w: Weights,
81    objf: f64,
82    cycle: usize,
83    status: Status,
84    data: Data,
85    cyclelog: CycleLog,
86    error_models: AssayErrorModels,
87    prior: Theta,
88    cycles: usize,
89    progress: bool,
90}
91
92impl<E: Equation + Send + 'static> NCNPAG<E> {
93    pub(crate) fn from_parts(
94        equation: E,
95        data: Data,
96        error_models: AssayErrorModels,
97        theta: Theta,
98        config: NcnpagConfig,
99    ) -> Result<Self> {
100        Ok(Self {
101            equation,
102            psi: Psi::new(),
103            theta: theta.clone(),
104            w: Weights::default(),
105            objf: f64::INFINITY,
106            cycle: 0,
107            status: Status::Continue,
108            data,
109            cyclelog: CycleLog::new(),
110            error_models,
111            prior: theta,
112            cycles: config.cycles,
113            progress: config.progress,
114        })
115    }
116}
117
118/// Refine each support point with a full NPAG seeded from that single point,
119/// preserving the point's filter weight (NPAGFULL). Points whose refinement
120/// fails or produces no output are kept at their original location.
121fn refine_points<E: Equation + Send + 'static>(
122    equation: &E,
123    data: &Data,
124    error_models: &AssayErrorModels,
125    theta: &Theta,
126    weights: &Weights,
127    cycles: usize,
128    progress: bool,
129) -> Result<(Theta, Weights)> {
130    let parameter_space = theta.parameters().clone();
131    let n_points = theta.matrix().nrows();
132    let mut refined_points: Vec<Vec<f64>> = Vec::with_capacity(n_points);
133    let mut kept_weights: Vec<f64> = Vec::with_capacity(n_points);
134
135    for i in 0..n_points {
136        let point: Vec<f64> = theta.matrix().row(i).iter().copied().collect();
137        let single = Mat::from_fn(1, point.len(), |_r, c| point[c]);
138        let single_theta = Theta::from_parts(single, parameter_space.clone())?;
139
140        let npag_config = NpagConfig {
141            max_cycles: cycles,
142            progress,
143            ..Default::default()
144        };
145        let mut npag = NPAG::from_parts(
146            equation.clone(),
147            data.clone(),
148            error_models.clone(),
149            single_theta,
150            npag_config,
151        )?;
152
153        #[allow(clippy::while_let_loop)]
154        let run = npag.initialize().and_then(|_| {
155            loop {
156                match npag.next_cycle()? {
157                    Status::Continue => continue,
158                    Status::Stop(_) => break,
159                }
160            }
161            Ok(())
162        });
163
164        match run {
165            Ok(()) if npag.theta().matrix().nrows() > 0 => {
166                let refined: Vec<f64> = npag.theta().matrix().row(0).iter().copied().collect();
167                refined_points.push(refined);
168            }
169            Ok(()) => {
170                tracing::warn!(
171                    "NCNPAG: refinement produced no points for support point {} — keeping original",
172                    i + 1
173                );
174                refined_points.push(point);
175            }
176            Err(e) => {
177                tracing::warn!(
178                    "NCNPAG: refinement failed for support point {}: {} — keeping original",
179                    i + 1,
180                    e
181                );
182                refined_points.push(point);
183            }
184        }
185        kept_weights.push(weights[i]);
186    }
187
188    let n_params = parameter_space.len();
189    let matrix = Mat::from_fn(refined_points.len(), n_params, |r, c| refined_points[r][c]);
190    let refined_theta = Theta::from_parts(matrix, parameter_space)?;
191
192    let weight_sum: f64 = kept_weights.iter().sum();
193    let refined_weights = if weight_sum > 0.0 {
194        Weights::from_vec(kept_weights.iter().map(|w| w / weight_sum).collect())
195    } else {
196        Weights::uniform(refined_points.len())
197    };
198
199    Ok((refined_theta, refined_weights))
200}
201
202/// Marginal log-likelihood of the data under a discrete `(psi, weights)` model.
203fn marginal_loglik(psi: &Psi, w: &Weights) -> f64 {
204    let m = psi.matrix();
205    (0..m.nrows())
206        .map(|s| {
207            let acc: f64 = (0..m.ncols()).map(|j| *m.get(s, j) * w[j]).sum();
208            acc.max(f64::MIN_POSITIVE).ln()
209        })
210        .sum::<f64>()
211        + psi.log_scale()
212}
213
214impl<E: Equation + Send + 'static> NonParametricRunner<E> for NCNPAG<E> {
215    fn into_result(&self) -> Result<NonParametricResult<E>> {
216        NonParametricResult::new(
217            self.equation.clone(),
218            self.data.clone(),
219            self.error_models.clone(),
220            self.prior.clone(),
221            self.theta.clone(),
222            self.psi.clone(),
223            self.w.clone(),
224            self.objf,
225            self.cycle,
226            self.status.clone(),
227            self.cyclelog.clone(),
228        )
229    }
230
231    fn error_models(&self) -> &AssayErrorModels {
232        &self.error_models
233    }
234
235    fn equation(&self) -> &E {
236        &self.equation
237    }
238
239    fn data(&self) -> &Data {
240        &self.data
241    }
242
243    fn likelihood(&self) -> f64 {
244        self.objf
245    }
246
247    fn increment_cycle(&mut self) -> usize {
248        0
249    }
250
251    fn cycle(&self) -> usize {
252        0
253    }
254
255    fn set_theta(&mut self, theta: Theta) {
256        self.theta = theta;
257    }
258
259    fn theta(&self) -> &Theta {
260        &self.theta
261    }
262
263    fn psi(&self) -> &Psi {
264        &self.psi
265    }
266
267    fn set_status(&mut self, status: Status) {
268        self.status = status;
269    }
270
271    fn status(&self) -> &Status {
272        &self.status
273    }
274
275    fn evaluation(&mut self) -> Result<Status> {
276        self.status = Status::Stop(StopReason::Converged);
277        Ok(self.status.clone())
278    }
279
280    fn estimation(&mut self) -> Result<()> {
281        // Likelihood of each fixed support point for the data.
282        let psi = calculate_psi(
283            &self.equation,
284            &self.data,
285            &self.theta,
286            &self.error_models,
287            false,
288        )?;
289
290        // Flat (uniform) prior: postⱼ ∝ ∏ᵢ P(dataᵢ | θⱼ). Accumulate in log space.
291        let n_points = self.theta.matrix().nrows();
292        let mut log_weights = vec![f64::NEG_INFINITY; n_points];
293        for (j, slot) in log_weights.iter_mut().enumerate() {
294            let mut log_weight = 0.0; // ln(uniform prior) is constant, drops out on normalization
295            let mut is_zero = false;
296            for s in 0..psi.matrix().nrows() {
297                let likelihood = psi.matrix()[(s, j)];
298                if likelihood <= 0.0 {
299                    is_zero = true;
300                    break;
301                }
302                log_weight += likelihood.ln();
303            }
304            if !is_zero {
305                *slot = log_weight;
306            }
307        }
308
309        let max_log_weight = log_weights
310            .iter()
311            .copied()
312            .fold(f64::NEG_INFINITY, f64::max);
313        if !max_log_weight.is_finite() {
314            anyhow::bail!("NCNPAG: every support point has zero joint likelihood for the data");
315        }
316
317        let mut weights: Vec<f64> = log_weights
318            .iter()
319            .map(|&lw| {
320                if lw.is_finite() {
321                    (lw - max_log_weight).exp()
322                } else {
323                    0.0
324                }
325            })
326            .collect();
327        let total: f64 = weights.iter().sum();
328        if total <= 0.0 {
329            anyhow::bail!("NCNPAG: filtering produced non-positive posterior mass");
330        }
331        for w in &mut weights {
332            *w /= total;
333        }
334
335        // Non-collapsing filter: keep points within 1e-100 of the maximum weight.
336        let max_weight = weights.iter().copied().fold(f64::NEG_INFINITY, f64::max);
337        let threshold = 1e-100;
338        let keep: Vec<usize> = weights
339            .iter()
340            .enumerate()
341            .filter(|(_, w)| **w > threshold * max_weight)
342            .map(|(i, _)| i)
343            .collect();
344
345        // Filter theta and renormalize the surviving weights (NPAGFULL11).
346        self.theta.filter_indices(&keep);
347        let kept: Vec<f64> = keep.iter().map(|&i| weights[i]).collect();
348        let sum: f64 = kept.iter().sum();
349        self.w = Weights::from_vec(kept.iter().map(|w| w / sum).collect());
350
351        // NPAGFULL: refine each surviving point with a full NPAG seeded from it.
352        if self.cycles > 0 {
353            let (refined_theta, refined_weights) = refine_points(
354                &self.equation,
355                &self.data,
356                &self.error_models,
357                &self.theta,
358                &self.w,
359                self.cycles,
360                self.progress,
361            )?;
362            self.theta = refined_theta;
363            self.w = refined_weights;
364        }
365
366        // Recompute psi over the final support points so psi/weights stay aligned.
367        self.psi = calculate_psi(
368            &self.equation,
369            &self.data,
370            &self.theta,
371            &self.error_models,
372            false,
373        )?;
374
375        self.objf = marginal_loglik(&self.psi, &self.w);
376        Ok(())
377    }
378
379    fn condensation(&mut self) -> Result<()> {
380        Ok(())
381    }
382
383    fn optimizations(&mut self) -> Result<()> {
384        Ok(())
385    }
386
387    fn expansion(&mut self) -> Result<()> {
388        Ok(())
389    }
390
391    fn log_cycle_state(&mut self) {
392        let state = NPCycle::new(
393            self.cycle,
394            self.objf,
395            self.error_models.clone(),
396            self.theta.clone(),
397            self.w.clone(),
398            self.theta.nspp(),
399            0.0,
400            self.status.clone(),
401        );
402        self.cyclelog.push(state);
403    }
404
405    /// NCNPAG is a single-pass reweighting: it evaluates the likelihood of the
406    /// fixed prior support points once, rather than iterating cycles.
407    fn fit(&mut self) -> Result<NonParametricResult<E>> {
408        self.estimation()?;
409        self.evaluation()?;
410        self.log_cycle_state();
411
412        self.into_result()
413    }
414}
415
416#[cfg(test)]
417mod tests {
418    use super::*;
419    use approx::assert_relative_eq;
420    use ndarray::Array2;
421
422    #[test]
423    fn marginal_loglik_restores_subject_scaling() -> anyhow::Result<()> {
424        let log_likelihoods = Array2::from_shape_vec((2, 2), vec![-1000.0, -1001.0, -2.0, -4.0])?;
425        let psi = Psi::from_log_likelihoods(log_likelihoods)?;
426        let weights = Weights::from_vec(vec![0.6, 0.4]);
427
428        let expected = -1000.0 + (0.6 + 0.4 * (-1.0_f64).exp()).ln() - 2.0
429            + (0.6 + 0.4 * (-2.0_f64).exp()).ln();
430
431        assert_relative_eq!(marginal_loglik(&psi, &weights), expected, epsilon = 1e-12);
432        Ok(())
433    }
434}
pmcore/algorithms/nonparametric/ncnpag.rs

pmcore/algorithms/nonparametric/
ncnpag.rs
