def _get_rf_unique_attr(ring, attr, cavpts=None): freq = numpy.unique(_get_rf_attr(ring, attr, cavpts=cavpts)) if len(freq) == 0: raise AtError('No cavity found in the lattice') elif len(freq) > 1: raise AtError('{0} not equal for all cavities'.format(attr)) else: return freq[0]
def get_timelag_fromU0(ring, method=ELossMethod.INTEGRAL, cavpts=None): """ Get the TimeLag attribute of RF cavities based on frequency, voltage and energy loss per turn, so that the synchronous phase is zero. An error occurs if all cavities do not have the same frequency. Used in set_cavity_phase() PARAMETERS ring lattice description KEYWORDS method=ELossMethod.INTEGRAL method for energy loss computation. See "get_energy_loss". cavpts=None Cavity location. If None, use all cavities. This allows to ignore harmonic cavities. RETURN timelag """ if cavpts is None: cavpts = get_cells(ring, checktype(RFCavity)) u0 = get_energy_loss(ring, method=method) try: rfv = ring.get_rf_voltage(cavpts=cavpts) freq = ring.get_rf_frequency(cavpts=cavpts) tl0 = ring.get_rf_timelag(cavpts=cavpts) except AtError: freq = numpy.array([cav.Frequency for cav in ring.select(cavpts)]) rfv = numpy.array([cav.Voltage for cav in ring.select(cavpts)]) tl0 = numpy.array([cav.TimeLag for cav in ring.select(cavpts)]) if u0 > numpy.sum(rfv): raise AtError('Not enough RF voltage: unstable ring') ctmax = 1 / numpy.amin(freq) * clight / 2 tt0 = tl0[numpy.argmin(freq)] eq = lambda x: numpy.sum(rfv * numpy.sin(2 * pi * freq * (x + tl0) / clight)) - u0 deq = lambda x: numpy.sum(rfv * numpy.cos(2 * pi * freq * (x + tl0) / clight)) zero_diff = least_squares(deq, -tt0 + ctmax / 2, bounds=(-tt0, ctmax - tt0)).x[0] if numpy.sign(deq(zero_diff - 1.0e-6)) > 0: ts = least_squares(eq, (zero_diff - tt0) / 2, bounds=(-tt0, zero_diff)).x[0] else: ts = least_squares(eq, (ctmax - tt0 + zero_diff) / 2, bounds=(zero_diff, ctmax - tt0)).x[0] timelag = ts + tl0 else: if u0 > rfv: raise AtError('Not enough RF voltage: unstable ring') timelag = clight / (2 * pi * freq) * numpy.arcsin(u0 / rfv) ts = timelag - tl0 timelag *= numpy.ones(refpts_len(ring, cavpts)) return timelag, ts
def reswall(self, wcomp, length, rvac, conduct, beta, yokoya_factor=1): if wcomp is WakeComponent.Z: raise AtError('Resitive wall not available ' 'for WakeComponent: {}'.format(wcomp)) elif wcomp is (WakeComponent.DX or wcomp is WakeComponent.DY or wcomp is WakeComponent.QX or wcomp is WakeComponent.QY): return wake_functions.transverse_reswall(self._srange, yokoya_factor, length, rvac, conduct, beta) else: raise AtError('Invalid WakeComponent: {}'.format(wcomp))
def __init__(self, *args, **kwargs): """Lattice constructor""" iterator = kwargs.pop('iterator', None) scan = kwargs.pop('scan', False) for key in list(k for k in kwargs.keys() if k.startswith('_')): kwargs.pop(key) if iterator is None: elem_list, = args or [[]] # accept 0 or 1 argument if isinstance(elem_list, Lattice): elems = lattice_filter(kwargs, elem_list) elif scan: elems = params_filter(kwargs, type_filter, elem_list) else: elems = type_filter(kwargs, elem_list) else: elems = iterator(kwargs, *args) super(Lattice, self).__init__(elems) # set default values kwargs.setdefault('name', '') kwargs.setdefault('periodicity', 1) if 'energy' not in kwargs: raise AtError('Lattice energy is not defined') # set attributes self.update(kwargs)
def _orbit6(ring, cavpts=None, guess=None, keep_lattice=False, **kwargs): """Solver for 6D motion""" convergence = kwargs.pop('convergence', DConstant.OrbConvergence) max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter) xy_step = kwargs.pop('XYStep', DConstant.XYStep) dp_step = kwargs.pop('DPStep', DConstant.DPStep) method = kwargs.pop('method', ELossMethod.TRACKING) # Get revolution period l0 = get_s_pos(ring, len(ring)) cavities = [elm for elm in ring if isinstance(elm, elements.RFCavity)] if len(cavities) == 0: raise AtError('No cavity found in the lattice.') f_rf = cavities[0].Frequency harm_number = cavities[0].HarmNumber if guess is None: _, dt = get_timelag_fromU0(ring, method=method, cavpts=cavpts) ref_in = numpy.zeros((6,), order='F') ref_in[5] = -dt else: ref_in = numpy.copy(guess) theta = numpy.zeros((6,)) theta[5] = constants.speed_of_light * harm_number / f_rf - l0 scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0]) + \ dp_step * numpy.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0]) delta_matrix = numpy.asfortranarray( numpy.concatenate((numpy.diag(scaling), numpy.zeros((6, 1))), axis=1)) id6 = numpy.asfortranarray(numpy.identity(6)) change = 1 itercount = 0 while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice) # the reference particle after one turn ref_out = in_mat[:, 6] # 6x6 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j6 = (in_mat[:, :6] - in_mat[:, 6:]) / scaling a = j6 - id6 # f'(r_n) - 1 b = ref_out[:] - ref_in[:] - theta # b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1) b_over_a = numpy.linalg.solve(a, b) r_next = ref_in - b_over_a # determine if we are close enough change = numpy.linalg.norm(r_next - ref_in) itercount += 1 ref_in = r_next keep_lattice = True if itercount == max_iterations: warnings.warn(AtWarning('Maximum number of iterations reached. ' 'Possible non-convergence')) return ref_in
def a_matrix(tt): """A, eigval = a_matrix(T) Find the A matrix from one turn map matrix T such that: [Rotx 0 0 ] inv(A)*T*A=[ 0 Rotz 0 ] [ 0 0 Rots] Order it so that it is close to the order of x,y,z also ensure that positive modes are before negative so that one has proper symplecticity B. Nash July 18, 2013 INPUT T (m, m) transfer matrix for 1 turn OUTPUT A (m, m) A-matrix eigval (m,) vector of Eigen values of T """ nv = tt.shape[0] dms = int(nv / 2) jmt = jmatswap(dms) select = numpy.arange(0, nv, 2) rbase = numpy.stack((select, select), axis=1).flatten() # noinspection PyTupleAssignmentBalance lmbd, vv = eig(tt) # Compute the norms vp = numpy.dot(vv.conj().T, jmt) n = -0.5j * numpy.sum(vp.T * vv, axis=0) if any(abs(n) < 1.0E-12): raise AtError('Unstable ring') # Move positive before negatives order = rbase + (n < 0) vv = vv[:, order] n = n[order] lmbd = lmbd[order] # Normalize vectors vn = vv / numpy.sqrt(abs(n)).reshape((1, nv)) # find the vectors that project most onto x,y,z, and reorder # nn will have structure # n1x n1y n1z # n2x n2y n2z # n3x n3y n3z nn = 0.5 * abs(numpy.sqrt(-1.j * vn.conj().T.dot(jmt).dot(_vxyz[dms - 1]))) rows = list(select) order = [] for ixz in select: ind = numpy.argmax(nn[rows, ixz]) order.append(rows[ind]) del rows[ind] v_ordered = vn[:, order] lmbd = lmbd[order] aa = numpy.vstack((numpy.real(v_ordered), numpy.imag(v_ordered))).reshape( (nv, nv), order='F') return aa, lmbd
def set_cavity(ring, Voltage=None, Frequency=None, HarmNumber=None, TimeLag=None, cavpts=None, copy=False): """ PARAMETERS ring lattice description KEYWORDS Frequency=None RF frequency. The special enum value Frf.NOMINAL sets the frequency to the nominal value, given ring length and harmonic number. Voltage=None RF voltage for the full ring. HarmNumber=None Harmonic number for the full ring. TimeLag=None cavpts=None Cavity location. If None, use all cavities. This allows to ignore harmonic cavities. copy=False If True, returns a shallow copy of ring with new cavity elements. Otherwise, modify ring in-place """ def get_harm(hcell): if hcell is not None: return hcell * n_periods else: return ring.get_rf_harmonic_number(cavpts=cavpts) if cavpts is None: cavpts = get_cells(ring, checktype(elements.RFCavity)) n_cavities = ring.refcount(cavpts) if n_cavities < 1: raise AtError('No cavity found in the lattice') n_periods = ring.periodicity modif = {} if Frequency is not None: if Frequency is Frf.NOMINAL: Frequency = (clight / ring.circumference) * get_harm(HarmNumber) modif['Frequency'] = Frequency if HarmNumber is not None: modif['HarmNumber'] = HarmNumber / n_periods if TimeLag is not None: modif['TimeLag'] = TimeLag if Voltage is not None: modif['Voltage'] = Voltage / n_periods / n_cavities # noinspection PyShadowingNames @make_copy(copy) def apply(ring, cavpts, modif): for cavity in ring.select(cavpts): cavity.update(modif) return apply(ring, cavpts, modif)
def add(self, wtype, wcomp, *args, **kwargs): if wtype is WakeType.FILE: w = self.readwakefile(*args, **kwargs) elif wtype is WakeType.TABLE: w = self.resample(*args) elif wtype is WakeType.RESONATOR: w = self.resonator(wcomp, *args, **kwargs) elif wtype is WakeType.RESWALL: w = self.reswall(wcomp, *args, **kwargs) else: raise AtError('Invalid WakeType: {}'.format(wtype)) if self.components[wcomp] is None: self.components[wcomp] = w else: self.components[wcomp] += w
def resonator(self, wcomp, frequency, qfactor, rshunt, beta, yokoya_factor=1): if wcomp is WakeComponent.Z: return wake_functions.long_resonator(self._srange, frequency, qfactor, rshunt, beta) elif (wcomp is WakeComponent.DX or wcomp is WakeComponent.DY or wcomp is WakeComponent.QX or wcomp is WakeComponent.QY): return wake_functions.transverse_resonator(self._srange, frequency, qfactor, rshunt, yokoya_factor, beta) else: raise AtError('Invalid WakeComponent: {}'.format(wcomp))
def matfile_generator(params, mat_file): """ run through matlab cells and generate AT elements KEYWORDS mat_file name of the .mat file mat_key name of the Matlab variable containing the lattice. Default: Matlab variable name if there is only one, otherwise 'RING' check=True if False, skip the coherence tests quiet=False If True, suppress the warning for non-standard classes """ def mclean(data): if data.dtype.type is numpy.str_: # Convert strings in arrays back to strings. return str(data[0]) elif data.size == 1: v = data[0, 0] if issubclass(v.dtype.type, numpy.void): for f in v.dtype.fields: v[f] = mclean(v[f]) # Return a scalar return v else: # Remove any surplus dimensions in arrays. return numpy.squeeze(data) m = scipy.io.loadmat(params.setdefault('mat_file', mat_file)) matvars = [varname for varname in m if not varname.startswith('__')] default_key = matvars[0] if (len(matvars) == 1) else 'RING' key = params.setdefault('mat_key', default_key) if key not in m.keys(): kok = [k for k in m.keys() if '__' not in k] raise AtError('Selected mat_key does not exist, ' 'please select in: {}'.format(kok)) check = params.pop('check', True) quiet = params.pop('quiet', False) cell_array = m[key].flat for index, mat_elem in enumerate(cell_array): elem = mat_elem[0, 0] kwargs = {f: mclean(elem[f]) for f in elem.dtype.fields} yield element_from_dict(kwargs, index=index, check=check, quiet=quiet)
def _orbit6(ring, cavpts=None, guess=None, keep_lattice=False, **kwargs): """Solver for 6D motion""" def iscavity(elem): return isinstance(elem, elements.RFCavity) and \ elem.PassMethod.endswith('CavityPass') convergence = kwargs.pop('convergence', DConstant.OrbConvergence) max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter) xy_step = kwargs.pop('XYStep', DConstant.XYStep) dp_step = kwargs.pop('DPStep', DConstant.DPStep) method = kwargs.pop('method', ELossMethod.TRACKING) l0 = get_s_pos(ring, len(ring))[0] # Get the main RF frequency (the lowest) try: f_rf = min(elm.Frequency for elm in ring if iscavity(elm)) except ValueError: raise AtError('No cavity found in the lattice.') # gamma = self.energy / self.particle.mass # beta = math.sqrt(1.0 - 1.0 / gamma / gamma) # h = round(fmin*l0/beta/clight) harm_number = round(f_rf * l0 / clight) if guess is None: _, dt = get_timelag_fromU0(ring, method=method, cavpts=cavpts) # Getting timelag by tracking uses a different lattice, # so we cannot now use the same one again. if method is ELossMethod.TRACKING: keep_lattice = False ref_in = numpy.zeros((6, ), order='F') ref_in[5] = -dt else: ref_in = numpy.copy(guess) theta = numpy.zeros((6, )) theta[5] = clight * harm_number / f_rf - l0 scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0]) + \ dp_step * numpy.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0]) delta_matrix = numpy.asfortranarray( numpy.concatenate((numpy.diag(scaling), numpy.zeros((6, 1))), axis=1)) id6 = numpy.asfortranarray(numpy.identity(6)) change = 1 itercount = 0 while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice) # the reference particle after one turn ref_out = in_mat[:, 6] # 6x6 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j6 = (in_mat[:, :6] - in_mat[:, 6:]) / scaling a = j6 - id6 # f'(r_n) - 1 b = ref_out[:] - ref_in[:] - theta # b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1) b_over_a = numpy.linalg.solve(a, b) r_next = ref_in - b_over_a # determine if we are close enough change = numpy.linalg.norm(r_next - ref_in) itercount += 1 ref_in = r_next keep_lattice = True if itercount == max_iterations: warnings.warn( AtWarning('Maximum number of iterations reached. ' 'Possible non-convergence')) return ref_in
def get_energy_loss(ring, method=ELossMethod.INTEGRAL): """Compute the energy loss per turn [eV] PARAMETERS ring lattice description KEYWORDS method=ELossMethod.INTEGRAL method for energy loss computation The enum class ELossMethod declares 2 values INTEGRAL: The losses are obtained from Losses = Cgamma / 2pi * EGeV^4 * i2 Takes into account bending magnets and wigglers. TRACKING: The losses are obtained by tracking without cavities. Needs radiation ON, takes into account all radiating elements. """ def integral(ring): """Losses = Cgamma / 2pi * EGeV^4 * i2 """ def wiggler_i2(wiggler): rhoinv = wiggler.Bmax / Brho coefh = wiggler.By[1, :] coefv = wiggler.Bx[1, :] return wiggler.Length * (numpy.sum(coefh * coefh) + numpy.sum(coefv * coefv)) * rhoinv**2 / 2 def dipole_i2(dipole): return dipole.BendingAngle**2 / dipole.Length Brho = sqrt(ring.energy**2 - e_mass**2) / clight i2 = 0.0 for el in ring: if isinstance(el, Dipole): i2 += dipole_i2(el) elif isinstance(el, Wiggler) and el.PassMethod != 'DriftPass': i2 += wiggler_i2(el) e_loss = Cgamma / 2.0 / pi * ring.energy**4 * i2 return e_loss @check_radiation(True) def tracking(ring): """Losses from tracking """ ringtmp = ring.radiation_off(dipole_pass=None, quadrupole_pass=None, wiggler_pass=None, sextupole_pass=None, octupole_pass=None, copy=True) o0 = numpy.zeros(6) o6 = numpy.squeeze(lattice_pass(ringtmp, o0, refpts=len(ringtmp))) return -o6[4] * ring.energy if isinstance(method, str): method = ELossMethod[method.upper()] warn(FutureWarning('You should use {0!s}'.format(method))) if method is ELossMethod.INTEGRAL: return ring.periodicity * integral(ring) elif method == ELossMethod.TRACKING: return ring.periodicity * tracking(ring) else: raise AtError('Invalid method: {}'.format(method))
def find_orbit6(ring, refpts=None, guess=None, **kwargs): """find_orbit6 finds the closed orbit in the full 6-D phase space by numerically solving for a fixed point of the one turn map M calculated with lattice_pass (X, PX, Y, PY, DP, CT2 ) = M (X, PX, Y, PY, DP, CT1) with constraint CT2 - CT1 = C*HarmNumber(1/Frf - 1/Frf0) IMPORTANT!!! find_orbit6 is a realistic simulation 1. The Frf frequency in the RF cavities (may be different from Frf0) imposes the synchronous condition CT2 - CT1 = C*HarmNumber(1/Frf - 1/Frf0) 2. The algorithm numerically calculates 6-by-6 Jacobian matrix J6. In order for (J-E) matrix to be non-singular it is NECESSARY to use a realistic PassMethod for cavities with non-zero momentum kick (such as ThinCavityPass). 3. find_orbit6 can find orbits with radiation. In order for the solution to exist the cavity must supply adequate energy compensation. In the simplest case of a single cavity, it must have 'Voltage' field set so that Voltage > Erad - energy loss per turn 4. There is a family of solutions that correspond to different RF buckets They differ in the 6-th coordinate by C*Nb/Frf. Nb = 1 .. HarmNum-1 5. The value of the 6-th coordinate found at the cavity gives the equilibrium RF phase. If there is no radiation the phase is 0; PARAMETERS ring Sequence of AT elements refpts elements at which data is returned. It can be: 1) an integer in the range [-len(ring), len(ring)-1] selecting the element according to python indexing rules. As a special case, len(ring) is allowed and refers to the end of the last element, 2) an ordered list of such integers without duplicates, 3) a numpy array of booleans of maximum length len(ring)+1, where selected elements are True. Defaults to None, if refpts is None an empty array is returned for orbit. OUTPUT orbit0 ((6,) closed orbit vector at the entrance of the 1-st element (x,px,y,py) orbit (6, Nrefs) closed orbit vector at each location specified in refpts KEYWORDS keep_lattice Assume no lattice change since the previous tracking. Default: False guess Initial value for the closed orbit. It may help convergence. Default: (0, 0, 0, 0) convergence Convergence criterion. Default: 1.e-12 max_iterations Maximum number of iterations. Default: 20 step_size Step size. Default: 1.e-6 See also find_orbit4, find_sync_orbit. """ keep_lattice = kwargs.pop('keep_lattice', False) convergence = kwargs.pop('convergence', CONVERGENCE) max_iterations = kwargs.pop('max_iterations', MAX_ITERATIONS) step_size = kwargs.pop('step_size', STEP_SIZE) ref_in = numpy.zeros((6, ), order='F') if guess is None else guess # Get evolution period l0 = get_s_pos(ring, len(ring)) cavities = [elem for elem in ring if isinstance(elem, elements.RFCavity)] if len(cavities) == 0: raise AtError('No cavity found in the lattice.') f_rf = cavities[0].Frequency harm_number = cavities[0].HarmNumber theta = numpy.zeros((6, )) theta[5] = constants.speed_of_light * harm_number / f_rf - l0 delta_matrix = numpy.zeros((6, 7), order='F') for i in range(6): delta_matrix[i, i] = step_size id6 = numpy.asfortranarray(numpy.identity(6)) change = 1 itercount = 0 while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice) # the reference particle after one turn ref_out = in_mat[:, 6] # 6x6 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j6 = (in_mat[:, :6] - in_mat[:, 6:]) / step_size a = j6 - id6 # f'(r_n) - 1 b = ref_out[:] - ref_in[:] - theta b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1) r_next = ref_in - b_over_a # determine if we are close enough change = numpy.linalg.norm(r_next - ref_in) itercount += 1 ref_in = r_next keep_lattice = True if itercount == max_iterations: warnings.warn( AtWarning('Maximum number of iterations reached. ' 'Possible non-convergence')) uint32refs = uint32_refpts(refpts, len(ring)) # bug in numpy < 1.13 all_points = numpy.empty( (0, 6), dtype=float) if len(uint32refs) == 0 else numpy.squeeze( lattice_pass(ring, ref_in.copy(order='K'), refpts=uint32refs, keep_lattice=keep_lattice), axis=(1, 3)).T return ref_in, all_points