#!/usr/bin/env python
u"""
math.py
Written by Tyler Sutterley (04/2025)
Special functions of mathematical physics
PYTHON DEPENDENCIES:
numpy: Scientific Computing Tools For Python
https://numpy.org
https://numpy.org/doc/stable/user/numpy-for-matlab-users.html
scipy: Scientific Tools for Python
https://docs.scipy.org/doc/
UPDATE HISTORY:
Updated 04/2025: use numpy power function over using pow for consistency
Updated 01/2025: added function for fully-normalized Legendre polynomials
Updated 12/2024: added function to calculate an aliasing frequency
Written 11/2024
"""
from __future__ import annotations
import numpy as np
from scipy.special import factorial
__all__ = [
"polynomial_sum",
"normalize_angle",
"rotate",
"aliasing",
"legendre",
"assoc_legendre",
"sph_harm"
]
# PURPOSE: calculate the sum of a polynomial function of time
[docs]
def polynomial_sum(
coefficients: list | np.ndarray,
t: np.ndarray
):
"""
Calculates the sum of a polynomial function using Horner's method
:cite:p:`Horner:1819br`
Parameters
----------
coefficients: list or np.ndarray
leading coefficient of polynomials of increasing order
t: np.ndarray
delta time in units for a given astronomical longitudes calculation
"""
# convert time to array if importing a single value
t = np.atleast_1d(t)
return np.sum([c * (t ** i) for i, c in enumerate(coefficients)], axis=0)
[docs]
def normalize_angle(
theta: float | np.ndarray,
circle: float = 360.0
):
"""
Normalize an angle to a single rotation
Parameters
----------
theta: float or np.ndarray
Angle to normalize
circle: float, default 360.0
Circle of the angle
"""
return np.mod(theta, circle)
[docs]
def rotate(
theta: float | np.ndarray,
axis: str = 'x'
):
"""
Rotate a 3-dimensional matrix about a given axis
Parameters
----------
theta: float or np.ndarray
Angle of rotation in radians
axis: str, default 'x'
Axis of rotation (``'x'``, ``'y'``, or ``'z'``)
"""
# allocate for output rotation matrix
R = np.zeros((3, 3, len(np.atleast_1d(theta))))
if (axis.lower() == 'x'):
# rotate about x-axis
R[0,0,:] = 1.0
R[1,1,:] = np.cos(theta)
R[1,2,:] = np.sin(theta)
R[2,1,:] = -np.sin(theta)
R[2,2,:] = np.cos(theta)
elif (axis.lower() == 'y'):
# rotate about y-axis
R[0,0,:] = np.cos(theta)
R[0,2,:] = -np.sin(theta)
R[1,1,:] = 1.0
R[2,0,:] = np.sin(theta)
R[2,2,:] = np.cos(theta)
elif (axis.lower() == 'z'):
# rotate about z-axis
R[0,0,:] = np.cos(theta)
R[0,1,:] = np.sin(theta)
R[1,0,:] = -np.sin(theta)
R[1,1,:] = np.cos(theta)
R[2,2,:] = 1.0
else:
raise ValueError(f'Invalid axis {axis}')
# return the rotation matrix
return R
[docs]
def aliasing(
f: float,
fs: float
) -> float:
"""
Calculate the aliasing frequency of a signal
Parameters
----------
f: float
Frequency of the signal
fs: float
Sampling frequency of the signal
Returns
-------
fa: float
Aliasing frequency of the signal
"""
fa = np.abs(f - fs*np.round(f/fs))
return fa
[docs]
def legendre(
l: int,
x: np.ndarray,
m: int = 0
):
"""
Computes associated Legendre functions for a particular degree
and order :cite:p:`Munk:1966go,HofmannWellenhof:2006hy`
Parameters
----------
l: int
degree of the Legendre polynomials (0 to 3)
x: np.ndarray
elements ranging from -1 to 1
Typically ``cos(theta)``, where ``theta`` is the colatitude in radians
m: int, default 0
order of the Legendre polynomials (0 to ``l``)
Returns
-------
Plm: np.ndarray
Legendre polynomials of degree ``l`` and order ``m``
"""
# verify values are integers
l = np.int64(l)
m = np.int64(m)
# assert values
assert (l >= 0) and (l <= 3), 'Degree must be between 0 and 3'
assert (m >= 0) and (m <= l), 'Order must be between 0 and l'
# verify dimensions
singular_values = (np.ndim(x) == 0)
x = np.atleast_1d(x).flatten()
# if x is the cos of colatitude, u is the sine
u = np.sqrt(1.0 - x**2)
# size of the x array
nx = len(x)
# complete matrix of associated legendre functions
# up to degree and order 3
Plm = np.zeros((4, 4, nx), dtype=np.float64)
# since tides only use low-degree harmonics:
# functions are hard coded rather than using a recursion relation
Plm[0, 0, :] = 1.0
Plm[1, 0, :] = x
Plm[1, 1, :] = u
Plm[2, 0, :] = 0.5*(3.0*x**2 - 1.0)
Plm[2, 1, :] = 3.0*x*u
Plm[2, 2, :] = 3.0*u**2
Plm[3, 0, :] = 0.5*(5.0*x**3 - 3.0*x)
Plm[3, 1, :] = 1.5*(5.0*x**2 - 1.0)*u
Plm[3, 2, :] = 15.0*x*u**2
Plm[3, 3, :] = 15.0*u**3
# return values
if singular_values:
return np.power(-1.0, m)*Plm[l, m, 0]
else:
return np.power(-1.0, m)*Plm[l, m, :]
[docs]
def assoc_legendre(lmax, x):
"""
Computes fully-normalized associated Legendre Polynomials using a
standard forward-column method :cite:p:`Colombo:1981vh`
:cite:p:`HofmannWellenhof:2006hy`
Parameters
----------
lmax: int
maximum degree and order of Legendre polynomials
x: np.ndarray
elements ranging from -1 to 1
Typically ``cos(theta)``, where ``theta`` is the colatitude in radians
Returns
-------
Plm: np.ndarray
fully-normalized Legendre polynomials
"""
# verify values are integers
lmax = np.int64(lmax)
# verify dimensions
singular_values = (np.ndim(x) == 0)
x = np.atleast_1d(x).flatten()
# if x is the cos of colatitude, u is the sine
u = np.sqrt(1.0 - x**2)
# size of the x array
nx = len(x)
# allocate for associated legendre functions
Plm = np.zeros((lmax+1,lmax+1,nx))
# initial polynomials for the recursion
Plm[0,0,:] = 1.0
Plm[1,0,:] = np.sqrt(3.0)*x
Plm[1,1,:] = np.sqrt(3.0)*u
for l in range(2, lmax+1):
# normalization factor
norm = np.sqrt(2.0*l+1.0)
for m in range(0, l):
# zonal and tesseral terms (non-sectorial)
a = np.sqrt((2.0*l-1.0)/((l-m)*(l+m)))
b = np.sqrt((l+m-1.0)*(l-m-1.0)/((l-m)*(l+m)*(2.0*l-3.0)))
Plm[l,m,:] = a*norm*x*Plm[l-1,m,:] - b*norm*Plm[l-2,m,:]
# sectorial terms: serve as seed values for the recursion
# starting with P00 and P11 (outside the loop)
Plm[l,l,:] = u*norm*np.sqrt(1.0/(2.0*l))*Plm[l-1,l-1,:]
# return values
if singular_values:
return Plm[:, :, 0]
else:
return Plm
[docs]
def sph_harm(
l: int,
theta: np.ndarray,
phi: np.ndarray,
m: int = 0
):
"""
Computes the spherical harmonics for a particular degree
and order :cite:p:`Munk:1966go,HofmannWellenhof:2006hy`
Parameters
----------
l: int
degree of the spherical harmonics (0 to 3)
theta: np.ndarray
colatitude in radians
phi: np.ndarray
longitude in radians
m: int, default 0
order of the spherical harmonics (0 to ``l``)
Returns
-------
Ylm: np.ndarray
complex spherical harmonics of degree ``l`` and order ``m``
"""
# verify dimensions
singular_values = (np.ndim(theta) == 0)
theta = np.atleast_1d(theta).flatten()
phi = np.atleast_1d(phi).flatten()
# assert dimensions
assert len(theta) == len(phi), 'coordinates must have the same dimensions'
# normalize associated Legendre functions
# following Munk and Cartwright (1966)
norm = np.sqrt(factorial(l - m)/factorial(l + m))
Plm = norm*legendre(l, np.cos(theta), m=m)
# spherical harmonics of degree l and order m
dfactor = np.sqrt((2.0*l + 1.0)/(4.0*np.pi))
Ylm = dfactor*Plm*np.sin(theta)*np.exp(1j*m*phi)
# return values
if singular_values:
return Ylm[0]
else:
return Ylm