NetKet 3.22.0.dev4__tar.gz → 3.22.2__tar.gz

{netket-3.22.0.dev4 → netket-3.22.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.22.0.dev4
+Version: 3.22.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/#
@@ -19,14 +19,14 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Physics
 Requires-Python: >=3.11
 Requires-Dist: einops>=0.8.1
-Requires-Dist: equinox>=0.11.3
+Requires-Dist: equinox>=0.12.2
 Requires-Dist: flax>=0.10.6
 Requires-Dist: igraph>=0.10.0
 Requires-Dist: jax>=0.7.0
 Requires-Dist: numba>=0.59
 Requires-Dist: numpy<3.0,>=1.20
 Requires-Dist: optax<0.3,>=0.2.2
-Requires-Dist: orjson<4,>=3.4
+Requires-Dist: orjson<4,>=3.10.7
 Requires-Dist: plum-dispatch>=2.4
 Requires-Dist: rich>=12.0
 Requires-Dist: scipy<2,>=1.5.3
@@ -51,7 +51,7 @@ Requires-Dist: wheel>=0.38.1; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: ipywidgets>=8.1.7; extra == 'docs'
 Requires-Dist: myst-nb<1.5,>=1.0; extra == 'docs'
-Requires-Dist: myst-parser<5.1.0,>=2.0.0; extra == 'docs'
+Requires-Dist: myst-parser<5.2.0,>=2.0.0; extra == 'docs'
 Requires-Dist: nbconvert>=7.16.6; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints<4.0,>=1.22; extra == 'docs'
 Requires-Dist: sphinx-book-theme<1.3,>=1.1; extra == 'docs'
@@ -62,6 +62,7 @@ Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
 Provides-Extra: extra
 Requires-Dist: h5py>=3.7.0; extra == 'extra'
+Requires-Dist: mlflow>=2.0.0; extra == 'extra'
 Requires-Dist: openfermion>=1.0.0; extra == 'extra'
 Requires-Dist: qutip>=4; extra == 'extra'
 Requires-Dist: tensorboardx>=2.0.0; extra == 'extra'

{netket-3.22.0.dev4 → netket-3.22.2}/netket/__init__.py

@@ -22,6 +22,7 @@ __all__ = [
     "logging",
     "models",
     "nn",
+    "observable",
     "operator",
     "optimizer",
     "sampler",
@@ -55,6 +56,7 @@ from . import (
     callbacks,
     logging,
     models,
+    observable,
     sampler,
     vqs,
     optimizer,

{netket-3.22.0.dev4 → netket-3.22.2}/netket/_src/callbacks/auto_chunk_size.py

@@ -16,7 +16,7 @@ def get_forward_operator(driver):
     ``@get_forward_operator.dispatch`` for custom driver types.
 
     Example:
-       class VMC(AbstractVariationalDriver):
+       class VMC(AbstractOptimizationDriver):
             ...
 
         @get_forward_operator.dispatch
@@ -59,7 +59,7 @@ class AutoChunkSize(AbstractCallback):
 
         from netket._src.callbacks.auto_chunk_size import get_forward_operator
 
-        class MyDriver(AbstractVariationalDriver):
+        class MyDriver(AbstractOptimizationDriver):
             def __init__(self, hamiltonian, ...):
                 self._ham = hamiltonian
                 ...

{netket-3.22.0.dev4 → netket-3.22.2}/netket/_src/callbacks/base.py

@@ -6,7 +6,7 @@ class StopRun(Exception):
     Exception to be raised by callbacks to gracefully stop the optimisation loop.
 
     Raise this exception (or a subclass of it) from any callback hook to stop
-    the :meth:`~netket.driver.AbstractVariationalDriver.run` loop early.
+    the :meth:`~netket.driver.AbstractDriver.run` loop early.
     The driver will catch it, call
     :meth:`~netket.callbacks.AbstractCallback.on_run_end` on all callbacks, and
     then return normally — no traceback is printed and no exception propagates to
@@ -104,7 +104,7 @@ class AbstractCallback(struct.Pytree, mutable=True):
 
         At this point:
 
-        - The loss and its gradient have been computed by :meth:`~netket.driver.AbstractVariationalDriver.compute_loss_and_update`.
+        - The loss and its gradient have been computed by :meth:`~netket.driver.AbstractDriver.compute_loss_and_update`.
         - The step has been accepted (not rejected by :meth:`on_compute_update_end`).
         - ``driver.step_count`` still refers to the *current* step — it has not yet been incremented.
         - The variational state parameters have **not** yet changed.

{netket-3.22.0.dev4/netket → netket-3.22.2/netket/_src}/callbacks/convergence_stopping.py

@@ -18,26 +18,27 @@ import numpy as np
 
 from netket.utils import struct
 
+from netket._src.callbacks.base import AbstractCallback, StopRun
 
-class ConvergenceStopping(struct.Pytree, mutable=True):
+
+class ConvergenceStopping(AbstractCallback, mutable=True):
     """A simple callback to stop the optimisation when the monitored quantity gets
     below a certain threshold for at least `patience` steps.
     """
 
-    target: float = struct.field(serialize=False)
+    target: float = struct.field(pytree_node=False)
     """Target value for the monitored quantity. Training will stop if the driver drops below this value."""
-    monitor: str = struct.field(serialize=False)
+    monitor: str = struct.field(pytree_node=False)
     """Loss statistic to monitor. Should be one of 'mean', 'variance', 'error_of_mean'."""
-    smoothing_window: int = struct.field(serialize=False)
+    smoothing_window: int = struct.field(pytree_node=False)
     """The loss is smoothed over the last `smoothing_window` iterations to
-    reduce statistical fluctuations"""
-    patience: int = struct.field(serialize=False)
+    reduce statistical fluctuations."""
+    patience: int = struct.field(pytree_node=False)
     """The loss must be consistently below this value for this number of
     iterations in order to stop the optimisation."""
 
-    # caches
-    _loss_window: deque
-    _patience_counter: int
+    _loss_window: deque = struct.field(pytree_node=False, serialize=False)
+    _patience_counter: int = struct.field(pytree_node=False, serialize=False, default=0)
 
     def __init__(
         self,
@@ -68,21 +69,10 @@ class ConvergenceStopping(struct.Pytree, mutable=True):
         self.smoothing_window = smoothing_window
         self.patience = patience
 
-        self._loss_window: deque = deque([], maxlen=self.smoothing_window)
-        self._patience_counter: int = 0
-
-    def __call__(self, step, log_data, driver):
-        """
-        A boolean function that determines whether or not to stop training.
+        self._loss_window = deque([], maxlen=smoothing_window)
+        self._patience_counter = 0
 
-        Args:
-            step: An integer corresponding to the step (iteration or epoch) in training.
-            log_data: A dictionary containing log data for training.
-            driver: A NetKet variational driver.
-
-        Returns:
-            A boolean. If True, training continues, else, it does not.
-        """
+    def on_step_end(self, step, log_data, driver):
         loss = np.asarray(np.real(getattr(log_data[driver._loss_name], self.monitor)))
 
         self._loss_window.append(loss)
@@ -94,6 +84,7 @@ class ConvergenceStopping(struct.Pytree, mutable=True):
             self._patience_counter = 0
 
         if self._patience_counter > self.patience:
-            return False
-
-        return True
+            raise StopRun(
+                f"ConvergenceStopping: smoothed loss {loss_smooth:.6g} has been below "
+                f"target {self.target} for {self._patience_counter} steps."
+            )

{netket-3.22.0.dev4/netket → netket-3.22.2/netket/_src}/callbacks/early_stopping.py

@@ -12,43 +12,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import numpy as np
 
 from netket.utils import struct
 
+from netket._src.callbacks.base import AbstractCallback, StopRun
+
 
-class EarlyStopping(struct.Pytree, mutable=True):
+class EarlyStopping(AbstractCallback, mutable=True):
     """A simple callback to stop NetKet if there are no more improvements in the training.
     based on `driver._loss_name`.
     """
 
-    min_delta: float
+    min_delta: float = struct.field(pytree_node=False)
     """Minimum change in the monitored quantity to qualify as an improvement."""
-    min_reldelta: float
+    min_reldelta: float = struct.field(pytree_node=False)
     """Minimum relative change in the monitored quantity to qualify as an improvement.
 
     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: int | float
+    patience: int | float = struct.field(pytree_node=False)
     """Number of epochs with no improvement after which training will be stopped."""
-    baseline: float | None
+    baseline: float | None = struct.field(pytree_node=False)
     """Baseline value for the monitored quantity. Training will stop if the driver is above the baseline."""
-    monitor: str
+    monitor: str = struct.field(pytree_node=False)
     """Loss statistic to monitor. Should be one of 'mean', 'variance', 'error_of_mean'."""
-    start_from_step: int
+    start_from_step: int = struct.field(pytree_node=False)
     """Number of steps to wait before the callback has any effect."""
 
-    # The quantities below are internal and should not be edited directly
-    # by the user
-
-    _best_val: float = np.inf
-    """Best value of the loss observed up to this iteration. """
-    _best_iter: int
+    _best_val: float = struct.field(pytree_node=False, serialize=False, default=np.inf)
+    """Best value of the loss observed up to this iteration."""
+    _best_iter: int = struct.field(pytree_node=False, serialize=False, default=0)
     """Iteration at which the `_best_val` was observed."""
-    _best_patience_counter: int
-    """Stores the iteration at which we've seen the best loss so far"""
+    _best_patience_counter: int = struct.field(
+        pytree_node=False, serialize=False, default=0
+    )
+    """Stores how many steps have elapsed without an improvement."""
 
     def __init__(
         self,
@@ -87,21 +87,9 @@ class EarlyStopping(struct.Pytree, mutable=True):
         self._best_iter = 0
         self._best_patience_counter = 0
 
-    def __call__(self, step, log_data, driver):
-        """
-        A boolean function that determines whether or not to stop training.
-
-        Args:
-            step: An integer corresponding to the step (iteration or epoch) in training.
-            log_data: A dictionary containing log data for training.
-            driver: A NetKet variational driver.
-
-        Returns:
-            A boolean. If True, training continues, else, it does not.
-        """
-
+    def on_step_end(self, step, log_data, driver):
         if step < self.start_from_step:
-            return True
+            return
 
         loss = np.real(getattr(log_data[driver._loss_name], self.monitor))
 
@@ -113,19 +101,16 @@ class EarlyStopping(struct.Pytree, mutable=True):
             if self.baseline is None:
                 self._best_patience_counter = 0
             elif self._is_improvement(loss, self.baseline):
-                # If using baseline, update patience only if we are better than baseline
                 self._best_patience_counter = 0
 
         if self._best_patience_counter > self.patience:
-            return False
-
-        return True
+            raise StopRun(
+                f"EarlyStopping: no improvement detected after {self._best_patience_counter} steps "
+                f"(best value {self._best_val} at step {self._best_iter})."
+            )
 
     def _is_improvement(self, loss, target):
-        # minimal value for absolute and relative improvement
         abs_minval = target - self.min_delta
         rel_minval = target * (1 - self.min_reldelta)
-        # minimval value that qualify as an improvement
         minval = min(abs_minval, rel_minval)
-
         return loss < minval

{netket-3.22.0.dev4/netket → netket-3.22.2/netket/_src}/callbacks/invalid_loss_stopping.py

@@ -16,21 +16,21 @@ import numpy as np
 
 from netket.utils import struct
 
+from netket._src.callbacks.base import AbstractCallback, StopRun
 
-# Mark this class a NetKet dataclass so that it can automatically be serialized by Flax.
-class InvalidLossStopping(struct.Pytree, mutable=True):
+
+class InvalidLossStopping(AbstractCallback, mutable=True):
     """A simple callback to stop the optimisation when the monitored quantity becomes
     invalid for at least `patience` steps.
     """
 
-    monitor: str
+    monitor: str = struct.field(pytree_node=False)
     """Loss statistic to monitor. Should be one of 'mean', 'variance', 'error_of_mean'."""
-    patience: int | float
+    patience: int | float = struct.field(pytree_node=False)
     """Number of epochs with invalid loss after which training will be stopped."""
 
-    # caches
-    _last_valid_iter: int
-    """Last valid iteration, to check against patience"""
+    _last_valid_iter: int = struct.field(pytree_node=False, serialize=False, default=0)
+    """Last valid iteration, to check against patience."""
 
     def __init__(self, monitor: str = "mean", patience: int | float = 0):
         """
@@ -48,23 +48,10 @@ class InvalidLossStopping(struct.Pytree, mutable=True):
         """
         self.monitor = monitor
         self.patience = patience
-
-        # caches
         self._last_valid_iter = 0
 
-    def __call__(self, step, log_data, driver):
-        """
-        A boolean function that determines whether or not to stop training.
-
-        Args:
-            step: An integer corresponding to the step (iteration or epoch) in training.
-            log_data: A dictionary containing log data for training.
-            driver: A NetKet variational driver.
-
-        Returns:
-            A boolean. If True, training continues, else, it does not.
-        """
-        # clears the _last_valid_iter in case the driver was reset
+    def on_step_end(self, step, log_data, driver):
+        # reset if the driver was restarted
         if driver.step_count < self._last_valid_iter:
             self._last_valid_iter = 0
 
@@ -73,7 +60,9 @@ class InvalidLossStopping(struct.Pytree, mutable=True):
 
             if not np.isfinite(loss):
                 if driver.step_count - self._last_valid_iter >= self.patience:
-                    return False
+                    raise StopRun(
+                        f"InvalidLossStopping: loss is not finite ({loss}) "
+                        f"for {driver.step_count - self._last_valid_iter} steps."
+                    )
             else:
                 self._last_valid_iter = driver.step_count
-        return True

{netket-3.22.0.dev4/netket → netket-3.22.2/netket/_src}/callbacks/timeout.py

@@ -16,27 +16,27 @@ import time
 
 from netket.utils import struct
 
+from netket._src.callbacks.base import AbstractCallback, StopRun
 
-class Timeout(struct.Pytree, mutable=True):
+
+class Timeout(AbstractCallback, mutable=True):
     """A simple callback to stop NetKet after some time has passed.
 
     This callback monitors whether a driver has been training for more
     than a given timeout in order to hard stop training.
     """
 
-    timeout: float
+    timeout: float = struct.field(pytree_node=False)
     """Number of seconds to wait before the training will be stopped."""
 
-    # caches
-    _init_time: float | None = struct.field(serialize=False)
-    """
-    Internal field storing the time at which the first iteration has been
-    performed.
-    """
+    _init_time: float | None = struct.field(
+        pytree_node=False, serialize=False, default=None
+    )
+    """Internal field storing the time at which the run started."""
 
     def __init__(self, timeout: float):
         """
-        Stops the optimisation after a certain time itnerval.
+        Stops the optimisation after a certain time interval.
 
         Args:
             timeout: number of seconds after which the optimisation will
@@ -45,28 +45,13 @@ class Timeout(struct.Pytree, mutable=True):
         if not timeout > 0:
             raise ValueError("`timeout` must be larger than 0.")
         self.timeout = timeout
-
         self._init_time = None
 
-    def __call__(self, step, log_data, driver):
-        """
-        A boolean function that determines whether or not to stop training.
-
-        Args:
-            step: An integer corresponding to the step (iteration or epoch) in training.
-            log_data: A dictionary containing log data for training.
-            driver: A NetKet variational driver.
-
-        Returns:
-            A boolean. If True, training continues, else, it does not.
-
-        Note:
-            This callback does not make use of `step`, `log_data` nor `driver`.
-        """
-        if self._init_time is None:
-            self._init_time = time.time()
+    def on_run_start(self, step, driver):
+        self._init_time = time.time()
 
+    def on_step_end(self, step, log_data, driver):
         if time.time() - self._init_time >= self.timeout:
-            return False
-        else:
-            return True
+            raise StopRun(
+                f"Timeout: training stopped after {self.timeout:.1f} seconds."
+            )

{netket-3.22.0.dev4 → netket-3.22.2}/netket/_src/driver/abstract_optimization_driver.py

@@ -31,8 +31,9 @@ class AbstractOptimizationDriver(AbstractDriver):
     Adds an `optax <https://optax.readthedocs.io/en/latest/>`_ optimizer and
     implements :meth:`update_parameters` via :func:`apply_gradient`.
 
-    Subclass this (not :class:`AbstractVariationalDriver` directly) when
-    implementing a new optimization driver.
+    Subclass this when implementing a new optimization driver. The public
+    name :class:`AbstractVariationalDriver` is currently kept as a compatibility
+    alias for this class.
     """
 
     _optimizer: Optimizer = struct.field(pytree_node=False, serialize=False)

{netket-3.22.0.dev4 → netket-3.22.2}/netket/_src/driver/abstract_variational_driver.py

@@ -313,7 +313,7 @@ class AbstractDriver(struct.Pytree, mutable=True):
         Example of subclassing:
 
         .. code:: python
-            class MyDriver(AbstractVariationalDriver):
+            class MyDriver(AbstractDriver):
                 ...
                 def _default_callbacks(self, *args, **kwargs):
                     callbacks = super()._default_callbacks(*args, **kwargs)

{netket-3.22.0.dev4 → netket-3.22.2}/netket/_src/driver/abstract_variational_driver_old.py

@@ -197,7 +197,7 @@ class AbstractVariationalDriver(abc.ABC):
         self._optimizer = optimizer
         if optimizer is not None:
             self._optimizer_state = optimizer.init(self.state.parameters)
-            if config.netket_experimental_sharding:
+            if config.netket_sharding:
                 self._optimizer_state = jax.lax.with_sharding_constraint(
                     self._optimizer_state,
                     NamedSharding(jax.sharding.get_abstract_mesh(), P()),
@@ -433,7 +433,7 @@ def apply_gradient(optimizer_fun, optimizer_state, dp, params):
 
     new_params = optax.apply_updates(params, updates)
 
-    if config.netket_experimental_sharding:
+    if config.netket_sharding:
         sharding = NamedSharding(jax.sharding.get_abstract_mesh(), P())
         new_optimizer_state = jax.lax.with_sharding_constraint(
             new_optimizer_state, sharding

{netket-3.22.0.dev4 → netket-3.22.2}/netket/_src/driver/infidelity_sr.py

@@ -24,8 +24,8 @@ from netket._src.ngd.sr_srt_common import sr, srt, get_samples_and_pdf
 from netket._src.ngd.srt_onthefly import srt_onthefly
 from netket._src.nn.apply_operator.functional import make_logpsi_op_afun
 from netket._src.callbacks.auto_chunk_size import get_forward_operator
-from netket.experimental.observable.infidelity.expect import get_local_estimator
-from netket.experimental.observable.infidelity import InfidelityOperator
+from netket._src.observable.infidelity.expect import get_local_estimator
+from netket._src.observable.infidelity import InfidelityOperator
 
 
 ApplyFun = Callable[[PyTree, Array], Array]
@@ -344,7 +344,7 @@ class Infidelity_SR(AbstractOptimizationDriver):
     @timing.timed
     def _log_additional_data(self, log_dict: dict):
         """
-        Method to be implemented in sub-classes of AbstractVariationalDriver to
+        Method to be implemented in sub-classes of AbstractDriver to
         log additional data at every step.
         This method is called at every iteration when executing with `run`.
 

{netket-3.22.0.dev4 → netket-3.22.2}/netket/_src/driver/vmc_sr.py

@@ -512,14 +512,13 @@ class VMC_SR(AbstractOptimizationDriver):
     @timing.timed
     def _log_additional_data(self, log_dict: dict):
         """
-        Method to be implemented in sub-classes of AbstractVariationalDriver to
+        Method to be implemented in sub-classes of AbstractDriver to
         log additional data at every step.
         This method is called at every iteration when executing with `run`.
 
         Args:
             log_dict: The dictionary containing all logged data. It must be
                 **modified in-place** adding new keys.
-            step: the current step number.
         """
         super()._log_additional_data(log_dict)
 

{netket-3.22.0.dev4/netket → netket-3.22.2/netket/_src}/logging/base.py

@@ -18,8 +18,6 @@ from typing import Any, TYPE_CHECKING
 
 import jax
 
-from netket.vqs import VariationalState
-
 if TYPE_CHECKING:
     from netket.vqs import VariationalState