Bug fix in DE ssfp/bssfp + more consistent derivative wrt exchange

Author: mcencini

examples/scripts/02-derivatives.py

@@ -105,7 +105,7 @@ def crlb_finitediff_cost(flip, ESP, T1, T2):
 
 # plot derivative
 #%%plots 
-fsz = 20
+fsz = 14
 plt.figure()
 plt.subplot(2,2,1)
 plt.rcParams.update({'font.size': 0.5 * fsz})
@@ -120,7 +120,7 @@ def crlb_finitediff_cost(flip, ESP, T1, T2):
 plt.xlabel("Echo #", fontsize=fsz)
 plt.xlim([-1, 97])
 plt.ylabel(r"$\frac{\partial signal}{\partial T2}$ [a.u.]", fontsize=fsz)
-plt.legend(["Finite Diff", "Auto Diff"])
+plt.legend(["Finite Diff", "Auto Diff"], prop={"size": 14})
 
 
 plt.subplot(2,2,3)
@@ -129,7 +129,7 @@ def crlb_finitediff_cost(flip, ESP, T1, T2):
 plt.xlabel("Echo #", fontsize=fsz)
 plt.xlim([-1, 97])
 plt.ylabel(r"$\frac{\partial CRLB}{\partial FA}$ [a.u.]", fontsize=fsz)
-plt.legend(["Finite Diff", "Auto Diff"])
+plt.legend(["Finite Diff", "Auto Diff"], prop={"size": 14})
 
 plt.subplot(2,2,4)
 
@@ -148,8 +148,8 @@ def crlb_finitediff_cost(flip, ESP, T1, T2):
 # Add some text for labels, title and custom x-axis tick labels, etc.
 plt.ylabel('Execution Time [s]', fontsize=fsz)
 plt.xticks(x, labels, fontsize=fsz)
-# plt.ylim([0, 25])
-plt.legend()
+plt.ylim([0, 50])
+plt.legend(loc='upper left', prop={"size": 14})
 
 plt.bar_label(rects1, padding=3, fontsize=fsz)
 plt.bar_label(rects2, padding=3, fontsize=fsz)

src/epgtorchx/bloch/model/base.py

@@ -273,52 +273,68 @@ def __post_init__(self):  # noqa
                 ((1 - self.weight_mt) * weight_free, self.weight_mt), axis=-1
             )
 
+        # build exchange matrix
+        # if self.model == "bm":
+        #     # build exchange matrix rows
+        #     k0 = 0 * self.kbm
+        #     kiew = torch.cat(
+        #         (k0, self.kbm * self.weight[..., [0]]), axis=-1
+        #     )  # intra-/extra-cellular water
+        #     kmw = torch.cat(
+        #         (self.kbm * self.weight[..., [1]], k0), axis=-1
+        #     )  # myelin water
+        #     self.k = torch.stack((kiew, kmw), axis=-2)
+        # elif self.model == "mt":
+        #     k0 = 0 * self.kmt
+        #     # build exchange matrix rows
+        #     kfree = torch.cat(
+        #         (k0, self.kmt * self.weight[..., [0]]), axis=-1
+        #     )  # myelin water (exchange both with intra-/extra-cellular water and semisolid)
+        #     kbound = torch.cat(
+        #         (self.kmt * self.weight[..., [1]], k0), axis=-1
+        #     )  # semisolid (exchange with myelin water only)
+        #     self.k = torch.stack((kfree, kbound), axis=-2)
+        # elif self.model == "bm-mt":
+        #     k0 = 0 * self.kbm
+        #     # build exchange matrix rows
+        #     kiew = torch.cat(
+        #         (k0, self.kbm * self.weight[..., [1]], k0), axis=-1
+        #     )  # intra-/extra-cellular water (exchange with myelin water only)
+        #     kmw = torch.cat(
+        #         (self.kbm * self.weight[..., [0]], k0, self.kmt * self.weight[..., 2]),
+        #         axis=-1,
+        #     )  # myelin water (exchange both with intra-/extra-cellular water and semisolid)
+        #     kbound = torch.cat(
+        #         (k0, self.kmt * self.weight[..., [1]], k0), axis=-1
+        #     )  # semisolid (exchange with myelin water only)
+        #     self.k = np.stack((kiew, kmw, kbound), axis=-2)
+        # else:
+        #     self.k = None
+
+        # # finalize exchange
+        # if self.k is not None:
+        #     self.k = _particle_conservation(self.k)
+
+        #     # single pool voxels do not exchange
+        #     idx = (self.weight == 1).sum(axis=-1) == 1
+        #     self.k[idx, :, :] = 0.0
+        
         # build exchange matrix
         if self.model == "bm":
-            # build exchange matrix rows
-            k0 = 0 * self.kbm
-            kiew = torch.cat(
-                (k0, self.kbm * self.weight[..., [0]]), axis=-1
-            )  # intra-/extra-cellular water
-            kmw = torch.cat(
-                (self.kbm * self.weight[..., [1]], k0), axis=-1
-            )  # myelin water
-            self.k = torch.stack((kiew, kmw), axis=-2)
+            self.k = self.kbm
         elif self.model == "mt":
-            k0 = 0 * self.kmt
-            # build exchange matrix rows
-            kfree = torch.cat(
-                (k0, self.kmt * self.weight[..., [0]]), axis=-1
-            )  # myelin water (exchange both with intra-/extra-cellular water and semisolid)
-            kbound = torch.cat(
-                (self.kmt * self.weight[..., [1]], k0), axis=-1
-            )  # semisolid (exchange with myelin water only)
-            self.k = torch.stack((kfree, kbound), axis=-2)
+            self.k = self.kmt
         elif self.model == "bm-mt":
-            k0 = 0 * self.kbm
-            # build exchange matrix rows
-            kiew = torch.cat(
-                (k0, self.kbm * self.weight[..., [1]], k0), axis=-1
-            )  # intra-/extra-cellular water (exchange with myelin water only)
-            kmw = torch.cat(
-                (self.kbm * self.weight[..., [0]], k0, self.kmt * self.weight[..., 2]),
-                axis=-1,
-            )  # myelin water (exchange both with intra-/extra-cellular water and semisolid)
-            kbound = torch.cat(
-                (k0, self.kmt * self.weight[..., [1]], k0), axis=-1
-            )  # semisolid (exchange with myelin water only)
-            self.k = np.stack((kiew, kmw, kbound), axis=-2)
+            self.k = torch.cat((self.kbm, self.kmt), axis=-1)
         else:
             self.k = None
 
         # finalize exchange
         if self.k is not None:
-            self.k = _particle_conservation(self.k)
-
             # single pool voxels do not exchange
-            idx = (self.weight == 1).sum(axis=-1) == 1
-            self.k[idx, :, :] = 0.0
-            
+            idx = torch.isclose(self.weight, torch.tensor(1.0)).sum(axis=-1) == 1
+            self.k[idx, :] = 0.0
+                    
         # chemical shift
         if self.model is not None and "bm" in self.model:
             if self.chemshift is not None and self.chemshift_bm is None:
@@ -422,13 +438,16 @@ def get_sim_inputs(self, modelsig):  # noqa
     def reformat(self, input): # noqa
         # handle tuples
         if isinstance(input, (list, tuple)):
-            output = [item[..., 0, 0] + 1j * item[..., -1, 0] for item in input]
-
+            output = [item[..., 0, :] + 1j * item[..., -1, :] for item in input]
+            # output = [torch.diagonal(item, dim1=-2, dim2=-1) if len(item.shape) == 4 else item for item in output]
+            output = [item.reshape(*item.shape[:2], -1) if len(item.shape) == 4 else item for item in output]
+                      
             # stack
             if len(output) == 1:
                 output = output[0]
             else:
-                output = torch.stack(output, dim=1)
+                output = torch.concatenate(output, dim=-1)
+                output = output.permute(2, 0, 1)
         else:
             output = input[..., 0] + 1j * input[..., -1]
 
@@ -543,16 +562,16 @@ def __call__(self, **seq_kwargs):  # noqa
 
 
 # %% local utils
-def _particle_conservation(k):
-    """Adjust diagonal of exchange matrix by imposing particle conservation."""
-    # get shape
-    npools = k.shape[-1]
+# def _particle_conservation(k):
+#     """Adjust diagonal of exchange matrix by imposing particle conservation."""
+#     # get shape
+#     npools = k.shape[-1]
 
-    for n in range(npools):
-        k[..., n, n] = 0.0  # ignore existing diagonal
-        k[..., n, n] = -k[..., n].sum(dim=-1)
+#     for n in range(npools):
+#         k[..., n, n] = 0.0  # ignore existing diagonal
+#         k[..., n, n] = -k[..., n].sum(dim=-1)
 
-    return k
+#     return k
 
 def inspect_signature(input):
     return list(inspect.signature(input).parameters)

src/epgtorchx/bloch/model/bssfpmrf.py

@@ -234,7 +234,7 @@ def sequence(
         # prepare free precession period
         Xpre, Xpost = blocks.bSSFPFidStep(states, TE, TR, T1, T2, weight, k, chemshift)
 
-        for r in (nreps):
+        for r in range(nreps):
             # magnetization prep
             states = Prep(states)
 

src/epgtorchx/bloch/model/ssfpmrf.py

@@ -280,7 +280,7 @@ def sequence(
             states, TE, TR, T1, T2, weight, k, chemshift, D, v, grad_props
         )
 
-        for r in (nreps):
+        for r in range(nreps):
             # magnetization prep
             states = Prep(states)
 

src/epgtorchx/bloch/ops/_relaxation_op.py

@@ -48,6 +48,8 @@ def __init__(
             weight = torch.as_tensor(weight, dtype=torch.float32, device=device)
         if k is not None:
             k = torch.as_tensor(k, dtype=torch.float32, device=device)
+            k = _prepare_exchange(weight, k)
+            
         if df is not None:
             df = torch.as_tensor(df, dtype=torch.float32, device=device)
             df = torch.atleast_1d(df)
@@ -102,6 +104,37 @@ def apply(self, states):
 
 
 # %% local utils
+def _prepare_exchange(weight, k):
+    
+    # prepare
+    if k.shape[-1] == 1: # BM or MT
+        k0 = 0 * k
+        k1 = torch.cat((k0, k * weight[..., [0]]), axis=-1)
+        k2 = torch.cat((k * weight[..., [1]], k0), axis=-1)
+        k = torch.stack((k1, k2), axis=-2)
+    else: # BM-MT
+        k0 = 0 * k[..., [0]]
+        k1 = torch.cat((k0, k[...,[0]] * weight[..., [1]], k0), axis=-1)
+        k2 = torch.cat((k[..., [0]] * weight[..., [0]], k0, k[..., [1]] * weight[..., [2]]), axis=-1)
+        k3 = torch.cat((k0, k[...,[1]] * weight[..., [1]], k0), axis=-1)
+        k = torch.stack((k1, k2, k3), axis=-2)
+    
+
+    # finalize exchange
+    return _particle_conservation(k)
+    
+
+def _particle_conservation(k):
+    """Adjust diagonal of exchange matrix by imposing particle conservation."""
+    # get shape
+    npools = k.shape[-1]
+
+    for n in range(npools):
+        k[..., n, n] = 0.0  # ignore existing diagonal
+        k[..., n, n] = -k[..., n].sum(dim=-1)
+
+    return k
+
 def _transverse_relax_apply(states, E2):
     """
     Apply transverse relaxation operator.