Add ode45 to runge_kutta.h
This gives us a way to integrate with an adaptive step size for when we
don't know the time constants super well.
Change-Id: Ie6073c208ae9988957f0c4cd79f9519a4a978efe
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/frc971/control_loops/BUILD b/frc971/control_loops/BUILD
index f5a720f..ccaac40 100644
--- a/frc971/control_loops/BUILD
+++ b/frc971/control_loops/BUILD
@@ -382,9 +382,11 @@
name = "runge_kutta",
hdrs = [
"runge_kutta.h",
+ "runge_kutta_helpers.h",
],
target_compatible_with = ["@platforms//os:linux"],
deps = [
+ "@com_github_google_glog//:glog",
"@org_tuxfamily_eigen//:eigen",
],
)
diff --git a/frc971/control_loops/runge_kutta.h b/frc971/control_loops/runge_kutta.h
index ed5a359..173014e 100644
--- a/frc971/control_loops/runge_kutta.h
+++ b/frc971/control_loops/runge_kutta.h
@@ -1,8 +1,11 @@
#ifndef FRC971_CONTROL_LOOPS_RUNGE_KUTTA_H_
#define FRC971_CONTROL_LOOPS_RUNGE_KUTTA_H_
+#include "glog/logging.h"
#include <Eigen/Dense>
+#include "frc971/control_loops/runge_kutta_helpers.h"
+
namespace frc971::control_loops {
// Implements Runge Kutta integration (4th order). fn is the function to
@@ -66,6 +69,144 @@
return X + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4);
}
+// Integrates f(t, y) from t0 to t0 + dt using an explicit Runge Kutta 5(4) to
+// implement an adaptive step size. Translated from Scipy.
+//
+// This uses the Dormand-Prince pair of formulas. The error is controlled
+// assuming accuracy of the fourth-order method accuracy, but steps are taken
+// using the fifth-order accurate formula (local extrapolation is done). A
+// quartic interpolation polynomial is used for the dense output.
+//
+// fn(t, y) is the function to integrate. y0 is the initial y, t0 is the
+// initial time, dt is the duration to integrate, rtol is the relative
+// tolerance, and atol is the absolute tolerance.
+template <typename F, typename T>
+T AdaptiveRungeKutta(const F &fn, T y0, double t0, double dt,
+ double rtol = 1e-3, double atol = 1e-6) {
+ // Multiply steps computed from asymptotic behaviour of errors by this.
+ constexpr double SAFETY = 0.9;
+ // Minimum allowed decrease in a step size.
+ constexpr double MIN_FACTOR = 0.2;
+ // Maximum allowed increase in a step size.
+ constexpr double MAX_FACTOR = 10;
+
+ // Final time
+ const double t_bound = t0 + dt;
+
+ constexpr int order = 5;
+ constexpr int error_estimator_order = 4;
+ constexpr int n_stages = 6;
+ constexpr int states = y0.rows();
+ const double sqrt_rows = std::sqrt(static_cast<double>(states));
+ const Eigen::Matrix<double, 1, n_stages> C =
+ (Eigen::Matrix<double, 1, n_stages>() << 0, 1.0 / 5.0, 3.0 / 10.0,
+ 4.0 / 5.0, 8.0 / 9.0, 1.0)
+ .finished();
+
+ const Eigen::Matrix<double, n_stages, order> A =
+ (Eigen::Matrix<double, n_stages, order>() << 0.0, 0.0, 0.0, 0.0, 0.0,
+ 1.0 / 5.0, 0.0, 0.0, 0.0, 0.0, 3.0 / 40.0, 9.0 / 40.0, 0.0, 0.0, 0.0,
+ 44.0 / 45.0, -56.0 / 15.0, 32.0 / 9.0, 0.0, 0.0, 19372.0 / 6561.0,
+ -25360.0 / 2187.0, 64448.0 / 6561.0, -212.0 / 729.0, 0.0,
+ 9017.0 / 3168.0, -355.0 / 33.0, 46732.0 / 5247.0, 49.0 / 176.0,
+ -5103.0 / 18656.0)
+ .finished();
+
+ const Eigen::Matrix<double, 1, n_stages> B =
+ (Eigen::Matrix<double, 1, n_stages>() << 35.0 / 384.0, 0.0,
+ 500.0 / 1113.0, 125.0 / 192.0, -2187.0 / 6784.0, 11.0 / 84.0)
+ .finished();
+
+ const Eigen::Matrix<double, 1, n_stages + 1> E =
+ (Eigen::Matrix<double, 1, n_stages + 1>() << -71.0 / 57600.0, 0.0,
+ 71.0 / 16695.0, -71.0 / 1920.0, 17253.0 / 339200.0, -22.0 / 525.0,
+ 1.0 / 40.0)
+ .finished();
+
+ T f = fn(t0, y0);
+ double h_abs = SelectRungeKuttaInitialStep(fn, t0, y0, f,
+ error_estimator_order, rtol, atol);
+ Eigen::Matrix<double, n_stages + 1, states> K;
+
+ Eigen::Matrix<double, states, 1> y = y0;
+ const double error_exponent = -1.0 / (error_estimator_order + 1.0);
+
+ double t = t0;
+ while (true) {
+ if (t >= t_bound) {
+ return y;
+ }
+
+ // Step
+ double min_step =
+ 10 * (std::nextafter(t, std::numeric_limits<double>::infinity()) - t);
+
+ // TODO(austin): max_step if we care.
+ if (h_abs < min_step) {
+ h_abs = min_step;
+ }
+
+ bool step_accepted = false;
+ bool step_rejected = false;
+
+ double t_new;
+ Eigen::Matrix<double, states, 1> y_new;
+ Eigen::Matrix<double, states, 1> f_new;
+ while (!step_accepted) {
+ // TODO(austin): Tell the user rather than just explode?
+ CHECK_GE(h_abs, min_step);
+
+ double h = h_abs;
+ t_new = t + h;
+ if (t_new >= t_bound) {
+ t_new = t_bound;
+ }
+ h = t_new - t;
+ h_abs = std::abs(h);
+
+ std::tie(y_new, f_new) =
+ RKStep<states, n_stages, order>(fn, t, y, f, h, A, B, C, K);
+
+ const Eigen::Matrix<double, states, 1> scale =
+ atol + y.array().abs().max(y_new.array().abs()) * rtol;
+
+ double error_norm =
+ (((K.transpose() * E.transpose()) * h).array() / scale.array())
+ .matrix()
+ .norm() /
+ sqrt_rows;
+
+ if (error_norm < 1) {
+ double factor;
+ if (error_norm == 0) {
+ factor = MAX_FACTOR;
+ } else {
+ factor = std::min(MAX_FACTOR,
+ SAFETY * std::pow(error_norm, error_exponent));
+ }
+
+ if (step_rejected) {
+ factor = std::min(1.0, factor);
+ }
+
+ h_abs *= factor;
+
+ step_accepted = true;
+ } else {
+ h_abs *=
+ std::max(MIN_FACTOR, SAFETY * std::pow(error_norm, error_exponent));
+ step_rejected = true;
+ }
+ }
+
+ t = t_new;
+ y = y_new;
+ f = f_new;
+ }
+
+ return y;
+}
+
} // namespace frc971::control_loops
#endif // FRC971_CONTROL_LOOPS_RUNGE_KUTTA_H_
diff --git a/frc971/control_loops/runge_kutta_helpers.h b/frc971/control_loops/runge_kutta_helpers.h
new file mode 100644
index 0000000..f9a4ecf
--- /dev/null
+++ b/frc971/control_loops/runge_kutta_helpers.h
@@ -0,0 +1,72 @@
+#ifndef FRC971_CONTROL_LOOPS_RUNGE_KUTTA_HELPERS_H_
+#define FRC971_CONTROL_LOOPS_RUNGE_KUTTA_HELPERS_H_
+
+#include "glog/logging.h"
+#include <Eigen/Dense>
+
+namespace frc971::control_loops {
+
+// Returns a reasonable Runge Kutta initial step size. This is translated from
+// scipy.
+template <typename F, typename T>
+double SelectRungeKuttaInitialStep(const F &fn, size_t t0, T y0, T f0,
+ int error_estimator_order, double rtol,
+ double atol) {
+ constexpr int states = y0.rows();
+ const Eigen::Matrix<double, states, 1> scale =
+ atol + (y0.cwiseAbs().matrix() * rtol).array();
+ const double sqrt_rows = std::sqrt(static_cast<double>(states));
+ const double d0 = (y0.array() / scale.array()).matrix().norm() / sqrt_rows;
+ const double d1 = (f0.array() / scale.array()).matrix().norm() / sqrt_rows;
+ double h0;
+ if (d0 < 1e-5 || d1 < 1e-5) {
+ h0 = 1e-6;
+ } else {
+ h0 = 0.01 * d0 / d1;
+ }
+
+ const Eigen::Matrix<double, states, 1> y1 = y0 + h0 * f0;
+ const Eigen::Matrix<double, states, 1> f1 = fn(t0 + h0, y1);
+ const double d2 =
+ ((f1 - f0).array() / scale.array()).matrix().norm() / sqrt_rows / h0;
+
+ double h1;
+ if (d1 <= 1e-15 && d2 <= 1e-15) {
+ h1 = std::max(1e-6, h0 * 1e-3);
+ } else {
+ h1 = std::pow((0.01 / std::max(d1, d2)),
+ (1.0 / (error_estimator_order + 1.0)));
+ }
+
+ return std::min(100 * h0, h1);
+}
+
+// Performs a single step of Runge Kutta integration for the adaptive algorithm
+// below. This is translated from scipy.
+template <size_t N, size_t NStages, size_t Order, typename F>
+std::tuple<Eigen::Matrix<double, N, 1>, Eigen::Matrix<double, N, 1>> RKStep(
+ const F &fn, const double t, const Eigen::Matrix<double, N, 1> &y0,
+ const Eigen::Matrix<double, N, 1> &f0, const double h,
+ const Eigen::Matrix<double, NStages, Order> &A,
+ const Eigen::Matrix<double, 1, NStages> &B,
+ const Eigen::Matrix<double, 1, NStages> &C,
+ Eigen::Matrix<double, NStages + 1, N> &K) {
+ K.template block<N, 1>(0, 0) = f0;
+ for (size_t s = 1; s < NStages; ++s) {
+ Eigen::Matrix<double, N, 1> dy =
+ K.block(0, 0, s, N).transpose() * A.block(s, 0, 1, s).transpose() * h;
+ K.template block<1, N>(s, 0) = fn(t + C(0, s) * h, y0 + dy).transpose();
+ }
+
+ Eigen::Matrix<double, N, 1> y_new =
+ y0 + h * (K.template block<NStages, N>(0, 0).transpose() * B.transpose());
+ Eigen::Matrix<double, N, 1> f_new = fn(t + h, y_new);
+
+ K.template block<1, N>(NStages, 0) = f_new.transpose();
+
+ return std::make_tuple(y_new, f_new);
+}
+
+} // namespace frc971::control_loops
+
+#endif // FRC971_CONTROL_LOOPS_RUNGE_KUTTA_HELPERS_H_
diff --git a/frc971/control_loops/runge_kutta_test.cc b/frc971/control_loops/runge_kutta_test.cc
index 1022a47..357b7f5 100644
--- a/frc971/control_loops/runge_kutta_test.cc
+++ b/frc971/control_loops/runge_kutta_test.cc
@@ -90,4 +90,28 @@
EXPECT_NEAR(y1(0, 0), RungeKuttaTimeVaryingSolution(6.0)(0, 0), 1e-7);
}
+// Tests that integrating dx/dt = e^x works with RK45
+TEST(RungeKuttaTest, RungeKuttaTimeVaryingAdaptive) {
+ ::Eigen::Matrix<double, 1, 1> y0 = RungeKuttaTimeVaryingSolution(5.0);
+
+ size_t count = 0;
+
+ ::Eigen::Matrix<double, 1, 1> y1 = AdaptiveRungeKutta(
+ [&](double t, ::Eigen::Matrix<double, 1, 1> x) {
+ ++count;
+ return (::Eigen::Matrix<double, 1, 1>()
+ << x(0, 0) * (2.0 / (::std::exp(t) + 1.0) - 1.0))
+ .finished();
+ },
+ y0, 5.0, 1.0, 1e-6, 1e-9);
+ EXPECT_NEAR(y1(0, 0), RungeKuttaTimeVaryingSolution(6.0)(0, 0), 1e-7);
+
+ // Using sympy as a benchmark, we should expect to see the function called 38
+ // times.
+ EXPECT_EQ(count, 38);
+
+ LOG(INFO) << "Got " << y1(0, 0) << " vs expected "
+ << RungeKuttaTimeVaryingSolution(6.0)(0, 0);
+}
+
} // namespace frc971::control_loops::testing