Add n_samples to APIs

Author: sgreenbury

autoemulate/core/compare.py

@@ -76,6 +76,7 @@ def __init__(
         log_level: str = "progress_bar",
         tuning_metric: str | Metric = "r2",
         evaluation_metrics: list[str | Metric] | None = None,
+        n_samples: int = 1000,
     ):
         """
         Initialize the AutoEmulate class.
@@ -135,6 +136,9 @@ def __init__(
             Each entry can be a string shortcut or a MetricConfig object.
             IMPORTANT: The first metric in the list is used to
             determine the best model.
+        n_samples: int
+            Number of samples to generate to predict mean when emulator does not have a
+            mean directly available. Defaults to 1000.
         """
         Results.__init__(self)
         self.random_seed = random_seed
@@ -192,6 +196,7 @@ def __init__(
         # Set up logger and ModelSerialiser for saving models
         self.logger, self.progress_bar = get_configured_logger(log_level)
         self.model_serialiser = ModelSerialiser(self.logger)
+        self.n_samples = n_samples
 
         # Run compare
         self.compare()
@@ -489,6 +494,7 @@ def compare(self):
                                 n_bootstraps=self.n_bootstraps,
                                 device=self.device,
                                 metrics=self.evaluation_metrics,
+                                n_samples=self.n_samples,
                             )
                             test_metrics = bootstrap(
                                 transformed_emulator,
@@ -497,6 +503,7 @@ def compare(self):
                                 n_bootstraps=self.n_bootstraps,
                                 device=self.device,
                                 metrics=self.evaluation_metrics,
+                                n_samples=self.n_samples,
                             )
 
                             # Log all test metrics from test_metrics dictionary

autoemulate/core/model_selection.py

@@ -64,6 +64,7 @@ def cross_validate(
     device: DeviceLike = "cpu",
     random_seed: int | None = None,
     metrics: list[Metric] | None = None,
+    n_samples: int = 1000,
 ):
     """
     Cross validate model performance using the given `cv` strategy.
@@ -88,6 +89,9 @@ def cross_validate(
         Optional random seed for reproducibility.
     metrics: list[TorchMetrics] | None
         List of metrics to compute. If None, uses r2 and rmse.
+    n_samples: int
+        Number of samples to generate to predict mean when emulator does not have a
+        mean directly available. Defaults to 1000.
 
     Returns
     -------
@@ -146,7 +150,7 @@ def cross_validate(
         # compute and save results
         y_pred = transformed_emulator.predict(x_val)
         for metric in metrics:
-            score = evaluate(y_pred, y_val, metric)
+            score = evaluate(y_pred, y_val, metric, n_samples=n_samples)
             cv_results[metric.name].append(score)
     return cv_results