add py module

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