Repository: netket/netket

Issue:
How to use NumpySampler with VMC.estimate (or MCState.expect)?
Hello!
Trying the `VMC.estimate` with the `MetropolisSamplerNumpy` i got an error:
```python
  File "main.py", line 85, in <module>
    res = model.get_results()
  File "/home/sem/GitHub/XXModelNQS/project/src/model/nqs.py", line 95, in get_results
    [
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 295, in estimate
    return tree_map(self._estimate_stats, observables)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 109, in _estimate_stats
    return self.state.expect(observable)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/variational/mc_state.py", line 407, in expect
    mels,
jax._src.traceback_util.FilteredStackTrace: TypeError: Argument 'MetropolisSamplerNumpy(rule = LocalRuleNumpy(), n_chains = 64, machine_power = 2, n_sweeps = 8, dtype = <class 'numpy.float32'>)' of type <class 'netket.sampler.metropolis_numpy.MetropolisSamplerNumpy'> is not a valid JAX type.

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "main.py", line 85, in <module>
    res = model.get_results()
  File "/home/sem/GitHub/XXModelNQS/project/src/model/nqs.py", line 95, in get_results
    [
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 295, in estimate
    return tree_map(self._estimate_stats, observables)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/_src/tree_util.py", line 162, in tree_map
    return treedef.unflatten(map(f, leaves))
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 109, in _estimate_stats
    return self.state.expect(observable)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/variational/mc_state.py", line 407, in expect
    mels,
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/api.py", line 330, in cache_miss
    _check_arg(arg)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/api.py", line 2209, in _check_arg
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")
TypeError: Argument 'MetropolisSamplerNumpy(rule = LocalRuleNumpy(), n_chains = 64, machine_power = 2, n_sweeps = 8, dtype = <class 'numpy.float32'>)' of type <class 'netket.sampler.metropolis_numpy.MetropolisSamplerNumpy'> is not a valid JAX type.
```

Could you please help what is the right way to use `estimate` (or `MCState.expect`) methods with Numpy samplers?
NetKet version:
`3.0b1.post5`

P.S. Error only in `estimate`. `VMC.run` works correctly!
P.P.S. With `JAX` samplers it works correctly too but results are strange and significantly worse than results from `NetKet` 2.1 and same parameters...
How to use NumpySampler with VMC.estimate (or MCState.expect)?
Hello!
Trying the `VMC.estimate` with the `MetropolisSamplerNumpy` i got an error:
```python
  File "main.py", line 85, in <module>
    res = model.get_results()
  File "/home/sem/GitHub/XXModelNQS/project/src/model/nqs.py", line 95, in get_results
    [
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 295, in estimate
    return tree_map(self._estimate_stats, observables)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 109, in _estimate_stats
    return self.state.expect(observable)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/variational/mc_state.py", line 407, in expect
    mels,
jax._src.traceback_util.FilteredStackTrace: TypeError: Argument 'MetropolisSamplerNumpy(rule = LocalRuleNumpy(), n_chains = 64, machine_power = 2, n_sweeps = 8, dtype = <class 'numpy.float32'>)' of type <class 'netket.sampler.metropolis_numpy.MetropolisSamplerNumpy'> is not a valid JAX type.

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "main.py", line 85, in <module>
    res = model.get_results()
  File "/home/sem/GitHub/XXModelNQS/project/src/model/nqs.py", line 95, in get_results
    [
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 295, in estimate
    return tree_map(self._estimate_stats, observables)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/_src/tree_util.py", line 162, in tree_map
    return treedef.unflatten(map(f, leaves))
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 109, in _estimate_stats
    return self.state.expect(observable)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/variational/mc_state.py", line 407, in expect
    mels,
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/api.py", line 330, in cache_miss
    _check_arg(arg)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/api.py", line 2209, in _check_arg
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")
TypeError: Argument 'MetropolisSamplerNumpy(rule = LocalRuleNumpy(), n_chains = 64, machine_power = 2, n_sweeps = 8, dtype = <class 'numpy.float32'>)' of type <class 'netket.sampler.metropolis_numpy.MetropolisSamplerNumpy'> is not a valid JAX type.
```

Could you please help what is the right way to use `estimate` (or `MCState.expect`) methods with Numpy samplers?
NetKet version:
`3.0b1.post5`

P.S. Error only in `estimate`. `VMC.run` works correctly!
P.P.S. With `JAX` samplers it works correctly too but results are strange and significantly worse than results from `NetKet` 2.1 and same parameters...
How to use NumpySampler with VMC.estimate (or MCState.expect)?
Hello!
Trying the `VMC.estimate` with the `MetropolisSamplerNumpy` i got an error:
```python
  File "main.py", line 85, in <module>
    res = model.get_results()
  File "/home/sem/GitHub/XXModelNQS/project/src/model/nqs.py", line 95, in get_results
    [
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 295, in estimate
    return tree_map(self._estimate_stats, observables)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 109, in _estimate_stats
    return self.state.expect(observable)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/variational/mc_state.py", line 407, in expect
    mels,
jax._src.traceback_util.FilteredStackTrace: TypeError: Argument 'MetropolisSamplerNumpy(rule = LocalRuleNumpy(), n_chains = 64, machine_power = 2, n_sweeps = 8, dtype = <class 'numpy.float32'>)' of type <class 'netket.sampler.metropolis_numpy.MetropolisSamplerNumpy'> is not a valid JAX type.

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "main.py", line 85, in <module>
    res = model.get_results()
  File "/home/sem/GitHub/XXModelNQS/project/src/model/nqs.py", line 95, in get_results
    [
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 295, in estimate
    return tree_map(self._estimate_stats, observables)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/_src/tree_util.py", line 162, in tree_map
    return treedef.unflatten(map(f, leaves))
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/driver/abstract_variational_driver.py", line 109, in _estimate_stats
    return self.state.expect(observable)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/netket/variational/mc_state.py", line 407, in expect
    mels,
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/api.py", line 330, in cache_miss
    _check_arg(arg)
  File "/home/sem/.cache/pypoetry/virtualenvs/xxmodelnqs--vHqlMQ2-py3.7/lib/python3.7/site-packages/jax/api.py", line 2209, in _check_arg
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")
TypeError: Argument 'MetropolisSamplerNumpy(rule = LocalRuleNumpy(), n_chains = 64, machine_power = 2, n_sweeps = 8, dtype = <class 'numpy.float32'>)' of type <class 'netket.sampler.metropolis_numpy.MetropolisSamplerNumpy'> is not a valid JAX type.
```

Could you please help what is the right way to use `estimate` (or `MCState.expect`) methods with Numpy samplers?
NetKet version:
`3.0b1.post5`

P.S. Error only in `estimate`. `VMC.run` works correctly!
P.P.S. With `JAX` samplers it works correctly too but results are strange and significantly worse than results from `NetKet` 2.1 and same parameters...


Discussion:

Resolving patch:
diff --git a/netket/operator/_local_liouvillian.py b/netket/operator/_local_liouvillian.py
--- a/netket/operator/_local_liouvillian.py
+++ b/netket/operator/_local_liouvillian.py
@@ -103,6 +103,11 @@ def jump_operators(self) -> PyList[LocalOperator]:
         """The list of local operators in this Liouvillian"""
         return self._jump_ops
 
+    @property
+    def max_conn_size(self) -> int:
+        """The maximum number of non zero ⟨x|O|x'⟩ for every x."""
+        return self._max_conn_size
+
     def _compute_hnh(self):
         # There is no i here because it's inserted in the kernel
         Hnh = 1.0 * self._H
@@ -113,7 +118,7 @@ def _compute_hnh(self):
 
         self._Hnh = Hnh.collect()
 
-        max_conn_size = self._max_dissipator_conn_size + Hnh.max_conn_size
+        max_conn_size = self._max_dissipator_conn_size + 2 * Hnh.max_conn_size
         self._max_conn_size = max_conn_size
         self._xprime = np.empty((max_conn_size, self.hilbert.size))
         self._xr_prime = np.empty((max_conn_size, self.hilbert.physical.size))
diff --git a/netket/variational/mc_state.py b/netket/variational/mc_state.py
--- a/netket/variational/mc_state.py
+++ b/netket/variational/mc_state.py
@@ -368,8 +368,13 @@ def samples(self) -> jnp.ndarray:
 
     def log_value(self, σ: jnp.ndarray) -> jnp.ndarray:
         """
-        Evaluate the variational state for a batch of states and return
-        the logarithm of the probability amplitude for those inputs.
+        Evaluate the variational state for a batch of states and returns
+        the logarithm of the amplitude of the quantum state. For pure states,
+        this is :math:`log(<σ|ψ>)`, whereas for mixed states this is
+        :math:`log(<σr|ρ|σc>)`, where ψ and ρ are respectively a pure state
+        (wavefunction) and a mixed state (density matrix).
+        For the density matrix, the left and right-acting states (row and column)
+        are obtained as :code:`σr=σ[::,0:N]` and :code:`σc=σ[::,N:]`.
 
         Given a batch of inputs (Nb, N), returns a batch of outputs (Nb,).
         """
@@ -397,7 +402,7 @@ def expect(self, Ô: AbstractOperator) -> Stats:
         σp, mels = Ô.get_conn_padded(np.asarray(σ).reshape((-1, σ.shape[-1])))
 
         return _expect(
-            self.sampler,
+            self.sampler.machine_pow,
             self._apply_fun,
             kernel,
             self.parameters,
@@ -439,7 +444,6 @@ def expect_and_grad(
             if squared_operator:
                 if isinstance(Ô, AbstractSuperOperator):
                     Ō, Ō_grad, new_model_state = grad_expect_operator_Lrho2(
-                        self.sampler,
                         self._apply_fun,
                         mutable,
                         self.parameters,
@@ -450,7 +454,7 @@ def expect_and_grad(
                     )
                 else:
                     Ō, Ō_grad, new_model_state = grad_expect_operator_kernel(
-                        self.sampler,
+                        self.sampler.machine_pow,
                         self._apply_fun,
                         local_value_squared_kernel,
                         mutable,
@@ -473,7 +477,7 @@ def expect_and_grad(
                 )
         else:
             Ō, Ō_grad, new_model_state = grad_expect_operator_kernel(
-                self.sampler,
+                self.sampler.machine_pow,
                 self._apply_fun,
                 local_value_kernel,
                 mutable,
@@ -534,7 +538,7 @@ def __str__(self):
 
 @partial(jax.jit, static_argnums=(1, 2))
 def _expect(
-    sampler: Sampler,
+    machine_pow: int,
     model_apply_fun: Callable,
     local_value_kernel: Callable,
     parameters: PyTree,
@@ -551,8 +555,7 @@ def _expect(
     logpsi = lambda w, σ: model_apply_fun({"params": w, **model_state}, σ)
 
     log_pdf = (
-        lambda w, σ: sampler.machine_pow
-        * model_apply_fun({"params": w, **model_state}, σ).real
+        lambda w, σ: machine_pow * model_apply_fun({"params": w, **model_state}, σ).real
     )
 
     local_value_vmap = jax.vmap(
@@ -630,7 +633,7 @@ def grad_expect_hermitian(
 
 @partial(jax.jit, static_argnums=(1, 2, 3))
 def grad_expect_operator_kernel(
-    sampler: Sampler,
+    machine_pow: int,
     model_apply_fun: Callable,
     local_kernel: Callable,
     mutable: bool,
@@ -672,8 +675,7 @@ def grad_expect_operator_kernel(
         )[0]
 
     log_pdf = (
-        lambda w, σ: sampler.machine_pow
-        * model_apply_fun({"params": w, **model_state}, σ).real
+        lambda w, σ: machine_pow * model_apply_fun({"params": w, **model_state}, σ).real
     )
 
     def expect_closure(*args):
@@ -696,9 +698,8 @@ def expect_closure_pars(pars):
     )
 
 
-@partial(jax.jit, static_argnums=(1, 2))
+@partial(jax.jit, static_argnums=(0, 1))
 def grad_expect_operator_Lrho2(
-    sampler: Sampler,
     model_apply_fun: Callable,
     mutable: bool,
     parameters: PyTree,


Test patch:
diff --git a/Test/Variational/test_variational.py b/Test/Variational/test_variational.py
--- a/Test/Variational/test_variational.py
+++ b/Test/Variational/test_variational.py
@@ -151,6 +151,23 @@ def _f(x, y):
     jax.tree_multimap(_f, pars, pars2)
 
 
+@pytest.mark.parametrize(
+    "operator",
+    [
+        pytest.param(
+            op,
+            id=name,
+        )
+        for name, op in operators.items()
+    ],
+)
+def test_expect_numpysampler_works(vstate, operator):
+    sampl = nk.sampler.MetropolisLocalNumpy(vstate.hilbert)
+    vstate.sampler = sampl
+    out = vstate.expect(operator)
+    assert isinstance(out, nk.stats.Stats)
+
+
 @pytest.mark.parametrize(
     "operator",
     [
diff --git a/Test/Variational/test_variational_mixed.py b/Test/Variational/test_variational_mixed.py
--- a/Test/Variational/test_variational_mixed.py
+++ b/Test/Variational/test_variational_mixed.py
@@ -163,3 +163,20 @@ def test_serialization(vstate):
     assert vstate.n_discard == vstate_new.n_discard
     assert vstate.n_samples_diag == vstate_new.n_samples_diag
     assert vstate.n_discard_diag == vstate_new.n_discard_diag
+
+
+@pytest.mark.parametrize(
+    "operator",
+    [
+        pytest.param(
+            op,
+            id=name,
+        )
+        for name, op in operators.items()
+    ],
+)
+def test_expect_numpysampler_works(vstate, operator):
+    sampl = nk.sampler.MetropolisLocalNumpy(vstate.hilbert)
+    vstate.sampler = sampl
+    out = vstate.expect(operator)
+    assert isinstance(out, nk.stats.Stats)