From 01dc720e554a580dcd534424f2987a34b224d003 Mon Sep 17 00:00:00 2001
From: LM <hi@leima.is>
Date: Sat, 13 Apr 2024 18:28:43 +0200
Subject: [PATCH] fix(harmonic-oscillator): use ArrayLike instead of union of
 float and ndarray

#34
---
 hamilflow/models/harmonic_oscillator.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/hamilflow/models/harmonic_oscillator.py b/hamilflow/models/harmonic_oscillator.py
index 2547997..28cc02b 100644
--- a/hamilflow/models/harmonic_oscillator.py
+++ b/hamilflow/models/harmonic_oscillator.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 import pandas as pd
+from numpy.typing import ArrayLike
 from pydantic import BaseModel, computed_field, field_validator
 
 
@@ -89,7 +90,7 @@ def definition(self) -> Dict[str, float]:
         }
 
     @abstractmethod
-    def _x(self, t: np.typing.ArrayLike) -> np.typing.ArrayLike:
+    def _x(self, t: ArrayLike) -> ArrayLike:
         r"""Solution to simple harmonic oscillators."""
         ...
 
@@ -158,7 +159,7 @@ def __init__(
                 f"System is not a Simple Harmonic Oscillator: {self.system}"
             )
 
-    def _x(self, t: np.typing.ArrayLike) -> np.typing.ArrayLike:
+    def _x(self, t: ArrayLike) -> ArrayLike:
         r"""Solution to simple harmonic oscillators:
 
         $$