NetKet 3.12.4__tar.gz → 3.14.0__tar.gz

{netket-3.12.4 → netket-3.14.0}/.gitignore

@@ -16,8 +16,10 @@ netket.egg-info/
 pip-wheel-metadata
 .ipynb_checkpoints
 .pytest_cache
+.benchmarks
 .mypy_cache
 .virtual_documents
+.ruff_cache
 oldest_requirements.txt
 
 # Packages #

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: NetKet
-Version: 3.12.4
+Version: 3.14.0
 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/#
@@ -17,27 +17,23 @@ Classifier: Operating System :: POSIX :: Linux
 Classifier: Operating System :: Unix
 Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Physics
-Requires-Python: >=3.9
-Requires-Dist: beartype>=0.16.2
+Requires-Python: >=3.10
 Requires-Dist: equinox<0.12,>=0.11.3
-Requires-Dist: flax<0.9,>=0.7.1
+Requires-Dist: flax<0.10,>=0.7.1
 Requires-Dist: igraph<0.12.0,>=0.10.0
 Requires-Dist: jax<0.5,>=0.4.16
-Requires-Dist: jaxlib<0.5,>=0.4.16
-Requires-Dist: numba4jax<0.1,>=0.0.10.post1
 Requires-Dist: numba<0.61,>=0.52
-Requires-Dist: numpy~=1.20
+Requires-Dist: numpy<3.0,>=1.20
 Requires-Dist: optax<0.3,>=0.1.3
 Requires-Dist: orjson<4,>=3.4
+Requires-Dist: plum-dispatch<3,>=2.4
 Requires-Dist: rich>=10.0
 Requires-Dist: scipy<2,>=1.5.3
 Requires-Dist: tqdm<5,>=4.60
-Requires-Dist: typing-extensions; python_version <= '3.10'
 Provides-Extra: dev
 Requires-Dist: black==23.9.1; extra == 'dev'
 Requires-Dist: build; extra == 'dev'
 Requires-Dist: coverage>=5; extra == 'dev'
-Requires-Dist: ipython; extra == 'dev'
 Requires-Dist: matplotlib>=3; extra == 'dev'
 Requires-Dist: networkx<4,>=2.4; extra == 'dev'
 Requires-Dist: pre-commit>=2.7; extra == 'dev'
@@ -46,15 +42,15 @@ Requires-Dist: pytest-json-report>=1.3; extra == 'dev'
 Requires-Dist: pytest-xdist[psutil]>=2; extra == 'dev'
 Requires-Dist: pytest>=6; extra == 'dev'
 Requires-Dist: qutip>=5; extra == 'dev'
-Requires-Dist: ruff==0.4.2; extra == 'dev'
+Requires-Dist: ruff==0.6.2; extra == 'dev'
 Requires-Dist: wheel; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: myst-nb<1.2,>=1.0; extra == 'docs'
-Requires-Dist: myst-parser<3.1.0,>=2.0.0; extra == 'docs'
+Requires-Dist: myst-parser<4.1.0,>=2.0.0; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints<3.0,>=1.22; extra == 'docs'
 Requires-Dist: sphinx-book-theme~=1.1.0; extra == 'docs'
 Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
-Requires-Dist: sphinx<7.4,>=7.0; extra == 'docs'
+Requires-Dist: sphinx<8.1,>=7.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
 Provides-Extra: extra
@@ -64,7 +60,7 @@ Requires-Dist: qutip>=4; extra == 'extra'
 Requires-Dist: tensorboardx>=2.0.0; extra == 'extra'
 Provides-Extra: mpi
 Requires-Dist: mpi4jax<0.6,>=0.3.9; extra == 'mpi'
-Requires-Dist: mpi4py<4,>=3.0.1; extra == 'mpi'
+Requires-Dist: mpi4py<5,>=3.0.1; extra == 'mpi'
 Provides-Extra: pyscf
 Requires-Dist: pyscf>=2.0; extra == 'pyscf'
 Requires-Dist: sparse>=0.12.0; extra == 'pyscf'

netket-3.14.0/hatch.toml

@@ -0,0 +1,84 @@
+
+[envs.hatch-test]
+default-args = ["test"]
+features = [
+  "dev",
+  "extra",
+]
+extra-dependencies = [
+]
+randomize = false
+parallel = false
+
+[[envs.hatch-test.matrix]]
+distributed = ['nompi', 'mpi', 'sharding']
+
+[envs.hatch-test.overrides]
+matrix.distributed.extra-dependencies = [
+  { value = "mpi4py", if = ["mpi", "sharding"] },
+  { value = "mpi4jax", if = ["mpi"] },
+]
+matrix.distributed.scripts = [
+    { key = "run", value = "mpirun -np 2 pytest -n 0 {env:HATCH_TEST_ARGS:} {args}", if = ["mpi"] },
+    { key = "run-cov", value = "mpirun -np 2 coverage run -m pytest -n 0 {env:HATCH_TEST_ARGS:} {args}", if = ["mpi"] },
+    { key = "cov-combine", value = "coverage combine --append", if = ["mpi"] },
+
+    { key = "run", value = "MPITRAMPOLINE_LIB=\"{env:VIRTUAL_ENV}/MPIWrapper/build/libmpiwrapper.so\" \"{env:VIRTUAL_ENV}/MPIWrapper/build/bin/mpiwrapperexec\" -np 2 pytest -n 0 --jax-distributed-mpi {env:HATCH_TEST_ARGS:} {args}", if = ["sharding"] },
+]
+matrix.distributed.env-vars = [
+  { key = "NETKET_MPI", value = "0", if = ["nompi", "sharding"] },
+  { key = "NETKET_MPI", value = "1", if = ["mpi"] },
+  { key = "NETKET_EXPERIMENTAL_SHARDING", value = "1", if = ["sharding"] },
+  { key = "NETKET_MPI_WARNING", value = "0", if = ["sharding"] },
+]
+matrix.distributed.parallel = { value = true, if = ["nompi"] }
+matrix.distributed.set-default-args = { value = ["test_sharding/test_sharding_distributed.py"], if = ["sharding"] }
+matrix.distributed.post-install-commands = [
+    { value = "mkdir \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper\"", if = ["sharding"] },
+    { value = "git clone https://github.com/eschnett/MPIwrapper.git \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper\"", if = ["sharding"] },
+    { value = "mkdir \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper/build\"",if = ["sharding"] },
+    { value = "cd \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper/build\" && cmake ..",if = ["sharding"] },
+    { value = "cd \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper/build\" && make",if = ["sharding"] },
+    ]
+
+[envs.docs]
+features = [
+  "docs",
+]
+extra-dependencies = [
+  "sphinx-autobuild"
+]
+[envs.docs.scripts]
+clean = "cd {root}/docs && make clean"
+build = "sphinx-build -M html docs docs/_build"
+serve = "sphinx-autobuild docs docs/_build/html"
+
+
+[envs.sharding]
+features = [
+  "dev",
+  "extra",
+]
+extra-dependencies = [
+  "mpi4py"
+]
+post-install-commands = [
+    "mkdir \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper\"", 
+    "git clone https://github.com/eschnett/MPIwrapper.git \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper\"", 
+    "mkdir \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper/build\"",
+    "cd \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper/build\" && cmake ..",
+    "cd \"{env:VIRTUAL_ENV:DEFAULT}/MPIWrapper/build\" && make",
+    ]
+[envs.sharding.env-vars]
+NETKET_EXPERIMENTAL_SHARDING = "1"
+
+[envs.sharding.scripts]
+run = [
+    "MPITRAMPOLINE_LIB=\"{env:VIRTUAL_ENV}/MPIWrapper/build/libmpiwrapper.so\" \"{env:VIRTUAL_ENV}/MPIWrapper/build/bin/mpiwrapperexec\" -np 2 python {args}"
+]
+run3 = [
+    "MPITRAMPOLINE_LIB=\"{env:VIRTUAL_ENV}/MPIWrapper/build/libmpiwrapper.so\" \"{env:VIRTUAL_ENV}/MPIWrapper/build/bin/mpiwrapperexec\" -np 3 python {args}"
+]
+run4 = [
+    "MPITRAMPOLINE_LIB=\"{env:VIRTUAL_ENV}/MPIWrapper/build/libmpiwrapper.so\" \"{env:VIRTUAL_ENV}/MPIWrapper/build/bin/mpiwrapperexec\" -np 4 python {args}"
+]

{netket-3.12.4 → netket-3.14.0}/netket/_version.py

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

{netket-3.12.4 → netket-3.14.0}/netket/callbacks/earlystopping.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Union
 
 import numpy as np
 
@@ -32,9 +31,9 @@ class EarlyStopping(struct.Pytree, mutable=True):
     This behaves similarly to `min_delta` but is more useful for intensive quantities that
     converge to 0, where absolute tolerances might not be effective.
     """
-    patience: Union[int, float]
+    patience: int | float
     """Number of epochs with no improvement after which training will be stopped."""
-    baseline: Optional[float]
+    baseline: float | None
     """Baseline value for the monitored quantity. Training will stop if the driver hits the baseline."""
     monitor: str
     """Loss statistic to monitor. Should be one of 'mean', 'variance', 'sigma'."""
@@ -53,8 +52,8 @@ class EarlyStopping(struct.Pytree, mutable=True):
         self,
         min_delta: float = 0.0,
         min_reldelta: float = 0.0,
-        patience: Union[int, float] = 0,
-        baseline: Optional[float] = None,
+        patience: int | float = 0,
+        baseline: float | None = None,
         monitor: str = "mean",
     ):
         """

{netket-3.12.4 → netket-3.14.0}/netket/callbacks/invalidlossstopping.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union
 import numpy as np
 
 from netket.utils import struct
@@ -25,14 +24,14 @@ class InvalidLossStopping(struct.Pytree, mutable=True):
 
     monitor: str
     """Loss statistic to monitor. Should be one of 'mean', 'variance', 'sigma'."""
-    patience: Union[int, float]
+    patience: int | float
     """Number of epochs with invalid loss after which training will be stopped."""
 
     # caches
     _last_valid_iter: int
     """Last valid iteration, to check against patience"""
 
-    def __init__(self, monitor: str = "mean", patience: Union[int, float] = 0):
+    def __init__(self, monitor: str = "mean", patience: int | float = 0):
         """
         Construct a callback stopping theoptimisation when the monitored quantity
         becaomes invalid for at least `patience` steps.
@@ -63,10 +62,16 @@ class InvalidLossStopping(struct.Pytree, mutable=True):
         Returns:
             A boolean. If True, training continues, else, it does not.
         """
-        loss = np.real(getattr(log_data[driver._loss_name], self.monitor))
-        if not np.isfinite(loss):
-            if step - self._last_valid_iter >= self.patience:
-                return False
-        else:
-            self._last_valid_iter = step
+        # clears the _last_valid_iter in case the driver was reset
+        if driver.step_count < self._last_valid_iter:
+            self._last_valid_iter = 0
+
+        if driver._loss_stats is not None:
+            loss = np.real(getattr(driver._loss_stats, self.monitor))
+
+            if not np.isfinite(loss):
+                if driver.step_count - self._last_valid_iter >= self.patience:
+                    return False
+            else:
+                self._last_valid_iter = driver.step_count
         return True

{netket-3.12.4 → netket-3.14.0}/netket/callbacks/timeout.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
 import time
 
 from netket.utils import struct
@@ -29,7 +28,7 @@ class Timeout(struct.Pytree, mutable=True):
     """Number of seconds to wait before the training will be stopped."""
 
     # caches
-    _init_time: Optional[float]
+    _init_time: float | None
     """
     Internal field storing the time at which the first iteration has been
     performed.

{netket-3.12.4 → netket-3.14.0}/netket/driver/abstract_variational_driver.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Optional
-from collections.abc import Iterable
+from typing import TypeVar
+from collections.abc import Callable, Iterable
 
 import abc
 import numbers
@@ -29,19 +29,21 @@ from netket.utils import mpi, timing
 from netket.utils.types import Optimizer, PyTree
 from netket.vqs import VariationalState
 
+CallbackT = Callable[[int, dict, "AbstractVariationalDriver"], bool]
 
-def _to_iterable(maybe_iterable):
+T = TypeVar("T")
+
+
+def _to_iterable(maybe_iterable: T | Iterable[T]) -> tuple[T, ...]:
     """
     _to_iterable(maybe_iterable)
 
     Ensure the result is iterable. If the input is not iterable, it is wrapped into a tuple.
     """
-    if hasattr(maybe_iterable, "__iter__"):
-        surely_iterable = maybe_iterable
-    else:
-        surely_iterable = (maybe_iterable,)
+    if not isinstance(maybe_iterable, Iterable):
+        maybe_iterable = (maybe_iterable,)
 
-    return surely_iterable
+    return tuple(maybe_iterable)
 
 
 class AbstractVariationalDriver(abc.ABC):
@@ -161,6 +163,11 @@ class AbstractVariationalDriver(abc.ABC):
                 **modified in-place** adding new keys.
             step: the current step number.
         """
+        # Always log the acceptance.
+        if hasattr(self.state, "sampler_state"):
+            acceptance = getattr(self.state.sampler_state, "acceptance", None)  # type: ignore
+            if acceptance is not None:
+                log_dict["acceptance"] = acceptance
 
     def reset(self):
         """
@@ -216,12 +223,12 @@ class AbstractVariationalDriver(abc.ABC):
         """
         for _ in range(0, n_steps, step):
             for i in range(0, step):
-                dp = self._forward_and_backward()
+                self._dp = self._forward_and_backward()
                 if i == 0:
                     yield self.step_count
 
                 self._step_count += 1
-                self.update_parameters(dp)
+                self.update_parameters(self._dp)
 
     def advance(self, steps: int = 1):
         """
@@ -237,15 +244,13 @@ class AbstractVariationalDriver(abc.ABC):
     def run(
         self,
         n_iter: int,
-        out: Optional[Iterable[AbstractLog]] = (),
-        obs: Optional[dict[str, AbstractObservable]] = None,
+        out: AbstractLog | Iterable[AbstractLog] | str | None = (),
+        obs: dict[str, AbstractObservable] | None = None,
         step_size: int = 1,
         show_progress: bool = True,
         save_params_every: int = 50,  # for default logger
         write_every: int = 50,  # for default logger
-        callback: Callable[
-            [int, dict, "AbstractVariationalDriver"], bool
-        ] = lambda *x: True,
+        callback: CallbackT | Iterable[CallbackT] = lambda *x: True,
         timeit: bool = False,
     ):
         """
@@ -310,7 +315,7 @@ class AbstractVariationalDriver(abc.ABC):
         if self._is_root:
             loggers = _to_iterable(out)
         else:
-            loggers = tuple()
+            loggers: tuple[AbstractLog, ...] = tuple()
             show_progress = False
 
         callbacks = _to_iterable(callback)

{netket-3.12.4 → netket-3.14.0}/netket/driver/steady_state.py

@@ -12,16 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
-
-from inspect import signature
 
 from netket.operator import Squared, AbstractSuperOperator
 from netket.vqs import MCMixedState
 from netket.optimizer import (
     identity_preconditioner,
     PreconditionerT,
-    _DeprecatedPreconditionerSignature,
 )
 from netket.jax import tree_cast
 from netket.utils.types import Optimizer
@@ -115,13 +111,10 @@ class SteadyState(AbstractVariationalDriver):
         return self._preconditioner
 
     @preconditioner.setter
-    def preconditioner(self, val: Optional[PreconditionerT]):
+    def preconditioner(self, val: PreconditionerT | None):
         if val is None:
             val = identity_preconditioner
 
-        if len(signature(val).parameters) == 2:
-            val = _DeprecatedPreconditionerSignature(val)
-
         self._preconditioner = val
 
     @property

{netket-3.12.4 → netket-3.14.0}/netket/driver/vmc.py

@@ -12,10 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
 
 from textwrap import dedent
-from inspect import signature
 
 from netket.utils.types import PyTree, Optimizer
 from netket.operator import AbstractOperator
@@ -23,7 +21,6 @@ from netket.stats import Stats
 from netket.optimizer import (
     identity_preconditioner,
     PreconditionerT,
-    _DeprecatedPreconditionerSignature,
 )
 from netket.vqs import VariationalState
 from netket.jax import tree_cast
@@ -103,13 +100,10 @@ class VMC(AbstractVariationalDriver):
         return self._preconditioner
 
     @preconditioner.setter
-    def preconditioner(self, val: Optional[PreconditionerT]):
+    def preconditioner(self, val: PreconditionerT | None):
         if val is None:
             val = identity_preconditioner
 
-        if len(signature(val).parameters) == 2:
-            val = _DeprecatedPreconditionerSignature(val)
-
         self._preconditioner = val
 
     def _forward_and_backward(self):