Python: Envelope Class

BLAST-ImpactX · Feb 21, 2025 · 0d04230 · 0d04230
1 parent 865cbd1
commit 0d04230
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H
@@ -23,7 +23,7 @@ namespace impactx
     using CovarianceMatrix = Map6x6;
 
 
-    /** This struct stores the beam envelope attribues, including the 6x6 beam
+    /** This struct stores the beam envelope attributes, including the 6x6 beam
      *  covariance matrix.  Used during envelope tracking mode.
      */
     struct Envelope
@@ -50,8 +50,6 @@ namespace impactx
         amrex::ParticleReal
         beam_intensity () const
         {
-            using namespace amrex::literals;
-
             return m_beam_intensity;
         }
 
@@ -67,6 +65,14 @@ namespace impactx
             return *this;
          }
 
+        /** Get the 6x6 covariance matrix for envelope tracking
+         */
+        CovarianceMatrix
+        covariance_matrix () const
+        {
+            return m_env;
+        }
+
     };
 
 

diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp
@@ -131,5 +131,12 @@ void init_distribution(py::module& m)
              "A 6D Waterbag distribution"
         );
 
+    py::class_<Envelope>(m, "Envelope")
+        .def(py::init<>())
+        .def(py::init<CovarianceMatrix, amrex::ParticleReal>())
+        .def_property("envelope", &Envelope::covariance_matrix, &Envelope::set_covariance_matrix)
+        .def_property("beam_intensity", &Envelope::beam_intensity, &Envelope::set_beam_intensity)
+    ;
+
     m.def("create_envelope", &initialization::create_envelope);
 }