refactor(gradvac): literal group types, eps/beta rules, and plotter UX

- Use group_type "whole_model" | "all_layer" | "all_matrix" instead of 0/1/2
- Remove DEFAULT_GRADVAC_EPS from the public API; keep default 1e-8; allow eps=0
- Validate beta via setter; tighten GradVac repr/str expectations
- Fix all_layer leaf sizing via children() and parameters() instead of private fields
- Trim redundant GradVac.rst prose; align docs with the new API
- Tests: GradVac cases, value regression with torch.manual_seed for GradVac
- Plotter: factory dict + fresh aggregator instances per update; legend from
  selected keys; MathJax labels and live angle/length readouts in the sidebar

This commit includes GradVac implementation with Aggregator class.

Author: rkhosrowshahi

docs/source/docs/aggregation/gradvac.rst

@@ -3,19 +3,6 @@
 GradVac
 =======
 
-.. autodata:: torchjd.aggregation.DEFAULT_GRADVAC_EPS
-
-The constructor argument ``group_type`` (default ``0``) sets **parameter granularity** for the
-per-block cosine statistics in GradVac:
-
-* ``0`` — **whole model** (``whole_model``): one block per task gradient row. Omit ``encoder`` and
-  ``shared_params``.
-* ``1`` — **all layer** (``all_layer``): one block per leaf submodule with parameters under
-  ``encoder`` (same traversal as ``encoder.modules()`` in the reference formulation).
-* ``2`` — **all matrix** (``all_matrix``): one block per tensor in ``shared_params``, in order. Use
-  the same tensors as for the shared-parameter Jacobian columns (e.g. the parameters you would pass
-  to a shared-gradient helper).
-
 .. autoclass:: torchjd.aggregation.GradVac
     :members:
     :undoc-members:

src/torchjd/aggregation/__init__.py

@@ -66,7 +66,7 @@
 from ._dualproj import DualProj, DualProjWeighting
 from ._flattening import Flattening
 from ._graddrop import GradDrop
-from ._gradvac import DEFAULT_GRADVAC_EPS, GradVac
+from ._gradvac import GradVac
 from ._imtl_g import IMTLG, IMTLGWeighting
 from ._krum import Krum, KrumWeighting
 from ._mean import Mean, MeanWeighting
@@ -88,7 +88,6 @@
     "ConFIG",
     "Constant",
     "ConstantWeighting",
-    "DEFAULT_GRADVAC_EPS",
     "DualProj",
     "DualProjWeighting",
     "Flattening",

src/torchjd/aggregation/_gradvac.py

@@ -1,3 +1,5 @@
+from collections.abc import Callable
+
 import numpy as np
 import torch
 from plotly import graph_objects as go
@@ -7,14 +9,22 @@
 
 
 class Plotter:
-    def __init__(self, aggregators: list[Aggregator], matrix: torch.Tensor, seed: int = 0) -> None:
-        self.aggregators = aggregators
+    def __init__(
+        self,
+        aggregator_factories: dict[str, Callable[[], Aggregator]],
+        selected_keys: list[str],
+        matrix: torch.Tensor,
+        seed: int = 0,
+    ) -> None:
+        self._aggregator_factories = aggregator_factories
+        self.selected_keys = selected_keys
         self.matrix = matrix
         self.seed = seed
 
     def make_fig(self) -> Figure:
         torch.random.manual_seed(self.seed)
-        results = [agg(self.matrix) for agg in self.aggregators]
+        aggregators = [self._aggregator_factories[key]() for key in self.selected_keys]
+        results = [agg(self.matrix) for agg in aggregators]
 
         fig = go.Figure()
 
@@ -23,14 +33,19 @@ def make_fig(self) -> Figure:
         fig.add_trace(cone)
 
         for i in range(len(self.matrix)):
-            scatter = make_vector_scatter(self.matrix[i], "blue", f"g{i + 1}")
+            scatter = make_vector_scatter(
+                self.matrix[i],
+                "blue",
+                f"g{i + 1}",
+                textposition="top right",
+            )
             fig.add_trace(scatter)
 
         for i in range(len(results)):
             scatter = make_vector_scatter(
                 results[i],
                 "black",
-                str(self.aggregators[i]),
+                self.selected_keys[i],
                 showlegend=True,
                 dash=True,
             )

tests/plots/_utils.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import webbrowser
+from collections.abc import Callable
 from threading import Timer
 
 import numpy as np
@@ -12,11 +13,13 @@
 from torchjd.aggregation import (
     IMTLG,
     MGDA,
+    Aggregator,
     AlignedMTL,
     CAGrad,
     ConFIG,
     DualProj,
     GradDrop,
+    GradVac,
     Mean,
     NashMTL,
     PCGrad,
@@ -30,6 +33,14 @@
 MAX_LENGTH = 25.0
 
 
+def _format_angle_display(angle: float) -> str:
+    return f"{angle:.4f} rad ({np.degrees(angle):.1f}°)"
+
+
+def _format_length_display(r: float) -> str:
+    return f"{r:.4f}"
+
+
 def main() -> None:
     log = logging.getLogger("werkzeug")
     log.setLevel(logging.CRITICAL)
@@ -42,26 +53,30 @@ def main() -> None:
         ],
     )
 
-    aggregators = [
-        AlignedMTL(),
-        CAGrad(c=0.5),
-        ConFIG(),
-        DualProj(),
-        GradDrop(),
-        IMTLG(),
-        Mean(),
-        MGDA(),
-        NashMTL(n_tasks=matrix.shape[0]),
-        PCGrad(),
-        Random(),
-        Sum(),
-        TrimmedMean(trim_number=1),
-        UPGrad(),
-    ]
-
-    aggregators_dict = {str(aggregator): aggregator for aggregator in aggregators}
-
-    plotter = Plotter([], matrix)
+    n_tasks = matrix.shape[0]
+    aggregator_factories: dict[str, Callable[[], Aggregator]] = {
+        "AlignedMTL-min": lambda: AlignedMTL(scale_mode="min"),
+        "AlignedMTL-median": lambda: AlignedMTL(scale_mode="median"),
+        "AlignedMTL-RMSE": lambda: AlignedMTL(scale_mode="rmse"),
+        str(CAGrad(c=0.5)): lambda: CAGrad(c=0.5),
+        str(ConFIG()): lambda: ConFIG(),
+        str(DualProj()): lambda: DualProj(),
+        str(GradDrop()): lambda: GradDrop(),
+        str(GradVac()): lambda: GradVac(),
+        str(IMTLG()): lambda: IMTLG(),
+        str(Mean()): lambda: Mean(),
+        str(MGDA()): lambda: MGDA(),
+        str(NashMTL(n_tasks=n_tasks)): lambda: NashMTL(n_tasks=n_tasks),
+        str(PCGrad()): lambda: PCGrad(),
+        str(Random()): lambda: Random(),
+        str(Sum()): lambda: Sum(),
+        str(TrimmedMean(trim_number=1)): lambda: TrimmedMean(trim_number=1),
+        str(UPGrad()): lambda: UPGrad(),
+    }
+
+    aggregator_strings = list(aggregator_factories.keys())
+
+    plotter = Plotter(aggregator_factories, [], matrix)
 
     app = Dash(__name__)
 
@@ -96,7 +111,6 @@ def main() -> None:
         gradient_slider_inputs.append(Input(angle_input, "value"))
         gradient_slider_inputs.append(Input(r_input, "value"))
 
-    aggregator_strings = [str(aggregator) for aggregator in aggregators]
     checklist = dcc.Checklist(aggregator_strings, [], id="aggregator-checklist")
 
     control_div = html.Div(
@@ -115,32 +129,40 @@ def update_seed(value: int) -> Figure:
         plotter.seed = value
         return plotter.make_fig()
 
+    n_gradients = len(matrix)
+    gradient_value_outputs: list[Output] = []
+    for i in range(n_gradients):
+        gradient_value_outputs.append(Output(f"g{i + 1}-angle-value", "children"))
+        gradient_value_outputs.append(Output(f"g{i + 1}-length-value", "children"))
+
     @callback(
         Output("aggregations-fig", "figure", allow_duplicate=True),
+        *gradient_value_outputs,
         *gradient_slider_inputs,
         prevent_initial_call=True,
     )
-    def update_gradient_coordinate(*values: str) -> Figure:
+    def update_gradient_coordinate(*values: str) -> tuple[Figure, ...]:
         values_ = [float(value) for value in values]
 
+        display_parts: list[str] = []
         for j in range(len(values_) // 2):
             angle = values_[2 * j]
             r = values_[2 * j + 1]
             x, y = angle_to_coord(angle, r)
             plotter.matrix[j, 0] = x
             plotter.matrix[j, 1] = y
+            display_parts.append(_format_angle_display(angle))
+            display_parts.append(_format_length_display(r))
 
-        return plotter.make_fig()
+        return (plotter.make_fig(), *display_parts)
 
     @callback(
         Output("aggregations-fig", "figure", allow_duplicate=True),
         Input("aggregator-checklist", "value"),
         prevent_initial_call=True,
     )
     def update_aggregators(value: list[str]) -> Figure:
-        aggregator_keys = value
-        new_aggregators = [aggregators_dict[key] for key in aggregator_keys]
-        plotter.aggregators = new_aggregators
+        plotter.selected_keys = list(value)
         return plotter.make_fig()
 
     Timer(1, open_browser).start()
@@ -173,11 +195,56 @@ def make_gradient_div(
         style={"width": "250px"},
     )
 
+    label_style: dict[str, str | int] = {
+        "display": "inline-block",
+        "width": "52px",
+        "margin-right": "8px",
+        "vertical-align": "middle",
+    }
+    value_style: dict[str, str] = {
+        "display": "inline-block",
+        "margin-left": "10px",
+        "min-width": "140px",
+        "font-family": "monospace",
+        "font-size": "13px",
+        "vertical-align": "middle",
+    }
+    row_style: dict[str, str] = {"display": "block", "margin-bottom": "6px"}
     div = html.Div(
         [
-            html.P(f"g{i + 1}", style={"display": "inline-block", "margin-right": 20}),
-            angle_input,
-            r_input,
+            dcc.Markdown(
+                f"$g_{{{i + 1}}}$",
+                mathjax=True,
+                style={
+                    "margin": "0 0 6px 0",
+                    "font-weight": "bold",
+                    "display": "block",
+                },
+            ),
+            html.Div(
+                [
+                    html.Span("Angle", style=label_style),
+                    angle_input,
+                    html.Span(
+                        id=f"g{i + 1}-angle-value",
+                        children=_format_angle_display(angle),
+                        style=value_style,
+                    ),
+                ],
+                style=row_style,
+            ),
+            html.Div(
+                [
+                    html.Span("Length", style=label_style),
+                    r_input,
+                    html.Span(
+                        id=f"g{i + 1}-length-value",
+                        children=_format_length_display(r),
+                        style=value_style,
+                    ),
+                ],
+                style={**row_style, "margin-bottom": "12px"},
+            ),
         ],
     )
     return div, angle_input, r_input