From 553fbb48706765153cd686f6385b1c57a956c2e0 Mon Sep 17 00:00:00 2001
From: Seung Whan Chung <seung.chung@austin.utexas.edu>
Date: Fri, 13 Sep 2024 16:26:37 -0700
Subject: [PATCH] save/load poisson rom operator.

---
 src/Control.cc       |  4 +++
 src/read_config.cc   |  4 ++-
 src/rom_Control.h    |  2 ++
 src/rom_main.cc      |  7 +++++
 src/rom_workflows.cc | 66 +++++++++++++++++++++++++++++++++++++++++++-
 src/rom_workflows.h  |  3 ++
 6 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/src/Control.cc b/src/Control.cc
index 2231813e..b509d1f5 100644
--- a/src/Control.cc
+++ b/src/Control.cc
@@ -2033,6 +2033,8 @@ 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)
@@ -2057,6 +2059,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
@@ -2072,6 +2075,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/read_config.cc b/src/read_config.cc
index 40194678..4f3b34c1 100644
--- 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
index 8c3b4b2c..472b22a2 100644
--- a/src/rom_Control.h
+++ b/src/rom_Control.h
@@ -22,6 +22,7 @@ enum class ROMStage
     ONLINE,
     RESTORE,    // TODO(kevin): what stage is this?
     BUILD,
+    ONLINE_POISSON,
     TEST_POISSON,
     TEST_RHO,
     UNSUPPORTED
@@ -51,6 +52,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
index fa956221..2b810bdf 100644
--- 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);
diff --git a/src/rom_workflows.cc b/src/rom_workflows.cc
index 68349e90..bf7faf37 100644
--- a/src/rom_workflows.cc
+++ b/src/rom_workflows.cc
@@ -218,7 +218,7 @@ void buildROMPoissonOperator(MGmolInterface *mgmol_)
     // write the file from PE0 only
     if (MPIdata::onpe0)
     {
-        std::string rom_oper = "pot_rom_oper.h5";
+        std::string rom_oper = rom_options.pot_rom_file;
         CAROM::HDFDatabase h5_helper;
         h5_helper.create(rom_oper);
         h5_helper.putInteger("number_of_potential_basis", num_pot_basis);
@@ -242,6 +242,67 @@ void buildROMPoissonOperator(MGmolInterface *mgmol_)
     }
 }
 
+template <class OrbitalsType>
+void runPoissonROM(MGmolInterface *mgmol_)
+{
+    Control& ct              = *(Control::instance());
+    Mesh* mymesh             = Mesh::instance();
+    const pb::PEenv& myPEenv = mymesh->peenv();
+    MGmol_MPI& mmpi          = *(MGmol_MPI::instance());
+    const int rank           = mmpi.mypeGlobal();
+    const int nprocs         = mmpi.size();
+
+    ROMPrivateOptions rom_options = ct.getROMOptions();
+    /* type of variable we intend to run POD */
+    ROMVariable rom_var = rom_options.variable;
+    if (rom_var != ROMVariable::POTENTIAL)
+    {
+        std::cerr << "runPoissonROM error: ROM variable must be POTENTIAL to run this stage!\n" << std::endl;
+        MPI_Abort(MPI_COMM_WORLD, 0);
+    }
+
+    /* Load Hartree potential basis matrix */
+    std::string basis_file = rom_options.basis_file;
+    const int num_pot_basis = rom_options.num_potbasis;
+    CAROM::BasisReader basis_reader(basis_file);
+    CAROM::Matrix *pot_basis = basis_reader.getSpatialBasis(num_pot_basis);
+
+    /* initialize rom operator variables */
+    CAROM::Matrix pot_rom(num_pot_basis, num_pot_basis, false);
+    CAROM::Matrix pot_rom_inv(num_pot_basis, num_pot_basis, false);
+    CAROM::Matrix pot_rhs_rom(num_pot_basis, num_pot_basis, false);
+    CAROM::Vector pot_rhs_rescaler(num_pot_basis, false);
+    
+    /* Load ROM operator */
+    // read the file from PE0 only
+    if (MPIdata::onpe0)
+    {
+        std::string rom_oper = rom_options.pot_rom_file;
+        CAROM::HDFDatabase h5_helper;
+        h5_helper.open(rom_oper, "r");
+        int num_pot_basis_file = -1;
+        h5_helper.getInteger("number_of_potential_basis", num_pot_basis_file);
+        CAROM_VERIFY(num_pot_basis_file == num_pot_basis);
+
+        h5_helper.getDoubleArray("potential_rom_operator", pot_rom.getData(),
+                                num_pot_basis * num_pot_basis, false);
+
+        /* load the inverse as well */
+        h5_helper.getDoubleArray("potential_rom_inverse", pot_rom_inv.getData(),
+                                num_pot_basis * num_pot_basis, false);
+
+        /* load right-hand side hyper-reduction operator */
+        h5_helper.getDoubleArray("potential_rhs_rom_inverse", pot_rhs_rom.getData(),
+                                num_pot_basis * num_pot_basis, false);
+
+        /* load right-hand side rescaling operator */
+        h5_helper.getDoubleArray("potential_rhs_rescaler", pot_rhs_rescaler.getData(),
+                                num_pot_basis, false);
+
+        h5_helper.close();
+    }
+}
+
 /* test routines */
 
 template <class OrbitalsType>
@@ -727,6 +788,9 @@ template void readRestartFiles<ExtendedGridOrbitals>(MGmolInterface *mgmol_);
 template void buildROMPoissonOperator<LocGridOrbitals>(MGmolInterface *mgmol_);
 template void buildROMPoissonOperator<ExtendedGridOrbitals>(MGmolInterface *mgmol_);
 
+template void runPoissonROM<LocGridOrbitals>(MGmolInterface *mgmol_);
+template void runPoissonROM<ExtendedGridOrbitals>(MGmolInterface *mgmol_);
+
 template void testROMPoissonOperator<LocGridOrbitals>(MGmolInterface *mgmol_);
 template void testROMPoissonOperator<ExtendedGridOrbitals>(MGmolInterface *mgmol_);
 
diff --git a/src/rom_workflows.h b/src/rom_workflows.h
index fd00c687..b8022a8b 100644
--- a/src/rom_workflows.h
+++ b/src/rom_workflows.h
@@ -45,6 +45,9 @@ void readRestartFiles(MGmolInterface *mgmol_);
 template <class OrbitalsType>
 void buildROMPoissonOperator(MGmolInterface *mgmol_);
 
+template <class OrbitalsType>
+void runPoissonROM(MGmolInterface *mgmol_);
+
 template <class OrbitalsType>
 void testROMPoissonOperator(MGmolInterface *mgmol_);