Version: 1.1.2

Ax for Materials Science

Some optimization experiments, like the one described in this tutorial, can be conducted in a completely automated manner. Other experiments may require a human in the loop, for instance a scientist manually conducting and evaluating each trial in a lab. In this tutorial we demonstrate this ask-tell optimization in a human-in-the-loop setting by imagining the task of maximizing the strength of a 3D printed part using compression testing (i.e., crushing the part) where different print settings will have to be manually tried and evaluated.

Background

In 3D printing, several parameters can significantly affect the mechanical properties of the printed object:

Infill Density: The percentage of material used inside the object. Higher infill density generally increases strength but also weight and material usage.
Layer Height: The thickness of each layer of material. Smaller layer heights can improve surface finish and detail but increase print time.
Infill Type: The pattern used to fill the interior of the object. Different patterns (e.g., honeycomb, gyroid, lines, rectilinear) offer various balances of strength, speed, and material efficiency.
Strength Measurement: In this tutorial, we assume the strength of the 3D printed part is measured using compression testing, which evaluates how the object performs under compressive stress.

Learning Objectives

Understand black box optimization concepts
Define an optimization problem using Ax
Configure and run an experiment using Ax's Client
Analyze the results of the optimization

Prerequisites

Familiarity with Python and basic programming concepts
Understanding of adaptive experimentation and Bayesian optimization

Step 1: Import Necessary Modules

from ax.api.client import Client
from ax.api.configs import  RangeParameterConfig, ChoiceParameterConfig

Step 2: Initialize Client

Create an instance of the Client to manage the state of your experiment.

client = Client()

Step 3: Configure Experiment

Define the parameters for the 3D printing optimization problem. The infill density and layer height can take on any value within their respective bounds so we will configure both using RangeParameterConfigs. On the other hand, infill type be either have one of four distinct values: "honeycomb", "gyroid", "lines", or "rectilinear". We will use a ChoiceParameterConfig to represent it in the optimization.

infill_density = RangeParameterConfig(name="infill_density", parameter_type="float", bounds=(0, 100))
layer_height = RangeParameterConfig(name="layer_height", parameter_type="float", bounds=(0.1, 0.4))
infill_type = ChoiceParameterConfig(name="infill_type", parameter_type="str", values=["honeycomb", "gyroid", "lines", "rectilinear"])

client.configure_experiment(
    parameters=[infill_density, layer_height, infill_type],
    # The following arguments are only necessary when saving to the DB
    name="3d_print_strength_experiment",
    description="Maximize strength of 3D printed parts",
    owner="developer",
)

Output:
/home/runner/work/Ax/Ax/ax/api/utils/instantiation/from_config.py:76: AxParameterWarning:
is_ordered is not specified for ChoiceParameter "infill_type". Defaulting to False  since the parameter is a string with more than 2 choices.. To override this behavior (or avoid this warning), specify is_ordered during ChoiceParameter construction. Note that choice parameters with exactly 2 choices are always considered ordered and that the user-supplied is_ordered has no effect in this particular case.

Step 4: Configure Optimization

We want to maximize the compressive strength of our part, so we will set the objective to compressive_strength. However, we know that modifying the infill density, layer height, and infill type will affect the weight of the part as well. We'll include a requirement that the part must not weigh more than 10 grams by setting an outcome constraint when we call configure_experiment.

The following code will tell the Client that we intend to maximize compressive strength while keeping the weight less than 10 grams.

client.configure_optimization(objective="compressive_strength", outcome_constraints=["weight <= 10"])

Step 5: Run Trials

Now the Client has been configured we can begin conducting the experiment. Use attach_trial to attach any existing data, use get_next_trials to generate parameter suggestions, and use complete_trial to report manually observed results.

Attach Preexisting Trials

Sometimes in our optimization experiments we may already have some previously collected data from manual "trials" conducted before the Ax experiment began. This can be incredibly useful! If we attach this data as custom trials, Ax will be able to use the data points in its optimization algorithm and improve performance.

# Pairs of previously evaluated parameterizations and associated metric readings
preexisting_trials = [
    (
        {"infill_density": 10.43, "layer_height": 0.3, "infill_type": "gyroid"},
        {"compressive_strength": 1.74, "weight": 0.52},
    ),
    (
        {"infill_density": 55.54, "layer_height": 0.12, "infill_type": "lines"},
        {"compressive_strength": 4.63, "weight": 2.31},
    ),
    (
        {"infill_density": 99.43, "layer_height": 0.35, "infill_type": "rectilinear"},
        {"compressive_strength": 5.68, "weight": 2.84},
    ),
    (
        {"infill_density": 41.44, "layer_height": 0.21, "infill_type": "rectilinear"},
        {"compressive_strength": 3.95, "weight": 1.97},
    ),
    (
        {"infill_density": 27.23, "layer_height": 0.37, "infill_type": "honeycomb"},
        {"compressive_strength": 7.36, "weight": 3.31},
    ),
    (
        {"infill_density": 33.57, "layer_height": 0.24, "infill_type": "honeycomb"},
        {"compressive_strength": 13.99, "weight": 6.29},
    ),
]

for parameters, data in preexisting_trials:
    # Attach the parameterization to the Client as a trial and immediately complete it with the preexisting data
    trial_index = client.attach_trial(parameters=parameters)
    client.complete_trial(trial_index=trial_index, raw_data=data)

Output:
[INFO 09-09 19:39:07] ax.api.client: Trial 0 marked COMPLETED.
[INFO 09-09 19:39:07] ax.api.client: Trial 1 marked COMPLETED.
[INFO 09-09 19:39:07] ax.api.client: Trial 2 marked COMPLETED.
[INFO 09-09 19:39:07] ax.api.client: Trial 3 marked COMPLETED.
[INFO 09-09 19:39:07] ax.api.client: Trial 4 marked COMPLETED.
[INFO 09-09 19:39:07] ax.api.client: Trial 5 marked COMPLETED.

Ask for trials

Now, let's have Ax suggest which trials to evaluate so that we can find the optimal configuration more efficiently. We'll do this by calling get_next_trials. We'll make use of Ax's support for parallelism, i.e. suggesting more than one trial at a time -- this can allow us to conduct our experiment much faster! If our lab had three identical 3D printers, we could ask Ax for a batch of three trials and evaluate three different infill density, layer height, and infill types at once.

Note that there will always be a tradeoff between "parallelism" and optimization performance since the quality of a suggested trial is often proportional to the amount of data Ax has access to.

trials = client.get_next_trials(max_trials=3)
trials

Output:
[INFO 09-09 19:39:07] ax.api.client: GenerationStrategy(name='Center+Sobol+MBM:fast', nodes=[CenterGenerationNode(next_node_name='Sobol'), GenerationNode(node_name='Sobol', generator_specs=[GeneratorSpec(generator_enum=Sobol, model_key_override=None)], transition_criteria=[MinTrials(transition_to='MBM'), MinTrials(transition_to='MBM')]), GenerationNode(node_name='MBM', generator_specs=[GeneratorSpec(generator_enum=BoTorch, model_key_override=None)], transition_criteria=[])]) chosen based on user input and problem structure.
[INFO 09-09 19:39:08] ax.api.client: Generated new trial 6 with parameters {'infill_density': 50.0, 'layer_height': 0.25, 'infill_type': 'lines'} using GenerationNode CenterOfSearchSpace.
[INFO 09-09 19:39:08] ax.api.client: Generated new trial 7 with parameters {'infill_density': 40.796083, 'layer_height': 0.232923, 'infill_type': 'honeycomb'} using GenerationNode Sobol.
[INFO 09-09 19:39:08] ax.api.client: Generated new trial 8 with parameters {'infill_density': 5.072865, 'layer_height': 0.217759, 'infill_type': 'honeycomb'} using GenerationNode MBM.
{6: {'infill_density': 50.0, 'layer_height': 0.25, 'infill_type': 'lines'},
7: {'infill_density': 40.796083211898804,
'layer_height': 0.23292334079742436,
'infill_type': 'honeycomb'},
8: {'infill_density': 5.072864683160275,
'layer_height': 0.21775920859627712,
'infill_type': 'honeycomb'}}

Tell Ax the results

In a real-world scenerio we would print parts using the three suggested parameterizations and measure the compressive strength and weight manually, though in this tutorial we will simulate by calling a function. Once the data is collected we will tell Ax the result by calling complete_trial.

def evaluate(
    infill_density: float, layer_height: float, infill_type: str
) -> dict[str, float]:
    strength_map = {"lines": 1, "rectilinear": 2, "gyroid": 5, "honeycomb": 10}
    weight_map = {"lines": 1, "rectilinear": 2, "gyroid": 3, "honeycomb": 9}

    return {
        "compressive_strength": (
            infill_density / layer_height * strength_map[infill_type]
        )
        / 100,
        "weight": (infill_density / layer_height * weight_map[infill_type]) / 200,
    }


for trial_index, parameters in trials.items():
    client.complete_trial(trial_index=trial_index, raw_data=evaluate(**parameters))

Output:
[INFO 09-09 19:39:08] ax.api.client: Trial 6 marked COMPLETED.
[INFO 09-09 19:39:08] ax.api.client: Trial 7 marked COMPLETED.
[INFO 09-09 19:39:08] ax.api.client: Trial 8 marked COMPLETED.

We'll repeat this process a number of times. Typically experimentation will continue until a satisfactory combination has been found, experimentation resources (in this example our 3D printing filliment) have been exhausted, or we feel we have spent enough time on optimization.

# Ask Ax for the next trials
trials = client.get_next_trials(max_trials=3)
trials

Output:
[INFO 09-09 19:39:11] ax.api.client: Generated new trial 9 with parameters {'infill_density': 44.755989, 'layer_height': 0.12383, 'infill_type': 'honeycomb'} using GenerationNode MBM.
[INFO 09-09 19:39:11] ax.api.client: Generated new trial 10 with parameters {'infill_density': 47.141078, 'layer_height': 0.303541, 'infill_type': 'honeycomb'} using GenerationNode MBM.
[INFO 09-09 19:39:11] ax.api.client: Generated new trial 11 with parameters {'infill_density': 50.929388, 'layer_height': 0.196821, 'infill_type': 'honeycomb'} using GenerationNode MBM.
{9: {'infill_density': 44.75598884557233,
'layer_height': 0.12383033199520808,
'infill_type': 'honeycomb'},
10: {'infill_density': 47.141077917494776,
'layer_height': 0.3035409025152121,
'infill_type': 'honeycomb'},
11: {'infill_density': 50.92938823659888,
'layer_height': 0.19682062562123864,
'infill_type': 'honeycomb'}}

# Tell Ax the result of those trials
for trial_index, parameters in trials.items():
    client.complete_trial(trial_index=trial_index, raw_data=evaluate(**parameters))

Output:
[INFO 09-09 19:39:11] ax.api.client: Trial 9 marked COMPLETED.
[INFO 09-09 19:39:11] ax.api.client: Trial 10 marked COMPLETED.
[INFO 09-09 19:39:11] ax.api.client: Trial 11 marked COMPLETED.

# Ask Ax for the next trials
trials = client.get_next_trials(max_trials=3)
trials

Output:
[INFO 09-09 19:39:18] ax.api.client: Generated new trial 12 with parameters {'infill_density': 54.480037, 'layer_height': 0.225861, 'infill_type': 'honeycomb'} using GenerationNode MBM.
[INFO 09-09 19:39:18] ax.api.client: Generated new trial 13 with parameters {'infill_density': 86.548683, 'layer_height': 0.1, 'infill_type': 'honeycomb'} using GenerationNode MBM.
[INFO 09-09 19:39:18] ax.api.client: Generated new trial 14 with parameters {'infill_density': 23.941233, 'layer_height': 0.1, 'infill_type': 'honeycomb'} using GenerationNode MBM.
{12: {'infill_density': 54.48003661970755,
'layer_height': 0.22586148175170223,
'infill_type': 'honeycomb'},
13: {'infill_density': 86.5486828150697,
'layer_height': 0.1,
'infill_type': 'honeycomb'},
14: {'infill_density': 23.941233021244642,
'layer_height': 0.1,
'infill_type': 'honeycomb'}}

# Tell Ax the result of those trials
for trial_index, parameters in trials.items():
    client.complete_trial(trial_index=trial_index, raw_data=evaluate(**parameters))

Output:
[INFO 09-09 19:39:18] ax.api.client: Trial 12 marked COMPLETED.
[INFO 09-09 19:39:18] ax.api.client: Trial 13 marked COMPLETED.
[INFO 09-09 19:39:18] ax.api.client: Trial 14 marked COMPLETED.

Step 6: Analyze Results

At any time during the experiment you may analyze the results of the experiment. Most commonly this means extracting the parameterization from the best performing trial you conducted. The best trial will have the optimal objective value without violating any outcome constraints.

best_parameters, prediction, index, name = client.get_best_parameterization()
print("Best Parameters:", best_parameters)
print("Prediction (mean, variance):", prediction)

Output:
[WARNING 09-09 19:39:18] ax.adapter.cross_validation: Metric weight was unable to be reliably fit.
Best Parameters: {'infill_density': 40.796083211898804, 'layer_height': 0.23292334079742436, 'infill_type': 'honeycomb'}
Prediction (mean, variance): {'weight': (np.float64(8.081860059598545), np.float64(0.07491897290920471)), 'compressive_strength': (np.float64(17.797050719299325), np.float64(0.25837937569404634))}

Step 7: Compute Analyses

Ax can also produce a number of analyses to help interpret the results of the experiment via client.compute_analyses. Users can manually select which analyses to run, or can allow Ax to select which would be most relevant. In this case Ax selects the following:

Arm Effects Plots show the metric value for each arm on the experiment. Ax produces one plot using values from its internal surrogate model (this can be helpful for seeing the true effect of an arm when evaluations are noisy) and another using the raw metric values as observed.
Scatter Plot shows the effects of each trial on two metrics, and is useful for understanding the trade-off between the two outcomes
Summary lists all trials generated along with their parameterizations, observations, and miscellaneous metadata
Sensitivity Analysis Plot shows which parameters have the largest affect on the objective using Sobol Indicies
Slice Plot shows how the model predicts a single parameter effects the objective along with a confidence interval
Contour Plot shows how the model predicts a pair of parameters effects the objective as a 2D surface
Cross Validation helps to visualize how well the surrogate model is able to predict out of sample points

# display=True instructs Ax to sort then render the resulting analyses
cards = client.compute_analyses(display=True)

Output:
/home/runner/work/Ax/Ax/ax/analysis/utils.py:584: UserWarning:
Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at /pytorch/torch/csrc/utils/tensor_new.cpp:253.)
[WARNING 09-09 19:39:31] ax.adapter.base: TorchAdapter(generator=BoTorchGenerator) was not able to generate 100 unique candidates. Generated arms have the following weights, as there are repeats:
[1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

Modeled Arm Effects on compressive_strength

Modeled effects on compressive_strength. This plot visualizes predictions of the true metric changes for each arm based on Ax's model. This is the expected delta you would expect if you (re-)ran that arm. This plot helps in anticipating the outcomes and performance of arms based on the model's predictions. Note, flat predictions across arms indicate that the model predicts that there is no effect, meaning if you were to re-run the experiment, the delta you would see would be small and fall within the confidence interval indicated in the plot.

Observed Arm Effects on compressive_strength

Observed effects on compressive_strength. This plot visualizes the effects from previously-run arms on a specific metric, providing insights into their performance. This plot allows one to compare and contrast the effectiveness of different arms, highlighting which configurations have yielded the most favorable outcomes.

Modeled Arm Effects on weight

Modeled effects on weight. This plot visualizes predictions of the true metric changes for each arm based on Ax's model. This is the expected delta you would expect if you (re-)ran that arm. This plot helps in anticipating the outcomes and performance of arms based on the model's predictions. Note, flat predictions across arms indicate that the model predicts that there is no effect, meaning if you were to re-run the experiment, the delta you would see would be small and fall within the confidence interval indicated in the plot.

Observed Arm Effects on weight

Observed effects on weight. This plot visualizes the effects from previously-run arms on a specific metric, providing insights into their performance. This plot allows one to compare and contrast the effectiveness of different arms, highlighting which configurations have yielded the most favorable outcomes.

Modeled Effect on the Objective vs % Chance of Satisfying the Constraints

This plot displays the effects of each arm on the two selected metrics. It is useful for understanding the trade-off between the two metrics and for visualizing the Pareto frontier.

Modeled Effects: compressive_strength vs. weight

This plot displays the effects of each arm on the two selected metrics. It is useful for understanding the trade-off between the two metrics and for visualizing the Pareto frontier.

Summary for 3d_print_strength_experiment

High-level summary of the Trial-s in this Experiment

	trial_index	arm_name	trial_status	generation_node	weight	compressive_strength	infill_density	layer_height	infill_type
0	0	0_0	COMPLETED	nan	0.52	1.74	10.43	0.3	gyroid
1	1	1_0	COMPLETED	nan	2.31	4.63	55.54	0.12	lines
2	2	2_0	COMPLETED	nan	2.84	5.68	99.43	0.35	rectilinear
3	3	3_0	COMPLETED	nan	1.97	3.95	41.44	0.21	rectilinear
4	4	4_0	COMPLETED	nan	3.31	7.36	27.23	0.37	honeycomb
5	5	5_0	COMPLETED	nan	6.29	13.99	33.57	0.24	honeycomb
6	6	6_0	COMPLETED	CenterOfSearchSpace	1	2	50	0.25	lines
7	7	7_0	COMPLETED	Sobol	7.88166	17.5148	40.7961	0.232923	honeycomb
8	8	8_0	COMPLETED	MBM	1.04831	2.32958	5.07287	0.217759	honeycomb
9	9	9_0	COMPLETED	MBM	16.2643	36.143	44.756	0.12383	honeycomb
10	10	10_0	COMPLETED	MBM	6.98867	15.5304	47.1411	0.303541	honeycomb
11	11	11_0	COMPLETED	MBM	11.6442	25.876	50.9294	0.196821	honeycomb
12	12	12_0	COMPLETED	MBM	10.8544	24.121	54.48	0.225861	honeycomb
13	13	13_0	COMPLETED	MBM	38.9469	86.5487	86.5487	0.1	honeycomb
14	14	14_0	COMPLETED	MBM	10.7736	23.9412	23.9412	0.1	honeycomb

Sensitivity Analysis for compressive_strength

Understand how each parameter affects compressive_strength according to a second-order sensitivity analysis.

layer_height vs. compressive_strength

The slice plot provides a one-dimensional view of predicted outcomes for compressive_strength as a function of a single parameter, while keeping all other parameters fixed at their status_quo value (or mean value if status_quo is unavailable). This visualization helps in understanding the sensitivity and impact of changes in the selected parameter on the predicted metric outcomes.

infill_density vs. compressive_strength

infill_density, layer_height vs. compressive_strength

The contour plot visualizes the predicted outcomes for compressive_strength across a two-dimensional parameter space, with other parameters held fixed at their status_quo value (or mean value if status_quo is unavailable). This plot helps in identifying regions of optimal performance and understanding how changes in the selected parameters influence the predicted outcomes. Contour lines represent levels of constant predicted values, providing insights into the gradient and potential optima within the parameter space.

Sensitivity Analysis for weight

Understand how each parameter affects weight according to a second-order sensitivity analysis.

layer_height vs. weight

The slice plot provides a one-dimensional view of predicted outcomes for weight as a function of a single parameter, while keeping all other parameters fixed at their status_quo value (or mean value if status_quo is unavailable). This visualization helps in understanding the sensitivity and impact of changes in the selected parameter on the predicted metric outcomes.

infill_density vs. weight

infill_density, layer_height vs. weight

The contour plot visualizes the predicted outcomes for weight across a two-dimensional parameter space, with other parameters held fixed at their status_quo value (or mean value if status_quo is unavailable). This plot helps in identifying regions of optimal performance and understanding how changes in the selected parameters influence the predicted outcomes. Contour lines represent levels of constant predicted values, providing insights into the gradient and potential optima within the parameter space.

Cross Validation for weight

The cross-validation plot displays the model fit for each metric in the experiment. It employs a leave-one-out approach, where the model is trained on all data except one sample, which is used for validation. The plot shows the predicted outcome for the validation set on the y-axis against its actual value on the x-axis. Points that align closely with the dotted diagonal line indicate a strong model fit, signifying accurate predictions. Additionally, the plot includes 95% confidence intervals that provide insight into the noise in observations and the uncertainty in model predictions. A horizontal, flat line of predictions indicates that the model has not picked up on sufficient signal in the data, and instead is just predicting the mean.

Cross Validation for compressive_strength

Ax Constraints Feasibility Warning

The constraints feasibility health check utilizes samples drawn during the optimization process to assess the feasibility of constraints set on the experiment. Given these samples, the model believes there is at least a 0.95 probability that the constraints will be violated. We suggest relaxing the bounds for the constraints on this Experiment.

Conclusion

This tutorial demonstrates how to use Ax's Client for optimizing the strength of 3D printed parts in a human-in-the-loop setting. By iteratively collecting data and refining parameters, you can effectively apply black box optimization to real-world experiments.

Background​

Learning Objectives​

Prerequisites​

Step 1: Import Necessary Modules​

Step 2: Initialize Client​

Step 3: Configure Experiment​

Step 4: Configure Optimization​

Step 5: Run Trials​

Attach Preexisting Trials​

Ask for trials​

Tell Ax the results​

Step 6: Analyze Results​

Step 7: Compute Analyses​

Conclusion​