Allowing CylindricalMesh and SphericalMesh to be fully made via constructor+ type hints

openmc/mesh.py

@@ -1184,6 +1184,17 @@ class CylindricalMesh(StructuredMesh):
         Unique identifier for the mesh
     name : str
         Name of the mesh
+    origin : numpy.ndarray
+        1-D array of length 3 the (x,y,z) origin of the mesh in
+        cartesian coordinates
+    r_grid : numpy.ndarray
+        1-D array of mesh boundary points along the r-axis.
+        Requirement is r >= 0.
+    phi_grid : numpy.ndarray
+        1-D array of mesh boundary points along the phi-axis in radians.
+        The default value is [0, 2π], i.e. the full phi range.
+    z_grid : numpy.ndarray
+        1-D array of mesh boundary points along the z-axis.
 
     Attributes
     ----------
@@ -1204,22 +1215,27 @@ class CylindricalMesh(StructuredMesh):
         The default value is [0, 2π], i.e. the full phi range.
     z_grid : numpy.ndarray
         1-D array of mesh boundary points along the z-axis.
-    origin : numpy.ndarray
-        1-D array of length 3 the (x,y,z) origin of the mesh in
-        cartesian coordinates
     indices : Iterable of tuple
         An iterable of mesh indices for each mesh element, e.g. [(1, 1, 1),
         (2, 1, 1), ...]
 
     """
 
-    def __init__(self, mesh_id: int = None, name: str = ''):
+    def __init__(
+        self,
+        mesh_id: int = None,
+        name: str = '',
+        origin: typing.Tuple[float] = (0., 0., 0.),
+        r_grid: typing.Optional[np.ndarray] = None,
+        phi_grid: np.ndarray = np.array([0, 2*pi]),
+        z_grid: typing.Optional[np.ndarray] = None,
+    ):
         super().__init__(mesh_id, name)
 
-        self._r_grid = None
-        self._phi_grid = [0.0, 2*pi]
-        self._z_grid = None
-        self.origin = (0., 0., 0.)
+        self._r_grid = r_grid
+        self._phi_grid = phi_grid
+        self._z_grid = z_grid
+        self.origin = origin
 
     @property
     def dimension(self):
@@ -1325,7 +1341,7 @@ def from_domain(
         dimension: typing.Sequence[int] = (10, 10, 10),
         mesh_id: typing.Optional[int] = None,
         phi_grid_bounds: typing.Sequence[float] = (0.0, 2*pi),
-        name=''
+        name: str = ''
     ):
         """Creates a regular CylindricalMesh from an existing openmc domain.
 
@@ -1484,6 +1500,18 @@ class SphericalMesh(StructuredMesh):
         Unique identifier for the mesh
     name : str
         Name of the mesh
+    r_grid : numpy.ndarray
+        1-D array of mesh boundary points along the r-axis.
+        Requirement is r >= 0.
+    theta_grid : numpy.ndarray
+        1-D array of mesh boundary points along the theta-axis in radians.
+        The default value is [0, π], i.e. the full theta range.
+    phi_grid : numpy.ndarray
+        1-D array of mesh boundary points along the phi-axis in radians.
+        The default value is [0, 2π], i.e. the full phi range.
+    origin : numpy.ndarray
+        1-D array of length 3 the (x,y,z) origin of the mesh in
+        cartesian coordinates
 
     Attributes
     ----------
@@ -1514,13 +1542,21 @@ class SphericalMesh(StructuredMesh):
 
     """
 
-    def __init__(self, mesh_id=None, name=''):
+    def __init__(
+        self,
+        mesh_id: typing.Optional[int] = None,
+        name: str = '',
+        origin: typing.Tuple[float] = (0., 0., 0.),
+        r_grid: typing.Optional[np.ndarray] = None,
+        theta_grid: np.ndarray = np.array([0, pi]),
+        phi_grid: np.ndarray = np.array([0, 2*pi])
+    ):
         super().__init__(mesh_id, name)
 
-        self._r_grid = None
-        self._theta_grid = [0, pi]
-        self._phi_grid = [0, 2*pi]
-        self.origin = (0., 0., 0.)
+        self._r_grid = r_grid
+        self._theta_grid = theta_grid
+        self._phi_grid = phi_grid
+        self.origin = origin
 
     @property
     def dimension(self):

openmc/stats/multivariate.py

@@ -1,15 +1,17 @@
+import typing
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from math import pi, cos
+from math import cos, pi
 from numbers import Real
-import lxml.etree as ET
 
+import lxml.etree as ET
 import numpy as np
 
 import openmc.checkvalue as cv
+
 from .._xml import get_text
-from .univariate import Univariate, Uniform, PowerLaw
 from ..mesh import MeshBase
+from .univariate import PowerLaw, Uniform, Univariate
 
 
 class UnitSphere(ABC):
@@ -177,7 +179,7 @@ def to_xml_element(self):
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem: ET.Element):
         """Generate isotropic distribution from an XML element
 
         Parameters
@@ -209,7 +211,7 @@ class Monodirectional(UnitSphere):
 
     """
 
-    def __init__(self, reference_uvw=[1., 0., 0.]):
+    def __init__(self, reference_uvw: typing.Sequence[float] = [1., 0., 0.]):
         super().__init__(reference_uvw)
 
     def to_xml_element(self):
@@ -228,7 +230,7 @@ def to_xml_element(self):
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem: ET.Element):
         """Generate monodirectional distribution from an XML element
 
         Parameters
@@ -304,7 +306,12 @@ class CartesianIndependent(Spatial):
 
     """
 
-    def __init__(self, x, y, z):
+    def __init__(
+        self,
+        x: 'openmc.stats.Univariate',
+        y: 'openmc.stats.Univariate',
+        z: 'openmc.stats.Univariate'
+    ):
         self.x = x
         self.y = y
         self.z = z
@@ -353,7 +360,7 @@ def to_xml_element(self):
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem: ET.Element):
         """Generate spatial distribution from an XML element
 
         Parameters
@@ -477,7 +484,7 @@ def to_xml_element(self):
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem: ET.Element):
         """Generate spatial distribution from an XML element
 
         Parameters
@@ -599,7 +606,7 @@ def to_xml_element(self):
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem: ET.Element):
         """Generate spatial distribution from an XML element
 
         Parameters
@@ -772,7 +779,12 @@ class Box(Spatial):
 
     """
 
-    def __init__(self, lower_left, upper_right, only_fissionable=False):
+    def __init__(
+        self,
+        lower_left: typing.Sequence[float],
+        upper_right: typing.Sequence[float],
+        only_fissionable: bool = False
+    ):
         self.lower_left = lower_left
         self.upper_right = upper_right
         self.only_fissionable = only_fissionable
@@ -826,7 +838,7 @@ def to_xml_element(self):
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem: ET.Element):
         """Generate box distribution from an XML element
 
         Parameters
@@ -865,7 +877,7 @@ class Point(Spatial):
 
     """
 
-    def __init__(self, xyz=(0., 0., 0.)):
+    def __init__(self, xyz: typing.Sequence[float] = (0., 0., 0.)):
         self.xyz = xyz
 
     @property
@@ -894,7 +906,7 @@ def to_xml_element(self):
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem: ET.Element):
         """Generate point distribution from an XML element
 
         Parameters
@@ -912,8 +924,13 @@ def from_xml_element(cls, elem):
         return cls(xyz)
 
 
-def spherical_uniform(r_outer, r_inner=0.0, thetas=(0., pi), phis=(0., 2*pi),
-                      origin=(0., 0., 0.)):
+def spherical_uniform(
+        r_outer: float,
+        r_inner: float = 0.0,
+        thetas: typing.Sequence[float] = (0., pi),
+        phis: typing.Sequence[float] = (0., 2*pi),
+        origin: typing.Sequence[float] = (0., 0., 0.)
+    ):
     """Return a uniform spatial distribution over a spherical shell.
 
     This function provides a uniform spatial distribution over a spherical
@@ -926,15 +943,15 @@ def spherical_uniform(r_outer, r_inner=0.0, thetas=(0., pi), phis=(0., 2*pi),
     ----------
     r_outer : float
         Outer radius of the spherical shell in [cm]
-    r_inner : float, optional
+    r_inner : float
         Inner radius of the spherical shell in [cm]
-    thetas : iterable of float, optional
+    thetas : iterable of float
         Starting and ending theta coordinates (angle relative to
         the z-axis) in radius in a reference frame centered at `origin`
-    phis : iterable of float, optional
+    phis : iterable of float
         Starting and ending phi coordinates (azimuthal angle) in
         radians in a reference frame centered at `origin`
-    origin: iterable of float, optional
+    origin: iterable of float
         Coordinates (x0, y0, z0) of the center of the spherical
         reference frame for the distribution.