transform turns vectors and normals; the VTK algorithm the active ones

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.

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,14 +69,42 @@ 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]):
             raise ValueError("a projective matrix (last row not 0 0 0 1) is not affine")
         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]

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

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]