flow_gate_core/gate/
ellipsoid.rs

1use smallvec::SmallVec;
2
3use crate::error::FlowGateError;
4use crate::traits::{Gate, GateId, ParameterName};
5use crate::transform::TransformKind;
6
7#[derive(Debug, Clone)]
8pub struct EllipsoidDimension {
9    pub parameter: ParameterName,
10    pub transform: Option<TransformKind>,
11}
12
13#[derive(Debug, Clone)]
14pub struct EllipsoidCovariance {
15    metric: EllipsoidMetric,
16    matrix_full: Vec<f64>,
17    n: usize,
18}
19
20#[derive(Debug, Clone)]
21enum EllipsoidMetric {
22    /// Strict SPD mode using Cholesky factorization.
23    Cholesky { l: Vec<f64> },
24    /// Compatibility mode matching flowCore/flowutils behavior: quadratic form with inv(S).
25    GeneralInverse { inv: Vec<f64> },
26}
27
28impl EllipsoidCovariance {
29    pub fn from_upper_triangular(upper: &[f64], n: usize) -> Result<Self, FlowGateError> {
30        let expected = n * (n + 1) / 2;
31        if upper.len() != expected {
32            return Err(FlowGateError::DimensionMismatch(upper.len(), expected));
33        }
34
35        let mut s = vec![0.0_f64; n * n];
36        let mut idx = 0usize;
37        for i in 0..n {
38            for j in i..n {
39                let v = upper[idx];
40                idx += 1;
41                s[i * n + j] = v;
42                s[j * n + i] = v;
43            }
44        }
45
46        Self::from_symmetric_matrix(&s, n)
47    }
48
49    pub fn from_full_matrix(values: &[f64], n: usize) -> Result<Self, FlowGateError> {
50        let expected = n * n;
51        if values.len() != expected {
52            return Err(FlowGateError::DimensionMismatch(values.len(), expected));
53        }
54        let mut s = values.to_vec();
55        for row in 0..n {
56            for col in 0..n {
57                let a = s[row * n + col];
58                let b = s[col * n + row];
59                if !a.is_finite() || !b.is_finite() {
60                    return Err(FlowGateError::NotPositiveDefinite);
61                }
62                // Enforce symmetry with tolerance for minor rounding differences.
63                if (a - b).abs() > 1e-9 {
64                    return Err(FlowGateError::InvalidGate(
65                        "Ellipsoid covariance matrix is not symmetric".to_string(),
66                    ));
67                }
68                let avg = 0.5 * (a + b);
69                s[row * n + col] = avg;
70                s[col * n + row] = avg;
71            }
72        }
73        Self::from_symmetric_matrix(&s, n)
74    }
75
76    pub fn from_full_matrix_general(values: &[f64], n: usize) -> Result<Self, FlowGateError> {
77        let expected = n * n;
78        if values.len() != expected {
79            return Err(FlowGateError::DimensionMismatch(values.len(), expected));
80        }
81        if values.iter().any(|v| !v.is_finite()) {
82            return Err(FlowGateError::InvalidGate(
83                "Ellipsoid covariance matrix must contain only finite values".to_string(),
84            ));
85        }
86
87        let inv =
88            invert_square(values, n).map_err(|e| FlowGateError::InvalidGate(e.to_string()))?;
89        Ok(Self {
90            metric: EllipsoidMetric::GeneralInverse { inv },
91            matrix_full: values.to_vec(),
92            n,
93        })
94    }
95
96    fn from_symmetric_matrix(s: &[f64], n: usize) -> Result<Self, FlowGateError> {
97        let mut l = vec![0.0_f64; n * n];
98        let mut min_diag = f64::INFINITY;
99        let mut max_diag = 0.0_f64;
100        for i in 0..n {
101            for j in 0..=i {
102                let mut sum = s[i * n + j];
103                for k in 0..j {
104                    sum -= l[i * n + k] * l[j * n + k];
105                }
106                if i == j {
107                    if !sum.is_finite() || sum <= 0.0 {
108                        return Err(FlowGateError::NotPositiveDefinite);
109                    }
110                    let diag = sum.sqrt();
111                    if !diag.is_finite() || diag <= 0.0 {
112                        return Err(FlowGateError::NotPositiveDefinite);
113                    }
114                    min_diag = min_diag.min(diag);
115                    max_diag = max_diag.max(diag);
116                    l[i * n + j] = diag;
117                } else {
118                    let denom = l[j * n + j];
119                    if !denom.is_finite() || denom <= 0.0 {
120                        return Err(FlowGateError::NotPositiveDefinite);
121                    }
122                    l[i * n + j] = sum / denom;
123                }
124            }
125        }
126
127        if !min_diag.is_finite() || min_diag <= 0.0 || !max_diag.is_finite() {
128            return Err(FlowGateError::NotPositiveDefinite);
129        }
130        // Cholesky-based condition estimate: cond(S) ~= cond(L)^2.
131        // Reject highly ill-conditioned matrices as requested by spec tests.
132        let cond_estimate = (max_diag / min_diag).powi(2);
133        if !cond_estimate.is_finite() || cond_estimate > 1.0e10 {
134            return Err(FlowGateError::NotPositiveDefinite);
135        }
136
137        Ok(Self {
138            metric: EllipsoidMetric::Cholesky { l },
139            matrix_full: s.to_vec(),
140            n,
141        })
142    }
143
144    pub fn n(&self) -> usize {
145        self.n
146    }
147
148    pub fn uses_general_inverse(&self) -> bool {
149        matches!(self.metric, EllipsoidMetric::GeneralInverse { .. })
150    }
151
152    pub fn full_matrix(&self) -> &[f64] {
153        &self.matrix_full
154    }
155
156    pub fn to_upper_triangular(&self) -> Vec<f64> {
157        let n = self.n;
158        let s = &self.matrix_full;
159        let mut upper = Vec::with_capacity(n * (n + 1) / 2);
160        for row in 0..n {
161            for col in row..n {
162                upper.push(s[row * n + col]);
163            }
164        }
165        upper
166    }
167
168    pub fn mahalanobis_sq(&self, delta: &[f64]) -> f64 {
169        if delta.len() != self.n || delta.iter().any(|v| !v.is_finite()) {
170            return f64::NAN;
171        }
172
173        match &self.metric {
174            EllipsoidMetric::Cholesky { l } => {
175                let n = self.n;
176                let mut z = SmallVec::<[f64; 8]>::with_capacity(n);
177                for i in 0..n {
178                    z.push(0.0);
179                    let mut rhs = delta[i];
180                    for k in 0..i {
181                        rhs -= l[i * n + k] * z[k];
182                    }
183                    let diag = l[i * n + i];
184                    if !diag.is_finite() || diag <= 0.0 {
185                        return f64::NAN;
186                    }
187                    z[i] = rhs / diag;
188                }
189                z.iter().map(|v| v * v).sum()
190            }
191            EllipsoidMetric::GeneralInverse { inv } => {
192                let n = self.n;
193                let mut quad = 0.0_f64;
194                for row in 0..n {
195                    let mut proj = 0.0_f64;
196                    for col in 0..n {
197                        proj += inv[row * n + col] * delta[col];
198                    }
199                    quad += delta[row] * proj;
200                }
201                quad
202            }
203        }
204    }
205}
206
207#[derive(Debug, Clone)]
208pub struct EllipsoidGate {
209    id: GateId,
210    parent_id: Option<GateId>,
211    dimensions: Vec<EllipsoidDimension>,
212    dim_names: Vec<ParameterName>,
213    mean: Vec<f64>,
214    covariance: EllipsoidCovariance,
215    distance_sq: f64,
216}
217
218impl EllipsoidGate {
219    pub fn new(
220        id: GateId,
221        parent_id: Option<GateId>,
222        dimensions: Vec<EllipsoidDimension>,
223        mean: Vec<f64>,
224        covariance_upper: &[f64],
225        distance_sq: f64,
226    ) -> Result<Self, FlowGateError> {
227        let n = dimensions.len();
228        if n == 0 {
229            return Err(FlowGateError::InvalidGate(
230                "EllipsoidGate requires at least one dimension".to_string(),
231            ));
232        }
233        if mean.len() != n {
234            return Err(FlowGateError::DimensionMismatch(mean.len(), n));
235        }
236        if !distance_sq.is_finite() || distance_sq <= 0.0 {
237            return Err(FlowGateError::InvalidGate(
238                "EllipsoidGate distanceSquare must be finite and > 0".to_string(),
239            ));
240        }
241        let covariance = match covariance_upper.len() {
242            len if len == n * (n + 1) / 2 => {
243                EllipsoidCovariance::from_upper_triangular(covariance_upper, n)?
244            }
245            len if len == n * n => EllipsoidCovariance::from_full_matrix(covariance_upper, n)?,
246            len => return Err(FlowGateError::DimensionMismatch(len, n * (n + 1) / 2)),
247        };
248        let dim_names = dimensions.iter().map(|d| d.parameter.clone()).collect();
249        Ok(Self {
250            id,
251            parent_id,
252            dimensions,
253            dim_names,
254            mean,
255            covariance,
256            distance_sq,
257        })
258    }
259
260    /// Parser-targeted compatibility constructor for official Gating-ML corpus.
261    /// Accepts any finite, invertible full matrix and evaluates via inv(S).
262    pub fn new_general_covariance(
263        id: GateId,
264        parent_id: Option<GateId>,
265        dimensions: Vec<EllipsoidDimension>,
266        mean: Vec<f64>,
267        covariance_full: &[f64],
268        distance_sq: f64,
269    ) -> Result<Self, FlowGateError> {
270        let n = dimensions.len();
271        if n == 0 {
272            return Err(FlowGateError::InvalidGate(
273                "EllipsoidGate requires at least one dimension".to_string(),
274            ));
275        }
276        if mean.len() != n {
277            return Err(FlowGateError::DimensionMismatch(mean.len(), n));
278        }
279        if !distance_sq.is_finite() || distance_sq <= 0.0 {
280            return Err(FlowGateError::InvalidGate(
281                "EllipsoidGate distanceSquare must be finite and > 0".to_string(),
282            ));
283        }
284        let covariance = EllipsoidCovariance::from_full_matrix_general(covariance_full, n)?;
285        let dim_names = dimensions.iter().map(|d| d.parameter.clone()).collect();
286        Ok(Self {
287            id,
288            parent_id,
289            dimensions,
290            dim_names,
291            mean,
292            covariance,
293            distance_sq,
294        })
295    }
296
297    pub fn dimensions_def(&self) -> &[EllipsoidDimension] {
298        &self.dimensions
299    }
300
301    pub fn mean(&self) -> &[f64] {
302        &self.mean
303    }
304
305    pub fn covariance(&self) -> &EllipsoidCovariance {
306        &self.covariance
307    }
308
309    pub fn distance_sq(&self) -> f64 {
310        self.distance_sq
311    }
312}
313
314impl Gate for EllipsoidGate {
315    fn dimensions(&self) -> &[ParameterName] {
316        &self.dim_names
317    }
318
319    fn contains(&self, coords: &[f64]) -> bool {
320        let n = self.dimensions.len();
321        if coords.len() != n {
322            return false;
323        }
324        if coords.iter().any(|c| !c.is_finite() || c.abs() > 1e15) {
325            return false;
326        }
327
328        let mut delta = SmallVec::<[f64; 8]>::with_capacity(n);
329        for (i, &coord) in coords.iter().enumerate() {
330            delta.push(coord - self.mean[i]);
331        }
332        let d_sq = self.covariance.mahalanobis_sq(&delta);
333        d_sq.is_finite() && d_sq <= self.distance_sq
334    }
335
336    fn gate_id(&self) -> &GateId {
337        &self.id
338    }
339
340    fn parent_id(&self) -> Option<&GateId> {
341        self.parent_id.as_ref()
342    }
343}
344
345fn invert_square(values: &[f64], n: usize) -> Result<Vec<f64>, &'static str> {
346    if values.len() != n * n {
347        return Err("matrix size mismatch");
348    }
349
350    let mut a = values.to_vec();
351    let mut inv = vec![0.0_f64; n * n];
352    for i in 0..n {
353        inv[i * n + i] = 1.0;
354    }
355
356    for col in 0..n {
357        let mut pivot_row = col;
358        let mut best = a[col * n + col].abs();
359        for row in (col + 1)..n {
360            let cand = a[row * n + col].abs();
361            if cand > best {
362                best = cand;
363                pivot_row = row;
364            }
365        }
366
367        if !best.is_finite() || best <= 1.0e-14 {
368            return Err("Ellipsoid covariance matrix is singular");
369        }
370
371        if pivot_row != col {
372            for k in 0..n {
373                a.swap(col * n + k, pivot_row * n + k);
374                inv.swap(col * n + k, pivot_row * n + k);
375            }
376        }
377
378        let pivot = a[col * n + col];
379        for k in 0..n {
380            a[col * n + k] /= pivot;
381            inv[col * n + k] /= pivot;
382        }
383
384        for row in 0..n {
385            if row == col {
386                continue;
387            }
388            let factor = a[row * n + col];
389            if factor == 0.0 {
390                continue;
391            }
392            for k in 0..n {
393                a[row * n + k] -= factor * a[col * n + k];
394                inv[row * n + k] -= factor * inv[col * n + k];
395            }
396        }
397    }
398
399    if inv.iter().any(|v| !v.is_finite()) {
400        return Err("Ellipsoid covariance inverse contains non-finite values");
401    }
402    Ok(inv)
403}
flow_gate_core/gate/ellipsoid.rs

flow_gate_core/gate/
ellipsoid.rs
