Fields given their gradient; world derivatives from the field's view

changelog.d/given-gradients.md

@@ -0,0 +1,14 @@
+### Added
+
+- `Field(space, values, unit, gradient=)`: a field's world gradient where
+  something other than its basis gives it, such as a solver's own
+  reconstruction, as a field on the same topology (3 values each for a scalar,
+  3k row by row for a vector of k). `world_gradient`, `world_jet` and
+  `gradients` read it, so `flow_quantities` and every kernel built on them use
+  the solver's derivatives; the Hessian it does not determine is NaN.
+
+### Changed
+
+- `world_gradient` and `world_jet` ask the field's view, whose default
+  (on every field view whose basis gives the derivatives) differentiates its basis as
+  before. A space that computes its values can give its own.

packages/tack-vis/src/tack/dataset/algorithms/_common.py

@@ -160,9 +160,9 @@ def _with_basis(data, field):
 
     space = field.space
     if isinstance(space, Values) and space.on == "points":
-        return Field(H1(data), field.values, field.unit)
+        return Field(H1(data), field.values, field.unit, gradient=field.gradient)
     if isinstance(space, Values) and space.on == "cells":
-        return Field(Constant(data), field.values, field.unit)
+        return Field(Constant(data), field.values, field.unit, gradient=field.gradient)
     return field
 
 

packages/tack-vis/src/tack/dataset/algorithms/gradients.py

@@ -20,20 +20,33 @@
 
 @tack.func
 def world_gradient(cells, u, c, pc):
-    """`u`'s gradient in the world at parametric `pc` of cell `c`: its
-    parametric gradient through the geometry's Jacobian. For kernels that need a
-    field's gradient where `gradients` does not put it."""
-    # d(u)/d(x) from d(u)/d(pc), through the geometry's Jacobian J. A solid's is
-    # J^-T d(u)/d(pc). A 2D cell's lies in its tangent plane, spanned by J's
-    # columns a and b: g = alpha a + beta b with g . a and g . b the parametric
-    # derivatives, VTK's and Viskores' in-plane gradient (no area: NaN, as a
-    # solid of no volume). Lines and vertices: zero.
-    g = tack.vec3(0.0, 0.0, 0.0)
-    if cells.DIMENSION == 3:
-        g = cells.geometry_jacobian(c, pc).inverse().transpose() @ u.parametric_gradient(c, pc)
-    elif cells.DIMENSION == 2:
-        g = _in_plane(cells.geometry_jacobian(c, pc), u.parametric_gradient(c, pc))
-    return g
+    """`u`'s gradient in the world at parametric `pc` of cell `c`, as its view gives it:
+    from its basis by default (`_BasisGradient`), the gradient it was given
+    (`Field(..., gradient=)`), or a space's own. For kernels that need a field's
+    gradient where `gradients` does not put it."""
+    return u.world_gradient(cells, c, pc)
+
+
+@tack.func
+def world_hessian(cells, u, c, pc):
+    """`u`'s second derivatives in the world at parametric `pc` of solid cell `c`, a
+    symmetric 3x3 matrix, as its view gives them (`world_jet`). From the basis: by the
+    chain rule through the geometry, the parametric Hessian less the geometry's
+    curvature along the world gradient, taken to world axes on both sides by the
+    inverse Jacobian; analytic on every shape at orders 1 and 2, and at any order on
+    tensor shapes, simplices and wedges; zero on surfaces, lines and vertices. NaN for
+    a field given only its gradient."""
+    _, _, h = world_jet(cells, u, c, pc)
+    return h
+
+
+@tack.func
+def world_jet(cells, u, c, pc):
+    """Scalar `u`'s value, world gradient and world Hessian at parametric `pc` of cell
+    `c`, as its view gives them, from one pass over the cell's values where it can:
+    what a kernel needing several of them reads once. The gradient is
+    `world_gradient`'s, and the Hessian `world_hessian`'s."""
+    return u.world_jet(cells, c, pc)
 
 
 @tack.func
@@ -49,36 +62,57 @@ def _in_plane(m, d):
     return ((bb * d[0] - ab * d[1]) * a + (aa * d[1] - ab * d[0]) * b) / det
 
 
-
-@tack.func
-def world_hessian(cells, u, c, pc):
-    """`u`'s second derivatives in the world at parametric `pc` of solid cell `c`, a
-    symmetric 3x3 matrix: by the chain rule through the geometry, the parametric
-    Hessian less the geometry's curvature along the world gradient, taken to world
-    axes on both sides by the inverse Jacobian. Analytic from the basis on every
-    shape at orders 1 and 2, and at any order on tensor shapes, simplices and
-    wedges. Zero on surfaces, lines and vertices."""
-    _, _, h = world_jet(cells, u, c, pc)
-    return h
-
-
-@tack.func
-def world_jet(cells, u, c, pc):
-    """Scalar `u`'s value, world gradient and world Hessian at parametric `pc` of cell
-    `c`, from the field view's `parametric_jet`, one pass over the cell's values: what a
-    kernel needing several of them reads once. The gradient is `world_gradient`'s on
-    every cell, in the cell's plane on surfaces, and the Hessian `world_hessian`'s, zero
-    on surfaces, lines and vertices."""
-    v, pg, ph = u.parametric_jet(c, pc)
-    g = tack.vec3(0.0, 0.0, 0.0)
-    h = tack.mat3([0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0])
-    if cells.DIMENSION == 3:
-        inverse = cells.geometry_jacobian(c, pc).inverse()
-        g = inverse.transpose() @ pg
-        h = inverse.transpose() @ (ph - cells.geometry_curvature(c, pc, g)) @ inverse
-    elif cells.DIMENSION == 2:
-        g = _in_plane(cells.geometry_jacobian(c, pc), pg)
-    return v, g, h
+class _BasisGradient:
+    """A field view's world gradient from its basis: on every field view with a
+    ``parametric_gradient`` and no other source (``_derivative_mixins``)."""
+
+    @tack.func
+    def world_gradient(self, cells, c, pc):
+        # d(u)/d(x) from d(u)/d(pc), through the geometry's Jacobian J. A solid's is
+        # J^-T d(u)/d(pc). A 2D cell's lies in its tangent plane, spanned by J's
+        # columns a and b: g = alpha a + beta b with g . a and g . b the parametric
+        # derivatives, VTK's and Viskores' in-plane gradient (no area: NaN, as a
+        # solid of no volume). Lines and vertices: zero.
+        g = tack.vec3(0.0, 0.0, 0.0)
+        if cells.DIMENSION == 3:
+            g = cells.geometry_jacobian(c, pc).inverse().transpose() @ self.parametric_gradient(c, pc)
+        elif cells.DIMENSION == 2:
+            g = _in_plane(cells.geometry_jacobian(c, pc), self.parametric_gradient(c, pc))
+        return g
+
+
+class _BasisJet:
+    """A field view's value, world gradient and world Hessian from its basis, in one
+    pass: on every field view with a ``parametric_jet``."""
+
+    @tack.func
+    def world_jet(self, cells, c, pc):
+        v, pg, ph = self.parametric_jet(c, pc)
+        g = tack.vec3(0.0, 0.0, 0.0)
+        h = tack.mat3([0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0])
+        if cells.DIMENSION == 3:
+            inverse = cells.geometry_jacobian(c, pc).inverse()
+            g = inverse.transpose() @ pg
+            h = inverse.transpose() @ (ph - cells.geometry_curvature(c, pc, g)) @ inverse
+        elif cells.DIMENSION == 2:
+            g = _in_plane(cells.geometry_jacobian(c, pc), pg)
+        return v, g, h
+
+
+class _GivenGradient:
+    """A field given its gradient (``Field(..., gradient=)``): ``world_gradient`` is that
+    field's value at ``pc`` of cell ``c``, and ``world_jet`` pairs it with the field's
+    value and a Hessian of NaN, which a gradient alone does not determine."""
+
+    @tack.func
+    def world_gradient(self, cells, c, pc):
+        return self.given_gradient.value(c, pc)
+
+    @tack.func
+    def world_jet(self, cells, c, pc):
+        n = self.no_hessian
+        return (self.value(c, pc), self.given_gradient.value(c, pc),
+                tack.mat3([n, n, n], [n, n, n], [n, n, n]))
 
 
 @tack.kernel
@@ -278,6 +312,24 @@ def _polyhedral_scalar_gradients(data, field, at):
     return out
 
 
+@tack.kernel
+def _row(values, k, out):
+    for i in range(out.shape[0]):
+        v = values[i]
+        out[i] = tack.vec3(v[3 * k], v[3 * k + 1], v[3 * k + 2])
+
+
+def _gradient_row(gradient, k):
+    """Row ``k`` of a vector field's given gradient: its component ``k``'s."""
+    if gradient is None:
+        return None
+    values = materialize(gradient.values)
+    out = tack.empty((size_of(values),), tack.vec(3, arrays.dtype_of(values)))
+    if out.shape[0]:
+        _row(values, k, out)
+    return Field(gradient.space, out, gradient.unit)
+
+
 @tack.kernel
 def _component(values, k, out):
     for i in range(out.shape[0]):
@@ -355,6 +407,9 @@ def gradients(data, field, at="cells"):
     Viskores adds nothing but still counts the cell). ``flow_quantities`` derives divergence, vorticity and
     the Q-criterion from a vector field's gradient.
 
+    A field given its gradient (``Field(..., gradient=)``) gives that instead of its
+    basis', evaluated at the same points and averaged the same way.
+
     On a polyhedral topology there is no reference element. A cell's gradient
     of point data is then Green-Gauss over its boundary: the faces walked
     outward and fanned from their points' means (the triangles ``cell_geometry``
@@ -365,6 +420,9 @@ def gradients(data, field, at="cells"):
     get zero. At the points, the cells' gradients are averaged over the cells
     around each point."""
     field = _with_basis(data, _field(data, field))
+    if not getattr(data.topology, "reference_cells", True) and field.gradient is not None:
+        raise NotImplementedError("gradients of a field given its gradient on a mesh without "
+                                  "reference elements")
     if (not getattr(data.topology, "reference_cells", True)
             and not isinstance(field.space, (H1, Constant))):
         raise NotImplementedError(f"gradients of a field on {field.space!r} on a mesh "
@@ -394,7 +452,7 @@ def gradients(data, field, at="cells"):
         part = tack.empty((size_of(field.values),), arrays.dtype_of(field.values))
         if part.shape[0]:
             _component(values, k, part)
-        rows = _scalar_gradients(data, Field(field.space, part), at)
+        rows = _scalar_gradients(data, Field(field.space, part, gradient=_gradient_row(field.gradient, k)), at)
         if n:
             _place_rows(rows, k, out)
     return Field(space, out, unit)

packages/tack-vis/src/tack/dataset/dataset.py

@@ -54,9 +54,16 @@ class Field:
     ``unit`` is the values' physical unit, a string such as ``"m/s"`` kept in
     unyt's canonical spelling (``tack.dataset.units``), or ``None`` for unknown:
     meaning, which kernels never see. Filters that keep a quantity keep its
-    unit; anything else gives ``None`` rather than a guess."""
+    unit; anything else gives ``None`` rather than a guess.
 
-    def __init__(self, space, values, unit=None):
+    ``gradient``, a field on the same topology, is its world gradient where something
+    other than its basis gives it, such as a solver's own reconstruction: 3 values
+    each for a scalar, 3k row by row for a vector of k components (``gradients``'
+    layout). ``world_gradient``, ``world_jet`` and ``gradients`` read it, evaluated
+    at the point asked for as that field's view evaluates it; the Hessian, which a
+    gradient does not determine, is NaN."""
+
+    def __init__(self, space, values, unit=None, gradient=None):
         if not isinstance(space, Space):
             raise TypeError(f"a field's space is a Space on a topology, not {space!r}")
         views.storage(values)                       # refuses what is not an array
@@ -67,6 +74,16 @@ def __init__(self, space, values, unit=None):
         self.values = values
         self._unit = units.canonical(unit)
         _ = self._device                    # refuses values on another device
+        if gradient is not None:
+            if not isinstance(gradient, Field):
+                raise TypeError(f"a field's gradient is a tack.dataset.Field, not {gradient!r}")
+            if gradient.space.topology is not space.topology:
+                raise ValueError("a field's gradient belongs to another topology")
+            width = max(1, arrays.width_of(values) or 0)
+            if (arrays.width_of(gradient.values) or 1) != 3 * width:
+                raise ValueError(f"a field of {width} components has a gradient of {3 * width}, "
+                                 f"not {arrays.width_of(gradient.values) or 1}")
+        self.gradient = gradient
 
     @property
     def unit(self):
@@ -96,7 +113,8 @@ def __repr__(self):
         values = self.values
         stored = repr(values) if isinstance(values, arrays._Implicit) else "stored"
         unit = "" if self._unit is None else f", in {self._unit}"
-        return f"Field({self.space!r}, {arrays.size_of(values)} values, {stored}{unit})"
+        given = "" if self.gradient is None else ", with its gradient"
+        return f"Field({self.space!r}, {arrays.size_of(values)} values, {stored}{unit}{given})"
 
 
 def _is_cells(group):
@@ -133,6 +151,8 @@ def _space_parts(field, group, space_mixin=None, corners_only=False):
         # A geometry's incidence is the domain view's own.
         mixins.extend(space.extra_mixins(group))
     mixins.append(storage)
+    if not space_mixin:
+        mixins.extend(_derivative_mixins(field, group, mixins, attributes))
     # Reading only the points' values -- an order-2 geometry on faces and edges
     # -- needs none of the space's own layout.
     layout = {} if corners_only else space.attributes(group, dtype=_basis_type(field))
@@ -142,6 +162,26 @@ def _space_parts(field, group, space_mixin=None, corners_only=False):
     return mixins, attributes
 
 
+def _derivative_mixins(field, group, mixins, attributes):
+    """Where a field view's ``world_gradient`` and ``world_jet`` come from: the gradient
+    the field was given, a space's own mixin, or else the basis -- the parts of it the
+    view has, since every method of a view resolves when a kernel takes it."""
+    from tack.dataset.algorithms.gradients import _BasisGradient, _BasisJet, _GivenGradient
+
+    given = getattr(field, "gradient", None)
+    if given is not None:
+        attributes.update(given_gradient=given.view(group), no_hessian=float("nan"))
+        return [_GivenGradient]
+    out = []
+    if (any(hasattr(m, "parametric_gradient") for m in mixins)
+            and not any(hasattr(m, "world_gradient") for m in mixins)):
+        out.append(_BasisGradient)
+    if (any(hasattr(m, "parametric_jet") for m in mixins)
+            and not any(hasattr(m, "world_jet") for m in mixins)):
+        out.append(_BasisJet)
+    return out
+
+
 def _basis_type(field):
     """The float type ``field``'s basis is evaluated in: its own (``spaces._float_type``)."""
     from tack.dataset.spaces import _float_type
@@ -781,7 +821,8 @@ def moved(field):
             values = field.values.to_device(target)
             if isinstance(space, L2) and space.layout != field.space.layout:
                 values = _relaid(values, field.space, space)      # "auto" chose another
-            return Field(space, values, field.unit)
+            gradient = None if field.gradient is None else moved(field.gradient)
+            return Field(space, values, field.unit, gradient=gradient)
 
         sets = {name: Set._from_sorted_ids(topology, s.on, s.ids.to_device(target))
                 for name, s in self.sets.items()}

packages/tack-vis/tests/_vtk_algorithms_mpi.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) Kitware Inc.
+# SPDX-License-Identifier: BSD-3-Clause
+
 """``test_vtk_algorithms``' check across MPI ranks: ``mpiexec -n 3 python
 _vtk_algorithms_mpi.py comm|vtk``. Each rank holds one slab of VTK's wavelet, the
 slabs sharing their planes of points, in a ``vtkPartitionedDataSet`` of one