Add backwards compatible collect function

xbout/__init__.py

@@ -1,4 +1,4 @@
-from .load import open_boutdataset
+from .load import open_boutdataset, collect
 
 from . import geometries
 from .geometries import register_geometry, REGISTERED_GEOMETRIES

xbout/load.py

@@ -156,6 +156,98 @@ def open_boutdataset(datapath='./BOUT.dmp.*.nc', inputfilepath=None,
     return ds
 
 
+def collect(varname, xind=None, yind=None, zind=None, tind=None,
+            path=".", yguards=False, xguards=True, info=True, prefix="BOUT.dmp"):
+
+    from os.path import join
+
+    """
+
+    Extract the data pertaining to a specified variable in a BOUT++ data set
+
+
+    Parameters
+    ----------
+    varname : str
+        Name of the variable
+    xind, yind, zind, tind : int, slice or list of int, optional
+        Range of X, Y, Z or time indices to collect. Either a single
+        index to collect, a list containing [start, end] (inclusive
+        end), or a slice object (usual python indexing). Default is to
+        fetch all indices
+    path : str, optional
+        Path to data files (default: ".")
+    prefix : str, optional
+        File prefix (default: "BOUT.dmp")
+    yguards : bool, optional
+        Collect Y boundary guard cells? (default: False)
+    xguards : bool, optional
+        Collect X boundary guard cells? (default: True)
+        (Set to True to be consistent with the definition of nx)
+    info : bool, optional
+        Print information about collect? (default: True)
+
+    Notes
+    ----------
+    strict : This option found in boutdata.collect() is not present in this function
+             it is assumed that the varname given is correct, if variable does not exist
+             the function will fail
+    tind_auto : This option is not required when using _auto_open_mfboutdataset as an
+             automatic failure if datasets are different lengths is included
+
+    Returns
+    ----------
+    ds : numpy.ndarray
+
+    """
+
+    datapath = join(path, prefix + "*.nc")
+
+    ds = _auto_open_mfboutdataset(datapath, keep_xboundaries=xguards,
+                                  keep_yboundaries=yguards, info=info)
+
+    if varname not in ds:
+        raise KeyError("No variable, {} was found in {}.".format(varname, datapath))
+
+    dims = list(ds.dims)
+    inds = [tind, xind, yind, zind]
+
+    selection = {}
+
+    # Convert indexing values to an isel suitable format
+    for dim, ind in zip(dims, inds):
+
+        if isinstance(ind, int):
+            indexer = [ind]
+        elif isinstance(ind, list):
+            start, end = ind
+            indexer = slice(start, end+1)
+        elif ind is not None:
+            indexer = ind
+        else:
+            indexer = None
+
+        if indexer:
+            selection[dim] = indexer
+
+    try:
+        version = ds['BOUT_VERSION']
+    except KeyError:
+        # If BOUT Version is not saved in the dataset
+        version = 0
+
+    # Subtraction of z-dimensional data occurs in boutdata.collect
+    # if BOUT++ version is old - same feature added here
+    if (version < 3.5) and ('z' in dims):
+        zsize = int(ds['nz']) - 1
+        ds = ds.isel(z=slice(zsize))
+
+    if selection:
+        ds = ds.isel(selection)
+
+    return ds[varname].values
+
+
 def _is_dump_files(datapath):
     """
     If there is only one file, and it's not got a time dimension, assume it's a

xbout/tests/test_against_collect.py

@@ -1,14 +1,15 @@
 import pytest
-
 import numpy.testing as npt
 
-from xbout.load import open_boutdataset
+from xbout import open_boutdataset, collect as new_collect
+
 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
+old_collect = boutdata.collect
 
 class TestAccuracyAgainstOldCollect:
+
     def test_single_file(self, tmpdir_factory):
 
         # Create temp directory for files
@@ -19,14 +20,22 @@ def test_single_file(self, tmpdir_factory):
         generated_ds.to_netcdf(str(test_dir.join("BOUT.dmp.0.nc")))
 
         var = 'n'
-        expected = collect(var, path=test_dir, xguards=True, yguards=False)
+        expected = old_collect(var, path=test_dir, xguards=True, yguards=False)
 
+        # Test against new standard - open_boutdataset
         ds = open_boutdataset(test_dir.join("BOUT.dmp.0.nc"))
         actual = ds[var].values
 
         assert expected.shape == actual.shape
         npt.assert_equal(actual, expected)
 
+        # Test against backwards compatible collect function
+        actual = new_collect(var, path=test_dir, xguards=True, yguards=False)
+
+        assert expected.shape == actual.shape
+        npt.assert_equal(actual, expected)
+
+
     def test_multiple_files_along_x(self, tmpdir_factory):
 
         # Create temp directory for files
@@ -39,15 +48,22 @@ def test_multiple_files_along_x(self, tmpdir_factory):
             temp_ds.to_netcdf(str(test_dir.join(str(file_name))))
 
         var = 'n'
-        expected = collect(var, path=test_dir,
+        expected = old_collect(var, path=test_dir,
                            prefix='BOUT.dmp', xguards=True)
 
+        # Test against new standard - open_boutdataset
         ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
         actual = ds[var].values
 
         assert expected.shape == actual.shape
         npt.assert_equal(actual, expected)
 
+        # Test against backwards compatible collect function
+        actual = new_collect(var, path=test_dir, prefix='BOUT.dmp', xguards=True)
+
+        assert expected.shape == actual.shape
+        npt.assert_equal(actual, expected)
+
     def test_multiple_files_along_y(self, tmpdir_factory):
 
         # Create temp directory for files
@@ -60,15 +76,23 @@ def test_multiple_files_along_y(self, tmpdir_factory):
             temp_ds.to_netcdf(str(test_dir.join(str(file_name))))
 
         var = 'n'
-        expected = collect(var, path=test_dir,
+        expected = old_collect(var, path=test_dir,
                            prefix='BOUT.dmp', xguards=True)
 
+        # Test against new standard - .open_boutdataset
         ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
         actual = ds[var].values
 
         assert expected.shape == actual.shape
         npt.assert_equal(actual, expected)
 
+        # Test against backwards compatible collect function
+        actual = new_collect(var, path=test_dir, prefix='BOUT.dmp', xguards=True)
+
+        assert expected.shape == actual.shape
+        npt.assert_equal(actual, expected)
+
+
     def test_multiple_files_along_xy(self, tmpdir_factory):
 
         # Create temp directory for files
@@ -81,15 +105,23 @@ def test_multiple_files_along_xy(self, tmpdir_factory):
             temp_ds.to_netcdf(str(test_dir.join(str(file_name))))
 
         var = 'n'
-        expected = collect(var, path=test_dir,
+        expected = old_collect(var, path=test_dir,
                            prefix='BOUT.dmp', xguards=True)
 
+        # Test against new standard - .open_boutdataset
         ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
         actual = ds[var].values
 
         assert expected.shape == actual.shape
         npt.assert_equal(actual, expected)
 
+        # Test against backwards compatible collect function
+        actual = new_collect(var, path=test_dir, prefix='BOUT.dmp', xguards=True)
+
+        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")
@@ -101,10 +133,89 @@ def test_metadata(self, tmpdir_factory):
         ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
 
         for v in METADATA_VARS:
-            expected = collect(v, path=test_dir)
+            expected = old_collect(v, path=test_dir)
+            # Check metadata against new standard - open_boutdataset
             actual = ds.bout.metadata[v]
             npt.assert_equal(actual, expected)
 
+            # Check against backwards compatible collect function
+            actual = new_collect(v, path=test_dir)
+            npt.assert_equal(actual, expected)
+
+    def test_new_collect_indexing_int(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'
+        indexers = ["tind", "xind", "yind", "zind"]
+        ind_arg = 0
+
+        for kwarg in indexers:
+            # Extracting a the first index of each dimension for comparison
+            expected = old_collect(var, path=test_dir, **{kwarg: ind_arg})
+
+            # Test against backwards compatible collect function
+            actual = new_collect(var, path=test_dir, **{kwarg: ind_arg})
+
+            assert expected.shape == actual.shape
+            npt.assert_equal(actual, expected)
+
+    def test_new_collect_indexing_list(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'
+        indexers = ["tind", "xind", "yind", "zind"]
+        ind_arg = [0, 4]
+
+        for kwarg in indexers:
+            # Extracting a the first index of each dimension for comparison
+            expected = old_collect(var, path=test_dir, **{kwarg: ind_arg})
+
+            # Test against backwards compatible collect function
+            actual = new_collect(var, path=test_dir, **{kwarg: ind_arg})
+
+            assert expected.shape == actual.shape
+            npt.assert_equal(actual, expected)
+
+    def test_new_collect_indexing_slice(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'
+        indexers = ["tind", "xind", "yind", "zind"]
+        ind_list = [slice(0, 4, 2), slice(0, 4)]
+
+        for kwarg in indexers:
+            for ind_arg in ind_list:
+                # Extracting a the first index of each dimension for comparison
+                expected = old_collect(var, path=test_dir, **{kwarg: ind_arg})
+
+                # Test against backwards compatible collect function
+                actual = new_collect(var, path=test_dir, **{kwarg: ind_arg})
+
+                assert expected.shape == actual.shape
+                npt.assert_equal(actual, expected)
 
 @pytest.mark.skip
 class test_speed_against_old_collect: