Added MPS codes
This commit is contained in:
tankya2
155
src/qibotn/MPSContractionHelper.py
Normal file
155
src/qibotn/MPSContractionHelper.py
Normal file
@@ -0,0 +1,155 @@
|
||||
from cuquantum import contract, contract_path, CircuitToEinsum, tensor
|
||||
|
||||
class MPSContractionHelper:
|
||||
"""
|
||||
A helper class to compute various quantities for a given MPS.
|
||||
|
||||
Interleaved format is used to construct the input args for `cuquantum.contract`.
|
||||
A concrete example on how the modes are populated for a 7-site MPS is provided below:
|
||||
|
||||
0 2 4 6 8 10 12 14
|
||||
bra -----A-----B-----C-----D-----E-----F-----G-----
|
||||
| | | | | | |
|
||||
1| 3| 5| 7| 9| 11| 13|
|
||||
| | | | | | |
|
||||
ket -----a-----b-----c-----d-----e-----f-----g-----
|
||||
15 16 17 18 19 20 21 22
|
||||
|
||||
|
||||
The follwing compute quantities are supported:
|
||||
|
||||
- the norm of the MPS.
|
||||
- the equivalent state vector from the MPS.
|
||||
- the expectation value for a given operator.
|
||||
- the equivalent state vector after multiplying an MPO to an MPS.
|
||||
|
||||
Note that for the nth MPS tensor (rank-3), the modes of the tensor are expected to be `(i,p,j)`
|
||||
where i denotes the bonding mode with the (n-1)th tensor, p denotes the physical mode for the qubit and
|
||||
j denotes the bonding mode with the (n+1)th tensor.
|
||||
|
||||
Args:
|
||||
num_qubits: The number of qubits for the MPS.
|
||||
"""
|
||||
|
||||
def __init__(self, num_qubits):
|
||||
self.num_qubits = num_qubits
|
||||
self.path_cache = dict()
|
||||
self.bra_modes = [(2*i, 2*i+1, 2*i+2) for i in range(num_qubits)]
|
||||
offset = 2*num_qubits+1
|
||||
self.ket_modes = [(i+offset, 2*i+1, i+1+offset) for i in range(num_qubits)]
|
||||
|
||||
def contract_norm(self, mps_tensors, options=None):
|
||||
"""
|
||||
Contract the corresponding tensor network to form the norm of the MPS.
|
||||
|
||||
Args:
|
||||
mps_tensors: A list of rank-3 ndarray-like tensor objects.
|
||||
The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
|
||||
the physical mode, and then the bonding index to the i+1th tensor.
|
||||
options: Specify the contract and decompose options.
|
||||
|
||||
Returns:
|
||||
The norm of the MPS.
|
||||
"""
|
||||
interleaved_inputs = []
|
||||
for i, o in enumerate(mps_tensors):
|
||||
interleaved_inputs.extend([o, self.bra_modes[i], o.conj(), self.ket_modes[i]])
|
||||
interleaved_inputs.append([]) # output
|
||||
return self._contract('norm', interleaved_inputs, options=options).real
|
||||
|
||||
def contract_state_vector(self, mps_tensors, options=None):
|
||||
"""
|
||||
Contract the corresponding tensor network to form the state vector representation of the MPS.
|
||||
|
||||
Args:
|
||||
mps_tensors: A list of rank-3 ndarray-like tensor objects.
|
||||
The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
|
||||
the physical mode, and then the bonding index to the i+1th tensor.
|
||||
options: Specify the contract and decompose options.
|
||||
|
||||
Returns:
|
||||
An ndarray-like object as the state vector.
|
||||
"""
|
||||
interleaved_inputs = []
|
||||
for i, o in enumerate(mps_tensors):
|
||||
interleaved_inputs.extend([o, self.bra_modes[i]])
|
||||
output_modes = tuple([bra_modes[1] for bra_modes in self.bra_modes])
|
||||
interleaved_inputs.append(output_modes) # output
|
||||
return self._contract('sv', interleaved_inputs, options=options)
|
||||
|
||||
def contract_expectation(self, mps_tensors, operator, qubits, options=None, normalize=False):
|
||||
"""
|
||||
Contract the corresponding tensor network to form the state vector representation of the MPS.
|
||||
|
||||
Args:
|
||||
mps_tensors: A list of rank-3 ndarray-like tensor objects.
|
||||
The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
|
||||
the physical mode, and then the bonding index to the i+1th tensor.
|
||||
operator: A ndarray-like tensor object.
|
||||
The modes of the operator are expected to be output qubits followed by input qubits, e.g,
|
||||
``A, B, a, b`` where `a, b` denotes the inputs and `A, B'` denotes the outputs.
|
||||
qubits: A sequence of integers specifying the qubits that the operator is acting on.
|
||||
options: Specify the contract and decompose options.
|
||||
normalize: Whether to scale the expectation value by the normalization factor.
|
||||
|
||||
Returns:
|
||||
An ndarray-like object as the state vector.
|
||||
"""
|
||||
|
||||
interleaved_inputs = []
|
||||
extra_mode = 3 * self.num_qubits + 2
|
||||
operator_modes = [None] * len(qubits) + [self.bra_modes[q][1] for q in qubits]
|
||||
qubits = list(qubits)
|
||||
for i, o in enumerate(mps_tensors):
|
||||
interleaved_inputs.extend([o, self.bra_modes[i]])
|
||||
k_modes = self.ket_modes[i]
|
||||
if i in qubits:
|
||||
k_modes = (k_modes[0], extra_mode, k_modes[2])
|
||||
q = qubits.index(i)
|
||||
operator_modes[q] = extra_mode # output modes
|
||||
extra_mode += 1
|
||||
interleaved_inputs.extend([o.conj(), k_modes])
|
||||
interleaved_inputs.extend([operator, tuple(operator_modes)])
|
||||
interleaved_inputs.append([]) # output
|
||||
if normalize:
|
||||
norm = self.contract_norm(mps_tensors, options=options)
|
||||
else:
|
||||
norm = 1
|
||||
return self._contract(f'exp{qubits}', interleaved_inputs, options=options) / norm
|
||||
|
||||
def contract_mps_mpo_to_state_vector(self, mps_tensors, mpo_tensors, options=None):
|
||||
"""
|
||||
Contract the corresponding tensor network to form the output state vector from applying the MPO to the MPS.
|
||||
|
||||
Args:
|
||||
mps_tensors: A list of rank-3 ndarray-like tensor objects.
|
||||
The indices of the ith tensor are expected to be the bonding index to the i-1 tensor,
|
||||
the physical mode, and then the bonding index to the i+1th tensor.
|
||||
mpo_tensors: A list of rank-4 ndarray-like tensor objects.
|
||||
The indics of the ith tensor are expected to be the bonding index to the i-1 tensor,
|
||||
the output physical mode, the bonding index to the i+1th tensor and then the inputput physical mode.
|
||||
options: Specify the contract and decompose options.
|
||||
|
||||
Returns:
|
||||
An ndarray-like object as the output state vector.
|
||||
"""
|
||||
interleaved_inputs = []
|
||||
for i, o in enumerate(mps_tensors):
|
||||
interleaved_inputs.extend([o, self.bra_modes[i]])
|
||||
output_modes = []
|
||||
offset = 2 * self.num_qubits + 1
|
||||
for i, o in enumerate(mpo_tensors):
|
||||
mpo_modes = (2*i+offset, 2*i+offset+1, 2*i+offset+2, 2*i+1)
|
||||
output_modes.append(2*i+offset+1)
|
||||
interleaved_inputs.extend([o, mpo_modes])
|
||||
interleaved_inputs.append(output_modes)
|
||||
return self._contract('mps_mpo', interleaved_inputs, options=options)
|
||||
|
||||
def _contract(self, key, interleaved_inputs, options=None):
|
||||
"""
|
||||
Perform the contraction task given interleaved inputs. Path will be cached.
|
||||
"""
|
||||
if key not in self.path_cache:
|
||||
self.path_cache[key] = contract_path(*interleaved_inputs, options=options)[0]
|
||||
path = self.path_cache[key]
|
||||
return contract(*interleaved_inputs, options=options, optimize={'path':path})
|
||||
Reference in New Issue
Block a user