Lightning Flash Clustering: process lma data

process_lma_data.py

from clustering_algorithm.cluster_flashes import load_dat, to_ecef, cluster, get_flash_params, merge_flashes
import numpy as np
import os
from datetime import date, timedelta
import json
from pathlib import Path

all_files = os.listdir(data_dir)

for f_idx, fname in enumerate(all_files):
  file_stem = str(Path(fname)).rstrip(''.join(Path(fname).suffixes))#Path(fname).stem
  print(f"processing {file_stem} ({str(f_idx + 1)} / {str(len(all_files))})")

  # load data
  # ...

  flashes, n_removed = cluster(sources, xyz_scale, t_scale, n_grid_points, min_samples=min_samples, epsilon=epsilon, max_duration=max_duration)

  if len(flashes) == 0:
      continue

  flashes_merged = merge_flashes(flashes, t_threshold=merge_t_threshold, xyz_threshold=merge_xyz_threshold)
  merge_count = len(flashes) - len(flashes_merged)

  if save_flash_sources:
      flashes_merged_npy = np.zeros((sum([elem.shape[0] for elem in flashes_merged]), 1+flashes_merged[0].shape[1]))
      counter = 0
      for flash_idx, flash in enumerate(flashes_merged):
          flashes_merged_npy[counter:counter+flash.shape[0],:-1] = flash
          flashes_merged_npy[counter:counter+flash.shape[0],-1] = flash_idx
          counter += flash.shape[0]

      with open(os.path.join(out_dir, f'{file_stem}_sources.npy'), 'wb') as f:
          np.save(f, flashes_merged_npy)

  flash_params = get_flash_params(flashes_merged, start_time)

  with open(os.path.join(out_dir, f'{file_stem}.json'), 'w') as f:
      json.dump({'data_start_time': start_time.strftime("%m/%d/%y %H:%M:%S"), 'n_removed': n_removed, 'merge_count': merge_count, 'flash_params': flash_params}, f)