lgbm test fixes

Author: SamoraHunter

.github/workflows/notebook-test.yml

@@ -37,6 +37,11 @@ jobs:
           sudo apt-get install -y tzdata
           sudo dpkg-reconfigure -f noninteractive tzdata
 
+      - name: Install libgomp1 for LightGBM
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libgomp1
+          
       - name: Install Python 3.10, Git, and set CA environment
         run: |
           sudo apt-get update

ml_grid/model_classes/knn_wrapper_class.py

@@ -4,7 +4,8 @@
 import numpy as np
 import pandas as pd
 import torch
-from simbsig.neighbors import KNeighborsClassifier
+from simbsig.neighbors import KNeighborsClassifier as SimbsigKNeighborsClassifier
+from sklearn.neighbors import KNeighborsClassifier as SklearnKNeighborsClassifier
 from sklearn import metrics
 import logging
 
@@ -48,18 +49,22 @@ def __init__(
         self.p = p
         self.metric = metric
         self.metric_params = metric_params
+        self._init_device = device  # Store the original device parameter
+        self.device = device
+        
+        # Auto-detect device if not specified, or validate if specified
+        self._set_device(device)
 
-        # Auto-detect device
+        self.model: Optional[Union[SimbsigKNeighborsClassifier, SklearnKNeighborsClassifier]] = None
+
+    def _set_device(self, device: Optional[str]):
+        """Helper to set the device, falling back to CPU if GPU is not available."""
         gpu_available = torch.cuda.is_available()
         if device == "gpu" and not gpu_available:
-            logging.getLogger('ml_grid').warning("GPU requested for KNNWrapper, but torch.cuda.is_available() is False. Falling back to CPU.")
+            logging.getLogger('ml_grid').warning("GPU requested for KNNWrapper, but torch.cuda is not available. Falling back to CPU.")
             self.device = "cpu"
-        elif device:
-            self.device = device
         else:
-            self.device = "gpu" if gpu_available else "cpu"
-
-        self.model: Optional[KNeighborsClassifier] = None
+            self.device = device if device else ("gpu" if gpu_available else "cpu")
 
     def fit(
         self, X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray]
@@ -75,17 +80,31 @@ def fit(
         Returns:
             KNNWrapper: The fitted estimator.
         """
-        self.model = KNeighborsClassifier(
-            n_neighbors=self.n_neighbors,
-            weights=self.weights,
-            algorithm=self.algorithm,
-            leaf_size=self.leaf_size,
-            p=self.p,
-            metric=self.metric,
-            metric_params=self.metric_params,
-            device=self.device,
-        )
-
+        # If the device is CPU, use the standard scikit-learn implementation
+        # to completely avoid any simbsig/torch/cuda calls.
+        if self.device == 'cpu':
+            logging.getLogger('ml_grid').info("Using scikit-learn's KNeighborsClassifier for CPU execution.")
+            self.model = SklearnKNeighborsClassifier(
+                n_neighbors=self.n_neighbors,
+                weights=self.weights,
+                algorithm=self.algorithm,
+                leaf_size=self.leaf_size,
+                p=self.p,
+                metric=self.metric,
+                metric_params=self.metric_params,
+            )
+        else:
+            # If GPU is intended and available, use the simbsig implementation.
+            self.model = SimbsigKNeighborsClassifier(
+                n_neighbors=self.n_neighbors,
+                weights=self.weights,
+                algorithm=self.algorithm,
+                leaf_size=self.leaf_size,
+                p=self.p,
+                metric=self.metric,
+                metric_params=self.metric_params,
+                device=self.device,
+            )
         self.model.fit(X, y)
         return self
 
@@ -97,18 +116,17 @@ def get_params(self, deep: bool = False) -> Dict[str, Any]:
                 contained subobjects that are estimators.
 
         Returns:
-            Dict[str, Any]: Parameter names mapped to their values.
+            Dict[str, Any]: Parameter names mapped to their original values.
         """
         return {
-            "device": self.device,
+            "device": self._init_device,
             "n_neighbors": self.n_neighbors,
             "weights": self.weights,
             "algorithm": self.algorithm,
             "leaf_size": self.leaf_size,
             "p": self.p,
             "metric": self.metric,
             "metric_params": self.metric_params,
-            "n_neighbors": self.n_neighbors,
         }
 
     def predict(self, X: Union[pd.DataFrame, np.ndarray]) -> np.ndarray:
@@ -158,5 +176,10 @@ def set_params(self, **parameters: Any) -> "KNNWrapper":
             KNNWrapper: The instance with updated parameters.
         """
         for parameter, value in parameters.items():
+            # Special handling for device to re-validate availability
+            if parameter == 'device':
+                # Update both the initial and current device setting
+                self._init_device = value
+                self._set_device(value)
             setattr(self, parameter, value)
         return self

ml_grid/pipeline/grid_search_cross_validate.py

@@ -8,6 +8,7 @@
 import numpy as np
 import pandas as pd
 import tensorflow as tf
+import torch
 from IPython.display import clear_output
 from numpy import absolute, mean, std
 from scikeras.wrappers import KerasClassifier
@@ -150,6 +151,15 @@ def __init__(
         if "catboost" in method_name.lower() and hasattr(current_algorithm, 'set_params'):
             ml_grid_object.logger.info("Silencing CatBoost verbose output.")
             current_algorithm.set_params(verbose=0)
+
+        # Check for GPU availability and set device for torch-based models
+        if "simbsig" in str(type(algorithm_implementation)):
+            if not torch.cuda.is_available():
+                self.logger.info("No CUDA GPU detected. Forcing simbsig model to use CPU.")
+                if hasattr(current_algorithm, 'set_params'):
+                    current_algorithm.set_params(device='cpu')
+            else:
+                self.logger.info("CUDA GPU detected. Allowing simbsig model to use GPU.")
 
         self.logger.info(f"Algorithm implementation: {algorithm_implementation}")
 
@@ -250,6 +260,9 @@ def __init__(
         # Dynamically adjust KNN parameter space for small datasets
         if "kneighbors" in method_name.lower():
             self._adjust_knn_parameters(parameter_space)
+            self.logger.info(
+                "Adjusted KNN n_neighbors parameter space to prevent errors on small CV folds."
+            )
 
         # Instantiate and run the hyperparameter grid/random search
         search = HyperparameterSearch(