1 files changed, 139 insertions, 0 deletions
diff --git a/fits.hpp b/fits.hpp
new file mode 100644
index 0000000..c6e56fa
--- /dev/null
+++ b/fits.hpp
@@ -0,0 +1,139 @@
+#include "pcg-cpp/include/pcg_random.hpp"
+#include "randutils/randutils.hpp"
+
+#include "eigen/Eigen/Dense"
+
+using Rng = randutils::random_generator<pcg32>;
+
+using Real = double;
+using Vector = Eigen::Matrix<Real, Eigen::Dynamic, 1>;
+using Matrix = Eigen::Matrix<Real, Eigen::Dynamic, Eigen::Dynamic>;
+using Data = std::vector<std::tuple<Real, Real>>;
+
+#ifndef FITS_ABS_BASIS
+inline Real basisFunctions(unsigned N, unsigned i, Real x) {
+  return std::legendre(i, 2.0 * (x - 0.5));
+}
+#endif
+
+#ifdef FITS_ABS_BASIS
+inline Real basisFunctions(unsigned N, unsigned i, Real x) {
+  return abs(x - i / (N - 1.0));
+}
+#endif
+
+Real value(const Vector& coeffs, Real x) {
+  Real v = 0;
+
+  for (unsigned j = 0; j < coeffs.size(); j++) {
+    v += coeffs[j] * basisFunctions(coeffs.size(), j, x);
+  }
+
+  return v;
+}
+
+Real cost(Data::const_iterator batchBegin, Data::const_iterator batchEnd, const Vector& coeffs) {
+  Real c = 0;
+
+#pragma omp parallel for reduction(+:c)
+  for (Data::const_iterator it = batchBegin; it != batchEnd; it++) {
+    Real x = std::get<0>(*it);
+    Real y = std::get<1>(*it);
+    c += 0.5 * pow(value(coeffs, x) - y, 2);
+  }
+
+  return c;
+}
+
+Vector dCost(Data::const_iterator batchBegin, Data::const_iterator batchEnd, const Vector& coeffs) {
+  Vector dC = Vector::Zero(coeffs.size());
+
+  for (Data::const_iterator it = batchBegin; it != batchEnd; it++) {
+    Real x = std::get<0>(*it);
+    Real y = std::get<1>(*it);
+    Real Δc = value(coeffs, x) - y;
+
+#pragma omp parallel for
+    for (unsigned j = 0; j < coeffs.size(); j++) {
+      dC[j] += Δc * basisFunctions(coeffs.size(), j, x);
+    }
+  }
+
+  return dC;
+}
+
+Vector underSolve(const Data& data, unsigned N) {
+  unsigned M = data.size();
+  Matrix A(M, N);
+  Vector b(M);
+  for (unsigned i = 0; i < M; i++) {
+    Real x = std::get<0>(data[i]);
+    Real y = std::get<1>(data[i]);
+    b[i] = y;
+    for (unsigned j = 0; j < N; j++) {
+      A(i, j) = basisFunctions(N, j, x);
+    }
+  }
+
+  return A.colPivHouseholderQr().solve(b);
+}
+
+Vector stochasticGradientDescent(const Data& data, const Vector& a₀, unsigned nBatches, long unsigned maxSteps, Real ε = 1e-12) {
+  Vector xₜ = a₀;
+  Real Hₜ;
+  Real α = 1.0;
+  Real m;
+  Vector ∇H;
+  long unsigned steps = 0;
+
+  unsigned batchSize = data.size() / nBatches;
+  Data::const_iterator batchBegin = data.begin();
+  Data::const_iterator batchEnd = data.begin() + batchSize;
+  Data::const_iterator effectiveEnd = data.begin() + batchSize * nBatches;
+
+  while (
+    Hₜ = cost(batchBegin, batchEnd, xₜ),
+    ∇H = dCost(batchBegin, batchEnd, xₜ),
+    m = ∇H.squaredNorm(),
+    m > ε && steps < maxSteps
+  ) {
+    Vector xₜ₊₁;
+    Real Hₜ₊₁;
+
+    while (
+      xₜ₊₁ = xₜ - α * ∇H, Hₜ₊₁ = cost(batchBegin, batchEnd, xₜ₊₁),
+      Hₜ₊₁ > Hₜ - 0.25 * α * m
+    ) {
+      α /= 2;
+    }
+
+    xₜ = xₜ₊₁;
+    α *= 1.25;
+    steps++;
+
+    if (batchEnd == data.end()) {
+      batchBegin = data.begin();
+      batchEnd = data.begin() + batchSize;
+    } else if (batchEnd == effectiveEnd) {
+      batchBegin = effectiveEnd;
+      batchEnd = data.end();
+    } else {
+      batchBegin += batchSize;
+      batchEnd += batchSize;
+    }
+  }
+
+  return xₜ;
+}
+
+Data generateData(Real(*f)(Real), unsigned M, Real σ, Rng& r) {
+  Data data;
+  data.reserve(M);
+
+  for (unsigned i = 0; i < M; i++) {
+    Real x = ((Real)i) / (M - 1.0);
+    data.push_back({x, f(x) + r.variate<Real, std::normal_distribution>(0, σ)});
+  }
+
+  return data;
+}