diff --git a/examples/neural_quantum_states/.gitignore b/examples/neural_quantum_states/.gitignore new file mode 100644 index 0000000..da11c2d --- /dev/null +++ b/examples/neural_quantum_states/.gitignore @@ -0,0 +1,3 @@ +logger/ +run.sh +**/__pycache__/ diff --git a/examples/neural_quantum_states/README.md b/examples/neural_quantum_states/README.md index fb4f2c3..fb4119f 100644 --- a/examples/neural_quantum_states/README.md +++ b/examples/neural_quantum_states/README.md @@ -1,18 +1,20 @@ # AUTOREGRESSIVE NEURAL QUANTUM STATES FOR QUANTUM CHEMISTRY -This respository contains code jointly developed between the University of Michigan and SandboxAQ to implement the retentive network (RetNet) neural quantum states ansatz outlined in the paper, "Retentive Nueral Quantum States: Efficient Ansatze for Ab Initio Quantum Chemistry," by Oliver Knitter, Dan Zhao, James Stokes, Martin Ganahl, Stefan Leichenauer, and Shravan Veerapaneni. +This repository contains code jointly developed between the University of Michigan and SandboxAQ to implement the retentive network (RetNet) neural quantum states ansatz outlined in the paper, "Retentive Neural Quantum States: Efficient Ansatze for Ab Initio Quantum Chemistry," by Oliver Knitter, Dan Zhao, James Stokes, Martin Ganahl, Stefan Leichenauer, and Shravan Veerapaneni. Preprint available on the arXiv: https://arxiv.org/abs/2411.03900 Corresponding Author: Oliver Knitter, knitter@umich.edu -This repository is based off of the code Tianchen Zhao released (https://github.com/Ericolony/made-qchem) alongside the paper "Scalable neural quantum states architecture for quantum chemistry" (https://arxiv.org/abs/2208.05637). This new code uses neural quantum states (NQS), implemented in PyTorch to calculate electronic ground state energies for second quantized molecular Hamiltonians, which are calculated using PySCF and Tangelo (https://github.com/sandbox-quantum/Tangelo/). The RetNet ansatz implementation was made using the yet-another-retnet repository (https://github.com/fkodom/yet-another-retnet). Other ansatze available are the MADE and Transformer ansatze, which are implemented natively in PyTorch. The Hamiltonian expectation value estimates are calculated following the procedure outlined in Zhao et al.'s paper, but the modular structure of the code allows for relatively simple plug-and-play implementations of different models and Hamiltonian expectation estimate calculators. +This repository is based off of the code Tianchen Zhao released (https://github.com/Ericolony/made-qchem) alongside the paper "Scalable neural quantum states architecture for quantum chemistry" (https://arxiv.org/abs/2208.05637). This new code uses neural quantum states (NQS), implemented in PyTorch to calculate electronic ground state energies for second quantized molecular Hamiltonians. Though not exactly a quantum or hybrid algorithm, this workflow makes use of Tangelo (https://github.com/sandbox-quantum/Tangelo/) and PySCF to calculate Jordan--Wigner encodings of the electronic Hamiltonians, along with some classical chemistry benchmark values. These uses of Tangelo are mainly found in the subdirectory src/data/. + +The RetNet ansatz implementation was made using the yet-another-retnet repository (https://github.com/fkodom/yet-another-retnet). Other ansatze available are the MADE and Transformer ansatze, which are implemented natively in PyTorch. The Hamiltonian expectation value estimates are calculated following the procedure outlined in Zhao et al.'s paper, but the modular structure of the code allows for relatively simple plug-and-play implementations of different models and Hamiltonian expectation estimate calculators. ## Usage ### 1. Environment Setup -- The environment requirements are avaialable in the nqs-qchem.yml file. +- The environment requirements are available in the nqs-qchem.yml file. - Make sure your operating system meets the requirements for PyTorch 2.5.1 and CUDA 11.8. diff --git a/examples/neural_quantum_states/src/__pycache__/__init__.cpython-312.pyc b/examples/neural_quantum_states/src/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 8fd0082..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/__init__.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/__init__.cpython-39.pyc deleted file mode 100644 index 6b7de0f..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/__init__.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/complex.cpython-312.pyc b/examples/neural_quantum_states/src/__pycache__/complex.cpython-312.pyc deleted file mode 100644 index cf41a8a..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/complex.cpython-312.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/complex.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/complex.cpython-39.pyc deleted file mode 100644 index 01ff613..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/complex.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/data_loader.cpython-312.pyc b/examples/neural_quantum_states/src/__pycache__/data_loader.cpython-312.pyc deleted file mode 100644 index 273022e..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/data_loader.cpython-312.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/data_loader.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/data_loader.cpython-39.pyc deleted file mode 100644 index 5508fc6..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/data_loader.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/evaluate.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/evaluate.cpython-39.pyc deleted file mode 100644 index 4cde55a..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/evaluate.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/helper.cpython-312.pyc b/examples/neural_quantum_states/src/__pycache__/helper.cpython-312.pyc deleted file mode 100644 index 4ad21bb..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/helper.cpython-312.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/helper.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/helper.cpython-39.pyc deleted file mode 100644 index ce19ce8..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/helper.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/optimizer.cpython-312.pyc b/examples/neural_quantum_states/src/__pycache__/optimizer.cpython-312.pyc deleted file mode 100644 index 125564f..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/optimizer.cpython-312.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/optimizer.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/optimizer.cpython-39.pyc deleted file mode 100644 index 3a078c0..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/optimizer.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/scheduler.cpython-312.pyc b/examples/neural_quantum_states/src/__pycache__/scheduler.cpython-312.pyc deleted file mode 100644 index 89d6aa1..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/scheduler.cpython-312.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/scheduler.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/scheduler.cpython-39.pyc deleted file mode 100644 index febb6fa..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/scheduler.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/train.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/train.cpython-39.pyc deleted file mode 100644 index fdc0b6d..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/train.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/util.cpython-312.pyc b/examples/neural_quantum_states/src/__pycache__/util.cpython-312.pyc deleted file mode 100644 index 90a9bb6..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/util.cpython-312.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/__pycache__/util.cpython-39.pyc b/examples/neural_quantum_states/src/__pycache__/util.cpython-39.pyc deleted file mode 100644 index 8016643..0000000 Binary files a/examples/neural_quantum_states/src/__pycache__/util.cpython-39.pyc and /dev/null differ diff --git a/examples/neural_quantum_states/src/complex.py b/examples/neural_quantum_states/src/complex.py index 1d924a6..6bac7f0 100644 --- a/examples/neural_quantum_states/src/complex.py +++ b/examples/neural_quantum_states/src/complex.py @@ -32,4 +32,3 @@ def scalar_mult(x: Union[torch.Tensor, np.ndarray], y: Union[torch.Tensor, np.nd re = real(x) * real(y) - imag(x) * imag(y) im = real(x) * imag(y) + imag(x) * real(y) return torch.stack([re, im], dim=-1) - #return torch.stack([re, im], dim=-1) if torch.is_tensor(x) else np.stack([re, im], axis=-1) diff --git a/examples/neural_quantum_states/src/objective/adaptive_shadows.py b/examples/neural_quantum_states/src/objective/adaptive_shadows.py index 1878d92..d13d65b 100644 --- a/examples/neural_quantum_states/src/objective/adaptive_shadows.py +++ b/examples/neural_quantum_states/src/objective/adaptive_shadows.py @@ -13,8 +13,12 @@ class AdaptiveShadows(Hamiltonian): Args: hamiltonian_string: Pauli string representation of Hamiltonian num_sites: qubit number of system + sample_count: Maximum number of Pauli string samples desired + total_unique_samples: Total number of unique Pauli string samples desired + reset_prob: Probability to resample Pauli strings + flip_bs: Number of unique bit flip patterns processed at a time on each GPU ''' - def __init__(self, hamiltonian_string, num_sites, sample_count, total_unique_samples, reset_prob, flip_bs, **kwargs): + def __init__(self, hamiltonian_string: str, num_sites: int, sample_count: int, total_unique_samples: int, reset_prob: float, flip_bs: int, **kwargs): super(AdaptiveShadows, self).__init__(hamiltonian_string, num_sites) # product of identity operators by default, encoded as 0 self.coefficients = torch.stack((self.coefficients.real, self.coefficients.imag), dim=-1) @@ -37,6 +41,9 @@ def __init__(self, hamiltonian_string, num_sites, sample_count, total_unique_sam self.counter = 0 # Counter to keep track of which term to replace when updating sample list def generate_coefficients(self): + ''' + Generates coefficients in Hamiltonian + ''' for i in range(len(self.sample_Z_idx)): cover = self.covers[i] keys = list(cover.keys()) @@ -46,6 +53,9 @@ def generate_coefficients(self): self.sample_coeffs = scalar_mult(self.sample_coeffs, part1) def generate_loss_idxs(self): + ''' + Generate unique bit flip patterns and indices mapping them to each term in the Hamiltonian + ''' flip_idx = self.sample_X_idx + self.sample_Y_idx self.select_idx = self.sample_Y_idx + self.sample_Z_idx self.unique_flips, self.unique_indices = torch.unique(flip_idx, return_inverse=True, dim=0) @@ -53,6 +63,9 @@ def generate_loss_idxs(self): self.unique_num_terms = self.unique_flips.shape[1] def update_sample_batch(self): + ''' + Updates existing sample batch with a new Pauli string sample + ''' new_sample_x, new_sample_y, new_sample_z, new_cover = self.generate_sample_paulis(1) if self.sample_X_idx.shape[0] < self.total_unique_samples: self.sample_X_idx = torch.cat((self.sample_X_idx, new_sample_x.unsqueeze(0)), dim=0) @@ -74,7 +87,15 @@ def update_sample_batch(self): self.generate_loss_idxs() - def generate_sample_paulis(self, num_samples): + def generate_sample_paulis(self, num_samples: int) -> [torch.Tensor, torch.Tensor, torch.Tensor, Counter]: + ''' + Generates one (or several) Pauli string samples according to Adaptive Pauli Shadows sampling scheme + Args: + num_samples: Number of new samples desired + Returns: + samples: indexes specifying locations of X, Y, and Z matrices in sample strings; separate index arrays for each letter + covers: a Counter object showing which Hamiltonian terms are covered by each string according to Adaptive Pauli Shadows + ''' samples = torch.zeros(num_samples, self.input_dim) # 'X' is 1, 'Y' is 2, 'Z' is 3, 'I' is 0 covers = [Counter(dict(zip(np.arange(self.num_terms), np.ones(self.num_terms)))) for _ in range(num_samples)] # Dictionaries of potential covers for each Pauli sample for i in range(num_samples): @@ -106,6 +127,15 @@ def generate_sample_paulis(self, num_samples): return (samples==1).int(), (samples==2).int(), (samples==3).int(), covers def compute_local_energy(self, x, model): + ''' + Compute estimated local energy values of Hamiltonian, using Pauli string samples, w.r.t. batch of qubit spin configurations and an ansatz model + Args: + x: qubit spin configurations + model: NQS ansatz + Returns: + local_energy: local energy values (detached from computational graph) + log_psi: logarithms of ansatz statevector entries (attached to computational graph) + ''' # see appendix B of https://arxiv.org/pdf/1909.12852.pdf # x [bs, input_dim] bs = x.shape[0] @@ -136,7 +166,10 @@ def compute_local_energy(self, x, model): local_energy = scalar_mult(self.sample_coeffs.unsqueeze(0), scalar_mult(mtx_k, ratio)).sum(1) # [bs, 2] return local_energy.detach(), log_psi - def set_device(self, device): + def set_device(self, device: str): + ''' + Sets device of Hamiltonian instance. + ''' self.coefficients = self.coefficients.to(device) self.select_idx = self.select_idx.to(device) self.unique_flips = self.unique_flips.to(device)