Read pickup MDS (#351)

* llc regional grid binary writing

* pioneer metadata

* nesba domain

* nesba meta

* single mitgcm face

* aste1080 metadata; doesnt match aste metadata

* reduced nesb domain

* update nesb meta

* check for single element length vars

* rm custom domain metas

* avoid error when using MDS from pickups

* add pickup meta check

* clean up utils test

* handle variable data sizes in pickup

* fix CI?

* comment on update

* Update xmitgcm/test/test_utils.py

Co-authored-by: Martin Losch <[email protected]>

---------

Co-authored-by: Martin Losch <[email protected]>

Author: IvanaEscobar

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)

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."""

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