Refined Vincetti model and added finite difference

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"),
+        )
+    )

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:

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()