Is there a more efficent way of "expanding" a numpy array?

Question

An algorithm requires me to add an additional variable into a square matrix, this results in adding a row and column to the matrix. I can use the np.insert method as shown below, but was wondering if there was a better way to do this to an already existing allocated array?

I was trying to play around with a masked write, but have not been successful:

import numpy as np

m1 = np.array([[1,3,4,5],
               [11,13,14,15],
               [16,18,19,20],
               [21,23,24,25]
              ])

row_1 = np.array([6,8,9,10])
col_1 = np.array([2,7,12,17,22])

m2 = np.insert(m1, 1, row_1, axis=0)
m2 = np.insert(m2, 1, col_1, axis=1)
print(f'{m2=}')

# create a mask to copy data from m1
m1_mask = np.ones(tuple(np.array(m1.shape) + 1), dtype=int)
expanded_m1 = np.zeros(tuple(np.array(m1.shape) + 1))
m1_mask[:,1] = 0
m1_mask[1, :] = 0
print(m1_mask)
# only write the indices where there are 1's from the original m1 data
expanded_m1[m1_mask] = m1
print(expanded_m1)

Asked By: Jordan Larson

||

Source

Answer 1

Your expanded_m1[m1_mask] = m1 fails because the LHS refers to (5, 5, 5) array.

>>> expanded_m1[m1_mask].shape
(5, 5, 5)

It becomes a mask if you make m1_mask a boolean array with dtype=bool at its definition. After setting m1_mask[:, 1] = 0 and m1_mask[1, :] = 0, we have m1_mask as:

[[ True False  True  True  True]
 [False False False False False]
 [ True False  True  True  True]
 [ True False  True  True  True]
 [ True False  True  True  True]]

Then, expanded_m1[m1_mask] returns a vector of elements, because it has no way to know that the result is actually a non-ragged array.

>>> expanded_m1[m1_mask].shape
(16,)

>>> m1_mask.sum()
16

You can set these elements if you also flatten m1:

>>> expanded_m1[m1_mask] = m1.flatten()
>>> print(expanded_m1)
[[ 1.  0.  3.  4.  5.]
 [ 0.  0.  0.  0.  0.]
 [11.  0. 13. 14. 15.]
 [16.  0. 18. 19. 20.]
 [21.  0. 23. 24. 25.]]

Then, you can set your new row and column. Some shenanigans are required because row has only four elements, but that is easy to handle:

row_mask = m1_mask[0]
expanded_m1[row_mask] = row_1
expanded_m1[:, 1] = col_1

Answered By: Pranav Hosangadi

Is there a more efficent way of "expanding" a numpy array?

Question:

Answers: