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()
211}
212
213impl<E: Equation + Send + 'static> NonParametricRunner<E> for NCNPAG<E> {
214    fn into_result(&self) -> Result<NonParametricResult<E>> {
215        NonParametricResult::new(
216            self.equation.clone(),
217            self.data.clone(),
218            self.error_models.clone(),
219            self.prior.clone(),
220            self.theta.clone(),
221            self.psi.clone(),
222            self.w.clone(),
223            self.objf,
224            self.cycle,
225            self.status.clone(),
226            self.cyclelog.clone(),
227        )
228    }
229
230    fn error_models(&self) -> &AssayErrorModels {
231        &self.error_models
232    }
233
234    fn equation(&self) -> &E {
235        &self.equation
236    }
237
238    fn data(&self) -> &Data {
239        &self.data
240    }
241
242    fn likelihood(&self) -> f64 {
243        self.objf
244    }
245
246    fn increment_cycle(&mut self) -> usize {
247        0
248    }
249
250    fn cycle(&self) -> usize {
251        0
252    }
253
254    fn set_theta(&mut self, theta: Theta) {
255        self.theta = theta;
256    }
257
258    fn theta(&self) -> &Theta {
259        &self.theta
260    }
261
262    fn psi(&self) -> &Psi {
263        &self.psi
264    }
265
266    fn set_status(&mut self, status: Status) {
267        self.status = status;
268    }
269
270    fn status(&self) -> &Status {
271        &self.status
272    }
273
274    fn evaluation(&mut self) -> Result<Status> {
275        self.status = Status::Stop(StopReason::Converged);
276        Ok(self.status.clone())
277    }
278
279    fn estimation(&mut self) -> Result<()> {
280        // Likelihood of each fixed support point for the data.
281        let psi = calculate_psi(
282            &self.equation,
283            &self.data,
284            &self.theta,
285            &self.error_models,
286            false,
287        )?;
288
289        // Flat (uniform) prior: postⱼ ∝ ∏ᵢ P(dataᵢ | θⱼ). Accumulate in log space.
290        let n_points = self.theta.matrix().nrows();
291        let mut log_weights = vec![f64::NEG_INFINITY; n_points];
292        for (j, slot) in log_weights.iter_mut().enumerate() {
293            let mut log_weight = 0.0; // ln(uniform prior) is constant, drops out on normalization
294            let mut is_zero = false;
295            for s in 0..psi.matrix().nrows() {
296                let likelihood = psi.matrix()[(s, j)];
297                if likelihood <= 0.0 {
298                    is_zero = true;
299                    break;
300                }
301                log_weight += likelihood.ln();
302            }
303            if !is_zero {
304                *slot = log_weight;
305            }
306        }
307
308        let max_log_weight = log_weights
309            .iter()
310            .copied()
311            .fold(f64::NEG_INFINITY, f64::max);
312        if !max_log_weight.is_finite() {
313            anyhow::bail!("NCNPAG: every support point has zero joint likelihood for the data");
314        }
315
316        let mut weights: Vec<f64> = log_weights
317            .iter()
318            .map(|&lw| {
319                if lw.is_finite() {
320                    (lw - max_log_weight).exp()
321                } else {
322                    0.0
323                }
324            })
325            .collect();
326        let total: f64 = weights.iter().sum();
327        if total <= 0.0 {
328            anyhow::bail!("NCNPAG: filtering produced non-positive posterior mass");
329        }
330        for w in &mut weights {
331            *w /= total;
332        }
333
334        // Non-collapsing filter: keep points within 1e-100 of the maximum weight.
335        let max_weight = weights.iter().copied().fold(f64::NEG_INFINITY, f64::max);
336        let threshold = 1e-100;
337        let keep: Vec<usize> = weights
338            .iter()
339            .enumerate()
340            .filter(|(_, w)| **w > threshold * max_weight)
341            .map(|(i, _)| i)
342            .collect();
343
344        // Filter theta and renormalize the surviving weights (NPAGFULL11).
345        self.theta.filter_indices(&keep);
346        let kept: Vec<f64> = keep.iter().map(|&i| weights[i]).collect();
347        let sum: f64 = kept.iter().sum();
348        self.w = Weights::from_vec(kept.iter().map(|w| w / sum).collect());
349
350        // NPAGFULL: refine each surviving point with a full NPAG seeded from it.
351        if self.cycles > 0 {
352            let (refined_theta, refined_weights) = refine_points(
353                &self.equation,
354                &self.data,
355                &self.error_models,
356                &self.theta,
357                &self.w,
358                self.cycles,
359                self.progress,
360            )?;
361            self.theta = refined_theta;
362            self.w = refined_weights;
363        }
364
365        // Recompute psi over the final support points so psi/weights stay aligned.
366        self.psi = calculate_psi(
367            &self.equation,
368            &self.data,
369            &self.theta,
370            &self.error_models,
371            false,
372        )?;
373
374        self.objf = marginal_loglik(&self.psi, &self.w);
375        Ok(())
376    }
377
378    fn condensation(&mut self) -> Result<()> {
379        Ok(())
380    }
381
382    fn optimizations(&mut self) -> Result<()> {
383        Ok(())
384    }
385
386    fn expansion(&mut self) -> Result<()> {
387        Ok(())
388    }
389
390    fn log_cycle_state(&mut self) {
391        let state = NPCycle::new(
392            self.cycle,
393            self.objf,
394            self.error_models.clone(),
395            self.theta.clone(),
396            self.w.clone(),
397            self.theta.nspp(),
398            0.0,
399            self.status.clone(),
400        );
401        self.cyclelog.push(state);
402    }
403
404    /// NCNPAG is a single-pass reweighting: it evaluates the likelihood of the
405    /// fixed prior support points once, rather than iterating cycles.
406    fn fit(&mut self) -> Result<NonParametricResult<E>> {
407        self.estimation()?;
408        self.evaluation()?;
409        self.log_cycle_state();
410
411        self.into_result()
412    }
413}
pmcore/algorithms/nonparametric/ncnpag.rs

pmcore/algorithms/nonparametric/
ncnpag.rs
