gam-problem 0.3.156

Neutral solver/criterion contract types for the gam penalized-likelihood engine
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
//! Typed structured constraint carriers for large factored coefficient blocks.
//!
//! The dense [`LinearInequalityConstraints`] system stores every row
//! explicitly, which is exact and fine for the small monotone blocks (a
//! `p × p` identity cone). A Khatri-Rao tensor block is different: the
//! monotonicity cone of a conditional transformation `h(y|x) = Σ_k α_k(x)
//! v_k(y)` is `α_k(x_i) ≥ 0` for every observation row `i` and every shape
//! column `k` — `n · p_shape` rows over `p_resp · p_cov` coefficients whose
//! dense materialization is gigabytes (gam#2306), while every operation an
//! active-set method actually performs factors through the covariate design
//! `Ψ` (`n × p_cov`):
//!
//! * constraint values are the columns of `Γ = Ψ Aᵀ` (one `n × p_cov` GEMM
//!   per shape column),
//! * a single row is `(e_k ⊗ ψ_i)ᵀ` — gathered densely only for the (small)
//!   active set,
//! * row norms are `‖ψ_i‖`, shared by every shape column.
//!
//! [`ConstraintSet`] is the closed union the solver plumbing carries: the
//! dense system verbatim, or the factored cone. Semantics are IDENTICAL to
//! canonicalizing the equivalent dense system: every slack / violation is
//! measured on unit-normalized rows, so tolerances stay geometric.

use crate::linear_constraints::LinearInequalityConstraints;
use ndarray::{Array1, Array2, ArrayView1};
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use std::sync::Arc;

/// Primal-feasibility tolerance of the inequality-constrained active-set Newton
/// solver, measured in the unit-normalized row metric this module defines:
/// a point `β` is feasible for a [`ConstraintSet`] exactly when
///
/// ```text
/// max_r  (b_r − a_r·β) / ‖a_r‖  ≤  PRIMAL_FEASIBILITY_TOL
/// ```
///
/// over the non-vacuous rows — the quantity [`ConstraintSet::max_scaled_violation`]
/// returns. This is the ONE definition of "feasible" in the codebase; the solver
/// certifies its returned iterate against it, every entry gate admits against it,
/// and `ConstraintSet::max_contract_feasible_step` sizes steps against it.
///
/// It lives beside the metric rather than in the solver because the two are the
/// same statement: the metric says what is measured, this says at what resolution.
/// `gam_solve::active_set` re-exports it as `ACTIVE_SET_PRIMAL_FEASIBILITY_TOL`.
///
/// Any consumer that re-derives a RAW (un-scaled) feasibility tolerance from a
/// returned iterate must scale this value by the per-row normalization the
/// constraint builder applied; demanding tighter feasibility than this is
/// inconsistent with the solver contract and will spuriously reject valid
/// boundary solutions (gam#2719: a step rule that demanded exact feasibility
/// refused 314 steps that violated nothing at this tolerance).
pub const PRIMAL_FEASIBILITY_TOL: f64 = 1e-8;

/// Can this row's feasibility be DECIDED by comparison at all?
///
/// EVERY feasibility rule in this module decides with an ordering predicate on
/// per-row quantities — `slack < −tol`, `drift ≥ 0`, `t < step`,
/// `violation > worst` — and EVERY one of those is `false` for `NaN`. A row
/// carrying a `NaN` therefore contributes NOTHING to any of those minima and
/// maxima, and the rule answers with its neutral element: "take the whole
/// step", "nothing is violated". That is the exact opposite of the truth, and
/// it is gam#2721: a step with a `NaN` component was certified at `α = 1.0`,
/// and its caller rejects only `!α.is_finite() || α ≤ 0.0`, neither of which
/// `1.0` is.
///
/// The quantities are therefore tested BEFORE they are compared, and a row that
/// cannot be decided is refused by name rather than skipped. The predicate is
/// exported — rather than re-written at each site — because the defect WAS the
/// rule existing in several copies and being repaired in one of them: the two
/// fraction-to-boundary rules here, the violation sweep here, the saddle-escape
/// chord truncation in `gam-custom-family`, and the Bernoulli marginal-slope
/// segment cap in `gam-models` all decide with the same comparisons.
///
/// `NaN` is the value that cannot be compared, but it is not the only value
/// that must be refused. An infinite drift passes `drift ≥ 0` and an infinite
/// iterate value drives the violation to `−∞`, both of which read as "this row
/// does not object" for an argument that is not a point. And the carrier's own
/// constructor already holds the same line —
/// `LinearInequalityConstraints::new` rejects a non-finite `A` or `b` with the
/// identical reason — so requiring finiteness here keeps the row descriptors
/// and the per-iterate quantities under ONE rule rather than two.
///
/// A `NaN` `row_norm` additionally defeats the `norm <= 0.0` vacuity test that
/// would otherwise be the branch to catch it, which is why the norm is checked
/// here and not left to that branch.
///
/// This does NOT collide with the legitimately-vacuous row: `‖a‖ = 0` with a
/// bound at or below zero is finite, passes here, and keeps its own
/// disposition in each rule.
/// Where one sweep reads a row's scaling denominator and right-hand side.
///
/// `row_norm(row)?` and `bound(row)?` are `Result`-returning methods that
/// re-derive the row's carrier, slot and index on every call. The sweep calls
/// them once per row, and on the large-scale CTN cone (1.6 M rows) they were
/// 34 % of the profile — more than the products they scale. The factored cone
/// answers both from slices instead; every other carrier keeps its accessors,
/// so there is still exactly one scan and one decidability contract (gam#979).
enum RowMetrics<'a> {
    /// The factored cone: `norms[row % tile]` with `bounds` indexed by row,
    /// `0` when the cone is homogeneous.
    Tiled {
        norms: &'a [f64],
        tile: usize,
        bounds: Option<&'a [f64]>,
    },
    /// Anything else: through the carrier's own accessors.
    Carrier(&'a ConstraintSet),
}

impl RowMetrics<'_> {
    #[inline]
    fn read(&self, row: usize) -> Result<(f64, f64), String> {
        match self {
            RowMetrics::Tiled {
                norms,
                tile,
                bounds,
            } => {
                let norm = norms[row % tile];
                let bound = bounds.map_or(0.0, |values| values[row]);
                Ok((norm, bound))
            }
            RowMetrics::Carrier(set) => Ok((set.row_norm(row)?, set.bound(row)?)),
        }
    }
}

/// Why a [`ConstraintSet::max_scaled_violation`] sweep stops at a row.
///
/// The serial loop this replaces returned at the first such row in index
/// order; the parallel sweep carries the smallest index instead, so the
/// verdict does not depend on the row split.
#[derive(Clone, Debug)]
enum SweepTerminal {
    /// The row's own norm or bound could not be read; carries the carrier's
    /// own refusal so the sweep does not swallow it.
    RowUnavailable(String),
    /// `gam#2721`: a non-finite norm, bound or value. Feasibility of an
    /// iterate that is not a number is undefined, and every comparison in the
    /// sweep is false for `NaN`, so the row cannot be skipped.
    Undecidable { norm: f64, bound: f64, value: f64 },
    /// `0ᵀβ ≥ b` with `b > 0`: unsatisfiable by any `β`.
    VacuousRowWithPositiveBound,
}

/// Running state of the row sweep: the first terminal row in index order, and
/// the largest scaled violation with the smallest row that attains it.
struct ScaledViolationSweep {
    terminal: Option<(usize, SweepTerminal)>,
    worst: f64,
    worst_row: Option<usize>,
}

impl ScaledViolationSweep {
    fn none() -> Self {
        Self {
            terminal: None,
            worst: 0.0,
            worst_row: None,
        }
    }

    fn record_terminal(&mut self, row: usize, terminal: SweepTerminal) {
        let keep = match self.terminal {
            Some((seen, _)) => row < seen,
            None => true,
        };
        if keep {
            self.terminal = Some((row, terminal));
        }
    }

    fn record_violation(&mut self, row: usize, violation: f64) {
        // `>` alone reproduces the serial loop's smallest-index tie-break only
        // while the rows arrive in order; the merge below restores it across
        // chunks.
        if violation > self.worst {
            self.worst = violation;
            self.worst_row = Some(row);
        }
    }

    fn merge(mut self, other: Self) -> Self {
        if let Some((row, terminal)) = other.terminal {
            self.record_terminal(row, terminal);
        }
        let take_other = match (other.worst > self.worst, other.worst == self.worst) {
            (true, _) => true,
            (false, true) => match (other.worst_row, self.worst_row) {
                (Some(candidate), Some(held)) => candidate < held,
                (Some(_), None) => true,
                _ => false,
            },
            _ => false,
        };
        if take_other {
            self.worst = other.worst;
            self.worst_row = other.worst_row;
        }
        self
    }

    fn verdict(self) -> Result<(f64, Option<usize>), String> {
        match self.terminal {
            Some((row, SweepTerminal::RowUnavailable(error))) => Err(format!(
                "ConstraintSet::max_scaled_violation: row {row} has no readable norm or \
                 bound: {error}"
            )),
            Some((
                row,
                SweepTerminal::Undecidable {
                    norm,
                    bound,
                    value,
                },
            )) => Err(format!(
                "ConstraintSet::max_scaled_violation: row {row} cannot be decided \
                 (row norm {norm:.3e}, bound {bound:.3e}, value {value:.3e}); \
                 feasibility of a non-finite iterate is undefined and every \
                 comparison in the sweep is false for NaN, so the row cannot \
                 be skipped (gam#2721)"
            )),
            Some((row, SweepTerminal::VacuousRowWithPositiveBound)) => {
                Ok((f64::INFINITY, Some(row)))
            }
            None => Ok((self.worst, self.worst_row)),
        }
    }
}

pub fn feasibility_quantities_are_finite(quantities: &[f64]) -> bool {
    quantities.iter().all(|q| q.is_finite())
}

/// The contract-feasible ratio test itself, over already-evaluated constraint
/// values, so every carrier — the dense system, the factored cone, the
/// block-diagonal composition — runs the SAME arithmetic without any of them
/// having to be materialized as another.
///
/// `values[r]` is `a_r·β`, `directional[r]` is `a_r·δ` (the constraint
/// functional is linear, so its value at `δ` IS the directional derivative).
/// The rule is documented on [`ConstraintSet::max_contract_feasible_step`].
pub(crate) fn contract_feasible_step_over_rows<B, N>(
    values: &Array1<f64>,
    directional: &Array1<f64>,
    bound: B,
    row_norm: N,
) -> Result<ContractFeasibleStep, ContractFeasibleStepError>
where
    B: Fn(usize) -> Result<f64, String>,
    N: Fn(usize) -> Result<f64, String>,
{
    let tol = PRIMAL_FEASIBILITY_TOL;
    let mut limit = ContractFeasibleStep::UNLIMITED;
    for row in 0..values.len() {
        let norm = row_norm(row).map_err(ContractFeasibleStepError::Carrier)?;
        let bound = bound(row).map_err(ContractFeasibleStepError::Carrier)?;
        // A ROW THAT CANNOT BE COMPARED IS NOT A FEASIBLE ROW (gam#2721) — see
        // [`feasibility_quantities_are_finite`] for why skipping it certifies
        // the whole step. Refuse instead, and name the condition: "this row is
        // not a number" is a different condition from "the current iterate
        // violates this row", so it gets its own variant rather than being
        // reported through `InfeasibleIterate`.
        if !feasibility_quantities_are_finite(&[norm, bound, values[row], directional[row]]) {
            return Err(ContractFeasibleStepError::NonFinite {
                row,
                scaled_slack: (values[row] - bound) / norm,
                scaled_drift: directional[row] / norm,
            });
        }
        if !(norm.is_finite() && norm > 0.0) {
            // A vacuous row constrains nothing unless its bound is positive,
            // in which case the feasible set is empty and no step fraction
            // exists. Same disposition as the solver's own violation scan.
            if bound > 0.0 {
                return Err(ContractFeasibleStepError::InfeasibleIterate {
                    row,
                    scaled_slack: f64::NEG_INFINITY,
                });
            }
            continue;
        }
        let slack = (values[row] - bound) / norm;
        let drift = directional[row] / norm;
        if slack < -tol {
            return Err(ContractFeasibleStepError::InfeasibleIterate {
                row,
                scaled_slack: slack,
            });
        }
        if drift >= 0.0 {
            continue;
        }
        if slack + drift >= -tol {
            // The endpoint of the FULL step is feasible on this row to the
            // contract. Nothing to limit.
            continue;
        }
        // `slack ≥ −tol` and `slack + drift < −tol` give `slack < −drift`, so
        // this ratio is strictly below 1 and non-negative.
        let fraction = (slack.max(0.0) / -drift).clamp(0.0, 1.0);
        if fraction < limit.fraction {
            limit = ContractFeasibleStep {
                fraction,
                blocking_row: Some(row),
                blocking_scaled_slack: slack,
                blocking_scaled_drift: drift,
            };
        }
    }
    Ok(limit)
}

/// Result of the contract-feasible ratio test
/// (`ConstraintSet::max_contract_feasible_step`).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct ContractFeasibleStep {
    /// Largest fraction in `[0, 1]` such that `β + fraction·δ` is feasible at
    /// [`PRIMAL_FEASIBILITY_TOL`]. `1.0` means no row limits the step.
    ///
    /// `0.0` is a legitimate, non-exceptional answer: it says a row is active
    /// at `β` and `δ` points strictly out of it by more than round-off, so no
    /// positive multiple of `δ` is admissible. The remedy is a projection onto
    /// the active face, not a smaller `δ` — the ratio test is invariant under
    /// `δ ↦ cδ` once the numerator is zero, so shrinking a trust radius against
    /// it cannot converge (gam#2719).
    pub fraction: f64,
    /// Row that limited `fraction`, if any.
    pub blocking_row: Option<usize>,
    /// Scaled slack `(a·β − b)/‖a‖` of `blocking_row` at `β`.
    pub blocking_scaled_slack: f64,
    /// Scaled drift `(a·δ)/‖a‖` of `blocking_row` (strictly negative when a
    /// row blocks).
    pub blocking_scaled_drift: f64,
}

impl ContractFeasibleStep {
    /// The unlimited answer: the whole direction is admissible.
    pub const UNLIMITED: Self = Self {
        fraction: 1.0,
        blocking_row: None,
        blocking_scaled_slack: f64::INFINITY,
        blocking_scaled_drift: 0.0,
    };

}

/// Why the contract-feasible ratio test could not answer.
///
/// Every variant is a violated PRECONDITION of the ratio test, never a small
/// step: "no admissible step exists" is reported as
/// [`ContractFeasibleStep::fraction`] `== 0.0`, not as an error.
#[derive(Clone, Debug, PartialEq)]
pub enum ContractFeasibleStepError {
    /// `beta` / `direction` widths disagree with the constraint carrier.
    Dimension {
        beta: usize,
        direction: usize,
        expected: usize,
    },
    /// The CURRENT iterate violates a row by more than
    /// [`PRIMAL_FEASIBILITY_TOL`], so the ratio test has no feasible origin to
    /// step from. This is the genuine "infeasible iterate" condition and stays
    /// loud.
    InfeasibleIterate { row: usize, scaled_slack: f64 },
    /// A row cannot be DECIDED by comparison: a non-finite row norm, bound,
    /// `a·β` or `a·δ`. Reported rather than skipped: every comparison in the
    /// rule is false for NaN, so a skipped row would silently certify a step
    /// that is not a number as fully feasible (gam#2721). The reported slack
    /// and drift are the scaled quantities as computed, so the offending one is
    /// visible.
    NonFinite {
        row: usize,
        scaled_slack: f64,
        scaled_drift: f64,
    },
    /// The carrier could not evaluate `Aβ` / `Aδ` or a row descriptor.
    Carrier(String),
}

impl std::fmt::Display for ContractFeasibleStepError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ContractFeasibleStepError::Dimension {
                beta,
                direction,
                expected,
            } => write!(
                f,
                "constraint step dimension mismatch: beta={beta}, direction={direction}, constraints={expected}"
            ),
            ContractFeasibleStepError::InfeasibleIterate { row, scaled_slack } => write!(
                f,
                "current iterate violates constraint row {row}: scaled slack={scaled_slack:.3e} \
                 below the primal-feasibility contract {PRIMAL_FEASIBILITY_TOL:.3e}"
            ),
            ContractFeasibleStepError::NonFinite {
                row,
                scaled_slack,
                scaled_drift,
            } => write!(
                f,
                "constraint row {row} has a non-finite ratio test: scaled slack={scaled_slack:.3e}, \
                 scaled drift={scaled_drift:.3e}"
            ),
            ContractFeasibleStepError::Carrier(reason) => write!(f, "{reason}"),
        }
    }
}

/// Nonnegativity cone `(e_k ⊗ ψ_i)ᵀ β ≥ 0` for a row-major Khatri-Rao block.
///
/// The coefficient block is `β = vec(A)` with `A` reshaped row-major as
/// `p_left × p_cov` (coefficient `A[k, j] = β[k · p_cov + j]`). The cone
/// constrains the factored linear functionals `α_k(x_i) = ψ_iᵀ A_{k,:}` to be
/// non-negative for every observation row `i` of `factor` and every
/// `k ∈ coupled_rows`.
///
/// Row identifiers are stable and dense: row `r = s · n + i` where `s` indexes
/// into `coupled_rows` and `i` is the observation row. Active-set warm starts
/// therefore survive across iterations exactly as with the dense system.
#[derive(Clone, Debug)]
pub struct KhatriRaoConeConstraints {
    /// Covariate factor `Ψ` (`n × p_cov`).
    factor: Arc<Array2<f64>>,
    /// Euclidean norm of each `Ψ` row (unit-normalization denominators).
    factor_row_norms: Array1<f64>,
    /// Coefficient rows of `A` (indices into `0..p_left`) that carry the cone.
    coupled_rows: Vec<usize>,
    /// Total number of coefficient rows in the block reshape.
    p_left: usize,
    /// Per-row right-hand sides. The homogeneous cone has `b ≡ 0`; a
    /// delta-coordinate solve (`β = β₀ + δ`) shifts them to `−(rowᵀβ₀)`.
    /// Bounds are `O(nrows)` — cheap even when the matrix is not.
    bounds: Option<Array1<f64>>,
}

impl KhatriRaoConeConstraints {
    pub fn new(
        factor: Arc<Array2<f64>>,
        coupled_rows: Vec<usize>,
        p_left: usize,
    ) -> Result<Self, String> {
        if factor.nrows() == 0 || factor.ncols() == 0 {
            return Err("KhatriRaoConeConstraints: factor must be non-empty".to_string());
        }
        if factor.iter().any(|v| !v.is_finite()) {
            return Err("KhatriRaoConeConstraints: factor must be finite".to_string());
        }
        if coupled_rows.is_empty() {
            return Err(
                "KhatriRaoConeConstraints: at least one coupled coefficient row is required"
                    .to_string(),
            );
        }
        let mut seen = vec![false; p_left];
        for &k in &coupled_rows {
            if k >= p_left {
                return Err(format!(
                    "KhatriRaoConeConstraints: coupled row {k} out of range (p_left = {p_left})"
                ));
            }
            if seen[k] {
                return Err(format!(
                    "KhatriRaoConeConstraints: coupled row {k} is duplicated"
                ));
            }
            seen[k] = true;
        }
        let factor_row_norms =
            Array1::from_iter(factor.rows().into_iter().map(|row| row.dot(&row).sqrt()));
        Ok(Self {
            factor,
            factor_row_norms,
            coupled_rows,
            p_left,
            bounds: None,
        })
    }

    pub fn factor(&self) -> &Array2<f64> {
        self.factor.as_ref()
    }

    pub fn coupled_rows(&self) -> &[usize] {
        &self.coupled_rows
    }

    pub fn p_left(&self) -> usize {
        self.p_left
    }

    /// One coupled response-row slot as a standalone cone over a single
    /// `p_cov` coefficient block. The covariate factor remains shared by
    /// [`Arc`]; only the small row-norm vector and this slot's optional bounds
    /// are copied. This is the exact block decomposition of an identity-Hessian
    /// projection, not a reduced-data approximation.
    pub fn single_coupled_slot(&self, slot: usize) -> Result<Self, String> {
        if slot >= self.coupled_rows.len() {
            return Err(format!(
                "KhatriRaoConeConstraints: coupled slot {slot} out of range ({} slots)",
                self.coupled_rows.len()
            ));
        }
        let n = self.factor.nrows();
        let bounds = self
            .bounds
            .as_ref()
            .map(|all| all.slice(ndarray::s![slot * n..(slot + 1) * n]).to_owned());
        Ok(Self {
            factor: Arc::clone(&self.factor),
            factor_row_norms: self.factor_row_norms.clone(),
            coupled_rows: vec![0],
            p_left: 1,
            bounds,
        })
    }

    pub fn nrows(&self) -> usize {
        self.coupled_rows.len() * self.factor.nrows()
    }

    pub fn ncols(&self) -> usize {
        self.p_left * self.factor.ncols()
    }

    /// Decompose a row id into `(coupled-row slot, observation row)`.
    #[inline]
    fn split_row_id(&self, row: usize) -> Result<(usize, usize), String> {
        let n = self.factor.nrows();
        let slot = row / n;
        if slot >= self.coupled_rows.len() {
            return Err(format!(
                "KhatriRaoConeConstraints: row id {row} out of range ({} rows)",
                self.nrows()
            ));
        }
        Ok((slot, row % n))
    }

    /// Raw (un-normalized) constraint values `A β` for the full row set,
    /// laid out slot-major (`r = s·n + i`).
    ///
    /// Cost: one `n × p_cov · p_cov` product per coupled row — never the
    /// `nrows × ncols` dense system.
    pub fn values(&self, beta: ArrayView1<'_, f64>) -> Result<Array1<f64>, String> {
        let p_cov = self.factor.ncols();
        if beta.len() != self.ncols() {
            return Err(format!(
                "KhatriRaoConeConstraints: beta length {} != {}",
                beta.len(),
                self.ncols()
            ));
        }
        let n = self.factor.nrows();
        let slots = self.coupled_rows.len();
        // One `Ψ · B` for every slot at once. `Array2::dot(&Array1)` is
        // ndarray's per-row `unrolled_dot`; `Array2::dot(&Array2)` is a real
        // matrix product, and this runs on every feasibility sweep of every
        // trial point at `n = 320000` (gam#979).
        let mut blocks = Array2::<f64>::zeros((p_cov, slots));
        for (slot, &k) in self.coupled_rows.iter().enumerate() {
            blocks
                .column_mut(slot)
                .assign(&beta.slice(ndarray::s![k * p_cov..(k + 1) * p_cov]));
        }
        let alpha = self.factor.dot(&blocks);
        let mut out = Array1::<f64>::zeros(self.nrows());
        for slot in 0..slots {
            out.slice_mut(ndarray::s![slot * n..(slot + 1) * n])
                .assign(&alpha.column(slot));
        }
        Ok(out)
    }

    /// Unit-normalization denominator of one row (`‖ψ_i‖`, shared across
    /// coupled slots). Zero rows are vacuous (`0ᵀβ ≥ 0` always holds) exactly
    /// like the canonicalized dense system keeps them inert.
    pub fn row_norm(&self, row: usize) -> Result<f64, String> {
        let (_, i) = self.split_row_id(row)?;
        Ok(self.factor_row_norms[i])
    }

    /// The coefficient columns row `row` acts on, ascending.
    ///
    /// Row `(slot, i)` has normal `e_k ⊗ ψ_i` with `k = coupled_rows[slot]`, and
    /// [`Self::values`] reads exactly the block `β[k·p_cov .. (k+1)·p_cov]`, so
    /// the support is `k·p_cov + j` over the columns `j` where `ψ_{i,j} ≠ 0`.
    /// Every other coefficient has a structurally zero coefficient in this row.
    pub fn row_column_support(&self, row: usize) -> Result<Vec<usize>, String> {
        let (slot, i) = self.split_row_id(row)?;
        let p_cov = self.factor.ncols();
        let base = self.coupled_rows[slot] * p_cov;
        Ok((0..p_cov)
            .filter(|&j| self.factor[[i, j]] != 0.0)
            .map(|j| base + j)
            .collect())
    }

    /// Per-row right-hand side (`0` for the homogeneous cone, shifted values
    /// after [`ConstraintSet::shifted_to_delta`]).
    pub fn bound(&self, row: usize) -> Result<f64, String> {
        self.split_row_id(row)?;
        Ok(self.bounds.as_ref().map_or(0.0, |bounds| bounds[row]))
    }

    /// The `Ψ`-row norms, one per `i`, shared across every coupled slot.
    ///
    /// Read directly by [`RowMetrics::Tiled`]: `row_norm(row)?` is the same
    /// lookup behind a slot split, a range check and a `Result`, and the
    /// feasibility sweep does it once per row (gam#979).
    pub(crate) fn row_norms_slice(&self) -> &[f64] {
        self.factor_row_norms
            .as_slice()
            .expect("factor row norms are contiguous")
    }

    /// Rows per coupled slot, so `row % tile_rows()` is the `Ψ` row index.
    pub(crate) fn tile_rows(&self) -> usize {
        self.factor.nrows()
    }

    /// The per-row right-hand sides, `None` for the homogeneous cone.
    pub(crate) fn bounds_slice(&self) -> Option<&[f64]> {
        self.bounds
            .as_ref()
            .map(|bounds| bounds.as_slice().expect("bounds are contiguous"))
    }

    /// Materialize the requested rows as a dense system (active-set KKT use;
    /// the id order of `rows` is preserved). Rows come out RAW (un-normalized),
    /// matching the raw dense construction path; callers that need geometric
    /// tolerances canonicalize the gathered system.
    pub fn gather_rows(&self, rows: &[usize]) -> Result<LinearInequalityConstraints, String> {
        let p_cov = self.factor.ncols();
        let mut a = Array2::<f64>::zeros((rows.len(), self.ncols()));
        let mut b = Array1::<f64>::zeros(rows.len());
        for (out_row, &row) in rows.iter().enumerate() {
            let (slot, i) = self.split_row_id(row)?;
            let k = self.coupled_rows[slot];
            a.row_mut(out_row)
                .slice_mut(ndarray::s![k * p_cov..(k + 1) * p_cov])
                .assign(&self.factor.row(i));
            b[out_row] = self.bound(row)?;
        }
        LinearInequalityConstraints::new(a, b)
    }

    /// Exact dense equivalent of the ENTIRE cone. Test/oracle use only — this
    /// is the materialization the carrier exists to avoid.
    pub fn to_dense(&self) -> Result<LinearInequalityConstraints, String> {
        let all: Vec<usize> = (0..self.nrows()).collect();
        self.gather_rows(&all)
    }
}

/// A row index in a [`ConstraintSet`]'s OWN constraint-row space — the space
/// addressed by [`ConstraintSet::values`], [`ConstraintSet::bound`] and
/// [`ConstraintSet::row_norm`], i.e. `0..nrows()`.
///
/// This is NOT a coefficient (β) index. The two spaces have different sizes
/// (`nrows()` vs `ncols()`) and different meanings, and they coincide only in
/// the special case of a square carrier whose row `r` is exactly the box
/// `β_r ≥ 0`. A block-diagonal composition breaks that coincidence: its row ids
/// are the CONCATENATION of the member row counts while its columns are the
/// concatenation of the member column ranges, so as soon as one member has
/// `nrows() < ncols()` (a monotone sub-basis alongside unconstrained intercept /
/// covariate columns) row id `r` of a later block names a β coordinate owned by
/// an EARLIER block. The newtype exists so that mistake cannot be made silently;
/// to go from a row to the coefficients it acts on, call
/// `ConstraintSet::row_column_support`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ConstraintRowId(pub usize);

impl ConstraintRowId {
    /// The raw index, for addressing a `values()` / `bound()` / `row_norm()`
    /// result. Deliberately explicit: reach for this only when indexing
    /// something that really is in constraint-row space.
    #[inline]
    pub fn index(self) -> usize {
        self.0
    }
}

/// One block of a [`ConstraintSet::BlockDiagonal`] composition: an inner set
/// acting on the coefficient columns `[col_start, col_start + set.ncols())` of
/// the joint vector.
#[derive(Clone, Debug)]
pub struct PlacedConstraintBlock {
    pub col_start: usize,
    pub set: ConstraintSet,
}

/// Closed union of the constraint carriers the blockwise solvers accept.
#[derive(Clone, Debug)]
pub enum ConstraintSet {
    /// Explicit rows, exactly as today.
    Dense(LinearInequalityConstraints),
    /// Factored Khatri-Rao nonnegativity cone.
    KhatriRaoCone(KhatriRaoConeConstraints),
    /// Block-diagonal composition over disjoint column ranges of a joint
    /// coefficient vector (the multi-block joint-Newton assembly). Row ids
    /// are the concatenation of the member row ids in order.
    BlockDiagonal {
        blocks: Vec<PlacedConstraintBlock>,
        total_cols: usize,
    },
}

impl ConstraintSet {
    /// Validated block-diagonal composition: member column ranges must lie
    /// inside the joint width and must not overlap.
    pub fn block_diagonal(
        blocks: Vec<PlacedConstraintBlock>,
        total_cols: usize,
    ) -> Result<Self, String> {
        let mut ranges: Vec<(usize, usize)> = Vec::with_capacity(blocks.len());
        for block in &blocks {
            let end = block.col_start + block.set.ncols();
            if end > total_cols {
                return Err(format!(
                    "ConstraintSet::block_diagonal: block columns {}..{} exceed joint width {}",
                    block.col_start, end, total_cols
                ));
            }
            ranges.push((block.col_start, end));
        }
        ranges.sort_unstable();
        for pair in ranges.windows(2) {
            if pair[1].0 < pair[0].1 {
                return Err(format!(
                    "ConstraintSet::block_diagonal: overlapping column ranges {:?} and {:?}",
                    pair[0], pair[1]
                ));
            }
        }
        Ok(ConstraintSet::BlockDiagonal { blocks, total_cols })
    }

    /// Locate the member block owning a joint row id.
    fn block_for_row<'a>(
        blocks: &'a [PlacedConstraintBlock],
        row: usize,
    ) -> Result<(&'a PlacedConstraintBlock, usize), String> {
        let mut offset = 0usize;
        for block in blocks {
            let rows = block.set.nrows();
            if row < offset + rows {
                return Ok((block, row - offset));
            }
            offset += rows;
        }
        Err(format!(
            "ConstraintSet: row {row} out of range ({offset} rows)"
        ))
    }

    pub fn nrows(&self) -> usize {
        match self {
            ConstraintSet::Dense(dense) => dense.a.nrows(),
            ConstraintSet::KhatriRaoCone(cone) => cone.nrows(),
            ConstraintSet::BlockDiagonal { blocks, .. } => {
                blocks.iter().map(|block| block.set.nrows()).sum()
            }
        }
    }

    pub fn ncols(&self) -> usize {
        match self {
            ConstraintSet::Dense(dense) => dense.a.ncols(),
            ConstraintSet::KhatriRaoCone(cone) => cone.ncols(),
            ConstraintSet::BlockDiagonal { total_cols, .. } => *total_cols,
        }
    }

    /// Raw constraint values `Aβ` (dense) / factored functional values (cone).
    pub fn values(&self, beta: ArrayView1<'_, f64>) -> Result<Array1<f64>, String> {
        match self {
            ConstraintSet::Dense(dense) => {
                if beta.len() != dense.a.ncols() {
                    return Err(format!(
                        "ConstraintSet: beta length {} != {}",
                        beta.len(),
                        dense.a.ncols()
                    ));
                }
                Ok(dense.a.dot(&beta))
            }
            ConstraintSet::KhatriRaoCone(cone) => cone.values(beta),
            ConstraintSet::BlockDiagonal { blocks, total_cols } => {
                if beta.len() != *total_cols {
                    return Err(format!(
                        "ConstraintSet: beta length {} != {}",
                        beta.len(),
                        total_cols
                    ));
                }
                let mut out = Array1::<f64>::zeros(self.nrows());
                let mut offset = 0usize;
                for block in blocks {
                    let width = block.set.ncols();
                    let local = beta.slice(ndarray::s![block.col_start..block.col_start + width]);
                    let values = block.set.values(local)?;
                    let rows = values.len();
                    out.slice_mut(ndarray::s![offset..offset + rows])
                        .assign(&values);
                    offset += rows;
                }
                Ok(out)
            }
        }
    }

    /// Right-hand sides (`b` dense; cone bounds are zero unless delta-shifted).
    pub fn bound(&self, row: usize) -> Result<f64, String> {
        match self {
            ConstraintSet::Dense(dense) => dense.b.get(row).copied().ok_or_else(|| {
                format!(
                    "ConstraintSet: row {row} out of range ({} rows)",
                    dense.b.len()
                )
            }),
            ConstraintSet::KhatriRaoCone(cone) => cone.bound(row),
            ConstraintSet::BlockDiagonal { blocks, .. } => {
                let (block, local) = Self::block_for_row(blocks, row)?;
                block.set.bound(local)
            }
        }
    }

    pub fn row_norm(&self, row: usize) -> Result<f64, String> {
        match self {
            ConstraintSet::Dense(dense) => {
                if row >= dense.a.nrows() {
                    return Err(format!(
                        "ConstraintSet: row {row} out of range ({} rows)",
                        dense.a.nrows()
                    ));
                }
                let r = dense.a.row(row);
                Ok(r.dot(&r).sqrt())
            }
            ConstraintSet::KhatriRaoCone(cone) => cone.row_norm(row),
            ConstraintSet::BlockDiagonal { blocks, .. } => {
                let (block, local) = Self::block_for_row(blocks, row)?;
                block.set.row_norm(local)
            }
        }
    }

    /// The same constraint system expressed in delta coordinates around
    /// `beta`: `A(β + δ) ≥ b  ⇔  Aδ ≥ b − Aβ`. The matrix carrier is shared;
    /// only the `O(nrows)` bounds change.
    pub fn shifted_to_delta(&self, beta: ArrayView1<'_, f64>) -> Result<Self, String> {
        let values = self.values(beta)?;
        match self {
            ConstraintSet::Dense(dense) => Ok(ConstraintSet::Dense(
                LinearInequalityConstraints::new(dense.a.clone(), &dense.b - &values)?,
            )),
            ConstraintSet::KhatriRaoCone(cone) => {
                let mut shifted = cone.clone();
                let base = shifted
                    .bounds
                    .take()
                    .unwrap_or_else(|| Array1::zeros(values.len()));
                shifted.bounds = Some(&base - &values);
                Ok(ConstraintSet::KhatriRaoCone(shifted))
            }
            ConstraintSet::BlockDiagonal { blocks, total_cols } => {
                let mut shifted_blocks = Vec::with_capacity(blocks.len());
                for block in blocks {
                    let width = block.set.ncols();
                    let local = beta.slice(ndarray::s![block.col_start..block.col_start + width]);
                    shifted_blocks.push(PlacedConstraintBlock {
                        col_start: block.col_start,
                        set: block.set.shifted_to_delta(local)?,
                    });
                }
                Ok(ConstraintSet::BlockDiagonal {
                    blocks: shifted_blocks,
                    total_cols: *total_cols,
                })
            }
        }
    }

    /// Scaled violation sweep: `max_r (b_r − (Aβ)_r) / ‖a_r‖` restricted to
    /// non-vacuous rows, plus the arg-max row. Matches the canonicalized dense
    /// geometry (unit rows) without materializing it.
    ///
    /// This is THE feasibility metric: `β` is feasible exactly when the value
    /// returned here is at or below [`PRIMAL_FEASIBILITY_TOL`].
    ///
    /// A vacuous row (`‖a‖ = 0`) with a bound at or below zero is `0 ≥ b`, true
    /// for every `β`, and contributes nothing. A vacuous row with a POSITIVE
    /// bound is `0 ≥ b > 0`: no `β` satisfies it, so its violation is infinite
    /// and the feasible set is empty. Reporting that as `+∞` — rather than
    /// skipping the row — is what makes this metric agree with
    /// `ConstraintSetOps::scaled_slack`, which already answers `−∞` for exactly
    /// this row, and keeps a gate built on this metric from silently admitting
    /// an unsatisfiable system.
    ///
    /// A row that cannot be decided by comparison — a non-finite row norm,
    /// bound or `a·β` — is refused rather than skipped (gam#2721): feasibility
    /// of an iterate that is not a number is undefined, and `violation > worst`
    /// being false for `NaN` would report the neutral `0.0` — "nothing is
    /// violated" — for exactly the iterate this metric exists to catch.
    pub fn max_scaled_violation(
        &self,
        beta: ArrayView1<'_, f64>,
    ) -> Result<(f64, Option<usize>), String> {
        let values = self.values(beta)?;
        // The sweep is a max over independent rows, and it is THE feasibility
        // verdict of every active-set solve, so it runs on every trial point.
        // On the large-scale CTN cone it is 1.6 M rows, and profiling the
        // preprocessor's reduced-face solve put 92 % of the process inside this
        // one function on ONE core (gam#979). Rows fan across the pool.
        //
        // The serial loop it replaces returned at the FIRST row that ended the
        // scan — an undecidable row, or a vacuous row with a positive bound —
        // so the reduction below carries the smallest such row index rather
        // than whichever thread found one first, and the running maximum breaks
        // exact ties toward the smaller index. Both make the verdict, the named
        // row, and the refusal text independent of how the rows were split.
        let metrics = match self {
            ConstraintSet::KhatriRaoCone(cone) => RowMetrics::Tiled {
                norms: cone.row_norms_slice(),
                tile: cone.tile_rows(),
                bounds: cone.bounds_slice(),
            },
            _ => RowMetrics::Carrier(self),
        };
        let sweep = (0..values.len())
            .into_par_iter()
            .fold(ScaledViolationSweep::none, |mut sweep, row| {
                let value = values[row];
                let (norm, bound) = match metrics.read(row) {
                    Ok(pair) => pair,
                    Err(error) => {
                        sweep.record_terminal(row, SweepTerminal::RowUnavailable(error));
                        return sweep;
                    }
                };
                // Decidability before comparison (gam#2721): `violation > worst`
                // is FALSE for `NaN`, so an undecidable row would leave `worst`
                // at `0.0` and this metric — THE feasibility verdict — would
                // call an iterate that is not a number feasible. `norm <= 0.0`
                // is false for a `NaN` norm too, so the vacuous-row branch below
                // cannot be the one that catches it. Refuse, naming the row and
                // the quantities.
                if !feasibility_quantities_are_finite(&[norm, bound, value]) {
                    sweep.record_terminal(
                        row,
                        SweepTerminal::Undecidable {
                            norm,
                            bound,
                            value,
                        },
                    );
                    return sweep;
                }
                if norm <= 0.0 {
                    if bound > 0.0 {
                        sweep.record_terminal(row, SweepTerminal::VacuousRowWithPositiveBound);
                    }
                    return sweep;
                }
                sweep.record_violation(row, (bound - value) / norm);
                sweep
            })
            .reduce(ScaledViolationSweep::none, ScaledViolationSweep::merge);
        sweep.verdict()
    }

    /// Largest `t ∈ [0, 1]` with `β + t·δ` feasible for every row, together
    /// with the first blocking row (the EXACT ratio test of a primal
    /// active-set method — zero tolerance, raw slacks). Rows already violated
    /// at `β` are reported as blocking at `t = 0`.
    ///
    /// This is the *pivot* rule: it answers "where does this chord cross a
    /// hyperplane in exact arithmetic", and its consumers (the feasible-chord
    /// clipper) want exactly that. It is NOT the rule for sizing a Newton step
    /// — a globalization that demands exact feasibility rejects steps this
    /// carrier's own contract calls feasible. Use
    /// `ConstraintSet::max_contract_feasible_step` for that.
    ///
    /// Like the contract rule, this one is TOTAL (gam#2721): a row that cannot
    /// be decided by comparison — a non-finite row norm, bound, `a·β` or `a·δ`
    /// — and that was not explicitly skipped is refused, because every
    /// comparison it would otherwise feed is false for `NaN` and the answer
    /// would be an unlimited `t = 1`.
    pub fn max_feasible_step(
        &self,
        beta: ArrayView1<'_, f64>,
        delta: ArrayView1<'_, f64>,
        skip_rows: &[usize],
    ) -> Result<(f64, Option<usize>), String> {
        let values = self.values(beta)?;
        let directional = self.values(delta)?;
        let mut skip = vec![false; values.len()];
        for &row in skip_rows {
            if row < skip.len() {
                skip[row] = true;
            }
        }
        let mut step = 1.0_f64;
        let mut blocking = None;
        for row in 0..values.len() {
            if skip[row] {
                continue;
            }
            let norm = self.row_norm(row)?;
            let bound = self.bound(row)?;
            let value = values[row];
            let rate = directional[row];
            // Same decidability requirement as the contract rule (gam#2721): a
            // `NaN` fails `rate >= 0.0` AND `t < step`, so the row would be
            // skipped twice over and this exact ratio test would answer
            // `step = 1.0` — "the whole chord is feasible" — for a chord that
            // is not a point. The clipper built on it would then accept the
            // endpoint. Refuse before comparing.
            if !feasibility_quantities_are_finite(&[norm, bound, value, rate]) {
                return Err(format!(
                    "ConstraintSet::max_feasible_step: row {row} cannot be decided \
                     (row norm {norm:.3e}, bound {bound:.3e}, value {value:.3e}, \
                     drift {rate:.3e}); every comparison in the ratio test is false \
                     for NaN, so skipping the row would report the whole step \
                     feasible (gam#2721)"
                ));
            }
            if norm <= 0.0 {
                continue;
            }
            if rate >= 0.0 {
                continue;
            }
            let t = (value - bound) / (-rate);
            if t < step {
                step = t.max(0.0);
                blocking = Some(row);
            }
        }
        Ok((step, blocking))
    }

    /// Materialize the requested rows densely (KKT systems on the active set).
    pub fn gather_rows(&self, rows: &[usize]) -> Result<LinearInequalityConstraints, String> {
        match self {
            ConstraintSet::Dense(dense) => {
                let mut a = Array2::<f64>::zeros((rows.len(), dense.a.ncols()));
                let mut b = Array1::<f64>::zeros(rows.len());
                for (out_row, &row) in rows.iter().enumerate() {
                    if row >= dense.a.nrows() {
                        return Err(format!(
                            "ConstraintSet: row {row} out of range ({} rows)",
                            dense.a.nrows()
                        ));
                    }
                    a.row_mut(out_row).assign(&dense.a.row(row));
                    b[out_row] = dense.b[row];
                }
                LinearInequalityConstraints::new(a, b)
            }
            ConstraintSet::KhatriRaoCone(cone) => cone.gather_rows(rows),
            ConstraintSet::BlockDiagonal { blocks, total_cols } => {
                let mut a = Array2::<f64>::zeros((rows.len(), *total_cols));
                let mut b = Array1::<f64>::zeros(rows.len());
                for (out_row, &row) in rows.iter().enumerate() {
                    let (block, local) = Self::block_for_row(blocks, row)?;
                    let gathered = block.set.gather_rows(&[local])?;
                    a.row_mut(out_row)
                        .slice_mut(ndarray::s![
                            block.col_start..block.col_start + block.set.ncols()
                        ])
                        .assign(&gathered.a.row(0));
                    b[out_row] = gathered.b[0];
                }
                LinearInequalityConstraints::new(a, b)
            }
        }
    }

    /// Exact dense equivalent of the whole set (tests / small systems only).
    pub fn to_dense(&self) -> Result<LinearInequalityConstraints, String> {
        match self {
            ConstraintSet::Dense(dense) => Ok(dense.clone()),
            _ => {
                let all: Vec<usize> = (0..self.nrows()).collect();
                self.gather_rows(&all)
            }
        }
    }
}

impl From<LinearInequalityConstraints> for ConstraintSet {
    fn from(dense: LinearInequalityConstraints) -> Self {
        ConstraintSet::Dense(dense)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::array;

    fn cone_fixture() -> KhatriRaoConeConstraints {
        // Ψ: 3 observations × 2 covariate columns; A is 3 coefficient rows
        // (row 0 = location, rows 1..2 = shape) × 2 columns.
        let psi = array![[1.0_f64, 0.5], [2.0, -1.0], [0.0, 3.0]];
        KhatriRaoConeConstraints::new(Arc::new(psi), vec![1, 2], 3).expect("cone fixture")
    }

    fn beta_fixture() -> Array1<f64> {
        // vec(A) row-major, A = [[9, -4], [1, 2], [0.5, -0.25]]
        array![9.0_f64, -4.0, 1.0, 2.0, 0.5, -0.25]
    }

    #[test]
    fn cone_values_match_dense_system() {
        let cone = cone_fixture();
        let set = ConstraintSet::KhatriRaoCone(cone.clone());
        let dense = ConstraintSet::Dense(cone.to_dense().expect("dense"));
        let beta = beta_fixture();
        let via_cone = set.values(beta.view()).expect("cone values");
        let via_dense = dense.values(beta.view()).expect("dense values");
        assert_eq!(via_cone.len(), 6);
        for (a, b) in via_cone.iter().zip(via_dense.iter()) {
            assert!((a - b).abs() < 1e-14, "cone/dense mismatch: {a} vs {b}");
        }
        // Spot-check one functional exactly: slot 0 (A row 1), observation 1:
        // ψ = (2, −1), A_{1,:} = (1, 2) → 2·1 − 1·2 = 0.
        assert!((via_cone[1] - 0.0).abs() < 1e-15);
    }

    #[test]
    fn cone_row_norms_are_factor_row_norms_for_every_slot() {
        let cone = cone_fixture();
        let set = ConstraintSet::KhatriRaoCone(cone);
        let expected = [(1.0_f64 + 0.25).sqrt(), (4.0_f64 + 1.0).sqrt(), 3.0_f64];
        for slot in 0..2 {
            for i in 0..3 {
                let norm = set.row_norm(slot * 3 + i).expect("norm");
                assert!((norm - expected[i]).abs() < 1e-15);
            }
        }
    }

    #[test]
    fn max_scaled_violation_agrees_with_canonicalized_dense() {
        let cone = cone_fixture();
        let set = ConstraintSet::KhatriRaoCone(cone.clone());
        let beta = beta_fixture();
        let (violation, row) = set.max_scaled_violation(beta.view()).expect("violation");
        // Dense oracle: canonicalize, then measure b − Aβ on unit rows.
        let dense = cone
            .to_dense()
            .expect("dense")
            .canonicalized()
            .expect("canon");
        let values = dense.a.dot(&beta);
        let mut worst = 0.0_f64;
        let mut worst_row = None;
        for r in 0..values.len() {
            let v = dense.b[r] - values[r];
            if v > worst {
                worst = v;
                worst_row = Some(r);
            }
        }
        assert!((violation - worst).abs() < 1e-14);
        assert_eq!(row, worst_row);
        assert!(violation > 0.0, "fixture must have a violated row");
    }

    #[test]
    fn max_feasible_step_matches_scalar_ratio_test() {
        let cone = cone_fixture();
        let set = ConstraintSet::KhatriRaoCone(cone);
        // Feasible start: shape rows of A strictly positive functionals.
        // A = [[0, 0], [1, 0.1], [1, 0.1]] → α values Ψ·(1, 0.1):
        // (1.05, 1.9, 0.3) — all positive for both slots.
        let beta = array![0.0_f64, 0.0, 1.0, 0.1, 1.0, 0.1];
        // Direction pushing slot 0 observation 2 down: δA_{1,:} = (0, −1) →
        // rate = ψ_2 · (0, −1) = −3; slack = 0.3 → t = 0.1. All other rows
        // untouched (rate 0 for slot 1, rates −0.5/1 for slot 0 rows 0/1:
        // row 0 rate = ψ_0·(0,−1) = −0.5, slack 1.05 → t = 2.1).
        let delta = array![0.0_f64, 0.0, 0.0, -1.0, 0.0, 0.0];
        let (step, blocking) = set
            .max_feasible_step(beta.view(), delta.view(), &[])
            .expect("step");
        assert!((step - 0.1).abs() < 1e-14, "expected 0.1, got {step}");
        assert_eq!(blocking, Some(2));
        // Skipping the blocking row exposes the next ratio (row 0, t = 2.1 → clamped to 1).
        let (step_skipped, blocking_skipped) = set
            .max_feasible_step(beta.view(), delta.view(), &[2])
            .expect("step skipped");
        assert!((step_skipped - 1.0).abs() < 1e-14);
        assert_eq!(blocking_skipped, None);
    }

    #[test]
    fn gather_rows_places_factor_rows_in_the_coupled_slot() {
        let cone = cone_fixture();
        // Row id 4 = slot 1 (A row 2), observation 1 → ψ = (2, −1) in cols 4..6.
        let gathered = cone.gather_rows(&[4]).expect("gather");
        assert_eq!(gathered.a.nrows(), 1);
        assert_eq!(gathered.a.ncols(), 6);
        let expected = [0.0, 0.0, 0.0, 0.0, 2.0, -1.0];
        for (j, &e) in expected.iter().enumerate() {
            assert_eq!(gathered.a[[0, j]], e);
        }
        assert_eq!(gathered.b[0], 0.0);
    }

    #[test]
    fn constructor_rejects_bad_coupled_rows() {
        let psi = array![[1.0_f64, 0.0], [0.0, 1.0]];
        assert!(KhatriRaoConeConstraints::new(Arc::new(psi.clone()), vec![3], 3).is_err());
        assert!(KhatriRaoConeConstraints::new(Arc::new(psi.clone()), vec![1, 1], 3).is_err());
        assert!(KhatriRaoConeConstraints::new(Arc::new(psi), vec![], 3).is_err());
    }

    #[test]
    fn shifted_to_delta_matches_dense_shift() {
        let cone = cone_fixture();
        let set = ConstraintSet::KhatriRaoCone(cone);
        let beta = beta_fixture();
        let shifted = set.shifted_to_delta(beta.view()).expect("shift");
        // Oracle: dense shift b' = b − Aβ.
        let dense = set.to_dense().expect("dense");
        let expected_b = &dense.b - &dense.a.dot(&beta);
        for row in 0..set.nrows() {
            assert!(
                (shifted.bound(row).expect("bound") - expected_b[row]).abs() < 1e-14,
                "shifted bound mismatch at row {row}"
            );
        }
        // The delta system at δ = 0 has slack equal to the original at β.
        let zero = Array1::<f64>::zeros(set.ncols());
        let (viol_delta, row_delta) = shifted
            .max_scaled_violation(zero.view())
            .expect("delta violation");
        let (viol_orig, row_orig) = set.max_scaled_violation(beta.view()).expect("violation");
        assert!((viol_delta - viol_orig).abs() < 1e-14);
        assert_eq!(row_delta, row_orig);
    }

    #[test]
    fn block_diagonal_composes_ids_bounds_and_values() {
        // Block 0: dense 2-row system on columns 0..2; block 1: cone on 2..8.
        let dense = LinearInequalityConstraints::new(
            array![[1.0_f64, 0.0], [0.0, -2.0]],
            array![0.5_f64, -1.0],
        )
        .expect("dense block");
        let cone = cone_fixture();
        let joint = ConstraintSet::block_diagonal(
            vec![
                PlacedConstraintBlock {
                    col_start: 0,
                    set: ConstraintSet::Dense(dense.clone()),
                },
                PlacedConstraintBlock {
                    col_start: 2,
                    set: ConstraintSet::KhatriRaoCone(cone.clone()),
                },
            ],
            8,
        )
        .expect("joint");
        assert_eq!(joint.nrows(), 2 + 6);
        assert_eq!(joint.ncols(), 8);
        let mut beta = Array1::<f64>::zeros(8);
        beta[0] = 2.0;
        beta[1] = 1.0;
        beta.slice_mut(ndarray::s![2..8]).assign(&beta_fixture());
        let values = joint.values(beta.view()).expect("values");
        assert!((values[0] - 2.0).abs() < 1e-15);
        assert!((values[1] + 2.0).abs() < 1e-15);
        let cone_values = cone.values(beta_fixture().view()).expect("cone values");
        for (idx, &cv) in cone_values.iter().enumerate() {
            assert!((values[2 + idx] - cv).abs() < 1e-15);
        }
        assert_eq!(joint.bound(0).expect("b0"), 0.5);
        assert_eq!(joint.bound(2).expect("b2"), 0.0);
        // Gathered joint row 3 (= cone row 1) occupies columns 2 + [2..4).
        let gathered = joint.gather_rows(&[3]).expect("gather");
        assert_eq!(gathered.a.ncols(), 8);
        assert_eq!(gathered.a[[0, 4]], 2.0);
        assert_eq!(gathered.a[[0, 5]], -1.0);
        // Overlapping ranges are rejected.
        assert!(
            ConstraintSet::block_diagonal(
                vec![
                    PlacedConstraintBlock {
                        col_start: 0,
                        set: ConstraintSet::Dense(dense.clone()),
                    },
                    PlacedConstraintBlock {
                        col_start: 1,
                        set: ConstraintSet::Dense(dense),
                    },
                ],
                8,
            )
            .is_err()
        );
    }

    #[test]
    fn zero_factor_rows_are_vacuous_not_violations() {
        // Ψ with an all-zero observation row: 0ᵀβ ≥ 0 is vacuous and must be
        // skipped by violation and ratio sweeps (norm 0), matching the dense
        // canonicalization contract for zero rows with b ≤ 0.
        let psi = array![[0.0_f64, 0.0], [1.0, 1.0]];
        let cone = KhatriRaoConeConstraints::new(Arc::new(psi), vec![1], 2).expect("cone");
        let set = ConstraintSet::KhatriRaoCone(cone);
        let beta = array![0.0_f64, 0.0, -5.0, 4.0];
        // Slot 0: values (0, −1). Row 0 vacuous; row 1 violated by 1/√2.
        let (violation, row) = set.max_scaled_violation(beta.view()).expect("violation");
        assert_eq!(row, Some(1));
        assert!((violation - 1.0 / 2.0_f64.sqrt()).abs() < 1e-14);
    }

}