Compute homogenized materials over mesh elements (#2971)

openmc/deplete/abc.py

@@ -10,8 +10,6 @@
 from copy import deepcopy
 from inspect import signature
 from numbers import Real, Integral
-from contextlib import contextmanager
-import os
 from pathlib import Path
 import time
 from typing import Optional, Union, Sequence
@@ -22,6 +20,7 @@
 
 from openmc.checkvalue import check_type, check_greater_than, PathLike
 from openmc.mpi import comm
+from openmc.utility_funcs import change_directory
 from openmc import Material
 from .stepresult import StepResult
 from .chain import Chain
@@ -61,25 +60,6 @@
     # Can't set __doc__ on properties on Python 3.4
     pass
 
-@contextmanager
-def change_directory(output_dir):
-    """
-    Helper function for managing the current directory.
-
-    Parameters
-    ----------
-    output_dir : pathlib.Path
-        Directory to switch to.
-    """
-    orig_dir  = os.getcwd()
-    output_dir = Path(output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-    try:
-        os.chdir(output_dir)
-        yield
-    finally:
-        os.chdir(orig_dir)
-
 
 class TransportOperator(ABC):
     """Abstract class defining a transport operator

openmc/deplete/independent_operator.py

@@ -130,7 +130,6 @@ def __init__(self,
         # Validate micro-xs parameters
         check_type('materials', materials, openmc.Materials)
         check_type('micros', micros, Iterable, MicroXS)
-        check_type('fluxes', fluxes, Iterable, float)
 
         if not (len(fluxes) == len(micros) == len(materials)):
             msg = (f'The length of fluxes ({len(fluxes)}) should be equal to '

openmc/deplete/microxs.py

@@ -13,11 +13,11 @@
 
 from openmc.checkvalue import check_type, check_value, check_iterable_type, PathLike
 from openmc.exceptions import DataError
+from openmc.utility_funcs import change_directory
 from openmc import StatePoint
 from openmc.mgxs import GROUP_STRUCTURES
 from openmc.data import REACTION_MT
 import openmc
-from .abc import change_directory
 from .chain import Chain, REACTIONS
 from .coupled_operator import _find_cross_sections, _get_nuclides_with_data
 import openmc.lib
@@ -284,38 +284,37 @@ def from_multigroup_flux(
         # Create 2D array for microscopic cross sections
         microxs_arr = np.zeros((len(nuclides), len(mts)))
 
-        with TemporaryDirectory() as tmpdir:
-            # Create a material with all nuclides
-            mat_all_nucs = openmc.Material()
-            for nuc in nuclides:
-                if nuc in nuclides_with_data:
-                    mat_all_nucs.add_nuclide(nuc, 1.0)
-            mat_all_nucs.set_density("atom/b-cm", 1.0)
-
-            # Create simple model containing the above material
-            surf1 = openmc.Sphere(boundary_type="vacuum")
-            surf1_cell = openmc.Cell(fill=mat_all_nucs, region=-surf1)
-            model = openmc.Model()
-            model.geometry = openmc.Geometry([surf1_cell])
-            model.settings = openmc.Settings(
-                particles=1, batches=1, output={'summary': False})
-
-            with change_directory(tmpdir):
-                # Export model within temporary directory
-                model.export_to_model_xml()
-
-                with openmc.lib.run_in_memory(**init_kwargs):
-                    # For each nuclide and reaction, compute the flux-averaged
-                    # cross section
-                    for nuc_index, nuc in enumerate(nuclides):
-                        if nuc not in nuclides_with_data:
-                            continue
-                        lib_nuc = openmc.lib.nuclides[nuc]
-                        for mt_index, mt in enumerate(mts):
-                            xs = lib_nuc.collapse_rate(
-                                mt, temperature, energies, multigroup_flux
-                            )
-                            microxs_arr[nuc_index, mt_index] = xs
+        # Create a material with all nuclides
+        mat_all_nucs = openmc.Material()
+        for nuc in nuclides:
+            if nuc in nuclides_with_data:
+                mat_all_nucs.add_nuclide(nuc, 1.0)
+        mat_all_nucs.set_density("atom/b-cm", 1.0)
+
+        # Create simple model containing the above material
+        surf1 = openmc.Sphere(boundary_type="vacuum")
+        surf1_cell = openmc.Cell(fill=mat_all_nucs, region=-surf1)
+        model = openmc.Model()
+        model.geometry = openmc.Geometry([surf1_cell])
+        model.settings = openmc.Settings(
+            particles=1, batches=1, output={'summary': False})
+
+        with change_directory(tmpdir=True):
+            # Export model within temporary directory
+            model.export_to_model_xml()
+
+            with openmc.lib.run_in_memory(**init_kwargs):
+                # For each nuclide and reaction, compute the flux-averaged
+                # cross section
+                for nuc_index, nuc in enumerate(nuclides):
+                    if nuc not in nuclides_with_data:
+                        continue
+                    lib_nuc = openmc.lib.nuclides[nuc]
+                    for mt_index, mt in enumerate(mts):
+                        xs = lib_nuc.collapse_rate(
+                            mt, temperature, energies, multigroup_flux
+                        )
+                        microxs_arr[nuc_index, mt_index] = xs
 
         return cls(microxs_arr, nuclides, reactions)
 

openmc/mesh.py

@@ -7,7 +7,8 @@
 from math import pi, sqrt, atan2
 from numbers import Integral, Real
 from pathlib import Path
-from typing import Optional, Sequence, Tuple
+import tempfile
+from typing import Optional, Sequence, Tuple, List
 
 import h5py
 import lxml.etree as ET
@@ -16,6 +17,7 @@
 import openmc
 import openmc.checkvalue as cv
 from openmc.checkvalue import PathLike
+from openmc.utility_funcs import change_directory
 from ._xml import get_text
 from .mixin import IDManagerMixin
 from .surface import _BOUNDARY_TYPES
@@ -145,6 +147,94 @@ def from_xml_element(cls, elem: ET.Element):
         else:
             raise ValueError(f'Unrecognized mesh type "{mesh_type}" found.')
 
+    def get_homogenized_materials(
+            self,
+            model: openmc.Model,
+            n_samples: int = 10_000,
+            prn_seed: Optional[int] = None,
+            **kwargs
+    ) -> List[openmc.Material]:
+        """Generate homogenized materials over each element in a mesh.
+
+        .. versionadded:: 0.14.1
+
+        Parameters
+        ----------
+        model : openmc.Model
+            Model containing materials to be homogenized and the associated
+            geometry.
+        n_samples : int
+            Number of samples in each mesh element.
+        prn_seed : int, optional
+            Pseudorandom number generator (PRNG) seed; if None, one will be
+            generated randomly.
+        **kwargs
+            Keyword-arguments passed to :func:`openmc.lib.init`.
+
+        Returns
+        -------
+        list of openmc.Material
+            Homogenized material in each mesh element
+
+        """
+        import openmc.lib
+
+        with change_directory(tmpdir=True):
+            # In order to get mesh into model, we temporarily replace the
+            # tallies with a single mesh tally using the current mesh
+            original_tallies = model.tallies
+            new_tally = openmc.Tally()
+            new_tally.filters = [openmc.MeshFilter(self)]
+            new_tally.scores = ['flux']
+            model.tallies = [new_tally]
+
+            # Export model to XML
+            model.export_to_model_xml()
+
+            # Get material volume fractions
+            openmc.lib.init(**kwargs)
+            mesh = openmc.lib.tallies[new_tally.id].filters[0].mesh
+            mat_volume_by_element = [
+                [
+                    (mat.id if mat is not None else None, volume)
+                    for mat, volume in mat_volume_list
+                ]
+                for mat_volume_list in mesh.material_volumes(n_samples, prn_seed)
+            ]
+            openmc.lib.finalize()
+
+            # Restore original tallies
+            model.tallies = original_tallies
+
+        # Create homogenized material for each element
+        materials = model.geometry.get_all_materials()
+        homogenized_materials = []
+        for mat_volume_list in mat_volume_by_element:
+            material_ids, volumes = [list(x) for x in zip(*mat_volume_list)]
+            total_volume = sum(volumes)
+
+            # Check for void material and remove
+            try:
+                index_void = material_ids.index(None)
+            except ValueError:
+                pass
+            else:
+                material_ids.pop(index_void)
+                volumes.pop(index_void)
+
+            # Compute volume fractions
+            volume_fracs = np.array(volumes) / total_volume
+
+            # Get list of materials and mix 'em up!
+            mats = [materials[uid] for uid in material_ids]
+            homogenized_mat = openmc.Material.mix_materials(
+                mats, volume_fracs, 'vo'
+            )
+            homogenized_mat.volume = total_volume
+            homogenized_materials.append(homogenized_mat)
+
+        return homogenized_materials
+
 
 class StructuredMesh(MeshBase):
     """A base class for structured mesh functionality

openmc/model/model.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 from collections.abc import Iterable
-from contextlib import contextmanager
 from functools import lru_cache
 import os
 from pathlib import Path
@@ -18,18 +17,7 @@
 from openmc.executor import _process_CLI_arguments
 from openmc.checkvalue import check_type, check_value, PathLike
 from openmc.exceptions import InvalidIDError
-
-
-@contextmanager
-def _change_directory(working_dir):
-    """A context manager for executing in a provided working directory"""
-    start_dir = Path.cwd()
-    Path.mkdir(working_dir, parents=True, exist_ok=True)
-    os.chdir(working_dir)
-    try:
-        yield
-    finally:
-        os.chdir(start_dir)
+from openmc.utility_funcs import change_directory
 
 
 class Model:
@@ -395,7 +383,7 @@ def deplete(self, timesteps, method='cecm', final_step=True,
         # Store whether or not the library was initialized when we started
         started_initialized = self.is_initialized
 
-        with _change_directory(Path(directory)):
+        with change_directory(directory):
             with openmc.lib.quiet_dll(output):
                 # TODO: Support use of IndependentOperator too
                 depletion_operator = dep.CoupledOperator(self, **op_kwargs)
@@ -677,7 +665,7 @@ def run(self, particles=None, threads=None, geometry_debug=False,
         last_statepoint = None
 
         # Operate in the provided working directory
-        with _change_directory(Path(cwd)):
+        with change_directory(cwd):
             if self.is_initialized:
                 # Handle the run options as applicable
                 # First dont allow ones that must be set via init
@@ -767,7 +755,7 @@ def calculate_volumes(self, threads=None, output=True, cwd='.',
             raise ValueError("The Settings.volume_calculations attribute must"
                              " be specified before executing this method!")
 
-        with _change_directory(Path(cwd)):
+        with change_directory(cwd):
             if self.is_initialized:
                 if threads is not None:
                     msg = "Threads must be set via Model.is_initialized(...)"
@@ -829,7 +817,7 @@ def plot_geometry(self, output=True, cwd='.', openmc_exec='openmc'):
             raise ValueError("The Model.plots attribute must be specified "
                              "before executing this method!")
 
-        with _change_directory(Path(cwd)):
+        with change_directory(cwd):
             if self.is_initialized:
                 # Compute the volumes
                 openmc.lib.plot_geometry(output)

openmc/utility_funcs.py

@@ -0,0 +1,38 @@
+from contextlib import contextmanager
+import os
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Optional
+
+from .checkvalue import PathLike
+
+@contextmanager
+def change_directory(working_dir: Optional[PathLike] = None, *, tmpdir: bool = False):
+    """Context manager for executing in a provided working directory
+
+    Parameters
+    ----------
+    working_dir : path-like
+        Directory to switch to.
+    tmpdir : bool
+        Whether to use a temporary directory instead of a specific working directory
+
+    """
+    orig_dir = Path.cwd()
+
+    # Set up temporary directory if requested
+    if tmpdir:
+        tmp = TemporaryDirectory()
+        working_dir = tmp.name
+    elif working_dir is None:
+        raise ValueError('Must pass working_dir argument or specify tmpdir=True.')
+
+    working_dir = Path(working_dir)
+    working_dir.mkdir(parents=True, exist_ok=True)
+    os.chdir(working_dir)
+    try:
+        yield
+    finally:
+        os.chdir(orig_dir)
+        if tmpdir:
+            tmp.cleanup()