Search code examples
pythonfor-loopmatplotlibcartopy

Python nested for loop is not iterating properly through both loops


I am attempting to iterate over both a collection of 'datasets' as well as a collection of 'domains'. The idea is to pair the first dataset with the first domain, the second dataset with the second domain, and so on. Ultimately I want to plot variables from these datasets on a map. I tried using a nested for loop, but it refuses to loop through either the dataset list or the domain list. This ends up with the program plotting the same dataset on the same domain over and over again. The relevant code is here:

# This extracts files and directories out of the path
arr = os.listdir(path)

# Create list for WRF output files (since files in case directories are separated by time; i.e. 12Z, 13Z, 14Z)
wrf_out = []

# For loop to fill list with WRF output files from case directory
for file in arr:
    # os.path.join is REQUIRED because it joins the file to the path; otherwise you'll get the file name with NO data
    filepath = os.path.join(path, file)
    wrf_out.append(filepath)

# Sort WRF output files chronologically
wrf_out.sort()

# Test to make sure you have the right case/mpscheme/res/vert
#ds = Dataset(wrf_out[0])
#print(ds)

# Basic NWS colors for reflectivity
dbz_contour_levels = np.arange(0,80.,5.)
# Taken from UCAR MMM/COD NEXLAB Archive Mosiac's table
dbz_colors = ("#00ffff", "#0080ff", "#0000ff", "#00ff00", "#00c000", "#00800a", "#ffff00", "#ffc000", "#ff8000",
              "#ff0000", "#c00000", "#800000", "#ff00ff", "#8d67cd", "#000000")

domain_2015_06_03 = [[-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], 
                     [-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], 
                     [-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], [-108.714, -101.222, 40.776, 45.348], 
                     [-108.714, -101.222, 40.776, 45.348], [-105.858, -98.827, 41.373, 45.840], [-105.858, -98.827, 41.373, 45.840], [-104.012, -95.882, 40.843, 45.902], 
                     [-104.012, -95.882, 40.843, 45.902], [-104.012, -95.882, 40.843, 45.902], [-104.012, -95.882, 40.843, 45.902], [-104.012, -95.882, 40.843, 45.902], 
                     [-101.222, -94.564, 40.007, 46.389], [-101.222, -94.564, 40.007, 46.389], [-100.738, -92.081, 39.972, 46.721], [-100.738, -92.081, 39.972, 46.721], 
                     [-100.738, -92.081, 39.972, 46.721]]

# For loop to get variables, create levels, plot
for i in range(len(wrf_out)):
    
    # For loop to create a shifting domain for higher resolution of variables
    for j in domain_2015_06_03:
        lonmin = j[0]
        lonmax = j[1]
        latmin = j[2]
        latmax = j[3]
        
        zoom_domain = [lonmin, lonmax, latmin, latmax]
        
        # Get dataset
        ds = Dataset(wrf_out[i])
        
        # Get WRF Variables
        theta_2m = getvar(ds, 'TH2')
        t_pert = getvar(ds, 'T')
        # wspd_10m is for divergence and arrows
        wspd_10m = getvar(ds, 'uvmet10', units='m s-1')
        # wind_10m is for contour
        wind_10m = getvar(ds, 'uvmet10_wspd_wdir', units='m s-1')[0,:]
        lat = getvar(ds, 'lat')
        lon = getvar(ds, 'lon')
        # mdbz is for reflectivity
        mdbz = getvar(ds, 'mdbz')

        # Calculate components and wrf lat lons
        wrf_lat, wrf_lon = latlon_coords(mdbz)
        wrf_lons = to_np(wrf_lon)
        wrf_lats = to_np(wrf_lat)
        x = wrf_lon.values
        y = wrf_lat.values
        # u_comp and v_comp are for divergence
        u_comp = wspd_10m[0,:,:]
        v_comp = wspd_10m[1,:,:]
        # u and v are for arrows
        u = u_comp.values
        v = v_comp.values

        # Get grid deltas
        #dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat)

        # Calculate divergence
        #div = mpcalc.divergence(u_comp, v_comp, dx=dx, dy=dy)

        # Just need values for calculating contour levels; otherwise error thrown for values with dimensions
        #div_vals = div.values

        # Create theta levels
        max_level_theta = np.max(theta_2m)
        min_level_theta = np.min(theta_2m)
        step_level_theta = 1.0
        theta_levels = np.arange(min_level_theta, max_level_theta + step_level_theta, step_level_theta)

        # Create wind levels and ticks
        max_level = np.round(np.max(wind_10m))
        min_level = np.round(np.min(wind_10m))
        step_level = 0.5
        wind_levels = np.arange(min_level, max_level + step_level, step_level)
        wind_ticks = wind_levels

        # Normalize u and v components for quiver plot if desired
        u_norm = u / np.sqrt(u**2 + v**2)
        v_norm = v / np.sqrt(u**2 + v**2)

        # Timestamp
        timestamp = to_np(mdbz.Time).astype('M8[s]').astype('O').isoformat()
        time = mdbz.Time
        time_str = str(time.values)

        # Get CRS
        crs = get_cartopy(wrfin=ds)

        # Create a 4 panel image with four colorbars
        fig = plt.figure(figsize=(30,15))
        #gs = gridspec.GridSpec(2, 4, width_ratios = [0.05, 1, 1, 0.05], height_ratios = [1, 1], hspace=0.10, wspace=0.20)
        #ax1 = plt.subplot(gs[0, 1], projection=crs)
        #ax2 = plt.subplot(gs[0, 2], projection=crs)
        #ax3 = plt.subplot(gs[1, 1], projection=crs)
        #ax4 = plt.subplot(gs[1, 2], projection=crs)
        #cax1 = plt.subplot(gs[0,0])
        #cax2 = plt.subplot(gs[0,3])
        #cax3 = plt.subplot(gs[1,0])
        #cax4 = plt.subplot(gs[1,3])
        ax1 = fig.add_subplot(2,2,1, projection=crs)
        ax2 = fig.add_subplot(2,2,2, projection=crs)
        ax3 = fig.add_subplot(2,2,3, projection=crs)
        ax4 = fig.add_subplot(2,2,4, projection=crs)

        ################################################################################################################################################
        # Top Left Plot - Reflectivity
        dbz_contour = ax1.contourf(wrf_lons, wrf_lats, mdbz, colors=dbz_colors, levels=dbz_contour_levels, zorder=3,
                                transform=ccrs.PlateCarree())

        #ax1.set_xlim(cartopy_xlim(mdbz))
        #ax1.set_ylim(cartopy_ylim(mdbz))
        ax1.set_extent(zoom_domain)
        plot_background(ax1)

        colorbar_dbz = fig.colorbar(dbz_contour, ax=ax1, orientation='vertical')
        colorbar_dbz.set_label('dBZ', fontsize=12)
        #cax1.yaxis.set_ticks_position('left')
        #cax1.yaxis.set_label_position('left')
        ax1.set_title('Composite Reflectivity (dBZ) '+ time_str[:19])

        ################################################################################################################################################
        # Top Right Plot - 2m Theta
        theta_contour = ax2.contour(wrf_lons, wrf_lats, theta_2m, levels=theta_levels, linewidths=0.9, colors='black', transform=ccrs.PlateCarree())
        theta_contourf = ax2.contourf(wrf_lons, wrf_lats, theta_2m, levels=theta_levels, transform=ccrs.PlateCarree(), cmap=get_cmap('coolwarm'))

        #ax2.set_xlim(cartopy_xlim(theta_2m))
        #ax2.set_ylim(cartopy_ylim(theta_2m))
        ax2.set_extent(zoom_domain)
        plot_background(ax2)

        colorbar_theta = fig.colorbar(theta_contourf, ax=ax2, orientation='vertical')
        colorbar_theta.set_label(r"$\theta$ (K)", fontsize=12)
        #cax2.yaxis.set_ticks_position('right')
        #cax2.yaxis.set_label_position('right')
        ax2.set_title('2-meter Potential Temperature (K) '+ time_str[:19])

        ################################################################################################################################################
        # Bottom Left Plot - Divergence
        #div_contour = ax3.contour(wrf_lons, wrf_lats, div, colors='black', transform=ccrs.PlateCarree())
        #div_contourf = ax3.contourf(wrf_lons, wrf_lats, div, cmap=get_cmap('jet'), transform=ccrs.PlateCarree())

        #ax3.set_xlim(cartopy_xlim(theta_2m))
        #ax3.set_ylim(cartopy_ylim(theta_2m))
        ax3.set_extent(zoom_domain)
        plot_background(ax3)

        #colorbar_div = fig.colorbar(div_contourf, ax=ax3, orientation='vertical')
        #colorbar_div.set_label('Divergence (1/s)', fontsize=12)
        #cax3.yaxis.set_ticks_position('left')
        #cax3.yaxis.set_label_position('left')
        ax3.set_title('Surface Divergence (1/s) ' + time_str[:19])

        ################################################################################################################################################
        # Bottom Right Plot - 10m Wind Speed and Direction
        wind_contour = ax4.contour(wrf_lons, wrf_lats, wind_10m, levels=wind_levels, colors='black', transform=ccrs.PlateCarree())
        wind_contourf = ax4.contourf(wrf_lons, wrf_lats, wind_10m, levels=wind_levels, cmap=get_cmap('rainbow'), transform=ccrs.PlateCarree())
        wind_quiver = ax4.quiver(x, y, u_norm, v_norm, pivot='middle', scale=25, width=0.005, color='black', regrid_shape=42.5, transform=ccrs.PlateCarree())

        #ax4.set_xlim(cartopy_xlim(theta_2m))
        #ax4.set_ylim(cartopy_ylim(theta_2m))
        ax4.set_extent(zoom_domain)
        plot_background(ax4)

        colorbar_wind = fig.colorbar(wind_contourf, ax=ax4, orientation='vertical')
        colorbar_wind.set_label('Wind Speed (m/s)', fontsize=12)
        #cax4.yaxis.set_ticks_position('right')
        #cax4.yaxis.set_label_position('right')
        ax4.set_title('Wind Speed and Direction ' + time_str[:19])
        
        plt.show()

Here is the plot that it produces repeating images of: Repeating image

Don't be concerned about the bottom left axes, or the white space on the left side of some of the images.

What can I do to fix this? Is there a better way to do this besides a nested for loop? Thank you for the help!


Solution

  • Two nested for loops will not do what you described as your goal. You said you wanted to "pair the first dataset with the first domain, the second dataset with the second domain, and so on" but with nested loops it will instead try the first dataset with every domain, and then try the second datsaset with every domain, and so on. If there were 10 datasets and 10 domains, this would result in 100 plots.

    It sounds like you want to use a single loop. Your for i in range(len(wrf_out)): loop looks fine. But, then you want to use i as the index into the list of domains, so that when i is 0, you would be using dataset 0 and domain 0, and then when i is 1, you would be using dataset 1 and domain 1, and so on.

    For example:

     for i in range(len(wrf_out)):
         domain = domain_2015_06_03[i]
         lonmin = domain[0]
         lonmax = domain[1]
         latmin = domain[2]
         latmax = domain[3]
         ...