quantax 0.1.1__tar.gz

quantax-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Chen Ao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

quantax-0.1.1/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: quantax
+Version: 0.1.1
+Summary: Flexible neural quantum states based on QuSpin, JAX, and Equinox
+Author: Christopher Roth
+Author-email: Ao Chen <chenao.phys@gmail.com>
+Requires-Python: <=3.13,>=3.10
+License-File: LICENSE
+Requires-Dist: numpy>=2.0.0
+Requires-Dist: quspin>=1.0.0
+Requires-Dist: matplotlib>=3.8.0
+Requires-Dist: ml_dtypes>=0.4.0
+Requires-Dist: jax>=0.4.34
+Requires-Dist: equinox>=0.11.4
+Dynamic: license-file

quantax-0.1.1/README.md

@@ -0,0 +1,14 @@
+<h1 align='center'>quantax</h1>
+
+Flexible neural quantum states based on [QuSpin](https://github.com/QuSpin/QuSpin/tree/dev_0.3.8), [JAX](https://github.com/google/jax), and [Equinox](https://github.com/patrick-kidger/equinox).
+
+[Documentation](https://chenao-phys.github.io/quantax)
+
+## Installation
+
+Requires Python 3.10+, JAX 0.4.34+
+
+```
+git clone https://github.com/ChenAo-Phys/quantax.git
+pip install ./quantax
+```

quantax-0.1.1/docs/source/conf.py

@@ -0,0 +1,39 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# For the full list of built-in configuration values, see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+import pathlib
+import sys
+
+sys.path.insert(0, pathlib.Path(__file__).parents[2].resolve().as_posix())
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = "Quantax"
+copyright = "2024, Ao Chen"
+author = "Ao Chen, Christopher Roth"
+release = "0.1.1"
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = ["sphinx.ext.autodoc", "sphinx.ext.autosummary"]
+
+templates_path = ["_templates"]
+exclude_patterns = []
+
+autodoc_member_order = "bysource"
+autosummary_generate = True
+autosummary_generate_overwrite = False
+default_role = "py:obj"
+
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = "pydata_sphinx_theme"
+#html_static_path = ["_static"]
+html_use_relative_urls = True
+html_baseurl = 'https://chenao-phys.github.io/quantax/'

quantax-0.1.1/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "quantax"
+version = "0.1.1"
+description = "Flexible neural quantum states based on QuSpin, JAX, and Equinox"
+authors = [
+    { name = "Ao Chen", email = "chenao.phys@gmail.com" },
+    { name = "Christopher Roth" }
+]
+requires-python = ">=3.10,<=3.13"
+dependencies = [
+    "numpy>=2.0.0",
+    "quspin>=1.0.0",
+    "matplotlib>=3.8.0",
+    "ml_dtypes>=0.4.0",
+    "jax>=0.4.34",
+    "equinox>=0.11.4"
+]
+
+[tool.setuptools.packages.find]
+where = ["."]

quantax-0.1.1/quantax/__init__.py

@@ -0,0 +1,22 @@
+from .global_defs import (
+    set_random_seed,
+    set_default_dtype,
+    get_default_dtype,
+    get_real_dtype,
+    get_subkeys,
+    PARTICLE_TYPE,
+    get_sites,
+    get_lattice,
+)
+
+from . import (
+    sites,
+    symmetry,
+    operator,
+    nn,
+    model,
+    state,
+    sampler,
+    optimizer,
+    utils,
+)

quantax-0.1.1/quantax/global_defs.py

@@ -0,0 +1,131 @@
+from typing import Optional, Tuple
+from functools import partial
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+import equinox as eqx
+
+
+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:
+    """
+    Set the default data type for the computation in quantax.
+    Recommended to be ``jnp.float64`` or ``jnp.complex128``. Default to ``jnp.float64``.
+
+    .. note::
+        This doesn't alter the computation inside `quantax.model`.
+    """
+    if not (
+        jnp.issubdtype(dtype, jnp.floating)
+        or jnp.issubdtype(dtype, jnp.complexfloating)
+    ):
+        raise ValueError("'dtype' should be float or complex types")
+    global DTYPE
+    DTYPE = dtype
+
+
+def get_default_dtype() -> jnp.dtype:
+    """Return the default data type for the computation in quantax."""
+    return DTYPE
+
+
+def get_real_dtype() -> jnp.dtype:
+    """Return the default real data type for the computation in quantax."""
+    return jnp.finfo(DTYPE).dtype
+
+
+def is_default_cpl() -> bool:
+    """Return whether the default data type is complex."""
+    return jnp.issubdtype(DTYPE, jnp.complexfloating)
+
+
+KEY = None
+
+
+def set_random_seed(seed: int) -> None:
+    """
+    Set the initial random seed for the computation in jax.
+    Default to be a random number generated by numpy.
+    """
+    global KEY
+    KEY = jr.key(seed)
+
+
+@partial(jax.jit, static_argnums=1)
+def _gen_keys(key, num: Optional[int] = 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]
+    new_keys = new_keys[1] if num is None else new_keys[1:]
+    return key, new_keys
+
+
+def get_subkeys(num: Optional[int] = None) -> jax.Array:
+    """
+    Get a certain number of jax PRNG keys as an array.
+
+    :param num:
+        The number of returned keys.
+        If ``num is None`` (default), then return only 1 key instead of an array of keys.
+
+    .. warning::
+        This function is not jittable, because it reads and writes the global key
+        stored in quantax.
+    """
+    global KEY
+    if KEY is None:
+        set_random_seed(0)
+    KEY, new_keys = _gen_keys(KEY, num)
+    return new_keys
+
+
+class PARTICLE_TYPE(eqx.Enumeration):
+    r"""
+    The enums to distinguish different particle types.
+
+    0: spin
+
+    1: spinful_fermion
+
+    2: spinless_fermion
+
+    (Not implemented) 3: boson
+    """
+
+    spin = 0
+    spinful_fermion = 1
+    spinless_fermion = 2
+
+
+from .sites import Sites, Lattice
+
+
+def get_sites() -> Sites:
+    """
+    Get the `quantax.sites.Sites` used in the current program.
+
+    .. warning::
+        Unlike other NQS packages, in quantax the geometry graph and the hilbert space
+        is defined as a global constant which shouldn't be changed within a single program.
+    """
+    sites = Sites._SITES
+    if sites is None:
+        raise RuntimeError("The `Sites` hasn't been defined.")
+    return sites
+
+
+def get_lattice() -> Lattice:
+    """
+    Get the `quantax.sites.Lattice` used in the current program. This is similar to `get_sites`,
+    but will raise an error if the defined `Sites` is not a `Lattice`.
+    """
+    sites = get_sites()
+    if not isinstance(sites, Lattice):
+        raise RuntimeError("Require a `Lattice`, but got a general `Sites`")
+    return sites

quantax-0.1.1/quantax/model/__init__.py

@@ -0,0 +1,11 @@
+from .prod_nets import SingleDense, RBM_Dense, SingleConv, RBM_Conv, ResProd, SchmittNet
+from .sum_nets import ResSum, ResSumGconv
+from .transformer import ConvTransformer
+from .sign_nets import SgnNet, MarshallSign, StripeSign, Neel120
+from .triangular_nets import Triangular_Neighbor_Conv, Triangular_ResSum
+from .fermion_mf import (
+    Determinant,
+    Pfaffian,
+    PairProduct,
+)
+from .neural_fermion import NeuralJastrow, HiddenPfaffian