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