Basic refactor to make predictors inherit from BasePredictor class

Author: pchlenski

manify/predictors/_base.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 import torch
-from sklearn.base import BaseEstimator
+from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
 
 if TYPE_CHECKING:
     from beartype.typing import Literal
@@ -34,17 +34,28 @@ class BasePredictor(BaseEstimator, ABC):
     def __init__(
         self,
         pm: ProductManifold,
-        task: Literal["classification", "regression"],
+        task: Literal["classification", "regression", "link_prediction"],
         random_state: int | None = None,
         device: str | None = None,
     ) -> None:
         self.pm = pm
         self.task = task
         self.random_state = random_state
         self.device = pm.device if device is None else device
-        self.loss_history_: Dict[str, List[float]] = {}
+        self.loss_history_: dict[str, list[float]] = {}
         self.is_fitted_: bool = False
 
+        # Initialize appropriate base class depending on task
+        if task == "classification":
+            ClassifierMixin.__init__(self)
+        elif task == "regression":
+            RegressorMixin.__init__(self)
+        elif task == "link_prediction":
+            # For link prediction, we also use ClassifierMixin, as we think of it as binary classificaiton.
+            ClassifierMixin.__init__(self)
+        else:
+            raise ValueError(f"Unknown task type: {task}")
+
     @abstractmethod
     def fit(
         self, X: Float[torch.Tensor, "n_points n_features"], y: Float[torch.Tensor, "n_points n_classes"]

manify/predictors/decision_tree.py

@@ -8,13 +8,13 @@
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
-from sklearn.base import BaseEstimator, ClassifierMixin
 
 if TYPE_CHECKING:
     from beartype.typing import Any, Literal
     from jaxtyping import Bool, Float, Int, Real
 
 from ..manifolds import ProductManifold
+from ._base import BasePredictor
 from ._midpoint import midpoint
 
 
@@ -259,7 +259,7 @@ def __init__(
         self.right = right
 
 
-class ProductSpaceDT(BaseEstimator, ClassifierMixin):
+class ProductSpaceDT(BasePredictor):
     """Decision tree in the product space to handle hyperbolic, euclidean, and hyperspherical data."""
 
     def __init__(
@@ -274,7 +274,12 @@ def __init__(
         batch_size: int | None = None,
         n_features: Literal["d", "d_choose_2"] = "d",
         ablate_midpoints: bool = False,
+        random_state: int | None = None,
+        device: str | None = None,
     ):
+        # Initialize the base class
+        super().__init__(pm=pm, task=task, random_state=random_state, device=device)
+
         # Raise error if manifold is stereographic
         if pm.is_stereographic:
             raise ValueError("Stereographic manifolds are not supported. Use a different representation.")
@@ -637,7 +642,7 @@ def score(
             return ((self.predict(X) - y) ** 2 * sample_weight).mean()
 
 
-class ProductSpaceRF(BaseEstimator, ClassifierMixin):
+class ProductSpaceRF(BasePredictor):
     """Random Forest in the product space."""
 
     def __init__(
@@ -657,7 +662,19 @@ def __init__(
         batch_size: int | None = None,
         random_state: int | None = None,
         n_jobs: int = -1,
+        device: str | None = None,
     ):
+        # Initialize the base class
+        super().__init__(pm=pm, task=task, random_state=random_state, device=device)
+
+        # Raise error if manifold is stereographic
+        if pm.is_stereographic:
+            raise ValueError("Stereographic manifolds are not supported. Use a different representation.")
+        if task == "link_prediction":
+            raise ValueError(
+                "Link prediction is not supported for decision trees. Please use utils.link_prediction to reframe as classification"
+            )
+
         # Tree hyperparameters
         tree_kwargs: Dict[str, Any] = {}
         self.pm = tree_kwargs["pm"] = pm

manify/predictors/kappa_gcn.py

@@ -12,6 +12,7 @@
     from beartype.typing import Callable, Literal
     from jaxtyping import Float, Real
 
+from ._base import BasePredictor
 from ..manifolds import Manifold, ProductManifold
 
 # TQDM: notebook or regular
@@ -142,7 +143,7 @@ def forward(
         return AXW
 
 
-class KappaGCN(torch.nn.Module):
+class KappaGCN(torch.nn.Module, BasePredictor):
     """Implementation for the Kappa GCN.
 
     Parameters
@@ -159,8 +160,11 @@ def __init__(
         hidden_dims: list[int] | None = None,
         nonlinearity: Callable = torch.relu,
         task: Literal["classification", "regression", "link_prediction"] = "classification",
+        random_state: int | None = None,
+        device: str | None = None,
     ):
-        super().__init__()
+        torch.nn.Module.__init__(self)
+        BasePredictor.__init__(self, pm=pm, task=task, random_state=random_state, device=device)
         self.pm = pm
         self.task = task
 

manify/predictors/perceptron.py

@@ -5,16 +5,16 @@
 from typing import TYPE_CHECKING
 
 import torch
-from sklearn.base import BaseEstimator, ClassifierMixin
 
 if TYPE_CHECKING:
     from jaxtyping import Float, Int
 
 from ..manifolds import ProductManifold
+from ._base import BasePredictor
 from ._kernel import product_kernel
 
 
-class ProductSpacePerceptron(BaseEstimator, ClassifierMixin):
+class ProductSpacePerceptron(BasePredictor):
     """A product-space perceptron model for multiclass classification in the product manifold space."""
 
     def __init__(
@@ -23,7 +23,12 @@ def __init__(
         max_epochs: int = 1_000,
         patience: int = 5,
         weights: Float[torch.Tensor, "n_manifolds"] | None = None,
+        task: str = "classification",
+        random_state: int | None = None,
+        device: str | None = None,
     ):
+        # Initialize base class
+        super().__init__(pm, task=task, random_state=random_state, device=device)
         self.pm = pm  # ProductManifold instance
         self.max_epochs = max_epochs
         self.patience = patience  # Number of consecutive epochs without improvement to consider convergence

manify/predictors/svm.py

@@ -14,6 +14,7 @@
     from jaxtyping import Float, Int
 
 from ..manifolds import ProductManifold
+from ._base import BasePredictor
 from ._kernel import product_kernel
 
 
@@ -30,6 +31,9 @@ def __init__(
         task: Literal["classification", "regression"] = "classification",
         epsilon: float = 1e-5,
     ):
+        # Initialize base class
+        super().__init__(pm, task=task, random_state=random_state, device=device)
+
         self.pm = pm
         self.h_constraints = h_constraints
         self.s_constraints = s_constraints