mypy docstyle etc

Author: soulios-basf

molpipeline/experimental/uncertainty/__init__.py

@@ -1,3 +1,8 @@
+"""Experimental uncertainty wrappers for conformal prediction in MolPipeline.
+
+Provides CrossConformalCV and UnifiedConformalCV for robust uncertainty quantification.
+"""
+
 from molpipeline.experimental.uncertainty.conformal import (
     CrossConformalCV,
     UnifiedConformalCV,

molpipeline/experimental/uncertainty/conformal.py

@@ -37,8 +37,7 @@ def bin_targets(y: np.ndarray, n_bins: int = 10) -> np.ndarray:
 
 
 class UnifiedConformalCV(BaseEstimator):
-    """One wrapper to rule them all: conformal prediction for both classifiers and
-    regressors.
+    """One wrapper to rule them all: conformal prediction for both classifiers and regressors.
 
     Uses crepes under the hood, so you know it's sweet.
 
@@ -78,7 +77,29 @@ def __init__(
         n_jobs: int = 1,
         **kwargs: Any,
     ) -> None:
-        """Initialize UnifiedConformalCV."""
+        """Initialize UnifiedConformalCV.
+
+        Parameters
+        ----------
+        estimator : Any
+            The base estimator or pipeline to wrap.
+        mondrian : Any, optional
+            Mondrian calibration/grouping (default: False).
+        confidence_level : float, optional
+            Confidence level for prediction sets/intervals (default: 0.9).
+        estimator_type : str, optional
+            Type of estimator: 'classifier' or 'regressor' (default: 'classifier').
+        nonconformity : Any, optional
+            Nonconformity function for classification.
+        difficulty_estimator : Any, optional
+            Difficulty estimator for normalized conformal prediction (regression).
+        binning : Any, optional
+            Number of bins or binning function for Mondrian calibration (regression).
+        n_jobs : int, optional
+            Number of parallel jobs (default: 1).
+        **kwargs : Any
+            Additional keyword arguments for crepes.
+        """
         self.estimator = estimator
         self.mondrian = mondrian
         self.confidence_level = confidence_level
@@ -284,8 +305,7 @@ def predict_int(self, x: np.ndarray, confidence: float | None = None) -> Any:
 
 
 class CrossConformalCV(BaseEstimator):
-    """Cross-conformal prediction for both classifiers and regressors using
-    WrapClassifier/WrapRegressor.
+    """Cross-conformal prediction for both classifiers and regressors using WrapClassifier/WrapRegressor.
 
     Handles Mondrian (class_cond) logic as described.
 
@@ -315,7 +335,7 @@ class CrossConformalCV(BaseEstimator):
         Extra toppings for crepes.
 
     """
-    
+
     def __init__(
         self,
         estimator: Any,
@@ -328,7 +348,29 @@ def __init__(
         n_bins: int = 10,
         **kwargs: Any,
     ) -> None:
-        """Initialize CrossConformalCV."""
+        """Initialize CrossConformalCV.
+
+        Parameters
+        ----------
+        estimator : Any
+            The base estimator or pipeline to wrap.
+        n_folds : int, optional
+            Number of cross-validation folds (default: 5).
+        confidence_level : float, optional
+            Confidence level for prediction sets/intervals (default: 0.9).
+        mondrian : Any, optional
+            Mondrian calibration/grouping (default: False).
+        nonconformity : Any, optional
+            Nonconformity function for classification.
+        binning : Any, optional
+            Number of bins or binning function for Mondrian calibration (regression).
+        estimator_type : str, optional
+            Type of estimator: 'classifier' or 'regressor' (default: 'classifier').
+        n_bins : int, optional
+            Number of bins for stratified splitting in regression (default: 10).
+        **kwargs : Any
+            Additional keyword arguments for crepes.
+        """
         self.estimator = estimator
         self.n_folds = n_folds
         self.confidence_level = confidence_level

tests/test_experimental/test_uncertainty/__init__.py

@@ -1 +1,4 @@
+"""Unit tests for conformal prediction wrappers in molpipeline.experimental.uncertainty.conformal.
+"""
+
 "Uncertainty test module"

tests/test_pipeline.py

@@ -376,48 +376,51 @@ def test_calibrated_classifier(self) -> None:
         self.assertEqual(predicted_value_array.shape, (len(TEST_SMILES),))
         self.assertEqual(predicted_proba_array.shape, (len(TEST_SMILES), 2))
 
-def test_conformal_pipeline_classifier(self):
-    """Test conformal prediction with a pipeline on SMILES data."""
-    from molpipeline.experimental.uncertainty.conformal import UnifiedConformalCV, CrossConformalCV
-
-    # Use the global test data
-    smiles = TEST_SMILES
-    y = np.array(CONTAINS_OX)
-
-    # Build a pipeline: SMILES -> Mol -> MorganFP -> RF
-    smi2mol = SmilesToMol()
-    mol2morgan = MolToMorganFP(radius=2, n_bits=128)
-    rf = RandomForestClassifier(n_estimators=10, random_state=42)
-    pipeline = Pipeline([
-        ("smi2mol", smi2mol),
-        ("morgan", mol2morgan),
-        ("rf", rf)
-    ])
-
-    # Split data
-    from sklearn.model_selection import train_test_split
-    X_train, X_calib, y_train, y_calib = train_test_split(smiles, y, test_size=0.3, random_state=42)
-
-    # UnifiedConformalCV
-    cp = UnifiedConformalCV(pipeline, estimator_type="classifier")
-    cp.fit(X_train, y_train)
-    cp.calibrate(X_calib, y_calib)
-    preds = cp.predict(X_calib)
-    probs = cp.predict_proba(X_calib)
-    sets = cp.predict_conformal_set(X_calib)
-    self.assertEqual(len(preds), len(y_calib))
-    self.assertEqual(probs.shape[0], len(y_calib))
-    self.assertEqual(len(sets), len(y_calib))
-
-    # CrossConformalCV
-    ccp = CrossConformalCV(pipeline, estimator_type="classifier", n_folds=3)
-    ccp.fit(smiles, y)
-    preds_ccp = ccp.predict(smiles)
-    probs_ccp = ccp.predict_proba(smiles)
-    sets_ccp = ccp.predict_conformal_set(smiles)
-    self.assertEqual(len(preds_ccp), len(y))
-    self.assertEqual(probs_ccp.shape[0], len(y))
-    self.assertEqual(len(sets_ccp), len(y))
+    def test_conformal_pipeline_classifier(self) -> None:
+        """Test conformal prediction with a pipeline on SMILES data.
+
+        This test does not take any parameters and does not return a value.
+        """
+        from molpipeline.experimental.uncertainty.conformal import UnifiedConformalCV, CrossConformalCV
+
+        # Use the global test data
+        smiles = TEST_SMILES
+        y = np.array(CONTAINS_OX)
+
+        # Build a pipeline: SMILES -> Mol -> MorganFP -> RF
+        smi2mol = SmilesToMol()
+        mol2morgan = MolToMorganFP(radius=2, n_bits=128)
+        rf = RandomForestClassifier(n_estimators=10, random_state=42)
+        pipeline = Pipeline([
+            ("smi2mol", smi2mol),
+            ("morgan", mol2morgan),
+            ("rf", rf)
+        ])
+
+        # Split data
+        from sklearn.model_selection import train_test_split
+        X_train, X_calib, y_train, y_calib = train_test_split(smiles, y, test_size=0.3, random_state=42)
+
+        # UnifiedConformalCV
+        cp = UnifiedConformalCV(pipeline, estimator_type="classifier")
+        cp.fit(X_train, y_train)
+        cp.calibrate(X_calib, y_calib)
+        preds = cp.predict(X_calib)
+        probs = cp.predict_proba(X_calib)
+        sets = cp.predict_conformal_set(X_calib)
+        self.assertEqual(len(preds), len(y_calib))
+        self.assertEqual(probs.shape[0], len(y_calib))
+        self.assertEqual(len(sets), len(y_calib))
+
+        # CrossConformalCV
+        ccp = CrossConformalCV(pipeline, estimator_type="classifier", n_folds=3)
+        ccp.fit(smiles, y)
+        preds_ccp = ccp.predict(smiles)
+        probs_ccp = ccp.predict_proba(smiles)
+        sets_ccp = ccp.predict_conformal_set(smiles)
+        self.assertEqual(len(preds_ccp), len(y))
+        self.assertEqual(probs_ccp.shape[0], len(y))
+        self.assertEqual(len(sets_ccp), len(y))
 
 
 if __name__ == "__main__":