Adding Additional Functionality to Neural Quantum States (NQS) for Quantum Chemistry Repo (#49)

* Added implementations for transformer ansatz with minimal phase function and surrogate local energy Hamiltonian class.

* Addressed requested changes for merging into Tangelo-Examples.

Author: knitterAQ

examples/neural_quantum_states/src/models/alt_transformer.py

@@ -0,0 +1,201 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base import Base
+
+class AltTransformer(Base):
+    def __init__(self, num_sites: int, num_spin_up: int, num_spin_down: int, embedding_dim: int=16, nhead: int=2, dim_feedforward: int=64, num_layers: int=1, temperature: float=1.0, device: str=None, **kwargs):
+        '''
+        A Transformer-based autoregressive NQS Ansatz using the phase strategy from Bennewitz et al, where a single linear layer operators on the concated list of transformer hidden states in lieu of a seperate phase network.
+        Parent class args:
+            num_sites: number of qubits in the ansatz system
+            num_spin_up: total occupancy number of spin-up spin-orbitals
+            num_spin_down: total occupancy number of spin-down spin-orbitals
+            device: Device (CPU or Cuda) to store model
+        Child class specific args:
+            embedding_dim: dimension of transformer hidden states
+            nhead: number of attention heads
+            dim_feedforward: dimension of transformer feedforward layer
+            num_layers: number of transformer blocks
+            temperature: modulus network softmax temperature parameter
+            device: device to store model on
+        '''
+        super(AltTransformer, self).__init__('AltTransformer', num_sites, num_spin_up, num_spin_down, device)
+
+        # construct model
+        self.num_in, self.num_out = num_sites, num_sites*2
+        self.temperature = temperature
+        # Sample function samples spatial orbitals in reverse order, but spin-up orbitals are always sampled first. self.input_order calculates this order for sampling.
+        self.input_order = np.stack([np.arange(self.num_sites-2,-1,-2), np.arange(self.num_sites-1,-1,-2)],1).reshape(-1) # [4,5,2,3,0,1]
+        self.input_order = torch.Tensor(self.input_order).int().to(self.device)
+        # Calculate spatial orbital sampling order
+        self.shell_order = torch.arange(self.num_sites//2-1, -1, -1) # [2,1,0]
+        
+        transformer_layer = nn.TransformerEncoderLayer(embedding_dim, nhead, dim_feedforward=dim_feedforward, dropout=0.0, batch_first=True)
+        self.transformer = nn.TransformerEncoder(transformer_layer, num_layers)
+        self.fc = nn.Linear(embedding_dim, 4)
+        self.tok_emb = nn.Embedding(5, embedding_dim)
+        self.pos_emb = nn.Embedding(len(self.shell_order), embedding_dim)
+        self.apply(self._init_weights)
+        self.softmax = nn.Softmax(dim=-1)
+        self.log_softmax = nn.LogSoftmax(dim=-1)
+
+        self.net_phase = nn.Linear(in_features=embedding_dim*len(self.shell_order), out_features=4, bias=True)
+
+        self.mask = torch.zeros((len(self.shell_order), len(self.shell_order))).to(self.device)
+        for i in range(len(self.mask)):
+            for j in range(len(self.mask)):
+                if i < j:
+                    self.mask[i][j] = float('-inf') 
+        
+    def _init_weights(self, module: nn.Module):
+        '''
+        Performs weight initialization for each module in ansatz, dependent on module type
+        Args:
+            module: module to be initialized
+        '''
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.zeros_(module.bias)
+            torch.nn.init.ones_(module.weight)
+
+    def forward(self, x: torch.Tensor, sample_shell: int = 0) -> torch.Tensor:
+        '''
+        Forward function for Transformer ansatz (used for both sampling and training)
+        Args:
+            x: qubit spin configuration
+            sample_shell: shell index provided during sampling to avoid extraneous forward passes
+        Returns:
+            prob_cond/log_psi: either conditional probabilities of logarithms of statevector entries, depending on if sampling
+        '''
+        # x: [bs, num_sites]
+        shells = self.state2shell(x)[:, self.shell_order]
+        input = 4*torch.ones(shells.shape, dtype=torch.int64).to(self.device) 
+        input[:,1:] = shells[:,:-1]
+        pos = self.shell_order.to(self.device)
+
+        input = self.tok_emb(input) + self.pos_emb(pos)
+        # new x is of shape (batch_size, sequence_length, d_model)
+
+        if self.mask.device != self.device:
+            self.mask = self.mask.to(self.device)
+        output = self.transformer(input[:,:(len(self.shell_order) - sample_shell + 1)], mask=self.mask[:(len(self.shell_order) - sample_shell + 1),:(len(self.shell_order) - sample_shell + 1)], is_causal=True)
+        
+        if not self.sampling:
+            phase_input = output.reshape(output.shape[0], -1)
+        output = self.fc(output)
+        
+        if output.shape[1] < len(self.shell_order):
+            new_output = torch.zeros(output.shape[0], len(self.shell_order), output.shape[2]).to(self.device)
+            new_output[:,:output.shape[1],:] = output
+            output = new_output[:, self.shell_order]
+        else:
+            output = output[:, self.shell_order]
+        
+        if (self.num_spin_up + self.num_spin_down) >= 0:
+            logits_cls = self.apply_constraint(x, output)
+        logits_cls /= self.temperature
+        
+        if self.sampling:
+            prob_cond = self.softmax(logits_cls)
+            return prob_cond
+        else:
+            log_psi_cond = 0.5 * self.log_softmax(logits_cls)
+            idx = self.state2shell(x)
+            log_psi_real = log_psi_cond.gather(-1, idx.unsqueeze(-1)).sum(-1).sum(-1)
+            log_psi_imag = self.net_phase(phase_input).gather(-1, idx[:, -1].unsqueeze(-1)).squeeze()
+            if log_psi_real.shape[0] == 1:
+                log_psi_imag = log_psi_imag.reshape(log_psi_real.shape)
+            log_psi = torch.stack((log_psi_real, log_psi_imag), dim=-1)
+            return log_psi
+
+    def apply_constraint(self, inp: torch.Tensor, log_psi_cond: torch.Tensor) -> torch.Tensor:
+        '''
+        Applies constraints that enforce particle number and spin on ansatz network
+        Args:
+            inp: input spin configurations
+            log_psi_cond: unconstrained ansatz outputs
+        Returns:
+            log_psi_cond: ansatz outputs with constraint applied
+        '''
+        # convert [|-1,-1>, |1,-1>, |-1,1>, |1,1>] to [0, 1, 2, 3]
+        device = inp.device
+        N = inp.shape[-1] // 2
+        inp_up = inp[:, self.input_order][:, ::2].clone()
+        inp_down = inp[:, self.input_order][:, 1::2].clone()
+        inp_cumsum_up = torch.cat((torch.zeros((inp_up.shape[0],1)).to(device), ((1 + inp_up)/2).cumsum(-1)[:, :-1]), axis=-1)
+        inp_cumsum_down = torch.cat((torch.zeros((inp_down.shape[0],1)).to(device), ((1 + inp_down)/2).cumsum(-1)[:, :-1]), axis=-1)
+        upper_bound_up = self.num_spin_up
+        lower_bound_up = (self.num_spin_up - (N - torch.arange(1, N+1)))
+        condition1_up = (inp_cumsum_up < lower_bound_up.to(device)).float()
+        condition2_up = (inp_cumsum_up >= upper_bound_up).float()
+        upper_bound_down = self.num_spin_down
+        lower_bound_down = (self.num_spin_down - (N - torch.arange(1, N+1)))
+        condition1_down = (inp_cumsum_down < lower_bound_down.to(device)).float()
+        condition2_down = (inp_cumsum_down >= upper_bound_down).float()
+        idx = torch.sort(self.shell_order)[1]
+        # first entry must be down
+        log_psi_cond[:,:,0].masked_fill_(condition1_up[:,idx]==1, float('-inf'))
+        log_psi_cond[:,:,2].masked_fill_(condition1_up[:,idx]==1, float('-inf'))
+        # second entry must be down
+        log_psi_cond[:,:,0].masked_fill_(condition1_down[:,idx]==1, float('-inf'))
+        log_psi_cond[:,:,1].masked_fill_(condition1_down[:,idx]==1, float('-inf'))
+        # first entry must be up
+        log_psi_cond[:,:,1].masked_fill_(condition2_up[:,idx]==1, float('-inf'))
+        log_psi_cond[:,:,3].masked_fill_(condition2_up[:,idx]==1, float('-inf'))
+        # second entry must be up
+        log_psi_cond[:,:,2].masked_fill_(condition2_down[:,idx]==1, float('-inf'))
+        log_psi_cond[:,:,3].masked_fill_(condition2_down[:,idx]==1, float('-inf'))
+        return log_psi_cond
+
+    @torch.no_grad()
+    def sample(self, bs: int, num_samples: int) -> [torch.Tensor, torch.Tensor]:
+        '''
+        Generates a set of samples from the ansatz state vector distribution
+        Inputs:
+            bs: total number of unique samples desired
+            num_samples: total number of non-unique samples desired
+        Returns:
+            uniq_samples: unique spin sample set
+            uniq_counts: tensor of count values (summing to num_samples) corresponding with uniq_samples
+        '''
+        self.eval()
+        self.sampling = True
+        sample_multinomial = True
+        # random initialize a configuration of values +- 1
+        uniq_samples = (torch.randn(1, self.num_sites).to(self.device) > 0.0).float() * 2 - 1
+        uniq_count = torch.tensor([num_samples]).to(self.device)
+        for i in self.shell_order:
+            prob = self.forward(uniq_samples, i)[:, i] # num_uniq, 4
+            num_uniq = uniq_samples.shape[0]
+            uniq_samples = uniq_samples.repeat(4,1) # 4*num_uniq, num_sites
+            # convert [|-1,-1>, |1,-1>, |-1,1>, |1,1>] to [0, 1, 2, 3]
+            uniq_samples[:num_uniq, 2*i] = -1
+            uniq_samples[:num_uniq, 2*i+1] = -1
+            uniq_samples[num_uniq:2*num_uniq, 2*i] = 1
+            uniq_samples[num_uniq:2*num_uniq, 2*i+1] = -1
+            uniq_samples[2*num_uniq:3*num_uniq, 2*i] = -1
+            uniq_samples[2*num_uniq:3*num_uniq, 2*i+1] = 1
+            uniq_samples[3*num_uniq:4*num_uniq, 2*i] = 1
+            uniq_samples[3*num_uniq:4*num_uniq, 2*i+1] = 1
+            if sample_multinomial:
+                uniq_count = torch.tensor(self.multinomial_arr(uniq_count.long().data.cpu().numpy(), prob.data.cpu().numpy())).T.flatten().to(prob.device)
+            else:
+                uniq_count = (uniq_count.unsqueeze(-1)*prob).T.flatten().round()
+            keep_idx = uniq_count > 0
+            uniq_samples = uniq_samples[keep_idx]
+            uniq_count = uniq_count[keep_idx]
+            uniq_samples = uniq_samples[uniq_count.sort()[1][-2*bs:]]
+            uniq_count = uniq_count[uniq_count.sort()[1][-2*bs:]]
+        uniq_samples = uniq_samples[uniq_count.sort()[1][-bs:]]
+        uniq_count = uniq_count[uniq_count.sort()[1][-bs:]]
+        self.sampling = False
+        self.train()
+        return [uniq_samples, uniq_count]

examples/neural_quantum_states/src/models/base.py

@@ -9,14 +9,16 @@ class Base(nn.Module):
     '''
     Base template for all autoregressive NQS ansatze.
     Args:
+        name: name of specific model type
         num_sites: qubit number
         num_spin_up: number of spin up electrons
         num_spin_down: number of spin down electrons
         device: Device (CPU or Cuda) to store model
         **kwargs: nonspecific kwargs
     '''
-    def __init__(self, num_sites: int, num_spin_up: int, num_spin_down: int, device: str, **kwargs):
+    def __init__(self, name: str, num_sites: int, num_spin_up: int, num_spin_down: int, device: str, **kwargs):
         super().__init__()
+        self.name = name
         self.num_sites = num_sites
         self.num_spin_up = num_spin_up
         self.num_spin_down = num_spin_down
@@ -154,6 +156,9 @@ def get_model(model_name: str, device: str, print_model_info: bool, **kwargs) ->
     elif model_name == 'retnet':
         from .retnet import NNQSRetNet
         model = NNQSRetNet(**kwargs)
+    elif model_name == 'alt_transformer':
+        from .alt_transformer import AltTransformer
+        model = AltTransformer(**kwargs)
     else:
         raise ValueError(f"Unknown model_name: {model_name}")
     if print_model_info:

examples/neural_quantum_states/src/models/made.py

@@ -8,14 +8,19 @@
 class MADE(Base):
     '''
     Class implements MADE-based NQS ansatz
+    Parent class args:
+        num_sites: number of qubits in the ansatz system
+        num_spin_up: total occupancy number of spin-up spin-orbitals
+        num_spin_down: total occupancy number of spin-down spin-orbitals
+        device: Device (CPU or Cuda) to store model
     Child class specific args:
         made_width: width of modulus and phase network hidden layers
         made_depth: number of hidden layers in modulus and phase networks
         temperature: Temperature variable for modulus softmax
         **kwargs: nonspecific kwargs
     '''
     def __init__(self, num_sites: int, num_spin_up: int, num_spin_down: int, made_width: int=64, made_depth: int=2, temperature: float=1.0, device: str=None, **kwargs):
-        super(MADE, self).__init__(num_sites, num_spin_up, num_spin_down, device)
+        super(MADE, self).__init__('MADE', num_sites, num_spin_up, num_spin_down, device)
         self.temperature = temperature
         # construct model
         self.net = []

examples/neural_quantum_states/src/models/retnet.py

@@ -10,6 +10,11 @@ class NNQSRetNet(Base):
     def __init__(self, num_sites: int, num_spin_up: int, num_spin_down: int, made_width: int=64, made_depth: int=2, embedding_dim: int=16, nhead: int=2, dim_feedforward: int=64, num_layers: int=1, temperature: float=1.0, device: str=None, **kwargs):
         '''
         Retentive network (RetNet) NQS ansatz
+        Parent class args:
+            num_sites: number of qubits in the ansatz system
+            num_spin_up: total occupancy number of spin-up spin-orbitals
+            num_spin_down: total occupancy number of spin-down spin-orbitals
+            device: Device (CPU or Cuda) to store model
         Child class specific args:
             made_width: width of phase network hidden layers
             made_depth: number of phase network hidden layers
@@ -20,7 +25,7 @@ def __init__(self, num_sites: int, num_spin_up: int, num_spin_down: int, made_wi
             temperature: RetNet softmax temperature parameter
             device: device on which the model is stored
         '''
-        super(NNQSRetNet, self).__init__(num_sites, num_spin_up, num_spin_down, device)
+        super(NNQSRetNet, self).__init__('RetNet', num_sites, num_spin_up, num_spin_down, device)
 
         # construct model
         self.num_in, self.num_out = num_sites, num_sites*2

examples/neural_quantum_states/src/models/transformer.py

@@ -9,6 +9,11 @@ class NNQSTransformer(Base):
     def __init__(self, num_sites: int, num_spin_up: int, num_spin_down: int, made_width: int=64, made_depth: int=2, embedding_dim: int=16, nhead: int=2, dim_feedforward: int=64, num_layers: int=1, temperature: float=1.0, device: str=None, **kwargs):
         '''
         A Transformer-based autoregressive NQS Ansatz
+        Parent class args:
+            num_sites: number of qubits in the ansatz system
+            num_spin_up: total occupancy number of spin-up spin-orbitals
+            num_spin_down: total occupancy number of spin-down spin-orbitals
+            device: Device (CPU or Cuda) to store model
         Child class specific args:
             made_width: width of phase network hidden layers
             made_depth: number of phase network hidden layers
@@ -19,7 +24,7 @@ def __init__(self, num_sites: int, num_spin_up: int, num_spin_down: int, made_wi
             temperature: modulus network softmax temperature parameter
             device: device to store model on
         '''
-        super(NNQSTransformer, self).__init__(num_sites, num_spin_up, num_spin_down, device)
+        super(NNQSTransformer, self).__init__('NNQSTransformer', num_sites, num_spin_up, num_spin_down, device)
 
         # construct model
         self.num_in, self.num_out = num_sites, num_sites*2

examples/neural_quantum_states/src/objective/adaptive_shadows.py

@@ -19,7 +19,7 @@ class AdaptiveShadows(Hamiltonian):
         flip_bs: Number of unique bit flip patterns processed at a time on each GPU
     '''
     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)
+        super(AdaptiveShadows, self).__init__('adaptive_shadows', 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)
         self.coefficients_square = norm_square(self.coefficients)

examples/neural_quantum_states/src/objective/automatic.py

@@ -14,7 +14,7 @@ def __init__(self, hamiltonian_string: str, num_sites: int, flip_bs: int, **kwar
             num_sites: qubit number of system
             flip_bs: largest batch size of model input tensors that each GPU is expected to handle at once
         '''
-        super(Automatic, self).__init__(hamiltonian_string, num_sites)
+        super(Automatic, self).__init__('automatic', 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)
         self.flip_bs = flip_bs

examples/neural_quantum_states/src/objective/hamiltonian.py

@@ -4,8 +4,9 @@
 from src.complex import scalar_mult, real, imag
 
 class Hamiltonian(nn.Module):
-    def __init__(self, hamiltonian_string, num_sites):
+    def __init__(self, choice, hamiltonian_string, num_sites):
         super().__init__()
+        self.name = choice
         self.operators, self.coefficients = self.parse_hamiltonian_string(hamiltonian_string, num_sites)
         self.num_terms, self.input_dim = self.operators.shape
         print("Number of terms is {}.".format(self.num_terms))
@@ -83,6 +84,9 @@ def get_hamiltonian(hamiltonian_choice: str, hamiltonian_data: dict) -> nn.Modul
     elif hamiltonian_choice in ['exact']:
         from .automatic import Automatic
         return Automatic(**hamiltonian_data)
+    elif hamiltonian_choice in ['surrogate']:
+        from .surrogate import Surrogate
+        return Surrogate(**hamiltonian_data)
     else:
         raise Exception('Hamiltonian choice not recognized!')
 

examples/neural_quantum_states/src/objective/naive_sampler.py

@@ -10,7 +10,7 @@ def __init__(self, hamiltonian_string: str, num_sites: int, sample_count: int, t
         '''
         A variation of the Automatic class that stochastically estimates the input Hamiltonian with Pauli strings sampled from the distribution proportional to the absolute values of the scalar coefficients (simple to construct because the Pauli string coefficients are real for Hamiltonians). This estimated Hamiltonian can be used to create local energy estimates during NQS training for (ideally) lower computational cost.
         '''
-        super(NaiveSampler, self).__init__(hamiltonian_string, num_sites)
+        super(NaiveSampler, self).__init__('naive_sampler', hamiltonian_string, num_sites)
         self.flip_bs = flip_bs
         # product of identity operators by default, encoded as 0
         self.coefficients = torch.stack((self.coefficients.real, self.coefficients.imag), dim=-1)