Add fitting functionality using Ceres-Solver

CMakeLists.txt

@@ -26,6 +26,7 @@ find_package(PythonInterp 3)
 find_package(PythonLibs 3)
 find_package(FLAC)
 find_package(GSL)
+find_package(Ceres)
 
 find_package(OpenMP)
 if(OPENMP_FOUND)
@@ -67,6 +68,7 @@ add_library(so3g SHARED
   src/so_linterp.cxx
   src/exceptions.cxx
   src/array_ops.cxx
+  src/fitting_ops.cxx
 )
 
 # We could disable the lib prefix on the output library... but let's not.
@@ -87,6 +89,8 @@ target_link_libraries(so3g spt3g::core)
 # Link GSL
 target_include_directories(so3g PRIVATE ${GSL_INCLUDE_DIR})
 target_link_libraries(so3g ${GSL_LIBRARIES})
+# Link Ceres
+target_link_libraries(so3g Ceres::ceres Eigen3::Eigen)
 
 # You probably want to select openblas, so pass -DBLA_VENDOR=OpenBLAS
 find_package(BLAS REQUIRED)

Dockerfile

@@ -14,11 +14,31 @@ RUN apt update && apt install -y \
     gfortran \
     libopenblas-dev \
     libbz2-dev \
-    python-is-python3
+    python-is-python3 \
+    libgoogle-glog-dev \
+    libgflags-dev \
+    libmetis-dev \
+    libgtest-dev \
+    libabsl-dev \
+    libeigen3-dev
 
 # Set the working directory
 WORKDIR /app_lib/so3g
 
+# Fetch and install ceres-solver
+RUN git clone --depth 1 --branch 2.2.0 --recurse-submodules https://github.com/ceres-solver/ceres-solver
+
+WORKDIR /app_lib/so3g/ceres-solver
+
+RUN mkdir build \
+    && cd build \
+    && cmake .. -DBUILD_TESTING=OFF \
+    && make -j$(nproc) \
+    && make install
+
+# Set the working directory back to so3g
+WORKDIR /app_lib/so3g
+
 # Copy the current directory contents into the container
 ADD . /app_lib/so3g
 

cmake/FindCeres.cmake

@@ -0,0 +1,29 @@
+# - Find the Ceres Solver library
+#
+# CERES_FOUND
+# CERES_INCLUDE_DIRS
+# CERES_LIBRARIES
+# CERES_LIBRARY_DIRS
+
+# Look for the Ceres package
+find_path(CERES_INCLUDE_DIR NAMES ceres/ceres.h HINTS ENV CERES_DIR PATH_SUFFIXES include)
+find_library(CERES_LIBRARY NAMES ceres HINTS ENV CERES_DIR PATH_SUFFIXES lib)
+
+# Get Dependencies
+find_package(Eigen3 REQUIRED)
+find_package(Glog REQUIRED)
+find_package(Gflags REQUIRED)
+
+# Create the imported Ceres target
+add_library(Ceres::Ceres UNKNOWN IMPORTED)
+set_target_properties(Ceres::Ceres PROPERTIES
+    IMPORTED_LOCATION "${CERES_LIBRARY}"
+    INTERFACE_INCLUDE_DIRECTORIES "${CERES_INCLUDE_DIR};${EIGEN3_INCLUDE_DIR}"
+    INTERFACE_LINK_LIBRARIES "${GLOG_LIBRARIES};${GFLAGS_LIBRARIES};${CERES_LIBRARY}"
+)
+
+include (FindPackageHandleStandardArgs)
+find_package_handle_standard_args (Ceres DEFAULT_MSG CERES_LIBRARY CERES_INCLUDE_DIR)
+
+# Set the results so they can be used by the project
+mark_as_advanced(CERES_INCLUDE_DIR CERES_LIBRARY)

cmake/FindEigen3.cmake

@@ -0,0 +1,25 @@
+# - Find the Eigen library
+# 
+# EIGEN_FOUND
+# EIGEN_INCLUDE_DIRS
+# EIGEN_LIBRARIES
+# EIGEN_LIBRARY_DIRS
+
+if (EIGEN_INCLUDE_DIR)
+# Already in cache, be silent
+set (EIGEN_FIND_QUIETLY TRUE)
+endif (EIGEN_INCLUDE_DIR)
+
+find_path(EIGEN_INCLUDE_DIR "Eigen/Core"
+    HINTS ENV EIGEN_DIR
+    PATH_SUFFIXES eigen3
+)
+
+add_library(Eigen3::Eigen INTERFACE IMPORTED)
+set_target_properties(Eigen3::Eigen PROPERTIES 
+    INTERFACE_INCLUDE_DIRECTORIES "${EIGEN_INCLUDE_DIR}")
+
+include (FindPackageHandleStandardArgs)
+find_package_handle_standard_args (Eigen3 DEFAULT_MSG EIGEN_INCLUDE_DIR)
+
+mark_as_advanced(EIGEN_INCLUDE_DIR)

cmake/FindGFlags.cmake

@@ -0,0 +1,23 @@
+# - Find GFLAGS
+#
+# GFLAGS_FOUND
+# GFLAGS_INCLUDE_DIRS
+# GFLAGS_LIBRARIES
+# GFLAGS_LIBRARY_DIRS
+
+if (GFLAGS_INCLUDE_DIR)
+# Already in cache, be silent
+set (GFLAG_FIND_QUIETLY TRUE)
+endif (GFLAGS_INCLUDE_DIR)
+
+find_path(GFLAGS_INCLUDE_DIR NAME gflags/gflags.h HINTS ENV GFLAGS_DIR PATH_SUFFIXES include)
+find_library(GFLAGS_LIBRARY NAME gflags HINTS ENV GFLAGS_DIR PATH_SUFFIXES lib)
+
+set(GFLAGS_INCLUDE_DIRS ${GFLAGS_INCLUDE_DIR})
+set(GFLAGS_LIBRARIES ${GFLAGS_LIBRARY})
+
+include (FindPackageHandleStandardArgs)
+find_package_handle_standard_args (Gflags DEFAULT_MSG GFLAGS_LIBRARIES GFLAGS_INCLUDE_DIRS)
+
+mark_as_advanced(GFLAGS_LIBRARY_DEBUG GFLAGS_LIBRARY_RELEASE
+                 GFLAGS_LIBRARY GFLAGS_INCLUDE_DIR GFLAGS_ROOT_DIR)

cmake/FindGlog.cmake

@@ -0,0 +1,23 @@
+# - Find glog
+#
+# GLOG_FOUND
+# GLOG_INCLUDE_DIRS
+# GLOG_LIBRARIES
+# GLOG_LIBRARY_DIRS
+
+if (GLOG_INCLUDE_DIR)
+# Already in cache, be silent
+set (GLOG_FIND_QUIETLY TRUE)
+endif (GLOG_INCLUDE_DIR)
+
+find_path(GLOG_INCLUDE_DIR glog/logging.h HINTS ENV GLOG_DIR PATH_SUFFIXES include)
+find_library(GLOG_LIBRARY NAME glog HINTS ENV GLOG_DIR PATH_SUFFIXES lib)
+
+set(GLOG_INCLUDE_DIRS ${GLOG_INCLUDE_DIR})
+set(GLOG_LIBRARIES ${GLOG_LIBRARY})
+
+include (FindPackageHandleStandardArgs)
+find_package_handle_standard_args (Glog DEFAULT_MSG GLOG_LIBRARY GLOG_INCLUDE_DIR)
+
+mark_as_advanced(GLOG_LIBRARY_DEBUG GLOG_LIBRARY_RELEASE
+                 GLOG_LIBRARY GLOG_INCLUDE_DIR GLOG_ROOT_DIR)

include/array_ops.h

@@ -0,0 +1,6 @@
+#pragma once
+
+int get_dtype(const bp::object &);
+
+template <typename T>
+T _calculate_median(const T*, const int);

include/fitting_ops.h

@@ -0,0 +1,139 @@
+#pragma once
+
+#include <ceres/ceres.h>
+
+template <int Degree>
+struct PolynomialModel
+{
+    // Ceres requires number of params
+    // to be known at compile time
+    static constexpr int nparams = Degree + 1;
+
+    template <typename T>
+    static T eval(T x, const T* params)
+    {
+        const T p0 = params[0];
+        T result = p0;
+        for (int i = 1; i < nparams; ++i) {
+            const T p = params[i];
+            result += p * ceres::pow(x, T(i));
+        }
+
+        return result;
+    }
+    // Not needed for least squares as ceres 
+    // supports boundaries
+    template <typename T>
+    static bool check_bounds(const T* params)
+    {
+        return true;
+    }
+};
+
+struct NoiseModel
+{
+    // Ceres requires number of params
+    // to be known at compile time
+    static constexpr int nparams = 3;
+
+    template <typename T>
+    static T eval(T f, const T* params)
+    {
+        const T fknee = params[0];
+        const T w = params[1];
+        const T alpha = params[2];
+
+        return w * (1.0 + ceres::pow(fknee / f, alpha));
+    }
+
+    // Slightly hacky way of bounds checking but is
+    // suggested by Ceres to ensure it never goes
+    // out of bounds
+    template <typename T>
+    static bool check_bounds(const T* params)
+    {
+        const T w = params[1];
+        if (w <= 0.0) {
+            return false;
+        }
+        return true;
+    }
+};
+
+// Model independent cost function for least-squares fitting
+template <typename Model>
+struct CostFunction
+{
+    using model = Model;
+
+    CostFunction(int n, const double* x_data, const double* y_data)
+        : n_pts(n), x(x_data), y(y_data) {}
+
+    template <typename T>
+    bool operator()(const T* const params, T* residual) const {
+        for (int i = 0; i < n_pts; ++i) {
+            T model = Model::eval(T(x[i]), params);
+            residual[i] = T(y[i]) - model;
+        }
+        return true;
+    }
+
+    static ceres::Problem create(const int n, const double* xx,
+                                 const double* yy, double* p)
+    {
+        ceres::Problem problem;
+
+        problem.AddResidualBlock(
+            new ceres::AutoDiffCostFunction<CostFunction<Model>, 
+            ceres::DYNAMIC, Model::nparams>(
+                new CostFunction<Model>(n, xx, yy), n), nullptr, p);
+
+        return problem;
+    }
+
+private:
+    const int n_pts;
+    const double* x;
+    const double* y;
+};
+
+// Model independent Negative Log Likelihood for generalized 
+// unconstrained minimization
+template <typename Model>
+struct NegLogLikelihood
+{
+    using model = Model;
+
+    NegLogLikelihood(int n, const double* x_data, const double* y_data)
+        : n_pts(n), x(x_data), y(y_data) {}
+
+    template <typename T>
+    bool operator()(const T* const params, T* cost) const
+    {
+        // Check bounds (saves a lot of time)
+        if (!model::check_bounds(params)) {
+            return false;
+        }
+
+        cost[0] = T(0.);
+        for (int i = 0; i < n_pts; ++i) {
+            T model = Model::eval(T(x[i]), params);
+            cost[0] += ceres::log(model) + T(y[i]) / model;
+        }
+
+        return true;
+    }
+
+    static ceres::FirstOrderFunction* create(int n, const double* xx,
+                                             const double* yy)
+    {
+        // Ceres takes ownership of pointers so no cleanup is required
+        return new ceres::AutoDiffFirstOrderFunction<NegLogLikelihood<Model>,
+            Model::nparams>(new NegLogLikelihood<Model>(n, xx, yy));
+    }
+
+private:
+    const int n_pts;
+    const double* x;
+    const double* y;
+};

src/array_ops.cxx

@@ -25,6 +25,7 @@ extern "C" {
 #include "so3g_numpy.h"
 #include "numpy_assist.h"
 #include "Ranges.h"
+#include "array_ops.h"
 
 // TODO: Generalize to double precision too.
 // This implements Jon's noise model for ACT. It takes in
@@ -1116,6 +1117,9 @@ T _calculate_median(const T* data, const int n)
     return gsl_stats_median(data_copy.data(), 1, n);
 }
 
+template double _calculate_median<double>(const double* arr, int size);
+template float _calculate_median<float>(const float* arr, int size);
+
 template <typename T>
 void _detrend(T* data, const int ndets, const int nsamps, const int row_stride,
               const std::string & method, const int linear_ncount,
@@ -1254,7 +1258,6 @@ void detrend(bp::object & tod, const std::string & method, const int linear_ncou
     }
 }
 
-
 PYBINDINGS("so3g")
 {
     bp::def("nmat_detvecs_apply", nmat_detvecs_apply);