diff --git a/scripts/refreshers/benchmarking.py b/scripts/refreshers/benchmarking.py new file mode 100644 index 00000000..f263ea19 --- /dev/null +++ b/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() diff --git a/scripts/refreshers/featurization.py b/scripts/refreshers/featurization.py new file mode 100644 index 00000000..024761ca --- /dev/null +++ b/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) diff --git a/scripts/refreshers/featurization_2.py b/scripts/refreshers/featurization_2.py new file mode 100644 index 00000000..1a976567 --- /dev/null +++ b/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 diff --git a/scripts/refreshers/multi_task.py b/scripts/refreshers/multi_task.py new file mode 100644 index 00000000..e853a0b8 --- /dev/null +++ b/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}")