Quantum optimal transport#

The partial-trace constraint operator and a dedicated dual solver.

sdplab.special.qot.QOTConstraintOp

Partial-trace operator \(\mathcal{A}\) for quantum optimal transport.

sdplab.special.qot.solve_qot_dual

Solve the QOT dual SDP.
class sdplab.special.qot.QOTConstraintOp(*, d, N, atol=0.0, rtol=0.0, enforce_herm=True, ctx=None)[source]#

Bases: MatrixFreeConstraintOp

Partial-trace operator \(\mathcal{A}\) for quantum optimal transport.

This is the linear map

\[\mathcal{A}: \operatorname{Herm}(d^N) \to \operatorname{Herm}(d)^N.\]

Its domain \(\operatorname{dom}(\mathcal{A})\) contains global Hermitian matrices on N tensor factors. Its codomain \(\operatorname{cod}(\mathcal{A})\) contains N Hermitian d x d matrices, one per subsystem.

If \(\Gamma \in \operatorname{dom}(\mathcal{A})\) is a feasible QOT coupling and \(\gamma_k\) is the prescribed marginal for site k, then the equality constraint is

\[(\mathcal{A}\Gamma)_k = \operatorname{Tr}^k[\Gamma] = \gamma_k.\]
Parameters:

  • d (int)

  • N (int)

  • atol (float)

  • rtol (float)

  • enforce_herm (bool)

  • ctx (Context | str | None)

__init__(*, d, N, atol=0.0, rtol=0.0, enforce_herm=True, ctx=None)[source]#

Create \(\mathcal{A}: \operatorname{Herm}(d^N) \to \operatorname{Herm}(d)^N\).

Parameters:

  • d (int) – Local Hilbert-space dimension.

  • N (int) – Number of tensor factors/subsystems.

  • atol (float) – Absolute tolerance for Hermitian membership checks.

  • rtol (float) – Relative tolerance for Hermitian membership checks.

  • enforce_herm (bool) – Whether domain and codomain require Hermitian input.

  • ctx (Context | str | None) – Optional backend context.

apply(X)[source]#

Return \(\mathcal{A}\Gamma\), the one-body marginals of X.

X is the numerical array representing \(\Gamma \in \operatorname{dom}(\mathcal{A})\). It has shape (d^N, d^N), and the return value lies in \(\operatorname{cod}(\mathcal{A})\) with shape (N, d, d). The k-th block is \((\mathcal{A}\Gamma)_k = \operatorname{Tr}^k[\Gamma]\).

Parameters:

X (DenseArray)

Return type:

DenseArray

rapply(y)[source]#

Apply the adjoint \(\mathcal{A}^\dagger\) as a Kronecker sum.

For \(y = (y_0, \ldots, y_{N-1}) \in \operatorname{cod}(\mathcal{A})\), the adjoint is

\[\mathcal{A}^\dagger y = y_0 \oplus \cdots \oplus y_{N-1} = \sum_k I \otimes \cdots \otimes y_k \otimes \cdots \otimes I,\]

as an element of \(\operatorname{dom}(\mathcal{A})\). This identity is characterized by

\[\operatorname{Tr}[(\mathcal{A}\Gamma)y] = \operatorname{Tr}[\Gamma(\mathcal{A}^\dagger y)].\]

The decorator already asserts codomain membership on entry and domain membership on exit, so no explicit check is repeated here.

Parameters:

y (DenseArray)

Return type:

Any

to_sparse()[source]#

Materialize \(\mathcal{A}\) as a sparse coordinate matrix.

The shape is (prod(cod.shape), prod(dom.shape)) = (N d^2, d^{2N}), matching to_matrix(): row \((k, a, b)\) is the flattened tensor slice reading off \(\operatorname{Tr}^k\) at marginal entry \((a, b)\),

\[(\mathcal{A}\Gamma)_{k,ab} = \sum_{l,r} \Gamma_{(l,a,r),(l,b,r)},\]

so that row carries exactly \(d^{N-1}\) unit entries – the configurations of the traced-out subsystems. Storage is therefore \(N d^{N+1}\) nonzeros instead of the \(N d^{2N+2}\) of the dense form. The base class raises NotImplementedError here; the partial trace is a 0/1 incidence matrix, so it is worth providing.

Return type:

SparseArray

to_cvxpy()[source]#

Return the QOT constraints as a list of per-constraint sparse matrices.

This adapts the dense Kronecker construction of the QOT-to-SDP proof (th. 3.1 of https://arxiv.org/abs/2105.06922) into standard-form constraint matrices for a general SDP solver such as the CVXPY backend, whose equalities read \(\operatorname{Re}\operatorname{Tr}[A_i \Gamma] = b_i\).

Constraint i = (k, \alpha) is the Hermitian matrix

\[A_i = \mathcal{A}^\dagger(H_\alpha^{(k)}) = I \otimes \cdots \otimes H_\alpha \otimes \cdots \otimes I,\]

where \(H_\alpha\) ranges over the real-coordinate Hermitian generators of the k-th d x d block (see _herm_generators()). Because each \(H_\alpha\) is Hermitian, the matching right-hand side \(b_i = \operatorname{Tr}[H_\alpha^{(k)} \gamma_k]\) is real even though the marginals \(\gamma_k\) are complex Hermitian – naively flattening the marginal entries would instead give a complex b. The matching b is produced by rhs_to_cvxpy() and the dual is reassembled by dual_from_cvxpy().

The returned list has m = N d(d+1)/2 entries for a real context and m = N d^2 for a complex one, each a sparse (d^N, d^N) matrix so a solver can form trace(A_i @ Gamma) directly. Each generator embeds over d^{N-1} configurations of the traced-out subsystems, so storage stays sparse rather than the dense m d^{2N}.

Return type:

list[SparseArray]

dual_from_cvxpy(y)[source]#

Reassemble marginal dual blocks from per-constraint scalar duals.

Inverse of the row layout of to_cvxpy() / rhs_to_cvxpy(): constraint i = (k, \alpha) reads off generator \(H_\alpha\) of block k, so the marginal dual is \(U_k = \sum_\alpha y_{(k,\alpha)} H_\alpha\), a Hermitian d x d block. The result is the stacked (N, d, d) codomain element. The caller supplies y already carrying the intended dual sign.

Parameters:

y (DenseArray)

Return type:

DenseArray

rhs_to_cvxpy(rhs)[source]#

Return the real right-hand side b matching to_cvxpy().

rhs is the stacked codomain array (N, d, d) of Hermitian one-body marginals \(\gamma_k\). The returned real vector b has length m (the number of matrices from to_cvxpy()) with \(b_{(k,\alpha)} = \operatorname{Tr}[H_\alpha^{(k)} \gamma_k]\), laid out in the same generator order so that the SDP equality \(\operatorname{Re}\operatorname{Tr}[A_i \Gamma] = b_i\) holds.

Equivalently \(b = \operatorname{Re}\langle H_\alpha, \gamma_k\rangle\) read off the generator matrices, which is how it is evaluated: the per-entry re/im selection of _herm_generators() is exactly what tracing against the generator performs.

Parameters:

rhs (DenseArray)

Return type:

DenseArray

property A: Any#

Native numerical representation of this operator.

Concrete subclasses may choose the representation that best matches their storage model: for example, dense operators return a dense array while sparse operators return their sparse matrix. Matrix-free or lazy operators generally do not have such a representation and should leave this property unimplemented. Use to_dense() when a dense tensor materialization is explicitly required.

property H: LinOp#

Hermitian-adjoint view of this linear operator.

Returns:

Adjoint view satisfying \(\langle A x, y\rangle_Y = \langle x, A^* y\rangle_X\).

Return type:

LinOp

adjoint()#

Return the Hermitian-adjoint view of this linear operator.

Return type:

LinOp

adjoint_apply(y)#

Apply the adjoint of this linear operator to y.

Parameters:

y (Any)

Return type:

Any

assert_codomain(y)#

Raise if y is not in the codomain.

Parameters:

y (Any)

Return type:

None

assert_domain(x)#

Raise if x is not in the domain.

Parameters:

x (Any)

Return type:

None

property check_level: Literal['none', 'cheap', 'standard', 'strict']#

Return this object’s runtime validation level.

property codomain: Codomain#

Codomain space of this linear operator.

convert(new_ctx=None)#

Return this object represented in new_ctx.

Parameters:

new_ctx (Context | BackendFamily | str | None)

Return type:

Self

property ctx: Context#

Return the execution context bound to this object.

property domain: Domain#

Domain space of this linear operator.

property dtype: Any#

Return the default dtype associated with this object’s context.

classmethod from_linop(op)#
Parameters:

op (LinOp)

Return type:

ConstraintOp

fuse(*, materialize=False)#

Return an equivalent operator with fusible sub-expressions multiplied out.

Tier-2 lazy-algebra simplification ([ADR-021](021_lazy_operator_algebra_and_simplification.md)): collapse each maximal subtree of densely-fusible operators into a single materialized operator — for example, a composition of dense operators becomes one DenseLinOp holding the matrix product \(M_A M_B\) — while leaving matrix-free and other non-materializable leaves intact.

This is an explicit, opt-in materialization. The result is mathematically equal to self but only within floating-point rounding: fusing reassociates the arithmetic (multiplying matrices then applying differs from applying in sequence at the ulp level), so equality holds up to tolerance, not bit-for-bit. The fused operator preserves the domain, codomain, context, and scalar-field/dtype identity. A leaf operator returns itself.

Parameters:

materialize (bool, optional) – With the default False, a matrix-free operand ([ADR-008](008_linop_subclasses.md)) is never densified: it remains a lazy leaf and only breaks a fusible run. With True the caller explicitly accepts giving up the matrix-free contract: a matrix-free operand is densified into a DenseLinOp (via its to_dense basis probe, which may be expensive), allowing the enclosing expression to collapse to a single dense operator.

Returns:

A fused operator with the same action as self (up to rounding).

Return type:

LinOp

is_hermitian()#

Return whether this operator is structurally Hermitian when known.

Returns:

True or False when the subclass can verify the structure cheaply, otherwise None for unknown or matrix-free operators.

Return type:

bool | None

property ops: BackendOps#

Return backend operations associated with this object’s context.

rvapply(ys)#

Apply the adjoint over a leading batch axis. Input must have shape (N,) + codomain.shape; use moveaxis for other layouts.

Parameters:

ys (Any)

Return type:

Any

to_dense()#

Materialize this operator as a dense backend array.

The returned array has shape self.codomain.shape + self.domain.shape. The default implementation is intended for small problems, debugging, and tests. It materializes the full coordinate matrix, so subclasses that already store a dense or sparse matrix should override this method for efficiency.

Return type:

Any

to_matrix()#

Materialize this operator as a 2D dense coordinate matrix.

The returned array has shape (prod(self.codomain.shape), prod(self.domain.shape)). The default implementation builds a batch of standard basis vectors and calls vapply() once. If a space cannot batch-flatten or batch-unflatten its representation, it falls back to a safe Python loop. This method is for small/testing use; concrete storage-backed subclasses should override it when they can expose a matrix directly.

Return type:

Any

vapply(xs)#

Apply over a leading batch axis. Input must have shape (N,) + domain.shape; use moveaxis for other layouts.

Parameters:

xs (Any)

Return type:

Any

sdplab.special.qot.solve_qot_dual(qot, solver='MOSEK', verbose=False, *args, **kwargs)[source]#

Solve the QOT dual SDP.

For the QOT constraint operator \(\mathcal{A}\), the primal coupling satisfies \(\Gamma \in \operatorname{dom}(\mathcal{A}) = \operatorname{Herm}(d^N)\) and the marginal data satisfy \(\gamma = (\gamma_0, \ldots, \gamma_{N-1}) \in \operatorname{cod}(\mathcal{A}) = \operatorname{Herm}(d)^N\). The dual problem is

\[\max_{U \in \operatorname{cod}(\mathcal{A})}\quad \sum_k \operatorname{Tr}[U_k \gamma_k] \quad \text{s.t.}\quad \mathcal{A}^\dagger U \preceq C.\]

Here \(C \in \operatorname{dom}(\mathcal{A})\) is the cost matrix and \(U = (U_0, \ldots, U_{N-1})\) is the block dual variable. The adjoint constraint is

\[\mathcal{A}^\dagger U = U_0 \oplus \cdots \oplus U_{N-1} = \sum_k I \otimes \cdots \otimes U_k \otimes \cdots \otimes I \preceq C,\]

equivalently \(C - \mathcal{A}^\dagger U \succeq 0\).

Returns:

A pair (primal, dual). The primal value stores the coupling \(\Gamma\), represented as the positive-semidefinite multiplier for \(C - \mathcal{A}^\dagger U \succeq 0\). The dual value stores the optimized blocks \(U_k\) in \(\operatorname{cod}(\mathcal{A})\).

Parameters:
Return type:

tuple[Any, Any]