NetKet 3.14.0__tar.gz → 3.14.2__tar.gz

{netket-3.14.0 → netket-3.14.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: NetKet
-Version: 3.14.0
+Version: 3.14.2
 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.14.0 → netket-3.14.2}/netket/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.14.0'
-__version_tuple__ = version_tuple = (3, 14, 0)
+__version__ = version = '3.14.2'
+__version_tuple__ = version_tuple = (3, 14, 2)

{netket-3.14.0 → netket-3.14.2}/netket/experimental/dynamics/_rk_solver_structures.py

@@ -22,7 +22,7 @@ import jax.numpy as jnp
 import netket as nk
 from netket import config
 from netket.utils.mpi.primitives import mpi_all_jax
-from netket.utils.struct import dataclass, field
+from netket.utils import struct, KahanSum
 from netket.utils.types import Array, PyTree
 from netket.utils.numbers import dtype as _dtype
 
@@ -124,13 +124,13 @@ def maximum_norm(x: PyTree | Array):
         )
 
 
-@dataclass
-class RungeKuttaState:
+@struct.dataclass
+class RungeKuttaState(struct.Pytree):
     step_no: int
     """Number of successful steps since the start of the iteration."""
     step_no_total: int
     """Number of steps since the start of the iteration, including rejected steps."""
-    t: nk.utils.KahanSum
+    t: KahanSum
     """Current time."""
     y: Array
     """Solution at current time."""
@@ -143,6 +143,53 @@ class RungeKuttaState:
     flags: SolverFlags = SolverFlags.INFO_STEP_ACCEPTED
     """Flags containing information on the solver state."""
 
+    def __init__(
+        self,
+        y,
+        t: float,
+        dt: float,
+        *,
+        step_no=0,
+        step_no_total=0,
+        last_norm=None,
+        last_scaled_error=None,
+        flags=SolverFlags.INFO_STEP_ACCEPTED,
+    ):
+        step_dtype = jnp.int64 if jax.config.x64_enabled else jnp.int32
+        err_dtype = jnp.float64 if jax.config.x64_enabled else jnp.float32
+
+        self.step_no = jnp.asarray(step_no, dtype=step_dtype)
+        self.step_no_total = jnp.asarray(step_no_total, dtype=step_dtype)
+
+        if not isinstance(t, KahanSum):
+            t = KahanSum(t)
+        if not isinstance(dt, jax.Array):
+            dt = jnp.asarray(dt)
+        self.t = t
+        self.dt = dt
+
+        # To avoid creating problems, just convert to array if the leaves
+        # in the parameters are not jax arrays already.
+        def _maybe_asarray(x):
+            if isinstance(x, jax.Array):
+                return x
+            else:
+                return jnp.asarray(x)
+
+        self.y = jax.tree_util.tree_map(_maybe_asarray, y)
+
+        if last_norm is not None:
+            last_norm = jnp.asarray(last_norm, dtype=err_dtype)
+        self.last_norm = last_norm
+        if last_scaled_error is not None:
+            last_scaled_error = jnp.asarray(last_scaled_error, dtype=err_dtype)
+        self.last_scaled_error = last_scaled_error
+
+        # Todo: if making SolverFlag a proper pytree this can be restored.
+        # if not isinstance(flags, SolverFlags):
+        #    raise TypeError(f"flags must be SolverFlags but got {type(flags)} : {flag}")
+        self.flags = flags
+
     def __repr__(self):
         try:
             dt = "{self.dt:.2e}"
@@ -263,8 +310,8 @@ def general_time_step_adaptive(
                 lambda _: rk_state.dt,
                 None,
             ),
-            last_norm=norm_y,
-            last_scaled_error=scaled_err,
+            last_norm=norm_y.astype(rk_state.last_norm.dtype),
+            last_scaled_error=scaled_err.astype(rk_state.last_scaled_error.dtype),
             flags=flags | SolverFlags.INFO_STEP_ACCEPTED,
         ),
         # step rejected, repeat with lower dt
@@ -300,13 +347,13 @@ def general_time_step_fixed(
     )
 
 
-@dataclass(_frozen=False)
+@struct.dataclass(_frozen=False)
 class RungeKuttaIntegrator:
     tableau: rkt.NamedTableau
 
-    f: Callable = field(repr=False)
+    f: Callable = struct.field(repr=False)
     t0: float
-    y0: Array = field(repr=False)
+    y0: Array = struct.field(repr=False)
 
     initial_dt: float
 
@@ -338,10 +385,8 @@ class RungeKuttaIntegrator:
         setattr(self, "initial_dt", jnp.array(self.initial_dt, dtype=t_dtype))
 
         self._rkstate = RungeKuttaState(
-            step_no=0,
-            step_no_total=0,
-            t=nk.utils.KahanSum(self.t0),
             y=self.y0,
+            t=KahanSum(self.t0),
             dt=self.initial_dt,
             last_norm=0.0 if self.use_adaptive else None,
             last_scaled_error=0.0 if self.use_adaptive else None,

{netket-3.14.0 → netket-3.14.2}/netket/experimental/hilbert/spin_orbital_fermions.py

@@ -168,15 +168,16 @@ class SpinOrbitalFermions(HomogeneousHilbert):
                 if n_fermions is not None:
                     occupation_constraint = SumConstraint(n_fermions)
 
-        # Wrap the extra user provided cosntraint around our
-        # occupation constraint of the populations.
-        if constraint is not None:
-            constraint = ExtraConstraint(
-                base_constraint=occupation_constraint,
-                extra_constraint=constraint,
-            )
-        else:
-            constraint = occupation_constraint
+        if occupation_constraint is not None:
+            # Wrap the extra user provided cosntraint around our
+            # occupation constraint of the populations.
+            if constraint is not None:
+                constraint = ExtraConstraint(
+                    base_constraint=occupation_constraint,
+                    extra_constraint=constraint,
+                )
+            else:
+                constraint = occupation_constraint
 
         """Internal representation of this Hilbert space (Fock or TensorHilbert)."""
         # local states are the occupation numbers (0, 1)

{netket-3.14.0 → netket-3.14.2}/netket/hilbert/constraint/multiple.py

@@ -72,15 +72,11 @@ class ExtraConstraint(DiscreteHilbertConstraint):
         ):
             raise TypeError(
                 dedent(
-                    """The type parameters T1, T2 of ExtraConstraint[T1, T2] must be subtypes of
+                    """The arguments of ExtraConstraint must both be subtypes of
                     nk.hilbert.constraint.DiscreteHilbertConstraint, but one or more are not.
 
                     To verify if those are valid type parameters you can check it with
-                    issubclass(T1, nk.hilbert.constraint.DiscreteHilbertConstraint).
-
-                    For more information, look at the documentation of parametric classes in
-                    https://beartype.github.io/plum/parametric.html or open an issue on NetKet's
-                    GitHub discussion session.
+                    issubclass(arg, nk.hilbert.constraint.DiscreteHilbertConstraint).
                     """
                 )
             )
@@ -99,15 +95,11 @@ class ExtraConstraint(DiscreteHilbertConstraint):
         ):
             raise TypeError(
                 dedent(
-                    """The type parameters T1, T2 of ExtraConstraint[T1, T2] must be subtypes of
+                    """The arguments of ExtraConstraint must both be subtypes of
                     nk.hilbert.constraint.DiscreteHilbertConstraint, but one or more are not.
 
-                    To verify if those are valid type parameters you can check it with
-                    issubclass(T1, nk.hilbert.constraint.DiscreteHilbertConstraint).
-
-                    For more information, look at the documentation of parametric classes in
-                    https://beartype.github.io/plum/parametric.html or open an issue on NetKet's
-                    GitHub discussion session.
+                    To verify if those are valid arguments you can check it with
+                    isinstance(arg, nk.hilbert.constraint.DiscreteHilbertConstraint).
                     """
                 )
             )

{netket-3.14.0 → netket-3.14.2}/netket/utils/summation.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .struct import dataclass
-from .types import Scalar
+import jax.numpy as jnp
 
+from netket.utils import struct
+from netket.utils.types import Scalar
 
-@dataclass
-class KahanSum:
+
+@struct.dataclass
+class KahanSum(struct.Pytree):
     """
     Accumulator implementing Kahan summation [1], which reduces
     the effect of accumulated floating-point error.
@@ -26,10 +28,33 @@ class KahanSum:
     """
 
     value: Scalar
-    compensator: Scalar = 0.0
+    """
+    Current value stored in this array
+    """
+    compensator: Scalar
+    """
+    Compensator used to fix addition/subtraction
+    """
+
+    def __init__(self, value, compensator=None) -> None:
+        """
+        Constructs the Kahan Summation accumulator with a zero-initialized
+        compensator if unspecified.
+
+        Args:
+            value: The value to initialize this scalar.
+            compensator: The state of the compensator. 0 by default.
+        """
+        self.value = jnp.asarray(value)
+        if compensator is None:
+            compensator = jnp.zeros_like(self.value)
+        self.compensator = compensator
 
     def __add__(self, other: Scalar):
         delta = other - self.compensator
         new_value = self.value + delta
         new_compensator = (new_value - self.value) - delta
         return KahanSum(new_value, new_compensator)
+
+    def __jax_array__(self, dtype=None):
+        return self.value.astype(dtype)

{netket-3.14.0 → netket-3.14.2}/netket/vqs/full_summ/state.py

@@ -185,7 +185,7 @@ class FullSumState(VariationalState):
 
         key = nkjax.PRNGKey(seed)
 
-        dummy_input = jnp.zeros((1, self.hilbert.size), dtype=dtype)
+        dummy_input = self.hilbert.random_state(key, 1, dtype=dtype)
 
         variables = jit_evaluate(self._init_fun, {"params": key}, dummy_input)
         self.variables = variables

{netket-3.14.0 → netket-3.14.2}/netket/vqs/mc/mc_state/state.py

@@ -274,7 +274,7 @@ class MCState(VariationalState):
 
         key = nkjax.PRNGKey(seed)
 
-        dummy_input = jnp.zeros((1, self.hilbert.size), dtype=dtype)
+        dummy_input = self.hilbert.random_state(key, 1, dtype=dtype)
 
         variables = jit_evaluate(self._init_fun, {"params": key}, dummy_input)
         self.variables = variables

{netket-3.14.0 → netket-3.14.2}/pyproject.toml

@@ -101,6 +101,12 @@ source = "vcs"
 [tool.hatch.build.hooks.vcs]
 version-file = "netket/_version.py"
 
+[tool.uv]
+dev-dependencies = [
+  "netket[dev]",
+  "mpi4py",
+  "jax>=0.4.33"
+]
 
 # Development tools
 
@@ -145,7 +151,9 @@ filterwarnings = [
     "ignore:Defaulting to `holomorphic=False`,",
     "ignore::PendingDeprecationWarning: the matrix subclass",
     "ignore::netket.errors.HolomorphicUndeclaredWarning",
-    "ignore:netket.hilbert.CustomHilbert is deprecated"
+    "ignore:netket.hilbert.CustomHilbert is deprecated",
+    "ignore::netket.errors.UndeclaredSpinOderingWarning",
+    "ignore::netket.errors.UnoptimalSRtWarning",
 ]
 testpaths = [
     "test",