Merge pull request #846 from tjlaboss/get_surfs

Wrote method Geometry.get_all_surfaces()

openmc/geometry.py

@@ -257,7 +257,23 @@ class Geometry(object):
                     lattices[cell.fill.id] = cell.fill
 
         return lattices
+    
+    def get_all_surfaces(self):
+        """
+        Return all surfaces used in the geometry
 
+        Returns
+        -------
+        collections.OrderedDict
+            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
+
+        """
+        surfaces = OrderedDict()
+        
+        for cell in self.get_all_cells().values():
+            surfaces = cell.region.get_surfaces(surfaces)
+        return surfaces
+                
     def get_materials_by_name(self, name, case_sensitive=False, matching=False):
         """Return a list of materials with matching names.
 

openmc/region.py

@@ -1,5 +1,5 @@
 from abc import ABCMeta, abstractmethod
-from collections import Iterable
+from collections import Iterable, OrderedDict
 
 from six import add_metaclass
 import numpy as np
@@ -46,6 +46,27 @@ class Region(object):
     def __ne__(self, other):
         return not self == other
 
+    def get_surfaces(self, surfaces=None):
+        """
+        Recursively find all the surfaces referenced by a region and return them
+
+        Parameters
+        ----------
+        surfaces: collections.OrderedDict, optional
+            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
+            
+        Returns
+        -------
+        surfaces: collections.OrderedDict
+            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
+
+        """
+        if surfaces is None:
+            surfaces = OrderedDict()
+        for region in self:
+            surfaces = region.get_surfaces(surfaces)
+        return surfaces
+
     @staticmethod
     def from_expression(expression, surfaces):
         """Generate a region given an infix expression.
@@ -431,3 +452,24 @@ class Complement(Region):
         else:
             temp_region = ~self.node
         return temp_region.bounding_box
+
+    def get_surfaces(self, surfaces=None):
+        """
+        Recursively find and return all the surfaces referenced by the node
+
+        Parameters
+        ----------
+        surfaces: collections.OrderedDict, optional
+            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
+
+        Returns
+        -------
+        surfaces: collections.OrderedDict
+            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
+
+        """
+        if surfaces is None:
+            surfaces = OrderedDict()
+        for region in self.node:
+            surfaces = region.get_surfaces(surfaces)
+        return surfaces

openmc/surface.py

@@ -1,5 +1,5 @@
 from abc import ABCMeta
-from collections import Iterable
+from collections import Iterable, OrderedDict
 from numbers import Real, Integral
 from xml.etree import ElementTree as ET
 from math import sqrt
@@ -1880,6 +1880,27 @@ class Halfspace(Region):
     def __str__(self):
         return '-' + str(self.surface.id) if self.side == '-' \
             else str(self.surface.id)
+    
+    def get_surfaces(self, surfaces=None):
+        """
+        Returns the surface that this is a halfspace of.
+    
+        Parameters
+        ----------
+        surfaces: collections.OrderedDict, optional
+            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
+    
+        Returns
+        -------
+        surfaces: collections.OrderedDict
+            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
+    
+        """
+        if surfaces is None:
+            surfaces = OrderedDict()
+        
+        surfaces[self.surface.id] = self.surface
+        return surfaces
 
 
 def get_rectangular_prism(width, height, axis='z', origin=(0., 0.),