From 57aacafa83b4c471caa0a080f7f55d7e51cc2c69 Mon Sep 17 00:00:00 2001 From: rkhosrowshahi Date: Fri, 6 Nov 2026 14:54:28 -0500 Subject: [PATCH 1/6] test: Add Valley-Bowl, optimizer and LR options to trajectories Add the Valley-Bowl toy problem from CAGrad and Nash-MTL as the VB and VB01 objectives, and add --optimizer and --lr flags to optimize.py. --- tests/trajectories/_constants.py | 17 +++++++- tests/trajectories/_objectives.py | 65 +++++++++++++++++++++++++++++ tests/trajectories/_optimization.py | 12 +++++- tests/trajectories/optimize.py | 26 ++++++++++-- tests/trajectories/plot_params.py | 29 ++++++++++++- tests/trajectories/plot_values.py | 32 +++++++++++++- 6 files changed, 172 insertions(+), 9 deletions(-) diff --git a/tests/trajectories/_constants.py b/tests/trajectories/_constants.py index d4c25e9d8..0ed7a36b1 100644 --- a/tests/trajectories/_constants.py +++ b/tests/trajectories/_constants.py @@ -4,6 +4,7 @@ from torchjd.aggregation import ( IMTLG, MGDA, + Aggregator, AlignedMTL, CAGrad, DualProj, @@ -18,9 +19,10 @@ ElementWiseQuadratic, HomogenousQuadraticFunction, QuadraticFunction, + ValleyBowl, ) -AGGREGATORS = { +AGGREGATORS: dict[str, Aggregator] = { "upgrad": UPGrad(projector=QuadprogProjector(reg_eps=1e-7, norm_eps=1e-9)), "mgda": MGDA(), "cagrad": CAGrad(c=0.5), @@ -33,6 +35,10 @@ "random": Random(), "mean": Mean(), } +OPTIMIZERS: dict[str, type[torch.optim.SGD] | type[torch.optim.Adam]] = { + "sgd": torch.optim.SGD, + "adam": torch.optim.Adam, +} LR_MULTIPLIERS = { "upgrad": 1.0, "mgda": 2.0, @@ -91,6 +97,8 @@ } OBJECTIVES = { + "VB": ValleyBowl(), + "VB01": ValleyBowl(scale=0.1), "EWQ": ElementWiseQuadratic(2), "CQF": QuadraticFunction( As=[torch.tensor([[1.0, 0.2], [0.2, 0.05]]), torch.tensor([[3.0, -0.6], [-0.6, 0.2]])], @@ -103,11 +111,16 @@ ), } BASE_LEARNING_RATES = { + "VB": 0.001, + "VB01": 0.001, "EWQ": 0.075, "CQF": 0.125, "HQF": 0.005, } +_VB_INITIAL_POINTS = [[-8.5, 7.5], [0.0, 0.0], [9.0, 9.0], [-7.5, -0.5], [9.0, -1.0]] INITIAL_POINTS = { + "VB": _VB_INITIAL_POINTS, + "VB01": _VB_INITIAL_POINTS, "EWQ": [ [3.0, -2.0], [0.0, -3.0], @@ -128,6 +141,8 @@ ], } N_ITERS = { + "VB": 35000, + "VB01": 35000, "EWQ": 50, "CQF": 200, "HQF": 100, diff --git a/tests/trajectories/_objectives.py b/tests/trajectories/_objectives.py index 939f7e253..dffd3931f 100644 --- a/tests/trajectories/_objectives.py +++ b/tests/trajectories/_objectives.py @@ -1,6 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +from typing import cast import torch from torch import Tensor @@ -130,3 +131,67 @@ def __init__(self, n_values: int) -> None: def __call__(self, w: Tensor) -> Tensor: # noqa: ARG002 return torch.zeros(self.n_values) + + +class ValleyBowl(Objective): + """Two-objective toy problem from CAGrad (Liu et al., 2021), also used by Nash-MTL + (Navon et al., 2022). + + Above :math:`x_2 = 0`, each objective is a log-shaped valley; below it, a quadratic bowl. + + Reference implementations: + + - CAGrad: https://github.com/Cranial-XIX/CAGrad/blob/main/toy.py + - Nash-MTL: https://github.com/AvivNavon/nash-mtl/tree/main/experiments/toy + + :param scale: Multiplier for the first objective. + :param scale_both_losses: Multiplier applied to both objectives. + """ + + # Value used for the Pareto set in CAGrad's plotting code (``toy.py``, ``plot_contour``). + PARETO_SET_X2 = -8.3552 + + def __init__(self, scale: float = 1.0, scale_both_losses: float = 1.0) -> None: + super().__init__(n_params=2, n_values=2) + self.scale = scale + self.scale_both_losses = scale_both_losses + + def __call__(self, x: Tensor) -> Tensor: + if x.shape != (2,): + raise ValueError("x must have shape (2,).") + x1, x2 = x.unbind() + valley1 = (0.5 * (-x1 - 7) - torch.tanh(-x2)).abs().clamp_min(5e-6).log() + 6 + valley2 = (0.5 * (-x1 + 3) + torch.tanh(-x2) + 2).abs().clamp_min(5e-6).log() + 6 + above = torch.tanh(x2 * 0.5).clamp_min(0) + below = torch.tanh(-x2 * 0.5).clamp_min(0) + bowl1 = ((-x1 + 7).square() + 0.1 * (-x2 - 8).square()) / 10 - 20 + bowl2 = ((-x1 - 7).square() + 0.1 * (-x2 - 8).square()) / 10 - 20 + return ( + torch.stack( + ((valley1 * above + bowl1 * below) * self.scale, valley2 * above + bowl2 * below) + ) + * self.scale_both_losses + ) + + def jacobian(self, x: Tensor) -> Tensor: + return cast(Tensor, torch.func.jacrev(self)(x)) + + def __repr__(self) -> str: + return f"ValleyBowl(scale={self.scale}, scale_both_losses={self.scale_both_losses})" + + def pareto_set(self, n: int = 1000) -> Tensor: + """Sample the reference Pareto set used by the Valley-Bowl example. + + :param n: Number of points to sample along the reference set. + """ + return torch.stack( + (torch.linspace(-7.0, 7.0, n), torch.full((n,), self.PARETO_SET_X2)), dim=1 + ) + + @property + def optimum(self) -> Tensor: + """ + Pareto-optimal point that (approximately) minimizes the mean of the unscaled objectives. It + does not depend on ``scale`` or ``scale_both_losses``. + """ + return torch.tensor([0.0, self.PARETO_SET_X2]) diff --git a/tests/trajectories/_optimization.py b/tests/trajectories/_optimization.py index e7c0568f3..4a603a941 100644 --- a/tests/trajectories/_optimization.py +++ b/tests/trajectories/_optimization.py @@ -14,11 +14,21 @@ def optimize( aggregator: Aggregator, lr: float, n_iters: int, + optimizer_cls: type[torch.optim.SGD] | type[torch.optim.Adam] = torch.optim.SGD, ) -> tuple[list[Tensor], list[Tensor]]: + """Run an aggregator with the selected PyTorch optimizer. + + :param objective: Vector-valued objective to optimize. + :param initial_x: Starting point in parameter space. + :param aggregator: Aggregator applied to the objective Jacobian. + :param lr: Optimizer learning rate. + :param n_iters: Number of optimization steps. + :param optimizer_cls: Optimizer class to instantiate for this trajectory. + """ xs = [] ys = [] x = initial_x.clone().requires_grad_() - optimizer = torch.optim.SGD([x], lr=lr) + optimizer = optimizer_cls([x], lr=lr) for i in range(n_iters): print(f"{i + 1}/{n_iters}", end="\r") xs.append(x.detach().clone()) diff --git a/tests/trajectories/optimize.py b/tests/trajectories/optimize.py index cd58922c9..40e91fad0 100644 --- a/tests/trajectories/optimize.py +++ b/tests/trajectories/optimize.py @@ -3,7 +3,7 @@ spaces. Usage: - uv run python tests/trajectories/optimize.py ... + uv run python tests/trajectories/optimize.py ... [--optimizer sgd|adam] [--lr LR] Arguments: The key of the objective function (e.g., EWQ, CQF, HQF). @@ -28,6 +28,7 @@ LR_MULTIPLIERS, N_ITERS, OBJECTIVES, + OPTIMIZERS, ) from trajectories._optimization import optimize from trajectories._paths import get_params_dir, get_values_dir @@ -52,7 +53,19 @@ def main() -> None: metavar="aggregator", help=f"Keys of the aggregators to use. Choices: {list(AGGREGATORS)}", ) + parser.add_argument( + "--optimizer", + "--optim", + choices=list(OPTIMIZERS), + default="sgd", + help="Optimizer (default: sgd).", + ) + parser.add_argument( + "--lr", type=float, default=None, help="Use this learning rate for every aggregator." + ) args = parser.parse_args() + if args.lr is not None and (not np.isfinite(args.lr) or args.lr <= 0): + parser.error("--lr must be finite and positive") objective_key = args.objective if objective_key not in OBJECTIVES: @@ -71,8 +84,9 @@ def main() -> None: base_lr = BASE_LEARNING_RATES[objective_key] for key in LR_MULTIPLIERS: mult = lr_multiplier_overrides.get(key, LR_MULTIPLIERS[key]) - learning_rates[key] = mult * base_lr + learning_rates[key] = args.lr if args.lr is not None else mult * base_lr n_iters = N_ITERS[objective_key] + optimizer_cls = OPTIMIZERS[args.optimizer] torch.use_deterministic_algorithms(True) @@ -88,6 +102,7 @@ def main() -> None: "aggregator_reprs": {key: repr(AGGREGATORS[key]) for key in aggregator_keys}, "learning_rates": learning_rates, "initial_points": initial_points, + "optimizer": args.optimizer, } with open(TRAJECTORIES_RESULTS_DIR / objective_key / "metadata.json", "w") as f: json.dump(metadata, f) @@ -107,7 +122,12 @@ def main() -> None: initial_x = torch.tensor(initial_point) xs, ys = optimize( - objective, initial_x=initial_x, aggregator=aggregator, lr=lr, n_iters=n_iters + objective, + initial_x=initial_x, + aggregator=aggregator, + lr=lr, + n_iters=n_iters, + optimizer_cls=optimizer_cls, ) xs_list.append(torch.stack(xs)) diff --git a/tests/trajectories/plot_params.py b/tests/trajectories/plot_params.py index c58840358..896a0e2fc 100644 --- a/tests/trajectories/plot_params.py +++ b/tests/trajectories/plot_params.py @@ -13,6 +13,7 @@ import matplotlib.pyplot as plt import numpy as np +import torch from tests.paths import TRAJECTORIES_RESULTS_DIR from trajectories._constants import ( @@ -22,7 +23,7 @@ LATEX_NAMES, OBJECTIVES, ) -from trajectories._objectives import ElementWiseQuadratic, WithSPSMappingMixin +from trajectories._objectives import ElementWiseQuadratic, ValleyBowl, WithSPSMappingMixin from trajectories._optimization import compute_gradient_cosine_similarities from trajectories._pareto_utils import sample_2d_spss from trajectories._paths import get_param_plots_dir, get_params_dir @@ -84,6 +85,11 @@ def main() -> None: common_plotter = SquareBoxAspectSetter() + if isinstance(objective, ValleyBowl): + sps_points = objective.pareto_set().numpy() + main_content = np.concatenate([main_content, sps_points]) + common_plotter += SPSPlotter(sps_points) + if objective.n_values == 2 and isinstance(objective, WithSPSMappingMixin): sps_points = sample_2d_spss(objective).numpy() main_content = np.concatenate([main_content, sps_points]) @@ -97,7 +103,22 @@ def main() -> None: adjust_plotter = ContentLimAdjuster(main_content) common_plotter += adjust_plotter - if objective.n_values == 2: + if isinstance(objective, ValleyBowl): + x0 = torch.linspace(adjust_plotter.xlim[0], adjust_plotter.xlim[1], 200) + x1 = torch.linspace(adjust_plotter.ylim[0], adjust_plotter.ylim[1], 200) + grid = torch.cartesian_prod(x0, x1) + mean_loss = torch.vmap(objective)(grid).mean(dim=1).reshape(200, 200) + common_plotter += HeatmapPlotter( + values=mean_loss.numpy(), + x_min=adjust_plotter.xlim[0], + x_max=adjust_plotter.xlim[1], + y_min=adjust_plotter.ylim[0], + y_max=adjust_plotter.ylim[1], + vmin=float(mean_loss.min()), + vmax=float(mean_loss.max()), + cmap="viridis", + ) + elif objective.n_values == 2: similarities = compute_gradient_cosine_similarities( objective, x0_min=adjust_plotter.xlim[0], @@ -147,6 +168,10 @@ def main() -> None: plotter += YAxisLabeller("$x_2$") if j == 0 else YTicksClearer() plotter(axes[i][j]) + if isinstance(objective, ValleyBowl): + axes[i][j].scatter( + *objective.optimum.tolist(), marker="*", color="black", s=45, zorder=10 + ) fig.tight_layout(h_pad=-2.5) print("Saving figure") diff --git a/tests/trajectories/plot_values.py b/tests/trajectories/plot_values.py index 145aaf1cf..3c44a0f05 100644 --- a/tests/trajectories/plot_values.py +++ b/tests/trajectories/plot_values.py @@ -13,6 +13,7 @@ import matplotlib.pyplot as plt import numpy as np +import torch from tests.paths import TRAJECTORIES_RESULTS_DIR from trajectories._constants import ( @@ -22,7 +23,7 @@ OBJECTIVES, PLOT_VALUES_LIMS, ) -from trajectories._objectives import WithSPSMappingMixin +from trajectories._objectives import ValleyBowl, WithSPSMappingMixin from trajectories._pareto_utils import compute_normalized_2d_pf_distances, sample_2d_pf from trajectories._paths import get_value_plots_dir, get_values_dir from trajectories._plotters import ( @@ -82,6 +83,12 @@ def main() -> None: initial_values = first_agg_Y[:, 0, :] main_content = initial_values + if isinstance(objective, ValleyBowl): + front = torch.vmap(objective)(objective.pareto_set()) + pf_points_array = front.numpy() + common_plotter += PFPlotter(pf_points_array) + main_content = np.concatenate([main_content, pf_points_array]) + if isinstance(objective, WithSPSMappingMixin): pf_points_array = sample_2d_pf(objective).numpy() common_plotter += PFPlotter(pf_points_array) @@ -98,7 +105,25 @@ def main() -> None: xlim = adjust_plotter.xlim ylim = adjust_plotter.ylim - if isinstance(objective, WithSPSMappingMixin): + if isinstance(objective, ValleyBowl): + grid = torch.cartesian_prod( + torch.linspace(xlim[0], xlim[1], 200), torch.linspace(ylim[0], ylim[1], 200) + ) + distances = torch.cat( + [torch.cdist(chunk, front).min(dim=1).values for chunk in grid.split(512)] + ) + distances = distances.reshape(200, 200) + common_plotter += HeatmapPlotter( + values=distances.numpy(), + x_min=xlim[0], + x_max=xlim[1], + y_min=ylim[0], + y_max=ylim[1], + vmin=0, + vmax=float(distances.max()), + cmap="magma_r", + ) + elif isinstance(objective, WithSPSMappingMixin): distances = compute_normalized_2d_pf_distances( objective, y0_min=xlim[0], @@ -145,6 +170,9 @@ def main() -> None: plotter += YAxisLabeller("Objective $2$") if j == 0 else YTicksClearer() plotter(axes[i][j]) + if isinstance(objective, ValleyBowl): + optimum = objective(objective.optimum).numpy() + axes[i][j].scatter(*optimum, marker="*", color="black", s=45, zorder=10) fig.tight_layout(h_pad=-2.5) print("Saving figure") From 42a062b5c05d85b9c273abbe51c57c2f6978d50b Mon Sep 17 00:00:00 2001 From: rkhosrowshahi Date: Mon, 9 Nov 2026 14:56:58 -0500 Subject: [PATCH 2/6] Remove optimum marker from Valley-Bowl plots --- tests/trajectories/_objectives.py | 8 -------- tests/trajectories/plot_params.py | 4 ---- tests/trajectories/plot_values.py | 3 --- 3 files changed, 15 deletions(-) diff --git a/tests/trajectories/_objectives.py b/tests/trajectories/_objectives.py index dffd3931f..55401ef19 100644 --- a/tests/trajectories/_objectives.py +++ b/tests/trajectories/_objectives.py @@ -187,11 +187,3 @@ def pareto_set(self, n: int = 1000) -> Tensor: return torch.stack( (torch.linspace(-7.0, 7.0, n), torch.full((n,), self.PARETO_SET_X2)), dim=1 ) - - @property - def optimum(self) -> Tensor: - """ - Pareto-optimal point that (approximately) minimizes the mean of the unscaled objectives. It - does not depend on ``scale`` or ``scale_both_losses``. - """ - return torch.tensor([0.0, self.PARETO_SET_X2]) diff --git a/tests/trajectories/plot_params.py b/tests/trajectories/plot_params.py index 896a0e2fc..92fc735a4 100644 --- a/tests/trajectories/plot_params.py +++ b/tests/trajectories/plot_params.py @@ -168,10 +168,6 @@ def main() -> None: plotter += YAxisLabeller("$x_2$") if j == 0 else YTicksClearer() plotter(axes[i][j]) - if isinstance(objective, ValleyBowl): - axes[i][j].scatter( - *objective.optimum.tolist(), marker="*", color="black", s=45, zorder=10 - ) fig.tight_layout(h_pad=-2.5) print("Saving figure") diff --git a/tests/trajectories/plot_values.py b/tests/trajectories/plot_values.py index 3c44a0f05..bf1222e94 100644 --- a/tests/trajectories/plot_values.py +++ b/tests/trajectories/plot_values.py @@ -170,9 +170,6 @@ def main() -> None: plotter += YAxisLabeller("Objective $2$") if j == 0 else YTicksClearer() plotter(axes[i][j]) - if isinstance(objective, ValleyBowl): - optimum = objective(objective.optimum).numpy() - axes[i][j].scatter(*optimum, marker="*", color="black", s=45, zorder=10) fig.tight_layout(h_pad=-2.5) print("Saving figure") From b02cf21fca635bd25500f7f5ac8121d91666e4b8 Mon Sep 17 00:00:00 2001 From: rkhosrowshahi Date: Mon, 9 Nov 2026 16:51:03 -0500 Subject: [PATCH 3/6] Draw Valley-Bowl backgrounds as filled contours --- tests/trajectories/_plotters.py | 44 +++++++++++++++++++++++++++++++ tests/trajectories/plot_params.py | 6 ++--- tests/trajectories/plot_values.py | 6 ++--- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/tests/trajectories/_plotters.py b/tests/trajectories/_plotters.py index 21b8157c2..04cf34756 100644 --- a/tests/trajectories/_plotters.py +++ b/tests/trajectories/_plotters.py @@ -312,6 +312,50 @@ def __call__(self, ax: plt.Axes) -> None: ) +class FilledContourPlotter(Plotter): + """ + Plotter that draws filled contours of the given values on a regular grid extending between the + provided coordinates. + + :param values: Array of shape ``[n_x, n_y]``, where ``values[i, j]`` is the value at the i-th x + coordinate and the j-th y coordinate. + :param x_min: Smallest x coordinate of the grid. + :param x_max: Largest x coordinate of the grid. + :param y_min: Smallest y coordinate of the grid. + :param y_max: Largest y coordinate of the grid. + :param cmap: Name of the matplotlib colormap. + :param alpha: Opacity of the filled contours. + :param levels: Number of contour levels. + """ + + def __init__( + self, + values: np.ndarray, + x_min: float, + x_max: float, + y_min: float, + y_max: float, + cmap: str, + alpha: float, + levels: int = 35, + ) -> None: + self.values = values + self.x_min = x_min + self.x_max = x_max + self.y_min = y_min + self.y_max = y_max + self.cmap = cmap + self.alpha = alpha + self.levels = levels + + def __call__(self, ax: plt.Axes) -> None: + xs = np.linspace(self.x_min, self.x_max, self.values.shape[0]) + ys = np.linspace(self.y_min, self.y_max, self.values.shape[1]) + ax.contourf( + xs, ys, self.values.T, levels=self.levels, cmap=self.cmap, alpha=self.alpha, zorder=0 + ) + + class LimAdjuster(Plotter): """Plotter that adjusts the xlim and ylim of the plot to the specified xlim and ylim.""" diff --git a/tests/trajectories/plot_params.py b/tests/trajectories/plot_params.py index 92fc735a4..76b23e7de 100644 --- a/tests/trajectories/plot_params.py +++ b/tests/trajectories/plot_params.py @@ -31,6 +31,7 @@ AxesPlotter, ContentLimAdjuster, ContourCirclesPlotter, + FilledContourPlotter, HeatmapPlotter, LimAdjuster, MultiTrajPlotter, @@ -108,15 +109,14 @@ def main() -> None: x1 = torch.linspace(adjust_plotter.ylim[0], adjust_plotter.ylim[1], 200) grid = torch.cartesian_prod(x0, x1) mean_loss = torch.vmap(objective)(grid).mean(dim=1).reshape(200, 200) - common_plotter += HeatmapPlotter( + common_plotter += FilledContourPlotter( values=mean_loss.numpy(), x_min=adjust_plotter.xlim[0], x_max=adjust_plotter.xlim[1], y_min=adjust_plotter.ylim[0], y_max=adjust_plotter.ylim[1], - vmin=float(mean_loss.min()), - vmax=float(mean_loss.max()), cmap="viridis", + alpha=0.55, ) elif objective.n_values == 2: similarities = compute_gradient_cosine_similarities( diff --git a/tests/trajectories/plot_values.py b/tests/trajectories/plot_values.py index bf1222e94..b78c88e7c 100644 --- a/tests/trajectories/plot_values.py +++ b/tests/trajectories/plot_values.py @@ -28,6 +28,7 @@ from trajectories._paths import get_value_plots_dir, get_values_dir from trajectories._plotters import ( ContentLimAdjuster, + FilledContourPlotter, HeatmapPlotter, LimAdjuster, MultiTrajPlotter, @@ -113,15 +114,14 @@ def main() -> None: [torch.cdist(chunk, front).min(dim=1).values for chunk in grid.split(512)] ) distances = distances.reshape(200, 200) - common_plotter += HeatmapPlotter( + common_plotter += FilledContourPlotter( values=distances.numpy(), x_min=xlim[0], x_max=xlim[1], y_min=ylim[0], y_max=ylim[1], - vmin=0, - vmax=float(distances.max()), cmap="magma_r", + alpha=0.5, ) elif isinstance(objective, WithSPSMappingMixin): distances = compute_normalized_2d_pf_distances( From afb1ee55b8309af5002193de0b7d34f7473cee68 Mon Sep 17 00:00:00 2001 From: rkhosrowshahi Date: Mon, 9 Nov 2026 17:08:33 -0500 Subject: [PATCH 4/6] Draw contour lines on Valley-Bowl backgrounds --- tests/trajectories/_plotters.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/trajectories/_plotters.py b/tests/trajectories/_plotters.py index 04cf34756..0ed96ffa6 100644 --- a/tests/trajectories/_plotters.py +++ b/tests/trajectories/_plotters.py @@ -315,7 +315,7 @@ def __call__(self, ax: plt.Axes) -> None: class FilledContourPlotter(Plotter): """ Plotter that draws filled contours of the given values on a regular grid extending between the - provided coordinates. + provided coordinates, with thin contour lines at the boundaries between levels. :param values: Array of shape ``[n_x, n_y]``, where ``values[i, j]`` is the value at the i-th x coordinate and the j-th y coordinate. @@ -354,6 +354,11 @@ def __call__(self, ax: plt.Axes) -> None: ax.contourf( xs, ys, self.values.T, levels=self.levels, cmap=self.cmap, alpha=self.alpha, zorder=0 ) + # Draw the level boundaries explicitly, since vector outputs (e.g. PDF) render adjacent + # filled levels without any visible seam. + ax.contour( + xs, ys, self.values.T, levels=self.levels, cmap=self.cmap, linewidths=0.3, zorder=0 + ) class LimAdjuster(Plotter): From ce81897ad4a5d03cb89ce18a38ce3902a1d9f9f0 Mon Sep 17 00:00:00 2001 From: rkhosrowshahi Date: Mon, 9 Nov 2026 20:12:32 -0500 Subject: [PATCH 5/6] Use 6:5 subplots and fixed Valley-Bowl axis limits Make the subplots slightly wider than tall (box aspect 5/6), so that matplotlib's automatic ticks on the x-axis are as complete as on the y-axis instead of skipping values. --- tests/trajectories/_constants.py | 20 ++++++++++++++++++++ tests/trajectories/_plotters.py | 6 +++--- tests/trajectories/plot_distance_to_pf.py | 4 ++-- tests/trajectories/plot_params.py | 12 ++++++++---- tests/trajectories/plot_values.py | 6 +++--- 5 files changed, 36 insertions(+), 12 deletions(-) diff --git a/tests/trajectories/_constants.py b/tests/trajectories/_constants.py index 0ed7a36b1..2bfe5c018 100644 --- a/tests/trajectories/_constants.py +++ b/tests/trajectories/_constants.py @@ -94,6 +94,26 @@ "xlim": (-0.125, 2.625), "ylim": (-0.425, 8.925), }, + "VB": { + "xlim": (-21.0, 11.0), + "ylim": (-21.0, 11.0), + }, + "VB01": { + "xlim": (-2.1, 1.1), + "ylim": (-21.0, 11.0), + }, +} + +# Fixed xlim and ylim of the param plot, used instead of limits fitted to the content +PLOT_PARAMS_LIMS = { + "VB": { + "xlim": (-10.5, 10.5), + "ylim": (-10.5, 10.5), + }, + "VB01": { + "xlim": (-10.5, 10.5), + "ylim": (-10.5, 10.5), + }, } OBJECTIVES = { diff --git a/tests/trajectories/_plotters.py b/tests/trajectories/_plotters.py index 0ed96ffa6..3a0f76b1e 100644 --- a/tests/trajectories/_plotters.py +++ b/tests/trajectories/_plotters.py @@ -432,8 +432,8 @@ def __call__(self, ax: plt.Axes) -> None: ax.set_title(self.title) -class SquareBoxAspectSetter(Plotter): - """Plotter that sets a square box aspect.""" +class BoxAspectSetter(Plotter): + """Plotter that sets a slightly wider than tall box aspect.""" def __call__(self, ax: plt.Axes) -> None: - ax.set_box_aspect(1) + ax.set_box_aspect(5 / 6) diff --git a/tests/trajectories/plot_distance_to_pf.py b/tests/trajectories/plot_distance_to_pf.py index c8d0eee2c..a85785cad 100644 --- a/tests/trajectories/plot_distance_to_pf.py +++ b/tests/trajectories/plot_distance_to_pf.py @@ -20,8 +20,8 @@ from trajectories._pareto_utils import make_2d_pf_distance_fn from trajectories._paths import get_distance_to_pf_plots_dir, get_values_dir from trajectories._plotters import ( + BoxAspectSetter, MultiEvolutionPlotter, - SquareBoxAspectSetter, TitleSetter, XAxisLabeller, XTicksClearer, @@ -65,7 +65,7 @@ def main() -> None: objective_key = metadata["objective_key"] objective = OBJECTIVES[objective_key] - common_plotter = SquareBoxAspectSetter() + common_plotter = BoxAspectSetter() aggregator_keys = metadata["aggregator_keys"] aggregator_to_Y = {key: np.load(values_dir / f"{key}.npy") for key in aggregator_keys} diff --git a/tests/trajectories/plot_params.py b/tests/trajectories/plot_params.py index 76b23e7de..963569f38 100644 --- a/tests/trajectories/plot_params.py +++ b/tests/trajectories/plot_params.py @@ -22,6 +22,7 @@ INITIAL_POINTS, LATEX_NAMES, OBJECTIVES, + PLOT_PARAMS_LIMS, ) from trajectories._objectives import ElementWiseQuadratic, ValleyBowl, WithSPSMappingMixin from trajectories._optimization import compute_gradient_cosine_similarities @@ -29,6 +30,7 @@ from trajectories._paths import get_param_plots_dir, get_params_dir from trajectories._plotters import ( AxesPlotter, + BoxAspectSetter, ContentLimAdjuster, ContourCirclesPlotter, FilledContourPlotter, @@ -36,7 +38,6 @@ LimAdjuster, MultiTrajPlotter, SPSPlotter, - SquareBoxAspectSetter, TitleSetter, XAxisLabeller, XTicksClearer, @@ -84,7 +85,7 @@ def main() -> None: initial_points_array = np.stack([np.array(point) for point in initial_points]) main_content = initial_points_array - common_plotter = SquareBoxAspectSetter() + common_plotter = BoxAspectSetter() if isinstance(objective, ValleyBowl): sps_points = objective.pareto_set().numpy() @@ -101,7 +102,10 @@ def main() -> None: common_plotter += ContourCirclesPlotter() common_plotter += LimAdjuster(xlim=(-5.0, 5.0), ylim=(-5.0, 5.0)) else: - adjust_plotter = ContentLimAdjuster(main_content) + if objective_key in PLOT_PARAMS_LIMS: + adjust_plotter = LimAdjuster(**PLOT_PARAMS_LIMS[objective_key]) + else: + adjust_plotter = ContentLimAdjuster(main_content) common_plotter += adjust_plotter if isinstance(objective, ValleyBowl): @@ -144,7 +148,7 @@ def main() -> None: n_aggregators = len(aggregator_keys) n_rows, n_cols = compute_subplot_layout(n_aggregators) key_to_index = map_orders_to_indices(aggregator_keys, AGGREGATOR_ORDER) - fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 2, n_rows * 2.5)) + fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 3, n_rows * 3)) if n_rows == n_cols == 1: axes = np.array([[axes]]) elif n_rows == 1: diff --git a/tests/trajectories/plot_values.py b/tests/trajectories/plot_values.py index b78c88e7c..0262b5918 100644 --- a/tests/trajectories/plot_values.py +++ b/tests/trajectories/plot_values.py @@ -27,13 +27,13 @@ from trajectories._pareto_utils import compute_normalized_2d_pf_distances, sample_2d_pf from trajectories._paths import get_value_plots_dir, get_values_dir from trajectories._plotters import ( + BoxAspectSetter, ContentLimAdjuster, FilledContourPlotter, HeatmapPlotter, LimAdjuster, MultiTrajPlotter, PFPlotter, - SquareBoxAspectSetter, TitleSetter, XAxisLabeller, XTicksClearer, @@ -77,7 +77,7 @@ def main() -> None: if objective.n_values != 2: raise ValueError("Can only plot values trajectories for objectives with 2 values.") - common_plotter = SquareBoxAspectSetter() + common_plotter = BoxAspectSetter() aggregator_keys = metadata["aggregator_keys"] aggregator_to_Y = {key: np.load(values_dir / f"{key}.npy") for key in aggregator_keys} first_agg_Y = next(iter(aggregator_to_Y.values())) @@ -146,7 +146,7 @@ def main() -> None: n_aggregators = len(aggregator_keys) n_rows, n_cols = compute_subplot_layout(n_aggregators) key_to_index = map_orders_to_indices(aggregator_keys, AGGREGATOR_ORDER) - fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 2, n_rows * 2.5)) + fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 3, n_rows * 3)) if n_rows == n_cols == 1: axes = np.array([[axes]]) elif n_rows == 1: From 1f7ee35f32ca19506059e1656c8771d581e86a74 Mon Sep 17 00:00:00 2001 From: rkhosrowshahi Date: Mon, 9 Nov 2026 20:16:05 -0500 Subject: [PATCH 6/6] Write Valley-Bowl initial points one per line --- tests/trajectories/_constants.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/tests/trajectories/_constants.py b/tests/trajectories/_constants.py index 2bfe5c018..dbfd1ac70 100644 --- a/tests/trajectories/_constants.py +++ b/tests/trajectories/_constants.py @@ -137,10 +137,21 @@ "CQF": 0.125, "HQF": 0.005, } -_VB_INITIAL_POINTS = [[-8.5, 7.5], [0.0, 0.0], [9.0, 9.0], [-7.5, -0.5], [9.0, -1.0]] INITIAL_POINTS = { - "VB": _VB_INITIAL_POINTS, - "VB01": _VB_INITIAL_POINTS, + "VB": [ + [-8.5, 7.5], + [0.0, 0.0], + [9.0, 9.0], + [-7.5, -0.5], + [9.0, -1.0], + ], + "VB01": [ + [-8.5, 7.5], + [0.0, 0.0], + [9.0, 9.0], + [-7.5, -0.5], + [9.0, -1.0], + ], "EWQ": [ [3.0, -2.0], [0.0, -3.0],