add py module

flake.nix

@@ -24,6 +24,8 @@
             lib,
             eigen,
             gtest,
+            python313,
+            python313Packages,
             withTests ? false,
           }:
           let
@@ -39,6 +41,8 @@
                 ./include
                 ./CMakeLists.txt
                 ./tests
+                ./src
+                ./pysrc
               ];
             };
 
@@ -48,6 +52,13 @@
               cmake
               ninja
               eigen
+              (python313.withPackages (
+                ps: with ps; [
+                  nanobind
+                  matplotlib
+                  numpy
+                ]
+              ))
             ];
           };
       in

include/autoopt/btls.hpp

@@ -9,10 +9,10 @@ namespace autoopt {
 template <typename T>
 struct btls_parameters {
   T step_decrease = T{0.5};
-  T step_increase = T{1.5};
+  T step_increase = T{1.2};
   T sufficient_decrease = T{1e-2};
-  T tolerance = T{1e-9};
+  T tolerance = T{1e-10};
-  size_t max_iters = 1000;
+  size_t max_iters = 2000;
 };
 
 template <typename T>
@@ -24,20 +24,27 @@ void btls(optimization_problem<T>& problem,
   for (size_t iter = 0; iter < params.max_iters; ++iter) {
     T obj_value = problem.objective(x);
 
-    std::cout << "Iter " << iter << ": obj = " << obj_value
+    Eigen::VectorX<T> grad = problem.gradient(x);
-              << ", x = " << x.transpose() << ", step_size = " << step_size
-              << std::endl;
 
-    Eigen::VectorX<T> grad = -problem.gradient(x);
+    Eigen::MatrixX<T> hess = problem.hessian(x);
+    Eigen::VectorX<T> step_dir = -hess.ldlt().solve(grad).normalized();
 
-    Eigen::VectorX<T> step_dir = grad.normalized();
+    // Eigen::VectorX<T> step_dir = -grad.normalized();
-    while (problem.objective(x + step_size * step_dir) >
+
-           obj_value +
+    auto decrease_condition = [&] {
-               params.sufficient_decrease * step_size * grad.dot(step_dir)) {
+      T new_obj = problem.objective(x + step_size * step_dir);
+      return std::isnan(new_obj) ||
+             (new_obj > obj_value - std::abs(params.sufficient_decrease *
+                                             step_size * grad.dot(step_dir)));
+    };
+
+    while (decrease_condition()) {
       step_size *= params.step_decrease;
     }
     x += step_size * step_dir;
 
+    step_size = step_size * params.step_increase;
+
     if (step_size < params.tolerance) {
       break;
     }

include/autoopt/interface.hpp

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <eigen3/Eigen/Eigen>
+
+namespace autoopt {
+
+Eigen::VectorX<double> fit_ellipse(
+    const std::vector<std::pair<double, double>>& data,
+    const Eigen::VectorX<double>& inital_params,
+    const Eigen::VectorX<double>& delta);
+}  // namespace autoopt

include/autoopt/quadric.hpp

@@ -1,7 +1,7 @@
 #pragma once
 
-#include <cmath>
 #include <autoopt/derivative.hpp>
+#include <cmath>
 
 namespace autoopt {
 
@@ -34,6 +34,14 @@ struct quadric {
     return quadric(A_new, B_new, C_new, D_new, E_new, F_new);
   }
 
+  constexpr quadric translated_by(T x_o, T y_o) const {
+    T D_new = _D - T{2} * _A * x_o - _B * y_o;
+    T E_new = _E - _B * x_o - T{2} * _C * y_o;
+    T F_new = _F + _A * x_o * x_o + _B * x_o * y_o + _C * y_o * y_o - _D * x_o -
+              _E * y_o;
+    return quadric(_A, _B, _C, D_new, E_new, F_new);
+  }
+
   constexpr T at(T x) const {
     T sign = T{-1};
     T Bx_E = _B * x + _E;
@@ -50,10 +58,12 @@ struct quadric {
   }
 
   constexpr T slope_at(T x) const {
-    return derivative([&]<typename U>(U x_val) {
+    return derivative(
+        [&]<typename U>(U x_val) {
           quadric<U> q{U{_A}, U{_B}, U{_C}, U{_D}, U{_E}, U{_F}};
           return q.at(x_val);
-    }, x);
+        },
+        x);
   }
 };
 

pysrc/module.cpp

@@ -0,0 +1,200 @@
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/eigen/dense.h>
+
+#include <autoopt/btls.hpp>
+#include <autoopt/ellipse.hpp>
+#include <autoopt/interface.hpp>
+#include <autoopt/optimization_problem.hpp>
+#include <autoopt/util.hpp>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <sstream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace nb = nanobind;
+
+struct input_data {
+  size_t n_points;
+  double* data;
+
+  input_data(size_t n) : n_points(n) {
+    if (n_points == 0) {
+      data = nullptr;
+    } else {
+      data = new double[n_points * 2];
+    }
+  }
+
+  input_data(const input_data&) = delete;
+  input_data& operator=(const input_data&) = delete;
+
+  input_data(input_data&& other) noexcept
+      : n_points(other.n_points), data(other.data) {
+    other.data = nullptr;
+    other.n_points = 0;
+  }
+
+  input_data& operator=(input_data&& other) noexcept {
+    if (this != &other) {
+      delete[] data;
+      n_points = other.n_points;
+      data = other.data;
+      other.data = nullptr;
+      other.n_points = 0;
+    }
+    return *this;
+  }
+
+  ~input_data() { delete[] data; }
+
+  void set_point(size_t i, double x, double y) {
+    data[i * 2] = x;
+    data[i * 2 + 1] = y;
+  }
+};
+
+input_data read_data(const std::filesystem::path& filename) {
+  // check if file exists
+  if (!std::filesystem::exists(filename)) {
+    throw std::runtime_error("File does not exist: " + filename.string());
+  }
+
+  std::fstream file(filename, std::ios::in);
+
+  std::unordered_map<double, size_t> index_map;
+  std::vector<double> x_values;
+  std::vector<std::vector<double>> data_points;
+
+  std::string line;
+
+  size_t index = 0;
+
+  double x_avg = 0.0;
+
+  while (std::getline(file, line)) {
+    if (line.empty() || line[0] == '#') {
+      std::cout << "Skipped: " << line << std::endl;
+      continue;
+    }
+    std::istringstream iss(line);
+    double x, y;
+    iss >> x;
+    for (size_t i = 1; i < 5; ++i) {
+      iss >> y;  // skip unused columns
+    }
+
+    y = autoopt::arcsec2rad(-y);
+
+    if (index_map.find(x) == index_map.end()) {
+      index_map[x] = index++;
+      x_values.push_back(x);
+      data_points.emplace_back();
+      data_points.back().push_back(y);
+      x_avg += x;
+      continue;
+    }
+    data_points[index_map[x]].push_back(y);
+  }
+  x_avg /= x_values.size();
+
+  input_data result(index);
+  for (size_t i = 0; i < x_values.size(); ++i) {
+    double x = x_values[i];
+    double y_avg = 0.0;
+    for (double y : data_points[i]) {
+      y_avg += y;
+    }
+    y_avg /= data_points[i].size();
+    result.set_point(i, x - x_avg, y_avg);
+  }
+  return result;
+}
+
+struct ellipse_params {
+  double left_arm;
+  double right_arm;
+  double theta;
+};
+
+NB_MODULE(slopefit, m) {
+  m.def("import_dat_file", [](std::string filename) {
+    input_data data = read_data(filename);
+    std::cout << "First data point: (" << data.data[0] << ", " << data.data[1]
+              << ")" << std::endl;
+    nb::ndarray<double, nb::numpy> array(data.data, {data.n_points, 2});
+    return nb::cast(array);
+  });
+
+  nb::class_<ellipse_params>(m, "EllipseParams")
+      .def(nb::init<double, double, double>())
+      .def_rw("left_arm", &ellipse_params::left_arm)
+      .def_rw("right_arm", &ellipse_params::right_arm)
+      .def_rw("theta", &ellipse_params::theta)
+      .def("__repr__", [](const ellipse_params& params) {
+        std::ostringstream oss;
+        oss << "EllipseParams(left_arm=" << params.left_arm
+            << ", right_arm=" << params.right_arm << ", theta=" << params.theta
+            << ")";
+        return oss.str();
+      }).def("__call__", [](const ellipse_params& params, nb::ndarray<double, nb::ndim<1>> xs) {
+        auto e = autoopt::ellipse(params.left_arm, params.right_arm,
+                                  autoopt::deg2rad(params.theta));
+        autoopt::quadric<double> q = e.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",
+      [](nb::ndarray<double, nb::ndim<2>> data, ellipse_params initial_params,
+         ellipse_params delta) -> ellipse_params {
+        std::cout << "Fitting ellipse 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: "
+                  << "left_arm=" << initial_params.left_arm
+                  << ", right_arm=" << initial_params.right_arm
+                  << ", theta=" << initial_params.theta << std::endl;
+
+        double midpoint_y = data_vec[data_vec.size() / 2].second;
+
+        Eigen::VectorX<double> init_params(4);
+        init_params(0) = initial_params.left_arm;
+        init_params(1) = initial_params.right_arm;
+        init_params(2) = autoopt::deg2rad(initial_params.theta);
+        init_params(3) = midpoint_y;
+
+        Eigen::VectorX<double> deltas(4);
+        deltas(0) = delta.left_arm;
+        deltas(1) = delta.right_arm;
+        deltas(2) = autoopt::deg2rad(delta.theta);
+        deltas(3) = autoopt::deg2rad(0.1);
+
+        std::cout << "calculating fit..." << std::endl;
+        Eigen::VectorX<double> fitted_params =
+            autoopt::fit_ellipse(data_vec, init_params, deltas);
+        ellipse_params result;
+        result.left_arm = fitted_params(0);
+        result.right_arm = fitted_params(1);
+        result.theta = autoopt::rad2deg(fitted_params(2));
+        return result;
+      });
+}

src/interface.cpp

@@ -0,0 +1,48 @@
+#include <autoopt/btls.hpp>
+#include <autoopt/ellipse.hpp>
+#include <autoopt/interface.hpp>
+#include <autoopt/optimization_problem.hpp>
+#include <iomanip>
+#include <iostream>
+
+Eigen::VectorX<double> autoopt::fit_ellipse(
+    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) {
+    ellipse<T> ellip(params(0), params(1), params(2));
+    quadric<T> quad = ellip.to_quadric().rotated_by(params(3));
+    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;
+}
+

tests/ellipse.cpp

@@ -20,7 +20,7 @@ TEST(Ellipse, Slope) {
 
 TEST(Ellipse, ParamGradient) {
   std::vector<std::pair<double, double>> data_points = {
-      {-10.0, -0.001}, {0.0, 0.0}, {10, 0.0009}};
+      {-10.0, -0.00103}, {0.0, 0.0}, {10, 0.0009}};
 
   Eigen::VectorX<double> params(4);
   params << 100, 1000, deg2rad(1.0), 0.0;
@@ -43,30 +43,37 @@ TEST(Ellipse, ParamGradient) {
 
   std::cout << grad << std::endl;
 
-  EXPECT_NEAR(grad(0), -2.0789313126683308e-10, 1e-15);  // d/d(left_arm)
+  EXPECT_NEAR(grad(0), -3.54845114759293e-12, 1e-15);  // d/d(left_arm)
-  EXPECT_NEAR(grad(1), -1.7464984353858657e-12, 1e-15);  // d/d(right_arm)
+  EXPECT_NEAR(grad(1), -3.0016523630530093e-14, 1e-15);  // d/d(right_arm)
-  EXPECT_NEAR(grad(2), 1.2013025455499119e-06, 1e-15);   // d/d(entrance_angle)
+  EXPECT_NEAR(grad(2), 2.0569619167404501e-08, 1e-15);   // d/d(entrance_angle)
-  EXPECT_NEAR(grad(3), -2.0332702665822054e-05, 1e-15);  // d/d(rotation_angle)
+  EXPECT_NEAR(grad(3), -3.3267028547673413e-07, 1e-15);  // d/d(rotation_angle)
 
   auto hess = problem.hessian(params);
 
   std::cout << hess << std::endl;
 
   Eigen::VectorX<double> params_delta(4);
-  params_delta << 1.0, 1.0, deg2rad(0.1), deg2rad(0.1);
+  params_delta << 10.0, 10.0, deg2rad(0.1), deg2rad(0.1);
 
-  params(0) += 0.9;  // left_arm
 
   // log barrier
   log_barrier_optimization_problem<double> log_barrier_problem(
       problem,
       params_delta,
-      1e-3);
+      1e-5);
   
-  auto log_barrier_grad = log_barrier_problem.gradient(params);
+  while (log_barrier_problem._barrier_strength > 1e-20) {
+    btls(log_barrier_problem);
+    log_barrier_problem._barrier_strength *= 1e-2;
+  }
 
-  std::cout << "Log barrier gradient:" << std::endl;
+  std::cout << "Optimum params:" << std::endl;
-  std::cout << log_barrier_grad << std::endl;
+  std::cout << "left_arm:       " << log_barrier_problem.x()(0) << std::endl;
-
+  std::cout << "right_arm:      " << log_barrier_problem.x()(1) << std::endl;
-  btls(problem);
+  std::cout << "entrance_angle: " << rad2deg(log_barrier_problem.x()(2)) << std::endl;
+  std::cout << "rotation_angle: " << rad2deg(log_barrier_problem.x()(3)) << std::endl;
+  std::cout << "Optimum objective:" << std::endl;
+  std::cout << problem.objective(problem.x()) << std::endl;
+  std::cout << "Optimum grad:" << std::endl;
+  std::cout << problem.gradient(problem.x()) << std::endl;
 }

tests/test_fit.cpp

@@ -0,0 +1,82 @@
+#include <gtest/gtest.h>
+
+#include <autoopt/ellipse.hpp>
+#include <autoopt/optimization_problem.hpp>
+#include <autoopt/util.hpp>
+#include <autoopt/btls.hpp>
+#include <fstream>
+#include <unordered_map>
+#include <autoopt/interface.hpp>
+#include <iomanip>
+
+
+using namespace autoopt;
+
+std::vector<std::pair<double, double>> read_data(const std::string& filename) {
+  std::fstream file(filename, std::ios::in);
+
+  std::unordered_map<double, size_t> index_map;
+  std::vector<double> x_values;
+  std::vector<std::vector<double>> data_points;
+
+  std::string line;
+
+  size_t index = 0;
+
+  double x_avg = 0.0;
+
+  for(;std::getline(file, line);) {
+    if (line.empty() || line[0] == '#') {
+      continue;
+    }
+    std::istringstream iss(line);
+    double x, y;
+    iss >> x;
+    for (size_t i = 1; i < 5; ++i) {
+      iss >> y;  // skip unused columns
+    }
+
+    y = arcsec2rad(-y);
+
+    if (index_map.find(x) == index_map.end()) {
+      index_map[x] = index++;
+      x_values.push_back(x);
+      data_points.emplace_back();
+      data_points.back().push_back(y);
+      x_avg += x;
+      continue;
+    }
+    data_points[index_map[x]].push_back(y);
+  }
+  x_avg /= x_values.size();
+  std::vector<std::pair<double, double>> result;
+  for (size_t i = 40; i < x_values.size() - 40; ++i) {
+    double x = x_values[i];
+    double y_avg = 0.0;
+    for (double y : data_points[i]) {
+      y_avg += y;
+    }
+    y_avg /= data_points[i].size();
+    result.emplace_back(x - x_avg, y_avg);
+  }
+  return result;
+}
+
+TEST(Fit, EllipseFit) {
+  auto data_points = read_data("tests/input/ellipse.dat");
+
+  std::cout << "Read " << data_points.size() << " data points." << std::endl;
+
+  auto mid_point = data_points[data_points.size() / 2];
+
+  std::cout << "Mid point: (" << mid_point.first << ", " << mid_point.second << ")" << std::endl;
+
+  Eigen::VectorX<double> initial_params(4);
+  initial_params << 6900.0, 500.0, deg2rad(2.0), mid_point.second;
+  
+  Eigen::VectorX<double> delta(4);
+  delta << 10.0, 10.0, deg2rad(0.1), deg2rad(0.1);
+
+  auto res = fit_ellipse(data_points, initial_params, delta);
+  std::cout << res.transpose() << std::endl;
+}