use super::*;
impl SurvivalMarginalSlopeFamily {
#[inline]
pub(crate) fn probit_frailty_scale(&self) -> f64 {
probit_frailty_scale(self.gaussian_frailty_sd)
}
pub(crate) fn z_subsample_key(&self) -> Array1<f64> {
self.z.column(0).to_owned()
}
#[inline]
pub(crate) fn score_dim(&self) -> usize {
assert_eq!(self.score_covariance.dim(), self.z.ncols());
self.z.ncols()
}
pub(crate) fn score_warp_basis_dim(&self) -> usize {
self.score_warp
.as_ref()
.map_or(0, DeviationRuntime::basis_dim)
}
pub(crate) fn score_warp_coord_basis_index(
&self,
local_idx: usize,
) -> Result<(usize, usize), String> {
let basis_dim = self.score_warp_basis_dim();
if basis_dim == 0 {
return Err(
"survival score-warp coordinate lookup without score-warp runtime".to_string(),
);
}
let coord = local_idx / basis_dim;
if coord >= self.score_dim() {
return Err(SurvivalMarginalSlopeError::IncompatibleDimensions {
reason: format!(
"survival score-warp local index {local_idx} exceeds K={} per-z blocks with basis dim {basis_dim}",
self.score_dim()
),
}
.into());
}
Ok((coord, local_idx % basis_dim))
}
pub(crate) fn score_warp_beta_for_coord(
&self,
beta_h: &Array1<f64>,
coord: usize,
) -> Result<Array1<f64>, String> {
let runtime = self
.score_warp
.as_ref()
.ok_or_else(|| "missing survival score-warp runtime".to_string())?;
score_warp_component_beta(runtime, beta_h, coord)
}
#[inline]
pub(crate) fn zero_score_warp_span() -> exact_kernel::LocalSpanCubic {
exact_kernel::LocalSpanCubic {
left: 0.0,
right: 1.0,
c0: 0.0,
c1: 0.0,
c2: 0.0,
c3: 0.0,
}
}
pub(crate) fn add_local_span_cubic(
left: f64,
right: f64,
target: &mut exact_kernel::LocalSpanCubic,
span: exact_kernel::LocalSpanCubic,
) {
target.left = left;
target.right = right;
target.c0 += span.c0;
target.c1 += span.c1;
target.c2 += span.c2;
target.c3 += span.c3;
}
pub(crate) fn score_warp_local_cubic_at(
&self,
beta_h: Option<&Array1<f64>>,
value: f64,
) -> Result<exact_kernel::LocalSpanCubic, String> {
let (Some(runtime), Some(beta_h)) = (self.score_warp.as_ref(), beta_h) else {
return Ok(Self::zero_score_warp_span());
};
if self.score_dim() == 1 {
return runtime.local_cubic_at(beta_h, value);
}
let mut sum = Self::zero_score_warp_span();
for coord in 0..self.score_dim() {
let local_beta = score_warp_component_beta(runtime, beta_h, coord)?;
let span = runtime.local_cubic_at(&local_beta, value)?;
if coord == 0 {
sum = exact_kernel::LocalSpanCubic {
left: span.left,
right: span.right,
c0: 0.0,
c1: 0.0,
c2: 0.0,
c3: 0.0,
};
}
Self::add_local_span_cubic(span.left, span.right, &mut sum, span);
}
Ok(sum)
}
pub(crate) fn score_warp_observed_value(
&self,
row: usize,
beta_h: Option<&Array1<f64>>,
) -> Result<f64, String> {
let (Some(runtime), Some(beta_h)) = (self.score_warp.as_ref(), beta_h) else {
return Ok(0.0);
};
let mut value = 0.0;
for coord in 0..self.score_dim() {
let local_beta = score_warp_component_beta(runtime, beta_h, coord)?;
let z_coord = self.z[[row, coord]];
value += runtime
.local_cubic_at(&local_beta, z_coord)?
.evaluate(z_coord);
}
Ok(value)
}
pub(crate) fn observed_score_projection(&self, row: usize) -> f64 {
if self.score_dim() == 1 {
self.z[[row, 0]]
} else {
self.z.row(row).sum()
}
}
pub(crate) fn integration_score_basis_coefficients(
&self,
local_idx: usize,
z_basis: f64,
multiplier: f64,
) -> Result<[f64; 4], String> {
let runtime = self
.score_warp
.as_ref()
.ok_or_else(|| "missing survival score-warp runtime".to_string())?;
let (_, basis_idx) = self.score_warp_coord_basis_index(local_idx)?;
let basis_span = runtime.basis_cubic_at(basis_idx, z_basis)?;
Ok(exact_kernel::score_basis_cell_coefficients(
basis_span, multiplier,
))
}
pub(crate) fn observed_score_basis_coefficients(
&self,
row: usize,
local_idx: usize,
z_obs: f64,
multiplier: f64,
) -> Result<[f64; 4], String> {
let runtime = self
.score_warp
.as_ref()
.ok_or_else(|| "missing survival score-warp runtime".to_string())?;
if self.score_dim() == 1 {
let basis_span = runtime.basis_cubic_at(local_idx, z_obs)?;
return Ok(exact_kernel::score_basis_cell_coefficients(
basis_span, multiplier,
));
}
let (coord, basis_idx) = self.score_warp_coord_basis_index(local_idx)?;
let z_coord = self.z[[row, coord]];
let basis_span = runtime.basis_cubic_at(basis_idx, z_coord)?;
Ok([multiplier * basis_span.evaluate(z_coord), 0.0, 0.0, 0.0])
}
pub(crate) fn logslope_vector_for_row(
&self,
row: usize,
logslope_eta: &Array1<f64>,
) -> Result<Vec<f64>, String> {
let k = self.score_dim();
if self.logslope_surface_ranges.len() == k && k > 1 && logslope_eta.len() == self.n {
return Err(
"survival marginal-slope internal logslope vector requested scalar eta for a per-z surface layout"
.to_string(),
);
}
if logslope_eta.len() == self.n {
return Ok(vec![logslope_eta[row]; k]);
}
if logslope_eta.len() == self.n * k {
let start = row * k;
return Ok(logslope_eta.slice(s![start..start + k]).to_vec());
}
Err(SurvivalMarginalSlopeError::IncompatibleDimensions {
reason: format!(
"survival marginal-slope logslope eta length {} is incompatible with n={} and score dim K={k}",
logslope_eta.len(),
self.n
),
}
.into())
}
pub(crate) fn per_z_logslope_active(&self) -> bool {
self.score_dim() > 1 && self.logslope_surface_ranges.len() == self.score_dim()
}
pub(crate) fn logslope_surface_values_for_row(
&self,
row: usize,
beta_logslope: &Array1<f64>,
) -> Result<Vec<f64>, String> {
let k = self.score_dim();
if !self.per_z_logslope_active() {
return self.logslope_vector_for_row(row, beta_logslope);
}
let g_row = self
.logslope_design
.try_row_chunk(row..row + 1)
.map_err(|e| format!("logslope_surface_values_for_row logslope row chunk: {e}"))?;
let row_view = g_row.row(0);
let mut out = Vec::with_capacity(k);
for range in &self.logslope_surface_ranges {
out.push(
row_view
.slice(s![range.clone()])
.dot(&beta_logslope.slice(s![range.clone()])),
);
}
Ok(out)
}
pub(crate) fn logslope_surface_row(&self, row: usize) -> Result<Array1<f64>, String> {
let chunk = self
.logslope_design
.try_row_chunk(row..row + 1)
.map_err(|e| format!("logslope_surface_row logslope row chunk: {e}"))?;
Ok(chunk.row(0).to_owned())
}
pub(crate) fn shared_logslope_covariance_scale(&self) -> Result<f64, String> {
let k = self.score_dim();
if k == 1 {
return Ok(1.0);
}
let ones = vec![1.0; k];
self.score_covariance.quadratic_form(&ones).map_err(|err| {
format!("survival marginal-slope shared log-slope covariance scale: {err}")
})
}
pub(crate) fn exact_shared_score_summary(
&self,
row: usize,
block_states: &[ParameterBlockState],
context: &str,
) -> Result<(f64, f64), String> {
let k = self.score_dim();
if k == 1 {
return Ok((self.z[[row, 0]], 1.0));
}
let logslope_eta_len = block_states[2].eta.len();
if logslope_eta_len != self.n {
return Err(SurvivalMarginalSlopeError::IncompatibleDimensions {
reason: format!(
"{context}: survival marginal-slope exact shared-slope calculus for K={k} requires one log-slope eta per row (n={}); got eta len {logslope_eta_len}. Per-z log-slope derivatives require a {}-primary row kernel.",
self.n,
3 + k
),
}
.into());
}
Ok((
self.z.row(row).sum(),
self.shared_logslope_covariance_scale()?,
))
}
}