quantax 0.1.1__py3-none-any.whl

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/__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/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/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