simplify logic of logpsi_U wrapper (#11)

@alleSini99

netket_fidelity/infidelity/overlap/operator.py

@@ -2,15 +2,13 @@
 
 import jax.numpy as jnp
 
-import flax
 
-from netket import jax as nkjax
 from netket.operator import AbstractOperator, DiscreteJaxOperator
 from netket.utils.types import DType
 from netket.utils.numbers import is_scalar
 from netket.vqs import VariationalState, MCState, FullSumState
 
-from netket_fidelity.utils.sampling_Ustate import _logpsi_U
+from netket_fidelity.utils.sampling_Ustate import make_logpsi_U_afun
 
 
 class InfidelityOperatorStandard(AbstractOperator):
@@ -72,12 +70,12 @@ def InfidelityUPsi(
  "an instance of DiscreteJaxOperator."
  )
 
- logpsiU = nkjax.HashablePartial(_logpsi_U, state._apply_fun)
+ logpsiU, variables_U = make_logpsi_U_afun(state._apply_fun, U, state.variables)
  target = MCState(
  sampler=state.sampler,
  apply_fun=logpsiU,
  n_samples=state.n_samples,
- variables=flax.core.copy(state.variables, {"unitary": U}),
+ variables=variables_U,
  )
 
  return InfidelityOperatorStandard(target, cv_coeff=cv_coeff, dtype=dtype)

netket_fidelity/utils/__init__.py

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

netket_fidelity/utils/sampling_Ustate.py

@@ -1,9 +1,37 @@
 import jax
-
 import flax
 
+from netket import jax as nkjax
+
+
+def make_logpsi_U_afun(logpsi_fun, U, variables):
+ """Wraps an apply_fun into another one that multiplies it by an
+ Unitary transformation U.
+
+ This wrapper is made such that the Unitary is passed as the model_state
+ of the new wrapped function, and therefore changes to the angles/coefficients
+ of the Unitary should not trigger recompilation.
+
+ Args:
+ logpsi_fun: a function that takes as input variables and samples
+ U: a {class}`nk.operator.JaxDiscreteOperator`
+ variables: The variables used to call *logpsi_fun*
+
+ Returns:
+ A tuple, where the first element is a new function with the same signature as
+ the original **logpsi_fun** and a set of new variables to be used to call it.
+ """
+ # wrap apply_fun into logpsi logpsi_U
+ logpsiU_fun = nkjax.HashablePartial(_logpsi_U_fun, logpsi_fun)
+
+ # Insert a new 'model_state' key to store the Unitary. This only works
+ # if U is a pytree that can be flattened/unflattened.
+ new_variables = flax.core.copy(variables, {"unitary": U})
+
+ return logpsiU_fun, new_variables
+
 
-def _logpsi_U(apply_fun, variables, x, *args):
+def _logpsi_U_fun(apply_fun, variables, x, *args):
  """
  This should be used as a wrapper to the original apply function, adding
  to the `variables` dictionary (in model_state) a new key `unitary` with