Merge development (post-#186 merge) — combine theta and monotone_mode params

Author: lmoresi

src/underworld3/cython/petsc_types.pyx

@@ -1,9 +1,36 @@
-from libc.stdlib cimport malloc
+from libc.stdlib cimport malloc, free
 
 cdef class PtrContainer:
 
+    def __cinit__(self):
+        self.fns_residual = NULL
+        self.fns_bcs = NULL
+        self.fns_jacobian = NULL
+        self.fns_bd_residual = NULL
+        self.fns_bd_jacobian = NULL
+
+    def __dealloc__(self):
+        if self.fns_residual != NULL:
+            free(self.fns_residual)
+        if self.fns_bcs != NULL:
+            free(self.fns_bcs)
+        if self.fns_jacobian != NULL:
+            free(self.fns_jacobian)
+        if self.fns_bd_residual != NULL:
+            free(self.fns_bd_residual)
+        if self.fns_bd_jacobian != NULL:
+            free(self.fns_bd_jacobian)
+
     cpdef allocate(self, int n_res, int n_bcs, int n_jac, int n_bd_res, int n_bd_jac):
         """Allocate function pointer arrays of the given sizes."""
+
+        # Free existing memory if already allocated
+        if self.fns_residual != NULL: free(self.fns_residual)
+        if self.fns_bcs != NULL:      free(self.fns_bcs)
+        if self.fns_jacobian != NULL: free(self.fns_jacobian)
+        if self.fns_bd_residual != NULL: free(self.fns_bd_residual)
+        if self.fns_bd_jacobian != NULL: free(self.fns_bd_jacobian)
+
         self.fns_residual    = <PetscDSResidualFn*>   malloc(n_res    * sizeof(PetscDSResidualFn))
         self.fns_bcs         = <PetscDSResidualFn*>   malloc(n_bcs    * sizeof(PetscDSResidualFn))
         self.fns_jacobian    = <PetscDSJacobianFn*>   malloc(n_jac    * sizeof(PetscDSJacobianFn))

src/underworld3/discretisation/discretisation_mesh.py

@@ -285,6 +285,11 @@ def __init__(
         self._registered_swarms = weakref.WeakSet()
         self._registered_surfaces = weakref.WeakSet()  # Surfaces using this mesh
         self._registered_submeshes = weakref.WeakSet()  # Submeshes from extract_region
+
+        # _mesh_update_lock: Re-entrant lock to coordinate mesh deformation.
+        # Held by mesh_update_callback during _deform_mesh(). Checked by
+        # MeshVariable callbacks (blocking=False) to skip PETSc sync during
+        # sensitive coordinate changes.
         self._mesh_update_lock = threading.RLock()
 
         comm = PETSc.COMM_WORLD
@@ -594,14 +599,25 @@ class replacement_boundaries(Enum):
         # to handle that so we just wrap it here.
 
         def mesh_update_callback(array, change_context):
-            print(f"Mesh update callback - mesh deform")
-            coords = array.reshape(-1, array.owner.cdim)
-            self._deform_mesh(coords, verbose=True)
+            mesh = array.owner
+            if mesh is None:
+                # This guard handles cases where the array is accessed during
+                # object teardown (e.g. at application exit or during mesh
+                # replacement), where the owning Python mesh object has already
+                # been garbage collected but the NDArray proxy still exists.
+                return
+
+            if verbose:
+                uw.pprint(0, f"Mesh update callback - mesh deform")
+
+            coords = array.reshape(-1, mesh.cdim)
+            mesh._deform_mesh(coords, verbose=verbose)
 
             # Increment mesh version to notify registered swarms of coordinate changes
-            with self._mesh_update_lock:
-                self._mesh_version += 1
-                print(f"Mesh version incremented to {self._mesh_version}")
+            with mesh._mesh_update_lock:
+                mesh._mesh_version += 1
+                if verbose:
+                    uw.pprint(0, f"Mesh version incremented to {mesh._mesh_version}")
 
             return
 
@@ -1239,11 +1255,12 @@ def _build_vertex_map(self):
         Uses coordinate matching at extraction time (before any
         deformation). Cached permanently since topology doesn't change.
         """
-        if hasattr(self, '_vertex_map') and self._vertex_map is not None:
+        if hasattr(self, "_vertex_map") and self._vertex_map is not None:
             return self._vertex_map
 
-        tree = uw.kdtree.KDTree(self.X.coords)
-        dists, indices = tree.query(self.parent.X.coords, sqr_dists=False)
+        # Use cached KDTree from coordinate variable
+        tree = self.X._get_kdtree()
+        dists, indices = tree.query(self.parent.X.coords_nd, sqr_dists=False)
         matched = dists < 1.0e-10
 
         # parent_rows[i] -> sub_rows[i]: matched vertex pairs
@@ -1337,10 +1354,14 @@ def _re_extract_from_parent(self, verbose=False):
             )
 
             def mesh_update_callback(array, change_context):
-                coords = array.reshape(-1, array.owner.cdim)
-                self._deform_mesh(coords, verbose=False)
-                with self._mesh_update_lock:
-                    self._mesh_version += 1
+                mesh = array.owner
+                if mesh is None:
+                    return
+
+                coords = array.reshape(-1, mesh.cdim)
+                mesh._deform_mesh(coords, verbose=False)
+                with mesh._mesh_update_lock:
+                    mesh._mesh_version += 1
                 return
 
             self._coords.add_callback(mesh_update_callback)
@@ -1433,8 +1454,8 @@ def _build_dof_map(self, parent_var, sub_var):
         if key in self._dof_maps:
             return self._dof_maps[key]
 
-        tree = uw.kdtree.KDTree(sub_var.coords)
-        dists, indices = tree.query(parent_var.coords, sqr_dists=False)
+        tree = sub_var._get_kdtree()
+        dists, indices = tree.query(parent_var.coords_nd, sqr_dists=False)
         matched = dists < 1.0e-10
 
         # indices[matched] maps parent row → sub row
@@ -1777,42 +1798,56 @@ def _deform_mesh(self, new_coords: numpy.ndarray, verbose=False):
         The coord array that is passed in should match the shape of self.data
         """
 
-        coord_vec = self.dm.getCoordinatesLocal()
-        coords = coord_vec.array.reshape(-1, self.cdim)
-        coords[...] = new_coords[...]
-
-        self.dm.setCoordinatesLocal(coord_vec)
-        self.nuke_coords_and_rebuild()
-
-        # Rebuild the _coords array view.  nuke_coords_and_rebuild may
-        # replace the coordinate vector internally (createCoordinateSpace),
-        # leaving self._coords as a stale numpy view of the old buffer.
-        import underworld3.utilities
-        old_callbacks = getattr(self._coords, "_callbacks", [])
-        self._coords = underworld3.utilities.NDArray_With_Callback(
-            numpy.ndarray.view(
-                self.dm.getCoordinatesLocal().array.reshape(-1, self.cdim)
-            ),
-            owner=self,
-        )
-        for cb in old_callbacks:
-            self._coords.add_callback(cb)
-
-        # BUGFIX(#122): mark registered solvers for rebuild. Since PR #127
-        # ("Trust JIT cache: skip DM rebuild on constant-only parameter
-        # changes") a solver with is_setup=True trusts its cached PETSc DM
-        # / SNES assembly and skips rebuild on the next solve(). After a
-        # coordinate change the cached DM still carries pre-deform
-        # coordinates, so F(v_prev) ≈ 0 and the solver converges in 0
-        # iterations without updating the solution. mesh.adapt() already
-        # does this; _deform_mesh must match.
-        for solver in self._equation_systems_register:
-            if solver is not None and hasattr(solver, "is_setup"):
-                solver.is_setup = False
+        with self._mesh_update_lock:
+            coord_vec = self.dm.getCoordinatesLocal()
+            coords = coord_vec.array.reshape(-1, self.cdim)
+            coords[...] = new_coords[...]
+
+            self.dm.setCoordinatesLocal(coord_vec)
+            self.nuke_coords_and_rebuild()
+
+            # Rebuild the _coords array view.  nuke_coords_and_rebuild may
+            # replace the coordinate vector internally (createCoordinateSpace),
+            # leaving self._coords as a stale numpy view of the old buffer.
+            import underworld3.utilities
+            old_callbacks = getattr(self._coords, "_callbacks", [])
+            self._coords = underworld3.utilities.NDArray_With_Callback(
+                numpy.ndarray.view(
+                    self.dm.getCoordinatesLocal().array.reshape(-1, self.cdim)
+                ),
+                owner=self,
+            )
+            for cb in old_callbacks:
+                self._coords.add_callback(cb)
+
+            # BUGFIX(#122): mark registered solvers for rebuild. Since PR #127
+            # ("Trust JIT cache: skip DM rebuild on constant-only parameter
+            # changes") a solver with is_setup=True trusts its cached PETSc DM
+            # / SNES assembly and skips rebuild on the next solve(). After a
+            # coordinate change the cached DM still carries pre-deform
+            # coordinates, so F(v_prev) ≈ 0 and the solver converges in 0
+            # iterations without updating the solution. mesh.adapt() already
+            # does this; _deform_mesh must match.
+            for solver in self._equation_systems_register:
+                if solver is not None and hasattr(solver, "is_setup"):
+                    solver.is_setup = False
+
+            # Invalidate caches whose contents become stale when mesh
+            # coordinates change. Matches the cache hygiene already
+            # performed by mesh.adapt() and _legacy_access. Without
+            # these, uw.function.evaluate (and any user code that keys
+            # lookups off _topology_version) can return values
+            # computed against the pre-deform mesh.
+            self._evaluation_hash = None
+            self._evaluation_interpolated_results = None
+            if hasattr(self, '_dminterpolation_cache'):
+                self._dminterpolation_cache.invalidate_all(
+                    reason="mesh deformed")
+            self._topology_version += 1
 
-        # Propagate coordinate changes to registered submeshes
-        for submesh in self._registered_submeshes:
-            submesh.sync_coordinates_from_parent()
+            # Propagate coordinate changes to registered submeshes
+            for submesh in self._registered_submeshes:
+                submesh.sync_coordinates_from_parent()
 
         return
 
@@ -3455,11 +3490,24 @@ def _get_mesh_centroids(self):
     def _get_domain_centroids(self):
 
         import numpy as np
+        from underworld3.utilities import gather_data
 
         domain_centroid = self._centroids.mean(axis=0)
         all_centroids = gather_data(domain_centroid, bcast=True).reshape(-1, self.dim)
         return all_centroids
 
+    def _get_domain_kdtree(self):
+        import underworld3 as uw
+        if (
+            not hasattr(self, "_domain_kdtree")
+            or self._domain_kdtree is None
+            or getattr(self, "_domain_kdtree_version", -1) != self._mesh_version
+        ):
+            centroids = self._get_domain_centroids()
+            self._domain_kdtree = uw.kdtree.KDTree(centroids)
+            self._domain_kdtree_version = self._mesh_version
+        return self._domain_kdtree
+
     def get_min_radius_old(self) -> float:
         """
         This method returns the global minimum distance from any cell centroid to a face.
@@ -3777,12 +3825,15 @@ def adapt(self, metric_field, verbose=False):
 
             # Rebuild the callback for mesh deformation
             def mesh_update_callback(array, change_context):
-                print(f"Mesh update callback - mesh deform")
+                if verbose:
+                    uw.pprint(0, f"Mesh update callback - mesh deform")
+
                 coords = array.reshape(-1, array.owner.cdim)
-                self._deform_mesh(coords, verbose=True)
+                self._deform_mesh(coords, verbose=verbose)
                 with self._mesh_update_lock:
                     self._mesh_version += 1
-                    print(f"Mesh version incremented to {self._mesh_version}")
+                    if verbose:
+                        uw.pprint(0, f"Mesh version incremented to {self._mesh_version}")
                 return
 
             self._coords.add_callback(mesh_update_callback)