Ion density sampled evaluation (#274)

* save/load poisson rom operator. * test routine for ion density * ion density hyper-reduction verified. * added rom ion test in ci workflow. * fixed ci workflow * setting ion configuration using Ions::setPositions.
LLNL · Oct 22, 2024 · f2f1821 · f2f1821
1 parent b32d686
commit f2f1821
Show file tree

Hide file tree

Showing 13 changed files with 417 additions and 3 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -49,7 +49,11 @@ jobs:
             cd ${GITHUB_WORKSPACE}/tests/ROM/test_rom_poisson
             ln -s ${GITHUB_WORKSPACE}/build/src/mgmol-rom .
             ln -s ${GITHUB_WORKSPACE}/potentials/* .
-            mpirun -n 3 --oversubscribe ./mgmol-rom -c carbyne.cfg -i carbyne.in
+            mpirun -n 3 --oversubscribe ./mgmol-rom -c carbyne.poisson.cfg -i carbyne.in
+      - name: test ROM ion density evaluation
+        run: |
+            cd ${GITHUB_WORKSPACE}/tests/ROM/test_rom_poisson
+            mpirun -n 3 --oversubscribe ./mgmol-rom -c carbyne.ion.cfg -i carbyne.in
   # code-style:
   #   runs-on: ubuntu-latest
   #   needs: [docker-image]

diff --git a/src/Control.cc b/src/Control.cc
@@ -2034,10 +2034,14 @@ void Control::setROMOptions(const boost::program_options::variables_map& vm)
             rom_pri_option.rom_stage = ROMStage::ONLINE;
         else if (str.compare("build") == 0)
             rom_pri_option.rom_stage = ROMStage::BUILD;
+        else if (str.compare("online_poisson") == 0)
+            rom_pri_option.rom_stage = ROMStage::ONLINE_POISSON;
         else if (str.compare("test_poisson") == 0)
             rom_pri_option.rom_stage = ROMStage::TEST_POISSON;
         else if (str.compare("test_rho") == 0)
             rom_pri_option.rom_stage = ROMStage::TEST_RHO;
+        else if (str.compare("test_ion") == 0)
+            rom_pri_option.rom_stage = ROMStage::TEST_ION;
         else if (str.compare("none") == 0)
             rom_pri_option.rom_stage = ROMStage::UNSUPPORTED;
 
@@ -2058,6 +2062,7 @@ void Control::setROMOptions(const boost::program_options::variables_map& vm)
         rom_pri_option.librom_snapshot_freq = vm["ROM.offline.librom_snapshot_freq"].as<int>();
 
         rom_pri_option.num_potbasis = vm["ROM.basis.number_of_potential_basis"].as<int>();
+        rom_pri_option.pot_rom_file = vm["ROM.potential_rom_file"].as<std::string>();
     }  // onpe0
 
     // synchronize all processors
@@ -2073,6 +2078,7 @@ void Control::syncROMOptions()
 
     mmpi.bcast(rom_pri_option.restart_file_fmt, comm_global_);
     mmpi.bcast(rom_pri_option.basis_file, comm_global_);
+    mmpi.bcast(rom_pri_option.pot_rom_file, comm_global_);
 
     auto bcast_check = [](int mpirc) {
         if (mpirc != MPI_SUCCESS)

diff --git a/src/Electrostatic.h b/src/Electrostatic.h
@@ -47,6 +47,9 @@ class Electrostatic
     ~Electrostatic();
     static Timer solve_tm() { return solve_tm_; }
 
+    const bool isDielectric() { return diel_flag_; }
+    pb::GridFunc<RHODTYPE>* getRhoc() { return grhoc_; }
+
     Poisson* getPoissonSolver() { return poisson_solver_; }
 
     void setup(const short max_sweeps);

diff --git a/src/MGmol.h b/src/MGmol.h
@@ -187,6 +187,7 @@ class MGmol : public MGmolInterface
         return hamiltonian_;
     }
     std::shared_ptr<Rho<OrbitalsType>> getRho() { return rho_; }
+    std::shared_ptr<Ions> getIons() { return ions_; }
 
     void run() override;
 

diff --git a/src/Potentials.cc b/src/Potentials.cc
@@ -899,6 +899,97 @@ void Potentials::initBackground(Ions& ions)
     if (fabs(background_charge_) < 1.e-10) background_charge_ = 0.;
 }
 
+void Potentials::evalIonDensityOnSamplePts(
+    Ions& ions, const std::vector<int> &local_idx, std::vector<RHODTYPE> &sampled_rhoc)
+{
+    Mesh* mymesh           = Mesh::instance();
+    MGmol_MPI& mmpi = *(MGmol_MPI::instance());
+    const pb::Grid& mygrid = mymesh->grid();
+
+    const char flag_filter = pot_type(0);
+    if (flag_filter == 's')
+    {
+        if (onpe0)
+        {
+            cout << "Potentials::evalIonDensityOnSamplePts - flag_filter s is not supported"
+                << endl;
+        }
+        mmpi.abort();
+    }
+
+    // initialize output vector
+    sampled_rhoc.resize(local_idx.size());
+    for (int d = 0; d < sampled_rhoc.size(); d++)
+        sampled_rhoc[d] = 0.0;
+
+    // Loop over ions
+    for (auto& ion : ions.overlappingVL_ions())
+    {
+        const Species& sp(ion->getSpecies());
+
+        Vector3D position(ion->position(0), ion->position(1), ion->position(2));
+
+        initializeRadialDataOnSampledPts(position, sp, local_idx, sampled_rhoc);
+    }
+
+    return;
+}
+
+void Potentials::initializeRadialDataOnSampledPts(
+    const Vector3D& position, const Species& sp, const std::vector<int> &local_idx, std::vector<RHODTYPE> &sampled_rhoc)
+{
+    assert(local_idx.size() == sampled_rhoc.size());
+
+    Control& ct = *(Control::instance());
+
+    Mesh* mymesh           = Mesh::instance();
+    const pb::Grid& mygrid = mymesh->grid();
+
+    const int dim0 = mygrid.dim(0);
+    const int dim1 = mygrid.dim(1);
+    const int dim2 = mygrid.dim(2);
+
+    const double start0 = mygrid.start(0);
+    const double start1 = mygrid.start(1);
+    const double start2 = mygrid.start(2);
+
+    const double h0 = mygrid.hgrid(0);
+    const double h1 = mygrid.hgrid(1);
+    const double h2 = mygrid.hgrid(2);
+
+    Vector3D point(0., 0., 0.);
+
+    const double lrad = sp.lradius();
+
+    const RadialInter& lpot(sp.local_pot());
+    const Vector3D lattice(mygrid.ll(0), mygrid.ll(1), mygrid.ll(2));
+
+    for(int k  = 0; k < local_idx.size(); k++)
+    {
+        /*
+            local_idx provides offset.
+            offset = iz + iy * dim2 + ix * dim1 * dim2;
+            evaluate x,y,z indices backward.
+        */
+        int iz = local_idx[k] % dim2;
+        int ix = local_idx[k] / dim2;
+        int iy = ix % dim1;
+        ix /= dim1;
+
+        /* compute grid point position */
+        point[0] = start0 + ix * h0;
+        point[1] = start1 + iy * h1;
+        point[2] = start2 + iz * h2;
+
+        /* accumulate ion species density */
+        const double r = position.minimage(point, lattice, ct.bcPoisson);
+        if (r < lrad)
+            sampled_rhoc[k] += sp.getRhoComp(r);
+    }
+
+    return;
+}
+
 template void Potentials::setVxc<double>(
     const double* const vxc, const int iterativeIndex);
 template void Potentials::setVxc<float>(

diff --git a/src/Potentials.h b/src/Potentials.h
@@ -95,6 +95,9 @@ class Potentials
     void initializeSupersampledRadialDataOnMesh(
         const Vector3D& position, const Species& sp);
 
+    void initializeRadialDataOnSampledPts(
+        const Vector3D& position, const Species& sp, const std::vector<int> &local_idx, std::vector<RHODTYPE> &sampled_rhoc);
+
 public:
     Potentials(const bool vh_frozen = false);
 
@@ -159,6 +162,8 @@ class Potentials
 
     double getChargeInCell() const { return charge_in_cell_; }
 
+    const double getBackgroundCharge() const { return background_charge_; }
+
     /*!
      * initialize total potential as local pseudopotential
      */
@@ -196,6 +201,8 @@ class Potentials
     void initBackground(Ions& ions);
     void addBackgroundToRhoComp();
 
+    void evalIonDensityOnSamplePts(Ions& ions, const std::vector<int> &local_idx, std::vector<RHODTYPE> &sampled_rhoc);
+
 #ifdef HAVE_TRICUBIC
     void readExternalPot(const string filename, const char type);
     void setupVextTricubic();

diff --git a/src/read_config.cc b/src/read_config.cc
@@ -431,6 +431,8 @@ void setupROMConfigOption(po::options_description &rom_cfg)
         ("ROM.offline.variable", po::value<std::string>()->default_value(""),
             "FOM variable to perform POD: either orbitals or potential.")
         ("ROM.basis.number_of_potential_basis", po::value<int>()->default_value(-1),
-            "Number of potential POD basis to build Hartree potential ROM operator.");
+            "Number of potential POD basis to build Hartree potential ROM operator.")
+        ("ROM.potential_rom_file", po::value<std::string>()->default_value(""),
+            "File name to save/load potential ROM operators.");
 }
 #endif
diff --git a/src/rom_Control.h b/src/rom_Control.h
@@ -22,8 +22,10 @@ enum class ROMStage
     ONLINE,
     RESTORE,    // TODO(kevin): what stage is this?
     BUILD,
+    ONLINE_POISSON,
     TEST_POISSON,
     TEST_RHO,
+    TEST_ION,
     UNSUPPORTED
 };
 
@@ -51,6 +53,7 @@ struct ROMPrivateOptions
 
     /* options for ROM building */
     int num_potbasis = -1;
+    std::string pot_rom_file = "";
 };
 
 #endif  // ROM_CONTROL_H
diff --git a/src/rom_main.cc b/src/rom_main.cc
@@ -118,6 +118,13 @@ int main(int argc, char** argv)
                 buildROMPoissonOperator<ExtendedGridOrbitals>(mgmol);
         break;
 
+        case (ROMStage::ONLINE_POISSON):
+            if (ct.isLocMode())
+                runPoissonROM<LocGridOrbitals>(mgmol);
+            else
+                runPoissonROM<ExtendedGridOrbitals>(mgmol);
+        break;
+
         case (ROMStage::TEST_POISSON):
             if (ct.isLocMode())
                 testROMPoissonOperator<LocGridOrbitals>(mgmol);
@@ -130,6 +137,13 @@ int main(int argc, char** argv)
                 testROMRhoOperator<LocGridOrbitals>(mgmol);
             else
                 testROMRhoOperator<ExtendedGridOrbitals>(mgmol);
+
+        case (ROMStage::TEST_ION):
+            if (ct.isLocMode())
+                testROMIonDensity<LocGridOrbitals>(mgmol);
+            else
+                testROMIonDensity<ExtendedGridOrbitals>(mgmol);
+
         break;
 
         default: