feat(oscillator): #35 parse initial condition for sho

hamilflow/models/harmonic_oscillator/__init__.py

@@ -7,7 +7,10 @@
 from numpy.typing import ArrayLike
 from pydantic import BaseModel, computed_field, field_validator
 
-from hamilflow.models.harmonic_oscillator.initial_conditions import HarmonicOscillatorIC
+from hamilflow.models.harmonic_oscillator.initial_conditions import (
+    HarmonicOscillatorIC,
+    parse_ic_for_sho,
+)
 
 
 class HarmonicOscillatorSystem(BaseModel):
@@ -144,6 +147,7 @@ def __init__(
         system: Dict[str, float],
         initial_condition: Optional[Dict[str, float]] = {},
     ):
+        initial_condition = parse_ic_for_sho(system["omega"], **initial_condition)
         super().__init__(system, initial_condition)
         if self.system.type != "simple":
             raise ValueError(

hamilflow/models/harmonic_oscillator/initial_conditions.py

@@ -1,3 +1,6 @@
+import math
+from typing import Any, Dict, cast
+
 from pydantic import BaseModel, Field
 
 
@@ -12,3 +15,25 @@ class HarmonicOscillatorIC(BaseModel):
     x0: float = Field(default=1.0)
     v0: float = Field(default=0.0)
     phi: float = Field(default=0.0)
+
+
+def parse_ic_for_sho(omega: float, **kwargs: Any) -> Dict[str, float]:
+    "Support alternative initial conditions"
+    match keys := {*kwargs.keys()}:
+        case set() if keys <= {"x0", "v0", "phi"}:
+            ret = {str(k): float(v) for k, v in kwargs.items()}
+        case set() if keys == {"x0", "t0"}:
+            ret = dict(x0=float(kwargs["x0"]), v0=0.0, phi=-float(omega * kwargs["t0"]))
+        case set() if keys == {"E", "t0"}:
+            ene = cast(float, kwargs["E"])
+            ret = dict(
+                x0=math.sqrt(2 * ene) / omega, v0=0.0, phi=-float(omega * kwargs["t0"])
+            )
+        case _:
+            raise ValueError(
+                f"Unsupported variable names as an initial condition: {keys}"
+            )
+    if phi := ret.get("phi"):
+        ret["phi"] = phi % (2 * math.pi)
+
+    return ret

tests/test_models/test_harmonic_oscillators/test_initial_conditions.py

@@ -0,0 +1,35 @@
+import math
+
+import pytest
+
+from hamilflow.models.harmonic_oscillator.initial_conditions import parse_ic_for_sho
+
+
+@pytest.fixture()
+def omega() -> float:
+    return 2 * math.pi
+
+
+class TestParseICForSHO:
+    @pytest.mark.parametrize(
+        ("input", "expected"),
+        [
+            (
+                dict(x0=1.0, v0=1.0, phi=7.0),
+                dict(x0=1.0, v0=1.0, phi=7 % (2 * math.pi)),
+            ),
+            (dict(x0=1.0, t0=1.0), dict(x0=1.0, v0=0.0, phi=0.0)),
+            (
+                dict(E=1.0, t0=1.0),
+                dict(x0=math.sqrt(2.0) / (2 * math.pi), v0=0.0, phi=0.0),
+            ),
+        ],
+    )
+    def test_output(
+        self, omega: float, input: dict[str, float], expected: dict[str, float]
+    ) -> None:
+        assert parse_ic_for_sho(omega, **input) == expected
+
+    def test_raise(self, omega: float) -> None:
+        with pytest.raises(ValueError):
+            parse_ic_for_sho(omega, **dict(x0=1.0, E=2.0))