Plotting

Plotting Styles

The style submodules provide styles to be used with matplotlib, mostly tailored for my use, and for good results with the plotters in plotting. Feel free to use them anyway, as they might be useful to you when using the plotting submodules, or to be adapted.

pyhdtoolkit.plotting.styles.install_mpl_styles() → None[source]

Added in version 1.0.0.

Installs the styles defined in the styles submodules to disk as .mplstyle files. This way, they can be used with use without having to import them and update the rcParams directly.

Paper Styles

This module contains different styles to be used with matplotlib, tailored for figures to be included in my journal papers. The end results should insert themselves seamlessly with the rest of the paper. All plots are included as either a single or double column figure in these papers.

The following are available:

  • SINGLE_COLUMN: For plots to be included in a single column figure environment.

  • DOUBLE_COLUMN: For plots to be included in a full width (double column) figure environment.

Thesis Styles

This module contains different styles to be used with matplotlib, tailored for figures to be included in my thesis document. The end results should insert themselves seamlessly with the rest of the paper. All plots are included in a LaTeX figure environment, included with [width=0.99\columnwidth] (and potentially subfloats inside). The text body has a fontsize of 12 points, and the figures should integrate to have roughly the same text size.

The following are available:

  • SMALL: For small simple plots to be included in a multi-figure environment (e.g. a LaTeX figure with 2 subfloats). Two of these should render well when displayed side-by-side in the figure environment (\subfloat[.9\linewidth] then \hspace{0.3cm} then \subfloat[.9\linewidth]).

  • MEDIUM: For simple plots to be included alone in a LaTeX figure environment (e.g single axis line plots, or scatters with a colorbar like in tune).

  • LARGE: For more complex plots to be included alone in a LaTeX figure environment (e.g. multi-axes figures such as in lattice).

Aperture Plotters

Module with functions to create aperture plots through a Madx object.

pyhdtoolkit.plotting.aperture.plot_aperture(madx: Madx, /, title: str | None = None, xoffset: float = 0, xlimits: tuple[float, float] | None = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, aperture_ylim: tuple[float, float] | None = None, k0l_lim: tuple[float, float] | float | None = None, k1l_lim: tuple[float, float] | float | None = None, k2l_lim: tuple[float, float] | float | None = None, k3l_lim: tuple[float, float] | float | None = None, color: str | None = None, **kwargs) → None[source]

Added in version 1.0.0.

Creates a plot representing the lattice layout and the aperture tolerance across the machine. The tolerance is based on the n1 values in the aperture table. One can find an example use of this function in the machine aperture example gallery.

Important

This function assumes the user has previously made a call to the APERTURE command in MAD-X, as it will query relevant values from the aperture table.

Note

This function has some heavy logic behind it, especially in how it needs to order several axes. The easiest way to go about using it is to manually create and empty figure with the desired properties (size, etc) then call this function. See the example below or the gallery for more details.

Warning

Currently the function tries to plot legends for the different layout patches. The position of the different legends has been hardcoded in corners and might require users to tweak the axis limits (through k0l_lim, k1l_lim and k2l_lim) to ensure legend labels and plotted elements don’t overlap.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • title (str, optional) -- Title of the figure.

  • xoffset (float) -- An offset applied to the S coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at exactly \(s = 0\). Beware this offset is applied before applying the xlimits. Defaults to 0.

  • xlimits (tuple[float, float], optional) -- If given, will be used for the xlim (for the s coordinate), using the tuple passed.

  • plot_dipoles (bool) -- If True, dipole patches will be plotted on the layout subplot of the figure. Defaults to True. Dipoles are plotted in blue.

  • plot_dipole_k1 (bool) -- If True, dipole elements with a quadrupolar gradient will have this gradient plotted as a quadrupole patch. Defaults to False.

  • plot_quadrupoles (bool) -- If True, quadrupole patches will be plotted on the layout subplot of the figure. Defaults to True. Quadrupoles are plotted in red.

  • plot_bpms (bool) -- If True, additional patches will be plotted on the layout subplot to represent Beam Position Monitors. BPMs are plotted in dark grey. Defaults to False.

  • aperture_ylim (tuple[float, float], optional) -- If given, will be used as vertical axis limits for the aperture values. Defaults to None, to be determined by matplotlib based on the provided values.

  • k0l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k0l values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k0l values of the dipoles.

  • k1l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k1l values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k1l values of the quadrupoles.

  • k2l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k2l values of the sextupoles.

  • k3l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k3l values of the octupoles.

  • color (str, optional) -- The color argument given to the aperture lines. Defaults to None, in which case the first color found in your rcParams’s cycler will be used.

  • **kwargs -- Any keyword argument will be transmitted to plot_machine_layout, later on to _plot_lattice_series, and then Rectangle, such as lw etc.

Example

plt.figure(figsize=(16, 11))
plot_aperture(
    madx,
    plot_bpms=True,
    aperture_ylim=(0, 20),
    k0l_lim=(-4e-4, 4e-4),
    k1l_lim=(-0.08, 0.08),
    color="darkslateblue",
)
pyhdtoolkit.plotting.aperture.plot_physical_apertures(madx, /, plane: str, scale: float = 1, xoffset: float = 0, xlimits: tuple[float, float] | None = None, **kwargs) → None[source]

Added in version 1.2.0.

Determine and plot the “real” physical apertures of elements in the sequence. A data point is extrapolated at the beginning and the end of each element, with values based on the aper_1 and aper_2 columns in the TWISS table. One can find an example use of this function in the machine aperture example gallery. Original code from Elias Waagaard.

Important

This function assumes the user has previously made a call to the APERTURE command in MAD-X, as it will query relevant values from the aperture table.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • plane (str) -- The physical plane to plot for, Accepted values are either x, horizontal, y or vertical. Case insensitive.

  • scale (float) -- A scaling factor to apply to the beam orbit and beam enveloppe, for the user to adjust to their wanted scale. Defaults to 1 (meaning values are assumed in [m]).

  • xoffset (float) -- An offset applied to the S coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at exactly \(s = 0\). Beware this offset is applied before applying the xlimits. Defaults to 0.

  • xlimits (tuple[float, float], optional) -- If given, will be used for the xlim (for the s coordinate), using the tuple passed.

  • **kwargs -- Any keyword argument that can be given to the MAD-X TWISS command. If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Raises:

ValueError -- If the plane argument is not one of x, horizontal, y or vertical.

Examples

fig, ax = plt.subplots(figsize=(10, 9))
plot_physical_apertures(madx, "x")
plt.show()

In order to do the same plot but have all values in millimeters:

fig, ax = plt.subplots(figsize=(10, 9))
plot_physical_apertures(madx, "x", scale=1e3)
plt.setp(ax, xlabel="S [m]", ylabel="X [mm]")
plt.show()

Crossing Scheme Plotters

Module with functions to plot LHC crossing schemes through a Madx object.

pyhdtoolkit.plotting.crossing.plot_single_ir_crossing(axis: Axes, plot_df_b1: DataFrame, plot_df_b2: DataFrame, plot_column: str, scaling: float = 1, xlabel: str | None = None, ylabel: str | None = None, title: str | None = None) → None[source]

Added in version 1.0.0.

Plots the X or Y orbit for the IR on the given axis.

Warning

This function assumes the provided the plot_df_b1 and plot_df_b2 are already centered at 0 on the IP point!

Parameters:

  • axis (matplotlib.axes.Axes) -- The Axes on which to plot.

  • plot_df_b1 (pd.DataFrame | tfs.TfsDataFrame) -- The TWISS dataframe of the IR zone for beam 1 of the LHC, centered on 0 at IP position (this can be achieved very simply with df.s = df.s - ip_s).

  • plot_df_b2 (pd.DataFrame | tfs.TfsDataFrame) -- The TWISS dataframe of the IR zone for beam 2 of the LHC, centered on 0 at IP position (this can be achieved very simply with df.s = df.s - ip_s).

  • plot_column (str) -- Which column (should be x or y) to plot for the orbit.

  • scaling (float) -- Scaling factor to apply to the plotted data. Defaults to 1 (no change of data).

  • xlabel (str, optional) -- If given, will be used for the xlabel of the axis.

  • ylabel (str, optional) -- If given, will be used for the ylabel of the axis.

  • title (str, optional) -- If given, will be used for the title of the axis.

Example

plot_single_ir_crossing(
    plt.gca(),
    b1_df,
    b2_df,
    plot_column="x",
    scaling=1e3,
    ylabel="Orbit X $[mm]$",
)
pyhdtoolkit.plotting.crossing.plot_two_lhc_ips_crossings(madx: Madx, /, first_ip: int, second_ip: int, ir_limit: float = 275, highlight_mqx_and_mbx: bool = True) → None[source]

Added in version 1.0.0.

Creates a plot representing the crossing schemes at the two provided IPs. One can find an example use of this function in the LHC crossing schemes example gallery.

Note

This function has some heavy logic behind it, especially in how it needs to order several axes. The easiest way to go about using it is to manually create and empty figure with the desired properties (size, etc) then call this function. See the example below or the gallery for more details.

Note

This assumes the appropriate LHC sequence and opticsfile have been loaded, and both lhcb1 and lhcb2 beams are defined. It is very recommended to first re-cycle the sequences so that the desired IPs do not happen at beginning or end of the lattice.

Warning

This function will get TWISS tables for both beams, which means it will USE both the lhcb1 and lhcb2 sequences, erasing previously defined errors or orbit corrections. The second sequence USE will be called on is lhcb2, which may not be the one you were using before. Please re-use your wanted sequence after you have called this function!

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • first_ip (int) -- The first of the two IPs to plot crossing schemes for.

  • second_ip (int) -- The second of the two IPs to plot crossing schemes for.

  • ir_limit (float) -- The amount of meters to keep left and right of the IP point. Will also determine the xlimits of the plots. Defaults to 275.

  • highlight_mqx_and_mbx (bool) -- If True, will add patches highlighting the zones corresponding to MBX and MQX elements. Defaults to True.

  • **kwargs -- If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Examples

plt.figure(figsize=(18, 12))
plot_two_lhc_ips_crossings(madx, first_ip=1, second_ip=5)

plt.figure(figsize=(16, 11))
plot_two_lhc_ips_crossings(
    madx, first_ip=2, second_ip=8, highlight_mqx_and_mbx=False
)

Beam Enveloppe Plotters

Module with functions to create beam enveloppe plots through a Madx object.

pyhdtoolkit.plotting.envelope.plot_beam_envelope(madx: Madx, /, sequence: str, plane: str, nsigma: float = 1, scale: float = 1, xoffset: float = 0, xlimits: tuple[float, float] | None = None, **kwargs) → None[source]

Added in version 1.2.0.

Draws the beam enveloppe around the beam orbit on the given axis. The enveloppe is determined from the active sequence’s beam’s parameters.

One can find an example use of this function in the beam enveloppe example gallery.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • sequence (str) -- The name of the sequence to plot the beam enveloppe for, which should be the active sequence. Case insensitive.

  • plane (str) -- The physical plane to plot for, which should be either x, horizontal, y or vertical. Case insensitive.

  • nsigma (float) -- The standard deviation to use for the beam enveloppe calculation. For instance, providing a value of 3 will draw the 3 sigma beam enveloppe. Defaults to 1.

  • scale (float) -- A scaling factor to apply to the beam orbit and beam enveloppe, for the user to adjust to their wanted scale. Defaults to 1 (meaning values are assumed in [m]).

  • xoffset (float) -- An offset applied to the S coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at \(s = 0\). Beware this offset is applied before applying the xlimits. Defaults to 0.

  • xlimits (tuple[float, float], optional) -- If given, will be used for the xlim (for the s coordinate), using the tuple passed.

  • **kwargs -- Any keyword argument that can be given to the MAD-X TWISS command. If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Raises:

ValueError -- If the plane argument is not one of x, horizontal, y or vertical.

Examples

fig, ax = plt.subplots(figsize=(10, 9))
plot_beam_envelope(madx, "lhcb1", "x", nsigma=3)
plt.show()

In order to do the same plot but have all values in millimeters:

fig, ax = plt.subplots(figsize=(10, 9))
plot_beam_envelope(madx, "lhcb1", "x", nsigma=3, scale=1e3)
plt.setp(ax, xlabel="S [m]", ylabel="X [mm]")
plt.show()

Lattice Plotters

Module with functions to create lattice plots through a Madx object.

pyhdtoolkit.plotting.lattice.plot_latwiss(madx: Madx, /, title: str | None = None, xoffset: float = 0, xlimits: tuple[float, float] | None = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, disp_ylim: tuple[float, float] | float | None = None, beta_ylim: tuple[float, float] | float | None = None, k0l_lim: tuple[float, float] | float | None = None, k1l_lim: tuple[float, float] | float | None = None, k2l_lim: tuple[float, float] | float | None = None, k3l_lim: tuple[float, float] | float | None = None, **kwargs) → None[source]

Added in version 1.0.0.

Creates a plot on the current figure (gcf) representing the lattice layout and the \(\beta\)-functions along with the horizontal dispertion function. This is a very, very heavily refactored version of an initial implementation by Guido Sterbini. One can find example uses of this function in the machine lattice example gallery.

Note

This function has some heavy logic behind it, especially in how it needs to order several axes. The easiest way to go about using it is to manually create an empty figure with the desired properties (size, etc) then call this function. See the example below or the gallery for more details.

Important

At the moment, it is important to give this function symmetric limits for the k0l_lim, k1l_lim and k2l_lim arguments. Otherwise the element patches will show up vertically displaced from the axis’ center line.

Warning

Currently the function tries to plot legends for the different layout patches. The position of the different legends has been hardcoded in corners and might require users to tweak the axis limits (through k0l_lim, k1l_lim and k2l_lim) to ensure legend labels and plotted elements don’t overlap.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • title (str, optional) -- If provided, is set as title of the plot.

  • xoffset (float) -- An offset applied to the S coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at \(s = 0\). Beware this offset is applied before applying the xlimits. Defaults to 0.

  • xlimits (tuple[float, float], optional) -- If given, will be used for the xlim (for the s coordinate), using the tuple passed.

  • plot_dipoles (bool) -- If True, dipole patches will be plotted on the layout subplot of the figure. Defaults to True. Dipoles are plotted in blue.

  • plot_dipole_k1 (bool) -- If True, dipole elements with a quadrupolar gradient will have this gradient plotted as a quadrupole patch. Defaults to False.

  • plot_quadrupoles (bool) -- If True, quadrupole patches will be plotted on the layout subplot of the figure. Defaults to True. Quadrupoles are plotted in red.

  • plot_bpms (bool) -- If True, additional patches will be plotted on the layout subplot to represent Beam Position Monitors. BPMs are plotted in dark grey. Defaults to False.

  • disp_ylim (tuple[float, float] | float, optional) -- If If given, will be used as vertical axis limits for the dispersion values. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). Defaults to None, and will be determined by matplotlib based on the dispersion values.

  • beta_ylim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the betatron function values. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). Defaults to None, and will be determined by matplotlib based on the beta values.

  • k0l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k0l values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k0l values of the dipoles.

  • k1l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k1l values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k1l values of the quadrupoles.

  • k2l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k2l values of the sextupoles.

  • k3l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k3l values of the octupoles.

  • **kwargs -- Any keyword argument will be transmitted to plot_machine_layout, later on to _plot_lattice_series, and then Rectangle, such as lw etc.

Examples

title = "Machine Layout"
plt.figure(figsize=(16, 11))
plot_latwiss(
    madx,
    title=title,
    k0l_lim=(-0.15, 0.15),
    k1l_lim=(-0.08, 0.08),
    disp_ylim=(-10, 125),
    lw=3,
)

One can provide ylimits for the machine layout patches as single values:

title = "Machine Layout"
plt.figure(figsize=(16, 11))
plot_latwiss(
    madx,
    title=title,
    k0l_lim=0.15,  # identical to k0l_lim=(-0.15, 0.15)
    k1l_lim=0.08,  # identical to k1l_lim=(-0.08, 0.08)
    disp_ylim=(-10, 125),
    lw=3,
)
pyhdtoolkit.plotting.lattice.plot_machine_survey(madx: Madx, /, title: str | None = None, show_elements: bool = False, high_orders: bool = False, **kwargs) → Axes[source]

Added in version 1.0.0.

Creates a plot representing the lattice layout and the machine geometry in 2D. This is a very, very heavily refactored version of an initial implementation by Guido Sterbini. One can find an example use of this function in the machine survey example gallery.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • title (str, optional) -- If provided, is set as title of the plot.

  • show_elements (bool) -- If True, will try to plot by differentiating elements. Defaults to False.

  • high_orders (bool) -- If True, plots sextupoles and octupoles if show_elements is True, otherwise only up to quadrupoles. Defaults to False.

  • **kwargs -- Any keyword argument is transmitted to scatter If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns:

matplotlib.axes.Axes -- The Axes on which the survey is drawn.

Example

fig, ax = plt.subplots(figsize=(6, 6))
plot_machine_survey(
    madx, title="Machine Survey", show_elements=True, high_orders=True
)

Layout Plotters

Module with functions used to represent a machine’s elements in an Axes object, mostly used in different plotting modules.

pyhdtoolkit.plotting.layout.plot_machine_layout(madx: Madx, /, title: str | None = None, xoffset: float = 0, xlimits: tuple[float, float] | None = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, k0l_lim: tuple[float, float] | float | None = None, k1l_lim: tuple[float, float] | float | None = None, k2l_lim: tuple[float, float] | float | None = None, k3l_lim: tuple[float, float] | float | None = None, **kwargs) → None[source]

Added in version 1.0.0.

Draws patches elements representing the lattice layout on the given axis. This is the function that takes care of the machine layout axis in plot_latwiss and plot_aperture. Its results can be seen in the machine lattice and machine aperture example galleries.

Note

This current implementation can plot dipoles, quadrupoles, sextupoles, octupoles and BPMs.

Important

If not provided, the limits for the k0l_lim, k1l_lim will be auto-determined, which might not be the perfect choice for the plot. When providing these limits (also for k2l_lim), make sure to provide symmetric values around 0 (so [-x, x]) otherwise the element patches will show up vertically displaced from the axis’ center line.

Warning

Currently the function tries to plot legends for the different layout patches. The position of the different legends has been hardcoded in corners of the Axes and might require users to tweak the axis limits (through k0l_lim, k1l_lim and k2l_lim) to ensure legend labels and plotted elements don’t overlap.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • title (str, optional) -- If provided, is set as title of the plot.

  • xoffset (float) -- An offset applied to the S coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at \(s = 0\). Beware this offset is applied before applying the xlimits. Defaults to 0.

  • xlimits (tuple[float, float], optional) -- If given, will be used for the xlim (for the s coordinate), using the tuple passed.

  • plot_dipoles (bool) -- If True, dipole patches will be plotted on the layout subplot of the figure. Defaults to True. Dipoles are plotted in blue.

  • plot_dipole_k1 (bool) -- If True, dipole elements with a quadrupolar gradient will have this gradient plotted as a quadrupole patch. Defaults to False.

  • plot_quadrupoles (bool) -- If True, quadrupole patches will be plotted on the layout subplot of the figure. Defaults to True. Quadrupoles are plotted in red.

  • plot_bpms (bool) -- If True, additional patches will be plotted on the layout subplot to represent Beam Position Monitors. BPMs are plotted in dark grey. Defaults to False.

  • k0l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k0l values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k0l values of the dipoles.

  • k1l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k1l values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k1l values of the quadrupoles.

  • k2l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k2l values of the sextupoles.

  • k3l_lim (tuple[float, float] | float, optional) -- If given, will be used as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If None is given, then the limits will be determined automatically based on the k3l values of the octupoles.

  • **kwargs -- Any keyword argument will be transmitted to _plot_lattice_series, and then Rectangle, such as lw etc. If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on. By definition, the quadrupole elements will be drawn on said axis, and for each new element type to plot a call to twinx is made and the new elements will be drawn on the newly created twin Axes. If bpms_legend is given as False and BPMs are plotted, the BPM legend will not be plotted on the layout axis.

Example

fig, ax = plt.subplots(figsize=(6, 2))
plot_machine_layout(madx, title="Machine Elements", lw=3)
pyhdtoolkit.plotting.layout.scale_patches(scale: float, ylabel: str, **kwargs) → None[source]

Added in version 1.3.0.

This is a convenience function to update the scale of the elements layout patches as well as the corresponding y-axis label.

Parameters:

  • scale (float) -- The scale factor to apply to the patches. The new height of the patches will be scale * original_height.

  • ylabel (str) -- The new label for the y-axis.

  • **kwargs -- If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on, otherwise the current axis is used.

Example

fig, ax = plt.subplots(figsize=(6, 2))
plot_machine_layout(madx, title="Machine Elements", lw=3)
scale_patches(ax=fig.axes[0], scale=100, ylabel=r"$K_{1}L$ $[10^{-2} m^{-1}]$")

Phase Space Plotters

Module with functions to create phase space plots through a Madx object.

pyhdtoolkit.plotting.phasespace.plot_courant_snyder_phase_space(madx: Madx, /, u_coordinates: np.ndarray, pu_coordinates: np.ndarray, plane: str = 'Horizontal', title: str | None = None, **kwargs) → Axes[source]

Added in version 1.0.0.

Creates a plot representing the normalized Courant-Snyder phase space of a particle distribution when provided by position and momentum coordinates for a specific plane. One can find an example use of this function in the phase space example gallery.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • u_coordinates (numpy.ndarray) -- A ndarray of particles’ coordinates for the given plane. Here u_coordinates[0] should be the tracked coordinates for the first particle and so on.

  • pu_coordinates (numpy.ndarray) -- A ndarray of particles’ momentum coordinates for the given plane. Here pu_coordinates[0] should be the tracked momenta for the first particle and so on.

  • plane (str) -- The physical plane to plot, should be either horizontal or vertical`, case insensitive. Defaults to horizontal.

  • title (str, optional) -- If provided, is set as title of the plot.

  • **kwargs -- If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns:

matplotlib.axes.Axes -- The Axes on which the phase space is drawn.

Example

fig, ax = plt.subplots(figsize=(10, 9))
plot_courant_snyder_phase_space(madx, x_coords, px_coords, plane="Horizontal")
pyhdtoolkit.plotting.phasespace.plot_courant_snyder_phase_space_colored(madx: Madx, /, u_coordinates: np.ndarray, pu_coordinates: np.ndarray, plane: str = 'Horizontal', title: str | None = None, **kwargs) → Figure[source]

Added in version 1.0.0.

Creates a plot representing the normalized Courant-Snyder phase space of a particle distribution when provided by position and momentum coordinates for a specific plane. Each single particle trajectory has its own color on the plot, within the limit of pyplot’s 156 named colors, after which the colors loop back to the first entry again. One can find an example use of this function in the phase space example gallery.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • u_coordinates (numpy.ndarray) -- A ndarray of particles’ coordinates for the given plane. Here u_coordinates[0] should be the tracked coordinates for the first particle and so on.

  • pu_coordinates (numpy.ndarray) -- A ndarray of particles’ momentum coordinates for the given plane. Here pu_coordinates[0] should be the tracked momenta for the first particle and so on.

  • plane (str) -- The physical plane to plot, should be either horizontal or vertical`, case insensitive. Defaults to horizontal.

  • title (str, optional) -- If provided, is set as title of the plot.

  • **kwargs -- If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns:

matplotlib.axes.Axes -- The Axes on which the phase space is drawn.

Example

fig, ax = plt.subplots(figsize=(10, 9))
plot_courant_snyder_phase_space_colored(
    madx, x_coords, px_coords, plane="Horizontal"
)

Tune Diagram Plotters

Module with functions to create tune diagram plots. These provide functionality to draw Farey sequences up to a desired order.

pyhdtoolkit.plotting.tune.farey_sequence(order: int) → list[tuple[int, int]][source]

Added in version 1.0.0.

Returns the n-th farey_sequence sequence, ascending, where n is the provided order. Original code from Rogelio Tomás (see Numerical Methods 2018 CAS proceedings, Tomás et al. [TBC+18]).

Parameters:

order (int) -- The order up to which we want to calculate the sequence.

Returns:

list[tuple[int, int]] -- The sequence as a list of plottable 2D points.

pyhdtoolkit.plotting.tune.plot_resonance_lines_for_order(order: int, axis: Axes, **kwargs) → None[source]

Added in version 1.0.0.

Plot resonance lines from farey sequences of the given order on the provided Axes.

Parameters:

  • order (int) -- The order of the resonance.

  • axis (matplotlib.axes.Axes) -- The Axes on which to plot the resonance lines.

  • **kwargs -- Any keyword argument is given to plot.

Example

fig, ax = plt.subplots(figsize=(6, 6))
plot_resonance_lines_for_order(order=3, axis=ax, color="blue")
pyhdtoolkit.plotting.tune.plot_tune_diagram(title: str | None = None, legend_title: str | None = None, max_order: int = 6, differentiate_orders: bool = False, **kwargs) → Axes[source]

Added in version 1.0.0.

Creates a plot representing the tune diagram up to the given max_order. One can find an example use of this function in the tune diagram example gallery.

Note

The first order lines make up the [(0, 0), (0, 1), (1, 1), (1, 0)] square and will only be seen when redefining the limits of the figure, which are by default [0, 1] on each axis.

Parameters:

  • title (str, optional) -- If provided, is set as title of the plot.

  • legend_title (str, optional) -- If given, will be used as the title of the plot’s legend.

  • max_order (int) -- The order up to which to plot resonance lines for. This parameter value should not exceed 6. Defaults to 6.

  • differentiate_orders (bool) -- If True, the lines for each order will be of a different color. When set to False, there is still differentation through alpha, linewidth and linestyle. Defaults to False.

  • **kwargs -- Any keyword argument is given to plot. Be aware that alpha, ls, lw, color and label are already set by this function and providing them as kwargs might lead to errors. If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns:

matplotlib.axes.Axes -- The Axes on which the tune diagram is drawn.

Raises:

ValueError -- If the max_order is not between 1 and 6, included.

Example

fig, ax = plt.subplots(figsize=(6, 6))
plot_tune_diagram(ax=ax, max_order=4, differentiate_orders=True)

Segment-by-Segment Coupling

Functions to plot coupling components of Segment-by-Segment results.

pyhdtoolkit.plotting.sbs.coupling.plot_full_ip_rdt(b1_segment_df: tfs.TfsDataFrame, b2_segment_df: tfs.TfsDataFrame, b1_model: tfs.TfsDataFrame = None, b2_model: tfs.TfsDataFrame = None, ip: int | None = None, rdt: str = 'F1001', abs_ylimits: tuple[float, float] | None = None, real_ylimits: tuple[float, float] | None = None, imag_ylimits: tuple[float, float] | None = None, **kwargs) → Figure[source]

Added in version 0.19.0.

Plots all component of the given coupling rdt over the segment, for both Beam 1 and Beam 2. Optionally highlights the IP location. One can find an example use of this function in the segment-by-segment plotting example gallery.

Parameters:

  • b1_segment_df (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment coupling result for Beam 1 in the given segment.

  • b2_segment_df (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment coupling result for Beam 2 in the given segment.

  • b1_model (tfs.TfsDataFrame, optional) -- A TfsDataFrame of the Beam 1 model used in the analysis. If given then the IP location in the segment will be highlighted by a vertical grey line.

  • b2_model (tfs.TfsDataFrame, optional) -- A TfsDataFrame of the Beam 2 model used in the analysis. If given then the IP location in the segment will be highlighted by a vertical grey line.

  • ip (int, optional) -- The IP number of the segment. Used for the label of the vertical grey line. Requires to have provided the model dataframe.

  • rdt (str) -- The name of the coupling resonance driving term to plot, either F1001 or F1010. Case insensitive.

  • **kwargs -- Keyword arguments will be transmitted to the figure creation call to subplots. If b1_ylabel or b2_ylabel are found, they will be used as y-label for the respective beams axes. If bbox_to_anchor is found, it will be used to position the legend across the whole figure space.

Returns:

matplotlib.figure.Figure -- The Figure on which the plot is created.

Example

fig = plot_full_ip_rdt(
    couple_b1_tfs,
    couple_b2_tfs,
    b1_model_tfs,
    b2_model_tfs,
    ip=1,
    figsize=(18, 9),
    abs_ylimits=(5e-3, 6.5e-2),
    real_ylimits=(-1e-1, 1e-1),
    imag_ylimits=(-1e-1, 1e-1),
)
pyhdtoolkit.plotting.sbs.coupling.plot_rdt_component(b1_segment_df: tfs.TfsDataFrame, b2_segment_df: tfs.TfsDataFrame, b1_model: tfs.TfsDataFrame = None, b2_model: tfs.TfsDataFrame = None, ip: int | None = None, rdt: str = 'F1001', component: str = 'ABS', **kwargs) → Figure[source]

Added in version 0.19.0.

Plots for Beam 1 and Beam 2 vertically a component of the given coupling rdt over the segment. Optionally highlights the IP location. One can find an example use of this function in the segment-by-segment plotting example gallery.

Parameters:

  • b1_segment_df (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment coupling result for Beam 1 in the given segment.

  • b2_segment_df (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment coupling result for Beam 2 in the given segment.

  • b1_model (tfs.TfsDataFrame, optional) -- A TfsDataFrame of the Beam 1 model used in the analysis. If given then the IP location in the segment will be highlighted by a vertical grey line.

  • b2_model (tfs.TfsDataFrame, optional) -- A TfsDataFrame of the Beam 2 model used in the analysis. If given then the IP location in the segment will be highlighted by a vertical grey line.

  • ip (int, optional) -- The IP number of the segment. Used for the label of the vertical grey line. Requires to have provided the model dataframe.

  • rdt (str) -- The name of the coupling resonance driving term to plot, either F1001 or F1010. Case insensitive.

  • component (str) -- Which component of the RDT is considered, either ABS, RE or IM, for absolute value or real / imaginary part, respectively. Case insensitive.

  • **kwargs -- Keyword arguments will be transmitted to the figure creation call to subplots. If b1_ylabel or b2_ylabel are found, they will be used as y-label for the respective beams axes.

Returns:

matplotlib.figure.Figure -- The Figure on which the plot is created.

Example

fig = plot_rdt_component(
    b1_segment_df=tfs.read("B1/sbscouple_IP1.out"),
    b2_segment_df=tfs.read("B2/sbscouple_IP1.out"),
    b1_model=b1_model_tfs,
    b2_model=b2_model_tfs,
    ip=1,
    figsize=(8, 8),
    b1_ylabel=r"$\mathrm{Beam\ 1}$ $|f_{1001}|$",
    b2_ylabel=r"$\mathrm{Beam\ 2}$ $|f_{1001}|$",
)

Segment-by-Segment Phase

Functions to plot phase values of Segment-by-Segment results.

pyhdtoolkit.plotting.sbs.phase.plot_phase_segment(ax: Axes = None, segment_df: tfs.TfsDataFrame = None, model_df: tfs.TfsDataFrame = None, plane: str = 'x', ip: int | None = None) → None[source]

Added in version 0.19.0.

Plots a the phase for a given plane over the segment, optionally highlighting the IP location. One can find an example use of this function in the segment-by-segment plotting example gallery.

Parameters:

  • ax (matplotlib.axes.Axes, optional) -- The Axes to plot on. Will get the current axis if no Axes is given.

  • segment_df (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment coupling result for the given segment.

  • model_df (tfs.TfsDataFrame, optional) -- A TfsDataFrame of the model used in the analysis. If given, then the IP location in the segment will be determined from the two dataframes and will be highlighted in the plot by a vertical grey line.

  • plane (str) -- The plane the data is is for in the provided segment_df. Will be used for the ylabel. Should be either “x” or “y”, case-insensitive.

  • ip (int, optional) -- The IP number of the segment. Used for the label of the vertical grey line. Requires to have provided the model dataframe.

Example

plot_phase_segment(ax, segment_df, b1_model_tfs, plane="x", ip=1)
pyhdtoolkit.plotting.sbs.phase.plot_phase_segment_both_beams(b1_phase_x: tfs.TfsDataFrame, b1_phase_y: tfs.TfsDataFrame, b2_phase_x: tfs.TfsDataFrame, b2_phase_y: tfs.TfsDataFrame, b1_model: tfs.TfsDataFrame = None, b2_model: tfs.TfsDataFrame = None, ip: int | None = None, **kwargs) → Figure[source]

Added in version 0.19.0.

Plots the propagated measured phase and the propagated corrected phase for the given IP segment, for both planes and both beams. Optionally highlights the IP location in the segment. One can find an example use of this function in the segment-by-segment plotting example gallery.

Parameters:

  • b1_phase_x (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment phase result for the horizontal plane in the given segment, for Beam 1.

  • b1_phase_y (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment phase result for the vertical plane in the given segment, for Beam 1.

  • b2_phase_x (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment phase result for the horizontal plane in the given segment, for Beam 2.

  • b2_phase_x -- A TfsDataFrame of the segment-by-segment phase result for the vertical plane in the given segment, for Beam 2.

  • b1_model (tfs.TfsDataFrame, optional) -- A TfsDataFrame of the Beam 1 model used in the analysis. If given, then the IP location in the segment will be highlighted by a vertical grey line.

  • b2_model (tfs.TfsDataFrame, optional) -- A TfsDataFrame of the Beam 2 model used in the analysis. If given, then the IP location in the segment will be highlighted by a vertical grey line.

  • ip (int, optional) -- The IP number of the segment. Used for the label of the vertical grey line. Requires to have provided the model dataframe.

Returns:

matplotlib.figure.Figure -- The Figure on which the plot is created.

Example

fig = plot_phase_segment_both_beams(
    phasex_b1_tfs,
    phasey_b1_tfs,
    phasex_b2_tfs,
    phasey_b2_tfs,
    b1_model_tfs,
    b2_model_tfs,
    ip=1,
    figsize=(18, 9),
    bbox_to_anchor=(0.535, 0.94),
)
pyhdtoolkit.plotting.sbs.phase.plot_phase_segment_one_beam(phase_x: tfs.TfsDataFrame, phase_y: tfs.TfsDataFrame, model: tfs.TfsDataFrame = None, ip: int | None = None, **kwargs) → Figure[source]

Added in version 0.19.0.

Plots the propagated measured phase and the propagated corrected phase for the given IP segment, for both planes for a given beam. Optionally highlights the IP location in the segment. One can find an example use of this function in the segment-by-segment plotting example gallery.

Parameters:

  • phase_x (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment phase result for the horizontal plane in the given segment.

  • phase_y (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment phase result for the vertical plane in the given segment.

  • model (tfs.TfsDataFrame, optional) -- A TfsDataFrame of the model used in the analysis. If given, the IP location in the segment will be highlighted by a vertical grey line.

  • ip (int, optional) -- The IP number of the segment. Used for the label of the vertical grey line. Requires to have provided the model dataframe.

  • **kwargs -- Keyword arguments will be transmitted to the figure creation call to subplots. If bbox_to_anchor is found, it will be used to position the legend across the whole figure space.

Returns:

matplotlib.figure.Figure -- The Figure on which the plot is created.

Example

fig = plot_phase_segment_one_beam(
    sbs_phasex, sbs_phasey, model=b2_model_tfs, ip=5, figsize=(8, 8)
)

Plotting Utility Functions

Module with functions to used throught the different plotting modules.

pyhdtoolkit.plotting.utils.draw_confidence_ellipse(x: ArrayLike, y: ArrayLike, n_std: float = 3.0, facecolor='none', **kwargs) → Ellipse[source]

Added in version 1.2.0.

Plot the covariance confidence ellipse of x and y. Credits: this code is taken from the examples in the matplotlib gallery.

Note

One might want to provide the edgecolor to this function.

Parameters:

  • x (ArrayLike) -- Array-like, should be of shape (n,).

  • y (ArrayLike) -- Array-like, should be of shape (n,).

  • n_std (float) -- The number of standard deviations of the data to highlight, to determine the ellipse’s radiuses. Defaults to 3.0.

  • facecolor (str) -- The facecolor of the ellipse. Defaults to “none”.

  • **kwargs -- Any additional keyword arguments are transmitted to Ellipse. If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns:

matplotlib.patches.Ellipse -- The corresponding Ellipse object added to the axis.

Example

x = np.random.normal(size=1000)
y = np.random.normal(size=1000)
plt.plot(x, y, ".", markersize=0.8)
draw_confidence_ellipse(x, y, n_std=2.5, edgecolor="red")
pyhdtoolkit.plotting.utils.draw_ip_locations(ip_positions: dict[str, float] | None = None, lines: bool = True, location: str = 'outside', **kwargs) → None[source]

Added in version 1.0.0.

Plots the interaction points’ locations into the background of your Axes.

Parameters:

  • ip_positions (dict[str, float]) -- A dict containing IP names as keys and their longitudinal positions as values, as returned by get_lhc_ips_positions.

  • lines (bool) -- If True, will also draw vertical lines at the IP positions. Defaults to True.

  • location (str) -- Where to show the IP names on the provided axis, either inside (will draw text at the bottom of the axis) or outside (will draw text on top of the axis). If None is given, then no labels are drawn. Defaults to outside.

  • **kwargs -- If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Example

twiss_df = tfs.read("twiss_output.tfs", index="NAME")
twiss_df.plot(x="S", y=["BETX", "BETY"])
ips = get_lhc_ips_positions(twiss_df)
draw_ip_locations(ip_positions=ips)
pyhdtoolkit.plotting.utils.find_ip_s_from_segment_start(segment_df: TfsDataFrame, model_df: TfsDataFrame, ip: int) → float[source]

Added in version 0.19.0.

Finds the S-offset of the IP from the start of segment by comparing the S-values for the elements in the model.

Parameters:

  • segment_df (tfs.TfsDataFrame) -- A TfsDataFrame of the segment-by-segment result for the given segment.

  • model_df (tfs.TfsDataFrame) -- The TfsDataframe of the model’s TWISS, usually the twiss_elements.dat file.

  • ip (int) -- The LHC IP number.

Returns:

float -- The S-offset of the IP from the BPM at the start of segment.

Example

ip_offset_in_segment = find_ip_s_from_segment_start(
    segment_df=sbsphaseext_IP1, model_df=twiss_elements, ip=1
)
pyhdtoolkit.plotting.utils.get_lhc_ips_positions(dataframe: DataFrame) → dict[str, float][source]

Added in version 1.0.0.

Returns a dict of LHC IPs and their positions from the provided dataframe.

Important

This function expects the IP names to be in the dataframe’s index, and cased as the longitudinal coordinate column: aka uppercase names (IP1, IP2, etc) and S column; or lowercase names (ip1, ip2, etc) and s column.

Parameters:

dataframe (pandas.DataFrame) -- A DataFrame containing at least the IP positions. A typical example is a TWISS call output.

Returns:

dict[str, float] -- A dict with IP names as keys and their longitudinal locations as values.

Example

twiss_df = tfs.read("twiss_output.tfs", index="NAME")
ips = get_lhc_ips_positions(twiss_df)
pyhdtoolkit.plotting.utils.make_elements_groups(madx: Madx, /, xoffset: float = 0, xlimits: tuple[float, float] | None = None) → dict[str, DataFrame][source]

Added in version 1.0.0.

Provided with an active cpymad instance after having ran a script, will returns different portions of the twiss table’s dataframe for different magnetic elements.

Parameters:

  • madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

  • xoffset (float) -- An offset applied to the S coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at \(s = 0\). Beware this offset is applied before applying the xlimits. Defaults to 0.

  • xlimits (tuple[float, float], optional) -- If given, will be used for the xlim (for the s coordinate), using the tuple passed.

Returns:

dict[str, DataFrame] -- A dict containing a pd.DataFrame for dipoles, focusing quadrupoles, defocusing quadrupoles, sextupoles and octupoles. The keys are quite self-explanatory.

Example

element_dfs = make_elements_groups(madx)
pyhdtoolkit.plotting.utils.make_survey_groups(madx: Madx, /) → dict[str, DataFrame][source]

Added in version 1.0.0.

Provided with an active cpymad instance after having ran a script, will returns different portions of the survey table’s dataframe for different magnetic elements.

Parameters:

madx (cpymad.madx.Madx) -- An instanciated Madx object. Positional only.

Returns:

dict[str, DataFrame] -- A dict containing a pd.DataFrame for dipoles, focusing quadrupoles, defocusing quadrupoles, sextupoles and octupoles. The keys are quite self-explanatory.

Example

survey_dfs = make_survey_groups(madx)
pyhdtoolkit.plotting.utils.maybe_get_ax(**kwargs)[source]

Added in version 1.0.0.

Convenience function to get the axis, regardless of whether or not it is provided to the plotting function itself. It used to be that the first argument of plotting functions in this package had to be the ‘axis’ object, but that’s no longer the case.

Parameters:

  • *args -- The arguments passed to the plotting function.

  • **kwargs -- The keyword arguments passed to the plotting function.

Returns:

tuple[matplotlib.axes.Axes, tuple, dict] -- The Axes object to plot on, as well as the args and kwargs (without the ‘ax’ argument if it initially was present). If no axis was provided, then it will be created with a call to matplotlib.pyplot.gca.

Example

This is to be called at the beginning of your plotting functions:

def my_plotting_function(*args, **kwargs):
    ax, kwargs = maybe_get_ax(**kwargs)
    # do stuff with ax
    ax.plot(*args, **kwargs)
pyhdtoolkit.plotting.utils.set_arrow_label(label: str, arrow_position: tuple[float, float], label_position: tuple[float, float], color: str = 'k', arrow_arc_rad: float = -0.2, fontsize: int = 20, **kwargs) → Annotation[source]

Added in version 0.6.0.

Adds on the provided matplotlib.axes.Axes a label box with text and an arrow from the box to a specified position. Original code from Guido Sterbini.

Parameters:

  • label (str) -- The label text to print on the axis.

  • arrow_position (tuple[float, float]) -- Where on the plot to point the tip of the arrow

  • label_position (tuple[float, float]) -- Where on the plot the text label (and thus start of the arrow) is.

  • color (str) -- Color parameter for your arrow and label. Defaults to “k”, which is black.

  • arrow_arc_rad (float) -- Angle value defining the upwards / downwards shape of and bending of the arrow. Defaults to -0.2.

  • fontsize (int) -- Text size in the box. Defaults to 20.

  • **kwargs -- Any additional keyword arguments are transmitted to annotate. If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns:

matplotlib.text.Annotation -- A matplotlib.text.Annotation of the created annotation.

Example

set_arrow_label(
    label="Your label",
    arrow_position=(1, 2),
    label_position=(1.1 * some_value, 0.75 * another_value),
    color="indianred",
    arrow_arc_rad=0.3,
    fontsize=25,
)