From 6ad7b9bc8192535535122d2cd2f4d3564e404580 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juli=C3=A1n=20D=2E=20Ot=C3=A1lvaro?= Date: Mon, 20 Jul 2026 19:55:55 +0100 Subject: [PATCH 1/2] Fix likelihood underflow --- src/algorithms/nonparametric/ncnpag.rs | 23 +++++++- src/estimation/nonparametric/ipm.rs | 17 +++++- src/estimation/nonparametric/psi.rs | 76 +++++++++++++++++++++++--- 3 files changed, 105 insertions(+), 11 deletions(-) diff --git a/src/algorithms/nonparametric/ncnpag.rs b/src/algorithms/nonparametric/ncnpag.rs index 4c02f5645..860968ce7 100644 --- a/src/algorithms/nonparametric/ncnpag.rs +++ b/src/algorithms/nonparametric/ncnpag.rs @@ -207,7 +207,8 @@ fn marginal_loglik(psi: &Psi, w: &Weights) -> f64 { let acc: f64 = (0..m.ncols()).map(|j| *m.get(s, j) * w[j]).sum(); acc.max(f64::MIN_POSITIVE).ln() }) - .sum() + .sum::() + + psi.log_scale() } impl NonParametricRunner for NCNPAG { @@ -411,3 +412,23 @@ impl NonParametricRunner for NCNPAG { self.into_result() } } + +#[cfg(test)] +mod tests { + use super::*; + use approx::assert_relative_eq; + use ndarray::Array2; + + #[test] + fn marginal_loglik_restores_subject_scaling() -> anyhow::Result<()> { + let log_likelihoods = Array2::from_shape_vec((2, 2), vec![-1000.0, -1001.0, -2.0, -4.0])?; + let psi = Psi::from_log_likelihoods(log_likelihoods)?; + let weights = Weights::from_vec(vec![0.6, 0.4]); + + let expected = -1000.0 + (0.6 + 0.4 * (-1.0_f64).exp()).ln() - 2.0 + + (0.6 + 0.4 * (-2.0_f64).exp()).ln(); + + assert_relative_eq!(marginal_loglik(&psi, &weights), expected, epsilon = 1e-12); + Ok(()) + } +} diff --git a/src/estimation/nonparametric/ipm.rs b/src/estimation/nonparametric/ipm.rs index 719fe1b55..e30de75a0 100644 --- a/src/estimation/nonparametric/ipm.rs +++ b/src/estimation/nonparametric/ipm.rs @@ -7,6 +7,7 @@ use rayon::prelude::*; /// Applies Burke's Interior Point Method (IPM) to solve a convex optimization problem. pub fn burke(psi: &Psi) -> anyhow::Result<(Weights, f64)> { + let log_scale = psi.log_scale(); let mut psi = psi.matrix().to_owned(); psi.row_iter_mut().try_for_each(|row| { @@ -183,7 +184,7 @@ pub fn burke(psi: &Psi) -> anyhow::Result<(Weights, f64)> { } lam /= n_sub as f64; - let obj = (psi * &lam).iter().map(|x| x.ln()).sum(); + let obj = (psi * &lam).iter().map(|x| x.ln()).sum::() + log_scale; let lam_sum: f64 = lam.iter().sum(); lam = &lam / lam_sum; @@ -225,6 +226,20 @@ mod tests { } } + #[test] + fn test_burke_restores_log_likelihood_scale() -> anyhow::Result<()> { + let log_likelihoods = + ndarray::Array2::from_shape_vec((2, 2), vec![-1000.0, -1001.0, -2.0, -4.0])?; + let psi = Psi::from_log_likelihoods(log_likelihoods)?; + + let (weights, objective) = burke(&psi)?; + let expected_objective = -1000.0 + (weights[0] + (-1.0_f64).exp() * weights[1]).ln() - 2.0 + + (weights[0] + (-2.0_f64).exp() * weights[1]).ln(); + + assert_relative_eq!(objective, expected_objective, epsilon = 1e-8); + Ok(()) + } + #[test] fn test_burke_with_non_finite_values() { let n_sub = 10; diff --git a/src/estimation/nonparametric/psi.rs b/src/estimation/nonparametric/psi.rs index 82b8a7d37..281ec475d 100644 --- a/src/estimation/nonparametric/psi.rs +++ b/src/estimation/nonparametric/psi.rs @@ -1,7 +1,7 @@ use anyhow::bail; use anyhow::Result; use faer::Mat; -use ndarray::Array2; +use ndarray::{Array2, Axis}; use pharmsol::prelude::simulator::log_likelihood_matrix; use pharmsol::AssayErrorModels; use pharmsol::Data; @@ -14,17 +14,47 @@ use super::theta::Theta; #[derive(Debug, Clone, PartialEq)] pub struct Psi { matrix: Mat, + /// Sum of per-subject log-likelihood offsets removed before exponentiation. + /// Row scaling leaves the optimal weights unchanged; this offset restores + /// the unscaled log-likelihood objective returned by the estimator. + log_scale: f64, } impl Psi { pub fn new() -> Self { - Psi { matrix: Mat::new() } + Psi { + matrix: Mat::new(), + log_scale: 0.0, + } } pub fn matrix(&self) -> &Mat { &self.matrix } + pub(crate) fn log_scale(&self) -> f64 { + self.log_scale + } + + pub(crate) fn from_log_likelihoods(mut log_likelihoods: Array2) -> Result { + let mut log_scale = 0.0; + + for mut row in log_likelihoods.axis_iter_mut(Axis(0)) { + let row_max = row.iter().copied().fold(f64::NEG_INFINITY, f64::max); + if !row_max.is_finite() { + bail!("Each subject must have at least one finite log-likelihood"); + } + + log_scale += row_max; + row.mapv_inplace(|value| (value - row_max).exp()); + } + + let matrix = Mat::from_fn(log_likelihoods.nrows(), log_likelihoods.ncols(), |i, j| { + log_likelihoods[(i, j)] + }); + Ok(Self { matrix, log_scale }) + } + pub fn nspp(&self) -> usize { self.matrix.nrows() } @@ -95,7 +125,10 @@ impl Psi { let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]); - Ok(Psi { matrix: mat }) + Ok(Psi { + matrix: mat, + log_scale: 0.0, + }) } } @@ -108,20 +141,29 @@ impl Default for Psi { impl From> for Psi { fn from(array: Array2) -> Self { let matrix = Mat::from_fn(array.nrows(), array.ncols(), |i, j| array[(i, j)]); - Psi { matrix } + Psi { + matrix, + log_scale: 0.0, + } } } impl From> for Psi { fn from(matrix: Mat) -> Self { - Psi { matrix } + Psi { + matrix, + log_scale: 0.0, + } } } impl From<&Array2> for Psi { fn from(array: &Array2) -> Self { let matrix = Mat::from_fn(array.nrows(), array.ncols(), |i, j| array[(i, j)]); - Psi { matrix } + Psi { + matrix, + log_scale: 0.0, + } } } @@ -192,7 +234,10 @@ impl<'de> Deserialize<'de> for Psi { let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]); - Ok(Psi { matrix: mat }) + Ok(Psi { + matrix: mat, + log_scale: 0.0, + }) } } @@ -211,9 +256,8 @@ pub(crate) fn calculate_psi( let theta_ndarray = Array2::from_shape_fn((tm.nrows(), tm.ncols()), |(i, j)| tm[(i, j)]); let log_psi = log_likelihood_matrix(equation, subjects, &theta_ndarray, error_models, progress)?; - let psi_ndarray = log_psi.mapv(f64::exp); - Ok(Psi::from(psi_ndarray)) + Psi::from_log_likelihoods(log_psi) } #[cfg(test)] @@ -221,6 +265,20 @@ mod tests { use super::*; use ndarray::Array2; + #[test] + fn log_likelihood_rows_are_scaled_before_exponentiation() -> Result<()> { + let log_likelihoods = Array2::from_shape_vec((2, 2), vec![-1000.0, -1001.0, -2.0, -4.0])?; + + let psi = Psi::from_log_likelihoods(log_likelihoods)?; + + assert_eq!(psi.log_scale(), -1002.0); + assert_eq!(psi.matrix()[(0, 0)], 1.0); + assert_eq!(psi.matrix()[(1, 0)], 1.0); + assert!((psi.matrix()[(0, 1)] - (-1.0_f64).exp()).abs() < 1e-12); + assert!((psi.matrix()[(1, 1)] - (-2.0_f64).exp()).abs() < 1e-12); + Ok(()) + } + #[test] fn test_from_array2() { let array = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); From 6338b787f344cc2c33114fcc99dba49e3eeed85e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juli=C3=A1n=20D=2E=20Ot=C3=A1lvaro?= Date: Mon, 20 Jul 2026 20:21:31 +0100 Subject: [PATCH 2/2] Fix NPOD likelihood scale --- src/algorithms/nonparametric/npod.rs | 9 +++++- src/estimation/nonparametric/psi.rs | 43 +++++++++++++++++----------- 2 files changed, 35 insertions(+), 17 deletions(-) diff --git a/src/algorithms/nonparametric/npod.rs b/src/algorithms/nonparametric/npod.rs index 8348f8296..32aaa3372 100644 --- a/src/algorithms/nonparametric/npod.rs +++ b/src/algorithms/nonparametric/npod.rs @@ -336,7 +336,14 @@ impl NonParametricRunner for NPOD { fn expansion(&mut self) -> Result<()> { let pyl_col = self.psi().matrix().as_ref() * self.w.weights().as_ref(); - let pyl: Array1 = pyl_col.iter().copied().collect(); + let mut pyl: Array1 = pyl_col.iter().copied().collect(); + + // ParameterOptimizer currently evaluates raw likelihoods, so restore + // pyl to the same scale before calculating directional derivatives. + // TODO: Move NPOD and its parameter optimizer fully into log space. + for (pyl_i, row_log_scale) in pyl.iter_mut().zip(self.psi().row_log_scales()) { + *pyl_i *= row_log_scale.exp(); + } let error_model: AssayErrorModels = self.error_models.clone(); diff --git a/src/estimation/nonparametric/psi.rs b/src/estimation/nonparametric/psi.rs index 281ec475d..c6a744854 100644 --- a/src/estimation/nonparametric/psi.rs +++ b/src/estimation/nonparametric/psi.rs @@ -14,17 +14,17 @@ use super::theta::Theta; #[derive(Debug, Clone, PartialEq)] pub struct Psi { matrix: Mat, - /// Sum of per-subject log-likelihood offsets removed before exponentiation. - /// Row scaling leaves the optimal weights unchanged; this offset restores - /// the unscaled log-likelihood objective returned by the estimator. - log_scale: f64, + /// Per-subject log-likelihood offsets removed before exponentiation. + /// Row scaling leaves the optimal weights unchanged; these offsets restore + /// the original likelihood scale where required by an algorithm. + row_log_scales: Vec, } impl Psi { pub fn new() -> Self { Psi { matrix: Mat::new(), - log_scale: 0.0, + row_log_scales: Vec::new(), } } @@ -33,11 +33,15 @@ impl Psi { } pub(crate) fn log_scale(&self) -> f64 { - self.log_scale + self.row_log_scales.iter().sum() + } + + pub(crate) fn row_log_scales(&self) -> &[f64] { + &self.row_log_scales } pub(crate) fn from_log_likelihoods(mut log_likelihoods: Array2) -> Result { - let mut log_scale = 0.0; + let mut row_log_scales = Vec::with_capacity(log_likelihoods.nrows()); for mut row in log_likelihoods.axis_iter_mut(Axis(0)) { let row_max = row.iter().copied().fold(f64::NEG_INFINITY, f64::max); @@ -45,14 +49,17 @@ impl Psi { bail!("Each subject must have at least one finite log-likelihood"); } - log_scale += row_max; + row_log_scales.push(row_max); row.mapv_inplace(|value| (value - row_max).exp()); } let matrix = Mat::from_fn(log_likelihoods.nrows(), log_likelihoods.ncols(), |i, j| { log_likelihoods[(i, j)] }); - Ok(Self { matrix, log_scale }) + Ok(Self { + matrix, + row_log_scales, + }) } pub fn nspp(&self) -> usize { @@ -127,7 +134,7 @@ impl Psi { Ok(Psi { matrix: mat, - log_scale: 0.0, + row_log_scales: vec![0.0; nrows], }) } } @@ -140,29 +147,32 @@ impl Default for Psi { impl From> for Psi { fn from(array: Array2) -> Self { - let matrix = Mat::from_fn(array.nrows(), array.ncols(), |i, j| array[(i, j)]); + let nrows = array.nrows(); + let matrix = Mat::from_fn(nrows, array.ncols(), |i, j| array[(i, j)]); Psi { matrix, - log_scale: 0.0, + row_log_scales: vec![0.0; nrows], } } } impl From> for Psi { fn from(matrix: Mat) -> Self { + let nrows = matrix.nrows(); Psi { matrix, - log_scale: 0.0, + row_log_scales: vec![0.0; nrows], } } } impl From<&Array2> for Psi { fn from(array: &Array2) -> Self { - let matrix = Mat::from_fn(array.nrows(), array.ncols(), |i, j| array[(i, j)]); + let nrows = array.nrows(); + let matrix = Mat::from_fn(nrows, array.ncols(), |i, j| array[(i, j)]); Psi { matrix, - log_scale: 0.0, + row_log_scales: vec![0.0; nrows], } } } @@ -236,7 +246,7 @@ impl<'de> Deserialize<'de> for Psi { Ok(Psi { matrix: mat, - log_scale: 0.0, + row_log_scales: vec![0.0; nrows], }) } } @@ -272,6 +282,7 @@ mod tests { let psi = Psi::from_log_likelihoods(log_likelihoods)?; assert_eq!(psi.log_scale(), -1002.0); + assert_eq!(psi.row_log_scales(), &[-1000.0, -2.0]); assert_eq!(psi.matrix()[(0, 0)], 1.0); assert_eq!(psi.matrix()[(1, 0)], 1.0); assert!((psi.matrix()[(0, 1)] - (-1.0_f64).exp()).abs() < 1e-12);