diff --git a/mbtrack2/tracking/synchrotron.py b/mbtrack2/tracking/synchrotron.py index 41458eff38988e909c1f3c90b2d4863c0a3e12ae..87717d06aec2dbed49728178571168faa5bc4bc0 100644 --- a/mbtrack2/tracking/synchrotron.py +++ b/mbtrack2/tracking/synchrotron.py @@ -373,23 +373,78 @@ class Synchrotron: tuneS = phi / (2 * np.pi) return tuneS - def get_adts(self): + def get_adts(self, + xm=1e-4, + ym=1e-4, + npoints=9, + plot=False, + ax=None, + **kwargs): """ Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar componenet from AT lattice. + + Parameters + ---------- + xm : float, optional + Maximum horizontal amplitude in [m]. + Default is 1e-4. + ym : float, optional + Maximum vertical amplitude in [m]. + Default is 1e-4. + npoints : int, optional + Number of points in each plane. + Default is 9. + plot : bool, optional + If True, plot the computed tune shift with amplitude. + Default is False. + ax : array of shape (2,) of matplotlib.axes.Axes, optional + Axes where to plot the figures. + Default is None. + + Returns + ------- + ax : array of shape (2,) of matplotlib.axes.Axes + Only if plot is True. + + See at.physics.nonlinear.detuning for **kwargs """ import at if self.optics.use_local_values: raise ValueError("ADTS needs to be provided manualy as no AT" + " lattice file is loaded.") - - det = at.physics.nonlinear.gen_detuning_elem(self.optics.lattice) - coef_xx = np.array([det.A1 / 2, 0]) - coef_yx = np.array([det.A2 / 2, 0]) - coef_xy = np.array([det.A2 / 2, 0]) - coef_yy = np.array([det.A3 / 2, 0]) + r0, r1, x, q_dx, y, q_dy = at.physics.nonlinear.detuning( + self.optics.lattice.radiation_off(copy=True), + npoints=npoints, + xm=xm, + ym=ym, + **kwargs) + coef_xx = np.array([r1[0][0] / 2, 0]) + coef_yx = np.array([r1[0][1] / 2, 0]) + coef_xy = np.array([r1[0][1] / 2, 0]) + coef_yy = np.array([r1[1][1] / 2, 0]) self.adts = [coef_xx, coef_yx, coef_xy, coef_yy] + if plot: + if ax is None: + fig, ax = plt.subplots(2, 1) + + ax[0].scatter(x * 1e6, q_dx[:, 0]) + ax[0].scatter(y * 1e6, q_dy[:, 0]) + ax[0].set_ylabel("Delta Tune X") + ax[0].set_xlabel("Amplitude [um]") + ax[0].legend(['dqx/dx', 'dqx/dy']) + ax[0].grid() + + ax[1].scatter(x * 1e6, q_dx[:, 1]) + ax[1].scatter(y * 1e6, q_dy[:, 1]) + ax[1].set_ylabel("Delta Tune Y") + ax[1].set_xlabel("Amplitude [um]") + ax[1].legend(['dqy/dx', 'dqy/dy']) + ax[1].grid() + + return ax + def get_chroma(self, order=4, dpm=0.02, n_points=100): """ Compute chromaticity (linear and nonlinear) from AT lattice and update the property. diff --git a/tests/unit/tracking/test_synchrotron.py b/tests/unit/tracking/test_synchrotron.py index fa0d42e7d1bb0071bb22ad752d0d510e8f839f82..daf706ca590a2c2a81bed09572d102d2460b58d3 100644 --- a/tests/unit/tracking/test_synchrotron.py +++ b/tests/unit/tracking/test_synchrotron.py @@ -109,6 +109,10 @@ class TestSynchrotron: ring_with_at_lattice.get_adts() assert ring_with_at_lattice.adts is not None + def test_get_adts_plot(self, ring_with_at_lattice): + axes = ring_with_at_lattice.get_adts(plot=True) + assert axes is not None + def test_get_chroma(self, ring_with_at_lattice): ring_with_at_lattice.get_chroma() assert len(ring_with_at_lattice.chro) == 8