Merge pull request #91 from dnv-opensource/feature/64-importance_sampling_unit_test

Feature/64 importance sampling MarginalCDF test

Author: swinter1

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/data/importance_sampling/brute_force_solution.json

@@ -1,32 +1,48 @@
+# %%
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+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 torch.distributions import Gumbel
+from torch.distributions import Gumbel, MultivariateNormal
+from torch.utils.data import DataLoader
 
+from axtreme.data import FixedRandomSampler, ImportanceAddedWrapper, MinimalDataset
+from axtreme.eval.qoi_job import QoIJob
 from axtreme.qoi.marginal_cdf_extrapolation import MarginalCDFExtrapolation, acceptable_timestep_error, q_to_qtimestep
 
+torch.set_default_dtype(torch.float64)
+
 
 class TestMarginalCDFExtrapolation:
     """
 
     Plan:
-        - __call__:
-            - Unit test None
-            - Itegration tests:
-                - batcheed params and weights
-                - dtype:
-                    - float32 where safe
-                    - float32 where not safe
-        - _parameter_estimates:
-            - Integration test:
-                - batch produce same results as non batch
-                - 1 and multi (2) posterior samples.
-                OUT OF SCOPE:
-                    - Are importance weights working properly?
-
-
-    Other sub componets we potentially should test?
+        - without importance sampling:
+            - __call__:
+                - Unit test None
+                - Integration tests:
+                    - batched params and weights
+                    - dtype:
+                        - float32 where safe
+                        - float32 where not safe
+            - _parameter_estimates:
+                - Integration test:
+                    - batch produce same results as non batch
+                    - 1 and multi (2) posterior samples.
+        - with importance sampling:
+            - Unit test:
+                - _parameter_estimates: input and output weights are attached to the same samples
+            - System Test
+                - Show that that the QoI estimation with importance samples has less uncertainty in the result compared
+                to when using the regular approach of including the whole environment dataset.
+
+    Other sub components we potentially should test?
         - distributions with different batches get optimised/treated properly (e.g in optimisation_)
     """
 
@@ -48,7 +64,7 @@ def test_call_basic_example(
 
         Demonstrate both float32 and float64 can be used (with this short period)
 
-        Take 2 posterior samples to comfirm the shape can be supported throughout.
+        Take 2 posterior samples to confirm the shape can be supported throughout.
         """
         # Define the inputs
         quantile = torch.tensor(0.5, dtype=dtype)
@@ -75,7 +91,7 @@ def test_call_basic_example(
 
         # Calculated the expected value directly.
         # This relies on knowledge of the internals, specifically that the underlying distribution produce will be
-        # [Gumbel(0, 1e-6), Gumbel(1, 1e-6), Gumbel(2, 1e-6)]. The first two will be clipped to qauntile q= 1-finfo.eps
+        # [Gumbel(0, 1e-6), Gumbel(1, 1e-6), Gumbel(2, 1e-6)]. The first two will be clipped to quantile q= 1-finfo.eps
         #  as per the bounds of ApproximateMixture
         dist = Gumbel(torch.tensor(2, dtype=dtype), 1e-6)
         q_timestep = (dist.cdf(qoi[0]) + (1 - torch.finfo(dtype).eps) * 2) / 3
@@ -88,7 +104,7 @@ def test_call_insuffecient_numeric_precision(self, gp_passthrough_1p: GenericDet
 
         Demonstrate both float32 and float64 can be used (with this short period)
 
-        Take 2 posterior samples to comfirm the shape can be supported throughout.
+        Take 2 posterior samples to confirm the shape can be supported throughout.
         """
         # Define the inputs
         dtype = torch.float32
@@ -119,7 +135,7 @@ def test_call_insuffecient_numeric_precision(self, gp_passthrough_1p: GenericDet
     def test_parameter_estimates_env_batch_invariant(
         self, gp_passthrough_1p_sampler: IndexSampler, gp_passthrough_1p: GenericDeterministicModel
     ):
-        """Checks that batched and non-batch enve data produce the saem result.
+        """Checks that batched and non-batch env data produce the same result.
 
         Tests the when the data is split into smaller batches they are then combined into the same output.
 
@@ -159,6 +175,200 @@ def test_parameter_estimates_env_batch_invariant(
         torch.testing.assert_close(param_non_batch, param_batch)
         torch.testing.assert_close(weight_non_batch, weight_batch)
 
+    def test_parameter_estimates_consistency_of_weights(self, gp_passthrough_1p: GenericDeterministicModel):
+        """
+        Tests that the `_parameter_estimates` method in `MarginalCDFExtrapolation` correctly combines environment
+        samples and their associated importance weights, and that the outputs match expected values when using a
+        deterministic Gaussian Process (GP). The deterministic nature of the GP model leads to fully predictable
+        posterior samples, i.e. they are identical to the environment samples.
+
+        This test ensures that:
+        - Posterior samples returned by `_parameter_estimates` are correctly ordered.
+        - Importance weights are preserved and correctly matched to the associated samples.
+
+        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.
+        """
+        # MarginalCDFExtrapolation expects an iterable of env data. To use importance samples, each item is expected to
+        # be of the following form [env_samples, importance_weights], where:
+        # env_samples.shape = (batch_size,d), and importance_weights.shape = (batch_size,)
+        # Note: in practice a dataloader is typically used to achieve this.
+        env_and_importance_data = [
+            # data batch 1: [env_samples, importance_weights]
+            [torch.tensor([[1.0], [2.0]]), torch.Tensor([0.1, 0.2])],
+            # data batch 2: [env_samples, importance_weights]
+            [torch.tensor([[3.0], [4.0]]), torch.Tensor([0.3, 0.4])],
+        ]
+
+        # Run the method
+        qoi_estimator = MarginalCDFExtrapolation(
+            env_iterable=env_and_importance_data,
+            period_len=99999,  # not used in this test
+            posterior_sampler=IndexSampler(torch.Size([1])),  # draw 1 posterior sample to simplify comparison
+            quantile=torch.tensor(float("nan"), dtype=torch.float64),  # not used in this test
+            quantile_accuracy=torch.tensor(float("nan"), dtype=torch.float64),  # not used in this test
+            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
+        # For the deterministic GP gp_passthrough_1p the a dummy posterior mean function is used which sets `loc`
+        # equal to the `env_value` and `scale` to 1e-6.
+        expected_posterior = torch.tensor(
+            [[[1.0000e00, 1.0000e-06], [2.0000e00, 1.0000e-06], [3.0000e00, 1.0000e-06], [4.0000e00, 1.0000e-06]]]
+        )
+        assert torch.equal(expected_posterior, posterior_samples)
+
+        # The calculated importance weights should be the same as the input importance weights
+        assert torch.equal(torch.tensor([[0.1, 0.2, 0.3, 0.4]]), importance_weights_qoi)
+
+    @pytest.mark.system
+    @pytest.mark.non_deterministic
+    # Ruff does not allow default arguments in test functions. Using a decorator circumvents that.
+    @pytest.mark.parametrize("error_tolerance, visualise", [(1, False)])
+    def test_system_marginal_cdf_with_importance_sampling(
+        self,
+        error_tolerance: float,
+        *,
+        visualise: bool,
+    ):
+        """
+        There is variability in the specific samples in the env_data used to instantiate a QoiEstimator, which creates
+        the variability of the QoIEstimator estimate. This can be seen by inspecting how the estimates given change when
+        the QoiEstimator has been instantiated with a different dataset. This test shows how (good) importance samples
+        can reduce the variability between estimates of QoIEstimators instantiated with different env_data. The effect
+        of GP uncertainty is removed by using a deterministic model, meaning all uncertainty in the estimate comes from
+        the environment samples.
+
+        Args:
+            error_tolerance: The error allowed in assertions is multiplied by this number.
+            visualise: Bool to specify if the QoI results shall be plotted.
+
+        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.
+        """
+
+        # 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
+        def _true_underlying_func(x: torch.Tensor) -> torch.Tensor:
+            # Define loc function
+            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)
+            loc = torch.exp(dist.log_prob(x))
+
+            # Define scale function
+            scale = torch.ones(x.shape[0]) * 0.1
+
+            return torch.stack([loc, scale], 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 results in less uncertainty in
+        # the QoI compared to when 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. By using importance sampling the variance between multiple runs is significantly reduced
+        # compared to using the whole environment dataset. This is the main benefit of using importance sampling.
+        # We verify this by checking that the std of the QoI means is small with importance sampling.
+        std_qoi_means_importance_sampling = df_jobs.loc[df_jobs["dataset_name"] == "importance_sample", "mean"].std()
+        assert std_qoi_means_importance_sampling <= 0.25 * error_tolerance
+
+        # 2. The mean of the QoI means is close to the brute force solution.
+        mean_qoi_means_importance_sampling = df_jobs.loc[df_jobs["dataset_name"] == "importance_sample", "mean"].mean()
+        assert abs(mean_qoi_means_importance_sampling - brute_force_qoi) <= 0.2 * error_tolerance
+
+        if visualise:
+            _, ax = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True)
+
+            # Plot results for full env data
+            std_qoi_means_full = df_jobs.loc[df_jobs["dataset_name"] == "full", "mean"].std()
+            mean_qoi_means_full = df_jobs.loc[df_jobs["dataset_name"] == "full", "mean"].mean()
+
+            df_jobs[df_jobs["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(f"Dataset: full, mean={std_qoi_means_full:.2f}, std={mean_qoi_means_full:.2f}")
+            ax[0].legend()
+
+            # Plot results for importance sampling
+            df_jobs[df_jobs["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, "
+                f"mean={mean_qoi_means_importance_sampling:.2f}, "
+                f"std={std_qoi_means_importance_sampling:.2f}"
+            )
+            ax[0].legend()
+
 
 @pytest.mark.parametrize("dtype", [(torch.float32), (torch.float64)])
 @pytest.mark.parametrize(
@@ -175,7 +385,7 @@ def test_parameter_estimates_env_batch_invariant(
 def test_q_to_qtimestep_numerical_precision_of_timestep_conversion(
     dtype: torch.dtype, longterm_q: float, period_len: int
 ):
-    """Numerical stability of converting from period quantiles to timestep quantiles.
+    """Numerical stability of converting from period quantiles to time step quantiles.
 
     This serves as documentation to show standard operators do not cause numerical issues when converting.
     """
@@ -199,15 +409,15 @@ def test_q_to_qtimestep_numerical_precision_of_timestep_conversion(
 def test_q_to_qtimestep_numerical_precision_period_increase(q: float):
     """Estimate the numerical error introduced through this operation.
 
-    This calcutates the "round trip error" of going q_longterm -> q_timestep -> q_longterm, as this is easier to test.
+    This calculates the "round trip error" of going q_longterm -> q_timestep -> q_longterm, as this is easier to test.
     This error may be larger than conversion q_longterm -> q_timestep.
-    By default python uses flaot64 (on most machines). This has a precision of 1e-15.
+    By default python uses float64 (on most machines). This has a precision of 1e-15.
 
     NOTE:
         - abs = 1e-10: all tests pass
-        - abs = 1e-11: japprox half the tests fail.
+        - abs = 1e-11: approx half the tests fail.
 
-    By default python uses flaot64 (on most machines). This has a precision of 1e-15.
+    By default python uses float64 (on most machines). This has a precision of 1e-15.
     """
 
     period_len = int(1e13)
@@ -219,3 +429,10 @@ def test_acceptable_timestep_error_at_limits_of_precision():
     """When values reach the limits of precision check an error is thrown."""
     with pytest.raises(ValueError):
         _ = acceptable_timestep_error(0.5, int(1e6), atol=1e-10)
+
+
+# %% Can be used to get plots for test_system_marginal_cdf_with_importance_sampling
+if __name__ == "__main__":
+    # %%
+    MarginalCDF = TestMarginalCDFExtrapolation()
+    MarginalCDF.test_system_marginal_cdf_with_importance_sampling(1, visualise=True)