From 21d5c96f33624883836a31f03973d391b216a67a Mon Sep 17 00:00:00 2001
From: cmp0xff <cmp0xff@gmail.com>
Date: Thu, 22 Aug 2024 18:16:28 +0200
Subject: [PATCH] chore(ruff): #80 comment by @emptymalei
 https://github.com/kausalflow/hamilflow/pull/86#discussion_r1725718123

---
 docs/tutorials/harmonic_oscillator.py          |  4 +++-
 hamilflow/models/brownian_motion.py            |  4 +++-
 hamilflow/models/harmonic_oscillator.py        | 10 ++++++++--
 hamilflow/models/harmonic_oscillator_chain.py  |  8 ++++++--
 hamilflow/models/kepler_problem/numerics.py    |  4 +++-
 tests/models/test_harmonic_oscillator_chain.py |  4 +++-
 6 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/docs/tutorials/harmonic_oscillator.py b/docs/tutorials/harmonic_oscillator.py
index a39afb0..6717a86 100644
--- a/docs/tutorials/harmonic_oscillator.py
+++ b/docs/tutorials/harmonic_oscillator.py
@@ -95,7 +95,9 @@
         DampedHarmonicOscillator(system=s)(
             n_periods=n_periods,
             n_samples_per_period=n_samples_per_period,
-        ).assign(system=rf"{s_name} (omega = {s.get('omega')}, zeta = {s.get('zeta')})"),
+        ).assign(
+            system=rf"{s_name} (omega = {s.get('omega')}, zeta = {s.get('zeta')})",
+        ),
     )
 
 fig = px.line(
diff --git a/hamilflow/models/brownian_motion.py b/hamilflow/models/brownian_motion.py
index 8dc8ccc..e1bff11 100644
--- a/hamilflow/models/brownian_motion.py
+++ b/hamilflow/models/brownian_motion.py
@@ -147,7 +147,9 @@ class BrownianMotion:
     def __init__(
         self,
         system: Mapping[str, float],
-        initial_condition: Mapping[str, "Sequence[float] | npt.ArrayLike"] | None = None,
+        initial_condition: (
+            Mapping[str, "Sequence[float] | npt.ArrayLike"] | None
+        ) = None,
     ):
         initial_condition = initial_condition or {}
         self.system = BrownianMotionSystem.model_validate(system)
diff --git a/hamilflow/models/harmonic_oscillator.py b/hamilflow/models/harmonic_oscillator.py
index 9f29d3a..2199e82 100644
--- a/hamilflow/models/harmonic_oscillator.py
+++ b/hamilflow/models/harmonic_oscillator.py
@@ -253,7 +253,10 @@ def _x_under_damped(self, t: "Sequence[float] | npt.ArrayLike") -> npt.ArrayLike
         omega_damp = self.system.omega * np.sqrt(1 - self.system.zeta)
         return (
             self.initial_condition.x0 * np.cos(omega_damp * t)
-            + (self.system.zeta * self.system.omega * self.initial_condition.x0 + self.initial_condition.v0)
+            + (
+                self.system.zeta * self.system.omega * self.initial_condition.x0
+                + self.initial_condition.v0
+            )
             / omega_damp
             * np.sin(omega_damp * t)
         ) * np.exp(-self.system.zeta * self.system.omega * t)
@@ -296,7 +299,10 @@ def _x_over_damped(self, t: "Sequence[float] | npt.ArrayLike") -> npt.ArrayLike:
 
         return (
             self.initial_condition.x0 * np.cosh(gamma_damp * t)
-            + (self.system.zeta * self.system.omega * self.initial_condition.x0 + self.initial_condition.v0)
+            + (
+                self.system.zeta * self.system.omega * self.initial_condition.x0
+                + self.initial_condition.v0
+            )
             / gamma_damp
             * np.sinh(gamma_damp * t)
         ) * np.exp(-self.system.zeta * self.system.omega * t)
diff --git a/hamilflow/models/harmonic_oscillator_chain.py b/hamilflow/models/harmonic_oscillator_chain.py
index 7b95228..e0bbc19 100644
--- a/hamilflow/models/harmonic_oscillator_chain.py
+++ b/hamilflow/models/harmonic_oscillator_chain.py
@@ -77,14 +77,18 @@ def definition(
         self,
     ) -> dict[
         str,
-        float | dict[str, dict[str, float | list[float]]] | list[dict[str, dict[str, float | tuple[float, float]]]],
+        float
+        | dict[str, dict[str, float | list[float]]]
+        | list[dict[str, dict[str, float | tuple[float, float]]]],
     ]:
         """model params and initial conditions defined as a dictionary."""
         return dict(
             omega=self.omega,
             n_dof=self.n_dof,
             free_mode=self.free_mode.definition,
-            independent_csho_modes=[rwm.definition for rwm in self.independent_csho_modes],
+            independent_csho_modes=[
+                rwm.definition for rwm in self.independent_csho_modes
+            ],
         )
 
     def _z(
diff --git a/hamilflow/models/kepler_problem/numerics.py b/hamilflow/models/kepler_problem/numerics.py
index 0909e7a..015572f 100644
--- a/hamilflow/models/kepler_problem/numerics.py
+++ b/hamilflow/models/kepler_problem/numerics.py
@@ -66,7 +66,9 @@ def nsolve_u_from_tau_bisect(ecc: float, tau: "npt.ArrayLike") -> list[OptimizeR
         raise ValueError(f"Expect ecc > 0, ecc != 1, got {ecc}")
 
     def f(u: float, tau: float) -> np.float64:
-        return np.finfo(np.float64).max if u == -1 else np.float64(tau_of_u(ecc, u) - tau)
+        return (
+            np.finfo(np.float64).max if u == -1 else np.float64(tau_of_u(ecc, u) - tau)
+        )
 
     return [toms748(f, max(-1, -ecc), ecc, (ta,), 2, full_output=True) for ta in tau_s]
 
diff --git a/tests/models/test_harmonic_oscillator_chain.py b/tests/models/test_harmonic_oscillator_chain.py
index afae511..2a6b7f6 100644
--- a/tests/models/test_harmonic_oscillator_chain.py
+++ b/tests/models/test_harmonic_oscillator_chain.py
@@ -26,7 +26,9 @@ def free_mode(self, request: pytest.FixtureRequest) -> dict[str, int]:
         return dict(zip(("x0", "v0"), request.param))
 
     @pytest.fixture(
-        params=chain.from_iterable(product(_possible_wave_modes, repeat=r) for r in range(3)),
+        params=chain.from_iterable(
+            product(_possible_wave_modes, repeat=r) for r in range(3)
+        ),
     )
     def wave_modes(
         self,