Added MPS codes

This commit is contained in:
tankya2
2023-07-14 09:51:06 +08:00
parent 76f61bc9fe
commit 3cb0fec99c
4 changed files with 278 additions and 57 deletions

View 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})

78
src/qibotn/MPSUtils.py Normal file
View File

@@ -0,0 +1,78 @@
import cupy as cp
from cuquantum.cutensornet.experimental import contract_decompose
from cuquantum import contract
def get_initial_mps(num_qubits, dtype='complex128'):
"""
Generate the MPS with an initial state of |00...00>
"""
state_tensor = cp.asarray([1, 0], dtype=dtype).reshape(1,2,1)
mps_tensors = [state_tensor] * num_qubits
return mps_tensors
def mps_site_right_swap(
mps_tensors,
i,
algorithm=None,
options=None
):
"""
Perform the swap operation between the ith and i+1th MPS tensors.
"""
# contraction followed by QR decomposition
a, _, b = contract_decompose('ipj,jqk->iqj,jpk', *mps_tensors[i:i+2], algorithm=algorithm, options=options)
mps_tensors[i:i+2] = (a, b)
return mps_tensors
def apply_gate(
mps_tensors,
gate,
qubits,
algorithm=None,
options=None
):
"""
Apply the gate operand to the MPS tensors in-place.
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.
gate: A ndarray-like tensor object representing the gate operand.
The modes of the gate is 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 denoting the qubits that the gate is applied onto.
algorithm: The contract and decompose algorithm to use for gate application.
Can be either a `dict` or a `ContractDecomposeAlgorithm`.
options: Specify the contract and decompose options.
Returns:
The updated MPS tensors.
"""
n_qubits = len(qubits)
if n_qubits == 1:
# single-qubit gate
i = qubits[0]
mps_tensors[i] = contract('ipj,qp->iqj', mps_tensors[i], gate, options=options) # in-place update
elif n_qubits == 2:
# two-qubit gate
i, j = qubits
if i > j:
# swap qubits order
return apply_gate(mps_tensors, gate.transpose(1,0,3,2), (j, i), algorithm=algorithm, options=options)
elif i+1 == j:
# two adjacent qubits
a, _, b = contract_decompose('ipj,jqk,rspq->irj,jsk', *mps_tensors[i:i+2], gate, algorithm=algorithm, options=options)
mps_tensors[i:i+2] = (a, b) # in-place update
else:
# non-adjacent two-qubit gate
# step 1: swap i with i+1
mps_site_right_swap(mps_tensors, i, algorithm=algorithm, options=options)
# step 2: apply gate to (i+1, j) pair. This amounts to a recursive swap until the two qubits are adjacent
apply_gate(mps_tensors, gate, (i+1, j), algorithm=algorithm, options=options)
# step 3: swap back i and i+1
mps_site_right_swap(mps_tensors, i, algorithm=algorithm, options=options)
else:
raise NotImplementedError("Only one- and two-qubit gates supported")
return mps_tensors

View File

@@ -0,0 +1,24 @@
import cupy as cp
import numpy as np
from cuquantum import cutensornet as cutn
from QiboCircuitConvertor import QiboCircuitToEinsum
from MPSUtils import get_initial_mps, apply_gate
class QiboCircuitToMPS:
def __init__(self,circ_qibo, gate_algo, dtype = 'complex128',rand_seed=0,):
np.random.seed(rand_seed)
cp.random.seed(rand_seed)
self.num_qubits = circ_qibo.nqubits
self.handle = cutn.create()
self.options = {'handle': self.handle}
self.dtype = dtype
self.mps_tensors = get_initial_mps(self.num_qubits, dtype=dtype)
circuitconvertor = QiboCircuitToEinsum(circ_qibo)
for (gate, qubits) in circuitconvertor.gate_tensors:
# mapping from qubits to qubit indices
# apply the gate in-place
apply_gate(self.mps_tensors, gate, qubits, algorithm=gate_algo, options=self.options)

View File

@@ -5,68 +5,32 @@ import cuquantum
from cuquantum import cutensornet as cutn
import cupy as cp
import numpy as np
from qibo.models import QFT
from QiboCircuitToMPS import QiboCircuitToMPS
from MPSContractionHelper import MPSContractionHelper
def eval(qibo_circ, datatype):
myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype)
return contract(*myconvertor.state_vector_operands())
def eval_mps(qibo_circ, datatype):
#Create MPS
cutensornet.create()
return contract()
def eval_mps(qibo_circ, gate_algo, datatype):
myconvertor = QiboCircuitToMPS(qibo_circ, gate_algo, datatype)
mps_helper = MPSContractionHelper(myconvertor.num_qubits)
sv_mps = mps_helper.contract_state_vector(myconvertor.mps_tensors,myconvertor.options)
return sv_mps
if __name__ == "__main__":
print("cuTensorNet-vers:", cutn.get_version())
dev = cp.cuda.Device() # get current device
props = cp.cuda.runtime.getDeviceProperties(dev.id)
print("===== device info ======")
print("GPU-name:", props["name"].decode())
print("GPU-clock:", props["clockRate"])
print("GPU-memoryClock:", props["memoryClockRate"])
print("GPU-nSM:", props["multiProcessorCount"])
print("GPU-major:", props["major"])
print("GPU-minor:", props["minor"])
print("========================")
data_type = cuquantum.cudaDataType.CUDA_C_64F
compute_type = cuquantum.ComputeType.COMPUTE_64F
num_sites = 16
phys_extent = 2
max_virtual_extent = 12
## we initialize the MPS state as a product state |000...000>
initial_state = []
for i in range(num_sites):
# we create dummpy indices for MPS tensors on the boundary for easier bookkeeping
# we'll use Fortran layout throughout this example
tensor = cp.zeros((1,2,1), dtype=np.complex128, order="F")
tensor[0,0,0] = 1.0
initial_state.append(tensor)
mps_helper = MPSHelper(num_sites, phys_extent, max_virtual_extent, initial_state, data_type, compute_type)
##################################
# Setup options for gate operation
##################################
abs_cutoff = 1e-2
rel_cutoff = 1e-2
renorm = cutn.TensorSVDNormalization.L2
partition = cutn.TensorSVDPartition.UV_EQUAL
mps_helper.set_svd_config(abs_cutoff, rel_cutoff, renorm, partition)
gate_algo = cutn.GateSplitAlgo.REDUCED
mps_helper.set_gate_algorithm(gate_algo)
#####################################
# Workspace estimation and allocation
#####################################
free_mem, total_mem = dev.mem_info
worksize = free_mem *.7
required_workspace_size = mps_helper.compute_max_workspace_sizes()
work = cp.cuda.alloc(worksize)
print(f"Maximal workspace size requried: {required_workspace_size / 1024 ** 3:.3f} GB")
mps_helper.set_workspace(work, required_workspace_size)
num_qubits = 25
swaps = True
circ_qibo = QFT(num_qubits, swaps)
exact_gate_algorithm = {'qr_method': False,
'svd_method':{'partition': 'UV', 'abs_cutoff':1e-12}}
dtype="complex128"
sv_mps = eval_mps(circ_qibo, exact_gate_algorithm, dtype)
sv_reference = eval(circ_qibo, dtype)
state_vec = np.array(circ_qibo())
print(f"State vector difference: {abs(sv_mps-sv_reference).max():0.3e}")
assert cp.allclose(sv_mps, sv_reference)
assert cp.allclose(sv_mps.flatten(), state_vec)