Global evaluation rewrite

[lmoresi]

This branch is to provide the global evaluation of variables by using a swarm to scatter points from the local process to the places we wish to sample.

1. Need to fix up the storage of data / evaluation of functions. The evaluate methods should return data that is of the same shape as the evaluated function. We also need to be able to store this on a swarm and pass it around, then retrieve it safely.

2. Decide on the interface we present to the user. Everything needs to look like a sympy matrix. Scalars are (X,1,1) in shape, vectors are (X,1,dim), tensors are (X,dim.dim) even if symmetric. We can squeeze arrays to remove extra dimensions but there should be a consistent pattern behind the scenes.

3. Simplify the access methods - there are many redundant operations: variables are switched between readable and writable but the readable versions are never seen in the user space. There is a lot of room to reduce the footprint of this awkward pattern

[lmoresi]

Don't merge - this branch is still for testing

[jcgraciosa]

I am testing out the changes in this pull request. I found the following issues so far:

1. uw.function.evaluate() produces errors when evaluating a matrix with a constant inside:

meshbox = uw.meshing.UnstructuredSimplexBox(minCoords  = (0., 0.),
                                            maxCoords  = (1., 1.),
                                            cellSize = 1. / 16,
                                            regular = False)

x, y = meshbox.N.x, meshbox.N.y
vector_fn2 = sympy.Matrix((x*x, 0)) 
uw.function.evaluate(vector_fn2, meshbox.data, rbf = False)

2. The pack/unpack results to mismatches if an asymmetrical tensor is used:

meshbox = uw.meshing.UnstructuredSimplexBox(minCoords   = (0., 0.),
                                            maxCoords   = (1., 1.),
                                            cellSize    = 1. / 16,
                                            regular     = False)

swarm = uw.swarm.Swarm(meshbox)

t = uw.swarm.SwarmVariable(
    "T",
    swarm,
    vtype=uw.VarType.SYM_TENSOR,
    proxy_degree=1,
    proxy_continuous=True,
    varsymbol=r"\mathbf{T}",
)

swarm.populate(0)

x, y = meshbox.N.x, meshbox.N.y
tensor_fn = sympy.Matrix(((x*x, x*x*y), (x*y, y*y))) # FIXME: this produces mismatch between packed and unpacked 

with swarm.access(t):
    t_data = uw.function.evaluate(tensor_fn, swarm.particle_coordinates.data, rbf=False)

t.pack(t_data)

print(f"mismatches: {(t_data != t.unpack()).sum()}")

Perhaps this is due to the vtype of t set as uw.VarType.SYM_TENSOR.
Changing the vtype to uw.VarType.TENSOR results to some array size mismatches.

[lmoresi]

@jcgraciosa - your point (1) above is (perhaps) fixed - go ahead and try again to break it ! Your point (2) is intentional / expected. If you pass a non-symmetric tensor to set a symmetric tensor, the non-symmetric part is ignored and the lower-triangular section takes precedence. We probably need to introduce the operator to split into symmetric / anti-symmetric parts, but here I am assuming that the input is a mistake. Sympy knows about symmetric matrices, so we could introduce a warning or error if the input is non-symmetric.

[Copilot]

Pull Request Overview

This PR implements a significant refactoring of the global evaluation system in Underworld3, focusing on how variables are evaluated across processes using swarms for distributed data handling. The main objectives are to improve data storage/evaluation consistency, standardize the user interface for matrix-like operations, and simplify variable access patterns.

Key changes include:

• Standardizing function evaluation to return consistently shaped arrays with .squeeze() calls
• Refactoring swarm variable data access patterns with new .array property and pack/unpack methods
• Splitting PIC swarm functionality into separate module for better organization
• Updating KDTree interface for consistent distance calculation options

Reviewed Changes

Copilot reviewed 38 out of 49 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
tests/test_*.py	Updated test files to use .squeeze() on evaluation results and modernized data access patterns
src/underworld3/swarm.py	Major refactoring with new array access patterns, pack/unpack methods, and PIC swarm extraction
src/underworld3/pic_swarm.py	New module containing extracted PIC swarm classes for better code organization
src/underworld3/kdtree.py	Enhanced KDTree interface with sqr_dists parameter for distance calculation control
src/underworld3/visualisation.py	Updated to use .squeeze() for consistent array shapes
src/underworld3/systems/*.py	Modified to use new array access patterns and .squeeze() for consistent returns
src/underworld3/utilities/_jitextension.py	Updated variable type handling and error messaging

Comments suppressed due to low confidence (1)

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}