Merge pull request #663 from alan-turing-institute/467-update-docs

Add docs lint config and update docstrings (#467, #555)

Author: sgreenbury

autoemulate/experimental/calibration/bayes.py

@@ -13,6 +13,8 @@
 
 class BayesianCalibration(TorchDeviceMixin):
     """
+    Bayesian calibration using Markov Chain Monte Carlo (MCMC).
+
     Bayesian calibration estimates the probability distribution over input parameters
     given observed data, providing uncertainty estimates.
     """
@@ -120,9 +122,7 @@ def __init__(  # noqa: PLR0913
             raise ValueError(msg)
 
     def _get_kernel(self, sampler: str, **sampler_kwargs):
-        """
-        Get the appropriate MCMC kernel based on sampler choice.
-        """
+        """Get the appropriate MCMC kernel based on sampler choice."""
         sampler = sampler.lower()
 
         if sampler == "nuts":
@@ -158,7 +158,6 @@ def model(self, predict: bool = False):
             Whether to run the model with existing samples to generate posterior
             predictive distribution. Used with `pyro.infer.Predictive`.
         """
-
         # Pre-allocate tensor for all input parameters, shape [1, n_inputs]
         param_list = []
         # Each param is either sampled (if calibrated) or set to a constant value
@@ -234,7 +233,6 @@ def run_mcmc(
         MCMC
             The Pyro MCMC object. Methods include `summary()` and `get_samples()`.
         """
-
         # Check initial param values match number of chains
 
         if initial_params is not None:

autoemulate/experimental/calibration/history_matching.py

@@ -13,7 +13,9 @@
 
 
 class HistoryMatching(TorchDeviceMixin):
-    """
+    r"""
+    History Matching class for model calibration.
+
     History matching is a model calibration method, which uses observed data to
     rule out ``implausible`` parameter values. The implausibility metric is:
 
@@ -130,8 +132,9 @@ def get_nroy(
         self, implausibility: TensorLike, x: TensorLike | None = None
     ) -> TensorLike:
         """
-        Get indices of NROY points from implausibility scores. If `x`
-        is provided, returns parameter values at NROY indices.
+        Get indices of NROY points from implausibility scores.
+
+        If `x` is provided, returns parameter values at NROY indices.
 
         Parameters
         ----------
@@ -155,8 +158,9 @@ def get_ro(
         self, implausibility: TensorLike, x: TensorLike | None = None
     ) -> TensorLike:
         """
-        Get indices of RO points from implausibility scores. If `x`
-        is provided, returns parameter values at RO indices.
+        Get indices of RO points from implausibility scores.
+
+        If `x` is provided, returns parameter values at RO indices.
 
         Parameters
         ----------
@@ -255,6 +259,8 @@ def generate_param_bounds(
 
 class HistoryMatchingWorkflow(HistoryMatching):
     """
+    History Matching Workflow class.
+
     Run history matching workflow:
     - sample parameter values to test from the current NROY parameter space
     - use emulator to rule out implausible parameters and update NROY space
@@ -321,6 +327,8 @@ def __init__(  # noqa: PLR0913 allow too many arguments since all currently requ
 
     def generate_samples(self, n: int) -> tuple[TensorLike, TensorLike]:
         """
+        Generate parameter samples and evaluate implausibility.
+
         Draw `n` samples from the simulator min/max parameter bounds and
         evaluate implausability given emulator predictions.
 

autoemulate/experimental/calibration/history_matching_dashboard.py

@@ -9,6 +9,8 @@
 
 class HistoryMatchingDashboard:
     """
+    History Matching Dashboard.
+
     Interactive dashboard for exploring history matching with UI controls that adapt
     based on selected plot type.
     """
@@ -787,7 +789,6 @@ def _plot_implausibility_radar(self, df: pd.DataFrame, impl_scores: NumpyLike):
 
     def display(self):
         """Display the dashboard."""
-
         heading = widgets.HTML(value="<h2>History Matching Dashboard</h2>")
 
         # Display the heading and instructions first

autoemulate/experimental/callbacks/early_stopping.py

@@ -5,9 +5,7 @@
 
 
 class EarlyStoppingException(Exception):
-    """
-    Custom exception to signal early stopping during training.
-    """
+    """Custom exception to signal early stopping during training."""
 
     def __init__(
         self, message: str = "Training stopped early due to early stopping criteria."
@@ -17,6 +15,8 @@ def __init__(
 
 class EarlyStopping:
     """
+    Early stopping callback for PyTorch models.
+
     Stop training early if the training loss did not improve in `patience` number of
     epochs by at least `threshold` value. Can be used inside the training loop of any
     PyTorch model.
@@ -65,12 +65,14 @@ def __init__(
         self.load_best = load_best
 
     def __getstate__(self):
+        """Return state without pickling the best model weights."""
         # Avoids having to save the module_ weights twice when pickling the model
         state = self.__dict__.copy()
         state["best_model_weights_"] = None
         return state
 
     def on_train_begin(self):
+        """Initialize early stopping parameters at the start of training."""
         if self.threshold_mode not in ["rel", "abs"]:
             raise ValueError(f"Invalid threshold mode: '{self.threshold_mode}'")
         self.misses_ = 0
@@ -91,7 +93,6 @@ def on_epoch_end(self, model: nn.Module, curr_epoch: int, curr_score: float):
         curr_score: float
             The current training loss.
         """
-
         if not self._is_score_improved(curr_score):
             self.misses_ += 1
         else:

autoemulate/experimental/compare.py

@@ -31,6 +31,15 @@
 
 
 class AutoEmulate(ConversionMixin, TorchDeviceMixin, Results):
+    """
+    Automated emulator fitting.
+
+    The AutoEmulate class is the main class of the AutoEmulate package.
+    It is used to set up and compare different emulator models on a given dataset.
+    It can also be used to summarise and visualise results, and to save and load models.
+
+    """
+
     def __init__(  # noqa: PLR0913
         self,
         x: InputLike,
@@ -48,10 +57,7 @@ def __init__(  # noqa: PLR0913
         log_level: str = "progress_bar",
     ):
         """
-        The AutoEmulate class is the main class of the AutoEmulate package.
-        It is used to set up and compare different emulator models on a given dataset.
-        It can also be used to summarise and visualise results,
-        and to save and load models.
+        Initialize the AutoEmulate class.
 
         Parameters
         ----------
@@ -143,13 +149,12 @@ def __init__(  # noqa: PLR0913
 
     @staticmethod
     def all_emulators() -> list[type[Emulator]]:
+        """Return a list of all available emulators."""
         return ALL_EMULATORS
 
     @staticmethod
     def list_emulators() -> pd.DataFrame:
-        """
-        Return a dataframe with the model_name and short_name
-        of all available emulators.
+        """Return a dataframe with model names of all available emulators.
 
         Returns
         -------
@@ -167,6 +172,7 @@ def list_emulators() -> pd.DataFrame:
     def get_models(
         self, models: list[type[Emulator] | str] | None = None
     ) -> list[type[Emulator]]:
+        """Return a list of the model classes for comparisons."""
         if models is None:
             return self.all_emulators()
 
@@ -186,6 +192,7 @@ def get_models(
     def get_transforms(
         self, transforms: list[AutoEmulateTransform | dict[str, object]]
     ) -> list[AutoEmulateTransform]:
+        """Process and return a list of transforms."""
         processed_transforms = []
         for transform in transforms:
             if isinstance(transform, dict):
@@ -201,6 +208,7 @@ def get_transforms(
     def filter_models_if_multioutput(
         self, models: list[type[Emulator]], warn: bool
     ) -> list[type[Emulator]]:
+        """Filter models to only include those that support multi-output data."""
         updated_models = []
         for model in models:
             if not model.is_multioutput():
@@ -223,6 +231,7 @@ def log_compare(  # noqa: PLR0913
         r2_score,
         rmse_score,
     ):
+        """Log the comparison results."""
         msg = (
             "Comparison results:\n"
             f"Best Model: {best_model_name}, "
@@ -236,8 +245,16 @@ def log_compare(  # noqa: PLR0913
 
     def compare(self):
         """
-        Tune hyperparameters of all emulators using the train/validation data
-        and evaluate performance of all tuned emulators on the test data.
+        Compare different models on the provided dataset.
+
+        The method will:
+        - Loop over all combinations of x and y transforms and models.
+        - Set up the tuner with the training/validation data.
+        - Tune hyperparameters for each model.
+        - Fit the best model with the tuned hyperparameters.
+        - Evaluate the performance of the best model on the test data.
+        - Log the results.
+        - Save the best model and its configuration.
         """
         tuner = Tuner(self.train_val, y=None, n_iter=self.n_iter, device=self.device)
         self.logger.info(
@@ -508,7 +525,7 @@ def save(
         path: str | Path | None = None,
         use_timestamp: bool = True,
     ) -> Path:
-        """Saves model to disk.
+        """Save model to disk.
 
         Parameters
         ----------
@@ -549,7 +566,7 @@ def save(
         return self.model_serialiser._save_model(model, filename, path)
 
     def load(self, path: str | Path) -> Emulator | Result:
-        """Loads a stored model or result from disk.
+        """Load a stored model or result from disk.
 
         Parameters
         ----------

autoemulate/experimental/data/utils.py

@@ -13,19 +13,15 @@
 
 
 class ConversionMixin:
-    """
-    Mixin class to convert input data to pytorch Datasets and DataLoaders.
-    """
+    """Mixin class to convert input data to PyTorch Datasets and DataLoaders."""
 
     @classmethod
     def _convert_to_dataset(
         cls,
         x: InputLike,
         y: InputLike | None = None,
     ) -> Dataset:
-        """
-        Convert input data to pytorch Dataset.
-        """
+        """Convert input data to PyTorch Dataset."""
         # Convert input to Dataset if not already
         if isinstance(x, np.ndarray):
             x = torch.tensor(x, dtype=torch.float32)
@@ -58,9 +54,7 @@ def _convert_to_dataloader(
         batch_size: int = 16,
         shuffle: bool = True,
     ) -> DataLoader:
-        """
-        Convert input data to pytorch DataLoaders.
-        """
+        """Convert input data to PyTorch DataLoaders."""
         if isinstance(x, DataLoader) and y is None:
             dataloader = x
         elif isinstance(x, DataLoader) and y is not None:
@@ -79,9 +73,7 @@ def _convert_to_tensors(
         y: InputLike | None = None,
         dtype: torch.dtype = torch.float32,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """
-        Convert InputLike x, y to Tensor or tuple of Tensors.
-        """
+        """Convert InputLike x, y to Tensor or tuple of Tensors."""
         dataset = cls._convert_to_dataset(x, y)
 
         # Handle Subset of TensorDataset
@@ -133,9 +125,7 @@ def _convert_to_numpy(
         x: InputLike,
         y: InputLike | None = None,
     ) -> tuple[np.ndarray, np.ndarray | None]:
-        """
-        Convert InputLike x, y to tuple of numpy arrays.
-        """
+        """Convert InputLike x, y to tuple of numpy arrays."""
         if isinstance(x, np.ndarray) and (y is None or isinstance(y, np.ndarray)):
             return x, y
 
@@ -227,6 +217,7 @@ def set_random_seed(seed: int = 42, deterministic: bool = True):
 class ValidationMixin:
     """
     Mixin class for validation methods.
+
     This class provides static methods for checking the types and shapes of
     input and output data, as well as validating specific tensor shapes.
     """
@@ -235,9 +226,9 @@ class ValidationMixin:
     def _check(x: TensorLike, y: TensorLike | None):
         """
         Check the types and shape are correct for the input data.
+
         Checks are equivalent to sklearn's check_array.
         """
-
         if not isinstance(x, TensorLike):
             raise ValueError(f"Expected x to be TensorLike, got {type(x)}")
 
@@ -271,10 +262,7 @@ def _check(x: TensorLike, y: TensorLike | None):
 
     @staticmethod
     def _check_output(output: OutputLike):
-        """
-        Check the types and shape are correct
-        for the output data.
-        """
+        """Check the types and shape are correct for the output data."""
         if not isinstance(output, OutputLike):
             raise ValueError(f"Expected OutputLike, got {type(output)}")
 
@@ -424,6 +412,8 @@ def trace(Sigma: TensorLike, d: int) -> TensorLike:
     @staticmethod
     def logdet(Sigma: TensorLike, dim: int) -> TensorLike:
         """
+        Return the log-determinant of the covariance matrix.
+
         Compute the log-determinant of the covariance matrix (D-optimal design
         criterion).
 
@@ -455,6 +445,8 @@ def logdet(Sigma: TensorLike, dim: int) -> TensorLike:
     @staticmethod
     def max_eigval(Sigma: TensorLike) -> TensorLike:
         """
+        Return the maximum eigenvalue of the covariance matrix.
+
         Compute the maximum eigenvalue of the covariance matrix (E-optimal design
         criterion).