2D line plot a 3D Numpy matrix / array given a chosen axis

Question

How can I plot a 2D line from a chosen axis of a Numpy Array quickly?

An analogy: when sum an arbitrary matrix sigma with respect to axis = 0, I would write:

import numpy as np
import matplotlib.pyplot as plt
sigma = np.array([
 [[0. , 0.9, 0.6],
 [0. , 0. , 0.4],
 [0. , 0. , 0. ]],
 [[0. , 0.8, 0.5],
 [0. , 0. , 0.3],
 [0. , 0. , 0. ]],
 [[0. , 0.7, 0.4],
 [0. , 0. , 0.2],
 [0. , 0. , 0. ]]
 
 ])
np.sum(sigma, axis=0)

with result:

array([[0. , 2.4, 1.5],
 [0. , 0. , 0.9],
 [0. , 0. , 0. ]])

I am seeking an equivalent straight forward method to plot axis=0, suggestively similar to:

plt.plot(sigma, axis=0)

This means, I will plot the depth of the matrix at each corresponding position. In the plot I will see three lines, one line starting at 0.9 in value at x =1, and 0.8 at x=2, and 0.7 at x-3. Similarly, for lines two and three, [0.6, 0.5, 0.4] ; [0.4, 0.3, 0.2].

I could find examples of plot 3d and a method (involving slice and len) for plot 2d that would yield in a solution similar to:

plt.plot(sigma[:,:,2])

However, I cannot get it to plot against the x-axis (x = 1..3, representing each layer of array)

How do I do it?

Update: a solutions seems to be:

plt.plot(sigma[:,:,:].reshape((3, 9)))

• yann ziselman: It's hard to understand what you need exactly. I'm not sure what you mean by plotting an axis. it looks to me that you're trying to plot three points from each row of a certain slice (:,:,2) in the array. what is the functionality you want the argument 'axis' to have in this situation? What do you mean by "a 1, 2, 3 x-axis"? do you mean that you want each point in y to be placed in the positions [1, 2, 3] on the x-axis?
• Jaco: Hi Yan, with axis I mean how you would cut the 3D matrix into a 2D shape, Assume you have three layers, each layer represent a point in time, so at t=0 you have a 3x3 array, at t=1 you have another 3x3 array. How would you plot, eg. the value in the top right corner throughout this three point in time, e.g. plot 0.6, 0.5, and 0.4 (top right position in the array). When you sum (np.sum(sigma, axis-0), you will sum this dimension consisting of three values. Hope this is a bit more clear

Answer 1

If I understood your question, your first dimension is a time, for which you have a 2D array at each time point, and you want to see how a given index in that 2D array evolves.

One way to approach (so that you don't have to keep copying data, assuming you have a large dataset), is to flatten your original sigma array and index the 2D array locations.

>> sigma.flatten()
array([0. , 0.9, 0.6, 0. , 0. , 0.4, 0. , 0. , 0. , 0. , 0.8, 0.5, 0. ,
 0. , 0.3, 0. , 0. , 0. , 0. , 0.7, 0.4, 0. , 0. , 0.2, 0. , 0. ,
 0. ])

Then, for each timestep in your 3x3 case, you could get the:

• [0, 0] index by indexing the data at locations [0, 9, 18]
• [0, 1] index by indexing [1, 10, 19]

etc of the flattened array.

A quick demo based on your example data:

import numpy as np
import matplotlib.pyplot as plt
sigma = np.array([
 [[0., 0.9, 0.6],
 [0., 0., 0.4],
 [0., 0., 0.]],
 [[0., 0.8, 0.5],
 [0., 0., 0.3],
 [0., 0., 0.]],
 [[0., 0.7, 0.4],
 [0., 0., 0.2],
 [0., 0., 0.]]
])
n, a, b = sigma.shape
n_ar = a * b # the len of a 2D array
x = np.arange(n) # The 2D array slice indices, [0, 1, 2]
sigma_flat = sigma.flatten() # Flatten into 1D array and index for points
fig, ax = plt.subplots() # Set up a figure and axes
for i in range(n_ar):
 idxs = x * n_ar + i # Get indices of flattened array
 ax.plot(x+1, sigma_flat[idxs], label=f'Loc {i}')
fig.legend()
plt.show()

Returns:

Example plot

Is that what you were trying to do?

• gboffi: plt_ij = lambda i, j: plt.plot(sigma[:, i, j])
• Jaco: Thanks, I also managed to do it by: plt.plot(sigma[:,:,:].reshape((3, 9)))