pmcore/algorithms/
mod.rs

1use std::fs;
2use std::path::Path;
3
4use crate::estimation::nonparametric::{NonParametricResult, Psi, Theta};
5use crate::estimation::{EstimationProblem, Framework};
6use crate::results::FitResult;
7
8use anyhow::Context;
9use anyhow::Result;
10use ndarray::parallel::prelude::{IntoParallelIterator, ParallelIterator};
11
12use pharmsol::prelude::{data::Data, simulator::Equation};
13
14use pharmsol::{Predictions, Subject};
15use serde::{Deserialize, Serialize};
16
17/// Defines an algorithm that can fit an [`EstimationProblem`] to produce a result.
18///
19/// Implementors are the lightweight, user-facing configuration structs (e.g.
20/// `NpagConfig`). The heavy, mutable execution state used while fitting is an
21/// internal implementation detail.
22pub trait Algorithm<E: Equation, F: crate::estimation::Framework> {
23    /// The specific result struct (e.g. `NonParametricResult<E>`).
24    type Output: FitResult;
25
26    /// Consumes the configuration and the problem, runs the optimization to
27    /// completion, and returns the strictly-typed result.
28    fn fit(self, problem: EstimationProblem<E, F>) -> Result<Self::Output>;
29}
30
31// Module organization for algorithm types
32pub mod nonparametric;
33pub mod parametric;
34
35impl<E: Equation, F: Framework> EstimationProblem<E, F> {
36    /// Consumes the problem and an algorithm configuration, runs the fit to
37    /// completion, and returns the result.
38    pub fn fit_with<A>(self, algorithm: A) -> Result<A::Output>
39    where
40        A: Algorithm<E, F>,
41    {
42        algorithm.fit(self)
43    }
44}
45
46pub trait NonParametricRunner<E: Equation + Send + 'static>: Sync + Send + 'static {
47    /// Identify subjects whose total probability given the model is zero or
48    /// non-finite.
49    ///
50    /// Each row of [`Psi`] holds the likelihood of a subject across every
51    /// support point, so a subject's probability is the sum across its row. A
52    /// subject is flagged when that sum is zero or not finite, meaning the model
53    /// cannot explain the subject's data. When any subject is flagged, detailed
54    /// per-subject diagnostics are logged and an error is returned.
55    fn check_zero_probability_subjects(&self) -> Result<()> {
56        let psi = self.psi().matrix();
57
58        // Report non-finite entries; these propagate into the row sums below.
59        let nonfinite = psi
60            .row_iter()
61            .flat_map(|row| row.iter().copied())
62            .filter(|v| !v.is_finite())
63            .count();
64        if nonfinite > 0 {
65            tracing::warn!(
66                "Psi matrix contains {} non-finite value(s) of {} total",
67                nonfinite,
68                psi.nrows() * psi.ncols()
69            );
70        }
71
72        // A subject's probability is the sum across its row.
73        let subjects = self.data().subjects();
74        let flagged: Vec<usize> = (0..psi.nrows())
75            .filter(|&i| {
76                let probability: f64 = (0..psi.ncols()).map(|j| psi[(i, j)]).sum();
77                !probability.is_finite() || probability == 0.0
78            })
79            .collect();
80
81        if flagged.is_empty() {
82            return Ok(());
83        }
84
85        tracing::error!(
86            "{}/{} subjects have zero probability given the model",
87            flagged.len(),
88            psi.nrows()
89        );
90
91        for &i in &flagged {
92            self.log_zero_probability_subject(subjects[i]);
93        }
94
95        let ids: Vec<&String> = flagged.iter().map(|&i| subjects[i].id()).collect();
96        Err(anyhow::anyhow!(
97            "The probability of {}/{} subjects is zero given the model. Affected subjects: {:?}",
98            flagged.len(),
99            psi.nrows(),
100            ids
101        ))
102    }
103
104    /// Log detailed likelihood diagnostics for a single subject whose
105    /// probability given the model is zero or non-finite.
106    fn log_zero_probability_subject(&self, subject: &Subject) {
107        tracing::debug!("Subject with zero probability: {}", subject.id());
108
109        let error_model = self.error_models().clone();
110
111        // Simulate every support point for this subject in parallel.
112        let mut results: Vec<_> = self
113            .theta()
114            .matrix()
115            .row_iter()
116            .enumerate()
117            .collect::<Vec<_>>()
118            .into_par_iter()
119            .map(|(i, spp)| {
120                let support_point: Vec<f64> = spp.iter().copied().collect();
121                let (pred, ll) = self
122                    .equation()
123                    .simulate_subject_dense(subject, &support_point, Some(&error_model))
124                    .unwrap(); //TODO: Handle error
125                (i, support_point, pred.get_predictions(), ll)
126            })
127            .collect();
128
129        // Summarise the distribution of likelihood values.
130        let mut nan = 0;
131        let mut pos_inf = 0;
132        let mut neg_inf = 0;
133        let mut zero = 0;
134        let mut valid = 0;
135        for (_, _, _, ll) in &results {
136            match ll {
137                Some(v) if v.is_nan() => nan += 1,
138                Some(v) if v.is_infinite() && v.is_sign_positive() => pos_inf += 1,
139                Some(v) if v.is_infinite() => neg_inf += 1,
140                Some(v) if *v == 0.0 => zero += 1,
141                Some(_) => valid += 1,
142                None => nan += 1,
143            }
144        }
145
146        let total = results.len();
147        let pct = |n: usize| 100.0 * n as f64 / total as f64;
148        tracing::debug!(
149            "\tLikelihood analysis for subject {} ({} support points):",
150            subject.id(),
151            total
152        );
153        tracing::debug!("\tNaN likelihoods: {} ({:.1}%)", nan, pct(nan));
154        tracing::debug!("\t+Inf likelihoods: {} ({:.1}%)", pos_inf, pct(pos_inf));
155        tracing::debug!("\t-Inf likelihoods: {} ({:.1}%)", neg_inf, pct(neg_inf));
156        tracing::debug!("\tZero likelihoods: {} ({:.1}%)", zero, pct(zero));
157        tracing::debug!("\tValid likelihoods: {} ({:.1}%)", valid, pct(valid));
158
159        // Show the most likely support points to aid debugging.
160        results.sort_by(|a, b| {
161            b.3.unwrap_or(f64::NEG_INFINITY)
162                .partial_cmp(&a.3.unwrap_or(f64::NEG_INFINITY))
163                .unwrap_or(std::cmp::Ordering::Equal)
164        });
165
166        const TAKE: usize = 3;
167        tracing::debug!("Top {} most likely support points:", TAKE);
168        for (i, support_point, preds, ll) in results.iter().take(TAKE) {
169            tracing::debug!("\tSupport point #{}: {:?}", i, support_point);
170            tracing::debug!("\t\tLog-likelihood: {:?}", ll);
171            tracing::debug!(
172                "\t\tTimes: {:?}",
173                preds.iter().map(|x| x.time()).collect::<Vec<f64>>()
174            );
175            tracing::debug!(
176                "\t\tObservations: {:?}",
177                preds
178                    .iter()
179                    .map(|x| x.observation())
180                    .collect::<Vec<Option<f64>>>()
181            );
182            tracing::debug!(
183                "\t\tPredictions: {:?}",
184                preds.iter().map(|x| x.prediction()).collect::<Vec<f64>>()
185            );
186            tracing::debug!(
187                "\t\tOuteqs: {:?}",
188                preds.iter().map(|x| x.outeq()).collect::<Vec<usize>>()
189            );
190            tracing::debug!(
191                "\t\tStates: {:?}",
192                preds
193                    .iter()
194                    .map(|x| x.state().to_vec())
195                    .collect::<Vec<Vec<f64>>>()
196            );
197        }
198        tracing::debug!("=====================");
199    }
200
201    fn error_models(&self) -> &pharmsol::prelude::data::AssayErrorModels;
202    /// Get the equation used in the algorithm
203    fn equation(&self) -> &E;
204    /// Get the data used in the algorithm
205    fn data(&self) -> &Data;
206
207    /// Increment the cycle counter and return the new value
208    fn increment_cycle(&mut self) -> usize;
209    /// Get the current cycle number
210    fn cycle(&self) -> usize;
211    /// Set the current [Theta]
212    fn set_theta(&mut self, theta: Theta);
213    /// Get the current [Theta]
214    fn theta(&self) -> &Theta;
215    /// Get the current [Psi]
216    fn psi(&self) -> &Psi;
217    /// Get the current likelihood
218    fn likelihood(&self) -> f64;
219    /// Get the current negative two log-likelihood
220    fn n2ll(&self) -> f64 {
221        -2.0 * self.likelihood()
222    }
223    /// Get the current [Status] of the algorithm
224    fn status(&self) -> &Status;
225    /// Set the current [Status] of the algorithm
226    fn set_status(&mut self, status: Status);
227    /// Evaluate convergence criteria and update status
228    fn evaluation(&mut self) -> Result<Status>;
229
230    /// Create and log a cycle state with the current algorithm state
231    fn log_cycle_state(&mut self);
232
233    /// Initialize the algorithm, setting up initial [Theta] and [Status]
234    fn initialize(&mut self) -> Result<()> {
235        // If a stop file exists in the current directory, remove it
236        if Path::new("stop").exists() {
237            tracing::info!("Removing existing stop file prior to run");
238            fs::remove_file("stop").context("Unable to remove previous stop file")?;
239        }
240        self.set_status(Status::Continue);
241
242        Ok(())
243    }
244    fn estimation(&mut self) -> Result<()>;
245    /// Performs condensation of [Theta] and updates [Psi]
246    ///
247    /// This step reduces the number of support points in [Theta] based on the current weights,
248    /// and updates the [Psi] matrix accordingly to reflect the new set of support points.
249    /// It is typically performed after the estimation step in each cycle of the algorithm.
250    fn condensation(&mut self) -> Result<()>;
251
252    /// Performs optimizations on the current `AssayErrorModels` and updates [Psi] accordingly
253    ///
254    /// This step refines the error model parameters to better fit the data,
255    /// and subsequently updates the [Psi] matrix to reflect these changes.
256    fn optimizations(&mut self) -> Result<()>;
257
258    /// Performs expansion of [Theta]
259    ///
260    /// This step increases the number of support points in [Theta] based on the current distribution,
261    /// allowing for exploration of the parameter space.
262    fn expansion(&mut self) -> Result<()>;
263
264    /// Proceed to the next cycle of the algorithm
265    ///
266    /// This method increments the cycle counter, performs expansion if necessary,
267    /// and then runs the estimation, condensation, optimization, logging, and evaluation steps
268    /// in sequence. It returns the current [Status] of the algorithm after completing these steps.
269    fn next_cycle(&mut self) -> Result<Status> {
270        let cycle = self.increment_cycle();
271
272        if cycle > 1 {
273            self.expansion()?;
274        }
275
276        let span = tracing::info_span!("", "{}", format!("Cycle {}", self.cycle()));
277        let _enter = span.enter();
278        self.estimation()?;
279        self.condensation()?;
280        self.optimizations()?;
281        self.evaluation()
282    }
283
284    /// Fit the model until convergence or stopping criteria are met
285    ///
286    /// This method runs the full fitting process, starting with initialization,
287    /// followed by iterative cycles of estimation, condensation, optimization, and evaluation
288    /// until the algorithm converges or meets a stopping criteria.
289    fn fit(&mut self) -> Result<NonParametricResult<E>> {
290        self.initialize()?;
291        while let Status::Continue = self.next_cycle()? {}
292        self.into_result()
293    }
294
295    #[allow(clippy::wrong_self_convention)]
296    fn into_result(&self) -> Result<NonParametricResult<E>>;
297}
298
299/// Where a fit stands: still running, or stopped (and why).
300#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
301pub enum Status {
302    Continue,
303    Stop(StopReason),
304}
305
306impl Status {
307    /// Whether the fit is still running.
308    pub fn is_continue(&self) -> bool {
309        matches!(self, Status::Continue)
310    }
311
312    /// Whether the fit has stopped.
313    pub fn is_stop(&self) -> bool {
314        matches!(self, Status::Stop(_))
315    }
316
317    /// Why the fit stopped, or `None` if it's still running.
318    pub fn stop_reason(&self) -> Option<&StopReason> {
319        match self {
320            Status::Stop(reason) => Some(reason),
321            Status::Continue => None,
322        }
323    }
324
325    /// Whether the fit stopped because it converged, rather than being cut short.
326    pub fn converged(&self) -> bool {
327        matches!(self, Status::Stop(StopReason::Converged))
328    }
329}
330
331impl std::fmt::Display for Status {
332    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
333        match self {
334            Status::Continue => write!(f, "Continue"),
335            Status::Stop(reason) => write!(f, "Stopped ({reason})"),
336        }
337    }
338}
339
340/// Why a fit stopped.
341#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
342pub enum StopReason {
343    /// The convergence criteria were met.
344    Converged,
345    /// Hit the cycle limit before converging.
346    MaxCycles,
347    /// A `stop` file was found on disk.
348    StopFile,
349    /// Stopped from code — [`request_stop`](crate::algorithms::nonparametric::FitController::request_stop)
350    /// or an observer returning [`CycleFlow::Stop`](crate::algorithms::nonparametric::CycleFlow::Stop).
351    Aborted,
352}
353
354impl std::fmt::Display for StopReason {
355    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
356        let reason = match self {
357            StopReason::Converged => "converged",
358            StopReason::MaxCycles => "maximum cycles reached",
359            StopReason::StopFile => "stop file detected",
360            StopReason::Aborted => "aborted",
361        };
362        f.write_str(reason)
363    }
364}
pmcore/algorithms/mod.rs

pmcore/algorithms/
mod.rs
