Add function to transform from physical space to voxel space

Author: finsberg

pyproject.toml

@@ -17,7 +17,9 @@ dependencies = [
   "numpy",
   "rich-argparse",
   "nibabel",
-  "pandas"]
+  "pandas",
+  "scipy",
+]
 
 [project.optional-dependencies]
 show = [

src/mritk/data/orientation.py

@@ -6,9 +6,71 @@
 
 
 import numpy as np
+
 from .base import MRIData
 
 
+def physical_to_voxel_indices(physical_coordinates: np.ndarray, affine: np.ndarray, round_coords: bool = True) -> np.ndarray:
+    """Transform physical coordinates to voxel indices using an affine matrix.
+
+    This function maps coordinates from physical space (e.g., FEM degrees of freedom)
+    back to the image voxel space by applying the inverse of the provided affine
+    transformation matrix.
+
+    Args:
+        physical_coordinates: (N, 3) array of coordinates in physical space (world coordinates).
+        affine: (4, 4) affine transformation matrix mapping voxel indices to physical space.
+        round_coords: If True, rounds the resulting voxel coordinates to the nearest
+            integer and casts them to `int`. If False, returns floating-point voxel coordinates.
+            Defaults to True.
+
+    Returns:
+        (N, 3) array of voxel indices (or coordinates).
+    """
+    # Note: Assumes apply_affine is available in the scope or imported
+    img_space_coords = apply_affine(np.linalg.inv(affine), physical_coordinates)
+    if round_coords:
+        return np.rint(img_space_coords).astype(int)
+    return img_space_coords
+
+
+def find_nearest_valid_voxels(query_indices: np.ndarray, mask: np.ndarray, k: int) -> np.ndarray:
+    """Find the nearest valid voxels in a mask for a set of query indices.
+
+    Uses a KDTree to find the `k` nearest neighbors for each point in `query_indices`
+    where the neighbors are restricted to positions where `mask` is True.
+
+    Args:
+        query_indices: (N, 3) array of voxel indices (or coordinates) to find neighbors for.
+        mask: Boolean array of shape (X, Y, Z). Neighbors will only be selected from
+            coordinates where this mask is True.
+        k: The number of nearest neighbors to find for each query point.
+
+    Returns:
+        Array of nearest neighbor indices.
+        - If k=1: Returns shape (3, 1, N) containing the coordinates of the single nearest neighbor.
+        - If k>1: Returns shape (3, k, N) containing the coordinates of the k nearest neighbors.
+
+    Raises:
+        ValueError: If the provided mask contains no valid (True) entries.
+    """
+    import scipy.spatial
+
+    valid_inds = np.argwhere(mask)
+    if len(valid_inds) == 0:
+        raise ValueError("No valid indices found in mask.")
+
+    tree = scipy.spatial.KDTree(valid_inds)
+    _, indices = tree.query(query_indices, k=k)
+
+    # Transpose to match the expected output shape (3, k, N) or (3, 1, N)
+    dof_neighbours = valid_inds[indices].T
+
+    if k == 1:
+        dof_neighbours = dof_neighbours[:, np.newaxis, :]
+    return dof_neighbours
+
+
 def apply_affine(T: np.ndarray, X: np.ndarray) -> np.ndarray:
     """Apply a homogeneous affine transformation matrix to a set of points.
 

test/test_data_orientation.py

@@ -0,0 +1,146 @@
+import numpy as np
+import pytest
+import mritk.data.orientation
+
+
+def test_apply_affine_identity():
+    """Test that applying an identity matrix returns the original points."""
+    points = np.array([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]])
+    identity_affine = np.eye(4)
+
+    result = mritk.data.orientation.apply_affine(identity_affine, points)
+
+    np.testing.assert_array_equal(result, points)
+
+
+def test_apply_affine_translation():
+    """Test translation logic: x' = x + t."""
+    points = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
+    affine = np.eye(4)
+    # Set translation vector [Tx, Ty, Tz]
+    translation = np.array([5.0, 10.0, -5.0])
+    affine[:3, 3] = translation
+
+    expected = points + translation
+    result = mritk.data.orientation.apply_affine(affine, points)
+
+    np.testing.assert_array_almost_equal(result, expected)
+
+
+def test_apply_affine_scaling():
+    """Test scaling logic: x' = s * x."""
+    points = np.array([[1.0, 2.0, 3.0]])
+    # Scale x*2, y*0.5, z*-1
+    affine = np.diag([2.0, 0.5, -1.0, 1.0])
+
+    expected = np.array([[2.0, 1.0, -3.0]])
+    result = mritk.data.orientation.apply_affine(affine, points)
+
+    np.testing.assert_array_almost_equal(result, expected)
+
+
+def test_physical_to_voxel_indices_basic_translation():
+    """
+    Test transforming world coordinates back to voxel coordinates.
+    Scenario: The affine translates voxel space by +10.
+    Therefore, a world coordinate of 10 should map back to voxel 0.
+    """
+    # World coordinates (DOFs)
+    dof_coords = np.array([[10.0, 10.0, 10.0], [11.0, 11.0, 11.0]])
+
+    # Affine that adds 10 to everything
+    affine = np.eye(4)
+    affine[:3, 3] = [10.0, 10.0, 10.0]
+
+    # We expect the function to apply the INVERSE of (+10) -> (-10)
+    # 10 - 10 = 0
+    # 11 - 10 = 1
+    expected = np.array([[0, 0, 0], [1, 1, 1]])
+
+    result = mritk.data.orientation.physical_to_voxel_indices(dof_coords, affine, round_coords=True)
+
+    np.testing.assert_array_equal(result, expected)
+    assert result.dtype == int
+
+
+def test_physical_to_voxel_indices_no_rounding():
+    """Test that floating point results are returned when rint=False."""
+    dof_coords = np.array([[10.5, 10.5, 10.5]])
+    affine = np.eye(4)  # Identity
+
+    result = mritk.data.orientation.physical_to_voxel_indices(dof_coords, affine, round_coords=False)
+
+    np.testing.assert_array_almost_equal(result, dof_coords)
+    assert np.issubdtype(result.dtype, np.floating)
+
+
+def test_physical_to_voxel_indices_rounding_behavior():
+    """Test that rint rounds correctly."""
+    # 10.1 -> 10, 10.9 -> 11
+    dof_coords = np.array([[10.1, 10.1, 10.1], [10.9, 10.9, 10.9]])
+    affine = np.eye(4)
+
+    expected = np.array([[10, 10, 10], [11, 11, 11]])
+
+    result = mritk.data.orientation.physical_to_voxel_indices(dof_coords, affine, round_coords=True)
+    np.testing.assert_array_equal(result, expected)
+
+
+def test_find_nearest_valid_voxels_1_neighbor():
+    """Test finding the single closest point in a 2D mask."""
+    # Define a mask with valid pixels only at (0,0) and (5,5)
+    mask = np.zeros((6, 6), dtype=bool)
+    mask[0, 0] = True
+    mask[5, 5] = True
+
+    # Point A is close to (0,0), Point B is close to (5,5)
+    dof_inds = np.array([[0.1, 0.1], [4.9, 4.9]])
+
+    # Function output shape is (ndim, N_neighbors, N_points)
+    result = mritk.data.orientation.find_nearest_valid_voxels(dof_inds, mask, k=1)
+
+    # Verify shape: (2 dims, 1 neighbor, 2 query points)
+    assert result.shape == (2, 1, 2)
+
+    # First point (0.1, 0.1) -> Neighbor should be (0, 0)
+    np.testing.assert_array_equal(result[:, 0, 0], [0, 0])
+    # Second point (4.9, 4.9) -> Neighbor should be (5, 5)
+    np.testing.assert_array_equal(result[:, 0, 1], [5, 5])
+
+
+def test_find_nearest_valid_voxels_N_neighbors():
+    """Test finding multiple neighbors (N=2) in 3D."""
+    mask = np.zeros((10, 10, 10), dtype=bool)
+    # Two valid points close to each other
+    mask[1, 1, 1] = True
+    mask[1, 1, 2] = True
+    # One valid point far away
+    mask[9, 9, 9] = True
+
+    # Query point right next to the cluster at (1,1,1)
+    dof_inds = np.array([[1.0, 1.0, 1.1]])
+
+    result = mritk.data.orientation.find_nearest_valid_voxels(dof_inds, mask, k=2)
+
+    # Shape should be (3 dims, 2 neighbors, 1 point)
+    assert result.shape == (3, 2, 1)
+
+    # Get the neighbors for the first (and only) query point
+    neighbors = result[:, :, 0].T  # Transpose to get list of coords: shape (2, 3)
+
+    # We expect (1,1,1) and (1,1,2) to be the neighbors.
+    # KDTree returns sorted by distance.
+    # Distance to (1,1,1) is 0.1
+    # Distance to (1,1,2) is 0.9
+    # So (1,1,1) should be first.
+    np.testing.assert_array_equal(neighbors[0], [1, 1, 1])
+    np.testing.assert_array_equal(neighbors[1], [1, 1, 2])
+
+
+def test_find_nearest_valid_voxels_empty_mask_error():
+    """Test behavior when no valid points exist (should raise ValueError from KDTree)."""
+    mask = np.zeros((5, 5), dtype=bool)
+    dof_inds = np.array([[1, 1]])
+
+    with pytest.raises(ValueError):
+        mritk.data.orientation.find_nearest_valid_voxels(dof_inds, mask, k=1)