|
| 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 |
0 commit comments