diff --git a/openmc/polynomial.py b/openmc/polynomial.py index abcfd8057e..82a6756c82 100644 --- a/openmc/polynomial.py +++ b/openmc/polynomial.py @@ -1,3 +1,4 @@ +import math import numpy as np import openmc from collections.abc import Iterable @@ -47,8 +48,6 @@ class ZernikeRadial(Polynomial): 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 ---------- @@ -65,7 +64,7 @@ class ZernikeRadial(Polynomial): 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) + norm_vec = (2 * np.arange(len(self.coef)) + 1) / (math.pi * radius**2) self._norm_coef = norm_vec * self.coef @property @@ -75,7 +74,67 @@ class ZernikeRadial(Polynomial): def __call__(self, r): import openmc.lib as lib if isinstance(r, Iterable): - return [np.sum(self._norm_coef * lib.calc_zn_rad(self.order, r_i)) + return [np.sum(self._norm_coef * lib.calc_zn_rad(self.order, r_i / self.radius)) for r_i in r] else: - return np.sum(self._norm_coef * lib.calc_zn_rad(self.order, r)) + return np.sum(self._norm_coef * lib.calc_zn_rad(self.order, r / self.radius)) + + +class Zernike(Polynomial): + r"""Create Zernike polynomials given coefficients and domain. + + The azimuthal 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. + + 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 + Azimuthal 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) + # Solve order from number of coefficients + # N = (order + 1)(order + 2) / 2 + self._order = int((math.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 + else: + norm_vec[j] = 2*n + 2 + norm_vec /= (math.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)) diff --git a/tests/unit_tests/test_polynomials.py b/tests/unit_tests/test_polynomials.py index f6ac2711fd..03f1fd5a46 100644 --- a/tests/unit_tests/test_polynomials.py +++ b/tests/unit_tests/test_polynomials.py @@ -24,7 +24,7 @@ def test_zernike_radial(): ref_vals = np.sum(norm_coeff * raw_zn) - test_vals = zn_rad(rho) + test_vals = zn_rad(rho * zn_rad.radius) assert ref_vals == test_vals @@ -39,6 +39,87 @@ def test_zernike_radial(): ref_vals = [np.sum(norm_coeff * raw_zn1), np.sum(norm_coeff * raw_zn2)] - test_vals = zn_rad(rho) + test_vals = zn_rad([i * zn_rad.radius for i in rho]) assert np.allclose(ref_vals, test_vals) + + +def test_zernike(): + import openmc.lib as lib + + coeff = np.asarray([1.1e-1, -3.2e2, 5.3, 7.4, -9.5, 0.005]) + zn_azimuthal = openmc.Zernike(coeff) + assert zn_azimuthal.order == 2 + assert zn_azimuthal.radius == 1 + + coeff = np.asarray([1.5, -3.6, 9.7e-1, -6.8e-1, 0.11, 0.33e2, 0.002, 13.75, + 3.1, -7.3, 7.8e-1, -1.1e-1, 2.56, 5.25e3, 0.123]) + zn_azimuthal = openmc.Zernike(coeff, 0.392) + assert zn_azimuthal.order == 4 + assert zn_azimuthal.radius == 0.392 + norm_vec = np.array([1, 4, 4, 6, 3, 6, 8, 8, 8, 8, + 10, 10, 5, 10, 10]) / (np.pi * 0.392 ** 2) + norm_coeff = norm_vec * coeff + + rho = 0.5 + + theta = np.radians(45) + # Reference solution from running the C API for calc_zn + raw_zn = lib.calc_zn(zn_azimuthal.order, rho, theta) + + ref_vals = np.sum(norm_coeff * raw_zn) + + test_vals = zn_azimuthal(rho * zn_azimuthal.radius, theta) + + assert ref_vals == test_vals + + rho = [0.2, 0.5] + + theta = np.radians(30) + #Reference solution from running the C API for calc_zn + raw_zn1 = lib.calc_zn(zn_azimuthal.order, rho[0], theta) + + raw_zn2 = lib.calc_zn(zn_azimuthal.order, rho[1], theta) + + ref_vals = [np.sum(norm_coeff * raw_zn1), np.sum(norm_coeff * raw_zn2)] + + test_vals = zn_azimuthal([i * zn_azimuthal.radius for i in rho], theta) + + assert np.allclose(ref_vals, test_vals) + + rho = 0.2 + + theta = np.radians([30, 60]) + #Reference solution from running the C API for calc_zn + raw_zn1 = lib.calc_zn(zn_azimuthal.order, rho, theta[0]) + + raw_zn2 = lib.calc_zn(zn_azimuthal.order, rho, theta[1]) + + ref_vals = [np.sum(norm_coeff * raw_zn1), np.sum(norm_coeff * raw_zn2)] + + test_vals = zn_azimuthal(rho * zn_azimuthal.radius, [j for j in theta]) + + assert np.allclose(ref_vals, test_vals) + + rho = [0.2, 0.5] + + theta = np.radians([30, 60]) + #Reference solution from running the C API for calc_zn + raw_zn1 = lib.calc_zn(zn_azimuthal.order, rho[0], theta[0]) + + raw_zn2 = lib.calc_zn(zn_azimuthal.order, rho[1], theta[0]) + + raw_zn3 = lib.calc_zn(zn_azimuthal.order, rho[0], theta[1]) + + raw_zn4 = lib.calc_zn(zn_azimuthal.order, rho[1], theta[1]) + + ref_vals = [np.sum(norm_coeff * raw_zn1), np.sum(norm_coeff * raw_zn2), + np.sum(norm_coeff * raw_zn3), np.sum(norm_coeff * raw_zn4)] + + test_vals = zn_azimuthal([i * zn_azimuthal.radius for i in rho], [j for j in theta]) + + test_vals = np.ravel(test_vals) + + assert np.allclose(ref_vals, test_vals) + +