flash_sort_and_grid.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d394e12",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python\n",
    "\"\"\"\n",
    "Example script showing how to use lmatools to sort LMA ASCII data into flashes\n",
    "and how to create gridded imagery from those flashes.\n",
    "\n",
    "The block of code at the bottom shows how to call sorting and gridding functions.\n",
    "The params dictionary controls the flash sorting parameters.\n",
    "\n",
    "USAGE: \n",
    "python flash_sort_and_grid.py /path/to/output/ lmatools/sampledata/ASCII_solutions/LYLOUT_140526_*.dat.gz\n",
    "\n",
    "Directories for the HDF5 files, grids, and plots are created within the \n",
    "directory indicated by /path/to/output/ \n",
    "\"\"\"\n",
    "\n",
    "import sys, os, glob\n",
    "from datetime import datetime, timedelta\n",
    "import subprocess\n",
    "\n",
    "from lmatools.flashsort.gen_autorun import logger_setup, sort_files\n",
    "from lmatools.flashsort.gen_sklearn import DBSCANFlashSorter\n",
    "\n",
    "from lmatools.grid.make_grids import grid_h5flashfiles, dlonlat_at_grid_center, write_cf_netcdf_latlon, write_cf_netcdf_3d_latlon\n",
    "from lmatools.vis.multiples_nc import make_plot, make_plot_3d, read_file_3d\n",
    "\n",
    "\n",
    "def tfromfile(name):\n",
    "    parts = name.split('_')\n",
    "    y, m, d = list(map(int, (parts[-3][0:2], parts[-3][2:4], parts[-3][4:6])))\n",
    "    H, M, S = list(map(int, (parts[-2][0:2], parts[-2][2:4], parts[-2][4:6])))\n",
    "    return y+2000,m,d,H,M,S    \n",
    "\n",
    "def sort_flashes(files, base_sort_dir, params):\n",
    "    \"\"\" Given a list of LMA ASCII data files, created HDF5 flash-sorted data \n",
    "        files in base_sort_dir/h5_files.\n",
    "    \n",
    "        params is a dictionary with the following format\n",
    "        params = {'stations':(6,99), # range of allowable numbers of contributing stations\n",
    "                  'chi2':(0,1.0),    # range of allowable chi-sq values\n",
    "                  'distance':3000.0, 'thresh_critical_time':0.15, # space and time grouping thresholds\n",
    "                  'thresh_duration':3.0, # maximum expected flash duration\n",
    "                  'mask_length':6, # length of the hexadecimal station mask column in HDF5 \n",
    "                  }\n",
    "        \n",
    "    \"\"\"\n",
    "    # base_sort_dir = outpath\n",
    "    logger_setup(base_sort_dir)\n",
    "    # files = lmatools.testing.test_gen_autorun_DBSCAN.get_sample_data_list()\n",
    "    h5_dir = os.path.join(base_sort_dir, 'h5_files') \n",
    "    \n",
    "    y,m,d,H,M,S = tfromfile(files[0])\n",
    "    date = datetime(y,m,d, H,M,S)\n",
    "\n",
    "    # Create HDF5 flash files\n",
    "    base_out_dir = (h5_dir+\"/20%s\" %(date.strftime('%y/%b/%d')))\n",
    "    if os.path.exists(base_out_dir) == False:\n",
    "        os.makedirs(base_out_dir)\n",
    "        subprocess.call(['chmod', 'a+w', base_out_dir, h5_dir+'/20%s' %(date.strftime('%y/%b')), h5_dir+'/20%s' %(date.strftime('%y'))])\n",
    "\n",
    "    tag = ''\n",
    "    outdir = os.path.join(base_out_dir, tag) \n",
    "    info = open(os.path.join(outdir, 'input_params.py'), 'w')\n",
    "    info.write(str(params))\n",
    "    info.close()\n",
    "    \n",
    "    if True:\n",
    "        cluster = DBSCANFlashSorter(params).cluster\n",
    "        sort_files(files, outdir, cluster)\n",
    "    # Figure out which HDF5 files were created\n",
    "    h5_filenames = glob.glob(h5_dir+'/20%s/LYLOUT*.dat.flash.h5' %(date.strftime('%y/%b/%d')))\n",
    "    h5_filenames.sort()\n",
    "    return h5_filenames\n",
    "    \n",
    "def grid_and_plot(h5_filenames, base_sort_dir, dx=1.0e3, dy=1.0e3, dz=1.0e3, frame_interval=60.0,\n",
    "                  x_bnd=(-200.0e3, 200.0e3), y_bnd=(-200.0e3, 200.0e3), z_bnd=(0.0e3, 20.0e3),\n",
    "                  ctr_lat=33.5, ctr_lon=-101.5, center_ID='WTLMA',\n",
    "                  n_cols=2, base_date=None\n",
    "                  ):\n",
    "    \"\"\" Given a list of HDF5 filenames (sorted by time order) in h5_filenames, \n",
    "        create 2D and 3D NetCDF grids with spacing dx, dy, dz in meters, \n",
    "        frame_interval in seconds, and tuples of grid edges \n",
    "        x_bnd, y_bnd, and z_bnd in meters\n",
    "                  \n",
    "        The actual grids are in regular lat,lon coordinates, with spacing at the \n",
    "        grid center matched to the dx, dy values given.\n",
    "                                    \n",
    "        n_cols controls how many columns are plotted on each page.\n",
    "                  \n",
    "        Grids and plots are written to base_sort_dir/grid_files/ and  base_sort_dir/plots/\n",
    "\t\t\t\t  \n",
    "\t\tbase_date is used to optionally set a common reference time for each of the NetCDF grids.\n",
    "    \"\"\"              \n",
    "    # not really in km, just a different name to distinguish from similar variables below.\n",
    "    dx_km=dx\n",
    "    dy_km=dy\n",
    "    x_bnd_km = x_bnd\n",
    "    y_bnd_km = y_bnd\n",
    "    z_bnd_km = z_bnd\n",
    "\n",
    "    grid_dir = os.path.join(base_sort_dir, 'grid_files')\n",
    "    plot_dir = os.path.join(base_sort_dir, 'plots')\n",
    "\n",
    "    # There are similar functions in lmatools to grid on a regular x,y grid in some map projection.\n",
    "    dx, dy, x_bnd, y_bnd = dlonlat_at_grid_center(ctr_lat, ctr_lon, \n",
    "                                dx=dx_km, dy=dy_km,\n",
    "                                x_bnd = x_bnd_km, y_bnd = y_bnd_km )\n",
    "    # print(\"dx, dy = {0}, {1} deg\".format(dx,dy))\n",
    "    # print(\"lon_range = {0} deg\".format(x_bnd))\n",
    "    # print(\"lat_range = {0} deg\".format(y_bnd))\n",
    "    \n",
    "    for f in h5_filenames:\n",
    "        y,m,d,H,M,S = tfromfile(f)\n",
    "        # print y,m,d,H,M,S\n",
    "        start_time = datetime(y,m,d, H,M,S)\n",
    "        end_time   = start_time + timedelta(0,600)\n",
    "        date = start_time\n",
    "        # print start_time, end_time\n",
    "    \n",
    "        outpath = grid_dir+'/20%s' %(date.strftime('%y/%b/%d'))\n",
    "        if os.path.exists(outpath) == False:\n",
    "            os.makedirs(outpath)\n",
    "            subprocess.call(['chmod', 'a+w', outpath, grid_dir+'/20%s' %(date.strftime('%y/%b')), grid_dir+'/20%s' %(date.strftime('%y'))])\n",
    "        if True:\n",
    "            grid_h5flashfiles(h5_filenames, start_time, end_time, frame_interval=frame_interval, proj_name='latlong',\n",
    "                    base_date = base_date, energy_grids=True,\n",
    "\t\t\t\t\tdx=dx, dy=dy, x_bnd=x_bnd, y_bnd=y_bnd, z_bnd=z_bnd_km,\n",
    "                    ctr_lon=ctr_lon, ctr_lat=ctr_lat, outpath = outpath,\n",
    "                    output_writer = write_cf_netcdf_latlon, output_writer_3d = write_cf_netcdf_3d_latlon,\n",
    "                    output_filename_prefix=center_ID, spatial_scale_factor=1.0\n",
    "                    )\n",
    "        \n",
    "    # Create plots\n",
    "    mapping = { 'source':'lma_source',\n",
    "                'flash_extent':'flash_extent',\n",
    "                'flash_init':'flash_initiation',\n",
    "                'footprint':'flash_footprint',\n",
    "                'specific_energy':'specific_energy',\n",
    "                'flashsize_std':'flashsize_std',\n",
    "                'total_energy': 'total_energy'\n",
    "               }\n",
    "    \n",
    "    nc_names = glob.glob(grid_dir+'/20%s/*.nc' %(date.strftime('%y/%b/%d')))\n",
    "    nc_names_3d = glob.glob(grid_dir+'/20%s/*_3d.nc' %(date.strftime('%y/%b/%d')))\n",
    "    nc_names_2d = list(set(nc_names) - set(nc_names_3d))\n",
    "    nc_names_2d.sort()\n",
    "    nc_names_3d.sort()\n",
    "    outpath = plot_dir+'/20%s' %(date.strftime('%y/%b/%d'))\n",
    "    if os.path.exists(outpath) == False:\n",
    "        os.makedirs(outpath)\n",
    "        subprocess.call(['chmod', 'a+w', outpath, plot_dir+'/20%s' %(date.strftime('%y/%b')), plot_dir+'/20%s' %(date.strftime('%y'))])\n",
    "\n",
    "    for f in nc_names_2d:\n",
    "        gridtype = f.split('dx_')[-1].replace('.nc', '')\n",
    "        var = mapping[gridtype]\n",
    "        make_plot(f, var, n_cols=n_cols, x_name='longitude', y_name='latitude', outpath = outpath)\n",
    "    \n",
    "    for f in nc_names_3d:\n",
    "        gridtype = f.split('dx_')[-1].replace('.nc', '').replace('_3d', '')\n",
    "        var = mapping[gridtype]\n",
    "        # grid_range = range_mapping[gridtype]\n",
    "        ###Read grid files, then plot in either 2d or 3d space###\n",
    "        grid, grid_name, x, y, z, t, grid_t_idx, grid_x_idx, grid_z_idx = read_file_3d(f, var, x_name='longitude', y_name='latitude', z_name='altitude')\n",
    "        make_plot_3d(grid, grid_name, x, y, z, t,\n",
    "                     grid_t_idx, grid_x_idx, grid_z_idx,\n",
    "                     n_cols = n_cols, outpath = outpath)\n",
    "        #, grid_range=grid_range)\n",
    "        \n",
    "    return nc_names_2d, nc_names_3d\n",
    "        \n",
    "        \n",
    "if __name__ == '__main__':\n",
    "    params = {'stations':(6,99), # range of allowable numbers of contributing stations\n",
    "              'chi2':(0,1.0),    # range of allowable chi-sq values\n",
    "              'distance':3000.0, 'thresh_critical_time':0.15, # space and time grouping thresholds\n",
    "              'thresh_duration':3.0, # maximum expected flash duration\n",
    "              'ctr_lat':33.5, 'ctr_lon':-101.5, #center lat/lon to use for flash sorting, gridding\n",
    "              # 'mask_length':6, # length of the hexadecimal station mask column in the LMA ASCII files\n",
    "              }\n",
    "    center_ID='WTLMA'\n",
    "    \n",
    "    data_out = sys.argv[1]\n",
    "    filenames = sys.argv[2:]\n",
    "    h5_filenames = sort_flashes(filenames, data_out, params) \n",
    "    \n",
    "    # other keyword arguments control the grid spacing ... see the function definition\n",
    "    nc_names_2d, nc_names_3d = grid_and_plot(h5_filenames, data_out, base_date=datetime(2012, 1, 1),\n",
    "        ctr_lat=params['ctr_lat'], ctr_lon=params['ctr_lon'], center_ID=center_ID)\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}