├── .clang-format ├── .gitignore ├── .gitmodules ├── CMakeLists.txt ├── LICENSE ├── MANIFEST.in ├── README.md ├── Singularity ├── benchmark ├── benchmark_apply.py ├── benchmark_jacobian.py └── benchmark_zanella.py ├── build_conda_package.sh ├── build_locally.sh ├── cartesius ├── build_pytorch.sh └── conda-cartesius.yml ├── cbits ├── common │ ├── CMakeLists.txt │ ├── accumulator.cpp │ ├── accumulator.hpp │ ├── bits512.cpp │ ├── bits512.hpp │ ├── common.hpp.backup │ ├── config.hpp │ ├── errors.cpp │ ├── errors.hpp │ ├── metropolis.cpp │ ├── metropolis.hpp │ ├── parallel.cpp │ ├── parallel.hpp │ ├── polynomial.cpp │ ├── polynomial.hpp │ ├── random.cpp │ ├── random.hpp │ ├── tensor_info.hpp │ ├── unpack.cpp │ ├── unpack.hpp │ ├── wrappers.cpp │ ├── wrappers.hpp │ ├── zanella.cpp │ └── zanella.hpp ├── cpu │ ├── CMakeLists.txt │ ├── kernels.cpp │ └── kernels.hpp ├── gpu │ ├── CMakeLists.txt │ ├── unpack.cu │ └── unpack.hpp └── python │ ├── CMakeLists.txt │ ├── bind_cuda.cpp │ ├── bind_cuda.hpp │ ├── bind_heisenberg.cpp │ ├── bind_heisenberg.hpp │ ├── bind_jacobian.cpp │ ├── bind_jacobian.hpp │ ├── bind_metropolis.cpp │ ├── bind_metropolis.hpp │ ├── bind_operator.cpp │ ├── bind_operator.hpp │ ├── bind_polynomial.cpp │ ├── bind_polynomial.hpp │ ├── bind_polynomial_state.cpp │ ├── bind_polynomial_state.hpp │ ├── bind_quantum_state.cpp │ ├── bind_quantum_state.hpp │ ├── bind_spin_basis.cpp │ ├── bind_spin_basis.hpp │ ├── bind_symmetry.cpp │ ├── bind_symmetry.hpp │ ├── bind_v2.cpp │ ├── bind_v2.hpp │ ├── init.cpp │ ├── pybind11_helpers.hpp │ ├── trim.cpp │ └── trim.hpp ├── conda-build.yml ├── conda-cpu.yml ├── conda-devel.yml ├── conda-gpu.yml ├── conda ├── cpu │ ├── build.sh │ ├── conda_build_config.yaml │ └── meta.yaml └── gpu │ ├── build.sh │ ├── conda_build_config.yaml │ └── meta.yaml ├── devel.yml ├── distributed_example.py ├── example ├── 1x10 │ ├── amplitude.py │ ├── amplitude_wip.py │ ├── gradient_descend.py │ ├── sign.py │ ├── sign_wip.py │ ├── stochastic_reconfiguration.py │ ├── supervised_amplitude.json │ └── supervised_sign.json ├── carleo2017 │ ├── from_scratch.py │ └── models.py ├── chain │ └── diagonalise.py ├── heisenberg_chain │ └── from_scratch.py ├── heisenberg_square │ ├── gradient_descend.py │ ├── heisenberg_square_36_positive.yaml │ └── pre_training.py ├── triangleperiodic │ └── 16 │ │ ├── amplitude.py │ │ └── sign.py └── wip │ └── square.py ├── nqs_playground ├── __init__.py ├── _extension.py ├── _jacobian.py ├── autoregressive.py ├── cbits │ ├── accumulator.cpp │ ├── accumulator.hpp │ ├── hedley.h │ ├── init.cpp │ ├── parallel.hpp │ ├── zanella.cpp │ └── zanella.hpp ├── core.py ├── distributed.py ├── hamiltonian.py ├── lbfgs.py ├── runner.py ├── sampling.py ├── sgd.py ├── sr.py ├── supervised.py └── swo.py ├── setup.py └── test ├── basis_5x5.pickle ├── difficult_to_sample_5x5.pt ├── test_acceptance.py ├── test_apply.py ├── test_basis.py ├── test_hamiltonian.py ├── test_jacobian.py ├── test_local_values.py ├── test_monte_carlo.py ├── test_polynomial.py ├── test_sampling.py ├── test_symmetry.py └── test_unpack.py /.clang-format: -------------------------------------------------------------------------------- 1 | --- 2 | Language: Cpp 3 | # BasedOnStyle: LLVM 4 | 5 | AccessModifierOffset: -2 6 | AlignAfterOpenBracket: Align 7 | AlignConsecutiveAssignments: true 8 | AlignConsecutiveDeclarations: true 9 | AlignEscapedNewlinesLeft: false 10 | AlignOperands: true 11 | AlignTrailingComments: true 12 | AllowAllParametersOfDeclarationOnNextLine: true 13 | AllowShortBlocksOnASingleLine: true 14 | AllowShortCaseLabelsOnASingleLine: true 15 | AllowShortFunctionsOnASingleLine: All 16 | AllowShortIfStatementsOnASingleLine: true 17 | AllowShortLoopsOnASingleLine: false 18 | AlwaysBreakAfterReturnType: None 19 | AlwaysBreakBeforeMultilineStrings: false 20 | AlwaysBreakTemplateDeclarations: MultiLine 21 | BinPackArguments: true 22 | BinPackParameters: true 23 | BraceWrapping: 24 | AfterClass: false 25 | AfterControlStatement: false 26 | AfterEnum: false 27 | AfterFunction: true 28 | AfterNamespace: false 29 | AfterStruct: false 30 | AfterUnion: false 31 | AfterExternBlock: false 32 | BeforeCatch: true 33 | BeforeElse: true 34 | IndentBraces: false 35 | SplitEmptyFunction: false 36 | SplitEmptyRecord: false 37 | SplitEmptyNamespace: true 38 | BreakBeforeBinaryOperators: NonAssignment 39 | BreakBeforeBraces: Custom 40 | BreakBeforeInheritanceComma: true 41 | BreakBeforeTernaryOperators: true 42 | BreakConstructorInitializers: BeforeComma 43 | BreakStringLiterals: true 44 | ColumnLimit: 100 45 | CommentPragmas: '^ IWYU pragma:' 46 | CompactNamespaces: false 47 | ConstructorInitializerAllOnOneLineOrOnePerLine: true 48 | ConstructorInitializerIndentWidth: 4 49 | ContinuationIndentWidth: 4 50 | Cpp11BracedListStyle: true 51 | DerivePointerAlignment: false 52 | DisableFormat: false 53 | ExperimentalAutoDetectBinPacking: false 54 | FixNamespaceComments: true 55 | ForEachMacros: [ foreach, Q_FOREACH, BOOST_FOREACH ] 56 | IncludeCategories: 57 | - Regex: '^"detail/' 58 | Priority: 1 59 | - Regex: '^<(boost|gsl|spdlog|tl)' 60 | Priority: 4 61 | - Regex: '^<.*\.h>$' 62 | Priority: 5 63 | - Regex: '^<' 64 | Priority: 6 65 | - Regex: '.*' 66 | Priority: 2 67 | IncludeIsMainRegex: '$' 68 | IndentCaseLabels: false 69 | IndentPPDirectives: AfterHash 70 | IndentWidth: 4 71 | IndentWrappedFunctionNames: false 72 | KeepEmptyLinesAtTheStartOfBlocks: true 73 | MacroBlockBegin: '' 74 | MacroBlockEnd: '' 75 | MaxEmptyLinesToKeep: 1 76 | NamespaceIndentation: Inner 77 | PenaltyBreakAssignment: 0 78 | PenaltyBreakBeforeFirstCallParameter: 20 79 | PenaltyBreakComment: 300 80 | PenaltyBreakFirstLessLess: 120 81 | PenaltyBreakString: 400 82 | PenaltyExcessCharacter: 12345 83 | PenaltyReturnTypeOnItsOwnLine: 60 84 | PointerAlignment: Left 85 | ReflowComments: false 86 | SortIncludes: true 87 | SortUsingDeclarations: true 88 | SpaceAfterCStyleCast: false 89 | SpaceAfterTemplateKeyword: true 90 | SpaceBeforeAssignmentOperators: true 91 | SpaceBeforeParens: ControlStatements 92 | SpaceInEmptyParentheses: false 93 | SpacesBeforeTrailingComments: 1 94 | SpacesInAngles: false 95 | SpacesInCStyleCastParentheses: false 96 | SpacesInParentheses: false 97 | SpacesInSquareBrackets: false 98 | Standard: Cpp11 99 | TabWidth: 4 100 | UseTab: Never 101 | ... 102 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # C++ stuff 10 | .clangd/ 11 | cbits/backup 12 | cbits/backup_v2 13 | compile_commands.json 14 | 15 | # Distribution / packaging 16 | .Python 17 | build/ 18 | develop-eggs/ 19 | dist/ 20 | downloads/ 21 | eggs/ 22 | .eggs/ 23 | lib/ 24 | lib64/ 25 | parts/ 26 | sdist/ 27 | var/ 28 | wheels/ 29 | *.egg-info/ 30 | .installed.cfg 31 | *.egg 32 | MANIFEST 33 | 34 | # PyInstaller 35 | # Usually these files are written by a python script from a template 36 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 37 | *.manifest 38 | *.spec 39 | 40 | # Installer logs 41 | pip-log.txt 42 | pip-delete-this-directory.txt 43 | 44 | # Unit test / coverage reports 45 | htmlcov/ 46 | .tox/ 47 | .coverage 48 | .coverage.* 49 | .cache 50 | nosetests.xml 51 | coverage.xml 52 | *.cover 53 | .hypothesis/ 54 | .pytest_cache/ 55 | 56 | # Translations 57 | *.mo 58 | *.pot 59 | 60 | # Django stuff: 61 | *.log 62 | local_settings.py 63 | db.sqlite3 64 | 65 | # Flask stuff: 66 | instance/ 67 | .webassets-cache 68 | 69 | # Scrapy stuff: 70 | .scrapy 71 | 72 | # Sphinx documentation 73 | docs/_build/ 74 | 75 | # PyBuilder 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # pyenv 82 | .python-version 83 | 84 | # celery beat schedule file 85 | celerybeat-schedule 86 | 87 | # SageMath parsed files 88 | *.sage.py 89 | 90 | # Environments 91 | .env 92 | .venv 93 | env/ 94 | venv/ 95 | ENV/ 96 | env.bak/ 97 | venv.bak/ 98 | 99 | # Spyder project settings 100 | .spyderproject 101 | .spyproject 102 | 103 | # Rope project settings 104 | .ropeproject 105 | 106 | # mkdocs documentation 107 | /site 108 | 109 | # mypy 110 | .mypy_cache/ 111 | 112 | # simulation output 113 | **/runs/ 114 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "external/pybind11"] 2 | path = external/pybind11 3 | url = https://github.com/pybind/pybind11.git 4 | [submodule "external/gsl-lite"] 5 | path = external/gsl-lite 6 | url = https://github.com/martinmoene/gsl-lite.git 7 | [submodule "external/fmt"] 8 | path = external/fmt 9 | url = https://github.com/fmtlib/fmt.git 10 | [submodule "external/SG14"] 11 | path = external/SG14 12 | url = https://github.com/WG21-SG14/SG14.git 13 | [submodule "external/vectorclass/version2"] 14 | path = external/vectorclass/version2 15 | url = https://github.com/vectorclass/version2 16 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2018, Tom Westerhout 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | * Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 30 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.rst 2 | include LICENSE 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | PyTorch-based implementation of SR and SWO for NQS. 2 | 3 | 4 | **Contents** 5 | - [Installation](#installation) 6 | - [Conda](#conda) 7 | - [Building from source](#building-from-source) 8 | 9 | 10 | ## Installation 11 | 12 | 13 | ## Conda 14 | 15 | > **WARNING:** The version available on Conda is currently out of date. 16 | > Please, build from source for the latest features. 17 | 18 | The simplest way to get started using `nqs_playground` package is to install it 19 | using [Conda](https://docs.conda.io/en/latest/): 20 | ```sh 21 | conda install -c twesterhout nqs_playground 22 | ``` 23 | 24 | 25 | ## Building from source 26 | 27 | ### CPU-only version 28 | 29 | If you do not have access or do not wish to use a GPU you can use cpu-only 30 | version PyToch and nqs_playground. For this, first clone the repository: 31 | 32 | ```sh 33 | git clone https://github.com/twesterhout/nqs-playground.git 34 | ``` 35 | 36 | Now just run [`build_locally_cpu.sh`](./build_locally_cpu.sh): 37 | 38 | ```sh 39 | ./build_locally_cpu.sh 40 | ``` 41 | 42 | ### Full version 43 | 44 | TODO 45 | -------------------------------------------------------------------------------- /Singularity: -------------------------------------------------------------------------------- 1 | Bootstrap: docker 2 | From: nvidia/cuda:11.0-devel-ubuntu20.04 3 | 4 | %post -c /bin/bash 5 | set -e 6 | export LC_ALL=C 7 | export PYTHON_VERSION=3.8 8 | export PYTORCH_VERSION=1.7 9 | export LATTICE_SYMMETRIES_VERSION=0.4.0 10 | 11 | apt-get update 12 | apt-get install -y --no-install-recommends \ 13 | ca-certificates \ 14 | pkg-config \ 15 | curl 16 | apt-get clean 17 | apt-get autoclean 18 | rm -rf /var/lib/apt/lists/* 19 | 20 | mkdir -p /workdir 21 | cd /workdir 22 | 23 | curl -o miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh 24 | chmod +x miniconda.sh 25 | ./miniconda.sh -b -p /opt/conda 26 | rm miniconda.sh 27 | 28 | . /opt/conda/etc/profile.d/conda.sh 29 | # We only match major and minor versions of cudatoolkit 30 | conda install -y python=$PYTHON_VERSION cudatoolkit=${CUDA_VERSION%.*} \ 31 | anaconda-client conda-build conda-verify \ 32 | gcc_linux-64 gxx_linux-64 cmake ninja 33 | conda install -y -c pytorch pytorch=$PYTORCH_VERSION 34 | conda install -y -c twesterhout lattice-symmetries=$LATTICE_SYMMETRIES_VERSION 35 | conda clean -ya 36 | 37 | 38 | %environment 39 | export LC_ALL=C 40 | export TERM=xterm-256color 41 | 42 | 43 | %runscript 44 | exec /bin/bash 45 | -------------------------------------------------------------------------------- /benchmark/benchmark_apply.py: -------------------------------------------------------------------------------- 1 | import lattice_symmetries as ls 2 | from loguru import logger 3 | import nqs_playground as nqs 4 | import time 5 | import torch 6 | from torch import Tensor 7 | 8 | 9 | def load_basis_and_hamiltonian( 10 | filename: str = "/vol/tcm28/westerhout_tom/papers/clever-sampling/exact_diagonalization/heisenberg_square_36.yaml", 11 | ): 12 | import yaml 13 | 14 | logger.info("Loading basis from '{}'...", filename) 15 | with open(filename, "r") as f: 16 | config = yaml.load(f, Loader=yaml.SafeLoader) 17 | basis = ls.SpinBasis.load_from_yaml(config["basis"]) 18 | logger.info("Loading Hamiltonian from '{}'...", filename) 19 | hamiltonian = ls.Operator.load_from_yaml(config["hamiltonian"], basis) 20 | return basis, hamiltonian 21 | 22 | 23 | def main(): 24 | number_spins = 36 25 | device = torch.device("cuda:0") 26 | amplitude = torch.nn.Sequential( 27 | nqs.Unpack(number_spins), 28 | torch.nn.Linear(number_spins, 144), 29 | torch.nn.ReLU(), 30 | torch.nn.Linear(144, 1, bias=False), 31 | ).to(device) 32 | phase = torch.nn.Sequential( 33 | nqs.Unpack(number_spins), 34 | torch.nn.Linear(number_spins, 144), 35 | torch.nn.ReLU(), 36 | torch.nn.Linear(144, 1, bias=False), 37 | ).to(device) 38 | combined_state = nqs.combine_amplitude_and_phase(amplitude, phase) 39 | basis, hamiltonian = load_basis_and_hamiltonian() 40 | basis.build() 41 | 42 | states, _, info = nqs.sample_some( 43 | amplitude, 44 | basis, 45 | nqs.SamplingOptions(number_samples=1, number_chains=1, device=device, mode="full") 46 | ) 47 | 48 | for n in [2000000]: 49 | tick = time.time() 50 | local_energies = nqs.local_values( 51 | states[:n], 52 | hamiltonian, 53 | combined_state, 54 | batch_size=2048, # 2048 55 | ) 56 | tock = time.time() 57 | logger.info("For n = {} local values were computed in {:.2f} seconds", n, tock - tick) 58 | 59 | 60 | if __name__ == "__main__": 61 | main() 62 | -------------------------------------------------------------------------------- /benchmark/benchmark_jacobian.py: -------------------------------------------------------------------------------- 1 | import timeit 2 | import numpy as np 3 | import torch 4 | from torch import Tensor 5 | 6 | from nqs_playground import * 7 | 8 | BATCH_SIZE = 1000 9 | REPEAT = 5 10 | NUMBER = 5 11 | 12 | def make_network(): 13 | # return torch.jit.script(torch.nn.Sequential( 14 | # torch.nn.Linear(5, 6), 15 | # torch.nn.ReLU(), 16 | # torch.nn.Linear(6, 1, bias=False), 17 | # )) 18 | return torch.jit.script(torch.nn.Sequential( 19 | Unpack(30), 20 | torch.nn.Linear(30, 128), 21 | torch.nn.ReLU(), 22 | torch.nn.Linear(128, 128), 23 | torch.nn.ReLU(), 24 | torch.nn.Linear(128, 128), 25 | torch.nn.ReLU(), 26 | torch.nn.Linear(128, 1, bias=False), 27 | )) 28 | 29 | 30 | def run_gpu(): 31 | module = make_network() 32 | module.to(device='cuda:0') 33 | inputs = torch.rand((BATCH_SIZE, 30), device='cuda:0') 34 | r = np.array(timeit.repeat(lambda: nqs_playground._jacobian.jacobian_cuda(module, inputs, devices=['cuda:0', 'cuda:1']), repeat=REPEAT, number=NUMBER)) 35 | # r = np.array(timeit.repeat(lambda: jacobian(module, inputs), repeat=REPEAT, number=NUMBER)) 36 | r *= 1000 / (BATCH_SIZE * NUMBER) 37 | return r 38 | 39 | 40 | def _jacobian(module: torch.nn.Module, inputs: Tensor) -> Tensor: 41 | r"""Trivial implementation of ``jacobian``. It is used to assess 42 | correctness of fancier techniques. 43 | """ 44 | parameters = list(module.parameters()) 45 | out = inputs.new_empty( 46 | [inputs.size(0), sum(map(torch.numel, parameters))], dtype=parameters[0].dtype 47 | ) 48 | outputs = module(inputs) 49 | for i in range(inputs.size(0)): 50 | dws = torch.autograd.grad([outputs[i]], parameters, retain_graph=True) 51 | torch.cat([dw.flatten() for dw in dws], out=out[i]) 52 | return out 53 | 54 | 55 | def run_cpu(): 56 | module = make_network() 57 | inputs = torch.randint(0, 1<<30 - 1, size=(BATCH_SIZE, 8), dtype=torch.int64, device='cpu') 58 | 59 | r1 = _jacobian(module, inputs) 60 | r2 = jacobian_simple(module, inputs) 61 | assert torch.isclose(r1, r2, rtol=1e-4, atol=1e-6).all() 62 | 63 | def f(): 64 | # return _jacobian(module, inputs) 65 | return jacobian_simple(module, inputs) 66 | 67 | r = np.array(timeit.repeat(f, repeat=REPEAT, number=NUMBER)) 68 | r *= 1000 / (BATCH_SIZE * NUMBER) 69 | return r 70 | 71 | 72 | # m = make_network() 73 | # xs = torch.rand((BATCH_SIZE, 5), device='cpu') 74 | # j1 = nqs_playground._jacobian.jacobian_simple(m, xs) 75 | # 76 | # m.cuda() 77 | # xs = xs.cuda() 78 | # j2 = nqs_playground._jacobian.jacobian_cuda(m, xs, devices=['cuda:0', 'cuda:1']) 79 | # j2 = j2.cpu() 80 | # 81 | # for i in range(j1.size(0)): 82 | # if not torch.allclose(j1[i], j2[i]): 83 | # print(i) 84 | # print(torch.allclose(j1, j2)) 85 | 86 | # print(run_gpu()) 87 | print(run_cpu()) 88 | -------------------------------------------------------------------------------- /benchmark/benchmark_zanella.py: -------------------------------------------------------------------------------- 1 | import json 2 | import lattice_symmetries as ls 3 | from loguru import logger 4 | import nqs_playground as nqs 5 | import numpy as np 6 | import os 7 | import time 8 | import torch 9 | from scipy import stats 10 | import subprocess 11 | import yaml 12 | 13 | 14 | def get_processor_name(): 15 | result = subprocess.run(["lscpu", "-J"], check=False, capture_output=True) 16 | if result.returncode != 0: 17 | logger.warn( 18 | "Failed to get processor name: {} returned error code {}: {}", 19 | result.args, 20 | result.returncode, 21 | result.stderr, 22 | ) 23 | return None 24 | for obj in json.loads(result.stdout)["lscpu"]: 25 | if obj["field"].startswith("Model name"): 26 | return obj["data"] 27 | 28 | 29 | def get_gpu_name(): 30 | result = subprocess.run(["nvidia-smi", "-L"], check=False, capture_output=True) 31 | if result.returncode != 0: 32 | logger.warn( 33 | "Failed to get processor name: {} returned error code {}: {}", 34 | result.args, 35 | result.returncode, 36 | result.stderr, 37 | ) 38 | return None 39 | for line in filter( 40 | lambda s: s.startswith("GPU 0:"), result.stdout.decode("utf-8").split("\n") 41 | ): 42 | line = line.strip("GPU 0:").strip(" ") 43 | return line 44 | 45 | 46 | def make_basis(number_spins: int) -> ls.SpinBasis: 47 | parity = (number_spins - 1) - np.arange(number_spins, dtype=np.int32) 48 | translation = (np.arange(number_spins, dtype=np.int32) + 1) % number_spins 49 | return ls.SpinBasis( 50 | ls.Group([ls.Symmetry(parity, sector=0), ls.Symmetry(translation, sector=0)]), 51 | number_spins=number_spins, 52 | hamming_weight=number_spins // 2, 53 | spin_inversion=1 if number_spins % 2 == 0 else None, 54 | ) 55 | 56 | 57 | def make_network(number_spins: int) -> torch.nn.Module: 58 | module = torch.nn.Sequential( 59 | nqs.Unpack(number_spins), 60 | torch.nn.Linear(number_spins, 64), 61 | torch.nn.ReLU(inplace=True), 62 | torch.nn.Linear(64, 64), 63 | torch.nn.ReLU(inplace=True), 64 | torch.nn.Linear(64, 1, bias=False), 65 | ) 66 | module = torch.jit.script(module) 67 | return module 68 | 69 | 70 | def profile_one(number_spins: int, device: torch.device, number_chains: int = 10): 71 | basis = make_basis(number_spins) 72 | log_amplitude_fn = make_network(basis.number_spins) 73 | log_amplitude_fn.to(device=device) 74 | options = nqs.SamplingOptions( 75 | number_chains=number_chains, 76 | number_samples=1, 77 | sweep_size=1, 78 | number_discarded=10, 79 | device=device, 80 | ) 81 | # chain_lengths = [1, 400, 800, 1200, 1400, 1600, 1800, 2000, 2200, 2400] 82 | chain_lengths = [1, 800, 1600, 2400] 83 | results = {} 84 | for method in ["zanella"]: 85 | line = [] 86 | logger.debug("Measuring '{}'...", method) 87 | for number_samples in chain_lengths: 88 | t1 = time.time() 89 | _ = nqs.sample_some( 90 | log_amplitude_fn, 91 | basis, 92 | options._replace(number_samples=number_samples, mode=method), 93 | ) 94 | t2 = time.time() 95 | line.append(t2 - t1) 96 | 97 | regression = stats.linregress(chain_lengths[1:], line[1:]) 98 | results[method] = { 99 | "raw": line, 100 | "slope": regression.slope, 101 | "slope_err": regression.stderr, 102 | "intercept": regression.intercept, 103 | "intercept_err": regression.intercept_stderr, 104 | } 105 | return results 106 | 107 | 108 | def profile_all(number_chains: int = 1): 109 | device = "cuda" if torch.cuda.is_available() else "cpu" 110 | logger.info("Running on '{}'...", device) 111 | 112 | header = "Date: {}\n".format(time.asctime()) 113 | cpu = get_processor_name() 114 | if cpu is not None: 115 | header += "CPU: {}\n".format(cpu) 116 | gpu = get_gpu_name() 117 | if gpu is not None: 118 | header += "GPU: {}\n".format(gpu) 119 | header += "number_spins\tmetropolis\tmetropolis_error\tzanella\tzanella_error" 120 | 121 | results = [] 122 | for system_size in [16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192]: 123 | r = profile_one(system_size, device) 124 | results.append( 125 | ( 126 | system_size, 127 | r["metropolis"]["slope"], 128 | r["metropolis"]["slope_err"], 129 | r["zanella"]["slope"], 130 | r["zanella"]["slope_err"], 131 | ) 132 | ) 133 | 134 | results = np.asarray(results, dtype=np.float64) 135 | script_dir = os.path.dirname(os.path.realpath(__file__)) 136 | np.savetxt( 137 | os.path.join(script_dir, "data", "timing_chain_{}.dat".format(number_chains)), 138 | results, 139 | header=header, 140 | ) 141 | 142 | 143 | def main(): 144 | print(profile_one(65, "cuda", 2)) 145 | # profile_all() 146 | 147 | 148 | if __name__ == "__main__": 149 | main() 150 | -------------------------------------------------------------------------------- /build_conda_package.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -p gpu_short 3 | #SBATCH -n 1 -c 16 4 | #SBATCH -t 1:00:00 5 | 6 | module load 2019 7 | module load GCC/7.3.0-2.30 8 | module load CUDA/10.0.130-GCC-7.3.0-2.30 9 | export PATH=/sw/arch/RedHatEnterpriseServer7/EB_production/2019/software/CUDA/10.0.130-GCC-7.3.0-2.30/bin:/sw/arch/RedHatEnterpriseServer7/EB_production/2019/software/binutils/2.30-GCCcore-7.3.0/bin:/sw/arch/RedHatEnterpriseServer7/EB_production/2019/software/GCCcore/7.3.0/bin:/home/twesterh/conda/condabin:/hpc/sw/hpc/bin:/hpc/sw/hpc/sbin:/usr/lib64/qt-3.3/bin:/hpc/eb/modules-tcl-1.923/bin:/usr/local/bin:/usr/bin:/usr/local/sbin:/usr/sbin 10 | export LD_LIBRARY_PATH=/sw/arch/RedHatEnterpriseServer7/EB_production/2019/software/CUDA/10.0.130-GCC-7.3.0-2.30/lib64 11 | 12 | . ~/conda/etc/profile.d/conda.sh 13 | conda activate nqs_dev 14 | 15 | pushd conda/gpu 16 | # rm -rf * 17 | # cmake -GNinja .. 18 | conda build -c defaults -c conda-forge . 19 | popd 20 | -------------------------------------------------------------------------------- /build_locally.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -e 4 | set -o pipefail 5 | 6 | USE_CUDA=0 7 | 8 | get_repo_root() { 9 | git rev-parse --show-toplevel 10 | } 11 | 12 | activate_environment() { 13 | declare -r HOST=$(hostname) 14 | case $HOST in 15 | int*.bullx | gcn*.bullx) 16 | # We're on Cartesius 17 | declare -r env_name="cartesius_devel" 18 | declare -r env_file="cartesius/conda-cartesius.yml" 19 | ;; 20 | *) 21 | if [ $USE_CUDA -eq 0 ]; then 22 | declare -r env_name="nqs_devel_cpu" 23 | declare -r env_file="conda-cpu.yml" 24 | else 25 | declare -r env_name="nqs_devel_gpu" 26 | declare -r env_file="conda-gpu.yml" 27 | fi 28 | ;; 29 | esac 30 | if ! conda env list | grep -q "$env_name"; then 31 | echo "You do not have $env_name Conda environment. Creating it..." 32 | conda env create --file "$env_file" 33 | fi 34 | if ! echo "$CONDA_DEFAULT_ENV" | grep -q "$env_name"; then 35 | echo "Activating $env_name environment..." 36 | # if ! which activate; then 37 | # . $(dirname "$CONDA_EXE")/../etc/profile.d/conda.sh 38 | # fi 39 | conda activate "$env_name" 40 | fi 41 | } 42 | 43 | build_cxx_code() { 44 | echo "Building C++ extension code..." 45 | mkdir -vp build 46 | pushd build 47 | declare -r site_packages_dir=$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])') 48 | cmake -GNinja \ 49 | $CMAKE_ARGS \ 50 | -DCMAKE_CUDA_FLAGS="-cudart shared --compiler-options -march=nehalem" \ 51 | -DCMAKE_CXX_FLAGS="$CXXFLAGS -march=nehalem" \ 52 | -DCMAKE_C_FLAGS="$CFLAGS -march=nehalem" \ 53 | -DCMAKE_PREFIX_PATH="$site_packages_dir/torch/share/cmake" \ 54 | -DCMAKE_BUILD_TYPE=Release \ 55 | -DNQS_PLAYGROUND_USE_CUDA=$USE_CUDA \ 56 | .. 57 | cmake --build . --target install 58 | popd 59 | echo "Done building C++ code!" 60 | } 61 | 62 | install_python_package() { 63 | echo "Installing Python package..." 64 | python3 -m pip install -e . 65 | } 66 | 67 | print_help() { 68 | echo "" 69 | echo "Usage: ./build_locally.sh [--help] [--cuda]" 70 | echo "" 71 | echo "This script builds and installs nqs_playground locally." 72 | echo "" 73 | echo "Options:" 74 | echo " --help Display this message." 75 | echo " --cuda Compile with GPU support." 76 | echo "" 77 | } 78 | 79 | main() { 80 | while [ $# -gt 0 ]; do 81 | key="$1" 82 | case $key in 83 | --cuda) 84 | USE_CUDA=1 85 | shift 86 | ;; 87 | --help) 88 | print_help 89 | exit 0 90 | ;; 91 | *) 92 | echo "Error: unexpected argument '$1'" 93 | print_help 94 | exit 1 95 | ;; 96 | esac 97 | done 98 | cd "$(get_repo_root)" 99 | activate_environment 100 | build_cxx_code 101 | install_python_package 102 | } 103 | 104 | main "$@" 105 | -------------------------------------------------------------------------------- /cartesius/build_pytorch.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -p gpu -n 1 -c 16 3 | #SBATCH --time 12:00:00 --mem 60G 4 | 5 | set -e 6 | 7 | export PYTORCH_VERSION=v1.8.0 # 9dfbfe9 # a7cf04ec40e487286ad3f8068fa18321f3474dd2 # master 8 | export VISION_VERSION=v0.9.0 # 9dfbfe9 # a7cf04ec40e487286ad3f8068fa18321f3474dd2 # master 9 | 10 | module load 2020 11 | module load GCC/9.3.0 CUDA/11.0.2-GCC-9.3.0 cuDNN/8.0.3.33-gcccuda-2020a 12 | 13 | . $HOME/conda/etc/profile.d/conda.sh 14 | conda activate cartesius_devel 15 | 16 | export CFLAGS="-march=native -ftree-vectorize -fPIC -fstack-protector-strong -fno-plt -O3 -ffunction-sections" 17 | export CXXFLAGS="-fvisibility-inlines-hidden -fmessage-length=0 -march=native -ftree-vectorize -fPIC -fstack-protector-strong -fno-plt -O3 -ffunction-sections" 18 | export LDFLAGS="-Wl,-O3 -Wl,--sort-common -Wl,--as-needed -Wl,-z,relro -Wl,-z,now -Wl,--disable-new-dtags -Wl,--gc-sections" 19 | export TORCH_CUDA_ARCH_LIST="3.5" 20 | export TORCH_NVCC_FLAGS="-Xfatbin -compress-all" 21 | export CMAKE_PREFIX_PATH="$CONDA_PREFIX:$CMAKE_PREFIX_PATH" 22 | export USE_CUDNN=1 23 | export USE_FBGEMM=1 24 | export USE_KINETO=0 25 | export USE_NUMPY=1 26 | export BUILD_TEST=0 27 | # export USE_MKL=0 28 | export USE_MKLDNN=1 29 | # export USE_MKLDNN_CBLAS=1 30 | export MKLDNN_CPU_RUNTIME="OMP" 31 | export USE_NNPACK=1 32 | export USE_QNNPACK=0 33 | export USE_XNNPACK=0 34 | export USE_DISTRIBUTED=1 35 | export USE_TENSORPIPE=1 36 | export USE_GLOO=1 37 | export USE_MPI=0 38 | export USE_SYSTEM_NCCL=0 39 | export BUILD_CAFFE2_OPS=1 40 | export BUILD_CAFFE2=1 41 | export USE_IBVERBS=0 42 | export USE_OPENCV=0 43 | export USE_OPENMP=1 44 | export USE_FFMPEG=0 45 | export USE_LEVELDB=0 46 | export USE_LMDB=0 47 | export USE_REDIS=0 48 | export USE_ZSTD=0 49 | export BLAS="MKL" 50 | export MKL_THREADING="OMP" 51 | export ATEN_THREADING="OMP" 52 | 53 | export USE_STATIC_CUDNN=1 54 | export USE_STATIC_NCCL=1 55 | 56 | export CMAKE_ARGS="-DTORCH_CUDA_ARCH_LIST=3.5 $CMAKE_ARGS" 57 | # export INTEL_COMPILER_DIR="/nonexistant" 58 | # export INTEL_MKL_DIR="/nonexistant" 59 | # export INTEL_OMP_DIR="/nonexistant" 60 | 61 | if [ ! -e pytorch ]; then git clone https://github.com/pytorch/pytorch.git; fi 62 | if [ ! -e vision ]; then git clone https://github.com/pytorch/vision.git; fi 63 | # pushd pytorch 64 | # git checkout $PYTORCH_VERSION 65 | # git submodule sync 66 | # git submodule update --init --recursive 67 | # python3 -m pip install -v . 68 | # popd 69 | 70 | pushd vision 71 | git checkout $VISION_VERSION 72 | git submodule sync 73 | git submodule update --init --recursive 74 | python3 setup.py install # build 75 | # python3 -m pip install -v . 76 | popd 77 | -------------------------------------------------------------------------------- /cartesius/conda-cartesius.yml: -------------------------------------------------------------------------------- 1 | name: cartesius_devel 2 | channels: 3 | - twesterhout 4 | - pytorch 5 | - conda-forge 6 | dependencies: 7 | - python=3.8 8 | - pip 9 | # Pytorch dependencies 10 | - numpy 11 | - pyyaml 12 | - scipy 13 | - mkl 14 | - mkl-include 15 | - cmake 16 | - ninja 17 | - magma-cuda110 18 | # Dependencies for Python code 19 | - pip: 20 | - pynvim # Python support for Neovim 21 | - py-spy # For performance analysis 22 | - loguru 23 | - lattice-symmetries 24 | - tensorboard 25 | - psutil 26 | # Code formatting 27 | - black 28 | -------------------------------------------------------------------------------- /cbits/common/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # General Python-independent code 2 | add_library(nqs SHARED 3 | bits512.hpp 4 | bits512.cpp 5 | errors.hpp 6 | errors.cpp 7 | metropolis.hpp 8 | metropolis.cpp 9 | zanella.hpp 10 | zanella.cpp 11 | accumulator.hpp 12 | accumulator.cpp 13 | polynomial.hpp 14 | polynomial.cpp 15 | unpack.cpp 16 | parallel.hpp 17 | parallel.cpp 18 | random.hpp 19 | random.cpp 20 | wrappers.hpp 21 | wrappers.cpp 22 | ) 23 | nqs_cbits_add_low_level_flags(nqs) 24 | target_compile_options(nqs PRIVATE ${TCM_WARNING_FLAGS}) 25 | target_link_libraries(nqs PUBLIC nqs_cbits_Common 26 | nqs_gpu 27 | nqs_cpu_kernels_avx2 28 | nqs_cpu_kernels_avx 29 | nqs_cpu_kernels_sse4) 30 | -------------------------------------------------------------------------------- /cbits/common/accumulator.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include "config.hpp" 32 | #include 33 | #include 34 | #include 35 | #include 36 | 37 | TCM_NAMESPACE_BEGIN 38 | 39 | using OperatorT = 40 | stdext::inplace_function coeff, 41 | gsl::span, gsl::span>) 42 | ->uint64_t, 43 | /*capacity=*/32, /*alignment=*/8>; 44 | 45 | using ForwardT = stdext::inplace_functiontorch::Tensor, 46 | /*capacity=*/32, /*alignment=*/8>; 47 | 48 | struct QuantumOperator { 49 | OperatorT function; 50 | uint64_t max_states; 51 | }; 52 | 53 | auto apply(torch::Tensor spins, OperatorT op, ForwardT psi, uint64_t max_required_size, 54 | uint32_t batch_size) -> torch::Tensor; 55 | 56 | TCM_NAMESPACE_END 57 | -------------------------------------------------------------------------------- /cbits/common/bits512.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020-2021, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #include "bits512.hpp" 30 | 31 | TCM_EXPORT auto operator==(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool 32 | { 33 | for (auto i = 0; i < static_cast(std::size(x.words)); ++i) { 34 | if (x.words[i] != y.words[i]) { return false; } 35 | } 36 | return true; 37 | } 38 | 39 | TCM_EXPORT auto operator!=(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool 40 | { 41 | return !(x == y); 42 | } 43 | 44 | TCM_EXPORT auto operator<(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool 45 | { 46 | for (auto i = 0; i < static_cast(std::size(x.words)); ++i) { 47 | if (x.words[i] < y.words[i]) { return true; } 48 | if (x.words[i] > y.words[i]) { return false; } 49 | } 50 | return false; 51 | } 52 | 53 | TCM_EXPORT auto operator>(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool 54 | { 55 | return y < x; 56 | } 57 | 58 | TCM_EXPORT auto operator<=(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool 59 | { 60 | return !(x > y); 61 | } 62 | 63 | TCM_EXPORT auto operator>=(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool 64 | { 65 | return !(x < y); 66 | } 67 | -------------------------------------------------------------------------------- /cbits/common/bits512.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020-2021, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include "errors.hpp" 32 | #include 33 | #include 34 | 35 | auto operator==(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool; 36 | auto operator!=(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool; 37 | auto operator<(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool; 38 | auto operator>(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool; 39 | auto operator<=(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool; 40 | auto operator>=(ls_bits512 const& x, ls_bits512 const& y) noexcept -> bool; 41 | 42 | TCM_NAMESPACE_BEGIN 43 | 44 | constexpr auto toggle_bit(uint64_t& bits, unsigned const i) noexcept -> void 45 | { 46 | TCM_ASSERT(i < 64U, "index out of bounds"); 47 | bits ^= uint64_t{1} << uint64_t{i}; 48 | } 49 | constexpr auto toggle_bit(ls_bits512& bits, unsigned const i) noexcept -> void 50 | { 51 | TCM_ASSERT(i < 512U, "index out of bounds"); 52 | return toggle_bit(bits.words[i / 64U], i % 64U); 53 | } 54 | constexpr auto test_bit(uint64_t const bits, unsigned const i) noexcept -> bool 55 | { 56 | TCM_ASSERT(i < 64U, "index out of bounds"); 57 | return static_cast((bits >> i) & 1U); 58 | } 59 | constexpr auto test_bit(ls_bits512 const& bits, unsigned const i) noexcept -> bool 60 | { 61 | TCM_ASSERT(i < 512U, "index out of bounds"); 62 | return test_bit(bits.words[i / 64U], i % 64U); 63 | } 64 | constexpr auto set_zero(uint64_t& bits) noexcept -> void { bits = 0UL; } 65 | constexpr auto set_zero(ls_bits512& bits) noexcept -> void 66 | { 67 | for (auto& w : bits.words) { 68 | set_zero(w); 69 | } 70 | } 71 | 72 | TCM_NAMESPACE_END 73 | 74 | // Formatting {{{ 75 | /// Formatting of bits512 using fmtlib facilities 76 | /// 77 | /// Used only for error reporting. 78 | namespace fmt { 79 | template <> struct formatter : formatter { 80 | // parse is inherited from formatter. 81 | 82 | template 83 | auto format(ls_bits512 const& type, FormatContext& ctx) 84 | -> decltype(formatter::format(std::declval(), 85 | std::declval())) 86 | { 87 | // TODO(twesterhout): This is probably the stupidest possible way to 88 | // implement it... But since we only use it for debugging and error 89 | // reporting, who cares? 90 | auto s = fmt::format("[{}]", fmt::join(type.words, ",")); 91 | return formatter::format(s, ctx); 92 | } 93 | }; 94 | } // namespace fmt 95 | // }}} 96 | -------------------------------------------------------------------------------- /cbits/common/errors.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #include "errors.hpp" 30 | #include 31 | 32 | TCM_NAMESPACE_BEGIN 33 | 34 | namespace detail { 35 | TCM_EXPORT auto assert_fail(char const* expr, char const* file, 36 | size_t const line, char const* function, 37 | std::string const& msg) noexcept -> void 38 | { 39 | std::fprintf( 40 | stderr, 41 | TCM_BUG_MESSAGE 42 | "\nAssertion failed at %s:%zu: %s: \"%s\" evaluated to false: %s\n", 43 | file, line, function, expr, msg.c_str()); 44 | std::terminate(); 45 | } 46 | 47 | TCM_EXPORT auto assert_fail(char const* expr, char const* file, 48 | size_t const line, char const* function, 49 | char const* msg) noexcept -> void 50 | { 51 | std::fprintf( 52 | stderr, 53 | TCM_BUG_MESSAGE 54 | "\nAssertion failed at %s:%zu: %s: \"%s\" evaluated to false: %s\n", 55 | file, line, function, expr, msg); 56 | std::terminate(); 57 | } 58 | 59 | TCM_EXPORT auto failed_to_construct_the_message() noexcept -> void 60 | { 61 | std::fprintf(stderr, 62 | "Failed to construct the message. Calling terminate..."); 63 | std::terminate(); 64 | } 65 | 66 | TCM_EXPORT auto make_what_message(char const* file, size_t const line, 67 | char const* function, 68 | std::string const& description) -> std::string 69 | { 70 | return fmt::format("{}:{}: {}: {}", file, line, function, description); 71 | } 72 | 73 | } // namespace detail 74 | 75 | TCM_NAMESPACE_END 76 | -------------------------------------------------------------------------------- /cbits/common/metropolis.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include "bits512.hpp" 32 | #include "random.hpp" 33 | #include "tensor_info.hpp" 34 | #include 35 | #include 36 | #include 37 | 38 | TCM_NAMESPACE_BEGIN 39 | 40 | class MetropolisGenerator { 41 | private: 42 | gsl::not_null _basis; 43 | gsl::not_null _generator; 44 | 45 | public: 46 | MetropolisGenerator(ls_spin_basis const& basis, 47 | RandomGenerator& generator = global_random_generator()); 48 | ~MetropolisGenerator(); 49 | 50 | MetropolisGenerator(MetropolisGenerator const&) = default; 51 | MetropolisGenerator(MetropolisGenerator&&) noexcept = default; 52 | MetropolisGenerator& operator=(MetropolisGenerator const&) = default; 53 | MetropolisGenerator& operator=(MetropolisGenerator&&) noexcept = default; 54 | 55 | auto operator()(torch::Tensor x, c10::ScalarType dtype) const 56 | -> std::tuple; 57 | 58 | private: 59 | template 60 | auto generate(TensorInfo src, TensorInfo dst, 61 | TensorInfo norm) const -> void; 62 | }; 63 | 64 | TCM_NAMESPACE_END 65 | -------------------------------------------------------------------------------- /cbits/common/parallel.cpp: -------------------------------------------------------------------------------- 1 | #include "parallel.hpp" 2 | 3 | TCM_NAMESPACE_BEGIN 4 | 5 | TCM_EXPORT ThreadPool::ThreadPool() : worker{}, tasks{}, queue_mutex{}, condition{}, stop{false} 6 | { 7 | worker = std::thread{[this] { 8 | for (;;) { 9 | std::function task; 10 | 11 | { 12 | std::unique_lock lock(queue_mutex); 13 | condition.wait(lock, [this] { return stop || !tasks.empty(); }); 14 | if (stop && tasks.empty()) return; 15 | task = std::move(tasks.front()); 16 | tasks.pop(); 17 | } 18 | 19 | task(); 20 | } 21 | }}; 22 | } 23 | 24 | TCM_EXPORT ThreadPool::~ThreadPool() 25 | { 26 | { 27 | std::unique_lock lock(queue_mutex); 28 | stop = true; 29 | } 30 | condition.notify_all(); 31 | worker.join(); 32 | } 33 | 34 | #if 0 35 | namespace detail { 36 | TCM_EXPORT _ThreadPoolBase::_ThreadPoolBase() 37 | : _tasks{}, _queue_mutex{}, _condition{}, _stop{false} 38 | {} 39 | 40 | TCM_EXPORT auto _ThreadPoolBase::run_one_task() -> bool 41 | { 42 | std::function task; 43 | 44 | { 45 | std::unique_lock lock(_queue_mutex); 46 | _condition.wait(lock, [this] { return _stop || !_tasks.empty(); }); 47 | if (_stop && _tasks.empty()) return true; 48 | task = std::move(_tasks.front()); 49 | _tasks.pop(); 50 | } 51 | 52 | task(); 53 | return false; 54 | } 55 | 56 | TCM_EXPORT auto _ThreadPoolBase::stop() -> void 57 | { 58 | { 59 | std::unique_lock lock(_queue_mutex); 60 | _stop = true; 61 | } 62 | _condition.notify_all(); 63 | } 64 | } // namespace detail 65 | #endif 66 | 67 | namespace detail { 68 | TCM_EXPORT auto global_thread_pool() noexcept -> ThreadPool& 69 | { 70 | static ThreadPool pool{}; 71 | return pool; 72 | } 73 | } // namespace detail 74 | 75 | TCM_NAMESPACE_END 76 | -------------------------------------------------------------------------------- /cbits/common/polynomial.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019-2020, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include "accumulator.hpp" 32 | #include 33 | #include 34 | #include 35 | 36 | TCM_NAMESPACE_BEGIN 37 | 38 | class Polynomial { 39 | private: 40 | QuantumOperator _op; 41 | std::vector> _roots; 42 | uint64_t _max_states; 43 | bool _normalising; 44 | 45 | public: 46 | Polynomial(QuantumOperator op, std::vector roots, bool normalising); 47 | Polynomial(Polynomial const&) = delete; 48 | Polynomial(Polynomial&& other) noexcept = default; 49 | Polynomial& operator=(Polynomial const&) = delete; 50 | Polynomial& operator=(Polynomial&&) = delete; 51 | 52 | auto degree() const noexcept -> uint64_t; 53 | auto roots() const noexcept -> gsl::span const>; 54 | 55 | auto operator()(ls_bits512 const& spin, complex_type coeff, gsl::span out_spins, 56 | gsl::span out_coeffs) const -> uint64_t; 57 | 58 | constexpr auto max_states() const noexcept -> uint64_t { return _max_states; } 59 | 60 | private: 61 | struct Buffer; 62 | auto iteration(complex_type root, Buffer& buffer) const -> void; 63 | }; 64 | 65 | TCM_NAMESPACE_END 66 | -------------------------------------------------------------------------------- /cbits/common/random.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #include "random.hpp" 30 | #include 31 | #include 32 | 33 | TCM_NAMESPACE_BEGIN 34 | 35 | namespace detail { 36 | namespace { 37 | auto really_need_that_random_seed_now() -> uint64_t 38 | { 39 | std::random_device random_device; 40 | std::uniform_int_distribution dist; 41 | auto const seed = dist(random_device); 42 | return seed; 43 | } 44 | } // namespace 45 | } // namespace detail 46 | 47 | TCM_EXPORT auto global_random_generator() -> RandomGenerator& 48 | { 49 | static thread_local RandomGenerator generator{detail::really_need_that_random_seed_now()}; 50 | return generator; 51 | } 52 | 53 | TCM_EXPORT auto manual_seed(uint64_t const seed) -> void 54 | { 55 | #pragma omp parallel default(none) firstprivate(seed) 56 | { 57 | auto i = omp_get_thread_num(); 58 | global_random_generator().seed(seed + static_cast(i)); 59 | } 60 | } 61 | 62 | TCM_NAMESPACE_END 63 | -------------------------------------------------------------------------------- /cbits/common/unpack.cpp: -------------------------------------------------------------------------------- 1 | #include "unpack.hpp" 2 | #include "errors.hpp" 3 | 4 | #include "../cpu/kernels.hpp" 5 | #if defined(TCM_USE_CUDA) 6 | # include "../gpu/unpack.hpp" 7 | #endif 8 | 9 | #include 10 | 11 | TCM_NAMESPACE_BEGIN 12 | 13 | namespace { 14 | auto unpack_impl(TensorInfo const& spins, int64_t const number_spins, 15 | c10::Device const device) -> torch::Tensor 16 | { 17 | TCM_CHECK(number_spins > 0, std::domain_error, 18 | fmt::format("invalid number_spins: {}; expected a positive integer", number_spins)); 19 | TCM_CHECK(number_spins <= 64 * spins.size<1>(), std::domain_error, 20 | fmt::format("number_spins too big: {}; expected <= {}", number_spins, 21 | 64 * spins.size<1>())); 22 | auto out = torch::empty(std::initializer_list{spins.size<0>(), number_spins}, 23 | torch::TensorOptions{}.device(device).dtype(torch::kFloat32)); 24 | auto out_info = tensor_info(out); 25 | switch (device.type()) { 26 | case c10::DeviceType::CPU: unpack_cpu(spins, out_info); break; 27 | #if defined(TCM_USE_CUDA) 28 | case c10::DeviceType::CUDA: unpack_cuda(spins, out_info, device); break; 29 | #endif 30 | default: { 31 | #if defined(TCM_USE_CUDA) 32 | auto const error_msg = "'spins' tensor resides on an unsupported " 33 | "device: {}; expected either CPU or CUDA"; 34 | #else 35 | auto const error_msg = "'spins' tensor resides on an unsupported device: {}; expected CPU"; 36 | #endif 37 | TCM_ERROR(std::domain_error, fmt::format(error_msg, c10::DeviceTypeName(device.type()))); 38 | } 39 | } // end switch 40 | return out; 41 | } 42 | 43 | auto hamming_weight_impl(TensorInfo const& spins, int64_t const number_spins, 44 | c10::Device const device) -> torch::Tensor 45 | { 46 | TCM_CHECK(number_spins > 0, std::domain_error, 47 | fmt::format("invalid number_spins: {}; expected a positive integer", number_spins)); 48 | TCM_CHECK(number_spins <= 64 * spins.size<1>(), std::domain_error, 49 | fmt::format("number_spins too big: {}; expected <= {}", number_spins, 50 | 64 * spins.size<1>())); 51 | auto out = torch::empty(std::initializer_list{spins.size<0>()}, 52 | torch::TensorOptions{}.device(device).dtype(torch::kFloat32)); 53 | auto out_info = tensor_info(out); 54 | switch (device.type()) { 55 | case c10::DeviceType::CPU: hamming_weight_cpu(spins, out_info); break; 56 | #if defined(TCM_USE_CUDA) 57 | case c10::DeviceType::CUDA: 58 | TCM_ERROR(std::domain_error, "hamming_weight() is not yet implemented for CUDA"); 59 | #endif 60 | default: { 61 | #if defined(TCM_USE_CUDA) 62 | auto const error_msg = "'spins' tensor resides on an unsupported " 63 | "device: {}; expected either CPU or CUDA"; 64 | #else 65 | auto const error_msg = "'spins' tensor resides on an unsupported device: {}; expected CPU"; 66 | #endif 67 | TCM_ERROR(std::domain_error, fmt::format(error_msg, c10::DeviceTypeName(device.type()))); 68 | } 69 | } // end switch 70 | return out; 71 | } 72 | 73 | } // namespace 74 | 75 | TCM_EXPORT auto unpack(torch::Tensor spins, int64_t const number_spins) -> torch::Tensor 76 | { 77 | auto const shape = spins.sizes(); 78 | auto const device = spins.device(); 79 | switch (shape.size()) { 80 | case 1: spins = torch::unsqueeze(spins, /*dim=*/1); TCM_FALLTHROUGH; 81 | case 2: 82 | return unpack_impl(tensor_info(spins, "spins"), number_spins, device); 83 | default: 84 | TCM_ERROR(std::domain_error, fmt::format("spins has wrong shape: [{}]; expected either a " 85 | "one- or two-dimensional tensor", 86 | fmt::join(shape, ", "))); 87 | } // end switch 88 | } 89 | 90 | TCM_EXPORT auto hamming_weight(torch::Tensor spins, int64_t const number_spins) -> torch::Tensor 91 | { 92 | auto const shape = spins.sizes(); 93 | auto const device = spins.device(); 94 | switch (shape.size()) { 95 | case 1: spins = torch::unsqueeze(spins, /*dim=*/1); TCM_FALLTHROUGH; 96 | case 2: 97 | return hamming_weight_impl(tensor_info(spins, "spins"), number_spins, 98 | device); 99 | default: 100 | TCM_ERROR(std::domain_error, fmt::format("spins has wrong shape: [{}]; expected either a " 101 | "one- or two-dimensional tensor", 102 | fmt::join(shape, ", "))); 103 | } // end switch 104 | } 105 | 106 | static auto torch_script_operators = 107 | torch::RegisterOperators{} 108 | .op("tcm::unpack", 109 | torch::RegisterOperators::options() 110 | .catchAllKerneltorch::Tensor, &unpack>()) 111 | .op("tcm::hamming_weight", 112 | torch::RegisterOperators::options() 113 | .catchAllKerneltorch::Tensor, &hamming_weight>()); 114 | 115 | TCM_NAMESPACE_END 116 | -------------------------------------------------------------------------------- /cbits/common/unpack.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "config.hpp" 4 | #include 5 | 6 | TCM_NAMESPACE_BEGIN 7 | 8 | auto unpack(torch::Tensor, int64_t) -> torch::Tensor; 9 | auto hamming_weight(torch::Tensor, int64_t) -> torch::Tensor; 10 | 11 | TCM_NAMESPACE_END 12 | -------------------------------------------------------------------------------- /cbits/common/wrappers.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #include "wrappers.hpp" 30 | #include "bits512.hpp" 31 | #include "errors.hpp" 32 | 33 | TCM_NAMESPACE_BEGIN 34 | 35 | TCM_EXPORT auto check_status_code(ls_error_code const code) -> void 36 | { 37 | if (TCM_UNLIKELY(code != LS_SUCCESS)) { 38 | auto deleter = [](auto const* s) { ls_destroy_string(s); }; 39 | auto c_str = 40 | std::unique_ptr{ls_error_to_string(code), deleter}; 41 | TCM_ERROR( 42 | std::runtime_error, 43 | fmt::format("lattice_symmetries failed with error code {}: {}", code, c_str.get())); 44 | } 45 | } 46 | 47 | TCM_EXPORT auto view_as_operator(ls_operator const& op) -> std::tuple 48 | { 49 | auto max_required_size = ls_operator_max_buffer_size(&op); 50 | auto action = [&op, max_required_size](ls_bits512 const& spin, std::complex coeff, 51 | gsl::span out_spins, 52 | gsl::span> out_coeffs) { 53 | TCM_CHECK(out_spins.size() >= max_required_size, std::runtime_error, 54 | fmt::format("out_spins buffer is too short: {}; expected at least {}", 55 | out_spins.size(), max_required_size)); 56 | TCM_CHECK(out_coeffs.size() >= max_required_size, std::runtime_error, 57 | fmt::format("out_coeffs buffer is too short: {}; expected at least {}", 58 | out_coeffs.size(), max_required_size)); 59 | struct cxt_t { 60 | ls_bits512* spins_ptr; 61 | std::complex* coeffs_ptr; 62 | uint64_t offset; 63 | uint64_t max_size; 64 | } cxt{out_spins.data(), out_coeffs.data(), 0, max_required_size}; 65 | auto callback = [](ls_bits512 const* _spin, void const* _coeff, void* cxt_raw) { 66 | auto* _cxt = static_cast(cxt_raw); 67 | TCM_CHECK(_cxt->offset < _cxt->max_size, std::runtime_error, ""); 68 | _cxt->spins_ptr[_cxt->offset] = *_spin; 69 | _cxt->coeffs_ptr[_cxt->offset] = *static_cast const*>(_coeff); 70 | ++(_cxt->offset); 71 | return LS_SUCCESS; 72 | }; 73 | check_status_code(ls_operator_apply(&op, &spin, callback, &cxt)); 74 | TCM_CHECK(cxt.offset <= max_required_size, std::runtime_error, "buffer overflow"); 75 | return cxt.offset; 76 | }; 77 | return {std::move(action), max_required_size}; 78 | } 79 | 80 | TCM_EXPORT auto random_spin(ls_spin_basis const& basis, RandomGenerator& generator) -> ls_bits512 81 | { 82 | auto const number_spins = ls_get_number_spins(&basis); 83 | ls_bits512 spin; 84 | set_zero(spin); 85 | if (ls_get_hamming_weight(&basis) == -1) { 86 | auto i = 0U; 87 | for (; i < number_spins / 64U; ++i) { 88 | spin.words[i] = std::uniform_int_distribution{}(generator); 89 | } 90 | auto const rest = number_spins % 64U; 91 | if (rest != 0) { 92 | spin.words[i] = 93 | std::uniform_int_distribution{0, (uint64_t{1} << rest) - 1}(generator); 94 | } 95 | } 96 | else { 97 | auto const m = static_cast(ls_get_hamming_weight(&basis)); 98 | std::vector buffer(number_spins); 99 | std::fill(std::begin(buffer), std::next(std::begin(buffer), m), uint8_t{1}); 100 | std::fill(std::next(std::begin(buffer), m), std::end(buffer), uint8_t{0}); 101 | std::shuffle(std::begin(buffer), std::end(buffer), generator); 102 | auto word = 0U; 103 | auto i = 0U; 104 | for (auto const b : buffer) { 105 | if (b) { spin.words[word] |= static_cast(b) << i; } 106 | ++i; 107 | if (i == 64) { 108 | i = 0U; 109 | ++word; 110 | } 111 | } 112 | } 113 | ls_bits512 repr; 114 | set_zero(repr); 115 | std::complex character; 116 | double norm; 117 | ls_get_state_info(&basis, &spin, &repr, &character, &norm); 118 | if (norm > 0.0) { return repr; } 119 | return random_spin(basis, generator); 120 | } 121 | 122 | TCM_NAMESPACE_END 123 | -------------------------------------------------------------------------------- /cbits/common/wrappers.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020-2021, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include "accumulator.hpp" 32 | #include "random.hpp" 33 | #include 34 | 35 | TCM_NAMESPACE_BEGIN 36 | 37 | auto check_status_code(ls_error_code code) -> void; 38 | auto view_as_operator(ls_operator const& op) -> std::tuple; 39 | 40 | auto random_spin(ls_spin_basis const& basis, RandomGenerator& generator = global_random_generator()) 41 | -> ls_bits512; 42 | 43 | TCM_NAMESPACE_END 44 | -------------------------------------------------------------------------------- /cbits/common/zanella.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020-2021, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include "tensor_info.hpp" 32 | #include 33 | #include 34 | #include 35 | 36 | TCM_NAMESPACE_BEGIN 37 | 38 | class ZanellaGenerator { 39 | private: 40 | gsl::not_null _basis; 41 | std::vector> _edges; 42 | 43 | public: 44 | ZanellaGenerator(ls_spin_basis const& basis, std::vector> edges); 45 | ~ZanellaGenerator(); 46 | 47 | auto operator()(torch::Tensor x) const -> std::tuple; 48 | 49 | auto max_states() const noexcept -> uint64_t 50 | { 51 | auto const number_spins = ls_get_number_spins(_basis); 52 | auto const hamming_weight = ls_get_hamming_weight(_basis); 53 | if (hamming_weight != -1) { 54 | auto const m = static_cast(hamming_weight); 55 | return m * (number_spins - m); 56 | } 57 | // By construction, number_spins > 1 58 | return number_spins * (number_spins - 1); 59 | } 60 | 61 | private: 62 | auto generate_general(ls_bits512 const& spin, TensorInfo out) const -> int64_t; 63 | auto project_states(TensorInfo spins, int64_t count, int num_threads) const -> int64_t; 64 | 65 | // auto generate(ls_bits512 const& spin, TensorInfo out) const -> unsigned; 66 | }; 67 | 68 | auto zanella_choose_samples(torch::Tensor weights, int64_t number_samples, double time_step, 69 | c10::Device device) -> torch::Tensor; 70 | 71 | TCM_NAMESPACE_END 72 | -------------------------------------------------------------------------------- /cbits/cpu/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Library of CPU specific kernels 2 | 3 | foreach(my_arch sse4 avx avx2) 4 | set(my_target "nqs_cpu_kernels_${my_arch}") 5 | add_library(${my_target} OBJECT 6 | kernels.hpp 7 | kernels.cpp 8 | ) 9 | nqs_cbits_add_low_level_flags(${my_target}) 10 | target_compile_options(${my_target} PRIVATE ${TCM_WARNING_FLAGS}) 11 | target_compile_definitions(${my_target} PRIVATE TCM_COMPILE_${my_arch}) 12 | target_link_libraries(${my_target} PUBLIC nqs_cbits_Common) 13 | target_link_libraries(${my_target} PRIVATE vectorclass) 14 | endforeach() 15 | target_compile_options(nqs_cpu_kernels_sse4 PRIVATE -msse4 -msse4.1 -msse4.2 -mpopcnt) 16 | target_compile_definitions(nqs_cpu_kernels_sse4 PRIVATE TCM_ADD_DISPATCH_CODE) 17 | target_compile_options(nqs_cpu_kernels_avx PRIVATE -mavx) 18 | target_compile_options(nqs_cpu_kernels_avx2 PRIVATE -mavx2 -mbmi2 -mfma) 19 | -------------------------------------------------------------------------------- /cbits/cpu/kernels.cpp: -------------------------------------------------------------------------------- 1 | #include "kernels.hpp" 2 | #include "../common/config.hpp" 3 | #include "../common/errors.hpp" 4 | 5 | #include 6 | #include 7 | 8 | #if defined(TCM_COMPILE_avx2) 9 | # define unpack_one_simd unpack_one_avx2 10 | #elif defined(TCM_COMPILE_avx) 11 | # define unpack_one_simd unpack_one_avx 12 | #elif defined(TCM_COMPILE_sse4) 13 | # define unpack_one_simd unpack_one_sse4 14 | #else 15 | # error "One of TCM_COMPILE_avx2, TCM_COMPILE_avx, or TCM_COMPILE_sse4 must be defined" 16 | #endif 17 | 18 | TCM_NAMESPACE_BEGIN 19 | 20 | namespace { 21 | auto unpack_byte(uint8_t const bits) noexcept -> vcl::Vec8f 22 | { 23 | auto const up = vcl::Vec8f{1.0f}; 24 | auto const down = vcl::Vec8f{-1.0f}; 25 | vcl::Vec8fb mask; 26 | mask.load_bits(bits); 27 | return vcl::select(mask, up, down); 28 | } 29 | 30 | auto unpack_word(uint64_t x, float* out) noexcept -> void 31 | { 32 | for (auto i = 0; i < 8; ++i, out += 8, x >>= 8) { 33 | unpack_byte(static_cast(x & 0xFF)).store(out); 34 | } 35 | } 36 | } // namespace 37 | 38 | TCM_EXPORT auto unpack_one_simd(uint64_t const x[], unsigned const number_spins, 39 | float* out) noexcept -> void 40 | { 41 | constexpr auto block = 64U; 42 | auto const words = number_spins / block; 43 | for (auto i = 0U; i < words; ++i, out += block) { 44 | unpack_word(x[i], out); 45 | } 46 | auto const rest_words = number_spins % block; 47 | if (rest_words != 0) { 48 | auto const bytes = rest_words / 8U; 49 | auto const rest_bytes = rest_words % 8U; 50 | auto y = x[words]; 51 | for (auto i = 0U; i < bytes; ++i, out += 8, y >>= 8U) { 52 | unpack_byte(static_cast(y & 0xFF)).store(out); 53 | } 54 | if (rest_bytes != 0) { 55 | auto const t = unpack_byte(static_cast(y & 0xFF)); 56 | t.store_partial(static_cast(rest_bytes), out); 57 | } 58 | } 59 | } 60 | 61 | namespace { 62 | auto popcount(unsigned x) noexcept { return __builtin_popcount(x); } 63 | auto popcount(unsigned long x) noexcept { return __builtin_popcountl(x); } 64 | auto popcount(unsigned long long x) noexcept { return __builtin_popcountll(x); } 65 | } // namespace 66 | 67 | #if defined(TCM_ADD_DISPATCH_CODE) 68 | template 69 | auto hamming_weight_cpu(TensorInfo const& spins, 70 | TensorInfo const& out) noexcept -> void 71 | { 72 | constexpr auto block = 64; 73 | auto const words = (block * spins.size<1>() + block - 1) / block; 74 | for (auto i = int64_t{0}; i < spins.size<0>(); ++i) { 75 | auto acc = 0; 76 | auto const current = row(spins, i); 77 | for (auto j = int64_t{0}; j < words; ++j) { 78 | acc += popcount(current[j]); 79 | } 80 | out[i] = static_cast(acc); 81 | } 82 | } 83 | 84 | template TCM_EXPORT auto hamming_weight_cpu(TensorInfo const&, 85 | TensorInfo const&) noexcept -> void; 86 | 87 | TCM_EXPORT auto unpack_cpu(TensorInfo const& src_info, 88 | TensorInfo const& dst_info) -> void 89 | { 90 | if (src_info.size<0>() == 0) { return; } 91 | TCM_CHECK(dst_info.strides[0] == dst_info.size<1>(), std::invalid_argument, 92 | "unpack_cpu does not support strided output tensors"); 93 | TCM_CHECK(dst_info.size<1>() <= 64 * src_info.size<1>(), std::invalid_argument, 94 | "dst_info is too wide"); 95 | 96 | using unpack_one_fn_t = auto (*)(uint64_t const[], unsigned, float*) noexcept->void; 97 | auto unpack_ptr = []() -> unpack_one_fn_t { 98 | if (__builtin_cpu_supports("avx2")) { return &unpack_one_avx2; } 99 | if (__builtin_cpu_supports("avx")) { return &unpack_one_avx; } 100 | return &unpack_one_sse4; 101 | }(); 102 | 103 | auto const number_spins = static_cast(dst_info.sizes[1]); 104 | for (auto i = int64_t{0}; i < src_info.size<0>(); ++i) { 105 | (*unpack_ptr)(src_info.data + i * src_info.stride<0>(), number_spins, 106 | dst_info.data + i * dst_info.stride<0>()); 107 | } 108 | } 109 | #endif 110 | 111 | TCM_NAMESPACE_END 112 | -------------------------------------------------------------------------------- /cbits/cpu/kernels.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019-2021, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include "../common/tensor_info.hpp" 32 | #include 33 | #include 34 | 35 | TCM_NAMESPACE_BEGIN 36 | 37 | auto unpack_cpu(TensorInfo const& spins, TensorInfo const& out) 38 | -> void; 39 | 40 | template 41 | auto hamming_weight_cpu(TensorInfo const& spins, 42 | TensorInfo const& out) noexcept -> void; 43 | 44 | auto unpack_one_avx2(uint64_t const[], unsigned, float*) noexcept -> void; 45 | auto unpack_one_avx(uint64_t const[], unsigned, float*) noexcept -> void; 46 | auto unpack_one_sse4(uint64_t const[], unsigned, float*) noexcept -> void; 47 | 48 | TCM_NAMESPACE_END 49 | -------------------------------------------------------------------------------- /cbits/gpu/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Determine if we have nvcc 2 | 3 | if(NQS_PLAYGROUND_USE_CUDA) 4 | include(CheckLanguage) 5 | check_language(CUDA) 6 | endif() 7 | if(CMAKE_CUDA_COMPILER) 8 | enable_language(CUDA) 9 | if(NOT DEFINED CMAKE_CUDA_STANDARD) 10 | set(CMAKE_CUDA_STANDARD 14) 11 | set(CMAKE_CUDA_STANDARD_REQUIRED ON) 12 | endif() 13 | else() 14 | message(WARNING "[nqs_cbits] Compiling without CUDA support") 15 | endif() 16 | 17 | # Library of CUDA specific kernels 18 | if(CMAKE_CUDA_COMPILER) 19 | add_library(nqs_gpu OBJECT 20 | unpack.hpp 21 | unpack.cu 22 | ) 23 | nqs_cbits_add_low_level_flags(nqs_gpu) 24 | target_compile_options(nqs_gpu PRIVATE --compiler-options "-W -Wall -Wextra") 25 | target_compile_options(nqs_gpu PRIVATE 26 | $<$:-Xfatbin -compress-all 27 | -gencode arch=compute_35,code=sm_35>) 28 | target_link_libraries(nqs_gpu PUBLIC nqs_cbits_Common) 29 | target_compile_definitions(nqs_gpu PUBLIC TCM_USE_CUDA) 30 | else() 31 | add_library(nqs_gpu INTERFACE) 32 | endif() 33 | -------------------------------------------------------------------------------- /cbits/gpu/unpack.cu: -------------------------------------------------------------------------------- 1 | #include "unpack.hpp" 2 | #include 3 | #include 4 | 5 | TCM_NAMESPACE_BEGIN 6 | 7 | namespace detail { 8 | 9 | struct SpinsInfo { 10 | uint64_t const* const data; 11 | int32_t const stride; 12 | }; 13 | 14 | struct OutInfo { 15 | float* const data; 16 | int32_t const shape[2]; 17 | int32_t const stride[2]; 18 | }; 19 | 20 | __device__ inline auto unpack_rest(uint64_t bits, int32_t const count, float* const out, 21 | int32_t const stride) noexcept -> void 22 | { 23 | for (auto i = 0; i < count; ++i, bits >>= 1) { 24 | out[i * stride] = 2.0f * static_cast(bits & 0x01) - 1.0f; 25 | } 26 | } 27 | 28 | __device__ inline auto unpack_word(uint64_t bits, float* const out, int32_t const stride) noexcept 29 | -> void 30 | { 31 | for (auto i = 0; i < 64; ++i, bits >>= 1) { 32 | out[i * stride] = 2.0f * static_cast(bits & 0x01) - 1.0f; 33 | } 34 | } 35 | 36 | __device__ inline auto unpack_full(uint64_t const bits[], int32_t const count, float* out, 37 | int32_t const stride) noexcept -> void 38 | { 39 | constexpr auto block = 64; 40 | 41 | auto i = 0; 42 | for (; i < count / block; ++i, out += block * stride) { 43 | unpack_word(bits[i], out, stride); 44 | } 45 | { 46 | auto const rest = count % block; 47 | if (rest != 0) { 48 | unpack_rest(bits[i], rest, out, stride); 49 | // out += rest * stride; 50 | } 51 | } 52 | } 53 | 54 | __global__ auto unpack_kernel_cuda(TensorInfo const spins, 55 | TensorInfo const out) -> void 56 | { 57 | auto const idx = blockIdx.x * blockDim.x + threadIdx.x; 58 | auto const stride = blockDim.x * gridDim.x; 59 | for (auto i = idx; i < out.sizes[0]; i += stride) { 60 | unpack_full(spins.data + i * spins.strides[0], out.sizes[1], out.data + i * out.strides[0], 61 | out.strides[1]); 62 | } 63 | } 64 | } // namespace detail 65 | 66 | auto unpack_cuda(TensorInfo const& spins, TensorInfo const& out, 67 | c10::Device const device) -> void 68 | { 69 | // clang-format off 70 | cudaSetDevice(device.index()); 71 | auto stream = at::cuda::getCurrentCUDAStream(); 72 | detail::unpack_kernel_cuda<<>>(spins, out); 74 | // clang-format on 75 | } 76 | 77 | TCM_NAMESPACE_END 78 | -------------------------------------------------------------------------------- /cbits/gpu/unpack.hpp: -------------------------------------------------------------------------------- 1 | #include "../common/bits512.hpp" 2 | #include "../common/tensor_info.hpp" 3 | #include 4 | 5 | TCM_NAMESPACE_BEGIN 6 | 7 | auto unpack_cuda(TensorInfo const& spins, TensorInfo const& out, 8 | c10::Device device) -> void; 9 | 10 | TCM_NAMESPACE_END 11 | -------------------------------------------------------------------------------- /cbits/python/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib") 2 | # Python bindings 3 | pybind11_add_module(_C MODULE NO_EXTRAS 4 | trim.cpp 5 | init.cpp 6 | ) 7 | nqs_cbits_add_low_level_flags(_C) 8 | target_compile_definitions(_C PUBLIC 9 | TORCH_API_INCLUDE_EXTENSION_H=1 10 | TORCH_EXTENSION_NAME="_C") 11 | target_link_libraries(_C PRIVATE nqs ${TORCH_PYTHON_LIBRARY}) 12 | -------------------------------------------------------------------------------- /cbits/python/bind_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include "bind_cuda.hpp" 2 | #include "../trim.hpp" 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | TCM_NAMESPACE_BEGIN 12 | 13 | auto bind_cuda(PyObject* _module) -> void 14 | { 15 | namespace py = pybind11; 16 | auto m = py::module{py::reinterpret_borrow(_module)}; 17 | 18 | } 19 | 20 | TCM_NAMESPACE_END 21 | -------------------------------------------------------------------------------- /cbits/python/bind_cuda.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_cuda(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_heisenberg.cpp: -------------------------------------------------------------------------------- 1 | #include "bind_heisenberg.hpp" 2 | #include "../heisenberg.hpp" 3 | #include "../trim.hpp" 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | TCM_NAMESPACE_BEGIN 13 | 14 | template auto make_call_function() 15 | { 16 | namespace py = pybind11; 17 | return [](Heisenberg const& self, py::array_t x, 18 | std::optional> y) { 19 | TCM_CHECK( 20 | x.ndim() == 1, std::domain_error, 21 | fmt::format("x has invalid shape: [{}]; expected a vector", 22 | fmt::join(x.shape(), x.shape() + x.ndim(), ", "))); 23 | auto src = 24 | gsl::span{x.data(), static_cast(x.shape()[0])}; 25 | auto out = y.value_or(py::array_t{src.size()}); 26 | if (y.has_value()) { 27 | TCM_CHECK(y->ndim() == 1, std::domain_error, 28 | fmt::format( 29 | "y has invalid shape: [{}]; expected a vector", 30 | fmt::join(y->shape(), y->shape() + y->ndim(), ", "))); 31 | } 32 | auto dst = gsl::span{out.mutable_data(), 33 | static_cast(out.shape()[0])}; 34 | self(src, dst); 35 | return out; 36 | }; 37 | } 38 | 39 | auto bind_heisenberg(PyObject* _module) -> void 40 | { 41 | namespace py = pybind11; 42 | auto m = py::module{py::reinterpret_borrow(_module)}; 43 | 44 | py::options options; 45 | options.disable_function_signatures(); 46 | 47 | std::vector keep_alive; 48 | #define DOC(str) trim(keep_alive, str) 49 | 50 | py::class_>(m, "Heisenberg") 51 | .def( 52 | py::init>(), 53 | DOC(R"EOF( 54 | Constructs the Heisenberg Hamiltonian with given exchange couplings. 55 | Full Hamiltonian is: ``∑Jᵢⱼ·σᵢ⊗σⱼ``, where ``σ`` are vector spin 56 | operators. 57 | 58 | :param specs: is a list of couplings. Each element is a tuple 59 | ``(Jᵢⱼ, i, j)``. This defines a term ``Jᵢⱼ·σᵢ⊗σⱼ``. 60 | :param basis: specifies the Hilbert space basis on which the 61 | Hamiltonian is defined.)EOF")) 62 | .def_property_readonly( 63 | "edges", 64 | [](Heisenberg const& self) { 65 | auto edges = self.edges(); 66 | return std::vector{edges.begin(), 67 | edges.end()}; 68 | }, 69 | DOC(R"EOF(Returns the list of couplings)EOF")) 70 | .def_property_readonly("basis", &Heisenberg::basis, DOC(R"EOF( 71 | Returns the Hilbert space basis on which the Hamiltonian is defined.)EOF")) 72 | .def_property_readonly("is_real", &Heisenberg::is_real, DOC(R"EOF( 73 | Returns whether the Hamiltonian is real, i.e. whether all couplings 74 | ``{cᵢⱼ}`` are real.)EOF")) 75 | .def("__call__", make_call_function(), py::arg{"x"}.noconvert(), 76 | py::arg{"y"}.noconvert() = py::none()) 77 | .def("__call__", make_call_function(), py::arg{"x"}.noconvert(), 78 | py::arg{"y"}.noconvert() = py::none()) 79 | .def("__call__", make_call_function>(), 80 | py::arg{"x"}.noconvert(), py::arg{"y"}.noconvert() = py::none()) 81 | .def("__call__", make_call_function>(), 82 | py::arg{"x"}.noconvert(), py::arg{"y"}.noconvert() = py::none()) 83 | .def("to_csr", [](Heisenberg const& self) { 84 | auto [values, indices] = self._to_sparse(); 85 | auto const n = values.size(0); 86 | auto const csr_matrix = 87 | py::module::import("scipy.sparse").attr("csr_matrix"); 88 | auto data = py::cast(values).attr("numpy")(); 89 | auto row_indices = 90 | py::cast(torch::narrow(indices, /*dim=*/1, /*start=*/0, 91 | /*length=*/1)) 92 | .attr("squeeze")() 93 | .attr("numpy")(); 94 | auto col_indices = 95 | py::cast(torch::narrow(indices, /*dim=*/1, /*start=*/1, 96 | /*length=*/1)) 97 | .attr("squeeze")() 98 | .attr("numpy")(); 99 | auto const* basis = 100 | static_cast(self.basis().get()); 101 | return csr_matrix( 102 | std::make_tuple(data, 103 | std::make_tuple(row_indices, col_indices)), 104 | std::make_tuple(basis->number_states(), 105 | basis->number_states())); 106 | }); 107 | 108 | #undef DOC 109 | } 110 | 111 | TCM_NAMESPACE_END 112 | -------------------------------------------------------------------------------- /cbits/python/bind_heisenberg.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_heisenberg(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_jacobian.cpp: -------------------------------------------------------------------------------- 1 | #include "bind_jacobian.hpp" 2 | #include "../jacobian.hpp" 3 | #include "../trim.hpp" 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | TCM_NAMESPACE_BEGIN 11 | 12 | auto bind_jacobian(PyObject* _module) -> void 13 | { 14 | namespace py = pybind11; 15 | auto m = py::module{py::reinterpret_borrow(_module)}; 16 | 17 | std::vector keep_alive; 18 | #define DOC(str) trim(keep_alive, str) 19 | 20 | m.def("_jacobian", &jacobian, DOC(R"EOF( 21 | Computes the jacobian of ``module(inputs)`` with respect to module 22 | parameters and stores the result in ``out``. 23 | 24 | :param module: Jacobian of this module is be computed. 25 | :param inputs: Points at which the Jacobian is evaluated. 26 | :param out: Output tensor where the Jacobian is stored.)EOF"), 27 | py::arg{"module"}.noconvert(), py::arg{"inputs"}.noconvert(), 28 | py::arg{"out"}.noconvert() = py::none(), 29 | py::arg{"num_threads"}.noconvert() = 1, 30 | // NOTE: releasing GIL here is important, see 31 | // https://github.com/pytorch/pytorch/issues/32045 32 | py::call_guard()); 33 | 34 | #undef DOC 35 | } 36 | 37 | TCM_NAMESPACE_END 38 | -------------------------------------------------------------------------------- /cbits/python/bind_jacobian.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_jacobian(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_metropolis.cpp: -------------------------------------------------------------------------------- 1 | #include "bind_metropolis.hpp" 2 | #include "../metropolis.hpp" 3 | #include "../spin_basis.hpp" 4 | #include "../tabu.hpp" 5 | #include "../trim.hpp" 6 | #include "pybind11_helpers.hpp" 7 | 8 | #include "../cpu/kernels.hpp" 9 | 10 | #include 11 | #include 12 | #include 13 | 14 | TCM_NAMESPACE_BEGIN 15 | 16 | auto bind_metropolis(PyObject* _module) -> void 17 | { 18 | namespace py = pybind11; 19 | auto m = py::module{py::reinterpret_borrow(_module)}; 20 | 21 | std::vector keep_alive; 22 | #define DOC(str) trim(keep_alive, str) 23 | 24 | // Currently, we only access the generator from one thread 25 | m.def( 26 | "manual_seed", 27 | [](uint64_t const seed) { global_random_generator().seed(seed); }, 28 | DOC(R"EOF(Seed the random number generator used by nqs_playground.)EOF"), 29 | py::arg{"seed"}); 30 | 31 | py::class_(m, "MetropolisKernel") 32 | .def(py::init>(), DOC(R"EOF( 33 | Constructs Metropolis transition kernel. 34 | 35 | :param basis: specifies the Hilbert space basis.)EOF")) 36 | .def_property_readonly("basis", &MetropolisKernel::basis, 37 | DOC(R"EOF(Returns the Hilbert space basis.)EOF")) 38 | .def( 39 | "__call__", 40 | [](MetropolisKernel const& self, torch::Tensor x) { 41 | return self(std::move(x)); 42 | }, 43 | py::arg{"x"}.noconvert()); 44 | 45 | py::class_(m, "_ProposalGenerator") 46 | .def(py::init>(), DOC(R"EOF( 47 | :param basis: specifies the Hilbert space basis.)EOF")) 48 | .def_property_readonly("basis", &ProposalGenerator::basis, 49 | DOC(R"EOF(Returns the Hilbert space basis.)EOF")) 50 | .def( 51 | "__call__", 52 | [](ProposalGenerator const& self, torch::Tensor x) { 53 | return self(std::move(x)); 54 | }, 55 | py::arg{"x"}.noconvert()); 56 | 57 | m.def("tabu_process", &tabu_process); 58 | #if 0 59 | m.def("zanella_next_state_index", &zanella_next_state_index, DOC(R"EOF( 60 | 61 | )EOF"), 62 | py::arg{"jump_rates"}.noconvert(), py::arg{"counts"}.noconvert(), 63 | py::arg{"out"}.noconvert() = py::none()); 64 | #endif 65 | 66 | m.def( 67 | "zanella_next_state_index", 68 | [](torch::Tensor jump_rates, torch::Tensor jump_rates_sum, 69 | std::vector const& counts, py::object device) { 70 | return zanella_next_state_index( 71 | std::move(jump_rates), std::move(jump_rates_sum), counts, 72 | torch::python::detail::py_object_to_device(device)); 73 | }, 74 | py::arg{"jump_rates"}.noconvert(), 75 | py::arg{"jump_rates_sum"}.noconvert(), py::arg{"counts"}, 76 | py::arg{"device"}.noconvert()); 77 | 78 | m.def("zanella_jump_rates", &zanella_jump_rates, DOC(R"EOF( 79 | 80 | )EOF"), 81 | py::arg{"current_log_prob"}.noconvert(), 82 | py::arg{"possible_log_prob"}.noconvert(), py::arg{"counts"}); 83 | 84 | m.def("zanella_waiting_time", &zanella_waiting_time, DOC(R"EOF( 85 | Calculate waiting time for current state. 86 | 87 | :param rates: Λs 88 | :param out: If specified, the result is stored into it. 89 | :return:)EOF"), 90 | py::arg{"rates"}.noconvert(), py::arg{"out"}.noconvert() = py::none(), 91 | py::call_guard()); 92 | 93 | m.def( 94 | "zanella_choose_samples", 95 | [](torch::Tensor weights, int64_t number_samples, double time_step, 96 | py::object device) { 97 | return zanella_choose_samples( 98 | std::move(weights), number_samples, time_step, 99 | torch::python::detail::py_object_to_device(device)); 100 | }, 101 | DOC(R"EOF( 102 | )EOF"), 103 | py::arg{"weights"}.noconvert(), py::arg{"number_samples"}, 104 | py::arg{"time_step"}.noconvert(), py::arg{"device"}.noconvert()); 105 | 106 | #undef DOC 107 | } 108 | 109 | TCM_NAMESPACE_END 110 | -------------------------------------------------------------------------------- /cbits/python/bind_metropolis.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_metropolis(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_operator.cpp: -------------------------------------------------------------------------------- 1 | #include "bind_operator.hpp" 2 | #include "../operator.hpp" 3 | #include "../trim.hpp" 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | TCM_NAMESPACE_BEGIN 13 | 14 | namespace { 15 | template auto make_call_function() 16 | { 17 | namespace py = pybind11; 18 | return [](Operator const& self, py::array_t x, 19 | std::optional> y) { 20 | TCM_CHECK( 21 | x.ndim() == 1, std::domain_error, 22 | fmt::format("x has invalid shape: [{}]; expected a vector", 23 | fmt::join(x.shape(), x.shape() + x.ndim(), ", "))); 24 | auto src = 25 | gsl::span{x.data(), static_cast(x.shape()[0])}; 26 | auto out = [&y, n = src.size()]() { 27 | if (y.has_value()) { 28 | TCM_CHECK( 29 | y->ndim() == 1, std::domain_error, 30 | fmt::format( 31 | "y has invalid shape: [{}]; expected a vector", 32 | fmt::join(y->shape(), y->shape() + y->ndim(), ", "))); 33 | return std::move(*y); 34 | } 35 | return py::array_t{n}; 36 | }(); 37 | auto dst = gsl::span{out.mutable_data(), 38 | static_cast(out.shape()[0])}; 39 | self(src, dst); 40 | return out; 41 | }; 42 | } 43 | } // namespace 44 | 45 | auto bind_operator(PyObject* _module) -> void 46 | { 47 | namespace py = pybind11; 48 | auto m = py::module{py::reinterpret_borrow(_module)}; 49 | 50 | py::options options; 51 | options.disable_function_signatures(); 52 | 53 | std::vector keep_alive; 54 | #define DOC(str) trim(keep_alive, str) 55 | 56 | py::class_(m, "Interaction") 57 | .def(py::init, 4>, 58 | std::vector>()) 59 | .def_property_readonly("is_real", &Interaction::is_real) 60 | .def_property_readonly("max_index", &Interaction::max_index) 61 | .def_property_readonly( 62 | "matrix", 63 | [](Interaction const& self) { 64 | return py::array_t{ 65 | /*shape=*/{4, 4}, /*data=*/self.matrix()}; 66 | }) 67 | .def_property_readonly("edges", 68 | [](Interaction const& self) { 69 | using std::begin, std::end; 70 | auto const edges = self.edges(); 71 | return std::vector{ 72 | begin(edges), end(edges)}; 73 | }) 74 | .def("diag", [](Interaction const& self, bits512 const& spin) { 75 | return self.diag(spin); 76 | }); 77 | 78 | #if 1 79 | py::class_>(m, "Operator") 80 | .def(py::init, 81 | std::shared_ptr>()) 82 | .def_property_readonly( 83 | "interactions", 84 | [](Operator const& self) { 85 | using std::begin, std::end; 86 | auto is = self.interactions(); 87 | return std::vector{begin(is), end(is)}; 88 | }, 89 | DOC(R"EOF(Returns the list of 'Interaction's)EOF")) 90 | .def_property_readonly("basis", &Operator::basis, DOC(R"EOF( 91 | Returns the Hilbert space basis on which the Operator is defined.)EOF")) 92 | .def_property_readonly("is_real", &Operator::is_real, DOC(R"EOF( 93 | Returns whether the Operator is real.)EOF")) 94 | .def("__call__", make_call_function(), py::arg{"x"}.noconvert(), 95 | py::arg{"y"}.noconvert() = py::none()) 96 | .def("__call__", make_call_function(), py::arg{"x"}.noconvert(), 97 | py::arg{"y"}.noconvert() = py::none()) 98 | .def("__call__", make_call_function>(), 99 | py::arg{"x"}.noconvert(), py::arg{"y"}.noconvert() = py::none()) 100 | .def("__call__", make_call_function>(), 101 | py::arg{"x"}.noconvert(), py::arg{"y"}.noconvert() = py::none()) 102 | .def("to_csr", [](Operator const& self) { 103 | auto [values, indices] = self._to_sparse(); 104 | auto const n = values.size(0); 105 | auto const csr_matrix = 106 | py::module::import("scipy.sparse").attr("csr_matrix"); 107 | auto data = py::cast(values).attr("numpy")(); 108 | auto row_indices = 109 | py::cast(torch::narrow(indices, /*dim=*/1, /*start=*/0, 110 | /*length=*/1)) 111 | .attr("squeeze")() 112 | .attr("numpy")(); 113 | auto col_indices = 114 | py::cast(torch::narrow(indices, /*dim=*/1, /*start=*/1, 115 | /*length=*/1)) 116 | .attr("squeeze")() 117 | .attr("numpy")(); 118 | auto const* basis = 119 | static_cast(self.basis().get()); 120 | return csr_matrix( 121 | std::make_tuple(data, 122 | std::make_tuple(row_indices, col_indices)), 123 | std::make_tuple(basis->number_states(), 124 | basis->number_states())); 125 | }); 126 | #endif 127 | 128 | #undef DOC 129 | } 130 | 131 | TCM_NAMESPACE_END 132 | -------------------------------------------------------------------------------- /cbits/python/bind_operator.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_operator(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_polynomial.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_polynomial(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_polynomial_state.cpp: -------------------------------------------------------------------------------- 1 | #include "bind_polynomial_state.hpp" 2 | #include "../polynomial_state.hpp" 3 | #include "../simple_accumulator.hpp" 4 | #include "../trim.hpp" 5 | #include "pybind11_helpers.hpp" 6 | 7 | #include 8 | #include 9 | 10 | TCM_NAMESPACE_BEGIN 11 | 12 | // namespace { 13 | // auto make_forward_function(torch::jit::script::Method method, 14 | // std::optional number_spins) -> v2::ForwardT 15 | // { 16 | // static_cast(number_spins); 17 | // // auto const n = 18 | // // number_spins.has_value() ? static_cast(*number_spins) : -1L; 19 | // return [f = std::move(method)](auto x) mutable { 20 | // // if (x.scalar_type == torch::kInt64) { x = unpack(x, n); } 21 | // return f({std::move(x)}).toTensor(); 22 | // }; 23 | // } 24 | // 25 | // auto make_forward_function(pybind11::object method) -> v2::ForwardT {} 26 | // } // namespace 27 | 28 | auto bind_polynomial_state(PyObject* _module) -> void 29 | { 30 | namespace py = pybind11; 31 | auto m = py::module{py::reinterpret_borrow(_module)}; 32 | 33 | std::vector keep_alive; 34 | #define DOC(str) trim(keep_alive, str) 35 | 36 | m.def( 37 | "apply", 38 | [](torch::Tensor spins, Heisenberg const& hamiltonian, 39 | v2::ForwardT forward, uint32_t const batch_size) { 40 | return apply(std::move(spins), hamiltonian, std::move(forward), 41 | batch_size); 42 | }, 43 | py::call_guard()); 44 | 45 | m.def("apply", [](torch::Tensor spins, Polynomial& polynomial, 46 | v2::ForwardT forward, uint32_t const batch_size) { 47 | return apply(std::move(spins), polynomial, std::move(forward), 48 | batch_size); 49 | }); 50 | m.def("apply_new", [](torch::Tensor spins, Polynomial& polynomial, 51 | v2::ForwardT psi, uint32_t batch_size, 52 | int32_t num_threads) { 53 | return apply(std::move(spins), std::ref(polynomial), std::move(psi), 54 | polynomial.max_states(), batch_size, num_threads); 55 | }); 56 | 57 | m.def( 58 | "diag", 59 | [](torch::Tensor spins, Heisenberg const& hamiltonian) { 60 | return diag(std::move(spins), hamiltonian); 61 | }, 62 | DOC(R"EOF( 63 | Computes diagonal elements ``⟨s|H|s⟩``. 64 | 65 | :param spins: a tensor of packed spin configurations for which to 66 | compute diagonal matrix elements. 67 | :param hamiltonian: operator which matrix elements to compute.)EOF"), 68 | py::arg{"spins"}, py::arg{"hamiltonian"}); 69 | 70 | #undef DOC 71 | } 72 | 73 | TCM_NAMESPACE_END 74 | -------------------------------------------------------------------------------- /cbits/python/bind_polynomial_state.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_polynomial_state(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_quantum_state.cpp: -------------------------------------------------------------------------------- 1 | #include "bind_heisenberg.hpp" 2 | #include "../heisenberg.hpp" 3 | #include "../trim.hpp" 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | TCM_NAMESPACE_BEGIN 12 | 13 | template auto make_call_function() 14 | { 15 | namespace py = pybind11; 16 | return [](Heisenberg const& self, py::array_t x, 17 | std::optional> y) { 18 | TCM_CHECK(x.ndim() == 1, std::domain_error, 19 | fmt::format("x has invalid shape: [{}]; expected a vector", 20 | fmt::join(x.shape(), x.shape() + x.ndim()))); 21 | auto src = 22 | gsl::span{x.data(), static_cast(x.shape()[0])}; 23 | auto out = y.value_or(py::array_t{src.size()}); 24 | if (y.has_value()) { 25 | TCM_CHECK( 26 | y->ndim() == 1, std::domain_error, 27 | fmt::format("y has invalid shape: [{}]; expected a vector", 28 | fmt::join(y->shape(), y->shape() + y->ndim()))); 29 | } 30 | auto dst = gsl::span{out.mutable_data(), 31 | static_cast(out.shape()[0])}; 32 | self(src, dst); 33 | return out; 34 | }; 35 | } 36 | 37 | auto bind_heisenberg(PyObject* _module) -> void 38 | { 39 | namespace py = pybind11; 40 | auto m = py::module{py::reinterpret_borrow(_module)}; 41 | 42 | std::vector keep_alive; 43 | #define DOC(str) trim(keep_alive, str) 44 | 45 | py::class_>(m, "Heisenberg") 46 | .def( 47 | py::init>(), 48 | DOC(R"EOF( 49 | Constructs the Heisenberg Hamiltonian. 50 | 51 | :param specs: is a list of couplings. Each element is a tuple 52 | ``(cᵢⱼ, i, j)``. This defines a term ``cᵢⱼ·σᵢ⊗σⱼ``. 53 | :param basis: specifies the Hilbert space basis on which the 54 | Hamiltonian is defined.)EOF")) 55 | .def_property_readonly( 56 | "edges", 57 | [](Heisenberg const& self) { 58 | auto edges = self.edges(); 59 | return std::vector{edges.begin(), 60 | edges.end()}; 61 | }, 62 | DOC(R"EOF(Returns the list of couplings)EOF")) 63 | .def_property_readonly("basis", &Heisenberg::basis, DOC(R"EOF( 64 | Returns the Hilbert space basis on which the Hamiltonian is defined.)EOF")) 65 | .def_property_readonly("is_real", &Heisenberg::is_real, DOC(R"EOF( 66 | Returns whether the Hamiltonian is real, i.e. whether all couplings 67 | ``{cᵢⱼ}`` are real.)EOF")) 68 | .def("__call__", make_call_function(), py::arg{"x"}.noconvert(), 69 | py::arg{"y"}.noconvert() = py::none()) 70 | .def("__call__", make_call_function(), py::arg{"x"}.noconvert(), 71 | py::arg{"y"}.noconvert() = py::none()) 72 | .def("__call__", make_call_function>(), 73 | py::arg{"x"}.noconvert(), py::arg{"y"}.noconvert() = py::none()) 74 | .def("__call__", make_call_function>(), 75 | py::arg{"x"}.noconvert(), py::arg{"y"}.noconvert() = py::none()); 76 | 77 | #undef DOC 78 | } 79 | 80 | TCM_NAMESPACE_END 81 | -------------------------------------------------------------------------------- /cbits/python/bind_quantum_state.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_heisenberg(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_spin_basis.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_spin_basis(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_symmetry.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_symmetry(PyObject* module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/bind_v2.cpp: -------------------------------------------------------------------------------- 1 | #include "bind_v2.hpp" 2 | #include "../metropolis.hpp" 3 | #include "../trim.hpp" 4 | #include "../v2/accumulator.hpp" 5 | #include "../v2/zanella.hpp" 6 | #include "../wrappers.hpp" 7 | #include "pybind11_helpers.hpp" 8 | 9 | namespace py = pybind11; 10 | 11 | TCM_NAMESPACE_BEGIN 12 | 13 | auto bind_v2(PyObject* _module) -> void 14 | { 15 | auto m = py::module{py::reinterpret_borrow(_module)}; 16 | 17 | std::vector keep_alive; 18 | #define DOC(str) trim(keep_alive, str) 19 | 20 | m.def("is_operator_real", [](ls_operator const& op) { return ls_operator_is_real(&op); }); 21 | 22 | m.def( 23 | "log_apply", 24 | [](torch::Tensor spins, ls_operator const& _op, v2::ForwardT psi, uint64_t batch_size) { 25 | auto [op, max_required_size] = view_as_operator(_op); 26 | return v2::apply(std::move(spins), std::move(op), std::move(psi), max_required_size, 27 | batch_size); 28 | }, 29 | py::call_guard()); 30 | 31 | // Currently, we only access the generator from one thread 32 | m.def( 33 | "manual_seed", [](uint64_t const seed) { global_random_generator().seed(seed); }, 34 | DOC(R"EOF(Seed the random number generator used by nqs_playground.)EOF"), py::arg{"seed"}); 35 | 36 | m.def( 37 | "random_spin", [](ls_spin_basis const& basis) { return random_spin(basis); }, 38 | py::arg{"basis"}.noconvert()); 39 | 40 | py::class_(m, "MetropolisGenerator") 41 | .def(py::init(), py::arg{"basis"}.noconvert(), DOC(R"EOF( 42 | :param basis: specifies the Hilbert space basis.)EOF")) 43 | .def( 44 | "__call__", 45 | [](MetropolisGenerator const& self, torch::Tensor x, py::object dtype) { 46 | return self(std::move(x), torch::python::detail::py_object_to_dtype(dtype)); 47 | }, 48 | py::arg{"x"}.noconvert(), py::arg{"dtype"}.noconvert()); 49 | 50 | py::class_(m, "ZanellaGenerator") 51 | .def(py::init(), py::arg{"basis"}.noconvert(), DOC(R"EOF( 52 | :param basis: specifies the Hilbert space basis.)EOF")) 53 | .def( 54 | "__call__", 55 | [](ZanellaGenerator const& self, torch::Tensor x) { return self(std::move(x)); }, 56 | py::arg{"x"}.noconvert()); 57 | 58 | m.def( 59 | "zanella_choose_samples", 60 | [](torch::Tensor weights, int64_t const number_samples, double const time_step, 61 | py::object device) { 62 | return zanella_choose_samples(std::move(weights), number_samples, time_step, 63 | torch::python::detail::py_object_to_device(device)); 64 | }, 65 | DOC(R"EOF( 66 | )EOF"), 67 | py::arg{"weights"}.noconvert(), py::arg{"number_samples"}, py::arg{"time_step"}.noconvert(), 68 | py::arg{"device"}.noconvert()); 69 | 70 | #undef DOC 71 | } 72 | 73 | TCM_NAMESPACE_END 74 | -------------------------------------------------------------------------------- /cbits/python/bind_v2.hpp: -------------------------------------------------------------------------------- 1 | #include "../config.hpp" 2 | #include 3 | 4 | TCM_NAMESPACE_BEGIN 5 | 6 | auto bind_v2(PyObject* _module) -> void; 7 | 8 | TCM_NAMESPACE_END 9 | -------------------------------------------------------------------------------- /cbits/python/trim.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #include "trim.hpp" 30 | 31 | TCM_NAMESPACE_BEGIN 32 | 33 | namespace detail { 34 | template 35 | constexpr auto for_each_line(std::string_view str, First first_fn, 36 | Middle middle_fn, Last last_fn) -> void 37 | { 38 | auto prev = std::string_view::size_type{0}; 39 | auto pos = str.find_first_of('\n'); 40 | if (pos != std::string_view::npos) { 41 | first_fn(str.substr(prev, pos - prev)); 42 | prev = pos + 1; 43 | while ((pos = str.find_first_of('\n', prev)) 44 | != std::string_view::npos) { 45 | middle_fn(str.substr(prev, pos - prev)); 46 | prev = pos + 1; 47 | } 48 | } 49 | last_fn(str.substr(prev)); 50 | } 51 | 52 | constexpr auto excess_indent(std::string_view const str) -> size_t 53 | { 54 | auto max = std::string_view::npos; 55 | auto const update = [&max](auto const line) { 56 | auto const n = line.find_first_not_of(' '); 57 | if (n != 0 && n != std::string_view::npos && n < max) { max = n; } 58 | }; 59 | for_each_line( 60 | str, [](auto) {}, update, update); 61 | return max != std::string_view::npos ? max : 0; 62 | } 63 | } // namespace detail 64 | 65 | TCM_EXPORT auto trim(std::vector& keep_alive, std::string_view raw) 66 | -> char const* 67 | { 68 | auto const n = detail::excess_indent(raw); 69 | if (n == 0) { return raw.data(); } 70 | 71 | std::string out; 72 | out.reserve(raw.size()); 73 | detail::for_each_line( 74 | raw, 75 | [&out](auto const line) { 76 | if (!line.empty()) { 77 | out.append(line.data(), line.size()); 78 | out.push_back('\n'); 79 | } 80 | }, 81 | [&out, n](auto line) { 82 | line.remove_prefix(std::min(n, line.size())); 83 | out.append(line.data(), line.size()); 84 | out.push_back('\n'); 85 | }, 86 | [&out, n](auto line) { 87 | line.remove_prefix(std::min(n, line.size())); 88 | out.append(line.data(), line.size()); 89 | }); 90 | keep_alive.push_back(std::move(out)); 91 | return keep_alive.back().c_str(); 92 | } 93 | 94 | TCM_NAMESPACE_END 95 | -------------------------------------------------------------------------------- /cbits/python/trim.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include "../common/config.hpp" 32 | #include 33 | #include 34 | #include 35 | 36 | TCM_NAMESPACE_BEGIN 37 | 38 | /// Removes excessive indents from the doc string \p raw. 39 | /// 40 | /// pybind11 needs `char const*` for doc strings, so we keep a vector of strings 41 | /// which need to remain alive before pybind11 converts them to Python strings 42 | /// (copying the data). 43 | auto trim(std::vector& keep_alive, std::string_view raw) -> char const*; 44 | 45 | TCM_NAMESPACE_END 46 | -------------------------------------------------------------------------------- /conda-build.yml: -------------------------------------------------------------------------------- 1 | name: nqs_conda_build 2 | channels: 3 | - defaults 4 | dependencies: 5 | - python=3.8 # Anaconda client fails with Python 3.9, see https://github.com/Anaconda-Platform/anaconda-client/issues/555 6 | - anaconda-client 7 | - conda-build 8 | - conda-verify 9 | -------------------------------------------------------------------------------- /conda-cpu.yml: -------------------------------------------------------------------------------- 1 | name: nqs_devel_cpu 2 | channels: 3 | - twesterhout 4 | - pytorch 5 | - defaults 6 | - conda-forge 7 | dependencies: 8 | - python=3.8 9 | - pip 10 | - black # For formatting code 11 | # Dependencies for Python code 12 | - lattice-symmetries 13 | - cpuonly 14 | - pytorch=1.8.1 15 | - numpy 16 | - scipy 17 | - h5py 18 | - pyyaml 19 | - tensorboard 20 | - pip: 21 | - loguru 22 | - pynvim # Python support for Neovim 23 | - py-spy # For performance analysis 24 | # For running experiments 25 | - pyyaml 26 | - h5py 27 | # Stuff to compile the package locally 28 | - gcc_linux-64 29 | - gxx_linux-64 30 | - cmake 31 | - ninja 32 | -------------------------------------------------------------------------------- /conda-devel.yml: -------------------------------------------------------------------------------- 1 | name: nqs_playground_devel 2 | channels: 3 | - defaults 4 | - pytorch 5 | - twesterhout 6 | dependencies: 7 | - python=3.8 8 | - pip 9 | - black # For formatting code 10 | # Dependencies for Python code 11 | - lattice-symmetries 12 | - numba 13 | - numpy 14 | - scipy 15 | - pip: 16 | - pynvim # Python support for Neovim 17 | - loguru 18 | - py-spy # For performance analysis 19 | - pytorch=1.7 20 | - cudatoolkit 21 | - tensorboard 22 | # Tools for generating Conda packages 23 | - anaconda-client 24 | - conda-build 25 | - conda-verify 26 | # Stuff to compile the package locally 27 | - gcc_linux-64 28 | - gxx_linux-64 29 | - cmake 30 | - ninja 31 | -------------------------------------------------------------------------------- /conda-gpu.yml: -------------------------------------------------------------------------------- 1 | name: nqs_devel_gpu 2 | channels: 3 | - pytorch 4 | - twesterhout 5 | - defaults 6 | dependencies: 7 | - python 8 | - pip 9 | - pip: 10 | - pynvim # Python support for Neovim 11 | - py-spy # For performance analysis 12 | - black # For formatting code 13 | - loguru # Logging 14 | - pytorch=1.9.1 15 | - cudatoolkit=10.2 16 | - numpy>=1.19 17 | - pyyaml 18 | - h5py 19 | # Stuff to compile extensions 20 | - gcc_linux-64 21 | - gxx_linux-64 22 | - ninja 23 | # Tools for generating Conda packages 24 | - anaconda-client 25 | - conda-build 26 | - conda-verify 27 | # Our own packages 28 | - lattice-symmetries 29 | - unpack-bits 30 | -------------------------------------------------------------------------------- /conda/cpu/build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -ex 4 | 5 | export CMAKE_LIBRARY_PATH=$PREFIX/lib # :$CMAKE_LIBRARY_PATH 6 | export CMAKE_PREFIX_PATH=$PREFIX:$SP_DIR/torch # :$CMAKE_PREFIX_PATH 7 | 8 | which cmake 9 | cmake --version 10 | 11 | mkdir -p build 12 | pushd build 13 | rm -rf * 14 | cmake -GNinja \ 15 | $CMAKE_ARGS \ 16 | -DCMAKE_CXX_FLAGS="$CXXFLAGS -march=nehalem" \ 17 | -DCMAKE_C_FLAGS="$CFLAGS -march=nehalem" \ 18 | -DCMAKE_BUILD_TYPE=Release .. 19 | cmake --build . --target install 20 | popd 21 | 22 | find nqs_playground/ -name "*.so*" -maxdepth 1 -type f | while read sofile; do 23 | echo "Setting rpath of $sofile to " '$ORIGIN:$ORIGIN/../torch/lib' 24 | patchelf --set-rpath '$ORIGIN:$ORIGIN/../torch/lib' --force-rpath $sofile 25 | patchelf --print-rpath $sofile 26 | done 27 | 28 | $PYTHON setup.py install 29 | -------------------------------------------------------------------------------- /conda/cpu/conda_build_config.yaml: -------------------------------------------------------------------------------- 1 | python: 2 | - 3.8 3 | pytorch: 4 | - 1.7.1 5 | ls: 6 | - 0.5.0 7 | -------------------------------------------------------------------------------- /conda/cpu/meta.yaml: -------------------------------------------------------------------------------- 1 | {% set data = load_setup_py_data() %} 2 | 3 | package: 4 | name: nqs-playground 5 | version: {{ data.get('version') }} 6 | 7 | source: 8 | path: ../../ 9 | 10 | build: 11 | number: 1 12 | string: py{{ python }}_torch{{ pytorch }}_ls{{ ls }}_cpu 13 | 14 | requirements: 15 | build: 16 | - {{ compiler('c') }} 17 | - {{ compiler('cxx') }} 18 | 19 | host: 20 | - python {{ python }} 21 | - cmake 22 | - ninja 23 | - cpuonly 24 | - pytorch {{ pytorch }} 25 | - lattice-symmetries {{ ls }} 26 | 27 | run: 28 | - _openmp_mutex 29 | - python 30 | - numpy 31 | - scipy 32 | - psutil 33 | - loguru 34 | - tensorboard >=2.0 35 | - {{ pin_compatible('pytorch', min_pin='x.x.x', max_pin='x.x.x') }} 36 | - {{ pin_compatible('lattice-symmetries', min_pin='x.x.x', max_pin='x.x.x') }} 37 | -------------------------------------------------------------------------------- /conda/gpu/build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -ex 4 | 5 | export CMAKE_LIBRARY_PATH=$PREFIX/lib:$CMAKE_LIBRARY_PATH 6 | export CMAKE_PREFIX_PATH=$PREFIX:$SP_DIR/torch:$CMAKE_PREFIX_PATH 7 | 8 | mkdir -p build 9 | pushd build 10 | rm -rf * 11 | cmake -GNinja \ 12 | -DCMAKE_CUDA_FLAGS="-cudart shared --compiler-options -march=nehalem" \ 13 | -DCMAKE_CXX_FLAGS="$CXXFLAGS -march=nehalem" \ 14 | -DCMAKE_C_FLAGS="$CFLAGS -march=nehalem" \ 15 | -DCMAKE_BUILD_TYPE=Release .. 16 | cmake --build . --target install 17 | popd 18 | 19 | find nqs_playground/ -name "*.so*" -maxdepth 1 -type f | while read sofile; do 20 | echo "Setting rpath of $sofile to " '$ORIGIN:$ORIGIN/../torch/lib' 21 | patchelf --set-rpath '$ORIGIN:$ORIGIN/../torch/lib' --force-rpath $sofile 22 | patchelf --print-rpath $sofile 23 | done 24 | 25 | $PYTHON setup.py install 26 | -------------------------------------------------------------------------------- /conda/gpu/conda_build_config.yaml: -------------------------------------------------------------------------------- 1 | python: 2 | - 3.6 3 | - 3.7 4 | - 3.8 5 | - 3.9 6 | cudatoolkit: 7 | - 11.1 8 | torch: 9 | - 1.8.1 # - 1.7.1 10 | ls: 11 | - 0.7.0 12 | -------------------------------------------------------------------------------- /conda/gpu/meta.yaml: -------------------------------------------------------------------------------- 1 | {% set data = load_setup_py_data() %} 2 | 3 | package: 4 | name: nqs-playground 5 | version: {{ data.get('version') }} 6 | 7 | source: 8 | path: ../../ 9 | 10 | build: 11 | string: py{{ python }}_cuda{{ cudatoolkit }}_torch{{ torch }}_ls{{ ls }}_1 12 | ignore_run_exports: 13 | - cudnn 14 | 15 | requirements: 16 | build: 17 | - {{ compiler('c') }} 18 | - {{ compiler('cxx') }} 19 | 20 | host: 21 | - python {{ python }} 22 | - cmake 23 | - ninja 24 | - cudatoolkit {{ cudatoolkit }} 25 | - pytorch {{ torch }} 26 | - cudnn >=7.6.0,<8.1.0 27 | - lattice-symmetries {{ ls }} 28 | 29 | run: 30 | - _openmp_mutex 31 | - python 32 | - numpy 33 | - scipy 34 | - psutil 35 | - loguru 36 | - tensorboard >=2.0 37 | - cudatoolkit 38 | - {{ pin_compatible('pytorch', min_pin='x.x.x', max_pin='x.x.x') }} 39 | - lattice-symmetries 40 | -------------------------------------------------------------------------------- /devel.yml: -------------------------------------------------------------------------------- 1 | name: nqs_dev 2 | channels: 3 | - defaults 4 | dependencies: 5 | - python 6 | - pip 7 | # Dependencies for Python code 8 | - numpy 9 | - scipy 10 | - typing_extensions 11 | - pip: 12 | - loguru 13 | # We choose to use GPU variant of PyTorch even on machines that don't have a 14 | # GPU. It will all just work! 15 | - cudatoolkit ==10.0.130 16 | - _pytorch_select ==0.2 17 | - pytorch 18 | - tensorboard 19 | # Tools for generating Conda packages 20 | - anaconda-client 21 | - conda-build 22 | - conda-verify 23 | # Stuff to compile the package locally 24 | - gcc_linux-64 25 | - gxx_linux-64 26 | - cmake 27 | - ninja 28 | -------------------------------------------------------------------------------- /distributed_example.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import socket 4 | import torch 5 | import torch.distributed 6 | # import torch.multipro 7 | 8 | def init_slurm(fn, backend="gloo"): 9 | slurm_nodelist = os.environ["SLURM_NODELIST"] 10 | root_node = slurm_nodelist.split(" ")[0].split(",")[0] 11 | if "[" in root_node: 12 | name, numbers = root_node.split("[", maxsplit=1) 13 | number = numbers.split(",", maxsplit=1)[0] 14 | if "-" in number: 15 | number = number.split("-")[0] 16 | number = re.sub("[^0-9]", "", number) 17 | root_node = name + number 18 | os.environ["MASTER_ADDR"] = root_node 19 | 20 | port = os.environ["SLURM_JOB_ID"] 21 | port = port[-4:] # use the last 4 numbers in the job id as the id 22 | port = int(port) + 15000 # all ports should be in the 10k+ range 23 | os.environ["MASTER_PORT"] = str(port) 24 | 25 | rank = int(os.environ["SLURM_PROCID"]) 26 | world_size = int(os.environ["SLURM_NTASKS"]) 27 | torch.distributed.init_process_group(backend, rank=rank, world_size=world_size) 28 | fn() 29 | 30 | def _local_init_process(rank, size, fn, backend): 31 | os.environ["MASTER_ADDR"] = "127.0.0.1" 32 | os.environ["MASTER_PORT"] = "12910" 33 | torch.distributed.init_process_group(backend, rank=rank, world_size=size) 34 | fn() 35 | 36 | def init_local(size, fn, backend="gloo"): 37 | import torch.multiprocessing 38 | 39 | processes = [] 40 | torch.multiprocessing.set_start_method("spawn") 41 | for rank in range(size): 42 | p = torch.multiprocessing.Process(target=_local_init_process, args=(rank, size, fn, backend)) 43 | p.start() 44 | processes.append(p) 45 | 46 | for p in processes: 47 | p.join() 48 | 49 | 50 | def run(): 51 | print( 52 | "Hello from process {} on node {} out of {}" 53 | "".format(torch.distributed.get_rank(), socket.gethostname(), torch.distributed.get_world_size()) 54 | ) 55 | 56 | def main(): 57 | if "SLURM_JOB_ID" in os.environ: 58 | init_slurm(run) 59 | else: 60 | init_local(torch.cuda.device_count(), run) 61 | 62 | if __name__ == '__main__': 63 | main() 64 | -------------------------------------------------------------------------------- /example/1x10/amplitude.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | class Net(torch.nn.Module): 4 | def __init__(self, n: int): 5 | super().__init__() 6 | assert n == 10 7 | self._conv1 = torch.nn.Conv1d( 8 | 1, 6, 10, stride=1, padding=0, dilation=1, groups=1, bias=True 9 | ) 10 | self._conv2 = torch.nn.Conv1d( 11 | 6, 6, 8, stride=1, padding=0, dilation=1, groups=1, bias=True 12 | ) 13 | self._conv3 = torch.nn.Conv1d( 14 | 6, 6, 6, stride=1, padding=0, dilation=1, groups=1, bias=True 15 | ) 16 | self._dense6 = torch.nn.Linear(6, 1, bias=False) 17 | 18 | def forward(self, x): 19 | x = x.view([x.shape[0], 1, 10]) 20 | 21 | x = torch.cat([x, x[:, :, :9]], dim=2) 22 | x = self._conv1(x) 23 | x = torch.nn.functional.relu(x) 24 | 25 | x = torch.cat([x, x[:, :, :7]], dim=2) 26 | x = self._conv2(x) 27 | x = torch.nn.functional.relu(x) 28 | 29 | x = torch.cat([x, x[:, :, :5]], dim=2) 30 | x = self._conv3(x) 31 | x = torch.nn.functional.relu(x) 32 | 33 | x = x.view([x.shape[0], 6, -1]) 34 | x = x.mean(dim=2) 35 | 36 | x = self._dense6(x) 37 | x = torch.exp(x) 38 | return x 39 | 40 | -------------------------------------------------------------------------------- /example/1x10/amplitude_wip.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class Exp(torch.nn.Module): 5 | def __init__(self): 6 | super().__init__() 7 | 8 | def forward(self, x): 9 | return torch.exp(x) 10 | 11 | 12 | Net = lambda n: torch.nn.Sequential( 13 | torch.nn.Linear(n, 16), 14 | torch.nn.ReLU(), 15 | torch.nn.Linear(16, 16), 16 | torch.nn.ReLU(), 17 | torch.nn.Linear(16, 1, bias=False), 18 | Exp(), 19 | ) 20 | -------------------------------------------------------------------------------- /example/1x10/gradient_descend.py: -------------------------------------------------------------------------------- 1 | from loguru import logger 2 | import lattice_symmetries as ls 3 | import numpy as np 4 | import torch 5 | 6 | np.random.seed(52339877) 7 | torch.manual_seed(9218823294) 8 | 9 | # try: 10 | import nqs_playground 11 | import nqs_playground.sgd 12 | import nqs_playground.autoregressive 13 | 14 | # except ImportError: 15 | # # For local development when we only compile the C++ extension, but don't 16 | # # actually install the package using pip 17 | # import os 18 | # import sys 19 | # 20 | # sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..")) 21 | # import nqs_playground 22 | 23 | 24 | def make_amplitude(arch="dense"): 25 | if arch == "dense": 26 | return torch.nn.Sequential( 27 | nqs_playground.Unpack(10), 28 | torch.nn.Linear(10, 10), 29 | torch.nn.ReLU(), 30 | torch.nn.Linear(10, 1, bias=False), 31 | ) 32 | if arch == "nade": 33 | return nqs_playground.autoregressive.NADE(10, 100) 34 | raise ValueError("invalid arch: {}".format(arch)) 35 | 36 | 37 | def make_phase(): 38 | class Phase(torch.nn.Module): 39 | def __init__(self): 40 | super().__init__() 41 | 42 | def forward(self, spins): 43 | return torch.zeros((spins.size(0), 1), dtype=torch.float32, device=spins.device) 44 | 45 | return Phase() 46 | 47 | 48 | def main(): 49 | use_autoregressive = True 50 | basis = ls.SpinBasis(ls.Group([]), number_spins=10, hamming_weight=None) 51 | basis.build() 52 | 53 | # fmt: off 54 | matrix = np.array([[1, 0, 0, 0], 55 | [0, -1, -2, 0], 56 | [0, -2, -1, 0], 57 | [0, 0, 0, 1]]) 58 | # fmt: on 59 | edges = [(i, (i + 1) % basis.number_spins) for i in range(basis.number_spins)] 60 | operator = ls.Operator(basis, [ls.Interaction(matrix, edges)]) 61 | # E, ground_state = ls.diagonalize(operator) 62 | # logger.info("E = {}", E) 63 | # logger.info("v₀ = {}", ground_state[:, 0]) 64 | 65 | amplitude = make_amplitude("nade" if use_autoregressive else "dense") 66 | phase = make_phase() 67 | optimizer = torch.optim.SGD( 68 | list(amplitude.parameters()) + list(phase.parameters()), 69 | lr=1e-1, 70 | momentum=0.9, 71 | weight_decay=1e-4, 72 | ) 73 | # optimizer = torch.optim.Adam(list(amplitude.parameters()) + list(phase.parameters()), lr=1e-2) 74 | options = nqs_playground.sgd.Config( 75 | amplitude=amplitude, 76 | phase=phase, 77 | hamiltonian=operator, 78 | output="1x10.result", 79 | epochs=600, 80 | sampling_options=nqs_playground.SamplingOptions(number_samples=128, number_chains=1), 81 | sampling_mode="autoregressive" if use_autoregressive else "exact", 82 | exact=None, # ground_state[:, 0], 83 | constraints={"hamming_weight": lambda i: 0.1}, 84 | optimizer=optimizer, 85 | inference_batch_size=8192, 86 | ) 87 | runner = nqs_playground.sgd.Runner(options) 88 | for i in range(options.epochs): 89 | runner.step() 90 | if runner._epoch % 200 == 0: 91 | for g in options.optimizer.param_groups: 92 | g["lr"] /= 2 93 | g["momentum"] /= 2 94 | runner.config = options._replace( 95 | sampling_options=nqs_playground.SamplingOptions(number_samples=200000, number_chains=1), 96 | ) 97 | runner.step() 98 | 99 | 100 | if __name__ == "__main__": 101 | main() 102 | -------------------------------------------------------------------------------- /example/1x10/sign.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | Net = lambda n: torch.nn.Sequential( 4 | torch.nn.Linear(n, 10), 5 | torch.nn.ReLU(), 6 | torch.nn.Linear(10, 8), 7 | torch.nn.ReLU(), 8 | torch.nn.Linear(8, 2) 9 | ) 10 | -------------------------------------------------------------------------------- /example/1x10/sign_wip.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | Net = lambda n: torch.nn.Sequential( 4 | torch.nn.Linear(n, 16), 5 | torch.nn.ReLU(), 6 | torch.nn.Linear(16, 16), 7 | torch.nn.ReLU(), 8 | torch.nn.Linear(16, 2) 9 | ) 10 | -------------------------------------------------------------------------------- /example/1x10/stochastic_reconfiguration.py: -------------------------------------------------------------------------------- 1 | from loguru import logger 2 | import lattice_symmetries as ls 3 | import numpy as np 4 | import torch 5 | 6 | np.random.seed(52339877) 7 | torch.manual_seed(9218823294) 8 | 9 | import nqs_playground as nqs 10 | import nqs_playground.sr 11 | import nqs_playground.autoregressive 12 | 13 | 14 | def make_amplitude(arch="dense"): 15 | if arch == "dense": 16 | return torch.nn.Sequential( 17 | nqs_playground.Unpack(10), 18 | torch.nn.Linear(10, 10), 19 | torch.nn.ReLU(), 20 | torch.nn.Linear(10, 1, bias=False), 21 | ) 22 | if arch == "nade": 23 | return nqs_playground.autoregressive.NADE(10, 100) 24 | raise ValueError("invalid arch: {}".format(arch)) 25 | 26 | 27 | def make_phase(): 28 | class Phase(torch.nn.Module): 29 | def __init__(self): 30 | super().__init__() 31 | 32 | def forward(self, spins): 33 | return torch.zeros((spins.size(0), 1), dtype=torch.float32, device=spins.device) 34 | 35 | return Phase() 36 | 37 | 38 | def main(): 39 | use_autoregressive = False 40 | basis = ls.SpinBasis(ls.Group([]), number_spins=10, hamming_weight=None) 41 | basis.build() 42 | 43 | # fmt: off 44 | matrix = np.array([[1, 0, 0, 0], 45 | [0, -1, -2, 0], 46 | [0, -2, -1, 0], 47 | [0, 0, 0, 1]]) 48 | # fmt: on 49 | edges = [(i, (i + 1) % basis.number_spins) for i in range(basis.number_spins)] 50 | operator = ls.Operator(basis, [ls.Interaction(matrix, edges)]) 51 | # E, ground_state = ls.diagonalize(operator) 52 | # logger.info("E = {}", E) 53 | # logger.info("v₀ = {}", ground_state[:, 0]) 54 | 55 | amplitude = make_amplitude("dense") 56 | phase = make_phase() 57 | optimizer = torch.optim.SGD(list(amplitude.parameters()) + list(phase.parameters()), lr=1e-2) 58 | options = nqs_playground.sr.Config( 59 | amplitude=amplitude, 60 | phase=phase, 61 | hamiltonian=operator, 62 | optimizer=optimizer, 63 | epochs=100, 64 | output="1x10.result", 65 | exact=None, 66 | sampling_options=nqs.SamplingOptions(number_samples=2000, number_chains=1), 67 | sampling_mode="exact", 68 | linear_system_kwargs={"rcond": 1e-4}, 69 | inference_batch_size=8192, 70 | ) 71 | runner = nqs_playground.sr.Runner(options) 72 | for i in range(options.epochs): 73 | runner.step() 74 | 75 | 76 | if __name__ == "__main__": 77 | main() 78 | -------------------------------------------------------------------------------- /example/1x10/supervised_amplitude.json: -------------------------------------------------------------------------------- 1 | { 2 | "target": "amplitude", 3 | "dataset": "../../data/1x10/ground_state.pickle", 4 | "model": "amplitude.py", 5 | "output": "runs/amplitude/1", 6 | "optimiser": "lambda m: torch.optim.RMSprop(m.parameters(), lr=5e-4, weight_decay=2e-4)", 7 | "train_fraction": 0.60, 8 | "val_fraction": 0.10, 9 | "train_batch_size": 16, 10 | "patience": 50, 11 | "max_epochs": 1000, 12 | "sampling": "quadratic", 13 | "replacement": true 14 | } 15 | -------------------------------------------------------------------------------- /example/1x10/supervised_sign.json: -------------------------------------------------------------------------------- 1 | { 2 | "dataset": "../../data/1x10/ground_state.pickle", 3 | "model": "sign.py", 4 | "optimiser": "lambda m: torch.optim.RMSprop(m.parameters(), lr=1e-3, weight_decay=1e-4)", 5 | "output": "runs/sign/1", 6 | "target": "sign", 7 | "train_fraction": 0.60, 8 | "val_fraction": 0.10, 9 | "train_batch_size": 16, 10 | "patience": 50, 11 | "max_epochs": 1000, 12 | "sampling": "quadratic", 13 | "replacement": true 14 | } 15 | -------------------------------------------------------------------------------- /example/carleo2017/from_scratch.py: -------------------------------------------------------------------------------- 1 | from loguru import logger 2 | import lattice_symmetries as ls 3 | import numpy as np 4 | import torch 5 | import torch.nn.functional as F 6 | import nqs_playground as nqs 7 | import nqs_playground.sr 8 | 9 | np.random.seed(77) 10 | torch.manual_seed(18) 11 | if torch.cuda.is_available(): 12 | torch.cuda.manual_seed(18) 13 | nqs.manual_seed(87) 14 | 15 | class SpinRBM(torch.nn.Linear): 16 | def __init__(self, number_visible: int, number_hidden: int): 17 | super().__init__(number_visible, number_hidden, bias=True) 18 | 19 | def forward(self, x: torch.Tensor) -> torch.Tensor: 20 | x = nqs.unpack(x, self.in_features).to(self.weight.dtype) 21 | y = F.linear(x, self.weight, self.bias) 22 | y = y - 2.0 * F.softplus(y, beta=-2.0) 23 | return y.sum(dim=1, keepdim=True) 24 | 25 | class Phase(torch.nn.Module): 26 | def __init__(self): 27 | super().__init__() 28 | 29 | def forward(self, spins): 30 | return torch.zeros((spins.size(0), 1), dtype=torch.float32, device=spins.device) 31 | 32 | 33 | def heisenberg1d_model(number_spins: int): 34 | assert number_spins > 0 35 | if number_spins % 2 != 0: 36 | raise NotImplementedError() 37 | sites = np.arange(number_spins) 38 | basis = ls.SpinBasis( 39 | ls.Group([ls.Symmetry(sites[::-1], sector=0)]), 40 | number_spins=number_spins, 41 | hamming_weight=number_spins // 2, 42 | spin_inversion=1, 43 | ) 44 | # fmt: off 45 | matrix = np.array([[1, 0, 0, 0], 46 | [0, -1, -2, 0], 47 | [0, -2, -1, 0], 48 | [0, 0, 0, 1]]) 49 | # fmt: on 50 | edges = [(i, (i + 1) % number_spins) for i in range(number_spins)] 51 | hamiltonian = ls.Operator(basis, [ls.Interaction(matrix, edges)]) 52 | return hamiltonian 53 | 54 | 55 | def analyze_checkpoint(filename: str, number_spins: int = 40, device=None): 56 | hamiltonian = heisenberg1d_model(number_spins) 57 | basis = hamiltonian.basis 58 | if number_spins <= 16: 59 | basis.build() 60 | if device is None: 61 | device = "cuda" if torch.cuda.is_available() else "cpu" 62 | 63 | amplitude = LogAmplitude(number_spins, number_blocks=4, number_channels=16) 64 | # amplitude = torch.nn.Sequential( 65 | # nqs.Unpack(number_spins), 66 | # torch.nn.Linear(number_spins, 13), 67 | # torch.nn.ReLU(inplace=True), 68 | # torch.nn.Linear(13, 1, bias=False) 69 | # ) 70 | amplitude.to(device) 71 | logger.info( 72 | "Amplitude network contains {} parameters", sum(t.numel() for t in amplitude.parameters()) 73 | ) 74 | amplitude.load_state_dict(torch.load(filename)["amplitude"]) 75 | phase = Phase() 76 | phase.to(device) 77 | 78 | sampling_options = nqs.SamplingOptions( 79 | number_samples=1000, number_chains=4, sweep_size=1, number_discarded=10, device=device 80 | ) 81 | states, _, info = nqs.sample_some(amplitude, basis, sampling_options, mode="zanella") 82 | logger.info("Info from the sampler: {}", info) 83 | 84 | combined_state = nqs.combine_amplitude_and_phase(amplitude, phase) 85 | local_energies = nqs.local_values(states, hamiltonian, combined_state, batch_size=8192,) 86 | logger.info("Energy: {}", local_energies.mean(dim=0).cpu()) 87 | logger.info("Energy variance: {}", local_energies.var(dim=0).cpu()) 88 | 89 | 90 | def optimize_with_sr(number_spins: int = 40): 91 | hamiltonian = heisenberg1d_model(number_spins) 92 | basis = hamiltonian.basis 93 | if basis.number_spins <= 16: 94 | basis.build() 95 | 96 | device = "cuda" if torch.cuda.is_available() else "cpu" 97 | amplitude = SpinRBM(number_spins, number_spins).to(device) 98 | # amplitude = torch.jit.script( 99 | # torch.nn.Sequential( 100 | # nqs.Unpack(number_spins), 101 | # torch.nn.Linear(number_spins, number_spins), 102 | # torch.nn.PReLU(), 103 | # torch.nn.Linear(number_spins, number_spins), 104 | # torch.nn.PReLU(), 105 | # torch.nn.Linear(number_spins, 1, bias=False), 106 | # ) 107 | # ).to(device) 108 | # amplitude.load_state_dict( 109 | # torch.load("runs_1x40/05/checkpoints/state_dict_0651_000.pt")["amplitude"] 110 | # ) 111 | # amplitude = LogAmplitude(number_spins, number_blocks=2, number_channels=16).to(device) 112 | num_parameters = sum(t.numel() for t in amplitude.parameters()) 113 | logger.info("Amplitude network contains {} parameters", num_parameters) 114 | phase = Phase().to(device) 115 | # optimizer = torch.optim.SGD( 116 | # amplitude.parameters(), 117 | # lr=1e-4, 118 | # momentum=0.9, 119 | # # weight_decay=1e-4, 120 | # ) 121 | # optimizer = torch.optim.Adam(amplitude.parameters(), lr=1e-3) 122 | parts = [ 123 | # (5000, 1, torch.optim.SGD(amplitude.parameters(), lr=1e-3)), 124 | (250, 1, torch.optim.SGD(amplitude.parameters(), lr=1e-2)), 125 | # (512 * number_spins, 1, torch.optim.SGD(amplitude.parameters(), lr=1e-2)) 126 | ] 127 | options = nqs.sr.Config( 128 | amplitude=amplitude, 129 | phase=phase, 130 | hamiltonian=hamiltonian, 131 | output="runs_1x{}/09".format(number_spins), 132 | epochs=0, 133 | sampling_options=nqs.SamplingOptions( 134 | number_samples=1, 135 | number_chains=4, 136 | sweep_size=5, 137 | number_discarded=10, 138 | mode="zanella", 139 | other={"edges": [(i, (i + 1) % number_spins) for i in range(number_spins)]}, 140 | ), 141 | exact=None, 142 | optimizer=None, 143 | scheduler=None, 144 | linear_system_kwargs={"rcond": 1e-4}, 145 | inference_batch_size=16 * 1024, 146 | ) 147 | runner = nqs.sr.Runner(options) 148 | 149 | for (batch_size, number_inner, optimizer) in parts: 150 | runner.config = runner.config._replace( 151 | sampling_options=runner.config.sampling_options._replace(number_samples=batch_size), 152 | optimizer=optimizer, 153 | epochs=runner.config.epochs + 1000, 154 | ) 155 | runner.run(number_inner=number_inner) 156 | 157 | 158 | if __name__ == "__main__": 159 | optimize_with_sr(40) 160 | -------------------------------------------------------------------------------- /example/carleo2017/models.py: -------------------------------------------------------------------------------- 1 | from loguru import logger 2 | import numpy as np 3 | import torch 4 | from torch.nn import functional as F 5 | from typing import Tuple 6 | import nqs_playground as nqs 7 | from unpack_bits import unpack 8 | import lattice_symmetries as ls 9 | 10 | # np.random.seed(52) 11 | # torch.manual_seed(92) 12 | 13 | 14 | class SpinRBM(torch.nn.Linear): 15 | def __init__(self, number_visible: int, number_hidden: int): 16 | super().__init__(number_visible, number_hidden, bias=True) 17 | 18 | def forward(self, x: torch.Tensor) -> torch.Tensor: 19 | x = unpack(x, self.in_features).to(self.weight.dtype) 20 | y = F.linear(x, self.weight, self.bias) 21 | y = y - 2.0 * F.softplus(y, beta=-2.0) 22 | return y.sum(dim=1, keepdim=True) 23 | 24 | 25 | class Phase(torch.nn.Module): 26 | def __init__(self, dtype=torch.float64): 27 | super().__init__() 28 | self.dtype = dtype 29 | 30 | def forward(self, spins): 31 | return torch.zeros((spins.size(0), 1), dtype=self.dtype, device=spins.device) 32 | 33 | 34 | @torch.no_grad() 35 | def _load_vector(stream, count: int): 36 | out = torch.empty(count, dtype=torch.float64) 37 | for i in range(count): 38 | out[i] = float(stream.readline().strip("()").split(",")[0]) 39 | return out 40 | 41 | 42 | @torch.no_grad() 43 | def load_rbm_weights(filename: str) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 44 | with open(filename, "r") as input: 45 | number_visible = int(input.readline()) 46 | number_hidden = int(input.readline()) 47 | a = _load_vector(input, number_visible) 48 | b = _load_vector(input, number_hidden) 49 | w = _load_vector(input, number_visible * number_hidden).view(number_visible, number_hidden) 50 | return a, b, w 51 | 52 | 53 | @torch.no_grad() 54 | def load_rbm(filename: str) -> torch.nn.Module: 55 | a, b, w = load_rbm_weights(filename) 56 | assert torch.all(a == 0.0) 57 | model = SpinRBM(a.numel(), b.numel()) 58 | model.to(torch.float64) 59 | model.weight.data.copy_(w.t()) 60 | model.bias.data.copy_(b) 61 | return torch.jit.script(model) 62 | 63 | 64 | def heisenberg2d_model(): 65 | L = 10 66 | basis = ls.SpinBasis(ls.Group([]), number_spins=L * L, hamming_weight=L * L // 2) 67 | # fmt: off 68 | matrix = np.array([[1, 0, 0, 0], 69 | [0, -1, -2, 0], 70 | [0, -2, -1, 0], 71 | [0, 0, 0, 1]]) 72 | # fmt: on 73 | edges = [] 74 | for y in range(L): 75 | for x in range(L): 76 | edges.append((x + L * y, (x + 1) % L + L * y)) 77 | edges.append((x + L * y, x + L * ((y + 1) % L))) 78 | hamiltonian = ls.Operator(basis, [ls.Interaction(matrix, edges)]) 79 | return hamiltonian 80 | 81 | 82 | def heisenberg1d_model(number_spins: int = 40): 83 | basis = ls.SpinBasis(ls.Group([]), number_spins=number_spins, hamming_weight=number_spins // 2) 84 | # fmt: off 85 | matrix = np.array([[1, 0, 0, 0], 86 | [0, -1, -2, 0], 87 | [0, -2, -1, 0], 88 | [0, 0, 0, 1]]) 89 | # fmt: on 90 | edges = [(i, (i + 1) % number_spins) for i in range(number_spins)] 91 | hamiltonian = ls.Operator(basis, [ls.Interaction(matrix, edges)]) 92 | return hamiltonian 93 | 94 | 95 | def reproduce_results_of_askar(filename: str = "amplitude_weights_439.pt"): 96 | device = "cuda" if torch.cuda.is_available() else "cpu" 97 | logger.debug("Running the simulation on '{}'...", device) 98 | 99 | hamiltonian = heisenberg2d_model() 100 | basis = hamiltonian.basis 101 | log_ψ = torch.jit.script( 102 | torch.nn.Sequential( 103 | nqs.Unpack(basis.number_spins), 104 | torch.nn.Linear(basis.number_spins, 144), 105 | torch.nn.ReLU(inplace=True), 106 | torch.nn.Linear(144, 1, bias=False), 107 | ) 108 | ) 109 | log_ψ.to(device) 110 | log_ψ.load_state_dict(torch.load(filename)) 111 | 112 | sampling_options = nqs.SamplingOptions( 113 | number_samples=2000, number_chains=4, number_discarded=100, sweep_size=1, device=device 114 | ) 115 | states, _, info = nqs.sample_some(log_ψ, basis, sampling_options, mode="zanella") 116 | logger.info("Info from the sampler: {}", info) 117 | 118 | combined_state = nqs.combine_amplitude_and_phase(log_ψ, Phase(dtype=torch.float32)) 119 | local_energies = nqs.local_values(states, hamiltonian, combined_state,) 120 | logger.info("Energy: {}", local_energies.mean(dim=0).cpu()) 121 | logger.info("Energy variance: {}", local_energies.var(dim=0).cpu()) 122 | 123 | 124 | def reproduce_results_of_carleo2017(filename: str): 125 | logger.debug("Loading RBM weights from '{}'...", filename) 126 | log_ψ = load_rbm(filename) 127 | device = "cuda" if torch.cuda.is_available() else "cpu" 128 | logger.debug("Running Monte Carlo sampling on '{}'...", device) 129 | log_ψ.to(device) 130 | 131 | if "Heisenberg2d" in filename: 132 | hamiltonian = heisenberg2d_model() 133 | elif "Heisenberg1d" in filename: 134 | hamiltonian = heisenberg1d_model(log_ψ.in_features) 135 | else: 136 | raise NotImplementedError() 137 | 138 | basis = hamiltonian.basis 139 | sampling_options = nqs.SamplingOptions( 140 | number_samples=2000, 141 | number_chains=4, 142 | number_discarded=100, 143 | sweep_size=1, 144 | mode="zanella", 145 | device=device, 146 | ) 147 | states, _, _, info = nqs.sample_some(log_ψ, basis, sampling_options) 148 | logger.info("Info from the sampler: {}", info) 149 | 150 | combined_state = nqs.combine_amplitude_and_phase(log_ψ, Phase(dtype=torch.float64)) 151 | local_energies = nqs.local_values(states, hamiltonian, combined_state, batch_size=8192) 152 | local_energies = local_energies.real 153 | logger.info("Energy: {}", local_energies.mean(dim=0).cpu()) 154 | logger.info("Energy variance: {}", local_energies.var(dim=0).cpu()) 155 | 156 | 157 | if __name__ == "__main__": 158 | # reproduce_results_of_askar() 159 | reproduce_results_of_carleo2017("Nqs/Ground/Heisenberg1d_40_1_2.wf") 160 | # reproduce_results_of_carleo2017("Nqs/Ground/Heisenberg2d_100_1_32.wf") 161 | -------------------------------------------------------------------------------- /example/chain/diagonalise.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | from nqs_playground import * 5 | 6 | import nqs_playground 7 | from nqs_playground._C import Interaction, Operator 8 | 9 | USE_SYMMETRIES = False 10 | if USE_SYMMETRIES: 11 | FILENAME_PATTERN = "ground_state_symm_1x{}.npy" 12 | else: 13 | FILENAME_PATTERN = "ground_state_1x{}.npy" 14 | CHECK_STATES = True 15 | 16 | __matrix = lambda data: np.array(data, dtype=np.complex128) 17 | # fmt: off 18 | σ_0 = __matrix([ [1, 0] 19 | , [0, 1] ]) 20 | σ_x = __matrix([ [0, 1] 21 | , [1, 0] ]) 22 | σ_y = __matrix([ [0 , 1j] 23 | , [-1j, 0] ]) 24 | σ_z = __matrix([ [-1, 0] 25 | , [0 , 1] ]) 26 | # fmt: on 27 | σ_p = σ_x + 1j * σ_y 28 | σ_m = σ_x - 1j * σ_y 29 | 30 | 31 | def make_basis(n: int): 32 | return SpinBasis([], number_spins=n, hamming_weight=n // 2) 33 | 34 | 35 | def make_hamiltonian(basis): 36 | if USE_SYMMETRIES: 37 | raise NotImplementedError() 38 | n = basis.number_spins 39 | return Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis) 40 | 41 | 42 | def make_operator(basis): 43 | matrix = lambda data: np.array(data, dtype=np.complex128) 44 | # fmt: off 45 | σ_x = matrix([ [0, 1] 46 | , [1, 0] ]) 47 | σ_y = matrix([ [0 , -1j] 48 | , [1j, 0] ]) 49 | σ_z = matrix([ [1, 0] 50 | , [0, -1] ]) 51 | # fmt: on 52 | op = np.kron(σ_x, σ_x) + np.kron(σ_y, σ_y) + np.kron(σ_z, σ_z) 53 | n = basis.number_spins 54 | edges = [(i, (i + 1) % n) for i in range(n)] 55 | return Operator([Interaction(op, edges)], basis) 56 | 57 | 58 | def check_state(H, y): 59 | Hy = H(y) 60 | E = np.dot(y.conj(), Hy) 61 | close = np.allclose(E * y, Hy) 62 | return close, E 63 | 64 | 65 | def make_exact(hamiltonian): 66 | filename = FILENAME_PATTERN.format(hamiltonian.basis.number_spins) 67 | if not os.path.exists(filename): 68 | print("Information :: Diagonalising...") 69 | energy, ground_state = diagonalise(hamiltonian, k=1) 70 | ground_state = ground_state.squeeze() 71 | np.save(filename, ground_state, allow_pickle=False) 72 | print(energy) 73 | elif CHECK_STATES: 74 | print("Information :: Checking...") 75 | ground_state = np.load(filename) 76 | close, E = check_state(hamiltonian, ground_state) 77 | if close: 78 | print(E) 79 | else: 80 | raise ValueError("'{}' contains an invalid eigenstate".format(filename)) 81 | 82 | 83 | def example0_quspin(): 84 | # Example #0 from QuSpin Documentation 85 | # https://weinbe58.github.io/QuSpin/examples/example0.html 86 | from quspin.operators import hamiltonian 87 | from quspin.basis import spin_basis_1d 88 | 89 | L = 10 # system size 90 | Jxy = np.sqrt(2.0) # xy interaction 91 | Jzz = 1.0 # zz interaction 92 | hz = 1.0 / np.sqrt(3.0) # z external field 93 | 94 | basis = spin_basis_1d(L, pauli=True) 95 | J_zz = [[Jzz, i, i + 1] for i in range(L - 1)] # OBC 96 | J_xy = [[Jxy / 2.0, i, i + 1] for i in range(L - 1)] # OBC 97 | h_z = [[hz, i] for i in range(L)] 98 | static = [["+-", J_xy], ["-+", J_xy], ["zz", J_zz], ["z", h_z]] 99 | dynamic = [] 100 | H_xxz = hamiltonian(static, dynamic, basis=basis, dtype=np.float64) 101 | E, V = H_xxz.eigh() 102 | # Returns the first 10 eigenstates. We reverse the order of states since 103 | # QuSpin stores basis.states in decreasing order and nqs_playground stores 104 | # them in increasing order. 105 | return E[:10], V[::-1, :10] 106 | 107 | 108 | def example0_nqs_playground(): 109 | L = 10 # system size 110 | Jxy = np.sqrt(2.0) # xy interaction 111 | Jzz = 1.0 # zz interaction 112 | hz = 1.0 / np.sqrt(3.0) # z external field 113 | 114 | basis = SpinBasis([], L) 115 | basis.build() 116 | 117 | # XXZ part 118 | xxz_op = Jxy / 2 * (np.kron(σ_p, σ_m) + np.kron(σ_m, σ_p)) + Jzz * np.kron(σ_z, σ_z) 119 | xxz_edges = [(i, i + 1) for i in range(L - 1)] 120 | # Field part. This is a trick. Since nqs_playground only supports 2-local 121 | # operators, we add an identity matrix. 122 | field_op = hz * np.kron(σ_z, σ_0) 123 | # we want the first index to run from 0 to L - 1. The second index is not 124 | # important as long as it's different from the first. 125 | field_edges = [(i, i + 1) for i in range(L - 1)] + [(L - 1, L - 2)] 126 | 127 | hamiltonian = Operator( 128 | [Interaction(xxz_op, xxz_edges), Interaction(field_op, field_edges)], basis 129 | ) 130 | energy, ground_state = diagonalise(hamiltonian, k=10) 131 | return energy, ground_state 132 | 133 | 134 | def example0_compare(): 135 | energy_quspin, states_quspin = example0_quspin() 136 | energy_nqs, states_nqs = example0_nqs_playground() 137 | assert np.allclose(energy_quspin, energy_nqs) 138 | for i in range(states_quspin.shape[1]): 139 | assert np.allclose(states_quspin[:, i], states_nqs[:, i]) or np.allclose( 140 | states_quspin[:, i], -states_nqs[:, i] 141 | ) 142 | 143 | 144 | def run(n): 145 | print("Information :: Creating basis...") 146 | basis = make_basis(n) 147 | print("Information :: Building list of representatives...") 148 | basis.build() 149 | print("Information :: Creating Hamiltonian...") 150 | # operator = make_hamiltonian(basis) 151 | operator = make_operator(basis) 152 | make_exact(operator) 153 | 154 | 155 | def main(): 156 | for n in [10, 12, 14, 16]: 157 | run(n) 158 | 159 | 160 | if __name__ == "__main__": 161 | main() 162 | -------------------------------------------------------------------------------- /example/heisenberg_square/gradient_descend.py: -------------------------------------------------------------------------------- 1 | from loguru import logger 2 | import lattice_symmetries as ls 3 | import numpy as np 4 | import torch 5 | import h5py 6 | import yaml 7 | 8 | np.random.seed(52339877) 9 | torch.manual_seed(9218823294) 10 | 11 | # try: 12 | import nqs_playground 13 | import nqs_playground.sgd 14 | import nqs_playground.autoregressive 15 | 16 | # except ImportError: 17 | # # For local development when we only compile the C++ extension, but don't 18 | # # actually install the package using pip 19 | # import os 20 | # import sys 21 | # 22 | # sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..")) 23 | # import nqs_playground 24 | 25 | 26 | def make_amplitude(arch="dense"): 27 | if arch == "dense": 28 | return torch.nn.Sequential( 29 | nqs_playground.Unpack(36), 30 | torch.nn.Linear(36, 64), 31 | torch.nn.ReLU(), 32 | torch.nn.Linear(64, 1, bias=False), 33 | ) 34 | if arch == "nade": 35 | return nqs_playground.autoregressive.PixelCNN((6, 6)) # 36, 100) 36 | # return nqs_playground.autoregressive.NADE(36, 100) 37 | raise ValueError("invalid arch: {}".format(arch)) 38 | 39 | 40 | def make_phase(): 41 | class Phase(torch.nn.Module): 42 | def __init__(self): 43 | super().__init__() 44 | 45 | def forward(self, spins): 46 | return torch.zeros((spins.size(0), 1), dtype=torch.float32, device=spins.device) 47 | 48 | return Phase() 49 | 50 | 51 | def load_ground_state(filename: str): 52 | logger.info("Loading ground state from '{}'...", filename) 53 | with h5py.File(filename, "r") as f: 54 | ground_state = f["/hamiltonian/eigenvectors"][:, 0] 55 | energy = f["/hamiltonian/eigenvalues"][0] 56 | basis_representatives = f["/basis/representatives"][:] 57 | logger.info("Ground state energy is {}", energy) 58 | return ground_state, energy, basis_representatives 59 | 60 | 61 | def load_basis_and_hamiltonian(filename: str): 62 | import yaml 63 | 64 | logger.info("Loading basis from '{}'...", filename) 65 | with open(filename, "r") as f: 66 | config = yaml.load(f, Loader=yaml.SafeLoader) 67 | basis = ls.SpinBasis.load_from_yaml(config["basis"]) 68 | basis = ls.SpinBasis(ls.Group([]), number_spins=basis.number_spins, hamming_weight=None) 69 | logger.info("Loading Hamiltonian from '{}'...", filename) 70 | hamiltonian = ls.Operator.load_from_yaml(config["hamiltonian"], basis) 71 | return basis, hamiltonian 72 | 73 | 74 | def main(): 75 | use_autoregressive = True 76 | 77 | basis, hamiltonian = load_basis_and_hamiltonian("heisenberg_square_36_positive.yaml") 78 | ground_state, energy, representatives = load_ground_state("data/heisenberg_square_36.h5") 79 | # basis.build(representatives=representatives) 80 | del representatives 81 | 82 | amplitude = make_amplitude("nade" if use_autoregressive else "dense") 83 | phase = make_phase() 84 | optimizer = torch.optim.Adam( 85 | list(amplitude.parameters()) + list(phase.parameters()), 86 | lr=5e-3, 87 | ) 88 | # optimizer = torch.optim.SGD( 89 | # list(amplitude.parameters()) + list(phase.parameters()), 90 | # lr=1e-3, 91 | # momentum=0.9, 92 | # weight_decay=1e-4, 93 | # ) 94 | options = nqs_playground.sgd.Config( 95 | amplitude=amplitude, 96 | phase=phase, 97 | hamiltonian=hamiltonian, 98 | output="data/6x6", 99 | epochs=500, 100 | sampling_options=nqs_playground.SamplingOptions(number_samples=4096, number_chains=1), 101 | sampling_mode="autoregressive" if use_autoregressive else "exact", 102 | exact=None, # ground_state[:, 0], 103 | constraints=dict(), # {"hamming_weight": lambda i: 0.2}, 104 | optimizer=optimizer, 105 | inference_batch_size=8192, 106 | ) 107 | runner = nqs_playground.sgd.Runner(options) 108 | for i in range(options.epochs): 109 | runner.step() 110 | # if runner._epoch % 200 == 0: 111 | # for g in options.optimizer.param_groups: 112 | # g["lr"] /= 2 113 | # g["momentum"] /= 2 114 | # runner.config = options._replace( 115 | # sampling_options=nqs_playground.SamplingOptions(number_samples=200000, number_chains=1), 116 | # ) 117 | # runner.step() 118 | 119 | 120 | if __name__ == "__main__": 121 | main() 122 | -------------------------------------------------------------------------------- /example/heisenberg_square/heisenberg_square_36_positive.yaml: -------------------------------------------------------------------------------- 1 | # Hilbert space dimension: 15804956 2 | # Ground state energy: -97.75758959723598 3 | # Time on Intel(R) Xeon(R) CPU E5-2650 v2 @ 2.60GHz: ~1.5 hours 4 | basis: 5 | number_spins: 36 6 | hamming_weight: 18 7 | spin_inversion: 1 8 | symmetries: 9 | # Initially: 10 | # 11 | # 0, 1, 2, 3, 4, 5, 12 | # 6, 7, 8, 9, 10, 11, 13 | # 12, 13, 14, 15, 16, 17, 14 | # 18, 19, 20, 21, 22, 23, 15 | # 24, 25, 26, 27, 28, 29, 16 | # 30, 31, 32, 33, 34, 35, 17 | # 18 | # Translation along x-direction 19 | - permutation: [1, 2, 3, 4, 5, 0, 20 | 7, 8, 9, 10, 11, 6, 21 | 13, 14, 15, 16, 17, 12, 22 | 19, 20, 21, 22, 23, 18, 23 | 25, 26, 27, 28, 29, 24, 24 | 31, 32, 33, 34, 35, 30] 25 | sector: 0 26 | # Translation along y-direction 27 | - permutation: [6, 7, 8, 9, 10, 11, 28 | 12, 13, 14, 15, 16, 17, 29 | 18, 19, 20, 21, 22, 23, 30 | 24, 25, 26, 27, 28, 29, 31 | 30, 31, 32, 33, 34, 35, 32 | 0, 1, 2, 3, 4, 5] 33 | sector: 0 34 | # 90-degree rotation 35 | - permutation: [5, 11, 17, 23, 29, 35, 36 | 4, 10, 16, 22, 28, 34, 37 | 3, 9, 15, 21, 27, 33, 38 | 2, 8, 14, 20, 26, 32, 39 | 1, 7, 13, 19, 25, 31, 40 | 0, 6, 12, 18, 24, 30] 41 | sector: 0 42 | # Parity along x-direction 43 | - permutation: [5, 4, 3, 2, 1, 0, 44 | 11, 10, 9, 8, 7, 6, 45 | 17, 16, 15, 14, 13, 12, 46 | 23, 22, 21, 20, 19, 18, 47 | 29, 28, 27, 26, 25, 24, 48 | 35, 34, 33, 32, 31, 30] 49 | sector: 0 50 | # Parity along y-direction 51 | - permutation: [30, 31, 32, 33, 34, 35, 52 | 24, 25, 26, 27, 28, 29, 53 | 18, 19, 20, 21, 22, 23, 54 | 12, 13, 14, 15, 16, 17, 55 | 6, 7, 8, 9, 10, 11, 56 | 0, 1, 2, 3, 4, 5] 57 | sector: 0 58 | hamiltonian: 59 | name: "Heisenberg Hamiltonian" 60 | terms: 61 | ##################################################################### 62 | ## IMPORTANT ## 63 | ## ## 64 | ## Here is the only difference between this file and ## 65 | ## heisenberg_square_36.yaml: we flip off-diagonal elements in the ## 66 | ## interaction matrix. This allows us to avoid the sign problem. ## 67 | ##################################################################### 68 | # Nearest neighbours 69 | - matrix: [[1, 0, 0, 0], 70 | [0, -1, -2, 0], 71 | [0, -2, -1, 0], 72 | [0, 0, 0, 1]] 73 | sites: [ 74 | [ 0, 1], [ 0, 6], [ 1, 2], [ 1, 7], [ 2, 3], [ 2, 8], [ 3, 4], [ 3, 9], [ 4, 5], [ 4, 10], [ 5, 0], [ 5, 11], 75 | [ 6, 7], [ 6, 12], [ 7, 8], [ 7, 13], [ 8, 9], [ 8, 14], [ 9, 10], [ 9, 15], [10, 11], [10, 16], [11, 6], [11, 17], 76 | [12, 13], [12, 18], [13, 14], [13, 19], [14, 15], [14, 20], [15, 16], [15, 21], [16, 17], [16, 22], [17, 12], [17, 23], 77 | [18, 19], [18, 24], [19, 20], [19, 25], [20, 21], [20, 26], [21, 22], [21, 27], [22, 23], [22, 28], [23, 18], [23, 29], 78 | [24, 25], [24, 30], [25, 26], [25, 31], [26, 27], [26, 32], [27, 28], [27, 33], [28, 29], [28, 34], [29, 24], [29, 35], 79 | [30, 31], [30, 0], [31, 32], [31, 1], [32, 33], [32, 2], [33, 34], [33, 3], [34, 35], [34, 4], [35, 30], [35, 5] 80 | ] 81 | observables: [] 82 | number_vectors: 1 83 | output: "data/heisenberg_square_36.h5" 84 | max_primme_block_size: 4 85 | max_primme_basis_size: 30 86 | 87 | -------------------------------------------------------------------------------- /example/triangleperiodic/16/amplitude.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | class Net(torch.jit.ScriptModule): 4 | def __init__(self, n: int): 5 | super().__init__() 6 | self._conv1 = torch.nn.Conv2d( 7 | 1, 10, 4, stride=1, padding=0, dilation=1, groups=1, bias=True 8 | ) 9 | self._conv2 = torch.nn.Conv2d( 10 | 10, 10, 3, stride=1, padding=0, dilation=1, groups=1, bias=True 11 | ) 12 | self._conv3 = torch.nn.Conv2d( 13 | 10, 10, 2, stride=1, padding=0, dilation=1, groups=1, bias=True 14 | ) 15 | self._dense6 = torch.nn.Linear(10, 1, bias=False) 16 | 17 | @torch.jit.script_method 18 | def forward(self, x): 19 | x = x.view([x.shape[0], 1, 4, 4]) 20 | 21 | x = torch.cat([x, x[:, :, :3, :]], dim=2) 22 | x = torch.cat([x, x[:, :, :, :3]], dim=3) 23 | x = self._conv1(x) 24 | x = torch.nn.functional.relu(x) 25 | 26 | x = torch.cat([x, x[:, :, :2, :]], dim=2) 27 | x = torch.cat([x, x[:, :, :, :2]], dim=3) 28 | x = self._conv2(x) 29 | x = torch.nn.functional.relu(x) 30 | 31 | x = torch.cat([x, x[:, :, :1, :]], dim=2) 32 | x = torch.cat([x, x[:, :, :, :1]], dim=3) 33 | x = self._conv3(x) 34 | x = torch.nn.functional.relu(x) 35 | 36 | x = x.view([x.shape[0], 10, -1]) 37 | x = x.mean(dim=2) 38 | 39 | x = self._dense6(x) 40 | # x = torch.nn.functional.softplus(x) 41 | x = torch.exp(x) 42 | return x 43 | -------------------------------------------------------------------------------- /example/triangleperiodic/16/sign.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class Net(torch.jit.ScriptModule): 5 | def __init__(self, n: int): 6 | super().__init__() 7 | self._conv1 = torch.nn.Conv2d( 8 | 1, 8, 4, stride=1, padding=0, dilation=1, groups=1, bias=True 9 | ) 10 | self._conv2 = torch.nn.Conv2d( 11 | 8, 6, 3, stride=1, padding=0, dilation=1, groups=1, bias=True 12 | ) 13 | self._conv3 = torch.nn.Conv2d( 14 | 6, 4, 2, stride=1, padding=0, dilation=1, groups=1, bias=True 15 | ) 16 | self._dense6 = torch.nn.Linear(4, 2, bias=True) 17 | 18 | @torch.jit.script_method 19 | def forward(self, x): 20 | x = x.view([x.shape[0], 1, 4, 4]) 21 | 22 | x = torch.cat([x, x[:, :, :3, :]], dim=2) 23 | x = torch.cat([x, x[:, :, :, :3]], dim=3) 24 | x = self._conv1(x) 25 | x = torch.nn.functional.relu(x) 26 | 27 | x = torch.cat([x, x[:, :, :2, :]], dim=2) 28 | x = torch.cat([x, x[:, :, :, :2]], dim=3) 29 | x = self._conv2(x) 30 | x = torch.nn.functional.relu(x) 31 | 32 | x = torch.cat([x, x[:, :, :1, :]], dim=2) 33 | x = torch.cat([x, x[:, :, :, :1]], dim=3) 34 | x = self._conv3(x) 35 | x = torch.nn.functional.relu(x) 36 | 37 | x = x.view([x.shape[0], 4, -1]) 38 | x = x.mean(dim=2) 39 | 40 | x = self._dense6(x) 41 | return x 42 | -------------------------------------------------------------------------------- /nqs_playground/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | __version__ = '1.0.1' 3 | 4 | # NOTE: We need to import PyTorch first, because we don't explicitly link 5 | # against it in C++ code. 6 | # import torch as __torch 7 | 8 | from .core import * 9 | from .sampling import * 10 | # from ._extension import PACKAGE_DIR, lib 11 | from .hamiltonian import * 12 | # from ._jacobian import * 13 | # from ._C import manual_seed, random_spin 14 | 15 | from .runner import * 16 | 17 | # This operator becomes available only after loading _C 18 | # unpack = __torch.ops.tcm.unpack 19 | # hamming_weight = __torch.ops.tcm.hamming_weight 20 | -------------------------------------------------------------------------------- /nqs_playground/_extension.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import torch 4 | from torch.utils.cpp_extension import load as _load 5 | 6 | PACKAGE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | 8 | lib = _load( 9 | "_nqs_playground_cpp", 10 | [os.path.join(PACKAGE_DIR, "cbits", "zanella.cpp"), 11 | os.path.join(PACKAGE_DIR, "cbits", "accumulator.cpp"), 12 | os.path.join(PACKAGE_DIR, "cbits", "init.cpp") 13 | ], 14 | extra_cflags=[ 15 | "-fvisibility=hidden", 16 | "-std=c++17", 17 | "-g", 18 | # "-march=native", 19 | # "-O3", 20 | # "-ftree-vectorize", 21 | "-fopenmp", 22 | "-Wall", 23 | "-Wextra" 24 | ], 25 | extra_include_paths=[os.path.join(sys.prefix, "include")], 26 | extra_ldflags=["-fopenmp", "-L" + os.path.join(sys.prefix, "lib"), "-llattice_symmetries"], 27 | verbose=True 28 | ) 29 | 30 | -------------------------------------------------------------------------------- /nqs_playground/_jacobian.py: -------------------------------------------------------------------------------- 1 | # Copyright Tom Westerhout (c) 2020-2021 2 | # 3 | # All rights reserved. 4 | # 5 | # Redistribution and use in source and binary forms, with or without 6 | # modification, are permitted provided that the following conditions are met: 7 | # 8 | # * Redistributions of source code must retain the above copyright 9 | # notice, this list of conditions and the following disclaimer. 10 | # 11 | # * Redistributions in binary form must reproduce the above 12 | # copyright notice, this list of conditions and the following 13 | # disclaimer in the documentation and/or other materials provided 14 | # with the distribution. 15 | # 16 | # * Neither the name of Tom Westerhout nor the names of other 17 | # contributors may be used to endorse or promote products derived 18 | # from this software without specific prior written permission. 19 | # 20 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 21 | # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 22 | # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 23 | # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 24 | # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 25 | # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 26 | # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 27 | # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 28 | # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 30 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 31 | 32 | __all__ = ["jacobian"] 33 | 34 | 35 | from typing import List, Optional 36 | 37 | import torch 38 | from torch import Tensor 39 | 40 | if torch.has_cuda: 41 | import threading 42 | from torch.nn.parallel.scatter_gather import scatter, gather 43 | from torch.nn.parallel.replicate import replicate 44 | from torch.cuda._utils import _get_device_index 45 | from torch._utils import ExceptionWrapper 46 | 47 | 48 | def jacobian(module: torch.nn.Module, parameters: List[Tensor], inputs: Tensor) -> Tensor: 49 | r"""Given a ``torch.nn.Module`` and a ``torch.Tensor`` of inputs, computes 50 | the Jacobian ∂module(inputs)/∂W where W are module's parameters. 51 | 52 | It is assumed that if ``inputs`` has shape ``(batch_size, in_features)``, 53 | then ``module(inputs)`` has shape ``(batch_size, 1)``. 54 | """ 55 | return jacobian_simple(module, parameters, inputs) 56 | # if inputs.device.type == "cuda": 57 | # return jacobian_cuda(module, inputs) 58 | # elif inputs.device.type == "cpu": 59 | # return jacobian_cpu(module, inputs) 60 | # else: 61 | # raise ValueError( 62 | # "'inputs' tensor resides on an unsupported device: {}; expected either " 63 | # "'cpu' or 'cuda'".format(inputs.device.type) 64 | # ) 65 | 66 | 67 | def jacobian_simple(module: torch.nn.Module, parameters: List[Tensor], inputs: Tensor) -> Tensor: 68 | r"""Trivial implementation of ``jacobian``. It is used to assess 69 | correctness of fancier techniques. 70 | """ 71 | out = inputs.new_empty( 72 | [inputs.size(0), sum(map(torch.numel, parameters))], dtype=parameters[0].dtype 73 | ) 74 | for i in range(inputs.size(0)): 75 | dws = torch.autograd.grad([module(inputs[[i]])], parameters) 76 | torch.cat([dw.flatten() for dw in dws], out=out[i]) 77 | return out 78 | 79 | 80 | # def jacobian_cpu(module: torch.jit.ScriptModule, inputs: Tensor, num_threads: int = -1) -> Tensor: 81 | # r"""Jacobian computation on CPU.""" 82 | # return _jacobian(module._c, inputs, num_threads=num_threads) 83 | 84 | 85 | # def jacobian_cuda( 86 | # module: torch.jit.ScriptModule, 87 | # inputs: Tensor, 88 | # devices: Optional[List[torch.device]] = None, 89 | # output_device: Optional[torch.device] = None, 90 | # parallel: Optional[bool] = True, 91 | # ) -> Tensor: 92 | # r"""Jacobian computation on (multiple) GPUs.""" 93 | # if not parallel: 94 | # return _jacobian(module._c, inputs) 95 | # if devices is None: 96 | # device_ids = list(range(torch.cuda.device_count())) 97 | # else: 98 | # device_ids = list(map(lambda x: _get_device_index(x, True), devices)) 99 | # 100 | # if output_device is None: 101 | # output_device = inputs.device 102 | # inputs = scatter(inputs, device_ids, dim=0) 103 | # replicas = replicate(module, device_ids, detach=False) 104 | # outputs = _parallel_apply_jacobian(replicas, inputs) 105 | # return gather(outputs, output_device, dim=0) 106 | 107 | 108 | # def _parallel_apply_jacobian( 109 | # replicas: List[torch.jit.ScriptModule], inputs: List[Tensor] 110 | # ) -> List[Tensor]: 111 | # results = [None] * len(inputs) 112 | # 113 | # def _worker(i, module, x): 114 | # try: 115 | # results[i] = _jacobian(module._c, x) 116 | # except Exception: 117 | # results[i] = ExceptionWrapper(where="in replica {} on device {}".format(i, x.device)) 118 | # 119 | # threads = [ 120 | # threading.Thread(target=_worker, args=(i, m, x)) 121 | # for i, (m, x) in enumerate(zip(replicas, inputs)) 122 | # ] 123 | # for thread in threads: 124 | # thread.start() 125 | # for thread in threads: 126 | # thread.join() 127 | # 128 | # for result in results: 129 | # if isinstance(result, ExceptionWrapper): 130 | # result.reraise() 131 | # return results 132 | -------------------------------------------------------------------------------- /nqs_playground/cbits/accumulator.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #pragma once 30 | 31 | #include 32 | #include 33 | #include 34 | 35 | namespace tcm { 36 | 37 | using ForwardT = std::function torch::Tensor>; 38 | 39 | auto log_apply(torch::Tensor spins, ls_operator const& op, ForwardT fn, int64_t batch_size) -> torch::Tensor; 40 | 41 | } // namespace tcm 42 | -------------------------------------------------------------------------------- /nqs_playground/cbits/init.cpp: -------------------------------------------------------------------------------- 1 | #include "zanella.hpp" 2 | #include "accumulator.hpp" 3 | #include 4 | #include 5 | 6 | namespace pybind11::detail { 7 | template <> struct type_caster { 8 | private: 9 | ls_spin_basis* payload; 10 | 11 | public: 12 | auto load(handle src, bool /*convert*/) -> bool 13 | { 14 | if (src.is_none()) { return false; } 15 | auto basis_type = module_::import("lattice_symmetries").attr("SpinBasis"); 16 | if (!isinstance(src, basis_type)) { return false; } 17 | payload = 18 | reinterpret_cast(src.attr("_payload").attr("value").cast()); 19 | return true; 20 | } 21 | 22 | operator ls_spin_basis*() { return payload; } 23 | operator ls_spin_basis&() { return *payload; } 24 | 25 | static constexpr auto name = _("lattice_symmetries.SpinBasis"); 26 | template using cast_op_type = pybind11::detail::cast_op_type; 27 | }; 28 | 29 | template <> struct type_caster { 30 | private: 31 | ls_operator* payload; 32 | 33 | public: 34 | auto load(handle src, bool /*convert*/) -> bool 35 | { 36 | if (src.is_none()) { return false; } 37 | auto operator_type = module_::import("lattice_symmetries").attr("Operator"); 38 | if (!isinstance(src, operator_type)) { return false; } 39 | payload = 40 | reinterpret_cast(src.attr("_payload").attr("value").cast()); 41 | return true; 42 | } 43 | 44 | operator ls_operator*() { return payload; } 45 | operator ls_operator&() { return *payload; } 46 | 47 | static constexpr auto name = _("lattice_symmetries.Operator"); 48 | template using cast_op_type = pybind11::detail::cast_op_type; 49 | }; 50 | } // namespace pybind11::detail 51 | 52 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 53 | { 54 | using namespace tcm; 55 | py::class_(m, "ZanellaGenerator") 56 | .def(py::init>>(), 57 | py::arg{"basis"}.noconvert(), py::arg{"edges"}, R"EOF( 58 | :param basis: specifies the Hilbert space basis.)EOF") 59 | .def( 60 | "__call__", 61 | [](ZanellaGenerator const& self, torch::Tensor x) { return self(std::move(x)); }, 62 | py::arg{"x"}.noconvert(), py::call_guard()); 63 | 64 | m.def("log_apply", &log_apply, py::call_guard()); 65 | } 66 | -------------------------------------------------------------------------------- /nqs_playground/cbits/parallel.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | namespace tcm { 14 | 15 | // Copyright (c) 2012 Jakob Progsch, Václav Zeman 16 | // 17 | // This software is provided 'as-is', without any express or implied 18 | // warranty. In no event will the authors be held liable for any damages 19 | // arising from the use of this software. 20 | // 21 | // Permission is granted to anyone to use this software for any purpose, 22 | // including commercial applications, and to alter it and redistribute it 23 | // freely, subject to the following restrictions: 24 | // 25 | // 1. The origin of this software must not be misrepresented; you must not 26 | // claim that you wrote the original software. If you use this software 27 | // in a product, an acknowledgment in the product documentation would be 28 | // appreciated but is not required. 29 | // 30 | // 2. Altered source versions must be plainly marked as such, and must not be 31 | // misrepresented as being the original software. 32 | // 33 | // 3. This notice may not be removed or altered from any source 34 | // distribution. 35 | // 36 | // The following is a small adaptation of https://github.com/progschj/ThreadPool 37 | // for a single worker thread. 38 | class ThreadPool { 39 | public: 40 | ThreadPool() : worker{}, tasks{}, queue_mutex{}, condition{}, stop{false} 41 | { 42 | worker = std::thread{[this] { 43 | for (;;) { 44 | std::function task; 45 | 46 | { 47 | std::unique_lock lock(queue_mutex); 48 | condition.wait(lock, [this] { return stop || !tasks.empty(); }); 49 | if (stop && tasks.empty()) return; 50 | task = std::move(tasks.front()); 51 | tasks.pop(); 52 | } 53 | 54 | task(); 55 | } 56 | }}; 57 | } 58 | 59 | template auto enqueue(F&& f) -> std::future::type> 60 | { 61 | using return_type = typename std::result_of::type; 62 | auto task = std::make_shared>(std::forward(f)); 63 | std::future res = task->get_future(); 64 | { 65 | std::unique_lock lock(queue_mutex); 66 | // don't allow enqueueing after stopping the pool 67 | if (stop) { throw std::runtime_error{"enqueue on stopped ThreadPool"}; } 68 | tasks.emplace([p = std::move(task)]() { (*p)(); }); 69 | } 70 | condition.notify_one(); 71 | return res; 72 | } 73 | 74 | ~ThreadPool() 75 | { 76 | { 77 | std::unique_lock lock(queue_mutex); 78 | stop = true; 79 | } 80 | condition.notify_all(); 81 | worker.join(); 82 | } 83 | 84 | private: 85 | // need to keep track of threads so we can join them 86 | std::thread worker; 87 | // task queue 88 | std::queue> tasks; 89 | // synchronization 90 | std::mutex queue_mutex; 91 | std::condition_variable condition; 92 | bool stop; 93 | }; 94 | 95 | 96 | 97 | template auto sync_task(F&& f) -> std::future::type> 98 | { 99 | using return_type = typename std::result_of::type; 100 | std::packaged_task task{std::forward(f)}; 101 | std::future res = task.get_future(); 102 | task(); 103 | return res; 104 | } 105 | 106 | } // namespace tcm 107 | -------------------------------------------------------------------------------- /nqs_playground/cbits/zanella.hpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, Tom Westerhout 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright notice, this 8 | // list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright notice, 11 | // this list of conditions and the following disclaimer in the documentation 12 | // and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of the copyright holder nor the names of its 15 | // contributors may be used to endorse or promote products derived from 16 | // this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | #include 30 | #include 31 | 32 | namespace tcm { 33 | 34 | class ZanellaGenerator { 35 | private: 36 | ls_spin_basis const* _basis; 37 | std::vector> _edges; 38 | 39 | auto generate(ls_bits512 const& spin, ls_bits512* out) const -> int64_t; 40 | auto project(ls_bits512* spins, int64_t count) const -> int64_t; 41 | 42 | public: 43 | ZanellaGenerator(ls_spin_basis const& basis, std::vector> edges); 44 | ~ZanellaGenerator(); 45 | 46 | auto operator()(torch::Tensor x) const -> std::tuple; 47 | auto max_states() const noexcept -> uint64_t; 48 | }; 49 | 50 | } // namespace tcm 51 | -------------------------------------------------------------------------------- /nqs_playground/distributed.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import socket 4 | import torch 5 | import torch.distributed 6 | from loguru import logger 7 | # import torch.multipro 8 | 9 | def init_slurm(fn, *args, backend="gloo", **kwargs): 10 | slurm_nodelist = os.environ["SLURM_NODELIST"] 11 | root_node = slurm_nodelist.split(" ")[0].split(",")[0] 12 | if "[" in root_node: 13 | name, numbers = root_node.split("[", maxsplit=1) 14 | number = numbers.split(",", maxsplit=1)[0] 15 | if "-" in number: 16 | number = number.split("-")[0] 17 | number = re.sub("[^0-9]", "", number) 18 | root_node = name + number 19 | os.environ["MASTER_ADDR"] = root_node 20 | 21 | port = os.environ["SLURM_JOB_ID"] 22 | port = port[-4:] # use the last 4 numbers in the job id as the id 23 | port = int(port) + 15000 # all ports should be in the 10k+ range 24 | os.environ["MASTER_PORT"] = str(port) 25 | 26 | rank = int(os.environ["SLURM_PROCID"]) 27 | world_size = int(os.environ["SLURM_NTASKS"]) 28 | torch.distributed.init_process_group(backend, rank=rank, world_size=world_size) 29 | fn(*args, ) 30 | 31 | def _local_init_process(rank, size, fn, backend, *args, **kwargs): 32 | os.environ["MASTER_ADDR"] = "127.0.0.1" 33 | os.environ["MASTER_PORT"] = "12910" 34 | torch.distributed.init_process_group(backend, rank=rank, world_size=size) 35 | fn() 36 | 37 | def init_local(size, fn, backend="gloo"): 38 | import torch.multiprocessing 39 | 40 | processes = [] 41 | torch.multiprocessing.set_start_method("spawn") 42 | for rank in range(size): 43 | p = torch.multiprocessing.Process(target=_local_init_process, args=(rank, size, fn, backend)) 44 | p.start() 45 | processes.append(p) 46 | 47 | for p in processes: 48 | p.join() 49 | 50 | 51 | def run(): 52 | print( 53 | "Hello from process {} on node {} out of {}" 54 | "".format(torch.distributed.get_rank(), socket.gethostname(), torch.distributed.get_world_size()) 55 | ) 56 | 57 | def main(): 58 | if "SLURM_JOB_ID" in os.environ: 59 | init_slurm(run) 60 | else: 61 | init_local(torch.cuda.device_count(), run) 62 | 63 | if __name__ == '__main__': 64 | main() 65 | -------------------------------------------------------------------------------- /nqs_playground/lbfgs.py: -------------------------------------------------------------------------------- 1 | # Copyright Tom Westerhout (c) 2021 2 | # 3 | # All rights reserved. 4 | # 5 | # Redistribution and use in source and binary forms, with or without 6 | # modification, are permitted provided that the following conditions are met: 7 | # 8 | # * Redistributions of source code must retain the above copyright 9 | # notice, this list of conditions and the following disclaimer. 10 | # 11 | # * Redistributions in binary form must reproduce the above 12 | # copyright notice, this list of conditions and the following 13 | # disclaimer in the documentation and/or other materials provided 14 | # with the distribution. 15 | # 16 | # * Neither the name of Tom Westerhout nor the names of other 17 | # contributors may be used to endorse or promote products derived 18 | # from this software without specific prior written permission. 19 | # 20 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 21 | # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 22 | # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 23 | # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 24 | # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 25 | # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 26 | # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 27 | # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 28 | # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 30 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 31 | 32 | from LBFGS import LBFGS 33 | 34 | from collections import namedtuple 35 | import time 36 | import os 37 | from loguru import logger 38 | import numpy as np 39 | import torch 40 | from torch import Tensor 41 | from torch.utils.tensorboard import SummaryWriter 42 | 43 | from . import * 44 | 45 | Config = namedtuple( 46 | "Config", 47 | [ 48 | "amplitude", 49 | "phase", 50 | "hamiltonian", 51 | "output", 52 | "epochs", 53 | "sampling_options", 54 | "optimizer", 55 | "scheduler", 56 | "exact", 57 | "constraints", 58 | "inference_batch_size", 59 | "checkpoint_every", 60 | ], 61 | defaults=[], 62 | ) 63 | 64 | 65 | class Runner(RunnerBase): 66 | def __init__(self, config): 67 | super().__init__(config) 68 | self.combined_state = combine_amplitude_and_phase( 69 | self.config.amplitude, self.config.phase, use_jit=False 70 | ) 71 | 72 | def outer_iteration(self, number_inner: int): 73 | (states, log_probs, weights) = self.perform_sampling() 74 | if log_probs is None: 75 | log_probs = self.compute_log_probs(states) 76 | assert not torch.any(torch.isnan(log_probs)) 77 | weights = weights.to(torch.float64) 78 | log_probs = log_probs.to(torch.float64) 79 | 80 | original_log_probs = log_probs 81 | log_original_weights = weights.to(torch.float64, copy=True).log_() 82 | 83 | def closure(): 84 | nonlocal states 85 | log_probs = self.compute_log_probs(states).to(torch.float64) 86 | weights = recompute_weights(log_probs, original_log_probs, log_original_weights) 87 | local_energies, energy, _, info = local_values_with_extras( 88 | (states, None, weights), 89 | self.config.hamiltonian, 90 | self.combined_state, 91 | self.config.inference_batch_size 92 | ) 93 | log_stuff_to_tensorboard(info, self.global_index, self.tb_writer) 94 | 95 | weights = weights.view(-1) 96 | local_energies = local_energies.real.view(-1) 97 | grad = 2 * weights * (local_energies - energy.real) 98 | grad = grad.view(-1, 1) 99 | 100 | self.config.optimizer.zero_grad() 101 | self.config.amplitude.train() 102 | output = torch.sum(self.config.amplitude(states.view(-1, states.size(-1))) * grad) 103 | print(output) 104 | assert output == energy.real 105 | return output 106 | # forward_fn = self.amplitude_forward_fn 107 | # for (states_chunk, grad_chunk) in split_into_batches((, grad), 1024): 108 | # output = forward_fn(states_chunk) 109 | # output.backward(grad_chunk) 110 | 111 | # grad = self.config.optimizer._gather_flat_grad() 112 | # return grad, energy.real 113 | 114 | 115 | for i in range(number_inner): 116 | 117 | grad, energy = compute_gradient(should_recompute_weights=i > 0) 118 | p = self.config.optimizer.two_loop_recursion(-grad) 119 | 120 | # @torch.no_grad() 121 | # def closure(): 122 | # log_probs = self.compute_log_probs(states).to(torch.float64) 123 | # weights = recompute_weights(log_probs, original_log_probs, log_original_weights) 124 | # _, energy, _, info = local_values_with_extras( 125 | # (states, None, weights), 126 | # self.config.hamiltonian, 127 | # self.combined_state, 128 | # self.config.inference_batch_size 129 | # ) 130 | # log_stuff_to_tensorboard(info, self.global_index, self.tb_writer) 131 | # return energy.real 132 | 133 | options = {'closure': closure, 'current_loss': energy} 134 | energy, grad, lr, _, _, _, _, _ = self.config.optimizer.step(p, grad, options=options) 135 | self.config.optimizer.curvature_update(grad) 136 | 137 | self.global_index += 1 138 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | import os 3 | import re 4 | 5 | 6 | def get_version(package): 7 | pwd = os.path.dirname(os.path.realpath(__file__)) 8 | with open(os.path.join(pwd, package, "__init__.py"), "r") as input: 9 | result = re.search(r'__version__\s*=\s*[\'"]([^\'"]*)[\'"]', input.read()) 10 | if not result: 11 | raise ValueError("failed to determine {} version".format(package)) 12 | return result.group(1) 13 | 14 | 15 | setup( 16 | name="nqs-playground", 17 | version=get_version("nqs_playground"), 18 | description="PyTorch-based implementation of SR and SWO for NQS", 19 | classifiers=[ 20 | "Development Status :: 2 - Pre-Alpha", 21 | "Environment :: Console", 22 | "Intended Audience :: Science/Research", 23 | "License :: OSI Approved :: BSD License", 24 | "Programming Language :: Python :: 3 :: Only", 25 | "Topic :: Scientific/Engineering :: Physics", 26 | ], 27 | url="http://github.com/twesterhout/nqs-playground", 28 | author="Tom Westerhout", 29 | author_email="14264576+twesterhout@users.noreply.github.com", 30 | license="BSD3", 31 | packages=["nqs_playground"], 32 | # This will break on OSX and Windows... 33 | package_data={"nqs_playground": ["_C.*.so", "libnqs.so"]}, 34 | install_requires=[], # "torch", "numpy", "scipy", "loguru"], 35 | zip_safe=False, 36 | ) 37 | -------------------------------------------------------------------------------- /test/basis_5x5.pickle: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/twesterhout/nqs-playground/c972002910e68ed20be4d8d89a0ec1848a203384/test/basis_5x5.pickle -------------------------------------------------------------------------------- /test/difficult_to_sample_5x5.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/twesterhout/nqs-playground/c972002910e68ed20be4d8d89a0ec1848a203384/test/difficult_to_sample_5x5.pt -------------------------------------------------------------------------------- /test/test_acceptance.py: -------------------------------------------------------------------------------- 1 | from math import pi as π 2 | from math import sqrt 3 | import numpy as np 4 | import torch 5 | from torch import Tensor 6 | 7 | from nqs_playground import * 8 | from nqs_playground.core import forward_with_batches 9 | 10 | 11 | def make_distribution(basis): 12 | n = basis.number_states 13 | xs = torch.arange(0, n, dtype=torch.float64) 14 | indices = torch.randperm(n) 15 | 16 | gauss = lambda x, μ, σ: torch.exp(- 0.5 * ((x - μ) / σ)**2) 17 | 18 | def _make_p(): 19 | ys = gauss(xs, n / 2, n / 20) 20 | ys /= torch.sum(ys) 21 | def distribution(states: Tensor) -> Tensor: 22 | if len(states.shape) == 2: 23 | states = states[:, 0] 24 | return ys[indices[basis.index(states)]].squeeze() 25 | return distribution 26 | 27 | def _make_E(): 28 | ys = 6 * gauss(xs, (0.5 - 0.05) * n, n / 16) + gauss(xs, (0.5 + 0.15) * n, n / 30) 29 | def observable(states: Tensor) -> Tensor: 30 | if len(states.shape) == 2: 31 | states = states[:, 0] 32 | return ys[indices[basis.index(states)]].squeeze() 33 | return observable 34 | 35 | return _make_p(), _make_E() 36 | 37 | def simple(): 38 | basis = SpinBasis([], 28, 14) 39 | basis.build() 40 | all_states = torch.from_numpy(basis.states.view(np.int64)) 41 | print(basis.number_states) 42 | p, E = make_distribution(basis) 43 | print(torch.sum(p(all_states))) 44 | print(torch.sum(E(all_states))) 45 | exact = torch.dot(p(all_states), E(all_states)) 46 | print(exact, torch.dot(p(all_states), (E(all_states) - exact)**2)) 47 | 48 | with torch.no_grad(): 49 | i = np.random.choice(basis.number_states, size=basis.number_spins * 2000, p=p(all_states).numpy(), replace=True) 50 | print(torch.mean(E(all_states[i])), torch.var(E(all_states[i]))) 51 | 52 | options = SamplingOptions( 53 | number_chains=basis.number_spins, 54 | number_samples=1000, 55 | device="cpu" 56 | ) 57 | x, y, r = sample_some(lambda x: 0.5 * torch.log(p(x)), basis, options, 58 | mode="monte_carlo") 59 | print(torch.mean(E(x.view(-1, 8))), torch.var(E(x.view(-1, 8)))) 60 | print(r) 61 | 62 | def make_basis(): 63 | Lx = 6 64 | Ly = 6 65 | NUMBER_SPINS = Lx * Ly 66 | sites = np.arange(NUMBER_SPINS) # site labels [0,1,2,....] 67 | x = sites%Lx # x positions for sites 68 | y = sites//Lx # y positions for sites 69 | # 70 | T_x = (x+1)%Lx + Lx*y # translation along x-direction 71 | T_y = x +Lx*((y+1)%Ly) # translation along y-direction 72 | # 73 | P_x = x + Lx*(Ly-y-1) # reflection about x-axis 74 | P_y = (Lx-x-1) + Lx*y # reflection about y-axis 75 | # 76 | symmetry_group = make_group( 77 | [ 78 | # Translation 79 | Symmetry( 80 | list(T_x), sector= 0 #NUMBER_SPINS // 2 81 | ), 82 | Symmetry( 83 | list(T_y), sector= 0 #NUMBER_SPINS // 2 84 | ), 85 | # Reflections 86 | Symmetry( 87 | list(P_x), sector= 0 #NUMBER_SPINS // 2 88 | ), 89 | Symmetry( 90 | list(P_y), sector= 0 #NUMBER_SPINS // 2 91 | ), 92 | ] 93 | ) 94 | return SpinBasis(symmetry_group, number_spins=NUMBER_SPINS, 95 | hamming_weight=NUMBER_SPINS // 2) 96 | 97 | 98 | def andrey(): 99 | import pickle 100 | with open("basis_6x6.pickle", "rb") as f: 101 | basis = pickle.load(f) 102 | # with open("basis_6x6.pickle", "wb") as f: 103 | # pickle.dump(basis, f) 104 | 105 | log_ψ = torch.jit.script(torch.nn.Sequential( 106 | torch.nn.Linear(basis.number_spins, 64), 107 | torch.nn.ReLU(), 108 | torch.nn.Linear(64, 64), 109 | torch.nn.ReLU(), 110 | torch.nn.Linear(64, 1, bias=False), 111 | )) 112 | log_ψ.load_state_dict(torch.load("/home/twesterh/mc/Fail2/amplitude_weights.pt")) 113 | options = SamplingOptions( 114 | number_chains=basis.number_spins, 115 | number_samples=1000, 116 | device="cpu" 117 | ) 118 | 119 | ys = forward_with_batches( 120 | lambda x: log_ψ(unpack(x, basis.number_spins)), 121 | torch.from_numpy(basis.states.view(np.int64)), 122 | batch_size=81920).squeeze() 123 | ys = torch.sort(ys)[0] 124 | print(ys[:10]) 125 | print(ys[-10:]) 126 | print(torch.histc(ys).detach().numpy().tolist()) 127 | 128 | x, y, r = sample_some(lambda x: log_ψ(unpack(x, basis.number_spins)), basis, options, mode="monte_carlo") 129 | print(r) 130 | 131 | 132 | if __name__ == '__main__': 133 | andrey() 134 | -------------------------------------------------------------------------------- /test/test_apply.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import lattice_symmetries as ls 4 | import nqs_playground as nqs 5 | 6 | from nqs_playground._extension import lib as _C 7 | 8 | def test_apply(use_jit=False): 9 | # ls.enable_logging() 10 | basis = ls.SpinBasis(ls.Group([]), number_spins=10, hamming_weight=None) 11 | basis.build() 12 | matrix = np.array([[1, 0, 0, 0], [0, -1, -2, 0], [0, -2, -1, 0], [0, 0, 0, 1]]) 13 | edges = [(i, (i + 1) % basis.number_spins) for i in range(basis.number_spins)] 14 | op = ls.Operator(basis, [ls.Interaction(matrix, edges)]) 15 | 16 | class MyModule(torch.nn.Module): 17 | def __init__(self, n): 18 | super().__init__() 19 | self.fn = torch.nn.Sequential( 20 | nqs.Unpack(n), 21 | torch.nn.Linear(n, 50), 22 | torch.nn.ReLU(inplace=True), 23 | torch.nn.Linear(50, 2, bias=False), 24 | ) 25 | 26 | def forward(self, x): 27 | y = self.fn(x) 28 | return torch.complex(y[:, 0], y[:, 1]) 29 | 30 | log_psi = MyModule(basis.number_spins) 31 | if use_jit == True: 32 | log_psi = torch.jit.script(log_psi) 33 | 34 | devices = ["cpu"] 35 | if torch.cuda.is_available(): 36 | devices.append("cuda") 37 | for device in devices: 38 | for batch_size in range(1, 10): 39 | for inference_batch_size in range(1, 20): 40 | states = basis.states[np.random.permutation(basis.number_states)[:batch_size]] 41 | states = torch.from_numpy(states.view(np.int64)) 42 | states = nqs.pad_states(states) 43 | predicted = _C.log_apply(states.to(device), op, log_psi.to(device), batch_size) 44 | expected = nqs.reference_log_apply(states.cpu(), op, log_psi.cpu(), batch_size) 45 | assert torch.allclose(predicted.cpu(), expected) 46 | 47 | 48 | test_apply() 49 | -------------------------------------------------------------------------------- /test/test_basis.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from nqs_playground import * 3 | 4 | 5 | def test_construction(): 6 | _ = SpinBasis([], number_spins=10, hamming_weight=5) 7 | _ = SpinBasis([], number_spins=200, hamming_weight=100) 8 | _ = SpinBasis([], number_spins=20) 9 | _ = SpinBasis( 10 | make_group( 11 | [ 12 | Symmetry([1, 2, 3, 4, 5, 0], sector=3), 13 | Symmetry([5, 4, 3, 2, 1, 0], sector=1), 14 | ] 15 | ), 16 | 6, 17 | ) 18 | extend = lambda x: x + list(range(len(x), 500)) 19 | _ = SpinBasis( 20 | make_group( 21 | [ 22 | Symmetry(extend([1, 2, 3, 4, 5, 0]), sector=3), 23 | Symmetry(extend([5, 4, 3, 2, 1, 0]), sector=1), 24 | ] 25 | ), 26 | 500, 27 | ) 28 | 29 | 30 | def test_building(): 31 | basis = SpinBasis([], number_spins=6, hamming_weight=3) 32 | basis.build() 33 | assert basis.number_spins == 6 34 | assert basis.hamming_weight == 3 35 | assert basis.number_states == 20 36 | assert basis.index(int("111000", base=2)) == 19 37 | assert basis.index(int("000111", base=2)) == 0 38 | assert basis.index(int("001011", base=2)) == 1 39 | assert basis.index(int("010101", base=2)) == 5 40 | 41 | 42 | # def to_int(k) -> int: 43 | # acc = int(k[7]) 44 | # for i in range(6, -1, -1): 45 | # acc <<= 64 46 | # acc |= k[i] 47 | # return acc 48 | 49 | 50 | def test_full_info(): 51 | L_x, L_y = 6, 4 52 | indices = np.arange(L_x * L_y).reshape(L_y, L_x)[::-1] 53 | x = indices % L_x 54 | y = indices // L_x 55 | T_x = (x + 1) % L_x + y * L_x 56 | T_y = x % L_x + ((y + 1) % L_y) * L_x 57 | 58 | T_x = T_x.reshape(-1) 59 | T_y = T_y.reshape(-1) 60 | 61 | indices = np.arange(24) 62 | np.random.shuffle(indices) 63 | def shuffle(xs): 64 | return [xs[i] for i in indices] 65 | 66 | basis_small = SpinBasis( 67 | shuffle(make_group([Symmetry(T_x, sector=1), Symmetry(T_y, sector=0)])) 68 | number_spins=L_x * L_y, 69 | hamming_weight=(L_x * L_y) // 2, 70 | ) 71 | e = lambda p: p.tolist() + list(range(L_x * L_y, 100)) 72 | basis_big = SpinBasis( 73 | shuffle(make_group([Symmetry(e(T_x), sector=1), Symmetry(e(T_y), sector=0)])), 74 | number_spins=L_x * L_y, 75 | hamming_weight=(L_x * L_y) // 2, 76 | ) 77 | 78 | basis_small.build() 79 | for state in basis_small.states: 80 | r1, e1, n1 = basis_small.full_info(state) 81 | r2, e2, n2 = basis_big.full_info(state) 82 | assert r1 == r2 83 | assert e1 == e2 84 | assert n1 == n2 85 | 86 | basis_dummy = SpinBasis([], L_x * L_y, (L_x * L_y) // 2) 87 | basis_dummy.build() 88 | for state in basis_dummy.states: 89 | r1, e1, n1 = basis_small.full_info(state) 90 | r2, e2, n2 = basis_big.full_info(state) 91 | assert r1 == r2 92 | assert n1 == n2 93 | if n1 != 0.0: 94 | assert e1 == e2 95 | else: 96 | assert np.isnan(e2) 97 | -------------------------------------------------------------------------------- /test/test_hamiltonian.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from nqs_playground import * 4 | from nqs_playground import _C 5 | 6 | try: 7 | from typing_extensions import Final 8 | except ImportError: 9 | # If you don't have `typing_extensions` installed, you can use a 10 | # polyfill from `torch.jit`. 11 | from torch.jit import Final 12 | 13 | 14 | def test_exact_diagonalisation(): 15 | n = 10 16 | basis = SpinBasis([], number_spins=n, hamming_weight=n // 2) 17 | hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis) 18 | energy, eigenvector = diagonalise(hamiltonian) 19 | assert np.isclose(energy, -18.06178542) 20 | basis = SpinBasis( 21 | make_group( 22 | [ 23 | Symmetry(list(range(1, n)) + [0], sector=n // 2), 24 | Symmetry(list(range(n))[::-1], sector=1), 25 | ] 26 | ), 27 | number_spins=n, 28 | hamming_weight=n // 2, 29 | ) 30 | hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis) 31 | energy, eigenvector = diagonalise(hamiltonian) 32 | assert np.isclose(energy, -18.06178542) 33 | 34 | 35 | def test_sparse_diagonalisation(): 36 | import scipy.sparse 37 | import scipy.sparse.linalg 38 | 39 | n = 10 40 | basis = SpinBasis([], number_spins=n, hamming_weight=n // 2) 41 | basis.build() 42 | hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis) 43 | m = hamiltonian.to_csr() 44 | energy, _ = scipy.sparse.linalg.eigsh(m, k=1) 45 | assert np.isclose(energy, -18.06178542) 46 | 47 | 48 | def reference_apply(spins, hamiltonian, psi): 49 | with torch.no_grad(): 50 | basis = hamiltonian.basis 51 | states = torch.from_numpy(basis.states.view(np.int64)) 52 | psi = psi(states) 53 | scale = torch.max(psi[:, 0]).item() 54 | psi[:, 0] -= scale 55 | psi = np.exp(psi.numpy().view(np.complex64)) 56 | hamiltonian = hamiltonian.to_csr() 57 | out = np.empty(len(spins), dtype=np.complex64) 58 | for i, s in enumerate(spins): 59 | state = np.zeros(basis.number_states, dtype=np.complex64) 60 | state[basis.index(s)] = 1.0 61 | state = hamiltonian @ state 62 | out[i] = scale + np.log(np.dot(state.conj(), psi)) 63 | return out 64 | 65 | 66 | def reference_diag(spins, hamiltonian): 67 | basis = hamiltonian.basis 68 | hamiltonian = hamiltonian.to_csr() 69 | out = np.empty(len(spins), dtype=np.complex64) 70 | for i, s in enumerate(spins): 71 | state = np.zeros(basis.number_states, dtype=np.complex64) 72 | state[basis.index(s)] = 1.0 73 | state = hamiltonian @ state 74 | out[i] = state[basis.index(s)] 75 | return out 76 | 77 | 78 | class Unpack(torch.nn.Module): 79 | n: Final[int] 80 | 81 | def __init__(self, n: int): 82 | super().__init__() 83 | self.n = n 84 | 85 | def forward(self, x): 86 | return unpack(x, self.n) 87 | 88 | 89 | def test_apply(): 90 | n = 20 91 | basis = SpinBasis([], number_spins=n, hamming_weight=n // 2) 92 | basis.build() 93 | hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis) 94 | for i in range(5): 95 | spins = torch.from_numpy(basis.states.view(np.int64))[ 96 | torch.randperm(basis.number_states)[:500] 97 | ] 98 | psi = torch.jit.script( 99 | torch.nn.Sequential( 100 | Unpack(n), 101 | torch.nn.Linear(n, 5), 102 | torch.nn.Tanh(), 103 | torch.nn.Linear(5, 2, bias=False), 104 | ) 105 | ) 106 | expected = reference_apply(spins, hamiltonian, psi) 107 | spins_512 = torch.cat( 108 | [spins.unsqueeze(dim=1), torch.zeros(spins.size(0), 7, dtype=torch.int64)], 109 | dim=1, 110 | ) 111 | predicted = ( 112 | _C.apply(spins_512, hamiltonian, psi._c._get_method("forward")) 113 | .numpy() 114 | .view(np.complex64) 115 | .squeeze() 116 | ) 117 | assert np.allclose(predicted, expected) 118 | 119 | 120 | def test_diag(): 121 | n = 20 122 | basis = SpinBasis([], number_spins=n, hamming_weight=n // 2) 123 | basis.build() 124 | hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis) 125 | spins = torch.from_numpy(basis.states.view(np.int64))[ 126 | torch.randperm(basis.number_states)[:1000] 127 | ] 128 | expected = reference_diag(spins, hamiltonian) 129 | spins_512 = torch.cat( 130 | [spins.unsqueeze(dim=1), torch.zeros(spins.size(0), 7, dtype=torch.int64)], 131 | dim=1, 132 | ) 133 | predicted = _C.diag(spins_512, hamiltonian).numpy().view(np.complex64).squeeze() 134 | assert np.allclose(predicted, expected) 135 | -------------------------------------------------------------------------------- /test/test_jacobian.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | 4 | from nqs_playground import jacobian_simple, jacobian_cpu, jacobian_cuda 5 | 6 | 7 | 8 | def make_dense_network(): 9 | return torch.jit.script(torch.nn.Sequential( 10 | torch.nn.Linear(20, 32), 11 | torch.nn.ReLU(), 12 | torch.nn.Linear(32, 32), 13 | torch.nn.ReLU(), 14 | torch.nn.Linear(32, 1, bias=False), 15 | )) 16 | 17 | def test_cpu(): 18 | for i in range(7): 19 | m = make_dense_network() 20 | for j in range(7): 21 | x = torch.rand(57, 20) 22 | j_expected = jacobian_simple(m, x) 23 | j_predicted = jacobian_cpu(m, x) 24 | assert torch.allclose(j_expected, j_predicted) 25 | 26 | def test_cuda(): 27 | for i in range(7): 28 | m = make_dense_network() 29 | for j in range(7): 30 | x = torch.rand(57, 20) 31 | j_expected = jacobian_simple(m, x) 32 | m.cuda() 33 | x = x.cuda() 34 | j_predicted = jacobian_cuda(m, x).cpu() 35 | m.cpu() 36 | x = x.cpu() 37 | try: 38 | assert torch.allclose(j_expected, j_predicted) 39 | except AssertionError: 40 | torch.jit.save(m, "bad_network.pth") 41 | torch.save(x, "bad_input.pth") 42 | raise 43 | 44 | # torch.backends.cudnn.deterministic = True 45 | # torch.backends.cudnn.benchmark = False 46 | # 47 | # def debug(): 48 | # m = torch.jit.load("bad_network.pth") 49 | # x = torch.load("bad_input.pth") 50 | # m.cuda() 51 | # x = x.cuda() 52 | # j_expected = jacobian_simple(m, x).cpu() 53 | # j_predicted = jacobian_cuda(m, x).cpu() 54 | # for i in range(j_expected.size(0)): 55 | # if not torch.allclose(j_expected[i], j_predicted[i]): 56 | # for j in range(j_expected.size(1)): 57 | # if not torch.allclose(j_expected[i, j], j_predicted[i, j]): 58 | # print(i, j, j_expected[i, j].tolist(), j_predicted[i, j].tolist()) 59 | # 60 | # 61 | # debug() 62 | -------------------------------------------------------------------------------- /test/test_local_values.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | import nqs_playground 5 | 6 | 7 | class SlowPolynomialState: 8 | def __init__(self, hamiltonian, roots, log_psi): 9 | self.hamiltonian = hamiltonian 10 | self.basis = hamiltonian.basis 11 | self.basis.build() 12 | self.state, self.scale = self._make_state(log_psi) 13 | self.roots = roots 14 | 15 | def _make_state(self, log_psi): 16 | with torch.no_grad(): 17 | spins = torch.from_numpy(self.basis.states.view(np.int64)) 18 | out = log_psi(torch.ops.tcm.unpack(spins, self.basis.number_spins)) 19 | scale = torch.max(out[:, 0]).item() 20 | out[:, 0] -= scale 21 | out = out.numpy().view(np.complex64).squeeze() 22 | out = np.exp(out) 23 | return out, scale 24 | 25 | def _forward_one(self, x): 26 | basis = self.hamiltonian.basis 27 | i = basis.index(x) 28 | vector = np.zeros((basis.number_states,), dtype=np.complex64) 29 | vector[i] = 1.0 30 | for r in self.roots: 31 | vector = self.hamiltonian(vector) - r * vector 32 | return self.scale + np.log(np.dot(self.state, vector)) 33 | 34 | def __call__(self, spins): 35 | if isinstance(spins, torch.Tensor): 36 | assert spins.dtype == torch.int64 37 | spins = spins.numpy().view(np.uint64) 38 | return torch.from_numpy( 39 | np.array([self._forward_one(x) for x in spins], dtype=np.complex64) 40 | .view(np.float32) 41 | .reshape(-1, 2) 42 | ) 43 | 44 | 45 | basis = nqs_playground.SpinBasis([], 16, 8) 46 | basis.build() 47 | hamiltonian = nqs_playground.Heisenberg( 48 | [(1.0, i, (i + 1) % basis.number_spins) for i in range(basis.number_spins)], basis 49 | ) 50 | m = torch.jit.script( 51 | torch.nn.Sequential( 52 | torch.nn.Linear(basis.number_spins, 64), 53 | torch.nn.ReLU(), 54 | torch.nn.Linear(64, 64), 55 | torch.nn.ReLU(), 56 | torch.nn.Linear(64, 64), 57 | torch.nn.ReLU(), 58 | torch.nn.Linear(64, 2), 59 | ) 60 | ) 61 | 62 | spins = torch.from_numpy(basis.states.view(np.int64)) 63 | 64 | A = SlowPolynomialState(hamiltonian, [0.0], m)(spins) 65 | B = nqs_playground._C.apply(spins, hamiltonian, m._c._get_method("forward")) 66 | if not torch.allclose(A, B): 67 | for i in range(A.size(0)): 68 | if not torch.allclose(A[i], B[i]): 69 | print(i, A[i], B[i]) 70 | print(A[i - 3 : i + 3]) 71 | print(B[i - 3 : i + 3]) 72 | break 73 | else: 74 | print("[+] passed!") 75 | -------------------------------------------------------------------------------- /test/test_monte_carlo.py: -------------------------------------------------------------------------------- 1 | from math import sqrt 2 | import pathlib 3 | import time 4 | import numpy as np 5 | import torch 6 | from nqs_playground import * 7 | from nqs_playground.core import Unpack 8 | from nqs_playground._tabu_sampler import _sample_using_zanella 9 | from nqs_playground.monte_carlo import _log_amplitudes_to_probabilities 10 | 11 | np.random.seed(1183472) 12 | torch.manual_seed(14346284) 13 | 14 | CURRENT_DIR = pathlib.Path(__file__).parent.absolute() 15 | 16 | @torch.no_grad() 17 | def histogram(spins, basis): 18 | if spins.dim() == 2: 19 | assert spins.size(1) == 8 20 | spins = spins[:, 0] 21 | spins = spins.cpu() 22 | 23 | r = torch.zeros(basis.number_states, dtype=torch.int64) 24 | spins, counts = torch.unique(spins, sorted=True, return_counts=True) 25 | r[basis.index(spins)] += counts 26 | return r 27 | 28 | @torch.no_grad() 29 | def are_close_l1(n, basis, sample, exact, eps, sweep_size=None, device=None): 30 | exact, order = torch.sort(exact) 31 | s = np.searchsorted(torch.cumsum(exact, dim=0).numpy(), eps / 8.0) 32 | ms = np.random.poisson(n, size=16) 33 | options = SamplingOptions(number_chains=len(ms), number_samples=max(ms), 34 | sweep_size=sweep_size, device=None) 35 | qs, _ = sample(options) 36 | qs = [histogram(qs[:m, i], basis)[order] for i, m in enumerate(ms)] 37 | 38 | def analyze(x, k): 39 | v = ((x - k * exact)**2 - x) * exact**(-2/3) 40 | w = exact**(2/3) 41 | cond1 = torch.sum(v[s:-1]) > 4 * k * torch.sum(w[s:-1])**(1/2) 42 | cond2 = torch.sum(x[:s]) > 3 / 16 * eps * k 43 | return not (cond1 or cond2) 44 | 45 | return [analyze(x, k) for x, k in zip(qs, ms)] 46 | 47 | 48 | def calculate_exact_probabilities(basis, log_ψ, device="cpu"): 49 | xs = torch.from_numpy(basis.states.view(np.int64)).to(device) 50 | ys = forward_with_batches(log_ψ, xs, batch_size=8192).squeeze() 51 | return _log_amplitudes_to_probabilities(ys) 52 | 53 | 54 | def test_simple(): 55 | basis = SpinBasis([], number_spins=10, hamming_weight=5) 56 | basis.build() 57 | 58 | log_ψ = torch.nn.Sequential( 59 | Unpack(basis.number_spins), 60 | torch.nn.Linear(basis.number_spins, 64), 61 | torch.nn.ReLU(), 62 | torch.nn.Linear(64, 64), 63 | torch.nn.ReLU(), 64 | torch.nn.Linear(64, 1), 65 | ) 66 | 67 | exact = calculate_exact_probabilities(basis, log_ψ) 68 | 69 | def sample_exact(options): 70 | return sample_some(log_ψ, basis, options, mode="monte_carlo")[:2] 71 | 72 | def sample_metropolis(options): 73 | return sample_some(log_ψ, basis, options, mode="monte_carlo")[:2] 74 | 75 | def sample_zanella(options): 76 | return _sample_using_zanella(log_ψ, basis, options, thin_rate=1e-1) 77 | 78 | for f in [sample_exact, sample_metropolis, sample_zanella]: 79 | r = are_close_l1(1000, basis, f, exact, eps=1e-5) 80 | print(r) 81 | assert sum(r) > len(r) / 2 82 | 83 | 84 | def test_tricky(): 85 | L = 20 86 | x = np.arange(L, dtype=np.int32) 87 | T = (x + 1) % L 88 | P = L - 1 - x 89 | G = make_group([Symmetry(T, sector=0), Symmetry(P, sector=0)]) 90 | basis = SpinBasis(G, number_spins=L, hamming_weight=L // 2) 91 | basis.build() 92 | 93 | log_ψ = torch.nn.Sequential( 94 | Unpack(L), 95 | torch.nn.Linear(L, L // 2), 96 | torch.nn.Tanh(), 97 | torch.nn.Linear(L // 2, 1, bias=False) 98 | ) 99 | exact = calculate_exact_probabilities(basis, log_ψ) 100 | 101 | def sample_exact(options): 102 | return sample_some(log_ψ, basis, options, mode="exact")[:2] 103 | 104 | def sample_metropolis(options): 105 | return sample_some(log_ψ, basis, options, mode="monte_carlo")[:2] 106 | 107 | def sample_zanella(options): 108 | return _sample_using_zanella(log_ψ, basis, options, thin_rate=2e-2) 109 | 110 | for f in [sample_exact, sample_metropolis, sample_zanella]: 111 | r = are_close_l1(2000, basis, f, exact, eps=1e-5) 112 | print(r) 113 | assert sum(r) > len(r) / 2 114 | 115 | 116 | def test_5x5_with_symmetries(): 117 | import pickle 118 | basis = with_file_like(CURRENT_DIR / "basis_5x5.pickle", "rb", pickle.load) 119 | log_ψ = torch.jit.load(str(CURRENT_DIR / "difficult_to_sample_5x5.pt")) 120 | exact = calculate_exact_probabilities(basis, log_ψ) 121 | 122 | def sample_exact(options): 123 | return sample_some(log_ψ, basis, options, mode="exact")[:2] 124 | 125 | def sample_metropolis(options): 126 | return sample_some(log_ψ, basis, options, mode="monte_carlo")[:2] 127 | 128 | def sample_zanella(options): 129 | return _sample_using_zanella(log_ψ, basis, options, thin_rate=1e0) 130 | 131 | for f in [ sample_exact, 132 | # sample_metropolis, 133 | sample_zanella, 134 | ]: 135 | r = are_close_l1(10000, basis, f, exact, eps=1e-2, sweep_size=100) 136 | print(r) 137 | assert sum(r) > len(r) / 2 138 | return 139 | 140 | def profile_5x5_with_symmetries(): 141 | import pickle 142 | import pprofile 143 | basis = with_file_like(CURRENT_DIR / "basis_5x5.pickle", "rb", pickle.load) 144 | log_ψ = torch.jit.load(str(CURRENT_DIR / "difficult_to_sample_5x5.pt")) 145 | options = SamplingOptions(number_chains=16, number_samples=100, 146 | sweep_size=100, device=None) 147 | prof = pprofile.Profile() 148 | with prof(): 149 | # sample_some(log_ψ, basis, options, mode="exact") 150 | # sample_some(log_ψ, basis, options, mode="monte_carlo") 151 | _sample_using_zanella(log_ψ, basis, options, thin_rate=1e0) 152 | prof.print_stats() 153 | 154 | # test_simple() 155 | # test_tricky() 156 | # test_zanella() 157 | test_5x5_with_symmetries() 158 | # profile_5x5_with_symmetries() 159 | -------------------------------------------------------------------------------- /test/test_symmetry.py: -------------------------------------------------------------------------------- 1 | from nqs_playground import * 2 | from nqs_playground import _C 3 | import numpy as np 4 | 5 | # Reference implementation 6 | def _btfly_step(x: int, m: int, d: int) -> int: 7 | y = (x ^ (x >> d)) & m 8 | return x ^ y ^ (y << d) 9 | 10 | 11 | def test_butterfly_step(): 12 | def _to_int(k: np.ndarray) -> int: 13 | acc = int(k[7]) 14 | for i in range(6, -1, -1): 15 | acc <<= 64 16 | acc |= k[i] 17 | return acc 18 | 19 | def _to_array(k: int) -> np.ndarray: 20 | return np.array([((k >> (64 * i)) & 0xFFFFFFFFFFFFFFFF) for i in range(8)], dtype=np.uint64) 21 | 22 | for i in range(9): 23 | d = 1 << i 24 | xs = np.random.randint(0, 0xFFFFFFFFFFFFFFFF, size=(10000, 8), dtype=np.uint64) 25 | ms = np.random.randint(0, 0xFFFFFFFFFFFFFFFF, size=(10000, 8), dtype=np.uint64) 26 | for x, m in zip(xs, ms): 27 | predicted = _C._btfly(x, m, d) 28 | expected = _to_array(_btfly_step(_to_int(x), _to_int(m), d)).tolist() 29 | assert predicted == expected 30 | 31 | def test_benes(n, bits=None): 32 | import numpy as np 33 | 34 | def make_perm_fn_simple(p: List[int]): 35 | def fn(x): 36 | assert 0 <= x and x < (1 << len(p)) 37 | s = "{1:0{0}b}".format(len(p), x)[::-1] 38 | s = "".join(s[i] for i in p) 39 | return int(s[::-1], base=2) 40 | 41 | return fn 42 | 43 | p = np.arange(n) 44 | for i in range(100): 45 | np.random.shuffle(p) 46 | benes = make_perm_fn(p, bits) 47 | simple = make_perm_fn_simple(p) 48 | if n <= 8: 49 | for x in range(1 << n): 50 | assert benes(x) == simple(x) 51 | else: 52 | for x in np.random.randint(low=0, high=1 << n, size=1000, dtype=np.uint64): 53 | assert benes(x) == simple(x) 54 | 55 | # def test_btfly(): 56 | # x = np.random.randint(0, 0xFFFFFFFFFFFFFFFF, size=8, dtype=np.uint64) 57 | # m = np.random.randint(0, 0xFFFFFFFFFFFFFFFF, size=8, dtype=np.uint64) 58 | # d = 128 59 | # 60 | # def _to_int(k: np.ndarray) -> int: 61 | # acc = int(k[7]) 62 | # for i in range(6, -1, -1): 63 | # acc <<= 64 64 | # acc |= k[i] 65 | # return acc 66 | # 67 | # def _to_array(k: int) -> np.ndarray: 68 | # return np.array([((k >> (64 * i)) & 0xFFFFFFFFFFFFFFFF) for i in range(8)], dtype=np.uint64) 69 | # 70 | # assert np.all(_to_array(_to_int(x)) == x) 71 | # assert np.all(_to_array(_to_int(m)) == m) 72 | # 73 | # print(x, bin(x[7])) 74 | # print(m) 75 | # return _to_array(_btfly_step(_to_int(x), _to_int(m), d)).tolist(), _C._btfly(x, m, d) 76 | 77 | test_butterfly_step() 78 | # a, b = test_btfly() 79 | # print(a) 80 | # print(b) 81 | 82 | 83 | -------------------------------------------------------------------------------- /test/test_unpack.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | try: 5 | from nqs_playground import * 6 | except ImportError: 7 | # For local development when we only compile the C++ extension, but don't 8 | # actually install the package using pip 9 | import os 10 | import sys 11 | 12 | sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")) 13 | from nqs_playground import * 14 | 15 | rng = np.random.default_rng(seed=52339877) 16 | torch.manual_seed(9218823294) 17 | # manual_seed(3362121853) 18 | 19 | 20 | @torch.no_grad() 21 | def unpack_complete(spins): 22 | def unpack_one(bits, out): 23 | bits = bits.numpy().view(np.uint64) 24 | for i in range(8): 25 | word = int(bits[i]) 26 | for j in range(64): 27 | out[64 * i + j] = float(2 * ((word >> j) & 1) - 1) 28 | 29 | device = spins.device 30 | spins = spins.cpu() 31 | out = torch.empty(spins.size(0), 512, dtype=torch.float32) 32 | for i in range(spins.size(0)): 33 | unpack_one(spins[i], out[i]) 34 | return out.to(device) 35 | 36 | 37 | @torch.no_grad() 38 | def generate(batch_size, device): 39 | data = rng.integers(1 << 64 - 1, size=(batch_size, 8), dtype=np.uint64, endpoint=True) 40 | data = torch.from_numpy(data.view(np.int64)).clone().to(device) 41 | return data 42 | 43 | 44 | @torch.no_grad() 45 | def test_unpack(device): 46 | for batch_size in [1, 5, 128, 437]: 47 | packed = generate(batch_size, device) 48 | full = unpack_complete(packed) 49 | for number_spins in [1, 2, 32, 64, 65, 400]: 50 | predicted = unpack(packed, number_spins) 51 | expected = full[:, :number_spins] 52 | assert (predicted == expected).all() 53 | 54 | 55 | test_unpack("cpu") 56 | if torch.cuda.is_available(): 57 | test_unpack("cuda") 58 | --------------------------------------------------------------------------------