Repository: stfc/PSyclone

Issue:
Splitting and renaming DynDofMaps from `dynamo0p3.py`
Sub-issue of #2235

Discussion:

Patch:
diff --git a/changelog b/changelog
index 95edbee8d..bbe810bae 100644
--- a/changelog
+++ b/changelog
@@ -4,7 +4,10 @@
 	for the Routines contained in a Module.
 
 	3) PR #2435 for #2422. Add PSyIR GenericInterface to capture the
-	information about which routines implement a specific interface symbol.  
+	information about which routines implement a specific interface symbol.
+
+	4) PR #2457 for #2455. Splitting and renaming DynDofMaps from
+	dynamo0p3.py.
 
 release 2.5.0 14th of February 2024
 
diff --git a/src/psyclone/domain/lfric/__init__.py b/src/psyclone/domain/lfric/__init__.py
index 07e572b72..1a3fe3220 100644
--- a/src/psyclone/domain/lfric/__init__.py
+++ b/src/psyclone/domain/lfric/__init__.py
@@ -71,6 +71,7 @@ from psyclone.domain.lfric.lfric_invokes import LFRicInvokes
 from psyclone.domain.lfric.lfric_scalar_args import LFRicScalarArgs
 from psyclone.domain.lfric.lfric_loop_bounds import LFRicLoopBounds
 from psyclone.domain.lfric.lfric_kern_metadata import LFRicKernMetadata
+from psyclone.domain.lfric.lfric_dofmaps import LFRicDofmaps
 from psyclone.domain.lfric.lfric_stencils import LFRicStencils
 
 
@@ -84,6 +85,7 @@ __all__ = [
     'LFRicArgDescriptor',
     'LFRicCollection',
     'LFRicConstants',
+    'LFRicDofmaps',
     'LFRicExtractDriverCreator',
     'LFRicFields',
     'LFRicInvoke',
diff --git a/src/psyclone/domain/lfric/lfric_builtins.py b/src/psyclone/domain/lfric/lfric_builtins.py
index 45068b037..b9582d02a 100644
--- a/src/psyclone/domain/lfric/lfric_builtins.py
+++ b/src/psyclone/domain/lfric/lfric_builtins.py
@@ -192,7 +192,7 @@ class LFRicBuiltIn(BuiltIn, metaclass=abc.ABCMeta):
 
         :param meta_args: a list of 'meta_args' metadata.
         :type meta_args: List[subclass of \
-            :py:class:`psyclone.domain.lifric.kernel.CommonArgMetadata`]
+            :py:class:`psyclone.domain.lfric.kernel.CommonArgMetadata`]
 
         :returns: LFRic kernel metadata for this built-in.
         :rtype: :py:class:`psyclone.domain.lfric.kernel.LFRicKernelMetadata`
diff --git a/src/psyclone/domain/lfric/lfric_dofmaps.py b/src/psyclone/domain/lfric/lfric_dofmaps.py
new file mode 100644
index 000000000..3716fb5b2
--- /dev/null
+++ b/src/psyclone/domain/lfric/lfric_dofmaps.py
@@ -0,0 +1,298 @@
+# -----------------------------------------------------------------------------
+# BSD 3-Clause License
+#
+# Copyright (c) 2017-2024, Science and Technology Facilities Council.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+# -----------------------------------------------------------------------------
+# Authors R. W. Ford, A. R. Porter and S. Siso, STFC Daresbury Lab
+# Modified I. Kavcic, A. Coughtrie, L. Turner and O. Brunt, Met Office
+# Modified J. Henrichs, Bureau of Meteorology
+# Modified A. B. G. Chalk and N. Nobre, STFC Daresbury Lab
+
+'''
+This module contains the LFRicDofmaps class which holds all the
+information and methods for the dofmaps required by an invoke such as:
+generating the calls to the LFRic infrastructure that look-up the necessary
+dofmaps; declaring all unique function space dofmaps in the PSy layer as
+pointers to integer arrays; and adding dofmap-related declarations to a
+Kernel stub.
+
+LFRicDofmaps is used in the LFRicInvoke module.
+'''
+
+from collections import OrderedDict
+
+from psyclone import psyGen
+from psyclone.configuration import Config
+from psyclone.domain.lfric import LFRicCollection
+from psyclone.errors import GenerationError, InternalError
+from psyclone.f2pygen import AssignGen, CommentGen, DeclGen
+
+
+class LFRicDofmaps(LFRicCollection):
+    '''
+    Holds all information on the dofmaps (including column-banded and
+    indirection) required by an invoke.
+
+    :param node: Kernel or Invoke for which to manage dofmaps.
+    :type node: :py:class:`psyclone.domain.lfric.LFRicKern` or \
+                :py:class:`psyclone.dynamo0p3.LFRicInvoke`
+
+    '''
+    def __init__(self, node):
+        # pylint: disable=too-many-branches
+        super().__init__(node)
+
+        # Look at every kernel call in this invoke and generate a list
+        # of the unique function spaces involved.
+        # We create a dictionary whose keys are the map names and entries
+        # are the corresponding field objects.
+        self._unique_fs_maps = OrderedDict()
+        # We also create a dictionary of column-banded dofmaps. Entries
+        # in this one are themselves dictionaries containing two entries:
+        # "argument" - the object holding information on the CMA kernel
+        #              argument
+        # "direction" - whether the dofmap is required for the "to" or
+        #               "from" function space of the operator.
+        self._unique_cbanded_maps = OrderedDict()
+        # A dictionary of required CMA indirection dofmaps. As with the
+        # column-banded dofmaps, each entry is itself a dictionary with
+        # "argument" and "direction" entries.
+        self._unique_indirection_maps = OrderedDict()
+
+        for call in self._calls:
+            # We only need a dofmap if the kernel operates on a cell_column
+            # or the domain.
+            if call.iterates_over in ["cell_column", "domain"]:
+                for unique_fs in call.arguments.unique_fss:
+                    # We only need a dofmap if there is a *field* on this
+                    # function space. If there is then we use it to look
+                    # up the dofmap.
+                    fld_arg = unique_fs.field_on_space(call.arguments)
+                    if fld_arg:
+                        map_name = unique_fs.map_name
+                        if map_name not in self._unique_fs_maps:
+                            self._unique_fs_maps[map_name] = fld_arg
+                if call.cma_operation == "assembly":
+                    # A kernel that assembles a CMA operator requires
+                    # column-banded dofmaps for its 'to' and 'from'
+                    # function spaces
+                    cma_args = psyGen.args_filter(
+                        call.arguments.args,
+                        arg_types=["gh_columnwise_operator"])
+
+                    # Sanity check - we expect only one CMA argument
+                    if len(cma_args) != 1:
+                        raise GenerationError(
+                            f"Internal error: there should only be one CMA "
+                            f"operator argument for a CMA assembly kernel but "
+                            f"found {len(cma_args)}")
+
+                    map_name = \
+                        cma_args[0].function_space_to.cbanded_map_name
+                    if map_name not in self._unique_cbanded_maps:
+                        self._unique_cbanded_maps[map_name] = {
+                            "argument": cma_args[0],
+                            "direction": "to"}
+                    map_name = \
+                        cma_args[0].function_space_from.cbanded_map_name
+                    if map_name not in self._unique_cbanded_maps:
+                        self._unique_cbanded_maps[map_name] = {
+                            "argument": cma_args[0],
+                            "direction": "from"}
+                elif call.cma_operation == "apply":
+                    # A kernel that applies (or applies the inverse of) a
+                    # CMA operator requires the indirection dofmaps for the
+                    # to- and from-spaces of the operator.
+                    cma_args = psyGen.args_filter(
+                        call.arguments.args,
+                        arg_types=["gh_columnwise_operator"])
+
+                    # Sanity check - we expect only one CMA argument
+                    if len(cma_args) != 1:
+                        raise GenerationError(
+                            f"Internal error: there should only be one CMA "
+                            f"operator argument for a kernel that applies a "
+                            f"CMA operator but found {len(cma_args)}")
+
+                    map_name = cma_args[0].function_space_to\
+                        .cma_indirection_map_name
+                    if map_name not in self._unique_indirection_maps:
+                        self._unique_indirection_maps[map_name] = {
+                            "argument": cma_args[0],
+                            "direction": "to"}
+                    map_name = cma_args[0].function_space_from\
+                        .cma_indirection_map_name
+                    if map_name not in self._unique_indirection_maps:
+                        self._unique_indirection_maps[map_name] = {
+                            "argument": cma_args[0],
+                            "direction": "from"}
+
+    def initialise(self, parent):
+        ''' Generates the calls to the LFRic infrastructure that
+        look-up the necessary dofmaps. Adds these calls as children
+        of the supplied parent node. This must be an appropriate
+        f2pygen object. '''
+
+        # If we've got no dofmaps then we do nothing
+        if self._unique_fs_maps:
+            parent.add(CommentGen(parent, ""))
+            parent.add(CommentGen(parent,
+                                  " Look-up dofmaps for each function space"))
+            parent.add(CommentGen(parent, ""))
+
+            for dmap, field in self._unique_fs_maps.items():
+                parent.add(AssignGen(parent, pointer=True, lhs=dmap,
+                                     rhs=field.proxy_name_indexed +
+                                     "%" + field.ref_name() +
+                                     "%get_whole_dofmap()"))
+        if self._unique_cbanded_maps:
+            parent.add(CommentGen(parent, ""))
+            parent.add(CommentGen(parent,
+                                  " Look-up required column-banded dofmaps"))
+            parent.add(CommentGen(parent, ""))
+
+            for dmap, cma in self._unique_cbanded_maps.items():
+                parent.add(AssignGen(parent, pointer=True, lhs=dmap,
+                                     rhs=cma["argument"].proxy_name_indexed +
+                                     "%column_banded_dofmap_" +
+                                     cma["direction"]))
+
+        if self._unique_indirection_maps:
+            parent.add(CommentGen(parent, ""))
+            parent.add(CommentGen(parent,
+                                  " Look-up required CMA indirection dofmaps"))
+            parent.add(CommentGen(parent, ""))
+
+            for dmap, cma in self._unique_indirection_maps.items():
+                parent.add(AssignGen(parent, pointer=True, lhs=dmap,
+                                     rhs=cma["argument"].proxy_name_indexed +
+                                     "%indirection_dofmap_"+cma["direction"]))
+
+    def _invoke_declarations(self, parent):
+        '''
+        Declare all unique function space dofmaps in the PSy layer as pointers
+        to integer arrays of rank 2.
+
+        :param parent: the f2pygen node to which to add the declarations.
+        :type parent: :py:class:`psyclone.f2pygen.SubroutineGen`
+
+        '''
+        api_config = Config.get().api_conf("dynamo0.3")
+
+        # Function space dofmaps
+        decl_map_names = \
+            [dmap+"(:,:) => null()" for dmap in sorted(self._unique_fs_maps)]
+
+        if decl_map_names:
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               pointer=True, entity_decls=decl_map_names))
+
+        # Column-banded dofmaps
+        decl_bmap_names = \
+            [dmap+"(:,:) => null()" for dmap in self._unique_cbanded_maps]
+        if decl_bmap_names:
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               pointer=True, entity_decls=decl_bmap_names))
+
+        # CMA operator indirection dofmaps
+        decl_ind_map_names = \
+            [dmap+"(:) => null()" for dmap in self._unique_indirection_maps]
+        if decl_ind_map_names:
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               pointer=True, entity_decls=decl_ind_map_names))
+
+    def _stub_declarations(self, parent):
+        '''
+        Add dofmap-related declarations to a Kernel stub.
+
+        :param parent: node in the f2pygen AST representing the Kernel stub.
+        :type parent: :py:class:`psyclone.f2pygen.SubroutineGen`
+
+        '''
+        api_config = Config.get().api_conf("dynamo0.3")
+
+        # Function space dofmaps
+        for dmap in sorted(self._unique_fs_maps):
+            # We declare ndf first as some compilers require this
+            ndf_name = \
+                self._unique_fs_maps[dmap].function_space.ndf_name
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               intent="in", entity_decls=[ndf_name]))
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               intent="in", dimension=ndf_name,
+                               entity_decls=[dmap]))
+        # Column-banded dofmaps
+        for dmap, cma in self._unique_cbanded_maps.items():
+            if cma["direction"] == "to":
+                ndf_name = cma["argument"].function_space_to.ndf_name
+            elif cma["direction"] == "from":
+                ndf_name = cma["argument"].function_space_from.ndf_name
+            else:
+                raise InternalError(
+                    f"Invalid direction ('{cma['''direction''']}') found for "
+                    f"CMA operator when collecting column-banded dofmaps. "
+                    f"Should be either 'to' or 'from'.")
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               intent="in", entity_decls=[ndf_name]))
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               intent="in",
+                               dimension=",".join([ndf_name, "nlayers"]),
+                               entity_decls=[dmap]))
+        # CMA operator indirection dofmaps
+        for dmap, cma in self._unique_indirection_maps.items():
+            if cma["direction"] == "to":
+                dim_name = cma["argument"].name + "_nrow"
+            elif cma["direction"] == "from":
+                dim_name = cma["argument"].name + "_ncol"
+            else:
+                raise InternalError(
+                    f"Invalid direction ('{cma['''direction''']}') found for "
+                    f"CMA operator when collecting indirection dofmaps. "
+                    f"Should be either 'to' or 'from'.")
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               intent="in", entity_decls=[dim_name]))
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               intent="in", dimension=dim_name,
+                               entity_decls=[dmap]))
+
+
+# The list of module members that we wish AutoAPI to generate
+# documentation for. (See https://psyclone-ref.readthedocs.io)
+__all__ = ['LFRicDofmaps']
diff --git a/src/psyclone/domain/lfric/lfric_invoke.py b/src/psyclone/domain/lfric/lfric_invoke.py
index 6b2f34bbf..0e14a9582 100644
--- a/src/psyclone/domain/lfric/lfric_invoke.py
+++ b/src/psyclone/domain/lfric/lfric_invoke.py
@@ -69,6 +69,7 @@ class LFRicInvoke(Invoke):
 
     '''
     # pylint: disable=too-many-instance-attributes
+    # pylint: disable=too-many-locals
     def __init__(self, alg_invocation, idx, invokes):
         if not alg_invocation and not idx:
             # This 'if' test is added to support pyreverse
@@ -93,16 +94,15 @@ class LFRicInvoke(Invoke):
 
         # Import here to avoid circular dependency
         # pylint: disable=import-outside-toplevel
-        from psyclone.dynamo0p3 import (DynFunctionSpaces, DynDofmaps,
-                                        DynLMAOperators, DynGlobalSum,
+        from psyclone.dynamo0p3 import (DynFunctionSpaces, DynGlobalSum,
+                                        DynLMAOperators, DynReferenceElement,
                                         DynCMAOperators, DynBasisFunctions,
                                         DynMeshes, DynBoundaryConditions,
                                         DynProxies, DynCellIterators,
-                                        DynReferenceElement,
                                         LFRicMeshProperties)
         from psyclone.domain.lfric import (LFRicLoopBounds, LFRicRunTimeChecks,
                                            LFRicScalarArgs, LFRicFields,
-                                           LFRicStencils)
+                                           LFRicDofmaps, LFRicStencils)
 
         self.scalar_args = LFRicScalarArgs(self)
 
@@ -114,7 +114,7 @@ class LFRicInvoke(Invoke):
 
         # Initialise the object holding all information on the dofmaps
         # required by this Invoke
-        self.dofmaps = DynDofmaps(self)
+        self.dofmaps = LFRicDofmaps(self)
 
         # Initialise information on all of the fields accessed in this Invoke
         self.fields = LFRicFields(self)
diff --git a/src/psyclone/domain/lfric/lfric_kern.py b/src/psyclone/domain/lfric/lfric_kern.py
index 2f1fd28ca..0196c08c2 100644
--- a/src/psyclone/domain/lfric/lfric_kern.py
+++ b/src/psyclone/domain/lfric/lfric_kern.py
@@ -233,6 +233,7 @@ class LFRicKern(CodedKern):
 
     def _setup(self, ktype, module_name, args, parent, check=True):
         # pylint: disable=too-many-arguments
+        # pylint: disable=too-many-branches, too-many-locals
         '''Internal setup of kernel information.
 
         :param ktype: object holding information on the parsed metadata for \
@@ -253,7 +254,6 @@ class LFRicKern(CodedKern):
         # Import here to avoid circular dependency
         # pylint: disable=import-outside-toplevel
         from psyclone.dynamo0p3 import DynKernelArguments, FSDescriptors
-        # pylint: disable=too-many-branches, too-many-locals
         super().__init__(DynKernelArguments,
                          KernelCall(module_name, ktype, args),
                          parent, check)
@@ -609,14 +609,13 @@ class LFRicKern(CodedKern):
         # Add all the declarations
         # Import here to avoid circular dependency
         # pylint: disable=import-outside-toplevel
-        from psyclone.domain.lfric import (LFRicFields, LFRicScalarArgs,
-                                           LFRicStencils)
-        from psyclone.dynamo0p3 import (DynCellIterators, DynDofmaps,
-                                        DynFunctionSpaces, DynCMAOperators,
-                                        DynBoundaryConditions,
+        from psyclone.domain.lfric import (LFRicScalarArgs, LFRicFields,
+                                           LFRicDofmaps, LFRicStencils)
+        from psyclone.dynamo0p3 import (DynCellIterators, DynFunctionSpaces,
+                                        DynCMAOperators, DynBoundaryConditions,
                                         DynLMAOperators, LFRicMeshProperties,
                                         DynBasisFunctions, DynReferenceElement)
-        for entities in [DynCellIterators, DynDofmaps, DynFunctionSpaces,
+        for entities in [DynCellIterators, LFRicDofmaps, DynFunctionSpaces,
                          DynCMAOperators, LFRicScalarArgs, LFRicFields,
                          DynLMAOperators, LFRicStencils, DynBasisFunctions,
                          DynBoundaryConditions, DynReferenceElement,
diff --git a/src/psyclone/dynamo0p3.py b/src/psyclone/dynamo0p3.py
index 77c81bcdf..c9045fdea 100644
--- a/src/psyclone/dynamo0p3.py
+++ b/src/psyclone/dynamo0p3.py
@@ -1180,243 +1180,6 @@ class DynReferenceElement(LFRicCollection):
                     f"faces({self._face_out_normals_symbol.name})"))
 
 
-class DynDofmaps(LFRicCollection):
-    '''
-    Holds all information on the dofmaps (including column-banded and
-    indirection) required by an invoke.
-
-    :param node: Kernel or Invoke for which to manage dofmaps.
-    :type node: :py:class:`psyclone.domain.lfric.LFRicKern` or \
-                :py:class:`psyclone.dynamo0p3.LFRicInvoke`
-
-    '''
-    def __init__(self, node):
-        # pylint: disable=too-many-branches
-        super().__init__(node)
-
-        # Look at every kernel call in this invoke and generate a list
-        # of the unique function spaces involved.
-        # We create a dictionary whose keys are the map names and entries
-        # are the corresponding field objects.
-        self._unique_fs_maps = OrderedDict()
-        # We also create a dictionary of column-banded dofmaps. Entries
-        # in this one are themselves dictionaries containing two entries:
-        # "argument" - the object holding information on the CMA kernel
-        #              argument
-        # "direction" - whether the dofmap is required for the "to" or
-        #               "from" function space of the operator.
-        self._unique_cbanded_maps = OrderedDict()
-        # A dictionary of required CMA indirection dofmaps. As with the
-        # column-banded dofmaps, each entry is itself a dictionary with
-        # "argument" and "direction" entries.
-        self._unique_indirection_maps = OrderedDict()
-
-        for call in self._calls:
-            # We only need a dofmap if the kernel operates on a cell_column
-            # or the domain.
-            if call.iterates_over in ["cell_column", "domain"]:
-                for unique_fs in call.arguments.unique_fss:
-                    # We only need a dofmap if there is a *field* on this
-                    # function space. If there is then we use it to look
-                    # up the dofmap.
-                    fld_arg = unique_fs.field_on_space(call.arguments)
-                    if fld_arg:
-                        map_name = unique_fs.map_name
-                        if map_name not in self._unique_fs_maps:
-                            self._unique_fs_maps[map_name] = fld_arg
-                if call.cma_operation == "assembly":
-                    # A kernel that assembles a CMA operator requires
-                    # column-banded dofmaps for its 'to' and 'from'
-                    # function spaces
-                    cma_args = psyGen.args_filter(
-                        call.arguments.args,
-                        arg_types=["gh_columnwise_operator"])
-
-                    # Sanity check - we expect only one CMA argument
-                    if len(cma_args) != 1:
-                        raise GenerationError(
-                            f"Internal error: there should only be one CMA "
-                            f"operator argument for a CMA assembly kernel but "
-                            f"found {len(cma_args)}")
-
-                    map_name = \
-                        cma_args[0].function_space_to.cbanded_map_name
-                    if map_name not in self._unique_cbanded_maps:
-                        self._unique_cbanded_maps[map_name] = {
-                            "argument": cma_args[0],
-                            "direction": "to"}
-                    map_name = \
-                        cma_args[0].function_space_from.cbanded_map_name
-                    if map_name not in self._unique_cbanded_maps:
-                        self._unique_cbanded_maps[map_name] = {
-                            "argument": cma_args[0],
-                            "direction": "from"}
-                elif call.cma_operation == "apply":
-                    # A kernel that applies (or applies the inverse of) a
-                    # CMA operator requires the indirection dofmaps for the
-                    # to- and from-spaces of the operator.
-                    cma_args = psyGen.args_filter(
-                        call.arguments.args,
-                        arg_types=["gh_columnwise_operator"])
-
-                    # Sanity check - we expect only one CMA argument
-                    if len(cma_args) != 1:
-                        raise GenerationError(
-                            f"Internal error: there should only be one CMA "
-                            f"operator argument for a kernel that applies a "
-                            f"CMA operator but found {len(cma_args)}")
-
-                    map_name = cma_args[0].function_space_to\
-                        .cma_indirection_map_name
-                    if map_name not in self._unique_indirection_maps:
-                        self._unique_indirection_maps[map_name] = {
-                            "argument": cma_args[0],
-                            "direction": "to"}
-                    map_name = cma_args[0].function_space_from\
-                        .cma_indirection_map_name
-                    if map_name not in self._unique_indirection_maps:
-                        self._unique_indirection_maps[map_name] = {
-                            "argument": cma_args[0],
-                            "direction": "from"}
-
-    def initialise(self, parent):
-        ''' Generates the calls to the LFRic infrastructure that
-        look-up the necessary dofmaps. Adds these calls as children
-        of the supplied parent node. This must be an appropriate
-        f2pygen object. '''
-
-        # If we've got no dofmaps then we do nothing
-        if self._unique_fs_maps:
-            parent.add(CommentGen(parent, ""))
-            parent.add(CommentGen(parent,
-                                  " Look-up dofmaps for each function space"))
-            parent.add(CommentGen(parent, ""))
-
-            for dmap, field in self._unique_fs_maps.items():
-                parent.add(AssignGen(parent, pointer=True, lhs=dmap,
-                                     rhs=field.proxy_name_indexed +
-                                     "%" + field.ref_name() +
-                                     "%get_whole_dofmap()"))
-        if self._unique_cbanded_maps:
-            parent.add(CommentGen(parent, ""))
-            parent.add(CommentGen(parent,
-                                  " Look-up required column-banded dofmaps"))
-            parent.add(CommentGen(parent, ""))
-
-            for dmap, cma in self._unique_cbanded_maps.items():
-                parent.add(AssignGen(parent, pointer=True, lhs=dmap,
-                                     rhs=cma["argument"].proxy_name_indexed +
-                                     "%column_banded_dofmap_" +
-                                     cma["direction"]))
-
-        if self._unique_indirection_maps:
-            parent.add(CommentGen(parent, ""))
-            parent.add(CommentGen(parent,
-                                  " Look-up required CMA indirection dofmaps"))
-            parent.add(CommentGen(parent, ""))
-
-            for dmap, cma in self._unique_indirection_maps.items():
-                parent.add(AssignGen(parent, pointer=True, lhs=dmap,
-                                     rhs=cma["argument"].proxy_name_indexed +
-                                     "%indirection_dofmap_"+cma["direction"]))
-
-    def _invoke_declarations(self, parent):
-        '''
-        Declare all unique function space dofmaps in the PSy layer as pointers
-        to integer arrays of rank 2.
-
-        :param parent: the f2pygen node to which to add the declarations.
-        :type parent: :py:class:`psyclone.f2pygen.SubroutineGen`
-
-        '''
-        api_config = Config.get().api_conf("dynamo0.3")
-
-        # Function space dofmaps
-        decl_map_names = \
-            [dmap+"(:,:) => null()" for dmap in sorted(self._unique_fs_maps)]
-
-        if decl_map_names:
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               pointer=True, entity_decls=decl_map_names))
-
-        # Column-banded dofmaps
-        decl_bmap_names = \
-            [dmap+"(:,:) => null()" for dmap in self._unique_cbanded_maps]
-        if decl_bmap_names:
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               pointer=True, entity_decls=decl_bmap_names))
-
-        # CMA operator indirection dofmaps
-        decl_ind_map_names = \
-            [dmap+"(:) => null()" for dmap in self._unique_indirection_maps]
-        if decl_ind_map_names:
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               pointer=True, entity_decls=decl_ind_map_names))
-
-    def _stub_declarations(self, parent):
-        '''
-        Add dofmap-related declarations to a Kernel stub.
-
-        :param parent: node in the f2pygen AST representing the Kernel stub.
-        :type parent: :py:class:`psyclone.f2pygen.SubroutineGen`
-
-        '''
-        api_config = Config.get().api_conf("dynamo0.3")
-
-        # Function space dofmaps
-        for dmap in sorted(self._unique_fs_maps):
-            # We declare ndf first as some compilers require this
-            ndf_name = \
-                self._unique_fs_maps[dmap].function_space.ndf_name
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               intent="in", entity_decls=[ndf_name]))
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               intent="in", dimension=ndf_name,
-                               entity_decls=[dmap]))
-        # Column-banded dofmaps
-        for dmap, cma in self._unique_cbanded_maps.items():
-            if cma["direction"] == "to":
-                ndf_name = cma["argument"].function_space_to.ndf_name
-            elif cma["direction"] == "from":
-                ndf_name = cma["argument"].function_space_from.ndf_name
-            else:
-                raise InternalError(
-                    f"Invalid direction ('{cma['''direction''']}') found for "
-                    f"CMA operator when collecting column-banded dofmaps. "
-                    f"Should be either 'to' or 'from'.")
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               intent="in", entity_decls=[ndf_name]))
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               intent="in",
-                               dimension=",".join([ndf_name, "nlayers"]),
-                               entity_decls=[dmap]))
-        # CMA operator indirection dofmaps
-        for dmap, cma in self._unique_indirection_maps.items():
-            if cma["direction"] == "to":
-                dim_name = cma["argument"].name + "_nrow"
-            elif cma["direction"] == "from":
-                dim_name = cma["argument"].name + "_ncol"
-            else:
-                raise InternalError(
-                    f"Invalid direction ('{cma['''direction''']}') found for "
-                    f"CMA operator when collecting indirection dofmaps. "
-                    f"Should be either 'to' or 'from'.")
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               intent="in", entity_decls=[dim_name]))
-            parent.add(DeclGen(parent, datatype="integer",
-                               kind=api_config.default_kind["integer"],
-                               intent="in", dimension=dim_name,
-                               entity_decls=[dmap]))
-
-
 class DynFunctionSpaces(LFRicCollection):
     '''
     Handles the declaration and initialisation of all function-space-related
@@ -6440,7 +6203,6 @@ class DynACCEnterDataDirective(ACCEnterDataDirective):
 __all__ = [
     'DynFuncDescriptor03',
     'DynamoPSy',
-    'DynDofmaps',
     'DynFunctionSpaces',
     'DynProxies',
     'DynCellIterators',


Test patch:
diff --git a/src/psyclone/tests/domain/lfric/lfric_dofmaps_test.py b/src/psyclone/tests/domain/lfric/lfric_dofmaps_test.py
new file mode 100644
index 000000000..0c8180160
--- /dev/null
+++ b/src/psyclone/tests/domain/lfric/lfric_dofmaps_test.py
@@ -0,0 +1,215 @@
+# -----------------------------------------------------------------------------
+# BSD 3-Clause License
+#
+# Copyright (c) 2017-2024, Science and Technology Facilities Council
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+# -----------------------------------------------------------------------------
+# Authors: R. W. Ford and A. R. Porter, STFC Daresbury Lab
+# Modified: I. Kavcic, L. Turner and O. Brunt, Met Office
+# Modified: J. Henrichs, Bureau of Meteorology
+
+''' This module tests the LFRicDofmaps class found within the LFRic
+domain. Tests here have been pulled from the dynamo0p3_cma_test.py
+test file. '''
+
+import os
+import pytest
+
+from psyclone.domain.lfric import LFRicDofmaps
+from psyclone.errors import GenerationError, InternalError
+from psyclone.f2pygen import ModuleGen
+from psyclone.gen_kernel_stub import generate
+from psyclone.parse.algorithm import parse
+from psyclone.psyGen import PSyFactory
+
+# Constants
+BASE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                         os.pardir, os.pardir, "test_files", "dynamo0p3")
+
+TEST_API = "dynamo0.3"
+
+# Error tests
+def test_lfricdofmap_stubdecln_err():
+    '''
+    Check that LFRicDofmaps._stub_declarations raises the expected errors
+    if the stored CMA information is invalid.
+
+    '''
+    _, invoke_info = parse(os.path.join(BASE_PATH,
+                                        "20.5_multi_cma_invoke.f90"),
+                           api=TEST_API)
+    psy = PSyFactory(TEST_API, distributed_memory=False).create(invoke_info)
+    dofmaps = LFRicDofmaps(psy.invokes.invoke_list[0])
+    mod = ModuleGen(name="test_module")
+    for cma in dofmaps._unique_indirection_maps.values():
+        cma["direction"] = "not-a-direction"
+    with pytest.raises(InternalError) as err:
+        dofmaps._stub_declarations(mod)
+    assert ("Invalid direction ('not-a-direction') found for CMA operator "
+            "when collecting indirection dofmaps" in str(err.value))
+    for cma in dofmaps._unique_cbanded_maps.values():
+        cma["direction"] = "not-a-direction"
+    with pytest.raises(InternalError) as err:
+        dofmaps._stub_declarations(mod)
+    assert ("Invalid direction ('not-a-direction') found for CMA operator "
+            "when collecting column-banded dofmaps" in str(err.value))
+
+
+def test_cma_asm_cbanded_dofmap_error():
+    '''
+    Check that we raise expected internal error if LFRicDofmaps
+    encounters an assembly kernel that has more than one CMA op argument.
+
+    Tests cma operations flagged as "assembly" have only one operator,
+    the "gh_columnwise_operator".
+
+    '''
+    _, invoke_info = parse(
+        os.path.join(BASE_PATH,
+                     "20.0_cma_assembly.f90"),
+        api=TEST_API)
+    psy = PSyFactory(TEST_API,
+                     distributed_memory=True).create(invoke_info)
+    invoke = psy.invokes.invoke_list[0]
+    calls = invoke.schedule.kernels()
+    # We must go in and make the internal state inconsistent in order
+    # to trigger the error. So, we set the type of all the arguments
+    # in the kernel call to be CMA operators...
+    for arg in calls[0].arguments.args:
+        arg._argument_type = 'gh_columnwise_operator'
+    with pytest.raises(GenerationError) as excinfo:
+        invoke.dofmaps.__init__(invoke)
+    assert ("Internal error: there should only be one CMA operator argument "
+            "for a CMA assembly kernel but found 2") in str(excinfo.value)
+
+
+def test_cma_apply_indirection_dofmap_error():
+    '''
+    Check that we raise expected internal error if LFRicDofmaps
+    encounters an apply kernel that has more than one CMA op argument.
+
+    Tests cma operations flagged as "apply" have only one operator,
+    the "gh_columnwise_operator".
+
+    '''
+    _, invoke_info = parse(
+        os.path.join(BASE_PATH,
+                     "20.1_cma_apply.f90"),
+        api=TEST_API)
+    psy = PSyFactory(TEST_API,
+                     distributed_memory=True).create(invoke_info)
+    invoke = psy.invokes.invoke_list[0]
+    calls = invoke.schedule.kernels()
+    # We must go in and make the internal state inconsistent in order
+    # to trigger the error. So, we set the type of all the arguments
+    # in the kernel cal to be CMA operators...
+    for arg in calls[0].arguments.args:
+        arg._argument_type = 'gh_columnwise_operator'
+    with pytest.raises(GenerationError) as excinfo:
+        invoke.dofmaps.__init__(invoke)
+    assert ("Internal error: there should only be one CMA operator argument "
+            "for a kernel that applies a CMA operator but found 3") in str(excinfo.value)
+
+# Generation tests
+def test_cbanded_test_comments():
+    '''
+    Check that LFRicDofmaps generates the correct comments for an "assembly"
+    cma operation.
+
+    '''
+    _, invoke_info = parse(os.path.join(BASE_PATH, "20.0_cma_assembly.f90"),
+                           api=TEST_API)
+    psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
+    code = str(psy.gen)
+
+    output = (
+        "      !\n"
+        "      ! Look-up required column-banded dofmaps\n"
+        "      !\n"
+    )
+
+    assert output in code
+
+
+def test_unique_fs_comments():
+    '''
+    Check that LFRicDofmaps generates the correct comments for an "apply"
+    cma operation.
+
+    '''
+    _, invoke_info = parse(os.path.join(BASE_PATH,"20.1_cma_apply.f90"),
+                           api=TEST_API)
+    psy = PSyFactory(TEST_API,distributed_memory=True).create(invoke_info)
+    code = str(psy.gen)
+
+    output = (
+        "      !\n"
+        "      ! Look-up dofmaps for each function space\n"
+        "      !\n"
+        )
+
+    assert output in code
+
+
+def test_stub_decl_dofmaps():
+    '''
+    Check that LFRicDofmaps generates the expected declarations in the stub.
+
+    '''
+
+    result = generate(os.path.join(BASE_PATH,
+                                   "columnwise_op_asm_kernel_mod.F90"),
+                      api=TEST_API)
+
+    assert ("INTEGER(KIND=i_def), intent(in) :: cma_op_2_nrow, cma_op_2_ncol"
+            in str(result))
+
+
+def test_lfricdofmaps_stub_gen():
+    '''
+    Test the kernel-stub generator for a CMA apply kernel. This has
+    two fields and one CMA operator as arguments.
+
+    '''
+    result = generate(os.path.join(BASE_PATH,
+                                   "columnwise_op_app_kernel_mod.F90"),
+                      api=TEST_API)
+
+    expected = (
+        "    SUBROUTINE columnwise_op_app_kernel_code(cell, ncell_2d, "
+        "field_1_aspc1_field_1, field_2_aspc2_field_2, cma_op_3, "
+        "cma_op_3_nrow, cma_op_3_ncol, cma_op_3_bandwidth, cma_op_3_alpha, "
+        "cma_op_3_beta, cma_op_3_gamma_m, cma_op_3_gamma_p, "
+        "ndf_aspc1_field_1, undf_aspc1_field_1, map_aspc1_field_1, "
+        "cma_indirection_map_aspc1_field_1, ndf_aspc2_field_2, "
+        "undf_aspc2_field_2, map_aspc2_field_2, "
+        "cma_indirection_map_aspc2_field_2)\n"
+        )
+    assert expected in str(result)
diff --git a/src/psyclone/tests/dynamo0p3_cma_test.py b/src/psyclone/tests/dynamo0p3_cma_test.py
index 85fceb075..fc4a50701 100644
--- a/src/psyclone/tests/dynamo0p3_cma_test.py
+++ b/src/psyclone/tests/dynamo0p3_cma_test.py
@@ -48,9 +48,7 @@ from psyclone.configuration import Config
 from psyclone.core.access_type import AccessType
 from psyclone.domain.lfric import (LFRicArgDescriptor, LFRicConstants,
                                    LFRicKernMetadata)
-from psyclone.dynamo0p3 import DynDofmaps
-from psyclone.errors import GenerationError, InternalError
-from psyclone.f2pygen import ModuleGen
+from psyclone.errors import InternalError
 from psyclone.gen_kernel_stub import generate
 from psyclone.parse.algorithm import parse
 from psyclone.parse.utils import ParseError
@@ -811,28 +809,6 @@ def test_cma_mdata_matrix_vector_error():
             "'gh_columnwise_operator * 3'." in str(excinfo.value))
 
 
-def test_cma_asm_cbanded_dofmap_error():
-    ''' Check that we raise expected internal error if DynInvokeDofmaps
-    encounters an assembly kernel that has more than one CMA op argument '''
-    _, invoke_info = parse(
-        os.path.join(BASE_PATH,
-                     "20.0_cma_assembly.f90"),
-        api=TEST_API)
-    psy = PSyFactory(TEST_API,
-                     distributed_memory=True).create(invoke_info)
-    invoke = psy.invokes.invoke_list[0]
-    calls = invoke.schedule.kernels()
-    # We must go in and make the internal state inconsistent in order
-    # to trigger the error. So, we set the type of all the arguments
-    # in the kernel cal to be CMA operators...
-    for arg in calls[0].arguments.args:
-        arg._argument_type = 'gh_columnwise_operator'
-    with pytest.raises(GenerationError) as excinfo:
-        invoke.dofmaps.__init__(invoke)
-    assert ("Internal error: there should only be one CMA operator argument "
-            "for a CMA assembly kernel but found 2") in str(excinfo.value)
-
-
 def test_cma_asm(tmpdir, dist_mem):
     ''' Test that we generate correct code for an invoke containing
     a kernel that assembles a CMA operator.
@@ -1011,29 +987,6 @@ def test_cma_asm_field_same_fs(dist_mem, tmpdir):
     assert LFRicBuild(tmpdir).code_compiles(psy)
 
 
-def test_cma_apply_indirection_dofmap_error():
-    ''' Check that we raise expected internal error if DynInvokeDofmaps
-    encounters an apply kernel that has more than one CMA op argument '''
-    _, invoke_info = parse(
-        os.path.join(BASE_PATH,
-                     "20.1_cma_apply.f90"),
-        api=TEST_API)
-    psy = PSyFactory(TEST_API,
-                     distributed_memory=True).create(invoke_info)
-    invoke = psy.invokes.invoke_list[0]
-    calls = invoke.schedule.kernels()
-    # We must go in and make the internal state inconsistent in order
-    # to trigger the error. So, we set the type of all the arguments
-    # in the kernel cal to be CMA operators...
-    for arg in calls[0].arguments.args:
-        arg._argument_type = 'gh_columnwise_operator'
-    with pytest.raises(GenerationError) as excinfo:
-        invoke.dofmaps.__init__(invoke)
-    assert ("Internal error: there should only be one CMA "
-            "operator argument for a kernel that applies a "
-            "CMA operator but found 3") in str(excinfo.value)
-
-
 def test_cma_apply(tmpdir, dist_mem):
     ''' Test that we generate correct code for a kernel that applies
     a CMA operator. '''
@@ -1352,32 +1305,8 @@ def test_cma_multi_kernel(tmpdir, dist_mem):
 
     assert LFRicBuild(tmpdir).code_compiles(psy)
 
-# Tests for the kernel-stub generator
-
-
-def test_dyndofmap_stubdecln_err():
-    ''' Check that DynDofmaps._stub_declarations raises the expected errors
-    if the stored CMA information is invalid. '''
-    _, invoke_info = parse(os.path.join(BASE_PATH,
-                                        "20.5_multi_cma_invoke.f90"),
-                           api=TEST_API)
-    psy = PSyFactory(TEST_API, distributed_memory=False).create(invoke_info)
-    dofmaps = DynDofmaps(psy.invokes.invoke_list[0])
-    mod = ModuleGen(name="test_module")
-    for cma in dofmaps._unique_indirection_maps.values():
-        cma["direction"] = "not-a-direction"
-    with pytest.raises(InternalError) as err:
-        dofmaps._stub_declarations(mod)
-    assert ("Invalid direction ('not-a-direction') found for CMA operator "
-            "when collecting indirection dofmaps" in str(err.value))
-    for cma in dofmaps._unique_cbanded_maps.values():
-        cma["direction"] = "not-a-direction"
-    with pytest.raises(InternalError) as err:
-        dofmaps._stub_declarations(mod)
-    assert ("Invalid direction ('not-a-direction') found for CMA operator "
-            "when collecting column-banded dofmaps" in str(err.value))
-
 
+# Tests for the kernel-stub generator
 def test_cma_asm_stub_gen():
     ''' Test the kernel-stub generator for CMA operator assembly.