From fe90b8e22063b7c3ca309d60ebb8996b1c8dbd2c Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 29 Sep 2022 15:23:00 -0500
Subject: [PATCH] Adding properties for number of DAGMC cells/surfaces

---
 openmc/universe.py             | 40 ++++++++++++++++++++++++++++++++++
 tests/unit_tests/dagmc/test.py |  8 +++++++
 2 files changed, 48 insertions(+)

diff --git a/openmc/universe.py b/openmc/universe.py
index 2ca4ad388..6b5e777ec 100644
--- a/openmc/universe.py
+++ b/openmc/universe.py
@@ -730,6 +730,46 @@ class DAGMCUniverse(UniverseBase):
     def get_all_materials(self, memo=None):
         return OrderedDict()
 
+    def _n_geom_elements(self, geom_type):
+        """
+        Helper function for retrieving the number geometric entities in a DAGMC
+        file
+
+        Parameters
+        ----------
+
+        geom_type : str
+            The type of geometric entity to count. One of {'Volume', 'Surface'}. Returns
+            the runtime number of voumes in the DAGMC model (includes implicit complement).
+        """
+        cv.check_value('geometry type', geom_type, ('volume', 'surface'))
+
+        def decode_str_tag(tag_val):
+            return tag_val.tobytes().decode().replace('\x00', '')
+
+        with h5py.File(self.filename) as dagmc_file:
+            category_data = dagmc_file['tstt/tags/CATEGORY/values']
+            category_strs = map(decode_str_tag, category_data)
+            # add one assuming implicit complement doesn't exist
+            n = sum([v == geom_type.capitalize() for v in category_strs])
+
+            # check for presence of an implicit complement in the file and
+            # increment the number of cells if it doesn't exist
+            if geom_type == 'volume':
+                name_data = dagmc_file['tstt/tags/NAME/values']
+                name_strs = map(decode_str_tag, name_data)
+                if not sum(['impl_complement' in n for n in name_strs]):
+                    n += 1
+        return n
+
+    @property
+    def n_cells(self):
+        return self._n_geom_elements('volume')
+
+    @property
+    def n_surfaces(self):
+        return self._n_geom_elements('surface')
+
     def create_xml_subelement(self, xml_element, memo=None):
         if memo and self in memo:
             return
diff --git a/tests/unit_tests/dagmc/test.py b/tests/unit_tests/dagmc/test.py
index b3e2c390f..352fe8895 100644
--- a/tests/unit_tests/dagmc/test.py
+++ b/tests/unit_tests/dagmc/test.py
@@ -33,6 +33,14 @@ def dagmc_model(request):
     dagmc_universe = openmc.DAGMCUniverse('dagmc.h5m')
     model.geometry = openmc.Geometry(dagmc_universe)
 
+    # check number of surfaces and volumes for this pincell model there should
+    # be 5 volumes: two fuel regions, water, graveyard, implicit complement (the
+    # implicit complement cell is created automatically at runtime)
+    # and 21 surfaces: 3 cylinders (9 surfaces) and a bounding cubic shell
+    # (12 surfaces)
+    assert dagmc_universe.n_cells == 5
+    assert dagmc_universe.n_surfaces == 21
+
     # tally
     tally = openmc.Tally()
     tally.scores = ['total']