diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py index 2b9895775b059ef0b389421465c9cb79fcd12eec..32016cd63afc3e2fd7379c3b3ba5b5f46ed05548 100644 --- a/meshmode/mesh/processing.py +++ b/meshmode/mesh/processing.py @@ -35,6 +35,7 @@ __doc__ = """ .. autofunction:: perform_flips .. autofunction:: find_bounding_box .. autofunction:: merge_disjoint_meshes +.. autofunction:: map_mesh .. autofunction:: affine_map """ @@ -271,34 +272,46 @@ def merge_disjoint_meshes(meshes, skip_tests=False): # }}} -# {{{ affine map +# {{{ map -def affine_map(mesh, A=None, b=None): # noqa - """Apply the affine map *f(x)=Ax+b* to the geometry of *mesh*.""" +def map_mesh(mesh, f): # noqa + """Apply the map *f* to the mesh. *f* needs to accept and return arrays of + shape ``(ambient_dim, npoints)``.""" - if A is None: - A = np.eye(mesh.ambient_dim) # noqa - - if b is None: - b = np.zeros(A.shape[0]) - - vertices = np.einsum("ij,jv->iv", A, mesh.vertices) + b[:, np.newaxis] + vertices = f(mesh.vertices) # {{{ assemble new groups list new_groups = [] for group in mesh.groups: - nodes = ( - np.einsum("ij,jen->ien", A, group.nodes) - + b[:, np.newaxis, np.newaxis]) - new_groups.append(group.copy(nodes=nodes)) + new_groups.append(group.copy(nodes=f(group.nodes))) # }}} from meshmode.mesh import Mesh return Mesh(vertices, new_groups, skip_tests=True, - nodal_adjacency=mesh.nodal_adjacency_init_arg()) + nodal_adjacency=mesh.nodal_adjacency_init_arg(), + facial_adjacency_groups=mesh._facial_adjacency_groups) + +# }}} + + +# {{{ affine map + +def affine_map(mesh, A=None, b=None): # noqa + """Apply the affine map *f(x)=Ax+b* to the geometry of *mesh*.""" + + if A is None: + A = np.eye(mesh.ambient_dim) # noqa + + if b is None: + b = np.zeros(A.shape[0]) + + def f(x): + return np.einsum("ij,jv->iv", A, x) + b[:, np.newaxis] + + return map_mesh(mesh, f) # }}}