use super::*;
#[derive(Debug)]
pub enum GamlssError {
DimensionMismatch { reason: String },
InvalidInput { reason: String },
NonFinite { reason: String },
UnsupportedConfiguration { reason: String },
ConstraintViolation { reason: String },
NumericalFailure { reason: String },
}
crate::impl_reason_error_boilerplate! {
GamlssError {
DimensionMismatch,
InvalidInput,
NonFinite,
UnsupportedConfiguration,
ConstraintViolation,
NumericalFailure,
}
}
impl From<crate::families::block_layout::block_count::BlockCountMismatch> for GamlssError {
fn from(err: crate::families::block_layout::block_count::BlockCountMismatch) -> GamlssError {
GamlssError::DimensionMismatch {
reason: err.message(),
}
}
}
pub(crate) const MIN_PROB: f64 = 1e-10;
pub(crate) const MIN_DERIV: f64 = 1e-8;
use crate::types::MIN_WEIGHT;
pub(crate) const ETA_HARD_CLAMP: f64 = 30.0;
#[inline]
pub(crate) fn saturated_exp_eta(eta: f64) -> f64 {
eta.clamp(-ETA_HARD_CLAMP, ETA_HARD_CLAMP)
.exp()
.max(MIN_WEIGHT)
}
pub(crate) const WARMSTART_LOG_LAMBDA_FLOOR: f64 = 1e-12;
pub(crate) const EXACT_DENSE_BLOCK_BUDGET_BYTES: usize = 512 * 1024 * 1024;
pub(crate) const EXACT_DENSE_TOTAL_BUDGET_BYTES: usize = 2 * 1024 * 1024 * 1024;
pub(crate) const GAMLSS_ROWWISE_PAR_MIN_N: usize = 4096;
pub(crate) const GAMLSS_PROJECTED_TRACE_TARGET_BYTES: usize = 32 * 1024 * 1024;
pub(crate) const GAMLSS_PROJECTED_TRACE_MIN_CHUNK_ROWS: usize = 64;
pub(crate) const GAMLSS_PROJECTED_TRACE_MAX_CHUNK_ROWS: usize = 8192;
pub(crate) fn gamlss_projected_trace_chunk_rows(
rank: usize,
projected_channel_count: usize,
gram_column_count: usize,
) -> usize {
let per_row_values = rank
.saturating_mul(projected_channel_count.max(1))
.saturating_add(gram_column_count.max(1))
.max(1);
let per_row_bytes = per_row_values.saturating_mul(std::mem::size_of::<f64>());
let rows = GAMLSS_PROJECTED_TRACE_TARGET_BYTES / per_row_bytes.max(1);
rows.clamp(
GAMLSS_PROJECTED_TRACE_MIN_CHUNK_ROWS,
GAMLSS_PROJECTED_TRACE_MAX_CHUNK_ROWS,
)
}
pub(crate) fn gamlss_rowwise_map<F>(n: usize, f: F) -> Array1<f64>
where
F: Fn(usize) -> f64 + Sync,
{
if n >= GAMLSS_ROWWISE_PAR_MIN_N {
Array1::from((0..n).into_par_iter().map(&f).collect::<Vec<f64>>())
} else {
Array1::from_iter((0..n).map(f))
}
}
pub(crate) fn gamlss_rowwise_map_result<F>(n: usize, f: F) -> Result<Array1<f64>, String>
where
F: Fn(usize) -> Result<f64, String> + Sync,
{
if n >= GAMLSS_ROWWISE_PAR_MIN_N {
let values: Result<Vec<f64>, String> = (0..n).into_par_iter().map(&f).collect();
Ok(Array1::from(values?))
} else {
let mut out = Array1::<f64>::zeros(n);
for i in 0..n {
out[i] = f(i)?;
}
Ok(out)
}
}
pub(crate) enum DenseOrOperator<'a> {
Borrowed(&'a Array2<f64>),
Owned(Array2<f64>),
Operator(DesignMatrix),
}
impl DenseOrOperator<'_> {
pub(crate) fn nrows(&self) -> usize {
match self {
Self::Borrowed(dense) => dense.nrows(),
Self::Owned(dense) => dense.nrows(),
Self::Operator(design) => design.nrows(),
}
}
pub(crate) fn ncols(&self) -> usize {
match self {
Self::Borrowed(dense) => dense.ncols(),
Self::Owned(dense) => dense.ncols(),
Self::Operator(design) => design.ncols(),
}
}
pub(crate) fn row_chunk(&self, rows: std::ops::Range<usize>) -> Result<Array2<f64>, String> {
match self {
Self::Borrowed(dense) => Ok(dense.slice(s![rows, ..]).to_owned()),
Self::Owned(dense) => Ok(dense.slice(s![rows, ..]).to_owned()),
Self::Operator(design) => design.try_row_chunk(rows).map_err(|e| e.to_string()),
}
}
pub(crate) fn dot(&self, beta: ArrayView1<'_, f64>) -> Array1<f64> {
let n = self.nrows();
let p = self.ncols();
assert_eq!(beta.len(), p);
match self {
Self::Borrowed(dense) => fast_av(*dense, &beta),
Self::Owned(dense) => fast_av(dense, &beta),
Self::Operator(design) => {
let mut out = Array1::<f64>::zeros(n);
for rows in exact_design_row_chunks(n, p) {
let chunk = design
.try_row_chunk(rows.clone())
.expect("gamlss DesignSlot::dot: design row chunk materialization failed");
out.slice_mut(s![rows]).assign(&fast_av(&chunk, &beta));
}
out
}
}
}
}
pub(crate) fn dense_block_from_spec<'a>(
spec: &'a ParameterBlockSpec,
material_policy: &crate::resource::MaterializationPolicy,
materialization_label: &str,
) -> Result<Cow<'a, Array2<f64>>, String> {
match spec.design.as_dense_ref() {
Some(d) => Ok(Cow::Borrowed(d)),
None => Ok(Cow::Owned(
spec.design
.try_to_dense_with_policy(material_policy, "gamlss dense_block_from_spec")
.map_err(|e| format!("{materialization_label}: {e}"))?
.as_ref()
.clone(),
)),
}
}
pub(crate) fn dense_locscale_block_designs_fromspecs<'a>(
specs: &'a [ParameterBlockSpec],
expected_count: usize,
family_name: &str,
short_family_name: &str,
primary_block_idx: usize,
log_sigma_block_idx: usize,
primary_label: &str,
material_policy: &crate::resource::MaterializationPolicy,
) -> Result<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>), String> {
if specs.len() != expected_count {
return Err(GamlssError::DimensionMismatch {
reason: format!(
"{family_name} expects {expected_count} specs, got {}",
specs.len()
),
}
.into());
}
let primary = dense_block_from_spec(
&specs[primary_block_idx],
material_policy,
&format!("{short_family_name} dense_block_designs_fromspecs {primary_label}"),
)?;
let log_sigma = dense_block_from_spec(
&specs[log_sigma_block_idx],
material_policy,
&format!("{short_family_name} dense_block_designs_fromspecs log_sigma"),
)?;
Ok((primary, log_sigma))
}
pub(crate) fn dense_locscale_block_designs_cached<'a>(
primary_design: Option<&'a DesignMatrix>,
log_sigma_design: Option<&'a DesignMatrix>,
family_name: &str,
short_family_name: &str,
primary_label: &str,
material_policy: &crate::resource::MaterializationPolicy,
) -> Result<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>), String> {
let primary_design = primary_design
.ok_or_else(|| format!("{family_name} exact path is missing {primary_label} design"))?;
let log_sigma_design = log_sigma_design
.ok_or_else(|| format!("{family_name} exact path is missing log-sigma design"))?;
let primary = match primary_design.as_dense_ref() {
Some(d) => Cow::Borrowed(d),
None => Cow::Owned(
primary_design
.try_to_dense_with_policy(material_policy, "gamlss dense_locscale_block_designs")
.map_err(|e| {
format!("{short_family_name} dense_block_designs {primary_label}: {e}")
})?
.as_ref()
.clone(),
),
};
let log_sigma = match log_sigma_design.as_dense_ref() {
Some(d) => Cow::Borrowed(d),
None => Cow::Owned(
log_sigma_design
.try_to_dense_with_policy(material_policy, "gamlss dense_locscale_block_designs")
.map_err(|e| format!("{short_family_name} dense_block_designs log_sigma: {e}"))?
.as_ref()
.clone(),
),
};
Ok((primary, log_sigma))
}
pub(crate) struct LocScalePsiDirectionParts {
pub(crate) block_idx: usize,
pub(crate) local_idx: usize,
pub(crate) primary_psi: PsiDesignMap,
pub(crate) log_sigma_psi: PsiDesignMap,
pub(crate) primary_z: Array1<f64>,
pub(crate) log_sigma_z: Array1<f64>,
}
pub(crate) fn locscale_joint_psi_direction_parts(
block_states: &[ParameterBlockState],
derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
psi_index: usize,
n: usize,
p_primary: usize,
p_log_sigma: usize,
primary_block_idx: usize,
log_sigma_block_idx: usize,
expected_blocks: usize,
family_name: &str,
primary_label: &str,
policy: &crate::resource::ResourcePolicy,
) -> Result<Option<LocScalePsiDirectionParts>, String> {
validate_block_count::<GamlssError>(family_name, expected_blocks, block_states.len())?;
if derivative_blocks.len() != expected_blocks {
return Err(GamlssError::DimensionMismatch {
reason: format!(
"{family_name} joint psi direction expects {expected_blocks} derivative block lists, got {}",
derivative_blocks.len()
),
}
.into());
}
let beta_primary = &block_states[primary_block_idx].beta;
let beta_log_sigma = &block_states[log_sigma_block_idx].beta;
let mut global = 0usize;
for (block_idx, block_derivs) in derivative_blocks.iter().enumerate() {
for (local_idx, deriv) in block_derivs.iter().enumerate() {
if global == psi_index {
let primary_psi;
let log_sigma_psi;
let primary_z;
let log_sigma_z;
if block_idx == primary_block_idx {
primary_psi = resolve_custom_family_x_psi_map(
deriv,
n,
p_primary,
0..n,
&format!("{family_name} {primary_label}"),
policy,
)?;
primary_z = primary_psi
.forward_mul(beta_primary.view())
.map_err(|e| format!("{family_name} {primary_label} forward_mul: {e}"))?;
log_sigma_psi = PsiDesignMap::Zero {
nrows: n,
ncols: p_log_sigma,
};
log_sigma_z = Array1::<f64>::zeros(n);
} else if block_idx == log_sigma_block_idx {
log_sigma_psi = resolve_custom_family_x_psi_map(
deriv,
n,
p_log_sigma,
0..n,
&format!("{family_name} log-sigma"),
policy,
)?;
log_sigma_z = log_sigma_psi
.forward_mul(beta_log_sigma.view())
.map_err(|e| format!("{family_name} log-sigma forward_mul: {e}"))?;
primary_psi = PsiDesignMap::Zero {
nrows: n,
ncols: p_primary,
};
primary_z = Array1::<f64>::zeros(n);
} else {
return Ok(None);
}
return Ok(Some(LocScalePsiDirectionParts {
block_idx,
local_idx,
primary_psi,
log_sigma_psi,
primary_z,
log_sigma_z,
}));
}
global += 1;
}
}
Ok(None)
}
pub(crate) struct LocScalePsiDriftConfig<'a> {
pub(crate) n: usize,
pub(crate) p_primary: usize,
pub(crate) p_log_sigma: usize,
pub(crate) primary_block_idx: usize,
pub(crate) log_sigma_block_idx: usize,
pub(crate) family_name: &'a str,
pub(crate) primary_label: &'a str,
pub(crate) policy: &'a crate::resource::ResourcePolicy,
}
pub(crate) fn locscale_joint_psisecond_design_drifts(
block_states: &[ParameterBlockState],
derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
psi_a: &LocationScaleJointPsiDirection,
psi_b: &LocationScaleJointPsiDirection,
cfg: LocScalePsiDriftConfig<'_>,
) -> Result<LocationScaleJointPsiSecondDrifts, String> {
let beta_primary = &block_states[cfg.primary_block_idx].beta;
let beta_log_sigma = &block_states[cfg.log_sigma_block_idx].beta;
let mut primary_ab_action = None;
let mut log_sigma_ab_action = None;
let mut primary_ab = None;
let mut log_sigma_ab = None;
if psi_a.block_idx == psi_b.block_idx {
let deriv = &derivative_blocks[psi_a.block_idx][psi_a.local_idx];
let deriv_b = &derivative_blocks[psi_b.block_idx][psi_b.local_idx];
if psi_a.block_idx == cfg.primary_block_idx {
let (action, matrix) = psi_psi_map_to_drift_slots(
deriv,
deriv_b,
psi_b.local_idx,
cfg.n,
cfg.p_primary,
&format!("{} {}", cfg.family_name, cfg.primary_label),
cfg.policy,
)?;
primary_ab_action = action;
primary_ab = matrix;
} else if psi_a.block_idx == cfg.log_sigma_block_idx {
let (action, matrix) = psi_psi_map_to_drift_slots(
deriv,
deriv_b,
psi_b.local_idx,
cfg.n,
cfg.p_log_sigma,
&format!("{} log-sigma", cfg.family_name),
cfg.policy,
)?;
log_sigma_ab_action = action;
log_sigma_ab = matrix;
}
}
let z_primary_ab = second_psi_linear_map(
primary_ab_action.as_ref(),
primary_ab.as_ref(),
cfg.n,
cfg.p_primary,
)
.forward_mul(beta_primary.view());
let z_ls_ab = second_psi_linear_map(
log_sigma_ab_action.as_ref(),
log_sigma_ab.as_ref(),
cfg.n,
cfg.p_log_sigma,
)
.forward_mul(beta_log_sigma.view());
Ok(LocationScaleJointPsiSecondDrifts {
x_primary_ab_action: primary_ab_action,
x_ls_ab_action: log_sigma_ab_action,
x_primary_ab: primary_ab,
x_ls_ab: log_sigma_ab,
z_primary_ab,
z_ls_ab,
})
}
pub(crate) fn psi_psi_map_to_drift_slots(
deriv: &crate::custom_family::CustomFamilyBlockPsiDerivative,
deriv_b: &crate::custom_family::CustomFamilyBlockPsiDerivative,
local_idx_b: usize,
n: usize,
p: usize,
label: &str,
policy: &crate::resource::ResourcePolicy,
) -> Result<
(
Option<crate::custom_family::CustomFamilyPsiSecondDesignAction>,
Option<Array2<f64>>,
),
String,
> {
match resolve_custom_family_x_psi_psi_map(
deriv,
deriv_b,
local_idx_b,
n,
p,
0..n,
label,
policy,
)? {
crate::custom_family::PsiDesignMap::Second { action } => Ok((Some(action), None)),
crate::custom_family::PsiDesignMap::Dense { matrix } => Ok((None, Some((*matrix).clone()))),
crate::custom_family::PsiDesignMap::Zero { .. } => Ok((None, None)),
crate::custom_family::PsiDesignMap::First { .. } => {
Err(GamlssError::UnsupportedConfiguration {
reason: format!("{label}: unexpected First variant from _psi_psi_map"),
}
.into())
}
}
}
pub(crate) fn dense_block_or_operator<'a>(
design: &'a DesignMatrix,
n: usize,
p: usize,
budget_bytes: usize,
policy: &crate::resource::ResourcePolicy,
) -> DenseOrOperator<'a> {
if let Some(dense) = design.as_dense_ref() {
return DenseOrOperator::Borrowed(dense);
}
let dense_bytes = 8usize.saturating_mul(n).saturating_mul(p);
if dense_bytes <= budget_bytes
&& let Ok(arc) = design
.try_to_dense_with_policy(&policy.material_policy(), "gamlss dense_block_or_operator")
{
return DenseOrOperator::Owned(arc.as_ref().clone());
}
DenseOrOperator::Operator(design.clone())
}
pub(crate) fn dense_blocks_planned_budget(blocks: &[&DesignMatrix]) -> Vec<usize> {
let mut planned = vec![0; blocks.len()];
let mut total = 0usize;
for (idx, design) in blocks.iter().enumerate() {
if design.as_dense_ref().is_some() {
continue;
}
let bytes = 8usize
.saturating_mul(design.nrows())
.saturating_mul(design.ncols());
if bytes <= EXACT_DENSE_BLOCK_BUDGET_BYTES
&& total.saturating_add(bytes) <= EXACT_DENSE_TOTAL_BUDGET_BYTES
{
planned[idx] = bytes;
total += bytes;
}
}
planned
}
pub(crate) fn exact_design_row_chunks(
n: usize,
p: usize,
) -> impl Iterator<Item = std::ops::Range<usize>> {
const TARGET_BYTES: usize = 8 * 1024 * 1024;
const MIN_ROWS: usize = 512;
const MAX_ROWS: usize = 131_072;
let rows = (TARGET_BYTES / (p.max(1) * 8))
.clamp(MIN_ROWS, MAX_ROWS)
.min(n.max(1));
(0..n)
.step_by(rows)
.map(move |start| start..(start + rows).min(n))
}
pub(crate) fn design_weighted_column_squares(
design: &DesignMatrix,
weights: &Array1<f64>,
) -> Result<Array1<f64>, String> {
let n = design.nrows();
let p = design.ncols();
if weights.len() != n {
return Err(GamlssError::DimensionMismatch {
reason: format!(
"design weighted column squares dimension mismatch: weights={}, rows={}",
weights.len(),
n
),
}
.into());
}
let mut out = Array1::<f64>::zeros(p);
for rows in exact_design_row_chunks(n, p) {
let chunk = design.try_row_chunk(rows.clone()).map_err(|e| {
format!("design weighted column squares row chunk materialization failed: {e}")
})?;
for (local_i, row) in chunk.outer_iter().enumerate() {
let w = weights[rows.start + local_i];
if w == 0.0 {
continue;
}
for j in 0..p {
let x = row[j];
out[j] += w * x * x;
}
}
}
Ok(out)
}
#[inline]
pub(crate) fn floor_positiveweight(rawweight: f64, minweight: f64) -> f64 {
if !rawweight.is_finite() || rawweight <= 0.0 {
0.0
} else {
rawweight.max(minweight)
}
}
#[inline]
pub(crate) fn logb_dlog_sigma_deta(sigma: f64, d_sigma_deta: f64) -> f64 {
if d_sigma_deta.is_infinite() {
1.0
} else {
let value = d_sigma_deta / sigma;
if value.is_finite() {
value.clamp(0.0, 1.0)
} else {
0.0
}
}
}
#[inline]
pub(crate) fn gaussian_log_sigma_irlsinfo_directional_derivative(
weight: f64,
sigma: f64,
d_sigma_deta: f64,
d_eta: f64,
) -> f64 {
if weight == 0.0 || d_eta == 0.0 || !sigma.is_finite() || sigma <= 0.0 {
return 0.0;
}
let g = logb_dlog_sigma_deta(sigma, d_sigma_deta);
if !g.is_finite() || !(0.0..1.0).contains(&g) {
return 0.0;
}
let rawinfo = 2.0 * weight * g * g;
if !rawinfo.is_finite() || rawinfo <= MIN_WEIGHT {
return 0.0;
}
let dg_deta = g * (1.0 - g);
let dw = 4.0 * weight * g * dg_deta * d_eta;
if dw.is_finite() { dw } else { 0.0 }
}
#[derive(Clone, Copy)]
pub(crate) struct GaussianDiagonalRowKernel {
pub(crate) log_likelihood: f64,
pub(crate) location_working_weight: f64,
pub(crate) location_working_shift: f64,
pub(crate) log_sigma_working_weight: f64,
pub(crate) log_sigma_working_response: f64,
}
#[inline]
pub(crate) fn gaussian_diagonal_row_kernel(
y: f64,
location_eta: f64,
eta_log_sigma: f64,
obs_weight: f64,
ln2pi: f64,
) -> GaussianDiagonalRowKernel {
if obs_weight == 0.0 {
return GaussianDiagonalRowKernel {
log_likelihood: 0.0,
location_working_weight: 0.0,
location_working_shift: 0.0,
log_sigma_working_weight: 0.0,
log_sigma_working_response: eta_log_sigma,
};
}
let SigmaJet1 { sigma, d1 } = logb_sigma_jet1_scalar(eta_log_sigma);
let inv_s2 = (sigma * sigma).recip();
let residual = y - location_eta;
let location_working_weight = floor_positiveweight(obs_weight * inv_s2, MIN_WEIGHT);
let dlog_sigma_deta = logb_dlog_sigma_deta(sigma, d1);
let log_sigma_working_weight = floor_positiveweight(
2.0 * obs_weight * dlog_sigma_deta * dlog_sigma_deta,
MIN_WEIGHT,
);
let log_sigma_score = obs_weight * (residual * residual * inv_s2 - 1.0) * dlog_sigma_deta;
let log_sigma_working_response = if log_sigma_working_weight == 0.0 {
eta_log_sigma
} else {
eta_log_sigma + log_sigma_score / log_sigma_working_weight
};
GaussianDiagonalRowKernel {
log_likelihood: obs_weight
* (-0.5 * (residual * residual * inv_s2 + ln2pi + 2.0 * sigma.ln())),
location_working_weight,
location_working_shift: residual,
log_sigma_working_weight,
log_sigma_working_response,
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ParameterLink {
Identity,
Log,
Logit,
Probit,
InverseLink,
Wiggle,
}