Added a 4D Kurth distribution and example. (#114)

* Added a 4D Kurth distribution and example.
ECP-WarpX · Apr 20, 2022 · c8c8d0e · c8c8d0e
1 parent a369232
commit c8c8d0e
Show file tree

Hide file tree

Showing 7 changed files with 329 additions and 0 deletions.
diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst
@@ -12,6 +12,7 @@ This section allows you to **download input files** that correspond to different
    examples/chicane/README.rst
    examples/cfchannel/README.rst
    examples/kurth/README.rst
+   examples/fodo_rf/README.rst
 
 
 Test Cases / Benchmarks

diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -135,6 +135,20 @@ Initial Beam Distributions
         * ``<distribution>.muypy`` (``float``, dimensionless, default: ``0``) correlation Y-Py
         * ``<distribution>.mutpt`` (``float``, dimensionless, default: ``0``) correlation T-Pt
 
+    * ``kurth4d`` for initial 4D Kurth distribution in the transverse plane.
+      The distribution is uniform in t and Gaussian in pt.
+      With additional parameters:
+
+        * ``<distribution>.sigx`` (``float``, in meters) rms X
+        * ``<distribution>.sigy`` (``float``, in meters) rms Y
+        * ``<distribution>.sigt`` (``float``, in radian) rms normalized time difference T
+        * ``<distribution>.sigpx`` (``float``, in momentum) rms Px
+        * ``<distribution>.sigpy`` (``float``, in momentum) rms Py
+        * ``<distribution>.sigpt`` (``float``, in energy deviation) rms Pt
+        * ``<distribution>.muxpx`` (``float``, dimensionless, default: ``0``) correlation X-Px
+        * ``<distribution>.muypy`` (``float``, dimensionless, default: ``0``) correlation Y-Py
+        * ``<distribution>.mutpt`` (``float``, dimensionless, default: ``0``) correlation T-Pt
+
 .. _running-cpp-parameters-lattice:
 
 Lattice Elements

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -133,3 +133,17 @@ add_test(NAME FODO_RF.analysis
 )
 
 set_tests_properties(FODO_RF.analysis PROPERTIES DEPENDS "FODO_RF.run")
+
+
+# 4D Kurth Distribution Test ############################################################
+#
+add_test(NAME kurth4d.run
+         COMMAND $<TARGET_FILE:app> ${ImpactX_SOURCE_DIR}/examples/distgen/input_kurth4d.in
+         WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+)
+add_test(NAME kurth4d.analysis
+         COMMAND ${ImpactX_SOURCE_DIR}/examples/distgen/analysis_kurth4d.py
+         WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+)
+
+set_tests_properties(kurth4d.analysis PROPERTIES DEPENDS "kurth4d.run")
diff --git a/examples/distgen/analysis_kurth4d.py b/examples/distgen/analysis_kurth4d.py
@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+
+import glob
+
+import numpy as np
+import pandas as pd
+from scipy.stats import moment
+
+
+def get_moments(beam):
+    """Calculate standard deviations of beam position & momenta
+    and emittance values
+
+    Returns
+    -------
+    sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
+    """
+    sigx = moment(beam["x"], moment=2)**0.5  # variance -> std dev.
+    sigpx = moment(beam["px"], moment=2)**0.5
+    sigy = moment(beam["y"], moment=2)**0.5
+    sigpy = moment(beam["py"], moment=2)**0.5
+    sigt = moment(beam["t"], moment=2)**0.5
+    sigpt = moment(beam["pt"], moment=2)**0.5
+
+    epstrms = beam.cov(ddof=0)
+    emittance_x = (sigx**2 * sigpx**2 - epstrms["x"]["px"]**2)**0.5
+    emittance_y = (sigy**2 * sigpy**2 - epstrms["y"]["py"]**2)**0.5
+    emittance_t = (sigt**2 * sigpt**2 - epstrms["t"]["pt"]**2)**0.5
+
+    return (
+        sigx, sigy, sigt,
+        emittance_x, emittance_y, emittance_t)
+
+
+def read_all_files(file_pattern):
+    """Read in all CSV files from each MPI rank (and potentially OpenMP
+    thread). Concatenate into one Pandas dataframe.
+
+    Returns
+    -------
+    pandas.DataFrame
+    """
+    return pd.concat(
+        (
+            pd.read_csv(filename, delimiter=r"\s+")
+            for filename in glob.glob(file_pattern)
+        ),
+        axis=0,
+        ignore_index=True,
+    )
+
+
+# initial/final beam on rank zero
+initial = read_all_files("diags/initial_beam.txt.*")
+final = read_all_files("diags/output_beam.txt.*")
+
+# compare number of particles
+num_particles = 10000
+assert num_particles == len(initial)
+assert num_particles == len(final)
+
+print("Initial Beam:")
+sigx, sigy, sigt, \
+    emittance_x, emittance_y, emittance_t = get_moments(initial)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}")
+
+atol = 1.0  # a big number
+rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt,
+     emittance_x, emittance_y, emittance_t],
+    [9.970922e-04, 9.908808e-04, 9.992460e-04,
+     9.878659e-07, 9.966353e-07, 1.994764e-06],
+    rtol=rtol,
+    atol=atol
+)
+
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, \
+    emittance_x, emittance_y, emittance_t = get_moments(final)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}")
+
+atol = 1.0  # a big number
+rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt,
+     emittance_x, emittance_y, emittance_t],
+    [9.885251e-04, 1.006606e-03, 1.103184e-03,
+     9.878658e-07, 9.966353e-07, 1.994764e-06],
+    rtol=rtol,
+    atol=atol
+)
diff --git a/examples/distgen/input_kurth4d.in b/examples/distgen/input_kurth4d.in
@@ -0,0 +1,36 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.energy = 2.0e3
+beam.charge = 0.0
+beam.particle = proton
+beam.distribution = kurth4d
+beam.sigmaX = 1.0e-3
+beam.sigmaY = 1.0e-3
+beam.sigmaT = 1.0e-3
+beam.sigmaPx = 1.0e-3
+beam.sigmaPy = 1.0e-3
+beam.sigmaPt = 2.0e-3
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+
+lattice.elements = constf1
+
+constf1.type = constf
+constf1.ds = 2.0
+constf1.kx = 1.0
+constf1.ky = 1.0
+constf1.kt = 1.0e-4
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp
@@ -10,6 +10,7 @@
 #include "particles/distribution/Kurth6D.H"
 #include "particles/distribution/Gaussian.H"
 #include "particles/distribution/KVdist.H"
+#include "particles/distribution/Kurth4D.H"
 
 #include <AMReX.H>
 #include <AMReX_REAL.H>
@@ -195,6 +196,26 @@ namespace impactx
 
           generate_add_particles(*m_particle_container, qm, bunch_charge, kvDist, npart);
 
+        } else if (distribution_type == "kurth4d") {
+          amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
+          amrex::ParticleReal muxpx = 0.0, muypy = 0.0, mutpt = 0.0;
+          pp_dist.get("sigmaX", sigx);
+          pp_dist.get("sigmaY", sigy);
+          pp_dist.get("sigmaT", sigt);
+          pp_dist.get("sigmaPx", sigpx);
+          pp_dist.get("sigmaPy", sigpy);
+          pp_dist.get("sigmaPt", sigpt);
+          pp_dist.query("muxpx", muxpx);
+          pp_dist.query("muypy", muypy);
+          pp_dist.query("mutpt", mutpt);
+
+          impactx::distribution::Kurth4D kurth4D(
+            sigx, sigy, sigt,
+            sigpx, sigpy, sigpt,
+            muxpx, muypy, mutpt);
+
+          generate_add_particles(*m_particle_container, qm, bunch_charge, kurth4D, npart);
+
         } else {
             amrex::Abort("Unknown distribution: " + distribution_type);
         }

diff --git a/src/particles/distribution/Kurth4D.H b/src/particles/distribution/Kurth4D.H
@@ -0,0 +1,138 @@
+/* Copyright 2022 Chad Mitchell, Axel Huebl
+ *
+ * This file is part of ImpactX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_DISTRIBUTION_KURTH4D
+#define IMPACTX_DISTRIBUTION_KURTH4D
+
+#include <AMReX_Random.H>
+#include <AMReX_REAL.H>
+
+#include <cmath>
+
+
+namespace impactx
+{
+namespace distribution
+{
+    struct Kurth4D
+    {
+        /** A 4D Kurth distribution transversely + a uniform distribution
+         *  it t + a Gaussian distribution in pt
+         *
+         * Return sampling from a 4D Kurth + uniform distribution.
+         *
+         * @param sigx,sigy,sigt for zero correlation, these are the related
+         *                                   RMS sizes (in meters)
+         * @param sigpx,sigpy,sigpt RMS momentum
+         * @param muxpx,muypy,mutpt correlation length-momentum
+         */
+         Kurth4D(amrex::ParticleReal const sigx, amrex::ParticleReal const sigy,
+         amrex::ParticleReal const sigt,amrex::ParticleReal const sigpx,
+         amrex::ParticleReal const sigpy,amrex::ParticleReal const sigpt,
+         amrex::ParticleReal const muxpx=0.0, amrex::ParticleReal const muypy=0.0,
+         amrex::ParticleReal const mutpt=0.0
+         )
+          : m_sigmaX(sigx),m_sigmaY(sigy),m_sigmaT(sigt),m_sigmaPx(sigpx),m_sigmaPy(sigpy),
+            m_sigmaPt(sigpt),m_muxpx(muxpx),m_muypy(muypy),m_mutpt(mutpt)
+          {
+          }
+
+        /** Return 1 6D particle coordinate
+         *
+         * @param x particle position in x
+         * @param y particle position in y
+         * @param t particle position in t
+         * @param px particle momentum in x
+         * @param py particle momentum in y
+         * @param pt particle momentum in t
+         * @param engine a random number engine (with associated state)
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                amrex::ParticleReal & x,
+                amrex::ParticleReal & y,
+                amrex::ParticleReal & t,
+                amrex::ParticleReal & px,
+                amrex::ParticleReal & py,
+                amrex::ParticleReal & pt,
+                amrex::RandomEngine const& engine) const {
+
+            using namespace amrex::literals;
+
+            amrex::ParticleReal v,phi,r,u1,u2,ln1;
+            amrex::ParticleReal alpha,u,Lz,pmax,pr,pphi;
+            amrex::ParticleReal root,a1,a2;
+
+            constexpr amrex::ParticleReal pi = 3.14159265358979_prt;
+
+            // Sample and transform to define (x,y):
+            v = amrex::Random(engine);
+            phi = amrex::Random(engine);
+            phi = 2_prt*pi*phi;
+            r = sqrt(v);
+            x = r*cos(phi);
+            y = r*sin(phi);
+
+            // Random samples used to define Lz:
+            u = amrex::Random(engine);
+            Lz = r*(2.0_prt*u-1.0_prt);
+
+            // Random samples used to define pr:
+            alpha = amrex::Random(engine);
+            alpha = pi*alpha;
+            pmax = 1.0_prt - pow((Lz/r),2) - pow(r,2) + pow(Lz,2);
+            pmax = sqrt(pmax);
+            pr = pmax*cos(alpha);
+            pphi = Lz/r;
+
+            // Transformations used to obtain (px,py):
+            px = pr*cos(phi)-pphi*sin(phi);
+            py = pr*sin(phi)+pphi*cos(phi);
+
+            // Sample and transform to define (t,pt):
+            t = amrex::Random(engine);
+            t = 2.0_prt*(t-0.5_prt);
+            u1 = amrex::Random(engine);
+            u2 = amrex::Random(engine);
+            ln1 = sqrt(-2_prt*log(u1));
+            pt = ln1*cos(2_prt*pi*u2);
+
+            // Scale to produce the identity covariance matrix:
+            amrex::ParticleReal c = sqrt(3.0_prt);
+            x = 2_prt*x;
+            y = 2_prt*y;
+            t = c*t;
+            px = 2_prt*px;
+            py = 2_prt*py;
+            // pt = pt;
+
+            // Transform to produce the desired second moments/correlations:
+            root = sqrt(1.0_prt-m_muxpx*m_muxpx);
+            a1 = m_sigmaX*x/root;
+            a2 = m_sigmaPx*(-m_muxpx*x/root+px);
+            x = a1;
+            px = a2;
+            root = sqrt(1.0_prt-m_muypy*m_muypy);
+            a1 = m_sigmaY*y/root;
+            a2 = m_sigmaPy*(-m_muypy*y/root+py);
+            y = a1;
+            py = a2;
+            root = sqrt(1.0_prt-m_mutpt*m_mutpt);
+            a1 = m_sigmaT*t/root;
+            a2 = m_sigmaPt*(-m_mutpt*t/root+pt);
+            t = a1;
+            pt = a2;
+        }
+    private:
+        amrex::ParticleReal m_sigmaX,m_sigmaY,m_sigmaT; //! related RMS sizes (length)
+        amrex::ParticleReal m_sigmaPx,m_sigmaPy,m_sigmaPt; //! RMS momentum
+        amrex::ParticleReal m_muxpx,m_muypy,m_mutpt; //! correlation length-momentum
+    };
+
+} // namespace distribution
+} // namespace impactx
+
+#endif // IMPACTX_DISTRIBUTION_KURTH4D