diff --git a/tests/test_local_momentum_acceptance.py b/tests/test_local_momentum_acceptance.py new file mode 100644 index 000000000..b4fd47b4b --- /dev/null +++ b/tests/test_local_momentum_acceptance.py @@ -0,0 +1,397 @@ +# copyright ############################### # +# This file is part of the Xtrack Package. # +# Copyright (c) CERN, 2025. # +# ######################################### # +import pytest +import numpy as np +import xtrack as xt + + +@pytest.fixture(scope="module") +def toy_ring(temp_context_default_mod): + """Build a toy ring and Twiss it once for the entire test session.""" + lbend = 3 + angle = np.pi / 2 + + env = xt.Environment() + + line = env.new_line(components=[ + env.new('mqf.1', xt.Quadrupole, length=0.3, k1=0.1), + env.new('d1.1', xt.Drift, length=1), + env.new('mb1.1', xt.Bend, length=lbend, angle=angle), + env.new('d2.1', xt.Drift, length=1), + + env.new('mqd.1', xt.Quadrupole, length=0.3, k1=-0.7), + env.new('d3.1', xt.Drift, length=1), + env.new('mb2.1', xt.Bend, length=lbend, angle=angle), + env.new('d4.1', xt.Drift, length=1), + + env.new('mqf.2', xt.Quadrupole, length=0.3, k1=0.1), + env.new('d1.2', xt.Drift, length=1), + env.new('mb1.2', xt.Bend, length=lbend, angle=angle), + env.new('d2.2', xt.Drift, length=1), + + env.new('mqd.2', xt.Quadrupole, length=0.3, k1=-0.7), + env.new('d3.2', xt.Drift, length=1), + env.new('mb2.2', xt.Bend, length=lbend, angle=angle), + env.new('d4.2', xt.Drift, length=1), + ]) + + line.set_particle_ref('electron', p0c=1e9) + line.configure_bend_model(core='full', edge=None) + + tt = line.get_table() + needs_aperture = tt.rows.match(element_type='Bend|Quadrupole| ').name + aper_size = 0.040 # m + + placements = [] + for nn in needs_aperture: + env.new( + f'{nn}_aper_entry', xt.LimitRect, + min_x=-aper_size, max_x=aper_size, + min_y=-aper_size, max_y=aper_size + ) + placements.append(env.place(f'{nn}_aper_entry', at=f'{nn}@start')) + + env.new( + f'{nn}_aper_exit', xt.LimitRect, + min_x=-aper_size, max_x=aper_size, + min_y=-aper_size, max_y=aper_size + ) + placements.append(env.place(f'{nn}_aper_exit', at=f'{nn}@end')) + + line.insert(placements) + + tw = line.twiss(method="4d") + tw.particle_on_co.move(_context=line._context) + + return {'line': line, 'twiss': tw} + + +def test_parameter_validation(): + """ + Validate that invalid input parameters are rejected with clear error messages. + """ + env = xt.Environment() + line = env.new_line(components=[env.new('d', xt.Drift, length=1.0)]) + line.set_particle_ref('electron', p0c=1e9) + + cases = [ + (dict(delta_negative_limit=0.0), ValueError, r"delta_negative_limit must be < 0"), + (dict(delta_positive_limit=0.0), ValueError, r"delta_positive_limit must be > 0"), + (dict(delta_step_size=0.0), ValueError, r"delta_step_size must be > 0"), + (dict(n_turns=0), ValueError, r"n_turns must be > 0"), + ] + + for kwargs, exc, pattern in cases: + with pytest.raises(exc, match=pattern): + line.get_local_momentum_acceptance( + nemitt_x=1e-6, + nemitt_y=1e-6, + method="4d", + **kwargs, + ) + + +def test_elements_validation(): + """ + Validate that invalid `elements` arguments are rejected with clear error messages. + """ + env = xt.Environment() + line = env.new_line(components=[env.new('d', xt.Drift, length=1.0)]) + line.set_particle_ref('electron', p0c=1e9) + + # Not iterable / scalar + with pytest.raises(ValueError, match=r"`elements` must be an iterable of strings"): + line.get_local_momentum_acceptance( + elements=42, + nemitt_x=1e-6, + nemitt_y=1e-6, + method="4d", + ) + + # Bare string (not a list of strings) + with pytest.raises(ValueError, match=r"`elements` must be an iterable of strings"): + line.get_local_momentum_acceptance( + elements='d', + nemitt_x=1e-6, + nemitt_y=1e-6, + method="4d", + ) + + # Contains non-strings + with pytest.raises(ValueError, match=r"All entries in `elements` must be strings"): + line.get_local_momentum_acceptance( + elements=['d', 123], + nemitt_x=1e-6, + nemitt_y=1e-6, + method="4d", + ) + + # Element not in the line + with pytest.raises(ValueError, match=r"The following elements were not found in the line"): + line.get_local_momentum_acceptance( + elements=['does_not_exist'], + nemitt_x=1e-6, + nemitt_y=1e-6, + method="4d", + ) + + +def test_nemitt_not_provided(): + """ + nemitt_x and nemitt_y are required; omitting either must raise. + """ + env = xt.Environment() + line = env.new_line(components=[env.new('d', xt.Drift, length=1.0)]) + line.set_particle_ref('electron', p0c=1e9) + + with pytest.raises(ValueError, match=r"nemitt_x and nemitt_y must be provided"): + line.get_local_momentum_acceptance(method="4d") + + +def test_no_particle_ref_raises(): + """ + If the line has no particle_ref set, the method must raise immediately. + """ + env = xt.Environment() + line = env.new_line(components=[env.new('d', xt.Drift, length=1.0)]) + # Deliberately no set_particle_ref + + with pytest.raises(ValueError, match=r"Line.particle_ref must be set"): + line.get_local_momentum_acceptance( + nemitt_x=1e-6, + nemitt_y=1e-6, + ) + + +def test_mutual_exclusivity_errors(toy_ring): + """ + x/y physical offsets must be mutually exclusive with normalized offsets. + """ + line = toy_ring['line'] + + with pytest.raises(ValueError, match=r"Provide either x_offset or x_norm_offset"): + line.get_local_momentum_acceptance( + nemitt_x=1e-6, + nemitt_y=1e-6, + method="4d", + x_offset=1e-3, + x_norm_offset=1.0, + ) + + with pytest.raises(ValueError, match=r"Provide either y_offset or y_norm_offset"): + line.get_local_momentum_acceptance( + nemitt_x=1e-6, + nemitt_y=1e-6, + method="4d", + y_offset=1e-3, + y_norm_offset=1.0, + ) + + +def test_elements_selects_subset(toy_ring): + """ + Passing a filtered list of element names should restrict the output to + only those elements. + """ + line = toy_ring['line'] + tw = toy_ring['twiss'] + + tt = line.get_table() + expected_names = tt.rows['^mqf.*_aper_entry$'].name + + out = line.get_local_momentum_acceptance( + elements=list(expected_names), + twiss=tw, + nemitt_x=1e-5, + nemitt_y=1e-7, + delta_negative_limit=-0.005, + delta_positive_limit=+0.005, + delta_step_size=0.001, + n_turns=512, + ) + + assert set(out.name) == set(expected_names) + assert len(out.name) == len(expected_names) + assert "s" in out.cols + assert "deltan" in out.cols + assert "deltap" in out.cols + + +def test_s_window_selects_subset(toy_ring): + """ + Passing elements pre-filtered by s window should restrict the output + to elements within the window. + """ + line = toy_ring['line'] + tw = toy_ring['twiss'] + + tt = line.get_table() + s_end = tt.s[-1] / 2.0 + tab_in_window = tt.rows[0.0:s_end:'s'] + tab_aper_in_window = tab_in_window.rows[tab_in_window.element_type == 'LimitRect'] + + out = line.get_local_momentum_acceptance( + elements=list(tab_aper_in_window.name), + twiss=tw, + nemitt_x=1e-5, + nemitt_y=1e-7, + delta_negative_limit=-0.005, + delta_positive_limit=+0.005, + delta_step_size=0.001, + n_turns=512, + ) + + assert len(out.s) == len(tab_aper_in_window.name) + assert set(out.name) == set(tab_aper_in_window.name) + + +def test_norm_offset_all_survive(toy_ring): + """ + A small normalized transverse offset should not cause additional losses. + """ + line = toy_ring['line'] + tw = toy_ring['twiss'] + + tt = line.get_table() + tt_aper = tt.rows[tt.element_type == 'LimitRect'] + + delta_neg = -0.005 + delta_pos = +0.005 + delta_step = 0.001 + + out = line.get_local_momentum_acceptance( + elements=list(tt_aper.name), + twiss=tw, + nemitt_x=1e-5, + nemitt_y=1e-7, + x_norm_offset=0.1, + y_norm_offset=0.1, + delta_negative_limit=delta_neg, + delta_positive_limit=delta_pos, + delta_step_size=delta_step, + n_turns=512, + ) + + delta_co = np.array([tw["delta", nn] for nn in tt_aper.name], dtype=float) + assert np.allclose(out.deltan, delta_co + delta_neg, atol=1e-12) + assert np.allclose(out.deltap, delta_co + delta_pos, atol=1e-12) + + +def test_all_survive(toy_ring): + """ + For a small delta scan all particles survive; deltan/deltap must match + the scan bounds around the local closed-orbit delta. + """ + line = toy_ring['line'] + tw = toy_ring['twiss'] + + tt = line.get_table() + tt_aper = tt.rows[tt.element_type == 'LimitRect'] + + delta_neg = -0.005 + delta_pos = +0.005 + delta_step = 0.001 + + out = line.get_local_momentum_acceptance( + elements=list(tt_aper.name), + twiss=tw, + nemitt_x=1e-5, + nemitt_y=1e-7, + delta_negative_limit=delta_neg, + delta_positive_limit=delta_pos, + delta_step_size=delta_step, + n_turns=512, + with_progress=False, + verbose=False, + ) + + assert "s" in out.cols + assert "deltan" in out.cols + assert "deltap" in out.cols + + assert len(out.s) == len(tt_aper.name) + assert np.allclose(out.s, tt_aper.s, atol=1e-12) + + delta_co = np.array([tw["delta", nn] for nn in tt_aper.name], dtype=float) + assert np.all(out.deltan <= delta_co + 1e-12) + assert np.all(out.deltap >= delta_co - 1e-12) + assert np.allclose(out.deltan, delta_co + delta_neg, atol=1e-12) + assert np.allclose(out.deltap, delta_co + delta_pos, atol=1e-12) + + +def test_all_lost(toy_ring): + """ + For a large delta scan all particles are lost; deltan=deltap=0 everywhere. + """ + line = toy_ring['line'] + tw = toy_ring['twiss'] + + tt = line.get_table() + tt_aper = tt.rows[tt.element_type == 'LimitRect'] + + out = line.get_local_momentum_acceptance( + elements=list(tt_aper.name), + twiss=tw, + nemitt_x=1e-5, + nemitt_y=1e-7, + delta_negative_limit=-0.2, + delta_positive_limit=+0.2, + delta_step_size=0.2, + with_progress=False, + verbose=False, + ) + + assert "s" in out.cols + assert "deltan" in out.cols + assert "deltap" in out.cols + assert len(out.s) == len(tt_aper.name) + assert np.allclose(out.s, tt_aper.s, atol=1e-12) + assert np.all(out.deltan <= 0) + assert np.all(out.deltap >= 0) + assert np.allclose(out.deltan, 0, atol=1e-12) + assert np.allclose(out.deltap, 0, atol=1e-12) + + +@pytest.mark.parametrize( + "x_offset, y_offset", + [ + pytest.param(0.050, 0.0, id="lost_by_x_offset"), + pytest.param(0.0, 0.050, id="lost_by_y_offset"), + ], +) +def test_all_lost_offset(toy_ring, x_offset, y_offset): + """ + A transverse offset larger than the aperture causes all particles to be lost. + """ + line = toy_ring['line'] + tw = toy_ring['twiss'] + + tt = line.get_table() + tt_aper = tt.rows[tt.element_type == 'LimitRect'] + + out = line.get_local_momentum_acceptance( + elements=list(tt_aper.name), + twiss=tw, + nemitt_x=1e-5, + nemitt_y=1e-7, + x_offset=x_offset, + y_offset=y_offset, + delta_negative_limit=-0.001, + delta_positive_limit=+0.001, + delta_step_size=0.0001, + n_turns=512, + with_progress=False, + verbose=False, + ) + + assert "s" in out.cols + assert "deltan" in out.cols + assert "deltap" in out.cols + assert len(out.s) == len(tt_aper.name) + assert np.allclose(out.s, tt_aper.s, atol=1e-12) + assert np.all(out.deltan <= 0) + assert np.all(out.deltap >= 0) + assert np.allclose(out.deltan, 0.0, atol=1e-12) + assert np.allclose(out.deltap, 0.0, atol=1e-12) \ No newline at end of file diff --git a/xtrack/line.py b/xtrack/line.py index de6e14e4c..fdfc6a30b 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -1937,6 +1937,242 @@ def track( multi_element_monitor_at=multi_element_monitor_at, **kwargs) + @doc_group("Tracking and Analysis") + def get_local_momentum_acceptance( + self, + *, + elements=None, + twiss=None, + scattering='off', + x_offset: float = 0.0, + y_offset: float = 0.0, + x_norm_offset: float = 0.0, + y_norm_offset: float = 0.0, + nemitt_x=None, + nemitt_y=None, + delta_negative_limit: float = -0.10, + delta_positive_limit: float = +0.10, + delta_step_size: float = 0.01, + n_turns: int = 512, + with_progress: bool | int = False, + verbose: bool = False, + **kwargs): + """ + Compute the local momentum acceptance (LMA) along the line by tracking a + grid of momentum offsets (δ) from the **entrance** of selected + elements and reporting the largest surviving negative and positive δ. + + The δ grid is centered on the local closed orbit at each element, and offsets + can be applied (either physical x/y or normalized x/y in σ units). + + Parameters + ---------- + elements : list of str or array-like of str, optional + Names of the elements at whose entrance the LMA is evaluated. + If ``None`` (default), all elements in the line are used. + If multiple elements share the same ``s``, only the first encountered + is used. + twiss : xt.TwissTable, optional + Twiss table to define the closed orbit and optics. By default, + a 6D solution is computed with `self.twiss(method='6d')`. You can + override the method with `method=...` in `**kwargs`. + scattering : str, optional + Wheter scattering has been enabled or not (`'on'` or `'off'`). + x_offset : float, default 0.0 + Horizontal physical offset in meters. Mutually exclusive with + `x_norm_offset`. + y_offset : float, default 0.0 + Vertical physical offset in meters. Mutually exclusive with + `y_norm_offset`. + x_norm_offset : float, default 0.0 + Horizontal normalized offset in units of σx (rms). Mutually exclusive + with `x_offset`. + y_norm_offset : float, default 0.0 + Vertical normalized offset in units of σy (rms). Mutually exclusive + with `y_offset`. + nemitt_x : float + Horizontal normalized emittance (m·rad, rms). + nemitt_y : float + Vertical normalized emittance (m·rad, rms). + delta_negative_limit : float, default -0.10 + Lower bound of the δ scan (inclusive). Must be < 0. + delta_positive_limit : float, default +0.10 + Upper bound of the δ scan (inclusive). Must be > 0. + delta_step_size : float, default 0.01 + Step for the δ grid. Must be > 0. The positive end is included + with a half-step guard to reduce floating-point exclusion. + n_turns : int, default 512 + Number of turns to track. + with_progress : bool | int, default False + If truthy, shows a per-element progress bar. + verbose : bool, default False + If True, enables tracker progress for each element scan. + **kwargs + Passed through to `self.twiss` and `build_particles`. + + Selection semantics + ------------------- + - LMA is evaluated at the **entrance** of each element. + - If multiple elements share the same `s`, only the first encountered is used. + + Algorithm (per selected element) + -------------------------------- + 1. Build particles on closed orbit with the requested (normalized or physical) offsets. + 2. Apply the δ grid by shifting the initial δ around `delta_co`. + 3. Track for `n_turns` turns from the element to itself. + 4. Among surviving particles, report: + - `deltan` = min of the *initial* δ of survivors, + - `deltap` = max of the *initial* δ of survivors. + If none survive, `deltan` = `deltap` = 0.0 + + Returns + ------- + xt.Table + Table indexed by `'name'` with columns: + - `name` (str): Element name. + - `s` (float): Element entrance position (m). + - `deltan` (float): Largest surviving negative δ (may be 0). + - `deltap` (float): Largest surviving positive δ (may be 0). + """ + if self.particle_ref is None: + raise ValueError("Line.particle_ref must be set to build probe particles.") + + # Mutual exclusivity: physical vs normalized offsets + if x_offset != 0.0 and x_norm_offset != 0.0: + raise ValueError("Provide either x_offset or x_norm_offset, not both.") + if y_offset != 0.0 and y_norm_offset != 0.0: + raise ValueError("Provide either y_offset or y_norm_offset, not both.") + + if nemitt_x is None or nemitt_y is None: + raise ValueError("nemitt_x and nemitt_y must be provided.") + + if delta_negative_limit >= 0: + raise ValueError("delta_negative_limit must be < 0") + if delta_positive_limit <= 0: + raise ValueError("delta_positive_limit must be > 0") + if delta_step_size <= 0: + raise ValueError("delta_step_size must be > 0") + if n_turns <= 0: + raise ValueError("n_turns must be > 0") + + if elements is not None: + if not hasattr(elements, '__iter__') or isinstance(elements, str): + raise ValueError("`elements` must be an iterable of strings, not a scalar.") + elements = list(elements) + if not all(isinstance(e, str) for e in elements): + raise ValueError("All entries in `elements` must be strings.") + invalid = [e for e in elements if e not in self.element_names] + if invalid: + raise ValueError( + f"The following elements were not found in the line: {invalid}") + + if not self._has_valid_tracker(): + self.build_tracker() + + # Compute twiss (use 6D by default, overridable via kwargs['method']) + twiss_method = kwargs.pop('method', '6d') + if twiss is None: + if scattering == 'on': + self.scattering.disable() + twiss = self.twiss(method=twiss_method, reverse=False) + if scattering == 'on': + self.scattering.enable() + + if elements is None: + tt = self.get_table() + elements = tt.name + + # Delta grid + deltas = np.arange(delta_negative_limit, delta_positive_limit + 0.5 * delta_step_size, + delta_step_size) + n_part = len(deltas) + + rows = [] + seen_s = set() + + iterable = progress(elements, desc="Local Momentum Acceptance") if with_progress else elements + + for ii, ee in enumerate(iterable): + s_here = float(twiss['s', ee]) + + # Some elements may share the same s + if s_here in seen_s: + continue + seen_s.add(s_here) + + ## Prepare test particles + # The longitudinal closed orbit need to be manually supplied + zeta_co = twiss['zeta', ee] + delta_co = twiss['delta', ee] + + if scattering == 'on': + self.scattering.disable() + + # Extract W_matrix and particle_on_co from the already-computed twiss + tw_init = twiss.get_twiss_init(at_element=ee) + + idx_at_element = self.element_names.index(ee) + + # On-momentum, matched, test particles + particles = self.build_particles( + _context=self._context, + num_particles=n_part, + x_norm=x_norm_offset, + y_norm=y_norm_offset, + zeta=zeta_co, + delta=delta_co, + nemitt_x=nemitt_x, + nemitt_y=nemitt_y, + W_matrix=tw_init.W_matrix, + particle_on_co=tw_init.particle_on_co, + ) + particles.at_element[:] = idx_at_element + particles.s[:] = s_here + particles.start_tracking_at_element = -1 + + if scattering == 'on': + self.scattering.enable() + + # Add the delta grid + delta_temp = particles.delta.copy() + delta_temp += deltas + particles.update_delta(delta_temp) + + initial_deltas = particles.delta.copy() + + # Apply absolute offsets, if any + particles.x += x_offset + particles.y += y_offset + + print(f"\nTrack test particles from reference point #{ii}") + self.track( + particles, + ele_start=ee, + ele_stop=ee, + num_turns=n_turns, + with_progress=1 if verbose else 0 + ) + + mask_alive = (particles.state == 1) + if np.any(mask_alive): + surviving_pids = particles.filter(mask_alive).particle_id + deltan = float(np.min(initial_deltas[surviving_pids])) + deltap = float(np.max(initial_deltas[surviving_pids])) + else: + deltan = float(0.0) + deltap = float(0.0) + + rows.append({ + 'name': ee, + 's': s_here, + 'deltan': deltan, + 'deltap': deltap + }) + + cols = {k: np.array([r[k] for r in rows]) for k in rows[0].keys()} + + return xt.Table(cols, index='name') + @doc_group("Line Editing") def slice_thick_elements(self, slicing_strategies): """