Python Bindings

Author: ax3l

src/CompensatedReal.H

@@ -213,9 +213,9 @@ namespace impactx
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
         CompensatedReal operator+(CompensatedReal const& rhs) const {
-            CompensatedReal result = *this;
+            CompensatedReal result(*this);
             result += rhs;
-            return result;
+            return CompensatedReal(result);
         }
 
         /** Subtraction operator with CompensatedReal
@@ -225,9 +225,9 @@ namespace impactx
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
         CompensatedReal operator-(CompensatedReal const& rhs) const {
-            CompensatedReal result = *this;
+            CompensatedReal result(*this);
             result -= rhs;
-            return result;
+            return CompensatedReal(result);
         }
 
         /** Equality comparison

src/python/CMakeLists.txt

@@ -4,6 +4,7 @@
 #
 target_sources(pyImpactX
   PRIVATE
+    CompensatedReal.cpp
     distribution.cpp
     elements.cpp
     ImpactX.cpp

src/python/CompensatedReal.cpp

@@ -0,0 +1,106 @@
+/* Copyright 2022-2024 The ImpactX Community
+ *
+ * Authors: Axel Huebl, Chad Mitchell
+ * License: BSD-3-Clause-LBNL
+ */
+#include "pyImpactX.H"
+
+#include <CompensatedReal.H>
+#include <AMReX_REAL.H>
+
+#include <cmath>
+
+namespace py = pybind11;
+using namespace impactx;
+
+
+void init_compensatedreal(py::module& m)
+{
+    // Expose the CompensatedParticleReal typedef
+    py::class_<CompensatedParticleReal>(m, "CompensatedParticleReal")
+        .def(py::init<>(),
+             "Default (zero) constructor for compensated ParticleReal type"
+        )
+        .def(py::init<amrex::ParticleReal>(),
+             py::arg("value"),
+             "Constructor from a ParticleReal value"
+        )
+        .def(py::init<CompensatedParticleReal const&>(),
+             py::arg("other"),
+             "Copy constructor"
+        )
+        .def("assign", &CompensatedParticleReal::assign,
+             py::arg("value"),
+             py::return_value_policy::reference_internal,
+             "Assign a new value and reset compensation"
+        )
+
+        // Arithmetic operators
+        .def(py::self += amrex::ParticleReal())
+        .def(py::self -= amrex::ParticleReal())
+        .def(py::self + amrex::ParticleReal())
+        .def(py::self - amrex::ParticleReal())
+        .def(amrex::ParticleReal() + py::self)
+        .def(amrex::ParticleReal() - py::self)
+        .def(py::self += py::self)
+        .def(py::self -= py::self)
+        .def(py::self + py::self)
+        .def(py::self - py::self)
+
+        // Comparison operators
+        .def(py::self == amrex::ParticleReal())
+        .def(py::self != amrex::ParticleReal())
+        .def(py::self < amrex::ParticleReal())
+        .def(py::self <= amrex::ParticleReal())
+        .def(py::self > amrex::ParticleReal())
+        .def(py::self >= amrex::ParticleReal())
+        .def(py::self == py::self)
+        .def(py::self != py::self)
+        .def(py::self < py::self)
+        .def(py::self <= py::self)
+        .def(py::self > py::self)
+        .def(py::self >= py::self)
+
+        // Value access
+        .def_property_readonly("value", &CompensatedParticleReal::value,
+             "Get the final compensated value (main + first-order + second-order compensation)"
+        )
+
+        // conversion to float for compatibility
+        .def("__float__", &CompensatedParticleReal::value,
+             "Convert to Python float"
+        )
+        .def("__abs__", [](CompensatedParticleReal const& self) {
+            return std::abs(self.value());
+        }, "Convert to absolute value")
+
+        // Utility methods
+        .def("reset_compensation", &CompensatedParticleReal::reset_compensation,
+             py::return_value_policy::reference_internal,
+             "Reset compensation variables to zero"
+        )
+
+        // String representation
+        .def("__repr__", [](CompensatedParticleReal const& self) {
+            return "CompensatedParticleReal(" + std::to_string(self.value()) + ")";
+        })
+        .def("__str__", [](CompensatedParticleReal const& self) {
+            return std::to_string(self.value());
+        })
+
+        .def_readonly_static("__doc__",
+             "A compensated floating-point type specialized for ParticleReal values.\n"
+             "\n"
+             "This provides high-precision arithmetic for particle tracking calculations\n"
+             "using the second-order iterative Kahan-Babuska algorithm.\n"
+             "\n"
+             "The algorithm follows Klein (2006) to provide high-precision arithmetic that avoids\n"
+             "floating-point precision errors when dealing with large numbers of small values.\n"
+             "\n"
+             "References:\n"
+             "- https://en.wikipedia.org/wiki/Kahan_summation_algorithm#Further_enhancements\n"
+             "- Klein (2006). \"A generalized Kahan–Babuška-Summation-Algorithm\". in\n"
+             "  Computing. 76 (3–4). Springer-Verlag: 279–293. doi:10.1007/s00607-005-0139-x"
+        )
+    ;
+}

src/python/pyImpactX.H

@@ -12,6 +12,7 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/functional.h>
 // include <pybind11/numpy.h>  // not yet used
+#include <pybind11/operators.h>
 #include <pybind11/stl.h>
 #include <pybind11/stl_bind.h>
 

src/python/pyImpactX.cpp

@@ -20,6 +20,7 @@ using namespace impactx;
 
 
 // forward declarations of exposed classes
+void init_compensatedreal(py::module&);
 void init_distribution(py::module&);
 void init_elements(py::module&);
 void init_ImpactX(py::module&);
@@ -45,6 +46,7 @@ PYBIND11_MODULE(impactx_pybind, m) {
     )pbdoc";
 
     // note: order from parent to child classes
+    init_compensatedreal(m);
     init_distribution(m);
     init_refparticle(m);
     init_impactxparticlecontainer(m);