Fixed changes for Tangelo-Examples PR.

examples/neural_quantum_states/.gitignore

@@ -0,0 +1,3 @@
+logger/
+run.sh
+**/__pycache__/

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, [email protected]
 
-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.
 

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)

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,13 +53,19 @@ 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)
         self.unique_flips = 1 - 2*self.unique_flips.unsqueeze(0)
         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)