Merge pull request #77 from StagPython/field-meta

snap.fields items come with their metadata

Author: amorison

docs/sources/stagyydata.rst

@@ -189,8 +189,14 @@ Vector and scalar fields are accessible through :attr:`Step.fields
 <stagpy._step.Step.fields>` using their name as key. For example, the
 temperature field of the 100th snapshot is obtained with
 ``sdat.snaps[100].fields['T']``.  Valid names of fields can be obtained by
-running ``% stagpy var``. Fields are four dimensional arrays, with indices in
-the order x, y, z and block.
+running ``% stagpy var``. Items are named tuples with two elements:
+
+- :data:`values`: the field itself, a four dimensional array with indices in
+  the order x, y, z and block;
+- :data:`meta`: metadata of the field, also a named tuple with:
+
+    - :data:`description`: explanation of what the field is;
+    - :data:`dim`: the dimension of the field (if applicable) in SI units.
 
 Tracers data
 ------------

docs/sources/tuto.rst

@@ -91,7 +91,7 @@ Instantiating and using this class is rather simple::
 
     # temperature field of the 10000th time step
     # (will be None if no snapshot is available at this timestep)
-    temp_field = sdat.steps[10000].fields['T']
+    temp_field = sdat.steps[10000].fields['T'].values
 
     # iterate through snaps 100, 105, 110... up to the last one
     for snap in sdat.snaps[100::5]:

stagpy/_step.py

@@ -224,13 +224,18 @@ def at_r(self, rval):
         return np.argmin(np.abs(self.r_centers - rval))
 
 
+Field = namedtuple('Field', ['values', 'meta'])
+
+
 class _Fields(Mapping):
     """Fields data structure.
 
     The :attr:`Step.fields` attribute is an instance of this class.
 
     :class:`_Fields` inherits from :class:`collections.abc.Mapping`. Keys are
-    fields names defined in :data:`stagpy.phyvars.[S]FIELD[_EXTRA]`.
+    fields names defined in :data:`stagpy.phyvars.[S]FIELD[_EXTRA]`.  Each item
+    is a name tuple ('values', 'meta'), respectively the field itself, and a
+    :class:`stagpy.phyvars.Varf` instance with relevant metadata.
 
     Attributes:
         step (:class:`Step`): the step object owning the :class:`_Fields`
@@ -252,7 +257,10 @@ def __getitem__(self, name):
         if name in self._vars:
             fld_names, parsed_data = self._get_raw_data(name)
         elif name in self._extra:
-            self._data[name] = self._extra[name].description(self.step)
+            meta = self._extra[name]
+            field = meta.description(self.step)
+            meta = phyvars.Varf(_helpers.baredoc(meta.description), meta.dim)
+            self._data[name] = Field(field, meta)
             return self._data[name]
         else:
             raise error.UnknownFieldVarError(name)
@@ -322,7 +330,7 @@ def _set(self, name, fld):
             while len(col_fld) > sdat.nfields_max:
                 istep, fld_name = col_fld.pop(0)
                 del sdat.steps[istep].fields[fld_name]
-        self._data[name] = fld
+        self._data[name] = Field(fld, self._vars[name])
 
     def __delitem__(self, name):
         if name in self._data:

stagpy/field.py

@@ -46,10 +46,10 @@ def _threed_extract(step, var, walls=False):
         i_x = i_y = slice(None)
         varx, vary = var + '1', var + '2'
     if is_vector:
-        data = (step.fields[varx][i_x, i_y, i_z, 0],
-                step.fields[vary][i_x, i_y, i_z, 0])
+        data = (step.fields[varx].values[i_x, i_y, i_z, 0],
+                step.fields[vary].values[i_x, i_y, i_z, 0])
     else:
-        data = step.fields[var][i_x, i_y, i_z, 0]
+        data = step.fields[var].values[i_x, i_y, i_z, 0]
     return (xcoord, ycoord), data
 
 
@@ -77,11 +77,11 @@ def get_meshes_fld(step, var, walls=False):
         var (str): scalar field name.
         walls (bool): consider the walls as the relevant mesh.
     Returns:
-        tuple of :class:`numpy.array`: xmesh, ymesh, fld
-            2D arrays containing respectively the x position, y position, and
-            the value of the requested field.
+        tuple of :class:`numpy.array`: xmesh, ymesh, fld, meta
+            2D arrays containing respectively the x position, y position, the
+            values and the metadata of the requested field.
     """
-    fld = step.fields[var]
+    fld, meta = step.fields[var]
     hwalls = (walls or fld.shape[0] != step.geom.nxtot or
               fld.shape[1] != step.geom.nytot)
     if step.geom.threed and step.geom.cartesian:
@@ -99,7 +99,7 @@ def get_meshes_fld(step, var, walls=False):
         fld = fld[0, :, :, 0]
     if step.geom.cartesian:
         xmesh, ymesh = np.meshgrid(xcoord, ycoord, indexing='ij')
-    return xmesh, ymesh, fld
+    return xmesh, ymesh, fld, meta
 
 
 def get_meshes_vec(step, var):
@@ -119,17 +119,17 @@ def get_meshes_vec(step, var):
         (xcoord, ycoord), (vec1, vec2) = _threed_extract(step, var)
     elif step.geom.twod_xz:
         xcoord, ycoord = step.geom.x_walls, step.geom.z_centers
-        vec1 = step.fields[var + '1'][:, 0, :, 0]
-        vec2 = step.fields[var + '3'][:, 0, :, 0]
+        vec1 = step.fields[var + '1'].values[:, 0, :, 0]
+        vec2 = step.fields[var + '3'].values[:, 0, :, 0]
     elif step.geom.cartesian and step.geom.twod_yz:
         xcoord, ycoord = step.geom.y_walls, step.geom.z_centers
-        vec1 = step.fields[var + '2'][0, :, :, 0]
-        vec2 = step.fields[var + '3'][0, :, :, 0]
+        vec1 = step.fields[var + '2'].values[0, :, :, 0]
+        vec2 = step.fields[var + '3'].values[0, :, :, 0]
     else:  # spherical yz
         pcoord = step.geom.p_walls
         pmesh = np.outer(pcoord, np.ones(step.geom.nrtot))
-        vec_phi = step.fields[var + '2'][0, :, :, 0]
-        vec_r = step.fields[var + '3'][0, :, :, 0]
+        vec_phi = step.fields[var + '2'].values[0, :, :, 0]
+        vec_r = step.fields[var + '3'].values[0, :, :, 0]
         vec1 = vec_r * np.cos(pmesh) - vec_phi * np.sin(pmesh)
         vec2 = vec_phi * np.cos(pmesh) + vec_r * np.sin(pmesh)
         pcoord, rcoord = step.geom.p_walls, step.geom.r_centers
@@ -162,17 +162,12 @@ def plot_scalar(step, var, field=None, axis=None, **extra):
             :func:`~matplotlib.axes.Axes.pcolormesh`, and the colorbar returned
             by :func:`matplotlib.pyplot.colorbar`.
     """
-    if var in phyvars.FIELD:
-        meta = phyvars.FIELD[var]
-    else:
-        meta = phyvars.FIELD_EXTRA[var]
-        meta = phyvars.Varf(_helpers.baredoc(meta.description), meta.dim)
     if step.geom.threed and step.geom.spherical:
         raise NotAvailableError(
             'plot_scalar not implemented for 3D spherical geometry')
 
-    xmesh, ymesh, fld = get_meshes_fld(step, var,
-                                       walls=not conf.field.interpolate)
+    xmesh, ymesh, fld, meta = get_meshes_fld(step, var,
+                                             walls=not conf.field.interpolate)
     if conf.field.interpolate and \
        step.geom.spherical and step.geom.twod_yz and \
        fld.shape[0] == step.geom.nytot:
@@ -254,7 +249,7 @@ def plot_iso(axis, step, var, **extra):
         extra (dict): options that will be passed on to
             :func:`matplotlib.axes.Axes.contour`.
     """
-    xmesh, ymesh, fld = get_meshes_fld(step, var)
+    xmesh, ymesh, fld, _ = get_meshes_fld(step, var)
     if conf.field.shift:
         fld = np.roll(fld, conf.field.shift, axis=0)
     extra_opts = dict(linewidths=1)
@@ -297,11 +292,8 @@ def _findminmax(sdat, sovs):
     for step in sdat.walk.filter(snap=True):
         for var in sovs:
             if var in step.fields:
-                if var in phyvars.FIELD:
-                    dim = phyvars.FIELD[var].dim
-                else:
-                    dim = phyvars.FIELD_EXTRA[var].dim
-                field, _ = sdat.scale(step.fields[var], dim)
+                field, meta = step.fields[var]
+                field, _ = sdat.scale(field, meta.dim)
                 if var in minmax:
                     minmax[var] = (min(minmax[var][0], np.nanmin(field)),
                                    max(minmax[var][1], np.nanmax(field)))

stagpy/plates.py

@@ -13,9 +13,9 @@
 
 def detect_plates_vzcheck(step, vz_thres_ratio=0):
     """Detect plates and check with vz and plate size."""
-    v_z = step.fields['v3'][0, :-1, :, 0]
-    v_x = step.fields['v2'][0, :, :, 0]
-    tcell = step.fields['T'][0, :, :, 0]
+    v_z = step.fields['v3'].values[0, :-1, :, 0]
+    v_x = step.fields['v2'].values[0, :, :, 0]
+    tcell = step.fields['T'].values[0, :, :, 0]
     n_z = step.geom.nztot
     nphi = step.geom.nptot
     r_c = step.geom.r_centers
@@ -68,7 +68,7 @@ def detect_plates_vzcheck(step, vz_thres_ratio=0):
 
 def detect_plates(step, vrms_surface, fids, time):
     """Detect plates using derivative of horizontal velocity."""
-    vphi = step.fields['v2'][0, :, :, 0]
+    vphi = step.fields['v2'].values[0, :, :, 0]
     ph_coord = step.geom.p_centers
 
     if step.sdat.par['boundaries']['air_layer']:
@@ -123,7 +123,7 @@ def detect_plates(step, vrms_surface, fids, time):
 
     dv_ridge = dvph2[arggreat_dv]
     if 'age' in conf.plates.plot:
-        agefld = step.fields['age'][0, :, :, 0]
+        agefld = step.fields['age'].values[0, :, :, 0]
         age_surface = np.ma.masked_where(agefld[:, indsurf] < 0.00001,
                                          agefld[:, indsurf])
         age_surface_dim = age_surface * vrms_surface *\
@@ -180,9 +180,9 @@ def plot_plate_limits_field(axis, rcmb, ridges, trenches):
 def plot_plates(step, time, vrms_surface, trench, ridge, agetrench,
                 topo, fids):
     """Handle plotting stuff."""
-    vphi = step.fields['v2'][0, :, :, 0]
-    tempfld = step.fields['T'][0, :, :, 0]
-    concfld = step.fields['c'][0, :, :, 0]
+    vphi = step.fields['v2'].values[0, :, :, 0]
+    tempfld = step.fields['T'].values[0, :, :, 0]
+    concfld = step.fields['c'].values[0, :, :, 0]
     timestep = step.isnap
 
     if step.sdat.par['boundaries']['air_layer']:
@@ -296,7 +296,7 @@ def plot_plates(step, time, vrms_surface, trench, ridge, agetrench,
 
     # plotting velocity and second invariant of stress
     if 'str' in conf.plates.plot:
-        stressfld = step.fields['sII'][0, :, :, 0]
+        stressfld = step.fields['sII'].values[0, :, :, 0]
         fig0, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(12, 8))
         ax1.plot(step.geom.p_walls, vphi[:, indsurf], label='Vel')
         ax1.axhline(y=0, xmin=0, xmax=2 * np.pi,
@@ -344,7 +344,7 @@ def plot_plates(step, time, vrms_surface, trench, ridge, agetrench,
 
     # plotting velocity and age at surface
     if 'age' in conf.plates.plot:
-        agefld = step.fields['age'][0, :, :, 0]
+        agefld = step.fields['age'].values[0, :, :, 0]
         age_surface = np.ma.masked_where(
             agefld[:, indsurf] < 0.00001, agefld[:, indsurf])
         age_surface_dim = (age_surface * vrms_surface * conf.scaling.ttransit /
@@ -467,7 +467,7 @@ def lithospheric_stress(step, trench, ridge, time):
     timestep = step.isnap
     base_lith = step.geom.rcmb + 1 - 0.105
 
-    stressfld = step.fields['sII'][0, :, :, 0]
+    stressfld = step.fields['sII'].values[0, :, :, 0]
     r_centers = np.outer(np.ones(stressfld.shape[0]), step.geom.r_centers)
     stressfld = np.ma.masked_where(r_centers < base_lith, stressfld)
 
@@ -485,10 +485,10 @@ def lithospheric_stress(step, trench, ridge, time):
     saveplot(fig, 'lith', timestep)
 
     # velocity
-    vphi = step.fields['v2'][0, :, :, 0]
+    vphi = step.fields['v2'].values[0, :, :, 0]
 
     # position of continents
-    concfld = step.fields['c'][0, :, :, 0]
+    concfld = step.fields['c'].values[0, :, :, 0]
     if step.sdat.par['boundaries']['air_layer']:
         # we are a bit below the surface; delete "-some number"
         # to be just below
@@ -612,12 +612,12 @@ def main_plates(sdat):
                         agetrenches, topo, fids)
 
             # prepare for continent plotting
-            concfld = step.fields['c'][0, :, :, 0]
+            concfld = step.fields['c'].values[0, :, :, 0]
             continentsfld = np.ma.masked_where(
                 concfld < 3, concfld)  # plotting continents, to-do
             continentsfld = continentsfld / continentsfld
 
-            temp = step.fields['T'][0, :, :, 0]
+            temp = step.fields['T'].values[0, :, :, 0]
             tgrad = (temp[:, isurf - 1] - temp[:, isurf]) /\
                 (step.geom.r_centers[isurf] - step.geom.r_centers[isurf - 1])
 
@@ -626,7 +626,7 @@ def main_plates(sdat):
             io_surface(timestep, time, fids[4], topo[:, 1])
             if 'age' in conf.plates.plot:
                 io_surface(timestep, time, fids[5],
-                           step.fields['age'][0, :, isurf, 0])
+                           step.fields['age'].values[0, :, isurf, 0])
 
             # plot viscosity field with position of trenches and ridges
             etamin, _ = sdat.scale(1e-2, 'Pa')

stagpy/processing.py

@@ -281,13 +281,13 @@ def stream_function(step):
     """
     if step.geom.twod_yz:
         x_coord = step.geom.y_walls
-        v_x = step.fields['v2'][0, :, :, 0]
-        v_z = step.fields['v3'][0, :, :, 0]
+        v_x = step.fields['v2'].values[0, :, :, 0]
+        v_z = step.fields['v3'].values[0, :, :, 0]
         shape = (1, v_x.shape[0], v_x.shape[1], 1)
     elif step.geom.twod_xz and step.geom.cartesian:
         x_coord = step.geom.x_walls
-        v_x = step.fields['v1'][:, 0, :, 0]
-        v_z = step.fields['v3'][:, 0, :, 0]
+        v_x = step.fields['v1'].values[:, 0, :, 0]
+        v_z = step.fields['v3'].values[:, 0, :, 0]
         shape = (v_x.shape[0], 1, v_x.shape[1], 1)
     else:
         raise NotAvailableError('Stream function only implemented in '

tests/test_field.py

@@ -27,13 +27,26 @@ def test_valid_field_var_invalid():
     assert not stagpy.field.valid_field_var('dummyfieldvar')
 
 
-def test_get_meshes_fld(step):
-    xmesh, ymesh, fld = stagpy.field.get_meshes_fld(step, 'T')
+def test_get_meshes_fld_no_walls(step):
+    xmesh, ymesh, fld, meta = stagpy.field.get_meshes_fld(step, 'T',
+                                                          walls=False)
     assert len(fld.shape) == 2
-    assert xmesh.shape == ymesh.shape == fld.shape
+    assert xmesh.shape[0] == ymesh.shape[0] == fld.shape[0]
+    assert xmesh.shape[1] == ymesh.shape[1] == fld.shape[1]
+    assert meta.description == 'Temperature'
+
+
+def test_get_meshes_fld_walls(step):
+    xmesh, ymesh, fld, meta = stagpy.field.get_meshes_fld(step, 'T',
+                                                          walls=True)
+    assert len(fld.shape) == 2
+    assert xmesh.shape[0] == ymesh.shape[0] == fld.shape[0] + 1
+    assert xmesh.shape[1] == ymesh.shape[1] == fld.shape[1] + 1
+    assert meta.description == 'Temperature'
 
 
 def test_get_meshes_vec(step):
     xmesh, ymesh, vec1, vec2 = stagpy.field.get_meshes_vec(step, 'v')
     assert len(vec1.shape) == 2
-    assert xmesh.shape == ymesh.shape == vec1.shape == vec2.shape
+    assert xmesh.shape[0] == ymesh.shape[0] == vec1.shape[0] == vec2.shape[0]
+    assert xmesh.shape[1] == ymesh.shape[1] == vec1.shape[1] == vec2.shape[1]

tests/test_processing.py

@@ -52,4 +52,4 @@ def test_energy_prof(step):
 
 def test_stream_function(step):
     psi = processing.stream_function(step)
-    assert psi.shape[1:3] == step.fields['v3'].shape[1:3]
+    assert psi.shape[1:3] == step.fields['v3'].values.shape[1:3]