REF: Refactor motion-related data computation functions

Refactor motion-related data computation functions into reusable code
from the PET notebook.

Rename the existing FD function to distinguish it from the refactored
function.

Author: jhlegarreta

docs/notebooks/PET_example.ipynb

@@ -2445,33 +2445,7 @@
     "import numpy as np\n",
     "import pandas as pd\n",
     "\n",
-    "def extract_motion_parameters(affine):\n",
-    "    \"\"\"Extract translation (mm) and rotation (degrees) parameters from an affine matrix.\"\"\"\n",
-    "    translation = affine[:3, 3]\n",
-    "    rotation_rad = np.arctan2(\n",
-    "        [affine[2, 1], affine[0, 2], affine[1, 0]],\n",
-    "        [affine[2, 2], affine[0, 0], affine[1, 1]]\n",
-    "    )\n",
-    "    rotation_deg = np.rad2deg(rotation_rad)\n",
-    "    return (*translation, *rotation_deg)\n",
-    "\n",
-    "def compute_fd(motion_parameters):\n",
-    "    \"\"\"Compute Framewise Displacement from motion parameters.\"\"\"\n",
-    "    translations = motion_parameters[:, :3]\n",
-    "    rotations_deg = motion_parameters[:, 3:]\n",
-    "    rotations_rad = np.deg2rad(rotations_deg)\n",
-    "    \n",
-    "    # Compute differences between consecutive frames\n",
-    "    d_translations = np.vstack([np.zeros((1, 3)), np.diff(translations, axis=0)])\n",
-    "    d_rotations = np.vstack([np.zeros((1, 3)), np.diff(rotations_rad, axis=0)])\n",
-    "    \n",
-    "    # Convert rotations from radians to displacement on a sphere (radius 50 mm)\n",
-    "    radius = 50  # typical head radius in mm\n",
-    "    rotation_displacement = d_rotations * radius\n",
-    "    \n",
-    "    # Compute FD as sum of absolute differences\n",
-    "    fd = np.sum(np.abs(d_translations) + np.abs(rotation_displacement), axis=1)\n",
-    "    return fd\n",
+    "from nifreeze.registration.utils import displacement_framewise_motion, extract_motion_parameters\n",
     "\n",
     "# Assume 'affines' is the list of affine matrices you computed earlier\n",
     "motion_parameters = []\n",
@@ -2481,7 +2455,7 @@
     "    motion_parameters.append([tx, ty, tz, rx, ry, rz])\n",
     "\n",
     "motion_parameters = np.array(motion_parameters)\n",
-    "fd = compute_fd(motion_parameters)\n",
+    "fd = displacement_framewise_motion(motion_parameters)\n",
     "\n",
     "# Creating a DataFrame for better visualization\n",
     "df_motion = pd.DataFrame({\n",

docs/notebooks/bold_realignment.ipynb

@@ -478,7 +478,7 @@
    "source": [
     "from nitransforms.resampling import apply\n",
     "\n",
-    "from nifreeze.registration.utils import displacement_framewise\n",
+    "from nifreeze.registration.utils import displacement_framewise_transformation\n",
     "\n",
     "afni_fd = {}\n",
     "nitransforms_fd = {}\n",
@@ -501,7 +501,9 @@
     "    ]\n",
     "\n",
     "    nii = nb.load(DATA_PATH / bold_run)\n",
-    "    nitransforms_fd[str(bold_run)] = np.array([displacement_framewise(nii, xfm) for xfm in xfms])\n",
+    "    nitransforms_fd[str(bold_run)] = np.array(\n",
+    "        [displacement_framewise_transformation(nii, xfm) for xfm in xfms]\n",
+    "    )\n",
     "\n",
     "    hmc_xfm = nt.linear.LinearTransformsMapping(xfms)\n",
     "    out_nitransforms = (\n",
@@ -520,7 +522,7 @@
     "        OUTPUT_DIR / bold_run.parent / f\"{bold_run.name.rsplit('_', 1)[0]}_desc-hmc_xfm.txt\"\n",
     "    )\n",
     "    afni_fd[str(bold_run)] = np.array(\n",
-    "        [displacement_framewise(nii, afni_xfms[i]) for i in range(len(afni_xfms))]\n",
+    "        [displacement_framewise_transformation(nii, afni_xfms[i]) for i in range(len(afni_xfms))]\n",
     "    )\n",
     "\n",
     "    out_afni = (\n",

src/nifreeze/registration/utils.py

@@ -32,11 +32,15 @@
 from __future__ import annotations
 
 from itertools import product
+from typing import Tuple
 
 import nibabel as nb
 import nitransforms as nt
 import numpy as np
 
+RADIUS = 50.0
+"""Typical radius (in mm) of a sphere mimicking the size of a typical human brain."""
+
 
 def displacements_within_mask(
     mask_img: nb.spatialimages.SpatialImage,
@@ -79,11 +83,11 @@ def displacements_within_mask(
     return np.linalg.norm(diffs, axis=-1)
 
 
-def displacement_framewise(
+def displacement_framewise_transform(
     img: nb.spatialimages.SpatialImage,
     test_xfm: nt.base.BaseTransform,
-    radius: float = 50.0,
-):
+    radius: float = RADIUS,
+) -> float:
     """
     Compute the framewise displacement (FD) for a given transformation.
 
@@ -95,7 +99,6 @@ def displacement_framewise(
         The transformation to test. Applied to coordinates around the image center.
     radius : :obj:`float`, optional
         The radius (in mm) of the spherical neighborhood around the center of the image.
-        Default is 50.0 mm.
 
     Returns
     -------
@@ -112,3 +115,63 @@ def displacement_framewise(
     fd_coords = np.array(list(product(*((radius, -radius),) * 3))) + center_xyz
     # Compute the average displacement from the test transformation
     return np.mean(np.linalg.norm(test_xfm.map(fd_coords) - fd_coords, axis=-1))
+
+
+def displacement_framewise_motion(
+    motion_parameters: np.ndarray, radius: float = RADIUS
+) -> np.ndarray:
+    """Compute framewise displacement (FD) from motion parameters.
+
+    Each row in the motion parameters represents one frame, and columns
+    represent each coordinate axis ``x``, `y``, and ``z``. Translation
+    parameters are followed by rotation parameters column-wise.
+
+    Parameters
+    ----------
+    motion_parameters : :obj:`numpy.ndarray`
+        Motion parameters.
+    radius : :obj:`float`, optional
+        Radius (in mm) of a sphere mimicking the size of a typical human brain.
+
+    Returns
+    -------
+    :obj:`numpy.ndarray`
+        The framewise displacement (FD) as the sum of absolute differences
+        between consecutive frames.
+    """
+
+    translations = motion_parameters[:, :3]
+    rotations_deg = motion_parameters[:, 3:]
+    rotations_rad = np.deg2rad(rotations_deg)
+
+    # Compute differences between consecutive frames
+    d_translations = np.vstack([np.zeros((1, 3)), np.diff(translations, axis=0)])
+    d_rotations = np.vstack([np.zeros((1, 3)), np.diff(rotations_rad, axis=0)])
+
+    # Convert rotations from radians to displacement on a sphere
+    rotation_displacement = d_rotations * radius
+
+    # Compute FD as sum of absolute differences
+    return np.sum(np.abs(d_translations) + np.abs(rotation_displacement), axis=1)
+
+
+def extract_motion_parameters(affine: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Extract translation (mm) and rotation (degrees) parameters from an affine matrix.
+
+    Parameters
+    ----------
+    affine : :obj:`~numpy.ndarray`
+        The affine transformation matrix.
+
+    Returns
+    -------
+    :obj:`tuple`
+        Extracted translation and rotation parameters.
+    """
+
+    translation = affine[:3, 3]
+    rotation_rad = np.arctan2(
+        [affine[2, 1], affine[0, 2], affine[1, 0]], [affine[2, 2], affine[0, 0], affine[1, 1]]
+    )
+    rotation_deg = np.rad2deg(rotation_rad)
+    return *translation, *rotation_deg