diff --git a/python/taichi/math/__init__.py b/python/taichi/math/__init__.py index c924fc6c9dd45..e0d05130f5286 100644 --- a/python/taichi/math/__init__.py +++ b/python/taichi/math/__init__.py @@ -5,5 +5,6 @@ from ._complex import * from .mathimpl import * # pylint: disable=W0622 +from .polynomial import * del mathimpl diff --git a/python/taichi/math/polynomial.py b/python/taichi/math/polynomial.py new file mode 100644 index 0000000000000..0b724e6edecc5 --- /dev/null +++ b/python/taichi/math/polynomial.py @@ -0,0 +1,589 @@ +""" +Math polynomial module. +""" + +from taichi.lang.kernel_impl import func +from taichi.types import template +from taichi.lang.ops import cos, sin, exp, sqrt +from taichi.lang.impl import static + + +@func +def lagval(x_field: template(), c_field: template()): + """ + Evaluate a Laguerre series in-place on a 1D Taichi array. + + For coefficients ``c_field`` of length ``n + 1``, this computes + + .. math:: + p(x) = \\sum_{k=0}^{n} c_k L_k(x) + + for every element in ``x_field`` and writes the result back into + ``x_field``. + + This function is intended to be called inside a Taichi kernel, + much like other math functions in Taichi's math module. + The evaluation uses Clenshaw recursion, and closely + follows the numpy lagval implementation. + + Parameters + ---------- + x_field : template + 1D mutable container (for example ``ti.ndarray`` or templated field) + containing the input ``x`` values. Updated in-place. + c_field : template + 1D coefficient container where ``c_field[k]`` is the coefficient + of :math:`L_k(x)`. + + Returns + ------- + template + The same container as ``x_field`` after in-place update. + """ + + c_len = c_field.shape[0] + + if c_len == 1: + for j in range(x_field.shape[0]): + x_field[j] = c_field[0] + elif c_len == 2: + for j in range(x_field.shape[0]): + x_field[j] = c_field[0] + c_field[1] * (1 - x_field[j]) + else: + for j in range(x_field.shape[0]): + nd = c_len + c0 = c_field[c_len - 2] + c1 = c_field[c_len - 1] + for i in range(3, c_len + 1): + tmp = c0 + nd = nd - 1 + c0 = c_field[c_len - i] - (c1 * (nd - 1)) / nd + c1 = tmp + (c1 * ((2 * nd - 1) - x_field[j])) / nd + x_field[j] = c0 + c1 * (1 - x_field[j]) + + return x_field + + +@func +def hermval(x_field: template(), c_field: template()): + """ + Evaluate an Hermite series in-place on a 1D Taichi array. + + For coefficients ``c_field`` of length ``n + 1``, this computes + + .. math:: + p(x) = \\sum_{k=0}^{n} c_k H_k(x) + + for every element in ``x_field`` and writes the result back into + ``x_field``. + + This function is intended to be called inside a Taichi kernel, + much like other math functions in Taichi's math module. + The evaluation uses Clenshaw recursion, and closely + follows the numpy hermval implementation. + + Parameters + ---------- + x_field : template + 1D mutable container (for example ``ti.ndarray`` or templated field) + containing the input ``x`` values. Updated in-place. + c_field : template + 1D coefficient container where ``c_field[k]`` is the coefficient + of :math:`H_k(x)`. + + Returns + ------- + template + The same container as ``x_field`` after in-place update. + + """ + c_len = c_field.shape[0] + if c_len == 1: + for j in range(x_field.shape[0]): + x_field[j] = c_field[0] + elif c_len == 2: + for j in range(x_field.shape[0]): + x_field[j] = c_field[0] + c_field[1] * x_field[j] * 2 + else: + for j in range(x_field.shape[0]): + nd = c_len + c0 = c_field[c_len - 2] + c1 = c_field[c_len - 1] + for i in range(3, c_len + 1): + tmp = c0 + nd = nd - 1 + c0 = c_field[c_len - i] - c1 * (2 * (nd - 1)) + c1 = tmp + c1 * x_field[j] * 2 + x_field[j] = c0 + c1 * x_field[j] * 2 + return x_field + + +@func +def chebval(x_field: template(), c_field: template()): + """ + Evaluate a Chebyshev series in-place on a 1D Taichi array. + + For coefficients ``c_field`` of length ``n + 1``, this computes + + .. math:: p(x) = \\sum_{k=0}^{n} c_k T_k(x) + + for every element in ``x_field`` and writes the result back into + ``x_field``. + + This function is intended to be called inside a Taichi kernel, + much like other math functions in Taichi's math module. + The evaluation uses Clenshaw recursion, and closely + follows the numpy chebval implementation. + + Parameters + ---------- + x_field : template + 1D mutable container (for example ``ti.ndarray`` or templated field) + containing the input ``x`` values. Updated in-place. + c_field : template + 1D coefficient container where ``c_field[k]`` is the coefficient + of :math:`T_k(x)`. + + Returns + ------- + template + The same container as ``x_field`` after in-place update. + + """ + c_len = c_field.shape[0] + + if c_len == 1: + for j in range(x_field.shape[0]): + x_field[j] = c_field[0] + elif c_len == 2: + for j in range(x_field.shape[0]): + x_field[j] = c_field[0] + c_field[1] * x_field[j] + else: + for j in range(x_field.shape[0]): + c0 = c_field[c_len - 2] + c1 = c_field[c_len - 1] + for i in range(3, c_len + 1): + tmp = c0 + c0 = c_field[c_len - i] - c1 + c1 = tmp + c1 * 2 * x_field[j] + x_field[j] = c0 + c1 * x_field[j] + return x_field + + +@func +def legval(x_field: template(), c_field: template()): + """ + Evaluate a Legendre series in-place on a 1D Taichi array. + + For coefficients ``c_field`` of length ``n + 1``, this computes + + .. math:: p(x) = \\sum_{k=0}^{n} c_k P_k(x) + + for every element in ``x_field`` and writes the result back into + ``x_field``. + + This function is intended to be called inside a Taichi kernel, + much like other math functions in Taichi's math module. + The evaluation uses Clenshaw recursion, and closely + follows the numpy legval implementation. + + Parameters + ---------- + x_field : template + 1D mutable container (for example ``ti.ndarray`` or templated field) + containing the input ``x`` values. Updated in-place. + c_field : template + 1D coefficient container where ``c_field[k]`` is the coefficient + of :math:`P_k(x)`. + + Returns + ------- + template + The same container as ``x_field`` after in-place update. + + """ + c_len = c_field.shape[0] + + if c_len == 1: + for j in range(x_field.shape[0]): + x_field[j] = c_field[0] + elif c_len == 2: + for j in range(x_field.shape[0]): + x_field[j] = c_field[0] + c_field[1] * x_field[j] + else: + for j in range(x_field.shape[0]): + nd = c_len + c0 = c_field[c_len - 2] + c1 = c_field[c_len - 1] + for i in range(3, c_len + 1): + tmp = c0 + nd = nd - 1 + c0 = c_field[c_len - i] - (c1 * (nd - 1)) / nd + c1 = tmp + (c1 * x_field[j] * (2 * nd - 1)) / nd + x_field[j] = c0 + c1 * x_field[j] + + return x_field + + +@func +def polyval(x_field: template(), c_field: template()): + """ + Evaluate a power series in-place on a 1D Taichi array. + + For coefficients ``c_field`` of length ``n + 1``, this computes + + .. math:: p(x) = \\sum_{k=0}^{n} c_k x^k + + for every element in ``x_field`` and writes the result back into + ``x_field``. + + This function is intended to be called inside a Taichi kernel, + much like other math functions in Taichi's math module. + The evaluation uses Horner's method, and closely + follows the numpy polyval implementation. + + Parameters + ---------- + x_field : template + 1D mutable container (for example ``ti.ndarray`` or templated field) + containing the input ``x`` values. Updated in-place. + c_field : template + 1D coefficient container where ``c_field[k]`` is the coefficient + of :math:`x^k`. + + Returns + ------- + template + The same container as ``x_field`` after in-place update. + + """ + + c_len = c_field.shape[0] + + for j in range(x_field.shape[0]): + c0 = c_field[c_len - 1] + for i in range(2, c_len + 1): + c0 = c_field[c_len - i] + c0 * x_field[j] + x_field[j] = c0 + + return x_field + + +@func +def fourval(x_field: template(), c_field: template()): + """ + Evaluate a Fourier series in-place on a 1D Taichi array. + + Coefficients are interpreted in the order: + ``[c0, cos(1*x), sin(1*x), cos(2*x), sin(2*x), ...]``. + For every element in ``x_field``, this computes the Fourier series value + and writes it back into ``x_field``. + + Parameters + ---------- + x_field : template + 1D mutable container (for example ``ti.ndarray`` or templated field) + containing the input ``x`` values. Updated in-place. + c_field : template + 1D coefficient container in interleaved cosine/sine order. + + Returns + ------- + template + The same container as ``x_field`` after in-place update. + """ + + c_len = c_field.shape[0] + + if c_len == 1: + for j in range(x_field.shape[0]): + x_field[j] = 0.5 * c_field[0] + else: + k_max = c_len // 2 + for j in range(x_field.shape[0]): + cx = cos(x_field[j]) + sx = sin(x_field[j]) + + bc1 = 0.0 + bc2 = 0.0 + bs1 = 0.0 + bs2 = 0.0 + + for k in range(k_max): + ia = 2 * (k_max - k) + + ak = c_field[ia - 1] if (ia - 1) < c_len else 0.0 + bk = c_field[ia] if ia < c_len else 0.0 + + bc0 = ak + 2.0 * cx * bc1 - bc2 + bs0 = bk + 2.0 * cx * bs1 - bs2 + + bc2 = bc1 + bc1 = bc0 + bs2 = bs1 + bs1 = bs0 + + x_field[j] = 0.5 * c_field[0] + (bc1 * cx - bc2) + bs1 * sx + + return x_field + + +@func +def lagmatrix(x_field: template(), use_orth_weight: template(), matrix_field: template()): + """ + Build a Laguerre pseudo-Vandermonde matrix in-place. + + For each input sample ``x_field[i]`` and degree ``j``, this writes + + .. math:: + \\text{matrix\\_field}[i, j] = w(x_i) L_j(x_i), \\quad 0 \\le j \\le deg, + + where ``deg = matrix_field.shape[1] - 1``, + ``w(x) = exp(-x/2)`` when ``use_orth_weight`` is ``True``, and + ``w(x) = 1`` otherwise. + + Parameters + ---------- + x_field : template + 1D input container of ``x`` values. + use_orth_weight : bool + If ``True``, apply the Laguerre orthogonality weight ``exp(-x/2)``. + matrix_field : template + 2D output container written in-place. Row count should match + ``x_field.shape[0]``. This implementation assumes at least 2 columns. + Reason: column 0 is the trivial base term (``w(x)`` or ``1``). + + Returns + ------- + template + The same container as ``matrix_field`` after in-place update. + """ + + for i in range(matrix_field.shape[0]): + if static(use_orth_weight): + matrix_field[i, 0] = exp(-0.5 * x_field[i]) + matrix_field[i, 1] = matrix_field[i, 0] * (1.0 - x_field[i]) + else: + matrix_field[i, 0] = 1.0 + matrix_field[i, 1] = 1.0 - x_field[i] + + for j in range(2, matrix_field.shape[1]): + matrix_field[i, j] = ( + (2 * j - 1 - x_field[i]) * matrix_field[i, j - 1] - (j - 1) * matrix_field[i, j - 2] + ) / j + + return matrix_field + + +@func +def hermmatrix(x_field: template(), use_orth_weight: template(), matrix_field: template()): + """ + Build a Hermite pseudo-Vandermonde matrix in-place. + + For each input sample ``x_field[i]`` and degree ``j``, this writes + + .. math:: + \\text{matrix\\_field}[i, j] = w(x_i) H_j(x_i), \\quad 0 \\le j \\le deg, + + where ``deg = matrix_field.shape[1] - 1``, + ``w(x) = exp(-x^2/2)`` when ``use_orth_weight`` is ``True``, and + ``w(x) = 1`` otherwise. + + Parameters + ---------- + x_field : template + 1D input container of ``x`` values. + use_orth_weight : bool + If ``True``, apply the Hermite root-weight ``exp(-x^2/2)``. + matrix_field : template + 2D output container written in-place. Row count should match + ``x_field.shape[0]``. This implementation assumes at least 2 columns. + + Returns + ------- + template + The same container as ``matrix_field`` after in-place update. + """ + for i in range(matrix_field.shape[0]): + matrix_field[i, 0] = exp(-0.5 * x_field[i] * x_field[i]) if static(use_orth_weight) else 1.0 + matrix_field[i, 1] = matrix_field[i, 0] * (2.0 * x_field[i]) + + for j in range(2, matrix_field.shape[1]): + matrix_field[i, j] = matrix_field[i, j - 1] * (2.0 * x_field[i]) - matrix_field[i, j - 2] * (2 * (j - 1)) + + return matrix_field + + +@func +def chebmatrix(x_field: template(), use_orth_weight: template(), matrix_field: template()): + """ + Build a Chebyshev pseudo-Vandermonde matrix in-place. + + For each input sample ``x_field[i]`` and degree ``j``, this writes + + .. math:: + \\text{matrix\\_field}[i, j] = w(x_i) T_j(x_i), \\quad 0 \\le j \\le deg, + + where ``deg = matrix_field.shape[1] - 1``, + ``w(x) = (1 - x^2)^(-1/4)`` when ``use_orth_weight`` is ``True``, and + ``w(x) = 1`` otherwise. + + Parameters + ---------- + x_field : template + 1D input container of ``x`` values. + use_orth_weight : template + If ``True``, apply the Chebyshev orthogonality weight + ``(1 - x^2)^(-1/4)``. + matrix_field : template + 2D output container written in-place. Row count should match + ``x_field.shape[0]``. This implementation assumes at least 2 columns. + + Returns + ------- + template + The same container as ``matrix_field`` after in-place update. + """ + for i in range(matrix_field.shape[0]): + if static(use_orth_weight): + matrix_field[i, 0] = 1.0 / sqrt(sqrt(1.0 - x_field[i] * x_field[i])) + matrix_field[i, 1] = matrix_field[i, 0] * x_field[i] + else: + matrix_field[i, 0] = 1.0 + matrix_field[i, 1] = x_field[i] + + for j in range(2, matrix_field.shape[1]): + matrix_field[i, j] = matrix_field[i, j - 1] * (2.0 * x_field[i]) - matrix_field[i, j - 2] + + return matrix_field + + +@func +def legmatrix(x_field: template(), matrix_field: template()): + """ + Build a Legendre pseudo-Vandermonde matrix in-place. + + For each input sample ``x_field[i]`` and degree ``j``, this writes + + .. math:: + \\text{matrix\\_field}[i, j] = P_j(x_i), \\quad 0 \\le j \\le deg, + + where ``deg = matrix_field.shape[1] - 1``. + + Parameters + ---------- + x_field : template + 1D input container of ``x`` values. + matrix_field : template + 2D output container written in-place. Row count should match + ``x_field.shape[0]``. This implementation assumes at least 2 columns. + + Returns + ------- + template + The same container as ``matrix_field`` after in-place update. + """ + for i in range(matrix_field.shape[0]): + matrix_field[i, 0] = 1.0 + matrix_field[i, 1] = x_field[i] + + for j in range(2, matrix_field.shape[1]): + matrix_field[i, j] = ( + matrix_field[i, j - 1] * x_field[i] * (2 * j - 1) - matrix_field[i, j - 2] * (j - 1) + ) / j + + return matrix_field + + +@func +def polymatrix(x_field: template(), matrix_field: template()): + """ + Build a power-basis Vandermonde matrix in-place. + + For each input sample ``x_field[i]`` and degree ``j``, this writes + + .. math:: + \\text{matrix\\_field}[i, j] = x_i^j, \\quad 0 \\le j \\le deg, + + where ``deg = matrix_field.shape[1] - 1``. + + Parameters + ---------- + x_field : template + 1D input container of ``x`` values. + matrix_field : template + 2D output container written in-place. Row count should match + ``x_field.shape[0]``. This implementation assumes at least 2 columns. + + Returns + ------- + template + The same container as ``matrix_field`` after in-place update. + """ + for i in range(matrix_field.shape[0]): + matrix_field[i, 0] = 1.0 + matrix_field[i, 1] = x_field[i] + + for j in range(2, matrix_field.shape[1]): + matrix_field[i, j] = matrix_field[i, j - 1] * x_field[i] + + return matrix_field + + +@func +def fourmatrix(x_field: template(), matrix_field: template()): + """ + Build a trigonometric design matrix in-place (Fourier basis). + + Column order matches ``fourval`` coefficients: + ``[c0, cos(1*x), sin(1*x), cos(2*x), sin(2*x), ...]``, + with the constant term represented as ``0.5 * c0``. + + Parameters + ---------- + x_field : template + 1D input container of ``x`` values. + matrix_field : template + 2D output container written in-place. Row count should match + ``x_field.shape[0]``. + + Returns + ------- + template + The same container as ``matrix_field`` after in-place update. + """ + for i in range(matrix_field.shape[0]): + matrix_field[i, 0] = 0.5 + + if matrix_field.shape[1] > 1: + matrix_field[i, 1] = cos(x_field[i]) + + if matrix_field.shape[1] > 2: + matrix_field[i, 2] = sin(x_field[i]) + + for k in range(2, matrix_field.shape[1] // 2 + 1): + matrix_field[i, 2 * k - 1] = ( + matrix_field[i, 2 * k - 3] * matrix_field[i, 1] - matrix_field[i, 2 * k - 2] * matrix_field[i, 2] + ) + if 2 * k < matrix_field.shape[1]: + matrix_field[i, 2 * k] = ( + matrix_field[i, 2 * k - 2] * matrix_field[i, 1] + matrix_field[i, 2 * k - 3] * matrix_field[i, 2] + ) + + return matrix_field + + +__all__ = [ + "lagval", + "hermval", + "chebval", + "legval", + "polyval", + "fourval", + "lagmatrix", + "hermmatrix", + "chebmatrix", + "legmatrix", + "polymatrix", + "fourmatrix", +] diff --git a/tests/python/basis_functions_ref.py b/tests/python/basis_functions_ref.py new file mode 100644 index 0000000000000..8ab950ed17bfb --- /dev/null +++ b/tests/python/basis_functions_ref.py @@ -0,0 +1,176 @@ +import numpy as np + +# Imports/Code to evaluate basis functions *series* for a defined set of coefficients +from numpy.polynomial.laguerre import lagval +from numpy.polynomial.hermite import hermval +from numpy.polynomial.chebyshev import chebval +from numpy.polynomial.legendre import legval + + +def fourval(x, coeffs): + """ + Evaluate the Fourier series of a given set of coefficients. + + Parameters: + x: array-like - Input values (stock prices or other variables) + coeffs: array-like - Coefficients of the Fourier series + + Returns: + array - Fourier basis functions evaluated at x + """ + coeffs = np.asarray(coeffs) + F = np.zeros_like(x, dtype=np.result_type(x, coeffs)) + + if coeffs.shape[0] > 0: # We have more than 0 coeffs + F = F + 0.5 * coeffs[0] # First term is constant + for i in range(1, coeffs.shape[0]): + k = (i + 1) // 2 + if i % 2 == 1: + F = F + coeffs[i] * np.cos(k * x) + else: + F = F + coeffs[i] * np.sin(k * x) + return F + + +# Code to evaluate (weighted) basis function *matrices* +def lagmatrix(x, x_length, num_basis_functions, use_orth_weight): + """ + Compute a set of Laguerre polynomials with custom weight. + + Parameters: + x: array-like - Input values (stock prices or other variables) + x_length: int - Length of the input array + num_basis_functions: int - Degree of the Laguerre polynomial to compute + use_orth_weight: bool - Whether to use the orthogonality weight + + Returns: + array - Laguerre basis functions evaluated at x + """ + weight = 1 if not use_orth_weight else np.exp(-x / 2) + L = np.zeros((x_length, num_basis_functions)) # Create an array for the basis functions + for i in range(num_basis_functions): + L[:, i] = weight * lagval( + x, [0] * i + [1] + ) # Evaluate Laguerre polynomial i. If I remove np.exp(-x/2) * then I get the same Laguerre polynomials as QuantLib. + return L + + +def hermmatrix(x, x_length, num_basis_functions, use_orth_weight): + """ + Compute a set of Hermite polynomials with custom weight. + + Parameters: + x: array-like - Input values (stock prices or other variables) + x_length: int - Length of the input array + num_basis_functions: int - Degree of the Laguerre polynomial to compute + use_orth_weight: bool - Whether to use the orthogonality weight + + Returns: + array - Laguerre basis functions evaluated at x + """ + weight = 1 if not use_orth_weight else np.exp(-(x**2) / 2) + L = np.zeros((x_length, num_basis_functions)) # Create an array for the basis functions + for i in range(num_basis_functions): + L[:, i] = weight * hermval( + x, [0] * i + [1] + ) # Evaluate Laguerre polynomial i. If I remove np.exp(-x/2) * then I get the same Laguerre polynomials as QuantLib. + return L + + +def chebmatrix(x, x_length, num_basis_functions, use_orth_weight): + """ + Compute a set of Chebyshev polynomials with custom weight. + + Parameters: + x: array-like - Input values (stock prices or other variables) + x_length: int - Length of the input array + num_basis_functions: int - Degree of the Laguerre polynomial to compute + use_orth_weight: bool - Whether to use the orthogonality weight + + Returns: + array - Laguerre basis functions evaluated at x + """ + weight = 1 if not use_orth_weight else (1 - x**2) ** (-1 / 4) + L = np.zeros((x_length, num_basis_functions)) # Create an array for the basis functions + for i in range(num_basis_functions): + L[:, i] = weight * chebval( + x, [0] * i + [1] + ) # Evaluate Laguerre polynomial i. If I remove np.exp(-x/2) * then I get the same Laguerre polynomials as QuantLib. + return L + + +def legmatrix(x, x_length, num_basis_functions, _): + """ + Compute a set of Legendre polynomials. + + Parameters: + x: array-like - Input values (stock prices or other variables) + num_basis_functions: int - Degree of the Laguerre polynomial to compute + + Returns: + array - Laguerre basis functions evaluated at x + """ + L = np.zeros((x_length, num_basis_functions)) # Create an array for the basis functions + for i in range(num_basis_functions): + L[:, i] = legval(x, [0] * i + [1]) + return L + + +def polymatrix(x, x_length, num_basis_functions, _): + """ + Compute a set of monomials. + + Parameters: + x: array-like - Input values (stock prices or other variables) + num_basis_functions: int - Degree of the Laguerre polynomial to compute + + Returns: + array - Laguerre basis functions evaluated at x + """ + L = np.zeros((x_length, num_basis_functions)) # Create an array for the basis functions + for i in range(num_basis_functions): + L[:, i] = x**i + return L + + +def fourmatrix(x, x_length, num_basis_functions, _): + """ + Construct the first `num_basis_functions` Fourier basis functions sampled at x: + + [constant, + cos(1⋅x), sin(1⋅x), + cos(2⋅x), sin(2⋅x), + …] + + Parameters + ---------- + x : array-like, shape (N,) + Array of input locations at which to evaluate basis functions (does not need to be in a particular range). + x_length : int + Number of points in x (for shape of output array). + num_basis_functions : int + Number of basis functions to include (length of the returned basis). + + Returns + ------- + F : ndarray, shape (N, num_basis_functions) + Matrix of basis functions evaluated at `x`. Columns are ordered: + - F[:, 0] = constant term (1/2), + - F[:, 1] = cos(1⋅x), + - F[:, 2] = sin(1⋅x), + - F[:, 3] = cos(2⋅x), + - F[:, 4] = sin(2⋅x), + ... + This is the design matrix for projecting onto the truncated Fourier basis. + """ + F = np.zeros((x_length, num_basis_functions)) + + # constant term (first basis function) + if num_basis_functions > 0: + F[:, 0] = 1 / 2 + + # Subsequent basis: cos/sin pairs + for i in range(1, num_basis_functions): + F[:, i] = fourval(x, [0] * i + [1]) + + return F diff --git a/tests/python/test_basis_matrices_eval.py b/tests/python/test_basis_matrices_eval.py new file mode 100644 index 0000000000000..66839ca13b468 --- /dev/null +++ b/tests/python/test_basis_matrices_eval.py @@ -0,0 +1,117 @@ +import pytest +from tests import test_utils +from types import SimpleNamespace + +import numpy as np +from .basis_functions_ref import ( + lagmatrix as np_lagmatrix, + hermmatrix as np_hermmatrix, + chebmatrix as np_chebmatrix, + legmatrix as np_legmatrix, + polymatrix as np_polymatrix, + fourmatrix as np_fourmatrix, +) +import taichi as ti +from taichi.math.polynomial import ( + lagmatrix as ti_lagmatrix, + hermmatrix as ti_hermmatrix, + chebmatrix as ti_chebmatrix, + legmatrix as ti_legmatrix, + polymatrix as ti_polymatrix, + fourmatrix as ti_fourmatrix, +) + +np_matrices_eval = SimpleNamespace( + laguerre=np_lagmatrix, + hermite=np_hermmatrix, + legendre=np_legmatrix, + chebyshev=np_chebmatrix, + monomial=np_polymatrix, + fourier=np_fourmatrix, +) + +ti_matrices_eval = SimpleNamespace( + laguerre=ti_lagmatrix, + hermite=ti_hermmatrix, + legendre=ti_legmatrix, + chebyshev=ti_chebmatrix, + monomial=ti_polymatrix, + fourier=ti_fourmatrix, +) + + +def _test_basis_matrices(dt, family, num_basis_functions, use_orth_weight): + + # Numpy logic to get expected values + np_dt = np.float32 if dt == ti.f32 else np.float64 + + x = np.linspace(0.0, 0.999, 37, dtype=np_dt) + np_basis_func = getattr(np_matrices_eval, family) + expected = np_basis_func(x, x.shape[0], num_basis_functions, use_orth_weight).astype(np_dt) + + # Taichi logic to get actual values + ti_x = ti.field(dt, shape=x.shape) + ti_x.from_numpy(x) + ti_matrix = ti.field(dt, shape=(x.shape[0], num_basis_functions)) + ti_basis_func = getattr(ti_matrices_eval, family) + + if use_orth_weight is not None: + + @ti.kernel + def basis_test(ti_x: ti.template(), use_orth_weight: ti.template(), ti_matrix: ti.template()): + ti_basis_func(ti_x, use_orth_weight, ti_matrix) + + basis_test(ti_x, use_orth_weight, ti_matrix) + else: + + @ti.kernel + def basis_test(ti_x: ti.template(), ti_matrix: ti.template()): + ti_basis_func(ti_x, ti_matrix) + + basis_test(ti_x, ti_matrix) + + actual = ti_matrix.to_numpy() + + # Compare expected and actual values + tol = 1e-5 if dt == ti.f32 else 1e-12 + np.testing.assert_allclose(actual, expected, rtol=tol, atol=tol) + + +@pytest.mark.parametrize( + "family,use_orth_weight", + [ + pytest.param("laguerre", False), + pytest.param("laguerre", True), + pytest.param("hermite", False), + pytest.param("hermite", True), + pytest.param("chebyshev", False), + pytest.param("chebyshev", True), + pytest.param("legendre", None), + pytest.param("monomial", None), + pytest.param("fourier", None), + ], +) +@pytest.mark.parametrize("num_basis_functions", [2, 4, 7, 9]) +@test_utils.test(default_fp=ti.f32, fast_math=False) +def test_basis_matrices_f32(family, use_orth_weight, num_basis_functions): + _test_basis_matrices(ti.f32, family, num_basis_functions, use_orth_weight) + + +@pytest.mark.parametrize( + "family,use_orth_weight", + [ + pytest.param("laguerre", False), + pytest.param("laguerre", True), + pytest.param("hermite", False), + pytest.param("hermite", True), + pytest.param("chebyshev", False), + pytest.param("chebyshev", True), + pytest.param("legendre", None), + pytest.param("monomial", None), + pytest.param("fourier", None), + ], +) +@pytest.mark.parametrize("num_basis_functions", [2, 4, 7, 9]) +@test_utils.test(require=ti.extension.data64, default_fp=ti.f64, fast_math=False) +def test_basis_matrices_f64(family, use_orth_weight, num_basis_functions): + _test_basis_matrices(ti.f64, family, num_basis_functions, use_orth_weight) diff --git a/tests/python/test_basis_series_eval.py b/tests/python/test_basis_series_eval.py new file mode 100644 index 0000000000000..57ad19fa92d31 --- /dev/null +++ b/tests/python/test_basis_series_eval.py @@ -0,0 +1,83 @@ +import pytest +from tests import test_utils +from types import SimpleNamespace + +import numpy as np +from numpy.polynomial.chebyshev import chebval as np_chebval +from numpy.polynomial.hermite import hermval as np_hermval +from numpy.polynomial.laguerre import lagval as np_lagval +from numpy.polynomial.legendre import legval as np_legval +from numpy.polynomial.polynomial import polyval as np_polyval + +import taichi as ti +from .basis_functions_ref import fourval as np_fourval +from taichi.math.polynomial import ( + lagval as ti_lagval, + hermval as ti_hermval, + chebval as ti_chebval, + legval as ti_legval, + polyval as ti_polyval, + fourval as ti_fourval, +) + + +np_series_eval = SimpleNamespace( + laguerre=np_lagval, + hermite=np_hermval, + legendre=np_legval, + chebyshev=np_chebval, + monomial=np_polyval, + fourier=np_fourval, +) + +ti_series_eval = SimpleNamespace( + laguerre=ti_lagval, + hermite=ti_hermval, + legendre=ti_legval, + chebyshev=ti_chebval, + monomial=ti_polyval, + fourier=ti_fourval, +) + + +def _test_basis_series_eval(dt, family, degree): + + # Numpy logic to get expected values + np_dt = np.float32 if dt == ti.f32 else np.float64 + + x = np.linspace(-2.0, 2.0, 37, dtype=np_dt) + coeffs = np.random.default_rng(1000 + degree).normal(size=degree + 1).astype(np_dt) + np_basis_func = getattr(np_series_eval, family) + expected = np_basis_func(x, coeffs).astype(np_dt) + + # Taichi logic to get actual values + ti_x = ti.field(dt, shape=x.shape) + ti_x.from_numpy(x) + ti_coeffs = ti.field(dt, shape=coeffs.shape) + ti_coeffs.from_numpy(coeffs) + ti_basis_func = getattr(ti_series_eval, family) + + @ti.kernel + def basis_test(ti_x: ti.template(), ti_coeffs: ti.template()): + ti_basis_func(ti_x, ti_coeffs) + + basis_test(ti_x, ti_coeffs) + actual = ti_x.to_numpy() + + # Compare expected and actual values + tol = 1e-5 if dt == ti.f32 else 1e-12 + np.testing.assert_allclose(np.asarray(actual, dtype=np_dt), expected, rtol=tol, atol=tol) + + +@pytest.mark.parametrize("family", ["laguerre", "hermite", "chebyshev", "legendre", "monomial", "fourier"]) +@pytest.mark.parametrize("degree", [0, 1, 2, 4, 7, 10]) +@test_utils.test(default_fp=ti.f32, fast_math=False) +def test_basis_series_eval_f32(family, degree): + _test_basis_series_eval(ti.f32, family, degree) + + +@pytest.mark.parametrize("family", ["laguerre", "hermite", "chebyshev", "legendre", "monomial", "fourier"]) +@pytest.mark.parametrize("degree", [0, 1, 2, 4, 7, 10]) +@test_utils.test(require=ti.extension.data64, default_fp=ti.f64, fast_math=False) +def test_basis_series_eval_f64(family, degree): + _test_basis_series_eval(ti.f64, family, degree)