xfields 0.22.2__tar.gz → 0.23.0__tar.gz

xfields-0.23.0/.gitignore

@@ -0,0 +1,13 @@
+*.DS_Store
+*.swp
+*.so
+*.pyc
+.project
+.pydevproject
+*.png
+norepository_*
+*~
+build
+CyFPPS.cpp
+*.egg-info
+.idea/**

{xfields-0.22.2 → xfields-0.23.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xfields
-Version: 0.22.2
+Version: 0.23.0
 Summary: Field Maps and Particle In Cell
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xfields

xfields-0.23.0/contributors_and_credits.txt

@@ -0,0 +1,16 @@
+The following people contributed to the development of this package:
+
+CERN contributors:
+ - H. Bartosik
+ - X. Buffat
+ - R. De Maria
+ - L. Giacomel
+ - G. Iadarola
+ - P. Kicsiny
+ - S. Lopaciuk
+ - K. Paraschou
+ - M. Schwinzerl
+ - F. Soubelet
+ - E. Waagaard
+
+The Particle In Cell implementation is largely based on the PyPIC package (https://github.com/pycomplete/pypic)

xfields-0.23.0/examples/000_sphere/.gitignore

@@ -0,0 +1 @@
+picsphere.pkl

xfields-0.23.0/examples/000_sphere/000_sphere_xfields3D.py

@@ -0,0 +1,167 @@
+# copyright ################################# #
+# This file is part of the Xfields Package.   #
+# Copyright (c) CERN, 2021.                   #
+# ########################################### #
+
+import time
+
+import numpy as np
+from numpy.random import rand
+from scipy.constants import epsilon_0
+from numpy import pi
+
+from xobjects import ContextCpu, ContextCupy, ContextPyopencl
+from xfields import TriLinearInterpolatedFieldMap
+
+context = ContextCpu()
+context = ContextCupy(default_block_size=256)
+#context = ContextPyopencl()
+
+print(repr(context))
+
+import matplotlib.pyplot as plt
+plt.close('all')
+
+center_xyz = np.array([0.1, 0.2, -0.3])
+radius = .5
+n_part_cube = 10000000
+
+x_cube = radius*(2. * rand(n_part_cube) - 1.) + center_xyz[0]
+y_cube = radius*(2. * rand(n_part_cube) - 1.) + center_xyz[1]
+z_cube = radius*(2. * rand(n_part_cube) - 1.) + center_xyz[2]
+
+mask_sphere = ((x_cube - center_xyz[0])**2
+             + (y_cube - center_xyz[1])**2
+             + (z_cube - center_xyz[2])**2) < radius**2
+x = x_cube[mask_sphere]
+y = y_cube[mask_sphere]
+z = z_cube[mask_sphere]
+pweights = 1. + 0*x
+
+
+# PIC
+x_lim = (-1.1, 1.)
+y_lim = (-1.2, 1.)
+z_lim = (-1.3, 1.)
+
+nx = 128
+ny = 128
+nz = 200
+
+theta_test = 70. * np.pi/360
+phi_test = 300 * np.pi/360.
+r_test = 2.5*radius*np.linspace(-1, 1., 1000)
+
+x0_test = r_test * np.sin(theta_test) * np.cos(phi_test)
+y0_test = r_test * np.sin(theta_test) * np.sin(phi_test)
+z0_test = r_test * np.cos(theta_test)
+
+x_test = center_xyz[0] + x0_test
+y_test = center_xyz[1] + y0_test
+z_test = center_xyz[2] + z0_test
+
+# Moving particles and test coordinates to GPU (if needed)
+np2platf = context.nparray_to_context_array
+x_dev = np2platf(x)
+y_dev = np2platf(y)
+z_dev = np2platf(z)
+pweights_dev = np2platf(pweights)
+x_test_dev = np2platf(x_test)
+y_test_dev = np2platf(y_test)
+z_test_dev = np2platf(z_test)
+
+###############
+# Actual test #
+###############
+
+# Build fieldmap object
+fmap = TriLinearInterpolatedFieldMap(
+        _context=context,
+        x_range=x_lim, nx=nx,
+        y_range=y_lim, ny=ny,
+        z_range=z_lim, nz=nz,
+        solver='FFTSolver3D')
+
+
+# Compute potential
+n_rep = 4
+for _ in range(n_rep):
+    t1 = time.time()
+    fmap.update_from_particles(x_p=x_dev, y_p=y_dev, z_p=z_dev,
+            ncharges_p=pweights_dev, q0_coulomb=1.)
+
+    # Check on the x axis
+    (rho_test_dev, phi_test_dev, dx_test_dev, dy_test_dev,
+            dz_test_dev) = fmap.get_values_at_points(
+                x=x_test_dev, y=y_test_dev, z=z_test_dev)
+    context.synchronize()
+    t2 = time.time()
+    print(f'Time: {t2-t1:.2e}')
+
+# Copy back for plotting
+platf2np = context.nparray_from_context_array
+rho_test =  platf2np(rho_test_dev)
+phi_test = platf2np(phi_test_dev)
+dx_test = platf2np(dx_test_dev)
+dy_test = platf2np(dy_test_dev)
+dz_test = platf2np(dz_test_dev)
+
+# pickle for other tests
+import pickle
+# with open('picsphere.pkl', 'wb') as fid:
+#     pickle.dump({
+#         'fmap': fmap,
+#         'x_test': x_test,
+#         'y_test': y_test,
+#         'z_test': z_test,
+#         }, fid)
+
+####################
+# Plots and checks #
+####################
+
+ex_test = -dx_test
+ey_test = -dy_test
+ez_test = -dz_test
+
+plt.figure(100)
+plt.plot(r_test, rho_test)
+plt.axhline(y=len(x)/(4/3*np.pi*radius**3))
+
+plt.figure(101)
+plt.plot(r_test, phi_test)
+
+e_ref = len(x)/(4*pi*epsilon_0) * (
+        r_test/radius**3*(np.abs(r_test)<radius)
+      + np.sign(r_test)/r_test**2*(np.abs(r_test)>=radius))
+plt.figure(102)
+plt.plot(r_test, ex_test*x0_test/r_test
+        + ey_test*y0_test/r_test + ez_test*z0_test/r_test)
+plt.plot(r_test, e_ref)
+plt.grid(True)
+
+# Check integral
+int_rho = np.sum(platf2np(fmap.rho))*fmap.dx*fmap.dy*fmap.dz
+assert np.isclose(int_rho, len(x))
+
+
+fig1 = plt.figure(1)
+ax1 = fig1.add_subplot(111)
+ax1.pcolormesh(fmap.x_grid, fmap.y_grid,
+        np.sum(platf2np(fmap.rho), axis=2).T, shading='gouraud')
+ax1.set_aspect('equal')
+ax1.add_patch(plt.Circle((center_xyz[0], center_xyz[1]), radius,
+                         color='w', fill=False))
+ax1.set_xlabel('x')
+ax1.set_ylabel('y')
+
+fig2 = plt.figure(2)
+ax2 = fig2.add_subplot(111)
+ax2.pcolormesh(fmap.y_grid, fmap.z_grid,
+        np.sum(platf2np(fmap.rho), axis=0).T, shading='gouraud')
+ax2.set_aspect('equal')
+ax2.add_patch(plt.Circle((center_xyz[1], center_xyz[2]), radius,
+                         color='w', fill=False))
+ax2.set_xlabel('y')
+ax2.set_ylabel('z')
+plt.show()

xfields-0.23.0/examples/000_sphere/001_sphere_xfields2p5D.py

@@ -0,0 +1,170 @@
+# copyright ################################# #
+# This file is part of the Xfields Package.   #
+# Copyright (c) CERN, 2021.                   #
+# ########################################### #
+
+import time
+
+import numpy as np
+from numpy.random import rand
+from scipy.constants import epsilon_0
+from numpy import pi
+
+from xobjects import ContextCpu, ContextCupy, ContextPyopencl
+from xfields import TriLinearInterpolatedFieldMap
+
+context = ContextCpu()
+#context = ContextCupy(default_block_size=256)
+#context = ContextPyopencl()
+
+print(repr(context))
+
+import matplotlib.pyplot as plt
+plt.close('all')
+
+center_xyz = np.array([0.1, 0.2, -0.3])
+radius = .5
+n_part_cube = 10000000
+
+x_cube = radius*(2. * rand(n_part_cube) - 1.) + center_xyz[0]
+y_cube = radius*(2. * rand(n_part_cube) - 1.) + center_xyz[1]
+z_cube = radius*(2. * rand(n_part_cube) - 1.) + center_xyz[2]
+
+mask_sphere = ((x_cube - center_xyz[0])**2
+             + (y_cube - center_xyz[1])**2
+             + (z_cube - center_xyz[2])**2) < radius**2
+x = x_cube[mask_sphere]
+y = y_cube[mask_sphere]
+z = z_cube[mask_sphere]
+pweights = 1. + 0*x
+
+
+# PIC
+x_lim = (-1.1, 1.)
+y_lim = (-1.2, 1.)
+z_lim = (-1.3, 1.)
+
+nx = 256
+ny = 256
+nz = 50
+
+theta_test = 70. * np.pi/360
+r_test = 2.5*radius*np.linspace(-1, 1., 1000)
+z_test = 0.4
+
+x0_test = r_test * np.cos(theta_test)
+y0_test = r_test * np.sin(theta_test)
+z0_test = z_test + 0*r_test
+
+x_test = center_xyz[0] + x0_test
+y_test = center_xyz[1] + y0_test
+z_test = center_xyz[2] + z0_test
+
+# Moving particles and test coordinates to GPU (if needed)
+np2platf = context.nparray_to_context_array
+x_dev = np2platf(x)
+y_dev = np2platf(y)
+z_dev = np2platf(z)
+pweights_dev = np2platf(pweights)
+x_test_dev = np2platf(x_test)
+y_test_dev = np2platf(y_test)
+z_test_dev = np2platf(z_test)
+
+###############
+# Actual test #
+###############
+
+# Build fieldmap object
+fmap = TriLinearInterpolatedFieldMap(
+    _context=context,
+    x_range=x_lim, nx=nx,
+    y_range=y_lim, ny=ny, z_range=z_lim, nz=nz,
+    solver='FFTSolver2p5D')
+
+
+# Compute potential
+n_rep = 10
+for _ in range(n_rep):
+    t1 = time.time()
+    fmap.update_from_particles(x_p=x_dev, y_p=y_dev, z_p=z_dev,
+            ncharges_p=pweights_dev, q0_coulomb=1.)
+
+    # Check on the x axis
+    (rho_test_dev, phi_test_dev, dx_test_dev, dy_test_dev,
+            dz_test_dev) = fmap.get_values_at_points(
+                x=x_test_dev, y=y_test_dev, z=z_test_dev)
+    context.synchronize()
+    t2 = time.time()
+    print(f'Time: {t2-t1:.2e}')
+
+# Copy back for plotting
+platf2np = context.nparray_from_context_array
+rho_test =  platf2np(rho_test_dev)
+phi_test = platf2np(phi_test_dev)
+dx_test = platf2np(dx_test_dev)
+dy_test = platf2np(dy_test_dev)
+dz_test = platf2np(dz_test_dev)
+
+# pickle for other tests
+import pickle
+# with open('picsphere.pkl', 'wb') as fid:
+#     pickle.dump({
+#         'fmap': fmap,
+#         'x_test': x_test,
+#         'y_test': y_test,
+#         'z_test': z_test,
+#         }, fid)
+
+####################
+# Plots and checks #
+####################
+
+ex_test = -dx_test
+ey_test = -dy_test
+ez_test = -dz_test
+
+plt.figure(100)
+plt.plot(r_test, rho_test)
+plt.axhline(y=len(x)/(4/3*np.pi*radius**3))
+
+plt.figure(101)
+plt.plot(r_test, phi_test)
+
+assert z0_test[0]<radius
+r_loc = np.sqrt(radius**2 - z0_test[0]**2)
+rho_ref = len(x)/(4/3*pi*radius**3)
+lam_loc = rho_ref*pi*r_loc**2
+e_ref = rho_ref/(epsilon_0) * (
+        r_test/2*(np.abs(r_test)<r_loc)
+      + r_loc**2/r_test/2*(np.abs(r_test)>=r_loc))
+plt.figure(102)
+plt.plot(r_test, ex_test*x0_test/r_test
+        + ey_test*y0_test/r_test)
+plt.plot(r_test, e_ref)
+plt.grid(True)
+
+# Check integral
+int_rho = np.sum(platf2np(fmap.rho))*fmap.dx*fmap.dy*fmap.dz
+assert np.isclose(int_rho, len(x))
+
+
+fig1 = plt.figure(1)
+ax1 = fig1.add_subplot(111)
+ax1.pcolormesh(fmap.x_grid, fmap.y_grid,
+        np.sum(platf2np(fmap.rho), axis=2).T, shading='gouraud')
+ax1.set_aspect('equal')
+ax1.add_patch(plt.Circle((center_xyz[0], center_xyz[1]), radius,
+                         color='w', fill=False))
+ax1.set_xlabel('x')
+ax1.set_ylabel('y')
+
+fig2 = plt.figure(2)
+ax2 = fig2.add_subplot(111)
+ax2.pcolormesh(fmap.y_grid, fmap.z_grid,
+        np.sum(platf2np(fmap.rho), axis=0).T, shading='gouraud')
+ax2.set_aspect('equal')
+ax2.add_patch(plt.Circle((center_xyz[1], center_xyz[2]), radius,
+                         color='w', fill=False))
+ax2.set_xlabel('y')
+ax2.set_ylabel('z')
+plt.show()

xfields-0.23.0/examples/001_spacecharge/000_spacecharge_example.py

@@ -0,0 +1,152 @@
+# copyright ################################# #
+# This file is part of the Xfields Package.   #
+# Copyright (c) CERN, 2021.                   #
+# ########################################### #
+
+import time
+
+import numpy as np
+
+from xobjects import ContextCpu, ContextCupy, ContextPyopencl
+import xtrack as xt
+import xpart as xp
+
+###################
+# Choose context #
+###################
+
+context = ContextCpu(omp_num_threads=0) # no omp
+#context = ContextCpu(omp_num_threads=48) # omp
+#context = ContextCupy(default_block_size=256)
+#context = ContextPyopencl('0.0')
+
+print(repr(context))
+
+#################################
+# Generate particles and probes #
+#################################
+
+n_macroparticles = int(1e6)
+bunch_intensity = 2.5e11
+sigma_x = 3e-3
+sigma_y = 2e-3
+sigma_z = 30e-2
+p0c = 25.92e9
+mass = xp.pmass
+theta_probes = 30 * np.pi/180
+r_max_probes = 2e-2
+z_probes = 1.2*sigma_z
+n_probes = 1000
+
+from xfields.test_support.temp_makepart import generate_particles_object
+(particles_dtk, r_probes, x_probes,
+        y_probes, z_probes) = generate_particles_object(
+                            n_macroparticles,
+                            bunch_intensity,
+                            sigma_x,
+                            sigma_y,
+                            sigma_z,
+                            p0c,
+                            mass,
+                            n_probes,
+                            r_max_probes,
+                            z_probes,
+                            theta_probes)
+particles = xp.Particles(
+        _context=context, **particles_dtk.to_dict())
+
+
+######################
+# Space charge (PIC) #
+######################
+
+x_lim = 5.*sigma_x
+y_lim = 5.*sigma_y
+z_lim = 5.*sigma_z
+
+from xfields import SpaceCharge3D
+spcharge = SpaceCharge3D(
+        _context=context,
+        length=1, update_on_track=True, apply_z_kick=False,
+        x_range=(-x_lim, x_lim),
+        y_range=(-y_lim, y_lim),
+        z_range=(-z_lim, z_lim),
+        nx=256, ny=256, nz=100,
+        # solver='FFTSolver2p5D',
+        solver='FFTSolver2p5DAveraged',
+        gamma0=particles_dtk.gamma0[0])
+
+spcharge.track(particles)
+
+
+#############################
+# Compare against ducktrack #
+#############################
+
+
+p2np = context.nparray_from_context_array
+
+from ducktrack import SCQGaussProfile, TestParticles
+scdtk = SCQGaussProfile(
+        number_of_particles = bunch_intensity,
+        bunchlength_rms=sigma_z,
+        sigma_x=sigma_x,
+        sigma_y=sigma_y,
+        length=spcharge.length,
+        x_co=0.,
+        y_co=0.)
+
+p_dtk = TestParticles(p0c=p0c,
+        mass=mass,
+        x=x_probes.copy(),
+        y=y_probes.copy(),
+        zeta=z_probes.copy())
+
+scdtk.track(p_dtk)
+
+mask_inside_grid = ((np.abs(x_probes)<0.9*x_lim) &
+                    (np.abs(y_probes)<0.9*y_lim))
+
+
+import matplotlib.pyplot as plt
+plt.close('all')
+plt.figure(1)
+plt.subplot(211)
+plt.plot(r_probes, p_dtk.px, color='red')
+plt.plot(r_probes, p2np(particles.px[:n_probes]), color='blue',
+        linestyle='--')
+plt.subplot(212)
+plt.plot(r_probes, p_dtk.py, color='red')
+plt.plot(r_probes, p2np(particles.py[:n_probes]), color='blue',
+        linestyle='--')
+
+plt.figure(2)
+plt.subplot(111)
+tr_profile = spcharge.fieldmap.rho.sum(axis=0)
+for ii in range(tr_profile.shape[1]):
+    tr_profile[:,ii] = tr_profile[:,ii]/tr_profile[:,ii].max()
+tr_profile = context.nparray_from_context_array(tr_profile)
+plt.plot(tr_profile)
+
+###########
+# Time it #
+###########
+
+n_rep = 5
+
+for _ in range(n_rep):
+    t1 = time.time()
+    spcharge.track(particles)
+    context.synchronize()
+    t2 = time.time()
+    print(f'Time (full PIC): {(t2-t1)*1e3:.2f} ms')
+
+spcharge.update_on_track = False
+for _ in range(n_rep):
+    t1 = time.time()
+    spcharge.track(particles)
+    context.synchronize()
+    t2 = time.time()
+    print(f'Time (interp only): {(t2-t1)*1e3:.2f} ms')
+
+plt.show()

xfields-0.23.0/examples/001_spacecharge/001_shared_fieldmap.py

@@ -0,0 +1,152 @@
+# copyright ################################# #
+# This file is part of the Xfields Package.   #
+# Copyright (c) CERN, 2021.                   #
+# ########################################### #
+
+import time
+
+import numpy as np
+
+from xline.particles import Particles
+from xobjects import ContextCpu, ContextCupy, ContextPyopencl
+import xtrack as xt
+
+###################
+# Choose context #
+###################
+
+#context = ContextCpu(omp_num_threads=0) # no omp
+context = ContextCpu(omp_num_threads=1) # omp
+#context = ContextCpu(omp_num_threads=48) # omp
+context = ContextCupy(default_block_size=256)
+#context = ContextPyopencl('0.0')
+
+print(repr(context))
+
+#################################
+# Generate particles and probes #
+#################################
+
+n_macroparticles = int(1e6)
+bunch_intensity = 2.5e11
+sigma_x = 3e-3
+sigma_y = 2e-3
+sigma_z = 30e-2
+p0c = 25.92e9
+mass = Particles.pmass
+theta_probes = 30 * np.pi/180
+r_max_probes = 2e-2
+z_probes = 1.2*sigma_z
+n_probes = 1000
+
+from xfields.test_support.temp_makepart import generate_particles_object
+(particles_pyst, r_probes, x_probes,
+        y_probes, z_probes) = generate_particles_object(
+                            n_macroparticles,
+                            bunch_intensity,
+                            sigma_x,
+                            sigma_y,
+                            sigma_z,
+                            p0c,
+                            mass,
+                            n_probes,
+                            r_max_probes,
+                            z_probes,
+                            theta_probes)
+particles = xt.Particles(
+        _context=context, **particles_pyst.to_dict())
+
+
+######################
+# Space charge (PIC) #
+######################
+
+x_lim = 5.*sigma_x
+y_lim = 5.*sigma_y
+z_lim = 5.*sigma_z
+
+from xfields import SpaceCharge3D
+spcharge_parent = SpaceCharge3D(
+        _context=context,
+        length=1, update_on_track=True, apply_z_kick=False,
+        x_range=(-x_lim, x_lim),
+        y_range=(-y_lim, y_lim),
+        z_range=(-z_lim, z_lim),
+        nx=256, ny=256, nz=100,
+        solver='FFTSolver2p5D',
+        gamma0=particles_pyst.gamma0)
+
+spcharge= SpaceCharge3D(
+        _buffer=spcharge_parent._buffer,
+        length=1, update_on_track=True, apply_z_kick=False,
+        fieldmap=spcharge_parent.fieldmap)
+
+spcharge.fieldmap.rho[0,0,0] = 100.
+assert np.isclose(spcharge_parent.fieldmap.rho[0,0,0], 100.)
+
+spcharge.track(particles)
+
+
+#########################
+# Compare against xline #
+#########################
+
+
+p2np = context.nparray_from_context_array
+
+from xline.elements import SCQGaussProfile
+scpyst = SCQGaussProfile(
+        number_of_particles = bunch_intensity,
+        bunchlength_rms=sigma_z,
+        sigma_x=sigma_x,
+        sigma_y=sigma_y,
+        length=spcharge.length,
+        x_co=0.,
+        y_co=0.)
+
+p_pyst = Particles(p0c=p0c,
+        mass=mass,
+        x=x_probes.copy(),
+        y=y_probes.copy(),
+        zeta=z_probes.copy())
+
+scpyst.track(p_pyst)
+
+mask_inside_grid = ((np.abs(x_probes)<0.9*x_lim) &
+                    (np.abs(y_probes)<0.9*y_lim))
+
+
+import matplotlib.pyplot as plt
+plt.close('all')
+plt.figure()
+plt.subplot(211)
+plt.plot(r_probes, p_pyst.px, color='red')
+plt.plot(r_probes, p2np(particles.px[:n_probes]), color='blue',
+        linestyle='--')
+plt.subplot(212)
+plt.plot(r_probes, p_pyst.py, color='red')
+plt.plot(r_probes, p2np(particles.py[:n_probes]), color='blue',
+        linestyle='--')
+
+###########
+# Time it #
+###########
+
+n_rep = 5
+
+for _ in range(n_rep):
+    t1 = time.time()
+    spcharge.track(particles)
+    context.synchronize()
+    t2 = time.time()
+    print(f'Time (full PIC): {(t2-t1)*1e3:.2f} ms')
+
+spcharge.update_on_track = False
+for _ in range(n_rep):
+    t1 = time.time()
+    spcharge.track(particles)
+    context.synchronize()
+    t2 = time.time()
+    print(f'Time (interp only): {(t2-t1)*1e3:.2f} ms')
+
+plt.show()

xfields-0.23.0/examples/001_spacecharge/002_spacecharge_with_losses.py

@@ -0,0 +1,33 @@
+import numpy as np
+from scipy.constants import e as qe
+
+import xfields as xf
+import xpart as xp
+
+sc = xf.SpaceCharge3D(
+    x_range=(-0.01, 0.01),
+    y_range=(-0.01, 0.01),
+    z_range=(-1, 1),
+    nx=128,
+    ny=128,
+    nz=128,
+    solver='FFTSolver2p5D'
+)
+
+particles = xp.Particles(
+    x=[0.002, 0.003, 0.004],
+    y=[0.002, 0.003, 0.004],
+    zeta=[-0.5, 0., 0.05],
+    mass0=xp.PROTON_MASS_EV,
+    p0c=1e9,
+    weight=[1., 1., 1.])
+
+particles.state = np.array([1, 1, -20])
+
+sc.track(particles)
+
+dx = sc.fieldmap.dx
+dy = sc.fieldmap.dy
+dz = sc.fieldmap.dz
+
+charges = np.sum(sc.fieldmap.rho[:]) * dx * dy * dz / qe