autoopt-cpp/pysrc/module.cpp

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