Baseline indexer part1

CMakeLists.txt

@@ -49,11 +49,40 @@ FetchContent_Declare(
 )
 FetchContent_MakeAvailable(spdlog)
 
+# Handle inclusion of dx2 dependency in a choice of several ways
+set(dx2_SOURCE "" CACHE FILEPATH "Location of dx2 source directory")
+set(_dx2_already_populated "${dx2_SOURCE_DIR}")
+if(dx2_SOURCE)
+    # Explicitly told where to find dx2
+    add_subdirectory(${dx2_SOURCE} dx2)
+    set(dx2_SOURCE_DIR "${dx2_SOURCE}")
+elseif(EXISTS ${CMAKE_SOURCE_DIR}/dx2)
+    # dx2 exists as a local checkout
+    add_subdirectory(dx2)
+    set(dx2_SOURCE_DIR "${CMAKE_SOURCE_DIR}/dx2")
+else()
+    # Try finding an installed dx2, or fetch
+    FetchContent_Declare(
+        dx2
+        GIT_REPOSITORY  https://github.com/dials/dx2
+        GIT_TAG 2654929bc73e0d699dfdfcc2d924c3454e434ccf
+        FIND_PACKAGE_ARGS
+    )
+    FetchContent_MakeAvailable(dx2)
+endif()
+# Only print a message when as a new action
+if (NOT dx2_SOURCE_DIR STREQUAL _dx2_already_populated)
+    message(STATUS "Found dx2: ${dx2_SOURCE_DIR}")
+endif()
+
 # Make a small library that we can link to to get the version
 project(version CXX)
 configure_file(version.cc.in version.cc @ONLY)
 add_library(version STATIC version.cc)
 
+enable_testing()
+
 add_subdirectory(h5read)
 add_subdirectory(baseline)
 add_subdirectory(spotfinder)
+add_subdirectory(baseline_indexer)

baseline_indexer/CMakeLists.txt

@@ -0,0 +1,48 @@
+project(indexer CXX)
+
+set(CMAKE_CXX_STANDARD 20)
+
+find_package (Threads)
+
+# Automatic Dependencies
+set(FETCHCONTENT_QUIET OFF)
+include(FetchContent)
+FetchContent_Declare(
+    Eigen3
+    GIT_REPOSITORY https://gitlab.com/libeigen/eigen.git
+    GIT_TAG 3.4.0
+    EXCLUDE_FROM_ALL
+    FIND_PACKAGE_ARGS
+)
+FetchContent_Declare(
+    pocketfft
+    GIT_REPOSITORY https://github.com/mreineck/pocketfft
+    GIT_TAG cpp
+)
+FetchContent_MakeAvailable(Eigen3)
+FetchContent_MakeAvailable(pocketfft)
+
+FetchContent_GetProperties(pocketfft)
+add_library(pocketfft INTERFACE)
+target_include_directories(pocketfft INTERFACE ${pocketfft_SOURCE_DIR})
+
+# GTest could have been made available under a different name
+if(NOT TARGET GTest::gtest_main)
+    FetchContent_MakeAvailable(GTest)
+endif()
+
+add_subdirectory(tests)
+add_executable(baseline_indexer
+    indexer.cc
+)
+target_link_libraries(baseline_indexer
+    PRIVATE
+    fmt
+    dx2
+    Eigen3::Eigen
+    pocketfft
+    argparse
+    nlohmann_json::nlohmann_json
+    spdlog::spdlog
+    ${CMAKE_THREAD_LIBS_INIT}
+)

baseline_indexer/fft3d.cc

@@ -0,0 +1,177 @@
+#include <assert.h>
+#include <math.h>
+#include <pocketfft_hdronly.h>
+
+#include <Eigen/Dense>
+#include <algorithm>
+#include <chrono>
+#include <map>
+#include <stack>
+#include <tuple>
+
+#include "common.hpp"
+
+using Eigen::Matrix3d;
+using Eigen::Vector3d;
+using Eigen::Vector3i;
+
+#define _USE_MATH_DEFINES
+#include <cmath>
+
+using namespace pocketfft;
+
+/**
+ * @brief map reciprocal space vectors onto a grid of size n_points^3.
+ * @param reciprocal_space_vectors Reciprocal space vectors to be mapped.
+ * @param data_in The vector (grid) which the data will be mapped to.
+ * @param selection The vector of the selection of points mapped to the grid.
+ * @param d_min A resolution limit for mapping to the grid.
+ * @param b_iso The isotropic B-factor used to weight the points as a function of resolution.
+ * @param n_points The size of each dimension of the FFT grid.
+ */
+void map_centroids_to_reciprocal_space_grid(
+  std::vector<Vector3d> const &reciprocal_space_vectors,
+  std::vector<std::complex<double>> &data_in,
+  std::vector<bool> &selection,
+  double d_min,
+  double b_iso = 0,
+  uint32_t n_points = 256) {
+    assert(data_in.size() == n_points * n_points * n_points);
+    // Determine the resolution span of the grid so we know how to map
+    // each coordinate to the grid.
+    // -dmin to +dmin spans the coordinate range from 0 to N (in each dimension),
+    // but the discrete FFT requires input coefficients 0 to N-1, and so our
+    // grid goes from 0 .. N-1 and we don't use data that maps to grid point N
+    // (which is equivalent to grid point 0)
+
+    const double rlgrid = 2 / (d_min * n_points);
+    const double one_over_rlgrid = 1 / rlgrid;
+    const int half_n_points = n_points / 2;
+    int count = 0;
+
+    for (int i = 0; i < reciprocal_space_vectors.size(); i++) {
+        const Vector3d v = reciprocal_space_vectors[i];
+        const double v_length = v.norm();
+        const double d_spacing = 1 / v_length;
+        if (d_spacing < d_min) {
+            selection[i] = false;
+            continue;
+        }
+        Vector3i coord;
+        // map to the nearest point in each dimension.
+        for (int j = 0; j < 3; j++) {
+            coord[j] = static_cast<int>(round(v[j] * one_over_rlgrid)) + half_n_points;
+        }
+        if ((coord.maxCoeff() >= n_points) || coord.minCoeff() < 0) {
+            // We can get maxcoeff == n_points if we anear +dmin, see comment above.
+            selection[i] = false;
+            continue;
+        }
+        // Use the b_iso to determine the weight for each coordinate.
+        double T;
+        if (b_iso != 0) {
+            T = std::exp(-b_iso * v_length * v_length / 4.0);
+        } else {
+            T = 1;
+        }
+        // unravel to the 1d index and write the complex<double> value.
+        size_t index =
+          coord[2] + (n_points * coord[1]) + (n_points * n_points * coord[0]);
+        if (!data_in[index].real()) {
+            count++;
+        }
+        data_in[index] = {T, 0.0};
+    }
+    logger->info("Number of centroids used: {}", count);
+}
+
+/**
+ * @brief Perform a 3D FFT of the reciprocal space coordinates (spots).
+ * @param reciprocal_space_vectors The input vector of reciprocal space coordinates.
+ * @param real_out The (real) array that the FFT result will be written to.
+ * @param d_min Cut the data at this resolution limit for the FFT
+ * @param b_iso The isotropic B-factor used to weight the points as a function of resolution.
+ * @param n_points The size of each dimension of the FFT grid.
+ * @returns A boolean array indicating which coordinates were used for the FFT.
+ */
+std::vector<bool> fft3d(std::vector<Vector3d> const &reciprocal_space_vectors,
+                        std::vector<double> &real_out,
+                        double d_min,
+                        double b_iso = 0,
+                        uint32_t n_points = 256,
+                        size_t nthreads = 1) {
+    auto start = std::chrono::system_clock::now();
+    assert(real_out.size() == n_points * n_points * n_points);
+
+    // We want to write out the real part of the FFT, but the pocketfft functions require
+    // complex vectors (we are using c2c i.e. complex to complex), so initialise these vectors.
+    // Note we should be able to use c2r rather than c2c, but I couldn't get this to work with
+    // the output ordering in c2r (JBE).
+    std::vector<std::complex<double>> complex_data_in(n_points * n_points * n_points);
+
+    // A boolean array of whether the vectors were used for the FFT.
+    std::vector<bool> used_in_indexing(reciprocal_space_vectors.size(), true);
+    auto t1 = std::chrono::system_clock::now();
+
+    // Map the vectors onto a discrete grid. The values of the grid points are weights
+    // determined by b_iso.
+    map_centroids_to_reciprocal_space_grid(reciprocal_space_vectors,
+                                           complex_data_in,
+                                           used_in_indexing,
+                                           d_min,
+                                           b_iso,
+                                           n_points);
+    auto t2 = std::chrono::system_clock::now();
+
+    // Prepare the required objects for the FFT.
+    shape_t shape_in{n_points, n_points, n_points};
+    int stride_x = sizeof(std::complex<double>);
+    int stride_y = static_cast<int>(sizeof(std::complex<double>) * n_points);
+    int stride_z = static_cast<int>(sizeof(std::complex<double>) * n_points * n_points);
+    stride_t stride_in{
+      stride_x, stride_y, stride_z};  // must have the size of each element. Must have
+                                      // size() equal to shape_in.size()
+    stride_t stride_out{
+      stride_x, stride_y, stride_z};  // must have the size of each element. Must
+                                      // have size() equal to shape_in.size()
+    shape_t axes{0, 1, 2};            // 0 to shape.size()-1 inclusive
+    bool forward{FORWARD};
+
+    double fct{1.0f};
+    // note, threads can be higher than the number of hardware threads.
+    // It is not clear what the best value is for this.
+    logger->info("Performing FFT with nthreads={}", nthreads);
+    // Do the FFT.
+    c2c(
+      shape_in,
+      stride_in,
+      stride_out,
+      axes,
+      forward,
+      complex_data_in.data(),
+      complex_data_in
+        .data(),  // this is the output array, we are going to overwrite the input as we don't need it
+      fct,
+      nthreads);
+    auto t3 = std::chrono::system_clock::now();
+
+    // Take the square of the real part as the output.
+    for (int i = 0; i < real_out.size(); ++i) {
+        real_out[i] = std::pow(complex_data_in[i].real(), 2);
+    }
+    auto t4 = std::chrono::system_clock::now();
+
+    std::chrono::duration<double> elapsed_seconds = t4 - start;
+    std::chrono::duration<double> elapsed_map = t2 - t1;
+    std::chrono::duration<double> elapsed_make_arrays = t1 - start;
+    std::chrono::duration<double> elapsed_c2c = t3 - t2;
+    std::chrono::duration<double> elapsed_square = t4 - t3;
+    logger->debug("Total time for fft3d: {:.5f}s", elapsed_seconds.count());
+    logger->debug("elapsed time for making data arrays: {:.5f}s",
+                  elapsed_make_arrays.count());
+    logger->debug("elapsed time for map_to_recip: {:.5f}s", elapsed_map.count());
+    logger->debug("elapsed time for c2c: {:.5f}s", elapsed_c2c.count());
+    logger->debug("elapsed time for squaring: {:.5f}s", elapsed_square.count());
+
+    return used_in_indexing;
+}