run_model.py

#!/usr/bin/env python
"""Discovered events and visualizes each events.

It loops through DMSP files and using an OMNIweb file it finds events and writes
an output CSV file of events, and plots to disk.

This uses case files generated by make_case_file.py.
"""

import argparse
from dataclasses import dataclass
from datetime import datetime
import json
from matplotlib import MatplotlibDeprecationWarning
from matplotlib.colors import LogNorm
from matplotlib.dates import num2date
import pandas as pd
import numpy as np
from numpy.typing import NDArray
import os
import pylab as plt
import cdflib
from termcolor import cprint
from typing import Dict, List
import warnings

import lib_dasilva2022  # Single Dispersion
import lib_dasilva2024  # Double Dispersion
import lib_util


@dataclass
class DetectionResult:
    """Holds information associated with a detection."""
    start_time: datetime     # Start time associated with event
    end_time: datetime       # End time associated with event

    times: NDArray
    integrands: List[NDArray]
    energy_curves: List[NDArray]
    
    dmsp_flux_file: str      # Name of DMSP flux file
    dmsp_flux_fh: Dict       # DMSP data, as returned by lib_util
    Bx: float                # Magnetic field X component
    By: float                # Magnetic field Y component
    Bz: float                # Magnetic field Z component

    @classmethod
    def list_to_dataframe(cls, detection_results):
        """Convert a list of DetectionResult instances to DataFrame.

        Args
          detection_results: List of DetectionResult instances
        Returns
          df: Pandas dataframe
        """
        df_rows = []

        for detection_result in detection_results:
            df_rows.append(pd.Series(dict(
                start_time=detection_result.start_time,
                end_time=detection_result.end_time,
                dmsp_flux_file=detection_result.dmsp_flux_file,
                Bx=detection_result.Bx,
                By=detection_result.By,
                Bz=detection_result.Bz,                
            )))

        return pd.DataFrame(df_rows)
    

def main():
    """Main routine of the program. Run with --help for description of arguments."""
    # Parse command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', metavar='CASE_FILE', required=True,
                        help='Path to case file')
    parser.add_argument('--threshold', type=float, default=-1,
                        help='Detection threshold')
    parser.add_argument('-D', '--double-dispersion', action='store_true',
                        help='Search double dispersion')    
    parser.add_argument('--no-plot', action='store_true',
                        help='Set to disable plotting')
    parser.add_argument('--simple-plots', action='store_true',
                        help='Do simple plots, without no algorithm information.')
    
    args = parser.parse_args()
    cprint(f'Parameters: {args}', 'green')
    
    # Load case file ---------------------------------------------------------
    with open(args.i) as fh:
        case_file = json.load(fh)

    # Read OMNIWeb data all at once ------------------------------------------
    omniweb_fh = lib_util.read_omniweb_files(
        case_file['OMNIWEB_FILES']
    )
        
    # Loop through files and call detection routine
    # ------------------------------------------------------------------------   
    detection_results = []
    
    for i, dmsp_flux_file in enumerate(sorted(case_file['DMSP_FLUX_FILES'])):
        cprint(f'Processing {i+1}/{len(case_file["DMSP_FLUX_FILES"])} :: '
               f'{dmsp_flux_file}', 'green')

        if args.double_dispersion:
            cur_detection_results = search_double_dispersion(
                dmsp_flux_file=dmsp_flux_file,
                omniweb_fh=omniweb_fh,
                reverse_effect=case_file['REVERSE_EFFECT'],
                inverse_effect=case_file['INVERSE_EFFECT'],
                integral_threshold=args.threshold,
            )        
        else:
            cur_detection_results = search_single_dispersion(
                dmsp_flux_file=dmsp_flux_file,
                omniweb_fh=omniweb_fh,
                reverse_effect=case_file['REVERSE_EFFECT'],
                inverse_effect=case_file['INVERSE_EFFECT'],
                integral_threshold=args.threshold,
            )
        
        if cur_detection_results:
            df = DetectionResult.list_to_dataframe(cur_detection_results)
            print(df.to_string(index=0))

        detection_results.extend(cur_detection_results)


    # Sort by time
    detection_results = sorted(
        detection_results,
        key=lambda res: res.start_time
    )

    # Do plotting ------------------------------------------------------------
    if not args.no_plot:
        for detection_result in detection_results:
            write_plot(
                detection_result, args.simple_plots, case_file['PLOT_OUTPUT']
            )
    
    # Write event list to console and output file ----------------------------
    df = DetectionResult.list_to_dataframe(detection_results)

    cprint('Discovered events:', 'green')
    print(df.to_string(index=0))

    cprint('Writing event output (' + str(len(df.index)) + ' events) to '
           + case_file['EVENT_OUTPUT'], 'green')
    
    df.to_csv(case_file['EVENT_OUTPUT'], index=0)

    
def search_single_dispersion(
    dmsp_flux_file, omniweb_fh, reverse_effect, inverse_effect,
    integral_threshold
):
    """Search for single dispersion events in a DMSP file.
    
    Args
      dmsp_flux_file: Path to HDF5 DMSP file holding spectrogram data (daily)
      omniweb_fh: Loaded omniweb data in a dictionary
      reverse_effect: Search for effects in the opposite direction with a
        magnetic field set to the opposite of the coded threshold.
      integration: detection threhsold value, set to -1 for default.
    Returns
      detection_results: List of DetectionResult instances describing
        events found and data necessary for plotting.
    """
    if integral_threshold < 0:
        integral_threshold = lib_dasilva2022.DEFAULT_INTEGRAL_THRESHOLD
    
    # Do computation --------------------------------------------------
    try:
        dmsp_flux_fh = lib_util.read_dmsp_flux_file(dmsp_flux_file)
    except (FileNotFoundError, OSError):
        return []

    dEicdt_smooth, Eic_smooth, Eic = (
        lib_dasilva2022.estimate_log_Eic_smooth_derivative(dmsp_flux_fh)
    )
    
    df_walk, integrand, _, _ = lib_dasilva2022.walk_and_integrate(
        dmsp_flux_fh, omniweb_fh, dEicdt_smooth, Eic_smooth,
        lib_dasilva2022.INTERVAL_LENGTH, integral_threshold,
        reverse_effect=reverse_effect, inverse_effect=inverse_effect,
        return_integrand=True
    )

    # Convert to list of detection result instances --------------------
    detection_results = []

    for _, row in df_walk.iterrows():        
        # Clean up Eic curve for plot
        i = dmsp_flux_fh['t'].searchsorted(row['start_time'])
        j = dmsp_flux_fh['t'].searchsorted(row['end_time'])

        delta_e = np.log10(dmsp_flux_fh['ch_energy'][-1]) - \
            np.log10(dmsp_flux_fh['ch_energy'][-2])
        energy_curve = Eic[i:j].copy()
        energy_curve += 0.5 * delta_e
        top_e = np.log10(lib_dasilva2022.MAX_SHEATH_ENERGY)
        energy_curve[energy_curve > top_e] = np.nan
        
        detection_results.append(DetectionResult(
            start_time=row['start_time'],
            end_time=row['end_time'],
            times=[dmsp_flux_fh['t'][i:j]],
            integrands=[integrand[i:j]],
            energy_curves=[energy_curve],
            dmsp_flux_file=dmsp_flux_file,
            dmsp_flux_fh=dmsp_flux_fh,
            Bx=row['Bx_mean'],
            By=row['By_mean'],
            Bz=row['Bz_mean'],
        ))
    
    return detection_results



def search_double_dispersion(
    dmsp_flux_file, omniweb_fh, reverse_effect, inverse_effect,
    integral_threshold
):
    """Search for double dispersion events in a DMSP file.
    
    Args
      dmsp_flux_file: Path to HDF5 DMSP file holding spectrogram data (daily)
      omniweb_fh: Loaded omniweb data in a dictionary
      reverse_effect: Search for effects in the opposite direction with a
        magnetic field set to the opposite of the coded threshold.
      integration: detection threhsold value, set to -1 for default.
    Returns
      detection_results: List of DetectionResult instances describing
        events found and data necessary for plotting.
    """
    if integral_threshold < 0:
        integral_threshold = lib_dasilva2024.DEFAULT_INTEGRAL_THRESHOLD
    
    # Do computation --------------------------------------------------
    try:
        dmsp_flux_fh = lib_util.read_dmsp_flux_file(dmsp_flux_file)
    except (FileNotFoundError, OSError):
        return []

    df_walk = lib_dasilva2024.walk_and_integrate(
            dmsp_flux_fh, omniweb_fh, reverse_effect, integral_threshold,
    )

    # Convert to list of detection result instances --------------------
    detection_results = []

    for _, row in df_walk.iterrows():
        detection_results.append(DetectionResult(
            start_time=row['start_time'],
            end_time=row['end_time'],
            times=[row['t'], row['t']],
            integrands=[row['lower_integrand'], row['upper_integrand']],
            energy_curves=[row['lower_Ep'], row['upper_Ep']],
            dmsp_flux_file=dmsp_flux_file,
            dmsp_flux_fh=dmsp_flux_fh,
            Bx=row['Bx_mean'],
            By=row['By_mean'],
            Bz=row['Bz_mean'],
        ))
    
    return detection_results


def write_plot(
    detection_result, simple_plots, plot_out_dir,
):
    """Visualizes detection and writes plot to disk.

    Args
      detection_result: Instance of DetectionResult describing the detection
        associated data for plotting
      simple_plots: Set to true to enable simpler version of plots
      plot_out_dir: Directory to write plot in
    """
    dmsp_flux_file = detection_result.dmsp_flux_file
    dmsp_flux_fh = detection_result.dmsp_flux_fh
    start_time = detection_result.start_time
    end_time = detection_result.end_time
    integrands = detection_result.integrands
    energy_curves = detection_result.energy_curves
    Bx = detection_result.Bx
    By = detection_result.By
    Bz = detection_result.Bz
        
    orig_i = dmsp_flux_fh['t'].searchsorted(start_time)
    orig_j = dmsp_flux_fh['t'].searchsorted(end_time)

    delta_index = int(1.2 * (orig_j - orig_i))  # make plot wider
    delta_index = int(2 * (orig_j - orig_i))  # make plot wider,dd,rd
    i = max(orig_i - delta_index, 0)
    j = min(orig_j + delta_index, dmsp_flux_fh['t'].size - 1)
    
    if simple_plots:
        fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=True)
    else:
        small_size, med_size, big_size = 16, 18, 20
        plt.rc('font', size=small_size)          # controls default text sizes
        plt.rc('axes', titlesize=small_size)     # fontsize of the axes title
        plt.rc('axes', labelsize=med_size)       # fontsize of the x and y labels
        plt.rc('xtick', labelsize=small_size)    # fontsize of the tick labels
        plt.rc('ytick', labelsize=small_size)    # fontsize of the tick labels
        plt.rc('legend', fontsize=small_size)    # legend fontsize
        plt.rc('figure', titlesize=big_size)     # fontsize of the figure title
        fig, axes = plt.subplots(3, 1, figsize=(18, 9), sharex=True, dpi=600)
        
    # Plot title
    nonzero = (integrands[0][orig_i:orig_j] > 0.01).nonzero()[0]
    if nonzero.size > 0:
        first_pos_i = nonzero[0] - 1
        orig_i += first_pos_i

    sat = None
    for sat_num in range(10, 100):
        if f'F{sat_num}' in dmsp_flux_file:
            sat = f'F{sat_num}'
                        
    if simple_plots:
        title = (
            f'DMSP {sat} Plasma Spectrograms'
        )
    else:            
        title = (
            f'DMSP {sat} Dispersion Event during '
            "$B_{IMF}$" + f' = ({Bx:.2f}, {By:.2f}, {Bz:.2f}) nT\n'
            f"{start_time.isoformat()} - "
            f"{end_time.isoformat()}"
        )
            
    axes[0].set_title(title)
        
    # Ion spectrogram
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", MatplotlibDeprecationWarning)
        im = axes[0].pcolor(
            dmsp_flux_fh['t'][i:j],
            np.log10(dmsp_flux_fh['ch_energy']),
            dmsp_flux_fh['ion_d_ener'][:, i:j], 
            norm=LogNorm(vmin=1e3, vmax=1e8), cmap='jet'
        )
        plt.colorbar(im, ax=axes[0]).set_label('Energy Flux')
        
    if not simple_plots:
        pairs = zip(detection_result.times, energy_curves)
        
        for color_idx, (cur_t, energy_curve) in enumerate(pairs):
            color = 'bg'[color_idx]
            axes[0].plot(cur_t, energy_curve, f'*{color}-')
            
        axes[0].axhline(
            np.log10(lib_dasilva2022.MAX_EIC_ENERGY),
            color='black', linestyle='dashed'
        )
    
    axes[0].set_ylabel('Ions Energy [eV]')
    adjust_axis_energy_yticks(axes[0])
        
    # Electron Spectrogram
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", MatplotlibDeprecationWarning)
        im = axes[1].pcolor(
            dmsp_flux_fh['t'][i:j],
            np.log10(dmsp_flux_fh['ch_energy']),
            dmsp_flux_fh['el_d_ener'][:, i:j], 
            norm=LogNorm(vmin=1e5, vmax=1e10), cmap='jet'
        )
    plt.colorbar(im, ax=axes[1]).set_label('Energy Flux')
    axes[1].set_ylabel('Electrons Energy [eV]')
    adjust_axis_energy_yticks(axes[1])
        
    # Scoring function
    if not simple_plots:
        score_range = [-.25, .25]

        if len(integrands) == 1:
            axes[2].fill_between(detection_result.times[0], 0, integrands[0])
        else:
            pairs = zip(detection_result.times, integrands)
            for color_idx, (cur_t, integrand) in enumerate(pairs):
                color = 'bg'[color_idx]
                axes[2].plot(cur_t, integrand, color=color)
    
        axes[2].axhline(0, color='black', linestyle='dashed')            
        axes[2].set_ylabel('D(t) [Log(eV)/s]')
        axes[2].set_ylim(score_range)
        plt.colorbar(im, ax=axes[2]).set_label('')
        
    add_multirow_xticks(axes[-1], dmsp_flux_fh, simple_plots)
        
    # Plot spacings
    plt.subplots_adjust(hspace=.05)
    plt.tight_layout()
    
    # Save image
    out_name = plot_out_dir + '/'
    out_name += f'{os.path.basename(dmsp_flux_file)}_'
    out_name += f"{start_time.isoformat().replace(':', '')}_"
    out_name += f"{end_time.isoformat().replace(':', '')}.png"
    
    os.makedirs(plot_out_dir, exist_ok=True)
    plt.savefig(out_name)
    plt.close()
    
    cprint('Wrote plot ' + out_name, 'green')         


def add_multirow_xticks(ax, dmsp_flux_fh, simple_plots):
    """Add multirow tickmarks to the bottom axis as is common in the
    magnetospheres community.

    Args
      ax: matplotlib axes
      dmsp_flux_fh: file handle (as returned by read_dmsp_flux_file)
    """
    xticks = ax.get_xticks()
    new_labels = []
    
    for time_float in xticks:
        time = num2date(time_float)
        i = dmsp_flux_fh['t'].searchsorted(time)
        if i == dmsp_flux_fh['t'].size:
            continue
        mlat = dmsp_flux_fh['mlat'][i]
        mlt = dmsp_flux_fh['mlt'][i]
        
        new_label = '%s\n%.1f\n%.1f' % (
            time.strftime('%H:%M:%S'), mlat, mlt
        )
        new_labels.append(new_label)
        
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", UserWarning)
        ax.set_xticklabels(new_labels)
        
        if simple_plots:
            ax.text(-0.075, -0.075, 'Time', transform=ax.transAxes)
            ax.text(-0.075, -0.15, 'MLAT', transform=ax.transAxes)
            ax.text(-0.075, -0.225, 'MLT', transform=ax.transAxes)
        else:            
            ax.text(-0.09, -0.12, 'Time', transform=ax.transAxes)
            ax.text(-0.09, -0.24, 'MLAT', transform=ax.transAxes)
            ax.text(-0.09, -0.36, 'MLT', transform=ax.transAxes)


def adjust_axis_energy_yticks(ax):
    """Adjust yticklabels for axes with y-axis being energy.
    
    Sets them to terms like eV and keV.

    Args
      ax: matplotlib axes    
    """
    yticks = 10**ax.get_yticks()
    labels = []

    for ytick in yticks:
        if ytick < 1000:
            labels.append('%d eV' % ytick)
        else:
            labels.append('%.1f keV' % (ytick/1000))

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", UserWarning)
        
        ax.set_yticklabels(labels)


if __name__ == '__main__':
    main()