CODEWITHSUNDEEP
pythonnumpymatplotlib
Jaco
Jaco

Reputation: 1694

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

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)))

Upvotes: 0

Views: 679

Answers (1)

Tim Jim
Tim Jim

Reputation: 670

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?

Upvotes: 1

Related Questions