feat(model): #35 add an interface of integrals of motion

hamilflow/models/harmonic_oscillator.py → hamilflow/models/harmonic_oscillator/__init__.py

@@ -7,6 +7,11 @@
 from numpy.typing import ArrayLike
 from pydantic import BaseModel, computed_field, field_validator
 
+from hamilflow.models.harmonic_oscillator.initial_conditions import (
+    HarmonicOscillatorIC,
+    parse_ic_for_sho,
+)
+
 
 class HarmonicOscillatorSystem(BaseModel):
     """The params for the harmonic oscillator
@@ -54,19 +59,6 @@ def check_zeta_non_negative(cls, v: float) -> float:
         return v
 
 
-class HarmonicOscillatorIC(BaseModel):
-    """The initial condition for a harmonic oscillator
-
-    :cvar x0: the initial displacement
-    :cvar v0: the initial velocity
-    :cvar phi: initial phase
-    """
-
-    x0: float = 1.0
-    v0: float = 0.0
-    phi: float = 0.0
-
-
 class HarmonicOscillatorBase(ABC):
     r"""Base class to generate time series data
     for a [harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator).
@@ -155,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

@@ -0,0 +1,39 @@
+import math
+from typing import Any, Dict, cast
+
+from pydantic import BaseModel, Field
+
+
+class HarmonicOscillatorIC(BaseModel):
+    """The initial condition for a harmonic oscillator
+
+    :cvar x0: initial displacement
+    :cvar v0: initial velocity
+    :cvar phi: initial phase
+    """
+
+    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

pyproject.toml

@@ -37,6 +37,7 @@ plotly = "^5.19.0"
 
 [tool.poetry.group.dev.dependencies]
 commitizen = "*"
+pre-commit = "*"
 
 
 [tool.commitizen]

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))