Issue
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.
Solution
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]]
Answered By - Claudio
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.