Skip to content

Commit e8cbf05

Browse files
committed
Merge pull request #85 from matthew-brett/add-apply-affine
NF - add affines code from nipy, use for trackvis Adds - in particular - an apply_affine method useful for applying affines to arrays of points.
2 parents 2a96c40 + 46d3e71 commit e8cbf05

File tree

3 files changed

+238
-7
lines changed

3 files changed

+238
-7
lines changed

nibabel/affines.py

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
2+
# vi: set ft=python sts=4 ts=4 sw=4 et:
3+
""" Utility routines for working with points and affine transforms
4+
"""
5+
6+
import numpy as np
7+
8+
9+
def apply_affine(aff, pts):
10+
""" Apply affine matrix `aff` to points `pts`
11+
12+
Returns result of application of `aff` to the *right* of `pts`. The
13+
coordinate dimension of `pts` should be the last.
14+
15+
For the 3D case, `aff` will be shape (4,4) and `pts` will have final axis
16+
length 3 - maybe it will just be N by 3. The return value is the transformed
17+
points, in this case::
18+
19+
res = np.dot(aff[:3,:3], pts.T) + aff[:3,3:4]
20+
transformed_pts = res.T
21+
22+
Notice though, that this routine is more general, in that `aff` can have any
23+
shape (N,N), and `pts` can have any shape, as long as the last dimension is
24+
for the coordinates, and is therefore length N-1.
25+
26+
Parameters
27+
----------
28+
aff : (N, N) array-like
29+
Homogenous affine, for 3D points, will be 4 by 4. Contrary to first
30+
appearance, the affine will be applied on the left of `pts`.
31+
pts : (..., N-1) array-like
32+
Points, where the last dimension contains the coordinates of each point.
33+
For 3D, the last dimension will be length 3.
34+
35+
Returns
36+
-------
37+
transformed_pts : (..., N-1) array
38+
transformed points
39+
40+
Examples
41+
--------
42+
>>> aff = np.array([[0,2,0,10],[3,0,0,11],[0,0,4,12],[0,0,0,1]])
43+
>>> pts = np.array([[1,2,3],[2,3,4],[4,5,6],[6,7,8]])
44+
>>> apply_affine(aff, pts)
45+
array([[14, 14, 24],
46+
[16, 17, 28],
47+
[20, 23, 36],
48+
[24, 29, 44]])
49+
50+
Just to show that in the simple 3D case, it is equivalent to:
51+
52+
>>> (np.dot(aff[:3,:3], pts.T) + aff[:3,3:4]).T
53+
array([[14, 14, 24],
54+
[16, 17, 28],
55+
[20, 23, 36],
56+
[24, 29, 44]])
57+
58+
But `pts` can be a more complicated shape:
59+
60+
>>> pts = pts.reshape((2,2,3))
61+
>>> apply_affine(aff, pts)
62+
array([[[14, 14, 24],
63+
[16, 17, 28]],
64+
<BLANKLINE>
65+
[[20, 23, 36],
66+
[24, 29, 44]]])
67+
"""
68+
aff = np.asarray(aff)
69+
pts = np.asarray(pts)
70+
shape = pts.shape
71+
pts = pts.reshape((-1, shape[-1]))
72+
# rzs == rotations, zooms, shears
73+
rzs = aff[:-1,:-1]
74+
trans = aff[:-1,-1]
75+
res = np.dot(pts, rzs.T) + trans[None,:]
76+
return res.reshape(shape)
77+
78+
79+
def append_diag(aff, steps, starts=()):
80+
""" Add diagonal elements `steps` and translations `starts` to affine
81+
82+
Typical use is in expanding 4x4 affines to larger dimensions. Nipy is the
83+
main consumer because it uses NxM affines, whereas we generally only use 4x4
84+
affines; the routine is here for convenience.
85+
86+
Parameters
87+
----------
88+
aff : 2D array
89+
N by M affine matrix
90+
steps : scalar or sequence
91+
diagonal elements to append.
92+
starts : scalar or sequence
93+
elements to append to last column of `aff`, representing translations
94+
corresponding to the `steps`. If empty, expands to a vector of zeros
95+
of the same length as `steps`
96+
97+
Returns
98+
-------
99+
aff_plus : 2D array
100+
Now P by Q where L = ``len(steps)`` and P == N+L, Q=N+L
101+
102+
Examples
103+
--------
104+
>>> aff = np.eye(4)
105+
>>> aff[:3,:3] = np.arange(9).reshape((3,3))
106+
>>> append_diag(aff, [9, 10], [99,100])
107+
array([[ 0., 1., 2., 0., 0., 0.],
108+
[ 3., 4., 5., 0., 0., 0.],
109+
[ 6., 7., 8., 0., 0., 0.],
110+
[ 0., 0., 0., 9., 0., 99.],
111+
[ 0., 0., 0., 0., 10., 100.],
112+
[ 0., 0., 0., 0., 0., 1.]])
113+
"""
114+
aff = np.asarray(aff)
115+
steps = np.atleast_1d(steps)
116+
starts = np.atleast_1d(starts)
117+
n_steps = len(steps)
118+
if len(starts) == 0:
119+
starts = np.zeros(n_steps, dtype=steps.dtype)
120+
elif len(starts) != n_steps:
121+
raise ValueError('Steps should have same length as starts')
122+
old_n_out, old_n_in = aff.shape[0]-1, aff.shape[1]-1
123+
# make new affine
124+
aff_plus = np.zeros((old_n_out + n_steps + 1,
125+
old_n_in + n_steps + 1), dtype=aff.dtype)
126+
# Get stuff from old affine
127+
aff_plus[:old_n_out,:old_n_in] = aff[:old_n_out, :old_n_in]
128+
aff_plus[:old_n_out,-1] = aff[:old_n_out,-1]
129+
# Add new diagonal elements
130+
for i, el in enumerate(steps):
131+
aff_plus[old_n_out+i, old_n_in+i] = el
132+
# Add translations for new affine, plus last 1
133+
aff_plus[old_n_out:,-1] = list(starts) + [1]
134+
return aff_plus

nibabel/tests/test_affines.py

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
2+
# vi: set ft=python sts=4 ts=4 sw=4 et:
3+
4+
import numpy as np
5+
6+
from ..affines import apply_affine, append_diag
7+
8+
from nose.tools import assert_equal, assert_raises
9+
from numpy.testing import assert_array_equal, assert_almost_equal, \
10+
assert_array_almost_equal
11+
12+
13+
def validated_apply_affine(T, xyz):
14+
# This was the original apply_affine implementation that we've stashed here
15+
# to test against
16+
xyz = np.asarray(xyz)
17+
shape = xyz.shape[0:-1]
18+
XYZ = np.dot(np.reshape(xyz, (np.prod(shape), 3)), T[0:3,0:3].T)
19+
XYZ[:,0] += T[0,3]
20+
XYZ[:,1] += T[1,3]
21+
XYZ[:,2] += T[2,3]
22+
XYZ = np.reshape(XYZ, shape+(3,))
23+
return XYZ
24+
25+
26+
def test_apply_affine():
27+
rng = np.random.RandomState(20110903)
28+
aff = np.diag([2, 3, 4, 1])
29+
pts = rng.uniform(size=(4,3))
30+
assert_array_equal(apply_affine(aff, pts),
31+
pts * [[2, 3, 4]])
32+
aff[:3,3] = [10, 11, 12]
33+
assert_array_equal(apply_affine(aff, pts),
34+
pts * [[2, 3, 4]] + [[10, 11, 12]])
35+
aff[:3,:] = rng.normal(size=(3,4))
36+
exp_res = np.concatenate((pts.T, np.ones((1,4))), axis=0)
37+
exp_res = np.dot(aff, exp_res)[:3,:].T
38+
assert_array_equal(apply_affine(aff, pts), exp_res)
39+
# Check we get the same result as the previous implementation
40+
assert_almost_equal(validated_apply_affine(aff, pts), apply_affine(aff, pts))
41+
# Check that lists work for inputs
42+
assert_array_equal(apply_affine(aff.tolist(), pts.tolist()), exp_res)
43+
# Check that it's the same as a banal implementation in the simple case
44+
aff = np.array([[0,2,0,10],[3,0,0,11],[0,0,4,12],[0,0,0,1]])
45+
pts = np.array([[1,2,3],[2,3,4],[4,5,6],[6,7,8]])
46+
exp_res = (np.dot(aff[:3,:3], pts.T) + aff[:3,3:4]).T
47+
assert_array_equal(apply_affine(aff, pts), exp_res)
48+
# That points can be reshaped and you'll get the same shape output
49+
pts = pts.reshape((2,2,3))
50+
exp_res = exp_res.reshape((2,2,3))
51+
assert_array_equal(apply_affine(aff, pts), exp_res)
52+
# That ND also works
53+
for N in range(2,6):
54+
aff = np.eye(N)
55+
nd = N-1
56+
aff[:nd,:nd] = rng.normal(size=(nd,nd))
57+
pts = rng.normal(size=(2,3,nd))
58+
res = apply_affine(aff, pts)
59+
# crude apply
60+
new_pts = np.ones((N,6))
61+
new_pts[:-1,:] = np.rollaxis(pts, -1).reshape((nd,6))
62+
exp_pts = np.dot(aff, new_pts)
63+
exp_pts = np.rollaxis(exp_pts[:-1,:], 0, 2)
64+
exp_res = exp_pts.reshape((2,3,nd))
65+
assert_array_almost_equal(res, exp_res)
66+
67+
68+
def test_append_diag():
69+
# Routine for appending diagonal elements
70+
assert_array_equal(append_diag(np.diag([2,3,1]), [1]),
71+
np.diag([2,3,1,1]))
72+
assert_array_equal(append_diag(np.diag([2,3,1]), [1,1]),
73+
np.diag([2,3,1,1,1]))
74+
aff = np.array([[2,0,0],
75+
[0,3,0],
76+
[0,0,1],
77+
[0,0,1]])
78+
assert_array_equal(append_diag(aff, [5], [9]),
79+
[[2,0,0,0],
80+
[0,3,0,0],
81+
[0,0,0,1],
82+
[0,0,5,9],
83+
[0,0,0,1]])
84+
assert_array_equal(append_diag(aff, [5,6], [9,10]),
85+
[[2,0,0,0,0],
86+
[0,3,0,0,0],
87+
[0,0,0,0,1],
88+
[0,0,5,0,9],
89+
[0,0,0,6,10],
90+
[0,0,0,0,1]])
91+
aff = np.array([[2,0,0,0],
92+
[0,3,0,0],
93+
[0,0,0,1]])
94+
assert_array_equal(append_diag(aff, [5], [9]),
95+
[[2,0,0,0,0],
96+
[0,3,0,0,0],
97+
[0,0,0,5,9],
98+
[0,0,0,0,1]])
99+
# Length of starts has to match length of steps
100+
assert_raises(ValueError, append_diag, aff, [5,6], [9])

nibabel/trackvis.py

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from .volumeutils import (native_code, swapped_code, endian_codes,
1212
allopen, rec2dict)
1313
from .orientations import aff2axcodes
14+
from .affines import apply_affine
1415

1516
# Definition of trackvis header structure.
1617
# See http://www.trackvis.org/docs/?subsect=fileformat
@@ -170,8 +171,6 @@ def read(fileobj, as_generator=False, points_space=None):
170171
affine = hdr['vox_to_ras']
171172
tv2vx = np.diag((1. / zooms).tolist() + [1])
172173
tv2mm = np.dot(affine, tv2vx).astype('f4')
173-
rzs = tv2mm[:3,:3].T
174-
trans = tv2mm[:3,3][None, :]
175174
n_s = hdr['n_scalars']
176175
n_p = hdr['n_properties']
177176
f4dt = np.dtype(endianness + 'f4')
@@ -211,7 +210,7 @@ def track_gen():
211210
if points_space == 'voxel':
212211
xyz = xyz / zooms
213212
elif points_space == 'rasmm':
214-
xyz = np.dot(xyz, rzs) + trans
213+
xyz = apply_affine(tv2mm, pts)
215214
if n_s:
216215
scalars = pts[:,3:]
217216
yield (xyz, scalars, ps)
@@ -368,8 +367,6 @@ def write(fileobj, streamlines, hdr_mapping=None, endianness=None,
368367
vx2tv = np.diag(zooms.tolist() + [1])
369368
mm2vx = npl.inv(affine)
370369
mm2tv = np.dot(vx2tv, mm2vx).astype('f4')
371-
rzs = mm2tv[:3,:3].T
372-
trans = mm2tv[:3,3][None, :]
373370
# write header
374371
fileobj = allopen(fileobj, mode='wb')
375372
fileobj.write(hdr.tostring())
@@ -385,7 +382,7 @@ def write(fileobj, streamlines, hdr_mapping=None, endianness=None,
385382
if points_space == 'voxel':
386383
pts = pts * zooms
387384
elif points_space == 'rasmm':
388-
pts = np.dot(pts, rzs) + trans
385+
pts = apply_affine(mm2tv, pts)
389386
# This call ensures that the data are 32-bit floats, and that
390387
# the endianness is OK.
391388
if pts.dtype != f4dt:
@@ -701,7 +698,7 @@ def aff_to_hdr(affine, trk_hdr, pos_vox=None, set_order=None):
701698
set_order = False
702699
try:
703700
version = trk_hdr['version']
704-
except KeyError, ValueError: # dict or structured array
701+
except (KeyError, ValueError): # dict or structured array
705702
version = 2
706703
if version == 2:
707704
trk_hdr['vox_to_ras'] = affine

0 commit comments

Comments
 (0)