FOM orbital snapshot saving workflow

src/CMakeLists.txt

@@ -159,6 +159,7 @@ set(SOURCES
  magma_singleton.cc
  ChebyshevApproximation.cc
  ChebyshevApproximationInterface.cc
+ rom.cc
 )
 
 add_library(mgmol_src ${SOURCES})

src/MGmol.cc

@@ -1148,7 +1148,19 @@ void MGmol<OrbitalsType>::cleanup()
 
         if (ierr < 0)
             os_ << "WARNING: writing restart data failed!!!" << std::endl;
-    }
+
+#ifdef MGMOL_HAS_LIBROM
+        // Save orbital snapshots
+        if (ct.getROMOptions().save_librom_snapshot > 0 && ct.AtomsDynamic() == AtomsDynamicType::Quench)
+        {
+            ierr = save_orbital_snapshot(
+                filename, *current_orbitals_);
+
+            if (ierr < 0)
+                os_ << "WARNING: writing ROM snapshot data failed!!!" << std::endl;
+        }
+#endif
+    }   // if (ct.out_restart_info > 0 && !ct.AtomsMove())
 
     MPI_Barrier(comm_);
     closing_tm_.stop();

src/MGmol.h

@@ -10,6 +10,8 @@
 #ifndef MGMOL_H
 #define MGMOL_H
 
+#include "mgmol_config.h"
+
 #include "Energy.h"
 #include "GridFuncVector.h"
 #include "Hamiltonian.h"
@@ -303,6 +305,9 @@ class MGmol : public MGmolInterface
     }
 
     OrbitalsType* loadOrbitalFromRestartFile(const std::string filename);
+#ifdef MGMOL_HAS_LIBROM
+    int save_orbital_snapshot(std::string snapshot_dir, OrbitalsType& orbitals);
+#endif
 };
 // Instantiate static variables here to avoid clang warnings
 template <class OrbitalsType>

src/md.cc

@@ -641,6 +641,18 @@ void MGmol<OrbitalsType>::md(OrbitalsType** orbitals, Ions& ions)
                         count++;
                     }
 
+#ifdef MGMOL_HAS_LIBROM
+                    // Save orbital snapshots
+                    if (ct.getROMOptions().save_librom_snapshot > 0)
+                    {
+                        int ierr = save_orbital_snapshot(
+                            ct.md_print_filename + "_mdstep" + std::to_string(mdstep), **orbitals);
+
+                        if (ierr < 0)
+                            os_ << "WARNING md(): writing ROM snapshot data failed!!!" << std::endl;
+                    }
+#endif
+
                     printWithTimeStamp("dumped restart file...", std::cout);
                 }
 

src/rom.cc

@@ -0,0 +1,61 @@
+// Copyright (c) 2017, Lawrence Livermore National Security, LLC and
+// UT-Battelle, LLC.
+// Produced at the Lawrence Livermore National Laboratory and the Oak Ridge
+// National Laboratory.
+// LLNL-CODE-743438
+// All rights reserved.
+// This file is part of MGmol. For details, see https://github.com/llnl/mgmol.
+// Please also read this link https://github.com/llnl/mgmol/LICENSE
+
+#include "mgmol_config.h"
+#ifdef MGMOL_HAS_LIBROM
+
+#include "LocGridOrbitals.h"
+#include "MGmol.h"
+
+#include "librom.h"
+
+#include <string>
+#include <iostream>
+#include <fstream>
+#include <sys/stat.h>
+
+// Save the wavefunction snapshots
+template <class OrbitalsType>
+int MGmol<OrbitalsType>::save_orbital_snapshot(std::string file_path, OrbitalsType& orbitals)
+{
+    std::string snapshot_filename = file_path;
+    struct stat s; 
+    if (stat(file_path.c_str(), &s) == 0)
+    {
+        if (s.st_mode & S_IFDIR)
+            snapshot_filename = file_path + "/orbital";
+        else if (s.st_mode & S_IFREG)
+        {
+            snapshot_filename = file_path + "_orbital";
+        }
+        else
+        {
+            std::cout << file_path << " exists but is not a directory or a file." << std::endl;
+            return 1;
+        }
+    }
+
+    const int dim = orbitals.getLocNumpt();
+    const int totalSamples = orbitals.chromatic_number();
+
+    CAROM::Options svd_options(dim, totalSamples, 1);
+    CAROM::BasisGenerator basis_generator(svd_options, false, snapshot_filename);
+
+    for (int i = 0; i < totalSamples; ++i)
+        basis_generator.takeSample(orbitals.getPsi(i));
+
+    basis_generator.writeSnapshot();
+
+    return 0;
+}
+
+template class MGmol<LocGridOrbitals>;
+template class MGmol<ExtendedGridOrbitals>;
+
+#endif  // MGMOL_HAS_LIBROM