aboutsummaryrefslogtreecommitdiff
path: root/Source/Python/Particles
diff options
context:
space:
mode:
Diffstat (limited to 'Source/Python/Particles')
-rw-r--r--Source/Python/Particles/CMakeLists.txt17
-rw-r--r--Source/Python/Particles/MultiParticleContainer.cpp46
-rw-r--r--Source/Python/Particles/ParticleBoundaryBuffer.cpp59
-rw-r--r--Source/Python/Particles/PinnedMemoryParticleContainer.cpp23
-rw-r--r--Source/Python/Particles/WarpXParticleContainer.cpp119
5 files changed, 264 insertions, 0 deletions
diff --git a/Source/Python/Particles/CMakeLists.txt b/Source/Python/Particles/CMakeLists.txt
new file mode 100644
index 000000000..eed1bb07c
--- /dev/null
+++ b/Source/Python/Particles/CMakeLists.txt
@@ -0,0 +1,17 @@
+###############################################################################
+# These are the files equivalent to the primary C++ implementation - but here
+# we define how they will appear in our Python module (aka Python bindings).
+#
+foreach(D IN LISTS WarpX_DIMS)
+ warpx_set_suffix_dims(SD ${D})
+ if(WarpX_PYTHON)
+ target_sources(pyWarpX_${SD}
+ PRIVATE
+ # pybind11
+ ParticleBoundaryBuffer.cpp
+ MultiParticleContainer.cpp
+ PinnedMemoryParticleContainer.cpp
+ WarpXParticleContainer.cpp
+ )
+ endif()
+endforeach()
diff --git a/Source/Python/Particles/MultiParticleContainer.cpp b/Source/Python/Particles/MultiParticleContainer.cpp
new file mode 100644
index 000000000..e709f0950
--- /dev/null
+++ b/Source/Python/Particles/MultiParticleContainer.cpp
@@ -0,0 +1,46 @@
+/* Copyright 2021-2022 The WarpX Community
+ *
+ * Authors: Axel Huebl, Remi Lehe
+ * License: BSD-3-Clause-LBNL
+ */
+
+#include "Python/pyWarpX.H"
+
+#include <Particles/MultiParticleContainer.H>
+
+#include <AMReX_GpuContainers.H>
+#include <AMReX_REAL.H>
+
+
+void init_MultiParticleContainer (py::module& m)
+{
+ py::class_<MultiParticleContainer>(m, "MultiParticleContainer")
+ .def("get_particle_container_from_name",
+ &MultiParticleContainer::GetParticleContainerFromName,
+ py::arg("name"),
+ py::return_value_policy::reference_internal
+ )
+
+ .def("set_plasma_lens_strength",
+ [](MultiParticleContainer& mpc, int i_lens, amrex::Real strength_E, amrex::Real strength_B) {
+ mpc.h_repeated_plasma_lens_strengths_E.at(i_lens) = strength_E;
+ mpc.h_repeated_plasma_lens_strengths_B.at(i_lens) = strength_B;
+ amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice,
+ mpc.h_repeated_plasma_lens_strengths_E.begin(), mpc.h_repeated_plasma_lens_strengths_E.end(),
+ mpc.d_repeated_plasma_lens_strengths_E.begin());
+ amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice,
+ mpc.h_repeated_plasma_lens_strengths_B.begin(), mpc.h_repeated_plasma_lens_strengths_B.end(),
+ mpc.d_repeated_plasma_lens_strengths_B.begin());
+ amrex::Gpu::synchronize();
+ },
+ py::arg("i_lens"), py::arg("strength_E"), py::arg("strength_B"),
+ R"pbdoc(Set the strength of the `i_lens`-th lens
+Parameters
+----------
+i_lens: int
+ Index of the lens to be modified
+strength_E, strength_B: floats
+ The electric and magnetic focusing strength of the lens)pbdoc"
+ )
+ ;
+}
diff --git a/Source/Python/Particles/ParticleBoundaryBuffer.cpp b/Source/Python/Particles/ParticleBoundaryBuffer.cpp
new file mode 100644
index 000000000..2a35faece
--- /dev/null
+++ b/Source/Python/Particles/ParticleBoundaryBuffer.cpp
@@ -0,0 +1,59 @@
+/* Copyright 2021-2023 The WarpX Community
+ *
+ * Authors: Axel Huebl, Remi Lehe, Roelof Groenewald
+ * License: BSD-3-Clause-LBNL
+ */
+
+#include "Python/pyWarpX.H"
+
+#include <Particles/ParticleBoundaryBuffer.H>
+
+namespace warpx {
+ class BoundaryBufferParIter
+ : public amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>
+ {
+ public:
+ using amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>::ParIter;
+
+ BoundaryBufferParIter(ContainerType& pc, int level) :
+ amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>(pc, level) {}
+ };
+}
+
+void init_BoundaryBufferParIter (py::module& m)
+{
+ py::class_<
+ warpx::BoundaryBufferParIter,
+ amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>
+ >(m, "BoundaryBufferParIter")
+ .def(py::init<amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>::ContainerType&, int>(),
+ py::arg("particle_container"), py::arg("level")
+ )
+ ;
+}
+
+void init_ParticleBoundaryBuffer (py::module& m)
+{
+ py::class_<ParticleBoundaryBuffer>(m, "ParticleBoundaryBuffer")
+ .def(py::init<>())
+ .def("clear_particles",
+ [](ParticleBoundaryBuffer& pbb) { pbb.clearParticles(); }
+ )
+ .def("get_particle_container",
+ [](ParticleBoundaryBuffer& pbb,
+ const std::string species_name, int boundary) {
+ return &pbb.getParticleBuffer(species_name, boundary);
+ },
+ py::arg("species_name"), py::arg("boundary"),
+ py::return_value_policy::reference_internal
+ )
+ .def("get_num_particles_in_container",
+ [](ParticleBoundaryBuffer& pbb,
+ const std::string species_name, int boundary, bool local)
+ {
+ return pbb.getNumParticlesInContainer(species_name, boundary, local);
+ },
+ py::arg("species_name"), py::arg("boundary"), py::arg("local")=true
+ )
+ ;
+}
diff --git a/Source/Python/Particles/PinnedMemoryParticleContainer.cpp b/Source/Python/Particles/PinnedMemoryParticleContainer.cpp
new file mode 100644
index 000000000..d048b4d9c
--- /dev/null
+++ b/Source/Python/Particles/PinnedMemoryParticleContainer.cpp
@@ -0,0 +1,23 @@
+/* Copyright 2021-2023 The WarpX Community
+ *
+ * Authors: Axel Huebl, Remi Lehe, Roelof Groenewald
+ * License: BSD-3-Clause-LBNL
+ */
+
+#include "Python/pyWarpX.H"
+
+#include <Particles/PinnedMemoryParticleContainer.H>
+
+
+void init_PinnedMemoryParticleContainer (py::module& m)
+{
+ py::class_<
+ PinnedMemoryParticleContainer,
+ amrex::ParticleContainer<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>
+ > pmpc (m, "PinnedMemoryParticleContainer");
+ pmpc
+ .def("num_int_comps",
+ [](PinnedMemoryParticleContainer& pc) { return pc.NumIntComps(); }
+ )
+ ;
+}
diff --git a/Source/Python/Particles/WarpXParticleContainer.cpp b/Source/Python/Particles/WarpXParticleContainer.cpp
new file mode 100644
index 000000000..b5d3b1626
--- /dev/null
+++ b/Source/Python/Particles/WarpXParticleContainer.cpp
@@ -0,0 +1,119 @@
+/* Copyright 2021-2022 The WarpX Community
+ *
+ * Authors: Axel Huebl, Remi Lehe
+ * License: BSD-3-Clause-LBNL
+ */
+
+#include "Python/pyWarpX.H"
+
+#include <Particles/WarpXParticleContainer.H>
+
+
+void init_WarpXParIter (py::module& m)
+{
+ py::class_<
+ WarpXParIter, amrex::ParIter<0,0,PIdx::nattribs>
+ >(m, "WarpXParIter")
+ .def(py::init<amrex::ParIter<0,0,PIdx::nattribs>::ContainerType&, int>(),
+ py::arg("particle_container"), py::arg("level"))
+ .def(py::init<amrex::ParIter<0,0,PIdx::nattribs>::ContainerType&, int, amrex::MFItInfo&>(),
+ py::arg("particle_container"), py::arg("level"),
+ py::arg("info"))
+ ;
+}
+
+void init_WarpXParticleContainer (py::module& m)
+{
+ py::class_<
+ WarpXParticleContainer,
+ amrex::ParticleContainer<0, 0, PIdx::nattribs, 0>
+ > wpc (m, "WarpXParticleContainer");
+ wpc
+ .def("add_real_comp",
+ [](WarpXParticleContainer& pc, const std::string& name, bool const comm) { pc.AddRealComp(name, comm); },
+ py::arg("name"), py::arg("comm")
+ )
+ .def("add_n_particles",
+ [](WarpXParticleContainer& pc, int lev,
+ int n, py::array_t<double> &x,
+ py::array_t<double> &y,
+ py::array_t<double> &z,
+ py::array_t<double> &ux,
+ py::array_t<double> &uy,
+ py::array_t<double> &uz,
+ const int nattr_real, py::array_t<double> &attr_real,
+ const int nattr_int, py::array_t<int> &attr_int,
+ int uniqueparticles, int id
+ ) {
+ amrex::Vector<amrex::ParticleReal> xp(x.data(), x.data() + n);
+ amrex::Vector<amrex::ParticleReal> yp(y.data(), y.data() + n);
+ amrex::Vector<amrex::ParticleReal> zp(z.data(), z.data() + n);
+ amrex::Vector<amrex::ParticleReal> uxp(ux.data(), ux.data() + n);
+ amrex::Vector<amrex::ParticleReal> uyp(uy.data(), uy.data() + n);
+ amrex::Vector<amrex::ParticleReal> uzp(uz.data(), uz.data() + n);
+
+ // create 2d arrays of real and in attributes
+ amrex::Vector<amrex::Vector<amrex::ParticleReal>> attr;
+ const double *attr_data = attr_real.data();
+ for (int ii=0; ii<nattr_real; ii++) {
+ amrex::Vector<amrex::ParticleReal> attr_ii(n);
+ for (int jj=0; jj<n; jj++) {
+ attr_ii[jj] = attr_data[ii + jj*nattr_real];
+ }
+ attr.push_back(attr_ii);
+ }
+
+ amrex::Vector<amrex::Vector<int>> iattr;
+ const int *iattr_data = attr_int.data();
+ for (int ii=0; ii<nattr_int; ii++) {
+ amrex::Vector<int> attr_ii(n);
+ for (int jj=0; jj<n; jj++) {
+ attr_ii[jj] = iattr_data[ii + jj*nattr_int];
+ }
+ iattr.push_back(attr_ii);
+ }
+
+ pc.AddNParticles(
+ lev, n, xp, yp, zp, uxp, uyp, uzp, nattr_real, attr,
+ nattr_int, iattr, uniqueparticles, id
+ );
+ },
+ py::arg("lev"), py::arg("n"),
+ py::arg("x"), py::arg("y"), py::arg("z"),
+ py::arg("ux"), py::arg("uy"), py::arg("uz"),
+ py::arg("nattr_real"), py::arg("attr_real"),
+ py::arg("nattr_int"), py::arg("attr_int"),
+ py::arg("uniqueparticles"), py::arg("id")=-1
+ )
+ .def("num_real_comps", &WarpXParticleContainer::NumRealComps)
+ .def("get_comp_index",
+ [](WarpXParticleContainer& pc, std::string comp_name)
+ {
+ auto particle_comps = pc.getParticleComps();
+ return particle_comps.at(comp_name);
+ },
+ py::arg("comp_name")
+ )
+ .def("num_local_tiles_at_level",
+ &WarpXParticleContainer::numLocalTilesAtLevel,
+ py::arg("level")
+ )
+ .def("total_number_of_particles",
+ &WarpXParticleContainer::TotalNumberOfParticles,
+ py::arg("valid_particles_only"), py::arg("local")
+ )
+ .def("deposit_charge",
+ [](WarpXParticleContainer& pc,
+ amrex::MultiFab* rho, const int lev)
+ {
+ for (WarpXParIter pti(pc, lev); pti.isValid(); ++pti)
+ {
+ const long np = pti.numParticles();
+ auto& wp = pti.GetAttribs(PIdx::w);
+ pc.DepositCharge(pti, wp, nullptr, rho, 0, 0, np, 0, lev, lev);
+ }
+ },
+ py::arg("rho"), py::arg("lev")
+ )
+ ;
+}
generated by cgit v1.2.3 (git 2.46.0) at 2025-08-12 11:34:01 +0000