Merge pull request #1043 from hanzhuoran/math

Add math functions for FET reconstruction

docs/source/pythonapi/base.rst

@@ -178,6 +178,23 @@ Post-processing
    openmc.StatePoint
    openmc.Summary
 
+The following classes and functions are used for functional expansion reconstruction.
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: myclass.rst 
+
+   openmc.ZernikeRadial
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: myfunction.rst
+
+   openmc.legendre_from_expcoef
+
+
 Various classes may be created when performing tally slicing and/or arithmetic:
 
 .. autosummary::

openmc/__init__.py

@@ -28,6 +28,7 @@ from openmc.particle_restart import *
 from openmc.mixin import *
 from openmc.plotter import *
 from openmc.search import *
+from openmc.polynomial import *
 from . import examples
 
 # Import a few convencience functions that used to be here

openmc/polynomial.py

@@ -0,0 +1,77 @@
+import numpy as np
+import openmc
+import openmc.capi as capi
+
+
+def legendre_from_expcoef(coef, domain= (-1,1)):
+    """Return a Legendre series object based on expansion coefficients.
+
+    Given a list of coefficients from FET tally and a array of down, return
+    the numpy Legendre object.
+
+    Parameters
+    ----------
+    coef : Iterable of float
+        A list of coefficients of each term in Legendre polynomials
+    domain : (2,) List of float
+        Domain of the Legendre polynomial
+
+    Returns
+    -------
+    numpy.polynomial.Legendre
+        A numpy Legendre series class
+
+    """
+
+    n = np.arange(len(coef))
+    c = (2*n + 1) * np.asarray(coef) / (domain[1] - domain[0])
+    return np.polynomial.Legendre(c, domain)
+
+
+class Polynomial(object):
+    """Abstract Polynomial Class for creating polynomials.
+    """
+    def __init__(self, coef):
+        self.coef = np.asarray(coef)
+
+
+class ZernikeRadial(Polynomial):
+    """Create radial only Zernike polynomials given coefficients and domain.
+
+    The radial only Zernike polynomials are defined as in
+    :class:`ZernikeRadialFilter`.
+
+    Parameters
+    ----------
+    coef : Iterable of float
+        A list of coefficients of each term in radial only Zernike polynomials
+    radius : float
+        Domain of Zernike polynomials to be applied on. Default is 1.
+    r : float
+        Position to be evaluated, normalized on radius [0,1]
+
+    Attributes
+    ----------
+    order : int
+        The maximum (even) order of Zernike polynomials.
+    radius : float
+        Domain of Zernike polynomials to be applied on. Default is 1.
+    norm_coef : iterable of float
+        The list of coefficients of each term in the polynomials after
+        normailization.
+
+    """
+    def __init__(self, coef, radius=1):
+        super().__init__(coef)
+        self._order = 2 * (len(self.coef) - 1)
+        self.radius = radius
+        norm_vec = (2 * np.arange(len(self.coef)) + 1) / (np.pi * radius**2)
+        self._norm_coef = norm_vec * self.coef
+
+    @property
+    def order(self):
+        return self._order
+
+    def __call__(self, r):
+        zn_rad = capi.calc_zn_rad(self.order, r)
+        return np.sum(self._norm_coef * zn_rad)

src/math_functions.h

@@ -93,13 +93,18 @@ extern "C" void calc_rn_c(int n, const double uvw[3], double rn[]);
 extern "C" void calc_zn_c(int n, double rho, double phi, double zn[]);
 
 //==============================================================================
-//! Calculate only the even order components of n-th order modified Zernike 
-//! polynomial moment with azimuthal dependency m = 0 for a given radial (rho)
-//! location on the unit disk.
+//! Calculate only the even radial components of n-th order modified Zernike 
+//! polynomial moment with azimuthal dependency m = 0 for a given angle 
+//! (rho, theta) location on the unit disk.
 //!
 //! Since m = 0, n could only be even orders. Z_q0 = R_q0
 //!
-//! See calc_zn_c for methodology.
+//! This procedure uses the modified Kintner's method for calculating Zernike
+//! polynomials as outlined in Chong, C. W., Raveendran, P., & Mukundan,
+//! R. (2003). A comparative analysis of algorithms for fast computation of
+//! Zernike moments. Pattern Recognition, 36(3), 731-742.
+//! The normalization of the polynomials is such that the integral of Z_pq^2
+//! over the unit disk is exactly pi.
 //!
 //! @param n       The maximum order requested
 //! @param rho     The radial parameter to specify location on the unit disk

tests/unit_tests/test_math.py

@@ -137,6 +137,20 @@ def test_calc_zn():
     assert np.allclose(ref_vals, test_vals)
 
 
+def test_calc_zn_rad():
+    n = 10
+    rho = 0.5
+
+    # Reference solution from running the Fortran implementation
+    ref_vals = np.array([
+        1.00000000e+00, -5.00000000e-01, -1.25000000e-01,
+        4.37500000e-01, -2.89062500e-01,-8.98437500e-02])
+
+    test_vals = openmc.capi.math.calc_zn_rad(n, rho)
+
+    assert np.allclose(ref_vals, test_vals)
+
+
 def test_rotate_angle():
     uvw0 = np.array([1., 0., 0.])
     phi = 0.

tests/unit_tests/test_polynomials.py

@@ -0,0 +1,29 @@
+import numpy as np
+
+import openmc
+
+
+def test_zernike_radial():
+    coeff = np.asarray([1.3, -3.0, 9e-1, -6e-1, 0.11])
+    zn_rad = openmc.ZernikeRadial(coeff)
+    assert zn_rad.order == 8
+    assert zn_rad.radius == 1
+
+    coeff = np.asarray([1.3, -3.0, 9e-1, -6e-1, 0.11, 0.222])
+    zn_rad = openmc.ZernikeRadial(coeff, 0.392)
+    assert zn_rad.order == 10
+    assert zn_rad.radius == 0.392
+    norm_vec = (2 * np.arange(6) + 1) / (np.pi * 0.392 ** 2)
+    norm_coeff = norm_vec * coeff
+
+    rho = 0.5
+    # Reference solution from running the Fortran implementation
+    raw_zn = np.array([
+        1.00000000e+00, -5.00000000e-01, -1.25000000e-01,
+        4.37500000e-01, -2.89062500e-01, -8.98437500e-02])
+
+    ref_vals = np.sum(norm_coeff * raw_zn)
+
+    test_vals = zn_rad(rho)
+
+    assert ref_vals == test_vals