Subspace: project with conj(basis), expand with basis.T, so complex signals round-trip

CHANGELOG.md

@@ -10,6 +10,16 @@
   along each; `refine=False` stops after the coarse pass. A 3 % (T1, T2) grid
   of about 37 000 atoms at `coarse=3` scores about 4 200 + 49 atoms per voxel.
 
+### Fixed
+
+- **`Subspace` round-trips complex signals.** `project` is `signals @ basis.conj()`
+  and `expand` is `coefficients @ basis.mT`, so a signal in the span comes back
+  unchanged; `modes` is `basis.mH`, which mri-nufft's `MRISubspace` expands
+  with its conjugate. Before, a complex dictionary whose atoms do not share one
+  phase was projected onto the conjugate of its span (96 % error on random
+  signals of the span). Real dictionaries are unaffected. bartorch's `basis=`
+  and BART's `pics -B` take `basis.mT`.
+
 ### Changed
 
 - **Every GPU, a bounded share of it, by default.** `execution="auto"` is the

docs/explanations/model-based-reconstruction.md

@@ -75,7 +75,8 @@ axis 1, as deepinv and mri-nufft lay it out.
 
 An estimator fitted with `rank` carries a {class}`~blochsim.Subspace`.
 {attr}`~blochsim.Subspace.modes` gives the basis as `(rank, contrasts)`, the
-layout mri-nufft's `MRISubspace` and BART's `pics -B` read. The reconstruction
+layout mri-nufft's `MRISubspace` reads; BART's `pics -B` and bartorch's
+`basis=` contract without its conjugate and read `subspace.basis.mT`. The reconstruction
 is linear in the coefficients, has one minimum and needs no starting point.
 {meth}`~blochsim.estimators.DictionaryMatcher.from_coefficients` then maps the coefficients
 without projecting them again. A signal that needs many components, such as a

src/blochsim/_subspace.py

@@ -63,16 +63,22 @@ def retained(self) -> float:
 
     @property
     def modes(self) -> torch.Tensor:
-        """The basis transposed, ``(rank, contrasts)``.
+        """The basis conjugate-transposed, ``(rank, contrasts)``: mri-nufft's ``MRISubspace`` layout.
 
-        This is the layout of mri-nufft's ``MRISubspace`` and BART's
-        ``pics -B``. It is a plain transpose: expansion is
-        ``image_t = sum_k conj(modes[k, t]) c_k``, as in :meth:`expand`.
+        ``MRISubspace`` expands ``image_t = sum_k conj(modes[k, t]) c_k`` and
+        projects ``c_k = sum_t modes[k, t] image_t``, which are :meth:`expand`
+        and :meth:`project`. BART's ``pics -B`` and bartorch's ``basis=``
+        contract without the conjugate, ``image_t = sum_k B[k, t] c_k``, so
+        they take ``basis.mT``; for a real basis the two are the same.
         """
-        return self.basis.mT
+        return self.basis.mH.resolve_conj()
 
     def project(self, signals: torch.Tensor) -> torch.Tensor:
-        """Project ``signals`` onto the basis, giving ``signals @ basis``.
+        r"""Project ``signals`` onto the basis, giving ``signals @ basis.conj()``.
+
+        The coefficients are the inner products of each signal with the basis
+        vectors, :math:`c_k = \langle b_k, s \rangle = \sum_t \overline{b_{tk}} s_t`,
+        so :meth:`expand` returns a signal of the span unchanged.
 
         Parameters
         ----------
@@ -87,10 +93,13 @@ def project(self, signals: torch.Tensor) -> torch.Tensor:
             ``(..., rank)`` coefficients.
         """
         dtype = torch.promote_types(signals.dtype, self.basis.dtype)
-        return signals.to(dtype) @ self.basis.to(dtype).to(signals.device)
+        return signals.to(dtype) @ self.basis.conj().to(dtype).to(signals.device)
 
     def expand(self, coefficients: torch.Tensor) -> torch.Tensor:
-        """Expand coefficients to contrasts, giving ``coefficients @ basis.mH``.
+        r"""Expand coefficients to contrasts, giving ``coefficients @ basis.mT``.
+
+        :math:`s_t = \sum_k b_{tk} c_k`, the inverse of :meth:`project` on the
+        span of the basis.
 
         Parameters
         ----------
@@ -103,7 +112,7 @@ def expand(self, coefficients: torch.Tensor) -> torch.Tensor:
             ``(..., contrasts)``, on the coefficients' device.
         """
         dtype = torch.promote_types(coefficients.dtype, self.basis.dtype)
-        return coefficients.to(dtype) @ self.basis.mH.to(dtype).to(coefficients.device)
+        return coefficients.to(dtype) @ self.basis.mT.to(dtype).to(coefficients.device)
 
     @classmethod
     def fit(cls, signals: torch.Tensor, rank: int) -> Subspace:

tests/estimators/test_subspace.py

@@ -96,3 +96,33 @@ def test_the_leading_axes_are_flattened() -> None:
     shaped = Subspace.fit(dictionary.reshape(6, 10, 32), 5)
 
     assert torch.allclose(flat.basis, shaped.basis, atol=1e-6)
+
+
+def test_complex_signals_of_the_span_round_trip() -> None:
+    """``expand`` undoes ``project`` on the span for complex signals too.
+
+    A real dictionary cannot tell ``s @ B`` from ``s @ conj(B)``; a complex one,
+    whose atoms do not share one phase, can.
+    """
+    generator = torch.Generator().manual_seed(0)
+    basis = torch.linalg.qr(
+        torch.randn(50, 3, dtype=torch.complex64, generator=generator)
+    )[0]
+    signals = torch.randn(10, 3, dtype=torch.complex64, generator=generator) @ basis.mT
+    subspace = Subspace.fit(signals, 3)
+
+    coefficients = subspace.project(signals)
+
+    torch.testing.assert_close(
+        subspace.expand(coefficients), signals, rtol=1e-5, atol=1e-5
+    )
+    torch.testing.assert_close(
+        coefficients, signals @ subspace.basis.conj(), rtol=1e-6, atol=1e-6
+    )
+    # mri-nufft's MRISubspace expands with conj(modes) and projects with modes.
+    torch.testing.assert_close(
+        coefficients @ subspace.modes.conj(), signals, rtol=1e-5, atol=1e-5
+    )
+    torch.testing.assert_close(
+        signals @ subspace.modes.mT, coefficients, rtol=1e-5, atol=1e-5
+    )

tests/operator/test_subspace_route.py

@@ -19,13 +19,20 @@
 RANK = 3
 
 
-@pytest.fixture(scope="module")
-def basis():
-    """A temporal basis for a multi-echo decay, from the model itself."""
+@pytest.fixture(scope="module", params=["real", "off-resonant"])
+def basis(request):
+    """A temporal basis for a multi-echo decay, from the model itself.
+
+    Off resonance turns each atom's phase along the echoes at a rate of its
+    own, so the basis is complex and a conjugate in the wrong place shows.
+    """
     acquisition = MultiEchoSimulator(TE=TE_MS)
     signals = torch.as_tensor(
         acquisition.simulate(T2=torch.linspace(20.0, 300.0, 64))
     ).to(torch.complex64)
+    if request.param == "off-resonant":
+        hertz = torch.linspace(-8.0, 8.0, 64)[:, None]
+        signals = signals * torch.exp(2j * torch.pi * hertz * TE_MS[None] / 1000)
     return Subspace.fit(signals, RANK)