diff --git a/CMakeLists.txt b/CMakeLists.txt index 144e2506..03944455 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/Dockerfile b/Dockerfile index 49abf5ad..8b2d746b 100644 --- a/Dockerfile +++ b/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 diff --git a/cmake/FindCeres.cmake b/cmake/FindCeres.cmake new file mode 100644 index 00000000..686c4463 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/cmake/FindEigen3.cmake b/cmake/FindEigen3.cmake new file mode 100644 index 00000000..303863c7 --- /dev/null +++ b/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) diff --git a/cmake/FindGFlags.cmake b/cmake/FindGFlags.cmake new file mode 100644 index 00000000..6b2197f9 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/cmake/FindGlog.cmake b/cmake/FindGlog.cmake new file mode 100644 index 00000000..6b380c61 --- /dev/null +++ b/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) diff --git a/include/array_ops.h b/include/array_ops.h new file mode 100644 index 00000000..366dc810 --- /dev/null +++ b/include/array_ops.h @@ -0,0 +1,6 @@ +#pragma once + +int get_dtype(const bp::object &); + +template +T _calculate_median(const T*, const int); \ No newline at end of file diff --git a/include/fitting_ops.h b/include/fitting_ops.h new file mode 100644 index 00000000..c8886062 --- /dev/null +++ b/include/fitting_ops.h @@ -0,0 +1,139 @@ +#pragma once + +#include + +template +struct PolynomialModel +{ + // Ceres requires number of params + // to be known at compile time + static constexpr int nparams = Degree + 1; + + template + 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 + 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 + 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 + 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 +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 + 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, + ceres::DYNAMIC, Model::nparams>( + new CostFunction(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 +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 + 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, + Model::nparams>(new NegLogLikelihood(n, xx, yy)); + } + +private: + const int n_pts; + const double* x; + const double* y; +}; diff --git a/src/array_ops.cxx b/src/array_ops.cxx index ba1a88b6..2befb440 100644 --- a/src/array_ops.cxx +++ b/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(const double* arr, int size); +template float _calculate_median(const float* arr, int size); + template 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); diff --git a/src/fitting_ops.cxx b/src/fitting_ops.cxx new file mode 100644 index 00000000..8d09731d --- /dev/null +++ b/src/fitting_ops.cxx @@ -0,0 +1,463 @@ +#define NO_IMPORT_ARRAY + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include +#include "so3g_numpy.h" +#include "numpy_assist.h" +#include "Ranges.h" +#include "array_ops.h" +#include "fitting_ops.h" + + +template +void _least_squares(const double* x, const double* y, double* p, const int n, + const int nthreads=1, const double* lower_bounds=nullptr, + const double* upper_bounds=nullptr, double* c=nullptr) +{ + // Set up ceres problem + ceres::Problem problem = CostFunction::create(n, x, y, p); + + ceres::Solver::Options options; + options.linear_solver_type = ceres::LinearSolverType::DENSE_QR; + options.logging_type = ceres::LoggingType::SILENT; + options.minimizer_progress_to_stdout = false; + options.num_threads = nthreads; + + // Set bounds + if (lower_bounds) { + for (int i = 0; i < CostFunction::model::nparams; ++i) { + if (!std::isnan(lower_bounds[i])) { + problem.SetParameterLowerBound(p, i, lower_bounds[i]); + } + } + } + if (upper_bounds) { + for (int i = 0; i < CostFunction::model::nparams; ++i) { + if (!std::isnan(upper_bounds[i])) { + problem.SetParameterUpperBound(p, i, upper_bounds[i]); + } + } + } + + ceres::Solver::Summary summary; + ceres::Solve(options, &problem, &summary); + + if (summary.IsSolutionUsable()) { + if (c) { + ceres::Covariance::Options covariance_options; + covariance_options.sparse_linear_algebra_library_type = + ceres::SparseLinearAlgebraLibraryType::EIGEN_SPARSE; + covariance_options.algorithm_type = + ceres::CovarianceAlgorithmType::DENSE_SVD; + covariance_options.null_space_rank = -1; + + ceres::Covariance covariance(covariance_options); + std::vector> covariance_blocks; + covariance_blocks.emplace_back(p, p); + if (covariance.Compute(covariance_blocks, &problem)) { + covariance.GetCovarianceBlock(p, p, c); + } + else { + for (int i = 0; i < CostFunction::model::nparams; ++i) { + p[i] = std::numeric_limits::quiet_NaN(); + c[i] = std::numeric_limits::quiet_NaN(); + } + } + } + } + else { + for (int i = 0; i < CostFunction::model::nparams; ++i) { + p[i] = std::numeric_limits::quiet_NaN(); + if (c) { + c[i] = std::numeric_limits::quiet_NaN(); + } + } + } +} + +template +bool _invert_matrix(const T* matrix, T* inverse, const int n) { + // Calculate inverse with Gaussian elimination (faster than GSL + // for small numbers of parameters) + std::unique_ptr aug_matrix = std::make_unique(n * 2 * n); + + for (int i = 0; i < n; ++i) { + for (int j = 0; j < n; ++j) { + aug_matrix[i * 2 * n + j] = matrix[i * n + j]; + } + for (int j = 0; j < n; ++j) { + aug_matrix[i * 2 * n + (n + j)] = (i == j) ? 1.0 : 0.0; + } + } + + for (int i = 0; i < n; ++i) { + T pivot = aug_matrix[i * 2 * n + i]; + if (pivot == 0.0) { + return false; + } + + for (int j = 0; j < 2 * n; ++j) { + aug_matrix[i * 2 * n + j] /= pivot; + } + + for (int k = 0; k < n; ++k) { + if (k == i) continue; + T factor = aug_matrix[k * 2 * n + i]; + for (int j = 0; j < 2 * n; ++j) { + aug_matrix[k * 2 * n + j] -= factor * aug_matrix[i * 2 * n + j]; + } + } + } + + for (int i = 0; i < n; ++i) { + for (int j = 0; j < n; ++j) { + inverse[i * n + j] = aug_matrix[i * 2 * n + (n + j)]; + } + } + + return true; +} + +// Compute the Hessian matrix using finite differences +template +void _compute_hessian(ceres::GradientProblem& problem, double* p, + double* hessian, const double epsilon=1e-5) { + + double gradient_plus[Likelihood::model::nparams]; + double gradient_minus[Likelihood::model::nparams]; + double perturbed_parameters[Likelihood::model::nparams]; + double cost; + + // Loop over each parameter + for (int i = 0; i < Likelihood::model::nparams; ++i) { + // Perturb in the positive direction + std::copy(p, p + Likelihood::model::nparams, perturbed_parameters); + perturbed_parameters[i] += epsilon; + problem.Evaluate(perturbed_parameters, &cost, gradient_plus); + + // Perturb in the negative direction + std::copy(p, p + Likelihood::model::nparams, perturbed_parameters); + perturbed_parameters[i] -= epsilon; + problem.Evaluate(perturbed_parameters, &cost, gradient_minus); + + // Second derivative approximation + for (int j = 0; j < Likelihood::model::nparams; ++j) { + hessian[i * Likelihood::model::nparams + j] = + (gradient_plus[j] - gradient_minus[j]) / (2 * epsilon); + } + } +} + +// General minimization function +template +void _minimize(const double* x, const double* y, double* p, const int n, + const double tol, const int niter, double* c=nullptr, + const double epsilon=1e-5) +{ + // Create problem + ceres::GradientProblem problem(Likelihood::create(n, x, y)); + + // Options for solving + ceres::GradientProblemSolver::Options options; + options.logging_type = ceres::LoggingType::SILENT; + options.minimizer_progress_to_stdout = false; + options.line_search_direction_type = ceres::BFGS; + options.max_num_iterations = niter; + options.function_tolerance = tol; + + ceres::GradientProblemSolver::Summary summary; + + // Run the solver + ceres::Solve(options, problem, p, &summary); + + // Calculate uncertainties. Ceres does not include a way to calculate + // uncertainties for gradient problems but this seems fast enough. + if (summary.IsSolutionUsable()) { + if (c) { + double hessian[Likelihood::model::nparams * Likelihood::model::nparams]; + double inv_hessian[Likelihood::model::nparams * Likelihood::model::nparams]; + + _compute_hessian(problem, p, hessian, epsilon); + bool not_singular = _invert_matrix(hessian, inv_hessian, + Likelihood::model::nparams); + + if (not_singular) { + for (int i = 0; i < Likelihood::model::nparams; ++i) { + c[i] = std::sqrt(inv_hessian[i * Likelihood::model::nparams + i]); + } + } + else { + for (int i = 0; i < Likelihood::model::nparams; ++i) { + c[i] = std::numeric_limits::quiet_NaN(); + } + } + } + } + else { + for (int i = 0; i < Likelihood::model::nparams; ++i) { + p[i] = std::numeric_limits::quiet_NaN(); + if (c) { + c[i] = std::numeric_limits::quiet_NaN(); + } + } + } +} + +auto _get_frequency_limits(const double* f, const double lowf, + const double fwhite_lower, + const double fwhite_upper, + const int nsamps) +{ + // Index for f = lowf + int lowf_i = 0; + for (int i = 0; i < nsamps; ++i) { + if (f[i] > lowf) { + lowf_i = i; + break; + } + } + + std::vector fwhite_i; + + // index for f > lower fwhite + for (int i = 0; i < nsamps; ++i) { + if (f[i] > fwhite_lower) { + fwhite_i.push_back(i); + break; + } + } + + // index for f < upper fwhite + for (int i = nsamps - 1; i >= 0; --i) { + if (f[i] < fwhite_upper) { + fwhite_i.push_back(i); + break; + } + } + + int fwhite_size = fwhite_i[1] - fwhite_i[0] + 1; + + return std::make_tuple(lowf_i, fwhite_i, fwhite_size); +} + +template +void _fit_noise(const double* f, const double* log_f, const double* pxx, + T* p, T* c, const int ndets, const int nsamps, + const int lowf_i, const std::vector fwhite_i, + const int fwhite_size, const double tol, const int niter, + const double epsilon) +{ + double log_pxx[nsamps]; + for (int i = 0; i < nsamps; ++i) { + log_pxx[i] = std::log10(pxx[i]); + } + + // Estimate of white noise + double wnest = _calculate_median(pxx + fwhite_i[0], fwhite_size); + + // Fit 1/f to line in logspace + double pfit[2] = {1., -1.}; + + // Some weak bounds (w > 0, alpha < 0) + double lower_bounds[2] = {0., std::numeric_limits::quiet_NaN()}; + double upper_bounds[2] = {std::numeric_limits::quiet_NaN(), 0.}; + + _least_squares(log_f, log_pxx, pfit, lowf_i + 1, 1, + lower_bounds, upper_bounds); + + // Don't run minimization if least squares fit failed + if (std::isnan(pfit[0]) || std::isnan(pfit[1])) { + // Implict cast from double to T + for (int i = 0; i < Likelihood::model::nparams; ++i) { + p[i] = std::numeric_limits::quiet_NaN(); + c[i] = std::numeric_limits::quiet_NaN(); + } + } + else { + // Find guess for fknee + int fi = 0; + double min_f = std::numeric_limits::max(); + for (int i = 0; i < nsamps; ++i) { + double model_val = CostFunc::model::eval(log_f[i], pfit); + double diff = std::abs(std::pow(10., model_val) - wnest); + if (diff <= min_f) { + min_f = diff; + fi = i; + } + } + + // Initial parameters + double p0[Likelihood::model::nparams] = {f[fi], wnest, -pfit[1]}; + double c0[Likelihood::model::nparams]; + + // Minimize + _minimize(f, pxx, p0, nsamps, tol, niter, c0, epsilon); + + // Implict cast from double to T + for (int i = 0; i < Likelihood::model::nparams; ++i) { + p[i] = p0[i]; + c[i] = c0[i]; + } + } +} + +template +void _fit_noise_buffer(const bp::object & f, const bp::object & pxx, + bp::object & p, bp::object & c, const double lowf, + const double fwhite_lower, const double fwhite_upper, + const double tol, const int niter, const double epsilon) +{ + using Likelihood = NegLogLikelihood; + using CostFunc = CostFunction>; + + // Disable google logging to prevent failed fit warnings for each detector + google::InitGoogleLogging("logger"); + FLAGS_minloglevel = google::FATAL + 1; + FLAGS_logtostderr = false; + + BufferWrapper pxx_buf ("pxx", pxx, false, std::vector{-1, -1}); + if (pxx_buf->strides[1] != pxx_buf->itemsize) + throw ValueError_exception("Argument 'pxx' must be contiguous in last axis."); + T* pxx_data = (T*)pxx_buf->buf; + const int ndets = pxx_buf->shape[0]; + const int nsamps = pxx_buf->shape[1]; + const int row_stride = pxx_buf->strides[0] / sizeof(T); + + BufferWrapper f_buf ("f", f, false, std::vector{nsamps}); + if (f_buf->strides[0] != f_buf->itemsize) + throw ValueError_exception("Argument 'f' must be a C-contiguous 1d array."); + T* f_data = (T*)f_buf->buf; + + BufferWrapper p_buf ("p", p, false, std::vector{ndets, Likelihood::model::nparams}); + if (p_buf->strides[1] != p_buf->itemsize) + throw ValueError_exception("Argument 'p' must be contiguous in last axis."); + T* p_data = (T*)p_buf->buf; + const int p_stride = p_buf->strides[0] / sizeof(T); + + BufferWrapper c_buf ("c", c, false, std::vector{ndets, Likelihood::model::nparams}); + if (c_buf->strides[1] != c_buf->itemsize) + throw ValueError_exception("Argument 'c' must be contiguous in last axis."); + T* c_data = (T*)c_buf->buf; + const int c_stride = c_buf->strides[0] / sizeof(T); + + if constexpr (std::is_same::value) { + // Copy f to double + double f_double[nsamps]; + + std::transform(f_data, f_data + nsamps, f_double, + [](float value) { return static_cast(value); }); + + // Get frequency bounds + auto [lowf_i, fwhite_i, fwhite_size] = + _get_frequency_limits(f_double, lowf, fwhite_lower, fwhite_upper, nsamps); + + // Fit in logspace + double log_f[nsamps]; + for (int i = 0; i < nsamps; ++i) { + log_f[i] = std::log10(f_double[i]); + } + + #pragma omp parallel for + for (int i = 0; i < ndets; ++i) { + int ioff = i * row_stride; + int poff = i * p_stride; + int coff = i * c_stride; + + // Copy pxx row to double + double pxx_det[nsamps]; + + std::transform(pxx_data + ioff, pxx_data + ioff + nsamps, pxx_det, + [](float value) { return static_cast(value); }); + + // Cast implicitly on assignment + T* p_det = p_data + poff; + T* c_det = c_data + coff; + + _fit_noise(f_double, log_f, pxx_det, p_det, + c_det, ndets, nsamps, lowf_i, fwhite_i, + fwhite_size, tol, niter, epsilon); + } + } + else if constexpr (std::is_same::value) { + // Get frequency bounds + auto [lowf_i, fwhite_i, fwhite_size] = + _get_frequency_limits(f_data, lowf, fwhite_lower, fwhite_upper, nsamps); + + // Fit in logspace + double log_f[nsamps]; + for (int i = 0; i < nsamps; ++i) { + log_f[i] = std::log10(f_data[i]); + } + + #pragma omp parallel for + for (int i = 0; i < ndets; ++i) { + int ioff = i * row_stride; + int poff = i * p_stride; + int coff = i * c_stride; + + double* pxx_det = pxx_data + ioff; + T* p_det = p_data + poff; + T* c_det = c_data + coff; + + _fit_noise(f_data, log_f, pxx_det, p_det, c_det, + ndets, nsamps, lowf_i, fwhite_i, fwhite_size, + tol, niter, epsilon); + } + } +} + +void fit_noise(const bp::object & f, const bp::object & pxx, bp::object & p, bp::object & c, + const double lowf, const double fwhite_lower, const double fwhite_upper, + const double tol, const int niter, const double epsilon) +{ + // Get data type + int dtype = get_dtype(pxx); + + if (dtype == NPY_FLOAT) { + _fit_noise_buffer(f, pxx, p, c, lowf, fwhite_lower, fwhite_upper, tol, niter, epsilon); + } + else if (dtype == NPY_DOUBLE) { + _fit_noise_buffer(f, pxx, p, c, lowf, fwhite_lower, fwhite_upper, tol, niter, epsilon); + } + else { + throw TypeError_exception("Only float32 or float64 arrays are supported."); + } +} + + +PYBINDINGS("so3g") +{ + bp::def("fit_noise", fit_noise, + "fit_noise(f, pxx, p, c, lowf, fwhite_lower, fwhite_upper, tol, niter, epsilon)" + "\n" + "Fits noise model with white and 1/f components to the PSD of signal. Uses a MLE " + "method that minimizes a log likelihood. OMP is used to parallelize across dets (rows)." + "\n" + "Args:\n" + " f: frequency array (float32/64) with dimensions (nsamps,).\n" + " pxx: PSD array (float32/64) with dimensions (ndets, nsamps).\n" + " p: parameter array (float32/64) with dimensions (ndets, nparams). This is modified in place.\n" + " c: uncertaintiy array (float32/64) with dimensions (ndets, nparams). This is modified in place.\n" + " lowf: Frequency below which the 1/f noise index and fknee are estimated for initial " + " guess passed to least_squares fit (float64).\n" + " fwhite_lower: Lower frequency used to estimate white noise for initial guess passed to " + " least_squares fit (float64).\n" + " fwhite_upper: Upper frequency used to estimate white noise for initial guess passed to " + " least_squares fit (float64).\n" + " tol: absolute tolerance for minimization (float64).\n" + " niter: total number of iterations to run minimization for (int).\n" + " epsilon: Value to perturb gradients by when calculating uncertainties with the inverse " + " Hessian matrix (float64). Affects minimization only.\n"); +} \ No newline at end of file diff --git a/test/test_fitting_ops.py b/test/test_fitting_ops.py new file mode 100644 index 00000000..c7d5d69c --- /dev/null +++ b/test/test_fitting_ops.py @@ -0,0 +1,67 @@ +import unittest + +import so3g +import numpy as np +from scipy.stats import chi2 + + +class TestFitting(unittest.TestCase): + """Test fitting noise model.""" + + def test_00_fit_noise(self): + + def noise_model(f, p): + fknee, w, alpha = p[0], p[1], p[2] + return w * (1 + (fknee / f) ** alpha) + + ndets = 1; + nsamps = 1024 // 2 + 1 # Assume nperseg = 1024 for psd + dtype = "float32" + order = "C" + + lowf = 1. + fwhite = [10., 100.] + + p0 = np.array([10., 2., 0.7]) # fknee, w, alpha + nparams = len(p0) + + tol = 1e-8 # so3g minimization tolerance + niter = 200*nparams # so3g max iterations + epsilon = 1e-5 # so3g gradient perturbation epsilon + + f = np.linspace(0.01, 200., nsamps, dtype=dtype) + pxx = np.zeros((ndets, nsamps), dtype=dtype, order=order) + + for i in range(ndets): + pxx[i,:] = noise_model(f, p0) + + so3g_params = np.zeros((ndets, nparams), dtype=dtype, order=order) + so3g_sigmas = np.zeros((ndets, nparams), dtype=dtype, order=order) + + so3g.fit_noise(f, pxx, so3g_params, so3g_sigmas, lowf, fwhite[0], + fwhite[1], tol, niter, epsilon) + + ddof = nsamps - nparams + cval = chi2.ppf(0.95, ddof) + + for i in range(ndets): + fk, w, alpha = so3g_params[i] + sfk, sw, salpha = so3g_sigmas[i] + + # Calculate model error + fkf = (1 + fk / f) + + dmdw = fkf**alpha + dmdfk = w * alpha * fkf**(alpha - 1) * f**-1 + dmdalpha = w * (fkf**alpha) * np.log(alpha) + + err = np.sqrt((dmdw*sw)**2 + (dmdfk*sfk)**2 + (dmdalpha*salpha)**2) + + model = noise_model(f, so3g_params[i]) + chi_sq = np.sum((pxx[i] - model)**2 / (err[i])**2) + + np.testing.assert_array_less(chi_sq, cval) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file