From 13467f59146732336a735b1a2743e64e30f5fb8a Mon Sep 17 00:00:00 2001
From: Adrien Morison <adrien.morison@gmail.com>
Date: Tue, 8 Sep 2020 14:27:46 +0100
Subject: [PATCH] Add hdf5 surface fields reading

---
 stagpy/_step.py         | 16 ++++++++++++----
 stagpy/phyvars.py       | 14 ++++++++++++++
 stagpy/stagyyparsers.py | 12 ++++++++++--
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/stagpy/_step.py b/stagpy/_step.py
index 491289e..36d243e 100644
--- a/stagpy/_step.py
+++ b/stagpy/_step.py
@@ -235,10 +235,10 @@ def __getitem__(self, name):
         for fld_name, fld in zip(fld_names, fields):
             if self._header['xyp'] == 0:
                 if not self.geom.twod_yz:
-                    newline = (fld[:1, :, :, :] + fld[-1:, :, :, :]) / 2
+                    newline = (fld[:1, ...] + fld[-1:, ...]) / 2
                     fld = np.concatenate((fld, newline), axis=0)
                 if not self.geom.twod_xz:
-                    newline = (fld[:, :1, :, :] + fld[:, -1:, :, :]) / 2
+                    newline = (fld[:, :1, ...] + fld[:, -1:, ...]) / 2
                     fld = np.concatenate((fld, newline), axis=1)
             self._set(fld_name, fld)
         return self._data[name]
@@ -277,8 +277,16 @@ def _get_raw_data(self, name):
             for filestem, list_fvar in self._filesh5.items():
                 if name in list_fvar:
                     break
+            if filestem in phyvars.SFIELD_FILES_H5:
+                xmff = 'Data{}.xmf'.format(
+                    'Bottom' if name.endswith('bot') else 'Surface')
+                _ = self.geom
+                header = self._header
+            else:
+                xmff = 'Data.xmf'
+                header = None
             parsed_data = stagyyparsers.read_field_h5(
-                self.step.sdat.hdf5 / 'Data.xmf', filestem, self.step.isnap)
+                self.step.sdat.hdf5 / xmff, filestem, self.step.isnap, header)
         return list_fvar, parsed_data
 
     def _set(self, name, fld):
@@ -406,7 +414,7 @@ def __init__(self, istep, sdat):
         self.fields = _Fields(self, phyvars.FIELD, phyvars.FIELD_EXTRA,
                               phyvars.FIELD_FILES, phyvars.FIELD_FILES_H5)
         self.sfields = _Fields(self, phyvars.SFIELD, [],
-                               phyvars.SFIELD_FILES, [])
+                               phyvars.SFIELD_FILES, phyvars.SFIELD_FILES_H5)
         self.tracers = _Tracers(self)
         self._isnap = UNDETERMINED
 
diff --git a/stagpy/phyvars.py b/stagpy/phyvars.py
index 2924a7f..2f326c6 100644
--- a/stagpy/phyvars.py
+++ b/stagpy/phyvars.py
@@ -107,6 +107,20 @@
     ('cr', ['crust']),
 ))
 
+SFIELD_FILES_H5 = OrderedDict((
+    ('BottomTopography', ['topo_bot']),
+    ('SurfaceTopography', ['topo_top']),
+    ('BottomGeoid', ['geoid_bot']),
+    ('TopGeoid', ['geoid_top']),
+    ('BottomCSGeoid', ['topo_g_bot']),
+    ('TopCSGeoid', ['topo_g_top']),
+    ('BottomHeatFlux', ['fbot']),
+    ('TopHeatFlux', ['ftop']),
+    ('BottomHFSpectrum', ['fsbot']),
+    ('TopHFSpectrum', ['fstop']),
+    ('CrustThickness', ['crust']),
+))
+
 Varr = namedtuple('Varr', ['description', 'kind', 'dim'])
 RPROF = OrderedDict((
     ('r', Varr('Radial coordinate', 'Radius', 'm')),
diff --git a/stagpy/stagyyparsers.py b/stagpy/stagyyparsers.py
index 1cbd57e..86c100a 100644
--- a/stagpy/stagyyparsers.py
+++ b/stagpy/stagyyparsers.py
@@ -16,7 +16,7 @@
 import h5py
 
 from .error import ParsingError
-from .phyvars import FIELD_FILES_H5
+from .phyvars import FIELD_FILES_H5, SFIELD_FILES_H5
 
 
 def _tidy_names(names, nnames, extra_names=None):
@@ -744,7 +744,7 @@ def _flds_shape(fieldname, header):
     """Compute shape of flds variable."""
     shp = list(header['nts'])
     shp.append(header['ntb'])
-    if len(FIELD_FILES_H5[fieldname]) == 3:
+    if len(FIELD_FILES_H5.get(fieldname, [])) == 3:
         shp.insert(0, 3)
         # extra points
         header['xp'] = int(header['nts'][0] != 1)
@@ -813,10 +813,15 @@ def read_field_h5(xdmf_file, fieldname, snapshot, header=None):
                 fld = fld.reshape((shp[0], shp[1], 1, shp[2]))
                 if header['rcmb'] < 0:
                     fld = fld[(1, 2, 0), ...]
+            elif fieldname in SFIELD_FILES_H5:
+                fld = fld.reshape((1, npc[0], npc[1], 1))
             elif header['nts'][1] == 1:  # cart XZ
                 fld = fld.reshape((1, shp[0], 1, shp[1]))
             ifs = [icore // np.prod(header['ncs'][:i]) % header['ncs'][i] *
                    npc[i] for i in range(3)]
+            if fieldname in SFIELD_FILES_H5:
+                ifs[2] = 0
+                npc[2] = 1
             if header['zp']:  # remove top row
                 fld = fld[:, :, :, :-1]
             flds[:,
@@ -828,6 +833,9 @@ def read_field_h5(xdmf_file, fieldname, snapshot, header=None):
 
     flds = _post_read_flds(flds, header)
 
+    if fieldname in SFIELD_FILES_H5:
+        # remove z component
+        flds = flds[..., 0, :]
     return (header, flds) if data_found else None