[GENERAL SUPPORT]: Selecting Highest Acquisition Function Value

[XZaitzeff]

Question

I am fairly new with Ax and am doing some onboarding type stuff but it is all from an older build of Ax so the examples do not work anymore. What I am trying to do is take a discrete set of featurized data points and evaluate these in the acquisition function then select the one with the highest value to test on my actual function. The model is trained on some data. However I cannot figure out from the documentation how to access my model's acquisition function to evaluate my points. Any help?

Please provide any relevant code snippet if applicable.

Code of Conduct

• I agree to follow this Ax's Code of Conduct

[mgarrard]

Hi @XZaitzeff could you provide a bit more information, notably:

1. Why are you sing an older build, is there something about the new Ax versions that don't fit your needs?
2. Could you provide a code snippet for us to be able to better help you?
   Thank you!

[XZaitzeff]

The older build allows me to access the acquisition function a lot easier. I was trying to replicate this behavior in the newer version but I could not. What I am trying to do is approximate the hardness of materials based on a lot of features. I have some already tested data so I attach trials from the preexisting data to start the approximation. Then using the features from materials not tested, run them through the acquisition function, and choose the one with the highest value and measure its hardness. Repeat this for a certain amount of trials. I was able to get it working on the older version but could not find the same functionality in the newer function. Here is the code that works on the older function.

gs = GenerationStrategy(
    steps=[
        GenerationStep(
            model=Models.BOTORCH_MODULAR,
            num_trials=-1,  # No limitation on how many trials should be produced from this step
            max_parallelism=3,  # Parallelism limit for this step, often lower than for Sobol
        ),
    ]
)

ax_client = AxClient(generation_strategy=gs, random_seed=42)

feature_names = X_train.columns

ax_client.create_experiment(
    parameters=[
        {"name": "avg_crystal_radius", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]}, # MOD: update param
        {"name": "mode_crystal_radius", "type": "range", 'value_type':'float', "bounds":[0.0, 200.0]},
        {"name": "avg_Miracle_Radius_[pm]", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]},
        {"name": "dev_Covalent_Radius", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]},
        {"name": "mode_Atomic_Number", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]},
        {"name": "avg_families", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]},
        {"name": "avg_ionic_radius", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]},
        {"name": "dev_heat_of_vaporization_(kJ/mol)_",'value_type':'float', "type": "range", "bounds":[0.0, 200.0]},
        {"name": "avg_Pauling_Electronegativity",'value_type':'float', "type": "range", "bounds":[0.0, 200.0]},
        {"name": "dev_ionic_radius", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]},
        {"name": "dev_families", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]},
        {"name": "avg_polarizability(A^3)",'value_type':'float', "type": "range", "bounds":[0.0, 200.0]},
        {"name": "load", "type": "range",'value_type':'float', "bounds":[0.0, 200.0]},
    ],
    objectives={
        obj1_name: ObjectiveProperties(minimize=False),
    },
)


initial_data = [({name : X_train[name].iloc[i] for name in feature_names}, {"hardness": y_train.iloc[i]}) for i in range(len(X_train))]

for parameters, raw_data in initial_data:
    # First attach the trial and note the trial index
    trial_index, trial = ax_client.attach_trial(
        parameters=parameters,
    )

    # Then complete the trial with the existing data
    ax_client.complete_trial(
        trial_index=trial, raw_data=raw_data,
    )

candidates = [{name : X_candidates[name].iloc[i] for name in feature_names} for i in range(len(X_candidates))]

obs_feat = [ObservationFeatures(parameters=candidate) for candidate in candidates]
candidates_tested = np.zeros(25)

for i in range(25):
    ax_client.fit_model()
    model = ax_client.generation_strategy.model
    acqf_values = model.evaluate_acquisition_function(obs_feat)

    candidate_index = np.argmax(acqf_values)
    candidate = candidates[candidate_index]
        
    candidates_tested[i] = candidate_index
    trial_index, trial = ax_client.attach_trial(
        parameters=candidate,
    )
    obs_feat.pop(candidate_index)
    result = measure_hardness(candidate_index)
    
    ax_client.complete_trial(trial_index=trial, raw_data={"hardness":result})```

Keep in mind this is a onboarding type example so the data is faked and the measuring function just grabs a hardness.
Thank you for your response!

[mgarrard]

Awesome thanks for the example @XZaitzeff -- I think it may actually be a matter of updating some of the naming of some components here, since we updated our param names in the most recent version.

Models --> Generators (ie Models.BOTORCH_MODULAR now is Generators.BOTORCH_MODULAR)
Model --> adapter (ie client.generation_strategy.model now is client.generation_strategy.adapater)

We also use the trial_index instead of the full trial object in complete_trial().

Given this, I think your code block should look something like this using the most recent API verison -- I am unable to run to validate given that I don't have access X_train. It is still a bit out of standard scope given the need to view the acquisition func etc, but this should be much more aligned. Lmk if it answers your question, or if you have follow ups :)

import numpy as np

from ax.adapter.registry import Generators
from ax.core.observation import ObservationFeatures
from ax.generation_strategy.generation_node import GenerationStep
from ax.generation_strategy.generation_strategy import GenerationStrategy
from ax.service.ax_client import AxClient
from ax.service.utils.instantiation import ObjectiveProperties


gs = GenerationStrategy(
    steps=[
        GenerationStep(
            generator=Generators.BOTORCH_MODULAR,
            num_trials=-1,  # No limitation on how many trials should be produced from this step
            max_parallelism=3,  # Parallelism limit for this step
        ),
    ]
)

ax_client = AxClient(generation_strategy=gs, random_seed=42)

feature_names = X_train.columns

ax_client.create_experiment(
    parameters=[
        {"name": "avg_crystal_radius", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "mode_crystal_radius", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "avg_Miracle_Radius_[pm]", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "dev_Covalent_Radius", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "mode_Atomic_Number", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "avg_families", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "avg_ionic_radius", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "dev_heat_of_vaporization_(kJ/mol)_", "value_type": "float", "type": "range", "bounds": [0.0, 200.0]},
        {"name": "avg_Pauling_Electronegativity", "value_type": "float", "type": "range", "bounds": [0.0, 200.0]},
        {"name": "dev_ionic_radius", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "dev_families", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
        {"name": "avg_polarizability(A^3)", "value_type": "float", "type": "range", "bounds": [0.0, 200.0]},
        {"name": "load", "type": "range", "value_type": "float", "bounds": [0.0, 200.0]},
    ],
    objectives={
        obj1_name: ObjectiveProperties(minimize=False),
    },
)


initial_data = [
    ({name: X_train[name].iloc[i] for name in feature_names}, {"hardness": y_train.iloc[i]})
    for i in range(len(X_train))
]

for parameters, raw_data in initial_data:
    # attach_trial returns (parameters, trial_index), not (trial_index, trial)
    _, trial_index = ax_client.attach_trial(parameters=parameters)

    ax_client.complete_trial(
        trial_index=trial_index,
        raw_data=raw_data,
    )

candidates = [{name: X_candidates[name].iloc[i] for name in feature_names} for i in range(len(X_candidates))]

obs_feat = [ObservationFeatures(parameters=candidate) for candidate in candidates]
candidates_tested = np.zeros(25)

for i in range(25):
    ax_client.fit_model()
    
    # Use adapter instead of model - the property has been renamed
    adapter = ax_client.generation_strategy.adapter
    acqf_values = adapter.evaluate_acquisition_function(observation_features=obs_feat)

    candidate_index = np.argmax(acqf_values)
    candidate = candidates[candidate_index]

    candidates_tested[i] = candidate_index
    _, trial_index = ax_client.attach_trial(parameters=candidate)
    obs_feat.pop(candidate_index)
    result = measure_hardness(candidate_index)

    ax_client.complete_trial(trial_index=trial_index, raw_data={"hardness": result})