xcoll 0.5.12__py3-none-any.whl → 0.6.1__py3-none-any.whl

xcoll/plot.py

@@ -0,0 +1,109 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+import numpy as np
+from numbers import Number
+
+
+_NORMS = ["total", "coll_max", "max", "none", "raw"]
+
+
+def _plot_lossmap_base(lossmap: dict, *, norm="total", ax=None, xlim=None, ylim=None,
+                       legend=True, grid=True, energy=False, show=True, savefig=None):
+    import matplotlib.pyplot as plt
+    if norm in _NORMS:
+        norm = norm.lower()
+    elif not isinstance(norm, Number):
+        raise ValueError(f"Norm must be one of {_NORMS} or a number, not {norm}.")
+
+    coll_s = lossmap['collimator']['s']
+    coll_val = lossmap['collimator']['e'] if energy else lossmap['collimator']['n']
+    if lossmap['interpolation']:
+        cold_s = lossmap['aperture']['s_bins']
+        cold_val = lossmap['aperture']['e_bins'] if energy else lossmap['aperture']['n_bins']
+    else:
+        cold_s = lossmap['aperture']['s']
+        cold_val = lossmap['aperture']['e'] if energy else lossmap['aperture']['n']
+    warm_s = np.array([])  # Placeholder for warm losses, if needed
+    warm_val = np.array([])  # Placeholder for warm losses, if needed
+    if norm != "raw":
+        if lossmap['interpolation']:
+            cold_val = cold_val / lossmap['interpolation']
+            # warm_val = warm_val / lossmap['interpolation']
+        coll_val = coll_val / lossmap['collimator']['length']
+        if norm == "total":
+            scale = coll_val.sum() + cold_val.sum() + warm_val.sum()
+        elif norm == "coll_max":
+            scale = coll_val.max()
+        elif norm == "max":
+            scale = max(coll_val.max(), cold_val.max(), warm_val.max())
+        elif norm == "none":
+            scale = 1
+        else:
+            scale = norm
+        coll_val = coll_val / scale
+        cold_val = cold_val / scale
+        warm_val = warm_val / scale
+        if lossmap['interpolation']:
+            cold_val = cold_val / lossmap['aperture']['length_bins']
+
+    L = lossmap['machine_length']
+    xlim = xlim if xlim else [-0.01*L, 1.01*L]
+    ylim = ylim if ylim else [1.e-7, 1.e1]
+
+    font = {'family': 'serif', 'size': 17}
+    format_dict = {f"font.{prop}": font[prop] for prop in font}
+    with plt.rc_context(format_dict):
+        if ax is None:
+            _, ax = plt.subplots(figsize=(16, 4))
+
+        if grid:
+            ax.grid(axis="y", which="major", zorder=0)
+
+        bar_common_kwargs = dict(width = 0.8, lw = 1, bottom = 1.e-9)
+        ax.bar(coll_s, coll_val, color="k", edgecolor="k", label="Collimator", zorder=10, **bar_common_kwargs)
+        ax.bar(cold_s, cold_val, color="b", edgecolor="b", label="Cold",  zorder=11, **bar_common_kwargs)
+        # ax.bar(warm_s, warm_val, color="r", edgecolor="r", label="Warm",  zorder=12, **bar_common_kwargs)
+
+        ax.set_yscale("log")
+        ax.set_ylim(ylim)
+        ax.set_xlim(xlim)
+
+        ax.set_yticks([10**i for i in range(int(np.log10(ylim[0])), 1 + int(np.log10(ylim[1])))])
+        ax.set_xlabel("s [m]")
+        ax.set_ylabel(_label_from_norm(norm, energy))
+
+        if legend:
+            ax.legend(loc='best', fancybox=True, framealpha=0.8, fontsize="small")
+
+        plt.tight_layout()
+        if savefig:
+            plt.savefig(savefig, dpi=300, bbox_inches='tight')
+        if show:
+            plt.show()
+        else:
+            plt.close()
+        return ax.figure, ax
+
+
+def _label_from_norm(norm, energy):
+    if norm == "none":
+        if energy:
+            return "Deposited energy [eV/m]"
+        else:
+            return "Particles absorbed [1/m]"
+    elif norm == "raw":
+        if energy:
+            return "Deposited energy [eV]"
+        else:
+            return "Particles absorbed [-]"
+    elif norm == "total":
+        return "Norm. inefficiency [1/m]"
+    elif norm == "max":
+        return "Inefficiency (norm. by max) [-]"
+    elif norm == "coll_max":
+        return "Inefficiency (norm. by coll max) [-]"
+    else:
+        return f"Inefficiency (norm. by {norm}) [1/m]"

xcoll/rf_sweep.py

@@ -16,6 +16,7 @@ class RFSweep:
         self.line = line
         self._get_base_frequency()
         self._install_zeta_shift()
+        self.reset()
 
     def _get_base_frequency(self):
         self.cavities = self.line.get_elements_of_type(xt.Cavity)[1]
@@ -179,6 +180,11 @@ class RFSweep:
             else:
                 self.line.tracker.time_last_track = None
 
+
+    def reset(self):
+        self.line['rf_sweep'].dzeta = 0
+
+
     def _tracking_func(self, particles, rf_shift_per_turn, num_turns=1, *args, **kwargs):
         existing_sweep = self.current_sweep_value
         for i in range(num_turns):