(Towards #868) Metadata support for fields with nlayers and ndata

doc/user_guide/lfric.rst

@@ -162,24 +162,32 @@ least rank (number of dimensions) one. Scalar arrays are identified with
 Field
 +++++
 
-LFRic API fields, identified with ``GH_FIELD`` metadata, represent
-FEM discretisations of various dynamical core prognostic and diagnostic
+LFRic API fields, identified with ``GH_FIELD`` metadata, represent FEM
+discretisations of various dynamical core prognostic and diagnostic
 variables. In FEM, variables are discretised by placing them into a
 function space (see :ref:`lfric-function-space`) from which they
 inherit a polynomial expansion via the basis functions of that space.
 Field values at points within a cell are evaluated as the sum of a set
-of basis functions multiplied by coefficients which are the data points.
-Points of evaluation are determined by a quadrature object
+of basis functions multiplied by coefficients which are the data
+points.  Points of evaluation are determined by a quadrature object
 (:ref:`lfric-quadrature`) and are independent of the function space
-the field is on. Placement of field data points, also called degrees of
-freedom (hereafter "DoFs"), is determined by the function space the field
-is on.
+the field is on. Placement of field data points, also called degrees
+of freedom (hereafter "DoFs"), is determined by the function space the
+field is on. An LFRic multi-data field can have more than one value
+associated with each data point.
+
 LFRic fields passed as arguments to any :ref:`LFRic kernel
 <lfric-kernel-valid-data-type>` can be of ``real`` or ``integer``
 primitive type. In the LFRic infrastructure, these fields are
 represented by instances of the ``field_type`` and ``integer_field_type``
 classes, respectively.
 
+Different fields may be defined on different numbers of vertical layers.
+The the number of layers can be as few as one (a 2D field). Unfortunately,
+the number of layers affects the numbering of the DoFs of a field. Therefore,
+a distinct DoF map is required for each unique combination of function
+space and number of vertical levels.
+
 .. _lfric-field-vector:
 
 Field Vector
@@ -919,8 +927,8 @@ All three CMA-related kernel types must obey the following rules:
 1) Since a CMA operator only acts within a single column of data,
    stencil operations are not permitted.
 
-2) No vector quantities (e.g. ``GH_FIELD*3`` - see below) are
-   permitted as arguments.
+2) No vector quantities (e.g. ``GH_FIELD*3`` - see below) or
+   multi-data fields are permitted as arguments.
 
 3) The kernel must operate on cell-columns.
 
@@ -1450,7 +1458,7 @@ Supported Function Spaces
 As mentioned in the :ref:`lfric-field` and :ref:`lfric-field-vector`
 sections, the function space of an argument specifies how it maps
 onto the underlying topology and, additionally, whether the data at a
-point is a vector. In LFRic API the dimension of the basis function
+point is a vector. In the LFRic API the dimension of the basis function
 set for the scalar function spaces is 1 and for the vector function spaces
 is 3 (see the table in :ref:`lfric-stub-generation-rules` for the
 dimensions of the basis and differential basis functions).
@@ -1636,7 +1644,7 @@ to have stencil accesses, these two options are mutually exclusive.
 The metadata for each case is described in the following sections.
 
 Stencil Metadata
-________________
+""""""""""""""""
 
 
 Stencil metadata specifies that the corresponding field argument is accessed
@@ -1716,8 +1724,7 @@ be found in ``examples/lfric/eg5``.
 .. _lfric-intergrid-mdata:
 
 Inter-Grid Metadata
-___________________
-
+"""""""""""""""""""
 
 The alternative form of the optional fifth metadata argument for a
 field specifies which mesh the associated field is on.  This is
@@ -1748,6 +1755,44 @@ meshes cannot be on the same function space while those on the same
 mesh must also be on the same function space.
 
 
+Number of Layers Metadata
+"""""""""""""""""""""""""
+
+If a particular field/operator kernel argument has a number of vertical
+levels that is *not* the same as the first field/operator argument then
+this must be specified using the ``NLAYERS`` option to ``GH_FIELD``/
+``GH_OPERATOR``, e.g.::
+
+  arg_type(GH_FIELD, GH_REAL, GH_READ, W3, NLAYERS=1)
+
+The value specified for ``NLAYERS`` may be an integer literal if it is known
+at compile time. Alternatively, it may be given a name (e.g.
+``GH_NLAYERS_SHIFTED``). If two or more field/operator
+arguments are on the same function space and have the same number
+of layers (whether a literal or a name) then only one dofmap (that of the
+first such field listed in the metadata) is passed to the kernel for
+those arguments.
+
+(Since the value of ``NLAYERS`` is looked-up from the corresponding kernel
+argument at run time, the labels given in the kernel metadata are just that
+- they do not have to correspond to anything in the LFRic infrastructure.)
+
+Multi-Data Metadata
+"""""""""""""""""""
+
+A multi-data field is the same as a standard field apart from having multiple
+values associated with each DoF. This is indicated in the field metadata by
+the optional ``NDATA`` argument to GH_FIELD, e.g.::
+
+  arg_type(GH_FIELD, GH_REAL, GH_READ, W2, NDATA=4)
+
+The value specified for ``NDATA`` may be a literal if it is known at
+compile time. Alternatively, it may be given the special value
+``GH_RUNTIME`` which means that the number of data values at each DoF
+is to be determined at runtime by querying the field object (in the
+generated PSy layer).
+
+
 Column-wise Operators (CMA)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -2080,7 +2125,12 @@ conventions, are:
       4) If the field entry stencil access is of type ``XORY1D`` then
          add an additional ``integer`` direction argument of kind
          ``i_def`` and with intent ``in``.
-
+      5) If the field is multi-data then the kernel must be passed the
+	 value of ``NDATA``: add an additional ``integer``, scalar
+	 argument of kind ``i_def`` and intent ``in``.
+      6) If the field has a custom number of vertical levels then pass this as
+	 an additional ``integer``, scalar argument of kind ``i_def`` and
+	 intent ``in``.
    3) If the current entry is a field vector then for each dimension
       of the vector, include a field array. The field array name is
       specified as
@@ -2106,21 +2156,26 @@ conventions, are:
       the data type and kind specified in the metadata. The ScalarArray
       must be denoted with intent ``in`` to match its read-only nature.
 
-4) For each function space in the order they appear in the metadata arguments
-   (the ``to`` function space of an operator is considered to be before the
+4) DoF maps for function spaces are handled in the order they appear in the
+   metadata arguments (the ``to`` function space of an operator is considered
+   to be before the
    ``from`` function space of the same operator as it appears first in
-   lexicographic order)
+   lexicographic order). Note that if two fields on a given function space have
+   differing numbers of vertical layers, then each requires that a
+   dofmap be supplied (because the number of vertical layers alters the
+   *values* within the map). For each required DoF map:
 
    1) Include the number of local degrees of freedom (i.e. number per-cell)
       for the function space. This is an ``integer`` of kind ``i_def`` and
       has intent ``in``. The name of this argument is
       ``"ndf_"<field_function_space>``.
+
    2) If there is a field on this space
 
       1) Include the unique number of degrees of freedom for the function
          space. This is an ``integer`` of kind ``i_def`` and has intent ``in``.
          The name of this argument is ``"undf_"<field_function_space>``.
-      2) Include the **dofmap** for this function space. This is an ``integer``
+      2) Include the **dofmap** itself. This is an ``integer``
          array of kind ``i_def`` with intent ``in``. It has one dimension
          sized by the local degrees of freedom for the function space.
 
@@ -2443,8 +2498,9 @@ dofmap for both the to- and from-function spaces of the CMA
 operator. Since it does not have any LMA operator arguments it does
 not require the ``ncell_3d`` and ``nlayers`` scalar arguments. (Since a
 column-wise operator is, by definition, assembled for a whole column,
-there is no loop over levels when applying it.)
-The full set of rules is then:
+there is no loop over levels when applying it.) Note that fields with
+non-standard ``nlayers`` or ``ndata > 1`` cannot be supplied as
+arguments to CMA kernels. The full set of rules is:
 
 1) Include the ``cell`` argument. ``cell`` is an ``integer`` of kind
    ``i_def`` and has intent ``in``.

src/psyclone/domain/lfric/lfric_arg_descriptor.py

@@ -51,7 +51,8 @@
 from psyclone.domain.lfric.lfric_constants import LFRicConstants
 from psyclone.errors import InternalError
 import psyclone.expression as expr
-from psyclone.parse.kernel import Descriptor, get_stencil, get_mesh
+from psyclone.parse.kernel import (
+    Descriptor, get_stencil, get_mesh, get_nlevels)
 from psyclone.parse.utils import ParseError
 
 # API configuration
@@ -109,6 +110,7 @@ def __init__(self, arg_type, operates_on, metadata_index):
         self._function_space2 = None
         self._stencil = None
         self._mesh = None
+        self._nlevels = ""
         self._nargs = 0
 
         # Check for the correct argument type descriptor
@@ -400,10 +402,11 @@ def _init_field(self, arg_type, operates_on):
                 f"'{arg_type.args[prop_ind].name}' in '{arg_type}'.")
         self._function_space1 = arg_type.args[prop_ind].name
 
-        # The optional 5th argument is either a stencil specification
-        # or a mesh identifier (for inter-grid kernels)
-        prop_ind = 4
-        if self._nargs == nargs_field_max:
+        num_args = len(arg_type.args)
+        if num_args > 4:
+            # The optional 5th argument is either a stencil specification
+            # or a mesh identifier (for inter-grid kernels)
+            prop_ind = 4
             try:
                 if "stencil" in str(arg_type.args[prop_ind]):
                     self._stencil = get_stencil(
@@ -412,15 +415,33 @@ def _init_field(self, arg_type, operates_on):
                 elif "mesh" in str(arg_type.args[prop_ind]):
                     self._mesh = get_mesh(arg_type.args[prop_ind],
                                           const.VALID_MESH_TYPES)
+                elif "nlevels" in str(arg_type.args[prop_ind]):
+                    self._nlevels = get_nlevels(arg_type.args[prop_ind])
                 else:
                     raise ParseError("Unrecognised metadata entry")
             except ParseError as err:
                 raise ParseError(
                     f"In the LFRic API argument {prop_ind+1} of a 'meta_arg' "
                     f"field entry must be either a valid stencil specification"
-                    f" or a mesh identifier (for inter-grid kernels). However,"
-                    f" entry '{arg_type}' raised the following error: "
-                    f"{err}.") from err
+                    f", a number of levels or a mesh identifier (for inter-"
+                    f"grid kernels). However, entry '{arg_type}' raised the "
+                    f"following error: {err}.") from err
+
+        if num_args > 5:
+            # If there are this many arguments then the last one must be
+            # nlevels.
+            prop_ind = 5
+            try:
+                if "nlevels" in str(arg_type.args[prop_ind]):
+                    self._nlevels = get_nlevels(arg_type.args[prop_ind])
+                else:
+                    raise ParseError("Unrecognised metadata entry")
+            except ParseError as err:
+                raise ParseError(
+                    f"In the LFRic API, argument {prop_ind+1} of a 'meta_arg' "
+                    f"field entry must be a number of levels. However entry "
+                    f"'{arg_type}' raised the following error: {err}."
+                ) from err
 
         # Test allowed accesses for fields
         field_disc_accesses = [AccessType.READ, AccessType.WRITE,
@@ -794,6 +815,14 @@ def function_spaces(self):
         raise InternalError(f"Expected a valid argument type but got "
                             f"'{self._argument_type}'.")
 
+    @property
+    def nlevels(self) -> str:
+        '''
+        :returns: a label identifying the number of vertical levels
+                  associated with this argument.
+        '''
+        return self._nlevels
+
     @property
     def vector_size(self):
         '''

src/psyclone/parse/kernel.py

@@ -400,6 +400,27 @@ def get_stencil(metadata, valid_types):
     return {"type": stencil_type, "extent": stencil_extent}
 
 
+def get_nlevels(metadata: expr.NamedArg) -> str:
+    '''
+    Returns the number of levels described by the supplied meta-data
+
+    :param metadata: node in fparser1 ast holding the meta-data.
+
+    :return: a label identifying the number of vertical levels.
+
+    :raises ParseError: if the supplied ast does not correspond to
+                        `nlevels="some-label"`.
+    '''
+    if (not isinstance(metadata, expr.NamedArg) or
+            metadata.name.lower() != "nlevels"):
+        raise ParseError(
+            f"{metadata} is not a valid mesh identifier (expected "
+            f"nlevels='label')")
+    mesh = metadata.value.lower()
+
+    return mesh
+
+
 class Descriptor():
     '''
     A description of how a kernel argument is accessed, constructed from

src/psyclone/tests/domain/lfric/lfric_field_mdata_test.py

@@ -43,7 +43,6 @@
 
 import os
 import pytest
-import fparser
 from fparser import api as fpapi
 from psyclone.core.access_type import AccessType
 from psyclone.domain.lfric import (LFRicArgDescriptor, LFRicConstants,
@@ -91,12 +90,6 @@
 '''
 
 
-@pytest.fixture(name="disable_fparser_logging", scope="function", autouse=True)
-def disable_fparser_logging_fixture():
-    '''Fixture to automate disabling of fparser logging.'''
-    fparser.logging.disable(fparser.logging.CRITICAL)
-
-
 def test_ad_fld_type_1st_arg():
     ''' Tests that an error is raised when the first argument descriptor
     metadata for a field is invalid. '''
@@ -399,6 +392,24 @@ def test_arg_descriptor_field():
     assert field_descriptor.vector_size == 1
 
 
+def test_fld_nlevels():
+    '''
+    Test a field argument with the optional 'nlevels' metatadata.
+    '''
+    code = FIELD_CODE.replace(
+        "arg_type(gh_scalar, gh_integer, gh_read)",
+        "arg_type(gh_field, gh_real, gh_read, w3, nlevels='double')", 1)
+    ast = fpapi.parse(code, ignore_comments=False)
+    name = "testkern_field_type"
+    mdata = LFRicKernMetadata(ast, name=name)
+    # By default, nlevels is left as an empty string.
+    field_descriptor = mdata.arg_descriptors[5]
+    assert field_descriptor.nlevels == ""
+    # The seventh argument has nlevels specified as "double"
+    field_descriptor = mdata.arg_descriptors[6]
+    assert field_descriptor.nlevels == "double"
+
+
 def test_invalid_vector_operator():
     ''' Tests that an error is raised when a field vector does not
     use "*" as its operator. '''