Add embedded BC in RZ. (#2602)

* Start to add embedded BC in RZ.

* Add .

* Remove scaling

* Fix compilation error

* Update

* Can compile.

* Add call linop.setRZ(true).

* Remove lines 264 to 312 and 343 to 345.

* Add assert.

* Remove an assert.

* Add an automated test.

* Change to MLEBNodeFDLaplacian.

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

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

* Fix compilation error

* Update the test selection

* Correct compilation error

* Move test to another worker

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

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

* Add new required argument

* Update Examples/Tests/ElectrostaticSphereEB/inputs_rz

* Update Examples/Tests/ElectrostaticSphereEB/analysis_rz.py

* Update analysis script

Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

11 files changed, 184 insertions, 74 deletions

diff --git a/Examples/Tests/ElectrostaticSphereEB/analysis_rz.py b/Examples/Tests/ElectrostaticSphereEB/analysis_rz.py
new file mode 100755
index 000000000..bb170f921
--- /dev/null
+++ b/Examples/Tests/ElectrostaticSphereEB/analysis_rz.py
@@ -0,0 +1,67 @@
+#!/usr/bin/env python
+
+# Copyright 2019-2021 Yinjian Zhao
+#
+# This file is part of WarpX.
+#
+# License: BSD-3-Clause-LBNL
+
+# This script tests the embedded boundary in RZ.
+# A cylindrical surface (r=0.1) has a fixed potential 1 V.
+# The outer surface has 0 V fixed.
+# Thus the analytical solution has the form:
+# phi(r) = A+B*log(r), Er(r) = -B/r.
+
+# Possible errors:
+# tolerance: 0.004
+# Possible running time: < 1 s
+
+import os
+import sys
+
+import numpy as np
+import yt
+
+sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
+import checksumAPI
+
+tolerance = 0.004
+
+fn = sys.argv[1]
+ds = yt.load( fn )
+
+all_data_level_0 = ds.covering_grid(level=0,left_edge=ds.domain_left_edge, dims=ds.domain_dimensions)
+phi = all_data_level_0['boxlib', 'phi'].v.squeeze()
+Er = all_data_level_0['boxlib', 'Ex'].v.squeeze()
+
+Dx = ds.domain_width/ds.domain_dimensions
+dr = Dx[0]
+rmin = ds.domain_left_edge[0]
+rmax = ds.domain_right_edge[0]
+nr = phi.shape[0]
+r = np.linspace(rmin+dr/2.,rmax-dr/2.,nr)
+B = 1.0/np.log(0.1/0.5)
+A = -B*np.log(0.5)
+
+err = 0.0
+errmax_phi = 0.0
+errmax_Er = 0.0
+for i in range(len(r)):
+    if r[i]>=0.1:
+        phi_theory = A+B*np.log(r[i])
+        Er_theory = -B/float(r[i])
+        err = abs( phi_theory - phi[i,:] ).max() / phi_theory
+        if err>errmax_phi:
+            errmax_phi = err
+        err = abs( Er_theory - Er[i,:] ).max() / Er_theory
+        # Exclude the last inaccurate interpolation.
+        if err>errmax_Er and i<len(r)-1:
+            errmax_Er = err
+
+print('max error of phi = ', errmax_phi)
+print('max error of Er = ', errmax_Er)
+print('tolerance = ', tolerance)
+assert(errmax_phi<tolerance and errmax_Er<tolerance)
+
+test_name = os.path.split(os.getcwd())[1]
+checksumAPI.evaluate_checksum(test_name, fn, do_particles=False)
diff --git a/Examples/Tests/ElectrostaticSphereEB/inputs_rz b/Examples/Tests/ElectrostaticSphereEB/inputs_rz
new file mode 100644
index 000000000..af5f7fcbf
--- /dev/null
+++ b/Examples/Tests/ElectrostaticSphereEB/inputs_rz
@@ -0,0 +1,29 @@
+max_step = 1
+amr.n_cell = 64 64
+amr.max_level = 0
+amr.blocking_factor = 8
+amr.max_grid_size = 128
+boundary.field_lo = none periodic
+boundary.field_hi = pec periodic
+boundary.potential_lo_x = 0
+boundary.potential_hi_x = 0
+boundary.potential_lo_y = 0
+boundary.potential_hi_y = 0
+boundary.potential_lo_z = 0
+boundary.potential_hi_z = 0
+geometry.dims = RZ
+geometry.prob_lo     = 0.0 -0.5
+geometry.prob_hi     = 0.5  0.5
+warpx.const_dt = 1e-6
+
+warpx.do_electrostatic = labframe
+warpx.eb_implicit_function = "-(x**2-0.1**2)"
+warpx.eb_potential(x,y,z,t) = "1."
+warpx.self_fields_required_precision = 1.e-7
+
+algo.field_gathering = momentum-conserving
+
+diagnostics.diags_names = diag1
+diag1.intervals = 1
+diag1.diag_type = Full
+diag1.fields_to_plot = Ex phi
diff --git a/Regression/Checksum/benchmarks_json/ElectrostaticSphereEB_RZ.json b/Regression/Checksum/benchmarks_json/ElectrostaticSphereEB_RZ.json
new file mode 100644
index 000000000..0599ba186
--- /dev/null
+++ b/Regression/Checksum/benchmarks_json/ElectrostaticSphereEB_RZ.json
@@ -0,0 +1,6 @@
+{
+  "lev=0": {
+    "Ex": 8497.669853722688,
+    "phi": 1235.661848443515
+  }
+}
diff --git a/Regression/WarpX-tests.ini b/Regression/WarpX-tests.ini
index a74dbdca8..e31614465 100644
--- a/Regression/WarpX-tests.ini
+++ b/Regression/WarpX-tests.ini
@@ -2206,6 +2206,22 @@ doVis = 0
 compareParticles = 0
 analysisRoutine = Examples/analysis_default_regression.py
 
+[ElectrostaticSphereEB_RZ]
+buildDir = .
+inputFile = Examples/Tests/ElectrostaticSphereEB/inputs_rz
+runtime_params =
+dim = 2
+addToCompileString = USE_EB=TRUE USE_RZ=TRUE
+restartTest = 0
+useMPI = 1
+numprocs = 2
+useOMP = 1
+numthreads = 1
+compileTest = 0
+doVis = 0
+compareParticles = 0
+analysisRoutine = Examples/Tests/ElectrostaticSphereEB/analysis_rz.py
+
 [Python_ElectrostaticSphereEB]
 buildDir = .
 inputFile = Examples/Tests/ElectrostaticSphereEB/PICMI_inputs_3d.py
diff --git a/Regression/prepare_file_ci.py b/Regression/prepare_file_ci.py
index 03accd185..142263866 100644
--- a/Regression/prepare_file_ci.py
+++ b/Regression/prepare_file_ci.py
@@ -154,6 +154,7 @@ if ci_qed:
     test_blocks = select_tests(test_blocks, ['QED=TRUE'], True)
 
 if ci_eb:
+    test_blocks = select_tests(test_blocks, ['USE_RZ=TRUE'], False)
     test_blocks = select_tests(test_blocks, ['USE_EB=TRUE'], True)
 
 # - Add the selected test blocks to the text
diff --git a/Source/EmbeddedBoundary/WarpXFaceExtensions.cpp b/Source/EmbeddedBoundary/WarpXFaceExtensions.cpp
index b52a392af..cde8d1dc2 100644
--- a/Source/EmbeddedBoundary/WarpXFaceExtensions.cpp
+++ b/Source/EmbeddedBoundary/WarpXFaceExtensions.cpp
@@ -141,6 +141,7 @@ amrex::Array1D<int, 0, 2>
 WarpX::CountExtFaces() {
     amrex::Array1D<int, 0, 2> sums{0, 0, 0};
 #ifdef AMREX_USE_EB
+#ifndef WARPX_DIM_RZ
 
 #ifdef WARPX_DIM_XZ
     // In 2D we change the extrema of the for loop so that we only have the case idim=1
@@ -167,6 +168,7 @@ WarpX::CountExtFaces() {
 
     amrex::ParallelDescriptor::ReduceIntSum(&(sums(0)), AMREX_SPACEDIM);
 #endif
+#endif
     return sums;
 }
 
@@ -369,6 +371,7 @@ ComputeNBorrowEightFacesExtension(const amrex::Dim3 cell, const amrex::Real S_ex
 void
 WarpX::ComputeOneWayExtensions() {
 #ifdef AMREX_USE_EB
+#ifndef WARPX_DIM_RZ
     auto const eb_fact = fieldEBFactory(maxLevel());
 
     auto const &cell_size = CellSize(maxLevel());
@@ -484,12 +487,14 @@ WarpX::ComputeOneWayExtensions() {
     }
 
 #endif
+#endif
 }
 
 
 void
 WarpX::ComputeEightWaysExtensions() {
 #ifdef AMREX_USE_EB
+#ifndef WARPX_DIM_RZ
     auto const &cell_size = CellSize(maxLevel());
 
     const amrex::Real dx = cell_size[0];
@@ -640,6 +645,7 @@ WarpX::ComputeEightWaysExtensions() {
         }
     }
 #endif
+#endif
 }
 
 
diff --git a/Source/EmbeddedBoundary/WarpXInitEB.cpp b/Source/EmbeddedBoundary/WarpXInitEB.cpp
index 1c1931837..8a62308ba 100644
--- a/Source/EmbeddedBoundary/WarpXInitEB.cpp
+++ b/Source/EmbeddedBoundary/WarpXInitEB.cpp
@@ -82,10 +82,6 @@ WarpX::InitEB ()
 #ifdef AMREX_USE_EB
     BL_PROFILE("InitEB");
 
-#if !(defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ))
-    amrex::Abort("InitEB: Embedded Boundaries are only implemented in 2D3V and 3D3V");
-#endif
-
     amrex::ParmParse pp_warpx("warpx");
     std::string impf;
     pp_warpx.query("eb_implicit_function", impf);
@@ -117,6 +113,7 @@ WarpX::InitEB ()
 void
 WarpX::ComputeEdgeLengths () {
 #ifdef AMREX_USE_EB
+#ifndef WARPX_DIM_RZ
     BL_PROFILE("ComputeEdgeLengths");
 
     auto const eb_fact = fieldEBFactory(maxLevel());
@@ -178,12 +175,14 @@ WarpX::ComputeEdgeLengths () {
         }
     }
 #endif
+#endif
 }
 
 
 void
 WarpX::ComputeFaceAreas () {
 #ifdef AMREX_USE_EB
+#ifndef WARPX_DIM_RZ
     BL_PROFILE("ComputeFaceAreas");
 
     auto const eb_fact = fieldEBFactory(maxLevel());
@@ -239,12 +238,14 @@ WarpX::ComputeFaceAreas () {
         }
     }
 #endif
+#endif
 }
 
 
 void
 WarpX::ScaleEdges () {
 #ifdef AMREX_USE_EB
+#ifndef WARPX_DIM_RZ
     BL_PROFILE("ScaleEdges");
 
     auto const &cell_size = CellSize(maxLevel());
@@ -269,11 +270,13 @@ WarpX::ScaleEdges () {
         }
     }
 #endif
+#endif
 }
 
 void
 WarpX::ScaleAreas() {
 #ifdef AMREX_USE_EB
+#ifndef WARPX_DIM_RZ
     BL_PROFILE("ScaleAreas");
 
     auto const& cell_size = CellSize(maxLevel());
@@ -321,12 +324,14 @@ WarpX::ScaleAreas() {
         }
     }
 #endif
+#endif
 }
 
 
 void
 WarpX::MarkCells(){
 #ifdef AMREX_USE_EB
+#ifndef WARPX_DIM_RZ
     auto const &cell_size = CellSize(maxLevel());
 
 #ifdef WARPX_DIM_3D
@@ -404,6 +409,7 @@ WarpX::MarkCells(){
         }
     }
 #endif
+#endif
 }
 
 
diff --git a/Source/FieldSolver/ElectrostaticSolver.H b/Source/FieldSolver/ElectrostaticSolver.H
index 725407407..f791379ce 100644
--- a/Source/FieldSolver/ElectrostaticSolver.H
+++ b/Source/FieldSolver/ElectrostaticSolver.H
@@ -10,13 +10,9 @@
 #include "Utils/WarpXUtil.H"
 
 #include <AMReX_Array.H>
-#ifdef WARPX_DIM_RZ
-#    include <AMReX_MLNodeLaplacian.H>
-#else
-#    include <AMReX_MLNodeTensorLaplacian.H>
-#    ifdef AMREX_USE_EB
-#        include <AMReX_MLEBNodeFDLaplacian.H>
-#    endif
+#include <AMReX_MLNodeTensorLaplacian.H>
+#if defined(AMREX_USE_EB) || defined(WARPX_DIM_RZ)
+#    include <AMReX_MLEBNodeFDLaplacian.H>
 #endif
 #include <AMReX_Parser.H>
 #include <AMReX_REAL.H>
diff --git a/Source/FieldSolver/ElectrostaticSolver.cpp b/Source/FieldSolver/ElectrostaticSolver.cpp
index b8cfe3ca3..98cbd2946 100644
--- a/Source/FieldSolver/ElectrostaticSolver.cpp
+++ b/Source/FieldSolver/ElectrostaticSolver.cpp
@@ -261,56 +261,6 @@ WarpX::computePhiRZ (const amrex::Vector<std::unique_ptr<amrex::MultiFab> >& rho
         geom_scaled[lev].define(geom_lev.Domain(), &rb);
     }
 
-    // Setup the sigma = radius
-    // sigma must be cell centered
-    amrex::Vector<std::unique_ptr<amrex::MultiFab> > sigma(max_level+1);
-    for (int lev = 0; lev <= max_level; ++lev) {
-        const amrex::Real * problo = geom_scaled[lev].ProbLo();
-        const amrex::Real * dx = geom_scaled[lev].CellSize();
-        const amrex::Real rmin = problo[0];
-        const amrex::Real dr = dx[0];
-
-        amrex::BoxArray nba = boxArray(lev);
-        nba.enclosedCells(); // Get cell centered array (correct?)
-        sigma[lev] = std::make_unique<MultiFab>(nba, DistributionMap(lev), 1, 0);
-        for ( MFIter mfi(*sigma[lev], TilingIfNotGPU()); mfi.isValid(); ++mfi )
-        {
-            const amrex::Box& tbx = mfi.tilebox();
-            const amrex::Dim3 lo = amrex::lbound(tbx);
-            const int irmin = lo.x;
-            Array4<amrex::Real> const& sigma_arr = sigma[lev]->array(mfi);
-            amrex::ParallelFor( tbx,
-                [=] AMREX_GPU_DEVICE (int i, int j, int /*k*/) {
-                    sigma_arr(i,j,0) = rmin + (i - irmin + 0.5_rt)*dr;
-                }
-            );
-        }
-
-        // Also, multiply rho by radius (rho is node centered)
-        // Note that this multiplication is not undone since rho is
-        // a temporary array.
-        for ( MFIter mfi(*rho[lev], TilingIfNotGPU()); mfi.isValid(); ++mfi )
-        {
-            const amrex::Box& tbx = mfi.tilebox();
-            const amrex::Dim3 lo = amrex::lbound(tbx);
-            const int irmin = lo.x;
-            int const ncomp = rho[lev]->nComp(); // This should be 1!
-            Array4<Real> const& rho_arr = rho[lev]->array(mfi);
-            amrex::ParallelFor(tbx, ncomp,
-            [=] AMREX_GPU_DEVICE (int i, int j, int /*k*/, int icomp)
-            {
-                amrex::Real r = rmin + (i - irmin)*dr;
-                if (r == 0.) {
-                    // dr/3 is used to be consistent with the finite volume formulism
-                    // that is used to solve Poisson's equation
-                    rho_arr(i,j,0,icomp) *= dr/3._rt;
-                } else {
-                    rho_arr(i,j,0,icomp) *= r;
-                }
-            });
-        }
-    }
-
     // get the potential at the current time
     amrex::Array<amrex::Real,AMREX_SPACEDIM> phi_bc_values_lo;
     amrex::Array<amrex::Real,AMREX_SPACEDIM> phi_bc_values_hi;
@@ -320,15 +270,6 @@ WarpX::computePhiRZ (const amrex::Vector<std::unique_ptr<amrex::MultiFab> >& rho
     // set the boundary potential values if needed
     setPhiBC(phi, phi_bc_values_lo, phi_bc_values_hi);
 
-    // Define the linear operator (Poisson operator)
-    MLNodeLaplacian linop( geom_scaled, boxArray(), DistributionMap() );
-
-    linop.setRZCorrection(true);
-
-    for (int lev = 0; lev <= max_level; ++lev) {
-        linop.setSigma( lev, *sigma[lev] );
-    }
-
     amrex::Real max_norm_b = 0.0;
     for (int lev=0; lev < rho.size(); lev++){
         rho[lev]->mult(-1._rt/PhysConst::ep0);
@@ -344,14 +285,55 @@ WarpX::computePhiRZ (const amrex::Vector<std::unique_ptr<amrex::MultiFab> >& rho
         );
     }
 
+    LPInfo info;
+
+#ifndef AMREX_USE_EB
+    // Define the linear operator (Poisson operator)
+    MLEBNodeFDLaplacian linop( geom_scaled, boxArray(), DistributionMap(), info );
+#else
+    // With embedded boundary: extract EB info
+    Vector<EBFArrayBoxFactory const*> eb_factory;
+    eb_factory.resize(max_level+1);
+    for (int lev = 0; lev <= max_level; ++lev) {
+        eb_factory[lev] = &WarpX::fieldEBFactory(lev);
+    }
+    MLEBNodeFDLaplacian linop( Geom(), boxArray(), DistributionMap(), info, eb_factory);
+    // if the EB potential only depends on time, the potential can be passed
+    // as a float instead of a callable
+    if (field_boundary_handler.phi_EB_only_t) {
+        linop.setEBDirichlet(field_boundary_handler.potential_eb_t(gett_new(0)));
+    }
+    else linop.setEBDirichlet(field_boundary_handler.getPhiEB(gett_new(0)));
+
+#endif
+
     // Solve the Poisson equation
     linop.setDomainBC( field_boundary_handler.lobc, field_boundary_handler.hibc );
+    linop.setRZ(true);
     MLMG mlmg(linop);
     mlmg.setVerbose(verbosity);
     mlmg.setMaxIter(max_iters);
     mlmg.setAlwaysUseBNorm(always_use_bnorm);
     mlmg.solve( GetVecOfPtrs(phi), GetVecOfConstPtrs(rho),
                 required_precision, absolute_tolerance );
+
+#ifdef AMREX_USE_EB
+    // use amrex to directly calculate the electric field since with EB's the
+    // simple finite difference scheme in WarpX::computeE sometimes fails
+    if (do_electrostatic == ElectrostaticSolverAlgo::LabFrame)
+    {
+        for (int lev = 0; lev <= max_level; ++lev) {
+            mlmg.getGradSolution(
+                {amrex::Array<amrex::MultiFab*,2>{
+                    get_pointer_Efield_fp(lev, 0),get_pointer_Efield_fp(lev, 2)
+                    }}
+            );
+            get_pointer_Efield_fp(lev, 0)->mult(-1._rt);
+            get_pointer_Efield_fp(lev, 2)->mult(-1._rt);
+        }
+    }
+#endif
+
 }
 
 #else
diff --git a/Source/WarpX.H b/Source/WarpX.H
index 2e6c46173..a5f30188a 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -65,10 +65,6 @@
 #include <string>
 #include <vector>
 
-#if defined(AMREX_USE_EB) && defined(WARPX_DIM_RZ)
-static_assert(false, "Embedded boundaries are not supported in RZ mode.");
-#endif
-
 enum struct PatchType : int
 {
     fine,
diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp
index 785debe79..585369d70 100644
--- a/Source/WarpX.cpp
+++ b/Source/WarpX.cpp
@@ -630,6 +630,11 @@ WarpX::ReadParameters ()
 
         do_electrostatic = GetAlgorithmInteger(pp_warpx, "do_electrostatic");
 
+#if defined(AMREX_USE_EB) && defined(WARPX_DIM_RZ)
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(do_electrostatic!=ElectrostaticSolverAlgo::None,
+        "Currently, the embedded boundary in RZ only works for electrostatic solvers.");
+#endif
+
         if (do_electrostatic == ElectrostaticSolverAlgo::LabFrame) {
             // Note that with the relativistic version, these parameters would be
             // input for each species.