Feature/64 importance sampling MarginalCDF test

examples/crest_heights_north_sea/importance_sampling.py

@@ -103,6 +103,12 @@ def importance_sampling_distribution_uniform_region(
             Importance sampling weights for each sample. Shape (n_samples,)
 
     Details:
+        The mathematical justification for this algorithm is given in
+
+        "Efficient Long-Term Structural Reliability Estimation with Non-Gaussian
+        Stochastic Models: A Design of Experiments Approach.” arXiv, March 3, 2025.
+        https://doi.org/10.48550/arXiv.2503.01566."
+
         The algorithm works as follows:
         1. For a chosen `threshold` c for the real environment distribution.
 
@@ -152,7 +158,7 @@ def importance_sampling_distribution_uniform_region(
                 - This is a reasonable approximation if p(x_i) is considered to be close enough to 0.
 
     Todo: TODO
-    - (sw 2024_09_17): Reference the paper here (or preprint) once its published Issue #177.
+    - (ak 2025-07-09): add seeded option for testing purposes
 
     """
     uniform_dist = torch.distributions.Uniform(region[0], region[1])

tests/qoi/test_marginal_cdf_extrapolation_with_importance_sampling.py

@@ -0,0 +1,282 @@
+"""
+Tests for MarginalCDFExtrapolation with importance sampling.
+
+Plan:
+    - Unit test:
+        - _parameter_estimates: input and output weights are attached to the same samples
+    - Integration test:
+        - successful run of minimal version of the qoi using a deterministic model with importance sampling
+    - System Test
+        - Show that that the QoI estimation requires less samples with importance samples compared to when using
+        the regular approach of including the whole environment dataset.
+"""
+
+# %%
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+from botorch.models.deterministic import GenericDeterministicModel
+from botorch.sampling.index_sampler import IndexSampler
+from numpy.typing import NDArray
+from torch.distributions import MultivariateNormal
+from torch.utils.data import DataLoader
+
+from axtreme.data import FixedRandomSampler, ImportanceAddedWrapper, MinimalDataset
+from axtreme.eval.qoi_job import QoIJob
+from axtreme.qoi import MarginalCDFExtrapolation
+
+
+def test_parameter_estimates_consistency_of_weights(gp_passthrough_1p: GenericDeterministicModel):
+    """
+    Run a minimal version of the qoi using a deterministic model and a short period_len to test that the importance
+    weights are still connected to the correct samples when using MarginalCDFExtrapolation._parameter_estimates.
+
+    Args:
+        gp_passthrough_1p is defined in conftest.py. It creates a deterministic GP which always produces identical
+        posterior samples. The output location is a direct pass through of the given input data, and the scale is set
+        to 1e-6.
+
+    Notes:
+    -   _parameter_estimates needs the importance dataset with specific dimensions by using ImportanceAddedWrapper and
+        DataLoader the correct format is ensured and in addition it can be tested that these two functions do not change
+        the relation between samples and weights.
+    """
+    # Define simple importance samples and weights
+    importance_samples = torch.Tensor([[1], [2], [3], [4]])
+    importance_weights = torch.Tensor([0.1, 0.2, 0.3, 0.4])
+
+    # Wrap them together
+    importance_dataset = ImportanceAddedWrapper(MinimalDataset(importance_samples), MinimalDataset(importance_weights))
+
+    # The batch size should be smaller than the number of input samples to ensure that the importance_dataset gets
+    # split up into smaller batches. By doing this we can test that the batching in _parameter_estimates does not mess
+    # up the relation between samples and weights.
+    dataloader = DataLoader(importance_dataset, batch_size=2)
+
+    # Run the method
+    qoi_estimator = MarginalCDFExtrapolation(
+        env_iterable=dataloader,
+        period_len=10,
+        posterior_sampler=IndexSampler(torch.Size([1])),  # draw 1 posterior samples to simplify comparison
+        quantile=torch.tensor(0.5, dtype=torch.float64),
+        quantile_accuracy=torch.tensor(0.1, dtype=torch.float64),
+        dtype=torch.float64,
+    )
+
+    # Get posterior samples and connected weights for given QoI estimator
+    posterior_samples, importance_weights_qoi = qoi_estimator._parameter_estimates(gp_passthrough_1p)
+
+    # As a deterministic GP is used the posterior samples should not only have the same values as the
+    # input samples but their order should be the same as well
+    assert torch.equal(torch.flatten(importance_samples), posterior_samples[..., 0][0])
+    assert torch.equal(torch.flatten(importance_weights), torch.flatten(importance_weights_qoi))
+
+
+def test_qoi_runs_with_importance_sampling(gp_passthrough_1p: GenericDeterministicModel):
+    """
+    Run a minimal version of the qoi using a deterministic model and a short period_len to test that it successfully
+    runs with importance sampling.
+
+    Args:
+        gp_passthrough_1p is defined in conftest.py. It creates a deterministic GP which always produces identical
+        posterior samples. The output location is a direct pass through of the given input data, and the scale is set
+        to 1e-6.
+
+    Notes:
+    -   The difference to test_parameter_estimates_consistency_of_weights is that the model is passed directly
+        to the QoI estimator not qoi_estimator._parameter_estimates.
+    """
+    # Define simple importance samples and weights
+    importance_samples = torch.Tensor([[1], [2], [3], [4]])
+    importance_weights = torch.Tensor([0.1, 0.2, 0.3, 0.4])
+
+    # Wrap them in a dataloader
+    importance_dataset = ImportanceAddedWrapper(MinimalDataset(importance_samples), MinimalDataset(importance_weights))
+    dataloader = DataLoader(importance_dataset, batch_size=10)
+
+    # Run the method
+    qoi_estimator = MarginalCDFExtrapolation(
+        env_iterable=dataloader,
+        period_len=10,
+        posterior_sampler=IndexSampler(torch.Size([1])),  # draw 1 posterior samples to simplify comparison
+        quantile=torch.tensor(0.5, dtype=torch.float64),
+        quantile_accuracy=torch.tensor(0.1, dtype=torch.float64),
+        dtype=torch.float64,
+    )
+
+    # This only tests if the functions runs successfully and does not concern itself with the output
+    _ = qoi_estimator(gp_passthrough_1p)
+
+
+# These are helper functions to define the loc and scale which are used in the simulator used for the system test
+def _true_loc_func(x: NDArray[np.float64]) -> NDArray[np.float64]:
+    # For this toy example we use a MultivariateNormal distribution
+    dist_mean, dist_cov = torch.tensor([1, 1]), torch.tensor([[0.03, 0], [0, 0.03]])
+
+    dist = MultivariateNormal(loc=dist_mean, covariance_matrix=dist_cov)
+
+    return np.exp(dist.log_prob(torch.tensor(x)).numpy())
+
+
+def _true_scale_func(x: NDArray[np.float64]) -> NDArray[np.float64]:
+    # For this toy example we use a constant scale for simplicity
+    return np.ones(x.shape[0]) * 0.1
+
+
+# Ruff does not allow default arguments in test functions. Using a decorator circumvents that.
+@pytest.mark.parametrize("params", [(1, False)])
+def test_system_marginal_cdf_with_importance_sampling(
+    params: tuple[float, bool],
+) -> tuple[pd.DataFrame, float] | None:
+    """
+    Show that that the QoI estimation requires less samples with importance samples compared to when using
+    the regular approach of including the whole environment dataset.
+
+    Args:
+        params: tuple of
+            - The error allowed in assertions is multiplied by this number.
+            - Bool to specify if the QoI results shall be returned for further testing or plotting.
+
+    Returns:
+        If return_results==True: the QoI results are returned in a pandas Dataframe
+
+    A deterministic GP is used to remove most of the uncertainty related to the GP. Some uncertainty remains as the GP
+    is not fit on the whole dataset.
+
+    Expectation:
+    As we use a deterministic GP the QoI estimator will not be a distribution but a point representing the mean.
+    The std of the means of several runs of the QoI estimator should be lower with uncertainty sampling than without.
+    By visual inspection a threshold for the std for importance sampling is chosen.
+    """
+    error_tolerance, return_results = params
+
+    # Load precalculated importance samples and weights
+    importance_samples = torch.load(
+        Path(__file__).parent / "data" / "importance_sampling" / "importance_samples.pt", weights_only=True
+    )
+    importance_weights = torch.load(
+        Path(__file__).parent / "data" / "importance_sampling" / "importance_weights.pt", weights_only=True
+    )
+
+    importance_dataset = ImportanceAddedWrapper(MinimalDataset(importance_samples), MinimalDataset(importance_weights))
+
+    # Load environment data
+    # The environment is chosen such that it is concentrated in one region of the search space and only few samples
+    # cover the subregion where the extreme response occurs
+    env_data = np.load(Path(__file__).parent / "data" / "importance_sampling" / "environment_distribution.npy")
+
+    # Get brute force estimate
+    brute_force_path = Path(__file__).parent / "data" / "importance_sampling" / "brute_force_solution.json"
+    with brute_force_path.open() as file:
+        brute_force_qoi = json.load(file)["statistics"]["median"]
+
+    # Set up a deterministic GP which uses the true underlying functions of loc and scale used in the simulator
+    def true_underlying_func(x: torch.Tensor) -> torch.Tensor:
+        locs = torch.from_numpy(_true_loc_func(x.numpy())).unsqueeze(-1)
+        scales = torch.from_numpy(_true_scale_func(x.numpy())).unsqueeze(-1)
+        return torch.concat([locs, scales], dim=-1)
+
+    gp_deterministic = GenericDeterministicModel(true_underlying_func, num_outputs=2)
+
+    # Create jobs with with and without importance sampling
+    qoi_jobs = []
+    datasets = {"full": env_data, "importance_sample": importance_dataset}
+    for dataset_name, dataset in datasets.items():
+        for i in range(200):
+            # A fixed random sampler selects the same samples if the seed is the same which allows the results to be
+            # compared if this function is run multiple times
+            dataset_size = 800
+            sampler = FixedRandomSampler(dataset, num_samples=dataset_size, seed=i, replacement=True)
+            dataloader = DataLoader(dataset, sampler=sampler, batch_size=100)
+
+            qoi_estimator = MarginalCDFExtrapolation(
+                env_iterable=dataloader,
+                period_len=1_000,
+                quantile=torch.tensor(0.5),
+                quantile_accuracy=torch.tensor(0.01),
+                # IndexSampler needs to be used with GenericDeterministicModel. Each sample just selects the mean.
+                # As we use a deterministic model all posterior samples are identical and hence we only use one.
+                posterior_sampler=IndexSampler(torch.Size([1])),
+            )
+
+            qoi_jobs.append(
+                QoIJob(
+                    name=f"qoi_{dataset_name}_{dataset_size}_{i}",
+                    qoi=qoi_estimator,
+                    model=gp_deterministic,
+                    tags={
+                        "dataset_name": dataset_name,
+                        "dataset_size": dataset_size,
+                    },
+                )
+            )
+
+    jobs_output_file = None
+    qoi_results = [job(output_file=jobs_output_file) for job in qoi_jobs]
+
+    df_jobs = pd.json_normalize([item.to_dict() for item in qoi_results], max_level=1)
+    df_jobs.columns = df_jobs.columns.str.removeprefix("tags.")
+
+    # The aim of this test is to show that QoI estimation with importance sampling requires less samples
+    # (i.e. num_samples) to converge to the solution than with using the full environment dataset.
+    # All absolute values in the following asserts are chosen based on visual inspection of the QoI results.
+    # There are two criteria we use to judge the convergence:
+    # 1. the mean of the QoI means is close to the brute force solution and in particular the one with
+    # importance sampling is closer than the one without
+    assert (
+        abs(df_jobs.loc[df_jobs["dataset_name"] == "importance_sample", "mean"].mean() - brute_force_qoi)
+        <= 0.2 * error_tolerance
+    )
+
+    assert abs(df_jobs.loc[df_jobs["dataset_name"] == "importance_sample", "mean"].mean() - brute_force_qoi) <= abs(
+        df_jobs.loc[df_jobs["dataset_name"] == "full", "mean"].mean() - brute_force_qoi
+    )
+
+    # 2. the std of the QoI means is small with importance sampling
+    assert df_jobs.loc[df_jobs["dataset_name"] == "importance_sample", "mean"].std() <= 0.25 * error_tolerance
+
+    if return_results:
+        return df_jobs, brute_force_qoi
+
+    return None
+
+
+# %% Can be used to get plots for test_system_marginal_cdf_with_importance_sampling
+if __name__ == "__main__":
+    # %% Calculate QoI
+    from typing import cast
+
+    qoi_results, brute_force_qoi = cast(
+        "tuple[pd.DataFrame, float]", test_system_marginal_cdf_with_importance_sampling((1, True))
+    )
+
+    # %% Plot results
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True)
+
+    # Plot results for full env data
+    qoi_results[qoi_results["dataset_name"] == "full"].hist(column="mean", ax=ax[0], grid=False)
+    ax[0].axvline(brute_force_qoi, c="orange", label=f"Brute force ({brute_force_qoi:.2f})")
+    ax[0].set_title(
+        "Dataset: full, "
+        "mean={qoi_results.loc[qoi_results['dataset_name'] == 'full', 'mean'].mean():.2f}, "
+        "std={qoi_results.loc[qoi_results['dataset_name'] == 'full', 'mean'].std():.2f}"
+    )
+    ax[0].legend()
+
+    # Plot results for importance sampling
+    qoi_results[qoi_results["dataset_name"] == "importance_sample"].hist(column="mean", ax=ax[1], grid=False)
+    ax[1].axvline(brute_force_qoi, c="orange", label=f"Brute force ({brute_force_qoi:.2f})")
+    ax[1].set_title(
+        "Dataset: importance sample, "
+        "mean={qoi_results.loc[qoi_results['dataset_name'] == 'importance_sample', 'mean'].mean():.2f}, "
+        "std={qoi_results.loc[qoi_results['dataset_name'] == 'importance_sample', 'mean'].std():.2f}"
+    )
+    ax[0].legend()
+
+# %%