MWEs

scripts/refreshers/benchmarking.py

@@ -0,0 +1,61 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from ax.service.ax_client import AxClient, ObjectiveProperties
+
+
+def branin(x1, x2):
+    # Branin function
+    a = 1.0
+    b = 5.1 / (4.0 * np.pi**2)
+    c = 5.0 / np.pi
+    r = 6.0
+    s = 10.0
+    t = 1.0 / (8.0 * np.pi)
+    return a * (x2 - b * x1**2 + c * x1 - r) ** 2 + s * (1 - t) * np.cos(x1) + s
+
+
+seed_list = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+
+# Using a list of lists to store traces
+EI_traces = []
+
+for seed in seed_list:
+    ax_client = AxClient(verbose_logging=False, random_seed=seed)
+    ax_client.create_experiment(
+        parameters=[
+            {"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
+            {"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
+        ],
+        objectives={"Branin": ObjectiveProperties(minimize=True)},
+    )
+
+    trace = []
+    for _ in range(20):
+        parameterization, trial_index = ax_client.get_next_trial()
+        x1, x2 = parameterization["x1"], parameterization["x2"]
+        results = branin(x1, x2)
+        ax_client.complete_trial(trial_index=trial_index, raw_data=results)
+        trace.append(ax_client.experiment.trials[trial_index].objective_mean)
+
+    EI_traces.append(trace)
+
+# Example analysis
+mins_EI = [np.minimum.accumulate(trace) for trace in EI_traces]
+trace_mean_EI = np.mean(mins_EI, axis=0)
+trace_std_EI = np.std(mins_EI, axis=0)
+
+# Plotting
+plt.plot(trace_mean_EI, label="Mean EI")
+plt.fill_between(
+    range(20),
+    trace_mean_EI - trace_std_EI,
+    trace_mean_EI + trace_std_EI,
+    alpha=0.3,
+    label="Std EI",
+)
+
+plt.xlabel("Trial")
+plt.ylabel("Objective Value")
+plt.legend()
+plt.title("Expected Improvement Performance on Branin Function")
+plt.show()

scripts/refreshers/featurization.py

@@ -0,0 +1,87 @@
+import numpy as np
+from ax.core.observation import ObservationFeatures
+from ax.service.ax_client import AxClient, ObjectiveProperties
+
+
+# Branin function definition
+def branin(x1, x2):
+    return (
+        (x2 - 5.1 / (4 * np.pi**2) * x1**2 + 5.0 / np.pi * x1 - 6.0) ** 2
+        + 10 * (1 - 1.0 / (8 * np.pi)) * np.cos(x1)
+        + 10
+    )
+
+
+def featurize_data(x1, x2):
+    # An example featurization, but often featurization is non-invertible
+    feat1 = (x1 + x2) ** 2
+    feat2 = (x1 - x2) ** 2
+    return feat1, feat2
+
+
+candidates = [
+    {"x1": 1.2025241239855058, "x2": 3.2867298479978357},
+    {"x1": 1.6182455989582625, "x2": 1.864807635595827},
+    {"x1": -1.6757337515082393, "x2": 4.319210309745688},
+    {"x1": -3.8917807965311626, "x2": 5.708143723440155},
+    {"x1": 3.0404671700080605, "x2": 11.538410880980019},
+]  # typically many more candidates (e.g., 1e6) to approximate the continuous space
+
+featurized_candidates = [featurize_data(c["x1"], c["x2"]) for c in candidates]
+
+
+ax_client = AxClient(verbose_logging=False, random_seed=42)
+
+# The original parameters, for context
+# parameters = [
+#     {"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
+#     {"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
+# ]
+
+featurized_parameters = [
+    {"name": "feat1", "type": "range", "bounds": [25, 625]},  # depends on features
+    {"name": "feat2", "type": "range", "bounds": [100, 400]},  # depends on features
+]
+
+ax_client.create_experiment(
+    parameters=featurized_parameters,
+    objectives={"y": ObjectiveProperties(minimize=True)},
+)
+
+num_params = len(featurized_parameters)
+
+rng = np.random.default_rng(seed=42)
+
+# Initialize random trials, sampling from predefined candidates
+for i in range(2 * num_params):
+    parameterization = rng.choice(featurized_candidates)
+    _, trial_index = ax_client.attach_trial(parameterization)
+    x1 = parameterization["feat1"]
+    x2 = parameterization["feat2"]
+    results = branin(x1, x2)
+    ax_client.complete_trial(trial_index=trial_index, raw_data=results)
+    continue
+
+# Optimization loop
+for i in range(25):
+    ax_client.fit_model()  # requires data to fit model
+    model = ax_client.generation_strategy.model
+
+    obs_feat = [
+        ObservationFeatures({"feat1": row["feat1"], "feat2": row["feat2"]})
+        for _, row in candidates.iterrows()
+    ]
+    acqf_values = np.array(
+        model.evaluate_acquisition_function(observation_features=obs_feat)
+    )
+    best_index = np.argmax(acqf_values)
+
+    best_params = candidates.iloc[best_index].to_dict()
+    result = branin(best_params["feat1"], best_params["feat2"])
+    ax_client.attach_trial(best_params)
+    ax_client.complete_trial(trial_index=i, raw_data=result)
+
+# Report the optimal composition and associated objective value
+best_parameters, metrics = ax_client.get_best_parameters()
+
+print(best_parameters)

scripts/refreshers/featurization_2.py

@@ -0,0 +1,120 @@
+import numpy as np
+import pandas as pd
+from ax.core.observation import ObservationFeatures
+from ax.modelbridge.factory import Models
+from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
+from ax.service.ax_client import AxClient, ObjectiveProperties
+
+# Set random seed for reproducibility
+rng = np.random.default_rng(seed=42)
+
+
+# Branin function definition
+def branin(x1, x2):
+    return (
+        (x2 - 5.1 / (4 * np.pi**2) * x1**2 + 5.0 / np.pi * x1 - 6.0) ** 2
+        + 10 * (1 - 1.0 / (8 * np.pi)) * np.cos(x1)
+        + 10
+    )
+
+
+def featurize_data(x1, x2):
+    # An example featurization, but often featurization is non-invertible
+    feat1 = (x1 + x2) ** 2
+    feat2 = (x1 - x2) ** 2
+    return feat1, feat2
+
+
+# Generate synthetic training data
+train_parameterizations = pd.DataFrame(
+    {"x1": np.random.uniform(-5, 10, 10), "x2": np.random.uniform(0, 15, 10)}
+)
+
+train_data_feat = pd.DataFrame(
+    [
+        featurize_data(row["x1"], row["x2"])
+        for _, row in train_parameterizations.iterrows()
+    ],
+    columns=["feat1", "feat2"],
+)
+
+train_data_feat["y"] = train_parameterizations.apply(
+    lambda row: branin(row["x1"], row["x2"]), axis=1
+)
+
+# Generate candidate data
+# typically many more candidates (e.g., 1e6) to approximate the continuous space
+num_candidates = 10000
+
+candidate_params = pd.DataFrame(
+    {
+        "x1": np.random.uniform(-5, 10, num_candidates),
+        "x2": np.random.uniform(0, 15, num_candidates),
+    }
+)
+
+candidate_features = pd.DataFrame(
+    [featurize_data(row["x1"], row["x2"]) for _, row in candidate_params.iterrows()],
+    columns=["feat1", "feat2"],
+)
+
+gs = GenerationStrategy(
+    steps=[  # skip Sobol generation step
+        GenerationStep(model=Models.BOTORCH_MODULAR, num_trials=-1, max_parallelism=3)
+    ]
+)
+
+# The original parameters, for context
+# parameters = [
+#     {"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
+#     {"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
+# ]
+
+# depends on features, sometimes needs to be estimated
+featurized_parameters = [
+    {"name": "feat1", "type": "range", "bounds": [0.0, 625.0]},
+    {"name": "feat2", "type": "range", "bounds": [0.0, 400.0]},
+]
+
+ax_client = AxClient(generation_strategy=gs, verbose_logging=False, random_seed=42)
+ax_client.create_experiment(
+    parameters=featurized_parameters,
+    objectives={"y": ObjectiveProperties(minimize=True)},
+)
+
+# Add training data to the ax client
+for _, row in train_data_feat.iterrows():
+    _, trial_index = ax_client.attach_trial(row[["feat1", "feat2"]].to_dict())
+    ax_client.complete_trial(trial_index=trial_index, raw_data=row["y"])
+
+# Optimization loop
+for _ in range(10):
+    ax_client.fit_model()
+    model = ax_client.generation_strategy.model
+
+    obs_feat = [
+        ObservationFeatures(row[["feat1", "feat2"]].to_dict())
+        for _, row in candidate_features.iterrows()
+    ]
+    acqf_values = np.array(
+        model.evaluate_acquisition_function(observation_features=obs_feat)
+    )
+    best_index = np.argmax(acqf_values)
+    best_params = candidate_params.iloc[best_index].to_dict()
+    result = branin(best_params["x1"], best_params["x2"])
+
+    best_feat = candidate_features.iloc[best_index].to_dict()
+    _, trial_index = ax_client.attach_trial(best_feat)
+    ax_client.complete_trial(trial_index=trial_index, raw_data=result)
+
+# Report the optimal composition and associated objective value
+best_features, metrics = ax_client.get_best_parameters()
+print(best_features)
+
+# Find the original parameters from the best_features using lookup
+best_params = candidate_params[
+    (candidate_features["feat1"] == best_features["feat1"])
+    & (candidate_features["feat2"] == best_features["feat2"])
+]
+
+1 + 1

scripts/refreshers/multi_task.py

@@ -0,0 +1,89 @@
+from ax.core.observation import ObservationFeatures
+from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
+from ax.modelbridge.registry import Models
+from ax.modelbridge.transforms.task_encode import TaskEncode
+from ax.modelbridge.transforms.unit_x import UnitX
+from ax.service.ax_client import AxClient, ObjectiveProperties
+from utils import measure_uniformity_A, measure_uniformity_B, set_seeds
+
+set_seeds()  # setting the random seed for reproducibility
+
+transforms = [TaskEncode, UnitX]
+
+gs = GenerationStrategy(
+    name="MultiTaskOp",
+    steps=[
+        GenerationStep(
+            model=Models.SOBOL,
+            num_trials=10,
+            model_kwargs={"deduplicate": True, "transforms": transforms},
+        ),
+        GenerationStep(
+            model=Models.BOTORCH_MODULAR,
+            num_trials=-1,
+            model_kwargs={"transforms": transforms},
+        ),
+    ],
+)
+
+ax_client = AxClient(generation_strategy=gs, random_seed=42, verbose_logging=False)
+
+ax_client.create_experiment(
+    name="MultiTaskOp",
+    parameters=[
+        {
+            "name": "Temperature",
+            "type": "range",
+            "value_type": "float",
+            "bounds": [600.0, 1100.0],
+        },
+        {
+            "name": "Pressure",
+            "type": "range",
+            "value_type": "float",
+            "bounds": [5.0, 300.0],
+        },
+        {
+            "name": "Gas_Flow",
+            "type": "range",
+            "value_type": "float",
+            "bounds": [10.0, 700.0],
+        },
+        {
+            "name": "Task",
+            "type": "choice",
+            "values": ["A", "B"],
+            "is_task": True,
+            "target_value": "B",
+        },
+    ],
+    objectives={"Uniformity": ObjectiveProperties(minimize=False)},
+)
+
+for i in range(40):
+    p, trial_index = ax_client.get_next_trial(
+        fixed_features=ObservationFeatures({"Task": "A" if i % 2 else "B"})
+    )
+
+    if p["Task"] == "A":
+        u = measure_uniformity_A(p["Temperature"], p["Pressure"], F=p["Gas_Flow"])
+    else:
+        u = measure_uniformity_B(p["Temperature"], p["Pressure"], F=p["Gas_Flow"])
+
+    ax_client.complete_trial(trial_index=trial_index, raw_data={"Uniformity": u})
+
+df = ax_client.get_trials_data_frame()
+df_A = df[df["Task"] == "A"]
+df_B = df[df["Task"] == "B"]
+
+# return the parameters as a dict for the row with the hihest uniformity
+optimal_parameters_A = df_A.loc[df_A["Uniformity"].idxmax()].to_dict()
+optimal_parameters_B = df_B.loc[df_B["Uniformity"].idxmax()].to_dict()
+
+uniformity_A = optimal_parameters_A["Uniformity"]
+uniformity_B = optimal_parameters_B["Uniformity"]
+
+print(f"Optimal parameters for reactor A: {optimal_parameters_A}")
+print(f"Optimal parameters for reactor B: {optimal_parameters_B}")
+print(f"Uniformity for reactor A: {uniformity_A}")
+print(f"Uniformity for reactor B: {uniformity_B}")