From fd738bed3fce5ff0860a1c6c2446f56a78e0678d Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 14 May 2026 09:03:50 +1000 Subject: [PATCH] Make evaluate() work on tensor expressions with derivatives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _project_to_work_variable's matrix-expression branch was dead code with two bugs marked by a TODO: 1. MeshVariable(num_components=rows*cols) (flat int) couldn't infer vtype, raising ValueError on every call. 2. SNES_Projection(scalar_component=c) — that kwarg doesn't exist on SNES_Projection.__init__; the loop would have raised TypeError if it ever reached. So evaluate(strain_tensor, pts) — or any tensor expression containing derivatives — bounced out via the ValueError, and callers like SemiLagrangian.update_pre_solve hit their outer ``except Exception`` and fell back to their own projection solver. After PR #183 the DDt fallback works correctly, but the wasted try-fail-fallback was running every NS step and the spurious ValueError polluted stack traces. Replace the dead block with a real multi-component projection: - Cache a flat (1, Nc) MATRIX work var as the projector target (mirrors the existing pattern in DDt._setup_projections for SYM_TENSOR/TENSOR). - Use SNES_MultiComponent_Projection (one solve, shared DM) instead of a per-component cycling loop. - Cache a tensor-shaped (rows, cols) work var so the caller's ``work_var.sym`` preserves the input expression's shape. - After solve, fan flat var values into the tensor work var. Adds tests/test_0506_tensor_evaluate.py covering: - evaluate(strain_tensor, pts) returns (n_pts, dim, dim) shape - A known shear field produces the expected non-zero strain entries - A NavierStokes solve completes without entering the DDt fallback Related: issue #180, PR #183 (which fixes the fallback path so it works correctly when other expressions still hit it). Underworld development team with AI support from Claude Code --- src/underworld3/function/_function.pyx | 83 ++++++++++++------- tests/test_0506_tensor_evaluate.py | 108 +++++++++++++++++++++++++ 2 files changed, 164 insertions(+), 27 deletions(-) create mode 100644 tests/test_0506_tensor_evaluate.py diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 9dd71e15d..c17126cd0 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -549,44 +549,73 @@ def _project_to_work_variable(expr, mesh, smoothing=1e-6): Work variable containing projected nodal values """ import underworld3 as uw + import sympy - # Handle matrix expressions - need multi-component work variable - # TODO(BUG): This fails for dim×dim matrices (e.g. 2×2 stress tensor) - # because MeshVariable can't infer vtype from num_components=4 (flat int). - # Needs: pass num_components=(rows,cols) with vtype=TENSOR, and use - # Tensor_Projection instead of per-component scalar Projection. - # Currently, callers like SemiLagrangian.update_pre_solve fall back to - # their own projection solver via except-clause when this fails. + # Handle matrix/tensor expressions — need a multi-component projection. + # Implementation: project the (rows, cols) expression into a flat + # (1, Nc) MATRIX work variable using SNES_MultiComponent_Projection + # (one solve, shared DM), then fan the result out into a (rows, cols) + # shaped work variable so the caller's ``work_var.sym`` preserves the + # original tensor shape. if hasattr(expr, 'shape') and expr.shape != (1, 1): rows, cols = expr.shape n_components = rows * cols - # Get or create multi-component work variable - cache_key = f'_eval_work_{n_components}' + shape_key = f'{rows}x{cols}' + cache_key = f'_eval_work_{shape_key}' + flat_key = f'{cache_key}_flat' + projector_key = f'{cache_key}_projector' + if not hasattr(mesh, cache_key): + # Tensor-shaped result variable that the caller sees via .sym + if rows == mesh.dim and cols == mesh.dim: + vtype_out = uw.VarType.TENSOR + else: + vtype_out = uw.VarType.MATRIX work_var = uw.discretisation.MeshVariable( - cache_key, mesh, num_components=n_components, degree=1 + cache_key, + mesh, + (rows, cols), + vtype=vtype_out, + degree=1, ) setattr(mesh, cache_key, work_var) - # Create projectors for each component - projectors = [] - for c in range(n_components): - proj = uw.systems.Projection(mesh, work_var, scalar_component=c) - proj.petsc_options["snes_rtol"] = 1e-6 - projectors.append(proj) - setattr(mesh, f'{cache_key}_projectors', projectors) + + # Flat (1, Nc) target for the multi-component projector + flat_var = uw.discretisation.MeshVariable( + flat_key, + mesh, + (1, n_components), + vtype=uw.VarType.MATRIX, + degree=1, + ) + setattr(mesh, flat_key, flat_var) + + projector = uw.systems.MultiComponent_Projection( + mesh, + u_Field=flat_var, + n_components=n_components, + degree=1, + ) + projector.petsc_options["snes_rtol"] = 1e-6 + setattr(mesh, projector_key, projector) work_var = getattr(mesh, cache_key) - projectors = getattr(mesh, f'{cache_key}_projectors') - - # Project each component - idx = 0 - for i in range(rows): - for j in range(cols): - projectors[idx].uw_function = expr[i, j] - projectors[idx].smoothing = smoothing - projectors[idx].solve(zero_init_guess=False) - idx += 1 + flat_var = getattr(mesh, flat_key) + projector = getattr(mesh, projector_key) + + # Build flat row-matrix source: [[expr[0,0], expr[0,1], ..., expr[r-1,c-1]]] + flat_source = sympy.Matrix( + [[expr[i, j] for i in range(rows) for j in range(cols)]] + ) + projector.uw_function = flat_source + projector.smoothing = smoothing + projector.solve(zero_init_guess=False) + + # Fan flat result back to the tensor work variable + for idx in range(n_components): + i, j = divmod(idx, cols) + work_var.array[:, i, j] = flat_var.array[:, 0, idx] return work_var diff --git a/tests/test_0506_tensor_evaluate.py b/tests/test_0506_tensor_evaluate.py new file mode 100644 index 000000000..b297e0192 --- /dev/null +++ b/tests/test_0506_tensor_evaluate.py @@ -0,0 +1,108 @@ +""" +Regression test for ``uw.function.evaluate`` on tensor-valued expressions +containing derivatives. + +Background +---------- +``_project_to_work_variable`` in ``_function.pyx`` is reached from +``evaluate_nd`` whenever the input expression contains derivatives that +can't be evaluated directly at points (e.g. ``grad(u)``, ``strain_tensor``, +viscous flux). For non-scalar expressions it allocates a multi-component +work mesh variable and projects each component into it. + +Until this PR the matrix branch was dead code with a ``TODO(BUG)`` marker +— it called ``MeshVariable(..., num_components=rows*cols)`` (flat int) +which can't infer ``vtype``, raising ``ValueError`` on every invocation. +That bubbled up to ``DDt.update_pre_solve``'s outer ``except Exception``, +forcing a wasteful try-fail-fallback on every NavierStokes step. + +The fix uses ``SNES_MultiComponent_Projection`` against a flat +``(1, Nc) MATRIX`` work var, then fans the flat result into a tensor-shaped +work var so the caller's ``work_var.sym`` keeps the original shape. + +See: issue #180, PR comment by @ss2098 pinpointing the SLCN projection +path. +""" +import numpy as np +import pytest +import sympy +import underworld3 as uw + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_evaluate_strain_tensor_returns_correct_shape(): + """``evaluate(strain_tensor, pts)`` returns shape (n_pts, dim, dim).""" + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.2 + ) + v = uw.discretisation.MeshVariable("U", mesh, mesh.dim, degree=2) + + # Pure rotation: strain tensor is identically zero (antisymmetric gradient). + v.array[:, 0, 0] = v.coords[:, 1] + v.array[:, 0, 1] = -v.coords[:, 0] + + strain = mesh.vector.strain_tensor(v.sym) + assert strain.shape == (2, 2) + + pts = np.array([[0.5, 0.5], [0.3, 0.7]]) + result = np.asarray(uw.function.evaluate(strain, pts)) + + assert result.shape == (2, 2, 2), f"expected (n_pts, 2, 2), got {result.shape}" + # Strain rate of pure rotation is zero up to projection noise. + assert np.allclose(result, 0.0, atol=1e-6) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_evaluate_shear_field_strain_nonzero(): + """A simple shear field produces a known non-zero off-diagonal strain.""" + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.2 + ) + v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) + + # Simple shear u = (y, 0) — strain has off-diagonal entries = 0.5. + v.array[:, 0, 0] = v.coords[:, 1] + v.array[:, 0, 1] = 0.0 + + strain = mesh.vector.strain_tensor(v.sym) + result = np.asarray(uw.function.evaluate(strain, np.array([[0.5, 0.5]]))) + # Off-diagonal entries should be ~0.5 (projection has some smoothing). + assert abs(result[0, 0, 1] - 0.5) < 0.05 + assert abs(result[0, 1, 0] - 0.5) < 0.05 + + +@pytest.mark.level_2 +@pytest.mark.tier_a +def test_navier_stokes_solve_does_not_trigger_ddt_fallback(): + """A NavierStokes solve completes without entering the DDt projection fallback. + + The DDt fallback in ``update_pre_solve`` exists for expressions where + ``uw.function.evaluate`` raises. Once tensor evaluate works, the + fallback should not be needed for the NS viscous flux. + """ + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1 / 8, regular=False + ) + + v = uw.discretisation.MeshVariable("u", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + + ns = uw.systems.NavierStokes( + mesh, velocityField=v, pressureField=p, rho=1.0, order=2 + ) + ns.constitutive_model = uw.constitutive_models.ViscousFlowModel + ns.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + ns.bodyforce = sympy.Matrix([0, 0]).T + ns.add_dirichlet_bc((1.0, 0.0), "Top") + ns.add_dirichlet_bc((0.0, 0.0), "Bottom") + ns.add_dirichlet_bc((0.0, 0.0), "Left") + ns.add_dirichlet_bc((0.0, 0.0), "Right") + + # Pre-fix: this raised ShapeError("(1, 3) + (2, 2)") because + # _project_to_work_variable failed (ValueError) and the DDt fallback + # used the wrong source shape. Post-fix: evaluate succeeds, fallback + # isn't reached, solve completes. + ns.solve(timestep=0.01, zero_init_guess=True) + ns.solve(timestep=0.01, zero_init_guess=False)