Repository: netket/netket

Issue:
`QGTJacobian***(holomorphic=True)` for R->C models returns wrong output. Should error.
cc @attila-i-szabo

I agree that `holomorphic=True` is wrong in that case, but it should rather throw an error than give wrong results...
It took me a while to notice that this was the cause of some wrong optimisations.

```python
import netket as nk
import jax

L = 20
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)
hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)
ha = nk.operator.Ising(hilbert=hi, graph=g, h=1.0)
ma = nk.models.RBMModPhase(alpha=1, param_dtype=float)
sa = nk.sampler.MetropolisLocal(hi, n_chains=16)
vs = nk.vqs.MCState(sa, ma, n_samples=1000, n_discard_per_chain=100)

qgt_holo = nk.optimizer.qgt.QGTJacobianDense(holomorphic=True)
qgt_nonholo = nk.optimizer.qgt.QGTJacobianDense(holomorphic=False)
qgt_otf = nk.optimizer.qgt.QGTOnTheFly()

S_holo = vs.quantum_geometric_tensor(qgt_holo)
S_nonholo = vs.quantum_geometric_tensor(qgt_nonholo)
S_otf = vs.quantum_geometric_tensor(qgt_otf)

F = vs.parameters

r_holo = S_holo@F
r_nonholo = S_nonholo@F
r_otf = S_otf@F

jax.tree_map(lambda x,y:x-y, r_holo, r_nonholo)
```
`QGTJacobian***(holomorphic=True)` for R->C models returns wrong output. Should error.
cc @attila-i-szabo

I agree that `holomorphic=True` is wrong in that case, but it should rather throw an error than give wrong results...
It took me a while to notice that this was the cause of some wrong optimisations.

```python
import netket as nk
import jax

L = 20
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)
hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)
ha = nk.operator.Ising(hilbert=hi, graph=g, h=1.0)
ma = nk.models.RBMModPhase(alpha=1, param_dtype=float)
sa = nk.sampler.MetropolisLocal(hi, n_chains=16)
vs = nk.vqs.MCState(sa, ma, n_samples=1000, n_discard_per_chain=100)

qgt_holo = nk.optimizer.qgt.QGTJacobianDense(holomorphic=True)
qgt_nonholo = nk.optimizer.qgt.QGTJacobianDense(holomorphic=False)
qgt_otf = nk.optimizer.qgt.QGTOnTheFly()

S_holo = vs.quantum_geometric_tensor(qgt_holo)
S_nonholo = vs.quantum_geometric_tensor(qgt_nonholo)
S_otf = vs.quantum_geometric_tensor(qgt_otf)

F = vs.parameters

r_holo = S_holo@F
r_nonholo = S_nonholo@F
r_otf = S_otf@F

jax.tree_map(lambda x,y:x-y, r_holo, r_nonholo)
```


Discussion:
This would be particularly important if we get the TDVP for real-parameters working...
This would be particularly important if we get the TDVP for real-parameters working...

Resolving patch:
diff --git a/netket/optimizer/qgt/qgt_jacobian_common.py b/netket/optimizer/qgt/qgt_jacobian_common.py
--- a/netket/optimizer/qgt/qgt_jacobian_common.py
+++ b/netket/optimizer/qgt/qgt_jacobian_common.py
@@ -24,21 +24,37 @@
 @partial(jax.jit, static_argnums=(0, 4, 5))
 def _choose_jacobian_mode(apply_fun, pars, model_state, samples, mode, holomorphic):
     homogeneous_vars = nkjax.tree_ishomogeneous(pars)
+    leaf_iscomplex = nkjax.tree_leaf_iscomplex(pars)
 
     if holomorphic is True:
-        if not homogeneous_vars:
+        if homogeneous_vars and leaf_iscomplex:
+            ## all complex parameters
+            mode = "holomorphic"
+        elif homogeneous_vars and not leaf_iscomplex:
+            # all real parameters
+            raise ValueError(
+                dedent(
+                    """
+                A function with real parameters cannot be holomorphic. 
+                
+                Please remove the kw-arg `holomorphic=True`.
+                """
+                )
+            )
+        else:
+            # mixed complex and real parameters
             warnings.warn(
                 dedent(
                     """The ansatz has non homogeneous variables, which might not behave well with the
                        holomorhic implementation.
+
                        Use `holomorphic=False` or mode='complex' for more accurate results but
                        lower performance.
                     """
                 )
             )
-        mode = "holomorphic"
+            mode = "holomorphic"
     else:
-        leaf_iscomplex = nkjax.tree_leaf_iscomplex(pars)
         complex_output = nkjax.is_complex(
             jax.eval_shape(
                 apply_fun,


Test patch:
diff --git a/test/dynamics/test_driver.py b/test/dynamics/test_driver.py
--- a/test/dynamics/test_driver.py
+++ b/test/dynamics/test_driver.py
@@ -72,9 +72,9 @@ def _stop_after_one_step(step, *_):
 
 nqs_models = [
     pytest.param(nk.models.RBM(alpha=1, dtype=np.complex128), id="RBM(complex128)"),
-    pytest.param(
-        nk.models.RBMModPhase(alpha=1, dtype=np.float64), id="RBMModPhase(float64)"
-    ),
+    # pytest.param(
+    #    nk.models.RBMModPhase(alpha=1, dtype=np.float64), id="RBMModPhase(float64)"
+    # ),
 ]
 
 
@@ -83,11 +83,12 @@ def _stop_after_one_step(step, *_):
 @pytest.mark.parametrize("propagation_type", ["real", "imag"])
 def test_one_fixed_step(model, integrator, propagation_type):
     ha, vstate, _ = _setup_system(L=2, model=model)
+    holomorphic = jnp.issubdtype(vstate.model.param_dtype, jnp.complexfloating)
     te = nkx.TDVP(
         ha,
         vstate,
         integrator,
-        qgt=nk.optimizer.qgt.QGTJacobianDense(holomorphic=True),
+        qgt=nk.optimizer.qgt.QGTJacobianDense(holomorphic=holomorphic),
         propagation_type=propagation_type,
     )
     te.run(T=0.01, callback=_stop_after_one_step)
diff --git a/test/optimizer/test_qgt_itersolve.py b/test/optimizer/test_qgt_itersolve.py
--- a/test/optimizer/test_qgt_itersolve.py
+++ b/test/optimizer/test_qgt_itersolve.py
@@ -125,6 +125,16 @@ def vstate(request, model, chunk_size):
     return vstate
 
 
+def is_complex_failing(vstate, qgt_partial):
+    """
+    returns true if this qgt should error on construction
+    """
+    if not jnp.issubdtype(vstate.model.param_dtype, jnp.complexfloating):
+        if qgt_partial.keywords.get("holomorphic", False):
+            return True
+    return False
+
+
 @pytest.mark.parametrize(
     "qgt",
     [pytest.param(sr, id=name) for name, sr in QGT_objects.items()],
@@ -135,7 +145,13 @@ def vstate(request, model, chunk_size):
 )
 @pytest.mark.parametrize("chunk_size", [None, 16])
 def test_qgt_solve(qgt, vstate, solver, _mpi_size, _mpi_rank):
-    S = qgt(vstate)
+    if is_complex_failing(vstate, qgt):
+        with pytest.raises(ValueError):
+            S = qgt(vstate)
+        return
+    else:
+        S = qgt(vstate)
+
     x, _ = S.solve(solver, vstate.parameters)
 
     rtol, atol = solvers_tol[solver, nk.jax.dtype_real(vstate.model.dtype)]
@@ -173,6 +189,9 @@ def test_qgt_solve(qgt, vstate, solver, _mpi_size, _mpi_rank):
 )
 @pytest.mark.parametrize("chunk_size", [None])
 def test_qgt_solve_with_x0(qgt, vstate):
+    if is_complex_failing(vstate, qgt):
+        return
+
     solver = jax.scipy.sparse.linalg.gmres
     x0 = jax.tree_map(jnp.zeros_like, vstate.parameters)
 
@@ -186,6 +205,8 @@ def test_qgt_solve_with_x0(qgt, vstate):
 )
 @pytest.mark.parametrize("chunk_size", [None, 16])
 def test_qgt_matmul(qgt, vstate, _mpi_size, _mpi_rank):
+    if is_complex_failing(vstate, qgt):
+        return
 
     rtol, atol = matmul_tol[nk.jax.dtype_real(vstate.model.dtype)]
 
@@ -240,6 +261,8 @@ def check_same_dtype(x, y):
 )
 @pytest.mark.parametrize("chunk_size", [None, 16])
 def test_qgt_dense(qgt, vstate, _mpi_size, _mpi_rank):
+    if is_complex_failing(vstate, qgt):
+        return
 
     rtol, atol = dense_tol[nk.jax.dtype_real(vstate.model.dtype)]
 
@@ -289,6 +312,9 @@ def test_qgt_dense(qgt, vstate, _mpi_size, _mpi_rank):
     ],
 )
 def test_qgt_pytree_diag_shift(qgt, vstate):
+    if is_complex_failing(vstate, qgt):
+        return
+
     v = vstate.parameters
     S = qgt(vstate)
     expected = S @ v
@@ -304,3 +330,19 @@ def test_qgt_pytree_diag_shift(qgt, vstate):
     S = S.replace(diag_shift=diag_shift_tree)
     res = S @ v
     jax.tree_map(lambda a, b: np.testing.assert_allclose(a, b), res, expected)
+
+
+@pytest.mark.skipif_mpi
+def test_qgt_holomorphic_real_pars_throws():
+    hi = nk.hilbert.Spin(1 / 2, 5)
+    vstate = nk.vqs.MCState(
+        nk.sampler.MetropolisLocal(hi),
+        nk.models.RBM(param_dtype=float),
+    )
+
+    with pytest.raises(ValueError):
+        vstate.quantum_geometric_tensor(qgt.QGTJacobianPyTree(holomorphic=True))
+    with pytest.raises(ValueError):
+        vstate.quantum_geometric_tensor(qgt.QGTJacobianDense(holomorphic=True))
+
+    return vstate