From f9b34363f66ba7c9615dedcdb788f80b0caaf0fb Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 21 Jul 2024 14:25:16 +0200 Subject: [PATCH 001/159] Trying to get the log working with multitracker --- xtrack/tracker.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/xtrack/tracker.py b/xtrack/tracker.py index 37a0de0eb..7279025ef 100644 --- a/xtrack/tracker.py +++ b/xtrack/tracker.py @@ -877,7 +877,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 +942,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, @@ -997,6 +999,10 @@ def _track_with_collective( particles.update_p0c_and_energy_deviations(p0c) if log is not None: + 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: @@ -1004,7 +1010,7 @@ def _track_with_collective( self.line.log_last_track[kk].append(self.line.vv[kk]) else: ff = log[kk] - val = ff(self.line, particles) + val = ff(self.line, plog) if hasattr(ff, '_store'): for nn in ff._store: if nn not in self.line.log_last_track: @@ -1073,7 +1079,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) From 48329af22f2f7a7d80294cad60c3292e91190aaa Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 21 Jul 2024 22:04:04 +0200 Subject: [PATCH 002/159] Log only at start turn --- xtrack/tracker.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/xtrack/tracker.py b/xtrack/tracker.py index 7279025ef..f16baf1d5 100644 --- a/xtrack/tracker.py +++ b/xtrack/tracker.py @@ -958,8 +958,8 @@ 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 @@ -998,7 +998,7 @@ def _track_with_collective( p0c = self.line.particle_ref._xobject.p0c[0] particles.update_p0c_and_energy_deviations(p0c) - if log is not None: + if log is not None and not _is_resume_within_turn(tt, tt_resume): if _session_to_resume is not None: plog = _session_to_resume['particles'] else: @@ -1594,4 +1594,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 From 3b3f0c07b6cc498980fbab89ea5161dc0b0940c0 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 21 Jul 2024 22:37:33 +0200 Subject: [PATCH 003/159] Refactor log and time-dependent vars --- xtrack/tracker.py | 116 +++++++++++++++++++++++++--------------------- 1 file changed, 62 insertions(+), 54 deletions(-) diff --git a/xtrack/tracker.py b/xtrack/tracker.py index f16baf1d5..ea807e5c1 100644 --- a/xtrack/tracker.py +++ b/xtrack/tracker.py @@ -964,62 +964,11 @@ def _track_with_collective( # 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] - - 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) + self._handle_time_dependent_vars(particles=particles) if log is not None and not _is_resume_within_turn(tt, tt_resume): - 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) + self._handle_log(_session_to_resume=_session_to_resume, + particles=particles, log=log) moveback_to_buffer = None moveback_to_offset = None @@ -1489,6 +1438,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): From dbd1773286449bb6802aa0fa8ec31fdfb59ffc65 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 21 Jul 2024 23:15:14 +0200 Subject: [PATCH 004/159] Correct typo --- xtrack/pipeline/manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/pipeline/manager.py b/xtrack/pipeline/manager.py index ff0157aab..c7d8e8f9b 100644 --- a/xtrack/pipeline/manager.py +++ b/xtrack/pipeline/manager.py @@ -61,7 +61,7 @@ def is_ready_to_send(self,element_name,sender_name,reciever_name,turn,internal_t return False if not self._pending_requests[key].Test(): if self.verbose: - _print(f'Pipeline manager {element_name}: {sender_name} at rank {self.get_particles_rank(sender_name)} previous message to {reciever_name} at rank {self.get_particles_rank(reciever_name)} with tag {tag} was not receviced yet') + _print(f'Pipeline manager {element_name}: {sender_name} at rank {self.get_particles_rank(sender_name)} previous message to {reciever_name} at rank {self.get_particles_rank(reciever_name)} with tag {tag} was not received yet') return False return True From 0acd5ddec362e7183403094eb6f082410f285283 Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 22 Jul 2024 08:37:46 +0200 Subject: [PATCH 005/159] Default start end for open twiss --- xtrack/twiss.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 119725ec4..7920e979e 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -330,6 +330,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: + # is open twiss + start = start or xt.START + end = end or xt.END + if num_turns != 1: # Untested cases assert num_turns > 0 From cea45126ddc5722a97b2d520df4f1d921dfdc6fc Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 28 Jul 2024 14:41:16 +0200 Subject: [PATCH 006/159] Illustration preparation --- examples/coasting/e000_illustration.py | 38 ++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 examples/coasting/e000_illustration.py diff --git a/examples/coasting/e000_illustration.py b/examples/coasting/e000_illustration.py new file mode 100644 index 000000000..6e0581038 --- /dev/null +++ b/examples/coasting/e000_illustration.py @@ -0,0 +1,38 @@ +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 = 35 +t = np.linspace(0, t_max, 10000) + +s = np.zeros((len(l), len(t))) +for il, ll in enumerate(l): + s[il, :] = np.mod(l_ref/2 + 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) +sp1 = plt.subplot(311) +sp2 = plt.subplot(312, sharex=sp1) +sp3 = plt.subplot(313, sharex=sp1) + +sp1.plot(t, passing[0, :]) +sp2.plot(t, passing[1, :]) +sp3.plot(t, passing[2, :]) + +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) + +plt.show() + From f16005a388a832cf14d35a4b35346ce1a99169a3 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 3 Aug 2024 15:33:07 +0200 Subject: [PATCH 007/159] Plot for presentation --- examples/coasting/001_frev_meas.py | 49 +++++++++++++++++--------- examples/coasting/e000_illustration.py | 25 +++++++++---- xtrack/synctime.py | 4 +-- 3 files changed, 53 insertions(+), 25 deletions(-) diff --git a/examples/coasting/001_frev_meas.py b/examples/coasting/001_frev_meas.py index 285694372..5cc80f74b 100644 --- a/examples/coasting/001_frev_meas.py +++ b/examples/coasting/001_frev_meas.py @@ -5,6 +5,20 @@ from scipy.constants import c as clight delta0 = 1e-2 +delta_range = 0 +num_turns=200 +num_particles = 50000 + +delta0 = -1e-2 +delta_range = 0 +num_turns=20 +num_particles = 50000 + +# # 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'] @@ -136,9 +150,9 @@ def y_mean_hist(line, particles): print('f_measured: ', f_measured, ' Hz') print('Error: ', f_measured - f_expected, 'Hz') -assert np.isclose(f_expected, f_measured, rtol=0, atol=2) # 2 Hz tolerance -assert np.isclose(np.mean(inten), inten_exp, rtol=1e-2, atol=0) -assert np.allclose(p.at_turn, num_turns*0.9, rtol=3e-2, atol=0) #beta1 defaults to 0.1 +# assert np.isclose(f_expected, f_measured, rtol=0, atol=2) # 2 Hz tolerance +# assert np.isclose(np.mean(inten), inten_exp, rtol=1e-2, atol=0) +# assert np.allclose(p.at_turn, num_turns*0.9, rtol=3e-2, atol=0) #beta1 defaults to 0.1 tt = line.get_table() tt_synch = tt.rows[tt.element_type=='SyncTime'] @@ -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 index 6e0581038..ed5198af1 100644 --- a/examples/coasting/e000_illustration.py +++ b/examples/coasting/e000_illustration.py @@ -8,31 +8,42 @@ l = np.array([0.95, 1, 1.05]) -t_max = 35 +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(l_ref/2 + v * t * ll/l_ref, l_ref) + s[il, :] = np.mod(0.62*l_ref + v * t * ll/l_ref, l_ref) -passing = np.diff(s, prepend=0)< 0 +passing = np.diff(s, prepend=0) < 0 import matplotlib.pyplot as plt plt.close('all') -plt.figure(1) +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.plot(t, passing[0, :]) -sp2.plot(t, passing[1, :]) -sp3.plot(t, passing[2, :]) +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/xtrack/synctime.py b/xtrack/synctime.py index d17dbe01a..abdc1a317 100644 --- a/xtrack/synctime.py +++ b/xtrack/synctime.py @@ -47,8 +47,8 @@ def track(self, particles): # Check if some particles are too fast mask_too_fast = mask_alive & ( particles.zeta > zeta_min + self.circumference * beta0_beta1) - if mask_too_fast.any(): - raise ValueError('Some particles move faster than the time window') + # if mask_too_fast.any(): + # raise ValueError('Some particles move faster than the time window') # Update zeta for particles that are stopped particles.zeta[mask_stop] += beta0_beta1 * self.circumference From bc48fcb6496a93c2c629eb2f586a63d6f18692f5 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 7 Aug 2024 14:30:09 +0200 Subject: [PATCH 008/159] Re-enable check --- examples/coasting/001_frev_meas.py | 14 +++++++------- xtrack/synctime.py | 11 +++++++++-- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/examples/coasting/001_frev_meas.py b/examples/coasting/001_frev_meas.py index 5cc80f74b..0b4df59ab 100644 --- a/examples/coasting/001_frev_meas.py +++ b/examples/coasting/001_frev_meas.py @@ -4,15 +4,15 @@ from scipy.constants import c as clight -delta0 = 1e-2 +delta0 = 0 #-1e-2 delta_range = 0 -num_turns=200 -num_particles = 50000 +num_turns=100 +num_particles = 100_000 -delta0 = -1e-2 -delta_range = 0 -num_turns=20 -num_particles = 50000 +# delta0 = 1e-2 +# delta_range = 0 +# num_turns=20 +# num_particles = 5_000_000 # # To see the different number of turns # delta0 = 0e-2 diff --git a/xtrack/synctime.py b/xtrack/synctime.py index abdc1a317..78d473c8c 100644 --- a/xtrack/synctime.py +++ b/xtrack/synctime.py @@ -47,8 +47,15 @@ def track(self, particles): # Check if some particles are too fast mask_too_fast = mask_alive & ( particles.zeta > zeta_min + self.circumference * beta0_beta1) - # if mask_too_fast.any(): - # raise ValueError('Some particles move faster than the time window') + 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 From eea21f12f0ebf3a9709ea0bd9fbe94518994ca7e Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 16 Aug 2024 15:04:19 +0200 Subject: [PATCH 009/159] adapt test to xboject PR139 --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 520e12852..ca1397cbe 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,7 +11,7 @@ def assert_context_empty(context): An active buffer is a buffer that has not been freed yet. """ gc.collect() - alive_buffer_count = sum(b.alive for b in context.buffers) + alive_buffer_count = len(context._buffers) if alive_buffer_count > 0: pytest.fail(f"There were {alive_buffer_count} active buffers after a " f"test run, which points to a memory leak during the test " From b3e1b54143e829029f903ac200f1af8c61991ba4 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Mon, 26 Aug 2024 13:48:06 +0200 Subject: [PATCH 010/159] refactor plot --- xtrack/twiss.py | 26 +++++++++---- xtrack/twissplot.py | 94 +++++++++++++++++++++++++-------------------- 2 files changed, 71 insertions(+), 49 deletions(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 7c184da22..4c37bc5c2 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3361,8 +3361,13 @@ def plot(self,yl="",yr="",x='s', 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, + ): """ Plot columns of the TwissTable @@ -3402,15 +3407,22 @@ 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 + ) if labels is not None: mask=self.mask[labels] diff --git a/xtrack/twissplot.py b/xtrack/twissplot.py index 389b11aed..48262fd69 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,12 @@ 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, ): import matplotlib.pyplot as plt @@ -94,34 +103,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) + 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) @@ -133,14 +142,16 @@ def __init__( self.right.yaxis.set_ticks_position("right") # 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.78) + 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): @@ -204,11 +215,10 @@ 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.ax.set_xlabel(_mylbl(self.axlabel, self.x)) + self.ax.set_xlim(min(self.xaxis), max(self.xaxis)) self.figure.legend(self.lines, self.legends, loc="upper right") - ca.grid(True) + self.ax.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) @@ -239,7 +249,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: @@ -288,17 +298,17 @@ 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): + 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) return self From 61bdceb00643d681d540237aef1132bf51473236 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Tue, 27 Aug 2024 13:58:51 +0200 Subject: [PATCH 011/159] restructure plot --- xtrack/twiss.py | 4 ++-- xtrack/twissplot.py | 26 +++++++++++++++++++++++--- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 4c37bc5c2..4b2ba778e 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3356,7 +3356,7 @@ 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, @@ -3393,7 +3393,7 @@ 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' diff --git a/xtrack/twissplot.py b/xtrack/twissplot.py index 48262fd69..840286b75 100644 --- a/xtrack/twissplot.py +++ b/xtrack/twissplot.py @@ -116,7 +116,7 @@ def __init__( self.clist.append(self.color[i]) if lattice and x == "s": self.lattice = self._new_axis(axlattice) - self.lattice.set_frame_on(False) + #self.lattice.set_frame_on(False) # self.lattice.set_autoscale_on(False) self.lattice.yaxis.set_visible(False) if yl: @@ -217,7 +217,11 @@ def run(self): self._column(i, self.right, self.color[i]) self.ax.set_xlabel(_mylbl(self.axlabel, self.x)) self.ax.set_xlim(min(self.xaxis), max(self.xaxis)) - self.figure.legend(self.lines, self.legends, loc="upper right") + self.ax.legend( + self.lines,self.legends, + loc='upper right', + bbox_to_anchor=(1.3, 1.1)) + #self.figure.legend(self.lines, self.legends, loc="upper right") self.ax.grid(True) # self.figure.canvas.mpl_connect('button_release_event',self.button_press) self.figure.canvas.mpl_connect("pick_event", self.pick) @@ -298,7 +302,8 @@ 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): + 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 @@ -311,4 +316,19 @@ def ylim(self, left_lo=None, left_hi=None, right_lo=None, right_hi=None): if right_hi is None: right_hi = 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 + + + From 5433e54a93e5ab0203d8cee13ce9662e462a1b59 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Tue, 27 Aug 2024 14:31:48 +0200 Subject: [PATCH 012/159] add example --- examples/twiss/031_twiss_plot.py | 33 ++++++++++++++++++++++++++ xtrack/twiss.py | 8 ++++++- xtrack/twissplot.py | 40 +++++++++++++++++++------------- 3 files changed, 64 insertions(+), 17 deletions(-) create mode 100644 examples/twiss/031_twiss_plot.py diff --git a/examples/twiss/031_twiss_plot.py b/examples/twiss/031_twiss_plot.py new file mode 100644 index 000000000..6b6738bfa --- /dev/null +++ b/examples/twiss/031_twiss_plot.py @@ -0,0 +1,33 @@ +# 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() + +tw.plot(figlabel="tw") +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/xtrack/twiss.py b/xtrack/twiss.py index 4b2ba778e..570af035b 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3356,7 +3356,9 @@ def target(self, tars=None, value=None, at=None, **kwargs): action=self._action, **kwargs) return tarset - def plot(self,yl=None,yr=None,x='s', + def plot(self, + yl=None, + yr=None,x='s', lattice=True, mask=None, labels=None, @@ -3396,6 +3398,10 @@ def plot(self,yl=None,yr=None,x='s', 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 lattice and 'length' not in self.keys(): self.add_strengths() diff --git a/xtrack/twissplot.py b/xtrack/twissplot.py index 840286b75..a2ae56069 100644 --- a/xtrack/twissplot.py +++ b/xtrack/twissplot.py @@ -103,11 +103,11 @@ def __init__( idx, clist, ) - self.ax=ax + self.ax = ax self.used_ax = False if ax is not None: self.figure = ax.figure - elif figure is None: + elif figure is None: self.figure = plt.figure(num=figlabel) if figlabel is not None: self.figure.clf() @@ -116,7 +116,7 @@ def __init__( self.clist.append(self.color[i]) if lattice and x == "s": self.lattice = self._new_axis(axlattice) - #self.lattice.set_frame_on(False) + # self.lattice.set_frame_on(False) # self.lattice.set_autoscale_on(False) self.lattice.yaxis.set_visible(False) if yl: @@ -146,12 +146,12 @@ def _new_axis(self, ax=None): if self.ax is None: out = self.figure.add_subplot(111) self.figure.subplots_adjust(right=0.78) - self.ax=out + self.ax = out if self.used_ax: out = self.ax.twinx() else: out = self.ax - self.used_ax=True + self.used_ax = True return out def __repr__(self): @@ -218,10 +218,9 @@ def run(self): 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.3, 1.1)) - #self.figure.legend(self.lines, self.legends, loc="upper right") + self.lines, self.legends, loc="upper right", bbox_to_anchor=(1.3, 1.1) + ) + # self.figure.legend(self.lines, self.legends, loc="upper right") self.ax.grid(True) # self.figure.canvas.mpl_connect('button_release_event',self.button_press) self.figure.canvas.mpl_connect("pick_event", self.pick) @@ -302,8 +301,15 @@ 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, - lattice_lo=None, lattice_hi=None): + 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 @@ -324,11 +330,13 @@ def ylim(self, left_lo=None, left_hi=None, right_lo=None, right_hi=None, 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) + 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 From 4d6a1e1ca8f42da0e6aac126d1fa5f70deae1a16 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Tue, 27 Aug 2024 14:32:32 +0200 Subject: [PATCH 013/159] black --- examples/twiss/031_twiss_plot.py | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/examples/twiss/031_twiss_plot.py b/examples/twiss/031_twiss_plot.py index 6b6738bfa..c01154f38 100644 --- a/examples/twiss/031_twiss_plot.py +++ b/examples/twiss/031_twiss_plot.py @@ -8,8 +8,7 @@ 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 = 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 @@ -18,16 +17,11 @@ tw.plot(figlabel="tw") 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 = 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) - - - - - +pl1.move_legend(1.5, 1.0) +pl2.move_legend(1.5, 1.0) From 2d98733434d99013566102eed6d56962991e3f1c Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Tue, 27 Aug 2024 16:09:40 +0200 Subject: [PATCH 014/159] add hover --- examples/twiss/031_twiss_plot.py | 4 ++-- xtrack/twiss.py | 4 +++- xtrack/twissplot.py | 35 ++++++++++++++++++++++++-------- 3 files changed, 31 insertions(+), 12 deletions(-) diff --git a/examples/twiss/031_twiss_plot.py b/examples/twiss/031_twiss_plot.py index c01154f38..26e5683db 100644 --- a/examples/twiss/031_twiss_plot.py +++ b/examples/twiss/031_twiss_plot.py @@ -14,8 +14,8 @@ # Twiss tw = line.twiss4d() -tw.plot(figlabel="tw") -tw.plot(figlabel="tw") +pl = tw.plot(figlabel="tw") +pl = tw.plot(figlabel="tw") fig = plt.figure() ax1 = plt.subplot(211) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 570af035b..9720f6c27 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3369,6 +3369,7 @@ def plot(self, axleft=None, axright=None, axlattice=None, + hover=False ): """ Plot columns of the TwissTable @@ -3427,7 +3428,8 @@ def plot(self, ax=ax, axleft=axleft, axright=axright, - axlattice=axlattice + axlattice=axlattice, + hover=hover, ) if labels is not None: diff --git a/xtrack/twissplot.py b/xtrack/twissplot.py index a2ae56069..d9889d793 100644 --- a/xtrack/twissplot.py +++ b/xtrack/twissplot.py @@ -84,6 +84,7 @@ def __init__( axleft=None, axright=None, axlattice=None, + hover=False, ): import matplotlib.pyplot as plt @@ -140,6 +141,8 @@ 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_axis(self, ax=None): @@ -193,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=[] @@ -220,21 +224,32 @@ def run(self): self.ax.legend( self.lines, self.legends, loc="upper right", bbox_to_anchor=(1.3, 1.1) ) - # self.figure.legend(self.lines, self.legends, loc="upper right") self.ax.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.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]) @@ -279,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] @@ -295,6 +311,7 @@ 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): From 6cc2cde58428883c36a497bbc930be6c69b91a71 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 09:57:27 +0200 Subject: [PATCH 015/159] Draft the change on the C side --- xtrack/tracker.py | 36 +++++++++++++++++------------------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/xtrack/tracker.py b/xtrack/tracker.py index ea807e5c1..558681143 100644 --- a/xtrack/tracker.py +++ b/xtrack/tracker.py @@ -545,19 +545,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 +608,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 +623,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); From 92b43a0685bab89017f44e87bce99b298c692f28 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 10:50:15 +0200 Subject: [PATCH 016/159] Seems to work :-) --- xtrack/line.py | 12 ++++++--- xtrack/twiss.py | 66 ++++++++++++++++++++++++++++++------------------- 2 files changed, 48 insertions(+), 30 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 1c8204600..228e3143c 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -1560,7 +1560,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 +1636,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 +1843,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 +1884,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): diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 38887e844..4bb8344e4 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -1274,14 +1274,15 @@ def _compute_chromatic_functions(line, init, delta_chrom, steps_r_matrix, 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) + 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', @@ -1763,7 +1764,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') @@ -1797,6 +1797,7 @@ 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) @@ -1834,13 +1835,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')) @@ -2065,7 +2064,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 @@ -2079,6 +2079,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 @@ -2091,6 +2092,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') @@ -2166,7 +2168,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 @@ -2175,7 +2178,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]) @@ -2205,7 +2209,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] @@ -2221,6 +2225,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]}') @@ -2233,11 +2242,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], @@ -2246,8 +2255,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], @@ -2256,8 +2265,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], @@ -2272,7 +2281,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: @@ -2327,9 +2337,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) @@ -2340,6 +2355,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) @@ -3878,8 +3897,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 From cd4a8c9d82a6d59ca4d8d952ae1aefca000900c7 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 10:55:11 +0200 Subject: [PATCH 017/159] ddx is working --- xtrack/twiss.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 4bb8344e4..683fad31d 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -1270,10 +1270,7 @@ 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, + 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 From a9003aede8e13a00cec592899cd2123c5ef1cf62 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 11:46:40 +0200 Subject: [PATCH 018/159] Adapted test to check second order dispersion --- examples/symm_twiss_and_match/t000_test_twiss_symm.py | 4 ++++ tests/test_periodic_symmetric_twiss_and_match.py | 4 ++++ 2 files changed, 8 insertions(+) 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..21fa4f40a 100644 --- a/examples/symm_twiss_and_match/t000_test_twiss_symm.py +++ b/examples/symm_twiss_and_match/t000_test_twiss_symm.py @@ -134,6 +134,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) diff --git a/tests/test_periodic_symmetric_twiss_and_match.py b/tests/test_periodic_symmetric_twiss_and_match.py index 0719bf6fb..0c3ad52ee 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,6 +148,8 @@ 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) From 15e1112b5f5c902995e6e51e5bb03915871bfa6e Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 12:15:24 +0200 Subject: [PATCH 019/159] Add check for off momentum orbit --- .../t000_test_twiss_symm.py | 6 ++ .../t001_test_off_momentum.py | 77 +++++++++++++++++++ ...test_periodic_symmetric_twiss_and_match.py | 6 ++ xtrack/twiss.py | 9 --- 4 files changed, 89 insertions(+), 9 deletions(-) create mode 100644 examples/symm_twiss_and_match/t001_test_off_momentum.py 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 21fa4f40a..4ac82d77a 100644 --- a/examples/symm_twiss_and_match/t000_test_twiss_symm.py +++ b/examples/symm_twiss_and_match/t000_test_twiss_symm.py @@ -153,3 +153,9 @@ 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/examples/symm_twiss_and_match/t001_test_off_momentum.py b/examples/symm_twiss_and_match/t001_test_off_momentum.py new file mode 100644 index 000000000..a1d5f5243 --- /dev/null +++ b/examples/symm_twiss_and_match/t001_test_off_momentum.py @@ -0,0 +1,77 @@ +import xtrack as xt +import xobjects as xo +import numpy as np + +delta0 = 1e-3 + +# Build line with half a cell +half_cell = xt.Line( + elements={ + 'start_cell': xt.Marker(), + 'drift0': xt.Drift(length=1.), + 'qf1': xt.Quadrupole(k1=0.027/2, length=1.), + 'drift1_1': xt.Drift(length=1), + 'bend1': xt.Bend(k0=3e-4, h=3e-4, length=45.), + 'drift1_2': xt.Drift(length=1.), + 'qd1': xt.Quadrupole(k1=-0.0271/2, length=1.), + 'drift2': xt.Drift(length=1), + 'mid_cell': xt.Marker(), + } +) +half_cell.particle_ref = xt.Particles(p0c=2e9) + +tw_half_cell = half_cell.twiss4d(init='periodic_symmetric', # <--- periodic-symmetric boundary + strengths=True, # to get the strengths in table + delta0=delta0 + ) + +cell = xt.Line( + elements={ + 'start_cell': xt.Marker(), + 'drift0': xt.Drift(length=1.), + 'qf1': xt.Quadrupole(k1=0.027/2, length=1.), + 'drift1': xt.Drift(length=1), + 'bend1': xt.Bend(k0=3e-4, h=3e-4, length=45.), + 'drift2': xt.Drift(length=1.), + 'qd1': xt.Quadrupole(k1=-0.0271/2, length=1.), + 'drift3': xt.Drift(length=1), + 'mid_cell': xt.Marker(), + 'drift4': xt.Replica('drift3'), + 'qd2': xt.Replica('qd1'), + 'drift5': xt.Replica('drift2'), + 'bend2': xt.Replica('bend1'), + 'drift6': xt.Replica('drift1'), + 'qf2': xt.Replica('qf1'), + 'drift7': xt.Replica('drift0'), + 'end_cell': xt.Marker(), + } +) +cell.particle_ref = xt.Particles(p0c=2e9) +tw_cell = cell.twiss4d(strengths=True, delta0=delta0) + +xo.assert_allclose(tw_half_cell.betx[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].betx, atol=0, rtol=1e-8) +xo.assert_allclose(tw_half_cell.bety[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].bety, atol=0, rtol=1e-8) +xo.assert_allclose(tw_half_cell.alfx[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].alfx, atol=1e-8, rtol=0) +xo.assert_allclose(tw_half_cell.alfy[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].alfy, atol=1e-8, rtol=0) +xo.assert_allclose(tw_half_cell.dx[:-1], # remove '_end_point' + 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.ax_chrom[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].ax_chrom, atol=1e-5, rtol=0) +xo.assert_allclose(tw_half_cell.ay_chrom[:-1], # remove '_end_point' + tw_cell.rows[:'mid_cell'].ay_chrom, atol=1e-5, rtol=0) +xo.assert_allclose(tw_half_cell.bx_chrom[:-1], # remove '_end_point' + 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.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) diff --git a/tests/test_periodic_symmetric_twiss_and_match.py b/tests/test_periodic_symmetric_twiss_and_match.py index 0c3ad52ee..b51de4b06 100644 --- a/tests/test_periodic_symmetric_twiss_and_match.py +++ b/tests/test_periodic_symmetric_twiss_and_match.py @@ -156,3 +156,9 @@ def test_periodic_symmetric_twiss_and_match(): 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/xtrack/twiss.py b/xtrack/twiss.py index 683fad31d..bcb73cc7b 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -1774,7 +1774,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 @@ -1799,14 +1798,6 @@ def _find_periodic_solution(line, particle_on_co, particle_ref, method, 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 From 561ab5ee01427dcff162c19093f9a29a04a89ef7 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 12:18:41 +0200 Subject: [PATCH 020/159] Clean up --- .../t001_test_off_momentum.py | 77 ------------------- 1 file changed, 77 deletions(-) delete mode 100644 examples/symm_twiss_and_match/t001_test_off_momentum.py diff --git a/examples/symm_twiss_and_match/t001_test_off_momentum.py b/examples/symm_twiss_and_match/t001_test_off_momentum.py deleted file mode 100644 index a1d5f5243..000000000 --- a/examples/symm_twiss_and_match/t001_test_off_momentum.py +++ /dev/null @@ -1,77 +0,0 @@ -import xtrack as xt -import xobjects as xo -import numpy as np - -delta0 = 1e-3 - -# Build line with half a cell -half_cell = xt.Line( - elements={ - 'start_cell': xt.Marker(), - 'drift0': xt.Drift(length=1.), - 'qf1': xt.Quadrupole(k1=0.027/2, length=1.), - 'drift1_1': xt.Drift(length=1), - 'bend1': xt.Bend(k0=3e-4, h=3e-4, length=45.), - 'drift1_2': xt.Drift(length=1.), - 'qd1': xt.Quadrupole(k1=-0.0271/2, length=1.), - 'drift2': xt.Drift(length=1), - 'mid_cell': xt.Marker(), - } -) -half_cell.particle_ref = xt.Particles(p0c=2e9) - -tw_half_cell = half_cell.twiss4d(init='periodic_symmetric', # <--- periodic-symmetric boundary - strengths=True, # to get the strengths in table - delta0=delta0 - ) - -cell = xt.Line( - elements={ - 'start_cell': xt.Marker(), - 'drift0': xt.Drift(length=1.), - 'qf1': xt.Quadrupole(k1=0.027/2, length=1.), - 'drift1': xt.Drift(length=1), - 'bend1': xt.Bend(k0=3e-4, h=3e-4, length=45.), - 'drift2': xt.Drift(length=1.), - 'qd1': xt.Quadrupole(k1=-0.0271/2, length=1.), - 'drift3': xt.Drift(length=1), - 'mid_cell': xt.Marker(), - 'drift4': xt.Replica('drift3'), - 'qd2': xt.Replica('qd1'), - 'drift5': xt.Replica('drift2'), - 'bend2': xt.Replica('bend1'), - 'drift6': xt.Replica('drift1'), - 'qf2': xt.Replica('qf1'), - 'drift7': xt.Replica('drift0'), - 'end_cell': xt.Marker(), - } -) -cell.particle_ref = xt.Particles(p0c=2e9) -tw_cell = cell.twiss4d(strengths=True, delta0=delta0) - -xo.assert_allclose(tw_half_cell.betx[:-1], # remove '_end_point' - tw_cell.rows[:'mid_cell'].betx, atol=0, rtol=1e-8) -xo.assert_allclose(tw_half_cell.bety[:-1], # remove '_end_point' - tw_cell.rows[:'mid_cell'].bety, atol=0, rtol=1e-8) -xo.assert_allclose(tw_half_cell.alfx[:-1], # remove '_end_point' - tw_cell.rows[:'mid_cell'].alfx, atol=1e-8, rtol=0) -xo.assert_allclose(tw_half_cell.alfy[:-1], # remove '_end_point' - tw_cell.rows[:'mid_cell'].alfy, atol=1e-8, rtol=0) -xo.assert_allclose(tw_half_cell.dx[:-1], # remove '_end_point' - 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.ax_chrom[:-1], # remove '_end_point' - tw_cell.rows[:'mid_cell'].ax_chrom, atol=1e-5, rtol=0) -xo.assert_allclose(tw_half_cell.ay_chrom[:-1], # remove '_end_point' - tw_cell.rows[:'mid_cell'].ay_chrom, atol=1e-5, rtol=0) -xo.assert_allclose(tw_half_cell.bx_chrom[:-1], # remove '_end_point' - 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.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) From b532353e4d9588ad4f2b41f0248c400ced3a1da0 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 16:42:41 +0200 Subject: [PATCH 021/159] Some more plotting --- .../t000_test_twiss_symm.py | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) 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 4ac82d77a..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' @@ -159,3 +170,38 @@ 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() From 33ecebde0cc3b998baff4e5eacea0d62570d11df Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 17:13:10 +0200 Subject: [PATCH 022/159] Managing to define an element --- examples/lattice_design_shortcuts/000_dev.py | 33 ++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 examples/lattice_design_shortcuts/000_dev.py diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py new file mode 100644 index 000000000..eab81a79a --- /dev/null +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -0,0 +1,33 @@ +import xtrack as xt + +line = xt.Line() + +class Expr: + def __init__(self, expr): + self.expr = expr + +# line._xdeps_vref._eval("a * b") + +def _eval(line, expr): + return line._xdeps_vref._eval(expr) + +def _new_element(line, name, cls, **kwargs): + + evaluated_kwargs = {} + value_kwargs = {} + for kk in kwargs: + if isinstance(kwargs[kk], Expr): + evaluated_kwargs[kk] = _eval(line, kwargs[kk].expr) + value_kwargs[kk] = evaluated_kwargs[kk]._value + else: + evaluated_kwargs[kk] = kwargs[kk] + value_kwargs[kk] = kwargs[kk] + + element = cls(**value_kwargs) + line.element_dict[name] = element + for kk in kwargs: + setattr(line.element_refs[name], kk, evaluated_kwargs[kk]) + +line.vars['a'] = 2. + +_new_element(line, 'd', xt.Drift, length=Expr('a')) From 913763ca5c99de237f3dd23dcaf9a2b83f2ae24e Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 3 Sep 2024 22:36:04 +0200 Subject: [PATCH 023/159] Better... --- examples/lattice_design_shortcuts/000_dev.py | 57 +++++++++++++++++--- 1 file changed, 50 insertions(+), 7 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index eab81a79a..26b532fdc 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -1,6 +1,6 @@ import xtrack as xt - -line = xt.Line() +import xdeps as xd +import xobjects as xo class Expr: def __init__(self, expr): @@ -8,16 +8,23 @@ def __init__(self, expr): # line._xdeps_vref._eval("a * b") -def _eval(line, expr): - return line._xdeps_vref._eval(expr) def _new_element(line, name, cls, **kwargs): + _eval = line._eval_obj.eval + evaluated_kwargs = {} value_kwargs = {} for kk in kwargs: if isinstance(kwargs[kk], Expr): - evaluated_kwargs[kk] = _eval(line, kwargs[kk].expr) + evaluated_kwargs[kk] = _eval(kwargs[kk].expr) + value_kwargs[kk] = evaluated_kwargs[kk]._value + elif hasattr(kwargs[kk], '_value'): + evaluated_kwargs[kk] = kwargs[kk] + value_kwargs[kk] = kwargs[kk]._value + elif (isinstance(kwargs[kk], str) and hasattr(cls, '_xofields') + and kk in cls._xofields and cls._xofields[kk].__name__ != 'String'): + evaluated_kwargs[kk] = _eval(kwargs[kk]) value_kwargs[kk] = evaluated_kwargs[kk]._value else: evaluated_kwargs[kk] = kwargs[kk] @@ -28,6 +35,42 @@ def _new_element(line, name, cls, **kwargs): for kk in kwargs: setattr(line.element_refs[name], kk, evaluated_kwargs[kk]) -line.vars['a'] = 2. +def _call_vars(vars, *args, **kwargs): + _eval = vars.line._eval_obj.eval + if len(args) > 0: + assert len(kwargs) == 0 + assert len(args) == 1 + if isinstance(args[0], str): + return vars[args[0]] + elif isinstance(args[0], dict): + kwargs.update(args[0]) + else: + raise ValueError('Invalid argument') + for kk in kwargs: + if isinstance(kwargs[kk], Expr): + vars[kk] = _eval(kwargs[kk].expr) + elif isinstance(kwargs[kk], str): + vars[kk] = _eval(kwargs[kk]) + else: + vars[kk] = kwargs[kk] + +xt.Line.newele = _new_element +xt.line.LineVars.__call__ = _call_vars + +line = xt.Line() + +line._eval_obj = xd.madxutils.MadxEval(variables=line._xdeps_vref, + functions=line._xdeps_fref, + elements=line.element_dict) + +line.vars({ + 'kqf': 0.027, + 'kqd': -0.0271, + 'kqf.1': 'kqf / 2', + 'kqd.1': 'kqd / 2', +}) + +line.newele('qf.1', xt.Quadrupole, k1='kqf.1', length=1.) +line.newele('qd.1', xt.Quadrupole, k1='kqd.1', length=1.) + -_new_element(line, 'd', xt.Drift, length=Expr('a')) From 303dc45d08f00355dd17b2ba26c5b5b1123f5cc2 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 4 Sep 2024 12:34:51 +0200 Subject: [PATCH 024/159] Getting somewhere --- examples/lattice_design_shortcuts/000_dev.py | 83 ++++++++++++++++++-- 1 file changed, 76 insertions(+), 7 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 26b532fdc..bc3facad2 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -2,6 +2,8 @@ import xdeps as xd import xobjects as xo +import numpy as np + class Expr: def __init__(self, expr): self.expr = expr @@ -35,6 +37,8 @@ def _new_element(line, name, cls, **kwargs): for kk in kwargs: setattr(line.element_refs[name], kk, evaluated_kwargs[kk]) + return name + def _call_vars(vars, *args, **kwargs): _eval = vars.line._eval_obj.eval if len(args) > 0: @@ -54,23 +58,88 @@ def _call_vars(vars, *args, **kwargs): else: vars[kk] = kwargs[kk] -xt.Line.newele = _new_element +def _flatten_components(components): + flatten_components = [] + for nn in components: + if isinstance(nn, Section): + flatten_components += _flatten_components(nn.components) + else: + flatten_components.append(nn) + return flatten_components + +class Section: + def __init__(self, line, components, name=None): + self.line = line + self.components = _flatten_components(components) + self.name = name + + def mirror(self): + self.components = self.components[::-1] + + def replicate(self, name): + new_components = [] + for nn in self.components: + new_nn = nn + '.' + name + self.line.element_dict[new_nn] = xt.Replica(nn) + new_components.append(new_nn) + return Section(self.line, new_components, name=name) + + def twiss4d(self, **kwargs): + temp_line = xt.Line(elements=self.line.element_dict, + element_names=self.components) + temp_line.particle_ref = self.line.particle_ref + return temp_line.twiss4d(**kwargs) + +def _section(line, components, name=None): + return Section(line, components, name=name) + +def _append_section(line, section): + line.element_names += section.components + +xt.Line.new_element = _new_element xt.line.LineVars.__call__ = _call_vars +xt.Line.new_section = _section +xt.Line.append_section = _append_section line = xt.Line() +line.particle_ref = xt.Particles(p0c=2e9) line._eval_obj = xd.madxutils.MadxEval(variables=line._xdeps_vref, functions=line._xdeps_fref, elements=line.element_dict) line.vars({ - 'kqf': 0.027, - 'kqd': -0.0271, - 'kqf.1': 'kqf / 2', - 'kqd.1': 'kqd / 2', + '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': 45., + 'angle.mb': 2 * np.pi / 1232. * 3, + 'k0.mb': 'angle.mb / l.mb', }) -line.newele('qf.1', xt.Quadrupole, k1='kqf.1', length=1.) -line.newele('qd.1', xt.Quadrupole, k1='kqd.1', length=1.) +halfcell = line.new_section(components=[ + line.new_element('start', xt.Marker), + line.new_element('drift.1', xt.Drift, length='l.mq / 2'), + line.new_element('qf', xt.Quadrupole, k1='kqf.1', length='l.mq'), + line.new_element('drift.2', xt.Replica, parent_name='drift.1'), + line.new_element('mb.1', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb / 4'), + line.new_element('mb.2', xt.Replica, parent_name='mb.1'), + line.new_element('mb.3', xt.Replica, parent_name='mb.1'), + line.new_element('mb.4', xt.Replica, parent_name='mb.1'), + line.new_element('drift.3', xt.Replica, parent_name='drift.1'), + line.new_element('qd', xt.Quadrupole, k1='kqd.1', length='l.mq'), + line.new_element('drift.4', xt.Replica, parent_name='drift.1'), +]) + +hcell_left = halfcell.replicate('l') +hcell_right = halfcell.replicate('r') +hcell_right.mirror() + +cell= line.new_section(components=[hcell_left, hcell_right]) + +cell.twiss4d().plot() + From 299ac37b8f33fc34d891226707213feeb7ede6be Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 4 Sep 2024 14:05:30 +0200 Subject: [PATCH 025/159] Section now inherits from Line --- examples/lattice_design_shortcuts/000_dev.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index bc3facad2..7a86ae5fc 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -67,14 +67,24 @@ def _flatten_components(components): flatten_components.append(nn) return flatten_components -class Section: +class Section(xt.Line): def __init__(self, line, components, name=None): self.line = line - self.components = _flatten_components(components) - self.name = name + xt.Line.__init__(self, elements=line.element_dict, + element_names=_flatten_components(components)) + self._var_management = line._var_management + self._name = name + + @property + def name(self): + return self._name + + @property + def components(self): + return self.element_names def mirror(self): - self.components = self.components[::-1] + self.element_names = self.element_names[::-1] def replicate(self, name): new_components = [] From 32ddd6c69d1a1a74c416a70bc57e9656e36834d6 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 4 Sep 2024 15:48:52 +0200 Subject: [PATCH 026/159] A nice cell --- examples/lattice_design_shortcuts/000_dev.py | 30 ++++++++++++-------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 7a86ae5fc..01d0e6209 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -79,6 +79,14 @@ def __init__(self, line, components, name=None): def name(self): return self._name + @property + def particle_ref(self): + return self.line.particle_ref + + @particle_ref.setter + def particle_ref(self, value): + assert value is None + @property def components(self): return self.element_names @@ -94,12 +102,6 @@ def replicate(self, name): new_components.append(new_nn) return Section(self.line, new_components, name=name) - def twiss4d(self, **kwargs): - temp_line = xt.Line(elements=self.line.element_dict, - element_names=self.components) - temp_line.particle_ref = self.line.particle_ref - return temp_line.twiss4d(**kwargs) - def _section(line, components, name=None): return Section(line, components, name=name) @@ -124,20 +126,18 @@ def _append_section(line, section): 'l.mq': 0.5, 'kqf.1': 'k1l.qf / l.mq', 'kqd.1': 'k1l.qd / l.mq', - 'l.mb': 45., - 'angle.mb': 2 * np.pi / 1232. * 3, + 'l.mb': 12, + 'angle.mb': 2 * np.pi / 48 , 'k0.mb': 'angle.mb / l.mb', }) halfcell = line.new_section(components=[ - line.new_element('start', xt.Marker), line.new_element('drift.1', xt.Drift, length='l.mq / 2'), line.new_element('qf', xt.Quadrupole, k1='kqf.1', length='l.mq'), line.new_element('drift.2', xt.Replica, parent_name='drift.1'), - line.new_element('mb.1', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb / 4'), + line.new_element('mb.1', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb'), line.new_element('mb.2', xt.Replica, parent_name='mb.1'), line.new_element('mb.3', xt.Replica, parent_name='mb.1'), - line.new_element('mb.4', xt.Replica, parent_name='mb.1'), line.new_element('drift.3', xt.Replica, parent_name='drift.1'), line.new_element('qd', xt.Quadrupole, k1='kqd.1', length='l.mq'), line.new_element('drift.4', xt.Replica, parent_name='drift.1'), @@ -147,7 +147,13 @@ def _append_section(line, section): hcell_right = halfcell.replicate('r') hcell_right.mirror() -cell= line.new_section(components=[hcell_left, hcell_right]) +cell= line.new_section(components=[ + line.new_element('start', xt.Marker), + hcell_left, + line.new_element('mid', xt.Marker), + hcell_right, + line.new_element('end', xt.Marker), +]) cell.twiss4d().plot() From 205dfff9b3e3e02ae5a48ed405b54389ae07bff7 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 4 Sep 2024 16:38:39 +0200 Subject: [PATCH 027/159] Nice triangular ring --- examples/lattice_design_shortcuts/000_dev.py | 50 +++++++++++++++++++- 1 file changed, 48 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 01d0e6209..b2bf0691b 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -72,6 +72,7 @@ def __init__(self, line, components, name=None): self.line = line xt.Line.__init__(self, elements=line.element_dict, element_names=_flatten_components(components)) + self._element_dict = line.element_dict # Avoid copying self._var_management = line._var_management self._name = name @@ -120,6 +121,13 @@ def _append_section(line, section): functions=line._xdeps_fref, elements=line.element_dict) +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 + + line.vars({ 'k1l.qf': 0.027 / 2, 'k1l.qd': -0.0271 / 2, @@ -127,7 +135,7 @@ def _append_section(line, section): 'kqf.1': 'k1l.qf / l.mq', 'kqd.1': 'k1l.qd / l.mq', 'l.mb': 12, - 'angle.mb': 2 * np.pi / 48 , + 'angle.mb': 2 * np.pi / n_bends, 'k0.mb': 'angle.mb / l.mb', }) @@ -147,7 +155,7 @@ def _append_section(line, section): hcell_right = halfcell.replicate('r') hcell_right.mirror() -cell= line.new_section(components=[ +cell = line.new_section(components=[ line.new_element('start', xt.Marker), hcell_left, line.new_element('mid', xt.Marker), @@ -155,6 +163,44 @@ def _append_section(line, section): line.new_element('end', xt.Marker), ]) +arc = line.new_section(components=[ + cell.replicate(name='cell.1'), + cell.replicate(name='cell.2'), + cell.replicate(name='cell.3'), +]) + +cell_ss = cell.replicate('ss') +line.new_element('drift_ss', xt.Drift, length='l.mb') +for ii, nn in enumerate(cell_ss.components): + if nn.startswith('mb'): + cell_ss.components[ii] = line.new_element( + f'drift.{ii}.ss', xt.Replica, parent_name='drift_ss') + +ss = line.new_section(components=[ + cell_ss.replicate('cell.1'), + # cell_ss.replicate('cell.2'), + # cell_ss.replicate('cell.3'), + # cell_ss.replicate('cell.3'), +]) + +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') + +line.append_section(arc1) +line.append_section(ss1) +line.append_section(arc2) +line.append_section(ss2) +line.append_section(arc3) +line.append_section(ss3) + + + + cell.twiss4d().plot() From 7c677a1f28513c78ec0b4e2c1a4abc556e2a96bd Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 4 Sep 2024 23:55:30 +0200 Subject: [PATCH 028/159] Strange insertion --- examples/lattice_design_shortcuts/000_dev.py | 128 +++++++++++++++++-- 1 file changed, 114 insertions(+), 14 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index b2bf0691b..a6ff4e3a2 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -106,13 +106,31 @@ def replicate(self, name): def _section(line, components, name=None): return Section(line, components, name=name) -def _append_section(line, section): +def _append(line, section): line.element_names += section.components +def _replace_replica(line, name): + name_parent = line[name].resolve(line, get_name=True) + line.element_dict[name] = line[name_parent].copy() + + pars_with_expr = list( + line._xdeps_manager.tartasks[line.element_refs[name_parent]].keys()) + + for rr in pars_with_expr: + assert isinstance(rr, xd.refs.AttrRef) + setattr(line.element_refs[name], rr._key, rr._expr) + +def _replace_all_replicas(line): + for nn in line.element_names: + if isinstance(line[nn], xt.Replica): + _replace_replica(line, nn) + xt.Line.new_element = _new_element xt.line.LineVars.__call__ = _call_vars xt.Line.new_section = _section -xt.Line.append_section = _append_section +xt.Line.append = _append +xt.Line.replace_replica = _replace_replica +xt.Line.replace_all_replicas = _replace_all_replicas line = xt.Line() line.particle_ref = xt.Particles(p0c=2e9) @@ -151,8 +169,8 @@ def _append_section(line, section): line.new_element('drift.4', xt.Replica, parent_name='drift.1'), ]) -hcell_left = halfcell.replicate('l') -hcell_right = halfcell.replicate('r') +hcell_left = halfcell.replicate(name='l') +hcell_right = halfcell.replicate(name='r') # could add mirror=True hcell_right.mirror() cell = line.new_section(components=[ @@ -178,9 +196,7 @@ def _append_section(line, section): ss = line.new_section(components=[ cell_ss.replicate('cell.1'), - # cell_ss.replicate('cell.2'), - # cell_ss.replicate('cell.3'), - # cell_ss.replicate('cell.3'), + cell_ss.replicate('cell.2'), ]) arc1 = arc.replicate(name='arc.1') @@ -191,17 +207,101 @@ def _append_section(line, section): ss2 = ss.replicate(name='ss.2') ss3 = ss.replicate(name='ss.3') -line.append_section(arc1) -line.append_section(ss1) -line.append_section(arc2) -line.append_section(ss2) -line.append_section(arc3) -line.append_section(ss3) +line.append(arc1) # Rename to append +line.append(ss1) +line.append(arc2) +line.append(ss2) +line.append(arc3) +line.append(ss3) +line.replace_all_replicas() +opt = cell.match( + method='4d', + vary=xt.VaryList(['k1l.qf', 'k1l.qd'], step=1e-5), + targets=xt.TargetSet( + qx=0.333, + qy=0.333, + )) +tw = line.twiss4d() -cell.twiss4d().plot() +line.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', +}) + +ss_left = line.new_section(components=[ + line.new_element('ip', xt.Marker), + line.new_element('dd.r', xt.Drift, length=24), + line.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), + line.new_element('dd.1', xt.Drift, length=8), + line.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), + line.new_element('dd.2', xt.Drift, length=8), + line.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), + line.new_element('dd.3', xt.Drift, length=11), + line.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), + line.new_element('dd.4', xt.Drift, length=12), + line.new_element('mq.5', xt.Quadrupole, k1='k1l.q5', length='l.mq'), + line.new_element('dd.5', xt.Drift, length=10.5), + line.new_element('e.ss.r', xt.Marker), +]) +ss_left.build_tracker() + +tw_arc = arc.twiss4d() + +bet_ip = 100. + +opt = ss_left.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( + betx=bet_ip, bety=bet_ip, alfx=0, alfy=0, + at='ip' + )) + +opt.step(40) +opt.targets[0].value=50. +opt.targets[1].value=50. +opt.step(40) +opt.targets[0].value=10. +opt.targets[1].value=10. +opt.step(40) +opt.targets[0].value=5. +opt.targets[1].value=5. +opt.step(40) +opt.targets[0].value=2. +opt.targets[1].value=2. +opt.step(40) + +prrrr + +ss_arc = line.new_section(components=[ss_left, arc]) +ss_arc.cut_at_s(np.arange(0, ss_arc.get_length(), 0.5)) +tw_ss_arc = ss_arc.twiss4d(betx=2, bety=2) +tw_ss_arc.plot() + +import matplotlib.pyplot as plt +plt.close('all') +fig = plt.figure(1) +ax1 = fig.add_subplot(2, 1, 1) +tw.plot('betx bety', ax=ax1) +ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) +tw.plot('dx', ax=ax2) + +plt.show() From b4e96e84afebd3626b1cdb3be076a1f24e9fe13d Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 00:15:45 +0200 Subject: [PATCH 029/159] Insertion.... --- examples/lattice_design_shortcuts/000_dev.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index a6ff4e3a2..f122f4fba 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -243,22 +243,24 @@ def _replace_all_replicas(line): line.new_element('ip', xt.Marker), line.new_element('dd.r', xt.Drift, length=24), line.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), - line.new_element('dd.1', xt.Drift, length=8), - line.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), - line.new_element('dd.2', xt.Drift, length=8), + line.new_element('dd.1', xt.Drift, length=2), + line.new_element('mq.2a', xt.Quadrupole, k1='k1l.q2', length='l.mq'), + line.new_element('dd.2a', xt.Drift, length=2), + line.new_element('mq.2b', xt.Quadrupole, k1='k1l.q2', length='l.mq'), + line.new_element('dd.2b', xt.Drift, length=2), line.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), - line.new_element('dd.3', xt.Drift, length=11), + line.new_element('dd.3', xt.Drift, length=13), line.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), line.new_element('dd.4', xt.Drift, length=12), line.new_element('mq.5', xt.Quadrupole, k1='k1l.q5', length='l.mq'), - line.new_element('dd.5', xt.Drift, length=10.5), + line.new_element('dd.5', xt.Drift, length=12), line.new_element('e.ss.r', xt.Marker), ]) ss_left.build_tracker() tw_arc = arc.twiss4d() -bet_ip = 100. +bet_ip = 300. opt = ss_left.match( solve=False, @@ -272,6 +274,10 @@ def _replace_all_replicas(line): at='ip' )) + +opt.step(40) +opt.targets[0].value=200. +opt.targets[1].value=200. opt.step(40) opt.targets[0].value=50. opt.targets[1].value=50. From af2d69ce294c9b2affedfd71fd954bf0e7c16807 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 00:29:43 +0200 Subject: [PATCH 030/159] Still playing with straight --- examples/lattice_design_shortcuts/000_dev.py | 30 +++++++------------- 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index f122f4fba..13694a943 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -241,19 +241,17 @@ def _replace_all_replicas(line): ss_left = line.new_section(components=[ line.new_element('ip', xt.Marker), - line.new_element('dd.r', xt.Drift, length=24), + line.new_element('dd.0', xt.Drift, length=26), line.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), - line.new_element('dd.1', xt.Drift, length=2), - line.new_element('mq.2a', xt.Quadrupole, k1='k1l.q2', length='l.mq'), - line.new_element('dd.2a', xt.Drift, length=2), - line.new_element('mq.2b', xt.Quadrupole, k1='k1l.q2', length='l.mq'), - line.new_element('dd.2b', xt.Drift, length=2), + line.new_element('dd.1', xt.Drift, length=10), + line.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), + line.new_element('dd.2', xt.Drift, length=10), line.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), - line.new_element('dd.3', xt.Drift, length=13), + line.new_element('dd.3', xt.Drift, length=10), line.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), - line.new_element('dd.4', xt.Drift, length=12), + line.new_element('dd.4', xt.Drift, length=10), line.new_element('mq.5', xt.Quadrupole, k1='k1l.q5', length='l.mq'), - line.new_element('dd.5', xt.Drift, length=12), + line.new_element('dd.5', xt.Drift, length=7.5), line.new_element('e.ss.r', xt.Marker), ]) ss_left.build_tracker() @@ -279,24 +277,16 @@ def _replace_all_replicas(line): opt.targets[0].value=200. opt.targets[1].value=200. opt.step(40) -opt.targets[0].value=50. +opt.targets[0].value=100. opt.targets[1].value=50. opt.step(40) -opt.targets[0].value=10. -opt.targets[1].value=10. -opt.step(40) -opt.targets[0].value=5. -opt.targets[1].value=5. -opt.step(40) -opt.targets[0].value=2. -opt.targets[1].value=2. -opt.step(40) + prrrr ss_arc = line.new_section(components=[ss_left, arc]) ss_arc.cut_at_s(np.arange(0, ss_arc.get_length(), 0.5)) -tw_ss_arc = ss_arc.twiss4d(betx=2, bety=2) +tw_ss_arc = ss_arc.twiss4d(betx=opt.targets[0].value, bety=opt.targets[1].value) tw_ss_arc.plot() import matplotlib.pyplot as plt From 6b21be19567af49730c835125cb2fbf24f60162a Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 07:52:59 +0200 Subject: [PATCH 031/159] This is a nice one! --- examples/lattice_design_shortcuts/000_dev.py | 27 ++++++++++---------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 13694a943..26fd070a6 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -241,15 +241,15 @@ def _replace_all_replicas(line): ss_left = line.new_section(components=[ line.new_element('ip', xt.Marker), - line.new_element('dd.0', xt.Drift, length=26), + line.new_element('dd.0', xt.Drift, length=10), line.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), - line.new_element('dd.1', xt.Drift, length=10), + line.new_element('dd.1', xt.Drift, length=6), line.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), - line.new_element('dd.2', xt.Drift, length=10), + line.new_element('dd.2', xt.Drift, length=18), line.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), - line.new_element('dd.3', xt.Drift, length=10), + line.new_element('dd.3', xt.Drift, length=14), line.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), - line.new_element('dd.4', xt.Drift, length=10), + line.new_element('dd.4', xt.Drift, length=14), line.new_element('mq.5', xt.Quadrupole, k1='k1l.q5', length='l.mq'), line.new_element('dd.5', xt.Drift, length=7.5), line.new_element('e.ss.r', xt.Marker), @@ -268,27 +268,26 @@ def _replace_all_replicas(line): start='ip', end='e.ss.r', vary=xt.VaryList(['k1l.q1', 'k1l.q2', 'k1l.q3', 'k1l.q4'], step=1e-5), targets=xt.TargetSet( - betx=bet_ip, bety=bet_ip, alfx=0, alfy=0, + alfx=0, alfy=0, at='ip' )) opt.step(40) -opt.targets[0].value=200. -opt.targets[1].value=200. -opt.step(40) -opt.targets[0].value=100. -opt.targets[1].value=50. -opt.step(40) -prrrr + +tw_arc = arc.twiss4d() ss_arc = line.new_section(components=[ss_left, arc]) ss_arc.cut_at_s(np.arange(0, ss_arc.get_length(), 0.5)) -tw_ss_arc = ss_arc.twiss4d(betx=opt.targets[0].value, bety=opt.targets[1].value) +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) tw_ss_arc.plot() +prrrr + import matplotlib.pyplot as plt plt.close('all') fig = plt.figure(1) From 08a2cfd2787523d7cf894a22bbf9d60a38e22faf Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 07:56:17 +0200 Subject: [PATCH 032/159] This is nice --- examples/lattice_design_shortcuts/000_dev.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 26fd070a6..63eb1b885 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -267,10 +267,11 @@ def _replace_all_replicas(line): 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' - )) + targets=[ + xt.TargetSet(alfx=0, alfy=0, at='ip'), + xt.Target(lambda tw: tw.betx[0] - tw.bety[0], 0), + ] + ) opt.step(40) From 93891b9d8a91c91c516ee6c8f787005bec1f8ae9 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 08:11:46 +0200 Subject: [PATCH 033/159] Closed with insertion --- examples/lattice_design_shortcuts/000_dev.py | 43 +++++++++++--------- 1 file changed, 24 insertions(+), 19 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 63eb1b885..31d32d6f7 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -207,12 +207,7 @@ def _replace_all_replicas(line): ss2 = ss.replicate(name='ss.2') ss3 = ss.replicate(name='ss.3') -line.append(arc1) # Rename to append -line.append(ss1) -line.append(arc2) -line.append(ss2) -line.append(arc3) -line.append(ss3) + line.replace_all_replicas() @@ -224,8 +219,6 @@ def _replace_all_replicas(line): qy=0.333, )) -tw = line.twiss4d() - line.vars({ 'k1l.q1': 0.012, 'k1l.q2': -0.012, @@ -239,7 +232,7 @@ def _replace_all_replicas(line): 'k1.q5': 'k1l.q5 / l.mq', }) -ss_left = line.new_section(components=[ +half_straight = line.new_section(components=[ line.new_element('ip', xt.Marker), line.new_element('dd.0', xt.Drift, length=10), line.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), @@ -247,20 +240,20 @@ def _replace_all_replicas(line): line.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), line.new_element('dd.2', xt.Drift, length=18), line.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), - line.new_element('dd.3', xt.Drift, length=14), + line.new_element('dd.3', xt.Drift, length=16), line.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), - line.new_element('dd.4', xt.Drift, length=14), + line.new_element('dd.4', xt.Drift, length=16), line.new_element('mq.5', xt.Quadrupole, k1='k1l.q5', length='l.mq'), line.new_element('dd.5', xt.Drift, length=7.5), line.new_element('e.ss.r', xt.Marker), ]) -ss_left.build_tracker() +half_straight.build_tracker() tw_arc = arc.twiss4d() bet_ip = 300. -opt = ss_left.match( +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], @@ -276,18 +269,30 @@ def _replace_all_replicas(line): opt.step(40) +half_straight_left = half_straight.replicate('ss.l') +half_straight_left.mirror() +half_straight_right = half_straight.replicate('ss.r') +straight = line.new_section(components=[half_straight_left, half_straight_right]) - - -tw_arc = arc.twiss4d() -ss_arc = line.new_section(components=[ss_left, arc]) -ss_arc.cut_at_s(np.arange(0, ss_arc.get_length(), 0.5)) +ss_arc = line.new_section(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) tw_ss_arc.plot() -prrrr +line.discard_tracker() +line.append(straight) +line.append(arc1) +line.append(ss2) +line.append(arc2) +line.append(ss3) +line.append(arc3) + +tw_open = line.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) + +tw = line.twiss4d() import matplotlib.pyplot as plt plt.close('all') From 54c891deb5588521d4ac3ebc9408597c339c630a Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 09:12:16 +0200 Subject: [PATCH 034/159] With floor plot --- examples/lattice_design_shortcuts/000_dev.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 31d32d6f7..7afe2a7be 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -288,10 +288,12 @@ def _replace_all_replicas(line): line.append(ss3) line.append(arc3) -tw_open = line.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) +line.replace_all_replicas() +line.build_tracker() +sv = line.survey() +line.discard_tracker() +line.cut_at_s(np.arange(0, line.get_length(), 0.5)) tw = line.twiss4d() import matplotlib.pyplot as plt @@ -302,6 +304,9 @@ def _replace_all_replicas(line): ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) tw.plot('dx', ax=ax2) +import xplt +xplt.FloorPlot(sv, line, element_width=10) + plt.show() From 337b87eba690879c6e84d01986276a95425b0918 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 09:19:46 +0200 Subject: [PATCH 035/159] Longer straight --- examples/lattice_design_shortcuts/000_dev.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 7afe2a7be..9dbdbf277 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -234,20 +234,21 @@ def _replace_all_replicas(line): half_straight = line.new_section(components=[ line.new_element('ip', xt.Marker), - line.new_element('dd.0', xt.Drift, length=10), + line.new_element('dd.0', xt.Drift, length=20), line.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), line.new_element('dd.1', xt.Drift, length=6), line.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), - line.new_element('dd.2', xt.Drift, length=18), + line.new_element('dd.2', xt.Drift, length=11), line.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), - line.new_element('dd.3', xt.Drift, length=16), + line.new_element('dd.3', xt.Drift, length=18), line.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), - line.new_element('dd.4', xt.Drift, length=16), + line.new_element('dd.4', xt.Drift, length=18), line.new_element('mq.5', xt.Quadrupole, k1='k1l.q5', length='l.mq'), - line.new_element('dd.5', xt.Drift, length=7.5), + line.new_element('dd.5', xt.Drift, length=0.5), line.new_element('e.ss.r', xt.Marker), ]) half_straight.build_tracker() +print(f'Half straight length: {half_straight.get_length()}') tw_arc = arc.twiss4d() From 8c80ca3602efe57f2afe3d039f13714ea25aac45 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 09:21:12 +0200 Subject: [PATCH 036/159] Another solution --- examples/lattice_design_shortcuts/000_dev.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 9dbdbf277..5f92deb61 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -236,9 +236,9 @@ def _replace_all_replicas(line): line.new_element('ip', xt.Marker), line.new_element('dd.0', xt.Drift, length=20), line.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), - line.new_element('dd.1', xt.Drift, length=6), + line.new_element('dd.1', xt.Drift, length=5), line.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), - line.new_element('dd.2', xt.Drift, length=11), + line.new_element('dd.2', xt.Drift, length=12), line.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), line.new_element('dd.3', xt.Drift, length=18), line.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), From be8d2a9c4bcb276e0ab5d47ba07431ded62128ee Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 10:43:07 +0200 Subject: [PATCH 037/159] Keep in same buffer --- examples/lattice_design_shortcuts/000_dev.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 5f92deb61..46b682254 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -252,8 +252,6 @@ def _replace_all_replicas(line): tw_arc = arc.twiss4d() -bet_ip = 300. - opt = half_straight.match( solve=False, betx=tw_arc.betx[0], bety=tw_arc.bety[0], @@ -264,6 +262,7 @@ def _replace_all_replicas(line): 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)), ] ) @@ -293,10 +292,14 @@ def _replace_all_replicas(line): line.build_tracker() sv = line.survey() +buffer = line._buffer line.discard_tracker() line.cut_at_s(np.arange(0, line.get_length(), 0.5)) +line.build_tracker(_buffer=buffer) tw = line.twiss4d() +two = line.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) From ddc9ab2d8f8ec15ffcc077722943f18206d3c120 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 12:38:25 +0200 Subject: [PATCH 038/159] Going to different straight --- examples/lattice_design_shortcuts/000_dev.py | 57 ++++++++++++++++++-- 1 file changed, 52 insertions(+), 5 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 46b682254..0b21e611d 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -203,9 +203,9 @@ def _replace_all_replicas(line): 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') +# ss1 = ss.replicate(name='ss.1') +# ss2 = ss.replicate(name='ss.2') +# ss3 = ss.replicate(name='ss.3') @@ -278,14 +278,61 @@ def _replace_all_replicas(line): 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) + +line.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 = line.new_section(components=[ + line.new_element('ss.start', xt.Marker), + line.new_element('dd.ss.1.l', xt.Drift, length='l.mq'), + line.new_element('qfss.l', xt.Quadrupole, k1='kqfss.1', length='l.mq'), + + line.new_element('dd.ss.3.l', xt.Drift, length='3 *l.mb'), + + line.new_element('qdss.l', xt.Quadrupole, k1='kqdss.1', length='l.mq'), + line.new_element('dd.ss.5.l', xt.Drift, length='l.mq'), + + line.new_element('dd.ss.5.r', xt.Drift, length='l.mq'), + line.new_element('qdss.r', xt.Quadrupole, k1='kqdss.1', length='l.mq'), + + line.new_element('dd.ss.3.r', xt.Drift, length='3 *l.mb'), + + line.new_element('qfss.r', xt.Quadrupole, k1='kqfss.1', length='l.mq'), + line.new_element('dd.ss.1.r', xt.Drift, length='l.mq'), + +]) + +cell1_ss = cell_ss.replicate('cell.1') +cell2_ss = cell_ss.replicate('cell.2') +std_ss = line.new_section(components=[cell1_ss, cell2_ss]) + +ss1 = std_ss.replicate('ss.1') +ss2 = std_ss.replicate('ss.2') + +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() line.discard_tracker() line.append(straight) line.append(arc1) -line.append(ss2) +line.append(ss1) line.append(arc2) -line.append(ss3) +line.append(ss2) line.append(arc3) line.replace_all_replicas() From 4ddd60b105839fc41842ebc384a3d7d11e92450f Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 12:47:46 +0200 Subject: [PATCH 039/159] A clean state --- examples/lattice_design_shortcuts/000_dev.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 0b21e611d..d282ed539 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -307,13 +307,6 @@ def _replace_all_replicas(line): ]) -cell1_ss = cell_ss.replicate('cell.1') -cell2_ss = cell_ss.replicate('cell.2') -std_ss = line.new_section(components=[cell1_ss, cell2_ss]) - -ss1 = std_ss.replicate('ss.1') -ss2 = std_ss.replicate('ss.2') - opt = cell_ss.match( solve=False, method='4d', @@ -327,7 +320,16 @@ def _replace_all_replicas(line): tw_ss_arc.plot() + +cell1_ss = cell_ss.replicate('cell.1') +cell2_ss = cell_ss.replicate('cell.2') +std_ss = line.new_section(components=[cell1_ss, cell2_ss]) + +ss1 = std_ss.replicate('ss.1') +ss2 = std_ss.replicate('ss.2') line.discard_tracker() + + line.append(straight) line.append(arc1) line.append(ss1) From 0c1924b3bcbd04f0b01fbab89c20eed430251eab Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 13:06:33 +0200 Subject: [PATCH 040/159] Better cancellation of dispersion but funky insertion --- examples/lattice_design_shortcuts/000_dev.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index d282ed539..4e6608d42 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -215,8 +215,8 @@ def _replace_all_replicas(line): method='4d', vary=xt.VaryList(['k1l.qf', 'k1l.qd'], step=1e-5), targets=xt.TargetSet( - qx=0.333, - qy=0.333, + qx=0.333333, + qy=0.333333, )) line.vars({ From 2339551b819c23e634c441be6fde59e40b63178d Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 13:16:03 +0200 Subject: [PATCH 041/159] Nice... --- examples/lattice_design_shortcuts/000_dev.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 4e6608d42..ba02d5c14 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -263,6 +263,9 @@ def _replace_all_replicas(line): 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)), ] ) @@ -329,12 +332,11 @@ def _replace_all_replicas(line): ss2 = std_ss.replicate('ss.2') line.discard_tracker() - -line.append(straight) -line.append(arc1) line.append(ss1) -line.append(arc2) +line.append(arc1) line.append(ss2) +line.append(arc2) +line.append(straight) line.append(arc3) line.replace_all_replicas() @@ -351,11 +353,14 @@ def _replace_all_replicas(line): import matplotlib.pyplot as plt plt.close('all') -fig = plt.figure(1) +fig = plt.figure(1, figsize=(6.4*1.2, 4.8)) ax1 = fig.add_subplot(2, 1, 1) -tw.plot('betx bety', ax=ax1) +pltbet = tw.plot('betx bety', ax=ax1) ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) -tw.plot('dx', ax=ax2) +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, line, element_width=10) From e5f770e3f339c97e0dc830f37ecd8b4a1ccde76e Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 14:31:18 +0200 Subject: [PATCH 042/159] Moved first method --- examples/lattice_design_shortcuts/000_dev.py | 29 -------------------- xtrack/line.py | 29 ++++++++++++++++++++ 2 files changed, 29 insertions(+), 29 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index ba02d5c14..cd18f46b2 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -11,34 +11,6 @@ def __init__(self, expr): # line._xdeps_vref._eval("a * b") -def _new_element(line, name, cls, **kwargs): - - _eval = line._eval_obj.eval - - evaluated_kwargs = {} - value_kwargs = {} - for kk in kwargs: - if isinstance(kwargs[kk], Expr): - evaluated_kwargs[kk] = _eval(kwargs[kk].expr) - value_kwargs[kk] = evaluated_kwargs[kk]._value - elif hasattr(kwargs[kk], '_value'): - evaluated_kwargs[kk] = kwargs[kk] - value_kwargs[kk] = kwargs[kk]._value - elif (isinstance(kwargs[kk], str) and hasattr(cls, '_xofields') - and kk in cls._xofields and cls._xofields[kk].__name__ != 'String'): - evaluated_kwargs[kk] = _eval(kwargs[kk]) - value_kwargs[kk] = evaluated_kwargs[kk]._value - else: - evaluated_kwargs[kk] = kwargs[kk] - value_kwargs[kk] = kwargs[kk] - - element = cls(**value_kwargs) - line.element_dict[name] = element - for kk in kwargs: - setattr(line.element_refs[name], kk, evaluated_kwargs[kk]) - - return name - def _call_vars(vars, *args, **kwargs): _eval = vars.line._eval_obj.eval if len(args) > 0: @@ -125,7 +97,6 @@ def _replace_all_replicas(line): if isinstance(line[nn], xt.Replica): _replace_replica(line, nn) -xt.Line.new_element = _new_element xt.line.LineVars.__call__ = _call_vars xt.Line.new_section = _section xt.Line.append = _append diff --git a/xtrack/line.py b/xtrack/line.py index 228e3143c..99fd0b4ea 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3342,6 +3342,35 @@ def transform_compound(self, *args, **kwargs): 'https://xsuite.readthedocs.io/en/latest/line.html#apply-transformations-tilt-shift-to-elements' ) + def new_element(line, name, cls, **kwargs): + + _eval = line._eval_obj.eval + + assert cls in [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, + xt.Marker, xt.Replica], ( + 'Only Drift, Dipole, Quadrupole, Sextupole, Octupole, Marker, and Replica ' + 'elements are allowed in `new_element` for now.') + evaluated_kwargs = {} + value_kwargs = {} + for kk in kwargs: + if hasattr(kwargs[kk], '_value'): + evaluated_kwargs[kk] = kwargs[kk] + value_kwargs[kk] = kwargs[kk]._value + elif (isinstance(kwargs[kk], str) and hasattr(cls, '_xofields') + and kk in cls._xofields and cls._xofields[kk].__name__ != 'String'): + evaluated_kwargs[kk] = _eval(kwargs[kk]) + value_kwargs[kk] = evaluated_kwargs[kk]._value + else: + evaluated_kwargs[kk] = kwargs[kk] + value_kwargs[kk] = kwargs[kk] + + element = cls(**value_kwargs) + line.element_dict[name] = element + for kk in kwargs: + setattr(line.element_refs[name], kk, evaluated_kwargs[kk]) + + return name + def __len__(self): return len(self.element_names) From 0dafa372b17426421730f4a4008fb7a8cafa423f Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 14:35:16 +0200 Subject: [PATCH 043/159] Move more methods --- examples/lattice_design_shortcuts/000_dev.py | 18 --------------- xtrack/line.py | 24 ++++++++++++++++---- 2 files changed, 20 insertions(+), 22 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index cd18f46b2..987545aa3 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -81,27 +81,9 @@ def _section(line, components, name=None): def _append(line, section): line.element_names += section.components -def _replace_replica(line, name): - name_parent = line[name].resolve(line, get_name=True) - line.element_dict[name] = line[name_parent].copy() - - pars_with_expr = list( - line._xdeps_manager.tartasks[line.element_refs[name_parent]].keys()) - - for rr in pars_with_expr: - assert isinstance(rr, xd.refs.AttrRef) - setattr(line.element_refs[name], rr._key, rr._expr) - -def _replace_all_replicas(line): - for nn in line.element_names: - if isinstance(line[nn], xt.Replica): - _replace_replica(line, nn) - xt.line.LineVars.__call__ = _call_vars xt.Line.new_section = _section xt.Line.append = _append -xt.Line.replace_replica = _replace_replica -xt.Line.replace_all_replicas = _replace_all_replicas line = xt.Line() line.particle_ref = xt.Particles(p0c=2e9) diff --git a/xtrack/line.py b/xtrack/line.py index 99fd0b4ea..9f419283f 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3342,9 +3342,9 @@ def transform_compound(self, *args, **kwargs): 'https://xsuite.readthedocs.io/en/latest/line.html#apply-transformations-tilt-shift-to-elements' ) - def new_element(line, name, cls, **kwargs): + def new_element(self, name, cls, **kwargs): - _eval = line._eval_obj.eval + _eval = self._eval_obj.eval assert cls in [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, xt.Marker, xt.Replica], ( @@ -3365,12 +3365,28 @@ def new_element(line, name, cls, **kwargs): value_kwargs[kk] = kwargs[kk] element = cls(**value_kwargs) - line.element_dict[name] = element + self.element_dict[name] = element for kk in kwargs: - setattr(line.element_refs[name], kk, evaluated_kwargs[kk]) + setattr(self.element_refs[name], kk, evaluated_kwargs[kk]) return name + def replace_replica(self, name): + name_parent = self[name].resolve(self, get_name=True) + 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) + setattr(self.element_refs[name], 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 __len__(self): return len(self.element_names) From e47d919591a1054316f5a4a191885e93731408d3 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 14:44:26 +0200 Subject: [PATCH 044/159] Move more stuff --- examples/lattice_design_shortcuts/000_dev.py | 23 --------------- xtrack/line.py | 31 +++++++++++++++++++- 2 files changed, 30 insertions(+), 24 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 987545aa3..c4c700cb3 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -11,25 +11,6 @@ def __init__(self, expr): # line._xdeps_vref._eval("a * b") -def _call_vars(vars, *args, **kwargs): - _eval = vars.line._eval_obj.eval - if len(args) > 0: - assert len(kwargs) == 0 - assert len(args) == 1 - if isinstance(args[0], str): - return vars[args[0]] - elif isinstance(args[0], dict): - kwargs.update(args[0]) - else: - raise ValueError('Invalid argument') - for kk in kwargs: - if isinstance(kwargs[kk], Expr): - vars[kk] = _eval(kwargs[kk].expr) - elif isinstance(kwargs[kk], str): - vars[kk] = _eval(kwargs[kk]) - else: - vars[kk] = kwargs[kk] - def _flatten_components(components): flatten_components = [] for nn in components: @@ -81,16 +62,12 @@ def _section(line, components, name=None): def _append(line, section): line.element_names += section.components -xt.line.LineVars.__call__ = _call_vars xt.Line.new_section = _section xt.Line.append = _append line = xt.Line() line.particle_ref = xt.Particles(p0c=2e9) -line._eval_obj = xd.madxutils.MadxEval(variables=line._xdeps_vref, - functions=line._xdeps_fref, - elements=line.element_dict) n_bends_per_cell = 6 n_cells_par_arc = 3 diff --git a/xtrack/line.py b/xtrack/line.py index 9f419283f..9bba81789 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3344,7 +3344,7 @@ def transform_compound(self, *args, **kwargs): def new_element(self, name, cls, **kwargs): - _eval = self._eval_obj.eval + _eval = self._xdeps_eval.eval assert cls in [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, xt.Marker, xt.Replica], ( @@ -3533,6 +3533,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'): @@ -4612,6 +4624,23 @@ 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] + class ActionVars(Action): def __init__(self, line): From a6ce6b12c3df5d2d7fb0d12b1b116d340f51d25b Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 15:01:02 +0200 Subject: [PATCH 045/159] Clean up --- examples/lattice_design_shortcuts/000_dev.py | 65 -------------------- xtrack/line.py | 49 +++++++++++++++ 2 files changed, 49 insertions(+), 65 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index c4c700cb3..1b77d5381 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -1,74 +1,9 @@ import xtrack as xt -import xdeps as xd -import xobjects as xo - import numpy as np -class Expr: - def __init__(self, expr): - self.expr = expr - -# line._xdeps_vref._eval("a * b") - - -def _flatten_components(components): - flatten_components = [] - for nn in components: - if isinstance(nn, Section): - flatten_components += _flatten_components(nn.components) - else: - flatten_components.append(nn) - return flatten_components - -class Section(xt.Line): - def __init__(self, line, components, name=None): - self.line = line - xt.Line.__init__(self, elements=line.element_dict, - element_names=_flatten_components(components)) - self._element_dict = line.element_dict # Avoid copying - self._var_management = line._var_management - self._name = name - - @property - def name(self): - return self._name - - @property - def particle_ref(self): - return self.line.particle_ref - - @particle_ref.setter - def particle_ref(self, value): - assert value is None - - @property - def components(self): - return self.element_names - - def mirror(self): - self.element_names = self.element_names[::-1] - - def replicate(self, name): - new_components = [] - for nn in self.components: - new_nn = nn + '.' + name - self.line.element_dict[new_nn] = xt.Replica(nn) - new_components.append(new_nn) - return Section(self.line, new_components, name=name) - -def _section(line, components, name=None): - return Section(line, components, name=name) - -def _append(line, section): - line.element_names += section.components - -xt.Line.new_section = _section -xt.Line.append = _append - line = xt.Line() line.particle_ref = xt.Particles(p0c=2e9) - n_bends_per_cell = 6 n_cells_par_arc = 3 n_arcs = 3 diff --git a/xtrack/line.py b/xtrack/line.py index 9bba81789..3b94717f9 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3387,6 +3387,12 @@ def replace_all_replicas(self): if isinstance(self[nn], xt.Replica): self.replace_replica(nn) + def new_section(self, components, name=None): + return Section(self, components, name=name) + + def append(self, section): + self.element_names += section.components + def __len__(self): return len(self.element_names) @@ -5025,4 +5031,47 @@ def _rot_s_from_attr(attr): return rot_s_rad +def _flatten_components(components): + flatten_components = [] + for nn in components: + if isinstance(nn, Section): + flatten_components += _flatten_components(nn.components) + else: + flatten_components.append(nn) + return flatten_components + +class Section(Line): + def __init__(self, line, components, name=None): + self.line = line + xt.Line.__init__(self, elements=line.element_dict, + element_names=_flatten_components(components)) + self._element_dict = line.element_dict # Avoid copying + self._var_management = line._var_management + self._name = name + + @property + def name(self): + return self._name + + @property + def particle_ref(self): + return self.line.particle_ref + @particle_ref.setter + def particle_ref(self, value): + assert value is None + + @property + def components(self): + return self.element_names + + def mirror(self): + self.element_names = self.element_names[::-1] + + def replicate(self, name): + new_components = [] + for nn in self.components: + new_nn = nn + '.' + name + self.line.element_dict[new_nn] = xt.Replica(nn) + new_components.append(new_nn) + return Section(self.line, new_components, name=name) From f8b08854fea91a0edfd195ba0aed952f02c102ff Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 15:30:56 +0200 Subject: [PATCH 046/159] Adapt plot --- examples/lattice_design_shortcuts/000_dev.py | 3 +-- xtrack/line.py | 14 ++++++++++++-- xtrack/twissplot.py | 6 +++--- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 1b77d5381..7152f3df8 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -35,8 +35,7 @@ ]) hcell_left = halfcell.replicate(name='l') -hcell_right = halfcell.replicate(name='r') # could add mirror=True -hcell_right.mirror() +hcell_right = halfcell.replicate(name='r', mirror=True) cell = line.new_section(components=[ line.new_element('start', xt.Marker), diff --git a/xtrack/line.py b/xtrack/line.py index 3b94717f9..52fc2eb6d 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -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) @@ -5068,10 +5073,15 @@ def components(self): def mirror(self): self.element_names = self.element_names[::-1] - def replicate(self, name): + def replicate(self, name, mirror=False): new_components = [] for nn in self.components: new_nn = nn + '.' + name self.line.element_dict[new_nn] = xt.Replica(nn) new_components.append(new_nn) - return Section(self.line, new_components, name=name) + out = Section(self.line, new_components, name=name) + + if mirror: + out.mirror() + + return out diff --git a/xtrack/twissplot.py b/xtrack/twissplot.py index d9889d793..e21763e95 100644 --- a/xtrack/twissplot.py +++ b/xtrack/twissplot.py @@ -109,7 +109,7 @@ def __init__( if ax is not None: self.figure = ax.figure elif figure is None: - self.figure = plt.figure(num=figlabel) + 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: @@ -148,7 +148,7 @@ def __init__( def _new_axis(self, ax=None): if self.ax is None: out = self.figure.add_subplot(111) - self.figure.subplots_adjust(right=0.78) + self.figure.subplots_adjust(right=0.75) self.ax = out if self.used_ax: out = self.ax.twinx() @@ -222,7 +222,7 @@ def run(self): 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.3, 1.1) + self.lines, self.legends, loc="upper right", bbox_to_anchor=(1.35, 1.) ) self.ax.grid(True) self.figure.canvas.draw() From f21346d30352830a7459defa4895af47a4934677 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 5 Sep 2024 21:26:17 +0200 Subject: [PATCH 047/159] Fix --- xtrack/line.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xtrack/line.py b/xtrack/line.py index 52fc2eb6d..951cdea96 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -1323,6 +1323,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(): From 48da9d3d03ed88703b23cbc9b8050f1def35bba6 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 6 Sep 2024 11:56:14 +0200 Subject: [PATCH 048/159] Clean up --- xtrack/twiss.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 41dd4fb97..d29ca4a52 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -2774,15 +2774,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() From 2fa18eeb0dfbf839e3d76a6fa1da0b9c269c8a7b Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 6 Sep 2024 11:56:32 +0200 Subject: [PATCH 049/159] Clean --- xtrack/twiss.py | 1 - 1 file changed, 1 deletion(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index d29ca4a52..df194c63d 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -2821,7 +2821,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): From edf598e0981e96f8a3cfcf42ef116e84d18b1d74 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 6 Sep 2024 15:49:13 +0200 Subject: [PATCH 050/159] We have an environment! --- .../lattice_design_shortcuts/001_dev_env.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 examples/lattice_design_shortcuts/001_dev_env.py diff --git a/examples/lattice_design_shortcuts/001_dev_env.py b/examples/lattice_design_shortcuts/001_dev_env.py new file mode 100644 index 000000000..f329b67c8 --- /dev/null +++ b/examples/lattice_design_shortcuts/001_dev_env.py @@ -0,0 +1,64 @@ +import xtrack as xt +import numpy as np + +class Environment: + def __init__(self, element_dict=None, particle_ref=None): + self._element_dict = element_dict or {} + self.particle_ref = particle_ref + + self._init_var_management() + +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.new_element = xt.Line.new_element +Environment.new_section = xt.Line.new_section + +env = Environment(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', +}) + +halfcell = env.new_section(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_section(components=[ + env.new_element('start', xt.Marker), + hcell_left, + env.new_element('mid', xt.Marker), + hcell_right, + env.new_element('end', xt.Marker), +]) \ No newline at end of file From eb0cb9a5211418dfda3df49ee8ce899575b4622a Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 6 Sep 2024 16:16:08 +0200 Subject: [PATCH 051/159] Seems to start working --- .../lattice_design_shortcuts/001_dev_env.py | 25 ++++++- xtrack/line.py | 72 ++++++------------- 2 files changed, 43 insertions(+), 54 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_env.py b/examples/lattice_design_shortcuts/001_dev_env.py index f329b67c8..49259c5ec 100644 --- a/examples/lattice_design_shortcuts/001_dev_env.py +++ b/examples/lattice_design_shortcuts/001_dev_env.py @@ -1,6 +1,15 @@ import xtrack as xt import numpy as np +def _flatten_components(components): + flatten_components = [] + for nn in components: + if 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): self._element_dict = element_dict or {} @@ -8,6 +17,17 @@ def __init__(self, element_dict=None, particle_ref=None): self._init_var_management() + def new_line(self, components, name=None): + out = xt.Line() + out.particle_ref = self.particle_ref + out.line = self + out._element_dict = self.element_dict # Avoid copying + out.element_names = _flatten_components(components) + out._var_management = self._var_management + out._name = name + + return out + Environment.element_dict = xt.Line.element_dict Environment._init_var_management = xt.Line._init_var_management Environment._xdeps_vref = xt.Line._xdeps_vref @@ -19,7 +39,6 @@ def __init__(self, element_dict=None, particle_ref=None): Environment.varval = xt.Line.varval Environment.vv = xt.Line.vv Environment.new_element = xt.Line.new_element -Environment.new_section = xt.Line.new_section env = Environment(particle_ref=xt.Particles(p0c=2e9)) @@ -40,7 +59,7 @@ def __init__(self, element_dict=None, particle_ref=None): 'k0.mb': 'angle.mb / l.mb', }) -halfcell = env.new_section(components=[ +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'), @@ -55,7 +74,7 @@ def __init__(self, element_dict=None, particle_ref=None): hcell_left = halfcell.replicate(name='l') hcell_right = halfcell.replicate(name='r', mirror=True) -cell = env.new_section(components=[ +cell = env.new_line(components=[ env.new_element('start', xt.Marker), hcell_left, env.new_element('mid', xt.Marker), diff --git a/xtrack/line.py b/xtrack/line.py index 951cdea96..fccfcdeb0 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3379,6 +3379,26 @@ def new_element(self, name, cls, **kwargs): return name + def mirror(self): + self._frozen_check() + self.element_names = list(reversed(self.element_names)) + + def replicate(self, name, mirror=False): + new_element_names = [] + for nn in self.element_names: + new_nn = nn + '.' + name + self.line.element_dict[new_nn] = xt.Replica(nn) + new_element_names.append(new_nn) + out = Line() + out.element_names = new_element_names + out._element_dict = self.element_dict # to make sure that the dict is not copied + out._name = name + + if mirror: + out.mirror() + + return out + def replace_replica(self, name): name_parent = self[name].resolve(self, get_name=True) self.element_dict[name] = self[name_parent].copy() @@ -3442,7 +3462,7 @@ def name(self): if vv is self: return kk else: - return None + return getattr(self, '_name', None) @property def iscollective(self): @@ -5038,53 +5058,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 - -def _flatten_components(components): - flatten_components = [] - for nn in components: - if isinstance(nn, Section): - flatten_components += _flatten_components(nn.components) - else: - flatten_components.append(nn) - return flatten_components - -class Section(Line): - def __init__(self, line, components, name=None): - self.line = line - xt.Line.__init__(self, elements=line.element_dict, - element_names=_flatten_components(components)) - self._element_dict = line.element_dict # Avoid copying - self._var_management = line._var_management - self._name = name - - @property - def name(self): - return self._name - - @property - def particle_ref(self): - return self.line.particle_ref - - @particle_ref.setter - def particle_ref(self, value): - assert value is None - - @property - def components(self): - return self.element_names - - def mirror(self): - self.element_names = self.element_names[::-1] - - def replicate(self, name, mirror=False): - new_components = [] - for nn in self.components: - new_nn = nn + '.' + name - self.line.element_dict[new_nn] = xt.Replica(nn) - new_components.append(new_nn) - out = Section(self.line, new_components, name=name) - - if mirror: - out.mirror() - - return out From 7b7c73e2d16705b23e4b1c3fb26dddb61dc64991 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 6 Sep 2024 16:33:45 +0200 Subject: [PATCH 052/159] A few fixes --- xtrack/__init__.py | 2 +- xtrack/line.py | 50 +++++++++++++++++++++++++++++++++++++++------- 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/xtrack/__init__.py b/xtrack/__init__.py index 4adde00ec..ec482e8a4 100644 --- a/xtrack/__init__.py +++ b/xtrack/__init__.py @@ -12,7 +12,7 @@ from .beam_elements import * from .random import * from .tracker_data import TrackerData -from .line import Line, Node, freeze_longitudinal, _temp_knobs, EnergyProgram +from .line import Line, Node, freeze_longitudinal, _temp_knobs, EnergyProgram, Environment from .tracker import Tracker, Log from .match import (Vary, Target, TargetList, VaryList, TargetInequality, Action, TargetRelPhaseAdvance, TargetSet, GreaterThan, LessThan, diff --git a/xtrack/line.py b/xtrack/line.py index fccfcdeb0..b41c916b6 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3317,7 +3317,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 @@ -3415,11 +3415,8 @@ def replace_all_replicas(self): if isinstance(self[nn], xt.Replica): self.replace_replica(nn) - def new_section(self, components, name=None): - return Section(self, components, name=name) - - def append(self, section): - self.element_names += section.components + def extend(self, line): + self.element_names.extend(line.element_names) def __len__(self): return len(self.element_names) @@ -4629,7 +4626,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 @@ -5058,3 +5055,42 @@ 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 + +def _flatten_components(components): + flatten_components = [] + for nn in components: + if isinstance(nn, 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): + self._element_dict = element_dict or {} + self.particle_ref = particle_ref + + self._init_var_management() + + def new_line(self, components, name=None): + out = Line() + out.particle_ref = self.particle_ref + out.line = self + out._element_dict = self.element_dict # Avoid copying + out.element_names = _flatten_components(components) + out._var_management = self._var_management + out._name = name + + return out + +Environment.element_dict = Line.element_dict +Environment._init_var_management = Line._init_var_management +Environment._xdeps_vref = Line._xdeps_vref +Environment._xdeps_fref = Line._xdeps_fref +Environment._xdeps_manager = Line._xdeps_manager +Environment._xdeps_eval = Line._xdeps_eval +Environment.element_refs = Line.element_refs +Environment.vars = Line.vars +Environment.varval = Line.varval +Environment.vv = Line.vv +Environment.new_element = Line.new_element From 93bbea756ece03f1dc4d4a81d8ff047bfaa28d7a Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 6 Sep 2024 16:42:05 +0200 Subject: [PATCH 053/159] Adapt example and a few fixes --- examples/lattice_design_shortcuts/000_dev.py | 143 +++++++++---------- xtrack/line.py | 6 +- 2 files changed, 75 insertions(+), 74 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 7152f3df8..d3765974f 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -1,8 +1,8 @@ import xtrack as xt import numpy as np -line = xt.Line() -line.particle_ref = xt.Particles(p0c=2e9) +env = xt.Environment() +env.particle_ref = xt.Particles(p0c=2e9) n_bends_per_cell = 6 n_cells_par_arc = 3 @@ -11,7 +11,7 @@ n_bends = n_bends_per_cell * n_cells_par_arc * n_arcs -line.vars({ +env.vars({ 'k1l.qf': 0.027 / 2, 'k1l.qd': -0.0271 / 2, 'l.mq': 0.5, @@ -22,43 +22,43 @@ 'k0.mb': 'angle.mb / l.mb', }) -halfcell = line.new_section(components=[ - line.new_element('drift.1', xt.Drift, length='l.mq / 2'), - line.new_element('qf', xt.Quadrupole, k1='kqf.1', length='l.mq'), - line.new_element('drift.2', xt.Replica, parent_name='drift.1'), - line.new_element('mb.1', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb'), - line.new_element('mb.2', xt.Replica, parent_name='mb.1'), - line.new_element('mb.3', xt.Replica, parent_name='mb.1'), - line.new_element('drift.3', xt.Replica, parent_name='drift.1'), - line.new_element('qd', xt.Quadrupole, k1='kqd.1', length='l.mq'), - line.new_element('drift.4', xt.Replica, parent_name='drift.1'), +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 = line.new_section(components=[ - line.new_element('start', xt.Marker), +cell = env.new_line(components=[ + env.new_element('start', xt.Marker), hcell_left, - line.new_element('mid', xt.Marker), + env.new_element('mid', xt.Marker), hcell_right, - line.new_element('end', xt.Marker), + env.new_element('end', xt.Marker), ]) -arc = line.new_section(components=[ +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') -line.new_element('drift_ss', xt.Drift, length='l.mb') -for ii, nn in enumerate(cell_ss.components): +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.components[ii] = line.new_element( + cell_ss.element_names[ii] = env.new_element( f'drift.{ii}.ss', xt.Replica, parent_name='drift_ss') -ss = line.new_section(components=[ +ss = env.new_line(components=[ cell_ss.replicate('cell.1'), cell_ss.replicate('cell.2'), ]) @@ -73,8 +73,6 @@ -line.replace_all_replicas() - opt = cell.match( method='4d', vary=xt.VaryList(['k1l.qf', 'k1l.qd'], step=1e-5), @@ -83,7 +81,7 @@ qy=0.333333, )) -line.vars({ +env.vars({ 'k1l.q1': 0.012, 'k1l.q2': -0.012, 'k1l.q3': 0.012, @@ -96,20 +94,20 @@ 'k1.q5': 'k1l.q5 / l.mq', }) -half_straight = line.new_section(components=[ - line.new_element('ip', xt.Marker), - line.new_element('dd.0', xt.Drift, length=20), - line.new_element('mq.1', xt.Quadrupole, k1='k1l.q1', length='l.mq'), - line.new_element('dd.1', xt.Drift, length=5), - line.new_element('mq.2', xt.Quadrupole, k1='k1l.q2', length='l.mq'), - line.new_element('dd.2', xt.Drift, length=12), - line.new_element('mq.3', xt.Quadrupole, k1='k1l.q3', length='l.mq'), - line.new_element('dd.3', xt.Drift, length=18), - line.new_element('mq.4', xt.Quadrupole, k1='k1l.q4', length='l.mq'), - line.new_element('dd.4', xt.Drift, length=18), - line.new_element('mq.5', xt.Quadrupole, k1='k1l.q5', length='l.mq'), - line.new_element('dd.5', xt.Drift, length=0.5), - line.new_element('e.ss.r', xt.Marker), +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()}') @@ -139,14 +137,14 @@ half_straight_left = half_straight.replicate('ss.l') half_straight_left.mirror() half_straight_right = half_straight.replicate('ss.r') -straight = line.new_section(components=[half_straight_left, half_straight_right]) +straight = env.new_line(components=[half_straight_left, half_straight_right]) -ss_arc = line.new_section(components=[arc1, straight, arc2]) +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) -line.vars({ +env.vars({ 'k1l.qfss': 0.027 / 2, 'k1l.qdss': -0.0271 / 2, 'kqfss.1': 'k1l.qfss / l.mq', @@ -154,23 +152,23 @@ 'angle.mb': 2 * np.pi / n_bends, 'k0.mb': 'angle.mb / l.mb', }) -cell_ss = line.new_section(components=[ - line.new_element('ss.start', xt.Marker), - line.new_element('dd.ss.1.l', xt.Drift, length='l.mq'), - line.new_element('qfss.l', xt.Quadrupole, k1='kqfss.1', length='l.mq'), +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'), - line.new_element('dd.ss.3.l', xt.Drift, length='3 *l.mb'), + env.new_element('dd.ss.3.l', xt.Drift, length='3 *l.mb'), - line.new_element('qdss.l', xt.Quadrupole, k1='kqdss.1', length='l.mq'), - line.new_element('dd.ss.5.l', xt.Drift, length='l.mq'), + 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'), - line.new_element('dd.ss.5.r', xt.Drift, length='l.mq'), - line.new_element('qdss.r', xt.Quadrupole, k1='kqdss.1', 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'), - line.new_element('dd.ss.3.r', xt.Drift, length='3 *l.mb'), + env.new_element('dd.ss.3.r', xt.Drift, length='3 *l.mb'), - line.new_element('qfss.r', xt.Quadrupole, k1='kqfss.1', length='l.mq'), - line.new_element('dd.ss.1.r', xt.Drift, length='l.mq'), + 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'), ]) @@ -190,30 +188,31 @@ cell1_ss = cell_ss.replicate('cell.1') cell2_ss = cell_ss.replicate('cell.2') -std_ss = line.new_section(components=[cell1_ss, cell2_ss]) +std_ss = env.new_line(components=[cell1_ss, cell2_ss]) ss1 = std_ss.replicate('ss.1') ss2 = std_ss.replicate('ss.2') -line.discard_tracker() -line.append(ss1) -line.append(arc1) -line.append(ss2) -line.append(arc2) -line.append(straight) -line.append(arc3) +ring = env.new_line() + +ring.extend(ss1) +ring.extend(arc1) +ring.extend(ss2) +ring.extend(arc2) +ring.extend(straight) +ring.extend(arc3) -line.replace_all_replicas() -line.build_tracker() -sv = line.survey() +ring.replace_all_replicas() +ring.build_tracker() +sv = ring.survey() -buffer = line._buffer -line.discard_tracker() -line.cut_at_s(np.arange(0, line.get_length(), 0.5)) -line.build_tracker(_buffer=buffer) -tw = line.twiss4d() +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 = line.twiss(start=xt.START, betx=tw_arc.betx[-1], bety=tw_arc.bety[-1]) +two = ring.twiss(start=xt.START, betx=tw_arc.betx[-1], bety=tw_arc.bety[-1]) import matplotlib.pyplot as plt plt.close('all') @@ -227,7 +226,7 @@ pltdx.move_legend(1.2,1) import xplt -xplt.FloorPlot(sv, line, element_width=10) +xplt.FloorPlot(sv, ring, element_width=10) plt.show() diff --git a/xtrack/line.py b/xtrack/line.py index b41c916b6..13d32c443 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3387,7 +3387,7 @@ def replicate(self, name, mirror=False): new_element_names = [] for nn in self.element_names: new_nn = nn + '.' + name - self.line.element_dict[new_nn] = xt.Replica(nn) + self.element_dict[new_nn] = xt.Replica(nn) new_element_names.append(new_nn) out = Line() out.element_names = new_element_names @@ -5072,11 +5072,13 @@ def __init__(self, element_dict=None, particle_ref=None): self._init_var_management() - def new_line(self, components, name=None): + def new_line(self, components=None, name=None): out = Line() out.particle_ref = self.particle_ref out.line = self out._element_dict = self.element_dict # Avoid copying + if components is None: + components = [] out.element_names = _flatten_components(components) out._var_management = self._var_management out._name = name From 861ecd8abf9ce672ee036d3903d306820c37124f Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 17:04:14 +0200 Subject: [PATCH 054/159] add gzip support --- xtrack/json_utils.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 xtrack/json_utils.py 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 From 54d45dd9f2f669f0453d46da999565d7358dea9e Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 17:08:55 +0200 Subject: [PATCH 055/159] add gzip support --- xtrack/line.py | 18 +++++------------- xtrack/multiline/multiline.py | 28 +++++++++------------------- 2 files changed, 14 insertions(+), 32 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 1c8204600..c04a47f1b 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 @@ -235,6 +234,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 +245,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 +622,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 +635,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(file, self.to_dict(**kwargs), indent=indent) def _to_table_dict(self): diff --git a/xtrack/multiline/multiline.py b/xtrack/multiline/multiline.py index 9027249fc..c981a65aa 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(file, self.to_dict(**kwargs), 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 From 7bd2db1ee8c3a2281be40dc6de241f20e416a3c2 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 17:11:48 +0200 Subject: [PATCH 056/159] fix --- xtrack/line.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/line.py b/xtrack/line.py index c04a47f1b..54b5de0c8 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -19,7 +19,7 @@ from scipy.constants import c as clight from . import linear_normal_form as lnf -from .. import json_utils +from . import json_utils import xobjects as xo import xtrack as xt From 57117bb04e4fe742e2083c573d41b4bef24483c6 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 17:14:04 +0200 Subject: [PATCH 057/159] fix --- xtrack/line.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/line.py b/xtrack/line.py index 54b5de0c8..8984f5c88 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -635,7 +635,7 @@ def to_json(self, file, indent=1, **kwargs): ''' - json_utils.to_json(file, self.to_dict(**kwargs), indent=indent) + json_utils.to_json(self.to_dict(**kwargs), file, indent=indent) def _to_table_dict(self): From 0dc264900585c5b6cef3ebc0fd16b3866bf36597 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 17:21:33 +0200 Subject: [PATCH 058/159] adding tests --- tests/test_line.py | 48 ++++++++++++++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 17 deletions(-) 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) From 674159667c967fd49f6134d5815717f4a7f07cf2 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 17:23:24 +0200 Subject: [PATCH 059/159] fix --- xtrack/multiline/multiline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/multiline/multiline.py b/xtrack/multiline/multiline.py index c981a65aa..b7c4df137 100644 --- a/xtrack/multiline/multiline.py +++ b/xtrack/multiline/multiline.py @@ -138,7 +138,7 @@ def to_json(self, file, indent=1, **kwargs): **kwargs: dict Additional keyword arguments are passed to the `Line.to_dict` method. ''' - json_utils.to_json(file, self.to_dict(**kwargs), indent=indent) + json_utils.to_json(self.to_dict(**kwargs), file, indent=indent) @classmethod From dd929e40bcc3069f12cb76cb99cd821c0e3ff248 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 18:11:23 +0200 Subject: [PATCH 060/159] fix double strength --- xtrack/line.py | 2 +- xtrack/twiss.py | 46 ++++++++++++++++++++++++++-------------------- 2 files changed, 27 insertions(+), 21 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 8984f5c88..f7430cf15 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -737,7 +737,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] diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 7c184da22..3419bf201 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__) @@ -3263,7 +3265,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] @@ -3807,25 +3809,29 @@ 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 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 + + for kk in SIGN_FLIP_FOR_ATTR_REVERSE: + if kk not in out: + val=out[kk] + val*=1 def _W_phys2norm(x, px, y, py, zeta, pzeta, W_matrix, co_dict, nemitt_x=None, nemitt_y=None, nemitt_zeta=None): From eda527711bf6a05ca40f4496bb5c4a6a64c43a9e Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 18:13:12 +0200 Subject: [PATCH 061/159] fix --- xtrack/twiss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 3419bf201..f9866a0d0 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3831,7 +3831,7 @@ def _reverse_strengths(out): for kk in SIGN_FLIP_FOR_ATTR_REVERSE: if kk not in out: val=out[kk] - val*=1 + val*=-1 def _W_phys2norm(x, px, y, py, zeta, pzeta, W_matrix, co_dict, nemitt_x=None, nemitt_y=None, nemitt_zeta=None): From 0807dbd07da983762e6f397c1659138501c25701 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 18:18:27 +0200 Subject: [PATCH 062/159] fix --- xtrack/twiss.py | 22 ++-------------------- 1 file changed, 2 insertions(+), 20 deletions(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index f9866a0d0..2152aa1f5 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3810,27 +3810,9 @@ 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 - for kk in SIGN_FLIP_FOR_ATTR_REVERSE: - if kk not in out: - val=out[kk] + if kk in out: + val=out[kk] #avoid passing by setitem val*=-1 def _W_phys2norm(x, px, y, py, zeta, pzeta, W_matrix, co_dict, nemitt_x=None, nemitt_y=None, nemitt_zeta=None): From 425acd497e1877108d90798897b38f8226e3dc24 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 6 Sep 2024 18:34:00 +0200 Subject: [PATCH 063/159] improvement --- xtrack/twiss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 2152aa1f5..8b6c8a1e3 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3812,8 +3812,8 @@ def _reverse_strengths(out): ### 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 - val*=-1 + 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): From 634e7c63c490f1b2ab74e175add73eb7bed475d8 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 6 Sep 2024 22:03:22 +0200 Subject: [PATCH 064/159] Does not break with multipole --- examples/lattice_design_shortcuts/000_dev.py | 5 ++ xtrack/line.py | 48 +++++++++++++++----- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index d3765974f..b04348f51 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -151,11 +151,16 @@ 'kqdss.1': 'k1l.qdss / l.mq', 'angle.mb': 2 * np.pi / n_bends, 'k0.mb': 'angle.mb / l.mb', + 'k0l.corrector': 0., + 'k1sl.corrector': 0., + }) 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('corrector.l', xt.Multipole, knl=['k0l.corrector', 0], + ksl=[0, 'k1sl.corrector'], length='l.mq'), env.new_element('dd.ss.3.l', xt.Drift, length='3 *l.mb'), diff --git a/xtrack/line.py b/xtrack/line.py index 13d32c443..c5a2d3e19 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3355,27 +3355,49 @@ def new_element(self, name, cls, **kwargs): _eval = self._xdeps_eval.eval assert cls in [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, - xt.Marker, xt.Replica], ( - 'Only Drift, Dipole, Quadrupole, Sextupole, Octupole, Marker, and Replica ' + xt.Multipole, xt.Marker, xt.Replica], ( + 'Only Drift, Dipole, Quadrupole, Sextupole, Octupole, Multipole, Marker, and Replica ' 'elements are allowed in `new_element` for now.') - evaluated_kwargs = {} + ref_kwargs = {} value_kwargs = {} for kk in kwargs: if hasattr(kwargs[kk], '_value'): - evaluated_kwargs[kk] = kwargs[kk] + 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'): - evaluated_kwargs[kk] = _eval(kwargs[kk]) - value_kwargs[kk] = evaluated_kwargs[kk]._value + ref_kwargs[kk] = _eval(kwargs[kk]) + value_kwargs[kk] = ref_kwargs[kk]._value else: - evaluated_kwargs[kk] = kwargs[kk] value_kwargs[kk] = kwargs[kk] element = cls(**value_kwargs) self.element_dict[name] = element - for kk in kwargs: - setattr(self.element_refs[name], kk, evaluated_kwargs[kk]) + for kk in ref_kwargs: + if isinstance(ref_kwargs[kk], list): + for ii, vvv in enumerate(ref_kwargs[kk]): + if vvv is not None: + getattr(self.element_refs[name], kk)[ii] = vvv + else: + setattr(self.element_refs[name], kk, ref_kwargs[kk]) return name @@ -3407,8 +3429,12 @@ def replace_replica(self, name): self._xdeps_manager.tartasks[self.element_refs[name_parent]].keys()) for rr in pars_with_expr: - assert isinstance(rr, xd.refs.AttrRef) - setattr(self.element_refs[name], rr._key, rr._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): + rr._owner[rr._key] = rr._expr def replace_all_replicas(self): for nn in self.element_names: From c3b9e06372b723241744943bdd74265ef7ab6e6d Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 6 Sep 2024 22:09:42 +0200 Subject: [PATCH 065/159] Better --- xtrack/line.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/line.py b/xtrack/line.py index c5a2d3e19..3361bdfc9 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3434,7 +3434,7 @@ def replace_replica(self, name): if isinstance(rr, xd.refs.AttrRef): setattr(self.element_refs[name], rr._key, rr._expr) elif isinstance(rr, xd.refs.ItemRef): - rr._owner[rr._key] = rr._expr + getattr(self.element_refs[name], rr._owner._key)[rr._key] = rr._expr def replace_all_replicas(self): for nn in self.element_names: From a4bced7f94e01bb7627e54fbcaf86b3d0e03c45e Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 7 Sep 2024 07:47:11 +0200 Subject: [PATCH 066/159] Correctors are working --- examples/lattice_design_shortcuts/000_dev.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index b04348f51..3b402d0ce 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -20,6 +20,8 @@ 'l.mb': 12, 'angle.mb': 2 * np.pi / n_bends, 'k0.mb': 'angle.mb / l.mb', + 'k0l.corrector': 0, + 'k1sl.corrector': 0, }) halfcell = env.new_line(components=[ @@ -39,8 +41,12 @@ 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), ]) @@ -151,16 +157,11 @@ 'kqdss.1': 'k1l.qdss / l.mq', 'angle.mb': 2 * np.pi / n_bends, 'k0.mb': 'angle.mb / l.mb', - 'k0l.corrector': 0., - 'k1sl.corrector': 0., - }) 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('corrector.l', xt.Multipole, knl=['k0l.corrector', 0], - ksl=[0, 'k1sl.corrector'], length='l.mq'), env.new_element('dd.ss.3.l', xt.Drift, length='3 *l.mb'), From 7f55f9bda9d121d98fc35e156d06ece9995f595f Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 7 Sep 2024 10:21:22 +0200 Subject: [PATCH 067/159] Ensure tracker consistency --- .../lattice_design_shortcuts/001_dev_env.py | 83 ------------------- xtrack/line.py | 16 +++- 2 files changed, 14 insertions(+), 85 deletions(-) delete mode 100644 examples/lattice_design_shortcuts/001_dev_env.py diff --git a/examples/lattice_design_shortcuts/001_dev_env.py b/examples/lattice_design_shortcuts/001_dev_env.py deleted file mode 100644 index 49259c5ec..000000000 --- a/examples/lattice_design_shortcuts/001_dev_env.py +++ /dev/null @@ -1,83 +0,0 @@ -import xtrack as xt -import numpy as np - -def _flatten_components(components): - flatten_components = [] - for nn in components: - if 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): - self._element_dict = element_dict or {} - self.particle_ref = particle_ref - - self._init_var_management() - - def new_line(self, components, name=None): - out = xt.Line() - out.particle_ref = self.particle_ref - out.line = self - out._element_dict = self.element_dict # Avoid copying - out.element_names = _flatten_components(components) - out._var_management = self._var_management - out._name = name - - return out - -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.new_element = xt.Line.new_element - -env = Environment(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', -}) - -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), - hcell_left, - env.new_element('mid', xt.Marker), - hcell_right, - env.new_element('end', xt.Marker), -]) \ No newline at end of file diff --git a/xtrack/line.py b/xtrack/line.py index 3361bdfc9..fb76dc62b 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -10,6 +10,7 @@ import uuid import os from collections import defaultdict +from weakref import WeakSet from contextlib import contextmanager from copy import deepcopy @@ -841,12 +842,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): @@ -5097,20 +5102,27 @@ def __init__(self, element_dict=None, particle_ref=None): self.particle_ref = particle_ref self._init_var_management() + self._lines = WeakSet() def new_line(self, components=None, name=None): out = Line() out.particle_ref = self.particle_ref - out.line = self + out.env = self out._element_dict = self.element_dict # Avoid copying if components is None: components = [] out.element_names = _flatten_components(components) out._var_management = self._var_management out._name = name + self._lines.add(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() + Environment.element_dict = Line.element_dict Environment._init_var_management = Line._init_var_management Environment._xdeps_vref = Line._xdeps_vref From a72ccd0ab387269a3e4846dc67f71c5e639260db Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 7 Sep 2024 10:23:56 +0200 Subject: [PATCH 068/159] White list in replace_replica --- xtrack/line.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/xtrack/line.py b/xtrack/line.py index fb76dc62b..4e12c84fd 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3428,6 +3428,11 @@ def replicate(self, name, mirror=False): 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( From 20b6abc44b810add92b54b396507d491d46bbf62 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 8 Sep 2024 20:07:07 +0200 Subject: [PATCH 069/159] Add plot method to survey --- examples/lattice_design_shortcuts/000_dev.py | 1 - xtrack/survey.py | 6 ++++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 3b402d0ce..f94c10b00 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -10,7 +10,6 @@ n_bends = n_bends_per_cell * n_cells_par_arc * n_arcs - env.vars({ 'k1l.qf': 0.027 / 2, 'k1l.qd': -0.0271 / 2, diff --git a/xtrack/survey.py b/xtrack/survey.py index bd7de99d1..02891fe46 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -167,6 +167,11 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, return out + def plot(self, **kwargs): + import xplt + xplt.FloorPlot(self, self.line, **kwargs) + + # ================================================== # Main function @@ -224,6 +229,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 From fcba1bd3e53c63ade1445afdbaf12c4723cf97a0 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 8 Sep 2024 20:30:47 +0200 Subject: [PATCH 070/159] Plot in table --- xtrack/survey.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/xtrack/survey.py b/xtrack/survey.py index 02891fe46..3fe8ad316 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -167,9 +167,14 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, return out - def plot(self, **kwargs): + def plot(self, element_width=None, **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, **kwargs) + xplt.FloorPlot(self, self.line, element_width=element_width, **kwargs) # ================================================== From 4a1480c06d9d7308af0eb6ca36a3f50eba4b73ae Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 8 Sep 2024 20:35:13 +0200 Subject: [PATCH 071/159] Legend in survey plot --- xtrack/survey.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/xtrack/survey.py b/xtrack/survey.py index 3fe8ad316..bf487d765 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -167,7 +167,7 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, return out - def plot(self, element_width=None, **kwargs): + 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) @@ -175,6 +175,9 @@ def plot(self, element_width=None, **kwargs): 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() # ================================================== From 344c198e6d90f49b5f6a2dc2ce9fe33323d1ff44 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 8 Sep 2024 21:08:56 +0200 Subject: [PATCH 072/159] Implement Line.get_setion --- xtrack/line.py | 24 ++++++++++++++++++++++-- xtrack/tracker.py | 4 ---- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 4e12c84fd..3cb286cf4 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3451,6 +3451,22 @@ def replace_all_replicas(self): if isinstance(self[nn], xt.Replica): self.replace_replica(nn) + def get_section(self, start=None, end=None, name=None): + + tt = self.get_table().rows[start:end] + if tt.name[-1] == '_end_point': + tt = tt.rows[:-1] + 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) + + out = self.env.new_line(components=list(tt.name), name=name) + + return out + + def extend(self, line): self.element_names.extend(line.element_names) @@ -5102,11 +5118,15 @@ def _flatten_components(components): return flatten_components class Environment: - def __init__(self, element_dict=None, particle_ref=None): + def __init__(self, element_dict=None, particle_ref=None, _var_management=None): self._element_dict = element_dict or {} self.particle_ref = particle_ref - self._init_var_management() + if _var_management is not None: + self._var_management = _var_management + else: + self._init_var_management() + self._lines = WeakSet() def new_line(self, components=None, name=None): diff --git a/xtrack/tracker.py b/xtrack/tracker.py index 558681143..982d92285 100644 --- a/xtrack/tracker.py +++ b/xtrack/tracker.py @@ -26,10 +26,6 @@ logger = logging.getLogger(__name__) - - - - class Tracker: ''' From 164b9128410ed6ba9863140aaad724bab7feae29 Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 09:10:27 +0200 Subject: [PATCH 073/159] A first resolution --- .../lattice_design_shortcuts/001_dev_at_s.py | 36 +++++++++++++++++++ xtrack/line.py | 8 ++++- 2 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 examples/lattice_design_shortcuts/001_dev_at_s.py 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..6c38521d4 --- /dev/null +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -0,0 +1,36 @@ +import xtrack as xt + +def _plot_line(line): + tt = line.get_table(attr=True) + xt.twiss.TwissTable.plot(tt, yl='', yr='') + +class SPlace: + def __init__(self, name, s, from_=None): + self.name = name + self.s = s + self.from_ = from_ + +env = xt.Environment() + +seq = [ + SPlace(env.new_element('ip', xt.Marker), s=10), + SPlace(env.new_element('left', xt.Marker), s=-5, from_='ip'), + SPlace(env.new_element('right', xt.Marker),s=+5, from_='ip'), +] + +s_dct = {} +n_resolved = 0 +n_resolved_prev = -1 +while n_resolved != n_resolved_prev: + for ss in seq: + if ss.from_ is None: + s_dct[ss.name] = ss.s + n_resolved += 1 + else: + if ss.from_ in s_dct: + s_dct[ss.name] = s_dct[ss.from_] + ss.s + n_resolved += 1 + n_resolved_prev = n_resolved + + + diff --git a/xtrack/line.py b/xtrack/line.py index 3cb286cf4..993ce6e01 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3451,7 +3451,13 @@ def replace_all_replicas(self): if isinstance(self[nn], xt.Replica): self.replace_replica(nn) - def get_section(self, start=None, end=None, name=None): + 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': From fbbf634b844cc6c0cfb05bd13bebc517de227d9e Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 13:30:02 +0200 Subject: [PATCH 074/159] Mock up --- examples/lattice_design_shortcuts/000_dev.py | 26 ++++++++ .../lattice_design_shortcuts/001_dev_at_s.py | 9 +-- examples/lattice_design_shortcuts/002_play.py | 63 +++++++++++++++++++ 3 files changed, 94 insertions(+), 4 deletions(-) create mode 100644 examples/lattice_design_shortcuts/002_play.py diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index f94c10b00..f5c44188d 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -23,6 +23,31 @@ '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'), @@ -50,6 +75,7 @@ env.new_element('end', xt.Marker), ]) + arc = env.new_line(components=[ cell.replicate(name='cell.1'), cell.replicate(name='cell.2'), diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 6c38521d4..ccf499bfb 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -4,7 +4,7 @@ def _plot_line(line): tt = line.get_table(attr=True) xt.twiss.TwissTable.plot(tt, yl='', yr='') -class SPlace: +class Place: def __init__(self, name, s, from_=None): self.name = name self.s = s @@ -13,9 +13,10 @@ def __init__(self, name, s, from_=None): env = xt.Environment() seq = [ - SPlace(env.new_element('ip', xt.Marker), s=10), - SPlace(env.new_element('left', xt.Marker), s=-5, from_='ip'), - SPlace(env.new_element('right', xt.Marker),s=+5, from_='ip'), + Place(env.new_element('ip', xt.Marker), s=10), + Place(env.new_element('left', xt.Marker), s=-5, from_='ip'), + env.new_element('after_left', xt.Marker), + Place(env.new_element('right', xt.Marker),s=+5, from_='ip'), ] s_dct = {} diff --git a/examples/lattice_design_shortcuts/002_play.py b/examples/lattice_design_shortcuts/002_play.py new file mode 100644 index 000000000..6d07faff5 --- /dev/null +++ b/examples/lattice_design_shortcuts/002_play.py @@ -0,0 +1,63 @@ +import xtrack as xt +env = xt.Environment() + + +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', + SPlace('qf', s + 'drift.2', + 'mb.1', + 'mb.2', + 'mb.3', + 'drift.3', + 'qd', + 'drift.4', +]) + +mbx = env.new_element('mbxw', xt.Bend, k0='k0.mb', h=0, length='l.mbxw') + +d1 = env.new_line('d1', components=[ + env.new_element('lmbxw.start', parent=xt.Marker), # shortcut env.new_element('lmbxw.start') + env.new_element('mbxw.a4@start', parent=xt.Replica, parent_name='mbxw', at=0.5), + env.new_element('mbxw.b4@start', parent='mbxw', _from='mbxw.a4@end'), + env.new_element('lmbxw.end', parent=xt.Marker, at=0.5, _from='mbxw.b4@end'), +]) + +d2 = env.new_element('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_element('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_element('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'), + + +]) \ No newline at end of file From 10d46c4ab79e6b21b023d2b0afa14641424793ca Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 16:02:37 +0200 Subject: [PATCH 075/159] Handle subsequence --- .../lattice_design_shortcuts/001_dev_at_s.py | 93 ++++++++++++++++--- examples/lattice_design_shortcuts/002_play.py | 18 +++- 2 files changed, 96 insertions(+), 15 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index ccf499bfb..03aaa8dda 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -1,37 +1,106 @@ import xtrack as xt +import numpy as np def _plot_line(line): tt = line.get_table(attr=True) xt.twiss.TwissTable.plot(tt, yl='', yr='') class Place: - def __init__(self, name, s, from_=None): + 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.s = s + self.at = at self.from_ = from_ + self.anchor = anchor + self.from_anchor = from_anchor env = xt.Environment() seq = [ - Place(env.new_element('ip', xt.Marker), s=10), - Place(env.new_element('left', xt.Marker), s=-5, from_='ip'), + Place(env.new_element('ip', xt.Marker), at=10), + # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + ( + env.new_element('before_before_right', xt.Marker), + env.new_element('before_right', xt.Quadrupole, length=1), + Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + env.new_element('after_right', xt.Marker), + env.new_element('after_right2', xt.Marker), + ), + Place(env.new_element('before_left', xt.Marker), at='__before__'), + Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), env.new_element('after_left', xt.Marker), - Place(env.new_element('right', xt.Marker),s=+5, from_='ip'), + env.new_element('after_left2', xt.Marker), ] -s_dct = {} +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].at = '__before__' + seq_all_places.extend(ss_aux) + else: + seq_all_places.append(Place(ss, at=None, from_=None)) + return seq_all_places + +seq_all_places = _all_places(seq) + + +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 while n_resolved != n_resolved_prev: - for ss in seq: - if ss.from_ is None: - s_dct[ss.name] = ss.s + 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 or ss.at == '__after__': + ss_prev = seq_all_places[ii-1] + if ss_prev.name in s_center_dct: + 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 == '__before__': + ss_next = seq_all_places[ii+1] + if ss_next.name in s_center_dct: + 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 + elif ss.from_ is None: + s_center_dct[ss.name] = ss.at n_resolved += 1 else: - if ss.from_ in s_dct: - s_dct[ss.name] = s_dct[ss.from_] + ss.s + if ss.from_ in s_center_dct: + s_center_dct[ss.name] = s_center_dct[ss.from_] + ss.at n_resolved += 1 - n_resolved_prev = n_resolved +assert n_resolved == len(seq_all_places) + +aux_s = np.array([s_center_dct[nn] for nn in aux_tt.name[:-1]]) +i_sorted = np.argsort(aux_s, stable=True) +name_sorted = [str(aux_tt.name[ii]) for ii in i_sorted] \ No newline at end of file diff --git a/examples/lattice_design_shortcuts/002_play.py b/examples/lattice_design_shortcuts/002_play.py index 6d07faff5..5919261c1 100644 --- a/examples/lattice_design_shortcuts/002_play.py +++ b/examples/lattice_design_shortcuts/002_play.py @@ -29,8 +29,8 @@ d1 = env.new_line('d1', components=[ env.new_element('lmbxw.start', parent=xt.Marker), # shortcut env.new_element('lmbxw.start') env.new_element('mbxw.a4@start', parent=xt.Replica, parent_name='mbxw', at=0.5), - env.new_element('mbxw.b4@start', parent='mbxw', _from='mbxw.a4@end'), - env.new_element('lmbxw.end', parent=xt.Marker, at=0.5, _from='mbxw.b4@end'), + env.new_element('mbxw.b4@start', parent='mbxw', from_='mbxw.a4@end'), + env.new_element('lmbxw.end', parent=xt.Marker, at=0.5, from_='mbxw.b4@end'), ]) d2 = env.new_element('d2.b1', xt.Bend, k0='k0.mb', h=0, length='l.mbxw', dx=0.188/2) @@ -60,4 +60,16 @@ place('ir2', ir.replicate('2'), at=s_ip2, patch=True, reference='ip.2'), -]) \ No newline at end of file +]) + +seq = [ + Place(env.new_element('ip', xt.Marker), at=10), + Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), + env.new_element('after_left', xt.Marker), + env.new_element('after_left2', xt.Marker), + Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + Place([env.new_element('before_right', xt.Quadrupole, length=1), + env.new_element('before_right2', xt.Marker)], at=0, from_='right@start', + anchor='before_right@center'), + Place(env.new_element('righter', xt.Quadrupole, length=1), at=+5, from_='before_right'), +] \ No newline at end of file From 6920e7238679403a645c737bc4a1861960aff9ea Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 16:16:52 +0200 Subject: [PATCH 076/159] Hide before --- .../lattice_design_shortcuts/001_dev_at_s.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 03aaa8dda..558c8d176 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -18,6 +18,7 @@ def __init__(self, name, at=None, from_=None, anchor=None, from_anchor=None): self.from_ = from_ self.anchor = anchor self.from_anchor = from_anchor + self._before = False env = xt.Environment() @@ -31,12 +32,25 @@ def __init__(self, name, at=None, from_=None, anchor=None, from_anchor=None): env.new_element('after_right', xt.Marker), env.new_element('after_right2', xt.Marker), ), - Place(env.new_element('before_left', xt.Marker), at='__before__'), Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), env.new_element('after_left', xt.Marker), env.new_element('after_left2', xt.Marker), ] +# seq = [ +# Place(env.new_element('ip', xt.Marker), at=10), +# # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), +# env.new_element('before_before_right', xt.Marker, at='__before__'), +# env.new_element('before_right', xt.Quadrupole, length=1, at='__before__'), +# Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), +# env.new_element('after_right', xt.Marker), +# env.new_element('after_right2', xt.Marker), +# Place(env.new_element('before_left', xt.Marker), at='__before__'), +# Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), +# env.new_element('after_left', xt.Marker), +# env.new_element('after_left2', xt.Marker), +# ] + def _all_places(seq): seq_all_places = [] for ss in seq: @@ -53,7 +67,7 @@ def _all_places(seq): raise ValueError('No Place in sequence') ss_aux = _all_places(ss) for ii in range(i_first): - ss_aux[ii].at = '__before__' + ss_aux[ii]._before = True seq_all_places.extend(ss_aux) else: seq_all_places.append(Place(ss, at=None, from_=None)) @@ -75,14 +89,14 @@ def _all_places(seq): for ii, ss in enumerate(seq_all_places): if ss.name in s_center_dct: continue - if ss.at is None or ss.at == '__after__': + if ss.at is None: ss_prev = seq_all_places[ii-1] if ss_prev.name in s_center_dct: 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 == '__before__': + elif ss._before: ss_next = seq_all_places[ii+1] if ss_next.name in s_center_dct: s_center_dct[ss.name] = (s_center_dct[ss_next.name] From a68259117d86bfcad20bbdba9fb176e0382f85de Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 16:55:09 +0200 Subject: [PATCH 077/159] Generated tt_sorted --- .../lattice_design_shortcuts/001_dev_at_s.py | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 558c8d176..e65d31a22 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -37,6 +37,21 @@ def __init__(self, name, at=None, from_=None, anchor=None, from_anchor=None): env.new_element('after_left2', xt.Marker), ] +seq = [ + Place(env.new_element('ip', xt.Marker), at=10), + # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + ( + env.new_element('before_before_right', xt.Marker), + env.new_element('before_right', xt.Quadrupole, length=1), + Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + env.new_element('after_right', xt.Marker), + env.new_element('after_right2', xt.Marker), + ), + Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), + env.new_element('after_left', xt.Marker), + env.new_element('after_left2', xt.Marker), +] + # seq = [ # Place(env.new_element('ip', xt.Marker), at=10), # # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), @@ -117,4 +132,8 @@ def _all_places(seq): i_sorted = np.argsort(aux_s, stable=True) -name_sorted = [str(aux_tt.name[ii]) for ii in i_sorted] \ No newline at end of file +name_sorted = [str(aux_tt.name[ii]) for ii in i_sorted] + +tt_sorted = aux_tt.rows[name_sorted] + +assert np.all(tt_sorted.name == np.array(name_sorted)) \ No newline at end of file From 7b2554d11b3d01c6360575ff2b9040e897749265 Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 17:10:54 +0200 Subject: [PATCH 078/159] Compute magnet lengths --- examples/lattice_design_shortcuts/001_dev_at_s.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index e65d31a22..e14f1b07d 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -128,12 +128,19 @@ def _all_places(seq): assert n_resolved == len(seq_all_places) -aux_s = np.array([s_center_dct[nn] for nn in aux_tt.name[:-1]]) +aux_s_center = np.array([s_center_dct[nn] for nn in aux_tt.name[:-1]]) +aux_tt['s_center'] = np.concatenate([aux_s_center, [0]]) -i_sorted = np.argsort(aux_s, stable=True) +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)) \ No newline at end of file From 722381cdb6ad37faa0a396cb432b62a69ac91d58 Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 17:33:28 +0200 Subject: [PATCH 079/159] Layout seems ok --- .../lattice_design_shortcuts/001_dev_at_s.py | 56 ++++++++----------- xtrack/line.py | 9 +++ 2 files changed, 31 insertions(+), 34 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index e14f1b07d..8d5abef85 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -20,52 +20,26 @@ def __init__(self, name, at=None, from_=None, anchor=None, from_anchor=None): self.from_anchor = from_anchor self._before = False -env = xt.Environment() + def __repr__(self): + return f'Place({self.name}, at={self.at}, from_={self.from_})' -seq = [ - Place(env.new_element('ip', xt.Marker), at=10), - # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), - ( - env.new_element('before_before_right', xt.Marker), - env.new_element('before_right', xt.Quadrupole, length=1), - Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), - env.new_element('after_right', xt.Marker), - env.new_element('after_right2', xt.Marker), - ), - Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), - env.new_element('after_left', xt.Marker), - env.new_element('after_left2', xt.Marker), -] +env = xt.Environment() seq = [ Place(env.new_element('ip', xt.Marker), at=10), # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), ( env.new_element('before_before_right', xt.Marker), - env.new_element('before_right', xt.Quadrupole, length=1), + env.new_element('before_right', xt.Sextupole, length=1), Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), env.new_element('after_right', xt.Marker), env.new_element('after_right2', xt.Marker), ), Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), env.new_element('after_left', xt.Marker), - env.new_element('after_left2', xt.Marker), + env.new_element('after_left2', xt.Bend, length=0.5), ] -# seq = [ -# Place(env.new_element('ip', xt.Marker), at=10), -# # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), -# env.new_element('before_before_right', xt.Marker, at='__before__'), -# env.new_element('before_right', xt.Quadrupole, length=1, at='__before__'), -# Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), -# env.new_element('after_right', xt.Marker), -# env.new_element('after_right2', xt.Marker), -# Place(env.new_element('before_left', xt.Marker), at='__before__'), -# Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), -# env.new_element('after_left', xt.Marker), -# env.new_element('after_left2', xt.Marker), -# ] - def _all_places(seq): seq_all_places = [] for ss in seq: @@ -104,14 +78,14 @@ def _all_places(seq): for ii, ss in enumerate(seq_all_places): if ss.name in s_center_dct: continue - if ss.at is None: + if ss.at is None and not ss._before: ss_prev = seq_all_places[ii-1] if ss_prev.name in s_center_dct: 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._before: + elif ss.at is None and ss._before: ss_next = seq_all_places[ii+1] if ss_next.name in s_center_dct: s_center_dct[ss.name] = (s_center_dct[ss_next.name] @@ -142,5 +116,19 @@ def _all_places(seq): 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)) + +s_tol = 1e-12 +names_with_drifts = [] +# Create drifts +for nn in name_sorted: + 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() + drift = env.new_element(drift_name, xt.Drift, length=ds_upstream) + names_with_drifts.append(drift_name) + names_with_drifts.append(nn) + +line = env.new_line(components=names_with_drifts) -assert np.all(tt_sorted.name == np.array(name_sorted)) \ No newline at end of file diff --git a/xtrack/line.py b/xtrack/line.py index 993ce6e01..9f45aefb6 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -5134,6 +5134,7 @@ def __init__(self, element_dict=None, particle_ref=None, _var_management=None): self._init_var_management() self._lines = WeakSet() + self._drift_counter = 0 def new_line(self, components=None, name=None): out = Line() @@ -5154,6 +5155,14 @@ def _ensure_tracker_consistency(self, buffer): 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() + Environment.element_dict = Line.element_dict Environment._init_var_management = Line._init_var_management Environment._xdeps_vref = Line._xdeps_vref From ef9e7c94e0dbfb7b242b9b0b130a7d692b46bf62 Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 20:55:51 +0200 Subject: [PATCH 080/159] Handle element placed at the beginning of the line --- examples/lattice_design_shortcuts/001_dev_at_s.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 8d5abef85..224d6b673 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -26,6 +26,8 @@ def __repr__(self): env = xt.Environment() seq = [ + env.new_element('b1', xt.Bend, length=0.5), + env.new_element('q1', xt.Quadrupole, length=0.5), Place(env.new_element('ip', xt.Marker), at=10), # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), ( @@ -73,6 +75,11 @@ def _all_places(seq): 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: + s_center_dct[seq_all_places[0].name] = aux_tt['length', seq_all_places[0].name] / 2 + n_resolved += 1 + while n_resolved != n_resolved_prev: n_resolved_prev = n_resolved for ii, ss in enumerate(seq_all_places): From 3048b479a5f1ca526de5d35352f5a86d5b7207ce Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 21:08:00 +0200 Subject: [PATCH 081/159] Organize in functions --- .../lattice_design_shortcuts/001_dev_at_s.py | 194 ++++++++++-------- 1 file changed, 107 insertions(+), 87 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 224d6b673..f8e382eeb 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -23,24 +23,7 @@ def __init__(self, name, at=None, from_=None, anchor=None, from_anchor=None): def __repr__(self): return f'Place({self.name}, at={self.at}, from_={self.from_})' -env = xt.Environment() -seq = [ - env.new_element('b1', xt.Bend, length=0.5), - env.new_element('q1', xt.Quadrupole, length=0.5), - Place(env.new_element('ip', xt.Marker), at=10), - # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), - ( - env.new_element('before_before_right', xt.Marker), - env.new_element('before_right', xt.Sextupole, length=1), - Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), - env.new_element('after_right', xt.Marker), - env.new_element('after_right2', xt.Marker), - ), - Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), - env.new_element('after_left', xt.Marker), - env.new_element('after_left2', xt.Bend, length=0.5), -] def _all_places(seq): seq_all_places = [] @@ -64,78 +47,115 @@ def _all_places(seq): seq_all_places.append(Place(ss, at=None, from_=None)) return seq_all_places -seq_all_places = _all_places(seq) -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: - s_center_dct[seq_all_places[0].name] = aux_tt['length', seq_all_places[0].name] / 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: - 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: - 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 - elif ss.from_ is None: - s_center_dct[ss.name] = ss.at - n_resolved += 1 - else: - if ss.from_ in s_center_dct: - s_center_dct[ss.name] = s_center_dct[ss.from_] + ss.at + +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: + s_center_dct[seq_all_places[0].name] = aux_tt['length', seq_all_places[0].name] / 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: + 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: + 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 + elif ss.from_ is None: + s_center_dct[ss.name] = ss.at n_resolved += 1 + else: + if ss.from_ in s_center_dct: + s_center_dct[ss.name] = s_center_dct[ss.from_] + ss.at + n_resolved += 1 + + assert n_resolved == len(seq_all_places), 'Not all positions resolved' + + aux_s_center = np.array([s_center_dct[nn] for nn in aux_tt.name[:-1]]) + 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)) + + return tt_sorted + +def _generate_line_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_element(drift_name, xt.Drift, length=ds_upstream) + names_with_drifts.append(drift_name) + names_with_drifts.append(nn) + + line = env.new_line(components=names_with_drifts) + + return line + + + + + +env = xt.Environment() + +seq = [ + env.new_element('b1', xt.Bend, length=0.5), + env.new_element('q1', xt.Quadrupole, length=0.5), + Place(env.new_element('ip', xt.Marker), at=10), + # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + ( + env.new_element('before_before_right', xt.Marker), + env.new_element('before_right', xt.Sextupole, length=1), + Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + env.new_element('after_right', xt.Marker), + env.new_element('after_right2', xt.Marker), + ), + Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), + env.new_element('after_left', xt.Marker), + env.new_element('after_left2', xt.Bend, length=0.5), +] + + + + +seq_all_places = _all_places(seq) +tab_sorted = _resolve_s_positions(seq_all_places, env) +line = _generate_line_with_drifts(env, tab_sorted) -assert n_resolved == len(seq_all_places) - -aux_s_center = np.array([s_center_dct[nn] for nn in aux_tt.name[:-1]]) -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)) - -s_tol = 1e-12 -names_with_drifts = [] -# Create drifts -for nn in name_sorted: - 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() - drift = env.new_element(drift_name, xt.Drift, length=ds_upstream) - names_with_drifts.append(drift_name) - names_with_drifts.append(nn) - -line = env.new_line(components=names_with_drifts) From d7fde72de172ade062481cb051578ccb91831038 Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 21:11:14 +0200 Subject: [PATCH 082/159] Clean --- examples/lattice_design_shortcuts/001_dev_at_s.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index f8e382eeb..58874649f 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -23,8 +23,6 @@ def __init__(self, name, at=None, from_=None, anchor=None, from_anchor=None): 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: From d0cab92b57fbea1d5a3c70c706b1f50758045543 Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 9 Sep 2024 21:18:25 +0200 Subject: [PATCH 083/159] Clean up --- examples/lattice_design_shortcuts/001_dev_at_s.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 58874649f..e34a78d1d 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -109,7 +109,7 @@ def _resolve_s_positions(seq_all_places, env): return tt_sorted -def _generate_line_with_drifts(env, tt_sorted, s_tol=1e-12): +def _generate_element_names_with_drifts(env, tt_sorted, s_tol=1e-12): names_with_drifts = [] # Create drifts @@ -122,9 +122,7 @@ def _generate_line_with_drifts(env, tt_sorted, s_tol=1e-12): names_with_drifts.append(drift_name) names_with_drifts.append(nn) - line = env.new_line(components=names_with_drifts) - - return line + return list(map(str, names_with_drifts)) @@ -154,6 +152,12 @@ def _generate_line_with_drifts(env, tt_sorted, s_tol=1e-12): seq_all_places = _all_places(seq) tab_sorted = _resolve_s_positions(seq_all_places, env) -line = _generate_line_with_drifts(env, tab_sorted) +names = _generate_element_names_with_drifts(env, tab_sorted) + +line = env.new_line(components=names) +import matplotlib.pyplot as plt +plt.close('all') +line.survey().plot() +plt.show() \ No newline at end of file From d77f82154cc7738b0134a2df34f03a32d2cfd253 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 06:44:44 +0200 Subject: [PATCH 084/159] Skip s eval in case of pure line case --- .../lattice_design_shortcuts/001_dev_at_s.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index e34a78d1d..46aa419dd 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -46,8 +46,6 @@ def _all_places(seq): return seq_all_places - - 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) @@ -124,10 +122,6 @@ def _generate_element_names_with_drifts(env, tt_sorted, s_tol=1e-12): return list(map(str, names_with_drifts)) - - - - env = xt.Environment() seq = [ @@ -148,12 +142,19 @@ def _generate_element_names_with_drifts(env, tt_sorted, s_tol=1e-12): ] +def handle_s_places(seq): + + places_found = np.array([isinstance(ss, Place) for ss in seq]).any() + if not places_found: + 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) -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 +names = handle_s_places(seq) line = env.new_line(components=names) import matplotlib.pyplot as plt From f3fc20d1282200fc8577a00c6b0c78d4bfe76008 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 09:27:11 +0200 Subject: [PATCH 085/159] Status --- .../lattice_design_shortcuts/001_dev_at_s.py | 97 +++++++++++++------ xtrack/handle_s_places.py | 0 2 files changed, 68 insertions(+), 29 deletions(-) create mode 100644 xtrack/handle_s_places.py diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 46aa419dd..5fc63269f 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -45,6 +45,18 @@ def _all_places(seq): seq_all_places.append(Place(ss, at=None, from_=None)) return seq_all_places +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] @@ -57,7 +69,8 @@ def _resolve_s_positions(seq_all_places, env): n_resolved_prev = -1 if seq_all_places[0].at is None and not seq_all_places[0]._before: - 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] = 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: @@ -69,27 +82,42 @@ def _resolve_s_positions(seq_all_places, env): ss_prev = seq_all_places[ii-1] if ss_prev.name in s_center_dct: s_center_dct[ss.name] = (s_center_dct[ss_prev.name] - + aux_tt['length', ss_prev.name] / 2 - + aux_tt['length', ss.name] / 2) + # + aux_tt['length', ss_prev.name] / 2 + + _length_expr_or_val(ss_prev.name, aux_line) / 2 + # + aux_tt['length', ss.name] / 2) + + _length_expr_or_val(ss.name, aux_line) / 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: s_center_dct[ss.name] = (s_center_dct[ss_next.name] - - aux_tt['length', ss_next.name] / 2 - - aux_tt['length', ss.name] / 2) + # - aux_tt['length', ss_next.name] / 2 + - _length_expr_or_val(ss_next.name, aux_line) / 2 + # - aux_tt['length', ss.name] / 2) + - _length_expr_or_val(ss.name, aux_line) / 2) n_resolved += 1 - elif ss.from_ is None: - s_center_dct[ss.name] = ss.at - n_resolved += 1 else: - if ss.from_ in s_center_dct: - s_center_dct[ss.name] = s_center_dct[ss.from_] + ss.at + 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 = np.array([s_center_dct[nn] for nn in aux_tt.name[:-1]]) + 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) @@ -122,38 +150,49 @@ def _generate_element_names_with_drifts(env, tt_sorted, s_tol=1e-12): return list(map(str, names_with_drifts)) +def handle_s_places(seq): + + places_found = np.array([isinstance(ss, Place) for ss in seq]).any() + if not places_found: + 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 + + + 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, +}) seq = [ - env.new_element('b1', xt.Bend, length=0.5), - env.new_element('q1', xt.Quadrupole, length=0.5), - Place(env.new_element('ip', xt.Marker), at=10), + env.new_element('b1', xt.Bend, length='l.b1'), + env.new_element('q1', xt.Quadrupole, length='l.q1'), + Place(env.new_element('ip', xt.Marker), at='s.ip'), # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), ( env.new_element('before_before_right', xt.Marker), env.new_element('before_right', xt.Sextupole, length=1), - Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + Place(env.new_element('right',xt.Quadrupole, length=1), at='s.right', from_='ip'), env.new_element('after_right', xt.Marker), env.new_element('after_right2', xt.Marker), ), - Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), + Place(env.new_element('left', xt.Quadrupole, length=1), at='s.left', from_='ip'), env.new_element('after_left', xt.Marker), - env.new_element('after_left2', xt.Bend, length=0.5), + env.new_element('after_left2', xt.Bend, length='l.after_left2'), ] -def handle_s_places(seq): - - places_found = np.array([isinstance(ss, Place) for ss in seq]).any() - if not places_found: - 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 - names = handle_s_places(seq) line = env.new_line(components=names) diff --git a/xtrack/handle_s_places.py b/xtrack/handle_s_places.py new file mode 100644 index 000000000..e69de29bb From e15dacdfd72e73fee7925c238859a040b2df3d10 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 09:33:32 +0200 Subject: [PATCH 086/159] With expressions --- examples/lattice_design_shortcuts/001_dev_at_s.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 5fc63269f..57deabbb2 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -133,6 +133,8 @@ def _resolve_s_positions(seq_all_places, env): 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): @@ -160,7 +162,7 @@ def handle_s_places(seq): tab_sorted = _resolve_s_positions(seq_all_places, env) names = _generate_element_names_with_drifts(env, tab_sorted) - return names + return names, tab_sorted @@ -192,8 +194,7 @@ def handle_s_places(seq): env.new_element('after_left2', xt.Bend, length='l.after_left2'), ] - -names = handle_s_places(seq) +names, tab_sorted = handle_s_places(seq) line = env.new_line(components=names) import matplotlib.pyplot as plt From c4ef72ec2053f4bcd7849bb006fe14ee77e0d18a Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 15:38:04 +0200 Subject: [PATCH 087/159] Disable expressions in s --- .../lattice_design_shortcuts/001_dev_at_s.py | 44 +++++++++++-------- 1 file changed, 25 insertions(+), 19 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 57deabbb2..f7a2bcdf5 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -45,17 +45,18 @@ def _all_places(seq): seq_all_places.append(Place(ss, at=None, from_=None)) return seq_all_places -def _length_expr_or_val(name, line): - if isinstance(line[name], xt.Replica): - name = line[name].resolve(line, get_name=True) +# 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 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 +# 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): @@ -69,8 +70,9 @@ def _resolve_s_positions(seq_all_places, env): n_resolved_prev = -1 if seq_all_places[0].at is None and not seq_all_places[0]._before: - # 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 + # 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: @@ -81,20 +83,24 @@ def _resolve_s_positions(seq_all_places, env): 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 - + _length_expr_or_val(ss_prev.name, aux_line) / 2 - # + aux_tt['length', ss.name] / 2) - + _length_expr_or_val(ss.name, aux_line) / 2) + + 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 - - _length_expr_or_val(ss_next.name, aux_line) / 2 - # - aux_tt['length', ss.name] / 2) - - _length_expr_or_val(ss.name, aux_line) / 2) + - aux_tt['length', ss_next.name] / 2 + - aux_tt['length', ss.name] / 2) n_resolved += 1 else: if isinstance(ss.at, str): From f142b9b9b706e0e786c4ab7bd584d6fee89bd1ee Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 16:53:57 +0200 Subject: [PATCH 088/159] Start preparing integration --- xtrack/handle_s_places.py | 168 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 168 insertions(+) diff --git a/xtrack/handle_s_places.py b/xtrack/handle_s_places.py index e69de29bb..81978f0f4 100644 --- a/xtrack/handle_s_places.py +++ b/xtrack/handle_s_places.py @@ -0,0 +1,168 @@ +import xtrack as xt +import numpy as np + +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_element(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): + + places_found = np.array([isinstance(ss, Place) for ss in seq]).any() + if not places_found: + 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, tab_sorted From 1f84b78629c1eca42d99d25e149630630317577e Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 17:03:35 +0200 Subject: [PATCH 089/159] Move environment --- xtrack/__init__.py | 3 +- xtrack/{handle_s_places.py => environment.py} | 63 +++++++++++++++++++ xtrack/line.py | 62 ------------------ 3 files changed, 65 insertions(+), 63 deletions(-) rename xtrack/{handle_s_places.py => environment.py} (76%) diff --git a/xtrack/__init__.py b/xtrack/__init__.py index ec482e8a4..3f0b15d07 100644 --- a/xtrack/__init__.py +++ b/xtrack/__init__.py @@ -12,7 +12,8 @@ from .beam_elements import * from .random import * from .tracker_data import TrackerData -from .line import Line, Node, freeze_longitudinal, _temp_knobs, EnergyProgram, Environment +from .line import Line, Node, freeze_longitudinal, _temp_knobs, EnergyProgram +from .environment import Environment from .tracker import Tracker, Log from .match import (Vary, Target, TargetList, VaryList, TargetInequality, Action, TargetRelPhaseAdvance, TargetSet, GreaterThan, LessThan, diff --git a/xtrack/handle_s_places.py b/xtrack/environment.py similarity index 76% rename from xtrack/handle_s_places.py rename to xtrack/environment.py index 81978f0f4..19fc32c52 100644 --- a/xtrack/handle_s_places.py +++ b/xtrack/environment.py @@ -1,5 +1,68 @@ import xtrack as xt import numpy as np +from weakref import WeakSet + + +def _flatten_components(components): + flatten_components = [] + for nn in components: + if 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 = WeakSet() + self._drift_counter = 0 + + 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 = [] + out.element_names = _flatten_components(components) + out._var_management = self._var_management + out._name = name + self._lines.add(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() + +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.new_element = xt.Line.new_element class Place: diff --git a/xtrack/line.py b/xtrack/line.py index 9f45aefb6..b0b5c6600 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -10,7 +10,6 @@ import uuid import os from collections import defaultdict -from weakref import WeakSet from contextlib import contextmanager from copy import deepcopy @@ -5113,64 +5112,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 - -def _flatten_components(components): - flatten_components = [] - for nn in components: - if isinstance(nn, 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 = WeakSet() - self._drift_counter = 0 - - def new_line(self, components=None, name=None): - out = Line() - out.particle_ref = self.particle_ref - out.env = self - out._element_dict = self.element_dict # Avoid copying - if components is None: - components = [] - out.element_names = _flatten_components(components) - out._var_management = self._var_management - out._name = name - self._lines.add(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() - -Environment.element_dict = Line.element_dict -Environment._init_var_management = Line._init_var_management -Environment._xdeps_vref = Line._xdeps_vref -Environment._xdeps_fref = Line._xdeps_fref -Environment._xdeps_manager = Line._xdeps_manager -Environment._xdeps_eval = Line._xdeps_eval -Environment.element_refs = Line.element_refs -Environment.vars = Line.vars -Environment.varval = Line.varval -Environment.vv = Line.vv -Environment.new_element = Line.new_element From bdca496b708b1efd075c5458b5efb1e677fe46c3 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 17:25:48 +0200 Subject: [PATCH 090/159] Integrated --- .../lattice_design_shortcuts/001_dev_at_s.py | 177 +----------------- xtrack/__init__.py | 2 +- xtrack/environment.py | 8 +- 3 files changed, 11 insertions(+), 176 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index f7a2bcdf5..1e387a646 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -1,176 +1,12 @@ import xtrack as xt import numpy as np +Place = xt.Place + def _plot_line(line): tt = line.get_table(attr=True) xt.twiss.TwissTable.plot(tt, yl='', yr='') -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_element(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): - - places_found = np.array([isinstance(ss, Place) for ss in seq]).any() - if not places_found: - 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, tab_sorted - - env = xt.Environment() @@ -183,7 +19,9 @@ def handle_s_places(seq): 'l.before_right': 1, 'l.after_left2': 0.5, }) -seq = [ + +# names, tab_sorted = handle_s_places(seq) +line = env.new_line(components=[ env.new_element('b1', xt.Bend, length='l.b1'), env.new_element('q1', xt.Quadrupole, length='l.q1'), Place(env.new_element('ip', xt.Marker), at='s.ip'), @@ -198,10 +36,7 @@ def handle_s_places(seq): Place(env.new_element('left', xt.Quadrupole, length=1), at='s.left', from_='ip'), env.new_element('after_left', xt.Marker), env.new_element('after_left2', xt.Bend, length='l.after_left2'), -] - -names, tab_sorted = handle_s_places(seq) -line = env.new_line(components=names) +]) import matplotlib.pyplot as plt plt.close('all') diff --git a/xtrack/__init__.py b/xtrack/__init__.py index 3f0b15d07..ec861dc05 100644 --- a/xtrack/__init__.py +++ b/xtrack/__init__.py @@ -13,7 +13,7 @@ from .random import * from .tracker_data import TrackerData from .line import Line, Node, freeze_longitudinal, _temp_knobs, EnergyProgram -from .environment import Environment +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 index 19fc32c52..7d84615cc 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -32,7 +32,8 @@ def new_line(self, components=None, name=None): out._element_dict = self.element_dict # Avoid copying if components is None: components = [] - out.element_names = _flatten_components(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) @@ -220,12 +221,11 @@ def _generate_element_names_with_drifts(env, tt_sorted, s_tol=1e-12): def handle_s_places(seq, env): - places_found = np.array([isinstance(ss, Place) for ss in seq]).any() - if not places_found: + 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, tab_sorted + return names From 519741f903ddb9ad77ab39c9a8581ea8d473ab74 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 17:28:34 +0200 Subject: [PATCH 091/159] Move new_element to environment --- xtrack/environment.py | 53 ++++++++++++++++++++++++++++++++++++++++++- xtrack/line.py | 51 ----------------------------------------- 2 files changed, 52 insertions(+), 52 deletions(-) diff --git a/xtrack/environment.py b/xtrack/environment.py index 7d84615cc..96eb6c50c 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -1,4 +1,5 @@ import xtrack as xt +import xobjects as xo import numpy as np from weakref import WeakSet @@ -53,6 +54,57 @@ def _get_a_drift_name(self): else: return self._get_a_drift_name() + def new_element(self, name, cls, **kwargs): + + _eval = self._xdeps_eval.eval + + 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.') + 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]) + value_kwargs[kk] = ref_kwargs[kk]._value + else: + value_kwargs[kk] = kwargs[kk] + + element = cls(**value_kwargs) + self.element_dict[name] = element + for kk in ref_kwargs: + if isinstance(ref_kwargs[kk], list): + for ii, vvv in enumerate(ref_kwargs[kk]): + if vvv is not None: + getattr(self.element_refs[name], kk)[ii] = vvv + else: + setattr(self.element_refs[name], kk, ref_kwargs[kk]) + + return name + Environment.element_dict = xt.Line.element_dict Environment._init_var_management = xt.Line._init_var_management Environment._xdeps_vref = xt.Line._xdeps_vref @@ -63,7 +115,6 @@ def _get_a_drift_name(self): Environment.vars = xt.Line.vars Environment.varval = xt.Line.varval Environment.vv = xt.Line.vv -Environment.new_element = xt.Line.new_element class Place: diff --git a/xtrack/line.py b/xtrack/line.py index b0b5c6600..ce5f815a4 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3354,57 +3354,6 @@ def transform_compound(self, *args, **kwargs): 'https://xsuite.readthedocs.io/en/latest/line.html#apply-transformations-tilt-shift-to-elements' ) - def new_element(self, name, cls, **kwargs): - - _eval = self._xdeps_eval.eval - - 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.') - 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]) - value_kwargs[kk] = ref_kwargs[kk]._value - else: - value_kwargs[kk] = kwargs[kk] - - element = cls(**value_kwargs) - self.element_dict[name] = element - for kk in ref_kwargs: - if isinstance(ref_kwargs[kk], list): - for ii, vvv in enumerate(ref_kwargs[kk]): - if vvv is not None: - getattr(self.element_refs[name], kk)[ii] = vvv - else: - setattr(self.element_refs[name], kk, ref_kwargs[kk]) - - return name - def mirror(self): self._frozen_check() self.element_names = list(reversed(self.element_names)) From 3705bd649013a603cd9f90943bc5c9b1b4da5b2e Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 17:33:05 +0200 Subject: [PATCH 092/159] Place disappears --- examples/lattice_design_shortcuts/001_dev_at_s.py | 8 +++----- xtrack/environment.py | 6 +++++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 1e387a646..bac9c8e26 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -1,7 +1,6 @@ import xtrack as xt import numpy as np -Place = xt.Place def _plot_line(line): tt = line.get_table(attr=True) @@ -24,16 +23,15 @@ def _plot_line(line): line = env.new_line(components=[ env.new_element('b1', xt.Bend, length='l.b1'), env.new_element('q1', xt.Quadrupole, length='l.q1'), - Place(env.new_element('ip', xt.Marker), at='s.ip'), - # Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), + env.new_element('ip', xt.Marker, at='s.ip'), ( env.new_element('before_before_right', xt.Marker), env.new_element('before_right', xt.Sextupole, length=1), - Place(env.new_element('right',xt.Quadrupole, length=1), at='s.right', from_='ip'), + env.new_element('right',xt.Quadrupole, length=1, at='s.right', from_='ip'), env.new_element('after_right', xt.Marker), env.new_element('after_right2', xt.Marker), ), - Place(env.new_element('left', xt.Quadrupole, length=1), at='s.left', from_='ip'), + env.new_element('left', xt.Quadrupole, length=1, at='s.left', from_='ip'), env.new_element('after_left', xt.Marker), env.new_element('after_left2', xt.Bend, length='l.after_left2'), ]) diff --git a/xtrack/environment.py b/xtrack/environment.py index 96eb6c50c..2d477e009 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -54,7 +54,11 @@ def _get_a_drift_name(self): else: return self._get_a_drift_name() - def new_element(self, name, cls, **kwargs): + def new_element(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_element(name, cls, **kwargs)) _eval = self._xdeps_eval.eval From 8b83bbe69552e7cfc781368cfefce821e852fb87 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 21:27:42 +0200 Subject: [PATCH 093/159] Refactoring example --- examples/lattice_design_shortcuts/000_dev.py | 67 ++--- .../002_ring_sequence.py | 265 ++++++++++++++++++ .../{002_play.py => 010_play.py} | 0 3 files changed, 296 insertions(+), 36 deletions(-) create mode 100644 examples/lattice_design_shortcuts/002_ring_sequence.py rename examples/lattice_design_shortcuts/{002_play.py => 010_play.py} (100%) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index f5c44188d..9410684fb 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -16,48 +16,34 @@ 'l.mq': 0.5, 'kqf.1': 'k1l.qf / l.mq', 'kqd.1': 'k1l.qd / l.mq', - 'l.mb': 12, + '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, }) +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') + env.new_element('mid', xt.Marker, at='l.halfcell'), -halfcell = env.new_line(components=[ - 'drift.1', - 'qf', - 'drift.2', - 'mb.1', - 'mb.2', - 'mb.3', - 'drift.3', - 'qd', - 'drift.4', -]) + env.new_element('mb.2', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb', at='l.halfcell / 2'), + env.new_element('mb.1', xt.Replica, parent_name='mb.2', at='-l.mb - 1', from_='mb.2'), + env.new_element('mb.3', xt.Replica, parent_name='mb.2', at='l.mb + 1', from_='mb.2'), + env.new_element('mq.f', xt.Quadrupole, k1='kqf.1', length='l.mq', at = '0.5 + l.mq / 2'), + env.new_element('mq.d', xt.Quadrupole, k1='kqd.1', length='l.mq', at = 'l.halfcell - l.mq / 2 - 0.5'), + env.new_element('corrector.h', xt.Multipole, at=1.5), + env.new_element('corrector.v', xt.Multipole, at='l.halfcell - 1.5'), + + env.new_element('ms.f', xt.Sextupole, length='l.ms', k2='k2sf', at=2.), + env.new_element('ms.d', xt.Sextupole, length='l.ms', k2='k2sd', at='l.halfcell - 2.'), -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') @@ -65,16 +51,21 @@ 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), ]) +halfcell_ss = env.new_line(components=[ + + env.new_element('mid', xt.Marker, at='l.halfcell'), + + env.new_element('mq.f', xt.Quadrupole, k1='kqf.1', length='l.mq', at = '0.5 + l.mq / 2'), + env.new_element('mq.d', xt.Quadrupole, k1='kqd.1', length='l.mq', at = 'l.halfcell - l.mq / 2 - 0.5'), + + env.new_element('corrector.h', xt.Multipole, at=1.5), +]) + arc = env.new_line(components=[ cell.replicate(name='cell.1'), @@ -82,6 +73,10 @@ cell.replicate(name='cell.3'), ]) + +cell_ss = env + + cell_ss = cell.replicate('ss') env.new_element('drift_ss', xt.Drift, length='l.mb') for ii, nn in enumerate(cell_ss.element_names): 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/002_play.py b/examples/lattice_design_shortcuts/010_play.py similarity index 100% rename from examples/lattice_design_shortcuts/002_play.py rename to examples/lattice_design_shortcuts/010_play.py From e9f9520007b27fb0e4d95bfee19b08acd734b7a3 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 21:50:28 +0200 Subject: [PATCH 094/159] Getting the ring back --- examples/lattice_design_shortcuts/000_dev.py | 93 ++++++++++++-------- 1 file changed, 55 insertions(+), 38 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 9410684fb..8fc4aca8c 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -11,11 +11,9 @@ 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', + 'kqf': 0.027, + 'kqd': -0.0271, 'l.mb': 10, 'l.ms': 0.3, 'k2sf': 0.001, @@ -35,19 +33,19 @@ env.new_element('mb.1', xt.Replica, parent_name='mb.2', at='-l.mb - 1', from_='mb.2'), env.new_element('mb.3', xt.Replica, parent_name='mb.2', at='l.mb + 1', from_='mb.2'), - env.new_element('mq.f', xt.Quadrupole, k1='kqf.1', length='l.mq', at = '0.5 + l.mq / 2'), - env.new_element('mq.d', xt.Quadrupole, k1='kqd.1', length='l.mq', at = 'l.halfcell - l.mq / 2 - 0.5'), + env.new_element('mq.d', xt.Quadrupole, k1='kqd', length='l.mq', at = '0.5 + l.mq / 2'), + env.new_element('mq.f', xt.Quadrupole, k1='kqf', length='l.mq', at = 'l.halfcell - l.mq / 2 - 0.5'), - env.new_element('corrector.h', xt.Multipole, at=1.5), - env.new_element('corrector.v', xt.Multipole, at='l.halfcell - 1.5'), + env.new_element('corrector.v', xt.Multipole, at=1.5), + env.new_element('corrector.h', xt.Multipole, at='l.halfcell - 1.5'), - env.new_element('ms.f', xt.Sextupole, length='l.ms', k2='k2sf', at=2.), - env.new_element('ms.d', xt.Sextupole, length='l.ms', k2='k2sd', at='l.halfcell - 2.'), + env.new_element('ms.d', xt.Sextupole, length='l.ms', k2='k2sf', at=2.), + env.new_element('ms.f', xt.Sextupole, length='l.ms', k2='k2sd', at='l.halfcell - 2.'), ]) -hcell_left = halfcell.replicate(name='l') -hcell_right = halfcell.replicate(name='r', mirror=True) +hcell_left = halfcell.replicate(name='l', mirror=True) +hcell_right = halfcell.replicate(name='r') cell = env.new_line(components=[ env.new_element('start', xt.Marker), @@ -56,16 +54,49 @@ env.new_element('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_element('mid', xt.Marker, at='l.halfcell'), - env.new_element('mq.f', xt.Quadrupole, k1='kqf.1', length='l.mq', at = '0.5 + l.mq / 2'), - env.new_element('mq.d', xt.Quadrupole, k1='kqd.1', length='l.mq', at = 'l.halfcell - l.mq / 2 - 0.5'), + env.new_element('mq.ss.d', xt.Quadrupole, k1='kqd.ss', length='l.mq', at = '0.5 + l.mq / 2'), + env.new_element('mq.ss.f', xt.Quadrupole, k1='kqf.ss', length='l.mq', at = 'l.halfcell - l.mq / 2 - 0.5'), - env.new_element('corrector.h', xt.Multipole, at=1.5), + env.new_element('corrector.ss.v', xt.Multipole, at=1.5), + env.new_element('corrector.ss.h', xt.Multipole, at='l.halfcell - 1.5'), ]) +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_element('start.ss', xt.Marker), + hcell_left_ss, + hcell_right_ss, + env.new_element('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'), @@ -74,38 +105,24 @@ ]) -cell_ss = env - - -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') +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'), +]) -# ss1 = ss.replicate(name='ss.1') -# ss2 = ss.replicate(name='ss.2') -# ss3 = ss.replicate(name='ss.3') +pttt -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, From e95fe68c5b72fbb3e401025fa0ce69d2767671ff Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 22:13:43 +0200 Subject: [PATCH 095/159] Still issues with the slicing --- examples/lattice_design_shortcuts/000_dev.py | 153 ++++++------------- 1 file changed, 44 insertions(+), 109 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 8fc4aca8c..078d7411a 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -119,51 +119,41 @@ ss.replicate(name='ss.3'), ]) - -pttt - - +## Insretion 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', + 'k1.q1': 0.025, + 'k1.q2': -0.025, + 'k1.q3': 0.025, + 'k1.q4': -0.02, + 'k1.q5': 0.025, }) -half_straight = env.new_line(components=[ +half_insertion = 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('mq.1', xt.Quadrupole, k1='k1.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('mq.2', xt.Quadrupole, k1='k1.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('mq.3', xt.Quadrupole, k1='k1.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('mq.4', xt.Quadrupole, k1='k1.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('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq'), env.new_element('dd.5', xt.Drift, length=0.5), - env.new_element('e.ss.r', xt.Marker), + env.new_element('e.insertion', xt.Marker), ]) -half_straight.build_tracker() -print(f'Half straight length: {half_straight.get_length()}') tw_arc = arc.twiss4d() -opt = half_straight.match( +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.ss.r', - start='ip', end='e.ss.r', - vary=xt.VaryList(['k1l.q1', 'k1l.q2', 'k1l.q3', 'k1l.q4'], step=1e-5), + 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), @@ -172,104 +162,49 @@ 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() -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'), +insertion = env.new_line([ + half_insertion.replicate('l', mirror=True), + half_insertion.replicate('r')]) - 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'), +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') ]) -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() +ring2_sliced = ring2.select() +# ring2_sliced.cut_at_s(np.arange(0, ring2.get_length(), 0.5)) -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) +import matplotlib.pyplot as plt +plt.close('all') +for ii, rr in enumerate([ring, ring2]): -ring.replace_all_replicas() -ring.build_tracker() -sv = ring.survey() + tw = rr.twiss4d() -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() + 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) -two = ring.twiss(start=xt.START, betx=tw_arc.betx[-1], bety=tw_arc.bety[-1]) +ring2.survey().plot() -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() From f38aef7758e462fc376ed3624eb9c685d66a0553 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 10 Sep 2024 22:21:50 +0200 Subject: [PATCH 096/159] Refactored 000_dev.py --- examples/lattice_design_shortcuts/000_dev.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 078d7411a..194fd1621 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -131,18 +131,12 @@ half_insertion = 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='k1.q1', length='l.mq'), - env.new_element('dd.1', xt.Drift, length=5), - env.new_element('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq'), - env.new_element('dd.2', xt.Drift, length=12), - env.new_element('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq'), - env.new_element('dd.3', xt.Drift, length=18), - env.new_element('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq'), - env.new_element('dd.4', xt.Drift, length=18), - env.new_element('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq'), - env.new_element('dd.5', xt.Drift, length=0.5), - env.new_element('e.insertion', xt.Marker), + env.new_element('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), + env.new_element('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), + env.new_element('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), + env.new_element('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), + env.new_element('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq', at=73), + env.new_element('e.insertion', xt.Marker, at=76), ]) tw_arc = arc.twiss4d() From 7994aa2da218aff02dc862188448f2dd1cb28bbc Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 06:45:13 +0200 Subject: [PATCH 097/159] Move all elements in the dict when building a tracker --- examples/lattice_design_shortcuts/000_dev.py | 9 ++++----- xtrack/line.py | 4 ++++ xtrack/tracker_data.py | 2 +- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 194fd1621..8a153d80d 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -175,16 +175,15 @@ ss.replicate(name='sss.3') ]) -ring2_sliced = ring2.select() -# ring2_sliced.cut_at_s(np.arange(0, ring2.get_length(), 0.5)) - - +# # 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]): +for ii, rr in enumerate([ring, ring2_sliced]): tw = rr.twiss4d() diff --git a/xtrack/line.py b/xtrack/line.py index ce5f815a4..9c7416b52 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -2062,6 +2062,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 diff --git a/xtrack/tracker_data.py b/xtrack/tracker_data.py index f4fa3af37..a99c1ba18 100644 --- a/xtrack/tracker_data.py +++ b/xtrack/tracker_data.py @@ -176,7 +176,7 @@ 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 ee._buffer is not buffer: if allow_move: ee.move(_buffer=buffer) From 7832dd9da05f12ffdaffc8cf741aad5fc3f198cf Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 09:53:30 +0200 Subject: [PATCH 098/159] Did not break anything --- xtrack/environment.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/xtrack/environment.py b/xtrack/environment.py index 2d477e009..a91c5f795 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -97,8 +97,18 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): else: value_kwargs[kk] = kwargs[kk] - element = cls(**value_kwargs) - self.element_dict[name] = element + 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) + for kk in value_kwargs: + setattr(self.element_dict[name], kk, value_kwargs[kk]) + else: + # Instantiate a new element + element = cls(**value_kwargs) + self.element_dict[name] = element + for kk in ref_kwargs: if isinstance(ref_kwargs[kk], list): for ii, vvv in enumerate(ref_kwargs[kk]): @@ -119,6 +129,7 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): Environment.vars = xt.Line.vars Environment.varval = xt.Line.varval Environment.vv = xt.Line.vv +Environment.replace_replica = xt.Line.replace_replica class Place: From 2981ceda9ffd744f6890e10fff687aa374d60c57 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 10:07:52 +0200 Subject: [PATCH 099/159] Capability to use another element as parent --- examples/lattice_design_shortcuts/000_dev.py | 4 +++- xtrack/environment.py | 19 +++++++++++++------ 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 8a153d80d..86720a988 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -25,11 +25,13 @@ 'l.halfcell': 38, }) +env.new_element('mb', xt.Bend) + halfcell = env.new_line(components=[ env.new_element('mid', xt.Marker, at='l.halfcell'), - env.new_element('mb.2', xt.Bend, k0='k0.mb', h='k0.mb', length='l.mb', at='l.halfcell / 2'), + env.new_element('mb.2', 'mb', k0='k0.mb', h='k0.mb', length='l.mb', at='l.halfcell / 2'), env.new_element('mb.1', xt.Replica, parent_name='mb.2', at='-l.mb - 1', from_='mb.2'), env.new_element('mb.3', xt.Replica, parent_name='mb.2', at='l.mb + 1', from_='mb.2'), diff --git a/xtrack/environment.py b/xtrack/environment.py index a91c5f795..bb207d835 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -62,10 +62,20 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): _eval = self._xdeps_eval.eval - assert cls in [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, + 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_element` 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 = {} for kk in kwargs: @@ -97,11 +107,7 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): else: value_kwargs[kk] = kwargs[kk] - 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) + if isinstance(cls_input, str): for kk in value_kwargs: setattr(self.element_dict[name], kk, value_kwargs[kk]) else: @@ -130,6 +136,7 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): Environment.varval = xt.Line.varval Environment.vv = xt.Line.vv Environment.replace_replica = xt.Line.replace_replica +Environment.__getitem__ = xt.Line.__getitem__ class Place: From 081dedfaf4cc57d03529fdd0be215c204bfe88e3 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 11:10:22 +0200 Subject: [PATCH 100/159] Inherit parameter --- examples/lattice_design_shortcuts/000_dev.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 86720a988..f257aa91a 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -25,13 +25,13 @@ 'l.halfcell': 38, }) -env.new_element('mb', xt.Bend) +env.new_element('mb', xt.Bend, length='l.mb') halfcell = env.new_line(components=[ env.new_element('mid', xt.Marker, at='l.halfcell'), - env.new_element('mb.2', 'mb', k0='k0.mb', h='k0.mb', length='l.mb', at='l.halfcell / 2'), + env.new_element('mb.2', 'mb', k0='k0.mb', h='k0.mb', at='l.halfcell / 2'), env.new_element('mb.1', xt.Replica, parent_name='mb.2', at='-l.mb - 1', from_='mb.2'), env.new_element('mb.3', xt.Replica, parent_name='mb.2', at='l.mb + 1', from_='mb.2'), From 3fc6d10673cb6db60a09edbb746a762b65f38210 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 11:28:27 +0200 Subject: [PATCH 101/159] Parents everywhere --- examples/lattice_design_shortcuts/000_dev.py | 52 ++++++++++++++------ 1 file changed, 36 insertions(+), 16 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index f257aa91a..4d61de979 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -25,24 +25,32 @@ 'l.halfcell': 38, }) -env.new_element('mb', xt.Bend, length='l.mb') +env.new_element('mb', xt.Bend, length='l.mb', k0='k0.mb', h='k0.mb') +env.new_element('mq', xt.Quadrupole, length='l.mq') +env.new_element('ms', xt.Sextupole, length='l.ms') +env.new_element('corrector', xt.Multipole, knl=[0], ksl=[0]) halfcell = env.new_line(components=[ + # End of the half cell (will be mid of the cell) env.new_element('mid', xt.Marker, at='l.halfcell'), - env.new_element('mb.2', 'mb', k0='k0.mb', h='k0.mb', at='l.halfcell / 2'), - env.new_element('mb.1', xt.Replica, parent_name='mb.2', at='-l.mb - 1', from_='mb.2'), - env.new_element('mb.3', xt.Replica, parent_name='mb.2', at='l.mb + 1', from_='mb.2'), + # Bends + env.new_element('mb.2', 'mb', at='l.halfcell / 2'), + env.new_element('mb.1', 'mb', at='-l.mb - 1', from_='mb.2'), + env.new_element('mb.3', 'mb', at='l.mb + 1', from_='mb.2'), - env.new_element('mq.d', xt.Quadrupole, k1='kqd', length='l.mq', at = '0.5 + l.mq / 2'), - env.new_element('mq.f', xt.Quadrupole, k1='kqf', length='l.mq', at = 'l.halfcell - l.mq / 2 - 0.5'), + # Quads + env.new_element('mq.d', 'mq', k1='kqd', at = '0.5 + l.mq / 2'), + env.new_element('mq.f', 'mq', k1='kqf', at = 'l.halfcell - l.mq / 2 - 0.5'), - env.new_element('corrector.v', xt.Multipole, at=1.5), - env.new_element('corrector.h', xt.Multipole, at='l.halfcell - 1.5'), + # Sextupoles + env.new_element('ms.d', 'ms', k2='k2sf', at=1.2, from_='mq.d'), + env.new_element('ms.f', 'ms', k2='k2sd', at=-1.2, from_='mq.f'), - env.new_element('ms.d', xt.Sextupole, length='l.ms', k2='k2sf', at=2.), - env.new_element('ms.f', xt.Sextupole, length='l.ms', k2='k2sd', at='l.halfcell - 2.'), + # Dipole correctors + env.new_element('corrector.v', 'corrector', at=0.75, from_='mq.d'), + env.new_element('corrector.h', 'corrector', at=-0.75, from_='mq.f') ]) @@ -75,11 +83,11 @@ env.new_element('mid', xt.Marker, at='l.halfcell'), - env.new_element('mq.ss.d', xt.Quadrupole, k1='kqd.ss', length='l.mq', at = '0.5 + l.mq / 2'), - env.new_element('mq.ss.f', xt.Quadrupole, k1='kqf.ss', length='l.mq', at = 'l.halfcell - l.mq / 2 - 0.5'), + env.new_element('mq.ss.d', 'mq', k1='kqd.ss', at = '0.5 + l.mq / 2'), + env.new_element('mq.ss.f', 'mq', k1='kqf.ss', at = 'l.halfcell - l.mq / 2 - 0.5'), - env.new_element('corrector.ss.v', xt.Multipole, at=1.5), - env.new_element('corrector.ss.h', xt.Multipole, at='l.halfcell - 1.5'), + env.new_element('corrector.ss.v', 'corrector', at=0.75, from_='mq.ss.d'), + env.new_element('corrector.ss.h', 'corrector', at=-0.75, from_='mq.ss.f') ]) hcell_left_ss = halfcell_ss.replicate(name='l', mirror=True) @@ -121,7 +129,7 @@ ss.replicate(name='ss.3'), ]) -## Insretion +## Insertion env.vars({ 'k1.q1': 0.025, @@ -132,13 +140,25 @@ }) half_insertion = env.new_line(components=[ + + # Start-end markers env.new_element('ip', xt.Marker), + env.new_element('e.insertion', xt.Marker, at=76), + + # Quads env.new_element('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), env.new_element('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), env.new_element('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), env.new_element('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), env.new_element('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq', at=73), - env.new_element('e.insertion', xt.Marker, at=76), + + # Dipole correctors (will use h and v on the same corrector) + env.new_element('corrector.ss.1', 'corrector', at=0.75, from_='mq.1'), + env.new_element('corrector.ss.2', 'corrector', at=-0.75, from_='mq.2'), + env.new_element('corrector.ss.3', 'corrector', at=0.75, from_='mq.3'), + env.new_element('corrector.ss.4', 'corrector', at=-0.75, from_='mq.4'), + env.new_element('corrector.ss.5', 'corrector', at=0.75, from_='mq.5'), + ]) tw_arc = arc.twiss4d() From ab32460f767c36ddc8c08463d4518aa74fcfb51e Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 21:13:17 +0200 Subject: [PATCH 102/159] Test place --- examples/lattice_design_shortcuts/000_dev.py | 7 +++++-- xtrack/environment.py | 3 +++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 4d61de979..0d895bcb3 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -30,6 +30,9 @@ env.new_element('ms', xt.Sextupole, length='l.ms') env.new_element('corrector', xt.Multipole, knl=[0], ksl=[0]) +env.new_element('mq.f', 'mq', k1='kqf') +env.new_element('mq.d', 'mq', k1='kqd') + halfcell = env.new_line(components=[ # End of the half cell (will be mid of the cell) @@ -41,8 +44,8 @@ env.new_element('mb.3', 'mb', at='l.mb + 1', from_='mb.2'), # Quads - env.new_element('mq.d', 'mq', k1='kqd', at = '0.5 + l.mq / 2'), - env.new_element('mq.f', 'mq', k1='kqf', at = 'l.halfcell - l.mq / 2 - 0.5'), + 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_element('ms.d', 'ms', k2='k2sf', at=1.2, from_='mq.d'), diff --git a/xtrack/environment.py b/xtrack/environment.py index bb207d835..801cf51a7 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -125,6 +125,9 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): 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 From 72271ba35e5b618c414f1425364093ee24706938 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 22:49:19 +0200 Subject: [PATCH 103/159] Force movable --- .../lattice_design_shortcuts/000a_subline.py | 236 ++++++++++++++++++ xtrack/environment.py | 13 +- xtrack/line.py | 1 + xtrack/slicing.py | 1 + 4 files changed, 250 insertions(+), 1 deletion(-) create mode 100644 examples/lattice_design_shortcuts/000a_subline.py diff --git a/examples/lattice_design_shortcuts/000a_subline.py b/examples/lattice_design_shortcuts/000a_subline.py new file mode 100644 index 000000000..0b385875d --- /dev/null +++ b/examples/lattice_design_shortcuts/000a_subline.py @@ -0,0 +1,236 @@ +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_element('mb', xt.Bend, length='l.mb', k0='k0.mb', h='k0.mb') +env.new_element('mq', xt.Quadrupole, length='l.mq') +env.new_element('ms', xt.Sextupole, length='l.ms') +env.new_element('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.replicate(name='f') +girder_d = girder.replicate(name='d', mirror=True) + +girder_f.replace_all_replicas() +girder_d.replace_all_replicas() + +girder_f.element_refs['mq.f'].k1 = env.vars['kqf'] +girder_d.element_refs['mq.d'].k1 = env.vars['kqd'] + +halfcell = env.new_line(components=[ + + # End of the half cell (will be mid of the cell) + env.new_element('mid', xt.Marker, at='l.halfcell'), + + # Bends + env.new_element('mb.2', 'mb', at='l.halfcell / 2'), + env.new_element('mb.1', 'mb', at='-l.mb - 1', from_='mb.2'), + env.new_element('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_element('start', xt.Marker), + hcell_left, + hcell_right, + env.new_element('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_element('mid', xt.Marker, at='l.halfcell'), + + env.new_element('mq.ss.d', 'mq', k1='kqd.ss', at = '0.5 + l.mq / 2'), + env.new_element('mq.ss.f', 'mq', k1='kqf.ss', at = 'l.halfcell - l.mq / 2 - 0.5'), + + env.new_element('corrector.ss.v', 'corrector', at=0.75, from_='mq.ss.d'), + env.new_element('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_element('start.ss', xt.Marker), + hcell_left_ss, + hcell_right_ss, + env.new_element('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_element('ip', xt.Marker), + env.new_element('e.insertion', xt.Marker, at=76), + + # Quads + env.new_element('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), + env.new_element('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), + env.new_element('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), + env.new_element('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), + env.new_element('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_element('corrector.ss.1', 'corrector', at=0.75, from_='mq.1'), + env.new_element('corrector.ss.2', 'corrector', at=-0.75, from_='mq.2'), + env.new_element('corrector.ss.3', 'corrector', at=0.75, from_='mq.3'), + env.new_element('corrector.ss.4', 'corrector', at=-0.75, from_='mq.4'), + env.new_element('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/xtrack/environment.py b/xtrack/environment.py index 801cf51a7..2d656086e 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -7,7 +7,18 @@ def _flatten_components(components): flatten_components = [] for nn in components: - if isinstance(nn, xt.Line): + if isinstance(nn, Place) and isinstance(nn.name, xt.Line): + line = nn.name + components = 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) diff --git a/xtrack/line.py b/xtrack/line.py index 85978083e..aa6a860af 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3363,6 +3363,7 @@ def replicate(self, name, mirror=False): out = Line() out.element_names = new_element_names out._element_dict = self.element_dict # to make sure that the dict is not copied + out._var_management = self._var_management out._name = name if mirror: diff --git a/xtrack/slicing.py b/xtrack/slicing.py index 8a4e60812..157776f95 100644 --- a/xtrack/slicing.py +++ b/xtrack/slicing.py @@ -413,6 +413,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) From da27a4ab6d5ef09f23abb1ea193ad54ccff381b6 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 22:51:46 +0200 Subject: [PATCH 104/159] Small fix --- xtrack/tracker_data.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xtrack/tracker_data.py b/xtrack/tracker_data.py index a99c1ba18..34ba21443 100644 --- a/xtrack/tracker_data.py +++ b/xtrack/tracker_data.py @@ -177,6 +177,8 @@ def check_elements_in_common_buffer(self, buffer, allow_move=False): there. """ 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) From f6193de0ea95a547a6009c6373c1dbf7b99fc405 Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 22:57:43 +0200 Subject: [PATCH 105/159] Another fix --- xtrack/slicing.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/xtrack/slicing.py b/xtrack/slicing.py index 157776f95..88a1bb943 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) @@ -393,6 +394,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 +405,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) @@ -425,6 +428,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) From b2ec24d4c694d7ae8a7087ada9960b2da513ddec Mon Sep 17 00:00:00 2001 From: giadarol Date: Wed, 11 Sep 2024 23:47:04 +0200 Subject: [PATCH 106/159] Hamdle replica reallocation --- tests/test_slice_and_insert_with_replicas.py | 16 ++++++++++++++++ xtrack/slicing.py | 1 + xtrack/tracker_data.py | 1 + 3 files changed, 18 insertions(+) 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/slicing.py b/xtrack/slicing.py index 88a1bb943..6ede147ff 100644 --- a/xtrack/slicing.py +++ b/xtrack/slicing.py @@ -383,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) diff --git a/xtrack/tracker_data.py b/xtrack/tracker_data.py index 34ba21443..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): From 7115fe8b14d6dea8c60cfd7ebae883cc49b1b49c Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 06:09:24 +0200 Subject: [PATCH 107/159] Fix --- xtrack/twiss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 3154a09bc..b3ca0fb75 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -2996,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, From f8910b4e17d45c19c41a36c2faede1a8a3b58d4b Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 06:38:03 +0200 Subject: [PATCH 108/159] Fix --- xtrack/twiss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index b3ca0fb75..00bec4f03 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -331,9 +331,9 @@ 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: + if (init is not None or betx is not None or bety is not None) and start is None: # is open twiss - start = start or xt.START + start = xt.START end = end or xt.END if num_turns != 1: From a3be6b53492bbb98d790c7c97982b7deaa7a12a9 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 10:19:48 +0200 Subject: [PATCH 109/159] Refactor --- .../lattice_design_shortcuts/000a_subline.py | 4 +- xtrack/environment.py | 65 ++++++++++--------- 2 files changed, 37 insertions(+), 32 deletions(-) diff --git a/examples/lattice_design_shortcuts/000a_subline.py b/examples/lattice_design_shortcuts/000a_subline.py index 0b385875d..cd1403136 100644 --- a/examples/lattice_design_shortcuts/000a_subline.py +++ b/examples/lattice_design_shortcuts/000a_subline.py @@ -37,9 +37,9 @@ ]) girder_f = girder.replicate(name='f') -girder_d = girder.replicate(name='d', mirror=True) - girder_f.replace_all_replicas() + +girder_d = girder.replicate(name='d', mirror=True) girder_d.replace_all_replicas() girder_f.element_refs['mq.f'].k1 = env.vars['kqf'] diff --git a/xtrack/environment.py b/xtrack/environment.py index 2d656086e..3b446aea2 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -87,36 +87,7 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): self.replace_replica(name) cls = type(self.element_dict[name]) - 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]) - value_kwargs[kk] = ref_kwargs[kk]._value - else: - value_kwargs[kk] = kwargs[kk] + ref_kwargs, value_kwargs = _parse_kwargs(cls, kwargs, _eval) if isinstance(cls_input, str): for kk in value_kwargs: @@ -316,3 +287,37 @@ def handle_s_places(seq, 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]) + value_kwargs[kk] = ref_kwargs[kk]._value + else: + value_kwargs[kk] = kwargs[kk] + + return ref_kwargs, value_kwargs \ No newline at end of file From 15a487c78d7d4887dd62d6f50785052d069fdd09 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 10:32:55 +0200 Subject: [PATCH 110/159] Refactor --- xtrack/environment.py | 36 +++++++++++++++++++++--------------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/xtrack/environment.py b/xtrack/environment.py index 3b446aea2..367b8e65d 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -89,21 +89,11 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): ref_kwargs, value_kwargs = _parse_kwargs(cls, kwargs, _eval) - if isinstance(cls_input, str): - for kk in value_kwargs: - setattr(self.element_dict[name], kk, value_kwargs[kk]) - else: - # Instantiate a new element - element = cls(**value_kwargs) - self.element_dict[name] = element + if not isinstance(cls_input, str): # Parent is a class and not another element + self.element_dict[name] = cls(**value_kwargs) - for kk in ref_kwargs: - if isinstance(ref_kwargs[kk], list): - for ii, vvv in enumerate(ref_kwargs[kk]): - if vvv is not None: - getattr(self.element_refs[name], kk)[ii] = vvv - else: - setattr(self.element_refs[name], kk, ref_kwargs[kk]) + _set_kwargs(name=name, ref_kwargs=ref_kwargs, value_kwargs=value_kwargs, + element_dict=self.element_dict, element_refs=self.element_refs) return name @@ -320,4 +310,20 @@ def _parse_kwargs(cls, kwargs, _eval): else: value_kwargs[kk] = kwargs[kk] - return ref_kwargs, value_kwargs \ No newline at end of file + 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]) + From 968895832271b683dfadd750e4900dd097980777 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 10:40:14 +0200 Subject: [PATCH 111/159] Set is working --- examples/lattice_design_shortcuts/000a_subline.py | 5 +++-- xtrack/environment.py | 9 +++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/000a_subline.py b/examples/lattice_design_shortcuts/000a_subline.py index cd1403136..fb1e722ed 100644 --- a/examples/lattice_design_shortcuts/000a_subline.py +++ b/examples/lattice_design_shortcuts/000a_subline.py @@ -42,8 +42,9 @@ girder_d = girder.replicate(name='d', mirror=True) girder_d.replace_all_replicas() -girder_f.element_refs['mq.f'].k1 = env.vars['kqf'] -girder_d.element_refs['mq.d'].k1 = env.vars['kqd'] +env.set('mq.f', k1='kqf') +env.set('mq.d', k1='kqd') + halfcell = env.new_line(components=[ diff --git a/xtrack/environment.py b/xtrack/environment.py index 367b8e65d..85af9f3e8 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -97,6 +97,15 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): return name + def set(self, name, **kwargs): + _eval = self._xdeps_eval.eval + + ref_kwargs, value_kwargs = _parse_kwargs( + type(self.element_dict[name]), kwargs, _eval) + + _set_kwargs(name=name, ref_kwargs=ref_kwargs, value_kwargs=value_kwargs, + element_dict=self.element_dict, element_refs=self.element_refs) + 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) From 193c9f72dfe15a0990e38645e7688e14de48cc68 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 11:01:55 +0200 Subject: [PATCH 112/159] Rename new_element -> new --- examples/lattice_design_shortcuts/000_dev.py | 70 +++++++++---------- .../lattice_design_shortcuts/000a_subline.py | 58 +++++++-------- examples/lattice_design_shortcuts/010_play.py | 50 ++++++------- xtrack/environment.py | 8 +-- 4 files changed, 93 insertions(+), 93 deletions(-) diff --git a/examples/lattice_design_shortcuts/000_dev.py b/examples/lattice_design_shortcuts/000_dev.py index 0d895bcb3..2de3da73f 100644 --- a/examples/lattice_design_shortcuts/000_dev.py +++ b/examples/lattice_design_shortcuts/000_dev.py @@ -25,35 +25,35 @@ 'l.halfcell': 38, }) -env.new_element('mb', xt.Bend, length='l.mb', k0='k0.mb', h='k0.mb') -env.new_element('mq', xt.Quadrupole, length='l.mq') -env.new_element('ms', xt.Sextupole, length='l.ms') -env.new_element('corrector', xt.Multipole, knl=[0], ksl=[0]) +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_element('mq.f', 'mq', k1='kqf') -env.new_element('mq.d', 'mq', k1='kqd') +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_element('mid', xt.Marker, at='l.halfcell'), + env.new('mid', xt.Marker, at='l.halfcell'), # Bends - env.new_element('mb.2', 'mb', at='l.halfcell / 2'), - env.new_element('mb.1', 'mb', at='-l.mb - 1', from_='mb.2'), - env.new_element('mb.3', 'mb', at='l.mb + 1', from_='mb.2'), + 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_element('ms.d', 'ms', k2='k2sf', at=1.2, from_='mq.d'), - env.new_element('ms.f', 'ms', k2='k2sd', at=-1.2, from_='mq.f'), + 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_element('corrector.v', 'corrector', at=0.75, from_='mq.d'), - env.new_element('corrector.h', 'corrector', at=-0.75, from_='mq.f') + env.new('corrector.v', 'corrector', at=0.75, from_='mq.d'), + env.new('corrector.h', 'corrector', at=-0.75, from_='mq.f') ]) @@ -61,10 +61,10 @@ hcell_right = halfcell.replicate(name='r') cell = env.new_line(components=[ - env.new_element('start', xt.Marker), + env.new('start', xt.Marker), hcell_left, hcell_right, - env.new_element('end', xt.Marker), + env.new('end', xt.Marker), ]) opt = cell.match( @@ -84,22 +84,22 @@ halfcell_ss = env.new_line(components=[ - env.new_element('mid', xt.Marker, at='l.halfcell'), + env.new('mid', xt.Marker, at='l.halfcell'), - env.new_element('mq.ss.d', 'mq', k1='kqd.ss', at = '0.5 + l.mq / 2'), - env.new_element('mq.ss.f', 'mq', k1='kqf.ss', at = 'l.halfcell - l.mq / 2 - 0.5'), + 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_element('corrector.ss.v', 'corrector', at=0.75, from_='mq.ss.d'), - env.new_element('corrector.ss.h', 'corrector', at=-0.75, from_='mq.ss.f') + 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_element('start.ss', xt.Marker), + env.new('start.ss', xt.Marker), hcell_left_ss, hcell_right_ss, - env.new_element('end.ss', xt.Marker), + env.new('end.ss', xt.Marker), ]) opt = cell_ss.match( @@ -145,22 +145,22 @@ half_insertion = env.new_line(components=[ # Start-end markers - env.new_element('ip', xt.Marker), - env.new_element('e.insertion', xt.Marker, at=76), + env.new('ip', xt.Marker), + env.new('e.insertion', xt.Marker, at=76), # Quads - env.new_element('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), - env.new_element('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), - env.new_element('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), - env.new_element('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), - env.new_element('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq', at=73), + 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_element('corrector.ss.1', 'corrector', at=0.75, from_='mq.1'), - env.new_element('corrector.ss.2', 'corrector', at=-0.75, from_='mq.2'), - env.new_element('corrector.ss.3', 'corrector', at=0.75, from_='mq.3'), - env.new_element('corrector.ss.4', 'corrector', at=-0.75, from_='mq.4'), - env.new_element('corrector.ss.5', 'corrector', at=0.75, from_='mq.5'), + 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'), ]) diff --git a/examples/lattice_design_shortcuts/000a_subline.py b/examples/lattice_design_shortcuts/000a_subline.py index fb1e722ed..25d19cf21 100644 --- a/examples/lattice_design_shortcuts/000a_subline.py +++ b/examples/lattice_design_shortcuts/000a_subline.py @@ -25,10 +25,10 @@ 'l.halfcell': 38, }) -env.new_element('mb', xt.Bend, length='l.mb', k0='k0.mb', h='k0.mb') -env.new_element('mq', xt.Quadrupole, length='l.mq') -env.new_element('ms', xt.Sextupole, length='l.ms') -env.new_element('corrector', xt.Multipole, knl=[0], length=0.1) +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), @@ -49,12 +49,12 @@ halfcell = env.new_line(components=[ # End of the half cell (will be mid of the cell) - env.new_element('mid', xt.Marker, at='l.halfcell'), + env.new('mid', xt.Marker, at='l.halfcell'), # Bends - env.new_element('mb.2', 'mb', at='l.halfcell / 2'), - env.new_element('mb.1', 'mb', at='-l.mb - 1', from_='mb.2'), - env.new_element('mb.3', 'mb', at='l.mb + 1', from_='mb.2'), + 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), @@ -67,10 +67,10 @@ hcell_right = halfcell.replicate(name='r') cell = env.new_line(components=[ - env.new_element('start', xt.Marker), + env.new('start', xt.Marker), hcell_left, hcell_right, - env.new_element('end', xt.Marker), + env.new('end', xt.Marker), ]) opt = cell.match( @@ -90,22 +90,22 @@ halfcell_ss = env.new_line(components=[ - env.new_element('mid', xt.Marker, at='l.halfcell'), + env.new('mid', xt.Marker, at='l.halfcell'), - env.new_element('mq.ss.d', 'mq', k1='kqd.ss', at = '0.5 + l.mq / 2'), - env.new_element('mq.ss.f', 'mq', k1='kqf.ss', at = 'l.halfcell - l.mq / 2 - 0.5'), + 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_element('corrector.ss.v', 'corrector', at=0.75, from_='mq.ss.d'), - env.new_element('corrector.ss.h', 'corrector', at=-0.75, from_='mq.ss.f') + 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_element('start.ss', xt.Marker), + env.new('start.ss', xt.Marker), hcell_left_ss, hcell_right_ss, - env.new_element('end.ss', xt.Marker), + env.new('end.ss', xt.Marker), ]) opt = cell_ss.match( @@ -152,22 +152,22 @@ half_insertion = env.new_line(components=[ # Start-end markers - env.new_element('ip', xt.Marker), - env.new_element('e.insertion', xt.Marker, at=76), + env.new('ip', xt.Marker), + env.new('e.insertion', xt.Marker, at=76), # Quads - env.new_element('mq.1', xt.Quadrupole, k1='k1.q1', length='l.mq', at = 20), - env.new_element('mq.2', xt.Quadrupole, k1='k1.q2', length='l.mq', at = 25), - env.new_element('mq.3', xt.Quadrupole, k1='k1.q3', length='l.mq', at=37), - env.new_element('mq.4', xt.Quadrupole, k1='k1.q4', length='l.mq', at=55), - env.new_element('mq.5', xt.Quadrupole, k1='k1.q5', length='l.mq', at=73), + 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_element('corrector.ss.1', 'corrector', at=0.75, from_='mq.1'), - env.new_element('corrector.ss.2', 'corrector', at=-0.75, from_='mq.2'), - env.new_element('corrector.ss.3', 'corrector', at=0.75, from_='mq.3'), - env.new_element('corrector.ss.4', 'corrector', at=-0.75, from_='mq.4'), - env.new_element('corrector.ss.5', 'corrector', at=0.75, from_='mq.5'), + 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'), ]) diff --git a/examples/lattice_design_shortcuts/010_play.py b/examples/lattice_design_shortcuts/010_play.py index 5919261c1..9ccda1e03 100644 --- a/examples/lattice_design_shortcuts/010_play.py +++ b/examples/lattice_design_shortcuts/010_play.py @@ -2,15 +2,15 @@ env = xt.Environment() -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') +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', @@ -24,26 +24,26 @@ 'drift.4', ]) -mbx = env.new_element('mbxw', xt.Bend, k0='k0.mb', h=0, length='l.mbxw') +mbx = env.new('mbxw', xt.Bend, k0='k0.mb', h=0, length='l.mbxw') d1 = env.new_line('d1', components=[ - env.new_element('lmbxw.start', parent=xt.Marker), # shortcut env.new_element('lmbxw.start') - env.new_element('mbxw.a4@start', parent=xt.Replica, parent_name='mbxw', at=0.5), - env.new_element('mbxw.b4@start', parent='mbxw', from_='mbxw.a4@end'), - env.new_element('lmbxw.end', parent=xt.Marker, at=0.5, from_='mbxw.b4@end'), + 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_element('d2.b1', xt.Bend, k0='k0.mb', h=0, length='l.mbxw', dx=0.188/2) +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_element('ip1') + 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_element('ip') + env.new('ip') ir_left.replicate('.r1') ]) @@ -63,13 +63,13 @@ ]) seq = [ - Place(env.new_element('ip', xt.Marker), at=10), - Place(env.new_element('left', xt.Quadrupole, length=1), at=-5, from_='ip'), - env.new_element('after_left', xt.Marker), - env.new_element('after_left2', xt.Marker), - Place(env.new_element('right',xt.Quadrupole, length=1), at=+5, from_='ip'), - Place([env.new_element('before_right', xt.Quadrupole, length=1), - env.new_element('before_right2', xt.Marker)], at=0, from_='right@start', + 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_element('righter', xt.Quadrupole, length=1), at=+5, from_='before_right'), + Place(env.new('righter', xt.Quadrupole, length=1), at=+5, from_='before_right'), ] \ No newline at end of file diff --git a/xtrack/environment.py b/xtrack/environment.py index 85af9f3e8..0fa59addb 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -65,11 +65,11 @@ def _get_a_drift_name(self): else: return self._get_a_drift_name() - def new_element(self, name, cls, at=None, from_=None, **kwargs): + 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_element(name, cls, **kwargs)) + name=self.new(name, cls, **kwargs)) _eval = self._xdeps_eval.eval @@ -77,7 +77,7 @@ def new_element(self, name, cls, at=None, from_=None, **kwargs): 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.') + 'elements are allowed in `new` for now.') cls_input = cls if isinstance(cls, str): @@ -270,7 +270,7 @@ def _generate_element_names_with_drifts(env, tt_sorted, s_tol=1e-12): 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_element(drift_name, xt.Drift, length=ds_upstream) + env.new(drift_name, xt.Drift, length=ds_upstream) names_with_drifts.append(drift_name) names_with_drifts.append(nn) From 7798d86924bdaf0b1f641f5bb2f889d5b59388f5 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 11:13:08 +0200 Subject: [PATCH 113/159] Implement clone --- .../lattice_design_shortcuts/000a_subline.py | 9 ++------ xtrack/line.py | 21 ++++++++++++++----- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/examples/lattice_design_shortcuts/000a_subline.py b/examples/lattice_design_shortcuts/000a_subline.py index 25d19cf21..9ba1b745c 100644 --- a/examples/lattice_design_shortcuts/000a_subline.py +++ b/examples/lattice_design_shortcuts/000a_subline.py @@ -36,16 +36,11 @@ env.place('corrector', at=-0.8, from_='mq'), ]) -girder_f = girder.replicate(name='f') -girder_f.replace_all_replicas() - -girder_d = girder.replicate(name='d', mirror=True) -girder_d.replace_all_replicas() - +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) diff --git a/xtrack/line.py b/xtrack/line.py index aa6a860af..058a9c277 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3355,6 +3355,9 @@ def mirror(self): 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 @@ -3371,6 +3374,11 @@ def replicate(self, name, mirror=False): 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__ @@ -3407,16 +3415,19 @@ def select(self, start=None, end=None, name=None): tt = self.get_table().rows[start:end] if tt.name[-1] == '_end_point': tt = tt.rows[:-1] - 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) + + self._env_if_needed() out = self.env.new_line(components=list(tt.name), name=name) return out + 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) From 8fc6fef08328d6036d990cca1eb30f82f255f723 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 16:32:36 +0200 Subject: [PATCH 114/159] A few checks on the variables --- .../lattice_design_shortcuts/t000_test_env.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 examples/lattice_design_shortcuts/t000_test_env.py 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..9373adec0 --- /dev/null +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -0,0 +1,24 @@ +import xtrack as xt +import xobjects as xo + +env = xt.Environment() + +env.vars({ + 'k.1': 1., + 'a': 2., + 'b': '2 * a + k.1', +}) + +assert env.vv['b'] == 2 * 2 + 1 + +env.vars['a'] = env.vars['k.1'] +assert env.vv['b'] == 2 * 1 + 1 + +env.vars(a=3.) +env.vars({'k.1': 'a'}) +assert env.vv['k.1'] == 3. +assert env.vv['b'] == 2 * 3 + 3. + +# line = env.new_line([ +# env.new +# ] \ No newline at end of file From dfb533e877e980c3d6d3d5180819d8cb181a5a48 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 16:46:12 +0200 Subject: [PATCH 115/159] Status --- .../lattice_design_shortcuts/t000_test_env.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index 9373adec0..b957fcdc5 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -19,6 +19,18 @@ assert env.vv['k.1'] == 3. assert env.vv['b'] == 2 * 3 + 3. -# line = env.new_line([ -# env.new -# ] \ No newline at end of file +env.vars['k.1'] = 2 * env.vars['a'] + 5 +assert env.vv['k.1'] == 2 * 3 + 5 +assert env.vv['b'] == 2 * 3 + 2 * 3 + 5 + + +env.vars({ + 'a': 4., + 'b': '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. \ No newline at end of file From 8d448fa587db279befa9aee397bd6b32b9ce5baa Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 16:57:14 +0200 Subject: [PATCH 116/159] More checks --- .../lattice_design_shortcuts/t000_test_env.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index b957fcdc5..5611879a5 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -33,4 +33,18 @@ h=5.) assert env['bb'].k0 == 2 * (2 * 4 + 5) assert env['bb'].length == 3 + 4 + 2 * 4 + 5 -assert env['bb'].h == 5. \ No newline at end of file +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') +]) + +assert line['bb1'].length == 6 +assert line['bb1'].k0 == 2 * (2 * 2 + 5) +assert line['bb1'].h == 5. + From b48cfd85acaf29ae7f34b4e86d460cc742586b16 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 17:12:55 +0200 Subject: [PATCH 117/159] Status --- examples/lattice_design_shortcuts/t000_test_env.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index 5611879a5..d5e73628f 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -41,10 +41,17 @@ assert env['bb'].h == 5. line = env.new_line([ - env.new('bb1', 'bb', length=3*env.vars['a'], at='2*a') + env.new('bb1', 'bb', length=3*env.vars['a'], at='2*a'), + env.place('bb', at=10 * env.vars['a']) ]) assert line['bb1'].length == 6 assert line['bb1'].k0 == 2 * (2 * 2 + 5) assert line['bb1'].h == 5. +assert env['bb'].k0 == 2 * (2 * 2 + 5) +assert env['bb'].length == 3 + 2 + 2 * 2 + 5 +assert env['bb'].h == 5. + + + From 0eac842879c423b4dfc48e598f39ad67fb3ff94a Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 17:20:08 +0200 Subject: [PATCH 118/159] Checks --- .../lattice_design_shortcuts/t000_test_env.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index d5e73628f..02d07ccfc 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -42,16 +42,28 @@ line = env.new_line([ env.new('bb1', 'bb', length=3*env.vars['a'], at='2*a'), - env.place('bb', at=10 * env.vars['a']) + env.place('bb', at=10 * env.vars['a'], from_='bb1'), ]) +assert line['bb1'] is not env['bb'] +assert line['bb'] is env['bb'] + assert line['bb1'].length == 6 assert line['bb1'].k0 == 2 * (2 * 2 + 5) assert line['bb1'].h == 5. -assert env['bb'].k0 == 2 * (2 * 2 + 5) -assert env['bb'].length == 3 + 2 + 2 * 2 + 5 -assert env['bb'].h == 5. +assert line['bb'].k0 == 2 * (2 * 2 + 5) +assert line['bb'].length == 3 + 2 + 2 * 2 + 5 +assert line['bb'].h == 5. + +tt = line.get_table(attr=True) +tt['s_center'] = tt['s'] + tt['length']/2 + +a = env.vv['a'] +assert tt['s_center', 'bb1'] == 2*a +assert tt['s_center', 'bb'] - tt['s_center', 'bb1'] == 10*a + + From 25bbeff5f39e1a8c67d3498c7f5624d16e85690f Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 18:47:35 +0200 Subject: [PATCH 119/159] Finish a test --- .../lattice_design_shortcuts/t000_test_env.py | 30 +++++++++++++++---- 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index 02d07ccfc..c0e31549c 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -1,5 +1,5 @@ import xtrack as xt -import xobjects as xo +import numpy as np env = xt.Environment() @@ -48,21 +48,39 @@ assert line['bb1'] is not env['bb'] assert line['bb'] is env['bb'] -assert line['bb1'].length == 6 -assert line['bb1'].k0 == 2 * (2 * 2 + 5) +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 * 2 + 5) -assert line['bb'].length == 3 + 2 + 2 * 2 + 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 -a = env.vv['a'] +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'] From 2e408d9f31ce4ae69e45dc13ccfa8dcbf90fa038 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 20:53:38 +0200 Subject: [PATCH 120/159] Add LineVars.set --- examples/lattice_design_shortcuts/t000_test_env.py | 3 +++ xtrack/line.py | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index c0e31549c..d063b3c5c 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -23,6 +23,9 @@ assert env.vv['k.1'] == 2 * 3 + 5 assert env.vv['b'] == 2 * 3 + 2 * 3 + 5 +env.vars.set('a', 4.) +assert env.vv['k.1'] == 2 * 4 + 5 +assert env.vv['b'] == 2 * 4 + 2 * 4 + 5 env.vars({ 'a': 4., diff --git a/xtrack/line.py b/xtrack/line.py index 058a9c277..c04587f1c 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -4686,6 +4686,12 @@ def __call__(self, *args, **kwargs): 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): From bec607ee9f6802d6f7ec82c3403651f50b73af9a Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 20:56:53 +0200 Subject: [PATCH 121/159] More checks --- examples/lattice_design_shortcuts/t000_test_env.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index d063b3c5c..ac9a54085 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -27,6 +27,15 @@ assert env.vv['k.1'] == 2 * 4 + 5 assert env.vv['b'] == 2 * 4 + 2 * 4 + 5 +env.vars.set('k.1', '2*a + 5') +assert env.vv['k.1'] == 2 * 4 + 5 +assert env.vv['b'] == 2 * 4 + 2 * 4 + 5 + +env.vars.set('k.1', 3 * env.vars['a'] + 6) +assert env.vv['k.1'] == 3 * 4 + 6 +assert env.vv['b'] == 2 * 4 + 3 * 4 + 6 + + env.vars({ 'a': 4., 'b': '2 * a + 5', From 23104576e39d477708c4970efd57019f0be12318 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 21:45:59 +0200 Subject: [PATCH 122/159] Checks --- .../t001_check_with_subline.py | 244 ++++++++++++++++++ 1 file changed, 244 insertions(+) create mode 100644 examples/lattice_design_shortcuts/t001_check_with_subline.py 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..b26903a4d --- /dev/null +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -0,0 +1,244 @@ +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) +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-15, rtol=0) +xo.assert_allclose(tt_girder['s_center', 'ms'] - tt_girder['s_center', 'mq'], 0.8, + atol=1e-15, rtol=0) +xo.assert_allclose( + tt_girder['s_center', 'corrector'] - tt_girder['s_center', 'mq'], -0.8, + atol=1e-15, rtol=0) + +prrr + +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() + + + From 07e27bf1df0058d0219794412f6c7ab86a1555be Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 12 Sep 2024 21:51:54 +0200 Subject: [PATCH 123/159] Working on checks --- .../t001_check_with_subline.py | 23 +++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index b26903a4d..315621d88 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -45,14 +45,33 @@ xo.assert_allclose( tt_girder['s_center', 'corrector'] - tt_girder['s_center', 'mq'], -0.8, atol=1e-15, rtol=0) - -prrr +assert np.all(tt_girder.name == np.array( + ['drift_1', 'corrector', 'drift_2', 'mq', 'drift_3', 'ms', '_end_point'])) 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) +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-15, rtol=0) +xo.assert_allclose(tt_girder_f['s_center', 'ms.f'] - tt_girder_f['s_center', 'mq.f'], 0.8, + atol=1e-15, rtol=0) +xo.assert_allclose( + tt_girder_f['s_center', 'corrector.f'] - tt_girder_f['s_center', 'mq.f'], -0.8, + atol=1e-15, rtol=0) + +# Check clone mirror +... + + + + +prrrr + halfcell = env.new_line(components=[ # End of the half cell (will be mid of the cell) From 45f4896849d6146e6fd39630f827f76de2fe2f59 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 06:04:43 +0200 Subject: [PATCH 124/159] More checks --- .../t001_check_with_subline.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 315621d88..4fa867770 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -38,6 +38,8 @@ ]) 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-15, rtol=0) xo.assert_allclose(tt_girder['s_center', 'ms'] - tt_girder['s_center', 'mq'], 0.8, @@ -45,8 +47,7 @@ xo.assert_allclose( tt_girder['s_center', 'corrector'] - tt_girder['s_center', 'mq'], -0.8, atol=1e-15, rtol=0) -assert np.all(tt_girder.name == np.array( - ['drift_1', 'corrector', 'drift_2', 'mq', 'drift_3', 'ms', '_end_point'])) + girder_f = girder.clone(name='f') girder_d = girder.clone(name='d', mirror=True) @@ -55,6 +56,8 @@ # Check clone tt_girder_f = girder_f.get_table(attr=True) +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-15, rtol=0) @@ -65,8 +68,17 @@ atol=1e-15, rtol=0) # Check clone mirror -... - +tt_girder_d = girder_d.get_table(attr=True) +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-15, rtol=0) +xo.assert_allclose(tt_girder_d['s_center', 'ms.d'] - tt_girder_d['s_center', 'mq.d'], + -0.8, atol=1e-15, rtol=0) +xo.assert_allclose(tt_girder_d['s_center', 'corrector.d'] - tt_girder_d['s_center', 'mq.d'], + 0.8, atol=1e-15, rtol=0) From c2150ea178debdd84789d5133ffd0e88b81ac5ca Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 06:09:18 +0200 Subject: [PATCH 125/159] Fix --- .../lattice_design_shortcuts/t001_check_with_subline.py | 9 ++++++--- xtrack/environment.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 4fa867770..5fe6e4d2c 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -81,9 +81,6 @@ 0.8, atol=1e-15, rtol=0) - -prrrr - halfcell = env.new_line(components=[ # End of the half cell (will be mid of the cell) @@ -100,6 +97,12 @@ ]) +l_hc = env.vv['l.halfcell'] +xo.assert_allclose(l_hc, l_hc, atol=1e-15, rtol=0) +tt_hc = halfcell.get_table(attr=True) + +prrrr + hcell_left = halfcell.replicate(name='l', mirror=True) hcell_right = halfcell.replicate(name='r') diff --git a/xtrack/environment.py b/xtrack/environment.py index 0fa59addb..a5effb674 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -9,7 +9,7 @@ def _flatten_components(components): for nn in components: if isinstance(nn, Place) and isinstance(nn.name, xt.Line): line = nn.name - components = line.element_names.copy() + components = list(line.element_names).copy() if nn.at is not None: if isinstance(nn.at, str): at = line._xdeps_eval.eval(nn.at) From be8c91430987a2a03b343cf796f535273beb1130 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 06:15:52 +0200 Subject: [PATCH 126/159] Another check --- .../lattice_design_shortcuts/t001_check_with_subline.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 5fe6e4d2c..cb0bd30a2 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -100,6 +100,15 @@ l_hc = env.vv['l.halfcell'] xo.assert_allclose(l_hc, l_hc, atol=1e-15, 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'])) +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-15, rtol=0) prrrr From 6762d570a7448df7285fdd089d019638c7fff242 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 06:21:38 +0200 Subject: [PATCH 127/159] More checks --- .../t001_check_with_subline.py | 32 ++++++++++++------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index cb0bd30a2..ee92147f6 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -41,12 +41,12 @@ 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-15, rtol=0) +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-15, rtol=0) + atol=1e-14, rtol=0) xo.assert_allclose( tt_girder['s_center', 'corrector'] - tt_girder['s_center', 'mq'], -0.8, - atol=1e-15, rtol=0) + atol=1e-14, rtol=0) girder_f = girder.clone(name='f') @@ -60,12 +60,12 @@ ['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-15, rtol=0) +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-15, rtol=0) + 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-15, rtol=0) + atol=1e-14, rtol=0) # Check clone mirror tt_girder_d = girder_d.get_table(attr=True) @@ -74,11 +74,11 @@ ['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-15, rtol=0) +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-15, rtol=0) + -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-15, rtol=0) + 0.8, atol=1e-14, rtol=0) halfcell = env.new_line(components=[ @@ -98,7 +98,7 @@ ]) l_hc = env.vv['l.halfcell'] -xo.assert_allclose(l_hc, l_hc, atol=1e-15, rtol=0) +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', @@ -108,8 +108,16 @@ 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-15, rtol=0) - + 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) prrrr From 08f2364eee8ebbe9f7738ccdd7578c7a2094ba26 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 06:45:26 +0200 Subject: [PATCH 128/159] More checks --- .../t001_check_with_subline.py | 78 +++++++++++++++---- 1 file changed, 62 insertions(+), 16 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index ee92147f6..2d193d4f4 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -40,13 +40,14 @@ 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']) +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) + tt_girder['s_center', 'corrector'] - tt_girder['s_center', 'mq'], -0.8, + atol=1e-14, rtol=0) girder_f = girder.clone(name='f') @@ -59,13 +60,14 @@ 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'])) + + 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) + 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) + 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) @@ -73,8 +75,9 @@ 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) + + 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'], @@ -105,20 +108,32 @@ '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'])) -tt_hc['s_center'] = tt_hc['s'] + tt_hc['length']/2 * np.float64(tt_hc['isthick']) +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) + 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) + 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) + 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) -prrrr +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) @@ -131,6 +146,37 @@ env.new('end', xt.Marker), ]) +tt_cell = cell.get_table(attr=True) +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])) + +prrrr + opt = cell.match( method='4d', vary=xt.VaryList(['kqf', 'kqd'], step=1e-5), From 75cdc3f39404e0e7527e18b39965b1cbf48cc4e2 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 07:03:00 +0200 Subject: [PATCH 129/159] Check second half of cell --- examples/lattice_design_shortcuts/t001_check_with_subline.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 2d193d4f4..1f8e484c4 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -175,6 +175,11 @@ 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 -= tt_cell_second_half.s[0] +tt_hc_stripped = tt_hc.rows[:-1] # Remove _end_point +xo.assert_allclose(tt_cell_second_half.s, tt_hc_stripped.s, atol=1e-14, rtol=0) prrrr opt = cell.match( From 2688d739227410bee4f21332b0dad5f73c7b90f8 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 13 Sep 2024 08:53:53 +0200 Subject: [PATCH 130/159] _sigma table return dict --- xtrack/twiss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 00bec4f03..00c8239d7 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3675,7 +3675,7 @@ def _build_sigma_table(Sigma, s=None, name=None): res_data['Sigma52'] = Sigma[:, 4, 1] - return Table(res_data) + return res_data def compute_T_matrix_line(line, start, end, particle_on_co=None, steps_t_matrix=None): From 25af20c9539b56d8f74d1704f6ec45fec547a33a Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 13 Sep 2024 10:06:08 +0200 Subject: [PATCH 131/159] new env syntax --- .../t010_new_syntax.py | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) create mode 100644 examples/lattice_design_shortcuts/t010_new_syntax.py 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 + + From 9d97878a100b936b30c3170951b6440dcf814301 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 13 Sep 2024 10:12:31 +0200 Subject: [PATCH 132/159] refix table --- xtrack/twiss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 00c8239d7..00bec4f03 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3675,7 +3675,7 @@ def _build_sigma_table(Sigma, s=None, name=None): res_data['Sigma52'] = Sigma[:, 4, 1] - return res_data + return Table(res_data) def compute_T_matrix_line(line, start, end, particle_on_co=None, steps_t_matrix=None): From 1263ba1fab18debcaa3b3e6a116aebfaa373c742 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 10:56:16 +0200 Subject: [PATCH 133/159] Keep track of named lines --- .../lattice_design_shortcuts/t001_check_with_subline.py | 4 ++++ xtrack/environment.py | 5 ++++- xtrack/line.py | 7 ++----- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 1f8e484c4..aa62ec5db 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -180,6 +180,10 @@ tt_cell_second_half.s -= tt_cell_second_half.s[0] tt_hc_stripped = tt_hc.rows[:-1] # Remove _end_point xo.assert_allclose(tt_cell_second_half.s, tt_hc_stripped.s, atol=1e-14, rtol=0) +tt_cell_first_half = tt_cell_stripped.rows[:len(tt_cell_stripped)//2] +s_mirrored_first_half = tt_cell_first_half.s[::-1] - tt_cell_first_half.length[::-1] +s_mirrored_first_half -= s_mirrored_first_half[0] +xo.assert_allclose(s_mirrored_first_half, tt_hc_stripped.s, atol=1e-14, rtol=0) prrrr opt = cell.match( diff --git a/xtrack/environment.py b/xtrack/environment.py index a5effb674..905947d13 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -34,6 +34,7 @@ def __init__(self, element_dict=None, particle_ref=None, _var_management=None): else: self._init_var_management() + self.lines = {} self._lines = WeakSet() self._drift_counter = 0 @@ -48,7 +49,9 @@ def new_line(self, components=None, name=None): out.element_names = handle_s_places(flattened_components, self) out._var_management = self._var_management out._name = name - self._lines.add(out) + self._lines.add(out) # Weak references + if name is not None: + self.lines[name] = out return out diff --git a/xtrack/line.py b/xtrack/line.py index c04587f1c..6b2fd372f 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3363,11 +3363,8 @@ def replicate(self, name, mirror=False): new_nn = nn + '.' + name self.element_dict[new_nn] = xt.Replica(nn) new_element_names.append(new_nn) - out = Line() - out.element_names = new_element_names - out._element_dict = self.element_dict # to make sure that the dict is not copied - out._var_management = self._var_management - out._name = name + + out = self.env.new_line(components=new_element_names, name=name) if mirror: out.mirror() From 866d786d2d02f91ac9e7277cb080b6af97dba330 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 11:15:24 +0200 Subject: [PATCH 134/159] Added setitem --- xtrack/line.py | 28 ++++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 6b2fd372f..0441f8dd6 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3774,20 +3774,32 @@ 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 + elif np.isscalar(value): + self.vars[key] = value + else: + self.element_dict[key] = value + def _get_non_collective_line(self): if not self.iscollective: return self From 0968a7091d5eedb634fe1ad71198abd8f2fa4ce6 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 11:19:38 +0200 Subject: [PATCH 135/159] __setitem__ in env --- xtrack/environment.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xtrack/environment.py b/xtrack/environment.py index 905947d13..5063c7a6d 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -124,6 +124,7 @@ def place(self, name, at=None, from_=None, anchor=None, from_anchor=None): Environment.vv = xt.Line.vv Environment.replace_replica = xt.Line.replace_replica Environment.__getitem__ = xt.Line.__getitem__ +Environment.__setitem__ = xt.Line.__setitem__ class Place: From a40e262e09a1351beac22ff793090fbd5169b42c Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 13:36:10 +0200 Subject: [PATCH 136/159] Add Line.ref and Environment.ref --- xtrack/environment.py | 34 ++++++++++++++++++++++++++++++++++ xtrack/line.py | 6 ++++++ 2 files changed, 40 insertions(+) diff --git a/xtrack/environment.py b/xtrack/environment.py index 5063c7a6d..5af3fb744 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -37,6 +37,7 @@ def __init__(self, element_dict=None, particle_ref=None, _var_management=None): self.lines = {} self._lines = WeakSet() self._drift_counter = 0 + self.ref = EnvRef(self) def new_line(self, components=None, name=None): out = xt.Line() @@ -340,3 +341,36 @@ def _set_kwargs(name, ref_kwargs, value_kwargs, element_dict, element_refs): 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/line.py b/xtrack/line.py index 0441f8dd6..2705b94ae 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -151,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=()): @@ -3792,12 +3793,17 @@ def __getitem__(self, key): 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 elif 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('There is already a variable with name {key}') self.element_dict[key] = value def _get_non_collective_line(self): From 7766bae19bf0c58a6e534105c851be3841e58c1d Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 15:30:10 +0200 Subject: [PATCH 137/159] Extend set --- xtrack/environment.py | 10 +--------- xtrack/line.py | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/xtrack/environment.py b/xtrack/environment.py index 5af3fb744..45f91b5b7 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -101,15 +101,6 @@ def new(self, name, cls, at=None, from_=None, **kwargs): return name - def set(self, name, **kwargs): - _eval = self._xdeps_eval.eval - - ref_kwargs, value_kwargs = _parse_kwargs( - type(self.element_dict[name]), kwargs, _eval) - - _set_kwargs(name=name, ref_kwargs=ref_kwargs, value_kwargs=value_kwargs, - element_dict=self.element_dict, element_refs=self.element_refs) - 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) @@ -126,6 +117,7 @@ def place(self, name, at=None, from_=None, anchor=None, from_anchor=None): Environment.replace_replica = xt.Line.replace_replica Environment.__getitem__ = xt.Line.__getitem__ Environment.__setitem__ = xt.Line.__setitem__ +Environment.set = xt.Line.set class Place: diff --git a/xtrack/line.py b/xtrack/line.py index 2705b94ae..63ad9f39b 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3420,6 +3420,32 @@ def select(self, start=None, end=None, name=None): return out + def set(self, name, *args, **kwargs): + _eval = self._xdeps_eval.eval + + 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) + elif name in self.vars: + 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[name] = _eval(value) + else: + self[name] = value + elif hasattr(self, 'lines') and name in self.lines: + raise ValueError('Cannot set a line') + else: + raise KeyError(f'Name {name} not found.') + def _env_if_needed(self): if not hasattr(self, 'env') or self.env is None: self.env = xt.Environment(element_dict=self.element_dict, From 4fd71bbd2626657af794923f9f8322c271720c1e Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 16:25:10 +0200 Subject: [PATCH 138/159] Add more checks --- examples/lattice_design_shortcuts/t000_test_env.py | 14 ++++++++++++++ xtrack/line.py | 4 ++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index ac9a54085..3e70aadb0 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -35,10 +35,24 @@ assert env.vv['k.1'] == 3 * 4 + 6 assert env.vv['b'] == 2 * 4 + 3 * 4 + 6 +env.set('a', 0.) +assert env.vv['k.1'] == 3 * 0 + 6 +assert env.vv['b'] == 2 * 0 + 3 * 0 + 6 + +env.set('a', 2.) +env.set('k.1', '2 * a + 5') +assert env.vv['k.1'] == 2 * 2 + 5 +assert env.vv['b'] == 2 * 2 + 2 * 2 + 5 + +env.set('k.1', 3 * env.vars['a'] + 6) +assert env.vv['k.1'] == 3 * 2 + 6 +assert env.vv['b'] == 2 * 2 + 3 * 2 + 6 + env.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'], diff --git a/xtrack/line.py b/xtrack/line.py index 63ad9f39b..074e7a88e 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3438,9 +3438,9 @@ def set(self, name, *args, **kwargs): raise ValueError(f'A value must be provided when setting a variable') value = args[0] if isinstance(value, str): - self[name] = _eval(value) + self.vars[name] = _eval(value) else: - self[name] = value + self.vars[name] = value elif hasattr(self, 'lines') and name in self.lines: raise ValueError('Cannot set a line') else: From cc81843deea38b0c14af488ee9c32703ff21e781 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 13 Sep 2024 17:47:02 +0200 Subject: [PATCH 139/159] more robust --- xtrack/twiss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 00bec4f03..fae84b9df 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3286,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 From 4d3576c8405d2dfefe6633cf0557e3cadf9e701e Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 18:36:38 +0200 Subject: [PATCH 140/159] Update --- examples/lattice_design_shortcuts/t000_test_env.py | 3 --- xtrack/line.py | 8 ++++++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index 3e70aadb0..1ae9dae08 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -108,6 +108,3 @@ tt_new['length', 'drift_1'] == tt['length', 'drift_1'] tt_new['length', 'drift_2'] == tt['length', 'drift_2'] - - - diff --git a/xtrack/line.py b/xtrack/line.py index 074e7a88e..9c312e5a8 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3823,13 +3823,17 @@ 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 - elif np.isscalar(value): + + 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('There is already a variable with name {key}') + raise ValueError(f'There is already a variable with name {key}') self.element_dict[key] = value def _get_non_collective_line(self): From b32a52f890e18839484a40dc82c0e8e0a69c442b Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 18:50:07 +0200 Subject: [PATCH 141/159] More checks --- examples/lattice_design_shortcuts/t000_test_env.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index 1ae9dae08..4e5e7364b 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -48,6 +48,16 @@ assert env.vv['k.1'] == 3 * 2 + 6 assert env.vv['b'] == 2 * 2 + 3 * 2 + 6 +assert hasattr(env.ref['k.1'], '_value') # is a Ref + +env.ref['a'] = 0 +assert env.vv['k.1'] == 3 * 0 + 6 +assert env.vv['b'] == 2 * 0 + 3 * 0 + 6 + +env.ref['a'] = 2 +env.ref['k.1'] = 2 * env.ref['a'] + 5 +assert env.vv['k.1'] == 2 * 2 + 5 +assert env.vv['b'] == 2 * 2 + 2 * 2 + 5 env.vars({ 'a': 4., From 73a47997250d45582bd69ec960611230035bd1f4 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 18:54:25 +0200 Subject: [PATCH 142/159] Status --- .../lattice_design_shortcuts/t000_test_env.py | 83 ++++++++++--------- 1 file changed, 45 insertions(+), 38 deletions(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index 4e5e7364b..036ecd4bc 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -9,62 +9,69 @@ 'b': '2 * a + k.1', }) -assert env.vv['b'] == 2 * 2 + 1 +line = env.new_line([]) -env.vars['a'] = env.vars['k.1'] -assert env.vv['b'] == 2 * 1 + 1 +ee = env # Test Environment +ee = line # Test Line -env.vars(a=3.) -env.vars({'k.1': 'a'}) -assert env.vv['k.1'] == 3. -assert env.vv['b'] == 2 * 3 + 3. +assert ee.vv['b'] == 2 * 2 + 1 -env.vars['k.1'] = 2 * env.vars['a'] + 5 -assert env.vv['k.1'] == 2 * 3 + 5 -assert env.vv['b'] == 2 * 3 + 2 * 3 + 5 +ee.vars['a'] = ee.vars['k.1'] +assert ee.vv['b'] == 2 * 1 + 1 -env.vars.set('a', 4.) -assert env.vv['k.1'] == 2 * 4 + 5 -assert env.vv['b'] == 2 * 4 + 2 * 4 + 5 +ee.vars(a=3.) +ee.vars({'k.1': 'a'}) +assert ee.vv['k.1'] == 3. +assert ee.vv['b'] == 2 * 3 + 3. -env.vars.set('k.1', '2*a + 5') -assert env.vv['k.1'] == 2 * 4 + 5 -assert env.vv['b'] == 2 * 4 + 2 * 4 + 5 +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 -env.vars.set('k.1', 3 * env.vars['a'] + 6) -assert env.vv['k.1'] == 3 * 4 + 6 -assert env.vv['b'] == 2 * 4 + 3 * 4 + 6 +ee.vars.set('a', 4.) +assert ee.vv['k.1'] == 2 * 4 + 5 +assert ee.vv['b'] == 2 * 4 + 2 * 4 + 5 -env.set('a', 0.) -assert env.vv['k.1'] == 3 * 0 + 6 -assert env.vv['b'] == 2 * 0 + 3 * 0 + 6 +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 -env.set('a', 2.) -env.set('k.1', '2 * a + 5') -assert env.vv['k.1'] == 2 * 2 + 5 -assert env.vv['b'] == 2 * 2 + 2 * 2 + 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('k.1', 3 * env.vars['a'] + 6) -assert env.vv['k.1'] == 3 * 2 + 6 -assert env.vv['b'] == 2 * 2 + 3 * 2 + 6 +ee.set('a', 0.) +assert ee.vv['k.1'] == 3 * 0 + 6 +assert ee.vv['b'] == 2 * 0 + 3 * 0 + 6 -assert hasattr(env.ref['k.1'], '_value') # is a Ref +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 -env.ref['a'] = 0 -assert env.vv['k.1'] == 3 * 0 + 6 -assert env.vv['b'] == 2 * 0 + 3 * 0 + 6 +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 -env.ref['a'] = 2 -env.ref['k.1'] = 2 * env.ref['a'] + 5 -assert env.vv['k.1'] == 2 * 2 + 5 -assert env.vv['b'] == 2 * 2 + 2 * 2 + 5 +assert hasattr(ee.ref['k.1'], '_value') # is a Ref -env.vars({ +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) From 2193ad6da7c73c7f53d1feae351a154f077c803f Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 19:00:59 +0200 Subject: [PATCH 143/159] More checks --- .../lattice_design_shortcuts/t000_test_env.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/examples/lattice_design_shortcuts/t000_test_env.py b/examples/lattice_design_shortcuts/t000_test_env.py index 036ecd4bc..ca28b0d47 100644 --- a/examples/lattice_design_shortcuts/t000_test_env.py +++ b/examples/lattice_design_shortcuts/t000_test_env.py @@ -1,6 +1,9 @@ import xtrack as xt import numpy as np +to_test = 'line' +to_test = 'env' + env = xt.Environment() env.vars({ @@ -11,8 +14,7 @@ line = env.new_line([]) -ee = env # Test Environment -ee = line # Test Line +ee = {'env': env, 'line': line}[to_test] assert ee.vv['b'] == 2 * 2 + 1 @@ -64,14 +66,14 @@ 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) @@ -88,6 +90,16 @@ 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'] From ccd48656013e38206c90d69fbb97860ec1663809 Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Fri, 13 Sep 2024 19:09:49 +0200 Subject: [PATCH 144/159] plot more robust --- xtrack/twiss.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xtrack/twiss.py b/xtrack/twiss.py index fae84b9df..143d5958b 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -3417,6 +3417,9 @@ def plot(self, if yr is None: yr="" + if not hasattr(self,"_action"): + lattice=False + if lattice and 'length' not in self.keys(): self.add_strengths() From ed94a326263b4d0d4a61139b956969242c463663 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 19:34:04 +0200 Subject: [PATCH 145/159] Checks --- .../t001_check_with_subline.py | 13 +++++++------ xtrack/twiss.py | 2 +- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index aa62ec5db..fd841e31b 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -147,6 +147,8 @@ ]) 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', @@ -177,14 +179,13 @@ 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 -= tt_cell_second_half.s[0] +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, tt_hc_stripped.s, atol=1e-14, rtol=0) +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_mirrored_first_half = tt_cell_first_half.s[::-1] - tt_cell_first_half.length[::-1] -s_mirrored_first_half -= s_mirrored_first_half[0] -xo.assert_allclose(s_mirrored_first_half, tt_hc_stripped.s, atol=1e-14, rtol=0) -prrrr +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) opt = cell.match( method='4d', diff --git a/xtrack/twiss.py b/xtrack/twiss.py index fae84b9df..f1670b3d2 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -58,7 +58,7 @@ 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'] +SIGN_FLIP_FOR_ATTR_REVERSE=['k0l', 'k2l', 'k4l', 'k1sl', 'k3sl', 'k5sl', 'vkick', 'angle_rad'] log = logging.getLogger(__name__) From 141c03d8da0f6f72dea0d3d9f0bdad92f5a51a55 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 19:49:15 +0200 Subject: [PATCH 146/159] Match at the end --- .../t001_check_with_subline.py | 76 ++++++++++--------- 1 file changed, 41 insertions(+), 35 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index fd841e31b..5d50d9f2a 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -187,14 +187,7 @@ 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) -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({ @@ -222,14 +215,7 @@ 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=[ @@ -285,26 +271,7 @@ ]) -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), @@ -326,6 +293,45 @@ 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() + import matplotlib.pyplot as plt plt.close('all') From a80de3f649a8fca474ed69fc5ab5f2f83147813c Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 20:25:10 +0200 Subject: [PATCH 147/159] More checks --- .../t001_check_with_subline.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 5d50d9f2a..64d03c292 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -57,6 +57,7 @@ # 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'] @@ -71,6 +72,7 @@ # 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'])) @@ -187,9 +189,6 @@ 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, @@ -216,14 +215,23 @@ ]) - - 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' +prrr ss = env.new_line(components=[ cell_ss.replicate('cell.1'), From 0e85b7da24aa97128d76812b4ed3b14f03f6794d Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 21:52:12 +0200 Subject: [PATCH 148/159] Some more checks --- .../t001_check_with_subline.py | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 64d03c292..0dc25e100 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -231,7 +231,14 @@ 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' -prrr + +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'), @@ -246,6 +253,19 @@ 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) + +prrrr ## Insertion From f07379788c6adf1a897d21586b9215bae189aae1 Mon Sep 17 00:00:00 2001 From: giadarol Date: Fri, 13 Sep 2024 22:15:57 +0200 Subject: [PATCH 149/159] More checks --- .../t001_check_with_subline.py | 29 ++++++++++--------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 0dc25e100..3b1d49f15 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -265,8 +265,6 @@ xo.assert_allclose(sv_ring.angle.sum(), 2*np.pi, atol=1e-12, rtol=0) -prrrr - ## Insertion env.vars({ @@ -299,21 +297,17 @@ ]) - - 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'), + env['arc.1'], + env['ss.1'], + env['arc.2'], insertion, - arc.replicate(name='arcc.3'), - ss.replicate(name='sss.3') + env['arc.3'], + env['ss.3'], ]) @@ -360,18 +354,25 @@ opt.step(40) opt.solve() +tw = ring2.twiss4d() + +# Check that the cell is matched to the rest of the ring +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) + import matplotlib.pyplot as plt plt.close('all') for ii, rr in enumerate([ring, ring2_sliced]): - tw = rr.twiss4d() + ttww = 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) + pltbet = ttww.plot('betx bety', ax=ax1) ax2 = fig.add_subplot(2, 1, 2, sharex=ax1) - pltdx = tw.plot('dx', ax=ax2) + pltdx = ttww.plot('dx', ax=ax2) fig.subplots_adjust(right=.85) pltbet.move_legend(1.2,1) pltdx.move_legend(1.2,1) From fb0ca153af11a775c2742ec7e32639d21c4ee008 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 14 Sep 2024 07:14:27 +0200 Subject: [PATCH 150/159] Some more checks --- .../t001_check_with_subline.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 3b1d49f15..7cb2eed68 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -361,6 +361,27 @@ 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) +# 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() + import matplotlib.pyplot as plt plt.close('all') From b90bde7e4352d058657e7ed86da3b6c9afc3fceb Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 14 Sep 2024 07:23:29 +0200 Subject: [PATCH 151/159] More checks --- .../t001_check_with_subline.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/examples/lattice_design_shortcuts/t001_check_with_subline.py b/examples/lattice_design_shortcuts/t001_check_with_subline.py index 7cb2eed68..e78f765fc 100644 --- a/examples/lattice_design_shortcuts/t001_check_with_subline.py +++ b/examples/lattice_design_shortcuts/t001_check_with_subline.py @@ -354,13 +354,17 @@ opt.step(40) opt.solve() -tw = ring2.twiss4d() - # 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') @@ -382,6 +386,15 @@ 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') From 3c18e261ea8348b731430eab11b8693b9f5ddc11 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 14 Sep 2024 08:28:28 +0200 Subject: [PATCH 152/159] Another check --- .../lattice_design_shortcuts/001_dev_at_s.py | 44 ++++++++++++------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index bac9c8e26..05bb88b5d 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -1,12 +1,7 @@ import xtrack as xt +import xobjects as xo import numpy as np - -def _plot_line(line): - tt = line.get_table(attr=True) - xt.twiss.TwissTable.plot(tt, yl='', yr='') - - env = xt.Environment() env.vars({ @@ -21,21 +16,36 @@ def _plot_line(line): # names, tab_sorted = handle_s_places(seq) line = env.new_line(components=[ - env.new_element('b1', xt.Bend, length='l.b1'), - env.new_element('q1', xt.Quadrupole, length='l.q1'), - env.new_element('ip', xt.Marker, at='s.ip'), + 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_element('before_before_right', xt.Marker), - env.new_element('before_right', xt.Sextupole, length=1), - env.new_element('right',xt.Quadrupole, length=1, at='s.right', from_='ip'), - env.new_element('after_right', xt.Marker), - env.new_element('after_right2', xt.Marker), + 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_element('left', xt.Quadrupole, length=1, at='s.left', from_='ip'), - env.new_element('after_left', xt.Marker), - env.new_element('after_left2', xt.Bend, length='l.after_left2'), + 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) + import matplotlib.pyplot as plt plt.close('all') line.survey().plot() From 8bde6254f70a33b1625c35e8e758fe715d91e83c Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 14 Sep 2024 08:32:40 +0200 Subject: [PATCH 153/159] More checks --- examples/lattice_design_shortcuts/001_dev_at_s.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 05bb88b5d..668a4bfe8 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -45,6 +45,12 @@ 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) + import matplotlib.pyplot as plt plt.close('all') From 33e927078e1db76bd87d0d5d243b63cc1e944e2e Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 14 Sep 2024 09:34:29 +0200 Subject: [PATCH 154/159] More checks --- examples/lattice_design_shortcuts/001_dev_at_s.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/examples/lattice_design_shortcuts/001_dev_at_s.py b/examples/lattice_design_shortcuts/001_dev_at_s.py index 668a4bfe8..ea411ecd7 100644 --- a/examples/lattice_design_shortcuts/001_dev_at_s.py +++ b/examples/lattice_design_shortcuts/001_dev_at_s.py @@ -50,6 +50,12 @@ 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 From 372b20de66eaf99d157be363858831beb7f7b472 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 14 Sep 2024 09:58:23 +0200 Subject: [PATCH 155/159] Add test for environment --- tests/test_environment.py | 434 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 434 insertions(+) create mode 100644 tests/test_environment.py diff --git a/tests/test_environment.py b/tests/test_environment.py new file mode 100644 index 000000000..8c3031f6f --- /dev/null +++ b/tests/test_environment.py @@ -0,0 +1,434 @@ +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 + + 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='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() + + From 26b537438a0eb65a1378ccbf51e514009b84f604 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 14 Sep 2024 21:24:46 +0200 Subject: [PATCH 156/159] Fix --- xtrack/line.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 9c312e5a8..a176525a5 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -3423,6 +3423,9 @@ def select(self, start=None, end=None, name=None): 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') @@ -3431,7 +3434,7 @@ def set(self, name, *args, **kwargs): xt.environment._set_kwargs( name=name, ref_kwargs=ref_kwargs, value_kwargs=value_kwargs, element_dict=self.element_dict, element_refs=self.element_refs) - elif name in self.vars: + else: if len(kwargs) > 0: raise ValueError(f'Only a single value is allowed when setting variable') if len(args) != 1: @@ -3441,10 +3444,6 @@ def set(self, name, *args, **kwargs): self.vars[name] = _eval(value) else: self.vars[name] = value - elif hasattr(self, 'lines') and name in self.lines: - raise ValueError('Cannot set a line') - else: - raise KeyError(f'Name {name} not found.') def _env_if_needed(self): if not hasattr(self, 'env') or self.env is None: From 238437ea110c15d6ff28c33d178dace97626bd37 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 14 Sep 2024 21:26:41 +0200 Subject: [PATCH 157/159] Add check --- tests/test_environment.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_environment.py b/tests/test_environment.py index 8c3031f6f..7dd1df0c4 100644 --- a/tests/test_environment.py +++ b/tests/test_environment.py @@ -44,6 +44,11 @@ def test_vars_and_element_access_modes(container_type): 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 + ee.set('a', 0.) assert ee.vv['k.1'] == 3 * 0 + 6 assert ee.vv['b'] == 2 * 0 + 3 * 0 + 6 From 383df523ef146ad3c3bd39c6dba7304e944dd260 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 15 Sep 2024 22:25:12 +0200 Subject: [PATCH 158/159] Corner case: kwarg with numerical value passed as string --- tests/test_environment.py | 4 ++++ xtrack/environment.py | 5 ++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_environment.py b/tests/test_environment.py index 7dd1df0c4..df22fff29 100644 --- a/tests/test_environment.py +++ b/tests/test_environment.py @@ -48,6 +48,8 @@ def test_vars_and_element_access_modes(container_type): 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 @@ -247,6 +249,8 @@ def test_assemble_ring(): 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') diff --git a/xtrack/environment.py b/xtrack/environment.py index 45f91b5b7..17d077912 100644 --- a/xtrack/environment.py +++ b/xtrack/environment.py @@ -312,7 +312,10 @@ def _parse_kwargs(cls, kwargs, _eval): 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]) - value_kwargs[kk] = ref_kwargs[kk]._value + 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] From fbbf97b8ce5f01cc49c692818bb2f3ad7c4bfe83 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 17 Sep 2024 15:43:10 +0200 Subject: [PATCH 159/159] Re-enable asserts in coasting example --- examples/coasting/001_frev_meas.py | 6 +++--- .../radiation/009a_sps_with_vertical_bump.py | 20 +++++++++---------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/examples/coasting/001_frev_meas.py b/examples/coasting/001_frev_meas.py index 0b4df59ab..cdfa3c13b 100644 --- a/examples/coasting/001_frev_meas.py +++ b/examples/coasting/001_frev_meas.py @@ -150,9 +150,9 @@ def y_mean_hist(line, particles): print('f_measured: ', f_measured, ' Hz') print('Error: ', f_measured - f_expected, 'Hz') -# assert np.isclose(f_expected, f_measured, rtol=0, atol=2) # 2 Hz tolerance -# assert np.isclose(np.mean(inten), inten_exp, rtol=1e-2, atol=0) -# assert np.allclose(p.at_turn, num_turns*0.9, rtol=3e-2, atol=0) #beta1 defaults to 0.1 +assert np.isclose(f_expected, f_measured, rtol=0, atol=2) # 2 Hz tolerance +assert np.isclose(np.mean(inten), inten_exp, rtol=1e-2, atol=0) +assert np.allclose(p.at_turn, num_turns*0.9, rtol=3e-2, atol=0) #beta1 defaults to 0.1 tt = line.get_table() tt_synch = tt.rows[tt.element_type=='SyncTime'] 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(