diff --git a/meshmode/mesh/tools.py b/meshmode/mesh/tools.py index 6607d80807f07a6e1c1fa749c0080473ca06437e..df1c353cb5365d79367d131f45961a3718c94b81 100644 --- a/meshmode/mesh/tools.py +++ b/meshmode/mesh/tools.py @@ -23,6 +23,7 @@ THE SOFTWARE. """ import numpy as np +import numpy.linalg as la from pytools.spatial_btree import SpatialBinaryTreeBucket from six.moves import range @@ -141,4 +142,37 @@ def rand_rotation_matrix(ambient_dim, deflection=1.0, randnums=None): # }}} + +# {{{ AffineMap + +class AffineMap(object): + """An affine map ``A@x+b``represented by a matrix *A* and an offset vector *b*. + + .. attribute:: matrix + + A :class:`numpy.ndarray` representing the matrix *A*. + + .. attribute:: offset + + A :class:`numpy.ndarray` representing the vector *b*. + + .. autofunction:: inverted + .. autofunction:: __call__ + """ + + def __init__(self, matrix, offset): + self.matrix = matrix + self.offset = offset + + def inverted(self): + return AffineMap(la.inv(self.matrix), -la.solve(self.matrix, self.offset)) + + def __call__(self, vecs): + """Apply the affine map to an array *vecs* whose first axis + length matches ``matrix.shape[1]``. + """ + return (np.dot(self.matrix, vecs).T + self.offset).T + +# }}} + # vim: foldmethod=marker diff --git a/test/test_meshmode.py b/test/test_meshmode.py index 0ec9d3f9ca2a75ff8b7d7c97f1a1f1b31d66e259..915b89e8c0a569c064ba41c801743d223fc5bce3 100644 --- a/test/test_meshmode.py +++ b/test/test_meshmode.py @@ -1121,6 +1121,24 @@ def test_mesh_to_tikz(): # }}} +def test_affine_map(): + from meshmode.mesh.tools import AffineMap + for d in range(1, 5): + for i in range(100): + a = np.random.randn(d, d)+10*np.eye(d) + b = np.random.randn(d) + + m = AffineMap(a, b) + + assert la.norm(m.inverted().matrix - la.inv(a)) < 1e-10*la.norm(a) + + x = np.random.randn(d) + + m_inv = m.inverted() + + assert la.norm(x-m_inv(m(x))) < 1e-10 + + if __name__ == "__main__": import sys if len(sys.argv) > 1: