fix: Normalize psi to prevent likelihood underflow

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::<f64>()
+        + psi.log_scale()
 }
 
 impl<E: Equation + Send + 'static> NonParametricRunner<E> for NCNPAG<E> {
@@ -411,3 +412,23 @@ impl<E: Equation + Send + 'static> NonParametricRunner<E> for NCNPAG<E> {
         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(())
+    }
+}

src/algorithms/nonparametric/npod.rs

@@ -336,7 +336,14 @@ impl<E: Equation + Send + 'static> NonParametricRunner<E> for NPOD<E> {
 
     fn expansion(&mut self) -> Result<()> {
         let pyl_col = self.psi().matrix().as_ref() * self.w.weights().as_ref();
-        let pyl: Array1<f64> = pyl_col.iter().copied().collect();
+        let mut pyl: Array1<f64> = 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();
 

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::<f64>() + 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;

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<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<f64>,
 }
 
 impl Psi {
     pub fn new() -> Self {
-        Psi { matrix: Mat::new() }
+        Psi {
+            matrix: Mat::new(),
+            row_log_scales: Vec::new(),
+        }
     }
 
     pub fn matrix(&self) -> &Mat<f64> {
         &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<f64>) -> Result<Self> {
+        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<Array2<f64>> for Psi {
     fn from(array: Array2<f64>) -> 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<Mat<f64>> for Psi {
     fn from(matrix: Mat<f64>) -> Self {
-        Psi { matrix }
+        let nrows = matrix.nrows();
+        Psi {
+            matrix,
+            row_log_scales: vec![0.0; nrows],
+        }
     }
 }
 
 impl From<&Array2<f64>> for Psi {
     fn from(array: &Array2<f64>) -> 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,16 +266,30 @@ 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)]
 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();