diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst index 31f0e06ab5b..8c110590a54 100644 --- a/Docs/source/usage/parameters.rst +++ b/Docs/source/usage/parameters.rst @@ -840,6 +840,7 @@ mu0 vacuum permeability clight speed of light kb Boltzmann's constant (J/K) pi math constant pi +inf maximum finite value of type ``amrex::Real`` ======== =================== See ``Source/Utils/WarpXConst.H`` for the values. @@ -1022,8 +1023,146 @@ 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. + + * ``.q_tot`` (`float`, coulomb) Beam charge :math:`Q_\text{tot} = N q`. + + * ``.npart`` (`int`) Number of macroparticles. + + * ``.x0/y0/z0`` (`float`, meter) Initial beam position :math:`\mathbf{x}_0`. + + * ``.twiss.ellipsoidal_cut`` (list of 6 `floats`, optional) Cut particles with + + .. math:: + \Biggl[ \left( + \frac{x}{a_1} \right)^2 + + \left( \frac{y}{a_2} \right)^2 + + \left( \frac{\zeta}{a_3} \right)^2 + + \left( \frac{u_x}{a_4} \right)^2 + + \left( \frac{u_y}{a_5} \right)^2 + + \left( \frac{u_\zeta}{a_6} \right)^2 \Biggr] > 1 + + Parameters: + + * :math:`(a_1, a_2, a_3, a_4, a_5, a_6)` in units of :math:`(\sigma^\star_x, \sigma^\star_y, \sigma^\star_\zeta, \sigma^\star_{u_x}, \sigma^\star_{u_y}, \sigma^\star_{u_\zeta})` + + The charge of the uncut beam is ``.q_tot``, so that cutting the distribution generally results in lower total injected charge. Specifying ``inf`` for any parameter disables the corresponding cut. + + * ``.twiss.planar_cut`` (list of 6 `floats`, optional) Cut particles with + + .. math:: + \max \Biggl ( \frac{|x|}{a_1}, \frac{|y|}{a_2}, \frac{|\zeta|}{a_3}, \frac{|u_x|}{a_4}, \frac{|u_y|}{a_5}, \frac{|u_\zeta|}{a_6} \Biggr) > 1 + + Parameters: + + * :math:`(a_1, a_2, a_3, a_4, a_5, a_6)` in units of :math:`(\sigma^\star_x, \sigma^\star_y, \sigma^\star_\zeta, \sigma^\star_{u_x}, \sigma^\star_{u_y}, \sigma^\star_{u_\zeta})` + + The charge of the uncut beam is ``.q_tot``, so that cutting the distribution generally results in lower total injected charge. Specifying ``inf`` for any parameter disables the corresponding cut. + + * ``.twiss.u0`` (`float`) Mean normalized particle momentum :math:`u_0 = \beta_0 \gamma_0 > 0`. + + * ``.twiss.nz`` (list of 3 `floats`, default: `0 0 1`) Longitudinal beam direction :math:`\mathbf{\hat{n}}_z`. + + * ``.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`. + + * ``.twiss.euler`` (list of 3 `floats`, radian, 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}}''` + + * ``.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: + + * ``.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. + + * ``.twiss.sigma_x/y/zeta`` (`float`, meter) RMS beam spread :math:`\sigma^\star_{x_i}` at focus. + + * ``.twiss.sigma_ux/uy/uzeta`` (`float`) RMS normalized momentum spread :math:`\sigma^\star_{u_i}` at focus, in units of :math:`u_0`. + + * ``.twiss.emittance_x/y/zeta`` (`float`, radian meter) RMS normalized emittance :math:`\epsilon_i`. + + * ``.twiss.alpha_x/y/zeta`` (`float`, dimensionless) Twiss parameter :math:`\alpha_i`. + + * ``.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: * ``.injection_file`` (`string`) openPMD file name and @@ -1176,6 +1315,8 @@ Particle initialization ``.uy_th`` and ``.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, ``.flux_normal_axis`` and diff --git a/Source/Initialization/InjectorMomentum.H b/Source/Initialization/InjectorMomentum.H index 52fedcd779f..aef711a510a 100644 --- a/Source/Initialization/InjectorMomentum.H +++ b/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; diff --git a/Source/Initialization/InjectorMomentum.cpp b/Source/Initialization/InjectorMomentum.cpp index 4642d3a0cc6..bbbeb8344f1 100644 --- a/Source/Initialization/InjectorMomentum.cpp +++ b/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: diff --git a/Source/Initialization/PlasmaInjector.H b/Source/Initialization/PlasmaInjector.H index f14720d271c..f8e6b4bdf6b 100644 --- a/Source/Initialization/PlasmaInjector.H +++ b/Source/Initialization/PlasmaInjector.H @@ -115,6 +115,19 @@ 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::Vector twiss_planar_cut; + amrex::Vector twiss_ellipsoidal_cut; + 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 +209,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_ diff --git a/Source/Initialization/PlasmaInjector.cpp b/Source/Initialization/PlasmaInjector.cpp index 468d9e7e336..110975c348e 100644 --- a/Source/Initialization/PlasmaInjector.cpp +++ b/Source/Initialization/PlasmaInjector.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -37,6 +38,7 @@ #include #include +#include #include #include #include @@ -46,6 +48,45 @@ using namespace amrex::literals; +namespace +{ + inline amrex::XDim3 CrossProduct (const amrex::XDim3& a, const amrex::XDim3& b) + { + return { a.y*b.z - a.z*b.y, a.z*b.x - a.x*b.z, a.x*b.y - a.y*b.x }; + } + + inline amrex::Real DotProduct (const amrex::XDim3& a, const amrex::XDim3& b) + { + return a.x*b.x + a.y*b.y + a.z*b.z; + } + + inline amrex::Real Norm (const amrex::XDim3& a) + { + return std::sqrt(a.x*a.x + a.y*a.y + a.z*a.z); + } + + inline amrex::XDim3 Normalize (const amrex::XDim3& a) + { + const amrex::Real norm = std::sqrt(a.x*a.x + a.y*a.y + a.z*a.z); + return { a.x/norm, a.y/norm, a.z/norm }; + } + + // Projection of "a" in the plane with normal "n" + inline amrex::XDim3 Projection (const amrex::XDim3& a, const amrex::XDim3& n) + { + const amrex::Real a_dot_n = DotProduct(a, n); + return { + a.x - a_dot_n * n.x, + a.y - a_dot_n * n.y, + a.z - a_dot_n * n.z }; + } + + inline amrex::Real Square (const amrex::Real x) + { + return amrex::Math::powi<2>(x); + } +} + PlasmaInjector::PlasmaInjector (int ispecies, const std::string& name, const amrex::Geometry& geom, const std::string& src_name): species_id{ispecies}, species_name{name}, source_name{src_name} @@ -131,6 +172,8 @@ PlasmaInjector::PlasmaInjector (int ispecies, const std::string& name, return; } else if (injection_style == "gaussian_beam") { setupGaussianBeam(pp_species); + } else if (injection_style == "twiss") { + setupTwiss(pp_species); } else if (injection_style == "nrandompercell") { setupNRandomPerCell(pp_species); } else if (injection_style == "nfluxpercell") { @@ -261,6 +304,113 @@ void PlasmaInjector::setupGaussianBeam (amrex::ParmParse const& pp_species) #endif } +void PlasmaInjector::setupTwiss (amrex::ParmParse const& pp_species) +{ + utils::parser::getWithParser(pp_species, source_name, "x0", x0.x); + utils::parser::getWithParser(pp_species, source_name, "y0", x0.y); + utils::parser::getWithParser(pp_species, source_name, "z0", x0.z); + utils::parser::getWithParser(pp_species, source_name, "q_tot", q_tot); + utils::parser::getWithParser(pp_species, source_name, "npart", npart); + + if (pp_species.contains("twiss.planar_cut")) { + utils::parser::getArrWithParser( + pp_species, source_name, "twiss.planar_cut", twiss_planar_cut, 0, 6); + } else { + twiss_planar_cut.resize(6, std::numeric_limits::max()); + } + + if (pp_species.contains("twiss.ellipsoidal_cut")) { + utils::parser::getArrWithParser( + pp_species, source_name, "twiss.ellipsoidal_cut", twiss_ellipsoidal_cut, 0, 6); + } else { + twiss_ellipsoidal_cut.resize(6, std::numeric_limits::max()); + } + + utils::parser::getWithParser(pp_species, source_name, "twiss.u0", twiss_u0); + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + twiss_u0 > 0_rt, "twiss.u0 must be positive"); + + amrex::Vector temp; + if (pp_species.contains("twiss.euler")) { + utils::parser::getArrWithParser(pp_species, source_name, "twiss.euler", temp, 0, 3); + const amrex::Real ca = std::cos(temp[0]), sa = std::sin(temp[0]); + const amrex::Real cb = std::cos(temp[1]), sb = std::sin(temp[1]); + const amrex::Real cg = std::cos(temp[2]), sg = std::sin(temp[2]); + twiss_nx = { + ca*cg - cb*sa*sg, + cg*sa + ca*cb*sg, + sb*sg + }; + twiss_ny = { + -ca*sg - cb*cg*sa, + ca*cb*cg - sa*sg, + cg*sb + }; + twiss_nz = { + sa*sb, + -ca*sb, + cb + }; + } else { + if (pp_species.contains("twiss.nz")) { + utils::parser::getArrWithParser(pp_species, source_name, "twiss.nz", temp, 0, 3); + twiss_nz = Normalize({temp[0], temp[1], temp[2]}); + } else { + twiss_nz = { 0_rt, 0_rt, 1_rt }; + } + if (pp_species.contains("twiss.nx")) { + utils::parser::getArrWithParser(pp_species, source_name, "twiss.nx", temp, 0, 3); + twiss_nx = Normalize({temp[0], temp[1], temp[2]}); + } else { + twiss_nx = { 1_rt, 0_rt, 0_rt }; + } + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + Norm(CrossProduct(twiss_nz, twiss_nx)) > 0_rt, + "twiss.nz and twiss.nx must not be parallel"); + twiss_nx = Normalize(Projection(twiss_nx, twiss_nz)); + twiss_ny = CrossProduct(twiss_nz, twiss_nx); + } + + parseTwissParameters( + pp_species, "x", twiss_focal_distance.x, twiss_sigma_x.x, twiss_sigma_u.x); + parseTwissParameters( + pp_species, "y", twiss_focal_distance.y, twiss_sigma_x.y, twiss_sigma_u.y); + parseTwissParameters( + pp_species, "zeta", twiss_focal_distance.z, twiss_sigma_x.z, twiss_sigma_u.z); + +#if defined(WARPX_DIM_XZ) + WARPX_ALWAYS_ASSERT_WITH_MESSAGE(twiss_sigma_x.y > 0_rt, + "Error: twiss.sigma_y must be strictly greater than 0 in 2D " + "(it is used when computing the particles' weights from the total beam charge)"); +#elif defined(WARPX_DIM_1D_Z) + WARPX_ALWAYS_ASSERT_WITH_MESSAGE(twiss_sigma_x.x> 0_rt, + "Error: twiss.sigma_x must be strictly greater than 0 in 1D " + "(it is used when computing the particles' weights from the total beam charge)"); + WARPX_ALWAYS_ASSERT_WITH_MESSAGE(twiss_sigma_x.y > 0_rt, + "Error: twiss.sigma_y must be strictly greater than 0 in 1D " + "(it is used when computing the particles' weights from the total beam charge)"); +#endif + + utils::parser::queryWithParser( + pp_species, source_name, "twiss.symmetrization_order", twiss_symmetrization_order); + const std::set valid_symmetries = { 1, 8, 16 }; + WARPX_ALWAYS_ASSERT_WITH_MESSAGE(valid_symmetries.count(twiss_symmetrization_order), + "Error: Symmetrization only supported for orders 1, 8, or 16 "); + + twiss = true; + + // Momentum parameters have now been calculated. Append them to the amrex::ParmParse for use in SpeciesUtils::parseMomentum(). + amrex::ParmParse pp_mod(species_name); + pp_mod.add(std::string("twiss.sigma_ux").c_str(), twiss_sigma_u.x); + pp_mod.add(std::string("twiss.sigma_uy").c_str(), twiss_sigma_u.y); + pp_mod.add(std::string("twiss.sigma_uzeta").c_str(), twiss_sigma_u.z); + SpeciesUtils::parseMomentum( + species_name, source_name, "twiss", h_inj_mom, + ux_parser, uy_parser, uz_parser, + ux_th_parser, uy_th_parser, uz_th_parser, + h_mom_temp, h_mom_vel); +} + void PlasmaInjector::setupNRandomPerCell (amrex::ParmParse const& pp_species) { utils::parser::getWithParser(pp_species, source_name, "num_particles_per_cell", num_particles_per_cell); @@ -577,6 +727,98 @@ void PlasmaInjector::parseFlux (amrex::ParmParse const& pp_species) } +void PlasmaInjector::parseTwissParameters ( + amrex::ParmParse const& pp_species, const std::string& dir, + amrex::Real& _focal_distance, amrex::Real& sigma_x, amrex::Real& sigma_u) +{ + enum class TwissParameter { + FOCAL_DISTANCE, SIGMA_X, SIGMA_U, EMITTANCE, ALPHA, BETA, GAMMA + }; + using TP = TwissParameter; + + const std::map lookup = { + {TP::FOCAL_DISTANCE, std::string("twiss.focal_distance_") + dir}, + {TP::SIGMA_X, std::string("twiss.sigma_") + dir}, + {TP::SIGMA_U, std::string("twiss.sigma_u") + dir}, + {TP::EMITTANCE, std::string("twiss.emittance_") + dir}, + {TP::ALPHA, std::string("twiss.alpha_") + dir}, + {TP::BETA, std::string("twiss.beta_") + dir} + }; + + std::map vars; + for (const auto& pair : lookup) { + amrex::Real value; + if (utils::parser::queryWithParser(pp_species, source_name, pair.second.c_str(), value)) { + vars[pair.first] = value; + } + } + + // sigma_u specified in units of u0 + if (vars.count(TP::SIGMA_U)) { + vars[TP::SIGMA_U] *= twiss_u0; + } + + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + vars.size() == 3, + std::string("Must provide exactly 3 Twiss parameters (dir=") + dir + ")"); + + // dispersion of transverse velocity with respect to transverse momentum + amrex::Real eta = 1_rt; + if (dir == "zeta") { + // dispersion of longitudinal velocity with respect to longitudinal momentum + const amrex::Real gamma0 = std::sqrt(1_rt + Square(twiss_u0)); + eta = 1_rt/Square(gamma0); + } + + while ( + vars.count(TP::FOCAL_DISTANCE) + vars.count(TP::SIGMA_X) + vars.count(TP::SIGMA_U) < 3) { + if ( // beta gamma = 1 + alpha^2 + vars.count(TP::BETA) && vars.count(TP::ALPHA) && !vars.count(TP::GAMMA)) { + vars[TP::GAMMA] = (1_rt + Square(vars[TP::ALPHA]))/vars[TP::BETA]; + } else if ( + vars.count(TP::GAMMA) && vars.count(TP::BETA) && !vars.count(TP::ALPHA)) { + vars[TP::ALPHA] = std::sqrt(vars[TP::BETA]*vars[TP::GAMMA] - 1_rt); + } else if ( + vars.count(TP::ALPHA) && vars.count(TP::GAMMA) && !vars.count(TP::BETA)) { + vars[TP::BETA] = (1_rt+Square(vars[TP::ALPHA]))/vars[TP::GAMMA]; + } else if ( // alpha = eta L gamma / u0 + vars.count(TP::ALPHA) && vars.count(TP::FOCAL_DISTANCE) && !vars.count(TP::GAMMA)) { + vars[TP::GAMMA] = twiss_u0 * vars[TP::ALPHA] / (eta * vars[TP::FOCAL_DISTANCE]); + } else if ( + vars.count(TP::GAMMA) && vars.count(TP::ALPHA) && !vars.count(TP::FOCAL_DISTANCE)) { + vars[TP::FOCAL_DISTANCE] = twiss_u0 * vars[TP::ALPHA] / (eta * vars[TP::GAMMA]); + } else if ( + vars.count(TP::FOCAL_DISTANCE) && vars.count(TP::GAMMA) && !vars.count(TP::ALPHA)) { + vars[TP::ALPHA] = eta * vars[TP::FOCAL_DISTANCE] * vars[TP::GAMMA] / twiss_u0; + } else if ( // sigma_x sigma_u = emittance + vars.count(TP::SIGMA_X) && vars.count(TP::SIGMA_U) && !vars.count(TP::EMITTANCE)) { + vars[TP::EMITTANCE] = vars[TP::SIGMA_X]*vars[TP::SIGMA_U]; + } else if ( + vars.count(TP::EMITTANCE) && vars.count(TP::SIGMA_X) && !vars.count(TP::SIGMA_U)) { + vars[TP::SIGMA_U] = vars[TP::EMITTANCE]/vars[TP::SIGMA_X]; + } else if ( + vars.count(TP::SIGMA_U) && vars.count(TP::EMITTANCE) && !vars.count(TP::SIGMA_X)) { + vars[TP::SIGMA_X] = vars[TP::EMITTANCE]/vars[TP::SIGMA_U]; + } else if ( // gamma sigma_x^2 = emittance + vars.count(TP::GAMMA) && vars.count(TP::SIGMA_X) && !vars.count(TP::EMITTANCE)) { + vars[TP::EMITTANCE] = vars[TP::GAMMA]*Square(vars[TP::SIGMA_X]); + } else if ( + vars.count(TP::EMITTANCE) && vars.count(TP::GAMMA) && !vars.count(TP::SIGMA_X)) { + vars[TP::SIGMA_X] = std::sqrt(vars[TP::EMITTANCE]/vars[TP::GAMMA]); + } else if ( + vars.count(TP::SIGMA_X) && vars.count(TP::EMITTANCE) && !vars.count(TP::GAMMA)) { + vars[TP::GAMMA] = vars[TP::EMITTANCE]/Square(vars[TP::SIGMA_X]); + } else { + WARPX_ABORT_WITH_MESSAGE( + std::string("Invalid set of Twiss parameters (dir=") + dir + ")"); + } + } + + _focal_distance = vars[TP::FOCAL_DISTANCE]; + sigma_x = vars[TP::SIGMA_X]; + sigma_u = vars[TP::SIGMA_U]; +} + amrex::XDim3 PlasmaInjector::getMomentum (amrex::Real x, amrex::Real y, amrex::Real z) const noexcept diff --git a/Source/Particles/PhysicalParticleContainer.H b/Source/Particles/PhysicalParticleContainer.H index 18880239183..0be7fcbe34a 100644 --- a/Source/Particles/PhysicalParticleContainer.H +++ b/Source/Particles/PhysicalParticleContainer.H @@ -192,6 +192,8 @@ public: void AddGaussianBeam (PlasmaInjector const& plasma_injector); + void AddTwiss (PlasmaInjector const& plasma_injector); + /** Load a particle beam from an external file * @param[in] the PlasmaInjector instance holding the input parameters * @param[in] q_tot total charge of the particle species to be initialized diff --git a/Source/Particles/PhysicalParticleContainer.cpp b/Source/Particles/PhysicalParticleContainer.cpp index c973e9afafa..4a2010144d9 100644 --- a/Source/Particles/PhysicalParticleContainer.cpp +++ b/Source/Particles/PhysicalParticleContainer.cpp @@ -211,6 +211,11 @@ namespace idcpu[ip] = amrex::ParticleIdCpus::Invalid; } + + inline amrex::Real Square (const amrex::Real x) + { + return amrex::Math::powi<2>(x); + } } PhysicalParticleContainer::PhysicalParticleContainer (AmrCore* amr_core, int ispecies, @@ -676,6 +681,177 @@ PhysicalParticleContainer::AddGaussianBeam (PlasmaInjector const& plasma_injecto 1, attr, 0, attr_int, 1); } +void PhysicalParticleContainer::AddTwiss (PlasmaInjector const& plasma_injector) +{ + const amrex::XDim3 x0 = plasma_injector.x0; + const amrex::Real q_tot = plasma_injector.q_tot; + const long npart = plasma_injector.npart; + const int symmetrization_order = plasma_injector.twiss_symmetrization_order; + const amrex::Real u0 = plasma_injector.twiss_u0; + const amrex::XDim3 nx = plasma_injector.twiss_nx; + const amrex::XDim3 ny = plasma_injector.twiss_ny; + const amrex::XDim3 nz = plasma_injector.twiss_nz; + const amrex::XDim3 focal_distance = plasma_injector.twiss_focal_distance; + const amrex::XDim3 sigma_x = plasma_injector.twiss_sigma_x; + const amrex::XDim3 sigma_u = plasma_injector.twiss_sigma_u; + const amrex::Real v0 = (u0 / std::sqrt(1_rt + u0*u0)) * PhysConst::c; + const amrex::Vector pcut = { + plasma_injector.twiss_planar_cut[0] * sigma_x.x, + plasma_injector.twiss_planar_cut[1] * sigma_x.y, + plasma_injector.twiss_planar_cut[2] * sigma_x.z, + plasma_injector.twiss_planar_cut[3] * sigma_u.x, + plasma_injector.twiss_planar_cut[4] * sigma_u.y, + plasma_injector.twiss_planar_cut[5] * sigma_u.z + }; + const amrex::Vector ecut = { + plasma_injector.twiss_ellipsoidal_cut[0] * sigma_x.x, + plasma_injector.twiss_ellipsoidal_cut[1] * sigma_x.y, + plasma_injector.twiss_ellipsoidal_cut[2] * sigma_x.z, + plasma_injector.twiss_ellipsoidal_cut[3] * sigma_u.x, + plasma_injector.twiss_ellipsoidal_cut[4] * sigma_u.y, + plasma_injector.twiss_ellipsoidal_cut[5] * sigma_u.z + }; + const amrex::XDim3 delta_t { + focal_distance.x / v0, + focal_distance.y / v0, + focal_distance.z / v0, + }; + + // Declare temporary vectors on the CPU + Gpu::HostVector particle_x; + Gpu::HostVector particle_y; + Gpu::HostVector particle_z; + Gpu::HostVector particle_ux; + Gpu::HostVector particle_uy; + Gpu::HostVector particle_uz; + Gpu::HostVector particle_w; + + if (ParallelDescriptor::IOProcessor()) { + for (long i=0; i pcut[0] || std::abs(y) > pcut[1] || std::abs(z) > pcut[2] || + std::abs(u.x) > pcut[3] || std::abs(u.y) > pcut[4] || std::abs(u.z) > pcut[5]) { + continue; + } + + if (Square(x/ecut[0]) + Square(y/ecut[1]) + Square(z/ecut[2]) + + Square(u.x/ecut[3]) + Square(u.y/ecut[4]) + Square((u.z-u0)/ecut[5]) > 1_rt) { + continue; + } + + const Real gamma = std::sqrt(1_rt + (u.x*u.x + u.y*u.y + u.z*u.z)); + + auto transform_and_push = [&]( + amrex::Real _x, amrex::Real _y, amrex::Real _z, + amrex::Real ux, amrex::Real uy, amrex::Real uz) { + + const amrex::XDim3 v { + (ux / gamma) * PhysConst::c, + (uy / gamma) * PhysConst::c, + (uz / gamma) * PhysConst::c + }; + +#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_RZ) + _x -= delta_t.x * v.x; + _y -= delta_t.y * v.y; + _z -= delta_t.z * (v.z - v0); +#elif defined(WARPX_DIM_XZ) + _x -= delta_t.x * v.x; + _z -= delta_t.z * (v.z - v0); +#elif defined(WARPX_DIM_1D_Z) + _z -= delta_t.z * (v.z - v0); +#endif + const amrex::XDim3 xbar { + (_x*nx.x + _y*ny.x + _z*nz.x) + x0.x, + (_x*nx.y + _y*ny.y + _z*nz.y) + x0.y, + (_x*nx.z + _y*ny.z + _z*nz.z) + x0.z + }; + + if (!(plasma_injector.insideBounds(xbar.x, xbar.y, xbar.z))) { + return; + } + + const amrex::XDim3 ubar { + ux*nx.x + uy*ny.x + uz*nz.x, + ux*nx.y + uy*ny.y + uz*nz.y, + ux*nx.z + uy*ny.z + uz*nz.z + }; + + CheckAndAddParticle( + xbar.x, xbar.y, xbar.z, + ubar.x*PhysConst::c, ubar.y*PhysConst::c, ubar.z*PhysConst::c, weight, + particle_x, particle_y, particle_z, + particle_ux, particle_uy, particle_uz, + particle_w); + }; + + // Discrete symmetries of transverse phase-space + switch (symmetrization_order) { + case 16: + // One-axis reflections + transform_and_push(-x, y, z, u.x, u.y, u.z); + transform_and_push( x,-y, z, u.x, u.y, u.z); + transform_and_push( x, y, z,-u.x, u.y, u.z); + transform_and_push( x, y, z, u.x,-u.y, u.z); + // Three-axis reflections + transform_and_push( x,-y, z,-u.x,-u.y, u.z); + transform_and_push(-x, y, z,-u.x,-u.y, u.z); + transform_and_push(-x,-y, z, u.x,-u.y, u.z); + transform_and_push(-x,-y, z,-u.x, u.y, u.z); + [[fallthrough]]; + case 8: + // Two-axis reflections + transform_and_push(-x,-y, z, u.x, u.y, u.z); + transform_and_push(-x, y, z,-u.x, u.y, u.z); + transform_and_push(-x, y, z, u.x,-u.y, u.z); + transform_and_push( x,-y, z,-u.x, u.y, u.z); + transform_and_push( x,-y, z, u.x,-u.y, u.z); + transform_and_push( x, y, z,-u.x,-u.y, u.z); + // Four-axis reflection + transform_and_push(-x,-y, z,-u.x,-u.y, u.z); + [[fallthrough]]; + case 1: + // Identity + transform_and_push( x, y, z, u.x, u.y, u.z); + } + } + } + // Add the temporary CPU vectors to the particle structure + auto const np = static_cast(particle_z.size()); + + const amrex::Vector xp(particle_x.data(), particle_x.data() + np); + const amrex::Vector yp(particle_y.data(), particle_y.data() + np); + const amrex::Vector zp(particle_z.data(), particle_z.data() + np); + const amrex::Vector uxp(particle_ux.data(), particle_ux.data() + np); + const amrex::Vector uyp(particle_uy.data(), particle_uy.data() + np); + const amrex::Vector uzp(particle_uz.data(), particle_uz.data() + np); + + amrex::Vector> attr; + const amrex::Vector wp(particle_w.data(), particle_w.data() + np); + attr.push_back(wp); + + const amrex::Vector> attr_int; + + AddNParticles(0, np, xp, yp, zp, uxp, uyp, uzp, 1, attr, 0, attr_int, 1); +} + void PhysicalParticleContainer::AddPlasmaFromFile(PlasmaInjector & plasma_injector, ParticleReal q_tot, @@ -909,6 +1085,10 @@ PhysicalParticleContainer::AddParticles (int lev) AddGaussianBeam(*plasma_injector); } + if (plasma_injector->twiss) { + AddTwiss(*plasma_injector); + } + if (plasma_injector->external_file) { AddPlasmaFromFile(*plasma_injector, plasma_injector->q_tot, diff --git a/Source/Utils/Parser/ParserUtils.cpp b/Source/Utils/Parser/ParserUtils.cpp index d017a6e019c..794d6829ec3 100644 --- a/Source/Utils/Parser/ParserUtils.cpp +++ b/Source/Utils/Parser/ParserUtils.cpp @@ -172,6 +172,7 @@ amrex::Parser utils::parser::makeParser ( {"m_u", PhysConst::m_u}, {"kb", PhysConst::kb}, {"pi", MathConst::pi}, + {"inf", std::numeric_limits::max()} }; for (auto it = symbols.begin(); it != symbols.end(); ) { diff --git a/Source/Utils/SpeciesUtils.cpp b/Source/Utils/SpeciesUtils.cpp index 56c269c2efe..1322afb8fcf 100644 --- a/Source/Utils/SpeciesUtils.cpp +++ b/Source/Utils/SpeciesUtils.cpp @@ -135,6 +135,14 @@ namespace SpeciesUtils { mom_dist_s.end(), mom_dist_s.begin(), ::tolower); + + if (style == "twiss") { + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + mom_dist_s == "twiss", + "'injection_style = twiss' requires " + "'momentum_distribution_type = twiss'"); + } + if (mom_dist_s == "at_rest") { constexpr amrex::Real ux = 0._rt; constexpr amrex::Real uy = 0._rt; @@ -166,6 +174,14 @@ namespace SpeciesUtils { // Construct InjectorMomentum with InjectorMomentumGaussian. h_inj_mom.reset(new InjectorMomentum((InjectorMomentumGaussian*)nullptr, ux_m, uy_m, uz_m, ux_th, uy_th, uz_th)); + } else if (mom_dist_s == "twiss") { + amrex::Real u0; + amrex::XDim3 sigma_u; + utils::parser::getWithParser(pp_species, source_name, "twiss.u0", u0); + utils::parser::getWithParser(pp_species, source_name, "twiss.sigma_ux", sigma_u.x); + utils::parser::getWithParser(pp_species, source_name, "twiss.sigma_uy", sigma_u.y); + utils::parser::getWithParser(pp_species, source_name, "twiss.sigma_uzeta", sigma_u.z); + h_inj_mom.reset(new InjectorMomentum((InjectorMomentumTwiss*) nullptr, u0, sigma_u)); } else if (mom_dist_s == "gaussianflux") { WARPX_ALWAYS_ASSERT_WITH_MESSAGE(style == "nfluxpercell", "Error: gaussianflux can only be used with injection_style = NFluxPerCell");