diff --git a/openmc/model/model.py b/openmc/model/model.py index 9ff574ec63..c19c8ac9f1 100644 --- a/openmc/model/model.py +++ b/openmc/model/model.py @@ -67,8 +67,14 @@ class Model: """ - def __init__(self, geometry=None, materials=None, settings=None, - tallies=None, plots=None): + def __init__( + self, + geometry: openmc.Geometry | None = None, + materials: openmc.Materials = None, + settings: openmc.Settings | None = None, + tallies: openmc.Tallies | None = None, + plots: openmc.Plots | None = None, + ): self.geometry = openmc.Geometry() if geometry is None else geometry self.materials = openmc.Materials() if materials is None else materials self.settings = openmc.Settings() if settings is None else settings @@ -192,24 +198,29 @@ class Model: return result @classmethod - def from_xml(cls, geometry='geometry.xml', materials='materials.xml', - settings='settings.xml', tallies='tallies.xml', - plots='plots.xml') -> Model: + def from_xml( + cls, + geometry: PathLike = "geometry.xml", + materials: PathLike = "materials.xml", + settings: PathLike = "settings.xml", + tallies: PathLike = "tallies.xml", + plots: PathLike = "plots.xml", + ) -> Model: """Create model from existing XML files Parameters ---------- - geometry : str + geometry : PathLike Path to geometry.xml file - materials : str + materials : PathLike Path to materials.xml file - settings : str + settings : PathLike Path to settings.xml file - tallies : str + tallies : PathLike Path to tallies.xml file .. versionadded:: 0.13.0 - plots : str + plots : PathLike Path to plots.xml file .. versionadded:: 0.13.0 @@ -229,14 +240,14 @@ class Model: return cls(geometry, materials, settings, tallies, plots) @classmethod - def from_model_xml(cls, path='model.xml'): + def from_model_xml(cls, path: PathLike = "model.xml") -> Model: """Create model from single XML file .. versionadded:: 0.13.3 Parameters ---------- - path : str or PathLike + path : PathLike Path to model.xml file """ parser = ET.XMLParser(huge_tree=True) @@ -262,8 +273,17 @@ class Model: return model - def init_lib(self, threads=None, geometry_debug=False, restart_file=None, - tracks=False, output=True, event_based=None, intracomm=None, directory=None): + def init_lib( + self, + threads: int | None = None, + geometry_debug: bool = False, + restart_file: PathLike | None = None, + tracks: bool = False, + output: bool = True, + event_based: bool | None = None, + intracomm=None, + directory: PathLike | None = None, + ): """Initializes the model in memory via the C API .. versionadded:: 0.13.0 @@ -278,7 +298,7 @@ class Model: variable). geometry_debug : bool, optional Turn on geometry debugging during simulation. Defaults to False. - restart_file : str, optional + restart_file : PathLike, optional Path to restart file to use tracks : bool, optional Enables the writing of particles tracks. The number of particle @@ -291,7 +311,7 @@ class Model: the Settings will be used. intracomm : mpi4py.MPI.Intracomm or None, optional MPI intracommunicator - directory : str or None, optional + directory : PathLike or None, optional Directory to write XML files to. Defaults to None. """ @@ -365,9 +385,15 @@ class Model: openmc.lib.finalize() - def deplete(self, timesteps, method='cecm', final_step=True, - operator_kwargs=None, directory='.', output=True, - **integrator_kwargs): + def deplete( + self, + method: str = "cecm", + final_step: bool = True, + operator_kwargs: dict | None = None, + directory: PathLike = ".", + output: bool = True, + **integrator_kwargs, + ): """Deplete model using specified timesteps/power .. versionchanged:: 0.13.0 @@ -376,10 +402,12 @@ class Model: Parameters ---------- - timesteps : iterable of float - Array of timesteps in units of [s]. Note that values are not - cumulative. - method : str, optional + timesteps : iterable of float or iterable of tuple + Array of timesteps. Note that values are not cumulative. The units are + specified by the `timestep_units` argument when `timesteps` is an + iterable of float. Alternatively, units can be specified for each step + by passing an iterable of (value, unit) tuples. + method : str Integration method used for depletion (e.g., 'cecm', 'predictor'). Defaults to 'cecm'. final_step : bool, optional @@ -388,14 +416,14 @@ class Model: operator_kwargs : dict Keyword arguments passed to the depletion operator initializer (e.g., :func:`openmc.deplete.Operator`) - directory : str, optional + directory : PathLike, optional Directory to write XML files to. If it doesn't exist already, it will be created. Defaults to the current working directory output : bool Capture OpenMC output from standard out integrator_kwargs : dict - Remaining keyword arguments passed to the depletion Integrator - initializer (e.g., :func:`openmc.deplete.integrator.cecm`). + Remaining keyword arguments passed to the depletion integrator + (e.g., :class:`openmc.deplete.CECMIntegrator`). """ @@ -426,8 +454,7 @@ class Model: check_value('method', method, dep.integrators.integrator_by_name.keys()) integrator_class = dep.integrators.integrator_by_name[method] - integrator = integrator_class(depletion_operator, timesteps, - **integrator_kwargs) + integrator = integrator_class(depletion_operator, **integrator_kwargs) # Now perform the depletion with openmc.lib.quiet_dll(output): @@ -456,7 +483,7 @@ class Model: Parameters ---------- - directory : str + directory : PathLike Directory to write XML files to. If it doesn't exist already, it will be created. remove_surfs : bool @@ -570,7 +597,7 @@ class Model: fh.write(ET.tostring(plots_element, encoding="unicode")) fh.write("\n") - def import_properties(self, filename): + def import_properties(self, filename: PathLike): """Import physical properties .. versionchanged:: 0.13.0 @@ -578,7 +605,7 @@ class Model: Parameters ---------- - filename : str + filename : PathLike Path to properties HDF5 file See Also @@ -631,11 +658,22 @@ class Model: C_mat = openmc.lib.materials[mat_id] C_mat.set_density(atom_density, 'atom/b-cm') - def run(self, particles=None, threads=None, geometry_debug=False, - restart_file=None, tracks=False, output=True, cwd='.', - openmc_exec='openmc', mpi_args=None, event_based=None, - export_model_xml=True, apply_tally_results=False, - **export_kwargs): + def run( + self, + particles: int | None = None, + threads: int | None = None, + geometry_debug: bool = False, + restart_file: PathLike | None = None, + tracks: bool = False, + output: bool = True, + cwd: PathLike = ".", + openmc_exec: PathLike = "openmc", + mpi_args: Iterable[str] = None, + event_based: bool | None = None, + export_model_xml: bool = True, + apply_tally_results: bool = False, + **export_kwargs, + ) -> Path: """Run OpenMC If the C API has been initialized, then the C API is used, otherwise, @@ -767,10 +805,17 @@ class Model: return last_statepoint - def calculate_volumes(self, threads=None, output=True, cwd='.', - openmc_exec='openmc', mpi_args=None, - apply_volumes=True, export_model_xml=True, - **export_kwargs): + def calculate_volumes( + self, + threads: int | None = None, + output: bool = True, + cwd: PathLike = ".", + openmc_exec: PathLike = "openmc", + mpi_args: list[str] | None = None, + apply_volumes: bool = True, + export_model_xml: bool = True, + **export_kwargs, + ): """Runs an OpenMC stochastic volume calculation and, if requested, applies volumes to the model @@ -1116,8 +1161,14 @@ class Model: """ self.tallies.add_results(statepoint) - def plot_geometry(self, output=True, cwd='.', openmc_exec='openmc', - export_model_xml=True, **export_kwargs): + def plot_geometry( + self, + output: bool = True, + cwd: PathLike = ".", + openmc_exec: PathLike = "openmc", + export_model_xml: bool = True, + **export_kwargs, + ): """Creates plot images as specified by the Model.plots attribute .. versionadded:: 0.13.0 @@ -1126,10 +1177,10 @@ class Model: ---------- output : bool, optional Capture OpenMC output from standard out - cwd : str, optional + cwd : PathLike, optional Path to working directory to run in. Defaults to the current working directory. - openmc_exec : str, optional + openmc_exec : PathLike, optional Path to OpenMC executable. Defaults to 'openmc'. This only applies to the case when not using the C API. export_model_xml : bool, optional @@ -1159,8 +1210,14 @@ class Model: openmc.plot_geometry(output=output, openmc_exec=openmc_exec, path_input=path_input) - def _change_py_lib_attribs(self, names_or_ids, value, obj_type, - attrib_name, density_units='atom/b-cm'): + def _change_py_lib_attribs( + self, + names_or_ids: Iterable[str] | Iterable[int], + value: float | Iterable[float], + obj_type: str, + attrib_name: str, + density_units: str = "atom/b-cm", + ): # Method to do the same work whether it is a cell or material and # a temperature or volume check_type('names_or_ids', names_or_ids, Iterable, (Integral, str)) @@ -1239,7 +1296,9 @@ class Model: else: setattr(lib_obj, attrib_name, value) - def rotate_cells(self, names_or_ids, vector): + def rotate_cells( + self, names_or_ids: Iterable[str] | Iterable[int], vector: Iterable[float] + ): """Rotate the identified cell(s) by the specified rotation vector. The rotation is only applied to cells filled with a universe. @@ -1261,7 +1320,9 @@ class Model: self._change_py_lib_attribs(names_or_ids, vector, 'cell', 'rotation') - def translate_cells(self, names_or_ids, vector): + def translate_cells( + self, names_or_ids: Iterable[str] | Iterable[int], vector: Iterable[float] + ): """Translate the identified cell(s) by the specified translation vector. The translation is only applied to cells filled with a universe. @@ -1284,7 +1345,12 @@ class Model: self._change_py_lib_attribs(names_or_ids, vector, 'cell', 'translation') - def update_densities(self, names_or_ids, density, density_units='atom/b-cm'): + def update_densities( + self, + names_or_ids: Iterable[str] | Iterable[int], + density: float, + density_units: str = "atom/b-cm", + ): """Update the density of a given set of materials to a new value .. note:: If applying this change to a name that is not unique, then @@ -1307,7 +1373,9 @@ class Model: self._change_py_lib_attribs(names_or_ids, density, 'material', 'density', density_units) - def update_cell_temperatures(self, names_or_ids, temperature): + def update_cell_temperatures( + self, names_or_ids: Iterable[str] | Iterable[int], temperature: float + ): """Update the temperature of a set of cells to the given value .. note:: If applying this change to a name that is not unique, then @@ -1328,7 +1396,9 @@ class Model: self._change_py_lib_attribs(names_or_ids, temperature, 'cell', 'temperature') - def update_material_volumes(self, names_or_ids, volume): + def update_material_volumes( + self, names_or_ids: Iterable[str] | Iterable[int], volume: float + ): """Update the volume of a set of materials to the given value .. note:: If applying this change to a name that is not unique, then @@ -1449,7 +1519,14 @@ class Model: self.geometry.get_all_materials().values() ) - def _generate_infinite_medium_mgxs(self, groups, nparticles, mgxs_path, correction, directory): + def _generate_infinite_medium_mgxs( + self, + groups: openmc.mgxs.EnergyGroups, + nparticles: int, + mgxs_path: PathLike, + correction: str | None, + directory: PathLike, + ): """Generate a MGXS library by running multiple OpenMC simulations, each representing an infinite medium simulation of a single isolated material. A discrete source is used to sample particles, with an equal @@ -1567,7 +1644,11 @@ class Model: mgxs_file.export_to_hdf5(mgxs_path) @staticmethod - def _create_stochastic_slab_geometry(materials, cell_thickness=1.0, num_repeats=100): + def _create_stochastic_slab_geometry( + materials: Sequence[openmc.Material], + cell_thickness: float = 1.0, + num_repeats: int = 100, + ) -> tuple[openmc.Geometry, openmc.stats.Box]: """Create a geometry representing a stochastic "sandwich" of materials in a layered slab geometry. To reduce the impact of the order of materials in the slab, the materials are applied to 'num_repeats' different randomly @@ -1636,7 +1717,14 @@ class Model: return geometry, box - def _generate_stochastic_slab_mgxs(self, groups, nparticles, mgxs_path, correction, directory) -> None: + def _generate_stochastic_slab_mgxs( + self, + groups: openmc.mgxs.EnergyGroups, + nparticles: int, + mgxs_path: PathLike, + correction: str | None, + directory: PathLike, + ) -> None: """Generate MGXS assuming a stochastic "sandwich" of materials in a layered slab geometry. While geometry-specific spatial shielding effects are not captured, this method can be useful when the geometry has materials only @@ -1741,7 +1829,14 @@ class Model: mgxs_file = mgxs_lib.create_mg_library(xs_type='macro', xsdata_names=names) mgxs_file.export_to_hdf5(mgxs_path) - def _generate_material_wise_mgxs(self, groups, nparticles, mgxs_path, correction, directory) -> None: + def _generate_material_wise_mgxs( + self, + groups: openmc.mgxs.EnergyGroups, + nparticles: int, + mgxs_path: PathLike, + correction: str | None, + directory: PathLike, + ) -> None: """Generate a material-wise MGXS library for the model by running the original continuous energy OpenMC simulation of the full material geometry and source, and tally MGXS data for each material. This method @@ -1758,12 +1853,12 @@ class Model: Energy group structure for the MGXS. nparticles : int Number of particles to simulate per batch when generating MGXS. - mgxs_path : str + mgxs_path : PathLike Filename for the MGXS HDF5 file. correction : str Transport correction to apply to the MGXS. Options are None and "P0". - directory : str + directory : PathLike Directory to run the simulation in, so as to contain XML files. """ openmc.reset_auto_ids() @@ -1831,9 +1926,15 @@ class Model: xs_type='macro', xsdata_names=names) mgxs_file.export_to_hdf5(mgxs_path) - def convert_to_multigroup(self, method="material_wise", groups='CASMO-2', - nparticles=2000, overwrite_mgxs_library=False, - mgxs_path: PathLike = "mgxs.h5", correction=None): + def convert_to_multigroup( + self, + method: str = "material_wise", + groups: str = "CASMO-2", + nparticles: int = 2000, + overwrite_mgxs_library: bool = False, + mgxs_path: PathLike = "mgxs.h5", + correction: str | None = None, + ): """Convert all materials from continuous energy to multigroup. If no MGXS data library file is found, generate one using one or more @@ -1868,7 +1969,7 @@ class Model: self.sync_dagmc_universes() self.finalize_lib() break - + # Make sure all materials have a name, and that the name is a valid HDF5 # dataset name for material in self.materials: