Longitudinal parabolic distribution

examples/longitudinal/006_generate_parabolic.py

@@ -0,0 +1,76 @@
+# copyright ############################### #
+# This file is part of the Xpart Package.   #
+# Copyright (c) CERN, 2021.                 #
+# ######################################### #
+
+"""
+Simple example on how to generate a parabolic particle distribution
+"""
+
+import json
+import xpart as xp
+import xtrack as xt
+import xobjects as xo
+
+from xpart import parabolic_longitudinal_distribution
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Load the reference particle
+filename = xt._pkg_root.parent.joinpath('test_data/lhc_no_bb/line_and_particle.json')
+with open(filename, 'r') as fid:
+    input_data = json.load(fid)
+line = xt.Line.from_dict(input_data['line'])
+line.build_tracker()
+
+# Specify the beam parameters
+num_part = 1000000
+sigma_z = 0.05
+
+# Build a reference particle
+p0 = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3,
+                  delta=[10])
+
+# Built a set of three particles with different x coordinates
+particles = parabolic_longitudinal_distribution(
+						num_particles=num_part,
+						nemitt_x=3e-6, 
+						nemitt_y=3e-6, 
+						sigma_z=sigma_z,
+						particle_ref=p0, 
+						total_intensity_particles=1e10,
+						line=line
+						)
+
+
+# Make a histogram
+bin_heights, bin_borders = np.histogram(particles.zeta, bins=300)
+bin_widths = np.diff(bin_borders)
+bin_centers = bin_borders[:-1] + bin_widths / 2
+
+
+# Generate the plots
+plt.close('all')
+
+fig1 = plt.figure(1, figsize=(6.4, 7))
+ax21 = fig1.add_subplot(3,1,1)
+ax22 = fig1.add_subplot(3,1,2)
+ax23 = fig1.add_subplot(3,1,3)
+ax21.plot(particles.x*1000, particles.px, '.', markersize=1)
+ax21.set_xlabel(r'x [mm]')
+ax21.set_ylabel(r'px [-]')
+ax22.plot(particles.y*1000, particles.py, '.', markersize=1)
+ax22.set_xlabel(r'y [mm]')
+ax22.set_ylabel(r'py [-]')
+ax23.plot(particles.zeta, particles.delta*1000, '.', markersize=1)
+ax23.set_xlabel(r'z [-]')
+ax23.set_ylabel(r'$\delta$ [1e-3]')
+fig1.subplots_adjust(bottom=.08, top=.93, hspace=.33, left=.18,
+                     right=.96, wspace=.33)
+
+fig2, ax2 = plt.subplots(1, 1, figsize = (6,5))
+ax2.bar(bin_centers, bin_heights, width=bin_widths)
+ax2.set_ylabel('Counts')
+ax2.set_xlabel(r'z [-]')
+
+plt.show()

tests/test_build_particles_parabolic.py

@@ -0,0 +1,66 @@
+# copyright ############################### #
+# This file is part of the Xpart Package.   #
+# Copyright (c) CERN, 2021.                 #
+# ######################################### #
+
+import json
+
+import numpy as np
+
+import xpart as xp
+import xtrack as xt
+import xobjects as xo
+
+from xpart.longitudinal import parabolic_longitudinal_distribution
+
+from xobjects.test_helpers import for_all_test_contexts
+
+
+@for_all_test_contexts
+def test_build_particles_parabolic(test_context):
+    for ctx_ref in [test_context, None]:
+        # Build a reference particle
+        p0 = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3,
+                          delta=[10], _context=ctx_ref)
+
+	# Parameters for the test 
+        num_part = 1000000
+        parabolic_parameter = 0.05
+
+        # Load machine model (from pymask)
+        filename = xt._pkg_root.parent.joinpath('test_data/lhc_no_bb/line_and_particle.json')
+        with open(filename, 'r') as fid:
+            input_data = json.load(fid)
+        line = xt.Line.from_dict(input_data['line'])
+        line.build_tracker(_context=test_context)
+
+	# Built a set of three particles with different x coordinates
+        particles, matcher = parabolic_longitudinal_distribution(
+                                                        num_particles=num_part,
+                                                        nemitt_x=3e-6, 
+                                                        nemitt_y=3e-6, 
+                                                        sigma_z=parabolic_parameter,
+                                                        particle_ref=p0, 
+                                                        total_intensity_particles=1e10,
+                                                        line=line,
+                                                        return_matcher=True
+                                                                )
+
+        dct = particles.to_dict() 
+        assert np.all(dct['p0c'] == 7e12)
+        tw = line.twiss(particle_ref=p0)
+
+	# Test if longitudinal coordinates match with Single
+	# Generate distribution from RF matcher
+        tau = particles.zeta / p0.beta0[0]
+        tau_distr_y = matcher.tau_distr_y
+        tau_distr_x = matcher.tau_distr_x
+        dx = tau_distr_x[1] - tau_distr_x[0]
+        hist, edges = np.histogram(tau,
+                    range=(tau_distr_x[0]-dx/2., tau_distr_x[-1]+dx/2.),
+                    bins=len(tau_distr_x))
+        hist = hist / sum(hist) * sum(tau_distr_y)
+
+        assert np.all(np.isclose(hist, tau_distr_y, atol=5.e-2, rtol=1.e-2))
+
+

xpart/__init__.py

@@ -20,6 +20,7 @@
 from .transverse_generators import generate_2D_gaussian
 
 from .longitudinal import generate_longitudinal_coordinates
+from .longitudinal import parabolic_longitudinal_distribution
 from .longitudinal.generate_longitudinal import _characterize_line
 
 from .monitors import PhaseMonitor

xpart/longitudinal/__init__.py

@@ -6,3 +6,4 @@
 from .generate_longitudinal import (generate_longitudinal_coordinates,
                                     _characterize_line)
 from .single_rf_harmonic_matcher import SingleRFHarmonicMatcher
+from .generate_parabolic_longitudinal_distribution import parabolic_longitudinal_distribution

xpart/longitudinal/generate_parabolic_longitudinal_distribution.py

@@ -0,0 +1,78 @@
+from xpart.longitudinal import generate_longitudinal_coordinates
+from xpart import build_particles
+import numpy as np
+
+def parabolic_longitudinal_distribution(_context=None, 
+					num_particles=None,
+                			nemitt_x=None, 
+                			nemitt_y=None, 
+                			sigma_z=None,
+                			particle_ref=None, 
+                			total_intensity_particles=None,
+                			tracker=None,
+                			line=None,
+                			return_matcher=False
+                			):
+
+	"""
+	Function to generate a parabolic longitudinal distribution 
+	"""
+
+	if line is not None and tracker is not None:
+		raise ValueError(
+		    'line and tracker cannot be provided at the same time.')
+
+	if tracker is not None:
+		_print('Warning! '
+		    "The argument tracker is deprecated. Please use line instead.")
+		line = tracker.line
+
+	if line is not None:
+		assert line.tracker is not None, ("The line must have a tracker, "
+		    "please call Line.build_tracker() first.")
+
+	if num_particles is None:
+		raise ValueError(
+				'Number of particles must be provided')
+	if sigma_z is None:
+		raise ValueError(
+				'Bunch length sigma_z must be provided')
+
+	if particle_ref is None:
+		raise ValueError(
+				'Reference particle must be provided')
+
+	# If emittances are not provided, set them to default value of one
+	if nemitt_x is None:
+		nemitt_x = 1.0
+
+	if nemitt_y is None:
+		nemitt_y = 1.0
+
+	# Generate longitudinal coordinates s
+	zeta, delta, matcher = generate_longitudinal_coordinates(line=line, distribution='parabolic', 
+							num_particles=num_particles, 
+							engine='single-rf-harmonic', sigma_z=sigma_z,
+							particle_ref=particle_ref, return_matcher=True)
+
+	# Initiate normalized coordinates 
+	x_norm = np.random.normal(size=num_particles)
+	px_norm = np.random.normal(size=num_particles)
+	y_norm = np.random.normal(size=num_particles)
+	py_norm = np.random.normal(size=num_particles)
+
+	# If not provided, use number of particles as intensity 
+	if total_intensity_particles is None:   
+		total_intensity_particles = num_particles
+
+	particles = build_particles(_context=None, particle_ref=particle_ref,
+				zeta=zeta, delta=delta, 
+				x_norm=x_norm, px_norm=px_norm,
+				y_norm=y_norm, py_norm=py_norm,
+				nemitt_x=nemitt_x, nemitt_y=nemitt_y,
+				weight=total_intensity_particles/num_particles, line=line)
+
+	if return_matcher:
+		return particles, matcher
+	else:
+		return particles