Feature Request: PyCharm support for viz

[HUSTYuChen]

1. I can't plot in Pycharm using the built-in functions in Ax. All the tutorials on the site are there about drawing with Jupyter, is there an easy way to draw with Pycharm? (some functions can't work when using them in Pycharm, such as render)
   
2. When I code by using develop API, all the metrics (Indicators to be optimized) are calculated one after another. If run function run_bdsim (a physics simulation software) with the above method, which is too time-consuming. How do I solve this problem at this point? Run the run_bdsim once to get all the metrics. (run_bdsim implements the physical simulation and get_beam_statistical gets the associated results.)
   So I run the function run_bdsim in a single metric class, which seems to work. However, the Sobol algorithm in the develop API seems to sample random points (multiple points) and then calculate the associated values at the same time. Is it possible to sample a point and calculate it once (only run once run_bdsim and get all the metrics), as the function (run_bdsim) I wrote can't calculate it in parallel?
   
3. The develop API and the service API are actually the same (I'm not sure if I'm understanding this correctly.), I can combine different models in the service API to achieve the same effect as tutorials of developer API on the website?
   For instance：
   gs = GenerationStrategy(steps=[GenerationStep(model=Models.SOBOL, num_trials=NUM_SOBOL_STEPS),
   GenerationStep(model=Models.MOO, num_trials=-1)])
   ax_client = AxClient(generation_strategy=gs, random_seed=1234, verbose_logging=True)
   The above code is same qNEHVI as the website：https://ax.dev/tutorials/multiobjective_optimization.html

I've attached my own code to the attachment.

from ax.plot.pareto_utils import compute_posterior_pareto_frontier
from ax.utils.notebook.plotting import render, init_notebook_plotting
import pandas as pd
from ax import *
import numpy as np
from ax.metrics.noisy_function import NoisyFunctionMetric
from ax.service.utils.report_utils import exp_to_df
from ax.runners.synthetic import SyntheticRunner
from ax.modelbridge.factory import get_MOO_EHVI, get_MOO_PAREGO
from ax.modelbridge.modelbridge_utils import observed_hypervolume
import os
import pybdsim

from ax.storage.json_store.save import save_experiment
def run_bdsim(parameters, magnets_name_list, gantry_components_file_name, particles_number):
    Q1G2 = parameters[0]
    Q2G2 = parameters[1]
    Q3G2 = parameters[2]
    Q4G2 = parameters[3]
    Q5G2 = parameters[4]
    Q6G2 = parameters[5]
    k_list = [Q1G2, Q2G2, Q3G2, Q4G2, Q5G2, Q6G2]
    with open(gantry_components_file_name, 'r', encoding='utf-8', errors='ignore') as fd:
        newText_gmad = fd.read()
        # replace degrader Thickness
        for i in range(6):
            start_name = newText_gmad.find(magnets_name_list[i])
            start = newText_gmad.find('k1', start_name)
            end = newText_gmad.find(',', start)
            comment_old = newText_gmad[start: end + 1]
            comment_new = "k1=" + str(k_list[i]) + ","
            newText_gmad = newText_gmad.replace(comment_old, comment_new)
    with open(gantry_components_file_name, 'w', encoding='utf-8', errors='ignore') as fd:
        fd.write(newText_gmad)
    pybdsim.Run.Bdsim(gmadpath="bdsim_model_1/gantry.gmad", outfile="output_1", ngenerate=particles_number, batch=True, silent=True)
    os.system("rebdsimOptics output_1.root output-optics_1.root --emittanceOnTheFly")


def get_beam_statistical(root_file_name, inintial_particles_number):
    load_data = pybdsim.Data.Load(root_file_name)
    transmission_efficiency_list = load_data.optics.GetColumn("Npart") / inintial_particles_number
    x_envelop_list = abs(load_data.optics.GetColumn("Sigma_x") * 1000)
    xp_envelop_list = abs(load_data.optics.GetColumn("Sigma_xp") * 1000)
    y_envelop_list = abs(load_data.optics.GetColumn("Sigma_y") * 1000)
    yp_envelop_list = abs(load_data.optics.GetColumn("Sigma_yp") * 1000)
    dispersion_x_list = abs(load_data.optics.GetColumn("Disp_x"))
    dispersion_xp_list = abs(load_data.optics.GetColumn("Disp_xp"))
    x_envelope_end = x_envelop_list[-1]
    y_envelope_end = y_envelop_list[-1]
    xp_envelope_end = xp_envelop_list[-1]
    yp_envelope_end = yp_envelop_list[-1]
    dispersion_x_end = dispersion_x_list[-2]
    dispersion_xp_end = dispersion_xp_list[-2]
    max_x_envelope = max(x_envelop_list)
    max_y_envelope = max(y_envelop_list)
    transmission_efficiency = transmission_efficiency_list[-1]
    return (transmission_efficiency, x_envelope_end, y_envelope_end, xp_envelope_end, yp_envelope_end, max_x_envelope,
            max_y_envelope, dispersion_x_end, dispersion_xp_end)


def inintial_magnet_parameters():
    element_list = ["GantryDR1", "GantryDR2", "GantryDR3", "Q1G2", "GantryDR4", "B1G21", "B1G22", "GantryDR5", "Q2G2",
                    "GantryDR6", "GantryDR7", "Q3G2", "GantryDR8", "GantryDR9", "Q4G2", "GantryDR10", "GantryDR11",
                    "Q5G2", "GantryDR12", "GantryDR13", "Q6G2", "GantryDR14", "B2G21", "B2G22",
                    "GantryDR15", "GantryDR16", "B3G21", "B3G22", "B3G23", "GantryDR17"]
    inintial_K_dict = {"Q1G2":-3.0182, "Q2G2":3.5537, "Q3G2":-5.4026, "Q4G2":6.5711, "Q5G2":-5.3012, "Q6G2":4.8104}
    magnets_name_list = ["Q1G2", "Q2G2", "Q3G2", "Q4G2", "Q5G2", "Q6G2"]
    return (element_list, inintial_K_dict, magnets_name_list)

def inintial_BO_parameters(inintial_K_dict, magnets_name_list):
    class Metric_transmission_efficiency(NoisyFunctionMetric):
        def f(self, x: np.ndarray) -> float:
            run_bdsim(x, magnets_name_list, "bdsim_model_1/gantry_components.gmad", 200)
            (transmission_efficiency, x_envelope_end, y_envelope_end, xp_envelope_end, yp_envelope_end, max_x_envelope, max_y_envelope, dispersion_x_end, dispersion_xp_end) \
                = get_beam_statistical("output-optics_1.root", 2000)
            return float(transmission_efficiency)

    class Metric_x_envelope_end(NoisyFunctionMetric):
        def f(self, x: np.ndarray) -> float:
            (transmission_efficiency, x_envelope_end, y_envelope_end, xp_envelope_end, yp_envelope_end, max_x_envelope,
             max_y_envelope, dispersion_x_end, dispersion_xp_end) \
                = get_beam_statistical("output-optics_1.root", 2000)
            return float(x_envelope_end)

    class Metric_y_envelope_end(NoisyFunctionMetric):
        def f(self, x: np.ndarray) -> float:

            (transmission_efficiency, x_envelope_end, y_envelope_end, xp_envelope_end, yp_envelope_end, max_x_envelope,
             max_y_envelope, dispersion_x_end, dispersion_xp_end) \
                = get_beam_statistical("output-optics_1.root", 2000)
            return float(y_envelope_end)

    class Metric_xp_envelope_end(NoisyFunctionMetric):
        def f(self, x: np.ndarray) -> float:
            (transmission_efficiency, x_envelope_end, y_envelope_end, xp_envelope_end, yp_envelope_end, max_x_envelope,
             max_y_envelope, dispersion_x_end, dispersion_xp_end) \
                = get_beam_statistical("output-optics_1.root", 2000)
            return float(xp_envelope_end)

    class Metric_yp_envelope_end(NoisyFunctionMetric):
        def f(self, x: np.ndarray) -> float:
            (transmission_efficiency, x_envelope_end, y_envelope_end, xp_envelope_end, yp_envelope_end, max_x_envelope,
             max_y_envelope, dispersion_x_end, dispersion_xp_end) \
                = get_beam_statistical("output-optics_1.root", 2000)
            return float(yp_envelope_end)

    class Metric_dispersion_x_end(NoisyFunctionMetric):
        def f(self, x: np.ndarray) -> float:
            (transmission_efficiency, x_envelope_end, y_envelope_end, xp_envelope_end, yp_envelope_end, max_x_envelope,
             max_y_envelope, dispersion_x_end, dispersion_xp_end) \
                = get_beam_statistical("output-optics_1.root", 2000)
            return float(dispersion_x_end)

    class Metric_dispersion_xp_end(NoisyFunctionMetric):
        def f(self, x: np.ndarray) -> float:

            (transmission_efficiency, x_envelope_end, y_envelope_end, xp_envelope_end, yp_envelope_end, max_x_envelope,
             max_y_envelope, dispersion_x_end, dispersion_xp_end) \
                = get_beam_statistical("output-optics_1.root", 2000)
            return float(dispersion_xp_end)

    # define search space
    parameter_list = []
    constant_rate = 0.1
    for i in range(len(magnets_name_list)):
        parameter_list.append(RangeParameter(name=magnets_name_list[i], lower=inintial_K_dict[magnets_name_list[i]]-abs(inintial_K_dict[magnets_name_list[i]])*constant_rate, upper= inintial_K_dict[magnets_name_list[i]]+abs(inintial_K_dict[magnets_name_list[i]])*constant_rate, parameter_type=ParameterType.FLOAT))
    search_space = SearchSpace(parameters=parameter_list)

    # instance metrics
    metric_transmission_efficiency = Metric_transmission_efficiency("transmission_efficiency", magnets_name_list, noise_sd=0.0, lower_is_better=False)
    metric_x_envelope_end = Metric_x_envelope_end("x_envelope_end", magnets_name_list, noise_sd=0.0, lower_is_better=True)
    metric_y_envelope_end = Metric_y_envelope_end("y_envelope_end", magnets_name_list, noise_sd=0.0, lower_is_better=True)
    metric_xp_envelope_end = Metric_xp_envelope_end("xp_envelope_end", magnets_name_list, noise_sd=0.0, lower_is_better=True)
    metric_yp_envelope_end = Metric_yp_envelope_end("yp_envelope_end", magnets_name_list, noise_sd=0.0, lower_is_better=True)
    metric_dispersion_x_end = Metric_dispersion_x_end("dispersion_x_end", magnets_name_list, noise_sd=0.0, lower_is_better=True)
    metric_dispersion_xp_end = Metric_dispersion_xp_end("dispersion_xp_end", magnets_name_list, noise_sd=0.0, lower_is_better=True)

    mo = MultiObjective(objectives=[Objective(metric=metric_transmission_efficiency, minimize=False),
                                    Objective(metric=metric_x_envelope_end, minimize=True),
                                    Objective(metric=metric_y_envelope_end, minimize=True),
                                    Objective(metric=metric_xp_envelope_end, minimize=True),
                                    Objective(metric=metric_yp_envelope_end, minimize=True),
                                    Objective(metric=metric_dispersion_x_end, minimize=True),
                                    Objective(metric=metric_dispersion_xp_end, minimize=True)])
    # bound is should to use prior experiences
    objective_thresholds = [ObjectiveThreshold(metric=metric_transmission_efficiency, bound=0.95, relative=False),
                            ObjectiveThreshold(metric=metric_x_envelope_end, bound=3.0, relative=False),
                            ObjectiveThreshold(metric=metric_y_envelope_end, bound=3.0, relative=False),
                            ObjectiveThreshold(metric=metric_xp_envelope_end, bound=2.5, relative=False),
                            ObjectiveThreshold(metric=metric_yp_envelope_end, bound=2.5, relative=False),
                            ObjectiveThreshold(metric=metric_dispersion_x_end, bound=0.1, relative=False),
                            ObjectiveThreshold(metric=metric_dispersion_xp_end, bound=0.1, relative=False)]

    optimization_config = MultiObjectiveOptimizationConfig(objective=mo, objective_thresholds=objective_thresholds)

    experiment = Experiment(name="gantry_experiment", search_space=search_space, optimization_config=optimization_config,
                            runner=SyntheticRunner(),
    )



    return experiment

def run_BO_Optimization(experiment, N_INIT, N_BATCH):
    def initialize_experiment(experiment, N_INIT):
        sobol = Models.SOBOL(search_space=experiment.search_space, seed=1234)

        for i in range(N_INIT):
            print(f"-------------SOBOL:{i}/{N_INIT}------------------")
            experiment.new_trial(sobol.gen(1)).run()
        return experiment.fetch_data()

    ehvi_experiment = experiment
    ehvi_data = initialize_experiment(ehvi_experiment, N_INIT)
    ehvi_hv_list = []
    for i in range(N_BATCH):
        print(f"--------------EHVI:{i}/{N_BATCH}------------------")
        ehvi_model = get_MOO_EHVI(experiment=ehvi_experiment, data=ehvi_data,)
        generator_run = ehvi_model.gen(1)
        trial = ehvi_experiment.new_trial(generator_run=generator_run)
        trial.run()
        ehvi_data = Data.from_multiple_data([ehvi_data, trial.fetch_data()])
        exp_df = exp_to_df(ehvi_experiment)
        try:
            hv = observed_hypervolume(modelbridge=ehvi_model)
        except:
            hv = 0
            print("Failed to compute hv")
        ehvi_hv_list.append(hv)
        print(f"Iteration: {i}, HV: {hv}")
    exp_df.to_excel("./data_save_1.xlsx")


if __name__ == '__main__':
    (element_list, inintial_K_dict, magnets_name_list) = inintial_magnet_parameters()
    experiment= inintial_BO_parameters(inintial_K_dict, magnets_name_list)
    run_BO_Optimization(experiment, 2, 2)

[lena-kashtelyan]

Hi @HUSTYuChen!

1. Please see this issue: https://github.com/facebook/Ax/issues?q=is%3Aissue+is%3Aopen+visualization for a workaround.
2. You could use the Service API, which would let you calculate the metrics however is most convenient for you, before submitting their data back via AxClient.complete_trial. Does this work for you? An alternative would be to implement a custom subclass of Metric with customization for batch-fetching, but I would just suggest going with the Service API route.
3. Not sure that I understand this:

The develop API and the service API are actually the same (I'm not sure if I'm understanding this correctly.)

But I can answer this question:

I can combine different models in the service API to achieve the same effect as tutorials of developer API on the website?

Yes, indeed you can! See examples of how to make a generation strategy here: https://ax.dev/tutorials/generation_strategy.html (we also recommend just using our auto-selection logic as illustrated in section 1B in that tutorial; this auto-selection will happen automatically if you use the Service API and don't specify a generation_strategy argument to AxClient).

So in short, if you use AxClient, you should not need to configure your own generation strategy at all. You can specify multiple objectives in AxClient.create_experiment, and it will select the qNEHVI acquisition function for you automatically. See example for that here at the top of the multi-objective tutorial: https://ax.dev/tutorials/multiobjective_optimization.html#Using-the-Service-API.

[HUSTYuChen]

Hi @lena-kashtelyan !
Thank you for answering my questions. I'm actually currently using Service Api to solve the problem. But there are a lot of advanced algorithms on tutorials that are implemented by Develop Api. I originally thought it could only be reproduced by Develop Api. Based on your answer, I now know that I can use Service Api to combine different mods to realize these advanced algorithms.
According to the relevant answer you provided me, the current built-in plotting functions are only supported to run in Jupyter Notebook, and it would be nice to be able to run them in Pycharm in the future.
Special thanks for taking time out of your busy schedule to answer my questions.

[lena-kashtelyan]

Hi again @HUSTYuChen, happy to help! Let us know if you run into any issues setting up the models you wanted to use, in Service API.

I actually don't know much about PyCharm, we'll look into it. I'll rename this issue accordingly and leave it open but put on the Wishlist, since the PyCharm aspect is the only open part of this I believe.