Fix test against collect

.travis.yml

@@ -3,7 +3,7 @@ python:
   - "3.6"
   - "3.7"
 install:
-  - pip install --upgrade setuptools pip pytest pytest-cov coverage codecov
+  - pip install --upgrade setuptools pip pytest pytest-cov coverage codecov boutdata "xarray!=0.14.0"
   - pip install -r requirements.txt
   - pip install -e .
 script:

pytest.ini

@@ -0,0 +1,4 @@
+[pytest]
+filterwarnings =
+    ignore:No geometry type found, no coordinates will be added:UserWarning
+    ignore:deallocating CachingFileManager.*, but file is not already closed. This may indicate a bug\.:RuntimeWarning

xbout/boutdataarray.py

@@ -20,7 +20,6 @@ def __init__(self, da):
         self.data = da
         self.metadata = da.attrs.get('metadata')  # None if just grid file
         self.options = da.attrs.get('options')  # None if no inp file
-        self.grid = da.attrs.get('grid')  # None if no grid file
 
     def __str__(self):
         """
@@ -35,8 +34,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 animate2D(self, animate_over='t', x='x', y='y', animate=True,

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

@@ -16,7 +16,7 @@ class UnregisteredGeometryError(Exception):
 # desired
 
 
-def apply_geometry(ds, geometry_name):
+def apply_geometry(ds, geometry_name, *, coordinates=None, grid=None):
     """
 
     Parameters
@@ -25,6 +25,15 @@ def apply_geometry(ds, geometry_name):
         Dataset (from
     geometry_name : str
         Name under which the desired geometry function was registered
+    coordinates : dict of str, optional
+        Names to give the physical coordinates corresponding to 'x', 'y' and 'z'; values
+        corresponding to 'x', 'y' and 'z' keys in the passed dict are used as the names
+        of the dimensions. Any not passed are given default values. If not specified,
+        default names are chosen.
+    grid : Dataset, optional
+        Dataset containing extra geometrical information not stored in the dump files
+        that is needed to add coordinates for the geometry being applied. For example,
+        should contain 2d arrays Rxy, Zxy and psixy for toroidal geometry.
 
     Returns
     -------
@@ -43,7 +52,16 @@ def apply_geometry(ds, geometry_name):
                          have been registered.""".format(geometry_name))
         raise UnregisteredGeometryError(message)
 
-    updated_ds = add_geometry_coords(ds)
+    # User-registered functions may accept 'coordinates' and 'grid' arguments, but do not
+    # have to as long as they are not used
+    if coordinates is not None and grid is not None:
+        updated_ds = add_geometry_coords(ds, coordinates=coordinates, grid=grid)
+    elif coordinates is not None:
+        updated_ds = add_geometry_coords(ds, coordinates=coordinates)
+    elif grid is not None:
+        updated_ds = add_geometry_coords(ds, grid=grid)
+    else:
+        updated_ds = add_geometry_coords(ds)
     return updated_ds
 
 
@@ -92,7 +110,7 @@ def _set_default_toroidal_coordinates(coordinates):
 
 
 @register_geometry('toroidal')
-def add_toroidal_geometry_coords(ds, coordinates=None):
+def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
 
     coordinates = _set_default_toroidal_coordinates(coordinates)
 
@@ -104,6 +122,16 @@ def add_toroidal_geometry_coords(ds, coordinates=None):
                          "It may be useful to use the 'coordinates' argument to "
                          "add_toroidal_geometry_coords() for this.".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 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] = grid[v]
+
     # Change names of dimensions to Orthogonal Toroidal ones
     ds = ds.rename(y=coordinates['y'])
 
@@ -137,11 +165,22 @@ def add_toroidal_geometry_coords(ds, coordinates=None):
 
 
 @register_geometry('s-alpha')
-def add_s_alpha_geometry_coords(ds, coordinates=None):
+def add_s_alpha_geometry_coords(ds, *, coordinates=None, grid=None):
 
     coordinates = _set_default_toroidal_coordinates(coordinates)
 
-    ds = add_toroidal_geometry_coords(ds, coordinates=coordinates)
+    # Add 'hthe' from grid file, needed below for radial coordinate
+    if 'hthe' not in ds:
+        hthe_from_grid = True
+        if 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'] = grid['hthe']
+    else:
+        hthe_from_grid = False
+
+    ds = add_toroidal_geometry_coords(ds, coordinates=coordinates, grid=grid)
 
     # Add 1D radial coordinate
     if 'r' in ds:
@@ -152,7 +191,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,8 +128,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:
+            grid = _open_grid(gridfilepath, chunks=chunks,
+                              keep_xboundaries=keep_xboundaries,
+                              keep_yboundaries=keep_yboundaries,
+                              mxg=ds.metadata['MXG'])
+        else:
+            grid = None
+
         # Update coordinates to match particular geometry of grid
-        ds = geometries.apply_geometry(ds, geometry)
+        ds = geometries.apply_geometry(ds, geometry, grid=grid)
     else:
         if info:
             warn("No geometry type found, no coordinates will be added")
@@ -199,7 +211,7 @@ def _expand_filepaths(datapath):
              " `file_cache_maxsize` global option to {} to accommodate this. "
              "Recommend using `xr.set_options(file_cache_maxsize=NUM)`"
              " to explicitly set this to a large enough value."
-             .format(str(len(filepaths))), UserWarning)
+             .format(str(len(filepaths))))
         xr.set_options(file_cache_maxsize=len(filepaths))
 
     return filepaths, filetype
@@ -393,10 +405,13 @@ def _get_limit(side, dim, keep_boundaries, boundaries, guards):
     else:
         limit = None
 
+    if limit == 0:
+        # 0 would give incorrect result as an upper limit
+        limit = None
     return limit
 
 
-def _open_grid(datapath, chunks, keep_xboundaries, keep_yboundaries):
+def _open_grid(datapath, chunks, keep_xboundaries, keep_yboundaries, mxg=2):
     """
     Opens a single grid file. Implements slightly different logic for
     boundaries to deal with different conventions in a BOUT grid file.
@@ -415,12 +430,11 @@ def _open_grid(datapath, chunks, keep_xboundaries, keep_yboundaries):
         # pytest warnings capture - doesn't match strings containing brackets
         warn(
             "Will drop all variables containing the dimensions {} because "
-            "they are not recognised".format(str(unrecognised_dims)[1:-1]),
-            UserWarning)
+            "they are not recognised".format(str(unrecognised_dims)[1:-1]))
         grid = grid.drop_dims(unrecognised_dims)
 
     if not keep_xboundaries:
-        xboundaries = int(grid.metadata['MXG'])
+        xboundaries = mxg
         if xboundaries > 0:
             grid = grid.isel(x=slice(xboundaries, -xboundaries, None))
     if not keep_yboundaries:

xbout/tests/test_against_collect.py

@@ -2,45 +2,108 @@
 
 import numpy.testing as npt
 
-from xbout.load import _auto_open_mfboutdataset
+from xbout.load import open_boutdataset
+from .test_load import create_bout_ds, create_bout_ds_list, METADATA_VARS
 
+boutdata = pytest.importorskip("boutdata", reason="boutdata is not available")
+collect = boutdata.collect
 
 class TestAccuracyAgainstOldCollect:
-    @pytest.mark.skip
-    def test_single_file(self):
-        from boutdata import collect
+    def test_single_file(self, tmpdir_factory):
+
+        # Create temp directory for files
+        test_dir = tmpdir_factory.mktemp("test_data")
+
+        # Generate some test data
+        generated_ds = create_bout_ds(syn_data_type="linear")
+        generated_ds.to_netcdf(str(test_dir.join("BOUT.dmp.0.nc")))
+
         var = 'n'
-        expected = collect(var, path='./tests/data/dump_files/single',
-                           prefix='equilibrium', xguards=False)
+        expected = collect(var, path=test_dir, xguards=True, yguards=False)
 
-        ds, metadata = _auto_open_mfboutdataset('./tests/data/dump_files/single/equilibrium.nc')
-        print(ds)
+        ds = open_boutdataset(test_dir.join("BOUT.dmp.0.nc"))
         actual = ds[var].values
 
         assert expected.shape == actual.shape
         npt.assert_equal(actual, expected)
 
-    @pytest.mark.skip
-    def test_multiple_files_along_x(self):
-        from boutdata import collect
+    def test_multiple_files_along_x(self, tmpdir_factory):
+
+        # Create temp directory for files
+        test_dir = tmpdir_factory.mktemp("test_data")
+
+        # Generate some test data
+        ds_list, file_list = create_bout_ds_list("BOUT.dmp", nxpe=3, nype=1,
+                                                 syn_data_type="linear")
+        for temp_ds, file_name in zip(ds_list, file_list):
+            temp_ds.to_netcdf(str(test_dir.join(str(file_name))))
+
         var = 'n'
-        expected = collect(var, path='./tests/data/dump_files/',
-                           prefix='BOUT.dmp', xguards=False)
+        expected = collect(var, path=test_dir,
+                           prefix='BOUT.dmp', xguards=True)
 
-        ds, metadata = _auto_open_mfboutdataset('./tests/data/dump_files/BOUT.dmp.*.nc')
+        ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
         actual = ds[var].values
 
         assert expected.shape == actual.shape
         npt.assert_equal(actual, expected)
 
+    def test_multiple_files_along_y(self, tmpdir_factory):
+
+        # Create temp directory for files
+        test_dir = tmpdir_factory.mktemp("test_data")
+
+        # Generate some test data
+        ds_list, file_list = create_bout_ds_list("BOUT.dmp", nxpe=1, nype=3,
+                                                 syn_data_type="linear")
+        for temp_ds, file_name in zip(ds_list, file_list):
+            temp_ds.to_netcdf(str(test_dir.join(str(file_name))))
+
+        var = 'n'
+        expected = collect(var, path=test_dir,
+                           prefix='BOUT.dmp', xguards=True)
+
+        ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
+        actual = ds[var].values
+
+        assert expected.shape == actual.shape
+        npt.assert_equal(actual, expected)
+
+    def test_multiple_files_along_xy(self, tmpdir_factory):
+
+        # Create temp directory for files
+        test_dir = tmpdir_factory.mktemp("test_data")
+
+        # Generate some test data
+        ds_list, file_list = create_bout_ds_list("BOUT.dmp", nxpe=3, nype=3,
+                                                 syn_data_type="linear")
+        for temp_ds, file_name in zip(ds_list, file_list):
+            temp_ds.to_netcdf(str(test_dir.join(str(file_name))))
+
+        var = 'n'
+        expected = collect(var, path=test_dir,
+                           prefix='BOUT.dmp', xguards=True)
+
+        ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
+        actual = ds[var].values
+
+        assert expected.shape == actual.shape
+        npt.assert_equal(actual, expected)
+
+    def test_metadata(self, tmpdir_factory):
+        # Create temp directory for files
+        test_dir = tmpdir_factory.mktemp("test_data")
+
+        # Generate some test data
+        generated_ds = create_bout_ds(syn_data_type="linear")
+        generated_ds.to_netcdf(str(test_dir.join("BOUT.dmp.0.nc")))
 
-    @pytest.mark.skip
-    def test_multiple_files_along_x(self):
-        ...
+        ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
 
-    @pytest.mark.skip
-    def test_metadata(self):
-        ...
+        for v in METADATA_VARS:
+            expected = collect(v, path=test_dir)
+            actual = ds.bout.metadata[v]
+            npt.assert_equal(actual, expected)
 
 
 @pytest.mark.skip

xbout/tests/test_animate.py

@@ -19,7 +19,7 @@ def create_test_file(tmpdir_factory):
     for ds, file_name in zip(ds_list, file_list):
         ds.to_netcdf(str(save_dir.join(str(file_name))))
 
-    ds = open_boutdataset(save_dir.join("BOUT.dmp.*.nc")) # Open test data
+    ds = open_boutdataset(save_dir.join("BOUT.dmp.*.nc"))  # Open test data
 
     return save_dir, ds
 

xbout/tests/test_grid.py

@@ -57,10 +57,6 @@ def test_open_grid_extra_dims(self, create_example_grid_file):
         assert_equal(result, example_grid)
         result.close()
 
-    @pytest.mark.skip
-    def test_open_grid_merge_ds(self):
-        ...
-
     def test_open_grid_apply_geometry(self, create_example_grid_file):
         @register_geometry(name="Schwarzschild")
         def add_schwarzschild_coords(ds, coordinates=None):