diff --git a/MANIFEST.in b/MANIFEST.in index 829d9280..f989391d 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -6,5 +6,8 @@ include *.md # Include the license file include LICENSE.txt +# Include the notice +include NOTICE + recursive-include xfields *.h -recursive-include xfields *.clh +recursive-include xfields *.clh \ No newline at end of file diff --git a/NOTICE b/NOTICE new file mode 100644 index 00000000..8c0c5460 --- /dev/null +++ b/NOTICE @@ -0,0 +1,7 @@ +This package incorporates portions adapted from Argonne National Laboratory’s “Elegant” +and from the SDDS Toolkit. © 2002 The University of Chicago; © 2002 The Regents of the +University of California. Distributed under their respective Software License Agreements. +See: xfields/third_party/elegant/LICENSE and xfields/third_party/SDDS/LICENSE. + +This package also incorporates LAPACK’s DLARAN algorithm (modified BSD license). +See: xfields/third_party/lapack/LICENSE. \ No newline at end of file diff --git a/contributors_and_credits.txt b/contributors_and_credits.txt index bf72b891..3af431d3 100644 --- a/contributors_and_credits.txt +++ b/contributors_and_credits.txt @@ -2,6 +2,7 @@ The following people contributed to the development of this package: CERN contributors: - H. Bartosik + - G. Broggi - X. Buffat - R. De Maria - L. Giacomel @@ -18,3 +19,8 @@ External contributors: The Particle In Cell implementation is largely based on the PyPIC package (https://github.com/pycomplete/pypic) + +The Touschek scattering routine contains portions adapted from Elegant and the SDDS Toolkit. +© 2002 The University of Chicago; © 2002 The Regents of the University of California. +Distributed subject to their Software License Agreements. See third_party/elegant/LICENSE and third_party/SDDS/LICENSE. +The RNG core used in the Touschek scattering routine incorporates LAPACK’s DLARAN (modified BSD license). See third_party/lapack/LICENSE. \ No newline at end of file diff --git a/examples/006_touschek/000_touschek_toy_ring.py b/examples/006_touschek/000_touschek_toy_ring.py new file mode 100644 index 00000000..190c19ed --- /dev/null +++ b/examples/006_touschek/000_touschek_toy_ring.py @@ -0,0 +1,231 @@ +# copyright ############################### # +# This file is part of the Xtrack Package. # +# Copyright (c) CERN, 2025. # +# ######################################### # +import numpy as np +import matplotlib.pyplot as plt + +import xobjects as xo +import xtrack as xt +import xfields as xf + +###################################################### +# Beam parameters +###################################################### +nemitt_x = 1e-5 +nemitt_y = 1e-7 + +sigma_z = 4e-3 +sigma_delta = 1e-3 + +bunch_population = 4e9 + +###################################################### +# Build a toy ring +###################################################### +lbend = 3 +angle = np.pi / 2 + +lquad = 0.3 +k1qf = 0.1 +k1qd = 0.7 + +# Create environment +env = xt.Environment() + +# Define the line (toy ring) +line = env.new_line(components=[ + env.new('mqf.1', xt.Quadrupole, length=lquad, k1=k1qf), + 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=lquad, k1=-k1qd), + 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=lquad, k1=k1qf), + 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=lquad, k1=-k1qd), + 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), +]) + +# Set the reference particle +line.set_particle_ref('electron', p0c=1e9) + +# Configure the bend model +line.configure_bend_model(core='full', edge=None) + +###################################################### +# Insert Touschek scattering centers +###################################################### +# We insert Touschek scattering centers in the middle of each magnet +# to have good coverage of variations of the optical functions +tab = line.get_table() +tab_bends_quads = tab.rows[(tab.element_type == 'Bend') | (tab.element_type == 'Quadrupole')] + +placements = [] +for ii, nn in enumerate(tab_bends_quads.name): + tscatter_name = f'TScatter.{ii}' + env.elements[tscatter_name] = xf.TouschekScattering() + placements.append(env.place(tscatter_name, at=0.0, from_=nn)) + +# The last TouschekScattering element has to be placed at the end of the line +tscatter_name = f'TScatter.{ii+1}' +env.elements[tscatter_name] = xf.TouschekScattering() +placements.append(env.place(tscatter_name, at=tab.s[-1])) + +line.insert(placements) + +###################################################### +# Install apertures +###################################################### +tab = line.get_table() +needs_aperture = tab.rows.match_not(element_type='Drift.*|Marker|').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) + +###################################################### +# Evaluate local momentum acceptance profile +###################################################### +# Evaluate local momentum aperture at the touschek scattering centers +tab = line.get_table() +elements = tab.rows.match(element_type="TouschekScattering").name +lma = line.get_local_momentum_acceptance( + # twiss=tw, + elements=elements, + nemitt_x=nemitt_x, + nemitt_y=nemitt_y, + y_offset=1e-9, + delta_negative_limit=-0.012, + delta_positive_limit=0.012, + delta_step_size=1e-4, + n_turns=1000, + method="4d" +) + +###################################################### +# Plot +###################################################### +plt.plot(lma.s, lma.deltan*100, c='r') +plt.plot(lma.s, lma.deltap*100, c='r') +plt.title('Toy ring: local momentum acceptance profile') +plt.xlabel('s [m]') +plt.ylabel(r'$\delta$ [%]') +plt.grid() +plt.show() + +###################################################### +# Touschek simulation +###################################################### +# Parameters +local_momentum_acceptance_scale = 0.85 # scaling factor for local momentum acceptance +n_simulated = 5e6 # number of simulated scattering events with delta > delta_min +nturns = 1000 # number of turns to simulate + +touschek_manager = xf.TouschekManager( + line, + local_momentum_acceptance=lma, + local_momentum_acceptance_scale=local_momentum_acceptance_scale, + nemitt_x=nemitt_x, + nemitt_y=nemitt_y, + sigma_z=sigma_z, + sigma_delta=sigma_delta, + bunch_population=bunch_population, + n_simulated=n_simulated, + nx=3, ny=3, nz=3, + ignored_portion=0.01, + seed=1997, + method='4d' +) + +# Initialise the Touschek simulation. +# Computes the integrated Piwinski scattering rate over each lattice section +# between consecutive TouschekScattering elements (using the trapezoidal rule), +# and configures each element with its local optics, beam parameters, and +# momentum acceptance. The integrated rate is later used to assign the +# correct weights to the Touschek-scattered macro-particles. +touschek_manager.initialise_touschek() + +# Get the list of TouschekScattering elements in the line +touschek_elements = tab.rows[tab.element_type == 'TouschekScattering'].name + +# Build a CPU tracker with OpenMP multithreading to speed up tracking +line.discard_tracker() +line.build_tracker(_context=xo.ContextCpu(omp_num_threads='auto')) + +# For each TouschekScattering element: +# 1. Generate Touschek-scattered macro-particles at that element +# 2. Track them around the ring for `nturns` turns, starting and ending at that element +particles_list = [] +for ii, element in enumerate(touschek_elements): + s_start_elem = tab.rows[tab.name == element].s[0] + + # Generate Touschek-scattered macro-particles + particles = line[element].scatter() + + # Track + print(f"\nTracking particles scattered at {element} (s = {s_start_elem:.2f} m)") + line.track(particles, ele_start=element, ele_stop=element, num_turns=nturns, with_progress=1) + + particles_list.append(particles) + +# Merge the macro-particle sets from all scattering elements into a single collection +particles = xt.Particles.merge(particles_list) + +# Optional: Refine loss location to improve loss map accuracy +loss_loc_refinement = xt.LossLocationRefinement(line, + n_theta = 360, # Angular resolution in the polygonal approximation of the aperture + r_max = 0.5, # Maximum transverse aperture in m + dr = 50e-6, # Transverse loss refinement accuracy [m] + ds = 0.1, # Longitudinal loss refinement accuracy [m] + ) + +loss_loc_refinement.refine_loss_location(particles) + +###################################################### +# Compute Touschek lifetime +###################################################### +# Keep lost particles only +particles = particles.filter(particles.state == 0) +# Compute total Touschek loss rate +loss_rate = sum(particles.weight) +# Compute Touschek lifetime +touschek_lifetime = bunch_population / loss_rate + +###################################################### +# Plot: Toy ring Touschek loss map +###################################################### +circumference = line.get_length() +binwidth = 0.1 # m + +plt.title(f'Toy ring Touschek loss map (Touschek lifetime: {touschek_lifetime/60:.2f} min)') +plt.hist(particles.s, bins=np.arange(0, circumference + binwidth, binwidth), weights=particles.weight*1e-3) +plt.xlabel('s [m]') +plt.ylabel('Loss rate [kHz]') +plt.grid() +plt.show() \ No newline at end of file diff --git a/tests/test_touschek.py b/tests/test_touschek.py new file mode 100644 index 00000000..05826a2e --- /dev/null +++ b/tests/test_touschek.py @@ -0,0 +1,525 @@ +# copyright ################################# # +# This file is part of the Xfields Package. # +# Copyright (c) CERN, 2025. # +# ########################################### # +import numpy as np +import pytest + +import xobjects as xo +import xtrack as xt +import xfields as xf + + +############################################################# +# Shared beam parameters +############################################################# +NEMITT_X = 1e-5 +NEMITT_Y = 1e-7 +SIGMA_Z = 4e-3 +SIGMA_DELTA = 1e-3 +BUNCH_POPULATION = 4e9 + + +############################################################# +# Module-level fixture: toy ring, twiss, and LMA +############################################################# +@pytest.fixture(scope='module') +def toy_ring(): + """ + Build a FODO-like toy ring with TouschekScattering elements and apertures, + run Twiss and the local momentum acceptance exactly once for the whole + test session. + + Yields a dict with keys: + line - the xtrack.Line (with tracker already built) + twiss - the 4d Twiss table + lma - the xt.Table of local momentum acceptance + """ + lbend = 3.0 + angle = np.pi / 2 + lquad = 0.3 + + env = xt.Environment() + line = env.new_line(components=[ + env.new('mqf.1', xt.Quadrupole, length=lquad, k1= 0.1), + env.new('d1.1', xt.Drift, length=1.0), + env.new('mb1.1', xt.Bend, length=lbend, angle=angle), + env.new('d2.1', xt.Drift, length=1.0), + + env.new('mqd.1', xt.Quadrupole, length=lquad, k1=-0.7), + env.new('d3.1', xt.Drift, length=1.0), + env.new('mb2.1', xt.Bend, length=lbend, angle=angle), + env.new('d4.1', xt.Drift, length=1.0), + + env.new('mqf.2', xt.Quadrupole, length=lquad, k1= 0.1), + env.new('d1.2', xt.Drift, length=1.0), + env.new('mb1.2', xt.Bend, length=lbend, angle=angle), + env.new('d2.2', xt.Drift, length=1.0), + + env.new('mqd.2', xt.Quadrupole, length=lquad, k1=-0.7), + env.new('d3.2', xt.Drift, length=1.0), + env.new('mb2.2', xt.Bend, length=lbend, angle=angle), + env.new('d4.2', xt.Drift, length=1.0), + ]) + + line.set_particle_ref('electron', p0c=1e9) + line.configure_bend_model(core='full', edge=None) + + # Insert one TouschekScattering at the entrance of every magnet, + # plus one at the very end of the line — all in a single batch insert. + tab = line.get_table() + tab_bends_quads = tab.rows[ + (tab.element_type == 'Bend') | (tab.element_type == 'Quadrupole') + ] + placements = [] + for ii, nn in enumerate(tab_bends_quads.name): + tscatter_name = f'TScatter.{ii}' + env.elements[tscatter_name] = xf.TouschekScattering() + placements.append(env.place(tscatter_name, at=0.0, from_=nn)) + + # Last TouschekScattering at the end of the line + tscatter_name = f'TScatter.{ii+1}' + env.elements[tscatter_name] = xf.TouschekScattering() + placements.append(env.place(tscatter_name, at=tab.s[-1])) + + line.insert(placements) + + # Rectangular apertures around every non-drift/non-marker element + aper_size = 0.040 # m + tab = line.get_table() + needs_aperture = tab.rows.match_not(element_type='Drift.*|Marker|').name + 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) + + # Twiss + tw = line.twiss(method='4d') + tw.particle_on_co.move(_context=line._context) + + # LMA + tab = line.get_table() + elements = tab.rows.match(element_type='TouschekScattering').name + lma = line.get_local_momentum_acceptance( + elements=elements, + twiss=tw, + nemitt_x=NEMITT_X, + nemitt_y=NEMITT_Y, + y_offset=1e-12, + delta_negative_limit=-0.012, + delta_positive_limit=+0.012, + delta_step_size=0.001, + n_turns=64, + method='4d', + with_progress=False, + verbose=False, + ) + + # Build a tracker once so tracking tests can reuse it + line.build_tracker(_context=xo.ContextCpu(omp_num_threads=1)) + + yield dict(line=line, twiss=tw, lma=lma) + + +############################################################# +# Helper: fresh (unscaled) copy of LMA and new TouschekManager +############################################################# +def _fresh_lma(toy_ring_data): + """Return an independent copy of the module-level LMA table.""" + lma = toy_ring_data['lma'] + return xt.Table({ + 'name': lma.name.copy(), + 's': lma.s.copy(), + 'deltan': lma.deltan.copy(), + 'deltap': lma.deltap.copy(), + }) + +def _build_manager(line, lma, twiss=None, *, n_simulated=int(1e6), **kwargs): + """ + Convenience factory for TouschekManager with sensible test defaults. + """ + defaults = dict( + local_momentum_acceptance=lma, + local_momentum_acceptance_scale=0.85, + nemitt_x=NEMITT_X, + nemitt_y=NEMITT_Y, + sigma_z=SIGMA_Z, + sigma_delta=SIGMA_DELTA, + bunch_population=BUNCH_POPULATION, + n_simulated=n_simulated, + nx=3, ny=3, nz=3, + ignored_portion=0.01, + seed=1997, + method='4d', + ) + defaults.update(kwargs) + return xf.TouschekManager(line, twiss=twiss, **defaults) + + +############################################################# +# Tests +############################################################# +class TestTouschekManagerInit: + """Unit-tests for TouschekManager constructor validation.""" + + def test_construction_with_normalised_emittances(self, toy_ring): + """Manager constructed with nemitt_x/y converts to geometric emittances.""" + line = toy_ring['line'] + lma = _fresh_lma(toy_ring) + tw = toy_ring['twiss'] + tm = _build_manager(line, lma, tw) + beta0 = line.particle_ref.beta0[0] + gamma0 = line.particle_ref.gamma0[0] + assert tm.gemitt_x == pytest.approx(NEMITT_X / (beta0 * gamma0)) + assert tm.gemitt_y == pytest.approx(NEMITT_Y / (beta0 * gamma0)) + + def test_construction_with_geometric_emittances(self, toy_ring): + """Manager constructed with gemitt_x/y stores them directly.""" + line = toy_ring['line'] + lma = _fresh_lma(toy_ring) + tw = toy_ring['twiss'] + beta0 = line.particle_ref.beta0[0] + gamma0 = line.particle_ref.gamma0[0] + gx = NEMITT_X / (beta0 * gamma0) + gy = NEMITT_Y / (beta0 * gamma0) + tm = _build_manager(line, lma, tw, gemitt_x=gx, gemitt_y=gy, + nemitt_x=None, nemitt_y=None) + assert tm.gemitt_x == pytest.approx(gx) + assert tm.gemitt_y == pytest.approx(gy) + + def test_lma_is_scaled_in_place(self, toy_ring): + """LMA columns must be multiplied by local_momentum_acceptance_scale.""" + line = toy_ring['line'] + tw = toy_ring['twiss'] + lma_fresh = _fresh_lma(toy_ring) + deltan_before = lma_fresh.deltan.copy() + deltap_before = lma_fresh.deltap.copy() + scale = 0.85 + _build_manager(line, lma_fresh, tw, + local_momentum_acceptance_scale=scale) + np.testing.assert_allclose(lma_fresh.deltan, deltan_before * scale) + np.testing.assert_allclose(lma_fresh.deltap, deltap_before * scale) + + def test_raises_on_missing_line(self, toy_ring): + with pytest.raises(ValueError, match=r'`line` is required'): + xf.TouschekManager( + line=None, + local_momentum_acceptance=_fresh_lma(toy_ring), + sigma_z=SIGMA_Z, sigma_delta=SIGMA_DELTA, + bunch_population=BUNCH_POPULATION, n_simulated=10, + nemitt_x=NEMITT_X, nemitt_y=NEMITT_Y, + ) + + @pytest.mark.parametrize('missing_key', [ + 'sigma_z', 'sigma_delta', 'bunch_population', 'n_simulated', + ]) + def test_raises_on_missing_required_kwarg(self, toy_ring, missing_key): + """Each required kwarg should raise ValueError when absent.""" + line = toy_ring['line'] + required = dict( + local_momentum_acceptance=_fresh_lma(toy_ring), + sigma_z=SIGMA_Z, sigma_delta=SIGMA_DELTA, + bunch_population=BUNCH_POPULATION, n_simulated=10, + nemitt_x=NEMITT_X, nemitt_y=NEMITT_Y, + ) + kw = {k: v for k, v in required.items() if k != missing_key} + with pytest.raises(ValueError): + xf.TouschekManager(line, **kw) + + def test_raises_on_both_nemitt_and_gemitt(self, toy_ring): + line = toy_ring['line'] + lma = _fresh_lma(toy_ring) + tw = toy_ring['twiss'] + with pytest.raises(ValueError, match=r'not both'): + _build_manager(line, lma, tw, gemitt_x=1e-9, gemitt_y=1e-11) + + def test_raises_on_neither_nemitt_nor_gemitt(self, toy_ring): + line = toy_ring['line'] + with pytest.raises(ValueError, match=r'must provide'): + xf.TouschekManager( + line, + local_momentum_acceptance=_fresh_lma(toy_ring), + sigma_z=SIGMA_Z, sigma_delta=SIGMA_DELTA, + bunch_population=BUNCH_POPULATION, n_simulated=10, + ) + + def test_raises_on_wrong_lma_type(self, toy_ring): + """Non-xt.Table local_momentum_acceptance must raise TypeError.""" + import pandas as pd + line = toy_ring['line'] + bad = pd.DataFrame({'name': ['a'], 's': [0.0], + 'deltan': [-0.01], 'deltap': [0.01]}) + with pytest.raises(TypeError, match=r'xt\.Table'): + xf.TouschekManager( + line, + local_momentum_acceptance=bad, + sigma_z=SIGMA_Z, sigma_delta=SIGMA_DELTA, + bunch_population=BUNCH_POPULATION, n_simulated=10, + nemitt_x=NEMITT_X, nemitt_y=NEMITT_Y, + ) + + def test_raises_on_missing_lma_columns(self, toy_ring): + """LMA table missing required columns must raise ValueError.""" + line = toy_ring['line'] + lma = toy_ring['lma'] + # Build an xt.Table without 'deltap' + bad = xt.Table({'name': lma.name, 's': lma.s, 'deltan': lma.deltan}) + with pytest.raises(ValueError, match=r'missing columns'): + xf.TouschekManager( + line, + local_momentum_acceptance=bad, + sigma_z=SIGMA_Z, sigma_delta=SIGMA_DELTA, + bunch_population=BUNCH_POPULATION, n_simulated=10, + nemitt_x=NEMITT_X, nemitt_y=NEMITT_Y, + ) + + def test_raises_on_lma_nan_values(self, toy_ring): + """LMA table with NaN values must raise ValueError.""" + line = toy_ring['line'] + lma = toy_ring['lma'] + deltan_bad = lma.deltan.copy().astype(float) + deltan_bad[0] = np.nan + bad = xt.Table({'name': lma.name, 's': lma.s, + 'deltan': deltan_bad, 'deltap': lma.deltap}) + with pytest.raises(ValueError, match=r'NaN'): + xf.TouschekManager( + line, + local_momentum_acceptance=bad, + sigma_z=SIGMA_Z, sigma_delta=SIGMA_DELTA, + bunch_population=BUNCH_POPULATION, n_simulated=10, + nemitt_x=NEMITT_X, nemitt_y=NEMITT_Y, + ) + + def test_raises_when_no_touschek_elements(self, toy_ring): + """A line with no TouschekScattering must raise ValueError.""" + env2 = xt.Environment() + bare = env2.new_line(components=[env2.new('d1', xt.Drift, length=1.0)]) + bare.set_particle_ref('electron', p0c=1e9) + with pytest.raises(ValueError, match=r'TouschekScattering'): + xf.TouschekManager( + bare, + local_momentum_acceptance=toy_ring['lma'], + sigma_z=SIGMA_Z, sigma_delta=SIGMA_DELTA, + bunch_population=BUNCH_POPULATION, n_simulated=10, + nemitt_x=NEMITT_X, nemitt_y=NEMITT_Y, + ) + + +class TestTouschekManagerInitialise: + """Tests for TouschekManager.initialise_touschek().""" + + @pytest.fixture(autouse=True) + def _setup(self, toy_ring): + """Build a fresh manager and initialise it.""" + self.line = toy_ring['line'] + self.tw = toy_ring['twiss'] + tab = self.line.get_table() + self.tnames = tab.rows.match(element_type='TouschekScattering').name + lma = _fresh_lma(toy_ring) + self.tm = _build_manager(self.line, lma, self.tw) + self.tm.initialise_touschek() + + def test_all_elements_are_configured(self): + assert len(self.tnames) > 0 + for nn in self.tnames: + assert isinstance(self.line[nn], xf.TouschekScattering) + + def test_p0c_positive_and_finite(self): + for nn in self.tnames: + el = self.line[nn] + assert np.isfinite(el.p0c) and el.p0c > 0 + + def test_integrated_piwinski_rate_nonnegative(self): + for nn in self.tnames: + el = self.line[nn] + assert np.isfinite(el.integrated_piwinski_rate) + assert el.integrated_piwinski_rate >= 0 + + def test_local_piwinski_rate_nonnegative(self): + for nn in self.tnames: + el = self.line[nn] + assert np.isfinite(el.piwinski_rate) + assert el.piwinski_rate >= 0 + + def test_lma_sign_convention(self): + """deltaN ≤ 0 and deltaP ≥ 0 must hold at every element.""" + for nn in self.tnames: + el = self.line[nn] + assert el.deltaN <= 0, f'{nn}: deltaN={el.deltaN} must be ≤ 0' + assert el.deltaP >= 0, f'{nn}: deltaP={el.deltaP} must be ≥ 0' + + def test_beam_params_stored_correctly(self): + for nn in self.tnames: + el = self.line[nn] + assert el.bunch_population == pytest.approx(BUNCH_POPULATION) + assert el.sigma_z == pytest.approx(SIGMA_Z) + assert el.sigma_delta == pytest.approx(SIGMA_DELTA) + + def test_nz_not_increased_by_nz_eff_logic(self): + """The nz-clamp must never raise nz above the requested value of 3.""" + for nn in self.tnames: + el = self.line[nn] + assert el.nz <= 3.0 + 1e-12, \ + f'{nn}: nz={el.nz} exceeds requested nz=3.0' + + def test_twiss_populated_after_initialise(self): + assert self.tm.twiss is not None + + def test_partial_initialise_single_element(self): + """ + Calling initialise_touschek(element=nn) should (re-)configure only + that element. Doubling bunch_population must increase the rate + because the Piwinski rate scales as N_b^2. + """ + nn = self.tnames[0] + el = self.line[nn] + old_rate = el.integrated_piwinski_rate + self.tm.bunch_population *= 2 + self.tm.initialise_touschek(element=nn) + assert el.integrated_piwinski_rate > old_rate + + def test_partial_initialise_raises_on_wrong_type(self): + with pytest.raises(TypeError, match=r'string'): + self.tm.initialise_touschek(element=42) + + +class TestTouschekScattering: + """Tests for the TouschekScattering.scatter() method.""" + + @pytest.fixture(autouse=True) + def _setup(self, toy_ring): + self.line = toy_ring['line'] + tw = toy_ring['twiss'] + lma = _fresh_lma(toy_ring) + tm = _build_manager(self.line, lma, tw) + tm.initialise_touschek() + + # Scatter from the first TouschekScattering element only (fast) + tab = self.line.get_table() + tnames = tab.rows.match(element_type='TouschekScattering').name + self.nn = tnames[0] + self.el = self.line[self.nn] + self.parts = self.el.scatter() + self.alive = self.parts.filter(self.parts.state == 1) + + def test_scatter_returns_particles(self): + assert isinstance(self.parts, xt.Particles) + + def test_some_particles_selected(self): + assert len(self.alive.x) > 0 + + def test_coordinates_finite(self): + for attr in ('x', 'px', 'y', 'py', 'zeta', 'delta'): + vals = getattr(self.alive, attr) + assert np.all(np.isfinite(vals)), f'Non-finite values in {attr}' + + def test_weights_finite_and_nonnegative(self): + assert np.all(np.isfinite(self.alive.weight)) + assert np.all(self.alive.weight >= 0) + + def test_at_element_matches_line_index(self): + expected = self.line.element_names.index(self.nn) + assert int(self.alive.at_element[0]) == expected + + def test_scattered_delta_outside_lma(self): + """ + The C kernel selects only particles whose delta lies outside the LMA; + every returned particle must satisfy delta < deltaN or delta > deltaP. + """ + d = self.alive.delta + outside = (d < self.el.deltaN) | (d > self.el.deltaP) + assert np.all(outside), 'Some scattered particles have delta inside the LMA' + + def test_total_mc_rate_recorded(self): + assert np.isfinite(self.el.total_mc_rate) + assert self.el.total_mc_rate >= 0 + + def test_ignored_rate_fraction(self): + """ignored_rate must equal ignored_portion * total_mc_rate.""" + assert self.el.ignored_rate == pytest.approx( + self.el.ignored_portion * self.el.total_mc_rate, rel=1e-9) + + def test_theta_log_populated(self): + """theta_log must map particle IDs → scattering angles in (0, π).""" + assert isinstance(self.el.theta_log, dict) + assert len(self.el.theta_log) > 0 + for pid, theta in self.el.theta_log.items(): + assert np.isfinite(theta) + assert 0 < theta < np.pi, f'theta={theta} out of (0, π)' + + def test_weight_sum_is_finite_and_positive(self): + assert np.isfinite(self.alive.weight.sum()) + assert self.alive.weight.sum() > 0 + + +class TestPiwinskiIntegral: + """ + Unit tests for TouschekCalculator._compute_piwinski_integral. + """ + + @pytest.fixture(autouse=True) + def _setup(self, toy_ring): + lma = _fresh_lma(toy_ring) + tm = _build_manager(toy_ring['line'], lma, toy_ring['twiss']) + self.calc = tm.touschek + + def test_integral_positive(self): + val = self.calc._compute_piwinski_integral(0.01, B1=5.0, B2=3.0) + assert val > 0 + + def test_integral_decreases_with_larger_tm(self): + """A larger momentum cut-off (larger tm) must give a smaller integral.""" + B1, B2 = 5.0, 3.0 + assert (self.calc._compute_piwinski_integral(0.001, B1, B2) > + self.calc._compute_piwinski_integral(0.10, B1, B2)) + + def test_integral_decreases_with_larger_B1(self): + """Larger B1 (tighter beam) must suppress the integral.""" + tm, B2 = 0.01, 2.0 + assert (self.calc._compute_piwinski_integral(tm, B1= 3.0, B2=B2) > + self.calc._compute_piwinski_integral(tm, B1=10.0, B2=B2)) + + def test_integral_finite_for_large_B2_t(self): + """The asymptotic I0 branch (B2*t > 500) must return a finite positive value.""" + val = self.calc._compute_piwinski_integral(0.001, B1=600.0, B2=599.0) + assert np.isfinite(val) and val > 0 + + +class TestEndToEndLifetime: + """ + Minimal end-to-end smoke test: scatter, track, merge, lifetime. + Verifies that the result is a physically plausible positive finite number. + """ + + def test_lifetime_positive_finite(self, toy_ring): + line = toy_ring['line'] + tw = toy_ring['twiss'] + lma = _fresh_lma(toy_ring) + + tm = _build_manager(line, lma, tw, n_simulated=int(2e5)) + tm.initialise_touschek() + + tab = line.get_table() + tnames = tab.rows.match(element_type='TouschekScattering').name + + particles_list = [] + for nn in tnames: + parts = line[nn].scatter() + line.track(parts, ele_start=nn, ele_stop=nn, num_turns=128) + particles_list.append(parts) + + merged = xt.Particles.merge(particles_list) + lost = merged.filter(merged.state == 0) + loss_rate = float(sum(lost.weight)) + + assert loss_rate > 0, 'No particles were lost — something is wrong' + + lifetime = BUNCH_POPULATION / loss_rate + assert np.isfinite(lifetime) and lifetime > 0 \ No newline at end of file diff --git a/xfields/__init__.py b/xfields/__init__.py index 2e009768..6796eb40 100644 --- a/xfields/__init__.py +++ b/xfields/__init__.py @@ -16,6 +16,8 @@ from .solvers.fftsolvers import FFTSolver3D +from .touschek.manager import TouschekManager + from .beam_elements.spacecharge import SpaceCharge3D, SpaceChargeBiGaussian from .beam_elements.beambeam2d import BeamBeamBiGaussian2D from .beam_elements.beambeam2d import ConfigForUpdateBeamBeamBiGaussian2D @@ -25,6 +27,7 @@ from .beam_elements.temp_slicer import TempSlicer from .beam_elements.electroncloud import ElectronCloud from .beam_elements.electronlens_interpolated import ElectronLensInterpolated +from .beam_elements.touschek import TouschekScattering from .beam_elements.transverse_damper import TransverseDamper from .beam_elements.collective_monitor import CollectiveMonitor from .beam_elements.waketracker import WakeTracker diff --git a/xfields/beam_elements/touschek.py b/xfields/beam_elements/touschek.py new file mode 100644 index 00000000..891d8286 --- /dev/null +++ b/xfields/beam_elements/touschek.py @@ -0,0 +1,379 @@ +# copyright ################################# # +# This file is part of the Xfields Package. # +# Copyright (c) CERN, 2021. # +# ########################################### # + +import xobjects as xo +import xtrack as xt +import numpy as np + +class TouschekScattering(xt.BeamElement): + """ + Beam element that performs a Monte Carlo Touschek scattering simulation + at a single location in a lattice. + + Each element represents one scattering center along the lattice. When + :meth:`scatter` is called it draws macro-particle pairs from the local + 6D phase-space (Gaussian) distribution, applies the Møller cross-section, + boosts scattered pairs back to the lab frame, and returns the subset of + macro-particles whose momentum deviation exceeds the local momentum + acceptance (LMA). + + The element is *passive* during normal tracking (``track`` is a no-op); + all physics happens inside :meth:`scatter`. + + The Monte Carlo kernel is implemented in C99 and follows the ELEGANT + algorithm of Xiao & Borland (PRSTAB 13, 074201, 2010). + + Parameters + ---------- + s : float, optional + Longitudinal position of the element in the lattice [m]. Default 0. + particle_ref : xtrack.Particles, optional + Reference particle carrying. + element_index : int, optional + Index of this element in the line. + bunch_population : float, optional + Number of particles in one bunch. + alfx, betx : float, optional + Horizontal Twiss parameters at the element. + alfy, bety : float, optional + Vertical Twiss parameters at the element. + dx, dpx : float, optional + Horizontal dispersion and its derivative at the element. + dy, dpy : float, optional + Vertical dispersion and its derivative at the element. + x_co, px_co : float, optional + Horizontal closed-orbit position and normalised momentum at the + element. + y_co, py_co : float, optional + Vertical closed-orbit position and normalised momentum at the + element. + zeta_co, delta_co : float, optional + Longitudinal closed-orbit coordinate and relative momentum + deviation. + deltaN : float, optional + Negative local momentum acceptance (scaled). + deltaP : float, optional + Positive local momentum acceptance (scaled). + gemitt_x : float, optional + Horizontal geometric emittance [m·rad]. + gemitt_y : float, optional + Vertical geometric emittance [m·rad]. + sigma_z : float, optional + RMS bunch length [m]. + sigma_delta : float, optional + RMS relative momentum spread. + n_simulated : int, optional + Number of macro-particles (scattered candidates) to generate in the + Monte Carlo loop. Larger values reduce statistical noise but + increase CPU time. + nx, ny, nz : float, optional + Truncation of the Gaussian distribution in units of + :math:`\\sqrt{\\varepsilon}` for the transverse planes and + :math:`\\sigma` for the longitudinal plane. The sampling window is + :math:`\\pm n_x \\sqrt{\\varepsilon_x}`, etc. ``nz`` may be + reduced automatically by :class:`TouschekManager` to prevent + particles being drawn outside the LMA before scattering. + theta_min, theta_max : float, optional + Lower and upper limits of the centre-of-mass scattering angle + :math:`\\theta^*` [rad]. In practice set to + :math:`0.00005\\pi` and :math:`0.99995\\pi` to avoid the + forward/backward divergence of the Møller cross-section. + piwinski_rate : float, optional + Local Piwinski scattering rate [Hz] evaluated at this element. + Stored for diagnostics; not used in the Monte Carlo kernel. + ignored_portion : float, optional + Fraction of the total simulated scattering weight that is discarded + before tracking. Only the highest-weight particles whose cumulative + weight reaches ``(1 - ignored_portion)`` of the total are retained + and tracked; the remaining low-weight, low-probability events are + dropped. The default value of ``0.01`` retains 99 % of the total + weight while significantly reducing the number of particles that must + be tracked, providing a good accuracy–efficiency trade-off. + Setting this to ``0`` keeps all simulated particles. + integrated_piwinski_rate : float, optional + Piwinski rate integrated (trapezoidal rule) over the lattice section + preceding this element and divided by :math:`c` and + :math:`T_{\\text{rev}}` to give a per-bunch, per-turn rate [1/s]. + Set by :meth:`TouschekManager.initialise_touschek`; used to weight + the scattered macro-particles. + seed : int, optional + Seed for the ELEGANT-compatible 48-bit LCG random number generator. + Using the same seed reproduces the ELEGANT Monte Carlo sequence + exactly. Default 1997. + inhibit_permute : int, optional + If non-zero, the random-order permutation step (``randomizeOrder``) + is skipped. Intended for reproducibility testing only. + + Attributes + ---------- + piwinski_rate : float + Local Piwinski scattering rate [Hz] at this element. + total_mc_rate : float + Total Monte Carlo scattering rate [Hz] returned by the last call to + :meth:`scatter`. + ignored_rate : float + Scattering rate [Hz] associated with the low-weight particles + discarded by ``pickPart`` in the last call to :meth:`scatter`, + i.e. the fraction ``ignored_portion`` of ``total_mc_rate`` that + is not represented in the tracked particle set. + theta_log : dict + Mapping ``{particle_id: theta}`` of centre-of-mass scattering + angles [rad] for the particles returned by the last call to + :meth:`scatter`. + + Notes + ----- + **Physics summary** + + The Monte Carlo loop follows Xiao & Borland (PRSTAB 13, 074201, 2010): + + 1. Two particles are drawn from the local 6-D Gaussian distribution + using ``selectPartGauss`` with a truncated range of + :math:`\\pm n \\sqrt{\\varepsilon}`. + 2. The pair is boosted to the centre-of-mass (CM) frame + (``bunch2cm``). + 3. A scattering angle :math:`\\theta^*` is drawn uniformly in + :math:`[\\theta_{\\min},\\,\\theta_{\\max}]` and a random azimuthal + angle :math:`\\phi` is drawn uniformly in :math:`[0,\\,\\pi]`. + 4. The Møller cross-section ``moeller`` is evaluated at + :math:`\\theta^*`. + 5. Scattered momenta are rotated (``eulertrans``) and boosted back to + the lab frame (``cm2bunch``). + 6. A particle is selected for tracking only if its resulting + :math:`\\delta` falls outside the LMA: + :math:`\\delta < \\delta_N` or :math:`\\delta > \\delta_P`. + 7. ``pickPart`` retains only the highest-weight particles whose + cumulative weight reaches ``(1 - ignored_portion)`` of the + total simulated weight. The remaining low-weight, + low-probability events are discarded. With the default + ``ignored_portion = 0.01``, 99 % of the total weight is + retained, significantly reducing the number of particles + that must be tracked. + + Macro-particle weights are normalised so that + :math:`\\sum_i w_i` equals the per-turn loss rate in the + corresponding lattice section (in particles/turn). + + References + ---------- + .. [1] A. Xiao and M. Borland, "Monte Carlo simulation of Touschek + effect", Phys. Rev. ST Accel. Beams **13**, 074201 (2010). + https://doi.org/10.1103/PhysRevSTAB.13.074201 + .. [2] M. Borland, "elegant: A Flexible SDDS-Compliant Code for + Accelerator Simulation", APS LS-287 (2000). + """ + + _xofields = { + 'p0c': xo.Float64, + 'bunch_population': xo.Float64, + 'gemitt_x': xo.Float64, + 'gemitt_y': xo.Float64, + 'alfx': xo.Float64, + 'betx': xo.Float64, + 'alfy': xo.Float64, + 'bety': xo.Float64, + 'dx': xo.Float64, + 'dpx': xo.Float64, + 'dy': xo.Float64, + 'dpy': xo.Float64, + 'deltaN': xo.Float64, + 'deltaP': xo.Float64, + 'sigma_z': xo.Float64, + 'sigma_delta': xo.Float64, + 'n_simulated': xo.Int64, + 'nx': xo.Float64, + 'ny': xo.Float64, + 'nz': xo.Float64, + 'theta_min': xo.Float64, + 'theta_max': xo.Float64, + 'ignored_portion': xo.Float64, + 'integrated_piwinski_rate': xo.Float64, + 'seed': xo.Int64, + 'inhibit_permute': xo.Int64 + } + + # allow_track = False + _depends_on = [xt.RandomUniformAccurate] + + _extra_c_sources = [ + '#include "xfields/beam_elements/touschek_src/touschek.h"' + ] + + _per_particle_kernels = { + '_scatter': xo.Kernel( + c_name='TouschekScatter', + args=[ + xo.Arg(xo.Float64, name='x_out', pointer=True), + xo.Arg(xo.Float64, name='px_out', pointer=True), + xo.Arg(xo.Float64, name='y_out', pointer=True), + xo.Arg(xo.Float64, name='py_out', pointer=True), + xo.Arg(xo.Float64, name='zeta_out', pointer=True), + xo.Arg(xo.Float64, name='delta_out', pointer=True), + xo.Arg(xo.Float64, name='theta_out', pointer=True), + xo.Arg(xo.Float64, name='weight_out', pointer=True), + xo.Arg(xo.Float64, name='totalMCRate_out', pointer=True), + xo.Arg(xo.Int64, name='n_selected_out', pointer=True), + ], + ), + } + + def __init__(self, s=0.0, + particle_ref=xt.Particles(), + element_index=0, + bunch_population=0.0, + alfx=0.0, betx=0.0, alfy=0.0, bety=0.0, + dx=0.0, dpx=0.0, dy=0.0, dpy=0.0, + x_co=0.0, px_co=0.0, y_co=0.0, py_co=0.0, + zeta_co=0.0, delta_co=0.0, + deltaN=0.0, deltaP=0.0, + gemitt_x=0.0, gemitt_y=0.0, + sigma_z=0.0, sigma_delta=0.0, + n_simulated=0, nx=0.0, ny=0.0, nz=0.0, + theta_min=0.0, theta_max=0.0, + piwinski_rate=0.0, + ignored_portion=0.0, + integrated_piwinski_rate=0.0, + seed=1997, + inhibit_permute=0, + **kwargs): + + # This gives AttributeError: 'TouschekScattering' object has no attribute '_xobject' + # if not isinstance(self._context, xo.ContextCpu) or self._context.openmp_enabled: + # raise ValueError('TouschekScattering only enabled on CPU.') + + if '_xobject' in kwargs.keys(): + self.xoinitialize(**kwargs) + return + + super().__init__(**kwargs) + + self.s = s + self.particle_ref = particle_ref + self.element_index = element_index + self.bunch_population = bunch_population + self.alfx = alfx + self.betx = betx + self.alfy = alfy + self.bety = bety + self.dx = dx + self.dpx = dpx + self.dy = dy + self.dpy = dpy + self.x_co = x_co + self.px_co = px_co + self.y_co = y_co + self.py_co = py_co + self.zeta_co = zeta_co + self.delta_co = delta_co + self.deltaN = deltaN + self.deltaP = deltaP + self.gemitt_x = gemitt_x + self.gemitt_y = gemitt_y + self.sigma_z = sigma_z + self.sigma_delta = sigma_delta + self.n_simulated = n_simulated + self.nx = nx + self.ny = ny + self.nz = nz + self.theta_min = theta_min + self.theta_max = theta_max + self.ignored_portion = ignored_portion + self.integrated_piwinski_rate = integrated_piwinski_rate + self.piwinski_rate = piwinski_rate + self.seed = seed + self.inhibit_permute = inhibit_permute + + def _configure(self, **kwargs): + config_allowed = { + "s", "particle_ref", "element_index", + "bunch_population", + "gemitt_x", "gemitt_y", + "alfx", "betx", "alfy", "bety", + "dx", "dpx", "dy", "dpy", + "x_co", "px_co", "y_co", "py_co", + "zeta_co", "delta_co", + "deltaN", "deltaP", + "sigma_z", "sigma_delta", + "n_simulated", "nx", "ny", "nz", + "theta_min", "theta_max", + "ignored_portion", "piwinski_rate", + "integrated_piwinski_rate", + "seed", "inhibit_permute" + } + + unknown = set(kwargs) - config_allowed + if unknown: + bad = ", ".join(sorted(unknown)) + raise KeyError(f"Unsupported configure() keys: {bad}") + + for kk, vv in kwargs.items(): + setattr(self, kk, vv) + if kk == "particle_ref": + self.p0c = self.particle_ref.p0c[0] + + def scatter(self): + context = self._context + particles = xt.Particles(_context=context) + + if not particles._has_valid_rng_state(): + particles._init_random_number_generator() + + x_out = context.zeros(shape=(self.n_simulated,), dtype=np.float64) + px_out = context.zeros(shape=(self.n_simulated,), dtype=np.float64) + y_out = context.zeros(shape=(self.n_simulated,), dtype=np.float64) + py_out = context.zeros(shape=(self.n_simulated,), dtype=np.float64) + zeta_out = context.zeros(shape=(self.n_simulated,), dtype=np.float64) + delta_out = context.zeros(shape=(self.n_simulated,), dtype=np.float64) + theta_out = context.zeros(shape=(self.n_simulated,), dtype=np.float64) + weight_out = context.zeros(shape=(self.n_simulated,), dtype=np.float64) + totalMCRate_out = context.zeros(shape=(1,), dtype=np.float64) + n_selected_out = context.zeros(shape=(1,), dtype=np.int64) + + self._scatter(particles=particles, + x_out=x_out, px_out=px_out, + y_out=y_out, py_out=py_out, + zeta_out=zeta_out, delta_out=delta_out, + theta_out=theta_out, + weight_out=weight_out, + totalMCRate_out=totalMCRate_out, + n_selected_out=n_selected_out) + + n = n_selected_out[0] + # Create particle object for tracking + # TODO: add at_element, start_tracking_at_element, ... + part = xt.Particles(_capacity=2*n, + p0c=self.p0c, + mass0=self.particle_ref.mass0, + q0=self.particle_ref.q0, + pdg_id=self.particle_ref.pdg_id, + x=x_out[:n], px=px_out[:n], + y=y_out[:n], py=py_out[:n], + zeta=zeta_out[:n], delta=delta_out[:n], + weight=weight_out[:n], + s=getattr(self, 's', 0.0)) + + # Shift Touschek scattered particles around the closed orbit + part.x[:n] += self.x_co + part.px[:n] += self.px_co + part.y[:n] += self.y_co + part.py[:n] += self.py_co + part.zeta[:n] += self.zeta_co + + delta_temp = part.delta.copy() + delta_temp[:n] += self.delta_co + part.update_delta(delta_temp) + + part.at_element = self.element_index + + part_ids = part.filter(part.state == 1).particle_id + self.theta_log = dict(zip(part_ids.astype(int), theta_out[:n].astype(float))) + + self.total_mc_rate = totalMCRate_out[0] + self.ignored_rate = self.ignored_portion * self.total_mc_rate + + return part + + def track(self, particles): + super().track(particles) \ No newline at end of file diff --git a/xfields/beam_elements/touschek_src/touschek.h b/xfields/beam_elements/touschek_src/touschek.h new file mode 100644 index 00000000..4ad8f39f --- /dev/null +++ b/xfields/beam_elements/touschek_src/touschek.h @@ -0,0 +1,560 @@ +/* + * touschek.h — Touschek scattering kernel (C99, header-only) + * + * Overview + * -------- + * Header-only C99 implementation of the Monte Carlo Touschek scattering + * routine for Xsuite/xfields. The physics and selection logic follow the + * Elegant implementation by A. Xiao and M. Borland (PRSTAB 13, 074201, 2010), + * with a simplified API suitable for xobjects; SDDS I/O is omitted. + * + * Provenance (portions adapted from) + * ---------------------------------- + * - Elegant: src/touschekScatter.c (A. Xiao, M. Borland) + * - SDDS Toolkit utilities: mdbmth/drand.c, mdbmth/dlaran.c + * - LAPACK DLARAN (48-bit LCG RNG core) + * + * RNG compatibility + * ----------------- + * RNG streams are chosen to reproduce Elegant’s sequences. + * See: xfields/xfields/headers/elegant_rng.h + * (uses random_1_elegant, random_4, randomizeOrder, and DLARAN). + * + * Summary of modifications vs Elegant + * ----------------------------------- + * - Converted to a header-only C99 kernel; simplified API; no SDDS I/O. + * - Integrated with xobjects; clarified slope (xp, yp) vs momentum usage. + * - Minor safety/cleanup (bounds checks, allocations, comments). + * - Physics and selection logic preserved. + * + * References + * ---------- + * - M. Borland, “elegant: A Flexible SDDS-Compliant Code for Accelerator Simulation,” + * APS LS-287 (2000). + * - A. Xiao and M. Borland, "Monte Carlo simulation of Touschek effect", + * Phys. Rev. ST Accel. Beams **13**, 074201 (2010). + */ +#ifndef XTRACK_TOUSCHEK_H +#define XTRACK_TOUSCHEK_H + +#include "xtrack/headers/track.h" +#include "xfields/headers/elegant_rng.h" + +#include +#include +#include +#include +#include + +static inline double sqr(double x){ return x*x; } + +/*gpufun*/ +void TouschekScattering_track_local_particle(TouschekScatteringData el, LocalParticle* part0) { + (void)el; (void)part0; + return; +} + +/* Adapted from Elegant touschekScatter.c: selectPartGauss (logic unchanged). */ +void selectPartGauss(double *p1, double *p2, + double *dens1, double *dens2, + const double *ran1, + const double *range, + const double *alfa, + const double *beta, + const double *disp, + const double *gemitt) { + int i; + double U[3], V1[3], V2[3], densa[3], densb[3]; + + /* Select random particle coordinates in normalized phase space */ + for (i = 0; i < 3; i++) { + U[i] = (ran1[i] - 0.5) * range[i] * sqrt(gemitt[i]); + V1[i] = (ran1[i + 3] - 0.5) * range[i] * sqrt(gemitt[i]); + V2[i] = (ran1[i + 6] - 0.5) * range[i] * sqrt(gemitt[i]); + densa[i] = exp(-0.5 * (U[i] * U[i] + V1[i] * V1[i]) / gemitt[i]); + } + densb[2] = exp(-0.5 * (U[2] * U[2] + V2[2] * V2[2]) / gemitt[2]); + /* Transform particle coordinates from normalized to real phase space */ + for (i = 0; i < 3; i++) { + p1[i] = p2[i] = sqrt(beta[i]) * U[i]; + p1[i + 3] = (V1[i] - alfa[i] * U[i]) / sqrt(beta[i]); + p2[i + 3] = (V2[i] - alfa[i] * U[i]) / sqrt(beta[i]); + } + /* Dispersion correction */ + p1[0] = p1[0] + p1[5] * disp[0]; + p1[1] = p1[1] + p1[5] * disp[1]; + p1[3] = p1[3] + p1[5] * disp[2]; + p1[4] = p1[4] + p1[5] * disp[3]; + + p2[0] = p1[0] - p2[5] * disp[0]; + p2[1] = p1[1] - p2[5] * disp[1]; + U[0] = p2[0] / sqrt(beta[0]); + U[1] = p2[1] / sqrt(beta[1]); + p2[3] = (V2[0] - alfa[0] * U[0]) / sqrt(beta[0]); + p2[4] = (V2[1] - alfa[1] * U[1]) / sqrt(beta[1]); + densb[0] = exp(-0.5 * (U[0] * U[0] + V2[0] * V2[0]) / gemitt[0]); + densb[1] = exp(-0.5 * (U[1] * U[1] + V2[1] * V2[1]) / gemitt[1]); + + p2[0] = p1[0]; + p2[1] = p1[1]; + p2[3] = p2[3] + p2[5] * disp[2]; + p2[4] = p2[4] + p2[5] * disp[3]; + + *dens1 = densa[0] * densa[1] * densa[2]; + *dens2 = densb[0] * densb[1] * densb[2]; + + return; +} + +/* From Elegant touschekScatter.c: bunch2cm */ +void bunch2cm(double *p1, double *p2, double *q, double *beta, double *gamma) { + double pp1, pp2, e1, e2, ee; + int i; + double bb, betap1, factor; + + pp1 = 0.0; + pp2 = 0.0; + for (i = 3; i < 6; i++) { + pp1 = pp1 + sqr(p1[i]); + pp2 = pp2 + sqr(p2[i]); + } + e1 = sqrt(ELECTRON_MASS_EV * ELECTRON_MASS_EV + pp1); + e2 = sqrt(ELECTRON_MASS_EV * ELECTRON_MASS_EV + pp2); + ee = e1 + e2; + + betap1 = 0.0; + bb = 0.0; + for (i = 0; i < 3; i++) { + beta[i] = (p1[i + 3] + p2[i + 3]) / ee; + betap1 = betap1 + beta[i] * p1[i + 3]; + bb = bb + beta[i] * beta[i]; + } + + *gamma = 1. / sqrt(1. - bb); + factor = ((*gamma) - 1.) * betap1 / bb; + + for (i = 0; i < 3; i++) { + q[i] = p1[i + 3] + factor * beta[i] - (*gamma) * e1 * beta[i]; + } + + return; +} + + +/* Rotate scattered p in c.o.m system */ +/* From Elegant touschekScatter.c: eulertrans*/ +void eulertrans(double *v0, double theta, double phi, double *v1, double *v) { + double th, ph, s1, s2, c1, c2; + double x0, y0, z0; + + *v = sqrt(v0[0] * v0[0] + v0[1] * v0[1] + v0[2] * v0[2]); + th = acos(v0[2] / (*v)); + ph = atan2(v0[1], v0[0]); + + s1 = sin(th); + s2 = sin(ph); + c1 = cos(th); + c2 = cos(ph); + + x0 = cos(theta); + y0 = sin(theta) * cos(phi); + z0 = sin(theta) * sin(phi); + + v1[0] = (*v) * (s1 * c2 * x0 - s2 * y0 - c1 * c2 * z0); + v1[1] = (*v) * (s1 * s2 * x0 + c2 * y0 - c1 * s2 * z0); + v1[2] = (*v) * (c1 * x0 + s1 * z0); + + return; +} + +/* From Elegant touschekScatter.c: cm2bunch*/ +void cm2bunch(double *p1, double *p2, double *q, double *beta, double *gamma) { + int i; + double pq, e, betaq, bb, factor; + + pq = 0.0; + for (i = 0; i < 3; i++) { + pq = pq + q[i] * q[i]; + } + + e = sqrt(ELECTRON_MASS_EV * ELECTRON_MASS_EV + pq); + + betaq = 0.0; + bb = 0.0; + for (i = 0; i < 3; i++) { + betaq = betaq + beta[i] * q[i]; + bb = bb + beta[i] * beta[i]; + } + + factor = ((*gamma) - 1) * betaq / bb; + for (i = 0; i < 3; i++) { + p1[i + 3] = q[i] + (*gamma) * beta[i] * e + factor * beta[i]; + p2[i + 3] = -q[i] + (*gamma) * beta[i] * e - factor * beta[i]; + } + + return; +} + +/* From Elegant touschekScatter.c: moeller */ +double moeller(double beta0, double theta) { + double cross; + double beta2, st2; + + beta2 = beta0 * beta0; + st2 = sqr(sin(theta)); + + cross = (1. - beta2) * (sqr(1. + 1. / beta2) * (4. / st2 / st2 - 3. / st2) + 1. + 4. / st2); + + return cross; +} + +/* From Elegant touschekScatter.c: pickPart */ +void pickPart(double *weight, long *index, long start, long end, + long *iTotal, double *wTotal, double weight_limit, double weight_ave) { + long i, i1, i2, N; + double w1, w2; + long *index1, *index2; + double *weight1, *weight2; + + i1 = i2 = 0; + w1 = w2 = 0.; + N = end - start; + if (N < 3) + return; /* scattered particles normally appear in pair */ + index2 = (long *)malloc(sizeof(long) * N); + weight2 = (double *)malloc(sizeof(double) * N); + index1 = (long *)malloc(sizeof(long) * N); + weight1 = (double *)malloc(sizeof(double) * N); + + for (i = start; i < end; i++) { + if (weight[i] > weight_ave) { + weight2[i2] = weight[i]; + index2[i2++] = index[i]; + w2 += weight[i]; + } else { + weight1[i1] = weight[i]; + index1[i1++] = index[i]; + w1 += weight[i]; + } + } + if ((w2 + (*wTotal)) > weight_limit) { + weight_ave = w2 / (double)i2; + for (i = 0; i < i2; i++) { + index[start + i] = index2[i]; + weight[start + i] = weight2[i]; + } + free(weight1); + free(index1); + free(weight2); + free(index2); + pickPart(weight, index, start, start + i2, + iTotal, wTotal, weight_limit, weight_ave); + return; + } + + *iTotal += i2; + *wTotal += w2; + weight_ave = w1 / (double)i1; + for (i = 0; i < i2; i++) { + index[start + i] = index2[i]; + weight[start + i] = weight2[i]; + } + for (i = 0; i < i1; i++) { + index[start + i2 + i] = index1[i]; + weight[start + i2 + i] = weight1[i]; + } + free(weight1); + free(index1); + free(weight2); + free(index2); + pickPart(weight, index, i2 + start, end, + iTotal, wTotal, weight_limit, weight_ave); + return; +} + +/* Adapted from Elegant touschekScatter.c: TouschekDistribution (logic unchanged) */ +void TouschekScatter(TouschekScatteringData el, + LocalParticle* part0, + double* x_out, + double* px_out, + double* y_out, + double* py_out, + double* zeta_out, + double* delta_out, + double* theta_out, + double* weight_out, + double* totalMCRate_out, + int64_t* n_selected_out){ + + const double p0c = TouschekScatteringData_get_p0c(el); + const double bunch_population = TouschekScatteringData_get_bunch_population(el); + const double gemitt_x = TouschekScatteringData_get_gemitt_x(el); + const double gemitt_y = TouschekScatteringData_get_gemitt_y(el); + const double alfx = TouschekScatteringData_get_alfx(el); + const double betx = TouschekScatteringData_get_betx(el); + const double alfy = TouschekScatteringData_get_alfy(el); + const double bety = TouschekScatteringData_get_bety(el); + const double dx = TouschekScatteringData_get_dx(el); + const double dpx = TouschekScatteringData_get_dpx(el); + const double dy = TouschekScatteringData_get_dy(el); + const double dpy = TouschekScatteringData_get_dpy(el); + const double deltaN = TouschekScatteringData_get_deltaN(el); + const double deltaP = TouschekScatteringData_get_deltaP(el); + const double sigma_z = TouschekScatteringData_get_sigma_z(el); + const double sigma_delta = TouschekScatteringData_get_sigma_delta(el); + const double n_simulated = TouschekScatteringData_get_n_simulated(el); + const double nx = TouschekScatteringData_get_nx(el); + const double ny = TouschekScatteringData_get_ny(el); + const double nz = TouschekScatteringData_get_nz(el); + const double theta_min = TouschekScatteringData_get_theta_min(el); + const double theta_max = TouschekScatteringData_get_theta_max(el); + const double ignoredPortion = TouschekScatteringData_get_ignored_portion(el); + const double integrated_piwinski_rate = TouschekScatteringData_get_integrated_piwinski_rate(el); + + long i, j, total_event, simuCount, iTotal; + double ran1[11]; + + long *index = NULL; + double *weight = (double*)malloc(sizeof(double) * n_simulated); + + const double twissAlpha[3] = { alfx, alfy, 0.0 }; + + const double bets = sigma_z/sigma_delta; + const double twissBeta[3] = { betx, bety, bets }; + + const double twissDisp[4] = { dx, dy, dpx, dpy }; + + const double gemitt_z = sigma_z*sigma_delta; + const double gemitt[3] = { gemitt_x, gemitt_y, gemitt_z }; + const double range[3] = { 2.0*nx, 2.0*ny, 2.0*nz }; + + double totalWeight = 0.0; + double totalMCRate = 0.0; + + double pTemp[6], p1[6], p2[6], dens1, dens2; + double theta, phi, qa[3], qb[3], beta[3], qabs, gamma; + double beta0, cross, temp; + + double weight_limit, weight_ave, wTotal; + + const double sigxyz = sqrt(twissBeta[0]*gemitt[0]) * sqrt(twissBeta[1]*gemitt[1]) * sigma_z; + temp = sqr(bunch_population) * sqr(PI) * sqr(RADIUS_ELECTRON) * C_LIGHT / 4.; + double factor = temp * pow(range[0], 3.0) * pow(range[1], 3.0) * pow(range[2], 3.0) / pow(2 * PI, 6.0) / sigxyz; + + double *xtemp = (double*)malloc(sizeof(double) * n_simulated); + double *pxtemp = (double*)malloc(sizeof(double) * n_simulated); + double *ytemp = (double*)malloc(sizeof(double) * n_simulated); + double *pytemp = (double*)malloc(sizeof(double) * n_simulated); + double *zetatemp = (double*)malloc(sizeof(double) * n_simulated); + double *deltatemp = (double*)malloc(sizeof(double) * n_simulated); + double *thetatemp = (double*)malloc(sizeof(double) * n_simulated); + + static int seeded_once = 0; + if (!seeded_once){ + long seed = TouschekScatteringData_get_seed(el); + short inhibit = (short)TouschekScatteringData_get_inhibit_permute(el); + seedElegantRandomNumbers(seed, inhibit); + seeded_once = 1; + } + + i = 0; + j = 0; + total_event = 0; + simuCount = 0; + + while (1) { + if (i >= n_simulated) + break; + + /* Select 11 random numbers, then mix them. */ + + // These 11 random numbers are assigned to: + // particle 1 (p1) as: { x, y, px, py, zeta, delta} + // particle 2 (p2) as: { -, -, px, py, -, delta} + // scattering angles in the cm frame: theta and phi + + // In ELEGANT the 11 random numbers are assigned to: + // particle 1 (p1) as: { x, y, xp, yp, zeta, delta} + // particle 2 (p2) as: { -, -, xp, yp, -, delta} + // scattering angles in the cm frame: theta and phi + + // NOTE: ELEGANT uses slopes xp=dx/ds, yp=dy/ds instead of the normalized momentum components px=Px/p0c, py=Py/p0c + for (j = 0; j < 11; j++) { + // ran1[j] = RandomUniformAccurate_generate(part0); // Does not match with ELEGANT + ran1[j] = random_1_elegant(1); + } + randomizeOrder((char*)ran1, sizeof(ran1[0]), 11, 0, random_4); // like ELEGANT + + total_event++; + + selectPartGauss(p1, p2, &dens1, &dens2, ran1, range, twissAlpha, twissBeta, twissDisp, gemitt); + + if (!dens1 || !dens2) { + continue; + } + /* Here ELEGANT changes from slopes to momentum components */ + // Since we use already the normalized momentum components {px, py} instead of the slopes {xp, yp} + // here we just unormalize the momentum components: Px=px*p0c, Py=py*p0c + for (j = 3; j < 5; j++) { + p1[j] *= p0c; + p2[j] *= p0c; + } + // p1[5] = (p1[5] + 1) * p0c; + // p2[5] = (p2[5] + 1) * p0c; + // Use exact formula to compute p1[5]=Pz1 and p2[5]=Pz2 + p1[5] = sqrt(sqr(p0c)*sqr(1. + p1[5]) - sqr(p1[3]) - sqr(p1[4])); + p2[5] = sqrt(sqr(p0c)*sqr(1. + p2[5]) - sqr(p2[3]) - sqr(p2[4])); + + bunch2cm(p1, p2, qa, beta, &gamma); + + theta = theta_min + (theta_max - theta_min) * ran1[9]; + phi = ran1[10] * PI; + + temp = dens1 * dens2 * sin(theta); + eulertrans(qa, theta, phi, qb, &qabs); + cm2bunch(p1, p2, qb, beta, &gamma); + // p1[5] = (p1[5] - p0c) / p0c; + // p2[5] = (p2[5] - p0c) / p0c; + // Use exact formula to compute p1[5]=delta1 and p2[5]=delta2 + p1[5] = (sqrt(sqr(p1[3]) + sqr(p1[4]) + sqr(p1[5])) - p0c) / p0c; + p2[5] = (sqrt(sqr(p2[3]) + sqr(p2[4]) + sqr(p2[5])) - p0c) / p0c; + + if (p1[5] > p2[5]) { + for (j = 0; j < 6; j++) { + pTemp[j] = p2[j]; + p2[j] = p1[j]; + p1[j] = pTemp[j]; + } + } + + if (p1[5] < deltaN || p2[5] > deltaP) { + beta0 = qabs / sqrt(qabs * qabs + ELECTRON_MASS_EV * ELECTRON_MASS_EV); + cross = moeller(beta0, theta); + temp *= cross * beta0 / gamma / gamma; + + if (p1[5] < deltaN) { + totalWeight += temp; + p1[3] /= p0c; + p1[4] /= p0c; + simuCount++; + + xtemp[i] = p1[0]; + pxtemp[i] = p1[3]; + ytemp[i] = p1[1]; + pytemp[i] = p1[4]; + zetatemp[i] = p1[2]; + deltatemp[i] = p1[5]; + thetatemp[i] = theta; + weight[i] = temp; + i++; + } + + if (i >= n_simulated) + break; + + if (p2[5] > deltaP) { + totalWeight += temp; + p2[3] /= p0c; + p2[4] /= p0c; + simuCount++; + + xtemp[i] = p2[0]; + pxtemp[i] = p2[3]; + ytemp[i] = p2[1]; + pytemp[i] = p2[4]; + zetatemp[i] = p2[2]; + deltatemp[i] = p2[5]; + thetatemp[i] = theta; + weight[i] = temp; + i++; + } + } + } + factor = factor / (double)(total_event); + totalMCRate = totalWeight * factor; + + /* Pick tracking particles from the simulated scattered particles */ + index = (long *)malloc(sizeof(long) * simuCount); + for (i = 0; i < simuCount; i++) { + index[i] = i; + } + + if (ignoredPortion <= 1e-9) { + iTotal = simuCount; + wTotal = totalWeight; + for (long k = 0; k < iTotal; ++k) { + x_out[k] = xtemp[k]; + px_out[k] = pxtemp[k]; + y_out[k] = ytemp[k]; + py_out[k] = pytemp[k]; + zeta_out[k] = zetatemp[k]; + delta_out[k] = deltatemp[k]; + theta_out[k] = thetatemp[k]; + weight_out[k] = weight[k]; + } + } else { + iTotal = 0; + wTotal = 0.; + weight_limit = totalWeight * (1 - ignoredPortion); + weight_ave = totalWeight / simuCount; + + pickPart(weight, index, 0, simuCount, + &iTotal, &wTotal, weight_limit, weight_ave); + + for (long k = 0; k < iTotal; ++k) { + long src = index[k]; + x_out[k] = xtemp[src]; + px_out[k] = pxtemp[src]; + y_out[k] = ytemp[src]; + py_out[k] = pytemp[src]; + zeta_out[k] = zetatemp[src]; + delta_out[k] = deltatemp[src]; + theta_out[k] = thetatemp[src]; + weight_out[k] = weight[src]; + } + } + + // printf("Total number of random numbers used: %ld\n", total_event * 11); + // fflush(stdout); + if (total_event * 11 > (long)2e9) { + printf("The total number of random numbers used exceeded 2e9. Use smaller n_simulated or smaller delta."); + fflush(stdout); + } + + // if (simuCount == 0) { + // printf("It appears that the Touschek lifetime is extremely wrong and there is no need to perform Touschek simulation.\n If you think this is a wrong statement, send input to developers for evaluation.\n"); + // } + + if (total_event / simuCount > 20) { + if (nx < 5 || ny < 5) { + printf("The Touschek scattering rate is low. Please use >=5 sigma beam for better simulation."); + fflush(stdout); + } else { + printf("The Touschek scattering rate is very low. Please ignore the rate from Monte Carlo simulation. Use Piwinski rate only."); + fflush(stdout); + } + } + + printf("%ld of %ld particles selected for tracking\n", iTotal, simuCount); + fflush(stdout); + + // Update weight_out to match ELEGANT + for (long k = 0; k < iTotal; ++k) { + weight_out[k] *= (factor / totalMCRate) * integrated_piwinski_rate; + } + + *n_selected_out = iTotal; + *totalMCRate_out = totalMCRate; + + free(index); + free(thetatemp); + free(weight); + free(xtemp); + free(pxtemp); + free(ytemp); + free(pytemp); + free(zetatemp); + free(deltatemp); +} + +#endif // XTRACK_TOUSCHEK_H \ No newline at end of file diff --git a/xfields/headers/constants.h b/xfields/headers/constants.h index 591e50f5..7f988af9 100644 --- a/xfields/headers/constants.h +++ b/xfields/headers/constants.h @@ -16,6 +16,10 @@ #define MELECTRON_GEV (0.00051099895000) #endif +#if !defined( MELECTRON_EV ) + #define MELECTRON_EV (510998.95) +#endif + #if !defined( MELECTRON_KG ) #define MELECTRON_KG (9.1093837015e-31) #endif diff --git a/xfields/headers/elegant_rng.h b/xfields/headers/elegant_rng.h new file mode 100644 index 00000000..a831c679 --- /dev/null +++ b/xfields/headers/elegant_rng.h @@ -0,0 +1,333 @@ +/* + * elegant_rng.h — Elegant-compatible RNG utilities (C99) + * + * Overview + * -------- + * Re-implements the random-number utilities used by Elegant/SDDS so that + * kernels compiled via xobjects can reproduce (where applicable) the same + * sequences in Xsuite/xfields. Includes the LAPACK DLARAN core (48-bit LCG) + * and the streams random_1..random_6 with Elegant-compatible seeding rules. + * Provides random_1_elegant, randomizeOrder, and related helpers. + * + * Provenance (portions adapted from) + * ---------------------------------- + * - SDDS: mdbmth/drand.c (random_* streams, randomizeOrder, seeding) + * - SDDS: mdbmth/dlaran.c (C translation of LAPACK's DLARAN, via f2c) + * - Elegant: src/drand_oag.c (random_1_elegant and seed behavior) + * - LAPACK: DLARAN (48-bit LCG RNG core) + * + * Purpose / Exposed API + * --------------------- + * - LAPACK-compatible DLARAN core (48-bit, 4×12-bit seed) + * - Streams: random_1 .. random_6 (Elegant-style seeding conventions) + * - random_1_elegant() and seedElegantRandomNumbers() (for Elegant compatibility) + * - permuteSeedBitOrder(), inhibitRandomSeedPermutation() + * - randomizeOrder() (qsort + random keys to match Elegant’s consumption) + * + * Usage Notes + * ----------- + * - Single-threaded, static state; synchronize if used from multiple threads. + * - Calls with negative seeds reinitialize the stream; non-negative seeds consume. + * - MPI seed diversification is intentionally omitted here. + * - Special behavior for seed 987654321 (permute inhibition) is preserved. + * - Goal: reproduce Elegant sequences (bitwise where possible). + * + * References + * ---------- + * - M. Borland, “elegant: A Flexible SDDS-Compliant Code for Accelerator Simulation,” + * APS LS-287 (2000). + */ +#ifndef ELEGANT_RNG_H +#define ELEGANT_RNG_H + +#include +#include +#include +#include + +// ----------------------------------------------------------------------------- +// LAPACK-style DLARAN core +// ------------------------ +// This is the 48-bit multiplicative LCG used by DLARAN, with the 48-bit state +// stored in 4 integers of 12 bits each. The last chunk (iseed[3]) must be odd. +// Returns a uniform double in (0,1). The recurrence and normalization constants +// reproduce LAPACK DLARAN bit-for-bit. +// ----------------------------------------------------------------------------- +static inline double dlaran_core(int32_t iseed[4]) { + // Seed chunks: iseed[0..3], iseed[3] must be odd + int32_t it1, it2, it3, it4; + it4 = iseed[3] * 2549; + it3 = it4 / 4096; + it4 -= (it3 << 12); + it3 = it3 + iseed[2] * 2549 + iseed[3] * 2508; + it2 = it3 / 4096; + it3 -= (it2 << 12); + it2 = it2 + iseed[1] * 2549 + iseed[2] * 2508 + iseed[3] * 322; + it1 = it2 / 4096; + it2 -= (it1 << 12); + it1 = it1 + iseed[0] * 2549 + iseed[1] * 2508 + iseed[2] * 322 + iseed[3] * 494; + it1 %= 4096; + + iseed[0] = it1; + iseed[1] = it2; + iseed[2] = it3; + iseed[3] = it4; + + const double twoneg12 = 2.44140625e-4; // 2^-12 + return ( (double)it1 + + ( (double)it2 + ( (double)it3 + (double)it4 * twoneg12 ) * twoneg12 ) * twoneg12 + ) * twoneg12; // in (0,1) +} + +/* ----------------------------------------------------------------------------- + Seed permutation control (Elegant-compatible) + -------------------------------------------- + Elegant permutes the bit order of integer seeds before packing 12-bit chunks. + This helps decorrelate “close” seeds. A global flag can inhibit this step. + + - inhibitRandomSeedPermutation(state>=0) sets the global flag. + - permuteSeedBitOrder(x) permutes bit positions unless inhibited. + - Elegant disables permutation when random_number_seed == 987654321. +----------------------------------------------------------------------------- */ +static short g_inhibitPermute = 0; +static inline short inhibitRandomSeedPermutation(short state){ + if (state >= 0) g_inhibitPermute = state; + return g_inhibitPermute; +} + +static inline uint32_t permuteSeedBitOrder(uint32_t input0){ + if (g_inhibitPermute) return input0; + uint32_t input = input0; + uint32_t newValue = 0u; + uint32_t offset = input0 % 1000u; + static const uint32_t bitMask[32] = { + 0x00000001u,0x00000002u,0x00000004u,0x00000008u, + 0x00000010u,0x00000020u,0x00000040u,0x00000080u, + 0x00000100u,0x00000200u,0x00000400u,0x00000800u, + 0x00001000u,0x00002000u,0x00004000u,0x00008000u, + 0x00010000u,0x00020000u,0x00040000u,0x00080000u, + 0x00100000u,0x00200000u,0x00400000u,0x00800000u, + 0x01000000u,0x02000000u,0x04000000u,0x08000000u, + 0x10000000u,0x20000000u,0x40000000u,0x80000000u + }; + for (int i=0;i<31;i++) + if (input & bitMask[i]) newValue |= bitMask[(i + offset) % 31]; + if (newValue == input){ + offset++; + newValue = 0u; + for (int i=0;i<31;i++) + if (input & bitMask[i]) newValue |= bitMask[(i + offset) % 31]; + } + return newValue; +} + +/* ----------------------------------------------------------------------------- + Packing helper: split a 32-bit integer into 4×12-bit chunks (DLARAN format). + The last chunk must be odd for DLARAN to work correctly (Elegant behavior). +----------------------------------------------------------------------------- */ +static inline void seed_from_long(int32_t seed[4], long iseed_in, int force_odd_last){ + uint32_t s = (uint32_t)(iseed_in < 0 ? -iseed_in : iseed_in); + s = permuteSeedBitOrder(s); + // pack into 4x12-bit chunks, last must be odd + seed[3] = (int32_t)(s & 4095u); s >>= 12; + if (force_odd_last) seed[3] = (seed[3] | 1); // ensure odd + seed[2] = (int32_t)(s & 4095u); s >>= 12; + seed[1] = (int32_t)(s & 4095u); s >>= 12; + seed[0] = (int32_t)(s & 4095u); +} + +/* ----------------------------------------------------------------------------- + RNG streams random_1 .. random_6 + -------------------------------- + These reproduce the SDDS/Elegant API and seeding semantics: + + - Calling with a *negative* iseed re-initializes that stream from |iseed|. + - Calling with a non-negative iseed consumes the next variate. + - random_1 (on (re)seed) also re-seeds streams 2..6 using |base|+{2,4,6,8,10}. + - All streams use the same DLARAN core, independent 48-bit states. + + Important: + * This file intentionally omits MPI diversification (modes 1..4 in Elegant). + * Keep static state: not thread-safe by design (matches Elegant). +----------------------------------------------------------------------------- */ +double random_2(long iseed); +double random_3(long iseed); +double random_4(long iseed); +double random_5(long iseed); +double random_6(long iseed); + +double random_1(long iseed){ + static int initialized = 0; + static int32_t seed[4] = {0,0,0,0}; + if (!initialized || iseed < 0){ + long base = (iseed < 0) ? -iseed : iseed; // abs + base = (long)permuteSeedBitOrder((uint32_t)base); // permute + // SDDS-like reseed + random_2(-(base + 2)); + random_3(-(base + 4)); + random_4(-(base + 6)); + random_5(-(base + 8)); + random_6(-(base + 10)); + // force odd and split in 4×12 bit + base = (base/2)*2 + 1; + uint32_t s = (uint32_t)base; + seed[3] = (int32_t)(s & 4095u); s >>= 12; + seed[2] = (int32_t)(s & 4095u); s >>= 12; + seed[1] = (int32_t)(s & 4095u); s >>= 12; + seed[0] = (int32_t)(s & 4095u); + initialized = 1; + } + return dlaran_core(seed); +} + +double random_2(long iseed){ + static int initialized = 0; + static int32_t seed[4] = {0,0,0,0}; + if (!initialized || iseed < 0){ + if (iseed >= 0) iseed = -1; + seed_from_long(seed, -iseed, /*force_odd_last=*/1); + initialized = 1; + } + return dlaran_core(seed); +} +double random_3(long iseed){ + static int initialized = 0; + static int32_t seed[4] = {0,0,0,0}; + if (!initialized || iseed < 0){ + if (iseed >= 0) iseed = -1; + seed_from_long(seed, -iseed, /*force_odd_last=*/1); + initialized = 1; + } + return dlaran_core(seed); +} +double random_4(long iseed){ + static int initialized = 0; + static int32_t seed[4] = {0,0,0,0}; + if (!initialized || iseed < 0){ + if (iseed >= 0) iseed = -1; + seed_from_long(seed, -iseed, /*force_odd_last=*/1); + initialized = 1; + } + return dlaran_core(seed); +} +double random_5(long iseed){ + static int initialized = 0; + static int32_t seed[4] = {0,0,0,0}; + if (!initialized || iseed < 0){ + if (iseed >= 0) iseed = -1; + seed_from_long(seed, -iseed, /*force_odd_last=*/1); + initialized = 1; + } + return dlaran_core(seed); +} +double random_6(long iseed){ + static int initialized = 0; + static int32_t seed[4] = {0,0,0,0}; + if (!initialized || iseed < 0){ + if (iseed >= 0) iseed = -1; + seed_from_long(seed, -iseed, /*force_odd_last=*/1); + initialized = 1; + } + return dlaran_core(seed); +} + +static inline double random_1_elegant(long iseed) { + static int initialized = 0; + static int32_t seed[4] = {0,0,0,0}; + if (!initialized || iseed < 0){ + long base = (iseed < 0) ? -iseed : iseed; // abs + // Respect current inhibit flag (seedElegantRandomNumbers sets it) + base = (long)permuteSeedBitOrder((uint32_t)base); // no-op if inhibited + base = (base/2)*2 + 1; // force odd + uint32_t s = (uint32_t)base; + seed[3] = (int32_t)(s & 4095u); s >>= 12; + seed[2] = (int32_t)(s & 4095u); s >>= 12; + seed[1] = (int32_t)(s & 4095u); s >>= 12; + seed[0] = (int32_t)(s & 4095u); + initialized = 1; + } + return dlaran_core(seed); +} + +/* ----------------------------------------------------------------------------- + randomizeOrder + ------------------------------ + Elegant shuffles arrays by: + 1) creating an array of (buffer copy, random key) pairs, + 2) sorting by the random key (qsort), + 3) copying buffers back in sorted order. + + This consumes RNG like Elegant does and matches its order exactly. + It allocates O(N*size) memory. +----------------------------------------------------------------------------- */ +typedef struct RANDOMIZATION_HOLDER_ { + void* buffer; + double randomValue; +} RANDOMIZATION_HOLDER; + +static int randomizeOrderCmp(const void *p1, const void *p2) { + const RANDOMIZATION_HOLDER *rh1 = (const RANDOMIZATION_HOLDER *)p1; + const RANDOMIZATION_HOLDER *rh2 = (const RANDOMIZATION_HOLDER *)p2; + if (rh1->randomValue > rh2->randomValue) return 1; + if (rh1->randomValue < rh2->randomValue) return -1; + return 0; +} + +static long randomizeOrder(char *ptr, long size, long length, + long iseed, double (*urandom)(long iseed1)) { + if (!ptr || size<=0 || !urandom) return 0; + if (length < 2) return 1; + if (iseed < 0) urandom(iseed); + + RANDOMIZATION_HOLDER *rh = + (RANDOMIZATION_HOLDER*)malloc(sizeof(*rh) * (size_t)length); + if (!rh) return 0; + + for (long i=0; i inhibit permutation globally. + + Call this exactly once (per process) before any RNG use to mimic Elegant’s + &run_setup random_number_seed behavior. +----------------------------------------------------------------------------- */ +static inline void seedElegantRandomNumbers(long seed, short inhibit_permute){ + long s0 = labs(seed), s1 = labs(seed + 2), s2 = labs(seed + 4), s3 = labs(seed + 6); + + if (s0 == 987654321) inhibitRandomSeedPermutation(1); + else inhibitRandomSeedPermutation(inhibit_permute ? 1 : 0); + + random_1_elegant(-s0); + random_2(-s1); + random_3(-s2); + random_4(-s3); +} + +#endif // ELEGANT_RNG_H \ No newline at end of file diff --git a/xfields/touschek/manager.py b/xfields/touschek/manager.py new file mode 100644 index 00000000..ff3d8d34 --- /dev/null +++ b/xfields/touschek/manager.py @@ -0,0 +1,619 @@ +# copyright ################################# # +# This file is part of the Xfields Package. # +# Copyright (c) CERN, 2021. # +# ########################################### # + +import xtrack as xt +import xfields as xf + +import numpy as np +import pandas as pd +from scipy.integrate import quad +from scipy.special import i0 +from scipy.constants import physical_constants + +ELECTRON_MASS_EV = xt.ELECTRON_MASS_EV +C_LIGHT_VACUUM = physical_constants['speed of light in vacuum'][0] +CLASSICAL_ELECTRON_RADIUS = physical_constants['classical electron radius'][0] + +class TouschekCalculator: + ''' + Internal helper that evaluates Piwinski scattering rates and integrates + them along the lattice. + + This class is instantiated and owned by :class:`TouschekManager`; it + should not be constructed directly by users. + + Parameters + ---------- + manager : TouschekManager + The owning manager, from which beam parameters, optics, and the + local momentum acceptance table are read. + + Attributes + ---------- + twiss : xtrack.TwissTable or None + Twiss table for the ring. Injected by + :meth:`TouschekManager.initialise_touschek` before any rate + evaluation is performed. + + References + ---------- + .. [1] A. Piwinski, "The Touschek Effect in Strong Focusing Storage + Rings", arXiv:physics/9903034 (1999). + .. [2] A. Xiao and M. Borland, "Monte Carlo simulation of Touschek + effect", Phys. Rev. ST Accel. Beams **13**, 074201 (2010). + https://doi.org/10.1103/PhysRevSTAB.13.074201 + .. [3] M. Borland, "elegant: A Flexible SDDS-Compliant Code for + Accelerator Simulation", APS LS-287 (2000). + ''' + def __init__(self, manager): + self.manager = manager + self.twiss = None + + def _compute_piwinski_integral(self, tm, B1, B2): + """ + Compute Piwinski integral for Touschek scattering rate calculation. + """ + from math import atan, tan, sqrt, exp, log, pi + + km = atan(sqrt(tm)) + + def int_piwinski(k): + t = np.tan(k) ** 2 + fact = ( + (2*t + 1)**2 * (t/tm / (1+t) - 1) / t + t - sqrt(t*tm * (1 + t)) + - (2 + 1 / (2*t)) * log(t/tm / (1+t)) + ) + if B2 * t < 500: + intp = fact * exp(-B1*t) * i0(B2*t) * sqrt(1+t) + else: + intp = ( + fact + * exp(B2*t - B1*t) + / sqrt(2*pi * B2*t) + * sqrt(1+t) + ) + return intp + + val, _ = quad( + int_piwinski, + km, + pi / 2, + epsabs=1e-16, + epsrel=1e-12 + ) + + return val + + def _compute_piwinski_scattering_rate(self, element): + """ + Compute Piwinski Touschek scattering rate. + """ + p0c = self.manager.particle_ref.p0c[0] + bunch_population = self.manager.bunch_population + local_momentum_acceptance = self.manager.local_momentum_acceptance + gemitt_x = self.manager.gemitt_x + gemitt_y = self.manager.gemitt_y + twiss = self.twiss + alfx = twiss['alfx', element] + betx = twiss['betx', element] + alfy = twiss['alfy', element] + bety = twiss['bety', element] + sigma_z = self.manager.sigma_z + sigma_delta = self.manager.sigma_delta + delta = twiss['delta', element] + dx = twiss['dx', element] + dpx = twiss['dpx', element] + dxt = alfx * dx + betx * dpx # dxt: dx tilde + dy = twiss['dy', element] + dpy = twiss['dpy', element] + dyt = alfy * dy + bety * dpy # dyt: dy tilde + + try: + s = self.twiss.rows[element].s[0] + except: + s = self.manager.line.get_s_position(element) + + deltaN = np.interp(s, local_momentum_acceptance.s, local_momentum_acceptance.deltan) + deltaP = np.interp(s, local_momentum_acceptance.s, local_momentum_acceptance.deltap) + + sigmab_x = np.sqrt(gemitt_x * betx) # Horizontal betatron beam size + sigma_x = np.sqrt(gemitt_x * betx + dx**2 * sigma_delta**2) # Horizontal beam size + + sigmab_y = np.sqrt(gemitt_y * bety) # Vertical betatron beam size + sigma_y = np.sqrt(gemitt_y * bety + dy**2 * sigma_delta**2) # Vertical beam size + + sigma_h = (sigma_delta**-2 + (dx**2 + dxt**2)/sigmab_x**2 + (dy**2 + dyt**2)/sigmab_y**2)**(-0.5) + + p = p0c * (1 + delta) + gamma = np.sqrt(1 + p**2 / ELECTRON_MASS_EV**2) + beta = np.sqrt(1 - gamma**-2) + + B1 = betx**2 / (2 * beta**2 * gamma**2 * sigmab_x**2) * (1 - sigma_h**2 * dxt**2 / sigmab_x**2) \ + + bety**2 / (2 * beta**2 * gamma**2 * sigmab_y**2) * (1 - sigma_h**2 * dyt**2 / sigmab_y**2) + + B2 = np.sqrt(B1**2 - betx**2 * bety**2 * sigma_h**2 / (beta**4 * gamma**4 * sigmab_x**4 * sigmab_y**4 * sigma_delta**2) \ + * (sigma_x**2 * sigma_y**2 - sigma_delta**4 * dx**2 * dy**2)) + + tmN = beta**2 * (deltaN**2) + tmP = beta**2 * (deltaP**2) + + piwinski_integralN = self._compute_piwinski_integral(tmN, B1, B2) + piwinski_integralP = self._compute_piwinski_integral(tmP, B1, B2) + + rateN = CLASSICAL_ELECTRON_RADIUS**2 * C_LIGHT_VACUUM * bunch_population**2 \ + / (8*np.pi * gamma**2 * sigma_z * np.sqrt(sigma_x**2 * sigma_y**2 - sigma_delta**4 * dx**2 * dy**2)) \ + * 2 * np.sqrt(np.pi * (B1**2 - B2**2)) * piwinski_integralN + + rateP = CLASSICAL_ELECTRON_RADIUS**2 * C_LIGHT_VACUUM * bunch_population**2 \ + / (8*np.pi * gamma**2 * sigma_z * np.sqrt(sigma_x**2 * sigma_y**2 - sigma_delta**4 * dx**2 * dy**2)) \ + * 2 * np.sqrt(np.pi * (B1**2 - B2**2)) * piwinski_integralP + + rate = (rateN + rateP) / 2 + + return rate + + def _compute_integrated_piwinski_rates(self, element): + """ + Integrate the Piwinski Touschek scattering rate along the line using + the trapezoidal rule, between successive TouschekScattering elements. + + For each TouschekScattering element, the method stores the integrated + rate per bunch over the preceding section of the line. This per-bunch + rate is later used to assign the correct weights to Touschek-scattered + particles at the corresponding element. + """ + def _get_s(name): + try: + return tab.rows[name].s[0] + except (KeyError, AttributeError, IndexError, TypeError): + return self.manager.line.get_s_position(name) + + def _step(name, s_before, rate_before, integrated): + s = _get_s(name) + ds = s - s_before + if ds > 0.0: + rate = self._compute_piwinski_scattering_rate(name) + integrated += 0.5 * (rate_before + rate) * ds + return s, rate, integrated + else: + return s_before, rate_before, integrated + + line = self.manager.line + tab = line.get_table() + t_rev0 = float(self.twiss.t_rev0) + + # Indexes of the TouschekScatterings + ii_t = [ii for ii, nn in enumerate(tab.name[:-1]) if isinstance(line[nn], xf.TouschekScattering)] + + integrated = 0.0 + + if element is None: + ii_current = 0 + s0 = 0.0 + r0 = self._compute_piwinski_scattering_rate(tab.name[0]) + else: + import re + ii_current = int(re.search(r'\d+', element).group()) + tscatter_before = tab.name[ii_t[ii_current - 1]] if ii_current != 0 else tab.name[0] + s0 = _get_s(tscatter_before) + r0 = self._compute_piwinski_scattering_rate(tscatter_before) + + s_before = s0 + rate_before = r0 + + if element is None: + # Configure all the TouschekScattering elements + for ii, nn in enumerate(tab.name): + s_before, rate_before, integrated = _step(nn, s_before, rate_before, integrated) + + if ii == ii_t[ii_current]: + # divide by c and by t_rev0 --> per-bunch rate + integrated_piwinski_rate = integrated / C_LIGHT_VACUUM / t_rev0 + elem = line[nn] # xf.TouschekScattering + # print(f'Integrated Piwinski rate at {nn}: {integrated_piwinski_rate*1e-3} [kHz]') + elem._configure(integrated_piwinski_rate=integrated_piwinski_rate) + integrated = 0.0 + ii_current += 1 + if ii_current == len(ii_t): + break + else: + # Configure only the TouschekScattering element named `element` + subtab = tab.rows[tscatter_before:element] + for nn in subtab.name: + s_before, rate_before, integrated = _step(nn, s_before, rate_before, integrated) + + if nn == element: + # divide by c and by t_rev0 --> per-bunch rate + integrated_piwinski_rate = integrated / C_LIGHT_VACUUM / t_rev0 + elem = line[nn] # xf.TouschekScattering + # print(f'Integrated Piwinski rate at {nn}: {integrated_piwinski_rate*1e-3} [kHz]') + elem._configure(integrated_piwinski_rate=integrated_piwinski_rate) + break + + +class TouschekManager: + ''' + High-level manager that orchestrates a full Touschek scattering simulation. + + The manager: + + 1. Computes the Piwinski scattering rate at every + :class:`TouschekScattering` element in the line using the local beam + optics, emittances, and the local momentum acceptance (LMA). + 2. Integrates those rates over each lattice section between consecutive + scattering elements (trapezoidal rule) to obtain the per-bunch, + per-turn loss probability in each section. + 3. Configures each :class:`TouschekScattering` element with all local + parameters so that it can generate and weight Monte Carlo + macro-particles. + + Parameters + ---------- + line : xtrack.Line + The accelerator lattice. Must contain at least one + :class:`TouschekScattering` element and a ``particle_ref``. + twiss : xtrack.TwissTable or None, optional + Pre-computed Twiss table. If ``None``, it is computed internally + by :meth:`initialise_touschek` using the ``method`` keyword + argument (default ``"6d"``). + local_momentum_acceptance : xtrack.Table + Table returned by ``line.get_local_momentum_acceptance()``. Must + contain the columns ``name``, ``s``, ``deltan`` (negative LMA, + :math:`\\delta_N < 0`), and ``deltap`` (positive LMA, + :math:`\\delta_P > 0`). Values are scaled in-place by + ``local_momentum_acceptance_scale`` upon construction. + nemitt_x : float, optional + Horizontal normalised emittance [m·rad]. Mutually exclusive with + ``gemitt_x``. + nemitt_y : float, optional + Vertical normalised emittance [m·rad]. Mutually exclusive with + ``gemitt_y``. + gemitt_x : float, optional + Horizontal geometric emittance [m·rad]. Mutually exclusive with + ``nemitt_x``. + gemitt_y : float, optional + Vertical geometric emittance [m·rad]. Mutually exclusive with + ``nemitt_y``. + sigma_z : float + RMS bunch length [m]. + sigma_delta : float + RMS relative momentum spread :math:`\\sigma_\\delta`. + bunch_population : float + Number of real particles per bunch, :math:`N_b`. + n_simulated : int + Number of scattered macro-particle candidates to generate per + scattering element. Larger values improve statistics at the cost + of CPU time. Values of :math:`10^6`–:math:`10^7` are typical. + nx : float, optional + Truncation of the transverse-horizontal Gaussian sampling window in + units of :math:`\\sqrt{\\varepsilon_x}`. Default 3. + ny : float, optional + Truncation of the transverse-vertical Gaussian sampling window in + units of :math:`\\sqrt{\\varepsilon_y}`. Default 3. + nz : float, optional + Truncation of the longitudinal Gaussian sampling window in units of + :math:`\\sigma_\\zeta` and :math:`\\sigma_\\delta`. May be reduced + element-by-element (see Notes) to prevent drawing initial particles + outside the LMA. + Default 3. + local_momentum_acceptance_scale : float, optional + Multiplicative safety factor applied to the LMA on construction. + A value of 0.85 (default) ensures that particles with momentum + deviation smaller than the LMA but with nonzero betatron amplitude + are accepted. This avoids that particles that are eventually lost + are not considered for tracking + ignored_portion : float, optional + Fraction of the total simulated scattering weight that is discarded + before tracking. Only the highest-weight particles whose cumulative + weight reaches ``(1 - ignored_portion)`` of the total are retained + and tracked; the remaining low-weight, low-probability events are + dropped. The default value of ``0.01`` retains 99 % of the total + weight while significantly reducing the number of particles that + must be tracked, providing a good accuracy–efficiency trade-off. + Setting this to ``0`` keeps all simulated particles. + seed : int, optional + RNG seed for the ELEGANT-compatible 48-bit LCG generator used + inside :class:`TouschekScattering`. Default 1997. + method : str, optional + Twiss method forwarded to ``line.twiss()`` when ``twiss`` is + ``None``. Accepted values: ``"4d"``, ``"6d"``. Default ``"6d"``. + + Attributes + ---------- + line : xtrack.Line + The accelerator lattice passed at construction. + particle_ref : xtrack.Particles + Reference particle extracted from ``line.particle_ref``. + twiss : xtrack.TwissTable or None + Twiss table (populated by :meth:`initialise_touschek` if not + provided at construction). + gemitt_x, gemitt_y : float + Geometric emittances [m·rad] derived from the input normalised or + geometric emittances. + local_momentum_acceptance : xtrack.Table + The (scaled) LMA table. + sigma_z, sigma_delta : float + Bunch length [m] and momentum spread. + bunch_population : float + Bunch population :math:`N_b`. + n_simulated : int + Number of simulated scattered candidates per element. + nx, ny, nz : float + Phase-space sampling truncation parameters. + seed : int + RNG seed. + touschek : TouschekCalculator + The calculator instance used internally to evaluate Piwinski rates. + + References + ---------- + .. [1] A. Xiao and M. Borland, "Monte Carlo simulation of Touschek + effect", Phys. Rev. ST Accel. Beams **13**, 074201 (2010). + https://doi.org/10.1103/PhysRevSTAB.13.074201 + .. [2] M. Borland, "elegant: A Flexible SDDS-Compliant Code for + Accelerator Simulation", APS LS-287 (2000). + ''' + def __init__(self, line=None, twiss=None, local_momentum_acceptance=None, + nemitt_x=None, nemitt_y=None, + sigma_z=None, sigma_delta=None, bunch_population=None, + n_simulated=None, gemitt_x=None, gemitt_y=None, + local_momentum_acceptance_scale=0.85, ignored_portion=0.01, + seed=1997, nx=3, ny=3, nz=3, **kwargs): + + # Input validation + if line is None: + raise ValueError("`line` is required.") + if not hasattr(line, "particle_ref"): + raise ValueError("`line` must have a `particle_ref`.") + if local_momentum_acceptance is None: + raise ValueError("`local_momentum_acceptance` is required.") + if sigma_z is None: + raise ValueError("`sigma_z` is required.") + if sigma_delta is None: + raise ValueError("`sigma_delta` is required.") + if bunch_population is None: + raise ValueError("`bunch_population` is required.") + if n_simulated is None: + raise ValueError("`n_simulated` is required.") + + # Local momentum acceptnace validation + required_cols = {"name", "s", "deltan", "deltap"} + if not isinstance(local_momentum_acceptance, xt.Table): + raise TypeError("`local_momentum_acceptance` must be an `xt.Table` object.") + missing = required_cols - set(local_momentum_acceptance._col_names) + if missing: + raise ValueError(f"`local_momentum_acceptance` missing columns: {sorted(missing)}") + + for col in ("s", "deltan", "deltap"): + try: + vals = np.asarray(local_momentum_acceptance[col], dtype=float) + except Exception: + raise TypeError(f"`{col}` column must be numeric (cannot coerce to float).") + nan_mask = np.isnan(vals) + if nan_mask.any(): + bad = list(np.where(nan_mask)[0][:5]) + raise ValueError(f"`{col}` contains NaN at indices {bad}.") + inf_mask = np.isinf(vals) + if inf_mask.any(): + bad = list(np.where(inf_mask)[0][:5]) + raise ValueError(f"`{col}` contains inf at indices {bad}.") + + self.line = line + self.particle_ref = line.particle_ref + self.twiss = twiss + + # Check that the line contains TouschekScatterings + tab = line.get_table() + try: + has = "TouschekScattering" in set(np.unique(tab.element_type)) + except Exception: + has = "TouschekScattering" in set(getattr(tab, "element_type", [])) + if not has: + raise ValueError("The line does not contain any TouschekScattering. " + "Please add them before initializing the TouschekManager.") + + # Local momentum acceptance + local_momentum_acceptance.deltan *= local_momentum_acceptance_scale + local_momentum_acceptance.deltap *= local_momentum_acceptance_scale + self.local_momentum_acceptance = local_momentum_acceptance + + self.sigma_z = sigma_z + self.sigma_delta = sigma_delta + self.bunch_population = bunch_population + self.n_simulated = n_simulated + self.ignored_portion = ignored_portion + self.seed = seed + self.nx = nx + self.ny = ny + self.nz = nz + + # Limits from ELEGANT + self._theta_min = 0.00005*np.pi + self._theta_max = 0.99995*np.pi + + # Emittance validation + nemitt_given = nemitt_x is not None and nemitt_y is not None + gemitt_given = gemitt_x is not None and gemitt_y is not None + + if nemitt_given and gemitt_given: + raise ValueError("Provide either normalized emittances (nemitt_x, nemitt_y) " + "OR geometric emittances (gemitt_x, gemitt_y), not both.") + if not (nemitt_given or gemitt_given): + raise ValueError("You must provide either both normalized emittances (nemitt_x, nemitt_y) " + "OR both geometric emittances (gemitt_x, gemitt_y).") + + if nemitt_given: + beta0 = line.particle_ref.beta0[0] + gamma0 = line.particle_ref.gamma0[0] + self.gemitt_x = nemitt_x / (beta0 * gamma0) + self.gemitt_y = nemitt_y / (beta0 * gamma0) + else: + self.gemitt_x = gemitt_x + self.gemitt_y = gemitt_y + + self.kwargs = kwargs + + self.touschek = TouschekCalculator(self) + + def initialise_touschek(self, element=None): + ''' + Compute and configure all Piwinski rates in the lattice. + + For each :class:`TouschekScattering` element this method: + + 1. Evaluates the local Piwinski scattering rate using the Twiss + parameters, beam emittances, and the local momentum acceptance. + 2. Integrates the rate over the preceding lattice section using the + trapezoidal rule to obtain the per-bunch, per-turn loss probability. + 3. Stores the integrated rate and all local optics parameters on the + element via :meth:`TouschekScattering._configure` so that + :meth:`TouschekScattering.scatter` can weight the Monte Carlo + macro-particles correctly. + + Parameters + ---------- + element : str or None, optional + If ``None`` (default), all :class:`TouschekScattering` elements in + the line are initialised in one pass. If a string, only the named + element is (re-)initialised; the Piwinski rate is integrated only + over the lattice section between that element and the preceding + scattering centre. + ''' + line = self.line + tab = line.get_table() + + local_momentum_acceptance = self.local_momentum_acceptance + + if self.twiss is None: + twiss_method = self.kwargs.get("method", "6d") + self.twiss = self.line.twiss(method=twiss_method) + + # Pass the twiss to the TouschekCalculator + self.touschek.twiss = self.twiss + + import time + t0 = time.time() + self.touschek._compute_integrated_piwinski_rates(element) + print(f"Computed integrated piwinski rates in {time.time() - t0:.2f} s.") + + # Helper to config all fields to a single TouschekScattering + def _config(nn): + try: + s = tab.rows[nn].s[0] + except Exception: + s = self.line.get_s_position(nn) + + twiss = self.twiss + alfx = twiss["alfx", nn]; betx = twiss["betx", nn] + alfy = twiss["alfy", nn]; bety = twiss["bety", nn] + dx = twiss["dx", nn]; dpx = twiss["dpx", nn] + dy = twiss["dy", nn]; dpy = twiss["dpy", nn] + dN = np.interp(s, local_momentum_acceptance.s, local_momentum_acceptance.deltan) + dP = np.interp(s, local_momentum_acceptance.s, local_momentum_acceptance.deltap) + + x_co = twiss["x", nn]; px_co = twiss["px", nn] + y_co = twiss["y", nn]; py_co = twiss["py", nn] + zeta_co = twiss["zeta", nn]; delta_co = twiss["delta", nn] + + # Adjust the effective longitudinal sampling cutoff (nz_eff) to prevent + # generation of initial particles that are already outside + # the local momentum aperture (LMA) before Touschek scattering. + # + # Background: + # In the Touschek Monte Carlo routine, initial particle coordinates + # are drawn from a truncated Gaussian distribution with cutoffs + # {nx, ny, nz} in the transverse and longitudinal planes. The longitudinal + # cutoff nz sets the maximum |δ| ≈ nz*σδ. If nz*σδ exceeds the local + # momentum acceptance at this location, some particles + # are sampled outside the LMA (|δ| > LMA) even before any scattering event occurs. + # + # These particles create a pathological situation: + # • For very small scattering angles (θ* --> 0 in the CM frame), + # such particles are flagged as "candidates for loss" and passed to tracking, + # even though their state is essentially unchanged by scattering. + # • Because the Møller differential cross-section diverges at + # θ* --> 0, the corresponding particle weights become extremely large. + # • The pickPart routine then tends to select a handful of these + # high-weight pathological particles, distorting both the local + # scattering rate (RMC/RP diverges) and the overall Touschek + # lifetime estimate. + # + # Mitigation: + # To eliminate these spurious contributions, we dynamically reduce + # the longitudinal cutoff at each TouschekScattering element: + # + # nz_eff = min(nz, 0.85 * min(|δN|, δP) / σδ) + # + # where δN, δP are the negative/positive momentum aperture limits + # (scaled by local_momentum_acceptance_scale). This ensures that the sampled + # longitudinal range ±nz_eff*σδ always lies strictly inside the local + # momentum aperture, with a small safety factor (0.85). As a result, + # only pathological large-weight events are avoided, and the Monte Carlo + # rate remains consistent with the Piwinski formula. + # + # NOTE: nz_eff is determined independently at each scattering element, + # so tighter cutoffs are applied only where the local momentum aperture + # is restrictive, while wider cutoffs are retained elsewhere. + min_dNdP = min(abs(dN), dP) + nz_eff = min(self.nz, 0.85 * min_dNdP / self.sigma_delta) + + if nz_eff < self.nz: + print(f""" + *********************************************************************************************** + [TouschekManager] Warning: longitudinal cutoff reduced at element '{nn}' (s={s:.2f} m). + + Using nz_eff={nz_eff:.2f} instead of nz={self.nz:.2f}. + This ensures that particles are sampled strictly within the local momentum aperture. + *********************************************************************************************** + """) + + piwinski_rate = self.touschek._compute_piwinski_scattering_rate(nn) + + elem = line[nn] # xf.TouschekScattering + element_index = line.element_names.index(nn) + + elem._configure( + s=s, + particle_ref=self.particle_ref, + element_index=element_index, + bunch_population=self.bunch_population, + gemitt_x=self.gemitt_x, + gemitt_y=self.gemitt_y, + alfx=alfx, betx=betx, + alfy=alfy, bety=bety, + dx=dx, dpx=dpx, + dy=dy, dpy=dpy, + x_co=x_co, px_co=px_co, + y_co=y_co, py_co=py_co, + zeta_co=zeta_co, delta_co=delta_co, + deltaN=dN, deltaP=dP, + sigma_z=self.sigma_z, + sigma_delta=self.sigma_delta, + n_simulated=self.n_simulated, + nx=self.nx, ny=self.ny, nz=nz_eff, + theta_min=self._theta_min, theta_max=self._theta_max, + ignored_portion=self.ignored_portion, + piwinski_rate=piwinski_rate, + seed=self.seed, inhibit_permute=0 + ) + + if element is None: + for nn in tab.name[:-1]: # Avoid the last tab.name which is _end_point + if isinstance(line[nn], xf.TouschekScattering): + print(f'Initialising TouschekScattering for {nn}') + _config(nn) + else: + if not isinstance(element, str): + raise TypeError(f"`element` must be a string (got {type(element).__name__}).") + if element not in set(tab.name): + raise ValueError( + f"`element='{element}'` is not present in the line provided to the TouschekManager." + ) + if not isinstance(line[element], xf.TouschekScattering): + raise TypeError( + f"`line['{element}']` is not a TouschekScattering (got {type(line[element]).__name__})." + ) + print(f'Initialising TouschekScattering for {element}') + _config(element) \ No newline at end of file