| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // Copyright (C) 2025 Tiago V. L. Amorim (@tiagovla) | ||
| 2 | // | ||
| 3 | // This file is part of oiseau (https://github.com/tiagovla/oiseau) | ||
| 4 | // | ||
| 5 | // SPDX-License-Identifier: GPL-3.0-or-later | ||
| 6 | |||
| 7 | #pragma once | ||
| 8 | |||
| 9 | #include <algorithm> | ||
| 10 | #include <cmath> | ||
| 11 | #include <concepts> | ||
| 12 | #include <limits> | ||
| 13 | #include <ranges> | ||
| 14 | #include <xtensor/containers/xarray.hpp> | ||
| 15 | |||
| 16 | namespace oiseau::utils { | ||
| 17 | |||
| 18 | /// Concept that checks for contiguous floating-point containers. | ||
| 19 | template <class Container> | ||
| 20 | concept FloatingArrayLike = std::ranges::contiguous_range<Container> && | ||
| 21 | std::is_floating_point_v<std::ranges::range_value_t<Container>>; | ||
| 22 | |||
| 23 | /// Computes the Jacobi polynomial P_n^{(alpha,beta)}(x) at a single point. | ||
| 24 | template <std::floating_point Real> | ||
| 25 | 2369105 | Real jacobi(unsigned n, Real alpha, Real beta, Real x) { | |
| 26 |
2/2✓ Branch 0 taken 525117 times.
✓ Branch 1 taken 1843988 times.
|
2369105 | if (n == 0) { |
| 27 | 525117 | return Real(1); | |
| 28 | } | ||
| 29 | |||
| 30 | 1843988 | Real y0 = 1; | |
| 31 | 1843988 | Real y1 = (alpha + 1) + (alpha + beta + 2) * (x - 1) / Real(2); | |
| 32 | |||
| 33 | 1843988 | Real yk = y1; | |
| 34 | 1843988 | unsigned k = 2; | |
| 35 | 1843988 | Real k_max = n * (1 + std::numeric_limits<Real>::epsilon()); | |
| 36 | |||
| 37 |
2/2✓ Branch 0 taken 2945563 times.
✓ Branch 1 taken 1843988 times.
|
4789551 | while (k < k_max) { |
| 38 | 2945563 | Real denom = 2 * k * (k + alpha + beta) * (2 * k + alpha + beta - 2); | |
| 39 | 2945563 | Real gamma1 = | |
| 40 | 2945563 | (2 * k + alpha + beta - 1) * | |
| 41 | 2945563 | ((2 * k + alpha + beta) * (2 * k + alpha + beta - 2) * x + alpha * alpha - beta * beta); | |
| 42 | 2945563 | Real gamma0 = -2 * (k + alpha - 1) * (k + beta - 1) * (2 * k + alpha + beta); | |
| 43 | 2945563 | yk = (gamma1 * y1 + gamma0 * y0) / denom; | |
| 44 | 2945563 | y0 = y1; | |
| 45 | 2945563 | y1 = yk; | |
| 46 | 2945563 | ++k; | |
| 47 | } | ||
| 48 | 1843988 | return yk; | |
| 49 | } | ||
| 50 | |||
| 51 | /// Computes the normalized Jacobi polynomial at each entry of a container. | ||
| 52 | template <std::floating_point Real, FloatingArrayLike Container> | ||
| 53 | 22812 | Container jacobi_p(unsigned n, Real alpha, Real beta, const Container& input) { | |
| 54 |
1/2✓ Branch 1 taken 22812 times.
✗ Branch 2 not taken.
|
22812 | Container v = input; |
| 55 | 22812 | Real norm = std::pow(2, alpha + beta + 1) / (2 * n + alpha + beta + 1); | |
| 56 | 22812 | norm *= std::tgamma(n + alpha + 1) * std::tgamma(n + beta + 1); | |
| 57 | 22812 | norm /= (std::tgamma(n + 1) * std::tgamma(n + alpha + beta + 1)); | |
| 58 | 4761022 | auto partial = [n, alpha, beta, norm](auto& k) { | |
| 59 | 2369105 | k = jacobi(n, alpha, beta, k) / std::sqrt(norm); | |
| 60 | }; | ||
| 61 |
1/2✓ Branch 3 taken 22812 times.
✗ Branch 4 not taken.
|
22812 | std::for_each(v.begin(), v.end(), partial); |
| 62 | 22812 | return v; | |
| 63 | ✗ | } | |
| 64 | |||
| 65 | /// Computes the gradient of the normalized Jacobi polynomial. | ||
| 66 | template <std::floating_point Real, FloatingArrayLike Container> | ||
| 67 | 7842 | Container grad_jacobi_p(int n, Real alpha, Real beta, const Container& input) { | |
| 68 |
1/2✓ Branch 1 taken 7842 times.
✗ Branch 2 not taken.
|
7842 | Container v = input; |
| 69 |
2/2✓ Branch 0 taken 1998 times.
✓ Branch 1 taken 5844 times.
|
7842 | if (n == 0) { |
| 70 |
1/2✓ Branch 3 taken 1998 times.
✗ Branch 4 not taken.
|
1998 | std::fill(v.begin(), v.end(), Real(0)); |
| 71 | 1998 | return v; | |
| 72 | } | ||
| 73 | 5844 | Real norm = std::sqrt(n * (n + alpha + beta + 1)); | |
| 74 |
1/2✓ Branch 1 taken 5844 times.
✗ Branch 2 not taken.
|
5844 | Container output = jacobi_p(n - 1, alpha + 1, beta + 1, v); |
| 75 | 5844 | std::for_each(output.begin(), output.end(), [norm](auto& k) { k *= norm; }); | |
| 76 | 5844 | return output; | |
| 77 | 7842 | } | |
| 78 | |||
| 79 | } // namespace oiseau::utils | ||
| 80 |