#!/usr/bin/env python
# coding: utf-8

# # Notebook

# A Coronal Mass Ejection (CME) throws magnetic flux and plasma from the Sun into interplanetary space. These eruptions are actually related to solar flares -- in fact, CMEs and solar flares are considered “a single magnetically driven event” ([Webb & Howard 2012](http://adsabs.harvard.edu/abs/2012LRSP....9....3W)), wherein a flare unassociated with a CME is called a confined or compact flare.

#
# In general, the more energetic a flare, the more likely it is to be associated with a CME ([Yashiro et al. 2005](http://adsabs.harvard.edu/abs/2005JGRA..11012S05Y)) -- but this is not, by any means, a rule. For example, [Sun et al. (2015)](http://adsabs.harvard.edu/abs/2015ApJ...804L..28S) found that the largest active region in the last 24 years, shown below, produced 6 X-class flares but not a single observed CME.

#
# In this notebook, we will be predicting whether or not a flaring active region will also emit a CME using a machine learning algorithm from the scikit-learn package called Support Vector Machine.
#
# The analysis that follows is published in [Bobra & Ilonidis, 2016,  Astrophysical Journal, 821, 127](http://adsabs.harvard.edu/abs/2016ApJ...821..127B). If you use any of this code, we ask that you cite Bobra & Ilonidis (2016).
#
# Here is a video that explains the difference between a flare and a CME:

# In[1]:

from IPython.display import YouTubeVideo
YouTubeVideo("TWjtYSRlOUI")

# To do this analysis, we'll look at every active region observed by the Helioseismic and Magnetic Imager instrument on NASA's Solar Dynamics Observatory (SDO) satellite over the last eight years. Each active region is characterized by a bunch of features. These features describe the magnetic field at the solar surface. One feature, for example, is the total energy contained within an active region. Another is the total flux through an active region. We have 18 features, all of which are calculated every 12 minutes throughout an active region's lifetime. See [Bobra et al., 2014](http://link.springer.com/article/10.1007%2Fs11207-014-0529-3) for more information on how we calculate these features.

#
# We'll then ascribe each active region to one of two classes:
#
# 1. The positive class contains flaring active regions that did produce a CME.
# 2. The negative class contains flaring active regions that did not produce a CME.

# First, we'll import some modules.

# In[2]:

import numpy as np
import matplotlib.pylab as plt
import matplotlib.mlab as mlab
import pandas as pd
import scipy.stats
import requests
import urllib
import json
from datetime import datetime as dt_obj
from datetime import timedelta
from sklearn import svm
from sklearn.model_selection import StratifiedKFold
from sunpy.time import TimeRange
import sunpy.instr.goes
import lime
import lime.lime_tabular
pd.set_option('display.max_rows', 500)
get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")

# Now we'll gather the data. The data come from three different places:
#
# 1. CME data from SOHO/LASCO and STEREO/SECCHI coronographs, which can be accesed from the [DONKI database](http://kauai.ccmc.gsfc.nasa.gov/DONKI/) at NASA Goddard. This tells us if an active region has produced a CME or not.
# 2. Flare data from the GOES flare catalog at NOAA, which can be accessed with the `sunpy.instr.goes.get_event_list()` function. This tells us if an active region produced a flare or not.
# 3. Active region data from the Solar Dynamics Observatory's Heliosesmic and Magnetic Imager instrument, which can be accessed from the [JSOC database](http://jsoc.stanford.edu/) via a JSON API. This gives us the features characterizing each active region.

# ## Step 1: Gathering data for the positive class

# Let's first query the [DONKI database](http://kauai.ccmc.gsfc.nasa.gov/DONKI/) to get the data associated with the positive class. Be forewarned: there's a lot of data cleaning involved with building the positive class.

# In[3]:

# request the data
baseurl = "https://kauai.ccmc.gsfc.nasa.gov/DONKI/WS/get/FLR?"
t_start = "2010-05-01"
t_end = "2018-04-01"
url = baseurl+"startDate="+t_start+"&endDate="+t_end

# if there's no response at this time, print warning
response = requests.get(url)
if response.status_code != 200:
    print('cannot successfully get an http response')

# In[4]:

# read the data

print("Getting data from", url)
df = pd.read_json(url)

# select flares associated with a linked event (SEP or CME), and
# select only M or X-class flares
events_list = df.loc[df['classType'].str.contains(
    "M|X") & ~df['linkedEvents'].isnull()]

# drop all rows that don't satisfy the above conditions
events_list = events_list.reset_index(drop=True)

# In[5]:

# drop the rows that aren't linked to CME events
for i in range(events_list.shape[0]):
    value = events_list.loc[i]['linkedEvents'][0]['activityID']
    if not "CME" in value:
        print(value, "not a CME, dropping row")
        events_list = events_list.drop([i])
events_list = events_list.reset_index(drop=True)

# Convert the `peakTime` column in the `events_list` dataframe from a string into a datetime object:

# In[6]:

def parse_tai_string(tstr):
    year = int(tstr[:4])
    month = int(tstr[5:7])
    day = int(tstr[8:10])
    hour = int(tstr[11:13])
    minute = int(tstr[14:16])
    return dt_obj(year, month, day, hour, minute)

for i in range(events_list.shape[0]):
    events_list['peakTime'].iloc[i] = parse_tai_string(
        events_list['peakTime'].iloc[i])

# Check for Case 1: In this case, the CME and flare exist but NOAA active region number does not exist in the DONKI database.

# In[7]:

# Case 1: CME and Flare exist but NOAA active region number does not exist in DONKI database

number_of_donki_mistakes = 0  # count the number of DONKI mistakes
# create an empty array to hold row numbers to drop at the end
event_list_drops = []

for i in range(events_list.shape[0]):
    if (np.isnan(events_list.loc[i]['activeRegionNum'])):
        time = events_list['peakTime'].iloc[i]
        time_range = TimeRange(time, time)
        listofresults = sunpy.instr.goes.get_goes_event_list(time_range, 'M1')
        if (listofresults[0]['noaa_active_region'] == 0):
            print(events_list.loc[i]['activeRegionNum'], events_list.loc[i]
                  ['classType'], "has no match in the GOES flare database ; dropping row.")
            event_list_drops.append(i)
            number_of_donki_mistakes += 1
            continue
        else:
            print("Missing NOAA number:", events_list['activeRegionNum'].iloc[i], events_list['classType'].iloc[i],
                  events_list['peakTime'].iloc[i], "should be", listofresults[0]['noaa_active_region'], "; changing now.")
            events_list['activeRegionNum'].iloc[i] = listofresults[0]['noaa_active_region']
            number_of_donki_mistakes += 1

# Drop the rows for which there is no active region number in both the DONKI and GOES flare databases
events_list = events_list.drop(event_list_drops)
events_list = events_list.reset_index(drop=True)
print('There are', number_of_donki_mistakes, 'DONKI mistakes so far.')

# Now we grab all the data from the GOES database in preparation for checking Cases 2 and 3.

# In[8]:

# Grab all the data from the GOES database
time_range = TimeRange(t_start, t_end)
listofresults = sunpy.instr.goes.get_goes_event_list(time_range, 'M1')
print('Grabbed all the GOES data; there are', len(listofresults), 'events.')

# Check for Case 2: In this case, the NOAA active region number is wrong in the DONKI database.

# In[9]:

# Case 2: NOAA active region number is wrong in DONKI database

# collect all the peak flares times in the NOAA database
peak_times_noaa = [item["peak_time"] for item in listofresults]

for i in range(events_list.shape[0]):
    # check if a particular DONKI flare peak time is also in the NOAA database
    peak_time_donki = events_list['peakTime'].iloc[i]
    if peak_time_donki in peak_times_noaa:
        index = peak_times_noaa.index(peak_time_donki)
    else:
        continue
    # ignore NOAA active region numbers equal to zero
    if (listofresults[index]['noaa_active_region'] == 0):
        continue
    # if yes, check if the DONKI and NOAA active region numbers match up for this peak time
    # if they don't, flag this peak time and replace the DONKI number with the NOAA number
    if (listofresults[index]['noaa_active_region'] != int(events_list['activeRegionNum'].iloc[i])):
        print('Messed up NOAA number:', int(events_list['activeRegionNum'].iloc[i]), events_list['classType'].iloc[i],
              events_list['peakTime'].iloc[i], "should be", listofresults[index]['noaa_active_region'], "; changing now.")
        events_list['activeRegionNum'].iloc[i] = listofresults[index]['noaa_active_region']
        number_of_donki_mistakes += 1
print('There are', number_of_donki_mistakes, 'DONKI mistakes so far.')

# Check for Case 3: In this case, the flare peak time is wrong in the DONKI database.

# In[10]:

# Case 3: The flare peak time is wrong in the DONKI database.

# create an empty array to hold row numbers to drop at the end
event_list_drops = []

active_region_numbers_noaa = [item["noaa_active_region"]
                              for item in listofresults]
flare_classes_noaa = [item["goes_class"] for item in listofresults]

for i in range(events_list.shape[0]):
    # check if a particular DONKI flare peak time is also in the NOAA database
    peak_time_donki = events_list['peakTime'].iloc[i]
    if not peak_time_donki in peak_times_noaa:
        active_region_number_donki = int(
            events_list['activeRegionNum'].iloc[i])
        flare_class_donki = events_list['classType'].iloc[i]
        flare_class_indices = [i for i, x in enumerate(
            flare_classes_noaa) if x == flare_class_donki]
        active_region_indices = [i for i, x in enumerate(
            active_region_numbers_noaa) if x == active_region_number_donki]
        common_indices = list(
            set(flare_class_indices).intersection(active_region_indices))
        if common_indices:
            print("Messed up time:", int(events_list['activeRegionNum'].iloc[i]), events_list['classType'].iloc[i],
                  events_list['peakTime'].iloc[i], "should be", peak_times_noaa[common_indices[0]], "; changing now.")
            events_list['peakTime'].iloc[i] = peak_times_noaa[common_indices[0]]
            number_of_donki_mistakes += 1
        if not common_indices:
            print("DONKI flare peak time",
                  events_list['peakTime'].iloc[i], "has no match; dropping row.")
            event_list_drops.append(i)
            number_of_donki_mistakes += 1

# Drop the rows for which the NOAA active region number and flare class associated with
# the messed-up flare peak time in the DONKI database has no match in the GOES flare database
events_list = events_list.drop(event_list_drops)
events_list = events_list.reset_index(drop=True)

# Create a list of corrected flare peak times
peak_times_donki = [events_list['peakTime'].iloc[i]
                    for i in range(events_list.shape[0])]

print('There are', number_of_donki_mistakes, 'DONKI mistakes so far.')

# This is our final table of events that fall into the positive class:

# In[11]:

events_list

# Now let's query the JSOC database to see if there are active region parameters at the time of the flare. First read the following file to map NOAA active region numbers to HARPNUMs (a HARP, or an HMI Active Region Patch, is the preferred numbering system for the HMI active regions as they appear in the magnetic field data before NOAA observes them in white light):

# In[12]:

answer = pd.read_csv(
    'http://jsoc.stanford.edu/doc/data/hmi/harpnum_to_noaa/all_harps_with_noaa_ars.txt', sep=' ')

# Now, let's determine at which time we'd like to predict CMEs. In general, many people try to predict a CME either 24 or 48 hours before it happens. We can report both in this study by setting a variable called `timedelayvariable`:

# In[13]:

timedelayvariable = 24

# Now, we'll convert subtract `timedelayvariable` from the GOES Peak Time and re-format the datetime object into a string that JSOC can understand:

# In[14]:

t_rec = [(events_list['peakTime'].iloc[i] - timedelta(hours=timedelayvariable)
          ).strftime('%Y.%m.%d_%H:%M_TAI') for i in range(events_list.shape[0])]

# Now we can grab the SDO data from the JSOC database by executing the JSON queries. We are selecting data that satisfies several criteria: The data has to be [1] disambiguated with a version of the disambiguation module greater than 1.1, [2] taken while the orbital velocity of the spacecraft is less than 3500 m/s, [3] of a high quality, and [4] within 70 degrees of central meridian. If the data pass all these tests, they are stuffed into one of two lists: one for the positive class (called CME_data) and one for the negative class (called no_CME_data).

# In[15]:

def get_the_jsoc_data(event_count, t_rec):
    """
    Parameters
    ----------
    event_count: number of events
                 int

    t_rec:       list of times, one associated with each event in event_count
                 list of strings in JSOC format ('%Y.%m.%d_%H:%M_TAI')

    """

    catalog_data = []
    classification = []

    for i in range(event_count):

        print("=====", i, "=====")
        # next match NOAA_ARS to HARPNUM
        idx = answer[answer['NOAA_ARS'].str.contains(
            str(int(listofactiveregions[i])))]

        # if there's no HARPNUM, quit
        if (idx.empty == True):
            print('skip: there are no matching HARPNUMs for',
                  str(int(listofactiveregions[i])))
            continue

        # construct jsoc_info queries and query jsoc database; we are querying for 25 keywords
        url = "http://jsoc.stanford.edu/cgi-bin/ajax/jsoc_info?ds=hmi.sharp_720s["+str(
            idx.HARPNUM.values[0])+"]["+t_rec[i]+"][? (CODEVER7 !~ '1.1 ') and (abs(OBS_VR)< 3500) and (QUALITY<65536) ?]&op=rs_list&key=USFLUX,MEANGBT,MEANJZH,MEANPOT,SHRGT45,TOTUSJH,MEANGBH,MEANALP,MEANGAM,MEANGBZ,MEANJZD,TOTUSJZ,SAVNCPP,TOTPOT,MEANSHR,AREA_ACR,R_VALUE,ABSNJZH"
        response = requests.get(url)

        # if there's no response at this time, quit
        if response.status_code != 200:
            print('skip: cannot successfully get an http response')
            continue

        # read the JSON output
        data = response.json()

        # if there are no data at this time, quit
        if data['count'] == 0:
            print('skip: there are no data for HARPNUM',
                  idx.HARPNUM.values[0], 'at time', t_rec[i])
            continue

        # check to see if the active region is too close to the limb
        # we can compute the latitude of an active region in stonyhurst coordinates as follows:
        # latitude_stonyhurst = CRVAL1 - CRLN_OBS
        # for this we have to query the CEA series (but above we queried the other series as the CEA series does not have CODEVER5 in it)

        url = "http://jsoc.stanford.edu/cgi-bin/ajax/jsoc_info?ds=hmi.sharp_cea_720s["+str(
            idx.HARPNUM.values[0])+"]["+t_rec[i]+"][? (abs(OBS_VR)< 3500) and (QUALITY<65536) ?]&op=rs_list&key=CRVAL1,CRLN_OBS"
        response = requests.get(url)

        # if there's no response at this time, quit
        if response.status_code != 200:
            print('skip: failed to find CEA JSOC data for HARPNUM',
                  idx.HARPNUM.values[0], 'at time', t_rec[i])
            continue

        # read the JSON output
        latitude_information = response.json()

        # if there are no data at this time, quit
        if latitude_information['count'] == 0:
            print('skip: there are no data for HARPNUM',
                  idx.HARPNUM.values[0], 'at time', t_rec[i])
            continue

        CRVAL1 = float(latitude_information['keywords'][0]['values'][0])
        CRLN_OBS = float(latitude_information['keywords'][1]['values'][0])
        if (np.absolute(CRVAL1 - CRLN_OBS) > 70.0):
            print('skip: latitude is out of range for HARPNUM',
                  idx.HARPNUM.values[0], 'at time', t_rec[i])
            continue

        if ('MISSING' in str(data['keywords'])):
            print('skip: there are some missing keywords for HARPNUM',
                  idx.HARPNUM.values[0], 'at time', t_rec[i])
            continue

        print('accept NOAA Active Region number', str(int(
            listofactiveregions[i])), 'and HARPNUM', idx.HARPNUM.values[0], 'at time', t_rec[i])

        individual_flare_data = []
        for j in range(18):
            individual_flare_data.append(
                float(data['keywords'][j]['values'][0]))

        catalog_data.append(list(individual_flare_data))

        single_class_instance = [idx.HARPNUM.values[0], str(
            int(listofactiveregions[i])), listofgoesclasses[i], t_rec[i]]
        classification.append(single_class_instance)

    return catalog_data, classification

# Now we prepare the data to be fed into the function:

# In[16]:

listofactiveregions = list(events_list['activeRegionNum'].values.flatten())
listofgoesclasses = list(events_list['classType'].values.flatten())

# And call the function:

# In[17]:

positive_result = get_the_jsoc_data(events_list.shape[0], t_rec)

# Here is the number of events associated with the positive class:

# In[18]:

CME_data = positive_result[0]
positive_class = positive_result[1]
print("There are", len(CME_data), "CME events in the positive class.")

# ## Step 2: Gathering data for the negative class

# To gather the examples for the negative class, we only need to:
#
# 1. Query the GOES database for all the M- and X-class flares during our time of interest, and
# 2. Select the ones that are not associated with a CME.

# In[19]:

# select peak times that belong to both classes
all_peak_times = np.array([(listofresults[i]['peak_time'])
                           for i in range(len(listofresults))])

negative_class_possibilities = []
counter_positive = 0
counter_negative = 0
for i in range(len(listofresults)):
    this_peak_time = all_peak_times[i]
    if (this_peak_time in peak_times_donki):
        counter_positive += 1
    else:
        counter_negative += 1
        this_instance = [listofresults[i]['noaa_active_region'],
                         listofresults[i]['goes_class'], listofresults[i]['peak_time']]
        negative_class_possibilities.append(this_instance)
print("There are", counter_positive, "events in the positive class.")
print("There are", counter_negative, "events in the negative class.")

# Again, we compute times that are one day before the flare peak time and convert it into a string that JSOC can understand:

# In[20]:

t_rec = np.array([(negative_class_possibilities[i][2] - timedelta(hours=timedelayvariable)
                   ).strftime('%Y.%m.%d_%H:%M_TAI') for i in range(len(negative_class_possibilities))])

# And again, we query the JSOC database to see if these data are present:

# In[21]:

listofactiveregions = list(
    negative_class_possibilities[i][0] for i in range(counter_negative))
listofgoesclasses = list(
    negative_class_possibilities[i][1] for i in range(counter_negative))

# In[22]:

negative_result = get_the_jsoc_data(counter_negative, t_rec)

# Here is the number of events associated with the negative class:

# In[23]:

no_CME_data = negative_result[0]
negative_class = negative_result[1]
print("There are", len(no_CME_data), "no-CME events in the negative class.")

# ## Step 3: Feature selection

# Some of the features within a data set may be powerful for distinguishing between the positive and negative class, whereas others may be redundant or irrelevant. To identify features in the former category, we use a univariate feature selection method, which is implemented in the feature selection module of the scikit-learn library, for feature scoring.

# To improve the performance of the feature selection algorithm, we'll normalize each feature so that they lie within similar ranges. To do this, we subtract from every feature its median value and divide by its standard deviation.

# In[24]:

CME_data = np.array(CME_data)
no_CME_data = np.array(no_CME_data)

# In[25]:

def normalize_the_data(flare_data):
    flare_data = np.array(flare_data)
    n_elements = flare_data.shape[0]
    for j in range(flare_data.shape[1]):
        standard_deviation_of_this_feature = np.std(flare_data[:, j])
        median_of_this_feature = np.median(flare_data[:, j])
        for i in range(n_elements):
            flare_data[i, j] = (
                flare_data[i, j] - median_of_this_feature) / (standard_deviation_of_this_feature)
    return flare_data

no_CME_data = normalize_the_data(no_CME_data)
CME_data = normalize_the_data(CME_data)

print("There are", no_CME_data.shape[0], "flares with no associated CMEs.")
print("There are", CME_data.shape[0], "flares with associated CMEs.")

# Let's look at the distribution of one feature for the active regions that both flared and produced a CME (green) and for the active regions that flared but did not produce a CME (red). You can change the value of `i` in the code block below to see that some features are totally useless as there is barely any difference in the distributions for the positive and negative class. As such, we can throw such features out of our sample. It's a good idea to do some feature selection before running the SVM, so as to reduce noise (in this case, with only 18 features, there's not too much noise to begin with).

# In[26]:

sharps = ['Total unsigned flux', 'Mean gradient of total field',
          'Mean current helicity (Bz contribution)', 'Mean photospheric magnetic free energy',
          'Fraction of Area with Shear > 45 deg', 'Total unsigned current helicity',
          'Mean gradient of horizontal field', 'Mean characteristic twist parameter, alpha',
          'Mean angle of field from radial', 'Mean gradient of vertical field',
          'Mean vertical current density', 'Total unsigned vertical current',
          'Sum of the modulus of the net current per polarity',
          'Total photospheric magnetic free energy density', 'Mean shear angle',
          'Area of strong field pixels in the active region', 'Sum of flux near polarity inversion line',
          'Absolute value of the net current helicity']

i = 2

# For the positive class (green)
mu_fl = np.mean(CME_data[:, i])
sigma_fl = np.std(CME_data[:, i])
num_bins = 15
n_fl, bins_fl, patches_fl = plt.hist(
    CME_data[:, i], num_bins, normed=1, facecolor='green', alpha=0.5)
y_fl = scipy.stats.norm.pdf(bins_fl, mu_fl, sigma_fl)
plt.plot(bins_fl, y_fl, 'g--', label='positive class')

# For the negative class (red)
mu_nofl = np.mean(no_CME_data[:, i])
sigma_nofl = np.std(no_CME_data[:, i])
n_nofl, bins_nofl, patches_nofl = plt.hist(
    no_CME_data[:, i], num_bins, normed=1, facecolor='red', alpha=0.5)
y_nofl = scipy.stats.norm.pdf(bins_nofl, mu_nofl, sigma_nofl)
plt.plot(bins_nofl, y_nofl, 'r--', label='negative class')

text_style = dict(fontsize=16, fontdict={'family': 'monospace'})
plt.xlabel('Normalized '+sharps[i], **text_style)
plt.ylabel('Number (normalized)', labelpad=20, **text_style)
fig = plt.gcf()
fig.set_size_inches(10, 5)
fig.savefig('fscore_tmp.png', bbox_inches='tight')
legend = plt.legend(loc='upper right', fontsize=12, framealpha=0.0, title='')
legend.get_frame().set_linewidth(0.0)

# Now we will compute the Univariate F-score for feature selection. It is a very simple method: the F-score measures the distance between the two distributions for a given feature (inter-class distance), divided by the sum of the variances for this feature (intra-class distance). We can use the `sklearn.feature_selection` module to do this:

# In[27]:

# import the feature selection method
from sklearn.feature_selection import SelectKBest, f_classif
# select the number of features
N_features = 18
Nfl = CME_data.shape[0]
Nnofl = no_CME_data.shape[0]
yfl = np.ones(Nfl)
ynofl = np.zeros(Nnofl)
# k is the number of features
selector = SelectKBest(f_classif, k=N_features)
selector.fit(np.concatenate((CME_data, no_CME_data), axis=0),
             np.concatenate((yfl, ynofl), axis=0))
scores = selector.scores_
print(scores)

# It's not easy to interpret the scores in this fashion, so let's plot the results. The higher the Univariate Fisher Score, the more predictive the feature.

# In[28]:

plt.clf()
order = np.argsort(scores)
orderedsharps = [sharps[i] for i in order]
y_pos2 = np.arange(18)
plt.barh(y_pos2, sorted(scores/np.max(scores)))
plt.ylim((-1, 19))
plt.yticks(y_pos2, orderedsharps, fontsize=12)
plt.xlabel('Normalized Fisher Score', fontsize=15)
plt.title('Normalized Univariate Fisher Score Per Feature', fontsize=15)
plt.subplots_adjust(left=0.5, right=1.0)
fig = plt.gcf()
fig.set_size_inches(9, 5)
plt.show()

# ## Step 4: The support vector machine

# Now we initialize the support vector machine on the data. The SVM uses non-linear decision functions to map the feature space into a higher-dimensional space, where the positive and negative examples can be separated linearly by a hyperplane.

#
# This is incredibly non-intuitive. But we can think of a simpler example. Suppose we had two classes: CME-producing and non-CME producing active regions. And suppose we had two features: the total flux in these regions, and the total area of these regions. We could construct a two-dimentional feature space, where we plot the flux against the area of each active region. Positive examples could be indicated by an X and negatives ones by an O. In theory, if our data behaved well, we could draw a line between these classess.

#
# Since we have 18 features, the SVM constructs an 18-dimensional feature space. In this feature space, the decision boundary separating the positive and negative examples may be non-linear. As such, the algorithm then enlarges this 18-dimensional feature space (using the function indicated by the `kernel` parameter in the [`svm.SVC`](http://scikit-learn.org/stable/modules/generated/sklearn.svm.SVC.html) function) into a higher-dimensional feature space wherein it is possible to linearly separate the positive and negatives classes. There are lots of people trying to work on how to [visualize these multi-dimensional feature spaces](https://github.com/tmadl/highdimensional-decision-boundary-plot), which is an active area of research.

# In[29]:

number_of_examples = Nfl + Nnofl
C = 4.0
gamma = 0.075
class_weight = {1: 6.5}
clf = svm.SVC(C=C, gamma=gamma, kernel='rbf', class_weight=class_weight,
              cache_size=500, max_iter=-1, shrinking=True, tol=1e-8, probability=True)

# ## Step 5: Stratified k-folds cross-validation

# Now we run and evaluate the performance of the SVM. There are lots of different ways to evaluate the performance of a classifier, which we discuss in Section 4 of [Bobra & Couvidat (2015)](https://arxiv.org/abs/1411.1405). We're going to choose a metric called the True Skill Score, or the TSS, which we can calculate from four quantities: true positives, true negatives, false positives, and false negatives. We prefer the TSS to all the other metrics as it is insensitive to the class imbalance ratio and thus best for comparison to other groups. The TSS is symmetrically distributed about 0: i.e., it goes from [-1, 1] where 0 represents no skill and a negative value represents a perverse prediction. Thus we are able to predict CMEs in a fashion better than randomly guessing. Here we define a confusion table to measure the performance of our binary classification:

# In[30]:

def confusion_table(pred, labels):
    """
    computes the number of TP, TN, FP, FN events given the arrays with predictions and true labels
    and returns the true skill score

    Args:
    pred: np array with predictions (1 for flare, 0 for nonflare)
    labels: np array with true labels (1 for flare, 0 for nonflare)

    Returns: true negative, false positive, true positive, false negative
    """
    Nobs = len(pred)
    TN = 0.
    TP = 0.
    FP = 0.
    FN = 0.
    for i in range(Nobs):
        if (pred[i] == 0 and labels[i] == 0):
            TN += 1
        elif (pred[i] == 1 and labels[i] == 0):
            FP += 1
        elif (pred[i] == 1 and labels[i] == 1):
            TP += 1
        elif (pred[i] == 0 and labels[i] == 1):
            FN += 1
        else:
            print("Error! Observation could not be classified.")
    return TN, FP, TP, FN

# Now we run the SVM on our data and cross-validate our results. In our case, the positive sample size is quite small (both objectively and compared to the negative sample size). Therefore, we use a stratified k-folds cross-validation method, which makes k partitions of the data set and uses k-1 folds for training the SVM and 1 fold for testing the trained SVM. The stratification preserves the ratio of positive to negative examples per fold. Then we can permute over the partitions such that each partition eventually makes its way into the testing set. For each individual testing set, we can calculate a skill score. Then we can average the skill scores over the total number of testing sets.
#
# To compute the TSS, we must first select a value of k. k can be arbitrarily defined and take any value between 2 and `number_of_examples`, so we can explore this parameter space. As k approaches `number_of_examples`, the k-fold method reduces to the Leave One Out method, in which only one example is in the testing set and all other examples are in the training set. The literature suggests this method is not the best, so we can stray away from high values of k. Many studies (e.g. [Kohavi, 1995](http://web.cs.iastate.edu/~jtian/cs573/Papers/Kohavi-IJCAI-95.pdf)) recommend the stratified 10-fold cross-validation to reduce variance and bias. Here, we test their recommendation by computing the TSS using 50 k values, ranging from 2 to 52.

# In[31]:

# lists to hold the TSS and standard deviation of the TSS
array_of_avg_TSS = np.ndarray([50])
array_of_std_TSS = np.ndarray([50])

# xdata are the examples
# ydata are the labels
xdata = np.concatenate((CME_data, no_CME_data), axis=0)
ydata = np.concatenate((np.ones(Nfl), np.zeros(Nnofl)), axis=0)

# mdata contain metadata about the active region that will be useful
# when we interpret the results using LIME
mdata = np.concatenate((positive_class, negative_class), axis=0)

# compute the TSS for a variety of k ranging from 2 to 52
# this is to see how the TSS varies as a function of k, and to test if k=10 really makes sense
for k in range(2, 52):
    skf = StratifiedKFold(n_splits=k, shuffle=True)
    these_TSS_for_this_k = []
    for train_index, test_index in skf.split(xdata, ydata):
        # xtrain are the examples in the training set
        xtrain = xdata[train_index]
        # ytrain are the labels in the training set
        ytrain = ydata[train_index]
        # xtest are the examples in the testing set
        xtest = xdata[test_index]
        ytest = ydata[test_index]    # ytest are the labels in the testing set
        # metadata useful for interpreting with LIME
        mtrain = mdata[train_index]
        # metadata useful for interpreting with LIME
        mtest = mdata[test_index]
        clf.fit(xtrain, ytrain)
        TN, FP, TP, FN = confusion_table(clf.predict(xtest), ytest)
        if (((TP+FN) == 0.0) or (FP+TN) == 0.0):
            these_TSS_for_this_k.append(np.nan)
            continue
        else:
            these_TSS_for_this_k.append(TP/(TP+FN) - FP/(FP+TN))
    TSS_k = np.array(these_TSS_for_this_k)
    array_of_avg_TSS[k-2] = np.mean(TSS_k)
    array_of_std_TSS[k-2] = np.std(TSS_k)

# Now we can plot the mean TSS per k, using the standard deviation as the error in the TSS. We see that for high values of k, the standard deviation in the TSS can be greater than the mean. These points are indicated in blue.

# In[32]:

fig, ax = plt.subplots(figsize=(10, 8))      # define the size of the figure
orangered = (1.0, 0.27, 0, 1.0)              # create an orange-red color
cornblue = (0.39, 0.58, 0.93, 1.0)           # create a cornflower-blue color

# define some style elements
marker_style_red = dict(linestyle='', markersize=8,
                        fillstyle='full', color=orangered, markeredgecolor=orangered)
marker_style_blue = dict(linestyle='', markersize=8,
                         fillstyle='full', color=cornblue, markeredgecolor=cornblue)
text_style = dict(fontsize=16, fontdict={'family': 'monospace'})

# ascribe the data to the axes
k = np.arange(50)+2
for i in range(50):
    if (array_of_avg_TSS[i] > array_of_std_TSS[i]):
        ax.errorbar(k[i], array_of_avg_TSS[i],
                    yerr=array_of_std_TSS[i], linestyle='', color=orangered)
        ax.plot(k[i], array_of_avg_TSS[i], 'o', **marker_style_red)
    if (array_of_avg_TSS[i] <= array_of_std_TSS[i]):
        ax.errorbar(k[i], array_of_avg_TSS[i],
                    yerr=array_of_std_TSS[i], linestyle='', color=cornblue)
        ax.plot(k[i], array_of_avg_TSS[i], 'o', **marker_style_blue)

# set plot limits
plt.xlim(xmax=52, xmin=0)
plt.ylim(ymax=1.1, ymin=0)

# label the axes and the plot
ax.set_xlabel('k', **text_style)
ax.set_ylabel('TSS', labelpad=20, **text_style)
plt.title(r'TSS per k using stratified k-fold cross-validation', **text_style)
fig = plt.gcf()
fig.set_size_inches(10, 5)

# As such, we confirm that high k-values result in a high variance. We find it reasonable to use the stratified 10-fold cross-validation method to compute the TSS and will follow this recommendation. Therefore we report this score as our final result:

# In[33]:

print("The TSS equals", array_of_avg_TSS[9],
      "plus or minus", array_of_std_TSS[9], ".")

# ## Addendum : Local Interpretable Model-Agnostic Explanations (LIME)

# Machine-learning is a powerful technique that can help us predict CMEs. However, our goal is not only to predict CMEs, but also to quantitatively understand which signatures indicate the imminent eruption of a CME. But the practical successes of machine-learning algorithms are often not matched by successes in understanding, and this has become an issue within the machine-learning community ([Rahimi and Recht, 2017](http://www.argmin.net/2017/12/11/alchemy-addendum/)).
#
# The SVM is a good model to start with, because it is (relatively) simple and we can use the Univariate Fisher Score to identify the most predictive features. But it would also be useful to figure out why each individual active region was classed as positive or negative. To do this, we can use a tool called [LIME](https://github.com/marcotcr/lime) (or Local Interpretable Model-Agnostic Explanations).

#
# First, we initialize the LIME explainer:

# In[34]:

explainer = lime.lime_tabular.LimeTabularExplainer(
    xtrain, feature_names=sharps, class_names=['CME', 'no CME'], discretize_continuous=True)

# Then we use the explainer to explain its choice for a particular active region. To do this, the LIME module generates neighborhood data by randomly perturbing the values of the features associated with this active region. If, for any given feature, this perturbation does not change the outcome of the prediction, this feature isn't useful along the perturbed dimension. If, for any given feature, the perturbation does change the outcome of the prediction, this feature is useful along the perturbed dimension. Thus the explainer can determine which features are useful under which conditions.

# In[35]:

i = np.random.randint(0, xtest.shape[0])
exp = explainer.explain_instance(xtest[i], clf.predict_proba, num_features=8)

# Now we can visualize the results. The bottom left panel shows the probabilities assigned to this particular example (which are computed by the SVM via the `probability=True` parameter). The right panel plots the weights per feature (and indicates the values of these weights at the end of each horizontal bar). The text describes the conditions under which this feature  is predictive.

# In[36]:

print("Here is the prediction explanation for NOAA Active Region",
      mtest[i][1], "(HARPNUM ", mtest[i][0], "),\n which produced a", mtest[i][2], "class flare on", mtest[i][3], ".")
exp.show_in_notebook(show_table=False, show_all=False)

# Here is the same information in words:

# In[37]:

explained_list = exp.as_list()
for i in range(len(explained_list)):
    if (explained_list[i][1]) < 0:
        feature_sign = 'no CME'
    else:
        feature_sign = 'CME'
    print("The following condition:", explained_list[i][0], "\n predicts",
          feature_sign, "with a model weight of", abs(explained_list[i][1]), ".")


import abc
import numpy as np

class BaseOptimizer(abc.ABC):
    @abc.abstractmethod
    def optimize(self, x: np.ndarray, y: np.ndarray, epochs: int=5, save_hist: bool=False):
        # NOTE when save_hist is True, must return histories where histories[0] is the history of w
        pass


""" Unit Testing for lambdata package"""

import pandas as pd
import numpy as np
import unittest
import example_module as em

class LambdataTests(unittest.TestCase):
    def setUp(self):
        self.test_df = pd.DataFrame([[1, np.NaN], [np.NaN, np.NaN]])
        self.test_list = [1, 1]

    """ Testing lambdata package """
    def test_number_of_nulls(self):
        """ Testing number of nulls method """
        self.assertEqual(em.number_of_nulls(self.test_df), 3)

    def test_add_list_to_df(self):
        """ Testing add list to df """
        self.assertEqual(self.test_df.shape[1] + 1, em.add_list_to_df(self.test_list, self.test_df).shape[1])

if __name__ == "__main__":
    unittest.main()


"""
  Name     : c7_31_mrege_01.py
  Book     : Python for Finance (2nd ed.)
  Publisher: Packt Publishing Ltd.
  Author   : Yuxing Yan
  Date     : 6/6/2017
  email    : yany@canisius.edu
             paulyxy@hotmail.com
"""
import scipy as sp
import pandas as pd
#
x= pd.DataFrame({'key': ['K0', 'K1', 'K2', 'K3'],
                 'A': ['A0', 'A1', 'A2', 'A3'],
                 'B': ['B0', 'B1', 'B2', 'B3']})
y = pd.DataFrame({'key': ['K0', 'K1', 'K2', 'K6'],
                  'C': ['C0', 'C1', 'C2', 'C3'],
                  'D': ['D0', 'D1', 'D2', 'D3']})
print(sp.shape(x))
print(sp.shape(y))
print(x)
print(y)

result = pd.merge(x,y, on='key')
print(result)

result2=pd.merge(x,y)
print(result2)


"""
Copyright (C) 2014 The HDF Group
Copyright (C) 2014 John Evans

This example code illustrates how to access and visualize an NSIDC/ICESat/GLAS
L2 HDF5 file in Python.

If you have any questions, suggestions, or comments on this example, please use
the HDF-EOS Forum (http://hdfeos.org/forums).  If you would like to see an
example of any other NASA HDF/HDF-EOS data product that is not listed in the
HDF-EOS Comprehensive Examples page (http://hdfeos.org/zoo), feel free to
contact us at eoshelp@hdfgroup.org or post it at the HDF-EOS Forum
(http://hdfeos.org/forums).

Usage:  save this script and run

    python GLAH13_633_2103_001_1317_0_01_0001_a.py

The HDF file must either be in your current working directory or in a directory
specified by the environment variable HDFEOS_ZOO_DIR.
"""

import os

import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import numpy as np

# Can do this using either netCDF4 or h5py.
USE_NETCDF4 = False

def run(FILE_NAME):
    if USE_NETCDF4:

        from netCDF4 import Dataset

        nc = Dataset(FILE_NAME)

        latvar = nc.groups['Data_1HZ'].groups['Geolocation'].variables['d_lat']
        latitude = latvar[:]
        lat_vr = [latvar.valid_min, latvar.valid_max]

        lonvar = nc.groups['Data_1HZ'].groups['Geolocation'].variables['d_lon']
        longitude = lonvar[:]
        lon_vr = [lonvar.valid_min, lonvar.valid_max]

        tempvar = nc.groups['Data_1HZ'].groups['Atmosphere'].variables['d_Surface_temp']
        temp = tempvar[:]
        temp_vr = [tempvar.valid_min, tempvar.valid_max]
        units = tempvar.units
        longname = tempvar.long_name

        timevar = nc.groups['Data_1HZ'].groups['Time'].variables['d_UTCTime_1']
        time = timevar[:]

    else:

        import h5py

        with h5py.File(FILE_NAME, mode='r') as f:

            latvar = f['/Data_1HZ/Geolocation/d_lat']
            latitude = latvar[:]
            lat_vr = [latvar.attrs['valid_min'], latvar.attrs['valid_max']]

            lonvar = f['/Data_1HZ/Geolocation/d_lon']
            longitude = lonvar[:]
            lon_vr = [lonvar.attrs['valid_min'], lonvar.attrs['valid_max']]

            tempvar = f['/Data_1HZ/Atmosphere/d_Surface_temp']
            temp = tempvar[:]
            temp_vr = [tempvar.attrs['valid_min'], tempvar.attrs['valid_max']]
            units = tempvar.attrs['units']
            longname = tempvar.attrs['long_name']

            time = f['/Data_1HZ/Time/d_UTCTime_1'][:]

    # Apply attribute restrictions.
    latitude[latitude < lat_vr[0]] = np.nan
    latitude[latitude > lat_vr[1]] = np.nan
    longitude[longitude < lon_vr[0]] = np.nan
    longitude[longitude > lon_vr[1]] = np.nan
    temp[temp < temp_vr[0]] = np.nan
    temp[temp > temp_vr[1]] = np.nan

    # Just use a small subset.
    idx = slice(0, 600)

    # Make a split window plot.  First plot is time vs. temperature.
    fig = plt.figure(figsize = (15, 6))
    ax1 = plt.subplot(1, 2, 1)
    elapsed_time = (time - time[0])/60
    ax1.plot(elapsed_time[idx], temp[idx], 'b-')
    ax1.set_xlabel('Elapsed Time (minutes)')
    ax1.set_ylabel(units)

    basename = os.path.basename(FILE_NAME)
    ax1.set_title('{0}\n{1}'.format(basename, longname))

    # The 2nd plot is the trajectory.
    # Use a north polar azimuthal equal area projection.
    ax2 = plt.subplot(1, 2, 2)
    m = Basemap(projection='nplaea', resolution='l',
                boundinglat=52, lon_0=0)
    m.drawcoastlines(linewidth=0.5)
    m.drawparallels(np.arange(0., 91., 10.), labels=[0, 0, 0, 1])
    m.drawmeridians(np.arange(-180, 180., 30.), labels=[0, 1, 0, 0])
    m.plot(longitude[idx], latitude[idx], linestyle='None', marker='.',
           color='blue', latlon=True)
    plt.title('Trajectory of Flight Path')

    fig = plt.gcf()
    # plt.show()
    pngfile = "{0}.a.py.png".format(basename)
    fig.savefig(pngfile)

if __name__ == "__main__":

    # If a certain environment variable is set, look there for the input
    # file, otherwise look in the current directory.
    hdffile = 'GLAH13_633_2103_001_1317_0_01_0001.h5'

    try:
        hdffile = os.path.join(os.environ['HDFEOS_ZOO_DIR'], hdffile)
    except KeyError:
        pass

    run(hdffile)


import os
from pyBigstick.nucleus import Nucleus
import streamlit as st
import numpy as np
import plotly.express as px
from barChartPlotly import plotly_barcharts_3d
from PIL import Image

he4_image = Image.open('assets/he4.png')
nucl_image = Image.open('assets/nucl_symbol.png')
table_image = Image.open('assets/table.jpg')
scattering_image = Image.open('assets/scattering.jpeg')
deexcitation_image = Image.open('assets/deexcitation.png')
lvl_image = Image.open('assets/Energy_levels.png')
shells_image = Image.open('assets/shells.png')

bs = os.getcwd() +'/src/bigstick.x'

header_container = st.container()
intro_container = st.container()
bs_container = st.container()
states_container = st.container()
densities_container = st.container()

hide_table_row_index = """
            <style>
            tbody th {display:none}
            .blank {display:none}
            </style>
            """

light_nuclei = ['F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl']

with header_container:
    st.title('pyBigstick')
    st.markdown("""This streamlit app visualizes the nuclear transitions calculated by [pyBigstick](https://github.com/noctildon/pyBigstick),
        including the energy levels and the density matrices.""")

with intro_container:
    st.subheader("Basic knowledge about nuclear physics")
    st.markdown('Physicists usually use a symbol + 2 number to represent a unique nucleus')
    st.image(nucl_image, width=500)
    st.markdown('For example, the following is identical to He (Helium) with mass number 4 and atomic number 2, or equivalently 2 protons and 2 neutrons.')
    st.image(he4_image,width=300)
    st.markdown('Atomic number can be determined by element symbol uniquely, so sometimes it is skipped and ignored.')

    st.text('And this is the well-known periodic table')
    st.image(table_image,width=800)

    st.markdown('Experimentalists use neutrinos (an extremely small and light particle) to hit the nucleus. This process is called "scattering".')
    st.image(scattering_image,width=800)
    st.markdown("""Before scattering the nucleus has lowest possbile energy (ground state). After scattering nucleus gain some energy from the neutrinos,
        being called "excited nucleus" or "excited state". Then there is a chance that the excited nucleus would drop back to the ground state by emitting gamma ray.
    """)

    col1, col2 = st.columns(2)
    with col1:
        st.image(deexcitation_image,width=400)
    with col2:
        st.image(lvl_image,width=400)

    st.markdown("""What happen in the scattering is that some of the nucleons get excited to the orbit with high energy.
        The core algorithm of pyBigstick is to iterate all possible combinations and transitions of the nucleons.
        And the density matrices describe how nucleons move among the orbits by talking us a probability-like value.
    """)
    st.image(shells_image,width=700)

with bs_container:
    st.subheader("Control panel")
    st.markdown("""Input the info of the interested nucleus, eg. F19, Be10.
        Not sure which nucleus to pick? check out [this](https://periodictable.com/Isotopes/009.19/index.html).
        (Not all of the nucleus is possible to calculate).""")

    col1_1, col1_2 = st.columns(2)
    with col1_1:
        s1 = st.selectbox('The nucleus to calculate', light_nuclei, index=0)
    with col1_2:
        s2 = st.selectbox('Mass number', range(9,41), index=10)

    col2_1, col2_2 = st.columns(2)
    with col2_1:
        n_states = st.selectbox('Number of states to calculate (more states always need more time)', range(1,7), index=2)
    with col2_2:
        maxiter = st.selectbox('Number of iteration (higher iteration is more accurate on results, but takes longer)', np.arange(50,510,10), index=5)

    s1 = s1.lower()
    nucl_name = f'{s1}{s2}'
    st.text(f'Calculate {nucl_name}...')
    nu = Nucleus(nucl_name, n_states=n_states, maxiter=maxiter)

    if st.button('Clean the previous result and rerun'):
        st.write(f'Successfully clean nucleus {nucl_name}. Running {nucl_name} again...')
        nu.clean()

    if not nu.check():
        nu.script_save()
        nu.prepare()
        nu.run(bs=bs)

    nu.save_results()

with states_container:
    st.subheader('Energy level states')
    st.markdown("""When the scattering happens to a nucleus, the nucleus could be excited to higher state.
        In general, initially the nucleus is in ground state (the state with the lowest energy).
        Then after scattering, it is excited to some higher state with energy higher than ground state.
        We also label the state with n. Ground state has n=1. First excited state has n=2. Second excited has n=3, and so on.""")

    fig = px.bar(nu.states, x='state', y='Ex',
        labels={
            'state': 'State',
            'Ex': 'Excitation energy (MeV)'
    })
    st.plotly_chart(fig, use_container_width=True)

    if st.checkbox('Show all states data'):
        st.text('States')
        st.write(nu.states)

with densities_container:
    st.subheader('Density matrices')
    st.markdown("""The amp (transition amplitdue) in the last column below is (to some degree) proportional to the probability that
        a nucleon moves from one orbit to another, given the condition that the nucleus jumps from one state to another (say from n=1 to n=2).
        Jt and Tt are the spin and isospin of the transition, respectively. They are the attributes of a transition.
        A transition could have multiple values of Jt. Tt can be either 0 or 1. Most of the amp is zero.""")

    col1, col2, col3, col4 = st.columns(4)
    with col1:
        statei = st.selectbox('Initial state', nu.states['state'])
    with col2:
        statej = st.selectbox('Final state', nu.states['state'])
    with col3:
        Jt = st.selectbox('Jt', np.unique(nu.densities['Jt']))
    with col4:
        Tt = st.selectbox('Tt', [0,1])

    filter_densities = nu.densities.loc[(nu.densities['statei']==statei) & (nu.densities['statej']==statej) &\
        (nu.densities['Jt']==Jt) & (nu.densities['Tt']==Tt)]

    st.subheader('Non-zero elements')
    st.markdown(hide_table_row_index, unsafe_allow_html=True)
    st.table(filter_densities)

    st.subheader('3D plot of the density matrices')
    st.text('The plot only shows non-zero elements.')

    if not filter_densities.empty:
        fig = plotly_barcharts_3d(filter_densities['orba'], filter_densities['orbb'], filter_densities['amp'],
                x_title='orbit a', y_title='orbit b', z_title='amp')
        fig.update_layout(width=700, height=700, yaxis = dict(scaleanchor = 'x'))
        st.plotly_chart(fig, use_container_width=True)
    else:
        st.text('All elements are zero, so the plot is skipped.')

    if st.checkbox('Show all raw densities data'):
        st.text('Density matrices')
        st.write(nu.densities)


"""
ASPIRE-Python Introduction
==========================

In this notebook we will introduce some code from ASPIRE-Python
that corresponds to topics from MATH586.
"""

# %%
# Imports
# -------
# First we import some of the usual suspects. In addition, we import some classes from
# the ASPIRE package that we will use throughout this tutorial.

# %%
# Homework Task 0
# ^^^^^^^^^^^^^^^
#
# Attempt to install ASPIRE on your machine. ASPIRE can generally install on
# Linux, Mac, and Windows
# under Anaconda Python, by following the instructions in the README.
# `The instructions for developers is the most comprehensive
# <https://github.com/ComputationalCryoEM/ASPIRE-Python/blob/master/README.md#for-developers>`_.
# Linux is the most tested platform.
#
# ASPIRE requires some resources to run, so if you wouldn't run typical data
# science codes on your machine (maybe a netbook for example),
# you may use TigerCPU. After logging into TigerCPU,
# ``module load anaconda3/2020.7`` and follow the anaconda instructions for
# developers in the link above.
# Those instructions should create a working environment for tinkering with
# ASPIRE code found in this notebook.

import logging
import os

import matplotlib.pyplot as plt
import mrcfile
import numpy as np

from aspire.abinitio import CLSyncVoting
from aspire.noise import AnisotropicNoiseEstimator, WhiteNoiseEstimator
from aspire.operators import FunctionFilter, RadialCTFFilter, ScalarFilter
from aspire.source import RelionSource, Simulation
from aspire.utils import (
    Rotation,
    get_aligned_rotations,
    get_rots_mse,
    register_rotations,
)
from aspire.volume import Volume

logger = logging.getLogger(__name__)

# %%
# ``Image`` Class
# ---------------
#
# The `Image <https://computationalcryoem.github.io/ASPIRE-Python/aspire.image.html#aspire.image.image.Image>`_ class
# is a thin wrapper over numpy arrays for a stack containing 1 or more images.
# In this notebook we won't be working directly with the ``Image`` class a lot, but it will be one of the fundemental structures behind the scenes.
# A lot of ASPIRE code passes around ``Image`` and ``Volume`` classes.
#
# Examples of using the Image class can be found in:
#
# - ``gallery/tutorials/basic_image_array.py``
#
# - ``gallery/tutorials/image_class.py``

# %%
# ``Volume`` Class
# ----------------
#
# Like ``Image``, the `Volume <https://computationalcryoem.github.io/ASPIRE-Python/aspire.volume.html#aspire.volume.Volume>`_ class
# is a thin wrapper over numpy arrays that provides specialized methods for a stack containing 1 or more volumes.
#
# Here we will instantiate a Volume using a numpy array and use it to downsample to a desired resolution (64 should be good).
# For the data source I chose to download a real volume density map from `EMDB <https://www.ebi.ac.uk/pdbe/entry/emdb/EMD-2660>`_.
# The download was uncompressed in my local directory.  The notebook defaults to a small low resolution sample file you may use to sanity check.
# Unfortunately real data can be quite large so we do not ship it with the repo.

# %%
# Homework Task 1
# ^^^^^^^^^^^^^^^
#
# - Starting from `EMPIAR <https://www.ebi.ac.uk/pdbe/emdb/empiar/>`_ find a molecule of interest and try to
# find if it has a corresponding volume density map from `EMDB <https://www.ebi.ac.uk/pdbe/emdb>`_.
#
# - Download such a map and use it in the following experiments where I have used 2660.
#
# Helpful friendly hint:
# `mrcfile` will typically open `.map` files provided by EMDB, corresponding to an EMPIAR entry.
# This was not obvious to me, but you may `read more about the format here <https://www.emdataresource.org/mapformat.html>`_.

# %%
# Initialize Volume
# -----------------

# A low res example file is included in the repo as a sanity check.
DATA_DIR = "data"
infile = mrcfile.open(os.path.join(DATA_DIR, "clean70SRibosome_vol_65p.mrc"))

# More interesting data requires downloading locally.
# infile = mrcfile.open("EMD-2660/map/emd_2660.map")

d = infile.data
logger.info(f"map data shape: {d.shape} dtype:{d.dtype}")
v = Volume(d)

# Downsample the volume to a desired resolution
img_size = 64
# Volume.downsample() Returns a new Volume instance.
#   We will use this lower resolution volume later.
v2 = v.downsample((img_size, img_size, img_size))
L = v2.resolution

# %%
# Contour Plot of Data
# --------------------

# Alternatively, for quick sanity checking purposes we can view as a contour plot.
#   We'll use three orthographic projections, one per axis.
for axis in range(3):
    plt.contourf(np.arange(L), np.arange(L), np.sum(v2[0], axis=axis), cmap="gray")
    plt.show()

# %%
# Scatter Plot
# ------------

# We can attempt a 3d scatter plot, but the results aren't very good.
x, y, z = np.meshgrid(np.arange(L), np.arange(L), np.arange(L))
ax = plt.axes(projection="3d")
ax.scatter3D(x, y, z, c=np.log10(v2.flatten()), cmap="Greys_r")
plt.show()

# %%
# Homework Task 2
# ^^^^^^^^^^^^^^^
#
# Above I have used a simple log transform with a scatter plot to peek at the 3D data.
# This was mainly just to make sure the data was in the neighborhood of what I was looking for.
# More commonly we will want to construct an ``isosurface`` plot.
# Try to create a better plot of the volume (this will probably require more advanced tools than matplotlib).

# %%
# ``Rotation`` Class - Generating Random Rotations
# ------------------------------------------------
#
# To get general projections this brings us to generating random rotations which we will apply to our volume.
#
# While you may bring your own 3x3 matrices or generate manually (say from your own Euler angles),
# ASPIRE has a `Rotation class <https://computationalcryoem.github.io/ASPIRE-Python/aspire.utils.html#module-aspire.utils.rotation>`_
# which can do this random rotation generation for us.  It also has some other utility methods if you would want to compare with something manual.
#
# The following code will generate some random rotations, and use the ``Volume.project()`` method to return an ``Image`` instance representing the stack of projections.
# We can display projection images using the ``Image.show()`` method.

num_rotations = 2
rots = Rotation.generate_random_rotations(n=num_rotations, seed=12345)

# We can access the numpy array holding the actual stack of 3x3 matrices:
logger.info(rots)
logger.info(rots.matrices)

# Using the first (and in this case, only) volume, compute projections using the stack of rotations:
projections = v.project(0, rots)

# project() returns an Image instance.
logger.info(projections)
projections.show()
# Neat, we've generated random projections of some real data.

# %%
# The ``source`` Package
# ----------------------
#
# `aspire.source <https://computationalcryoem.github.io/ASPIRE-Python/aspire.source.html#module-aspire.source.simulation>`_
# package contains a collection of data source interfaces.
# The idea is that we can design an experiment using a synthetic ``Simulation`` source or our own provided array via ``ArrayImageSource``;
# then later swap out the source for a large experimental data set using something like ``RelionSource``.
#
# We do this because the experimental datasets are too large to fit in memory.
# They cannot be provided as a massive large array, and instead require methods to orchestrate batching.
# Depending on the application, they may also require corresponding batched algorithms.
# The ``Source`` classes try to make most of this opaque to an end user.  Ideally we can swap one source for another.
#
# For now we will build up to the creation and application of synthetic data set based on the real volume data used previously.

# %%
# ``Simulation`` Class
# --------------------
#
# Generating realistic synthetic data sources is a common task.
# The process of generating then projecting random rotations is integrated into the
# `Simulation <https://computationalcryoem.github.io/ASPIRE-Python/aspire.source.html#module-aspire.source.simulation>`_ class.
# Using ``Simulation``, we can generate arbitrary numbers of projections for use in experiments.
# Later we will demonstrate additional features which allow us to create more realistic data sources.

num_imgs = 100  # How many images in our source.
# Generate a Simulation instance based on the original volume data.
sim = Simulation(L=v.resolution, n=num_imgs, vols=v)
# Display the first 10 images
sim.images(0, 10).show()  # Hi Res

# Repeat for the lower resolution (downsampled) volume v2.
sim2 = Simulation(L=v2.resolution, n=num_imgs, vols=v2)
sim2.images(0, 10).show()  # Lo Res

# Note both of those simulations have the same rotations
#   because they had the same seed by default,
# We can set our own seed to get a different random samples (of rotations).
sim_seed = Simulation(L=v.resolution, n=num_imgs, vols=v, seed=42)
sim_seed.images(0, 10).show()

# We can also view the rotations used to create these projections
# logger.info(sim2.rots)  # Commented due to long output

# %%
# Simulation with Noise - Filters
# -------------------------------
#
# Filters
# ^^^^^^^
#
# `Filters <https://computationalcryoem.github.io/ASPIRE-Python/aspire.operators.html#module-aspire.operators.filters>`_
# are a collection of classes which once configured can be applied to ``Source`` pipelines.
# Common filters we might use are ``ScalarFilter``, ``PowerFilter``, ``FunctionFilter``, and ``CTFFilter``.
#
# Adding to Simulation
# ^^^^^^^^^^^^^^^^^^^^
#
# We can customize Sources by adding stages to their generation pipeline.
# In this case of a Simulation source, we want to corrupt the projection images with significant noise.
#
# First we create a constant two dimension filter (constant value set to our desired noise variance).
# Then when used in the ``noise_filter``, this scalar will be multiplied by a random sample.
# Similar to before, if you require a different sample, this would be controlled via a ``seed``.

# Using the sample variance, we'll compute a target noise variance
var = np.var(sim2.images(0, sim2.n).asnumpy())
logger.info(f"sim2 clean sample var {var}")
noise_variance = 100.0 * var
logger.info(f"noise var {noise_variance}")

# Then create a constant filter based on that variance
white_noise_filter = ScalarFilter(dim=2, value=noise_variance)
# We can create a similar simulation with this additional noise_filter argument:
sim3 = Simulation(L=v2.resolution, n=num_imgs, vols=v2, noise_filter=white_noise_filter)
sim3.images(0, 10).show()
# These should be rather noisy now ...

# %%
# Common Line Estimation
# ----------------------
#
# Now we can create a CL instance for estimating orientation of projections using the Common Line with synchronization method.
#
# We will import `CLSyncVoting <(https://computationalcryoem.github.io/ASPIRE-Python/aspire.abinitio.html?highlight=clsyncvoting#aspire.abinitio.commonline_sync.CLSyncVoting>`_,
# then several helper utilities fron the ``coor_trans`` package to help verify our estimates.
#
# For each iteration in the loop:
# - Save the true rotations
# - Compute orientation estimate using CLSyncVoting method
# - Compare the estimated vs true rotations
#
# Each iteration will logger.info some diagnostic information that contains the top eigenvalues found.
# From class we learned that a healthy eigendistribution should have a significant gap after the third eigenvalue.
# It is clear we have such eigenspacing for the clean images, but not for the noisy images.

for desc, _sim in [
    ("High Res", sim),
    ("Downsampled", sim2),
    ("Downsampled with Noise", sim3),
]:
    logger.info(desc)
    true_rotations = _sim.rots  # for later comparison

    orient_est = CLSyncVoting(_sim, n_theta=36)
    # Get the estimated rotations
    orient_est.estimate_rotations()
    rots_est = orient_est.rotations

    # Compare with known true rotations
    Q_mat, flag = register_rotations(rots_est, true_rotations)
    regrot = get_aligned_rotations(rots_est, Q_mat, flag)
    mse_reg = get_rots_mse(regrot, true_rotations)
    logger.info(
        f"MSE deviation of the estimated rotations using register_rotations : {mse_reg}\n"
    )

# %%
# Homework Task 3
# ^^^^^^^^^^^^^^^
#
# We confirmed a dramatic change in the eigenspacing when we add a lot of noise.
# Compute the SNR in this case using the formula described from class.
# Repeat the experiment with varying levels of SNR to find at what level the character of the eigenspacing changes.
# This will require changing the Simulation Source's noise_filter.
# How does this compare with the levels discussed in lecture?

# %%
# More Advanced Noise - Whitening
# -------------------------------
#
# We can estimate the noise across the stack of images
#
# The ``noise`` Package
# ^^^^^^^^^^^^^^^^^^^^^
#
# The `aspire.noise <https://computationalcryoem.github.io/ASPIRE-Python/aspire.noise.html>`_
# package contains several useful classes for generating and estimating different types of noise.
#
# In this case, we know the noise to be white, so we can proceed directly to
# `WhiteNoiseEstimator <https://computationalcryoem.github.io/ASPIRE-Python/aspire.noise.html#aspire.noise.noise.WhiteNoiseEstimator>`_.  The noise estimators consume from a ``Source``.
#
# The white noise estimator should log a diagnostic variance value. How does this compare with the known noise variance above?

# Create another Simulation source to tinker with.
sim_wht = Simulation(
    L=v2.resolution, n=num_imgs, vols=v2, noise_filter=white_noise_filter
)

# Estimate the white noise.
noise_estimator = WhiteNoiseEstimator(sim_wht)

# %%
# A Custom ``FunctionFilter``
# ---------------------------
#
# We will now apply some more interesting noise, using a custom function, and then apply a ``whitening`` process to our data.
#
# Using ``FunctionFilter`` we can create our own custom functions to apply in a pipeline.
# Here we want to apply a custom filter as a noise adder.  We can use a function of two variables for example.

def noise_function(x, y):
    return 1e-7 * np.exp(-(x * x + y * y) / (2 * 0.3**2))

# In python, functions are first class objects.
# We take advantage of that to pass this function around as a variable.
# It will be evaluated later...
custom_noise_filter = FunctionFilter(noise_function)

# Create yet another Simulation source to tinker with.
sim4 = Simulation(
    L=v2.resolution, n=num_imgs, vols=v2, noise_filter=custom_noise_filter
)
sim4.images(0, 10).show()

# %%
# Noise Whitening
# ---------------
#
# Applying the ``Simulation.whiten()`` method just requires passing the filter corresponding to the estimated noise instance.
# Then we can inspect some of the whitened images.  While noise is still present, we can see a dramatic change.

# Estimate noise.
aiso_noise_estimator = AnisotropicNoiseEstimator(sim4)

# Whiten based on the estimated noise
sim4.whiten(aiso_noise_estimator.filter)

# What do the whitened images look like...
sim4.images(0, 10).show()

# %%
# Homework Task 4
# ^^^^^^^^^^^^^^^
#
# Try some other image preprocessing methods exposed by the ``Simulation``/``ImageSource`` classes.
#
# Try some other custom function to add noise or other corruptions to the images.

# %%
# Real Experimental Data - ``RelionSource``
# -----------------------------------------
#
# Now that we know our experiment code seems to run,
# we can try to replace the simulation with a real experimental data source.
#
# Lets attempt the same CL experiment, but with a ``RelionSource``.

src = RelionSource(
    "data/sample_relion_data.star",
    data_folder="",
    pixel_size=5.0,
    max_rows=1024,
)

# Data resides on Tiger Cluster
# Please make sure you are using a compute node once you've installed ASPIRE, not the head node...
# src =  RelionSource(
#    "/tigress/gbwright/data/cryo-em/CryoEMdata/empiar10028/shiny_2sets.star", data_folder="", pixel_size=5.0, max_rows=100
# )
src.downsample(img_size)

src.images(0, 10).show()

noise_estimator = WhiteNoiseEstimator(src)
src.whiten(noise_estimator.filter)

orient_est = CLSyncVoting(src, n_theta=36)
orient_est.estimate_rotations()
rots_est = orient_est.rotations

# %%
# We can see that the code can easily run with experimental data by subsituting the ``Source`` class.
# However, we have hit the practical limitation that requires class averaging of images....

# %%
# CTF Filter
# ----------
#
# Here we can use the ``RadialCTFFilter`` subclass of
# `CTFFilter <https://computationalcryoem.github.io/ASPIRE-Python/aspire.operators.html?highlight=ctffilter#aspire.operators.filters.CTFFilter>`_
# to generate some simulated images with CTF effects.
#
# We use the ``unique_filter`` argument of the ``Simulation`` class to apply a collection of several CTFs with different defocus.
# The defocus values are generated from the ``np.linspace`` method. We end up with a list of filters.
#
# By combining CTFFilters, noise, and other filters ASPIRE can generate repeatable rich data sets with controlled parameters.
# The ``Simulation`` class will attempt to apply transforms ``on the fly`` to batches of our images, allowing us to generate arbitrarily long stacks of data.

# Specify the CTF parameters not used for this example
# but necessary for initializing the simulation object
pixel_size = 5  # Pixel size of the images (in angstroms)
voltage = 200  # Voltage (in KV)
defocus_min = 1.5e4  # Minimum defocus value (in angstroms)
defocus_max = 2.5e4  # Maximum defocus value (in angstroms)
defocus_ct = 7  # Number of defocus groups.
Cs = 2.0  # Spherical aberration
alpha = 0.1  # Amplitude contrast

# Initialize simulation object with CTF filters.
# Create CTF filters
filters = [
    RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
    for d in np.linspace(defocus_min, defocus_max, defocus_ct)
]

sim5 = Simulation(L=v2.resolution, n=num_imgs, vols=v2, unique_filters=filters)
sim5.images(0, 10).show()

# Here we will combine CTF and noise features to our projections.
sim6 = Simulation(
    L=v2.resolution,
    n=num_imgs,
    vols=v2,
    unique_filters=filters,
    noise_filter=custom_noise_filter,
)
sim6.images(0, 10).show()

# Estimate noise.
aiso_noise_estimator = AnisotropicNoiseEstimator(sim6)

# Whiten based on the estimated noise
sim6.whiten(aiso_noise_estimator.filter)
sim6.images(0, 10).show()


import pandas as pd
import numpy as np

def load_and_proccess(rawData):
    rawData.columns = ['Age','Workclass','fnlwgt','Education','Education-Num','Marital-Status','Occupation','Realtionship','Race','Sex','Capital-Gain','Capital-Loss','Hours-Per-Week','Native-Country','Income']
    return rawData.drop_duplicates()

def describedf(data):
    return data.describe()

def nuniquedf(data):
    return data.nunique(axis=0)

def columnsdf(data):
    return data.columns

def shapedf(data):
    return data.shape


# Use an import statement to import statsmodels
import statsmodels

# Import statsmodels under the alias sm
import statsmodels as sm

# Use an import statement to import seaborn with alias sns
import seaborn as sns

# Fix the import of numpy to run without errors
import numpy as np


import argparse
import contextlib
import glob
import io
import math
import numbers
import os
import pickle
import platform
import re
import sys
import types

import colorama
import nbformat
import numpy as np
import pkg_resources
from nbconvert import PythonExporter

def print_warning(message):
    print(colorama.Fore.YELLOW + colorama.Style.BRIGHT + 'WARNING: ' + message + colorama.Style.RESET_ALL)

def check_var(stud_val, sol_val, var_name_style):
    def unify_vector(vec):
        # Remove single dimensions
        vec = np.squeeze(vec)

        # Convert column or row matrices to vectors
        if len(vec.shape) == 2 and (vec.shape[0] == 1 or vec.shape[1] == 1):
            return vec.reshape(-1)
        else:
            return vec

    if type(stud_val) != type(sol_val):
        print_warning('The variable %s has the type %s but the type %s is expected. Even though it is possible that the other checks still work, it can also lead to errors.' % (var_name_style, type(stud_val), type(sol_val)))

    if type(stud_val) == list:
        assert len(stud_val) == len(sol_val), 'The list %s does not have the correct length\nGot: %d\nExpected: %d' % (var_name_style, len(stud_val), len(sol_val))

        if any([isinstance(val, np.ndarray) or type(val) == list for val in stud_val]):
            # Check nested lists recursively
            for v1, v2 in zip(stud_val, sol_val):
                check_var(v1, v2, var_name_style)
            return
        elif all([isinstance(val, numbers.Number) for val in stud_val]):
            # The list contains only numbers --> unify to numpy arrays
            stud_val = np.asarray(stud_val)
            sol_val = np.asarray(sol_val)
        else:
            assert stud_val == sol_val, 'The content of the list %s does not match with the solution\nGot: %s\nExpected: %s' % (var_name_style, stud_val, sol_val)
            return

    if isinstance(stud_val, np.ndarray):
        stud_val = unify_vector(stud_val)
        sol_val = unify_vector(sol_val)

        if stud_val.dtype != sol_val.dtype:
            print_warning('The type of your array %s has the type %s but the type %s is expected. Even though it is possible that the other checks still work, it can also lead to errors.' % (var_name_style, stud_val.dtype, sol_val.dtype))

        assert stud_val.shape == sol_val.shape, 'Your array %s does not have the correct dimensions\nGot: %s\nExpected: %s' % (var_name_style, stud_val.shape, sol_val.shape)

    if isinstance(stud_val, np.ndarray) or isinstance(stud_val, (np.floating, float)):
        assert np.allclose(stud_val, sol_val, rtol=1e-04, atol=1e-05), 'The content of the array %s does not match with the solution (tolerance considered)\nGot: %s\nExpected: %s' % (var_name_style, stud_val, sol_val)
    else:
        assert stud_val == sol_val, 'The variable %s is not set to the correct value\nGot: %s\nExpected: %s' % (var_name_style, stud_val, sol_val)

def check_vars(stud_vars, sol_vars):
    def function_name(function):
        # Returns the name of a function object
        match = re.search(r'function\s+([^\s]+)\s+at', str(function))
        if match:
            return match.group(1)
        else:
            return ''

    # Compare all variables
    error_messages = []
    for var_name in sol_vars.keys():
        if var_name == 'sol_vars_code':
            continue

        try:
            var_name_style = colorama.Style.NORMAL + var_name + colorama.Style.BRIGHT
            assert var_name in stud_vars, 'You do not have a variable with the name %s in your script.' % var_name_style

            stud_val = stud_vars[var_name]
            sol_val = sol_vars[var_name]

            if isinstance(stud_val, types.FunctionType):
                # For functions, compare only the names
                stud_val = function_name(stud_val)

            check_var(stud_val, sol_val, var_name_style)
        except AssertionError as failure:
            error_messages.append(colorama.Fore.RED + colorama.Style.BRIGHT + str(failure) + colorama.Style.RESET_ALL)
            continue

    # Compare all code snippets
    if 'sol_vars_code' in sol_vars:
        for var_code in sol_vars['sol_vars_code']:
            try:
                var_name_style = colorama.Style.NORMAL + var_code['name'] + colorama.Style.BRIGHT

                stud_val = eval(var_code['code'], stud_vars)    # Runs the code in the student's environment
                sol_val = var_code['result']

                check_var(stud_val, sol_val, var_name_style)
            except AssertionError as failure:
                error_messages.append(colorama.Fore.RED + colorama.Style.BRIGHT + str(failure) + '\nDescription: ' + var_code['description'] + '\nNote: This is a generated variable.' + colorama.Style.RESET_ALL)
                continue

    return error_messages

def check_notebook(filename, all_sol_vars):
    """
    Compares the variables of a Jupyter notebook to a reference solution.

    :param filename: of the notebook to check.
    :param all_sol_vars: dict which contains the reference solutions.
    :return: True when all checks were successful or no checks could be performed.
    """
    print('Testing notebook ' + filename)

    # Find all solutions for the current notebook (usually more than one for TensorFlow code)
    solutions = sorted([solution for solution in all_sol_vars.keys() if solution.startswith(filename)])

    if not solutions:
        print('No matching solution for the file %s found. The notebook is skipped.\n' % filename)
        return True

    # Load the notebook file
    with open(filename) as file:
        nb = nbformat.read(file, as_version=4)

    # Keep only the code cells (and especially remove the raw cells)
    nb.cells[:] = [cell for cell in nb.cells if cell.cell_type == 'code']

    # Convert the notebook to a Python script
    exporter = PythonExporter()
    source, meta = exporter.from_notebook_node(nb)
    source = re.sub('(.*?)get_ipython', r'#\1get_ipython', source)  # Comment out code lines which are only available in ipython

    # Run the student's solution
    stud_vars = {}
    with open(os.devnull, "w") as stream, contextlib.redirect_stdout(stream):
        # print commands are prevented (warnings will still be shown)
        exec(source, stud_vars)

    # Test each solution
    messages = {}
    correct_solution = ''
    for solution in solutions:
        messages[solution] = check_vars(stud_vars, all_sol_vars[solution])

        if not messages[solution]:
            correct_solution = solution

            # The first check was already successful; no need to check other possible solutions
            break

    if correct_solution:
        print(f'{colorama.Fore.GREEN}The test for the notebook %s was successful (checked against %s). No errors found.{colorama.Style.RESET_ALL}\n' % (filename, correct_solution))
        return True
    else:
        # All solutions are incorrect. Show the errors for the first solution (arbitrary)
        print('The test for the notebook %s was not successful. At least one variable does not contain the expected result (errors compared to the solution %s are shown).' % (filename, solutions[0]))

        for message in messages[solutions[0]]:
            print(message)

        return False

def version_info():
    packages = ['numpy', 'scikit-learn', 'tensorflow']
    versions = {}
    for package in packages:
        versions[package] = pkg_resources.get_distribution(package).version

    versions['python'] = platform.python_version()

    return versions

def check_folder():
    folder_name = os.path.basename(os.getcwd())
    print('Testing all notebooks in the folder ' + folder_name)

    solution_filename = 'solution_vars.pickle'
    if not os.path.exists(solution_filename):
        print('No pickle file containing the solution data found. The folder %s is skipped.' % folder_name)
        return

    # Load the official solution data
    all_sol_vars = pickle.load(open(solution_filename, 'rb'))

    # First, check whether the same packages were used
    for package, version in version_info().items():
        if version != all_sol_vars['versions'][package]:
            print_warning('You are using %s in version %s but the solution variables were created with version %s. Results may differ.' % (package, version, all_sol_vars['versions'][package]))

    # All notebooks in the current folder
    notebook_filenames = glob.glob('*.ipynb')

    # Test each notebook where a corresponding solution is available
    checks = []
    for filename in notebook_filenames:
        successful = check_notebook(filename, all_sol_vars)
        checks.append(successful)

    # Find all available solution names
    solutions = set()
    for sol in all_sol_vars.keys():
        match = re.search(r'^[^.]+\.ipynb', sol)
        if match:
            solutions.add(match.group(0))

    # Check if the students provided a notebook for each solution
    all_sols_checked = True
    for sol in solutions:
        if sol not in notebook_filenames:
            print_warning(f'A notebook with the name {colorama.Style.NORMAL}%s{colorama.Style.BRIGHT} could not be found in the folder %s but is part of this assignment. You have either not finished your implementation yet or you used the wrong name for your notebook.' % (sol, folder_name))
            all_sols_checked = False

    if all(checks) and all_sols_checked:
        print(f'{colorama.Fore.GREEN}Congratulations! The checked notebooks in the folder %s match with the reference implementation.{colorama.Style.RESET_ALL}\n' % folder_name)

colorama.init()
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                 description='This script tests all notebooks in one or more folders. Requirement is that the respective folder contains a file with the name solution_vars.pickle which contains the solution variables of the notebooks.')
parser.add_argument('folders', type=str, nargs='*', default='.', help='Folders to check (each with its own solution file). If no folder name is given, the current folder is used.')
__version__ = '0.7.3'
parser.add_argument('--version', action='version', version='%(prog)s ' + __version__)

args = parser.parse_args()
print(os.path.basename(__file__) + ' ' + __version__)

# Check all notebooks in the given folder name(s)
for folder in args.folders:
    prev_cwd = os.getcwd()
    os.chdir(folder)  # Notebooks should be executed in their corresponding folder

    check_folder()

    os.chdir(prev_cwd)


# This code is part of Qiskit.
#
# (C) Copyright IBM 2021.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""
A Tester for the RB utils module
"""

import numpy as np
from ddt import ddt, data, unpack
from qiskit.test import QiskitTestCase
from qiskit import QuantumCircuit
from qiskit.circuit.library import (
    IGate,
    XGate,
    YGate,
    ZGate,
    HGate,
    SGate,
    SdgGate,
    CXGate,
    CZGate,
    SwapGate,
)
import qiskit_experiments.library.randomized_benchmarking as rb

@ddt
class TestRBUtilities(QiskitTestCase):
    """
    A test class for additional functionality provided by the StandardRB
    class.
    """

    instructions = {
        "i": IGate(),
        "x": XGate(),
        "y": YGate(),
        "z": ZGate(),
        "h": HGate(),
        "s": SGate(),
        "sdg": SdgGate(),
        "cx": CXGate(),
        "cz": CZGate(),
        "swap": SwapGate(),
    }
    seed = 42

    @data(
        [1, {((0,), "x"): 3, ((0,), "y"): 2, ((0,), "h"): 1}],
        [5, {((1,), "x"): 3, ((4,), "y"): 2, ((1,), "h"): 1, ((1, 4), "cx"): 7}],
    )
    @unpack
    def test_count_ops(self, num_qubits, expected_counts):
        """Testing the count_ops utility function
        this function receives a circuit and counts the number of gates
        in it, counting gates for different qubits separately"""
        circuit = QuantumCircuit(num_qubits)
        gates_to_add = []
        for gate, count in expected_counts.items():
            gates_to_add += [gate for _ in range(count)]
        rng = np.random.default_rng(self.seed)
        rng.shuffle(gates_to_add)
        for qubits, gate in gates_to_add:
            circuit.append(self.instructions[gate], qubits)
        counts = rb.RBUtils.count_ops(circuit)
        self.assertDictEqual(expected_counts, counts)

    def test_calculate_1q_epg(self):
        """Testing the calculation of 1 qubit error per gate
        The EPG is computed based on the error per clifford determined
        in the RB experiment, the gate counts, and an estimate about the
        relations between the errors of different gate types
        """
        epc_1_qubit = 0.0037
        qubits = [0]
        gate_error_ratio = {((0,), "id"): 1, ((0,), "rz"): 0, ((0,), "sx"): 1, ((0,), "x"): 1}
        gates_per_clifford = {((0,), "rz"): 10.5, ((0,), "sx"): 8.15, ((0,), "x"): 0.25}
        epg = rb.RBUtils.calculate_1q_epg(epc_1_qubit, qubits, gate_error_ratio, gates_per_clifford)
        error_dict = {
            ((0,), "rz"): 0,
            ((0,), "sx"): 0.0004432101747785104,
            ((0,), "x"): 0.0004432101747785104,
        }

        for gate in ["x", "sx", "rz"]:
            expected_epg = error_dict[((0,), gate)]
            actual_epg = epg[0][gate]
            self.assertTrue(np.allclose(expected_epg, actual_epg, atol=0.001))

    def test_calculate_2q_epg(self):
        """Testing the calculation of 2 qubit error per gate
        The EPG is computed based on the error per clifford determined
        in the RB experiment, the gate counts, and an estimate about the
        relations between the errors of different gate types
        """
        epc_2_qubit = 0.034184849962675984
        qubits = [1, 4]
        gate_error_ratio = {
            ((1,), "id"): 1,
            ((4,), "id"): 1,
            ((1,), "rz"): 0,
            ((4,), "rz"): 0,
            ((1,), "sx"): 1,
            ((4,), "sx"): 1,
            ((1,), "x"): 1,
            ((4,), "x"): 1,
            ((4, 1), "cx"): 1,
            ((1, 4), "cx"): 1,
        }
        gates_per_clifford = {
            ((1, 4), "barrier"): 1.032967032967033,
            ((1,), "rz"): 15.932967032967033,
            ((1,), "sx"): 12.382417582417583,
            ((4,), "rz"): 18.681946624803768,
            ((4,), "sx"): 14.522605965463109,
            ((1, 4), "cx"): 1.0246506515936569,
            ((4, 1), "cx"): 0.5212064090480678,
            ((4,), "x"): 0.24237661112857592,
            ((1,), "measure"): 0.01098901098901099,
            ((4,), "measure"): 0.01098901098901099,
            ((1,), "x"): 0.2525918944392083,
        }
        epg_1_qubit = {
            1: {"rz": 0.0, "sx": 0.00036207066403884814, "x": 0.00036207066403884814},
            4: {"rz": 0.0, "sx": 0.0005429962529239195, "x": 0.0005429962529239195},
        }
        epg = rb.RBUtils.calculate_2q_epg(
            epc_2_qubit, qubits, gate_error_ratio, gates_per_clifford, epg_1_qubit
        )
        error_dict = {
            ((1, 4), "cx"): 0.012438847900902494,
        }

        expected_epg = error_dict[((1, 4), "cx")]
        actual_epg = epg[(1, 4)]["cx"]
        self.assertTrue(np.allclose(expected_epg, actual_epg, atol=0.001))

    def test_coherence_limit(self):
        """Test coherence_limit."""
        t1 = 100.0
        t2 = 100.0
        gate_2_qubits = 0.5
        gate_1_qubit = 0.1
        twoq_coherence_err = rb.RBUtils.coherence_limit(2, [t1, t1], [t2, t2], gate_2_qubits)

        oneq_coherence_err = rb.RBUtils.coherence_limit(1, [t1], [t2], gate_1_qubit)

        self.assertAlmostEqual(oneq_coherence_err, 0.00049975, 6, "Error: 1Q Coherence Limit")

        self.assertAlmostEqual(twoq_coherence_err, 0.00597, 5, "Error: 2Q Coherence Limit")


"""
Helper functions for when working with any image data

Arthur McCray
amccray@anl.gov
"""

import os

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import skimage
from ipywidgets import interact
from scipy import ndimage as ndi
from scipy.signal import tukey
from scipy.spatial.transform import Rotation as R

###
### Functions for displaying images
###

def show_im(
    image,
    title=None,
    simple=False,
    origin="upper",
    cbar=True,
    cbar_title="",
    scale=None,
    save=None,
    **kwargs,
):
    """Display an image on a new axis.

    Takes a 2D array and displays the image in grayscale with optional title on
    a new axis. In general it's nice to have things on their own axes, but if
    too many are open it's a good idea to close with plt.close('all').

    Args:
        image (2D array): Image to be displayed.
        title (str): (`optional`) Title of plot.
        simple (bool): (`optional`) Default output or additional labels.

            - True, will just show image.
            - False, (default) will show a colorbar with axes labels, and will adjust the
              contrast range for images with a very small range of values (<1e-12).

        origin (str): (`optional`) Control image orientation.

            - 'upper': (default) (0,0) in upper left corner, y-axis goes down.
            - 'lower': (0,0) in lower left corner, y-axis goes up.

        cbar (bool): (`optional`) Choose to display the colorbar or not. Only matters when
            simple = False.
        cbar_title (str): (`optional`) Title attached to the colorbar (indicating the
            units or significance of the values).
        scale (float): Scale of image in nm/pixel. Axis markers will be given in
            units of nanometers.

    Returns:
        None
    """
    _fig, ax = plt.subplots()
    image = np.array(image)
    if "cmap" not in kwargs:
        kwargs["cmap"] = "gray"
    if image.dtype == "bool":
        image = image.astype("int")
    if not simple and np.max(image) - np.min(image) < 1e-12:
        # adjust coontrast range
        vmin = np.min(image) - 1e-12
        vmax = np.max(image) + 1e-12
        im = ax.matshow(image, origin=origin, vmin=vmin, vmax=vmax, **kwargs)
    else:
        im = ax.matshow(image, origin=origin, **kwargs)

    if title is not None:
        ax.set_title(str(title), pad=0)

    if simple:
        plt.axis("off")
    else:
        plt.tick_params(axis="x", top=False)
        ax.xaxis.tick_bottom()
        ax.tick_params(direction="in")
        if scale is None:
            ticks_label = "pixels"
        else:

            def mjrFormatter(x):
                return f"{scale*x:.3g}"

            fov = scale * max(image.shape[0], image.shape[1])

            if fov < 4e3:  # if fov < 4um use nm scale
                ticks_label = " nm "
            elif fov > 4e6:  # if fov > 4mm use m scale
                ticks_label = "  m  "
                scale /= 1e9
            else:  # if fov between the two, use um
                ticks_label = r" $\mu$m "
                scale /= 1e3

            ax.yaxis.set_major_formatter(mpl.ticker.FuncFormatter(mjrFormatter))
            ax.xaxis.set_major_formatter(mpl.ticker.FuncFormatter(mjrFormatter))

        if origin == "lower":
            ax.text(y=0, x=0, s=ticks_label, rotation=-45, va="top", ha="right")
        elif origin == "upper":  # keep label in lower left corner
            ax.text(
                y=image.shape[0], x=0, s=ticks_label, rotation=-45, va="top", ha="right"
            )

        if cbar:
            plt.colorbar(im, ax=ax, pad=0.02, format="%.2g", label=str(cbar_title))

    if save:
        print("saving: ", save)
        plt.savefig(save, dpi=400, bbox_inches="tight")

    plt.show()
    return

def show_im_points(im=None, points=None, points2=None, size=None, title=None, **kwargs):
    """
    points an array [[y1,x1], [y2,x2], ...]
    """
    _fig, ax = plt.subplots()
    if im is not None:
        ax.matshow(im, cmap="gray", **kwargs)
    if points is not None:
        points = np.array(points)
        ax.plot(
            points[:, 1],
            points[:, 0],
            c="r",
            alpha=0.9,
            ms=size,
            marker="o",
            fillstyle="none",
            linestyle="none",
        )
    if points2 is not None:
        points2 = np.array(points2)
        ax.plot(
            points2[:, 1],
            points2[:, 0],
            c="b",
            alpha=0.9,
            ms=size,
            marker="o",
            fillstyle="none",
            linestyle="none",
        )
    ax.set_aspect(1)
    if title is not None:
        ax.set_title(str(title), pad=0)
    plt.show()


#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
import sys
sys.path.append("..")
import unittest
import numpy as np
from op_test_xpu import XPUOpTest
import paddle
from paddle import enable_static
import paddle.fluid as fluid
import paddle.fluid.core as core
from paddle.fluid.op import Operator
from paddle.fluid.tests.unittests.op_test import (
    OpTest, convert_float_to_uint16, convert_uint16_to_float)
from paddle import _C_ops

paddle.enable_static()

class TestSumOp(XPUOpTest):
    def setUp(self):
        self.op_type = "sum"
        self.init_kernel_type()
        self.init_kernel_type()
        x0 = np.random.random((3, 40)).astype(self.dtype)
        x1 = np.random.random((3, 40)).astype(self.dtype)
        x2 = np.random.random((3, 40)).astype(self.dtype)
        self.inputs = {"X": [("x0", x0), ("x1", x1), ("x2", x2)]}
        y = x0 + x1 + x2
        self.outputs = {'Out': y}

    def init_kernel_type(self):
        self.dtype = np.float32

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['x0'], 'Out')

#----------- test fp16 -----------
class TestFP16SumOp(TestSumOp):
    def init_kernel_type(self):
        self.dtype = np.float16

    def test_check_output(self):
        place = core.XPUPlace(0)
        # if core.is_float16_supported(place):
        self.check_output_with_place(place, atol=2e-2)

    # FIXME: Because of the precision fp16, max_relative_error
    # should be 0.15 here.
    def test_check_grad(self):
        place = core.XPUPlace(0)
        # if core.is_float16_supported(place):
        self.check_grad_with_place(
            place, ['x0'], 'Out', max_relative_error=0.15)

def create_test_sum_fp16_class(parent):
    class TestSumFp16Case(parent):
        def init_kernel_type(self):
            self.dtype = np.float16

        def test_w_is_selected_rows(self):
            place = core.XPUPlace(0)
            # if core.is_float16_supported(place):
            for inplace in [True, False]:
                self.check_with_place(place, inplace)

    cls_name = "{0}_{1}".format(parent.__name__, "SumFp16Test")
    TestSumFp16Case.__name__ = cls_name
    globals()[cls_name] = TestSumFp16Case

class API_Test_Add_n(unittest.TestCase):
    def test_api(self):
        with fluid.program_guard(fluid.Program(), fluid.Program()):
            input0 = fluid.layers.fill_constant(
                shape=[2, 3], dtype='int64', value=5)
            input1 = fluid.layers.fill_constant(
                shape=[2, 3], dtype='int64', value=3)
            expected_result = np.empty((2, 3))
            expected_result.fill(8)
            sum_value = paddle.add_n([input0, input1])
            exe = fluid.Executor(fluid.XPUPlace(0))
            result = exe.run(fetch_list=[sum_value])

            self.assertEqual((result == expected_result).all(), True)

        with fluid.dygraph.guard():
            input0 = paddle.ones(shape=[2, 3], dtype='float32')
            expected_result = np.empty((2, 3))
            expected_result.fill(2)
            sum_value = paddle.add_n([input0, input0])

            self.assertEqual((sum_value.numpy() == expected_result).all(), True)

class TestRaiseSumError(unittest.TestCase):
    def test_errors(self):
        def test_type():
            fluid.layers.sum([11, 22])

        self.assertRaises(TypeError, test_type)

        def test_dtype():
            data1 = fluid.data(name="input1", shape=[10], dtype="int8")
            data2 = fluid.data(name="input2", shape=[10], dtype="int8")
            fluid.layers.sum([data1, data2])

        self.assertRaises(TypeError, test_dtype)

        def test_dtype1():
            data1 = fluid.data(name="input1", shape=[10], dtype="int8")
            fluid.layers.sum(data1)

        self.assertRaises(TypeError, test_dtype1)

class TestRaiseSumsError(unittest.TestCase):
    def test_errors(self):
        def test_type():
            fluid.layers.sums([11, 22])

        self.assertRaises(TypeError, test_type)

        def test_dtype():
            data1 = fluid.data(name="input1", shape=[10], dtype="int8")
            data2 = fluid.data(name="input2", shape=[10], dtype="int8")
            fluid.layers.sums([data1, data2])

        self.assertRaises(TypeError, test_dtype)

        def test_dtype1():
            data1 = fluid.data(name="input1", shape=[10], dtype="int8")
            fluid.layers.sums(data1)

        self.assertRaises(TypeError, test_dtype1)

        def test_out_type():
            data1 = fluid.data(name="input1", shape=[10], dtype="flaot32")
            data2 = fluid.data(name="input2", shape=[10], dtype="float32")
            fluid.layers.sums([data1, data2], out=[10])

        self.assertRaises(TypeError, test_out_type)

        def test_out_dtype():
            data1 = fluid.data(name="input1", shape=[10], dtype="flaot32")
            data2 = fluid.data(name="input2", shape=[10], dtype="float32")
            out = fluid.data(name="out", shape=[10], dtype="int8")
            fluid.layers.sums([data1, data2], out=out)

        self.assertRaises(TypeError, test_out_dtype)

class TestSumOpError(unittest.TestCase):
    def test_errors(self):
        def test_empty_list_input():
            with fluid.dygraph.guard():
                fluid._C_ops.sum([])

        def test_list_of_none_input():
            with fluid.dygraph.guard():
                fluid._C_ops.sum([None])

        self.assertRaises(Exception, test_empty_list_input)
        self.assertRaises(Exception, test_list_of_none_input)

if __name__ == "__main__":
    enable_static()
    unittest.main()


import numpy as np

def assert_equal(A: np.ndarray, B: np.ndarray):
    """Asserts that the arrays A and B are the same. Verifies and throws
    exceptions if the following are not the same:

    * Dimensions;

    * Shape;

    * Values
    """

    if A.ndim != B.ndim:
        raise ValueError("A has different dimension of B.")

    if A.shape != B.shape:
        raise ValueError("A has different shape of B.")

    if not np.allclose(A, B):
        raise ValueError("A is not equal to B.")


# -*- coding: utf-8 -*-
"""
Created on Mon Jun 18 17:24:04 2018

@author: SilverDoe
"""

'''
>>Pandas is an open-source Python Library used for high-performance data manipulation
 and data analysis using its powerful data structures.

>> It can load, organize, manipulate, model, and analyse the data.

>> Python deals with the following data structures :

    1. Series - 1D labeled homogeneous array, sizeimmutable, value mutable.
    2. DataFrame - General 2D labeled, value and size-mutable tabular structure with potentially heterogeneously typed columns.
    3. Panel - General 3D labeled, value and size-mutable array.

    higher dimensional data structure is a container of its lower dimensional data structure
    ========================================================================================
    >>DataFrame is a container of Series.
    >>Panel is a container of DataFrame.

============================== Series =========================================

pandas.Series( data, index, dtype, copy)
pandas.Series(data=None, index=None, dtype=None, name=None, copy=False, fastpath=False)

parameters
==========

1. data : data takes various forms like ndarray, list, constants.

2. index : Index values must be unique and hashable, same length as data. Default
           np.arrange(n) if no index is passed.

3. 	dtype : dtype is for data type. If None, data type will be inferred.

4. copy : Copy data. Default False
'''
#=============== Empty Series =================================================

import pandas as pd
s = pd.Series()
print(s)

#============= Series from ndarray ============================================

# no index passed
import pandas as pd
import numpy as np
data = np.array(['a','b','c','d'])
s = pd.Series(data)
print(s)

# index passed
import pandas as pd
import numpy as np
data = np.array(['a','b','c','d'])
s = pd.Series(data,index=[100,101,102,103])
print(s)

#========== Series from Dictionary ============================================

#No index given. Dictionary keys are used to construct index
import pandas as pd
data = {'a' : 0., 'b' : 1., 'c' : 2.}
s = pd.Series(data)
print(s)

#Index given. Index order is persisted and the missing element is filled with NaN
import pandas as pd
data = {'a' : 0., 'b' : 1., 'c' : 2.}
s = pd.Series(data,index=['b','c','d','a'])
print(s)

#=========== Series from Scalar ===============================================

import pandas as pd
s = pd.Series(5, index=[0, 1, 2, 3])
print(s)

#=========== Accessing data from series with position or index ================
import pandas as pd
s = pd.Series([1,2,3,4,5],index = ['a','b','c','d','e'])

#retrieve data
print(s[0]) # using position
print(s['b']) # using index
print(s[:3]) # retreiving the first 3 elements in the series
print(s[-3:]) # last three elements
print(s[['a','c','d']]) # retreiving multiple values