Lightning Flash Clustering: plot bucket statistics

plot_bucket_stats.py

import matplotlib.pyplot as plt

n_buckets = int(max_dist // bucket_size)
param_names = ('duration', 'hull_area', 'mean_power', 'n_sources')
param_units = ('(ms)', '(square km)', '(dBW)', '')
stats_buckets = [{param: [] for param in param_names} for i in range(n_buckets)]

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']
    total_flash_count += len(flashes)

    for flash in flashes:
        bucket_idx = int(flash['dist_from_center'] // bucket_size)
        grid_x_vals, grid_y_vals = [pt[0] for pt in flash['grid_points']], [pt[1] for pt in flash['grid_points']]

        if not ( (grid_min_x <= min(grid_x_vals) and grid_max_x > max(grid_x_vals)) and (grid_min_y <= min(grid_y_vals) and grid_max_y > max(grid_y_vals)) ):
            continue

        if bucket_idx > len(stats_buckets):
            continue
        for param in param_names:
            stats_buckets[bucket_idx][param].append(flash[param])

percentiles = (5, 50, 95)
all_flash_durations = []
for elem in stats_buckets:
    all_flash_durations += elem['duration']
all_flash_durations = np.array(all_flash_durations)

for p in percentiles:
    d = np.percentile(all_flash_durations, p) * 1000
    print(f'percentile {p} of flash durations (milliseconds): {round(d)}')

plt.rc('xtick', labelsize=8)
fig, axs = plt.subplots(2, 2)

for param_idx, param in enumerate(param_names):
    scale_val = 1e3 if param == 'duration' else 1.0 # convert duration to ms
    param_vals = [np.array(elem[param]) * scale_val for elem in stats_buckets]

    if param != 'mean_power':
        axs[param_idx//2, param_idx % 2].set_yscale('log')

    axs[param_idx//2, param_idx % 2].boxplot(param_vals, tick_labels=[str(i) for i in range(bucket_size, bucket_size*(len(param_vals)+1), bucket_size)], showfliers=False)
    axs[param_idx//2, param_idx % 2].set(xlabel='distance from LMA center (km)', ylabel=f'{param} {param_units[param_idx]}')
    axs[param_idx//2, param_idx % 2].set_title(param)

plt.tight_layout()
plt.savefig(os.path.join(plot_dir, 'flash_params.png'))
plt.close()