Merge branch 'development' into Use_amrex_GpuComplex_in_diagnostics

.github/workflows/macos.yml

@@ -58,11 +58,10 @@ jobs:
         brew --cache
         set +e
 
-        brew install --overwrite python
-        brew link --force --overwrite python
-
-        brew install --overwrite python@3.12
-        brew link --force --overwrite python@3.12
+        rm -rf /usr/local/bin/2to3*
+        rm -rf /usr/local/bin/idle3*
+        rm -rf /usr/local/bin/pydoc3*
+        rm -rf /usr/local/bin/python3*
 
         brew install fftw --only-dependencies --force
         brew install fftw
@@ -89,9 +88,6 @@ jobs:
         brew install openpmd-api
         brew link openpmd-api
 
-        python3 -m pip install matplotlib
-        python3 -m pip install numpy
-        python3 -m pip install scipy
     - name: build HiPACE++
       run: |
         cmake -S . -B build_sp          \

src/mg_solver/HpMultiGrid.H

@@ -174,6 +174,8 @@ private:
     int m_num_mg_levels;
     /** Number of single-block-kernel levels */
     int m_num_single_block_levels;
+    /** If the single block kernel should be used */
+    bool m_use_single_block_kernel = true;
 
     /** Alias to the solution argument passed in solve() */
     amrex::FArrayBox m_sol;

src/mg_solver/HpMultiGrid.cpp

@@ -693,6 +693,16 @@ MultiGrid::MultiGrid (Real dx, Real dy, Box a_domain)
                               { return b.volume() <= n_cell_single*n_cell_single; });
     m_single_block_level_begin = std::distance(std::begin(m_domain), r);
     m_single_block_level_begin = std::max(1, m_single_block_level_begin);
+    if (m_single_block_level_begin > m_max_level) {
+        m_single_block_level_begin = m_max_level;
+        m_use_single_block_kernel = false;
+        amrex::Print() << "hpmg: WARNING domain of size "
+            << a_domain_len[0] << " " << a_domain_len[1]
+            << " cannot be coarsened enough times to be solved efficiently.\n"
+            << "hpmg: Size of the final MG level: "
+            << m_domain[m_max_level].length(0) << " " << m_domain[m_max_level].length(1) << ".\n"
+            << "hpmg: Please consider using a domain size of the form '2^n', '3*2^n', '2^n+1' or '3*n^2+1'.\n";
+    }
 #else
     m_single_block_level_begin = m_max_level;
 #endif
@@ -1099,62 +1109,67 @@ MultiGrid::bottomsolve ()
     Real dx0 = m_dx * fac;
     Real dy0 = m_dy * fac;
 #if defined(AMREX_USE_GPU)
-    Array4<amrex::Real> const* acf = m_acf_a;
-    Array4<amrex::Real> const* res = m_res_a;
-    Array4<amrex::Real> const* cor = m_cor_a;
-    Array4<amrex::Real> const* rescor = m_rescor_a;
-    int nlevs = m_num_single_block_levels;
-    int const corner_offset = m_domain[0].cellCentered() ? 0 : 1;
-
-    if (m_system_type == 1) {
-        bottomsolve_gpu<nsweeps>(dx0, dy0, acf, res, cor, rescor, nlevs, corner_offset,
-            [=] AMREX_GPU_DEVICE (int i, int j, int ilo, int jlo, int ihi, int jhi,
-                                  Array4<Real> const& phi, Real rhs0, Real rhs1,
-                                  Array4<Real> const& acf, Real facx, Real facy)
-            {
-                Real a = acf(i,j,0);
-                gs1(i, j, 0, ilo, jlo, ihi, jhi, phi, rhs0, a, facx, facy);
-                gs1(i, j, 1, ilo, jlo, ihi, jhi, phi, rhs1, a, facx, facy);
-            },
-            [=] AMREX_GPU_DEVICE (int i, int j, Real& res0, Real& res1,
-                                  int ilo, int jlo, int ihi, int jhi,
-                                  Array4<Real> const& phi, Real rhs0, Real rhs1,
-                                  Array4<Real> const& acf, Real facx, Real facy)
-            {
-                Real a = acf(i,j,0);
-                res0 = residual1(i, j, 0, ilo, jlo, ihi, jhi, phi, rhs0, a, facx, facy);
-                res1 = residual1(i, j, 1, ilo, jlo, ihi, jhi, phi, rhs1, a, facx, facy);
-            });
-    } else {
-        bottomsolve_gpu<nsweeps>(dx0, dy0, acf, res, cor, rescor, nlevs, corner_offset,
-            [=] AMREX_GPU_DEVICE (int i, int j, int ilo, int jlo, int ihi, int jhi,
-                                  Array4<Real> const& phi, Real rhs0, Real rhs1,
-                                  Array4<Real> const& acf, Real facx, Real facy)
-            {
-                Real ar = acf(i,j,0,0);
-                Real ai = acf(i,j,0,1);
-                gs2(i, j, ilo, jlo, ihi, jhi, phi, rhs0, rhs1, ar, ai, facx, facy);
-            },
-            [=] AMREX_GPU_DEVICE (int i, int j, Real& res0, Real& res1,
-                                  int ilo, int jlo, int ihi, int jhi,
-                                  Array4<Real> const& phi, Real rhs_r, Real rhs_i,
-                                  Array4<Real> const& acf, Real facx, Real facy)
-            {
-                Real ar = acf(i,j,0,0);
-                Real ai = acf(i,j,0,1);
-                res0 = residual2r(i, j, ilo, jlo, ihi, jhi, phi, rhs_r, ar, ai, facx, facy);
-                res1 = residual2i(i, j, ilo, jlo, ihi, jhi, phi, rhs_i, ar, ai, facx, facy);
-            });
-    }
-#else
-    const int ilev = m_single_block_level_begin;
-    m_cor[ilev].setVal(Real(0.));
-    for (int is = 0; is < nsweeps; ++is) {
-        gsrb(is, m_domain[ilev], m_cor[ilev].array(),
-             m_res[ilev].const_array(), m_acf[ilev].const_array(), dx0, dy0,
-             m_system_type);
-    }
+    if (m_use_single_block_kernel) {
+        Array4<amrex::Real> const* acf = m_acf_a;
+        Array4<amrex::Real> const* res = m_res_a;
+        Array4<amrex::Real> const* cor = m_cor_a;
+        Array4<amrex::Real> const* rescor = m_rescor_a;
+        int nlevs = m_num_single_block_levels;
+        int const corner_offset = m_domain[0].cellCentered() ? 0 : 1;
+
+        if (m_system_type == 1) {
+            bottomsolve_gpu<nsweeps>(dx0, dy0, acf, res, cor, rescor, nlevs, corner_offset,
+                [=] AMREX_GPU_DEVICE (int i, int j, int ilo, int jlo, int ihi, int jhi,
+                                    Array4<Real> const& phi, Real rhs0, Real rhs1,
+                                    Array4<Real> const& acf, Real facx, Real facy)
+                {
+                    Real a = acf(i,j,0);
+                    gs1(i, j, 0, ilo, jlo, ihi, jhi, phi, rhs0, a, facx, facy);
+                    gs1(i, j, 1, ilo, jlo, ihi, jhi, phi, rhs1, a, facx, facy);
+                },
+                [=] AMREX_GPU_DEVICE (int i, int j, Real& res0, Real& res1,
+                                    int ilo, int jlo, int ihi, int jhi,
+                                    Array4<Real> const& phi, Real rhs0, Real rhs1,
+                                    Array4<Real> const& acf, Real facx, Real facy)
+                {
+                    Real a = acf(i,j,0);
+                    res0 = residual1(i, j, 0, ilo, jlo, ihi, jhi, phi, rhs0, a, facx, facy);
+                    res1 = residual1(i, j, 1, ilo, jlo, ihi, jhi, phi, rhs1, a, facx, facy);
+                });
+        } else {
+            bottomsolve_gpu<nsweeps>(dx0, dy0, acf, res, cor, rescor, nlevs, corner_offset,
+                [=] AMREX_GPU_DEVICE (int i, int j, int ilo, int jlo, int ihi, int jhi,
+                                    Array4<Real> const& phi, Real rhs0, Real rhs1,
+                                    Array4<Real> const& acf, Real facx, Real facy)
+                {
+                    Real ar = acf(i,j,0,0);
+                    Real ai = acf(i,j,0,1);
+                    gs2(i, j, ilo, jlo, ihi, jhi, phi, rhs0, rhs1, ar, ai, facx, facy);
+                },
+                [=] AMREX_GPU_DEVICE (int i, int j, Real& res0, Real& res1,
+                                    int ilo, int jlo, int ihi, int jhi,
+                                    Array4<Real> const& phi, Real rhs_r, Real rhs_i,
+                                    Array4<Real> const& acf, Real facx, Real facy)
+                {
+                    Real ar = acf(i,j,0,0);
+                    Real ai = acf(i,j,0,1);
+                    res0 = residual2r(i, j, ilo, jlo, ihi, jhi, phi, rhs_r, ar, ai, facx, facy);
+                    res1 = residual2i(i, j, ilo, jlo, ihi, jhi, phi, rhs_i, ar, ai, facx, facy);
+                });
+        }
+    } else
 #endif
+    {
+        const int ilev = m_single_block_level_begin;
+        m_cor[ilev].setVal<amrex::RunOn::Device>(Real(0.));
+        // Use numsweeps equal to the box length rounded up to an even number for large boxes
+        int numsweeps = std::max(nsweeps, (m_cor[ilev].box().length().max() + 1) / 2 * 2);
+        for (int is = 0; is < numsweeps; ++is) {
+            gsrb(is, m_domain[ilev], m_cor[ilev].array(),
+                m_res[ilev].const_array(), m_acf[ilev].const_array(), dx0, dy0,
+                m_system_type);
+        }
+    }
 }
 
 #if defined(AMREX_USE_GPU)

src/particles/beam/BeamParticleContainer.cpp

@@ -378,8 +378,7 @@ BeamParticleContainer::intializeSlice (int slice, int which_slice) {
                 ptd.rdata(BeamIdx::uy)[ip] = ptd_init.rdata(BeamIdx::uy)[idx_src];
                 ptd.rdata(BeamIdx::uz)[ip] = ptd_init.rdata(BeamIdx::uz)[idx_src];
 
-                ptd.id(ip) = ptd_init.id(idx_src);
-                ptd.cpu(ip) = 0;
+                ptd.idcpu(ip) = ptd_init.idcpu(idx_src);
                 ptd.idata(BeamIdx::nsubcycles)[ip] = 0;
             }
         );
@@ -472,7 +471,7 @@ BeamParticleContainer::InSituComputeDiags (int islice)
 
             const amrex::Real uz_inv = uz == 0._rt ? 0._rt : 1._rt / uz;
 
-            if (ptd.id(ip) < 0 || x*x + y*y > insitu_radius_sq) {
+            if (!ptd.id(ip).is_valid() || x*x + y*y > insitu_radius_sq) {
                 return amrex::IdentityTuple(ReduceTuple{}, reduce_op);
             }
             const amrex::Real gamma = std::sqrt(1.0_rt + ux*ux + uy*uy + uz*uz);

src/particles/beam/BeamParticleContainerInit.cpp

@@ -53,6 +53,7 @@ namespace
         ptd.rdata(BeamIdx::w  )[ip] = std::abs(weight);
 
         ptd.id(ip) = pid > 0 ? pid + ip : pid;
+        ptd.cpu(ip) = 0;
     }
 
     /** \brief Adds a single beam particle into the per-slice BeamTile

src/particles/deposition/BeamDepositCurrent.cpp

@@ -99,7 +99,7 @@ DepositCurrentSlice (BeamParticleContainer& beam, Fields& fields,
 
             // Skip invalid particles and ghost particles not in the last slice
             // beam deposits only up to its finest level
-            if (ptd.id(ip) < 0 || (only_highest ? (ptd.cpu(ip)!=lev) : (ptd.cpu(ip)<lev))) return;
+            if (!ptd.id(ip).is_valid() || (only_highest ? (ptd.cpu(ip)!=lev) : (ptd.cpu(ip)<lev))) return;
 
             // --- Get particle quantities
             const amrex::Real ux = ptd.rdata(BeamIdx::ux)[ip];

src/particles/deposition/ExplicitDeposition.cpp

@@ -103,7 +103,7 @@ ExplicitDeposition (PlasmaParticleContainer& plasma, Fields& fields, const Multi
 
                 ip += idx_begin;
 
-                if (ptd.id(ip) < 0 || (lev != 0 && ptd.cpu(ip) < lev)) return;
+                if (!ptd.id(ip).is_valid() || (lev != 0 && ptd.cpu(ip) < lev)) return;
                 const amrex::Real psi_inv = 1._rt/ptd.rdata(PlasmaIdx::psi)[ip];
                 const amrex::Real xp = ptd.pos(0, ip);
                 const amrex::Real yp = ptd.pos(1, ip);

src/particles/deposition/PlasmaDepositCurrent.cpp

@@ -190,7 +190,7 @@ DepositCurrent (PlasmaParticleContainer& plasma, Fields & fields, const MultiLas
                 const int ox = idx % outer_depos_order_x_1;
 
                 // only deposit plasma currents on or below their according MR level
-                if (ptd.id(ip) < 0 || (lev != 0 && ptd.cpu(ip) < lev)) return;
+                if (!ptd.id(ip).is_valid() || (lev != 0 && ptd.cpu(ip) < lev)) return;
 
                 const amrex::Real psi_inv = 1._rt/ptd.rdata(PlasmaIdx::psi)[ip];
                 const amrex::Real xp = ptd.pos(0, ip);
@@ -233,7 +233,7 @@ DepositCurrent (PlasmaParticleContainer& plasma, Fields & fields, const MultiLas
                     // This particle violates the QSA, discard it and do not deposit its current
                     amrex::Gpu::Atomic::Add(p_n_qsa_violation, 1);
                     ptd.rdata(PlasmaIdx::w)[ip] = 0.0_rt;
-                    ptd.id(ip) = -std::abs(ptd.id(ip));
+                    ptd.id(ip).make_invalid();
                     return;
                 }
 

src/particles/plasma/PlasmaParticleContainer.cpp

@@ -484,7 +484,7 @@ PlasmaParticleContainer::InSituComputeDiags (int islice)
                 const amrex::Real uy = ptd.rdata(PlasmaIdx::uy)[ip] * clight_inv;
                 const amrex::Real psi = ptd.rdata(PlasmaIdx::psi)[ip];
 
-                if (ptd.id(ip) < 0 || x*x + y*y > insitu_radius_sq) {
+                if (!ptd.id(ip).is_valid() || x*x + y*y > insitu_radius_sq) {
                     return amrex::IdentityTuple(ReduceTuple{}, reduce_op);
                 }
                 // particle's Lorentz factor

src/particles/pusher/BeamParticleAdvance.cpp

@@ -129,7 +129,7 @@ AdvanceBeamParticlesSlice (
         beam.getNumParticlesIncludingSlipped(WhichBeamSlice::This),
         [=] AMREX_GPU_DEVICE (int ip, auto depos_order, auto c_use_external_fields) {
 
-            if (ptd.id(ip) < 0) return;
+            if (!ptd.id(ip).is_valid()) return;
 
             amrex::Real xp = ptd.pos(0, ip);
             amrex::Real yp = ptd.pos(1, ip);

src/particles/pusher/GetAndSetPosition.H

@@ -72,7 +72,7 @@ struct EnforceBCandSetPos
         const bool invalid = (shifted && !m_is_per);
         if (invalid) {
             ptd.rdata(PlasmaIdx::w)[ip] = 0.0_rt;
-            p.id() = -std::abs(p.id());
+            p.id().make_invalid();
         }
         x = p.pos(0);
         y = p.pos(1);
@@ -103,7 +103,7 @@ struct EnforceBCandSetPos
         const bool invalid = (shifted && !m_is_per);
         if (invalid) {
             ptd.rdata(BeamIdx::w)[ip] = 0.0_rt;
-            p.id() = -std::abs(p.id());
+            p.id().make_invalid();
         }
         x = p.pos(0);
         y = p.pos(1);

src/particles/pusher/PlasmaParticleAdvance.cpp

@@ -107,7 +107,7 @@ AdvancePlasmaParticles (PlasmaParticleContainer& plasma, const Fields & fields,
                 [=] AMREX_GPU_DEVICE (int idx, auto depos_order, auto use_laser) {
                     const int ip = idx + idx_begin;
                     // only push plasma particles on their according MR level
-                    if (ptd.id(ip) < 0 || ptd.cpu(ip) != lev) return;
+                    if (!ptd.id(ip).is_valid() || ptd.cpu(ip) != lev) return;
 
                     // define field at particle position reals
                     amrex::Real ExmByp = 0._rt, EypBxp = 0._rt, Ezp = 0._rt;