Merge pull request #324 from effigies/rf/revert_to_old_grid

RF: Generate the b-spline design matrix directly for efficiency

Author: effigies

.gitignore

@@ -110,3 +110,6 @@ venv.bak/
 
 # Mac OSX
 .DS_Store
+
+.*.swp
+line-contributors.txt

sdcflows/tests/test_transform.py

@@ -161,3 +161,23 @@ def test_conversions(tmpdir, testdata_dir, pe_dir):
         fmap_nii.get_fdata(dtype="float32"),
         new_nii.get_fdata(dtype="float32"),
     )
+
+
+def test_grid_bspline_weights():
+    target_shape = (10, 10, 10)
+    target_aff = [[1, 0, 0, -5], [0, 1, 0, -5], [0, 0, 1, -5], [0, 0, 0, 1]]
+    ctrl_shape = (4, 4, 4)
+    ctrl_aff = [[3, 0, 0, -6], [0, 3, 0, -6], [0, 0, 3, -6], [0, 0, 0, 1]]
+
+    weights = tf.grid_bspline_weights(
+        nb.Nifti1Image(np.zeros(target_shape), target_aff),
+        nb.Nifti1Image(np.zeros(ctrl_shape), ctrl_aff),
+    ).tocsr()
+    assert weights.shape == (64, 1000)
+    # Empirically determined numbers intended to indicate that something
+    # significant has changed. If it turns out we've been doing this wrong,
+    # these numbers will probably change.
+    assert np.isclose(weights[0, 0], 0.18919244)
+    assert np.isclose(weights[-1, -1], 0.18919244)
+    assert np.isclose(weights.sum(axis=1).max(), 26.833675)
+    assert np.isclose(weights.sum(axis=1).min(), 1.5879614)

sdcflows/transform.py

@@ -27,8 +27,8 @@
 import numpy as np
 from warnings import warn
 from scipy import ndimage as ndi
-from scipy.interpolate import BSpline
-from scipy.sparse import vstack as sparse_vstack, kron
+from scipy.signal import cubic
+from scipy.sparse import vstack as sparse_vstack, kron, lil_array
 
 import nibabel as nb
 import nitransforms as nt
@@ -405,18 +405,18 @@ def grid_bspline_weights(target_nii, ctrl_nii, dtype="float32"):
         coords[axis] = np.arange(sample_shape[axis], dtype=dtype)
 
         # Calculate the index component of samples w.r.t. B-Spline knots along current axis
-        x = nb.affines.apply_affine(target_to_grid, coords.T)[:, axis]
-        pad_left = max(int(-np.rint(x.min())), 0)
-        pad_right = max(int(np.rint(x.max()) - knots_shape[axis]), 0)
-
-        # BSpline.design_matrix requires all x be within -4 and 4 padding
-        # This padding results from the B-Spline degree (3) plus one
-        t = np.arange(-4 - pad_left, knots_shape[axis] + 4 + pad_right, dtype=dtype)
-
-        # Calculate K x N collocation matrix (discarding extra padding)
-        colloc_ax = BSpline.design_matrix(x, t, 3)[:, (2 + pad_left):-(2 + pad_right)]
-        # Design matrix returns K x N and we want N x K
-        wd.append(colloc_ax.T.tocsr())
+        locs = nb.affines.apply_affine(target_to_grid, coords.T)[:, axis]
+        knots = np.arange(knots_shape[axis], dtype=dtype)
+
+        distance = np.abs(locs[np.newaxis, ...] - knots[..., np.newaxis])
+        within_support = distance < 2.0
+        d_vals, d_idxs = np.unique(distance[within_support], return_inverse=True)
+        bs_w = cubic(d_vals)
+
+        colloc_ax = lil_array((knots_shape[axis], sample_shape[axis]), dtype=dtype)
+        colloc_ax[within_support] = bs_w[d_idxs]
+
+        wd.append(colloc_ax)
 
     # Calculate the tensor product of the three design matrices
     return kron(kron(wd[0], wd[1]), wd[2]).astype(dtype)