some diverse hints

dynax/estimation.py

@@ -430,7 +430,7 @@ def transfer_function(sys: DynamicalSystem, to_states: bool = False, **kwargs):
     linsys = sys.linearize(**kwargs)
     A, B, C, D = linsys.A, linsys.B, linsys.C, linsys.D
 
-    def H(s):
+    def H(s: complex):
         """Transfer-function at s."""
         identity = np.eye(linsys.n_states)
         phi_B = jnp.linalg.solve(s * identity - A, B)

dynax/evolution.py

@@ -5,7 +5,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from numpy.typing import ArrayLike
+from jaxtyping import Array, ArrayLike
 
 from .interpolation import spline_it
 from .system import DynamicalSystem
@@ -44,7 +44,7 @@ def __call__(
         ucoeffs: Optional[ArrayLike] = None,
         squeeze: bool = True,
         **diffeqsolve_kwargs,
-    ):
+    ) -> tuple[Array, Array]:
         """Solve initial value problem for state and output trajectories."""
         t = jnp.asarray(t)
         x0 = jnp.asarray(x0)

dynax/system.py

@@ -126,9 +126,9 @@ def output(self, x, u=None, t=None):
     def linearize(self, x0=None, u0=None, t=None) -> "LinearSystem":
         """Compute the approximate linearized system around a point."""
         if x0 is None:
-            x0 = np.zeros(self.n_states)
+            x0 = jnp.zeros(self.n_states)
         if u0 is None:
-            u0 = np.zeros(self.n_inputs)
+            u0 = jnp.zeros(self.n_inputs)
         A, B, C, D = _linearize(self.vector_field, self.output, x0, u0, t)
         # jax creates empty arrays
         if B.size == 0: