initialize the distributions on the GPU (#498)

* initialize the distributions on the GPU

* Fix Thermal: `const` operator

* Python Support, Cleaning

* Cleanup: Unused, Resize

Fix the missing resize before redistribute.
Remove the old, unusued, serial initialization logic.

* Thermal Distribution: GPU Memory

Remove stack memory usage and use GPU memory for pre-computed CDF.

* Cosmetic Cleaning

* Waterbag: Clean C++ Math

Avoid promoting to double in C functions.

* Cosmetic Cleaning for Clarity

* Fix CUDA Compile

On older NVCC

* Analysis: Slightly Increase Sampling Tolerance

Run on local GPU (RTX A2000).

---------

Co-authored-by: Axel Huebl <[email protected]>

docs/source/usage/python.rst

@@ -237,15 +237,14 @@ Particles
 
    This class stores particles, distributed over MPI ranks.
 
-   .. py:method:: add_n_particles(lev, x, y, t, px, py, pt, qm, bchchg)
+   .. py:method:: add_n_particles(x, y, t, px, py, pt, qm, bchchg)
 
       Add new particles to the container for fixed s.
 
       Note: This can only be used *after* the initialization (grids) have
             been created, meaning after the call to :py:meth:`ImpactX.init_grids`
             has been made in the ImpactX class.
 
-      :param lev: mesh-refinement level
       :param x: positions in x
       :param y: positions in y
       :param t: positions as time-of-flight in c*t

examples/alignment/analysis_alignment.py

@@ -76,7 +76,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = 1.8 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.0 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/apochromatic/analysis_apochromatic.py

@@ -108,7 +108,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = 15.0 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 19.0 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/chicane/analysis_chicane.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/compression/analysis_compression.py

@@ -69,7 +69,7 @@ def get_moments(beam):
 print(f"  gamma={initial_gamma_ref:e}")
 
 atol = 0.0  # ignored
-rtol = 1.8 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.9 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -99,7 +99,7 @@ def get_moments(beam):
 
 
 atol = 0.0  # ignored
-rtol = 1.8 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.9 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/cyclotron/analysis_cyclotron.py

@@ -86,7 +86,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = 1.8 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.0 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/distgen/analysis_gaussian.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.3 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.3 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/distgen/analysis_kurth4d.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.3 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.3 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/distgen/analysis_semigaussian.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.3 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.3 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/expanding_beam/analysis_expanding.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/fodo/analysis_fodo.py

@@ -59,7 +59,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -86,7 +86,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/fodo_chromatic/analysis_fodo_chr.py

@@ -59,7 +59,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -86,7 +86,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/fodo_rf/analysis_fodo_rf.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/iota_lattice/analysis_iotalattice.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # a big number
-rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.9 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -78,7 +78,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # a big number
-rtol = 1.8 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.9 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/iota_lattice/analysis_iotalattice_sdep.py

@@ -56,7 +56,7 @@ def read_all_files(file_pattern):
 print(f"  meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
 
 atol = 0.0  # a big number
-rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.6 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -73,7 +73,7 @@ def read_all_files(file_pattern):
 print(f"  meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
 
 atol = 0.0  # a big number
-rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.6 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -94,14 +94,14 @@ def read_all_files(file_pattern):
 Hrms = np.sqrt(sigH**2 + meanH**2)
 Irms = np.sqrt(sigI**2 + meanI**2)
 
-atol = 4.0e-3 * Hrms
+atol = 4.2e-3 * Hrms
 rtol = 0.0  # large number
 print()
 print(f"  atol={atol} (ignored: rtol~={rtol})")
 print(f"  dH_max={beam_joined['dH'].max()}")
 assert np.allclose(beam_joined["dH"], 0.0, rtol=rtol, atol=atol)
 
-atol = 5.5e-3 * Irms
+atol = 5.7e-3 * Irms
 rtol = 0.0
 print()
 print(f"  atol={atol} (ignored: rtol~={rtol})")

examples/kurth/analysis_kurth_periodic.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.7 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.7 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/multipole/analysis_multipole.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.5 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/quadrupole_softedge/analysis_quadrupole_softedge.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

examples/solenoid/analysis_solenoid.py

@@ -58,7 +58,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.3 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(
@@ -85,7 +85,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.3 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(