updated problem formulation for basic_example

examples/demo2d/problem/env_data.py

@@ -5,10 +5,48 @@
 
 import numpy as np
 import pandas as pd
+import torch
+from torch.distributions import MultivariateNormal
+
+
+def generate_and_save_data(n_samples: int = 50_000, seed: int = 42) -> None:
+    """Generate environment data using multivariate normal distribution and save to file."""
+    # Define mean and covariance (using the distribution from junk.py)
+    mean = torch.tensor([0.4, 0.4])
+    cov = torch.tensor([[0.2, 0], [0, 0.2]])
+
+    # Create distribution and sample
+    env_mvn = MultivariateNormal(mean, covariance_matrix=cov)
+    with torch.random.fork_rng():
+        _ = torch.manual_seed(seed)
+        samples = env_mvn.sample(torch.Size([n_samples]))
+
+    # Convert to numpy and filter to [0,1] range
+    env_data = samples.numpy()
+    env_data = env_data[(env_data > 0).all(axis=1)]  # remove negative values
+    env_data = env_data[(env_data < 1).all(axis=1)]  # remove values > 1
+
+    # Save to file
+    current_dir = Path(__file__).parent
+    data_dir = current_dir / "data"
+    data_dir.mkdir(exist_ok=True)
+
+    np.save(data_dir / "environment_distribution.npy", env_data)
 
 
 def collect_data() -> pd.DataFrame:
     """Returns a dataframe of the env data."""
     current_dir = Path(__file__).parent
-    numpy = np.load(current_dir / "data/environment_distribution.npy")
-    return pd.DataFrame(numpy, columns=["x1", "x2"])
+    data_file = current_dir / "data/environment_distribution.npy"
+
+    # Generate data if the data file does not exist
+    if not data_file.exists():
+        generate_and_save_data()
+
+    numpy_data = np.load(data_file)
+    return pd.DataFrame(numpy_data, columns=["x1", "x2"])
+
+
+# %%
+if __name__ == "__main__":
+    generate_and_save_data(n_samples=10_000_000, seed=42)

examples/demo2d/problem/simulator.py

@@ -8,6 +8,7 @@
 from botorch.test_functions import BraninCurrin
 from numpy.typing import NDArray
 from scipy.stats import gumbel_r
+from torch.distributions import Categorical, MixtureSameFamily, MultivariateNormal
 
 from axtreme.simulator.base import Simulator
 
@@ -17,13 +18,29 @@
 
 
 # %%
-# These are helpers for our dummy simulator, and would not be available in a real problme
+# These are helpers for our dummy simulator, and would not be available in a real problem
 def _true_loc_func(x: NDArray[np.float64]) -> NDArray[np.float64]:
-    return ((_branin_currin(torch.tensor(x)) / 20)[..., 0]).numpy()
+    # For this toy example we use a Mixture distribution of a MultivariateNormal distribution
+    dist1_mean, dist1_cov = torch.tensor([0.8, 0.8]), torch.tensor([[0.03, 0], [0, 0.03]])
+    dist2_mean, dist2_cov = torch.tensor([0.2, 0.8]), torch.tensor([[0.04, 0.01], [0.01, 0.04]])
+    dist3_mean, dist3_cov = torch.tensor([0.5, 0.2]), torch.tensor([[0.06, 0], [0, 0.06]])
+
+    locs = torch.stack([dist1_mean, dist2_mean, dist3_mean])
+    covs = torch.stack([dist1_cov, dist2_cov, dist3_cov])
+    component_dist = MultivariateNormal(loc=locs, covariance_matrix=covs)
+
+    mix = Categorical(
+        torch.ones(
+            3,
+        )
+    )
+    gmm = MixtureSameFamily(mix, component_dist)
+    return np.exp(gmm.log_prob(torch.tensor(x)).numpy())
 
 
 def _true_scale_func(x: NDArray[np.float64]) -> NDArray[np.float64]:
-    return ((_branin_currin(1 - torch.tensor(x)) * 0.6)[..., 1]).numpy()
+    # For this toy example we use a constant scale for simplicity
+    return np.ones(x.shape[0]) * 0.1
 
 
 def dummy_simulator_function(x: NDArray[np.float64]) -> NDArray[np.float64]: