Python: Lattice Elements

src/ImpactX.H

@@ -92,6 +92,7 @@ namespace impactx
         //! Delete level data
         void ClearLevel (int lev) override;
 
+      public:
         /** Resize the mesh, based on the extent of the bunch of particle
          *
          * This only changes the physical extent of the mesh, but not the
@@ -107,6 +108,8 @@ namespace impactx
 
         /** these are elements defining the lattice */
         std::list<KnownElements> m_lattice;
+
+        std::list<std::variant<int>> m_lattice_test;
     };
 
 } // namespace impactx

src/particles/ImpactXParticleContainer.H

@@ -160,6 +160,16 @@ namespace impactx
          */
         void SetParticleShape ();
 
+        /** Set Particle Shape order
+         *
+         * Note: this can only be called once. All later calls are a logic error.
+         * The reason for that is that subsequent calls would need to change
+         * the guard size of all our MultiFabs, which is not implemented.
+         *
+         * @param order the order of the particle shape
+         */
+        void SetParticleShape (int const order);
+
         /** Compute the min and max of the particle position in each dimension
          *
          * @returns x_min, y_min, z_min, x_max, y_max, z_max

src/particles/ImpactXParticleContainer.cpp

@@ -29,23 +29,28 @@ namespace impactx
         SetParticleSize();
     }
 
-    void ImpactXParticleContainer::SetParticleShape ()
-    {
-        if (m_particle_shape.has_value()) {
+    void ImpactXParticleContainer::SetParticleShape (int const order) {
+        if (m_particle_shape.has_value())
+        {
             throw std::logic_error(
                 "ImpactXParticleContainer::SetParticleShape This was already called before and cannot be changed.");
         } else
         {
-            amrex::ParmParse pp_algo("algo");
-            int v = 0;
-            pp_algo.get("particle_shape", v);
-            m_particle_shape = v;
-            if (m_particle_shape.value() < 1 || m_particle_shape.value() > 3) {
-                amrex::Abort("algo.particle_shape can be only 1, 2, or 3");
+            if (order < 1 || order > 3) {
+                amrex::Abort("algo.particle_shape order can be only 1, 2, or 3");
             }
+            m_particle_shape = order;
         }
     }
 
+    void ImpactXParticleContainer::SetParticleShape ()
+    {
+        amrex::ParmParse pp_algo("algo");
+        int v = 0;
+        pp_algo.get("particle_shape", v);
+        SetParticleShape(v);
+    }
+
     void
     ImpactXParticleContainer::AddNParticles (int lev,
                                              amrex::Vector<amrex::ParticleReal> const & x,

src/particles/elements/All.H

@@ -17,15 +17,18 @@
 #include "ConstF.H"
 #include "ShortRF.H"
 #include "Multipole.H"
+#include "None.H"
 #include "NonlinearLens.H"
 
 #include <variant>
 
 
 namespace impactx
 {
-    using KnownElements = std::variant<Drift, Sbend, Quad, DipEdge, ConstF,
-                                       ShortRF, Multipole, NonlinearLens>;
+    using KnownElements = std::variant<
+        None, /* must be first, so KnownElements creates a default constructor */
+        ConstF, DipEdge, Drift, Multipole, NonlinearLens,
+        Quad, Sbend, ShortRF>;
 
 } // namespace impactx
 

src/particles/elements/None.H

@@ -0,0 +1,67 @@
+/* Copyright 2022 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_NONE_H
+#define IMPACTX_NONE_H
+
+#include "particles/ImpactXParticleContainer.H"
+
+#include <AMReX_Extension.H>
+#include <AMReX_REAL.H>
+
+
+namespace impactx
+{
+    struct None
+    {
+        using PType = ImpactXParticleContainer::ParticleType;
+
+        /** This element does nothing.
+         */
+        None ()
+        {
+        }
+
+        /** Does nothing to a particle.
+         *
+         * @param p Particle AoS data for positions and cpu/id
+         * @param px particle momentum in x
+         * @param py particle momentum in y
+         * @param pt particle momentum in t
+         * @param refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                [[maybe_unused]] PType& AMREX_RESTRICT p,
+                [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT px,
+                [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT py,
+                [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT pt,
+                [[maybe_unused]] RefPart const refpart) const
+        {
+            // nothing to do
+        }
+
+
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                [[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const {
+
+            // nothing to do: this is a zero-length element
+
+        }
+
+    };
+
+} // namespace impactx
+
+#endif // IMPACTX_NONE_H

src/python/CMakeLists.txt

@@ -4,5 +4,6 @@
 #
 target_sources(pyImpactX
   PRIVATE
+    elements.cpp
     ImpactX.cpp
 )

src/python/ImpactX.cpp

@@ -1,5 +1,9 @@
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+/* Copyright 2021-2022 The ImpactX Community
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "pyImpactX.H"
 
 #include <ImpactX.H>
 #include <AMReX.H>
@@ -19,6 +23,12 @@ namespace {
 
 void init_ImpactX(py::module& m)
 {
+    /*
+    py::bind_map<
+        std::unordered_map<int, amrex::MultiFab>
+    >(m, "MultiFabPerLevel");
+    */
+
     py::class_<ImpactX, amrex::AmrCore>(m, "ImpactX")
         .def(py::init<>())
 
@@ -28,7 +38,7 @@ void init_ImpactX(py::module& m)
                 // note: only in debug, since this is costly for the file
                 // system for highly parallel simulations with MPI
                 // possible improvement:
-                // - rank 0 tests file & broadcasts existance/failure
+                // - rank 0 tests file & broadcasts existence/failure
                 bool inputs_file_exists = false;
                 if (FILE *fp = fopen(filename.c_str(), "r")) {
                     fclose(fp);
@@ -38,13 +48,27 @@ void init_ImpactX(py::module& m)
                     "load_inputs_file: invalid filename");
 #endif
 
-                // TODO: needs https://github.com/AMReX-Codes/amrex/pull/2842
                 amrex::ParmParse::addfile(filename);
             })
+        .def("set_particle_shape",
+            [](ImpactX & ix, int const order) {
+                AMREX_ALWAYS_ASSERT_WITH_MESSAGE(ix.m_particle_container,
+                    "particle container not initialized");
+                // todo: why does that not work?
+                //ix.m_particle_container->SetParticleShape(order);
+
+                amrex::ParmParse pp_ago("algo");
+                pp_ago.add("particle_shape", order);
+            })
         .def("init_grids", &ImpactX::initGrids)
         .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs)
         .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs)
         .def("evolve", &ImpactX::evolve, py::arg("num_steps"))
+
+        //.def_property("particle_container", &ImpactX::m_particle_container)
+        //.def_readwrite("rho", &ImpactX::m_rho)
+        .def_readwrite("lattice", &ImpactX::m_lattice)
+        //.def_readwrite("lattice", &ImpactX::m_lattice_test)
     ;
 
     py::class_<Config>(m, "Config")

src/python/elements.cpp

@@ -0,0 +1,116 @@
+/* Copyright 2021-2022 The ImpactX Community
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "pyImpactX.H"
+
+#include <particles/elements/All.H>
+#include <AMReX.H>
+
+namespace py = pybind11;
+using namespace impactx;
+
+
+void init_elements(py::module& m)
+{
+    py::module_ me = m.def_submodule(
+        "elements",
+        "Accelerator lattice elements in ImpactX"
+    );
+
+    using KnownElementsList = std::list<KnownElements>;
+    py::class_<KnownElementsList> kel(me, "KnownElementsList");
+    kel
+        .def(py::init<>())
+        .def(py::init<KnownElements>())
+        .def(py::init([](py::list l){
+            auto v = new KnownElementsList;
+            for (auto &handle : l)
+                v->push_back(handle.cast<KnownElements>());
+            return v;
+        }))
+
+        .def("append", [](KnownElementsList &v, KnownElements el) { v.emplace_back(el); })
+
+        .def("extend", [](KnownElementsList &v, KnownElementsList l) {
+            for (auto &el : l)
+                v.push_back(el);
+            return v;
+        })
+        .def("extend", [](KnownElementsList &v, py::list l) {
+            for (auto &handle : l)
+            {
+                auto el = handle.cast<KnownElements>();
+                v.push_back(el);
+            }
+            return v;
+        })
+
+        .def("clear", &KnownElementsList::clear)
+        .def("pop_back", &KnownElementsList::pop_back)
+        .def("__len__", [](const KnownElementsList &v) { return v.size(); })
+        .def("__iter__", [](KnownElementsList &v) {
+            return py::make_iterator(v.begin(), v.end());
+        }, py::keep_alive<0, 1>()) /* Keep list alive while iterator is used */
+    ;
+
+    py::class_<ConstF>(me, "ConstF")
+        .def(py::init<
+                amrex::ParticleReal const,
+                amrex::ParticleReal const,
+                amrex::ParticleReal const,
+                amrex::ParticleReal const>(),
+             py::arg("ds"), py::arg("kx"), py::arg("ky"), py::arg("kt")
+        );
+
+    py::class_<DipEdge>(me, "DipEdge")
+        .def(py::init<
+                amrex::ParticleReal const,
+                amrex::ParticleReal const,
+                amrex::ParticleReal const,
+                amrex::ParticleReal const>(),
+             py::arg("psi"), py::arg("rc"), py::arg("g"), py::arg("K2")
+        );
+
+    py::class_<Drift>(me, "Drift")
+        .def(py::init<amrex::ParticleReal const>(),
+             py::arg("ds")
+        );
+
+    py::class_<Multipole>(me, "Multipole")
+        .def(py::init<
+                int const,
+                amrex::ParticleReal const,
+                amrex::ParticleReal const>(),
+             py::arg("multiple"), py::arg("K_normal"), py::arg("K_skew")
+        );
+
+    py::class_<NonlinearLens>(me, "NonlinearLens")
+        .def(py::init<
+                amrex::ParticleReal const,
+                amrex::ParticleReal const>(),
+             py::arg("knll"), py::arg("cnll")
+        );
+
+    py::class_<Sbend>(me, "Sbend")
+        .def(py::init<
+                amrex::ParticleReal const,
+                amrex::ParticleReal const>(),
+             py::arg("ds"), py::arg("rc")
+        );
+
+    py::class_<ShortRF>(me, "ShortRF")
+        .def(py::init<
+                amrex::ParticleReal const,
+                amrex::ParticleReal const>(),
+             py::arg("V"), py::arg("k")
+        );
+
+    py::class_<Quad>(me, "Quad")
+        .def(py::init<
+                amrex::ParticleReal const,
+                amrex::ParticleReal const>(),
+             py::arg("ds"), py::arg("k")
+        );
+}

src/python/pyImpactX.H

@@ -0,0 +1,21 @@
+/* Copyright 2021-2022 The ImpactX Community
+ *
+ * This header is used to centrally define classes that shall not violate the
+ * C++ one-definition-rule (ODR) for various Python translation units.
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
+#include <pybind11/functional.h>
+
+#include <particles/elements/All.H>
+
+#include <list>
+
+namespace py = pybind11;
+using namespace impactx;
+
+PYBIND11_MAKE_OPAQUE(std::list<KnownElements>)