diff --git a/.gitignore b/.gitignore index 73cd142fab..cc963e99c6 100644 --- a/.gitignore +++ b/.gitignore @@ -110,3 +110,6 @@ venv.bak/ # Mac OSX .DS_Store + +.*.swp +line-contributors.txt diff --git a/sdcflows/tests/test_transform.py b/sdcflows/tests/test_transform.py index b4571a07cb..153a71ca6e 100644 --- a/sdcflows/tests/test_transform.py +++ b/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) diff --git a/sdcflows/transform.py b/sdcflows/transform.py index ef630d8e63..742c643e34 100644 --- a/sdcflows/transform.py +++ b/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)