fix NDArray types

src/sparselm/_utils/validation.py

@@ -4,7 +4,10 @@
 from numpy.typing import NDArray
 
 
-def _check_groups(groups: NDArray | None, n_features: int) -> None:
+def _check_groups(
+    groups: NDArray[np.floating | np.integer] | list[int | float] | None,
+    n_features: int,
+) -> None:
     """Check that groups are 1D and of the correct length.
 
     Args:
@@ -30,7 +33,9 @@ def _check_groups(groups: NDArray | None, n_features: int) -> None:
         )
 
 
-def _check_group_weights(group_weights: NDArray | None, n_groups: int) -> None:
+def _check_group_weights(
+    group_weights: NDArray[np.floating] | None, n_groups: int
+) -> None:
     """Check that group weights are 1D and of the correct length.
 
     Args:

src/sparselm/model/_adaptive_lasso.py

@@ -304,9 +304,9 @@ class AdaptiveGroupLasso(AdaptiveLasso, GroupLasso):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         alpha: float = 1.0,
-        group_weights: NDArray | None = None,
+        group_weights: NDArray[np.floating] | None = None,
         max_iter: int = 3,
         eps: float = 1e-6,
         tol: float = 1e-10,
@@ -453,7 +453,7 @@ def __init__(
         self,
         group_list: list[list[int]] | None = None,
         alpha: float = 1.0,
-        group_weights: NDArray | None = None,
+        group_weights: NDArray[np.floating] | None = None,
         max_iter: int = 3,
         eps: float = 1e-6,
         tol: float = 1e-10,
@@ -599,10 +599,10 @@ class AdaptiveSparseGroupLasso(AdaptiveLasso, SparseGroupLasso):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         l1_ratio: float = 0.5,
         alpha: float = 1.0,
-        group_weights: NDArray | None = None,
+        group_weights: NDArray[np.floating] | None = None,
         max_iter: int = 3,
         eps: float = 1e-6,
         tol: float = 1e-10,
@@ -801,10 +801,10 @@ class AdaptiveRidgedGroupLasso(AdaptiveGroupLasso, RidgedGroupLasso):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         alpha: float = 1.0,
         delta: NDArray | Sequence = (1.0,),
-        group_weights: NDArray | None = None,
+        group_weights: NDArray[np.floating] | None = None,
         max_iter: int = 3,
         eps: float = 1e-6,
         tol: float = 1e-10,

src/sparselm/model/_base.py

@@ -143,7 +143,7 @@ def fit(
         self,
         X: NDArray,
         y: NDArray,
-        sample_weight: NDArray | None = None,
+        sample_weight: NDArray[np.floating] | None = None,
         *args,
         **kwargs,
     ):
@@ -205,7 +205,7 @@ def fit(
         return self
 
     def _preprocess_data(
-        self, X: NDArray, y: NDArray, sample_weight: NDArray | None = None
+        self, X: NDArray, y: NDArray, sample_weight: NDArray[np.floating] | None = None
     ) -> tuple[NDArray, NDArray, NDArray, NDArray, NDArray]:
         """Preprocess data for fitting."""
         if sample_weight is not None:
@@ -416,7 +416,7 @@ def generate_problem(
         X: NDArray,
         y: NDArray,
         preprocess_data: bool = True,
-        sample_weight: NDArray | None = None,
+        sample_weight: NDArray[np.floating] | None = None,
     ) -> None:
         """Generate regression problem and auxiliary cvxpy objects.
 

src/sparselm/model/_lasso.py

@@ -182,9 +182,9 @@ class GroupLasso(Lasso):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         alpha: float = 1.0,
-        group_weights: NDArray | None = None,
+        group_weights: NDArray[np.floating] | None = None,
         standardize: bool = False,
         fit_intercept: bool = False,
         copy_X: bool = True,
@@ -345,7 +345,7 @@ def __init__(
         self,
         group_list: list[list[int]] | None = None,
         alpha: float = 1.0,
-        group_weights: NDArray | None = None,
+        group_weights: NDArray[np.floating] | None = None,
         standardize: bool = False,
         fit_intercept: bool = False,
         copy_X: bool = True,
@@ -409,7 +409,7 @@ def generate_problem(
         X: NDArray,
         y: NDArray,
         preprocess_data: bool = True,
-        sample_weight: NDArray | None = None,
+        sample_weight: NDArray[np.floating] | None = None,
     ) -> None:
         """Initialize cvxpy problem from the generated objective function.
 
@@ -567,10 +567,10 @@ class SparseGroupLasso(GroupLasso):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         l1_ratio: float = 0.5,
         alpha: float = 1.0,
-        group_weights: NDArray | None = None,
+        group_weights: NDArray[np.floating] | None = None,
         standardize: bool = False,
         fit_intercept: bool = False,
         copy_X: bool = True,
@@ -716,10 +716,10 @@ class RidgedGroupLasso(GroupLasso):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         alpha: float = 1.0,
         delta: NDArray | Sequence = (1.0,),
-        group_weights: NDArray | None = None,
+        group_weights: NDArray[np.floating] | None = None,
         standardize: bool = False,
         fit_intercept: bool = False,
         copy_X: bool = True,

src/sparselm/model/_miqp/_base.py

@@ -76,7 +76,7 @@ class MIQPl0(CVXRegressor, metaclass=ABCMeta):
     @abstractmethod  # force inspect.isabstract to return True
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         big_M: int = 100,
         hierarchy: list[list[int]] | None = None,
         ignore_psd_check: bool = True,

src/sparselm/model/_miqp/_best_subset.py

@@ -99,7 +99,7 @@ class BestSubsetSelection(MIQPl0):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         sparse_bound=100,
         big_M: int = 100,
         hierarchy: list[list[int]] | None = None,
@@ -219,7 +219,7 @@ class RidgedBestSubsetSelection(TikhonovMixin, BestSubsetSelection):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         sparse_bound: int = 100,
         eta: float = 1.0,
         big_M: int = 100,

src/sparselm/model/_miqp/_regularized_l0.py

@@ -119,7 +119,7 @@ class RegularizedL0(MIQPl0):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         alpha: float = 1.0,
         big_M: int = 100,
         hierarchy: list[list[int]] | None = None,
@@ -171,7 +171,7 @@ class MixedL0(RegularizedL0, metaclass=ABCMeta):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         alpha: float = 1.0,
         eta: float = 1.0,
         big_M: int = 100,
@@ -342,7 +342,7 @@ class L1L0(MixedL0):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         alpha: float = 1.0,
         eta: float = 1.0,
         big_M: int = 100,
@@ -501,7 +501,7 @@ class L2L0(TikhonovMixin, MixedL0):
 
     def __init__(
         self,
-        groups: NDArray | None = None,
+        groups: NDArray[np.floating | np.integer] | None = None,
         alpha: float = 1.0,
         eta: float = 1.0,
         big_M: int = 100,

src/sparselm/stepwise.py

@@ -151,7 +151,7 @@ def fit(
         self,
         X: NDArray,
         y: NDArray,
-        sample_weight: NDArray | None = None,
+        sample_weight: NDArray[np.floating] | None = None,
         *args,
         **kwargs,
     ):

src/sparselm/tools.py

@@ -13,8 +13,8 @@
 
 def constrain_coefficients(
     indices: NDArray,
-    high: NDArray | None = None,
-    low: NDArray | None = None,
+    high: NDArray[np.floating] | float | None = None,
+    low: NDArray[np.floating] | float | None = None,
 ):
     """Constrain a fit method to keep coefficients within a specified range.
 
@@ -98,7 +98,10 @@ def wrapped(X, y, *args, **kwargs):
 
 
 def r2_score_to_cv_error(
-    score: float, y: NDArray, y_pred: NDArray, weights: NDArray | None = None
+    score: float,
+    y: NDArray,
+    y_pred: NDArray,
+    weights: NDArray[np.floating] | None = None,
 ):
     """Convert r2 score to cross-validation error.