1from collections import OrderedDict, defaultdict
2from collections.abc import Iterable
3from copy import deepcopy
4from pathlib import Path
5from xml.etree import ElementTree as ET
6
7import openmc
8import openmc._xml as xml
9from .checkvalue import check_type
10
11
12class Geometry:
13    """Geometry representing a collection of surfaces, cells, and universes.
14
15    Parameters
16    ----------
17    root : openmc.Universe or Iterable of openmc.Cell, optional
18        Root universe which contains all others, or an iterable of cells that
19        should be used to create a root universe.
20
21    Attributes
22    ----------
23    root_universe : openmc.Universe
24        Root universe which contains all others
25    bounding_box : 2-tuple of numpy.array
26        Lower-left and upper-right coordinates of an axis-aligned bounding box
27        of the universe.
28
29    """
30
31    def __init__(self, root=None):
32        self._root_universe = None
33        self._offsets = {}
34        if root is not None:
35            if isinstance(root, openmc.Universe):
36                self.root_universe = root
37            else:
38                univ = openmc.Universe()
39                for cell in root:
40                    univ.add_cell(cell)
41                self._root_universe = univ
42
43    @property
44    def root_universe(self):
45        return self._root_universe
46
47    @property
48    def bounding_box(self):
49        return self.root_universe.bounding_box
50
51    @root_universe.setter
52    def root_universe(self, root_universe):
53        check_type('root universe', root_universe, openmc.Universe)
54        self._root_universe = root_universe
55
56    def add_volume_information(self, volume_calc):
57        """Add volume information from a stochastic volume calculation.
58
59        Parameters
60        ----------
61        volume_calc : openmc.VolumeCalculation
62            Results from a stochastic volume calculation
63
64        """
65        if volume_calc.domain_type == 'cell':
66            for cell in self.get_all_cells().values():
67                if cell.id in volume_calc.volumes:
68                    cell.add_volume_information(volume_calc)
69        elif volume_calc.domain_type == 'material':
70            for material in self.get_all_materials().values():
71                if material.id in volume_calc.volumes:
72                    material.add_volume_information(volume_calc)
73        elif volume_calc.domain_type == 'universe':
74            for universe in self.get_all_universes().values():
75                if universe.id in volume_calc.volumes:
76                    universe.add_volume_information(volume_calc)
77
78    def export_to_xml(self, path='geometry.xml', remove_surfs=False):
79        """Export geometry to an XML file.
80
81        Parameters
82        ----------
83        path : str
84            Path to file to write. Defaults to 'geometry.xml'.
85        remove_surfs : bool
86            Whether or not to remove redundant surfaces from the geometry when
87            exporting
88
89            .. versionadded:: 0.12
90
91        """
92        # Find and remove redundant surfaces from the geometry
93        if remove_surfs:
94            self.remove_redundant_surfaces()
95
96        # Create XML representation
97        root_element = ET.Element("geometry")
98        self.root_universe.create_xml_subelement(root_element, memo=set())
99
100        # Sort the elements in the file
101        root_element[:] = sorted(root_element, key=lambda x: (
102            x.tag, int(x.get('id'))))
103
104        # Clean the indentation in the file to be user-readable
105        xml.clean_indentation(root_element)
106
107        # Check if path is a directory
108        p = Path(path)
109        if p.is_dir():
110            p /= 'geometry.xml'
111
112        # Write the XML Tree to the geometry.xml file
113        xml.reorder_attributes(root_element)  # TODO: Remove when support is Python 3.8+
114        tree = ET.ElementTree(root_element)
115        tree.write(str(p), xml_declaration=True, encoding='utf-8')
116
117    @classmethod
118    def from_xml(cls, path='geometry.xml', materials=None):
119        """Generate geometry from XML file
120
121        Parameters
122        ----------
123        path : str, optional
124            Path to geometry XML file
125        materials : openmc.Materials or None
126            Materials used to assign to cells. If None, an attempt is made to
127            generate it from the materials.xml file.
128
129        Returns
130        -------
131        openmc.Geometry
132            Geometry object
133
134        """
135        # Helper function for keeping a cache of Universe instances
136        universes = {}
137        def get_universe(univ_id):
138            if univ_id not in universes:
139                univ = openmc.Universe(univ_id)
140                universes[univ_id] = univ
141            return universes[univ_id]
142
143        tree = ET.parse(path)
144        root = tree.getroot()
145
146        # Get surfaces
147        surfaces = {}
148        periodic = {}
149        for surface in root.findall('surface'):
150            s = openmc.Surface.from_xml_element(surface)
151            surfaces[s.id] = s
152
153            # Check for periodic surface
154            other_id = xml.get_text(surface, 'periodic_surface_id')
155            if other_id is not None:
156                periodic[s.id] = int(other_id)
157
158        # Apply periodic surfaces
159        for s1, s2 in periodic.items():
160            surfaces[s1].periodic_surface = surfaces[s2]
161
162        # Dictionary that maps each universe to a list of cells/lattices that
163        # contain it (needed to determine which universe is the root)
164        child_of = defaultdict(list)
165
166        for elem in root.findall('lattice'):
167            lat = openmc.RectLattice.from_xml_element(elem, get_universe)
168            universes[lat.id] = lat
169            if lat.outer is not None:
170                child_of[lat.outer].append(lat)
171            for u in lat.universes.ravel():
172                child_of[u].append(lat)
173
174        for elem in root.findall('hex_lattice'):
175            lat = openmc.HexLattice.from_xml_element(elem, get_universe)
176            universes[lat.id] = lat
177            if lat.outer is not None:
178                child_of[lat.outer].append(lat)
179            if lat.ndim == 2:
180                for ring in lat.universes:
181                    for u in ring:
182                        child_of[u].append(lat)
183            else:
184                for axial_slice in lat.universes:
185                    for ring in axial_slice:
186                        for u in ring:
187                            child_of[u].append(lat)
188
189        # Create dictionary to easily look up materials
190        if materials is None:
191            filename = Path(path).parent / 'materials.xml'
192            materials = openmc.Materials.from_xml(str(filename))
193        mats = {str(m.id): m for m in materials}
194        mats['void'] = None
195
196        for elem in root.findall('cell'):
197            c = openmc.Cell.from_xml_element(elem, surfaces, mats, get_universe)
198            if c.fill_type in ('universe', 'lattice'):
199                child_of[c.fill].append(c)
200
201        # Determine which universe is the root by finding one which is not a
202        # child of any other object
203        for u in universes.values():
204            if not child_of[u]:
205                return cls(u)
206        else:
207            raise ValueError('Error determining root universe.')
208
209    def find(self, point):
210        """Find cells/universes/lattices which contain a given point
211
212        Parameters
213        ----------
214        point : 3-tuple of float
215            Cartesian coordinates of the point
216
217        Returns
218        -------
219        list
220            Sequence of universes, cells, and lattices which are traversed to
221            find the given point
222
223        """
224        return self.root_universe.find(point)
225
226    def get_instances(self, paths):
227        """Return the instance number(s) for a cell/material in a geometry path.
228
229        The instance numbers are used as indices into distributed
230        material/temperature arrays and tally distribcell filter arrays.
231
232        Parameters
233        ----------
234        paths : str or iterable of str
235            The path traversed through the CSG tree to reach a cell or material
236            instance. For example, 'u0->c10->l20(2,2,1)->u5->c5' would indicate
237            the cell instance whose first level is universe 0 and cell 10,
238            second level is lattice 20 position (2,2,1), and third level is
239            universe 5 and cell 5.
240
241        Returns
242        -------
243        int or list of int
244            Instance number(s) for the given path(s)
245
246        """
247        # Make sure we are working with an iterable
248        return_list = (isinstance(paths, Iterable) and
249                       not isinstance(paths, str))
250        path_list = paths if return_list else [paths]
251
252        indices = []
253        for p in path_list:
254            # Extract the cell id from the path
255            last_index = p.rfind('>')
256            last_path = p[last_index+1:]
257            uid = int(last_path[1:])
258
259            # Get corresponding cell/material
260            if last_path[0] == 'c':
261                obj = self.get_all_cells()[uid]
262            elif last_path[0] == 'm':
263                obj = self.get_all_materials()[uid]
264
265            # Determine index in paths array
266            try:
267                indices.append(obj.paths.index(p))
268            except ValueError:
269                indices.append(None)
270
271        return indices if return_list else indices[0]
272
273    def get_all_cells(self):
274        """Return all cells in the geometry.
275
276        Returns
277        -------
278        collections.OrderedDict
279            Dictionary mapping cell IDs to :class:`openmc.Cell` instances
280
281        """
282        if self.root_universe is not None:
283            return self.root_universe.get_all_cells(memo=set())
284        else:
285            return OrderedDict()
286
287    def get_all_universes(self):
288        """Return all universes in the geometry.
289
290        Returns
291        -------
292        collections.OrderedDict
293            Dictionary mapping universe IDs to :class:`openmc.Universe`
294            instances
295
296        """
297        universes = OrderedDict()
298        universes[self.root_universe.id] = self.root_universe
299        universes.update(self.root_universe.get_all_universes())
300        return universes
301
302    def get_all_materials(self):
303        """Return all materials within the geometry.
304
305        Returns
306        -------
307        collections.OrderedDict
308            Dictionary mapping material IDs to :class:`openmc.Material`
309            instances
310
311        """
312        if self.root_universe is not None:
313            return self.root_universe.get_all_materials(memo=set())
314        else:
315            return OrderedDict()
316
317    def get_all_material_cells(self):
318        """Return all cells filled by a material
319
320        Returns
321        -------
322        collections.OrderedDict
323            Dictionary mapping cell IDs to :class:`openmc.Cell` instances that
324            are filled with materials or distributed materials.
325
326        """
327        material_cells = OrderedDict()
328
329        for cell in self.get_all_cells().values():
330            if cell.fill_type in ('material', 'distribmat'):
331                if cell not in material_cells:
332                    material_cells[cell.id] = cell
333
334        return material_cells
335
336    def get_all_material_universes(self):
337        """Return all universes having at least one material-filled cell.
338
339        This method can be used to find universes that have at least one cell
340        that is filled with a material or is void.
341
342        Returns
343        -------
344        collections.OrderedDict
345            Dictionary mapping universe IDs to :class:`openmc.Universe`
346            instances with at least one material-filled cell
347
348        """
349        material_universes = OrderedDict()
350
351        for universe in self.get_all_universes().values():
352            for cell in universe.cells.values():
353                if cell.fill_type in ('material', 'distribmat', 'void'):
354                    if universe not in material_universes:
355                        material_universes[universe.id] = universe
356
357        return material_universes
358
359    def get_all_lattices(self):
360        """Return all lattices defined
361
362        Returns
363        -------
364        collections.OrderedDict
365            Dictionary mapping lattice IDs to :class:`openmc.Lattice` instances
366
367        """
368        lattices = OrderedDict()
369
370        for cell in self.get_all_cells().values():
371            if cell.fill_type == 'lattice':
372                if cell.fill.id not in lattices:
373                    lattices[cell.fill.id] = cell.fill
374
375        return lattices
376
377    def get_all_surfaces(self):
378        """
379        Return all surfaces used in the geometry
380
381        Returns
382        -------
383        collections.OrderedDict
384            Dictionary mapping surface IDs to :class:`openmc.Surface` instances
385
386        """
387        surfaces = OrderedDict()
388
389        for cell in self.get_all_cells().values():
390            if cell.region is not None:
391                surfaces = cell.region.get_surfaces(surfaces)
392        return surfaces
393
394    def get_redundant_surfaces(self):
395        """Return all of the topologically redundant surface IDs
396
397        .. versionadded:: 0.12
398
399        Returns
400        -------
401        dict
402            Dictionary whose keys are the ID of a redundant surface and whose
403            values are the topologically equivalent :class:`openmc.Surface`
404            that should replace it.
405
406        """
407        tally = defaultdict(list)
408        for surf in self.get_all_surfaces().values():
409            coeffs = tuple(surf._coefficients[k] for k in surf._coeff_keys)
410            key = (surf._type,) + coeffs
411            tally[key].append(surf)
412        return {replace.id: keep
413                for keep, *redundant in tally.values()
414                for replace in redundant}
415
416    def _get_domains_by_name(self, name, case_sensitive, matching, domain_type):
417        if not case_sensitive:
418            name = name.lower()
419
420        domains = []
421
422        func = getattr(self, 'get_all_{}s'.format(domain_type))
423        for domain in func().values():
424            domain_name = domain.name if case_sensitive else domain.name.lower()
425            if domain_name == name:
426                domains.append(domain)
427            elif not matching and name in domain_name:
428                domains.append(domain)
429
430        domains.sort(key=lambda x: x.id)
431        return domains
432
433    def get_materials_by_name(self, name, case_sensitive=False, matching=False):
434        """Return a list of materials with matching names.
435
436        Parameters
437        ----------
438        name : str
439            The name to match
440        case_sensitive : bool
441            Whether to distinguish upper and lower case letters in each
442            material's name (default is False)
443        matching : bool
444            Whether the names must match completely (default is False)
445
446        Returns
447        -------
448        list of openmc.Material
449            Materials matching the queried name
450
451        """
452        return self._get_domains_by_name(name, case_sensitive, matching, 'material')
453
454    def get_cells_by_name(self, name, case_sensitive=False, matching=False):
455        """Return a list of cells with matching names.
456
457        Parameters
458        ----------
459        name : str
460            The name to search match
461        case_sensitive : bool
462            Whether to distinguish upper and lower case letters in each
463            cell's name (default is False)
464        matching : bool
465            Whether the names must match completely (default is False)
466
467        Returns
468        -------
469        list of openmc.Cell
470            Cells matching the queried name
471
472        """
473        return self._get_domains_by_name(name, case_sensitive, matching, 'cell')
474
475    def get_cells_by_fill_name(self, name, case_sensitive=False, matching=False):
476        """Return a list of cells with fills with matching names.
477
478        Parameters
479        ----------
480        name : str
481            The name to match
482        case_sensitive : bool
483            Whether to distinguish upper and lower case letters in each
484            cell's name (default is False)
485        matching : bool
486            Whether the names must match completely (default is False)
487
488        Returns
489        -------
490        list of openmc.Cell
491            Cells with fills matching the queried name
492
493        """
494
495        if not case_sensitive:
496            name = name.lower()
497
498        cells = set()
499
500        for cell in self.get_all_cells().values():
501            names = []
502            if cell.fill_type in ('material', 'universe', 'lattice'):
503                names.append(cell.fill.name)
504            elif cell.fill_type == 'distribmat':
505                for mat in cell.fill:
506                    if mat is not None:
507                        names.append(mat.name)
508
509            for fill_name in names:
510                if not case_sensitive:
511                    fill_name = fill_name.lower()
512
513                if fill_name == name:
514                    cells.add(cell)
515                elif not matching and name in fill_name:
516                    cells.add(cell)
517
518        return sorted(cells, key=lambda x: x.id)
519
520    def get_universes_by_name(self, name, case_sensitive=False, matching=False):
521        """Return a list of universes with matching names.
522
523        Parameters
524        ----------
525        name : str
526            The name to match
527        case_sensitive : bool
528            Whether to distinguish upper and lower case letters in each
529            universe's name (default is False)
530        matching : bool
531            Whether the names must match completely (default is False)
532
533        Returns
534        -------
535        list of openmc.Universe
536            Universes matching the queried name
537
538        """
539        return self._get_domains_by_name(name, case_sensitive, matching, 'universe')
540
541    def get_lattices_by_name(self, name, case_sensitive=False, matching=False):
542        """Return a list of lattices with matching names.
543
544        Parameters
545        ----------
546        name : str
547            The name to match
548        case_sensitive : bool
549            Whether to distinguish upper and lower case letters in each
550            lattice's name (default is False)
551        matching : bool
552            Whether the names must match completely (default is False)
553
554        Returns
555        -------
556        list of openmc.Lattice
557            Lattices matching the queried name
558
559        """
560        return self._get_domains_by_name(name, case_sensitive, matching, 'lattice')
561
562    def remove_redundant_surfaces(self):
563        """Remove redundant surfaces from the geometry"""
564
565        # Get redundant surfaces
566        redundant_surfaces = self.get_redundant_surfaces()
567
568        # Iterate through all cells contained in the geometry
569        for cell in self.get_all_cells().values():
570            # Recursively remove redundant surfaces from regions
571            if cell.region:
572                cell.region.remove_redundant_surfaces(redundant_surfaces)
573
574    def determine_paths(self, instances_only=False):
575        """Determine paths through CSG tree for cells and materials.
576
577        This method recursively traverses the CSG tree to determine each unique
578        path that reaches every cell and material. The paths are stored in the
579        :attr:`Cell.paths` and :attr:`Material.paths` attributes.
580
581        Parameters
582        ----------
583        instances_only : bool, optional
584            If true, this method will only determine the number of instances of
585            each cell and material.
586
587        """
588        # (Re-)initialize all cell instances to 0
589        for cell in self.get_all_cells().values():
590            cell._paths = []
591            cell._num_instances = 0
592        for material in self.get_all_materials().values():
593            material._paths = []
594            material._num_instances = 0
595
596        # Recursively traverse the CSG tree to count all cell instances
597        self.root_universe._determine_paths(instances_only=instances_only)
598
599    def clone(self):
600        """Create a copy of this geometry with new unique IDs for all of its
601        enclosed materials, surfaces, cells, universes and lattices."""
602
603        clone = deepcopy(self)
604        clone.root_universe = self.root_universe.clone()
605        return clone
606