Repository: ad12/meddlr (language: python)

Issue:
Handle arbitrary dimensional inputs into nn.Modules
Medical images have a heterogenous number of dimensions. For example, MRI volumes can be treated as 2D slices, 3D volumes, or nD tensors (e.g. multi-echo, time, etc.). It is important to have built-in modules be able to handle tensors of arbitrary dimensions with some (but little) opinion about the behavior of the model across these different dimensions.

## Proposal
Currently, frameworks require the user to keep track of the order of dimensions for tensors. We may be able to simplify this by having a tensor `dim_order` value that is managed and manipulated by the meddlr ecosystem.

## TODOs
- [ ] Update `forward.mri.SenseModel` as a sample use case for managing multiple dimensions


Discussion / hints:

Fix patch:
diff --git a/meddlr/forward/mri.py b/meddlr/forward/mri.py
index c2d91506..39a985e8 100644
--- a/meddlr/forward/mri.py
+++ b/meddlr/forward/mri.py
@@ -52,7 +52,13 @@ def _adjoint_op(self, kspace):
             image = cplx.mul(image.unsqueeze(-2), cplx.conj(self.maps))  # [B,...,#coils,#maps,2]
             return image.sum(-3)
         else:
-            image = cplx.mul(image.unsqueeze(-1), cplx.conj(self.maps))  # [B,...,#coils,#maps,1]
+            # This is a hacky solution managing multi-channel inputs.
+            # Note multi-channel inputs are only supported in complex tensors.
+            # TODO (arjundd, issue #18): Fix with tensor ordering.
+            if image.ndim != self.maps.ndim:
+                image = image.unsqueeze(-1)
+
+            image = cplx.mul(image, cplx.conj(self.maps))  # [B,...,#coils,#maps,1]
             return image.sum(-2)
 
     def _forward_op(self, image):
@@ -67,7 +73,13 @@ def _forward_op(self, image):
             kspace = self.weights * oF.fft2c(kspace.sum(-2), channels_last=True)  # [B,...,#coils,2]
         else:
             kspace = cplx.mul(image.unsqueeze(-2), self.maps)
-            kspace = self.weights * oF.fft2c(kspace.sum(-1), channels_last=True)
+            # This is a hacky solution managing multi-channel inputs.
+            # Note this change means that multiple maps with multi-channel inputs
+            # is not supported for forward operations. This will change in future updates.
+            # TODO (arjundd, issue #18): Fix with tensor ordering.
+            if image.shape[-1] == self.maps.shape[-1]:
+                kspace = kspace.sum(-1)
+            kspace = self.weights * oF.fft2c(kspace, channels_last=True)
         return kspace
 
     def forward(self, input: torch.Tensor, adjoint: bool = False):
diff --git a/meddlr/modeling/loss_computer.py b/meddlr/modeling/loss_computer.py
index ade1bae6..2398ab9c 100644
--- a/meddlr/modeling/loss_computer.py
+++ b/meddlr/modeling/loss_computer.py
@@ -100,7 +100,7 @@ def __call__(self, input, output):
         target = output["target"].to(pred.device)
 
         if self.renormalize_data:
-            normalization_args = {k: input.get(k, output[k]) for k in ["mean", "std"]}
+            normalization_args = {k: input.get(k, output.get(k, None)) for k in ["mean", "std"]}
             normalized = self._normalizer.undo(
                 image=pred,
                 target=target,
diff --git a/meddlr/modeling/meta_arch/cs_model.py b/meddlr/modeling/meta_arch/cs_model.py
index 3afdf87c..2c44cd2d 100644
--- a/meddlr/modeling/meta_arch/cs_model.py
+++ b/meddlr/modeling/meta_arch/cs_model.py
@@ -19,8 +19,8 @@
 
 import meddlr.ops.complex as cplx
 from meddlr.config.config import configurable
+from meddlr.forward.mri import SenseModel
 from meddlr.utils.general import move_to_device
-from meddlr.utils.transforms import SenseModel
 
 from .build import META_ARCH_REGISTRY
 
diff --git a/meddlr/modeling/meta_arch/denoising.py b/meddlr/modeling/meta_arch/denoising.py
index 503c77f0..07c85c6b 100644
--- a/meddlr/modeling/meta_arch/denoising.py
+++ b/meddlr/modeling/meta_arch/denoising.py
@@ -7,10 +7,10 @@
 import meddlr.ops.complex as cplx
 from meddlr.config.config import configurable
 from meddlr.data.transforms.noise import NoiseModel
+from meddlr.forward.mri import SenseModel
 from meddlr.modeling.meta_arch.build import META_ARCH_REGISTRY, build_model
 from meddlr.utils.events import get_event_storage
 from meddlr.utils.general import move_to_device
-from meddlr.utils.transforms import SenseModel
 
 __all__ = ["DenoisingModel"]
 
diff --git a/meddlr/modeling/meta_arch/ssdu.py b/meddlr/modeling/meta_arch/ssdu.py
index cd18be21..5b16d24f 100644
--- a/meddlr/modeling/meta_arch/ssdu.py
+++ b/meddlr/modeling/meta_arch/ssdu.py
@@ -4,13 +4,13 @@
 
 import meddlr.ops as oF
 from meddlr.config.config import configurable
+from meddlr.forward.mri import SenseModel
 from meddlr.modeling.meta_arch.build import META_ARCH_REGISTRY, build_model
 from meddlr.ops import complex as cplx
 from meddlr.transforms.base.mask import KspaceMaskTransform
 from meddlr.transforms.gen.mask import RandomKspaceMask
 from meddlr.utils.events import get_event_storage
 from meddlr.utils.general import move_to_device
-from meddlr.utils.transforms import SenseModel
 
 
 @META_ARCH_REGISTRY.register()
diff --git a/meddlr/modeling/meta_arch/unet.py b/meddlr/modeling/meta_arch/unet.py
index f702d2c7..87e9f127 100644
--- a/meddlr/modeling/meta_arch/unet.py
+++ b/meddlr/modeling/meta_arch/unet.py
@@ -11,9 +11,9 @@
 
 import meddlr.ops.complex as cplx
 from meddlr.config.config import configurable
+from meddlr.forward.mri import SenseModel
 from meddlr.utils import transforms as T
 from meddlr.utils.events import get_event_storage
-from meddlr.utils.transforms import SenseModel
 
 from .build import META_ARCH_REGISTRY
 
@@ -337,7 +337,7 @@ def forward(self, input, return_pp=False, vis_training=False):
             output = zf_image.reshape(zf_image.shape[:-2] + (-1,)).unsqueeze(-2)
         else:
             output = zf_image
-        output = zf_image.permute(0, 4, 1, 2, 3).squeeze(-1)
+        output = output.permute(0, 4, 1, 2, 3).squeeze(-1)
 
         # Run U-Net.
         output = self.base_forward(output)
diff --git a/meddlr/modeling/meta_arch/unrolled.py b/meddlr/modeling/meta_arch/unrolled.py
index fdda4760..688a1638 100644
--- a/meddlr/modeling/meta_arch/unrolled.py
+++ b/meddlr/modeling/meta_arch/unrolled.py
@@ -7,9 +7,9 @@
 
 import meddlr.ops.complex as cplx
 from meddlr.config.config import configurable
+from meddlr.forward.mri import SenseModel
 from meddlr.utils.events import get_event_storage
 from meddlr.utils.general import move_to_device
-from meddlr.utils.transforms import SenseModel
 
 from ..layers.layers2D import ResNet
 from .build import META_ARCH_REGISTRY, build_model
@@ -134,7 +134,7 @@ def forward(self, inputs, return_pp=False, vis_training=False):
         A = inputs.get("signal_model", None)
         maps = inputs["maps"]
         num_maps_dim = -2 if cplx.is_complex_as_real(maps) else -1
-        if self.num_emaps != maps.size()[num_maps_dim]:
+        if self.num_emaps != maps.size()[num_maps_dim] and maps.size()[num_maps_dim] != 1:
             raise ValueError("Incorrect number of ESPIRiT maps! Re-prep data...")
 
         # Move step sizes to the right device.
diff --git a/meddlr/ops/complex.py b/meddlr/ops/complex.py
index 4c763f3f..03a72007 100644
--- a/meddlr/ops/complex.py
+++ b/meddlr/ops/complex.py
@@ -66,7 +66,7 @@ def is_complex_as_real(x):
         has a size of ``2`` because it is the real-imaginary channel or
         for some other reason.
     """
-    return x.size(-1) == 2
+    return not is_complex(x) and x.size(-1) == 2
 
 
 def conj(x):
diff --git a/meddlr/ops/fft.py b/meddlr/ops/fft.py
index c30cc791..b323e702 100644
--- a/meddlr/ops/fft.py
+++ b/meddlr/ops/fft.py
@@ -196,7 +196,23 @@ def ifftshift(x, dim=None):
 def _fft_template(
     data: torch.Tensor, kind, dim=None, norm="ortho", is_real: bool = None, centered: bool = True
 ) -> torch.Tensor:
-    """Template for fft operations."""
+    """Template for fft operations.
+
+    Args:
+        data (torch.Tensor): A tensor.
+        kind (str): Either ``'fft'`` or ``'ifft'``.
+        dim (int(s), optional): The dimension(s) along which to apply the operation.
+            Defaults to all dimensions.
+        norm (str, optional): The normalization method. Defaults to ``'ortho'``.
+        is_real (bool, optional): If ``True``, ``input`` is treated like a real-valued
+            tensor. If not specified, this is ``True`` only if ``data`` is not complex
+            and data is not inferred to be a real view of a complex tensor
+            (i.e. ``data.shape[-1] != 2``).
+        centered (bool, optional): If ``True``, apply centered FFT. Defaults to ``True``.
+
+    Returns:
+        torch.Tensor: The FFT (or IFFT) of the input.
+    """
     if isinstance(dim, int):
         dim = (dim,)
     if norm is True:


Test patch:
diff --git a/tests/forward/test_mri.py b/tests/forward/test_mri.py
index 4be68cc7..34d011a0 100644
--- a/tests/forward/test_mri.py
+++ b/tests/forward/test_mri.py
@@ -49,3 +49,33 @@ def test_reproducibility(self):
         out = A_new(kspace, adjoint=True)
         assert torch.allclose(out, expected)
         assert torch.allclose(A_new(out, adjoint=False), A(expected, adjoint=False))
+
+    def test_multichannel(self):
+        """Test multi-channel inputs."""
+        ky = 20
+        kz = 20
+        nc = 8
+        num_channels = 3
+        nm = 1
+        bsz = 5
+
+        kspace = torch.randn(bsz, ky, kz, nc, num_channels, dtype=torch.complex64)
+        maps = torch.rand(bsz, ky, kz, nc, nm, dtype=torch.complex64)
+        maps = maps / cplx.rss(maps, dim=-2).unsqueeze(-2)
+
+        A = SenseModel(maps)
+
+        expected = []
+        for c in range(num_channels):
+            expected.append(A(kspace[..., c], adjoint=True))
+        expected = torch.cat(expected, dim=-1)
+        out_image = A(kspace, adjoint=True)
+        assert torch.allclose(out_image, expected)
+
+        expected = []
+        for c in range(num_channels):
+            expected.append(A(out_image[..., c : c + 1], adjoint=False))
+        expected = torch.stack(expected, dim=-1)
+        out_kspace = A(out_image, adjoint=False)
+        # both clauses required for CI to pass on python 3.7 - torch.allclose does not work
+        assert torch.allclose(out_kspace, expected, atol=1e-5)