Using arrays more effectively with numpy

Question

I previously posted some code I've been working on and had a fantastic answer, but as I didn't post the full code I was then unable to bring it all together using the much faster numpy implementation, so a review of the full code is much appreciated. I now realise that using vectorised operations would be much, much faster, but have little experience.

import numpy as np
from numpy import linalg as LA

def g(x, y):
    gx = -x * (x ** 2 + y ** 2)
    gy = -y * (x ** 2 + y ** 2)
    return np.array((gx, gy))

def wendland(r):
    if r < 1:
        return ((1-r)**4)*(4*r+1)
    else:
        return 0

def make_collocation_points(h, x1, x2, y1, y2):
    i = 0
    x = []
    y = []

    for j in range(x1,x2+1):
        for k in np.arange(y1, y2, h):
            if LA.norm(g(j*h, k*h) - (j*h, k*h)) > 0.00001: # Don't accept values on the chain recurrent set!
                x.append(j*h)
                y.append(k*h)

                i += 1

    return x, y, i

def get_aij(x, y, n):
    a = np.zeros((n,n))
    for j in range(0, n):
        for k in range(0, n):
            a[j, k] = (
                wendland(LA.norm(g(x[j], y[j]) - g(x[k], y[k])))
                - wendland(LA.norm(g(x[j], y[j]) - np.array([x[k], y[k]])))
                - wendland(LA.norm([x[j], y[j]] - g(x[k], y[k])))
                + wendland(LA.norm(np.array([x[j], y[j]]) - np.array([x[k], y[k]])))
            )
    return a

def solve_for_B(a, n):
    r = -1 * np.ones(n)
    return LA.solve(a, r)


def find_v_and_OD(x, y, n):
    x_range = np.arange(x1, x2 + h, h)
    y_range = np.arange(y1, y2 + h, h)
    v_val = np.zeros((len(x_range), len(y_range)))
    delV = np.zeros((len(x_range), len(y_range)))

    for i, r in enumerate(x_range):
        for j, q in enumerate(y_range):
            v_val[i, j] = v(r * h, q * h, x, y)
            delV[i, j] = orbital_derivative(r * h, q * h, x, y)
    return v_val, delV


def v(s, t, x, y):
    out = np.multiply(
        b[0],
        (
            wendland(LA.norm(np.array([s, t]) - g(x[0], y[0])))
            - wendland(LA.norm(np.array([s, t]) - np.array(x[0], y[0])))
        ),
    )
    for k in range(1, len(b)):
        out += np.multiply(
            b[k],
            (
                wendland(LA.norm(np.array([s, t]) - g(x[k], y[k])))
                - wendland(LA.norm(np.array([s, t]) - np.array((x[k], y[k]))))
            ),
        )
    return out


def orbital_derivative(s, t, x, y):
    out = np.multiply(
        b[0],
        (
            wendland(LA.norm(g(s, t) - g(x[0], y[0])))
            - wendland(LA.norm(g(s, t) - np.array((x[0], y[0]))))
            - wendland(LA.norm(np.array([s, t]) - g(x[0], y[0])))
            + wendland(LA.norm(np.array([s, t]) - np.array([x[0], y[0]])))
        ),
    )
    for k in range(1, len(b)):
        out += np.multiply(
            b[k],
            (
                wendland(LA.norm(g(s, t) - g(x[k], y[k])))
                - wendland(LA.norm(g(s, t) - np.array((x[k], y[k]))))
                - wendland(LA.norm(np.array([s, t]) - g(x[k], y[k])))
                + wendland(LA.norm(np.array([s, t]) - np.array([x[k], y[k]])))
            ),
        )
    return out

if __name__ == "__main__":
    h = 0.11
    x1, x2 = -15, 15
    y1, y2 = -15, 15

    x, y, n = make_collocation_points(h, x1, x2, y1, y2)
    a = get_aij(x, y, n)
    b = solve_for_B(a, n)
    v_vals, delV = find_v_and_OD(x, y, n)

Answer 1

Gosh, this took a very long time to research and write. Apparently this is an implementation of Computing complete Lyapunov functions for discrete-time dynamical systems. It's a nasty calculation.

Tensors

Your code disguises the fact that many of your variables - a, b, x, y, etc. - should be treated as tensors. This starts in make_collation_points (more on that below), and continues through the rest of the code. The flattened length of your original x, y disguises the fact that these are actually two-dimensional matrices.

One of the more important effects is that solve_for_b shouldn't actually use solve on a two-dimensional matrix at all; instead, it should use tensor_solve on four-dimensional left-hand and two-dimensional right-hand tensors.

Frobenius norm

The norm using np.linalg.norm() is not the fastest option for large tensors when we understand the last dimension to be the only dimension of aggregation. Using an Einstein sum to effectively calculate \$x^2 + y^2\$ followed by a manual root is faster.

Wendland polynomial

The Wendland polynomial expands to

$$1 - 10 r^2 + 20 r^3 - 15 r^4 + 4 r^5$$

implementable as

WENDLAND_POLY = np.polynomial.Polynomial((1, 0, -10, 20, -15, 4), symbol='r')

but using this doesn't particularly make a performance difference.

Chain recurrence

Vectorise make_collocation_points(); an exact equivalent is

def make_collocation_points(
    h: float, x1: int, x2: int, y1: int, y2: int,
    chain_tol: float = 1e-5,
) -> tuple[
    np.ndarray,  # x
    np.ndarray,  # y
]:
    j = np.linspace(start=h*x1, stop=h*x2, num=x2 - x1 + 1)

    '''
    To avoid the O(n) post-multiplication, alternatives for k are:
        k = np.arange(h*y1, h*y2, h**2)
    Or:
        nk = int(np.ceil((y2 - y1)/h))
        k = np.linspace(start=h*y1, stop=h*(y1 + h*(nk - 1)), num=nk)
    These are equivalent to within machine precision, but since there are errors at that order of
    magnitude that prevent the current regression tests from passing, don't use them for now.
    '''
    k = h*np.arange(y1, y2, h)

    x = j[:, np.newaxis]
    y = k
    grid = np.stack(np.meshgrid(k, j))

    # Don't accept values on the chain recurrent set
    norms = np.linalg.norm(grid*(x**2 + y**2 + 1), axis=0)
    y, x = grid[:, norms > chain_tol]  # Off the chain!
    return x, y

However, your chain recurrence check is what forces you to flatten your x, y; and never actually fires for your input range of [-15, 15]. Since it never actually fires, and since the flattening operation loses the import tensor-dimensional information, I'm getting rid of it. You can reintroduce it at the end once the rest of your calculation is done, in case you're paranoid.

Numerical stability

Add regression tests. During regression testing, since the old code is so ponderously slow, I reduce to the much smaller problem of

x1, x2 = -1, 1
y1, y2 = -1, 1

which still produces x, y of 57 elements each, enough to meaningfully test.

The test flags the following change: moving make_collation_points to use the faster np.arange(h*y1, h*y2, h**2) and vectorising get_aij introduce differences at the order of magnitude of machine precision, and this causes the solution of linear tensor system \$Ab = r\$ in solve_for_b to introduce a relative difference of about 2e-5. This linear system is exactly-determined and sensitive to change. I consider this change acceptable in context; so the options are either to keep a low tolerance and replace the expected values of b in the regression test, or replace the tolerance of 1e-8 with 3e-5. I do the former. The dependent variables of v_vals and del_v are not as sensitive and do not need to be replaced when this change is performed.

Gigantic data

For the full 15,15 problem, the fully-broadcast tensor subs alone is 572,978,952 elements, so over 4 GB. This is big enough that you likely need to worry about memory tricks like in-place operations and strategic divide-and-conquer. Despite having 24 GiB my poor little laptop can't run the naive vectorised implementation at full scale. To keep things as simple as possible I don't fix this in the demo code.

Performance

The fully-vectorised implementation is guaranteed faster by many orders of magnitude. How much faster, it isn't really worth testing. I have patience for an outermost dimension of 4, with a new-implementation duration of about a second. The longest tensor shape for this case is

(9*73)**2 * 2**3 == 3_453_192

The time complexity is strongly superlinear.

Suggested

Whether it's a good idea or not, this preserves most of the original function signatures for clarity. Again, if this really needs to run with an outer dimension of 15, the larger tensors probably need to be broken down from those shown here.

import cProfile
import pstats

import numpy as np

# Used for last-step Einstein product on four Wendland terms
TOEPLITZ = np.array((
    ( 1., -1.),
    (-1.,  1.),
))
TOEPLITZ.flags.writeable = False

WENDLAND_POLY = np.polynomial.Polynomial((1, 0, -10, 20, -15, 4), symbol='r')


def g(x: float, y: float) -> np.ndarray:
    return -(x**2 + y**2)*np.array((x, y))


def wendland(r: np.ndarray) -> np.ndarray:
    # w = WENDLAND_POLY(r)
    w = (1 + 4*r)*(1 - r)**4
    w[r >= 1] = 0.
    return w


def ein_norm(x: np.ndarray) -> np.ndarray:
    # linalg.norm(x, axis=-1)  # slower
    inner = np.einsum('...i,...i->...', x, x)
    return np.sqrt(inner)


def make_collocation_points(
    h: float, x1: int, x2: int, y1: int, y2: int,
) -> tuple[
    np.ndarray,  # x
    np.ndarray,  # y
]:
    x = np.linspace(start=h*x1, stop=h*x2, num=x2 - x1 + 1)
    y = np.arange(h*y1, h*y2, h**2)
    return x, y


def get_aij(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    m = x.size
    n = y.size

    '''
    The indices and dimensions in this function need clarity, because this operates in up to
    seven dimensions.
    
    Whenever (x or y) is written, that's one dimension of size 2 that represents a stack
    with one position for all x coordinates and one position for all y coordinates.
    
    The calculation requires that we build up a large tensor (subs) that is then aggregated
    from the rightmost dimension leftward. subs is 7-dimensional; norm() reduces this to 6;
    einsum() reduces this to 4.
    '''

    from numpy import newaxis as na
    x, y = xy = np.stack(np.meshgrid(y, x))  # (x or y = 2)*m*n
    xy = xy.transpose((1, 2, 0))        # m*n*(x or y)
    gxy = g(x, y).transpose((1, 2, 0))  # m*n*(x or y)
    # (gxy or xy) is of length 2, and maps to either of the Wendland term dimensions i,j
    gxyxy = np.stack((gxy, xy), axis=2)  # m*n*(gxy or xy)*(x or y)

    # Wendland minuends
    lhs = np.broadcast_to(
        #      m,n   m,n   i,j  (x or y)
        gxyxy[ :,:, na,na, :,na],
        shape=(m,n,  1,1,  2,1, 2),
    )
    # subtrahend: m,n,(gxy or xy)...,  i,j, (x or y)
    rhs = gxyxy[..., na,:, :]
    norms = ein_norm(lhs - rhs)
    wend = wendland(norms)  # m,n m,n i,j

    # Einstein product over Wendland terms on inner two dimensions
    # output is m,n outer, m,n inner
    return np.einsum('ij,...ij->...', TOEPLITZ, wend)


def solve_for_b(a: np.ndarray) -> np.ndarray:
    r = np.full(shape=a.shape[:2], fill_value=-1, dtype=a.dtype)
    return np.linalg.tensorsolve(a, r)


def find_v_and_od(
    x: np.ndarray, y: np.ndarray,
    x1: int, x2: int, y1: int, y2: int,
    b: np.ndarray, h: float,
) -> tuple[
    np.ndarray,  # v_val
    np.ndarray,  # del_v
]:
    x_range = np.arange(h*x1, h*(x2 + h), h*h, dtype=x.dtype)  # a,
    y_range = np.arange(h*y1, h*(y2 + h), h*h, dtype=y.dtype)  # b,
    v_val = v(s=x_range, t=y_range, x=x, y=y, b=b)                   # a,b
    del_v = orbital_derivative(s=x_range, t=y_range, x=x, y=y, b=b)  # a,b
    return v_val, del_v


def v(s: np.ndarray, t: np.ndarray, x: np.ndarray, y: np.ndarray, b: np.ndarray) -> np.ndarray:
    from numpy import newaxis as na
    y, x = np.meshgrid(y, x)           # each m,n
    xy = np.stack((x, y), axis=-1)  # m,n,2

    lhs = np.stack(
        np.meshgrid(s, t), axis=-1,
    )[..., na,na, na, :]  # a,b, 1,1, 1, (x or y=2)
    rhs = np.stack(  # m,n, (wend term=2), (x or y=2)
        (g(x, y).transpose((1, 2, 0)), xy), axis=2,
    )

    norms = np.linalg.norm(lhs - rhs, axis=-1)  # a,b, m,n, wendterm
    wend = wendland(norms)                      # a,b, m,n, wendterm
    return np.einsum('abijk,ij,k->ba', wend, b, (1, -1))  # b,a


def orbital_derivative(s: np.ndarray, t: np.ndarray, x: np.ndarray, y: np.ndarray, b: np.ndarray) -> np.ndarray:
    from numpy import newaxis as na
    y, x = np.meshgrid(y, x)           # each m,n
    xy = np.stack((x, y), axis=-1)  # m,n,2

    ss, tt = np.meshgrid(s, t)        # each a,b
    gst = g(ss, tt).transpose((1, 2, 0))  # a,b,(s or t=2)
    st = np.stack((ss, tt), axis=-1)  # a,b,(s or t=2)
    gstst = np.stack((gst, st), axis=2)  # a,b,k=2,(s or t=2)

    lhs = np.broadcast_to(
        #          a,b     m,n   k,l   (s or t)
        gstst[     :,:,   na,na, :,na, :],
        gstst.shape[:2] + (1,1,  2,1,  2),
    )
    rhs = np.stack(
        (g(x, y).transpose((1, 2, 0)), xy), axis=2,
    )[:,:, na,:, :]  # m,n, k,l, (x or y)

    norms = ein_norm(lhs - rhs)  # a,b, m,n, k,l
    wend = wendland(norms)       # a,b, m,n, k,l
    # In this notation, ij index into mn; the output is b,a
    return np.einsum('abijkl,ij,kl->ba', wend, b, TOEPLITZ)


def regression_test() -> None:
    h = 0.11
    # x1, x2 = -15, 15  # too slow
    # y1, y2 = -15, 15
    x1, x2 = -1, 1
    y1, y2 = -1, 1

    x, y = make_collocation_points(h=h, x1=x1, x2=x2, y1=y1, y2=y2)
    assert np.allclose(x, (-0.11, 0., 0.11), rtol=0, atol=1e-14)
    assert np.allclose(
        y,
        (
            -0.11, -0.0979, -0.0858, -0.0737, -0.0616, -0.0495, -0.0374,
            -0.0253, -0.0132, -0.0011, 0.011, 0.0231, 0.0352, 0.0473,
            0.0594, 0.0715, 0.0836, 0.0957, 0.1078,
        ), atol=0, rtol=1e-14,
    )

    a = get_aij(x, y)
    assert a.shape == (3, 19, 3, 19)
    b = solve_for_b(a)
    b_expect = [
        [-1.13417555e+02, 1.86990410e+02, -5.90872756e+01,
         6.84671203e+00, -9.52021273e+00, -3.22056034e+00,
         -2.93455717e+00, -1.54528024e+00, -1.36274174e+00,
         -1.48840698e+00, -1.12241753e+00, 3.40039011e-01,
         2.72073810e+00, 4.66361029e+00, 7.68873628e+00,
         3.94670241e+00, 2.26685665e+01, -4.47530131e+01,
         -4.90079716e-01],
        [1.29238665e+02, -1.32344416e+02, 5.92473688e+01,
         -2.12968486e+02, 7.03532353e+02, -2.66787550e+03,
         9.92364052e+03, -3.70538082e+04, 1.38283297e+05,
         -2.19167742e+06, -1.86620295e+05, 4.95516837e+04,
         -1.32808648e+04, 3.55457893e+03, -9.55760927e+02,
         2.74109445e+02, -6.56675655e+01, 2.09491484e+02,
         -1.68829793e+02],
        [-1.13417555e+02, 1.86990410e+02, -5.90872756e+01,
         6.84671201e+00, -9.52021272e+00, -3.22056035e+00,
         -2.93455717e+00, -1.54528025e+00, -1.36274173e+00,
         -1.48840698e+00, -1.12241754e+00, 3.40039035e-01,
         2.72073808e+00, 4.66361030e+00, 7.68873629e+00,
         3.94670239e+00, 2.26685665e+01, -4.47530131e+01,
         -4.90079717e-01],
    ]
    assert np.allclose(b, b_expect, atol=0, rtol=1e-8)

    v_vals, del_v = find_v_and_od(x=x, y=y, x1=x1, x2=x2, y1=y1, y2=y2, b=b, h=h)

    v_expect = np.array([
        [-2.66156654, -2.57725947, -2.46182213, -2.33092178, -2.19623345,
          -2.06524427, -1.9415408, -1.82536358, -1.71427504, -1.60384795,
          -1.48834134, -1.36138632, -1.21674426, -1.04920059, -0.85559453,
          -0.63587836, -0.39403177, -0.13869265, 0.11657553, 0.36811772],
         [-2.66087142, -2.56415052, -2.44558203, -2.31353315, -2.18015795,
          -2.05316124, -1.9364033, -1.83006417, -1.73120529, -1.63436377,
          -1.53247407, -1.4180553, -1.28427149, -1.12567545, -0.93892681,
          -0.72375974, -0.48393153, -0.22807345, 0.03263469, 0.29243954],
         [-2.64877226, -2.53764291, -2.41343994, -2.2813445, -2.15083587,
          -2.02941114, -1.92107085, -1.8261619, -1.74141945, -1.66015215,
          -1.57345745, -1.47252641, -1.35033145, -1.20179358, -1.02370505,
          -0.81548015, -0.58028303, -0.32537788, -0.05988961, 0.20893837],
         [-2.62897029, -2.50368915, -2.37142562, -2.23754602, -2.10965528,
          -1.99425286, -1.89509196, -1.81279752, -1.7442255, -1.68123272,
          -1.61181539, -1.52520273, -1.41505517, -1.27760961, -1.10979993,
          -0.9101071, -0.68030633, -0.42608434, -0.15567925, 0.12249201],
         [-2.60530586, -2.46574541, -2.32379372, -2.18611662, -2.05965525,
          -1.94976298, -1.85943333, -1.78971328, -1.73910359, -1.69843257,
          -1.64901507, -1.57621537, -1.47720902, -1.35145327, -1.19538131,
          -1.00531248, -0.78087606, -0.52637081, -0.25003432, 0.03814834],
         [-2.58114617, -2.42692975, -2.27403827, -2.13104459, -2.00492596,
          -1.89980384, -1.81701537, -1.75727422, -1.72447222, -1.71376376,
          -1.68866639, -1.62470166, -1.5333506, -1.41990085, -1.27733019,
          -1.09794271, -0.87858008, -0.62236359, -0.33840438, -0.03915869],
         [-2.55959473, -2.3903678, -2.22568139, -2.07658599, -1.95038309,
          -1.84988893, -1.77373763, -1.71814361, -1.69293471, -1.73142764,
          -1.74060287, -1.66549608, -1.57718273, -1.47860214, -1.35179922,
          -1.1840766, -0.96939884, -0.70973098, -0.41595268, -0.10452108],
         [-2.54345285, -2.35924823, -2.18272009, -2.02778984, -1.90161596,
          -1.80654989, -1.73702065, -1.68936072, -1.6159707, -1.76063436,
          -1.83246668, -1.67451495, -1.60434262, -1.52401206, -1.41449816,
          -1.25834664, -1.04803916, -0.78330197, -0.47732005, -0.1532218],
         [-2.53480818, -2.33673828, -2.15028496, -1.99142513, -1.86474025,
          -1.77775477, -1.70298543, -1.72106046, -1.45068754, -1.81549842,
          -1.98886535, -1.59812162, -1.63175361, -1.55183706, -1.46024094,
          -1.31235801, -1.10661151, -0.83663808, -0.51677291, -0.18179429],
         [-2.53339713, -2.32635416, -2.1360735, -1.97537768, -1.85006706,
          -1.76137204, -1.70333423, -1.67886762, -1.64584532, -1.69736329,
          -1.68196799, -1.644176, -1.62508522, -1.56914688, -1.47775235,
          -1.33566403, -1.13224279, -0.86080056, -0.53024915, -0.19099517],
         [-2.53410222, -2.33381462, -2.14606588, -1.98678959, -1.85996434,
          -1.77443599, -1.69742099, -1.7306203, -1.42770604, -1.82566247,
          -2.0033213, -1.58140306, -1.63717775, -1.55491365, -1.46605793,
          -1.31915006, -1.11419801, -0.84354381, -0.52118223, -0.18477867],
         [-2.54128235, -2.35442402, -2.17585969, -2.02001823, -1.89383298,
          -1.7999687, -1.73102503, -1.68910178, -1.59235928, -1.76845082,
          -1.85731061, -1.6681283, -1.60838944, -1.53056261, -1.4242594,
          -1.2699828, -1.06047521, -0.79471881, -0.48628973, -0.15997222],
         [-2.55618275, -2.38421099, -2.21733921, -2.06712392, -1.94091461,
          -1.84131707, -1.76642672, -1.71111848, -1.68377088, -1.73547974,
          -1.75305589, -1.67081551, -1.58328796, -1.48793083, -1.36420889,
          -1.1986576, -0.98478898, -0.7243027, -0.42847988, -0.1147441],
         [-2.57694652, -2.42002173, -2.26503389, -2.12096408, -1.99483926,
          -1.89054241, -1.80898011, -1.75048337, -1.72040852, -1.71662391,
          -1.69672354, -1.63299648, -1.54245815, -1.4314229, -1.29154596,
          -1.11421759, -0.89573742, -0.63902817, -0.35345839, -0.05207456],
         [-2.6008561, -2.45866217, -2.31480517, -2.17625567, -2.0499225,
          -1.94094477, -1.85209782, -1.78448688, -1.73724277, -1.70127447,
          -1.65587615, -1.58527633, -1.48800993, -1.36442346, -1.21065738,
          -1.02244913, -0.79896066, -0.54426208, -0.26668274, 0.02342526],
         [-2.62485178, -2.49699451, -2.36307707, -2.22864658, -2.10111244,
          -1.98676835, -1.8892733, -1.80932156, -1.74390302, -1.68460861,
          -1.61859984, -1.53460667, -1.426613, -1.29126119, -1.12547594,
          -0.92745738, -0.69864596, -0.44445967, -0.17307365, 0.10685485],
         [-2.64562366, -2.53191206, -2.40639077, -2.2741288, -2.14415888,
          -2.0238306, -1.91713238, -1.82445166, -1.74249562, -1.6643232,
          -1.58059966, -1.48222741, -1.36221214, -1.2156248, -1.03929297,
          -0.83252304, -0.59828522, -0.34355011, -0.07720702, 0.19329018],
         [-2.65943779, -2.56007688, -2.44072333, -2.3087142, -2.175799,
          -2.04972668, -1.93440211, -1.83002645, -1.73358642, -1.63941234,
          -1.5401538, -1.42811061, -1.29637592, -1.13952088, -0.9542144,
          -0.74015362, -0.50104138, -0.24532778, 0.01628601, 0.27770617],
         [-2.66228719, -2.57644585, -2.46029589, -2.3289334, -2.19429033,
          -2.063878, -1.94131615, -1.82680861, -1.71781464, -1.60973077,
          -1.49660212, -1.37186783, -1.22915662, -1.06316964, -0.87067924,
          -0.65158394, -0.40982021, -0.15408333, 0.10219424, 0.35512225],
         [-2.65411109, -2.5705731, -2.46078991, -2.33329649, -2.19963478,
          -2.06689952, -1.9387606, -1.81573109, -1.69590434, -1.57559107,
          -1.44990043, -1.31328753, -1.16022145, -0.98607427, -0.78816892,
          -0.56683993, -0.32620562, -0.07491521, 0.17592464, 0.42191875],
    ])
    assert np.allclose(v_vals, v_expect, atol=0, rtol=1e-6)

    assert np.allclose(
        del_v,
        [[-1., -1., -1., -1., -1.,
          -1., -1., -1., -1., -1.,
          -1., -1., -1., -1., -1.,
          -1., -1., -1., -1., -1.01462631],
         [-1.03431034, -1.01856243, -1.00127954, -0.99011636, -0.98419595,
          -0.98244512, -0.98350584, -0.98628558, -0.99015487, -0.99506395,
          -1.00125302, -1.00865335, -1.01647273, -1.02329191, -1.02746664,
          -1.02745804, -1.02205102, -1.0102714, -0.99406302, -0.98622636],
         [-1.05203371, -1.02600875, -0.99766217, -0.97773682, -0.96734186,
          -0.96460189, -0.96689254, -0.97185504, -0.97819515, -0.98622667,
          -0.99723707, -1.01164782, -1.02779282, -1.04242352, -1.05189686,
          -1.05301531, -1.04353214, -1.02279202, -0.99424635, -0.96917811],
         [-1.05654846, -1.02402083, -0.99077753, -0.96696178, -0.95469558,
          -0.95202808, -0.95547543, -0.96138484, -0.96748985, -0.97489063,
          -0.98745741, -1.00708883, -1.03077927, -1.05287771, -1.06786651,
          -1.07125106, -1.06017089, -1.0344337, -0.99792749, -0.96062188],
         [-1.0516481, -1.01686835, -0.98378809, -0.96065473, -0.94933021,
          -0.94795637, -0.95277454, -0.9587986, -0.96139315, -0.96217464,
          -0.97137474, -0.99454929, -1.02491429, -1.05299736, -1.07247657,
          -1.07880032, -1.06869468, -1.04189696, -1.00211902, -0.95774033],
         [-1.04128577, -1.00852302, -0.97977648, -0.96072923, -0.95231765,
          -0.95295365, -0.95961937, -0.9668755, -0.96408988, -0.94815262,
          -0.94689351, -0.97566828, -1.01338683, -1.04456208, -1.06558725,
          -1.07417607, -1.06704315, -1.04285138, -1.00445841, -0.95814092],
         [-1.02877243, -1.0018602, -0.98061913, -0.96763666, -0.96274873,
          -0.96486055, -0.97289958, -0.98528409, -0.98495943, -0.93034276,
          -0.90567016, -0.95681227, -1.00320737, -1.03177988, -1.04956851,
          -1.0581948, -1.05489009, -1.03654448, -1.00396196, -0.96040155],
         [-1.01647055, -0.99838365, -0.98641905, -0.97979171, -0.9778775,
          -0.97942868, -0.98700827, -0.99843832, -1.05373146, -0.90061119,
          -0.82109136, -0.96307574, -0.99965889, -1.01879414, -1.02855463,
          -1.03475405, -1.03452741, -1.0242471, -1.00142544, -0.9641884],
         [-1.00604695, -0.99825616, -0.99487767, -0.99274012, -0.9934085,
          -0.99006944, -1.00683903, -0.95694344, -1.21386334, -0.84533341,
          -0.66905721, -1.04868632, -0.98649941, -1.01056962, -1.00812459,
          -1.01188866, -1.01219619, -1.00966308, -0.99935379, -0.96934585],
         [-1.00006585, -0.99996515, -0.99996435, -0.99991924, -1.00002834,
          -0.999624, -1.0011621, -0.99547038, -1.01677978, -0.96330814,
          -0.97756512, -1.00603231, -0.99841953, -1.00048098, -0.99995316,
          -1.00012156, -1.00009019, -1.00010732, -0.99996825, -0.97295567],
         [-1.00449625, -0.99852587, -0.99630165, -0.99465665, -0.9957135,
          -0.99129194, -1.01076683, -0.94625699, -1.23625276, -0.83512015,
          -0.65509718, -1.06645227, -0.98269343, -1.00971829, -1.00518514,
          -1.00864909, -1.00876814, -1.0072396, -0.99927026, -0.97028719],
         [-1.01438992, -0.99812645, -0.98787572, -0.98227111, -0.98082621,
          -0.98189105, -0.98977662, -0.99646942, -1.07617124, -0.89270399,
          -0.79725014, -0.9715785, -0.9988849, -1.01674437, -1.0245998,
          -1.03024507, -1.03034947, -1.02161324, -1.00092047, -0.9650485],
         [-1.02646545, -1.00096934, -0.98133528, -0.96954294, -0.96525629,
          -0.96746314, -0.97551135, -0.98905909, -0.99242592, -0.92620241,
          -0.89477553, -0.95474351, -1.00211759, -1.02932884, -1.04597662,
          -1.05431383, -1.05167737, -1.03468056, -1.00359926, -0.96097317],
         [-1.03908851, -1.00712362, -0.97953925, -0.96149241, -0.9537182,
          -0.95470676, -0.96172029, -0.96957411, -0.96607409, -0.94533357,
          -0.94107808, -0.97189789, -1.01118526, -1.04242471, -1.06325179,
          -1.07204731, -1.06558887, -1.04223683, -1.00458073, -0.95843647],
         [-1.05006159, -1.01533827, -0.98276216, -0.96015778, -0.9492525,
          -0.94821883, -0.95333973, -0.95941795, -0.96106919, -0.95974907,
          -0.96766668, -0.99149676, -1.02308064, -1.05199662, -1.07203955,
          -1.07887565, -1.06918777, -1.04260052, -1.00273255, -0.95762598],
         [-1.0562381, -1.02297751, -0.98941403, -0.96541139, -0.95312074,
          -0.95060174, -0.95426982, -0.9602288, -0.96593751, -0.97265892,
          -0.9850275, -1.00538597, -1.03033135, -1.05366143, -1.06959091,
          -1.07348812, -1.06240323, -1.0361869, -0.99873461, -0.95975538],
         [-1.05373613, -1.02623014, -0.99657315, -0.97556425, -0.96461201,
          -0.96179555, -0.96429831, -0.9695294, -0.97605765, -0.98431198,
          -0.9958993, -1.01141657, -1.02901411, -1.04506266, -1.05555052,
          -1.05699876, -1.04702648, -1.02507759, -0.99474375, -0.96707428],
         [-1.03865624, -1.02077813, -1.00096169, -0.98790377, -0.9810249,
          -0.97903284, -0.9803214, -0.983572, -0.98803897, -0.99370577,
          -1.00093171, -1.00968803, -1.01903191, -1.02724098, -1.03232305,
          -1.03243577, -1.02614032, -1.01245585, -0.99370768, -0.9823681],
         [-1.0077488, -1.00394675, -1.00068633, -0.99864236, -0.9975126,
          -0.99713315, -0.99727989, -0.9977778, -0.99851502, -0.99945456,
          -1.00059239, -1.00188335, -1.00318911, -1.00428589, -1.00491446,
          -1.00482165, -1.00380004, -1.00162325, -0.99835853, -1.00847627],
         [-0.95526052, -0.97935174, -0.99284054, -1.00215138, -1.00744539,
          -1.00963186, -1.00943266, -1.00752499, -1.00438605, -1.0003134,
          -0.99553571, -0.99039507, -0.98547564, -0.98159809, -0.97971535,
          -0.98076808, -0.98568129, -0.99512324, -1.01170936, -1.0487705]],
        atol=0, rtol=1e-6,
    )


def big_benchmark() -> None:
    n = 4
    x1, x2 = -n, n
    y1, y2 = -n, n
    h = 0.11

    with cProfile.Profile(builtins=False) as profile:
        x, y = make_collocation_points(h=h, x1=x1, x2=x2, y1=y1, y2=y2)
        a = get_aij(x, y)
        b = solve_for_b(a)
        v_vals, del_v = find_v_and_od(x=x, y=y, x1=x1, x2=x2, y1=y1, y2=y2, b=b, h=h)

    stats = pstats.Stats(profile).sort_stats(pstats.SortKey.CUMULATIVE)
    stats.print_stats(12)


if __name__ == '__main__':
    regression_test()
    big_benchmark()

         238 function calls in 0.787 seconds

   Ordered by: cumulative time
   List reduced from 50 to 20 due to restriction <20>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.002    0.002    0.720    0.720 268055.py:84(find_v_and_od)
        1    0.204    0.204    0.416    0.416 268055.py:116(orbital_derivative)
        1    0.088    0.088    0.303    0.303 268055.py:99(v)
        3    0.213    0.071    0.213    0.071 268055.py:18(wendland)
        1    0.139    0.139    0.139    0.139 venv/lib/python3.12/site-packages/numpy/linalg/_linalg.py:2624(norm)
        5    0.098    0.020    0.098    0.020 venv/lib/python3.12/site-packages/numpy/_core/einsumfunc.py:1057(einsum)
        1    0.035    0.035    0.059    0.059 268055.py:35(get_aij)
        1    0.000    0.000    0.008    0.008 268055.py:79(solve_for_b)
        1    0.000    0.000    0.008    0.008 venv/lib/python3.12/site-packages/numpy/linalg/_linalg.py:237(tensorsolve)
        1    0.008    0.008    0.008    0.008 venv/lib/python3.12/site-packages/numpy/linalg/_linalg.py:320(solve)
        5    0.000    0.000    0.000    0.000 venv/lib/python3.12/site-packages/numpy/lib/_function_base_impl.py:5077(meshgrid)
        5    0.000    0.000    0.000    0.000 venv/lib/python3.12/site-packages/numpy/lib/_stride_tricks_impl.py:495(broadcast_arrays)
       12    0.000    0.000    0.000    0.000 venv/lib/python3.12/site-packages/numpy/lib/_stride_tricks_impl.py:350(_broadcast_to)
        9    0.000    0.000    0.000    0.000 venv/lib/python3.12/site-packages/numpy/_core/shape_base.py:377(stack)
        4    0.000    0.000    0.000    0.000 268055.py:14(g)
        1    0.000    0.000    0.000    0.000 268055.py:24(make_collocation_points)
        1    0.000    0.000    0.000    0.000 venv/lib/python3.12/site-packages/numpy/_core/function_base.py:25(linspace)
       15    0.000    0.000    0.000    0.000 venv/lib/python3.12/site-packages/numpy/lib/_function_base_impl.py:5231(<genexpr>)
        2    0.000    0.000    0.000    0.000 venv/lib/python3.12/site-packages/numpy/lib/_stride_tricks_impl.py:377(broadcast_to)
        1    0.000    0.000    0.000    0.000 .pyenv/versions/3.12.0/lib/python3.12/cProfile.py:119(__exit__)