complex-function-utils.py

import numpy as np
import scipy as sp

def finite_difference_coefficients(k, xbar, x):
    """
    based on the maths in Finite Difference Methods for Ordinary and Partial Differential Equations
    :param k: the order of the derivative
    :param xbar: the point at which the derivative is to be evaluated
    :param x: the set of points at which the function is known
    :return: the coefficients of the finite difference approximation
    """
    x = np.array(x)
    n = len(x)
    factorial_values = 1 / sp.special.factorial(np.arange(n))
    matrix = np.einsum('i,ji->ij', factorial_values, np.vander(x - xbar, n, increasing=True), dtype=np.complex128)

    b = np.zeros(len(x))
    b[k] = 1
    return sp.linalg.solve(matrix, b)


# divide by zero to nan

def nth_derivative(f, n: int, step_size=1e-5):
    """
    the nth derivative of f numerically around x using finite differences
    does not evaluate f at x
    """

    def diff(x):
        xbar = x
        # sample a circle of complex oints around x but not including x
        number_of_samples = 2 * n + 1
        x = np.array(
            [xbar + step_size * np.exp(2j * np.pi * i / number_of_samples) for i in range(0, number_of_samples)])
        return np.dot(finite_difference_coefficients(n, xbar, x), f(x))

    return diff


def residue(f, pole, order):
    return nth_derivative(lambda z: ((z - pole) ** order * f(z)), order - 1)(pole) / np.math.factorial(order - 1)

@nb.njit
def z(theta):  # funny bromwich contour
    return (
            .1446 + 3.0232 * theta ** 2 / (theta ** 2 - 3.0767 * np.pi ** 2)
            + 0.2339 * 1j * theta
    )


@nb.njit
def z_dash(theta, h=.001):
    return (z(theta + h) - z(theta - h)) / (2 * h)



terms = 200
k = np.arange(1, terms)

h = 2 * np.pi / terms
theta = -np.pi + (k - .5) * h
z_computed = z(theta)
z_dash_computed = z_dash(theta)

t1 = np.exp(z_computed * terms) * z_dash_computed / 1j


@nb.njit(fastmath=True, cache=True)
def inverse_laplace_transform_(F, t):
    return np.sum(t1 * F(terms / t * z_computed) / t, axis=-1)


def inverse_laplace_transform(F):
    F = nb.njit(F, fastmath=True, cache=True)
    return lambda t: inverse_laplace_transform_(F, t)