Expand matrix based on vector - numpy

Question

I want to turn matrix A into matrix B.
Is there a better/more efficient approach with NumPy than the following?
import numpy as np
a = np.array([[0.02, 0.05, 0.05],
[0.35, 0.10, 0.45],
[0.08, 0.25, 0.15]])
w = np.array([0.75, 0.25])
B = np.insert(a, 9, a[2, :]).reshape(4, 3)
B = np.insert(B.T, 12, B[:, 2]).reshape(4, 4).T
B[2:4, :] = np.multiply(B[2:4, :].T, w).T

Answer 1

.insert isn't a good choice here because numpy needs to allocate memory to create a whole new array every time you do so. Instead, just pre-allocate the size of array you need, and then assign to its slices.
a = np.array([[0.02, 0.05, 0.05],
[0.35, 0.10, 0.45],
[0.08, 0.25, 0.15]])
w = np.array([0.75, 0.25])
b_shape = tuple(s + 1 for s in a.shape) # We need one more row and column than a
b = np.zeros(b_shape) # Create zero array of required shape
b[:a.shape[0], :a.shape[1]] = a # Set a in the top left corner
b[:, -1] = b[:, -2] # Set last column from second-last column
b[-1, :] = b[-2, :] # Set last row from second-last row
b[-w.shape[0]:, :] = b[-w.shape[0]:, :] * w[:, None] # Multiply last two rows with `w`
w[:, None] makes w a column vector (a 2x1 matrix), and numpy broadcasts the shapes to do the correct elementwise multiplication.
This gives us the required b:
array([[0.02 , 0.05 , 0.05 , 0.05 ],
[0.35 , 0.1 , 0.45 , 0.45 ],
[0.06 , 0.1875, 0.1125, 0.1125],
[0.02 , 0.0625, 0.0375, 0.0375]])
Putting this in a function to compare runtimes against your approach:
import numpy as np
import timeit
from matplotlib import pyplot as plt
#% Define functions
def func_insert(a, w):
B = np.insert(a, a.size, a[-1, :]).reshape(a.shape[0]+1, a.shape[1])
B = np.insert(B.T, B.size, B[:, -1]).reshape(a.shape[0]+1, a.shape[1]+1).T
B[-w.shape[0]:, :] = np.multiply(B[-w.shape[0]:, :].T, w).T
return B
def func_prealloc(a, w):
b_shape = tuple(s + 1 for s in a.shape)
b = np.zeros(b_shape)
b[:a.shape[0], :a.shape[1]] = a
b[:, -1] = b[:, -2]
b[-1, :] = b[-2, :]
b[-w.shape[0]:, :] = b[-w.shape[0]:, :] * w[:, None]
return b
#% Time function calls
sizes = [3, 10, 50, 100, 500, 1000, 5000, 10_000]
times = np.zeros((len(sizes), 2))
for i, size in enumerate(sizes):
a = np.random.random((size, size))
w = np.random.random((2,))
times[i, 0] = timeit.timeit("func_insert(a, w)", globals=globals(), number=10) / 10
print(".")
times[i, 1] = timeit.timeit("func_prealloc(a, w)", globals=globals(), number=10) / 10
print("x")
#% Plot results
fig, ax = plt.subplots()
ax.plot(sizes, times[:, 0], label="Insert")
ax.plot(sizes, times[:, 1], label="Prealloc")
ax.set_xscale('log')
ax.set_yscale('log')
ax.legend()
ax.set_xlabel('Array size (NxN)')
ax.set_ylabel('Time per function call (s)')
ax.grid(True)
fig.tight_layout()
]
There's a consistent 3-5x speedup by preallocating.