qmb - npm Package Compare versions

Comparing version 0.0.9 to 0.0.10

.github/workflows/pre-commit.yml

+name: Pre-commit Hooks
+
+on: [push, pull_request]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - uses: pre-commit/action@v3.0.1

.github/workflows/wheels.yml

+name: Build
+
+on: [push, pull_request]
+
+jobs:
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up QEMU
+        if: runner.os == 'Linux'
+        uses: docker/setup-qemu-action@v3
+        with:
+          platforms: all
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install cibuildwheel
+        run: python -m pip install cibuildwheel==2.21.1
+
+      - name: Build wheels
+        run: python -m cibuildwheel --output-dir wheelhouse
+        env:
+          CIBW_TEST_SKIP: "*-win_arm64"
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
+          path: ./wheelhouse/*.whl
+
+  build_sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build sdist
+        run: pipx run build --sdist
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-sdist
+          path: dist/*.tar.gz
+
+  upload_pypi:
+    name: Upload wheels to pypi
+    runs-on: ubuntu-latest
+    needs: [build_wheels, build_sdist]
+    environment:
+      name: release
+      url: https://pypi.org/project/qmb/
+    permissions:
+      id-token: write
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
+
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: cibw-*
+          path: dist
+          merge-multiple: true
+      - uses: pypa/gh-action-pypi-publish@release/v1

.gitignore

+__pycache__
+build
+env
+dist
+_version.py
+.idea
+*.so
+*.hdf5
+*.npy
+*.pt
+*.log
+*.egg-info

.pre-commit-config.yaml

+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.6.0
+  hooks:
+  - id: check-added-large-files
+  - id: check-ast
+  - id: check-byte-order-marker
+  - id: check-builtin-literals
+  - id: check-case-conflict
+  - id: check-docstring-first
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-shebang-scripts-are-executable
+  - id: pretty-format-json
+  - id: check-merge-conflict
+  - id: check-symlinks
+  - id: check-toml
+  - id: check-vcs-permalinks
+  - id: check-xml
+  - id: check-yaml
+  - id: debug-statements
+  - id: destroyed-symlinks
+  # - id: detect-aws-credentials
+  - id: detect-private-key
+  # - id: double-quote-string-fixer
+  - id: end-of-file-fixer
+  - id: file-contents-sorter
+  - id: fix-byte-order-marker
+  # - id: fix-encoding-pragma
+  - id: forbid-new-submodules
+  - id: forbid-submodules
+  - id: mixed-line-ending
+  - id: name-tests-test
+  # - id: no-commit-to-branch
+  - id: requirements-txt-fixer
+  - id: sort-simple-yaml
+  - id: trailing-whitespace
+
+- repo: https://github.com/pre-commit/mirrors-clang-format
+  rev: v19.1.0
+  hooks:
+  - id: clang-format
+    types_or: [c++, c, cuda]
+
+- repo: https://github.com/pre-commit/mirrors-jshint
+  rev: v2.13.6
+  hooks:
+  - id: jshint
+
+- repo: https://github.com/google/yapf
+  rev: v0.40.2
+  hooks:
+  - id: yapf
+
+- repo: https://github.com/Lucas-C/pre-commit-hooks
+  rev: v1.5.5
+  hooks:
+  - id: remove-crlf
+  - id: remove-tabs

PKG-INFO

+Metadata-Version: 2.1
+Name: qmb
+Version: 0.0.10
+Summary: Quantum Manybody Problem
+Author-email: Hao Zhang <[e-mail]>
+License: GPLv3
+Project-URL: Homepage, https://github.com/USTC-KnowledgeComputingLab/qmb
+Project-URL: Documentation, https://github.com/USTC-KnowledgeComputingLab/qmb
+Project-URL: Repository, https://github.com/USTC-KnowledgeComputingLab/qmb.git
+Project-URL: Issues, https://github.com/USTC-KnowledgeComputingLab/qmb/issues
+Keywords: quantum,manybody,quantum-chemistry,quantum-simulation,molecular-simulation,algorithms,simulation,wave-function,ground-state,ansatz,torch,pytorch
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
+Classifier: Programming Language :: C++
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: torch
+Requires-Dist: tyro
+Requires-Dist: openfermion
+
+# Quantum-Many-Body

pyproject.toml

+[build-system]
+requires = ["setuptools>=69.5", "setuptools_scm>=8.1", "pybind11>=2.13"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "qmb"
+dynamic = ["version"]
+dependencies = ["numpy", "scipy", "torch", "tyro", "openfermion"]
+requires-python = ">=3.12"
+authors = [{ name = "Hao Zhang", email = "[e-mail]" }]
+description = "Quantum Manybody Problem"
+readme = "README.md"
+license = {text = "GPLv3"}
+keywords = ["quantum", "manybody", "quantum-chemistry", "quantum-simulation", "molecular-simulation", "algorithms", "simulation", "wave-function", "ground-state", "ansatz", "torch", "pytorch"]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Environment :: Console",
+  "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)",
+  "Programming Language :: C++",
+  "Programming Language :: Python :: 3",
+  "Topic :: Scientific/Engineering :: Mathematics",
+  "Topic :: Scientific/Engineering :: Physics",
+  "Topic :: Scientific/Engineering :: Chemistry",
+  "Topic :: Software Development :: Libraries :: Python Modules",
+]
+
+[project.urls]
+Homepage = "https://github.com/USTC-KnowledgeComputingLab/qmb"
+Documentation = "https://github.com/USTC-KnowledgeComputingLab/qmb"
+Repository = "https://github.com/USTC-KnowledgeComputingLab/qmb.git"
+Issues = "https://github.com/USTC-KnowledgeComputingLab/qmb/issues"
+
+[project.scripts]
+qmb = "qmb.__main__:main"
+
+[tool.setuptools.packages.find]
+include = ["qmb"]
+
+[tool.setuptools_scm]
+version_file = "qmb/_version.py"
+version_scheme = "no-guess-dev"
+
+[tool.yapf]
+based_on_style = "google"
+column_limit = 200
+
+[tool.cibuildwheel.linux]
+archs = ["x86_64", "i686", "aarch64", "ppc64le", "s390x"]
+
+[tool.cibuildwheel.macos]
+archs = ["x86_64", "arm64", "universal2"]
+
+[tool.cibuildwheel.windows]
+archs = ["AMD64", "x86", "ARM64"]

qmb.egg-info/dependency_links.txt

+

qmb.egg-info/entry_points.txt

+[console_scripts]
+qmb = qmb.__main__:main

qmb.egg-info/PKG-INFO

+Metadata-Version: 2.1
+Name: qmb
+Version: 0.0.10
+Summary: Quantum Manybody Problem
+Author-email: Hao Zhang <[e-mail]>
+License: GPLv3
+Project-URL: Homepage, https://github.com/USTC-KnowledgeComputingLab/qmb
+Project-URL: Documentation, https://github.com/USTC-KnowledgeComputingLab/qmb
+Project-URL: Repository, https://github.com/USTC-KnowledgeComputingLab/qmb.git
+Project-URL: Issues, https://github.com/USTC-KnowledgeComputingLab/qmb/issues
+Keywords: quantum,manybody,quantum-chemistry,quantum-simulation,molecular-simulation,algorithms,simulation,wave-function,ground-state,ansatz,torch,pytorch
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
+Classifier: Programming Language :: C++
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: torch
+Requires-Dist: tyro
+Requires-Dist: openfermion
+
+# Quantum-Many-Body

qmb.egg-info/requires.txt

+numpy
+scipy
+torch
+tyro
+openfermion

qmb.egg-info/SOURCES.txt

+.clang-format
+.gitignore
+.pre-commit-config.yaml
+LICENSE.md
+README.md
+pyproject.toml
+setup.py
+.github/workflows/pre-commit.yml
+.github/workflows/wheels.yml
+qmb/__init__.py
+qmb/_version.py
+qmb/version.py
+qmb.egg-info/PKG-INFO
+qmb.egg-info/SOURCES.txt
+qmb.egg-info/dependency_links.txt
+qmb.egg-info/entry_points.txt
+qmb.egg-info/requires.txt
+qmb.egg-info/top_level.txt

qmb.egg-info/top_level.txt

+qmb

README.md

+# Quantum-Many-Body

setup.cfg

+[egg_info]
+tag_build = 
+tag_date = 0
+

setup.py

+import os
+import glob
+from setuptools import setup
+from pybind11.setup_helpers import Pybind11Extension
+
+cpp_files = glob.glob(os.path.join("qmb", "*.cpp"))
+
+ext_modules = [Pybind11Extension(
+    f"qmb.{os.path.splitext(os.path.basename(cpp_file))[0]}",
+    [cpp_file],
+    cxx_std=17,
+) for cpp_file in cpp_files]
+
+setup(ext_modules=ext_modules)

qmb/_version.py

@@ -15,3 +15,3 @@ # file generated by setuptools_scm
 
-__version__ = version = '0.0.9'
-__version_tuple__ = version_tuple = (0, 0, 9)
+__version__ = version = '0.0.10'
+__version_tuple__ = version_tuple = (0, 0, 10)

entry_points.txt

-[console_scripts]
-qmb = qmb.__main__:main

METADATA

-Metadata-Version: 2.1
-Name: qmb
-Version: 0.0.9
-Summary: Quantum Manybody Problem
-Author-email: Hao Zhang <[e-mail]>
-License: GPLv3
-Project-URL: Homepage, https://github.com/USTC-KnowledgeComputingLab/qmb
-Project-URL: Documentation, https://github.com/USTC-KnowledgeComputingLab/qmb
-Project-URL: Repository, https://github.com/USTC-KnowledgeComputingLab/qmb.git
-Project-URL: Issues, https://github.com/USTC-KnowledgeComputingLab/qmb/issues
-Keywords: quantum,manybody,quantum-chemistry,quantum-simulation,molecular-simulation,algorithms,simulation,wave-function,ground-state,ansatz,torch,pytorch
-Classifier: Development Status :: 4 - Beta
-Classifier: Environment :: Console
-Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
-Classifier: Programming Language :: C++
-Classifier: Programming Language :: Python :: 3
-Classifier: Topic :: Scientific/Engineering :: Mathematics
-Classifier: Topic :: Scientific/Engineering :: Physics
-Classifier: Topic :: Scientific/Engineering :: Chemistry
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.12
-Description-Content-Type: text/markdown
-License-File: LICENSE.md
-Requires-Dist: numpy
-Requires-Dist: scipy
-Requires-Dist: torch
-Requires-Dist: tyro
-Requires-Dist: openfermion
-
-# Quantum-Many-Body

qmb/__main__.py

-import tyro
-from . import learn as _
-from . import vmc as _
-from . import iter as _
-from .subcommand_dict import subcommand_dict
-
-
-def main():
-    tyro.extras.subcommand_cli_from_dict(subcommand_dict).main()
-
-
-if __name__ == "__main__":
-    main()

qmb/_binary_tree.hpp

-#ifndef QMB_BINARY_TREE
-#define QMB_BINARY_TREE
-
-#include <memory>
-
-// Binary tree, with each leaf associated with a value,
-// and we can set and get leaf value by an iterator.
-template<typename T, T miss>
-class Tree {
-    std::unique_ptr<Tree<T, miss>> _left;
-    std::unique_ptr<Tree<T, miss>> _right;
-    std::unique_ptr<T> _value;
-
-    T& value() {
-        if (!_value) {
-            _value = std::make_unique<T>(miss);
-        }
-        return *_value;
-    }
-
-
-  public:
-    template<typename It>
-    void set(It begin, It end, T v) {
-        if (begin == end) {
-            value() = v;
-        } else {
-            if (*(begin++)) {
-                if (!_right) {
-                    _right = std::make_unique<Tree<T, miss>>();
-                }
-                _right->set(begin, end, v);
-            } else {
-                if (!_left) {
-                    _left = std::make_unique<Tree<T, miss>>();
-                }
-                _left->set(begin, end, v);
-            }
-        }
-    }
-
-    template<typename It>
-    T get(It begin, It end) {
-        if (begin == end) {
-            return value();
-        } else {
-            if (*(begin++)) {
-                if (_right) {
-                    return _right->get(begin, end);
-                } else {
-                    return miss;
-                }
-            } else {
-                if (_left) {
-                    return _left->get(begin, end);
-                } else {
-                    return miss;
-                }
-            }
-        }
-    }
-};
-
-#endif

qmb/_ising.cpp

-#include <cstdint>
-#include <pybind11/complex.h>
-#include <pybind11/numpy.h>
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
-#include <tuple>
-#include <vector>
-
-#include "_binary_tree.hpp"
-
-namespace py = pybind11;
-
-// Hamiltonian handle.
-class Hamiltonian {
-    using Coef = std::complex<double>;
-
-    Coef X, Y, Z, XX, YY, ZZ;
-
-    // Convert c++ vector to numpy array, with given shape.
-    template<typename T>
-    static py::array_t<T> vector_to_array(const std::vector<T>& vec, std::vector<int64_t> shape) {
-        py::array_t<T> result(shape);
-        auto result_buffer = result.request();
-        T* ptr = static_cast<T*>(result_buffer.ptr);
-        std::copy(vec.begin(), vec.end(), ptr);
-        return result;
-    }
-
-  public:
-    Hamiltonian(Coef _X, Coef _Y, Coef _Z, Coef _XX, Coef _YY, Coef _ZZ) : X(_X), Y(_Y), Z(_Z), XX(_XX), YY(_YY), ZZ(_ZZ) { }
-
-    template<bool outside>
-    auto call(const py::array_t<int64_t, py::array::c_style>& configs) {
-        // Configs is usually a numpy array, with rank of 2.
-        py::buffer_info configs_buf = configs.request();
-        const int64_t batch = configs_buf.shape[0];
-        const int64_t sites = configs_buf.shape[1];
-        int64_t* configs_ptr = static_cast<int64_t*>(configs_buf.ptr);
-
-        // config dict map every config to a index, default is -1 for missing configs.
-        Tree<int64_t, -1> config_dict;
-        // the prime configs count, which is at least batch size, since we will insert given configs first of all.
-        int64_t prime_count = batch;
-        // The prime configs array, used in outside mode.
-        std::vector<int64_t> config_j_pool;
-        // The indices_i_and_j and coefs is the main body of sparse matrix.
-        std::vector<int64_t> indices_i_and_j;
-        std::vector<std::complex<double>> coefs;
-        // config j in temperary vector, allocate it here for better performance
-        std::vector<int64_t> config_j(sites);
-
-        // Set input configs index in configs dict.
-        for (int64_t i = 0; i < batch; ++i) {
-            config_dict.set(&configs_ptr[i * sites], &configs_ptr[(i + 1) * sites], i);
-        }
-        // In outside mode, we need prime config array, so create it by copy the given configuration as the first batch size configs.
-        if constexpr (outside) {
-            config_j_pool.resize(batch * sites);
-            for (int64_t i = 0; i < batch * sites; ++i) {
-                config_j_pool[i] = configs_ptr[i];
-            }
-        }
-
-        // Loop over every batch and every hamiltonian term
-        for (int64_t index_i = 0; index_i < batch; ++index_i) {
-            for (int64_t type = 0; type < 6; ++type) {
-                for (int64_t site = 0; site < sites; ++site) {
-                    // type 0, 1, 2: X, Y, Z on site
-                    // type 3, 4, 5: XX, YY, ZZ on site and site-1
-                    if ((type >= 3) && (site == 0)) {
-                        continue;
-                    }
-                    Coef coef;
-                    switch (type) {
-                    case (0):
-                        coef = X;
-                        break;
-                    case (1):
-                        coef = Y;
-                        break;
-                    case (2):
-                        coef = Z;
-                        break;
-                    case (3):
-                        coef = XX;
-                        break;
-                    case (4):
-                        coef = YY;
-                        break;
-                    case (5):
-                        coef = ZZ;
-                        break;
-                    }
-                    if (coef == Coef()) {
-                        continue;
-                    }
-                    // Prepare config j to be operated by hamiltonian term
-                    for (int64_t i = 0; i < sites; ++i) {
-                        config_j[i] = configs_ptr[index_i * sites + i];
-                    }
-                    Coef param;
-                    switch (type) {
-                    case (0):
-                        param = 1;
-                        config_j[site] = 1 - config_j[site];
-                        break;
-                    case (1):
-                        param = config_j[site] == 0 ? Coef(0, +1) : Coef(0, -1);
-                        config_j[site] = 1 - config_j[site];
-                        break;
-                    case (2):
-                        param = config_j[site] == 0 ? +1 : -1;
-                        break;
-                    case (3):
-                        param = 1;
-                        config_j[site] = 1 - config_j[site];
-                        config_j[site - 1] = 1 - config_j[site - 1];
-                        break;
-                    case (4):
-                        param = config_j[site - 1] == config_j[site] ? -1 : +1;
-                        config_j[site] = 1 - config_j[site];
-                        config_j[site - 1] = 1 - config_j[site - 1];
-                        break;
-                    case (5):
-                        param = config_j[site - 1] == config_j[site] ? +1 : -1;
-                        break;
-                    }
-
-                    // Find the index j first
-                    int64_t index_j = config_dict.get(config_j.begin(), config_j.end());
-                    if (index_j == -1) {
-                        // If index j not found
-                        if constexpr (outside) {
-                            // Insert it to prime config pool in outside mode
-                            int64_t size = config_j_pool.size();
-                            config_j_pool.resize(size + sites);
-                            for (int64_t i = 0; i < sites; ++i) {
-                                config_j_pool[i + size] = config_j[i];
-                            }
-                            index_j = prime_count;
-                            config_dict.set(config_j.begin(), config_j.end(), prime_count++);
-                        } else {
-                            // Continue to next batch or hamiltonian term in inside mode.
-                            continue;
-                        }
-                    }
-                    indices_i_and_j.push_back(index_i);
-                    indices_i_and_j.push_back(index_j);
-                    coefs.push_back(coef * param);
-                }
-            }
-        }
-
-        int64_t term_count = coefs.size();
-        if constexpr (outside) {
-            return std::make_tuple(
-                vector_to_array(indices_i_and_j, {term_count, 2}),
-                vector_to_array(coefs, {term_count}),
-                vector_to_array(config_j_pool, {prime_count, sites})
-            );
-        } else {
-            return std::make_tuple(vector_to_array(indices_i_and_j, {term_count, 2}), vector_to_array(coefs, {term_count}));
-        }
-    }
-};
-
-PYBIND11_MODULE(_ising, m) {
-    py::class_<Hamiltonian>(m, "Hamiltonian", py::module_local())
-        .def(
-            py::init<
-                std::complex<double>,
-                std::complex<double>,
-                std::complex<double>,
-                std::complex<double>,
-                std::complex<double>,
-                std::complex<double>>(),
-            py::arg("X"),
-            py::arg("Y"),
-            py::arg("Z"),
-            py::arg("XX"),
-            py::arg("YY"),
-            py::arg("ZZ")
-        )
-        .def("inside", &Hamiltonian::call<false>, py::arg("configs"))
-        .def("outside", &Hamiltonian::call<true>, py::arg("configs"));
-}

qmb/_ising.pyi

-import numpy as np
-
-__all__ = ['Hamiltonian']
-
-
-class Hamiltonian:
-
-    def __init__(self, X: complex, Y: complex, Z: complex, XX: complex, YY: complex, ZZ: complex) -> None:
-        ...
-
-    def inside(self, configs: np.ndarray[np.int64]) -> tuple[np.ndarray[np.int64], np.ndarray[np.complex128]]:
-        ...
-
-    def outside(self, configs: np.ndarray[np.int64]) -> tuple[np.ndarray[np.int64], np.ndarray[np.complex128], np.ndarray[np.int64]]:
-        ...

qmb/_openfermion.cpp

-#include <array>
-#include <cstdint>
-#include <pybind11/complex.h>
-#include <pybind11/numpy.h>
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
-#include <tuple>
-#include <utility>
-#include <vector>
-
-#include "_binary_tree.hpp"
-
-namespace py = pybind11;
-
-// Hamiltonian handle for openfermion data.
-// Every term of hamiltonian is operators less or equal than 4,
-// so we use std::pair<Site, Type> to represent the term.
-class Hamiltonian {
-    using Coef = std::complex<double>;
-    using Site = int16_t;
-    using Type = int16_t; // 0 for empty, 1 for annihilation, 2 for creation
-    using Op = std::pair<Site, Type>;
-    using Ops = std::array<Op, 4>;
-    using Term = std::pair<Ops, Coef>;
-    std::vector<Term> terms;
-
-    // Convert c++ vector to numpy array, with given shape.
-    template<typename T>
-    static py::array_t<T> vector_to_array(const std::vector<T>& vec, std::vector<int64_t> shape) {
-        py::array_t<T> result(shape);
-        auto result_buffer = result.request();
-        T* ptr = static_cast<T*>(result_buffer.ptr);
-        std::copy(vec.begin(), vec.end(), ptr);
-        return result;
-    }
-
-  public:
-    // Analyze openfermion data, which is converted from python,
-    // It may be slow, but we only need to call this constructor once every process so it is ok.
-    Hamiltonian(const std::vector<std::tuple<std::vector<std::pair<int, int>>, std::complex<double>>>& openfermion_hamiltonian) {
-        for (const auto& [openfermion_ops, coef] : openfermion_hamiltonian) {
-            Ops ops;
-            size_t i = 0;
-            for (; i < openfermion_ops.size(); ++i) {
-                ops[i].first = openfermion_ops[i].first;
-                ops[i].second = 1 + openfermion_ops[i].second; // in openfermion, 0 for annihilation, 1 for creation
-            }
-            for (; i < 4; ++i) {
-                ops[i].second = 0; // 0 empty
-            }
-            terms.emplace_back(ops, coef);
-        }
-    }
-
-    template<bool outside>
-    auto call(const py::array_t<int64_t, py::array::c_style>& configs) {
-        // Configs is usually a numpy array, with rank of 2.
-        py::buffer_info configs_buf = configs.request();
-        const int64_t batch = configs_buf.shape[0];
-        const int64_t sites = configs_buf.shape[1];
-        int64_t* configs_ptr = static_cast<int64_t*>(configs_buf.ptr);
-
-        // config dict map every config to a index, default is -1 for missing configs.
-        Tree<int64_t, -1> config_dict;
-        // the prime configs count, which is at least batch size, since we will insert given configs first of all.
-        int64_t prime_count = batch;
-        // The prime configs array, used in outside mode.
-        std::vector<int64_t> config_j_pool;
-        // The indices_i_and_j and coefs is the main body of sparse matrix.
-        std::vector<int64_t> indices_i_and_j;
-        std::vector<std::complex<double>> coefs;
-        // config j in temperary vector, allocate it here for better performance
-        std::vector<int64_t> config_j(sites);
-
-        // Set input configs index in configs dict.
-        for (int64_t i = 0; i < batch; ++i) {
-            config_dict.set(&configs_ptr[i * sites], &configs_ptr[(i + 1) * sites], i);
-        }
-        // In outside mode, we need prime config array, so create it by copy the given configuration as the first batch size configs.
-        if constexpr (outside) {
-            config_j_pool.resize(batch * sites);
-            for (int64_t i = 0; i < batch * sites; ++i) {
-                config_j_pool[i] = configs_ptr[i];
-            }
-        }
-
-        // Loop over every batch and every hamiltonian term
-        for (int64_t index_i = 0; index_i < batch; ++index_i) {
-            for (const auto& [ops, coef] : terms) {
-                // Prepare config j to be operated by hamiltonian term
-                for (int64_t i = 0; i < sites; ++i) {
-                    config_j[i] = configs_ptr[index_i * sites + i];
-                }
-                bool success = true;
-                bool parity = false;
-                // Apply operator one by one
-                for (auto i = 4; i-- > 0;) {
-                    auto [site, operation] = ops[i];
-                    if (operation == 0) {
-                        // Empty operator, nothing happens
-                        continue;
-                    } else if (operation == 1) {
-                        // Annihilation operator
-                        if (config_j[site] != 1) {
-                            success = false;
-                            break;
-                        }
-                        config_j[site] = 0;
-                        if (std::accumulate(config_j.begin(), config_j.begin() + site, 0) % 2 == 1) {
-                            parity ^= true;
-                        }
-                    } else {
-                        // Creation operator
-                        if (config_j[site] != 0) {
-                            success = false;
-                            break;
-                        }
-                        config_j[site] = 1;
-                        if (std::accumulate(config_j.begin(), config_j.begin() + site, 0) % 2 == 1) {
-                            parity ^= true;
-                        }
-                    }
-                }
-
-                if (success) {
-                    // Success, insert this term to sparse matrix
-                    // Find the index j first
-                    int64_t index_j = config_dict.get(config_j.begin(), config_j.end());
-                    if (index_j == -1) {
-                        // If index j not found
-                        if constexpr (outside) {
-                            // Insert it to prime config pool in outside mode
-                            int64_t size = config_j_pool.size();
-                            config_j_pool.resize(size + sites);
-                            for (int64_t i = 0; i < sites; ++i) {
-                                config_j_pool[i + size] = config_j[i];
-                            }
-                            index_j = prime_count;
-                            config_dict.set(config_j.begin(), config_j.end(), prime_count++);
-                        } else {
-                            // Continue to next batch or hamiltonian term in inside mode.
-                            continue;
-                        }
-                    }
-                    indices_i_and_j.push_back(index_i);
-                    indices_i_and_j.push_back(index_j);
-                    coefs.push_back(parity ? -coef : +coef);
-                }
-            }
-        }
-
-        int64_t term_count = coefs.size();
-        if constexpr (outside) {
-            return std::make_tuple(
-                vector_to_array(indices_i_and_j, {term_count, 2}),
-                vector_to_array(coefs, {term_count}),
-                vector_to_array(config_j_pool, {prime_count, sites})
-            );
-        } else {
-            return std::make_tuple(vector_to_array(indices_i_and_j, {term_count, 2}), vector_to_array(coefs, {term_count}));
-        }
-    }
-};
-
-PYBIND11_MODULE(_openfermion, m) {
-    py::class_<Hamiltonian>(m, "Hamiltonian", py::module_local())
-        .def(py::init<std::vector<std::tuple<std::vector<std::pair<int, int>>, std::complex<double>>>>(), py::arg("openfermion_hamiltonian"))
-        .def("inside", &Hamiltonian::call<false>, py::arg("configs"))
-        .def("outside", &Hamiltonian::call<true>, py::arg("configs"));
-}

qmb/_openfermion.pyi

-import numpy as np
-
-__all__ = ['Hamiltonian']
-
-
-class Hamiltonian:
-
-    def __init__(self, openfermion_hamiltonian: list[tuple[list[tuple[int, int]], complex]]) -> None:
-        ...
-
-    def inside(self, configs: np.ndarray[np.int64]) -> tuple[np.ndarray[np.int64], np.ndarray[np.complex128]]:
-        ...
-
-    def outside(self, configs: np.ndarray[np.int64]) -> tuple[np.ndarray[np.int64], np.ndarray[np.complex128], np.ndarray[np.int64]]:
-        ...

qmb/attention.py

-# This file implements auto regressive transformers network for electron system.
-
-import torch
-
-
-class FeedForward(torch.nn.Module):
-    """
-    Feed forward layer in transformers.
-    """
-
-    def __init__(self, embedding_dim: int, hidden_dim: int) -> None:
-        super().__init__()
-        self.model: torch.nn.Module = torch.nn.Sequential(
-            torch.nn.LayerNorm(embedding_dim),
-            torch.nn.Linear(embedding_dim, hidden_dim),
-            torch.nn.GELU(),
-            torch.nn.Linear(hidden_dim, embedding_dim),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: batch * site * embedding
-        x = self.model(x)
-        # x: batch * site * embedding
-        return x
-
-
-class SelfAttention(torch.nn.Module):
-    """
-    Self attention unit with kv cache support.
-    """
-
-    def __init__(self, embedding_dim: int, heads_num: int) -> None:
-        super().__init__()
-
-        self.heads_num: int = heads_num
-        self.heads_dim: int = embedding_dim // heads_num
-        assert self.heads_num * self.heads_dim == embedding_dim
-
-        self.norm: torch.nn.Module = torch.nn.LayerNorm(embedding_dim)
-
-        self.qkv: torch.nn.Module = torch.nn.Linear(embedding_dim, 3 * embedding_dim)
-        self.out: torch.nn.Module = torch.nn.Linear(embedding_dim, embedding_dim)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        kv_cache: tuple[torch.Tensor, torch.Tensor] | None,
-        mask: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
-        # x: batch * site * embedding
-        x = self.norm(x)
-        # x: batch * site * embedding
-        batch_size, sites, embedding_dim = x.shape
-        q, k, v = self.qkv(x).split(embedding_dim, dim=-1)
-        # q, k, v: batch * site * embedding
-        q = q.view([batch_size, sites, self.heads_num, self.heads_dim])
-        k = k.view([batch_size, sites, self.heads_num, self.heads_dim])
-        v = v.view([batch_size, sites, self.heads_num, self.heads_dim])
-        # q, k, v: batch, site, heads_num, heads_dim
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-        # q, k, v: batch, heads_num, site, heads_dim
-        if kv_cache is not None:
-            k_cache, v_cache = kv_cache
-            k = torch.cat([k_cache, k], dim=2)
-            v = torch.cat([v_cache, v], dim=2)
-        # q: batch, heads_num, site, heads_dim
-        # k, v: batch, heads_num, total_site, heads_dim
-        if mask is None:
-            total_sites = k.shape[-2]
-            mask = torch.ones(sites, total_sites, dtype=torch.bool, device=x.device).tril(diagonal=total_sites - sites)
-        # call scaled dot product attention
-        attn = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-        # attn: batch, heads_num, site, heads_dim
-        out = attn.transpose(1, 2).reshape([batch_size, sites, embedding_dim])
-        # out: batch, site, embedding_dim
-        return self.out(out), (k, v)
-
-
-class DecoderUnit(torch.nn.Module):
-    """
-    Decoder unit in transformers, containing self attention and feed forward.
-    """
-
-    def __init__(self, embedding_dim: int, heads_num: int, feed_forward_dim: int) -> None:
-        super().__init__()
-        self.attention: torch.nn.Module = SelfAttention(embedding_dim, heads_num)
-        self.feed_forward: torch.nn.Module = FeedForward(embedding_dim, feed_forward_dim)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        kv_cache: tuple[torch.Tensor, torch.Tensor] | None,
-        mask: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
-        # x, y: batch * site * embedding
-        y, result_cache = self.attention(x, kv_cache, mask)
-        x = x + y
-        y = self.feed_forward(x)
-        x = x + y
-        return x, result_cache
-
-
-class Transformers(torch.nn.Module):
-
-    def __init__(self, embedding_dim: int, heads_num: int, feed_forward_dim: int, depth: int) -> None:
-        super().__init__()
-        self.layers: torch.nn.ModuleList = torch.nn.ModuleList(DecoderUnit(embedding_dim, heads_num, feed_forward_dim) for _ in range(depth))
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        kv_cache: list[tuple[torch.Tensor, torch.Tensor]] | None,
-        mask: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, torch.Tensor]]]:
-        # x: batch * site * embedding
-        result_cache: list[tuple[torch.Tensor, torch.Tensor]] = []
-        for i, layer in enumerate(self.layers):
-            if kv_cache is None:
-                x, cache = layer(x, None, mask)
-            else:
-                x, cache = layer(x, kv_cache[i], mask)
-            result_cache.append(cache)
-        return x, result_cache
-
-
-class Tail(torch.nn.Module):
-    """
-    The tail layer for transformers.
-    """
-
-    def __init__(self, embedding_dim: int, hidden_dim: int, output_dim: int) -> None:
-        super().__init__()
-        self.model: torch.nn.Module = torch.nn.Sequential(
-            torch.nn.LayerNorm(embedding_dim),
-            torch.nn.Linear(embedding_dim, hidden_dim),
-            torch.nn.GELU(),
-            torch.nn.Linear(hidden_dim, output_dim),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: batch * site * embedding
-        x = self.model(x)
-        # x: batch * site * embedding
-        return x
-
-
-class Embedding(torch.nn.Module):
-    """
-    Embedding layer for transformers.
-    """
-
-    def __init__(self, sites: int, physical_dim: int, embedding_dim: int) -> None:
-        super().__init__()
-        self.parameter: torch.nn.Parameter = torch.nn.Parameter(torch.randn([sites, physical_dim, embedding_dim]))
-
-    def forward(self, x: torch.Tensor, base: int) -> torch.Tensor:
-        # x: batch * sites
-        x = x.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, -1, self.parameter.shape[-1])
-        # x: batch * sites * config=1 * embedding
-
-        # param: sites * config * embedding
-        parameter = self.parameter[base:][:x.shape[1]].unsqueeze(0).expand(x.shape[0], -1, -1, -1)
-        # param: batch * sites * config * embedding
-
-        result = torch.gather(parameter, -2, x)
-        # result: batch * site * 1 * embedding
-
-        return result.squeeze(-2)
-
-
-class WaveFunction(torch.nn.Module):
-
-    def __init__(
-            self,
-            *,
-            double_sites: int,  # qubits number, qubits are grouped by two for each site in naqs, so name it as double sites
-            physical_dim: int,  # is always 2 for naqs
-            is_complex: bool,  # is always true for naqs
-            spin_up: int,  # spin up number
-            spin_down: int,  # spin down number
-            embedding_dim: int,  # embedding dim in transformers
-            heads_num: int,  # heads number in transformers
-            feed_forward_dim: int,  # feed forward dim and tail dim in transformers
-            depth: int,  # depth of transformers
-            ordering: int | list[int],  # ordering of sites +1 for normal order, -1 for reversed order, or the order list directly
-    ) -> None:
-        super().__init__()
-        assert double_sites % 2 == 0
-        self.double_sites: int = double_sites
-        self.sites: int = double_sites // 2
-        assert physical_dim == 2
-        assert is_complex == True
-        self.spin_up: int = spin_up
-        self.spin_down: int = spin_down
-        self.embedding_dim: int = embedding_dim
-        self.heads_num: int = heads_num
-        self.feed_forward_dim: int = feed_forward_dim
-        self.depth: int = depth
-
-        # embedding configurations on every sites(each sites contain two qubits)
-        self.embedding: torch.nn.Module = Embedding(self.sites, 4, self.embedding_dim)  # spin_up * spin_down
-        # main body
-        self.transformers: torch.nn.Module = Transformers(self.embedding_dim, self.heads_num, self.feed_forward_dim, self.depth)
-        # tail, mapping from embedding space to amplitude and phase space.
-        self.tail: torch.nn.Module = Tail(self.embedding_dim, self.feed_forward_dim, 8)  # (amplitude and phase) * (4 configs)
-
-        # ordering of sites +1 for normal order, -1 for reversed order
-        if isinstance(ordering, int) and ordering == +1:
-            ordering = list(range(self.sites))
-        if isinstance(ordering, int) and ordering == -1:
-            ordering = list(reversed(range(self.sites)))
-        self.register_buffer('ordering', torch.tensor(ordering, dtype=torch.int64))
-        self.register_buffer('ordering_reversed', torch.scatter(torch.zeros(self.sites, dtype=torch.int64), 0, self.ordering, torch.arange(self.sites, dtype=torch.int64)))
-
-        # used to get device and dtype
-        self.dummy_param = torch.nn.Parameter(torch.empty(0))
-
-    @torch.jit.export
-    def mask(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Determined whether we could append spin up or spin down after uncompleted configurations.
-        """
-        # x : batch_size * current_site * 2
-        # x are the uncompleted configurations
-        batch_size = x.shape[0]
-        current_site = x.shape[1]
-        # number : batch_size * 2
-        # number is the total electron number of uncompleted configurations
-        number = x.sum(dim=1)
-
-        # up/down_electron/hole : batch_size
-        # the electron and hole number of uncompleted configurations
-        up_electron = number[:, 0]
-        down_electron = number[:, 1]
-        up_hole = current_site - up_electron
-        down_hole = current_site - down_electron
-
-        # add_up/down_electron/hole : batch_size
-        # whether able to append up/down electron/hole
-        add_up_electron = up_electron < self.spin_up
-        add_down_electron = down_electron < self.spin_down
-        add_up_hole = up_hole < self.sites - self.spin_up
-        add_down_hole = down_hole < self.sites - self.spin_down
-
-        # add_up : batch_size * 2 * 1
-        # add_down : batch_size * 1 * 2
-        add_up = torch.stack([add_up_hole, add_up_electron], dim=-1).unsqueeze(-1)
-        add_down = torch.stack([add_down_hole, add_down_electron], dim=-1).unsqueeze(-2)
-        # add : batch_size * 2 * 2
-        add = torch.logical_and(add_up, add_down)
-        # add contains whether to append up/down electron/hole after uncompleted configurations
-        # add[_, 0, 0] means we could add up hole and down hole
-        # add[_, 0, 1] means we could add up hole and down electron
-        # add[_, 1, 0] means we could add up electron and down hole
-        # add[_, 1, 1] means we could add up electron and down electron
-        return add
-
-    @torch.jit.export
-    def normalize_amplitude(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Normalize uncompleted log amplitude.
-        """
-        # x : ... * 2 * 2
-        # param :  ...
-        param = (2 * x).exp().sum(dim=[-2, -1]).log() / 2
-        x = x - param.unsqueeze(-1).unsqueeze(-1)
-        # 1 = param = sqrt(sum(x.exp()^2)) now
-        return x
-
-    @torch.jit.export
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Calculate psi of given configurations.
-        """
-        device: torch.device = self.dummy_param.device
-        dtype: torch.dtype = self.dummy_param.dtype
-
-        batch_size: int = x.shape[0]
-        # x : batch_size * sites * 2
-        x = x.reshape([batch_size, self.sites, 2])
-        # apply ordering
-        x = torch.index_select(x, 1, self.ordering_reversed)
-
-        # prepare x4 as the input of the network
-        # the last one is not needed, and zeros vector prepend at the beginning
-        x4: torch.Tensor = x[:, :-1, 0] * 2 + x[:, :-1, 1]
-        x4 = torch.cat([torch.zeros([batch_size, 1], device=device, dtype=torch.int64), x4], dim=1)
-        # x4: batch_size * sites
-        emb: torch.Tensor = self.embedding(x4, 0)
-        # emb: batch_size * sites * embedding
-        # in embedding layer, 0 for bos, 1 for site 0, ...
-        post_transformers, _ = self.transformers(emb, None, None)  # batch * sites * embedding
-        tail: torch.Tensor = self.tail(post_transformers)  # batch * sites * 8
-        # amplitude/phase : batch * sites * 2 * 2
-        amplitude: torch.Tensor = tail[:, :, :4].reshape(batch_size, self.sites, 2, 2)
-        phase: torch.Tensor = tail[:, :, 4:].reshape(batch_size, self.sites, 2, 2)
-        # filter mask for amplitude
-        amplitude: torch.Tensor = amplitude + torch.stack([torch.where(self.mask(x[:, :i]), 0, -torch.inf) for i in range(self.sites)], dim=1)
-        # normalize amplitude
-        amplitude: torch.Tensor = self.normalize_amplitude(amplitude)
-
-        # batch/sites_indices: batch * sites
-        batch_indices: torch.Tensor = torch.arange(batch_size).unsqueeze(1).expand(-1, self.sites)
-        sites_indices: torch.Tensor = torch.arange(self.sites).unsqueeze(0).expand(batch_size, -1)
-
-        # amplitude/phase: batch * sites first, after sum over dim=1, they are amplitude/phase: batch
-        amplitude: torch.Tensor = amplitude[batch_indices, sites_indices, x[:, :, 0], x[:, :, 1]].sum(dim=1)
-        phase: torch.Tensor = phase[batch_indices, sites_indices, x[:, :, 0], x[:, :, 1]].sum(dim=1)
-        return torch.view_as_complex(torch.stack([amplitude, phase], dim=-1)).exp()
-
-    @torch.jit.export
-    def binomial(self, count: torch.Tensor, probability: torch.Tensor) -> torch.Tensor:
-        """
-        Binomial sampling with given count and probability
-        """
-        # clamp probability
-        probability = torch.clamp(probability, min=0, max=1)
-        # set probability to zero for count = 0 since it may be nan when count = 0
-        probability = torch.where(count == 0, 0, probability)
-        # create dist and sample
-        result = torch.binomial(count, probability).to(dtype=torch.int64)
-        # numerical error since result is cast from float.
-        return torch.clamp(result, min=torch.zeros_like(count), max=count)
-
-    @torch.jit.export
-    def generate_unique(self, batch_size: int) -> tuple[torch.Tensor, torch.Tensor, None, None]:
-        """
-        Generate configurations uniquely.
-        see https://arxiv.org/pdf/2408.07625
-        """
-        device: torch.device = self.dummy_param.device
-        dtype: torch.dtype = self.dummy_param.dtype
-
-        cache: list[tuple[torch.Tensor, torch.Tensor]] | None = None
-
-        # x : local_batch_size * current_site * 2
-        x: torch.Tensor = torch.zeros([1, 1, 2], device=device, dtype=torch.int64)  # site=1, since the first is bos
-        # (un)perturbed_log_probability : local_batch_size
-        unperturbed_probability: torch.Tensor = torch.tensor([0], dtype=dtype, device=device)
-        perturbed_probability: torch.Tensor = torch.tensor([0], dtype=dtype, device=device)
-        for i in range(self.sites):
-            local_batch_size: int = x.size(0)
-
-            # xi: batch * sites=1 * 2
-            xi: torch.Tensor = x[:, -1:]
-            # xi4: batch * sites=1
-            xi4: torch.Tensor = xi[:, :, 0] * 2 + xi[:, :, 1]
-            # emb: batch * sites=1 * embedding
-            emb: torch.Tensor = self.embedding(xi4, i)
-            # post_transformers: batch * sites=1 * embedding
-            post_transformers, cache = self.transformers(emb, cache, None)
-            assert cache is not None
-            # tail: batch * sites=1 * 8
-            tail: torch.Tensor = self.tail(post_transformers)
-
-            # the first 4 item are amplitude
-            # delta_amplitude: batch * 2 * 2
-            delta_amplitude: torch.Tensor = tail[:, :, :4].reshape([local_batch_size, 2, 2])
-            # filter mask for amplitude
-            delta_amplitude: torch.Tensor = delta_amplitude + torch.where(self.mask(x[:, 1:]), 0, -torch.inf)
-            # normalize amplitude
-            delta_amplitude: torch.Tensor = self.normalize_amplitude(delta_amplitude)
-
-            # delta unperturbed prob for all batch and 4 adds
-            l: torch.Tensor = (2 * delta_amplitude).reshape([local_batch_size, 4])
-            # and add to get the current unperturbed prob
-            l: torch.Tensor = unperturbed_probability.view([-1, 1]) + l
-            # get perturbed prob
-            L: torch.Tensor = l - (-torch.rand_like(l).log()).log()
-            # get max perturbed prob
-            Z: torch.Tensor = L.max(dim=-1).values.reshape([-1, 1])
-            # evaluate the conditioned prob
-            L: torch.Tensor = -torch.log(torch.exp(-perturbed_probability.view([-1, 1])) - torch.exp(-Z) + torch.exp(-L))
-
-            # calculate appended configurations for 4 adds
-            # local_batch_size * current_site * 2 + local_batch_size * 1 * 2
-            x0: torch.Tensor = torch.cat([x, torch.tensor([[0, 0]], device=device).expand(local_batch_size, -1, -1)], dim=1)
-            x1: torch.Tensor = torch.cat([x, torch.tensor([[0, 1]], device=device).expand(local_batch_size, -1, -1)], dim=1)
-            x2: torch.Tensor = torch.cat([x, torch.tensor([[1, 0]], device=device).expand(local_batch_size, -1, -1)], dim=1)
-            x3: torch.Tensor = torch.cat([x, torch.tensor([[1, 1]], device=device).expand(local_batch_size, -1, -1)], dim=1)
-
-            # cat all configurations to get x : new_local_batch_size * current_size * 2
-            # (un)perturbed prob : new_local_batch_size
-            x = torch.cat([x0, x1, x2, x3])
-            unperturbed_probability = l.permute(1, 0).reshape([-1])
-            perturbed_probability = L.permute(1, 0).reshape([-1])
-            cache = [(kv[0].repeat(4, 1, 1, 1), kv[1].repeat(4, 1, 1, 1)) for kv in cache]
-            # kv cache: batch * heads_num * site * heads_dim, so just repeat first dimension
-
-            # filter new data, only used largest batch_size ones
-            selected = perturbed_probability.sort(descending=True).indices[:batch_size]
-            x = x[selected]
-            unperturbed_probability = unperturbed_probability[selected]
-            perturbed_probability = perturbed_probability[selected]
-            cache = [(kv[0][selected], kv[1][selected]) for kv in cache]
-
-            # if prob = 0, filter it forcely
-            selected = perturbed_probability != -torch.inf
-            x = x[selected]
-            unperturbed_probability = unperturbed_probability[selected]
-            perturbed_probability = perturbed_probability[selected]
-            cache = [(kv[0][selected], kv[1][selected]) for kv in cache]
-
-        # apply ordering
-        x = torch.index_select(x[:, 1:], 1, self.ordering)
-        # flatten site part of x
-        x = x.reshape([x.size(0), self.double_sites])
-        # it should return configurations, amplitudes, probabilities and multiplicities
-        # but it is unique generator, so last two field is none
-        return x, self(x), None, None

qmb/common.py

-import os
-import sys
-import logging
-import typing
-import pathlib
-import dataclasses
-import torch
-import tyro
-from . import openfermion
-from . import openfermion_operator
-from . import ising
-
-model_dict = {
-    "openfermion": openfermion.Model,
-    "openfermion_operator": openfermion_operator.Model,
-    "ising": ising.Model,
-}
-
-
[e-mail]
-class CommonConfig:
-    # The model name
-    model_name: typing.Annotated[str, tyro.conf.Positional, tyro.conf.arg(metavar="MODEL")]
-    # The network name
-    network_name: typing.Annotated[str, tyro.conf.Positional, tyro.conf.arg(metavar="NETWORK")]
-    # Arguments for physical model
-    physics_args: typing.Annotated[tuple[str, ...], tyro.conf.arg(aliases=["-P"]), tyro.conf.UseAppendAction] = ()
-    # Arguments for network
-    network_args: typing.Annotated[tuple[str, ...], tyro.conf.arg(aliases=["-N"]), tyro.conf.UseAppendAction] = ()
-
-    # The job name used in checkpoint and log, leave empty to use the preset job name given by the model and network
-    job_name: typing.Annotated[str | None, tyro.conf.arg(aliases=["-J"])] = None
-    # The checkpoint path
-    checkpoint_path: typing.Annotated[pathlib.Path, tyro.conf.arg(aliases=["-C"])] = pathlib.Path("checkpoints")
-    # The log path
-    log_path: typing.Annotated[pathlib.Path, tyro.conf.arg(aliases=["-L"])] = pathlib.Path("logs")
-    # The manual random seed, leave empty for set seed automatically
-    random_seed: typing.Annotated[int | None, tyro.conf.arg(aliases=["-S"])] = None
-
-    def main(self):
-        if "-h" in self.network_args or "--help" in self.network_args:
-            getattr(model_dict[self.model_name], self.network_name)(object(), self.network_args)
-        default_job_name = model_dict[self.model_name].preparse(self.physics_args)
-        if self.job_name is None:
-            self.job_name = default_job_name
-
-        os.makedirs(self.checkpoint_path, exist_ok=True)
-        os.makedirs(self.log_path, exist_ok=True)
-
-        logging.basicConfig(
-            handlers=[logging.StreamHandler(), logging.FileHandler(f"{self.log_path}/{self.job_name}.log")],
-            level=logging.INFO,
-            format=f"[%(process)d] %(asctime)s {self.job_name}({self.network_name}) %(levelname)s: %(message)s",
-        )
-
-        logging.info("%s script start, with %a", os.path.splitext(os.path.basename(sys.argv[0]))[0], sys.argv)
-        logging.info("model name: %s, network name: %s, job name: %s", self.model_name, self.network_name, self.job_name)
-        logging.info("log path: %s, checkpoint path: %s", self.log_path, self.checkpoint_path)
-        logging.info("arguments will be passed to network parser: %a", self.network_args)
-        logging.info("arguments will be passed to physics parser: %a", self.physics_args)
-
-        if self.random_seed is not None:
-            logging.info("setting random seed to %d", self.random_seed)
-            torch.manual_seed(self.random_seed)
-        else:
-            logging.info("random seed not set, using %d", torch.seed())
-
-        logging.info("disabling torch default gradient behavior")
-        torch.set_grad_enabled(False)
-
-        logging.info("loading %s model as physical model", self.model_name)
-        model = model_dict[self.model_name].parse(self.physics_args)
-        logging.info("the physical model has been loaded")
-
-        logging.info("loading network %s and create network with physical model and args %s", self.network_name, self.network_args)
-        network = getattr(model, self.network_name)(self.network_args)
-        logging.info("network created")
-
-        logging.info("trying to load checkpoint")
-        if os.path.exists(f"{self.checkpoint_path}/{self.job_name}.pt"):
-            logging.info("checkpoint found, loading")
-            network.load_state_dict(torch.load(f"{self.checkpoint_path}/{self.job_name}.pt", map_location="cpu", weights_only=True))
-            logging.info("checkpoint loaded")
-        else:
-            logging.info("checkpoint not found")
-        logging.info("moving model to cuda")
-        network.cuda()
-        logging.info("model has been moved to cuda")
-
-        return model, network

qmb/ising.py

-# This file declare transverse field ising model.
-
-import typing
-import logging
-import dataclasses
-import torch
-import tyro
-from . import _ising
-from . import naqs as naqs_m
-
-
[e-mail]
-class ModelConfig:
-    # The length of the ising chain
-    length: typing.Annotated[int, tyro.conf.Positional]
-    # The coefficient of X
-    X: typing.Annotated[float, tyro.conf.arg(aliases=["-x"])] = 0
-    # The coefficient of Y
-    Y: typing.Annotated[float, tyro.conf.arg(aliases=["-y"])] = 0
-    # The coefficient of Z
-    Z: typing.Annotated[float, tyro.conf.arg(aliases=["-z"])] = 0
-    # The coefficient of XX
-    XX: typing.Annotated[float, tyro.conf.arg(aliases=["-X"])] = 0
-    # The coefficient of YY
-    YY: typing.Annotated[float, tyro.conf.arg(aliases=["-Y"])] = 0
-    # The coefficient of ZZ
-    ZZ: typing.Annotated[float, tyro.conf.arg(aliases=["-Z"])] = 0
-
-
[e-mail]
-class NaqsConfig:
-    # The hidden widths of the network
-    hidden: typing.Annotated[tuple[int, ...], tyro.conf.arg(aliases=["-w"])] = (512,)
-
-
-class Model:
-
-    @classmethod
-    def preparse(cls, input_args):
-        args = tyro.cli(ModelConfig, args=input_args)
-        return f"Ising_L{args.length}_X{args.X}_Y{args.Y}_Z{args.Z}_XX{args.XX}_YY{args.YY}_ZZ{args.ZZ}"
-
-    @classmethod
-    def parse(cls, input_args):
-        logging.info("parsing args %a by ising model", input_args)
-        args = tyro.cli(ModelConfig, args=input_args)
-        logging.info("length: %d, X: %.10f, Y: %.10f, Z: %.10f, XX: %.10f, YY: %.10f, ZZ: %.10f", args.length, args.X, args.Y, args.Z, args.XX, args.YY, args.ZZ)
-
-        return cls(args.length, args.X, args.Y, args.Z, args.XX, args.YY, args.ZZ)
-
-    def __init__(self, length, X, Y, Z, XX, YY, ZZ):
-        self.length = length
-        self.X = X
-        self.Y = Y
-        self.Z = Z
-        self.XX = XX
-        self.YY = YY
-        self.ZZ = ZZ
-        logging.info("creating ising model with length = %d, X = %.10f, Y = %.10f, Z = %.10f, XX = %.10f, YY = %.10f, ZZ = %.10f", self.length, self.X, self.Y, self.Z, self.XX, self.YY, self.ZZ)
-
-        self.ref_energy = torch.nan
-
-        logging.info("creating ising hamiltonian handle")
-        self.hamiltonian = _ising.Hamiltonian(self.X, self.Y, self.Z, self.XX, self.YY, self.ZZ)
-        logging.info("hamiltonian handle has been created")
-
-    def inside(self, configs):
-        return self.hamiltonian.inside(configs)
-
-    def outside(self, configs):
-        return self.hamiltonian.outside(configs)
-
-    def naqs(self, input_args):
-        logging.info("parsing args %a by network naqs", input_args)
-        args = tyro.cli(NaqsConfig, args=input_args)
-        logging.info("hidden: %a", args.hidden)
-
-        network = naqs_m.WaveFunctionNormal(
-            sites=self.length,
-            physical_dim=2,
-            is_complex=True,
-            hidden_size=args.hidden,
-            ordering=+1,
-        ).double()
-
-        return torch.jit.script(network)

qmb/iter.py

-import logging
-import typing
-import dataclasses
-import numpy
-import scipy
-import torch
-import tyro
-from .common import CommonConfig
-from .subcommand_dict import subcommand_dict
-
-
[e-mail]
-class IterConfig:
-    common: typing.Annotated[CommonConfig, tyro.conf.OmitArgPrefixes]
-
-    # The sampling count
-    sampling_count: typing.Annotated[int, tyro.conf.arg(aliases=["-n"])] = 128
-
-    # The selected extended sampling count
-    selected_sampling_count: typing.Annotated[int, tyro.conf.arg(aliases=["-s"])] = 4000
-
-    def main(self):
-        model, network = self.common.main()
-
-        logging.info(
-            "sampling count: %d, selected sampling count: %d",
-            self.sampling_count,
-            self.selected_sampling_count,
-        )
-
-        logging.info("first sampling core configurations")
-        configs_core, psi_core, _, _ = network.generate_unique(self.sampling_count)
-        configs_core = configs_core.cpu()
-        psi_core = psi_core.cpu()
-        logging.info("core configurations sampled")
-
-        while True:
-            sampling_count_core = len(configs_core)
-            logging.info("core configurations count is %d", sampling_count_core)
-
-            logging.info("calculating extended configurations")
-            indices_i_and_j, values, configs_extended = model.outside(configs_core)
-            logging.info("extended configurations created")
-            sampling_count_extended = len(configs_extended)
-            logging.info("extended configurations count is %d", sampling_count_extended)
-
-            logging.info("converting sparse extending matrix data to sparse matrix")
-            hamiltonian = torch.sparse_coo_tensor(indices_i_and_j.T, values, [sampling_count_core, sampling_count_extended], dtype=torch.complex128).to_sparse_csr()
-            logging.info("sparse extending matrix created")
-
-            logging.info("estimating the importance of extended configurations")
-            importance = (psi_core.conj() * psi_core).abs() @ (hamiltonian.conj() * hamiltonian).abs()
-            importance[:sampling_count_core] += importance.max()
-            logging.info("importance of extended configurations created")
-
-            logging.info("selecting extended configurations by importance")
-            selected_indices = importance.sort(descending=True).indices[:self.selected_sampling_count].sort().values
-            logging.info("extended configurations selected indices prepared")
-
-            logging.info("selecting extended configurations")
-            configs_extended = configs_extended[selected_indices]
-            logging.info("extended configurations selected")
-            sampling_count_extended = len(configs_extended)
-            logging.info("selected extended configurations count is %d", sampling_count_extended)
-
-            logging.info("calculating sparse data of hamiltonian on extended configurations")
-            indices_i_and_j, values = model.inside(configs_extended)
-            logging.info("converting sparse matrix data to sparse matrix")
-            hamiltonian = scipy.sparse.coo_matrix((values, indices_i_and_j.T), [sampling_count_extended, sampling_count_extended], dtype=numpy.complex128).tocsr()
-            logging.info("sparse matrix on extended configurations created")
-
-            logging.info("preparing initial psi used in lobpcg")
-            psi_extended = numpy.pad(psi_core, (0, sampling_count_extended - sampling_count_core)).reshape([-1, 1])
-            logging.info("initial psi used in lobpcg has been created")
-
-            logging.info("calculating minimum energy on extended configurations")
-            energy, psi_extended = scipy.sparse.linalg.lobpcg(hamiltonian, psi_extended, largest=False, maxiter=1024)
-            logging.info("energy on extended configurations is %.10f, ref energy is %.10f, error is %.10f", energy.item(), model.ref_energy, energy.item() - model.ref_energy)
-
-            logging.info("calculating indices of new core configurations")
-            indices = numpy.argsort(numpy.abs(psi_extended).flatten())[-self.sampling_count:]
-            logging.info("indices of new core configurations has been obtained")
-
-            logging.info("update new core configurations")
-            configs_core = configs_extended[indices]
-            psi_core = torch.tensor(psi_extended[indices].flatten())
-            logging.info("new core configurations has been updated")
-
-
-subcommand_dict["iter"] = IterConfig

qmb/learn.py

-import logging
-import typing
-import dataclasses
-import numpy
-import scipy
-import torch
-import tyro
-from . import losses
-from .common import CommonConfig
-from .subcommand_dict import subcommand_dict
-
-
[e-mail]
-class LearnConfig:
-    common: typing.Annotated[CommonConfig, tyro.conf.OmitArgPrefixes]
-
-    # sampling count
-    sampling_count: typing.Annotated[int, tyro.conf.arg(aliases=["-n"])] = 4000
-    # learning rate for the local optimizer
-    learning_rate: typing.Annotated[float, tyro.conf.arg(aliases=["-r"], help_behavior_hint="(default: 1e-3 for Adam, 1 for LBFGS)")] = -1
-    # step count for the local optimizer
-    local_step: typing.Annotated[int, tyro.conf.arg(aliases=["-s"], help_behavior_hint="(default: 1000 for Adam, 400 for LBFGS)")] = -1
-    # early break loss threshold for local optimization
-    local_loss: typing.Annotated[float, tyro.conf.arg(aliases=["-t"])] = 1e-8
-    # psi count to be printed after local optimizer
-    logging_psi: typing.Annotated[int, tyro.conf.arg(aliases=["-p"])] = 30
-    # the loss function to be used
-    loss_name: typing.Annotated[str, tyro.conf.arg(aliases=["-l"])] = "log"
-    # Use LBFGS instead of Adam
-    use_lbfgs: typing.Annotated[bool, tyro.conf.arg(aliases=["-2"])] = False
-
-    def __post_init__(self):
-        if self.learning_rate == -1:
-            self.learning_rate = 1 if self.use_lbfgs else 1e-3
-        if self.local_step == -1:
-            self.local_step = 400 if self.use_lbfgs else 1000
-
-    def main(self):
-        model, network = self.common.main()
-
-        logging.info(
-            "sampling count: %d, learning rate: %f, local step: %d, local loss: %f, logging psi: %d, loss name: %s, use_lbfgs: %a",
-            self.sampling_count,
-            self.learning_rate,
-            self.local_step,
-            self.local_loss,
-            self.logging_psi,
-            self.loss_name,
-            self.use_lbfgs,
-        )
-
-        logging.info("main looping")
-        while True:
-            logging.info("sampling configurations")
-            configs, pre_amplitudes, _, _ = network.generate_unique(self.sampling_count)
-            logging.info("sampling done")
-            unique_sampling_count = len(configs)
-            logging.info("unique sampling count is %d", unique_sampling_count)
-
-            logging.info("generating hamiltonian data to create sparse matrix")
-            indices_i_and_j, values = model.inside(configs.cpu())
-            logging.info("sparse matrix data created")
-            logging.info("converting sparse matrix data to sparse matrix")
-            hamiltonian = scipy.sparse.coo_matrix((values, indices_i_and_j.T), [unique_sampling_count, unique_sampling_count], dtype=numpy.complex128).tocsr()
-            logging.info("sparse matrix created")
-            logging.info("estimating ground state")
-            target_energy, targets = scipy.sparse.linalg.lobpcg(hamiltonian, pre_amplitudes.cpu().reshape([-1, 1]).detach().numpy(), largest=False, maxiter=1024)
-            logging.info("estimiated, target energy is %.10f, ref energy is %.10f", target_energy.item(), model.ref_energy)
-            logging.info("preparing learning targets")
-            targets = torch.tensor(targets).view([-1]).cuda()
-            max_index = targets.abs().argmax()
-            targets = targets / targets[max_index]
-
-            logging.info("choosing loss function as %s", self.loss_name)
-            loss_func = getattr(losses, self.loss_name)
-
-            if self.use_lbfgs:
-                optimizer = torch.optim.LBFGS(network.parameters(), lr=self.learning_rate)
-            else:
-                optimizer = torch.optim.Adam(network.parameters(), lr=self.learning_rate)
-
-            def closure():
-                optimizer.zero_grad()
-                amplitudes = network(configs)
-                amplitudes = amplitudes / amplitudes[max_index]
-                loss = loss_func(amplitudes, targets)
-                loss.backward()
-                loss.amplitudes = amplitudes
-                return loss
-
-            logging.info("local optimization starting")
-            for i in range(self.local_step):
-                loss = optimizer.step(closure)
-                logging.info("local optimizing, step %d, loss %.10f", i, loss.item())
-                if loss < self.local_loss:
-                    logging.info("local optimization stop since local loss reached")
-                    break
-
-            logging.info("local optimization finished")
-            logging.info("saving checkpoint")
-            torch.save(network.state_dict(), f"{self.common.checkpoint_path}/{self.common.job_name}.pt")
-            logging.info("checkpoint saved")
-            logging.info("calculating current energy")
-            torch.enable_grad(closure)()
-            amplitudes = loss.amplitudes.cpu().detach().numpy()
-            final_energy = ((amplitudes.conj() @ (hamiltonian @ amplitudes)) / (amplitudes.conj() @ amplitudes)).real
-            logging.info(
-                "loss = %.10f during local optimization, final energy %.10f, target energy %.10f, ref energy %.10f",
-                loss.item(),
-                final_energy.item(),
-                target_energy.item(),
-                model.ref_energy,
-            )
-            logging.info("printing several largest amplitudes")
-            indices = targets.abs().sort(descending=True).indices
-            for index in indices[:self.logging_psi]:
-                logging.info("config %s, target %s, final %s", "".join(map(str, configs[index].cpu().numpy())), f"{targets[index].item():.8f}", f"{amplitudes[index].item():.8f}")
-
-
-subcommand_dict["learn"] = LearnConfig

qmb/losses.py

-import torch
-
-
-def log(a, b):
-    # b is target, a is learnable
-    error = (a.log() - b.log()) / (2 * torch.pi)
-    error_real = error.real
-    error_imag = error.imag - error.imag.round()
-    loss = error_real**2 + error_imag**2
-    return loss.mean()
-
-
-def target_reweighted_log(a, b):
-    # b is target, a is learnable
-    error = (a.log() - b.log()) / (2 * torch.pi)
-    error_real = error.real
-    error_imag = error.imag - error.imag.round()
-    loss = error_real**2 + error_imag**2
-    abs_b = b.abs()
-    loss = loss * abs_b
-    return loss.mean()
-
-
-def target_filtered_log(a, b):
-    # b is target, a is learnable
-    error = (a.log() - b.log()) / (2 * torch.pi)
-    error_real = error.real
-    error_imag = error.imag - error.imag.round()
-    loss = error_real**2 + error_imag**2
-    abs_b = b.abs()
-    loss = loss / (1 + 1e-10 / abs_b)
-    # This function scale only for very small abs value.
-    # I think we could ignore those definitly for amplitude less than 1e-10.
-    return loss.mean()
-
-
-def sum_filtered_log(a, b):
-    # b is target, a is learnable
-    error = (a.log() - b.log()) / (2 * torch.pi)
-    error_real = error.real
-    error_imag = error.imag - error.imag.round()
-    loss = error_real**2 + error_imag**2
-    sum_a_b = a.abs() + b.abs()
-    loss = loss / (1 + 1e-10 / sum_a_b)
-    # This function scale only for very small abs value.
-    # I think we could ignore those definitly for amplitude less than 1e-10.
-    return loss.mean()
-
-
-def direct(a, b):
-    # b is target, a is learnable
-    error = a - b
-    loss = (error.conj() * error).real
-    return loss.mean()

qmb/Makefile

-# This Makefile is used for temporary compilation during development.
-# For release, the configuration in setup.py should be used for compilation.
-
-CXX := g++
-CXXFLAGS := -O3 -Wall -shared -std=c++17 -fPIC
-PYTHON_INCLUDE := $(shell python3 -m pybind11 --includes)
-PYTHON_SUFFIX := $(shell python3-config --extension-suffix)
-SOURCES := $(wildcard _*.cpp)
-MODULES := $(basename $(SOURCES))
-TARGETS := $(addsuffix $(PYTHON_SUFFIX), $(MODULES))
-
-all: $(TARGETS)
-
-$(TARGETS): %$(PYTHON_SUFFIX): %.cpp
-	$(CXX) $(CXXFLAGS) $(PYTHON_INCLUDE) $< -o $@
-
-clean:
-	rm -f $(TARGETS)
-
-.PHONY: all clean

qmb/naqs.py

-# This file implements naqs from https://arxiv.org/pdf/2109.12606
-
-import torch
-
-
-class FakeLinear(torch.nn.Module):
-    """
-    Fake linear layer where dim_in = 0, avoid pytorch warning in initialization
-    """
-
-    def __init__(self, dim_in: int, dim_out: int) -> None:
-        super().__init__()
-        assert dim_in == 0
-        self.bias: torch.nn.Parameter = torch.nn.Parameter(torch.zeros([dim_out]))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        batch, zero = x.shape
-        return self.bias.view([1, -1]).expand([batch, -1])
-
-
-def Linear(dim_in: int, dim_out: int) -> torch.nn.Module:
-    # avoid torch warning when initialize a linear layer with dim_in = 0
-    if dim_in == 0:
-        return FakeLinear(dim_in, dim_out)
-    else:
-        return torch.nn.Linear(dim_in, dim_out)
-
-
-class MLP(torch.nn.Module):
-    """
-    This module implements multiple layers MLP with given dim_input, dim_output and hidden_size
-    """
-
-    def __init__(self, dim_input: int, dim_output: int, hidden_size: tuple[int, ...]) -> None:
-        super().__init__()
-        self.dim_input: int = dim_input
-        self.dim_output: int = dim_output
-        self.hidden_size: tuple[int, ...] = hidden_size
-        self.depth: int = len(hidden_size)
-
-        dimensions: list[int] = [dim_input] + list(hidden_size) + [dim_output]
-        linears: list[torch.nn.Module] = [Linear(i, j) for i, j in zip(dimensions[:-1], dimensions[1:])]
-        modules: list[torch.nn.Module] = [module for linear in linears for module in (linear, torch.nn.SiLU())][:-1]
-        self.model: torch.nn.Module = torch.nn.Sequential(*modules)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.model(x)
-
-
-class WaveFunction(torch.nn.Module):
-    """
-    This module implements naqs from https://arxiv.org/pdf/2109.12606
-    """
-
-    def __init__(
-            self,
-            *,
-            double_sites: int,  # qubits number, qubits are grouped by two for each site in naqs, so name it as double sites
-            physical_dim: int,  # is always 2 for naqs
-            is_complex: bool,  # is always true for naqs
-            spin_up: int,  # spin up number
-            spin_down: int,  # spin down number
-            hidden_size: tuple[int, ...],  # hidden size for MLP
-            ordering: int | list[int],  # ordering of sites +1 for normal order, -1 for reversed order, or the order list directly
-    ) -> None:
-        super().__init__()
-        assert double_sites % 2 == 0
-        self.double_sites: int = double_sites
-        self.sites: int = double_sites // 2
-        assert physical_dim == 2
-        assert is_complex == True
-        self.spin_up: int = spin_up
-        self.spin_down: int = spin_down
-        self.hidden_size: tuple[int, ...] = hidden_size
-
-        # The amplitude and phase network for each site
-        # each of them accept qubits before them and output vector with dimension of 4 as the configuration of two qubits on the current site.
-        self.amplitude: torch.nn.ModuleList = torch.nn.ModuleList([MLP(i * 2, 4, self.hidden_size) for i in range(self.sites)])
-        self.phase: torch.nn.ModuleList = torch.nn.ModuleList([MLP(i * 2, 4, self.hidden_size) for i in range(self.sites)])
-
-        # ordering of sites +1 for normal order, -1 for reversed order
-        if isinstance(ordering, int) and ordering == +1:
-            ordering = list(range(self.sites))
-        if isinstance(ordering, int) and ordering == -1:
-            ordering = list(reversed(range(self.sites)))
-        self.register_buffer('ordering', torch.tensor(ordering, dtype=torch.int64))
-        self.register_buffer('ordering_reversed', torch.scatter(torch.zeros(self.sites, dtype=torch.int64), 0, self.ordering, torch.arange(self.sites, dtype=torch.int64)))
-
-        # used to get device and dtype
-        self.dummy_param = torch.nn.Parameter(torch.empty(0))
-
-    @torch.jit.export
-    def mask(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Determined whether we could append spin up or spin down after uncompleted configurations.
-        """
-        # x : batch_size * current_site * 2
-        # x are the uncompleted configurations
-        batch_size = x.shape[0]
-        current_site = x.shape[1]
-        # number : batch_size * 2
-        # number is the total electron number of uncompleted configurations
-        number = x.sum(dim=1)
-
-        # up/down_electron/hole : batch_size
-        # the electron and hole number of uncompleted configurations
-        up_electron = number[:, 0]
-        down_electron = number[:, 1]
-        up_hole = current_site - up_electron
-        down_hole = current_site - down_electron
-
-        # add_up/down_electron/hole : batch_size
-        # whether able to append up/down electron/hole
-        add_up_electron = up_electron < self.spin_up
-        add_down_electron = down_electron < self.spin_down
-        add_up_hole = up_hole < self.sites - self.spin_up
-        add_down_hole = down_hole < self.sites - self.spin_down
-
-        # add_up : batch_size * 2 * 1
-        # add_down : batch_size * 1 * 2
-        add_up = torch.stack([add_up_hole, add_up_electron], dim=-1).unsqueeze(-1)
-        add_down = torch.stack([add_down_hole, add_down_electron], dim=-1).unsqueeze(-2)
-        # add : batch_size * 2 * 2
-        add = torch.logical_and(add_up, add_down)
-        # add contains whether to append up/down electron/hole after uncompleted configurations
-        # add[_, 0, 0] means we could add up hole and down hole
-        # add[_, 0, 1] means we could add up hole and down electron
-        # add[_, 1, 0] means we could add up electron and down hole
-        # add[_, 1, 1] means we could add up electron and down electron
-        return add
-
-    @torch.jit.export
-    def normalize_amplitude(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Normalize uncompleted log amplitude.
-        """
-        # x : batch_size * 2 * 2
-        # param :  batch_size
-        param = (2 * x).exp().sum(dim=[-2, -1]).log() / 2
-        x = x - param.unsqueeze(-1).unsqueeze(-1)
-        # 1 = param = sqrt(sum(x.exp()^2)) now
-        return x
-
-    @torch.jit.export
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Calculate psi of given configurations.
-        """
-        device: torch.device = self.dummy_param.device
-        dtype: torch.dtype = self.dummy_param.dtype
-
-        batch_size: int = x.shape[0]
-        # x : batch_size * sites * 2
-        x = x.reshape([batch_size, self.sites, 2])
-        # apply ordering
-        x = torch.index_select(x, 1, self.ordering_reversed)
-
-        x_float: torch.Tensor = x.to(dtype=dtype)
-        arange: torch.Tensor = torch.arange(batch_size, device=device)
-        total_amplitude: torch.Tensor = torch.zeros([batch_size], device=device, dtype=dtype)
-        total_phase: torch.Tensor = torch.zeros([batch_size], device=device, dtype=dtype)
-        for i, amplitude_phase_m in enumerate(zip(self.amplitude, self.phase)):
-            amplitude_m, phase_m = amplitude_phase_m
-            # delta_amplitude/phase : batch * 2 * 2
-            # delta amplitude and phase for the configurations at new site.
-            delta_amplitude: torch.Tensor = amplitude_m(x_float[:, :i].reshape([batch_size, 2 * i])).reshape([batch_size, 2, 2])
-            delta_phase: torch.Tensor = phase_m(x_float[:, :i].reshape([batch_size, 2 * i])).reshape([batch_size, 2, 2])
-            # filter mask for amplitude
-            delta_amplitude: torch.Tensor = delta_amplitude + torch.where(self.mask(x[:, :i]), 0, -torch.inf)
-            # normalize amplitude
-            delta_amplitude: torch.Tensor = self.normalize_amplitude(delta_amplitude)
-            # delta_amplitude/phase : batch
-            delta_amplitude: torch.Tensor = delta_amplitude[arange, x[:, i, 0], x[:, i, 1]]
-            delta_phase: torch.Tensor = delta_phase[arange, x[:, i, 0], x[:, i, 1]]
-            # calculate total amplitude and phase
-            total_amplitude = total_amplitude + delta_amplitude
-            total_phase = total_phase + delta_phase
-        return torch.view_as_complex(torch.stack([total_amplitude, total_phase], dim=-1)).exp()
-
-    @torch.jit.export
-    def binomial(self, count: torch.Tensor, probability: torch.Tensor) -> torch.Tensor:
-        """
-        Binomial sampling with given count and probability
-        """
-        # clamp probability
-        probability = torch.clamp(probability, min=0, max=1)
-        # set probability to zero for count = 0 since it may be nan when count = 0
-        probability = torch.where(count == 0, 0, probability)
-        # create dist and sample
-        result = torch.binomial(count, probability).to(dtype=torch.int64)
-        # numerical error since result is cast from float.
-        return torch.clamp(result, min=torch.zeros_like(count), max=count)
-
-    @torch.jit.export
-    def generate_unique(self, batch_size: int) -> tuple[torch.Tensor, torch.Tensor, None, None]:
-        """
-        Generate configurations uniquely.
-        see https://arxiv.org/pdf/2408.07625
-        """
-        device: torch.device = self.dummy_param.device
-        dtype: torch.dtype = self.dummy_param.dtype
-
-        # x : local_batch_size * current_site * 2
-        x: torch.Tensor = torch.empty([1, 0, 2], device=device, dtype=torch.int64)
-        # (un)perturbed_log_probability : local_batch_size
-        unperturbed_probability: torch.Tensor = torch.tensor([0], dtype=dtype, device=device)
-        perturbed_probability: torch.Tensor = torch.tensor([0], dtype=dtype, device=device)
-        for i, amplitude_m in enumerate(self.amplitude):
-            local_batch_size: int = x.shape[0]
-            x_float: torch.Tensor = x.to(dtype=dtype)
-            # delta_amplitude : batch * 2 * 2
-            # delta amplitude for the configurations at new site.
-            delta_amplitude: torch.Tensor = amplitude_m(x_float.reshape([local_batch_size, 2 * i])).reshape([local_batch_size, 2, 2])
-            # filter mask for amplitude
-            delta_amplitude: torch.Tensor = delta_amplitude + torch.where(self.mask(x), 0, -torch.inf)
-            # normalize amplitude
-            delta_amplitude: torch.Tensor = self.normalize_amplitude(delta_amplitude)
-
-            # delta unperturbed prob for all batch and 4 adds
-            l: torch.Tensor = (2 * delta_amplitude).reshape([local_batch_size, 4])
-            # and add to get the current unperturbed prob
-            l: torch.Tensor = unperturbed_probability.view([-1, 1]) + l
-            # get perturbed prob
-            L: torch.Tensor = l - (-torch.rand_like(l).log()).log()
-            # get max perturbed prob
-            Z: torch.Tensor = L.max(dim=-1).values.reshape([-1, 1])
-            # evaluate the conditioned prob
-            L: torch.Tensor = -torch.log(torch.exp(-perturbed_probability.view([-1, 1])) - torch.exp(-Z) + torch.exp(-L))
-
-            # calculate appended configurations for 4 adds
-            # local_batch_size * current_site * 2 + local_batch_size * 1 * 2
-            x0: torch.Tensor = torch.cat([x, torch.tensor([[0, 0]], device=device).expand(local_batch_size, -1, -1)], dim=1)
-            x1: torch.Tensor = torch.cat([x, torch.tensor([[0, 1]], device=device).expand(local_batch_size, -1, -1)], dim=1)
-            x2: torch.Tensor = torch.cat([x, torch.tensor([[1, 0]], device=device).expand(local_batch_size, -1, -1)], dim=1)
-            x3: torch.Tensor = torch.cat([x, torch.tensor([[1, 1]], device=device).expand(local_batch_size, -1, -1)], dim=1)
-
-            # cat all configurations to get x : new_local_batch_size * current_size * 2
-            # (un)perturbed prob : new_local_batch_size
-            x = torch.cat([x0, x1, x2, x3])
-            unperturbed_probability = l.permute(1, 0).reshape([-1])
-            perturbed_probability = L.permute(1, 0).reshape([-1])
-
-            # filter new data, only used largest batch_size ones
-            selected = perturbed_probability.sort(descending=True).indices[:batch_size]
-            x = x[selected]
-            unperturbed_probability = unperturbed_probability[selected]
-            perturbed_probability = perturbed_probability[selected]
-
-            # if prob = 0, filter it forcely
-            selected = perturbed_probability != -torch.inf
-            x = x[selected]
-            unperturbed_probability = unperturbed_probability[selected]
-            perturbed_probability = perturbed_probability[selected]
-
-        # apply ordering
-        x = torch.index_select(x, 1, self.ordering)
-        # flatten site part of x
-        x = x.reshape([x.size(0), self.double_sites])
-        # it should return configurations, amplitudes, probabilities and multiplicities
-        # but it is unique generator, so last two field is none
-        return x, self(x), None, None
-
-
-class WaveFunctionNormal(torch.nn.Module):
-    """
-    This module implements naqs from https://arxiv.org/pdf/2109.12606
-    No subspace restrictor.
-    """
-
-    def __init__(
-            self,
-            *,
-            sites: int,  # qubits number
-            physical_dim: int,  # is always 2 for naqs
-            is_complex: bool,  # is always true for naqs
-            hidden_size: tuple[int, ...],  # hidden size for MLP
-            ordering: int | list[int],  # ordering of sites +1 for normal order, -1 for reversed order, or the order list directly
-    ) -> None:
-        super().__init__()
-        self.sites: int = sites
-        assert physical_dim == 2
-        assert is_complex == True
-        self.hidden_size: tuple[int, ...] = hidden_size
-
-        # The amplitude and phase network for each site
-        # each of them accept qubits before them and output vector with dimension of 2 as the configuration of the qubit on the current site.
-        self.amplitude: torch.nn.ModuleList = torch.nn.ModuleList([MLP(i, 2, self.hidden_size) for i in range(self.sites)])
-        self.phase: torch.nn.ModuleList = torch.nn.ModuleList([MLP(i, 2, self.hidden_size) for i in range(self.sites)])
-
-        # ordering of sites +1 for normal order, -1 for reversed order
-        if isinstance(ordering, int) and ordering == +1:
-            ordering = list(range(self.sites))
-        if isinstance(ordering, int) and ordering == -1:
-            ordering = list(reversed(range(self.sites)))
-        self.register_buffer('ordering', torch.tensor(ordering, dtype=torch.int64))
-        self.register_buffer('ordering_reversed', torch.scatter(torch.zeros(self.sites, dtype=torch.int64), 0, self.ordering, torch.arange(self.sites, dtype=torch.int64)))
-
-        # used to get device and dtype
-        self.dummy_param = torch.nn.Parameter(torch.empty(0))
-
-    @torch.jit.export
-    def normalize_amplitude(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Normalize uncompleted log amplitude.
-        """
-        # x : batch_size * 2
-        # param :  batch_size
-        param = (2 * x).exp().sum(dim=[-1]).log() / 2
-        x = x - param.unsqueeze(-1)
-        # 1 = param = sqrt(sum(x.exp()^2)) now
-        return x
-
-    @torch.jit.export
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Calculate psi of given configurations.
-        """
-        device: torch.device = self.dummy_param.device
-        dtype: torch.dtype = self.dummy_param.dtype
-
-        batch_size: int = x.shape[0]
-        # x : batch_size * sites
-        x = x.reshape([batch_size, self.sites])
-        # apply ordering
-        x = torch.index_select(x, 1, self.ordering_reversed)
-
-        x_float: torch.Tensor = x.to(dtype=dtype)
-        arange: torch.Tensor = torch.arange(batch_size, device=device)
-        total_amplitude: torch.Tensor = torch.zeros([batch_size], device=device, dtype=dtype)
-        total_phase: torch.Tensor = torch.zeros([batch_size], device=device, dtype=dtype)
-        for i, amplitude_phase_m in enumerate(zip(self.amplitude, self.phase)):
-            amplitude_m, phase_m = amplitude_phase_m
-            # delta_amplitude/phase : batch * 2
-            # delta amplitude and phase for the configurations at new site.
-            delta_amplitude: torch.Tensor = amplitude_m(x_float[:, :i].reshape([batch_size, i])).reshape([batch_size, 2])
-            delta_phase: torch.Tensor = phase_m(x_float[:, :i].reshape([batch_size, i])).reshape([batch_size, 2])
-            # normalize amplitude
-            delta_amplitude: torch.Tensor = self.normalize_amplitude(delta_amplitude)
-            # delta_amplitude/phase : batch
-            delta_amplitude: torch.Tensor = delta_amplitude[arange, x[:, i]]
-            delta_phase: torch.Tensor = delta_phase[arange, x[:, i]]
-            # calculate total amplitude and phase
-            total_amplitude = total_amplitude + delta_amplitude
-            total_phase = total_phase + delta_phase
-        return torch.view_as_complex(torch.stack([total_amplitude, total_phase], dim=-1)).exp()
-
-    @torch.jit.export
-    def binomial(self, count: torch.Tensor, probability: torch.Tensor) -> torch.Tensor:
-        """
-        Binomial sampling with given count and probability
-        """
-        # clamp probability
-        probability = torch.clamp(probability, min=0, max=1)
-        # set probability to zero for count = 0 since it may be nan when count = 0
-        probability = torch.where(count == 0, 0, probability)
-        # create dist and sample
-        result = torch.binomial(count, probability).to(dtype=torch.int64)
-        # numerical error since result is cast from float.
-        return torch.clamp(result, min=torch.zeros_like(count), max=count)
-
-    @torch.jit.export
-    def generate_unique(self, batch_size: int) -> tuple[torch.Tensor, torch.Tensor, None, None]:
-        """
-        Generate configurations uniquely.
-        see https://arxiv.org/pdf/2408.07625
-        """
-        device: torch.device = self.dummy_param.device
-        dtype: torch.dtype = self.dummy_param.dtype
-
-        # x : local_batch_size * current_site
-        x: torch.Tensor = torch.empty([1, 0], device=device, dtype=torch.int64)
-        # (un)perturbed_log_probability : local_batch_size
-        unperturbed_probability: torch.Tensor = torch.tensor([0], dtype=dtype, device=device)
-        perturbed_probability: torch.Tensor = torch.tensor([0], dtype=dtype, device=device)
-        for i, amplitude_m in enumerate(self.amplitude):
-            local_batch_size: int = x.shape[0]
-            x_float: torch.Tensor = x.to(dtype=dtype)
-            # delta_amplitude : batch * 2
-            # delta amplitude for the configurations at new site.
-            delta_amplitude: torch.Tensor = amplitude_m(x_float.reshape([local_batch_size, i])).reshape([local_batch_size, 2])
-            # normalize amplitude
-            delta_amplitude: torch.Tensor = self.normalize_amplitude(delta_amplitude)
-
-            # delta unperturbed prob for all batch and 2 adds
-            l: torch.Tensor = (2 * delta_amplitude).reshape([local_batch_size, 2])
-            # and add to get the current unperturbed prob
-            l: torch.Tensor = unperturbed_probability.view([-1, 1]) + l
-            # get perturbed prob
-            L: torch.Tensor = l - (-torch.rand_like(l).log()).log()
-            # get max perturbed prob
-            Z: torch.Tensor = L.max(dim=-1).values.reshape([-1, 1])
-            # evaluate the conditioned prob
-            L: torch.Tensor = -torch.log(torch.exp(-perturbed_probability.view([-1, 1])) - torch.exp(-Z) + torch.exp(-L))
-
-            # calculate appended configurations for 2 adds
-            # local_batch_size * current_site + local_batch_size * 1
-            x0: torch.Tensor = torch.cat([x, torch.tensor([0], device=device).expand(local_batch_size, -1)], dim=1)
-            x1: torch.Tensor = torch.cat([x, torch.tensor([1], device=device).expand(local_batch_size, -1)], dim=1)
-
-            # cat all configurations to get x : new_local_batch_size * current_size * 2
-            # (un)perturbed prob : new_local_batch_size
-            x = torch.cat([x0, x1])
-            unperturbed_probability = l.permute(1, 0).reshape([-1])
-            perturbed_probability = L.permute(1, 0).reshape([-1])
-
-            # filter new data, only used largest batch_size ones
-            selected = perturbed_probability.sort(descending=True).indices[:batch_size]
-            x = x[selected]
-            unperturbed_probability = unperturbed_probability[selected]
-            perturbed_probability = perturbed_probability[selected]
-
-            # if prob = 0, filter it forcely
-            selected = perturbed_probability != -torch.inf
-            x = x[selected]
-            unperturbed_probability = unperturbed_probability[selected]
-            perturbed_probability = perturbed_probability[selected]
-
-        # apply ordering
-        x = torch.index_select(x, 1, self.ordering)
-        # flatten site part of x
-        x = x.reshape([x.size(0), self.sites])
-        # it should return configurations, amplitudes, probabilities and multiplicities
-        # but it is unique generator, so last two field is none
-        return x, self(x), None, None

qmb/openfermion_operator.py

-# This file implements interface to openfermion model, but read the fermion operators directly.
-
-import re
-import logging
-import numpy
-import torch
-from .openfermion import Model as OpenFermionModel
-from . import _openfermion
-
-
-class Model(OpenFermionModel):
-
-    def __init__(self, model_name, model_path):
-        self.model_name = model_name
-        self.model_path = model_path
-        self.model_file_name = f"{self.model_path}/{self.model_name}.npy"
-        logging.info("loading operator of openfermion model %s from %s", self.model_name, self.model_file_name)
-        self.openfermion = numpy.load(self.model_file_name, allow_pickle=True).item()
-        logging.info("operator of openfermion model %s loaded", self.model_name)
-
-        n_electrons, n_qubits = re.match(r"\w*_(\d*)_(\d*)", model_name).groups()
-        self.n_qubits = int(n_qubits)
-        self.n_electrons = int(n_electrons)
-        logging.info("n_qubits: %d, n_electrons: %d", self.n_qubits, self.n_electrons)
-
-        self.ref_energy = torch.nan
-        logging.info("reference energy is unknown")
-
-        logging.info("converting openfermion handle to hamiltonian handle")
-        self.hamiltonian = _openfermion.Hamiltonian(list(self.openfermion.terms.items()))
-        logging.info("hamiltonian handle has been created")

qmb/openfermion.py

-# This file implements interface to openfermion model.
-
-import typing
-import logging
-import pathlib
-import dataclasses
-import torch
-import tyro
-import openfermion
-from . import _openfermion
-from . import naqs as naqs_m
-from . import attention as attention_m
-
-
[e-mail]
-class ModelConfig:
-    # The openfermion model name
-    model_name: typing.Annotated[str, tyro.conf.Positional, tyro.conf.arg(metavar="MODEL")]
-    # The path of models folder
-    model_path: typing.Annotated[pathlib.Path, tyro.conf.arg(aliases=["-M"])] = pathlib.Path("models")
-
-
[e-mail]
-class NaqsConfig:
-    # The hidden widths of the network
-    hidden: typing.Annotated[tuple[int, ...], tyro.conf.arg(aliases=["-w"])] = (512,)
-
-
[e-mail]
-class AttentionConfig:
-    # Embedding dimension
-    embedding_dim: typing.Annotated[int, tyro.conf.arg(aliases=["-e"])] = 512
-    # Heads number
-    heads_num: typing.Annotated[int, tyro.conf.arg(aliases=["-m"])] = 8
-    # Feedforward dimension
-    feed_forward_dim: typing.Annotated[int, tyro.conf.arg(aliases=["-f"])] = 2048
-    # Network depth
-    depth: typing.Annotated[int, tyro.conf.arg(aliases=["-d"])] = 6
-
-
-class Model:
-
-    @classmethod
-    def preparse(cls, input_args):
-        args = tyro.cli(ModelConfig, args=input_args)
-        return args.model_name
-
-    @classmethod
-    def parse(cls, input_args):
-        logging.info("parsing args %a by openfermion model", input_args)
-        args = tyro.cli(ModelConfig, args=input_args)
-        logging.info("model name: %s, model path: %s", args.model_name, args.model_path)
-
-        return cls(args.model_name, args.model_path)
-
-    def __init__(self, model_name, model_path):
-        self.model_name = model_name
-        self.model_path = model_path
-        self.model_file_name = f"{self.model_path}/{self.model_name}.hdf5"
-        logging.info("loading openfermion model %s from %s", self.model_name, self.model_file_name)
-        self.openfermion = openfermion.MolecularData(filename=self.model_file_name)
-        logging.info("openfermion model %s loaded", self.model_name)
-
-        self.n_qubits = self.openfermion.n_qubits
-        self.n_electrons = self.openfermion.n_electrons
-        logging.info("n_qubits: %d, n_electrons: %d", self.n_qubits, self.n_electrons)
-
-        self.ref_energy = self.openfermion.fci_energy.item()
-        logging.info("reference energy in openfermion data is %.10f", self.ref_energy)
-
-        logging.info("converting openfermion handle to hamiltonian handle")
-        self.hamiltonian = _openfermion.Hamiltonian(list(openfermion.transforms.get_fermion_operator(self.openfermion.get_molecular_hamiltonian()).terms.items()))
-        logging.info("hamiltonian handle has been created")
-
-    def inside(self, configs):
-        return self.hamiltonian.inside(configs)
-
-    def outside(self, configs):
-        return self.hamiltonian.outside(configs)
-
-    def naqs(self, input_args):
-        logging.info("parsing args %a by network naqs", input_args)
-        args = tyro.cli(NaqsConfig, args=input_args)
-        logging.info("hidden: %a", args.hidden)
-
-        network = naqs_m.WaveFunction(
-            double_sites=self.n_qubits,
-            physical_dim=2,
-            is_complex=True,
-            spin_up=self.n_electrons // 2,
-            spin_down=self.n_electrons // 2,
-            hidden_size=args.hidden,
-            ordering=+1,
-        ).double()
-
-        return torch.jit.script(network)
-
-    def attention(self, input_args):
-        logging.info("parsing args %a by network attention", input_args)
-        args = tyro.cli(AttentionConfig, args=input_args)
-        logging.info("embedding dim: %d, heads_num: %d, feed forward dim: %d, depth: %d", args.embedding_dim, args.heads_num, args.feed_forward_dim, args.depth)
-
-        network = attention_m.WaveFunction(
-            double_sites=self.n_qubits,
-            physical_dim=2,
-            is_complex=True,
-            spin_up=self.n_electrons // 2,
-            spin_down=self.n_electrons // 2,
-            embedding_dim=args.embedding_dim,
-            heads_num=args.heads_num,
-            feed_forward_dim=args.feed_forward_dim,
-            depth=args.depth,
-            ordering=+1,
-        ).double()
-
-        return torch.jit.script(network)

qmb/subcommand_dict.py

-subcommand_dict = {}

qmb/vmc.py

-import logging
-import typing
-import dataclasses
-import torch
-import tyro
-from .common import CommonConfig
-from .subcommand_dict import subcommand_dict
-
-
[e-mail]
-class VmcConfig:
-    common: typing.Annotated[CommonConfig, tyro.conf.OmitArgPrefixes]
-
-    # sampling count
-    sampling_count: typing.Annotated[int, tyro.conf.arg(aliases=["-n"])] = 4000
-    # learning rate for the local optimizer
-    learning_rate: typing.Annotated[float, tyro.conf.arg(aliases=["-r"], help_behavior_hint="(default: 1e-3 for Adam, 1 for LBFGS)")] = -1
-    # step count for the local optimizer
-    local_step: typing.Annotated[int, tyro.conf.arg(aliases=["-s"])] = 1000
-    # calculate all psi(s)')
-    include_outside: typing.Annotated[bool, tyro.conf.arg(aliases=["-o"])] = False
-    # Use deviation instead of energy
-    deviation: typing.Annotated[bool, tyro.conf.arg(aliases=["-d"])] = False
-    # Fix outside phase when optimizing outside deviation
-    fix_outside: typing.Annotated[bool, tyro.conf.arg(aliases=["-f"])] = False
-    # Use LBFGS instead of Adam
-    use_lbfgs: typing.Annotated[bool, tyro.conf.arg(aliases=["-2"])] = False
-    # Do not calculate deviation when optimizing energy
-    omit_deviation: typing.Annotated[bool, tyro.conf.arg(aliases=["-i"])] = False
-
-    def __post_init__(self):
-        if self.learning_rate == -1:
-            self.learning_rate = 1 if self.use_lbfgs else 1e-3
-
-    def main(self):
-        model, network = self.common.main()
-
-        logging.info(
-            "sampling count: %d, learning rate: %f, local step: %d, include outside: %a, use deviation: %a, fix outside: %a, use lbfgs: %a, omit deviation: %a",
-            self.sampling_count,
-            self.learning_rate,
-            self.local_step,
-            self.include_outside,
-            self.deviation,
-            self.fix_outside,
-            self.use_lbfgs,
-            self.omit_deviation,
-        )
-
-        logging.info("main looping")
-        while True:
-            logging.info("sampling configurations")
-            configs_i, _, _, _ = network.generate_unique(self.sampling_count)
-            logging.info("sampling done")
-            unique_sampling_count = len(configs_i)
-            logging.info("unique sampling count is %d", unique_sampling_count)
-
-            if self.include_outside:
-                logging.info("generating hamiltonian data to create sparse matrix outsidely")
-                indices_i_and_j, values, configs_j = model.outside(configs_i.cpu())
-                logging.info("sparse matrix data created")
-                outside_count = len(configs_j)
-                logging.info("outside configs count is %d", outside_count)
-                logging.info("converting sparse matrix data to sparse matrix")
-                hamiltonian = torch.sparse_coo_tensor(indices_i_and_j.T, values, [unique_sampling_count, outside_count], dtype=torch.complex128).to_sparse_csr().cuda()
-                logging.info("sparse matrix created")
-                logging.info("moving configs j to cuda")
-                configs_j = torch.tensor(configs_j).cuda()
-                logging.info("configs j has been moved to cuda")
-            else:
-                logging.info("generating hamiltonian data to create sparse matrix insidely")
-                indices_i_and_j, values = model.inside(configs_i.cpu())
-                logging.info("sparse matrix data created")
-                logging.info("converting sparse matrix data to sparse matrix")
-                hamiltonian = torch.sparse_coo_tensor(indices_i_and_j.T, values, [unique_sampling_count, unique_sampling_count], dtype=torch.complex128).to_sparse_csr().cuda()
-                logging.info("sparse matrix created")
-
-            if self.use_lbfgs:
-                optimizer = torch.optim.LBFGS(network.parameters(), lr=self.learning_rate)
-            else:
-                optimizer = torch.optim.Adam(network.parameters(), lr=self.learning_rate)
-
-            if self.deviation:
-
-                def closure():
-                    # Optimizing deviation
-                    optimizer.zero_grad()
-                    # Calculate amplitudes i and amplitudes j
-                    # When including outside, amplitudes j should be calculated individually, otherwise, it equals to amplitudes i
-                    # It should be notices that sometimes we do not want to optimize small configurations
-                    # So we calculate amplitudes j in no grad mode
-                    # but the first several configurations in amplitudes j are duplicated with those in amplitudes i
-                    # So cat them manually
-                    amplitudes_i = network(configs_i)
-                    if self.include_outside:
-                        if self.fix_outside:
-                            with torch.no_grad():
-                                amplitudes_j = network(configs_j)
-                            amplitudes_j = torch.cat([amplitudes_i[:unique_sampling_count], amplitudes_j[unique_sampling_count:]])
-                        else:
-                            amplitudes_j = network(configs_j)
-                    else:
-                        amplitudes_j = amplitudes_i
-                    # <s|H|psi> will be used multiple times, calculate it first
-                    # as we want to optimize deviation, every value should be calculated in grad mode, so we do not detach anything
-                    hamiltonian_amplitudes_j = hamiltonian @ amplitudes_j
-                    # energy is just <psi|s> <s|H|psi> / <psi|s> <s|psi>
-                    energy = (amplitudes_i.conj() @ hamiltonian_amplitudes_j) / (amplitudes_i.conj() @ amplitudes_i)
-                    # we want to estimate variance of E_s - E with weight <psi|s><s|psi>
-                    # where E_s = <s|H|psi>/<s|psi>
-                    # the variance is (E_s - E).conj() @ (E_s - E) * <psi|s> <s|psi> / ... = (E_s <s|psi> - E <s|psi>).conj() @ (E_s <s|psi> - E <s|psi>) / ...
-                    # so we calculate E_s <s|psi> - E <s|psi> first, which is just <s|H|psi> - <s|psi> E, we name it as `difference'
-                    difference = hamiltonian_amplitudes_j - amplitudes_i * energy
-                    # the numerator calculated, the following is the variance
-                    variance = (difference.conj() @ difference) / (amplitudes_i.conj() @ amplitudes_i)
-                    # calculate the deviation
-                    deviation = variance.real.sqrt()
-                    deviation.backward()
-                    # As we have already calculated energy, embed it in deviation for logging
-                    deviation.energy = energy.real
-                    return deviation
-
-                logging.info("local optimization for deviation starting")
-                for i in range(self.local_step):
-                    deviation = optimizer.step(closure)
-                    logging.info("local optimizing, step: %d, energy: %.10f, deviation: %.10f", i, deviation.energy.item(), deviation.item())
-            else:
-
-                def closure():
-                    # Optimizing energy
-                    optimizer.zero_grad()
-                    # Calculate amplitudes i and amplitudes j
-                    # When including outside, amplitudes j should be calculated individually, otherwise, it equals to amplitudes i
-                    # Because of gradient formula, we always calculate amplitudes j in no grad mode
-                    amplitudes_i = network(configs_i)
-                    if self.include_outside:
-                        with torch.no_grad():
-                            amplitudes_j = network(configs_j)
-                    else:
-                        amplitudes_j = amplitudes_i.detach()
-                    # <s|H|psi> will be used multiple times, calculate it first
-                    # it should be notices that this <s|H|psi> is totally detached, since both hamiltonian and amplitudes j is detached
-                    hamiltonian_amplitudes_j = hamiltonian @ amplitudes_j
-                    # energy is just <psi|s> <s|H|psi> / <psi|s> <s|psi>
-                    # we only calculate gradient on <psi|s>, both <s|H|psi> and <s|psi> should be detached
-                    # since <s|H|psi> has been detached already, we detach <s|psi> here manually
-                    energy = (amplitudes_i.conj() @ hamiltonian_amplitudes_j) / (amplitudes_i.conj() @ amplitudes_i.detach())
-                    # Calculate deviation
-                    # The variance is (E_s <s|psi> - E <s|psi>).conj() @ (E_s <s|psi> - E <s|psi>) / <psi|s> <s|psi>
-                    # Calculate E_s <s|psi> - E <s|psi> first and name it as difference
-                    if self.omit_deviation:
-                        deviation = torch.tensor(torch.nan)
-                    else:
-                        with torch.no_grad():
-                            difference = hamiltonian_amplitudes_j - amplitudes_i * energy
-                            variance = (difference.conj() @ difference) / (amplitudes_i.conj() @ amplitudes_i)
-                            deviation = variance.real.sqrt()
-                    energy = energy.real
-                    energy.backward()
-                    # Embed the deviation which has been calculated in energy for logging
-                    energy.deviation = deviation
-                    return energy
-
-                logging.info("local optimization for energy starting")
-                for i in range(self.local_step):
-                    energy = optimizer.step(closure)
-                    logging.info("local optimizing, step: %d, energy: %.10f, deviation: %.10f", i, energy.item(), energy.deviation.item())
-
-            logging.info("local optimization finished")
-            logging.info("saving checkpoint")
-            torch.save(network.state_dict(), f"{self.common.checkpoint_path}/{self.common.job_name}.pt")
-            logging.info("checkpoint saved")
-
-
-subcommand_dict["vmc"] = VmcConfig

RECORD

-qmb/_openfermion.cpython-312-darwin.so,sha256=H6cAEiWgTnRabc5Y2WAdXlS4pOk_LChoIMfNvFU0B3Y,497576
-qmb/iter.py,sha256=N_94Mxd85xXp0_x2E_yyQdu6f4ObTk3RBnIA3foO4Io,4453
-qmb/attention.py,sha256=T4HB5Ls6FdrJ3zBcYcAyITvW6uNndgmVmFz859VpyZ8,18701
-qmb/ising.py,sha256=7WsM3XvWyFkMvI3eo_SybE_CKgS1iuiLdwmG8p6QCDA,2958
-qmb/_binary_tree.hpp,sha256=78Eus5lGQeJc45_dmWtevisBlSGqQDQZauhDmPcNyXs,1538
-qmb/version.py,sha256=TwRi2oMKFivr-wPPu9jL2dNTutt5H5w2riUVZ_WoZFE,108
-qmb/_version.py,sha256=VUKDhsDLOPKj9zNw_322lJTB3k9JKH6BOu4nkEc42KY,411
-qmb/Makefile,sha256=at602IizZ30TK0QbFR5SxnhsFYlihB7VOe4r8kK6J8A,588
-qmb/naqs.py,sha256=v3v4Ms6MtWPUZBgYb8SD7zgevGPRbmlLCHJBpaaPcwE,20977
-qmb/openfermion_operator.py,sha256=Vi9aP7XG09E3VRuaFK5mcS9P8g2RIZk75uycABLx-9g,1303
-qmb/__init__.py,sha256=3rjaLnjJirLdTgG2c2sy0R5rcs2Jos_mcO_G23uO9iI,80
-qmb/subcommand_dict.py,sha256=aG_seFnDYVxiVOAzBjtxsZUAocReDc28ycPSik89WIM,21
-qmb/_ising.pyi,sha256=SwZm3HFtAnoVxvu7meJo3tcDBPhh-ia-14jJdf8ssb0,458
-qmb/common.py,sha256=ysKZaCpvDhDowsCgJ2LSw5Zp5GUMUrHn8eEWQMjqK38,4143
-qmb/_openfermion.pyi,sha256=iNYzzeGlsi2pCVTuP9b3ebG1fzf8lW6sui8iQnuPCkg,453
-qmb/learn.py,sha256=FmK8_9TbBwhhG2AcW8tFjauy0SbsU9rSxvcCLNIluug,5642
-qmb/losses.py,sha256=DAU0MAMIa5c5-Mkd4BCmSWi7I7q9uqufQb7z5NVYus4,1600
-qmb/_ising.cpython-312-darwin.so,sha256=Qg51iDg-GSLFuSohlrbzokyJdnBoviVsiUAQBFTK8gc,461136
-qmb/openfermion.py,sha256=oyMloMv_5Pfkx8IKdw7K3UoE-o_2e8OTfT4ZVQBe4CU,4262
-qmb/vmc.py,sha256=JdPjlb-w0xZGJ-5hVsuvn6l8V9hYdyDqeEV0x29nMWg,9938
-qmb/__main__.py,sha256=KHKu7RBoi8EbQ6yMOBBuCqDm9_jmwRDzhLwmRuuLnQE,248
-qmb/_ising.cpp,sha256=y5NhEu5HAdP7WiC4-lqFViTF8J3av5w6mMXILHMKS1o,7526
-qmb/_openfermion.cpp,sha256=gwdizAqao2mVE9nJj0elbzeZQqKqgyuAf5J75SDXmrA,7449
-qmb-0.0.9.dist-info/LICENSE.md,sha256=zFRw_u1mGSOH8GrpOu0L1P765aX9fB5UpKz06mTxAos,34893
-qmb-0.0.9.dist-info/RECORD,,
-qmb-0.0.9.dist-info/WHEEL,sha256=JjfPN3tuNVktXeVx2CdYSsvuQYIjI_7f2R0S3yE0Gks,115
-qmb-0.0.9.dist-info/entry_points.txt,sha256=-j_dDu8UGn1_9iOs18ZRouGXcFoXo-5m11RgRtW_n7Q,42
-qmb-0.0.9.dist-info/top_level.txt,sha256=P0yWQTYqagwkL4IQ256SLpcuAnNYQglhOg1OzIBD7Rg,4
-qmb-0.0.9.dist-info/METADATA,sha256=mv4iLJ1KPn6E_PzzZbiGZuMuB2txSkxlFu4dLXTNhqQ,1324

top_level.txt

-qmb

WHEEL

-Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
-Root-Is-Purelib: false
-Tag: cp312-cp312-macosx_10_13_universal2
-

Alert delta unavailable

Currently unable to show alert delta for PyPI packages.

Improved metrics

Total package byte prevSize

47937

Infinity%

Number of package files

20

Infinity%

Lines of code

31

Infinity%