Hadrons: A2A vectors I/O

Hadrons/A2AVectors.hpp

@@ -36,7 +36,7 @@ See the full license in the file "LICENSE" in the top level distribution directo
 BEGIN_HADRONS_NAMESPACE
 
 /******************************************************************************
- *               Classes to generate V & W all-to-all vectors                 *
+ *                 Class to generate V & W all-to-all vectors                 *
  ******************************************************************************/
 template <typename FImpl>
 class A2AVectorsSchurDiagTwo
@@ -70,6 +70,42 @@ private:
     SchurDiagTwoOperator<FMat, FermionField> op_;
 };
 
+/******************************************************************************
+ *                  Methods for V & W all-to-all vectors I/O                  *
+ ******************************************************************************/
+class A2AVectorsIo
+{
+public:
+    struct Record: Serializable
+    {
+        GRID_SERIALIZABLE_CLASS_MEMBERS(Record,
+                                        unsigned int, index);
+        Record(void): index(0) {}
+    };
+public:
+    template <typename Field>
+    static void write(const std::string fileStem, std::vector<Field> &vec, 
+                      const bool multiFile, const int trajectory = -1);
+    template <typename Field>
+    static void read(std::vector<Field> &vec, const std::string fileStem,
+                     const bool multiFile, const int trajectory = -1);
+private:
+    static inline std::string vecFilename(const std::string stem, const int traj, 
+                                          const bool multiFile)
+    {
+        std::string t = (traj < 0) ? "" : ("." + std::to_string(traj));
+
+        if (multiFile)
+        {
+            return stem + t;
+        }
+        else
+        {
+            return stem + t + ".bin";
+        }
+    }
+};
+
 /******************************************************************************
  *               A2AVectorsSchurDiagTwo template implementation               *
  ******************************************************************************/
@@ -217,6 +253,90 @@ void A2AVectorsSchurDiagTwo<FImpl>::makeHighModeW5D(FermionField &wout_4d,
     }
 }
 
+/******************************************************************************
+ *               all-to-all vectors I/O template implementation               *
+ ******************************************************************************/
+template <typename Field>
+void A2AVectorsIo::write(const std::string fileStem, std::vector<Field> &vec, 
+                         const bool multiFile, const int trajectory)
+{
+    Record       record;
+    GridBase     *grid = vec[0]._grid;
+    ScidacWriter binWriter(grid->IsBoss());
+    std::string  filename = vecFilename(fileStem, multiFile, trajectory);
+
+    if (multiFile)
+    {
+        std::string fullFilename;
+
+        for (unsigned int i = 0; i < vec.size(); ++i)
+        {
+            fullFilename = filename + "/elem" + std::to_string(i) + ".bin";
+
+            LOG(Message) << "Writing vector " << i << std::endl;
+            makeFileDir(fullFilename, grid);
+            binWriter.open(fullFilename);
+            record.index = i;
+            binWriter.writeScidacFieldRecord(vec[i], record);
+            binWriter.close();
+        }
+    }
+    else
+    {
+        makeFileDir(filename, grid);
+        binWriter.open(filename);
+        for (unsigned int i = 0; i < vec.size(); ++i)
+        {
+            LOG(Message) << "Writing vector " << i << std::endl;
+            record.index = i;
+            binWriter.writeScidacFieldRecord(vec[i], record);
+        }
+        binWriter.close();
+    }
+}
+
+template <typename Field>
+void A2AVectorsIo::read(std::vector<Field> &vec, const std::string fileStem, 
+                        const bool multiFile, const int trajectory)
+{
+    Record       record;
+    ScidacReader binReader;
+    std::string  filename = vecFilename(fileStem, multiFile, trajectory);
+
+    if (multiFile)
+    {
+        std::string fullFilename;
+
+        for (unsigned int i = 0; i < vec.size(); ++i)
+        {
+            fullFilename = filename + "/elem" + std::to_string(i) + ".bin";
+
+            LOG(Message) << "Reading vector " << i << std::endl;
+            binReader.open(fullFilename);
+            binReader.readScidacFieldRecord(vec[i], record);
+            binReader.close();
+            if (record.index != i)
+            {
+                HADRONS_ERROR(Io, "vector index mismatch");
+            }
+        }
+    }
+    else
+    {
+        binReader.open(filename);
+        for (unsigned int i = 0; i < vec.size(); ++i)
+        {
+            LOG(Message) << "Reading vector " << i << std::endl;
+            binReader.readScidacFieldRecord(vec[i], record);
+            if (record.index != i)
+            {
+                HADRONS_ERROR(Io, "vector index mismatch");
+            }
+        }
+        binReader.close();
+    }
+}
+
 END_HADRONS_NAMESPACE
 
 #endif // A2A_Vectors_hpp_

Hadrons/Modules.hpp

@@ -65,6 +65,7 @@
 #include <Hadrons/Modules/MScalarSUN/TrKinetic.hpp>
 #include <Hadrons/Modules/MIO/LoadEigenPack.hpp>
 #include <Hadrons/Modules/MIO/LoadNersc.hpp>
+#include <Hadrons/Modules/MIO/LoadA2AVectors.hpp>
 #include <Hadrons/Modules/MIO/LoadCosmHol.hpp>
 #include <Hadrons/Modules/MIO/LoadCoarseEigenPack.hpp>
 #include <Hadrons/Modules/MIO/LoadBinary.hpp>

Hadrons/Modules/MIO/LoadA2AVectors.cc

@@ -0,0 +1,7 @@
+#include <Hadrons/Modules/MIO/LoadA2AVectors.hpp>
+
+using namespace Grid;
+using namespace Hadrons;
+using namespace MIO;
+
+template class Grid::Hadrons::MIO::TLoadA2AVectors<FIMPL>;

Hadrons/Modules/MIO/LoadA2AVectors.hpp

@@ -0,0 +1,93 @@
+#ifndef Hadrons_MIO_LoadA2AVectors_hpp_
+#define Hadrons_MIO_LoadA2AVectors_hpp_
+
+#include <Hadrons/Global.hpp>
+#include <Hadrons/Module.hpp>
+#include <Hadrons/ModuleFactory.hpp>
+#include <Hadrons/A2AVectors.hpp>
+
+BEGIN_HADRONS_NAMESPACE
+
+/******************************************************************************
+ *                    Module to load all-to-all vectors                       *
+ ******************************************************************************/
+BEGIN_MODULE_NAMESPACE(MIO)
+
+class LoadA2AVectorsPar: Serializable
+{
+public:
+    GRID_SERIALIZABLE_CLASS_MEMBERS(LoadA2AVectorsPar,
+                                    std::string,  filestem,
+                                    bool,         multiFile,
+                                    unsigned int, size);
+};
+
+template <typename FImpl>
+class TLoadA2AVectors: public Module<LoadA2AVectorsPar>
+{
+public:
+    FERM_TYPE_ALIASES(FImpl,);
+public:
+    // constructor
+    TLoadA2AVectors(const std::string name);
+    // destructor
+    virtual ~TLoadA2AVectors(void) {};
+    // dependency relation
+    virtual std::vector<std::string> getInput(void);
+    virtual std::vector<std::string> getOutput(void);
+    // setup
+    virtual void setup(void);
+    // execution
+    virtual void execute(void);
+};
+
+MODULE_REGISTER_TMP(LoadA2AVectors, TLoadA2AVectors<FIMPL>, MIO);
+
+/******************************************************************************
+ *                      TLoadA2AVectors implementation                        *
+ ******************************************************************************/
+// constructor /////////////////////////////////////////////////////////////////
+template <typename FImpl>
+TLoadA2AVectors<FImpl>::TLoadA2AVectors(const std::string name)
+: Module<LoadA2AVectorsPar>(name)
+{}
+
+// dependencies/products ///////////////////////////////////////////////////////
+template <typename FImpl>
+std::vector<std::string> TLoadA2AVectors<FImpl>::getInput(void)
+{
+    std::vector<std::string> in;
+    
+    return in;
+}
+
+template <typename FImpl>
+std::vector<std::string> TLoadA2AVectors<FImpl>::getOutput(void)
+{
+    std::vector<std::string> out = {getName()};
+    
+    return out;
+}
+
+// setup ///////////////////////////////////////////////////////////////////////
+template <typename FImpl>
+void TLoadA2AVectors<FImpl>::setup(void)
+{
+    envCreate(std::vector<FermionField>, getName(), 1, par().size, 
+              envGetGrid(FermionField));
+}
+
+// execution ///////////////////////////////////////////////////////////////////
+template <typename FImpl>
+void TLoadA2AVectors<FImpl>::execute(void)
+{
+    auto &vec = envGet(std::vector<FermionField>, getName());
+
+    A2AVectorsIo::read(vec, par().filestem, par().multiFile, vm().getTrajectory());
+}
+
+END_MODULE_NAMESPACE
+
+END_HADRONS_NAMESPACE
+
+#endif // Hadrons_MIO_LoadA2AVectors_hpp_

Hadrons/Modules/MSolver/A2AVectors.hpp

@@ -51,7 +51,9 @@ public:
                                   std::string, noise,
                                   std::string, action,
                                   std::string, eigenPack,
-                                  std::string, solver);
+                                  std::string, solver,
+                                  std::string, output,
+                                  bool,        multiFile);
 };
 
 template <typename FImpl, typename Pack>
@@ -236,6 +238,13 @@ void TA2AVectors<FImpl, Pack>::execute(void)
         }
         stopTimer("W high mode");
     }
+
+    // I/O if necessary
+    if (!par().output.empty())
+    {
+        A2AVectorsIo::write(par().output + "_w", w, par().multiFile, vm().getTrajectory());
+        A2AVectorsIo::write(par().output + "_v", v, par().multiFile, vm().getTrajectory());
+    }
 }
 
 END_MODULE_NAMESPACE

Hadrons/modules.inc

@@ -64,7 +64,8 @@ modules_cc =\
   Modules/MIO/LoadBinary.cc \
   Modules/MIO/LoadNersc.cc \
   Modules/MIO/LoadCoarseEigenPack.cc \
-  Modules/MIO/LoadCosmHol.cc
+  Modules/MIO/LoadCosmHol.cc \
+  Modules/MIO/LoadA2AVectors.cc
 
 modules_hpp =\
   Modules/MContraction/Baryon.hpp \
@@ -134,6 +135,7 @@ modules_hpp =\
   Modules/MScalarSUN/TrKinetic.hpp \
   Modules/MIO/LoadEigenPack.hpp \
   Modules/MIO/LoadNersc.hpp \
+  Modules/MIO/LoadA2AVectors.hpp \
   Modules/MIO/LoadCosmHol.hpp \
   Modules/MIO/LoadCoarseEigenPack.hpp \
   Modules/MIO/LoadBinary.hpp