LinearMap: ds bookkeeping

And reference particle pushing (drifting), if thick.

src/particles/elements/LinearMap.H

@@ -14,7 +14,6 @@
 #include "mixin/alignment.H"
 #include "mixin/beamoptic.H"
 #include "mixin/lineartransport.H"
-#include "mixin/thin.H"
 #include "mixin/named.H"
 #include "mixin/nofinalize.H"
 
@@ -29,7 +28,6 @@ namespace impactx
     struct LinearMap
     : public elements::Named,
       public elements::BeamOptic<LinearMap>,
-      public elements::Thin,
       public elements::Alignment,
       public elements::LinearTransport,
       public elements::NoFinalize
@@ -43,13 +41,15 @@ namespace impactx
          *       px_final = R(2,1)*x + R(2,2)*px + R(2,3)*y + ..., etc.
          *
          * @param R user-provided transport map
+         * @param ds Segment length in m
          * @param dx horizontal translation error in m
          * @param dy vertical translation error in m
          * @param rotation_degree rotation error in the transverse plane [degrees]
          * @param name a user defined and not necessarily unique name of the element
          */
         LinearMap (
             LinearTransport::Map6x6 const & R,
+            amrex::ParticleReal ds = 0,
             amrex::ParticleReal dx = 0,
             amrex::ParticleReal dy = 0,
             amrex::ParticleReal rotation_degree = 0,
@@ -59,6 +59,7 @@ namespace impactx
           Alignment(dx, dy, rotation_degree)
         {
             m_transport_map = R;
+            m_ds = ds;
         }
 
         /** Push all particles */
@@ -112,10 +113,68 @@ namespace impactx
             shift_out(x, y, px, py);
         }
 
-        /** This pushes the reference particle. */
-        using Thin::operator();
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (RefPart & AMREX_RESTRICT refpart) const
+        {
+            if (m_ds > 0)  // Drift
+            {
+                using namespace amrex::literals; // for _rt and _prt
+
+                // assign input reference particle values
+                amrex::ParticleReal const x = refpart.x;
+                amrex::ParticleReal const px = refpart.px;
+                amrex::ParticleReal const y = refpart.y;
+                amrex::ParticleReal const py = refpart.py;
+                amrex::ParticleReal const z = refpart.z;
+                amrex::ParticleReal const pz = refpart.pz;
+                amrex::ParticleReal const t = refpart.t;
+                amrex::ParticleReal const pt = refpart.pt;
+                amrex::ParticleReal const s = refpart.s;
+
+                // length of the current slice
+                amrex::ParticleReal const slice_ds = m_ds / nslice();
+
+                // assign intermediate parameter
+                amrex::ParticleReal const step = slice_ds /std::sqrt(std::pow(pt,2)-1.0_prt);
+
+                // advance position and momentum (drift)
+                refpart.x = x + step*px;
+                refpart.y = y + step*py;
+                refpart.z = z + step*pz;
+                refpart.t = t - step*pt;
+
+                // advance integrated path length
+                refpart.s = s + slice_ds;
+            }
+            // else nothing to do for a zero-length element
+        }
+
+        /** Number of slices used for the application of space charge
+         *
+         * @return one, because we do not support slicing of this element
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        int nslice () const
+        {
+            return 1;
+        }
+
+        /** Return the segment length
+         *
+         * @return by default zero, but users can set a corresponding ds for bookkeeping
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        amrex::ParticleReal ds () const
+        {
+            return m_ds;
+        }
 
         LinearTransport::Map6x6 m_transport_map;  // 6x6 transport map
+        amrex::ParticleReal m_ds;  // finite ds allowed for bookkeeping, but we do not allow slicing
     };
 
 } // namespace impactx

src/python/elements.cpp

@@ -216,7 +216,7 @@ void init_elements(py::module& m)
         )
     ;
 
-    py::class_<elements::LinearTransport>(me, "LinearTransport")
+    py::class_<elements::LinearTransport>(mx, "LinearTransport")
         .def(py::init<>(),
              "Mixin class for linear transport approximation via matrices."
         )
@@ -1573,7 +1573,7 @@ void init_elements(py::module& m)
     ;
     register_beamoptics_push(py_TaperedPL);
 
-    py::class_<LinearMap, elements::Named, elements::Thin, elements::Alignment, elements::LinearTransport> py_LinearMap(me, "LinearMap");
+    py::class_<LinearMap, elements::Named, elements::Alignment, elements::LinearTransport> py_LinearMap(me, "LinearMap");
     py_LinearMap
         .def("__repr__",
              [](LinearMap const & linearmap) {
@@ -1589,9 +1589,11 @@ void init_elements(py::module& m)
                 amrex::ParticleReal,
                 amrex::ParticleReal,
                 amrex::ParticleReal,
+                amrex::ParticleReal,
                 std::optional<std::string>
              >(),
              py::arg("R"),
+             py::arg("ds") = 0,
              py::arg("dx") = 0,
              py::arg("dy") = 0,
              py::arg("rotation") = 0,
@@ -1603,6 +1605,15 @@ void init_elements(py::module& m)
             [](LinearMap & linearmap, elements::LinearTransport::Map6x6 R) { linearmap.m_transport_map = R; },
             "linear map as a 6x6 transport matrix"
         )
+        .def_property("ds",
+            [](LinearMap & linearmap) { return linearmap.m_ds; },
+            [](LinearMap & linearmap, amrex::ParticleReal ds) { linearmap.m_ds = ds; },
+            "segment length in m"
+        )
+        .def_property_readonly("nslice",
+            [](LinearMap & linearmap) { return linearmap.nslice(); },
+            "one, because we do not support slicing of this element"
+        )
      ;
      register_beamoptics_push(py_LinearMap);