CODEWITHSUNDEEP
pythonarraysnumpy
James Hoyland
James Hoyland

Reputation: 225

Numpy meshes of vectors

I am trying to set up some code for doing some numerical calculations on 3D vector fields such as electric or magnetic fields. I am having trouble setting up my meshes in the way I would like. Consider this program:

import numpy as np

X1, Y1, Z1 = np.meshgrid(np.linspace(-10,10,10),np.linspace(-10,10,10),np.linspace(-10,10,10))

def scalarf(x,y,z):
    return x**2 + y**2 + z**2

def vectorf(x,y,z):
    return np.array([y,x,z])


def afunc(p,v):
    return np.cross(p,v)

V = scalarf(X1,Y1,Z1)
F = vectorf(X1,Y1,Z1)
# This line clearly not working
F2 = afunc(F,np.array([1,0,0]))

print (V.shape)
print (F.shape)

The output of this gives (10,10,10) for V and (3,10,10,10) for F. So V is a 10x10x10 array of scalar values as intended. But for F what I wanted was a 10x10x10 array of 3 element arrays representing mathematical 3D vectors. Instead I have a 3 element array containing a 10x10x10 array as each of its elements. So I'm guessing I want a (10,10,10,3) shape. Ultimately I want to be able to (for ex) run functions like afun in the above. Again here the intention is that F2 would now be a new 10x10x10 array of vectors. At the moment it just fails because I guess its trying to perform a cross product with the 10x10x10 array and the fixed 3d vector in the function call.

Am I going about this in remotely the right way? Is there another way of creating a 3D array of vectors? BTW I have also tried using mgrids with much the same result.

Any help pointing me in the right direction much appreciated.

Upvotes: 0

Views: 455

Answers (1)

Daniel F
Daniel F

Reputation: 14409

Use np.stack with the axis keyword

def vectorf(x, y, z):
    return np.stack((y, x, z), axis = -1)

Upvotes: 1

Related Questions