From 27220d6f5c44e99e1dc84e19e4d26118231413bc Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Mon, 28 Oct 2019 12:40:12 +0000 Subject: [PATCH 01/14] Added backwards compatible collect function to replicate the behaviour of boutdata.collect(). --- xbout/load.py | 62 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/xbout/load.py b/xbout/load.py index 8017a378..b490b4f4 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -144,6 +144,68 @@ 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) + + Returns + ---------- + ds : numpy.ndarray + + """ + + datapath = join(path, prefix + "*.nc") + + ds = _auto_open_mfboutdataset(datapath, keep_xboundaries=xguards, + keep_yboundaries=yguards, info=info) + + dims = ('t', 'x', 'y', 'z') + indexers = (tind, xind, yind, zind) + + selection = {} + + for i in range(len(dims)): + + if indexers[i] != None: + + if isinstance(indexers[i], int): + selection[dims[i]] = [indexers[i]] + else: + selection[dims[i]] = indexers[i] + + 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 From 00c2c840d662ec64c5f268b56220fb0460374989 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Mon, 28 Oct 2019 12:49:21 +0000 Subject: [PATCH 02/14] PEP8 fixes --- xbout/load.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/load.py b/xbout/load.py index b490b4f4..80c65af3 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -193,7 +193,7 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, for i in range(len(dims)): - if indexers[i] != None: + if indexers[i] not None: if isinstance(indexers[i], int): selection[dims[i]] = [indexers[i]] From d6ed549c9248869e2976994e43114b9f8f3f1c3d Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Mon, 28 Oct 2019 12:55:38 +0000 Subject: [PATCH 03/14] Further fixes --- xbout/load.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index 80c65af3..27b6e7b2 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -186,14 +186,14 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, ds = _auto_open_mfboutdataset(datapath, keep_xboundaries=xguards, keep_yboundaries=yguards, info=info) - dims = ('t', 'x', 'y', 'z') - indexers = (tind, xind, yind, zind) + dims = ['t', 'x', 'y', 'z'] + indexers = [tind, xind, yind, zind] selection = {} for i in range(len(dims)): - if indexers[i] not None: + if indexers[i] is not None: if isinstance(indexers[i], int): selection[dims[i]] = [indexers[i]] From 7c2b93720ee69a5ddd01007939d53754fb15ac89 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Mon, 28 Oct 2019 13:37:19 +0000 Subject: [PATCH 04/14] Simplify if statements and add list check --- xbout/load.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index 27b6e7b2..9c5ee293 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -187,18 +187,21 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, keep_yboundaries=yguards, info=info) dims = ['t', 'x', 'y', 'z'] - indexers = [tind, xind, yind, zind] + inds = [tind, xind, yind, zind] selection = {} - for i in range(len(dims)): + for dim, ind in zip(dims, inds): - if indexers[i] is not None: + if isinstance(ind, int): + indexer = [ind] + elif isinstance(ind, list): + start, end = ind + indexer = slice(start, end) + elif ind is not None: + indexer = ind - if isinstance(indexers[i], int): - selection[dims[i]] = [indexers[i]] - else: - selection[dims[i]] = indexers[i] + selection[dim] = indexer if selection: ds = ds.isel(selection) From d5ece6460c58c0e3298a403a4fbe1da59b088256 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Mon, 28 Oct 2019 14:11:47 +0000 Subject: [PATCH 05/14] Notes and variable check added --- xbout/load.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/xbout/load.py b/xbout/load.py index 9c5ee293..4b9f4b53 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -175,6 +175,14 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, 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 @@ -186,6 +194,9 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, 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 = ['t', 'x', 'y', 'z'] inds = [tind, xind, yind, zind] From 6d39a4799d1d9291dd6157dbcc3db6ca9d9e9ff0 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Tue, 29 Oct 2019 15:16:57 +0000 Subject: [PATCH 06/14] Holding commit many more edits to be made - z-dim indexing option updated and tests added (tests will fail if xguards and yguards parameters change, issue being looked at in #61) --- xbout/load.py | 20 +++++++++- xbout/tests/test_against_collect.py | 59 +++++++++++++++++++---------- 2 files changed, 57 insertions(+), 22 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index 4b9f4b53..9b43d588 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -197,11 +197,12 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, if varname not in ds: raise KeyError("No variable, {} was found in {}.".format(varname, datapath)) - dims = ['t', 'x', 'y', 'z'] + 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): @@ -211,8 +212,23 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, indexer = slice(start, end) elif ind is not None: indexer = ind + else: + indexer = None + + if indexer: + selection[dim] = 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 + selection['z'] = slice(zsize) if selection: ds = ds.isel(selection) diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py index 45a932e1..b5ab6eea 100644 --- a/xbout/tests/test_against_collect.py +++ b/xbout/tests/test_against_collect.py @@ -2,42 +2,61 @@ import numpy.testing as npt -from xbout.load import _auto_open_mfboutdataset +from xbout import open_boutdataset +from .test_load import create_bout_ds_list + +from xbout.load import collect as new_collect +from boutdata import collect as old_collect + + +@pytest.fixture +def create_test_file(tmpdir_factory): + def _foo(nxpe, nype): + # Create temp dir for test files + save_dir = tmpdir_factory.mktemp("test_data") + + # Generate test data + ds_list, file_list = create_bout_ds_list("data", nxpe=nxpe, nype=nype, + syn_data_type="linear") + + for ds, file_name in zip(ds_list, file_list): + ds.to_netcdf(str(save_dir.join(str(file_name)))) + + return save_dir + return _foo class TestAccuracyAgainstOldCollect: - @pytest.mark.skip - def test_single_file(self): - from boutdata import collect + # @pytest.mark.skip + def test_single_file(self, create_test_file): + + save_dir = create_test_file(nxpe=1,nype=1) + var = 'n' - expected = collect(var, path='./tests/data/dump_files/single', - prefix='equilibrium', xguards=False) + expected = old_collect(var, path=save_dir, prefix='data', xguards=False) - ds, metadata = _auto_open_mfboutdataset('./tests/data/dump_files/single/equilibrium.nc') - print(ds) - actual = ds[var].values + actual = new_collect(var, path=save_dir, prefix='data', xguards=False) assert expected.shape == actual.shape npt.assert_equal(actual, expected) - @pytest.mark.skip - def test_multiple_files_along_x(self): - from boutdata import collect + + # @pytest.mark.skip + def test_multiple_files_along_x(self, create_test_file): + + save_dir = create_test_file(nxpe=3, nype=3) + var = 'n' - expected = collect(var, path='./tests/data/dump_files/', - prefix='BOUT.dmp', xguards=False) + expected = old_collect(var, path=save_dir, prefix='data', xguards=False) - ds, metadata = _auto_open_mfboutdataset('./tests/data/dump_files/BOUT.dmp.*.nc') - actual = ds[var].values + actual = new_collect(var, path=save_dir, prefix='data', xguards=False) + + print(expected.shape, actual.shape) assert expected.shape == actual.shape npt.assert_equal(actual, expected) - @pytest.mark.skip - def test_multiple_files_along_x(self): - ... - @pytest.mark.skip def test_metadata(self): ... From 9457bbd211fce2b4ae825baebadb1740e78eb965 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Wed, 4 Dec 2019 10:36:01 +0000 Subject: [PATCH 07/14] PEP-8 Fixes --- xbout/load.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index 0ffa7ded..a3fc77ef 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -231,10 +231,10 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, selection[dim] = indexer try: - version = ds['BOUT_VERSION'] + version = ds['BOUT_VERSION'] except KeyError: - # If BOUT Version is not saved in the dataset - version = 0 + # 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 From b052d070d8ee8b391563266a47d84dd61b78336f Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Thu, 5 Dec 2019 09:39:13 +0000 Subject: [PATCH 08/14] Allow for collect to be imported from xbout --- xbout/__init__.py | 2 +- xbout/tests/test_against_collect.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/xbout/__init__.py b/xbout/__init__.py index dbdd4ebf..f7fed880 100644 --- a/xbout/__init__.py +++ b/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 diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py index 952f4d45..f04840fc 100644 --- a/xbout/tests/test_against_collect.py +++ b/xbout/tests/test_against_collect.py @@ -1,8 +1,7 @@ import pytest import numpy.testing as npt -from xbout import open_boutdataset -from xbout.load import collect as new_collect +from xbout import open_boutdataset, collect as new_collect from .test_load import create_bout_ds, create_bout_ds_list, METADATA_VARS From f2257b52701bfdaddf684fcfbbd4ce0732feddb3 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Thu, 5 Dec 2019 10:30:33 +0000 Subject: [PATCH 09/14] Fix indexing in collect function --- xbout/load.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index a3fc77ef..aa8ab9bc 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -172,7 +172,7 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, 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 + index to collect, a list containing [start, end, step (optional)] (inclusive end), or a slice object (usual python indexing). Default is to fetch all indices path : str, optional @@ -220,8 +220,12 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, if isinstance(ind, int): indexer = [ind] elif isinstance(ind, list): - start, end = ind - indexer = slice(start, end) + try: + start, end, step = ind + indexer = slice(start, end, step) + except ValueError: + start, end = ind + indexer = slice(start, end) elif ind is not None: indexer = ind else: From 6fb3a0737484023b34f7a39e57154da90b27c2d9 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Thu, 5 Dec 2019 10:31:28 +0000 Subject: [PATCH 10/14] Fix application of zsize --- xbout/load.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/load.py b/xbout/load.py index aa8ab9bc..f7df0dcb 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -244,7 +244,7 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, # if BOUT++ version is old - same feature added here if (version < 3.5) and ('z' in dims): zsize = int(ds['nz']) - 1 - selection['z'] = slice(zsize) + ds = ds.isel(z=slice(zsize)) if selection: ds = ds.isel(selection) From e3dc96940c492a26055b414bc62f80d714499317 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Thu, 5 Dec 2019 11:43:22 +0000 Subject: [PATCH 11/14] Added tests to checking indexing --- xbout/load.py | 4 +- xbout/tests/test_against_collect.py | 76 +++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index f7df0dcb..d45a9604 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -172,7 +172,7 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, 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, step (optional)] (inclusive + index to collect, a list containing [start, end, step] or [start, end] (inclusive end), or a slice object (usual python indexing). Default is to fetch all indices path : str, optional @@ -225,7 +225,7 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, indexer = slice(start, end, step) except ValueError: start, end = ind - indexer = slice(start, end) + indexer = slice(start, end+1) elif ind is not None: indexer = ind else: diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py index f04840fc..aff30447 100644 --- a/xbout/tests/test_against_collect.py +++ b/xbout/tests/test_against_collect.py @@ -142,6 +142,82 @@ def test_metadata(self, tmpdir_factory): 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_list = [[0, 4, 2], [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) + + 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: ... From c80f0dc21816d0158e307434c28f00eb601165f8 Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Thu, 5 Dec 2019 11:50:25 +0000 Subject: [PATCH 12/14] PEP8 fixes --- xbout/tests/test_against_collect.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py index aff30447..730a9adc 100644 --- a/xbout/tests/test_against_collect.py +++ b/xbout/tests/test_against_collect.py @@ -149,7 +149,8 @@ def test_new_collect_indexing_int(self, tmpdir_factory): # Generate some test data ds_list, file_list = create_bout_ds_list("BOUT.dmp", nxpe=3, nype=3, - syn_data_type="linear") + 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)))) @@ -159,10 +160,10 @@ def test_new_collect_indexing_int(self, tmpdir_factory): for kwarg in indexers: # Extracting a the first index of each dimension for comparison - expected = old_collect(var, path=test_dir, **{kwarg:ind_arg}) + 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}) + actual = new_collect(var, path=test_dir, **{kwarg: ind_arg}) assert expected.shape == actual.shape npt.assert_equal(actual, expected) @@ -174,7 +175,7 @@ def test_new_collect_indexing_list(self, tmpdir_factory): # Generate some test data ds_list, file_list = create_bout_ds_list("BOUT.dmp", nxpe=3, nype=3, - syn_data_type="linear") + 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)))) @@ -185,10 +186,10 @@ def test_new_collect_indexing_list(self, tmpdir_factory): 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}) + 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}) + actual = new_collect(var, path=test_dir, **{kwarg: ind_arg}) assert expected.shape == actual.shape npt.assert_equal(actual, expected) @@ -199,7 +200,8 @@ def test_new_collect_indexing_slice(self, tmpdir_factory): # Generate some test data ds_list, file_list = create_bout_ds_list("BOUT.dmp", nxpe=3, nype=3, - syn_data_type="linear") + 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)))) @@ -210,10 +212,10 @@ def test_new_collect_indexing_slice(self, tmpdir_factory): 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}) + 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}) + actual = new_collect(var, path=test_dir, **{kwarg: ind_arg}) assert expected.shape == actual.shape npt.assert_equal(actual, expected) From 87f5a1b37962baadfac7fc4111f7f7dbec70fb5c Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Thu, 5 Dec 2019 11:53:10 +0000 Subject: [PATCH 13/14] More PEP8 fixes --- xbout/tests/test_against_collect.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py index 730a9adc..c2f4efc6 100644 --- a/xbout/tests/test_against_collect.py +++ b/xbout/tests/test_against_collect.py @@ -142,7 +142,6 @@ def test_metadata(self, tmpdir_factory): 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") @@ -168,7 +167,6 @@ def test_new_collect_indexing_int(self, tmpdir_factory): 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") @@ -207,7 +205,7 @@ def test_new_collect_indexing_slice(self, tmpdir_factory): var = 'n' indexers = ["tind", "xind", "yind", "zind"] - ind_list = [slice(0,4,2), slice(0,4)] + ind_list = [slice(0, 4, 2), slice(0, 4)] for kwarg in indexers: for ind_arg in ind_list: From 18db7f02ffc2d6157197088f17ed6b6625418c9f Mon Sep 17 00:00:00 2001 From: Rhys Doyle Date: Thu, 5 Dec 2019 16:10:51 +0000 Subject: [PATCH 14/14] Removed 3-element list form --- xbout/load.py | 10 +++------- xbout/tests/test_against_collect.py | 15 +++++++-------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index d45a9604..80d25760 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -172,7 +172,7 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, 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, step] or [start, end] (inclusive + 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 @@ -220,12 +220,8 @@ def collect(varname, xind=None, yind=None, zind=None, tind=None, if isinstance(ind, int): indexer = [ind] elif isinstance(ind, list): - try: - start, end, step = ind - indexer = slice(start, end, step) - except ValueError: - start, end = ind - indexer = slice(start, end+1) + start, end = ind + indexer = slice(start, end+1) elif ind is not None: indexer = ind else: diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py index c2f4efc6..41d5c6b6 100644 --- a/xbout/tests/test_against_collect.py +++ b/xbout/tests/test_against_collect.py @@ -179,18 +179,17 @@ def test_new_collect_indexing_list(self, tmpdir_factory): var = 'n' indexers = ["tind", "xind", "yind", "zind"] - ind_list = [[0, 4, 2], [0, 4]] + ind_arg = [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}) + # 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}) + # 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) + assert expected.shape == actual.shape + npt.assert_equal(actual, expected) def test_new_collect_indexing_slice(self, tmpdir_factory): # Create temp directory for files