Metadata-Version: 2.4
Name: quantumsolax
Version: 0.1.0
Summary: A Python solver for fermionic quantum systems with neural network support
Author-email: Pavlo Bilous <pavlo.bilous@proton.me>
License: Creative Commons Legal Code
        
        CC0 1.0 Universal
        
            CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE
            LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN
            ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS
            INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES
            REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS
            PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM
            THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED
            HEREUNDER.
        
        Statement of Purpose
        
        The laws of most jurisdictions throughout the world automatically confer
        exclusive Copyright and Related Rights (defined below) upon the creator
        and subsequent owner(s) (each and all, an "owner") of an original work of
        authorship and/or a database (each, a "Work").
        
        Certain owners wish to permanently relinquish those rights to a Work for
        the purpose of contributing to a commons of creative, cultural and
        scientific works ("Commons") that the public can reliably and without fear
        of later claims of infringement build upon, modify, incorporate in other
        works, reuse and redistribute as freely as possible in any form whatsoever
        and for any purposes, including without limitation commercial purposes.
        These owners may contribute to the Commons to promote the ideal of a free
        culture and the further production of creative, cultural and scientific
        works, or to gain reputation or greater distribution for their Work in
        part through the use and efforts of others.
        
        For these and/or other purposes and motivations, and without any
        expectation of additional consideration or compensation, the person
        associating CC0 with a Work (the "Affirmer"), to the extent that he or she
        is an owner of Copyright and Related Rights in the Work, voluntarily
        elects to apply CC0 to the Work and publicly distribute the Work under its
        terms, with knowledge of his or her Copyright and Related Rights in the
        Work and the meaning and intended legal effect of CC0 on those rights.
        
        1. Copyright and Related Rights. A Work made available under CC0 may be
        protected by copyright and related or neighboring rights ("Copyright and
        Related Rights"). Copyright and Related Rights include, but are not
        limited to, the following:
        
          i. the right to reproduce, adapt, distribute, perform, display,
             communicate, and translate a Work;
         ii. moral rights retained by the original author(s) and/or performer(s);
        iii. publicity and privacy rights pertaining to a person's image or
             likeness depicted in a Work;
         iv. rights protecting against unfair competition in regards to a Work,
             subject to the limitations in paragraph 4(a), below;
          v. rights protecting the extraction, dissemination, use and reuse of data
             in a Work;
         vi. database rights (such as those arising under Directive 96/9/EC of the
             European Parliament and of the Council of 11 March 1996 on the legal
             protection of databases, and under any national implementation
             thereof, including any amended or successor version of such
             directive); and
        vii. other similar, equivalent or corresponding rights throughout the
             world based on applicable law or treaty, and any national
             implementations thereof.
        
        2. Waiver. To the greatest extent permitted by, but not in contravention
        of, applicable law, Affirmer hereby overtly, fully, permanently,
        irrevocably and unconditionally waives, abandons, and surrenders all of
        Affirmer's Copyright and Related Rights and associated claims and causes
        of action, whether now known or unknown (including existing as well as
        future claims and causes of action), in the Work (i) in all territories
        worldwide, (ii) for the maximum duration provided by applicable law or
        treaty (including future time extensions), (iii) in any current or future
        medium and for any number of copies, and (iv) for any purpose whatsoever,
        including without limitation commercial, advertising or promotional
        purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each
        member of the public at large and to the detriment of Affirmer's heirs and
        successors, fully intending that such Waiver shall not be subject to
        revocation, rescission, cancellation, termination, or any other legal or
        equitable action to disrupt the quiet enjoyment of the Work by the public
        as contemplated by Affirmer's express Statement of Purpose.
        
        3. Public License Fallback. Should any part of the Waiver for any reason
        be judged legally invalid or ineffective under applicable law, then the
        Waiver shall be preserved to the maximum extent permitted taking into
        account Affirmer's express Statement of Purpose. In addition, to the
        extent the Waiver is so judged Affirmer hereby grants to each affected
        person a royalty-free, non transferable, non sublicensable, non exclusive,
        irrevocable and unconditional license to exercise Affirmer's Copyright and
        Related Rights in the Work (i) in all territories worldwide, (ii) for the
        maximum duration provided by applicable law or treaty (including future
        time extensions), (iii) in any current or future medium and for any number
        of copies, and (iv) for any purpose whatsoever, including without
        limitation commercial, advertising or promotional purposes (the
        "License"). The License shall be deemed effective as of the date CC0 was
        applied by Affirmer to the Work. Should any part of the License for any
        reason be judged legally invalid or ineffective under applicable law, such
        partial invalidity or ineffectiveness shall not invalidate the remainder
        of the License, and in such case Affirmer hereby affirms that he or she
        will not (i) exercise any of his or her remaining Copyright and Related
        Rights in the Work or (ii) assert any associated claims and causes of
        action with respect to the Work, in either case contrary to Affirmer's
        express Statement of Purpose.
        
        4. Limitations and Disclaimers.
        
         a. No trademark or patent rights held by Affirmer are waived, abandoned,
            surrendered, licensed or otherwise affected by this document.
         b. Affirmer offers the Work as-is and makes no representations or
            warranties of any kind concerning the Work, express, implied,
            statutory or otherwise, including without limitation warranties of
            title, merchantability, fitness for a particular purpose, non
            infringement, or the absence of latent or other defects, accuracy, or
            the present or absence of errors, whether or not discoverable, all to
            the greatest extent permissible under applicable law.
         c. Affirmer disclaims responsibility for clearing rights of other persons
            that may apply to the Work or any use thereof, including without
            limitation any person's Copyright and Related Rights in the Work.
            Further, Affirmer disclaims responsibility for obtaining any necessary
            consents, permissions or other rights required for any use of the
            Work.
         d. Affirmer understands and acknowledges that Creative Commons is not a
            party to this document and has no duty or obligation with respect to
            this CC0 or use of the Work.
        
Project-URL: Homepage, https://github.com/pavlo-bilous/quantumsolax
Project-URL: Repository, https://github.com/pavlo-bilous/quantumsolax
Project-URL: Documentation, https://quantumsolax.readthedocs.io/
Project-URL: Paper, https://www.scipost.org/SciPostPhysCodeb.51
Keywords: quantum,fermions,second quantization,jax,neural networks
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: CC0 1.0 Universal (CC0 1.0) Public Domain Dedication
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: scipy
Requires-Dist: pandas
Requires-Dist: jax
Requires-Dist: flax
Requires-Dist: optax
Requires-Dist: orbax-checkpoint
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: pytest-cov; extra == "test"
Provides-Extra: lint
Requires-Dist: ruff; extra == "lint"
Provides-Extra: docs
Requires-Dist: sphinx; extra == "docs"
Requires-Dist: furo; extra == "docs"
Dynamic: license-file

# quantumsolax

[![tests](https://github.com/pavlo-bilous/quantumsolax/actions/workflows/tests.yml/badge.svg)](https://github.com/pavlo-bilous/quantumsolax/actions/workflows/tests.yml)
[![docs](https://readthedocs.org/projects/quantumsolax/badge/?version=latest)](https://quantumsolax.readthedocs.io/en/latest/?badge=latest)

A [JAX](https://github.com/jax-ml/jax)-based Python library for solving fermionic quantum many-body systems with neural network support. The framework allows to efficiently encode and manipulate bases of Slater determinants, quantum states and operators within the second quantization formalism. Operators can be converted to matrices on a given basis for subsequent diagonalization. In case the basis is too large to treat directly, neural-network-assisted support can be leveraged to select the most important Slater determinants. See the paper for the full design and physics background: [SciPost Phys. Codebases 51](https://www.scipost.org/SciPostPhysCodeb.51). Full documentation is at [quantumsolax.readthedocs.io](https://quantumsolax.readthedocs.io/).

## Installation

```bash
pip install quantumsolax
```

This installs a CPU-only JAX version. For GPU acceleration, additionally install the CUDA build of JAX matching your system:

```bash
pip install -U "jax[cuda12]"
```

Requires Python >=3.10.

For development (running the test suite, linting):

```bash
git clone https://github.com/pavlo-bilous/quantumsolax.git
cd quantumsolax
pip install -e ".[test,lint]"
```

## Quick start

From [`JupyterNotebooks/main_solax_classes.ipynb`](JupyterNotebooks/main_solax_classes.ipynb) (a hopping term plus a density-density interaction, combined into a Hamiltonian and represented as a matrix in a 4-determinant basis):

```python
import numpy as np
import solax as sx

# One-directional hopping term: a_0^dagger a_2 + a_1^dagger a_3
V0 = sx.OperatorTerm((1, 0), np.array([[0, 2], [1, 3]]), np.array([1.0, 1.0]))

# Density-density interaction term
U = sx.Operator((1, 0, 1, 0), np.array([[0, 0, 1, 1], [2, 2, 3, 3]]), np.array([0.25, 0.75]))

# Adding a bare number introduces a constant/identity term
H = 1 + V0 + V0.hconj + U

basis = sx.Basis(["1001", "1100", "0110", "0011"])
matrix = H.build_matrix(basis)
print(matrix.to_scipy().todense())
# [[ 1.    1.    0.    1.  ]
#  [ 1.    1.25 -1.    0.  ]
#  [ 0.   -1.    1.   -1.  ]
#  [ 1.    0.   -1.    1.75]]
```

`Basis`/`State`/`Operator`/`OperatorTerm`/`OperatorMatrix` are the core building blocks; see the docstrings (`help(sx.Basis)`, etc., or the generated API docs) and the rest of the notebooks in [`JupyterNotebooks/`](JupyterNotebooks/) for a full walkthrough, including a Single Impurity Anderson Model example and the neural-network-assisted basis-selection workflow (`sx.BasisClassifier`, `sx.BigBasisManager`).

## Features

- **`quantum_core`** — `sx.Basis`, `sx.State`, `sx.Operator`, `sx.OperatorTerm`, `sx.OperatorMatrix`: build second-quantized operators, apply them to a basis/state, and construct sparse Hamiltonian matrices, batched and JAX-accelerated (including optional multi-GPU support).
- **`big_basis_management`** — `sx.BasisClassifier`/`sx.BigBasisManager`: train a small neural-network classifier to predict which determinants in an intractably large basis are likely important, so you can restrict further work to a tractable subset.
- **`save_load`** — `sx.save`/`sx.load`: persist and restore solax objects (or nested dicts mixing them with NumPy arrays and plain Python values) to/from disk, without pickle.

## Advanced use

The article [SciPost Phys. Codebases 51](https://www.scipost.org/SciPostPhysCodeb.51) describes the quantumsolax functionality usually necessary for fermionic many-body computations. This section documents additional functionality not covered there.

### Squeezing (deduplication) control

`Basis` and `State` normally deduplicate their determinants automatically — on construction, and after operations (`+`, applying an `OperatorTerm`) that might introduce repeats. `is_squeezed` reports whether an object currently holds no duplicates, and `squeeze()` returns a deduplicated copy (merging coefficients for `State`, keeping the first occurrence for `Basis`). To build up an intermediate result across several steps without paying for deduplication after each one, wrap the steps in `manual_squeezing()` (`solax.quantum_core.mode_ctrl`), which suspends all automatic squeezing for its duration — then call `squeeze()` explicitly once at the end.

### Tracking determinants through an operator application

`Operator`/`OperatorTerm.__call__` accept a `det_tracking=True` keyword: alongside the usual result, it returns a 1D integer array mapping each determinant in the *output* back to the index of the determinant in the *input* it came from — useful when you need to know, not just compute, which input determinant produced which output.

### Building a custom NN-assisted tool with `neural_framework`

`BasisClassifier`/`BigBasisManager` are one particular application (basis-importance classification) of a smaller, generic Flax/JAX training layer, `solax.neural_framework`, importable and reusable directly for other "features → labels" tasks:

- `NeuralModel(call_on_entry, loss_fn, post_transform=identity)` wraps a per-entry architecture function, loss, and optional output transform into a trainable model; call `.initialize(key, dummy_features, optimizer)` once before use.
- `train_on_data(key, model, train_data, *, val_data=None, batch_size=None, epochs=1, train_metrics=None, val_metrics=None, ...)` runs the batched training loop, with optional validation and early stopping via a `MetricsMonitor`/`Guard`.
- `predict_on_data(model, features, *, batch_size=None)` runs batched inference; the model is also directly callable on a batch of features.

`LeastSqRegressor`/`SoftmaxClassifier` (`neural_framework.ready_classes`) are worked examples of subclassing `NeuralModel` for a specific loss/output configuration, paired with `LossMonitor`/`AccuracyMonitor` for tracking.

## For developers

This section contains remarks for future developers of code based on quantumsolax.

### Making a custom class savable/loadable

Users can implement their own classes from scratch, or by inheriting from solax's existing ones. Objects of such a custom class can be made compatible with `sx.save`/`sx.load` by registering their class with `save_load_registry` (`solax.save_load.registration`):

```python
from solax.save_load.registration import save_load_registry

save_load_registry.register(label, cls, init_from_attr)
```

- `label`: a string identifying the class in the saved data (by convention, the class's own name).
- `cls`: the class itself.
- `init_from_attr`: a callable that reconstructs an instance from the class's dictified attributes, passed as keyword arguments. If `__init__` already accepts its attributes by name, `cls` itself works here — this is the common case.

By convention, a class registers itself immediately after its own definition, in the same module — solax's own classes follow this pattern: `Basis`, `State`, `Operator`, `OperatorTerm`, `OperatorMatrix`, and `RandomKeys` each call `register(...)` right after their class body. Follow the same convention for a new class.

If an attribute isn't itself a solax object, a NumPy array, or a Python primitive (e.g. it holds a JAX array, or a dict keyed by something other than strings), define `__pre_dictify__`/`__post_undictify__` on the class: `__pre_dictify__()` returns a surrogate instance of the same class with only dictifiable attributes (this is what actually gets saved), and `__post_undictify__()` converts that surrogate back after loading. `solax/random_keys.py`'s `RandomKeys` is a worked example — it holds a JAX key, converted to/from plain NumPy for saving.

A class can instead fully implement its own saving/loading mechanism by defining `__save__`/`__load__` on it: `sx.save`/`sx.load` then step aside for such a class entirely, recording only which class it was rather than calling these methods themselves — actually invoking `__save__`/`__load__` and doing the work is left to the class's own code. Not used by any class in the current quantumsolax version.

### Running tests and linting

solax uses `pytest` for its test suite and `ruff` for linting. Most of the suite runs on a single device; a few tests specifically exercise JAX's multi-device (`pmap`) batching path and are skipped unless multiple devices are actually available, and the neural-network training tests are marked `slow` since they actually train small models.

```bash
pytest                                     # full suite (multi_device auto-skips below 2 devices)
pytest -m "not slow"                       # skip the NN-training tests, for a quick check
XLA_FLAGS=--xla_force_host_platform_device_count=2 pytest -m multi_device  # fake 2 CPU devices to run the pmap tests
ruff check solax/ tests/                   # lint
```

## Authorship

The code in this repository was written by Pavlo Bilous. At the packaging stage, assistance of Claude Code was used.

## License

[CC0 1.0 Universal](LICENSE) (public domain dedication).
