Skip to content

Commit ac7fb58

Browse files
ChilippChilipp
committed
Added support for multifile datasets in from_dict
1 parent cd6507f commit ac7fb58

3 files changed

Lines changed: 155 additions & 32 deletions

File tree

‎CHANGELOG.rst‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,12 @@ Added
1111
* Added ``requires_replot`` attribute for the ``Formatoption`` class. If this
1212
attribute is True and the formatoption is contained in an update, it is the
1313
same as calling ``Plotter.update(replot=True))``.
14+
* We added support for multifile datasets when saving a project.
15+
Multifile datasets are datasets that have been opened with, e.g.
16+
``psyplot.data.open_mfdataset`` or
17+
``psyplot.project.plot.<plotmethod>(..., mfmode=True)``. This however does
18+
not always work with datasets opened with ``xarray.open_mfdataset``. In these
19+
cases, you have to set the ``Dataset.psy._concat_dim`` attribute manually
1420

1521

1622
v1.0.0

‎psyplot/data.py‎

Lines changed: 83 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import division
22
import os
3+
import os.path as osp
34
from threading import Thread
45
from functools import partial
56
from glob import glob
@@ -467,11 +468,28 @@ def tmp_it():
467468
store_mod, store_cls = ds.psy.data_store
468469
if store_mod is not None:
469470
store = ds._file_obj
470-
# try several datasets
471-
for func in get_fname_funcs:
472-
fname = func(store)
473-
if fname is not None:
474-
break
471+
# try several engines
472+
if hasattr(store, 'file_objs'):
473+
fname = []
474+
store_mod = []
475+
store_cls = []
476+
for obj in store.file_objs: # mfdataset
477+
_fname = None
478+
for func in get_fname_funcs:
479+
if _fname is None:
480+
_fname = func(obj)
481+
if _fname is not None:
482+
fname.append(_fname)
483+
store_mod.append(obj.__module__)
484+
store_cls.append(obj.__class__.__name__)
485+
fname = tuple(fname)
486+
store_mod = tuple(store_mod)
487+
store_cls = tuple(store_cls)
488+
else:
489+
for func in get_fname_funcs:
490+
fname = func(store)
491+
if fname is not None:
492+
break
475493
# check if paths is provided and if yes, save the file
476494
if fname is None and paths is not None:
477495
fname = next(paths, None)
@@ -1722,8 +1740,8 @@ def open_dataset(filename_or_obj, decode_cf=True, decode_times=True,
17221740
xarray.Dataset
17231741
The dataset that contains the variables from `filename_or_obj`"""
17241742
# use the absolute path name (is saver when saving the project)
1725-
if isstring(filename_or_obj) and os.path.exists(filename_or_obj):
1726-
filename_or_obj = os.path.abspath(filename_or_obj)
1743+
if isstring(filename_or_obj) and osp.exists(filename_or_obj):
1744+
filename_or_obj = osp.abspath(filename_or_obj)
17271745
if engine == 'gdal':
17281746
from psyplot.gdal_store import GdalStore
17291747
filename_or_obj = GdalStore(filename_or_obj)
@@ -1795,9 +1813,10 @@ def open_mfdataset(paths, decode_cf=True, decode_times=True,
17951813
paths, decode_cf=decode_cf, decode_times=decode_times, engine=engine,
17961814
decode_coords=False, **kwargs)
17971815
if decode_cf:
1798-
return CFDecoder.decode_ds(ds, gridfile=gridfile, inplace=True,
1799-
decode_coords=decode_coords,
1800-
decode_times=decode_times)
1816+
ds = CFDecoder.decode_ds(ds, gridfile=gridfile, inplace=True,
1817+
decode_coords=decode_coords,
1818+
decode_times=decode_times)
1819+
ds.psy._concat_dim = kwargs.get('concat_dim')
18011820
return ds
18021821

18031822

@@ -2387,7 +2406,7 @@ def filter_attrs(item):
23872406
queues[0].task_done()
23882407
self._new_dims = {}
23892408
self.onupdate.emit()
2390-
except:
2409+
except Exception:
23912410
self._finish_all(queues)
23922411
raise
23932412
return InteractiveBase.start_update(self, draw=draw, queues=queues)
@@ -2873,15 +2892,20 @@ def sel_method(key, dims, name=None):
28732892
return instance
28742893

28752894
@classmethod
2876-
def _get_dsnames(cls, data, ignore_keys=['attrs', 'plotter', 'ds']):
2895+
def _get_dsnames(cls, data, ignore_keys=['attrs', 'plotter', 'ds'],
2896+
concat_dim=False):
28772897
"""Recursive method to get all the file names out of a dictionary
28782898
`data` created with the :meth`array_info` method"""
28792899
def filter_ignores(item):
28802900
return item[0] not in ignore_keys and isinstance(item[1], dict)
28812901
if 'fname' in data:
2882-
return {(data['fname'], data['store'])}
2883-
return set(chain(*map(cls._get_dsnames, dict(
2884-
filter(filter_ignores, six.iteritems(data))).values())))
2902+
return {tuple(
2903+
[data['fname'], data['store']] +
2904+
([data.get('concat_dim')] if concat_dim else []))}
2905+
return set(chain(*map(partial(cls._get_dsnames, concat_dim=concat_dim,
2906+
ignore_keys=ignore_keys),
2907+
dict(filter(filter_ignores,
2908+
six.iteritems(data))).values())))
28852909

28862910
@classmethod
28872911
def _get_ds_descriptions(
@@ -3022,29 +3046,38 @@ def only_filter(arr_name, info):
30223046
return arr_name in save_only
30233047
save_only = only
30243048
only = None
3049+
3050+
def get_fname_use(fname):
3051+
squeeze = isstring(fname)
3052+
fname = safe_list(fname)
3053+
ret = tuple(f if utils.is_remote_url(f) or osp.isabs(f) else
3054+
osp.join(pwd, f)
3055+
for f in fname)
3056+
return ret[0] if squeeze else ret
3057+
30253058
if not isinstance(alternative_paths, dict):
30263059
it = iter(alternative_paths)
30273060
alternative_paths = defaultdict(partial(next, it, None))
30283061
# first open all datasets if not already done
30293062
if datasets is None:
3030-
names_and_stores = cls._get_dsnames(d)
3063+
replace_concat_dim = 'concat_dim' not in kwargs
3064+
3065+
names_and_stores = cls._get_dsnames(d, concat_dim=True)
30313066
datasets = {}
3032-
for fname, (store_mod, store_cls) in names_and_stores:
3067+
for fname, (store_mod, store_cls), concat_dim in names_and_stores:
30333068
fname_use = fname
30343069
got = True
3070+
if replace_concat_dim and concat_dim is not None:
3071+
kwargs['concat_dim'] = concat_dim
3072+
elif replace_concat_dim and concat_dim is None:
3073+
kwargs.pop('concat_dim', None)
30353074
try:
30363075
fname_use = alternative_paths[fname]
30373076
except KeyError:
30383077
got = False
30393078
if not got or not fname_use:
30403079
if fname is not None:
3041-
if utils.is_remote_url(fname):
3042-
fname_use = fname
3043-
else:
3044-
if os.path.isabs(fname):
3045-
fname_use = fname
3046-
else:
3047-
fname_use = os.path.join(pwd, fname)
3080+
fname_use = get_fname_use(fname)
30483081
if fname_use is not None:
30493082
datasets[fname] = _open_ds_from_store(
30503083
fname_use, store_mod, store_cls, **kwargs)
@@ -3176,7 +3209,7 @@ def array_info(self, dump=None, paths=None, attrs=True,
31763209
saved_ds = kwargs.pop('_saved_ds', {})
31773210

31783211
def get_alternative(f):
3179-
return next(filter(lambda t: os.path.samefile(f, t[0]),
3212+
return next(filter(lambda t: osp.samefile(f, t[0]),
31803213
six.iteritems(alternative_paths)), [False, f])
31813214

31823215
if copy:
@@ -3240,15 +3273,17 @@ def copy_obj(obj):
32403273
else:
32413274
found, f = get_alternative(f)
32423275
if use_rel_paths:
3243-
f = os.path.relpath(f, pwd)
3276+
f = osp.relpath(f, pwd)
32443277
else:
3245-
f = os.path.abspath(f)
3278+
f = osp.abspath(f)
32463279
d['fname'].append(f)
32473280
if fname is None or isinstance(fname,
32483281
six.string_types):
32493282
d['fname'] = d['fname'][0]
32503283
else:
32513284
d['fname'] = tuple(safe_list(fname))
3285+
if arr.psy.base.psy._concat_dim is not None:
3286+
d['concat_dim'] = arr.psy.base.psy._concat_dim
32523287
if 'ds' in ds_description:
32533288
if full_ds:
32543289
d['ds'] = copy_obj(arr.psy.base)
@@ -3680,6 +3715,9 @@ class DatasetAccessor(object):
36803715
_num = None
36813716
_plot = None
36823717

3718+
#: The concatenation dimension for datasets opened with open_mfdataset
3719+
_concat_dim = None
3720+
36833721
@property
36843722
def num(self):
36853723
"""A unique number for the dataset"""
@@ -3721,7 +3759,7 @@ def filename(self):
37213759
"""The name of the file that stores this dataset"""
37223760
fname = self._filename
37233761
if fname is None:
3724-
fname, store_mod, store_cls = get_filename_ds(self.ds, dump=False)
3762+
fname = get_filename_ds(self.ds, dump=False)[0]
37253763
return fname
37263764

37273765
@filename.setter
@@ -3889,7 +3927,7 @@ def start_update(self, draw=None, queues=None):
38893927
for arr in self:
38903928
arr.psy.start_update(draw=False)
38913929
self.onupdate.emit()
3892-
except:
3930+
except Exception:
38933931
self._finish_all(queues)
38943932
raise
38953933
if queues is not None:
@@ -3987,6 +4025,22 @@ def _open_ds_from_store(fname, store_mod=None, store_cls=None, **kwargs):
39874025
"""Open a dataset and return it"""
39884026
if isinstance(fname, xr.Dataset):
39894027
return fname
4028+
if not isstring(fname):
4029+
try: # test iterable
4030+
fname[0]
4031+
except TypeError:
4032+
pass
4033+
else:
4034+
if store_mod is not None and store_cls is not None:
4035+
if isstring(store_mod):
4036+
store_mod = repeat(store_mod)
4037+
if isstring(store_cls):
4038+
store_cls = repeat(store_cls)
4039+
fname = [getattr(import_module(sm), sc)(f)
4040+
for sm, sc, f in zip(store_mod, store_cls, fname)]
4041+
kwargs['engine'] = None
4042+
kwargs['lock'] = False
4043+
return open_mfdataset(fname, **kwargs)
39904044
if store_mod is not None and store_cls is not None:
39914045
fname = getattr(import_module(store_mod), store_cls)(fname)
39924046
return open_dataset(fname, **kwargs)

‎tests/test_data.py‎

Lines changed: 66 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import _base_testing as bt
1111
import numpy as np
1212
from collections import OrderedDict
13+
import tempfile
1314

1415
try:
1516
import PyNio
@@ -529,6 +530,16 @@ def test_update_06_2variables(self):
529530
class TestArrayList(unittest.TestCase):
530531
"""Test the :class:`psyplot.data.ArrayList` class"""
531532

533+
_created_files = set()
534+
535+
def setUp(self):
536+
self._created_files = set()
537+
538+
def tearDown(self):
539+
for f in self._created_files:
540+
os.remove(f)
541+
self._created_files.clear()
542+
532543
list_class = psyd.ArrayList
533544

534545
def test_setup_coords(self):
@@ -642,7 +653,6 @@ def test_filter_6_ax(self):
642653
from psyplot.plotter import Plotter
643654
ds = self._filter_test_ds
644655
l = self.list_class.from_dataset(ds, ydim=[0, 1], name='v0')
645-
print(l)
646656
axes = plt.subplots(1, 2)[1]
647657
for i, arr in enumerate(l):
648658
Plotter(arr, ax=axes[i])
@@ -1025,6 +1035,48 @@ def test_from_dict_02_only(self):
10251035
).array_info(),
10261036
l[1:].array_info())
10271037

1038+
def test_from_dict_03_mfdataset(self):
1039+
"""Test opening a multifile dataset"""
1040+
ds = xr.Dataset(*self._from_dataset_test_variables)
1041+
ds1 = ds.isel(time=slice(0, 2))
1042+
ds2 = ds.isel(time=slice(2, None))
1043+
fname1 = tempfile.NamedTemporaryFile(suffix='.nc',
1044+
prefix='tmp_psyplot_').name
1045+
ds1.to_netcdf(fname1)
1046+
self._created_files.add(fname1)
1047+
fname2 = tempfile.NamedTemporaryFile(suffix='.nc',
1048+
prefix='tmp_psyplot_').name
1049+
ds2.to_netcdf(fname2)
1050+
self._created_files.add(fname2)
1051+
1052+
# now open the mfdataset
1053+
ds = psyd.open_mfdataset([fname1, fname2])
1054+
l = self.list_class.from_dataset(ds, name=['v0'], time=[0, 3])
1055+
self.assertEqual(
1056+
self.list_class.from_dict(l.array_info()).array_info(),
1057+
l.array_info())
1058+
1059+
def test_from_dict_04_concat_dim(self):
1060+
"""Test opening a multifile dataset that requires a ``concat_dim``"""
1061+
ds = xr.Dataset(*self._from_dataset_test_variables)
1062+
ds1 = ds.isel(time=0)
1063+
ds2 = ds.isel(time=1)
1064+
fname1 = tempfile.NamedTemporaryFile(suffix='.nc',
1065+
prefix='tmp_psyplot_').name
1066+
ds1.to_netcdf(fname1)
1067+
self._created_files.add(fname1)
1068+
fname2 = tempfile.NamedTemporaryFile(suffix='.nc',
1069+
prefix='tmp_psyplot_').name
1070+
ds2.to_netcdf(fname2)
1071+
self._created_files.add(fname2)
1072+
1073+
# now open the mfdataset
1074+
ds = psyd.open_mfdataset([fname1, fname2], concat_dim='time')
1075+
l = self.list_class.from_dataset(ds, name=['v0'], time=[0, 1])
1076+
self.assertEqual(
1077+
self.list_class.from_dict(l.array_info()).array_info(),
1078+
l.array_info())
1079+
10281080
def test_logger(self):
10291081
"""Test whether one can access the logger"""
10301082
import logging
@@ -1065,6 +1117,16 @@ def test_to_dataframe(self):
10651117
class AbsoluteTimeTest(unittest.TestCase, AlmostArrayEqualMixin):
10661118
"""TestCase for loading and storing absolute times"""
10671119

1120+
_created_files = set()
1121+
1122+
def setUp(self):
1123+
self._created_files = set()
1124+
1125+
def tearDown(self):
1126+
for f in self._created_files:
1127+
os.remove(f)
1128+
self._created_files.clear()
1129+
10681130
@property
10691131
def _test_ds(self):
10701132
import xarray as xr
@@ -1079,9 +1141,10 @@ def _test_ds(self):
10791141
def test_to_netcdf(self):
10801142
"""Test whether the data is stored correctly"""
10811143
import netCDF4 as nc
1082-
import tempfile
10831144
ds = self._test_ds
1084-
fname = tempfile.NamedTemporaryFile().name
1145+
fname = tempfile.NamedTemporaryFile(suffix='.nc',
1146+
prefix='tmp_psyplot_').name
1147+
self._created_files.add(fname)
10851148
psyd.to_netcdf(ds, fname)
10861149
with nc.Dataset(fname) as nco:
10871150
self.assertAlmostArrayEqual(

0 commit comments

Comments
 (0)