Skip to main content

linfa_logistic/
lib.rs

1//! # Logistic Regression
2//!
3//! ## The Big Picture
4//!
5//! `linfa-logistic` is a crate in the [`linfa`](https://crates.io/crates/linfa) ecosystem, an effort to create a toolkit for classical Machine Learning implemented in pure Rust, akin to Python's `scikit-learn`.
6//!
7//! ## Current state
8//! `linfa-logistic` provides a pure Rust implementation of a [binomial logistic regression model](LogisticRegression) and a [multinomial logistic regression model](MultiLogisticRegression).
9//!
10//! ## Examples
11//!
12//! There is an usage example in the `examples/` directory. To run, use:
13//!
14//! ```bash
15//! $ cargo run --example winequality
16//! ```
17//!
18
19pub mod error;
20
21use crate::error::{Error, Result};
22use argmin::core::{CostFunction, Executor, Gradient, IterState, OptimizationResult, Solver};
23use argmin::solver::linesearch::MoreThuenteLineSearch;
24use argmin::solver::quasinewton::LBFGS;
25use linfa::dataset::AsSingleTargets;
26use linfa::prelude::DatasetBase;
27use linfa::traits::{Fit, PredictInplace};
28use ndarray::{
29    s, Array, Array1, Array2, ArrayBase, ArrayView, ArrayView2, Axis, CowArray, Data, DataMut,
30    Dimension, IntoDimension, Ix1, Ix2, RemoveAxis, Slice, Zip,
31};
32use ndarray_stats::QuantileExt;
33use std::default::Default;
34
35#[cfg(feature = "serde")]
36use serde_crate::de::DeserializeOwned;
37#[cfg(feature = "serde")]
38use serde_crate::{Deserialize, Serialize};
39
40mod argmin_param;
41mod float;
42mod hyperparams;
43
44use argmin_param::*;
45use float::Float;
46use hyperparams::{LogisticRegressionParams, LogisticRegressionValidParams};
47
48/// A two-class logistic regression model.
49///
50/// Logistic regression combines linear models with
51/// the sigmoid function `sigm(x) = 1/(1+exp(-x))`
52/// to learn a family of functions that map the feature space to `[0,1]`.
53///
54/// Logistic regression is used in binary classification
55/// by interpreting the predicted value as the probability that the sample
56/// has label `1`. A threshold can be set in the [fitted model](FittedLogisticRegression) to decide the minimum
57/// probability needed to classify a sample as `1`, which defaults to `0.5`.
58///
59/// In this implementation any binary set of labels can be used, not necessarily `0` and `1`.
60///
61/// l2 regularization is used by this algorithm and is weighted by parameter `alpha`. Setting `alpha`
62/// close to zero removes regularization and the problem solved minimizes only the
63/// empirical risk. On the other hand, setting `alpha` to a high value increases
64/// the weight of the l2 norm of the linear model coefficients in the cost function.
65///
66/// ## Examples
67///
68/// Here's an example on how to train a logistic regression model on the `winequality` dataset
69/// ```rust
70/// use linfa::traits::{Fit, Predict};
71/// use linfa_logistic::LogisticRegression;
72///
73/// // Example on using binary labels different from 0 and 1
74/// let dataset = linfa_datasets::winequality().map_targets(|x| if *x > 6 { "good" } else { "bad" });
75/// let model = LogisticRegression::default().fit(&dataset).unwrap();
76/// let prediction = model.predict(&dataset);
77/// ```
78pub type LogisticRegression<F> = LogisticRegressionParams<F, Ix1>;
79
80/// Validated version of `LogisticRegression`
81pub type ValidLogisticRegression<F> = LogisticRegressionValidParams<F, Ix1>;
82
83/// A multinomial class logistic regression model.
84///
85/// The output labels can map to any discrete feature space, since the algorithm calculates the
86/// likelihood of a feature vector corresponding to any given outcome using the softmax function
87/// `softmax(x) = exp(x) / sum(exp(xi))`
88///
89/// l2 regularization is used by this algorithm and is weighted by parameter `alpha`. Setting `alpha`
90/// close to zero removes regularization and the problem solved minimizes only the
91/// empirical risk. On the other hand, setting `alpha` to a high value increases
92/// the weight of the l2 norm of the linear model coefficients in the cost function.
93pub type MultiLogisticRegression<F> = LogisticRegressionParams<F, Ix2>;
94
95/// Validated version of `MultiLogisticRegression`
96pub type ValidMultiLogisticRegression<F> = LogisticRegressionValidParams<F, Ix2>;
97
98impl<F: Float, D: Dimension> Default for LogisticRegressionParams<F, D> {
99    fn default() -> Self {
100        LogisticRegressionParams::new()
101    }
102}
103
104type LBFGSType<F, D> = LBFGS<
105    MoreThuenteLineSearch<ArgminParam<F, D>, ArgminParam<F, D>, F>,
106    ArgminParam<F, D>,
107    ArgminParam<F, D>,
108    F,
109>;
110type LBFGSType1<F> = LBFGSType<F, Ix1>;
111type LBFGSType2<F> = LBFGSType<F, Ix2>;
112
113type IterStateType<F, D> = IterState<ArgminParam<F, D>, ArgminParam<F, D>, (), (), (), F>;
114
115impl<F: Float, D: Dimension> LogisticRegressionValidParams<F, D> {
116    /// Create the initial parameters, either from a user supplied array
117    /// or an array of 0s
118    fn setup_init_params(&self, dims: D::Pattern) -> ArgminParam<F, D> {
119        if let Some(params) = self.initial_params.as_ref() {
120            ArgminParam(params.clone())
121        } else {
122            let mut dims = dims.into_dimension();
123            dims.as_array_view_mut()[0] += self.fit_intercept as usize;
124            ArgminParam(Array::zeros(dims))
125        }
126    }
127
128    /// Ensure that `x` and `y` have the right shape and that all data and
129    /// configuration parameters are finite.
130    fn validate_data<A: Data<Elem = F>, B: Data<Elem = F>>(
131        &self,
132        x: &ArrayBase<A, Ix2>,
133        y: &ArrayBase<B, D>,
134    ) -> Result<()> {
135        if x.shape()[0] != y.shape()[0] {
136            return Err(Error::MismatchedShapes(x.shape()[0], y.shape()[0]));
137        }
138        if x.iter().any(|x| !x.is_finite()) || y.iter().any(|y| !y.is_finite()) {
139            return Err(Error::InvalidValues);
140        }
141        self.validate_init_dims(x.shape()[1], y.shape().get(1).copied())?;
142        Ok(())
143    }
144
145    fn validate_init_dims(&self, mut n_features: usize, n_classes: Option<usize>) -> Result<()> {
146        if let Some(params) = self.initial_params.as_ref() {
147            let shape = params.shape();
148            n_features += self.fit_intercept as usize;
149            if n_features != shape[0] {
150                return Err(Error::InitialParameterFeaturesMismatch {
151                    n_features,
152                    rows: shape[0],
153                });
154            }
155            if let Some(n_classes) = n_classes {
156                if n_classes != shape[1] {
157                    return Err(Error::InitialParameterClassesMismatch {
158                        n_classes,
159                        cols: shape[1],
160                    });
161                }
162            }
163        }
164        Ok(())
165    }
166
167    /// Create a `LogisticRegressionProblem`.
168    fn setup_problem<'a, A: Data<Elem = F>>(
169        &self,
170        x: &'a ArrayBase<A, Ix2>,
171        target: Array<F, D>,
172    ) -> LogisticRegressionProblem<'a, F, A, D> {
173        LogisticRegressionProblem {
174            x,
175            target,
176            alpha: self.alpha,
177            offset: self.offset.clone(),
178        }
179    }
180
181    /// Create the LBFGS solver using MoreThuenteLineSearch and set gradient
182    /// tolerance.
183    fn setup_solver(&self) -> LBFGSType<F, D> {
184        let linesearch = MoreThuenteLineSearch::new();
185        LBFGS::new(linesearch, 10)
186            .with_tolerance_grad(self.gradient_tolerance)
187            .unwrap()
188    }
189}
190
191impl<
192        F: Float,
193        #[cfg(feature = "serde")] D: Dimension + Serialize + DeserializeOwned,
194        #[cfg(not(feature = "serde"))] D: Dimension,
195    > LogisticRegressionValidParams<F, D>
196{
197    /// Run the LBFGS solver until it converges or runs out of iterations.
198    fn run_solver<P: SolvableProblem<F, D>>(
199        &self,
200        problem: P,
201        solver: P::Solver,
202        init_params: ArgminParam<F, D>,
203    ) -> Result<OptimizationResult<P, P::Solver, IterStateType<F, D>>> {
204        Executor::new(problem, solver)
205            .configure(|state| state.param(init_params).max_iters(self.max_iterations))
206            .run()
207            .map_err(move |err| err.into())
208    }
209}
210
211impl<C: Ord + Clone, F: Float, D: Data<Elem = F>, T: AsSingleTargets<Elem = C>>
212    Fit<ArrayBase<D, Ix2>, T, Error> for ValidLogisticRegression<F>
213{
214    type Object = FittedLogisticRegression<F, C>;
215
216    /// Given a 2-dimensional feature matrix array `x` with shape
217    /// (n_samples, n_features) and an array of target classes to predict,
218    /// create a `FittedLinearRegression` object which allows making
219    /// predictions.
220    ///
221    /// The array of target classes `y` must have exactly two discrete values, (e.g. 0 and 1, "cat"
222    /// and "dog", ...), which represent the two different classes the model is supposed to
223    /// predict.
224    ///
225    /// The array `y` must also have exactly `n_samples` items, i.e.
226    /// exactly as many items as there are rows in the feature matrix `x`.
227    ///
228    /// This method returns an error if any of the preconditions are violated,
229    /// i.e. any values are `Inf` or `NaN`, `y` doesn't have as many items as
230    /// `x` has rows, or if other parameters (gradient_tolerance, alpha) have
231    /// been set to inalid values.
232    fn fit(&self, dataset: &DatasetBase<ArrayBase<D, Ix2>, T>) -> Result<Self::Object> {
233        let (x, y) = (dataset.records(), dataset.targets());
234        let (labels, target) = label_classes(y)?;
235        self.validate_data(x, &target)?;
236
237        if let Some(ref offset) = self.offset {
238            if offset.len() != x.nrows() {
239                return Err(Error::OffsetLengthMismatch {
240                    offset_len: offset.len(),
241                    n_samples: x.nrows(),
242                });
243            }
244        }
245
246        let problem = self.setup_problem(x, target);
247        let solver = self.setup_solver();
248        let init_params = self.setup_init_params(x.ncols());
249        let result = self.run_solver(problem, solver, init_params)?;
250
251        let params = result
252            .state
253            .best_param
254            .unwrap_or(self.setup_init_params(x.ncols()));
255        let (w, intercept) = convert_params(x.ncols(), params.as_array());
256        Ok(FittedLogisticRegression::new(
257            *intercept.view().into_scalar(),
258            w.to_owned(),
259            labels,
260        ))
261    }
262}
263
264impl<C: Ord + Clone, F: Float, D: Data<Elem = F>, T: AsSingleTargets<Elem = C>>
265    Fit<ArrayBase<D, Ix2>, T, Error> for ValidMultiLogisticRegression<F>
266{
267    type Object = MultiFittedLogisticRegression<F, C>;
268
269    /// Given a 2-dimensional feature matrix array `x` with shape
270    /// (n_samples, n_features) and an array of target classes to predict,
271    /// create a `MultiFittedLogisticRegression` object which allows making
272    /// predictions. The target classes can have any number of discrete values.
273    ///
274    /// This method returns an error if any of the preconditions are violated,
275    /// i.e. any values are `Inf` or `NaN`, `y` doesn't have as many items as
276    /// `x` has rows, or if other parameters (gradient_tolerance, alpha) have
277    /// been set to inalid values. The input features are also strongly recommended to be
278    /// normalized to ensure numerical stability.
279    fn fit(&self, dataset: &DatasetBase<ArrayBase<D, Ix2>, T>) -> Result<Self::Object> {
280        let (x, y) = (dataset.records(), dataset.targets());
281        let (classes, target) = label_classes_multi(y)?;
282        self.validate_data(x, &target)?;
283        let problem = self.setup_problem(x, target);
284        let solver = self.setup_solver();
285        let init_params = self.setup_init_params((x.ncols(), classes.len()));
286        let result = self.run_solver(problem, solver, init_params)?;
287
288        let params = result
289            .state
290            .best_param
291            .unwrap_or(self.setup_init_params((x.ncols(), classes.len())));
292        let (w, intercept) = convert_params(x.ncols(), params.as_array());
293        Ok(MultiFittedLogisticRegression::new(
294            intercept.to_owned(),
295            w.to_owned(),
296            classes,
297        ))
298    }
299}
300
301/// Identify the distinct values of the classes `y` and associate
302/// the target labels `-1.0` and `1.0` to it. -1.0 always labels the
303/// smaller class (by PartialOrd) and 1.0 always labels the larger
304/// class.
305///
306/// It is an error to have more than two classes.
307fn label_classes<F, T, C>(y: T) -> Result<(BinaryClassLabels<F, C>, Array1<F>)>
308where
309    F: Float,
310    T: AsSingleTargets<Elem = C>,
311    C: Ord + Clone,
312{
313    let y = y.as_single_targets();
314
315    let mut binary_classes = [None, None];
316    for class in y {
317        binary_classes = match binary_classes {
318            [None, None] => [Some((class, 1)), None],
319            [Some((c, count)), c2] if c == class => [Some((class, count + 1)), c2],
320            [c1, Some((c, count))] if c == class => [c1, Some((class, count + 1))],
321            [Some(c1), None] => [Some(c1), Some((class, 1))],
322            [None, Some(_)] => unreachable!("impossible binary class array"),
323            [Some(_), Some(_)] => return Err(Error::TooManyClasses),
324        };
325    }
326
327    let (class_a, class_b) = match binary_classes {
328        [Some(a), Some(b)] => (a, b),
329        _ => return Err(Error::TooFewClasses),
330    };
331
332    // Sort by label value (Ord), not by encounter order or count.
333    // The smaller label is always negative (-1),
334    // the larger label is always positive (+1).
335    let (neg_class, pos_class) = if class_a.0 < class_b.0 {
336        (class_a, class_b)
337    } else {
338        (class_b, class_a)
339    };
340
341    let target_array = y
342        .into_iter()
343        .map(|x| {
344            if x == pos_class.0 {
345                F::POSITIVE_LABEL
346            } else {
347                F::NEGATIVE_LABEL
348            }
349        })
350        .collect::<Array1<_>>();
351
352    Ok((
353        BinaryClassLabels {
354            pos: ClassLabel {
355                class: pos_class.0.clone(),
356                label: F::POSITIVE_LABEL,
357            },
358            neg: ClassLabel {
359                class: neg_class.0.clone(),
360                label: F::NEGATIVE_LABEL,
361            },
362        },
363        target_array,
364    ))
365}
366
367/// Identify the distinct values of the classes in `y` and map each value to an integer. Smaller
368/// classes (by `PartialOrd`) map to smaller integers. Returns the mapping along with a one-hot
369/// encoding of the numerical labels corresponding to `y`.
370fn label_classes_multi<F, T, C>(y: T) -> Result<(Vec<C>, Array2<F>)>
371where
372    F: Float,
373    T: AsSingleTargets<Elem = C>,
374    C: Ord + Clone,
375{
376    let y_single_target = y.as_single_targets();
377    let mut classes = y_single_target.to_vec();
378    // Dedup the list of classes
379    classes.sort();
380    classes.dedup();
381
382    let mut onehot = Array2::zeros((y_single_target.len(), classes.len()));
383    Zip::from(onehot.rows_mut())
384        .and(&y_single_target)
385        .for_each(|mut oh_row, cls| {
386            let idx = classes.binary_search(cls).unwrap();
387            oh_row[idx] = F::one();
388        });
389    Ok((classes, onehot))
390}
391
392/// Conditionally split the feature vector `w` into parameter vector and
393/// intercept parameter.
394/// Dimensions of `w` are either (f) or (f, n_classes)
395fn convert_params<F: Float, D: Dimension + RemoveAxis>(
396    n_features: usize,
397    w: &Array<F, D>,
398) -> (ArrayView<'_, F, D>, CowArray<'_, F, D::Smaller>) {
399    let nrows = w.shape()[0];
400    if n_features == nrows {
401        (
402            w.view(),
403            Array::zeros(w.raw_dim().remove_axis(Axis(0))).into(),
404        )
405    } else if n_features + 1 == nrows {
406        (
407            w.slice_axis(Axis(0), Slice::from(..n_features)),
408            w.index_axis(Axis(0), n_features).into(),
409        )
410    } else {
411        panic!(
412            "Unexpected length of parameter vector `w`, exected {} or {}, found {}",
413            n_features,
414            n_features + 1,
415            nrows
416        );
417    }
418}
419
420/// The logistic function
421fn logistic<F: linfa::Float>(x: F) -> F {
422    F::one() / (F::one() + (-x).exp())
423}
424
425/// A numerically stable version of the log of the logistic function.
426///
427/// Taken from scikit-learn
428/// https://github.com/scikit-learn/scikit-learn/blob/0.23.1/sklearn/utils/_logistic_sigmoid.pyx
429///
430/// See the blog post describing this implementation:
431/// http://fa.bianp.net/blog/2013/numerical-optimizers-for-logistic-regression/
432fn log_logistic<F: linfa::Float>(x: F) -> F {
433    if x > F::zero() {
434        -(F::one() + (-x).exp()).ln()
435    } else {
436        x - (F::one() + x.exp()).ln()
437    }
438}
439
440/// Finds the log of the sum of exponents across a specific axis in a numerically stable way. More
441/// specifically, computes `ln(exp(x1) + exp(x2) + exp(e3) + ...)` across an axis.
442///
443/// Based off this implementation: https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.logsumexp.html
444fn log_sum_exp<F: linfa::Float, A: Data<Elem = F>>(
445    m: &ArrayBase<A, Ix2>,
446    axis: Axis,
447) -> Array<F, Ix1> {
448    // Find max value of the array
449    let max = m.iter().copied().reduce(F::max).unwrap();
450    // Computes `max + ln(exp(x1-max) + exp(x2-max) + exp(x3-max) + ...)`, which is equal to the
451    // log_sum_exp formula
452    let reduced = m.fold_axis(axis, F::zero(), |acc, elem| *acc + (*elem - max).exp());
453    reduced.mapv_into(|e| e.max(F::cast(1e-15)).ln() + max)
454}
455
456/// Computes `exp(n - max) / sum(exp(n- max))`, which is a numerically stable version of softmax
457fn softmax_inplace<F: linfa::Float, A: DataMut<Elem = F>>(v: &mut ArrayBase<A, Ix1>) {
458    let max = v.iter().copied().reduce(F::max).unwrap();
459    v.mapv_inplace(|n| (n - max).exp());
460    let sum = v.sum();
461    v.mapv_inplace(|n| n / sum);
462}
463
464/// Computes the logistic loss assuming the training labels $y \in {-1, 1}$
465///
466/// Because the logistic function fullfills $\sigma(-z) = 1 - \sigma(z)$
467/// we can write $P(y=1|z) = \sigma(z) = \sigma(yz)$ and
468/// $P(y=-1|z) = 1 - P(y=1|z) = 1 - \sigma(z) = \sigma(-z) = \sigma(yz)$, so
469/// $P(y|z) = \sigma(yz)$ for both $y=1$ and $y=-1$.
470///
471/// Thus, the log loss can be written as
472/// $$-\sum_{i=1}^{N} \log(\sigma(y_i z_i)) + \frac{\alpha}{2}\text{params}^T\text{params}$$
473fn logistic_loss<F: Float, A: Data<Elem = F>>(
474    x: &ArrayBase<A, Ix2>,
475    y: &Array1<F>,
476    alpha: F,
477    w: &Array1<F>,
478    offset: Option<&Array1<F>>,
479) -> F {
480    let n_features = x.shape()[1];
481    let (params, intercept) = convert_params(n_features, w);
482    let yz = x.dot(&params.into_shape_with_order((params.len(), 1)).unwrap()) + intercept;
483    let len = yz.len();
484    let mut yz = yz.into_shape_with_order(len).unwrap();
485
486    if let Some(off) = offset {
487        yz += off;
488    }
489
490    yz *= y;
491    yz.mapv_inplace(log_logistic);
492    -yz.sum() + F::cast(0.5) * alpha * params.dot(&params)
493}
494
495/// Computes the gradient of the logistic loss function
496fn logistic_grad<F: Float, A: Data<Elem = F>>(
497    x: &ArrayBase<A, Ix2>,
498    y: &Array1<F>,
499    alpha: F,
500    w: &Array1<F>,
501    offset: Option<&Array1<F>>,
502) -> Array1<F> {
503    let n_features = x.shape()[1];
504    let (params, intercept) = convert_params(n_features, w);
505    let yz = x.dot(&params.into_shape_with_order((params.len(), 1)).unwrap()) + intercept;
506    let len = yz.len();
507    let mut yz = yz.into_shape_with_order(len).unwrap();
508
509    if let Some(off) = offset {
510        yz += off;
511    }
512
513    yz *= y;
514    yz.mapv_inplace(logistic);
515    yz -= F::one();
516    yz *= y;
517    if w.len() == n_features + 1 {
518        let mut grad = Array::zeros(w.len());
519        grad.slice_mut(s![..n_features])
520            .assign(&(x.t().dot(&yz) + (&params * alpha)));
521        grad[n_features] = yz.sum();
522        grad
523    } else {
524        x.t().dot(&yz) + (&params * alpha)
525    }
526}
527
528/// Compute the log of probabilities, which is `log(softmax(H))`, where H is `X . W + b`. Also
529/// returns `W` without the intercept.
530/// `Y` is the output (n_samples * n_classes), `X` is the input (n_samples * n_features), `W` is the
531/// params (n_features * n_classes), `b` is the intercept vector (n_classes).
532fn multi_logistic_prob_params<'a, F: Float, A: Data<Elem = F>>(
533    x: &ArrayBase<A, Ix2>,
534    w: &'a Array2<F>, // This parameter includes `W` and `b`
535) -> (Array2<F>, ArrayView2<'a, F>) {
536    let n_features = x.shape()[1];
537    let (params, intercept) = convert_params(n_features, w);
538    // Compute H
539    let h = x.dot(&params) + intercept;
540    // This computes `H - log(sum(exp(H)))`, which is equal to
541    // `log(softmax(H)) = log(exp(H) / sum(exp(H)))`
542    let log_prob = &h
543        - log_sum_exp(&h, Axis(1))
544            .into_shape_with_order((h.nrows(), 1))
545            .unwrap();
546    (log_prob, params)
547}
548
549/// Computes loss function of `-sum(Y * log(softmax(H))) + alpha/2 * norm(W)`
550fn multi_logistic_loss<F: Float, A: Data<Elem = F>>(
551    x: &ArrayBase<A, Ix2>,
552    y: &Array2<F>,
553    alpha: F,
554    w: &Array2<F>,
555) -> F {
556    let (log_prob, params) = multi_logistic_prob_params(x, w);
557    // Calculate loss
558    -elem_dot(&log_prob, y) + F::cast(0.5) * alpha * elem_dot(&params, &params)
559}
560
561/// Computes multinomial gradients for `W` and `b` and combine them.
562/// Gradient for `W` is `Xt . (softmax(H) - Y) + alpha * W`.
563/// Gradient for `b` is `sum(softmax(H) - Y)`.
564fn multi_logistic_grad<F: Float, A: Data<Elem = F>>(
565    x: &ArrayBase<A, Ix2>,
566    y: &Array2<F>,
567    alpha: F,
568    w: &Array2<F>,
569) -> Array2<F> {
570    let (log_prob, params) = multi_logistic_prob_params(x, w);
571    let (n_features, n_classes) = params.dim();
572    let intercept = w.nrows() > n_features;
573    let mut grad = Array::zeros((n_features + intercept as usize, n_classes));
574
575    // This value is `softmax(H)`
576    let prob = log_prob.mapv_into(num_traits::Float::exp);
577    let diff = prob - y;
578    // Compute gradient for `W` and place it at start of the grad matrix
579    let dw = x.t().dot(&diff) + (&params * alpha);
580    grad.slice_mut(s![..n_features, ..]).assign(&dw);
581    // Compute gradient for `b` and place it at end of grad matrix
582    if intercept {
583        grad.row_mut(n_features).assign(&diff.sum_axis(Axis(0)));
584    }
585    grad
586}
587
588/// A fitted logistic regression which can make predictions
589#[derive(Debug, Clone, PartialEq)]
590#[cfg_attr(
591    feature = "serde",
592    derive(Serialize, Deserialize),
593    serde(crate = "serde_crate")
594)]
595pub struct FittedLogisticRegression<F: Float, C: PartialOrd + Clone> {
596    threshold: F,
597    intercept: F,
598    params: Array1<F>,
599    labels: BinaryClassLabels<F, C>,
600}
601
602impl<F: Float, C: PartialOrd + Clone> FittedLogisticRegression<F, C> {
603    fn new(
604        intercept: F,
605        params: Array1<F>,
606        labels: BinaryClassLabels<F, C>,
607    ) -> FittedLogisticRegression<F, C> {
608        FittedLogisticRegression {
609            threshold: F::cast(0.5),
610            intercept,
611            params,
612            labels,
613        }
614    }
615
616    /// Set the probability threshold for which the 'positive' class will be
617    /// predicted. Defaults to 0.5.
618    pub fn set_threshold(mut self, threshold: F) -> FittedLogisticRegression<F, C> {
619        if threshold < F::zero() || threshold > F::one() {
620            panic!("FittedLogisticRegression::set_threshold: threshold needs to be between 0.0 and 1.0");
621        }
622        self.threshold = threshold;
623        self
624    }
625
626    pub fn intercept(&self) -> F {
627        self.intercept
628    }
629
630    pub fn params(&self) -> &Array1<F> {
631        &self.params
632    }
633
634    /// Get the model positive and negative classes mapped to their
635    /// corresponding problem input labels.
636    pub fn labels(&self) -> &BinaryClassLabels<F, C> {
637        &self.labels
638    }
639
640    /// Given a feature matrix, predict the probabilities that a sample
641    /// should be classified as the larger of the two classes learned when the
642    /// model was fitted.
643    pub fn predict_probabilities<A: Data<Elem = F>>(&self, x: &ArrayBase<A, Ix2>) -> Array1<F> {
644        let mut probs = x.dot(&self.params) + self.intercept;
645        probs.mapv_inplace(logistic);
646        probs
647    }
648}
649
650impl<C: PartialOrd + Clone + Default, F: Float, D: Data<Elem = F>>
651    PredictInplace<ArrayBase<D, Ix2>, Array1<C>> for FittedLogisticRegression<F, C>
652{
653    /// Given a feature matrix, predict the classes learned when the model was
654    /// fitted.
655    fn predict_inplace(&self, x: &ArrayBase<D, Ix2>, y: &mut Array1<C>) {
656        assert_eq!(
657            x.nrows(),
658            y.len(),
659            "The number of data points must match the number of output targets."
660        );
661        assert_eq!(
662            x.ncols(),
663            self.params.len(),
664            "Number of data features must match the number of features the model was trained with."
665        );
666
667        let pos_class = &self.labels.pos.class;
668        let neg_class = &self.labels.neg.class;
669        Zip::from(&self.predict_probabilities(x))
670            .and(y)
671            .for_each(|prob, out| {
672                *out = if *prob >= self.threshold {
673                    pos_class.clone()
674                } else {
675                    neg_class.clone()
676                }
677            });
678    }
679
680    fn default_target(&self, x: &ArrayBase<D, Ix2>) -> Array1<C> {
681        Array1::default(x.nrows())
682    }
683}
684
685/// A fitted multinomial logistic regression which can make predictions
686#[derive(Debug, Clone, PartialEq, Eq)]
687#[cfg_attr(
688    feature = "serde",
689    derive(Serialize, Deserialize),
690    serde(crate = "serde_crate")
691)]
692pub struct MultiFittedLogisticRegression<F, C: PartialOrd + Clone> {
693    intercept: Array1<F>,
694    params: Array2<F>,
695    classes: Vec<C>,
696}
697
698impl<F: Float, C: PartialOrd + Clone> MultiFittedLogisticRegression<F, C> {
699    fn new(intercept: Array1<F>, params: Array2<F>, classes: Vec<C>) -> Self {
700        Self {
701            intercept,
702            params,
703            classes,
704        }
705    }
706
707    pub fn intercept(&self) -> &Array1<F> {
708        &self.intercept
709    }
710
711    pub fn params(&self) -> &Array2<F> {
712        &self.params
713    }
714
715    /// Return non-normalized probabilities (n_samples * n_classes)
716    fn predict_nonorm_probabilities<A: Data<Elem = F>>(&self, x: &ArrayBase<A, Ix2>) -> Array2<F> {
717        x.dot(&self.params) + &self.intercept
718    }
719
720    /// Return normalized probabilities for each output class. The output dimensions are (n_samples
721    /// * n_classes).
722    pub fn predict_probabilities<A: Data<Elem = F>>(&self, x: &ArrayBase<A, Ix2>) -> Array2<F> {
723        let mut probs = self.predict_nonorm_probabilities(x);
724        probs
725            .rows_mut()
726            .into_iter()
727            .for_each(|mut row| softmax_inplace(&mut row));
728        probs
729    }
730
731    /// Get the list of class labels, which maps the numerical class indices to the labels
732    pub fn classes(&self) -> &[C] {
733        &self.classes
734    }
735}
736
737impl<C: PartialOrd + Clone + Default, F: Float, D: Data<Elem = F>>
738    PredictInplace<ArrayBase<D, Ix2>, Array1<C>> for MultiFittedLogisticRegression<F, C>
739{
740    /// Given a feature matrix, predict the classes learned when the model was
741    /// fitted.
742    fn predict_inplace(&self, x: &ArrayBase<D, Ix2>, y: &mut Array1<C>) {
743        assert_eq!(
744            x.nrows(),
745            y.len(),
746            "The number of data points must match the number of output targets."
747        );
748        assert_eq!(
749            x.ncols(),
750            self.params.nrows(),
751            "Number of data features must match the number of features the model was trained with."
752        );
753
754        let probs = self.predict_nonorm_probabilities(x);
755        Zip::from(probs.rows()).and(y).for_each(|prob_row, out| {
756            let idx = prob_row.argmax().unwrap();
757            *out = self.classes[idx].clone();
758        });
759    }
760
761    fn default_target(&self, x: &ArrayBase<D, Ix2>) -> Array1<C> {
762        Array1::default(x.nrows())
763    }
764}
765
766#[derive(Debug, Clone, PartialEq)]
767#[cfg_attr(
768    feature = "serde",
769    derive(Serialize, Deserialize),
770    serde(crate = "serde_crate")
771)]
772pub struct ClassLabel<F, C: PartialOrd> {
773    pub class: C,
774    pub label: F,
775}
776
777#[derive(Debug, Clone, PartialEq)]
778#[cfg_attr(
779    feature = "serde",
780    derive(Serialize, Deserialize),
781    serde(crate = "serde_crate")
782)]
783pub struct BinaryClassLabels<F, C: PartialOrd> {
784    pub pos: ClassLabel<F, C>,
785    pub neg: ClassLabel<F, C>,
786}
787
788/// Internal representation of a logistic regression problem.
789/// This data structure exists to be handed to Argmin.
790struct LogisticRegressionProblem<'a, F: Float, A: Data<Elem = F>, D: Dimension> {
791    x: &'a ArrayBase<A, Ix2>,
792    target: Array<F, D>,
793    alpha: F,
794    offset: Option<Array1<F>>,
795}
796
797type LogisticRegressionProblem1<'a, F, A> = LogisticRegressionProblem<'a, F, A, Ix1>;
798type LogisticRegressionProblem2<'a, F, A> = LogisticRegressionProblem<'a, F, A, Ix2>;
799
800impl<F: Float, A: Data<Elem = F>> CostFunction for LogisticRegressionProblem1<'_, F, A> {
801    type Param = ArgminParam<F, Ix1>;
802    type Output = F;
803
804    /// Apply the cost function to a parameter `p`
805    fn cost(&self, p: &Self::Param) -> std::result::Result<Self::Output, argmin::core::Error> {
806        let w = p.as_array();
807        let cost = logistic_loss(self.x, &self.target, self.alpha, w, self.offset.as_ref());
808        Ok(cost)
809    }
810}
811
812impl<F: Float, A: Data<Elem = F>> Gradient for LogisticRegressionProblem1<'_, F, A> {
813    type Param = ArgminParam<F, Ix1>;
814    type Gradient = ArgminParam<F, Ix1>;
815
816    /// Compute the gradient at parameter `p`.
817    fn gradient(&self, p: &Self::Param) -> std::result::Result<Self::Param, argmin::core::Error> {
818        let w = p.as_array();
819        let grad = ArgminParam(logistic_grad(
820            self.x,
821            &self.target,
822            self.alpha,
823            w,
824            self.offset.as_ref(),
825        ));
826        Ok(grad)
827    }
828}
829
830impl<F: Float, A: Data<Elem = F>> CostFunction for LogisticRegressionProblem2<'_, F, A> {
831    type Param = ArgminParam<F, Ix2>;
832    type Output = F;
833
834    /// Apply the cost function to a parameter `p`
835    fn cost(&self, p: &Self::Param) -> std::result::Result<Self::Output, argmin::core::Error> {
836        let w = p.as_array();
837        let cost = multi_logistic_loss(self.x, &self.target, self.alpha, w);
838        Ok(cost)
839    }
840}
841
842impl<F: Float, A: Data<Elem = F>> Gradient for LogisticRegressionProblem2<'_, F, A> {
843    type Param = ArgminParam<F, Ix2>;
844    type Gradient = ArgminParam<F, Ix2>;
845
846    /// Compute the gradient at parameter `p`.
847    fn gradient(&self, p: &Self::Param) -> std::result::Result<Self::Param, argmin::core::Error> {
848        let w = p.as_array();
849        let grad = ArgminParam(multi_logistic_grad(self.x, &self.target, self.alpha, w));
850        Ok(grad)
851    }
852}
853
854trait SolvableProblem<F: Float, D: Dimension>: Gradient + Sized {
855    type Solver: Solver<Self, IterStateType<F, D>>;
856}
857
858impl<F: Float, A: Data<Elem = F>> SolvableProblem<F, Ix1> for LogisticRegressionProblem1<'_, F, A> {
859    type Solver = LBFGSType1<F>;
860}
861
862impl<F: Float, A: Data<Elem = F>> SolvableProblem<F, Ix2> for LogisticRegressionProblem2<'_, F, A> {
863    type Solver = LBFGSType2<F>;
864}
865
866#[cfg(test)]
867mod test {
868    extern crate linfa;
869
870    use super::Error;
871    use super::*;
872    use approx::{assert_abs_diff_eq, assert_relative_eq, AbsDiffEq};
873    use linfa::prelude::*;
874    use ndarray::{array, Array2, Dim, Ix};
875
876    #[test]
877    fn autotraits() {
878        fn has_autotraits<T: Send + Sync + Sized + Unpin>() {}
879        has_autotraits::<LogisticRegressionParams<f64, Dim<[Ix; 0]>>>();
880        has_autotraits::<LogisticRegressionValidParams<f64, Dim<[Ix; 0]>>>();
881        has_autotraits::<ArgminParam<f64, Dim<[Ix; 0]>>>();
882    }
883
884    /// Test that the logistic loss function works as expected.
885    /// The expected values were obtained from running sklearn's
886    /// _logistic_loss_and_grad function.
887    #[test]
888    fn test_logistic_loss() {
889        let x = array![
890            [0.0],
891            [1.0],
892            [2.0],
893            [3.0],
894            [4.0],
895            [5.0],
896            [6.0],
897            [7.0],
898            [8.0],
899            [9.0]
900        ];
901        let y = array![0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
902        let ws = [
903            array![0.0, 0.0],
904            array![0.0, 1.0],
905            array![1.0, 0.0],
906            array![1.0, 1.0],
907            array![0.0, -1.0],
908            array![-1.0, 0.0],
909            array![-1.0, -1.0],
910        ];
911        let alphas = &[0.0, 1.0, 10.0];
912        let expecteds = vec![
913            6.931471805599453,
914            6.931471805599453,
915            6.931471805599453,
916            4.652158847349118,
917            4.652158847349118,
918            4.652158847349118,
919            2.8012999588008323,
920            3.3012999588008323,
921            7.801299958800833,
922            2.783195429782239,
923            3.283195429782239,
924            7.783195429782239,
925            10.652158847349117,
926            10.652158847349117,
927            10.652158847349117,
928            41.80129995880083,
929            42.30129995880083,
930            46.80129995880083,
931            47.78319542978224,
932            48.28319542978224,
933            52.78319542978224,
934        ];
935
936        for ((w, alpha), exp) in ws
937            .iter()
938            .flat_map(|w| alphas.iter().map(move |&alpha| (w, alpha)))
939            .zip(&expecteds)
940        {
941            assert_abs_diff_eq!(logistic_loss(&x, &y, alpha, w, None), *exp);
942        }
943    }
944
945    /// Test that the logistic grad function works as expected.
946    /// The expected values were obtained from running sklearn's
947    /// _logistic_loss_and_grad function.
948    #[test]
949    fn test_logistic_grad() {
950        let x = array![
951            [0.0],
952            [1.0],
953            [2.0],
954            [3.0],
955            [4.0],
956            [5.0],
957            [6.0],
958            [7.0],
959            [8.0],
960            [9.0]
961        ];
962        let y = array![0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
963        let ws = [
964            array![0.0, 0.0],
965            array![0.0, 1.0],
966            array![1.0, 0.0],
967            array![1.0, 1.0],
968            array![0.0, -1.0],
969            array![-1.0, 0.0],
970            array![-1.0, -1.0],
971        ];
972        let alphas = &[0.0, 1.0, 10.0];
973        let expecteds = vec![
974            array![-19.5, -3.],
975            array![-19.5, -3.],
976            array![-19.5, -3.],
977            array![-10.48871543, -1.61364853],
978            array![-10.48871543, -1.61364853],
979            array![-10.48871543, -1.61364853],
980            array![-0.13041554, -0.02852148],
981            array![0.86958446, -0.02852148],
982            array![9.86958446, -0.02852148],
983            array![-0.04834401, -0.01058067],
984            array![0.95165599, -0.01058067],
985            array![9.95165599, -0.01058067],
986            array![-28.51128457, -4.38635147],
987            array![-28.51128457, -4.38635147],
988            array![-28.51128457, -4.38635147],
989            array![-38.86958446, -5.97147852],
990            array![-39.86958446, -5.97147852],
991            array![-48.86958446, -5.97147852],
992            array![-38.95165599, -5.98941933],
993            array![-39.95165599, -5.98941933],
994            array![-48.95165599, -5.98941933],
995        ];
996
997        for ((w, alpha), exp) in ws
998            .iter()
999            .flat_map(|w| alphas.iter().map(move |&alpha| (w, alpha)))
1000            .zip(&expecteds)
1001        {
1002            let actual = logistic_grad(&x, &y, alpha, w, None);
1003            assert!(actual.abs_diff_eq(exp, 1e-8));
1004        }
1005    }
1006
1007    #[test]
1008    fn simple_example_1() {
1009        let log_reg = LogisticRegression::default();
1010        let x = array![[-1.0], [-0.01], [0.01], [1.0]];
1011        let y = array![0, 0, 1, 1];
1012        let dataset = Dataset::new(x, y);
1013        let res = log_reg.fit(&dataset).unwrap();
1014        assert_abs_diff_eq!(res.intercept(), 0.0);
1015        assert!(res.params().abs_diff_eq(&array![0.681], 1e-3));
1016        assert_eq!(
1017            &res.predict(dataset.records()),
1018            dataset.targets().as_single_targets()
1019        );
1020    }
1021
1022    #[test]
1023    fn simple_example_1_cats_dogs() {
1024        let log_reg = LogisticRegression::default();
1025        let x = array![[0.01], [1.0], [-1.0], [-0.01]];
1026        let y = array!["dog", "dog", "cat", "cat"];
1027        let dataset = Dataset::new(x, y);
1028        let res = log_reg.fit(&dataset).unwrap();
1029        assert_abs_diff_eq!(res.intercept(), 0.0);
1030        assert!(res.params().abs_diff_eq(&array![0.681], 1e-3));
1031        assert!(res
1032            .predict_probabilities(dataset.records())
1033            .abs_diff_eq(&array![0.501, 0.664, 0.335, 0.498], 1e-3));
1034        assert_eq!(
1035            &res.predict(dataset.records()),
1036            dataset.targets().as_single_targets()
1037        );
1038        assert_eq!(res.labels().pos.class, "dog");
1039        assert_eq!(res.labels().neg.class, "cat");
1040    }
1041
1042    #[test]
1043    fn simple_example_2() {
1044        let log_reg = LogisticRegression::default().alpha(1.0);
1045        let x = array![
1046            [0.0],
1047            [1.0],
1048            [2.0],
1049            [3.0],
1050            [4.0],
1051            [5.0],
1052            [6.0],
1053            [7.0],
1054            [8.0],
1055            [9.0]
1056        ];
1057        let y = array![0, 0, 0, 0, 1, 1, 1, 1, 1, 1];
1058        let dataset = Dataset::new(x, y);
1059        let res = log_reg.fit(&dataset).unwrap();
1060        assert_eq!(
1061            &res.predict(dataset.records()),
1062            dataset.targets().as_single_targets()
1063        );
1064    }
1065
1066    #[test]
1067    fn simple_example_3() {
1068        let x = array![[1.0], [0.0], [1.0], [0.0]];
1069        let y = array![1, 0, 1, 0];
1070        let dataset = DatasetBase::new(x, y);
1071        let model = LogisticRegression::default().fit(&dataset).unwrap();
1072
1073        let pred = model.predict(&dataset.records);
1074        assert_eq!(dataset.targets(), pred);
1075    }
1076
1077    #[test]
1078    fn rejects_mismatching_x_y() {
1079        let log_reg = LogisticRegression::default();
1080        let x = array![[-1.0], [-0.01], [0.01]];
1081        let y = array![0, 0, 1, 1];
1082        let res = log_reg.fit(&Dataset::new(x, y));
1083        assert!(matches!(res.unwrap_err(), Error::MismatchedShapes(3, 4)));
1084    }
1085
1086    #[test]
1087    fn rejects_inf_values() {
1088        let infs = &[f64::INFINITY, f64::NEG_INFINITY, f64::NAN];
1089        let inf_xs: Vec<_> = infs.iter().map(|&inf| array![[1.0], [inf]]).collect();
1090        let log_reg = LogisticRegression::default();
1091        let normal_x = array![[-1.0], [1.0]];
1092        let y = array![0, 1];
1093        for inf_x in &inf_xs {
1094            let res = log_reg.fit(&DatasetBase::new(inf_x.view(), &y));
1095            assert!(matches!(res.unwrap_err(), Error::InvalidValues));
1096        }
1097        for inf in infs {
1098            let log_reg = LogisticRegression::default().alpha(*inf);
1099            let res = log_reg.fit(&DatasetBase::new(normal_x.view(), &y));
1100            assert!(matches!(res.unwrap_err(), Error::InvalidAlpha));
1101        }
1102        let mut non_positives = infs.to_vec();
1103        non_positives.push(-1.0);
1104        non_positives.push(0.0);
1105        for inf in &non_positives {
1106            let log_reg = LogisticRegression::default().gradient_tolerance(*inf);
1107            let res = log_reg.fit(&Dataset::new(normal_x.to_owned(), y.to_owned()));
1108            assert!(matches!(res.unwrap_err(), Error::InvalidGradientTolerance));
1109        }
1110    }
1111
1112    #[test]
1113    fn validates_initial_params() {
1114        let infs = &[f64::INFINITY, f64::NEG_INFINITY, f64::NAN];
1115        let normal_x = array![[-1.0], [1.0]];
1116        let normal_y = array![0, 1];
1117        let dataset = Dataset::new(normal_x, normal_y);
1118        for inf in infs {
1119            let log_reg = LogisticRegression::default().initial_params(array![*inf, 0.0]);
1120            let res = log_reg.fit(&dataset);
1121            assert!(matches!(res.unwrap_err(), Error::InvalidInitialParameters));
1122        }
1123        {
1124            let log_reg = LogisticRegression::default().initial_params(array![0.0, 0.0, 0.0]);
1125            let res = log_reg.fit(&dataset);
1126            assert!(matches!(
1127                res.unwrap_err(),
1128                Error::InitialParameterFeaturesMismatch {
1129                    rows: 3,
1130                    n_features: 2
1131                }
1132            ));
1133        }
1134        {
1135            let log_reg = LogisticRegression::default()
1136                .with_intercept(false)
1137                .initial_params(array![0.0, 0.0]);
1138            let res = log_reg.fit(&dataset);
1139            assert!(matches!(
1140                res.unwrap_err(),
1141                Error::InitialParameterFeaturesMismatch {
1142                    rows: 2,
1143                    n_features: 1
1144                }
1145            ));
1146        }
1147    }
1148
1149    #[test]
1150    fn uses_initial_params() {
1151        let params = array![1.2, -4.12];
1152        let log_reg = LogisticRegression::default()
1153            .initial_params(params)
1154            .max_iterations(5);
1155        let x = array![
1156            [0.0],
1157            [1.0],
1158            [2.0],
1159            [3.0],
1160            [4.0],
1161            [5.0],
1162            [6.0],
1163            [7.0],
1164            [8.0],
1165            [9.0]
1166        ];
1167        let y = array![0, 0, 0, 0, 1, 1, 1, 1, 1, 1];
1168        let dataset = Dataset::new(x, y);
1169        let res = log_reg.fit(&dataset).unwrap();
1170        assert!(res.intercept().abs_diff_eq(&-4.124, 1e-3));
1171        assert!(res.params().abs_diff_eq(&array![1.181], 1e-3));
1172        assert_eq!(
1173            &res.predict(dataset.records()),
1174            dataset.targets().as_single_targets()
1175        );
1176
1177        // Test serialization
1178        #[cfg(feature = "serde")]
1179        {
1180            let ser = rmp_serde::to_vec(&res).unwrap();
1181            let unser: FittedLogisticRegression<f32, f32> = rmp_serde::from_slice(&ser).unwrap();
1182
1183            let x = array![[1.0]];
1184            let y_hat = unser.predict(&x);
1185
1186            assert!(y_hat[0] == 0.0);
1187        }
1188    }
1189
1190    #[test]
1191    fn works_with_f32() {
1192        let log_reg = LogisticRegression::default();
1193        let x: Array2<f32> = array![[-1.0], [-0.01], [0.01], [1.0]];
1194        let y = array![0, 0, 1, 1];
1195        let dataset = Dataset::new(x, y);
1196        let res = log_reg.fit(&dataset).unwrap();
1197        assert_abs_diff_eq!(res.intercept(), 0.0_f32);
1198        assert!(res.params().abs_diff_eq(&array![0.682_f32], 1e-3));
1199        assert_eq!(
1200            &res.predict(dataset.records()),
1201            dataset.targets().as_single_targets()
1202        );
1203    }
1204
1205    #[test]
1206    fn test_log_sum_exp() {
1207        let data = array![[3.3, 0.4, -2.1], [0.4, 2.2, -0.1], [1., 0., -1.]];
1208        let out = log_sum_exp(&data, Axis(1));
1209        assert_abs_diff_eq!(out, array![3.35783, 2.43551, 1.40761], epsilon = 1e-5);
1210    }
1211
1212    #[test]
1213    fn test_softmax() {
1214        let mut data = array![3.3, 5.5, 0.1, -4.4, 8.0];
1215        softmax_inplace(&mut data);
1216        assert_relative_eq!(
1217            data,
1218            array![0.0083324, 0.075200047, 0.000339647, 0.000003773, 0.91612413],
1219            epsilon = 1e-8
1220        );
1221        assert_abs_diff_eq!(data.sum(), 1.0);
1222    }
1223
1224    #[test]
1225    fn test_multi_logistic_loss_grad() {
1226        let x = array![
1227            [0.0, 0.5],
1228            [1.0, -1.0],
1229            [2.0, -2.0],
1230            [3.0, -3.0],
1231            [4.0, -4.0],
1232            [5.0, -5.0],
1233            [6.0, -6.0],
1234            [7.0, -7.0],
1235        ];
1236        let y = array![
1237            [1.0, 0.0, 0.0],
1238            [1.0, 0.0, 0.0],
1239            [0.0, 1.0, 0.0],
1240            [0.0, 1.0, 0.0],
1241            [0.0, 1.0, 0.0],
1242            [0.0, 0.0, 1.0],
1243            [0.0, 0.0, 1.0],
1244            [0.0, 0.0, 1.0],
1245        ];
1246        let params1 = array![[4.4, -1.2, 3.3], [3.4, 0.1, 0.0]];
1247        let params2 = array![[0.001, -3.2, 2.9], [0.1, 4.5, 5.7], [4.5, 2.2, 1.7]];
1248        let alpha = 0.6;
1249
1250        {
1251            let (log_prob, w) = multi_logistic_prob_params(&x, &params1);
1252            assert_abs_diff_eq!(
1253                log_prob,
1254                array![
1255                    [-3.18259845e-01, -1.96825985e+00, -2.01825985e+00],
1256                    [-2.40463987e+00, -4.70463987e+00, -1.04639868e-01],
1257                    [-4.61010168e+00, -9.21010168e+00, -1.01016809e-02],
1258                    [-6.90100829e+00, -1.38010083e+01, -1.00829256e-03],
1259                    [-9.20010104e+00, -1.84001010e+01, -1.01044506e-04],
1260                    [-1.15000101e+01, -2.30000101e+01, -1.01301449e-05],
1261                    [-1.38000010e+01, -2.76000010e+01, -1.01563199e-06],
1262                    [-1.61000001e+01, -3.22000001e+01, -1.01826043e-07],
1263                ],
1264                epsilon = 1e-6
1265            );
1266            assert_abs_diff_eq!(w, params1);
1267            let loss = multi_logistic_loss(&x, &y, alpha, &params1);
1268            assert_abs_diff_eq!(loss, 57.11212197835295, epsilon = 1e-6);
1269            let grad = multi_logistic_grad(&x, &y, alpha, &params1);
1270            assert_abs_diff_eq!(
1271                grad,
1272                array![
1273                    [1.7536815, -9.71074369, 11.85706219],
1274                    [2.79002537, 9.12059357, -9.81061893]
1275                ],
1276                epsilon = 1e-6
1277            );
1278        }
1279
1280        {
1281            let (log_prob, w) = multi_logistic_prob_params(&x, &params2);
1282            assert_abs_diff_eq!(
1283                log_prob,
1284                array![
1285                    [-1.06637742e+00, -1.16637742e+00, -1.06637742e+00],
1286                    [-4.12429463e-03, -9.90512429e+00, -5.50512429e+00],
1287                    [-2.74092305e-04, -1.75022741e+01, -8.20227409e+00],
1288                    [-1.84027855e-05, -2.51030184e+01, -1.09030184e+01],
1289                    [-1.23554225e-06, -3.27040012e+01, -1.36040012e+01],
1290                    [-8.29523046e-08, -4.03050001e+01, -1.63050001e+01],
1291                    [-5.56928016e-09, -4.79060000e+01, -1.90060000e+01],
1292                    [-3.73912013e-10, -5.55070000e+01, -2.17070000e+01]
1293                ],
1294                epsilon = 1e-6
1295            );
1296            assert_abs_diff_eq!(w, params2.slice(s![..params2.nrows() - 1, ..]));
1297            let loss = multi_logistic_loss(&x, &y, alpha, &params2);
1298            assert_abs_diff_eq!(loss, 154.8177958366479, epsilon = 1e-6);
1299            let grad = multi_logistic_grad(&x, &y, alpha, &params2);
1300            assert_abs_diff_eq!(
1301                grad,
1302                array![
1303                    [26.99587549, -10.91995003, -16.25532546],
1304                    [-27.26314882, 11.85569669, 21.58745213],
1305                    [5.33984376, -2.68845675, -2.65138701]
1306                ],
1307                epsilon = 1e-6
1308            );
1309        }
1310    }
1311
1312    #[test]
1313    fn simple_multi_example() {
1314        let x = array![[-1., 0.], [0., 1.], [1., 1.]];
1315        let y = array![2, 1, 0];
1316        let log_reg = MultiLogisticRegression::default()
1317            .alpha(0.1)
1318            .initial_params(Array::zeros((3, 3)));
1319        let dataset = Dataset::new(x, y);
1320        let res = log_reg.fit(&dataset).unwrap();
1321        assert_eq!(res.params().dim(), (2, 3));
1322        assert_eq!(res.intercept().dim(), 3);
1323        assert_eq!(
1324            &res.predict(dataset.records()),
1325            dataset.targets().as_single_targets()
1326        );
1327    }
1328
1329    #[test]
1330    fn simple_multi_example_2() {
1331        let x = array![[1.0], [0.0], [1.0], [0.0]];
1332        let y = array![1, 0, 1, 0];
1333        let dataset = DatasetBase::new(x, y);
1334        let model = MultiLogisticRegression::default().fit(&dataset).unwrap();
1335
1336        let pred = model.predict(&dataset.records);
1337        assert_eq!(dataset.targets(), pred);
1338    }
1339
1340    #[test]
1341    fn simple_multi_example_text() {
1342        let log_reg = MultiLogisticRegression::default().alpha(0.1);
1343        let x = array![[0.1], [1.0], [-1.0], [-0.1]];
1344        let y = array!["dog", "ape", "rocket", "cat"];
1345        let dataset = Dataset::new(x, y);
1346        let res = log_reg.fit(&dataset).unwrap();
1347        assert_eq!(res.params().dim(), (1, 4));
1348        assert_eq!(res.intercept().dim(), 4);
1349        assert_eq!(
1350            &res.predict(dataset.records()),
1351            dataset.targets().as_single_targets()
1352        );
1353    }
1354
1355    #[test]
1356    fn multi_on_binary_problem() {
1357        let log_reg = MultiLogisticRegression::default().alpha(1.0);
1358        let x = array![
1359            [0.0],
1360            [1.0],
1361            [2.0],
1362            [3.0],
1363            [4.0],
1364            [5.0],
1365            [6.0],
1366            [7.0],
1367            [8.0],
1368            [9.0]
1369        ];
1370        let y = array![0, 0, 0, 0, 1, 1, 1, 1, 1, 1];
1371        let dataset = Dataset::new(x, y);
1372        let res = log_reg.fit(&dataset).unwrap();
1373        assert_eq!(res.params().dim(), (1, 2));
1374        assert_eq!(res.intercept().dim(), 2);
1375        assert_eq!(
1376            &res.predict(dataset.records()),
1377            dataset.targets().as_single_targets()
1378        );
1379    }
1380
1381    #[test]
1382    fn reject_num_class_mismatch() {
1383        let n_samples = 4;
1384        let n_classes = 3;
1385        let n_features = 1;
1386        let x = Array2::<f64>::zeros((n_samples, n_features));
1387        let y = array![0, 1, 2, 0];
1388        let dataset = Dataset::new(x, y);
1389
1390        let log_reg = MultiLogisticRegression::default()
1391            .with_intercept(false)
1392            .initial_params(Array::zeros((n_features, n_classes - 1)));
1393        assert!(matches!(
1394            log_reg.fit(&dataset).unwrap_err(),
1395            Error::InitialParameterClassesMismatch {
1396                cols: 2,
1397                n_classes: 3,
1398            }
1399        ));
1400    }
1401
1402    #[test]
1403    fn label_order_independent() {
1404        let x1 = array![[-1.0], [1.0], [-0.5], [0.5]];
1405        let y1 = array!["cat", "dog", "cat", "dog"];
1406
1407        let x2 = array![[1.0], [-1.0], [0.5], [-0.5]];
1408        let y2 = array!["dog", "cat", "dog", "cat"];
1409
1410        let model1 = LogisticRegression::default()
1411            .fit(&Dataset::new(x1, y1))
1412            .unwrap();
1413        let model2 = LogisticRegression::default()
1414            .fit(&Dataset::new(x2, y2))
1415            .unwrap();
1416
1417        assert_eq!(model1.labels().pos.class, "dog");
1418        assert_eq!(model1.labels().neg.class, "cat");
1419        assert_eq!(model2.labels().pos.class, "dog");
1420        assert_eq!(model2.labels().neg.class, "cat");
1421
1422        assert_abs_diff_eq!(model1.intercept(), model2.intercept());
1423        assert!(model1.params().abs_diff_eq(model2.params(), 1e-6));
1424    }
1425
1426    #[test]
1427    fn rejects_mismatched_offset_length() {
1428        let log_reg = LogisticRegression::default().offset(array![1.0, 2.0, 3.0]);
1429        let x = array![[-1.0], [-0.01], [0.01], [1.0]];
1430        let y = array![0, 0, 1, 1];
1431        let res = log_reg.fit(&Dataset::new(x, y));
1432        assert!(matches!(
1433            res.unwrap_err(),
1434            Error::OffsetLengthMismatch {
1435                offset_len: 3,
1436                n_samples: 4,
1437            }
1438        ));
1439    }
1440
1441    #[test]
1442    fn zero_offset_same_as_no_offset() {
1443        let x = array![[-1.0], [-0.01], [0.01], [1.0]];
1444        let y = array![0, 0, 1, 1];
1445
1446        let model_none = LogisticRegression::default()
1447            .fit(&Dataset::new(x.clone(), y.clone()))
1448            .unwrap();
1449
1450        let model_zero = LogisticRegression::default()
1451            .offset(array![0.0, 0.0, 0.0, 0.0])
1452            .fit(&Dataset::new(x, y))
1453            .unwrap();
1454
1455        assert_abs_diff_eq!(model_none.intercept(), model_zero.intercept());
1456        assert!(model_none.params().abs_diff_eq(model_zero.params(), 1e-6));
1457    }
1458
1459    #[test]
1460    fn offset_changes_model() {
1461        let x = array![[-1.0], [-0.01], [0.01], [1.0]];
1462        let y = array![0, 0, 1, 1];
1463
1464        let model_none = LogisticRegression::default()
1465            .fit(&Dataset::new(x.clone(), y.clone()))
1466            .unwrap();
1467
1468        let model_offset = LogisticRegression::default()
1469            .offset(array![1.0, 1.0, -1.0, -1.0])
1470            .fit(&Dataset::new(x, y))
1471            .unwrap();
1472
1473        assert!(
1474            !model_none.params().abs_diff_eq(model_offset.params(), 1e-3),
1475            "Offset should change the learned parameters"
1476        );
1477    }
1478}