R/046_atoms_affine_partial_trace.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/atoms/affine/partial_trace.R
#####

## CVXPY SOURCE: atoms/affine/partial_trace.py
## partial_trace -- trace out a subsystem from a tensor product expression
##
## This is a factory function (not an Atom class). It constructs an Expression
## by summing (I kron <j| kron I) X (I kron |j> kron I) over j in the traced-out
## subsystem.

#' Partial trace of a tensor product expression
#'
#' Assumes `expr` is a 2D square matrix representing a Kronecker product
#' of `length(dims)` subsystems. Returns the partial trace over the
#' subsystem at index `axis` (1-indexed).
#'
#' @param expr An Expression (2D square matrix)
#' @param dims Integer vector of subsystem dimensions
#' @param axis Integer (1-indexed) subsystem to trace out
#' @returns An Expression representing the partial trace
#' @export
partial_trace <- function(expr, dims, axis = 1L) {
  expr <- as_expr(expr)
  if (length(expr@shape) < 2L || expr@shape[1L] != expr@shape[2L]) {
    cli_abort("{.fn partial_trace} only supports 2-d square arrays.")
  }
  if (expr@shape[1L] != prod(dims)) {
    cli_abort("Dimension of system doesn't correspond to dimension of subsystems.")
  }
  ## Normalize axis (handle negative indices, 1-indexed)
  n_sys <- length(dims)
  if (axis < 0L) axis <- axis + n_sys + 1L
  if (axis < 1L || axis > n_sys) {
    cli_abort("axis {axis} is out of bounds for {n_sys} subsystems.")
  }

  ## Build each term: (I kron <j| kron I) expr (I kron |j> kron I)
  ## and sum over j = 0, ..., dims[axis]-1
  terms <- vector("list", dims[axis])
  for (j in seq_len(dims[axis]) - 1L) {  ## 0-indexed j
    terms[[j + 1L]] <- .partial_trace_term(expr, j, dims, axis)
  }
  Reduce("+", terms)
}

## Helper: compute jth term in the partial trace sum
## CVXPY: partial_trace.py _term()
.partial_trace_term <- function(expr, j, dims, axis) {
  ## a = (I kron <j| kron I), b = (I kron |j> kron I)
  ## Start with 1x1 identity
  a <- Matrix::sparseMatrix(i = 1L, j = 1L, x = 1.0, dims = c(1L, 1L))
  b <- Matrix::sparseMatrix(i = 1L, j = 1L, x = 1.0, dims = c(1L, 1L))
  for (i_axis in seq_along(dims)) {
    dim_k <- dims[i_axis]
    if (i_axis == axis) {  ## both 1-indexed
      ## v = |j> (column vector with 1 at position j, 0-indexed)
      v <- Matrix::sparseMatrix(i = j + 1L, j = 1L, x = 1.0,
                                dims = c(dim_k, 1L))
      a <- kronecker(a, t(v))   ## a kron <j| (row vector)
      b <- kronecker(b, v)      ## b kron |j> (column vector)
    } else {
      eye_k <- Matrix::Diagonal(dim_k)
      a <- kronecker(a, eye_k)
      b <- kronecker(b, eye_k)
    }
  }
  ## a @ expr @ b -- a and b are numeric sparse matrices
  Constant(a) %*% expr %*% Constant(b)
}