GraphOTTOutput.transport_matrix returns the expanded (n+m)² matrix, not the (n, m) block

Split out of #846, where it was explicitly left out of scope.

What I see

GraphOTTOutput wraps an output solved on a joint graph geometry over source and target cells, so the underlying ott problem has shape (n + m, n + m). It therefore overrides shape with the logical block it represents (src/moscot/backends/ott/output.py:307-308):

@property
def shape(self) -> Tuple[int, int]:  # noqa: D102
    return self._shape          # (n, m), passed in by the solver

and slices its applies back to that block (:317-321, via _expand_data):

def _apply(self, x: ArrayLike, *, forward: bool) -> ArrayLike:
    x_expanded = self._expand_data(x, forward=forward)
    res = self._output.apply(x_expanded.T, axis=1 - forward, lse_mode=False).T
    return res[len(x):] if forward else res[:-len(x)]

But transport_matrix is not overridden, so it inherits OTTOutput's (:191-193):

@property
def transport_matrix(self) -> ArrayLike:  # noqa: D102
    return self._output.matrix

which is the full (n + m, n + m) matrix of the expanded problem. So output.transport_matrix.shape != output.shape for this class, and the returned matrix is indexed in expanded coordinates — row i of it is not row i of the coupling the user asked for.

GraphOTTOutput is constructed in src/moscot/backends/ott/solver.py:167 whenever the linear term is a Geodesic geometry (i.e. after set_graph_xy(..., cost="geodesic")).

Why it matters

Anything that consumes the materialized matrix rather than push/pull gets the wrong object. The clearest downstream case is cellrank.kernels.RealTimeKernel.from_moscot, which does

coupling = solution.transport_matrix
couplings[t1, t2] = AnnData(coupling, obs=adata_src.obs, var=adata_tgt.obs)

so a graph-based temporal problem would either raise on the shape mismatch or silently build a coupling from expanded coordinates.

push/pull, and therefore sparsify, are unaffected — they go through _apply, which slices correctly.

Caveat on my evidence

This is a code-reading finding. I could not get a GraphOTTOutput instance to reproduce it at runtime: passing a joint k-NN graph to set_graph_xy(..., cost="geodesic", t=1.0) on a TemporalProblem gave me a plain OTTOutput over a dense Geometry (shape (60, 60), matching shape), so I never reached the isinstance(prob.geom, geodesic.Geodesic) branch in _solve. Someone who knows the configuration that actually produces a Geodesic linear term could confirm in a couple of lines:

out = ...  # solved graph-based problem
assert out.transport_matrix.shape == out.shape

Suggested fix

Override transport_matrix on GraphOTTOutput to return the (n, m) block — either by slicing self._output.matrix[:n, n:] (whichever block corresponds), or by materializing it through the same _apply path the class already uses for push/pull, so expanded coordinates never escape the class. A test asserting transport_matrix.shape == shape and that its marginals match a/b would pin it.

🤖 Generated with Claude Code

Start in src/moscot/backends/ott/output.py at GraphOTTOutput and compare its shape and _apply behavior with OTTOutput.transport_matrix. Check the GraphOTTOutput construction branch in src/moscot/backends/ott/solver.py, then reproduce a graph-based solve if possible. Done means transport_matrix has shape (n, m) and its marginals match a and b, with a regression test covering the behavior.