# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# # Data Exploration & Preparation
#
# This notebook is used to explore the in-situ data for the entire list of STEREO A and B ICMEs.
#
# Parameters of interest:
# - Magnetic field components
# - Magnetic field strength
# - Proton number density
# - Proton speed
# - Proton temperature

# +
from collections import defaultdict
import datetime as dt
import json

import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.units as munits
from matplotlib.backends.backend_pdf import PdfPages
import numpy as np
import pandas as pd
from pyts.image import RecurrencePlot

from sklearn.preprocessing import StandardScaler
from tqdm.notebook import tqdm

converter = mdates.ConciseDateConverter()
munits.registry[np.datetime64] = converter
munits.registry[dt.date] = converter
munits.registry[dt.datetime] = converter
# -

# ### Parse the full helcats ICME list and extract all of the stereo A and B ICMEs

# +
with open('../ICME_WP4_V10.json', 'r') as fobj:
    json_data = json.load(fobj)
    
df = pd.DataFrame(json_data['data'], columns=json_data['columns'])

sta_icme_df = df[df['SC_INSITU'].str.contains('STEREO-A')]
stb_icme_df = df[df['SC_INSITU'].str.contains('STEREO-B')]

sta_icme_df.index = pd.DatetimeIndex(sta_icme_df.loc[:,'ICME_START_TIME']).tz_localize(None)
stb_icme_df.index = pd.DatetimeIndex(stb_icme_df.loc[:,'ICME_START_TIME']).tz_localize(None)
# -

sta_icme_df.to_csv('../data/sta_icme_list.txt', header=True, index=True)
stb_icme_df.to_csv('../data/stb_icme_list.txt', header=True, index=True)


def read_stereo_datasets(fname):
    """Function for reading in stereo datasets"""
    with open(fname, 'r') as fobj:
        lines = fobj.readlines()

    colnames = lines[0].split()
    tmp = lines[1].split()
    units = []
    units.append(' '.join(tmp[:2]))
    units += tmp[2:]

    for col, unit in zip(colnames, units):
        print(col, unit)
        
    data = []
    index = []
    for line in tqdm(lines[2:]):
        lsplit = line.split()
        index.append(dt.datetime.strptime(' '.join(lsplit[:2]), '%d-%m-%Y %H:%M:%S.%f'))
        data.append(list(map(float, lsplit[2:])))
    
    df = pd.DataFrame(data, columns=colnames[1:], index=pd.DatetimeIndex(index))
    return df


# ### STEREO A dataset

sta_icme_df.info()

sta_data_df = read_stereo_datasets('../data/sta_l2_magplasma.txt')

sta_data_df.index[0], sta_data_df.index[-1]

sta_data_df[sta_data_df['BTOTAL'].gt(-1e30)].sort_index().rolling('20D', center=True).mean().plot(y='BTOTAL')

for col in sta_data_df.columns:
    print(col)

cols_of_interest = [
    'BTOTAL', 
    'BX(RTN)',
    'BY(RTN)',
    'BZ(RTN)', 
    'VP_RTN',
    'NP',
    'TEMPERATURE',
    'BETA'
]

sta_data_cut_df = sta_data_df[cols_of_interest]

sta_data_cut_df = sta_data_cut_df[sta_data_cut_df.gt(-1e30)].dropna().sort_index()

# Remove all rows where the number density, temperature, or beta values are negative since they are unphysical

sta_data_cut_df = sta_data_cut_df[~sta_data_cut_df['NP'].lt(0)]
sta_data_cut_df = sta_data_cut_df[~sta_data_cut_df['TEMPERATURE'].lt(0)]
sta_data_cut_df = sta_data_cut_df[~sta_data_cut_df['BETA'].lt(0)]

sta_data_cut_df.describe()

(~sta_data_cut_df['NP'].lt(0)).sum()

sta_data_cut_df.to_csv("../data/sta_dataset_cleaned.txt", header=True, index=True)

sta_data_cut_df.info()

sta_data_cut_df.head()

cols_of_interest


def quality_check_plot(stereo_df, icme_date, window_size=dt.timedelta(days=5), cols=[], normalize=True):
    fig, axes = plt.subplots(nrows=len(cols), ncols=1, figsize=(5,10), gridspec_kw={'hspace':0.1}, sharex=True)
    icme_window = slice(icme_date - window_size, icme_date + window_size)
    icme_data = stereo_df[icme_window][cols]
    for col, ax in zip(cols, axes):
        x = icme_data.index
        if normalize:
            y = StandardScaler().fit_transform(icme_data[col].values.reshape(-1,1)).flatten()
        else:
            y = icme_data[col]
        ax.plot(x, y, lw=0.8)
        ax.set_ylabel(col)
#         mean, std = icme_data[col].mean(), icme_data[col].std()
#         if 'temp' in col.lower() and not normalize:
#             ax.set_yscale('log')
        ax.grid(True, lw=0.8, ls='--', alpha=0.5, c='k')
        ax.axvline(icme_date, ls='-', c='r', lw=1.25)
    axes[0].set_title(f'Normalized ICME Measurements\n ICME Start time: {icme_date}')
    return fig


sta_icme_df.head()

sta_icme_df.head()

sta_icme_df.index[0]

fig = quality_check_plot(sta_data_cut_df, sta_icme_df.index[1], window_size=dt.timedelta(days=1), cols=cols_of_interest, normalize=False)

sta_data_cut_df['2014'].plot(y='BTOTAL')

sta_data_cut_detrend_df = sta_data_cut_df - sta_data_cut_df.rolling('1D', center=True).mean()

errors = []
with PdfPages('icmes_stereoA_4day_window_detrended.pdf', 'w') as pdf:
    for date in tqdm(sta_icme_df.index):
        try:
            fig = quality_check_plot(
                sta_data_cut_detrend_df, 
                date, 
                window_size=dt.timedelta(days=2), 
                cols=cols_of_interest,
                normalize=False
            )
        except Exception as e:
            print(e)
            errors.append(date)
        else:
            pdf.savefig(fig)
            plt.close(fig)

sta_data_cut_detrend_df.index[0].strftime('%Y-%m-%d_%H:%M:%S')

sta_data_cut_detrend_df.index[0].strftime


def check_sampling_freq(mag_df, min_sep=None, verbose=False):
    """Determine the sampling frequency from the data

    Compute a weighted-average of the sampling frequency
    present in the time-series data. This is done by taking
    the rolling difference between consecutive datetime indices
    and then binning them up using a method of pd.Series objects.
    Also computes some statistics describing the distribution of
    sampling frequencies.

    Parameters
    ----------
    mag_df : pd.DataFrame
        Pandas dataframe containing the magnetometer data

    min_sep : float
        Minimum separation between two consecutive observations 
        to be consider usable for discontinuity identification

    verbose : boolean
        Specifies information on diverging sampling frequencies
    
    Returns
    -------
    avg_sampling_freq : float
        Weighted average of the sampling frequencies in the dataset

    stats : dict
        Some descriptive statistics for the interval
    
    """
    # Boolean flag for quality of data in interval
    # Assume its not bad and set to True if it is
    bad = False

    # Compute the time difference between consecutive measurements
    # a_i - a_{i-1} and save the data as dt.timedelta objects
    # rounded to the nearest milisecond
    diff_dt = mag_df.index.to_series().diff(1).round('ms')
    sampling_freqs = diff_dt.value_counts()
    sampling_freqs /= sampling_freqs.sum()

    avg_sampling_freq = 0
    for t, percentage in sampling_freqs.items():
        avg_sampling_freq += t.total_seconds() * percentage

    # Compute the difference in units of seconds so we can compute the RMS
    diff_s = np.array(
                list(
                    map(lambda val: val.total_seconds(), diff_dt)
                )
            )

    # Compute the RMS of the observation times to look for gaps in 
    # in the observation period
    t_rms = np.sqrt(
                np.nanmean(
                    np.square(diff_s)
                )
            )
    # flag that the gaps larger the min_sep. 
    if min_sep is None:
        min_sep = 5 * t_rms

    gap_indices = np.where(diff_s > min_sep)[0]
    n_gaps = len(gap_indices)
    
    try:
        previous_indices = gap_indices - 1
    except TypeError as e:
#         LOG.warning(e)
        print(e)
        total_missing = 0
    else:
        interval_durations = mag_df.index[gap_indices] \
                                - mag_df.index[previous_indices]
        total_missing = sum(interval_durations.total_seconds())

    # Compute the duration of the entire interval and determine the coverage
    total_duration = (mag_df.index[-1] - mag_df.index[0]).total_seconds()
    coverage = 1 - total_missing / total_duration

    if verbose and coverage < 0.5:
        msg = (
            f"\n Observational coverage: {coverage:0.2%}\n"
            f"Number of data gaps: {n_gaps:0.0f}\n"
            f"Average sampling rate: {avg_sampling_freq:0.5f}"
            )
#         LOG.warning(msg)
        print(msg)
        bad = True

    stats_data = {}
    stats_data['average_freq'] = avg_sampling_freq
    stats_data['max_freq'] = sampling_freqs.index.max().total_seconds()
    stats_data['min_freq'] = sampling_freqs.index.min().total_seconds()
    stats_data['n_gaps'] = len(gap_indices)
    stats_data['starttime_gaps'] = [mag_df.index[previous_indices]]
    stats_data['total_missing'] = total_missing
    stats_data['coverage'] = coverage

    return avg_sampling_freq, stats_data, bad


sta_data_cut_df.head()

check_sampling_freq(sta_data_cut_df, min_sep=120)

# ### STEREO B dataset (not used)

cols_of_interest = [
    'BTOTAL', 
    'BX(RTN)',
    'BY(RTN)',
    'BZ(RTN)', 
    'VP_RTN',
    'TEMPERATURE',
    'BETA',
    'Np'
]


def read_stereo_datasets(fname):
    with open(fname, 'r') as fobj:
        lines = fobj.readlines()

    colnames = lines[0].split()
    tmp = lines[1].split()
    units = []
    units.append(' '.join(tmp[:2]))
    units += tmp[2:]

    for col, unit in zip(colnames, units):
        print(col, unit)
        
    data = []
    index = []
    for line in tqdm(lines[2:]):
        lsplit = line.split()
        index.append(dt.datetime.strptime(' '.join(lsplit[:2]), '%d-%m-%Y %H:%M:%S.%f'))
        data.append(list(map(float, lsplit[2:])))
    
    df = pd.DataFrame(data, columns=colnames[1:], index=pd.DatetimeIndex(index))
    return df


stb_data_df = read_stereo_datasets('../data/stb_l2_magplasma.txt')

stb_data_df.head()

stb_data_cut_df = stb_data_df[stb_data_df.gt(-1e30)]

stb_data_cut_df.info()

stb_data_cut_df = stb_data_cut_df.dropna()

stb_data_cut_df.info()

stb_data_cut_detrend_df = stb_data_cut_df - stb_data_cut_df.rolling('1D', center=True).mean()

errors = []
with PdfPages('icmes_stereoB_4day_window_detrended.pdf', 'w') as pdf:
    for date in tqdm(stb_icme_df.index):
        try:
            fig = quality_check_plot(
                stb_data_cut_detrend_df, 
                date, 
                window_size=dt.timedelta(days=2), 
                cols=cols_of_interest,
                normalize=False
            )
        except Exception as e:
            print(e)
            errors.append(date)
        else:
            pdf.savefig(fig)
            plt.close(fig)

StarcoderdataJupyterScriptsDedupFiltered