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/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/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..c6a744854 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,54 @@ use super::theta::Theta; #[derive(Debug, Clone, PartialEq)] pub struct Psi { matrix: Mat, + /// 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() } + Psi { + matrix: Mat::new(), + row_log_scales: Vec::new(), + } } pub fn matrix(&self) -> &Mat { &self.matrix } + pub(crate) fn log_scale(&self) -> f64 { + 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 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); + if !row_max.is_finite() { + bail!("Each subject must have at least one finite log-likelihood"); + } + + 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, + row_log_scales, + }) + } + pub fn nspp(&self) -> usize { self.matrix.nrows() } @@ -95,7 +132,10 @@ impl Psi { let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]); - Ok(Psi { matrix: mat }) + Ok(Psi { + matrix: mat, + row_log_scales: vec![0.0; nrows], + }) } } @@ -107,21 +147,33 @@ 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 } + let nrows = array.nrows(); + let matrix = Mat::from_fn(nrows, array.ncols(), |i, j| array[(i, j)]); + Psi { + matrix, + row_log_scales: vec![0.0; nrows], + } } } impl From> for Psi { fn from(matrix: Mat) -> Self { - Psi { matrix } + let nrows = matrix.nrows(); + Psi { + matrix, + 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)]); - Psi { matrix } + let nrows = array.nrows(); + let matrix = Mat::from_fn(nrows, array.ncols(), |i, j| array[(i, j)]); + Psi { + matrix, + row_log_scales: vec![0.0; nrows], + } } } @@ -192,7 +244,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, + row_log_scales: vec![0.0; nrows], + }) } } @@ -211,9 +266,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 +275,21 @@ 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.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); + 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();