From 49cf7c15ef0d9aac3e193c38ed46b7b39d4c0bcc Mon Sep 17 00:00:00 2001 From: Jeff Lee Date: Tue, 3 Nov 2026 19:54:25 +0000 Subject: [PATCH 1/3] Fields given their gradient; world derivatives from the field's view A solver's derivatives are its own: a finite-volume reconstruction, or a lattice Boltzmann solver's from its populations. Field(..., gradient=) takes that gradient as a field on the same topology (gradients' layout), and world_gradient and world_jet ask the field's view, which reads it, or a space's own derivatives, or else differentiates the basis as before. gradients splits a vector's given gradient into its rows, so flow_quantities and every kernel on world_gradient use the solver's derivatives. --- changelog.d/given-gradients.md | 14 ++ .../src/tack/dataset/algorithms/_common.py | 4 +- .../src/tack/dataset/algorithms/gradients.py | 148 ++++++++++++------ packages/tack-vis/src/tack/dataset/dataset.py | 49 +++++- .../tack-vis/tests/test_given_gradients.py | 114 ++++++++++++++ 5 files changed, 278 insertions(+), 51 deletions(-) create mode 100644 changelog.d/given-gradients.md create mode 100644 packages/tack-vis/tests/test_given_gradients.py diff --git a/changelog.d/given-gradients.md b/changelog.d/given-gradients.md new file mode 100644 index 000000000..9e3ba6b8c --- /dev/null +++ b/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. diff --git a/packages/tack-vis/src/tack/dataset/algorithms/_common.py b/packages/tack-vis/src/tack/dataset/algorithms/_common.py index 5af7990a6..612c2b058 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/_common.py +++ b/packages/tack-vis/src/tack/dataset/algorithms/_common.py @@ -157,9 +157,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 diff --git a/packages/tack-vis/src/tack/dataset/algorithms/gradients.py b/packages/tack-vis/src/tack/dataset/algorithms/gradients.py index 74efcd1b9..523ba137b 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/gradients.py +++ b/packages/tack-vis/src/tack/dataset/algorithms/gradients.py @@ -17,20 +17,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 @@ -46,36 +59,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 @@ -275,6 +309,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]): @@ -352,6 +404,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`` @@ -362,6 +417,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 " @@ -391,7 +449,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) diff --git a/packages/tack-vis/src/tack/dataset/dataset.py b/packages/tack-vis/src/tack/dataset/dataset.py index e005fb4b9..cf4bd65b9 100644 --- a/packages/tack-vis/src/tack/dataset/dataset.py +++ b/packages/tack-vis/src/tack/dataset/dataset.py @@ -51,9 +51,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 @@ -64,6 +71,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): @@ -93,7 +110,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): @@ -130,6 +148,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)) @@ -139,6 +159,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 @@ -778,7 +818,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()} diff --git a/packages/tack-vis/tests/test_given_gradients.py b/packages/tack-vis/tests/test_given_gradients.py new file mode 100644 index 000000000..b01df215c --- /dev/null +++ b/packages/tack-vis/tests/test_given_gradients.py @@ -0,0 +1,114 @@ +"""A field given its gradient: world_gradient, world_jet and gradients read it instead of +differentiating the basis.""" + +import numpy as np +import pytest + +import tack +import tack.dataset as td + + +@tack.func +def _u(x): + return x[0] * x[0] + 2.0 * x[1] - x[2] + + +@tack.func +def _v(x): + return tack.vec3(x[0] * x[1], x[2], -x[0]) + + +@tack.kernel +def _gradients_at(cells, u, out): + for c in cells: + out[cells.entity_id(c)] = td.world_gradient(cells, u, c, cells.parametric_center()) + + +@tack.kernel +def _jets_at(cells, u, values, grads, hess): + for c in cells: + v, g, h = td.world_jet(cells, u, c, cells.parametric_center()) + e = cells.entity_id(c) + values[e] = v + grads[e] = g + hess[e] = h[0, 0] + + +def _data(): + return td.uniform_grid((5, 4, 3), origin=(0.1, 0.2, 0.3), spacing=(0.3, 0.25, 0.4)) + + +def _vectors(rows): + rows = np.asarray(rows, np.float32) + out = tack.empty((len(rows),), tack.vec(rows.shape[1], tack.f32)) + tack.copyto(out, rows) + return out + + +def _per_cell(data, rows): + return td.Field(td.Constant(data), _vectors(rows)) + + +def test_world_gradient_and_gradients_read_a_given_gradient(backend): + data = _data() + given = np.random.default_rng(1).uniform(-1, 1, (data.num_cells, 3)).astype(np.float32) + u = td.interpolate(data, _u, td.H1(data)) + data.fields["u"] = td.Field(u.space, u.values, gradient=_per_cell(data, given)) + out = tack.zeros((data.num_cells,), tack.vec3(tack.f32)) + td.for_each(_gradients_at, data, "cells", data.fields["u"], out) + np.testing.assert_allclose(out.numpy(), given) + np.testing.assert_allclose(td.gradients(data, "u").values.numpy(), given) + plain = td.gradients(data, td.Field(u.space, u.values), at="points").values.numpy() + at_points = td.gradients(data, "u", at="points").values.numpy() + assert not np.allclose(plain, at_points) + + +def test_world_jet_pairs_the_given_gradient_with_the_value(backend): + data = _data() + given = np.tile(np.float32([1.0, 2.0, 3.0]), (data.num_cells, 1)) + u = td.interpolate(data, _u, td.H1(data)) + field = td.Field(u.space, u.values, gradient=_per_cell(data, given)) + values = tack.zeros((data.num_cells,), tack.f32) + grads = tack.zeros((data.num_cells,), tack.vec3(tack.f32)) + hess = tack.zeros((data.num_cells,), tack.f32) + td.for_each(_jets_at, data, "cells", field, values, grads, hess) + centers = td.cell_centers(data).values.numpy() + np.testing.assert_allclose(grads.numpy(), given) + np.testing.assert_allclose(values.numpy(), _u_numpy(centers), atol=0.05) + assert np.isnan(hess.numpy()).all() + + +def _u_numpy(x): + return x[:, 0] ** 2 + 2.0 * x[:, 1] - x[:, 2] + + +def test_a_vector_field_gives_its_rows(backend): + data = _data() + jacobian = np.random.default_rng(2).uniform(-1, 1, (data.num_cells, 9)).astype(np.float32) + v = td.interpolate(data, _v, td.H1(data), dtype=tack.vec3(tack.f32)) + data.fields["v"] = td.Field(v.space, v.values, + gradient=td.Field(td.Constant(data), _vectors(jacobian))) + np.testing.assert_allclose(td.gradients(data, "v").values.numpy(), jacobian) + q = td.flow_quantities(data, td.gradients(data, "v")) + J = jacobian.reshape(-1, 3, 3) + np.testing.assert_allclose(q[0].values.numpy(), np.trace(J, axis1=1, axis2=2), rtol=1e-5, atol=1e-6) + + +def test_a_given_gradient_is_checked(backend): + data = _data() + u = td.interpolate(data, _u, td.H1(data)) + with pytest.raises(ValueError): + td.Field(u.space, u.values, gradient=td.Field(td.Constant(data), tack.zeros((data.num_cells,), tack.f32))) + with pytest.raises(ValueError): + other = _data() + td.Field(u.space, u.values, gradient=_per_cell(other, np.zeros((other.num_cells, 3)))) + with pytest.raises(TypeError): + td.Field(u.space, u.values, gradient=np.zeros(3)) + + +def test_point_data_keeps_its_given_gradient(backend): + data = _data() + u = td.interpolate(data, _u, td.H1(data)) + given = np.tile(np.float32([4.0, 5.0, 6.0]), (data.num_cells, 1)) + data.fields["u"] = td.Field(td.Values(data, "points"), u.values, gradient=_per_cell(data, given)) + np.testing.assert_allclose(td.gradients(data, "u").values.numpy(), given) From a273ae699289812402244f0b55551d33460338c8 Mon Sep 17 00:00:00 2001 From: Jeff Lee Date: Wed, 4 Nov 2026 20:43:06 +0000 Subject: [PATCH 2/3] SPDX header for test_given_gradients.py --- packages/tack-vis/tests/test_given_gradients.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/tack-vis/tests/test_given_gradients.py b/packages/tack-vis/tests/test_given_gradients.py index b01df215c..55c8d2d7a 100644 --- a/packages/tack-vis/tests/test_given_gradients.py +++ b/packages/tack-vis/tests/test_given_gradients.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) Kitware Inc. +# SPDX-License-Identifier: BSD-3-Clause + """A field given its gradient: world_gradient, world_jet and gradients read it instead of differentiating the basis.""" From 97ae286f4d3b1c1b3e0ebf1bf9473aab12022d67 Mon Sep 17 00:00:00 2001 From: Jeff Lee Date: Wed, 4 Nov 2026 21:42:45 +0000 Subject: [PATCH 3/3] SPDX header for _vtk_algorithms_mpi.py --- packages/tack-vis/tests/_vtk_algorithms_mpi.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/tack-vis/tests/_vtk_algorithms_mpi.py b/packages/tack-vis/tests/_vtk_algorithms_mpi.py index 6029609d4..60c76f411 100644 --- a/packages/tack-vis/tests/_vtk_algorithms_mpi.py +++ b/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