diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 279fdbd..699ce9f 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -952,7 +952,12 @@ def _if_not_none(a, b): data['RF'][sl], _VAR_METADATA[zv]['attrs']) - ds = ds.update(variables) + # Save vars as dataArray instead of Variables, avoid dim mismatch error + vars_da = { + name: xr.DataArray(v.data, dims=v.dims, attrs=v.attrs) + for name, v in variables.items() + } + ds = ds.assign(**vars_da) if grid_vars_to_coords: ds = ds.set_coords(gridlist) diff --git a/xmitgcm/test/test_utils.py b/xmitgcm/test/test_utils.py index 35183f3..11f5a1e 100644 --- a/xmitgcm/test/test_utils.py +++ b/xmitgcm/test/test_utils.py @@ -20,6 +20,20 @@ nrecords = [ 1 ]; """ +_pk_meta_content = """ simulation = { 'global_ocean.32x15' }; + nDims = [ 2 ]; + dimList = [ + 192, 1, 192, + 32, 1, 32 + ]; + dataprec = [ 'float64' ]; + nrecords = [ 123 ]; + timeStepNumber = [ 72000 ]; + nFlds = [ 11 ]; + fldList = { + 'Uvel ' 'GuNm1 ' 'Vvel ' 'GvNm1 ' 'Theta ' 'GtNm1 ' 'Salt ' 'GsNm1 ' 'EtaN ' 'dEtaHdt ' 'EtaH ' + }; +""" def test_parse_meta(tmpdir): """Check the parsing of MITgcm .meta into python dictionary.""" @@ -41,6 +55,38 @@ def test_parse_meta(tmpdir): assert result[k] == v +def test_parse_pickup_meta(tmpdir): + """Check the parsing of MITgcm pickup .meta into python dictionary.""" + + from xmitgcm.utils import parse_meta_file + p = tmpdir.join("pickup.0000000000.meta") + p.write(_pk_meta_content) + fname = str(p) + result = parse_meta_file(fname) + expected = { + 'basename': 'pickup', + 'nDims': 2, + 'dimList': [[192, 1, 192], [32, 1, 32]], + 'dataprec': np.dtype('float64'), + 'nrecords': 123, + 'timeStepNumber': '72000', + 'nFlds': 11, + 'fldList': ['Uvel', + 'GuNm1', + 'Vvel', + 'GvNm1', + 'Theta', + 'GtNm1', + 'Salt', + 'GsNm1', + 'EtaN', + 'dEtaHdt', + 'EtaH'] + } + for k, v in expected.items(): + assert result[k] == v + + @pytest.mark.parametrize("dtype", [np.dtype('f8'), np.dtype('f4'), np.dtype('i4')]) def test_read_raw_data(tmpdir, dtype): """Check our utility for reading raw data.""" diff --git a/xmitgcm/utils.py b/xmitgcm/utils.py index d05146c..c4ee541 100644 --- a/xmitgcm/utils.py +++ b/xmitgcm/utils.py @@ -60,7 +60,11 @@ def parse_meta_file(fname): if 'fldList' in flds: flds['fldList'] = [re.match(r"'*(\w+)", g).groups()[0] for g in re.split(r"'\s+'", flds['fldList'])] - assert flds['nrecords'] == len(flds['fldList']) + if 'nFlds' in flds: + flds['nFlds'] = int(flds['nFlds']) + assert flds['nFlds'] == len(flds['fldList']) + else: + assert flds['nrecords'] == len(flds['fldList']) return flds def _get_useful_info_from_meta_file(metafile): @@ -81,6 +85,10 @@ def _get_useful_info_from_meta_file(metafile): name = meta['basename'] fldlist = None + if 'pickup' in metafile: # if 'nFlds' in meta: + # reset UVel to be a single 3D field + shape[0]=int( (nrecs-3)/(meta['nFlds']-3) ) + return nrecs, shape, name, dtype, fldlist @@ -287,13 +295,22 @@ def read_mds(fname, iternum=None, use_mmap=None, endian='>', shape=None, # convert list into dictionary out = {} for n, name in enumerate(file_metadata['fldList']): - if ndims == 3: - out[name] = d[n] - elif ndims == 2: - if use_mmap: - out[name] = d[n].reshape((ny,nx)) + if 'pickup' in file_metadata['basename']: # pickup file + if "Eta" in name: # 2D variables + if use_mmap: + out[name] = d[n].reshape((ny,nx)) + else: + out[name] = d[n][:,0,:] else: - out[name] = d[n][:,0,:] + out[name] = d[n] + else: + if ndims == 2: + if use_mmap: + out[name] = d[n].reshape((ny,nx)) + else: + out[name] = d[n][:,0,:] + elif ndims == 3: + out[name] = d[n] # --------------- LEGACY -------------------------- # from legacy code (needs to be phased out) @@ -785,16 +802,38 @@ def read_all_variables(variable_list, file_metadata, use_mmap=False, """ + local_metadata = file_metadata out = [] for variable in variable_list: + if 'basename' in local_metadata: + if 'pickup' in local_metadata['basename']: # pickup file + if "Eta" in variable: + chunks="2D" + local_metadata['nDims'] = 2 + local_metadata['nz'] = 1 + local_metadata['dimList'] = [ + file_metadata['dimList'][0], + file_metadata['dimList'][1] + ] + else: + chunks="3D" + nz=(local_metadata['nrecords']-3) / (local_metadata['nFlds']-3) + local_metadata['nDims'] = 3 + local_metadata['nz'] = nz + local_metadata['dimList'] = [ + file_metadata['dimList'][0], + file_metadata['dimList'][1], + [nz,1,nz] + ] + if chunks == "2D": - out.append(read_2D_chunks(variable, file_metadata, + out.append(read_2D_chunks(variable, local_metadata, use_mmap=use_mmap, use_dask=use_dask)) elif chunks == "3D": - out.append(read_3D_chunks(variable, file_metadata, + out.append(read_3D_chunks(variable, local_metadata, use_mmap=use_mmap, use_dask=use_dask)) elif chunks == "CS": - out.append(read_CS_chunks(variable, file_metadata, + out.append(read_CS_chunks(variable, local_metadata, use_mmap=use_mmap, use_dask=use_dask)) return out