Create 2D numpy array applying custom logic

serdar_bay 0
Viewed 28

I have a list of numbers as follows:

list = [5,6,4,7,8]

I need to build a 2d NumPy array using the list above applying the following logic:

arr = np.array([[k1+k2, -k2, 0, 0, 0],
                [-k2, k2+k3, -k3, 0, 0],
                [0, -k3, k3+k4, -k4, 0],
                [0, 0, -k4, k4+k5, -k5],
                [0, 0, 0, -k5, k5]])

k - are the elements in the list, for ex: k1=5, k2=6 and so on.

So, the expected output in this example should look like:

arr = np.array([[11, -6, 0, 0, 0],
                [-6, 10, -4, 0, 0],
                [0, -4, 11, -7, 0],
                [0, 0, -7, 15, -8],
                [0, 0, 0, -8, 8]])

Appreciate any help to build this logic.

pandas numpy python
2 Answers

The best came to my mind is to use several arrays built using np.diagflat() function and sum them up.

lst = [5,6,4,7,8]
res = np.diagflat(lst)
res += np.diagflat(lst[1:] + [0])
res -= np.diagflat(lst[1:], k=-1)
res -= np.diagflat(lst[1:], k=1)
Boris Silantev0

Along with explanation of the logic behind in the other answer provided solution I would like to give here an improved version of Boris Silantevs code for actual use. The improvement reduces the amount of times a full 2D array is created from four times to one time keeping the advantage of numpy vectorization:

L = [5,6,4,7,8]
lenL = len(L)
np_L = np.array(L)
arr  = np.diagflat( np_L + np.append(np_L[1:],0) )
i_diagl = np.arange(lenL-1)
i_subdg = i_diagl + 1 
arr[i_diagl, i_subdg] = arr[i_subdg, i_diagl] = -np_L[1:]

And here the explanation of the logic:

The first code block given below explains the logic behind the numpy vectorized solution providing the equivalent pure Python code using a loop and the second block provides comments on numpy np.diagflat() method for diagonal writing of values to a numpy array:

res = np.array([[11, -6,  0,  0,  0],
                [-6, 10, -4,  0,  0],
                [ 0, -4, 11, -7,  0],
                [ 0,  0, -7, 15, -8],
                [ 0,  0,  0, -8,  8]])
# Appreciate any help to build this logic: 
L = [5,6,4,7,8]
s_L = len(L) 
arr = np.zeros((s_L,s_L))
L += [0] # extend L to allow iterating up to s_L == len(L)
for i in range(s_L):
    arr[i, i] = L[i]+L[i+1] # <-- needs L += [0]
    if i < s_L-1:
        arr[i+1, i] = arr[i,i+1] = -L[i+1]
print (arr)
assert np.all(arr == res)

And below the numpy way of doing the same as the code above:

L = [5,6,4,7,8]
arr = np.diagflat(L) # writes along the diagonal of the array
print(arr) # gives: 
#  [[5 0 0 0 0]
#   [0 6 0 0 0]
#   [0 0 4 0 0]
#   [0 0 0 7 0]
#   [0 0 0 0 8]]
arr += np.diagflat(L[1:] + [0]) # equivalent to L+=[0] in first code block
print(arr) # gives: 
#  [[11  0  0  0  0]
#   [ 0 10  0  0  0]
#   [ 0  0 11  0  0]
#   [ 0  0  0 15  0]
#   [ 0  0  0  0  8]]
arr -= np.diagflat(L[1:], k=-1) # k=-1 writes to diagonal below the main one
print(arr) # gives: 
#  [[11  0  0  0  0]
#   [-6 10  0  0  0]
#   [ 0 -4 11  0  0]
#   [ 0  0 -7 15  0]
#   [ 0  0  0 -8  8]]
arr -= np.diagflat(L[1:], k=1) # k=1 writes to diagonal above the main one
print(arr) # gives: 
#  [[11 -6  0  0  0]
#   [-6 10 -4  0  0]
#   [ 0 -4 11 -7  0]
#   [ 0  0 -7 15 -8]
#   [ 0  0  0 -8  8]]
user95mw0
Related