home/notes/product-space-states

Working with product-space-states

We’d like to build up a computational framework for working with states in a general Hilbert space comprising all product states assembled from a number of subsystems that may be independent or entangled. Let’s start with some definitions.

Suppose we have subsystems

where each is spanned by basis vectors labeled by a quantum number (where denotes the set of all possible values of ).

The complete Hilbert space is, of course, then

for which . And, given some , we are interested in recovering the states of the individual subsystems. Since they may be entangled, we need to work with states as density matrices . The partial trace operation to pare off subsystems is easy to describe with kets and bras. To trace out the -th subsystem, one computes

But when doing this with a program, it is often much more convenient to represent the subsystem vectors more simply as -element arrays (rather than arrays whose length is equal to the total dimension ). As such, these arrays’ dimensions won’t agree with those of .

To remedy this, we have to find a way to build non-square matrices from the kets and bras of the subsystems we’re tracing over, since we need this operation to take our original matrix down to a matrix. It turns out that all we need to do is insert some identities.

In the above expression for , the single kets and bras are being used as linear operators that select out particular components of . Because they’re not ket-bras, their action as operators is to perform a projection of the subspace onto the trivial vector space , while leaving all other subspaces untouched. This means they’re all implicitly tensored with the identity operators for every other subspace–because computers don’t do so well with implicit things, we have to make them explicit in our computer program and build up to the proper projection operators on the full space:

where is said identity on subspace . Then the desired partial trace operation is accomplished as

Here’s some quick-and-dirty Python for qubits:

import numpy as np

type Dens = np.ndarray[(..., ...), complex]

# qubit |0⟩ state
q0 = np.array([[1, 0]]).T
# qubit |1⟩ state
q1 = np.array([[0, 1]]).T

# trinary kronecker product
kron3 = lambda a, b, c: np.kron(np.kron(a, b), c)

def trace_qubit(n: int, k: int, rho: Dens) -> Dens
    """
    Compute the partial trace over the `k`-th qubit out of `n` total qubits for
    a given density matrix `rho`.
    """
    T0 = kron3(np.eye(2 ** k), q0, np.eye(2 ** (n - k - 1)))
    T1 = kron3(np.eye(2 ** k), q1, np.eye(2 ** (n - k - 1)))
    return (T0.T @ rho @ T0) + (T1.T @ rho @ T1)

def isolate_qubit(n: int, k: int, rho: Dens) -> Dens:
    """
    Compute the reduced density matrix for only the `k`-th qubit out of `n`
    total qubits.
    """
    if n == 1:
        return rho
    elif k == 0:
        return isolate_qubit(n - 1, k, trace_qubit(n, n - 1, rho))
    else:
        return isolate_qubit(n - 1, k - 1, trace_qubit(n, 0, rho))

But this is also pretty computationally awkward. At the end of the day, we still have to store an exponentially large array with which we conjugate our density matrix. Each “tracer” is a sparse array with only 1′s and 0′s and while we could cache computations of each in a program to save time with the tensor products, we then still have to carry them around and do large matrix multiplication.

A better (and perhaps altogether more physical) way is to treat our base object as a tensor, specifically one with exactly two indices (one on the ket side and one on the bra side) for each subsystem. This turns our matrix into a many-dimensional array of numbers where each gives the density matrix element. If we have some logical layout (e.g. row-major or column-major) for the data giving the original matrix representing , we can implement this conversion as a simple reshaping operation,

# tuple of #Q_k subspace dimensions
DIMS: tuple[int, ...] = (...)

# original rho as a proper matrix (two-dimensional array)
rho: np.ndarray[tuple[int, int], np.dtype[complex]] = np.array(...)

# rho as a tensor (N-dimensional array)
rho_tens: np.ndarray[tuple[int, ..., int], np.dtype[complex]] = (
  rho.reshape((*DIMS, *DIMS))
)

Hence, we have a way to “talk to” each quantum number individually–it’s just the corresponding axis of rho_tens. Then since the partial trace really just picks out certain elements of (i.e. without any additional scaling or association between distinct subsystems), we can implement it as a simple sum over the axes associated with and :

# we'll trace out the `k`-th subsystem
k: int = ...
dim_k: int = DIMS[k]
dims_no_k = (*(DIMS[:k]), *(DIMS[k + 1:]))

# original rho as a proper matrix (two-dimensional array)
rho: np.ndarray[tuple[int, int], np.dtype[complex]] = np.array(...)

# slicing object: unconstrained slice for subsystems `j != k` and quantum
# number array for `j == k`
sl = tuple([slice() if j != k else np.arange(dim_k) for j in range(len(DIMS))])

rho_no_k = (
  rho.reshape((*DIMS, *DIMS))        # convert to a tensor
  [(*sl, *sl)]                       # slicing moves index `k` to the 0-th axis
  .sum(axis=0)                       # perform the trace
  .reshape((*dims_no_k, *dims_no_k)) # convert back to a matrix
)