NetKet 3.9.1__tar.gz → 3.10__tar.gz

{netket-3.9.1 → netket-3.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: NetKet
-Version: 3.9.1
+Version: 3.10
 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,23 +17,26 @@ Classifier: Operating System :: POSIX :: Linux
 Classifier: Operating System :: Unix
 Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Physics
-Requires-Python: >=3.8
-Requires-Dist: flax<0.8,>=0.6.5
-Requires-Dist: igraph<0.11.0,>=0.10.0
+Requires-Python: >=3.9
+Requires-Dist: beartype>=0.16.2
+Requires-Dist: flax<0.8,>=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: numba4jax<0.1,>=0.0.10.post1
-Requires-Dist: numba<0.58,>=0.52
+Requires-Dist: numba<0.59,>=0.52
 Requires-Dist: numpy~=1.20
 Requires-Dist: optax<0.2,>=0.1.3
 Requires-Dist: orjson<4,>=3.4
-Requires-Dist: plum-dispatch<3,>=1.5.1
+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==22.10.0; 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'
@@ -41,25 +44,28 @@ Requires-Dist: pytest-cov>=2.10.1; extra == 'dev'
 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: ruff==0.0.278; extra == 'dev'
+Requires-Dist: ruff==0.1.0; extra == 'dev'
 Requires-Dist: wheel; extra == 'dev'
 Provides-Extra: docs
-Requires-Dist: myst-nb~=0.17.1; extra == 'docs'
+Requires-Dist: myst-nb~=1.0.0rc0; extra == 'docs'
 Requires-Dist: myst-parser<2.1.0,>=0.18.1; extra == 'docs'
-Requires-Dist: nbsphinx~=0.9.1; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints~=1.22; extra == 'docs'
 Requires-Dist: sphinx-book-theme~=1.0.1; extra == 'docs'
-Requires-Dist: sphinx<7.1,>=5.3; extra == 'docs'
+Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
+Requires-Dist: sphinx<7.3,>=6.2; extra == 'docs'
 Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
-Requires-Dist: sphinxcontrib-napoleon==0.7; extra == 'docs'
 Provides-Extra: extra
 Requires-Dist: h5py>=3.7.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'
 Provides-Extra: mpi
-Requires-Dist: mpi4jax<0.4,>=0.3.9; extra == 'mpi'
+Requires-Dist: mpi4jax<0.5,>=0.3.9; extra == 'mpi'
 Requires-Dist: mpi4py<4,>=3.0.1; extra == 'mpi'
+Provides-Extra: pyscf
+Requires-Dist: pyscf>=2.0; extra == 'pyscf'
+Requires-Dist: sparse>=0.12.0; extra == 'pyscf'
 Description-Content-Type: text/markdown
 
 <div align="center">
@@ -158,7 +164,7 @@ If you want MPI support, please follow the discussion in [mpi4jax](https://githu
 
 ## Getting Started
 
-To get started with NetKet, we recommend you give a look at our [tutorials page](https://netket.readthedocs.io/en/latest/tutorials/gs-ising.htmls), by running them on your computer or on [Google Colaboratory](https://colab.research.google.com).
+To get started with NetKet, we recommend you give a look at our [tutorials page](https://netket.readthedocs.io/en/latest/tutorials/gs-ising.html), by running them on your computer or on [Google Colaboratory](https://colab.research.google.com).
 There are also many example scripts that you can download, run and edit that showcase some use-cases of NetKet, although they are not commented.
 
 If you want to get in touch with us, feel free to open an issue or a discussion here on GitHub, or to join the MLQuantum slack group where several people involved with NetKet hang out. To join the slack channel just accept [this invitation](https://join.slack.com/t/mlquantum/shared_invite/zt-19wibmfdv-LLRI6i43wrLev6oQX0OfOw)

{netket-3.9.1 → netket-3.10}/README.md

@@ -94,7 +94,7 @@ If you want MPI support, please follow the discussion in [mpi4jax](https://githu
 
 ## Getting Started
 
-To get started with NetKet, we recommend you give a look at our [tutorials page](https://netket.readthedocs.io/en/latest/tutorials/gs-ising.htmls), by running them on your computer or on [Google Colaboratory](https://colab.research.google.com).
+To get started with NetKet, we recommend you give a look at our [tutorials page](https://netket.readthedocs.io/en/latest/tutorials/gs-ising.html), by running them on your computer or on [Google Colaboratory](https://colab.research.google.com).
 There are also many example scripts that you can download, run and edit that showcase some use-cases of NetKet, although they are not commented.
 
 If you want to get in touch with us, feel free to open an issue or a discussion here on GitHub, or to join the MLQuantum slack group where several people involved with NetKet hang out. To join the slack channel just accept [this invitation](https://join.slack.com/t/mlquantum/shared_invite/zt-19wibmfdv-LLRI6i43wrLev6oQX0OfOw)

{netket-3.9.1 → netket-3.10}/netket/__init__.py

@@ -12,18 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# enable x64 on jax
-# must be done at startup.
-from jax.config import config
-
-config.update("jax_enable_x64", True)
-del config
-
 from ._version import version as __version__  # noqa: F401
 
-from . import utils
 from .utils import config
 
+from . import utils
 from . import errors
 
 __all__ = [

netket-3.10/netket/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '3.10'
+__version_tuple__ = version_tuple = (3, 10)

{netket-3.9.1 → netket-3.10}/netket/driver/abstract_variational_driver.py

@@ -22,6 +22,7 @@ import jax
 from jax.tree_util import tree_map
 
 from netket.logging import JsonLog
+from netket.operator import AbstractOperator
 from netket.utils import mpi
 
 
@@ -108,19 +109,18 @@ class AbstractVariationalDriver(abc.ABC):
     def reset(self):
         """
         Resets the driver.
-        Concrete drivers should also call super().reset() to ensure that the step
-        count is set to 0.
+
+        Subclasses should make sure to call :code:`super().reset()` to ensure
+        that the step count is set to 0.
         """
         self.state.reset()
         self._step_count = 0
-        pass
 
     @abc.abstractmethod
     def info(self, depth=0):
         """
         Returns an info string used to print information to screen about this driver.
         """
-        pass  # pragma: no cover
 
     @property
     def state(self):
@@ -156,9 +156,10 @@ class AbstractVariationalDriver(abc.ABC):
         after every `step_size` steps.
 
         Args:
-            n_iter: The total number of steps to perform.
-            step_size: The number of internal steps the simulation
-                is advanced every turn.
+            n_steps: The total number of steps to perform (this is
+                equivalent to the length of the iterator)
+            step: The number of internal steps the simulation
+                is advanced between yielding from the iterator
 
         Yields:
             int: The current step.
@@ -176,7 +177,9 @@ class AbstractVariationalDriver(abc.ABC):
         """
         Performs `steps` optimization steps.
 
-        steps: (Default=1) number of steps
+        Args:
+            steps: (Default=1) number of steps.
+
         """
         for _ in self.iter(steps):
             pass
@@ -245,12 +248,15 @@ class AbstractVariationalDriver(abc.ABC):
         callbacks = _to_iterable(callback)
         callback_stop = False
 
-        with tqdm(total=n_iter, disable=not show_progress) as pbar:
+        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)
 
@@ -302,7 +308,14 @@ class AbstractVariationalDriver(abc.ABC):
             A pytree of the same structure as the input, containing MCMC statistics
             for the corresponding operators as leaves.
         """
-        return tree_map(self._estimate_stats, observables)
+
+        # Do not unpack operators, even if they are pytrees!
+        # this is necessary to support jax operators.
+        return tree_map(
+            self._estimate_stats,
+            observables,
+            is_leaf=lambda x: isinstance(x, AbstractOperator),
+        )
 
     def update_parameters(self, dp):
         """
@@ -329,7 +342,6 @@ class AbstractVariationalDriver(abc.ABC):
         Returns:
             Nothing. The log dictionary should be modified in place.
         """
-        pass  # pragma: no cover
 
 
 @partial(jax.jit, static_argnums=0)

{netket-3.9.1 → netket-3.10}/netket/driver/steady_state.py

@@ -119,8 +119,8 @@ class SteadyState(AbstractVariationalDriver):
         argument is the step, used to change some parameters along the
         optimisation.
 
-        Often, this is taken to be :func:`nk.optimizer.SR`. If it is set to
-        `None`, then the identity is used.
+        Often, this is taken to be :func:`~netket.optimizer.SR`. If it is
+        set to `None`, then the identity is used.
         """
         return self._preconditioner
 
@@ -158,4 +158,4 @@ class SteadyState(AbstractVariationalDriver):
                 ("SR solver   ", self.sr),
             ]
         ]
-        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self)] + lines)
+        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self), *lines])

{netket-3.9.1 → netket-3.10}/netket/driver/vmc.py

@@ -103,8 +103,8 @@ class VMC(AbstractVariationalDriver):
         argument is the step, used to change some parameters along the
         optimisation.
 
-        Often, this is taken to be :func:`nk.optimizer.SR`. If it is set to
-        `None`, then the identity is used.
+        Often, this is taken to be :func:`~netket.optimizer.SR`. If it is
+        set to `None`, then the identity is used.
         """
         return self._preconditioner
 
@@ -169,4 +169,4 @@ class VMC(AbstractVariationalDriver):
                 ("State          ", self.state),
             ]
         ]
-        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self)] + lines)
+        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self), *lines])

{netket-3.9.1 → netket-3.10}/netket/errors.py

@@ -599,6 +599,29 @@ class RealQGTComplexDomainError(Exception):
         )
 
 
+class UnoptimalSRtWarning(NetketWarning):
+    """
+    SRt should be used when the number of parameters exceed the number of samples.
+    If this is not the case, employing `netket.driver.VMC` with the `nk.optimizer.SR` preconditioner
+    is a more efficient option while maintaining the same parameter dynamics.
+
+    .. note::
+        a  detailed discussion can be found in the documentation of
+        :func:`netket.experimental.driver.VMC_SRt`).
+
+    """
+
+    def __init__(self, n_parameters, n_samples):
+        super().__init__(
+            f"""
+            You are in the case n_samples > num_params ({n_samples} > {n_parameters}),
+            for which the `VMC_SRt` is not optimal. Consider using `netket.driver.VMC`
+            with the preconditioner `nk.optimizer.SR` to achieve the same parameter dynamics,
+            but with improved speed.
+            """
+        )
+
+
 #################################################
 # Functions to throw errors                     #
 #################################################

{netket-3.9.1 → netket-3.10}/netket/exact.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
+
 import numpy as _np
 from scipy.sparse.linalg import bicgstab as _bicgstab
 from scipy.sparse.linalg import LinearOperator as _LinearOperator
@@ -26,7 +28,7 @@ def lanczos_ed(
     k: int = 1,
     compute_eigenvectors: bool = False,
     matrix_free: bool = False,
-    scipy_args: dict = None,
+    scipy_args: Optional[dict] = None,
 ):
     r"""Computes `first_n` smallest eigenvalues and, optionally, eigenvectors
     of a Hermitian operator using :meth:`scipy.sparse.linalg.eigsh`.

{netket-3.9.1 → netket-3.10}/netket/experimental/__init__.py

@@ -18,9 +18,11 @@ __all__ = [
     "sampler",
     "vqs",
     "TDVP",
+    "models",
     "hilbert",
     "operator",
     "logging",
+    "observable",
 ]
 
 from . import hilbert
@@ -28,9 +30,11 @@ from . import operator
 from . import driver
 from . import dynamics
 from . import sampler
+from . import models
 from . import vqs
 from . import logging
 from . import qsr
+from . import observable
 
 from .driver import TDVP
 from .qsr import QSR

{netket-3.9.1 → netket-3.10}/netket/experimental/driver/__init__.py

@@ -14,7 +14,7 @@
 
 from .tdvp import TDVP
 from .tdvp_schmitt import TDVPSchmitt
-
+from .vmc_srt import VMC_SRt
 from netket.utils import _hide_submodules
 
 _hide_submodules(__name__)

{netket-3.9.1 → netket-3.10}/netket/experimental/driver/tdvp.py

@@ -135,7 +135,7 @@ class TDVP(TDVPBaseDriver):
                 ("state         ", self.state),
             ]
         ]
-        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self)] + lines)
+        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self), *lines])
 
 
 @odefun.dispatch
@@ -179,7 +179,6 @@ def odefun_tdvp(  # noqa: F811
 
 @partial(jax.jit, static_argnums=(3, 4))
 def _map_parameters(forces, parameters, loss_grad_factor, propagation_type, state_T):
-
     forces = jax.tree_map(
         lambda x, target: loss_grad_factor * x,
         forces,

{netket-3.9.1 → netket-3.10}/netket/experimental/driver/tdvp_common.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Optional, Sequence, Union
+from typing import Callable, Optional, Union
+from collections.abc import Sequence
 import warnings
 
 import jax
@@ -205,6 +206,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
                 at which this method will stop and yield. By default, a stop is performed
                 after each time step (at potentially varying step size if an adaptive
                 integrator is used).
+
         Yields:
             The current step count.
         """
@@ -214,7 +216,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
         self,
         T: float,
         tstops: Optional[Sequence[float]] = None,
-        callback: Callable = None,
+        callback: Optional[Callable] = None,
     ):
         """
         Implementation of :code:`iter`. This method accepts and additional `callback` object, which
@@ -328,6 +330,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
             total=t_end,
             disable=not show_progress,
             unit_scale=True,
+            dynamic_ncols=True,
         ) as pbar:
             first_step = True
 

{netket-3.9.1 → netket-3.10}/netket/experimental/driver/tdvp_schmitt.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Union
+from typing import Callable, Union, Optional
 
 from functools import partial
 
@@ -108,9 +108,9 @@ class TDVPSchmitt(TDVPBaseDriver):
         *,
         t0: float = 0.0,
         propagation_type="real",
-        holomorphic: bool = None,
+        holomorphic: Optional[bool] = None,
         diag_shift: float = 0.0,
-        diag_scale: float = None,
+        diag_scale: Optional[float] = None,
         error_norm: Union[str, Callable] = "qgt",
         rcond: float = 1e-14,
         rcond_smooth: float = 1e-8,
@@ -192,7 +192,7 @@ class TDVPSchmitt(TDVPBaseDriver):
                 ("state         ", self.state),
             ]
         ]
-        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self)] + lines)
+        return "\n{}".format(" " * 3 * (depth + 1)).join([str(self), *lines])
 
 
 # Copyright notice:
@@ -204,7 +204,7 @@ class TDVPSchmitt(TDVPBaseDriver):
 @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)
-    ΔE_loc = E_loc.T.reshape(-1, 1) - E.mean
+    ΔE_loc = E_loc.reshape(-1, 1) - E.mean
 
     stack_jacobian = S.mode == "complex"
 
@@ -295,7 +295,6 @@ def odefun_schmitt(state: MCState, self: TDVPSchmitt, t, w, *, stage=0):  # noqa
 
 @partial(jax.jit, static_argnums=(3, 4))
 def _map_parameters(forces, parameters, loss_grad_factor, propagation_type, state_T):
-
     forces = jax.tree_map(
         lambda x, target: loss_grad_factor * x,
         forces,