quantax 0.2.1__tar.gz → 0.3.0__tar.gz

{quantax-0.2.1/quantax.egg-info → quantax-0.3.0}/PKG-INFO

@@ -1,22 +1,28 @@
 Metadata-Version: 2.4
 Name: quantax
-Version: 0.2.1
+Version: 0.3.0
 Summary: Flexible neural quantum states based on JAX
 Author: Christopher Roth
 Author-email: Ao Chen <chenao.phys@gmail.com>
-Requires-Python: <=3.13,>=3.10
+Requires-Python: <3.14,>=3.11
 License-File: LICENSE
-Requires-Dist: numpy>=1.26.0
-Requires-Dist: scipy>=1.11.1
-Requires-Dist: jax>=0.6.1
+Requires-Dist: numpy>=2.0.0
+Requires-Dist: scipy>=1.14.0
+Requires-Dist: jax>=0.8.0
 Requires-Dist: equinox>=0.11.4
 Requires-Dist: lrux>=0.1.0
 Provides-Extra: full
 Requires-Dist: numpy>=2.0.0; extra == "full"
-Requires-Dist: scipy>=1.13.1; extra == "full"
+Requires-Dist: scipy>=1.14.0; extra == "full"
 Requires-Dist: matplotlib>=3.8.0; extra == "full"
 Requires-Dist: quspin>=1.0.0; extra == "full"
-Requires-Dist: jax>=0.6.1; extra == "full"
+Requires-Dist: jax>=0.8.0; extra == "full"
 Requires-Dist: equinox>=0.11.4; extra == "full"
+Requires-Dist: lineax>=0.1.0; extra == "full"
 Requires-Dist: lrux>=0.1.0; extra == "full"
+Requires-Dist: jaxmg>=0.0.6; extra == "full"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: myst-nb; extra == "docs"
+Requires-Dist: pydata-sphinx-theme; extra == "docs"
 Dynamic: license-file

{quantax-0.2.1 → quantax-0.3.0}/README.md

@@ -1,7 +1,7 @@
 <h1 align='center'>Quantax</h1>
 
 <p align="center">
-  <strong>Flexible and scalable neural quantum states built on JAX</strong>
+  <strong>Flexible and scalable neural quantum states in JAX</strong>
 </p>
 
 <p align="center">
@@ -22,7 +22,7 @@ Quantax is a research-oriented Python package for quantum many-body physics, wit
 
 ## ⚙️Installation
 
-Requires Python 3.10+, JAX 0.6.1+
+Requires Python 3.11+, JAX 0.8.0+
 
 First, ensure that a correct JAX version is installed. For details, check [JAX Installation](https://docs.jax.dev/en/latest/installation.html).
 
@@ -75,3 +75,15 @@ energy_data.plot(baseline=E)
 plt.show()
 ```
 ![quick_start](images/quick_start.svg)
+
+## 📚Citation
+
+```bibtex
+@software{quantax,
+  author  = {Chen, Ao and Roth, Christopher},
+  title   = {Quantax: Flexible neural quantum states in JAX},
+  year    = {2024},
+  url     = {https://github.com/ChenAo-Phys/quantax},
+  version = {0.3.0}
+}
+```

{quantax-0.2.1 → quantax-0.3.0}/pyproject.toml

@@ -4,17 +4,17 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "quantax"
-version = "0.2.1"
+version = "0.3.0"
 description = "Flexible neural quantum states based on JAX"
 authors = [
     { name = "Ao Chen", email = "chenao.phys@gmail.com" },
     { name = "Christopher Roth" }
 ]
-requires-python = ">=3.10,<=3.13"
+requires-python = ">=3.11,<3.14"
 dependencies = [
-    "numpy>=1.26.0",
-    "scipy>=1.11.1",
-    "jax>=0.6.1",
+    "numpy>=2.0.0",
+    "scipy>=1.14.0",
+    "jax>=0.8.0",
     "equinox>=0.11.4",
     "lrux>=0.1.0",
 ]
@@ -22,14 +22,30 @@ dependencies = [
 [project.optional-dependencies]
 full = [
     "numpy>=2.0.0",
-    "scipy>=1.13.1",
+    "scipy>=1.14.0",
     "matplotlib>=3.8.0",
     "quspin>=1.0.0",
-    "jax>=0.6.1",
+    "jax>=0.8.0",
     "equinox>=0.11.4",
+    "lineax>=0.1.0",
     "lrux>=0.1.0",
+    "jaxmg>=0.0.6",
+]
+docs = [
+    "sphinx",
+    "myst-nb",
+    "pydata-sphinx-theme",
 ]
 
 [tool.setuptools.packages.find]
 where = ["."]
 include = ["quantax*"]
+
+[tool.black]
+target-version = ["py311", "py312", "py313"]
+
+[tool.pyright]
+include = ["quantax"]
+pythonVersion = "3.12"
+typeCheckingMode = "standard"
+

{quantax-0.2.1 → quantax-0.3.0}/quantax/__init__.py

@@ -1,3 +1,10 @@
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("quantax")
+except PackageNotFoundError:
+    __version__ = "unknown"
+
 from .global_defs import (
     set_random_seed,
     set_default_dtype,

{quantax-0.2.1 → quantax-0.3.0}/quantax/global_defs.py

@@ -1,31 +1,32 @@
-from typing import Optional, Tuple
 from functools import partial
 from enum import Enum
 import jax
 import jax.numpy as jnp
 import jax.random as jr
+from jax.typing import DTypeLike
 
+DTYPE = jnp.dtype(jnp.float32)
 
-jax.config.update("jax_enable_x64", True)
-jax.config.update("jax_default_matmul_precision", "float32")
 
-
-DTYPE = jnp.float64
-
-
-def set_default_dtype(dtype: jnp.dtype) -> None:
+def set_default_dtype(dtype: DTypeLike) -> None:
     """
     Set the default data type in Quantax.
-    Recommended to be ``jnp.float64`` or ``jnp.complex128``. Default to ``jnp.float64``.
+    Recommended to be ``jnp.float32`` or ``jnp.complex64``. Default to ``jnp.float32``.
 
     .. note::
         This doesn't alter the computation inside ``quantax.model``.
     """
+    dtype = jnp.dtype(dtype)
+
     if not (
         jnp.issubdtype(dtype, jnp.floating)
         or jnp.issubdtype(dtype, jnp.complexfloating)
     ):
         raise ValueError("'dtype' should be float or complex types")
+
+    if dtype == jnp.float64 or dtype == jnp.complex128:
+        jax.config.update("jax_enable_x64", True)
+
     global DTYPE
     DTYPE = dtype
 
@@ -58,13 +59,13 @@ def set_random_seed(seed: int) -> None:
     global KEY
     KEY = jr.key(seed)
 
-    from .utils import to_replicate_array
+    from .utils import local_to_replicated
 
-    KEY = to_replicate_array(KEY)
+    KEY = local_to_replicated(KEY)
 
 
 @partial(jax.jit, static_argnums=1)
-def _gen_keys(key, num: Optional[int] = None) -> Tuple[jax.Array, jax.Array]:
+def _gen_keys(key, num: int | None = None) -> tuple[jax.Array, jax.Array]:
     nkeys = 2 if num is None else num + 1
     new_keys = jr.split(key, nkeys)
     key = new_keys[0]
@@ -72,7 +73,7 @@ def _gen_keys(key, num: Optional[int] = None) -> Tuple[jax.Array, jax.Array]:
     return key, new_keys
 
 
-def get_subkeys(num: Optional[int] = None) -> jax.Array:
+def get_subkeys(num: int | None = None) -> jax.Array:
     """
     Get jax keys stored in Quantax. The keys are replicated across all devices.
 

{quantax-0.2.1 → quantax-0.3.0}/quantax/model/__init__.py

@@ -11,6 +11,5 @@ from .fermion_mf import (
     MultiPf,
     PartialPair,
 )
-from .transformer import Transformer, Embedding, MHSA, FFN
-from .jastrow import NeuralJastrow
+from .jastrow import Jastrow, GeneralJastrow
 from .backflow import DetBackflow, PfBackflow

{quantax-0.2.1 → quantax-0.3.0}/quantax/model/backflow.py

@@ -1,8 +1,8 @@
-from typing import Optional, Tuple, Union
+from typing import Any, Callable, overload, Literal
 import jax
 import jax.numpy as jnp
 import jax.random as jr
-import equinox as eqx
+from jax.typing import DTypeLike
 import lrux
 from .fermion_mf import GeneralDet, MF_Internal, _init_spinless_orbs
 from ..global_defs import get_sites, get_subkeys
@@ -18,23 +18,23 @@ from ..utils import LogArray
 
 
 class DetBackflow(RefModel):
-    net: eqx.Module
-    U0: jax.Array
-    W: jax.Array
-    dtype: jnp.dtype
-
     r"""
     Determinant backflow model.
     :math:`\psi(n) = \mathrm{det}(n \star (U_0 + U_1(n)))`,
     where :math:`\star` denotes the operation slicing the rows of the matrix.
     """
 
+    net: Callable[[jax.Array], jax.Array]
+    U0: jax.Array
+    W: jax.Array
+    dtype: DTypeLike
+
     def __init__(
         self,
-        net: eqx.Module,
+        net: Callable[[jax.Array], jax.Array],
         d: int,
-        U0: Optional[jax.Array] = None,
-        dtype: jnp.dtype = jnp.float64,
+        U0: jax.Array | None = None,
+        dtype: DTypeLike = jnp.float64,
     ):
         r"""
         Initialize the determinant backflow model.
@@ -71,8 +71,8 @@ class DetBackflow(RefModel):
             d //= 2
         self.W = lecun_normal(get_subkeys(), (sites.Ntotal, d), dtype=dtype) / 10
 
-    def __call__(self, s: jax.Array) -> jax.Array:
-        x = self.net(s)
+    def __call__(self, s: jax.Array) -> LogArray:
+        x = jnp.asarray(self.net(s))
 
         idx = fermion_idx(s)
         x = x.reshape(-1, get_sites().Nfmodes).astype(self.dtype)
@@ -89,43 +89,78 @@ class DetBackflow(RefModel):
         backflow correction is larger than the total number of particles.
         """
         Ntotal = get_sites().Ntotal
+        if Ntotal is None:
+            raise ValueError
         rank = self.W.shape[1]
         return rank < Ntotal
 
-    def init_internal(self, s) -> Optional[MF_Internal]:
+    def init_internal(self, s: jax.Array) -> tuple[LogArray, MF_Internal | None]:
         """
-        Initialize internal values for given input configurations.
+        Return wavefunction and internal values for given input configurations.
         See `~quantax.nn.RefModel` for details.
         """
         if not self.use_ref:
-            return None
+            return self(s), None
 
         idx = fermion_idx(s)
         orbs = self.U0[idx, :]
         inv = jnp.linalg.inv(orbs)
         sign, logabs = jnp.linalg.slogdet(orbs)
         psi = LogArray(sign, logabs) * fermion_inverse_sign(s)
-        return MF_Internal(idx, inv, psi)
+        return psi, MF_Internal(idx, inv, psi)
+
+    @property
+    def required_update_modes(self) -> tuple[str, ...]:
+        """
+        The required update modes for accelerated ref_forward pass.
+        """
+        return ("nflips",)
+
+    @overload
+    def ref_forward(
+        self,
+        s: jax.Array,
+        s_old: jax.Array,
+        update_mode: dict[str, Any],
+        internal: MF_Internal | None,
+        return_update: Literal[False] = False,
+    ) -> LogArray: ...
+
+    @overload
+    def ref_forward(
+        self,
+        s: jax.Array,
+        s_old: jax.Array,
+        update_mode: dict[str, Any],
+        internal: MF_Internal | None,
+        return_update: Literal[True],
+    ) -> tuple[LogArray, MF_Internal | None]: ...
 
     def ref_forward(
         self,
         s: jax.Array,
         s_old: jax.Array,
-        nflips: int,
-        internal: Optional[MF_Internal],
+        update_mode: dict[str, Any],
+        internal: MF_Internal | None,
         return_update: bool = False,
-    ) -> Union[LogArray, Tuple[LogArray, MF_Internal]]:
+    ) -> LogArray | tuple[LogArray, MF_Internal | None]:
         """
         Accelerated forward pass through local updates and internal quantities.
         See `~quantax.nn.RefModel` for details.
         """
-        if not self.use_ref:
+        if (not self.use_ref) or (internal is None):
             psi = self(s)
             if return_update:
                 return psi, internal
             else:
                 return psi
 
+        if not isinstance(internal.idx, jax.Array):
+            raise ValueError
+        if not isinstance(internal.inv, jax.Array):
+            raise ValueError
+
+        nflips = update_mode["nflips"]
         nhops = nflips // 2
         idx_annihilate, idx_create = changed_inds(s, s_old, nhops)
         idx = internal.idx
@@ -135,7 +170,7 @@ class DetBackflow(RefModel):
         new_idx = idx.at[row_update_idx].set(idx_create)
         row_update = self.U0[idx_create] - self.U0[idx_annihilate]
 
-        x = self.net(s)
+        x = jnp.asarray(self.net(s))
         x = x.reshape(-1, get_sites().Nfmodes).astype(self.dtype)
         x = x.T[new_idx]
 
@@ -156,25 +191,25 @@ class DetBackflow(RefModel):
 
 
 class PfBackflow(RefModel):
-    net: eqx.Module
-    U0: jax.Array
-    J0: jax.Array
-    W: jax.Array
-    dtype: jnp.dtype
-
     r"""
     Pfaffian backflow model.
     :math:`\psi(n) = \mathrm{pf}(n \star (U_0 + U_1(n)) J_0 (U_0 + U_1(n))^T)`,
     where :math:`\star` denotes the operation slicing the rows and columns of the matrix.
     """
 
+    net: Callable[[jax.Array], jax.Array]
+    U0: jax.Array
+    J0: jax.Array
+    W: jax.Array
+    dtype: DTypeLike
+
     def __init__(
         self,
-        net: eqx.Module,
+        net: Callable[[jax.Array], jax.Array],
         d: int,
-        U0: Optional[jax.Array] = None,
-        J0: Optional[jax.Array] = None,
-        dtype: jnp.dtype = jnp.float64,
+        U0: jax.Array | None = None,
+        J0: jax.Array | None = None,
+        dtype: DTypeLike = jnp.float64,
     ):
         r"""
         Initialize the Pfaffian backflow model.
@@ -206,6 +241,9 @@ class PfBackflow(RefModel):
                 U1 = U0.conj() if jnp.issubdtype(dtype, jnp.complexfloating) else U0
                 O = jnp.zeros_like(U0)
                 U0 = jnp.block([[U0, O], [O, U1]])
+            # Break the degeneracy of the clean Fermi sea, otherwise the spinful
+            # block structure makes the mean-field Pfaffian U0 J0 U0^T singular.
+            U0 += jr.normal(get_subkeys(), U0.shape, U0.dtype) * jnp.std(U0) * 0.1
         elif U0.shape != (M, M):
             raise ValueError(f"U0 must have shape {(M, M)}, got {U0.shape}")
         U0 /= jnp.std(U0)
@@ -213,7 +251,7 @@ class PfBackflow(RefModel):
 
         if J0 is None:
             if sites.is_spinful:
-                J0 = lrux.skew_eye(M // 2, dtype)
+                J0 = lrux.skew_eye(M // 2, dtype)  # type: ignore
             else:
                 J0 = jr.normal(get_subkeys(), (M, M), dtype=dtype)
                 J0 = (J0 - J0.T) / 2
@@ -227,6 +265,10 @@ class PfBackflow(RefModel):
 
     @property
     def J0_full(self) -> jax.Array:
+        """
+        Returns the full antisymmetric mean-field pairing matrix J0, reconstructed
+        from its stored upper-triangular entries.
+        """
         M = self.U0.shape[0]
         J_full = jnp.zeros((M, M), dtype=self.dtype)
         triu_indices = jnp.triu_indices(M, k=1)
@@ -234,8 +276,8 @@ class PfBackflow(RefModel):
         J_full = J_full - J_full.T
         return J_full
 
-    def __call__(self, s: jax.Array) -> jax.Array:
-        x = self.net(s)
+    def __call__(self, s: jax.Array) -> LogArray:
+        x = jnp.asarray(self.net(s))
 
         idx = fermion_idx(s)
         x = x.reshape(-1, get_sites().Nfmodes).astype(self.dtype)
@@ -253,16 +295,18 @@ class PfBackflow(RefModel):
         backflow correction is larger than the total number of particles.
         """
         Ntotal = get_sites().Ntotal
+        if Ntotal is None:
+            raise ValueError
         rank = self.W.shape[1] * 2
         return rank < Ntotal
 
-    def init_internal(self, s) -> Optional[MF_Internal]:
+    def init_internal(self, s: jax.Array) -> tuple[LogArray, MF_Internal | None]:
         """
-        Initialize internal values for given input configurations.
+        Return wavefunction and internal values for given input configurations.
         See `~quantax.nn.RefModel` for details.
         """
         if not self.use_ref:
-            return None
+            return self(s), None
 
         idx = fermion_idx(s)
         U = self.U0[idx, :]
@@ -270,27 +314,60 @@ class PfBackflow(RefModel):
         inv = jnp.linalg.inv(F)
         sign, logabs = lrux.slogpf(F)
         psi = LogArray(sign, logabs) * fermion_inverse_sign(s)
-        return MF_Internal(idx, inv, psi)
+        return psi, MF_Internal(idx, inv, psi)
+
+    @property
+    def required_update_modes(self) -> tuple[str, ...]:
+        """
+        The required update modes for accelerated ref_forward pass.
+        """
+        return ("nflips",)
 
+    @overload
     def ref_forward(
         self,
         s: jax.Array,
         s_old: jax.Array,
-        nflips: int,
-        internal: Optional[MF_Internal],
+        update_mode: dict[str, Any],
+        internal: MF_Internal | None,
+        return_update: Literal[False] = False,
+    ) -> LogArray: ...
+
+    @overload
+    def ref_forward(
+        self,
+        s: jax.Array,
+        s_old: jax.Array,
+        update_mode: dict[str, Any],
+        internal: MF_Internal | None,
+        return_update: Literal[True],
+    ) -> tuple[LogArray, MF_Internal | None]: ...
+
+    def ref_forward(
+        self,
+        s: jax.Array,
+        s_old: jax.Array,
+        update_mode: dict[str, Any],
+        internal: MF_Internal | None,
         return_update: bool = False,
-    ) -> Union[LogArray, Tuple[LogArray, MF_Internal]]:
+    ) -> LogArray | tuple[LogArray, MF_Internal | None]:
         """
         Accelerated forward pass through local updates and internal quantities.
         See `~quantax.nn.RefModel` for details.
         """
-        if not self.use_ref:
+        if (not self.use_ref) or (internal is None):
             psi = self(s)
             if return_update:
                 return psi, internal
             else:
                 return psi
 
+        if not isinstance(internal.idx, jax.Array):
+            raise ValueError
+        if not isinstance(internal.inv, jax.Array):
+            raise ValueError
+
+        nflips = update_mode["nflips"]
         nhops = nflips // 2
         idx_annihilate, idx_create = changed_inds(s, s_old, nhops)
         idx = internal.idx
@@ -305,7 +382,7 @@ class PfBackflow(RefModel):
         U_mean = (U0[new_idx, :] + U0[idx, :]) / 2
         x = jnp.einsum("im,mn,jn->ij", U_diff, J0, U_mean).T
 
-        U1 = self.net(s)
+        U1 = jnp.asarray(self.net(s))
         U1 = U1.reshape(-1, get_sites().Nfmodes).astype(self.dtype)
         U1 = U1.T[new_idx, :]
         W = self.W