I have a first 2D array that represents locations of spatial points. Let's name it A. My first array has a shape of (N, 2) where N represents the number of spatial points and 2 corresponds to both variables x and y. Then, I have a second 2D array corresponding to a time series B of shape (T, 3) where T is the number of timesteps and 3 corresponds to time t and two other variables that I will name y1, y2. What I want is to build a 3D array C of shape (N, T, 2 + 3) that associates for each spatial point the second array B so that C have the following columns: (x, y, t, y1, y2) or (t, y1, y2, x, y). I can do it by loop but with large data, it becomes uneffiient. Do you have a more efficient way to do it? Below, an example:
A = np.array([
[1, 2],
[3, 4]
])
B = np.array([
[-1, -1],
[-2, -2],
[-3, -3]
])
C = np.zeros((2, 3, 4))
for j in range(2):
C[j, :, :2] = np.tile(A[j:j + 1, :], (3, 1))
C[j, :, 2:] = B
print(A)
print(B)
print(C)
Outputs:
A
[[1 2]
[3 4]]
B
[[-1 -1]
[-2 -2]
[-3 -3]]
C
[[[ 1. 2. -1. -1.]
[ 1. 2. -2. -2.]
[ 1. 2. -3. -3.]]
[[ 3. 4. -1. -1.]
[ 3. 4. -2. -2.]
[ 3. 4. -3. -3.]]]
NumPy broadcasting!
First, you need to reshape A and B in such a way that they can be broadcasted together. You can achieve this by adding new axes to both A and B:
A = np.array([
[1, 2],
[3, 4]
])
B = np.array([
[-1, -1],
[-2, -2],
[-3, -3]
])
A_reshaped = A[:, np.newaxis, :] # shape (N, 1, 2)
B_reshaped = B[np.newaxis, :, :] # shape (1, T, 3)
Now, you can use NumPy broadcasting to combine the reshaped A and B arrays then concatenate the broadcasted arrays along the last axis.
A_broadcasted, B_broadcasted = np.broadcast_arrays(A_reshaped, B_reshaped)
C = np.concatenate((A_broadcasted, B_broadcasted), axis=2) # shape (N, T, 2 + 3)
The output will be the expected result:
[[[ 1 2 -1 -1]
[ 1 2 -2 -2]
[ 1 2 -3 -3]]
[[ 3 4 -1 -1]
[ 3 4 -2 -2]
[ 3 4 -3 -3]]]