ferrolearn-linear 0.5.0

Linear models for the ferrolearn ML framework
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
//! Logistic Regression with built-in cross-validated C selection.
//!
//! This module provides [`LogisticRegressionCV`], which automatically selects
//! the optimal regularization strength `C` by running k-fold cross-validation
//! over a grid of candidate values.
//!
//! ## REQ status (per `.design/linear/logistic_regression_cv.md`, mirrors `sklearn/linear_model/_logistic.py` @ 1.5.2)
//!
//! Mirrors `sklearn.linear_model.LogisticRegressionCV` (`_logistic.py:1464`): Cs=10 log-spaced
//! 1e-4..1e4, StratifiedKFold CV, accuracy scoring, best-C refit. Grid + scoring + the stratified
//! fold partition match sklearn exactly; exact `C_` is gated by the inner-LR LBFGS convergence #412.
//!
//! | REQ | Status | Evidence |
//! |---|---|---|
//! | REQ-1 (Cs grid 10 log-spaced 1e-4..1e4) | SHIPPED | matches `np.logspace(-4,4,10)` (`_logistic.py:276`) to ≤3.4e-13. |
//! | REQ-2 (stratified k-fold C-selection + refit) | SHIPPED | per-C accuracy over folds, argmax, refit `LogisticRegression`. Consumer: `pub use LogisticRegressionCV` (boundary API). C_ in sklearn's regularization region (exact C_ gated by #412). |
//! | REQ-3 (predict / predict_proba) | SHIPPED | `Predict`/`predict_proba for FittedLogisticRegressionCV`. |
//! | REQ-4 (fit_intercept / HasCoefficients / HasClasses) | SHIPPED | original-label classes_. |
//! | REQ-5 (sklearn StratifiedKFold partition) | SHIPPED | #456 fixed (45221f0) — `stratified_kfold_split` now uses sklearn's contiguous-block allocation (`_split.py:746-806`); deterministic unit test matches sklearn partition exactly. Was round-robin `i%k`. |
//! | REQ-6 (accuracy scoring matching sklearn scoring=None) | SHIPPED | per-fold `correct/count` = `accuracy_score`. |
//! | REQ-7..19 NOT-STARTED | Cs as int param (#457), CV attrs scores_/Cs_/coefs_paths_ (#458-#460), refit=False (#461), l1_ratios (#462), penalty (#442), solver variants (#443), multi_class=ovr (#444), class_weight (#445), sample_weight (#451), random_state/n_jobs (#452), ferray substrate (#453). |
//! | REQ-20 (non-finite input rejected) | SHIPPED | `Fit for LogisticRegressionCV::fit` rejects any NaN/+/-inf in X BEFORE the stratified k-fold split with `FerroError::InvalidParameter`, mirroring sklearn `LogisticRegressionCV.fit` `_validate_data(force_all_finite=True)` (`_logistic.py:1868`) → `ValueError("Input X contains NaN.")` / `"... contains infinity ..."`. `y` is `Array1<usize>` (integer class labels — never non-finite) and `Fit::fit` takes only `(x, y)` (no `sample_weight` in the trait surface), so X is the sole checkable input. `.iter().any(|v| !v.is_finite())` catches both NaN and Inf; the finite path is byte-identical (best_c_/coef_ unchanged). Verified vs the live sklearn 1.5.2 oracle (R-CHAR-3): NaN/+inf/-inf in X all raise `ValueError`; all-finite `C_=[0.1]` unchanged (`tests/divergence_linear_nonfinite_batch5.rs::logistic_regression_cv_*`). Non-test consumer: the existing `pub use logistic_regression_cv::LogisticRegressionCV` boundary API. (#2265) |
//!
//! acto-critic: Cs grid (3e-13), accuracy scoring, and (after #456) the StratifiedKFold partition
//! all match sklearn; #456 (round-robin → contiguous-block stratified folds) found + fixed — moved
//! C_ from the wrong regularization side to sklearn's region. Exact `C_` is bounded by the inner-LR
//! LBFGS convergence #412 (intercept precision at strong regularization). Two states only per R-DEFER-2.
//!
//! # Examples
//!
//! ```
//! use ferrolearn_linear::logistic_regression_cv::LogisticRegressionCV;
//! use ferrolearn_core::{Fit, Predict};
//! use ndarray::{array, Array2};
//!
//! let x = Array2::from_shape_vec(
//!     (12, 2),
//!     vec![
//!         1.0, 1.0, 1.0, 2.0, 2.0, 1.0, 2.0, 2.0, 1.5, 1.5, 1.0, 1.8,
//!         8.0, 8.0, 8.0, 9.0, 9.0, 8.0, 9.0, 9.0, 8.5, 8.5, 8.0, 8.8,
//!     ],
//! ).unwrap();
//! let y = array![0usize, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1];
//!
//! let model = LogisticRegressionCV::<f64>::new();
//! let fitted = model.fit(&x, &y).unwrap();
//! let preds = fitted.predict(&x).unwrap();
//! assert_eq!(preds.len(), 12);
//! ```

use ferrolearn_core::error::FerroError;
use ferrolearn_core::introspection::{HasClasses, HasCoefficients};
use ferrolearn_core::traits::{Fit, Predict};
use ndarray::{Array1, Array2, ScalarOperand};
use num_traits::Float;

use crate::logistic_regression::LogisticRegression;

/// Logistic Regression with cross-validated C selection.
///
/// Evaluates a grid of C values using k-fold cross-validation and refits
/// on the full dataset with the best C.
///
/// # Type Parameters
///
/// - `F`: The floating-point type (`f32` or `f64`).
#[derive(Debug, Clone)]
pub struct LogisticRegressionCV<F> {
    /// Candidate C values to evaluate.
    pub cs: Vec<F>,
    /// Number of cross-validation folds.
    pub cv: usize,
    /// Maximum number of L-BFGS iterations per fit.
    pub max_iter: usize,
    /// Convergence tolerance per fit.
    pub tol: F,
}

impl<F: Float> LogisticRegressionCV<F> {
    /// Create a new `LogisticRegressionCV` with default settings.
    ///
    /// Defaults: `cs` = 10 log-spaced values from 1e-4 to 1e4,
    /// `cv = 5`, `max_iter = 1000`, `tol = 1e-4`.
    #[must_use]
    pub fn new() -> Self {
        // Build log-spaced C values from 1e-4 to 1e4.
        let mut cs = Vec::with_capacity(10);
        for i in 0..10 {
            let exp = F::from(-4.0 + i as f64 * 8.0 / 9.0).unwrap();
            let base = F::from(10.0).unwrap();
            cs.push(base.powf(exp));
        }

        Self {
            cs,
            cv: 5,
            max_iter: 1000,
            tol: F::from(1e-4).unwrap(),
        }
    }

    /// Set the candidate C values.
    #[must_use]
    pub fn with_cs(mut self, cs: Vec<F>) -> Self {
        self.cs = cs;
        self
    }

    /// Set the number of cross-validation folds.
    #[must_use]
    pub fn with_cv(mut self, cv: usize) -> Self {
        self.cv = cv;
        self
    }

    /// Set the maximum number of iterations.
    #[must_use]
    pub fn with_max_iter(mut self, max_iter: usize) -> Self {
        self.max_iter = max_iter;
        self
    }

    /// Set the convergence tolerance.
    #[must_use]
    pub fn with_tol(mut self, tol: F) -> Self {
        self.tol = tol;
        self
    }
}

impl<F: Float> Default for LogisticRegressionCV<F> {
    fn default() -> Self {
        Self::new()
    }
}

/// Fitted Logistic Regression CV model.
///
/// Stores the model fitted with the best C value. Exposes `best_c()`.
#[derive(Debug, Clone)]
pub struct FittedLogisticRegressionCV<F> {
    /// The best C value found by cross-validation.
    best_c: F,
    /// The inner fitted logistic regression model.
    inner: crate::logistic_regression::FittedLogisticRegression<F>,
    /// CV scores for each C value (accuracy).
    cv_scores: Vec<F>,
    /// The C values evaluated.
    cs_evaluated: Vec<F>,
}

impl<F: Float> FittedLogisticRegressionCV<F> {
    /// Returns the best C value selected by cross-validation.
    #[must_use]
    pub fn best_c(&self) -> F {
        self.best_c
    }

    /// Returns the CV accuracy scores, one per C value.
    #[must_use]
    pub fn cv_scores(&self) -> &[F] {
        &self.cv_scores
    }

    /// Returns the C values that were evaluated.
    #[must_use]
    pub fn cs_evaluated(&self) -> &[F] {
        &self.cs_evaluated
    }
}

impl<F: Float + ndarray::ScalarOperand + Send + Sync + 'static> FittedLogisticRegressionCV<F> {
    /// Predict per-class probabilities. Mirrors sklearn
    /// `LogisticRegressionCV.predict_proba`. Delegates to the inner
    /// fitted LogisticRegression at the best C value.
    ///
    /// # Errors
    ///
    /// Forwards any error from the inner `predict_proba`.
    pub fn predict_proba(&self, x: &Array2<F>) -> Result<Array2<F>, FerroError> {
        self.inner.predict_proba(x)
    }

    /// Element-wise log of [`predict_proba`](Self::predict_proba).
    ///
    /// # Errors
    ///
    /// Forwards any error from [`predict_proba`](Self::predict_proba).
    pub fn predict_log_proba(&self, x: &Array2<F>) -> Result<Array2<F>, FerroError> {
        let proba = self.predict_proba(x)?;
        Ok(crate::log_proba(&proba))
    }

    /// Raw decision-function scores delegated to the inner fitted
    /// LogisticRegression at the best C value. Mirrors sklearn
    /// `LogisticRegressionCV.decision_function`.
    ///
    /// # Errors
    ///
    /// Forwards any error from the inner `decision_function`.
    pub fn decision_function(&self, x: &Array2<F>) -> Result<Array2<F>, FerroError> {
        self.inner.decision_function(x)
    }
}

/// StratifiedKFold-style split, replicating scikit-learn's non-shuffled
/// `StratifiedKFold._make_test_folds` (`sklearn/model_selection/_split.py:746-806`),
/// the default CV splitter for `LogisticRegressionCV` (#346, #456).
///
/// Each class's samples are assigned to folds in CONTIGUOUS blocks (not
/// round-robin `i % k`). The per-fold class allocation is the round-robin
/// distribution over the *sorted* class labels:
///   `allocation[j][c] = bincount(y_order[j::k])[c]` where `y_order = sort(y_encoded)`
/// (sklearn `_split.py:786-792`). Then, for each class `c`, its samples (in
/// original input order) are assigned to folds via
///   `folds_for_class = arange(k).repeat(allocation[:, c])`
/// so fold 0 takes the first `allocation[0][c]` of class-`c`'s samples, fold 1
/// the next `allocation[1][c]`, etc. (sklearn `_split.py:794-805`).
///
/// Returns `(train_indices, test_indices)` for fold number `fold`.
fn stratified_kfold_split(y: &Array1<usize>, k: usize, fold: usize) -> (Vec<usize>, Vec<usize>) {
    let n = y.len();

    // Distinct classes, sorted ascending — this is sklearn's class encoding
    // (`np.unique`); the encoded label of sample `i` is its position here.
    let mut classes: Vec<usize> = y.iter().copied().collect();
    classes.sort_unstable();
    classes.dedup();
    let n_classes = classes.len();
    let encode = |label: usize| classes.binary_search(&label).unwrap_or(0);

    // y_order = the class-encoded labels sorted ascending (all class-0 samples,
    // then class-1, ...). sklearn: `y_order = np.sort(y_encoded)`.
    let mut y_order: Vec<usize> = y.iter().map(|&label| encode(label)).collect();
    y_order.sort_unstable();

    // allocation[j][c] = bincount(y_order[j::k], minlength=n_classes)[c]
    // = number of class-c samples assigned to fold j (balances fold sizes).
    let mut allocation = vec![vec![0usize; n_classes]; k];
    for (j, alloc_row) in allocation.iter_mut().enumerate() {
        let mut idx = j;
        while idx < y_order.len() {
            alloc_row[y_order[idx]] += 1;
            idx += k;
        }
    }

    // For each class c, assign its original-order samples to folds in
    // contiguous blocks per `allocation[:, c]`, then collect those landing in
    // the requested `fold` as the test set.
    let mut test_indices: Vec<usize> = Vec::new();
    for (c, &cls) in classes.iter().enumerate() {
        let cls_indices: Vec<usize> = (0..n).filter(|&i| y[i] == cls).collect();
        let mut pos = 0usize;
        for (j, alloc_row) in allocation.iter().enumerate() {
            let count = alloc_row[c];
            if j == fold {
                for &idx in &cls_indices[pos..pos + count] {
                    test_indices.push(idx);
                }
            }
            pos += count;
        }
    }

    test_indices.sort_unstable();
    let test_set: std::collections::HashSet<usize> = test_indices.iter().copied().collect();
    let train_indices: Vec<usize> = (0..n).filter(|i| !test_set.contains(i)).collect();
    (train_indices, test_indices)
}

/// Extract rows from a 2D array by index.
fn select_rows<F: Float>(x: &Array2<F>, indices: &[usize]) -> Array2<F> {
    let n_features = x.ncols();
    let n_rows = indices.len();
    let mut result = Array2::<F>::zeros((n_rows, n_features));
    for (r, &i) in indices.iter().enumerate() {
        for j in 0..n_features {
            result[[r, j]] = x[[i, j]];
        }
    }
    result
}

/// Extract elements from a 1D array by index.
fn select_elements(y: &Array1<usize>, indices: &[usize]) -> Array1<usize> {
    Array1::from_vec(indices.iter().map(|&i| y[i]).collect())
}

impl<F: Float + Send + Sync + ScalarOperand + 'static> Fit<Array2<F>, Array1<usize>>
    for LogisticRegressionCV<F>
{
    type Fitted = FittedLogisticRegressionCV<F>;
    type Error = FerroError;

    /// Fit logistic regression with cross-validated C selection.
    ///
    /// For each candidate C, runs k-fold CV and computes accuracy.
    /// Selects the C with the highest mean accuracy and refits on the
    /// full dataset.
    ///
    /// # Errors
    ///
    /// - [`FerroError::ShapeMismatch`] — sample count mismatch.
    /// - [`FerroError::InvalidParameter`] — empty cs or cv < 2.
    /// - [`FerroError::InsufficientSamples`] — fewer than 2 classes or
    ///   too few samples for the number of folds.
    fn fit(
        &self,
        x: &Array2<F>,
        y: &Array1<usize>,
    ) -> Result<FittedLogisticRegressionCV<F>, FerroError> {
        let (n_samples, _n_features) = x.dim();

        if n_samples != y.len() {
            return Err(FerroError::ShapeMismatch {
                expected: vec![n_samples],
                actual: vec![y.len()],
                context: "y length must match number of samples in X".into(),
            });
        }

        if self.cs.is_empty() {
            return Err(FerroError::InvalidParameter {
                name: "cs".into(),
                reason: "must have at least one candidate C value".into(),
            });
        }

        if self.cv < 2 {
            return Err(FerroError::InvalidParameter {
                name: "cv".into(),
                reason: "must be at least 2".into(),
            });
        }

        // Non-finite input validation (#2265 batch5, ordering #2267). sklearn
        // `LogisticRegressionCV.fit` calls `self._validate_data(X, y, ...)`
        // (`_logistic.py:1868`) keeping the default `force_all_finite=True`, so
        // any NaN or +/-inf in X raises a `ValueError` at the very TOP of `fit`,
        // BEFORE `check_classification_targets` / `LabelEncoder().fit(y)`
        // (`_logistic.py:1876`-`:1889`) and BEFORE the stratified k-fold split.
        // So a NaN in X with `n_samples < cv` OR with a SINGLE class still raises
        // the non-finite error first — the finiteness check must therefore
        // precede BOTH the `n_samples < cv` guard and the class-count guard below
        // (#2267). `y` is `Array1<usize>` (integer class labels) which can never
        // be non-finite, and ferrolearn's `Fit::fit` takes only `(x, y)` (no
        // `sample_weight` in the trait surface), so X is the sole
        // non-finite-checkable input. `.iter().any(|v| !v.is_finite())` rejects
        // both NaN and Inf (bounds-safe, no panic, R-CODE-2), matching the crate
        // idiom. The finite path is byte-identical (the guard never fires on
        // finite input), and the FINITE-X `n_samples < cv` / single-class cases
        // still hit their respective `InsufficientSamples` guards below.
        if x.iter().any(|v| !v.is_finite()) {
            return Err(FerroError::InvalidParameter {
                name: "X".into(),
                reason: "Input X contains NaN or infinity.".into(),
            });
        }

        if n_samples < self.cv {
            return Err(FerroError::InsufficientSamples {
                required: self.cv,
                actual: n_samples,
                context: "need at least as many samples as CV folds".into(),
            });
        }

        let mut classes: Vec<usize> = y.to_vec();
        classes.sort_unstable();
        classes.dedup();

        if classes.len() < 2 {
            return Err(FerroError::InsufficientSamples {
                required: 2,
                actual: classes.len(),
                context: "LogisticRegressionCV requires at least 2 distinct classes".into(),
            });
        }

        // Evaluate each C via k-fold CV.
        let mut cv_scores = Vec::with_capacity(self.cs.len());
        let mut best_c = self.cs[0];
        let mut best_score = F::neg_infinity();

        for &c in &self.cs {
            let mut total_correct = 0usize;
            let mut total_count = 0usize;
            let mut fold_failed = false;

            for fold in 0..self.cv {
                let _ = n_samples; // computed via y.len() in the splitter
                let (train_idx, test_idx) = stratified_kfold_split(y, self.cv, fold);

                let x_train = select_rows(x, &train_idx);
                let y_train = select_elements(y, &train_idx);
                let x_test = select_rows(x, &test_idx);
                let y_test = select_elements(y, &test_idx);

                // Check that training set has at least 2 classes.
                let mut train_classes: Vec<usize> = y_train.to_vec();
                train_classes.sort_unstable();
                train_classes.dedup();
                if train_classes.len() < 2 {
                    fold_failed = true;
                    break;
                }

                let lr = LogisticRegression::<F>::new()
                    .with_c(c)
                    .with_max_iter(self.max_iter)
                    .with_tol(self.tol);

                match lr.fit(&x_train, &y_train) {
                    Ok(fitted) => match fitted.predict(&x_test) {
                        Ok(preds) => {
                            let correct = preds
                                .iter()
                                .zip(y_test.iter())
                                .filter(|(p, a)| p == a)
                                .count();
                            total_correct += correct;
                            total_count += y_test.len();
                        }
                        Err(_) => {
                            fold_failed = true;
                            break;
                        }
                    },
                    Err(_) => {
                        fold_failed = true;
                        break;
                    }
                }
            }

            let score = if fold_failed || total_count == 0 {
                F::zero()
            } else {
                F::from(total_correct).unwrap() / F::from(total_count).unwrap()
            };

            cv_scores.push(score);

            if score > best_score {
                best_score = score;
                best_c = c;
            }
        }

        // Refit on the full dataset with the best C.
        let lr = LogisticRegression::<F>::new()
            .with_c(best_c)
            .with_max_iter(self.max_iter)
            .with_tol(self.tol);

        let inner = lr.fit(x, y)?;

        Ok(FittedLogisticRegressionCV {
            best_c,
            inner,
            cv_scores,
            cs_evaluated: self.cs.clone(),
        })
    }
}

impl<F: Float + Send + Sync + ScalarOperand + 'static> Predict<Array2<F>>
    for FittedLogisticRegressionCV<F>
{
    type Output = Array1<usize>;
    type Error = FerroError;

    /// Predict class labels using the refitted model with the best C.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if the number of features
    /// does not match the fitted model.
    fn predict(&self, x: &Array2<F>) -> Result<Array1<usize>, FerroError> {
        self.inner.predict(x)
    }
}

impl<F: Float + Send + Sync + ScalarOperand + 'static> HasCoefficients<F>
    for FittedLogisticRegressionCV<F>
{
    fn coefficients(&self) -> &Array1<F> {
        self.inner.coefficients()
    }

    fn intercept(&self) -> F {
        self.inner.intercept()
    }
}

impl<F: Float + Send + Sync + ScalarOperand + 'static> HasClasses
    for FittedLogisticRegressionCV<F>
{
    fn classes(&self) -> &[usize] {
        self.inner.classes()
    }

    fn n_classes(&self) -> usize {
        self.inner.n_classes()
    }
}

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

    /// #456: `stratified_kfold_split` must replicate sklearn's non-shuffled
    /// `StratifiedKFold._make_test_folds` (contiguous blocks), NOT the old
    /// `i % k` round-robin. Expected test-fold memberships come from the live
    /// oracle (R-CHAR-3):
    ///   `StratifiedKFold(3).split(zeros, y)` for
    ///   `y=[0,0,0,0,1,1,1,1,1]` -> [[0,1,4],[2,5,6],[3,7,8]]
    /// and for the 18-sample test dataset
    ///   `y=[1,1,0,1,1,0,0,0,0,1,1,0,0,0,1,1,0,1]` ->
    ///   [[0,1,2,3,5,6],[4,7,8,9,10,11],[12,13,14,15,16,17]].
    #[test]
    fn stratified_kfold_split_matches_sklearn_partition() {
        // y=[0,0,0,0,1,1,1,1,1], k=3 (the goal.md verification case).
        let y = Array1::from_vec(vec![0usize, 0, 0, 0, 1, 1, 1, 1, 1]);
        let expected: [Vec<usize>; 3] = [vec![0, 1, 4], vec![2, 5, 6], vec![3, 7, 8]];
        for (fold, exp) in expected.iter().enumerate() {
            let (_train, test) = stratified_kfold_split(&y, 3, fold);
            assert_eq!(&test, exp, "fold {fold} test indices mismatch vs sklearn");
        }

        // The 18-sample, 2-class dataset used by the divergence integration test.
        let y2 = Array1::from_vec(vec![
            1usize, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1,
        ]);
        let expected2: [Vec<usize>; 3] = [
            vec![0, 1, 2, 3, 5, 6],
            vec![4, 7, 8, 9, 10, 11],
            vec![12, 13, 14, 15, 16, 17],
        ];
        for (fold, exp) in expected2.iter().enumerate() {
            let (train, test) = stratified_kfold_split(&y2, 3, fold);
            assert_eq!(&test, exp, "18-sample fold {fold} test mismatch vs sklearn");
            // train is the disjoint complement.
            assert_eq!(train.len() + test.len(), y2.len());
        }
    }

    #[test]
    fn test_default_constructor() {
        let m = LogisticRegressionCV::<f64>::new();
        assert_eq!(m.cv, 5);
        assert_eq!(m.max_iter, 1000);
        assert_eq!(m.cs.len(), 10);
    }

    #[test]
    fn test_builder() {
        let m = LogisticRegressionCV::<f64>::new()
            .with_cs(vec![0.1, 1.0, 10.0])
            .with_cv(3)
            .with_max_iter(500)
            .with_tol(1e-6);
        assert_eq!(m.cs.len(), 3);
        assert_eq!(m.cv, 3);
        assert_eq!(m.max_iter, 500);
    }

    #[test]
    fn test_binary_cv() {
        let x = Array2::from_shape_vec(
            (12, 2),
            vec![
                1.0, 1.0, 1.0, 2.0, 2.0, 1.0, 2.0, 2.0, 1.5, 1.5, 1.0, 1.8, 8.0, 8.0, 8.0, 9.0,
                9.0, 8.0, 9.0, 9.0, 8.5, 8.5, 8.0, 8.8,
            ],
        )
        .unwrap();
        let y = array![0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1];

        let model = LogisticRegressionCV::<f64>::new()
            .with_cs(vec![0.1, 1.0, 10.0])
            .with_cv(3);
        let fitted = model.fit(&x, &y).unwrap();

        assert!(fitted.best_c() > 0.0);
        assert_eq!(fitted.cv_scores().len(), 3);
        assert_eq!(fitted.cs_evaluated().len(), 3);

        let preds = fitted.predict(&x).unwrap();
        let correct: usize = preds.iter().zip(y.iter()).filter(|(p, a)| p == a).count();
        assert!(correct >= 8, "expected at least 8 correct, got {correct}");
    }

    #[test]
    fn test_multiclass_cv() {
        let x = Array2::from_shape_vec(
            (12, 2),
            vec![
                0.0, 0.0, 0.5, 0.0, 0.0, 0.5, 0.5, 0.5, 10.0, 0.0, 10.5, 0.0, 10.0, 0.5, 10.5, 0.5,
                0.0, 10.0, 0.5, 10.0, 0.0, 10.5, 0.5, 10.5,
            ],
        )
        .unwrap();
        let y = array![0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2];

        let model = LogisticRegressionCV::<f64>::new()
            .with_cs(vec![1.0, 10.0])
            .with_cv(2);
        let fitted = model.fit(&x, &y).unwrap();

        assert_eq!(fitted.n_classes(), 3);
        assert_eq!(fitted.classes(), &[0, 1, 2]);

        let preds = fitted.predict(&x).unwrap();
        assert_eq!(preds.len(), 12);
    }

    #[test]
    fn test_best_c_in_cs() {
        let x = Array2::from_shape_vec(
            (10, 2),
            vec![
                1.0, 1.0, 1.0, 2.0, 2.0, 1.0, 2.0, 2.0, 1.5, 1.5, 8.0, 8.0, 8.0, 9.0, 9.0, 8.0,
                9.0, 9.0, 8.5, 8.5,
            ],
        )
        .unwrap();
        let y = array![0, 0, 0, 0, 0, 1, 1, 1, 1, 1];

        let cs = vec![0.01, 0.1, 1.0, 10.0, 100.0];
        let model = LogisticRegressionCV::<f64>::new()
            .with_cs(cs.clone())
            .with_cv(2);
        let fitted = model.fit(&x, &y).unwrap();

        assert!(cs.contains(&fitted.best_c()));
    }

    #[test]
    fn test_shape_mismatch() {
        let x = Array2::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]).unwrap();
        let y = array![0, 1]; // Wrong length

        let model = LogisticRegressionCV::<f64>::new();
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_empty_cs() {
        let x = Array2::from_shape_vec((4, 1), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let y = array![0, 0, 1, 1];

        let model = LogisticRegressionCV::<f64>::new().with_cs(vec![]);
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_cv_too_small() {
        let x = Array2::from_shape_vec((4, 1), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let y = array![0, 0, 1, 1];

        let model = LogisticRegressionCV::<f64>::new().with_cv(1);
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_single_class_error() {
        let x = Array2::from_shape_vec((4, 1), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let y = array![0, 0, 0, 0];

        let model = LogisticRegressionCV::<f64>::new()
            .with_cs(vec![1.0])
            .with_cv(2);
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_has_coefficients() {
        let x = Array2::from_shape_vec(
            (10, 2),
            vec![
                1.0, 1.0, 1.0, 2.0, 2.0, 1.0, 2.0, 2.0, 1.5, 1.5, 8.0, 8.0, 8.0, 9.0, 9.0, 8.0,
                9.0, 9.0, 8.5, 8.5,
            ],
        )
        .unwrap();
        let y = array![0, 0, 0, 0, 0, 1, 1, 1, 1, 1];

        let fitted = LogisticRegressionCV::<f64>::new()
            .with_cs(vec![1.0])
            .with_cv(2)
            .fit(&x, &y)
            .unwrap();
        assert_eq!(fitted.coefficients().len(), 2);
    }

    #[test]
    fn test_stratified_kfold_split() {
        // 10 samples, 5 folds, balanced binary labels.
        let y: Array1<usize> = ndarray::array![0, 0, 0, 0, 0, 1, 1, 1, 1, 1];
        let (train, test) = stratified_kfold_split(&y, 5, 0);
        // Each fold gets one sample per class.
        assert_eq!(test.len(), 2);
        assert_eq!(train.len(), 8);
        // The test set must contain one of each class.
        let n0 = test.iter().filter(|&&i| y[i] == 0).count();
        let n1 = test.iter().filter(|&&i| y[i] == 1).count();
        assert_eq!(n0, 1);
        assert_eq!(n1, 1);
    }

    #[test]
    fn test_stratified_kfold_uneven() {
        // 7 samples, 3 folds, 4 of class 0 + 3 of class 1.
        let y: Array1<usize> = ndarray::array![0, 0, 0, 0, 1, 1, 1];
        let (train, test) = stratified_kfold_split(&y, 3, 0);
        // train + test must partition the 7 samples.
        assert_eq!(train.len() + test.len(), 7);
        // Class-0 distributed: 4 / 3 ≈ 1-2 per fold; class-1: 3 / 3 = 1 per fold.
        let n1 = test.iter().filter(|&&i| y[i] == 1).count();
        assert!(n1 >= 1, "class 1 must be present in stratified test fold");
    }
}