Add twiss injection source

Docs/source/usage/parameters.rst

@@ -951,8 +951,134 @@ Particle initialization
 
           \sigma_{x,y}(z) &= \sigma^*_{x,y} \sqrt{1 + \left( \frac{z - z^*}{\beta^*_{x,y}} \right)^2}
 
+    * ``twiss``: Inject particle beam, whose particles have mean position :math:`\mathbf{x}_0` and mean normalized momentum :math:`u_0 \mathbf{\hat{n}}_z`, with distribution function
 
-    * ``external_file``: Inject macroparticles with properties (mass, charge, position, and momentum - :math:`\gamma \beta m c`) read from an external openPMD file.
+      .. math::
+         f \left( \mathbf{R} \cdot \mathbf{x} + \mathbf{x}_0, \mathbf{R} \cdot (\mathbf{u} + u_0 \mathbf{\hat{z}}) \right) = N f_x(x, u_x) f_y(y, u_z) f_\zeta(\zeta, u_\zeta),
+      where the rotation matrix is
+
+      .. math::
+         \mathbf{R} = \begin{pmatrix}
+           \mathbf{\hat{n}}_x \cdot \mathbf{\hat{x}} &
+           \mathbf{\hat{n}}_y \cdot \mathbf{\hat{x}} &
+           \mathbf{\hat{n}}_z \cdot \mathbf{\hat{x}} \\
+           \mathbf{\hat{n}}_x \cdot \mathbf{\hat{y}} &
+           \mathbf{\hat{n}}_y \cdot \mathbf{\hat{y}} &
+           \mathbf{\hat{n}}_z \cdot \mathbf{\hat{y}} \\
+           \mathbf{\hat{n}}_x \cdot \mathbf{\hat{z}} &
+           \mathbf{\hat{n}}_y \cdot \mathbf{\hat{z}} &
+           \mathbf{\hat{n}}_z \cdot \mathbf{\hat{z}} \end{pmatrix},
+
+      and the initialization coordinates are
+
+      .. math::
+         \mathbf{x} = x \mathbf{\hat{x}} + y \mathbf{\hat{y}} + \zeta \mathbf{\hat{z}}~~~~~\text{and}~~~~
+         \mathbf{u} = u_x \mathbf{\hat{x}} + u_y \mathbf{\hat{y}} + u_\zeta \mathbf{\hat{z}}.
+
+      Each factor of the distribution function is a bivariate normal distribution, conventionally described by the Courant-Snyder (Twiss) parameters,
+
+      .. math::
+         f_i(x_i, u_i) = \frac{1}{2 \pi \epsilon_i} \exp \left( - \frac{\gamma_i x_i^2 + 2 \alpha_i x_i u_i + \beta_i u_i^2}{2 \epsilon_i} \right),
+
+      where :math:`\epsilon_i` is the RMS normalized emittance.
+
+      * ``<species_name>.q_tot`` (`float`, coulomb) Beam charge :math:`Q_\text{tot} = N q`.
+
+      * ``<species_name>.npart`` (`int`) Number of macroparticles.
+
+      * ``<species_name>.x0/y0/z0`` (`float`, meter) Initial beam position :math:`\mathbf{x}_0`.
+
+      * ``<species_name>.twiss.planar_cut_x/y/zeta`` (`float`, optional)  Cut particles with
+
+        .. math::
+           |x| > a_x ~~\text{or}~~ |y| > a_y ~~\text{or}~~ |\zeta| > a_\zeta.
+
+        The charge of the uncut beam is ``<species_name>.q_tot``, so that cutting the distribution generally results in lower total injected charge.  Specify :math:`a_i` in units of :math:`\sigma_{x_i}`.
+
+      * ``<species_name>.twiss.ellipsoidal_cut_x/y/zeta`` (`float`, optional) Cut particles with
+
+        .. math::
+           \left( \frac{x}{a_x} \right)^2
+           + \left( \frac{y}{a_y} \right)^2
+           + \left( \frac{\zeta}{a_\zeta} \right)^2 > 1
+
+        The charge of the uncut beam is ``<species_name>.q_tot``, so that cutting the distribution generally results in lower total injected charge.  Specify :math:`a_i` in units of :math:`\sigma_{x_i}`.
+
+      * ``<species_name>.twiss.u0`` (`float`) Mean normalized particle momentum :math:`u_0 = \beta_0 \gamma_0 > 0`.
+
+      * ``<species_name>.twiss.nz`` (list of 3 `floats`, default: `0 0 1`) Longitudinal beam direction :math:`\mathbf{\hat{n}}_z`.
+
+      * ``<species_name>.twiss.nx`` (list of 3 `floats`, default: `1 0 0`) Transverse beam direction :math:`\mathbf{\hat{n}}_x`.  In the code, this is projected into the plane perpendicular to :math:`\mathbf{\hat{n}}_z`, so that we only require :math:`\mathbf{\hat{n}}_z` and :math:`\mathbf{\hat{n}}_x` not be parallel.  Both unit vectors are then normalized and used to calculate :math:`\mathbf{\hat{n}}_y = \mathbf{\hat{n}}_z \times \mathbf{\hat{n}}_x`.
+
+      * ``<species_name>.twiss.euler`` (list of 3 `floats`, optional) If given, provides an alternative way to specify :math:`\mathbf{\hat{n}}_z` and :math:`\mathbf{\hat{n}}_x` in terms of an intrinsic :math:`Z\text{-}X\text{-}Z` Euler rotation.  Parameters are:
+
+        * :math:`\mathbf{\alpha}` = right-handed rotation about :math:`\mathbf{\hat{z}}`
+
+        * :math:`\mathbf{\beta}` = right-handed rotation about :math:`\mathbf{\hat{x}}'`
+
+        * :math:`\mathbf{\gamma}` = right-handed rotation about :math:`\mathbf{\hat{z}}''`
+
+      * ``<species_name>.twiss.symmetrization_order`` (`int`, default=1) Symmetrization of 4D uncorrelated transverse phase-space.  Allowed values:
+
+        * ``1`` Include only the identity transformation:
+
+        .. math::
+             (x,u_x;y,u_y) \mapsto (x,u_x;y,u_y)
+
+        * ``8`` In addition to the identity transformation, include all remaining even-parity symmetry transformations:
+
+          .. math::
+             \begin{align*}
+             (x,u_x;y,u_y) \mapsto
+               & (-x,-u_x;y,u_y), (-x,u_x;-y,u_y), (-x,u_x;y,-u_y), \\
+               &  (x,-u_x;-y,u_y), (x,-u_x;y,-u_y), (x,u_x;-y,-u_y) \\
+               &  (-x,-u_x;-y,-u_y)
+             \end{align*}
+
+        * ``16`` In addition to all the even-parity symmetry transformations, include all the odd-parity symmetry transformations:
+
+          .. math::
+             \begin{align*}
+             (x,u_x;y,u_y) \mapsto
+             & (-x,u_x;y,u_y), (x,-u_x;y,u_y), (x,u_x;-y,u_y), \\
+             & (x,u_x;y,-u_y), (x,-u_x;-y,-u_y), (-x,u_x;-y,-u_y) \\
+             & (-x,-u_x;y,-u_y), (-x,-u_x;y,u_y)
+             \end{align*}
+
+      Each of the three bivariate distributions is separately specified by three of the following parameters:
+
+      * ``<species_name>.twiss.focal_distance_x/y/zeta`` (`float`, meter) Focal distance :math:`L_i`.  Can be positive, negative, or zero.  A negative value corresponds to the focus occurring prior to the initialization time.
+
+      * ``<species_name>.twiss.sigma_x/y/zeta`` (`float`, meter) RMS beam spread :math:`\sigma^\star_{x_i}` at focus.
+
+      * ``<species_name>.twiss.sigma_ux/uy/uzeta`` (`float`) RMS normalized momentum spread :math:`\sigma^\star_{u_i}` at focus, in units of :math:`u_0`.
+
+      * ``<species_name>.twiss.emittance_x/y/zeta`` (`float`, radian meter) RMS normalized emittance :math:`\epsilon_i`.
+
+      * ``<species_name>.twiss.alpha_x/y/zeta`` (`float`, dimensionless) Twiss parameter :math:`\alpha_i`.
+
+      * ``<species_name>.twiss.beta_x/y/zeta`` (`float`, meter) Twiss parameter :math:`\beta_i`.
+
+      For each Cartesian direction, any set of three parameters may be used as long as it can be used to determine the primary set of initialization parameters :math:`\{L_i, \sigma^\star_{x_i}, \sigma^\star_{u_i}\}` using the relations
+
+      .. math::
+         \beta_i \gamma_i &= 1 + \alpha_i^2
+
+         \alpha_i &= L_i \gamma_i
+
+         \sigma^\star_{x_i} \sigma^\star_{u_i} &= \epsilon_i
+
+         \gamma_i \sigma^{\star 2}_{x_i} &= \epsilon_i
+
+      In addition to the primary set, common choices of parameters include :math:`\{L_i, \sigma^\star_{x_i}, \epsilon_i\}` and :math:`\{\alpha_i, \beta_i, \epsilon_i\}`.
+
+      To use ``injection_style = twiss``, you must also set ``momentum_distribution_type = twiss``.
+
+    ..
+         gaussian distribution in space in all directions. This requires additional parameters:
+
+
+      * ``external_file``: Inject macroparticles with properties (mass, charge, position, and momentum - :math:`\gamma \beta m c`) read from an external openPMD file.
       With it users can specify the additional arguments:
 
       * ``<species_name>.injection_file`` (`string`) openPMD file name and
@@ -1105,6 +1231,8 @@ Particle initialization
       ``<species_name>.uy_th`` and ``<species_name>.uz_th``.
       These 6 parameters are all ``0.`` by default.
 
+    * ``twiss``: Only valid for ``injection_style = twiss``, from which the momentum parameters are calculated.  No further parameters are specified here.
+
     * ``gaussianflux``: Gaussian momentum flux distribution, which is Gaussian in the plane and v*Gaussian normal to the plane.
       It can only be used when ``injection_style = NFluxPerCell``.
       This can be controlled with the additional arguments to specify the plane's orientation, ``<species_name>.flux_normal_axis`` and

Source/Initialization/InjectorMomentum.H

@@ -91,6 +91,43 @@ private:
     amrex::Real m_ux_th, m_uy_th, m_uz_th;
 };
 
+struct InjectorMomentumTwiss
+{
+    InjectorMomentumTwiss (amrex::Real a_u0, amrex::XDim3 a_sigma_u) noexcept
+        : m_u0(a_u0), m_sigma_u(a_sigma_u)
+    {
+    }
+
+    [[nodiscard]]
+    AMREX_GPU_HOST_DEVICE
+    amrex::XDim3
+    getMomentum (
+        amrex::Real /*x*/, amrex::Real /*y*/, amrex::Real /*z*/,
+        amrex::RandomEngine const& engine) const noexcept
+    {
+        using namespace amrex::literals;
+
+        return amrex::XDim3 {
+            amrex::RandomNormal(0_rt, m_sigma_u.x, engine),
+            amrex::RandomNormal(0_rt, m_sigma_u.y, engine),
+            m_u0 + amrex::RandomNormal(0_rt, m_sigma_u.z, engine)
+        };
+    }
+
+    [[nodiscard]]
+    AMREX_GPU_HOST_DEVICE
+    amrex::XDim3
+    getBulkMomentum (amrex::Real /*x*/, amrex::Real /*y*/, amrex::Real /*z*/) const noexcept
+    {
+        using namespace amrex::literals;
+        return amrex::XDim3{ 0_rt, 0_rt, m_u0 };
+    }
+
+private:
+    const amrex::Real m_u0;
+    const amrex::XDim3 m_sigma_u;
+};
+
 // struct whose getMomentum returns momentum for 1 particle, from random
 // gaussian flux distribution in the specified direction.
 // Along the normal axis, the distribution is v*Gaussian,
@@ -540,6 +577,13 @@ struct InjectorMomentum
           object(t,a_ux_m,a_uy_m,a_uz_m,a_ux_th,a_uy_th,a_uz_th)
     { }
 
+    // This constructor stores a InjectorMomentumGaussian in union object.
+    InjectorMomentum (
+        InjectorMomentumTwiss* t, amrex::Real a_u0, amrex::XDim3 a_sigma_u)
+        : type(Type::twiss),
+          object(t, a_u0, a_sigma_u)
+    { }
+
     // This constructor stores a InjectorMomentumGaussianFlux in union object.
     InjectorMomentum (InjectorMomentumGaussianFlux* t,
                       amrex::Real a_ux_m, amrex::Real a_uy_m, amrex::Real a_uz_m,
@@ -607,6 +651,10 @@ struct InjectorMomentum
         {
             return object.gaussian.getMomentum(x,y,z,engine);
         }
+        case Type::twiss:
+        {
+            return object.twiss.getMomentum(x,y,z,engine);
+        }
         case Type::gaussianparser:
         {
             return object.gaussianparser.getMomentum(x,y,z,engine);
@@ -660,6 +708,10 @@ struct InjectorMomentum
         {
             return object.gaussian.getBulkMomentum(x,y,z);
         }
+        case Type::twiss:
+        {
+            return object.twiss.getBulkMomentum(x,y,z);
+        }
         case Type::gaussianparser:
         {
             return object.gaussianparser.getBulkMomentum(x,y,z);
@@ -696,7 +748,7 @@ struct InjectorMomentum
         }
     }
 
-    enum struct Type { constant, gaussian, gaussianflux, uniform, boltzmann, juttner, radial_expansion, parser, gaussianparser };
+    enum struct Type { constant, gaussian, twiss, gaussianflux, uniform, boltzmann, juttner, radial_expansion, parser, gaussianparser };
     Type type;
 
 private:
@@ -713,6 +765,9 @@ private:
                 amrex::Real a_uz_m, amrex::Real a_ux_th,
                 amrex::Real a_uy_th, amrex::Real a_uz_th) noexcept
             : gaussian(a_ux_m,a_uy_m,a_uz_m,a_ux_th,a_uy_th,a_uz_th) {}
+        Object (InjectorMomentumTwiss*,
+                amrex::Real a_u0, amrex::XDim3 a_sigma_u) noexcept
+            : twiss(a_u0, a_sigma_u) {}
         Object (InjectorMomentumGaussianFlux*,
                 amrex::Real a_ux_m, amrex::Real a_uy_m,
                 amrex::Real a_uz_m, amrex::Real a_ux_th,
@@ -749,6 +804,7 @@ private:
                              a_ux_th_parser, a_uy_th_parser, a_uz_th_parser) {}
         InjectorMomentumConstant constant;
         InjectorMomentumGaussian gaussian;
+        InjectorMomentumTwiss twiss;
         InjectorMomentumGaussianFlux gaussianflux;
         InjectorMomentumUniform uniform;
         InjectorMomentumBoltzmann boltzmann;

Source/Initialization/InjectorMomentum.cpp

@@ -15,6 +15,7 @@ void InjectorMomentum::clear () // NOLINT(readability-make-member-function-const
     {
     case Type::parser:
     case Type::gaussian:
+    case Type::twiss:
     case Type::gaussianparser:
     case Type::gaussianflux:
     case Type::uniform:

Source/Initialization/PlasmaInjector.H

@@ -115,6 +115,27 @@ public:
     bool do_focusing = false;
     amrex::Real focal_distance;
 
+    bool twiss = false;
+    amrex::XDim3 x0;
+    amrex::Real twiss_u0;
+    amrex::XDim3 twiss_nx;
+    amrex::XDim3 twiss_ny;
+    amrex::XDim3 twiss_nz;
+    amrex::XDim3 twiss_focal_distance;
+    amrex::XDim3 twiss_sigma_x;
+    amrex::XDim3 twiss_sigma_u;
+    amrex::XDim3 twiss_planar_cut = {
+        std::numeric_limits<amrex::Real>::max(),
+        std::numeric_limits<amrex::Real>::max(),
+        std::numeric_limits<amrex::Real>::max()
+    };
+    amrex::XDim3 twiss_ellipsoidal_cut = {
+        std::numeric_limits<amrex::Real>::max(),
+        std::numeric_limits<amrex::Real>::max(),
+        std::numeric_limits<amrex::Real>::max()
+    };
+    int twiss_symmetrization_order = 1;
+
     bool external_file = false; //! initialize from an openPMD file
     amrex::Real z_shift = 0.0; //! additional z offset for particle positions
 #ifdef WARPX_USE_OPENPMD
@@ -196,12 +217,16 @@ protected:
     void setupSingleParticle (amrex::ParmParse const& pp_species);
     void setupMultipleParticles (amrex::ParmParse const& pp_species);
     void setupGaussianBeam (amrex::ParmParse const& pp_species);
+    void setupTwiss (amrex::ParmParse const& pp_species);
     void setupNRandomPerCell (amrex::ParmParse const& pp_species);
     void setupNFluxPerCell (amrex::ParmParse const& pp_species);
     void setupNuniformPerCell (amrex::ParmParse const& pp_species);
     void setupExternalFile (amrex::ParmParse const& pp_species);
 
     void parseFlux (amrex::ParmParse const& pp_species);
+    void parseTwissParameters (
+        amrex::ParmParse const& pp_species, const std::string& dir,
+        amrex::Real& _focal_distance, amrex::Real& sigma_x, amrex::Real& sigma_u);
 };
 
 #endif //WARPX_PLASMA_INJECTOR_H_