Added extra plotting of extreme response locations and example convergence to basic example

examples/demo2d/problem/brute_force.py

@@ -25,14 +25,13 @@
 from scipy.stats import gumbel_r
 from torch.utils.data import DataLoader, RandomSampler, TensorDataset
 
+torch.set_default_dtype(torch.float64)
+
 # This allows us to run as interactive and as a module.
 if __name__ == "__main__":
     from simulator import _true_loc_func, _true_scale_func  # type: ignore[import-not-found]
 else:
     from .simulator import _true_loc_func, _true_scale_func
-
-torch.set_default_dtype(torch.float64)
-
 # for typing
 _: Any
 
@@ -49,12 +48,13 @@ class ResultsObject:
     # statistics are optional
     statistics: dict[str, float]
     samples: list[float]
+    env_data: list[float]
 
     @classmethod
-    def from_samples(cls, samples: torch.Tensor) -> "ResultsObject":
+    def from_samples(cls, samples: torch.Tensor, env_data: torch.Tensor) -> "ResultsObject":
         """Create the object directly from samples."""
         statistics = {"median": float(samples.median()), "mean": float(samples.mean())}
-        return ResultsObject(statistics=statistics, samples=samples.tolist())
+        return ResultsObject(statistics=statistics, samples=samples.tolist(), env_data=env_data.tolist())
 
 
 # %%
@@ -63,7 +63,7 @@ def _result_file_name(period_length: int) -> str:
     return f"n_sample_per_period_{period_length}"
 
 
-def collect_or_calculate_results(period_length: int, num_estimates: int = 2_000) -> torch.Tensor:
+def collect_or_calculate_results(period_length: int, num_estimates: int = 2_000) -> tuple[torch.Tensor, torch.Tensor]:
     """Return a saved result for the desired length of time if available, otherwise calculate the result.
 
     New results will also be saved within this directory.
@@ -74,30 +74,40 @@ def collect_or_calculate_results(period_length: int, num_estimates: int = 2_000)
         num_estimates: The number of brute force estimates of the QoI. A new period is drawn for each estimate.
 
     Returns:
-        The QoI values calculated for each period. Shape (num_estimates,)
+        Tuple of:
+            ERD samples: (num_estimates,) samples of the ERD for that period length. QoIs can be calculated from this.
+            X_max: (num_estimates, d) The location in the environments space that produced the ERD sample.
     """
     results_path = _results_dir / f"{_result_file_name(period_length)}.json"
 
     samples = torch.tensor([])
+    max_location = torch.tensor([])
 
     if results_path.exists():
         with results_path.open() as fp:
             results = json.load(fp)
             samples = torch.tensor(results["samples"])
+            max_location = torch.tensor(results["env_data"])
 
     # make any additional samples required
     if len(samples) < num_estimates:
-        new_samples = brute_force(period_length, num_estimates - len(samples))
+        new_samples, new_max_location = brute_force(period_length, num_estimates - len(samples))
 
         samples = torch.concat([samples, new_samples])
+        max_location = torch.concat([max_location, new_max_location])
+
         # save results:
         with results_path.open("w") as fp:
-            json.dump(asdict(ResultsObject.from_samples(samples)), fp)
+            json.dump(asdict(ResultsObject.from_samples(samples, max_location)), fp)
 
-    return samples
+    elif len(samples) > num_estimates:
+        samples = samples[:num_estimates]
+        max_location = max_location[:num_estimates]
 
+    return samples, max_location
 
-def brute_force(period_length: int, num_estimates: int = 2_000) -> torch.Tensor:
+
+def brute_force(period_length: int, num_estimates: int = 2_000) -> tuple[torch.Tensor, torch.Tensor]:
     """Produces brute force samples of the Extreme Response Distibtuion.
 
     Args:
@@ -119,7 +129,9 @@ def brute_force(period_length: int, num_estimates: int = 2_000) -> torch.Tensor:
     return _brute_force_calc(dataloader, num_estimates)
 
 
-def _brute_force_calc(dataloader: DataLoader[tuple[torch.Tensor, ...]], num_estimates: int = 2_000) -> torch.Tensor:
+def _brute_force_calc(
+    dataloader: DataLoader[tuple[torch.Tensor, ...]], num_estimates: int = 2_000
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Calculate the QOI by brute force.
 
     Args:
@@ -131,9 +143,15 @@ def _brute_force_calc(dataloader: DataLoader[tuple[torch.Tensor, ...]], num_esti
         num_estimates: The number of brute force estimates of the QoI. A new period is drawn for each estimate.
 
     Returns:
-        The QoI values calculated for each period. Shape (num_estimates,)
+        Tuple of:
+            ERD samples: (num_estimates,) samples of the ERD for that period length. QoIs can be calculated from this.
+            X_max: (num_estimates, d) The location in the environments space that produced the ERD sample.
     """
     maxs = torch.zeros(num_estimates)
+
+    _, d = next(iter(dataloader))[0].shape
+    maxs_location = torch.zeros(num_estimates, d)
+
     for i in tqdm.tqdm(range(num_estimates)):
         current_max = float("-inf")
 
@@ -144,11 +162,17 @@ def _brute_force_calc(dataloader: DataLoader[tuple[torch.Tensor, ...]], num_esti
 
             gumbel_samples: np.ndarray[tuple[int,], Any] = gumbel_r.rvs(loc=loc, scale=scale)  # type: ignore  # noqa: PGH003
 
-            current_max = max(current_max, gumbel_samples.max())
+            simulator_samples_max = gumbel_samples.max()
+            if simulator_samples_max > current_max:
+                current_max = simulator_samples_max
+
+                # Get env data corresponding to max(c_max)
+                max_index = np.argmax(gumbel_samples)
+                maxs_location[i] = torch.tensor(samples[max_index, :])
 
         maxs[i] = current_max
 
-    return maxs
+    return maxs, maxs_location
 
 
 # %%
@@ -161,7 +185,7 @@ def _brute_force_calc(dataloader: DataLoader[tuple[torch.Tensor, ...]], num_esti
     N_ENV_SAMPLES_PER_PERIOD = N_YEARS_IN_PERIOD * N_SECONDS_IN_YEAR // N_SECONDS_IN_TIME_STEP
     N_ENV_SAMPLES_PER_PERIOD = 1000
 
-    samples = collect_or_calculate_results(N_ENV_SAMPLES_PER_PERIOD, 300_000)
+    samples, x_max = collect_or_calculate_results(N_ENV_SAMPLES_PER_PERIOD, 300_000)
 
     _ = plt.hist(samples, bins=100, density=True)
     _ = plt.title(
@@ -176,6 +200,9 @@ def _brute_force_calc(dataloader: DataLoader[tuple[torch.Tensor, ...]], num_esti
     )
     plt.show()
 
+    _ = plt.scatter(x_max[:, 0], x_max[:, 1])
+    plt.show()
+
     # %%
     """Here we explore the noise in the brute force result as more samples are collected.
 

src/axtreme/plotting/doe.py

@@ -62,6 +62,7 @@ def plot_qoi_estimates_from_experiment(
     ax: None | Axes = None,
     points_between_ests: int = 1,
     name: str | None = None,
+    trial_index: int | None = None,
     **kwargs: Any,  # noqa: ANN401
 ) -> Axes:
     """Plot how the QoI estimates changes over the DoE process from a given experiment with the QoI metric attached.
@@ -72,6 +73,7 @@ def plot_qoi_estimates_from_experiment(
         points_between_ests: This should be used if multiple DoE iterations are used between qoi estimates
             (e.g if the estimate is expensive). It adjusts the scale of the x axis.
         name: optional name that should be added to the legend information for this plot
+        trial_index: If provided, only plot data up to this trial index
         kwargs: kwargs that should be passed to matplotlib. Must be applicable to `ax.plot` and `ax.fill_between`
 
     Returns:
@@ -80,8 +82,13 @@ def plot_qoi_estimates_from_experiment(
     metrics = experiment.fetch_data()
     qoi_metrics = metrics.df[metrics.df["metric_name"] == "QoIMetric"]
 
+    # Filter by trial index if provided
+    if trial_index is not None:
+        qoi_metrics = qoi_metrics[qoi_metrics["trial_index"] <= trial_index]
+
     qoi_means = qoi_metrics["mean"]
     qoi_sems = qoi_metrics["sem"]
+
     if ax is None:
         _, ax = plt.subplots()
 

tests/qoi/test_gp_brute_force_system.py

@@ -210,11 +210,13 @@ def test_qoi_brute_force_system_test(  # noqa: C901, PLR0913, PLR0912, PLR0915
     Use comparable budgets allows apples-to-apples comparison of the confidence and variance of different methods.
     """
     # bookkeeping parameters
-    jobs_output_file = output_dir / "qoi_job_results.json" if output_dir else None
-
-    # Problem constants # TODO(sw): come back with a cleaner way to do this
+    jobs_output_file = (
+        output_dir / "qoi_job_results.json" if output_dir else None
+    )  # Problem constants # TODO(sw): come back with a cleaner way to do this
     brute_force_qoi: float = float(
-        brute_force.collect_or_calculate_results(period_length=N_ENV_SAMPLES_PER_PERIOD, num_estimates=300_000).median()
+        brute_force.collect_or_calculate_results(period_length=N_ENV_SAMPLES_PER_PERIOD, num_estimates=300_000)[
+            0
+        ].median()
     )
 
     _data = env_data.collect_data()
@@ -754,7 +756,7 @@ def ground_truth_estimate(
             - samples: The samples produced by the QoIeEstimator for that run.
             - name: Name of the estimator. This is used to group these results together.
     """
-    brute_force_erd_samples = brute_force.collect_or_calculate_results(
+    brute_force_erd_samples, _ = brute_force.collect_or_calculate_results(
         period_length=N_ENV_SAMPLES_PER_PERIOD, num_estimates=300_000
     )
 

tests/qoi/test_marginal_cdf_extrapolation_system.py

@@ -198,11 +198,13 @@ def test_qoi_brute_force_system_test(  # noqa: C901, PLR0912, PLR0913, PLR0915
     """
     #### Problem parameters
     # bookkeeping parameters
-    jobs_output_file = output_dir / "qoi_job_results.json" if output_dir else None
-
-    # Problem constants # TODO(sw): come back with a cleaner way to do this
+    jobs_output_file = (
+        output_dir / "qoi_job_results.json" if output_dir else None
+    )  # Problem constants # TODO(sw): come back with a cleaner way to do this
     brute_force_qoi: float = float(
-        brute_force.collect_or_calculate_results(period_length=N_ENV_SAMPLES_PER_PERIOD, num_estimates=300_000).median()
+        brute_force.collect_or_calculate_results(period_length=N_ENV_SAMPLES_PER_PERIOD, num_estimates=300_000)[
+            0
+        ].median()
     )
 
     _data = env_data.collect_data()
@@ -813,7 +815,7 @@ def ground_truth_estimate(
             - samples: The samples produced by the QoIeEstimator for that run.
             - name: Name of the estimator. This is used to group these results together.
     """
-    brute_force_erd_samples = brute_force.collect_or_calculate_results(
+    brute_force_erd_samples, _ = brute_force.collect_or_calculate_results(
         period_length=N_ENV_SAMPLES_PER_PERIOD, num_estimates=300_000
     )