Basic PyQT application for exploring an OpenMC model

openmc_tree.py

#!python
import sys
from PySide6.QtWidgets import (
    QApplication, QTabWidget, QVBoxLayout, QWidget, QTreeWidget, QTreeWidgetItem, QMessageBox
)
from PySide6.QtCore import Qt, QThread, Signal
import openmc


class LoadGeometryThread(QThread):
    geometry_loaded = Signal(object)

    def run(self):
        try:
            model = openmc.Model.from_model_xml()
        except OSError:
            model = openmc.Model.from_xml()
        self.geometry_loaded.emit(model.geometry)


class GeometryTreeWidget(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("OpenMC Geometry Viewer")
        self.setGeometry(300, 100, 600, 400)

        layout = QVBoxLayout()
        self.setLayout(layout)

        # Create the tab widget
        self.tabs = QTabWidget()
        self.full_view = QTreeWidget()
        self.cells_view = QTreeWidget()
        self.lattices_view = QTreeWidget()

        self.setup_tree_widget(self.full_view, ['Geometry', 'Info'])
        self.setup_tree_widget(self.cells_view, ['Cell ID', 'Details'])
        self.setup_tree_widget(self.lattices_view, ['Lattice ID', 'Details'])

        # Add tabs to the main window
        self.tabs.addTab(self.full_view, "Full Geometry")
        self.tabs.addTab(self.cells_view, "Cells")
        self.tabs.addTab(self.lattices_view, "Lattices")
        layout.addWidget(self.tabs)

    def setup_tree_widget(self, tree_widget, headers):
        tree_widget.setHeaderLabels(headers)
        tree_widget.itemDoubleClicked.connect(self.display_cell_info)

    def populate_tree(self, geometry):
        try:
            for universe_id, universe in geometry.get_all_universes().items():
                # Full Geometry tab
                universe_item = QTreeWidgetItem(self.full_view, [f"Universe {universe_id}", repr(universe)])
                self.add_cells_to_tree(universe_item, universe)

                # Cells and Lattices Views
                for cell_index, cell in enumerate(universe.cells.values()):
                    # Cells View
                    cell_item = QTreeWidgetItem(self.cells_view, [f"Cell {cell.id}", repr(cell)])
                    self.add_cell_details(cell_item, cell_index, cell)

                    # Lattices View
                    if isinstance(cell.fill, openmc.Lattice):
                        lattice_item = QTreeWidgetItem(self.lattices_view, [f"Lattice {cell.fill.id}", repr(cell.fill)])
                        self.add_lattice_details(lattice_item, cell_index, cell.fill)

            # Adjust column widths for readability
            for view in [self.full_view, self.cells_view, self.lattices_view]:
                view.resizeColumnToContents(0)
                view.resizeColumnToContents(1)

        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to populate tree: {str(e)}")

    def add_cells_to_tree(self, parent_item, universe):
        for cell_index, cell in enumerate(universe.cells.values()):
            cell_item = QTreeWidgetItem(parent_item, [f"Cell {cell.id}", repr(cell)])
            self.add_cell_details(cell_item, cell_index, cell)

    def add_cell_details(self, cell_item, cell_index, cell):
        # Alternate background colors for readability
        if cell_index % 2 == 1:
            cell_item.setBackground(0, Qt.lightGray)
            cell_item.setBackground(1, Qt.lightGray)

        try:
            if isinstance(cell.fill, openmc.Material):
                QTreeWidgetItem(cell_item, [f"Material {cell.fill.id}: {cell.fill.name}", ""])
            elif isinstance(cell.fill, openmc.Universe):
                self.add_cells_to_tree(cell_item, cell.fill)
            elif isinstance(cell.fill, openmc.Lattice):
                lattice_item = QTreeWidgetItem(cell_item, [f"Filled with Lattice {cell.fill.id}", ""])
                self.add_lattice_details(lattice_item, cell_index, cell.fill)
        except Exception as e:
            QMessageBox.warning(self, "Warning", f"Error loading cell info: {str(e)}")

    def add_lattice_details(self, lattice_item, lattice_index, lattice):
        # Lattice-specific logic if needed
        # Currently just adds a placeholder text, further expand as needed
        if lattice_index % 2 == 1:
            lattice_item.setBackground(0, Qt.lightGray)
            lattice_item.setBackground(1, Qt.lightGray)

    def display_cell_info(self, item):
        if "Cell" in item.text(0):
            try:
                cell_id = int(item.text(0).split()[1])
                cell = next(
                    (c for u in self.geometry.get_all_universes().values() for c in u.cells.values() if c.id == cell_id),
                    None
                )
                if cell:
                    surfaces_info = "\n".join(str(surface) for surface in cell.region.get_surfaces().values())
                    msg = QMessageBox(self)
                    msg.setWindowTitle(f"Cell {cell.id} Information")
                    msg.setText(f"Cell ID: {cell.id}\n\nFill:\n\n{cell.fill}\n\nSurfaces:\n{surfaces_info}")
                    msg.exec()
            except Exception as e:
                QMessageBox.critical(self, "Error", f"Could not display cell info: {str(e)}")


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = GeometryTreeWidget()

    # Load geometry in a separate thread to keep UI responsive
    geometry_thread = LoadGeometryThread()
    geometry_thread.geometry_loaded.connect(window.populate_tree)
    geometry_thread.start()

    window.show()
    sys.exit(app.exec())