diff --git a/changelog.d/transform-vectors-normals.md b/changelog.d/transform-vectors-normals.md new file mode 100644 index 000000000..1292d918a --- /dev/null +++ b/changelog.d/transform-vectors-normals.md @@ -0,0 +1,13 @@ +### Added + +- `transform(..., vectors=, normals=)`: fields that turn with the map, its linear + part applied to a vector (a velocity, a displacement) and its inverse transpose + to a normal, made unit length again so it stays normal to the moved surface. + +### Fixed + +- The VTK algorithm `Transform` carried the input's vectors and normals as they + were, so after a rotation they pointed the old way (and normals lit a moved + surface wrongly). It now transforms the input's active vectors and normals, on + points and cells, as VTK's `vtkTransformFilter` does, unless `vectors` or + `normals` say otherwise. Found running VTK's Python tests on Tack. diff --git a/packages/tack-vis/src/tack/dataset/algorithms/transform.py b/packages/tack-vis/src/tack/dataset/algorithms/transform.py index 8d1969774..4e24a4587 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/transform.py +++ b/packages/tack-vis/src/tack/dataset/algorithms/transform.py @@ -23,8 +23,42 @@ def _affine(positions, m00, m01, m02, m03, m10, m11, m12, m13, m20, m21, m22, m2 m20 * p[0] + m21 * p[1] + m22 * p[2] + m23] +@tack.kernel +def _directions(values, a00, a01, a02, a10, a11, a12, a20, a21, a22, unit, out): + # Each 3-vector under a 3x3 map; with ``unit`` made unit length again (a zero + # one stays zero). + for i in range(out.shape[0]): + v = values[i] + w = [a00 * v[0] + a01 * v[1] + a02 * v[2], + a10 * v[0] + a11 * v[1] + a12 * v[2], + a20 * v[0] + a21 * v[1] + a22 * v[2]] + if unit != 0: + length = sqrt(w[0] * w[0] + w[1] * w[1] + w[2] * w[2]) + if length > 0.0: + w = w / length + out[i] = w + + +def _mapped(data, key, linear, unit): + """Field ``key`` of ``data`` with each 3-vector under ``linear`` (3x3), made unit + length again if ``unit``: a new field in the same space and unit.""" + from tack.arrays import dtype_of, materialize, width_of + from tack.dataset.dataset import Field + + field = data.fields[key] + values = materialize(field.values) + if width_of(values) != 3: + raise TypeError(f"transform maps 3-vector fields; {key!r} is not one") + n = size_of(values) + kind = dtype_of(values) + out = _vectors(n, getattr(kind, "scalar_dtype", kind)) + if n: + _directions(values, *map(float, np.asarray(linear).reshape(-1)), int(unit), out) + return Field(field.space, out, field.unit) + + @on_inputs_device -def transform(data, matrix, fields=None): +def transform(data, matrix, fields=None, vectors=(), normals=()): """The geometry under an affine map: a 4x4 matrix (its last row 0 0 0 1) or 3x4, applied to each position as a column (Viskores' PointTransform, VTK's vtkTransformFilter). On a dataset with an origin the map moves the origin, @@ -35,7 +69,15 @@ def transform(data, matrix, fields=None): (``UniformCoordinates``; the direction may be any invertible matrix, a shear too), in double precision, and no position is stored -- as VTK's image data with a direction matrix holds it. A singular map (a projection) gives the - positions themselves.""" + positions themselves. + + Fields are carried as they are, but ``vectors`` and ``normals`` (field names, + or one): a vector -- a velocity, a displacement -- turns with the map, its + linear part applied to it, and a normal takes the inverse transpose of that + part, made unit length again, so it stays normal to the moved surface. VTK's + vtkTransformFilter does so to the active vectors and normals, which the VTK + algorithm ``Transform`` takes by default; Tack keeps no roles, so here they + are named. Each is a 3-vector field; normals need an invertible map.""" m = np.asarray(matrix, float) if m.shape == (4, 4): if not np.allclose(m[3], [0, 0, 0, 1]): @@ -43,6 +85,26 @@ def transform(data, matrix, fields=None): m = m[:3] if m.shape != (3, 4): raise ValueError(f"transform takes a 4x4 or 3x4 matrix, not {m.shape}") + out = _placed(data, m, fields) + from tack.dataset.dataset import is_field_key + + vectors = [vectors] if is_field_key(vectors) else list(vectors or ()) + normals = [normals] if is_field_key(normals) else list(normals or ()) + linear = m[:, :3] + for key in vectors: + out.fields[key] = _mapped(out, key, linear, unit=False) + if normals: + if np.linalg.det(linear) == 0.0: + raise ValueError("a singular map (a projection) has no normals to give") + inverse_transpose = np.linalg.inv(linear).T + for key in normals: + out.fields[key] = _mapped(out, key, inverse_transpose, unit=True) + return out + + +def _placed(data, m, fields): + """``data`` with its geometry under the affine map ``m`` (3x4), its fields + carried as they are.""" origin = None if any(data.origin): origin = m[:, :3] @ np.asarray(data.origin) + m[:, 3] diff --git a/packages/tack-vis/src/tack/interop/vtk/algorithms.py b/packages/tack-vis/src/tack/interop/vtk/algorithms.py index 6abe1e557..7cb69fb28 100644 --- a/packages/tack-vis/src/tack/interop/vtk/algorithms.py +++ b/packages/tack-vis/src/tack/interop/vtk/algorithms.py @@ -1520,18 +1520,52 @@ def _matrix(value): class Transform(TackAlgorithm): """``td.transform``: the geometry under an affine ``matrix`` (4x4 or 3x4, a - ``vtkMatrix4x4`` or a ``vtkTransform``); an image stays an image. VTK's - ``vtkTransformFilter``.""" + ``vtkMatrix4x4`` or a ``vtkTransform``); an image stays an image. The input's + active vectors turn with it and its active normals stay normal, on the points + and on the cells, as VTK's ``vtkTransformFilter`` transforms them, unless + ``vectors`` or ``normals`` name others (``[]``: none).""" OUTPUT = "structure" FUNCTION = "transform" + vectors = Parameter(None, "the 3-vector fields that turn with the map; by default the " + "input's active vectors") + normals = Parameter(None, "the 3-vector fields that stay normal to the moved surface; " + "by default the input's active normals") + + def _unit(self, vtk_data): + self._active = _active_vectors(vtk_data) + return super()._unit(vtk_data) + def arguments(self): values = super().arguments() values["matrix"] = _matrix(values["matrix"]) + active = getattr(self, "_active", {"vectors": [], "normals": []}) + for name in ("vectors", "normals"): + given = getattr(self, name) + values[name] = active[name] if given is None else [ + _field_argument(v) for v in ([given] if isinstance(given, (str, tuple)) + else given)] return values +def _active_vectors(vtk_data): + """The active 3-component vectors and normals of a VTK dataset (a partitioned one's + first partition), on its points and cells, as Tack's field keys.""" + block = vtk_data + if isinstance(vtk_data, vtkPartitionedDataSet): + block = next(iter(vtk_data), None) + out = {"vectors": [], "normals": []} + if block is None: + return out + for attributes, where in ((block.GetPointData(), "points"), (block.GetCellData(), "cells")): + for role, array in (("vectors", attributes.GetVectors()), + ("normals", attributes.GetNormals())): + if array is not None and array.GetName() and array.GetNumberOfComponents() == 3: + out[role].append(_field_argument((array.GetName(), where))) + return out + + class Warp(TackAlgorithm): """``td.warp``: the points moved along ``direction`` (a 3-vector field, or one vector) times ``scale`` (and the ``scale_by`` field). VTK's ``vtkWarpVector`` and diff --git a/packages/tack-vis/tests/test_transforms.py b/packages/tack-vis/tests/test_transforms.py index 400f28663..5031e4cc9 100644 --- a/packages/tack-vis/tests/test_transforms.py +++ b/packages/tack-vis/tests/test_transforms.py @@ -140,6 +140,65 @@ def test_transform_is_vtks(backend): td.transform(vtk_to_dataset(grid), np.ones((4, 4))) +@pytest.mark.parametrize("shear", [False, True]) +def test_vectors_turn_and_normals_stay_normal(backend, shear): + """Named vectors take the map's linear part, normals its inverse transpose made + unit length: a surface's normal is still normal to it after a shear, where the + linear part would tilt it.""" + m = _affine(4, shear=shear) + linear = m[:3, :3] + grid = td.uniform_grid((3, 3, 2)) + n = grid.num_points + rng = np.random.default_rng(5) + v = rng.normal(size=(n, 3)).astype(np.float32) + normal = rng.normal(size=(n, 3)) + normal /= np.linalg.norm(normal, axis=1)[:, None] + grid.fields["v"] = td.Field(td.H1(grid), tack.array(v, tack.vec3()), "m/s") + grid.fields["n"] = td.Field(td.H1(grid), tack.array(normal.astype(np.float32), + tack.vec3())) + grid.fields["kept"] = td.Field(td.H1(grid), tack.array(v, tack.vec3())) + out = td.transform(grid, m, vectors="v", normals=["n"]) + np.testing.assert_allclose(out.fields["v"].values.numpy(), v @ linear.T, rtol=1e-5, + atol=1e-5) + assert out.fields["v"].unit == "m/s" + expected = normal @ np.linalg.inv(linear) + expected /= np.linalg.norm(expected, axis=1)[:, None] + np.testing.assert_allclose(out.fields["n"].values.numpy(), expected, atol=1e-5) + # Normal still: perpendicular to every direction the map moves the plane's to. + tangents = np.cross(normal, [1.0, 0.0, 0.0]) + np.testing.assert_allclose(np.einsum("ij,ij->i", out.fields["n"].values.numpy(), + tangents @ linear.T), 0.0, atol=1e-5) + np.testing.assert_array_equal(out.fields["kept"].values.numpy(), v) + with pytest.raises(TypeError, match="3-vector"): + grid.fields["s"] = td.Field(td.H1(grid), tack.array(v[:, 0].copy())) + td.transform(grid, m, vectors="s") + with pytest.raises(ValueError, match="singular"): + td.transform(grid, np.diag([1.0, 1.0, 0.0, 1.0]), normals="n") + + +@needs_vtk +def test_the_vtk_algorithm_transforms_the_active_vectors_and_normals_as_vtk(): + from tack.interop.vtk import algorithms as ta + + sphere = vtk.vtkSphereSource()() + n = sphere.number_of_points + sphere.point_data["v"] = np.random.default_rng(6).normal(size=(n, 3)) + sphere.GetPointData().SetActiveVectors("v") + t = vtk.vtkTransform() + t.RotateWXYZ(33.0, 0.2, 1.0, 0.4) + t.Scale(1.5, 0.7, 2.0) + theirs = vtk.vtkTransformFilter(transform=t)(sphere) + ours = ta.Transform(matrix=t)(sphere) + for name in ("v", "Normals"): + np.testing.assert_allclose(np.asarray(ours.point_data[name]), + np.asarray(theirs.point_data[name]), atol=1e-5) + np.testing.assert_allclose(np.asarray(ours.points.data), np.asarray(theirs.points.data), + atol=1e-5) + untouched = ta.Transform(matrix=t, vectors=[], normals=[])(sphere) + np.testing.assert_array_equal(np.asarray(untouched.point_data["Normals"]), + np.asarray(sphere.point_data["Normals"])) + + @tack.func def _wave(x): return tack.sin(x[0]) * x[1] + 0.5 * x[2] * x[2]