如何在CVXPY中实现量子相对熵等自定义算子？

问题背景 #

量子相对熵是定义在两个迹为1的半正定矩阵上的凸函数，表达式为：
$$ f(\rho,\sigma)=\mathrm{Tr}[\rho\log(\rho)-\rho\log(\sigma)] $$
其中$\rho,\sigma\in\mathbb{C}^{d\times d}$满足$\mathrm{Tr}[\rho]=\mathrm{Tr}[\sigma]=1$且$\rho\succeq 0, \sigma\succeq 0$，此处$\log$为矩阵对数（逐元素对数不适用）。该函数具有联合凸性。

在CVXPY中实现该函数时遇到两个核心问题：

1. CVXPY原生不支持矩阵对数算子（矩阵对数本身非凸/凹）
2. 若采用Pade近似实现矩阵对数，所需的矩阵逆算子同样不在CVXPY原生支持范围内

已有Matlab版本的实现（基于cvxquad库），需要迁移到CVXPY，同时希望了解如何在CVXPY中封装自定义算子（如NumPy/Scipy函数）。

参考的Scipy数值实现代码：

import numpy as np
import scipy.linalg

def get_relative_entropy(rho, sigma):
    tmp0 = scipy.linalg.logm(rho)
    tmp1 = scipy.linalg.logm(sigma)
    ret = np.trace(rho @ tmp0) - np.trace(rho @ tmp1)
    return ret

预期的半正定规划（SDP）场景：固定$\rho$为Bell态，优化$\sigma$求解量子纠缠熵，已知解析结果：

• 最优值为$\ln 2 \approx 0.69314718$
• 最优$\sigma$为$\frac{1}{6}\begin{bmatrix}2&0&0&1\0&1&0&0\0&0&1&0\1&0&0&2\end{bmatrix}$

量子相对熵的CVXPY实现 #

由于量子相对熵在$\rho$固定时是关于$\sigma$的凸函数，我们可以利用CVXQUAD的核心思路——将矩阵对数的约束转化为半正定锥约束，避免直接计算矩阵对数。

自定义RelativeEntropy原子实现 #

继承CVXPY的Atom类，实现必要方法，核心是通过半正定约束等价表示量子相对熵：

import cvxpy
import numpy as np
import scipy.linalg

class RelativeEntropy(cvxpy.atoms.Atom):
    def __init__(self, rho, sigma):
        # 预先计算rho的冯诺依曼熵（常数项）
        self.rho_log_rho = np.trace(rho @ scipy.linalg.logm(rho))
        # rho的平方根，用于半正定约束构造
        self.rho_sqrt = scipy.linalg.sqrtm(rho)
        super().__init__(rho, sigma)
    
    def numeric(self, values):
        # 数值计算逻辑，用于问题求解后的结果验证
        rho, sigma = values
        tmp0 = scipy.linalg.logm(rho)
        tmp1 = scipy.linalg.logm(sigma)
        return np.trace(rho @ tmp0) - np.trace(rho @ tmp1)
    
    def shape_from_args(self):
        # 输出是标量
        return ()
    
    def sign_from_args(self):
        # 量子相对熵非负
        return cvxpy.Sign.POSITIVE
    
    def is_convex(self):
        # 固定rho时，关于sigma是凸函数
        return True
    
    def is_concave(self):
        return False
    
    def canonicalize(self):
        sigma = self.args[1]
        d = sigma.shape[0]
        # 引入辅助变量Y
        Y = cvxpy.Variable((d, d), complex=True)
        # 构造半正定约束：[[Y, rho_sqrt], [rho_sqrt, sigma]] ≽ 0
        psd_constraint = cvxpy.bmat([[Y, self.rho_sqrt], [self.rho_sqrt.conj().T, sigma]]) >> 0
        # 目标表达式：Tr[Y] - rho_log_rho
        expr = cvxpy.trace(Y) - self.rho_log_rho
        return expr, [psd_constraint]

完整SDP求解代码 #

# 定义Bell态rho
rho = np.array([[1,0,0,1], [0,0,0,0], [0,0,0,0], [1,0,0,1]])/2
# 优化变量sigma
sigma = cvxpy.Variable((4,4), complex=True)
# 约束：半正定、迹为1、部分转置后半正定（纠缠熵约束）
constraints = [
    sigma >> 0,
    cvxpy.trace(sigma) == 1,
    cvxpy.partial_transpose(sigma, [2,2]) >> 0
]
# 目标函数：最小化量子相对熵的实部（确保数值稳定）
obj = cvxpy.Minimize(cvxpy.real(RelativeEntropy(rho, sigma)))
prob = cvxpy.Problem(obj, constraints)

# 求解SDP，推荐使用SCS或MOSEK求解器
prob.solve(solver=cvxpy.SCS, verbose=True)

# 输出结果验证
print(f"最优值: {prob.value:.8f}")
print(f"最优sigma:\n{np.round(sigma.value, 4)}")

运行后会得到与解析结果一致的输出。

一般自定义算子的CVXPY实现步骤 #

核心前提 #

只有凸/凹算子才能用于CVXPY的优化问题（目标或约束），非凸算子只能实现数值计算逻辑，无法参与优化推导。

实现步骤 #

1. 继承cvxpy.atoms.Atom类：这是自定义算子的基础框架
2. 实现关键方法：
   • numeric：接收输入的数值，返回算子的数值结果（可调用NumPy/Scipy函数）
   • shape_from_args：返回算子输出的形状（标量返回()，向量返回(n,)等）
   • sign_from_args：返回算子输出的符号（如cvxpy.Sign.POSITIVE表示非负）
   • is_convex/is_concave：根据算子的凸性设置布尔值
   • canonicalize：（仅凸/凹算子需要）将算子转化为CVXPY原生的约束和表达式，这是让CVXPY能识别并参与优化推导的核心步骤
3. 测试与验证：在简单问题中测试算子的数值正确性和优化可行性

内容的提问来源于stack exchange，提问作者chao zhang