From 54f3a2b8fec6fe2b05541025e1630b7e053d57fd Mon Sep 17 00:00:00 2001
From: cmp0xff <cmp0xff@gmail.com>
Date: Fri, 16 Aug 2024 00:18:22 +0200
Subject: [PATCH] feat(test): #76 unit test for tau of u

---
 .../kepler_problem/{math.py => dynamics.py}   |  4 +-
 hamilflow/models/kepler_problem/model.py      |  2 +-
 tests/models/kepler_problem/test_dynamics.py  | 44 +++++++++++++++++++
 3 files changed, 47 insertions(+), 3 deletions(-)
 rename hamilflow/models/kepler_problem/{math.py => dynamics.py} (91%)
 create mode 100644 tests/models/kepler_problem/test_dynamics.py

diff --git a/hamilflow/models/kepler_problem/math.py b/hamilflow/models/kepler_problem/dynamics.py
similarity index 91%
rename from hamilflow/models/kepler_problem/math.py
rename to hamilflow/models/kepler_problem/dynamics.py
index ffc4fe4..f25554b 100644
--- a/hamilflow/models/kepler_problem/math.py
+++ b/hamilflow/models/kepler_problem/dynamics.py
@@ -27,8 +27,8 @@ def tau_of_u_hyperbolic(ecc: float, u: "npt.ArrayLike") -> "npt.ArrayLike":
     )
 
 
-def tau_of_u_prime(e: float, u: "npt.ArrayLike") -> "npt.ArrayLike":
-    return 1 / (1 + u) ** 2 / np.sqrt(e**2 - u**2)
+def tau_of_u_prime(ecc: float, u: "npt.ArrayLike") -> "npt.ArrayLike":
+    return -1 / (1 + u) ** 2 / np.sqrt(ecc**2 - u**2)
 
 
 def solve_u_of_tau(
diff --git a/hamilflow/models/kepler_problem/model.py b/hamilflow/models/kepler_problem/model.py
index 3b482bc..0d1e2cc 100644
--- a/hamilflow/models/kepler_problem/model.py
+++ b/hamilflow/models/kepler_problem/model.py
@@ -6,7 +6,7 @@
 import pandas as pd
 from pydantic import BaseModel, Field, field_validator
 
-from .math import (
+from .dynamics import (
     acos_with_shift,
     solve_u_of_tau,
     tau_of_u_elliptic,
diff --git a/tests/models/kepler_problem/test_dynamics.py b/tests/models/kepler_problem/test_dynamics.py
new file mode 100644
index 0000000..8b63edc
--- /dev/null
+++ b/tests/models/kepler_problem/test_dynamics.py
@@ -0,0 +1,44 @@
+from functools import partial
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pytest
+from numpy.testing import assert_array_almost_equal
+from scipy.integrate import quad
+
+from hamilflow.models.kepler_problem.dynamics import (
+    tau_of_u_elliptic,
+    tau_of_u_hyperbolic,
+    tau_of_u_parabolic,
+    tau_of_u_prime,
+)
+
+if TYPE_CHECKING:
+    from numpy import typing as npt
+
+_EPS = 0.05
+
+
+@pytest.mark.parametrize("ecc", [1 / 3, 1 / 2, 5 / 7, 1.0, 12 / 11, 27 / 13])
+def test_tau_of_u(ecc: float) -> None:
+    def integrand(u: float) -> float:
+        return tau_of_u_prime(ecc, u)
+
+    if 0 < ecc < 1:
+        tau_of_u = tau_of_u_elliptic
+        cosqr = 1 - ecc**2
+        const = -(ecc + np.arcsin(ecc) / np.sqrt(cosqr)) / cosqr
+    elif ecc == 1:
+        tau_of_u = tau_of_u_parabolic
+        const = 2 / 3
+    elif ecc > 1:
+        tau_of_u = tau_of_u_hyperbolic
+        cosqr = ecc**2 - 1
+        const = (ecc - np.arccosh(ecc) / np.sqrt(cosqr)) / cosqr
+    else:
+        raise ValueError(f"Expected ecc > 0, got {ecc}")
+
+    u_s = np.linspace(max(-1, -ecc) + _EPS, ecc - _EPS, 5)
+    rets = [quad(integrand, 0, u) for u in u_s]
+    integrals = np.array([ret[0] for ret in rets]) + const
+    assert_array_almost_equal(integrals, tau_of_u(ecc, u_s))