Merge pull request #58 from boutproject/test-geometries

Test geometries

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):

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

xbout/load.py

@@ -38,7 +38,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):
     """
@@ -67,6 +67,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
@@ -125,6 +128,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:

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):
         ...