Add a category property to DAGSet

dagmc/dagnav.py

@@ -3,9 +3,10 @@
 from functools import cached_property
 from itertools import chain
 from typing import Optional, Dict
+from warnings import warn
 
 import numpy as np
-from pymoab import core, types, rng
+from pymoab import core, types, rng, tag
 
 
 class DAGModel:
@@ -76,7 +77,7 @@ class DAGSet:
     """
     Generic functionality for a DAGMC EntitySet.
     """
-    def __init__(self, model: DAGModel, handle):
+    def __init__(self, model: DAGModel, handle: np.uint64):
         self.model = model
         self.handle = handle
 
@@ -86,23 +87,40 @@ def __eq__(self, other):
     def __repr__(self):
         return f'{type(self).__name__} {self.id}, {self.num_triangles()} triangles'
 
+    def _tag_get_data(self, tag: tag.Tag):
+        return self.model.mb.tag_get_data(tag, self.handle, flat=True)[0]
+
+    def _tag_set_data(self, tag: tag.Tag, value):
+        self.model.mb.tag_set_data(tag, self.handle, value)
+
     @property
-    def id(self):
-        """Return the DAGMC set's ID.
-        """
-        return self.model.mb.tag_get_data(self.model.id_tag, self.handle, flat=True)[0]
+    def id(self) -> int:
+        """Return the DAGMC set's ID."""
+        return self._tag_get_data(self.model.id_tag)
 
     @id.setter
-    def id(self, i):
-        """Set the DAGMC set's ID.
-        """
-        self.model.mb.tag_set_data(self.model.id_tag, self.handle, i)
+    def id(self, i: int):
+        """Set the DAGMC set's ID."""
+        self._tag_set_data(self.model.id_tag, i)
 
     @property
-    def geom_dimension(self):
-        """Return the DAGMC set's geometry dimension.
-        """
-        return self.model.mb.tag_get_data(self.model.geom_dimension_tag, self.handle, flat=True)[0]
+    def geom_dimension(self) -> int:
+        """Return the DAGMC set's geometry dimension."""
+        return self._tag_get_data(self.model.geom_dimension_tag)
+
+    @geom_dimension.setter
+    def geom_dimension(self, dimension: int):
+        self._tag_set_data(self.model.geom_dimension_tag, dimension)
+
+    @property
+    def category(self) -> str:
+        """Return the DAGMC set's category."""
+        return self._tag_get_data(self.model.category_tag)
+
+    @category.setter
+    def category(self, category: str):
+        """Set the DAGMC set's category."""
+        self._tag_set_data(self.model.category_tag, category)
 
     @abstractmethod
     def _get_triangle_sets(self):
@@ -210,6 +228,16 @@ def get_triangle_coordinate_mapping(self, compress=False):
 
 class Surface(DAGSet):
 
+    _category = 'Surface'
+    _geom_dimension = 2
+
+    def __init__(self, model: DAGModel, handle: np.uint64):
+        super().__init__(model, handle)
+        if self.geom_dimension != self._geom_dimension:
+            warn(f"DAGMC surface with global ID {self.id} does not have geom_dimension=2.")
+        if self.category != self._category:
+            warn(f"DAGMC surface with global ID {self.id} does not have category='Surface'.")
+
     def get_volumes(self):
         """Get the parent volumes of this surface.
         """
@@ -225,6 +253,16 @@ def _get_triangle_sets(self):
 
 class Volume(DAGSet):
 
+    _category: str = 'Volume'
+    _geom_dimension: int = 3
+
+    def __init__(self, model: DAGModel, handle: np.uint64):
+        super().__init__(model, handle)
+        if self.geom_dimension != self._geom_dimension:
+            warn(f"DAGMC volume with global ID {self.id} does not have geom_dimension=3.")
+        if self.category != self._category:
+            warn(f"DAGMC volume with global ID {self.id} does not have category='Volume'.")
+
     @property
     def groups(self) -> list[Group]:
         """Get list of groups containing this volume."""
@@ -272,8 +310,17 @@ def num_triangles(self):
     def _get_triangle_sets(self):
         return [s.handle for s in self.get_surfaces().values()]
 
+
 class Group(DAGSet):
 
+    _category: str = 'Group'
+    _geom_dimension: int = 4
+
+    def __init__(self, model: DAGModel, handle: np.uint64):
+        super().__init__(model, handle)
+        if self.geom_dimension != self._geom_dimension:
+            warn(f'DAGMC group with global ID {self.id} does not have geom_dimension=4.')
+
     def __contains__(self, ent_set: DAGSet):
         return any(vol.handle == ent_set.handle for vol in chain(
             self.get_volumes().values(), self.get_surfaces().values()))
@@ -373,16 +420,17 @@ def merge(self, other_group):
         other_group.handle = self.handle
 
     @classmethod
-    def create(cls, model, name=None, group_id=None) -> Group:
+    def create(cls, model: DAGModel, name: Optional[str] = None, group_id: Optional[int] = None) -> Group:
         """Create a new group instance with the given name"""
-        mb = model.mb
         # add necessary tags for this meshset to be identified as a group
-        group_handle = mb.create_meshset()
-        mb.tag_set_data(model.category_tag, group_handle, 'Group')
-        mb.tag_set_data(model.geom_dimension_tag, group_handle, 4)
-        group = cls(model, group_handle)
+        ent_set = DAGSet(model, model.mb.create_meshset())
+        ent_set.category = cls._category
+        ent_set.geom_dimension = cls._geom_dimension
+        if group_id is not None:
+            ent_set.id = group_id
+
+        # Now that entity set has proper tags, create Group, assign name, and return
+        group = cls(model, ent_set.handle)
         if name is not None:
             group.name = name
-        if group_id is not None:
-            group.id = group_id
         return group