TEST: More comprehensive testing of nonlinear module

tests/test_transform/test_nonlinear.py

@@ -36,6 +36,32 @@ def _random_field():
 
     return nonlinear.DisplacementField(field, src=src, xform=aff)
 
+def _random_affine_field():
+
+    src = _random_image()
+    ref = _random_image()
+
+    # our test field just encodes an affine
+    xform = transform.compose(
+        np.random.randint(2, 5, 3),
+        np.random.randint(1, 10, 3),
+        np.random.random(3))
+
+    rx, ry, rz = np.meshgrid(np.arange(ref.shape[0]),
+                             np.arange(ref.shape[1]),
+                             np.arange(ref.shape[2]), indexing='ij')
+
+    rvoxels  = np.vstack((rx.flatten(), ry.flatten(), rz.flatten())).T
+    rcoords  = transform.transform(rvoxels, ref.voxToScaledVoxMat)
+    scoords  = transform.transform(rcoords, xform)
+
+    field    = np.zeros(list(ref.shape[:3]) + [3])
+    field[:] = (scoords - rcoords).reshape(*it.chain(ref.shape, [3]))
+    field    = nonlinear.DisplacementField(field, src, ref,
+                                           header=ref.header,
+                                           dispType='relative')
+    return field, xform
+
 
 def _field_coords(field):
     vx, vy, vz = field.shape[ :3]
@@ -52,58 +78,94 @@ def test_detectDisplacementType():
     relfield = _random_field()
     coords   = _field_coords(relfield)
     absfield = nonlinear.DisplacementField(
-        relfield.data + coords, xform=relfield.voxToWorldMat)
+        relfield.data + coords,
+        src=relfield.src,
+        xform=relfield.voxToWorldMat)
 
     assert nonlinear.detectDisplacementType(relfield) == 'relative'
     assert nonlinear.detectDisplacementType(absfield) == 'absolute'
 
 
-def test_convertDisplacemenyType():
-    relfield = nonlinear.DisplacementField(_random_field())
-    coords   = _field_coords(relfield)
-    absfield = fslimage.Image(
-        relfield.data + coords, xform=relfield.voxToWorldMat)
+def test_convertDisplacementType():
 
-    relfield = nonlinear.DisplacementField(relfield)
-    absfield = nonlinear.DisplacementField(absfield)
+    relfield = _random_field()
+    coords   = _field_coords(relfield)
+    absfield = nonlinear.DisplacementField(
+        relfield.data + coords,
+        src=relfield.src,
+        xform=relfield.voxToWorldMat)
 
     gotconvrel1 = nonlinear.convertDisplacementType(relfield)
     gotconvabs1 = nonlinear.convertDisplacementType(absfield)
     gotconvrel2 = nonlinear.convertDisplacementType(relfield, 'absolute')
     gotconvabs2 = nonlinear.convertDisplacementType(absfield, 'relative')
 
-    assert np.all(np.isclose(gotconvrel1, absfield.data))
-    assert np.all(np.isclose(gotconvabs1, relfield.data))
-    assert np.all(np.isclose(gotconvrel2, absfield.data))
-    assert np.all(np.isclose(gotconvabs2, relfield.data))
+    tol = dict(atol=1e-5, rtol=1e-5)
 
-def test_DisplacementField_transform():
+    assert np.all(np.isclose(gotconvrel1, absfield.data, **tol))
+    assert np.all(np.isclose(gotconvabs1, relfield.data, **tol))
+    assert np.all(np.isclose(gotconvrel2, absfield.data, **tol))
+    assert np.all(np.isclose(gotconvabs2, relfield.data, **tol))
 
-    src = _random_image()
-    ref = _random_image()
 
-    # our test field just encodes an affine
-    xform = transform.compose(
-        np.random.randint(2, 5, 3),
-        np.random.randint(1, 10, 3),
-        np.random.random(3))
+def test_convertDisplacementSpace():
+
+    basefield, xform = _random_affine_field()
+    src              = basefield.src
+    ref              = basefield.ref
+
+    # generate reference fsl->fsl coordinate mappings
+
+    # For each combination of srcspace->tospace
+    # Generate random coordinates, check that
+    # displacements are correct
+    spaces = ['fsl', 'voxel', 'world']
+    spaces = list(it.combinations_with_replacement(spaces, 2))
+    spaces = spaces + [(r, s) for s, r in spaces]
+    spaces = list(set(spaces))
+
+    for from_, to in spaces:
+
+        refcoords = [np.random.randint(0, basefield.shape[0], 5),
+                     np.random.randint(0, basefield.shape[1], 5),
+                     np.random.randint(0, basefield.shape[2], 5)]
+        refcoords = np.array(refcoords, dtype=np.int).T
+        refcoords = transform.transform(refcoords, ref.voxToScaledVoxMat)
+        srccoords = basefield.transform(refcoords)
+
+        field   = nonlinear.convertDisplacementSpace(basefield, from_, to)
+        premat  = ref.getAffine('fsl', from_)
+        postmat = src.getAffine('fsl', to)
+
+        input  = transform.transform(refcoords, premat)
+        expect = transform.transform(srccoords, postmat)
+
+        got  = field.transform(input)
+        enan = np.isnan(expect)
+        gnan = np.isnan(got)
+
+        assert np.all(np.isclose(enan, gnan))
+        assert np.all(np.isclose(expect[~enan], got[~gnan]))
+
+
+def test_DisplacementField_transform():
+
+    relfield, xform = _random_affine_field()
+    src             = relfield.src
+    ref             = relfield.ref
 
     rx, ry, rz = np.meshgrid(np.arange(ref.shape[0]),
                              np.arange(ref.shape[1]),
                              np.arange(ref.shape[2]), indexing='ij')
-
     rvoxels  = np.vstack((rx.flatten(), ry.flatten(), rz.flatten())).T
     rcoords  = transform.transform(rvoxels, ref.voxToScaledVoxMat)
     scoords  = transform.transform(rcoords, xform)
     svoxels  = transform.transform(scoords, src.scaledVoxToVoxMat)
 
-    relfield    = np.zeros(list(ref.shape[:3]) + [3])
-    relfield[:] = (scoords - rcoords).reshape(*it.chain(ref.shape, [3]))
-    relfield    = nonlinear.DisplacementField(relfield, src, ref,
-                                              dispType='relative')
     absfield    = np.zeros(list(ref.shape[:3]) + [3])
     absfield[:] = scoords.reshape(*it.chain(ref.shape, [3]))
     absfield    = nonlinear.DisplacementField(absfield, src, ref,
+                                              header=ref.header,
                                               dispType='absolute')
 
     got = relfield.transform(rcoords)