# Calculating a function over a large array

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)


## 1 Answer

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