Repository: scientificcomputing/scifem
Pull request #107: Update read_mri_data_to_function
Author: login: cdaversin
Created: 2025-04-07T08:34:21Z | Merged: 2025-04-11T12:03:33Z

Updates to read_mri_data_to_function :

-  Can read nifti Mri data (.nii / .nii.gz), in addition to mgh images (.mgz)
-  Checks the orientation of the given data (has to be RAS)
-  Can take an optional vox2ras transform, if one needs a specific one (default is affine transform from the given MRI data)
-  Can take `use_tkr`as input (optional), to be used in `apply_mri_transform`


Diff:
diff --git a/src/scifem/biomedical.py b/src/scifem/biomedical.py
index f67cdf3..bade023 100644
--- a/src/scifem/biomedical.py
+++ b/src/scifem/biomedical.py
@@ -3,20 +3,26 @@
 from pathlib import Path
 import numpy as np
 import numpy.typing as npt
+import warnings
 
 
 def apply_mri_transform(
-    path: Path, coordinates: npt.NDArray[np.floating], use_tkr: bool = False
+    path: Path,
+    coordinates: npt.NDArray[np.floating],
+    vox2ras_transform: npt.NDArray[np.floating] | None = None,
+    use_tkr: bool = False,
 ) -> npt.NDArray[np.inexact]:
     """Given a path to a set of MRI voxel data, return the data evaluated at the given coordinates.
 
     Args:
         path: Path to the MRI data, should be a file format supported by nibabel.
         coordinates: Coordinates to evaluate the MRI data at.
+        vox2ras_transform: Optional transformation matrix to convert from voxel to ras coordinates.
+            If None, use the transformation matrix from the given MRI data.
         use_tkr: If true, use the old freesurfer `tkregister`, see:
             https://www.mail-archive.com/freesurfer@nmr.mgh.harvard.edu/msg69541.html
             for more details. Else use the standard VOX2RAS transform (equivalent to attaching
-            an affine map to a nibabel image.
+            an affine map to a nibabel image.)
 
     Returns:
         The data evaluated at these points.
@@ -28,21 +34,43 @@ def apply_mri_transform(
         raise ImportError("This function requires the nibabel package to be installed")
     image = nibabel.load(path)
     data = image.get_fdata()
+    # Check if the image is in the correct orientation
+    orientation = nibabel.aff2axcodes(image.affine)
+    if orientation != ("R", "A", "S"):
+        warnings.warn(f"Image orientation is {orientation}, expected (R, A, S). ")
+
+    # Define mgh_header depending on MRI data types (i.e. nifti or mgz)
+    if isinstance(image, nibabel.freesurfer.mghformat.MGHImage):
+        mgh_header = image.header
+    elif isinstance(image, nibabel.nifti1.Nifti1Image):
+        mgh = nibabel.MGHImage(image.dataobj, image.affine)
+        mgh_header = mgh.header
+    else:
+        raise ValueError(
+            f"Unsupported image type: {type(image)} - Supported types are mgz and nifti"
+        )
+
     # This depends on how one uses FreeSurfer
-    if use_tkr:
-        vox2ras = image.header.get_vox2ras_tkr()
+    if vox2ras_transform is not None:
+        vox2ras = vox2ras_transform
     else:
-        # VOX to ras explanation: https://surfer.nmr.mgh.harvard.edu/ftp/articles/vox2ras.pdf
-        vox2ras = image.header.get_vox2ras()
+        if use_tkr:
+            vox2ras = mgh_header.get_vox2ras_tkr()
+        else:
+            # VOX to ras explanation: https://surfer.nmr.mgh.harvard.edu/ftp/articles/vox2ras.pdf
+            vox2ras = mgh_header.get_vox2ras()
+    # Check shape
+    if vox2ras.shape != (4, 4):
+        raise ValueError(f"vox2ras transform must be a 4x4 matrix, got shape {vox2ras.shape}")
+
     ras2vox = np.linalg.inv(vox2ras)
-    ijk_vectorized = naff.apply_affine(ras2vox, coordinates)
 
-    # Round indices to nearest integer
+    ijk_vectorized = naff.apply_affine(ras2vox, coordinates)
     # The file standard assumes that the voxel coordinates refer to the center of each voxel.
     # https://brainder.org/2012/09/23/the-nifti-file-format/
     ijk_rounded = np.rint(ijk_vectorized).astype("int")
-
     assert np.all(ijk_rounded >= 0)
+
     return data[ijk_rounded[:, 0], ijk_rounded[:, 1], ijk_rounded[:, 2]]
 
 
@@ -51,6 +79,8 @@ def read_mri_data_to_tag(
     mesh: dolfinx.mesh.Mesh,
     edim: int,
     entities: npt.NDArray[np.int32] | None = None,
+    vox2ras_transform: npt.NDArray[np.floating] | None = None,
+    use_tkr: bool = False,
 ) -> dolfinx.mesh.MeshTags:
     """Read in MRI data over a set of entities in the mesh and attach it to a mesh tag.
 
@@ -71,7 +101,7 @@ def read_mri_data_to_tag(
         entities = np.arange(num_entities, dtype=np.int32)
 
     midpoints = dolfinx.mesh.compute_midpoints(mesh, edim, entities)
-    data = apply_mri_transform(Path(mri_data_path), midpoints)
+    data = apply_mri_transform(Path(mri_data_path), midpoints, vox2ras_transform, use_tkr=use_tkr)
     et = dolfinx.mesh.meshtags(mesh, edim, entities, data.astype(np.int32))
     return et
 
@@ -80,8 +110,10 @@ def read_mri_data_to_function(
     mri_data_path: str | Path,
     mesh: dolfinx.mesh.Mesh,
     cells: npt.NDArray[np.int32] | None = None,
+    vox2ras_transform: npt.NDArray[np.floating] | None = None,
     degree: int = 0,
     dtype: np.inexact = dolfinx.default_scalar_type,
+    use_tkr: bool = False,
 ) -> dolfinx.fem.Function:
     """Read in MRI data over a set of cells in the mesh and attach it to an appropriate function.
 
@@ -89,10 +121,17 @@ def read_mri_data_to_function(
         mri_data_path: Path to MRI data
         mesh: The mesh to attach the MRI data to.
         cells: Subset of cells to evaluate the MRI data at. If None, all cells are used.
+        vox2ras_transform: Optional transformation matrix to convert from voxel to ras coordinates.
+            If None, use the transformation matrix from the given MRI data.
         degree: Degree of the (quadrature) function space to attach the MRI data to. Defaults to 0.
             If degree is 0, use a DG-0 space instead of a quadrature space, to simplify
             post-processing.
         dtype: Data type used for input data. Can be used for rounding data.
+        use_tkr: If true, use the old freesurfer `tkregister`, see:
+            https://www.mail-archive.com/freesurfer@nmr.mgh.harvard.edu/msg69541.html
+            for more details. Else use the standard VOX2RAS transform (equivalent to attaching
+            an affine map to a nibabel image.)
+
 
     Raises:
         ValueError: If degree is negative.
@@ -117,7 +156,9 @@ def read_mri_data_to_function(
     dof_positions = Vh.tabulate_dof_coordinates()
     dofmap = Vh.dofmap.list
     dofmap_pos = dofmap[cells].flatten()
-    data = apply_mri_transform(Path(mri_data_path), dof_positions[dofmap_pos])
+    data = apply_mri_transform(
+        Path(mri_data_path), dof_positions[dofmap_pos], vox2ras_transform, use_tkr=use_tkr
+    )
 
     v = dolfinx.fem.Function(Vh)
     v.x.array[dofmap_pos] = data.astype(dtype)
diff --git a/tests/test_read_mri_data.py b/tests/test_read_mri_data.py
index 4333c0b..2110c23 100644
--- a/tests/test_read_mri_data.py
+++ b/tests/test_read_mri_data.py
@@ -13,7 +13,12 @@
 @pytest.mark.parametrize("Nz", [9, 17])
 @pytest.mark.parametrize("theta", [np.pi / 3, 0, -np.pi])
 @pytest.mark.parametrize("translation", [np.array([0, 0, 0]), np.array([2.1, 1.3, 0.4])])
-def test_read_mri_data_to_function(degree, M, Nx, Ny, Nz, theta, translation, tmp_path):
+@pytest.mark.parametrize("mri_data_format", ["nifti", "mgh"])
+@pytest.mark.parametrize("external_affine", [True, False])
+@pytest.mark.parametrize("use_tkr", [False])
+def test_read_mri_data_to_function(
+    degree, M, Nx, Ny, Nz, theta, translation, tmp_path, mri_data_format, external_affine, use_tkr
+):
     nibabel = pytest.importorskip("nibabel")
 
     # Generate rotation and scaling matrix equivalent to a unit cube
@@ -25,14 +30,25 @@ def test_read_mri_data_to_function(degree, M, Nx, Ny, Nz, theta, translation, tm
     # Generate the affine mapping for nibabel
     A = np.append(np.dot(rotation_matrix_3D, scale_matrix), (translation).reshape(3, 1), axis=1)
     A = np.vstack([A, [0, 0, 0, 1]])
+    Id = np.identity(4)
 
     # Write transformed data to file
-    data = np.arange(1, M**3 + 1, dtype=np.float64).reshape(M, M, M)
-    image = nibabel.nifti1.Nifti1Image(data, affine=A)
-
     path = MPI.COMM_WORLD.bcast(tmp_path, root=0)
+    if mri_data_format == "nifti":
+        data = np.arange(1, M**3 + 1, dtype=np.float64).reshape(M, M, M)
+        image = nibabel.nifti1.Nifti1Image(data, affine=Id if external_affine else A)
+    elif mri_data_format == "mgh":
+        data = np.arange(1, M**3 + 1, dtype=np.int32).reshape(M, M, M)
+        image = nibabel.freesurfer.mghformat.MGHImage(data, affine=Id if external_affine else A)
+
+    # Reorient the image to RAS
+    orig_ornt = nibabel.io_orientation(image.affine)
+    targ_ornt = nibabel.orientations.axcodes2ornt("RAS")
+    transform = nibabel.orientations.ornt_transform(orig_ornt, targ_ornt)
+    image = image.as_reoriented(transform)
     filename = path.with_suffix(".mgz")
 
+    # Save the image to file
     if MPI.COMM_WORLD.rank == 0:
         nibabel.save(image, filename)
     MPI.COMM_WORLD.Barrier()
@@ -69,7 +85,14 @@ def test_read_mri_data_to_function(degree, M, Nx, Ny, Nz, theta, translation, tm
     mesh.geometry.x[:] = np.dot(rotation_matrix_3D, shifted_coords.T).T + translation
 
     # # Read data to function
-    func = read_mri_data_to_function(filename, mesh, degree=degree, dtype=np.float64)
+    if external_affine:
+        func = read_mri_data_to_function(
+            filename, mesh, degree=degree, dtype=np.float64, vox2ras_transform=A, use_tkr=use_tkr
+        )
+    else:
+        func = read_mri_data_to_function(
+            filename, mesh, degree=degree, dtype=np.float64, use_tkr=use_tkr
+        )
 
     np.testing.assert_allclose(func.x.array, reference_data)
 
@@ -78,5 +101,17 @@ def test_read_mri_data_to_function(degree, M, Nx, Ny, Nz, theta, translation, tm
         num_cells_local = cell_map.size_local + cell_map.num_ghosts
         # Pick every second cell
         cells = np.arange(num_cells_local, dtype=np.int32)[::2]
-        tag = read_mri_data_to_tag(filename, mesh, edim=mesh.topology.dim, entities=cells)
+        if external_affine:
+            tag = read_mri_data_to_tag(
+                filename,
+                mesh,
+                edim=mesh.topology.dim,
+                entities=cells,
+                vox2ras_transform=A,
+                use_tkr=use_tkr,
+            )
+        else:
+            tag = read_mri_data_to_tag(
+                filename, mesh, edim=mesh.topology.dim, entities=cells, use_tkr=use_tkr
+            )
         np.testing.assert_allclose(tag.values, reference_data[::2])


Reviews:
user: login: jorgensd
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-07T08:44:46Z
user: login: jorgensd
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-07T08:45:45Z
user: login: jorgensd
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-07T08:47:17Z
user: login: cdaversin
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T09:13:45Z
user: login: jorgensd
body: Looks good from my end.
@finsberg do you mind having a look?
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-07T11:05:12Z
user: login: jorgensd
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-11T12:03:26Z

Review comments:
diff_hunk: @@ -16,7 +19,7 @@ def apply_mri_transform(
         use_tkr: If true, use the old freesurfer `tkregister`, see:
             https://www.mail-archive.com/freesurfer@nmr.mgh.harvard.edu/msg69541.html
             for more details. Else use the standard VOX2RAS transform (equivalent to attaching
-            an affine map to a nibabel image.
+            an affine map to a nibabel image.)
path: src/scifem/biomedical.py
user: login: jorgensd
body: Here we need to add documentation  of `vox2ras_transform`.


created_at: 2025-04-07T08:44:45Z
author_association: MEMBER
diff_hunk: @@ -28,21 +31,41 @@ def apply_mri_transform(
         raise ImportError("This function requires the nibabel package to be installed")
     image = nibabel.load(path)
     data = image.get_fdata()
+    # Check if the image is in the correct orientation
+    orientation = nibabel.aff2axcodes(image.affine)
+    assert orientation == ("R", "A", "S")
+
+    # Define mgh_header depending on MRI data types (i.e. nifti or mgz)
+    if isinstance(image, nibabel.freesurfer.mghformat.MGHImage):
+        mgh_header = image.header
+    elif isinstance(image, nibabel.nifti1.Nifti1Image):
+        mgh = nibabel.MGHImage(image.dataobj, image.affine)
+        orientation = nibabel.aff2axcodes(mgh.affine)
path: src/scifem/biomedical.py
user: login: jorgensd
body: Is the orientation needed here?
created_at: 2025-04-07T08:45:45Z
author_association: MEMBER
diff_hunk: @@ -80,19 +103,28 @@ def read_mri_data_to_function(
     mri_data_path: str | Path,
     mesh: dolfinx.mesh.Mesh,
     cells: npt.NDArray[np.int32] | None = None,
+    vox2ras_transform: npt.NDArray[np.floating] | None = None,
     degree: int = 0,
     dtype: np.inexact = dolfinx.default_scalar_type,
+    use_tkr: bool = False,
 ) -> dolfinx.fem.Function:
     """Read in MRI data over a set of cells in the mesh and attach it to an appropriate function.
 
     Args:
         mri_data_path: Path to MRI data
         mesh: The mesh to attach the MRI data to.
         cells: Subset of cells to evaluate the MRI data at. If None, all cells are used.
+        vox2ras_transform: Optional transformation matrix to convert from voxel to ras coordinates.
path: src/scifem/biomedical.py
user: login: jorgensd
body: Should we sanity check the size of this transform (4x4), or let nibabel handle it?
created_at: 2025-04-07T08:47:16Z
author_association: MEMBER
diff_hunk: @@ -28,21 +31,41 @@ def apply_mri_transform(
         raise ImportError("This function requires the nibabel package to be installed")
     image = nibabel.load(path)
     data = image.get_fdata()
+    # Check if the image is in the correct orientation
+    orientation = nibabel.aff2axcodes(image.affine)
+    assert orientation == ("R", "A", "S")
+
+    # Define mgh_header depending on MRI data types (i.e. nifti or mgz)
+    if isinstance(image, nibabel.freesurfer.mghformat.MGHImage):
+        mgh_header = image.header
+    elif isinstance(image, nibabel.nifti1.Nifti1Image):
+        mgh = nibabel.MGHImage(image.dataobj, image.affine)
+        orientation = nibabel.aff2axcodes(mgh.affine)
path: src/scifem/biomedical.py
user: login: cdaversin
body: Nope, I removed it :)
created_at: 2025-04-07T09:13:45Z
author_association: CONTRIBUTOR

PR comments:
user: login: jorgensd
created_at: 2025-04-07T08:47:38Z
body: A test that illustrates how to use it would be nice 👼🏻 
author_association: MEMBER
user: login: finsberg
created_at: 2025-04-07T11:14:00Z
body: > Looks good from my end. @finsberg do you mind having a look?

Looks good, but would be nice to have some tests that test the following
- User provided `vox2ras_transform` give correct result
- User provided `vox2ras_transform` with wrong shape raises ValueError
- Test for both formats (i.e `nibabel.freesurfer.mghformat.MGHImage` and `nibabel.nifti1.Nifti1Image`) if possible.
author_association: MEMBER
user: login: cdaversin
created_at: 2025-04-07T12:08:38Z
body: I agree, I will implement some tests ASAP
author_association: CONTRIBUTOR
user: login: cdaversin
created_at: 2025-04-11T09:35:26Z
body: The tests are now updated with both formats (nifti and mgh), and with user provided vox2ras affine transform (nibabel already handles the case where vox2ras doesn't have the right shape).  All tests passed :partying_face: 

Question : I have added `use_tkr` as an input to `read_mri_data_to_function` since it was in `apply_mri_transform` but could not be set by users. The tests fail with `use_tkr = True` due to assertion `np.all(ijk_rounded >= 0)` which fails in this case. I guess this is because the data we build in the test has RAS orientation and not tkRAS. I don't think we need to specifically test the "tkr case" but we should still give the possibility to use this parameter ?
author_association: CONTRIBUTOR
user: login: jorgensd
created_at: 2025-04-11T09:39:17Z
body: > The tests are now updated with both formats (nifti and mgh), and with user provided vox2ras affine transform (nibabel already handles the case where vox2ras doesn't have the right shape). All tests passed 🥳
> 
> Question : I have added `use_tkr` as an input to `read_mri_data_to_function` since it was in `apply_mri_transform` but could not be set by users. The tests fail with `use_tkr = True` due to assertion `np.all(ijk_rounded >= 0)` which fails in this case. I guess this is because the data we build in the test has RAS orientation and not tkRAS. I don't think we need to specifically test the "tkr case" but we should still give the possibility to use this parameter ?

Yes, we can’t test it. But we should still give users the option to use it, I think!
author_association: MEMBER