Improving the numpy functionality of a 2D cosine function

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

1. 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!

>Solution :

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 by doing:

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