Standardise plot interface on all Segment plot functions
#604
Conversation
…eturning axes or figure object
…rning the respective one for remaining `Segment` plot functions
There was a problem hiding this comment.
Pull request overview
This PR standardizes the plotting interface for all Segment plot functions to accept optional axes or figure objects and return the axes/figure they used, following a pattern similar to Seaborn. This improves compatibility with interactive Python editors and makes the functions more flexible.
Key Changes
- Modified parameter order to move optional
ax/fig/axx/axyparameters to the end for better consistency - All plot functions now return the figure or axes object they used
- Enhanced docstrings to document the new optional parameters and return values
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 8 comments.
| File | Description |
|---|---|
| cheetah/accelerator/segment.py | Updated five plotting methods (plot_mean_and_std, plot_overview, plot_beam_attrs_over_lattice, plot_twiss, plot_twiss_over_lattice) to accept optional figure/axes parameters and return the figure/axes used |
| CHANGELOG.md | Added entry documenting the new plotting interface feature |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| vector_idx: tuple | None = None, | ||
| figsize=(8, 4), | ||
| fig: matplotlib.figure.Figure | None = None, | ||
| ) -> None: |
There was a problem hiding this comment.
The return type annotation is None, but the function now returns a matplotlib.figure.Figure object. This should be updated to matplotlib.figure.Figure.
| ) -> None: | |
| ) -> matplotlib.figure.Figure: |
| incoming: Beam, | ||
| figsize=(8, 4), | ||
| fig: matplotlib.figure.Figure | None = None, | ||
| ) -> None: |
There was a problem hiding this comment.
The return type annotation is None, but the function now returns a matplotlib.figure.Figure object. This should be updated to matplotlib.figure.Figure.
| ) -> None: | |
| ) -> matplotlib.figure.Figure: |
| if fig is None: | ||
| fig = plt.figure(figsize=figsize) |
There was a problem hiding this comment.
The figsize parameter is only used when creating a new figure (fig is None). If a user provides an existing fig, the figsize parameter will be silently ignored. Consider either:
- Documenting that
figsizeis ignored whenfigis provided - Raising a warning or error if both
figand a non-defaultfigsizeare provided - Removing
figsizefrom the parameter list and documenting that users should create their figure with the desired size if they want to customize it
| if fig is None: | ||
| fig = plt.figure(figsize=figsize) |
There was a problem hiding this comment.
The figsize parameter is only used when creating a new figure (fig is None). If a user provides an existing fig, the figsize parameter will be silently ignored. Consider either:
- Documenting that
figsizeis ignored whenfigis provided - Raising a warning or error if both
figand a non-defaultfigsizeare provided - Removing
figsizefrom the parameter list and documenting that users should create their figure with the desired size if they want to customize it
| axx.set_xlabel("s (m)") | ||
| axy.set_ylabel("y (m)") | ||
|
|
||
| return fig |
There was a problem hiding this comment.
The function returns fig unconditionally, but fig is only defined when axx is None or axy is None. When existing axes are passed in, fig will be undefined, causing a NameError.
Consider one of these approaches:
- Return a tuple of
(fig, axx, axy)or(axx, axy)instead of just the figure - Store
fig = Nonebefore the conditional and return it (allowingNoneas a return value) - Get the figure from the axes:
fig = axx.get_figure()when axes are provided
| vector_idx: tuple | None = None, | ||
| axx: plt.Axes | None = None, | ||
| axy: plt.Axes | None = None, | ||
| ) -> None: |
There was a problem hiding this comment.
The return type annotation is None, but the function now returns a matplotlib.figure.Figure object. This should be updated to matplotlib.figure.Figure or matplotlib.figure.Figure | None depending on the intended behavior when axes are provided.
| ) -> None: | |
| ) -> plt.Figure | None: |
| resolution: float | None = None, | ||
| vector_idx: tuple | None = None, | ||
| fig: matplotlib.figure.Figure | None = None, | ||
| ) -> None: |
There was a problem hiding this comment.
The return type annotation is None, but the function now returns a matplotlib.figure.Figure object. This should be updated to matplotlib.figure.Figure.
| ) -> None: | |
| ) -> matplotlib.figure.Figure: |
| vector_idx: tuple | None = None, | ||
| axx: plt.Axes | None = None, | ||
| axy: plt.Axes | None = None, | ||
| ) -> None: |
| resolution: float | None = None, | ||
| vector_idx: tuple | None = None, | ||
| fig: matplotlib.figure.Figure | None = None, | ||
| ) -> None: |
| figsize=(8, 4), | ||
| resolution: float | None = None, | ||
| vector_idx: tuple | None = None, | ||
| figsize=(8, 4), |
There was a problem hiding this comment.
Is the figsize argument still relevant if we return the created figure at the end?
| vector_idx: tuple | None = None, | ||
| figsize=(8, 4), | ||
| fig: matplotlib.figure.Figure | None = None, | ||
| ) -> None: |
| figsize=(8, 4), | ||
| fig: matplotlib.figure.Figure | None = None, | ||
| ) -> None: |
There was a problem hiding this comment.
Same comment about figsize and the return type as above
Description
All plotting functions in
Segmentnow accept an optional axes or figure object. If they are passed on, they use it, otherwise they create one themselves. Either way, they return the axes or figure they used.Motivation and Context
Closes #601.
This is the interface used by Seaborn and makes it very convenient to use these functions on their own, or to fill already created plots. It also improves compatibility with other interactive Python editors (see #601).
Types of changes
Checklist
flake8(required).pytesttests pass (required).pyteston a machine with a CUDA GPU and made sure all tests pass (required).Note: We are using a maximum length of 88 characters per line.