From fbf7713a335def236e6a51a74adf2ddd837825f7 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Tue, 21 Aug 2018 07:47:29 -0500 Subject: [PATCH] Split up Geometry.from_xml, rename clean_xml -> _xml --- openmc/_xml.py | 51 +++++++++++++++ openmc/cell.py | 59 +++++++++++++++++ openmc/clean_xml.py | 26 -------- openmc/cmfd.py | 4 +- openmc/data/library.py | 4 +- openmc/deplete/chain.py | 4 +- openmc/geometry.py | 138 ++++++---------------------------------- openmc/lattice.py | 109 +++++++++++++++++++++++++++++++ openmc/material.py | 4 +- openmc/plots.py | 4 +- openmc/settings.py | 4 +- openmc/tallies.py | 4 +- 12 files changed, 254 insertions(+), 157 deletions(-) create mode 100644 openmc/_xml.py delete mode 100644 openmc/clean_xml.py diff --git a/openmc/_xml.py b/openmc/_xml.py new file mode 100644 index 0000000000..9c6c219e2f --- /dev/null +++ b/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 diff --git a/openmc/cell.py b/openmc/cell.py index cbfd74b21b..492bdeeb68 100644 --- a/openmc/cell.py +++ b/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 + `` 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 diff --git a/openmc/clean_xml.py b/openmc/clean_xml.py deleted file mode 100644 index d1002070b3..0000000000 --- a/openmc/clean_xml.py +++ /dev/null @@ -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 diff --git a/openmc/cmfd.py b/openmc/cmfd.py index 5444334b20..c156cd00a3 100644 --- a/openmc/cmfd.py +++ b/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) diff --git a/openmc/data/library.py b/openmc/data/library.py index 58e5e4e16a..8d821e49fd 100644 --- a/openmc/data/library.py +++ b/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) diff --git a/openmc/deplete/chain.py b/openmc/deplete/chain.py index 1826ca9ca0..4635d99d15 100644 --- a/openmc/deplete/chain.py +++ b/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): diff --git a/openmc/geometry.py b/openmc/geometry.py index 3a2c18311b..4cbc60c249 100644 --- a/openmc/geometry.py +++ b/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(): diff --git a/openmc/lattice.py b/openmc/lattice.py index ef33247fca..8bf4e47571 100644 --- a/openmc/lattice.py +++ b/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 + `` 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 + `` 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. diff --git a/openmc/material.py b/openmc/material.py index d19a5f4b99..4ba4f03063 100644 --- a/openmc/material.py +++ b/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) diff --git a/openmc/plots.py b/openmc/plots.py index 8eff78b1d9..f99392def9 100644 --- a/openmc/plots.py +++ b/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) diff --git a/openmc/settings.py b/openmc/settings.py index 7629ea9afc..efe41b2a1b 100644 --- a/openmc/settings.py +++ b/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) diff --git a/openmc/tallies.py b/openmc/tallies.py index 74f3429034..1e4b142fe1 100644 --- a/openmc/tallies.py +++ b/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)