diff --git a/xmitgcm/mds_store.py b/xmitgcm/mds_store.py index 7534ad32..cc4ca654 100644 --- a/xmitgcm/mds_store.py +++ b/xmitgcm/mds_store.py @@ -503,6 +503,7 @@ def __init__(self, data_dir, grid_dir=None, self.grid_dir, self.layers) self._all_data_variables = _get_all_data_variables(self.data_dir, + self.grid_dir, self.layers) # The rest of the data has to be read from disk. @@ -815,17 +816,25 @@ def _recursively_replace(item, search, replace): return item -def _get_all_data_variables(data_dir, layers): +def _get_all_data_variables(data_dir, grid_dir, layers): """"Put all the relevant data metadata into one big dictionary.""" allvars = [state_variables] allvars.append(package_state_variables) + # add others from available_diagnostics.log - fname = os.path.join(data_dir, 'available_diagnostics.log') - if os.path.exists(fname): - diag_file = fname + # search in the data dir + fnameD = os.path.join(data_dir, 'available_diagnostics.log') + # and in the grid dir + fnameG = os.path.join(grid_dir, 'available_diagnostics.log') + # first look in the data dir + if os.path.exists(fnameD): + diag_file = fnameD + # then in the grid dir + elif os.path.exists(fnameG): + diag_file = fnameG else: warnings.warn("Couldn't find available_diagnostics.log " - "in %s. Using default version." % data_dir) + "in %s or %s. Using default version." % (data_dir, grid_dir)) from .default_diagnostics import diagnostics diag_file = StringIO(diagnostics) available_diags = parse_available_diagnostics(diag_file, layers) diff --git a/xmitgcm/test/test_mds_store.py b/xmitgcm/test/test_mds_store.py index 1bcfdfca..9ede4865 100644 --- a/xmitgcm/test/test_mds_store.py +++ b/xmitgcm/test/test_mds_store.py @@ -403,6 +403,26 @@ def test_default_diagnostics(mds_datadirs_with_diagnostics): mate = ds[diagname].attrs['mate'] assert ds[mate].attrs['mate'] == diagname +def test_avail_diags_in_grid_dir(mds_datadirs_with_diagnostics): + """Try reading dataset with diagnostics output.""" + dirname, expected = mds_datadirs_with_diagnostics + + diag_prefix, expected_diags = expected['diagnostics'] + iters = expected['test_iternum'] + + with hide_file(dirname, + *['XC.meta', 'XC.data', 'RC.meta', 'RC.data', + 'available_diagnostics.log']) as grid_dir: + ds = xmitgcm.open_mdsdataset( + dirname, grid_dir=grid_dir, iters=iters, prefix=[diag_prefix], + read_grid=False, geometry=expected['geometry']) + + for diagname in expected_diags: + assert diagname in ds + if 'mate' in ds[diagname].attrs: + mate = ds[diagname].attrs['mate'] + assert ds[mate].attrs['mate'] == diagname + def test_layers_diagnostics(layers_mds_datadirs): """Try reading dataset with layers output.""" dirname, expected = layers_mds_datadirs