chore: fix type hints

docs/source/conf.py

@@ -56,6 +56,7 @@
 intersphinx_mapping = {
     "sklearn": ("https://scikit-learn.org/stable", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
+    "scipy": ("https://docs.scipy.org/doc/scipy/", None),
 }
 
 # Myst

sortedl1/estimators.py

@@ -60,7 +60,7 @@ def __init__(
         self.max_iter = max_iter
         self.tol = tol
 
-    def fit(self, X: ArrayLike, y: ArrayLike):
+    def fit(self, X: np.ndarray | sparse.csc_array, y: ArrayLike):
         """
         Fit the model according to the given training data.
 
@@ -112,11 +112,11 @@ def fit(self, X: ArrayLike, y: ArrayLike):
         self.lambda_ = result[2]
         self.alpha_ = result[3]
         self.n_iter_ = result[4]
-        self.n_features_in_ = np.shape(X)[1]
+        self.n_features_in_ = X.shape[1]
 
         return self
 
-    def predict(self, X: ArrayLike) -> np.ndarray:
+    def predict(self, X: ArrayLike | sparse.sparray) -> np.ndarray:
         """
         Generate predictions for new data.
 

tests/test_basic.py

@@ -1,7 +1,7 @@
 """Basic tests for the sortedl1 package."""
 import numpy as np
 from numpy.random import default_rng
-from scipy.sparse import random
+from scipy.sparse import csc_array, random
 
 from sortedl1 import Slope
 
@@ -37,7 +37,7 @@ def test_simple_sparse_problem():
     p = 3
 
     rng = np.random.default_rng(4)
-    x = random(n, p, density=0.5, random_state=rng)
+    x = csc_array(random(n, p, density=0.5, random_state=rng))
     beta = np.array([1.0, 2, -0.9])
     y = x @ beta