119 lines
5.0 KiB
Python
119 lines
5.0 KiB
Python
from cuquantum import contract, contract_path
|
|
|
|
# Reference: https://github.com/NVIDIA/cuQuantum/blob/main/python/samples/cutensornet/tn_algorithms/mps_algorithms.ipynb
|
|
|
|
|
|
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`.
|
|
|
|
Reference: https://github.com/NVIDIA/cuQuantum/blob/main/python/samples/cutensornet/tn_algorithms/mps_algorithms.ipynb
|
|
|
|
The following 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.
|
|
|
|
Parameters:
|
|
num_qubits: The number of qubits for the MPS.
|
|
"""
|
|
|
|
def __init__(self, num_qubits):
|
|
self.num_qubits = num_qubits
|
|
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.
|
|
|
|
Parameters:
|
|
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(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.
|
|
|
|
Parameters:
|
|
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(interleaved_inputs, options=options)
|
|
|
|
def contract_expectation(
|
|
self, mps_tensors, operator, qubits, options=None, normalize=False
|
|
):
|
|
"""Contract the corresponding tensor network to form the expectation of
|
|
the MPS.
|
|
|
|
Parameters:
|
|
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(interleaved_inputs, options=options) / norm
|
|
|
|
def _contract(self, interleaved_inputs, options=None):
|
|
path = contract_path(*interleaved_inputs, options=options)[0]
|
|
|
|
return contract(*interleaved_inputs, options=options, optimize={"path": path})
|