diff --git a/.travis.yml b/.travis.yml index 7fffb52e..9e244f42 100644 --- a/.travis.yml +++ b/.travis.yml @@ -3,8 +3,7 @@ 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 + - pip install --upgrade setuptools pip pytest pytest-cov coverage codecov boutdata "xarray!=0.14.0" - pip install -r requirements.txt - pip install -e . script: 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 cf94060b..8ee6ce61 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): @@ -106,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.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.bout._grid[v] + # Change names of dimensions to Orthogonal Toroidal ones ds = ds.rename(y=coordinates['y']) @@ -143,6 +155,17 @@ 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 '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) # Add 1D radial coordinate @@ -154,7 +177,8 @@ 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) + if hthe_from_grid: + # remove hthe because it does not have correct metadata + del ds['hthe'] return ds diff --git a/xbout/load.py b/xbout/load.py index d0cc4881..61a9b353 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,14 @@ def open_boutdataset(datapath='./BOUT.dmp.*.nc', inputfilepath=None, if geometry: if info: print("Applying {} geometry conventions".format(geometry)) + + if gridfilepath is not None: + ds.bout._grid = _open_grid(gridfilepath, chunks=chunks, + keep_xboundaries=keep_xboundaries, + keep_yboundaries=keep_yboundaries) + else: + ds.bout._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 40709f67..cb0f40fc 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() @@ -169,7 +168,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 +184,12 @@ 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) + 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]))) @@ -342,6 +348,22 @@ 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. METADATA_VARS = ['BOUT_VERSION', 'NXPE', 'NYPE', 'NZPE', 'MXG', 'MYG', 'nx', 'ny', 'nz', @@ -364,7 +386,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 +441,26 @@ 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', grid='grid') + 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): ...