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]: New 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: cpymad.madx.Madx, /, title: Optional[str] = None, xoffset: float = 0, xlimits: Tuple[float, float] = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, aperture_ylim: Tuple[float, float] = None, k0l_lim: Union[Tuple[float, float], float, int] = None, k1l_lim: Union[Tuple[float, float], float, int] = None, k2l_lim: Union[Tuple[float, float], float, int] = None, k3l_lim: Union[Tuple[float, float], float, int] = None, color: str = None, **kwargs) → None[source]

New 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 (Optional[str]) -- 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 \(s = 0\). Beware this offset is applied before applying the xlimits. Defaults to 0.
xlimits (Tuple[float, float]) -- will implement xlim (for the s coordinate) if this is not None, 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.
aperture_ylim (Tuple[float, float]) -- vertical axis limits for the aperture values. Defaults to None, to be determined by matplotlib based on the provided values.
k0l_lim (Union[Tuple[float, float], float, int]) -- 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 (default) is given, then the limits will be auto-determined based on the k0l values of the dipoles in the plot.
k1l_lim (Union[Tuple[float, float], float, int]) -- 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 (default) is given, then the limits will be auto-determined based on the k0l values of the quadrupoles in the plot.
k2l_lim (Union[Tuple[float, float], float, int]) -- if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts 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).
k3l_lim (Union[Tuple[float, float], float, int]) -- if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts 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).
color (str) -- the color argument given to the aperture lines. Defaults to None, in which case the first color 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, **kwargs) → None[source]

New 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, should be either x, horizontal, y or vertical, and is 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 (values 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]) -- will implement xlim (for the s coordinate) if this is not None, using the tuple passed. Defaults to None.
**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: matplotlib.axes._axes.Axes, plot_df_b1: pandas.core.frame.DataFrame, plot_df_b2: pandas.core.frame.DataFrame, plot_column: str, scaling: float = 1, ylabel: str = None, xlabel: str = None, title: str = None) → None[source]

New 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 (Union[pd.DataFrame, tfs.TfsDataFrame]) -- TWISS dataframe of the IR zone for beam 1 of the LHC, centered on 0 at IP position (simply done with df.s = df.s - ip_s).
plot_df_b2 (Union[pd.DataFrame, tfs.TfsDataFrame]) -- TWISS dataframe of the IR zone for beam 2 of the LHC, centered on 0 at IP position (simply done 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) -- if given, will be used for the xlabel of the axis. Defaults to None.
ylabel (str) -- if given, will be used for the ylabel of the axis. Defaults to None.
title (str) -- if given, will be used for the title of the axis. Defaults to None.

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: cpymad.madx.Madx, /, first_ip: int, second_ip: int, ir_limit: float = 275, highlight_mqx_and_mbx: bool = True) → None[source]

New 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 calling 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: cpymad.madx.Madx, /, sequence: str, plane: str, nsigma: float = 1, scale: float = 1, xoffset: float = 0, xlimits: Tuple[float, float] = None, **kwargs) → None[source]

New 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, should be the active sequence. Case-insensitive.
plane (str) -- the physical plane to plot for, should be either x, horizontal, y or vertical, and is case-insensitive.
nsigma (float) -- the standard deviation to use for the beam enveloppe calculation. 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 (values 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]) -- will implement xlim (for the s coordinate) if this is not None, using the tuple passed. Defaults to None.
**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: cpymad.madx.Madx, /, title: Optional[str] = None, xoffset: float = 0, xlimits: Tuple[float, float] = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, disp_ylim: Union[Tuple[float, float], float, int] = None, beta_ylim: Union[Tuple[float, float], float, int] = None, k0l_lim: Union[Tuple[float, float], float, int] = None, k1l_lim: Union[Tuple[float, float], float, int] = None, k2l_lim: Union[Tuple[float, float], float, int] = None, k3l_lim: Union[Tuple[float, float], float, int] = None, **kwargs) → None[source]

New 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 and 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 (Optional[str]) -- if provided, is set as title of the plot. Defaults to None.
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]) -- will implement xlim (for the s coordinate) if this is not None, 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.
disp_ylim (Tuple[float, float]) -- 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 (-10, 125).
beta_ylim (Tuple[float, float]) -- 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, to be determined by matplotlib based on the plotted beta values.
k0l_lim (Union[Tuple[float, float], float, int]) -- 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 (default) is given, then the limits will be auto-determined based on the k0l values of the dipoles in the plot.
k1l_lim (Union[Tuple[float, float], float, int]) -- 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 (default) is given, then the limits will be auto-determined based on the k0l values of the quadrupoles in the plot.
k2l_lim (Union[Tuple[float, float], float, int]) -- if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts 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).
k3l_lim (Union[Tuple[float, float], float, int]) -- if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts 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).
**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

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: cpymad.madx.Madx, /, title: str = None, show_elements: bool = False, high_orders: bool = False, **kwargs) → matplotlib.axes._axes.Axes[source]

New 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 (Optional[str]) -- if provided, is set as title of the plot. Defaults to None.
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 will be transmitted to scatter calls later on. If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns

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’ elements in an Axes object, mostly used in different plotting modules.

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

New 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 you 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 (Optional[str]) -- if provided, is set as title of the plot. Defaults to None.
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 the \(s = 0\) position. Beware this offset is applied before applying the xlimits. Defaults to 0.
xlimits (Tuple[float, float]) -- will implement xlim (for the s coordinate) if this is not None, 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.
k0l_lim (Union[Tuple[float, float], float, int]) -- 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 (default) is given, then the limits will be auto-determined based on the k0l values of the dipoles in the plot.
k1l_lim (Union[Tuple[float, float], float, int]) -- 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 (default) is given, then the limits will be auto-determined based on the k0l values of the quadrupoles in the plot.
k2l_lim (Union[Tuple[float, float], float, int]) -- if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts 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).
k3l_lim (Union[Tuple[float, float], float, int]) -- if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts 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).
**kwargs -- any keyword argument will be transmitted to plot_machine_layout, later on 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]

New 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: cpymad.madx.Madx, /, u_coordinates: numpy.ndarray, pu_coordinates: numpy.ndarray, plane: str = 'Horizontal', title: str = None, **kwargs) → matplotlib.axes._axes.Axes[source]

New 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 (np.ndarray) -- 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 (np.ndarray) -- 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, and is case-insensitive. Defaults to Horizontal.
title (Optional[str]) -- if provided, is set as title of the plot. Defaults to None.
**kwargs -- If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns

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: cpymad.madx.Madx, /, u_coordinates: numpy.ndarray, pu_coordinates: numpy.ndarray, plane: str = 'Horizontal', title: str = None, **kwargs) → matplotlib.figure.Figure[source]

New 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 particle trajectory has its own color on the plot, within the limit of pyplot’s 156 named colors, after the function loops back to the first color 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 (np.ndarray) -- 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 (np.ndarray) -- 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.
savefig (str) -- if not None, will save the figure to file using the string value passed.
plane (str) -- the physical plane to plot, should be either Horizontal or Vertical, and is case-insensitive. Defaults to Horizontal.
title (Optional[str]) -- if provided, is set as title of the plot. Defaults to None.
**kwargs -- If either ax or axis is found in the kwargs, the corresponding value is used as the axis object to plot on.

Returns

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]

New 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: The sequence as a list of plottable 2D points.

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

New 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, legend_title: str = None, max_order: int = 6, differentiate_orders: bool = False, **kwargs) → matplotlib.axes._axes.Axes[source]

New 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 (Optional[str]) -- if provided, is set as title of the plot. Defaults to None.
legend_title (str) -- if given, will be used as the title of the plot’s legend. If set to None, then creating a legend for the figure will not be done by this function and left up to the user’s care (a call to legend will do). Defaults to None.
max_order (int) -- the order up to which to plot resonance lines for, 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 minimal differentation through alpha, linewidth and linestyle. Defaults to False.
**kwargs -- keyword arguments will be transmitted to the plot_resonance_lines_for_order function and later on 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

The Axes on which the tune diagram is drawn.

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.frame.TfsDataFrame, b2_segment_df: tfs.frame.TfsDataFrame, b1_model: tfs.frame.TfsDataFrame = None, b2_model: tfs.frame.TfsDataFrame = None, ip: int = None, rdt: str = 'F1001', abs_ylimits: Tuple[float, float] = None, real_ylimits: Tuple[float, float] = None, imag_ylimits: Tuple[float, float] = None, **kwargs) → matplotlib.figure.Figure[source]

New 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) -- A TfsDataFrame of the Beam 1 model used in the analysis, optional. If given, then the IP location in the segment will be highlighted by a vertical grey line.
b2_model (tfs.TfsDataFrame) -- A TfsDataFrame of the Beam 2 model used in the analysis, optional. If given, then the IP location in the segment will be highlighted by a vertical grey line.
ip (int) -- The IP number of the segment.
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

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.frame.TfsDataFrame, b2_segment_df: tfs.frame.TfsDataFrame, b1_model: tfs.frame.TfsDataFrame = None, b2_model: tfs.frame.TfsDataFrame = None, ip: int = None, rdt: str = 'F1001', component: str = 'ABS', **kwargs) → matplotlib.figure.Figure[source]

New 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

ax (matplotlib.axes.Axes) -- The Axes to plot on. Will get the current axis if no Axes is given.
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) -- A TfsDataFrame of the Beam 1 model used in the analysis, optional. If given, then the IP location in the segment will be highlighted by a vertical grey line.
b2_model (tfs.TfsDataFrame) -- A TfsDataFrame of the Beam 2 model used in the analysis, optional. If given, then the IP location in the segment will be highlighted by a vertical grey line.
ip (int) -- The IP number of the segment.
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

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: matplotlib.axes._axes.Axes = None, segment_df: tfs.frame.TfsDataFrame = None, model_df: tfs.frame.TfsDataFrame = None, plane: str = 'x', ip: int = None) → None[source]

New 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) -- 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) -- A TfsDataFrame of the model used in the analysis, optional. 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) -- The IP number of the segment.

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.frame.TfsDataFrame, b1_phase_y: tfs.frame.TfsDataFrame, b2_phase_x: tfs.frame.TfsDataFrame, b2_phase_y: tfs.frame.TfsDataFrame, b1_model: tfs.frame.TfsDataFrame = None, b2_model: tfs.frame.TfsDataFrame = None, ip: int = None, **kwargs) → matplotlib.figure.Figure[source]

New 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) -- A TfsDataFrame of the Beam 1 model used in the analysis, optional. If given, then the IP location in the segment will be highlighted by a vertical grey line.
b2_model (tfs.TfsDataFrame) -- A TfsDataFrame of the Beam 2 model used in the analysis, optional. If given, then the IP location in the segment will be highlighted by a vertical grey line.
ip (int) -- The IP number of the segment.
**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

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.frame.TfsDataFrame, phase_y: tfs.frame.TfsDataFrame, model: tfs.frame.TfsDataFrame = None, ip: int = None, **kwargs) → matplotlib.figure.Figure[source]

New 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) -- A TfsDataFrame of the model used in the analysis, optional. If given, then the IP location in the segment will be highlighted by a vertical grey line.
ip (int) -- The IP number of the segment.
**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

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) → matplotlib.patches.Ellipse[source]

New in version 1.2.0.

Plot the covariance confidence ellipse of x and y. This code is taken from 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.
facecolor (str) -- The facecolor of the ellipse.
**kwargs -- Any keyword argument will be forwarded 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

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, lines: bool = True, location: str = 'outside', **kwargs) → None[source]

New in version 1.0.0.

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

Parameters

ip_positions (dict) -- a dict containing IP names as keys and their longitudinal positions as values, as returned by get_lhc_ips_positions.
lines (bool) -- whether to also draw vertical lines at the IP positions. Defaults to True.
location -- 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: tfs.frame.TfsDataFrame, model_df: tfs.frame.TfsDataFrame, ip: int) → float[source]

New 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 twiss_elements.dat.
ip (int) -- The LHC IP number.

Returns

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: pandas.core.frame.DataFrame) → dict[str, float][source]

New 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 IP positions. A typical example is a TWISS call output.
Returns: 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: cpymad.madx.Madx, /, xoffset: float = 0, xlimits: tuple[float, float] = None) → dict[str, pandas.core.frame.DataFrame][source]

New 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 is you want to center a plot around a specific point or element, which would then become located at s = 0.
xlimits (tuple[float, float]) -- will only consider elements within xlim (for the s coordinate) if this is not None, using the tuple passed.

Returns

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

Example

element_dfs = make_elements_groups(madx)

pyhdtoolkit.plotting.utils.make_survey_groups(madx: cpymad.madx.Madx, /) → dict[str, pandas.core.frame.DataFrame][source]

New 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: A dict containing a pd.DataFrame for dipoles, focusing quadrupoles, defocusing quadrupoles, sextupoles and octupoles. The keys are self-explanatory.

Example

survey_dfs = make_survey_groups(madx)

pyhdtoolkit.plotting.utils.maybe_get_ax(**kwargs)[source]

New 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

The Axes object to plot on, the args and the kwargs (without the ‘ax’ argument if it initially was present). If no axis was provided, then it will be created with a call to gca.

Examples

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) → matplotlib.text.Annotation[source]

New 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

axis (matplotlib.axes.Axes) -- a matplotlib.axes.Axes to plot on.
label (str) -- 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”.
arrow_arc_rad (float) -- angle value defining the upwards / downwards shape of and bending of the arrow.
fontsize (int) -- text size in the box.
**kwargs -- 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

A matploblit.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,
)