Merge pull request #23 from burlen/sensei_insitu_dev

Sensei insitu dev

6 files changed, 376 insertions, 10 deletions

diff --git a/Source/Make.WarpX b/Source/Make.WarpX
index 44001eeaa..9dfe667e0 100644
--- a/Source/Make.WarpX
+++ b/Source/Make.WarpX
@@ -31,6 +31,11 @@ include $(AMREX_HOME)/Src/Particle/Make.package
 include $(AMREX_HOME)/Src/Boundary/Make.package
 include $(AMREX_HOME)/Src/AmrCore/Make.package
 
+ifeq ($(USE_SENSEI_INSITU),TRUE)
+  include $(AMREX_HOME)/Src/Amr/Make.package
+  include $(AMREX_HOME)/Src/Extern/SENSEI/Make.package
+endif
+
 include $(PICSAR_HOME)/src/Make.package
 
 DEFINES += -DPICSAR_NO_ASSUMED_ALIGNMENT
diff --git a/Source/WarpX.H b/Source/WarpX.H
index 32df63e19..b3c382cc7 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -27,6 +27,12 @@
 #include <fftw3.h>
 #endif
 
+#if defined(BL_USE_SENSEI_INSITU)
+namespace amrex {
+class AmrMeshInSituBridge;
+}
+#endif
+
 enum struct DtType : int
 {
     Full = 0,
@@ -170,6 +176,7 @@ public:
 
     void WriteCheckPointFile () const;
     void WritePlotFile () const;
+    void UpdateInSitu () const;
 
     void WritePlotFileES(const amrex::Vector<std::unique_ptr<amrex::MultiFab> >& rho,
                          const amrex::Vector<std::unique_ptr<amrex::MultiFab> >& phi,
@@ -568,6 +575,13 @@ private:
     void PushPSATD (int lev, amrex::Real dt);
 
 #endif
+
+#if defined(BL_USE_SENSEI_INSITU)
+    amrex::AmrMeshInSituBridge *insitu_bridge;
+#endif
+    int insitu_int;
+    int insitu_start;
+    std::string insitu_config;
 };
 
 #endif
diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp
index ae39f073c..44eff9cd4 100644
--- a/Source/WarpX.cpp
+++ b/Source/WarpX.cpp
@@ -18,6 +18,10 @@
 #include <WarpXWrappers.h>
 #include <WarpXUtil.H>
 
+#ifdef BL_USE_SENSEI_INSITU
+#include <AMReX_AmrMeshInSituBridge.H>
+#endif
+
 using namespace amrex;
 
 Vector<Real> WarpX::B_external(3, 0.0);
@@ -197,6 +201,10 @@ WarpX::WarpX ()
     comm_fft.resize(nlevs_max,MPI_COMM_NULL);
     color_fft.resize(nlevs_max,-1);
 #endif
+
+#ifdef BL_USE_SENSEI_INSITU
+    insitu_bridge = nullptr;
+#endif
 }
 
 WarpX::~WarpX ()
@@ -205,6 +213,10 @@ WarpX::~WarpX ()
     for (int lev = 0; lev < nlevs_max; ++lev) {
         ClearLevel(lev);
     }
+
+#ifdef BL_USE_SENSEI_INSITU
+    delete insitu_bridge;
+#endif
 }
 
 void
@@ -433,6 +445,17 @@ WarpX::ReadParameters ()
         pp.query("noz", noz_fft);
     }
 #endif
+
+    {
+        insitu_start = 0;
+        insitu_int = 0;
+        insitu_config = "";
+
+        ParmParse pp("insitu");
+        pp.query("int", insitu_int);
+        pp.query("start", insitu_start);
+        pp.query("config", insitu_config);
+    }
 }
 
 // This is a virtual function.
diff --git a/Source/WarpXEvolve.cpp b/Source/WarpXEvolve.cpp
index 092d5adfd..4a77fde43 100644
--- a/Source/WarpXEvolve.cpp
+++ b/Source/WarpXEvolve.cpp
@@ -9,6 +9,10 @@
 #include <WarpX_py.H>
 #endif
 
+#ifdef BL_USE_SENSEI_INSITU
+#include <AMReX_AmrMeshInSituBridge.H>
+#endif
+
 using namespace amrex;
 
 void
@@ -35,6 +39,7 @@ WarpX::EvolveEM (int numsteps)
     Real cur_time = t_new[0];
     static int last_plot_file_step = 0;
     static int last_check_file_step = 0;
+    static int last_insitu_step = 0;
 
     int numsteps_max;
     if (numsteps < 0) {  // Note that the default argument is numsteps = -1
@@ -133,7 +138,10 @@ WarpX::EvolveEM (int numsteps)
 
         bool to_make_plot = (plot_int > 0) && ((step+1) % plot_int == 0);
 
-        bool move_j = is_synchronized || to_make_plot;
+        bool do_insitu = ((step+1) >= insitu_start) &&
+             (insitu_int > 0) && ((step+1) % insitu_int == 0);
+
+        bool move_j = is_synchronized || to_make_plot || do_insitu;
         // If is_synchronized we need to shift j too so that next step we can evolve E by dt/2.
         // We might need to move j because we are going to make a plotfile.
 	MoveWindow(move_j);
@@ -166,7 +174,7 @@ WarpX::EvolveEM (int numsteps)
             myBFD->writeLabFrameData(cell_centered_data.get(), *mypc, geom[0], cur_time, dt[0]);
         }
 
-	if (to_make_plot)
+	if (to_make_plot || do_insitu)
         {
             FillBoundaryE();
             FillBoundaryB();
@@ -179,9 +187,16 @@ WarpX::EvolveEM (int numsteps)
             }
 
 	    last_plot_file_step = step+1;
-	    WritePlotFile();
+	    last_insitu_step = step+1;
+
+        if (to_make_plot)
+    	    WritePlotFile();
+
+        if (do_insitu)
+            UpdateInSitu();
 	}
 
+
 	if (check_int > 0 && (step+1) % check_int == 0) {
 	    last_check_file_step = step+1;
 	    WriteCheckPointFile();
@@ -198,7 +213,12 @@ WarpX::EvolveEM (int numsteps)
 	// End loop on time steps
     }
 
-    if (plot_int > 0 && istep[0] > last_plot_file_step && (max_time_reached || istep[0] >= max_step))
+    bool write_plot_file = plot_int > 0 && istep[0] > last_plot_file_step && (max_time_reached || istep[0] >= max_step);
+
+    bool do_insitu = (insitu_start >= istep[0]) && (insitu_int > 0) &&
+        (istep[0] > last_insitu_step) && (max_time_reached || istep[0] >= max_step);
+
+    if (write_plot_file || do_insitu)
     {
         FillBoundaryE();
         FillBoundaryB();
@@ -210,7 +230,11 @@ WarpX::EvolveEM (int numsteps)
                               *Bfield_aux[lev][0],*Bfield_aux[lev][1],*Bfield_aux[lev][2]);
         }
 
-	WritePlotFile();
+        if (write_plot_file)
+            WritePlotFile();
+
+        if (do_insitu)
+            UpdateInSitu();
     }
 
     if (check_int > 0 && istep[0] > last_check_file_step && (max_time_reached || istep[0] >= max_step)) {
@@ -220,6 +244,10 @@ WarpX::EvolveEM (int numsteps)
     if (do_boosted_frame_diagnostic) {
         myBFD->Flush(geom[0]);
     }
+
+#ifdef BL_USE_SENSEI_INSITU
+    insitu_bridge->finalize();
+#endif
 }
 
 /* /brief Perform one PIC iteration, without subcycling
diff --git a/Source/WarpXIO.cpp b/Source/WarpXIO.cpp
index d59720a33..4c673b1f5 100644
--- a/Source/WarpXIO.cpp
+++ b/Source/WarpXIO.cpp
@@ -7,6 +7,10 @@
 
 #include "AMReX_buildInfo.H"
 
+#ifdef BL_USE_SENSEI_INSITU
+#include <AMReX_AmrMeshInSituBridge.H>
+#endif
+
 using namespace amrex;
 
 namespace
@@ -454,6 +458,277 @@ WarpX::GetCellCenteredData() {
 }
 
 void
+WarpX::UpdateInSitu () const
+{
+#ifdef BL_USE_SENSEI_INSITU
+    BL_PROFILE("WarpX::UpdateInSitu()");
+
+    int numLevels = finest_level + 1;
+    Vector<std::string> varnames;
+    Vector<MultiFab> mf;
+
+    const int ncomp = 3*3
+        + static_cast<int>(plot_part_per_cell)
+        + static_cast<int>(plot_part_per_grid)
+        + static_cast<int>(plot_part_per_proc)
+        + static_cast<int>(plot_proc_number)
+        + static_cast<int>(plot_divb)
+        + static_cast<int>(plot_dive)
+        + static_cast<int>(plot_finepatch)*6
+        + static_cast<int>(plot_crsepatch)*6
+        + static_cast<int>(costs[0] != nullptr);
+
+    for (int lev = 0; lev <= finest_level; ++lev)
+    {
+        const int ngrow = 1;
+        mf.push_back(MultiFab(grids[lev], dmap[lev], ncomp, ngrow));
+
+        Vector<const MultiFab*> srcmf(AMREX_SPACEDIM);
+        PackPlotDataPtrs(srcmf, current_fp[lev]);
+        int dcomp = 0;
+        amrex::average_edge_to_cellcenter(mf[lev], dcomp, srcmf);
+
+#if (AMREX_SPACEDIM == 2)
+        MultiFab::Copy(mf[lev], mf[lev], dcomp+1, dcomp+2, 1, ngrow);
+            amrex::average_node_to_cellcenter(mf[lev], dcomp+1, *current_fp[lev][1], 0, 1);
+#endif
+        if (lev == 0)
+        {
+            varnames.push_back("jx");
+            varnames.push_back("jy");
+            varnames.push_back("jz");
+        }
+        dcomp += 3;
+
+        PackPlotDataPtrs(srcmf, Efield_aux[lev]);
+        amrex::average_edge_to_cellcenter(mf[lev], dcomp, srcmf);
+#if (AMREX_SPACEDIM == 2)
+        MultiFab::Copy(mf[lev], mf[lev], dcomp+1, dcomp+2, 1, ngrow);
+            amrex::average_node_to_cellcenter(mf[lev], dcomp+1, *Efield_aux[lev][1], 0, 1);
+#endif
+        if (lev == 0)
+        {
+            varnames.push_back("Ex");
+            varnames.push_back("Ey");
+            varnames.push_back("Ez");
+        }
+        dcomp += 3;
+
+        PackPlotDataPtrs(srcmf, Bfield_aux[lev]);
+        amrex::average_face_to_cellcenter(mf[lev], dcomp, srcmf);
+#if (AMREX_SPACEDIM == 2)
+        MultiFab::Copy(mf[lev], mf[lev], dcomp+1, dcomp+2, 1, ngrow);
+            MultiFab::Copy(mf[lev], *Bfield_aux[lev][1], 0, dcomp+1, 1, ngrow);
+#endif
+        if (lev == 0)
+        {
+            varnames.push_back("Bx");
+            varnames.push_back("By");
+            varnames.push_back("Bz");
+        }
+        dcomp += 3;
+
+        if (plot_part_per_cell)
+        {
+            MultiFab temp_dat(grids[lev],mf[lev].DistributionMap(),1,0);
+            temp_dat.setVal(0);
+
+            // MultiFab containing number of particles in each cell
+            mypc->Increment(temp_dat, lev);
+            MultiFab::Copy(mf[lev], temp_dat, 0, dcomp, 1, 0);
+            if (lev == 0)
+            {
+                varnames.push_back("part_per_cell");
+            }
+            dcomp += 1;
+        }
+
+        if (plot_part_per_grid || plot_part_per_proc)
+        {
+            const Vector<long>& npart_in_grid = mypc->NumberOfParticlesInGrid(lev);
+
+            if (plot_part_per_grid)
+            {
+#ifdef _OPENMP
+#pragma omp parallel
+#endif
+               for (MFIter mfi(mf[lev]); mfi.isValid(); ++mfi) {
+                   (mf[lev])[mfi].setVal(static_cast<Real>(npart_in_grid[mfi.index()]), dcomp);
+               }
+               if (lev == 0)
+               {
+                   varnames.push_back("part_per_grid");
+               }
+               dcomp += 1;
+           }
+
+           if (plot_part_per_proc)
+           {
+               long n_per_proc = 0;
+#ifdef _OPENMP
+#pragma omp parallel reduction(+:n_per_proc)
+#endif
+                for (MFIter mfi(mf[lev]); mfi.isValid(); ++mfi) {
+                    n_per_proc += npart_in_grid[mfi.index()];
+                }
+                mf[lev].setVal(static_cast<Real>(n_per_proc), dcomp,1);
+                if (lev == 0)
+                {
+                    varnames.push_back("part_per_proc");
+                }
+                dcomp += 1;
+            }
+        }
+
+        if (plot_proc_number)
+        {
+#ifdef _OPENMP
+#pragma omp parallel
+#endif
+            for (MFIter mfi(mf[lev]); mfi.isValid(); ++mfi) {
+                (mf[lev])[mfi].setVal(static_cast<Real>(ParallelDescriptor::MyProc()), dcomp);
+            }
+            if (lev == 0)
+            {
+                varnames.push_back("proc_number");
+            }
+            dcomp += 1;
+        }
+
+        if (plot_divb)
+        {
+            ComputeDivB(mf[lev], dcomp,
+                        {Bfield_aux[lev][0].get(),Bfield_aux[lev][1].get(),Bfield_aux[lev][2].get()},
+                        WarpX::CellSize(lev));
+            if (lev == 0)
+            {
+                varnames.push_back("divB");
+            }
+            dcomp += 1;
+        }
+
+        if (plot_dive)
+        {
+            const BoxArray& ba = amrex::convert(boxArray(lev),IntVect::TheUnitVector());
+            MultiFab dive(ba,DistributionMap(lev),1,0);
+            ComputeDivE(dive, 0,
+                        {Efield_aux[lev][0].get(), Efield_aux[lev][1].get(), Efield_aux[lev][2].get()},
+                        WarpX::CellSize(lev));
+            amrex::average_node_to_cellcenter(mf[lev], dcomp, dive, 0, 1);
+            if (lev == 0)
+            {
+                varnames.push_back("divE");
+            }
+            dcomp += 1;
+        }
+
+        if (plot_finepatch)
+        {
+            PackPlotDataPtrs(srcmf, Efield_fp[lev]);
+            amrex::average_edge_to_cellcenter(mf[lev], dcomp, srcmf);
+#if (AMREX_SPACEDIM == 2)
+            MultiFab::Copy(mf[lev], mf[lev], dcomp+1, dcomp+2, 1, ngrow);
+            amrex::average_node_to_cellcenter(mf[lev], dcomp+1, *Efield_fp[lev][1], 0, 1);
+#endif
+            if (lev == 0)
+            {
+                varnames.push_back("Ex_fp");
+                varnames.push_back("Ey_fp");
+                varnames.push_back("Ez_fp");
+            }
+            dcomp += 3;
+
+            PackPlotDataPtrs(srcmf, Bfield_fp[lev]);
+            amrex::average_face_to_cellcenter(mf[lev], dcomp, srcmf);
+#if (AMREX_SPACEDIM == 2)
+            MultiFab::Copy(mf[lev], mf[lev], dcomp+1, dcomp+2, 1, ngrow);
+            MultiFab::Copy(mf[lev], *Bfield_fp[lev][1], 0, dcomp+1, 1, ngrow);
+#endif
+            if (lev == 0)
+            {
+                varnames.push_back("Bx_fp");
+                varnames.push_back("By_fp");
+                varnames.push_back("Bz_fp");
+            }
+            dcomp += 3;
+        }
+
+        if (plot_crsepatch)
+        {
+            // First the electric field
+            if (lev == 0)
+            {
+                mf[lev].setVal(0.0, dcomp, 3, ngrow);
+            }
+            else
+            {
+                std::array<std::unique_ptr<MultiFab>, 3> E = getInterpolatedE(lev);
+                PackPlotDataPtrs(srcmf, E);
+                amrex::average_edge_to_cellcenter(mf[lev], dcomp, srcmf);
+#if (AMREX_SPACEDIM == 2)
+                MultiFab::Copy(mf[lev], mf[lev], dcomp+1, dcomp+2, 1, ngrow);
+                amrex::average_node_to_cellcenter(mf[lev], dcomp+1, *E[1], 0, 1);
+#endif
+            }
+            if (lev == 0)
+            {
+                varnames.push_back("Ex_cp");
+                varnames.push_back("Ey_cp");
+                varnames.push_back("Ez_cp");
+            }
+            dcomp += 3;
+
+            // now the magnetic field
+            if (lev == 0)
+            {
+                mf[lev].setVal(0.0, dcomp, 3, ngrow);
+            }
+            else
+            {
+                std::array<std::unique_ptr<MultiFab>, 3> B = getInterpolatedB(lev);
+                PackPlotDataPtrs(srcmf, B);
+                amrex::average_face_to_cellcenter(mf[lev], dcomp, srcmf);
+#if (AMREX_SPACEDIM == 2)
+                MultiFab::Copy(mf[lev], mf[lev], dcomp+1, dcomp+2, 1, ngrow);
+                MultiFab::Copy(mf[lev], *B[1], 0, dcomp+1, 1, ngrow);
+#endif
+            }
+            if (lev == 0)
+            {
+                varnames.push_back("Bx_cp");
+                varnames.push_back("By_cp");
+                varnames.push_back("Bz_cp");
+            }
+            dcomp += 3;
+        }
+
+        if (costs[0] != nullptr)
+        {
+            MultiFab::Copy(mf[lev], *costs[lev], 0, dcomp, 1, ngrow);
+            if (lev == 0)
+            {
+                varnames.push_back("costs");
+            }
+            dcomp += 1;
+        }
+
+        BL_ASSERT(dcomp == ncomp);
+    }
+
+    if (insitu_bridge->update(istep[0], t_new[0],
+        dynamic_cast<amrex::AmrMesh*>(const_cast<WarpX*>(this)),
+        {&mf}, {varnames}))
+    {
+        amrex::ErrorStream()
+            << "WarpXIO::UpdateInSitu : Failed to update the in situ bridge."
+            << std::endl;
+
+        amrex::Abort();
+    }
+#endif
+}
+
+void
 WarpX::WritePlotFile () const
 {
     BL_PROFILE("WarpX::WritePlotFile()");
diff --git a/Source/WarpXInitData.cpp b/Source/WarpXInitData.cpp
index 769df87f5..a6b63a92d 100644
--- a/Source/WarpXInitData.cpp
+++ b/Source/WarpXInitData.cpp
@@ -7,6 +7,10 @@
 #include <WarpX_f.H>
 #include <WarpXWrappers.h>
 
+#ifdef BL_USE_SENSEI_INSITU
+#include <AMReX_AmrMeshInSituBridge.H>
+#endif
+
 using namespace amrex;
 
 void
@@ -43,14 +47,31 @@ WarpX::InitData ()
         printGridSummary(std::cout, 0, finestLevel());
     }
 
+#ifdef BL_USE_SENSEI_INSITU
+    insitu_bridge = new amrex::AmrMeshInSituBridge;
+    insitu_bridge->setEnabled(insitu_int > 0 ? 1 : 0);
+    insitu_bridge->setConfig(insitu_config);
+    if (insitu_bridge->initialize())
+    {
+        amrex::ErrorStream()
+            << "WarpX::InitData : Failed to initialize the in situ bridge."
+            << std::endl;
+
+        amrex::Abort();
+    }
+    insitu_bridge->setFrequency(1);
+#endif
+
     if (restart_chkfile.empty())
     {
-	if (plot_int > 0) {
+        if (plot_int > 0)
             WritePlotFile();
-	}
-	if (check_int > 0) {
-	    WriteCheckPointFile();
-	}
+
+        if (check_int > 0)
+            WriteCheckPointFile();
+
+        if ((insitu_int > 0) && (insitu_start == 0))
+            UpdateInSitu();
     }
 }