From 5bb4b4bbdfb6c69917dad72a4f3c6509612d896b Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Thu, 20 Jan 2022 20:23:27 -0600
Subject: [PATCH] Implement Plots.from_xml classmethod

---
 openmc/model/model.py |  15 +++---
 openmc/plots.py       | 103 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 112 insertions(+), 6 deletions(-)

diff --git a/openmc/model/model.py b/openmc/model/model.py
index 5c9fefda7..cf082407c 100644
--- a/openmc/model/model.py
+++ b/openmc/model/model.py
@@ -5,7 +5,6 @@ import os
 from pathlib import Path
 from numbers import Integral
 from tempfile import NamedTemporaryFile
-import time
 
 import h5py
 
@@ -204,11 +203,10 @@ class Model:
 
     @classmethod
     def from_xml(cls, geometry='geometry.xml', materials='materials.xml',
-                 settings='settings.xml', tallies='tallies.xml'):
+                 settings='settings.xml', tallies='tallies.xml',
+                 plots='plots.xml'):
         """Create model from existing XML files
 
-        When initializing this way, the user must manually load plots.
-
         Parameters
         ----------
         geometry : str
@@ -220,6 +218,10 @@ class Model:
         tallies : str
             Path to tallies.xml file
 
+            .. versionadded:: 0.13.0
+        plots : str
+            Path to plots.xml file
+
             .. versionadded:: 0.13.0
 
         Returns
@@ -231,8 +233,9 @@ class Model:
         materials = openmc.Materials.from_xml(materials)
         geometry = openmc.Geometry.from_xml(geometry, materials)
         settings = openmc.Settings.from_xml(settings)
-        tallies = openmc.Tallies.from_xml(tallies)
-        return cls(geometry, materials, settings, tallies)
+        tallies = openmc.Tallies.from_xml(tallies) if Path(tallies).exists() else None
+        plots = openmc.Plots.from_xml(plots) if Path(plots).exists() else None
+        return cls(geometry, materials, settings, tallies, plots)
 
     def init_lib(self, threads=None, geometry_debug=False, restart_file=None,
                  tracks=False, output=True, event_based=None, intracomm=None):
diff --git a/openmc/plots.py b/openmc/plots.py
index 521be95ad..6c5a5d437 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -684,6 +684,85 @@ class Plot(IDManagerMixin):
 
         return element
 
+    @classmethod
+    def from_xml_element(cls, elem):
+        """Generate plot object from an XML element
+
+        Parameters
+        ----------
+        elem : xml.etree.ElementTree.Element
+            XML element
+
+        Returns
+        -------
+        openmc.Plot
+            Plot object
+
+        """
+        plot_id = int(elem.get("id"))
+        plot = cls(plot_id)
+        if "filename" in elem.keys():
+            plot.filename = elem.get("filename")
+        plot.color_by = elem.get("color_by")
+        plot.type = elem.get("type")
+        plot.basis = elem.get("basis")
+
+        # Helper function to get a tuple of values
+        def get_tuple(elem, name, dtype=int):
+            subelem = elem.find(name)
+            if subelem is not None:
+                return tuple([dtype(x) for x in subelem.text.split()])
+
+        plot.origin = get_tuple(elem, "origin", float)
+        plot.width = get_tuple(elem, "width", float)
+        plot.pixels = get_tuple(elem, "pixels")
+        plot._background = get_tuple(elem, "background")
+
+        # Set plot colors
+        colors = {}
+        for color_elem in elem.findall("color"):
+            uid = color_elem.get("id")
+            colors[uid] = tuple([int(x) for x in color_elem.get("rgb").split()])
+        # TODO: set colors (needs geometry information)
+
+        # Set masking information
+        mask_elem = elem.find("mask")
+        if mask_elem is not None:
+            mask_components = [int(x) for x in mask_elem.get("components").split()]
+            # TODO: set mask components (needs geometry information)
+            background = mask_elem.get("background")
+            if background is not None:
+                plot.mask_background = tuple([int(x) for x in background.split()])
+
+        # show overlaps
+        overlap_elem = elem.find("show_overlaps")
+        if overlap_elem is not None:
+            plot.show_overlaps = (overlap_elem.text in ('true', '1'))
+        overlap_color = get_tuple(elem, "overlap_color")
+        if overlap_color is not None:
+            plot.overlap_color = overlap_color
+
+        # Set universe level
+        level = elem.find("level")
+        if level is not None:
+            plot.level = int(level.text)
+
+        # Set meshlines
+        mesh_elem = elem.find("meshlines")
+        if mesh_elem is not None:
+            meshlines = {'type': mesh_elem.get('meshtype')}
+            if 'id' in mesh_elem.keys():
+                meshlines['id'] = int(mesh_elem.get('id'))
+            if 'linewidth' in mesh_elem.keys():
+                meshlines['linewidth'] = int(mesh_elem.get('linewidth'))
+            if 'color' in mesh_elem.keys():
+                meshlines['color'] = tuple(
+                    [int(x) for x in mesh_elem.get('color').split()]
+                )
+            plot.meshlines = meshlines
+
+        return plot
+
     def to_ipython_image(self, openmc_exec='openmc', cwd='.'):
         """Render plot as an image
 
@@ -848,3 +927,27 @@ class Plots(cv.CheckedList):
         reorder_attributes(self._plots_file)  # TODO: Remove when support is Python 3.8+
         tree = ET.ElementTree(self._plots_file)
         tree.write(str(p), xml_declaration=True, encoding='utf-8')
+
+    @classmethod
+    def from_xml(cls, path='plots.xml'):
+        """Generate plots collection from XML file
+
+        Parameters
+        ----------
+        path : str, optional
+            Path to plots XML file
+
+        Returns
+        -------
+        openmc.Plots
+            Plots collection
+
+        """
+        tree = ET.parse(path)
+        root = tree.getroot()
+
+        # Generate each plot
+        plots = cls()
+        for elem in root.findall('plot'):
+            plots.append(Plot.from_xml_element(elem))
+        return plots