From c1c6865026b687a2904e08eea38dc9329e8d40de Mon Sep 17 00:00:00 2001 From: Christina Cai Date: Wed, 12 Apr 2023 14:00:48 -0400 Subject: [PATCH] added type hinting for weight_windows.py --- openmc/weight_windows.py | 56 ++++++++++++++++++++++------------------ 1 file changed, 31 insertions(+), 25 deletions(-) diff --git a/openmc/weight_windows.py b/openmc/weight_windows.py index 65b9b7bb08..b825796dcd 100644 --- a/openmc/weight_windows.py +++ b/openmc/weight_windows.py @@ -1,10 +1,12 @@ from __future__ import annotations from collections.abc import Iterable from numbers import Real, Integral -from typing import Iterable, List +from typing import Iterable, List, Optional, Union +import pathlib from xml.etree import ElementTree as ET import numpy as np +import h5py from openmc.filter import _PARTICLES from openmc.mesh import MeshBase, RectilinearMesh, UnstructuredMesh @@ -18,7 +20,7 @@ class WeightWindows(IDManagerMixin): """Mesh-based weight windows This class enables you to specify weight window parameters that are used in - a simulation. Multiple sets of weight windows can be defined for different + a simulation. Multiple sets of weight windows can be defined for differen t meshes and different particles. An iterable of :class:`WeightWindows` instances can be assigned to the :attr:`openmc.Settings.weight_windows` attribute, which is then exported to XML. @@ -103,16 +105,20 @@ class WeightWindows(IDManagerMixin): next_id = 1 used_ids = set() - def __init__(self, mesh, lower_ww_bounds, - upper_ww_bounds=None, - upper_bound_ratio=None, - energy_bounds=None, - particle_type='neutron', - survival_ratio=3, - max_lower_bound_ratio=None, - max_split=10, - weight_cutoff=1.e-38, - id=None): + def __init__( + self, + mesh: MeshBase, + lower_ww_bounds: np.ndarray, + upper_ww_bounds: Optional[np.ndarray] = None, + upper_bound_ratio: Optional[float] = None, + energy_bounds: Optional[Iterable[Real]] = None, + particle_type: str = 'neutron', + survival_ratio: float = 3, + max_lower_bound_ratio: Optional[float] = None, + max_split: int = 10, + weight_cutoff: float = 1.e-38, + id: Optional[int] = None + ): self.mesh = mesh self.id = id self.particle_type = particle_type @@ -161,7 +167,7 @@ class WeightWindows(IDManagerMixin): string += '{: <16}=\t{}\n'.format('\tWeight Cutoff', self._weight_cutoff) return string - def __eq__(self, other) -> bool: + def __eq__(self, other: WeightWindows) -> bool: # ensure that `other` is a WeightWindows object if not isinstance(other, WeightWindows): return False @@ -195,7 +201,7 @@ class WeightWindows(IDManagerMixin): return self._mesh @mesh.setter - def mesh(self, mesh): + def mesh(self, mesh: MeshBase): cv.check_type('Weight window mesh', mesh, MeshBase) self._mesh = mesh @@ -204,7 +210,7 @@ class WeightWindows(IDManagerMixin): return self._particle_type @particle_type.setter - def particle_type(self, pt): + def particle_type(self, pt: str): cv.check_value('Particle type', pt, _PARTICLES) self._particle_type = pt @@ -213,7 +219,7 @@ class WeightWindows(IDManagerMixin): return self._energy_bounds @energy_bounds.setter - def energy_bounds(self, bounds): + def energy_bounds(self, bounds: Iterable[Real]): cv.check_type('Energy bounds', bounds, Iterable, Real) self._energy_bounds = np.asarray(bounds) @@ -228,7 +234,7 @@ class WeightWindows(IDManagerMixin): return self._lower_ww_bounds @lower_ww_bounds.setter - def lower_ww_bounds(self, bounds): + def lower_ww_bounds(self, bounds: Iterable[Real]): cv.check_iterable_type('Lower WW bounds', bounds, Real, @@ -247,7 +253,7 @@ class WeightWindows(IDManagerMixin): return self._upper_ww_bounds @upper_ww_bounds.setter - def upper_ww_bounds(self, bounds): + def upper_ww_bounds(self, bounds: Iterable[Real]): cv.check_iterable_type('Upper WW bounds', bounds, Real, @@ -266,7 +272,7 @@ class WeightWindows(IDManagerMixin): return self._survival_ratio @survival_ratio.setter - def survival_ratio(self, val): + def survival_ratio(self, val: float): cv.check_type('Survival ratio', val, Real) cv.check_greater_than('Survival ratio', val, 1.0, True) self._survival_ratio = val @@ -276,7 +282,7 @@ class WeightWindows(IDManagerMixin): return self._max_lower_bound_ratio @max_lower_bound_ratio.setter - def max_lower_bound_ratio(self, val): + def max_lower_bound_ratio(self, val: float): cv.check_type('Maximum lower bound ratio', val, Real) cv.check_greater_than('Maximum lower bound ratio', val, 1.0) self._max_lower_bound_ratio = val @@ -286,7 +292,7 @@ class WeightWindows(IDManagerMixin): return self._max_split @max_split.setter - def max_split(self, val): + def max_split(self, val: int): cv.check_type('Max split', val, Integral) self._max_split = val @@ -295,7 +301,7 @@ class WeightWindows(IDManagerMixin): return self._weight_cutoff @weight_cutoff.setter - def weight_cutoff(self, cutoff): + def weight_cutoff(self, cutoff: float): cv.check_type('Weight cutoff', cutoff, Real) cv.check_greater_than('Weight cutoff', cutoff, 0.0, True) self._weight_cutoff = cutoff @@ -343,7 +349,7 @@ class WeightWindows(IDManagerMixin): return element @classmethod - def from_xml_element(cls, elem, root) -> WeightWindows: + def from_xml_element(cls, elem: ET.Element, root: ET.Element) -> WeightWindows: """Generate weight window settings from an XML element Parameters @@ -398,7 +404,7 @@ class WeightWindows(IDManagerMixin): ) @classmethod - def from_hdf5(cls, group, meshes) -> WeightWindows: + def from_hdf5(cls, group: h5py.Group, meshes: dict) -> WeightWindows: """Create weight windows from HDF5 group Parameters @@ -443,7 +449,7 @@ class WeightWindows(IDManagerMixin): ) -def wwinp_to_wws(path) -> List[WeightWindows]: +def wwinp_to_wws(path: Union[str, pathlib.Path]) -> List[WeightWindows]: """Create WeightWindows instances from a wwinp file .. versionadded:: 0.13.1