Implement a bayesian opimization algorithm for CADET-Process based on ax

[flo-schu]

implements the developer API of Ax (https://ax.dev/tutorials/gpei_hartmann_developer.html) for CADET-Process.

new features

• single objective optimization
• multi-objective optimization
• implement different algorithms for multi-objective optimization https://ax.dev/tutorials/multiobjective_optimization.html
  • Subclass AxInterface for different algorithms
  • GPEI
• termination criteria (nice to have)
  • multi-objective (can be singular in the case of minimization)
• linear constraints
• nonlinear constraints
• unittests
  • single objective NRMSE
  • multi-objective (fit gaussian curve)
• parallelize evaluations https://ax.dev/tutorials/scheduler.html
• docstrings
• Parametrisierung (sollte einstellbar sein - API - genaue Werte über Forum lösen)
  • SEM: Forumsbeitrag um Wissen zu sammeln (sollte einstellbar sein)
  • hyperparameters: Forumsbeitrag exploration vs exploitation was sind gute Parameter für verschiedene Algortihmen

improvements

• increases the flexibility of setup_optimization_problem by allowing to pass custom obj_functions

infra

• Make ax optional dependency

[flo-schu]

@schmoelder: Right now multiobjective optimization works.

to-dos

1. finalize unittests for

• multi-objective
• single-objective --> this involves making single-objective work again (since it was broken for multi-objective)
• reproducible benchmarking tests for Ax. Ideally can be reused for other optimizers

because these follow different implementation. I will see if I can implement single obj. as a special case of multi-obj. or if I have to use different implementations dependent on the case

2. Ax provides a runner class in the developer API

After a short interlude of experimenting with the Service API, we decided that the Developer API of Ax is the right interface, because the Service API makes some passing of objects in the optimization problem very complicated such as string-based constraints. In addition, the Service API is relative inflexible when it comes to Generation Strategies. The reason for going for the Service API was increased simplicity of the interface, mainly because the Scheduling function seemed extremely complicated. Since parallelization is a task that can be handled very well by CADET-Process. This feature seems to be unnecessary. Instead the focus is now to improve the adapter by implementing an interface that CADET-Process can best interact with. In detail this involves:

• Implementation of a runner (which takes a trial with N arms as an argument) could be the right object for implementing the following steps. But maybe this can be solved differently as well.
  • after generation of a trial with $N$ arms, parameters can be extracted and passed to the optimization problem for parallel evaluation
  • These results would then need to be backconverted to a form that Ax can understand

3. Start writing docstrings for the classes and methods and clean up code

4. Implement Different Surrogate Models / Acquisition Functions in Ax

I'm currently understanding the framework of ax that deals with models. It is again complicated and not super easy to disentangle the different components, so I'm just writing things down here to track the evolution of my knowledge.

A model in ax is a combination of surrogate and acquisition function.

Surrogate can be passed to the optimizer in a pre-instantiated form. This allows to resume a fitting procedure, which is nice if evaluations are costly. @schmoelder is this a feature we want/need?

Acquisition functions in contrast are passed to the Model as classes that are instantiated at each iteration and deleted afterwards.

Then there is a Model Bridge component that links the experiment (problem, data, trials, arms) to the model (surrogate and acquisition) by applying a number of transforms to the data Model Stack in Ax and implements the methods predict and gen. The latter is needed to generate samples for forward simulations. So we can understand the ModelBrige as adapters between the external components (e.g. torch) which are described as Model and the interface of ax (Arms, Trials, ...).

A number of ready-made models are already available in the ax.modelbridge.registry module. These are specified in the global dictionary MODEL_KEY_TO_MODEL_SETUP like so. The Bridge class for initialization would be: RandomModelBridge so it would be possible

• to write a Bridge for Hopsy
• or simply use hopsy for the generation of initial samples and wrap this data in the Data class of Ax (which I would prefer) and then use it to initialize the true BO step.

MODEL_KEY_TO_MODEL_SETUP = {
    "MOO": ModelSetup(
        bridge_class=TorchModelBridge,
        model_class=MultiObjectiveBotorchModel,
        transforms=Cont_X_trans + Y_trans,
        standard_bridge_kwargs=STANDARD_TORCH_BRIDGE_KWARGS,
    ),
    "BoTorch": ModelSetup(
        bridge_class=TorchModelBridge,
        model_class=ModularBoTorchModel,
        transforms=Cont_X_trans + Y_trans,
        standard_bridge_kwargs=STANDARD_TORCH_BRIDGE_KWARGS,
    ),
}

the Cont_X_trans and Y_trans are lists of transformations that are equal as it seems for Single Objective and Multi objective

For single objective, a ModularBotorchModel can be specified as the model_class, while for multi objective this seems currently not possible. Here, as of now in the defaults only EHVI and NEHVI (noisy) expected hypervolume improvement are implemented. I think this means for multi-objective currently we are limited somewhat in the choice of models / acquisition function, unless we implement our own method. However, as I uncover more of the docu/examples of Ax, I could imagine that I discover more methods. Registration of our own models could be an interesting feature to give users some options, of which we know they work in the majority of cases. But we can also use simply the ModularBotorchModel and initialize it with the surrogate and acquisition function we want. Currently this behavior is implemented.

So the best way to go currently I think is to:

• implement alternatives for surrogates and acquisiton functions with model bridge utilities. DOcumentation (Acquisition Functions · BoTorch, https://ax.dev/docs/models.html, https://ax.dev/tutorials/modular_botax.html)

  • allow user specification of different acquisition functions for single objective and pass it to ModularBotorchModel
  • fix the EHVI/NEHVI for multiple objective (this is done implicitly by not specifying anything).

• implementation of a strategy to draw initial samples from Hopsy. (No Priority currently, may be more difficult to implement)

  • switch from SOBOL to HOPSY for initial evaluations (wrap output from hopsy in a form digestable for Ax or write a BridgeClass for Hopsy)
  • one possible approach could be to subclass Model and implement Hopsy sampling in it
  • @schmoelder: do you have a different implementation idea?
  • create_initial_values. The function is maybe a bit loaded with other stuff, perhaps this can be cleaned up.
  • sobol arranges hypervolumes so their midpoints are most evenly distributed (not random)

questions

1. Do you think you can work on making intermediate results of the evaluation-toolchain available, or give me a short Idea on how to do this? This would improve the efficiency of my callbacks, unless I'm missing something again 😄

can be solved by adding the flag requires=[simulator] to add_callback

2. post processing is inefficient

Hypothesis Jo:

• serialization may be inefficient. Try turn of by optimizer.optimize(save_results=False)

try profiling in vscode.

[flo-schu]

Tasks for this week

priority

• continue testing sweet for Ax. As benchmarking problems that are reproducibly solved

  • write test case for multi-objective optimization (Wednesday)
  • make sure that moo uses NEHVI or EHVI (print/warning statement) (Tuesday)
  • set seed in for MCMC sampling in Ax according to How to set a seed for the acquisition function facebook/Ax#1663 (comment)
  • write linear constraints TestProblem for single-objective optimization
  • wrap single objective without constraints in TestProblem class

• Test the different TestProblems for Scipy and Pymoo

  • also use for scipy (postponed to in 1--2 weeks)
  • pymoo (postponed to in 1--2 weeks)

• work on models in Ax

  • implement different acquisition functions / surrogate models --> model bridge

• hopsy for initial sampling. May be not so simple to implement, because I am not sure if passing the output of fetch_data to Model(..., data=experiment.fetch_data()) contains all information of the experiment. IN this case generating data and wrapping it in a ax.Data(pd.DataFrame(mygenerated_data)) fashion might not be enough. **However, this is currently just an assumption due to the form of the output of fetch_data, which just outputs the f-values and not the parameters. The best approach currently would be to see if the Data class of Ax also contains the parameters. If this is not the case. Probably approach 2 must be taken, which is: Implement HOPSY as a Model class with a RandomModelBridge. (Postponed)

[schmoelder]

If I can add:

• Allow custom initialization (e.g. using sample from hopsy)

[flo-schu]

@schmoelder, is there an assertion in your test suite that asserts the near equality of two arrays? If not should I write one?

[schmoelder]

I mostly just use np.testing.almost_equal.

[flo-schu]

@schmoelder do you know a good simple multi-objective optimization problem with linear constraints and a known pareto front, ideally implemented in python (not ax so I can use it for other optimization packages as well)?

If not, have you experience how to derive the pareto front for simple problems or do you know who to ask.

I asked chatGPT for a problem but I suspect it got the pareto front wrong. But maybe its also me and we can have a look tomorrow together

[schmoelder]

Maybe some of pymoo's test problems would work (see: https://pymoo.org/problems/test_problems.html#Multi-Objective).

[flo-schu]

@schmoelder i have written a test suite for Ax. This is not completely finalized, but I think its a good start to use this also for other optimizers. I wouldn't mind if you took a look at it and review if you're happy with the implementation. I don't mind changing parts of the implementation.

[schmoelder]

Thanks a lot! I will definitely have a look and review!

Btw, I took the liberty to rebase again. you might have to reset --hard (but, as always, be careful with that one! ;-) )

[flo-schu]

you might have to reset --hard (but, as always, be careful with that one! ;-) )

I used git merge --abort when the pull was not successful and updated strategy for reconciling diverging branches. This worked quite smoothly

[schmoelder]

This comment might get us some idea on how to address custom initialization:

facebook/Ax#768 (comment)

[flo-schu]

Locally all tests are passing (11 min).

• I have made a minor change to the OptimizerBase (added supports_single_objective flag) and added a corresponding check. @schmoelder Please review.
• Temporarily added ax-platform to [testing] optional dependencies, to test AxInterface on github. @schmoelder, maybe you can think about how you would like to integrate ax into the testing pipeline at this point.

Other than that all should be okay for now and ready for merge 🎉

[schmoelder]

If we don't need this anymore, please remove.

[schmoelder]

Great news, @flo-schu! Thanks for all your efforts!

I will try to run some problems locally to test performance of some "real" problems. Really excited!

Other than some minor remarks in the code base (see review comments), please make sure to clean up the git history a bit. I can also help you with that if you want.
Other than that, LGTM!

[flo-schu]

@schmoelder: Everything looks good. Tests are passing. IMO we can merge 🚀

[schmoelder]

This branch was merged manually in e5cbf17. Congrats @flo-schu and many thanks for this amazing contribution! 🥳