Set energies_iters to max compatible with diffrax

jpc/_utils.py

@@ -119,7 +119,7 @@ def loop_body(state):
         energies_iters = energies_iters.at[:, t].set(energies)
         return t + 1, energies_iters
 
-    energies_iters = zeros((len(model), 1000))
+    energies_iters = zeros((len(model), 4096))
     _, energies_iters = jax.lax.while_loop(
         lambda state: state[0] < t_max,
         loop_body,