added ax to hardness multi-task

scripts/04_01_multiTaskLearning.py

@@ -248,10 +248,10 @@
 # add the correlation to the plot
 ax.text(0.1, 0.9, f"Pearson Correlation: {fltCorrelation:.2f}", transform=ax.transAxes)
 
-# calculate the r2 value
-fltR2 = fltCorrelation**2
+import scipy.stats as stats
+slope, intercept, r_value, p_value, std_err = stats.linregress(dfExp_integratedHardness_wPred["integratedHardness"],dfExp_integratedHardness_wPred["vrh_pred"])
 # add the r2 value to the plot
-ax.text(0.1, 0.8, f"R2: {fltR2:.2f}", transform=ax.transAxes)
+ax.text(0.1, 0.8, f"R2: {r_value**2:.2f}", transform=ax.transAxes)
 
 # add a grid
 ax.grid(True)

scripts/examples/ax.py

@@ -0,0 +1,118 @@
+#%%
+import numpy as np
+import pandas as pd
+from ax.service.ax_client import AxClient, ObjectiveProperties
+from ax.modelbridge.factory import Models
+from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
+from ax.core.observation import ObservationFeatures
+
+rng = np.random.default_rng(seed=42)
+
+
+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
+
+
+# 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]},
+# ]
+
+# bounds depend on features, sometimes needs to be estimated / empirical (e.g., calculated based on min/max's from candidates)
+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"])
+]
+# %%

scripts/examples/ax_multiTask_test.py

@@ -0,0 +1,107 @@
+#%%
+import numpy as np
+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
+
+
+# Function to set seeds for reproducibility
+def set_seeds(seed=42):
+    np.random.seed(seed)
+
+
+# Branin function
+def branin(x1, x2):
+    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
+
+
+# Shifted and inverted Branin function
+def shifted_inverted_branin(x1, x2):
+    return -branin(x1 + 2.5, x2 + 2.5) + 300
+
+
+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": "x1",
+            "type": "range",
+            "value_type": "float",
+            "bounds": [-5.0, 10.0],
+        },
+        {
+            "name": "x2",
+            "type": "range",
+            "value_type": "float",
+            "bounds": [0.0, 15.0],
+        },
+        {
+            "name": "Task",
+            "type": "choice",
+            "values": ["A", "B"],
+            "is_task": True,
+            "target_value": "B",
+        },
+    ],
+    objectives={"Objective": 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 = branin(p["x1"], p["x2"])
+    else:
+        u = shifted_inverted_branin(p["x1"], p["x2"])
+
+    ax_client.complete_trial(trial_index=trial_index, raw_data={"Objective": 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 highest objective value
+optimal_parameters_A = df_A.loc[df_A["Objective"].idxmax()].to_dict()
+optimal_parameters_B = df_B.loc[df_B["Objective"].idxmax()].to_dict()
+
+objective_A = optimal_parameters_A["Objective"]
+objective_B = optimal_parameters_B["Objective"]
+
+print(f"Optimal parameters for task A: {optimal_parameters_A}")
+print(f"Optimal parameters for task B: {optimal_parameters_B}")
+print(f"Objective for task A: {objective_A}")
+print(f"Objective for task B: {objective_B}")
+# %%