Skip to content

Commit

Permalink
Merge pull request #40 from boutproject/y-boundaries-second-divertor
Browse files Browse the repository at this point in the history 
Allow keeping y boundary cells at second divertor
  • Loading branch information
@TomNicholas
TomNicholas authored Aug 7, 2019
2 parents 9502b3d + 94b16e7 commit 78a8b9f
Show file tree
Hide file tree
Showing 2 changed files with 156 additions and 29 deletions.
67 changes: 40 additions & 27 deletions xbout/load.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -262,50 +262,35 @@ def _trim(ds, *, guards, keep_boundaries, nxpe, nype):
# Relies on a change to xarray so datasets always have source encoding
# See xarray GH issue #2550
lower_boundaries, upper_boundaries = _infer_contains_boundaries(
ds.encoding['source'], nxpe, nype)
ds, nxpe, nype)
else:
lower_boundaries, upper_boundaries = {}, {}

selection = {}
for dim in ds.dims:
# Check for boundary cells, otherwise use guard cells, else leave alone
if keep_boundaries.get(dim, False):
if lower_boundaries.get(dim, False):
lower = None
else:
lower = guards[dim]
elif guards.get(dim, False):
lower = guards[dim]
else:
lower = None
if keep_boundaries.get(dim, False):
if upper_boundaries.get(dim, False):
upper = None
else:
upper = -guards[dim]
elif guards.get(dim, False):
upper = -guards[dim]
else:
upper = None
lower = _get_limit('lower', dim, keep_boundaries, lower_boundaries,
guards)
upper = _get_limit('upper', dim, keep_boundaries, upper_boundaries,
guards)
selection[dim] = slice(lower, upper)

trimmed_ds = ds.isel(**selection)

trimmed_ds = trimmed_ds.drop(_BOUT_TIMING_VARIABLES, errors='ignore')
return trimmed_ds.drop(_BOUT_TIMING_VARIABLES, errors='ignore')

return trimmed_ds


def _infer_contains_boundaries(filename, nxpe, nype):
def _infer_contains_boundaries(ds, nxpe, nype):
"""
Uses the name of the output file and the domain decomposition to work out
whether this dataset contains boundary cells, and on which side.
Uses the name of the output file, BOUT++'s topology indices, and the domain
decomposition to work out whether this dataset contains boundary cells, and on which
side.
Uses knowledge that BOUT names its output files as /folder/prefix.num.nc,
with a numbering scheme
num = nxpe*i + j, where i={0, ..., nype}, j={0, ..., nxpe}
"""

filename = ds.encoding['source']

*prefix, filenum, extension = Path(filename).suffixes
filenum = int(filenum.replace('.', ''))

Expand All @@ -317,4 +302,32 @@ def _infer_contains_boundaries(filename, nxpe, nype):
lower_boundaries['y'] = filenum < nxpe
upper_boundaries['y'] = filenum >= (nype-1)*nxpe

jyseps2_1 = int(ds['jyseps2_1'])
jyseps1_2 = int(ds['jyseps1_2'])
if jyseps1_2 > jyseps2_1:
# second divertor present
yproc = filenum // nxpe
ny_inner = int(ds['ny_inner'])
mysub = int(ds['MYSUB'])
if mysub*(yproc + 1) == ny_inner:
upper_boundaries['y'] = True
elif mysub*yproc == ny_inner:
lower_boundaries['y'] = True

return lower_boundaries, upper_boundaries


def _get_limit(side, dim, keep_boundaries, boundaries, guards):
# Check for boundary cells, otherwise use guard cells, else leave alone

if keep_boundaries.get(dim, False):
if boundaries.get(dim, False):
limit = None
else:
limit = guards[dim] if side is 'lower' else -guards[dim]
elif guards.get(dim, False):
limit = guards[dim] if side is 'lower' else -guards[dim]
else:
limit = None

return limit
118 changes: 116 additions & 2 deletions xbout/tests/test_load.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -430,9 +430,81 @@ def test_infer_boundaries_2d_parallelization(self, filenum, nxpe, nype,
-----> x
"""

filename = "folder0/BOUT.dmp." + str(filenum) + ".nc"
ds = create_test_data(0)
ds['jyseps2_1'] = 0
ds['jyseps1_2'] = 0
ds.encoding['source'] = "folder0/BOUT.dmp." + str(filenum) + ".nc"
actual_lower_boundaries, actual_upper_boundaries = _infer_contains_boundaries(
filename, nxpe, nype)
ds, nxpe, nype)

assert actual_lower_boundaries == lower_boundaries
assert actual_upper_boundaries == upper_boundaries

@pytest.mark.parametrize("filenum, nxpe, nype, lower_boundaries, upper_boundaries",
# 1d parallelization along y:
# Bottom
[(0, 1, 4, {'x': True, 'y': True},
{'x': True, 'y': False}),
# Lower Middle
(1, 1, 4, {'x': True, 'y': False},
{'x': True, 'y': True}),
# Upper Middle
(2, 1, 4, {'x': True, 'y': True},
{'x': True, 'y': False}),
# Top
(3, 1, 4, {'x': True, 'y': False},
{'x': True, 'y': True}),
# 2d parallelization:
# Bottom left corner
(0, 3, 4, {'x': True, 'y': True},
{'x': False, 'y': False}),
(1, 3, 4, {'x': False, 'y': True},
{'x': False, 'y': False}),
# Bottom right corner
(2, 3, 4, {'x': False, 'y': True},
{'x': True, 'y': False}),
(3, 3, 4, {'x': True, 'y': False},
{'x': False, 'y': True}),
(4, 3, 4, {'x': False, 'y': False},
{'x': False, 'y': True}),
(5, 3, 4, {'x': False, 'y': False},
{'x': True, 'y': True}),
(6, 3, 4, {'x': True, 'y': True},
{'x': False, 'y': False}),
(7, 3, 4, {'x': False, 'y': True},
{'x': False, 'y': False}),
(8, 3, 4, {'x': False, 'y': True},
{'x': True, 'y': False}),
# Top left corner
(9, 3, 4, {'x': True, 'y': False},
{'x': False, 'y': True}),
(10, 3, 4, {'x': False, 'y': False},
{'x': False, 'y': True}),
# Top right corner
(11, 3, 4, {'x': False, 'y': False},
{'x': True, 'y': True}),
])
def test_infer_boundaries_2d_parallelization_doublenull(
self, filenum, nxpe, nype, lower_boundaries, upper_boundaries):
"""
Numbering scheme for nxpe=3, nype=4
y 9 10 11
^ 6 7 8
| 3 4 5
| 0 1 2
-----> x
"""

ds = create_test_data(0)
ds['jyseps2_1'] = 3
ds['jyseps1_2'] = 11
ds['ny_inner'] = 8
ds['MYSUB'] = 4
ds.encoding['source'] = "folder0/BOUT.dmp." + str(filenum) + ".nc"
actual_lower_boundaries, actual_upper_boundaries = _infer_contains_boundaries(
ds, nxpe, nype)

assert actual_lower_boundaries == lower_boundaries
assert actual_upper_boundaries == upper_boundaries
Expand All @@ -444,10 +516,52 @@ def test_keep_xboundaries(self):
# Manually add filename - encoding normally added by xr.open_dataset
ds.encoding['source'] = 'folder0/BOUT.dmp.0.nc'

ds['jyseps2_1'] = 8
ds['jyseps1_2'] = 8

actual = _trim(ds, guards={'x': 2}, keep_boundaries={'x': True}, nxpe=1, nype=1)
expected = ds # Should be unchanged
xrt.assert_equal(expected, actual)

def test_keep_yboundaries(self):
ds = create_test_data(0)
ds = ds.rename({'dim2': 'y'})

# Manually add filename - encoding normally added by xr.open_dataset
ds.encoding['source'] = 'folder0/BOUT.dmp.0.nc'

ds['jyseps2_1'] = 8
ds['jyseps1_2'] = 8

actual = _trim(ds, guards={'y': 2}, keep_boundaries={'y': True}, nxpe=1, nype=1)
expected = ds # Should be unchanged
xrt.assert_equal(expected, actual)

@pytest.mark.parametrize("filenum, lower, upper",
[(0, True, False),
(1, False, True),
(2, True, False),
(3, False, True)])
def test_keep_yboundaries_doublenull(self, filenum, lower, upper):
ds = create_test_data(0)
ds = ds.rename({'dim2': 'y'})

# Manually add filename - encoding normally added by xr.open_dataset
ds.encoding['source'] = 'folder0/BOUT.dmp.'+str(filenum)+'.nc'

ds['jyseps2_1'] = 3
ds['jyseps1_2'] = 11
ds['ny_inner'] = 8
ds['MYSUB'] = 4

actual = _trim(ds, guards={'y': 2}, keep_boundaries={'y': True}, nxpe=1, nype=4)
expected = ds # Should be unchanged
if not lower:
expected = expected.isel(y=slice(2, None, None))
if not upper:
expected = expected.isel(y=slice(None, -2, None))
xrt.assert_equal(expected, actual)

def test_trim_timing_info(self):
ds = create_test_data(0)
from xbout.load import _BOUT_TIMING_VARIABLES
Expand Down

0 comments on commit 78a8b9f

Please sign in to comment.