path: root/Source/WarpXElectrostatic.cpp

#include <AMReX_MGT_Solver.H>
#include <AMReX_stencil_types.H>

#include <WarpX.H>
#include <WarpX_f.H>

namespace
{
    const std::string level_prefix {"Level_"};
}

using namespace amrex;

class NoOpPhysBC
    : public amrex::PhysBCFunctBase
{
public:
    NoOpPhysBC () {}
    virtual ~NoOpPhysBC () {}
    virtual void FillBoundary (amrex::MultiFab& mf, int, int, amrex::Real time) override { }
    using amrex::PhysBCFunctBase::FillBoundary;
};

void
WarpX::EvolveES (int numsteps) {

    amrex::Print() << "Running in electrostatic mode \n";

    BL_PROFILE("WarpX::EvolveES()");
    Real cur_time = t_new[0];
    static int last_plot_file_step = 0;
    static int last_check_file_step = 0;

    int numsteps_max;
    if (numsteps < 0) {  // Note that the default argument is numsteps = -1
        numsteps_max = max_step;
    } else {
        numsteps_max = std::min(istep[0]+numsteps, max_step);
    }

    bool max_time_reached = false;

    // nodal storage for thee electrostatic case
    const int num_levels = max_level + 1;
    Vector<std::unique_ptr<MultiFab> > rhoNodal(num_levels);
    Vector<std::unique_ptr<MultiFab> > phiNodal(num_levels);
    Vector<std::array<std::unique_ptr<MultiFab>, 3> > eFieldNodal(num_levels);
    const int ng = 1;
    for (int lev = 0; lev <= max_level; lev++) {
        BoxArray nba = boxArray(lev);
        nba.surroundingNodes();
        rhoNodal[lev].reset(new MultiFab(nba, dmap[lev], 1, ng));
        phiNodal[lev].reset(new MultiFab(nba, dmap[lev], 1, 2));

        eFieldNodal[lev][0].reset(new MultiFab(nba, dmap[lev], 1, ng));
        eFieldNodal[lev][1].reset(new MultiFab(nba, dmap[lev], 1, ng));
        eFieldNodal[lev][2].reset(new MultiFab(nba, dmap[lev], 1, ng));
    }

    const int lev = 0;
    for (int step = istep[0]; step < numsteps_max && cur_time < stop_time; ++step)
    {

	// Start loop on time steps
        amrex::Print() << "\nSTEP " << step+1 << " starts ...\n";

        // At initialization, particles have p^{n-1/2} and x^{n-1/2}.
        // Beyond one step, particles have p^{n-1/2} and x^{n}.
        if (is_synchronized) {
            // on first step, push X by 0.5*dt
            mypc->PushXES(0.5*dt[lev]);
            UpdatePlasmaInjectionPosition(0.5*dt[lev]);
            mypc->Redistribute();
            mypc->DepositCharge(rhoNodal);
            computePhi(rhoNodal, phiNodal);
            computeE(eFieldNodal, phiNodal);
            is_synchronized = false;
        }

        mypc->FieldGatherES(eFieldNodal, gather_masks);

        const std::string& ppltfile = amrex::Concatenate("particles", istep[0], 5);
        auto& pc = mypc->GetParticleContainer(0);
        pc.WriteAsciiFile(ppltfile);

        // Evolve particles to p^{n+1/2} and x^{n+1}
        mypc->EvolveES(eFieldNodal, rhoNodal, cur_time, dt[lev]);

        mypc->DepositCharge(rhoNodal);

        computePhi(rhoNodal, phiNodal);
        computeE(eFieldNodal, phiNodal);

        if (cur_time + dt[0] >= stop_time - 1.e-3*dt[0] || step == numsteps_max-1) {
            // on last step, push by only 0.5*dt to synchronize all at n+1/2
            mypc->PushXES(-0.5*dt[lev]);
            UpdatePlasmaInjectionPosition(-0.5*dt[lev]);
            is_synchronized = true;
        }

        mypc->Redistribute();

        ++istep[0];

        cur_time += dt[0];

        bool to_make_plot = (plot_int > 0) && ((step+1) % plot_int == 0);

        amrex::Print()<< "STEP " << step+1 << " ends." << " TIME = "
                      << cur_time << " DT = " << dt[0] << "\n";

	// sync up time
	for (int i = 0; i <= finest_level; ++i) {
	    t_new[i] = cur_time;
	}

	if (to_make_plot) {
            // replace with ES field Gather
            mypc->DepositCharge(rhoNodal);
            computePhi(rhoNodal, phiNodal);
            phiNodal[0]->FillBoundary(Geom(0).periodicity());
            computeE(eFieldNodal, phiNodal);
            mypc->FieldGatherES(eFieldNodal, gather_masks);
	    last_plot_file_step = step+1;
	    WritePlotFileES(rhoNodal, phiNodal, eFieldNodal);
	}

	if (check_int > 0 && (step+1) % check_int == 0) {
	    last_check_file_step = step+1;
	    WriteCheckPointFile();
	}

	if (cur_time >= stop_time - 1.e-3*dt[0]) {
	    max_time_reached = true;
	    break;
	}

	// End loop on time steps
    }

    if (plot_int > 0 && istep[0] > last_plot_file_step && (max_time_reached || istep[0] >= max_step)) {
        WritePlotFileES(rhoNodal, phiNodal, eFieldNodal);
    }

    if (check_int > 0 && istep[0] > last_check_file_step && (max_time_reached || istep[0] >= max_step)) {
	WriteCheckPointFile();
    }
}

void WarpX::zeroOutBoundary(amrex::MultiFab& input_data,
                            amrex::MultiFab& bndry_data,
                            const FabArray<BaseFab<int> >& mask) const {
    bndry_data.setVal(0.0, 1);
    for (MFIter mfi(input_data); mfi.isValid(); ++mfi) {
        const Box& bx = mfi.validbox();
        WRPX_ZERO_OUT_BNDRY(bx.loVect(), bx.hiVect(),
                            input_data[mfi].dataPtr(),
                            bndry_data[mfi].dataPtr(),
                            mask[mfi].dataPtr());
    }
    bndry_data.FillBoundary();
}

void WarpX::sumFineToCrseNodal(const amrex::MultiFab& fine,
                               amrex::MultiFab& crse,
                               const amrex::Geometry& cgeom,
                               const amrex::IntVect& ratio) {
    const BoxArray& fine_BA = fine.boxArray();
    const DistributionMapping& fine_dm = fine.DistributionMap();
    BoxArray coarsened_fine_BA = fine_BA;
    coarsened_fine_BA.coarsen(ratio);

    MultiFab coarsened_fine_data(coarsened_fine_BA, fine_dm, 1, 0);
    coarsened_fine_data.setVal(0.0);

    for (MFIter mfi(coarsened_fine_data); mfi.isValid(); ++mfi) {
        const Box& bx = mfi.validbox();
        const Box& crse_box = coarsened_fine_data[mfi].box();
        const Box& fine_box = fine[mfi].box();
        WRPX_SUM_FINE_TO_CRSE_NODAL(bx.loVect(), bx.hiVect(), ratio.getVect(),
                                    coarsened_fine_data[mfi].dataPtr(), crse_box.loVect(), crse_box.hiVect(),
                                    fine[mfi].dataPtr(), fine_box.loVect(), fine_box.hiVect());
    }

    crse.copy(coarsened_fine_data, cgeom.periodicity(), FabArrayBase::ADD);
}

void
WarpX::fixRHSForSolve(Vector<std::unique_ptr<MultiFab> >& rhs,
                      const Vector<std::unique_ptr<FabArray<BaseFab<int> > > >& masks) const {
    int num_levels = rhs.size();
    for (int lev = 0; lev < num_levels; ++lev) {
        MultiFab& fine_rhs = *rhs[lev];
        const FabArray<BaseFab<int> >& mask = *masks[lev];
        const BoxArray& fine_ba = fine_rhs.boxArray();
        const DistributionMapping& fine_dm = fine_rhs.DistributionMap();
        MultiFab fine_bndry_data(fine_ba, fine_dm, 1, 1);
        zeroOutBoundary(fine_rhs, fine_bndry_data, mask);
    }
}

void WarpX::getLevelMasks(Vector<std::unique_ptr<FabArray<BaseFab<int> > > >& masks,
                          const int nnodes) {
    int num_levels = grids.size();
    BL_ASSERT(num_levels == dmap.size());

    int covered = 0;
    int notcovered = 1;
    int physbnd = 1;
    int interior = 0;

    for (int lev = 0; lev < num_levels; ++lev) {
        BoxArray nba = grids[lev];
        nba.surroundingNodes();

        FabArray<BaseFab<int> > tmp_mask(nba, dmap[lev], 1, nnodes);
        tmp_mask.BuildMask(geom[lev].Domain(), geom[lev].periodicity(),
                           covered, notcovered, physbnd, interior);
        masks[lev].reset(new FabArray<BaseFab<int> >(nba, dmap[lev], 1, 0));
        for (MFIter mfi(tmp_mask); mfi.isValid(); ++mfi) {
            const Box& bx = mfi.validbox();
            WRPX_BUILD_MASK(bx.loVect(), bx.hiVect(),
                            tmp_mask[mfi].dataPtr(), (*masks[lev])[mfi].dataPtr(), &nnodes);
        }
    }
}


void WarpX::computePhi(const Vector<std::unique_ptr<MultiFab> >& rho,
                             Vector<std::unique_ptr<MultiFab> >& phi) const {


    int num_levels = rho.size();
    Vector<std::unique_ptr<MultiFab> > rhs(num_levels);
    for (int lev = 0; lev < num_levels; ++lev) {
        phi[lev]->setVal(0.0, 2);
        rhs[lev].reset(new MultiFab(rho[lev]->boxArray(), dmap[lev], 1, 0));
        MultiFab::Copy(*rhs[lev], *rho[lev], 0, 0, 1, 0);
        rhs[lev]->mult(-1.0/PhysConst::ep0, 0);
    }

    fixRHSForSolve(rhs, masks);

    bool nodal = true;
    bool have_rhcc = false;
    int  nc = 0;
    int Ncomp = 1;
    int stencil = ND_CROSS_STENCIL;
    int verbose = 0;
    Vector<int> mg_bc(2*AMREX_SPACEDIM, 1); // this means Dirichlet
    Real rel_tol = 1.0e-14;
    Real abs_tol = 1.0e-14;

    Vector<Geometry>            level_geom(1);
    Vector<BoxArray>            level_grids(1);
    Vector<DistributionMapping> level_dm(1);
    Vector<MultiFab*>           level_phi(1);
    Vector<MultiFab*>           level_rhs(1);

    for (int lev = 0; lev < num_levels; ++lev) {
        level_phi[0]   = phi[lev].get();
        level_rhs[0]   = rhs[lev].get();
        level_geom[0]  = geom[lev];
        level_grids[0] = grids[lev];
        level_dm[0]    = dmap[lev];

        MGT_Solver solver(level_geom, mg_bc.dataPtr(), level_grids,
                          level_dm, nodal,
                          stencil, have_rhcc, nc, Ncomp, verbose);

        solver.set_nodal_const_coefficients(1.0);

        solver.solve_nodal(level_phi, level_rhs, rel_tol, abs_tol);

        if (lev < num_levels-1) {

            NoOpPhysBC cphysbc, fphysbc;
#if AMREX_SPACEDIM == 3
            int lo_bc[] = {BCType::int_dir, BCType::int_dir, BCType::int_dir};
            int hi_bc[] = {BCType::int_dir, BCType::int_dir, BCType::int_dir};
#else
            int lo_bc[] = {BCType::int_dir, BCType::int_dir};
            int hi_bc[] = {BCType::int_dir, BCType::int_dir};
#endif
            Vector<BCRec> bcs(1, BCRec(lo_bc, hi_bc));
            NodeBilinear mapper;

            amrex::InterpFromCoarseLevel(*phi[lev+1], 0.0, *phi[lev],
                                         0, 0, 1, geom[lev], geom[lev+1],
                                         cphysbc, fphysbc,
                                         IntVect(AMREX_D_DECL(2, 2, 2)), &mapper, bcs);
        }
    }

    for (int lev = 0; lev < num_levels; ++lev) {
        const Geometry& gm = geom[lev];
        phi[lev]->FillBoundary(gm.periodicity());
    }
}

void WarpX::computeE(Vector<std::array<std::unique_ptr<MultiFab>, 3> >& E,
                     const Vector<std::unique_ptr<MultiFab> >& phi) const {

    const int num_levels = E.size();
    for (int lev = 0; lev < num_levels; ++lev) {
        const auto& gm = GetInstance().Geom(lev);
        const Real* dx = gm.CellSize();
        for (MFIter mfi(*phi[lev]); mfi.isValid(); ++mfi) {
            const Box& bx = mfi.validbox();

            WRPX_COMPUTE_E_NODAL(bx.loVect(), bx.hiVect(),
                                 (*phi[lev] )[mfi].dataPtr(),
                                 (*E[lev][0])[mfi].dataPtr(),
                                 (*E[lev][1])[mfi].dataPtr(),
#if AMREX_SPACEDIM == 3
                                 (*E[lev][2])[mfi].dataPtr(),
#endif
                                 dx);
        }

        E[lev][0]->FillBoundary(gm.periodicity());
        E[lev][1]->FillBoundary(gm.periodicity());
#if AMREX_SPACEDIM == 3
        E[lev][2]->FillBoundary(gm.periodicity());
#endif
    }
}

void
WarpX::
WritePlotFileES (const amrex::Vector<std::unique_ptr<amrex::MultiFab> >& rho,
                 const amrex::Vector<std::unique_ptr<amrex::MultiFab> >& phi,
                 const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>, 3> >& E)
{
    BL_PROFILE("WarpX::WritePlotFileES()");

    VisMF::Header::Version current_version = VisMF::GetHeaderVersion();
    VisMF::SetHeaderVersion(plotfile_headerversion);

    const std::string& plotfilename = amrex::Concatenate(plot_file,istep[0]);

    amrex::Print() << "  Writing plotfile " << plotfilename << "\n";

    const int nlevels = finestLevel()+1;

    {
	Vector<std::string> varnames;
	Vector<std::unique_ptr<MultiFab> > mf(finest_level+1);

	for (int lev = 0; lev <= finest_level; ++lev) {
            int ncomp = 5;
            const int ngrow = 0;
            mf[lev].reset(new MultiFab(grids[lev], dmap[lev], ncomp, ngrow));

            int dcomp = 0;
            amrex::average_node_to_cellcenter(*mf[lev], dcomp, *rho[lev], 0, 1);
            if (lev == 0) {
                varnames.push_back("rho");
            }
            dcomp += 1;

            amrex::average_node_to_cellcenter(*mf[lev], dcomp  , *E[lev][0], 0, 1);
            amrex::average_node_to_cellcenter(*mf[lev], dcomp+1, *E[lev][0], 0, 1);
            amrex::average_node_to_cellcenter(*mf[lev], dcomp+2, *E[lev][0], 0, 1);

            if (lev == 0) {
                varnames.push_back("Ex");
                varnames.push_back("Ey");
                varnames.push_back("Ez");
            }
            dcomp += 3;

            amrex::average_node_to_cellcenter(*mf[lev], dcomp, *phi[lev], 0, 1);
            if (lev == 0) {
                varnames.push_back("phi");
            }
            dcomp += 1;
        }

        Vector<std::string> rfs(1,"raw_fields"); // pre-build raw_fields/
        amrex::WriteMultiLevelPlotfile(plotfilename, finest_level+1,
                                       amrex::GetVecOfConstPtrs(mf),
                                       varnames, Geom(), t_new[0], istep, refRatio(),
                                       "HyperCLaw-V1.1",
                                       "Level_",
                                       "Cell",
                                       rfs);
    }

    {
        const std::string raw_plotfilename = plotfilename + "/raw_fields";
        const int nlevels = finestLevel()+1;
        for (int lev = 0; lev < nlevels; ++lev) {
            const DistributionMapping& dm = DistributionMap(lev);

            MultiFab Ex( E[lev][0]->boxArray(), dm, 1, 0);
            MultiFab Ey( E[lev][1]->boxArray(), dm, 1, 0);
            MultiFab Ez( E[lev][2]->boxArray(), dm, 1, 0);
            MultiFab charge_density(rho[lev]->boxArray(), dm, 1, 0);
            MultiFab potential(phi[lev]->boxArray(), dm, 1, 0);

            MultiFab::Copy(Ex, *E[lev][0], 0, 0, 1, 0);
            MultiFab::Copy(Ey, *E[lev][1], 0, 0, 1, 0);
            MultiFab::Copy(Ez, *E[lev][2], 0, 0, 1, 0);
            MultiFab::Copy(charge_density, *rho[lev], 0, 0, 1, 0);
            MultiFab::Copy(potential, *phi[lev], 0, 0, 1, 0);

            VisMF::Write(Ex, amrex::MultiFabFileFullPrefix(lev, raw_plotfilename, level_prefix, "Ex"));
            VisMF::Write(Ey, amrex::MultiFabFileFullPrefix(lev, raw_plotfilename, level_prefix, "Ey"));
            VisMF::Write(Ez, amrex::MultiFabFileFullPrefix(lev, raw_plotfilename, level_prefix, "Ez"));
            VisMF::Write(charge_density, amrex::MultiFabFileFullPrefix(lev, raw_plotfilename, level_prefix, "rho"));
            VisMF::Write(potential, amrex::MultiFabFileFullPrefix(lev, raw_plotfilename, level_prefix, "phi"));
        }
    }

    Vector<std::string> particle_varnames;
    particle_varnames.push_back("weight");

    particle_varnames.push_back("momentum_x");
    particle_varnames.push_back("momentum_y");
    particle_varnames.push_back("momentum_z");

    particle_varnames.push_back("Ex");
    particle_varnames.push_back("Ey");
    particle_varnames.push_back("Ez");

    particle_varnames.push_back("Bx");
    particle_varnames.push_back("By");
    particle_varnames.push_back("Bz");

    mypc->Checkpoint(plotfilename, "particle", true, particle_varnames);

    WriteJobInfo(plotfilename);

    WriteWarpXHeader(plotfilename);

    VisMF::SetHeaderVersion(current_version);
}