packages.ufo: add more data to unfold output

packages/biu/include/biu.hpp

@@ -11,4 +11,5 @@
 # include <biu/smartref.tpp>
 # include <biu/fft.tpp>
 # include <biu/yaml.tpp>
+# include <biu/serialize.tpp>
 # include <range/v3/all.hpp>

packages/biu/include/biu/common.hpp

@@ -95,14 +95,6 @@ namespace biu
     }
     template <detail_::ExecMode Mode = {}> detail_::ExecResult<Mode> exec(detail_::ExecInput<Mode> input);
 
-    static_assert(sizeof(char) == sizeof(std::byte));
-    template <typename Char = std::byte, typename T> requires (std::same_as<Char, std::byte>)
-      std::vector<std::byte> serialize(const T& data);
-    template <typename Char = std::byte, typename T> requires (std::same_as<Char, char>)
-      std::string serialize(const T& data);
-    template <typename T> T deserialize(const std::string& serialized_data);
-    template <typename T> T deserialize(const std::vector<std::byte>& serialized_data);
-
     template <typename Array> concurrencpp::generator<std::pair<Array, std::size_t>> sequence(Array from, Array to);
     template <typename Array> concurrencpp::generator<std::pair<Array, std::size_t>> sequence(Array to);
 
@@ -128,6 +120,6 @@ namespace biu
   }
   using common::hash, common::unused, common::block_forever, common::is_interactive, common::env, common::int128_t,
     common::uint128_t, common::Empty, common::CaseInsensitiveStringLessComparator, common::RemoveMemberPointer,
-    common::MoveQualifiers, common::FallbackIfNoTypeDeclared, common::exec, common::serialize, common::deserialize,
-    common::sequence, common::read, common::toLvalue;
+    common::MoveQualifiers, common::FallbackIfNoTypeDeclared, common::exec, common::sequence, common::read,
+    common::toLvalue;
 }

packages/biu/include/biu/common.tpp

@@ -23,31 +23,6 @@ namespace biu::common
     );
   }
 
-  template <typename Char, typename T> requires (std::same_as<Char, std::byte>)
-    std::vector<std::byte> serialize(const T& data)
-  {
-    auto [serialized_data, out] = zpp::bits::data_out();
-    out(data).or_throw();
-    return serialized_data;
-  }
-  template <typename Char, typename T> requires (std::same_as<Char, char>) std::string serialize(const T& data)
-  {
-    auto serialized_data = serialize<std::byte>(data);
-    return {reinterpret_cast<const char*>(serialized_data.data()), serialized_data.size()};
-  }
-  template <typename T> T deserialize(const std::vector<std::byte>& serialized_data)
-  {
-    auto in = zpp::bits::in(serialized_data);
-    T data;
-    in(data).or_throw();
-    return data;
-  }
-  template <typename T> T deserialize(const std::string& serialized_data)
-  {
-    auto begin = reinterpret_cast<const std::byte*>(serialized_data.data()), end = begin + serialized_data.size();
-    return deserialize<T>(std::vector<std::byte>{begin, end});
-  }
-
   template <typename Array> concurrencpp::generator<std::pair<Array, std::size_t>> sequence(Array from, Array to)
   {
 # ifndef NDEBUG

packages/biu/include/biu/serialize.hpp

@@ -0,0 +1,24 @@
+# pragma once
+# include <concurrencpp/concurrencpp.h>
+# include <type_traits>
+# include <fmt/format.h>
+# include <magic_enum_all.hpp>
+# include <biu/concepts.hpp>
+
+namespace biu
+{
+  static_assert(sizeof(char) == sizeof(std::byte));
+  template <typename Char = std::byte, typename T> requires (std::same_as<Char, std::byte>)
+    std::vector<std::byte> serialize(const T& data);
+  template <typename Char = std::byte, typename T> requires (std::same_as<Char, char>)
+    std::string serialize(const T& data);
+  template <typename T> T deserialize(const std::string& serialized_data);
+  template <typename T> T deserialize(const std::vector<std::byte>& serialized_data);
+}
+
+namespace std
+{
+  // TODO: remove at clang19
+  template <typename Complex> constexpr auto serialize(auto & archive, Complex& complex)
+    requires biu::SpecializationOf<std::remove_cvref_t<Complex>, std::complex>;
+}

packages/biu/include/biu/serialize.tpp

@@ -0,0 +1,51 @@
+# pragma once
+# include <zpp_bits.h>
+# include <biu/serialize.hpp>
+
+namespace biu
+{
+  template <typename Char, typename T> requires (std::same_as<Char, std::byte>)
+    std::vector<std::byte> serialize(const T& data)
+  {
+    auto [serialized_data, out] = zpp::bits::data_out();
+    out(data).or_throw();
+    return serialized_data;
+  }
+  template <typename Char, typename T> requires (std::same_as<Char, char>) std::string serialize(const T& data)
+  {
+    auto serialized_data = serialize<std::byte>(data);
+    return {reinterpret_cast<const char*>(serialized_data.data()), serialized_data.size()};
+  }
+  template <typename T> T deserialize(const std::vector<std::byte>& serialized_data)
+  {
+    auto in = zpp::bits::in(serialized_data);
+    T data;
+    in(data).or_throw();
+    return data;
+  }
+  template <typename T> T deserialize(const std::string& serialized_data)
+  {
+    auto begin = reinterpret_cast<const std::byte*>(serialized_data.data());
+    auto end = begin + serialized_data.size();
+    return deserialize<T>(std::vector<std::byte>{begin, end});
+  }
+}
+
+namespace std
+{
+  template <typename Complex> constexpr auto serialize(auto & archive, Complex& complex)
+    requires biu::SpecializationOf<std::remove_cvref_t<Complex>, std::complex>
+  {
+    if constexpr (std::integral<decltype(archive())>) return 2;
+    else
+    {
+      using archive_type = std::remove_cvref_t<decltype(archive)>;
+      std::array<typename Complex::value_type, 2> data;
+      if constexpr (archive_type::kind() == zpp::bits::kind::out) data = {complex.real(), complex.imag()};
+      zpp::bits::errc result;
+      if (result = archive(data); result.code != std::errc{}) [[unlikely]] return result; 
+      if constexpr (archive_type::kind() == zpp::bits::kind::in) complex = {data[0], data[1]};
+      return result;
+    }
+  }
+}

packages/biu/test/serialize.cpp

@@ -18,9 +18,11 @@ int main()
   {
     int x;
     std::string y;
+    std::complex<double> z;
   };
-  A a{ 123, "abc" };
+  A a{ 123, "abc", 3i };
   auto b = biu::deserialize<A>(biu::serialize(a));
   assert(a.x == b.x);
   assert(a.y == b.y);
+  assert(a.z == b.z);
 }

packages/ufo/include/ufo.hpp

@@ -16,10 +16,13 @@ namespace ufo
   void plot_point(std::string config_file);
 
   // unfold 和 plot 都需要用到这个，所以写出来
-  // TODO: 把输入的数据也保留进来
   struct UnfoldOutput
   {
     Eigen::Matrix3d PrimativeCell;
+    Eigen::Matrix3i SuperCellTransformation;
+    Eigen::Vector3i SuperCellMultiplier;
+    Eigen::Matrix3d SuperCellDeformation;
+    std::optional<std::vector<std::size_t>> SelectedAtoms;
 
     // 关于各个 Q 点的数据
     struct QpointDataType
@@ -42,5 +45,27 @@ namespace ufo
       std::vector<ModeDataType> ModeData;
     };
     std::vector<QpointDataType> QpointData;
+
+    struct MetaQpointDataType
+    {
+      // Q 点的坐标，单位为单胞的倒格矢
+      Eigen::Vector3d Qpoint;
+
+      // 关于这个 Q 点上各个模式的数据
+      struct ModeDataType
+      {
+        // 模式的频率，单位为 THz
+        double Frequency;
+        // 模式中各个原子的运动状态
+        // 这个数据应当是这样得到的：动态矩阵的 eigenvector 乘以 $\exp(-2 \pi i \vec q \cdot \vec r)$
+        // 这个数据可以认为是原子位移中, 关于超胞有周期性的那一部分, 再乘以原子质量的开方.
+        // 这个数据会在 unfold 时被归一化
+        Eigen::MatrixX3cd AtomMovement;
+      };
+      std::vector<ModeDataType> ModeData;
+    };
+    std::vector<MetaQpointDataType> MetaQpointData;
+
+    using serialize = zpp::bits::members<7>;
   };
-}
+}

packages/ufo/src/unfold.cpp

@@ -54,26 +54,6 @@ void ufo::unfold(std::string config_file)
     std::vector<QpointDataOutputFileType> QpointDataOutputFile;
   };
 
-  // 关于各个 Q 点的数据
-  struct QpointData
-  {
-    // Q 点的坐标，单位为超胞的倒格矢
-    Eigen::Vector3d Qpoint;
-
-    // 关于这个 Q 点上各个模式的数据
-    struct ModeDataType
-    {
-      // 模式的频率，单位为 THz
-      double Frequency;
-      // 模式中各个原子的运动状态
-      // 这个数据应当是这样得到的：动态矩阵的 eigenvector 乘以 $\exp(-2 \pi i \vec q \cdot \vec r)$
-      // 这个数据可以认为是原子位移中, 关于超胞有周期性的那一部分, 再乘以原子质量的开方.
-      // 这个数据在读入后会被立即归一化.
-      Eigen::MatrixX3cd AtomMovement;
-    };
-    std::vector<ModeDataType> ModeData;
-  };
-
   // 从文件中读取 QpointData
   auto read_qpoint_data = [](std::string filename)
   {
@@ -108,7 +88,7 @@ void ufo::unfold(std::string config_file)
 
     // 整理得到结果
     auto number_of_qpoints = frequency.size(), num_of_modes = frequency[0].size();
-    std::vector<QpointData> qpoint_data(number_of_qpoints);
+    std::vector<UnfoldOutput::MetaQpointDataType> qpoint_data(number_of_qpoints);
     for (unsigned i = 0; i < number_of_qpoints; i++)
     {
       qpoint_data[i].Qpoint = qpoint[i] | biu::toEigen<>;
@@ -306,7 +286,7 @@ void ufo::unfold(std::string config_file)
   (
     const std::vector<std::vector<Eigen::VectorXcd>>& basis,
     // 实际上只需要其中的 AtomMovement
-    const std::vector<QpointData>& qpoint_data,
+    const std::vector<UnfoldOutput::MetaQpointDataType>& qpoint_data,
     std::atomic<unsigned>& number_of_finished_modes
   )
   {
@@ -358,16 +338,21 @@ void ufo::unfold(std::string config_file)
   // 组装输出，即将投影系数应用到原始数据上
   auto construct_output = []
   (
-    const Eigen::Matrix3d& primative_cell,
+    const Input& input,
     const Eigen::Vector3i& super_cell_multiplier,
     const Eigen::Matrix3d& super_cell_deformation,
     const std::vector<std::vector<std::vector<double>>>& projection_coefficient,
-    const std::vector<QpointData>& qpoint_data,
+    const std::vector<UnfoldOutput::MetaQpointDataType>& qpoint_data,
     const std::optional<std::vector<std::size_t>>& selected_atoms
   )
   {
     UnfoldOutput output;
-    output.PrimativeCell = primative_cell;
+    output.PrimativeCell = input.PrimativeCell;
+    output.SuperCellTransformation = input.SuperCellTransformation;
+    output.SuperCellMultiplier = super_cell_multiplier;
+    output.SuperCellDeformation = super_cell_deformation;
+    output.SelectedAtoms = selected_atoms;
+    output.MetaQpointData = qpoint_data;
     for (unsigned i_of_meta_qpoint = 0; i_of_meta_qpoint < qpoint_data.size(); i_of_meta_qpoint++)
     {
       // 如果需要投影到特定的原子上，需要先计算当前 meta qpoint 的不同模式的投影系数
@@ -487,7 +472,7 @@ void ufo::unfold(std::string config_file)
   {
     auto output = construct_output
     (
-      input.PrimativeCell, super_cell_multiplier, super_cell_deformation,
+      input, super_cell_multiplier, super_cell_deformation,
       projection_coefficient, qpoint_data, output_file.SelectedAtoms
     );
     if (output_file.OutputAsYaml.value_or(false)) std::ofstream(output_file.Filename) << YAML::Node(output);