From ad234fe60364b2c044515718e1cdf9d672e70ea2 Mon Sep 17 00:00:00 2001 From: Felix Soubelet <19598248+fsoubelet@users.noreply.github.com> Date: Thu, 4 Mar 2021 20:55:15 +0000 Subject: [PATCH] make layout plot into its own function (#37) --- pyhdtoolkit/__init__.py | 2 +- pyhdtoolkit/cpymadtools/latwiss.py | 364 ++++++++++++++++++----------- pyproject.toml | 2 +- 3 files changed, 229 insertions(+), 139 deletions(-) diff --git a/pyhdtoolkit/__init__.py b/pyhdtoolkit/__init__.py index 395e31aa..b3abd03c 100644 --- a/pyhdtoolkit/__init__.py +++ b/pyhdtoolkit/__init__.py @@ -13,7 +13,7 @@ __title__ = "pyhdtoolkit" __description__ = "An all-in-one toolkit package to easy my Python work in my PhD." __url__ = "https://github.com/fsoubelet/PyhDToolkit" -__version__ = "0.8.3" +__version__ = "0.8.4" __author__ = "Felix Soubelet" __author_email__ = "felix.soubelet@cern.ch" __license__ = "MIT" diff --git a/pyhdtoolkit/cpymadtools/latwiss.py b/pyhdtoolkit/cpymadtools/latwiss.py index f1243417..92040a62 100644 --- a/pyhdtoolkit/cpymadtools/latwiss.py +++ b/pyhdtoolkit/cpymadtools/latwiss.py @@ -22,8 +22,214 @@ plt.rcParams.update(PLOT_PARAMS) +# ----- Plotters ----- # -# ----- Utilities ----- # + +def plot_latwiss( + madx: Madx, + title: str, + figsize: Tuple[int, int] = (18, 11), + savefig: str = None, + xoffset: float = 0, + xlimits: Tuple[float, float] = None, + plot_dipoles: bool = True, + plot_quadrupoles: bool = True, + plot_bpms: bool = False, + disp_ylim: Tuple[float, float] = (-10, 125), + beta_ylim: Tuple[float, float] = None, + k0l_lim: Tuple[float, float] = (-0.25, 0.25), + k1l_lim: Tuple[float, float] = (-0.08, 0.08), + k2l_lim: Tuple[float, float] = None, + **kwargs, +) -> matplotlib.figure.Figure: + """ + Provided with an active Cpymad class after having ran a script, will create a plot representing nicely + the lattice layout and the beta functions along with the horizontal dispertion function. This is very + heavily refactored code, inspired by code from Guido Sterbini. + + WARNING: This WILL FAIL if you have not included 'q' or 'Q' in your quadrupoles' names, and 'b' or 'B' + in your dipoles' names when defining your MAD-X sequence. + + Args: + madx (cpymad.madx.Madx): an instanciated cpymad Madx object. + title (str): title of your plot. + figsize (Tuple[int, int]): size of the figure, defaults to (16, 10). + savefig (str): will save the figure if this is not None, using the string value passed. + 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. + Beware this offset is applied before applying the `xlimits`. Offset defaults to 0 (no change). + 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_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. + Defaults to (-10, 125). + beta_ylim (Tuple[float, float]): vertical axis limits for the betatron function values. + Defaults to None, to be determined by matplotlib based on the provided beta values. + k0l_lim (Tuple[float, float]): vertical axis limits for the k0l values used for the + height of dipole patches. Defaults to (-0.25, 0.25). + k1l_lim (Tuple[float, float]): vertical axis limits for the k1l values used for the + height of quadrupole patches. Defaults to (-0.08, 0.08). + k2l_lim (Tuple[float, float]): if given, sextupole patches will be plotted on the layout subplot of + the figure, and the provided values act as vertical axis limits for the k2l values used for the + height of sextupole patches. + + Keyword Args: + Any keyword argument to be transmitted to `_plot_machine_layout`, later on to `plot_lattice_series` + and then `matplotlib.patches.Rectangle`, such as lw etc. + + 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. + + Returns: + The figure on which the plots are drawn. The underlying axes can be accessed with + 'fig.get_axes()'. Eventually saves the figure as a file. + """ + # pylint: disable=too-many-arguments + # Restrict the span of twiss_df to avoid plotting all elements then cropping when xlimits is given + logger.info("Plotting optics functions and machine layout") + logger.debug("Getting Twiss dataframe from cpymad") + twiss_df = madx.table.twiss.dframe().copy() + twiss_df.s = twiss_df.s - xoffset + twiss_df = twiss_df[twiss_df.s.between(xlimits[0], xlimits[1])] if xlimits else twiss_df + + # Create a subplot for the lattice patches (takes a third of figure) + figure = plt.figure(figsize=figsize) + quadrupole_patches_axis = plt.subplot2grid((3, 3), (0, 0), colspan=3, rowspan=1) + _plot_machine_layout( + madx, + quadrupole_patches_axis=quadrupole_patches_axis, + title=title, + xoffset=xoffset, + xlimits=xlimits, + plot_dipoles=plot_dipoles, + plot_quadrupoles=plot_quadrupoles, + plot_bpms=plot_bpms, + k0l_lim=k0l_lim, + k1l_lim=k1l_lim, + k2l_lim=k2l_lim, + **kwargs, + ) + + # Plotting beta functions on remaining two thirds of the figure + logger.debug("Setting up betatron functions subplot") + betatron_axis = plt.subplot2grid((3, 3), (1, 0), colspan=3, rowspan=2, sharex=quadrupole_patches_axis) + betatron_axis.plot(twiss_df.s, twiss_df.betx, label="$\\beta_x$", lw=1.5) + betatron_axis.plot(twiss_df.s, twiss_df.bety, label="$\\beta_y$", lw=1.5) + betatron_axis.legend(loc=2) + betatron_axis.set_ylabel("$\\beta$-functions [m]") + betatron_axis.set_xlabel("s [m]") + + logger.trace("Setting up dispersion functions subplot") + dispertion_axis = betatron_axis.twinx() + dispertion_axis.plot(twiss_df.s, twiss_df.dx, color="brown", label="$D_x$", lw=2) + dispertion_axis.plot(twiss_df.s, twiss_df.dy, ls="-.", color="sienna", label="$D_y$", lw=2) + dispertion_axis.legend(loc=1) + dispertion_axis.set_ylabel("Dispersions [m]", color="brown") + dispertion_axis.tick_params(axis="y", labelcolor="brown") + dispertion_axis.grid(False) + + if beta_ylim: + logger.debug("Setting ylim for betatron functions plot") + betatron_axis.set_ylim(beta_ylim) + + if disp_ylim: + logger.debug("Setting ylim for dispersion plot") + dispertion_axis.set_ylim(disp_ylim) + + if xlimits: + logger.debug("Setting xlim for longitudinal coordinate") + plt.xlim(xlimits) + + if savefig: + logger.info(f"Saving latwiss plot as {savefig}") + plt.savefig(savefig) + return figure + + +def plot_machine_survey( + madx: Madx, + title: str = "Machine Layout", + figsize: Tuple[int, int] = (16, 11), + savefig: str = None, + show_elements: bool = False, + high_orders: bool = False, +) -> matplotlib.figure.Figure: + """ + Provided with an active Cpymad class after having ran a script, will create a plot + representing the machine geometry in 2D. Original code is from Guido Sterbini. + + Args: + madx (cpymad.madx.Madx): an instanciated cpymad Madx object. + title (str): title of your plot. + figsize (Tuple[int, int]): size of the figure, defaults to (16, 10). + savefig (str): will save the figure if this is not None, using the string value passed. + show_elements (bool): if True, will try to plot by differentiating elements. + Experimental, defaults to False. + high_orders (bool): if True, plot sextupoles and octupoles when show_elements is True, + otherwise only up to quadrupoles. Defaults to False. + + Returns: + The figure on which the plots are drawn. The underlying axes can be accessed with + 'fig.get_axes()'. Eventually saves the figure as a file. + """ + logger.debug("Getting machine survey from cpymad") + madx.command.survey() + survey = madx.table.survey.dframe() + figure = plt.figure(figsize=figsize) + + if show_elements: + logger.debug("Plotting survey with elements differentiation") + element_dfs = _make_survey_groups(survey) + plt.scatter( + element_dfs["dipoles"].z, + element_dfs["dipoles"].x, + marker=".", + c=element_dfs["dipoles"].s, + cmap="copper", + label="Dipoles", + ) + plt.scatter( + element_dfs["quad_foc"].z, element_dfs["quad_foc"].x, marker="o", color="blue", label="QF", + ) + plt.scatter( + element_dfs["quad_defoc"].z, element_dfs["quad_defoc"].x, marker="o", color="red", label="QD", + ) + + if high_orders: + logger.debug("Plotting high order magnetic elements (up to octupoles)") + plt.scatter( + element_dfs["sextupoles"].z, element_dfs["sextupoles"].x, marker=".", color="m", label="MS", + ) + plt.scatter( + element_dfs["octupoles"].z, element_dfs["octupoles"].x, marker=".", color="cyan", label="MO", + ) + plt.legend(loc=2) + + else: + logger.debug("Plotting survey without elements differentiation") + plt.scatter(survey.z, survey.x, c=survey.s) + + plt.axis("equal") + plt.colorbar().set_label("s [m]") + plt.xlabel("z [m]") + plt.ylabel("x [m]") + plt.title(title) + + if savefig: + logger.info(f"Saving machine survey plot as {savefig}") + plt.savefig(savefig, format="pdf", dpi=500) + return figure + + +# ----- Utility plotters ----- # def _plot_lattice_series( @@ -63,35 +269,36 @@ def _plot_lattice_series( ) -def plot_latwiss( +def _plot_machine_layout( madx: Madx, + quadrupole_patches_axis: matplotlib.axes.Axes, title: str, - figsize: Tuple[int, int] = (18, 11), - savefig: str = None, + xoffset: float = 0, xlimits: Tuple[float, float] = None, plot_dipoles: bool = True, plot_quadrupoles: bool = True, plot_bpms: bool = False, - disp_ylim: Tuple[float, float] = (-10, 125), - beta_ylim: Tuple[float, float] = None, k0l_lim: Tuple[float, float] = (-0.25, 0.25), k1l_lim: Tuple[float, float] = (-0.08, 0.08), k2l_lim: Tuple[float, float] = None, **kwargs, -) -> matplotlib.figure.Figure: +) -> None: """ - Provided with an active Cpymad class after having ran a script, will create a plot representing nicely - the lattice layout and the beta functions along with the horizontal dispertion function. This is very - heavily refactored code, inspired by code from Guido Sterbini. + Provided with an active Cpymad class after having ran a script, will plot the lattice layout and the + on a given axis. This is the function that takes care of the machine layout in `plot_latwiss`, and + is in theory a private function, though if you know what you are doing you may use it individually. WARNING: This WILL FAIL if you have not included 'q' or 'Q' in your quadrupoles' names, and 'b' or 'B' in your dipoles' names when defining your MAD-X sequence. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. + quadrupole_patches_axis (matplotlib.axes.Axes): the axis on which to plot. Will also create the + appropriate new axes with `twinx()` to plot the element orders asked for. title (str): title of your plot. - figsize (Tuple[int, int]): size of the figure, defaults to (16, 10). - savefig (str): will save the figure if this is not None, using the string value passed. + 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. + Beware this offset is applied before applying the `xlimits`. Offset defaults to 0 (no change). 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 @@ -100,10 +307,6 @@ def plot_latwiss( 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. - Defaults to (-10, 125). - beta_ylim (Tuple[float, float]): vertical axis limits for the betatron function values. - Defaults to None, to be determined by matplotlib based on the provided beta values. k0l_lim (Tuple[float, float]): vertical axis limits for the k0l values used for the height of dipole patches. Defaults to (-0.25, 0.25). k1l_lim (Tuple[float, float]): vertical axis limits for the k1l values used for the @@ -118,26 +321,23 @@ def plot_latwiss( WARNING: Currently the function tries to plot legends for the different layout patches. The position of the - different legends has been hardcoded and can lead to messed-up layouts. User beware. - - Returns: - The figure on which the plots are drawn. The underlying axes can be accessed with - 'fig.get_axes()'. Eventually saves the figure as a file. + 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. """ # pylint: disable=too-many-arguments # Restrict the span of twiss_df to avoid plotting all elements then cropping when xlimits is given - logger.info("Plotting optics functions and machine layout") - logger.debug("Getting Twiss dataframe from cpymad") + logger.trace("Getting Twiss dataframe from cpymad") twiss_df = madx.table.twiss.dframe().copy() + twiss_df.s = twiss_df.s - xoffset twiss_df = twiss_df[twiss_df.s.between(xlimits[0], xlimits[1])] if xlimits else twiss_df - figure = plt.figure(figsize=figsize) - # Create a subplot for the lattice patches (takes a third of figure) - logger.trace("Setting up element patches subplots") - quadrupole_patches_axis = plt.subplot2grid((3, 3), (0, 0), colspan=3, rowspan=1) + logger.debug("Plotting machine layout") + logger.trace(f"Plotting from axis '{quadrupole_patches_axis}'") quadrupole_patches_axis.set_ylabel("1/f=K1L [m$^{-1}$]", color="red") # quadrupole in red quadrupole_patches_axis.tick_params(axis="y", labelcolor="red") quadrupole_patches_axis.set_ylim(k1l_lim) + quadrupole_patches_axis.set_xlim(xlimits) quadrupole_patches_axis.set_title(title) quadrupole_patches_axis.plot(twiss_df.s, 0 * twiss_df.s, "k") # 0-level line @@ -239,116 +439,6 @@ def plot_latwiss( plotted_elements += 1 bpm_patches_axis.legend(loc=4, fontsize=16) - # Plotting beta functions on remaining two thirds of the figure - logger.trace("Setting up betatron functions subplot") - betatron_axis = plt.subplot2grid((3, 3), (1, 0), colspan=3, rowspan=2, sharex=quadrupole_patches_axis) - betatron_axis.plot(twiss_df.s, twiss_df.betx, label="$\\beta_x$", lw=1.5) - betatron_axis.plot(twiss_df.s, twiss_df.bety, label="$\\beta_y$", lw=1.5) - betatron_axis.legend(loc=2) - betatron_axis.set_ylabel("$\\beta$-functions [m]") - betatron_axis.set_xlabel("s [m]") - - logger.trace("Setting up dispersion functions subplot") - dispertion_axis = betatron_axis.twinx() - dispertion_axis.plot(twiss_df.s, twiss_df.dx, color="brown", label="$D_x$", lw=2) - dispertion_axis.plot(twiss_df.s, twiss_df.dy, ls="-.", color="sienna", label="$D_y$", lw=2) - dispertion_axis.legend(loc=1) - dispertion_axis.set_ylabel("Dispersions [m]", color="brown") - dispertion_axis.tick_params(axis="y", labelcolor="brown") - dispertion_axis.grid(False) - - if beta_ylim: - logger.debug("Setting ylim for betatron functions plot") - betatron_axis.set_ylim(beta_ylim) - - if disp_ylim: - logger.debug("Setting ylim for dispersion plot") - dispertion_axis.set_ylim(disp_ylim) - - if xlimits: - logger.debug("Setting xlim for longitudinal coordinate") - plt.xlim(xlimits) - - if savefig: - logger.info(f"Saving latwiss plot as {savefig}") - plt.savefig(savefig) - return figure - - -def plot_machine_survey( - madx: Madx, - title: str = "Machine Layout", - figsize: Tuple[int, int] = (16, 11), - savefig: str = None, - show_elements: bool = False, - high_orders: bool = False, -) -> matplotlib.figure.Figure: - """ - Provided with an active Cpymad class after having ran a script, will create a plot - representing the machine geometry in 2D. Original code is from Guido Sterbini. - - Args: - madx (cpymad.madx.Madx): an instanciated cpymad Madx object. - title (str): title of your plot. - figsize (Tuple[int, int]): size of the figure, defaults to (16, 10). - savefig (str): will save the figure if this is not None, using the string value passed. - show_elements (bool): if True, will try to plot by differentiating elements. - Experimental, defaults to False. - high_orders (bool): if True, plot sextupoles and octupoles when show_elements is True, - otherwise only up to quadrupoles. Defaults to False. - - Returns: - The figure on which the plots are drawn. The underlying axes can be accessed with - 'fig.get_axes()'. Eventually saves the figure as a file. - """ - logger.debug("Getting machine survey from cpymad") - madx.command.survey() - survey = madx.table.survey.dframe() - figure = plt.figure(figsize=figsize) - - if show_elements: - logger.debug("Plotting survey with elements differentiation") - element_dfs = _make_survey_groups(survey) - plt.scatter( - element_dfs["dipoles"].z, - element_dfs["dipoles"].x, - marker=".", - c=element_dfs["dipoles"].s, - cmap="copper", - label="Dipoles", - ) - plt.scatter( - element_dfs["quad_foc"].z, element_dfs["quad_foc"].x, marker="o", color="blue", label="QF", - ) - plt.scatter( - element_dfs["quad_defoc"].z, element_dfs["quad_defoc"].x, marker="o", color="red", label="QD", - ) - - if high_orders: - logger.debug("Plotting high order magnetic elements (up to octupoles)") - plt.scatter( - element_dfs["sextupoles"].z, element_dfs["sextupoles"].x, marker=".", color="m", label="MS", - ) - plt.scatter( - element_dfs["octupoles"].z, element_dfs["octupoles"].x, marker=".", color="cyan", label="MO", - ) - plt.legend(loc=2) - - else: - logger.debug("Plotting survey without elements differentiation") - plt.scatter(survey.z, survey.x, c=survey.s) - - plt.axis("equal") - plt.colorbar().set_label("s [m]") - plt.xlabel("z [m]") - plt.ylabel("x [m]") - plt.title(title) - - if savefig: - logger.info(f"Saving machine survey plot as {savefig}") - plt.savefig(savefig, format="pdf", dpi=500) - return figure - # ----- Helpers ----- # diff --git a/pyproject.toml b/pyproject.toml index 73a67fe3..acd5f243 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pyhdtoolkit" -version = "0.8.3" +version = "0.8.4" description = "An all-in-one toolkit package to easy my Python work in my PhD." authors = ["Felix Soubelet "] license = "MIT"