Merge pull request #204 from mcflugen/mcflugen/add-mask-to-mesh-add-nodes

billsacks · web-flow · commit 02737916e6fe · 2024-07-29T17:08:44.000-06:00
ESMPY: Add ability to mask nodes of a Mesh

This pull request adds the ability for a user to associate a node mask with an ESMPY Mesh.

The Mesh.add_nodes method now accepts a keyword, node_mask, through which a user can pass a mask array of length n_nodes. Previously it was only possible to associate a mask with a mesh's elements. This new functionality follows the same pattern as Mesh.add_elements.

To make this work I had to add an extra argument, nodeMask, to ESMC_MeshAddNodes and, as such, modified all existing calls to this function to include the value NULL for this extra argument (i.e. no masking).
diff --git a/src/Infrastructure/Field/tests/ESMC_FieldGridRegrid2UTest.c b/src/Infrastructure/Field/tests/ESMC_FieldGridRegrid2UTest.c
@@ -300,7 +300,7 @@ int main(void){
   //EX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId, nodeCoord, nodeOwner);
+  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId, nodeCoord, nodeOwner, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
diff --git a/src/Infrastructure/Field/tests/ESMC_FieldGridRegridCsrv2UTest.c b/src/Infrastructure/Field/tests/ESMC_FieldGridRegridCsrv2UTest.c
@@ -337,7 +337,7 @@ int main(void){
   //EX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId, nodeCoord, nodeOwner);
+  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId, nodeCoord, nodeOwner, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
diff --git a/src/Infrastructure/Field/tests/ESMC_FieldGridRegridCsrvUTest.c b/src/Infrastructure/Field/tests/ESMC_FieldGridRegridCsrvUTest.c
@@ -370,7 +370,7 @@ int main(void){
   //EX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(dstmesh, num_node, nodeId, nodeCoord, nodeOwner);
+  rc = ESMC_MeshAddNodes(dstmesh, num_node, nodeId, nodeCoord, nodeOwner, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
diff --git a/src/Infrastructure/Field/tests/ESMC_FieldGridRegridParUTest.c b/src/Infrastructure/Field/tests/ESMC_FieldGridRegridParUTest.c
@@ -476,7 +476,7 @@ int main(void){
   //EX_UTest_Multi_Proc_Only
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(dstmesh, num_node, nodeId, nodeCoord, nodeOwner);
+  rc = ESMC_MeshAddNodes(dstmesh, num_node, nodeId, nodeCoord, nodeOwner, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
diff --git a/src/Infrastructure/Field/tests/ESMC_FieldGridRegridUTest.c b/src/Infrastructure/Field/tests/ESMC_FieldGridRegridUTest.c
@@ -281,7 +281,7 @@ int main(void){
   //EX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(dstmesh, num_node, nodeId, nodeCoord, nodeOwner);
+  rc = ESMC_MeshAddNodes(dstmesh, num_node, nodeId, nodeCoord, nodeOwner, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
diff --git a/src/Infrastructure/Field/tests/ESMC_FieldRegridCsrvUTest.c b/src/Infrastructure/Field/tests/ESMC_FieldRegridCsrvUTest.c
@@ -118,7 +118,7 @@ int main(void){
   //NEX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(srcmesh, num_node_s, nodeId_s, nodeCoord_s, nodeOwner_s);
+  rc = ESMC_MeshAddNodes(srcmesh, num_node_s, nodeId_s, nodeCoord_s, nodeOwner_s, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
@@ -229,7 +229,7 @@ int main(void){
   //NEX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(dstmesh, num_node_d, nodeId_d, nodeCoord_d, nodeOwner_d);
+  rc = ESMC_MeshAddNodes(dstmesh, num_node_d, nodeId_d, nodeCoord_d, nodeOwner_d, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
diff --git a/src/Infrastructure/Field/tests/ESMC_FieldRegridUTest.c b/src/Infrastructure/Field/tests/ESMC_FieldRegridUTest.c
@@ -127,7 +127,7 @@ int main(void){
   //NEX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(srcmesh, num_node_s, nodeId_s, nodeCoord_s, nodeOwner_s);
+  rc = ESMC_MeshAddNodes(srcmesh, num_node_s, nodeId_s, nodeCoord_s, nodeOwner_s, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
@@ -233,7 +233,7 @@ int main(void){
   //NEX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(dstmesh, num_node_d, nodeId_d, nodeCoord_d, nodeOwner_d);
+  rc = ESMC_MeshAddNodes(dstmesh, num_node_d, nodeId_d, nodeCoord_d, nodeOwner_d, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //--------------------------------------------------------------------------- -
 
diff --git a/src/Infrastructure/Field/tests/ESMC_FieldUTest.c b/src/Infrastructure/Field/tests/ESMC_FieldUTest.c
@@ -132,7 +132,7 @@ int main(void){
   // Add node information to the mesh
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId, nodeCoord, nodeOwner);
+  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId, nodeCoord, nodeOwner, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
diff --git a/src/Infrastructure/Mesh/include/ESMC_Mesh.h b/src/Infrastructure/Mesh/include/ESMC_Mesh.h
@@ -213,7 +213,8 @@ int ESMC_MeshAddNodes(
   int nodeCount,           // in
   int *nodeIds,            // in
   double *nodeCoords,      // in
-  int *nodeOwners          // in
+  int *nodeOwners,         // in
+  int *nodeMask            // in
 );
 
 // !RETURN VALUE:
diff --git a/src/Infrastructure/Mesh/interface/ESMC_Mesh.C b/src/Infrastructure/Mesh/interface/ESMC_Mesh.C
@@ -179,20 +179,22 @@ ESMC_Mesh ESMC_MeshCreateFromFile(const char *filename, int fileTypeFlag,
 #undef  ESMC_METHOD
 #define ESMC_METHOD "ESMC_MeshAddNodes"
 int ESMC_MeshAddNodes(ESMC_Mesh mesh, int nodeCount, int *nodeIds,
-                      double *nodeCoords, int *nodeOwners){
+                      double *nodeCoords, int *nodeOwners, int *nodeMask){
 
   // initialize return code; assume routine not implemented
   int localrc = ESMC_RC_NOT_IMPL;
   int rc = ESMC_RC_NOT_IMPL;
 
   MeshCap *mc = static_cast<MeshCap*> (mesh.ptr);
 
-  // Wrap node_owners in IntArray
+  // Wrap nodeOwners and nodeMask in IntArray
   InterArray<int> nodeOwnersIA(nodeOwners,nodeCount);
+  InterArray<int> nodeMaskIA(nodeMask, nodeCount);
 
   // call into ESMCI method
   mc->meshaddnodes(&nodeCount, nodeIds, nodeCoords, &nodeOwnersIA,
-                   NULL, &(mc->coordsys_mc), &(mc->sdim_mc), &localrc);
+                   &nodeMaskIA, &(mc->coordsys_mc), &(mc->sdim_mc), &localrc);
+
   if (ESMC_LogDefault.MsgFoundError(localrc, ESMCI_ERR_PASSTHRU, ESMC_CONTEXT,
     &rc)) return rc;
 
diff --git a/src/Infrastructure/Mesh/tests/ESMCI_MeshUTest.C b/src/Infrastructure/Mesh/tests/ESMCI_MeshUTest.C
@@ -125,7 +125,7 @@ int main(void){
   //NEX_UTest
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId_s, nodeCoord_s, nodeOwner_s);
+  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId_s, nodeCoord_s, nodeOwner_s, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
    //----------------------------------------------------------------------------
 
diff --git a/src/Infrastructure/Mesh/tests/ESMC_MeshUTest.c b/src/Infrastructure/Mesh/tests/ESMC_MeshUTest.c
@@ -381,7 +381,7 @@ int main(void){
   }
 
   // Add nodes to mesh                                                     
-  localrc = ESMC_MeshAddNodes(mesh, numNodes, nodeIds, nodeCoords, nodeOwners);
+  localrc = ESMC_MeshAddNodes(mesh, numNodes, nodeIds, nodeCoords, nodeOwners, NULL);
   if (localrc != ESMF_SUCCESS) rc=ESMF_FAILURE;
 
   // Add elements and finish the mesh   
diff --git a/src/Infrastructure/Mesh/tests/ESMC_MeshVTKUTest.c b/src/Infrastructure/Mesh/tests/ESMC_MeshVTKUTest.c
@@ -154,7 +154,7 @@ int main(void){
   // Add node information to the mesh
   strcpy(name, "MeshAddNodes");
   strcpy(failMsg, "Did not return ESMF_SUCCESS");
-  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId, nodeCoord, nodeOwner);
+  rc = ESMC_MeshAddNodes(mesh, num_node, nodeId, nodeCoord, nodeOwner, NULL);
   ESMC_Test((rc==ESMF_SUCCESS), name, failMsg, &result, __FILE__, __LINE__, 0);
   //----------------------------------------------------------------------------
 
diff --git a/src/addon/esmpy/src/esmpy/api/mesh.py b/src/addon/esmpy/src/esmpy/api/mesh.py
@@ -481,7 +481,8 @@ def add_elements(self, element_count,
     def add_nodes(self, node_count,
                   node_ids,
                   node_coords,
-                  node_owners):
+                  node_owners,
+                  node_mask=None):
         """
         Add nodes to a :class:`~esmpy.api.mesh.Mesh`, this must be done before adding elements.
 
@@ -494,6 +495,13 @@ def add_nodes(self, node_count,
             tuples into a numpy array that correspond to node_ids.
         :param ndarray node_owners: a numpy array of shape (node_count, 1) to
             specify the rank of the processor that owns each node.
+
+        *OPTIONAL:*
+
+        :param ndarray node_mask: a numpy array of shape
+            ``(node_count, 1)`` containing integer values to specify masked
+            nodes. The specific values that are masked are specified in the
+            :class:`~esmpy.api.regrid.Regrid` constructor.
         """
 
         self._node_count = node_count
@@ -509,10 +517,22 @@ def add_nodes(self, node_count,
             self._node_owners = np.array(node_owners, dtype=np.int32)
         else:
             self._node_owners = node_owners
- 
+
+        if node_mask is not None:
+            node_mask = np.asarray(node_mask).astype(np.int32, copy=False).reshape(-1)
+        self._node_mask = node_mask
+        self._mask[0] = self._node_mask
+
+
+        if node_mask is not None and len(node_mask) != node_count:
+            raise ValueError(
+                "`node_mask` must have the same length as the number of nodes"
+                f" (expected {node_count}, got {len(node_mask)})"
+            )
+
         # call into ctypes layer
         ESMP_MeshAddNodes(self, self.node_count, self.node_ids, 
-                          self.node_coords, self.node_owners)
+                          self.node_coords, self.node_owners, self._node_mask)
         # can't get the sizes until mesh is "committed" in element call
 
     def copy(self):
diff --git a/src/addon/esmpy/src/esmpy/interface/cbindings.py b/src/addon/esmpy/src/esmpy/interface/cbindings.py
@@ -1141,13 +1141,16 @@ def ESMP_MeshAddElements(mesh, elementCount,
                         '.    '+constants._errmsg)
 
 _ESMF.ESMC_MeshAddNodes.restype = ct.c_int
-_ESMF.ESMC_MeshAddNodes.argtypes = [ct.c_void_p, ct.c_int,
-                                    np.ctypeslib.ndpointer(dtype=np.int32),
-                                    np.ctypeslib.ndpointer(dtype=np.float64),
-                                    np.ctypeslib.ndpointer(dtype=np.int32)]
+_ESMF.ESMC_MeshAddNodes.argtypes = [
+    ct.c_void_p,  # mesh
+    ct.c_int,  # nodeCount
+    np.ctypeslib.ndpointer(dtype=np.int32),  # nodeIds
+    np.ctypeslib.ndpointer(dtype=np.float64),  # nodeCoords
+    np.ctypeslib.ndpointer(dtype=np.int32),  # nodeOwners
+    OptionalNumpyArrayInt32,  # nodeMask
+]
 
-def ESMP_MeshAddNodes(mesh, nodeCount,
-                      nodeIds, nodeCoords, nodeOwners):
+def ESMP_MeshAddNodes(mesh, nodeCount, nodeIds, nodeCoords, nodeOwners, nodeMask=None):
     """
     Preconditions: An ESMP_Mesh has been created.  'nodeIds' holds the
                    IDs of the nodes, 'nodeCoords' holds the coordinates
@@ -1169,9 +1172,9 @@ def ESMP_MeshAddNodes(mesh, nodeCount,
     nodeCoordsD = np.array(nodeCoords, dtype=np.float64)
     # this variant uses the ndarray.astype casting function
     nodeOwnersD = np.ndarray.astype(nodeOwners, np.int32)
-
-    rc = _ESMF.ESMC_MeshAddNodes(mesh.struct.ptr, lnc,
-                                 nodeIdsD, nodeCoordsD, nodeOwnersD)
+    rc = _ESMF.ESMC_MeshAddNodes(
+        mesh.struct.ptr, lnc, nodeIdsD, nodeCoordsD, nodeOwnersD, nodeMask
+    )
     if rc != constants._ESMP_SUCCESS:
         raise ValueError('ESMC_MeshAddNodes() failed with rc = '+str(rc)+'.    '+
                         constants._errmsg)
diff --git a/src/addon/esmpy/src/esmpy/test/test_api/test_mesh.py b/src/addon/esmpy/src/esmpy/test/test_api/test_mesh.py
@@ -1,17 +1,79 @@
-
 """
 mesh unit test file
 """
-
-import pytest
-
 import os
 
-from esmpy import *
+import numpy as np
+import pytest
+from esmpy.api.constants import _ESMF_NETCDF
+from esmpy.api.constants import _ESMF_PIO
+from esmpy.api.constants import CoordSys
+from esmpy.api.esmpymanager import local_pet
+from esmpy.api.esmpymanager import Manager
+from esmpy.api.esmpymanager import pet_count
+from esmpy.api.field import element
+from esmpy.api.field import node
+from esmpy.api.mesh import Mesh
 from esmpy.test.base import TestBase
-from esmpy.api.constants import _ESMF_NETCDF, _ESMF_PIO
-from esmpy.util.mesh_utilities import *
 from esmpy.util.cache_data import DATA_DIR
+from esmpy.util.mesh_utilities import *
+from numpy.testing import assert_array_equal
+
+from esmpy.api.constants import FileFormat
+
+# Note that the tests of adding nodes with a mask are outside the test class because some
+# are parametrized, and pytest doesn't allow parametrized test methods within a class.
+
+@pytest.mark.parametrize(
+    "mask_value",
+    (
+        None,
+        [[0, 0, 0, 0], [-1, -1, -1, -1], [1, 1, 1, 1]],
+        [[False, False, False, False], [True, False, True, False], [True, True, True, True]],
+    ),
+)
+def test_add_nodes_with_mask(mask_value):
+    """Check adding masked nodes to a mesh."""
+    if mask_value is None:
+        mask = None
+    else:
+        mask = np.array(mask_value)
+    mesh = Mesh(parametric_dim=2, spatial_dim=2, coord_sys=CoordSys.CART)
+
+    x_of_points, y_of_points = np.meshgrid([0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0])
+    xy_of_points = np.c_[x_of_points.flat, y_of_points.flat]
+
+    mesh.add_nodes(
+        len(xy_of_points),
+        np.arange(len(xy_of_points))+1, #PR204 esmpy currently uses 1-based indexing for mesh node ids
+        xy_of_points,
+        np.zeros(len(xy_of_points), dtype=int),
+        node_mask=mask,
+    )
+
+    assert mesh.mask[1] is None
+    if mask is None:
+        assert mesh.mask[0] is None
+    else:
+        assert_array_equal(mesh.mask[0], np.asarray(mask).flat)
+
+
+def test_add_nodes_with_mask_bad():
+    """Check that adding a mask with an incorrect size raises an error."""
+    mesh = Mesh(parametric_dim=2, spatial_dim=2, coord_sys=CoordSys.CART)
+
+    x_of_points, y_of_points = np.meshgrid([0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0])
+    xy_of_points = np.c_[x_of_points.flat, y_of_points.flat]
+
+    with pytest.raises(ValueError):
+        mesh.add_nodes(
+            len(xy_of_points),
+            np.arange(len(xy_of_points))+1, #PR204 esmpy currently uses 1-based indexing for mesh node ids
+            xy_of_points,
+            np.zeros(len(xy_of_points), dtype=int),
+            node_mask=[1, 2, 3],
+        )
+
 
 class TestMesh(TestBase):
 
diff --git a/src/addon/esmpy/src/esmpy/test/test_api/test_regrid.py b/src/addon/esmpy/src/esmpy/test/test_api/test_regrid.py
diff --git a/src/system_tests/ESMF_CompFortranAndC/user_CComponent.C b/src/system_tests/ESMF_CompFortranAndC/user_CComponent.C
