main.py

import argparse
import os
import json
from src.plot_utils import plotLightCurve
from src.CNN import classification as CNN
import src.normalizing_flows as NF
import src.nf_metrics as nf_met
import src.cnn_metrics as cnn_met
import sys
import shutil
import numpy as np

def plots(nf):
    passband2name = {0: 'g', 1: 'r'}
    df = nf.df # Accessing the dataframe from NF object
    objects = nf.objects # Accessing unique objects from data
    path = "Nf_run_plots/"
    #creates Nf_run_plots and saves plots at each run
    isExist = os.path.exists(path)
    if isExist:
        shutil.rmtree(path)
    os.makedirs(path)

    for i in range(len(objects)):
        obj_name = objects[i]
        df_obj = df.loc[df['object_id'] == obj_name] # select data for object=object_name
        plotLightCurve(obj_name, df_obj, nf.flux_pred[i], nf.aug_timestamps[i], passband2name)

# Run binary classification
def run_CNN(cnn_params, report_file, nf = 0, og = 0, run_nf = 0):
    
    X_matrix = []
    y_vector = []
    X_test = []
    y_test = []
    
    if (not run_nf):
        
        X_file = os.path.join("data", "images.json")
        y_file = os.path.join("data", "labels.json")
        with open(X_file, 'r') as f:
            X_matrix = json.load(f)
        with open(y_file, 'r') as f:
            y_vector = json.load(f)
        
    else:
        if og:
            X_file = os.path.join("data", "images.json")
            y_file = os.path.join("data", "labels.json")
            with open(X_file, 'r') as f:
                X_matrix = json.load(f)
            with open(y_file, 'r') as f:
                y_vector = json.load(f)
            X_test = nf.X_matrix
            y_test = nf.y_vector
            X_test = np.array(X_test)
            y_test = np.array(y_test)
            X_test = np.array((X_test - X_test.mean()) / X_test.std(), dtype = np.float32)
        else:
            X_matrix = nf.X_matrix
            y_vector = nf.y_vector
            
    X_matrix = np.array(X_matrix)
    y_vector = np.array(y_vector)
    # normalize input data
    X_matrix = np.array((X_matrix - X_matrix.mean()) / X_matrix.std(), dtype = np.float32)
    
    print('Training CNN')
    y_test, y_test_pred = CNN(cnn_params, X_matrix, y_vector, report_file, og, X_test, y_test)
        
    # Report
    report = cnn_met.gen_report(y_test, y_test_pred)
    print(report)

    original_stdout = sys.stdout
    with open(report_file, 'a') as f:
        sys.stdout = f
        print(report)
        sys.stdout = original_stdout

if __name__ == '__main__':
    
    # Command line arguments
    parser = argparse.ArgumentParser(description='Classification of Supernovae Light Curves by NF approximation')
    
    # Default values of parameters are defined
    parser.add_argument('--param', default = 'param/param.json', help='file containing hyperparameters')
    parser.add_argument('-nf', '--nf', help="input number of objects generated by normalizing flows", type = int, default = 0)
    parser.add_argument('-v', '--verbose', help="increase output verbosity", action="store_true")
    parser.add_argument('-p', '--plot', help="generate plots", action="store_true" )
    parser.add_argument('-s', '--shuffle', help="Shuffles the objects generated by Normalizing Flows", action="store_true")
    parser.add_argument('-o', '--og', help="Runs CNN to classify the objects generated by Normalizing Flows", action="store_true")
    
    args = parser.parse_args()
    num_objects = args.nf
    
    verbose = args.verbose
    plot = args.plot
    shuffle = args.shuffle
    og = args.og
    
    # Hyperparameters from json file
    with open(args.param) as paramfile:
        param = json.load(paramfile)
    
    # Creates a report file
    report_file = 'report.txt'
    
    if os.path.exists(report_file):
        os.remove(report_file)
    f = open(report_file, 'w')
    f.close()
    
    cnn_params = param["CNN"]
    
    # Runs CNN as default
    if (num_objects == 0):
        run_CNN(cnn_params, report_file)
    else:
        # Load data
        nf_params = param["NF"]
        data_dir = 'data/ANTARES_NEW.csv'

        print('Training NF')
        # Run Normalizing Flows to obtain the approximate light curve
        nf = NF.FitNF(data_dir, shuffle, num_objects, nf_params, report_file, verbose)

        metrics = nf_met.generate_NF_report(nf.flux, nf.flux_pred_metrics, nf.flux_err, nf.flux_err_pred_metrics)
        metrics.to_csv('nfmetrics.csv')
        
        original_stdout = sys.stdout
        with open(report_file, 'a') as f:
            sys.stdout = f
            print(metrics)
            sys.stdout = original_stdout
        
        if plot:
            plots(nf)
            
        run_CNN(cnn_params, report_file, nf, og, run_nf = 1)