packages.pocketfft: init

flake.lock

@@ -1265,6 +1265,22 @@
         "type": "github"
       }
     },
+    "pocketfft": {
+      "flake": false,
+      "locked": {
+        "lastModified": 1722531938,
+        "narHash": "sha256-1kfMW9xGHtYiiCD4Qh4rXEuAi9EIzJ4lR+YLGjqzoaU=",
+        "owner": "mreineck",
+        "repo": "pocketfft",
+        "rev": "bb87ca50df0478415a12d9011dc374eeed4e9d93",
+        "type": "github"
+      },
+      "original": {
+        "owner": "mreineck",
+        "repo": "pocketfft",
+        "type": "github"
+      }
+    },
     "poetry2nix": {
       "inputs": {
         "flake-utils": "flake-utils_4",
@@ -1393,6 +1409,7 @@
         "nur-xddxdd": "nur-xddxdd",
         "openxlsx": "openxlsx",
         "plasma-manager": "plasma-manager",
+        "pocketfft": "pocketfft",
         "poetry2nix": "poetry2nix",
         "py4vasp": "py4vasp",
         "qchem": "qchem",

flake.nix

@@ -70,6 +70,7 @@
     hextra = { url = "github:imfing/hextra"; flake = false; };
     nu-scripts = { url = "github:nushell/nu_scripts"; flake = false; };
     py4vasp = { url = "github:vasp-dev/py4vasp"; flake = false; };
+    pocketfft = { url = "github:/mreineck/pocketfft"; flake = false; };
 
     # does not support lfs yet
     # nixos-wallpaper = { url = "git+https://git.chn.moe/chn/nixos-wallpaper.git"; flake = false; };

packages/biu/CMakeLists.txt

@@ -25,11 +25,12 @@ find_path(LIBBACKTRACE_INCLUDE_DIR backtrace.h REQUIRED)
 find_library(LIBBACKTRACE_LIBRARY NAMES backtrace REQUIRED)
 find_package(HDF5 REQUIRED)
 find_package(concurrencpp REQUIRED)
+find_path(POCKETFFT_INCLUDE_DIR pocketfft.h REQUIRED)
 
 add_library(biu src/common.cpp src/hdf5.cpp src/logger.cpp src/string.cpp)
 target_include_directories(biu PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
   $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}> ${NAMEOF_INCLUDE_DIR} ${ZPP_BITS_INCLUDE_DIR}
-  ${LIBBACKTRACE_INCLUDE_DIR})
+  ${LIBBACKTRACE_INCLUDE_DIR} ${POCKETFFT_INCLUDE_DIR})
 target_link_libraries(biu PUBLIC magic_enum::magic_enum fmt::fmt Boost::headers Boost::iostreams Boost::filesystem
   range-v3::range-v3 Eigen3::Eigen HighFive TgBot::TgBot ${LIBBACKTRACE_LIBRARY} hdf5::hdf5 concurrencpp::concurrencpp)
 target_compile_features(biu PUBLIC cxx_std_23)

packages/biu/default.nix

@@ -1,11 +1,16 @@
 {
   stdenv, cmake, lib,
-  magic-enum, fmt, boost, eigen, range-v3, nameof, zpp-bits, highfive, tgbot-cpp, libbacktrace, hdf5, concurrencpp
+  magic-enum, fmt, boost, eigen, range-v3, nameof, zpp-bits, highfive, tgbot-cpp, libbacktrace, hdf5, concurrencpp,
+  pocketfft
 }: stdenv.mkDerivation rec
 {
   name = "biu";
   src = ./.;
-  buildInputs = [ magic-enum fmt boost range-v3 nameof zpp-bits eigen highfive tgbot-cpp libbacktrace hdf5 concurrencpp ];
+  buildInputs =
+  [
+    magic-enum fmt boost range-v3 nameof zpp-bits eigen highfive tgbot-cpp libbacktrace hdf5
+    concurrencpp pocketfft
+  ];
   propagatedBuildInputs = buildInputs;
   nativeBuildInputs = [ cmake ];
   doCheck = true;

packages/biu/include/biu.hpp

@@ -9,5 +9,6 @@
 # include <biu/hdf5.tpp>
 # include <biu/logger.tpp>
 # include <biu/smartref.tpp>
+# include <biu/fft.tpp>
 
 # include <range/v3/all.hpp>

packages/biu/include/biu/fft.hpp

@@ -0,0 +1,9 @@
+# pragma once
+# include <pocketfft.h>
+
+namespace biu::fft
+{
+  template <typename T> std::vector<std::complex<T>> forward(std::vector<T> input);
+  template <typename T> std::vector<T> backward
+    (std::vector<std::complex<T>> input, std::optional<std::size_t> output_size = std::nullopt);
+}

packages/biu/include/biu/fft.tpp

@@ -0,0 +1,29 @@
+# pragma once
+# include <biu/fft.hpp>
+
+namespace biu::fft
+{
+  template <typename T> std::vector<std::complex<T>> forward(std::vector<T> input)
+  {
+    std::vector<std::complex<T>> output(input.size() / 2 + 1);
+    pocketfft::r2c<T>
+    (
+      {input.size()}, {sizeof(T)}, {sizeof(std::complex<T>)},
+      0, pocketfft::FORWARD, input.data(), output.data(), 1
+    );
+    return output;
+  }
+  template <typename T> std::vector<T> backward
+    (std::vector<std::complex<T>> input, std::optional<std::size_t> output_size)
+  {
+    if (!output_size) output_size = (input.size() - 1) * 2;
+    else [[unlikely]] assert(*output_size / 2 + 1 == input.size());
+    std::vector<T> output(*output_size);
+    pocketfft::c2r<T>
+    (
+      {output.size()}, {sizeof(std::complex<T>)}, {sizeof(T)},
+      0, pocketfft::BACKWARD, input.data(), output.data(), 1
+    );
+    return output;
+  }
+}

packages/default.nix

@@ -64,7 +64,7 @@ inputs: rec
     { src = inputs.topInputs.kylin-virtual-keyboard; };
   biu = inputs.pkgs.callPackage ./biu
   {
-    inherit nameof zpp-bits tgbot-cpp concurrencpp;
+    inherit nameof zpp-bits tgbot-cpp concurrencpp pocketfft;
     stdenv = inputs.pkgs.clang18Stdenv;
     fmt = inputs.pkgs.fmt_11.overrideAttrs (prev: { patches = prev.patches or [] ++ [ ./biu/fmt.patch ]; });
   };
@@ -86,6 +86,7 @@ inputs: rec
   sockpp = inputs.pkgs.callPackage ./sockpp.nix { src = inputs.topInputs.sockpp; };
   git-lfs-transfer = inputs.pkgs.callPackage ./git-lfs-transfer.nix { src = inputs.topInputs.git-lfs-transfer; };
   py4vasp = inputs.pkgs.callPackage ./py4vasp.nix { src = inputs.topInputs.py4vasp; };
+  pocketfft = inputs.pkgs.callPackage ./pocketfft.nix { src = inputs.topInputs.pocketfft; };
 
   fromYaml = content: builtins.fromJSON (builtins.readFile
     (inputs.pkgs.runCommand "toJSON" {}

packages/pocketfft.nix

@@ -0,0 +1,13 @@
+{ lib, stdenv, src }: stdenv.mkDerivation
+{
+  name = "pocketfft";
+  inherit src;
+  phases = [ "installPhase" ];
+  installPhase =
+  ''
+    runHook preInstall
+    mkdir -p $out/include
+    cp -r $src/pocketfft_hdronly.h $out/include/pocketfft.h
+    runHook postInstall
+  '';
+}