bit of clean up

Author: ZoeLeibowitz

devito/petsc/equations.py

@@ -18,11 +18,8 @@ def lower_exprs_petsc(expressions, **kwargs):
 
 def constrain_essential_bcs(expressions, **kwargs):
     """
-    Expand ConstrainBC expressions to include halo regions by creating new
-    CustomDimensions for all relevant subdimensions and space dimensions,
-    then applying the mapper to all ConstrainBCs.
+    NOTE:
     """
-
     sregistry = kwargs['sregistry']
     new_exprs = []
 
@@ -40,7 +37,7 @@ def constrain_essential_bcs(expressions, **kwargs):
     space_dims = [d for d in all_dims if isinstance(d, SpaceDimension)]
 
     mapper = {}
-
+    # from IPython import embed; embed()
     for d in subdims:
         halo = halo_size[d]
 
@@ -52,7 +49,7 @@ def constrain_essential_bcs(expressions, **kwargs):
         )
 
         mapper[d] = CustomDimension(
-            name=sregistry.make_name(prefix=f"{d.name}_new"),
+            name=d.name,
             symbolic_min=Max(subdim_min, d.parent.symbolic_min - halo.left),
             symbolic_max=Min(subdim_max, d.parent.symbolic_max + halo.right),
         )
@@ -73,6 +70,7 @@ def constrain_essential_bcs(expressions, **kwargs):
         )
 
     # Apply mapper to all expressions
+    # from IPython import embed; embed()
     for e in expressions:
         if not isinstance(e, ConstrainBC):
             new_exprs.append(e)
@@ -83,7 +81,9 @@ def constrain_essential_bcs(expressions, **kwargs):
             new_exprs.append(e)
             continue
 
-        new_e = uxreplace(e, mapper)
+        # new_e = uxreplace(e, mapper)
+        new_e = e.subs(mapper)
+        
 
         if e.implicit_dims:
             new_e = new_e._rebuild(

examples/petsc/Poisson/ed_bueler_2d.py

@@ -103,7 +103,6 @@ def exact(x, y):
 bcs += [EssentialBC(u, bc, subdomain=sub4)]
 
 
-
 exprs = [eqn] + bcs
 petsc = petscsolve(
     exprs, target=u,
@@ -112,8 +111,6 @@ def exact(x, y):
     constrain_bcs=True
 )
 
-
-
 rank = grid.distributor.myrank
 # from IPython import embed; embed()
 

tests/test_petsc.py

@@ -2213,25 +2213,24 @@ def test_petsc_pi(self):
 class TestPetscSection:
     """
     These tests validate the use of `PetscSection` (from PETSc) to constrain essential
-    boundary nodes by removing them from the linear solver, rather than leaving them in
+    boundary nodes by removing them from the linear solver, rather than keeping them in
     the system as trivial equations.
 
     Users specify essential boundary conditions via the `EssentialBC` equation, with a specifed `SubDomain`.
     When `constrain_bcs=True` is passed to `petscsolve`, the Devito compiler generates code that
-    removes these degrees of freedom from the linear system. A requirement for this in PETSc is
-    that each MPI rank identifies ALL constrained nodes that lie within its local data region, including
+    removes these degrees of freedom from the linear system. A PETSc requirement is
+    that each MPI rank identifies ALL constrained nodes within its local data region, including
     non-owned (halo) nodes.
 
     To achieve this, the compiler creates new `EssentialBC`-like equations with modified (sub)dimensions
-    (to extend the loop bounds) that are utilised in two callback functions to constrain the nodes.
-    No non-owned (haloed) data is indexed into - the loops are purely used to specify the constrained "local" indices
-    on each rank.
+    (to extend the loop bounds), which are used in two callback functions to constrain the nodes.
+    No non-owned (halo) data is indexed into - the loops are only used to specify the constrained "local"
+    indices on each rank.
 
     Tests in this class use the following notation:
     - `x`   : a grid point
     - `[]`  : the `SubDomain` specified by the `EssentialBC` (the region being constrained)
     - `|`   : an MPI rank boundary
-
     """
     # first test that the loop generated is correct symbolically..
 
@@ -2272,10 +2271,7 @@ def _get_loop_bounds(self, shape, so, subdomain):
     @pytest.mark.parallel(mode=[1, 2, 4])
     def test_1_constrain_indices_1d(self, mode):
         # halo size 2
-        # 12 grid points in 1D with a subdomain spanning the region contained inside the brackets [], x represents a grid point, | represents the decomposition
-        # {rank: (min, max)} represents the loop bounds to constrain indices (locally) corresponding to the subdomain on this particular rank
-        # IMPROVE WORDING
-        # In petsc, when constraining indices, you constrain local (haloed) indices - INCLUDING NON OWNED INDICES
+        # 12 grid points
 
         # 1 rank:
         #            0   
@@ -2295,8 +2291,7 @@ def test_1_constrain_indices_1d(self, mode):
         #    0     1      2     3         
         # [x x x|x x x|x] x x|x x x
 
-        # Expected: {0: (0, 4), 1: (-2, 3), 2: (-2, 0), 3: (null loop - locally doesn't cover any of the subdomain)}
-        # rank 3 comes out as (-2,-3) -> null loop
+        # Expected: {0: (0, 4), 1: (-2, 3), 2: (-2, 0), 3: (-2, -3)}
 
         class Middle(SubDomain):
             name = 'submiddle'
@@ -2519,10 +2514,7 @@ def test_4_constrain_indices_1d(self, mode):
 
         # subdomain on the right side of the grid not left
         # halo size 2
-        # 24 grid points in 1D with a subdomain spanning the region contained inside the brackets [], x represents a grid point, | represents the decomposition
-        # {rank: (min, max)} represents the loop bounds to constrain indices (locally) corresponding to the subdomain on this particular rank
-        # IMPROVE WORDING
-        # In petsc, when constraining indices, you constrain local (haloed) indices - INCLUDING NON OWNED INDICES
+        # 24 grid points in 1D
 
         # 1 rank:
         #                        0   
@@ -2623,10 +2615,7 @@ def test_5_constrain_indices_1d(self, mode):
         # subdomain on the right side of the grid not left
         # halo size 4
         # same as test 4 but halo size 4 (so don't test 8 ranks)
-        # 24 grid points in 1D with a subdomain spanning the region contained inside the brackets [], x represents a grid point, | represents the decomposition
-        # {rank: (min, max)} represents the loop bounds to constrain indices (locally) corresponding to the subdomain on this particular rank
-        # IMPROVE WORDING
-        # In petsc, when constraining indices, you constrain local (haloed) indices - INCLUDING NON OWNED INDICES
+        # 24 grid points in 1D
 
         # 1 rank:
         #                        0   
@@ -2709,10 +2698,7 @@ def test_6_constrain_indices_1d(self, mode):
 
         # subdomain in the MIDDLE 
         # halo size 2
-        # 24 grid points in 1D with a subdomain spanning the region contained inside the brackets [], x represents a grid point, | represents the decomposition
-        # {rank: (min, max)} represents the loop bounds to constrain indices (locally) corresponding to the subdomain on this particular rank
-        # IMPROVE WORDING
-        # In petsc, when constraining indices, you constrain local (haloed) indices - INCLUDING NON OWNED INDICES
+        # 24 grid points in 1D
 
         # 1 rank:
         #                        0   
@@ -2811,10 +2797,7 @@ def test_7_constrain_indices_1d(self, mode):
 
         # subdomain in the MIDDLE 
         # halo size 4
-        # 24 grid points in 1D with a subdomain spanning the region contained inside the brackets [], x represents a grid point, | represents the decomposition
-        # {rank: (min, max)} represents the loop bounds to constrain indices (locally) corresponding to the subdomain on this particular rank
-        # IMPROVE WORDING
-        # In petsc, when constraining indices, you constrain local (haloed) indices - INCLUDING NON OWNED INDICES
+        # 24 grid points in 1D
 
         # 1 rank:
         #                        0