Given my code below:
import numpy as np
nx, ny, nz = (50, 50, 50)
x = np.linspace(1, 2, nx)
y = np.linspace(2, 3, ny)
z = np.linspace(0, 1, nz)
xg, yg, zg = np.meshgrid(x, y, z)
def fun(x, y, z, a, b, c):
r = np.array([x, y, z])
r0 = np.array([a, b, c])
R = ((x-a)**2 + (y-b)**2 + (z-c)**2)**(1/2)
return np.linalg.norm(r - r0)
# return R
evald_z = np.trapz(fun(xg, yg, zg, 1, 1, 1), zg)
evald_y = np.trapz(evald_z, y)
evald_x = np.trapz(evald_y, x)
print(evald_x)
I get the error:
ValueError: operands could not be broadcast together with shapes (3,50,50,50) (3,)
It works simply by changing the return statement (to return R):
1.7086941510502398
Obviously if the function is evaluated at some single (x, y, z) location, the function itself works. I understand why I get this error - when defining the numpy array into r, it is taking full (50, 50, 50) xg, yg and zg grids. I have other points in my code (not shown) where using the (1, 3) shapes for r and r0 come in handy (and make the code more readable). How can I alter the definitions of r and r0 (keeping the (1,3) shape) to perform elementwise operations (as done when using R instead of the norm)? And/or - is this possible, why/why not?
Thanks in advance!!
You can stack up additional dimensions and let Numpy broadcasting do the work:
def fun(x, y, z, a, b, c):
r = np.array([x, y, z])
r0 = np.array([a, b, c])
return np.linalg.norm(r - r0[..., np.newaxis, np.newaxis, np.newaxis], axis=0)