Repository: netket/netket

Issue:
Experimental dynamics RK4 Butcher tableau typo
Hi,

I was going through your new experimental ODE integrator code and I think I noticed a typo in one of the predefined Butcher tableaus for the `RK4` solver. You have:

```python
bt_rk4  = TableauRKExplicit(
                order = (4,),
                a = jnp.array([[0,   0,   0,   0],
                               [1/2, 0,   0,   0],
                               [0,   1/2, 0,   0],
                               [0,   0,   1,   1]], dtype=default_dtype),
                b = jnp.array( [1/6,  1/3,  1/3,  1/6], dtype=default_dtype),
                c = jnp.array( [0, 1/2, 1/2, 1], dtype=default_dtype),
                c_error = None,
                )
```

while I believe that the correct tableau should be

```python
bt_rk4  = TableauRKExplicit(
                order = (4,),
                a = jnp.array([[0,   0,   0,   0],
                               [1/2, 0,   0,   0],
                               [0,   1/2, 0,   0],
                               [0,   0,   1,   0]], dtype=default_dtype),
                b = jnp.array( [1/6,  1/3,  1/3,  1/6], dtype=default_dtype),
                c = jnp.array( [0, 1/2, 1/2, 1], dtype=default_dtype),
                c_error = None,
                )
```

(Note the very last matrix element in the `a` field should be 0, not 1, if I am reading [the relevant Wikipedia page](https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods) correctly.)

It should be a pretty easy thing to fix. :-)


Discussion:
Dammit! Thanks for the report. Btw @femtobit when you have time we should investigate why the test did not caught that…

The tableaus are copied from https://github.com/SciML/ODE.jl/blob/master/src/runge_kutta.jl and indeed there’s a typo.

Can you submit a PR fixing it? I’m on holiday for a while but can quickly merge it
Il 29 dic 2021, 13:26 +0100, Matematija ***@***.***>, ha scritto:
> Hi,
> I was going through your new experimental ODE integrator code and I think I noticed a typo in one of the predefined Butcher tableaus for the RK4 solver. You have:
> bt_rk4  = TableauRKExplicit(
>                order = (4,),
>                a = jnp.array([[0,   0,   0,   0],
>                               [1/2, 0,   0,   0],
>                               [0,   1/2, 0,   0],
>                               [0,   0,   1,   1]], dtype=default_dtype),
>                b = jnp.array( [1/6,  1/3,  1/3,  1/6], dtype=default_dtype),
>                c = jnp.array( [0, 1/2, 1/2, 1], dtype=default_dtype),
>                c_error = None,
>                )
> while I believe that the correct tableau should be
> bt_rk4  = TableauRKExplicit(
>                order = (4,),
>                a = jnp.array([[0,   0,   0,   0],
>                               [1/2, 0,   0,   0],
>                               [0,   1/2, 0,   0],
>                               [0,   0,   1,   0]], dtype=default_dtype),
>                b = jnp.array( [1/6,  1/3,  1/3,  1/6], dtype=default_dtype),
>                c = jnp.array( [0, 1/2, 1/2, 1], dtype=default_dtype),
>                c_error = None,
>                )
> (Note the very last matrix element in the a field should be 0, not 1, if I am reading the relevant Wikipedia page correctly.)
> It should be a pretty easy thing to fix. :-)
> —
> Reply to this email directly, view it on GitHub, or unsubscribe.
> You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

Sure, I can do a PR. I'll give it a day to go over other coefficients, maybe there are other typos that sneaked by tests.
Thanks for spotting the typo, @Matematija.

> Dammit! Thanks for the report. Btw @femtobit when you have time we should investigate why the test did not caught that…

The tests do not catch this typo because it does not actually affect the output of the solver. In RK tableaus, `a` is a lower triangular matrix with zero diagonal, because the slope `k[l]` only depends on `k[m]` where `m < l`. Otherwise, a slope would depend on itself (diagonal entries) or future slopes (upper diagonal entries).

In our RK step code, the slopes array `k` is zero-initialized, so by the time `k[l]` is computed, `k[m]` with `m >= l` is still zero and its coefficient does not affect the output. As a consequence, diagonal and upper triangle entries in the tableau can be arbitrary (though of course they should be zero for correctness).

I'll add a test to make some sanity checks to the tableaus.
@femtobit - I see, it's just about symbolic correctness if all intermediate slopes are zero-initialized. Thanks for the explanation.

@PhilipVinc - I can still do a PR if you want me to but if nothing is actually broken, then it's not urgent.

I am okay with closing this now unless you want to keep this issue around as a reminder.
Ah that’s good that nothing is broken.
But I’d still fix it. It’s an error after all
Il 29 dic 2021, 21:18 +0100, Matematija ***@***.***>, ha scritto:
> @femtobit - I see, it's just about symbolic correctness if all intermediate slopes are zero-initialized. Thanks for the explanation.
> @PhilipVinc - I can still do a PR if you want me to but if nothing is actually broken, then it's not urgent.
> I am okay with closing this now unless you want to keep this issue around as a reminder.
> —
> Reply to this email directly, view it on GitHub, or unsubscribe.
> You are receiving this because you were mentioned.Message ID: ***@***.***>


Resolving patch:
diff --git a/netket/experimental/dynamics/_rk_tableau.py b/netket/experimental/dynamics/_rk_tableau.py
--- a/netket/experimental/dynamics/_rk_tableau.py
+++ b/netket/experimental/dynamics/_rk_tableau.py
@@ -179,7 +179,7 @@ def __repr__(self) -> str:
 bt_feuler = TableauRKExplicit(
                 order = (1,),
                 a = jnp.zeros((1,1), dtype=default_dtype),
-                b = jnp.ones((1,1), dtype=default_dtype),
+                b = jnp.ones((1,), dtype=default_dtype),
                 c = jnp.zeros((1), dtype=default_dtype),
                 c_error = None,
                 )
@@ -213,7 +213,7 @@ def __repr__(self) -> str:
                 a = jnp.array([[0,   0,   0,   0],
                                [1/2, 0,   0,   0],
                                [0,   1/2, 0,   0],
-                               [0,   0,   1,   1]], dtype=default_dtype),
+                               [0,   0,   1,   0]], dtype=default_dtype),
                 b = jnp.array( [1/6,  1/3,  1/3,  1/6], dtype=default_dtype),
                 c = jnp.array( [0, 1/2, 1/2, 1], dtype=default_dtype),
                 c_error = None,


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
@@ -12,13 +12,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
+import jax.numpy as jnp
 import pytest
 import numpy as np
 
 import scipy.integrate as sci
 
 from netket.experimental.dynamics import Euler, Heun, Midpoint, RK4, RK12, RK23, RK45
+from netket.experimental.dynamics._rk_tableau import (
+    NamedTableau,
+    bt_feuler,
+    bt_heun,
+    bt_midpoint,
+    bt_rk12,
+    bt_rk23,
+    bt_rk4,
+    bt_rk4_dopri,
+    bt_rk4_fehlberg,
+)
+
+
+tableaus = {
+    "bt_feuler": bt_feuler,
+    "bt_heun": bt_heun,
+    "bt_midpoint": bt_midpoint,
+    "bt_rk12": bt_rk12,
+    "bt_rk23": bt_rk23,
+    "bt_rk4": bt_rk4,
+    "bt_rk4_dopri": bt_rk4_dopri,
+    "bt_rk4_fehlberg": bt_rk4_fehlberg,
+}
+
 
 explicit_fixed_step_solvers = {
     "Euler": Euler,
@@ -34,38 +58,70 @@
 }
 
 
-@pytest.mark.parametrize("solver", explicit_fixed_step_solvers)
-def test_ode_solver(solver):
-    if solver == "Euler":  # first order
+@pytest.mark.parametrize("tableau", tableaus)
+def test_tableau(tableau: str):
+    tableau = tableaus[tableau]  # type: NamedTableau
+    assert tableau.name != ""
+    td = tableau.data
+
+    for x in td.a, td.b, td.c:
+        assert np.all(np.isfinite(x))
 
-        def ode(_t, _x, **_):
-            return 1.0
+    assert td.a.ndim == 2
+    # a should be strictly upper triangular
+    np.testing.assert_array_equal(np.triu(td.a), np.zeros_like(td.a))
+    # c's should be in [0, 1]
+    assert np.all(td.c >= 0.0)
+    assert np.all(td.c <= 1.0)
 
-    else:  # quadratic function for higher-order solvers
+    assert len(td.order) in (1, 2)
+    assert len(td.order) == td.b.ndim
 
-        def ode(t, _x, **_):
-            return t
+    assert td.a.shape[0] == td.a.shape[1]
+    assert td.a.shape[0] == td.b.shape[-1]
+    assert td.a.shape[0] == td.c.shape[0]
+    if len(td.order) == 2:
+        assert td.b.shape[0] == 2
 
-    solver = explicit_fixed_step_solvers[solver]
+
+@pytest.mark.parametrize("method", explicit_fixed_step_solvers)
+def test_ode_solver(method):
+    def ode(t, x, **_):
+        return -t * x
+
+    dt = 0.01
+    n_steps = 100
+    solver = explicit_fixed_step_solvers[method](dt=dt)
 
     y0 = np.array([1.0])
-    times = np.linspace(0, 2, 10, endpoint=False)
+    times = np.linspace(0, n_steps * dt, n_steps, endpoint=False)
 
-    sol = sci.solve_ivp(ode, (0.0, 2.0), y0, t_eval=times)
+    sol = sci.solve_ivp(ode, (0.0, n_steps * dt), y0, t_eval=times)
     y_ref = sol.y[0]
 
-    solv = solver(dt=0.2)(ode, 0.0, y0)
+    solv = solver(ode, 0.0, y0)
 
     t = []
     y_t = []
-    for _ in range(10):
+    for _ in range(n_steps):
+        print(solv.t, solv.y)
         t.append(solv.t)
         y_t.append(solv.y)
         solv.step()
     y_t = np.asarray(y_t)
 
+    assert np.all(np.isfinite(t))
+    assert np.all(np.isfinite(y_t))
+
     np.testing.assert_allclose(t, times)
-    np.testing.assert_allclose(y_t[:, 0], y_ref)
+
+    # somewhat arbitrary tolerances, that may still help spot
+    # errors introduced later
+    rtol = {
+        "Euler": 1e-2,
+        "RK4": 5e-4,
+    }.get(method, 1e-3)
+    np.testing.assert_allclose(y_t[:, 0], y_ref, rtol=rtol)
 
 
 @pytest.mark.parametrize("solver", explicit_adaptive_solvers)