diff --git a/examples/coasting/001_frev_meas.py b/examples/coasting/001_frev_meas.py index 285694372..cdfa3c13b 100644 --- a/examples/coasting/001_frev_meas.py +++ b/examples/coasting/001_frev_meas.py @@ -4,7 +4,21 @@ from scipy.constants import c as clight -delta0 = 1e-2 +delta0 = 0 #-1e-2 +delta_range = 0 +num_turns=100 +num_particles = 100_000 + +# delta0 = 1e-2 +# delta_range = 0 +# num_turns=20 +# num_particles = 5_000_000 + +# # To see the different number of turns +# delta0 = 0e-2 +# delta_range = 10e-3 +# num_turns=5000 +# num_particles = 5000 line = xt.Line.from_json( '../../test_data/psb_injection/line_and_particle.json') @@ -42,15 +56,15 @@ zeta_min0 = -circumference/2*tw.beta0/beta1 zeta_max0 = circumference/2*tw.beta0/beta1 -num_particles = 50000 + p = line.build_particles( - delta=delta0 + 0 * np.random.uniform(-1, 1, num_particles), + delta=delta0 + delta_range * np.random.uniform(-1, 1, num_particles), x_norm=0, y_norm=0 ) # Need to take beta of actual particles to convert the distribution along the # circumference to a distribution in time -p.zeta = (np.random.uniform(0, circumference, num_particles) / p.rvv +p.zeta = (np.random.uniform(0, circumference, num_particles) / p.rvv * 0.999 + (zeta_max0 - circumference) / p.rvv) st.prepare_particles_for_sync_time(p, line) @@ -95,13 +109,13 @@ def y_mean_hist(line, particles): line.enable_time_dependent_vars = True -num_turns=200 + line.track(p, num_turns=num_turns, log=xt.Log(intensity=intensity, long_density=long_density, y_mean_hist=y_mean_hist, z_range=z_range, particles=particles - ), with_progress=10) + ), with_progress=2) inten = line.log_last_track['intensity'] @@ -165,8 +179,8 @@ def y_mean_hist(line, particles): plt.figure(2) plt.plot(p.delta, p.at_turn, '.') -plt.ylabel('Number of turns') -plt.xlabel(r'$\delta$') +plt.ylabel('Number of revolutions') +plt.xlabel(r'$\Delta P / P_0$') plt.figure(3) plt.plot([zz[1]-zz[0] for zz in line.log_last_track['z_range']]) @@ -201,18 +215,21 @@ def y_mean_hist(line, particles): plt.xlabel('z [m]') plt.ylabel('x [m]') -plt.figure(8) +f8 = plt.figure(8) ax1 = plt.subplot(2, 1, 1) plt.plot(t_unwrapped*1e6, y_vs_t, '-') plt.ylabel('y mean [m]') -plt.grid() ax2 = plt.subplot(2, 1, 2, sharex=ax1) -plt.plot(t_unwrapped*1e6, intensity_vs_t, '-') -plt.ylabel('intensity') +plt.plot(t_unwrapped*1e6, intensity_vs_t/np.mean(intensity_vs_t), '-') +plt.ylabel('Beam line density [a.u.]') plt.xlabel('t [us]') plt.ylim(bottom=0) -for tt in t_range_size * np.arange(0, hist_y.shape[0]): - ax1.axvline(x=tt*1e6, color='red', linestyle='--', alpha=0.5) +for tt in np.arange(0, t_unwrapped[-1], 1/f_nominal): + for ax in [ax1, ax2]: + ax.axvline(x=tt*1e6, color='red', linestyle='--', alpha=0.5) + +# zoom +ax1.set_xlim(0, 15) plt.show() \ No newline at end of file diff --git a/examples/coasting/e000_illustration.py b/examples/coasting/e000_illustration.py new file mode 100644 index 000000000..ed5198af1 --- /dev/null +++ b/examples/coasting/e000_illustration.py @@ -0,0 +1,49 @@ +import numpy as np +import matplotlib.pyplot as plt + +v = 1. +l_ref = 1 + +t_ref = l_ref/v + +l = np.array([0.95, 1, 1.05]) + +t_max = 80 +t = np.linspace(0, t_max, 10000) + +s = np.zeros((len(l), len(t))) +for il, ll in enumerate(l): + s[il, :] = np.mod(0.62*l_ref + v * t * ll/l_ref, l_ref) + +passing = np.diff(s, prepend=0) < 0 + +import matplotlib.pyplot as plt +plt.close('all') + +plt.figure(1, figsize=(6.4*1.8*0.8, 4.8*0.8)) +sp1 = plt.subplot(311) +sp2 = plt.subplot(312, sharex=sp1) +sp3 = plt.subplot(313, sharex=sp1) + +sp1.stem(t, passing[1, :], markerfmt=' ', basefmt='C0', linefmt='C0') +sp2.stem(t, passing[0, :], markerfmt=' ', basefmt='C1', linefmt='C1') +sp3.stem(t, passing[2, :], markerfmt=' ', basefmt='C2', linefmt='C2') + +t_turn = np.arange(0, t_max, t_ref) +for tt in t_turn: + for sp in [sp1, sp2, sp3]: + sp.axvline(tt, color='k', linestyle='--', alpha=0.4) + +sp3.set_xlim(0, t_max/t_ref) +sp3.set_xlabel(r'$t~/~T_0$') + +for sp in [sp1, sp2, sp3]: + sp.set_ylim(0, 1.1) + +plt.subplots_adjust(bottom=.14, top=.95, hspace=0.3) + +# For zoom +# plt.subplots_adjust(right=.6) + +plt.show() + diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py new file mode 100644 index 000000000..2de3da73f --- /dev/null +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -0,0 +1,230 @@ +import xtrack as xt +import numpy as np + +env = xt.Environment() +env.particle_ref = xt.Particles(p0c=2e9) + +n_bends_per_cell = 6 +n_cells_par_arc = 3 +n_arcs = 3 + +n_bends = n_bends_per_cell * n_cells_par_arc * n_arcs + +env.vars({ + 'l.mq': 0.5, + 'kqf': 0.027, + 'kqd': -0.0271, + 'l.mb': 10, + 'l.ms': 0.3, + 'k2sf': 0.001, + 'k2sd': -0.001, + 'angle.mb': 2 * np.pi / n_bends, + 'k0.mb': 'angle.mb / l.mb', + 'k0l.corrector': 0, + 'k1sl.corrector': 0, + 'l.halfcell': 38, +}) + +env.new('mb', xt.Bend, length='l.mb', k0='k0.mb', h='k0.mb') +env.new('mq', xt.Quadrupole, length='l.mq') +env.new('ms', xt.Sextupole, length='l.ms') +env.new('corrector', xt.Multipole, knl=[0], ksl=[0]) + +env.new('mq.f', 'mq', k1='kqf') +env.new('mq.d', 'mq', k1='kqd') + +halfcell = env.new_line(components=[ + + # End of the half cell (will be mid of the cell) + env.new('mid', xt.Marker, at='l.halfcell'), + + # Bends + env.new('mb.2', 'mb', at='l.halfcell / 2'), + env.new('mb.1', 'mb', at='-l.mb - 1', from_='mb.2'), + env.new('mb.3', 'mb', at='l.mb + 1', from_='mb.2'), + + # Quads + env.place('mq.d', at = '0.5 + l.mq / 2'), + env.place('mq.f', at = 'l.halfcell - l.mq / 2 - 0.5'), + + # Sextupoles + env.new('ms.d', 'ms', k2='k2sf', at=1.2, from_='mq.d'), + env.new('ms.f', 'ms', k2='k2sd', at=-1.2, from_='mq.f'), + + # Dipole correctors + env.new('corrector.v', 'corrector', at=0.75, from_='mq.d'), + env.new('corrector.h', 'corrector', at=-0.75, from_='mq.f') + +]) + +hcell_left = halfcell.replicate(name='l', mirror=True) +hcell_right = halfcell.replicate(name='r') + +cell = env.new_line(components=[ + env.new('start', xt.Marker), + hcell_left, + hcell_right, + env.new('end', xt.Marker), +]) + +opt = cell.match( + method='4d', + vary=xt.VaryList(['kqf', 'kqd'], step=1e-5), + targets=xt.TargetSet( + qx=0.333333, + qy=0.333333, + )) +tw_cell = cell.twiss4d() + + +env.vars({ + 'kqf.ss': 0.027 / 2, + 'kqd.ss': -0.0271 / 2, +}) + +halfcell_ss = env.new_line(components=[ + + env.new('mid', xt.Marker, at='l.halfcell'), + + env.new('mq.ss.d', 'mq', k1='kqd.ss', at = '0.5 + l.mq / 2'), + env.new('mq.ss.f', 'mq', k1='kqf.ss', at = 'l.halfcell - l.mq / 2 - 0.5'), + + env.new('corrector.ss.v', 'corrector', at=0.75, from_='mq.ss.d'), + env.new('corrector.ss.h', 'corrector', at=-0.75, from_='mq.ss.f') +]) + +hcell_left_ss = halfcell_ss.replicate(name='l', mirror=True) +hcell_right_ss = halfcell_ss.replicate(name='r') +cell_ss = env.new_line(components=[ + env.new('start.ss', xt.Marker), + hcell_left_ss, + hcell_right_ss, + env.new('end.ss', xt.Marker), +]) + +opt = cell_ss.match( + method='4d', + vary=xt.VaryList(['kqf.ss', 'kqd.ss'], step=1e-5), + targets=xt.TargetSet( + betx=tw_cell.betx[-1], bety=tw_cell.bety[-1], at='start.ss', + )) + + + +arc = env.new_line(components=[ + cell.replicate(name='cell.1'), + cell.replicate(name='cell.2'), + cell.replicate(name='cell.3'), +]) + + +ss = env.new_line(components=[ + cell_ss.replicate('cell.1'), + cell_ss.replicate('cell.2'), +]) + +ring = env.new_line(components=[ + arc.replicate(name='arc.1'), + ss.replicate(name='ss.1'), + arc.replicate(name='arc.2'), + ss.replicate(name='ss.2'), + arc.replicate(name='arc.3'), + ss.replicate(name='ss.3'), +]) + +## Insertion + +env.vars({ + 'k1.q1': 0.025, + 'k1.q2': -0.025, + 'k1.q3': 0.025, + 'k1.q4': -0.02, + 'k1.q5': 0.025, +}) + +half_insertion = env.new_line(components=[ + + # Start-end markers + env.new('ip', xt.Marker), + env.new('e.insertion', xt.Marker, at=76), + + # Quads + env.new('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), + env.new('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), + env.new('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), + env.new('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), + env.new('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq', at=73), + + # Dipole correctors (will use h and v on the same corrector) + env.new('corrector.ss.1', 'corrector', at=0.75, from_='mq.1'), + env.new('corrector.ss.2', 'corrector', at=-0.75, from_='mq.2'), + env.new('corrector.ss.3', 'corrector', at=0.75, from_='mq.3'), + env.new('corrector.ss.4', 'corrector', at=-0.75, from_='mq.4'), + env.new('corrector.ss.5', 'corrector', at=0.75, from_='mq.5'), + +]) + +tw_arc = arc.twiss4d() + +opt = half_insertion.match( + solve=False, + betx=tw_arc.betx[0], bety=tw_arc.bety[0], + alfx=tw_arc.alfx[0], alfy=tw_arc.alfy[0], + init_at='e.insertion', + start='ip', end='e.insertion', + vary=xt.VaryList(['k1.q1', 'k1.q2', 'k1.q3', 'k1.q4'], step=1e-5), + targets=[ + xt.TargetSet(alfx=0, alfy=0, at='ip'), + xt.Target(lambda tw: tw.betx[0] - tw.bety[0], 0), + xt.Target(lambda tw: tw.betx.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.bety.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.betx.min(), xt.GreaterThan(2)), + xt.Target(lambda tw: tw.bety.min(), xt.GreaterThan(2)), + ] +) +opt.step(40) +opt.solve() + +insertion = env.new_line([ + half_insertion.replicate('l', mirror=True), + half_insertion.replicate('r')]) + + + +ring2 = env.new_line(components=[ + arc.replicate(name='arcc.1'), + ss.replicate(name='sss.2'), + arc.replicate(name='arcc.2'), + insertion, + arc.replicate(name='arcc.3'), + ss.replicate(name='sss.3') +]) + + +# # Check buffer behavior +ring2_sliced = ring2.select() +ring2_sliced.cut_at_s(np.arange(0, ring2.get_length(), 0.5)) + + +import matplotlib.pyplot as plt +plt.close('all') +for ii, rr in enumerate([ring, ring2_sliced]): + + tw = rr.twiss4d() + + fig = plt.figure(ii, figsize=(6.4*1.2, 4.8)) + ax1 = fig.add_subplot(2, 1, 1) + pltbet = tw.plot('betx bety', ax=ax1) + ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) + pltdx = tw.plot('dx', ax=ax2) + fig.subplots_adjust(right=.85) + pltbet.move_legend(1.2,1) + pltdx.move_legend(1.2,1) + +ring2.survey().plot() + + +plt.show() + + + diff --git a/examples/lattice_design_shortcuts/000a_subline.py b/examples/lattice_design_shortcuts/000a_subline.py new file mode 100644 index 000000000..9ba1b745c --- /dev/null +++ b/examples/lattice_design_shortcuts/000a_subline.py @@ -0,0 +1,232 @@ +import xtrack as xt +import numpy as np + +env = xt.Environment() +env.particle_ref = xt.Particles(p0c=2e9) + +n_bends_per_cell = 6 +n_cells_par_arc = 3 +n_arcs = 3 + +n_bends = n_bends_per_cell * n_cells_par_arc * n_arcs + +env.vars({ + 'l.mq': 0.5, + 'kqf': 0.027, + 'kqd': -0.0271, + 'l.mb': 10, + 'l.ms': 0.3, + 'k2sf': 0.001, + 'k2sd': -0.001, + 'angle.mb': 2 * np.pi / n_bends, + 'k0.mb': 'angle.mb / l.mb', + 'k0l.corrector': 0, + 'k1sl.corrector': 0, + 'l.halfcell': 38, +}) + +env.new('mb', xt.Bend, length='l.mb', k0='k0.mb', h='k0.mb') +env.new('mq', xt.Quadrupole, length='l.mq') +env.new('ms', xt.Sextupole, length='l.ms') +env.new('corrector', xt.Multipole, knl=[0], length=0.1) + +girder = env.new_line(components=[ + env.place('mq', at=1), + env.place('ms', at=0.8, from_='mq'), + env.place('corrector', at=-0.8, from_='mq'), +]) + +girder_f = girder.clone(name='f') +girder_d = girder.clone(name='d', mirror=True) +env.set('mq.f', k1='kqf') +env.set('mq.d', k1='kqd') + +halfcell = env.new_line(components=[ + + # End of the half cell (will be mid of the cell) + env.new('mid', xt.Marker, at='l.halfcell'), + + # Bends + env.new('mb.2', 'mb', at='l.halfcell / 2'), + env.new('mb.1', 'mb', at='-l.mb - 1', from_='mb.2'), + env.new('mb.3', 'mb', at='l.mb + 1', from_='mb.2'), + + # Quadrupoles, sextupoles and correctors + env.place(girder_d, at=1.2), + env.place(girder_f, at='l.halfcell - 1.2'), + +]) + + +hcell_left = halfcell.replicate(name='l', mirror=True) +hcell_right = halfcell.replicate(name='r') + +cell = env.new_line(components=[ + env.new('start', xt.Marker), + hcell_left, + hcell_right, + env.new('end', xt.Marker), +]) + +opt = cell.match( + method='4d', + vary=xt.VaryList(['kqf', 'kqd'], step=1e-5), + targets=xt.TargetSet( + qx=0.333333, + qy=0.333333, + )) +tw_cell = cell.twiss4d() + + +env.vars({ + 'kqf.ss': 0.027 / 2, + 'kqd.ss': -0.0271 / 2, +}) + +halfcell_ss = env.new_line(components=[ + + env.new('mid', xt.Marker, at='l.halfcell'), + + env.new('mq.ss.d', 'mq', k1='kqd.ss', at = '0.5 + l.mq / 2'), + env.new('mq.ss.f', 'mq', k1='kqf.ss', at = 'l.halfcell - l.mq / 2 - 0.5'), + + env.new('corrector.ss.v', 'corrector', at=0.75, from_='mq.ss.d'), + env.new('corrector.ss.h', 'corrector', at=-0.75, from_='mq.ss.f') +]) + +hcell_left_ss = halfcell_ss.replicate(name='l', mirror=True) +hcell_right_ss = halfcell_ss.replicate(name='r') +cell_ss = env.new_line(components=[ + env.new('start.ss', xt.Marker), + hcell_left_ss, + hcell_right_ss, + env.new('end.ss', xt.Marker), +]) + +opt = cell_ss.match( + solve=False, + method='4d', + vary=xt.VaryList(['kqf.ss', 'kqd.ss'], step=1e-5), + targets=xt.TargetSet( + betx=tw_cell.betx[-1], bety=tw_cell.bety[-1], at='start.ss', + )) +opt.solve() + + +arc = env.new_line(components=[ + cell.replicate(name='cell.1'), + cell.replicate(name='cell.2'), + cell.replicate(name='cell.3'), +]) + + +ss = env.new_line(components=[ + cell_ss.replicate('cell.1'), + cell_ss.replicate('cell.2'), +]) + +ring = env.new_line(components=[ + arc.replicate(name='arc.1'), + ss.replicate(name='ss.1'), + arc.replicate(name='arc.2'), + ss.replicate(name='ss.2'), + arc.replicate(name='arc.3'), + ss.replicate(name='ss.3'), +]) + +## Insertion + +env.vars({ + 'k1.q1': 0.025, + 'k1.q2': -0.025, + 'k1.q3': 0.025, + 'k1.q4': -0.02, + 'k1.q5': 0.025, +}) + +half_insertion = env.new_line(components=[ + + # Start-end markers + env.new('ip', xt.Marker), + env.new('e.insertion', xt.Marker, at=76), + + # Quads + env.new('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), + env.new('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), + env.new('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), + env.new('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), + env.new('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq', at=73), + + # Dipole correctors (will use h and v on the same corrector) + env.new('corrector.ss.1', 'corrector', at=0.75, from_='mq.1'), + env.new('corrector.ss.2', 'corrector', at=-0.75, from_='mq.2'), + env.new('corrector.ss.3', 'corrector', at=0.75, from_='mq.3'), + env.new('corrector.ss.4', 'corrector', at=-0.75, from_='mq.4'), + env.new('corrector.ss.5', 'corrector', at=0.75, from_='mq.5'), + +]) + +tw_arc = arc.twiss4d() + +opt = half_insertion.match( + solve=False, + betx=tw_arc.betx[0], bety=tw_arc.bety[0], + alfx=tw_arc.alfx[0], alfy=tw_arc.alfy[0], + init_at='e.insertion', + start='ip', end='e.insertion', + vary=xt.VaryList(['k1.q1', 'k1.q2', 'k1.q3', 'k1.q4'], step=1e-5), + targets=[ + xt.TargetSet(alfx=0, alfy=0, at='ip'), + xt.Target(lambda tw: tw.betx[0] - tw.bety[0], 0), + xt.Target(lambda tw: tw.betx.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.bety.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.betx.min(), xt.GreaterThan(2)), + xt.Target(lambda tw: tw.bety.min(), xt.GreaterThan(2)), + ] +) +opt.step(40) +opt.solve() + +insertion = env.new_line([ + half_insertion.replicate('l', mirror=True), + half_insertion.replicate('r')]) + + + +ring2 = env.new_line(components=[ + arc.replicate(name='arcc.1'), + ss.replicate(name='sss.2'), + arc.replicate(name='arcc.2'), + insertion, + arc.replicate(name='arcc.3'), + ss.replicate(name='sss.3') +]) + + +# # Check buffer behavior +ring2_sliced = ring2.select() +ring2_sliced.cut_at_s(np.arange(0, ring2.get_length(), 0.5)) + + +import matplotlib.pyplot as plt +plt.close('all') +for ii, rr in enumerate([ring, ring2_sliced]): + + tw = rr.twiss4d() + + fig = plt.figure(ii, figsize=(6.4*1.2, 4.8)) + ax1 = fig.add_subplot(2, 1, 1) + pltbet = tw.plot('betx bety', ax=ax1) + ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) + pltdx = tw.plot('dx', ax=ax2) + fig.subplots_adjust(right=.85) + pltbet.move_legend(1.2,1) + pltdx.move_legend(1.2,1) + +ring2.survey().plot() + + +plt.show() + + + diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py new file mode 100644 index 000000000..ea411ecd7 --- /dev/null +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -0,0 +1,65 @@ +import xtrack as xt +import xobjects as xo +import numpy as np + +env = xt.Environment() + +env.vars({ + 'l.b1': 1.0, + 'l.q1': 0.5, + 's.ip': 10, + 's.left': -5, + 's.right': 5, + 'l.before_right': 1, + 'l.after_left2': 0.5, +}) + +# names, tab_sorted = handle_s_places(seq) +line = env.new_line(components=[ + env.new('b1', xt.Bend, length='l.b1'), + env.new('q1', xt.Quadrupole, length='l.q1'), + env.new('ip', xt.Marker, at='s.ip'), + ( + env.new('before_before_right', xt.Marker), + env.new('before_right', xt.Sextupole, length=1), + env.new('right',xt.Quadrupole, length=0.8, at='s.right', from_='ip'), + env.new('after_right', xt.Marker), + env.new('after_right2', xt.Marker), + ), + env.new('left', xt.Quadrupole, length=1, at='s.left', from_='ip'), + env.new('after_left', xt.Marker), + env.new('after_left2', xt.Bend, length='l.after_left2'), +]) + +tt = line.get_table(attr=True) +tt['s_center'] = tt['s'] + tt['length']/2 +assert np.all(tt.name == np.array([ + 'b1', 'q1', 'drift_1', 'left', 'after_left', 'after_left2', + 'drift_2', 'ip', 'drift_3', 'before_before_right', 'before_right', + 'right', 'after_right', 'after_right2', '_end_point'])) + +xo.assert_allclose(env['b1'].length, 1.0, rtol=0, atol=1e-14) +xo.assert_allclose(env['q1'].length, 0.5, rtol=0, atol=1e-14) +xo.assert_allclose(tt['s', 'ip'], 10, rtol=0, atol=1e-14) +xo.assert_allclose(tt['s', 'before_before_right'], tt['s', 'before_right'], + rtol=0, atol=1e-14) +xo.assert_allclose(tt['s_center', 'before_right'] - tt['s_center', 'right'], + -(1 + 0.8)/2, rtol=0, atol=1e-14) +xo.assert_allclose(tt['s_center', 'right'] - tt['s', 'ip'], 5, rtol=0, atol=1e-14) +xo.assert_allclose(tt['s_center', 'after_right'] - tt['s_center', 'right'], + 0.8/2, rtol=0, atol=1e-14) +xo.assert_allclose(tt['s_center', 'after_right2'] - tt['s_center', 'right'], + 0.8/2, rtol=0, atol=1e-14) +xo.assert_allclose(tt['s_center', 'left'] - tt['s_center', 'ip'], -5, + rtol=0, atol=1e-14) +xo.assert_allclose(tt['s_center', 'after_left'] - tt['s_center', 'left'], 1/2, + rtol=0, atol=1e-14) +xo.assert_allclose(tt['s_center', 'after_left2'] - tt['s_center', 'after_left'], + 0.5/2, rtol=0, atol=1e-14) + + +import matplotlib.pyplot as plt +plt.close('all') +line.survey().plot() + +plt.show() \ No newline at end of file diff --git a/examples/lattice_design_shortcuts/002_ring_sequence.py b/examples/lattice_design_shortcuts/002_ring_sequence.py new file mode 100644 index 000000000..f5c44188d --- /dev/null +++ b/examples/lattice_design_shortcuts/002_ring_sequence.py @@ -0,0 +1,265 @@ +import xtrack as xt +import numpy as np + +env = xt.Environment() +env.particle_ref = xt.Particles(p0c=2e9) + +n_bends_per_cell = 6 +n_cells_par_arc = 3 +n_arcs = 3 + +n_bends = n_bends_per_cell * n_cells_par_arc * n_arcs + +env.vars({ + 'k1l.qf': 0.027 / 2, + 'k1l.qd': -0.0271 / 2, + 'l.mq': 0.5, + 'kqf.1': 'k1l.qf / l.mq', + 'kqd.1': 'k1l.qd / l.mq', + 'l.mb': 12, + 'angle.mb': 2 * np.pi / n_bends, + 'k0.mb': 'angle.mb / l.mb', + 'k0l.corrector': 0, + 'k1sl.corrector': 0, +}) + + +env.new_element('drift.1', xt.Drift, length='l.mq / 2') +env.new_element('qf', xt.Quadrupole, k1='kqf.1', length='l.mq') +env.new_element('drift.2', xt.Replica, parent_name='drift.1') +env.new_element('mb.1', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb') +env.new_element('mb.2', xt.Replica, parent_name='mb.1') +env.new_element('mb.3', xt.Replica, parent_name='mb.1') +env.new_element('drift.3', xt.Replica, parent_name='drift.1') +env.new_element('qd', xt.Quadrupole, k1='kqd.1', length='l.mq') +env.new_element('drift.4', xt.Replica, parent_name='drift.1') + +halfcell = env.new_line(components=[ + 'drift.1', + 'qf', + 'drift.2', + 'mb.1', + 'mb.2', + 'mb.3', + 'drift.3', + 'qd', + 'drift.4', +]) + + + +halfcell = env.new_line(components=[ + env.new_element('drift.1', xt.Drift, length='l.mq / 2'), + env.new_element('qf', xt.Quadrupole, k1='kqf.1', length='l.mq'), + env.new_element('drift.2', xt.Replica, parent_name='drift.1'), + env.new_element('mb.1', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb'), + env.new_element('mb.2', xt.Replica, parent_name='mb.1'), + env.new_element('mb.3', xt.Replica, parent_name='mb.1'), + env.new_element('drift.3', xt.Replica, parent_name='drift.1'), + env.new_element('qd', xt.Quadrupole, k1='kqd.1', length='l.mq'), + env.new_element('drift.4', xt.Replica, parent_name='drift.1'), +]) + +hcell_left = halfcell.replicate(name='l') +hcell_right = halfcell.replicate(name='r', mirror=True) + +cell = env.new_line(components=[ + env.new_element('start', xt.Marker), + env.new_element('corrector.l', xt.Multipole, knl=['k0l.corrector', 0], + ksl=[0, 'k1sl.corrector']), + hcell_left, + env.new_element('mid', xt.Marker), + env.new_element('corrector.v', xt.Multipole, knl=['k0l.corrector', 0], + ksl=[0, 'k1sl.corrector']), + hcell_right, + env.new_element('end', xt.Marker), +]) + + +arc = env.new_line(components=[ + cell.replicate(name='cell.1'), + cell.replicate(name='cell.2'), + cell.replicate(name='cell.3'), +]) + +cell_ss = cell.replicate('ss') +env.new_element('drift_ss', xt.Drift, length='l.mb') +for ii, nn in enumerate(cell_ss.element_names): + if nn.startswith('mb'): + cell_ss.element_names[ii] = env.new_element( + f'drift.{ii}.ss', xt.Replica, parent_name='drift_ss') + +ss = env.new_line(components=[ + cell_ss.replicate('cell.1'), + cell_ss.replicate('cell.2'), +]) + +arc1 = arc.replicate(name='arc.1') +arc2 = arc.replicate(name='arc.2') +arc3 = arc.replicate(name='arc.3') + +# ss1 = ss.replicate(name='ss.1') +# ss2 = ss.replicate(name='ss.2') +# ss3 = ss.replicate(name='ss.3') + + + +opt = cell.match( + method='4d', + vary=xt.VaryList(['k1l.qf', 'k1l.qd'], step=1e-5), + targets=xt.TargetSet( + qx=0.333333, + qy=0.333333, + )) + +env.vars({ + 'k1l.q1': 0.012, + 'k1l.q2': -0.012, + 'k1l.q3': 0.012, + 'k1l.q4': -0.012, + 'k1l.q5': 0.012, + 'k1.q1': 'k1l.q1 / l.mq', + 'k1.q2': 'k1l.q2 / l.mq', + 'k1.q3': 'k1l.q3 / l.mq', + 'k1.q4': 'k1l.q4 / l.mq', + 'k1.q5': 'k1l.q5 / l.mq', +}) + +half_straight = env.new_line(components=[ + env.new_element('ip', xt.Marker), + env.new_element('dd.0', xt.Drift, length=20), + env.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), + env.new_element('dd.1', xt.Drift, length=5), + env.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), + env.new_element('dd.2', xt.Drift, length=12), + env.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), + env.new_element('dd.3', xt.Drift, length=18), + env.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), + env.new_element('dd.4', xt.Drift, length=18), + env.new_element('mq.5', xt.Quadrupole, k1='k1l.q5', length='l.mq'), + env.new_element('dd.5', xt.Drift, length=0.5), + env.new_element('e.ss.r', xt.Marker), +]) +half_straight.build_tracker() +print(f'Half straight length: {half_straight.get_length()}') + +tw_arc = arc.twiss4d() + +opt = half_straight.match( + solve=False, + betx=tw_arc.betx[0], bety=tw_arc.bety[0], + alfx=tw_arc.alfx[0], alfy=tw_arc.alfy[0], + init_at='e.ss.r', + start='ip', end='e.ss.r', + vary=xt.VaryList(['k1l.q1', 'k1l.q2', 'k1l.q3', 'k1l.q4'], step=1e-5), + targets=[ + xt.TargetSet(alfx=0, alfy=0, at='ip'), + xt.Target(lambda tw: tw.betx[0] - tw.bety[0], 0), + xt.Target(lambda tw: tw.betx.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.bety.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.betx.min(), xt.GreaterThan(2)), + xt.Target(lambda tw: tw.bety.min(), xt.GreaterThan(2)), + ] + ) + + +opt.step(40) + +half_straight_left = half_straight.replicate('ss.l') +half_straight_left.mirror() +half_straight_right = half_straight.replicate('ss.r') +straight = env.new_line(components=[half_straight_left, half_straight_right]) + +ss_arc = env.new_line(components=[arc1, straight, arc2]) +tw_ss_arc = ss_arc.twiss4d(betx=tw_arc.betx[-1], bety=tw_arc.bety[-1], + alfx=tw_arc.alfx[-1], alfy=tw_arc.alfy[-1], + init_at=xt.END) + +env.vars({ + 'k1l.qfss': 0.027 / 2, + 'k1l.qdss': -0.0271 / 2, + 'kqfss.1': 'k1l.qfss / l.mq', + 'kqdss.1': 'k1l.qdss / l.mq', + 'angle.mb': 2 * np.pi / n_bends, + 'k0.mb': 'angle.mb / l.mb', +}) +cell_ss = env.new_line(components=[ + env.new_element('ss.start', xt.Marker), + env.new_element('dd.ss.1.l', xt.Drift, length='l.mq'), + env.new_element('qfss.l', xt.Quadrupole, k1='kqfss.1', length='l.mq'), + + env.new_element('dd.ss.3.l', xt.Drift, length='3 *l.mb'), + + env.new_element('qdss.l', xt.Quadrupole, k1='kqdss.1', length='l.mq'), + env.new_element('dd.ss.5.l', xt.Drift, length='l.mq'), + + env.new_element('dd.ss.5.r', xt.Drift, length='l.mq'), + env.new_element('qdss.r', xt.Quadrupole, k1='kqdss.1', length='l.mq'), + + env.new_element('dd.ss.3.r', xt.Drift, length='3 *l.mb'), + + env.new_element('qfss.r', xt.Quadrupole, k1='kqfss.1', length='l.mq'), + env.new_element('dd.ss.1.r', xt.Drift, length='l.mq'), + +]) + +opt = cell_ss.match( + solve=False, + method='4d', + vary=xt.VaryList(['k1l.qfss', 'k1l.qdss'], step=1e-5), + targets=xt.TargetSet(at='ss.start', + betx=tw_arc.betx[-1], + bety=tw_arc.bety[-1], + )) +opt.step(40) +opt.solve() + +tw_ss_arc.plot() + + +cell1_ss = cell_ss.replicate('cell.1') +cell2_ss = cell_ss.replicate('cell.2') +std_ss = env.new_line(components=[cell1_ss, cell2_ss]) + +ss1 = std_ss.replicate('ss.1') +ss2 = std_ss.replicate('ss.2') + +ring = env.new_line() + +ring.extend(ss1) +ring.extend(arc1) +ring.extend(ss2) +ring.extend(arc2) +ring.extend(straight) +ring.extend(arc3) + +ring.replace_all_replicas() +ring.build_tracker() +sv = ring.survey() + +buffer = ring._buffer +ring.discard_tracker() +ring.cut_at_s(np.arange(0, ring.get_length(), 0.5)) +ring.build_tracker(_buffer=buffer) +tw = ring.twiss4d() + +two = ring.twiss(start=xt.START, betx=tw_arc.betx[-1], bety=tw_arc.bety[-1]) + +import matplotlib.pyplot as plt +plt.close('all') +fig = plt.figure(1, figsize=(6.4*1.2, 4.8)) +ax1 = fig.add_subplot(2, 1, 1) +pltbet = tw.plot('betx bety', ax=ax1) +ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) +pltdx = tw.plot('dx', ax=ax2) +fig.subplots_adjust(right=.85) +pltbet.move_legend(1.2,1) +pltdx.move_legend(1.2,1) + +import xplt +xplt.FloorPlot(sv, ring, element_width=10) + +plt.show() + + + diff --git a/examples/lattice_design_shortcuts/010_play.py b/examples/lattice_design_shortcuts/010_play.py new file mode 100644 index 000000000..9ccda1e03 --- /dev/null +++ b/examples/lattice_design_shortcuts/010_play.py @@ -0,0 +1,75 @@ +import xtrack as xt +env = xt.Environment() + + +env.new('drift.1', xt.Drift, length='l.mq / 2') +env.new('qf', xt.Quadrupole, k1='kqf.1', length='l.mq') +env.new('drift.2', xt.Replica, parent_name='drift.1') +env.new('mb.1', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb') +env.new('mb.2', xt.Replica, parent_name='mb.1') +env.new('mb.3', xt.Replica, parent_name='mb.1') +env.new('drift.3', xt.Replica, parent_name='drift.1') +env.new('qd', xt.Quadrupole, k1='kqd.1', length='l.mq') +env.new('drift.4', xt.Replica, parent_name='drift.1') + +halfcell = env.new_line(components=[ + 'drift.1', + SPlace('qf', s + 'drift.2', + 'mb.1', + 'mb.2', + 'mb.3', + 'drift.3', + 'qd', + 'drift.4', +]) + +mbx = env.new('mbxw', xt.Bend, k0='k0.mb', h=0, length='l.mbxw') + +d1 = env.new_line('d1', components=[ + env.new('lmbxw.start', parent=xt.Marker), # shortcut env.new('lmbxw.start') + env.new('mbxw.a4@start', parent=xt.Replica, parent_name='mbxw', at=0.5), + env.new('mbxw.b4@start', parent='mbxw', from_='mbxw.a4@end'), + env.new('lmbxw.end', parent=xt.Marker, at=0.5, from_='mbxw.b4@end'), +]) + +d2 = env.new('d2.b1', xt.Bend, k0='k0.mb', h=0, length='l.mbxw', dx=0.188/2) + +ir_left = env.new_line('ir_left', components=[ + env.new('ip1') + Splace('d1r1@start', d1, at=100, from_='ip1'), + Splace('d2@start', d2, at=200, from_='ip1'), +]) + +ir = evn.new_line(components=[ + ir_left.replicate('.l1', mirror=True), + env.new('ip') + ir_left.replicate('.r1') +]) + + +s_ip1 = 0 +s_ip2 = 2000 +s_ip3 = 4000 + + +lhc = env.new_line(components=[ + + + place('ir1', ir.replicate('1'), at=s_ip1, patch=True, reference=PatchReference('ip.1', x=0.1, xp=0.1)), + place('ir2', ir.replicate('2'), at=s_ip2, patch=True, reference='ip.2'), + + +]) + +seq = [ + Place(env.new('ip', xt.Marker), at=10), + Place(env.new('left', xt.Quadrupole, length=1), at=-5, from_='ip'), + env.new('after_left', xt.Marker), + env.new('after_left2', xt.Marker), + Place(env.new('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + Place([env.new('before_right', xt.Quadrupole, length=1), + env.new('before_right2', xt.Marker)], at=0, from_='right@start', + anchor='before_right@center'), + Place(env.new('righter', xt.Quadrupole, length=1), at=+5, from_='before_right'), +] \ No newline at end of file diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py new file mode 100644 index 000000000..ca28b0d47 --- /dev/null +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -0,0 +1,139 @@ +import xtrack as xt +import numpy as np + +to_test = 'line' +to_test = 'env' + +env = xt.Environment() + +env.vars({ + 'k.1': 1., + 'a': 2., + 'b': '2 * a + k.1', +}) + +line = env.new_line([]) + +ee = {'env': env, 'line': line}[to_test] + +assert ee.vv['b'] == 2 * 2 + 1 + +ee.vars['a'] = ee.vars['k.1'] +assert ee.vv['b'] == 2 * 1 + 1 + +ee.vars(a=3.) +ee.vars({'k.1': 'a'}) +assert ee.vv['k.1'] == 3. +assert ee.vv['b'] == 2 * 3 + 3. + +ee.vars['k.1'] = 2 * ee.vars['a'] + 5 +assert ee.vv['k.1'] == 2 * 3 + 5 +assert ee.vv['b'] == 2 * 3 + 2 * 3 + 5 + +ee.vars.set('a', 4.) +assert ee.vv['k.1'] == 2 * 4 + 5 +assert ee.vv['b'] == 2 * 4 + 2 * 4 + 5 + +ee.vars.set('k.1', '2*a + 5') +assert ee.vv['k.1'] == 2 * 4 + 5 +assert ee.vv['b'] == 2 * 4 + 2 * 4 + 5 + +ee.vars.set('k.1', 3 * ee.vars['a'] + 6) +assert ee.vv['k.1'] == 3 * 4 + 6 +assert ee.vv['b'] == 2 * 4 + 3 * 4 + 6 + +ee.set('a', 0.) +assert ee.vv['k.1'] == 3 * 0 + 6 +assert ee.vv['b'] == 2 * 0 + 3 * 0 + 6 + +ee.set('a', 2.) +ee.set('k.1', '2 * a + 5') +assert ee.vv['k.1'] == 2 * 2 + 5 +assert ee.vv['b'] == 2 * 2 + 2 * 2 + 5 + +ee.set('k.1', 3 * ee.vars['a'] + 6) +assert ee.vv['k.1'] == 3 * 2 + 6 +assert ee.vv['b'] == 2 * 2 + 3 * 2 + 6 + +assert hasattr(ee.ref['k.1'], '_value') # is a Ref + +ee.ref['a'] = 0 +assert ee.vv['k.1'] == 3 * 0 + 6 +assert ee.vv['b'] == 2 * 0 + 3 * 0 + 6 + +ee.ref['a'] = 2 +ee.ref['k.1'] = 2 * ee.ref['a'] + 5 +assert ee.vv['k.1'] == 2 * 2 + 5 +assert ee.vv['b'] == 2 * 2 + 2 * 2 + 5 + +#-------------------------------------------------- + +ee.vars({ + 'a': 4., + 'b': '2 * a + 5', + 'k.1': '2 * a + 5', +}) + +env.new('bb', xt.Bend, k0='2 * b', length=3+env.vars['a'] + env.vars['b'], + h=5.) +assert env['bb'].k0 == 2 * (2 * 4 + 5) +assert env['bb'].length == 3 + 4 + 2 * 4 + 5 +assert env['bb'].h == 5. + +env.vars['a'] = 2. +assert env['bb'].k0 == 2 * (2 * 2 + 5) +assert env['bb'].length == 3 + 2 + 2 * 2 + 5 +assert env['bb'].h == 5. + +line = env.new_line([ + env.new('bb1', 'bb', length=3*env.vars['a'], at='2*a'), + env.place('bb', at=10 * env.vars['a'], from_='bb1'), +]) + +assert hasattr(env.ref['bb1'].length, '_value') # is a Ref +assert not hasattr(env['bb1'].length, '_value') # a number +assert env.ref['bb1'].length._value == 3 * 2 +assert env['bb1'].length == 3 * 2 + +assert hasattr(env.ref['bb1'].length, '_value') # is a Ref +assert not hasattr(env['bb1'].length, '_value') # a number +assert env.ref['bb1'].length._value == 3 * 2 +assert env['bb1'].length == 3 * 2 + +assert line['bb1'] is not env['bb'] +assert line['bb'] is env['bb'] + +a = env.vv['a'] +assert line['bb1'].length == 3 * a +assert line['bb1'].k0 == 2 * (2 * a + 5) +assert line['bb1'].h == 5. + +assert line['bb'].k0 == 2 * (2 * a + 5) +assert line['bb'].length == 3 + a + 2 * a + 5 +assert line['bb'].h == 5. + +tt = line.get_table(attr=True) +tt['s_center'] = tt['s'] + tt['length']/2 + +assert np.all(tt.name == np.array(['drift_1', 'bb1', 'drift_2', 'bb', '_end_point'])) + +assert tt['s_center', 'bb1'] == 2*a +assert tt['s_center', 'bb'] - tt['s_center', 'bb1'] == 10*a + +old_a = a +line.vars['a'] = 3. +a = line.vv['a'] +assert line['bb1'].length == 3 * a +assert line['bb1'].k0 == 2 * (2 * a + 5) +assert line['bb1'].h == 5. + +assert line['bb'].k0 == 2 * (2 * a + 5) +assert line['bb'].length == 3 + a + 2 * a + 5 +assert line['bb'].h == 5. + +tt_new = line.get_table(attr=True) + +# Drifts are not changed: +tt_new['length', 'drift_1'] == tt['length', 'drift_1'] +tt_new['length', 'drift_2'] == tt['length', 'drift_2'] + diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py new file mode 100644 index 000000000..e78f765fc --- /dev/null +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -0,0 +1,420 @@ +import xtrack as xt +import xobjects as xo +import numpy as np + +env = xt.Environment() +env.particle_ref = xt.Particles(p0c=2e9) + +n_bends_per_cell = 6 +n_cells_par_arc = 3 +n_arcs = 3 + +n_bends = n_bends_per_cell * n_cells_par_arc * n_arcs + +env.vars({ + 'l.mq': 0.5, + 'kqf': 0.027, + 'kqd': -0.0271, + 'l.mb': 10, + 'l.ms': 0.3, + 'k2sf': 0.001, + 'k2sd': -0.001, + 'angle.mb': 2 * np.pi / n_bends, + 'k0.mb': 'angle.mb / l.mb', + 'k0l.corrector': 0, + 'k1sl.corrector': 0, + 'l.halfcell': 38, +}) + +env.new('mb', xt.Bend, length='l.mb', k0='k0.mb', h='k0.mb') +env.new('mq', xt.Quadrupole, length='l.mq') +env.new('ms', xt.Sextupole, length='l.ms') +env.new('corrector', xt.Multipole, knl=[0], length=0.1) + +girder = env.new_line(components=[ + env.place('mq', at=1), + env.place('ms', at=0.8, from_='mq'), + env.place('corrector', at=-0.8, from_='mq'), +]) + +tt_girder = girder.get_table(attr=True) +assert np.all(tt_girder.name == np.array( + ['drift_1', 'corrector', 'drift_2', 'mq', 'drift_3', 'ms', '_end_point'])) +tt_girder['s_center'] = tt_girder['s'] + \ + tt_girder['length']/2 * np.float64(tt_girder['isthick']) +xo.assert_allclose(tt_girder['s_center', 'mq'], 1., atol=1e-14, rtol=0) +xo.assert_allclose(tt_girder['s_center', 'ms'] - tt_girder['s_center', 'mq'], 0.8, + atol=1e-14, rtol=0) +xo.assert_allclose( + tt_girder['s_center', 'corrector'] - tt_girder['s_center', 'mq'], -0.8, + atol=1e-14, rtol=0) + + +girder_f = girder.clone(name='f') +girder_d = girder.clone(name='d', mirror=True) +env.set('mq.f', k1='kqf') +env.set('mq.d', k1='kqd') + +# Check clone +tt_girder_f = girder_f.get_table(attr=True) +assert (~(tt_girder_f.isreplica)).all() +assert np.all(tt_girder_f.name == np.array( + ['drift_1.f', 'corrector.f', 'drift_2.f', 'mq.f', 'drift_3.f', 'ms.f', '_end_point'])) +tt_girder_f['s_center'] = (tt_girder_f['s'] + + tt_girder_f['length']/2 * np.float64(tt_girder_f['isthick'])) +xo.assert_allclose(tt_girder_f['s_center', 'mq.f'], 1., atol=1e-14, rtol=0) +xo.assert_allclose(tt_girder_f['s_center', 'ms.f'] - tt_girder_f['s_center', 'mq.f'], 0.8, + atol=1e-14, rtol=0) +xo.assert_allclose( + tt_girder_f['s_center', 'corrector.f'] - + tt_girder_f['s_center', 'mq.f'], -0.8, + atol=1e-14, rtol=0) + +# Check clone mirror +tt_girder_d = girder_d.get_table(attr=True) +assert (~(tt_girder_d.isreplica)).all() +len_girder = tt_girder_d.s[-1] +assert np.all(tt_girder_d.name == np.array( + ['ms.d', 'drift_3.d', 'mq.d', 'drift_2.d', 'corrector.d', 'drift_1.d', '_end_point'])) +tt_girder_d['s_center'] = (tt_girder_d['s'] + + tt_girder_d['length']/2 * np.float64(tt_girder_d['isthick'])) +xo.assert_allclose(tt_girder_d['s_center', 'mq.d'], + len_girder - 1., atol=1e-14, rtol=0) +xo.assert_allclose(tt_girder_d['s_center', 'ms.d'] - tt_girder_d['s_center', 'mq.d'], + -0.8, atol=1e-14, rtol=0) +xo.assert_allclose(tt_girder_d['s_center', 'corrector.d'] - tt_girder_d['s_center', 'mq.d'], + 0.8, atol=1e-14, rtol=0) + + +halfcell = env.new_line(components=[ + + # End of the half cell (will be mid of the cell) + env.new('mid', xt.Marker, at='l.halfcell'), + + # Bends + env.new('mb.2', 'mb', at='l.halfcell / 2'), + env.new('mb.1', 'mb', at='-l.mb - 1', from_='mb.2'), + env.new('mb.3', 'mb', at='l.mb + 1', from_='mb.2'), + + # Quadrupoles, sextupoles and correctors + env.place(girder_d, at=1.2), + env.place(girder_f, at='l.halfcell - 1.2'), + +]) + +l_hc = env.vv['l.halfcell'] +xo.assert_allclose(l_hc, l_hc, atol=1e-14, rtol=0) +tt_hc = halfcell.get_table(attr=True) +assert np.all(tt_hc.name == np.array( + ['drift_4', 'ms.d', 'drift_3.d', 'mq.d', 'drift_2.d', 'corrector.d', + 'drift_1.d', 'drift_5', 'mb.1', 'drift_6', 'mb.2', 'drift_7', + 'mb.3', 'drift_8', 'drift_1.f', 'corrector.f', 'drift_2.f', 'mq.f', + 'drift_3.f', 'ms.f', 'drift_9', 'mid', '_end_point'])) +assert np.all(tt_hc.element_type == np.array( + ['Drift', 'Sextupole', 'Drift', 'Quadrupole', 'Drift', 'Multipole', + 'Drift', 'Drift', 'Bend', 'Drift', 'Bend', 'Drift', 'Bend', + 'Drift', 'Drift', 'Multipole', 'Drift', 'Quadrupole', 'Drift', + 'Sextupole', 'Drift', 'Marker', ''])) +assert np.all(tt_hc.isreplica == False) +tt_hc['s_center'] = ( + tt_hc['s'] + tt_hc['length'] / 2 * np.float64(tt_hc['isthick'])) +xo.assert_allclose(tt_hc['s_center', 'mq.d'], + 1.2 - tt_girder_d.s[-1] / 2 + + tt_girder_d['s_center', 'mq.d'], + atol=1e-14, rtol=0) +xo.assert_allclose(tt_hc['s_center', 'ms.f'] - tt_hc['s_center', 'mq.f'], 0.8, + atol=1e-14, rtol=0) +xo.assert_allclose( + tt_hc['s_center', 'corrector.f'] - tt_hc['s_center', 'mq.f'], -0.8, + atol=1e-14, rtol=0) +xo.assert_allclose(tt_hc['s_center', 'ms.d'] - tt_hc['s_center', 'mq.d'], + -0.8, atol=1e-14, rtol=0) +xo.assert_allclose(tt_hc['s_center', 'corrector.d'] - tt_hc['s_center', 'mq.d'], + 0.8, atol=1e-14, rtol=0) +xo.assert_allclose(tt_hc['s_center', 'mb.2'], l_hc / 2, atol=1e-14, rtol=0) +xo.assert_allclose(tt_hc['s_center', 'mb.1'], tt_hc['s_center', 'mb.2'] - env.vv['l.mb'] - 1, + atol=1e-14, rtol=0) +xo.assert_allclose(tt_hc['s_center', 'mb.3'], tt_hc['s_center', 'mb.2'] + env.vv['l.mb'] + 1, + atol=1e-14, rtol=0) + + +hcell_left = halfcell.replicate(name='l', mirror=True) +hcell_right = halfcell.replicate(name='r') + +cell = env.new_line(components=[ + env.new('start', xt.Marker), + hcell_left, + hcell_right, + env.new('end', xt.Marker), +]) + +tt_cell = cell.get_table(attr=True) +tt_cell['s_center'] = ( + tt_cell['s'] + tt_cell['length'] / 2 * np.float64(tt_cell['isthick'])) +assert np.all(tt_cell.name == np.array( + ['start', 'mid.l', 'drift_9.l', 'ms.f.l', 'drift_3.f.l', 'mq.f.l', + 'drift_2.f.l', 'corrector.f.l', 'drift_1.f.l', 'drift_8.l', + 'mb.3.l', 'drift_7.l', 'mb.2.l', 'drift_6.l', 'mb.1.l', + 'drift_5.l', 'drift_1.d.l', 'corrector.d.l', 'drift_2.d.l', + 'mq.d.l', 'drift_3.d.l', 'ms.d.l', 'drift_4.l', 'drift_4.r', + 'ms.d.r', 'drift_3.d.r', 'mq.d.r', 'drift_2.d.r', 'corrector.d.r', + 'drift_1.d.r', 'drift_5.r', 'mb.1.r', 'drift_6.r', 'mb.2.r', + 'drift_7.r', 'mb.3.r', 'drift_8.r', 'drift_1.f.r', 'corrector.f.r', + 'drift_2.f.r', 'mq.f.r', 'drift_3.f.r', 'ms.f.r', 'drift_9.r', + 'mid.r', 'end', '_end_point'])) +assert np.all(tt_cell.element_type == np.array( + ['Marker', 'Marker', 'Drift', 'Sextupole', 'Drift', 'Quadrupole', + 'Drift', 'Multipole', 'Drift', 'Drift', 'Bend', 'Drift', 'Bend', + 'Drift', 'Bend', 'Drift', 'Drift', 'Multipole', 'Drift', + 'Quadrupole', 'Drift', 'Sextupole', 'Drift', 'Drift', 'Sextupole', + 'Drift', 'Quadrupole', 'Drift', 'Multipole', 'Drift', 'Drift', + 'Bend', 'Drift', 'Bend', 'Drift', 'Bend', 'Drift', 'Drift', + 'Multipole', 'Drift', 'Quadrupole', 'Drift', 'Sextupole', 'Drift', + 'Marker', 'Marker', ''])) +assert np.all(tt_cell.isreplica == np.array( + [False, True, True, True, True, True, True, True, True, + True, True, True, True, True, True, True, True, True, + True, True, True, True, True, True, True, True, True, + True, True, True, True, True, True, True, True, True, + True, True, True, True, True, True, True, True, True, + False, False])) + +tt_cell_stripped = tt_cell.rows[1:-2] # Remove _end_point and markers added in cell +tt_cell_second_half = tt_cell_stripped.rows[len(tt_cell_stripped)//2 :] +tt_cell_second_half.s_center -= tt_cell_second_half.s[0] +tt_hc_stripped = tt_hc.rows[:-1] # Remove _end_point +xo.assert_allclose(tt_cell_second_half.s_center, tt_hc_stripped.s_center, atol=5e-14, rtol=0) +tt_cell_first_half = tt_cell_stripped.rows[:len(tt_cell_stripped)//2] +s_center_mirrored_first_half = ( + tt_cell_stripped['s', len(tt_cell_stripped)//2] - tt_cell_first_half.s_center[::-1]) +xo.assert_allclose(s_center_mirrored_first_half, tt_hc_stripped.s_center, atol=5e-14, rtol=0) + +env.vars({ + 'kqf.ss': 0.027 / 2, + 'kqd.ss': -0.0271 / 2, +}) + +halfcell_ss = env.new_line(components=[ + + env.new('mid', xt.Marker, at='l.halfcell'), + + env.new('mq.ss.d', 'mq', k1='kqd.ss', at = '0.5 + l.mq / 2'), + env.new('mq.ss.f', 'mq', k1='kqf.ss', at = 'l.halfcell - l.mq / 2 - 0.5'), + + env.new('corrector.ss.v', 'corrector', at=0.75, from_='mq.ss.d'), + env.new('corrector.ss.h', 'corrector', at=-0.75, from_='mq.ss.f') +]) + +hcell_left_ss = halfcell_ss.replicate(name='l', mirror=True) +hcell_right_ss = halfcell_ss.replicate(name='r') +cell_ss = env.new_line(components=[ + env.new('start.ss', xt.Marker), + hcell_left_ss, + hcell_right_ss, + env.new('end.ss', xt.Marker), +]) + + +arc = env.new_line(components=[ + cell.replicate(name='cell.1'), + cell.replicate(name='cell.2'), + cell.replicate(name='cell.3'), +]) + +assert 'cell.2' in env.lines +tt_cell2 = env.lines['cell.2'].get_table(attr=True) +assert np.all(tt_cell2.name[:-1] == np.array([ + nn+'.cell.2' for nn in tt_cell.name[:-1]])) +assert np.all(tt_cell2.s == tt_cell.s) +assert tt_cell2.isreplica[:-1].all() +assert tt_cell2['parent_name', 'mq.d.l.cell.2'] == 'mq.d.l' +assert tt_cell2['parent_name', 'mq.f.l.cell.2'] == 'mq.f.l' +assert tt_cell['parent_name', 'mq.d.l'] == 'mq.d' +assert tt_cell['parent_name', 'mq.f.l'] == 'mq.f' + +tt_arc = arc.get_table(attr=True) +assert len(tt_arc) == 3 * (len(tt_cell)-1) + 1 +n_cell = len(tt_cell) - 1 +assert np.all(tt_arc.name[n_cell:2*n_cell] == tt_cell2.name[:-1]) +for nn in tt_cell2.name[:-1]: + assert arc[nn] is env[nn] + assert arc[nn] is env['cell.2'][nn] + +ss = env.new_line(components=[ + cell_ss.replicate('cell.1'), + cell_ss.replicate('cell.2'), +]) + +ring = env.new_line(components=[ + arc.replicate(name='arc.1'), + ss.replicate(name='ss.1'), + arc.replicate(name='arc.2'), + ss.replicate(name='ss.2'), + arc.replicate(name='arc.3'), + ss.replicate(name='ss.3'), +]) +tt_ring = ring.get_table(attr=True) +# Check length +xo.assert_allclose(tt_ring.s[-1], 2*l_hc * (n_cells_par_arc * n_arcs + 2*n_arcs), + atol=1e-12, rtol=0) +# Check closure +sv_ring = ring.survey() +xo.assert_allclose(sv_ring.X[-1], 0, atol=1e-12, rtol=0) +xo.assert_allclose(sv_ring.Y[-1], 0, atol=1e-12, rtol=0) +xo.assert_allclose(sv_ring.Z[-1], 0, atol=1e-12, rtol=0) + +xo.assert_allclose(sv_ring.angle.sum(), 2*np.pi, atol=1e-12, rtol=0) + +## Insertion + +env.vars({ + 'k1.q1': 0.025, + 'k1.q2': -0.025, + 'k1.q3': 0.025, + 'k1.q4': -0.02, + 'k1.q5': 0.025, +}) + +half_insertion = env.new_line(components=[ + + # Start-end markers + env.new('ip', xt.Marker), + env.new('e.insertion', xt.Marker, at=76), + + # Quads + env.new('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), + env.new('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), + env.new('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), + env.new('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), + env.new('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq', at=73), + + # Dipole correctors (will use h and v on the same corrector) + env.new('corrector.ss.1', 'corrector', at=0.75, from_='mq.1'), + env.new('corrector.ss.2', 'corrector', at=-0.75, from_='mq.2'), + env.new('corrector.ss.3', 'corrector', at=0.75, from_='mq.3'), + env.new('corrector.ss.4', 'corrector', at=-0.75, from_='mq.4'), + env.new('corrector.ss.5', 'corrector', at=0.75, from_='mq.5'), + +]) + +insertion = env.new_line([ + half_insertion.replicate('l', mirror=True), + half_insertion.replicate('r')]) + +ring2 = env.new_line(components=[ + env['arc.1'], + env['ss.1'], + env['arc.2'], + insertion, + env['arc.3'], + env['ss.3'], +]) + + +# # Check buffer behavior +ring2_sliced = ring2.select() +ring2_sliced.cut_at_s(np.arange(0, ring2.get_length(), 0.5)) + +opt = cell.match( + method='4d', + vary=xt.VaryList(['kqf', 'kqd'], step=1e-5), + targets=xt.TargetSet( + qx=0.333333, + qy=0.333333, + )) +tw_cell = cell.twiss4d() + +opt = cell_ss.match( + solve=False, + method='4d', + vary=xt.VaryList(['kqf.ss', 'kqd.ss'], step=1e-5), + targets=xt.TargetSet( + betx=tw_cell.betx[-1], bety=tw_cell.bety[-1], at='start.ss', + )) +opt.solve() + +tw_arc = arc.twiss4d() + +opt = half_insertion.match( + solve=False, + betx=tw_arc.betx[0], bety=tw_arc.bety[0], + alfx=tw_arc.alfx[0], alfy=tw_arc.alfy[0], + init_at='e.insertion', + start='ip', end='e.insertion', + vary=xt.VaryList(['k1.q1', 'k1.q2', 'k1.q3', 'k1.q4'], step=1e-5), + targets=[ + xt.TargetSet(alfx=0, alfy=0, at='ip'), + xt.Target(lambda tw: tw.betx[0] - tw.bety[0], 0), + xt.Target(lambda tw: tw.betx.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.bety.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.betx.min(), xt.GreaterThan(2)), + xt.Target(lambda tw: tw.bety.min(), xt.GreaterThan(2)), + ] +) +opt.step(40) +opt.solve() + +# Check that the cell is matched to the rest of the ring +tw = ring.twiss4d() +tw_cell_from_ring = tw.rows['start.cell.3.arc.2':'end.cell.3.arc.2'] +xo.assert_allclose(tw_cell_from_ring.betx, tw_cell.betx[:-1], atol=0, rtol=2e-4) +xo.assert_allclose(tw_cell_from_ring.bety, tw_cell.bety[:-1], atol=0, rtol=2e-4) + +tw2 = ring2.twiss4d() +tw_cell_from_ring2 = tw2.rows['start.cell.3.arc.2':'end.cell.3.arc.2'] +xo.assert_allclose(tw_cell_from_ring2.betx, tw_cell.betx[:-1], atol=0, rtol=2e-4) +xo.assert_allclose(tw_cell_from_ring2.bety, tw_cell.bety[:-1], atol=0, rtol=2e-4) + +# Check select +cell3_select = ring2.select(start='start.cell.3.arc.2', end='end.cell.3.arc.2', + name='cell3_copy') +assert 'cell3_copy' in env.lines +assert env.lines['cell3_copy'] is cell3_select +assert cell3_select._element_dict is env.element_dict +assert cell3_select.element_names[0] == 'start.cell.3.arc.2' +assert cell3_select.element_names[-1] == 'end.cell.3.arc.2' +assert (np.array(cell3_select.element_names) == np.array( + tw.rows['start.cell.3.arc.2':'end.cell.3.arc.2'].name)).all() + +# Check that they share the _element_dict +assert cell._element_dict is env.element_dict +assert halfcell._element_dict is env.element_dict +assert halfcell_ss._element_dict is env.element_dict +assert cell_ss._element_dict is env.element_dict +assert insertion._element_dict is env.element_dict +assert ring2._element_dict is env.element_dict + +cell3_select.twiss4d() + +tw2_slice = ring2_sliced.twiss4d() +xo.assert_allclose(tw2_slice['betx', 'ip.l'], tw2['betx', 'ip.l'], atol=0, rtol=2e-4) +xo.assert_allclose(tw2_slice['bety', 'ip.l'], tw2['bety', 'ip.l'], atol=0, rtol=2e-4) +xo.assert_allclose(tw2_slice['alfx', 'ip.l'], 0, atol=1e-6, rtol=0) +xo.assert_allclose(tw2_slice['alfy', 'ip.l'], 0, atol=1e-6, rtol=0) +xo.assert_allclose(tw2_slice['dx', 'ip.l'], 0, atol=1e-4, rtol=0) +xo.assert_allclose(tw2_slice['dpx', 'ip.l'], 0, atol=1e-6, rtol=0) +xo.assert_allclose(tw2_slice['dy', 'ip.l'], 0, atol=1e-4, rtol=0) +xo.assert_allclose(tw2_slice['dpy', 'ip.l'], 0, atol=1e-6, rtol=0) + +import matplotlib.pyplot as plt +plt.close('all') +for ii, rr in enumerate([ring, ring2_sliced]): + + ttww = rr.twiss4d() + + fig = plt.figure(ii, figsize=(6.4*1.2, 4.8)) + ax1 = fig.add_subplot(2, 1, 1) + pltbet = ttww.plot('betx bety', ax=ax1) + ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) + pltdx = ttww.plot('dx', ax=ax2) + fig.subplots_adjust(right=.85) + pltbet.move_legend(1.2,1) + pltdx.move_legend(1.2,1) + +ring2.survey().plot() + + +plt.show() + + + diff --git a/examples/lattice_design_shortcuts/t010_new_syntax.py b/examples/lattice_design_shortcuts/t010_new_syntax.py new file mode 100644 index 000000000..4990a39d1 --- /dev/null +++ b/examples/lattice_design_shortcuts/t010_new_syntax.py @@ -0,0 +1,49 @@ +import xtrack as xt +import numpy as np +env = xt.Environment() + +# Variables +env["a"] = 3 +env["b"] = np.array([1, 3]) # to_dict issue, json reconstruc np.array natively +env["fun"]= math.sin + +env.ref["b"][0] = env.ref["a"] * 3 # +env.set("b[0]", "fun(a*3)") # +env["c"] = env.ref["b"][0] * 3 # + +#Elements +env.new("mq0","Quadrupole",l="b[0]") +env.ref["mq0"] + +#Lines +env["b1"] # line +env.b1 # line has the right to pollute namespace +env.ref["b1"] # ref of a line + +# General +env[] -> Python values +env.ref -> reference + +# METADATA +env["mq"].metadata={"slot_id":123123,"layout_id":21332141} +# OR +env.extra_element_attributes=["slot_id","layout_id","polarity"] + +#Containenr-like with benefits +env.vars # references duplicated by ref +env.elems # values like element_dict +env.lines # lines values + +env.vars.get_table() +env.funcs.get_table() +env.elems.get_table() +env.lines.get_table() +#introduce version and metadata on json + +#-------------------------------------------------- + +#to be deprecated but supported +env.vv +env.element_refs + + diff --git a/examples/radiation/009a_sps_with_vertical_bump.py b/examples/radiation/009a_sps_with_vertical_bump.py index e08ac56e4..0bfef821f 100644 --- a/examples/radiation/009a_sps_with_vertical_bump.py +++ b/examples/radiation/009a_sps_with_vertical_bump.py @@ -47,17 +47,17 @@ deferred_expressions=True) line.particle_ref = xt.Particles(mass0=xt.ELECTRON_MASS_EV, q0=-1, gamma0=mad.sequence.sps.beam.gamma) -line.cycle('bpv.11706_entry', inplace=True) +line.cycle('bpv.11706', inplace=True) -line.insert_element(element=line['actcse.31632'].copy(), index='bpv.11706_entry', +line.insert_element(element=line['actcse.31632'].copy(), index='bpv.11706', name='cav1') -line.insert_element(element=line['actcse.31632'].copy(), index='bpv.21508_entry', +line.insert_element(element=line['actcse.31632'].copy(), index='bpv.21508', name='cav2') -line.insert_element(element=line['actcse.31632'].copy(), index='bpv.41508_entry', +line.insert_element(element=line['actcse.31632'].copy(), index='bpv.41508', name='cav4') -line.insert_element(element=line['actcse.31632'].copy(), index='bpv.51508_entry', +line.insert_element(element=line['actcse.31632'].copy(), index='bpv.51508', name='cav5') -line.insert_element(element=line['actcse.31632'].copy(), index='bpv.61508_entry', +line.insert_element(element=line['actcse.31632'].copy(), index='bpv.61508', name='cav6') tt = line.get_table() @@ -105,10 +105,10 @@ line.element_refs['mdv.53507'].ksl[0] = line.vars['mdv.53507.ksl0'] # Kill sextupoles in the bump -line.element_refs['lsf.53205..0'].knl[2] = 0 -line.element_refs['lsd.53505..0'].knl[2] = 0 -line.element_refs['lsf.53605..0'].knl[2] = 0 -line.element_refs['lsd.60105..0'].knl[2] = 0 +line.element_refs['lsf.53205'].k2l = 0 +line.element_refs['lsd.53505'].k2l = 0 +line.element_refs['lsf.53605'].k2l = 0 +line.element_refs['lsd.60105'].k2l = 0 tw0 = line.twiss() opt_bump = line.match( diff --git a/examples/symm_twiss_and_match/t000_test_twiss_symm.py b/examples/symm_twiss_and_match/t000_test_twiss_symm.py index 4e54f9fb9..8905a0358 100644 --- a/examples/symm_twiss_and_match/t000_test_twiss_symm.py +++ b/examples/symm_twiss_and_match/t000_test_twiss_symm.py @@ -18,6 +18,11 @@ ) half_cell.particle_ref = xt.Particles(p0c=2e9) +# Add observation points every 1 m (to see betas inside bends) +half_cell.discard_tracker() +s_cut = np.arange(0, half_cell.get_length(), 1.) +half_cell.cut_at_s(s_cut) + tw_half_cell = half_cell.twiss4d(init='periodic_symmetric', # <--- periodic-symmetric boundary strengths=True # to get the strengths in table ) @@ -44,6 +49,12 @@ } ) cell.particle_ref = xt.Particles(p0c=2e9) + +# Add observation points every 1 m (to see betas inside bends) +cell.discard_tracker() +s_cut = np.arange(0, cell.get_length(), 1.) +cell.cut_at_s(s_cut) + tw_cell = cell.twiss4d(strengths=True) xo.assert_allclose(tw_half_cell.betx[:-1], # remove '_end_point' @@ -134,6 +145,10 @@ tw_cell.rows[:'mid_cell'].dx, atol=1e-8, rtol=0) xo.assert_allclose(tw_half_cell.dpx[:-1], # remove '_end_point' tw_cell.rows[:'mid_cell'].dpx, atol=1e-8, rtol=0) +xo.assert_allclose(tw_half_cell.ddx[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].ddx, atol=1e-7, rtol=0) +xo.assert_allclose(tw_half_cell.ddpx[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].ddpx, atol=1e-8, rtol=0) xo.assert_allclose(tw_half_cell.ax_chrom[:-1], # remove '_end_point' tw_cell.rows[:'mid_cell'].ax_chrom, atol=1e-5, rtol=0) @@ -149,3 +164,44 @@ xo.assert_allclose(tw_half_cell.dqx, tw_cell.dqx / 2, atol=1e-6, rtol=0) xo.assert_allclose(tw_half_cell.dqy, tw_cell.dqy / 2, atol=1e-6, rtol=0) +tw_off_mom_cell = cell.twiss4d(strengths=True, delta0=1e-3) +tw_off_mom_half_cell = half_cell.twiss4d( + init='periodic_symmetric', strengths=True, delta0=1e-3) + +xo.assert_allclose(tw_off_mom_half_cell.x[:-1], + tw_off_mom_cell.rows[:'mid_cell'].x, atol=1e-12, rtol=0) + +import matplotlib.pyplot as plt +plt.close('all') + +fig1 = plt.figure(1, figsize=(6.4*1.2, 4.8)) +ax1 = fig1.add_subplot(211) +plt_cell = tw_cell.plot(ax=ax1) +plt_cell.move_legend(left=1.4, bottom=1) +ax1.set_title('Full cell') + +ax2 = fig1.add_subplot(212, sharex=ax1, sharey=ax1) +plt_half_cell = tw_half_cell.plot(ax=ax2) +plt_half_cell.move_legend(left=1.4, bottom=1) +ax2.set_title('Half cell') + +ax1.set_xlim(0, tw_cell.s[-1]) +fig1.subplots_adjust(right=0.73, hspace=0.5) + +fig2 = plt.figure(2, figsize=(6.4*1.2, 4.8)) +ax1 = fig2.add_subplot(211) +plt_cell = tw_cell.plot('ddx', ax=ax1) +plt_cell.move_legend(left=1.4, bottom=1) +plt_cell.ax.set_title('Full cell') + +ax2 = fig2.add_subplot(212, sharex=ax1, sharey=ax1) +plt_half_cell = tw_half_cell.plot('ddx', ax=ax2) +plt_half_cell.move_legend(left=1.4, bottom=1) +plt_half_cell.ax.set_title('Half cell') + +ax1.set_xlim(0, tw_cell.s[-1]) +fig2.subplots_adjust(right=0.73, hspace=0.5) + + + +plt.show() diff --git a/examples/twiss/031_twiss_plot.py b/examples/twiss/031_twiss_plot.py new file mode 100644 index 000000000..26e5683db --- /dev/null +++ b/examples/twiss/031_twiss_plot.py @@ -0,0 +1,27 @@ +# copyright ############################### # +# This file is part of the Xtrack Package. # +# Copyright (c) CERN, 2021. # +# ######################################### # + +import xtrack as xt + +import matplotlib.pyplot as plt + +# Load a line and build tracker +line = xt.Line.from_json("../../test_data/hllhc15_thick/lhc_thick_with_knobs.json") +line.particle_ref = xt.Particles(mass0=xt.PROTON_MASS_EV, q0=1, energy0=7e12) + +# Twiss +tw = line.twiss4d() + +pl = tw.plot(figlabel="tw") +pl = tw.plot(figlabel="tw") + +fig = plt.figure() +ax1 = plt.subplot(211) +ax2 = plt.subplot(212, sharex=ax1) +pl1 = tw.plot(ax=ax1) +pl2 = tw.plot(yl="mux muy", ax=ax2) +fig.subplots_adjust(right=0.7) +pl1.move_legend(1.5, 1.0) +pl2.move_legend(1.5, 1.0) diff --git a/tests/test_environment.py b/tests/test_environment.py new file mode 100644 index 000000000..df22fff29 --- /dev/null +++ b/tests/test_environment.py @@ -0,0 +1,443 @@ +import xtrack as xt +import xobjects as xo +import numpy as np +import pytest + +@pytest.mark.parametrize('container_type', ['env', 'line']) +def test_vars_and_element_access_modes(container_type): + + env = xt.Environment() + + env.vars({ + 'k.1': 1., + 'a': 2., + 'b': '2 * a + k.1', + }) + + line = env.new_line([]) + + ee = {'env': env, 'line': line}[container_type] + + assert ee.vv['b'] == 2 * 2 + 1 + + ee.vars['a'] = ee.vars['k.1'] + assert ee.vv['b'] == 2 * 1 + 1 + + ee.vars(a=3.) + ee.vars({'k.1': 'a'}) + assert ee.vv['k.1'] == 3. + assert ee.vv['b'] == 2 * 3 + 3. + + ee.vars['k.1'] = 2 * ee.vars['a'] + 5 + assert ee.vv['k.1'] == 2 * 3 + 5 + assert ee.vv['b'] == 2 * 3 + 2 * 3 + 5 + + ee.vars.set('a', 4.) + assert ee.vv['k.1'] == 2 * 4 + 5 + assert ee.vv['b'] == 2 * 4 + 2 * 4 + 5 + + ee.vars.set('k.1', '2*a + 5') + assert ee.vv['k.1'] == 2 * 4 + 5 + assert ee.vv['b'] == 2 * 4 + 2 * 4 + 5 + + ee.vars.set('k.1', 3 * ee.vars['a'] + 6) + assert ee.vv['k.1'] == 3 * 4 + 6 + assert ee.vv['b'] == 2 * 4 + 3 * 4 + 6 + + env.set('c', '2*b') + assert env.vv['c'] == 2 * (2 * 4 + 3 * 4 + 6) + env.set('d', 6) + assert env.vv['d'] == 6 + env.set('d', '7') + assert env.vv['d'] == 7 + + ee.set('a', 0.) + assert ee.vv['k.1'] == 3 * 0 + 6 + assert ee.vv['b'] == 2 * 0 + 3 * 0 + 6 + + ee.set('a', 2.) + ee.set('k.1', '2 * a + 5') + assert ee.vv['k.1'] == 2 * 2 + 5 + assert ee.vv['b'] == 2 * 2 + 2 * 2 + 5 + + ee.set('k.1', 3 * ee.vars['a'] + 6) + assert ee.vv['k.1'] == 3 * 2 + 6 + assert ee.vv['b'] == 2 * 2 + 3 * 2 + 6 + + assert hasattr(ee.ref['k.1'], '_value') # is a Ref + + ee.ref['a'] = 0 + assert ee.vv['k.1'] == 3 * 0 + 6 + assert ee.vv['b'] == 2 * 0 + 3 * 0 + 6 + + ee.ref['a'] = 2 + ee.ref['k.1'] = 2 * ee.ref['a'] + 5 + assert ee.vv['k.1'] == 2 * 2 + 5 + assert ee.vv['b'] == 2 * 2 + 2 * 2 + 5 + + #-------------------------------------------------- + + ee.vars({ + 'a': 4., + 'b': '2 * a + 5', + 'k.1': '2 * a + 5', + }) + + env.new('bb', xt.Bend, k0='2 * b', length=3+env.vars['a'] + env.vars['b'], + h=5.) + assert env['bb'].k0 == 2 * (2 * 4 + 5) + assert env['bb'].length == 3 + 4 + 2 * 4 + 5 + assert env['bb'].h == 5. + + env.vars['a'] = 2. + assert env['bb'].k0 == 2 * (2 * 2 + 5) + assert env['bb'].length == 3 + 2 + 2 * 2 + 5 + assert env['bb'].h == 5. + + line = env.new_line([ + env.new('bb1', 'bb', length=3*env.vars['a'], at='2*a'), + env.place('bb', at=10 * env.vars['a'], from_='bb1'), + ]) + + assert hasattr(env.ref['bb1'].length, '_value') # is a Ref + assert not hasattr(env['bb1'].length, '_value') # a number + assert env.ref['bb1'].length._value == 3 * 2 + assert env['bb1'].length == 3 * 2 + + assert hasattr(env.ref['bb1'].length, '_value') # is a Ref + assert not hasattr(env['bb1'].length, '_value') # a number + assert env.ref['bb1'].length._value == 3 * 2 + assert env['bb1'].length == 3 * 2 + + assert line['bb1'] is not env['bb'] + assert line['bb'] is env['bb'] + + a = env.vv['a'] + assert line['bb1'].length == 3 * a + assert line['bb1'].k0 == 2 * (2 * a + 5) + assert line['bb1'].h == 5. + + assert line['bb'].k0 == 2 * (2 * a + 5) + assert line['bb'].length == 3 + a + 2 * a + 5 + assert line['bb'].h == 5. + + tt = line.get_table(attr=True) + tt['s_center'] = tt['s'] + tt['length']/2 + + assert np.all(tt.name == np.array(['drift_1', 'bb1', 'drift_2', 'bb', '_end_point'])) + + assert tt['s_center', 'bb1'] == 2*a + assert tt['s_center', 'bb'] - tt['s_center', 'bb1'] == 10*a + + old_a = a + line.vars['a'] = 3. + a = line.vv['a'] + assert line['bb1'].length == 3 * a + assert line['bb1'].k0 == 2 * (2 * a + 5) + assert line['bb1'].h == 5. + + assert line['bb'].k0 == 2 * (2 * a + 5) + assert line['bb'].length == 3 + a + 2 * a + 5 + assert line['bb'].h == 5. + + tt_new = line.get_table(attr=True) + + # Drifts are not changed: + tt_new['length', 'drift_1'] == tt['length', 'drift_1'] + tt_new['length', 'drift_2'] == tt['length', 'drift_2'] + +def test_element_placing_at_s(): + + env = xt.Environment() + + env.vars({ + 'l.b1': 1.0, + 'l.q1': 0.5, + 's.ip': 10, + 's.left': -5, + 's.right': 5, + 'l.before_right': 1, + 'l.after_left2': 0.5, + }) + + # names, tab_sorted = handle_s_places(seq) + line = env.new_line(components=[ + env.new('b1', xt.Bend, length='l.b1'), + env.new('q1', xt.Quadrupole, length='l.q1'), + env.new('ip', xt.Marker, at='s.ip'), + ( + env.new('before_before_right', xt.Marker), + env.new('before_right', xt.Sextupole, length=1), + env.new('right',xt.Quadrupole, length=0.8, at='s.right', from_='ip'), + env.new('after_right', xt.Marker), + env.new('after_right2', xt.Marker), + ), + env.new('left', xt.Quadrupole, length=1, at='s.left', from_='ip'), + env.new('after_left', xt.Marker), + env.new('after_left2', xt.Bend, length='l.after_left2'), + ]) + + tt = line.get_table(attr=True) + tt['s_center'] = tt['s'] + tt['length']/2 + assert np.all(tt.name == np.array([ + 'b1', 'q1', 'drift_1', 'left', 'after_left', 'after_left2', + 'drift_2', 'ip', 'drift_3', 'before_before_right', 'before_right', + 'right', 'after_right', 'after_right2', '_end_point'])) + + xo.assert_allclose(env['b1'].length, 1.0, rtol=0, atol=1e-14) + xo.assert_allclose(env['q1'].length, 0.5, rtol=0, atol=1e-14) + xo.assert_allclose(tt['s', 'ip'], 10, rtol=0, atol=1e-14) + xo.assert_allclose(tt['s', 'before_before_right'], tt['s', 'before_right'], + rtol=0, atol=1e-14) + xo.assert_allclose(tt['s_center', 'before_right'] - tt['s_center', 'right'], + -(1 + 0.8)/2, rtol=0, atol=1e-14) + xo.assert_allclose(tt['s_center', 'right'] - tt['s', 'ip'], 5, rtol=0, atol=1e-14) + xo.assert_allclose(tt['s_center', 'after_right'] - tt['s_center', 'right'], + 0.8/2, rtol=0, atol=1e-14) + xo.assert_allclose(tt['s_center', 'after_right2'] - tt['s_center', 'right'], + 0.8/2, rtol=0, atol=1e-14) + xo.assert_allclose(tt['s_center', 'left'] - tt['s_center', 'ip'], -5, + rtol=0, atol=1e-14) + xo.assert_allclose(tt['s_center', 'after_left'] - tt['s_center', 'left'], 1/2, + rtol=0, atol=1e-14) + xo.assert_allclose(tt['s_center', 'after_left2'] - tt['s_center', 'after_left'], + 0.5/2, rtol=0, atol=1e-14) + + + # import matplotlib.pyplot as plt + # plt.close('all') + # line.survey().plot() + + # plt.show() + +def test_assemble_ring(): + + env = xt.Environment() + env.particle_ref = xt.Particles(p0c=2e9) + + n_bends_per_cell = 6 + n_cells_par_arc = 3 + n_arcs = 3 + + n_bends = n_bends_per_cell * n_cells_par_arc * n_arcs + + env.vars({ + 'l.mq': 0.5, + 'kqf': 0.027, + 'kqd': -0.0271, + 'l.mb': 10, + 'l.ms': 0.3, + 'k2sf': 0.001, + 'k2sd': -0.001, + 'angle.mb': 2 * np.pi / n_bends, + 'k0.mb': 'angle.mb / l.mb', + 'k0l.corrector': 0, + 'k1sl.corrector': 0, + 'l.halfcell': 38, + }) + + env.new('mb', xt.Bend, length='l.mb', k0='k0.mb', h='k0.mb') + env.new('mq', xt.Quadrupole, length='l.mq') + env.new('ms', xt.Sextupole, length='l.ms') + env.new('corrector', xt.Multipole, knl=[0], length=0.1) + + girder = env.new_line(components=[ + env.place('mq', at=1), + env.place('ms', at=0.8, from_='mq'), + env.place('corrector', at=-0.8, from_='mq'), + ]) + + girder_f = girder.clone(name='f') + girder_d = girder.clone(name='d', mirror=True) + env.set('mq.f', k1='3') # Test string with value + assert env['mq.f'].k1 == 3. + env.set('mq.f', k1='kqf') + env.set('mq.d', k1='kqd') + + halfcell = env.new_line(components=[ + + # End of the half cell (will be mid of the cell) + env.new('mid', xt.Marker, at='l.halfcell'), + + # Bends + env.new('mb.2', 'mb', at='l.halfcell / 2'), + env.new('mb.1', 'mb', at='-l.mb - 1', from_='mb.2'), + env.new('mb.3', 'mb', at='l.mb + 1', from_='mb.2'), + + # Quadrupoles, sextupoles and correctors + env.place(girder_d, at=1.2), + env.place(girder_f, at='l.halfcell - 1.2'), + + ]) + + + hcell_left = halfcell.replicate(name='l', mirror=True) + hcell_right = halfcell.replicate(name='r') + + cell = env.new_line(components=[ + env.new('start', xt.Marker), + hcell_left, + hcell_right, + env.new('end', xt.Marker), + ]) + + opt = cell.match( + method='4d', + vary=xt.VaryList(['kqf', 'kqd'], step=1e-5), + targets=xt.TargetSet( + qx=0.333333, + qy=0.333333, + )) + tw_cell = cell.twiss4d() + + + env.vars({ + 'kqf.ss': 0.027 / 2, + 'kqd.ss': -0.0271 / 2, + }) + + halfcell_ss = env.new_line(components=[ + + env.new('mid', xt.Marker, at='l.halfcell'), + + env.new('mq.ss.d', 'mq', k1='kqd.ss', at = '0.5 + l.mq / 2'), + env.new('mq.ss.f', 'mq', k1='kqf.ss', at = 'l.halfcell - l.mq / 2 - 0.5'), + + env.new('corrector.ss.v', 'corrector', at=0.75, from_='mq.ss.d'), + env.new('corrector.ss.h', 'corrector', at=-0.75, from_='mq.ss.f') + ]) + + hcell_left_ss = halfcell_ss.replicate(name='l', mirror=True) + hcell_right_ss = halfcell_ss.replicate(name='r') + cell_ss = env.new_line(components=[ + env.new('start.ss', xt.Marker), + hcell_left_ss, + hcell_right_ss, + env.new('end.ss', xt.Marker), + ]) + + opt = cell_ss.match( + solve=False, + method='4d', + vary=xt.VaryList(['kqf.ss', 'kqd.ss'], step=1e-5), + targets=xt.TargetSet( + betx=tw_cell.betx[-1], bety=tw_cell.bety[-1], at='start.ss', + )) + opt.solve() + + + arc = env.new_line(components=[ + cell.replicate(name='cell.1'), + cell.replicate(name='cell.2'), + cell.replicate(name='cell.3'), + ]) + + + ss = env.new_line(components=[ + cell_ss.replicate('cell.1'), + cell_ss.replicate('cell.2'), + ]) + + ring = env.new_line(components=[ + arc.replicate(name='arc.1'), + ss.replicate(name='ss.1'), + arc.replicate(name='arc.2'), + ss.replicate(name='ss.2'), + arc.replicate(name='arc.3'), + ss.replicate(name='ss.3'), + ]) + + ## Insertion + + env.vars({ + 'k1.q1': 0.025, + 'k1.q2': -0.025, + 'k1.q3': 0.025, + 'k1.q4': -0.02, + 'k1.q5': 0.025, + }) + + half_insertion = env.new_line(components=[ + + # Start-end markers + env.new('ip', xt.Marker), + env.new('e.insertion', xt.Marker, at=76), + + # Quads + env.new('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), + env.new('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), + env.new('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), + env.new('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), + env.new('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq', at=73), + + # Dipole correctors (will use h and v on the same corrector) + env.new('corrector.ss.1', 'corrector', at=0.75, from_='mq.1'), + env.new('corrector.ss.2', 'corrector', at=-0.75, from_='mq.2'), + env.new('corrector.ss.3', 'corrector', at=0.75, from_='mq.3'), + env.new('corrector.ss.4', 'corrector', at=-0.75, from_='mq.4'), + env.new('corrector.ss.5', 'corrector', at=0.75, from_='mq.5'), + + ]) + + tw_arc = arc.twiss4d() + + opt = half_insertion.match( + solve=False, + betx=tw_arc.betx[0], bety=tw_arc.bety[0], + alfx=tw_arc.alfx[0], alfy=tw_arc.alfy[0], + init_at='e.insertion', + start='ip', end='e.insertion', + vary=xt.VaryList(['k1.q1', 'k1.q2', 'k1.q3', 'k1.q4'], step=1e-5), + targets=[ + xt.TargetSet(alfx=0, alfy=0, at='ip'), + xt.Target(lambda tw: tw.betx[0] - tw.bety[0], 0), + xt.Target(lambda tw: tw.betx.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.bety.max(), xt.LessThan(400)), + xt.Target(lambda tw: tw.betx.min(), xt.GreaterThan(2)), + xt.Target(lambda tw: tw.bety.min(), xt.GreaterThan(2)), + ] + ) + opt.step(40) + opt.solve() + + insertion = env.new_line([ + half_insertion.replicate('l', mirror=True), + half_insertion.replicate('r')]) + + + + ring2 = env.new_line(components=[ + arc.replicate(name='arcc.1'), + ss.replicate(name='sss.2'), + arc.replicate(name='arcc.2'), + insertion, + arc.replicate(name='arcc.3'), + ss.replicate(name='sss.3') + ]) + + + # # Check buffer behavior + ring2_sliced = ring2.select() + ring2_sliced.cut_at_s(np.arange(0, ring2.get_length(), 0.5)) + + + # import matplotlib.pyplot as plt + # plt.close('all') + # for ii, rr in enumerate([ring, ring2_sliced]): + + # tw = rr.twiss4d() + + # fig = plt.figure(ii, figsize=(6.4*1.2, 4.8)) + # ax1 = fig.add_subplot(2, 1, 1) + # pltbet = tw.plot('betx bety', ax=ax1) + # ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) + # pltdx = tw.plot('dx', ax=ax2) + # fig.subplots_adjust(right=.85) + # pltbet.move_legend(1.2,1) + # pltdx.move_legend(1.2,1) + + # ring2.survey().plot() + + # plt.show() + + diff --git a/tests/test_line.py b/tests/test_line.py index b80b4295f..559844567 100644 --- a/tests/test_line.py +++ b/tests/test_line.py @@ -656,36 +656,50 @@ def test_from_json_to_json(tmp_path): } line.metadata = example_metadata + def asserts(): + assert len(result.element_dict.keys()) == 2 + assert result.element_names == ['m', 'd', 'm', 'd'] + + assert isinstance(result['m'], xt.Multipole) + assert (result['m'].knl == [1, 2]).all() + + assert isinstance(result['d'], xt.Drift) + assert result['d'].length == 1 + + assert result.metadata == example_metadata + result.metadata['qx']['lhcb1'] = result.metadata['qx']['lhcb1'] + 1 + assert result.metadata != example_metadata + result.metadata['qx']['lhcb1'] = result.metadata['qx']['lhcb1'] - 1 + line.to_json(tmp_path / 'test.json') result = xt.Line.from_json(tmp_path / 'test.json') - assert len(result.element_dict.keys()) == 2 - assert result.element_names == ['m', 'd', 'm', 'd'] + asserts() - assert isinstance(result['m'], xt.Multipole) - assert (result['m'].knl == [1, 2]).all() + with open(tmp_path / 'test2.json', 'w') as f: + line.to_json(f) - assert isinstance(result['d'], xt.Drift) - assert result['d'].length == 1 + with open(tmp_path / 'test2.json', 'r') as f: + result = xt.Line.from_json(f) + + asserts() with open(tmp_path / 'test2.json', 'w') as f: - line.to_json(f) + line.to_json(f,indent=None) with open(tmp_path / 'test2.json', 'r') as f: result = xt.Line.from_json(f) - assert len(result.element_dict.keys()) == 2 - assert result.element_names == ['m', 'd', 'm', 'd'] + asserts() - assert isinstance(result['m'], xt.Multipole) - assert (result['m'].knl == [1, 2]).all() + with open(tmp_path / 'test2.json.gz', 'w') as f: + line.to_json(f,indent=2) - assert isinstance(result['d'], xt.Drift) - assert result['d'].length == 1 + with open(tmp_path / 'test2.json.gz', 'r') as f: + result = xt.Line.from_json(f) + + asserts() - assert result.metadata == example_metadata - result.metadata['qx']['lhcb1'] = result.metadata['qx']['lhcb1'] + 1 - assert result.metadata != example_metadata @for_all_test_contexts def test_config_propagation(test_context): @@ -1060,4 +1074,4 @@ def test_get_strengths(test_context): rtol=0, atol=1e-14) xo.assert_allclose(line['mbw.a6l3.b2'].h, str_table['angle_rad', 'mbw.a6l3.b2'] / str_table['length', 'mbw.a6l3.b2'], - rtol=0, atol=1e-14) \ No newline at end of file + rtol=0, atol=1e-14) diff --git a/tests/test_periodic_symmetric_twiss_and_match.py b/tests/test_periodic_symmetric_twiss_and_match.py index 0719bf6fb..b51de4b06 100644 --- a/tests/test_periodic_symmetric_twiss_and_match.py +++ b/tests/test_periodic_symmetric_twiss_and_match.py @@ -137,6 +137,8 @@ def test_periodic_symmetric_twiss_and_match(): tw_cell.rows[:'mid_cell'].dx, atol=1e-8, rtol=0) xo.assert_allclose(tw_half_cell.dpx[:-1], # remove '_end_point' tw_cell.rows[:'mid_cell'].dpx, atol=1e-8, rtol=0) + xo.assert_allclose(tw_half_cell.ddx[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].ddx, atol=1e-7, rtol=0) xo.assert_allclose(tw_half_cell.ax_chrom[:-1], # remove '_end_point' tw_cell.rows[:'mid_cell'].ax_chrom, atol=1e-5, rtol=0) @@ -146,9 +148,17 @@ def test_periodic_symmetric_twiss_and_match(): tw_cell.rows[:'mid_cell'].bx_chrom, atol=1e-5, rtol=0) xo.assert_allclose(tw_half_cell.by_chrom[:-1], # remove '_end_point' tw_cell.rows[:'mid_cell'].by_chrom, atol=1e-5, rtol=0) + xo.assert_allclose(tw_half_cell.ddx[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].ddx, atol=1e-7, rtol=0) xo.assert_allclose(tw_half_cell.qx, tw_cell.qx / 2, atol=1e-9, rtol=0) xo.assert_allclose(tw_half_cell.qy, tw_cell.qy / 2, atol=1e-9, rtol=0) xo.assert_allclose(tw_half_cell.dqx, tw_cell.dqx / 2, atol=1e-6, rtol=0) xo.assert_allclose(tw_half_cell.dqy, tw_cell.dqy / 2, atol=1e-6, rtol=0) + tw_off_mom_cell = cell.twiss4d(strengths=True, delta0=1e-3) + tw_off_mom_half_cell = half_cell.twiss4d( + init='periodic_symmetric', strengths=True, delta0=1e-3) + + xo.assert_allclose(tw_off_mom_half_cell.x[:-1], + tw_off_mom_cell.rows[:'mid_cell'].x, atol=1e-12, rtol=0) diff --git a/tests/test_slice_and_insert_with_replicas.py b/tests/test_slice_and_insert_with_replicas.py index 2430037d9..508a1b9a7 100644 --- a/tests/test_slice_and_insert_with_replicas.py +++ b/tests/test_slice_and_insert_with_replicas.py @@ -236,6 +236,8 @@ def test_slice_thick_and_insert_with_replicas(test_context): element_names=list(elements.keys())) line.build_tracker(_context=test_context) + assert line['e2']._movable + element_no_repl={ 'e0': xt.Bend(k0=0.3, h=0.31, length=1), 'e1': xt.Bend(k0=0.3, h=0.31, length=1), @@ -275,18 +277,26 @@ def test_slice_thick_and_insert_with_replicas(test_context): assert_allclose(p2.zeta, p1.zeta, rtol=0, atol=1e-14) assert_allclose(p2.delta, p1.delta, rtol=0, atol=1e-14) + assert line['e2']._movable + # line['e2']._mark = True + line.slice_thick_elements( slicing_strategies=[ xt.Strategy(None), xt.Strategy(xt.Teapot(3, mode='thick'), name='e2|e3|e4')]) + assert line['e2']._movable line.build_tracker(_context=test_context) + assert line['e2']._movable + line_no_repl.slice_thick_elements( slicing_strategies=[ xt.Strategy(None), xt.Strategy(xt.Teapot(3, mode='thick'), name='e2|e3|e4')]) line_no_repl.build_tracker(_context=test_context) + assert line['e2']._movable + tt = line.get_table() tt_no_repl = line_no_repl.get_table() @@ -341,13 +351,17 @@ def test_slice_thick_and_insert_with_replicas(test_context): assert_allclose(p2.zeta, p1.zeta, rtol=0, atol=1e-14) assert_allclose(p2.delta, p1.delta, rtol=0, atol=1e-14) + assert line['e2']._movable + line.discard_tracker() + assert line['e2']._movable line.insert_element(name='mkins1', element=xt.Marker(), at_s=0.5) line.insert_element(name='mkins2', element=xt.Marker(), at_s=1.5) line.insert_element(name='mkins3', element=xt.Marker(), at_s=2.5) line.insert_element(name='mkins4', element=xt.Marker(), at_s=3.5) line.insert_element(name='mkins5', element=xt.Marker(), at_s=4.5) line.build_tracker(_context=test_context) + assert line['e2']._movable line_no_repl.discard_tracker() line_no_repl.insert_element(name='mkins1', element=xt.Marker(), at_s=0.5) @@ -356,6 +370,7 @@ def test_slice_thick_and_insert_with_replicas(test_context): line_no_repl.insert_element(name='mkins4', element=xt.Marker(), at_s=3.5) line_no_repl.insert_element(name='mkins5', element=xt.Marker(), at_s=4.5) line_no_repl.build_tracker(_context=test_context) + assert line['e2']._movable tt = line.get_table() tt_no_repl = line_no_repl.get_table() @@ -428,6 +443,7 @@ def test_slice_thick_and_insert_with_replicas(test_context): line.track(p1) line_no_repl.track(p2) + assert line['e2']._movable assert_allclose(p2.x, p1.x, rtol=0, atol=1e-14) assert_allclose(p2.px, p1.px, rtol=0, atol=1e-14) diff --git a/xtrack/__init__.py b/xtrack/__init__.py index 4adde00ec..ec861dc05 100644 --- a/xtrack/__init__.py +++ b/xtrack/__init__.py @@ -13,6 +13,7 @@ from .random import * from .tracker_data import TrackerData from .line import Line, Node, freeze_longitudinal, _temp_knobs, EnergyProgram +from .environment import Environment, Place from .tracker import Tracker, Log from .match import (Vary, Target, TargetList, VaryList, TargetInequality, Action, TargetRelPhaseAdvance, TargetSet, GreaterThan, LessThan, diff --git a/xtrack/environment.py b/xtrack/environment.py new file mode 100644 index 000000000..17d077912 --- /dev/null +++ b/xtrack/environment.py @@ -0,0 +1,371 @@ +import xtrack as xt +import xobjects as xo +import numpy as np +from weakref import WeakSet + + +def _flatten_components(components): + flatten_components = [] + for nn in components: + if isinstance(nn, Place) and isinstance(nn.name, xt.Line): + line = nn.name + components = list(line.element_names).copy() + if nn.at is not None: + if isinstance(nn.at, str): + at = line._xdeps_eval.eval(nn.at) + else: + at = nn.at + at_first_element = at - line.get_length() / 2 + line[0].length / 2 + components[0] = Place(components[0], at=at_first_element, from_=nn.from_) + flatten_components += components + elif isinstance(nn, xt.Line): + flatten_components += nn.element_names + else: + flatten_components.append(nn) + return flatten_components + +class Environment: + def __init__(self, element_dict=None, particle_ref=None, _var_management=None): + self._element_dict = element_dict or {} + self.particle_ref = particle_ref + + if _var_management is not None: + self._var_management = _var_management + else: + self._init_var_management() + + self.lines = {} + self._lines = WeakSet() + self._drift_counter = 0 + self.ref = EnvRef(self) + + def new_line(self, components=None, name=None): + out = xt.Line() + out.particle_ref = self.particle_ref + out.env = self + out._element_dict = self.element_dict # Avoid copying + if components is None: + components = [] + flattened_components = _flatten_components(components) + out.element_names = handle_s_places(flattened_components, self) + out._var_management = self._var_management + out._name = name + self._lines.add(out) # Weak references + if name is not None: + self.lines[name] = out + + return out + + def _ensure_tracker_consistency(self, buffer): + for ln in self._lines: + if ln._has_valid_tracker() and ln._buffer is not buffer: + ln.discard_tracker() + + def _get_a_drift_name(self): + self._drift_counter += 1 + nn = f'drift_{self._drift_counter}' + if nn not in self.element_dict: + return nn + else: + return self._get_a_drift_name() + + def new(self, name, cls, at=None, from_=None, **kwargs): + + if from_ is not None or at is not None: + return Place(at=at, from_=from_, + name=self.new(name, cls, **kwargs)) + + _eval = self._xdeps_eval.eval + + assert isinstance(cls, str) or cls in [xt.Drift, xt.Bend, xt. + Quadrupole, xt.Sextupole, xt.Octupole, + xt.Multipole, xt.Marker, xt.Replica], ( + 'Only Drift, Dipole, Quadrupole, Sextupole, Octupole, Multipole, Marker, and Replica ' + 'elements are allowed in `new` for now.') + + cls_input = cls + if isinstance(cls, str): + # Clone an existing element + assert cls in self.element_dict, f'Element {cls} not found in environment' + self.element_dict[name] = xt.Replica(parent_name=cls) + self.replace_replica(name) + cls = type(self.element_dict[name]) + + ref_kwargs, value_kwargs = _parse_kwargs(cls, kwargs, _eval) + + if not isinstance(cls_input, str): # Parent is a class and not another element + self.element_dict[name] = cls(**value_kwargs) + + _set_kwargs(name=name, ref_kwargs=ref_kwargs, value_kwargs=value_kwargs, + element_dict=self.element_dict, element_refs=self.element_refs) + + return name + + def place(self, name, at=None, from_=None, anchor=None, from_anchor=None): + return Place(name, at=at, from_=from_, anchor=anchor, from_anchor=from_anchor) + +Environment.element_dict = xt.Line.element_dict +Environment._init_var_management = xt.Line._init_var_management +Environment._xdeps_vref = xt.Line._xdeps_vref +Environment._xdeps_fref = xt.Line._xdeps_fref +Environment._xdeps_manager = xt.Line._xdeps_manager +Environment._xdeps_eval = xt.Line._xdeps_eval +Environment.element_refs = xt.Line.element_refs +Environment.vars = xt.Line.vars +Environment.varval = xt.Line.varval +Environment.vv = xt.Line.vv +Environment.replace_replica = xt.Line.replace_replica +Environment.__getitem__ = xt.Line.__getitem__ +Environment.__setitem__ = xt.Line.__setitem__ +Environment.set = xt.Line.set + +class Place: + + def __init__(self, name, at=None, from_=None, anchor=None, from_anchor=None): + + if anchor is not None: + raise ValueError('anchor not implemented') + if from_anchor is not None: + raise ValueError('from_anchor not implemented') + + self.name = name + self.at = at + self.from_ = from_ + self.anchor = anchor + self.from_anchor = from_anchor + self._before = False + + def __repr__(self): + return f'Place({self.name}, at={self.at}, from_={self.from_})' + +def _all_places(seq): + seq_all_places = [] + for ss in seq: + if isinstance(ss, Place): + seq_all_places.append(ss) + elif not isinstance(ss, str) and hasattr(ss, '__iter__'): + # Find first place + i_first = None + for ii, sss in enumerate(ss): + if isinstance(sss, Place): + i_first = ii + break + if i_first is None: + raise ValueError('No Place in sequence') + ss_aux = _all_places(ss) + for ii in range(i_first): + ss_aux[ii]._before = True + seq_all_places.extend(ss_aux) + else: + seq_all_places.append(Place(ss, at=None, from_=None)) + return seq_all_places + +# In case we want to allow for the length to be an expression +# def _length_expr_or_val(name, line): +# if isinstance(line[name], xt.Replica): +# name = line[name].resolve(line, get_name=True) + +# if not line[name].isthick: +# return 0 + +# if line.element_refs[name]._expr is not None: +# return line.element_refs[name]._expr +# else: +# return line[name].length + + +def _resolve_s_positions(seq_all_places, env): + names_unsorted = [ss.name for ss in seq_all_places] + aux_line = env.new_line(components=names_unsorted) + aux_tt = aux_line.get_table() + aux_tt['length'] = np.diff(aux_tt._data['s'], append=0) + + s_center_dct = {} + n_resolved = 0 + n_resolved_prev = -1 + + if seq_all_places[0].at is None and not seq_all_places[0]._before: + # In case we want to allow for the length to be an expression + s_center_dct[seq_all_places[0].name] = aux_tt['length', seq_all_places[0].name] / 2 + # s_center_dct[seq_all_places[0].name] = _length_expr_or_val(seq_all_places[0].name, aux_line) / 2 + n_resolved += 1 + + while n_resolved != n_resolved_prev: + n_resolved_prev = n_resolved + for ii, ss in enumerate(seq_all_places): + if ss.name in s_center_dct: + continue + if ss.at is None and not ss._before: + ss_prev = seq_all_places[ii-1] + if ss_prev.name in s_center_dct: + # in case we want to allow for the length to be an expression + # s_center_dct[ss.name] = (s_center_dct[ss_prev.name] + # + _length_expr_or_val(ss_prev.name, aux_line) / 2 + # + _length_expr_or_val(ss.name, aux_line) / 2) + s_center_dct[ss.name] = (s_center_dct[ss_prev.name] + + aux_tt['length', ss_prev.name] / 2 + + aux_tt['length', ss.name] / 2) + n_resolved += 1 + elif ss.at is None and ss._before: + ss_next = seq_all_places[ii+1] + if ss_next.name in s_center_dct: + # in case we want to allow for the length to be an expression + # s_center_dct[ss.name] = (s_center_dct[ss_next.name] + # - _length_expr_or_val(ss_next.name, aux_line) / 2 + # - _length_expr_or_val(ss.name, aux_line) / 2) + s_center_dct[ss.name] = (s_center_dct[ss_next.name] + - aux_tt['length', ss_next.name] / 2 + - aux_tt['length', ss.name] / 2) + n_resolved += 1 + else: + if isinstance(ss.at, str): + at = aux_line._xdeps_eval.eval(ss.at) + else: + at = ss.at + + if ss.from_ is None: + s_center_dct[ss.name] = at + n_resolved += 1 + elif ss.from_ in s_center_dct: + s_center_dct[ss.name] = s_center_dct[ss.from_] + at + n_resolved += 1 + + assert n_resolved == len(seq_all_places), 'Not all positions resolved' + + aux_s_center_expr = np.array([s_center_dct[nn] for nn in aux_tt.name[:-1]]) + aux_s_center = [] + for ss in aux_s_center_expr: + if hasattr(ss, '_value'): + aux_s_center.append(ss._value) + else: + aux_s_center.append(ss) + aux_tt['s_center'] = np.concatenate([aux_s_center, [0]]) + + i_sorted = np.argsort(aux_s_center, stable=True) + + name_sorted = [str(aux_tt.name[ii]) for ii in i_sorted] + + tt_sorted = aux_tt.rows[name_sorted] + tt_sorted['s_entry'] = tt_sorted['s_center'] - tt_sorted['length'] / 2 + tt_sorted['s_exit'] = tt_sorted['s_center'] + tt_sorted['length'] / 2 + tt_sorted['ds_upstream'] = 0 * tt_sorted['s_entry'] + tt_sorted['ds_upstream'][1:] = tt_sorted['s_entry'][1:] - tt_sorted['s_exit'][:-1] + tt_sorted['ds_upstream'][0] = tt_sorted['s_entry'][0] + tt_sorted['s'] = tt_sorted['s_center'] + assert np.all(tt_sorted.name == np.array(name_sorted)) + + tt_sorted._data['s_center_dct'] = s_center_dct + + return tt_sorted + +def _generate_element_names_with_drifts(env, tt_sorted, s_tol=1e-12): + + names_with_drifts = [] + # Create drifts + for nn in tt_sorted.name: + ds_upstream = tt_sorted['ds_upstream', nn] + if np.abs(ds_upstream) > s_tol: + assert ds_upstream > 0, f'Negative drift length: {ds_upstream}, upstream of {nn}' + drift_name = env._get_a_drift_name() + env.new(drift_name, xt.Drift, length=ds_upstream) + names_with_drifts.append(drift_name) + names_with_drifts.append(nn) + + return list(map(str, names_with_drifts)) + +def handle_s_places(seq, env): + + if np.array([isinstance(ss, str) for ss in seq]).all(): + return [str(ss) for ss in seq] + + seq_all_places = _all_places(seq) + tab_sorted = _resolve_s_positions(seq_all_places, env) + names = _generate_element_names_with_drifts(env, tab_sorted) + + return names + +def _parse_kwargs(cls, kwargs, _eval): + ref_kwargs = {} + value_kwargs = {} + for kk in kwargs: + if hasattr(kwargs[kk], '_value'): + ref_kwargs[kk] = kwargs[kk] + value_kwargs[kk] = kwargs[kk]._value + elif (hasattr(cls, '_xofields') and kk in cls._xofields + and xo.array.is_array(cls._xofields[kk])): + assert hasattr(kwargs[kk], '__iter__'), ( + f'{kk} should be an iterable for {cls} element') + ref_vv = [] + value_vv = [] + for ii, vvv in enumerate(kwargs[kk]): + if hasattr(vvv, '_value'): + ref_vv.append(vvv) + value_vv.append(vvv._value) + elif isinstance(vvv, str): + ref_vv.append(_eval(vvv)) + value_vv.append(ref_vv[-1]._value) + else: + ref_vv.append(None) + value_vv.append(vvv) + ref_kwargs[kk] = ref_vv + value_kwargs[kk] = value_vv + elif (isinstance(kwargs[kk], str) and hasattr(cls, '_xofields') + and kk in cls._xofields and cls._xofields[kk].__name__ != 'String'): + ref_kwargs[kk] = _eval(kwargs[kk]) + if hasattr(ref_kwargs[kk], '_value'): + value_kwargs[kk] = ref_kwargs[kk]._value + else: + value_kwargs[kk] = ref_kwargs[kk] + else: + value_kwargs[kk] = kwargs[kk] + + return ref_kwargs, value_kwargs + +def _set_kwargs(name, ref_kwargs, value_kwargs, element_dict, element_refs): + for kk in value_kwargs: + if hasattr(value_kwargs[kk], '__iter__'): + len_value = len(value_kwargs[kk]) + getattr(element_dict[name], kk)[:len_value] = value_kwargs[kk] + if kk in ref_kwargs: + for ii, vvv in enumerate(value_kwargs[kk]): + if vvv is not None: + getattr(element_refs[name], kk)[ii] = vvv + else: + if kk in ref_kwargs: + setattr(element_refs[name], kk, ref_kwargs[kk]) + else: + setattr(element_dict[name], kk, value_kwargs[kk]) + +class EnvRef: + def __init__(self, env): + self.env = env + + def __getitem__(self, name): + if hasattr(self.env, 'lines') and name in self.env.lines: + return self.env.lines[name].ref + elif name in self.env.element_dict: + return self.env.element_refs[name] + elif name in self.env.vars: + return self.env.vars[name] + else: + raise KeyError(f'Name {name} not found.') + + def __setitem__(self, key, value): + if isinstance(value, xt.Line): + raise ValueError('Cannot set a Line, please use Envirnoment.new_line') + + if hasattr(value, '_value'): + val_ref = value + val_value = value._value + else: + val_ref = value + val_value = value + + if np.isscalar(val_value): + if key in self.env.element_dict: + raise ValueError(f'There is already an element with name {key}') + self.env.vars[key] = val_ref + else: + if key in self.env.vars: + raise ValueError(f'There is already a variable with name {key}') + self.element_refs[key] = val_ref \ No newline at end of file diff --git a/xtrack/json_utils.py b/xtrack/json_utils.py new file mode 100644 index 000000000..5ffc8abb1 --- /dev/null +++ b/xtrack/json_utils.py @@ -0,0 +1,39 @@ +import json +import io +from pathlib import Path +import gzip +from xobjects import JEncoder + + +def to_json(data, file, indent): + if isinstance(file, io.IOBase): + fh, close = file, False + elif (isinstance(file, str) and file.endswith(".gz")) or ( + isinstance(file, Path) and file.suffix == ".gz" + ): + fh, close = gzip.open(file, "wt"), True + else: + fh, close = open(file, "w"), True + + json.dump(data, fh, indent=indent, cls=JEncoder) + + if close: + fh.close() + + +def from_json(file): + if isinstance(file, io.IOBase): + fh, close = file, False + elif (isinstance(file, str) and file.endswith(".gz")) or ( + isinstance(file, Path) and file.suffix == ".gz" + ): + fh, close = gzip.open(file, "rt"), True + else: + fh, close = open(file, "r"), True + + data = json.load(fh) + + if close: + fh.close() + + return data diff --git a/xtrack/line.py b/xtrack/line.py index 1c8204600..a176525a5 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3,10 +3,8 @@ # Copyright (c) CERN, 2023. # # ######################################### # -import io import math import logging -import json import uuid import os from collections import defaultdict @@ -21,6 +19,7 @@ from scipy.constants import c as clight from . import linear_normal_form as lnf +from . import json_utils import xobjects as xo import xtrack as xt @@ -152,6 +151,7 @@ def __init__(self, elements=(), element_names=None, particle_ref=None, self._line_before_slicing_cache = None self._element_names_before_slicing = None + self.ref = xt.environment.EnvRef(self) @classmethod def from_dict(cls, dct, _context=None, _buffer=None, classes=()): @@ -235,6 +235,7 @@ def from_json(cls, file, **kwargs): ---------- file : str or file-like object Path to the json file or file-like object. + If filename ends with '.gz' file is decompressed. **kwargs : dict Additional keyword arguments passed to `Line.from_dict`. @@ -245,11 +246,7 @@ def from_json(cls, file, **kwargs): """ - if isinstance(file, io.IOBase): - dct = json.load(file) - else: - with open(file, 'r') as fid: - dct = json.load(fid) + dct = json_utils.from_json(file) if 'line' in dct.keys(): dct_line = dct['line'] @@ -626,7 +623,7 @@ def __setstate__(self, state): self.__dict__.update(state) - def to_json(self, file, **kwargs): + def to_json(self, file, indent=1, **kwargs): '''Save the line to a json file. Parameters @@ -639,11 +636,7 @@ def to_json(self, file, **kwargs): ''' - if isinstance(file, io.IOBase): - json.dump(self.to_dict(**kwargs), file, cls=xo.JEncoder) - else: - with open(file, 'w') as fid: - json.dump(self.to_dict(**kwargs), fid, cls=xo.JEncoder) + json_utils.to_json(self.to_dict(**kwargs), file, indent=indent) def _to_table_dict(self): @@ -745,7 +738,7 @@ def get_strengths(self, reverse=None): tab = xt.Table(out) if reverse: - xt.twiss._reverse_strengths(tab) # Change signs + xt.twiss._reverse_strengths(tab._data) # Change signs tab._data['reference_frame'] = { True: 'reverse', False: 'proper'}[reverse] @@ -841,12 +834,16 @@ def build_tracker( enable_pipeline_hold=enable_pipeline_hold, **kwargs) + if hasattr(self, 'env') and self.env is not None: + self.env._ensure_tracker_consistency(buffer=self._buffer) + return self.tracker @property def attr(self): - self._check_valid_tracker() + if not self._has_valid_tracker(): + self.build_tracker() if ('attr' not in self.tracker._tracker_data_base.cache.keys() or self.tracker._tracker_data_base.cache['attr'] is None): @@ -1323,6 +1320,9 @@ def match(self, vary, targets, solve=True, assert_within_tol=True, ''' + if not self._has_valid_tracker(): + self.build_tracker() + for old, new in zip(['ele_start', 'ele_stop', 'ele_init', 'twiss_init'], ['start', 'end', 'init_at', 'init']): if old in kwargs.keys(): @@ -1364,6 +1364,8 @@ def match_knob(self, knob_name, vary, targets, Value of the knob after the matching. Defaults to 1. ''' + if not self._has_valid_tracker(): + self.build_tracker() opt = match_knob_line(self, vary=vary, targets=targets, knob_name=knob_name, knob_value_start=knob_value_start, @@ -1401,6 +1403,9 @@ def survey(self,X0=0,Y0=0,Z0=0,theta0=0, phi0=0, psi0=0, Survey table. """ + if not self._has_valid_tracker(): + self.build_tracker() + if reverse is None: reverse = self.twiss_default.get('reverse', False) @@ -1560,7 +1565,8 @@ def find_closed_orbit(self, co_guess=None, particle_ref=None, num_turns=1, co_search_at=None, search_for_t_rev=False, - num_turns_search_t_rev=None): + num_turns_search_t_rev=None, + symmetrize=False): """ Find the closed orbit of the beamline. @@ -1635,7 +1641,8 @@ def find_closed_orbit(self, co_guess=None, particle_ref=None, start=start, end=end, num_turns=num_turns, co_search_at=co_search_at, search_for_t_rev=search_for_t_rev, - num_turns_search_t_rev=num_turns_search_t_rev) + num_turns_search_t_rev=num_turns_search_t_rev, + symmetrize=symmetrize) def compute_T_matrix(self, start=None, end=None, particle_on_co=None, steps_t_matrix=None): @@ -1841,7 +1848,8 @@ def compute_one_turn_matrix_finite_differences( steps_r_matrix=None, start=None, end=None, num_turns=1, - element_by_element=False, only_markers=False): + element_by_element=False, only_markers=False, + symmetrize=False): '''Compute the one turn matrix using finite differences. @@ -1881,7 +1889,8 @@ def compute_one_turn_matrix_finite_differences( steps_r_matrix, start=start, end=end, num_turns=num_turns, element_by_element=element_by_element, - only_markers=only_markers) + only_markers=only_markers, + symmetrize=symmetrize) def get_non_linear_chromaticity(self, delta0_range=(-1e-3, 1e-3), num_delta=5, fit_order=3, **kwargs): @@ -2046,6 +2055,10 @@ def _elements_intersecting_s( def cut_at_s(self, s: List[float], s_tol=1e-6): """Slice the line so that positions in s never fall inside an element.""" + + if self._has_valid_tracker(): + self.discard_tracker() + cuts_for_element = self._elements_intersecting_s(s, s_tol=s_tol) strategies = [Strategy(None)] # catch-all, ignore unaffected elements @@ -3305,7 +3318,7 @@ def get_line_with_second_order_maps(self, split_at): names_map_line.append(ele_cut_sorted[-1]) elements_map_line.append(self[ele_cut_sorted[-1]]) - line_maps = xt.Line(elements=elements_map_line, element_names=names_map_line) + line_maps = Line(elements=elements_map_line, element_names=names_map_line) line_maps.particle_ref = self.particle_ref.copy() return line_maps @@ -3338,6 +3351,110 @@ def transform_compound(self, *args, **kwargs): 'https://xsuite.readthedocs.io/en/latest/line.html#apply-transformations-tilt-shift-to-elements' ) + def mirror(self): + self._frozen_check() + self.element_names = list(reversed(self.element_names)) + + def replicate(self, name, mirror=False): + + self._env_if_needed() + + new_element_names = [] + for nn in self.element_names: + new_nn = nn + '.' + name + self.element_dict[new_nn] = xt.Replica(nn) + new_element_names.append(new_nn) + + out = self.env.new_line(components=new_element_names, name=name) + + if mirror: + out.mirror() + + return out + + def clone(self, name, mirror=False): + out = self.replicate(name=name, mirror=mirror) + out.replace_all_replicas() + return out + + def replace_replica(self, name): + name_parent = self[name].resolve(self, get_name=True) + cls = self.element_dict[name].__class__ + assert cls in [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, + xt.Multipole, xt.Marker, xt.Replica], ( + 'Only Drift, Dipole, Quadrupole, Sextupole, Octupole, Multipole, Marker, and Replica ' + 'elements are allowed in `new_element` for now.') + self.element_dict[name] = self[name_parent].copy() + + pars_with_expr = list( + self._xdeps_manager.tartasks[self.element_refs[name_parent]].keys()) + + for rr in pars_with_expr: + assert isinstance(rr, (xd.refs.AttrRef, xd.refs.ItemRef)), ( + 'Only AttrRef and ItemRef are supported for now') + if isinstance(rr, xd.refs.AttrRef): + setattr(self.element_refs[name], rr._key, rr._expr) + elif isinstance(rr, xd.refs.ItemRef): + getattr(self.element_refs[name], rr._owner._key)[rr._key] = rr._expr + + def replace_all_replicas(self): + for nn in self.element_names: + if isinstance(self[nn], xt.Replica): + self.replace_replica(nn) + + def select(self, start=None, end=None, name=None): + + if start is xt.START: + start = None + + if end is xt.END: + end = None + + tt = self.get_table().rows[start:end] + if tt.name[-1] == '_end_point': + tt = tt.rows[:-1] + + self._env_if_needed() + + out = self.env.new_line(components=list(tt.name), name=name) + + return out + + def set(self, name, *args, **kwargs): + _eval = self._xdeps_eval.eval + + if hasattr(self, 'lines') and name in self.lines: + raise ValueError('Cannot set a line') + + if name in self.element_dict: + if len(args) > 0: + raise ValueError(f'Only kwargs are allowed when setting element attributes') + ref_kwargs, value_kwargs = xt.environment._parse_kwargs( + type(self.element_dict[name]), kwargs, _eval) + xt.environment._set_kwargs( + name=name, ref_kwargs=ref_kwargs, value_kwargs=value_kwargs, + element_dict=self.element_dict, element_refs=self.element_refs) + else: + if len(kwargs) > 0: + raise ValueError(f'Only a single value is allowed when setting variable') + if len(args) != 1: + raise ValueError(f'A value must be provided when setting a variable') + value = args[0] + if isinstance(value, str): + self.vars[name] = _eval(value) + else: + self.vars[name] = value + + def _env_if_needed(self): + if not hasattr(self, 'env') or self.env is None: + self.env = xt.Environment(element_dict=self.element_dict, + particle_ref=self.particle_ref, + _var_management=self._var_management) + self.env._lines.add(self) + + def extend(self, line): + self.element_names.extend(line.element_names) + def __len__(self): return len(self.element_names) @@ -3379,7 +3496,7 @@ def name(self): if vv is self: return kk else: - return None + return getattr(self, '_name', None) @property def iscollective(self): @@ -3484,6 +3601,18 @@ def _xdeps_manager(self): if self._var_management is not None: return self._var_management['manager'] + @property + def _xdeps_eval(self): + try: + eva_obj = self._xdeps_eval_obj + except AttributeError: + eva_obj = xd.madxutils.MadxEval(variables=self._xdeps_vref, + functions=self._xdeps_fref, + elements=self.element_dict) + self._xdeps_eval_obj = eva_obj + + return eva_obj + @property def element_refs(self): if hasattr(self, '_in_multiline'): @@ -3671,20 +3800,41 @@ def steering_correctors_y(self): def steering_correctors_y(self, value): self._extra_config['steering_correctors_y'] = value - def __getitem__(self, ii): - if isinstance(ii, str): - - try: - return self.element_dict.__getitem__(ii) - except KeyError: - raise KeyError(f'No installed element with name {ii}') + def __getitem__(self, key): + if isinstance(key, str): + if key in self.element_dict: + return self.element_dict[key] + elif key in self.vars: + return self.vv[key] + elif hasattr(self, 'lines') and key in self.lines: # Want to reuse the method for the env + return self.lines[key] + else: + raise KeyError(f'Name {key} not found') else: - names = self.element_names.__getitem__(ii) + names = self.element_names.__getitem__(key) if isinstance(names, str): return self.element_dict.__getitem__(names) else: return [self.element_dict[nn] for nn in names] + def __setitem__(self, key, value): + + if isinstance(value, Line): + raise ValueError('Cannot set a Line, please use Envirnoment.new_line') + # Would need to make sure they refer to the same environment + + if hasattr(value, '_value'): + raise ValueError('Value cannot be a Ref. Please use Env.ref or Line.ref') + + if np.isscalar(value): + if key in self.element_dict: + raise ValueError(f'There is already an element with name {key}') + self.vars[key] = value + else: + if key in self.vars: + raise ValueError(f'There is already a variable with name {key}') + self.element_dict[key] = value + def _get_non_collective_line(self): if not self.iscollective: return self @@ -4534,7 +4684,7 @@ def set_from_madx_file(self, filename, mad_stdout=False): defined_vars = set(mad.globals.keys()) xt.general._print.suppress = True - dummy_line = xt.Line.from_madx_sequence(mad.sequence.dummy, + dummy_line = Line.from_madx_sequence(mad.sequence.dummy, deferred_expressions=True) xt.general._print.suppress = False @@ -4563,6 +4713,29 @@ def target(self, tar, value, **kwargs): action = ActionVars(self.line) return xt.Target(action=action, tar=tar, value=value, **kwargs) + def __call__(self, *args, **kwargs): + _eval = self.line._xdeps_eval.eval + if len(args) > 0: + assert len(kwargs) == 0 + assert len(args) == 1 + if isinstance(args[0], str): + return self[args[0]] + elif isinstance(args[0], dict): + kwargs.update(args[0]) + else: + raise ValueError('Invalid argument') + for kk in kwargs: + if isinstance(kwargs[kk], str): + self[kk] = _eval(kwargs[kk]) + else: + self[kk] = kwargs[kk] + + def set(self, name, value): + if isinstance(value, str): + self[name] = self.line._xdeps_eval.eval(value) + else: + self[name] = value + class ActionVars(Action): def __init__(self, line): @@ -4946,5 +5119,3 @@ def _rot_s_from_attr(attr): parent_cos_rot_s[has_parent_rot]) * attr._rot_and_shift_from_parent[has_parent_rot] return rot_s_rad - - diff --git a/xtrack/multiline/multiline.py b/xtrack/multiline/multiline.py index 9027249fc..b7c4df137 100644 --- a/xtrack/multiline/multiline.py +++ b/xtrack/multiline/multiline.py @@ -1,9 +1,8 @@ -import io -import json import pandas as pd import numpy as np from copy import deepcopy +from .. import json_utils from .shared_knobs import VarSharing from ..match import match_knob_line import xobjects as xo @@ -127,23 +126,20 @@ def from_dict(cls, dct): return new_multiline - def to_json(self, file, **kwargs): + def to_json(self, file, indent=1, **kwargs): '''Save the multiline to a json file. Parameters ---------- file: str or file-like object The file to save to. If a string is provided, a file is opened and - closed. If a file-like object is provided, it is used directly. + closed. If filename ends with '.gz' file is compressed. + If a file-like object is provided, it is used directly. **kwargs: dict Additional keyword arguments are passed to the `Line.to_dict` method. ''' + json_utils.to_json(self.to_dict(**kwargs), file, indent=indent) - if isinstance(file, io.IOBase): - json.dump(self.to_dict(**kwargs), file, cls=xo.JEncoder) - else: - with open(file, 'w') as fid: - json.dump(self.to_dict(**kwargs), fid, cls=xo.JEncoder) @classmethod def from_json(cls, file, **kwargs): @@ -153,7 +149,8 @@ def from_json(cls, file, **kwargs): ---------- file: str or file-like object The file to load from. If a string is provided, a file is opened and - closed. If a file-like object is provided, it is used directly. + closed. If the string endswith '.gz' the file is decompressed. + If a file-like object is provided, it is used directly. **kwargs: dict Returns @@ -161,14 +158,7 @@ def from_json(cls, file, **kwargs): new_multiline: Multiline The multiline object. ''' - - if isinstance(file, io.IOBase): - dct = json.load(file) - else: - with open(file, 'r') as fid: - dct = json.load(fid) - - return cls.from_dict(dct, **kwargs) + return cls.from_dict(json_utils.from_json(file), **kwargs) @classmethod def from_madx(cls, filename=None, madx=None, stdout=None, return_lines=False, **kwargs): @@ -622,4 +612,4 @@ def _dispatch_twiss_kwargs(kwargs, lines): f'Length of {arg_name} must be equal to the number of lines' kwargs_per_twiss[arg_name] = list(kwargs[arg_name]) kwargs.pop(arg_name) - return kwargs, kwargs_per_twiss \ No newline at end of file + return kwargs, kwargs_per_twiss diff --git a/xtrack/slicing.py b/xtrack/slicing.py index 8a4e60812..6ede147ff 100644 --- a/xtrack/slicing.py +++ b/xtrack/slicing.py @@ -364,6 +364,7 @@ def _make_slices(self, element, chosen_slicing, name): nn = f'{name}..entry_map' ee = element._entry_slice_class( _parent=element, _buffer=element._buffer) + element._movable = True # Force movable ee.parent_name = parent_name self._line.element_dict[nn] = ee slices_to_append.append(nn) @@ -382,6 +383,7 @@ def _make_slices(self, element, chosen_slicing, name): ee = type(element)( _parent=element._parent, _buffer=element._buffer, weight=weight * element.weight) + element._parent._movable = True # Force movable ee.parent_name = element.parent_name self._line.element_dict[nn] = ee slices_to_append.append(nn) @@ -393,6 +395,7 @@ def _make_slices(self, element, chosen_slicing, name): ee = element._drift_slice_class( _parent=element, _buffer=element._buffer, weight=weight) + element._movable = True # Force movable ee.parent_name = parent_name self._line.element_dict[nn] = ee slices_to_append.append(nn) @@ -403,6 +406,7 @@ def _make_slices(self, element, chosen_slicing, name): ee = element._thin_slice_class( _parent=element, _buffer=element._buffer, weight=weight) + element._movable = True # Force movable ee.parent_name = parent_name self._line.element_dict[nn] = ee slices_to_append.append(nn) @@ -413,6 +417,7 @@ def _make_slices(self, element, chosen_slicing, name): ee = element._thick_slice_class( _parent=element, _buffer=element._buffer, weight=weight) + element._movable = True # Force movable ee.parent_name = parent_name self._line.element_dict[nn] = ee slices_to_append.append(nn) @@ -424,6 +429,7 @@ def _make_slices(self, element, chosen_slicing, name): nn = f'{name}..exit_map' ee = element._exit_slice_class( _parent=element, _buffer=element._buffer) + element._movable = True # Force movable ee.parent_name = parent_name self._line.element_dict[nn] = ee slices_to_append.append(nn) diff --git a/xtrack/survey.py b/xtrack/survey.py index bd7de99d1..bf487d765 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -167,6 +167,19 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, return out + def plot(self, element_width=None, legend=True, **kwargs): + if element_width is None: + x_range = max(self.X) - min(self.X) + y_range = max(self.Y) - min(self.Y) + z_range = max(self.Z) - min(self.Z) + element_width = max([x_range, y_range, z_range]) * 0.03 + import xplt + xplt.FloorPlot(self, self.line, element_width=element_width, **kwargs) + if legend: + import matplotlib.pyplot as plt + plt.legend() + + # ================================================== # Main function @@ -224,6 +237,7 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, out = SurveyTable(data={**out_columns, **out_scalars}, # this is a merge col_names=out_columns.keys()) + out._data['line'] = line return out diff --git a/xtrack/synctime.py b/xtrack/synctime.py index d17dbe01a..78d473c8c 100644 --- a/xtrack/synctime.py +++ b/xtrack/synctime.py @@ -50,6 +50,13 @@ def track(self, particles): if mask_too_fast.any(): raise ValueError('Some particles move faster than the time window') + # For debugging (expected to be triggered) + # mask_out_of_circumference = mask_alive & ( + # (particles.zeta > self.circumference / 2) + # | (particles.zeta < -self.circumference / 2)) + # if mask_out_of_circumference.any(): + # raise ValueError('Some particles are out of the circumference') + # Update zeta for particles that are stopped particles.zeta[mask_stop] += beta0_beta1 * self.circumference diff --git a/xtrack/tracker.py b/xtrack/tracker.py index 37a0de0eb..982d92285 100644 --- a/xtrack/tracker.py +++ b/xtrack/tracker.py @@ -26,10 +26,6 @@ logger = logging.getLogger(__name__) - - - - class Tracker: ''' @@ -545,19 +541,19 @@ def _build_kernel( int64_t const ele_stop = ele_start + num_ele_track; - #ifndef XSUITE_BACKTRACK - if (flag_monitor==1){ - ParticlesMonitor_track_local_particle(tbt_monitor, &lpart); - } - int64_t elem_idx = ele_start; - int64_t const increm = 1; - #else + #if defined(XSUITE_BACKTRACK) || defined(XSUITE_MIRROR) int64_t elem_idx = ele_stop - 1; int64_t const increm = -1; if (flag_end_turn_actions>0){ increment_at_turn_backtrack(&lpart, flag_reset_s_at_end_turn, line_length, num_ele_line); } + #else + if (flag_monitor==1){ + ParticlesMonitor_track_local_particle(tbt_monitor, &lpart); + } + int64_t elem_idx = ele_start; + int64_t const increm = 1; #endif for (; ((elem_idx >= ele_start) && (elem_idx < ele_stop)); elem_idx+=increm){ @@ -608,11 +604,11 @@ def _build_kernel( break; } - #ifndef XSUITE_BACKTRACK - increment_at_element(&lpart, 1); - #else + #if defined(XSUITE_BACKTRACK) || defined(XSUITE_MIRROR) increment_at_element(&lpart, -1); - #endif //XSUITE_BACKTRACK + #else + increment_at_element(&lpart, 1); + #endif #endif //DANGER_SKIP_ACTIVE_CHECK_AND_SWAPS @@ -623,20 +619,18 @@ def _build_kernel( ParticlesMonitor_track_local_particle(tbt_monitor, &lpart); } - #ifndef XSUITE_BACKTRACK + #if defined(XSUITE_BACKTRACK) || defined(XSUITE_MIRROR) + if (flag_monitor==1){ + ParticlesMonitor_track_local_particle(tbt_monitor, &lpart); + } + # else if (flag_end_turn_actions>0){ if (isactive){ increment_at_turn(&lpart, flag_reset_s_at_end_turn); } } - #endif //XSUITE_BACKTRACK - + #endif - #ifdef XSUITE_BACKTRACK - if (flag_monitor==1){ - ParticlesMonitor_track_local_particle(tbt_monitor, &lpart); - } - #endif //XSUITE_BACKTRACK } // for turns LocalParticle_to_Particles(&lpart, particles, part_id, 1); @@ -877,7 +871,8 @@ def resume(self, session): """ Resume a track session that had been placed on hold. """ - return self._track_with_collective(particles=None, _session_to_resume=session) + return self._track_with_collective(particles=None, _session_to_resume=session, + _reset_log=False) def _track_with_collective( self, @@ -941,6 +936,7 @@ def _track_with_collective( '_context_needs_clean_active_lost_state'] tt_resume = _session_to_resume['tt'] ipp_resume = _session_to_resume['ipp'] + log = _session_to_resume['log'] _session_to_resume['resumed'] = True else: (ele_start, ele_stop, num_turns, flag_monitor, monitor, @@ -956,64 +952,17 @@ def _track_with_collective( if tt_resume is not None and tt < tt_resume: continue - if (flag_monitor and (ele_start == 0 or tt>0)): # second condition is for delayed start - if not(tt_resume is not None and tt == tt_resume): + if (flag_monitor and (ele_start == 0 or tt>0) # second condition is for delayed start + and not _is_resume_within_turn(tt, tt_resume)): monitor.track(particles) # Time dependent vars and energy ramping if self.line.enable_time_dependent_vars: - # Find first active particle - state = particles.state - if isinstance(particles._context, xo.ContextPyopencl): - state = state.get() - ii_first_active = int((state > 0).argmax()) - if ii_first_active == 0 and particles._xobject.state[0] <= 0: - # No active particles - at_turn = 0 # convenient for multi-turn injection - else: - at_turn = particles._xobject.at_turn[ii_first_active] + self._handle_time_dependent_vars(particles=particles) - if self.line.energy_program is not None: - t_turn = self.line.energy_program.get_t_s_at_turn(at_turn) - else: - beta0 = particles._xobject.beta0[ii_first_active] - t_turn = (at_turn * self._tracker_data_base.line_length - / (beta0 * clight)) - - # Clean leftover from previous trackings - if (self.line._t_last_update_time_dependent_vars and - t_turn < self.line._t_last_update_time_dependent_vars): - self.line._t_last_update_time_dependent_vars = None - - if (self.line._t_last_update_time_dependent_vars is None - or self.line.dt_update_time_dependent_vars is None - or t_turn > self.line._t_last_update_time_dependent_vars - + self.line.dt_update_time_dependent_vars): - self.line._t_last_update_time_dependent_vars = t_turn - self.vars['t_turn_s'] = t_turn - - if self.line.energy_program is not None: - p0c = self.line.particle_ref._xobject.p0c[0] - particles.update_p0c_and_energy_deviations(p0c) - - if log is not None: - for kk in log: - if log[kk] == None: - if kk not in self.line.log_last_track: - self.line.log_last_track[kk] = [] - self.line.log_last_track[kk].append(self.line.vv[kk]) - else: - ff = log[kk] - val = ff(self.line, particles) - if hasattr(ff, '_store'): - for nn in ff._store: - if nn not in self.line.log_last_track: - self.line.log_last_track[nn] = [] - self.line.log_last_track[nn].append(val[nn]) - else: - if kk not in self.line.log_last_track: - self.line.log_last_track[kk] = [] - self.line.log_last_track[kk].append(val) + if log is not None and not _is_resume_within_turn(tt, tt_resume): + self._handle_log(_session_to_resume=_session_to_resume, + particles=particles, log=log) moveback_to_buffer = None moveback_to_offset = None @@ -1073,7 +1022,8 @@ def _track_with_collective( 'moveback_to_offset': moveback_to_offset, 'ipp': ipp, 'tt': tt, - 'resumed': False + 'resumed': False, + 'log': log } return PipelineStatus(on_hold=True, data=session_on_hold) @@ -1482,6 +1432,65 @@ def check_compatibility_with_prebuilt_kernels(self): line_element_classes=self.line_element_classes, verbose=True) + def _handle_time_dependent_vars(self, particles): + + # Find first active particle + state = particles.state + if isinstance(particles._context, xo.ContextPyopencl): + state = state.get() + ii_first_active = int((state > 0).argmax()) + if ii_first_active == 0 and particles._xobject.state[0] <= 0: + # No active particles + at_turn = 0 # convenient for multi-turn injection + else: + at_turn = particles._xobject.at_turn[ii_first_active] + + if self.line.energy_program is not None: + t_turn = self.line.energy_program.get_t_s_at_turn(at_turn) + else: + beta0 = particles._xobject.beta0[ii_first_active] + t_turn = (at_turn * self._tracker_data_base.line_length + / (beta0 * clight)) + + # Clean leftover from previous trackings + if (self.line._t_last_update_time_dependent_vars and + t_turn < self.line._t_last_update_time_dependent_vars): + self.line._t_last_update_time_dependent_vars = None + + if (self.line._t_last_update_time_dependent_vars is None + or self.line.dt_update_time_dependent_vars is None + or t_turn > self.line._t_last_update_time_dependent_vars + + self.line.dt_update_time_dependent_vars): + self.line._t_last_update_time_dependent_vars = t_turn + self.vars['t_turn_s'] = t_turn + + if self.line.energy_program is not None: + p0c = self.line.particle_ref._xobject.p0c[0] + particles.update_p0c_and_energy_deviations(p0c) + + def _handle_log(self, _session_to_resume, particles, log): + if _session_to_resume is not None: + plog = _session_to_resume['particles'] + else: + plog = particles + for kk in log: + if log[kk] == None: + if kk not in self.line.log_last_track: + self.line.log_last_track[kk] = [] + self.line.log_last_track[kk].append(self.line.vv[kk]) + else: + ff = log[kk] + val = ff(self.line, plog) + if hasattr(ff, '_store'): + for nn in ff._store: + if nn not in self.line.log_last_track: + self.line.log_last_track[nn] = [] + self.line.log_last_track[nn].append(val[nn]) + else: + if kk not in self.line.log_last_track: + self.line.log_last_track[kk] = [] + self.line.log_last_track[kk].append(val) + class TrackerConfig(UserDict): @@ -1587,4 +1596,7 @@ def __init__(self, *args, **kwargs): self[arg] = None else: self[f'_unnamed_{unnamed_indx}'] = arg - unnamed_indx += 1 \ No newline at end of file + unnamed_indx += 1 + +def _is_resume_within_turn(tt, tt_resume): + return tt_resume is not None and tt == tt_resume \ No newline at end of file diff --git a/xtrack/tracker_data.py b/xtrack/tracker_data.py index f4fa3af37..919467a7d 100644 --- a/xtrack/tracker_data.py +++ b/xtrack/tracker_data.py @@ -131,6 +131,7 @@ def __init__( if this_parent._buffer is not self._element_dict[nn]._buffer: this_parent.move(_buffer=self._element_dict[nn]._buffer) self._element_dict[nn]._parent = this_parent + this_parent._movable = True assert self._element_dict[nn]._parent._offset == self._element_dict[nn]._xobject._parent._offset def common_buffer_for_elements(self): @@ -176,7 +177,9 @@ def check_elements_in_common_buffer(self, buffer, allow_move=False): Move all the elements to the common buffer, if they are not already there. """ - for ee in self._elements: + for nn, ee in self._element_dict.items(): + if not hasattr(ee, '_buffer'): + continue if ee._buffer is not buffer: if allow_move: ee.move(_buffer=buffer) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 7c184da22..6c1a43765 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -58,6 +58,8 @@ SKEW_STRENGTHS_FROM_ATTR=['k0sl', 'k1sl', 'k2sl', 'k3sl', 'k4sl', 'k5sl'] OTHER_FIELDS_FROM_ATTR=['angle_rad', 'rot_s_rad', 'hkick', 'vkick', 'ks', 'length'] OTHER_FIELDS_FROM_TABLE=['element_type', 'isthick', 'parent_name'] +SIGN_FLIP_FOR_ATTR_REVERSE=['k0l', 'k2l', 'k4l', 'k1sl', 'k3sl', 'k5sl', 'vkick', 'angle_rad'] + log = logging.getLogger(__name__) @@ -329,6 +331,11 @@ def twiss_line(line, particle_ref=None, method=None, kwargs = locals().copy() + if (init is not None or betx is not None or bety is not None) and start is None: + # is open twiss + start = xt.START + end = end or xt.END + if num_turns != 1: # Untested cases assert num_turns > 0 @@ -1265,18 +1272,16 @@ def _compute_chromatic_functions(line, init, delta_chrom, steps_r_matrix, particle_on_co=on_momentum_twiss_res.particle_on_co.copy(), nemitt_x=nemitt_x, nemitt_y=nemitt_y, W_matrix=tw_init_chrom.W_matrix) - if periodic_mode == 'periodic_symmetric': - part_chrom = part_guess.copy() # Finding closed orbit does not make sense in this case - else: - part_chrom = line.find_closed_orbit(delta0=dd, co_guess=part_guess, - start=start, end=end, num_turns=num_turns) + part_chrom = line.find_closed_orbit(delta0=dd, co_guess=part_guess, + start=start, end=end, num_turns=num_turns, + symmetrize=(periodic_mode == 'periodic_symmetric')) tw_init_chrom.particle_on_co = part_chrom RR_chrom = line.compute_one_turn_matrix_finite_differences( particle_on_co=tw_init_chrom.particle_on_co.copy(), start=start, end=end, num_turns=num_turns, - steps_r_matrix=steps_r_matrix)['R_matrix'] - if periodic_mode == 'periodic_symmetric': - RR_chrom = _compute_R_periodic_symmetric(RR_chrom) + steps_r_matrix=steps_r_matrix, + symmetrize=(periodic_mode == 'periodic_symmetric') + )['R_matrix'] (WW_chrom, _, _, _) = lnf.compute_linear_normal_form(RR_chrom, only_4d_block=method=='4d', @@ -1758,7 +1763,6 @@ def _find_periodic_solution(line, particle_on_co, particle_ref, method, assert periodic_mode in ['periodic', 'periodic_symmetric'] - if start is not None or end is not None: assert start is not None and end is not None, ( 'start and end must be both None or both not None') @@ -1772,7 +1776,6 @@ def _find_periodic_solution(line, particle_on_co, particle_ref, method, if periodic_mode == 'periodic_symmetric': assert R_matrix is None, 'R_matrix must be None for `periodic_symmetric`' assert W_matrix is None, 'W_matrix must be None for `periodic_symmetric`' - assert delta0 == 0, 'delta0 must be 0 for `periodic_symmetric`' if particle_on_co is not None: part_on_co = particle_on_co @@ -1792,18 +1795,11 @@ def _find_periodic_solution(line, particle_on_co, particle_ref, method, co_search_at=co_search_at, search_for_t_rev=search_for_t_rev, num_turns_search_t_rev=num_turns_search_t_rev, + symmetrize=(periodic_mode == 'periodic_symmetric'), ) if only_orbit: W_matrix = np.eye(6) - if periodic_mode == 'periodic_symmetric': - assert np.allclose(part_on_co.x[0], 0, atol=1e-12, rtol=0) - assert np.allclose(part_on_co.px[0], 0, atol=1e-12, rtol=0) - assert np.allclose(part_on_co.y[0], 0, atol=1e-12, rtol=0) - assert np.allclose(part_on_co.py[0], 0, atol=1e-12, rtol=0) - assert np.allclose(part_on_co.delta[0], 0, atol=1e-12, rtol=0) - assert np.allclose(part_on_co.zeta[0], 0, atol=1e-12, rtol=0) - if W_matrix is not None: W = W_matrix @@ -1829,13 +1825,11 @@ def _find_periodic_solution(line, particle_on_co, particle_ref, method, num_turns=num_turns, element_by_element=compute_R_element_by_element, only_markers=only_markers, + symmetrize=(periodic_mode == 'periodic_symmetric'), ) RR = RR_out['R_matrix'] RR_ebe = RR_out['R_matrix_ebe'] - if periodic_mode == 'periodic_symmetric': - RR = _compute_R_periodic_symmetric(RR) - if matrix_responsiveness_tol is not None: lnf._assert_matrix_responsiveness(RR, matrix_responsiveness_tol, only_4d=(method == '4d')) @@ -2060,7 +2054,8 @@ def find_closed_orbit_line(line, co_guess=None, particle_ref=None, co_search_at=None, search_for_t_rev=False, continue_on_closed_orbit_error=False, - num_turns_search_t_rev=None): + num_turns_search_t_rev=None, + symmetrize=False): if search_for_t_rev: assert line.particle_ref is not None @@ -2074,6 +2069,7 @@ def find_closed_orbit_line(line, co_guess=None, particle_ref=None, assert num_turns == 1, '`num_turns` not supported when `search_for_t_rev` is True' assert co_search_at is None, '`co_search_at` not supported when `search_for_t_rev` is True' assert continue_on_closed_orbit_error is False, '`continue_on_closed_orbit_error` not supported when `search_for_t_rev` is True' + assert symmetrize is False, '`symmetrize` not supported when `search_for_t_rev` is True' out = _find_closed_orbit_search_t_rev(line, num_turns_search_t_rev) return out @@ -2086,6 +2082,7 @@ def find_closed_orbit_line(line, co_guess=None, particle_ref=None, co_search_at = None # needs to be implemented if co_search_at is not None: + assert not symmetrize, 'Symmetrize not supported when `co_search_at` is provided' kwargs = locals().copy() kwargs.pop('start') kwargs.pop('end') @@ -2161,7 +2158,8 @@ def find_closed_orbit_line(line, co_guess=None, particle_ref=None, if np.all(np.abs(_error_for_co( x0, co_guess, line, delta_zeta, delta0, zeta0, start=start, end=end, - num_turns=num_turns)) < DEFAULT_CO_SEARCH_TOL): + num_turns=num_turns, + symmetrize=symmetrize)) < DEFAULT_CO_SEARCH_TOL): res = x0 fsolve_info = 'taken_guess' ier = 1 @@ -2170,7 +2168,8 @@ def find_closed_orbit_line(line, co_guess=None, particle_ref=None, opt = xt.match.opt_from_callable( lambda p: _error_for_co(p, co_guess, line, delta_zeta, delta0, zeta0, start=start, - end=end, num_turns=num_turns), + end=end, num_turns=num_turns, + symmetrize=symmetrize), x0=x0, steps=[1e-8, 1e-9, 1e-8, 1e-9, 1e-7, 1e-8], tar=[0., 0., 0., 0., 0., 0.], tols=[1e-12, 1e-12, 1e-12, 1e-12, 1e-12, 1e-12]) @@ -2200,7 +2199,7 @@ def find_closed_orbit_line(line, co_guess=None, particle_ref=None, return particle_on_co -def _one_turn_map(p, particle_ref, line, delta_zeta, start, end, num_turns): +def _one_turn_map(p, particle_ref, line, delta_zeta, start, end, num_turns, symmetrize): part = particle_ref.copy() part.x = p[0] part.px = p[1] @@ -2216,6 +2215,11 @@ def _one_turn_map(p, particle_ref, line, delta_zeta, start, end, num_turns): part.update_p0c_and_energy_deviations(p0c = part._xobject.p0c[0] + dp0c) line.track(part, ele_start=start, ele_stop=end, num_turns=num_turns) + if symmetrize: + assert num_turns == 1 + with xt.line._preserve_config(line): + line.config.XSUITE_MIRROR = True + line.track(part, ele_start=start, ele_stop=end, num_turns=1) if part.state[0] < 0: raise ClosedOrbitSearchError( f'Particle lost, p.state = {part.state[0]}') @@ -2228,11 +2232,11 @@ def _one_turn_map(p, particle_ref, line, delta_zeta, start, end, num_turns): part._xobject.delta[0]]) return p_res -def _error_for_co_search_6d(p, co_guess, line, delta_zeta, delta0, zeta0, start, end, num_turns): - return p - _one_turn_map(p, co_guess, line, delta_zeta, start, end, num_turns) +def _error_for_co_search_6d(p, co_guess, line, delta_zeta, delta0, zeta0, start, end, num_turns, symmetrize): + return p - _one_turn_map(p, co_guess, line, delta_zeta, start, end, num_turns, symmetrize) -def _error_for_co_search_4d_delta0(p, co_guess, line, delta_zeta, delta0, zeta0, start, end, num_turns): - one_turn_res = _one_turn_map(p, co_guess, line, delta_zeta, start, end, num_turns) +def _error_for_co_search_4d_delta0(p, co_guess, line, delta_zeta, delta0, zeta0, start, end, num_turns, symmetrize): + one_turn_res = _one_turn_map(p, co_guess, line, delta_zeta, start, end, num_turns, symmetrize) return np.array([ p[0] - one_turn_res[0], p[1] - one_turn_res[1], @@ -2241,8 +2245,8 @@ def _error_for_co_search_4d_delta0(p, co_guess, line, delta_zeta, delta0, zeta0, 0, p[5] - delta0]) -def _error_for_co_search_4d_zeta0(p, co_guess, line, delta_zeta, delta0, zeta0, start, end, num_turns): - one_turn_res = _one_turn_map(p, co_guess, line, delta_zeta, start, end, num_turns) +def _error_for_co_search_4d_zeta0(p, co_guess, line, delta_zeta, delta0, zeta0, start, end, num_turns, symmetrize): + one_turn_res = _one_turn_map(p, co_guess, line, delta_zeta, start, end, num_turns, symmetrize) return np.array([ p[0] - one_turn_res[0], p[1] - one_turn_res[1], @@ -2251,8 +2255,8 @@ def _error_for_co_search_4d_zeta0(p, co_guess, line, delta_zeta, delta0, zeta0, p[4] - zeta0, 0]) -def _error_for_co_search_4d_delta0_zeta0(p, co_guess, line, delta_zeta, delta0, zeta0, start, end, num_turns): - one_turn_res = _one_turn_map(p, co_guess, line, delta_zeta, start, end, num_turns) +def _error_for_co_search_4d_delta0_zeta0(p, co_guess, line, delta_zeta, delta0, zeta0, start, end, num_turns, symmetrize): + one_turn_res = _one_turn_map(p, co_guess, line, delta_zeta, start, end, num_turns, symmetrize) return np.array([ p[0] - one_turn_res[0], p[1] - one_turn_res[1], @@ -2267,7 +2271,8 @@ def compute_one_turn_matrix_finite_differences( start=None, end=None, num_turns=1, element_by_element=False, - only_markers=False): + only_markers=False, + symmetrize=True): import xpart if steps_r_matrix is None: @@ -2322,9 +2327,14 @@ def compute_one_turn_matrix_finite_differences( assert num_turns == 1, 'Not yet implemented' assert end is not None line.track(part_temp, ele_start=start, ele_stop=end) + if symmetrize: + with xt.line._preserve_config(line): + line.config.XSUITE_MIRROR = True + line.track(part_temp, ele_start=start, ele_stop=end) elif particle_on_co._xobject.at_element[0]>0: assert element_by_element is False, 'Not yet implemented' assert num_turns == 1, 'Not yet implemented' + assert symmetrize is False, 'Not yet implemented' i_start = particle_on_co._xobject.at_element[0] line.track(part_temp, ele_start=i_start) line.track(part_temp, num_elements=i_start) @@ -2335,6 +2345,10 @@ def compute_one_turn_matrix_finite_differences( monitor_setting = 'ONE_TURN_EBE' if element_by_element else None line.track(part_temp, num_turns=num_turns, turn_by_turn_monitor=monitor_setting) + if symmetrize: + with xt.line._preserve_config(line): + line.config.XSUITE_MIRROR = True + line.track(part_temp, num_turns=num_turns) temp_mat = np.zeros(shape=(6, 12), dtype=np.float64) temp_mat[0, :] = context.nparray_from_context_array(part_temp.x) @@ -2762,15 +2776,15 @@ def __setattr__(self, name, value): else: self.__dict__[name] = value - def get_normalized_coordinates(self, particles, nemitt_x=None, nemitt_y=None, + def get_normalized_coordinates(self, particles, nemitt_x=None, nemitt_y=None, nemitt_zeta=None): - + ctx2np = particles._context.nparray_from_context_array - + part_id = ctx2np(particles.particle_id).copy() at_element = ctx2np(particles.at_element).copy() at_turn = ctx2np(particles.at_element).copy() - x_norm = ctx2np(particles.x).copy() + x_norm = ctx2np(particles.x).copy() px_norm = x_norm.copy() y_norm = x_norm.copy() py_norm = x_norm.copy() @@ -2809,7 +2823,6 @@ def get_normalized_coordinates(self, particles, nemitt_x=None, nemitt_y=None, 'x_norm': x_norm, 'px_norm': px_norm, 'y_norm': y_norm, 'py_norm': py_norm, 'zeta_norm': zeta_norm, 'pzeta_norm': pzeta_norm}, index='particle_id') - @property def betx(self): @@ -2983,7 +2996,7 @@ def get_beam_covariance(self, Sigma = gemitt_x * Sigma1 + gemitt_y * Sigma2 + gemitt_zeta * Sigma3 res = _build_sigma_table(Sigma=Sigma, s=self.s, name=self.name) - return Table(res) + return res def get_ibs_growth_rates( self, @@ -3263,7 +3276,7 @@ def reverse(self): out.qs = 0 out.muzeta[:] = 0 - _reverse_strengths(out) + _reverse_strengths(out._data) out._data['reference_frame'] = { 'proper': 'reverse', 'reverse': 'proper'}[self.reference_frame] @@ -3273,7 +3286,7 @@ def reverse(self): ind_per_table = [] def add_strengths(self, line=None): - if line is None: + if line is None and hasattr(self,"_action"): line = self._action.line _add_strengths_to_twiss_res(self, line) return self @@ -3356,13 +3369,21 @@ def target(self, tars=None, value=None, at=None, **kwargs): action=self._action, **kwargs) return tarset - def plot(self,yl="",yr="",x='s', + def plot(self, + yl=None, + yr=None,x='s', lattice=True, mask=None, labels=None, clist="k r b g c m", + figure=None, + figlabel=None, ax=None, - figlabel=None): + axleft=None, + axright=None, + axlattice=None, + hover=False + ): """ Plot columns of the TwissTable @@ -3388,9 +3409,16 @@ def plot(self,yl="",yr="",x='s', label to use for the figure """ - if yl=="" and yr=="": + if yl is None and yr is None: yl='betx bety' yr='dx dy' + if yl is None: + yl="" + if yr is None: + yr="" + + if not hasattr(self,"_action"): + lattice=False if lattice and 'length' not in self.keys(): self.add_strengths() @@ -3402,15 +3430,23 @@ def plot(self,yl="",yr="",x='s', idx=mask else: idx=slice(None) - if ax is None: - newfig=True - else: - raise NotImplementedError self._is_s_begin=True - pl=TwissPlot(self, x=x, yl=yl, yr=yr, idx=idx, lattice=lattice, newfig=newfig, - figlabel=figlabel,clist=clist) + pl=TwissPlot(self, + x=x, + yl=yl, + yr=yr, + idx=idx, + lattice=lattice, + figure=figure, + figlabel=figlabel,clist=clist, + ax=ax, + axleft=axleft, + axright=axright, + axlattice=axlattice, + hover=hover, + ) if labels is not None: mask=self.mask[labels] @@ -3807,25 +3843,11 @@ def _find_closed_orbit_search_t_rev(line, num_turns_search_t_rev=None): def _reverse_strengths(out): - # Same convention as in MAD-X for reversing strengths - for kk in NORMAL_STRENGTHS_FROM_ATTR: - if kk not in out._col_names: - continue - ii = int(kk.split('k')[-1].split('l')[0]) - out[kk] *= (-1)**(ii+1) - - for kk in SKEW_STRENGTHS_FROM_ATTR: - if kk not in out._col_names: - continue - ii = int(kk.split('k')[-1].split('sl')[0]) - out[kk] *= (-1)**ii - - if 'vkick' in out._col_names: - out['vkick'] *= -1 - - if 'angle_rad' in out._col_names: - out['angle_rad'] *= -1 - + ### Same convention as in MAD-X for reversing strengths + for kk in SIGN_FLIP_FOR_ATTR_REVERSE: + if kk in out: + val=out[kk]#avoid passing by setitem + np.negative(val,val) def _W_phys2norm(x, px, y, py, zeta, pzeta, W_matrix, co_dict, nemitt_x=None, nemitt_y=None, nemitt_zeta=None): @@ -3873,8 +3895,3 @@ def _add_strengths_to_twiss_res(twiss_res, line): + OTHER_FIELDS_FROM_ATTR + OTHER_FIELDS_FROM_TABLE): twiss_res._col_names.append(kk) twiss_res._data[kk] = tt[kk].copy() - -def _compute_R_periodic_symmetric(RR): - inv_momenta = np.diag([1., -1., 1., -1., 1., 1.]) - RR_symm = inv_momenta @ np.linalg.inv(RR) @ inv_momenta @ RR - return RR_symm diff --git a/xtrack/twissplot.py b/xtrack/twissplot.py index 389b11aed..e21763e95 100644 --- a/xtrack/twissplot.py +++ b/xtrack/twissplot.py @@ -10,7 +10,7 @@ def _mylbl(d, x): - return d.get(x, r"$%s$" % x) + return d.get(x, r"%s" % x) class TwissPlot(object): @@ -21,12 +21,12 @@ class TwissPlot(object): "dy": r"$D_y$", "mux": r"$\mu_x$", "muy": r"$\mu_y$", - "Ax": "$A_x$", - "Ay": "$A_y$", - "Bx": "$B_x$", - "By": "$B_y$", - "wx": "$w_x$", - "wy": "$w_y$", + "ax_chrom": "$A_x$", + "ay_chrom": "$A_y$", + "bx_chrom": "$B_x$", + "by_chrom": "$B_y$", + "wx_chrom": "$W_x$", + "wy_chrom": "$W_y$", "sigx": r"$\sigma_x=\sqrt{\beta_x \epsilon}$", "sigy": r"$\sigma_y=\sqrt{\beta_y \epsilon}$", "sigdx": r"$\sigma_{D_x}=D_x \delta$", @@ -47,6 +47,12 @@ class TwissPlot(object): "sigx": r"$\sigma$ [mm]", "sigy": r"$\sigma$ [mm]", "sigdx": r"$\sigma$ [mm]", + "ax_chrom": "$A$", + "ay_chrom": "$A$", + "bx_chrom": "$B$", + "by_chrom": "$B$", + "wx_chrom": "$W$", + "wy_chrom": "$W$", "n1": r"Aperture [$\sigma$]", } autoupdate = [] @@ -54,8 +60,7 @@ class TwissPlot(object): def ani_autoupdate(self): from matplotlib.animation import FuncAnimation - self._ani = FuncAnimation( - self.figure, self.update, blit=False, interval=1000) + self._ani = FuncAnimation(self.figure, self.update, blit=False, interval=1000) def ani_stopupdate(self): del self._ani @@ -73,8 +78,13 @@ def __init__( idx=slice(None), clist="k r b g c m", lattice=None, - newfig=True, - figlabel=None + figure=None, + figlabel=None, + ax=None, + axleft=None, + axright=None, + axlattice=None, + hover=False, ): import matplotlib.pyplot as plt @@ -94,34 +104,34 @@ def __init__( idx, clist, ) + self.ax = ax + self.used_ax = False + if ax is not None: + self.figure = ax.figure + elif figure is None: + self.figure = plt.figure(num=figlabel, figsize=(6.4*1.2, 4.8)) + if figlabel is not None: + self.figure.clf() for i in self.yl + self.yr: self.color[i] = self.clist.pop(0) self.clist.append(self.color[i]) - if newfig is True: - self.figure = plt.figure(num=figlabel) - self.figure.clf() - elif newfig is False: - self.figure = plt.gcf() - self.figure.clf() - else: - self.figure = newfig - self.figure.clf() - if lattice and x=="s": - self.lattice = self._new_axes() + if lattice and x == "s": + self.lattice = self._new_axis(axlattice) + # self.lattice.set_frame_on(False) # self.lattice.set_autoscale_on(False) self.lattice.yaxis.set_visible(False) if yl: - self.left = self._new_axes() + self.left = self._new_axis(axleft) # self.left.set_autoscale_on(False) if yr: - self.right = self._new_axes() + self.right = self._new_axis(axright) # self.right.set_autoscale_on(False) self.left.yaxis.set_label_position("right") self.left.yaxis.set_ticks_position("right") # timeit('Setup') self.run() - if lattice and x=="s": + if lattice and x == "s": self.lattice.set_autoscale_on(False) if yl: self.left.set_autoscale_on(False) @@ -131,16 +141,20 @@ def __init__( self.right.set_autoscale_on(False) self.right.yaxis.set_label_position("right") self.right.yaxis.set_ticks_position("right") + if hover: + self.set_hover() # timeit('Update') - def _new_axes(self): - if self.figure.axes: - ax = self.figure.axes[-1] - out = self.figure.add_axes( - ax.get_position(), sharex=ax, frameon=False) + def _new_axis(self, ax=None): + if self.ax is None: + out = self.figure.add_subplot(111) + self.figure.subplots_adjust(right=0.75) + self.ax = out + if self.used_ax: + out = self.ax.twinx() else: - # adjust plot dimensions - out = self.figure.add_axes([0.17, 0.12, 0.6, 0.8]) + out = self.ax + self.used_ax = True return out def __repr__(self): @@ -182,6 +196,7 @@ def run(self): self.xaxis = self.ont[self.x][self.idx] self.lines = [] self.legends = [] + self.names = [] # self.figure.lines=[] # self.figure.patches=[] # self.figure.texts=[] @@ -204,24 +219,37 @@ def run(self): self.right.clear() for i in self.yr: self._column(i, self.right, self.color[i]) - ca = self.figure.gca() - ca.set_xlabel(_mylbl(self.axlabel, self.x)) - ca.set_xlim(min(self.xaxis), max(self.xaxis)) - self.figure.legend(self.lines, self.legends, loc="upper right") - ca.grid(True) - # self.figure.canvas.mpl_connect('button_release_event',self.button_press) - self.figure.canvas.mpl_connect("pick_event", self.pick) - # plt.interactive(is_ion) + self.ax.set_xlabel(_mylbl(self.axlabel, self.x)) + self.ax.set_xlim(min(self.xaxis), max(self.xaxis)) + self.ax.legend( + self.lines, self.legends, loc="upper right", bbox_to_anchor=(1.35, 1.) + ) + self.ax.grid(True) self.figure.canvas.draw() if hasattr(self, "on_run"): self.on_run(self) + def set_hover(self): + self.figure.canvas.mpl_connect("motion_notify_event", self.pick) + def pick(self, event): - pos = np.array([event.mouseevent.x, event.mouseevent.y]) - name = event.artist.elemname - prop = event.artist.elemprop - value = event.artist.elemvalue - print("\n %s.%s=%s" % (name, prop, value), end=" ") + for ii, ll in enumerate(self.lines): + _, data = ll.contains(event) + lgd = self.names[ii] + if "ind" in data: + xx = ll.get_xdata() + yy = ll.get_ydata() + for idx in data["ind"]: + if "name" in self.table._col_names: + name = self.table.name[idx] + else: + name = "" + print(f"{name:25}, s={xx[idx]:15.6g}, {lgd:>10}={yy[idx]:15.6g}") + # pos = np.array([event.mouseevent.x, event.mouseevent.y]) + # name = event.artist.elemname + # prop = event.artist.elemprop + # value = event.artist.elemvalue + # print("\n %s.%s=%s" % (name, prop, value), end=" ") # def button_press(self,mouseevent): # rel=np.array([mouseevent.x,mouseevent.y]) @@ -239,7 +267,7 @@ def _lattice(self, names, color, lbl): vd = 0 sp = self.lattice s = self.ont.s - l = np.diff(s,append=[s[-1]]) + l = np.diff(s, append=[s[-1]]) for name in names: myvd = self.ont._data.get(name, None) if myvd is not None: @@ -266,6 +294,7 @@ def _lattice(self, names, color, lbl): if bplt: self.lines.append(bplt[0]) self.legends.append(lbl) + self.names.append(lbl) row_names = self.ont.name for r, name in zip(bplt, c): r.elemname = row_names[name] @@ -282,23 +311,49 @@ def _column(self, name, sp, color): sp.set_ylabel(_mylbl(self.axlabel, name)) self.lines.append(bxp) self.legends.append(_mylbl(self.lglabel, name)) + self.names.append(name) sp.autoscale_view() def savefig(self, name): self.figure.savefig(name) return self - def ylim(self,left_lo=None,left_hi=None,right_lo=None,right_hi=None): - lo,hi=self.left.get_ylim() + def ylim( + self, + left_lo=None, + left_hi=None, + right_lo=None, + right_hi=None, + lattice_lo=None, + lattice_hi=None, + ): + lo, hi = self.left.get_ylim() if left_lo is None: left_lo = lo if left_hi is None: left_hi = hi - self.left.set_ylim(left_lo,left_hi) - lo,hi=self.right.get_ylim() + self.left.set_ylim(left_lo, left_hi) + lo, hi = self.right.get_ylim() if right_lo is None: right_lo = lo if right_hi is None: right_hi = hi - self.right.set_ylim(right_lo,right_hi) + self.right.set_ylim(right_lo, right_hi) + lo, hi = self.lattice.get_ylim() + if lattice_lo is None: + lattice_lo = lo + if lattice_hi is None: + lattice_hi = hi + self.lattice.set_ylim(lattice_lo, lattice_hi) + return self + + def set_s_label(self, regexp="ip.*"): + sel = self.table.rows[regexp] + self.ax.set_xticks(sel.s, sel.name) + self.ax.set_xlabel(None) + return self + + def move_legend(self, left=0, bottom=0, width=0, height=0): + """Uses ax.legend_.set_bbox_to_anchor""" + self.ax.legend_.set_bbox_to_anchor((left, bottom, width, height)) return self