[Python] `numpy.irfftn` returns incorrect shape after multiplying `numpy.rfft`-transformed...

I'm writing a code to perform differential operations on 3D fields by exploiting the fact that derivatives in configuration space (x, y, z) become products in Fourier space (kx, ky, kz).

This is my code:

ki = 2*np.pi/L*kvecs(Ng)
k = 2*np.pi/L*knorms(ki)
vel = np.zeros((*field.shape, 3))
for i in range(3):
velk = rfftn(field)
print(velk.shape)
for x in range(Ng):
for y in range(Ng):
for z in range(Ng//2+1):
velk[x, y, z] *= -1j*ki[x, y, z, i]/k[x, y, z]**2 if k[x, y, z]!=0 else 0

print(irfftn(velk).shape)
vel[..., i] = irfftn(velk)


where field is an array of shape (Ng, Ng, Ng) where Ng=5, ki is a (Ng, Ng, Ng, 3)-array with the (kx, ky, kz) components of each pixel, and k is the norm of ki (that is the square root of kx^2+ky^2+kz^2, so it has shape (Ng, Ng, Ng)). The code code clearly doesn't change the shape of velk at all, just multiplies some factors to its values, yet the two print statements give (5, 5, 3) (correct) and (5, 5, 4) (wrong, it should be the original shape of (5, 5, 5)), and the code errors out at the last line because it cannot broadcast the inverse transform.

If I comment out the last three nested loops, the code works, and if I swap the real fft for the complex one (while also changing the range for the z loop from range(Ng//2+1) to range(Ng)) the code works as expected.

Is this a bug or am I missing something here?

Continue reading...