Numpy Value Error: Elementwise Operations on Meshgrid

Question

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

score 0 · Accepted Answer · answered Feb 22 '23 at 15:13

0

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)

answered Feb 22 '23 at 15:13

Kuroneko

146
6

Can you explain a little more in depth what is actually happening here and why it works? I can't even say I'm very familiar with ellipses in Python – Sterling Butters Feb 22 '23 at 15:35
It is worth noting here that this method is about 2.3x slower (per evaluation for 1000 evaluations) than `R = ((x-a)**2 + (y-b)**2 + (z-c)**2)**(1/2)`. (Your answer is what I was looking for - just wanted to comment on the performance aspect) – Sterling Butters Feb 22 '23 at 15:57
@SterlingButters , I suggest reading the Numpy manual on indexing to better understand ellipsis and newaxis notations https://numpy.org/doc/stable/user/basics.indexing.html#dimensional-indexing-tools – Kuroneko Feb 22 '23 at 16:32
`norm` is using `(((r-r0[:,None,None,None])**2).sum(axis=0)`, where `r` is (3,50,50,50) array. You are doing 3 squares on (50,50,50) array. Often a few iterations on a complex operation can be faster than a single one on even larger arrays. It comes down to details of memory management for large arrays. – hpaulj Feb 22 '23 at 18:13

Numpy Value Error: Elementwise Operations on Meshgrid

1 Answers1