From 0bb5e281df95f96ca9b10a8a6b5e9103c11e47df Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 24 Oct 2019 14:49:30 +0100 Subject: [PATCH 01/10] Return result from geometries._set_default_toroidal_coordinates() Previously the return statement was missing, leading to a bug. --- xbout/geometries.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xbout/geometries.py b/xbout/geometries.py index cf94060b..3ec2c809 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -93,6 +93,8 @@ def _set_default_toroidal_coordinates(coordinates): coordinates['y'] = coordinates.get('y', 'theta') coordinates['z'] = coordinates.get('z', 'phi') + return coordinates + @register_geometry('toroidal') def add_toroidal_geometry_coords(ds, coordinates=None): From 807c8e3d0e10f335355ccfdc39e4765691cca201 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 24 Oct 2019 18:21:51 +0100 Subject: [PATCH 02/10] Add tests opening 'toroidal' and 's-alpha' geometries --- xbout/tests/test_load.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 40709f67..a67c430c 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -364,7 +364,7 @@ def test_strip_metadata(self): # TODO also test loading multiple files which have guard cells -class TestCombineNoTrim: +class TestOpen: def test_single_file(self, tmpdir_factory, bout_xyt_example_files): path = bout_xyt_example_files(tmpdir_factory, nxpe=1, nype=1, nt=1) actual = open_boutdataset(datapath=path, keep_xboundaries=False) @@ -419,6 +419,16 @@ def test_combine_along_xy(self, tmpdir_factory, bout_xyt_example_files): expected.drop(METADATA_VARS + _BOUT_PER_PROC_VARIABLES, errors='ignore')) + def test_toroidal(self, tmpdir_factory, bout_xyt_example_files): + path = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=3, nt=1, + syn_data_type='stepped') + actual = open_boutdataset(datapath=path, geometry='toroidal') + + def test_salpha(self, tmpdir_factory, bout_xyt_example_files): + path = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=3, nt=1, + syn_data_type='stepped') + actual = open_boutdataset(datapath=path, geometry='s-alpha') + @pytest.mark.skip def test_combine_along_tx(self): ... From f8fa440fd10685bc2c72ec885476282dfb7de8ed Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 24 Oct 2019 19:06:23 +0100 Subject: [PATCH 03/10] Add option to pass grid file May be needed to provide some variables that are not saved by default to BOUT++'s output files, e.g. psixy, Rxy, Zxy. --- xbout/geometries.py | 10 ++++++++++ xbout/load.py | 16 +++++++++++++++- xbout/tests/test_load.py | 35 ++++++++++++++++++++++++++++++----- 3 files changed, 55 insertions(+), 6 deletions(-) diff --git a/xbout/geometries.py b/xbout/geometries.py index 3ec2c809..c22bde0f 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -108,6 +108,16 @@ def add_toroidal_geometry_coords(ds, coordinates=None): "Use the 'coordinates' argument of open_boutdataset to provide " "alternative names".format(bad_names)) + # Get extra geometry information from grid file if it's not in the dump files + needed_variables = ['psixy', 'Rxy', 'Zxy'] + for v in needed_variables: + if v not in ds: + if ds._grid is None: + raise ValueError("Grid file is required to provide %s. Pass the grid " + "file name as the 'gridfilepath' argument to " + "open_boutdataset().") + ds[v] = ds._grid[v] + # Change names of dimensions to Orthogonal Toroidal ones ds = ds.rename(y=coordinates['y']) diff --git a/xbout/load.py b/xbout/load.py index d0cc4881..923b225c 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -39,7 +39,7 @@ def open_boutdataset(datapath='./BOUT.dmp.*.nc', inputfilepath=None, - geometry=None, chunks={}, + geometry=None, gridfilepath=None, chunks={}, keep_xboundaries=True, keep_yboundaries=False, run_name=None, info=True): """ @@ -68,6 +68,9 @@ def open_boutdataset(datapath='./BOUT.dmp.*.nc', inputfilepath=None, To define a new type of geometry you need to use the `register_geometry` decorator. You are encouraged to do this for your own BOUT++ physics module, to apply relevant normalisations. + gridfilepath : str, optional + The path to a grid file, containing any variables needed to apply the geometry + specified by the 'geometry' option, which are not contained in the dump files. keep_xboundaries : bool, optional If true, keep x-direction boundary cells (the cells past the physical edges of the grid, where boundary conditions are set); increases the @@ -126,6 +129,17 @@ def open_boutdataset(datapath='./BOUT.dmp.*.nc', inputfilepath=None, if geometry: if info: print("Applying {} geometry conventions".format(geometry)) + + if gridfilepath is not None: + print('here in load', ds.attrs) + ds.attrs["_grid"] = open_boutdataset(gridfilepath, chunks=chunks, + keep_xboundaries=keep_xboundaries, + keep_yboundaries=keep_yboundaries, + info=info) + print('after in load', ds.attrs) + else: + ds.attrs["_grid"] = None + # Update coordinates to match particular geometry of grid ds = geometries.apply_geometry(ds, geometry) else: diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index a67c430c..0a60c3ac 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -169,7 +169,8 @@ def bout_xyt_example_files(tmpdir_factory): def _bout_xyt_example_files(tmpdir_factory, prefix='BOUT.dmp', lengths=(6, 2, 4, 7), - nxpe=4, nype=2, nt=1, guards={}, syn_data_type='random'): + nxpe=4, nype=2, nt=1, guards={}, syn_data_type='random', + grid=None): """ Mocks up a set of BOUT-like netCDF files, and return the temporary test directory containing them. @@ -184,6 +185,13 @@ def _bout_xyt_example_files(tmpdir_factory, prefix='BOUT.dmp', lengths=(6, 2, 4, for ds, file_name in zip(ds_list, file_list): ds.to_netcdf(str(save_dir.join(str(file_name)))) + if grid is not None: + xsize = lengths[1]*nxpe + ysize = lengths[2]*nype + grid_ds = create_bout_grid_ds(xsize=xsize, ysize=ysize, guards=guards) + print('check grid_ds',xsize,ysize,grid_ds) + grid_ds.to_netcdf(str(save_dir.join(grid + ".nc"))) + # Return a glob-like path to all files created, which has all file numbers replaced with a single asterix path = str(save_dir.join(str(file_list[-1]))) @@ -341,6 +349,21 @@ def create_bout_ds(syn_data_type='random', lengths=(6, 2, 4, 7), num=0, nxpe=1, return ds +def create_bout_grid_ds(xsize=2, ysize=4, guards={}): + + # Set the shape of the data in this dataset + mxg = guards.get('x', 0) + myg = guards.get('y', 0) + xsize += 2*mxg + ysize += 2*myg + shape = (xsize, ysize) + + data = DataArray(np.ones(shape), dims=['x', 'y']) + + ds = Dataset({'psixy': data, 'Rxy': data, 'Zxy': data, 'hthe': data}) + + return ds + # Note, MYPE, PE_XIND and PE_YIND not included, since they are different for each # processor and so are dropped when loading datasets. @@ -421,13 +444,15 @@ def test_combine_along_xy(self, tmpdir_factory, bout_xyt_example_files): def test_toroidal(self, tmpdir_factory, bout_xyt_example_files): path = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=3, nt=1, - syn_data_type='stepped') - actual = open_boutdataset(datapath=path, geometry='toroidal') + syn_data_type='stepped', grid='grid') + actual = open_boutdataset(datapath=path, geometry='toroidal', + gridfilepath=Path(path).parent.joinpath('grid.nc')) def test_salpha(self, tmpdir_factory, bout_xyt_example_files): path = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=3, nt=1, - syn_data_type='stepped') - actual = open_boutdataset(datapath=path, geometry='s-alpha') + syn_data_type='stepped', grid='grid') + actual = open_boutdataset(datapath=path, geometry='s-alpha', + gridfilepath=Path(path).parent.joinpath('grid.nc')) @pytest.mark.skip def test_combine_along_tx(self): From 4fad704de1f82167bc25913c97e0a96223565671 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 24 Oct 2019 22:45:55 +0100 Subject: [PATCH 04/10] Add 'hthe' from grid in s-alpha geometry Need to add hthe before getting toroidal coordinates, as dimension names are changed only in the Dataset, not in the _grid member variable. 'r' coordinate is created as 1d, so selecting 'theta=0' part is an error. --- xbout/geometries.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/xbout/geometries.py b/xbout/geometries.py index c22bde0f..86deca49 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -155,6 +155,14 @@ def add_s_alpha_geometry_coords(ds, coordinates=None): coordinates = _set_default_toroidal_coordinates(coordinates) + # Add 'hthe' from grid file, needed below for radial coordinate + if not 'hthe' in ds: + if ds._grid is None: + raise ValueError("Grid file is required to provide %s. Pass the grid " + "file name as the 'gridfilepath' argument to " + "open_boutdataset().") + ds['hthe'] = ds._grid['hthe'] + ds = add_toroidal_geometry_coords(ds, coordinates=coordinates) # Add 1D radial coordinate @@ -166,7 +174,4 @@ def add_s_alpha_geometry_coords(ds, coordinates=None): ds = ds.set_coords('r') ds = ds.rename(x='r') - # Simplify psi to be radially-varying only - ds['r'] = ds['r'].isel({coordinates['y']: 0}).squeeze(drop=True) - return ds From 181a691807f2ed1031f60dd64a332cdce19add04 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 25 Oct 2019 09:08:15 +0100 Subject: [PATCH 05/10] Travis workaround to skip xarray-0.14.0, which breaks the tests --- .travis.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.travis.yml b/.travis.yml index 7fffb52e..ba5fefca 100644 --- a/.travis.yml +++ b/.travis.yml @@ -4,7 +4,7 @@ python: - "3.7" install: - pip install --upgrade vtk # fixes install of mayavi for Python-3.6, which is a dependency of boutdata - - pip install --upgrade setuptools pip pytest pytest-cov coverage codecov boutdata + - pip install --upgrade setuptools pip pytest pytest-cov coverage codecov boutdata "xarray!=0.14.0" - pip install -r requirements.txt - pip install -e . script: From 60b9d9a94da1450192fe18bc181df3cb097be53a Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 25 Oct 2019 18:27:55 +0100 Subject: [PATCH 06/10] Fix PEP8 issuses, remove debugging print statements --- xbout/geometries.py | 2 +- xbout/load.py | 2 -- xbout/tests/test_load.py | 3 +-- 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/xbout/geometries.py b/xbout/geometries.py index 86deca49..db6146c6 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -156,7 +156,7 @@ def add_s_alpha_geometry_coords(ds, coordinates=None): coordinates = _set_default_toroidal_coordinates(coordinates) # Add 'hthe' from grid file, needed below for radial coordinate - if not 'hthe' in ds: + if 'hthe' not in ds: if ds._grid is None: raise ValueError("Grid file is required to provide %s. Pass the grid " "file name as the 'gridfilepath' argument to " diff --git a/xbout/load.py b/xbout/load.py index 923b225c..10c8dc45 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -131,12 +131,10 @@ def open_boutdataset(datapath='./BOUT.dmp.*.nc', inputfilepath=None, print("Applying {} geometry conventions".format(geometry)) if gridfilepath is not None: - print('here in load', ds.attrs) ds.attrs["_grid"] = open_boutdataset(gridfilepath, chunks=chunks, keep_xboundaries=keep_xboundaries, keep_yboundaries=keep_yboundaries, info=info) - print('after in load', ds.attrs) else: ds.attrs["_grid"] = None diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 0a60c3ac..c7dbba2d 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -76,7 +76,6 @@ def test_no_files(self, tmpdir): with pytest.raises(IOError): path = Path(str(files_dir.join('run*/example.*.nc'))) actual_filepaths = _expand_filepaths(path) - print(actual_filepaths) @pytest.fixture() @@ -189,7 +188,6 @@ def _bout_xyt_example_files(tmpdir_factory, prefix='BOUT.dmp', lengths=(6, 2, 4, xsize = lengths[1]*nxpe ysize = lengths[2]*nype grid_ds = create_bout_grid_ds(xsize=xsize, ysize=ysize, guards=guards) - print('check grid_ds',xsize,ysize,grid_ds) grid_ds.to_netcdf(str(save_dir.join(grid + ".nc"))) # Return a glob-like path to all files created, which has all file numbers replaced with a single asterix @@ -349,6 +347,7 @@ def create_bout_ds(syn_data_type='random', lengths=(6, 2, 4, 7), num=0, nxpe=1, return ds + def create_bout_grid_ds(xsize=2, ysize=4, guards={}): # Set the shape of the data in this dataset From cfc039c32d4097e6b16cd3189b9faa68a182b4d6 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 26 Oct 2019 11:53:27 +0100 Subject: [PATCH 07/10] Don't store _grid in dataset attrs Doing this prevented the dataset being saved to netCDF file. Also use _open_grid instead of open_boutdataset to open the grid file so that the grid dataset does not have metadata added to the DataArray variables. These variables are added as coordinates and become members of the DataArrays representing simulation variables; if they have a metadata dict, it is not possible to save them to netCDF. --- xbout/boutdataset.py | 3 --- xbout/geometries.py | 8 ++++---- xbout/load.py | 9 ++++----- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index e4e47d30..364d0731 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -19,7 +19,6 @@ def __init__(self, ds): self.data = ds self.metadata = ds.attrs.get('metadata') # None if just grid file self.options = ds.attrs.get('options') # None if no inp file - self.grid = ds.attrs.get('grid') # None if no grid file def __str__(self): """ @@ -34,8 +33,6 @@ def __str__(self): "Metadata:\n{}\n".format(styled(self.metadata)) if self.options: text += "Options:\n{}".format(styled(self.options)) - if self.grid: - text += "Grid:\n{}".format(styled(self.grid)) return text #def __repr__(self): diff --git a/xbout/geometries.py b/xbout/geometries.py index db6146c6..e8c395df 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -112,11 +112,11 @@ def add_toroidal_geometry_coords(ds, coordinates=None): needed_variables = ['psixy', 'Rxy', 'Zxy'] for v in needed_variables: if v not in ds: - if ds._grid is None: + if ds.bout._grid is None: raise ValueError("Grid file is required to provide %s. Pass the grid " "file name as the 'gridfilepath' argument to " "open_boutdataset().") - ds[v] = ds._grid[v] + ds[v] = ds.bout._grid[v] # Change names of dimensions to Orthogonal Toroidal ones ds = ds.rename(y=coordinates['y']) @@ -157,11 +157,11 @@ def add_s_alpha_geometry_coords(ds, coordinates=None): # Add 'hthe' from grid file, needed below for radial coordinate if 'hthe' not in ds: - if ds._grid is None: + if ds.bout._grid is None: raise ValueError("Grid file is required to provide %s. Pass the grid " "file name as the 'gridfilepath' argument to " "open_boutdataset().") - ds['hthe'] = ds._grid['hthe'] + ds['hthe'] = ds.bout._grid['hthe'] ds = add_toroidal_geometry_coords(ds, coordinates=coordinates) diff --git a/xbout/load.py b/xbout/load.py index 10c8dc45..61a9b353 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -131,12 +131,11 @@ def open_boutdataset(datapath='./BOUT.dmp.*.nc', inputfilepath=None, print("Applying {} geometry conventions".format(geometry)) if gridfilepath is not None: - ds.attrs["_grid"] = open_boutdataset(gridfilepath, chunks=chunks, - keep_xboundaries=keep_xboundaries, - keep_yboundaries=keep_yboundaries, - info=info) + ds.bout._grid = _open_grid(gridfilepath, chunks=chunks, + keep_xboundaries=keep_xboundaries, + keep_yboundaries=keep_yboundaries) else: - ds.attrs["_grid"] = None + ds.bout._grid = None # Update coordinates to match particular geometry of grid ds = geometries.apply_geometry(ds, geometry) From 7ae0f18d4621c144323426eee34dfe86724dd1a5 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 26 Oct 2019 12:17:45 +0100 Subject: [PATCH 08/10] Test saving BoutDataset with geometry --- xbout/tests/test_load.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index c7dbba2d..cb0f40fc 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -447,12 +447,20 @@ def test_toroidal(self, tmpdir_factory, bout_xyt_example_files): actual = open_boutdataset(datapath=path, geometry='toroidal', gridfilepath=Path(path).parent.joinpath('grid.nc')) + # check dataset can be saved + save_dir = tmpdir_factory.mktemp('data') + actual.bout.save(str(save_dir.join('boutdata.nc'))) + def test_salpha(self, tmpdir_factory, bout_xyt_example_files): path = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=3, nt=1, syn_data_type='stepped', grid='grid') actual = open_boutdataset(datapath=path, geometry='s-alpha', gridfilepath=Path(path).parent.joinpath('grid.nc')) + # check dataset can be saved + save_dir = tmpdir_factory.mktemp('data') + actual.bout.save(str(save_dir.join('boutdata.nc'))) + @pytest.mark.skip def test_combine_along_tx(self): ... From 3fc4f02b89ab8a07aed39a7dfee808ba3782c329 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 26 Oct 2019 12:24:23 +0100 Subject: [PATCH 09/10] Do not save hthe in dataset for s-alpha geometry ...when hthe is loaded from the grid file. Variables from the grid file do not have all the correct metadata, so cause errors when trying to save to netCDF. --- xbout/geometries.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/xbout/geometries.py b/xbout/geometries.py index e8c395df..8ee6ce61 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -157,11 +157,14 @@ def add_s_alpha_geometry_coords(ds, coordinates=None): # Add 'hthe' from grid file, needed below for radial coordinate if 'hthe' not in ds: + hthe_from_grid = True if ds.bout._grid is None: raise ValueError("Grid file is required to provide %s. Pass the grid " "file name as the 'gridfilepath' argument to " "open_boutdataset().") ds['hthe'] = ds.bout._grid['hthe'] + else: + hthe_from_grid = False ds = add_toroidal_geometry_coords(ds, coordinates=coordinates) @@ -174,4 +177,8 @@ def add_s_alpha_geometry_coords(ds, coordinates=None): ds = ds.set_coords('r') ds = ds.rename(x='r') + if hthe_from_grid: + # remove hthe because it does not have correct metadata + del ds['hthe'] + return ds From 75f3072f5b93c191b5d99d963d63e1ff3e77fe02 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 26 Oct 2019 21:40:04 +0100 Subject: [PATCH 10/10] Remove workaround for mayavi dependency of boutdata Depedency has been made optional by an update to boututils. --- .travis.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.travis.yml b/.travis.yml index ba5fefca..9e244f42 100644 --- a/.travis.yml +++ b/.travis.yml @@ -3,7 +3,6 @@ python: - "3.6" - "3.7" install: - - pip install --upgrade vtk # fixes install of mayavi for Python-3.6, which is a dependency of boutdata - pip install --upgrade setuptools pip pytest pytest-cov coverage codecov boutdata "xarray!=0.14.0" - pip install -r requirements.txt - pip install -e .