Support 3D in ESPIRiT

Use batch 3D array_to_block instead of 4D array_to_block in ESPIRiT to support 3D ESPIRiT.
mikgroup · Sep 7, 2020 · a05b481 · a05b481
1 parent cb57c8f
commit a05b481
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 10 deletions.
diff --git a/sigpy/mri/app.py b/sigpy/mri/app.py
@@ -416,19 +416,22 @@ def __init__(self, ksp, calib_width=24,
 
         xp = self.device.xp
         with self.device:
-            # Get calibration matrix
-            kernel_shape = [num_coils] + [kernel_width] * img_ndim
-            kernel_strides = [1] * (img_ndim + 1)
-            mat = sp.array_to_blocks(calib, kernel_shape, kernel_strides)
-            mat = mat.reshape([-1, sp.prod(kernel_shape)])
+            # Get calibration matrix.
+            # Shape [num_coils] + num_blks + [kernel_width] * img_ndim
+            mat = sp.array_to_blocks(
+                calib, [kernel_width] * img_ndim, [1] * img_ndim)
+            mat = mat.reshape([num_coils, -1, kernel_width**img_ndim])
+            mat = mat.transpose([1, 0, 2])
+            mat = mat.reshape([-1, num_coils * kernel_width**img_ndim])
 
             # Perform SVD on calibration matrix
             _, S, VH = xp.linalg.svd(mat, full_matrices=False)
             VH = VH[S > thresh * S.max(), :]
 
             # Get kernels
             num_kernels = len(VH)
-            kernels = VH.reshape([num_kernels] + kernel_shape)
+            kernels = VH.reshape(
+                [num_kernels, num_coils] + [kernel_width] * img_ndim)
             img_shape = ksp.shape[1:]
 
             # Get covariance matrix in image domain

diff --git a/tests/mri/test_app.py b/tests/mri/test_app.py
@@ -96,20 +96,41 @@ def test_ones_JsenseRecon(self):
         npt.assert_allclose(mps, mps_rec, atol=1e-2, rtol=1e-2)
 
     def test_espirit_maps(self):
-        mps_shape = [8, 32, 32]
+        # 2D
+        mps_shape = [8, 16, 16]
         mps = sim.birdcage_maps(mps_shape)
         ksp = sp.fft(mps, axes=[-1, -2])
         mps_rec = app.EspiritCalib(ksp, show_pbar=False).run()
 
-        np.testing.assert_allclose(np.abs(mps)[:, 8:24, 8:24],
-                                   np.abs(mps_rec[:, 8:24, 8:24]),
+        np.testing.assert_allclose(np.abs(mps)[:, 4:-4, 4:-4],
+                                   np.abs(mps_rec[:, 4:-4, 4:-4]),
+                                   rtol=1e-2, atol=1e-2)
+
+        # 3D
+        mps_shape = [8, 16, 16, 16]
+        mps = sim.birdcage_maps(mps_shape)
+        ksp = sp.fft(mps, axes=[-1, -2, -3])
+        mps_rec = app.EspiritCalib(ksp, show_pbar=False).run()
+
+        np.testing.assert_allclose(np.abs(mps)[:, 4:-4, 4:-4, 4:-4],
+                                   np.abs(mps_rec[:, 4:-4, 4:-4, 4:-4]),
                                    rtol=1e-2, atol=1e-2)
 
     def test_espirit_maps_eig(self):
-        mps_shape = [8, 32, 32]
+        # 2D
+        mps_shape = [8, 16, 16]
         mps = sim.birdcage_maps(mps_shape)
         ksp = sp.fft(mps, axes=[-1, -2])
         mps_rec, eig_val = app.EspiritCalib(
             ksp, output_eigenvalue=True, show_pbar=False).run()
 
         np.testing.assert_allclose(eig_val, 1, rtol=0.01, atol=0.01)
+
+        # 3D
+        mps_shape = [8, 16, 16, 16]
+        mps = sim.birdcage_maps(mps_shape)
+        ksp = sp.fft(mps, axes=[-1, -2, -3])
+        mps_rec, eig_val = app.EspiritCalib(
+            ksp, output_eigenvalue=True, show_pbar=False).run()
+
+        np.testing.assert_allclose(eig_val, 1, rtol=0.01, atol=0.01)