NetKet 3.17__tar.gz → 3.18__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.17
+Version: 3.18
 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,7 +17,7 @@ Classifier: Operating System :: POSIX :: Linux
 Classifier: Operating System :: Unix
 Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Physics
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 Requires-Dist: equinox>=0.11.3
 Requires-Dist: flax>=0.10.2
 Requires-Dist: igraph>=0.10.0
@@ -29,6 +29,7 @@ Requires-Dist: orjson<4,>=3.4
 Requires-Dist: plum-dispatch>=2.4
 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: dev
 Requires-Dist: black==24.1.0; extra == 'dev'
@@ -50,9 +51,9 @@ Requires-Dist: myst-parser<4.1.0,>=2.0.0; 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-design>=0.4.0; extra == 'docs'
-Requires-Dist: sphinx<8.2,>=7.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
+Requires-Dist: sphinx~=8.2.3; extra == 'docs'
 Provides-Extra: extra
 Requires-Dist: h5py>=3.7.0; extra == 'extra'
 Requires-Dist: openfermion>=1.0.0; extra == 'extra'
@@ -63,7 +64,6 @@ 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'
-Requires-Dist: sparse>=0.12.0; extra == 'pyscf'
 Description-Content-Type: text/markdown
 
 <div align="center">

{netket-3.17 → netket-3.18}/netket/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.17'
-__version_tuple__ = version_tuple = (3, 17)
+__version__ = version = '3.18'
+__version_tuple__ = version_tuple = (3, 18)

{netket-3.17 → netket-3.18}/netket/experimental/__init__.py

@@ -15,6 +15,7 @@
 __all__ = [
     "driver",
     "dynamics",
+    "geometry",
     "sampler",
     "vqs",
     "TDVP",
@@ -23,18 +24,20 @@ __all__ = [
     "operator",
     "logging",
     "observable",
+    "QSR",
 ]
 
+from . import geometry
 from . import hilbert
 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 . import models
 
 from .driver import TDVP
 from .qsr import QSR

netket-3.18/netket/experimental/geometry/__init__.py

@@ -0,0 +1,8 @@
+"""Geometry classes for continuous Hilbert spaces."""
+
+__all__ = [
+    "Cell",
+    "FreeSpace",
+]
+
+from .cell import Cell, FreeSpace

netket-3.18/netket/experimental/geometry/cell.py

@@ -0,0 +1,88 @@
+from dataclasses import dataclass
+from typing import Sequence
+
+import jax.numpy as jnp
+import numpy as np
+from jax.tree_util import register_pytree_node_class
+
+
+@register_pytree_node_class
+@dataclass(frozen=True)
+class Cell:
+    """A finite region in continuous space."""
+
+    extent: tuple[float, ...]
+    pbc: tuple[bool, ...]
+
+    def __init__(
+        self,
+        d: int,
+        L: float | Sequence[float] = 1.0,
+        pbc: bool | Sequence[bool] = True,
+    ) -> None:
+        if d <= 0:
+            raise ValueError("d must be positive")
+
+        if hasattr(L, "__len__"):
+            if len(L) != d:
+                raise ValueError("Length of L must match d")
+            extent = tuple(float(x) for x in L)
+        else:
+            extent = (float(L),) * d
+
+        if isinstance(pbc, bool):
+            pbc_tuple = (pbc,) * d
+        else:
+            if len(pbc) != d:
+                raise ValueError("Length of pbc must match d")
+            pbc_tuple = tuple(bool(x) for x in pbc)
+
+        for e, b in zip(extent, pbc_tuple):
+            if np.isinf(e) and b:
+                raise ValueError(
+                    "Cannot combine periodic boundary conditions with infinite extent"
+                )
+
+        object.__setattr__(self, "extent", extent)
+        object.__setattr__(self, "pbc", pbc_tuple)
+
+    def tree_flatten(self):
+        data = (jnp.asarray(self.extent), jnp.asarray(self.pbc))
+        return data, None
+
+    @classmethod
+    def tree_unflatten(cls, aux, data):
+        extent, pbc = data
+        return cls(
+            d=len(extent),
+            L=tuple(float(x) for x in extent),
+            pbc=tuple(bool(x) for x in pbc),
+        )
+
+    @property
+    def dimension(self) -> int:
+        return len(self.extent)
+
+    def distance(
+        self, r1: Sequence[float] | jnp.ndarray, r2: Sequence[float] | jnp.ndarray
+    ) -> jnp.ndarray:
+        """Euclidean distance between ``r1`` and ``r2`` considering boundary conditions."""
+
+        r1 = jnp.asarray(r1)
+        r2 = jnp.asarray(r2)
+        diff = r1 - r2
+        ext = jnp.asarray(self.extent)
+        pbc = jnp.asarray(self.pbc)
+
+        shift = jnp.where(pbc, jnp.round(diff / ext) * ext, 0.0)
+        diff = diff - shift
+        return jnp.linalg.norm(diff, axis=-1)
+
+
+@register_pytree_node_class
+@dataclass(frozen=True)
+class FreeSpace(Cell):
+    """Infinite space without periodic boundary conditions."""
+
+    def __init__(self, d: int) -> None:
+        super().__init__(d=d, L=(np.inf,) * d, pbc=False)

{netket-3.17 → netket-3.18}/netket/experimental/hilbert/__init__.py

@@ -13,11 +13,13 @@
 # limitations under the License.
 
 
-__all__ = ["SpinOrbitalFermions"]
+__all__ = ["ContinuousHilbert", "Particle"]
 
+from .continuous_hilbert import ContinuousHilbert
+from .particle import Particle
+from . import random
 from netket.hilbert import SpinOrbitalFermions as _deprecated_SpinOrbitalFermions
 
-
 _deprecations = {
     # May 2024
     "SpinOrbitalFermions": (
@@ -27,11 +29,10 @@ _deprecations = {
     ),
 }
 
+from netket.utils import _auto_export
 from netket.utils.deprecation import deprecation_getattr as _deprecation_getattr
 
 __getattr__ = _deprecation_getattr(__name__, _deprecations)
-del _deprecation_getattr
-
-from netket.utils import _hide_submodules
+_auto_export(__name__)
 
-_hide_submodules(__name__)
+del _deprecation_getattr, _auto_export

netket-3.18/netket/experimental/hilbert/continuous_hilbert.py

@@ -0,0 +1,45 @@
+# Copyright 2021 The NetKet Authors - All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from netket.hilbert.abstract_hilbert import AbstractHilbert
+
+
+class ContinuousHilbert(AbstractHilbert):
+    """Abstract class for the Hilbert space of particles
+    in continuous space.
+
+    This class defines the common interface that
+    can be used to interact with particles defined
+    in continuous space.
+    """
+
+    def __init__(self, domain: tuple[float, ...]):
+        """
+        Constructs an Hilbert space with continuous degrees of freedom, given specifications of the space.
+
+        This object returns an Hilbert space.
+
+        Args:
+            domain: Tuple indicating the maximum of the continuous quantum number(s) in the configurations. Each entry
+                in the tuple corresponds to a different physical dimension.
+                A particle in a 3D box of size L would take `(L,L,L)`. A rotor model would take e.g. `(2pi,)`.
+        """
+        self._extent = tuple(domain)
+
+        super().__init__()
+
+    @property
+    def domain(self) -> tuple[float, ...]:
+        r"""Domain of the continuous variable, specified for each dimension"""
+        return self._extent

netket-3.18/netket/experimental/hilbert/particle.py

@@ -0,0 +1,85 @@
+# Copyright 2021 The NetKet Authors - All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .continuous_hilbert import ContinuousHilbert
+from netket.experimental.geometry import Cell
+
+
+class Particle(ContinuousHilbert):
+    r"""Hilbert space derived from ContinuousHilbert defining N particles
+    in continuous space with or without periodic boundary conditions."""
+
+    def __init__(
+        self,
+        N: int | tuple[int, ...],
+        *,
+        geometry: Cell,
+    ) -> None:
+        """
+        Constructs new ``Particles`` given specifications
+        of the continuous space they are defined in.
+
+        Args:
+            N: Number of particles. If int all have the same spin. If Tuple the entry indicates how many particles
+                there are with a certain spin-projection.
+            geometry: Instance of :class:`~netket.experimental.geometry.Cell` defining the geometry.
+        """
+
+        if not hasattr(N, "__len__"):
+            N = (N,)
+
+        if not isinstance(geometry, Cell):
+            raise TypeError(
+                "`geometry` must be an instance of `netket.experimental.geometry.Cell`."
+            )
+
+        self._N = sum(N)
+        self._n_per_spin = N
+        self._geometry = geometry
+
+        super().__init__(geometry.extent)
+
+    @property
+    def size(self) -> int:
+        return self._N * len(self.domain)
+
+    @property
+    def n_particles(self) -> int:
+        r"""The number of particles"""
+        return self._N
+
+    @property
+    def n_per_spin(self) -> tuple[int, ...]:
+        r"""Gives the number of particles in a specific spin
+        projection.
+
+        The length of this tuple indicates the total spin whereas
+        the position in the tuple indicates the spin projection.
+
+        Example: (10,5,3) describes 18 particles of total spin 1
+        where 10 of those have spin-projection -1, 5 have
+        spin-projection 0 and 3 have spin-projection 1.
+        """
+        return self._n_per_spin
+
+    @property
+    def geometry(self) -> Cell:
+        """Geometry of the continuous space."""
+        return self._geometry
+
+    @property
+    def _attrs(self):
+        return (self._N, self.geometry)
+
+    def __repr__(self):
+        return f"Particle(N={self.n_particles}, d={len(self.domain)})"

netket-3.18/netket/experimental/hilbert/random/__init__.py

@@ -0,0 +1,75 @@
+# Copyright 2021 The NetKet Authors - All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This module contains functions that act on hilbert spaces necessary
+for sampling.
+
+If you define custom hilbert spaces, and want to sample from it, you
+should read carefully.
+
+This module *exports* two functions:
+ - `random_state(hilbert, key, batch_size, dtype)` which generates a batch of random
+ states in the given hilbert space, in an array with the specified dtype.
+ - `flip_state(hilb, key, states, indices)` which, for every state σ in the batch of
+ states, considers σᵢ and returns a new state where that entry is different from the
+ previous. The new configuration is selected with uniform probability among the local
+ possible configurations.
+
+Hilbert spaces must at least implement a `random_state` to support sampling.
+`flip_state` is only necessary in order to use LocalRule samplers.
+All hilbert spaces in netket implement both.
+
+How to implement the two functions above
+----------------------------------------
+
+While the *exported* function acts on batches, sometimes it is hard to implement
+the function on batches. Therefore they can be implemented either as a scalar
+function, that gets jax.vmap-ed automatically, or directly as a batched rule.
+Of course, if you implement the batched rule you will most likely see better
+performance.
+
+In order to implement the scalar rule for your custom hilbert object `MyHilbert`
+you should define a function taking 4 inputs, the hilbert space, a jax PRNG key and
+the dtype of the desired result. For random operations you should use the key
+provided.
+
+@netket.utils.dispatch.dispatch
+def random_state(hilb: MyHilbert, key, dtype):
+    return mystate
+
+@netket.utils.dispatch.dispatch
+def random_state(hilb: MyHilbert, key, batches: int, dtype):
+    return mystate
+
+flip_state is implemented in the same way, through
+
+@netket.utils.dispatch.dispatch
+def flip_state_scalar(hilb: Fock, key, σ, idx):
+    return new_state, oldval
+
+There is a vmapped default fallback for the batched version.
+
+@netket.utils.dispatch.dispatch
+def flip_state_batch(hilb: Fock, key, σ, idx):
+    return new_states, oldvals
+
+"""
+
+from netket.utils import _hide_submodules
+
+from . import particle
+from netket.hilbert.random.base import flip_state, random_state
+
+_hide_submodules(__name__)

{netket-3.17/netket → netket-3.18/netket/experimental}/hilbert/random/particle.py

@@ -17,7 +17,7 @@ from jax import numpy as jnp
 import numpy as np
 
 from netket import jax as nkjax
-from netket.hilbert import ContinuousHilbert, Particle
+from netket.experimental.hilbert import ContinuousHilbert, Particle
 from netket.utils.dispatch import dispatch
 
 
@@ -38,10 +38,11 @@ def random_state(hilb: Particle, key, batches: int, *, dtype):
     If periodic boundary conditions are chosen only for certain dimensions, the periodic initialization is used for
     all of those dimensions and the free initialization is used for all the other ones.
     """
-    pbc = np.array(hilb.n_particles * hilb.pbc)
+    pbc = np.array(hilb.n_particles * hilb.geometry.pbc)
     boundary = np.tile(pbc, (batches, 1))
 
-    Ls = np.array(hilb.n_particles * hilb.extent)
+    # TODO genearalize to particles with diferent domains
+    Ls = np.array(hilb.n_particles * hilb.domain)
 
     # If we have PBCs this defines the size of the Gaussian noise in the particle positions.
     # Without PBCs this is never needed (initialization with random Gaussian)
@@ -65,9 +66,9 @@ def random_state(hilb: Particle, key, batches: int, *, dtype):
     noise = gaussian * width
 
     key = jax.random.split(key, num=batches)
-    sdim = len(hilb.extent)
+    sdim = len(hilb.domain)
     n = int(np.ceil(hilb.n_particles ** (1 / sdim)))
-    xs = jnp.linspace(0, min(hilb.extent), n)
+    xs = jnp.linspace(0, min(hilb.domain), n)
     uniform = jnp.array(jnp.meshgrid(*(sdim * [xs]))).reshape(-1, sdim)
     uniform = jnp.tile(uniform, (batches, 1, 1))
 

{netket-3.17 → netket-3.18}/netket/experimental/logging/hdf5_log.py

@@ -18,6 +18,7 @@ from flax.serialization import to_bytes
 from flax.core import pop as fpop, FrozenDict
 
 from netket.logging import AbstractLog
+from netket.jax.sharding import extract_replicated
 
 _mode_shorthands = {"write": "w", "append": "a", "fail": "x"}
 
@@ -172,22 +173,27 @@ class HDF5Log(AbstractLog):
         if self._writer is None:
             self._init_output_file()
 
+        if self._steps_notsaved_params % self._save_params_every == 0:
+            variables = variational_state.variables
+
+            # TODO: remove - FrozenDict are deprecated
+            if isinstance(variables, FrozenDict):
+                variables = variables.unfreeze()
+            variables = extract_replicated(variables)
+            self._steps_notsaved_params = 0
+        else:
+            variables = None
+
         if self._is_master_process:
             tree_log(log_data, "data", self._writer, iter=step)
 
-            if self._steps_notsaved_params % self._save_params_every == 0:
-                variables = variational_state.variables
-                # TODO: remove - FrozenDict are deprecated
-                if isinstance(variables, FrozenDict):
-                    variables = variables.unfreeze()
-
+            if variables is not None:
                 _, params = fpop(variables, "params")
                 binary_data = to_bytes(variables)
                 tree = {"model_state": binary_data, "parameters": params, "iter": step}
                 tree_log(tree, "variational_state", self._writer)
-                self._steps_notsaved_params = 0
-
             self._writer.flush()
+
         self._steps_notsaved_params += 1
 
     def flush(self, variational_state=None):

{netket-3.17 → netket-3.18}/netket/experimental/operator/__init__.py

@@ -16,6 +16,11 @@
 __all__ = ["FermionOperator2nd", "FermionOperator2ndJax"]
 
 from .pyscf import from_pyscf_molecule
+from ._particle_number_conserving_fermionic import (
+    ParticleNumberConservingFermioperator2nd,
+    ParticleNumberAndSpinConservingFermioperator2nd,
+    FermiHubbardJax,
+)
 
 
 from . import fermion