diff --git a/.gitignore b/.gitignore index 4bd2411b..8545b3b7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +.DS_Store +.idea build PyORBIT.egg-info .vscode diff --git a/examples/Envelope/plot.py b/examples/Envelope/plot.py new file mode 100644 index 00000000..f493b108 --- /dev/null +++ b/examples/Envelope/plot.py @@ -0,0 +1,129 @@ +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches as patches + + +def calc_rms_ellipse_params(cov_matrix: np.ndarray) -> tuple[float, float, float]: + """Return rms ellipse dimensions and orientation.""" + i, j = (0, 1) + + sii = cov_matrix[i, i] + sjj = cov_matrix[j, j] + sij = cov_matrix[i, j] + + angle = -0.5 * np.arctan2(2.0 * sij, sii - sjj) + + _sin = np.sin(angle) + _cos = np.cos(angle) + _sin2 = _sin**2 + _cos2 = _cos**2 + + c1 = np.sqrt(abs(sii * _cos2 + sjj * _sin2 - 2 * sij * _sin * _cos)) + c2 = np.sqrt(abs(sii * _sin2 + sjj * _cos2 + 2 * sij * _sin * _cos)) + + return (c1, c2, angle) + + +def plot_ellipse( + r1: float = 1.0, + r2: float = 1.0, + angle: float = 0.0, + center: tuple[float, float] = None, + ax=None, + **kws, +): + kws.setdefault("fill", False) + kws.setdefault("color", "black") + kws.setdefault("lw", 1.25) + + if center is None: + center = (0.0, 0.0) + + d1 = r1 * 2.0 + d2 = r2 * 2.0 + angle = -np.degrees(angle) + + ax.add_patch(patches.Ellipse(center, d1, d2, angle=angle, **kws)) + return ax + + +def plot_rms_ellipse( + cov_matrix: np.ndarray, + level: float = 1.0, + ax=None, + **ellipse_kws, +): + """Plot rms ellipse from 2 x 2 covariance matrix.""" + r1, r2, angle = calc_rms_ellipse_params(cov_matrix) + plot_ellipse(r1 * level, r2 * level, angle=angle, ax=ax, **ellipse_kws) + return ax + + +def plot_corner( + particles: np.ndarray, + limits: list[tuple[float, float]] = None, + bins: int = 64, + labels: list[str] = None, + blur: float = None, +) -> tuple: + """Generate corner plot.""" + ndim = particles.shape[1] + + if limits is None: + xmax = np.max(particles, axis=0) + xmin = np.min(particles, axis=0) + limits = list(zip(xmin, xmax)) + + if labels is None: + labels = ndim * [""] + + fig, axs = plt.subplots(ncols=ndim, nrows=ndim, sharex=None, sharey=None, figsize=(8, 8)) + for i in range(ndim): + for j in range(ndim): + axis = (j, i) + ax = axs[i, j] + if i > j: + values, edges = np.histogramdd( + particles[:, axis], bins=bins, range=[limits[k] for k in axis] + ) + if blur: + values = scipy.ndimage.gaussian_filter(values, sigma=blur) + ax.pcolormesh( + edges[0], + edges[1], + values.T, + linewidth=0.0, + rasterized=True, + shading="auto", + ) + elif i == j: + values, edges = np.histogram(particles[:, i], bins=bins, range=limits[i]) + if blur: + values = scipy.ndimage.gaussian_filter(values, sigma=blur) + ax.stairs(values, edges, lw=1.5, color="black") + else: + ax.axis("off") + + for i in range(0, ndim - 1): + for j in range(0, ndim): + axs[i, j].set_xticklabels([]) + for i in range(0, ndim): + for j in range(1, ndim): + axs[i, j].set_yticklabels([]) + + for ax in axs.flat: + for loc in ["top", "right"]: + ax.spines[loc].set_visible(False) + + for i, label in enumerate(labels): + axs[-1, i].set_xlabel(label) + for i, label in enumerate(labels[1:], start=1): + axs[i, 0].set_ylabel(label) + + axs[0, 0].set_yticklabels([]) + axs[0, 0].set_ylabel(None) + + fig.align_ylabels() + fig.align_xlabels() + + return fig, axs diff --git a/examples/Envelope/sns_linac/diagnostics.py b/examples/Envelope/sns_linac/diagnostics.py new file mode 100644 index 00000000..7edfebdd --- /dev/null +++ b/examples/Envelope/sns_linac/diagnostics.py @@ -0,0 +1,56 @@ +import numpy as np + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis + + +class BunchMonitor: + def __init__(self) -> None: + self.twiss_calc = BunchTwissAnalysis() + self.position_start = 0.0 + + self.history = {} + self.history["position"] = [] + self.history["rms_x"] = [] + self.history["rms_y"] = [] + self.history["rms_z"] = [] + self.history["kin_energy"] = [] + + def __call__(self, params_dict: dict) -> None: + bunch = params_dict["bunch"] + node = params_dict["node"] + position = params_dict["path_length"] + + if params_dict["old_pos"] == position: + return + if params_dict["old_pos"] + params_dict["pos_step"] > position: + return + params_dict["old_pos"] = position + params_dict["count"] += 1 + + sync_part = bunch.getSyncParticle() + + self.twiss_calc.analyzeBunch(bunch) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(6): + cov_matrix[i, j] = cov_matrix[j, i] = self.twiss_calc.getCorrelation(i, j) + + xrms = 1000.0 * np.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * np.sqrt(cov_matrix[2, 2]) + zrms = 1000.0 * np.sqrt(cov_matrix[4, 4]) + + message = "" + message += " s={:0.3f}".format(position + self.position_start) + message += " xrms={:0.3f}".format(xrms) + message += " yrms={:0.3f}".format(yrms) + message += " zrms={:0.3f}".format(zrms) + message += " node={}".format(node.getName()) + print(message) + + self.history["position"].append(position + self.position_start) + self.history["rms_x"].append(xrms) + self.history["rms_y"].append(yrms) + self.history["rms_z"].append(zrms) + self.history["kin_energy"].append(sync_part.kinEnergy()) diff --git a/examples/Envelope/sns_linac/inputs/sns_linac.xml b/examples/Envelope/sns_linac/inputs/sns_linac.xml new file mode 100644 index 00000000..62cdbf41 --- /dev/null +++ b/examples/Envelope/sns_linac/inputs/sns_linac.xml @@ -0,0 +1,13475 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/examples/Envelope/sns_linac/style.mplstyle b/examples/Envelope/sns_linac/style.mplstyle new file mode 100644 index 00000000..71f36605 --- /dev/null +++ b/examples/Envelope/sns_linac/style.mplstyle @@ -0,0 +1,8 @@ +axes.linewidth: 1.25 +axes.titlesize: "medium" +image.cmap: "Greys" +figure.constrained_layout.use: True +savefig.dpi: 300 +savefig.format: "png" +xtick.minor.visible: True +ytick.minor.visible: True diff --git a/examples/Envelope/sns_linac/test_sns_linac.py b/examples/Envelope/sns_linac/test_sns_linac.py new file mode 100755 index 00000000..045e6155 --- /dev/null +++ b/examples/Envelope/sns_linac/test_sns_linac.py @@ -0,0 +1,297 @@ +"""SNS linac benchmark. + +Note that there is no analytic benchmark with 3D space charge since there +is no 3D KV equilibrium distribution. For comparison to particle tracking, +we use the `SpaceChargeCalcUnifEllipse` space charge calculator, which +approximates the charge distribution as a uniform-density ellipsoid with +the same x-y-z covariance matrix as the real charge distribution. (It +currently assumes an upright ellipsoid.) +""" + +import argparse +import math +import os +import random +import sys + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.core.bunch import SyncParticle +from orbit.core.linac import BaseRfGap +from orbit.core.linac import MatrixRfGap +from orbit.core.spacecharge import SpaceChargeCalcUnifEllipse +from orbit.core.spacecharge import SpaceChargeCalc3D +from orbit.bunch_generators import TwissContainer +from orbit.bunch_generators import WaterBagDist3D +from orbit.bunch_generators import GaussDist3D +from orbit.bunch_generators import KVDist3D +from orbit.bunch_utils import collect_bunch +from orbit.envelope import Envelope +from orbit.envelope import EnvelopeTracker +from orbit.lattice import AccLattice +from orbit.lattice import AccNode +from orbit.lattice import AccActionsContainer +from orbit.py_linac.linac_parsers import SNS_LinacLatticeFactory +from orbit.py_linac.lattice import LinacAccLattice +from orbit.space_charge.sc3d import setSC3DAccNodes +from orbit.space_charge.sc3d import setUniformEllipsesSCAccNodes +from orbit.utils.consts import mass_proton +from orbit.utils.consts import mass_electron +from orbit.utils.consts import charge_electron + +# local +from diagnostics import BunchMonitor + +sys.path.append("..") +from plot import plot_corner +from plot import plot_rms_ellipse +from utils import project_cov_matrix + +plt.style.use("style.mplstyle") + + +# Parse arguments +# -------------------------------------------------------------------------------- + +parser = argparse.ArgumentParser() +parser.add_argument("--dist", type=str, default="kv") +parser.add_argument("--nparts", type=int, default=20_000) +parser.add_argument("--sc", type=int, default=0) +parser.add_argument("--sc-model", type=str, default="ellipsoid") +parser.add_argument("--sc-path-length-min", type=float, default=0.01) +parser.add_argument("--current", type=float, default=0.038) +parser.add_argument("--show", type=int, default=0) +parser.add_argument("--seq-stop", type=str, default="CCL2") +args = parser.parse_args() + + +# Setup +# -------------------------------------------------------------------------------- + +output_dir = "outputs" +os.makedirs(output_dir, exist_ok=True) + +random.seed(23) + + +# Bunch +# -------------------------------------------------------------------------------- + +kin_energy = 0.0025 # [GeV] +mass = mass_proton + 2.0 * mass_electron +frequency = 402.5e06 +charge = -1.0 +intensity = args.current / frequency / (math.fabs(charge) * charge_electron) + +bunch = Bunch() +bunch.mass(mass) +bunch.macroSize(intensity / args.nparts) +bunch.charge(charge) + +sync_part = bunch.getSyncParticle() +sync_part.kinEnergy(kin_energy) +sync_part.time(0.0) + +alpha_x, beta_x, eps_x = (-1.962, 0.183, 2.874e-06) +alpha_y, beta_y, eps_y = (+1.768, 0.162, 2.874e-06) +alpha_z, beta_z, eps_z = (-0.0196, 116.414, 1.651e-08) + +twiss_x = TwissContainer(alpha_x, beta_x, eps_x) +twiss_y = TwissContainer(alpha_y, beta_y, eps_y) +twiss_z = TwissContainer(alpha_z, beta_z, eps_z) + +if args.dist == "waterbag": + dist = WaterBagDist3D(twiss_x, twiss_y, twiss_z) +elif args.dist == "kv": + dist = KVDist3D(twiss_x, twiss_y, twiss_z) +elif args.dist == "gauss": + dist = GaussDist3D(twiss_x, twiss_y, twiss_z) +else: + raise ValueError("Unknown distribution '{}'".format(args.dist)) + +for _ in range(args.nparts): + bunch.addParticle(*dist.getCoordinates()) + +# Lattice +# -------------------------------------------------------------------------------- + +seq_names = [ + "MEBT", + "DTL1", + "DTL2", + "DTL3", + "DTL4", + "DTL5", + "DTL6", + "CCL1", + "CCL2", + "CCL3", + "CCL4", + "SCLMed", + "SCLHigh", + "HEBT1", + "HEBT2", +] +if args.seq_stop: + index = seq_names.index(args.seq_stop) + 1 + seq_names = seq_names[:index] + +sns_linac_factory = SNS_LinacLatticeFactory() +sns_linac_factory.setMaxDriftLength(args.sc_path_length_min) +lattice = sns_linac_factory.getLinacAccLattice(seq_names, "inputs/sns_linac.xml") + +for node in lattice.getNodes(): + try: + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + except: + pass + +rf_gaps = lattice.getRF_Gaps() +for rf_gap in rf_gaps: + rf_gap.setCppGapModel(MatrixRfGap()) + +for index, node in enumerate(lattice.getNodes()): + print(index, type(node), node.getName()) + +lattice.trackDesignBunch(bunch) + + +# Track envelope +# -------------------------------------------------------------------------------- + +twiss_calc = BunchTwissAnalysis() +twiss_calc.analyzeBunch(bunch) + +cov_matrix = np.zeros((6, 6)) +for i in range(6): + for j in range(6): + cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) + +envelope = Envelope(bunch=bunch, cov_matrix=cov_matrix, intensity=intensity) + +tracker = EnvelopeTracker(lattice, sc=("3d" if args.sc else None)) + +histories = {} +histories["envelope"] = tracker.track_history(envelope) + + +# Track bunch +# -------------------------------------------------------------------------------- + +lattice.trackDesignBunch(bunch) + +if args.sc: + if args.sc_model == "ellipsoid": + n_ellipsoids = 1 + sc_calc = SpaceChargeCalcUnifEllipse(n_ellipsoids) + sc_nodes = setUniformEllipsesSCAccNodes(lattice, args.sc_path_length_min, sc_calc) + if args.sc_model == "3d": + sc_calc = SpaceChargeCalc3D(64, 64, 64) + sc_nodes = setSC3DAccNodes(lattice, args.sc_path_length_min, sc_calc) + + +monitor = BunchMonitor() + +action_container = AccActionsContainer() +action_container.addAction(monitor, AccActionsContainer.ENTRANCE) +action_container.addAction(monitor, AccActionsContainer.EXIT) + +params_dict = {"old_pos": -1.0, "count": 0, "pos_step": args.sc_path_length_min} + +lattice.trackBunch(bunch, paramsDict=params_dict, actionContainer=action_container) + +histories["bunch"] = monitor.history + + +# Analysis +# -------------------------------------------------------------------------------- + + +# History: rms +for mode in histories: + for key in histories[mode]: + histories[mode][key] = np.array(histories[mode][key]) + +plot_kws = {} +plot_kws["bunch"] = dict(color="black", lw=0, marker=".", ms=2) +plot_kws["envelope"] = dict(color="red", lw=0, marker=".", ms=1) + +fig, axs = plt.subplots(nrows=3, figsize=(10, 5), sharex=True, constrained_layout=True) +for mode in ["bunch", "envelope"]: + history = histories[mode] + for ax, key in zip(axs, ["rms_x", "rms_y", "rms_z"]): + ax.plot(history["position"], history[key], **plot_kws[mode], label=mode) +for ax in axs: + ax.legend(loc="lower right") +axs[0].set_ylabel("x rms [mm]") +axs[1].set_ylabel("y rms [mm]") +axs[2].set_ylabel("z rms [mm]") +axs[2].set_xlabel("s [m]") +plt.savefig(os.path.join(output_dir, "fig_history_rms.png")) +if args.show: + plt.show() +plt.close() + +# History: energy +fig, ax = plt.subplots(figsize=(5, 3)) +for mode in ["bunch", "envelope"]: + history = histories[mode] + ax.plot(history["position"], history["kin_energy"], **plot_kws[mode], label=mode) +ax.legend(loc="lower right") +ax.set_ylabel("energy [GeV]") +ax.set_xlabel("s [m]") +plt.savefig(os.path.join(output_dir, "fig_history_energy.png")) +plt.close() + +# Collect bunch/envelope data +particles = collect_bunch(bunch)["coords"] +particles *= 1e3 + +env_cov_matrix = envelope.cov_matrix +env_cov_matrix *= 1e6 + +env_centroid = envelope.centroid +env_centroid *= 1e3 + +xmax = 4.0 * np.std(particles, axis=0) +limits = list(zip(-xmax, xmax)) +labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]", "z [mm]", "dE [MeV]"] + +# Plot x-x' +fig, ax = plt.subplots(figsize=(4, 4)) +ax.hist2d(particles[:, 0], particles[:, 1], bins=100, range=[limits[0], limits[1]]) +plot_rms_ellipse( + env_cov_matrix[0:2, 0:2], + center=(env_centroid[0], env_centroid[1]), + level=2.0, + color="red", + ax=ax, +) +ax.set_xlabel(labels[0]) +ax.set_ylabel(labels[1]) +plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) +plt.close() + +# Plot corner +fig, axs = plot_corner( + particles, + limits=limits, + bins=100, + labels=labels, +) +for i in range(6): + for j in range(i): + env_cov_matrix_proj = project_cov_matrix(env_cov_matrix, axis=(j, i)) + plot_rms_ellipse( + env_cov_matrix_proj, + center=(env_centroid[j], env_centroid[i]), + level=2.0, + color="red", + ax=axs[i, j], + ) +plt.savefig(os.path.join(output_dir, "fig_dist_corner")) +plt.close() diff --git a/examples/Envelope/sns_ring/inputs/sns_ring.lat b/examples/Envelope/sns_ring/inputs/sns_ring.lat new file mode 100644 index 00000000..cbfb8df1 --- /dev/null +++ b/examples/Envelope/sns_ring/inputs/sns_ring.lat @@ -0,0 +1,768 @@ +kvx12 := kdc; +kdc = -3.412580596; +kta11c12 := 0; +brho := 1e9*(pc/c); +pc := sqrt(ek*(ek+2*e0)); +ek := 0.8; +e0 := 0.93827231; +c := 299792458; +khx13 := kfc; +kfc = 3.161218767; +kta10b13 := 0; +kdhta13 := 0; +kdvta13 := 0; +vkck12 := 0; +hkck12 := 0; +vkck13 := 0; +hkck13 := 0; +kdvtb1 := 0; +ksc_b01 := 0; +kssb1_9 = 0; +kvx1 := kdee; +kdee = -2.579970923; +ktb1_9 := 0; +ksv1 := 0; +ksc_b02 := 0; +kssb2_8 = 0; +khx2 := kf; +kf = 3.388342139; +ktb2d8 := 0; +ksh2 := 0; +kdvtb3 := 0; +ksc_b03 := 0; +kvx3 := kd; +kd = -3.731394588; +ktb357 := 0; +ksv3 := chrm3; +chrm3 := 0; +kdhtb4 := 0; +khx4 := kf26; +kf26 := kf*(lf/lf26); +lf := 0.25; +lf26 := 0.2705; +kta4b6 := 0; +ksh4 := chrm4; +chrm4 := 0; +ksv5 := chrm5; +chrm5 := 0; +kvx5 := kd; +kdvtb5 := 0; +ksc_b05 := 0; +ksh6 := chrm6; +chrm6 := 0; +khx6 := kf26; +kdhtb6 := 0; +ksv7 := chrm7; +chrm7 := chrm3; +kvx7 := kd; +kdvtb7 := 0; +ksc_b07 := 0; +ko_b08 := 0; +khx8 := kf; +kdhtb8 := 0; +ksc_b08 := 0; +ko_b09 := 0; +kvx9 := kdee; +kdvtb9 := 0; +ksc_b09 := 0; +kdhtb10 := 0; +kdvtb10 := 0; +khx10 := kfc; +kvx11 := kdc; +ktb11d12 := 0; +kdhtb13 := 0; +kdvtb13 := 0; +kdvtc1 := 0; +ksc_c01 := 0; +kssc1_9 = 0; +ktc1_9 := 0; +kdhtc2 := 0; +kssc2_8 = 0; +ksc_c02 := 0; +kta2c8 := 0; +kdvtc3 := 0; +ksc_c03 := 0; +ktc357 := 0; +kdhtc4 := 0; +ktc4d6 := 0; +kdvtc5 := 0; +ksc_c05 := 0; +kdhtc6 := 0; +kdvtc7 := 0; +ksc_c07 := 0; +ko_c08 := 0; +kdhtc8 := 0; +ksc_c08 := 0; +ko_c09 := 0; +kdvtc9 := 0; +ksc_c09 := 0; +kdhtc10 := 0; +kdvtc10 := 0; +ktc10d13 := 0; +kdhtc13 := 0; +kdvtc13 := 0; +ksisol := 1.22174*kssol; +kssol := solfield/brho; +solfield = 0; +kdvtd1 := 0; +ksc_d01 := 0; +kssd1_9 = 0; +ktd1_9 := 0; +kdhtd2 := 0; +ksc_d02 := 0; +kssd2_8 = 0; +kdvtd3 := 0; +ksc_d03 := 0; +ktd357 := 0; +lsv3 := lsxt; +lsxt := 0.317; +kdhtd4 := 0; +kdvtd5 := 0; +ksc_d05 := 0; +kdhtd6 := 0; +kdvtd7 := 0; +ksc_d07 := 0; +ko_d08 := 0; +kdhtd8 := 0; +ksc_d08 := 0; +ko_d09 := 0; +kdvtd9 := 0; +ksc_d09 := 0; +kdhtd10 := 0; +kdvtd10 := 0; +kdhtd13 := 0; +kdvtd13 := 0; +kdvta1 := 0; +ksc_a01 := 0; +kssa1_9 = 0; +kta1_9 := 0; +kdhta2 := 0; +ksc_a02 := 0; +kssa2_8 = 0; +kdvta3 := 0; +ksc_a03 := 0; +kta357 := 0; +kdhta4 := 0; +kdvta5 := 0; +kdhta6 := 0; +kdvta7 := 0; +ksc_a07 := 0; +ko_a08 := 0; +kdhta8 := 0; +ksc_a08 := 0; +ko_a09 := 0; +kdvta9 := 0; +ksc_a09 := 0; +hkck10 := 0; +vkck10 := 0; +hkck11 := 0; +vkck11 := 0; +kdhta10 := 0; +kdvta10 := 0; +injm1: marker; +dh_a12: sbend,l:= 0.9903829659,angle:= 0.0436,e1:= -0.003,e2:= 0.0466; +dh_a13: sbend,l:= 0.8903221964,angle:= -0.0466,e1:= -0.0466,e2:= -6.938893904e-18; +qtv_a12: quadrupole,l:= 0.533,k1:=( kvx12 + kta11c12 ) / brho ; +injm4: marker; +qth_a13: quadrupole,l:= 0.673,k1:=( khx13 + kta10b13 ) / brho ; +bpm_a13: monitor; +dchv_a13: kicker,l:= 0,hkick:=kdhta13 ,vkick:=kdvta13 ; +ikickv_a12: vkicker,l:= 0.428,kick:=vkck12 ; +ikickh_a12: hkicker,l:= 0.428,kick:=hkck12 ; +ikickv_a13: vkicker,l:= 0.839,kick:=vkck13 ; +ikickh_a13: hkicker,l:= 0.839,kick:=hkck13 ; +injm2: marker; +dmcv_b01: vkicker,l:= 0,kick:=kdvtb1 ; +qsc_b01: multipole,knl:={ 0,ksc_b01 }; +ssxc_b01: multipole,ksl:={ 0, 0,kssb1_9 }; +bpm_b01: monitor; +qtv_b01: quadrupole,l:= 0.5,k1:=( kvx1 + ktb1_9 ) / brho ; +scv_b01: sextupole,l:= 0.354,k2:=ksv1 ; +dh_b01: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_b02: hkicker,l:= 0,kick:= 0; +qsc_b02: multipole,knl:={ 0,ksc_b02 }; +ssxc_b02: multipole,ksl:={ 0, 0,kssb2_8 }; +bpm_b02: monitor; +qth_b02: quadrupole,l:= 0.5,k1:=( khx2 + ktb2d8 ) / brho ; +sch_b02: sextupole,l:= 0.354,k2:=ksh2 ; +dh_b02: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmcv_b03: vkicker,l:= 0,kick:=kdvtb3 ; +qsc_b03: multipole,knl:={ 0,ksc_b03 }; +bpm_b03: monitor; +qtv_b03: quadrupole,l:= 0.5,k1:=( kvx3 + ktb357 ) / brho ; +sv_b03: sextupole,l:= 0.317,k2:=ksv3 ; +dh_b03: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_b04: hkicker,l:= 0,kick:=kdhtb4 ; +bpm_b04: monitor; +qth_b04: quadrupole,l:= 0.541,k1:=( khx4 - kta4b6 ) / brho ; +sh_b04: sextupole,l:= 0.33,k2:=ksh4 ; +dh_b04: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_b05: sextupole,l:= 0.317,k2:=ksv5 ; +qtv_b05: quadrupole,l:= 0.5,k1:=( kvx5 + ktb357 ) / brho ; +bpm_b05: monitor; +dmcv_b05: vkicker,l:= 0,kick:=kdvtb5 ; +qsc_b05: multipole,knl:={ 0,ksc_b05 }; +dh_b06: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sh_b06: sextupole,l:= 0.33,k2:=ksh6 ; +qth_b06: quadrupole,l:= 0.541,k1:=( khx6 - kta4b6 ) / brho ; +bpm_b06: monitor; +dmch_b06: hkicker,l:= 0,kick:=kdhtb6 ; +dh_b07: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_b07: sextupole,l:= 0.317,k2:=ksv7 ; +qtv_b07: quadrupole,l:= 0.5,k1:=( kvx7 + ktb357 ) / brho ; +bpm_b07: monitor; +dmcv_b07: vkicker,l:= 0,kick:=kdvtb7 ; +qsc_b07: multipole,knl:={ 0,ksc_b07 }; +dh_b08: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_b08: octupole,l:= 0.33,k3:=ko_b08 ; +qth_b08: quadrupole,l:= 0.5,k1:=( khx8 + ktb2d8 ) / brho ; +bpm_b08: monitor; +dmch_b08: hkicker,l:= 0,kick:=kdhtb8 ; +qsc_b08: multipole,knl:={ 0,ksc_b08 }; +ssxc_b08: multipole,ksl:={ 0, 0,kssb2_8 }; +dh_b09: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_b09: octupole,l:= 0.33,k3:=ko_b09 ; +qtv_b09: quadrupole,l:= 0.5,k1:=( kvx9 + ktb1_9 ) / brho ; +bpm_b09: monitor; +dmcv_b09: vkicker,l:= 0,kick:=kdvtb9 ; +qsc_b09: multipole,knl:={ 0,ksc_b09 }; +ssxc_b09: multipole,ksl:={ 0, 0,kssb1_9 }; +dchv_b10: kicker,l:= 0,hkick:=kdhtb10 ,vkick:=kdvtb10 ; +bpm_b10: monitor; +qth_b10: quadrupole,l:= 0.673,k1:=( khx10 - kta10b13 ) / brho ; +qtv_b11: quadrupole,l:= 0.533,k1:=( kvx11 + ktb11d12 ) / brho ; +qtv_b12: quadrupole,l:= 0.533,k1:=( kvx12 + ktb11d12 ) / brho ; +qth_b13: quadrupole,l:= 0.673,k1:=( khx13 - kta10b13 ) / brho ; +bpm_b13: monitor; +dchv_b13: kicker,l:= 0,hkick:=kdhtb13 ,vkick:=kdvtb13 ; +dampkicker1: marker; +dampkicker2: marker; +qmmkicker: marker; +tunekicker: marker; +dmcv_c01: vkicker,l:= 0,kick:=kdvtc1 ; +qsc_c01: multipole,knl:={ 0,ksc_c01 }; +ssxc_c01: multipole,ksl:={ 0, 0,kssc1_9 }; +bpm_c01: monitor; +qtv_c01: quadrupole,l:= 0.5,k1:=( kvx1 + ktc1_9 ) / brho ; +scv_c01: sextupole,l:= 0.354,k2:=ksv1 ; +dh_c01: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_c02: hkicker,l:= 0,kick:=kdhtc2 ; +ssxc_c02: multipole,ksl:={ 0, 0,kssc2_8 }; +qsc_c02: multipole,knl:={ 0,ksc_c02 }; +bpm_c02: monitor; +qth_c02: quadrupole,l:= 0.5,k1:=( khx2 + kta2c8 ) / brho ; +sch_c02: sextupole,l:= 0.354,k2:=ksh2 ; +dh_c02: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmcv_c03: vkicker,l:= 0,kick:=kdvtc3 ; +qsc_c03: multipole,knl:={ 0,ksc_c03 }; +bpm_c03: monitor; +qtv_c03: quadrupole,l:= 0.5,k1:=( kvx3 + ktc357 ) / brho ; +sv_c03: sextupole,l:= 0.317,k2:=ksv3 ; +dh_c03: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_c04: hkicker,l:= 0,kick:=kdhtc4 ; +bpm_c04: monitor; +qth_c04: quadrupole,l:= 0.541,k1:=( khx4 + ktc4d6 ) / brho ; +sh_c04: sextupole,l:= 0.33,k2:=ksh4 ; +dh_c04: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_c05: sextupole,l:= 0.317,k2:=ksv5 ; +qtv_c05: quadrupole,l:= 0.5,k1:=( kvx5 + ktc357 ) / brho ; +bpm_c05: monitor; +dmcv_c05: vkicker,l:= 0,kick:=kdvtc5 ; +qsc_c05: multipole,knl:={ 0,ksc_c05 }; +dh_c06: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sh_c06: sextupole,l:= 0.33,k2:=ksh6 ; +qth_c06: quadrupole,l:= 0.541,k1:=( khx6 + ktc4d6 ) / brho ; +bpm_c06: monitor; +dmch_c06: hkicker,l:= 0,kick:=kdhtc6 ; +dh_c07: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_c07: sextupole,l:= 0.317,k2:=ksv7 ; +qtv_c07: quadrupole,l:= 0.5,k1:=( kvx7 + ktc357 ) / brho ; +bpm_c07: monitor; +dmcv_c07: vkicker,l:= 0,kick:=kdvtc7 ; +qsc_c07: multipole,knl:={ 0,ksc_c07 }; +dh_c08: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_c08: octupole,l:= 0.33,k3:=ko_c08 ; +qth_c08: quadrupole,l:= 0.5,k1:=( khx8 + kta2c8 ) / brho ; +bpm_c08: monitor; +dmch_c08: hkicker,l:= 0,kick:=kdhtc8 ; +qsc_c08: multipole,knl:={ 0,ksc_c08 }; +ssxc_c08: multipole,ksl:={ 0,kssc2_8 }; +dh_c09: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_c09: octupole,l:= 0.33,k3:=ko_c09 ; +qtv_c09: quadrupole,l:= 0.5,k1:=( kvx9 + ktc1_9 ) / brho ; +bpm_c09: monitor; +dmcv_c09: vkicker,l:= 0,kick:=kdvtc9 ; +qsc_c09: multipole,knl:={ 0,ksc_c09 }; +ssxc_c09: multipole,ksl:={ 0, 0,kssc1_9 }; +ekick01: vkicker,l:= 0.4; +ekick02: vkicker,l:= 0.4; +ekick03: vkicker,l:= 0.4; +ekick04: vkicker,l:= 0.505; +ekick05: vkicker,l:= 0.505; +ekick06: vkicker,l:= 0.505; +ekick07: vkicker,l:= 0.505; +dchv_c10: kicker,l:= 0,hkick:=kdhtc10 ,vkick:=kdvtc10 ; +bpm_c10: monitor; +qth_c10: quadrupole,l:= 0.673,k1:=( khx10 + ktc10d13 ) / brho ; +qtv_c11: quadrupole,l:= 0.533,k1:=( kvx11 - kta11c12 ) / brho ; +ekick08: vkicker,l:= 0.4275; +ekick09: vkicker,l:= 0.4275; +ekick10: vkicker,l:= 0.4275; +ekick11: vkicker,l:= 0.4275; +ekick12: vkicker,l:= 0.39; +ekick13: vkicker,l:= 0.39; +ekick14: vkicker,l:= 0.39; +qtv_c12: quadrupole,l:= 0.533,k1:=( kvx12 - kta11c12 ) / brho ; +qth_c13: quadrupole,l:= 0.673,k1:=( khx13 + ktc10d13 ) / brho ; +bpm_c13: monitor; +dchv_c13: kicker,l:= 0,hkick:=kdhtc13 ,vkick:=kdvtc13 ; +scbdsol_c13a: solenoid,l:= 1.22174,ks:= 0,ksi:=ksisol ; +scbdsol_c13b: solenoid,l:= 1.22174,ks:= 0,ksi:=ksisol ; +dmcv_d01: vkicker,l:= 0,kick:=kdvtd1 ; +qsc_d01: multipole,knl:={ 0,ksc_d01 }; +ssxc_d01: multipole,ksl:={ 0,kssd1_9 }; +bpm_d01: monitor; +qtv_d01: quadrupole,l:= 0.5,k1:=( kvx1 + ktd1_9 ) / brho ; +scv_d01: sextupole,l:= 0.354,k2:=ksv1 ; +dh_d01: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_d02: hkicker,l:= 0,kick:=kdhtd2 ; +qsc_d02: multipole,knl:={ 0,ksc_d02 }; +ssxc_d02: multipole,ksl:={ 0, 0,kssd2_8 }; +bpm_d02: monitor; +qth_d02: quadrupole,l:= 0.5,k1:=( khx2 + ktb2d8 ) / brho ; +sch_d02: sextupole,l:= 0.354,k2:=ksh2 ; +dh_d02: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmcv_d03: vkicker,l:= 0,kick:=kdvtd3 ; +qsc_d03: multipole,knl:={ 0,ksc_d03 }; +bpm_d03: monitor; +qtv_d03: quadrupole,l:= 0.5,k1:=( kvx3 + ktd357 ) / brho ; +sv_d03: sextupole,l:=lsv3 ,k2:=ksv3 ; +dh_d03: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_d04: hkicker,l:= 0,kick:=kdhtd4 ; +bpm_d04: monitor; +qth_d04: quadrupole,l:= 0.541,k1:=( khx4 - ktc4d6 ) / brho ; +sh_d04: sextupole,l:= 0.33,k2:=ksh4 ; +dh_d04: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_d05: sextupole,l:= 0.317,k2:=ksv5 ; +qtv_d05: quadrupole,l:= 0.5,k1:=( kvx5 + ktd357 ) / brho ; +bpm_d05: monitor; +dmcv_d05: vkicker,l:= 0,kick:=kdvtd5 ; +qsc_d05: multipole,knl:={ 0,ksc_d05 }; +dh_d06: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sh_d06: sextupole,l:= 0.33,k2:=ksh6 ; +qth_d06: quadrupole,l:= 0.541,k1:=( khx6 - ktc4d6 ) / brho ; +bpm_d06: monitor; +dmch_d06: hkicker,l:= 0,kick:=kdhtd6 ; +dh_d07: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_d07: sextupole,l:= 0.317,k2:=ksv7 ; +qtv_d07: quadrupole,l:= 0.5,k1:=( kvx7 + ktd357 ) / brho ; +bpm_d07: monitor; +dmcv_d07: vkicker,l:= 0,kick:=kdvtd7 ; +qsc_d07: multipole,knl:={ 0,ksc_d07 }; +dh_d08: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_d08: octupole,l:= 0.33,k3:=ko_d08 ; +qth_d08: quadrupole,l:= 0.5,k1:=( khx8 + ktb2d8 ) / brho ; +bpm_d08: monitor; +dmch_d08: hkicker,l:= 0,kick:=kdhtd8 ; +qsc_d08: multipole,knl:={ 0,ksc_d08 }; +ssxc_d08: multipole,ksl:={ 0, 0,kssd2_8 }; +dh_d09: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_d09: octupole,l:= 0.33,k3:=ko_d09 ; +qtv_d09: quadrupole,l:= 0.5,k1:=( kvx9 + ktd1_9 ) / brho ; +bpm_d09: monitor; +dmcv_d09: vkicker,l:= 0,kick:=kdvtd9 ; +qsc_d09: multipole,knl:={ 0,ksc_d09 }; +ssxc_d09: multipole,ksl:={ 0, 0,kssd1_9 }; +tunepickup: marker; +qmmpickup: marker; +wcm: marker; +bcm: marker; +dchv_d10: kicker,l:= 0,hkick:=kdhtd10 ,vkick:=kdvtd10 ; +bpm_d10: monitor; +qth_d10: quadrupole,l:= 0.673,k1:=( khx10 - ktc10d13 ) / brho ; +qtv_d11: quadrupole,l:= 0.533,k1:=( kvx11 - ktb11d12 ) / brho ; +cav_01: rfcavity,l:= 2.1466,volt:= 0.0133,harmon:= 1; +cav_02: rfcavity,l:= 2.1466,volt:= 0.0133,harmon:= 1; +cav_03: rfcavity,l:= 2.1466,volt:= 0.0133,harmon:= 1; +cav_04: rfcavity,l:= 2.1466,volt:= -0.02,harmon:= 2; +qtv_d12: quadrupole,l:= 0.533,k1:=( kvx12 - ktb11d12 ) / brho ; +qth_d13: quadrupole,l:= 0.673,k1:=( khx13 - ktc10d13 ) / brho ; +bpm_d13: monitor; +dchv_d13: kicker,l:= 0,hkick:=kdhtd13 ,vkick:=kdvtd13 ; +haloscanner1: marker; +wirescanner1: marker; +ipm1: marker; +ipm2: marker; +wirescanner2: marker; +haloscanner2: marker; +dmcv_a01: vkicker,l:= 0,kick:=kdvta1 ; +qsc_a01: multipole,knl:={ 0,ksc_a01 }; +ssxc_a01: multipole,ksl:={ 0, 0,kssa1_9 }; +bpm_a01: monitor; +qtv_a01: quadrupole,l:= 0.5,k1:=( kvx1 + kta1_9 ) / brho ; +scv_a01: sextupole,l:= 0.354,k2:=ksv1 ; +dh_a01: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_a02: hkicker,l:= 0,kick:=kdhta2 ; +qsc_a02: multipole,knl:={ 0,ksc_a02 }; +ssxc_a02: multipole,ksl:={ 0, 0,kssa2_8 }; +bpm_a02: monitor; +qth_a02: quadrupole,l:= 0.5,k1:=( khx2 + kta2c8 ) / brho ; +sch_a02: sextupole,l:= 0.354,k2:=ksh2 ; +dh_a02: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmcv_a03: vkicker,l:= 0,kick:=kdvta3 ; +qsc_a03: multipole,knl:={ 0,ksc_a03 }; +bpm_a03: monitor; +qtv_a03: quadrupole,l:= 0.5,k1:=( kvx3 + kta357 ) / brho ; +sv_a03: sextupole,l:= 0.317,k2:=ksv3 ; +dh_a03: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_a04: hkicker,l:= 0,kick:=kdhta4 ; +bpm_a04: monitor; +qth_a04: quadrupole,l:= 0.541,k1:=( khx4 + kta4b6 ) / brho ; +sh_a04: sextupole,l:= 0.33,k2:=ksh4 ; +dh_a04: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_a05: sextupole,l:= 0.317,k2:=ksv5 ; +qtv_a05: quadrupole,l:= 0.5,k1:=( kvx5 + kta357 ) / brho ; +bpm_a05: monitor; +dmcv_a05: vkicker,l:= 0,kick:=kdvta5 ; +qsc_a05: multipole,knl:={ 0}; +dh_a06: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sh_a06: sextupole,l:= 0.33,k2:=ksh6 ; +qth_a06: quadrupole,l:= 0.541,k1:=( khx6 + kta4b6 ) / brho ; +bpm_a06: monitor; +dmch_a06: hkicker,l:= 0,kick:=kdhta6 ; +dh_a07: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_a07: sextupole,l:= 0.317,k2:=ksv7 ; +qtv_a07: quadrupole,l:= 0.5,k1:=( kvx7 + kta357 ) / brho ; +bpm_a07: monitor; +dmcv_a07: vkicker,l:= 0,kick:=kdvta7 ; +qsc_a07: multipole,knl:={ 0,ksc_a07 }; +dh_a08: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_a08: octupole,l:= 0.33,k3:=ko_a08 ; +qth_a08: quadrupole,l:= 0.5,k1:=( khx8 + kta2c8 ) / brho ; +bpm_a08: monitor; +dmch_a08: hkicker,l:= 0,kick:=kdhta8 ; +qsc_a08: multipole,knl:={ 0,ksc_a08 }; +ssxc_a08: multipole,ksl:={ 0, 0,kssa2_8 }; +dh_a09: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_a09: octupole,l:= 0.33,k3:=ko_a09 ; +qtv_a09: quadrupole,l:= 0.5,k1:=( kvx9 + kta1_9 ) / brho ; +bpm_a09: monitor; +dmcv_a09: vkicker,l:= 0,kick:=kdvta9 ; +qsc_a09: multipole,knl:={ 0,ksc_a09 }; +ssxc_a09: multipole,ksl:={ 0, 0,kssa1_9 }; +ikickh_a10: hkicker,l:= 0.839,kick:=hkck10 ; +ikickv_a10: vkicker,l:= 0.839,kick:=vkck10 ; +ikickh_a11: hkicker,l:= 0.428,kick:=hkck11 ; +ikickv_a11: vkicker,l:= 0.428,kick:=vkck11 ; +dchv_a10: kicker,l:= 0,hkick:=kdhta10 ,vkick:=kdvta10 ; +bpm_a10: monitor; +qth_a10: quadrupole,l:= 0.673,k1:=( khx10 + kta10b13 ) / brho ; +injm3: marker; +qtv_a11: quadrupole,l:= 0.533,k1:=( kvx11 + kta11c12 ) / brho ; +dh_a10: sbend,l:= 0.8632537742,angle:= -0.042,e1:= 0,e2:= -0.042; +dh_a11: sbend,l:= 0.8722394086,angle:= 0.045,e1:= 0.042,e2:= 0.003; +rnginjsol: sequence, l = 248.0098418; +injm1, at = 0; +dh_a12, at = 1.378195456; +dh_a13, at = 3.408731523; +qtv_a12, at = 7.004892621; +injm4, at = 7.693392621; +qth_a13, at = 8.029892621; +bpm_a13, at = 8.547265121; +dchv_a13, at = 8.683688621; +ikickv_a12, at = 10.59989262; +ikickh_a12, at = 11.13989262; +ikickv_a13, at = 12.66989262; +ikickh_a13, at = 13.82989262; +injm2, at = 14.24939262; +dmcv_b01, at = 14.92668062; +qsc_b01, at = 14.92668062; +ssxc_b01, at = 14.92668062; +bpm_b01, at = 15.12176002; +qtv_b01, at = 15.47989262; +scv_b01, at = 16.03310462; +dh_b01, at = 17.47998825; +dmch_b02, at = 18.92687188; +qsc_b02, at = 18.92687188; +ssxc_b02, at = 18.92687188; +bpm_b02, at = 19.11879438; +bpm_b02, at = 19.23008388; +qth_b02, at = 19.48008388; +sch_b02, at = 20.03329588; +dh_b02, at = 21.4801795; +dmcv_b03, at = 22.92706313; +qsc_b03, at = 22.92706313; +bpm_b03, at = 23.11546513; +qtv_b03, at = 23.48027513; +sv_b03, at = 24.01443713; +dh_b03, at = 25.48037076; +dmch_b04, at = 26.89474238; +bpm_b04, at = 27.11860888; +qth_b04, at = 27.48046638; +sh_b04, at = 28.08524038; +dh_b04, at = 29.48056201; +sv_b05, at = 30.94649564; +qtv_b05, at = 31.48065764; +bpm_b05, at = 31.84546764; +dmcv_b05, at = 32.03386964; +qsc_b05, at = 32.03386964; +dh_b06, at = 33.48075326; +sh_b06, at = 34.87607489; +qth_b06, at = 35.48084889; +bpm_b06, at = 35.84270639; +dmch_b06, at = 36.06657289; +dh_b07, at = 37.48094452; +sv_b07, at = 38.94687815; +qtv_b07, at = 39.48104015; +bpm_b07, at = 39.84585015; +dmcv_b07, at = 40.03425215; +qsc_b07, at = 40.03425215; +dh_b08, at = 41.48113577; +oct_b08, at = 42.9280194; +qth_b08, at = 43.4812314; +bpm_b08, at = 43.8425209; +dmch_b08, at = 44.0344434; +qsc_b08, at = 44.0344434; +ssxc_b08, at = 44.0344434; +dh_b09, at = 45.48132703; +oct_b09, at = 46.92821065; +qtv_b09, at = 47.48142265; +bpm_b09, at = 47.83955525; +dmcv_b09, at = 48.03463465; +qsc_b09, at = 48.03463465; +ssxc_b09, at = 48.03463465; +dchv_b10, at = 54.27762665; +bpm_b10, at = 54.41405015; +qth_b10, at = 54.93142265; +qtv_b11, at = 55.95642265; +qtv_b12, at = 69.00642265; +qth_b13, at = 70.03142265; +bpm_b13, at = 70.54879515; +dchv_b13, at = 70.68521865; +dampkicker1, at = 73.19024865; +dampkicker2, at = 73.69024765; +qmmkicker, at = 74.31523965; +tunekicker, at = 75.44023065; +dmcv_c01, at = 76.92821065; +qsc_c01, at = 76.92821065; +ssxc_c01, at = 76.92821065; +bpm_c01, at = 77.12329005; +qtv_c01, at = 77.48142265; +scv_c01, at = 78.03463465; +dh_c01, at = 79.48151828; +dmch_c02, at = 80.92840191; +ssxc_c02, at = 80.92840191; +qsc_c02, at = 80.92840191; +bpm_c02, at = 81.12032441; +qth_c02, at = 81.48161391; +sch_c02, at = 82.03482591; +dh_c02, at = 83.48170954; +dmcv_c03, at = 84.92859316; +qsc_c03, at = 84.92859316; +bpm_c03, at = 85.11699516; +qtv_c03, at = 85.48180516; +sv_c03, at = 86.01596716; +dh_c03, at = 87.48190079; +dmch_c04, at = 88.89627242; +bpm_c04, at = 89.12013892; +qth_c04, at = 89.48199642; +sh_c04, at = 90.08677042; +dh_c04, at = 91.48209204; +sv_c05, at = 92.94802567; +qtv_c05, at = 93.48218767; +bpm_c05, at = 93.84699767; +dmcv_c05, at = 94.03539967; +qsc_c05, at = 94.03539967; +dh_c06, at = 95.4822833; +sh_c06, at = 96.87760493; +qth_c06, at = 97.48237893; +bpm_c06, at = 97.84423643; +dmch_c06, at = 98.06810293; +dh_c07, at = 99.48247455; +sv_c07, at = 100.9484082; +qtv_c07, at = 101.4825702; +bpm_c07, at = 101.8473802; +dmcv_c07, at = 102.0357822; +qsc_c07, at = 102.0357822; +dh_c08, at = 103.4826658; +oct_c08, at = 104.9295494; +qth_c08, at = 105.4827614; +bpm_c08, at = 105.8440509; +dmch_c08, at = 106.0359734; +qsc_c08, at = 106.0359734; +ssxc_c08, at = 106.0359734; +dh_c09, at = 107.4828571; +oct_c09, at = 108.9297407; +qtv_c09, at = 109.4829527; +bpm_c09, at = 109.8410853; +dmcv_c09, at = 110.0361647; +qsc_c09, at = 110.0361647; +ssxc_c09, at = 110.0361647; +ekick01, at = 112.3299527; +ekick02, at = 112.8099527; +ekick03, at = 113.2899527; +ekick04, at = 113.8229527; +ekick05, at = 114.4079527; +ekick06, at = 114.9929527; +ekick07, at = 115.5779527; +dchv_c10, at = 116.2791567; +bpm_c10, at = 116.4155802; +qth_c10, at = 116.9329527; +qtv_c11, at = 117.9579527; +ekick08, at = 118.9059527; +ekick09, at = 119.4149527; +ekick10, at = 119.9249527; +ekick11, at = 120.4309527; +ekick12, at = 120.9202527; +ekick13, at = 121.3895527; +ekick14, at = 121.8606527; +qtv_c12, at = 131.0079527; +qth_c13, at = 132.0329527; +bpm_c13, at = 132.5503252; +dchv_c13, at = 132.6867487; +scbdsol_c13a, at = 134.3897627; +scbdsol_c13b, at = 138.0255227; +dmcv_d01, at = 138.9297407; +qsc_d01, at = 138.9297407; +ssxc_d01, at = 138.9297407; +bpm_d01, at = 139.1248201; +qtv_d01, at = 139.4829527; +scv_d01, at = 140.0361647; +dh_d01, at = 141.4830483; +dmch_d02, at = 142.9299319; +qsc_d02, at = 142.9299319; +ssxc_d02, at = 142.9299319; +bpm_d02, at = 143.1218544; +qth_d02, at = 143.4831439; +sch_d02, at = 144.0363559; +dh_d02, at = 145.4832396; +dmcv_d03, at = 146.9301232; +qsc_d03, at = 146.9301232; +bpm_d03, at = 147.1185252; +qtv_d03, at = 147.4833352; +sv_d03, at = 148.0174972; +dh_d03, at = 149.4834308; +dmch_d04, at = 150.8978024; +bpm_d04, at = 151.1216689; +qth_d04, at = 151.4835264; +sh_d04, at = 152.0883004; +dh_d04, at = 153.4836221; +sv_d05, at = 154.9495557; +qtv_d05, at = 155.4837177; +bpm_d05, at = 155.8485277; +dmcv_d05, at = 156.0369297; +qsc_d05, at = 156.0369297; +dh_d06, at = 157.4838133; +sh_d06, at = 158.879135; +qth_d06, at = 159.483909; +bpm_d06, at = 159.8457665; +dmch_d06, at = 160.069633; +dh_d07, at = 161.4840046; +sv_d07, at = 162.9499382; +qtv_d07, at = 163.4841002; +bpm_d07, at = 163.8489102; +dmcv_d07, at = 164.0373122; +qsc_d07, at = 164.0373122; +dh_d08, at = 165.4841958; +oct_d08, at = 166.9310795; +qth_d08, at = 167.4842915; +bpm_d08, at = 167.845581; +dmch_d08, at = 168.0375035; +qsc_d08, at = 168.0375035; +ssxc_d08, at = 168.0375035; +dh_d09, at = 169.4843871; +oct_d09, at = 170.9312707; +qtv_d09, at = 171.4844827; +bpm_d09, at = 171.8426153; +dmcv_d09, at = 172.0376947; +qsc_d09, at = 172.0376947; +ssxc_d09, at = 172.0376947; +tunepickup, at = 173.5472277; +qmmpickup, at = 174.6722197; +wcm, at = 175.7778927; +bcm, at = 176.7642687; +dchv_d10, at = 178.2806867; +bpm_d10, at = 178.4171102; +qth_d10, at = 178.9344827; +qtv_d11, at = 179.9594827; +cav_01, at = 183.0386827; +cav_02, at = 185.3358827; +cav_03, at = 187.6330827; +cav_04, at = 189.9302827; +qtv_d12, at = 193.0094827; +qth_d13, at = 194.0344827; +bpm_d13, at = 194.5518552; +dchv_d13, at = 194.6882787; +haloscanner1, at = 196.4528147; +wirescanner1, at = 196.6309197; +ipm1, at = 197.0093417; +ipm2, at = 199.3134267; +wirescanner2, at = 199.6918487; +haloscanner2, at = 199.8699527; +dmcv_a01, at = 200.9312707; +qsc_a01, at = 200.9312707; +ssxc_a01, at = 200.9312707; +bpm_a01, at = 201.1263501; +qtv_a01, at = 201.4844827; +scv_a01, at = 202.0376947; +dh_a01, at = 203.4845783; +dmch_a02, at = 204.931462; +qsc_a02, at = 204.931462; +ssxc_a02, at = 204.931462; +bpm_a02, at = 205.1233845; +qth_a02, at = 205.484674; +sch_a02, at = 206.037886; +dh_a02, at = 207.4847696; +dmcv_a03, at = 208.9316532; +qsc_a03, at = 208.9316532; +bpm_a03, at = 209.1200552; +qtv_a03, at = 209.4848652; +sv_a03, at = 210.0190272; +dh_a03, at = 211.4849609; +dmch_a04, at = 212.8993325; +bpm_a04, at = 213.123199; +qth_a04, at = 213.4850565; +sh_a04, at = 214.0898305; +dh_a04, at = 215.4851521; +sv_a05, at = 216.9510857; +qtv_a05, at = 217.4852477; +bpm_a05, at = 217.8500577; +dmcv_a05, at = 218.0384597; +qsc_a05, at = 218.0384597; +dh_a06, at = 219.4853434; +sh_a06, at = 220.880665; +qth_a06, at = 221.485439; +bpm_a06, at = 221.8472965; +dmch_a06, at = 222.071163; +dh_a07, at = 223.4855346; +sv_a07, at = 224.9514682; +qtv_a07, at = 225.4856302; +bpm_a07, at = 225.8504402; +dmcv_a07, at = 226.0388422; +qsc_a07, at = 226.0388422; +dh_a08, at = 227.4857259; +oct_a08, at = 228.9326095; +qth_a08, at = 229.4858215; +bpm_a08, at = 229.847111; +dmch_a08, at = 230.0390335; +qsc_a08, at = 230.0390335; +ssxc_a08, at = 230.0390335; +dh_a09, at = 231.4859171; +oct_a09, at = 232.9328008; +qtv_a09, at = 233.4860128; +bpm_a09, at = 233.8441454; +dmcv_a09, at = 234.0392248; +qsc_a09, at = 234.0392248; +ssxc_a09, at = 234.0392248; +ikickh_a10, at = 235.1360128; +ikickv_a10, at = 236.2960128; +ikickh_a11, at = 237.8260128; +ikickv_a11, at = 238.3660128; +dchv_a10, at = 240.2822168; +bpm_a10, at = 240.4186403; +qth_a10, at = 240.9360128; +injm3, at = 241.2725128; +qtv_a11, at = 241.9610128; +dh_a10, at = 245.1911396; +dh_a11, at = 247.5737221; +endsequence; diff --git a/examples/Envelope/sns_ring/style.mplstyle b/examples/Envelope/sns_ring/style.mplstyle new file mode 100644 index 00000000..71f36605 --- /dev/null +++ b/examples/Envelope/sns_ring/style.mplstyle @@ -0,0 +1,8 @@ +axes.linewidth: 1.25 +axes.titlesize: "medium" +image.cmap: "Greys" +figure.constrained_layout.use: True +savefig.dpi: 300 +savefig.format: "png" +xtick.minor.visible: True +ytick.minor.visible: True diff --git a/examples/Envelope/sns_ring/test_sns_ring.py b/examples/Envelope/sns_ring/test_sns_ring.py new file mode 100644 index 00000000..b7e74716 --- /dev/null +++ b/examples/Envelope/sns_ring/test_sns_ring.py @@ -0,0 +1,389 @@ +"""Test envelope tracker in SNS ring.""" + +import argparse +import copy +import math +import os +import pathlib +import time +import sys + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.bunch_utils import collect_bunch +from orbit.envelope import Envelope +from orbit.envelope import EnvelopeTracker +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.space_charge.sc2p5d import setSC2p5DAccNodes +from orbit.teapot import TEAPOT_Ring +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.teapot import teapot +from orbit.utils.consts import mass_proton + +sys.path.append("..") +from plot import plot_rms_ellipse +from plot import plot_corner +from utils import gen_dist +from utils import build_rotation_matrix_xy +from utils import project_cov_matrix + +plt.style.use("style.mplstyle") + + +# Arguments +# ------------------------------------------------------------------------------ + +parser = argparse.ArgumentParser() +parser.add_argument("--bunch-length", type=float, default=120.0) +parser.add_argument("--kin-energy", type=float, default=1.300) +parser.add_argument("--intensity", type=float, default=2e14) + +parser.add_argument("--dist", type=str, default="kv", choices=["kv", "waterbag", "gauss"]) +parser.add_argument("--eps-x", type=float, default=25.0) +parser.add_argument("--eps-y", type=float, default=25.0) +parser.add_argument("--mismatch-x", type=float, default=0.0) +parser.add_argument("--mismatch-y", type=float, default=0.0) +parser.add_argument("--offset-x", type=float, default=0.0) +parser.add_argument("--offset-y", type=float, default=0.0) +parser.add_argument("--tilt", type=float, default=0) + +parser.add_argument("--nparts", type=int, default=100_000) +parser.add_argument("--turns", type=int, default=25) +parser.add_argument("--sol", type=int, default=0) +parser.add_argument("--sc", type=int, default=0) +parser.add_argument("--sc-grid", type=int, default=64) + +parser.add_argument("--handle-unknown", type=str, default=None, choices=["drift", "fit"]) +args = parser.parse_args() + +# Setup +# ------------------------------------------------------------------------------ + +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + +# Lattice +# ------------------------------------------------------------------------------ + +lattice = TEAPOT_Ring() +lattice.readMADX("inputs/sns_ring.lat", "rnginjsol") +lattice.initialize() + +for node in lattice.getNodes(): + if type(node) != teapot.TurnCounterTEAPOT: + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + if type(node) is teapot.BendTEAPOT: + node.setParam("ea1", 0.0) + node.setParam("ea2", 0.0) + +if args.sol: + for name in ["scbdsol_c13a", "scbdsol_c13b"]: + node = lattice.getNodeForName(name) + node.setParam("B", 0.15) + +for node in lattice.getNodes(): + max_length = 1.0 + if node.getLength() > max_length: + node.setnParts(1 + int(node.getLength() / max_length)) + +# Bunch +# ------------------------------------------------------------------------------ + +bunch = Bunch() +bunch.mass(mass_proton) +sync_part = bunch.getSyncParticle() +sync_part.kinEnergy(args.kin_energy) + +matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) +matrix_lattice_params = matrix_lattice.getRingParametersDict() +alpha_x = matrix_lattice_params["alpha x"] +alpha_y = matrix_lattice_params["alpha y"] +beta_x = matrix_lattice_params["beta x [m]"] +beta_y = matrix_lattice_params["beta y [m]"] +eps_x = args.eps_x * 1e-6 +eps_y = args.eps_y * 1e-6 + +cov_matrix = np.zeros((6, 6)) +cov_matrix[0, 0] = eps_x * beta_x +cov_matrix[2, 2] = eps_y * beta_y +cov_matrix[0, 1] = cov_matrix[1, 0] = -eps_x * alpha_x +cov_matrix[2, 3] = cov_matrix[3, 2] = -eps_y * alpha_y +cov_matrix[1, 1] = eps_x * (1.0 + alpha_x**2) / beta_x +cov_matrix[3, 3] = eps_y * (1.0 + alpha_y**2) / beta_y +cov_matrix[4, 4] = (args.bunch_length / 4.0) ** 2 +cov_matrix[5, 5] = 0.0 + +if args.tilt: + rot_matrix = np.identity(6) + rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi)) + cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T]) + +if args.mismatch_x or args.mismatch_y: + cov_matrix[0, 0] *= (1.0 + args.mismatch_x) ** 2 + cov_matrix[2, 2] *= (1.0 + args.mismatch_y) ** 2 + +centroid = np.zeros(6) +centroid[0] += args.offset_x +centroid[2] += args.offset_y + +rng = np.random.default_rng() +bunch_coords = np.zeros((args.nparts, 6)) +bunch_coords[:, :4] = gen_dist(size=args.nparts, cov_matrix=cov_matrix[0:4, 0:4], name=args.dist) +bunch_coords[:, 4] = args.bunch_length * rng.uniform(-0.5, 0.5, size=args.nparts) +bunch_coords[:, 5] *= 0.0 +bunch_coords += centroid[None, :6] + +for i in range(bunch_coords.shape[0]): + bunch.addParticle(*bunch_coords[i]) + +# Use covariance matrix from initial bunch, which is slightly +# different from the one used to generate the bunch due to +# finite statistics. +cov_matrix_init = np.cov(bunch_coords, rowvar=False) +centroid_init = np.mean(bunch_coords, axis=0) + +# Track envelope +# ------------------------------------------------------------------------------ + +print("TRACK ENVELOPE") + +envelope = Envelope( + bunch=bunch, + cov_matrix=cov_matrix_init, + centroid=centroid_init, + intensity=args.intensity, +) +tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) + +history_keys = [ + "rms_x", + "rms_y", + "avg_x", + "avg_y", + "eps_x", + "eps_y", +] +history = {key: [] for key in history_keys} + +start_time = time.time() + +for turn in range(args.turns + 1): + if turn > 0: + tracker.track_ring(envelope) + + cov_matrix = envelope.cov_matrix + centroid = envelope.centroid + + rms_x = 1000.0 * math.sqrt(cov_matrix[0, 0]) + rms_y = 1000.0 * math.sqrt(cov_matrix[2, 2]) + avg_x = 1000.0 * centroid[0] + avg_y = 1000.0 * centroid[2] + eps_x = 1e6 * np.sqrt(np.linalg.det(cov_matrix[0:2, 0:2])) + eps_y = 1e6 * np.sqrt(np.linalg.det(cov_matrix[2:4, 2:4])) + + message = "" + message += " turn={}".format(turn) + message += " time={:0.3f}".format(time.time() - start_time) + message += " eps_x={:0.2f}".format(eps_x) + message += " eps_y={:0.2f}".format(eps_y) + message += " xrms={:0.2f}".format(rms_x) + message += " yrms={:0.2f}".format(rms_y) + message += " xavg={:0.2f}".format(avg_x) + message += " yavg={:0.2f}".format(avg_y) + print(message) + + history["rms_x"].append(rms_x) + history["rms_y"].append(rms_y) + history["avg_x"].append(avg_x) + history["avg_y"].append(avg_y) + history["eps_x"].append(eps_x) + history["eps_y"].append(eps_y) + +histories = {} +histories["envelope"] = copy.deepcopy(history) + +# Track bunch +# ------------------------------------------------------------------------------ + +print("TRACK BUNCH") + +if args.sc: + sc_calc = SpaceChargeCalc2p5D(64, 64, 1) + sc_path_length_min = 1.00e-06 + sc_nodes = setSC2p5DAccNodes(lattice, sc_path_length_min, sc_calc) + + bunch_size = bunch.getSizeGlobal() + bunch.macroSize(args.intensity / bunch_size) + +history_keys = [ + "rms_x", + "rms_y", + "avg_x", + "avg_y", + "eps_x", + "eps_y", +] +history = {key: [] for key in history_keys} + +start_time = time.time() + +for turn in range(args.turns + 1): + if turn > 0: + lattice.trackBunch(bunch) + + twiss_calc = BunchTwissAnalysis() + twiss_calc.computeBunchMoments(bunch, 2, 0, 0) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(i + 1): + cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + cov_matrix[j, i] = cov_matrix[i, j] + + centroid = np.zeros(6) + for i in range(6): + centroid[i] = twiss_calc.getAverage(i) + + rms_x = 1000.0 * math.sqrt(cov_matrix[0, 0]) + rms_y = 1000.0 * math.sqrt(cov_matrix[2, 2]) + avg_x = 1000.0 * centroid[0] + avg_y = 1000.0 * centroid[2] + eps_x = 1e6 * np.sqrt(np.linalg.det(cov_matrix[0:2, 0:2])) + eps_y = 1e6 * np.sqrt(np.linalg.det(cov_matrix[2:4, 2:4])) + + message = "" + message += " turn={}".format(turn) + message += " time={:0.3f}".format(time.time() - start_time) + message += " eps_x={:0.2f}".format(eps_x) + message += " eps_y={:0.2f}".format(eps_y) + message += " xrms={:0.2f}".format(rms_x) + message += " yrms={:0.2f}".format(rms_y) + message += " xavg={:0.2f}".format(avg_x) + message += " yavg={:0.2f}".format(avg_y) + print(message) + + history["rms_x"].append(rms_x) + history["rms_y"].append(rms_y) + history["avg_x"].append(avg_x) + history["avg_y"].append(avg_y) + history["eps_x"].append(eps_x) + history["eps_y"].append(eps_y) + +histories["bunch"] = copy.deepcopy(history) + +# Analysis +# ------------------------------------------------------------------------------ + +for history in histories.values(): + for key in history: + history[key] = np.array(history[key]) + +# Print errors +for key in histories["envelope"]: + deltas = histories["bunch"][key] - histories["envelope"][key] + print("key:", key) + print("max_abs_delta:", np.max(np.abs(deltas))) + print("avg_abs_delta:", np.mean(np.abs(deltas))) + +# Plot rms bunch sizes +for key in ["rms_x", "rms_y"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(0.0, ax.get_ylim()[1] * 2.0) + ax.set_xlabel("Turn") + ax.set_ylabel(key + " [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + +# Plot rms emittances +color_cycle = plt.rcParams["axes.prop_cycle"].by_key()["color"] + +fig, ax = plt.subplots(figsize=(5, 3)) +for key, color in zip(["eps_x", "eps_y"], color_cycle): + for i, model in enumerate(["envelope", "bunch"]): + if model == "envelope": + plot_kws = dict() + else: + plot_kws = dict(marker=".", lw=0) + label = key + "_" + model + ax.plot(histories[model][key], color=color, label=label, **plot_kws) + +ymax = 2.0 * np.mean(histories["envelope"]["eps_x"]) +ax.set_ylim(0.0, ymax) +ax.set_xlabel("Turn") +ax.set_ylabel("[mm mrad]") +plt.savefig(os.path.join(output_dir, f"fig_emittances")) +plt.close() + +# Plot centroids +for key in ["avg_x", "avg_y"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(-5.0, 5.0) + ax.set_xlabel("Turn") + ax.set_ylabel(key + " [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + +# Collect bunch/envelope data on final turn. +particles = collect_bunch(bunch)["coords"] +particles[:, :4] *= 1000.0 + +env_cov_matrix = envelope.cov_matrix +env_cov_matrix[:4, :4] *= 1000.0**2 + +env_centroid = envelope.centroid +env_centroid[:4] *= 1000.0 + +xmax = 4.0 * np.std(particles, axis=0) +limits = list(zip(-xmax, xmax)) +labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]", "z [m]", "dE [GeV]"] + +# Plot x-x' +fig, ax = plt.subplots(figsize=(4, 4)) +ax.hist2d(particles[:, 0], particles[:, 1], bins=100, range=[limits[0], limits[1]]) +plot_rms_ellipse( + env_cov_matrix[0:2, 0:2], + center=(env_centroid[0], env_centroid[1]), + level=2.0, + color="red", + ax=ax, +) +ax.set_xlabel(labels[0]) +ax.set_ylabel(labels[1]) +plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) +plt.close() + +# Plot corner +fig, axs = plot_corner( + particles, + limits=limits, + bins=100, + labels=labels, +) +for i in range(6): + for j in range(i): + env_cov_matrix_proj = project_cov_matrix(env_cov_matrix, axis=(j, i)) + plot_rms_ellipse( + env_cov_matrix_proj, + center=(env_centroid[j], env_centroid[i]), + level=2.0, + color="red", + ax=axs[i, j], + ) +plt.savefig(os.path.join(output_dir, "fig_dist_corner")) +plt.close() diff --git a/examples/Envelope/sns_ring/test_sns_ring_speed.py b/examples/Envelope/sns_ring/test_sns_ring_speed.py new file mode 100644 index 00000000..73b84ff0 --- /dev/null +++ b/examples/Envelope/sns_ring/test_sns_ring_speed.py @@ -0,0 +1,143 @@ +"""Test envelope tracker speed in SNS ring.""" + +import argparse +import time +import cProfile +import pstats +import sys + +import numpy as np +from tqdm import trange + +from orbit.core.bunch import Bunch +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.envelope import Envelope +from orbit.envelope import EnvelopeTracker +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.space_charge.sc2p5d import setSC2p5DAccNodes +from orbit.teapot import TEAPOT_Ring +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + +sys.path.append("..") +from utils import gen_dist + +parser = argparse.ArgumentParser() +parser.add_argument("--bunch-length", type=float, default=120.0) +parser.add_argument("--kin-energy", type=float, default=1.300) +parser.add_argument("--intensity", type=float, default=2e14) + +parser.add_argument("--nparts", type=int, default=10_000) +parser.add_argument("--turns", type=int, default=25) +parser.add_argument("--sc", type=int, default=0) +parser.add_argument("--sc-grid", type=int, default=64) +args = parser.parse_args() + + +lattice = TEAPOT_Ring() +lattice.readMADX("inputs/sns_ring.lat", "rnginjsol") +lattice.initialize() + +for node in lattice.getNodes(): + try: + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + except: + pass + +for node in lattice.getNodes(): + max_length = 1.0 + if node.getLength() > max_length: + node.setnParts(1 + int(node.getLength() / max_length)) + +bunch = Bunch() +bunch.mass(mass_proton) +sync_part = bunch.getSyncParticle() +sync_part.kinEnergy(args.kin_energy) + +matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) +matrix_lattice_params = matrix_lattice.getRingParametersDict() +alpha_x = matrix_lattice_params["alpha x"] +alpha_y = matrix_lattice_params["alpha y"] +beta_x = matrix_lattice_params["beta x [m]"] +beta_y = matrix_lattice_params["beta y [m]"] +eps_x = 25.0e-06 +eps_y = eps_x + +cov_matrix = np.zeros((6, 6)) +cov_matrix[0, 0] = eps_x * beta_x +cov_matrix[2, 2] = eps_y * beta_y +cov_matrix[0, 1] = cov_matrix[1, 0] = -eps_x * alpha_x +cov_matrix[2, 3] = cov_matrix[3, 2] = -eps_y * alpha_y +cov_matrix[1, 1] = eps_x * (1.0 + alpha_x**2) / beta_x +cov_matrix[3, 3] = eps_y * (1.0 + alpha_y**2) / beta_y +cov_matrix[4, 4] = (args.bunch_length / 4.0) ** 2 +cov_matrix[5, 5] = 0.0 + +cov_matrix_init = np.copy(cov_matrix) + +# Track envelope +print("ENVELOPE") + +envelope = Envelope( + bunch=bunch, + cov_matrix=cov_matrix_init, + intensity=args.intensity, +) +tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) + +start_time = time.time() + +profiler = cProfile.Profile() +profiler.enable() + +for turn in trange(args.turns): + tracker.track_ring(envelope) + +time_per_turn = (time.time() - start_time) / args.turns + +profiler.disable() + +print("Time per turn:", time_per_turn) + +profiler_stats = pstats.Stats(profiler) +profiler_stats.sort_stats(pstats.SortKey.TIME) +profiler_stats.print_stats(10) + +# Track bunch +print("BUNCH") + +rng = np.random.default_rng() + +bunch_coords = np.zeros((args.nparts, 6)) +bunch_coords[:, :4] = gen_dist(size=args.nparts, cov_matrix=cov_matrix_init[0:4, 0:4], name="kv") +bunch_coords[:, 4] = args.bunch_length * rng.uniform(-0.5, 0.5, size=args.nparts) + +for i in range(bunch_coords.shape[0]): + bunch.addParticle(*bunch_coords[i]) + +if args.sc: + sc_calc = SpaceChargeCalc2p5D(64, 64, 1) + sc_path_length_min = 1.00e-06 + sc_nodes = setSC2p5DAccNodes(lattice, sc_path_length_min, sc_calc) + + bunch_size = bunch.getSizeGlobal() + bunch.macroSize(args.intensity / bunch_size) + +start_time = time.time() + +profiler = cProfile.Profile() +profiler.enable() + +for turn in trange(args.turns): + lattice.trackBunch(bunch) + +time_per_turn = (time.time() - start_time) / args.turns + +profiler.disable() + +print("Time per turn:", time_per_turn) + +profiler_stats = pstats.Stats(profiler) +profiler_stats.sort_stats(pstats.SortKey.TIME) +profiler_stats.print_stats(10) diff --git a/examples/Envelope/style.mplstyle b/examples/Envelope/style.mplstyle new file mode 100644 index 00000000..71f36605 --- /dev/null +++ b/examples/Envelope/style.mplstyle @@ -0,0 +1,8 @@ +axes.linewidth: 1.25 +axes.titlesize: "medium" +image.cmap: "Greys" +figure.constrained_layout.use: True +savefig.dpi: 300 +savefig.format: "png" +xtick.minor.visible: True +ytick.minor.visible: True diff --git a/examples/Envelope/test_env_2d_fodo.py b/examples/Envelope/test_env_2d_fodo.py new file mode 100644 index 00000000..de5a4444 --- /dev/null +++ b/examples/Envelope/test_env_2d_fodo.py @@ -0,0 +1,316 @@ +"""Test 2D envelope tracker in FODO lattice.""" + +import argparse +import copy +import math +import os +import pathlib + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.bunch_utils import collect_bunch +from orbit.envelope import Envelope +from orbit.envelope import EnvelopeTracker +from orbit.lattice import AccLattice +from orbit.lattice import AccNode +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.space_charge.sc2p5d import setSC2p5DAccNodes +from orbit.teapot import DriftTEAPOT +from orbit.teapot import QuadTEAPOT +from orbit.teapot import TEAPOT_Lattice +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + +from plot import plot_rms_ellipse +from plot import plot_corner +from utils import gen_dist +from utils import build_rotation_matrix_xy +from utils import project_cov_matrix + +plt.style.use("style.mplstyle") + + +def main(args: argparse.Namespace) -> None: + + # Setup + # ------------------------------------------------------------------------------ + + path = pathlib.Path(__file__) + output_dir = os.path.join("outputs", path.stem) + os.makedirs(output_dir, exist_ok=True) + + # Create lattice + # ------------------------------------------------------------------------------ + + nodes = [ + QuadTEAPOT(length=0.5, kq=+args.kq), + DriftTEAPOT(length=1.0), + QuadTEAPOT(length=1.0, kq=-args.kq), + DriftTEAPOT(length=1.0), + QuadTEAPOT(length=0.5, kq=+args.kq), + ] + + lattice = TEAPOT_Lattice() + for node in nodes: + node.setnParts(args.nslice) + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + lattice.addNode(node) + + lattice.initialize() + + # Create envelope + # ------------------------------------------------------------------------------ + + # Create bunch + bunch = Bunch() + bunch.mass(mass_proton) + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(args.kin_energy) + + # Find periodic lattice parameters + matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) + matrix_lattice_params = matrix_lattice.getRingParametersDict() + alpha_x = matrix_lattice_params["alpha x"] + alpha_y = matrix_lattice_params["alpha y"] + beta_x = matrix_lattice_params["beta x [m]"] + beta_y = matrix_lattice_params["beta y [m]"] + eps_x = 0.25e-06 + eps_y = eps_x + + # Generate covariance matrix + cov_matrix = np.zeros((6, 6)) + cov_matrix[0, 0] = eps_x * beta_x + cov_matrix[2, 2] = eps_y * beta_y + cov_matrix[0, 1] = cov_matrix[1, 0] = -eps_x * alpha_x + cov_matrix[2, 3] = cov_matrix[3, 2] = -eps_y * alpha_y + cov_matrix[1, 1] = eps_x * (1.0 + alpha_x**2) / beta_x + cov_matrix[3, 3] = eps_y * (1.0 + alpha_y**2) / beta_y + cov_matrix[4, 4] = args.bunch_length**2 / 12.0 + cov_matrix[5, 5] = 0.0 + + # Tilt + if args.tilt: + rot_matrix = np.identity(6) + rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi)) + cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T]) + + # Mismatch + cov_matrix[0, 0] *= (1.0 + args.mismatch_x) ** 2 + cov_matrix[2, 2] *= (1.0 + args.mismatch_y) ** 2 + cov_matrix_init = np.copy(cov_matrix) + + # Offset + centroid_init = np.zeros(6) + centroid_init[0] += args.offset_x + centroid_init[2] += args.offset_y + + # Create envelope + envelope = Envelope( + bunch=bunch, + cov_matrix=cov_matrix_init, + centroid=centroid_init, + intensity=args.intensity, + ) + + # Track envelope + # ------------------------------------------------------------------------------ + + print("TRACK ENVELOPE") + + tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) + + history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []} + for turn in range(args.turns): + if turn > 0: + tracker.track_ring(envelope) + + cov_matrix = envelope.cov_matrix + centroid = envelope.centroid + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + xavg = 1000.0 * centroid[0] + yavg = 1000.0 * centroid[2] + + print(f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} xavg={xavg:0.3f} yavg={yavg:0.3f}") + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["xavg"].append(xavg) + history["yavg"].append(yavg) + + histories = {} + histories["envelope"] = copy.deepcopy(history) + + # Track bunch + # ------------------------------------------------------------------------------ + + print("TRACK BUNCH") + + rng = np.random.default_rng() + + bunch_coords = np.zeros((args.nparts, 6)) + bunch_coords[:, :4] = gen_dist( + size=args.nparts, cov_matrix=cov_matrix_init[0:4, 0:4], name=args.dist + ) + bunch_coords[:, 4] = args.bunch_length * rng.uniform(-0.5, 0.5, size=args.nparts) + bunch_coords += centroid_init[None, :6] + + for i in range(bunch_coords.shape[0]): + bunch.addParticle(*bunch_coords[i]) + + if args.sc: + sc_calc = SpaceChargeCalc2p5D(128, 128, 1) + sc_path_length_min = 1.00e-06 + sc_nodes = setSC2p5DAccNodes(lattice, sc_path_length_min, sc_calc) + + bunch_size = bunch.getSizeGlobal() + bunch.macroSize(args.intensity / bunch_size) + + history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []} + for turn in range(args.turns): + if turn > 0: + lattice.trackBunch(bunch) + + twiss_calc = BunchTwissAnalysis() + twiss_calc.computeBunchMoments(bunch, 2, 0, 0) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(i + 1): + cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + cov_matrix[j, i] = cov_matrix[i, j] + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + xavg = 1000.0 * twiss_calc.getAverage(0) + yavg = 1000.0 * twiss_calc.getAverage(2) + + print(f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} xavg={xavg:0.3f} yavg={yavg:0.3f}") + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["xavg"].append(xavg) + history["yavg"].append(yavg) + + histories["bunch"] = copy.deepcopy(history) + + # Analysis + # ------------------------------------------------------------------------------ + + for history in histories.values(): + for key in history: + history[key] = np.array(history[key]) + + # Print errors + for key in histories["envelope"]: + deltas = histories["bunch"][key] - histories["envelope"][key] + print("key:", key) + print("max_abs_delta:", np.max(np.abs(deltas))) + print("avg_abs_delta:", np.mean(np.abs(deltas))) + + # Plot rms bunch sizes + for key in ["xrms", "yrms"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(0.0, ax.get_ylim()[1] * 2.0) + ax.set_xlabel("Turn") + ax.set_ylabel("RMS [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + + # Plot centroids + for key in ["xavg", "yavg"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(-5.0, 5.0) + ax.set_xlabel("Turn") + ax.set_ylabel("AVG [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + + # Collect bunch/envelope data on final turn. + particles = collect_bunch(bunch)["coords"] + particles[:, :4] *= 1000.0 + + env_cov_matrix = envelope.cov_matrix + env_cov_matrix[:4, :4] *= 1000.0**2 + + env_centroid = envelope.centroid + env_centroid[:4] *= 1000.0 + + xmax = 4.0 * np.std(particles, axis=0) + limits = list(zip(-xmax, xmax)) + labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]", "z [m]", "dE [GeV]"] + + # Plot x-x' + fig, ax = plt.subplots(figsize=(4, 4)) + ax.hist2d(particles[:, 0], particles[:, 1], bins=100, range=[limits[0], limits[1]]) + plot_rms_ellipse( + env_cov_matrix[0:2, 0:2], + center=(env_centroid[0], env_centroid[1]), + level=2.0, + color="red", + ax=ax, + ) + ax.set_xlabel(labels[0]) + ax.set_ylabel(labels[1]) + plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) + plt.close() + + # Plot corner + fig, axs = plot_corner( + particles, + limits=limits, + bins=100, + labels=labels, + ) + for i in range(6): + for j in range(i): + env_cov_matrix_proj = project_cov_matrix(env_cov_matrix, axis=(j, i)) + plot_rms_ellipse( + env_cov_matrix_proj, + center=(env_centroid[j], env_centroid[i]), + level=2.0, + color="red", + ax=axs[i, j], + ) + plt.savefig(os.path.join(output_dir, "fig_dist_corner")) + plt.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--bunch-length", type=float, default=5.0) + parser.add_argument("--kin-energy", type=float, default=0.0025) + parser.add_argument("--intensity", type=float, default=5e9) + + parser.add_argument("--dist", type=str, default="kv", choices=["kv", "waterbag", "gauss"]) + parser.add_argument("--mismatch-x", type=float, default=0.0) + parser.add_argument("--mismatch-y", type=float, default=0.0) + parser.add_argument("--offset-x", type=float, default=0.0) + parser.add_argument("--offset-y", type=float, default=0.0) + parser.add_argument("--tilt", type=float, default=0) + + parser.add_argument("--nslice", type=int, default=10) + parser.add_argument("--kq", type=float, default=0.25) + + parser.add_argument("--nparts", type=int, default=100_000) + parser.add_argument("--turns", type=int, default=25) + parser.add_argument("--sc", type=int, default=0) + args = parser.parse_args() + + main(args) diff --git a/examples/Envelope/test_env_2d_fodo_speed.py b/examples/Envelope/test_env_2d_fodo_speed.py new file mode 100644 index 00000000..2b85d185 --- /dev/null +++ b/examples/Envelope/test_env_2d_fodo_speed.py @@ -0,0 +1,145 @@ +"""Test envelope tracker speed in SNS ring.""" + +import argparse +import time +import cProfile +import pstats + +import numpy as np +from tqdm import trange + +from orbit.core.bunch import Bunch +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.envelope import Envelope +from orbit.envelope import EnvelopeTracker +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.space_charge.sc2p5d import setSC2p5DAccNodes +from orbit.teapot import QuadTEAPOT +from orbit.teapot import DriftTEAPOT +from orbit.teapot import TEAPOT_Lattice +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + +from utils import gen_dist + +parser = argparse.ArgumentParser() +parser.add_argument("--bunch-length", type=float, default=5.0) +parser.add_argument("--kin-energy", type=float, default=1.300) +parser.add_argument("--intensity", type=float, default=2e14) + +parser.add_argument("--nslice", type=int, default=10) +parser.add_argument("--kq", type=float, default=0.25) + +parser.add_argument("--nparts", type=int, default=10_000) +parser.add_argument("--turns", type=int, default=5000) +parser.add_argument("--sc", type=int, default=0) +parser.add_argument("--sc-grid", type=int, default=64) +args = parser.parse_args() + +nodes = [ + QuadTEAPOT(length=0.5, kq=+args.kq), + DriftTEAPOT(length=1.0), + QuadTEAPOT(length=1.0, kq=-args.kq), + DriftTEAPOT(length=1.0), + QuadTEAPOT(length=0.5, kq=+args.kq), +] + +lattice = TEAPOT_Lattice() +for node in nodes: + node.setnParts(args.nslice) + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + lattice.addNode(node) +lattice.initialize() + +bunch = Bunch() +bunch.mass(mass_proton) +sync_part = bunch.getSyncParticle() +sync_part.kinEnergy(args.kin_energy) + +matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) +matrix_lattice_params = matrix_lattice.getRingParametersDict() +alpha_x = matrix_lattice_params["alpha x"] +alpha_y = matrix_lattice_params["alpha y"] +beta_x = matrix_lattice_params["beta x [m]"] +beta_y = matrix_lattice_params["beta y [m]"] +eps_x = 25.0e-06 +eps_y = eps_x + +cov_matrix = np.zeros((6, 6)) +cov_matrix[0, 0] = eps_x * beta_x +cov_matrix[2, 2] = eps_y * beta_y +cov_matrix[0, 1] = cov_matrix[1, 0] = -eps_x * alpha_x +cov_matrix[2, 3] = cov_matrix[3, 2] = -eps_y * alpha_y +cov_matrix[1, 1] = eps_x * (1.0 + alpha_x**2) / beta_x +cov_matrix[3, 3] = eps_y * (1.0 + alpha_y**2) / beta_y +cov_matrix[4, 4] = args.bunch_length**2 / 12.0 +cov_matrix[5, 5] = 0.0 + +cov_matrix_init = np.copy(cov_matrix) + +# Track envelope +print("ENVELOPE") + +envelope = Envelope( + bunch=bunch, + cov_matrix=cov_matrix_init, + intensity=args.intensity, +) +tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) + +start_time = time.time() + +profiler = cProfile.Profile() +profiler.enable() + +for turn in trange(args.turns): + tracker.track_ring(envelope) + +time_per_turn = (time.time() - start_time) / args.turns + +profiler.disable() + +print("Time per turn:", time_per_turn) + +profiler_stats = pstats.Stats(profiler) +profiler_stats.sort_stats(pstats.SortKey.TIME) +profiler_stats.print_stats(10) + +# Track bunch +print("BUNCH") + +rng = np.random.default_rng() + +bunch_coords = np.zeros((args.nparts, 6)) +bunch_coords[:, :4] = gen_dist(size=args.nparts, cov_matrix=cov_matrix_init[0:4, 0:4], name="kv") +bunch_coords[:, 4] = args.bunch_length * rng.uniform(-0.5, 0.5, size=args.nparts) + +for i in range(bunch_coords.shape[0]): + bunch.addParticle(*bunch_coords[i]) + +if args.sc: + sc_calc = SpaceChargeCalc2p5D(64, 64, 1) + sc_path_length_min = 1.00e-06 + sc_nodes = setSC2p5DAccNodes(lattice, sc_path_length_min, sc_calc) + + bunch_size = bunch.getSizeGlobal() + bunch.macroSize(args.intensity / bunch_size) + +start_time = time.time() + +profiler = cProfile.Profile() +profiler.enable() + +for turn in trange(args.turns): + lattice.trackBunch(bunch) + +time_per_turn = (time.time() - start_time) / args.turns + +profiler.disable() + +print("Time per turn:", time_per_turn) + +profiler_stats = pstats.Stats(profiler) +profiler_stats.sort_stats(pstats.SortKey.TIME) +profiler_stats.print_stats(10) diff --git a/examples/Envelope/test_env_3d_drift.py b/examples/Envelope/test_env_3d_drift.py new file mode 100644 index 00000000..d66e592b --- /dev/null +++ b/examples/Envelope/test_env_3d_drift.py @@ -0,0 +1,275 @@ +"""Test 3D envelope tracker in drift + +The initial beam is a uniform-density ellipsoid in the x-y-z plane, with zero initial velocity. +The ellipsoid can have arbitrary size and orientation. +""" + +import argparse +import copy +import math +import os +import pathlib + +import numpy as np +import matplotlib.pyplot as plt +import scipy.spatial + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.core.spacecharge import SpaceChargeCalc3D +from orbit.bunch_utils import collect_bunch +from orbit.envelope import Envelope +from orbit.envelope import EnvelopeTracker +from orbit.space_charge.sc3d import setSC3DAccNodes +from orbit.teapot import DriftTEAPOT +from orbit.teapot import TEAPOT_Lattice +from orbit.utils.consts import mass_proton + +from plot import plot_rms_ellipse +from plot import plot_corner +from utils import gen_dist +from utils import project_cov_matrix + +plt.style.use("style.mplstyle") + + +def rotation_matrix_3d(angle_x: float, angle_y: float, angle_z: float) -> np.ndarray: + return scipy.spatial.transform.Rotation.from_euler( + "xyz", [angle_x, angle_y, angle_z] + ).as_matrix() + + +def build_cov_matrix_xyz(rms_sizes: np.ndarray, rotation_matrix: np.ndarray = None) -> np.ndarray: + cov_matrix = np.diag(np.square(rms_sizes)) + if rotation_matrix is None: + return cov_matrix + return rotation_matrix @ cov_matrix @ rotation_matrix.T + + +def main(args: argparse.Namespace) -> None: + + # Setup + # ------------------------------------------------------------------------------ + path = pathlib.Path(__file__) + output_dir = os.path.join("outputs", path.stem) + os.makedirs(output_dir, exist_ok=True) + + # Create lattice + # ------------------------------------------------------------------------------ + node = DriftTEAPOT(length=args.length) + node.setLength(args.length) + node.setnParts(args.nslice) + + lattice = TEAPOT_Lattice() + lattice.addNode(node) + lattice.initialize() + + # Create envelope + # ------------------------------------------------------------------------------ + bunch = Bunch() + bunch.mass(mass_proton) + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(args.kin_energy) + + cov_matrix_init = np.zeros((6, 6)) + + rotation_matrix = rotation_matrix_3d( + math.radians(args.rot_x), math.radians(args.rot_y), math.radians(args.rot_z) + ) + print(rotation_matrix) + + cov_matrix_xyz = build_cov_matrix_xyz( + [args.rms_x, args.rms_y, args.rms_z], rotation_matrix=rotation_matrix + ) + + lorentz_matrix = np.diag([1.0, 1.0, 1.0 / sync_part.gamma()]) + cov_matrix_xyz = lorentz_matrix @ cov_matrix_xyz @ lorentz_matrix.T + + spatial_idx = np.ix_([0, 2, 4], [0, 2, 4]) + cov_matrix_init[spatial_idx] = cov_matrix_xyz + + print(cov_matrix_xyz * 1e6) + print() + print(cov_matrix_init * 1e6) + + centroid_init = np.zeros(6) + + envelope = Envelope( + bunch=bunch, + cov_matrix=cov_matrix_init, + centroid=centroid_init, + intensity=args.intensity, + ) + + # Track envelope + # ------------------------------------------------------------------------------ + print("TRACK ENVELOPE") + + tracker = EnvelopeTracker(lattice, sc=("3d" if args.sc else None)) + + history = {"xrms": [], "yrms": [], "zrms": []} + for turn in range(args.turns): + if turn > 0: + tracker.track(envelope) + + cov_matrix = envelope.cov_matrix + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + zrms = 1000.0 * math.sqrt(cov_matrix[4, 4]) * envelope.gamma + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["zrms"].append(zrms) + + print(f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} zrms={zrms:0.3f}") + + histories = {} + histories["envelope"] = copy.deepcopy(history) + + # Track bunch + # ------------------------------------------------------------------------------ + print("TRACK BUNCH") + + bunch_coords = np.zeros((args.nparts, 6)) + bunch_coords[:, (0, 2, 4)] = gen_dist( + size=args.nparts, cov_matrix=cov_matrix_xyz, name="waterbag" + ) + + for x, xp, y, yp, z, dE in bunch_coords: + bunch.addParticle(x, xp, y, yp, z, dE) + + size_global = bunch.getSizeGlobal() + bunch.macroSize(args.intensity / size_global) + + if args.sc: + sc_calc = SpaceChargeCalc3D(args.sc_grid, args.sc_grid, args.sc_grid) + sc_path_length_min = 0.01 + sc_nodes = setSC3DAccNodes(lattice, sc_path_length_min, sc_calc) + + history = {"xrms": [], "yrms": [], "zrms": []} + for turn in range(args.turns): + if turn > 0: + lattice.trackBunch(bunch) + + twiss_calc = BunchTwissAnalysis() + twiss_calc.computeBunchMoments(bunch, 2, 0, 0) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(i + 1): + cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + cov_matrix[j, i] = cov_matrix[i, j] + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + zrms = 1000.0 * math.sqrt(cov_matrix[4, 4]) * bunch.getSyncParticle().gamma() + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["zrms"].append(zrms) + + print(f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} zrms={zrms:0.3f}") + + histories["bunch"] = copy.deepcopy(history) + + # Analysis + # ------------------------------------------------------------------------------ + for history in histories.values(): + for key in history: + history[key] = np.array(history[key]) + + # Print errors + for key in histories["envelope"]: + deltas = histories["bunch"][key] - histories["envelope"][key] + print("key:", key) + print("max_abs_delta:", np.max(np.abs(deltas))) + print("avg_abs_delta:", np.mean(np.abs(deltas))) + + # Plot rms bunch sizes + for key in ["xrms", "yrms", "zrms"]: + fig, ax = plt.subplots(figsize=(4, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(0.0, ax.get_ylim()[1] * 1.2) + ax.set_xlabel("s [mm]") + ax.set_ylabel("rms size [mm]") + ax.legend(loc="upper left") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + + # Collect bunch/envelope data on final turn. + particles = collect_bunch(bunch)["coords"] + particles *= 1e3 + + env_cov_matrix = envelope.cov_matrix + env_cov_matrix *= 1e6 + + env_centroid = envelope.centroid + env_centroid *= 1e3 + + xmax = 4.0 * np.std(particles, axis=0) + limits = list(zip(-xmax, xmax)) + labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]", "z [mm]", "dE [MeV]"] + + # Plot x-x' + fig, ax = plt.subplots(figsize=(4, 4)) + ax.hist2d(particles[:, 0], particles[:, 1], bins=100, range=[limits[0], limits[1]]) + plot_rms_ellipse( + env_cov_matrix[0:2, 0:2], + center=(env_centroid[0], env_centroid[1]), + level=2.0, + color="red", + ax=ax, + ) + ax.set_xlabel(labels[0]) + ax.set_ylabel(labels[1]) + plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) + plt.close() + + # Plot corner + fig, axs = plot_corner( + particles, + limits=limits, + bins=100, + labels=labels, + ) + for i in range(6): + for j in range(i): + env_cov_matrix_proj = project_cov_matrix(env_cov_matrix, axis=(j, i)) + plot_rms_ellipse( + env_cov_matrix_proj, + center=(env_centroid[j], env_centroid[i]), + level=2.0, + color="red", + ax=axs[i, j], + ) + plt.savefig(os.path.join(output_dir, "fig_dist_corner")) + plt.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--kin-energy", type=float, default=0.0025) + parser.add_argument("--intensity", type=float, default=3e10) + + parser.add_argument("--rms-x", type=float, default=0.010) + parser.add_argument("--rms-y", type=float, default=0.010) + parser.add_argument("--rms-z", type=float, default=0.010) + + parser.add_argument("--rot-x", type=float, default=0.0) + parser.add_argument("--rot-y", type=float, default=0.0) + parser.add_argument("--rot-z", type=float, default=0.0) + + parser.add_argument("--nslice", type=int, default=10) + parser.add_argument("--length", type=float, default=0.1) + parser.add_argument("--turns", type=int, default=20) + parser.add_argument("--sc-grid", type=int, default=64) + + parser.add_argument("--nparts", type=int, default=100_000) + parser.add_argument("--sc", type=int, default=0) + args = parser.parse_args() + + main(args) diff --git a/examples/Envelope/utils.py b/examples/Envelope/utils.py new file mode 100644 index 00000000..daf2a6c7 --- /dev/null +++ b/examples/Envelope/utils.py @@ -0,0 +1,60 @@ +import numpy as np + + +def gen_dist_gauss(size: int, dim: np.ndarray) -> np.ndarray: + return np.random.multivariate_normal( + mean=np.zeros(dim), + cov=np.eye(dim), + size=size, + ) + + +def gen_dist_kv(size: int, dim: int) -> np.ndarray: + X = np.random.normal(size=(size, dim)) + X /= np.linalg.norm(X, axis=1)[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist_waterbag(size: int, dim: int) -> np.ndarray: + X = gen_dist_kv(size, dim) + dim = X.shape[1] + r = np.random.uniform(size=size) ** (1.0 / dim) + X *= r[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist(size: int, cov_matrix: np.ndarray, name: str) -> np.ndarray: + dim = cov_matrix.shape[0] + if name == "kv": + X = gen_dist_kv(size, dim) + elif name == "waterbag": + X = gen_dist_waterbag(size, dim) + elif name == "gauss": + X = gen_dist_gauss(size, dim) + else: + raise ValueError(f"Invalid distribution name: {name}") + + L = np.linalg.cholesky(cov_matrix) + return np.matmul(X, L.T) + + +def build_rotation_matrix_xy(angle: float) -> np.ndarray: + cs = np.cos(angle) + sn = np.sin(angle) + + matrix = np.identity(4) + matrix[0, 0] = matrix[1, 1] = +cs + matrix[0, 2] = matrix[1, 3] = +sn + matrix[2, 0] = matrix[3, 1] = -sn + matrix[2, 2] = matrix[3, 3] = +cs + return matrix + + +def project_cov_matrix(cov_matrix: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: + cov_matrix_proj = np.zeros((len(axis), len(axis))) + for i in range(len(axis)): + for j in range(len(axis)): + cov_matrix_proj[i, j] = cov_matrix[axis[i], axis[j]] + return cov_matrix_proj diff --git a/py/orbit/envelope/__init__.py b/py/orbit/envelope/__init__.py new file mode 100644 index 00000000..7a4586ae --- /dev/null +++ b/py/orbit/envelope/__init__.py @@ -0,0 +1,2 @@ +from .envelope import Envelope +from .track import EnvelopeTracker \ No newline at end of file diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py new file mode 100644 index 00000000..2c181626 --- /dev/null +++ b/py/orbit/envelope/envelope.py @@ -0,0 +1,235 @@ +import numpy as np +import scipy.constants +import scipy.special + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.core.bunch import SyncParticle + +from .utils import convert_matrix_zp_to_dE +from .utils import gen_dist +from .utils import get_classical_radius +from .utils import proj_cov_matrix + + +def build_diag_matrix_from_xyz_eig(eigenvectors: np.ndarray) -> np.ndarray: + A = np.eye(7) + for i in range(eigenvectors.shape[0]): + for j in range(eigenvectors.shape[1]): + row = i * 2 + col = j * 2 + A[row, col] = A[row + 1, col + 1] = eigenvectors[i, j] + return A + + +class Envelope: + """Represents beam envelope/centroid. + + Attributes: + bunch: Bunch containing synchronous particle and (optionally) test particles. + cov_matrix: 6 x 6 covariance matrix + centroid: 6 x 1 centroid vector. + intensity: Total number of particles. + """ + + def __init__( + self, + bunch: Bunch, + cov_matrix: np.ndarray = None, + centroid: np.ndarray = None, + intensity: float = 0.0, + ) -> None: + + empty_bunch = Bunch() + bunch.copyEmptyBunchTo(empty_bunch) + + self.bunch = empty_bunch + self.sync_part = self.bunch.getSyncParticle() + + self.centroid = centroid + if self.centroid is None: + if bunch.getSize(): + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + self.centroid = np.zeros(6) + for i in range(6): + self.centroid[i] = twiss_calc.getAverage(i) + else: + self.centroid = np.zeros(6) + + self.cov_matrix = cov_matrix + if self.cov_matrix is None: + if bunch.getSize(): + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + self.cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(6): + self.cov_matrix[i, j] = twiss_calc.getCorrelation(i, j) + self.cov_matrix[j, i] = self.cov_matrix[i, j] + else: + self.cov_matrix = np.eye(6) + + self.intensity = intensity + self.classical_radius = get_classical_radius(self.charge, self.mass) + self.charge_sign = self.charge / abs(self.charge) + + # For a uniform one-dimensional distribution over length L, the standard + # deviation is L * sqrt(12). This quantity is used to calculate the line + # density for two-dimensional space charge kicks. + self.rms_bunch_length_factor = np.sqrt(12.0) + + def copy(self): + return Envelope( + bunch=self.bunch, + cov_matrix=self.cov_matrix, + centroid=self.centroid, + intensity=self.intensity + ) + + @property + def gamma(self) -> float: + return self.sync_part.gamma() + + @property + def beta(self) -> float: + return self.sync_part.beta() + + @property + def mass(self) -> float: + return self.sync_part.mass() + + @property + def charge(self) -> float: + return self.bunch.charge() + + @property + def momentum(self) -> float: + return self.sync_part.momentum() + + @property + def sc_factor(self) -> float: + return 2.0 * self.intensity * self.classical_radius / (self.beta ** 2 * self.gamma ** 3) + + def rms(self, axis: int = None) -> float | np.ndarray: + rms_arr = np.sqrt(np.diag(self.cov_matrix)) + return rms_arr[axis] + + def transform(self, matrix: np.ndarray) -> None: + m = matrix[:-1, :-1] + u = matrix[:-1, -1] + self.cov_matrix = m @ self.cov_matrix @ m.T + self.centroid = np.matmul(m, self.centroid) + u + + def sample(self, size: int, dist: str = "kv") -> np.ndarray: + particles = gen_dist(size=size, cov_matrix=self.cov_matrix, name=dist) + particles = particles + self.centroid + return particles + + def sc_matrix_2d(self, length: float) -> np.ndarray: + centroid = self.centroid + cov_matrix = self.cov_matrix + + # Calculate transfer matrix in normalized (upright) frame. + cov_xx = cov_matrix[0, 0] + cov_yy = cov_matrix[2, 2] + cov_xy = cov_matrix[0, 2] + + phi = -0.5 * np.arctan2(2 * cov_xy, cov_xx - cov_yy) + sin_phi = np.sin(phi) + cos_phi = np.cos(phi) + rx = 2.0 * np.sqrt(abs(cov_xx * cos_phi**2 + cov_yy * sin_phi**2 - 2.0 * cov_xy * sin_phi * cos_phi)) + ry = 2.0 * np.sqrt(abs(cov_xx * sin_phi**2 + cov_yy * cos_phi**2 + 2.0 * cov_xy * sin_phi * cos_phi)) + + bunch_length = self.rms_bunch_length_factor * np.sqrt(cov_matrix[4, 4]) + perveance = self.sc_factor / bunch_length + kappa_factor = 2.0 * perveance / (rx + ry) + + M = np.identity(7) + M[1, 0] = kappa_factor * length / rx + M[3, 2] = kappa_factor * length / ry + + # Build matrix A to transform out of normalized frame. + A = np.eye(7) + A[0, 0] = A[1, 1] = +cos_phi + A[0, 2] = A[1, 3] = +sin_phi + A[2, 0] = A[3, 1] = -sin_phi + A[2, 2] = A[3, 3] = +cos_phi + + A_inv = A.T + + # Build matrix T to shift to beam centroid. + T = np.identity(7) + T[0, -1] = centroid[0] + T[2, -1] = centroid[2] + + T_inv = np.copy(T) + T_inv[:-1, -1] = -T[:-1, -1] + + # Compute transfer matrix in lab frame. + return T @ A @ M @ A_inv @ T_inv + + def sc_matrix_3d(self, length: float) -> np.ndarray: + # Build Lorentz matrix: rest frame to lab frame. + # x -> x + # y -> y + # z -> z / gamma + # x' = dx/ds -> x' * gamma + # y' = dy/ds -> y' * gamma + # z' = dz/ds -> z' + gamma = self.gamma + gamma_inv = 1.0 / gamma + + L = np.identity(7) + L[1, 1] = gamma + L[3, 3] = gamma + L[4, 4] = gamma_inv + L_inv = np.diag(1.0 / np.diag(L)) + + # Get centroid in rest frame. + centroid = np.matmul(L_inv[:-1, :-1], self.centroid) + + # Get covariance matrix in rest frame. + cov_matrix = L_inv[:-1, :-1] @ self.cov_matrix @ L_inv[:-1, :-1].T + + # Project covariance matrix onto x-y-z plane. + cov_matrix_proj = proj_cov_matrix(cov_matrix, axis=(0, 2, 4)) + + # Compute eigenvalues and eigenvectors of x-y-z covariance matrix. + cov_eig_res = np.linalg.eigh(cov_matrix_proj) + cov_eig_vals = cov_eig_res.eigenvalues + cov_eig_vecs = cov_eig_res.eigenvectors + + # Build transfer matrix in upright frame. + cov_xx, cov_yy, cov_zz = cov_eig_vals + RDx = scipy.special.elliprd(cov_yy, cov_zz, cov_xx) + RDy = scipy.special.elliprd(cov_xx, cov_zz, cov_yy) + RDz = scipy.special.elliprd(cov_xx, cov_yy, cov_zz) + + factor = 0.5 * self.sc_factor * ((1.0 / 5.0) ** 1.5) + kappa_x = factor * RDx + kappa_y = factor * RDy + kappa_z = factor * RDz + + M = np.identity(7) + M[1, 0] = kappa_x * length + M[3, 2] = kappa_y * length + M[5, 4] = kappa_z * length + + # Build matrix to undo x-y-z diagonalization. + A = build_diag_matrix_from_xyz_eig(cov_eig_vecs) + A_inv = A.T + + # Build matrix for translation to centroid. + T = np.identity(7) + for i in (0, 2, 4): + T[i, -1] = centroid[i] + + T_inv = np.copy(T) + T_inv[:-1, -1] = -T[:-1, -1] + + # Compute transfer matrix in lab frame. + M = L @ T @ A @ M @ A_inv @ T_inv @ L_inv + + # Convert from z' to dE + return convert_matrix_zp_to_dE(M, self.sync_part) diff --git a/py/orbit/envelope/matrix.py b/py/orbit/envelope/matrix.py new file mode 100644 index 00000000..ee0cebdb --- /dev/null +++ b/py/orbit/envelope/matrix.py @@ -0,0 +1,443 @@ +import math + +import numpy as np + +from orbit.core.bunch import Bunch +from orbit.core.bunch import SyncParticle +from orbit.lattice import AccNode +from orbit.teapot import ApertureTEAPOT +from orbit.teapot import DriftTEAPOT +from orbit.teapot import BendTEAPOT +from orbit.teapot import KickTEAPOT +from orbit.teapot import MonitorTEAPOT +from orbit.teapot import MultipoleTEAPOT +from orbit.teapot import NodeTEAPOT +from orbit.teapot import QuadTEAPOT +from orbit.teapot import SolenoidTEAPOT +from orbit.teapot import FringeFieldTEAPOT +from orbit.teapot import BunchWrapTEAPOT +from orbit.teapot import TiltTEAPOT +from orbit.teapot import ContinuousLinearFocusingTEAPOT +from orbit.teapot import TurnCounterTEAPOT +from orbit.py_linac.lattice import MarkerLinacNode as MarkerLINAC +from orbit.py_linac.lattice import Drift as DriftLINAC +from orbit.py_linac.lattice import Quad as QuadLINAC +from orbit.py_linac.lattice import Bend as BendLINAC +from orbit.py_linac.lattice import DCorrectorH as DCorrectorHLINAC +from orbit.py_linac.lattice import DCorrectorV as DCorrectorVLINAC +from orbit.py_linac.lattice import Solenoid as SolenoidLINAC +from orbit.py_linac.lattice import TiltElement as TiltLINAC +from orbit.py_linac.lattice import FringeField as FringeFieldLINAC +from orbit.py_linac.lattice import BaseRF_Gap as BaseRF_Gap +from orbit.py_linac.lattice import LinacApertureNode as ApertureLINAC +from orbit.utils.consts import speed_of_light + +from .envelope import Envelope +from .utils import get_dp_p_coeff + + +IGNORE_NODE_TYPES = [ + NodeTEAPOT, + MonitorTEAPOT, + FringeFieldTEAPOT, + ApertureTEAPOT, + BunchWrapTEAPOT, + TurnCounterTEAPOT, + MarkerLINAC, + FringeFieldLINAC, +] + + +def get_matrix_tilt(angle: float) -> np.ndarray: + cos_phi = math.cos(angle) + sin_phi = math.sin(angle) + + M = np.identity(7) + M[0, 0] = M[1, 1] = +cos_phi + M[0, 2] = M[1, 3] = -sin_phi + M[2, 0] = M[3, 1] = +sin_phi + M[2, 2] = M[3, 3] = +cos_phi + return M + + +def get_matrix_kick(kx: float = 0.0, ky: float = 0.0, kE: float = 0.0) -> np.ndarray: + M = np.identity(7) + M[1, -1] = kx + M[3, -1] = ky + M[5, -1] = kE + return M + + +def get_matrix_drift(envelope: Envelope, length: float) -> np.ndarray: + sync_part = envelope.sync_part + + M = np.identity(7) + M[0, 1] = length + M[2, 3] = length + M[4, 5] = length / (sync_part.gamma() ** 2) + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_quad(envelope: Envelope, length: float, kq: float) -> np.ndarray: + if abs(kq) == 0: + return get_matrix_drift(envelope=envelope, length=length) + + sync_part = envelope.sync_part + + sqrt_abs_kq = math.sqrt(abs(kq)) + + M = np.identity(7) + if kq > 0: + cx = np.cos(sqrt_abs_kq * length) + sx = np.sin(sqrt_abs_kq * length) + cy = np.cosh(sqrt_abs_kq * length) + sy = np.sinh(sqrt_abs_kq * length) + M[0, 0] = cx + M[0, 1] = +sx / sqrt_abs_kq + M[1, 0] = -sx * sqrt_abs_kq + M[1, 1] = cx + M[2, 2] = cy + M[2, 3] = sy / sqrt_abs_kq + M[3, 2] = sy * sqrt_abs_kq + M[3, 3] = cy + elif kq < 0: + cx = np.cosh(sqrt_abs_kq * length) + sx = np.sinh(sqrt_abs_kq * length) + cy = np.cos(sqrt_abs_kq * length) + sy = np.sin(sqrt_abs_kq * length) + M[0, 0] = cx + M[0, 1] = sx / sqrt_abs_kq + M[1, 0] = sx * sqrt_abs_kq + M[1, 1] = cx + M[2, 2] = cy + M[2, 3] = +sy / sqrt_abs_kq + M[3, 2] = -sy * sqrt_abs_kq + M[3, 3] = cy + + M[4, 5] = length / (sync_part.gamma()**2) + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_bend(envelope: Envelope, length: float, theta: float) -> np.ndarray: + sync_part = envelope.sync_part + + rho = length / theta + cx = math.cos(theta) + sx = math.sin(theta) + + M = np.identity(7) + M[0, 0] = cx + M[0, 1] = rho * sx + M[0, 5] = rho * (1.0 - cx) + M[1, 0] = -sx / rho + M[1, 1] = cx + M[1, 5] = sx + M[2, 3] = length + M[4, 0] = -sx + M[4, 1] = -rho * (1.0 - cx) + M[4, 5] = -(sync_part.beta() ** 2) * length + rho * sx + M[:5, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_solenoid(envelope: Envelope, length: float, B: float) -> np.ndarray: + if B == 0: + return get_matrix_drift(envelope=envelope, length=length) + + sync_part = envelope.sync_part + + phase = B * length + + V = np.identity(7) + V[:4, :4] = 0.0 + V[0, 1] = -1.0 / B + V[0, 2] = 0.5 + V[1, 0] = 0.5 * B + V[1, 3] = 1.0 + V[2, 1] = 1.0 / B + V[2, 2] = 0.5 + V[3, 0] = -0.5 * B + V[3, 3] = 1.0 + + M = np.identity(7) + M[0, 0] = +1.0 + M[1, 1] = -1.0 + M[2, 2] = math.cos(phase) + M[2, 3] = math.sin(phase) / B + M[3, 2] = math.sin(phase) * B * -1.0 + M[3, 3] = math.cos(phase) + M[4, 5] = length / (sync_part.gamma()**2) + + M = np.linalg.inv(V) @ M @ V + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_cf(envelope: Envelope, length: float, kq: float) -> np.ndarray: + if kq == 0: + return get_matrix_drift(envelope=envelope, length=length) + + sync_part = envelope.sync_part + + sqrt_abs_kq = math.sqrt(abs(kq)) + + cx = math.cos(sqrt_abs_kq * length) + sx = math.sin(sqrt_abs_kq * length) + + M = np.identity(7) + M[0, 0] = M[2, 2] = cx + M[0, 1] = M[2, 3] = +sx / sqrt_abs_kq + M[1, 0] = M[3, 2] = -sx * sqrt_abs_kq + M[1, 1] = M[3, 3] = cx + M[4, 5] = length / (sync_part.gamma()**2) + M[4, 5] *= get_dp_p_coeff(sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_rf_gap(envelope: Envelope, frequency: float, E0TL: float, phase: float) -> np.ndarray: + sync_part = envelope.sync_part + + gamma = sync_part.gamma() + beta = sync_part.beta() + mass = sync_part.mass() + charge = envelope.charge + + kin_energy_in = sync_part.kinEnergy() + charge_E0TL_sin = charge * E0TL * math.sin(phase) + kin_energy_delta = charge * E0TL * math.cos(phase) + + # Calculate parameters in the center of the gap. + sync_part.momentum(sync_part.energyToMomentum(kin_energy_in + kin_energy_delta / 2.0)) + gamma_gap = sync_part.gamma() + beta_gap = sync_part.beta() + + # Move to the end of the gap. + kin_energy_out = kin_energy_in + kin_energy_delta + sync_part.momentum(sync_part.energyToMomentum(kin_energy_out)) + + # The base RF gap is simple - no phase correction. + gamma_out = sync_part.gamma() + beta_out = sync_part.beta() + prime_coeff = (beta * gamma) / (beta_out * gamma_out) + + # Wave momentum + k = 2.0 * math.pi * frequency / speed_of_light + phase_time_coeff = k / beta + + # Transverse focusing coefficient + kappa = -charge * E0TL * k / (2.0 * mass * beta_gap**2 * beta_out * gamma_gap**2 * gamma_out) + d_rp = kappa * math.sin(phase) + + M = np.eye(7) + M[5, 4] = charge_E0TL_sin * phase_time_coeff + M[4, 4] = beta_out / beta + M[1, 1] = prime_coeff + M[3, 3] = prime_coeff + M[1, 0] = d_rp + M[3, 2] = d_rp + return M + + +def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.ndarray | None: + """Calculate transfer matrix and update synchronous particle. + + This function maps various accelerator nodes to 7 x 7 transfer matrices + for envelope tracking. For non-accelerating, finite-length nodes, the + synchronous particle time is updated as in a drift. Accelerating nodes + such as RF gaps will update the synchronous particle energy. + + Args: + node: The accelerator node. + envelope: The beam envelope. + part_index: Index of the part within the node. An index of -1 returns + the transfer matrix for the entire node. + Returns: + 7 x 7 transfer matrix or None. If None, the node can be ignored during + envelope tracking. + """ + + node_type = type(node) + if node_type in IGNORE_NODE_TYPES: + return None + + length = node.getLength(part_index) + nparts = node.getnParts() + + if node_type is DriftTEAPOT: + if length <= 0: + return None + return get_matrix_drift(envelope=envelope, length=length) + + elif node_type is SolenoidTEAPOT: + if length <= 0: + return None + + B = node.getParam("B") + if node.waveform: + B *= node.waveform.getStrength() + B *= envelope.charge_sign + + return get_matrix_solenoid(envelope=envelope, length=length, B=B) + + elif node_type is MultipoleTEAPOT: + if length <= 0: + return None + + if np.all(np.abs(node.getParam("kls")) == 0): + return get_matrix_drift(envelope=envelope, length=length) + + elif node_type is QuadTEAPOT: + if length <= 0: + return None + + kq = node.getParam("kq") + if node.waveform: + kq *= node.waveform.getStrength() + kq *= envelope.charge_sign + + return get_matrix_quad(envelope=envelope, length=length, kq=kq) + + elif node_type is BendTEAPOT: + if length <= 0: + return None + + theta = node.getParam("theta") / (nparts - 1) + if part_index == 0 or part_index == nparts - 1: + theta *= 0.5 + theta *= envelope.charge_sign + + return get_matrix_bend(envelope=envelope, length=length, theta=theta) + + elif node_type is KickTEAPOT: + scale = 1.0 + if node.waveform is not None: + scale = node.waveform.getStrength() + + scale /= (nparts - 1) + kx = scale * node.getParam("kx") + ky = scale * node.getParam("ky") + kE = node.getParam("dE") + + if abs(kx) > 0 or abs(ky) > 0 or abs(kE) > 0: + return np.matmul( + get_matrix_kick(kx=kx, ky=ky, kE=kE), + get_matrix_drift(envelope=envelope, length=length), + ) + else: + return get_matrix_drift(envelope=envelope, length=length) + + elif node_type is TiltTEAPOT: + angle = node.getTiltAngle() + if angle == 0: + return None + return get_matrix_tilt(angle) + + elif node_type is ContinuousLinearFocusingTEAPOT: + if length <= 0: + return None + + kq = node.getParam("kq") + kq *= envelope.charge_sign + if node.waveform: + kq *= node.waveform.getStrength() + + return get_matrix_cf(envelope=envelope, length=length, kq=kq) + + elif node_type is DriftLINAC: + if length <= 0: + return None + return get_matrix_drift(envelope=envelope, length=length) + + elif node_type is QuadLINAC: + if length <= 0: + return None + + brho = 3.335640952 * envelope.momentum / envelope.charge + kq = node.getParam("dB/dr") / brho + return get_matrix_quad(envelope=envelope, length=length, kq=kq) + + elif node_type is BendLINAC: + if length <= 0: + return None + + theta = node.getParam("theta") / (nparts - 1) + if part_index == 0 or part_index == nparts - 1: + theta *= 0.5 + theta *= envelope.charge_sign + + return get_matrix_bend(envelope=envelope, length=length, theta=theta) + + elif node_type is DCorrectorHLINAC: + length = node.getParam("effLength") / nparts + field = node.getParam("B") + delta_xp = -field * envelope.charge * length * 0.299792 / envelope.momentum + if delta_xp == 0: + return None + return get_matrix_kick(kx=delta_xp, ky=0.0, kE=0.0) + + elif node_type is DCorrectorVLINAC: + length = node.getParam("effLength") / nparts + field = node.getParam("B") + delta_yp = -field * envelope.charge * length * 0.299792 / envelope.momentum + if delta_yp == 0: + return None + return get_matrix_kick(kx=0.0, ky=delta_yp, kE=0.0) + + elif node_type is SolenoidLINAC: + if length <= 0: + return None + B = node.getParam("B") * envelope.charge_sign + return get_matrix_solenoid(envelope=envelope, length=length, B=B) + + elif node_type is TiltLINAC: + angle = node.getTiltAngle() + if angle == 0: + return None + return get_matrix_tilt(angle=angle) + + elif node_type is BaseRF_Gap: + E0TL = node.getParam("E0TL") + mode_phase = node.getParam("mode") * math.pi + + cavity = node.getRF_Cavity() + frequency = cavity.getFrequency() + phase = cavity.getPhase() + mode_phase + amplitude = cavity.getAmp() + + sync_part = envelope.sync_part + arrival_time = sync_part.time() + arrival_time_design = cavity.getDesignArrivalTime() + + if node.isFirstRFGap(): + if cavity.isDesignSetUp(): + phase = math.fmod(frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase, 2.0 * math.pi) + else: + orbitFinalize("Run `trackDesign` first to initialize cavity phases.") + else: + phase = math.fmod(frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase,2.0 * math.pi) + + node.setGapPhase(phase) + + if amplitude == 0.0: + return None + + return get_matrix_rf_gap( + envelope=envelope, + frequency=frequency, + E0TL=(E0TL * amplitude), + phase=phase, + ) + + raise NotImplementedError(str(node)) \ No newline at end of file diff --git a/py/orbit/envelope/meson.build b/py/orbit/envelope/meson.build new file mode 100644 index 00000000..90eab844 --- /dev/null +++ b/py/orbit/envelope/meson.build @@ -0,0 +1,14 @@ +py_sources = files([ + '__init__.py', + 'envelope.py', + 'matrix.py', + 'track.py', + 'utils.py' +]) + +python.install_sources( + py_sources, + subdir: 'orbit/envelope', + # pure: true, +) + diff --git a/py/orbit/envelope/track.py b/py/orbit/envelope/track.py new file mode 100644 index 00000000..98653fe9 --- /dev/null +++ b/py/orbit/envelope/track.py @@ -0,0 +1,232 @@ +import numpy as np +import warnings + +from orbit.core.bunch import Bunch +from orbit.core.bunch import SyncParticle +from orbit.lattice import AccNode +from orbit.lattice import AccLattice +from orbit.teapot import BendTEAPOT +from orbit.py_linac.lattice import Bend as BendLINAC + +from .matrix import get_matrix +from .envelope import Envelope + + +ENTRANCE = AccNode.ENTRANCE +BODY = AccNode.BODY +EXIT = AccNode.EXIT + +BEFORE = AccNode.BEFORE +AFTER = AccNode.AFTER + + +class EnvelopeTracker: + def __init__(self, lattice: AccLattice, sc: str | None = None) -> None: + """Constructor. + + Args: + lattice: The accelerator lattice. + sc: Envelope space charge model {"2d", "3d", None}. + """ + self.lattice = lattice + self.sc = sc + + # For pre-computing elements + self.elements = [] + self.one_turn_matrix = None + # [TO DO] option to return one-turn matrix including linear space charge + + for node in self.lattice.getNodes(): + if type(node) in (BendTEAPOT, BendLINAC): + if node.getParam("ea1") != 0.0 or node.getParam("ea2") != 0.0: + message = f"Found bend ea1 or ea2 != 0.0 ({node.getName()}.)" + message += " Nonzero edge angles are not yet supported in envelope tracking." + message += " Setting ea1 and ea2 to 0.0." + warnings.warn(message) + + node.setParam("ea1", 0.0) + node.setParam("ea2", 0.0) + + def track(self, envelope: Envelope) -> None: + """Track envelope through lattice. + + This is not recursive, so grandchild nodes are not tracked. + """ + for node_index, node in enumerate(self.lattice.getNodes()): + for child_node in node.getChildNodes(ENTRANCE): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + envelope.transform(matrix) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + envelope.transform(matrix) + + matrix_sc = None + if self.sc: + length = node.getLength(part_index) + if length > 0: + if self.sc == "2d": + matrix_sc = envelope.sc_matrix_2d(length) + elif self.sc == "3d": + matrix_sc = envelope.sc_matrix_3d(length) + else: + raise ValueError + + matrix = get_matrix(node, envelope=envelope, part_index=part_index) + if matrix is not None: + if matrix_sc is not None: + matrix = matrix @ matrix_sc + envelope.transform(matrix) + + for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + envelope.transform(matrix) + + for child_node in node.getChildNodes(EXIT): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + envelope.transform(matrix) + + def track_history(self, envelope: Envelope) -> dict[str, list]: + """Track and return envelope parameters vs. position in lattice.""" + history = {} + history["position"] = [] + history["rms_x"] = [] + history["rms_y"] = [] + history["rms_z"] = [] + history["kin_energy"] = [] + + node_positions = self.lattice.getNodePositionsDict() + + history["position"].append(0.0) + history["rms_x"].append(1000.0 * envelope.rms(0)) + history["rms_y"].append(1000.0 * envelope.rms(2)) + history["rms_z"].append(1000.0 * envelope.rms(4)) + history["kin_energy"].append(envelope.sync_part.kinEnergy()) + + for node_index, node in enumerate(self.lattice.getNodes()): + for child_node in node.getChildNodes(ENTRANCE): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + envelope.transform(matrix) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + envelope.transform(matrix) + + matrix_sc = None + if self.sc: + length = node.getLength(part_index) + if length > 0: + if self.sc == "2d": + matrix_sc = envelope.sc_matrix_2d(length) + elif self.sc == "3d": + matrix_sc = envelope.sc_matrix_3d(length) + else: + raise ValueError + + matrix = get_matrix(node, envelope=envelope, part_index=part_index) + if matrix is not None: + if matrix_sc is not None: + matrix = matrix @ matrix_sc + envelope.transform(matrix) + + position_start, position_stop = node_positions[node] + position = position_start + node.getLength(part_index) * (part_index + 1) + + history["position"].append(position) + history["rms_x"].append(1000.0 * envelope.rms(0)) + history["rms_y"].append(1000.0 * envelope.rms(2)) + history["rms_z"].append(1000.0 * envelope.rms(4)) + history["kin_energy"].append(envelope.sync_part.kinEnergy()) + + for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + envelope.transform(matrix) + + for child_node in node.getChildNodes(EXIT): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + envelope.transform(matrix) + + return history + + def precompute_matrices(self, envelope: Envelope) -> None: + """Pre-compute transfer matrices for each node. + + For each node, return tuple (node, matrix). Mark space charge kicks as ("sc", length). + """ + self.elements = [] + for node_index, node in enumerate(self.lattice.getNodes()): + for child_node in node.getChildNodes(ENTRANCE): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + self.elements.append((child_node, matrix)) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + self.elements.append((child_node, matrix)) + + if self.sc: + length = node.getLength(part_index) + if length > 0: + self.elements.append(("sc", length)) + + matrix = get_matrix(node, envelope=envelope, part_index=part_index) + if matrix is not None: + self.elements.append((node, matrix)) + + for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + self.elements.append((node, matrix)) + + for child_node in node.getChildNodes(EXIT): + matrix = get_matrix(child_node, envelope=envelope) + if matrix is not None: + self.elements.append((node, matrix)) + + def track_ring(self, envelope: Envelope) -> None: + """Track using pre-computed transfer matrices. + + The method assumes that all nodes are static and that there is no + change in the synchronous particle energy. In this case the matrices + can be computed once and reused on each turn. If there is no space charge, + we track using the one-turn matrix. + """ + + # Pre-compute transfer matrices on the first turn. + if not self.elements: + self.precompute_matrices(envelope) + self.one_turn_matrix = None + + # If there is no space charge, apply the one-turn transfer matrix. + if not self.sc: + if self.one_turn_matrix is None: + self.one_turn_matrix = np.identity(7) + for (node, matrix) in self.elements: + self.one_turn_matrix = np.matmul(self.one_turn_matrix, matrix) + return envelope.transform(self.one_turn_matrix) + + # If there is space charge, apply the matrices one-by-one. + for element in self.elements: + if element[0] == "sc": + length = element[1] + if self.sc == "2d": + envelope.transform(envelope.sc_matrix_2d(length)) + elif self.sc == "3d": + envelope.transform(envelope.sc_matrix_3d(length)) + else: + raise ValueError + else: + node, matrix = element + envelope.transform(matrix) diff --git a/py/orbit/envelope/utils.py b/py/orbit/envelope/utils.py new file mode 100644 index 00000000..71499e31 --- /dev/null +++ b/py/orbit/envelope/utils.py @@ -0,0 +1,99 @@ +import math +import numpy as np +from scipy.constants import epsilon_0 + +from orbit.core.bunch import SyncParticle +from orbit.utils.consts import charge_electron + + +def get_classical_radius(charge: float, mass: float) -> float: + q = charge * charge_electron # [C] + rest_energy = mass * 1e9 * charge_electron # [J] + return q**2 / (4.0 * math.pi * epsilon_0 * rest_energy) + + +def get_dp_p_coeff(sync_part: SyncParticle) -> float: + # dE/E = (beta^2) * dp/p + # dE = (beta^2 * E) * dp/p + # dE = (beta^2 * gamma * m * c^2) * dp/p + beta = sync_part.beta() + gamma = sync_part.gamma() + rest_energy = sync_part.mass() # GeV + return 1.0 / (beta**2 * gamma * rest_energy) + + +def get_zp_coeff(sync_part: SyncParticle) -> float: + # dE/E = (beta^2) * dp/p = (beta^2) * (gamma^2) z' + # dE = (beta^2 * gamma^2 * E) * z' + # dE = (beta^2 * gamma^3 * m * c^2) * z' + beta = sync_part.beta() + gamma = sync_part.gamma() + rest_energy = sync_part.mass() + return 1.0 / (beta**2 * gamma**3 * rest_energy) + + +def convert_matrix_dp_p_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: + # v = [x, x', y, y', z, dp/p] + # w = [x, x', y, y', z, dE] + # v = A w + # v -> M v + # w -> A M A^-1 + dp_p_coeff = get_dp_p_coeff(sync_part) + matrix[:5, 5] *= dp_p_coeff + matrix[5, :5] /= dp_p_coeff + matrix[5, 6] /= dp_p_coeff # driving term + return matrix + + +def convert_matrix_zp_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: + zp_coeff = get_zp_coeff(sync_part) + matrix[:5, 5] *= zp_coeff + matrix[5, :5] /= zp_coeff + matrix[5, 6] /= zp_coeff # driving term + return matrix + + +def gen_dist_gauss(size: int, cov_matrix: np.ndarray) -> np.ndarray: + return np.random.multivariate_normal( + mean=np.zeros(cov_matrix.shape[0]), + cov=cov_matrix, + size=size, + ) + + +def gen_dist_kv(size: int, cov_matrix: np.ndarray) -> np.ndarray: + X = np.random.normal(size=(size, cov_matrix.shape[0])) + X /= np.linalg.norm(X, axis=1)[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist_waterbag(size: int, cov_matrix: np.ndarray) -> np.ndarray: + X = gen_dist_kv(size, cov_matrix) + dim = X.shape[1] + r = np.random.uniform(size=size) ** (1.0 / dim) + X *= r[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist(size: int, cov_matrix: np.ndarray, name: str) -> np.ndarray: + if name == "kv": + X = gen_dist_kv(size, cov_matrix) + elif name == "waterbag": + X = gen_dist_waterbag(size, cov_matrix) + elif name == "gauss": + X = gen_dist_gauss(size, cov_matrix) + else: + raise ValueError(f"Invalid distribution name: {name}") + + L = np.linalg.cholesky(cov_matrix) + return np.matmul(X, L.T) + + +def proj_cov_matrix(cov_matrix: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: + cov_matrix_proj = np.zeros((len(axis), len(axis))) + for i in range(len(axis)): + for j in range(len(axis)): + cov_matrix_proj[i, j] = cov_matrix[axis[i], axis[j]] + return cov_matrix_proj diff --git a/py/orbit/meson.build b/py/orbit/meson.build index 2ba16e73..651b14d5 100644 --- a/py/orbit/meson.build +++ b/py/orbit/meson.build @@ -24,6 +24,7 @@ subdir('space_charge') subdir('errors') subdir('matrix_lattice') subdir('teapot') +subdir('envelope') py_sources = files([ diff --git a/py/orbit/py_linac/lattice/__init__.py b/py/orbit/py_linac/lattice/__init__.py index 85a9df80..eb9ed985 100644 --- a/py/orbit/py_linac/lattice/__init__.py +++ b/py/orbit/py_linac/lattice/__init__.py @@ -13,6 +13,8 @@ from orbit.py_linac.lattice.LinacAccNodes import Solenoid from orbit.py_linac.lattice.LinacAccNodes import DCorrectorH, DCorrectorV from orbit.py_linac.lattice.LinacAccNodes import ThickKick +from orbit.py_linac.lattice.LinacAccNodes import TiltElement +from orbit.py_linac.lattice.LinacAccNodes import FringeField from orbit.py_linac.lattice.LinacRfGapNodes import BaseRF_Gap, AxisFieldRF_Gap, RF_AxisFieldsStore @@ -56,7 +58,6 @@ __all__.append("Bend") __all__.append("Solenoid") - __all__.append("LinacApertureNode") __all__.append("CircleLinacApertureNode") __all__.append("EllipseLinacApertureNode") @@ -64,7 +65,6 @@ __all__.append("LinacPhaseApertureNode") __all__.append("LinacEnergyApertureNode") - __all__.append("RF_Cavity") __all__.append("Sequence") @@ -91,3 +91,6 @@ __all__.append("LinacTrMatricesController") __all__.append("LinacBPM") + +__all__.append("FringeField") +__all__.append("TiltElement") diff --git a/py/orbit/teapot/__init__.py b/py/orbit/teapot/__init__.py index b77d9cc8..9f62adfc 100644 --- a/py/orbit/teapot/__init__.py +++ b/py/orbit/teapot/__init__.py @@ -18,6 +18,9 @@ from .teapot import TiltTEAPOT from .teapot import NodeTEAPOT from .teapot import ContinuousLinearFocusingTEAPOT +from .teapot import ApertureTEAPOT +from .teapot import MonitorTEAPOT +from .teapot import TurnCounterTEAPOT from .teapot import TPB @@ -41,3 +44,6 @@ __all__.append("ContinuousLinearFocusingTEAPOT") __all__.append("TPB") __all__.append("TEAPOT_MATRIX_Lattice") +__all__.append("ApertureTEAPOT") +__all__.append("MonitorTEAPOT") +__all__.append("TurnCounterTEAPOT") \ No newline at end of file diff --git a/tests/py/orbit/test_env.py b/tests/py/orbit/test_env.py new file mode 100644 index 00000000..16f1e5c7 --- /dev/null +++ b/tests/py/orbit/test_env.py @@ -0,0 +1,383 @@ +import numpy as np + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.core.linac import MatrixRfGap +from orbit.bunch_utils import collect_bunch +from orbit.envelope import Envelope +from orbit.envelope import EnvelopeTracker +from orbit.lattice import AccNode +from orbit.lattice import AccLattice +from orbit.py_linac.lattice import Drift +from orbit.py_linac.lattice import Quad +from orbit.py_linac.lattice import Bend +from orbit.py_linac.lattice import TiltElement +from orbit.py_linac.lattice import Solenoid +from orbit.teapot import BendTEAPOT +from orbit.teapot import ContinuousLinearFocusingTEAPOT +from orbit.teapot import DriftTEAPOT +from orbit.teapot import KickTEAPOT +from orbit.teapot import QuadTEAPOT +from orbit.teapot import SolenoidTEAPOT +from orbit.teapot import TiltTEAPOT +from orbit.teapot import TEAPOT_Lattice +from orbit.utils.consts import mass_proton + + +def get_lorentz_factors(kin_energy: float, mass: float) -> tuple[float, float]: + gamma = 1.0 + kin_energy / mass + beta = np.sqrt(1.0 - (1.0 / gamma) ** 2) + return (gamma, beta) + + +def calc_bunch_cov(bunch: Bunch) -> np.ndarray: + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(i + 1): + cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + cov_matrix[j, i] = cov_matrix[i, j] + return cov_matrix + + +def make_lattice(nodes: list[AccNode]) -> AccLattice: + lattice = TEAPOT_Lattice() + for node in nodes: + lattice.addNode(node) + + lattice.initialize() + + for node in lattice.getNodes(): + try: + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + except: + pass + + return lattice + + +def track_and_compare_rms( + lattice: AccLattice, + kin_energy: float, + cov_matrix: np.ndarray, + nparts: int = 100_000, + verbose: int = 1, +) -> dict: + """Track bunch/envelope and compare rms beam sizes. + + Args: + lattice: Accelerator lattice. + kin_energy: Synchronous particle kinetic energy [GeV]. + cov_matrix: 6 x 6 covariance matrix. + nparts: Number of particles in bunch. + verbose: Whether to print results. + """ + cov_scale = 1e6 + + data = {} + for k1 in ["env", "bunch"]: + data[k1] = {} + for k2 in ["rms", "cov"]: + data[k1][k2] = {} + for k3 in ["env", "bunch"]: + data[k1][k2][k3] = {} + + # Initialize bunch + bunch = Bunch() + bunch.mass(mass_proton) + bunch.getSyncParticle().kinEnergy(kin_energy) + + # Track bunch + particles = np.random.multivariate_normal(np.zeros(6), cov_matrix, size=nparts) + for x in particles: + bunch.addParticle(*x) + + # Covariance matrix calculated from particles will be slightly different. + cov_matrix = calc_bunch_cov(bunch) + + data["bunch"]["cov"]["in"] = cov_scale * calc_bunch_cov(bunch) + lattice.trackBunch(bunch) + data["bunch"]["cov"]["out"] = cov_scale * calc_bunch_cov(bunch) + + # Track envelope + envelope = Envelope(bunch=bunch, cov_matrix=cov_matrix) + envelope_tracker = EnvelopeTracker(lattice=lattice) + + data["env"]["cov"]["in"] = cov_scale * envelope.cov_matrix + envelope_tracker.track(envelope) + data["env"]["cov"]["out"] = cov_scale * envelope.cov_matrix + + # Compare + for mode in ["env", "bunch"]: + for loc in ["in", "out"]: + data[mode]["rms"][loc] = np.sqrt(np.diag(data[mode]["cov"][loc])) + + if verbose: + dims = ["x", "xp", "y", "yp", "z", "dE"] + for key in ["in", "out"]: + print(key.upper()) + for i in range(6): + print(" rms {}:".format(dims[i])) + print(" env: {}".format(data["env"]["rms"][key][i])) + print(" bunch: {}".format(data["bunch"]["rms"][key][i])) + + assert np.all(np.isclose(data["env"]["cov"]["in"], data["bunch"]["cov"]["in"])) + + atol = np.ones(6) + atol[0:4] = 1e-3 # [mm mrad] + atol[4] = 1e-3 # [mm] + atol[5] = 1e-3 # [MeV] + + for i in range(6): + x = data["env"]["rms"]["out"][i] + y = data["bunch"]["rms"]["out"][i] + assert np.abs(x - y) <= atol[i] + + +def make_default_cov_matrix( + rms_x: float = 1e-3, + rms_xp: float = 1e-3, + rms_y: float = 1e-3, + rms_yp: float = 1e-3, + rms_z: float = 1e-3, + rms_dE: float = 1e-5, +) -> np.ndarray: + return np.diag(np.square([rms_x, rms_xp, rms_y, rms_yp, rms_z, rms_dE])) + + +def test_drift_teapot( + kin_energy: float = 0.0025, + length: float = 1.0, + cov_matrix: np.ndarray = None, + nparts: int = 6, +) -> None: + node = DriftTEAPOT(length=length, nparts=nparts) + lattice = make_lattice([node]) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_drift_linac( + kin_energy: float = 0.0025, + length: float = 1.0, + cov_matrix: np.ndarray = None, + nparts: int = 6, +) -> None: + node = Drift() + node.setLength(length) + node.setnParts(nparts) + nodes = [node] + + lattice = make_lattice(nodes) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_quad_teapot( + kin_energy: float = 0.0025, + length: float = 1.0, + kq: float = 1.0, + cov_matrix: np.ndarray = None, + nparts: int = 10, +) -> None: + node = QuadTEAPOT(length=length, kq=kq, nparts=nparts) + lattice = make_lattice([node]) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_cf_teapot( + kin_energy: float = 0.0025, + length: float = 10.0, + kq: float = 1.0, + nparts: int = 10, +) -> None: + node = ContinuousLinearFocusingTEAPOT(length=length, kq=kq, nparts=nparts) + lattice = make_lattice([node]) + cov_matrix = np.diag(np.square([1e-3, 0, 1e-3, 0.0, 0.0, 0.0])) + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_quad_linac( + kin_energy: float = 0.0025, + length: float = 1.0, + field_grad: float = 0.23, + cov_matrix: np.ndarray = None, + nparts: int = 10, +) -> None: + node = Quad() + node.setLength(length) + node.setnParts(nparts) + node.setParam("dB/dr", field_grad) + nodes = [node] + + lattice = make_lattice(nodes) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_bend_teapot( + kin_energy: float = 0.0025, + length: float = 1.0, + theta: float = 20.0, + cov_matrix: np.ndarray = None, + nparts: int = 5, +) -> None: + node = BendTEAPOT(length=length, theta=np.radians(theta), nparts=nparts) + lattice = make_lattice([node]) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_bend_linac( + kin_energy: float = 0.0025, + length: float = 1.0, + theta: float = 20.0, + cov_matrix: np.ndarray = None, + nparts: int = 5, +) -> None: + node = Bend() + node.setLength(length) + node.setnParts(nparts) + node.setParam("theta", np.radians(theta)) + nodes = [node] + + lattice = make_lattice(nodes) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_kick_teapot( + kin_energy: float = 0.0025, + length: float = 0.1, + kx: float = 0.001, + ky: float = 0.001, + dE: float = 0.00001, + cov_matrix: np.ndarray = None, + nparts: int = 4, +) -> None: + node = KickTEAPOT(kx=kx, ky=ky, dE=dE, length=length, nparts=nparts) + lattice = make_lattice([node]) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_tilt_teapot( + kin_energy: float = 0.0025, + angle: float = 0.25 * np.pi, + cov_matrix: np.ndarray = None, +) -> None: + node = TiltTEAPOT(angle=angle) + lattice = make_lattice([node]) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_tilt_linac( + kin_energy: float = 0.0025, + angle: float = 0.25 * np.pi, + cov_matrix: np.ndarray = None, +) -> None: + node = TiltElement() + node.setTiltAngle(angle) + nodes = [node] + lattice = make_lattice(nodes) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_solenoid_teapot( + kin_energy: float = 0.0025, + length: float = 2.0, + B: float = 1.0, + cov_matrix: np.ndarray = None, + nparts: int = 10, +) -> None: + node = SolenoidTEAPOT(length=length, B=B, nparts=nparts) + lattice = make_lattice([node]) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_solenoid_linac( + kin_energy: float = 0.0025, + length: float = 2.0, + B: float = 1.0, + cov_matrix: np.ndarray = None, + nparts: int = 10, +) -> None: + node = Solenoid() + node.setLength(length) + node.setnParts(nparts) + node.setParam("B", B) + nodes = [node] + + lattice = make_lattice(nodes) + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy, cov_matrix) + + +def test_rf_gap_matrix( + kin_energy: float = 0.0025, + frequency: float = 402.5e06, + E0TL: float = 0.001, + phase: float = 0.0, + charge: float = -1.0, +) -> None: + cov_matrix = make_default_cov_matrix() + + bunch_in = Bunch() + bunch_in.mass(mass_proton) + bunch_in.charge(charge) + bunch_in.getSyncParticle().kinEnergy(kin_energy) + + coords_in = np.random.multivariate_normal(np.zeros(6), cov_matrix, size=10) + for x in coords_in: + bunch_in.addParticle(*x) + + bunch_out_1 = Bunch() + bunch_in.copyBunchTo(bunch_out_1) + + matrix_rf_gap = MatrixRfGap() + matrix_rf_gap.trackBunch(bunch_out_1, frequency, E0TL, phase) + + coords_out_1 = collect_bunch(bunch_out_1)["coords"] + + from orbit.envelope.matrix import get_matrix_rf_gap + + bunch_out_2 = Bunch() + bunch_in.copyBunchTo(bunch_out_2) + + envelope = Envelope(bunch=bunch_in) + matrix = get_matrix_rf_gap( + envelope=envelope, + frequency=frequency, + E0TL=E0TL, + phase=phase, + ) + coords_in = np.column_stack([coords_in, np.ones(coords_in.shape[0])]) + coords_out_2 = np.matmul(coords_in, matrix.T) + coords_out_2 = coords_out_2[:, :-1] + assert np.allclose(coords_out_1, coords_out_2) + + +def test_sc_3d_cold_expansion(): + # This should test expansion of cold uniform-density sphere + # (in rest frame). We can calculate the time to expand to + # twice the initial size. (See examples from A. Shishlo or + # from the ImpactX repo.) + pass