Transition to pyAMReX (#3474)

* pyAMReX: Build System

* CI Updates (Changed Options)

* Callback modernization (#4)

* refactor callbacks.py
* added binding code in `pyWarpX.cpp` to add or remove keys from the callback dictionary
* minor PR cleanups

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Added Python level reference to fetch the multifabs (#3)

* pyAMReX: Build System

* Added Python level reference to Ex_aux

* Now uses the multifab map

* Fix typo

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Add initialization and finalize routines (#5)

A basic PICMI input file will now run to completion.

* Regression Tests: WarpX_PYTHON=ON

* Update Imports to nD pyAMReX

* IPO/LTO Control

Although pybind11 relies heavily on IPO/LTO to create low-latency,
small-binary bindings, some compilers will have troubles with that.

Thus, we add a compile-time option to optionally disable it when
needed.

* Fix: Link Legacy WarpXWrappers.cpp

* Wrap WarpX instance and start multi particle container

* Fix test Python_pass_mpi_comm
* Start wrapper for multiparticle container
* Add return policy

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Update fields to use MultiFabs directly

Remove EOL white space

Removed old routines accessing MultiFabs

Update to use "node_centered"

* Fix compilation with Python

* Update fields.py to use modified MultiFab tag names

* Remove incorrect, unused code

* Add function to extract the WarpX MultiParticleContainer

* Complete class WarpXParticleContainer

* Wrap functions getNprocs / getMyProc

* restore `install___` callback API - could remove later if we want but should maintain backward compatibility for now

* add `gett_new` and `getistep` functions wrappers; fix typos in `callbacks.py`; avoid error in getting `rho` from `fields.py`

* Update callback call and `getNproc`/`getMyProc` function

* Replace function add_n_particles

* Fix setitem in fields.py for 1d and 2d

* also update `gett_new()` in `_libwarpx.py` in case we want to keep that API

* added binding for `WarpXParIter` - needed to port `libwarpx.depositChargeDensity()` which is an ongoing effort

* Wrap function num_real_comp

* added binding for `TotalNumberOfParticles` and continue progress on enabling 1d MCC test to run

* add `SyncRho()` binding and create helper function in `libwarpx.depositChargeDensity` to manage scope of the particle iter

* Clean up issues in fields.py

* update bindings for `get_particle_structs`

* Fix setitem in fields.py

* switch order of initialization for particle container and particle iterator

* switch deposit_charge loop to C++ code; bind `ApplyInverseVolumeScalingToChargeDensity`

* move `WarpXParticleContainer.cpp` and `MultiParticleContainer.cpp` to new Particles folder

* added binding for `ParticleBoundaryBuffer`

* More fixes for fields.py

* Fix: Backtraces from Python

Add the Python executable name with an absolute path, so backtraces
in AMReX work. See linked AMReX issue for details.

* Cleaning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix: Backtraces from Python Part II

Do not add Python script name - it confuses the AMReX ParmParser to
build its table.

* Fix: CMake Dependencies for Wheel

This fixes a racecondition during `pip_install`: it was possible
that not all dimensions where yet build from pybind before we
start packing them in the wheel for pip install.

* MCC Test: Install Callbacks before Run

Otherwise hangs in aquiring the gil during shutdown.

* addition of `Python/pywarpx/particle_containers.py` and various associated bindings

* Fix: CMake Superbuild w/ Shared AMReX

We MUST build AMReX as a shared (so/dll/dylib) library, otherwise
all the global state in it will cause split-brain situations, where
our Python modules operate on different stacks.

* add `clear_all()` to callbacks in order to remove all callbacks at finalize

* add `-DWarpX_PYTHON=ON` to CI tests that failed to build

* add `get_comp_index` and continue to port particle data bindings

* Add AMReX Module as `libwarpx_so.amr`

Attribute to pass through the already loaded AMReX module with the
matching dimensionality to the simulation.

* Fix for fields accounting for guard cells

* Fix handling of ghost cells in fields

* Update & Test: Particle Boundary Scraping

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* CI: Python Updates

- modernize Python setups
- drop CUDA 11.0 for good and go 11.3+ as documented already
```
Error #3246: Internal Compiler Error (codegen): "there was an error in verifying the lgenfe output!"
```

* CI: Python Updates (chmod)

* Add support for cupy in fields.py

* Add MultiFab reduction routines

* CI: CUDA 11.3 is <= Ubuntu 20.04

* changed `AddNParticles` to take `amrex::Vector` arguments

* setup.py: WarpX_PYTHON=ON

* update various 2d and rz tests with new APIs

* add `-DWarpX_PYTHON_IPO=OFF` to compile string for 2d and 3d Python CI tests to speed up linking

* CI: -DpyAMReX_IPO=OFF

* CI: -DpyAMReX_IPO=OFF

actually adding `=OFF`

* CI: Intel Python

* CI: macOS Python Executable

Ensure we always use the same `python3` executable, as specified
by the `PATH` priority.

* CMake: Python Multi-Config Build

Add support for multi-config generators, especially on Windows.

* __init__.py: Windows DLL Support

Python 3.8+ on Windows: DLL search paths for dependent
shared libraries
Refs.:
- https://github.com/python/cpython/issues/80266
- https://docs.python.org/3.8/library/os.html#os.add_dll_directory

* CI: pywarpx Update

our setup.py cannot install pyamrex yet as a dependency.

* ABLASTR: `ablastr/export.H`

Add a new header to export public globals that are not covered by
`WINDOWS_EXPORT_ALL_SYMBOLS`.
https://stackoverflow.com/questions/54560832/cmake-windows-export-all-symbols-does-not-cover-global-variables/54568678#54568678

* WarpX: EXPORT Globals in `.dll` files

WarpX still uses a lot of globals:
- `static` member variables
- `extern` global variables

These globals cannot be auto-exported with CMake's
`WINDOWS_EXPORT_ALL_SYMBOLS` helper and thus we need to mark them
manually for DLL export (and import) via the new ABLASTR
`ablastr/export.H` helper macros.

This starts to export symbols in the:
- WarpX and particle container classes
- callback hook database map
- ES solver

* CI: pywarpx Clang CXXFLAGS Down

Move CXXFLAGS (`-Werror ...`) down until deps are installed.

* GNUmake: Generate `ablastr/export.H`

* CMake: More Symbol Exports for Windows

* `WarpX-tests.ini`: Simplify Python no-IPO

Also avoids subtle differences in compilation that increase
compile time.

* Update PICMI_inputs_EB_API.py for embedded_boundary_python_API CI test

* Fix Python_magnetostatic_eb_3d

* Update: Python_restart_runtime_components

New Python APIs

* Windows: no dllimport for now

* CI: Skip `PYINSTALLOPTIONS` For Now

* CMake: Dependency Bump Min-Versions

for external packages picked up by `find_package`.

* Fix F and G_fp names in fields.py

* Tests: Disable `Python_pass_mpi_comm`

* Wrappers: Cleanup

* pyWarpX: Include Cleaning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fields.py: Fix F and G Wrappers

Correct MultiFab names (w/o components).

* Remove unused in fields.py

* Windows: New Export Headers

- ABLASTR: `ablastr/export.H`
- WarpX: `Utils/export.H`

* removed `WarpInterface.py` since that functionality is now in `particle_containers.py`; removed parts of `WarpXWrappers.cpp` that have been ported to pyamrex

* CMake: Link OBJECT Target PRIVATE

* CMake: Remove OBJECT Target

Simplify and make `app` link `lib` (default: static). Remove OBJECT
target.

* Fix in fields.py for the components index

* Update get_particle_id/cpu

As implemented in pyAMReX with
https://github.com/AMReX-Codes/pyamrex/pull/165

* WarpX: Update for Private Constructor

* Import AMReX Before pyd DLL Call

Importing AMReX will add the `add_dll_directory` to a potentially
shared amrex DLL on Windows.

* Windows CI: Set PATH to amrex_Nd.dll

* CMake: AMReX_INSTALL After Python

In superbuild, Python can modify `AMReX_BUILD_SHARED_LIBS`.

* Clang Win CI: Manually Install requirements

Sporadic error is:
```
...
Installing collected packages: pyparsing, numpy, scipy, periodictable, picmistandard
ERROR: Could not install packages due to an OSError: [WinError 32] The process cannot access the file because it is being used by another process: 'C:\\hostedtoolcache\\windows\\Python\\3.11.4\\x64\\Lib\\site-packages\\numpy\\polynomial\\__init__.py'
Consider using the `--user` option or check the permissions.
```

* Hopefully final fixes to fields.py

* Update getProbLo/getProbHi

* Set plasma length strength

Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>

* Fix fields method to remove CodeQL notice

* Update Comments & Some Finalize

* Move: set_plasma_lens_strength

to MPC

---------

Co-authored-by: Roelof Groenewald <40245517+roelof-groenewald@users.noreply.github.com>
Co-authored-by: David Grote <dpgrote@lbl.gov>
Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>
Co-authored-by: Dave Grote <grote1@llnl.gov>
Co-authored-by: Roelof Groenewald <regroenewald@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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")
+        )
+    ;
+}