Skip to content

Commit

Permalink
Allowing CylindricalMesh and SphericalMesh to be fully made via const…
Browse files Browse the repository at this point in the history
…ructor+ type hints (#2619)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
  • Loading branch information
shimwell and paulromano authored Aug 5, 2023
1 parent cd13d34 commit 557880a
Show file tree
Hide file tree
Showing 13 changed files with 366 additions and 207 deletions.
176 changes: 113 additions & 63 deletions openmc/mesh.py

Large diffs are not rendered by default.

57 changes: 37 additions & 20 deletions openmc/stats/multivariate.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,17 @@
from __future__ import annotations
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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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.
Expand Down
Loading

0 comments on commit 557880a

Please sign in to comment.