R/047_atoms_affine_partial_transpose.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/affine/partial_transpose.py
## partial_transpose -- transpose a subsystem in a tensor product expression
##
## This is a factory function (not an Atom class). It constructs an Expression
## by summing (I kron |i><j| kron I) X (I kron |i><j| kron I) over all i,j
## in the transposed subsystem.

#' Partial transpose of a tensor product expression
#'
#' Assumes `expr` is a 2D square matrix representing a Kronecker product
#' of `length(dims)` subsystems. Returns the partial transpose with
#' the transpose applied to 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 transpose
#' @returns An Expression representing the partial transpose
#' @export
partial_transpose <- function(expr, dims, axis = 1L) {
  expr <- as_expr(expr)
  if (length(expr@shape) < 2L || expr@shape[1L] != expr@shape[2L]) {
    cli_abort("{.fn partial_transpose} 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 |i><j| kron I) expr (I kron |i><j| kron I)
  ## and sum over all i,j = 0, ..., dims[axis]-1
  d_ax <- dims[axis]
  terms <- vector("list", d_ax * d_ax)
  idx <- 1L
  for (i in seq_len(d_ax) - 1L) {       ## 0-indexed
    for (j in seq_len(d_ax) - 1L) {     ## 0-indexed
      terms[[idx]] <- .partial_transpose_term(expr, i, j, dims, axis)
      idx <- idx + 1L
    }
  }
  Reduce("+", terms)
}

## Helper: compute (i,j)-th term in the partial transpose sum
## CVXPY: partial_transpose.py _term()
.partial_transpose_term <- function(expr, i, j, dims, axis) {
  ## a = (I kron |i><j| kron I)
  ## Start with 1x1 identity
  a <- 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
      ## |i><j| is a dim_k x dim_k matrix with 1 at position (i, j)
      v <- Matrix::sparseMatrix(i = i + 1L, j = j + 1L, x = 1.0,
                                dims = c(dim_k, dim_k))
      a <- kronecker(a, v)
    } else {
      eye_k <- Matrix::Diagonal(dim_k)
      a <- kronecker(a, eye_k)
    }
  }
  ## a @ expr @ a -- same matrix on both sides
  Constant(a) %*% expr %*% Constant(a)
}