Lightning Flash Clustering: plot grid statistics

plot_grid_stats.py

import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
import datetime

all_flash_areas = {}
all_flash_hours = {}

for fname in os.listdir(json_dir):
    with open(os.path.join(json_dir, fname), 'r') as f:
        data = json.load(f)

    flashes = data['flash_params']

    for flash in flashes:

        for grid_point in flash['grid_points']:
            if tuple(grid_point) not in all_flash_areas:
                all_flash_areas[tuple(grid_point)] = []
            all_flash_areas[tuple(grid_point)].append(flash['hull_area'])

            if tuple(grid_point) not in all_flash_hours:
                all_flash_hours[tuple(grid_point)] = []

            dt = datetime.datetime.strptime(flash['init_time'], "%m/%d/%y %H:%M:%S")
            all_flash_hours[tuple(grid_point)].append(dt.strftime("%m/%d/%y %H"))
            
area_grid = np.zeros((grid_size, grid_size))
lh_grid = np.zeros((grid_size, grid_size))
fph_grid = np.zeros((grid_size, grid_size))

for k, v in all_flash_areas.items():
    flashes_sorted = sorted(v)
    median = flashes_sorted[len(flashes_sorted) // 2]
    area_grid[k[0],k[1]] = median

for k, v in all_flash_hours.items():
    n_hours = len(list(set(v)))
    lh_grid[k[0],k[1]] = n_hours
    cntr = Counter(v)
    fph_grid[k[0],k[1]] = np.mean(np.array(list(cntr.values())))

grids = (area_grid, lh_grid, fph_grid)

fig, axs = plt.subplots(1, 3, figsize=(9, 3), subplot_kw={'projection': ccrs.PlateCarree()})

for param_idx, param in enumerate(('flash_area', 'lightning_hours', 'flashes_per_lightning_hour')):
    grid = grids[param_idx]
    vmax = 100.0 if param == 'flash_area' else np.max(grid)

    axs[param_idx].add_feature(cfeature.STATES, linewidth=1.4, edgecolor='white')
    plot_extent = [*grid_longitude_range, *grid_latitude_range]
    axs[param_idx].set_extent(plot_extent,crs=ccrs.PlateCarree())
    axs[param_idx].imshow(grid, origin='upper', cmap='jet', extent=plot_extent, vmax=vmax, transform=ccrs.PlateCarree())

    for landmark_idx in range(len(landmark_names)):
        landmark_coords = landmark_coordinates[landmark_idx]

        axs[param_idx].plot(landmark_coords[1], landmark_coords[0], 'wo', markersize=7, transform=ccrs.PlateCarree())
        axs[param_idx].text(landmark_coords[1]-0.001, landmark_coords[0]+0.2, landmark_names[landmark_idx], color='black', size=14, path_effects=[pe.withStroke(linewidth=2, foreground="white")], transform=ccrs.PlateCarree())
    axs[param_idx].set_title(param)

plt.savefig(os.path.join(plot_dir, 'grid_stats.png'))