diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index 3abd1dc..6fdc8ee 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -1,11 +1,12 @@ +import math from dataclasses import dataclass -from typing import Union +from functools import cache +from typing import Sequence, Union import numba import numpy as np from scipy.interpolate import interp1d, lagrange from scipy.special import jn_zeros -from functools import cache from scgenerator.cache import np_cache @@ -576,3 +577,84 @@ def cumulative_boole(x: np.ndarray) -> np.ndarray: for i in range(4, len(x)): out[i] = out[i - 4] + np.sum(c * x[i - 4 : i + 1]) return out + + +def stencil_coefficients(stencil_points: Sequence, order: int) -> np.ndarray: + """ + Reference + --------- + https://en.wikipedia.org/wiki/Finite_difference_coefficient#Arbitrary_stencil_points + """ + mat = np.power.outer(stencil_points, np.arange(len(stencil_points))).T + d = np.zeros(len(stencil_points)) + d[order] = math.factorial(order) + return np.linalg.solve(mat, d) + + +@cache +def central_stencil_coefficients(n: int, order: int) -> np.ndarray: + """ + returns the coefficients of a centered finite difference scheme + + Parameters + ---------- + n : int + number of points + order : int + order of differentiation + """ + return stencil_coefficients(np.arange(n * 2 + 1) - n, order) + + +@cache +def stencil_coefficients_set(n: int, order: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + coefficients for a forward, centered and backward finite difference differentation scheme of + order `order` that extends `n` points away from the evaluation point (`x_0`) + """ + central = central_stencil_coefficients(n, order) + left = stencil_coefficients(np.arange(n + 1), order) + right = stencil_coefficients(-np.arange(n + 1)[::-1], order) + return left, central, right + + +def differentiate_arr( + values: np.ndarray, diff_order: int, extent: int, correct_edges=True +) -> np.ndarray: + """ + takes a derivative of order `diff_order` using equally spaced values + **NOTE** : this function doesn't actually know about your grid spacing `h`, so you need to + divide the result by `h**diff_order` to get a correct result. + + Parameters + --------- + values : ndarray + equally spaced values + diff_order : int + order of differentiation + extent : int + how many points away from the center the scheme uses. This determines accuracy. + example: extent=6 means that 13 (6 on either side + center) points are used to evaluate the + derivative at each point. + correct_edges : bool, optional + avoid artifacts by using forward/backward schemes on the edges, by default True + + Reference + --------- + https://en.wikipedia.org/wiki/Finite_difference_coefficient + + """ + n_points = (diff_order + extent) // 2 + + if not correct_edges: + central_coefs = central_stencil_coefficients(n_points, diff_order) + return np.convolve(values, central_coefs[::-1], mode="same") + + left_coefs, central_coefs, right_coefs = stencil_coefficients_set(n_points, diff_order) + return np.concatenate( + ( + np.convolve(values[: 2 * n_points], left_coefs[::-1], mode="valid"), + np.convolve(values, central_coefs[::-1], mode="valid"), + np.convolve(values[-2 * n_points :], right_coefs[::-1], mode="valid"), + ) + ) diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index 33f76ea..61cbda7 100644 --- a/src/scgenerator/physics/fiber.py +++ b/src/scgenerator/physics/fiber.py @@ -973,7 +973,7 @@ def effective_core_radius_vincetti( def li_vincetti(f_2: T, f0_2: float) -> T: k = f0_2 - f_2 - return k / (k**2 + 9e-6 * f_2) + return k / (k**2 + 9e-4 * f_2) def cutoff_frequency_he_vincetti(mu: int, nu: int, t_ratio: float, n_clad_0: float) -> np.ndarray: diff --git a/testing/test_math.py b/testing/test_math.py index a1f8929..bd3fbd2 100644 --- a/testing/test_math.py +++ b/testing/test_math.py @@ -1,13 +1,15 @@ +from math import factorial + import numpy as np import pytest + import scgenerator.math as m -from math import factorial def test__power_fact_array(): x = np.random.rand(5) for i in range(5): - assert m._power_fact_array(x, i) == pytest.approx(x ** i / factorial(i)) + assert m._power_fact_array(x, i) == pytest.approx(x**i / factorial(i)) def test__power_fact_single(): @@ -98,3 +100,35 @@ def test_update_frequency_domain(): def test_wspace(): pass + + +def test_differentiate(): + x = np.linspace(-10, 10, 256) + y = np.exp(-((x / 3) ** 2)) * (1 + 0.2 * np.sin(x * 5)) + + y[100] = 1e4 + # true = np.exp(-(x/3)**2) * (x*(-0.4/9 * np.sin(5*x) - 2/9) + np.cos(5*x)) + true = np.exp(-((x / 3) ** 2)) * ( + x**2 * (0.00987654321 * np.sin(5 * x) + 0.0493827) + - 5.044444 * np.sin(5 * x) + - 0.44444 * x * np.cos(5 * x) + - 0.2222222 + ) + + import matplotlib.pyplot as plt + + h = x[1] - x[0] + + grad = np.gradient(np.gradient(y)) / h**2 + fine = m.differentiate_arr(y, 2, 6) / h**2 + + plt.plot(x, y) + plt.plot(x, grad, label="gradient") + plt.plot(x, fine, label="fine") + plt.plot(x, true, label="ture", ls=":") + plt.legend() + plt.show() + + +if __name__ == "__main__": + test_differentiate()