Adding ability to tag node ids in MPhys TacsBuilder (#287)

* Adding capability to specify node IDs in `get_tagged_indices`

* Make work for numpy integer types

* Add test for `get_tagged_indices`

* black formatting

* Fix integer check in `getCompNames`

* Remove negative indices

---------

Co-authored-by: Alasdair Gray <alachris@umich.edu>

tacs/mphys/builder.py

@@ -423,7 +423,8 @@ def get_tagged_indices(self, tags):
 
         Parameters
         ----------
-        tags : list[str]
+        tags : list[str, int]
+            list of component names or node IDs to include in body
 
         Returns
         -------
@@ -436,13 +437,29 @@ def get_tagged_indices(self, tags):
             # Select all node IDs
             masked_local_nodes = np.arange(nnodes)
 
-        # Get the compIDs associated with tags
+        # Get the node IDs associated with tags
         else:
-            tagged_comps = self.fea_assembler.selectCompIDs(include=tags)
-            # Select local node IDs for tags
+            # Pick out any component names in supplied tags
+            comp_names = [comp_name for comp_name in tags if isinstance(comp_name, str)]
+            tagged_comps = self.fea_assembler.selectCompIDs(include=comp_names)
+            # Select local node IDs for components
             masked_local_nodes = self.fea_assembler.getLocalNodeIDsForComps(
                 tagged_comps
             )
+            # Pick out any node IDs in supplied tags
+            global_node_ids = [
+                comp_name
+                for comp_name in tags
+                if isinstance(comp_name, (int, np.integer))
+            ]
+            # Select local node IDs from global node IDs
+            local_node_ids = self.fea_assembler.getLocalNodeIDsFromGlobal(
+                global_node_ids, nastranOrdering=True
+            )
+            # getLocalNodeIDsFromGlobal returns -1 for nodes not on this processor, so remove those
+            local_node_ids[:] = [id for id in local_node_ids if id >= 0]
+            masked_local_nodes += local_node_ids
+            masked_local_nodes = np.unique(masked_local_nodes)
 
         # Select local node IDs and multiplier node IDs
         local_mnodes = self.fea_assembler.getLocalMultiplierNodeIDs()

tacs/pytacs.py

@@ -683,7 +683,7 @@ def getCompNames(self, compIDs=None):
         if compIDs is None:
             return copy.deepcopy(self.compDescripts)
         # Convert to list
-        elif isinstance(compIDs, int):
+        elif isinstance(compIDs, (int, np.integer)):
             compIDs = [compIDs]
         # Make sure list is flat
         else:
@@ -744,6 +744,31 @@ def getLocalNodeIDsForComps(self, compIDs):
 
         return self.meshLoader.getLocalNodeIDsForComps(compIDs)
 
+    def getLocalNodeIDsFromGlobal(self, globalIDs, nastranOrdering=False):
+        """
+        Given a list of node IDs in global (non-partitioned) ordering
+        returns the local (partitioned) node IDs on each processor.
+        If a requested node is not included on this processor,
+        an entry of -1 will be returned.
+
+        Parameters
+        ----------
+        globalIDs : int or list[int]
+            List of global node IDs.
+
+        nastranOrdering : bool
+            Flag signaling whether globalIDs is in TACS (default) or NASTRAN (grid IDs in bdf file) ordering
+            Defaults to False.
+
+        Returns
+        -------
+        localIDs : list[int]
+            List of local node IDs for each entry in globalIDs.
+            If the node is not owned by this processor, its index is filled with a value of -1.
+        """
+
+        return self.meshLoader.getLocalNodeIDsFromGlobal(globalIDs, nastranOrdering)
+
     def initialize(self, elemCallBack=None):
         """
         This is the 'last' method to be called during the setup. The
@@ -857,10 +882,12 @@ def _checkNonlinearity(self) -> bool:
         # Check if (res2-res0) - 2 * (res1 - res0) is zero (or very close to it)
         resNorm = np.real(res1.norm())
         res2.axpy(-2.0, res1)
-        if resNorm == 0.0 or (np.real(res2.norm()) / resNorm) <= self.getOption("linearityTol"):
-            return False # not nonlinear case
+        if resNorm == 0.0 or (np.real(res2.norm()) / resNorm) <= self.getOption(
+            "linearityTol"
+        ):
+            return False  # not nonlinear case
         else:
-            return True # nonlinear case
+            return True  # nonlinear case
 
     def _elemCallBackFromBDF(self):
         """

tests/integration_tests/test_mphys_struct_plate.py

@@ -109,17 +109,19 @@ def constraint_setup(scenario_name, fea_assembler, constraints):
             constr_list = [constr]
             return constr_list
 
+        tacs_builder = tacs.mphys.TacsBuilder(
+            mesh_file=bdf_file,
+            element_callback=element_callback,
+            problem_setup=problem_setup,
+            constraint_setup=constraint_setup,
+            check_partials=True,
+            coupled=True,
+            write_solution=False,
+        )
+        self.tacs_builder = tacs_builder
+
         class Top(Multipoint):
             def setup(self):
-                tacs_builder = tacs.mphys.TacsBuilder(
-                    mesh_file=bdf_file,
-                    element_callback=element_callback,
-                    problem_setup=problem_setup,
-                    constraint_setup=constraint_setup,
-                    check_partials=True,
-                    coupled=True,
-                    write_solution=False,
-                )
                 tacs_builder.initialize(self.comm)
 
                 dvs = self.add_subsystem("dvs", om.IndepVarComp(), promotes=["*"])
@@ -154,3 +156,39 @@ def setup_funcs(self):
         to test their sensitivities with respect to.
         """
         return FUNC_REFS, wrt
+
+    def test_get_tagged_indices(self):
+        """
+        Test the get_tagged_indices method
+        """
+        prob = self.setup_problem(dtype=float)
+        prob.setup()
+
+        # We want to test that:
+        # - For each comp_id, get_tagged_indices returns the same indices as `getLocalNodeIDsForComps`
+        # - For a random set of the NASTRAN node IDs, get_tagged_indices returns the corresponding local indices
+        # - For a mix of comp_ids and NASTRAN node IDs, get_tagged_indices returns the correct local indices
+        FEAAssembler = self.tacs_builder.get_fea_assembler()
+        compIDs = FEAAssembler.selectCompIDs()[0]
+        meshloader = FEAAssembler.meshLoader
+
+        for compID in compIDs:
+            trueNodeIDs = FEAAssembler.getLocalNodeIDsForComps([compID])
+            compName = FEAAssembler.getCompNames(compID)
+            taggedIndIDs = self.tacs_builder.get_tagged_indices(compName)
+            self.assertEqual(sorted(trueNodeIDs), sorted(taggedIndIDs))
+
+            nastranIDs = meshloader.getGlobalNodeIDsForComps(
+                [compID], nastranOrdering=True
+            )
+            taggedIndIDs = self.tacs_builder.get_tagged_indices(nastranIDs)
+            self.assertEqual(sorted(trueNodeIDs), sorted(taggedIndIDs))
+
+        # now test a mix of comp_ids and NASTRAN node IDs, we'll use the name of the first compID and the NASTRAN node
+        # IDs of the last compID
+        tags = FEAAssembler.getCompNames(
+            compIDs[0]
+        ) + meshloader.getGlobalNodeIDsForComps([compIDs[-1]], nastranOrdering=True)
+        trueNodeIDs = FEAAssembler.getLocalNodeIDsForComps([compIDs[0], compIDs[-1]])
+        taggedIndIDs = self.tacs_builder.get_tagged_indices(tags)
+        self.assertEqual(sorted(trueNodeIDs), sorted(taggedIndIDs))