Repository: netket/netket

Issue:
Adaptive time step fails if t0 is an integer
Hi,

Adaptive time step integration fails if `t0` in `RungeKuttaIntegrator` is an integer (which for me is an easy mistake since `t0=0`).

Example code:
```python
import numpy as np
import netket.experimental.dynamics as nkd

def df(t, y, stage=None):
    return np.sin(t)**2*y

int_config = nkd.RK23(dt=0.04, adaptive=True, atol=1e-3, rtol=1e-3, dt_limits=[1e-3,1e-1])
integrator = int_config(df, 0, np.array([1.0]))  # <-- the second argument has to be a float

y_vec = [integrator.y]
t_vec = [0]
while integrator.t < 0.5:
    integrator.step()
    t_vec.append(integrator.t)
    y_vec.append(integrator.y)
```

An additional problem is that the traceback is confusing:
```
Traceback (most recent call last):
  File "/Users/macek/Work/Projects/TimeEvolution/DensityMatrix/script.py", line 13, in <module>
    integrator.step()
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 357, in step
    self._rkstate = self._do_step(self._rkstate, max_dt)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 369, in _do_step_adaptive
    return general_time_step_adaptive(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 46, in _maybe_jitted_fun
    res = fun(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 249, in general_time_step_adaptive
    return jax.lax.cond(
           ^^^^^^^^^^^^^
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 151, in __repr__
    f", {self.last_norm:.2e}" if self.last_norm is not None else "",
        ^^^^^^^^^^^^^^^^^^^^
ValueError: Unknown format code 'e' for object of type 'str'
```
The last `ValueError` is from `RungeKuttaState.__repr__`, so quite confusing.

Adaptive time step fails if t0 is an integer
Hi,

Adaptive time step integration fails if `t0` in `RungeKuttaIntegrator` is an integer (which for me is an easy mistake since `t0=0`).

Example code:
```python
import numpy as np
import netket.experimental.dynamics as nkd

def df(t, y, stage=None):
    return np.sin(t)**2*y

int_config = nkd.RK23(dt=0.04, adaptive=True, atol=1e-3, rtol=1e-3, dt_limits=[1e-3,1e-1])
integrator = int_config(df, 0, np.array([1.0]))  # <-- the second argument has to be a float

y_vec = [integrator.y]
t_vec = [0]
while integrator.t < 0.5:
    integrator.step()
    t_vec.append(integrator.t)
    y_vec.append(integrator.y)
```

An additional problem is that the traceback is confusing:
```
Traceback (most recent call last):
  File "/Users/macek/Work/Projects/TimeEvolution/DensityMatrix/script.py", line 13, in <module>
    integrator.step()
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 357, in step
    self._rkstate = self._do_step(self._rkstate, max_dt)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 369, in _do_step_adaptive
    return general_time_step_adaptive(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 46, in _maybe_jitted_fun
    res = fun(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 249, in general_time_step_adaptive
    return jax.lax.cond(
           ^^^^^^^^^^^^^
  File "/Users/macek/.pyenv/NetKet/lib/python3.12/site-packages/netket/experimental/dynamics/_rk_solver_structures.py", line 151, in __repr__
    f", {self.last_norm:.2e}" if self.last_norm is not None else "",
        ^^^^^^^^^^^^^^^^^^^^
ValueError: Unknown format code 'e' for object of type 'str'
```
The last `ValueError` is from `RungeKuttaState.__repr__`, so quite confusing.



Discussion:

Resolving patch:
diff --git a/netket/experimental/dynamics/_rk_solver_structures.py b/netket/experimental/dynamics/_rk_solver_structures.py
--- a/netket/experimental/dynamics/_rk_solver_structures.py
+++ b/netket/experimental/dynamics/_rk_solver_structures.py
@@ -24,6 +24,7 @@
 from netket.utils.mpi.primitives import mpi_all_jax
 from netket.utils.struct import dataclass, field
 from netket.utils.types import Array, PyTree
+from netket.utils.numbers import dtype as _dtype
 
 from . import _rk_tableau as rkt
 
@@ -143,13 +144,22 @@ class RungeKuttaState:
     """Flags containing information on the solver state."""
 
     def __repr__(self):
-        return "RKState(step_no(total)={}({}), t={}, dt={:.2e}{}{})".format(
+        try:
+            dt = "{self.dt:.2e}"
+            last_norm = f", {self.last_norm:.2e}" if self.last_norm is not None else ""
+            accepted = (f", {'A' if self.accepted else 'R'}",)
+        except (ValueError, TypeError):
+            dt = f"{self.dt}"
+            last_norm = f"{self.last_norm}"
+            accepted = f"{SolverFlags.INFO_STEP_ACCEPTED}"
+
+        return "RKState(step_no(total)={}({}), t={}, dt={}{}{})".format(
             self.step_no,
             self.step_no_total,
             self.t.value,
-            self.dt,
-            f", {self.last_norm:.2e}" if self.last_norm is not None else "",
-            f", {'A' if self.accepted else 'R'}",
+            dt,
+            last_norm,
+            accepted,
         )
 
     @property
@@ -330,6 +340,10 @@ def __post_init__(self):
         if self.dt_limits is None:
             self.dt_limits = (None, 10 * self.initial_dt)
 
+        t_dtype = jnp.result_type(_dtype(self.t0), _dtype(self.initial_dt))
+        setattr(self, "t0", jnp.array(self.t0, dtype=t_dtype))
+        setattr(self, "initial_dt", jnp.array(self.initial_dt, dtype=t_dtype))
+
         self._rkstate = RungeKuttaState(
             step_no=0,
             step_no_total=0,


Test patch:
diff --git a/test/dynamics/test_solvers.py b/test/dynamics/test_solvers.py
--- a/test/dynamics/test_solvers.py
+++ b/test/dynamics/test_solvers.py
@@ -15,6 +15,7 @@
 import pytest
 import numpy as np
 
+import jax
 import scipy.integrate as sci
 
 from netket.experimental.dynamics import Euler, Heun, Midpoint, RK4, RK12, RK23, RK45
@@ -138,6 +139,48 @@ def ode(t, x, **_):
     np.testing.assert_allclose(y_t[:, 0], y_ref, rtol=rtol)
 
 
+def test_ode_repr():
+    dt = 0.01
+
+    def ode(t, x, **_):
+        return -t * x
+
+    solver = RK23(dt=dt, adaptive=True)
+
+    y0 = np.array([1.0])
+    solv = solver(ode, 0.0, y0)
+
+    assert isinstance(repr(solv), str)
+    assert isinstance(repr(solv._rkstate), str)
+
+    @jax.jit
+    def _test_jit_repr(x):
+        assert isinstance(repr(x), str)
+        return 1
+
+    _test_jit_repr(solv._rkstate)
+    # _test_jit_repr(solv) # this is broken. should be fixed in the zukumft
+
+
+def test_solver_t0_is_integer():
+    # See issue netket/netket#1735
+    # https://github.com/netket/netket/issues/1735
+
+    def df(t, y, stage=None):
+        return np.sin(t) ** 2 * y
+
+    int_config = RK23(
+        dt=0.04, adaptive=True, atol=1e-3, rtol=1e-3, dt_limits=[1e-3, 1e-1]
+    )
+    integrator = int_config(
+        df, 0, np.array([1.0])
+    )  # <-- the second argument has to be a float
+
+    integrator.step()
+    assert integrator.t > 0.0
+    assert integrator.t.dtype == integrator.dt.dtype
+
+
 @pytest.mark.parametrize("solver", explicit_adaptive_solvers_params)
 def test_adaptive_solver(solver):
     tol = 1e-7