pmcore/iov/mod.rs
1//! SDE-based Inter-Occasion Variability (IOV) analysis.
2//!
3//! This module provides [`optimize_diffusion`](crate::iov::DiffusionOptimize::optimize_diffusion),
4//! which optimizes SDE diffusion (sigma) parameters for each support point independently, using the
5//! NelderMead algorithm. The optimization runs in parallel over support points via rayon.
6//!
7//! # Workflow
8//!
9//! 1. Fit an ODE model with NPAG/NPOD to obtain support points (Stage 1).
10//! 2. Add sigma parameter columns to the theta using
11//! [`Theta::with_added_parameter`](crate::estimation::nonparametric::Theta::with_added_parameter).
12//! 3. Construct an SDE model (user-provided) that maps sigma parameters to
13//! diffusion terms.
14//! 4. Call [`optimize_diffusion`](crate::iov::DiffusionOptimize::optimize_diffusion) to optimize sigma per support point.
15//!
16//! # Example
17//!
18//! ```ignore
19//! use pmcore::prelude::*;
20//! use pmcore::iov::DiffusionOptimize;
21//! use pmcore::iov::DiffusionConfig;
22//!
23//! let r_ode = problem.fit_with(NPAG::default())?;
24//! let mut joint = r_ode.get_theta()
25//! .with_added_parameter("ske", 1e-6, 1.0, 0.01)?;
26//!
27//! let diff = sde.optimize_diffusion(
28//! &r_ode.data(), &mut joint,
29//! &["ske"], &r_ode.error_models(),
30//! DiffusionConfig::default(),
31//! )?;
32//! ```
33
34mod optimizer;
35
36use anyhow::bail;
37use rayon::prelude::*;
38
39use pharmsol::prelude::data::AssayErrorModels;
40use pharmsol::{Data, SDE};
41
42use crate::estimation::nonparametric::Theta;
43
44/// Configuration for SDE diffusion parameter optimization.
45#[derive(Debug, Clone)]
46pub struct DiffusionConfig {
47 /// Maximum NelderMead iterations per support point.
48 ///
49 /// Default: 50. Set lower for speed, higher for difficult surfaces.
50 pub max_iter: usize,
51
52 /// Convergence tolerance on simplex standard deviation.
53 ///
54 /// NelderMead stops when the standard deviation of function values
55 /// across the simplex vertices falls below this threshold.
56 /// Default: 1e-3.
57 pub sd_tolerance: f64,
58
59 /// Fraction of the distance to the upper bound for simplex construction.
60 ///
61 /// The second simplex vertex is placed at `init + perturbation × (upper − init)`.
62 /// This spans the simplex toward the upper bound without overshooting.
63 /// Default: 0.15 (for init=0.01, bounds [0,0.5] → vertex at 0.084).
64 pub initial_perturbation: f64,
65
66 /// Number of resampled evaluations per cost function call.
67 ///
68 /// The SDE particle filter produces noisy likelihood estimates. Averaging
69 /// over multiple evaluations reduces variance and makes NelderMead decisions
70 /// reliable by giving every vertex the same precision.
71 ///
72 /// Default: 3. Set to 1 for speed (raw NelderMead), higher for precision.
73 pub resampling_samples: usize,
74}
75
76impl Default for DiffusionConfig {
77 fn default() -> Self {
78 Self {
79 max_iter: 50,
80 sd_tolerance: 1e-3,
81 initial_perturbation: 0.15,
82 resampling_samples: 3,
83 }
84 }
85}
86
87/// Results of SDE diffusion parameter optimization.
88#[derive(Debug, Clone)]
89pub struct DiffusionResult {
90 /// Final log-likelihood for each support point after optimization.
91 /// Length equals `theta.nspp()`.
92 pub per_point_likelihood: Vec<f64>,
93
94 /// Number of NelderMead iterations used for each point.
95 pub per_point_iterations: Vec<usize>,
96
97 /// Whether NelderMead converged within `max_iter` for each point.
98 pub per_point_converged: Vec<bool>,
99}
100
101/// Trait for SDEs that support diffusion parameter optimization.
102///
103/// This enables method-style calls: `sde.optimize_diffusion(...)`.
104pub trait DiffusionOptimize {
105 /// Optimize SDE diffusion parameters for each support point independently.
106 ///
107 /// Modifies `theta` **in-place**: for each support point, the sigma parameter
108 /// columns are replaced with values that maximize the log-likelihood of all
109 /// subjects under this SDE. Primary (non-sigma) parameter values are held fixed.
110 ///
111 /// If `posterior` is provided, subject contributions are weighted by their
112 /// posterior responsibility for each support point: `p(z_i=j)` from Stage 1.
113 /// If `None`, falls back to uniform weighting.
114 ///
115 /// # Panics
116 ///
117 /// Panics if any name in `sigma_params` is not found in `theta.parameters()`.
118 fn optimize_diffusion(
119 &self,
120 data: &Data,
121 theta: &mut Theta,
122 sigma_params: &[String],
123 error_models: &AssayErrorModels,
124 posterior: Option<&crate::estimation::nonparametric::Posterior>,
125 config: DiffusionConfig,
126 ) -> anyhow::Result<DiffusionResult>;
127}
128
129impl DiffusionOptimize for SDE {
130 fn optimize_diffusion(
131 &self,
132 data: &Data,
133 theta: &mut Theta,
134 sigma_params: &[String],
135 error_models: &AssayErrorModels,
136 posterior: Option<&crate::estimation::nonparametric::Posterior>,
137 config: DiffusionConfig,
138 ) -> anyhow::Result<DiffusionResult> {
139 optimize_diffusion(
140 self,
141 data,
142 theta,
143 sigma_params,
144 error_models,
145 posterior,
146 config,
147 )
148 }
149}
150
151/// Optimize SDE diffusion parameters for each support point independently.
152///
153/// Free-function form of [`DiffusionOptimize::optimize_diffusion`].
154/// Prefer `sde.optimize_diffusion(...)` for readability.
155///
156/// # Important: disable SDE caching
157///
158/// SDEs cache likelihood results by default. This optimization is a Monte Carlo
159/// method that requires fresh random evaluations every iteration. Ensure the SDE
160/// is constructed with `.disable_cache()` before passing it here. This function
161/// warns (but does not error) if caching may be enabled.
162pub(crate) fn optimize_diffusion(
163 sde: &SDE,
164 data: &Data,
165 theta: &mut Theta,
166 sigma_params: &[String],
167 error_models: &AssayErrorModels,
168 posterior: Option<&crate::estimation::nonparametric::Posterior>,
169 config: DiffusionConfig,
170) -> anyhow::Result<DiffusionResult> {
171 let n_spp = theta.nspp();
172 if n_spp == 0 {
173 bail!("theta has no support points");
174 }
175
176 // Resolve sigma parameter indices in theta
177 let sigma_indices: Vec<usize> = sigma_params
178 .iter()
179 .map(|name| {
180 theta
181 .parameters()
182 .iter()
183 .position(|p| p.name.as_str() == name.as_str())
184 .unwrap_or_else(|| {
185 panic!(
186 "sigma parameter '{}' not found in theta parameters: {:?}",
187 name,
188 theta.parameters().names()
189 )
190 })
191 })
192 .collect();
193
194 // Identify primary parameter indices (all others)
195 let n_total = theta.matrix().ncols();
196 let sigma_set: std::collections::HashSet<usize> = sigma_indices.iter().copied().collect();
197 let primary_indices: Vec<usize> = (0..n_total).filter(|i| !sigma_set.contains(i)).collect();
198
199 // Check for sigma initialized to zero
200 for &si in &sigma_indices {
201 for r in 0..n_spp {
202 if theta.matrix()[(r, si)] == 0.0 {
203 tracing::warn!(
204 "sigma parameter at column {} (support point {}) initialized to 0.0; \
205 the SDE degenerates to an ODE at sigma=0. Consider using a small \
206 non-zero initial value (e.g., 0.01)",
207 si,
208 r
209 );
210 }
211 }
212 }
213
214 // Extract sigma parameter bounds for simplex construction
215 let sigma_bounds: Vec<(f64, f64)> = sigma_indices
216 .iter()
217 .map(|&si| {
218 let bp = &theta.parameters().items[si];
219 (bp.lower, bp.upper)
220 })
221 .collect();
222
223 // Parallel optimization over support points — each SP optimized independently.
224 // If a Stage 1 posterior is provided, subject contributions are weighted by
225 // p(z_i=j), correctly modeling population structure without inner-loop Burke.
226 let results: Vec<optimizer::OptimizationOutcome> = (0..n_spp)
227 .into_par_iter()
228 .map(|i| {
229 let primary: Vec<f64> = primary_indices
230 .iter()
231 .map(|&pi| theta.matrix()[(i, pi)])
232 .collect();
233
234 let sigma_init: Vec<f64> = sigma_indices
235 .iter()
236 .map(|&si| theta.matrix()[(i, si)])
237 .collect();
238
239 // Extract posterior responsibilities for this SP (if available)
240 let responsibilities: Option<Vec<f64>> = posterior.map(|p| {
241 (0..data.subjects().len())
242 .map(|s| p.matrix()[(s, i)])
243 .collect()
244 });
245 let resp_slice: Option<&[f64]> = responsibilities.as_deref();
246
247 let cost = optimizer::SigmaCost::new(
248 sde,
249 data,
250 &primary,
251 &primary_indices,
252 &sigma_indices,
253 error_models,
254 resp_slice,
255 );
256
257 optimizer::optimize_sigma(cost, &sigma_init, &sigma_bounds, &config)
258 })
259 .collect();
260
261 // Update theta with optimized sigma values
262 let mut per_point_likelihood = Vec::with_capacity(n_spp);
263 let mut per_point_iterations = Vec::with_capacity(n_spp);
264 let mut per_point_converged = Vec::with_capacity(n_spp);
265
266 for (i, outcome) in results.iter().enumerate() {
267 for (j, &si) in sigma_indices.iter().enumerate() {
268 theta.matrix_mut()[(i, si)] = outcome.optimized_params[j];
269 }
270 per_point_likelihood.push(-outcome.final_cost);
271 per_point_iterations.push(outcome.iterations);
272 per_point_converged.push(outcome.converged);
273 }
274
275 let n_converged = per_point_converged.iter().filter(|&&c| c).count();
276 tracing::info!(
277 "SDE IOV optimization: {}/{} support points converged, \
278 mean iterations: {:.1}, mean log-likelihood: {:.2}",
279 n_converged,
280 n_spp,
281 per_point_iterations.iter().sum::<usize>() as f64 / n_spp.max(1) as f64,
282 per_point_likelihood.iter().sum::<f64>() / n_spp.max(1) as f64,
283 );
284
285 Ok(DiffusionResult {
286 per_point_likelihood,
287 per_point_iterations,
288 per_point_converged,
289 })
290}