dsl: Introduce ability to define Functions on Subdomains #2245

EdCaunt · 2023-10-24T16:52:24Z

Add Functions on SubDomains functionality.

Remaining todo:

Tutorial notebook showing
- Demo for seafloor-like configuration
- Demo for damping-like configuration
- Demo for snapshotting only one region of the grid
Change subdomain numbering to be more robust and prevent memory leaks
Tests for propagator on subdomain
- Including with MPI

codecov · 2023-10-24T17:27:17Z

Codecov Report

Attention: Patch coverage is 96.04830% with 36 lines in your changes are missing coverage. Please review.

Project coverage is 86.98%. Comparing base (7fca4a8) to head (c868c91).

Files	Patch %	Lines
devito/types/grid.py	80.85%	18 Missing ⚠️
devito/types/dimension.py	70.00%	4 Missing and 2 partials ⚠️
devito/mpi/distributed.py	95.72%	3 Missing and 2 partials ⚠️
devito/types/sparse.py	20.00%	4 Missing ⚠️
devito/data/decomposition.py	97.29%	0 Missing and 1 partial ⚠️
tests/test_interpolation.py	99.39%	0 Missing and 1 partial ⚠️
tests/test_mpi.py	98.82%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2245      +/-   ##
==========================================
+ Coverage   86.80%   86.98%   +0.17%     
==========================================
  Files         233      233              
  Lines       43755    44534     +779     
  Branches     8078     8237     +159     
==========================================
+ Hits        37983    38739     +756     
- Misses       5063     5082      +19     
- Partials      709      713       +4

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

FabioLuporini · 2023-10-27T12:20:52Z

devito/types/grid.py

+            return None
+
+    @property
+    def has_grid(self):


FabioLuporini · 2023-10-27T12:21:21Z

devito/types/grid.py

+
+    @property
+    def distributor(self):
+        if self.has_grid:


The pythonic way of doing this is:

try: return self.grid._distributor except AttributeError: return None

FabioLuporini · 2023-10-27T12:21:35Z

devito/types/grid.py

-        for k, v, s in zip(self.define(grid.dimensions).keys(),
-                           self.define(grid.dimensions).values(), grid.shape):
+
+        for k, v, s in zip(self.define(self.grid.dimensions).keys(),


no need for .keys()

FabioLuporini · 2023-10-27T12:21:54Z

devito/types/grid.py

@@ -483,12 +500,17 @@ class SubDomain(AbstractSubDomain):
    """

    def __subdomain_finalize__(self, grid, **kwargs):
+        if self.has_grid:


if self.grid:

devito/types/grid.py

review-notebook-app · 2023-11-30T15:55:25Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

FabioLuporini

For the amount of changes introduced, I'd prefer to see a lot more new tests.

Shouldn't all (most) of the old tests that @rhodrin wrote be here as well?

FabioLuporini · 2024-01-25T10:48:21Z

devito/mpi/distributed.py

+class SubDomainDistributor(DenseDistributor):
+
+    """
+    Decompose a subdomain over a set of MPI processes.


nitpicking: SubDomain

I don't think this class should be (or really is aside from needing a couple of tweaks) tied specifically to SubDomain's - why not make the slightly more more general SubDistributor class?

Yeah, I think just SubDistributor would be good

FabioLuporini · 2024-01-25T10:48:30Z

devito/mpi/distributed.py

+    Parameters
+    ----------
+    subdomain : SubDomain
+        The subdomain to be decomposed


nitpicking: full stop

FabioLuporini · 2024-01-25T10:56:48Z

devito/mpi/distributed.py

+        decomposition of the grid.
+        """
+        dist_interval = {dim: Interval(s.start, s.stop-1)
+                         for dim, s in self.parent.glb_slices.items()}


for comprehensions, use d instead of dim . Not just for homogeneity, but also to minimize verbosity

FabioLuporini · 2024-01-25T10:57:53Z

devito/mpi/distributed.py

+        bounds_map = {**{dim.symbolic_min: 0
+                         for dim in self.parent.dimensions},
+                      **{dim.symbolic_max: sha-1
+                         for dim, sha in zip(self.parent.dimensions,


perhaps you could add properties

subdomain dimensions -> subdims
self.parent.dimensions -> dimensions
self.parent.glb_shape -> glb_shape

I seem to remember that making SubDistributor.dimensions not match SubDomain.dimensions requires a bunch of extra complexity anywhere it's accessed, so I want to leave these properties intact. Have introduced par_dimensions and par_shape for the latter two however, as I agree it would be tidier with the attributes.

FabioLuporini · 2024-01-25T10:59:02Z

devito/mpi/distributed.py

+        # The domain decomposition
+        decompositions = []
+        for dec, dim in zip(self.parent._decomposition, self.parent.dimensions):
+            decompositions.append([d if sdim_interval[dim] is None


we need to find a way of decreasing verbosity because in its current form the deciphering overhead is too high

FabioLuporini · 2024-01-25T11:20:37Z

devito/types/dense.py

@@ -333,8 +353,17 @@ def _size_outhalo(self):
                                                    self._distributor.myrank,
                                                    min(self.grid.shape_local))
                if not self._distributor.is_boundary_rank:
+                    # FIXME: Will throw an error if a Function on a Subdomain doesn't
+                    # enter any boundary ranks -> Needs to not do this
+                    print("First warning")


FabioLuporini · 2024-01-25T11:21:41Z

devito/types/dense.py

@@ -552,13 +582,17 @@ def _data_in_region(self, region, dim, side):
        Typically, this accessor won't be used in user code to set or read
        data values.
        """
+        if self._distributor:


the changes in this method seem very odd to me

They were leftover from something previous I think? They weren't doing anything so have been reverted.

FabioLuporini · 2024-01-25T11:22:53Z

devito/types/equation.py

@@ -169,6 +188,9 @@ def __str__(self):
    __repr__ = __str__


+_uxreplace_registry.register(Eq)


ah it's here now 😂 why?

To prevent a circular import. However, given what Rhodri said about explicitly specifying the subdomain, this could probably be put back

FabioLuporini · 2024-01-25T11:24:14Z

devito/types/grid.py

+        """The counter for a SubDomain"""
+        if self.grid:
+            try:
+                _subdomain_count[self.grid].add(id(self))


We need to rework all this because:

it's too ugly, being based on a global variable

it can leak memory

@georgebisbas indeed..

FabioLuporini · 2024-01-25T11:25:29Z

devito/types/grid.py

            if isinstance(v, Dimension):
                sub_dimensions.append(v)
                sdshape.append(s)
            else:
+                name = 'i%d%s' % (self._counter, k.name)
                try:
                    # Case ('middle', int, int)
                    side, thickness_left, thickness_right = v


Could you change it into side, ltkn, rtkn = v

mloubout

quick pass

mloubout · 2024-01-17T17:02:55Z

devito/mpi/distributed.py

@@ -12,7 +12,8 @@

 from devito.data import LEFT, CENTER, RIGHT, Decomposition
 from devito.parameters import configuration
-from devito.tools import EnrichedTuple, as_tuple, ctypes_to_cstr, filter_ordered
+from devito.tools import EnrichedTuple, as_tuple, ctypes_to_cstr, filter_ordered, \


convention through devito is parenthesise not line break

mloubout · 2024-01-17T17:03:32Z

devito/mpi/distributed.py

+        return self._comm
+
+    @property
+    def myrank(self):


mloubout · 2024-01-17T17:03:43Z

devito/mpi/distributed.py

+            return 0
+
+    @property
+    def mycoords(self):


devito/mpi/distributed.py

devito/types/dimension.py

devito/types/grid.py

mloubout · 2024-03-11T16:03:32Z

tests/test_caching.py

@@ -831,10 +831,13 @@ def test_solve(self, operate_on_empty_cache):
        # created by the finite difference (u.dt, u.dx2). We would have had
        # three extra references to u(t + dt), u(x - h_x) and u(x + h_x).
        # But this is not the case anymore!
+        # TODO: This cache doesn't clear currently as retained legacy SubDomain API


This sounds quite dangerous as it might end up preventing Function and TimeFucntion to be reachable by the gc

georgebisbas · 2024-04-23T13:55:42Z

devito/data/decomposition.py

@@ -204,9 +204,19 @@ def index_glb_to_loc(self, *args, rel=True):
        >>> d.index_glb_to_loc((1, 6), rel=False)
        (5, 6)
        """
+        # Need to offset the loc_abs_min, loc_abs_min, glb_min, and glb_max


loc_abs_max

devito/data/decomposition.py

devito/types/grid.py

georgebisbas · 2024-04-23T14:06:58Z

devito/types/grid.py

@@ -22,6 +24,8 @@

 GlobalLocal = namedtuple('GlobalLocal', 'glb loc')

+_subdomain_count = {}  # Track attachment of grids to SubDomains for dimension labelling


Not sure if this is good to have! Need to check!

we have to fix it

georgebisbas · 2024-04-23T14:07:58Z

devito/types/grid.py

+        self._distributor = SubDistributor(self)
+
+        # Intervals of form Interval(n, n) automatically become FiniteSet
+        # +1 as intervals are in terms of indices (inclusive of endpoints)


what does this +1 mean here?

georgebisbas · 2024-04-23T14:10:25Z

devito/types/grid.py

+                raise ValueError("`SubDomain` %s has already been attached to a `Grid`"
+                                 % self)
+            else:
+                self._grid = grid


georgebisbas · 2024-04-23T14:12:04Z

devito/types/grid.py

+        if grid:
+            if self.grid:
+                raise ValueError("`SubDomain` %s has already been attached to a `Grid`"
+                                 % self)
+            else:
+                self._grid = grid


Can this be done in less lines?

Dropped a line. Could I do this with a decorator, since it's used at the start of SubDomain.__init_finalize__?

georgebisbas · 2024-04-23T14:22:45Z

tests/test_mpi.py

+            if spec[0] == 'middle':
+                start = spec[1]
+                end = size - spec[2] - 1
+            elif spec[0] == 'left':
+                start = 0
+                end = spec[1] - 1
+            else:
+                start = size - spec[1]
+                end = size - 1


This can also be parametrized, having an 'expected' argument with the expected to be asserted result

I think that would probably make it messier and less extensible, given the combinatorial use of subdomain specifications and their orientation vs the domain decomposition

FabioLuporini · 2024-05-09T10:01:23Z

devito/operations/interpolators.py

@@ -226,6 +237,24 @@ def _interp_idx(self, variables, implicit_dims=None, pos_only=()):
        """
        Generate interpolation indices for the DiscreteFunctions in ``variables``.
        """
+        # Check if any Functions are defined on SubDomains


this new chunk of code deserves an utility function

FabioLuporini · 2024-05-09T10:06:15Z

devito/operator/operator.py

@@ -585,25 +585,34 @@ def _prepare_arguments(self, autotune=None, **kwargs):
        discretizations = {getattr(kwargs[p.name], 'grid', None) for p in overrides}
        discretizations.update({getattr(p, 'grid', None) for p in defaults})
        discretizations.discard(None)
-        # Remove subgrids if multiple grids
+
+        # Remove subgrids if multiple Grids


this chunk of code that pulls the grids and sanitizes them to ensure there's just one Grid could/should be now moved to a separate utility function, as I feel it could be reused in some other places as well (surely I found myself in need to extract the unique grid from a list of something... I think IIRC cluster.grid is one such case)

devito/operations/interpolators.py

FabioLuporini · 2024-05-09T10:14:55Z

devito/mpi/distributed.py

+        return self._dimension_map
+
+    @cached_property
+    def _d_interval(self):


btw _domain_interval would be clearer, I wouldn't have guessed it w/o a _sd_interval

…ns now run

mloubout · 2024-05-22T15:35:12Z

devito/operations/interpolators.py

+        cdmapper = {}
+        # Need to adjust bounds if Function defined on a SubDomain
+        if subdomain:
+            adjust_interp_indices(subdomain, mapper, smapper, cdmapper)


This seems overlly intricated and overkill to build everything then redo a pass to replace everything with the input dimensions. What's the problem with doing self._rrdims(grid) and self._gdims(grid) and catching directly the right dimensions instead of rebuilding everything?

Yeah I was thinking something along those lines. I figured it would be best to get it working then refactor it

mloubout · 2024-05-22T15:46:20Z

devito/types/basic.py

            if d in a.free_symbols:
-                # Shift by origin d -> d - o.
-                indices.append(sympy.sympify(a.subs(d, d - o).xreplace(s)))
+                # Shift by origin and offset d -> d - o - of.


Would it make sense to have origin just directly return the right one, or maybe local_origin ?

EdCaunt requested review from FabioLuporini and rhodrin October 24, 2023 16:52

mloubout added the API api (symbolics, types, ...) label Oct 24, 2023

FabioLuporini reviewed Oct 27, 2023

View reviewed changes

devito/types/grid.py Outdated

return None

@property

def has_grid(self):

Copy link

Contributor

FabioLuporini Oct 27, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

overkill

FabioLuporini reviewed Oct 27, 2023

View reviewed changes

devito/types/grid.py Show resolved Hide resolved

EdCaunt force-pushed the funcs_on_subdims branch from 5fc54ee to a0742f0 Compare November 30, 2023 15:55

EdCaunt force-pushed the funcs_on_subdims branch from 9fe8916 to 1092e78 Compare December 13, 2023 14:14

EdCaunt added the WIP Still work in progress label Dec 13, 2023

EdCaunt force-pushed the funcs_on_subdims branch from 1d37846 to d462af1 Compare January 17, 2024 16:15

EdCaunt requested review from FabioLuporini and mloubout January 17, 2024 16:16

FabioLuporini reviewed Jan 25, 2024

View reviewed changes

EdCaunt force-pushed the funcs_on_subdims branch from 588c25e to 2cd952e Compare February 2, 2024 15:26

EdCaunt force-pushed the funcs_on_subdims branch from 2cd952e to 341ee08 Compare February 23, 2024 14:15

EdCaunt requested a review from FabioLuporini February 23, 2024 14:15

mloubout reviewed Mar 11, 2024

View reviewed changes

EdCaunt requested a review from mloubout April 4, 2024 13:01

EdCaunt removed the WIP Still work in progress label Apr 4, 2024

EdCaunt force-pushed the funcs_on_subdims branch from 4268c54 to af06ba9 Compare April 8, 2024 15:41

georgebisbas reviewed Apr 23, 2024

View reviewed changes

FabioLuporini reviewed May 9, 2024

View reviewed changes

devito/operations/interpolators.py Show resolved Hide resolved

FabioLuporini reviewed May 9, 2024

View reviewed changes

EdCaunt added 17 commits May 21, 2024 17:02

dsl: Operators with interpolation/injection and Functions on SubDomai…

e8927f9

…ns now run

tests: Added test for interpolation off Functions on SubDomains

50dc29a

tests: Add test for injection into Function on SubDomain

e6a621d

misc: Tidy up, fix broken tests, address PR comments

9e9344f

dsl: Remove subdomain intersection

8361c62

tests: Add test for derivatives of functions on subdomains

47a8efa

tests: Add parallel derivatives test

b782a97

tests: Start work on MPI functions on subdomains interpolation test

bbf3edd

tests: Add mode to MPI tests

89c6698

dsl: Tweak order of argument preparation + misc investigations

2850284

dsl: Experiments with SubDimension overrides

5b0b141

dsl: Add switch to prevent interpolation off empty ranks

f4bc252

tests: Add test for interpolation off functions on subdomains with MPI

3f481ff

tests: Added test for injection into functions on subdomains with MPI

92fdae2

tests: Finish test for injection into functions on subdomains with MPI

16f47cf

tests: Tweak subdomain tests

b72ede0

misc: Tidy up

12d7169

EdCaunt force-pushed the funcs_on_subdims branch from c6b95d7 to 12d7169 Compare May 21, 2024 16:03

EdCaunt added 3 commits May 22, 2024 15:01

misc: Apply comments and fix flake8

a7ee8b7

misc: Refactoring

728cb4b

misc: Flake8

05dee8d

mloubout reviewed May 22, 2024

View reviewed changes

EdCaunt added 7 commits May 22, 2024 17:44

dsl: Move subdomain counter to inside class

edbcdd0

dsl: Reinstate legacy subdomain init

17f57e9

dsl: Refactor subdomain finalization

a38c6ae

dsl: Improve subdomain counter

d24ada6

tests: Add pickling test for functions on subdomains

ccb5934

tests: Consolidate pickling tests

167d963

dsl: Fix funcs on subdomains pickling

c868c91

		@@ -169,6 +188,9 @@ def __str__(self):
		__repr__ = __str__


		_uxreplace_registry.register(Eq)

		@@ -22,6 +24,8 @@

		GlobalLocal = namedtuple('GlobalLocal', 'glb loc')

		_subdomain_count = {} # Track attachment of grids to SubDomains for dimension labelling

dsl: Introduce ability to define Functions on Subdomains #2245

Are you sure you want to change the base?

dsl: Introduce ability to define Functions on Subdomains #2245

Conversation

EdCaunt commented Oct 24, 2023 • edited

codecov bot commented Oct 24, 2023 • edited

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

FabioLuporini Oct 27, 2023 • edited

Choose a reason for hiding this comment

review-notebook-app bot commented Nov 30, 2023

FabioLuporini left a comment • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mloubout left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

EdCaunt commented Oct 24, 2023 •

edited

codecov bot commented Oct 24, 2023 •

edited

FabioLuporini Oct 27, 2023 •

edited

FabioLuporini left a comment •

edited