""".. _plotting-aperture:Aperture Plotters-----------------Module with functions to create aperture plotsthrough a `~cpymad.madx.Madx` object."""from__future__importannotationsfromtypingimportTYPE_CHECKINGimportmatplotlib.pyplotaspltimportnumpyasnpimportpandasaspdfromloguruimportloggerfrompyhdtoolkit.plotting.layoutimportplot_machine_layoutfrompyhdtoolkit.plotting.utilsimportmaybe_get_axifTYPE_CHECKING:fromcpymad.madximportMadx
[docs]defplot_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:""" .. versionadded:: 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 :ref:`machine aperture <demo-accelerator-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 `~cpymad.madx.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 :math:`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 `~.plotting.utils.plot_machine_layout`, later on to `~.plotting.utils._plot_lattice_series`, and then `~matplotlib.patches.Rectangle`, such as ``lw`` etc. Example ------- .. code-block:: python 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", ) """# pylint: disable=too-many-argumentslogger.debug("Plotting aperture limits and machine layout")logger.debug("Getting Twiss dataframe from cpymad")madx.command.twiss(centre=True)twiss_df:pd.DataFrame=madx.table.twiss.dframe()aperture_df=pd.DataFrame.from_dict(dict(madx.table.aperture))# slicing -> issues with .dframe()# Restrict the span of twiss_df to avoid plotting all elements then cropping when xlimits is giventwiss_df.s=twiss_df.s-xoffsetaperture_df.s=aperture_df.s-xoffsetxlimits=(twiss_df.s.min(),twiss_df.s.max())ifxlimitsisNoneelsexlimitstwiss_df=twiss_df[twiss_df.s.between(*xlimits)]ifxlimitselsetwiss_dfaperture_df=aperture_df[aperture_df.s.between(*xlimits)]ifxlimitselseaperture_df# Create a subplot for the lattice patches (takes a third of figure)# figure = plt.gcf()quadrupole_patches_axis=plt.subplot2grid((3,3),(0,0),colspan=3,rowspan=1)plot_machine_layout(madx,axis=quadrupole_patches_axis,title=title,xoffset=xoffset,xlimits=xlimits,plot_dipoles=plot_dipoles,plot_dipole_k1=plot_dipole_k1,plot_quadrupoles=plot_quadrupoles,plot_bpms=plot_bpms,k0l_lim=k0l_lim,k1l_lim=k1l_lim,k2l_lim=k2l_lim,k3l_lim=k3l_lim,**kwargs,)# Plotting aperture values on remaining two thirds of the figurelogger.debug("Plotting aperture values")aperture_axis=plt.subplot2grid((3,3),(1,0),colspan=3,rowspan=2,sharex=quadrupole_patches_axis)aperture_axis.plot(aperture_df.s,aperture_df.n1,marker=".",ls="-",lw=0.8,color=color,label="Aperture Limits")aperture_axis.fill_between(aperture_df.s,aperture_df.n1,aperture_df.n1.max(),interpolate=True,color=color)aperture_axis.legend()aperture_axis.set_ylabel(r"$n_{1} \ [\sigma]$")aperture_axis.set_xlabel(r"$S \ [m]$")ifaperture_ylim:logger.debug("Setting ylim for aperture plot")aperture_axis.set_ylim(aperture_ylim)ifxlimits:logger.debug("Setting xlim for longitudinal coordinate")plt.xlim(xlimits)
[docs]defplot_physical_apertures(madx,/,plane:str,scale:float=1,xoffset:float=0,xlimits:tuple[float,float]|None=None,**kwargs,)->None:""" .. versionadded:: 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 :ref:`machine aperture <demo-accelerator-aperture>` example gallery. Original code from :user:`Elias Waagaard <ewaagaard>`. 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 `~cpymad.madx.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 :math:`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 -------- .. code-block:: python 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: .. code-block:: python 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() """# pylint: disable=too-many-argumentsifplane.lower()notin("x","y","horizontal","vertical"):logger.error(f"'plane' argument should be 'x', 'horizontal', 'y' or 'vertical' not '{plane}'")msg="Invalid 'plane' argument."raiseValueError(msg)logger.debug("Plotting real element apertures")axis,kwargs=maybe_get_ax(**kwargs)positions,apertures=_get_positions_and_real_apertures(madx,plane,xlimits=xlimits,**kwargs)logger.trace(f"Applying scale ({scale}) and offset ({xoffset})")positions=positions-xoffsetapertures=apertures*scalelogger.trace("Plotting apertures")# previously drawstyle="steps", but do not use if entry and exit aperture offset differsaxis.fill_between(positions,apertures,0.2*scale,color="black",alpha=0.4,label="_nolegend_")axis.fill_between(positions,-1*apertures,-0.2*scale,color="black",alpha=0.4,label="_nolegend_")axis.plot(positions,apertures,"k",label="_nolegend_")axis.plot(positions,-1*apertures,"k",label="_nolegend_")ifxlimits:logger.trace("Setting xlim for longitudinal coordinate")axis.set_xlim(xlimits)
# ----- Helpers ----- #def_get_positions_and_real_apertures(madx,/,plane:str,xoffset:float=0,xlimits:tuple[float,float]|None=None,**kwargs)->tuple[np.ndarray,np.ndarray]:""" .. versionadded:: 1.2.0 Determines the "real" physical apertures of elements in the sequence. This is done by extrapolating a data point at the beginning and the end of each element, with values based on the ``aper_1`` and the ``aper_2`` columns in the ``TWISS`` table. Original code from :user:`Elias Waagaard <ewaagaard>`. 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 `~cpymad.madx.Madx` object. Positional only. plane : str The physical plane to plot for, Accepted values are either `x`, `horizontal`, `y` or `vertical`. Case insensitive. 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 :math:`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. Returns ------- tuple[numpy.ndarray, numpy.ndarray] A `~numpy.ndarray` of the longitudinal positions for the data points, and another `~numpy.ndarray` with the physical aperture values at these positions. """logger.debug("Getting Twiss dframe from MAD-X")madx.command.select(flag="twiss",column=["aper_1","aper_2"])# make sure we to get these twotwiss_df=madx.twiss(**kwargs).dframe()madx.command.select(flag="twiss",clear=True)# clean uptwiss_df.s=twiss_df.s-xoffsetlogger.trace("Determining aperture column to use")plane_number=1ifplane.lower()in("x","horizontal")else2apercol=f"aper_{plane_number:d}"ifxlimitsisnotNone:twiss_df=twiss_df[twiss_df.s.between(*xlimits)]# Initiate arrays for new aperture and positions,# these need to be lists as we will insert elementsnew_aper=twiss_df[apercol].tolist()new_pos=twiss_df.s.tolist()indices=[]logger.trace("Finding non-zero aperture elements")foriinrange(len(twiss_df[apercol])-1,0,-1):iftwiss_df[apercol].iloc[i]!=0:new_aper.insert(i,twiss_df[apercol].iloc[i])indices.append(i)indices=list(reversed(indices))logger.trace("Extrapolating data at beginning of elements")# counter keeps track of exact position in new array with counterforcounter,jinenumerate(indices):new_pos.insert(j+counter,(twiss_df.s.iloc[j]-twiss_df.l.iloc[j]))# Replace all zeros with Nanapertures=np.array(new_aper)apertures[apertures==0]=np.nanpositions=np.array(new_pos)returnpositions,apertures