ax-js Jupyter Demo¶
This notebook demonstrates interactive GP (Gaussian Process) model diagnostics for a DTLZ2 multi-objective experiment — a standard benchmark with 5 input parameters and 2 objectives.
All plots are rendered client-side using ax-js, a TypeScript library that runs BoTorch GP predictions directly in the browser. The model hyperparameters are exported from Ax/BoTorch and loaded into a JavaScript predictor — no Python server needed after export.
Pre-populated outputs — you can view results without re-running. Or execute the cells to regenerate from scratch.
from ax.api import Client
from ax.api.configs import RangeParameterConfig
from botorch.test_functions.multi_objective import DTLZ2
import torch
dim = 5
client = Client()
client.configure_experiment(
name='dtlz2_demo',
parameters=[
RangeParameterConfig(name=f'x{i+1}', parameter_type='float', bounds=(0.0, 1.0))
for i in range(dim)
],
)
client.configure_optimization(objective='f0, f1')
fn = DTLZ2(dim=dim, num_objectives=2, negate=True)
for i in range(30):
trials = client.get_next_trials(max_trials=1)
for trial_index, params in trials.items():
x = torch.tensor([[params[f'x{j+1}'] for j in range(dim)]])
y = fn(x).squeeze()
client.complete_trial(trial_index=trial_index, raw_data={'f0': y[0].item(), 'f1': y[1].item()})
print(f'Completed 30 trials')
import sys; sys.path.insert(0, 'python')
from axjs_jupyter import (
slice_plot, response_surface,
feature_importance, cross_validation, optimization_trace,
)
1D Slice Plots¶
Each subplot shows the GP posterior prediction (mean ± confidence band) as a single parameter varies, with all other parameters held fixed. This lets you see how each input individually affects the selected outcome.
How dimensions are chosen: Parameters are sorted by feature importance (derived from the GP kernel's lengthscales — shorter lengthscale = more important). The most influential parameters appear first.
Reading the plot: The blue curve is the posterior mean prediction. The shaded band shows the 95% confidence interval. Training points are shown as dots — hover over them to highlight their nearest neighbors across all subplots.
Interactivity:
- Outcome selector — switch between objectives (f0, f1) to see how each parameter affects different outcomes
- Sliders — adjust the "held fixed" values for other parameters and watch the curves update in real time
- Pivoting — click anywhere on the mean curve to "pivot" to that point. This sets the clicked parameter's value in all other subplots' sliders, so you can explore the response surface from that operating point. For example, clicking at x1=0.7 on the x1 subplot will fix x1=0.7 in all other subplots, showing you what the model predicts when x1 is held at 0.7. Points that are more "distant" in terms of (relevant) inputs are shown as more transparent. At a more technical level, opacity is proportional to the correlation between the points under the given GP kernel lengthscales.
- Training dots — click a training point to snap all sliders to that trial's parameter values, exploring the model's predictions at that observed configuration
slice_plot(client)
2D Response Surface¶
A heatmap of the GP posterior mean over two input parameters, with all others held fixed. This shows how pairs of parameters jointly influence the outcome — revealing interactions that 1D slices can't capture.
How dimensions are chosen: The two most important parameters (by kernel lengthscale) are selected by default. You can change them with the axis dropdowns.
Reading the plot: Color encodes the predicted outcome value (yellow = high, purple = low). Training points are overlaid as dots. A second panel shows the posterior standard deviation (model uncertainty) — brighter regions indicate where the model is less certain, suggesting where additional trials would be most informative.
Interactivity:
- Axis dropdowns — choose which two parameters to plot
- Outcome selector — switch between objectives
- Sliders — adjust the held-fixed values for all non-plotted parameters
response_surface(client)
Feature Importance¶
A bar chart ranking parameters by their influence on the selected outcome.
How it works: Importance is derived from Sobol sensitivity analysis on the GP model. Each bar shows two components: the first-order effect (how much variance the parameter explains on its own) and the total-order effect (including interactions with other parameters). A large gap between them means the parameter's effect depends on the values of other parameters. These importances are also embedded in the sliders of the Ax Explorer.
Interactivity:
- Outcome selector — switch between objectives to see which parameters matter most for each
feature_importance(client)
Leave-One-Out Cross-Validation¶
Each training point is predicted by a GP model trained on all other points. Observed values are plotted against these LOO predictions — points on the diagonal indicate good model fit.
Reading the plot: Each point has a vertical error bar showing the GP's 95% confidence interval for that prediction. Points far from the diagonal, or where the diagonal falls outside the CI, indicate regions where the model struggles. The R² value summarizes overall fit quality.
Interactivity:
- Outcome selector — switch between objectives
- Hover — see the trial index and exact observed/predicted values
cross_validation(client)
Optimization Trace¶
The observed outcome value at each trial, with a running "best so far" line showing optimization progress.
Reading the plot: Each dot is one trial's observed value. The step line tracks the best value seen up to that trial. For multi-objective problems, "best" is determined per-outcome using the direction inferred from the optimization config (maximize or minimize).
Interactivity:
- Outcome selector — switch between objectives to see convergence for each
optimization_trace(client)
Five one-liners from ax.api.Client to interactive GP diagnostics. All rendering happens in the browser via ax-js — the GP posterior is evaluated in JavaScript using exported hyperparameters, so plots update instantly without round-trips to Python.