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

openmc/mesh.py

@@ -1,19 +1,21 @@
+import typing
+import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from math import pi
-from numbers import Real, Integral
+from numbers import Integral, Real
 from pathlib import Path
-import typing
-import warnings
-import lxml.etree as ET
+from typing import Iterable, Optional, Sequence, Tuple
 
 import h5py
+import lxml.etree as ET
 import numpy as np
 
-import openmc.checkvalue as cv
 import openmc
-from ._xml import get_text
+import openmc.checkvalue as cv
 from openmc.checkvalue import PathLike
+
+from ._xml import get_text
 from .mixin import IDManagerMixin
 from .surface import _BOUNDARY_TYPES
 
@@ -40,7 +42,7 @@ class MeshBase(IDManagerMixin, ABC):
     next_id = 1
     used_ids = set()
 
-    def __init__(self, mesh_id: typing.Optional[int] = None, name: str = ''):
+    def __init__(self, mesh_id: Optional[int] = None, name: str = ''):
         # Initialize Mesh class attributes
         self.id = mesh_id
         self.name = name
@@ -264,7 +266,7 @@ def num_mesh_cells(self):
 
     def write_data_to_vtk(self,
                           filename: PathLike,
-                          datasets: typing.Optional[dict] = None,
+                          datasets: Optional[dict] = None,
                           volume_normalization: bool = True,
                           curvilinear: bool = False):
         """Creates a VTK object of the mesh
@@ -512,7 +514,7 @@ class RegularMesh(StructuredMesh):
 
     """
 
-    def __init__(self, mesh_id: typing.Optional[int] = None, name: str = ''):
+    def __init__(self, mesh_id: Optional[int] = None, name: str = ''):
         super().__init__(mesh_id, name)
 
         self._dimension = None
@@ -525,7 +527,7 @@ def dimension(self):
         return tuple(self._dimension)
 
     @dimension.setter
-    def dimension(self, dimension: typing.Iterable[int]):
+    def dimension(self, dimension: Iterable[int]):
         cv.check_type('mesh dimension', dimension, Iterable, Integral)
         cv.check_length('mesh dimension', dimension, 1, 3)
         self._dimension = dimension
@@ -542,7 +544,7 @@ def lower_left(self):
         return self._lower_left
 
     @lower_left.setter
-    def lower_left(self, lower_left: typing.Iterable[Real]):
+    def lower_left(self, lower_left: Iterable[Real]):
         cv.check_type('mesh lower_left', lower_left, Iterable, Real)
         cv.check_length('mesh lower_left', lower_left, 1, 3)
         self._lower_left = lower_left
@@ -562,7 +564,7 @@ def upper_right(self):
                 return [l + w * d for l, w, d in zip(ls, ws, dims)]
 
     @upper_right.setter
-    def upper_right(self, upper_right: typing.Iterable[Real]):
+    def upper_right(self, upper_right: Iterable[Real]):
         cv.check_type('mesh upper_right', upper_right, Iterable, Real)
         cv.check_length('mesh upper_right', upper_right, 1, 3)
         self._upper_right = upper_right
@@ -586,7 +588,7 @@ def width(self):
                 return [(u - l) / d for u, l, d in zip(us, ls, dims)]
 
     @width.setter
-    def width(self, width: typing.Iterable[Real]):
+    def width(self, width: Iterable[Real]):
         cv.check_type('mesh width', width, Iterable, Real)
         cv.check_length('mesh width', width, 1, 3)
         self._width = width
@@ -697,7 +699,7 @@ def from_rect_lattice(
         cls,
         lattice: 'openmc.RectLattice',
         division: int = 1,
-        mesh_id: typing.Optional[int] = None,
+        mesh_id: Optional[int] = None,
         name: str = ''
     ):
         """Create mesh from an existing rectangular lattice
@@ -736,8 +738,8 @@ def from_rect_lattice(
     def from_domain(
         cls,
         domain: typing.Union['openmc.Cell', 'openmc.Region', 'openmc.Universe', 'openmc.Geometry'],
-        dimension: typing.Sequence[int] = (10, 10, 10),
-        mesh_id: typing.Optional[int] = None,
+        dimension: Sequence[int] = (10, 10, 10),
+        mesh_id: Optional[int] = None,
         name: str = ''
     ):
         """Create mesh from an existing openmc cell, region, universe or
@@ -768,7 +770,7 @@ def from_domain(
             (openmc.Cell, openmc.Region, openmc.Universe, openmc.Geometry),
         )
 
-        mesh = cls(mesh_id, name)
+        mesh = cls(mesh_id=mesh_id, name=name)
         mesh.lower_left = domain.bounding_box[0]
         mesh.upper_right = domain.bounding_box[1]
         mesh.dimension = dimension
@@ -820,7 +822,7 @@ def from_xml_element(cls, elem: ET.Element):
 
         """
         mesh_id = int(get_text(elem, 'id'))
-        mesh = cls(mesh_id)
+        mesh = cls(mesh_id=mesh_id)
 
         mesh_type = get_text(elem, 'type')
         if mesh_type is not None:
@@ -844,7 +846,7 @@ def from_xml_element(cls, elem: ET.Element):
 
         return mesh
 
-    def build_cells(self, bc: typing.Optional[str] = None):
+    def build_cells(self, bc: Optional[str] = None):
         """Generates a lattice of universes with the same dimensionality
         as the mesh object.  The individual cells/universes produced
         will not have material definitions applied and so downstream code
@@ -1119,7 +1121,7 @@ def from_hdf5(cls, group: h5py.Group):
         mesh_id = int(group.name.split('/')[-1].lstrip('mesh '))
 
         # Read and assign mesh properties
-        mesh = cls(mesh_id)
+        mesh = cls(mesh_id=mesh_id)
         mesh.x_grid = group['x_grid'][()]
         mesh.y_grid = group['y_grid'][()]
         mesh.z_grid = group['z_grid'][()]
@@ -1141,8 +1143,8 @@ def from_xml_element(cls, elem: ET.Element):
             Rectilinear mesh object
 
         """
-        id = int(get_text(elem, 'id'))
-        mesh = cls(id)
+        mesh_id = int(get_text(elem, 'id'))
+        mesh = cls(mesh_id=mesh_id)
         mesh.x_grid = [float(x) for x in get_text(elem, 'x_grid').split()]
         mesh.y_grid = [float(y) for y in get_text(elem, 'y_grid').split()]
         mesh.z_grid = [float(z) for z in get_text(elem, 'z_grid').split()]
@@ -1184,6 +1186,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 +1217,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,
+        origin: Tuple[float] = (0., 0., 0.),
+        r_grid: Sequence[float] = (0, 10),
+        phi_grid: Sequence[float] = (0, 2*pi),
+        z_grid: Sequence[float] = (0, 10),
+        mesh_id: int = None,
+        name: str = '',
+    ):
         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):
@@ -1309,7 +1327,7 @@ def from_hdf5(cls, group: h5py.Group):
         mesh_id = int(group.name.split('/')[-1].lstrip('mesh '))
 
         # Read and assign mesh properties
-        mesh = cls(mesh_id)
+        mesh = cls(mesh_id=mesh_id)
         mesh.r_grid = group['r_grid'][()]
         mesh.phi_grid = group['phi_grid'][()]
         mesh.z_grid = group['z_grid'][()]
@@ -1322,10 +1340,10 @@ def from_hdf5(cls, group: h5py.Group):
     def from_domain(
         cls,
         domain: typing.Union['openmc.Cell', 'openmc.Region', 'openmc.Universe', 'openmc.Geometry'],
-        dimension: typing.Sequence[int] = (10, 10, 10),
-        mesh_id: typing.Optional[int] = None,
-        phi_grid_bounds: typing.Sequence[float] = (0.0, 2*pi),
-        name=''
+        dimension: Sequence[int] = (10, 10, 10),
+        mesh_id: Optional[int] = None,
+        phi_grid_bounds: Sequence[float] = (0.0, 2*pi),
+        name: str = ''
     ):
         """Creates a regular CylindricalMesh from an existing openmc domain.
 
@@ -1358,7 +1376,7 @@ def from_domain(
             (openmc.Cell, openmc.Region, openmc.Universe, openmc.Geometry),
         )
 
-        mesh = cls(mesh_id, name)
+        mesh = cls(mesh_id=mesh_id, name=name)
 
         # loaded once to avoid reading h5m file repeatedly
         cached_bb = domain.bounding_box
@@ -1433,7 +1451,7 @@ def from_xml_element(cls, elem: ET.Element):
         """
 
         mesh_id = int(get_text(elem, 'id'))
-        mesh = cls(mesh_id)
+        mesh = cls(mesh_id=mesh_id)
         mesh.r_grid = [float(x) for x in get_text(elem, "r_grid").split()]
         mesh.phi_grid = [float(x) for x in get_text(elem, "phi_grid").split()]
         mesh.z_grid = [float(x) for x in get_text(elem, "z_grid").split()]
@@ -1463,7 +1481,7 @@ def cartesian_vertices(self):
         return self._convert_to_cartesian(self.vertices, self.origin)
 
     @staticmethod
-    def _convert_to_cartesian(arr, origin: typing.Sequence[float]):
+    def _convert_to_cartesian(arr, origin: Sequence[float]):
         """Converts an array with xyz values in the first dimension (shape (3, ...))
         to Cartesian coordinates.
         """
@@ -1484,6 +1502,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 +1544,21 @@ class SphericalMesh(StructuredMesh):
 
     """
 
-    def __init__(self, mesh_id=None, name=''):
+    def __init__(
+        self,
+        origin: Tuple[float] = (0., 0., 0.),
+        r_grid: Sequence[float] = (0, 10),
+        phi_grid: Sequence[float] = (0, 2*pi),
+        theta_grid: Sequence[float] = (0, pi),
+        mesh_id: Optional[int] = None,
+        name: str = '',
+    ):
         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):
@@ -1610,7 +1648,7 @@ def from_hdf5(cls, group: h5py.Group):
         mesh_id = int(group.name.split('/')[-1].lstrip('mesh '))
 
         # Read and assign mesh properties
-        mesh = cls(mesh_id)
+        mesh = cls(mesh_id=mesh_id)
         mesh.r_grid = group['r_grid'][()]
         mesh.theta_grid = group['theta_grid'][()]
         mesh.phi_grid = group['phi_grid'][()]
@@ -1663,7 +1701,7 @@ def from_xml_element(cls, elem: ET.Element):
 
         """
         mesh_id = int(get_text(elem, 'id'))
-        mesh = cls(mesh_id)
+        mesh = cls(mesh_id=mesh_id)
         mesh.r_grid = [float(x) for x in get_text(elem, "r_grid").split()]
         mesh.theta_grid = [float(x) for x in get_text(elem, "theta_grid").split()]
         mesh.phi_grid = [float(x) for x in get_text(elem, "phi_grid").split()]
@@ -1693,7 +1731,7 @@ def cartesian_vertices(self):
         return self._convert_to_cartesian(self.vertices, self.origin)
 
     @staticmethod
-    def _convert_to_cartesian(arr, origin: typing.Sequence[float]):
+    def _convert_to_cartesian(arr, origin: Sequence[float]):
         """Converts an array with xyz values in the first dimension (shape (3, ...))
         to Cartesian coordinates.
         """
@@ -1768,7 +1806,7 @@ class UnstructuredMesh(MeshBase):
     _LINEAR_TET = 0
     _LINEAR_HEX = 1
 
-    def __init__(self, filename: PathLike, library: str, mesh_id: typing.Optional[int] = None,
+    def __init__(self, filename: PathLike, library: str, mesh_id: Optional[int] = None,
                  name: str = '', length_multiplier: float = 1.0):
         super().__init__(mesh_id, name)
         self.filename = filename
@@ -1830,7 +1868,7 @@ def volumes(self):
         return self._volumes
 
     @volumes.setter
-    def volumes(self, volumes: typing.Iterable[Real]):
+    def volumes(self, volumes: Iterable[Real]):
         cv.check_type("Unstructured mesh volumes", volumes, Iterable, Real)
         self._volumes = volumes
 
@@ -1940,8 +1978,8 @@ def write_vtk_mesh(self, **kwargs):
 
     def write_data_to_vtk(
             self,
-            filename: typing.Optional[PathLike] = None,
-            datasets: typing.Optional[dict] = None,
+            filename: Optional[PathLike] = None,
+            datasets: Optional[dict] = None,
             volume_normalization: bool = True
     ):
         """Map data to unstructured VTK mesh elements.
@@ -2040,7 +2078,7 @@ def from_hdf5(cls, group: h5py.Group):
         filename = group['filename'][()].decode()
         library = group['library'][()].decode()
 
-        mesh = cls(filename, library, mesh_id=mesh_id)
+        mesh = cls(filename=filename, library=library, mesh_id=mesh_id)
         vol_data = group['volumes'][()]
         mesh.volumes = np.reshape(vol_data, (vol_data.shape[0],))
         mesh.n_elements = mesh.volumes.size