Extract exporting (and importing) to Exporter. Currently complex is failing

Author: ProExpertProg

cpp/CMakeLists.txt

@@ -59,8 +59,9 @@ target_include_directories(adaptive_hermite_refinement PUBLIC ${CMAKE_CURRENT_SO
 target_link_libraries(adaptive_hermite_refinement argparse mdspan)
 
 add_library(src-lib OBJECT lib/Naive.cpp lib/Naive.hpp lib/HermiteRunner.cpp lib/HermiteRunner.hpp
-        lib/Transformer.cpp
-        lib/Transformer.hpp)
+        lib/Transformer.hpp lib/Transformer.cpp
+        lib/Exporter.hpp lib/Exporter.cpp
+)
 target_include_directories(src-lib PUBLIC lib/)
 target_link_libraries(src-lib mdspan fftw-cpp cnpy spdlog::spdlog)
 

cpp/lib/Exporter.cpp

@@ -0,0 +1,55 @@
+#include "Exporter.hpp"
+#include "Transformer.hpp"
+#include <cnpy.h>
+#include <filesystem>
+
+namespace ahr {
+namespace fs = std::filesystem;
+void Exporter::exportTo(fs::path const &filename, Grid::View::C_XY cView) {
+  // fft overwrites the input, so we need to copy it to a temporary buffer
+  auto tempK = grid.cBufXY();
+  auto temp = grid.rBufXY();
+
+  // Copy and also normalize
+  tf.normalize(cView, tempK);
+
+  // Backwards FFT
+  tf.bfft(tempK, temp);
+
+  // Write the real buffer to file
+  exportTo(filename, temp);
+}
+
+void Exporter::exportTo(fs::path const &filename, Grid::View::R_XY rView) {
+  // TODO Copy the data to a layout-right buffer
+  //  stdex::mdarray<Real, stdex::dextents<size_t, 2u>> rArray{rView.extents()};
+  //  grid.for_each_xy([&](Dim x, Dim y) { rArray(x, y) = rView(x, y); });
+  //  cnpy::npy_save(prefix_dir / filename, rArray.data(), {grid.X, grid.Y}, "w");
+
+  // Dimensions are switched because we use layout_left
+  auto const path = filename.is_absolute() ? filename : prefix_dir / filename;
+  cnpy::npy_save(path, rView.data_handle(), {grid.Y, grid.X}, "w");
+}
+
+NpyMdspan
+Exporter::importReal(const fs::path &filename) { // NOLINT(*-convert-member-functions-to-static)
+  auto const path = filename.is_absolute() ? filename : prefix_dir / filename;
+  return NpyMdspan{cnpy::npy_load(path)};
+}
+
+void Exporter::importReal(const fs::path &filename, Grid::View::R_XY rView) {
+  auto npy = importReal(filename);
+  if (!npy.valid()) { throw std::runtime_error("Invalid npy file"); }
+  if (npy.view().extents() != rView.extents()) { throw std::runtime_error("Incompatible extents"); }
+
+  // Copy the data
+  grid.for_each_xy([&](Dim x, Dim y) { rView(x, y) = npy.view()(x, y); });
+}
+
+Grid::Buf::R_XY Exporter::importRealBuf(const fs::path &filename) {
+  auto rBuf = grid.rBufXY();
+  importReal(filename, rBuf.to_mdspan());
+  return rBuf;
+}
+
+} // namespace ahr

cpp/lib/Exporter.hpp

@@ -0,0 +1,51 @@
+#pragma once
+
+#include "NpyMdspan.hpp"
+#include "grid.hpp"
+#include <filesystem>
+#include <optional>
+
+namespace ahr {
+namespace fs = std::filesystem;
+class Transformer;
+
+/// This class is responsible for exporting and importing .npy buffers.
+/// It can automatically transform complex buffers to real before exporting.
+/// Any relative path is interpreted as relative to the prefix_dir.
+class Exporter {
+  Grid const &grid;
+  Transformer const &tf;
+  fs::path prefix_dir;
+
+public:
+  Exporter(Grid const &grid, Transformer const &transformer,
+           std::optional<fs::path> prefix_dir = std::nullopt)
+      : grid(grid), tf(transformer) {
+    if (prefix_dir) {
+      this->prefix_dir = fs::canonical(*prefix_dir);
+    } else if (auto prefix_dir_env = std::getenv("EXPORT_PREFIX_DIR"); prefix_dir_env) {
+      this->prefix_dir = fs::canonical(prefix_dir_env);
+    } else {
+      this->prefix_dir = fs::current_path();
+    }
+
+    fs::create_directories(this->prefix_dir);
+  }
+
+  /// Export the complex view to file by transforming it to real first.
+  void exportTo(fs::path const &filename, Grid::View::C_XY cView);
+
+  /// Export the real buffer to file.
+  void exportTo(fs::path const &filename, Grid::View::R_XY rView);
+
+  /// Import a real buffer from file, and write it to the given view.
+  void importReal(const fs::path &filename, Grid::View::R_XY rView);
+
+  /// Import a real buffer from file, and return it as an owning NpyMdspan.
+  [[nodiscard]] NpyMdspan importReal(fs::path const &filename);
+
+  /// Import a real buffer from file, and return it as a real buffer.
+  [[nodiscard]] Grid::Buf::R_XY importRealBuf(fs::path const &filename);
+
+};
+} // namespace ahr

cpp/lib/Naive.cpp

@@ -452,15 +452,15 @@ void Naive::run(Dim N, Dim saveInterval) {
 void Naive::exportTimestep(Dim t) {
   std::ostringstream oss;
   oss << "a_par_t" << t << ".npy";
-  exportToNpy(oss.str(), Grid::sliceXY(moments_K, A_PAR));
+  exporter.exportTo(oss.str(), Grid::sliceXY(moments_K, A_PAR));
 
   oss.str("");
   oss << "phi_t" << t << ".npy";
-  exportToNpy(oss.str(), phi_K);
+  exporter.exportTo(oss.str(), phi_K);
 
   oss.str("");
   oss << "uekpar_t" << t << ".npy";
-  exportToNpy(oss.str(), ueKPar_K);
+  exporter.exportTo(oss.str(), ueKPar_K);
 }
 
 Real Naive::updateTimestep(Real dt, Real tempDt, bool noInc, Real relative_error) const {
@@ -519,24 +519,6 @@ Naive::Buf::C_XY Naive::halfBracket(Naive::DxDy<View::R_XY> derOp1,
   return br_K;
 }
 
-void Naive::exportToNpy(std::string path, View::R_XY view) const {
-  // Coordinates are flipped because we use layout_left
-  cnpy::npy_save(std::move(path), view.data_handle(), {g.Y, g.X}, "w");
-}
-
-void Naive::exportToNpy(std::string path, View::C_XY view) const {
-  // fft overwrites the input, so we need to copy it to a temporary buffer
-  auto tempK = g.cBufXY();
-  auto temp = g.rBufXY();
-
-  g.for_each_kxky([&](Dim kx, Dim ky) { tempK(kx, ky) = view(kx, ky); });
-
-  tf.bfft(tempK.to_mdspan(), temp.to_mdspan());
-  tf.normalize(temp.to_mdspan(), temp.to_mdspan());
-
-  exportToNpy(std::move(path), temp.to_mdspan());
-}
-
 Naive::Energies Naive::calculateEnergies() const {
   Energies e{};
   g.for_each_kxky([&](Dim kx, Dim ky) {

cpp/lib/Naive.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "Exporter.hpp"
 #include "HermiteRunner.hpp"
 #include "Transformer.hpp"
 #include "constants.hpp"
@@ -34,9 +35,10 @@ class Naive : public ahr::HermiteRunner {
   mdarray<Real, dextents<Dim, 2u>> getFinalAPar() override;
 
 
-private:
   Grid g;
   Transformer tf{g};
+  Exporter exporter{g, tf};
+private:
 
   using View = Grid::View;
   using Buf = Grid::Buf;
@@ -224,14 +226,6 @@ class Naive : public ahr::HermiteRunner {
 private:
   Real updateTimestep(Real dt, Real tempDt, bool noInc, Real relative_error) const;
 
-public:
-  // TODO(luka) separate exporting utility
-  void exportToNpy(std::string path, View::R_XY view) const;
-
-  // Will also normalize and inverseFFT
-  void exportToNpy(std::string path, View::C_XY view) const;
-
-private:
   void exportTimestep(Dim t);
 };
 }; // namespace ahr

cpp/lib/NpyMdspan.hpp

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "grid.hpp"
+#include <cnpy.h>
+#include <filesystem>
+
+namespace ahr {
+namespace fs = std::filesystem;
+
+/// Owning holder of a npy array with convenience to see it as an mdspan.
+/// We can use this to avoid needlessly copying into an mdarray.
+class NpyMdspan {
+  cnpy::NpyArray array_;
+
+public:
+  explicit NpyMdspan(cnpy::NpyArray array) : array_(std::move(array)) {}
+
+  // Layout-right
+  using ViewXY = stdex::mdspan<Real, stdex::dextents<size_t, 2u>>;
+
+  // TODO(luka) const view
+  Grid::View::R_XY view() {
+    // Reverse the dimensions
+    auto shape = array_.shape;
+    std::reverse(shape.begin(), shape.end());
+    std::span<size_t, 2> const extents{shape.data(), 2};
+
+    // Return a layout_left view
+    return Grid::View::R_XY{array_.data<Real>(), extents};
+  }
+
+  [[nodiscard]] bool valid() const { return array_.word_size == sizeof(Real); }
+};
+
+} // namespace ahr

cpp/test/CMakeLists.txt

@@ -20,6 +20,7 @@ endfunction()
 
 add_viriato_test(grid.cpp)
 add_viriato_test(transformer.cpp)
+add_viriato_test(exporter.cpp)
 
 add_viriato_test(naive-energies.cpp)
 add_viriato_test(naive-moments.cpp)

cpp/test/exporter.cpp

@@ -0,0 +1,63 @@
+#include "Exporter.hpp"
+#include "Transformer.hpp"
+#include "debug.hpp"
+#include "util.hpp"
+#include <gtest/gtest.h>
+
+namespace ahr {
+
+class TestExporter : public ::testing::Test {
+protected:
+  Grid grid{5, 16, 24};
+  Transformer tf{grid};
+  fs::path const tmp_dir{fs::temp_directory_path()};
+  Exporter exporter{grid, tf, tmp_dir};
+
+  TestExporter() { tf.init(); }
+};
+
+TEST_F(TestExporter, RoundTripReal) {
+
+  auto rBuf = grid.rBufXY();
+  grid.for_each_xy([&](Dim x, Dim y) { rBuf(x, y) = Real(x + y * grid.X); });
+
+  exporter.exportTo("test.npy", rBuf);
+
+  auto const path = tmp_dir / "test.npy";
+  ASSERT_TRUE(fs::exists(path));
+
+  // Import buffer and compare
+  auto rBuf2 = exporter.importRealBuf("test.npy");
+  // No math, tolerance is 0
+  EXPECT_THAT(rBuf2.to_mdspan(), MdspanElementsAllClose(rBuf.to_mdspan(), 0.0));
+
+  // Import into npy view and compare
+  auto rNpy = exporter.importReal("test.npy");
+  EXPECT_THAT(rNpy.view(), MdspanElementsAllClose(rBuf.to_mdspan(), 0.0));
+}
+
+TEST_F(TestExporter, RoundTripComplex) {
+  auto cBuf = grid.cBufXY();
+  grid.for_each_kxky([&](Dim kx, Dim ky) {
+    cBuf(kx, ky) = Complex(Real(kx + ky * grid.KX), Real(kx - ky * grid.KX));
+  });
+
+  // Use absolute path this time
+  exporter.exportTo(tmp_dir / "test.npy", cBuf);
+
+  auto const path = tmp_dir / "test.npy";
+  ASSERT_TRUE(fs::exists(path));
+
+  // Import buffer and compare
+  auto rBuf = exporter.importReal("test.npy");
+  auto rBuf2 = grid.rBufXY();
+  tf.bfft(cBuf, rBuf2);
+
+  // No math, tolerance is 0
+  EXPECT_THAT(rBuf.view(), MdspanElementsAllClose(rBuf2.to_mdspan(), 0.0))
+      << "rBuf:\n"
+      << rBuf.view() << "\nrBuf2:\n"
+      << rBuf2.to_mdspan();
+}
+
+} // namespace ahr

cpp/test/naive-moments.cpp

@@ -27,35 +27,13 @@ constexpr Real ATOL = LOW_PRECISION_ATOL;
 template <param_like Param> class NaiveMomentsBase : public NaiveTester<Param> {
 protected:
   using Base = NaiveTester<Param>;
-  std::string getFilename(Dim m) {
+  fs::path getFilename(Dim m) {
     auto const p = Base::GetParam();
 
     // Tests are run from inside the `test` directory
     auto const filename = p.to_param_str() + "_m" + std::to_string(m) + ".npy";
     return fs::current_path() / "_test_data" / filename;
   }
-
-  // Owning holder of a npy array with convenience to see it as an mdspan.
-  // We can use this to avoid needlessly copying into an mdarray.
-  class NpyMdspan {
-    cnpy::NpyArray array_;
-
-  public:
-    explicit NpyMdspan(cnpy::NpyArray array) : array_(std::move(array)) {}
-
-    // TODO(luka) const view
-    Grid::View::R_XY view() {
-      std::span<size_t, 2> const extents{array_.shape.data(), 2};
-      return Grid::View::R_XY{array_.data<Real>(), extents};
-    }
-
-    [[nodiscard]] bool valid() const { return array_.word_size == sizeof(Real); }
-  };
-
-  NpyMdspan readMoment(Dim m) {
-    auto const filename = getFilename(m);
-    return NpyMdspan{cnpy::npy_load(filename)};
-  }
 };
 
 using MomentParam = WithEquilibrium<WithDiffusion<NaiveParam>>;
@@ -79,9 +57,9 @@ TEST_P(NaiveMoments, CheckMoments) {
   for (Dim m = 0; m < p.M; m++) {
     // To update values, uncomment these 2 lines
     // std::cout << "WARNING!: Overwriting " << getFilename(m) << std::endl;
-    // naive.exportToNpy(getFilename(m), naive.getMoment(m));
+    // naive.exporter.exportTo(getFilename(m), naive.getMoment(m));
 
-    auto npy = readMoment(m);
+    auto npy = naive.exporter.importReal(getFilename(m));
     ASSERT_TRUE(npy.valid());
     auto const max_val = std::ranges::max(std::span(npy.view().data_handle(), npy.view().size()));
     auto const n = static_cast<Real>(p.N);