From 13e54858634e01da07ef813c228bb7f015bdfcf4 Mon Sep 17 00:00:00 2001
From: chn <chn@chn.moe>
Date: Tue, 13 Aug 2024 23:44:54 +0800
Subject: [PATCH] packages.biu: add Eigen serialization

---
 flake.nix                           |  10 ---
 packages/biu/include/biu.hpp        |   3 +-
 packages/biu/include/biu/eigen.hpp  |  88 +++++++++++++---------
 packages/biu/include/biu/eigen.tpp  | 111 +++++++++++++++++++++++-----
 packages/ufo/include/ufo/fold.hpp   |   1 +
 packages/ufo/include/ufo/solver.hpp |  13 +---
 6 files changed, 152 insertions(+), 74 deletions(-)

diff --git a/flake.nix b/flake.nix
index 23a9154d..83137c00 100644
--- a/flake.nix
+++ b/flake.nix
@@ -83,16 +83,6 @@
           system = "x86_64-linux";
           config.allowUnfree = true;
           overlays = [ inputs.self.overlays.default ];
-          crossOverlays = [(final: prev:
-          {
-            boost = prev.boost.override { zstd = null; };
-            magic-enum = prev.magic-enum.overrideAttrs (prev: { cmakeFlags = prev.cmakeFlags ++
-              [ "-DMAGIC_ENUM_OPT_BUILD_EXAMPLES=OFF" "-DMAGIC_ENUM_OPT_BUILD_TESTS=OFF" ]; });
-            range-v3 = prev.range-v3.overrideAttrs (prev: { cmakeFlags = prev.cmakeFlags ++
-              [ "-DRANGE_V3_DOCS=OFF" "-DRANGE_V3_TESTS=OFF" "-DRANGE_V3_EXAMPLES=OFF" ]; });
-            abseil-cpp = prev.abseil-cpp.overrideAttrs (prev: { buildInputs = prev.buildInputs ++
-              [ final.windows.pthreads ]; });
-          })];
         });
         default = inputs.nixpkgs.legacyPackages.x86_64-linux.writeText "systems"
           (builtins.concatStringsSep "\n" (builtins.map
diff --git a/packages/biu/include/biu.hpp b/packages/biu/include/biu.hpp
index d0801e3f..ac4a5fa3 100644
--- a/packages/biu/include/biu.hpp
+++ b/packages/biu/include/biu.hpp
@@ -5,7 +5,8 @@
 # include <biu/concepts.tpp>
 # include <biu/string.tpp>
 # include <biu/format.tpp>
+# include <biu/eigen.hpp>
 // # include <biu/logger.tpp>
 // # include <biu/smartref.tpp>
 
-// # include <biu/eigen.hpp>
+
diff --git a/packages/biu/include/biu/eigen.hpp b/packages/biu/include/biu/eigen.hpp
index a648bce7..a5959914 100644
--- a/packages/biu/include/biu/eigen.hpp
+++ b/packages/biu/include/biu/eigen.hpp
@@ -6,49 +6,69 @@
 
 namespace biu
 {
-  namespace detail_::eigen
+  namespace eigen
   {
-    // user-specified size of destination container: dynamic, unspecified(use default), or fixed
-    constexpr std::size_t dynamicSize = std::dynamic_extent, unspecifiedSize = std::dynamic_extent - 1;
-    static_assert(std::dynamic_extent == std::numeric_limits<std::size_t>::max());
+    namespace detail_
+    {
+      // user-specified size of destination container: dynamic, unspecified(use default), or fixed
+      constexpr std::size_t dynamicSize = std::dynamic_extent, unspecifiedSize = std::dynamic_extent - 1;
+      static_assert(std::dynamic_extent == std::numeric_limits<std::size_t>::max());
 
-    // supported types of standard containers
-		template <typename T, typename Scalar> struct SpecializationOfArrayHelper : std::false_type {};
-		template <typename Scalar, std::size_t N>
-      struct SpecializationOfArrayHelper<std::array<Scalar, N>, Scalar> : std::true_type {};
-		template <typename Scalar, std::size_t N>
-      struct SpecializationOfArrayHelper<std::array<Scalar, N>, void> : std::true_type {};
-    template <typename T, typename Scalar = void> concept SpecializationOfArray =
-      SpecializationOfArrayHelper<T, Scalar>::value;
-    template <typename T, typename Scalar> concept StandardContainer =
-      SpecializationOf<T, std::vector, Scalar> || SpecializationOfArray<T, Scalar>;
+      // supported types of standard containers
+      template <typename T, typename Scalar> struct SpecializationOfArrayHelper : std::false_type {};
+      template <typename Scalar, std::size_t N>
+        struct SpecializationOfArrayHelper<std::array<Scalar, N>, Scalar> : std::true_type {};
+      template <typename Scalar, std::size_t N>
+        struct SpecializationOfArrayHelper<std::array<Scalar, N>, void> : std::true_type {};
+      template <typename T, typename Scalar = void> concept SpecializationOfArray =
+        SpecializationOfArrayHelper<T, Scalar>::value;
+      template <typename T, typename Scalar> concept StandardContainer =
+        SpecializationOf<T, std::vector, Scalar> || SpecializationOfArray<T, Scalar>;
+
+      // deduce the size of the destination Eigen container
+      // if no size is specified, convert std::vector to dynamic-size Eigen::Vector,
+      //  std::array to fixed-size Eigen::Vector;
+      // if size is std::dynamic_extent, always convert to dynamic-size Eigen::Vector
+      // if size is specified as a number, convert to fixed-size Eigen::Vector if specified size equals the size of the
+      //  input, otherwise throw an error
+      // return deduced size if the size is deducible in compile time, otherwise return Empty
+      template <std::size_t ToSize, typename Container> auto deduce_eigen_size();
+
+      // helper operator| to specify the size of the destination container
+      template <std::size_t Row, std::size_t Col> struct ToEigenHelper {};
+    }
 
-    // helper operator| to specify the size of the destination container
     // usage: some_value | toEigen<Row, Col>
-    template <std::size_t Row, std::size_t Col> struct ToEigenHelper {};
-    template <std::size_t Row = unspecifiedSize, std::size_t Col = unspecifiedSize>
-      inline constexpr ToEigenHelper<Row, Col> toEigen;
-    // convert 1D standard container to Eigen::Vector
-    // if no size is specified, convert std::vector to dynamic-size Eigen::Vector,
-    //  std::array to fixed-size Eigen::Vector;
-    // if size is std::dynamic_extent, always convert to dynamic-size Eigen::Vector
-    // if size is specified as a number, convert to fixed-size Eigen::Vector if specified size equals the size of the
-    //  input, otherwise throw an error
-    template <template <int N> typename Callback, std::size_t ToSize> auto deduce_eigen_size(auto&& container);
-    template <Arithmetic T, StandardContainer<T> From, std::size_t ToSize> auto operator|
-      (const From&, const ToEigenHelper<ToSize, unspecifiedSize>&);
+    template <std::size_t Row = detail_::unspecifiedSize, std::size_t Col = detail_::unspecifiedSize>
+      inline constexpr detail_::ToEigenHelper<Row, Col> toEigen;
+
+    // convert 1D standard container to Eigen::Matrix, the second argument should always be unspecified
+    template <Arithmetic T, detail_::StandardContainer<T> From, std::size_t ToSize> auto operator|
+      (const From&, const detail_::ToEigenHelper<ToSize, detail_::unspecifiedSize>&);
+
     // convert 2D standard container to Eigen::Matrix
-    // the same rules as above apply
-    // besides, all rows must have the same size, otherwise throw an error
     template
     <
-      Arithmetic T, StandardContainer<T> FromPerRow, StandardContainer<FromPerRow> From,
+      Arithmetic T, detail_::StandardContainer<T> FromPerRow, detail_::StandardContainer<FromPerRow> From,
       std::size_t ToRow, std::size_t ToCol
     >
-      auto operator|(const From&, const ToEigenHelper<ToRow, ToCol>&);
+      auto operator|(const From&, const detail_::ToEigenHelper<ToRow, ToCol>&);
 
-    // TODO: implement fromEigen
+    // test if a class is an eigen matrix
+    namespace detail_
+    {
+      template <typename Matrix> class EigenMatrix : public std::false_type {};
+      template <typename Scalar, int Rows, int Cols, int Options>
+        class EigenMatrix<Eigen::Matrix<Scalar, Rows, Cols, Options>> : public std::true_type {};
+    }
+    template <typename Matrix> concept EigenMatrix = detail_::EigenMatrix<Matrix>::value;
   }
-
-  inline namespace eigen { using detail_::eigen::toEigen; using detail_::eigen::operator|; }
+  using eigen::toEigen, eigen::operator|, eigen::EigenMatrix;
+}
+
+// archive a matrix
+namespace Eigen
+{
+  template <typename Matrix> constexpr auto serialize(auto & archive, Matrix& matrix)
+    requires biu::EigenMatrix<std::remove_cvref_t<Matrix>>;
 }
diff --git a/packages/biu/include/biu/eigen.tpp b/packages/biu/include/biu/eigen.tpp
index b6c7fc6a..a7ba1428 100644
--- a/packages/biu/include/biu/eigen.tpp
+++ b/packages/biu/include/biu/eigen.tpp
@@ -1,29 +1,100 @@
 # pragma once
 # include <biu/eigen.hpp>
-// TODO: fix biu::logger
-// # include <biu/logger.hpp>
+# include <biu/common.hpp>
 # include <range/v3/view.hpp>
+# include <zpp_bits.h>
 
-namespace biu
+namespace biu::eigen
 {
-  namespace detail_::eigen
+  template <std::size_t ToSize, typename Container> auto detail_::deduce_eigen_size()
   {
-    template <template <int N> typename Callback, std::size_t ToSize> auto deduce_eigen_size(auto&& from)
-    {
-      if constexpr (ToSize == dynamicSize)
-        return Callback<Eigen::Dynamic>()(from.data(), from.size());
-      else if constexpr (ToSize == unspecifiedSize)
-        if constexpr (SpecializationOfArray<decltype(from)>)
-          return Callback<from.size()>()(from.data());
-        else
-          return Callback<Eigen::Dynamic>()(from.data(), from.size());
+    if constexpr (ToSize == detail_::dynamicSize) return Empty{};
+    else if constexpr (ToSize == detail_::unspecifiedSize)
+      if constexpr (SpecializationOfArray<Container>) return Container::size();
+      else return Empty{};
+    else
+      if constexpr (SpecializationOfArray<Container>)
+        if (Container::size() == ToSize) return ToSize;
+        else throw std::invalid_argument("The size of the destination Eigen container mismatches the input container");
+      else return ToSize;
+  }
+  template <Arithmetic T, detail_::StandardContainer<T> From, std::size_t ToSize> auto operator|
+    (const From& from, const detail_::ToEigenHelper<ToSize, detail_::unspecifiedSize>&)
+  {
+    // dynamic size
+    if (detail_::deduce_eigen_size<ToSize, From>() == Empty{})
+      return Eigen::Vector<T, Eigen::Dynamic>(from.data(), from.size());
+    // fixed size
+    else
+      // from vector, vector.size() == ToSize
+      if (SpecializationOf<From, std::vector, T> && from.size() == ToSize)
+        return Eigen::Vector<T, ToSize>(from.data());
       else
-        if (from.size() != ToSize)
-          // TODO: use biu::logger
-          throw std::invalid_argument("biu::toEigen: size mismatch");
-        else
-          return Callback<ToSize>()(from.data());
-    }
-    // TODO: implement 2D case
+        throw std::invalid_argument("The size of the destination Eigen container mismatches the input container");
+  }
+  template
+  <
+    Arithmetic T, detail_::StandardContainer<T> FromPerRow, detail_::StandardContainer<FromPerRow> From,
+    std::size_t ToRow, std::size_t ToCol
+  >
+    auto operator|(const From& from, const detail_::ToEigenHelper<ToRow, ToCol>&)
+  {
+    auto nRow = detail_::deduce_eigen_size<ToRow, From>();
+    auto nCol = detail_::deduce_eigen_size<ToCol, FromPerRow>();
+
+    // ensure all rows have the same size
+    std::vector<std::size_t> size_in_each_row;
+    // each row is a std::array, they must have the same size
+    if constexpr (!SpecializationOf<FromPerRow, std::vector, T>)
+      size_in_each_row.push_back(FromPerRow::size());
+    else
+      size_in_each_row = from
+        | ranges::views::transform([](const auto& row) { return row.size(); })
+        | ranges::views::unique
+        | ranges::to<std::vector<std::size_t>>;
+    if (size_in_each_row.size() > 1)
+      throw std::invalid_argument("The sizes of the rows of the input container are not the same");
+    else if (size_in_each_row.empty()) size_in_each_row.push_back(0);
+
+    // ensure specified size is consistent with the input
+    if (nRow != Empty{} && SpecializationOf<From, std::vector> && from.size() != nRow)
+      throw std::invalid_argument("The size of the destination Eigen container mismatches the input container");
+    if (nCol != Empty{} && size_in_each_row[0] != nCol)
+      throw std::invalid_argument("The size of the destination Eigen container mismatches the input container");
+
+    // copy all data into a single vector
+    auto data = from | ranges::views::join | ranges::to<std::vector<T>>;
+
+    // dynamic row and dynamic col
+    if constexpr (nRow == Empty{} && nCol == Empty{})
+      return Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor | Eigen::AutoAlign>
+        (data.data(), from.size(), size_in_each_row[0]);
+    // fixed row and dynamic col
+    else if constexpr (nRow != Empty{} && nCol == Empty{})
+      return Eigen::Matrix<T, nRow, Eigen::Dynamic, Eigen::RowMajor | Eigen::AutoAlign>
+        (data.data(), size_in_each_row[0]);
+    // dynamic row and fixed col
+    else if constexpr (nRow == Empty{} && nCol != Empty{})
+      return Eigen::Matrix<T, Eigen::Dynamic, nCol, Eigen::RowMajor | Eigen::AutoAlign>
+        (data.data(), from.size());
+    // fixed row and fixed col
+    else
+      return Eigen::Matrix<T, ToRow, ToCol, Eigen::RowMajor | Eigen::AutoAlign>
+        (data.data());
   }
 }
+template <typename Matrix> constexpr auto Eigen::serialize(auto & archive, Matrix& matrix)
+  requires biu::EigenMatrix<std::remove_cvref_t<Matrix>>
+  {
+    auto nRow = matrix.rows(), nCol = matrix.cols();
+    zpp::bits::errc result;
+    if (result.code == std::errc{} && Matrix::CompileTimeTraits::RowsAtCompileTime == Eigen::Dynamic)
+      result = archive(nRow);
+    if (result.code == std::errc{} && Matrix::CompileTimeTraits::ColsAtCompileTime == Eigen::Dynamic)
+      result = archive(nCol);
+    if (archive.kind() == zpp::bits::kind::in)
+      matrix.resize(nRow, nCol);
+    if (result.code == std::errc{})
+      result = archive(matrix.data(), matrix.size());
+    return result;
+  }
diff --git a/packages/ufo/include/ufo/fold.hpp b/packages/ufo/include/ufo/fold.hpp
index 6ce66dc6..4c132061 100644
--- a/packages/ufo/include/ufo/fold.hpp
+++ b/packages/ufo/include/ufo/fold.hpp
@@ -18,6 +18,7 @@ namespace ufo
       struct OutputType
       {
         std::vector<Eigen::Vector3d> Qpoints;
+        using serialize = zpp::bits::members<1>;
         void write(std::string filename) const;
       };
     protected:
diff --git a/packages/ufo/include/ufo/solver.hpp b/packages/ufo/include/ufo/solver.hpp
index 10ee2f1c..7f473dad 100644
--- a/packages/ufo/include/ufo/solver.hpp
+++ b/packages/ufo/include/ufo/solver.hpp
@@ -30,17 +30,12 @@
 
 namespace Eigen
 {
-  // template <typename Scalar, int Rows, int Cols> constexpr inline auto serialize
-  //   (auto & archive, Eigen::Matrix<Scalar, Rows, Cols>& matrix)
-  //   { return archive(std::span(matrix.data(), matrix.size())); }
-  constexpr inline auto serialize(auto & archive, Eigen::Matrix3d& matrix)
+  template <typename Matrix> constexpr inline auto serialize
+    (auto & archive, Matrix& matrix)
     { return archive(std::span(matrix.data(), matrix.size())); }
-  constexpr inline auto serialize(auto & archive, const Eigen::Matrix3d& matrix)
+  template <typename Scalar, int Rows, int Cols> constexpr inline auto serialize
+    (auto & archive, const Eigen::Matrix<Scalar, Rows, Cols>& matrix)
     { return archive(std::span(matrix.data(), matrix.size())); }
-  constexpr inline auto serialize(auto & archive, Eigen::Vector3d& vector)
-    { return archive(std::span(vector.data(), vector.size())); }
-  constexpr inline auto serialize(auto & archive, const Eigen::Vector3d& vector)
-    { return archive(std::span(vector.data(), vector.size())); }
 }
 
 namespace ufo