Ax Idea for interactive template generation: multi-objective, multi-fidelity, batch optimization, etc.

There's an idea I'm playing around with and thinking about implementing to make some of the advanced BO topics more accessible, especially in chemistry and materials science.

With the escience notebook tutorials that I put out (see @sp8rks's YouTube playlist), I started expanding them with progressively complicated names as I was incorporating different features:

A lot of these were based on content from other issues, often using the Service API, and occasionally delving into Developer API or BoTorch components. As I was creating these extra tutorials with a labmate in mind (@ramseyissa), it struck me how modular the functionality was (👏) as well as the sheer number of notebooks that would be required to cover every combination of these.

My idea is to create something somewhat similar to the interactive API snippets: except allowing for an interactive selection of the desired features.

For example, a user could choose from the following options to create the corresponding template for 2.8.0.4-ax_service_existing_data_saasbo_multi_objective_batch_equality.ipynb:

[x] multi-objective
[ ] continuous multi-fidelity
[ ] discrete multi-fidelity
[ ] multi-task
[x] high-dimensional
[x] batch optimization
[ ] asynchronous optimization (mutually exclusive with batch)
[ ] linear constraints
[x] equality constraints
[x] attach existing data (e.g., from CSV file)

Other categories:

Choice of acquisition function (e.g., minimizing model error https://github.com/facebook/Ax/issues/930)

And maybe in terms of the user experience, it would look and act like the PyTorch installation page:

https://user-images.githubusercontent.com/45469701/222261641-39e94351-d7c9-41d8-85f0-320daaa879b0.mp4

In terms of actually implementing the above, I wonder if ipywidgets might come in handy. I'm not very well versed in HTML and Javascript. I'm also having trouble finding the general language to describe this type of interface. An alternative would be a Google Colab notebook that someone can run based on the options they select (which is probably where I would start).

For PyTorch's graphical implementation, see What is the PyTorch installation interactive grid called and how do I make my own?.

The template snippets at the bottom could be generated in advance for all combinations, rather than trying to run something on-the-fly. Additionally, it would be possible to run unit tests for each of the snippet combinations (this could simply be ensuring that it runs without errors).

I wanted to make a post to get some feedback and also have a place where I can refer back to it. What are your thoughts?

EDIT: also gauging interest via Twitter and LinkedIn

Mar 01 '23 20:03 sgbaird

As a chemist / material scientist discovering Ax just right now and looking into how to have batches for single-objective optimization with ALEBO, I am seduced.

Mar 01 '23 23:03 CharlyEmpereurmot

This is a neat idea! We are always looking for ways to improve our documentation and tutorials to make learning Ax easier. We already have some plans for improving docs in the pipeline, and I will raise this concept to the team as well. As always, thanks for your dedication to improving Ax for all our users :)

Mar 08 '23 16:03 mpolson64

I'm fleshing this idea out more. I'm planning to:

Tease out the PyTorch installation interface code as a MWE that I'll adapt to Ax
Use Jinja2-based templating to create the scripts

Jun 24 '23 21:06 sgbaird

I created a repo called Honegumi (骨組み, pronounced "ho neh goo mee") which means skeletal framework in Japanese. I wrote "a gentle introduction to Jinja" Colab tutorial that will make it easier for others to contribute via general feedback, design principles, and features.

Just as a very basic example, the following Jinja template has a toggle for switching between single- and multi-objective, going from the objective_name and minimize kwargs to the objectives kwarg with ObjectiveProperties

from ax.service.ax_client import AxClient
from ax.utils.measurement.synthetic_functions import branin
{% if use_moo %}
from ax.service.utils.instantiation import ObjectiveProperties

obj1_name = "branin"
obj2_name = "neg_branin"

def branin_moo(x1, x2):
    """Multi-objective branin function

    The first objective is the normal branin value and the second
    objective is the negative branin value.
    """
    return {obj1_name: branin(x1, x2), obj2_name: -branin(x1, x2)}
{% endif %}

ax_client = AxClient()
ax_client.create_experiment(
    parameters=[
        {"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
        {"name": "x2", "type": "range", "bounds": [0.0, 10.0]},
    ],
{% if use_moo %}
    objectives={
        obj1_name: ObjectiveProperties(minimize=True, threshold=None),
        obj2_name: ObjectiveProperties(minimize=True, threshold=None),
    },
{% else %}
    objective_name="branin",
    minimize=True,
{% endif %}
)

for _ in range(15):
    parameters, trial_index = ax_client.get_next_trial()
    results = branin{% if use_moo %}_moo{% endif %}(
        parameters["x1"], parameters["x2"]
        )
    ax_client.complete_trial(trial_index=trial_index, raw_data=results)

{% if use_moo %}
pareto_results = ax_client.get_pareto_optimal_parameters()
{% else %}
best_parameters, metrics = ax_client.get_best_parameters()
{% endif %}

Hoping to define some good design principles and a roadmap. Lmk if you're interested in helping!

Jun 27 '23 09:06 sgbaird

@sgbaird, this is awesome, sorry we left this unanswered for a while! Any help you currently need?

Jul 25 '23 17:07 lena-kashtelyan

@lena-kashtelyan, thanks! No worries. Actually, I almost have a full working prototype with just a few options. I'll share that really soon! I'd love to get some help with it.

Jul 25 '23 18:07 sgbaird

Hi @lena-kashtelyan, what do you think of this kind of functionality in terms of the high-level interface? (nothing specific to Ax in this one - still need to connect the interface to the backend)

https://github.com/sgbaird/honegumi/blob/ax/docs/bootstrap_examples/bootstrap_gpt_help/main.html

(download and open in a browser)

Screencast: https://app.screencast.com/ywGS5w2FsLygO

Aug 01 '23 16:08 sgbaird

This is very cool! Where are you thinking of hosting this, @sgbaird? Is the thought to put these on the Ax website or somewhere else?

Aug 02 '23 14:08 lena-kashtelyan

@lena-kashtelyan, I've been wondering about that. For now, I plan to have a small Readthedocs page based on https://github.com/sgbaird/honegumi. While I'm focusing on Ax implementations, I also wanted to set it up so that solutions for other platforms could be added.

Aug 04 '23 19:08 sgbaird

Hi @lena-kashtelyan, I'm actively thinking about that and open to feedback. While I'm prototyping, I plan to expose it via a readthedocs generated from https://github.com/sgbaird/honegumi. I'm designing the tool in a way that can be extended to other platforms; however, the Ax functionality has its own namespace and therefore can be installed independently from other platforms. Right now, my focus is on Ax.

I have almost all of the functionality ready:

https://github.com/sgbaird/honegumi/blob/ax/scripts/generate_scripts.py
https://github.com/sgbaird/honegumi/tree/ax/docs/generated_scripts/ax
https://github.com/sgbaird/honegumi/tree/ax/tests/generated_scripts/ax
https://github.com/sgbaird/honegumi/actions/runs/5766340905
https://github.com/sgbaird/honegumi/tree/ax/docs/sandbox/jinja2

The last major piece of plumbing I need is programmatically running pytest and gathering the results within Python to pass it to the HTML file.

Aug 05 '23 10:08 sgbaird

Wow, very cool! @mpolson64 and @esantorella are thinking through a potential redesign of our website infra/deployment, and this seems like it could be an excellent fit to play into that project. I'll assign this issue to @mpolson64 to think that through : )

Aug 10 '23 15:08 lena-kashtelyan

Hi @lena-kashtelyan, @mpolson64, and @esantorella, I have a working example on the index page of https://honegumi.readthedocs.io/. I'm keen to get your thoughts! See the short screencast below:

https://github.com/facebook/Ax/assets/45469701/fa362011-5a74-4279-9a40-995cf39acec7

cc @Balandat @saitcakmak @bernardbeckerman @Ryan-Rhys

Aug 29 '23 18:08 sgbaird

Wow this is so cool!! cc @mpolson64 who is thinking through our potential website reset –– would be great to feature this!

Sep 12 '23 19:09 lena-kashtelyan

Thank you! I felt really good getting to this working example. I would love some help on continuing to expand/develop this and make sure I don't make some early decisions that make it really hard to maintain/scale later on, add features that aren't as useful, etc.

Sep 14 '23 20:09 sgbaird

Lots of progress made since then! https://honegumi.readthedocs.io/

Now there are 10 visible options
We added tooltips next to options
Tutorials and conceptual docs added (more to come)

May 15 '24 18:05 sgbaird

Thanks so much for all your hard work on Honegumi @sgbaird ! I'm going to close out this Issue for now since there are no action items for us, but we really appreciate all your contributions to the Ax ecosystem.

Relatedly, there are some changes coming to Ax in the near future that we hope will address many common pain points with new and inexperienced users. I'd love to get some initial feedback as its coming together -- I'll plan on reaching out at some point soon :)

Jul 10 '24 16:07 mpolson64