More docstring tweaking

ppdebreuck · Apr 28, 2021 · ee73681 · ee73681
1 parent 8523396
commit ee73681
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 9 deletions.
diff --git a/modnet/models/base.py b/modnet/models/base.py
@@ -37,7 +37,10 @@ class BaseMODNetModel(ABC):
 
     """
 
-    can_return_uncertainty = False
+    can_return_uncertainty: bool = False
+    """Whether or not this model supports the ``return_unc`` parameter
+    in its `predict` method, which will enable returning uncertainties.
+    """
 
     def __init__(
         self,
@@ -211,8 +214,17 @@ def fit_preset(
         """
 
     @abstractmethod
-    def evaluate(self, test_data):
-        pass
+    def evaluate(self, test_data: MODData) -> pd.DataFrame:
+        """Evaluates the target values for the passed `MODData` and returns the corresponding loss.
+
+        Parameters:
+            test_data: A featurized and feature-selected `MODData`
+                object containing the descriptors used in training.
+
+        Returns:
+            An array containing the defined losses for the model on the passed test data.
+
+        """
 
     def save(self, filename: str):
         """Save the `MODNetModel` to filename:

diff --git a/modnet/models/ensemble.py b/modnet/models/ensemble.py
@@ -169,15 +169,15 @@ class OR only return the most probable class.
             return df_mean
 
     def evaluate(self, test_data: MODData) -> pd.DataFrame:
-        """Evaluates the target values for the passed MODData by returning the corresponding loss.
+        """Evaluates the target values for the passed `MODData` and returns the corresponding loss.
 
         Parameters:
             test_data: A featurized and feature-selected `MODData`
                 object containing the descriptors used in training.
 
-
         Returns:
-            Loss score
+            An array containing the defined losses for the model on the passed test data.
+
         """
         all_losses = np.zeros(self.n_models)
         for i, m in enumerate(self.model):

diff --git a/modnet/models/vanilla.py b/modnet/models/vanilla.py
@@ -36,6 +36,8 @@ class MODNetModel(BaseMODNetModel):
 
     """
 
+    can_return_uncertainty = False
+
     def build_model(
         self,
         targets: List,
@@ -536,15 +538,15 @@ class OR only return the most probable class.
         return predictions
 
     def evaluate(self, test_data: MODData) -> pd.DataFrame:
-        """Evaluates the target values for the passed MODData by returning the corresponding loss.
+        """Evaluates the target values for the passed `MODData` and returns the corresponding loss.
 
         Parameters:
             test_data: A featurized and feature-selected `MODData`
                 object containing the descriptors used in training.
 
-
         Returns:
-            Loss score
+            An array containing the defined losses for the model on the passed test data.
+
         """
         # prevents Nan predictions if some features are inf
         x = (