Add LinearMap to 'Named' elements.

ECP-WarpX · Nov 2, 2024 · 2e57e91 · 2e57e91
1 parent 06718d8
commit 2e57e91
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 16 deletions.
diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H
@@ -15,6 +15,7 @@
 #include "mixin/beamoptic.H"
 #include "mixin/lineartransport.H"
 #include "mixin/thin.H"
+#include "mixin/named.H"
 #include "mixin/nofinalize.H"
 
 #include <AMReX_Extension.H>
@@ -26,7 +27,8 @@
 namespace impactx
 {
     struct LinearMap
-    : public elements::BeamOptic<LinearMap>,
+    : public elements::Named,
+      public elements::BeamOptic<LinearMap>,
       public elements::Thin,
       public elements::Alignment,
       public elements::LinearTransport,
@@ -49,9 +51,11 @@ namespace impactx
             LinearTransport::Map6x6 const & R,
             amrex::ParticleReal dx = 0,
             amrex::ParticleReal dy = 0,
-            amrex::ParticleReal rotation_degree = 0
+            amrex::ParticleReal rotation_degree = 0,
+            std::optional<std::string> name = std::nullopt
         )
-        : Alignment(dx, dy, rotation_degree)
+        : Named(name),
+        Alignment(dx, dy, rotation_degree)
         {
             for (int i=1; i<=6; ++i) {
                 for (int j = 1; j <= 6; ++j) {

diff --git a/src/python/elements.cpp b/src/python/elements.cpp
@@ -7,6 +7,7 @@
 
 #include <particles/Push.H>
 #include <particles/elements/All.H>
+#include <particles/elements/mixin/lineartransport.H>
 #include <AMReX.H>
 
 #include <optional>
@@ -1480,7 +1481,7 @@ void init_elements(py::module& m)
     ;
     register_beamoptics_push(py_TaperedPL);
 
-    py::class_<LinearMap, elements::Named, elements::Thin, elements::Alignment> py_LinearMap(me, "LinearMap");
+    py::class_<LinearMap, elements::Named, elements::Thin, elements::Alignment, elements::LinearTransport> py_LinearMap(me, "LinearMap");
     py_LinearMap
         .def("__repr__",
              [](LinearMap const & linearmap) {
@@ -1490,17 +1491,14 @@ void init_elements(py::module& m)
                  );
              }
         )
-        .def(py::init([](
-                elements::LinearTransport::Map6x6 R,
-                amrex::ParticleReal dx,
-                amrex::ParticleReal dy,
-                amrex::ParticleReal rotation_degree,
-                std::optional<std::string> name
-             )
-             {
-             }),
+        .def(py::init<
+                elements::LinearTransport::Map6x6,
+                amrex::ParticleReal,
+                amrex::ParticleReal,
+                amrex::ParticleReal,
+                std::optional<std::string>
+             >(),
              py::arg("R"),
-             py::arg("unit") = "dimensionless",
              py::arg("dx") = 0,
              py::arg("dy") = 0,
              py::arg("rotation") = 0,
@@ -1509,10 +1507,11 @@ void init_elements(py::module& m)
         )
         .def_property("R",
             [](LinearMap & linearmap) { return linearmap.m_transport_map; },
-            [](LinearMap & linearmap, LinearTransport::Map6x6 R) { linearmap.m_transport_map = linearmap; },
+            [](LinearMap & linearmap, elements::LinearTransport::Map6x6 R) { linearmap.m_transport_map = R; },
             "linear map as a 6x6 transport matrix"
         )
-
+     ;
+     register_beamoptics_push(py_LinearMap);
 
     // freestanding push function
     m.def("push", &Push,