Add support for current input, some refactoring.

examples/fodo_space_charge/input_fodo_envelope_sc.in

@@ -2,6 +2,7 @@
 # Particle Beam(s)
 ###############################################################################
 beam.kin_energy = 6.7
+beam.current = 0.5
 beam.particle = proton
 beam.distribution = kvdist_from_twiss
 beam.alphaX = 2.4685083

examples/fodo_space_charge/run_fodo_envelope_sc.py

@@ -27,6 +27,9 @@
 ref = sim.particle_container().ref_particle()
 ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
 
+#  beam current in A
+beam_current_A = 0.5
+
 #   particle bunch
 distr = distribution.KVdist(
     **twiss(
@@ -42,7 +45,7 @@
     )
 )
 
-sim.init_envelope(ref, distr)
+sim.init_envelope(ref, distr, beam_current_A)
 
 # add beam diagnostics
 monitor = elements.BeamMonitor("monitor", backend="h5")

src/initialization/AmrCoreData.H

@@ -89,7 +89,7 @@ namespace impactx::initialization
             std::optional<RefPart> m_ref;
 
             /** The 6x6 covariance matrix for the beam envelope, relative to the reference particle */
-            std::optional<CovarianceMatrix> m_cm;
+            std::optional<Envelope> m_env;
 
         } track_envelope;
 

src/initialization/InitDistribution.H

@@ -40,8 +40,9 @@ namespace impactx::initialization
 
     /** Ignore the shape of a distribution and use the 2nd moments to create a covariance matrix
      */
-    CovarianceMatrix
-    create_covariance_matrix (
+    Envelope
+    create_envelope (
+        amrex::ParticleReal const current,
         distribution::KnownDistributions const & distr
     );
 

src/initialization/InitDistribution.cpp

@@ -185,8 +185,9 @@ namespace impactx
         return dist;
     }
 
-    CovarianceMatrix
-    initialization::create_covariance_matrix (
+    Envelope
+    initialization::create_envelope (
+        amrex::ParticleReal const current,
         distribution::KnownDistributions const & distr
     )
     {
@@ -233,7 +234,11 @@ namespace impactx
             }
         }, distr);
 
-        return cv;
+        Envelope env;
+        env.set_beam_current_A(current);
+        env.set_covariance_matrix(cv);
+
+        return env;
     }
 
     void
@@ -471,17 +476,17 @@ namespace impactx
         }
         else if (track == "envelope")
         {
-            // For treating 3D space charge in bunched beams (not yet implemented in envelope mode)
-            //amrex::ParticleReal bunch_charge = 0.0;  // Bunch charge (C)
-            //pp_dist.getWithParser("charge", bunch_charge);
+            // relevant for envelope tracking with 3D space charge
+            amrex::ParticleReal bunch_charge = 0.0;  // Bunch charge (C)
+            pp_dist.query("charge", bunch_charge);
 
-            // For treating 2D space charge in unbunched beams
+            // relevant for envelope tracking with 2D space charge
             amrex::ParticleReal beam_current = 0.0;  // Beam current (A)
             pp_dist.query("current", beam_current);
 
             amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist);
             auto dist = initialization::read_distribution(pp_dist);
-            amr_data->track_envelope.m_cm = impactx::initialization::create_covariance_matrix(dist);
+            amr_data->track_envelope.m_env = impactx::initialization::create_envelope(beam_current,dist);
         }
         else if (track == "reference_orbit")
         {

src/particles/CovarianceMatrix.H

@@ -19,9 +19,60 @@ namespace impactx
     /** this is a 6x6 matrix */
     using Map6x6 = amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1>;
 
-    /** the covariance matrix is 6x6 */
+    /** this is the 6x6 covariance matrix */
     using CovarianceMatrix = Map6x6;
 
+
+    /** This struct stores the beam envelope attribues, including the 6x6 beam
+     *  covariance matrix.  Used during envelope tracking mode.
+     */
+    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
+
+        /** Set envelope beam current for 2D space charge
+         *
+         * @param beam_current beam current (A)
+         */
+        Envelope &
+        set_beam_current_A (amrex::ParticleReal const beam_current_A)
+        {
+            using namespace amrex::literals;
+
+            beam_current = beam_current_A;
+
+            return *this;
+         }
+
+        /** Get envelope beam current for 2D space charge
+         *
+         * @returns beam current in A
+         */ 
+        amrex::ParticleReal
+        beam_current_A ()
+        {
+            using namespace amrex::literals;
+
+            return beam_current;
+        }
+
+        /** Set 6x6 covariance matrix for envelope tracking
+         *
+         * @param Map6x6 beam 6x6 covariance matrix
+         */
+        Envelope &
+        set_covariance_matrix (CovarianceMatrix const covariance_matrix)
+        {
+            cm = covariance_matrix;
+
+            return *this;
+         }
+
+    };
+
+
 } // namespace impactx::distribution
 
 #endif // IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H

src/python/ImpactX.cpp

@@ -438,11 +438,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) {
+            [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr, amrex::Real current) {
                 ix.amr_data->track_envelope.m_ref = ref;
-                ix.amr_data->track_envelope.m_cm = initialization::create_covariance_matrix(distr);
+                ix.amr_data->track_envelope.m_env = initialization::create_envelope(current,distr);
             },
-            py::arg("ref"), py::arg("distr"),
+            py::arg("ref"), py::arg("distr"), py::arg("current"),
             "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/python/distribution.cpp

@@ -131,5 +131,5 @@ void init_distribution(py::module& m)
              "A 6D Waterbag distribution"
         );
 
-    m.def("create_covariance_matrix", &initialization::create_covariance_matrix);
+    m.def("create_covariance_matrix", &initialization::create_envelope);
 }

src/tracking/envelope.cpp

@@ -47,12 +47,14 @@ namespace impactx
         {
             throw std::runtime_error("track_envelope: Reference particle not set.");
         }
-        if (!amr_data->track_envelope.m_cm.has_value())
+        if (!amr_data->track_envelope.m_env.has_value())
         {
             throw std::runtime_error("track_envelope: Envelope (covariance matrix) not set.");
         }
         auto & ref = amr_data->track_envelope.m_ref.value();
-        auto & cm = amr_data->track_envelope.m_cm.value();
+        auto & env = amr_data->track_envelope.m_env.value();
+        auto & cm = env.cm;
+        auto & current = env.beam_current;
 
         // output of init state
         amrex::ParmParse pp_diag("diag");
@@ -127,7 +129,6 @@ namespace impactx
                     if (space_charge)
                     {
                         // push Covariance Matrix in 2D space charge fields
-                        amrex::ParticleReal current=0.5;  //TODO: This must be set.
                         spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds);
                     }