test(aggregation): Improve interactive plotter

[rkhosrowshahi]

Summary

Interactive plotter

• tests/plots/interactive_plotter.py and tests/plots/_utils.py: factory-based aggregators, checklist selection, improved gradient controls

[PierreQuinton]

Ideally, this should be split in 3 PRs, in order:

• grouping.rst
• dradvac
• interactive plotter

[ValerianRey]

Ideally, this should be split in 3 PRs, in order:

• grouping.rst
• dradvac
• interactive plotter

I merged main into this and removed grouping example (59e7942), so now this only does the improvement of the plotter. I'll also update the PR title and message accordingly.

[ValerianRey]

/opencode:Plan please review this.

[opencode-agent]

Now I have both versions. Let me analyze the changes:

Code Review: Interactive Plotter Improvements

Summary of Changes

tests/plots/_utils.py

Key Changes:

1. Factory pattern for aggregators: Changed from storing aggregator instances to storing factory functions (Callable[[], Aggregator])
2. Key-based selection: Uses selected_keys to track which aggregators to display
3. Dynamic instantiation: Aggregators are created fresh each time make_fig() is called
4. Added gradient vector text positioning: Now explicitly positions labels at "top right"

Benefits:

• Ensures fresh aggregator state for each plot (important for stateful aggregators like Random, GradDrop)
• Cleaner separation between available aggregators and selected ones

tests/plots/interactive_plotter.py

Key Changes:

1. Factory-based aggregator definitions: Each aggregator is now a lambda that creates a fresh instance
2. Added helper functions: _format_angle_display() and _format_length_display() for human-readable values
3. Improved gradient controls:
   • Added formatted value displays showing both radians and degrees for angles
   • Better styling with labels and monospace font for values
   • Uses MathJax for gradient labels (e.g., $g_{1}$)
4. Updated callback signatures: update_gradient_coordinate now also updates the displayed values

Issues Found

1. Missing Unpack import handling

The code imports Unpack from typing_extensions but this is only available in Python 3.11+. For older Python versions, this could cause issues.

Location: interactive_plotter.py, line 11

2. Unused variable in loop

In update_gradient_coordinate, the display_parts list is built but could be simplified.

Current code:

display_parts: list[str] = []
for j in range(len(values_) // 2):
    ...
    display_parts.append(_format_angle_display(angle))
    display_parts.append(_format_length_display(r))
return (plotter.make_fig(), *display_parts)

This works but is slightly verbose. Not a blocker.

3. Potential issue with TrimmedMean(trim_number=1)

With the default 3 tasks, TrimmedMean(trim_number=1) will trim 1 from each side (2 total), leaving only 1 gradient. This might produce unexpected results but is acceptable for testing.

Positive Aspects

1. Factory pattern is well-implemented: Solves the stateful aggregator issue properly
2. UI improvements are excellent: The formatted displays and better styling significantly improve usability
3. Clean code structure: The helper functions and styling dictionaries are well-organized
4. MathJax integration: Using LaTeX formatting for gradient labels is a nice touch

Recommendations

1. Consider adding a comment about why the factory pattern is used (stateful aggregators):

# Use factory functions to ensure fresh aggregator state for each plot
# (important for stateful aggregators like Random, GradDrop)

2. The typing_extensions dependency - Verify this is already in the project's dependencies since Unpack is used.

Overall Assessment

✅ Approve - This is a well-structured improvement that:

• Properly addresses state management for aggregators
• Significantly improves the UI/UX of the interactive plotter
• Follows the existing code style
• Makes the testing tool more useful and user-friendly

The changes are focused, minimal, and achieve the stated goals of factory-based aggregators, checklist selection, and improved gradient controls.

opencode session  |  github run

[ValerianRey]

The code imports Unpack from typing_extensions but this is only available in Python 3.11+. For older Python versions, this could cause issues.

The interactive plotter is only for torchjd devs (not people who pip install the library), so it's fine if it doesn't work with python 3.10. So IMO this is not an issue.

[ValerianRey]

Really nice. Few nitpicks and we can merge.

[ValerianRey]

It seems like it would be nice to change the __str__ method of AlignedMTL to also include the scale_mode. For example, str(AlignedMTL(scale_mode="min")) would become AlignedMTL-min.

But this is quite independent from this PR, and should come in a different PR if you care about that! @rkhosrowshahi

[ValerianRey]

We could make another PR that adds the possibility to change the number of gradients. A few things to keep in mind:

• Some aggregators like NashMTL will have to be re-created based on the number of selected gradients (so that the num_tasks parameters can match that).
• Some aggregators like TrimmedMean require at least 3 gradients, and there's probably some aggregators that require at least 2. They would have to be deactivated when the number of gradients is not enough.

So it's not trivial, but if you're interested in doing that, we could land the PR quite easily @rkhosrowshahi