NetKet 3.22.1__tar.gz → 3.22.3__tar.gz

{netket-3.22.1 → netket-3.22.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.22.1
+Version: 3.22.3
 Summary: Netket : Machine Learning toolbox for many-body quantum systems.
 Project-URL: homepage, https://www.netket.org
 Project-URL: documentation, https://netket.readthedocs.io/en/latest/#

{netket-3.22.1 → netket-3.22.3}/netket/_src/observable/vscore.py

@@ -20,6 +20,7 @@ import jax
 import jax.numpy as jnp
 
 from netket._src.stats.local_estimators import LocalEstimatorsBatch
+from netket.jax import HashablePartial
 from netket.operator import AbstractOperator
 from netket.operator._abstract_observable import AbstractObservable
 from netket.stats import Stats
@@ -34,27 +35,61 @@ from netket.vqs import (
 from netket.vqs.mc.common import local_estimators
 
 
-def _vscore_combinator(trace_diagonal: float):
-    def combinator(mu):
-        return (mu[1] - mu[0] ** 2) / (mu[0] - trace_diagonal) ** 2
+def _vscore_combinator(trace_diagonal: float, N: int, mu):
+    return N * (mu[1] - mu[0] ** 2) / (mu[0] - trace_diagonal) ** 2
 
-    return combinator
+
+def _vscore_n_dof(hilbert) -> int:
+    r"""Estimate the V-score normalisation :math:`N` (number of degrees of
+    freedom) for a given Hilbert space, following arXiv:2302.04919.
+
+    For spin, qubit and bosonic lattice models this is the number of sites/modes
+    (``hilbert.size``). For fermionic systems the V-score is normalised by the
+    total particle number, so we use ``SpinOrbitalFermions.n_fermions`` instead
+    of ``hilbert.size`` (which counts spin-orbitals).
+    """
+    from netket.hilbert import SpinOrbitalFermions
+
+    if isinstance(hilbert, SpinOrbitalFermions):
+        if hilbert.n_fermions is None:
+            raise ValueError(
+                "Cannot infer the V-score normalisation `N` for a "
+                "SpinOrbitalFermions space with unfixed particle number. "
+                "Pass `N=...` explicitly to VScore: the V-score is normalised by "
+                "the total particle number (see arXiv:2302.04919)."
+            )
+        return hilbert.n_fermions
+    return hilbert.size
 
 
 class VScore(AbstractObservable):
     r"""
-    Observable computing the V-score of a quantum operator :math:`H`:
+    Observable computing the V-score of a quantum operator :math:`H`, as defined
+    in `Wu et al., arXiv:2302.04919 <https://arxiv.org/abs/2302.04919>`_:
 
     .. math::
 
         V_{\mathrm{score}} =
-        \frac{\mathrm{Var}(H)}{(\langle H \rangle - \mathrm{trace\_diagonal})^2}
+        \frac{N\,\mathrm{Var}(H)}{(\langle H \rangle - E_\infty)^2}
         =
-        \frac{\langle H^2 \rangle - \langle H \rangle^2}
-        {(\langle H \rangle - \mathrm{trace\_diagonal})^2}.
+        \frac{N\,(\langle H^2 \rangle - \langle H \rangle^2)}
+        {(\langle H \rangle - E_\infty)^2}.
+
+    The ``trace_diagonal`` parameter is the infinite-temperature reference energy
+    :math:`E_\infty := \mathrm{Tr}(H) / \dim(\mathcal{H})` entering the denominator
+    (it is :math:`0` for traceless Hamiltonians, e.g. spin models built from Pauli
+    operators).
 
-    The ``trace_diagonal`` parameter is a mandatory diagonal-energy shift used in
-    the denominator.
+    The prefactor :math:`N` is the number of degrees of freedom. By default it is
+    inferred from the Hilbert space:
+
+    * for spin, qubit and bosonic lattice models, ``N = hilbert.size`` (number of
+      sites/modes);
+    * for fermionic systems (:class:`~netket.hilbert.SpinOrbitalFermions`),
+      ``N`` is the total particle number ``hilbert.n_fermions``.
+
+    For composite or exotic Hilbert spaces, or to use a different convention, pass
+    ``N`` explicitly.
     """
 
     def __init__(
@@ -62,6 +97,7 @@ class VScore(AbstractObservable):
         operator: AbstractOperator,
         *,
         trace_diagonal: float,
+        N: int | None = None,
     ):
         super().__init__(operator.hilbert)
 
@@ -72,6 +108,7 @@ class VScore(AbstractObservable):
         self._operator = operator
         self._operator_squared = operator @ operator
         self._trace_diagonal = float(trace_diagonal)
+        self._N = int(N) if N is not None else _vscore_n_dof(operator.hilbert)
 
     @property
     def operator(self) -> AbstractOperator:
@@ -85,8 +122,16 @@ class VScore(AbstractObservable):
     def trace_diagonal(self) -> float:
         return self._trace_diagonal
 
+    @property
+    def N(self) -> int:
+        """The number of degrees of freedom normalising the V-score."""
+        return self._N
+
     def __repr__(self):
-        return f"VScore(op={self.operator}, trace_diagonal={self.trace_diagonal})"
+        return (
+            f"VScore(op={self.operator}, trace_diagonal={self.trace_diagonal}, "
+            f"N={self.N})"
+        )
 
 
 @local_estimators.dispatch
@@ -119,7 +164,9 @@ def vscore_local_estimators(
     data = jnp.stack([O_loc, O2_loc], axis=-1).reshape(n_chains, -1, 2)
     return LocalEstimatorsBatch(
         data=data,
-        combinator=_vscore_combinator(vscore_op.trace_diagonal),
+        combinator=HashablePartial(
+            _vscore_combinator, vscore_op.trace_diagonal, vscore_op.N
+        ),
     )
 
 
@@ -144,6 +191,7 @@ def expect(vstate: FullSumState, vscore_op: VScore):  # noqa: F811
         operator_mtrx,
         operator_squared_mtrx,
         vscore_op.trace_diagonal,
+        vscore_op.N,
     )
 
 
@@ -156,6 +204,7 @@ def _expect_vscore_fs(
     operator_mtrx,
     operator_squared_mtrx,
     trace_diagonal,
+    N,
 ):
     W = {"params": params, **model_state}
 
@@ -164,6 +213,6 @@ def _expect_vscore_fs(
 
     E_mean = (state.conj() @ (operator_mtrx @ state)).real
     E2_mean = (state.conj() @ (operator_squared_mtrx @ state)).real
-    vscore = (E2_mean - E_mean**2) / (E_mean - trace_diagonal) ** 2
+    vscore = N * (E2_mean - E_mean**2) / (E_mean - trace_diagonal) ** 2
 
     return Stats(mean=vscore, error_of_mean=0.0, variance=0.0)

{netket-3.22.1 → netket-3.22.3}/netket/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '3.22.1'
-__version_tuple__ = version_tuple = (3, 22, 1)
+__version__ = version = '3.22.3'
+__version_tuple__ = version_tuple = (3, 22, 3)
 
 __commit_id__ = commit_id = None

{netket-3.22.1 → netket-3.22.3}/netket/utils/_tips.py

@@ -18,6 +18,8 @@ import random
 import sys
 
 TIPS = [
+    # News
+    "NetKet 3.22 has plenty of new stuff! Check it out at https://www.netket.org/posts/07-netket322/",
     # Drivers
     "Prefer the new nk.driver.VMC_SR over VMC which supports minSR and SPRING.",
     "Use driver.run(..., timeit=True) to know where your dominant cost is.",

{netket-3.22.1 → netket-3.22.3}/netket/utils/model_frameworks/base.py

@@ -31,6 +31,23 @@ class ModuleFramework(abc.ABC):
         """
         return False
 
+    @staticmethod
+    def wrapped_model_class(static_model: Any) -> type:
+        """
+        Returns the concrete user-defined model class, given the (possibly wrapped)
+        static module stored in :attr:`MCState._model`.
+
+        For flax-like frameworks the static module *is* the model, so this returns
+        its type. Frameworks that wrap the model into a static container (such as
+        ``flax.nnx`` or ``equinox``) override this to peek inside the wrapper and
+        return the original class.
+
+        This lets serialization code recover the original model class without
+        unwrapping (which would require the variables), e.g. to dispatch on whether
+        a custom serializer is registered for that class.
+        """
+        return type(static_model)
+
     @staticmethod
     @abc.abstractmethod
     def is_loaded() -> bool:

{netket-3.22.1 → netket-3.22.3}/netket/utils/model_frameworks/equinox.py

@@ -56,7 +56,14 @@ class EquinoxWrapper:
         module = self.recompose(variables)
         fun = getattr(module, method)
 
-        return fun(*args, key=rngs, **kwargs)
+        # Only forward a `key` to the model if one was actually provided.
+        # Passing `key=None` unconditionally breaks Equinox modules whose
+        # `__call__` does not declare a `key` argument (e.g. simple models
+        # that do not use randomness).
+        if rngs is not None:
+            kwargs["key"] = rngs
+
+        return fun(*args, **kwargs)
 
     def recompose(self, variables):
         import equinox as eqx
@@ -94,6 +101,11 @@ class EquinoxFramework(ModuleFramework):
         """
         return True
 
+    @staticmethod
+    def wrapped_model_class(static_model: "EquinoxWrapper") -> type:
+        # The static partition is an instance of the original equinox.Module class.
+        return type(static_model.static_module)
+
     @staticmethod
     def is_loaded() -> bool:
         return "equinox" in sys.modules

{netket-3.22.1 → netket-3.22.3}/netket/utils/model_frameworks/nnx.py

@@ -107,6 +107,11 @@ class NNXFramework(ModuleFramework):
         """
         return True
 
+    @staticmethod
+    def wrapped_model_class(static_model: "NNXWrapper") -> type:
+        # The original nnx.Module class is recorded in the graphdef.
+        return _get_graphdef_type(static_model.graphdef)
+
     @staticmethod
     def is_loaded() -> bool:
         # this should be not necessary, as netket requires and loads