1
\$\begingroup\$

I have written some code to determine an interpolation matrix over a grid of points. I originally wrote the code in matlab, where under the specified parameters it took just over 2 minutes. The equivalent code in Python however takes 40minutes!

I was expecting Python to be slower, but not be such an amount! Any feedback on the code is much appreciated.

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

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)
\$\endgroup\$
3
\$\begingroup\$

In both Matlab and Numpy, your indexed loops are not the right way to express vectorized operations. You should adopt a strong bias against loops and individual element indexing, learn how broadcasting works, and pass around ndarrays instead of lists.

Based on the above, a potential vectorized solution - with a regression test to ensure that the results have not changed - is

from numbers import Real
from typing import Sequence, Union

import numpy as np
from numpy.linalg import norm
from scipy.stats import describe


EPSILON = 1e-12


def g(xy: np.ndarray) -> np.ndarray:
    return -xy * np.sum(xy**2, axis=0)


def wendland(r: np.ndarray) -> np.ndarray:
    r = norm(r, axis=0)
    out = (1 - r)**4 * (1 + 4*r)
    out[r >= 1] = 0
    return out


def make_collocation_points(h: Real, x1: int, x2: int, y1: int, y2: int) -> np.ndarray:
    j = np.arange(x1, x2+1)
    k = np.arange(y1, y2, h)
    jk = np.stack(np.meshgrid(k, j)) * h

    # Don't accept values on the chain recurrent set!
    g_norm = norm(g(jk) - jk, axis=0)
    jk = jk[::-1, g_norm > 1e-5]

    return jk


def get_aij(xy: np.ndarray) -> np.ndarray:
    # broadcast to get the effect of xj, xk, yj, yk
    xyj, xyk = np.broadcast_arrays(
        xy[:, np.newaxis, :],
        xy[:, :, np.newaxis],
    )

    gj = g(xyj)
    gk = g(xyk)

    return (
        + wendland(gj - gk)  - wendland(xyj - gk)
        - wendland(gj - xyk) + wendland(xyj - xyk)
    )


def assert_close(
    a: Union[Real, np.ndarray],
    b: Union[Real, Sequence[Real]],
) -> None:
    assert np.all(np.isclose(a, b, rtol=0, atol=EPSILON))


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

    xy = make_collocation_points(h, x1, x2, y1, y2)
    assert xy.shape == (2, 185)
    stats = describe(xy.T)
    assert_close(stats.minmax[0], (-0.22, -0.22))
    assert_close(stats.minmax[1], (+0.22, +0.2156))
    assert_close(stats.mean, (0, -2.2e-3))
    assert_close(stats.variance, (0.02433152173913048, 0.016781450543478273))

    a = get_aij(xy)
    assert a.shape == (185, 185)
    stats = describe(a.flatten())
    assert_close(stats.minmax[0], -0.17266964737638246)
    assert_close(stats.minmax[1], +1.10925185890418380)
    assert_close(stats.mean, 0.11771392866408713)
    assert_close(stats.variance, 0.061445578720171375)


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

    xy = make_collocation_points(h, x1, x2, y1, y2)
    print(describe(xy.T))

    a = get_aij(xy)
    print(describe(a.flatten()))


if __name__ == "__main__":
    regression_test()
    demo()

I don't have a whole lot of patience for long execution runs, but for n=10 the demo finishes in about five seconds even though it produces 14,607,684 elements. A quick-and-dirty timeit produces:

old, n=2: 1.433 s
new, n=2: 0.007 s
old, n=3: 6.206 s
new, n=3: 0.035 s
old, n=4: 18.033 s
new, n=4: 0.099 s
old, n=5: 43.935 s
new, n=5: 0.224 s
old, n=6: 90.451 s
new, n=6: 0.459 s
old, n=7: 161.559 s
new, n=7: 0.878 s
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.