add reconstruction for Azimuthal Zernike

openmc/polynomial.py

@@ -81,3 +81,63 @@ class ZernikeRadial(Polynomial):
             return np.sum(self._norm_coef * lib.calc_zn_rad(self.order, r / self.radius))
 
 
+class Zernike(Polynomial):
+    """Create Zernike polynomials given coefficients and domain.
+    The Zernike polynomials are defined as in
+    :class:`ZernikeFilter`.
+    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]
+    theta : float 
+        Polar angular position to be evaluated, normalized on :math:`[0, 2*\pi]`
+        
+    Attributes
+    ----------
+    order : int
+        The maximum (even) order of Zernike polynomials.
+    radius : float
+        Domain of Zernike polynomials to be applied on. Default is 1.
+    theta : float 
+        Polar angle of Zernike polynomial to be applied on. Default is 0.    
+    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 = int(((np.sqrt(8 * len(self.coef) + 1)) -3) / 2)
+        self.radius = radius
+        norm_vec = np.ones((len(self.coef)))
+        for n in range(self._order + 1):
+            for m in range(-n,(n+1),2):
+                j = int((n * (n + 2) + m) / 2)
+                if m==0:
+                    norm_vec[j] = n + 1.0 
+                else:
+                    norm_vec[j] = 2 * n + 2
+        norm_vec /= (np.pi * radius**2)
+        self._norm_coef = norm_vec * self.coef
+
+    @property
+    def order(self):
+        return self._order
+
+    def __call__(self, r, theta=0.0):
+        import openmc.lib as lib
+        if isinstance(r, Iterable) and isinstance(theta, Iterable):
+            return [[np.sum(self._norm_coef * lib.calc_zn(self.order, r_i/self.radius, theta_i))
+                    for r_i in r] for theta_i in theta]
+        elif isinstance(r, Iterable) and not isinstance(theta, Iterable):
+            return [np.sum(self._norm_coef * lib.calc_zn(self.order, r_i/self.radius, theta))
+                    for r_i in r]
+        elif not isinstance(r, Iterable) and isinstance(theta, Iterable):
+            return [np.sum(self._norm_coef * lib.calc_zn(self.order, r/self.radius, theta_i))
+                    for theta_i in theta]
+        else:
+            return np.sum(self._norm_coef * lib.calc_zn(self.order, r/self.radius, theta))
+