From f95e0a2c8b896574c43da42d0a46171d7e0b55c1 Mon Sep 17 00:00:00 2001 From: Khush Date: Sat, 1 Aug 2026 12:31:57 -0400 Subject: [PATCH 1/6] feat: Add hosted interactive plotter webapp --- .github/workflows/deploy-visualization.yml | 36 +++ visualization/Dockerfile | 17 ++ visualization/README.md | 34 +++ visualization/_utils.py | 209 ++++++++++++++ visualization/app.py | 304 +++++++++++++++++++++ visualization/requirements.txt | 4 + 6 files changed, 604 insertions(+) create mode 100644 .github/workflows/deploy-visualization.yml create mode 100644 visualization/Dockerfile create mode 100644 visualization/README.md create mode 100644 visualization/_utils.py create mode 100644 visualization/app.py create mode 100644 visualization/requirements.txt diff --git a/.github/workflows/deploy-visualization.yml b/.github/workflows/deploy-visualization.yml new file mode 100644 index 000000000..c45649045 --- /dev/null +++ b/.github/workflows/deploy-visualization.yml @@ -0,0 +1,36 @@ +name: Deploy Visualization to Hugging Face Spaces + +on: + push: + branches: [main] + paths: + - visualization/** + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + deploy: + name: Deploy to Hugging Face Spaces + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v7 + + - name: Deploy visualization folder to HF Spaces + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + pip install --quiet huggingface-hub + python -c " + import os + from huggingface_hub import HfApi + api = HfApi(token=os.environ['HF_TOKEN']) + api.upload_folder( + folder_path='visualization', + repo_id='torchjd/interactive-plotter', + repo_type='space', + ) + print('Deployed successfully.') + " diff --git a/visualization/Dockerfile b/visualization/Dockerfile new file mode 100644 index 000000000..15fd978ad --- /dev/null +++ b/visualization/Dockerfile @@ -0,0 +1,17 @@ +FROM python:3.11-slim + +WORKDIR /app + +RUN apt-get update && apt-get install -y gcc g++ && rm -rf /var/lib/apt/lists/* + +# CPU-only torch keeps the image small +RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY . . + +EXPOSE 7860 + +CMD ["python", "app.py"] diff --git a/visualization/README.md b/visualization/README.md new file mode 100644 index 000000000..922ae6ff9 --- /dev/null +++ b/visualization/README.md @@ -0,0 +1,34 @@ +--- +title: TorchJD Interactive Plotter +emoji: 📊 +colorFrom: blue +colorTo: green +sdk: docker +app_port: 7860 +pinned: false +license: mit +--- + +# TorchJD Interactive Plotter + +Interactive visualization of gradient aggregation methods from [TorchJD](https://torchjd.org). + +Adjust the angle and length of each gradient vector and select aggregators to see how they combine +the gradients. The green region shows the dual cone — any descent direction must lie inside it. + +## URL parameters + +The app accepts query parameters so you can link to a specific configuration or embed it in +documentation with an aggregator pre-selected: + +| Parameter | Format | Example | +|-----------|--------|---------| +| `agg` | Comma-separated aggregator names | `?agg=Mean,MGDA` | +| `g1`, `g2`, `g3` | `angle_radians,length` | `?g1=1.5708,2.0` | +| `seed` | Integer | `?seed=42` | + +Example — open with MGDA pre-selected: + +``` +https://torchjd-interactive-plotter.hf.space/?agg=MGDA +``` diff --git a/visualization/_utils.py b/visualization/_utils.py new file mode 100644 index 000000000..5edca32b3 --- /dev/null +++ b/visualization/_utils.py @@ -0,0 +1,209 @@ +from collections.abc import Callable + +import numpy as np +import torch +from plotly import graph_objects as go +from plotly.graph_objs import Figure, Scatter + +from torchjd.aggregation import Aggregator + + +class Plotter: + 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) + aggregators = [self._aggregator_factories[key]() for key in self.selected_keys] + results = [agg(self.matrix) for agg in aggregators] + + fig = go.Figure() + + start_angle, opening = compute_2d_dual_cone(self.matrix.numpy()) + cone = make_cone_scatter(start_angle, opening, label="Dual cone") + fig.add_trace(cone) + + for i in range(len(self.matrix)): + 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", + self.selected_keys[i], + showlegend=True, + dash=True, + ) + fig.add_trace(scatter) + + all_x = [0] + self.matrix.numpy()[:, 0].tolist() + [res[0].item() for res in results] + all_y = [0] + self.matrix.numpy()[:, 1].tolist() + [res[1].item() for res in results] + min_x = min(all_x) + max_x = max(all_x) + min_y = min(all_y) + max_y = max(all_y) + len_x = max_x - min_x + len_y = max_y - min_y + margin_prop = 0.05 + range_x = [min_x - len_x * margin_prop, max_x + len_x * margin_prop] + range_y = [min_y - len_y * margin_prop, max_y + len_y * margin_prop] + + fig.update_layout(hovermode=False, width=1000, height=900) + fig.update_xaxes(scaleanchor="y", scaleratio=1, range=range_x) + fig.update_yaxes(range=range_y) + + return fig + + +def make_vector_scatter( + gradient: torch.Tensor, + color: str, + label: str, + showlegend: bool = False, + dash: bool = False, + textposition: str = "bottom center", + line_width: float = 2.5, + text_size: float = 12, + marker_size: float = 12, +) -> Scatter: + line = {"color": color, "width": line_width} + if dash: + line["dash"] = "dash" + + scatter = go.Scatter( + x=[0, gradient[0]], + y=[0, gradient[1]], + mode="lines+markers+text", + line=line, + marker={ + "symbol": "arrow", + "color": color, + "size": marker_size, + "angleref": "previous", + }, + name=label, + text=["", label], + textposition=textposition, + textfont={"color": color, "size": text_size}, + showlegend=showlegend, + ) + return scatter + + +def make_cone_scatter( + start_angle: float, + opening: float, + label: str, + scale: float = 100.0, + printable: bool = False, +) -> Scatter: + if opening < -1e-8: + cone_outline = np.zeros([0, 2]) + else: + middle_angle = start_angle + opening / 2 + end_angle = start_angle + opening + + start_vec = angle_to_coord(start_angle, scale) + end_vec = angle_to_coord(end_angle, scale) + + if np.abs(opening - 2 * np.pi) <= 0.01: + cone_outline = np.array( + [ + [0, 0], + start_vec, + end_vec, + [0, 0], + ], + ) + else: + middle_point = angle_to_coord(middle_angle, scale) + + cone_outline = np.array( + [ + [0, 0], + start_vec, + middle_point, + end_vec, + [0, 0], + ], + ) + + if printable: + fillpattern = { + "bgcolor": "white", + "shape": "\\", + "fgcolor": "rgba(0, 220, 0, 0.5)", + "size": 30, + "solidity": 0.15, + } + else: + fillpattern = None + + cone = go.Scatter( + x=cone_outline[:, 0], + y=cone_outline[:, 1], + fill="toself", + mode="lines", + fillcolor="rgba(0, 255, 0, 0.07)", + line={"color": "rgb(0, 220, 0)", "width": 2}, + name=label, + fillpattern=fillpattern, + ) + + return cone + + +def compute_2d_dual_cone(matrix: np.ndarray) -> tuple[float, float]: + row_angles = [coord_to_angle(*row)[0] for row in matrix] + start_angles = [(angle - np.pi / 2) % (2 * np.pi) for angle in row_angles] + + cone_start_angle = start_angles[0] + opening = np.pi + for hs_start_angle in start_angles[1:]: + cone_start_angle, opening = combine_bounds(cone_start_angle, opening, hs_start_angle) + + return cone_start_angle, opening + + +def combine_bounds( + cone_start_angle: float, + opening: float, + hs_start_angle: float, +) -> tuple[float, float]: + angle_between_starts = (hs_start_angle - cone_start_angle) % (2 * np.pi) + if angle_between_starts < np.pi: + cone_start_angle = hs_start_angle + opening = opening - angle_between_starts + else: + opening = min(opening, (hs_start_angle + np.pi - cone_start_angle) % (2 * np.pi)) + + return cone_start_angle, opening + + +def coord_to_angle(x: float, y: float) -> tuple[float, float]: + r = np.sqrt(x**2 + y**2) + if r == 0: + raise ValueError("No angle") + angle = np.arccos(x / r) if y >= 0 else 2 * np.pi - np.arccos(x / r) + return angle, r + + +def angle_to_coord(angle: float, r: float = 1.0) -> tuple[float, float]: + x = r * np.cos(angle) + y = r * np.sin(angle) + return x, y diff --git a/visualization/app.py b/visualization/app.py new file mode 100644 index 000000000..7cdc08a0a --- /dev/null +++ b/visualization/app.py @@ -0,0 +1,304 @@ +import logging +from urllib.parse import parse_qs, urlencode + +import numpy as np +import torch +from _utils import Plotter, angle_to_coord, coord_to_angle +from dash import Dash, Input, Output, State, dcc, html, no_update + +from torchjd.aggregation import ( + IMTLG, + MGDA, + AlignedMTL, + CAGrad, + ConFIG, + DualProj, + FairGrad, + GradDrop, + GradVac, + Mean, + NashMTL, + PCGrad, + Random, + Sum, + TrimmedMean, + UPGrad, +) +from torchjd.linalg import QuadprogProjector + +logging.getLogger("werkzeug").setLevel(logging.CRITICAL) + +MIN_LENGTH = 0.01 +MAX_LENGTH = 25.0 +N_TASKS = 3 + +DEFAULT_MATRIX = torch.tensor( + [ + [0.0, 1.0], + [1.0, -1.0], + [1.0, 0.0], + ] +) + +AGGREGATOR_FACTORIES = { + "AlignedMTL-min": lambda: AlignedMTL(scale_mode="min"), + "AlignedMTL-median": lambda: AlignedMTL(scale_mode="median"), + "AlignedMTL-RMSE": lambda: AlignedMTL(scale_mode="rmse"), + "CAGrad": lambda: CAGrad(c=0.5), + "ConFIG": lambda: ConFIG(), + "DualProj": lambda: DualProj(projector=QuadprogProjector(reg_eps=1e-7)), + "FairGrad": lambda: FairGrad(alpha=1.0), + "GradDrop": lambda: GradDrop(), + "GradVac": lambda: GradVac(), + "IMTLG": lambda: IMTLG(), + "Mean": lambda: Mean(), + "MGDA": lambda: MGDA(), + "NashMTL": lambda: NashMTL(n_tasks=N_TASKS), + "PCGrad": lambda: PCGrad(), + "Random": lambda: Random(), + "Sum": lambda: Sum(), + "TrimmedMean": lambda: TrimmedMean(trim_number=1), + "UPGrad": lambda: UPGrad(projector=QuadprogProjector(reg_eps=1e-7)), +} + +ALL_KEYS = list(AGGREGATOR_FACTORIES.keys()) + +# Default slider values derived from DEFAULT_MATRIX +_DEFAULT_ANGLES_RS: list[float] = [] +for _i in range(N_TASKS): + _x, _y = DEFAULT_MATRIX[_i, 0].item(), DEFAULT_MATRIX[_i, 1].item() + _a, _r = coord_to_angle(_x, _y) + _DEFAULT_ANGLES_RS.extend([_a, _r]) + + +def _format_angle(angle: float) -> str: + return f"{np.degrees(angle):.1f}°" + + +def _format_length(r: float) -> str: + return f"{r:.2f}" + + +def _make_gradient_div(i: int, angle: float, r: float) -> html.Div: + label_style = { + "display": "inline-block", + "width": "52px", + "margin-right": "8px", + "vertical-align": "middle", + } + value_style = { + "display": "inline-block", + "margin-left": "10px", + "min-width": "140px", + "font-family": "monospace", + "font-size": "13px", + "vertical-align": "middle", + } + row_style = {"display": "block", "margin-bottom": "6px"} + + return html.Div( + [ + 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), + dcc.Input( + id=f"g{i + 1}-angle", + type="range", + value=angle, + min=0, + max=2 * np.pi, + style={"width": "250px"}, + ), + html.Span( + id=f"g{i + 1}-angle-display", + children=_format_angle(angle), + style=value_style, + ), + ], + style=row_style, + ), + html.Div( + [ + html.Span("Length", style=label_style), + dcc.Input( + id=f"g{i + 1}-r", + type="range", + value=r, + min=MIN_LENGTH, + max=MAX_LENGTH, + style={"width": "250px"}, + ), + html.Span( + id=f"g{i + 1}-r-display", + children=_format_length(r), + style=value_style, + ), + ], + style={**row_style, "margin-bottom": "12px"}, + ), + ] + ) + + +_default_fig = Plotter(AGGREGATOR_FACTORIES, [], DEFAULT_MATRIX.clone(), 0).make_fig() + +_gradient_divs = [ + _make_gradient_div(i, _DEFAULT_ANGLES_RS[2 * i], _DEFAULT_ANGLES_RS[2 * i + 1]) + for i in range(N_TASKS) +] + +app = Dash(__name__) + +app.layout = html.Div( + [ + dcc.Location(id="url", refresh=False), + # Tracks whether URL params have been applied on first load + dcc.Store(id="initialized", data=False), + html.Div( + [dcc.Graph(id="figure", figure=_default_fig)], + style={"display": "inline-block"}, + ), + html.Div( + [ + html.Div( + [ + html.P("Seed", style={"display": "inline-block", "margin-right": 20}), + dcc.Input( + id="seed", + type="number", + value=0, + style={ + "display": "inline-block", + "border": "1px solid black", + "width": "25%", + }, + ), + ], + style={"display": "inline-block", "width": "100%"}, + ), + *_gradient_divs, + dcc.Checklist(ALL_KEYS, [], id="agg-checklist"), + ], + style={"display": "inline-block", "vertical-align": "top"}, + ), + ] +) + +# Inputs/outputs shared by both callbacks +_gradient_angle_r_inputs = [] +for _i in range(N_TASKS): + _gradient_angle_r_inputs.extend( + [ + Input(f"g{_i + 1}-angle", "value"), + Input(f"g{_i + 1}-r", "value"), + ] + ) + +_display_outputs = [] +for _i in range(N_TASKS): + _display_outputs.extend( + [ + Output(f"g{_i + 1}-angle-display", "children"), + Output(f"g{_i + 1}-r-display", "children"), + ] + ) + +_gradient_angle_r_outputs = [] +for _i in range(N_TASKS): + _gradient_angle_r_outputs.extend( + [ + Output(f"g{_i + 1}-angle", "value"), + Output(f"g{_i + 1}-r", "value"), + ] + ) + + +@app.callback( + *_gradient_angle_r_outputs, + Output("agg-checklist", "value"), + Output("seed", "value"), + Output("initialized", "data"), + Input("url", "search"), + State("initialized", "data"), + prevent_initial_call=False, +) +def init_from_url(search: str | None, initialized: bool) -> tuple: + """Reads URL query params once on page load and sets initial slider values.""" + n_outputs = N_TASKS * 2 + 3 # angles+rs + checklist + seed + initialized flag + if initialized: + return (*[no_update] * (n_outputs - 1), no_update) + + if not search: + return (*_DEFAULT_ANGLES_RS, [], 0, True) + + params = parse_qs(search.lstrip("?")) + + agg_param = params.get("agg", [""])[0] + selected = [a for a in agg_param.split(",") if a in AGGREGATOR_FACTORIES] if agg_param else [] + + seed = max(0, int(params.get("seed", ["0"])[0])) + + angles_rs: list[float] = [] + for i in range(N_TASKS): + default_angle = _DEFAULT_ANGLES_RS[2 * i] + default_r = _DEFAULT_ANGLES_RS[2 * i + 1] + g_param = params.get(f"g{i + 1}", [""])[0] + if g_param: + parts = g_param.split(",") + if len(parts) == 2: + try: + default_angle = float(parts[0]) + default_r = max(MIN_LENGTH, min(MAX_LENGTH, float(parts[1]))) + except ValueError: + pass + angles_rs.extend([default_angle, default_r]) + + return (*angles_rs, selected, seed, True) + + +@app.callback( + Output("figure", "figure"), + Output("url", "search"), + *_display_outputs, + *_gradient_angle_r_inputs, + Input("agg-checklist", "value"), + Input("seed", "value"), + prevent_initial_call=True, +) +def update_all(*args: object) -> tuple: + """Updates the figure and URL whenever any control changes.""" + gradient_values = args[: N_TASKS * 2] + selected = list(args[-2] or []) + seed = max(0, int(args[-1] or 0)) + + matrix = DEFAULT_MATRIX.clone() + search_params: dict[str, str] = {} + display_values: list[str] = [] + + for i in range(N_TASKS): + angle = float(gradient_values[2 * i]) + r = float(gradient_values[2 * i + 1]) + x, y = angle_to_coord(angle, r) + matrix[i, 0] = x + matrix[i, 1] = y + search_params[f"g{i + 1}"] = f"{angle},{r}" + display_values.extend([_format_angle(angle), _format_length(r)]) + + if selected: + search_params["agg"] = ",".join(selected) + search_params["seed"] = str(seed) + + plotter = Plotter(AGGREGATOR_FACTORIES, selected, matrix, seed) + fig = plotter.make_fig() + search = "?" + urlencode(search_params) + + return (fig, search, *display_values) + + +if __name__ == "__main__": + app.run(host="0.0.0.0", port=7860, debug=False) diff --git a/visualization/requirements.txt b/visualization/requirements.txt new file mode 100644 index 000000000..35a80faee --- /dev/null +++ b/visualization/requirements.txt @@ -0,0 +1,4 @@ +torchjd[full]>=0.16.0 +dash>=2.16.0 +plotly>=5.19.0 +numpy>=1.21.2 From 07a8f9e81ed5583ae1f560ba2fd2e99071a3f199 Mon Sep 17 00:00:00 2001 From: Khush Date: Sun, 2 Aug 2026 12:16:42 -0400 Subject: [PATCH 2/6] Converting Code to Gradio for Plotter Webapp --- visualization/Dockerfile | 17 --- visualization/README.md | 11 +- visualization/app.py | 267 +++++++-------------------------- visualization/requirements.txt | 2 +- 4 files changed, 62 insertions(+), 235 deletions(-) delete mode 100644 visualization/Dockerfile diff --git a/visualization/Dockerfile b/visualization/Dockerfile deleted file mode 100644 index 15fd978ad..000000000 --- a/visualization/Dockerfile +++ /dev/null @@ -1,17 +0,0 @@ -FROM python:3.11-slim - -WORKDIR /app - -RUN apt-get update && apt-get install -y gcc g++ && rm -rf /var/lib/apt/lists/* - -# CPU-only torch keeps the image small -RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu - -COPY requirements.txt . -RUN pip install --no-cache-dir -r requirements.txt - -COPY . . - -EXPOSE 7860 - -CMD ["python", "app.py"] diff --git a/visualization/README.md b/visualization/README.md index 922ae6ff9..8f02a244c 100644 --- a/visualization/README.md +++ b/visualization/README.md @@ -3,8 +3,9 @@ title: TorchJD Interactive Plotter emoji: 📊 colorFrom: blue colorTo: green -sdk: docker -app_port: 7860 +sdk: gradio +sdk_version: 4.0.0 +app_file: app.py pinned: false license: mit --- @@ -26,9 +27,3 @@ documentation with an aggregator pre-selected: | `agg` | Comma-separated aggregator names | `?agg=Mean,MGDA` | | `g1`, `g2`, `g3` | `angle_radians,length` | `?g1=1.5708,2.0` | | `seed` | Integer | `?seed=42` | - -Example — open with MGDA pre-selected: - -``` -https://torchjd-interactive-plotter.hf.space/?agg=MGDA -``` diff --git a/visualization/app.py b/visualization/app.py index 7cdc08a0a..10f37e07b 100644 --- a/visualization/app.py +++ b/visualization/app.py @@ -1,10 +1,9 @@ import logging -from urllib.parse import parse_qs, urlencode +import gradio as gr import numpy as np import torch from _utils import Plotter, angle_to_coord, coord_to_angle -from dash import Dash, Input, Output, State, dcc, html, no_update from torchjd.aggregation import ( IMTLG, @@ -63,242 +62,92 @@ ALL_KEYS = list(AGGREGATOR_FACTORIES.keys()) -# Default slider values derived from DEFAULT_MATRIX _DEFAULT_ANGLES_RS: list[float] = [] for _i in range(N_TASKS): _x, _y = DEFAULT_MATRIX[_i, 0].item(), DEFAULT_MATRIX[_i, 1].item() _a, _r = coord_to_angle(_x, _y) - _DEFAULT_ANGLES_RS.extend([_a, _r]) + _DEFAULT_ANGLES_RS.extend([float(_a), float(_r)]) -def _format_angle(angle: float) -> str: - return f"{np.degrees(angle):.1f}°" - - -def _format_length(r: float) -> str: - return f"{r:.2f}" - - -def _make_gradient_div(i: int, angle: float, r: float) -> html.Div: - label_style = { - "display": "inline-block", - "width": "52px", - "margin-right": "8px", - "vertical-align": "middle", - } - value_style = { - "display": "inline-block", - "margin-left": "10px", - "min-width": "140px", - "font-family": "monospace", - "font-size": "13px", - "vertical-align": "middle", - } - row_style = {"display": "block", "margin-bottom": "6px"} - - return html.Div( - [ - 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), - dcc.Input( - id=f"g{i + 1}-angle", - type="range", - value=angle, - min=0, - max=2 * np.pi, - style={"width": "250px"}, - ), - html.Span( - id=f"g{i + 1}-angle-display", - children=_format_angle(angle), - style=value_style, - ), - ], - style=row_style, - ), - html.Div( - [ - html.Span("Length", style=label_style), - dcc.Input( - id=f"g{i + 1}-r", - type="range", - value=r, - min=MIN_LENGTH, - max=MAX_LENGTH, - style={"width": "250px"}, - ), - html.Span( - id=f"g{i + 1}-r-display", - children=_format_length(r), - style=value_style, - ), - ], - style={**row_style, "margin-bottom": "12px"}, - ), - ] - ) - - -_default_fig = Plotter(AGGREGATOR_FACTORIES, [], DEFAULT_MATRIX.clone(), 0).make_fig() - -_gradient_divs = [ - _make_gradient_div(i, _DEFAULT_ANGLES_RS[2 * i], _DEFAULT_ANGLES_RS[2 * i + 1]) - for i in range(N_TASKS) -] - -app = Dash(__name__) - -app.layout = html.Div( - [ - dcc.Location(id="url", refresh=False), - # Tracks whether URL params have been applied on first load - dcc.Store(id="initialized", data=False), - html.Div( - [dcc.Graph(id="figure", figure=_default_fig)], - style={"display": "inline-block"}, - ), - html.Div( - [ - html.Div( - [ - html.P("Seed", style={"display": "inline-block", "margin-right": 20}), - dcc.Input( - id="seed", - type="number", - value=0, - style={ - "display": "inline-block", - "border": "1px solid black", - "width": "25%", - }, - ), - ], - style={"display": "inline-block", "width": "100%"}, - ), - *_gradient_divs, - dcc.Checklist(ALL_KEYS, [], id="agg-checklist"), - ], - style={"display": "inline-block", "vertical-align": "top"}, - ), - ] -) - -# Inputs/outputs shared by both callbacks -_gradient_angle_r_inputs = [] -for _i in range(N_TASKS): - _gradient_angle_r_inputs.extend( - [ - Input(f"g{_i + 1}-angle", "value"), - Input(f"g{_i + 1}-r", "value"), - ] - ) - -_display_outputs = [] -for _i in range(N_TASKS): - _display_outputs.extend( - [ - Output(f"g{_i + 1}-angle-display", "children"), - Output(f"g{_i + 1}-r-display", "children"), - ] - ) - -_gradient_angle_r_outputs = [] -for _i in range(N_TASKS): - _gradient_angle_r_outputs.extend( - [ - Output(f"g{_i + 1}-angle", "value"), - Output(f"g{_i + 1}-r", "value"), - ] - ) +def _build_matrix(angles_rs: list[float]) -> torch.Tensor: + matrix = DEFAULT_MATRIX.clone() + for i in range(N_TASKS): + x, y = angle_to_coord(angles_rs[2 * i], angles_rs[2 * i + 1]) + matrix[i, 0] = x + matrix[i, 1] = y + return matrix -@app.callback( - *_gradient_angle_r_outputs, - Output("agg-checklist", "value"), - Output("seed", "value"), - Output("initialized", "data"), - Input("url", "search"), - State("initialized", "data"), - prevent_initial_call=False, -) -def init_from_url(search: str | None, initialized: bool) -> tuple: - """Reads URL query params once on page load and sets initial slider values.""" - n_outputs = N_TASKS * 2 + 3 # angles+rs + checklist + seed + initialized flag - if initialized: - return (*[no_update] * (n_outputs - 1), no_update) +def update_plot(seed: float, *args: float | list[str]) -> gr.Plot: + gradient_values = args[: N_TASKS * 2] + selected = list(args[-1] or []) + angles_rs = [float(v) if v is not None else 0.0 for v in gradient_values] + matrix = _build_matrix(angles_rs) + plotter = Plotter(AGGREGATOR_FACTORIES, selected, matrix, int(seed or 0)) + return plotter.make_fig() - if not search: - return (*_DEFAULT_ANGLES_RS, [], 0, True) - params = parse_qs(search.lstrip("?")) +def load_from_url(request: gr.Request) -> list: + params = dict(request.query_params) - agg_param = params.get("agg", [""])[0] + agg_param = params.get("agg", "") selected = [a for a in agg_param.split(",") if a in AGGREGATOR_FACTORIES] if agg_param else [] - seed = max(0, int(params.get("seed", ["0"])[0])) + seed = max(0, int(params.get("seed", 0) or 0)) - angles_rs: list[float] = [] + angles_rs = list(_DEFAULT_ANGLES_RS) for i in range(N_TASKS): - default_angle = _DEFAULT_ANGLES_RS[2 * i] - default_r = _DEFAULT_ANGLES_RS[2 * i + 1] - g_param = params.get(f"g{i + 1}", [""])[0] + g_param = params.get(f"g{i + 1}", "") if g_param: parts = g_param.split(",") if len(parts) == 2: try: - default_angle = float(parts[0]) - default_r = max(MIN_LENGTH, min(MAX_LENGTH, float(parts[1]))) + angles_rs[2 * i] = float(parts[0]) + angles_rs[2 * i + 1] = max(MIN_LENGTH, min(MAX_LENGTH, float(parts[1]))) except ValueError: pass - angles_rs.extend([default_angle, default_r]) - return (*angles_rs, selected, seed, True) + matrix = _build_matrix(angles_rs) + plotter = Plotter(AGGREGATOR_FACTORIES, selected, matrix, seed) + fig = plotter.make_fig() + return [fig, float(seed), *[float(v) for v in angles_rs], selected] -@app.callback( - Output("figure", "figure"), - Output("url", "search"), - *_display_outputs, - *_gradient_angle_r_inputs, - Input("agg-checklist", "value"), - Input("seed", "value"), - prevent_initial_call=True, -) -def update_all(*args: object) -> tuple: - """Updates the figure and URL whenever any control changes.""" - gradient_values = args[: N_TASKS * 2] - selected = list(args[-2] or []) - seed = max(0, int(args[-1] or 0)) - matrix = DEFAULT_MATRIX.clone() - search_params: dict[str, str] = {} - display_values: list[str] = [] +with gr.Blocks(title="TorchJD Interactive Plotter") as demo: + with gr.Row(): + with gr.Column(scale=3): + plot = gr.Plot() + with gr.Column(scale=1): + seed_input = gr.Number(value=0, label="Seed", precision=0) - for i in range(N_TASKS): - angle = float(gradient_values[2 * i]) - r = float(gradient_values[2 * i + 1]) - x, y = angle_to_coord(angle, r) - matrix[i, 0] = x - matrix[i, 1] = y - search_params[f"g{i + 1}"] = f"{angle},{r}" - display_values.extend([_format_angle(angle), _format_length(r)]) + gradient_sliders: list[gr.Slider] = [] + for i in range(N_TASKS): + gr.Markdown(f"**$g_{{{i + 1}}}$**") + angle_slider = gr.Slider( + minimum=0, + maximum=2 * np.pi, + value=_DEFAULT_ANGLES_RS[2 * i], + step=0.01, + label=f"g{i + 1} angle (rad)", + ) + r_slider = gr.Slider( + minimum=MIN_LENGTH, + maximum=MAX_LENGTH, + value=_DEFAULT_ANGLES_RS[2 * i + 1], + step=0.01, + label=f"g{i + 1} length", + ) + gradient_sliders.extend([angle_slider, r_slider]) - if selected: - search_params["agg"] = ",".join(selected) - search_params["seed"] = str(seed) + agg_check = gr.CheckboxGroup(ALL_KEYS, label="Aggregators", value=[]) - plotter = Plotter(AGGREGATOR_FACTORIES, selected, matrix, seed) - fig = plotter.make_fig() - search = "?" + urlencode(search_params) + all_inputs = [seed_input, *gradient_sliders, agg_check] - return (fig, search, *display_values) + for component in all_inputs: + component.change(update_plot, inputs=all_inputs, outputs=plot) + demo.load(load_from_url, inputs=None, outputs=[plot, seed_input, *gradient_sliders, agg_check]) if __name__ == "__main__": - app.run(host="0.0.0.0", port=7860, debug=False) + demo.launch(server_name="0.0.0.0", server_port=7860) diff --git a/visualization/requirements.txt b/visualization/requirements.txt index 35a80faee..64a60c8da 100644 --- a/visualization/requirements.txt +++ b/visualization/requirements.txt @@ -1,4 +1,4 @@ torchjd[full]>=0.16.0 -dash>=2.16.0 +gradio>=4.0 plotly>=5.19.0 numpy>=1.21.2 From 118fddc1950fec15161e84f0354f7671b6dbca86 Mon Sep 17 00:00:00 2001 From: Khush Date: Sun, 2 Aug 2026 13:11:08 -0400 Subject: [PATCH 3/6] fix: update Gradio to 6.x and add audioop-lts for Python 3.13 compatibility Co-Authored-By: Claude Sonnet 4.6 --- visualization/README.md | 2 +- visualization/requirements.txt | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/visualization/README.md b/visualization/README.md index 8f02a244c..5eb243e27 100644 --- a/visualization/README.md +++ b/visualization/README.md @@ -4,7 +4,7 @@ emoji: 📊 colorFrom: blue colorTo: green sdk: gradio -sdk_version: 4.0.0 +sdk_version: 6.21.0 app_file: app.py pinned: false license: mit diff --git a/visualization/requirements.txt b/visualization/requirements.txt index 64a60c8da..57be7bf16 100644 --- a/visualization/requirements.txt +++ b/visualization/requirements.txt @@ -1,4 +1,5 @@ torchjd[full]>=0.16.0 -gradio>=4.0 +gradio>=6.0 +audioop-lts; python_version >= "3.13" plotly>=5.19.0 numpy>=1.21.2 From ec587defd3022ce5264ddd16f1cc6ee183c37dbb Mon Sep 17 00:00:00 2001 From: Khush Date: Sun, 2 Aug 2026 18:14:23 -0400 Subject: [PATCH 4/6] refactor: move webapp files to tests/plots to avoid code duplication Move app.py, README.md, and requirements.txt from visualization/ into tests/plots/ so they share _utils.py with the existing plotters. Update the deploy workflow to point at tests/plots/. Co-Authored-By: Claude Sonnet 4.6 --- .github/workflows/deploy-visualization.yml | 4 +- {visualization => tests/plots}/README.md | 2 +- {visualization => tests/plots}/app.py | 0 .../plots}/requirements.txt | 0 visualization/_utils.py | 209 ------------------ 5 files changed, 3 insertions(+), 212 deletions(-) rename {visualization => tests/plots}/README.md (86%) rename {visualization => tests/plots}/app.py (100%) rename {visualization => tests/plots}/requirements.txt (100%) delete mode 100644 visualization/_utils.py diff --git a/.github/workflows/deploy-visualization.yml b/.github/workflows/deploy-visualization.yml index c45649045..cafbd4fbe 100644 --- a/.github/workflows/deploy-visualization.yml +++ b/.github/workflows/deploy-visualization.yml @@ -4,7 +4,7 @@ on: push: branches: [main] paths: - - visualization/** + - tests/plots/** concurrency: group: ${{ github.workflow }}-${{ github.ref }} @@ -28,7 +28,7 @@ jobs: from huggingface_hub import HfApi api = HfApi(token=os.environ['HF_TOKEN']) api.upload_folder( - folder_path='visualization', + folder_path='tests/plots', repo_id='torchjd/interactive-plotter', repo_type='space', ) diff --git a/visualization/README.md b/tests/plots/README.md similarity index 86% rename from visualization/README.md rename to tests/plots/README.md index 5eb243e27..d79840e54 100644 --- a/visualization/README.md +++ b/tests/plots/README.md @@ -15,7 +15,7 @@ license: mit Interactive visualization of gradient aggregation methods from [TorchJD](https://torchjd.org). Adjust the angle and length of each gradient vector and select aggregators to see how they combine -the gradients. The green region shows the dual cone — any descent direction must lie inside it. +the gradients. The green region shows the dual cone: the set of vectors with a non-negative inner product with each gradient. ## URL parameters diff --git a/visualization/app.py b/tests/plots/app.py similarity index 100% rename from visualization/app.py rename to tests/plots/app.py diff --git a/visualization/requirements.txt b/tests/plots/requirements.txt similarity index 100% rename from visualization/requirements.txt rename to tests/plots/requirements.txt diff --git a/visualization/_utils.py b/visualization/_utils.py deleted file mode 100644 index 5edca32b3..000000000 --- a/visualization/_utils.py +++ /dev/null @@ -1,209 +0,0 @@ -from collections.abc import Callable - -import numpy as np -import torch -from plotly import graph_objects as go -from plotly.graph_objs import Figure, Scatter - -from torchjd.aggregation import Aggregator - - -class Plotter: - 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) - aggregators = [self._aggregator_factories[key]() for key in self.selected_keys] - results = [agg(self.matrix) for agg in aggregators] - - fig = go.Figure() - - start_angle, opening = compute_2d_dual_cone(self.matrix.numpy()) - cone = make_cone_scatter(start_angle, opening, label="Dual cone") - fig.add_trace(cone) - - for i in range(len(self.matrix)): - 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", - self.selected_keys[i], - showlegend=True, - dash=True, - ) - fig.add_trace(scatter) - - all_x = [0] + self.matrix.numpy()[:, 0].tolist() + [res[0].item() for res in results] - all_y = [0] + self.matrix.numpy()[:, 1].tolist() + [res[1].item() for res in results] - min_x = min(all_x) - max_x = max(all_x) - min_y = min(all_y) - max_y = max(all_y) - len_x = max_x - min_x - len_y = max_y - min_y - margin_prop = 0.05 - range_x = [min_x - len_x * margin_prop, max_x + len_x * margin_prop] - range_y = [min_y - len_y * margin_prop, max_y + len_y * margin_prop] - - fig.update_layout(hovermode=False, width=1000, height=900) - fig.update_xaxes(scaleanchor="y", scaleratio=1, range=range_x) - fig.update_yaxes(range=range_y) - - return fig - - -def make_vector_scatter( - gradient: torch.Tensor, - color: str, - label: str, - showlegend: bool = False, - dash: bool = False, - textposition: str = "bottom center", - line_width: float = 2.5, - text_size: float = 12, - marker_size: float = 12, -) -> Scatter: - line = {"color": color, "width": line_width} - if dash: - line["dash"] = "dash" - - scatter = go.Scatter( - x=[0, gradient[0]], - y=[0, gradient[1]], - mode="lines+markers+text", - line=line, - marker={ - "symbol": "arrow", - "color": color, - "size": marker_size, - "angleref": "previous", - }, - name=label, - text=["", label], - textposition=textposition, - textfont={"color": color, "size": text_size}, - showlegend=showlegend, - ) - return scatter - - -def make_cone_scatter( - start_angle: float, - opening: float, - label: str, - scale: float = 100.0, - printable: bool = False, -) -> Scatter: - if opening < -1e-8: - cone_outline = np.zeros([0, 2]) - else: - middle_angle = start_angle + opening / 2 - end_angle = start_angle + opening - - start_vec = angle_to_coord(start_angle, scale) - end_vec = angle_to_coord(end_angle, scale) - - if np.abs(opening - 2 * np.pi) <= 0.01: - cone_outline = np.array( - [ - [0, 0], - start_vec, - end_vec, - [0, 0], - ], - ) - else: - middle_point = angle_to_coord(middle_angle, scale) - - cone_outline = np.array( - [ - [0, 0], - start_vec, - middle_point, - end_vec, - [0, 0], - ], - ) - - if printable: - fillpattern = { - "bgcolor": "white", - "shape": "\\", - "fgcolor": "rgba(0, 220, 0, 0.5)", - "size": 30, - "solidity": 0.15, - } - else: - fillpattern = None - - cone = go.Scatter( - x=cone_outline[:, 0], - y=cone_outline[:, 1], - fill="toself", - mode="lines", - fillcolor="rgba(0, 255, 0, 0.07)", - line={"color": "rgb(0, 220, 0)", "width": 2}, - name=label, - fillpattern=fillpattern, - ) - - return cone - - -def compute_2d_dual_cone(matrix: np.ndarray) -> tuple[float, float]: - row_angles = [coord_to_angle(*row)[0] for row in matrix] - start_angles = [(angle - np.pi / 2) % (2 * np.pi) for angle in row_angles] - - cone_start_angle = start_angles[0] - opening = np.pi - for hs_start_angle in start_angles[1:]: - cone_start_angle, opening = combine_bounds(cone_start_angle, opening, hs_start_angle) - - return cone_start_angle, opening - - -def combine_bounds( - cone_start_angle: float, - opening: float, - hs_start_angle: float, -) -> tuple[float, float]: - angle_between_starts = (hs_start_angle - cone_start_angle) % (2 * np.pi) - if angle_between_starts < np.pi: - cone_start_angle = hs_start_angle - opening = opening - angle_between_starts - else: - opening = min(opening, (hs_start_angle + np.pi - cone_start_angle) % (2 * np.pi)) - - return cone_start_angle, opening - - -def coord_to_angle(x: float, y: float) -> tuple[float, float]: - r = np.sqrt(x**2 + y**2) - if r == 0: - raise ValueError("No angle") - angle = np.arccos(x / r) if y >= 0 else 2 * np.pi - np.arccos(x / r) - return angle, r - - -def angle_to_coord(angle: float, r: float = 1.0) -> tuple[float, float]: - x = r * np.cos(angle) - y = r * np.sin(angle) - return x, y From d19f44c11c7f882ab5d0e29cf414876e3b0de302 Mon Sep 17 00:00:00 2001 From: Khush Date: Sun, 2 Aug 2026 18:16:23 -0400 Subject: [PATCH 5/6] ci: trigger plotter deploy on release and workflow_dispatch --- .github/workflows/deploy-visualization.yml | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/.github/workflows/deploy-visualization.yml b/.github/workflows/deploy-visualization.yml index cafbd4fbe..1d8c82e7c 100644 --- a/.github/workflows/deploy-visualization.yml +++ b/.github/workflows/deploy-visualization.yml @@ -1,14 +1,9 @@ name: Deploy Visualization to Hugging Face Spaces on: - push: - branches: [main] - paths: - - tests/plots/** - -concurrency: - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: true + release: + types: [published] + workflow_dispatch: jobs: deploy: From 5d69829e33fd063dd50b682628449d09210e3efd Mon Sep 17 00:00:00 2001 From: Khush Date: Sun, 2 Aug 2026 18:18:42 -0400 Subject: [PATCH 6/6] ci: exclude tests/plots/app.py from ty type checking --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index bf60cdb5e..6987c55de 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -193,4 +193,4 @@ quote-style = "double" [tool.ty.src] include = ["src", "tests"] -exclude = ["src/torchjd/aggregation/_nash_mtl.py"] +exclude = ["src/torchjd/aggregation/_nash_mtl.py", "tests/plots/app.py"]