diff --git a/openmc/cell.py b/openmc/cell.py index 82a034b1c8c..b486954f328 100644 --- a/openmc/cell.py +++ b/openmc/cell.py @@ -605,6 +605,19 @@ def plot(self, *args, **kwargs): openmc.UniverseBase.next_id = next_id return u.plot(*args, **kwargs) + def voxel_plot(self, *args, **kwargs): + """Create a 3D voxel plot of the cell. + + .. versionadded:: 0.15.2 + """ + # Create dummy universe but preserve used_ids + next_id = openmc.UniverseBase.next_id + u = openmc.Universe(cells=[self]) + openmc.UniverseBase.used_ids.remove(u.id) + openmc.UniverseBase.next_id = next_id + return u.voxel_plot(*args, **kwargs) + + def create_xml_subelement(self, xml_element, memo=None): """Add the cell's xml representation to an incoming xml element diff --git a/openmc/geometry.py b/openmc/geometry.py index 8496fb23ad7..776d5d4acfc 100644 --- a/openmc/geometry.py +++ b/openmc/geometry.py @@ -781,3 +781,37 @@ def plot(self, *args, **kwargs): all_material_names.add(name) return model.plot(*args, **kwargs) + + def voxel_plot(self, *args, **kwargs): + """Create a 3D voxel plot of the geometry. + + .. versionadded:: 0.15.4 + """ + model = openmc.Model() + model.geometry = self + model.materials = self.get_all_materials().values() + + # collect all the material names from the geometry + all_material_names = {m.name for m in model.materials if m.name is not None} + + # makes a placeholder material for each material name if it isn't + # already present on the model. These materials are otherwise missing + # from the geometry and are needed for plotting. + for universe in model.geometry.get_all_universes().values(): + if not isinstance(universe, openmc.DAGMCUniverse): + continue + for name in universe.material_names: + # if this name is already present in the model, skip it + # (this can happen if the same material name is used in multiple + # universes) + if name in all_material_names: + continue + # if the material name is not present on the model, + # create a placeholder material with the same name + # and add it to the model + mat_dag = openmc.Material(name=name) + mat_dag.add_nuclide('H1', 1.0) + model.materials.append(mat_dag) + all_material_names.add(name) + + return model.voxel_plot(*args, **kwargs) diff --git a/openmc/model/model.py b/openmc/model/model.py index a9aaa481d5b..bf14b7f8051 100644 --- a/openmc/model/model.py +++ b/openmc/model/model.py @@ -1031,6 +1031,175 @@ def _set_plot_defaults( return origin, width, pixels + def _set_voxel_defaults( + self, + origin: Sequence[float] | None, + width: Sequence[float] | None, + pixels: int | Sequence[int], + ): + """Set default values for voxel plot parameters. + + Parameters + ---------- + origin : Sequence[float] or None + Origin (center) of the voxel plot + width : Sequence[float] or None + Width of the voxel plot in each dimension + pixels : int or Sequence[int] + Number of voxels in each direction + + Returns + ------- + tuple + (origin, width, pixels) with defaults applied + """ + bb = self.bounding_box + + # Check if bounding box contains inf values + if np.isinf(bb.extent['xyz']).any(): + if origin is None: + origin = (0.0, 0.0, 0.0) + if width is None: + width = (10.0, 10.0, 10.0) + else: + if origin is None: + # Replace NaN values with 0.0 + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + origin = np.nan_to_num(bb.center) + if width is None: + width = bb.width + + # Convert single int pixels to 3D tuple with cubic voxels + if isinstance(pixels, int): + # Calculate cubic root to get voxels per side + voxels_per_side = (pixels ** (1/3)) + + # Distribute voxels proportionally to width in each dimension + # to maintain cubic voxels + pixels = tuple( + max(1, int(round(voxels_per_side * w / max(width)))) + for w in width + ) + + return origin, width, pixels + + def voxel_plot( + self, + origin: Sequence[float] | None = None, + width: Sequence[float] | None = None, + pixels: int | Sequence[int] = 64000, + color_by: str = 'cell', + colors: dict | None = None, + seed: int | None = None, + openmc_exec: PathLike = 'openmc', + output: PathLike = 'voxel_plot.vti', + show_overlaps: bool = False, + overlap_color: Sequence[int] | str | None = None, + ) -> Path: + """Create a 3D voxel plot of the model. + + .. versionadded:: 0.15.4 + + Parameters + ---------- + origin : Sequence[float], optional + Origin (center) of the plot (3 values). If unspecified, defaults to + the center of the bounding box if available, otherwise (0, 0, 0). + width : Sequence[float], optional + Width of the plot in each dimension (3 values). If unspecified, + defaults to the bounding box width if available, otherwise + (10, 10, 10). + pixels : int | Sequence[int], optional + If an iterable of ints is provided then this directly sets the + number of voxels in each direction (3 values). If a single int is + provided then this sets the total number of voxels in the plot and + the number of voxels in each direction is calculated to maintain + cubic voxels based on the width. + color_by : {'cell', 'material'}, optional + Indicate whether the plot should be colored by cell or by material. + colors : dict, optional + Dictionary indicating that certain cells/materials should be + displayed with a particular color. + seed : int, optional + Pseudorandom number seed for plot coloring + openmc_exec : PathLike, optional + Path to OpenMC executable. Defaults to 'openmc'. + output : PathLike, optional + Path to the output file. File extension determines format: '.h5' + for HDF5 voxel file, '.vti' for VTK image data file. Defaults to + 'voxel_plot.vti'. + show_overlaps : bool, optional + Indicate whether or not overlapping regions are shown. + overlap_color : Sequence[int] or str, optional + Color to apply to overlapping regions. + + Returns + ------- + Path + Path to the generated voxel plot file (.h5 or .vti) + + """ + from openmc.plots import voxel_to_vtk + import shutil + + # Set defaults using helper function + origin, width, pixels = self._set_voxel_defaults(origin, width, pixels) + + # Validate output file extension + output_path = Path(output) + if output_path.suffix not in ('.h5', '.vti'): + raise ValueError( + f"Output file extension must be '.h5' or '.vti', " + f"got '{output_path.suffix}'" + ) + + with TemporaryDirectory() as tmpdir: + _plot_seed = self.settings.plot_seed + if seed is not None: + self.settings.plot_seed = seed + + # Create voxel plot object + voxel = openmc.VoxelPlot() + voxel.origin = origin + voxel.width = width + voxel.pixels = pixels + voxel.color_by = color_by + voxel.show_overlaps = show_overlaps + if overlap_color is not None: + voxel.overlap_color = overlap_color + if colors is not None: + voxel.colors = colors + + self.plots.append(voxel) + + # Run OpenMC in geometry plotting mode + self.plot_geometry(False, cwd=tmpdir, openmc_exec=openmc_exec) + + # Undo changes to model + self.plots.pop() + self.settings._plot_seed = _plot_seed + + # Determine base filename + base_name = output_path.stem + h5_filename = f'{base_name}.h5' + h5_src = Path(tmpdir) / h5_filename + + # Convert to VTK if .vti extension + if output_path.suffix == '.vti': + vti_filename = f'{base_name}.vti' + + # Use voxel_to_vtk function from openmc.plots + vti_src = voxel_to_vtk(h5_src, Path(tmpdir) / vti_filename) + + shutil.move(str(vti_src), output_path) + + return output_path + else: # .h5 extension + shutil.move(str(h5_src), output_path) + + return output_path + def id_map( self, origin: Sequence[float] | None = None, diff --git a/openmc/region.py b/openmc/region.py index cb9f3abd23a..8a55a9aa3cb 100644 --- a/openmc/region.py +++ b/openmc/region.py @@ -361,6 +361,22 @@ def plot(self, *args, **kwargs): openmc.Cell.next_id = next_id return c.plot(*args, **kwargs) + def voxel_plot(self, *args, **kwargs): + """Create a 3D voxel plot of the region. + + .. versionadded:: 0.15.2 + """ + for key in ('color_by', 'colors'): + if key in kwargs: + warnings.warn(f"The '{key}' argument is present but won't be applied in a region plot") + + # Create cell while not perturbing use of autogenerated IDs + next_id = openmc.Cell.next_id + c = openmc.Cell(region=self) + openmc.Cell.used_ids.remove(c.id) + openmc.Cell.next_id = next_id + return c.voxel_plot(*args, **kwargs) + class Intersection(Region, MutableSequence): r"""Intersection of two or more regions. diff --git a/openmc/universe.py b/openmc/universe.py index 0e64693ba8e..d83b8d14823 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -339,6 +339,15 @@ def plot(self, *args, **kwargs): model.geometry = openmc.Geometry(self) return model.plot(*args, **kwargs) + def voxel_plot(self, *args, **kwargs): + """Create a 3D voxel plot of the universe. + + .. versionadded:: 0.15.2 + """ + model = openmc.Model() + model.geometry = openmc.Geometry(self) + return model.voxel_plot(*args, **kwargs) + def get_nuclides(self): """Returns all nuclides in the universe diff --git a/tests/unit_tests/test_cell.py b/tests/unit_tests/test_cell.py index 60b20581868..731696f77f3 100644 --- a/tests/unit_tests/test_cell.py +++ b/tests/unit_tests/test_cell.py @@ -384,3 +384,36 @@ def test_plot(run_in_tmpdir): # ensure that calling the plot method doesn't # affect the universe ID space assert u_before.id + 1 == u_after.id + + +def test_voxel_plot(run_in_tmpdir): + """Test Cell.voxel_plot() method""" + + # Create a material + mat = openmc.Material(material_id=1, name='test_mat') + mat.set_density('g/cm3', 1.0) + mat.add_element('H', 1.0) + + # Create a simple cell with a sphere + sphere = openmc.Sphere(r=5.0, boundary_type='vacuum') + c = openmc.Cell(fill=mat, region=-sphere) + + # Create a universe before the voxel_plot + u_before = openmc.Universe() + + # Test VTI generation + vti_path = c.voxel_plot(pixels=1000, output='cell_test.vti') + assert vti_path.exists() + assert vti_path.suffix == '.vti' + + # Test H5 generation + h5_path = c.voxel_plot(pixels=1000, output='cell_test.h5') + assert h5_path.exists() + assert h5_path.suffix == '.h5' + + # Create a universe after the voxel_plot + u_after = openmc.Universe() + + # Ensure that calling the voxel_plot method doesn't + # affect the universe ID space + assert u_before.id + 1 == u_after.id diff --git a/tests/unit_tests/test_geometry.py b/tests/unit_tests/test_geometry.py index 6cc577c820c..0df5d9613f5 100644 --- a/tests/unit_tests/test_geometry.py +++ b/tests/unit_tests/test_geometry.py @@ -403,3 +403,45 @@ def test_redundant_surfaces(): geom = openmc.Geometry([c3]) redundant_surfs = geom.remove_redundant_surfaces() assert len(redundant_surfs) == 0 + + +def test_geometry_voxel_plot(run_in_tmpdir): + """Test Geometry.voxel_plot() method""" + + # Create materials + mat1 = openmc.Material(material_id=1, name='material_1') + mat1.set_density('g/cm3', 1.0) + mat1.add_element('H', 1.0) + + mat2 = openmc.Material(material_id=2, name='material_2') + mat2.set_density('g/cm3', 2.0) + mat2.add_element('Fe', 1.0) + + # Create simple geometry with two spheres + inner_sphere = openmc.Sphere(r=5.0) + outer_sphere = openmc.Sphere(r=10.0, boundary_type='vacuum') + + inner_cell = openmc.Cell(fill=mat1, region=-inner_sphere) + outer_cell = openmc.Cell(fill=mat2, region=+inner_sphere & -outer_sphere) + + geometry = openmc.Geometry([inner_cell, outer_cell]) + + # Test 1: Generate VTI file + vti_path = geometry.voxel_plot(pixels=1000, output='geom_test.vti') + assert vti_path.exists() + assert vti_path.suffix == '.vti' + + # Test 2: Generate H5 file + h5_path = geometry.voxel_plot(pixels=1000, output='geom_test.h5') + assert h5_path.exists() + assert h5_path.suffix == '.h5' + + # Test 3: Test with explicit pixel dimensions + explicit_path = geometry.voxel_plot(pixels=(10, 10, 10), output='geom_explicit.vti') + assert explicit_path.exists() + + # Test 4: Test default output + default_path = geometry.voxel_plot(pixels=1000) + assert default_path.exists() + assert default_path.name == 'voxel_plot.vti' + diff --git a/tests/unit_tests/test_model.py b/tests/unit_tests/test_model.py index d553af53c7e..819c7732546 100644 --- a/tests/unit_tests/test_model.py +++ b/tests/unit_tests/test_model.py @@ -659,6 +659,55 @@ def test_model_plot(): plt.close('all') +def test_model_voxel_plot(run_in_tmpdir): + """Test the Model.voxel_plot() method""" + + # Create a simple sphere geometry + mat = openmc.Material(material_id=1, name='test_material') + mat.set_density('g/cm3', 1.0) + mat.add_element('H', 1.0) + + sphere = openmc.Sphere(r=5.0, boundary_type='vacuum') + cell = openmc.Cell(fill=mat, region=-sphere) + geometry = openmc.Geometry([cell]) + + settings = openmc.Settings(particles=100, batches=5) + model = openmc.Model(geometry=geometry, settings=settings) + + # Test 1: Generate VTI file (default) + vti_path = model.voxel_plot(pixels=1000, output='test_voxel.vti') + assert vti_path.exists() + assert vti_path.suffix == '.vti' + assert vti_path.name == 'test_voxel.vti' + + # Test 2: Generate H5 file + h5_path = model.voxel_plot(pixels=1000, output='test_voxel.h5') + assert h5_path.exists() + assert h5_path.suffix == '.h5' + assert h5_path.name == 'test_voxel.h5' + + # Test 3: Use explicit pixel dimensions (tuple) + vti_path2 = model.voxel_plot(pixels=(10, 10, 10), output='explicit.vti') + assert vti_path2.exists() + + # Test 4: Test with color_by material + mat_path = model.voxel_plot( + pixels=1000, + color_by='material', + output='material_voxel.vti' + ) + assert mat_path.exists() + + # Test 5: Invalid file extension should raise ValueError + with pytest.raises(ValueError, match="must be '.h5' or '.vti'"): + model.voxel_plot(output='bad_extension.png') + + # Test 6: Test default output filename + default_path = model.voxel_plot(pixels=1000) + assert default_path.exists() + assert default_path.name == 'voxel_plot.vti' + + def test_model_id_map_initialization(run_in_tmpdir): model = openmc.examples.pwr_assembly() model.init_lib(output=False) diff --git a/tests/unit_tests/test_region.py b/tests/unit_tests/test_region.py index cb9fa171bb6..69d2a7dca95 100644 --- a/tests/unit_tests/test_region.py +++ b/tests/unit_tests/test_region.py @@ -255,3 +255,33 @@ def test_plot(): # Ensure that calling plot doesn't affect cell ID space c_after = openmc.Cell() assert c_after.id - 1 == c_before.id + + +def test_voxel_plot(run_in_tmpdir): + """Test Region.voxel_plot() method""" + + # Create a region (intersection of sphere and half-space) + sphere = openmc.Sphere(r=5.0, boundary_type='vacuum') + plane = openmc.ZPlane(z0=0.0) + region = -sphere & +plane # Upper hemisphere + + # Track cell IDs before voxel_plot + c_before = openmc.Cell() + + # Test VTI generation + vti_path = region.voxel_plot(pixels=1000, output='region_test.vti') + assert vti_path.exists() + assert vti_path.suffix == '.vti' + + # Test H5 generation + h5_path = region.voxel_plot(pixels=1000, output='region_test.h5') + assert h5_path.exists() + assert h5_path.suffix == '.h5' + + # Ensure that calling voxel_plot doesn't affect cell ID space + c_after = openmc.Cell() + assert c_after.id - 1 == c_before.id + + # Test that color_by parameter produces a warning + with pytest.warns(UserWarning, match="won't be applied"): + region.voxel_plot(pixels=1000, output='region_warning.vti', color_by='material') diff --git a/tests/unit_tests/test_universe.py b/tests/unit_tests/test_universe.py index efe8552a648..46ebc611a98 100644 --- a/tests/unit_tests/test_universe.py +++ b/tests/unit_tests/test_universe.py @@ -216,3 +216,40 @@ def test_get_nuclide_densities(): universe = openmc.Universe(cells=[cell]) with pytest.raises(RuntimeError): universe.get_nuclide_densities() + + +def test_voxel_plot(run_in_tmpdir): + """Test Universe.voxel_plot() method""" + + # Create materials + mat1 = openmc.Material(material_id=1, name='mat1') + mat1.set_density('g/cm3', 1.0) + mat1.add_element('H', 1.0) + + mat2 = openmc.Material(material_id=2, name='mat2') + mat2.set_density('g/cm3', 2.0) + mat2.add_element('Fe', 1.0) + + # Create cells + inner_sphere = openmc.Sphere(r=3.0) + outer_sphere = openmc.Sphere(r=5.0, boundary_type='vacuum') + + cell1 = openmc.Cell(fill=mat1, region=-inner_sphere) + cell2 = openmc.Cell(fill=mat2, region=+inner_sphere & -outer_sphere) + + # Create universe + u = openmc.Universe(cells=[cell1, cell2]) + + # Test VTI generation + vti_path = u.voxel_plot(pixels=1000, output='univ_test.vti') + assert vti_path.exists() + assert vti_path.suffix == '.vti' + + # Test H5 generation + h5_path = u.voxel_plot(pixels=1000, output='univ_test.h5') + assert h5_path.exists() + assert h5_path.suffix == '.h5' + + # Test with explicit pixel dimensions + explicit_path = u.voxel_plot(pixels=(10, 10, 10), output='univ_explicit.vti') + assert explicit_path.exists()