diff --git a/tracking/optics.py b/tracking/optics.py index 58049c53d30e36c3307eae3baf3d35c5a1a9d456..67fad0fc30351183337be5dfdcfc6fd0526f0796 100644 --- a/tracking/optics.py +++ b/tracking/optics.py @@ -7,7 +7,6 @@ parameters are defined. @date: 10/04/2020 """ -import at import numpy as np import matplotlib.pyplot as plt from scipy.interpolate import interp1d @@ -100,22 +99,28 @@ class Optics: Lattice periodicity, if not specified the AT lattice periodicity is used. """ + import at self.n_points = int(kwargs.get("n_points", 1e3)) periodicity = kwargs.get("periodicity") self.lattice = at.load_lattice(lattice_file) + if self.lattice.radiation: + self.lattice.radiation_off() lattice = self.lattice.slice(slices=self.n_points) refpts = np.arange(0, len(lattice)) - twiss0, tune, chrom, twiss = at.get_twiss(lattice, refpts=refpts, + twiss0, tune, chrom, twiss = at.linopt(lattice, refpts=refpts, get_chrom=True) if periodicity is None: self.periodicity = lattice.periodicity else: self.periodicity = periodicity - - for i in range(self.periodicity-1): - pos = np.append(twiss.s_pos, twiss.s_pos + twiss.s_pos[-1]*(i+1)) + + if self.periodicity > 1: + for i in range(self.periodicity-1): + pos = np.append(twiss.s_pos, twiss.s_pos + twiss.s_pos[-1]*(i+1)) + else: + pos = twiss.s_pos self.position = pos self.beta_array = np.tile(twiss.beta.T, self.periodicity)