def bump_4cors(lat: MagneticLattice, cor_list: List[Element], marker: Union[Marker, Monitor], x: float = 0.0001, xp: float = 0., energy: float = 14.) -> np.array: """ Bump with 4 correctors. Function calculates correctors angles (strength) to make closed bump with (x, x') or (y, y') at the marker position. Calculated angles are applied to corresponding correctors (MagneticLattice is mutated). :param lat: MagneticLattice :param cor_list: list of 4 correctors, All for correctors must be in the same plane :param marker: Marker or Monitor. It must be between 2d and 3d corrector :param x: x or y beam coordinate at Marker position [m] :param xp: x' or y' beam coordinate at Marker position [rad] :param energy: Beam energy [GeV] :return: numpy array of corrector angles in [rad] """ # lattice analysis, calculate R matrices between different devices if len(cor_list) != 4: raise TypeError("It must be 4 correctors") if len([cor for cor in cor_list if isinstance(cor, Hcor)]) == 4: x_idx, xp_idx = 0, 1 elif len([cor for cor in cor_list if isinstance(cor, Vcor)]) == 4: x_idx, xp_idx = 2, 3 else: raise TypeError("Not all correctors belong to Y- or X-plane") Rs = [] sorted_list = [cor_list[0], cor_list[1], marker, cor_list[2], cor_list[3]] for i in range(4): sequence = copy.deepcopy(lat.get_sequence_part(start=sorted_list[i], stop=sorted_list[i + 1])) sequence[0].l /= 2. sequence[-1].l /= 2. lat1 = MagneticLattice(sequence) R1 = lattice_transfer_map(lat1, energy) Rs.append(R1) R21 = np.dot(Rs[1], Rs[0]) M = np.array([[R21[x_idx, xp_idx], Rs[1][x_idx, xp_idx]], [R21[xp_idx, xp_idx], Rs[1][xp_idx, xp_idx]]]) Minv = np.linalg.inv(M) X = np.array([x, xp]) cor_strengths = np.dot(Minv, X) angle_1, angle_2 = cor_strengths R43 = np.dot(Rs[3], Rs[2]) Xf = np.array([0., 0.]) M43 = np.array([[R43[x_idx, x_idx], R43[x_idx, xp_idx]], [R43[xp_idx, x_idx], R43[xp_idx, xp_idx]]]) X4 = np.dot(M43, X) angle_3 = (Xf[0] - X4[0]) / Rs[3][x_idx, xp_idx] angle_4 = -(X4[1] + Rs[3][xp_idx, xp_idx] * angle_3 - Xf[1]) a = np.array([angle_1, angle_2, angle_3, angle_4]) for i, cor in enumerate(cor_list): cor.angle = a[i] lat.update_transfer_maps() return a
def closed_orbit(lattice, eps_xy=1.e-7, eps_angle=1.e-7, energy=0): __author__ = 'Sergey Tomin' """ Searching of initial coordinates (p0) by iteration method. For initial conditions p uses exact solution of equation p = M*p + B :param lattice: class MagneticLattice :param eps_xy: tolerance on coordinates of beam in the start and end of lattice :param eps_angle: tolerance on the angles of beam in the start and end of lattice :return: class Particle """ R = lattice_transfer_map(lattice, energy) smult = SecondOrderMult() ME = np.eye(4) - R[:4, :4] P = np.dot(np.linalg.inv(ME), lattice.B[:4]) def errf(x): X = np.array([[x[0]], [x[1]], [x[2]], [x[3]], [0], [0]]) smult.numpy_apply(X, R, lattice.T) X += lattice.B err = np.sum([1000 * (X[i, 0] - x[i])**2 for i in range(4)]) return err res = fmin(errf, P, xtol=1e-8, maxiter=2e3, maxfun=2.e3) return Particle(x=res[0], px=res[1], y=res[2], py=res[3])
def merger(lat, remaining_types=[], remaining_elems=[], init_energy=0.): """ Function to compress the lattice excluding elements by type or by individual elements :param lat: MagneticLattice :param remaining_types: list, the type of the elements which needed to be untouched others will be "compress" to Matrix element e.g. [Monitor, Quadrupole, Bend, Hcor] :param remaining_elems: list of elements (ids (str) or object) :param init_energy: initial energy :return: New MagneticLattice """ _logger.debug("element numbers before: " + str(len(lat.sequence))) lattice_analysis = [] merged_elems = [] for elem in lat.sequence: if elem.__class__ in remaining_types or elem.id in remaining_elems or elem in remaining_elems: lattice_analysis.append(merged_elems) merged_elems = [] lattice_analysis.append([elem]) elif elem.__class__ == Edge and ((Bend in remaining_types) or (SBend in remaining_types) or (RBend in remaining_types)): continue else: merged_elems.append(elem) if len(merged_elems) != 0: lattice_analysis.append(merged_elems) seq = [] E = init_energy for elem_list in lattice_analysis: if len(elem_list) == 1: E += elem_list[0].transfer_map.delta_e seq.append(elem_list[0]) elif len(elem_list) == 0: continue else: delta_e = np.sum([elem.transfer_map.delta_e for elem in elem_list]) lattice = MagneticLattice(elem_list, method=lat.method) R = lattice_transfer_map(lattice, energy=E) m = Matrix() m.r = lattice.R m.t = lattice.T m.b = lattice.B m.l = lattice.totalLen m.delta_e = delta_e E += delta_e seq.append(m) new_lat = MagneticLattice(seq, method=lat.method) _logger.debug("element numbers after: " + str(len(new_lat.sequence))) return new_lat
def error_func(x): for i in range(len(varz)): if varz[i].__class__ == Quadrupole: varz[i].k1 = x[i] varz[i].transfer_map = lat.method.create_tm( varz[i]) # create_transfer_map(varz[i]) R = lattice_transfer_map(lat, beam.E)[0:2, 0:2] # print # R err = np.linalg.norm(np.abs(R - target_matrix)**2) # print # 'iteration error: ', err return err
def errf(x): tw_loc = deepcopy(tw) tw0 = deepcopy(tw) ''' parameter to be varied is determined by variable class ''' for i in range(len(vars)): if isinstance(vars[i], Drift): if x[i] < 0: # print('negative length in match') return weights('negative_length') vars[i].l = x[i] vars[i].transfer_map = lat.method.create_tm(vars[i]) if isinstance(vars[i], Quadrupole): vars[i].k1 = x[i] vars[i].transfer_map = lat.method.create_tm(vars[i]) if isinstance(vars[i], Solenoid): vars[i].k = x[i] vars[i].transfer_map = lat.method.create_tm(vars[i]) if isinstance(vars[i], (RBend, SBend, Bend)): if vary_bend_angle: vars[i].angle = x[i] else: vars[i].k1 = x[i] vars[i].transfer_map = lat.method.create_tm(vars[i]) if isinstance(vars[i], list): if isinstance(vars[i][0], Twiss) and isinstance( vars[i][1], str): k = vars[i][1] tw_loc.__dict__[k] = x[i] if isinstance( vars[i], tuple ): # all quads strength in tuple varied simultaneously for v in vars[i]: v.k1 = x[i] v.transfer_map = lat.method.create_tm(v) err = 0.0 if "periodic" in constr.keys(): if constr["periodic"]: tw_loc = periodic_twiss(tw_loc, lattice_transfer_map(lat, tw.E)) tw0 = deepcopy(tw_loc) if tw_loc is None: print("########") return weights('periodic') # save reference points where equality is asked ref_hsh = {} # penalties on two-point inequalities for e in constr.keys(): if e == 'periodic': continue if e == 'total_len': continue for k in constr[e].keys(): if isinstance(constr[e][k], list): if constr[e][k][0] == '->': # print 'creating reference to', constr[e][k][1].id ref_hsh[constr[e][k][1]] = {k: 0.0} # evaluating global and point penalties tw_loc.s = 0 for e in lat.sequence: tw_loc = e.transfer_map * tw_loc if 'global' in constr.keys(): for c in constr['global'].keys(): if isinstance(constr['global'][c], list): v1 = constr['global'][c][1] if constr['global'][c][0] == '<': if tw_loc.__dict__[c] > v1: err = err + weights(k) * (tw_loc.__dict__[c] - v1)**2 if constr['global'][c][0] == '>': if tw_loc.__dict__[c] < v1: err = err + weights(k) * (tw_loc.__dict__[c] - v1)**2 if 'delta' in constr.keys(): if e in constr['delta'].keys(): tw_k = constr['delta'][e][0] constr['delta'][e][1] = tw_loc.__dict__[tw_k] if e in ref_hsh.keys(): ref_hsh[e] = deepcopy(tw_loc) if e in constr.keys(): for k in constr[e].keys(): if isinstance(constr[e][k], list): v1 = constr[e][k][1] if constr[e][k][0] == '<': if tw_loc.__dict__[k] > v1: err = err + weights(k) * (tw_loc.__dict__[k] - v1)**2 if constr[e][k][0] == '>': if tw_loc.__dict__[k] < v1: err = err + weights(k) * (tw_loc.__dict__[k] - v1)**2 if constr[e][k][0] == 'a<': if np.abs(tw_loc.__dict__[k]) > v1: err = err + weights(k) * (tw_loc.__dict__[k] - v1)**2 if constr[e][k][0] == 'a>': if np.abs(tw_loc.__dict__[k]) < v1: err = err + weights(k) * (tw_loc.__dict__[k] - v1)**2 if constr[e][k][0] == '->': try: if len(constr[e][k]) > 2: dv1 = float(constr[e][k][2]) else: dv1 = 0.0 err += (tw_loc.__dict__[k] - (ref_hsh[v1].__dict__[k] + dv1))**2 if tw_loc.__dict__[k] < v1: err = err + (tw_loc.__dict__[k] - v1)**2 except: print( 'constraint error: rval should precede lval in lattice' ) if tw_loc.__dict__[k] < 0: err += (tw_loc.__dict__[k] - v1)**2 elif isinstance(constr[e][k], str): pass else: err = err + weights(k) * (constr[e][k] - tw_loc.__dict__[k])**2 if "total_len" in constr.keys(): total_len = constr["periodic"] err = err + weights('total_len') * (tw_loc.s - total_len)**2 if 'delta' in constr.keys(): delta_dict = constr['delta'] elems = [] for e in delta_dict.keys(): if isinstance(e, Element): elems.append(e) delta_err = delta_dict["weight"] * (delta_dict[elems[0]][1] - delta_dict[elems[1]][1] - delta_dict["val"])**2 err = err + delta_err if min_i5: ''' evaluating integral parameters ''' I1, I2, I3, I4, I5 = radiation_integrals(lat, tw0, nsuperperiod=1) err += I5 * weights('i5') Je = 2 + I4 / I2 Jx = 1 - I4 / I2 Jy = 1 if Je < 0 or Jx < 0 or Jy < 0: err = 100000.0 # c1, c2 = natural_chromaticity(lat, tw0) # err += ( c1**2 + c2**2) * 1.e-6 if verbose: print('iteration error:', err) return err
def errf(x): p_array0 = deepcopy(p_array) tws = get_envelope(p_array0, bounds=bounds) tw_loc = deepcopy(tws) tw0 = deepcopy(tws) ''' parameter to be varied is determined by variable class ''' for i in range(len(vars)): if vars[i].__class__ == Drift: if x[i] < 0: # print('negative length in match') return weights('negative_length') pass vars[i].l = x[i] vars[i].transfer_map = lat.method.create_tm(vars[i]) if vars[i].__class__ == Quadrupole: vars[i].k1 = x[i] vars[i].transfer_map = lat.method.create_tm(vars[i]) if vars[i].__class__ == Solenoid: vars[i].k = x[i] vars[i].transfer_map = lat.method.create_tm(vars[i]) if vars[i].__class__ in [RBend, SBend, Bend]: if vary_bend_angle: vars[i].angle = x[i] else: vars[i].k1 = x[i] vars[i].transfer_map = lat.method.create_tm(vars[i]) if vars[i].__class__ == list: if vars[i][0].__class__ == Twiss and vars[i][ 1].__class__ == str: k = vars[i][1] tw_loc.__dict__[k] = x[i] if vars[i].__class__ == tuple: # all quads strength in tuple varied simultaneously for v in vars[i]: v.k1 = x[i] v.transfer_map = lat.method.create_tm(v) err = 0.0 if "periodic" in constr.keys(): if constr["periodic"] is True: tw_loc = periodic_twiss(tw_loc, lattice_transfer_map(lat, tw_loc.E)) tw0 = deepcopy(tw_loc) if tw_loc is None: print("########") return weights('periodic') # save reference points where equality is asked ref_hsh = {} # penalties on two-point inequalities for e in constr.keys(): if e == 'periodic': continue if e == 'total_len': continue for k in constr[e].keys(): if constr[e][k].__class__ == list: if constr[e][k][0] == '->': # print 'creating reference to', constr[e][k][1].id ref_hsh[constr[e][k][1]] = {k: 0.0} # print 'references:', ref_hsh.keys() # evaluating global and point penalties # tw_loc.s = 0 # print("start = ", get_envelope(p_array0)) navi.go_to_start() tws_list, p_array0 = track(lat, p_array0, navi, print_progress=False, bounds=bounds) s = np.array([tw.s for tw in tws_list]) # print("stop = ", tws_list[-1]) L = 0. for e in lat.sequence: indx = (np.abs(s - L)).argmin() L += e.l tw_loc = tws_list[indx] if 'global' in constr.keys(): # print 'there is a global constraint', constr['global'].keys() for c in constr['global'].keys(): if constr['global'][c].__class__ == list: # print 'list' v1 = constr['global'][c][1] if constr['global'][c][0] == '<': if tw_loc.__dict__[c] > v1: err = err + weights(k) * (tw_loc.__dict__[c] - v1)**2 if constr['global'][c][0] == '>': # print '> constr' if tw_loc.__dict__[c] < v1: err = err + weights(k) * (tw_loc.__dict__[c] - v1)**2 if e in ref_hsh.keys(): # print 'saving twiss for', e.id ref_hsh[e] = deepcopy(tw_loc) if e in constr.keys(): for k in constr[e].keys(): # print(k) if constr[e][k].__class__ == list: v1 = constr[e][k][1] if constr[e][k][0] == '<': if tw_loc.__dict__[k] > v1: err = err + weights(k) * (tw_loc.__dict__[k] - v1)**2 if constr[e][k][0] == '>': if tw_loc.__dict__[k] < v1: err = err + weights(k) * (tw_loc.__dict__[k] - v1)**2 if constr[e][k][0] == 'a<': if np.abs(tw_loc.__dict__[k]) > v1: err = err + weights(k) * (tw_loc.__dict__[k] - v1)**2 if constr[e][k][0] == 'a>': if np.abs(tw_loc.__dict__[k]) < v1: err = err + weights(k) * (tw_loc.__dict__[k] - v1)**2 if constr[e][k][0] == '->': try: # print 'huh', k, e.id, float(constr[e][k][2]) if len(constr[e][k]) > 2: dv1 = float(constr[e][k][2]) else: dv1 = 0.0 # print 'weiter' err += (tw_loc.__dict__[k] - (ref_hsh[v1].__dict__[k] + dv1))**2 if tw_loc.__dict__[k] < v1: err = err + (tw_loc.__dict__[k] - v1)**2 except: print( 'constraint error: rval should precede lval in lattice' ) if tw_loc.__dict__[k] < 0: # print 'negative constr (???)' err += (tw_loc.__dict__[k] - v1)**2 else: # print "safaf", constr[e][k] , tw_loc.__dict__[k], k, e.id, x err = err + weights(k) * (constr[e][k] - tw_loc.__dict__[k])**2 # print err for v in vars: print(v.id, v.k1) if "total_len" in constr.keys(): total_len = constr["periodic"] err = err + weights('total_len') * (tw_loc.s - total_len)**2 if min_i5: ''' evaluating integral parameters ''' I1, I2, I3, I4, I5 = radiation_integrals(lat, tw0, nsuperperiod=1) err += I5 * weights('i5') Je = 2 + I4 / I2 Jx = 1 - I4 / I2 Jy = 1 if Je < 0 or Jx < 0 or Jy < 0: err = 100000.0 # c1, c2 = natural_chromaticity(lat, tw0) # err += ( c1**2 + c2**2) * 1.e-6 if verbose: print('iteration error:', err) return err