Apply Changes from Review, Fix Tests

docs/source/usage/python.rst

@@ -180,14 +180,14 @@ Collective Effects & Overall Simulation Parameters
       :param distr: distribution function to draw from (object from :py:mod:`impactx.distribution`)
       :param int npart: number of particles to draw
 
-   .. py:method:: init_envelope(ref, distr, intensity)
+   .. py:method:: init_envelope(ref, distr, intensity=None)
 
       Envelope tracking mode:
       Create a 6x6 covariance matrix from a distribution and then initialize the the simulation for envelope tracking relative to a reference particle.
 
       :param ref: the reference particle (object from :py:class:`impactx.RefPart`)
       :param distr: distribution function (object from :py:mod:`impactx.distribution`)
-      :param float intensity: the beam intensity, given as bunch charge (C) or beam current (A)
+      :param float intensity: the beam intensity, given as bunch charge (C) for 3D or beam current (A) for 2D space charge
 
    .. py:method:: particle_container()
 

src/envelope/spacecharge/EnvelopeSpaceChargePush.H

@@ -13,10 +13,11 @@
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/CovarianceMatrix.H"
 
+#include <AMReX_Extension.H>  // for AMREX_RESTRICT
 #include <AMReX_REAL.H>
 
 
-namespace impactx::spacecharge
+namespace impactx::envelope::spacecharge
 {
      /** This function pushes the 6x6 beam covariance matrix for a slice
       *  of length ds, using the linear space charge fields in an rms
@@ -29,9 +30,9 @@ namespace impactx::spacecharge
       * @param[in] ds step size [m]
       */
      void
-     envelope_space_charge2D_push (
-         RefPart const & refpart,
-         Map6x6 & cm,
+     space_charge2D_push (
+         RefPart const & AMREX_RESTRICT refpart,
+         Map6x6 & AMREX_RESTRICT cm,
          amrex::ParticleReal current,
          amrex::ParticleReal ds
     );

src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp

@@ -9,20 +9,19 @@
  */
 #include "EnvelopeSpaceChargePush.H"
 
-#include <AMReX_BLProfiler.H>
 #include <AMReX_REAL.H>       // for Real
 #include <AMReX_SmallMatrix.H>
-#include <AMReX_ParmParse.H>
 #include <AMReX_Print.H>
 
 #include <cmath>
 
-namespace impactx::spacecharge
+
+namespace impactx::envelope::spacecharge
 {
-        void
-        envelope_space_charge2D_push (
-        [[maybe_unused]] RefPart const & refpart,
-        Map6x6 & cm,
+    void
+    space_charge2D_push (
+        RefPart const & AMREX_RESTRICT refpart,
+        Map6x6 & AMREX_RESTRICT cm,
         amrex::ParticleReal current,
         amrex::ParticleReal ds
     )
@@ -33,25 +32,25 @@ namespace impactx::spacecharge
         Map6x6 R = Map6x6::Identity();
 
         // physical constants and reference quantities
-        amrex::ParticleReal const c = ablastr::constant::SI::c;
-        amrex::ParticleReal const ep0 = ablastr::constant::SI::ep0;
-        amrex::ParticleReal const pi = ablastr::constant::math::pi;
+        using ablastr::constant::SI::c;
+        using ablastr::constant::SI::ep0;
+        using ablastr::constant::math::pi;
         amrex::ParticleReal const mass = refpart.mass;
         amrex::ParticleReal const charge = refpart.charge;
         amrex::ParticleReal const pt_ref = refpart.pt;
-        amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt;
+        amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1_prt;
         amrex::ParticleReal const betgam = std::sqrt(betgam2);
         amrex::ParticleReal const betgam3 = std::pow(betgam,3);
 
         // evaluate the beam space charge perveance from current
-        amrex::ParticleReal const IA = 4.0_prt*pi*ep0*mass*pow(c,3)/charge;
-        amrex::ParticleReal const Kpv = std::abs(current/IA) * 2.0_prt/betgam3;
+        amrex::ParticleReal const IA = 4_prt * pi * ep0 * mass * std::pow(c,3) / charge;
+        amrex::ParticleReal const Kpv = std::abs(current / IA) * 2_prt / betgam3;
 
         // evaluate the linear transfer map
-        amrex::ParticleReal const sigma2 = cm(1,1)*cm(3,3)-cm(1,3)*cm(1,3);
+        amrex::ParticleReal const sigma2 = cm(1,1) * cm(3,3) - cm(1,3) * cm(1,3);
         amrex::ParticleReal const sigma = std::sqrt(sigma2);
-        amrex::ParticleReal const D = (sigma + cm(1,1)) * (sigma + cm(3,3)) - cm(1,3)*cm(1,3);
-        amrex::ParticleReal const coeff = Kpv*ds/(2.0_prt*D);
+        amrex::ParticleReal const D = (sigma + cm(1,1)) * (sigma + cm(3,3)) - cm(1,3) * cm(1,3);
+        amrex::ParticleReal const coeff = Kpv * ds / (2_prt * D);
 
         R(2,1) = coeff * (sigma + cm(3,3));
         R(2,3) = coeff * (-cm(1,3));

src/initialization/InitDistribution.H

@@ -17,6 +17,7 @@
 #include <AMReX_Extension.H>  // for AMREX_RESTRICT
 #include <AMReX_REAL.H>
 
+#include <optional>  // for std::optional
 #include <utility>  // for std::move
 
 
@@ -39,11 +40,14 @@ namespace impactx::initialization
     read_distribution (amrex::ParmParse const & pp_dist);
 
     /** Ignore the shape of a distribution and use the 2nd moments to create a covariance matrix
+     *
+     * @param distr the distribution to draw from
+     * @param intensity the beam charge (in C, 3D space charge) or current (in A, 2D space charge)
      */
     Envelope
     create_envelope (
-        amrex::ParticleReal const current,
-        distribution::KnownDistributions const & distr
+        distribution::KnownDistributions const & distr,
+        std::optional<amrex::ParticleReal> intensity
     );
 
     /** Initialize a single particle's data using the given distribution

src/initialization/InitDistribution.cpp

@@ -187,8 +187,8 @@ namespace impactx
 
     Envelope
     initialization::create_envelope (
-        amrex::ParticleReal const current,
-        distribution::KnownDistributions const & distr
+        distribution::KnownDistributions const & distr,
+        std::optional<amrex::ParticleReal> intensity
     )
     {
         // zero out the 6x6 matrix
@@ -235,7 +235,7 @@ namespace impactx
         }, distr);
 
         Envelope env;
-        env.set_beam_current_A(current);
+        if (intensity) { env.set_beam_intensity(intensity.value()); }
         env.set_covariance_matrix(cv);
 
         return env;
@@ -491,7 +491,7 @@ namespace impactx
 
             amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist);
             auto dist = initialization::read_distribution(pp_dist);
-            amr_data->track_envelope.m_env = impactx::initialization::create_envelope(intensity,dist);
+            amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist, intensity);
         }
         else if (track == "reference_orbit")
         {

src/particles/CovarianceMatrix.H

@@ -28,44 +28,41 @@ namespace impactx
      */
     struct Envelope
     {
-        amrex::ParticleReal bunch_charge = 0.0; ///< total charge in C (for 3D space charge)
-        amrex::ParticleReal beam_current = 0.0; ///< current in A (for 2D space charge)
-        CovarianceMatrix cm; ///< the 6x6 beam covariance matrix
+        CovarianceMatrix m_env; ///< the 6x6 beam covariance matrix
+        amrex::ParticleReal m_beam_intensity = 0.0; ///< optional: charge in A (for 3D space charge) or current in A (for 2D space charge)
 
-        /** Set envelope beam current for 2D space charge
+        /** Set envelope beam charge/current for 3D/2D space charge
          *
-         * @param beam_current_A beam current (A)
+         * @param intensity beam charge (C) in 3D or beam current (A) in 2D
          */
         Envelope &
-        set_beam_current_A (amrex::ParticleReal const beam_current_A)
+        set_beam_intensity (amrex::ParticleReal const intensity)
         {
-            using namespace amrex::literals;
-
-            beam_current = beam_current_A;
+            m_beam_intensity = intensity;
 
             return *this;
          }
 
-        /** Get envelope beam current for 2D space charge
+        /** Get envelope beam charge/current for 3D/2D space charge
          *
-         * @returns beam current in A
+         * @returns 3D: beam charge in C; 2D: beam current in A
          */
         amrex::ParticleReal
-        beam_current_A ()
+        beam_intensity () const
         {
             using namespace amrex::literals;
 
-            return beam_current;
+            return m_beam_intensity;
         }
 
         /** Set 6x6 covariance matrix for envelope tracking
          *
          * @param covariance_matrix beam 6x6 covariance matrix
          */
         Envelope &
-        set_covariance_matrix (const CovarianceMatrix& covariance_matrix)
+        set_covariance_matrix (CovarianceMatrix const & covariance_matrix)
         {
-            cm = covariance_matrix;
+            m_env = covariance_matrix;
 
             return *this;
          }

src/python/ImpactX.cpp

@@ -16,6 +16,7 @@
 #if defined(AMREX_DEBUG) || defined(DEBUG)
 #   include <cstdio>
 #endif
+#include <optional>
 #include <string>
 
 
@@ -452,11 +453,11 @@ void init_ImpactX (py::module& m)
         .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs)
         .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs)
         .def("init_envelope",
-            [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr, amrex::Real current) {
+            [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr, std::optional<amrex::Real> intensity) {
                 ix.amr_data->track_envelope.m_ref = ref;
-                ix.amr_data->track_envelope.m_env = initialization::create_envelope(current,distr);
+                ix.amr_data->track_envelope.m_env = initialization::create_envelope(distr, intensity);
             },
-            py::arg("ref"), py::arg("distr"), py::arg("current") = 0.0,
+            py::arg("ref"), py::arg("distr"), py::arg("intensity") = py::none(),
             "Envelope tracking mode:"
             "Create a 6x6 covariance matrix from a distribution and then initialize "
             "the the simulation for envelope tracking relative to a reference particle."

src/tracking/envelope.cpp

@@ -53,8 +53,8 @@ namespace impactx
         }
         auto & ref = amr_data->track_envelope.m_ref.value();
         auto & env = amr_data->track_envelope.m_env.value();
-        auto & cm = env.cm;
-        auto & current = env.beam_current;
+        auto & cm = env.m_env;
+        auto & current = env.m_beam_intensity;
 
         // output of init state
         amrex::ParmParse pp_diag("diag");
@@ -129,7 +129,7 @@ namespace impactx
                     if (space_charge)
                     {
                         // push Covariance Matrix in 2D space charge fields
-                        spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds);
+                        envelope::spacecharge::space_charge2D_push(ref,cm,current,slice_ds);
                     }
 
                     std::visit([&ref, &cm](auto&& element)