Merge pull request #9 from eepeterson/add_volume_groups

Add ability to create multiple groups of volumes at once

dagmc/dagnav.py

@@ -111,6 +111,31 @@ def write_file(self, filename):
         """
         self.mb.write_file(filename)
 
+    def add_groups(self, group_map):
+        """Adds groups of DAGSets to the model.
+
+        Parameters
+        ----------
+        group_map : dict
+            A dictionary whose keys are 2-tuples of (str, int) containing the
+            group name and group ID respectively and whose values are iterables
+            of DAGSet objects or DAGSet ID numbers.
+        """
+        for (group_name, group_id), dagsets in group_map.items():
+            group = Group.create(self, name=group_name, group_id=group_id)
+
+            for dagset in dagsets:
+                if isinstance(dagset, DAGSet):
+                    group.add_set(dagset)
+                else:
+                    if dagset in self.volumes:
+                        group.add_set(self.volumes[dagset])
+                    elif dagset in self.surfaces:
+                        group.add_set(self.surfaces[dagset])
+                    else:
+                        raise ValueError(f"DAGSet ID={dagset} could not be "
+                                         "found in model volumes or surfaces.")
+
 
 class DAGSet:
     """

test/test_basic.py

@@ -325,4 +325,34 @@ def test_area(request):
                  10: np.pi * 10**2,
                  11: np.pi * 10**2 }
     for surf_id, exp_area in exp_areas.items():
-        pytest.approx(model.surfaces[surf_id].area, exp_area)
+        pytest.approx(model.surfaces[surf_id].area, exp_area)
+
+
+def test_add_groups(request):
+    test_file = str(request.path.parent / 'fuel_pin.h5m')
+    model = dagmc.DAGModel(test_file)
+    volumes = model.volumes
+    surfaces = model.surfaces
+
+    for group in model.groups.values():
+        group.delete()
+
+    assert len(model.groups) == 0
+
+    group_map = {("mat:fuel", 1): [1, 2],
+                 ("mat:Graveyard", 0): [volumes[6]],
+                 ("mat:41", 2): [3],
+                 ("boundary:Reflecting", 3): [27, 28, 29],
+                 ("boundary:Vacuum", 4): [surfaces[24], surfaces[25]]
+                 }
+
+    model.add_groups(group_map)
+
+    groups = model.groups
+
+    assert len(groups) == 5
+    assert [1, 2] == sorted(groups['mat:fuel'].get_volume_ids())
+    assert [6] == groups['mat:Graveyard'].get_volume_ids()
+    assert [3] == groups['mat:41'].get_volume_ids()
+    assert [27, 28, 29] == sorted(groups['boundary:Reflecting'].get_surface_ids())
+    assert [24, 25] == sorted(groups['boundary:Vacuum'].get_surface_ids())