From ee73681501f21710949854161404338d12e217f3 Mon Sep 17 00:00:00 2001
From: Matthew Evans <git@ml-evs.science>
Date: Wed, 28 Apr 2021 17:20:01 +0100
Subject: [PATCH] More docstring tweaking

---
 modnet/models/base.py     | 18 +++++++++++++++---
 modnet/models/ensemble.py |  6 +++---
 modnet/models/vanilla.py  |  8 +++++---
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/modnet/models/base.py b/modnet/models/base.py
index 28898662..30c63131 100644
--- 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
index a9d5ed46..ae418691 100644
--- 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
index d89e6b36..bbc9045f 100644
--- 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 = (