Fix code style

openmc/filter_expansion.py

@@ -325,8 +325,8 @@ class ZernikeFilter(ExpansionFilter):
 
     This filter allows scores to be multiplied by Zernike polynomials of the
     particle's position normalized to a given unit circle, up to a
-    user-specified order. The standard Zernike polynomials follow the definition by
-    Born and Wolf, *Principles of Optics* and are defined as
+    user-specified order. The standard Zernike polynomials follow the
+    definition by Born and Wolf, *Principles of Optics* and are defined as
 
     .. math::
         Z_n^m(\rho, \theta) = R_n^m(\rho) \cos (m\theta), \quad m > 0
@@ -342,7 +342,7 @@ class ZernikeFilter(ExpansionFilter):
         \frac{n+m}{2} - k)! (\frac{n-m}{2} - k)!} \rho^{n-2k}.
 
     With this definition, the integral of :math:`(Z_n^m)^2` over the unit disk
-    is :math:`\frac{\epsilon_m\pi}{2n+2}` for each polynomial where 
+    is :math:`\frac{\epsilon_m\pi}{2n+2}` for each polynomial where
     :math:`\epsilon_m` is 2 if :math:`m` equals 0 and 1 otherwise.
 
     Specifying a filter with order N tallies moments for all :math:`n` from 0
@@ -488,7 +488,7 @@ class ZernikeRadialFilter(ZernikeFilter):
     is :math:`\frac{\pi}{n+1}`.
 
     If there is only radial dependency, the polynomials are integrated over
-    the azimuthal angles. The only terms left are :math:`Z_n^{0}(\rho, \theta) 
+    the azimuthal angles. The only terms left are :math:`Z_n^{0}(\rho, \theta)
     = R_n^{0}(\rho)`. Note that :math:`n` could only be even orders.
     Therefore, for a radial Zernike polynomials up to order of :math:`n`,
     there are :math:`\frac{n}{2} + 1` terms in total. The indexing is from the

openmc/polynomial.py

@@ -1,9 +1,10 @@
 import numpy as np
 import openmc
 import openmc.capi as capi
+from collections import Iterable
 
 
-def legendre_from_expcoef(coef, domain= (-1,1)):
+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
@@ -73,5 +74,8 @@ class ZernikeRadial(Polynomial):
         return self._order
 
     def __call__(self, r):
-        zn_rad = capi.calc_zn_rad(self.order, r)
-        return np.sum(self._norm_coef * zn_rad)
+        if isinstance(r, Iterable):
+            return [np.sum(self._norm_coef * capi.calc_zn_rad(self.order, r_i))
+                    for r_i in r]
+        else:
+            return np.sum(self._norm_coef * capi.calc_zn_rad(self.order, r))

tests/unit_tests/test_polynomials.py

@@ -27,3 +27,21 @@ def test_zernike_radial():
     test_vals = zn_rad(rho)
 
     assert ref_vals == test_vals
+
+    rho = [0.2, 0.5]
+    # Reference solution from running the Fortran implementation
+    raw_zn1 = np.array([
+        1.00000000e+00, -9.20000000e-01, 7.69600000e-01,
+        -5.66720000e-01, 3.35219200e-01, -1.01747000e-01])
+    raw_zn2 = 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_zn1), np.sum(norm_coeff * raw_zn2)]
+
+    test_vals = zn_rad(rho)
+
+    print(ref_vals)
+    print(test_vals)
+
+    assert np.allclose(ref_vals, test_vals)