How to loop through axis for subplots in a while loop?

Question

How can i use my while loop conditions to loop through axis of subplots. I am very new to python. I have written a code where i plot the data based on some conditions in a while loop. The code works fine for plotting individual plots but when i have to plot all in one plot like subplots, i don't know how can i index that every time one round of plotting is finished the axis index should change and the next plotting is done on different index. As in the image it can be seen, the first row is plotted and rest everything is plotted all over again as the loop has no condition to go for axis indexing. How can i achieve is that every time the value of i,j.l,m in the code is incremented, the plots should move to next row in subplot figure.[ enter image description here ]

import matplotlib.pyplot as plt
import xarray as xr
file="/home/arjun2516/hlrn/JOBS/cy_topo_ref4hr2/OUTPUT/cy_topo_ref4hr2_av_3d.002.nc"
Data=xr.open_dataset(file)

l=150
m=300
i = 75
j = 175
while i and j < 700 and l and m < 800 :
    fig,axs = plt.subplots(ncols=3, nrows=3,figsize=(20,20))
    Data.zusi[i:j,75:175].plot.contourf(ax=axs[0,0])
    print(i,j)
    # plt.show()
    Data.zusi[l:m,250:400].plot.contourf(ax=axs[0,1])
    #  plt.show()
    Data.zusi[l:m,450:600].plot.contourf(ax=axs[0,2])
    # plt.show()
    i += 200
    j += 200
    l += 200
    m += 200
    print(i,j)
print('ok')

I tried to introduce a for loop inside the while loop but it was also producing same results.

score 0 · Accepted Answer · answered Jan 17 '23 at 21:48

There is several problems in your code:

You are creating a new figure object (containing the grid of subplots) in every loop iteration, so the plots from different iterations will end up in different figures. Move the plt.subplots command before the loop.
In order to plot onto the axis of a different row in each loop iteration, you need an axis index that starts at zero (that is, indexing the first row) and is incremented in each iteration.

With these changes, your code becomes:

l=150
m=300
i = 75
j = 175

fig,axs = plt.subplots(ncols=3, nrows=3,figsize=(20,20))
ax_idx = 0

while i and j < 700 and l and m < 800 :
    # Select axis based on the axis index
    Data.zusi[i:j,75:175].plot.contourf(ax=axs[ax_idx,0])
    print(i,j)
    # plt.show()
    Data.zusi[l:m,250:400].plot.contourf(ax=axs[ax_idx,1])
    #  plt.show()
    Data.zusi[l:m,450:600].plot.contourf(ax=axs[ax_idx,2])
    # plt.show()
    i += 200
    j += 200
    l += 200
    m += 200
    # Increase the axis index
    ax_idx += 1
    print(i,j)

Note that you could also simplify your code by using a for loop. I would also highly recommend using xarray's capabilities for label-based indexing, in this case isel. It makes the code a little bit more verbose, but much more understandable.

n_rows = 3
fig,axs = plt.subplots(ncols=3, nrows=n_rows, figsize=(20,20))
ax_idx = 0

for ax_idx in range(n_rows):
    # Compute the index values
    l = 150 + ax_idx * 200
    m = 300 + ax_idx * 200
    i = 75 + ax_idx * 200
    j = 175 + ax_idx * 200
    # Index array based on named dimensions and plot it
    Data.zusi.isel(x=slice(i, j), y=slice(75, 175)).plot.contourf(ax=axs[ax_idx, 0])
    Data.zusi.isel(x=slice(l, m), y=slice(250, 400)).plot.contourf(ax=axs[ax_idx, 1])
    Data.zusi.isel(x=slice(l, m), y=slice(450, 600)).plot.contourf(ax=axs[ax_idx, 2])
    print(i,j)

How to loop through axis for subplots in a while loop?

1 Answers1