Split up Geometry.from_xml, rename clean_xml -> _xml

openmc/_xml.py

@@ -0,0 +1,51 @@
+def clean_indentation(element, level=0, spaces_per_level=2):
+    """
+    copy and paste from http://effbot.org/zone/elementent-lib.htm#prettyprint
+    it basically walks your tree and adds spaces and newlines so the tree is
+    printed in a nice way
+    """
+
+    i = "\n" + level*spaces_per_level*" "
+
+    if len(element):
+
+        if not element.text or not element.text.strip():
+            element.text = i + spaces_per_level*" "
+
+        if not element.tail or not element.tail.strip():
+            element.tail = i
+
+        for sub_element in element:
+            clean_indentation(sub_element, level+1, spaces_per_level)
+
+        if not sub_element.tail or not sub_element.tail.strip():
+            sub_element.tail = i
+
+    else:
+        if level and (not element.tail or not element.tail.strip()):
+            element.tail = i
+
+
+def get_text(elem, name, default=None):
+    """Retrieve text of an attribute or subelement.
+
+    Parameters
+    ----------
+    elem : xml.etree.ElementTree.Element
+        Element from which to search
+    name : str
+        Name of attribute/subelement
+    default : object
+        A defult value to return if matching attribute/subelement exists
+
+    Returns
+    -------
+    str
+        Text of attribute or subelement
+
+    """
+    if name in elem.attrib:
+        return elem.get(name, default)
+    else:
+        child = elem.find(name)
+        return child.text if child is not None else default

openmc/cell.py

@@ -13,6 +13,7 @@ import openmc
 import openmc.checkvalue as cv
 from openmc.surface import Halfspace
 from openmc.region import Region, Intersection, Complement
+from openmc._xml import get_text
 from .mixin import IDManagerMixin
 
 
@@ -522,3 +523,61 @@ class Cell(IDManagerMixin):
             element.set("rotation", ' '.join(map(str, self.rotation)))
 
         return element
+
+    @classmethod
+    def from_xml_element(cls, elem, surfaces, materials, get_universe):
+        """Generate cell from XML element
+
+        Parameters
+        ----------
+        elem : xml.etree.ElementTree.Element
+            `<cell>` element
+        surfaces : dict
+            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
+        materials : dict
+            Dictionary mapping material IDs to :class:`openmc.Material`
+            instances (defined in :math:`openmc.Geometry.from_xml`)
+        get_universe : function
+            Function returning universe (defined in
+            :meth:`openmc.Geometry.from_xml`)
+
+        Returns
+        -------
+        Cell
+            Cell instance
+
+        """
+        cell_id = int(get_text(elem, 'id'))
+        name = get_text(elem, 'name')
+        c = cls(cell_id, name)
+
+        # Assign material/distributed materials or fill
+        mat_text = get_text(elem, 'material')
+        if mat_text is not None:
+            mat_ids = mat_text.split()
+            if len(mat_ids) > 1:
+                c.fill = [materials[i] for i in mat_ids]
+            else:
+                c.fill = materials[mat_ids[0]]
+        else:
+            fill_id = int(get_text(elem, 'fill'))
+            c.fill = get_universe(fill_id)
+
+        # Assign region
+        region = get_text(elem, 'region')
+        if region is not None:
+            c.region = Region.from_expression(region, surfaces)
+
+        # Check for other attributes
+        t = get_text(elem, 'temperature')
+        if t is not None:
+            c.temperature = float(t)
+        for key in ('temperature', 'rotation', 'translation'):
+            value = get_text(elem, key)
+            if value is not None:
+                setattr(c, key, [float(x) for x in value.split()])
+
+        # Add this cell to appropriate universe
+        univ_id = int(get_text(elem, 'universe', 0))
+        get_universe(univ_id).add_cell(c)
+        return c

openmc/clean_xml.py

@@ -1,26 +0,0 @@
-def clean_xml_indentation(element, level=0, spaces_per_level=2):
-    """
-    copy and paste from http://effbot.org/zone/elementent-lib.htm#prettyprint
-    it basically walks your tree and adds spaces and newlines so the tree is
-    printed in a nice way
-    """
-
-    i = "\n" + level*spaces_per_level*" "
-
-    if len(element):
-
-        if not element.text or not element.text.strip():
-            element.text = i + spaces_per_level*" "
-
-        if not element.tail or not element.tail.strip():
-            element.tail = i
-
-        for sub_element in element:
-            clean_xml_indentation(sub_element, level+1, spaces_per_level)
-
-        if not sub_element.tail or not sub_element.tail.strip():
-            sub_element.tail = i
-
-    else:
-        if level and (not element.tail or not element.tail.strip()):
-            element.tail = i

openmc/cmfd.py

@@ -15,7 +15,7 @@ from numbers import Real, Integral
 from xml.etree import ElementTree as ET
 import sys
 
-from openmc.clean_xml import clean_xml_indentation
+from openmc._xml import clean_indentation
 from openmc.checkvalue import (check_type, check_length, check_value,
                                check_greater_than, check_less_than)
 
@@ -512,7 +512,7 @@ class CMFD(object):
         self._create_write_matrices_subelement()
 
         # Clean the indentation in the file to be user-readable
-        clean_xml_indentation(self._cmfd_file)
+        clean_indentation(self._cmfd_file)
 
         # Write the XML Tree to the cmfd.xml file
         tree = ET.ElementTree(self._cmfd_file)

openmc/data/library.py

@@ -4,7 +4,7 @@ import xml.etree.ElementTree as ET
 import h5py
 
 from openmc.mixin import EqualityMixin
-from openmc.clean_xml import clean_xml_indentation
+from openmc._xml import clean_indentation
 from openmc.checkvalue import check_type
 
 
@@ -95,7 +95,7 @@ class DataLibrary(EqualityMixin):
             lib_element.set('type', library['type'])
 
         # Clean the indentation to be user-readable
-        clean_xml_indentation(root)
+        clean_indentation(root)
 
         # Write XML file
         tree = ET.ElementTree(root)

openmc/deplete/chain.py

@@ -23,7 +23,7 @@ except ImportError:
 import scipy.sparse as sp
 
 import openmc.data
-from openmc.clean_xml import clean_xml_indentation
+from openmc._xml import clean_indentation
 from .nuclide import Nuclide, DecayTuple, ReactionTuple
 
 
@@ -356,7 +356,7 @@ class Chain(object):
         if _have_lxml:
             tree.write(str(filename), encoding='utf-8', pretty_print=True)
         else:
-            clean_xml_indentation(root_elem)
+            clean_indentation(root_elem)
             tree.write(str(filename), encoding='utf-8')
 
     def form_matrix(self, rates):

openmc/geometry.py

@@ -7,7 +7,7 @@ from xml.etree import ElementTree as ET
 import numpy as np
 
 import openmc
-from openmc.clean_xml import clean_xml_indentation
+import openmc._xml as xml
 from openmc.checkvalue import check_type
 
 
@@ -95,7 +95,7 @@ class Geometry(object):
             x.tag, int(x.get('id'))))
 
         # Clean the indentation in the file to be user-readable
-        clean_xml_indentation(root_element)
+        xml.clean_indentation(root_element)
 
         # Write the XML Tree to the geometry.xml file
         tree = ET.ElementTree(root_element)
@@ -119,14 +119,6 @@ class Geometry(object):
             Geometry object
 
         """
-        # Helper function to get value from an attribute/element
-        def get(elem, name, default=None):
-            if name in elem.attrib:
-                return elem.get(name, default)
-            else:
-                child = elem.find(name)
-                return child.text if child is not None else default
-
         # Helper function for keeping a cache of Universe instances
         universes = {}
         def get_universe(univ_id):
@@ -146,7 +138,7 @@ class Geometry(object):
             surfaces[s.id] = s
 
             # Check for periodic surface
-            other_id = get(surface, 'periodic_surface_id')
+            other_id = xml.get_text(surface, 'periodic_surface_id')
             if other_id is not None:
                 periodic[s.id] = int(other_id)
 
@@ -159,84 +151,27 @@ class Geometry(object):
         child_of = defaultdict(list)
 
         for elem in root.findall('lattice'):
-            lat_id = int(get(elem, 'id'))
-            name = get(elem, 'name')
-            lat = openmc.RectLattice(lat_id, name)
-            lat.lower_left = [float(i) for i in get(elem, 'lower_left').split()]
-            lat.pitch = [float(i) for i in get(elem, 'pitch').split()]
-            outer = get(elem, 'outer')
-            if outer is not None:
-                lat.outer = get_universe(int(outer))
+            lat = openmc.RectLattice.from_xml_element(elem, get_universe)
+            universes[lat.id] = lat
+            if lat.outer is not None:
                 child_of[lat.outer].append(lat)
-            universes[lat_id] = lat
-
-            # Get array of universes
-            dimension = get(elem, 'dimension').split()
-            shape = np.array(dimension, dtype=int)[::-1]
-            uarray = np.array([get_universe(int(i)) for i in
-                               get(elem, 'universes').split()])
-            for u in uarray:
+            for u in lat.universes.ravel():
                 child_of[u].append(lat)
-            uarray.shape = shape
-            lat.universes = uarray
 
         for elem in root.findall('hex_lattice'):
-            lat_id = int(get(elem, 'id'))
-            name = get(elem, 'name')
-            lat = openmc.HexLattice(lat_id, name)
-            lat.center = [float(i) for i in get(elem, 'center').split()]
-            lat.pitch = [float(i) for i in get(elem, 'pitch').split()]
-            outer = get(elem, 'outer')
-            if outer is not None:
-                lat.outer = get_universe(int(outer))
+            lat = openmc.HexLattice.from_xml_element(elem, get_universe)
+            universes[lat.id] = lat
+            if lat.outer is not None:
                 child_of[lat.outer].append(lat)
-            universes[lat_id] = lat
-
-            # Get nested lists of universes
-            lat._num_rings = n_rings = int(get(elem, 'n_rings'))
-            lat._num_axial = n_axial = int(get(elem, 'n_axial', 1))
-
-            # Create empty nested lists for one axial level
-            univs = [[None for _ in range(max(6*(n_rings - 1 - r), 1))]
-                     for r in range(n_rings)]
-            if n_axial > 1:
-                univs = [deepcopy(univs) for i in range(n_axial)]
-
-            # Get flat array of universes numbers
-            uarray = np.array([get_universe(int(i)) for i in
-                               get(elem, 'universes').split()])
-            for u in uarray:
-                child_of[u].append(lat)
-
-            # Fill nested lists
-            j = 0
-            for z in range(n_axial):
-                # Get list for a single axial level
-                axial_level = univs[z] if n_axial > 1 else univs
-
-                # Start iterating from top
-                x, alpha = 0, n_rings - 1
-                while True:
-                    # Set entry in list based on (x,alpha,z) coordinates
-                    _, i_ring, i_within = lat.get_universe_index((x, alpha, z))
-                    axial_level[i_ring][i_within] = uarray[j]
-
-                    # Move to the right
-                    x += 2
-                    alpha -= 1
-                    if not lat.is_valid_index((x, alpha, z)):
-                        # Move down in y direction
-                        alpha += x - 1
-                        x = 1 - x
-                        if not lat.is_valid_index((x, alpha, z)):
-                            # Move to the right
-                            x += 2
-                            alpha -= 1
-                            if not lat.is_valid_index((x, alpha, z)):
-                                # Reached the bottom
-                                break
-                    j += 1
-            lat.universes = univs
+            if lat.ndim == 2:
+                for ring in lat.universes:
+                    for u in ring:
+                        child_of[u].append(lat)
+            else:
+                for axial_slice in lat.universes:
+                    for ring in axial_slice:
+                        for u in ring:
+                            child_of[u].append(lat)
 
         # Create dictionary to easily look up materials
         if materials is None:
@@ -246,41 +181,10 @@ class Geometry(object):
         mats['void'] = None
 
         for elem in root.findall('cell'):
-            cell_id = int(get(elem, 'id'))
-            name = get(elem, 'name')
-            c = openmc.Cell(cell_id, name)
-
-            # Assign material/distributed materials or fill
-            mat_text = get(elem, 'material')
-            if mat_text is not None:
-                mat_ids = mat_text.split()
-                if len(mat_ids) > 1:
-                    c.fill = [mats[i] for i in mat_ids]
-                else:
-                    c.fill = mats[mat_ids[0]]
-            else:
-                fill_id = int(get(elem, 'fill'))
-                c.fill = get_universe(fill_id)
+            c = openmc.Cell.from_xml_element(elem, surfaces, mats, get_universe)
+            if c.fill_type in ('universe', 'lattice'):
                 child_of[c.fill].append(c)
 
-            # Assign region
-            region = get(elem, 'region')
-            if region is not None:
-                c.region = openmc.Region.from_expression(region, surfaces)
-
-            # Check for other attributes
-            t = get(elem, 'temperature')
-            if t is not None:
-                c.temperature = float(t)
-            for key in ('temperature', 'rotation', 'translation'):
-                value = get(elem, key)
-                if value is not None:
-                    setattr(c, key, [float(x) for x in value.split()])
-
-            # Add this cell to appropriate universe
-            univ_id = int(get(elem, 'universe', 0))
-            get_universe(univ_id).add_cell(c)
-
         # Determine which universe is the root by finding one which is not a
         # child of any other object
         for u in universes.values():

openmc/lattice.py

@@ -10,6 +10,7 @@ import numpy as np
 
 import openmc.checkvalue as cv
 import openmc
+from openmc._xml import get_text
 from openmc.mixin import IDManagerMixin
 
 
@@ -768,6 +769,42 @@ class RectLattice(Lattice):
         # Append the XML subelement for this Lattice to the XML element
         xml_element.append(lattice_subelement)
 
+    @classmethod
+    def from_xml_element(cls, elem, get_universe):
+        """Generate rectangular lattice from XML element
+
+        Parameters
+        ----------
+        elem : xml.etree.ElementTree.Element
+            `<lattice>` element
+        get_universe : function
+            Function returning universe (defined in
+            :meth:`openmc.Geometry.from_xml`)
+
+        Returns
+        -------
+        RectLattice
+            Rectangular lattice
+
+        """
+        lat_id = int(get_text(elem, 'id'))
+        name = get_text(elem, 'name')
+        lat = cls(lat_id, name)
+        lat.lower_left = [float(i) for i in get_text(elem, 'lower_left').split()]
+        lat.pitch = [float(i) for i in get_text(elem, 'pitch').split()]
+        outer = get_text(elem, 'outer')
+        if outer is not None:
+            lat.outer = get_universe(int(outer))
+
+        # Get array of universes
+        dimension = get_text(elem, 'dimension').split()
+        shape = np.array(dimension, dtype=int)[::-1]
+        uarray = np.array([get_universe(int(i)) for i in
+                            get_text(elem, 'universes').split()])
+        uarray.shape = shape
+        lat.universes = uarray
+        return lat
+
 
 class HexLattice(Lattice):
     r"""A lattice consisting of hexagonal prisms.
@@ -1207,6 +1244,78 @@ class HexLattice(Lattice):
         # Append the XML subelement for this Lattice to the XML element
         xml_element.append(lattice_subelement)
 
+    @classmethod
+    def from_xml_element(cls, elem, get_universe):
+        """Generate hexagonal lattice from XML element
+
+        Parameters
+        ----------
+        elem : xml.etree.ElementTree.Element
+            `<hex_lattice>` element
+        get_universe : function
+            Function returning universe (defined in
+            :meth:`openmc.Geometry.from_xml`)
+
+        Returns
+        -------
+        HexLattice
+            Hexagonal lattice
+
+        """
+        lat_id = int(get_text(elem, 'id'))
+        name = get_text(elem, 'name')
+        lat = cls(lat_id, name)
+        lat.center = [float(i) for i in get_text(elem, 'center').split()]
+        lat.pitch = [float(i) for i in get_text(elem, 'pitch').split()]
+        outer = get_text(elem, 'outer')
+        if outer is not None:
+            lat.outer = get_universe(int(outer))
+
+        # Get nested lists of universes
+        lat._num_rings = n_rings = int(get_text(elem, 'n_rings'))
+        lat._num_axial = n_axial = int(get_text(elem, 'n_axial', 1))
+
+        # Create empty nested lists for one axial level
+        univs = [[None for _ in range(max(6*(n_rings - 1 - r), 1))]
+                    for r in range(n_rings)]
+        if n_axial > 1:
+            univs = [deepcopy(univs) for i in range(n_axial)]
+
+        # Get flat array of universes numbers
+        uarray = np.array([get_universe(int(i)) for i in
+                            get_text(elem, 'universes').split()])
+
+        # Fill nested lists
+        j = 0
+        for z in range(n_axial):
+            # Get list for a single axial level
+            axial_level = univs[z] if n_axial > 1 else univs
+
+            # Start iterating from top
+            x, alpha = 0, n_rings - 1
+            while True:
+                # Set entry in list based on (x,alpha,z) coordinates
+                _, i_ring, i_within = lat.get_universe_index((x, alpha, z))
+                axial_level[i_ring][i_within] = uarray[j]
+
+                # Move to the right
+                x += 2
+                alpha -= 1
+                if not lat.is_valid_index((x, alpha, z)):
+                    # Move down in y direction
+                    alpha += x - 1
+                    x = 1 - x
+                    if not lat.is_valid_index((x, alpha, z)):
+                        # Move to the right
+                        x += 2
+                        alpha -= 1
+                        if not lat.is_valid_index((x, alpha, z)):
+                            # Reached the bottom
+                            break
+                j += 1
+        lat.universes = univs
+        return lat
+
     def _repr_axial_slice(self, universes):
         """Return string representation for the given 2D group of universes.
 

openmc/material.py

@@ -9,7 +9,7 @@ import numpy as np
 import openmc
 import openmc.data
 import openmc.checkvalue as cv
-from openmc.clean_xml import clean_xml_indentation
+from openmc._xml import clean_indentation
 from .mixin import IDManagerMixin
 
 
@@ -1076,7 +1076,7 @@ class Materials(cv.CheckedList):
         self._create_material_subelements(root_element)
 
         # Clean the indentation in the file to be user-readable
-        clean_xml_indentation(root_element)
+        clean_indentation(root_element)
 
         # Write the XML Tree to the materials.xml file
         tree = ET.ElementTree(root_element)

openmc/plots.py

@@ -9,7 +9,7 @@ import numpy as np
 
 import openmc
 import openmc.checkvalue as cv
-from openmc.clean_xml import clean_xml_indentation
+from openmc._xml import clean_indentation
 from openmc.mixin import IDManagerMixin
 
 
@@ -816,7 +816,7 @@ class Plots(cv.CheckedList):
         self._create_plot_subelements()
 
         # Clean the indentation in the file to be user-readable
-        clean_xml_indentation(self._plots_file)
+        clean_indentation(self._plots_file)
 
         # Write the XML Tree to the plots.xml file
         tree = ET.ElementTree(self._plots_file)

openmc/settings.py

@@ -6,7 +6,7 @@ import sys
 
 import numpy as np
 
-from openmc.clean_xml import clean_xml_indentation
+from openmc._xml import clean_indentation
 import openmc.checkvalue as cv
 from openmc import VolumeCalculation, Source, Mesh
 
@@ -994,7 +994,7 @@ class Settings(object):
         self._create_log_grid_bins_subelement(root_element)
 
         # Clean the indentation in the file to be user-readable
-        clean_xml_indentation(root_element)
+        clean_indentation(root_element)
 
         # Write the XML Tree to the settings.xml file
         tree = ET.ElementTree(root_element)

openmc/tallies.py

@@ -15,7 +15,7 @@ import h5py
 
 import openmc
 import openmc.checkvalue as cv
-from openmc.clean_xml import clean_xml_indentation
+from openmc._xml import clean_indentation
 from .mixin import IDManagerMixin
 
 
@@ -3187,7 +3187,7 @@ class Tallies(cv.CheckedList):
         self._create_derivative_subelements(root_element)
 
         # Clean the indentation in the file to be user-readable
-        clean_xml_indentation(root_element)
+        clean_indentation(root_element)
 
         # Write the XML Tree to the tallies.xml file
         tree = ET.ElementTree(root_element)