Adding Multi-Task ElasticNet support (rust-ml#238)

* added block coordinate descent function

* added duality_gap_mtl computation

* ENH cd pass to be consistent with bcd

* added prox operator for MTL Enet

* added helper functions for tests

* working ent mtl penalties

* bcd lower objective test pass

* added MultiTaskEnet struct

* added MTENET documentation

* added API MTENET

* added variance, z-score, conf interval for multitask ENET

* added multi-task estimators

* added tests for MTL

* added tests for Enet and MTL

* WIP: made variance params generic over the number of tasks

* added z_score and confidence_95th for MTL

* WIP make compute_variance generic over the dimension

* Replace for loops in block_coordinate_descent with general_mat_mul calls

* Bring back generic compute_intercept

* Replace manual norm calculations with norm trait calls

* Add docs and derives to multi task types

* Add example for multitask_elasticnet

* Rename shape() calls to nrows and ncols

Co-authored-by: Pierre-Antoine Bannier <[email protected]>

Author: YuhanLiin

algorithms/linfa-elasticnet/Cargo.toml

@@ -40,6 +40,6 @@ thiserror = "1.0"
 linfa = { version = "0.6.0", path = "../.." }
 
 [dev-dependencies]
-linfa-datasets = { version = "0.6.0", path = "../../datasets", features = ["diabetes"] }
+linfa-datasets = { version = "0.6.0", path = "../../datasets", features = ["diabetes", "linnerud"] }
 ndarray-rand = "0.14"
 rand_xoshiro = "0.6"

algorithms/linfa-elasticnet/examples/elasticnet.rs

@@ -5,7 +5,7 @@ fn main() -> Result<()> {
     // load Diabetes dataset
     let (train, valid) = linfa_datasets::diabetes().split_with_ratio(0.90);
 
-    // train pure LASSO model with 0.1 penalty
+    // train pure LASSO model with 0.3 penalty
     let model = ElasticNet::params()
         .penalty(0.3)
         .l1_ratio(1.0)

algorithms/linfa-elasticnet/examples/multitask_elasticnet.rs

@@ -0,0 +1,24 @@
+use linfa::prelude::*;
+use linfa_elasticnet::{MultiTaskElasticNet, Result};
+
+fn main() -> Result<()> {
+    // load Diabetes dataset
+    let (train, valid) = linfa_datasets::linnerud().split_with_ratio(0.80);
+
+    // train pure LASSO model with 0.1 penalty
+    let model = MultiTaskElasticNet::params()
+        .penalty(0.1)
+        .l1_ratio(1.0)
+        .fit(&train)?;
+
+    println!("intercept:  {}", model.intercept());
+    println!("params: {}", model.hyperplane());
+
+    println!("z score: {:?}", model.z_score());
+
+    // validate
+    let y_est = model.predict(&valid);
+    println!("predicted variance: {}", y_est.r2(&valid)?);
+
+    Ok(())
+}

algorithms/linfa-elasticnet/src/algorithm.rs

@@ -25,6 +25,8 @@ pub enum ElasticNetError {
     InvalidPenalty(f32),
     #[error("invalid tolerance {0}")]
     InvalidTolerance(f32),
+    #[error("the target can either be a vector (ndim=1) or a matrix (ndim=2)")]
+    IncorrectTargetShape,
     #[error(transparent)]
     BaseCrate(#[from] linfa::Error),
 }

algorithms/linfa-elasticnet/src/error.rs

@@ -12,19 +12,26 @@ use super::Result;
     derive(Serialize, Deserialize),
     serde(crate = "serde_crate")
 )]
-/// A verified hyper-parameter set ready for the estimation of a ElasticNet regression model
-///
-/// See [`ElasticNetParams`](crate::ElasticNetParams) for more informations.
-#[derive(Clone, Debug, PartialEq)]
-pub struct ElasticNetValidParams<F> {
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ElasticNetValidParamsBase<F, const MULTI_TASK: bool> {
     penalty: F,
     l1_ratio: F,
     with_intercept: bool,
     max_iterations: u32,
     tolerance: F,
 }
 
-impl<F: Float> ElasticNetValidParams<F> {
+/// A verified hyper-parameter set ready for the estimation of a ElasticNet regression model
+///
+/// See [`ElasticNetParams`](crate::ElasticNetParams) for more information.
+pub type ElasticNetValidParams<F> = ElasticNetValidParamsBase<F, false>;
+
+/// A verified hyper-parameter set ready for the estimation of a multi-task ElasticNet regression model
+///
+/// See [`MultiTaskElasticNetParams`](crate::MultiTaskElasticNetParams) for more information.
+pub type MultiTaskElasticNetValidParams<F> = ElasticNetValidParamsBase<F, true>;
+
+impl<F: Float, const MULTI_TASK: bool> ElasticNetValidParamsBase<F, MULTI_TASK> {
     pub fn penalty(&self) -> F {
         self.penalty
     }
@@ -46,7 +53,12 @@ impl<F: Float> ElasticNetValidParams<F> {
     }
 }
 
-/// A hyper-parameter set during construction
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ElasticNetParamsBase<F, const MULTI_TASK: bool>(
+    ElasticNetValidParamsBase<F, MULTI_TASK>,
+);
+
+/// A hyper-parameter set for Elastic-Net
 ///
 /// Configures and minimizes the following objective function:
 /// ```ignore
@@ -57,18 +69,18 @@ impl<F: Float> ElasticNetValidParams<F> {
 ///
 /// The parameter set can be verified into a
 /// [`ElasticNetValidParams`](crate::hyperparams::ElasticNetValidParams) by calling
-/// [ParamGuard::check](Self::check). It is also possible to directly fit a model with
+/// [ParamGuard::check](Self::check()). It is also possible to directly fit a model with
 /// [Fit::fit](linfa::traits::Fit::fit) which implicitely verifies the parameter set prior to the
 /// model estimation and forwards any error.
 ///
 /// # Parameters
 /// | Name | Default | Purpose | Range |
 /// | :--- | :--- | :---| :--- |
-/// | [penalty](Self::penalty) | `1.0` | Overall parameter penalty | `[0, inf)` |
-/// | [l1_ratio](Self::l1_ratio) | `0.5` | Distribution of penalty to L1 and L2 regularizations | `[0.0, 1.0]` |
-/// | [with_intercept](Self::with_intercept) | `true` | Enable intercept | `false`, `true` |
-/// | [tolerance](Self::tolerance) | `1e-4` | Absolute change of any of the parameters | `(0, inf)` |
-/// | [max_iterations](Self::max_iterations) | `1000` | Maximum number of iterations | `[1, inf)` |
+/// | [penalty](Self::penalty()) | `1.0` | Overall parameter penalty | `[0, inf)` |
+/// | [l1_ratio](Self::l1_ratio()) | `0.5` | Distribution of penalty to L1 and L2 regularizations | `[0.0, 1.0]` |
+/// | [with_intercept](Self::with_intercept()) | `true` | Enable intercept | `false`, `true` |
+/// | [tolerance](Self::tolerance()) | `1e-4` | Absolute change of any of the parameters | `(0, inf)` |
+/// | [max_iterations](Self::max_iterations()) | `1000` | Maximum number of iterations | `[1, inf)` |
 ///
 /// # Errors
 ///
@@ -105,17 +117,55 @@ impl<F: Float> ElasticNetValidParams<F> {
 /// let model = checked_params.fit(&ds)?;
 /// # Ok::<(), ElasticNetError>(())
 /// ```
-#[derive(Clone, Debug, PartialEq)]
-pub struct ElasticNetParams<F>(ElasticNetValidParams<F>);
+pub type ElasticNetParams<F> = ElasticNetParamsBase<F, false>;
+
+/// A hyper-parameter set for multi-task Elastic-Net
+///
+/// The multi-task version (Y becomes a measurement matrix) is also supported and
+/// solves the following objective function:
+/// ```ignore
+/// 1 / (2 * n_samples) * || Y - XW ||^2_F
+///     + penalty * l1_ratio * ||W||_2,1
+///     + 0.5 * penalty * (1 - l1_ratio) * ||W||^2_F
+/// ```
+///
+/// See [`ElasticNetParams`](crate::ElasticNetParams) for information on parameters and return
+/// values.
+///
+/// # Example
+///
+/// ```rust
+/// use linfa_elasticnet::{MultiTaskElasticNetParams, ElasticNetError};
+/// use linfa::prelude::*;
+/// use ndarray::array;
+///
+/// let ds = Dataset::new(array![[1.0, 0.0], [0.0, 1.0]], array![[3.0, 1.1], [2.0, 2.2]]);
+///
+/// // create a new parameter set with penalty equals `1e-5`
+/// let unchecked_params = MultiTaskElasticNetParams::new()
+///     .penalty(1e-5);
+///
+/// // fit model with unchecked parameter set
+/// let model = unchecked_params.fit(&ds)?;
+///
+/// // transform into a verified parameter set
+/// let checked_params = unchecked_params.check()?;
+///
+/// // Regenerate model with the verified parameters, this only returns
+/// // errors originating from the fitting process
+/// let model = checked_params.fit(&ds)?;
+/// # Ok::<(), ElasticNetError>(())
+/// ```
+pub type MultiTaskElasticNetParams<F> = ElasticNetParamsBase<F, true>;
 
-impl<F: Float> Default for ElasticNetParams<F> {
+impl<F: Float, const MULTI_TASK: bool> Default for ElasticNetParamsBase<F, MULTI_TASK> {
     fn default() -> Self {
         Self::new()
     }
 }
 
 /// Configure and fit a Elastic Net model
-impl<F: Float> ElasticNetParams<F> {
+impl<F: Float, const MULTI_TASK: bool> ElasticNetParamsBase<F, MULTI_TASK> {
     /// Create default elastic net hyper parameters
     ///
     /// By default, an intercept will be fitted. To disable fitting an
@@ -124,8 +174,8 @@ impl<F: Float> ElasticNetParams<F> {
     /// To additionally normalize the feature matrix before fitting, call
     /// `fit_intercept_and_normalize()` before calling `fit()`. The feature
     /// matrix will not be normalized by default.
-    pub fn new() -> ElasticNetParams<F> {
-        Self(ElasticNetValidParams {
+    pub fn new() -> ElasticNetParamsBase<F, MULTI_TASK> {
+        Self(ElasticNetValidParamsBase {
             penalty: F::one(),
             l1_ratio: F::cast(0.5),
             with_intercept: true,
@@ -134,7 +184,7 @@ impl<F: Float> ElasticNetParams<F> {
         })
     }
 
-    /// Set the overall parameter penalty parameter of the elastic net.
+    /// Set the overall parameter penalty parameter of the elastic net, otherwise known as `alpha`.
     /// Use `l1_ratio` to configure how the penalty distributed to L1 and L2
     /// regularization.
     pub fn penalty(mut self, penalty: F) -> Self {
@@ -180,8 +230,8 @@ impl<F: Float> ElasticNetParams<F> {
     }
 }
 
-impl<F: Float> ParamGuard for ElasticNetParams<F> {
-    type Checked = ElasticNetValidParams<F>;
+impl<F: Float, const MULTI_TASK: bool> ParamGuard for ElasticNetParamsBase<F, MULTI_TASK> {
+    type Checked = ElasticNetValidParamsBase<F, MULTI_TASK>;
     type Error = ElasticNetError;
 
     /// Validate the hyper parameters

algorithms/linfa-elasticnet/src/hyperparams.rs

@@ -1,7 +1,7 @@
 #![doc = include_str!("../README.md")]
 
 use linfa::Float;
-use ndarray::Array1;
+use ndarray::{Array1, Array2};
 
 #[cfg(feature = "serde")]
 use serde_crate::{Deserialize, Serialize};
@@ -11,7 +11,10 @@ mod error;
 mod hyperparams;
 
 pub use error::{ElasticNetError, Result};
-pub use hyperparams::{ElasticNetParams, ElasticNetValidParams};
+pub use hyperparams::{
+    ElasticNetParams, ElasticNetParamsBase, ElasticNetValidParams, ElasticNetValidParamsBase,
+    MultiTaskElasticNetParams, MultiTaskElasticNetValidParams,
+};
 
 #[cfg_attr(
     feature = "serde",
@@ -66,3 +69,45 @@ impl<F: Float> ElasticNet<F> {
         ElasticNetParams::new().l1_ratio(F::one())
     }
 }
+
+#[cfg_attr(
+    feature = "serde",
+    derive(Serialize, Deserialize),
+    serde(crate = "serde_crate")
+)]
+/// MultiTask Elastic Net model
+///
+/// This struct contains the parameters of a fitted multi-task elastic net model. This includes the
+/// coefficients (a 2-dimensional array), (optionally) intercept (a 1-dimensional array), duality gaps
+/// and the number of steps needed in the computation.
+///
+/// ## Model implementation
+///
+/// The block coordinate descent is widely used to solve generalized linear models optimization problems,
+/// like Group Lasso, MultiTask Ridge or MultiTask Lasso. It cycles through a group of parameters and update
+/// the groups separately, holding all the others fixed. The optimization routine stops when a criterion is
+/// satisfied (dual sub-optimality gap or change in coefficients).
+#[derive(Debug, Clone)]
+pub struct MultiTaskElasticNet<F> {
+    hyperplane: Array2<F>,
+    intercept: Array1<F>,
+    duality_gap: F,
+    n_steps: u32,
+    variance: Result<Array1<F>>,
+}
+
+impl<F: Float> MultiTaskElasticNet<F> {
+    pub fn params() -> MultiTaskElasticNetParams<F> {
+        MultiTaskElasticNetParams::new()
+    }
+
+    /// Create a multi-task ridge only model
+    pub fn ridge() -> MultiTaskElasticNetParams<F> {
+        MultiTaskElasticNetParams::new().l1_ratio(F::zero())
+    }
+
+    /// Create a multi-task Lasso only model
+    pub fn lasso() -> MultiTaskElasticNetParams<F> {
+        MultiTaskElasticNetParams::new().l1_ratio(F::one())
+    }
+}