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/load.py b/xbout/load.py index 171b05e8..80d25760 100644 --- a/xbout/load.py +++ b/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 diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py index 58052f14..41d5c6b6 100644 --- a/xbout/tests/test_against_collect.py +++ b/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: