Improving the numpy functionality of a 2D cosine function

Question

I have functioning code but I would like help on two aspects:

Improve the numpy code so that the code runs faster (i.e. don’t use for loops)

I would like to first ask about number 1. Here is the code snippet I am inquiring about:

import numpy as np

k = np.linspace(0,4.76*10,2400)
kx,ky = np.meshgrid(k, k)

rx = [0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4]
ry = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.1, 1.2]

tensx = np.tensordot(rx,kx, axes = 0)
tensy = np.tensordot(ry,ky, axes = 0)
z = 0

for i in range(0, len(tensx)):
    z = z + 0.05*np.cos(2*np.pi*(tensx[i]+tensy[i]))

The code is currently not slow but when I start to increase the amount of points in the rx and ry vectors, it will be much slower. So my question: is there a "numpy-er" way to write this code?

Any help is greatly appreciated!

Asked By: mikanim

||

Source

Answer 1

You can re-write the for loop as:

0.05*np.cos(2*np.pi*(tensx+tensy[:-1])).sum(axis=0)

Or more generally:

0.05*np.cos(2*np.pi*(tensx+tensy[:tensx.shape[0]])).sum(axis=0)

You can also avoid the meshgrid, avoiding calculating unused elements of tensy, by doing:

tensx = np.tensordot(rx, k[None, :], axes=0)
tensy = np.tensordot(ry[:len(rx)], k[:, None], axes=0)

Timings:

def op():
    k = np.linspace(0,4.76*10,2400)
    kx,ky = np.meshgrid(k, k)

    rx = [0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4]
    ry = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.1, 1.2]

    tensx = np.tensordot(rx,kx, axes = 0)
    tensy = np.tensordot(ry,ky, axes = 0)
    z = 0
    for i in range(0, len(tensx)):
        z = z + 0.05*np.cos(2*np.pi*(tensx[i]+tensy[i]))
    return z

def Nin17():
    k = np.linspace(0,4.76*10,2400)

    rx = [0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4]
    ry = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.1, 1.2]
    tensx = np.tensordot(rx, k[None, :], axes=0)
    tensy = np.tensordot(ry[:len(rx)], k[:, None], axes=0)
    return 0.05*np.cos(2*np.pi*(tensx+tensy)).sum(axis=0)

assert np.isclose(Nin17(), op()).all()
%timeit Nin17()
%timeit op()

Output:

359 ms ± 2.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
410 ms ± 5.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answered By: Nin17

Improving the numpy functionality of a 2D cosine function

Question:

Answers: