add parabola

CMakeLists.txt

@@ -8,12 +8,25 @@ add_compile_options(-Wall -Werror -Wpedantic)
 
 find_package(Eigen3 REQUIRED)
 
-add_library(autoopt INTERFACE)
+find_package(Python 3.13 COMPONENTS Interpreter Development.Module REQUIRED)
-target_include_directories(autoopt INTERFACE include)
+
-target_link_libraries(autoopt INTERFACE Eigen3::Eigen)
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)
+find_package(nanobind CONFIG REQUIRED)
+
+file(GLOB_RECURSE SOURCES src/*.cpp)
+
+add_library(autoopt STATIC ${SOURCES})
+target_include_directories(autoopt PUBLIC include)
+target_link_libraries(autoopt PUBLIC Eigen3::Eigen)
 
 install(DIRECTORY include/autoopt DESTINATION include)
 
+nanobind_add_module(slopefit NOMINSIZE pysrc/module.cpp)
+target_link_libraries(slopefit PRIVATE autoopt)
+install(TARGETS slopefit DESTINATION lib/python3.13/site-packages)
+
 # add tests if gtest is found
 find_library(GTestPackage gtest QUIET)
 if(GTestPackage)

examples/test.py

@@ -0,0 +1,60 @@
+import sys
+import os.path as path
+
+build_dir = path.join(path.dirname(__file__), '..', 'build')
+print('Adding build directory to sys.path:', build_dir)
+sys.path.append(build_dir)
+
+import slopefit as sf
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+fl = 1200.0 # mm
+theta_mrad = 3.0 # mrad
+theta_rad = theta_mrad * 1e-3
+theta = theta_rad / np.pi * 180.0 # deg
+
+parab = sf.ParabolaParams(fl, theta)
+
+filename = 'microMAX_vsru_mispu_absum.dat'
+dirname = path.dirname(__file__)
+
+data = sf.import_dat_file(path.join(dirname, filename))
+
+# flip both x and y
+data[:, 0] = -data[:, 0]
+data[:, 1] = -data[:, 1]
+
+n_skip = 20
+data = data[n_skip:-n_skip, :]
+
+ref = parab(data[:, 0])
+
+# convert ref from slope to rad
+# ref = np.arctan(ref)
+# convert to arcsec
+# ref = ref / np.pi * 180.0 * 3600.0
+
+
+plt.plot(data[:, 0], data[:, 1], 'x')
+plt.plot(data[:, 0], ref, '-')
+plt.xlabel('x (mm)')
+plt.ylabel('slope (arcsec)')
+plt.show()
+
+d_parab = sf.ParabolaParams(100.0, 0.00001)
+
+fitted = sf.fit_parabola(data, parab, d_parab)
+
+print('Fitted parameters:')
+print('fl =', fitted.focal_length)
+print('theta =', fitted.theta * np.pi * 1e3 / 180.0, 'mrad')
+
+fitted_ref = fitted(data[:, 0])
+
+plt.plot(data[:, 0], data[:, 1], 'x')
+plt.plot(data[:, 0], fitted_ref, '-')
+plt.xlabel('x (mm)')
+plt.ylabel('slope (arcsec)')
+plt.show()

include/autoopt/interface.hpp

@@ -8,4 +8,9 @@ Eigen::VectorX<double> fit_ellipse(
     const std::vector<std::pair<double, double>>& data,
     const Eigen::VectorX<double>& inital_params,
     const Eigen::VectorX<double>& delta);
+
+Eigen::VectorX<double> fit_parabola(
+    const std::vector<std::pair<double, double>>& data,
+    const Eigen::VectorX<double>& inital_params,
+    const Eigen::VectorX<double>& delta);
 }  // namespace autoopt

include/autoopt/parabola.hpp

@@ -0,0 +1,27 @@
+#pragma once
+
+#include "autoopt/quadric.hpp"
+
+namespace autoopt
+{
+  
+template <typename T>
+struct parabola {
+  // T focal_length;
+  T exit_arm;
+  T entrance_angle;
+
+  quadric<T> to_quadric() const {
+    // T a = T{1} / (T{4} * focal_length);
+    // T x = T{2} * focal_length / std::tan(entrance_angle);
+    // T y = a * x * x;
+    T x = exit_arm * std::sin(T{2.0} * entrance_angle);
+    T f = T{0.5} * (T{1.0} - std::cos(T{2.0} * entrance_angle)) * exit_arm;
+    T a = T{1.0} / (T{4.0} * f);
+    T y = a * x * x;
+    quadric<T> q = quadric(a, T{0}, T{0}, T{0}, T{-1.0}, T{0});
+    return q.translated_by(-x, -y).rotated_by(entrance_angle - T{M_PI_2});
+  }
+};
+
+} // namespace autoopt

pysrc/module.cpp

@@ -5,6 +5,7 @@
 
 #include <autoopt/btls.hpp>
 #include <autoopt/ellipse.hpp>
+#include <autoopt/parabola.hpp>
 #include <autoopt/interface.hpp>
 #include <autoopt/optimization_problem.hpp>
 #include <autoopt/util.hpp>
@@ -121,6 +122,11 @@ struct ellipse_params {
   double theta;
 };
 
+struct parabola_params {
+  double focal_length;
+  double theta;
+};
+
 NB_MODULE(slopefit, m) {
   m.def("import_dat_file", [](std::string filename) {
     input_data data = read_data(filename);
@@ -152,6 +158,25 @@ NB_MODULE(slopefit, m) {
         return ys;
       })
       ;
+    nb::class_<parabola_params>(m, "ParabolaParams")
+      .def(nb::init<double, double>())
+      .def_rw("focal_length", &parabola_params::focal_length)
+      .def_rw("theta", &parabola_params::theta)
+      .def("__repr__", [](const parabola_params& params) {
+        std::ostringstream oss;
+        oss << "ParabolaParams(focal_length=" << params.focal_length
+            << ", theta=" << params.theta << ")";
+        return oss.str();
+      })
+      .def("__call__", [](const parabola_params& params, nb::ndarray<double, nb::ndim<1>> xs) {
+        auto p = autoopt::parabola(params.focal_length, autoopt::deg2rad(params.theta));
+        autoopt::quadric<double> q = p.to_quadric();
+        Eigen::VectorXd ys(xs.shape(0));
+        for (size_t i = 0; i < xs.shape(0); ++i) {
+          ys(i) = q.slope_at(xs(i));
+        }
+        return ys;
+      });
 
   m.def(
       "fit_ellipse",
@@ -197,4 +222,45 @@ NB_MODULE(slopefit, m) {
         result.theta = autoopt::rad2deg(fitted_params(2));
         return result;
       });
+
+  m.def(
+      "fit_parabola",
+      [](nb::ndarray<double, nb::ndim<2>> data, parabola_params initial_params,
+         parabola_params delta) -> parabola_params {
+        std::cout << "Fitting parabola to data with " << data.shape(0)
+                  << " points." << std::endl;
+
+        if (data.shape(1) != 2) {
+          throw std::runtime_error("Data array must have shape (n_points, 2)");
+        }
+
+        std::vector<std::pair<double, double>> data_vec;
+        for (size_t i = 0; i < data.shape(0); ++i) {
+          data_vec.emplace_back(data(i, 0), data(i, 1));
+        }
+
+        std::cout << "Initial parameters: "
+                  << "focal_length=" << initial_params.focal_length
+                  << ", theta=" << initial_params.theta  << std::endl;
+
+        double midpoint_y = data_vec[data_vec.size() / 2].second;
+
+        Eigen::VectorX<double> init_params(3);
+        init_params(0) = initial_params.focal_length;
+        init_params(1) = autoopt::deg2rad(initial_params.theta);
+        init_params(2) = midpoint_y;
+
+        Eigen::VectorX<double> deltas(3);
+        deltas(0) = delta.focal_length;
+        deltas(1) = autoopt::deg2rad(delta.theta);
+        deltas(2) = autoopt::deg2rad(0.1);
+
+        std::cout << "calculating fit..." << std::endl;
+        Eigen::VectorX<double> fitted_params =
+            autoopt::fit_parabola(data_vec, init_params, deltas);
+        parabola_params result;
+        result.focal_length = fitted_params(0);
+        result.theta = autoopt::rad2deg(fitted_params(1));
+        return result;
+      });
 }

src/interface.cpp

@@ -1,5 +1,6 @@
 #include <autoopt/btls.hpp>
 #include <autoopt/ellipse.hpp>
+#include <autoopt/parabola.hpp>
 #include <autoopt/interface.hpp>
 #include <autoopt/optimization_problem.hpp>
 #include <iomanip>
@@ -46,3 +47,43 @@ Eigen::VectorX<double> autoopt::fit_ellipse(
   return fitted_params;
 }
 
+Eigen::VectorX<double> autoopt::fit_parabola(
+    const std::vector<std::pair<double, double>>& data,
+    const Eigen::VectorX<double>& initial_params,
+    const Eigen::VectorX<double>& delta) {
+  auto opt_func = [&]<typename T>(const Eigen::VectorX<T>& params) {
+    parabola<T> parab(params(0), params(1));
+    quadric<T> quad = parab.to_quadric().rotated_by(params(2));
+    T error = T(0);
+    for (const auto& [x, y] : data) {
+      T y_fit = quad.slope_at(T(x));
+      error = error + (y_fit - T(y)) * (y_fit - T(y));
+    }
+    return error / T(data.size());
+  };
+
+  auto_diff_optimization_problem problem(opt_func, initial_params);
+
+  log_barrier_optimization_problem lb_problem(problem, delta, 1e-4);
+
+  while (lb_problem._barrier_strength > 1e-20) {
+    btls(lb_problem);
+    lb_problem._barrier_strength *= 1e-2;
+  }
+
+  Eigen::VectorX<double> fitted_params = problem.x();
+
+  std::cout << "Fitted parameters: " << std::setprecision(10)
+            << fitted_params.transpose() << std::endl;
+
+  double obj_value = problem.objective(fitted_params);
+  std::cout << "Objective value: " << obj_value << std::endl;
+
+  // rms in radians
+  std::cout << "RMS error: " << std::sqrt(obj_value) << " radians" << std::endl;
+  // rms in arcsec
+  double rms_arcsec = std::sqrt(obj_value) * (3600.0 * 180.0 / M_PI);
+  std::cout << "RMS error: " << rms_arcsec << " arcsec" << std::endl;
+
+  return fitted_params;
+}