add tests for underdamped

kausalflow · Feb 20, 2024 · 70d63ad · 70d63ad
1 parent 0415d0e
commit 70d63ad
Showing 1 changed file with 61 additions and 19 deletions.
diff --git a/tests/test_models/test_harmonic_oscillator.py b/tests/test_models/test_harmonic_oscillator.py
@@ -4,27 +4,24 @@
 from hamiltonian_flow.models.harmonic_oscillator import HarmonicOscillator
 
 
-@pytest.fixture
-def omega():
-    return 0.5
-
-
-@pytest.mark.parametrize("omega", [0.5])
 @pytest.mark.parametrize(
-    "expected",
+    "omega,expected",
     [
-        [
-            {"t": 0.0, "x": 1.0},
-            {"t": 1.2566370614359172, "x": 0.8090169943749475},
-            {"t": 2.5132741228718345, "x": 0.30901699437494745},
-            {"t": 3.7699111843077517, "x": -0.30901699437494734},
-            {"t": 5.026548245743669, "x": -0.8090169943749473},
-            {"t": 6.283185307179586, "x": -1.0},
-            {"t": 7.5398223686155035, "x": -0.8090169943749475},
-            {"t": 8.79645943005142, "x": -0.30901699437494756},
-            {"t": 10.053096491487338, "x": 0.30901699437494723},
-            {"t": 11.309733552923255, "x": 0.8090169943749473},
-        ]
+        (
+            0.5,
+            [
+                {"t": 0.0, "x": 1.0},
+                {"t": 1.2566370614359172, "x": 0.8090169943749475},
+                {"t": 2.5132741228718345, "x": 0.30901699437494745},
+                {"t": 3.7699111843077517, "x": -0.30901699437494734},
+                {"t": 5.026548245743669, "x": -0.8090169943749473},
+                {"t": 6.283185307179586, "x": -1.0},
+                {"t": 7.5398223686155035, "x": -0.8090169943749475},
+                {"t": 8.79645943005142, "x": -0.30901699437494756},
+                {"t": 10.053096491487338, "x": 0.30901699437494723},
+                {"t": 11.309733552923255, "x": 0.8090169943749473},
+            ],
+        )
     ],
 )
 def test_simple_harmonic_oscillator(omega, expected):
@@ -33,3 +30,48 @@ def test_simple_harmonic_oscillator(omega, expected):
     df = ho(n_periods=1, n_samples_per_period=10)
 
     pd.testing.assert_frame_equal(df, pd.DataFrame(expected))
+
+
+@pytest.mark.parametrize(
+    "omega,zeta,expected",
+    [
+        (
+            0.5,
+            0,
+            [
+                {"t": 0.0, "x": 1.0},
+                {"t": 1.2566370614359172, "x": 0.8090169943749475},
+                {"t": 2.5132741228718345, "x": 0.30901699437494745},
+                {"t": 3.7699111843077517, "x": -0.30901699437494734},
+                {"t": 5.026548245743669, "x": -0.8090169943749473},
+                {"t": 6.283185307179586, "x": -1.0},
+                {"t": 7.5398223686155035, "x": -0.8090169943749475},
+                {"t": 8.79645943005142, "x": -0.30901699437494756},
+                {"t": 10.053096491487338, "x": 0.30901699437494723},
+                {"t": 11.309733552923255, "x": 0.8090169943749473},
+            ],
+        ),
+        (
+            0.5,
+            0.5,
+            [
+                {"t": 0.0, "x": 1.0},
+                {"t": 1.2566370614359172, "x": 0.8814807646720196},
+                {"t": 2.5132741228718345, "x": 0.6291724461162962},
+                {"t": 3.7699111843077517, "x": 0.35960997861868876},
+                {"t": 5.026548245743669, "x": 0.1386644177513349},
+                {"t": 6.283185307179586, "x": -0.008951254750058093},
+                {"t": 7.5398223686155035, "x": -0.08578239268073515},
+                {"t": 8.79645943005142, "x": -0.10837036278920763},
+                {"t": 10.053096491487338, "x": -0.09717507559642745},
+                {"t": 11.309733552923255, "x": -0.07035823865195263},
+            ],
+        ),
+    ],
+)
+def test_underdamped_harmonic_oscillator(omega, zeta, expected):
+    ho = HarmonicOscillator(system={"omega": omega, "zeta": zeta})
+
+    df = ho(n_periods=1, n_samples_per_period=10)
+
+    pd.testing.assert_frame_equal(df, pd.DataFrame(expected))