use ndarray::{Array2, Array3, ArrayView1};
use std::sync::Arc;
use crate::normalize_fisher_rao_blocks;
#[derive(Clone)]
pub enum WeightField {
Identity,
Factored {
u: Arc<Array2<f64>>,
rank: usize,
p_out: usize,
},
}
impl std::fmt::Debug for WeightField {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
WeightField::Identity => f.write_str("Identity"),
WeightField::Factored { u, rank, p_out } => f
.debug_struct("Factored")
.field("shape", &format_args!("{}×{}", u.nrows(), u.ncols()))
.field("rank", rank)
.field("p_out", p_out)
.finish(),
}
}
}
impl WeightField {
pub fn project_jac_row_with_u(
u_row: &[f64],
jac_row: &[f64],
p: usize,
rank: usize,
d: usize,
) -> Array2<f64> {
let mut m = Array2::<f64>::zeros((rank, d));
for k in 0..rank {
for a in 0..d {
let mut s = 0.0;
for i in 0..p {
s += u_row[i * rank + k] * jac_row[i * d + a];
}
m[[k, a]] = s;
}
}
m
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum MetricProvenance {
Euclidean,
OutputFisher { rank: usize },
OutputFisherDownstream { rank: usize },
BehavioralFisher { probes: usize },
WhitenedStructured { factor_rank: usize },
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum FisherFactorKind {
ExactFull,
CertifiedPsdLowerBound,
UncertifiedApproximation,
}
impl FisherFactorKind {
pub const fn tag(self) -> &'static str {
match self {
Self::ExactFull => "exact_full",
Self::CertifiedPsdLowerBound => "certified_psd_lower_bound",
Self::UncertifiedApproximation => "uncertified_approximation",
}
}
pub fn from_tag(tag: &str) -> Result<Self, String> {
match tag {
"exact_full" => Ok(Self::ExactFull),
"certified_psd_lower_bound" => Ok(Self::CertifiedPsdLowerBound),
"uncertified_approximation" => Ok(Self::UncertifiedApproximation),
other => Err(format!(
"fisher_factor_kind must be 'exact_full', 'certified_psd_lower_bound', or \
'uncertified_approximation'; got {other:?}"
)),
}
}
}
#[derive(Clone, Debug)]
pub struct RowMetric {
provenance: MetricProvenance,
n_rows: usize,
p: usize,
rank: usize,
factors: Option<Arc<Array2<f64>>>,
solver_delta: f64,
traces: ndarray::Array1<f64>,
fisher_factor_kind: Option<FisherFactorKind>,
truncation_mass_residual: Option<Arc<ndarray::Array1<f64>>>,
}
impl RowMetric {
pub fn euclidean(n_rows: usize, p: usize) -> Result<Self, String> {
Ok(Self {
provenance: MetricProvenance::Euclidean,
n_rows,
p,
rank: p,
factors: None,
solver_delta: 0.0,
traces: ndarray::Array1::<f64>::from_elem(n_rows, p as f64),
fisher_factor_kind: None,
truncation_mass_residual: None,
})
}
pub fn output_fisher(u: Arc<Array2<f64>>, p: usize, rank: usize) -> Result<Self, String> {
Self::from_factors(MetricProvenance::OutputFisher { rank }, u, p, rank, 0.0)
}
pub fn output_fisher_downstream(
u: Arc<Array2<f64>>,
p: usize,
rank: usize,
) -> Result<Self, String> {
Self::from_factors(
MetricProvenance::OutputFisherDownstream { rank },
u,
p,
rank,
0.0,
)
}
pub fn behavioral_fisher(u: Arc<Array2<f64>>, p: usize, probes: usize) -> Result<Self, String> {
Self::from_factors(
MetricProvenance::BehavioralFisher { probes },
u,
p,
probes,
0.0,
)
}
pub fn whitened_structured(u: Arc<Array2<f64>>, p: usize, rank: usize) -> Result<Self, String> {
Self::from_factors(
MetricProvenance::WhitenedStructured { factor_rank: rank },
u,
p,
rank,
0.0,
)
}
fn from_factors(
provenance: MetricProvenance,
u: Arc<Array2<f64>>,
p: usize,
rank: usize,
solver_delta: f64,
) -> Result<Self, String> {
let n_rows = u.nrows();
if u.ncols() != p * rank {
return Err(format!(
"RowMetric::from_factors: factor matrix has {} cols; expected p*rank = {}*{} = {}",
u.ncols(),
p,
rank,
p * rank
));
}
if !u.iter().all(|v| v.is_finite()) {
return Err("RowMetric::from_factors: factors must be finite".to_string());
}
let mut traces = ndarray::Array1::<f64>::zeros(n_rows);
let mut full = Array3::<f64>::zeros((1, p, p));
for row in 0..n_rows {
for i in 0..p {
for j in 0..p {
let mut acc = 0.0;
for k in 0..rank {
acc += u[[row, i * rank + k]] * u[[row, j * rank + k]];
}
full[[0, i, j]] = acc;
}
}
normalize_fisher_rao_blocks(full.view().into_dyn(), 1, p)
.map_err(|e| format!("RowMetric::from_factors: row {row}: {e}"))?;
let mut tr = 0.0_f64;
for i in 0..p {
tr += full[[0, i, i]];
}
traces[row] = tr;
}
Ok(Self {
provenance,
n_rows,
p,
rank,
factors: Some(u),
solver_delta,
traces,
fisher_factor_kind: match provenance {
MetricProvenance::OutputFisher { .. }
| MetricProvenance::OutputFisherDownstream { .. }
| MetricProvenance::BehavioralFisher { .. } => {
Some(FisherFactorKind::UncertifiedApproximation)
}
MetricProvenance::Euclidean | MetricProvenance::WhitenedStructured { .. } => None,
},
truncation_mass_residual: None,
})
}
pub fn with_fisher_factor_kind(mut self, kind: FisherFactorKind) -> Result<Self, String> {
if self.fisher_factor_kind.is_none() {
return Err(
"RowMetric::with_fisher_factor_kind requires an output-Fisher metric".to_string(),
);
}
match kind {
FisherFactorKind::ExactFull if self.truncation_mass_residual.is_some() => {
return Err(
"RowMetric::with_fisher_factor_kind ExactFull forbids an omitted-trace record"
.to_string(),
);
}
FisherFactorKind::CertifiedPsdLowerBound if self.truncation_mass_residual.is_none() => {
return Err(
"RowMetric::with_fisher_factor_kind CertifiedPsdLowerBound requires an exact omitted-trace record"
.to_string(),
);
}
FisherFactorKind::ExactFull
| FisherFactorKind::CertifiedPsdLowerBound
| FisherFactorKind::UncertifiedApproximation => {}
}
self.fisher_factor_kind = Some(kind);
Ok(self)
}
pub fn with_truncation_mass_residual(
mut self,
residual: Arc<ndarray::Array1<f64>>,
) -> Result<Self, String> {
if self.factors.is_none() {
return Err(
"RowMetric::with_truncation_mass_residual requires a factored metric".to_string(),
);
}
if residual.len() != self.n_rows {
return Err(format!(
"RowMetric::with_truncation_mass_residual requires {} rows; got {}",
self.n_rows,
residual.len()
));
}
for (row, &value) in residual.iter().enumerate() {
if !(value.is_finite() && value >= 0.0) {
return Err(format!(
"RowMetric::with_truncation_mass_residual row {row} must be finite and non-negative; got {value}"
));
}
}
self.truncation_mass_residual = Some(residual);
Ok(self)
}
pub fn gather_rows(&self, rows: &[usize]) -> Result<Self, String> {
for (pos, &r) in rows.iter().enumerate() {
if r >= self.n_rows {
return Err(format!(
"RowMetric::gather_rows: row index {r} at position {pos} is out of bounds \
(n_rows = {})",
self.n_rows
));
}
}
match self.factors.as_ref() {
None => Self::euclidean(rows.len(), self.p),
Some(factors) => {
let cols = self.p * self.rank;
let mut sub = Array2::<f64>::zeros((rows.len(), cols));
for (pos, &r) in rows.iter().enumerate() {
sub.row_mut(pos).assign(&factors.row(r));
}
let mut metric = Self::from_factors(
self.provenance,
Arc::new(sub),
self.p,
self.rank,
self.solver_delta,
)?;
metric.fisher_factor_kind = self.fisher_factor_kind;
match self.truncation_mass_residual.as_ref() {
None => Ok(metric),
Some(residual) => {
let gathered =
ndarray::Array1::from_iter(rows.iter().map(|&row| residual[row]));
metric.with_truncation_mass_residual(Arc::new(gathered))
}
}
}
}
}
pub fn provenance(&self) -> MetricProvenance {
self.provenance
}
pub fn fisher_factor_kind(&self) -> Option<FisherFactorKind> {
self.fisher_factor_kind
}
pub fn whitens_likelihood(&self) -> bool {
matches!(
self.provenance,
MetricProvenance::WhitenedStructured { .. } | MetricProvenance::BehavioralFisher { .. }
)
}
pub fn drives_gauge(&self) -> bool {
!matches!(self.provenance, MetricProvenance::Euclidean)
}
pub fn n_rows(&self) -> usize {
self.n_rows
}
pub fn p_out(&self) -> usize {
self.p
}
pub fn metric_rank(&self) -> usize {
self.rank
}
pub fn row_traces(&self) -> ndarray::ArrayView1<'_, f64> {
self.traces.view()
}
pub fn truncation_mass_residual(&self, row: usize) -> Option<f64> {
self.truncation_mass_residual
.as_ref()
.map(|residual| residual[row])
}
pub fn truncation_mass_residual_fraction(&self, row: usize) -> Option<f64> {
self.truncation_mass_residual(row).map(|residual| {
let total = self.traces[row] + residual;
if total > 0.0 { residual / total } else { 0.0 }
})
}
pub fn whiten_residual_row(&self, row: usize, r: ArrayView1<'_, f64>) -> Vec<f64> {
match &self.factors {
None => r.iter().copied().collect(),
Some(u) => {
let mut out = vec![0.0_f64; self.rank];
for k in 0..self.rank {
let mut acc = 0.0;
for i in 0..self.p {
acc += u[[row, i * self.rank + k]] * r[i];
}
out[k] = acc;
}
out
}
}
}
#[inline]
pub fn factor_entry(&self, row: usize, i: usize, k: usize) -> f64 {
match &self.factors {
None => {
if i == k {
1.0
} else {
0.0
}
}
Some(u) => u[[row, i * self.rank + k]],
}
}
pub fn apply_metric_row(&self, row: usize, x: ArrayView1<'_, f64>) -> Vec<f64> {
match &self.factors {
None => x.iter().copied().collect(),
Some(u) => {
let mut w = vec![0.0_f64; self.rank];
for k in 0..self.rank {
let mut acc = 0.0;
for i in 0..self.p {
acc += u[[row, i * self.rank + k]] * x[i];
}
w[k] = acc;
}
let mut out = vec![0.0_f64; self.p];
for i in 0..self.p {
let mut acc = 0.0;
for k in 0..self.rank {
acc += u[[row, i * self.rank + k]] * w[k];
}
out[i] = acc;
}
out
}
}
}
pub fn pullback(&self, row: usize, j_row: &[f64], d: usize) -> Array2<f64> {
match &self.factors {
None => {
let mut g = Array2::<f64>::zeros((d, d));
for a in 0..d {
for b in a..d {
let mut acc = 0.0;
for i in 0..self.p {
acc += j_row[i * d + a] * j_row[i * d + b];
}
g[[a, b]] = acc;
g[[b, a]] = acc;
}
}
g
}
Some(u) => {
let mut m = Array2::<f64>::zeros((self.rank, d));
for k in 0..self.rank {
for a in 0..d {
let mut acc = 0.0;
for i in 0..self.p {
acc += u[[row, i * self.rank + k]] * j_row[i * d + a];
}
m[[k, a]] = acc;
}
}
let mut g = Array2::<f64>::zeros((d, d));
for a in 0..d {
for b in a..d {
let mut acc = 0.0;
for k in 0..self.rank {
acc += m[[k, a]] * m[[k, b]];
}
g[[a, b]] = acc;
g[[b, a]] = acc;
}
}
g
}
}
}
#[inline]
pub fn quad_form(&self, row: usize, r: ArrayView1<'_, f64>) -> f64 {
match &self.factors {
None => r.iter().map(|&v| v * v).sum(),
Some(_) => self
.whiten_residual_row(row, r)
.iter()
.map(|&w| w * w)
.sum(),
}
}
#[inline]
pub fn fisher_mass(&self, row: usize, x: ArrayView1<'_, f64>) -> f64 {
self.quad_form(row, x)
}
pub fn to_weight_field(&self) -> crate::WeightField {
use crate::WeightField;
match &self.factors {
None => WeightField::Identity,
Some(u) => WeightField::Factored {
u: Arc::clone(u),
rank: self.rank,
p_out: self.p,
},
}
}
}
pub fn pack_probe_factors(probes: ndarray::ArrayView3<'_, f64>) -> Result<Array2<f64>, String> {
let (n_rows, p, s) = probes.dim();
if s == 0 {
return Err("pack_probe_factors: need at least one probe (s == 0)".to_string());
}
if !probes.iter().all(|v| v.is_finite()) {
return Err("pack_probe_factors: probe entries must be finite".to_string());
}
let mut u = Array2::<f64>::zeros((n_rows, p * s));
for n in 0..n_rows {
for i in 0..p {
for k in 0..s {
u[[n, i * s + k]] = probes[[n, i, k]];
}
}
}
Ok(u)
}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::array;
#[test]
fn euclidean_metric_has_correct_dimensions() {
let m = RowMetric::euclidean(5, 3).unwrap();
assert_eq!(m.n_rows(), 5);
assert_eq!(m.p_out(), 3);
assert_eq!(m.metric_rank(), 3);
}
#[test]
fn euclidean_metric_traces_equal_p() {
let p = 4_usize;
let m = RowMetric::euclidean(3, p).unwrap();
for tr in m.row_traces().iter() {
assert!((*tr - p as f64).abs() < 1e-14, "trace {tr} != p={p}");
}
}
#[test]
fn euclidean_provenance_is_euclidean() {
let m = RowMetric::euclidean(1, 2).unwrap();
assert_eq!(m.provenance(), MetricProvenance::Euclidean);
}
#[test]
fn euclidean_does_not_whiten_likelihood() {
let m = RowMetric::euclidean(1, 2).unwrap();
assert!(!m.whitens_likelihood());
}
#[test]
fn euclidean_does_not_drive_gauge() {
let m = RowMetric::euclidean(1, 2).unwrap();
assert!(!m.drives_gauge());
}
#[test]
fn euclidean_to_weight_field_is_identity() {
let m = RowMetric::euclidean(1, 2).unwrap();
assert!(matches!(m.to_weight_field(), WeightField::Identity));
}
#[test]
fn euclidean_whiten_residual_is_passthrough() {
let m = RowMetric::euclidean(1, 3).unwrap();
let r = array![1.0_f64, 2.0, 3.0];
let w = m.whiten_residual_row(0, r.view());
assert_eq!(w, vec![1.0, 2.0, 3.0]);
}
#[test]
fn euclidean_factor_entry_is_identity() {
let m = RowMetric::euclidean(1, 3).unwrap();
assert_eq!(m.factor_entry(0, 0, 0), 1.0);
assert_eq!(m.factor_entry(0, 1, 1), 1.0);
assert_eq!(m.factor_entry(0, 2, 2), 1.0);
assert_eq!(m.factor_entry(0, 0, 1), 0.0);
assert_eq!(m.factor_entry(0, 1, 0), 0.0);
}
#[test]
fn euclidean_quad_form_is_squared_norm() {
let m = RowMetric::euclidean(1, 3).unwrap();
let r = array![1.0_f64, 2.0, 2.0];
assert!((m.quad_form(0, r.view()) - 9.0).abs() < 1e-14);
}
#[test]
fn behavioral_fisher_whitens_likelihood_and_drives_gauge() {
let u = Arc::new(array![[1.0_f64, 0.5]]); let m = RowMetric::behavioral_fisher(u, 1, 2).unwrap();
assert!(m.whitens_likelihood());
assert!(m.drives_gauge());
assert_eq!(
m.provenance(),
MetricProvenance::BehavioralFisher { probes: 2 }
);
assert_eq!(m.metric_rank(), 2);
}
#[test]
fn behavioral_fisher_quad_form_is_probe_sum() {
let u = Arc::new(array![[1.0_f64, 0.0, 0.0, 2.0]]);
let m = RowMetric::behavioral_fisher(u, 2, 2).unwrap();
let e = array![3.0_f64, 1.0];
assert!((m.quad_form(0, e.view()) - 13.0).abs() < 1e-12);
}
#[test]
fn behavioral_fisher_g_identity_reproduces_euclidean_quad_form() {
let p = 3;
let mut u = Array2::<f64>::zeros((1, p * p));
for i in 0..p {
u[[0, i * p + i]] = 1.0;
}
let bf = RowMetric::behavioral_fisher(Arc::new(u), p, p).unwrap();
let euc = RowMetric::euclidean(1, p).unwrap();
let e = array![1.5_f64, -2.0, 0.25];
assert_eq!(bf.metric_rank(), euc.metric_rank());
assert!((bf.quad_form(0, e.view()) - euc.quad_form(0, e.view())).abs() < 1e-14);
assert_eq!(bf.whiten_residual_row(0, e.view()), vec![1.5, -2.0, 0.25]);
}
#[test]
fn pack_probe_factors_matches_manual_layout() {
use ndarray::Array3;
let mut probes = Array3::<f64>::zeros((1, 2, 2));
probes[[0, 0, 0]] = 1.0; probes[[0, 1, 0]] = 3.0; probes[[0, 0, 1]] = 2.0; probes[[0, 1, 1]] = 4.0; let u = pack_probe_factors(probes.view()).unwrap();
assert_eq!(u.as_slice().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
let m = RowMetric::behavioral_fisher(Arc::new(u), 2, 2).unwrap();
let e = array![1.0_f64, 0.0];
assert!((m.quad_form(0, e.view()) - 5.0).abs() < 1e-12);
}
#[test]
fn pack_probe_factors_rejects_zero_probes() {
use ndarray::Array3;
let probes = Array3::<f64>::zeros((2, 3, 0));
assert!(pack_probe_factors(probes.view()).is_err());
}
#[test]
fn fisher_factor_status_is_never_inferred_from_zero_residual_2249() {
let factors = Arc::new(Array2::from_elem((1, 1), 2.0));
let metric = RowMetric::output_fisher(factors, 1, 1)
.unwrap()
.with_truncation_mass_residual(Arc::new(array![0.0]))
.unwrap();
assert_eq!(
metric.fisher_factor_kind(),
Some(FisherFactorKind::UncertifiedApproximation)
);
let certified = metric
.clone()
.with_fisher_factor_kind(FisherFactorKind::CertifiedPsdLowerBound)
.unwrap();
assert_eq!(
certified.fisher_factor_kind(),
Some(FisherFactorKind::CertifiedPsdLowerBound)
);
assert!(
metric
.with_fisher_factor_kind(FisherFactorKind::ExactFull)
.is_err(),
"an omitted-trace record is incompatible with an exact-full claim"
);
}
#[test]
fn project_jac_with_identity_returns_jac() {
let u_row = [1.0_f64, 0.0, 0.0, 1.0]; let j_row = [1.0_f64, 2.0, 3.0, 4.0]; let m = WeightField::project_jac_row_with_u(&u_row, &j_row, 2, 2, 2);
assert!((m[[0, 0]] - 1.0).abs() < 1e-14);
assert!((m[[0, 1]] - 2.0).abs() < 1e-14);
assert!((m[[1, 0]] - 3.0).abs() < 1e-14);
assert!((m[[1, 1]] - 4.0).abs() < 1e-14);
}
#[test]
fn project_jac_with_zeros_returns_zero_matrix() {
let u_row = [0.0_f64, 0.0];
let j_row = [1.0_f64, 2.0];
let m = WeightField::project_jac_row_with_u(&u_row, &j_row, 2, 1, 1);
assert_eq!(m[[0, 0]], 0.0);
}
}