def test_deep_remove(): l = [[[0], 1], 2] l = deep_remove(l, 1) assert_equal(l, [[[0]], 2]) # Harder case... l = [[[[0, 1, 2]], [3, 4], [5], [6, 7]]] l = deep_remove(l, 0, 5) assert l == [[[[1, 2]], [3, 4], [], [6, 7]]]
def tensor_contract(qobj, *pairs): """Contracts a qobj along one or more index pairs. Note that this uses dense representations and thus should *not* be used for very large Qobjs. Parameters ---------- pairs : tuple One or more tuples ``(i, j)`` indicating that the ``i`` and ``j`` dimensions of the original qobj should be contracted. Returns ------- cqobj : Qobj The original Qobj with all named index pairs contracted away. """ # Record and label the original dims. dims = qobj.dims dims_idxs = enumerate_flat(dims) tensor_dims = dims_to_tensor_shape(dims) # Convert to dense first, since sparse won't support the reshaping we need. qtens = qobj.data.toarray() # Reshape by the flattened dims. qtens = qtens.reshape(tensor_dims) # Contract out the indices from the flattened object. # Note that we need to feed pairs through dims_idxs_to_tensor_idxs # to ensure that we are contracting the right indices. qtens = _tensor_contract_dense(qtens, *dims_idxs_to_tensor_idxs(dims, pairs)) # Remove the contracted indexes from dims so we know how to # reshape back. # This concerns dims, and not the tensor indices, so we need # to make sure to use the original dims indices and not the ones # generated by dims_to_* functions. contracted_idxs = deep_remove(dims_idxs, *flatten(list(map(list, pairs)))) contracted_dims = unflatten(flatten(dims), contracted_idxs) # We don't need to check for tensor idxs versus dims idxs here, # as column- versus row-stacking will never move an index for the # vectorized operator spaces all the way from the left to the right. l_mtx_dims, r_mtx_dims = map(np.product, map(flatten, contracted_dims)) # Reshape back into a 2D matrix. qmtx = qtens.reshape((l_mtx_dims, r_mtx_dims)) # Return back as a qobj. return Qobj(qmtx, dims=contracted_dims, superrep=qobj.superrep)
def test_deep_remove(base, to_remove, expected): assert deep_remove(base, *to_remove) == expected