Differentiate materials in DAGMC universes

include/openmc/capi.h

@@ -29,6 +29,7 @@ int openmc_cell_set_temperature(
   int32_t index, double T, const int32_t* instance, bool set_contained = false);
 int openmc_cell_set_translation(int32_t index, const double xyz[]);
 int openmc_cell_set_rotation(int32_t index, const double rot[], size_t rot_len);
+int openmc_get_dagmc_cell_ids(int32_t univ_id, int32_t* ids, size_t* n);
 int openmc_energy_filter_get_bins(
   int32_t index, const double** energies, size_t* n);
 int openmc_energy_filter_set_bins(

include/openmc/cell.h

@@ -319,7 +319,6 @@ class Cell {
   int32_t universe_;        //!< Universe # this cell is in
   int32_t fill_;            //!< Universe # filling this cell
   int32_t n_instances_ {0}; //!< Number of instances of this cell
-  GeometryType geom_type_;  //!< Geometric representation type (CSG, DAGMC)
 
   //! \brief Index corresponding to this cell in distribcell arrays
   int distribcell_index_ {C_NONE};
@@ -349,6 +348,12 @@ class Cell {
   vector<double> rotation_;
 
   vector<int32_t> offset_; //!< Distribcell offset table
+
+  const GeometryType& geom_type() const { return geom_type_; }
+  GeometryType& geom_type() { return geom_type_; }
+
+private:
+  GeometryType geom_type_; //!< Geometric representation type (CSG, DAGMC)
 };
 
 struct CellInstanceItem {

include/openmc/dagmc.h

@@ -184,6 +184,7 @@ class DAGUniverse : public Universe {
                              //!< generate new material IDs for the universe
   bool has_graveyard_; //!< Indicates if the DAGMC geometry has a "graveyard"
                        //!< volume
+  std::map<std::string, vector<std::string>> instance_mat_assignment;
 };
 
 //==============================================================================

openmc/cell.py

@@ -175,6 +175,8 @@ def fill(self, fill):
                        f'non-Material or Universe fill "{fill}"')
                 raise ValueError(msg)
         self._fill = fill
+        if isinstance(self.fill, openmc.DAGMCUniverse):
+            self.fill._num_instances += 1
 
         # Info about atom content can now be invalid
         # (since fill has just changed)
@@ -780,3 +782,83 @@ def from_xml_element(cls, elem, surfaces, materials, get_universe):
         univ_id = int(get_text(elem, 'universe', 0))
         get_universe(univ_id).add_cell(c)
         return c
+
+
+class DAGSpeudoCell(Cell):
+    def __init__(self, cell_id=None, name='', fill=None, region=None):
+        super().__init__(cell_id, name, fill, region)
+
+    @property
+    def DAG_parent_universe(self):
+        """Get the parent universe of the cell."""
+        return self._parent_universe
+
+    @DAG_parent_universe.setter
+    def DAG_parent_universe(self, universe):
+        """Set the parent universe of the cell."""
+        self._parent_universe = universe.id
+
+    def boundingbox(self):
+        print("Warning: Bounding box is not available for cells in a DAGMC universe.")
+        return {}
+
+    def get_all_cells(self, memo=None):
+        return {}
+
+    def get_all_materials(self, memo=None):
+        """Return all materials that are contained within the cell
+
+        Returns
+        -------
+        materials : dict
+            Dictionary whose keys are material IDs and values are
+            :class:`Material` instances
+
+        """
+        materials = {}
+        if self.fill_type == 'material':
+            materials[self.fill.id] = self.fill
+        elif self.fill_type == 'distribmat':
+            for m in self.fill:
+                if m is not None:
+                    materials[m.id] = m
+        else:
+            # Append all Cells in each Cell in the Universe to the dictionary
+            cells = self.get_all_cells(memo)
+            for cell in cells.values():
+                materials.update(cell.get_all_materials(memo))
+
+        return materials
+
+    @property
+    def fill_type(self):
+        if isinstance(self.fill, openmc.Material):
+            return 'material'
+        elif isinstance(self.fill, openmc.UniverseBase):
+            return 'universe'
+        elif isinstance(self.fill, openmc.Lattice):
+            return 'lattice'
+        elif isinstance(self.fill, Iterable):
+            return 'distribmat'
+        else:
+            return 'void'
+
+    def get_all_universes(self, memo=None):
+        return {}
+
+    def clone(self, clone_materials=True, clone_regions=True, memo=None):
+        print("Warning: clone is not available for cells in a DAGMC universe.")
+        return None
+
+    def plot(self, *args, **kwargs):
+        print("Warning: plot is not available for cells in a DAGMC universe.")
+        return None
+
+    def create_xml_subelement(self, xml_element, memo=None):
+        print("Warning: create_xml_subelement is not available for cells in a DAGMC universe.")
+        return None
+
+    @classmethod
+    def from_xml_element(cls, elem, surfaces, materials, get_universe):
+        print("Warning: from_xml_element is not available for cells in a DAGMC universe.")
+        return None

openmc/geometry.py

@@ -404,6 +404,23 @@ def get_all_nuclides(self) -> list[str]:
         for material in self.get_all_materials().values():
             all_nuclides |= set(material.get_nuclides())
         return sorted(all_nuclides)
+
+    def get_all_dag_universes(self) -> typing.Dict[int, openmc.DAGMCUniverse]:
+        """Return all universes in the geometry.
+
+        Returns
+        -------
+        dict
+            Dictionary mapping universe IDs to :class:`openmc.DAGMCUniverse`
+            instances
+
+        """
+        universes = self.get_all_universes()
+        dag_universes = {}
+        for id, uni in universes.items():
+            if isinstance(uni, openmc.DAGMCUniverse):
+                dag_universes[id] = uni
+        return dag_universes
 
     def get_all_materials(self) -> dict[int, openmc.Material]:
         """Return all materials within the geometry.
@@ -737,6 +754,8 @@ def determine_paths(self, instances_only=False):
         for material in self.get_all_materials().values():
             material._paths = []
             material._num_instances = 0
+        for dag_uni in self.get_all_dag_universes().values():
+            dag_uni._num_instances = 0
 
         # Recursively traverse the CSG tree to count all cell instances
         self.root_universe._determine_paths(instances_only=instances_only)

openmc/lib/__init__.py

@@ -68,6 +68,7 @@ def _uwuw_enabled():
 from .math import *
 from .plot import *
 from .weight_windows import *
+from .dagmc import *
 
 # Flag to denote whether or not openmc.lib.init has been called
 # TODO: Establish and use a flag in the C++ code to represent the status of the

openmc/lib/dagmc.py

@@ -0,0 +1,44 @@
+import sys
+
+from ctypes import c_int, c_int32, POINTER, c_size_t
+
+import numpy as np
+
+from . import _dll
+from .error import _error_handler
+
+
+
+# DAGMC functions
+_dll.openmc_get_dagmc_cell_ids.argtypes = [c_int32, POINTER(c_int32), POINTER(c_size_t)]
+_dll.openmc_get_dagmc_cell_ids.restype = c_int
+_dll.openmc_get_dagmc_cell_ids.errcheck = _error_handler
+
+
+def get_dagmc_cell_ids(volume_id, n_cells):
+    """Get the DAGMC cell IDs for a volume.
+
+    Parameters
+    ----------
+    volume_id : int
+        ID of the volume to get DAGMC cell IDs for.
+    n_cells : int
+        Number of cells in the volume.
+
+    Returns
+    -------
+    numpy.ndarray
+        DAGMC cell IDs for the volume.
+
+    """
+    cell_ids = np.empty(n_cells, dtype=np.int32)
+    n = c_size_t()
+    _dll.openmc_get_dagmc_cell_ids(
+        volume_id,
+        cell_ids.ctypes.data_as(POINTER(c_int32)),
+        n
+    )
+    if n.value != n_cells:
+        raise ValueError(f"Number of cells obtained {n.value} from DAGMC does "
+                         f"not match the expected number of cells {n_cells}.")
+    return cell_ids

openmc/model/model.py

@@ -16,6 +16,7 @@
 from openmc.executor import _process_CLI_arguments
 from openmc.checkvalue import check_type, check_value, PathLike
 from openmc.exceptions import InvalidIDError
+import openmc.lib
 from openmc.utility_funcs import change_directory
 
 
@@ -322,6 +323,39 @@ def init_lib(self, threads=None, geometry_debug=False, restart_file=None,
         # the user-provided intracomm which will either be None or an mpi4py
         # communicator
         openmc.lib.init(args=args, intracomm=intracomm, output=output)
+        self.sync_dagmc_cells()
+
+    def sync_dagmc_cells(self):
+        """Synchronize DAGMC cell information between Python and C API
+
+        .. versionadded:: 0.13.0
+
+        """
+        if not self.is_initialized:
+            raise RuntimeError("Model must be initialized via Model.init_lib "
+                               "before calling this method.")
+
+        import openmc.lib
+
+        if self.materials:
+            mats = self.materials
+        else:
+            mats = self.geometry.get_all_materials().values()
+        mats_per_id = {mat.id: mat for mat in mats}
+
+        for cell in self.geometry.get_all_cells().values():
+            if isinstance(cell.fill, openmc.DAGMCUniverse):
+                for dag_cell_id in openmc.lib.dagmc.get_dagmc_cell_ids(cell.fill.id, cell.fill._n_geom_elements('volume')):
+                    dag_cell = openmc.lib.cells[dag_cell_id]
+                    if isinstance(dag_cell.fill, Iterable):
+                        fill = [mats_per_id[mat_id.id]
+                                for mat_id in dag_cell.fill]
+                    else:
+                        fill = mats_per_id[dag_cell.fill.id]
+                    dag_pseudo_cell = openmc.DAGSpeudoCell(
+                        cell_id=dag_cell_id, fill=fill
+                    )
+                    cell.fill.add_cell(dag_pseudo_cell)
 
     def finalize_lib(self):
         """Finalize simulation and free memory allocated for the C API
@@ -1026,13 +1060,51 @@ def differentiate_depletable_mats(self, diff_volume_method: str):
             volume equally between the new materials, 'match cell' sets the
             volume of the material to volume of the cell they fill.
         """
+        self.differentiate_mats(diff_volume_method, depletable_only=True)
+
+    def differentiate_mats(self,
+                           diff_volume_method: str,
+                           depletable_only: bool = True):
+        """Assign distribmats for each material
+
+        .. versionadded:: 0.14.0
+
+        Parameters
+        ----------
+        diff_volume_method : str
+            Specifies how the volumes of the new materials should be found.
+            Default is to 'divide equally' which divides the original material
+            volume equally between the new materials, 'match cell' sets the
+            volume of the material to volume of the cell they fill.
+        depletable_ony : bool
+            Differentiate depletable materials only.
+        """
+        # Check if there is some DAGMC universe
+        if len(self.geometry.get_all_dag_universes()) > 0:
+            # Reset num_instances of models.materials
+            # (Updated from dag-verses appearance or from CSG presences...)
+            for mat in self._materials:
+                mat._num_instances = 0
+
         # Count the number of instances for each cell and material
         self.geometry.determine_paths(instances_only=True)
 
+        # Get mat instances from the different dag-verses
+        dag_mats_n_inst = {}
+        if len(self.geometry.get_all_dag_universes()) > 0:
+            # Get material in the DAG-verses
+            for dag_verse in self.geometry.get_all_dag_universes().values():
+                for mat_name in dag_verse.material_names:
+                    dag_mats_n_inst[mat_name] = dag_mats_n_inst.get(mat_name, 0) + dag_verse.num_instances
+        # Account for the multiplicity of the dag-verses materials
+        for mat in self.materials:
+            if mat.name in dag_mats_n_inst:
+                mat._num_instances += dag_mats_n_inst[mat.name]
+
         # Extract all depletable materials which have multiple instances
         distribmats = set(
             [mat for mat in self.materials
-                if mat.depletable and mat.num_instances > 1])
+                if (mat.depletable or not depletable_only) and mat.num_instances > 1])
 
         if diff_volume_method == 'divide equally':
             for mat in distribmats:
@@ -1058,8 +1130,59 @@ def differentiate_depletable_mats(self, diff_volume_method: str):
                                     "diff_volume_method='match cell'."
                                 )
                             cell.fill.volume = cell.volume
-
+                    else:
+                        for _ in range(cell.num_instances):
+                            cell.fill = mat.clone()
+
+        dag_verse_mats = []
+        for dag_verse in self.geometry.get_all_dag_universes().values():
+            if dag_verse.num_instances > 1:
+                deplete_mat_dict = {}
+                for mat_name in dag_verse.material_names:
+                    mat_found = False
+                    for mat in self.materials:
+                        if mat.name == mat_name:
+                            mat_found = True
+                            if mat.depletable or not depletable_only:
+                                mats_clones = [mat.clone() for _ in range(dag_verse.num_instances)]
+                                for i, mat_clone in enumerate(mats_clones):
+                                    mat_clone.name += "_" + str(dag_verse.id) + "_" + str(i)
+                                    dag_verse_mats.append(mat_clone)
+                                    if diff_volume_method == 'match cell':
+                                        if self.is_initialized:
+                                            for cell in dag_verse._cells.values():
+                                                if not isinstance(cell.fill, list):
+                                                    if cell.fill.name == mat.name:
+                                                        cell.fill = mat_clone
+                                                        mat_clone.volume = cell.volume
+                                                    elif cell.fill.name.split("_")[0] == mat.name:
+                                                        mat_clone.volume = cell.volume
+                                                        cell.fill = [cell.fill]
+                                                        cell.fill.append(mat_clone)
+                                                else:
+                                                    if cell.fill[0].name.split("_")[0] == mat.name:
+                                                        mat_clone.volume = cell.volume
+                                                        cell.fill.append(mat_clone)
+                                        else:
+                                            raise NotImplementedError("differentiate mats with DAGMC universe and match cell option is only available when cpp_lib is initialized")
+                                deplete_mat_dict[mat_name.lower()] = [mat.name for mat in mats_clones]
+                            else:
+                                dag_verse_mats.append(mat)
+                    if not mat_found:
+                        raise ValueError(f"Material {mat_name} referenced in "
+                                         f"the dagmc Universe "
+                                         f"{dag_verse.filename} not found in "
+                                         f"the Model. Please add it to the "
+                                         f"Model. Please note that uwuw "
+                                         f"formalism is not compatible with "
+                                         f"differentiate_depletable_mats yet.")
+                dag_verse.mat_assignment = deplete_mat_dict
+
         if self.materials is not None:
             self.materials = openmc.Materials(
                 self.geometry.get_all_materials().values()
             )
+        if dag_verse_mats:
+            self.materials.extend(dag_verse_mats)
+        self.materials = list(set(self.materials))
+