ENH: Add displacement field interp

Author: larsoner

mne/tests/test_transforms.py

@@ -52,6 +52,7 @@
     _quat_to_affine,
     _compute_r2,
     _validate_pipeline,
+    _MatchedDisplacementFieldInterpolator,
 )
 
 data_path = testing.data_path(download=False)
@@ -630,3 +631,19 @@ def test_volume_registration():
         ],
         atol=0.001,
     )
+
+
+def test_displacement_field():
+    """Test that our matched point deformation works."""
+    to = np.array([[5, 4, 1], [6, 1, 0], [4, -1, 1], [3, 3, 0]], float)
+    fro = np.array([[0, 2, 2], [2, 2, 1], [2, 0, 2], [0, 0, 1]], float)
+    interp = _MatchedDisplacementFieldInterpolator(fro, to)
+    fro_t = interp(fro)
+    assert_allclose(to, fro_t, atol=1e-12)
+    # check midpoints (should all be decent)
+    for a in range(len(to)):
+        for b in range(a + 1, len(to)):
+            to_ = np.mean(to[[a, b]], axis=0)
+            fro_ = np.mean(fro[[a, b]], axis=0)
+            fro_t = interp(fro_)
+            assert_allclose(to_, fro_t, atol=1e-12)

mne/transforms.py

@@ -2073,3 +2073,47 @@ def apply_volume_registration_points(
     info2.set_montage(montage2)  # converts to head coordinates
 
     return info2, trans2
+
+
+class _MatchedDisplacementFieldInterpolator:
+    """Interpolate from matched points using a displacement field in ND.
+
+    For a demo, see
+    https://gist.github.com/larsoner/fbe32d57996848395854d5e59dff1e10
+    and related tests.
+    """
+
+    def __init__(self, fro, to):
+        from scipy.interpolate import LinearNDInterpolator
+
+        fro = np.array(fro, float)
+        to = np.array(to, float)
+        assert fro.shape == to.shape
+        assert fro.ndim == 2
+        # this restriction is only necessary because it's what
+        # _fit_matched_points requires
+        assert fro.shape[1] == 3
+
+        # Prealign using affine + uniform scaling
+        trans, scale = _fit_matched_points(fro, to, scale=True)
+        trans = _quat_to_affine(trans)
+        trans[:3, :3] *= scale
+        self._affine = trans
+        fro = apply_trans(trans, fro)
+
+        # Add points at extrema
+        delta = (to.max(axis=0) - to.min(axis=0)) / 2.0
+        extrema = np.array([fro.min(axis=0) - delta, fro.max(axis=0) + delta])
+        self._extrema = np.array(np.meshgrid(*extrema.T)).T.reshape(-1, fro.shape[-1])
+        fro_concat = np.concatenate((fro, self._extrema))
+        to_concat = np.concatenate((to, self._extrema))
+
+        # Compute the interpolator (which internally uses Delaunay)
+        self._interp = LinearNDInterpolator(fro_concat, to_concat)
+
+    def __call__(self, x):
+        assert x.ndim in (1, 2) and x.shape[-1] == 3
+        singleton = x.ndim == 1
+        out = self._interp(apply_trans(self._affine, x))
+        out = out[0] if singleton else out
+        return out