From f95e0a2c8b896574c43da42d0a46171d7e0b55c1 Mon Sep 17 00:00:00 2001 From: Khush Date: Sat, 1 Aug 2026 12:31:57 -0400 Subject: [PATCH] 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