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