Skip to content

Commit

Permalink
api: make SubDimension's side easy to check
Browse files Browse the repository at this point in the history
  • Loading branch information
mloubout committed Sep 20, 2023
1 parent ab160dd commit 94b5c07
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 14 deletions.
41 changes: 28 additions & 13 deletions devito/types/dimension.py
Original file line number Diff line number Diff line change
Expand Up @@ -586,7 +586,7 @@ def left(cls, name, parent, thickness, local=True):
return cls(name, parent,
left=parent.symbolic_min,
right=parent.symbolic_min+lst-1,
thickness=((lst, thickness), (rst, 0)),
thickness=((lst, thickness), (rst, None)),
local=local)

@classmethod
Expand All @@ -595,7 +595,7 @@ def right(cls, name, parent, thickness, local=True):
return cls(name, parent,
left=parent.symbolic_max-rst+1,
right=parent.symbolic_max,
thickness=((lst, 0), (rst, thickness)),
thickness=((lst, None), (rst, thickness)),
local=local)

@classmethod
Expand Down Expand Up @@ -628,6 +628,18 @@ def local(self):
def thickness(self):
return self._thickness

@property
def is_left(self):
return self.thickness.right[1] is None

@property
def is_right(self):
return self.thickness.left[1] is None

@property
def is_middle(self):
return not self.is_left and not self.is_right

@cached_property
def bound_symbols(self):
# Add thickness symbols
Expand Down Expand Up @@ -701,7 +713,7 @@ def _arg_values(self, interval, grid=None, **kwargs):
# However, arguments from the user are considered global
# So overriding the thickness to a nonzero value should not cause
# boundaries to exist between ranks where they did not before
requested_ltkn, requested_rtkn = (
r_ltkn, r_rtkn = (
kwargs.get(k.name, v) for k, v in self.thickness
)

Expand All @@ -710,19 +722,22 @@ def _arg_values(self, interval, grid=None, **kwargs):
if self.local:
# dimension is of type ``left``/right`` - compute the 'offset'
# and then add 1 to get the appropriate thickness
ltkn = grid.distributor.glb_to_loc(self.root, requested_ltkn-1, LEFT)
rtkn = grid.distributor.glb_to_loc(self.root, requested_rtkn-1, RIGHT)
ltkn = ltkn+1 if ltkn is not None else 0
rtkn = rtkn+1 if rtkn is not None else 0
if r_ltkn is not None:
ltkn = grid.distributor.glb_to_loc(self.root, r_ltkn-1, LEFT) + 1
else:
ltkn = 0

if r_rtkn is not None:
rtkn = grid.distributor.glb_to_loc(self.root, r_rtkn-1, RIGHT) + 1
else:
rtkn = 0
else:
# dimension is of type ``middle``
ltkn = grid.distributor.glb_to_loc(self.root, requested_ltkn,
LEFT) or 0
rtkn = grid.distributor.glb_to_loc(self.root, requested_rtkn,
RIGHT) or 0
ltkn = grid.distributor.glb_to_loc(self.root, r_ltkn, LEFT) or 0
rtkn = grid.distributor.glb_to_loc(self.root, r_rtkn, RIGHT) or 0
else:
ltkn = requested_ltkn
rtkn = requested_rtkn
ltkn = r_ltkn or 0
rtkn = r_rtkn or 0

return {i.name: v for i, v in zip(self._thickness_map, (ltkn, rtkn))}

Expand Down
9 changes: 8 additions & 1 deletion tests/test_dimension.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,10 @@ def test_subdim_middle(self, opt):
xi = SubDimension.middle(name='xi', parent=x,
thickness_left=1,
thickness_right=1)
assert xi.is_middle
assert not xi.is_left
assert not xi.is_right

eqs = [Eq(u.forward, u + 1)]
eqs = [e.subs(x, xi) for e in eqs]

Expand All @@ -261,6 +265,8 @@ def test_symbolic_size(self):
thickness = 4

xleft = SubDimension.left(name='xleft', parent=x, thickness=thickness)
assert xleft.is_left
assert not xleft.is_middle
assert xleft.symbolic_size == xleft.thickness.left[0]

xi = SubDimension.middle(name='xi', parent=x,
Expand Down Expand Up @@ -289,7 +295,8 @@ def test_bcs(self, opt):
xi = SubDimension.middle(name='xi', parent=x,
thickness_left=thickness, thickness_right=thickness)
xright = SubDimension.right(name='xright', parent=x, thickness=thickness)

assert xright.is_right
assert not xright.is_middle
yi = SubDimension.middle(name='yi', parent=y,
thickness_left=thickness, thickness_right=thickness)

Expand Down

0 comments on commit 94b5c07

Please sign in to comment.