types: Fix SubDomain pickling for Dask compatibility

Fixes circular reference pickling bug where Grid and SubDomain
objects reference each other, causing failures in Dask workflows.

Changes:
- Implement lazy SubDistributor initialization via property getter
- Restore grid references in Grid.__setstate__ for legacy API
- Update examples/seismic/model.py to use new SubDomain API
- Add regression tests for SubDomain pickling (new and legacy API)
- Simplify redundant check in test_interpolate_subdomain_mpi_mfe

The issue occurred because SubDomain.__setstate__ tried to create
a SubDistributor before Grid was fully unpickled, causing
AttributeError. Now the distributor is created lazily on first
access, after both objects are fully restored.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: ggorman

devito/types/grid.py

@@ -402,6 +402,10 @@ def __setstate__(self, state):
         for k, v in state.items():
             setattr(self, k, v)
         self._distributor = Distributor(self.shape, self.dimensions, MPI.COMM_SELF)
+        # Restore grid references in subdomains (needed for legacy API)
+        for sd in self._subdomains:
+            if hasattr(sd, '_grid'):
+                sd._grid = self
 
 
 class AbstractSubDomain(CartesianDiscretization):
@@ -510,7 +514,22 @@ def stepping_dim(self):
 
     @property
     def distributor(self):
-        """The Distributor used for MPI-decomposing the CartesianDiscretization."""
+        """
+        The Distributor used for MPI-decomposing the SubDomain.
+
+        Lazy-initialized to handle circular reference pickling issues.
+        """
+        if self._distributor is None and self._grid is not None:
+            # Lazy initialization after unpickling or late binding
+            # Check Grid is fully initialized before creating SubDistributor
+            if (hasattr(self._grid, '_dimensions') and
+                    hasattr(self._grid, 'distributor') and
+                    self._grid.distributor is not None):
+                try:
+                    self._distributor = SubDistributor(self)
+                except AttributeError:
+                    # Grid still not ready, return None for now
+                    pass
         return self._distributor
 
     def is_distributed(self, dim):
@@ -681,8 +700,9 @@ def __getstate__(self):
     def __setstate__(self, state):
         for k, v in state.items():
             setattr(self, k, v)
-        if self.grid:
-            self._distributor = SubDistributor(self)
+        # Don't create distributor here - will be lazy-initialized on first access
+        # to avoid race condition when Grid hasn't been fully unpickled yet
+        self._distributor = None
 
 
 class MultiSubDomain(AbstractSubDomain):

examples/seismic/model.py

@@ -59,10 +59,10 @@ class PhysicalDomain(SubDomain):
 
     name = 'physdomain'
 
-    def __init__(self, so, fs=False):
-        super().__init__()
+    def __init__(self, so, fs=False, grid=None):
         self.so = so
         self.fs = fs
+        super().__init__(grid=grid)
 
     def define(self, dimensions):
         map_d = {d: d for d in dimensions}
@@ -75,9 +75,9 @@ class FSDomain(SubDomain):
 
     name = 'fsdomain'
 
-    def __init__(self, so):
-        super().__init__()
+    def __init__(self, so, grid=None):
         self.size = so
+        super().__init__(grid=grid)
 
     def define(self, dimensions):
         """
@@ -104,12 +104,8 @@ def __init__(self, origin, spacing, shape, space_order, nbl=20,
         origin_pml = [dtype(o - s*nbl) for o, s in zip(origin, spacing)]
         shape_pml = np.array(shape) + 2 * self.nbl
 
-        # Model size depending on freesurface
-        physdomain = PhysicalDomain(space_order, fs=fs)
-        subdomains = subdomains + (physdomain,)
+        # Free surface adjustments
         if fs:
-            fsdomain = FSDomain(space_order)
-            subdomains = subdomains + (fsdomain,)
             origin_pml[-1] = origin[-1]
             shape_pml[-1] -= self.nbl
 
@@ -118,8 +114,19 @@ def __init__(self, origin, spacing, shape, space_order, nbl=20,
         if grid is None:
             # Physical extent is calculated per cell, so shape - 1
             extent = tuple(np.array(spacing) * (shape_pml - 1))
+            # Create grid first (new API - no subdomains parameter)
             self.grid = Grid(extent=extent, shape=shape_pml, origin=origin_pml,
-                             dtype=dtype, subdomains=subdomains, topology=topology)
+                             dtype=dtype, topology=topology)
+
+            # Create subdomains with grid parameter (new API)
+            physdomain = PhysicalDomain(space_order, fs=fs, grid=self.grid)
+            all_subdomains = list(subdomains) + [physdomain]
+            if fs:
+                fsdomain = FSDomain(space_order, grid=self.grid)
+                all_subdomains.append(fsdomain)
+
+            # Add subdomains to grid
+            self.grid._subdomains = self.grid._subdomains + tuple(all_subdomains)
         else:
             self.grid = grid
 

tests/test_interpolation.py

@@ -1150,23 +1150,11 @@ def test_interpolate_subdomain_mpi_mfe(self, mode):
         # when interpolating from points outside the subdomain
         rank = grid.distributor.myrank
 
-        # Check all values on this rank are finite
-        if len(sr0.data) > 0:
-            print(f"Rank {rank} has {len(sr0.data)} values: {sr0.data}")
-            # Convert to numpy array to check for non-finite values
-            vals = np.asarray(sr0.data)
-            if not np.all(np.isfinite(vals)):
-                bad_indices = np.where(~np.isfinite(vals))[0]
-                raise AssertionError(
-                    f"BUG REPRODUCED! Rank {rank}, points {bad_indices}: "
-                    f"non-finite values {vals[bad_indices]} detected. "
-                    f"Full data: {vals}. "
-                    f"This occurs when interpolating from a Function on a SubDomain "
-                    f"for coordinates outside the subdomain."
-                )
-            # Verify all values are reasonable (not huge garbage)
-            assert np.all(np.abs(vals) < 1000), \
-                f"Rank {rank}: Suspicious large values: {vals}"
+        # Check all values are finite and reasonable (not huge garbage)
+        assert np.all(np.isfinite(sr0.data)), \
+            f"Rank {rank}: sr0 has non-finite values: {sr0.data}"
+        assert np.all(np.abs(sr0.data) < 1000), \
+            f"Rank {rank}: Suspicious large values: {sr0.data}"
 
     @pytest.mark.parallel(mode=4)
     def test_interpolate_multiple_subdomains_unique_temps(self, mode):

tests/test_pickle.py

@@ -127,6 +127,54 @@ def test_function(self, pickle, on_sd):
         assert f.dtype == new_f.dtype
         assert f.shape == new_f.shape
 
+    @pytest.mark.parametrize('use_new_api', [False, True])
+    def test_grid_with_subdomain(self, pickle, use_new_api):
+        """
+        Test pickling Grid with SubDomains (both old and new API).
+
+        Regression test for circular reference pickling bug where Grid
+        contains SubDomains that reference back to the Grid.
+        """
+        if use_new_api:
+            # New API: create SubDomain with grid parameter
+            grid = Grid(shape=(10, 10, 10))
+            sd = SD(grid=grid)
+            grid._subdomains = grid._subdomains + (sd,)
+        else:
+            # Old API: pass subdomains to Grid constructor
+            sd = SD()
+            grid = Grid(shape=(10, 10, 10), subdomains=(sd,))
+
+        # Pickle the grid
+        pkl_grid = pickle.dumps(grid)
+        new_grid = pickle.loads(pkl_grid)
+
+        # Verify subdomains were pickled correctly
+        assert len(new_grid._subdomains) == len(grid._subdomains)
+
+        # Verify subdomain distributor is accessible (lazy-initialized)
+        new_sd = [s for s in new_grid._subdomains if s.name == 'sd'][0]
+        assert new_sd.distributor is not None
+        assert new_sd.grid is new_grid
+
+    def test_subdomain_with_grid_circular_ref(self, pickle):
+        """
+        Test pickling SubDomain with grid reference (circular reference).
+
+        This tests the lazy distributor initialization fix for unpickling
+        SubDomains that have circular references to their Grid.
+        """
+        grid = Grid(shape=(10, 10, 10))
+        sd = SD(grid=grid)
+
+        # Pickle the subdomain (which references the grid)
+        pkl_sd = pickle.dumps(sd)
+        new_sd = pickle.loads(pkl_sd)
+
+        # Verify grid reference and distributor
+        assert new_sd.grid is not None
+        assert new_sd.distributor is not None
+
     @pytest.mark.parametrize('interp', ['linear', 'sinc'])
     def test_sparse_function(self, pickle, interp):
         grid = Grid(shape=(3,))