inp: added cell visualization

Author: BitterB0NG0

src/pymcnp/Inp.py

@@ -1,7 +1,6 @@
 import re
 
 from . import inp
-from . import _show
 from . import _file
 from . import types
 from . import errors
@@ -142,25 +141,6 @@ def to_mcnp(self):
             + (self.other.value if self.other is not None else '')
         )
 
-    def draw(self) -> _show.Shape:
-        """
-        Generates ``Visualization`` from ``Inp``.
-
-        Returns:
-            ``Visualization`` for ``Inp``.
-        """
-
-        surfaces = list(filter(lambda surface: isinstance(surface, inp.Surface), self.surfaces))
-
-        if surfaces:
-            vis = surfaces[0].draw()
-            for surface in surfaces[1:]:
-                if isinstance(surface, inp.Surface):
-                    vis += surface.draw()
-            return vis
-
-        return None
-
     @staticmethod
     def _preprocess(source: str):
         """

src/pymcnp/Visualize.py

@@ -1,8 +1,11 @@
 import os
 import pathlib
 
+import numpy
 import pyvista
 
+from . import inp
+from . import _show
 from . import _doer
 from . import errors
 from .Inp import Inp
@@ -13,24 +16,30 @@ class Visualize(_doer.Doer):
     Visualizes INP files.
 
     Attributes:
-        inp: Files to visualize.
+        inpt: Files to visualize.
     """
 
-    def __init__(self, inp: Inp):
+    GRID = pyvista.ImageData(
+        dimensions=(_show.pyvista.RESOLUTION, _show.pyvista.RESOLUTION, _show.pyvista.RESOLUTION),
+        spacing=(_show.pyvista.BOUND / _show.pyvista.RESOLUTION, _show.pyvista.BOUND / _show.pyvista.RESOLUTION, _show.pyvista.BOUND / _show.pyvista.RESOLUTION),
+        origin=(-_show.pyvista.BOUND / 2, -_show.pyvista.BOUND / 2, -_show.pyvista.BOUND / 2),
+    )
+
+    def __init__(self, inpt: Inp):
         """
         Initializes ``Visualize``.
 
         Parameters:
-            inp: Files to visualize.
+            inpt: Files to visualize.
 
         Raises:
             CliCode: RUNTIME_DOER.
         """
 
-        if inp is None:
-            raise errors.CliError(errors.CliCode.RUNTIME_DOER, inp)
+        if inpt is None:
+            raise errors.CliError(errors.CliCode.RUNTIME_DOER, inpt)
 
-        self.inp = inp
+        self.inpt = inpt
 
     def to_show_cells(self) -> pyvista.Plotter:
         """
@@ -39,8 +48,23 @@ def to_show_cells(self) -> pyvista.Plotter:
 
         plot = pyvista.Plotter()
         plot.add_axes()
-        # vis = self.inp.draw()
-        # plot.add_mesh(vis.data)
+
+        surfaces = {str(surface.number): surface.to_show() for surface in self.inpt.surfaces if isinstance(surface, inp.Surface)}
+        cells = {}
+
+        for cell in self.inpt.cells:
+            if isinstance(cell, inp.Cell):
+                shape = cell.to_show(surfaces, cells)
+            elif isinstance(cell, inp.Like):
+                shape = cells[str(cell.cell)]
+            else:
+                continue
+
+            cells[str(cell.number)] = shape
+
+            self.GRID['cell'] = shape.cell(self.GRID.points).astype(numpy.float32)
+            plot.add_volume(self.GRID, scalars='cell', opacity=[0, 0, 0.01, 0.01])
+            plot.add_mesh(shape.surface, opacity=0.9)
 
         return plot
 
@@ -51,8 +75,73 @@ def to_show_surfaces(self) -> pyvista.Plotter:
 
         plot = pyvista.Plotter()
         plot.add_axes()
-        vis = self.inp.draw()
-        plot.add_mesh(vis.data)
+
+        for surface in self.inpt.surfaces:
+            if not isinstance(surface, inp.Surface):
+                continue
+
+            shape = surface.to_show()
+
+            plot.add_mesh(shape.surface)
+
+        return plot
+
+    def to_show_cell(self, number: str) -> pyvista.Plotter:
+        """
+        Visualizes INP cells.
+
+        Parameters:
+            number: Cell number to visualize.
+        """
+
+        plot = pyvista.Plotter()
+        plot.add_axes()
+
+        surfaces = {str(surface.number): surface.to_show() for surface in self.inpt.surfaces if isinstance(surface, inp.Surface)}
+        cells = {}
+
+        for cell in self.inpt.cells:
+            if isinstance(cell, inp.Cell):
+                shape = cell.to_show(surfaces, cells)
+            elif isinstance(cell, inp.Like):
+                shape = cells[str(cell.cell)]
+            else:
+                continue
+
+            cells[str(cell.number)] = shape
+
+            if str(cell.number) != number:
+                continue
+
+            self.GRID['cell'] = shape.cell(self.GRID.points).astype(numpy.float32)
+            plot.add_volume(self.GRID, scalars='cell', opacity=[0, 0, 0.01, 0.01])
+            plot.add_mesh(shape.surface, opacity=0.9)
+
+            break
+        else:
+            raise errors.CliError(errors.CliCode.RUNTIME_DOER, number)
+
+        return plot
+
+    def to_show_surface(self, number: str) -> pyvista.Plotter:
+        """
+        Visualizes INP surfaces.
+
+        Parameters:
+            number: Surface number to visualize.
+        """
+
+        plot = pyvista.Plotter()
+        plot.add_axes()
+
+        for surface in self.inpt.surfaces:
+            if not isinstance(surface, inp.Surface) or str(surface.number) != number:
+                continue
+
+            plot.add_mesh(surface.to_show().surface)
+            break
+        else:
+            raise errors.CliError(errors.CliCode.RUNTIME_DOER, number)
 
         return plot
 
@@ -64,11 +153,10 @@ def to_pdf_cells(self, path: str | pathlib.Path):
             path: Path to new pdf file.
         """
 
-        # plot = pyvista.Plotter()
-        # vis = self.inp.draw()
-        # plot.add_mesh(vis.data)
+        plot = self.to_show_cells()
 
-        # plot.save_graphic(str(path))
+        if 'PYTEST_CURRENT_TEST' not in os.environ:  # pragma: no cover
+            plot.save_graphic(str(path))
 
     def to_pdf_surfaces(self, path: str | pathlib.Path):
         """
@@ -78,10 +166,7 @@ def to_pdf_surfaces(self, path: str | pathlib.Path):
             path: Path to new pdf file.
         """
 
-        plot = pyvista.Plotter()
-        plot.add_axes()
-        vis = self.inp.draw()
-        plot.add_mesh(vis.data)
+        plot = self.to_show_surfaces()
 
         if 'PYTEST_CURRENT_TEST' not in os.environ:  # pragma: no cover
             plot.save_graphic(str(path))

src/pymcnp/_cli/visualize.py

@@ -1,11 +1,13 @@
 """
 Usage:
-    pymcnp visualize <inp> [ options ]
+    pymcnp visualize <inp> [ options... ]
 
 Options:
-    -c --cells          Visualize cells.
-    -s --surfaces       Visualize surfaces.
-    --pdf               Write PDF.
+    -c --cells                Visualize all cells.
+    -s --surfaces             Visualize all surfaces.
+    -c --cell=[<number>]      Visualize cell.
+    -s --surface=[<number>]   Visualize surface.
+    --pdf                     Write PDF.
 """
 
 import os
@@ -15,6 +17,7 @@
 
 from . import _io
 from .. import errors
+from .. import inp
 from ..Inp import Inp
 from ..Visualize import Visualize
 
@@ -32,15 +35,21 @@ def main() -> None:
 
     # Reading INP.
     try:
-        inp = Inp.from_file(file)
-        visualize = Visualize(inp)
+        inpt = Inp.from_file(file)
+        visualize = Visualize(inpt)
     except errors.InpError as err:
         _io.error(str(err))
         exit(1)
     except errors.CliError as err:
         _io.error(str(err))
         exit(2)
 
+    if not (args['--cells'] or args['--surfaces'] or args['--cell'] or args['--surface']):
+        args['--cells'] = True
+        args['--surfaces'] = True
+        args['--cell'] = [str(cell.number) for cell in inpt.cells if isinstance(cell, inp.Cell)]
+        args['--surface'] = [str(surface.number) for surface in inpt.surfaces if isinstance(surface, inp.Surface)]
+
     # Visualizing!
     if args['--cells']:
         if args['--pdf']:
@@ -51,7 +60,7 @@ def main() -> None:
             if 'PYTEST_CURRENT_TEST' not in os.environ:  # pragma: no cover
                 plot.show()
 
-    else:
+    if args['--surfaces']:
         if args['--pdf']:
             visualize.to_pdf_surfaces(_io.get_outfile(file, 'pdf', 'surfaces'))
         else:
@@ -60,4 +69,22 @@ def main() -> None:
             if 'PYTEST_CURRENT_TEST' not in os.environ:  # pragma: no cover
                 plot.show()
 
+    for number in args['--cell']:
+        if args['--pdf']:
+            visualize.to_pdf_cells(_io.get_outfile(file, 'pdf', f'cell_{number}'))
+        else:
+            plot = visualize.to_show_cell(number)
+
+            if 'PYTEST_CURRENT_TEST' not in os.environ:  # pragma: no cover
+                plot.show()
+
+    for number in args['--surface']:
+        if args['--pdf']:
+            visualize.to_pdf_cells(_io.get_outfile(file, 'pdf', f'surface_{number}'))
+        else:
+            plot = visualize.to_show_surface(number)
+
+            if 'PYTEST_CURRENT_TEST' not in os.environ:  # pragma: no cover
+                plot.show()
+
     _io.done()

src/pymcnp/_show/_shape.py

@@ -8,21 +8,24 @@ class Shape(metaclass=abc.ABCMeta):
     Represents generic visualizations.
 
     Attributes:
-        data: Underlying visualization.
+        surface: Underlying surface visualization.
+        cell: Underlying cell visualization.
     """
 
-    def __init__(self, data):
+    def __init__(self, surface, cell):
         """
         Initializes ``Shape``.
 
         Parameters:
-            data: Underlying visualization.
+            surface: Underlying surface visualization.
+            cell: Underlying cell visualization.
 
         Returns:
             ``Shape``
         """
 
-        self.data = data
+        self.surface = surface
+        self.cell = cell
 
     @abc.abstractmethod
     def __add__(a, b):  # pragma: no cover
@@ -54,34 +57,34 @@ def __and__(a, b):  # pragma: no cover
 
         raise NotImplementedError
 
-    # @abc.abstractmethod
-    # def __or__(a, b):  # pragma: no cover
-    # """
-    # Intersects ``Shape``.
+    @abc.abstractmethod
+    def __or__(a, b):  # pragma: no cover
+        """
+        Intersects ``Shape``.
 
-    # Parameters:
-    # a: Operand #1.
-    # b: Operand #2.
+        Parameters:
+            a: Operand #1.
+            b: Operand #2.
 
-    # Returns:
-    # ``Shape`` intersection.
-    # """
+        Returns:
+        ``Shape`` intersection.
+        """
 
-    # raise NotImplementedError
+        raise NotImplementedError
 
-    # @abc.abstractmethod
-    # def __invert__(a):  # pragma: no cover
-    # """
-    # Inverts ``Shape``.
+    @abc.abstractmethod
+    def __invert__(a):  # pragma: no cover
+        """
+        Complements ``Shape``.
 
-    # Parameters:
-    # a: Operand #1.
+        Parameters:
+            a: Operand #1.
 
-    # Returns:
-    # ``Shape`` complement.
-    # """
+        Returns:
+        ``Shape`` complement.
+        """
 
-    # raise NotImplementedError
+        raise NotImplementedError
 
     @abc.abstractmethod
     def rotate(self, axis: numpy.ndarray):  # pragma: no cover

src/pymcnp/_show/pyvista/Box.py

@@ -1,3 +1,4 @@
+import numpy
 import pyvista
 
 from . import _shape
@@ -19,7 +20,6 @@ def __init__(self, a: float, b: float, c: float):
         """
 
         super().__init__(
-            pyvista.Box(
-                bounds=(0, a, 0, b, 0, c),
-            )
+            pyvista.Box((0, a, 0, b, 0, c)),
+            lambda p: ~numpy.all((p >= 0) & (p <= [a, b, c]), axis=1),
         )