Repository: netket/netket

Issue:
Bug in TDVPSchmitt
### Noticed in https://github.com/orgs/netket/discussions/1549

TDVPSchmitt gives the wrong dynamics.

I bisected it to commit 5ccd622e8d8513d6c8209c28813f50b2d0fd99d1  (cc @inailuig ) which switched the order of the axes used to sample.



<div type='discussions-op-text'>

<sup>Originally posted by **whatisspin** August  1, 2023</sup>
Hello,

Thanks for adding this driver. However, I could not get it to work. Here is a simple example.

```python
import numpy as np
import netket as nk
import netket.experimental as nkx

L=10
g= nk.graph.Chain(L)
hi= nk.hilbert.Spin(0.5,L)
H= nk.operator.Ising(hi,g,h=0.5)
H1= nk.operator.Ising(hi,g,h=1.0)

modelRBM= nk.models.RBM(param_dtype=np.complex128,alpha=2)

sa=nk.sampler.MetropolisLocal(hi)
vs= nk.vqs.MCState(sampler=sa,model=modelRBM,n_samples=2**10)
print(vs.n_parameters)

vmc= nk.VMC(H,nk.optimizer.Sgd(0.01),variational_state=vs, preconditioner=nk.optimizer.SR(diag_shift=0.01,holomorphic=True))
vmc.run(2000)


obs= {"sx": sum(nk.operator.spin.sigmax(hi, i) for i in range(L))}

#Time evolution

integrator= nkx.dynamics.Euler(dt=0.001)
qgt= nk.optimizer.qgt.QGTJacobianDense(holomorphic=True)

te= nkx.driver.TDVPSchmitt(operator=H1, variational_state=vs, integrator=integrator, holomorphic=True)
te.run(T=1.0,out="Ising10_tdvp_schmitt",obs=obs)


te= nkx.TDVP(operator=H1, variational_state=vs, integrator=integrator, qgt=qgt)
te.run(T=1.0,out="Ising10_tdvp",obs=obs)
```
I also tried different values of `snr_atol`, but it did not work. What am I missing here?

![image](https://github.com/netket/netket/assets/82375136/27d949f1-9099-4128-9c76-efe495dc5fe2)
</div>
Bug in TDVPSchmitt
### Noticed in https://github.com/orgs/netket/discussions/1549

TDVPSchmitt gives the wrong dynamics.

I bisected it to commit 5ccd622e8d8513d6c8209c28813f50b2d0fd99d1  (cc @inailuig ) which switched the order of the axes used to sample.



<div type='discussions-op-text'>

<sup>Originally posted by **whatisspin** August  1, 2023</sup>
Hello,

Thanks for adding this driver. However, I could not get it to work. Here is a simple example.

```python
import numpy as np
import netket as nk
import netket.experimental as nkx

L=10
g= nk.graph.Chain(L)
hi= nk.hilbert.Spin(0.5,L)
H= nk.operator.Ising(hi,g,h=0.5)
H1= nk.operator.Ising(hi,g,h=1.0)

modelRBM= nk.models.RBM(param_dtype=np.complex128,alpha=2)

sa=nk.sampler.MetropolisLocal(hi)
vs= nk.vqs.MCState(sampler=sa,model=modelRBM,n_samples=2**10)
print(vs.n_parameters)

vmc= nk.VMC(H,nk.optimizer.Sgd(0.01),variational_state=vs, preconditioner=nk.optimizer.SR(diag_shift=0.01,holomorphic=True))
vmc.run(2000)


obs= {"sx": sum(nk.operator.spin.sigmax(hi, i) for i in range(L))}

#Time evolution

integrator= nkx.dynamics.Euler(dt=0.001)
qgt= nk.optimizer.qgt.QGTJacobianDense(holomorphic=True)

te= nkx.driver.TDVPSchmitt(operator=H1, variational_state=vs, integrator=integrator, holomorphic=True)
te.run(T=1.0,out="Ising10_tdvp_schmitt",obs=obs)


te= nkx.TDVP(operator=H1, variational_state=vs, integrator=integrator, qgt=qgt)
te.run(T=1.0,out="Ising10_tdvp",obs=obs)
```
I also tried different values of `snr_atol`, but it did not work. What am I missing here?

![image](https://github.com/netket/netket/assets/82375136/27d949f1-9099-4128-9c76-efe495dc5fe2)
</div>


Discussion:
don't we have tests for the TDVP ?
Ok I just found the bug. It was an extra transpose that should have been removed when we swapped the axes.

@gcarleo yes we do, but apparently not good enough. If someone has an idea on how to properly test this I would be happy to hear it.

don't we have tests for the TDVP ?
Ok I just found the bug. It was an extra transpose that should have been removed when we swapped the axes.

@gcarleo yes we do, but apparently not good enough. If someone has an idea on how to properly test this I would be happy to hear it.


Resolving patch:
diff --git a/netket/experimental/driver/tdvp_schmitt.py b/netket/experimental/driver/tdvp_schmitt.py
--- a/netket/experimental/driver/tdvp_schmitt.py
+++ b/netket/experimental/driver/tdvp_schmitt.py
@@ -204,7 +204,7 @@ def info(self, depth=0):
 @partial(jax.jit, static_argnames=("n_samples"))
 def _impl(parameters, n_samples, E_loc, S, rhs_coeff, rcond, rcond_smooth, snr_atol):
     E = stats.statistics(E_loc)
-    ΔE_loc = E_loc.T.reshape(-1, 1) - E.mean
+    ΔE_loc = E_loc.reshape(-1, 1) - E.mean
 
     stack_jacobian = S.mode == "complex"
 


Test patch: