NetKet 3.18__tar.gz → 3.19.0__tar.gz

{netket-3.18 → netket-3.19.0}/.gitignore

@@ -1,5 +1,9 @@
 netket/_version.py
 
+# AI stuff #
+############
+CLAUDE.md
+
 # Compiled source #
 ###################
 *.com

{netket-3.18 → netket-3.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.18
+Version: 3.19.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/#
@@ -18,10 +18,11 @@ Classifier: Operating System :: Unix
 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: flax>=0.10.2
 Requires-Dist: igraph>=0.10.0
-Requires-Dist: jax<0.6,>=0.3.37
+Requires-Dist: jax<0.6,>=0.4.35
 Requires-Dist: numba>=0.59
 Requires-Dist: numpy<3.0,>=1.20
 Requires-Dist: optax<0.3,>=0.2.2
@@ -31,8 +32,10 @@ Requires-Dist: rich>=12.0
 Requires-Dist: scipy<2,>=1.5.3
 Requires-Dist: sparse>=0.12.0
 Requires-Dist: tqdm<5,>=4.60
+Provides-Extra: cuda
+Requires-Dist: jax[cuda12]; (sys_platform == 'linux') and extra == 'cuda'
 Provides-Extra: dev
-Requires-Dist: black==24.1.0; extra == 'dev'
+Requires-Dist: black==25.1.0; extra == 'dev'
 Requires-Dist: build>=1; extra == 'dev'
 Requires-Dist: coverage>=5; extra == 'dev'
 Requires-Dist: matplotlib>=3; extra == 'dev'
@@ -43,13 +46,16 @@ 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.9.4; extra == 'dev'
+Requires-Dist: ruff==0.12.2; extra == 'dev'
 Requires-Dist: wheel>=0.38.1; extra == 'dev'
 Provides-Extra: docs
-Requires-Dist: myst-nb<1.3,>=1.0; extra == 'docs'
+Requires-Dist: ipywidgets>=8.1.7; extra == 'docs'
+Requires-Dist: myst-nb<1.4,>=1.0; extra == 'docs'
 Requires-Dist: myst-parser<4.1.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.1.0; extra == 'docs'
+Requires-Dist: sphinx-copybutton>=0.5.2; extra == 'docs'
 Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
@@ -59,9 +65,6 @@ 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.3.9; extra == 'mpi'
-Requires-Dist: mpi4py<5,>=3.0.1; extra == 'mpi'
 Provides-Extra: pyscf
 Requires-Dist: pyscf>=2.0; extra == 'pyscf'
 Description-Content-Type: text/markdown
@@ -78,8 +81,7 @@ Description-Content-Type: text/markdown
 [![codecov](https://codecov.io/gh/netket/netket/branch/master/graph/badge.svg?token=gzcOlpO5lB)](https://codecov.io/gh/netket/netket)
 [![Slack](https://img.shields.io/badge/slack-chat-green.svg)](https://join.slack.com/t/mlquantum/shared_invite/zt-19wibmfdv-LLRI6i43wrLev6oQX0OfOw)
 
-NetKet is an open-source project delivering cutting-edge methods for the study
-of many-body quantum systems with artificial neural networks and machine learning techniques.
+NetKet is an open-source project delivering cutting-edge methods for the study of many-body quantum systems with artificial neural networks and machine learning techniques.
 It is a Python library built on [JAX](https://github.com/google/jax).
 
 NetKet is an affiliated project to [numFOCUS](https://numfocus.org).
@@ -93,55 +95,24 @@ NetKet is an affiliated project to [numFOCUS](https://numfocus.org).
 
 ## Installation and Usage
 
-NetKet runs on MacOS and Linux. We recommend to install NetKet using `pip`, but it can also be installed with `conda`.
-It is often necessary to first update `pip` to a recent release (`>=20.3`) in order for upper compatibility bounds to be considered and avoid a broken installation.
-For instructions on how to install the latest stable/beta release of NetKet see the [Get Started](https://www.netket.org/get_started/) page of our website or run the following command (Apple M1 users, follow that link for more instructions):
+NetKet runs on MacOS and Linux and requires Python 3.11 or later. We recommend installing NetKet using `pip` or `uv`. **Do not use conda** as JAX has known issues when installed through conda.
 
 ```sh
 pip install --upgrade pip
-pip install --upgrade netket
+pip install netket
 ```
 
-If you wish to install the current development version of NetKet, which is the master branch of this GitHub repository, together with the additional dependencies, you can run the following command:
-
+**With GPU support (Linux only):**
 ```sh
-pip install --upgrade pip
-pip install 'git+https://github.com/netket/netket.git#egg=netket[all]'
+pip install 'netket[cuda]'
 ```
 
-To speed-up NetKet-computations, even on a single machine, you
-can install the MPI-related dependencies by using `[mpi]` between square brackets.
-
+**Development version:**
 ```sh
-pip install --upgrade pip
-pip install --upgrade "netket[mpi]"
+pip install git+https://github.com/netket/netket.git
 ```
 
-We recommend to install NetKet with all it's extra dependencies, which are documented below.
-However, if you do not have a working MPI compiler in your PATH this installation will most likely fail because
-it will attempt to install `mpi4py`, which enables MPI support in netket.
-
-The latest release of NetKet is always available on PyPi and can be installed with `pip`.
-NetKet is also available on conda-forge, however the version available through `conda install`
-can be slightly out of date compared to PyPi.
-To check what is the latest version released on both distributions you can inspect the badges at the top of this readme.
-
-### Extra dependencies
-When installing `netket` with pip, you can pass the following extra variants as square brakets. You can install several of them by separating them with a comma.
- - `"[dev]"`: installs development-related dependencies such as black, pytest and testing dependencies
- - `"[mpi]"`: Installs `mpi4py` to enable multi-process parallelism. Requires a working MPI compiler in your path
- - `"[extra]"`: Installs `tensorboardx` to enable logging to tensorboard, and openfermion to convert the QubitOperators.
- - `"[all]"`: Installs all extra dependencies
-
-### MPI Support
-To enable MPI support you must install [mpi4jax](https://github.com/PhilipVinc/mpi4jax). Please note that we advise to install mpi4jax  with the same tool (conda or pip) with which you install it's dependency `mpi4py`.
-
-To check whether MPI support is enabled, check the flags
-```python
->>> import netket
->>> netket.utils.mpi.available
-True
-```
+For detailed installation instructions including GPU setup, see our [installation guide](https://netket.readthedocs.io/en/latest/install.html).
 
 ## Getting Started
 

netket-3.19.0/README.md

@@ -0,0 +1,55 @@
+<div align="center">
+<img src="https://www.netket.org/logo/logo_simple.jpg" alt="logo" width="400"></img>
+</div>
+
+# __NetKet__
+
+[![Powered by NumFOCUS](https://img.shields.io/badge/powered%20by-NumFOCUS-orange.svg?style=flat&colorA=E1523D&colorB=007D8A)](https://numfocus.org)
+[![Release](https://img.shields.io/github/release/netket/netket.svg)](https://github.com/netket/netket/releases)
+[![Paper (v3)](https://img.shields.io/badge/paper%20%28v3%29-arXiv%3A2112.10526-B31B1B)](https://scipost.org/SciPostPhysCodeb.7/pdf)
+[![codecov](https://codecov.io/gh/netket/netket/branch/master/graph/badge.svg?token=gzcOlpO5lB)](https://codecov.io/gh/netket/netket)
+[![Slack](https://img.shields.io/badge/slack-chat-green.svg)](https://join.slack.com/t/mlquantum/shared_invite/zt-19wibmfdv-LLRI6i43wrLev6oQX0OfOw)
+
+NetKet is an open-source project delivering cutting-edge methods for the study of many-body quantum systems with artificial neural networks and machine learning techniques.
+It is a Python library built on [JAX](https://github.com/google/jax).
+
+NetKet is an affiliated project to [numFOCUS](https://numfocus.org).
+
+- **Homepage:** <https://www.netket.org>
+- **Citing:** <https://www.netket.org/cite/>
+- **Documentation:** <https://netket.readthedocs.io/en/latest/index.html>
+- **Tutorials:** <https://netket.readthedocs.io/en/latest/tutorials/gs-ising.html>
+- **Examples:** <https://github.com/netket/netket/tree/master/Examples>
+- **Source code:** <https://github.com/netket/netket>
+
+## Installation and Usage
+
+NetKet runs on MacOS and Linux and requires Python 3.11 or later. We recommend installing NetKet using `pip` or `uv`. **Do not use conda** as JAX has known issues when installed through conda.
+
+```sh
+pip install --upgrade pip
+pip install netket
+```
+
+**With GPU support (Linux only):**
+```sh
+pip install 'netket[cuda]'
+```
+
+**Development version:**
+```sh
+pip install git+https://github.com/netket/netket.git
+```
+
+For detailed installation instructions including GPU setup, see our [installation guide](https://netket.readthedocs.io/en/latest/install.html).
+
+## 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.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)
+
+## License
+
+[Apache License 2.0](https://github.com/netket/netket/blob/master/LICENSE)

netket-3.19.0/netket/_src/ngd/kwargs.py

@@ -0,0 +1,42 @@
+import inspect
+
+from collections.abc import Callable
+from functools import partial
+
+
+def ensure_accepts_kwargs(f: Callable, kwarg_name: str):
+    """
+    Add a keyword argument to a function if it is not already present.
+
+    This is useful to make a function compatible with a certain interface, even if it does not
+    need the additional argument.
+
+    In this case, we use it to make the linear solver functions compatible with the QGT interface.
+    Some QGT solvers require the vector of local energies as an additional argument. The standard
+    solver interface however only requires the matrix and the right-hand side (jacobian @ local_energies).
+
+    Args:
+        f: The function to concretize.
+        kwarg_name: The name of the keyword argument to add.
+
+    Returns:
+        A new function that accepts the keyword argument `kwarg_name`.
+    """
+    sig = inspect.signature(f)
+    params = sig.parameters
+
+    if kwarg_name in params:
+        return f
+    else:
+        to_concretize = f
+        if isinstance(f, partial):
+            to_concretize = f.func
+
+        def concretized_func(*args, **kwargs):
+            kwargs.pop(kwarg_name, None)
+            return to_concretize(*args, **kwargs)
+
+        if isinstance(f, partial):
+            concretized_func = partial(concretized_func, *f.args, **f.keywords)
+
+        return concretized_func

netket-3.19.0/netket/_src/ngd/sr.py

@@ -0,0 +1,70 @@
+from collections.abc import Callable
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax.sharding import PositionalSharding
+
+from netket import jax as nkjax
+from netket.utils.types import Array
+
+from netket._src.ngd.kwargs import ensure_accepts_kwargs
+
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "solver_fn",
+        "mode",
+    ),
+)
+def _compute_sr_update(
+    O_L,
+    dv,
+    *,
+    diag_shift: float | Array,
+    solver_fn: Callable[[Array, Array], Array],
+    mode: str,
+    proj_reg: float | Array | None = None,
+    momentum: float | Array | None = None,
+    old_updates: Array | None = None,
+    params_structure,
+):
+    # We concretize the solver function to ensure it accepts the additional argument `dv`.
+    # Typically solvers only accept the matrix and the right-hand side.
+    solver_fn = ensure_accepts_kwargs(solver_fn, "dv")
+
+    if (momentum is not None) or (old_updates is not None) or (proj_reg is not None):
+        raise ValueError("Not implemented")
+
+    # (np, #ns) x (#ns) -> (np) - where the sum over #ns is done automatically
+    F = O_L.T @ dv
+
+    # This does the contraction (np, #ns) x (#ns, np) -> (np, np).
+    matrix = O_L.T @ O_L
+    matrix_side = matrix.shape[-1]  # * it can be ns or 2*ns, depending on mode
+
+    shifted_matrix = jax.lax.add(
+        matrix, diag_shift * jnp.eye(matrix_side, dtype=matrix.dtype)
+    )
+    updates = solver_fn(shifted_matrix, F, dv=dv)
+
+    # Some solvers return a tuple, some others do not.
+    if isinstance(updates, tuple):
+        updates, info = updates
+        if info is None:
+            info = {}
+    else:
+        info = {}
+
+    # If complex mode and we have complex parameters, we need
+    # To repack the real coefficients in order to get complex updates
+    if mode == "complex" and nkjax.tree_leaf_iscomplex(params_structure):
+        num_p = updates.shape[-1] // 2
+        updates = updates[:num_p] + 1j * updates[num_p:]
+
+    updates = jax.lax.with_sharding_constraint(
+        updates, PositionalSharding(jax.devices()).replicate()
+    )
+
+    return updates, old_updates, info

netket-3.19.0/netket/_src/ngd/sr_srt_common.py

@@ -0,0 +1,153 @@
+from collections.abc import Callable
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax.flatten_util import ravel_pytree
+
+import netket.jax as nkjax
+from netket.utils import timing
+from netket.utils.types import Array, PyTree
+
+from netket._src.ngd.sr import _compute_sr_update
+from netket._src.ngd.srt import _compute_srt_update
+
+
+@partial(jax.jit, static_argnames=("mode",))
+def _prepare_input(
+    O_L,
+    local_grad,
+    *,
+    mode: str,
+) -> tuple[jax.Array, jax.Array]:
+    r"""
+    Prepare the input for the SR/SRt solvers.
+
+    The local eneriges and the jacobian are reshaped, centered and normalized by the number of Monte Carlo samples.
+    The complex case is handled by concatenating the real and imaginary parts of the jacobian and the local energies.
+
+    We use [Re_x1, Im_x1, Re_x2, Im_x2, ...] so that shards are contiguous, and jax can keep track of the sharding information.
+    This format is applied both to the jacobian and to the vector.
+
+    Args:
+        O_L: The jacobian of the ansatz.
+        local_grad: The local energies.
+        mode: The mode of the jacobian: `'real'` or `'complex'`.
+
+    Returns:
+        The reshaped jacobian and the reshaped local energies.
+    """
+    N_mc = O_L.shape[0]
+
+    local_grad = local_grad.flatten()
+    de = local_grad - jnp.mean(local_grad)
+
+    O_L = O_L / jnp.sqrt(N_mc)
+    dv = 2.0 * de / jnp.sqrt(N_mc)
+
+    if mode == "complex":
+        # Concatenate the real and imaginary derivatives of the ansatz
+        # (#ns, 2, np) -> (#ns*2, np)
+        O_L = jax.lax.collapse(O_L, 0, 2)
+
+        # (#ns, 2) -> (#ns*2)
+        dv2 = jnp.stack([jnp.real(dv), jnp.imag(dv)], axis=-1)
+        dv = jax.lax.collapse(dv2, 0, 2)
+    elif mode == "real":
+        dv = dv.real
+    else:
+        raise NotImplementedError()
+    return O_L, dv
+
+
+@timing.timed
+@partial(
+    jax.jit,
+    static_argnames=(
+        "log_psi",
+        "solver_fn",
+        "mode",
+        "chunk_size",
+        "use_ntk",
+    ),
+)
+def _sr_srt_common(
+    log_psi,
+    local_grad,
+    parameters,
+    model_state,
+    samples,
+    *,
+    diag_shift: float | Array,
+    solver_fn: Callable[[Array, Array], Array],
+    mode: str,
+    proj_reg: float | Array | None = None,
+    momentum: float | Array | None = None,
+    old_updates: PyTree | None = None,
+    chunk_size: int | None = None,
+    use_ntk: bool = False,
+):
+    r"""
+    Compute the SR/Natural gradient update for the model specified by
+    `log_psi({parameters, model_state}, samples)` and the local gradient contributions `local_grad`.
+
+    Uses a code equivalent to QGTJacobianDense by default, or with the NTK/MinSR if `use_ntk` is True.
+
+    Args:
+        log_psi: The log of the wavefunction.
+        local_grad: The local values of the estimator.
+        parameters: The parameters of the model.
+        model_state: The state of the model.
+        diag_shift: The diagonal shift of the stochastic reconfiguration matrix. Typical values are 1e-4 ÷ 1e-3. Can also be an optax schedule.
+        proj_reg: Weight before the matrix `1/N_samples \\bm{1} \\bm{1}^T` used to regularize the linear solver in SPRING.
+        momentum: Momentum used to accumulate updates in SPRING.
+        linear_solver_fn: Callable to solve the linear problem associated to the updates of the parameters.
+        mode: The mode used to compute the jacobian of the variational state. Can be `'real'` or `'complex'` (defaults to the dtype of the output of the model).
+
+    Returns:
+        The new parameters, the old updates, and the info dictionary.
+    """
+    _, unravel_params_fn = ravel_pytree(parameters)
+    _params_structure = jax.tree_util.tree_map(
+        lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), parameters
+    )
+
+    jacobians = nkjax.jacobian(
+        log_psi,
+        parameters,
+        samples,
+        model_state,
+        mode=mode,
+        dense=True,
+        center=True,
+        chunk_size=chunk_size,
+    )  # jacobian is centered
+
+    O_L, dv = _prepare_input(jacobians, local_grad, mode=mode)
+
+    if old_updates is None and momentum is not None:
+        old_updates = jnp.zeros(jacobians.shape[-1], dtype=jacobians.dtype)
+
+    compute_update = _compute_srt_update if use_ntk else _compute_sr_update
+
+    # TODO: Add support for proj_reg and momentum
+    # At the moment SR does not support momentum, proj_reg.
+    # We raise an error if they are passed with a value different from None.
+    updates, old_updates, info = compute_update(
+        O_L,
+        dv,
+        diag_shift=diag_shift,
+        solver_fn=solver_fn,
+        mode=mode,
+        proj_reg=proj_reg,
+        momentum=momentum,
+        old_updates=old_updates,
+        params_structure=_params_structure,
+    )
+
+    return unravel_params_fn(updates), old_updates, info
+
+
+sr = partial(_sr_srt_common, use_ntk=False)
+
+srt = partial(_sr_srt_common, use_ntk=True)

netket-3.19.0/netket/_src/ngd/srt.py

@@ -0,0 +1,93 @@
+from collections.abc import Callable
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax.sharding import PositionalSharding
+
+from netket import jax as nkjax
+from netket import config
+from netket.utils import timing
+from netket.utils.types import Array
+
+
+@timing.timed
+@partial(
+    jax.jit,
+    static_argnames=(
+        "solver_fn",
+        "mode",
+    ),
+)
+def _compute_srt_update(
+    O_L,
+    dv,
+    *,
+    diag_shift: float | Array,
+    solver_fn: Callable[[Array, Array], Array],
+    mode: str,
+    proj_reg: float | Array | None = None,
+    momentum: float | Array | None = None,
+    old_updates: Array | None = None,
+    params_structure,
+):
+    if momentum is not None:
+        dv -= momentum * (O_L @ old_updates)
+
+    # (#ns, np) -> (ns, #np)
+    O_LT = O_L
+    if config.netket_experimental_sharding:
+        nkjax.sharding.pad_axis_for_sharding(O_LT, axis=1, padding_value=0.0)
+        O_LT = jax.lax.with_sharding_constraint(
+            O_LT,
+            PositionalSharding(jax.devices()).reshape(1, -1),
+        )
+        dv = jax.lax.with_sharding_constraint(
+            dv, PositionalSharding(jax.devices()).replicate()
+        )
+
+    # This does the contraction (ns, #np) x (#np, ns) -> (ns, ns).
+    # When using sharding the sum over #ns is done automatically.
+    # When using MPI we need to do it manually with an allreduce_sum.
+    matrix = O_LT @ O_LT.T
+    matrix_side = matrix.shape[-1]  # * it can be ns or 2*ns, depending on mode
+
+    shifted_matrix = jax.lax.add(
+        matrix, diag_shift * jnp.eye(matrix_side, dtype=matrix.dtype)
+    )
+    # replicate
+
+    if proj_reg is not None:
+        shifted_matrix = jax.lax.add(
+            shifted_matrix, jnp.full_like(shifted_matrix, proj_reg / matrix_side)
+        )
+
+    aus_vector = solver_fn(shifted_matrix, dv)
+
+    # Some solvers return a tuple, some others do not.
+    if isinstance(aus_vector, tuple):
+        aus_vector, info = aus_vector
+        if info is None:
+            info = {}
+    else:
+        info = {}
+
+    # (np, #ns) x (#ns) -> (np).
+    updates = O_L.T @ aus_vector
+    if momentum is not None:
+        updates += momentum * old_updates
+        old_updates = updates
+
+    # If complex mode and we have complex parameters, we need
+    # To repack the real coefficients in order to get complex updates
+    if mode == "complex" and nkjax.tree_leaf_iscomplex(params_structure):
+        num_p = updates.shape[-1] // 2
+        updates = updates[:num_p] + 1j * updates[num_p:]
+
+    if config.netket_experimental_sharding:
+        out_shardings = (
+            PositionalSharding(jax.devices()).replicate().reshape((1,) * updates.ndim)
+        ).replicate()
+        updates = jax.lax.with_sharding_constraint(updates, out_shardings)
+
+    return updates, old_updates, info