Merge pull request #1033 from karpovantonme/docs/docstring-parameters-match-signatures

Fix docstring parameters that do not match the signatures

Author: radka-j

autoemulate/calibration/bayes.py

@@ -42,7 +42,7 @@ def __init__(
         ----------
         emulator: Emulator
             Fitted Emulator object.
-        parameters_range : dict[str, tuple[float, float]]
+        parameter_range : dict[str, tuple[float, float]]
             A dictionary mapping input parameter names to their (min, max) ranges.
         observations: dict[str, TensorLike]
             A dictionary of observations for each output.
@@ -70,10 +70,10 @@ def __init__(
         Notes
         -----
         The model assumes:
-        - Uniform priors for calibrated parameters (bounds given by `parameters_range`)
+        - Uniform priors for calibrated parameters (bounds given by `parameter_range`)
         - Gaussian likelihood with no correlation between outputs
         All non-calibrated parameters are set to a constant value. This is chosen as the
-        midpoint value of `parameters_range`.
+        midpoint value of `parameter_range`.
         """
         TorchDeviceMixin.__init__(self, device=device)
         self.parameter_range = parameter_range

autoemulate/core/results.py

@@ -185,11 +185,8 @@ def best_result(self, metric: str | Metric | None = None) -> Result:
         metric: str | Metric | None
             The name of the metric to use for comparison. If None, uses the first
             available metric found in the results. The metric should exist in the
-            test_metrics of the results.
-        metric_maximize: bool | None
-            Whether higher values are better for the metric. If None, defaults to True
-            (assumes higher is better). Set to False for metrics like RMSE or MAE where
-            lower is better.
+            test_metrics of the results. Whether higher or lower values are better is
+            taken from the metric itself (`Metric.maximize`).
 
         Returns
         -------

autoemulate/core/sensitivity_analysis.py

@@ -1236,20 +1236,20 @@ def _plot_sa_heatmap(
 
     Parameters
     ----------
-    results: pd.DataFrame
+    si_df: pd.DataFrame
         Sensitivity index dataframe with columns ['index', 'parameter',
         'output', 'value'].
     index: str
         The type of sensitivity index to plot (e.g., 'ST').
-    top_n: int | None
-        Number of top parameters to include. If None, returns all. Defaults to
-        None.
+    parameters: list[str]
+        Parameters to include, in the order they should appear as columns.
     cmap: str
         Matplotlib colormap. Defaults to 'coolwarm'.
     normalize: bool
-        Wheterto normalize values to [0, 1]. Defaults to True.
-    figsize: tuple | None
-        Figure size as (width, height) in inches. Defaults to None.
+        Whether to normalize values to [0, 1]. Defaults to True.
+    fig_size: tuple | None
+        Figure size as (width, height) in inches. If None, it is derived from
+        the heatmap layout. Defaults to None.
 
     Returns
     -------

autoemulate/emulators/gaussian_process/exact.py

@@ -81,6 +81,12 @@ def __init__(
             Input features, expected to be a 2D tensor of shape (n_samples, n_features).
         y: TensorLike
             Target values, expected to be a 2D tensor of shape (n_samples, n_tasks).
+        standardize_x: bool
+            Whether to standardize the input features with a `StandardizeTransform`
+            before fitting. Defaults to False.
+        standardize_y: bool
+            Whether to standardize the target values with a `StandardizeTransform`
+            before fitting. Defaults to True.
         likelihood_cls: type[MultitaskGaussianLikelihood]
             Likelihood class to use for the model. Defaults to
             `MultitaskGaussianLikelihood`.
@@ -370,6 +376,12 @@ def __init__(
             Input features, expected to be a 2D tensor of shape (n_samples, n_features).
         y: TensorLike
             Target values, expected to be a 2D tensor of shape (n_samples, n_tasks).
+        standardize_x: bool
+            Whether to standardize the input features with a `StandardizeTransform`
+            before fitting. Defaults to False.
+        standardize_y: bool
+            Whether to standardize the target values with a `StandardizeTransform`
+            before fitting. Defaults to True.
         likelihood_cls: type[MultitaskGaussianLikelihood]
             Likelihood class to use for the model. Defaults to
             `MultitaskGaussianLikelihood`.
@@ -389,9 +401,7 @@ def __init__(
             False, it will return the posterior distribution over the modelled function.
             Defaults to False.
         epochs: int
-            Number of training epochs.
-        activation: type[nn.Module]
-            Activation function to use in the model. Defaults to `nn.ReLU`.
+            Number of training epochs. Defaults to 50.
         lr: float
             Learning rate for the optimizer. Defaults to 2e-1.
         early_stopping: EarlyStopping | None
@@ -401,7 +411,8 @@ def __init__(
             GPU). Defaults to None.
         scheduler_cls: type[LRScheduler] | None
             Learning rate scheduler class. If None, no scheduler is used. Defaults to
-        scheduler_params: dict
+            None.
+        scheduler_params: dict | None
             Additional keyword arguments for the learning rate scheduler.
         """
         # Init device

autoemulate/learners/base.py

@@ -290,7 +290,7 @@ def query(
 
         Parameters
         ----------
-        *arg: TensorLike or None
+        x: TensorLike or None
             Optional input samples.
 
         Returns