[BUG] - Matrix Mismatch Regarding NavierStokesSLCN Solver

[Chawgai]

I am facing the following problem while using the attached test script. I run like mpirun -n 4 python nstt.py

nstt.txt

Screenshots
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/systems/solvers.py", line 2414, in <la
navier_stokes.solve(timestep=delta_t, zero_init_guess=False)
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/systems/solvers.py", line 3773, in solve
self.DFDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf)
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/systems/ddt.py", line 1771, in update_pre_solve
self._psi_star_projection_solver.solve(verbose=verbose)
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "src/underworld3/cython/petsc_generic_snes_solvers.pyx", line 3863, in underworld3.cython.generic_solvers.SNES_MultiComponent.solve
File "/Users/umbda>
lambda self: (self.u.sym - self.uw_function) * self.uw_weighting_function,
~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/core/decorators.py", line 118, in binary_op_wrapper
return func(self, other)
^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/matrices/matrixbase.py", line 3045, in sub
return self + (-a)
~~~~~^~~~~~
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/core/decorators.py", line 118, in binary_op_wrapper
return func(self, other)
^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/matrices/matrixbase.py", line 2765, in add
raise ShapeError(f"Matrix size mismatch: {self.shape} + {other.shape}.")
sympy.matrices.exceptions.ShapeError: Matrix size mismatch: (1, 3) + (2, 2).
8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "src/underworld3/cython/petsc_generic_snes_solvers.pyx", line 721, in underworld3.cython.generic_solvers.SolverBaseClass._build
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "src/underworld3/cython/petsc_generic_snes_solvers.pyx", line 3549, in underworld3.cython.generic_solvers.SNES_MultiComponent._setup_pointwise_functions
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/utilities/_api_tools.py", line 423, in get
updated_value = self.value_fn(obj)
^^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/systems/solvers.py", line 2414, in
lambda self: (self.u.sym - self.uw_function) * self.uw_weighting_function,
~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/core/decorators.py", line 118, in binary_op_wrapper
return func(self, other)
^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/matrices/matrixbase.py", line 3045, in sub
return self + (-a)
~~~~~^~~~~~
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/core/decorators.py", line 118, in binary_op_wrapper
return func(self, other)
^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/matrices/matrixbase.py", line 2765, in add
raise ShapeError(f"Matrix size mismatch: {self.shape} + {other.shape}.")
sympy.matrices.exceptions.ShapeError: Matrix size mismatch: (1, 3) + (2, 2).

nstt.txt

[ss2098]

I tested a diagnostic version of your script in the Underworld3 pixi environment. The first thing I would try is to simplify the solve setup and run it.

For example, from inside the underworld3 repository:

pixi run --manifest-path pixi.toml python nstt.py

Then, after the serial case works, try MPI:

pixi run --manifest-path pixi.toml mpirun -n 4 python nstt.py

or, depending on your local MPI setup:

mpirun -n 4 pixi run --manifest-path pixi.toml python nstt.py

In my diagnostic run, the basic symbolic shapes were consistent for a 2-D problem:

meshbox.dim = 2
v_soln.sym.shape = (1, 2)
p_soln.sym.shape = (1, 1)
bodyforce.shape = (1, 2)

So the issue does not look like a simple velocity/pressure/body-force dimension error. The traceback points instead to the NavierStokesSLCN semi-Lagrangian update/projection step.

I would make these changes first:

# Use a smaller mesh while debugging
resolution = 8

# Keep the solver setup minimal first
navier_stokes.petsc_options.setValue("snes_monitor", None)
navier_stokes.petsc_options.setValue("snes_converged_reason", None)

# Temporarily comment out custom field-split options
# navier_stokes.petsc_options["snes_type"] = "newtonls"
# navier_stokes.petsc_options["ksp_type"] = "fgmres"
# navier_stokes.petsc_options["pc_type"] = "fieldsplit"
# navier_stokes.petsc_options["pc_fieldsplit_type"] = "schur"
# navier_stokes.petsc_options["pc_fieldsplit_schur_fact_type"] = "full"
# navier_stokes.petsc_options["ksp_rtol"] = 1e-4

Then solve with an explicit timestep and a clean initial guess:

delta_x = meshbox.get_min_radius()
dt_advective = 0.5 * delta_x / vel
dt_diffusive = 0.5 * (delta_x**2) / kin_visc
delta_t = min(dt_advective, dt_diffusive)

navier_stokes.solve(
    timestep=delta_t,
    zero_init_guess=True,
)

I would not start with:

navier_stokes.solve(
    timestep=delta_t,
    zero_init_guess=False,
)

because the SLCN history/projection step appears to be where the shape mismatch is being generated.

If the reduced serial case works, then increase the resolution back to 16 and test MPI with:

pixi run --manifest-path pixi.toml mpirun -n 4 python nstt.py

If the reduced serial case still gives:

Matrix size mismatch: (1, 3) + (2, 2)

then I think this is likely an internal NavierStokesSLCN projection/history-field issue rather than a simple user-side shape error, because the visible 2-D fields have consistent shapes.

[ss2098]

One more point: if this example is intended to be a steady lid-driven cavity / Couette-style test, then NavierStokesSLCN may not be the simplest system to use, because it automatically enters the semi-Lagrangian time-derivative machinery. If you do not actually need the transient SLCN update, it may be better to use a steady Stokes/Navier-Stokes formulation instead. The error is coming from the SLCN update path, so avoiding that path is the cleanest solution for a steady benchmark.

[lmoresi]

@ss2098 - the steady state NS problem requires us to determine the characteristics of the flow and the easiest way, honestly, to do that is to solve the transient problem towards steady state (if that exists).

[ss2098]

@lmoresi Yes, that makes sense. I was not trying to avoid the transient formulation. I agree that, for a steady-state Navier-Stokes problem, evolving the transient system toward a steady solution is often the most practical and physically meaningful way to determine the flow, assuming such a steady state exists.

My concern was more about the startup/setup side of the computation. In this case, the error seems to arise very early in the NavierStokesSLCN semi-Lagrangian history/projection update, before we can really assess the flow behaviour. So I was trying to separate the mathematical/geodynamic issue from the numerical/API setup issue.

In other words, I agree that the transient solve is the right geodynamic strategy. My suggestion was not to replace it with a purely steady solve, but to make the first transient step as clean and diagnostic as possible: explicit timestep, clean initial guess, minimal PETSc options, and serial execution first. That way, if the model fails, we can tell whether the problem is with the physical formulation, the solver initialization, or the internal SLCN projection/history-field construction.

So perhaps the better phrasing of what I said could be : the transient approach is appropriate, but the initial transient setup needs to be reduced and checked carefully before adding MPI, custom PETSc options, or larger meshes.

[lmoresi]

Thanks @Chawgai and @ss2098 — diagnosed and fixed in #183. The root cause was an outdated API call that @ss2098's bisect to the SLCN projection/history-field path pointed at directly.

When SNES_MultiComponent_Projection was wired into the SemiLagrangian DDt earlier this year, the psi_fn setter was updated to flatten the source tensor to a (1, Nc) row matrix (via a new _build_projection_source helper), matching the multicomponent projector's shape contract. The fallback inside update_pre_solve — taken whenever uw.function.evaluate() raises, which it does for any expression containing derivatives — was missed during that migration and kept assigning the raw self.psi_fn directly.

Stokes never hits this fallback because its projection source has no derivatives. NavierStokes hits it every step, producing your exact error: (1, 3) + (2, 2).

Fix is one line in src/underworld3/systems/ddt.py. The regression test in PR #183 reproduces your nstt.py failure on a small mesh and passes with the fix. Once #183 lands, please re-run your original script and confirm.

[Chawgai]

Thanks @lmoresi and @ss2098 - After reinstalling the Underworld3 development version, I am still getting the same error,

Traceback (most recent call last):
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/systems/ddt.py", line 1754, in update_pre_solve
eval_result = uw.function.evaluate(
^^^^^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/function/functions_unit_system.py", line 315, in evaluate
raw_result_nondim = _evaluate_nd(
^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "src/underworld3/function/_function.pyx", line 817, in underworld3.function._function.evaluate_nd
File "src/underworld3/function/_function.pyx", line 567, in underworld3.function._function._project_to_work_variable
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/discretisation/enhanced_variables.py", line 130, in init
self._base_var = _BaseMeshVariable(
^^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/discretisation/discretisation_mesh_variables.py", line 245, in init
raise ValueError(
ValueError: Unable to infer variable type from num_components. Please explicitly set the vtype parameter.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/Users/u8085704/Library/CloudStorage/OneDrive-Personal/uw3d/nst/Finalised/CouetteDS.py", line 130, in
navier_stokes.solve(timestep=dt, zero_init_guess=False)
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/utilities/memprobe.py", line 269, in wrapper
return func(*args, **kw)
^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/systems/solvers.py", line 3776, in solve
self.DFDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf)
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/systems/ddt.py", line 1771, in update_pre_solve
self._psi_star_projection_solver.solve(verbose=verbose)
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "src/underworld3/cython/petsc_generic_snes_solvers.pyx", line 3863, in underworld3.cython.generic_solvers.SNES_MultiComponent.solve
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "src/underworld3/cython/petsc_generic_snes_solvers.pyx", line 721, in underworld3.cython.generic_solvers.SolverBaseClass._build
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/timing.py", line 326, in timed
result = routine(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "src/underworld3/cython/petsc_generic_snes_solvers.pyx", line 3549, in underworld3.cython.generic_solvers.SNES_MultiComponent._setup_pointwise_functions
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/utilities/_api_tools.py", line 423, in get
updated_value = self.value_fn(obj)
^^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/underworld3/systems/solvers.py", line 2416, in
lambda self: (self.u.sym - self.uw_function) * self.uw_weighting_function,
~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/core/decorators.py", line 118, in binary_op_wrapper
return func(self, other)
^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/matrices/matrixbase.py", line 3045, in sub
return self + (-a)
~~~~~^~~~~~
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/core/decorators.py", line 118, in binary_op_wrapper
return func(self, other)
^^^^^^^^^^^^^^^^^
File "/Users/u8085704/underworld3/.pixi/envs/dev/lib/python3.12/site-packages/sympy/matrices/matrixbase.py", line 2765, in add
raise ShapeError(f"Matrix size mismatch: {self.shape} + {other.shape}.")
sympy.matrices.exceptions.ShapeError: Matrix size mismatch: (1, 3) + (2, 2).
(underworld3:dev) (base) u8085704@arm64-apple-darwin20 Finalised %

[lmoresi]

Let me check. We made a couple of updates to the PR already, but maybe we did not catch your bug after all.

[lmoresi]

@Chawgai — thanks for re-running. From the traceback you're on the current development branch (your install path shows memprobe and your error is at ddt.py:1771, which is the pre-fix line number). The fix in PR #183 is not yet merged, so a plain git pull on development won't include it.

To test PR #183 directly, from your underworld3 checkout:

git fetch origin
git checkout bugfix/ns-ddt-projection-source
./uw build

Then re-run your script. Expected: the (1, 3) + (2, 2) ShapeError is gone.

A side note from your trace that's worth flagging separately: just before the fallback fires you have a

ValueError: Unable to infer variable type from `num_components`.
            Please explicitly set the `vtype` parameter.

raised from _project_to_work_variable in _function.pyx. That's what's causing your evaluate() to fail and drop into the projection fallback path on every step. It's a separate bug (probably fallout from one of yesterday's merges around mesh variable construction), and unrelated to the projection-shape fix in this PR. Once #183 lands and you can solve again, please re-file the vtype issue with a small repro so we can chase it cleanly — it shouldn't be raising during routine NS solves.

[ss2098]

That sounds right and thanks @Chawgai for confirming this. It looks like the projection-shape issue and the vtype inference message are probably two separate things. The branch @lmoresi suggested should be the right test for the (1, 3) + (2, 2) shape mismatch, but the earlier vtype message may still need to be isolated separately if it persists.

One small data point from my side: in a reduced diagnostic run, the user-defined 2-D fields looked dimensionally consistent:

meshbox.dim = 2
v_soln.sym.shape = (1, 2)
p_soln.sym.shape = (1, 1)
bodyforce.shape = (1, 2)

So I also did not see an obvious velocity/pressure/body-force shape problem in the user script. The failure seemed to occur later in the NavierStokesSLCN update/projection path.

[Chawgai]

Thanks @lmoresi for the updates - I am unable to incorporate the above changes, the output log is as under,

u8085704@its-089242 underworld3 % git fetch origin
u8085704@its-089242 underworld3 % git checkout bugfix/ns-ddt-projection-source
error: pathspec 'bugfix/ns-ddt-projection-source' did not match any file(s) known to git
u8085704@its-089242 underworld3 % git branch

• development
  u8085704@its-089242 underworld3 % git branch -r
  origin/development
  u8085704@its-089242 underworld3 % git log --oneline --decorate -5
  f029a02 (HEAD -> development, origin/development) Merge pull request Fix swarm particle loss across rank boundaries during advection (#175) #177 from underworldcode/bugfix/swarm-advection-migrate
  21b037c Merge pull request Add memprobe diagnostic module for memory-growth tracking (#176) #179 from underworldcode/bugfix/memprobe-diagnostics
  5f64de2 Fix dictionary-iteration race in UWexpression._ephemeral_expr_names
  389b1c9 Address Copilot review on PR Add memprobe diagnostic module for memory-growth tracking (#176) #179 + fix flaky KDTree count test
  9cce322 Add memprobe diagnostic module for tracking memory growth ([BUG] - memory leak #176)
  u8085704@its-089242 underworld3 % git grep live_count
  docs/developer/guides/memory-diagnostics.md:| KDTree live count | uw.kdtree.live_count() | free |
  docs/developer/guides/memory-diagnostics.md:Check uw.kdtree.live_count() directly, or look for the kdtree: live +N
  src/underworld3/ckdtree.pyx:# uw.utilities.memprobe.snapshot() or directly via uw.kdtree.live_count().
  src/underworld3/ckdtree.pyx:def live_count():
  src/underworld3/utilities/memprobe.py: "live": kdtree.live_count(),
  tests/test_0780_memprobe.py:def test_kdtree_live_count_tracks_construction_destruction():
  tests/test_0780_memprobe.py: before_live = uw.kdtree.live_count()
  tests/test_0780_memprobe.py: assert uw.kdtree.live_count() - before_live == 1
  tests/test_0780_memprobe.py: after_create_live = uw.kdtree.live_count()
  tests/test_0780_memprobe.py: assert uw.kdtree.live_count() - after_create_live == -1
  u8085704@its-089242 underworld3

My installation process was as under,

git clone -b development --single-branch https://github.com/underworldcode/underworld3.git
cd underworld3
./uw setup
./uw build

[ss2098]

@Chawgai - I reproduced the checkout issue locally. The problem is the --single-branch clone.

Because the repo was cloned with:

git clone -b development --single-branch https://github.com/underworldcode/underworld3.git

the local clone only knows about:

origin/development

so this fails:

git checkout bugfix/ns-ddt-projection-source

with:

error: pathspec 'bugfix/ns-ddt-projection-source' did not match any file(s) known to git

I tested the fix, and this works:

git fetch origin bugfix/ns-ddt-projection-source:bugfix/ns-ddt-projection-source
git checkout bugfix/ns-ddt-projection-source

After that, rebuild with:

./uw build

So the issue is not that the bugfix branch is unavailable; it is just that the --single-branch clone never fetched it.
Alternatively, a fresh clone without --single-branch should also make the branch visible after git fetch origin.

[Chawgai]

Thanks @ss2098 - That works for me as well. Thanks @lmoresi, the issue is resolved.

[lmoresi]

Thanks @ss2098 for helping out !