OpenMC/openmc/lib/math.py
lhchg 3bedd043d0
update math function unit test with catch2 (#2955)
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
2024-06-19 15:56:02 +00:00

61 lines
1.6 KiB
Python

from ctypes import c_int, c_double
import numpy as np
from numpy.ctypeslib import ndpointer
from . import _dll
_dll.calc_zn.restype = None
_dll.calc_zn.argtypes = [c_int, c_double, c_double, ndpointer(c_double)]
_dll.calc_zn_rad.restype = None
_dll.calc_zn_rad.argtypes = [c_int, c_double, ndpointer(c_double)]
def calc_zn(n, rho, phi):
""" Calculate the n-th order modified Zernike polynomial moment for a
given angle (rho, theta) location in the unit disk. The normalization of
the polynomials is such that the integral of Z_pq*Z_pq over the unit disk
is exactly pi
Parameters
----------
n : int
Maximum order
rho : float
Radial location in the unit disk
phi : float
Theta (radians) location in the unit disk
Returns
-------
numpy.ndarray
Corresponding resulting list of coefficients
"""
num_bins = ((n + 1) * (n + 2)) // 2
zn = np.zeros(num_bins, dtype=np.float64)
_dll.calc_zn(n, rho, phi, zn)
return zn
def calc_zn_rad(n, rho):
""" Calculate the even orders in n-th order modified Zernike polynomial
moment with no azimuthal dependency (m=0) for a given radial location in
the unit disk. The normalization of the polynomials is such that the
integral of Z_pq*Z_pq over the unit disk is exactly pi.
Parameters
----------
n : int
Maximum order
rho : float
Radial location in the unit disk
Returns
-------
numpy.ndarray
Corresponding resulting list of coefficients
"""
num_bins = n // 2 + 1
zn_rad = np.zeros(num_bins, dtype=np.float64)
_dll.calc_zn_rad(n, rho, zn_rad)
return zn_rad