How to declare categorical variables in custom metric

[Meganton]

Hi everyone,
I want to use the Jahs-Bench-Metric as a custom metric in the ax-platform with several Aquisiton and Surrogate-Functions. I created a custom metric for it and it works great, just when i try using it with the 'MixedSingleTaskGP'-Model, i get the error ValueError: Must specify categorical dimensions for MixedSingleTaskGP. Now i would like to declare which variables are categorical, as the error told me to do, but this is supposed to happen through the model_options-argument of Surrogate, which is a dictionary. But there is not documentation on how this dictionary is supposed to look, as in what keywords I have to use and what formats the variables have to have, just that it's a dict[str, any].
Would be great if someone could help me, thanks in advance!

My code-snippet:

exp = Experiment(
    name="experiment_jahsbench",
    search_space=SearchSpace(parameters=parameterlist),
    optimization_config=OptimizationConfig(
        objective=Objective(
            metric=JAHS_BENCH_Metric("jahs_metric"),
            minimize=True)),runner=MockRunner())

for _ in range(2):
    model_bridge_with_GPEI = Models.BOTORCH_MODULAR(
        experiment=exp,
        data=exp.fetch_data(),
        surrogate=Surrogate(
            MixedSingleTaskGP, model_options=???}
        ),
        botorch_acqf_class=ExpectedImprovement,
    )
    generator_run = model_bridge_with_GPEI.gen(1)
    trial = exp.new_trial(generator_run=generator_run)
    trial.run()
    trial.mark_completed()

[saitcakmak]

Hi @Meganton. Thanks for pointing out the lack of documentation. model_kwargs get passed to the construct_inputs method of the model class (MixedSingleTaskGP). In this case though, the cat_dims input should be passed down as SearchSpaceDigest.categorical_features, which you shouldn't have to include in model_kwargs -- it should get extracted from the search space in extract_search_space_digest.

My guess is that this isn't happening because we pass in the transformed search space to this function, and at that time your categorical features has been transformed either by one-hot-encoding (if specified as unordered) or by continuous relaxation (if ordered -- first converted into integers, which are then relaxed to floats). This happens because Models.BOTORCH_MODULAR specifies the following transforms by default:

Cont_X_trans: List[Type[Transform]] = [
    RemoveFixed,
    OrderedChoiceEncode,
    OneHot,
    IntToFloat,
    Log,
    Logit,
    UnitX,
]
Y_trans: List[Type[Transform]] = [IVW, Derelativize, StandardizeY]

To avoid the categoricals getting auto transformed, we need to change the list of transforms. If you're using an unordered choice parameter for the categorical dimension, removing the OneHot transform should do it. If it is ordered, than it'd be the OrderedChoiceEncode transform. You should be able to pass the list of transforms as another argument to Models.BOTORCH_MODULAR.

[Meganton]

I changed my modelbridge to this

model_bridge_with_GPEI = Models.BOTORCH_MODULAR(
        experiment=exp,
        data=exp.fetch_data(),
        surrogate=Surrogate(
            MixedSingleTaskGP
        ),
        botorch_acqf_class=ExpectedImprovement,transforms=[RemoveFixed,
    OrderedChoiceEncode,
    IntToFloat,
    Log,
    Logit,
    UnitX]+Y_trans
    )

I hope that's what you meant, leaving out the OneHot-Encoder, as i use ChoiceParameters with is_ordered=False. But now i get the error too many dimensions 'str'...

[saitcakmak]

I am not familiar with that error message. While working on something else, I noticed that we actually use the MixedSingleTaskGP by default in Ax (via choose_generation_strategy) in certain cases and that it is the default model chosen by MBM when there are categorical variables. This is encoded in the registry as BO_MIXED and the dispatching rule in choose_generation_strategy checks that there aren't too many combinations of possible choices (since this uses optimize_acqf_mixed under the hood, enumerating all categorical options): https://github.com/facebook/Ax/blob/main/ax/modelbridge/dispatch_utils.py#L202-L213. Could it be that you have too many choices for it to be handled by optimize_acqf_mixed?

The optimization with optimize_acqf_mixed is handled in Acquisition.optimize. If you'd like to customize that part, you can subclass Acquisition and pass that to the model constructor as acquisition_class=CustomAcquisition. I hope this helps!