NetKet 3.11.4__tar.gz → 3.12.0__tar.gz

{netket-3.11.4 → netket-3.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: NetKet
-Version: 3.11.4
+Version: 3.12.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/#
@@ -19,10 +19,11 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Physics
 Requires-Python: >=3.9
 Requires-Dist: beartype>=0.16.2
+Requires-Dist: equinox<0.12,>=0.11.3
 Requires-Dist: flax<0.9,>=0.7.1
 Requires-Dist: igraph<0.12.0,>=0.10.0
-Requires-Dist: jax<0.5,>=0.4.3
-Requires-Dist: jaxlib<0.5,>=0.4.3
+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.60,>=0.52
 Requires-Dist: numpy~=1.20
@@ -45,15 +46,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.1.0; extra == 'dev'
+Requires-Dist: ruff==0.4.2; extra == 'dev'
 Requires-Dist: wheel; extra == 'dev'
 Provides-Extra: docs
-Requires-Dist: myst-nb~=1.0.0; extra == 'docs'
-Requires-Dist: myst-parser<2.1.0,>=2.0.0; 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: 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.3,>=7.0; extra == 'docs'
+Requires-Dist: sphinx<7.4,>=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

{netket-3.11.4 → netket-3.12.0}/netket/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.11.4'
-__version_tuple__ = version_tuple = (3, 11, 4)
+__version__ = version = '3.12.0'
+__version_tuple__ = version_tuple = (3, 12, 0)

{netket-3.11.4 → netket-3.12.0}/netket/callbacks/convergence_stopping.py

@@ -12,33 +12,48 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
-
 from collections import deque
 
 import numpy as np
 
+from netket.utils import struct
+
 
-@dataclass
-class ConvergenceStopping:
+class ConvergenceStopping(struct.Pytree, mutable=True):
     """A simple callback to stop the optimisation if the loss gets below a certain threshold.
     based on `driver._loss_name`."""
 
     target: float
     """Target value for the monitored quantity. Training will stop if the driver hits the baseline."""
-    monitor: str = "mean"
+    monitor: str
     """Loss statistic to monitor. Should be one of 'mean', 'variance', 'sigma'."""
 
-    smoothing_window: int = 10
+    smoothing_window: int
     """
     The loss is smoothed over the last `smoothing_window` iterations to reduce statistical fluctuations
     """
-    patience: int = 10
+    patience: int
     """
     The loss must be consistently below this value for this number of iterations in order to stop the optimisation.
     """
 
-    def __post_init__(self):
+    # caches
+    _loss_window: deque
+    _patience_counter: int
+
+    def __init__(
+        self,
+        target: float,
+        monitor: str = "mean",
+        *,
+        smoothing_window: int = 10,
+        patience: int = 10,
+    ):
+        self.target = target
+        self.monitor = monitor
+        self.smoothing_window = smoothing_window
+        self.patience = patience
+
         self._loss_window: deque = deque([], maxlen=self.smoothing_window)
         self._patience_counter: int = 0
 

{netket-3.11.4 → netket-3.12.0}/netket/callbacks/earlystopping.py

@@ -19,26 +19,24 @@ import numpy as np
 from netket.utils import struct
 
 
-# Mark this class a NetKet dataclass so that it can automatically be serialized by Flax.
-@struct.dataclass(_frozen=False)
-class EarlyStopping:
+class EarlyStopping(struct.Pytree, 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 = 0.0
+    min_delta: float
     """Minimum change in the monitored quantity to qualify as an improvement."""
-    min_reldelta: float = 0.0
+    min_reldelta: float
     """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: Union[int, float] = 0
+    patience: Union[int, float]
     """Number of epochs with no improvement after which training will be stopped."""
-    baseline: Optional[float] = None
+    baseline: Optional[float]
     """Baseline value for the monitored quantity. Training will stop if the driver hits the baseline."""
-    monitor: str = "mean"
+    monitor: str
     """Loss statistic to monitor. Should be one of 'mean', 'variance', 'sigma'."""
 
     # The quantities below are internal and should not be edited directly
@@ -46,11 +44,44 @@ class EarlyStopping:
 
     _best_val: float = np.inf
     """Best value of the loss observed up to this iteration. """
-    _best_iter: int = 0
+    _best_iter: int
     """Iteration at which the `_best_val` was observed."""
-    _best_patience_counter: int = 0
+    _best_patience_counter: int
     """Stores the iteration at which we've seen the best loss so far"""
 
+    def __init__(
+        self,
+        min_delta: float = 0.0,
+        min_reldelta: float = 0.0,
+        patience: Union[int, float] = 0,
+        baseline: Optional[float] = None,
+        monitor: str = "mean",
+    ):
+        """
+        Construct an early stopping callback.
+
+        Args:
+            min_delta: Minimum change in the monitored quantity to
+                qualify as an improvement.
+            min_reldelta: Minimum relative change in the monitored
+                quantity to qualify as an improvement.
+            patience: Number of epochs with no improvement after which
+                training will be stopped.
+            baseline: Baseline value for the monitored quantity. Training
+                will stop if the driver hits the baseline.
+            monitor: Loss statistic to monitor. Should be one of
+                `mean`, `variance`, `sigma`.
+        """
+        self.min_delta = min_delta
+        self.min_reldelta = min_reldelta
+        self.patience = patience
+        self.baseline = baseline
+        self.monitor = monitor
+
+        self._best_val = np.infty
+        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.

{netket-3.11.4 → netket-3.12.0}/netket/callbacks/invalidlossstopping.py

@@ -19,19 +19,38 @@ from netket.utils import struct
 
 
 # Mark this class a NetKet dataclass so that it can automatically be serialized by Flax.
-@struct.dataclass(_frozen=False)
-class InvalidLossStopping:
+class InvalidLossStopping(struct.Pytree, mutable=True):
     """A simple callback to stop NetKet if there are no more improvements in the training.
     based on `driver._loss_name`."""
 
-    monitor: str = "mean"
+    monitor: str
     """Loss statistic to monitor. Should be one of 'mean', 'variance', 'sigma'."""
-    patience: Union[int, float] = 0
+    patience: Union[int, float]
     """Number of epochs with invalid loss after which training will be stopped."""
 
-    _last_valid_iter: int = 0
+    # caches
+    _last_valid_iter: int
     """Last valid iteration, to check against patience"""
 
+    def __init__(self, monitor: str = "mean", patience: Union[int, float] = 0):
+        """
+        Construct a callback stopping theoptimisation when the monitored quantity
+        becaomes invalid for at least `patience` steps.
+
+        Args:
+            monitor: a string with the name of the quantity to be monitored. This
+                is applied to the standard loss optimised by a driver, such as the
+                Energy for the VMC driver (default: 'mean').
+            patience: Number of steps to wait before stopping the execution after
+                the tracked quantity becomes invalid (default 0, meaning that it
+                stops immediately).
+        """
+        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.

{netket-3.11.4 → netket-3.12.0}/netket/callbacks/timeout.py

@@ -18,9 +18,7 @@ import time
 from netket.utils import struct
 
 
-# Mark this class a NetKet dataclass so that it can automatically be serialized by Flax.
-@struct.dataclass(_frozen=False)
-class Timeout:
+class Timeout(struct.Pytree, mutable=True):
     """A simple callback to stop NetKet after some time has passed.
 
     This callback monitors whether a driver has been training for more
@@ -30,19 +28,26 @@ class Timeout:
     timeout: float
     """Number of seconds to wait before the training will be stopped."""
 
-    _init_time: Optional[float] = None
+    # caches
+    _init_time: Optional[float]
     """
     Internal field storing the time at which the first iteration has been
     performed.
     """
 
-    def __post_init__(self):
-        if not self.timeout > 0:
+    def __init__(self, timeout: float):
+        """
+        Stops the optimisation after a certain time itnerval.
+
+        Args:
+            timeout: number of seconds after which the optimisation will
+                be stopped.
+        """
+        if not timeout > 0:
             raise ValueError("`timeout` must be larger than 0.")
+        self.timeout = timeout
 
-    def reset(self):
-        """Resets the initial time of the training"""
-        self.__init_time = None
+        self._init_time = None
 
     def __call__(self, step, log_data, driver):
         """

{netket-3.11.4 → netket-3.12.0}/netket/driver/abstract_variational_driver.py

@@ -19,13 +19,13 @@ import abc
 import numbers
 from functools import partial
 
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 import jax
 
 from netket.logging import AbstractLog, JsonLog
 from netket.operator._abstract_observable import AbstractObservable
-from netket.utils import mpi
+from netket.utils import mpi, timing
 from netket.utils.types import Optimizer, PyTree
 from netket.vqs import VariationalState
 
@@ -104,6 +104,7 @@ class AbstractVariationalDriver(abc.ABC):
         self._loss_stats = None
         self._loss_name = minimized_quantity_name
         self._step_count = 0
+        self._timer = None
 
         self._variational_state = variational_state
         self.optimizer = optimizer
@@ -245,6 +246,7 @@ class AbstractVariationalDriver(abc.ABC):
         callback: Callable[
             [int, dict, "AbstractVariationalDriver"], bool
         ] = lambda *x: True,
+        timeit: bool = False,
     ):
         """
         Runs this variational driver, updating the weights of the network stored in
@@ -287,6 +289,7 @@ class AbstractVariationalDriver(abc.ABC):
                 serialized to disk (ignored if logger is provided)
             write_every: Every how many steps the json data should be flushed to disk (ignored if
                 logger is provided)
+            timeit: If True, provide timing information.
         """
 
         if not isinstance(n_iter, numbers.Number):
@@ -313,52 +316,62 @@ class AbstractVariationalDriver(abc.ABC):
         callbacks = _to_iterable(callback)
         callback_stop = False
 
-        with tqdm(
-            total=n_iter,
-            disable=not show_progress,
-            dynamic_ncols=True,
-        ) as pbar:
-            old_step = self.step_count
-            first_step = True
-
-            for step in self.iter(n_iter, step_size):
-                log_data = self.estimate(obs)
-                self._log_additional_data(log_data, step)
-
-                # if the cost-function is defined then report it in the progress bar
-                if self._loss_stats is not None:
-                    pbar.set_postfix_str(self._loss_name + "=" + str(self._loss_stats))
-                    log_data[self._loss_name] = self._loss_stats
-
-                # Execute callbacks before loggers because they can append to log_data
-                for callback in callbacks:
-                    if not callback(step, log_data, self):
-                        callback_stop = True
-
-                for logger in loggers:
-                    logger(self.step_count, log_data, self.state)
-
-                if len(callbacks) > 0:
-                    if mpi.mpi_any(callback_stop):
-                        break
-
-                # Reset the timing of tqdm after the first step, to ignore compilation time
-                if first_step:
-                    first_step = False
-                    pbar.unpause()
-
-                # Update the progress bar
-                pbar.update(self.step_count - old_step)
+        with timing.timed_scope(force=timeit) as timer:
+            with tqdm(
+                total=n_iter,
+                disable=not show_progress,
+                dynamic_ncols=True,
+            ) as pbar:
                 old_step = self.step_count
-
-            # Final update so that it shows up filled.
-            pbar.update(self.step_count - old_step)
+                first_step = True
+
+                for step in self.iter(n_iter, step_size):
+                    with timing.timed_scope(name="observables"):
+                        log_data = self.estimate(obs)
+                        self._log_additional_data(log_data, step)
+
+                    # if the cost-function is defined then report it in the progress bar
+                    if self._loss_stats is not None:
+                        pbar.set_postfix_str(
+                            self._loss_name + "=" + str(self._loss_stats)
+                        )
+                        log_data[self._loss_name] = self._loss_stats
+
+                    # Execute callbacks before loggers because they can append to log_data
+                    for callback in callbacks:
+                        if not callback(step, log_data, self):
+                            callback_stop = True
+
+                    with timing.timed_scope(name="loggers"):
+                        for logger in loggers:
+                            logger(self.step_count, log_data, self.state)
+
+                    if len(callbacks) > 0:
+                        if mpi.mpi_any(callback_stop):
+                            break
+
+                    # Reset the timing of tqdm after the first step, to ignore compilation time
+                    if first_step:
+                        first_step = False
+                        pbar.unpause()
+
+                    # Update the progress bar
+                    pbar.update(self.step_count - old_step)
+                    old_step = self.step_count
+
+                # Final update so that it shows up filled.
+                pbar.update(self.step_count - old_step)
 
         # flush at the end of the evolution so that final values are saved to
         # file
         for logger in loggers:
             logger.flush(self.state)
 
+        if timeit:
+            self._timer = timer
+            if self._is_root:
+                print(timer)
+
         return loggers
 
     def estimate(self, observables):

{netket-3.11.4 → netket-3.12.0}/netket/experimental/driver/tdvp_common.py

@@ -19,7 +19,7 @@ import warnings
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 import netket as nk
 from netket import config
@@ -28,7 +28,7 @@ from netket.driver.abstract_variational_driver import _to_iterable
 from netket.jax import HashablePartial
 from netket.logging.json_log import JsonLog
 from netket.operator import AbstractOperator
-from netket.utils import mpi
+from netket.utils import mpi, timing
 from netket.utils.dispatch import dispatch
 from netket.utils.types import PyTree
 from netket.vqs import VariationalState
@@ -275,12 +275,13 @@ class TDVPBaseDriver(AbstractVariationalDriver):
     def run(
         self,
         T,
-        out=None,
+        out=(),
         obs=None,
         *,
         tstops=None,
         show_progress=True,
         callback=None,
+        timeit: bool = False,
     ):
         """
         Runs the time evolution.
@@ -301,6 +302,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
             show_progress: If true displays a progress bar (default=True)
             callback: Callable or list of callable callback functions to be executed at each
                 stopping time.
+            timeit: If True, provide timing information.
         """
         if obs is None:
             obs = {}
@@ -308,15 +310,15 @@ class TDVPBaseDriver(AbstractVariationalDriver):
         if callback is None:
             callback = lambda *_args, **_kwargs: True
 
+        # if out is a path, create an overwriting Json Log for output
+        if isinstance(out, str):
+            out = JsonLog(out, "w")
+        elif out is None:
+            out = ()
+
         # Log only non-root nodes
-        if self._mynode == 0:
-            if out is None:
-                loggers = ()
-            # if out is a path, create an overwriting Json Log for output
-            elif isinstance(out, str):
-                loggers = (JsonLog(out, "w"),)
-            else:
-                loggers = _to_iterable(out)
+        if self._is_root:
+            loggers = _to_iterable(out)
         else:
             loggers = tuple()
             show_progress = False
@@ -325,6 +327,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
         callback_stop = False
 
         t_end = np.asarray(self.t + T)
+
         with tqdm(
             total=t_end,
             disable=not show_progress,
@@ -353,33 +356,35 @@ class TDVPBaseDriver(AbstractVariationalDriver):
                 pbar.set_postfix(self._postfix)
                 pbar.refresh()
 
-            for step in self._iter(T, tstops=tstops, callback=update_progress_bar):
-                log_data = self.estimate(obs)
-                self._log_additional_data(log_data, step)
-
-                self._postfix = {"n": self.step_count}
-                # if the cost-function is defined then report it in the progress bar
-                if self._loss_stats is not None:
-                    self._postfix.update(
-                        {
-                            self._loss_name: str(self._loss_stats),
-                        }
-                    )
-                    log_data[self._loss_name] = self._loss_stats
-                pbar.set_postfix(self._postfix)
-
-                # Execute callbacks before loggers because they can append to log_data
-                for callback in callbacks:
-                    if not callback(step, log_data, self):
-                        callback_stop = True
-
-                for logger in loggers:
-                    logger(self.step_value, log_data, self.state)
-
-                if len(callbacks) > 0:
-                    if mpi.mpi_any(callback_stop):
-                        break
-                update_progress_bar()
+            with timing.timed_scope(force=timeit) as timer:
+                for step in self._iter(T, tstops=tstops, callback=update_progress_bar):
+                    with timing.timed_scope(name="observables"):
+                        log_data = self.estimate(obs)
+                        self._log_additional_data(log_data, step)
+
+                    self._postfix = {"n": self.step_count}
+                    # if the cost-function is defined then report it in the progress bar
+                    if self._loss_stats is not None:
+                        self._postfix.update(
+                            {
+                                self._loss_name: str(self._loss_stats),
+                            }
+                        )
+                        log_data[self._loss_name] = self._loss_stats
+                    pbar.set_postfix(self._postfix)
+
+                    # Execute callbacks before loggers because they can append to log_data
+                    for callback in callbacks:
+                        if not callback(step, log_data, self):
+                            callback_stop = True
+
+                    for logger in loggers:
+                        logger(self.step_value, log_data, self.state)
+
+                    if len(callbacks) > 0:
+                        if mpi.mpi_any(callback_stop):
+                            break
+                    update_progress_bar()
 
             # Final update so that it shows up filled.
             update_progress_bar()
@@ -389,6 +394,10 @@ class TDVPBaseDriver(AbstractVariationalDriver):
         for logger in loggers:
             logger.flush(self.state)
 
+        if timeit:
+            self._timer = timer
+            if self._is_root:
+                print(timer)
         return loggers
 
     def _log_additional_data(self, log_dict, step):

{netket-3.11.4 → netket-3.12.0}/netket/experimental/driver/tdvp_schmitt.py

@@ -25,6 +25,7 @@ from netket.optimizer.qgt import QGTJacobianDense
 from netket.optimizer.qgt.qgt_jacobian_dense import convert_tree_to_dense_format
 from netket.vqs import VariationalState, VariationalMixedState, MCState
 from netket.jax import tree_cast
+from netket.utils import timing
 
 from netket.experimental.dynamics import RKIntegratorConfig
 
@@ -190,6 +191,7 @@ class TDVPSchmitt(TDVPBaseDriver):
 # MIT License, Copyright (c) 2021 Markus Schmitt
 
 
+@timing.timed
 @partial(jax.jit, static_argnames=("n_samples"))
 def _impl(parameters, n_samples, E_loc, S, rhs_coeff, rcond, rcond_smooth, snr_atol):
     E = stats.statistics(E_loc)

{netket-3.11.4 → netket-3.12.0}/netket/experimental/driver/vmc_srt.py

@@ -28,13 +28,14 @@ from netket.driver import AbstractVariationalDriver
 from netket.errors import UnoptimalSRtWarning
 from netket.jax import sharding
 from netket.operator import AbstractOperator
-from netket.utils import mpi
+from netket.utils import mpi, timing
 from netket.utils.types import ScalarOrSchedule, Optimizer, PyTree
 from netket.vqs import MCState
 
 from jax.flatten_util import ravel_pytree
 
 
+@timing.timed
 @partial(jax.jit, static_argnames=("mode", "solver_fn"))
 def SRt(
     O_L, local_energies, diag_shift, *, mode, solver_fn, e_mean=None, params_structure
@@ -49,7 +50,7 @@ def SRt(
     local_energies = local_energies.flatten()
 
     if e_mean is None:
-        e_mean = mpi.mean(local_energies)
+        e_mean = nkstats.mean(local_energies)
     de = jnp.conj(local_energies - e_mean).squeeze()
 
     # * in this case O_L should be padded with zeros
@@ -102,7 +103,7 @@ def SRt(
     # To repack the real coefficients in order to get complex updates
     if mode == "complex" and nkjax.tree_leaf_iscomplex(params_structure):
         np = updates.shape[-1] // 2
-        updates = (updates[:np] + 1j * updates[np:]) / 2
+        updates = updates[:np] + 1j * updates[np:]
 
     return -updates
 
@@ -278,6 +279,7 @@ class VMC_SRt(AbstractVariationalDriver):
             mode=self.jacobian_mode,
             dense=True,
             center=True,
+            chunk_size=self.state.chunk_size,
         )  # jacobians is centered
 
         diag_shift = self.diag_shift

{netket-3.11.4 → netket-3.12.0}/netket/experimental/dynamics/_rk_solver_structures.py

@@ -153,14 +153,7 @@ class RungeKuttaState:
             last_norm = f"{self.last_norm}"
             accepted = f"{SolverFlags.INFO_STEP_ACCEPTED}"
 
-        return "RKState(step_no(total)={}({}), t={}, dt={}{}{})".format(
-            self.step_no,
-            self.step_no_total,
-            self.t.value,
-            dt,
-            last_norm,
-            accepted,
-        )
+        return f"RKState(step_no(total)={self.step_no}({self.step_no_total}), t={self.t.value}, dt={dt}{last_norm}{accepted})"
 
     @property
     def accepted(self):

{netket-3.11.4 → netket-3.12.0}/netket/experimental/hilbert/spin_orbital_fermions.py

@@ -23,6 +23,7 @@ from netket.hilbert.fock import Fock
 from netket.hilbert.tensor_hilbert_discrete import TensorDiscreteHilbert
 from netket.hilbert.homogeneous import HomogeneousHilbert
 from netket.utils.dispatch import dispatch
+from netket.utils import StaticRange
 
 
 class SpinOrbitalFermions(HomogeneousHilbert):
@@ -156,7 +157,7 @@ class SpinOrbitalFermions(HomogeneousHilbert):
         self._fock = hilbert
         """Internal representation of this Hilbert space (Fock or TensorHilbert)."""
         # local states are the occupation numbers (0, 1)
-        local_states = np.array((0.0, 1.0))
+        local_states = StaticRange(0.0, 1.0, 2, dtype=np.int8)
 
         # we use the constraints from the Fock spaces, and override `constrained`
         super().__init__(local_states, N=total_size, constraint_fn=None)
@@ -302,6 +303,11 @@ class SpinOrbitalFermions(HomogeneousHilbert):
         _str += ")"
         return _str
 
+    @property
+    def is_indexable(self) -> bool:
+        """Whether the space can be indexed with an integer"""
+        return self._fock.is_indexable
+
 
 @dispatch
 def random_state(hilb: SpinOrbitalFermions, key, batches: int, *, dtype):