cmr_rnd_diffusion/__init__.py

@@ -3,4 +3,4 @@ import cmr_rnd_diffusion.diff3d
 import cmr_rnd_diffusion.mesh_utils
 import cmr_rnd_diffusion.sampling
 import cmr_rnd_diffusion.vtk_utils
-from cmr_rnd_diffusion._cardiac_mesh_dataset import CardiacMeshDataset, SliceDataset
+from cmr_rnd_diffusion._cardiac_mesh_dataset import BaseDataset, CardiacMeshDataset, DiffusionDataset

cmr_rnd_diffusion/diff3d.py

@@ -134,7 +134,7 @@ def generate_ha(transmural_position: np.ndarray,
     return np.deg2rad(helix_angles)
 
 
-def compute_local_basis(positional_vectors: np.ndarray, mesh_numbers: Tuple[int, int, int]):
+def compute_local_basis(positional_vectors: np.ndarray, mesh_numbers: Tuple[int, int, int]) -> dict:
     """ Computes the local coordinate system (radial, circumferential, longitudinal) for
     a left ventricle mesh.
 
@@ -152,6 +152,7 @@ def compute_local_basis(positional_vectors: np.ndarray, mesh_numbers: Tuple[int,
     mesh_vecs = positional_vectors.reshape(mesh_numbers[0], points_per_shell, 3, order="C")[:, 1:]
     mesh_vecs = mesh_vecs.reshape(mesh_numbers[0], mesh_numbers[1], mesh_numbers[2], 3, order="C")
     mesh_vecs_aug = np.concatenate([mesh_vecs, mesh_vecs[:, :, 0:1]], axis=2)
+
     e_c = np.diff(mesh_vecs_aug, axis=2).reshape(mesh_numbers[0], -1, 3)
     e_c = np.concatenate([e_c, np.tile(np.array([[[0., 1., 0.]]]), [mesh_numbers[0], 1, 1])],
                          axis=1).reshape(-1, 3)
@@ -166,11 +167,14 @@ def compute_local_basis(positional_vectors: np.ndarray, mesh_numbers: Tuple[int,
     e_l = np.gradient(mesh_vecs, axis=1)
     e_l = np.concatenate([e_l[:, 0:1, 0], e_l[:, 1:].reshape(mesh_numbers[0], -1, 3)], axis=1)
     e_l = e_l.reshape(-1, 3)
-    e_l /= np.linalg.norm(e_l, axis=-1, keepdims=True)
+    norm_l = np.linalg.norm(e_l, axis=-1, keepdims=True)
+    e_l = np.divide(e_l, norm_l, where=norm_l > 0)
 
     # Compute radial vector as cross product
     e_r = np.cross(e_c, e_l)
-    e_r /= np.linalg.norm(e_r, axis=-1, keepdims=True)
+    norm_r = np.linalg.norm(e_r, axis=-1, keepdims=True)
+
+    e_r = np.divide(e_r, norm_r, where=norm_r > 0)
 
     return_dict = dict()
     return_dict["e_r"] = e_r
@@ -178,6 +182,7 @@ def compute_local_basis(positional_vectors: np.ndarray, mesh_numbers: Tuple[int,
     return_dict["e_c"] = e_c
     return return_dict
 
+
 def create_eigenbasis(local_basis: np.ndarray, helix_angles: np.ndarray,
                       e2_angles: np.ndarray) -> np.ndarray:
     """ Calculates the local diffusion tensor eigen-basis from helix angles and
@@ -207,7 +212,11 @@ def create_eigenbasis(local_basis: np.ndarray, helix_angles: np.ndarray,
     ev2 = - sin_sa * temp3 + cos_sa * temp2
     ev3 = cos_sa * temp3 + sin_sa * temp2
 
-    eigen_vectors = [ev / tf.linalg.norm(ev, axis=-1, keepdims=True) for ev in (ev1, ev2, ev3)]
+    norms = [tf.linalg.norm(ev, axis=-1, keepdims=True) for ev in (ev1, ev2, ev3)]
+    eigen_vectors = [np.divide(ev, n, where=n > 0) for ev, n in zip((ev1, ev2, ev3), norms)]
+    for ev in eigen_vectors:
+        ev[np.where(np.logical_not(np.isfinite(ev)))] = 0
+
     local_eigen_basis = tf.stack(eigen_vectors, axis=-1)
     return local_eigen_basis.numpy()