Simplify: remove expect_2distr (#15)

@alleSini99 what do you think? Isn't this much simpler ? (We should
check that it is as fast... but I think it is)

cc @lgravina1997

netket_fidelity/infidelity/overlap/expect.py

@@ -7,7 +7,6 @@
 from netket import jax as nkjax
 from netket.utils import mpi
 
-from netket_fidelity.utils import expect_2distr
 
 from .operator import InfidelityOperatorStandard
 
@@ -76,7 +75,10 @@ def infidelity_sampling_MCState(
     σ_t = sigma_t.reshape(-1, N)
 
     def expect_kernel(params):
-        def kernel_fun(params, params_t, σ, σ_t):
+        def kernel_fun(params_all, samples_all):
+            params, params_t = params_all
+            σ, σ_t = samples_all
+
             W = {"params": params, **model_state}
             W_t = {"params": params_t, **model_state_t}
 
@@ -91,14 +93,24 @@ def kernel_fun(params, params_t, σ, σ_t):
             lambda params, σ: 2 * afun_t({"params": params, **model_state_t}, σ).real
         )
 
-        return expect_2distr(
-            log_pdf,
-            log_pdf_t,
+        def log_pdf_joint(params_all, samples_all):
+            params, params_t = params_all
+            σ, σ_t = samples_all
+            log_pdf_vals = log_pdf(params, σ)
+            log_pdf_t_vals = log_pdf_t(params_t, σ_t)
+            return log_pdf_vals + log_pdf_t_vals
+
+        return nkjax.expect(
+            log_pdf_joint,
             kernel_fun,
-            params,
-            params_t,
-            σ,
-            σ_t,
+            (
+                params,
+                params_t,
+            ),
+            (
+                σ,
+                σ_t,
+            ),
             n_chains=n_chains_t,
         )
 

netket_fidelity/infidelity/overlap_U/exact.py

@@ -82,8 +82,8 @@ def expect_fun(params):
     F, F_vjp_fun = nkjax.vjp(expect_fun, params, conjugate=True)
 
     F_grad = F_vjp_fun(jnp.ones_like(F))[0]
-    F_grad = jax.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], F_grad)
-    I_grad = jax.tree_map(lambda x: -x, F_grad)
+    F_grad = jax.tree_util.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], F_grad)
+    I_grad = jax.tree_util.tree_map(lambda x: -x, F_grad)
     I_stats = Stats(mean=1 - F, error_of_mean=0.0, variance=0.0)
 
     return I_stats, I_grad

netket_fidelity/infidelity/overlap_U/expect.py

@@ -9,7 +9,6 @@
 from netket.vqs import MCState, expect, expect_and_grad, get_local_kernel_arguments
 from netket.utils import mpi
 
-from netket_fidelity.utils import expect_2distr
 
 from .operator import InfidelityOperatorUPsi
 
@@ -113,7 +112,10 @@ def infidelity_sampling_MCState(
     xp_t_ravel = jnp.vstack(xp_t_splitted)
 
     def expect_kernel(params):
-        def kernel_fun(params, params_t, σ, σ_t):
+        def kernel_fun(params_all, samples_all):
+            params, params_t = params_all
+            σ, σ_t = samples_all
+
             W = {"params": params, **model_state}
             W_t = {"params": params_t, **model_state_t}
 
@@ -139,14 +141,24 @@ def kernel_fun(params, params_t, σ, σ_t):
             lambda params, σ: 2 * afun_t({"params": params, **model_state_t}, σ).real
         )
 
-        return expect_2distr(
-            log_pdf,
-            log_pdf_t,
+        def log_pdf_joint(params_all, samples_all):
+            params, params_t = params_all
+            σ, σ_t = samples_all
+            log_pdf_vals = log_pdf(params, σ)
+            log_pdf_t_vals = log_pdf_t(params_t, σ_t)
+            return log_pdf_vals + log_pdf_t_vals
+
+        return nkjax.expect(
+            log_pdf_joint,
             kernel_fun,
-            params,
-            params_t,
-            σ,
-            σ_t,
+            (
+                params,
+                params_t,
+            ),
+            (
+                σ,
+                σ_t,
+            ),
             n_chains=n_chains_t,
         )
 

netket_fidelity/utils/__init__.py

@@ -1,4 +1,3 @@
-from .expect import expect_2distr
 from .sampling_Ustate import make_logpsi_U_afun, _logpsi_U_fun
 
 from netket.utils import _hide_submodules