test: Clean up trajectories

CONTRIBUTING.md

@@ -234,9 +234,9 @@ The `tests/trajectories/` directory contains scripts to generate and visualize o
 trajectories using various aggregators on simple multi-objective problems. They require the `plot`
 dependency group.
 
-Available objective keys: `EWQ`, `CQF`, `CQF2`, `HQF`, `MN2`, `MN20`.
+Available objective keys: `EWQ`, `CQF`, `HQF`.
 
-Available aggregator keys: `upgrad`, `mgda`, `cagrad`, `nashmtl`, `nashmtl20`, `graddrop`,
+Available aggregator keys: `upgrad`, `mgda`, `cagrad`, `nashmtl`, `graddrop`,
 `imtl_g`, `aligned_mtl`, `dualproj`, `pcgrad`, `random`, `mean`.
 
 **Step 1 — Optimize:** run the optimization for an objective and a selection of aggregators:
@@ -253,8 +253,24 @@ uv run python tests/trajectories/plot_values.py EWQ
 uv run python tests/trajectories/plot_distance_to_pf.py EWQ
 ```
 
-Replace `EWQ` with any other objective key. The three plot scripts produce PDFs saved to
-`tests/trajectories/results/<objective>/`.
+To run everything:
+```bash
+export MPLBACKEND=Agg
+uv run python tests/trajectories/optimize.py EWQ upgrad mean mgda cagrad dualproj graddrop imtl_g aligned_mtl nashmtl random
+uv run python tests/trajectories/plot_params.py EWQ
+uv run python tests/trajectories/plot_values.py EWQ
+uv run python tests/trajectories/plot_distance_to_pf.py EWQ
+uv run python tests/trajectories/optimize.py CQF upgrad mean mgda cagrad dualproj graddrop imtl_g aligned_mtl nashmtl random
+uv run python tests/trajectories/plot_params.py CQF
+uv run python tests/trajectories/plot_values.py CQF
+uv run python tests/trajectories/plot_distance_to_pf.py CQF
+uv run python tests/trajectories/optimize.py HQF upgrad mean mgda cagrad dualproj graddrop imtl_g aligned_mtl nashmtl random
+uv run python tests/trajectories/plot_params.py HQF
+uv run python tests/trajectories/plot_values.py HQF
+uv run python tests/trajectories/plot_distance_to_pf.py HQF
+```
+
+The three plot scripts produce PDFs saved to `tests/trajectories/results/<objective>/`.
 
 > [!NOTE]
 > The plot scripts require a LaTeX installation for rendering:

tests/trajectories/_constants.py

@@ -1,6 +1,3 @@
-from math import cos, sin
-
-import numpy as np
 import torch
 
 from torchjd._linalg import QuadprogProjector
@@ -18,19 +15,16 @@
     UPGrad,
 )
 from trajectories._objectives import (
-    ConvexQuadraticForm,
     ElementWiseQuadratic,
-    HomogenousQuadraticForm,
-    Multinorm,
-    QuadraticForm,
+    HomogenousQuadraticFunction,
+    QuadraticFunction,
 )
 
 AGGREGATORS = {
     "upgrad": UPGrad(projector=QuadprogProjector(reg_eps=1e-7, norm_eps=1e-9)),
     "mgda": MGDA(),
     "cagrad": CAGrad(c=0.5),
     "nashmtl": NashMTL(n_tasks=2, optim_niter=1),
-    "nashmtl20": NashMTL(n_tasks=20, optim_niter=1),
     "graddrop": GradDrop(),
     "imtl_g": IMTLG(),
     "aligned_mtl": AlignedMTL(),
@@ -44,7 +38,6 @@
     "mgda": 2.0,
     "cagrad": 1.0,
     "nashmtl": 2.0,
-    "nashmtl20": 2.0,
     "graddrop": 0.5,
     "imtl_g": 1.0,
     "aligned_mtl": 4.0,
@@ -61,14 +54,12 @@
         "imtl_g": 2.0,
     },
     "CQF": {"nashmtl": 0.5},
-    "CQF2": {"nashmtl": 0.5},
 }
 AGGREGATOR_ORDER = {
     "upgrad": 9,
     "mgda": 1,
     "cagrad": 5,
     "nashmtl": 7,
-    "nashmtl20": 7,
     "graddrop": 3,
     "imtl_g": 4,
     "aligned_mtl": 8,
@@ -82,7 +73,6 @@
     "mgda": r"$\mathcal A_{\mathrm{MGDA}}$",
     "cagrad": r"$\mathcal A_{\mathrm{CAGrad}}$",
     "nashmtl": r"$\mathcal A_{\mathrm{Nash-MTL}}$",
-    "nashmtl20": r"$\mathcal A_{\mathrm{Nash-MTL}}$",
     "graddrop": r"$\mathcal A_{\mathrm{GradDrop}}$",
     "imtl_g": r"$\mathcal A_{\mathrm{IMTL-G}}$",
     "aligned_mtl": r"$\mathcal A_{\mathrm{Aligned-MTL}}$",
@@ -98,44 +88,24 @@
         "xlim": (-0.125, 2.625),
         "ylim": (-0.425, 8.925),
     },
-    "CQF2": {
-        "xlim": (-0.125, 2.625),
-        "ylim": (-0.425, 8.925),
-    },
 }
 
-THETA = np.pi / 16
-
 OBJECTIVES = {
     "EWQ": ElementWiseQuadratic(2),
-    "CQF": ConvexQuadraticForm(
-        Bs=[
-            torch.tensor([[cos(THETA), -sin(THETA)], [sin(THETA), cos(THETA)]])
-            @ torch.diag(torch.tensor([1.0, 0.1])),
-            torch.tensor([[cos(THETA), sin(THETA)], [-sin(THETA), cos(THETA)]])
-            @ torch.diag(torch.tensor([torch.sqrt(torch.tensor(3.0)), 0.1])),
-        ],
-        us=[torch.tensor([1.0, 0.0]), torch.tensor([-1.0, 0.0])],
-    ),
-    "CQF2": QuadraticForm(
+    "CQF": QuadraticFunction(
         As=[torch.tensor([[1.0, 0.2], [0.2, 0.05]]), torch.tensor([[3.0, -0.6], [-0.6, 0.2]])],
         us=[torch.tensor([1.0, 0.0]), torch.tensor([-1.0, 0.0])],
     ),
-    "HQF": HomogenousQuadraticForm(
+    "HQF": HomogenousQuadraticFunction(
         A=torch.tensor([[2.0, -1.0], [-1.0, 2.0]]),
         scales=torch.tensor([1.0, 10.0]),
         us=[torch.tensor([1.0, 0.0]), torch.tensor([-10.0, 0.0])],
     ),
-    "MN2": Multinorm(torch.tensor([1.0, 10.0])),
-    "MN20": Multinorm(torch.arange(1, 21)),
 }
 BASE_LEARNING_RATES = {
     "EWQ": 0.075,
     "CQF": 0.125,
-    "CQF2": 0.125,
     "HQF": 0.005,
-    "MN2": 0.02,
-    "MN20": 0.005,
 }
 INITIAL_POINTS = {
     "EWQ": [
@@ -146,12 +116,6 @@
         [-3.5, -0.75],
     ],
     "CQF": [
-        [0.5, 0.5],
-        [-1.0, 7.0],
-        [0.0, 0.0],
-        [1.0, 6.0],
-    ],
-    "CQF2": [
         [0.5, 0.5],
         [-0.3, 7.0],
         [0.0, 0.0],
@@ -162,22 +126,9 @@
         [1.5, 2.0],
         [2.5, 5.5],
     ],
-    "MN2": [
-        [0.0, 0.0],
-        [-5.0, 5.0],
-        [10.0, 5.0],
-        [10.0, 0.0],
-        [20.0, 0.0],
-    ],
-    "MN20": [
-        [0.0] * 20,
-    ],
 }
 N_ITERS = {
     "EWQ": 50,
     "CQF": 200,
-    "CQF2": 200,
     "HQF": 100,
-    "MN2": 50,
-    "MN20": 500,
 }

tests/trajectories/_objectives.py

@@ -45,7 +45,7 @@ def sps_mapping(self) -> "WithSPSMappingMixin.SPSMapping":
         pass
 
 
-class QuadraticForm(Objective, WithSPSMappingMixin):
+class QuadraticFunction(Objective, WithSPSMappingMixin):
     def __init__(self, As: list[Tensor], us: list[Tensor]) -> None:
         if len(As) != len(us):
             raise ValueError("As and us must have the same length.")
@@ -86,11 +86,11 @@ def __call__(self, w: Tensor) -> Tensor:
             return torch.linalg.lstsq(G, b, driver="gelsd").solution
 
     @property
-    def sps_mapping(self) -> "QuadraticForm.SPSMapping":
+    def sps_mapping(self) -> "QuadraticFunction.SPSMapping":
         return self.SPSMapping(self.As, self.us)
 
 
-class HomogenousQuadraticForm(QuadraticForm):
+class HomogenousQuadraticFunction(QuadraticFunction):
     def __init__(self, A: Tensor, scales: Tensor, us: list[Tensor]) -> None:
         self.A = A
         self.scales = scales
@@ -101,15 +101,6 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}(A={self.A}, scales={self.scales}, us={self.us})"
 
 
-class ConvexQuadraticForm(QuadraticForm):
-    def __init__(self, Bs: list[Tensor], us: list[Tensor]) -> None:
-        self.Bs = Bs
-        super().__init__(As=[B @ B.T for B in self.Bs], us=us)
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(Bs={self.Bs}, us={self.us})"
-
-
 class ElementWiseQuadratic(Objective, WithSPSMappingMixin):
     def __init__(self, n_dim: int) -> None:
         super().__init__(n_params=n_dim, n_values=n_dim)
@@ -132,31 +123,3 @@ def __call__(self, w: Tensor) -> Tensor:  # noqa: ARG002
     @property
     def sps_mapping(self) -> "ElementWiseQuadratic.SPSMapping":
         return self.SPSMapping(self.n_values)
-
-
-class Multinorm(Objective, WithSPSMappingMixin):
-    def __init__(self, a: Tensor) -> None:
-        n = len(a)
-        super().__init__(n_params=n, n_values=n)
-        self.a = a
-
-    def __call__(self, x: Tensor) -> Tensor:
-        if len(x) != self.n_values:
-            raise ValueError("x must have the same length as the number of values.")
-
-        # f_i(x) = a_i * || x - a_i * e_i  ||²
-        return self.a * torch.norm(x.expand(len(x), len(x)) - torch.diag(self.a), dim=1) ** 2
-
-    def jacobian(self, x: Tensor) -> Tensor:
-        return self.a * 2 * (x.expand(len(x), len(x)) - torch.diag(self.a))
-
-    class SPSMapping(WithSPSMappingMixin.SPSMapping):
-        def __init__(self, a: Tensor) -> None:
-            self.a = a
-
-        def __call__(self, w: Tensor) -> Tensor:
-            return w * self.a
-
-    @property
-    def sps_mapping(self) -> "Multinorm.SPSMapping":
-        return self.SPSMapping(self.a)

tests/trajectories/optimize.py

@@ -6,7 +6,7 @@
     uv run python tests/trajectories/optimize.py <objective> <aggregator>...
 
 Arguments:
-    <objective>        The key of the objective function (e.g., EWQ, CQF, HQF, MN2, MN20).
+    <objective>        The key of the objective function (e.g., EWQ, CQF, HQF).
     <aggregator>...    The keys of the aggregators to use (e.g., upgrad, mean, mgda).
 """
 

tests/trajectories/plot_distance_to_pf.py

@@ -5,7 +5,7 @@
     uv run python tests/trajectories/plot_distance_to_pf.py <objective>
 
 Arguments:
-    <objective>    The key of the objective function (e.g., EWQ, CQF, HQF, MN2, MN20).
+    <objective>    The key of the objective function (e.g., EWQ, CQF, HQF).
 """
 
 import argparse

tests/trajectories/plot_params.py

@@ -5,7 +5,7 @@
     uv run python tests/trajectories/plot_params.py <objective>
 
 Arguments:
-    <objective>    The key of the objective function (e.g., EWQ, CQF, HQF, MN2, MN20).
+    <objective>    The key of the objective function (e.g., EWQ, CQF, HQF).
 """
 
 import argparse

tests/trajectories/plot_values.py

@@ -5,7 +5,7 @@
     uv run python tests/trajectories/plot_values.py <objective>
 
 Arguments:
-    <objective>    The key of the objective function (e.g., EWQ, CQF, HQF, MN2, MN20).
+    <objective>    The key of the objective function (e.g., EWQ, CQF, HQF).
 """
 
 import argparse