Commit 89bedcae authored by Michele Maris's avatar Michele Maris
Browse files

u

parent 2ae6a7f8

notebooks/.ipynb_checkpoints/numba_cubic_spline-checkpoint.ipynb

0 → 100644
+208 −0

File added.

Preview size limit exceeded, changes collapsed.

notebooks/numba_cubic_spline.ipynb

0 → 100644
+208 −0

File added.

Preview size limit exceeded, changes collapsed.

src/yapsut/numba_cubic_spline.py

0 → 100644
+307 −0
Original line number Diff line number Diff line 
__description__=""" 

   cubic spline with numba

   as produced by chat gpt

"""



import numpy as np

from numba import njit



@njit

def cubic_spline_natural_coeffs(x, y):

    n = len(x) - 1

    h = np.diff(x)

    alpha = np.zeros(n)

    for i in range(1, n):

        alpha[i] = (3/h[i]) * (y[i+1] - y[i]) - (3/h[i-1]) * (y[i] - y[i-1])

    

    l = np.ones(n+1)

    mu = np.zeros(n+1)

    z = np.zeros(n+1)

    

    for i in range(1, n):

        l[i] = 2*(x[i+1] - x[i-1]) - h[i-1]*mu[i-1]

        mu[i] = h[i]/l[i]

        z[i] = (alpha[i] - h[i-1]*z[i-1])/l[i]

    

    b = np.zeros(n)

    c = np.zeros(n+1)

    d = np.zeros(n)

    

    for j in range(n-1, -1, -1):

        c[j] = z[j] - mu[j]*c[j+1]

        b[j] = (y[j+1] - y[j])/h[j] - h[j]*(c[j+1] + 2*c[j])/3

        d[j] = (c[j+1] - c[j])/(3*h[j])

    

    return b, c[:-1], d  # a = y[:-1]



@njit

def evaluate_spline(x, y, b, c, d, x_eval):

    n = len(x) - 1

    y_eval = np.zeros_like(x_eval)

    for j in range(len(x_eval)):

        xi = x_eval[j]

        for i in range(n):

            if x[i] <= xi <= x[i+1]:

                dx = xi - x[i]

                y_eval[j] = y[i] + b[i]*dx + c[i]*dx**2 + d[i]*dx**3

                break

    return y_eval



@njit

def evaluate_spline_derivative(x, b, c, d, x_eval):

    n = len(x) - 1

    y_deriv = np.zeros_like(x_eval)

    for j in range(len(x_eval)):

        xi = x_eval[j]

        for i in range(n):

            if x[i] <= xi <= x[i+1]:

                dx = xi - x[i]

                y_deriv[j] = b[i] + 2 * c[i] * dx + 3 * d[i] * dx**2

                break

    return y_deriv

 

 

""" =====================   """



@njit

def _compute_coeffs(x, y):

    n = len(x) - 1

    h = np.diff(x)

    alpha = np.zeros(n)

    for i in range(1, n):

        alpha[i] = (3/h[i]) * (y[i+1] - y[i]) - (3/h[i-1]) * (y[i] - y[i-1])

    

    l = np.ones(n+1)

    mu = np.zeros(n+1)

    z = np.zeros(n+1)

    

    for i in range(1, n):

        l[i] = 2*(x[i+1] - x[i-1]) - h[i-1]*mu[i-1]

        mu[i] = h[i]/l[i]

        z[i] = (alpha[i] - h[i-1]*z[i-1])/l[i]

    

    b = np.zeros(n)

    c = np.zeros(n+1)

    d = np.zeros(n)

    

    for j in range(n-1, -1, -1):

        c[j] = z[j] - mu[j]*c[j+1]

        b[j] = (y[j+1] - y[j])/h[j] - h[j]*(c[j+1] + 2*c[j])/3

        d[j] = (c[j+1] - c[j])/(3*h[j])

    

    return b, c[:-1], d  # a = y[:-1]



@njit

def _evaluate(x, y, b, c, d, x_eval):

    n = len(x) - 1

    y_eval = np.zeros_like(x_eval)

    for j in range(len(x_eval)):

        xi = x_eval[j]

        for i in range(n):

            if x[i] <= xi <= x[i+1]:

                dx = xi - x[i]

                y_eval[j] = y[i] + b[i]*dx + c[i]*dx**2 + d[i]*dx**3

                break

    return y_eval



@njit

def _evaluate_derivative(x, b, c, d, x_eval):

    n = len(x) - 1

    y_deriv = np.zeros_like(x_eval)

    for j in range(len(x_eval)):

        xi = x_eval[j]

        for i in range(n):

            if x[i] <= xi <= x[i+1]:

                dx = xi - x[i]

                y_deriv[j] = b[i] + 2 * c[i] * dx + 3 * d[i] * dx**2

                break

    return y_deriv



# Classe wrapper

class NumbaNaturalCubicSpline:

    """ natural cubic spline with numba """

    def __init__(self, x, y):

        self.x = np.asarray(x)

        self.y = np.asarray(y)

        self.b, self.c, self.d = _compute_coeffs(self.x, self.y)

    

    def __call__(self, x_eval):

        x_eval = np.asarray(x_eval)

        return _evaluate(self.x, self.y, self.b, self.c, self.d, x_eval)

    

    def derivative(self, x_eval):

        x_eval = np.asarray(x_eval)

        return _evaluate_derivative(self.x, self.b, self.c, self.d, x_eval)



""" =====================   """



@njit

def _xt_compute_clamped_coeffs(x, y, fp0, fpn):

    n = len(x) - 1

    h = np.diff(x)

    alpha = np.zeros(n+1)



    alpha[0] = 3 * (y[1] - y[0]) / h[0] - 3 * fp0

    alpha[n] = 3 * fpn - 3 * (y[n] - y[n-1]) / h[n-1]



    for i in range(1, n):

        alpha[i] = (3/h[i]) * (y[i+1] - y[i]) - (3/h[i-1]) * (y[i] - y[i-1])



    l = np.ones(n+1)

    mu = np.zeros(n+1)

    z = np.zeros(n+1)



    l[0] = 2 * h[0]

    mu[0] = 0.5

    z[0] = alpha[0] / l[0]



    for i in range(1, n):

        l[i] = 2*(x[i+1] - x[i-1]) - h[i-1]*mu[i-1]

        mu[i] = h[i]/l[i]

        z[i] = (alpha[i] - h[i-1]*z[i-1])/l[i]



    l[n] = h[n-1]*(2 - mu[n-1])

    z[n] = (alpha[n] - h[n-1]*z[n-1])/l[n]



    c = np.zeros(n+1)

    b = np.zeros(n)

    d = np.zeros(n)



    c[n] = z[n]

    for j in range(n-1, -1, -1):

        c[j] = z[j] - mu[j]*c[j+1]

        b[j] = (y[j+1] - y[j])/h[j] - h[j]*(c[j+1] + 2*c[j])/3

        d[j] = (c[j+1] - c[j])/(3*h[j])



    return b, c[:-1], d



@njit

def _xt_evaluate(x, y, b, c, d, x_eval):

    n = len(x) - 1

    y_eval = np.zeros_like(x_eval)

    for j in range(len(x_eval)):

        xi = x_eval[j]

        for i in range(n):

            if x[i] <= xi <= x[i+1]:

                dx = xi - x[i]

                y_eval[j] = y[i] + b[i]*dx + c[i]*dx**2 + d[i]*dx**3

                break

    return y_eval



@njit

def _xt_evaluate_derivative(x, b, c, d, x_eval):

    n = len(x) - 1

    y_deriv = np.zeros_like(x_eval)

    for j in range(len(x_eval)):

        xi = x_eval[j]

        for i in range(n):

            if x[i] <= xi <= x[i+1]:

                dx = xi - x[i]

                y_deriv[j] = b[i] + 2*c[i]*dx + 3*d[i]*dx**2

                break

    return y_deriv



@njit

def _xt_evaluate_second_derivative(x, c, d, x_eval):

    n = len(x) - 1

    y_sec = np.zeros_like(x_eval)

    for j in range(len(x_eval)):

        xi = x_eval[j]

        for i in range(n):

            if x[i] <= xi <= x[i+1]:

                dx = xi - x[i]

                y_sec[j] = 2*c[i] + 6*d[i]*dx

                break

    return y_sec



class NumbaCubicSpline:

    def __init__(self, x, y, bc_type='natural', fp0=0.0, fpn=0.0):

        self.x = np.asarray(x)

        self.y = np.asarray(y)



        if bc_type == 'natural':

            self.b, self.c, self.d = _xt_compute_clamped_coeffs(self.x, self.y, 0.0, 0.0)

        elif bc_type == 'clamped':

            self.b, self.c, self.d = _xt_compute_clamped_coeffs(self.x, self.y, fp0, fpn)

        else:

            raise ValueError("bc_type must be 'natural' or 'clamped'")



    def __call__(self, x_eval):

        x_eval = np.asarray(x_eval)

        return _xt_evaluate(self.x, self.y, self.b, self.c, self.d, x_eval)



    def derivative(self, x_eval):

        x_eval = np.asarray(x_eval)

        return _xt_evaluate_derivative(self.x, self.b, self.c, self.d, x_eval)



    def second_derivative(self, x_eval):

        x_eval = np.asarray(x_eval)

        return _xt_evaluate_second_derivative(self.x, self.c, self.d, x_eval)



    def coefficients(self):

        return self.b.copy(), self.c.copy(), self.d.copy()



    def as_derivative(self):

        x_new = self.x.copy()

        y_new = self.derivative(self.x)

        return NumbaCubicSpline(x_new, y_new, bc_type='natural')



    def antiderivative(self):

        n = len(self.x) - 1

        x_new = self.x.copy()

        y_new = np.zeros_like(self.x)



        for i in range(n):

            h = self.x[i+1] - self.x[i]

            a = self.y[i]

            b = self.b[i]

            c = self.c[i]

            d = self.d[i]



            integral = (

                a * h +

                b * h**2 / 2 +

                c * h**3 / 3 +

                d * h**4 / 4

            )

            y_new[i+1] = y_new[i] + integral



        return NumbaCubicSpline(x_new, y_new, bc_type='natural')



    def integrate(self, a, b):

        if a == b:

            return 0.0

        if a > b:

            return -self.integrate(b, a)



        a = max(self.x[0], a)

        b = min(self.x[-1], b)



        total = 0.0

        n = len(self.x) - 1



        for i in range(n):

            x0 = self.x[i]

            x1 = self.x[i+1]

            if b <= x0 or a >= x1:

                continue

            xl = max(a, x0)

            xr = min(b, x1)

            dxl = xl - x0

            dxr = xr - x0



            a0 = self.y[i]

            b0 = self.b[i]

            c0 = self.c[i]

            d0 = self.d[i]



            def F(x):

                return (

                    a0 * x +

                    b0 * x**2 / 2 +

                    c0 * x**3 / 3 +

                    d0 * x**4 / 4

                )

            total += F(dxr) - F(dxl)



        return total