Add grpnet and notebook

adelie/__init__.py

@@ -2,4 +2,5 @@
 from . import bcd
 from . import matrix
 from . import state
-from . import solver
+from . import solver
+from .solver import grpnet

adelie/solver.py

@@ -1,5 +1,6 @@
 from . import adelie_core as core
 from . import logger
+from . import matrix
 import adelie as ad
 import numpy as np
 
@@ -165,13 +166,138 @@ def solve_basil(
     return state
 
 
-def grpnet():
-    """
-    TODO:
-    - cap max_strong_size by number of features
-    - decreasing order of lmdas
-    - cap max number of lambdas per iter
-    - alpha <= 0: add all variables to strong_set
-    - alpha > 0: add all variables with penalty <= 0 to strong_set
+def grpnet(
+    *,
+    X: np.ndarray | matrix.base,
+    y: np.ndarray,
+    groups: np.ndarray,
+    group_sizes: np.ndarray,
+    alpha: float,
+    penalty: np.ndarray,
+    lmda_path: np.ndarray =None,
+    max_iters: int =int(1e5),
+    tol: float =1e-12,
+    rsq_tol: float =0.9,
+    rsq_slope_tol: float =1e-3,
+    rsq_curv_tol: float =1e-3,
+    newton_tol: float =1e-12,
+    newton_max_iters: int =1000,
+    n_threads: int =1,
+    early_exit: bool =True,
+    intercept: bool =True,
+    strong_rule: str ="default",
+    min_ratio: float =1e-2,
+    lmda_path_size: int =100,
+    delta_lmda_path_size: int =5,
+    delta_strong_size: int =5,
+    max_strong_size: int =None,
+    check_state: bool =True,
+):
+    """TODO
     """
-    pass
+
+    if isinstance(X, np.ndarray):
+        X_raw = X
+        X = ad.matrix.naive_dense(X_raw, n_threads=n_threads)
+        _X = X.internal()
+    else:
+        assert (
+            isinstance(X, matrix.MatrixNaiveBase64) or
+            isinstance(X, matrix.MatrixNaiveBase32)
+        )
+        _X = X
+
+    dtype = (
+        np.float64
+        if isinstance(_X, matrix.MatrixNaiveBase64) else
+        np.float32
+    )
+
+    p = _X.cols()
+    G = len(groups)
+
+    actual_lmda_path_size = (
+        lmda_path_size
+        if lmda_path is None else
+        len(lmda_path)
+    )
+    delta_lmda_path_size = np.minimum(delta_lmda_path_size, actual_lmda_path_size)
+
+    max_strong_size = np.minimum(max_strong_size, G)
+
+    X_means = np.empty(p, dtype=dtype)
+    _X.means(X_means)
+
+    X_group_norms = np.empty(G, dtype=dtype)
+    _X.group_norms(
+        groups,
+        group_sizes,
+        X_means,
+        intercept,
+        X_group_norms,
+    )
+
+    y_mean = np.mean(y)
+    yc = y
+    if intercept:
+        yc = yc - y_mean
+    y_var = np.sum(yc ** 2)
+    resid = yc
+
+    strong_set = np.arange(G)[(penalty <= 0) | (alpha <= 0)]
+    strong_beta = np.zeros(np.sum(group_sizes[strong_set]), dtype=dtype)
+    strong_is_active = np.ones(strong_set.shape[0], dtype=bool)
+
+    rsq = 0
+    lmda = np.inf
+    grad = np.empty(p, dtype=dtype)
+    _X.bmul(0, p, resid, grad)
+
+    if not (lmda_path is None):
+        lmda_path = np.flip(np.sort(lmda_path))
+
+    state = ad.state.basil_naive(
+        X=X,
+        X_means=X_means,
+        X_group_norms=X_group_norms,
+        y_mean=y_mean,
+        y_var=y_var,
+        resid=resid,
+        groups=groups,
+        group_sizes=group_sizes,
+        alpha=alpha,
+        penalty=penalty,
+        strong_set=strong_set,
+        strong_beta=strong_beta,
+        strong_is_active=strong_is_active,
+        rsq=rsq,
+        lmda=lmda,
+        grad=grad,
+        lmda_path=lmda_path,
+        lmda_max=None,
+        edpp_safe_set=None,
+        edpp_v1_0=None,
+        edpp_resid_0=None,
+        max_iters=max_iters,
+        tol=tol,
+        rsq_tol=rsq_tol,
+        rsq_slope_tol=rsq_slope_tol,
+        rsq_curv_tol=rsq_curv_tol,
+        newton_tol=newton_tol,
+        newton_max_iters=newton_max_iters,
+        n_threads=n_threads,
+        early_exit=early_exit,
+        intercept=intercept,
+        strong_rule=strong_rule,
+        min_ratio=min_ratio,
+        lmda_path_size=lmda_path_size,
+        delta_lmda_path_size=delta_lmda_path_size,
+        delta_strong_size=delta_strong_size,
+        max_strong_size=max_strong_size,
+    )
+
+    if check_state:
+        assert isinstance(X_raw, np.ndarray)
+        state.check(X_raw, y, method="assert")
+
+    return solve_basil(state)

adelie/src/include/adelie_core/matrix/matrix_cov_base.hpp

@@ -9,8 +9,8 @@ class MatrixCovBase
 {
 public:
     using value_t = ValueType;
-    using rowvec_t = util::rowvec_type<value_t>;
-    using colmat_t = util::colmat_type<value_t>;
+    using vec_value_t = util::rowvec_type<value_t>;
+    using colmat_value_t = util::colmat_type<value_t>;
 
     virtual ~MatrixCovBase() {}
 
@@ -26,8 +26,8 @@ class MatrixCovBase
      */
     virtual void bmul(
         int i, int j, int p, int q, 
-        const Eigen::Ref<const rowvec_t>& v, 
-        Eigen::Ref<rowvec_t> out
+        const Eigen::Ref<const vec_value_t>& v, 
+        Eigen::Ref<vec_value_t> out
     ) =0;
 
     /**
@@ -41,7 +41,7 @@ class MatrixCovBase
      */
     virtual void to_dense(
         int i, int j, int p, int q,
-        Eigen::Ref<colmat_t> out
+        Eigen::Ref<colmat_value_t> out
     ) const =0;
 
     /**

adelie/src/include/adelie_core/matrix/matrix_cov_dense.hpp

@@ -12,8 +12,8 @@ class MatrixCovDense: public MatrixCovBase<typename DenseType::Scalar>
     using base_t = MatrixCovBase<typename DenseType::Scalar>;
     using dense_t = DenseType;
     using typename base_t::value_t;
-    using typename base_t::rowvec_t;
-    using typename base_t::colmat_t;
+    using typename base_t::vec_value_t;
+    using typename base_t::colmat_value_t;
 
 private:
     const Eigen::Map<const dense_t> _mat;   // underlying dense matrix
@@ -32,8 +32,8 @@ class MatrixCovDense: public MatrixCovBase<typename DenseType::Scalar>
 
     void bmul(
         int i, int j, int p, int q, 
-        const Eigen::Ref<const rowvec_t>& v, 
-        Eigen::Ref<rowvec_t> out
+        const Eigen::Ref<const vec_value_t>& v, 
+        Eigen::Ref<vec_value_t> out
     ) override
     {
         auto outm = out.matrix();
@@ -48,7 +48,7 @@ class MatrixCovDense: public MatrixCovBase<typename DenseType::Scalar>
 
     void to_dense(
         int i, int j, int p, int q,
-        Eigen::Ref<colmat_t> out
+        Eigen::Ref<colmat_value_t> out
     ) const override
     {
         out = _mat.block(i, j, p, q);

adelie/src/include/adelie_core/matrix/matrix_cov_lazy.hpp

@@ -14,8 +14,8 @@ class MatrixCovLazy:
 public: 
     using base_t = MatrixCovBase<typename std::decay_t<DenseType>::Scalar>;
     using typename base_t::value_t;
-    using typename base_t::rowvec_t;
-    using typename base_t::colmat_t;
+    using typename base_t::vec_value_t;
+    using typename base_t::colmat_value_t;
     using dense_t = DenseType;
     using index_t = Eigen::Index;
 
@@ -72,8 +72,8 @@ class MatrixCovLazy:
 
     void bmul(
         int i, int j, int p, int q, 
-        const Eigen::Ref<const rowvec_t>& v, 
-        Eigen::Ref<rowvec_t> out
+        const Eigen::Ref<const vec_value_t>& v, 
+        Eigen::Ref<vec_value_t> out
     ) override
     {
         if (i < 0 || j < 0 || p <= 0 || q <= 0) {
@@ -96,7 +96,7 @@ class MatrixCovLazy:
 
     void to_dense(
         int i, int j, int p, int q,
-        Eigen::Ref<colmat_t> out
+        Eigen::Ref<colmat_value_t> out
     ) const override
     {
         if (i < 0 || i > static_cast<int>(_index_map.size())) {

adelie/src/include/adelie_core/matrix/matrix_naive_base.hpp

@@ -4,13 +4,15 @@
 namespace adelie_core {
 namespace matrix {
 
-template <class ValueType> 
+template <class ValueType, class IndexType=int> 
 class MatrixNaiveBase
 {
 public:
     using value_t = ValueType;
-    using rowvec_t = util::rowvec_type<value_t>;
-    using colmat_t = util::colmat_type<value_t>;
+    using index_t = IndexType;
+    using vec_value_t = util::rowvec_type<value_t>;
+    using vec_index_t = util::rowvec_type<index_t>;
+    using colmat_value_t = util::colmat_type<value_t>;
 
     virtual ~MatrixNaiveBase() {}
 
@@ -23,7 +25,7 @@ class MatrixNaiveBase
      */
     virtual value_t cmul(
         int j, 
-        const Eigen::Ref<const rowvec_t>& v
+        const Eigen::Ref<const vec_value_t>& v
     ) const =0;
 
     /**
@@ -36,7 +38,7 @@ class MatrixNaiveBase
     virtual void ctmul(
         int j, 
         value_t v, 
-        Eigen::Ref<rowvec_t> out
+        Eigen::Ref<vec_value_t> out
     ) const =0;
 
     /**
@@ -49,8 +51,8 @@ class MatrixNaiveBase
      */
     virtual void bmul(
         int j, int q, 
-        const Eigen::Ref<const rowvec_t>& v, 
-        Eigen::Ref<rowvec_t> out
+        const Eigen::Ref<const vec_value_t>& v, 
+        Eigen::Ref<vec_value_t> out
     ) =0;
 
     /**
@@ -63,8 +65,8 @@ class MatrixNaiveBase
      */
     virtual void btmul(
         int j, int q, 
-        const Eigen::Ref<const rowvec_t>& v, 
-        Eigen::Ref<rowvec_t> out
+        const Eigen::Ref<const vec_value_t>& v, 
+        Eigen::Ref<vec_value_t> out
     ) =0;
 
     /**
@@ -76,7 +78,32 @@ class MatrixNaiveBase
      */
     virtual void to_dense(
         int j, int q,
-        Eigen::Ref<colmat_t> out
+        Eigen::Ref<colmat_value_t> out
+    ) const =0;
+
+    /**
+     * @brief Computes column-wise mean.
+     * 
+     * @param out   resulting column means.
+     */
+    virtual void means(
+        Eigen::Ref<vec_value_t> out
+    ) const =0;
+
+    /**
+     * @brief Computes group-wise column norms.
+     * 
+     * @param groups    mapping group number to starting position.
+     * @param group_sizes   mapping group number to group size.
+     * @param center        true to compute centered column norms.
+     * @param out           resulting group-wise column norms.
+     */
+    virtual void group_norms(
+        const Eigen::Ref<const vec_index_t>& groups,
+        const Eigen::Ref<const vec_index_t>& group_sizes,
+        const Eigen::Ref<const vec_value_t>& means,
+        bool center,
+        Eigen::Ref<vec_value_t> out
     ) const =0;
 
     /**