add parabola

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;
+      });
 }