R/052_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(.shape(expr)) < 2L || .shape(expr)[1L] != .shape(expr)[2L]) {
    cli_abort("{.fn partial_trace} only supports 2-d square arrays.")
  }
  if (.shape(expr)[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)
  }
  result <- Reduce("+", terms)

  ## Propagate the argument's matrix structure to the result.
  ## CVXPY SOURCE: partial_trace.py:87-93 (CVXPY 1.9.2).
  ## Partial trace PRESERVES positive semidefiniteness, symmetry and
  ## Hermitian-ness, but the sum of Constant %*% expr %*% Constant terms
  ## cannot infer that on its own -- without the wrap every attribute came
  ## back FALSE and DCP analysis lost the PSD-ness of a PSD input.
  ## Order matters: PSD is the strongest claim, so it is tested first.
  if (is_psd(expr)) return(psd_wrap(result))
  if (is_symmetric(expr)) return(symmetric_wrap(result))
  if (is_hermitian(expr)) return(hermitian_wrap(result))
  result
}

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