johnhenry
johnhenry

Reputation: 1333

Adding a new condition on a matrix - Python

After calculating the coordinates of the points on a three-dimensional grid, and storing the values obtained in the following arrays:

import numpy as np

x = np.zeros(4*4*8, float)
y = np.zeros(4*4*8, float)
z = np.zeros(4*4*8, float)

for l in range(4*4*8):
    x[l], y[l], z[l] = np.unravel_index(l, (8, 4, 4))

I define my matrix as

def getMatrix(kappa, x, y, z):


   nxnynz = 4*4*8
   nrange_x = np.arange(nxnynz)
   nrange_y = nrange_x
   w, r = np.meshgrid(nrange_x, nrange_y)
   delta_r = np.sqrt(((x[w]-x[r])**2)+((y[w]-y[r])**2)+((z[w]-z[r])**2)*)

   matrix = np.zeros(delta_r.shape)

   matrix[delta_r == 0] = 4.*kappa/(nxnynz**3.)
   matrix[delta_r != 0] = np.sin(kappa*2.*np.pi*delta_r[delta_r != 0])/(delta_r[delta_r != 0]*float(nxnynz))

where the crucial point is the different definition for those values when delta_r == 0 and when delta_r != 0. Now, in addition to this specification, I need to add a further one, namely the condition (in badly written code, just to give the idea)

if ((abs(w-r)) % 8 ) != 0:
    matrix[w][r] = 0.

where w and r are the same indices used before. How can I fit this new condition in the previous definition of the matrix?

Upvotes: 0

Views: 110

Answers (1)

hpaulj
hpaulj

Reputation: 231510

This should be a straight forward use of boolean indexing. For example:

In [401]: w,r=np.meshgrid(np.arange(3),np.arange(2))
In [402]: w
Out[402]: 
array([[0, 1, 2],
       [0, 1, 2]])
In [403]: r
Out[403]: 
array([[0, 0, 0],
       [1, 1, 1]])
In [404]: mat = np.ones(w.shape,int)
In [405]: mat
Out[405]: 
array([[1, 1, 1],
       [1, 1, 1]])

In [406]: abs(w-r)%8
Out[406]: 
array([[0, 1, 2],
       [1, 0, 1]], dtype=int32)
In [407]: (abs(w-r)%8)!=0
Out[407]: 
array([[False,  True,  True],
       [ True, False,  True]], dtype=bool)

In [408]: mat[(abs(w-r)%8)!=0]=0
In [409]: mat
Out[409]: 
array([[1, 0, 0],
       [0, 1, 0]])

Without fully studying your code, I think your matrix (bad name) has the same shape as w and r. I can construct a boolean array where the condition is satisfied, and easily apply that the matrix.

Upvotes: 1

Related Questions