
########################################
########### IMPORT PACKAGES ############
########################################

### standard packages
import os
import numpy as np
import glob as glob
from tqdm import tqdm 
from random import random
import argparse
import pandas as pd
import pdb

### torch packages
import torch
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch.autograd import Variable
import torch.nn as nn

### sklearn packages
from sklearn.metrics import precision_recall_curve
from sklearn.metrics import average_precision_score
from sklearn.metrics import confusion_matrix
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score

### plotting packages
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec


########################################
########### PARSE ARGUMENTS ############
########################################

### an example input to command line:
### python exonet.py 225 64 1e-5 '/data/kepler/new' '/data_sata1/ensembling/test'

### parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("n_epochs", help="number of epochs to use for training", type=int)
parser.add_argument("n_batches", help="number of matches to use for training", type=int)
parser.add_argument("r_learn", help="learning rate for Adam optimizer", type=float)
parser.add_argument("d_rate", help="Dropout rate to use during training", type=float)
parser.add_argument("d_path", help="path to data (should contain folders named val, test, train")
parser.add_argument("m_out", help="path for output model, files, and plots")
parser.add_argument("--fixed_seed", help="set if wanting to fix the seed", action="store_true")
parser.add_argument("--XS", help="use ExtranetXS model", action="store_true")
args = parser.parse_args()

### set manual seed
if args.fixed_seed:
    torch.cuda.manual_seed(42)


########################################
########### DEFINE CLASSES  ############
########################################

class KeplerDataLoader(Dataset):
    
    '''
    
    PURPOSE: DATA LOADER FOR KERPLER LIGHT CURVES
    INPUT: PATH TO DIRECTOR WITH LIGHT CURVES + INFO FILES
    OUTPUT: LOCAL + GLOBAL VIEWS, LABELS
    
    '''

    def __init__(self, filepath):

        ### list of global, local, and info files (assumes certain names of files)
        self.flist_global = np.sort(glob.glob(os.path.join(filepath, '*global.npy')))
        self.flist_local = np.sort(glob.glob(os.path.join(filepath, '*local.npy')))
        self.flist_info = np.sort(glob.glob(os.path.join(filepath, '*info.npy')))
        
        ### list of whitened centroid files
        self.flist_global_cen = np.sort(glob.glob(os.path.join(filepath, '*global_cen_w.npy')))
        self.flist_local_cen = np.sort(glob.glob(os.path.join(filepath, '*local_cen_w.npy')))
        
        ### ids = {TIC}_{TCE}
        self.ids = np.sort([(x.split('/')[-1]).split('_')[1] + '_' + (x.split('/')[-1]).split('_')[2] for x in self.flist_global])

    def __len__(self):

        return self.ids.shape[0]

    def __getitem__(self, idx):

        ### grab local and global views
        data_global = np.load(self.flist_global[idx])
        data_local = np.load(self.flist_local[idx])

        ### grab centroid views
        data_global_cen = np.load(self.flist_global_cen[idx])
        data_local_cen = np.load(self.flist_local_cen[idx])
        
        ### info file contains: [0]kic, [1]tce, [2]period, [3]epoch, [4]duration, [5]label)
        data_info = np.load(self.flist_info[idx])
        
        return (data_local, data_global, data_local_cen, data_global_cen, data_info[6:]), data_info[5]


class ExtranetModel(nn.Module):

    '''
    
    PURPOSE: DEFINE EXTRANET MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    '''
    
    def __init__(self, d_rate):

        ### initialize model
        super(ExtranetModel, self).__init__()
        droplayer = nn.Dropout(d_rate) if d_rate > 1e-3 else nn.Dropout(0.)

        ### define global convolutional lalyer
        self.fc_global = nn.Sequential(
            nn.Conv1d(2, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(16, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(5, stride=2),
            nn.Conv1d(16, 32, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(32, 32, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(5, stride=2),
            nn.Conv1d(32, 64, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(64, 64, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(5, stride=2),
            nn.Conv1d(64, 128, 5, stride=1, padding=2),
            droplayer,
            nn.ReLU(),
            nn.Conv1d(128, 128, 5, stride=1, padding=2),
            nn.MaxPool1d(5, stride=2),
            nn.Conv1d(128, 256, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(256, 256, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(5, stride=2),
        )

        ### define local convolutional lalyer
        self.fc_local = nn.Sequential(
            nn.Conv1d(2, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(16, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(7, stride=2),
            nn.Conv1d(16, 32, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(32, 32, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(7, stride=2),
        )

        ### define fully connected layer that combines both views
        self.final_layer = nn.Sequential(
            nn.Linear(16586, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            ### need output of 1 because using BCE for loss
            nn.Linear(512, 1),
            nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_star):

        ### concatonate light curve and centroid data
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        
        ### get outputs of global and local convolutional layers
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        
        ### flattening outputs from convolutional layers into vector
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)

        ### concatonate global and local views with stellar parameters
        out = torch.cat([out_global, out_local, x_star.squeeze(1)], dim=1)
        out = self.final_layer(out)

        return out


class ExtranetXSModel(nn.Module):

    '''
    
    PURPOSE: DEFINE EXTRANET-XS MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    '''

    def __init__(self):

        ### initializing the nn.Moduel (super) class
        ### (must do this first always)
        super(ExtranetXSModel, self).__init__()

        ### define global convolutional lalyer
        self.fc_global = nn.Sequential(
            nn.Conv1d(2, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(2, stride=2),
            nn.Conv1d(16, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(2, stride=2),
            nn.Conv1d(16, 32, 5, stride=1, padding=2),
            nn.ReLU(),
        )

        ### define the local convolutional layer
        self.fc_local = nn.Sequential(
            nn.Conv1d(2, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(2, stride=2),
            nn.Conv1d(16, 16, 5, stride=1, padding=2),
            nn.ReLU(),
        )
                
        ### define fully connected layer that combines both views
        self.final_layer = nn.Sequential(
            nn.Linear(58, 1),
            nn.Sigmoid())
        
    ### define how to move forward through model
    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_starpars):

        ### concatonate light curve and centroid data
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        
        ### do global pooling
        out_global = F.max_pool1d(out_global, out_global.shape[-1])
        out_local = F.max_pool1d(out_local, out_local.shape[-1])

        ### flattening outputs from convolutional layers into vector
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)

        ### concatonate global and local views with stellar parameters
        out = torch.cat([out_global, out_local, x_starpars.squeeze(1)], dim=1)
        out = self.final_layer(out)

        return out


########################################
########## DEFINE FUNCTIONS ############
########################################

def invert_tensor(tensor):

    '''

    PURPOSE: FLIP A 1D TENSOR ALONG ITS AXIS
    INPUT: 1D TENSOR
    OUTPUT: INVERTED 1D TENSOR

    '''

    idx = [i for i in range(tensor.size(0)-1, -1, -1)]
    idx = torch.cuda.LongTensor(idx)
    inverted_tensor = tensor.index_select(0, idx)

    return inverted_tensor


def train_model(n_epochs, kepler_data_loader, kepler_val_loader, model, criterion, optimizer):

    '''

    PURPOSE: TRAIN MODEL

    INPUTS:  num_epoch = number of epochs for training
             kepler_data_loader = data loader for Kepler dataset
             model = model use for training
             criterion = criterion for calculating loss

    OUTPUT:  epoch_{train/val}_loss = training set loss for each epoch
             epoch_val_acc = validation set accuracy for each epoch
             epoch_val_ap = validation set avg. precision for each epoch
             final_val_pred = validation predictions from final model
             final_val_gt = validation ground truths

    '''

    ### empty arrays to fill per-epoch outputs
    epoch_train_loss = []
    epoch_val_loss = []
    epoch_val_acc = []
    epoch_val_ap = []

    ### loop over number of epochs of training
    for epoch in tqdm(range(n_epochs)):

        ####################
        ### for training set
        
        ### loop over batches
        train_loss = torch.zeros(1).cuda()
        for x_train_data, y_train in kepler_data_loader:
            
            ### get local view, global view, and label for training
            x_train_local, x_train_global, x_train_local_cen, x_train_global_cen, x_train_star = x_train_data
            x_train_local = Variable(x_train_local).type(torch.FloatTensor).cuda()
            x_train_global = Variable(x_train_global).type(torch.FloatTensor).cuda()    
            x_train_local_cen = Variable(x_train_local_cen).type(torch.FloatTensor).cuda()
            x_train_global_cen = Variable(x_train_global_cen).type(torch.FloatTensor).cuda()
            x_train_star = Variable(x_train_star).type(torch.FloatTensor).cuda()        
            y_train = Variable(y_train).type(torch.FloatTensor).cuda()
                      
            ### randomly invert half of light curves
            for batch_ind in range(x_train_local.shape[0]):     
                
                ### add random gaussian noise
                sig_noise = np.random.uniform(0, 1.0)
                local_noise = Variable(x_train_local[batch_ind].data.new(x_train_local[batch_ind].size()).normal_(0.0, sig_noise))
                x_train_local[batch_ind] = x_train_local[batch_ind] + local_noise
                global_noise = Variable(x_train_global[batch_ind].data.new(x_train_global[batch_ind].size()).normal_(0.0, sig_noise))
                x_train_global[batch_ind] = x_train_global[batch_ind] + global_noise

                if random() < 0.5:
                    x_train_local[batch_ind] = invert_tensor(x_train_local[batch_ind])
                    x_train_global[batch_ind] = invert_tensor(x_train_global[batch_ind])
                    x_train_local_cen[batch_ind] = invert_tensor(x_train_local_cen[batch_ind])
                    x_train_global_cen[batch_ind] = invert_tensor(x_train_global_cen[batch_ind])

            ### fix dimensions for next steps
            x_train_local = x_train_local.unsqueeze(1)
            x_train_global = x_train_global.unsqueeze(1)
            x_train_local_cen = x_train_local_cen.unsqueeze(1)
            x_train_global_cen = x_train_global_cen.unsqueeze(1)
            x_train_star = x_train_star.unsqueeze(1)
            y_train = y_train.unsqueeze(1)

            ### calculate loss using model
            output_train = model(x_train_local, x_train_global, x_train_local_cen, x_train_global_cen, x_train_star)
            loss = criterion(output_train, y_train)
            train_loss += loss.data

            ### train model (zero gradients and back propogate results)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        ### record training loss for this epoch (divided by size of training dataset)
        epoch_train_loss.append(train_loss.cpu().numpy() / len(kepler_data_loader.dataset))
        
        ######################
        ### for validation set
        
        ### loop over batches
        val_pred, val_gt, val_loss, num_corr = [], [], 0, 0
        for x_val_data, y_val in kepler_val_loader:
                        
            ### get local view, global view, and label for validating
            x_val_local, x_val_global, x_val_local_cen, x_val_global_cen, x_val_star = x_val_data
            x_val_local = Variable(x_val_local).type(torch.FloatTensor).cuda()
            x_val_global = Variable(x_val_global).type(torch.FloatTensor).cuda()
            x_val_local_cen = Variable(x_val_local_cen).type(torch.FloatTensor).cuda()
            x_val_global_cen = Variable(x_val_global_cen).type(torch.FloatTensor).cuda()
            x_val_star = Variable(x_val_star).type(torch.FloatTensor).cuda()
 
            ### fix dimensions for next steps
            y_val = Variable(y_val).type(torch.FloatTensor).cuda()
            x_val_local = x_val_local.unsqueeze(1)
            x_val_global = x_val_global.unsqueeze(1)
            x_val_local_cen = x_val_local_cen.unsqueeze(1)
            x_val_global_cen = x_val_global_cen.unsqueeze(1)
            x_val_star = x_val_star.unsqueeze(1)
            y_val = y_val.unsqueeze(1)

            ### calculate loss & add to sum over all batches
            output_val = model(x_val_local, x_val_global, x_val_local_cen, x_val_global_cen, x_val_star)
            loss_val = criterion(output_val, y_val)
            val_loss += loss_val.data

            ### get number of correct predictions using threshold=0.5
            ### & sum over all batches
            output_pred = output_val >= 0.5
            num_corr += output_pred.eq(y_val.byte()).sum().item()
                        
            ### record predictions and ground truth by model
            ### (used for AP per epoch; reset at each epoch; final values output)
            val_pred.append(output_val.data.cpu().numpy())
            val_gt.append(y_val.data.cpu().numpy())
            
        ### record validation loss calculate for this epoch (divided by size of validation dataset)
        epoch_val_loss.append(val_loss.cpu().numpy() / len(kepler_val_loader.dataset))
        
        ### record validation accuracy (# correct predictions in val set) for this epoch
        epoch_val_acc.append(num_corr / len(kepler_val_loader.dataset))
        
        ### calculate average precision for this epoch
        epoch_val_ap.append(average_precision_score(np.concatenate(val_gt).ravel(), np.concatenate(val_pred).ravel(), average=None))            
        
    ### grab final predictions and ground truths for validation set
    final_val_pred = np.concatenate(val_pred).ravel()
    final_val_gt = np.concatenate(val_gt).ravel()

    return epoch_train_loss, epoch_val_loss, epoch_val_acc, epoch_val_ap, final_val_pred, final_val_gt


########################################
############ TRAIN MODEL  ##############
########################################

### setup screen output
print("\nTRAINING MODEL...\n")

### initialize model; cuda puts it on GPU
if args.XS:
    model = ExtranetXSModel().cuda()
else:
    ### dropout rate for training
    d_rate = args.d_rate
    model = ExtranetModel(d_rate).cuda()

### learning rate
lr  = args.r_learn

### specify optimizer for learning to use for training
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

### specify loss function to use for training
criterion = nn.BCELoss()

### specify batch size to use for training
batch_size = args.n_batches 

### number of epochs to use for training
n_epochs = args.n_epochs 

### grab data using data loader
kepler_train_data = KeplerDataLoader(filepath=os.path.join(args.d_path, 'train'))
kepler_val_data = KeplerDataLoader(filepath=os.path.join(args.d_path, 'test'))
kepler_data_loader = DataLoader(kepler_train_data, batch_size=batch_size, shuffle=True, num_workers=4)
kepler_val_loader = DataLoader(kepler_val_data, batch_size=batch_size, shuffle=False, num_workers=4)

### train model
loss_train_epoch, loss_val_epoch, acc_val_epoch, ap_val_epoch, pred_val_final, gt_val_final  = train_model(n_epochs, kepler_data_loader, kepler_val_loader, model, criterion, optimizer)


########################################
####### CALCULATE STATISTICS ###########
########################################

### setup screen output
print("\nCALCULATING METRICS...\n")

### calculate average precision & precision-recall curves
AP = average_precision_score(gt_val_final, pred_val_final, average=None)
print("   average precision = {0:0.4f}\n".format(AP))
 
### calculate precision-recall curve
P, R, _ = precision_recall_curve(gt_val_final, pred_val_final)

### calculate confusion matrix based on different thresholds 
thresh = [0.5, 0.6, 0.7, 0.8, 0.9]
prec_thresh, recall_thresh = np.zeros(len(thresh)), np.zeros(len(thresh))
for n, nval in enumerate(thresh):
    pred_byte = np.zeros(len(pred_val_final))
    for i, val in enumerate(pred_val_final):
        if val > nval:
            pred_byte[i] = 1.0
        else:
            pred_byte[i] = 0.0
    prec_thresh[n] = precision_score(gt_val_final, pred_byte)
    recall_thresh[n] = recall_score(gt_val_final, pred_byte)
    print("   thresh = {0:0.2f}, precision = {1:0.2f}, recall = {2:0.2f}".format(thresh[n], prec_thresh[n], recall_thresh[n]))
    tn, fp, fn, tp = confusion_matrix(gt_val_final, pred_byte).ravel()
    print("      TN = {0:0}, FP = {1:0}, FN = {2:0}, TP = {3:0}".format(tn, fp, fn, tp))


########################################
######### OUTPUT MODEL + STATS  ########
########################################

### transform from loss per sample to loss per batch (multiple by batch size to compare to Chris')
loss_train_batch = [x.item()* batch_size for x in loss_train_epoch]
loss_val_batch = [x.item()* batch_size for x in loss_val_epoch]

### setup output
run = 0

### output predictions & ground truth
pt_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-lr' + str(lr) + '-pt.csv')
while os.path.isfile(pt_fname):
    run +=1
    pt_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-lr' + str(lr) + '-pt.csv')
df = pd.DataFrame({"gt" : gt_val_final, "pred" : pred_val_final})
df.to_csv(pt_fname, index=False)

### output per-iteration values
epochs_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-lr' + str(lr) + '-epoch.csv')
df = pd.DataFrame({"loss_train":loss_train_batch, "loss_val":loss_val_batch, "acc_val":acc_val_epoch, "ap_val":ap_val_epoch})
df.to_csv(epochs_fname, index=False)

### save model
model_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-lr' + str(lr) + '-drop' + str(d_rate) + '-model.pth')
torch.save(model.state_dict(), model_fname)
print("\nOUTPUTTING MODEL + RESULTS @ " + os.path.join(args.m_out, model_fname) + "\n")


########################################
################ MAKE PLOTS ############
########################################

### setup figure
fig = plt.figure(figsize=(7, 7))
ax = gridspec.GridSpec(2,2)
ax.update(wspace = 0.4, hspace = 0.4)
ax1 = plt.subplot(ax[0,0])
ax2 = plt.subplot(ax[0,1])
ax3 = plt.subplot(ax[1,0])
ax4 = plt.subplot(ax[1,1])

### plot precision-recall curve
ax1.set_xlabel('Precision', fontsize=10, labelpad=10)
ax1.set_ylabel('Recall', fontsize=10)
ax1.set_xlim([0.0, 1.0])
ax1.set_ylim([0.0, 1.0])
ax1.plot(R, P, linewidth=3, color='black')

### plot loss curve for training and validation sets
ax2.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax2.set_ylabel('Loss', fontsize=10)
ax2.set_xlim([0.0, n_epochs])
ax2.set_ylim([0.0, np.max(loss_train_batch)*1.5])
ax2.plot(np.arange(len(loss_train_batch)), loss_train_batch, linewidth=3, color='cadetblue')
ax2.plot(np.arange(len(loss_val_batch)), loss_val_batch, linewidth=3, color='orangered')

### plot average precision per epoch
ax3.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax3.set_ylabel('Average Precision', fontsize=10)
ax3.plot(np.arange(len(ap_val_epoch)), ap_val_epoch, linewidth=1.0, color='orangered')
ax3.scatter(np.arange(len(ap_val_epoch)), ap_val_epoch, marker='o', edgecolor='orangered', facecolor='orangered', s=10, linewidth=0.5, alpha=0.5)

### plot accuracy per epoch
ax4.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax4.set_ylabel('Accuracy', fontsize=10)
ax4.plot(np.arange(len(acc_val_epoch)), acc_val_epoch, color='orangered', linewidth=1.0)
ax4.scatter(np.arange(len(acc_val_epoch)), acc_val_epoch, marker='o', edgecolor='orangered', facecolor='orangered', s=10, linewidth=0.5, alpha=0.5)

### save plot
plot_fname = 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-plot.pdf'
plt.savefig(os.path.join(args.m_out, plot_fname), bbox_inches='tight', dpi=200, rastersized=True, alpha=True)


3d97c943d12b39464dcfcc58b563dccd40f1bcac
<mask token>


class KeplerDataLoader(Dataset):
    <mask token>

    def __init__(self, filepath):
        self.flist_global = np.sort(glob.glob(os.path.join(filepath,
            '*global.npy')))
        self.flist_local = np.sort(glob.glob(os.path.join(filepath,
            '*local.npy')))
        self.flist_info = np.sort(glob.glob(os.path.join(filepath,
            '*info.npy')))
        self.flist_global_cen = np.sort(glob.glob(os.path.join(filepath,
            '*global_cen_w.npy')))
        self.flist_local_cen = np.sort(glob.glob(os.path.join(filepath,
            '*local_cen_w.npy')))
        self.ids = np.sort([(x.split('/')[-1].split('_')[1] + '_' + x.split
            ('/')[-1].split('_')[2]) for x in self.flist_global])

    def __len__(self):
        return self.ids.shape[0]

    def __getitem__(self, idx):
        data_global = np.load(self.flist_global[idx])
        data_local = np.load(self.flist_local[idx])
        data_global_cen = np.load(self.flist_global_cen[idx])
        data_local_cen = np.load(self.flist_local_cen[idx])
        data_info = np.load(self.flist_info[idx])
        return (data_local, data_global, data_local_cen, data_global_cen,
            data_info[6:]), data_info[5]


class ExtranetModel(nn.Module):
    """
    
    PURPOSE: DEFINE EXTRANET MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    """

    def __init__(self, d_rate):
        super(ExtranetModel, self).__init__()
        droplayer = nn.Dropout(d_rate) if d_rate > 0.001 else nn.Dropout(0.0)
        self.fc_global = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1,
            padding=2), nn.ReLU(), nn.Conv1d(16, 16, 5, stride=1, padding=2
            ), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.Conv1d(16, 32, 5,
            stride=1, padding=2), nn.ReLU(), nn.Conv1d(32, 32, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.Conv1d(32,
            64, 5, stride=1, padding=2), nn.ReLU(), nn.Conv1d(64, 64, 5,
            stride=1, padding=2), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.
            Conv1d(64, 128, 5, stride=1, padding=2), droplayer, nn.ReLU(),
            nn.Conv1d(128, 128, 5, stride=1, padding=2), nn.MaxPool1d(5,
            stride=2), nn.Conv1d(128, 256, 5, stride=1, padding=2), nn.ReLU
            (), nn.Conv1d(256, 256, 5, stride=1, padding=2), nn.ReLU(), nn.
            MaxPool1d(5, stride=2))
        self.fc_local = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1, padding
            =2), nn.ReLU(), nn.Conv1d(16, 16, 5, stride=1, padding=2), nn.
            ReLU(), nn.MaxPool1d(7, stride=2), nn.Conv1d(16, 32, 5, stride=
            1, padding=2), nn.ReLU(), nn.Conv1d(32, 32, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(7, stride=2))
        self.final_layer = nn.Sequential(nn.Linear(16586, 512), nn.ReLU(),
            nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 512), nn.ReLU(),
            nn.Linear(512, 1), nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_star):
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)
        out = torch.cat([out_global, out_local, x_star.squeeze(1)], dim=1)
        out = self.final_layer(out)
        return out


class ExtranetXSModel(nn.Module):
    """
    
    PURPOSE: DEFINE EXTRANET-XS MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    """

    def __init__(self):
        super(ExtranetXSModel, self).__init__()
        self.fc_global = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(16,
            16, 5, stride=1, padding=2), nn.ReLU(), nn.MaxPool1d(2, stride=
            2), nn.Conv1d(16, 32, 5, stride=1, padding=2), nn.ReLU())
        self.fc_local = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1, padding
            =2), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(16, 16, 5,
            stride=1, padding=2), nn.ReLU())
        self.final_layer = nn.Sequential(nn.Linear(58, 1), nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_starpars
        ):
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        out_global = F.max_pool1d(out_global, out_global.shape[-1])
        out_local = F.max_pool1d(out_local, out_local.shape[-1])
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)
        out = torch.cat([out_global, out_local, x_starpars.squeeze(1)], dim=1)
        out = self.final_layer(out)
        return out


<mask token>

<mask token>


class KeplerDataLoader(Dataset):
    """
    
    PURPOSE: DATA LOADER FOR KERPLER LIGHT CURVES
    INPUT: PATH TO DIRECTOR WITH LIGHT CURVES + INFO FILES
    OUTPUT: LOCAL + GLOBAL VIEWS, LABELS
    
    """

    def __init__(self, filepath):
        self.flist_global = np.sort(glob.glob(os.path.join(filepath,
            '*global.npy')))
        self.flist_local = np.sort(glob.glob(os.path.join(filepath,
            '*local.npy')))
        self.flist_info = np.sort(glob.glob(os.path.join(filepath,
            '*info.npy')))
        self.flist_global_cen = np.sort(glob.glob(os.path.join(filepath,
            '*global_cen_w.npy')))
        self.flist_local_cen = np.sort(glob.glob(os.path.join(filepath,
            '*local_cen_w.npy')))
        self.ids = np.sort([(x.split('/')[-1].split('_')[1] + '_' + x.split
            ('/')[-1].split('_')[2]) for x in self.flist_global])

    def __len__(self):
        return self.ids.shape[0]

    def __getitem__(self, idx):
        data_global = np.load(self.flist_global[idx])
        data_local = np.load(self.flist_local[idx])
        data_global_cen = np.load(self.flist_global_cen[idx])
        data_local_cen = np.load(self.flist_local_cen[idx])
        data_info = np.load(self.flist_info[idx])
        return (data_local, data_global, data_local_cen, data_global_cen,
            data_info[6:]), data_info[5]


class ExtranetModel(nn.Module):
    """
    
    PURPOSE: DEFINE EXTRANET MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    """

    def __init__(self, d_rate):
        super(ExtranetModel, self).__init__()
        droplayer = nn.Dropout(d_rate) if d_rate > 0.001 else nn.Dropout(0.0)
        self.fc_global = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1,
            padding=2), nn.ReLU(), nn.Conv1d(16, 16, 5, stride=1, padding=2
            ), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.Conv1d(16, 32, 5,
            stride=1, padding=2), nn.ReLU(), nn.Conv1d(32, 32, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.Conv1d(32,
            64, 5, stride=1, padding=2), nn.ReLU(), nn.Conv1d(64, 64, 5,
            stride=1, padding=2), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.
            Conv1d(64, 128, 5, stride=1, padding=2), droplayer, nn.ReLU(),
            nn.Conv1d(128, 128, 5, stride=1, padding=2), nn.MaxPool1d(5,
            stride=2), nn.Conv1d(128, 256, 5, stride=1, padding=2), nn.ReLU
            (), nn.Conv1d(256, 256, 5, stride=1, padding=2), nn.ReLU(), nn.
            MaxPool1d(5, stride=2))
        self.fc_local = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1, padding
            =2), nn.ReLU(), nn.Conv1d(16, 16, 5, stride=1, padding=2), nn.
            ReLU(), nn.MaxPool1d(7, stride=2), nn.Conv1d(16, 32, 5, stride=
            1, padding=2), nn.ReLU(), nn.Conv1d(32, 32, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(7, stride=2))
        self.final_layer = nn.Sequential(nn.Linear(16586, 512), nn.ReLU(),
            nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 512), nn.ReLU(),
            nn.Linear(512, 1), nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_star):
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)
        out = torch.cat([out_global, out_local, x_star.squeeze(1)], dim=1)
        out = self.final_layer(out)
        return out


class ExtranetXSModel(nn.Module):
    """
    
    PURPOSE: DEFINE EXTRANET-XS MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    """

    def __init__(self):
        super(ExtranetXSModel, self).__init__()
        self.fc_global = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(16,
            16, 5, stride=1, padding=2), nn.ReLU(), nn.MaxPool1d(2, stride=
            2), nn.Conv1d(16, 32, 5, stride=1, padding=2), nn.ReLU())
        self.fc_local = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1, padding
            =2), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(16, 16, 5,
            stride=1, padding=2), nn.ReLU())
        self.final_layer = nn.Sequential(nn.Linear(58, 1), nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_starpars
        ):
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        out_global = F.max_pool1d(out_global, out_global.shape[-1])
        out_local = F.max_pool1d(out_local, out_local.shape[-1])
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)
        out = torch.cat([out_global, out_local, x_starpars.squeeze(1)], dim=1)
        out = self.final_layer(out)
        return out


def invert_tensor(tensor):
    """

    PURPOSE: FLIP A 1D TENSOR ALONG ITS AXIS
    INPUT: 1D TENSOR
    OUTPUT: INVERTED 1D TENSOR

    """
    idx = [i for i in range(tensor.size(0) - 1, -1, -1)]
    idx = torch.cuda.LongTensor(idx)
    inverted_tensor = tensor.index_select(0, idx)
    return inverted_tensor


<mask token>

<mask token>


class KeplerDataLoader(Dataset):
    """
    
    PURPOSE: DATA LOADER FOR KERPLER LIGHT CURVES
    INPUT: PATH TO DIRECTOR WITH LIGHT CURVES + INFO FILES
    OUTPUT: LOCAL + GLOBAL VIEWS, LABELS
    
    """

    def __init__(self, filepath):
        self.flist_global = np.sort(glob.glob(os.path.join(filepath,
            '*global.npy')))
        self.flist_local = np.sort(glob.glob(os.path.join(filepath,
            '*local.npy')))
        self.flist_info = np.sort(glob.glob(os.path.join(filepath,
            '*info.npy')))
        self.flist_global_cen = np.sort(glob.glob(os.path.join(filepath,
            '*global_cen_w.npy')))
        self.flist_local_cen = np.sort(glob.glob(os.path.join(filepath,
            '*local_cen_w.npy')))
        self.ids = np.sort([(x.split('/')[-1].split('_')[1] + '_' + x.split
            ('/')[-1].split('_')[2]) for x in self.flist_global])

    def __len__(self):
        return self.ids.shape[0]

    def __getitem__(self, idx):
        data_global = np.load(self.flist_global[idx])
        data_local = np.load(self.flist_local[idx])
        data_global_cen = np.load(self.flist_global_cen[idx])
        data_local_cen = np.load(self.flist_local_cen[idx])
        data_info = np.load(self.flist_info[idx])
        return (data_local, data_global, data_local_cen, data_global_cen,
            data_info[6:]), data_info[5]


class ExtranetModel(nn.Module):
    """
    
    PURPOSE: DEFINE EXTRANET MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    """

    def __init__(self, d_rate):
        super(ExtranetModel, self).__init__()
        droplayer = nn.Dropout(d_rate) if d_rate > 0.001 else nn.Dropout(0.0)
        self.fc_global = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1,
            padding=2), nn.ReLU(), nn.Conv1d(16, 16, 5, stride=1, padding=2
            ), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.Conv1d(16, 32, 5,
            stride=1, padding=2), nn.ReLU(), nn.Conv1d(32, 32, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.Conv1d(32,
            64, 5, stride=1, padding=2), nn.ReLU(), nn.Conv1d(64, 64, 5,
            stride=1, padding=2), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.
            Conv1d(64, 128, 5, stride=1, padding=2), droplayer, nn.ReLU(),
            nn.Conv1d(128, 128, 5, stride=1, padding=2), nn.MaxPool1d(5,
            stride=2), nn.Conv1d(128, 256, 5, stride=1, padding=2), nn.ReLU
            (), nn.Conv1d(256, 256, 5, stride=1, padding=2), nn.ReLU(), nn.
            MaxPool1d(5, stride=2))
        self.fc_local = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1, padding
            =2), nn.ReLU(), nn.Conv1d(16, 16, 5, stride=1, padding=2), nn.
            ReLU(), nn.MaxPool1d(7, stride=2), nn.Conv1d(16, 32, 5, stride=
            1, padding=2), nn.ReLU(), nn.Conv1d(32, 32, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(7, stride=2))
        self.final_layer = nn.Sequential(nn.Linear(16586, 512), nn.ReLU(),
            nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 512), nn.ReLU(),
            nn.Linear(512, 1), nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_star):
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)
        out = torch.cat([out_global, out_local, x_star.squeeze(1)], dim=1)
        out = self.final_layer(out)
        return out


class ExtranetXSModel(nn.Module):
    """
    
    PURPOSE: DEFINE EXTRANET-XS MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    """

    def __init__(self):
        super(ExtranetXSModel, self).__init__()
        self.fc_global = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(16,
            16, 5, stride=1, padding=2), nn.ReLU(), nn.MaxPool1d(2, stride=
            2), nn.Conv1d(16, 32, 5, stride=1, padding=2), nn.ReLU())
        self.fc_local = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1, padding
            =2), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(16, 16, 5,
            stride=1, padding=2), nn.ReLU())
        self.final_layer = nn.Sequential(nn.Linear(58, 1), nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_starpars
        ):
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        out_global = F.max_pool1d(out_global, out_global.shape[-1])
        out_local = F.max_pool1d(out_local, out_local.shape[-1])
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)
        out = torch.cat([out_global, out_local, x_starpars.squeeze(1)], dim=1)
        out = self.final_layer(out)
        return out


def invert_tensor(tensor):
    """

    PURPOSE: FLIP A 1D TENSOR ALONG ITS AXIS
    INPUT: 1D TENSOR
    OUTPUT: INVERTED 1D TENSOR

    """
    idx = [i for i in range(tensor.size(0) - 1, -1, -1)]
    idx = torch.cuda.LongTensor(idx)
    inverted_tensor = tensor.index_select(0, idx)
    return inverted_tensor


def train_model(n_epochs, kepler_data_loader, kepler_val_loader, model,
    criterion, optimizer):
    """

    PURPOSE: TRAIN MODEL

    INPUTS:  num_epoch = number of epochs for training
             kepler_data_loader = data loader for Kepler dataset
             model = model use for training
             criterion = criterion for calculating loss

    OUTPUT:  epoch_{train/val}_loss = training set loss for each epoch
             epoch_val_acc = validation set accuracy for each epoch
             epoch_val_ap = validation set avg. precision for each epoch
             final_val_pred = validation predictions from final model
             final_val_gt = validation ground truths

    """
    epoch_train_loss = []
    epoch_val_loss = []
    epoch_val_acc = []
    epoch_val_ap = []
    for epoch in tqdm(range(n_epochs)):
        train_loss = torch.zeros(1).cuda()
        for x_train_data, y_train in kepler_data_loader:
            (x_train_local, x_train_global, x_train_local_cen,
                x_train_global_cen, x_train_star) = x_train_data
            x_train_local = Variable(x_train_local).type(torch.FloatTensor
                ).cuda()
            x_train_global = Variable(x_train_global).type(torch.FloatTensor
                ).cuda()
            x_train_local_cen = Variable(x_train_local_cen).type(torch.
                FloatTensor).cuda()
            x_train_global_cen = Variable(x_train_global_cen).type(torch.
                FloatTensor).cuda()
            x_train_star = Variable(x_train_star).type(torch.FloatTensor).cuda(
                )
            y_train = Variable(y_train).type(torch.FloatTensor).cuda()
            for batch_ind in range(x_train_local.shape[0]):
                sig_noise = np.random.uniform(0, 1.0)
                local_noise = Variable(x_train_local[batch_ind].data.new(
                    x_train_local[batch_ind].size()).normal_(0.0, sig_noise))
                x_train_local[batch_ind] = x_train_local[batch_ind
                    ] + local_noise
                global_noise = Variable(x_train_global[batch_ind].data.new(
                    x_train_global[batch_ind].size()).normal_(0.0, sig_noise))
                x_train_global[batch_ind] = x_train_global[batch_ind
                    ] + global_noise
                if random() < 0.5:
                    x_train_local[batch_ind] = invert_tensor(x_train_local[
                        batch_ind])
                    x_train_global[batch_ind] = invert_tensor(x_train_global
                        [batch_ind])
                    x_train_local_cen[batch_ind] = invert_tensor(
                        x_train_local_cen[batch_ind])
                    x_train_global_cen[batch_ind] = invert_tensor(
                        x_train_global_cen[batch_ind])
            x_train_local = x_train_local.unsqueeze(1)
            x_train_global = x_train_global.unsqueeze(1)
            x_train_local_cen = x_train_local_cen.unsqueeze(1)
            x_train_global_cen = x_train_global_cen.unsqueeze(1)
            x_train_star = x_train_star.unsqueeze(1)
            y_train = y_train.unsqueeze(1)
            output_train = model(x_train_local, x_train_global,
                x_train_local_cen, x_train_global_cen, x_train_star)
            loss = criterion(output_train, y_train)
            train_loss += loss.data
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        epoch_train_loss.append(train_loss.cpu().numpy() / len(
            kepler_data_loader.dataset))
        val_pred, val_gt, val_loss, num_corr = [], [], 0, 0
        for x_val_data, y_val in kepler_val_loader:
            (x_val_local, x_val_global, x_val_local_cen, x_val_global_cen,
                x_val_star) = x_val_data
            x_val_local = Variable(x_val_local).type(torch.FloatTensor).cuda()
            x_val_global = Variable(x_val_global).type(torch.FloatTensor).cuda(
                )
            x_val_local_cen = Variable(x_val_local_cen).type(torch.FloatTensor
                ).cuda()
            x_val_global_cen = Variable(x_val_global_cen).type(torch.
                FloatTensor).cuda()
            x_val_star = Variable(x_val_star).type(torch.FloatTensor).cuda()
            y_val = Variable(y_val).type(torch.FloatTensor).cuda()
            x_val_local = x_val_local.unsqueeze(1)
            x_val_global = x_val_global.unsqueeze(1)
            x_val_local_cen = x_val_local_cen.unsqueeze(1)
            x_val_global_cen = x_val_global_cen.unsqueeze(1)
            x_val_star = x_val_star.unsqueeze(1)
            y_val = y_val.unsqueeze(1)
            output_val = model(x_val_local, x_val_global, x_val_local_cen,
                x_val_global_cen, x_val_star)
            loss_val = criterion(output_val, y_val)
            val_loss += loss_val.data
            output_pred = output_val >= 0.5
            num_corr += output_pred.eq(y_val.byte()).sum().item()
            val_pred.append(output_val.data.cpu().numpy())
            val_gt.append(y_val.data.cpu().numpy())
        epoch_val_loss.append(val_loss.cpu().numpy() / len(
            kepler_val_loader.dataset))
        epoch_val_acc.append(num_corr / len(kepler_val_loader.dataset))
        epoch_val_ap.append(average_precision_score(np.concatenate(val_gt).
            ravel(), np.concatenate(val_pred).ravel(), average=None))
    final_val_pred = np.concatenate(val_pred).ravel()
    final_val_gt = np.concatenate(val_gt).ravel()
    return (epoch_train_loss, epoch_val_loss, epoch_val_acc, epoch_val_ap,
        final_val_pred, final_val_gt)


<mask token>

<mask token>
parser = argparse.ArgumentParser()
parser.add_argument('n_epochs', help='number of epochs to use for training',
    type=int)
parser.add_argument('n_batches', help=
    'number of matches to use for training', type=int)
parser.add_argument('r_learn', help='learning rate for Adam optimizer',
    type=float)
parser.add_argument('d_rate', help='Dropout rate to use during training',
    type=float)
parser.add_argument('d_path', help=
    'path to data (should contain folders named val, test, train')
parser.add_argument('m_out', help='path for output model, files, and plots')
parser.add_argument('--fixed_seed', help='set if wanting to fix the seed',
    action='store_true')
parser.add_argument('--XS', help='use ExtranetXS model', action='store_true')
args = parser.parse_args()
if args.fixed_seed:
    torch.cuda.manual_seed(42)


class KeplerDataLoader(Dataset):
    """
    
    PURPOSE: DATA LOADER FOR KERPLER LIGHT CURVES
    INPUT: PATH TO DIRECTOR WITH LIGHT CURVES + INFO FILES
    OUTPUT: LOCAL + GLOBAL VIEWS, LABELS
    
    """

    def __init__(self, filepath):
        self.flist_global = np.sort(glob.glob(os.path.join(filepath,
            '*global.npy')))
        self.flist_local = np.sort(glob.glob(os.path.join(filepath,
            '*local.npy')))
        self.flist_info = np.sort(glob.glob(os.path.join(filepath,
            '*info.npy')))
        self.flist_global_cen = np.sort(glob.glob(os.path.join(filepath,
            '*global_cen_w.npy')))
        self.flist_local_cen = np.sort(glob.glob(os.path.join(filepath,
            '*local_cen_w.npy')))
        self.ids = np.sort([(x.split('/')[-1].split('_')[1] + '_' + x.split
            ('/')[-1].split('_')[2]) for x in self.flist_global])

    def __len__(self):
        return self.ids.shape[0]

    def __getitem__(self, idx):
        data_global = np.load(self.flist_global[idx])
        data_local = np.load(self.flist_local[idx])
        data_global_cen = np.load(self.flist_global_cen[idx])
        data_local_cen = np.load(self.flist_local_cen[idx])
        data_info = np.load(self.flist_info[idx])
        return (data_local, data_global, data_local_cen, data_global_cen,
            data_info[6:]), data_info[5]


class ExtranetModel(nn.Module):
    """
    
    PURPOSE: DEFINE EXTRANET MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    """

    def __init__(self, d_rate):
        super(ExtranetModel, self).__init__()
        droplayer = nn.Dropout(d_rate) if d_rate > 0.001 else nn.Dropout(0.0)
        self.fc_global = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1,
            padding=2), nn.ReLU(), nn.Conv1d(16, 16, 5, stride=1, padding=2
            ), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.Conv1d(16, 32, 5,
            stride=1, padding=2), nn.ReLU(), nn.Conv1d(32, 32, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.Conv1d(32,
            64, 5, stride=1, padding=2), nn.ReLU(), nn.Conv1d(64, 64, 5,
            stride=1, padding=2), nn.ReLU(), nn.MaxPool1d(5, stride=2), nn.
            Conv1d(64, 128, 5, stride=1, padding=2), droplayer, nn.ReLU(),
            nn.Conv1d(128, 128, 5, stride=1, padding=2), nn.MaxPool1d(5,
            stride=2), nn.Conv1d(128, 256, 5, stride=1, padding=2), nn.ReLU
            (), nn.Conv1d(256, 256, 5, stride=1, padding=2), nn.ReLU(), nn.
            MaxPool1d(5, stride=2))
        self.fc_local = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1, padding
            =2), nn.ReLU(), nn.Conv1d(16, 16, 5, stride=1, padding=2), nn.
            ReLU(), nn.MaxPool1d(7, stride=2), nn.Conv1d(16, 32, 5, stride=
            1, padding=2), nn.ReLU(), nn.Conv1d(32, 32, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(7, stride=2))
        self.final_layer = nn.Sequential(nn.Linear(16586, 512), nn.ReLU(),
            nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 512), nn.ReLU(),
            nn.Linear(512, 1), nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_star):
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)
        out = torch.cat([out_global, out_local, x_star.squeeze(1)], dim=1)
        out = self.final_layer(out)
        return out


class ExtranetXSModel(nn.Module):
    """
    
    PURPOSE: DEFINE EXTRANET-XS MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    """

    def __init__(self):
        super(ExtranetXSModel, self).__init__()
        self.fc_global = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1,
            padding=2), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(16,
            16, 5, stride=1, padding=2), nn.ReLU(), nn.MaxPool1d(2, stride=
            2), nn.Conv1d(16, 32, 5, stride=1, padding=2), nn.ReLU())
        self.fc_local = nn.Sequential(nn.Conv1d(2, 16, 5, stride=1, padding
            =2), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(16, 16, 5,
            stride=1, padding=2), nn.ReLU())
        self.final_layer = nn.Sequential(nn.Linear(58, 1), nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_starpars
        ):
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        out_global = F.max_pool1d(out_global, out_global.shape[-1])
        out_local = F.max_pool1d(out_local, out_local.shape[-1])
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)
        out = torch.cat([out_global, out_local, x_starpars.squeeze(1)], dim=1)
        out = self.final_layer(out)
        return out


def invert_tensor(tensor):
    """

    PURPOSE: FLIP A 1D TENSOR ALONG ITS AXIS
    INPUT: 1D TENSOR
    OUTPUT: INVERTED 1D TENSOR

    """
    idx = [i for i in range(tensor.size(0) - 1, -1, -1)]
    idx = torch.cuda.LongTensor(idx)
    inverted_tensor = tensor.index_select(0, idx)
    return inverted_tensor


def train_model(n_epochs, kepler_data_loader, kepler_val_loader, model,
    criterion, optimizer):
    """

    PURPOSE: TRAIN MODEL

    INPUTS:  num_epoch = number of epochs for training
             kepler_data_loader = data loader for Kepler dataset
             model = model use for training
             criterion = criterion for calculating loss

    OUTPUT:  epoch_{train/val}_loss = training set loss for each epoch
             epoch_val_acc = validation set accuracy for each epoch
             epoch_val_ap = validation set avg. precision for each epoch
             final_val_pred = validation predictions from final model
             final_val_gt = validation ground truths

    """
    epoch_train_loss = []
    epoch_val_loss = []
    epoch_val_acc = []
    epoch_val_ap = []
    for epoch in tqdm(range(n_epochs)):
        train_loss = torch.zeros(1).cuda()
        for x_train_data, y_train in kepler_data_loader:
            (x_train_local, x_train_global, x_train_local_cen,
                x_train_global_cen, x_train_star) = x_train_data
            x_train_local = Variable(x_train_local).type(torch.FloatTensor
                ).cuda()
            x_train_global = Variable(x_train_global).type(torch.FloatTensor
                ).cuda()
            x_train_local_cen = Variable(x_train_local_cen).type(torch.
                FloatTensor).cuda()
            x_train_global_cen = Variable(x_train_global_cen).type(torch.
                FloatTensor).cuda()
            x_train_star = Variable(x_train_star).type(torch.FloatTensor).cuda(
                )
            y_train = Variable(y_train).type(torch.FloatTensor).cuda()
            for batch_ind in range(x_train_local.shape[0]):
                sig_noise = np.random.uniform(0, 1.0)
                local_noise = Variable(x_train_local[batch_ind].data.new(
                    x_train_local[batch_ind].size()).normal_(0.0, sig_noise))
                x_train_local[batch_ind] = x_train_local[batch_ind
                    ] + local_noise
                global_noise = Variable(x_train_global[batch_ind].data.new(
                    x_train_global[batch_ind].size()).normal_(0.0, sig_noise))
                x_train_global[batch_ind] = x_train_global[batch_ind
                    ] + global_noise
                if random() < 0.5:
                    x_train_local[batch_ind] = invert_tensor(x_train_local[
                        batch_ind])
                    x_train_global[batch_ind] = invert_tensor(x_train_global
                        [batch_ind])
                    x_train_local_cen[batch_ind] = invert_tensor(
                        x_train_local_cen[batch_ind])
                    x_train_global_cen[batch_ind] = invert_tensor(
                        x_train_global_cen[batch_ind])
            x_train_local = x_train_local.unsqueeze(1)
            x_train_global = x_train_global.unsqueeze(1)
            x_train_local_cen = x_train_local_cen.unsqueeze(1)
            x_train_global_cen = x_train_global_cen.unsqueeze(1)
            x_train_star = x_train_star.unsqueeze(1)
            y_train = y_train.unsqueeze(1)
            output_train = model(x_train_local, x_train_global,
                x_train_local_cen, x_train_global_cen, x_train_star)
            loss = criterion(output_train, y_train)
            train_loss += loss.data
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        epoch_train_loss.append(train_loss.cpu().numpy() / len(
            kepler_data_loader.dataset))
        val_pred, val_gt, val_loss, num_corr = [], [], 0, 0
        for x_val_data, y_val in kepler_val_loader:
            (x_val_local, x_val_global, x_val_local_cen, x_val_global_cen,
                x_val_star) = x_val_data
            x_val_local = Variable(x_val_local).type(torch.FloatTensor).cuda()
            x_val_global = Variable(x_val_global).type(torch.FloatTensor).cuda(
                )
            x_val_local_cen = Variable(x_val_local_cen).type(torch.FloatTensor
                ).cuda()
            x_val_global_cen = Variable(x_val_global_cen).type(torch.
                FloatTensor).cuda()
            x_val_star = Variable(x_val_star).type(torch.FloatTensor).cuda()
            y_val = Variable(y_val).type(torch.FloatTensor).cuda()
            x_val_local = x_val_local.unsqueeze(1)
            x_val_global = x_val_global.unsqueeze(1)
            x_val_local_cen = x_val_local_cen.unsqueeze(1)
            x_val_global_cen = x_val_global_cen.unsqueeze(1)
            x_val_star = x_val_star.unsqueeze(1)
            y_val = y_val.unsqueeze(1)
            output_val = model(x_val_local, x_val_global, x_val_local_cen,
                x_val_global_cen, x_val_star)
            loss_val = criterion(output_val, y_val)
            val_loss += loss_val.data
            output_pred = output_val >= 0.5
            num_corr += output_pred.eq(y_val.byte()).sum().item()
            val_pred.append(output_val.data.cpu().numpy())
            val_gt.append(y_val.data.cpu().numpy())
        epoch_val_loss.append(val_loss.cpu().numpy() / len(
            kepler_val_loader.dataset))
        epoch_val_acc.append(num_corr / len(kepler_val_loader.dataset))
        epoch_val_ap.append(average_precision_score(np.concatenate(val_gt).
            ravel(), np.concatenate(val_pred).ravel(), average=None))
    final_val_pred = np.concatenate(val_pred).ravel()
    final_val_gt = np.concatenate(val_gt).ravel()
    return (epoch_train_loss, epoch_val_loss, epoch_val_acc, epoch_val_ap,
        final_val_pred, final_val_gt)


print("""
TRAINING MODEL...
""")
if args.XS:
    model = ExtranetXSModel().cuda()
else:
    d_rate = args.d_rate
    model = ExtranetModel(d_rate).cuda()
lr = args.r_learn
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.BCELoss()
batch_size = args.n_batches
n_epochs = args.n_epochs
kepler_train_data = KeplerDataLoader(filepath=os.path.join(args.d_path,
    'train'))
kepler_val_data = KeplerDataLoader(filepath=os.path.join(args.d_path, 'test'))
kepler_data_loader = DataLoader(kepler_train_data, batch_size=batch_size,
    shuffle=True, num_workers=4)
kepler_val_loader = DataLoader(kepler_val_data, batch_size=batch_size,
    shuffle=False, num_workers=4)
(loss_train_epoch, loss_val_epoch, acc_val_epoch, ap_val_epoch,
    pred_val_final, gt_val_final) = (train_model(n_epochs,
    kepler_data_loader, kepler_val_loader, model, criterion, optimizer))
print("""
CALCULATING METRICS...
""")
AP = average_precision_score(gt_val_final, pred_val_final, average=None)
print('   average precision = {0:0.4f}\n'.format(AP))
P, R, _ = precision_recall_curve(gt_val_final, pred_val_final)
thresh = [0.5, 0.6, 0.7, 0.8, 0.9]
prec_thresh, recall_thresh = np.zeros(len(thresh)), np.zeros(len(thresh))
for n, nval in enumerate(thresh):
    pred_byte = np.zeros(len(pred_val_final))
    for i, val in enumerate(pred_val_final):
        if val > nval:
            pred_byte[i] = 1.0
        else:
            pred_byte[i] = 0.0
    prec_thresh[n] = precision_score(gt_val_final, pred_byte)
    recall_thresh[n] = recall_score(gt_val_final, pred_byte)
    print('   thresh = {0:0.2f}, precision = {1:0.2f}, recall = {2:0.2f}'.
        format(thresh[n], prec_thresh[n], recall_thresh[n]))
    tn, fp, fn, tp = confusion_matrix(gt_val_final, pred_byte).ravel()
    print('      TN = {0:0}, FP = {1:0}, FN = {2:0}, TP = {3:0}'.format(tn,
        fp, fn, tp))
loss_train_batch = [(x.item() * batch_size) for x in loss_train_epoch]
loss_val_batch = [(x.item() * batch_size) for x in loss_val_epoch]
run = 0
pt_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(
    n_epochs) + '-lr' + str(lr) + '-pt.csv')
while os.path.isfile(pt_fname):
    run += 1
    pt_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' +
        str(n_epochs) + '-lr' + str(lr) + '-pt.csv')
df = pd.DataFrame({'gt': gt_val_final, 'pred': pred_val_final})
df.to_csv(pt_fname, index=False)
epochs_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' +
    str(n_epochs) + '-lr' + str(lr) + '-epoch.csv')
df = pd.DataFrame({'loss_train': loss_train_batch, 'loss_val':
    loss_val_batch, 'acc_val': acc_val_epoch, 'ap_val': ap_val_epoch})
df.to_csv(epochs_fname, index=False)
model_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str
    (n_epochs) + '-lr' + str(lr) + '-drop' + str(d_rate) + '-model.pth')
torch.save(model.state_dict(), model_fname)
print("""
OUTPUTTING MODEL + RESULTS @ """ + os.path.join(args.m_out,
    model_fname) + '\n')
fig = plt.figure(figsize=(7, 7))
ax = gridspec.GridSpec(2, 2)
ax.update(wspace=0.4, hspace=0.4)
ax1 = plt.subplot(ax[0, 0])
ax2 = plt.subplot(ax[0, 1])
ax3 = plt.subplot(ax[1, 0])
ax4 = plt.subplot(ax[1, 1])
ax1.set_xlabel('Precision', fontsize=10, labelpad=10)
ax1.set_ylabel('Recall', fontsize=10)
ax1.set_xlim([0.0, 1.0])
ax1.set_ylim([0.0, 1.0])
ax1.plot(R, P, linewidth=3, color='black')
ax2.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax2.set_ylabel('Loss', fontsize=10)
ax2.set_xlim([0.0, n_epochs])
ax2.set_ylim([0.0, np.max(loss_train_batch) * 1.5])
ax2.plot(np.arange(len(loss_train_batch)), loss_train_batch, linewidth=3,
    color='cadetblue')
ax2.plot(np.arange(len(loss_val_batch)), loss_val_batch, linewidth=3, color
    ='orangered')
ax3.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax3.set_ylabel('Average Precision', fontsize=10)
ax3.plot(np.arange(len(ap_val_epoch)), ap_val_epoch, linewidth=1.0, color=
    'orangered')
ax3.scatter(np.arange(len(ap_val_epoch)), ap_val_epoch, marker='o',
    edgecolor='orangered', facecolor='orangered', s=10, linewidth=0.5,
    alpha=0.5)
ax4.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax4.set_ylabel('Accuracy', fontsize=10)
ax4.plot(np.arange(len(acc_val_epoch)), acc_val_epoch, color='orangered',
    linewidth=1.0)
ax4.scatter(np.arange(len(acc_val_epoch)), acc_val_epoch, marker='o',
    edgecolor='orangered', facecolor='orangered', s=10, linewidth=0.5,
    alpha=0.5)
plot_fname = 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-plot.pdf'
plt.savefig(os.path.join(args.m_out, plot_fname), bbox_inches='tight', dpi=
    200, rastersized=True, alpha=True)


########################################
########### IMPORT PACKAGES ############
########################################

### standard packages
import os
import numpy as np
import glob as glob
from tqdm import tqdm 
from random import random
import argparse
import pandas as pd
import pdb

### torch packages
import torch
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch.autograd import Variable
import torch.nn as nn

### sklearn packages
from sklearn.metrics import precision_recall_curve
from sklearn.metrics import average_precision_score
from sklearn.metrics import confusion_matrix
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score

### plotting packages
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec


########################################
########### PARSE ARGUMENTS ############
########################################

### an example input to command line:
### python exonet.py 225 64 1e-5 '/data/kepler/new' '/data_sata1/ensembling/test'

### parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("n_epochs", help="number of epochs to use for training", type=int)
parser.add_argument("n_batches", help="number of matches to use for training", type=int)
parser.add_argument("r_learn", help="learning rate for Adam optimizer", type=float)
parser.add_argument("d_rate", help="Dropout rate to use during training", type=float)
parser.add_argument("d_path", help="path to data (should contain folders named val, test, train")
parser.add_argument("m_out", help="path for output model, files, and plots")
parser.add_argument("--fixed_seed", help="set if wanting to fix the seed", action="store_true")
parser.add_argument("--XS", help="use ExtranetXS model", action="store_true")
args = parser.parse_args()

### set manual seed
if args.fixed_seed:
    torch.cuda.manual_seed(42)


########################################
########### DEFINE CLASSES  ############
########################################

class KeplerDataLoader(Dataset):
    
    '''
    
    PURPOSE: DATA LOADER FOR KERPLER LIGHT CURVES
    INPUT: PATH TO DIRECTOR WITH LIGHT CURVES + INFO FILES
    OUTPUT: LOCAL + GLOBAL VIEWS, LABELS
    
    '''

    def __init__(self, filepath):

        ### list of global, local, and info files (assumes certain names of files)
        self.flist_global = np.sort(glob.glob(os.path.join(filepath, '*global.npy')))
        self.flist_local = np.sort(glob.glob(os.path.join(filepath, '*local.npy')))
        self.flist_info = np.sort(glob.glob(os.path.join(filepath, '*info.npy')))
        
        ### list of whitened centroid files
        self.flist_global_cen = np.sort(glob.glob(os.path.join(filepath, '*global_cen_w.npy')))
        self.flist_local_cen = np.sort(glob.glob(os.path.join(filepath, '*local_cen_w.npy')))
        
        ### ids = {TIC}_{TCE}
        self.ids = np.sort([(x.split('/')[-1]).split('_')[1] + '_' + (x.split('/')[-1]).split('_')[2] for x in self.flist_global])

    def __len__(self):

        return self.ids.shape[0]

    def __getitem__(self, idx):

        ### grab local and global views
        data_global = np.load(self.flist_global[idx])
        data_local = np.load(self.flist_local[idx])

        ### grab centroid views
        data_global_cen = np.load(self.flist_global_cen[idx])
        data_local_cen = np.load(self.flist_local_cen[idx])
        
        ### info file contains: [0]kic, [1]tce, [2]period, [3]epoch, [4]duration, [5]label)
        data_info = np.load(self.flist_info[idx])
        
        return (data_local, data_global, data_local_cen, data_global_cen, data_info[6:]), data_info[5]


class ExtranetModel(nn.Module):

    '''
    
    PURPOSE: DEFINE EXTRANET MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    '''
    
    def __init__(self, d_rate):

        ### initialize model
        super(ExtranetModel, self).__init__()
        droplayer = nn.Dropout(d_rate) if d_rate > 1e-3 else nn.Dropout(0.)

        ### define global convolutional lalyer
        self.fc_global = nn.Sequential(
            nn.Conv1d(2, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(16, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(5, stride=2),
            nn.Conv1d(16, 32, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(32, 32, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(5, stride=2),
            nn.Conv1d(32, 64, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(64, 64, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(5, stride=2),
            nn.Conv1d(64, 128, 5, stride=1, padding=2),
            droplayer,
            nn.ReLU(),
            nn.Conv1d(128, 128, 5, stride=1, padding=2),
            nn.MaxPool1d(5, stride=2),
            nn.Conv1d(128, 256, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(256, 256, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(5, stride=2),
        )

        ### define local convolutional lalyer
        self.fc_local = nn.Sequential(
            nn.Conv1d(2, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(16, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(7, stride=2),
            nn.Conv1d(16, 32, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(32, 32, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(7, stride=2),
        )

        ### define fully connected layer that combines both views
        self.final_layer = nn.Sequential(
            nn.Linear(16586, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            ### need output of 1 because using BCE for loss
            nn.Linear(512, 1),
            nn.Sigmoid())

    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_star):

        ### concatonate light curve and centroid data
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        
        ### get outputs of global and local convolutional layers
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        
        ### flattening outputs from convolutional layers into vector
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)

        ### concatonate global and local views with stellar parameters
        out = torch.cat([out_global, out_local, x_star.squeeze(1)], dim=1)
        out = self.final_layer(out)

        return out


class ExtranetXSModel(nn.Module):

    '''
    
    PURPOSE: DEFINE EXTRANET-XS MODEL ARCHITECTURE
    INPUT: GLOBAL + LOCAL LIGHT CURVES AND CENTROID CURVES, STELLAR PARAMETERS
    OUTPUT: BINARY CLASSIFIER
    
    '''

    def __init__(self):

        ### initializing the nn.Moduel (super) class
        ### (must do this first always)
        super(ExtranetXSModel, self).__init__()

        ### define global convolutional lalyer
        self.fc_global = nn.Sequential(
            nn.Conv1d(2, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(2, stride=2),
            nn.Conv1d(16, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(2, stride=2),
            nn.Conv1d(16, 32, 5, stride=1, padding=2),
            nn.ReLU(),
        )

        ### define the local convolutional layer
        self.fc_local = nn.Sequential(
            nn.Conv1d(2, 16, 5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(2, stride=2),
            nn.Conv1d(16, 16, 5, stride=1, padding=2),
            nn.ReLU(),
        )
                
        ### define fully connected layer that combines both views
        self.final_layer = nn.Sequential(
            nn.Linear(58, 1),
            nn.Sigmoid())
        
    ### define how to move forward through model
    def forward(self, x_local, x_global, x_local_cen, x_global_cen, x_starpars):

        ### concatonate light curve and centroid data
        x_local_all = torch.cat([x_local, x_local_cen], dim=1)
        x_global_all = torch.cat([x_global, x_global_cen], dim=1)
        out_global = self.fc_global(x_global_all)
        out_local = self.fc_local(x_local_all)
        
        ### do global pooling
        out_global = F.max_pool1d(out_global, out_global.shape[-1])
        out_local = F.max_pool1d(out_local, out_local.shape[-1])

        ### flattening outputs from convolutional layers into vector
        out_global = out_global.view(out_global.shape[0], -1)
        out_local = out_local.view(out_local.shape[0], -1)

        ### concatonate global and local views with stellar parameters
        out = torch.cat([out_global, out_local, x_starpars.squeeze(1)], dim=1)
        out = self.final_layer(out)

        return out


########################################
########## DEFINE FUNCTIONS ############
########################################

def invert_tensor(tensor):

    '''

    PURPOSE: FLIP A 1D TENSOR ALONG ITS AXIS
    INPUT: 1D TENSOR
    OUTPUT: INVERTED 1D TENSOR

    '''

    idx = [i for i in range(tensor.size(0)-1, -1, -1)]
    idx = torch.cuda.LongTensor(idx)
    inverted_tensor = tensor.index_select(0, idx)

    return inverted_tensor


def train_model(n_epochs, kepler_data_loader, kepler_val_loader, model, criterion, optimizer):

    '''

    PURPOSE: TRAIN MODEL

    INPUTS:  num_epoch = number of epochs for training
             kepler_data_loader = data loader for Kepler dataset
             model = model use for training
             criterion = criterion for calculating loss

    OUTPUT:  epoch_{train/val}_loss = training set loss for each epoch
             epoch_val_acc = validation set accuracy for each epoch
             epoch_val_ap = validation set avg. precision for each epoch
             final_val_pred = validation predictions from final model
             final_val_gt = validation ground truths

    '''

    ### empty arrays to fill per-epoch outputs
    epoch_train_loss = []
    epoch_val_loss = []
    epoch_val_acc = []
    epoch_val_ap = []

    ### loop over number of epochs of training
    for epoch in tqdm(range(n_epochs)):

        ####################
        ### for training set
        
        ### loop over batches
        train_loss = torch.zeros(1).cuda()
        for x_train_data, y_train in kepler_data_loader:
            
            ### get local view, global view, and label for training
            x_train_local, x_train_global, x_train_local_cen, x_train_global_cen, x_train_star = x_train_data
            x_train_local = Variable(x_train_local).type(torch.FloatTensor).cuda()
            x_train_global = Variable(x_train_global).type(torch.FloatTensor).cuda()    
            x_train_local_cen = Variable(x_train_local_cen).type(torch.FloatTensor).cuda()
            x_train_global_cen = Variable(x_train_global_cen).type(torch.FloatTensor).cuda()
            x_train_star = Variable(x_train_star).type(torch.FloatTensor).cuda()        
            y_train = Variable(y_train).type(torch.FloatTensor).cuda()
                      
            ### randomly invert half of light curves
            for batch_ind in range(x_train_local.shape[0]):     
                
                ### add random gaussian noise
                sig_noise = np.random.uniform(0, 1.0)
                local_noise = Variable(x_train_local[batch_ind].data.new(x_train_local[batch_ind].size()).normal_(0.0, sig_noise))
                x_train_local[batch_ind] = x_train_local[batch_ind] + local_noise
                global_noise = Variable(x_train_global[batch_ind].data.new(x_train_global[batch_ind].size()).normal_(0.0, sig_noise))
                x_train_global[batch_ind] = x_train_global[batch_ind] + global_noise

                if random() < 0.5:
                    x_train_local[batch_ind] = invert_tensor(x_train_local[batch_ind])
                    x_train_global[batch_ind] = invert_tensor(x_train_global[batch_ind])
                    x_train_local_cen[batch_ind] = invert_tensor(x_train_local_cen[batch_ind])
                    x_train_global_cen[batch_ind] = invert_tensor(x_train_global_cen[batch_ind])

            ### fix dimensions for next steps
            x_train_local = x_train_local.unsqueeze(1)
            x_train_global = x_train_global.unsqueeze(1)
            x_train_local_cen = x_train_local_cen.unsqueeze(1)
            x_train_global_cen = x_train_global_cen.unsqueeze(1)
            x_train_star = x_train_star.unsqueeze(1)
            y_train = y_train.unsqueeze(1)

            ### calculate loss using model
            output_train = model(x_train_local, x_train_global, x_train_local_cen, x_train_global_cen, x_train_star)
            loss = criterion(output_train, y_train)
            train_loss += loss.data

            ### train model (zero gradients and back propogate results)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        ### record training loss for this epoch (divided by size of training dataset)
        epoch_train_loss.append(train_loss.cpu().numpy() / len(kepler_data_loader.dataset))
        
        ######################
        ### for validation set
        
        ### loop over batches
        val_pred, val_gt, val_loss, num_corr = [], [], 0, 0
        for x_val_data, y_val in kepler_val_loader:
                        
            ### get local view, global view, and label for validating
            x_val_local, x_val_global, x_val_local_cen, x_val_global_cen, x_val_star = x_val_data
            x_val_local = Variable(x_val_local).type(torch.FloatTensor).cuda()
            x_val_global = Variable(x_val_global).type(torch.FloatTensor).cuda()
            x_val_local_cen = Variable(x_val_local_cen).type(torch.FloatTensor).cuda()
            x_val_global_cen = Variable(x_val_global_cen).type(torch.FloatTensor).cuda()
            x_val_star = Variable(x_val_star).type(torch.FloatTensor).cuda()
 
            ### fix dimensions for next steps
            y_val = Variable(y_val).type(torch.FloatTensor).cuda()
            x_val_local = x_val_local.unsqueeze(1)
            x_val_global = x_val_global.unsqueeze(1)
            x_val_local_cen = x_val_local_cen.unsqueeze(1)
            x_val_global_cen = x_val_global_cen.unsqueeze(1)
            x_val_star = x_val_star.unsqueeze(1)
            y_val = y_val.unsqueeze(1)

            ### calculate loss & add to sum over all batches
            output_val = model(x_val_local, x_val_global, x_val_local_cen, x_val_global_cen, x_val_star)
            loss_val = criterion(output_val, y_val)
            val_loss += loss_val.data

            ### get number of correct predictions using threshold=0.5
            ### & sum over all batches
            output_pred = output_val >= 0.5
            num_corr += output_pred.eq(y_val.byte()).sum().item()
                        
            ### record predictions and ground truth by model
            ### (used for AP per epoch; reset at each epoch; final values output)
            val_pred.append(output_val.data.cpu().numpy())
            val_gt.append(y_val.data.cpu().numpy())
            
        ### record validation loss calculate for this epoch (divided by size of validation dataset)
        epoch_val_loss.append(val_loss.cpu().numpy() / len(kepler_val_loader.dataset))
        
        ### record validation accuracy (# correct predictions in val set) for this epoch
        epoch_val_acc.append(num_corr / len(kepler_val_loader.dataset))
        
        ### calculate average precision for this epoch
        epoch_val_ap.append(average_precision_score(np.concatenate(val_gt).ravel(), np.concatenate(val_pred).ravel(), average=None))            
        
    ### grab final predictions and ground truths for validation set
    final_val_pred = np.concatenate(val_pred).ravel()
    final_val_gt = np.concatenate(val_gt).ravel()

    return epoch_train_loss, epoch_val_loss, epoch_val_acc, epoch_val_ap, final_val_pred, final_val_gt


########################################
############ TRAIN MODEL  ##############
########################################

### setup screen output
print("\nTRAINING MODEL...\n")

### initialize model; cuda puts it on GPU
if args.XS:
    model = ExtranetXSModel().cuda()
else:
    ### dropout rate for training
    d_rate = args.d_rate
    model = ExtranetModel(d_rate).cuda()

### learning rate
lr  = args.r_learn

### specify optimizer for learning to use for training
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

### specify loss function to use for training
criterion = nn.BCELoss()

### specify batch size to use for training
batch_size = args.n_batches 

### number of epochs to use for training
n_epochs = args.n_epochs 

### grab data using data loader
kepler_train_data = KeplerDataLoader(filepath=os.path.join(args.d_path, 'train'))
kepler_val_data = KeplerDataLoader(filepath=os.path.join(args.d_path, 'test'))
kepler_data_loader = DataLoader(kepler_train_data, batch_size=batch_size, shuffle=True, num_workers=4)
kepler_val_loader = DataLoader(kepler_val_data, batch_size=batch_size, shuffle=False, num_workers=4)

### train model
loss_train_epoch, loss_val_epoch, acc_val_epoch, ap_val_epoch, pred_val_final, gt_val_final  = train_model(n_epochs, kepler_data_loader, kepler_val_loader, model, criterion, optimizer)


########################################
####### CALCULATE STATISTICS ###########
########################################

### setup screen output
print("\nCALCULATING METRICS...\n")

### calculate average precision & precision-recall curves
AP = average_precision_score(gt_val_final, pred_val_final, average=None)
print("   average precision = {0:0.4f}\n".format(AP))
 
### calculate precision-recall curve
P, R, _ = precision_recall_curve(gt_val_final, pred_val_final)

### calculate confusion matrix based on different thresholds 
thresh = [0.5, 0.6, 0.7, 0.8, 0.9]
prec_thresh, recall_thresh = np.zeros(len(thresh)), np.zeros(len(thresh))
for n, nval in enumerate(thresh):
    pred_byte = np.zeros(len(pred_val_final))
    for i, val in enumerate(pred_val_final):
        if val > nval:
            pred_byte[i] = 1.0
        else:
            pred_byte[i] = 0.0
    prec_thresh[n] = precision_score(gt_val_final, pred_byte)
    recall_thresh[n] = recall_score(gt_val_final, pred_byte)
    print("   thresh = {0:0.2f}, precision = {1:0.2f}, recall = {2:0.2f}".format(thresh[n], prec_thresh[n], recall_thresh[n]))
    tn, fp, fn, tp = confusion_matrix(gt_val_final, pred_byte).ravel()
    print("      TN = {0:0}, FP = {1:0}, FN = {2:0}, TP = {3:0}".format(tn, fp, fn, tp))


########################################
######### OUTPUT MODEL + STATS  ########
########################################

### transform from loss per sample to loss per batch (multiple by batch size to compare to Chris')
loss_train_batch = [x.item()* batch_size for x in loss_train_epoch]
loss_val_batch = [x.item()* batch_size for x in loss_val_epoch]

### setup output
run = 0

### output predictions & ground truth
pt_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-lr' + str(lr) + '-pt.csv')
while os.path.isfile(pt_fname):
    run +=1
    pt_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-lr' + str(lr) + '-pt.csv')
df = pd.DataFrame({"gt" : gt_val_final, "pred" : pred_val_final})
df.to_csv(pt_fname, index=False)

### output per-iteration values
epochs_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-lr' + str(lr) + '-epoch.csv')
df = pd.DataFrame({"loss_train":loss_train_batch, "loss_val":loss_val_batch, "acc_val":acc_val_epoch, "ap_val":ap_val_epoch})
df.to_csv(epochs_fname, index=False)

### save model
model_fname = os.path.join(args.m_out, 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-lr' + str(lr) + '-drop' + str(d_rate) + '-model.pth')
torch.save(model.state_dict(), model_fname)
print("\nOUTPUTTING MODEL + RESULTS @ " + os.path.join(args.m_out, model_fname) + "\n")


########################################
################ MAKE PLOTS ############
########################################

### setup figure
fig = plt.figure(figsize=(7, 7))
ax = gridspec.GridSpec(2,2)
ax.update(wspace = 0.4, hspace = 0.4)
ax1 = plt.subplot(ax[0,0])
ax2 = plt.subplot(ax[0,1])
ax3 = plt.subplot(ax[1,0])
ax4 = plt.subplot(ax[1,1])

### plot precision-recall curve
ax1.set_xlabel('Precision', fontsize=10, labelpad=10)
ax1.set_ylabel('Recall', fontsize=10)
ax1.set_xlim([0.0, 1.0])
ax1.set_ylim([0.0, 1.0])
ax1.plot(R, P, linewidth=3, color='black')

### plot loss curve for training and validation sets
ax2.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax2.set_ylabel('Loss', fontsize=10)
ax2.set_xlim([0.0, n_epochs])
ax2.set_ylim([0.0, np.max(loss_train_batch)*1.5])
ax2.plot(np.arange(len(loss_train_batch)), loss_train_batch, linewidth=3, color='cadetblue')
ax2.plot(np.arange(len(loss_val_batch)), loss_val_batch, linewidth=3, color='orangered')

### plot average precision per epoch
ax3.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax3.set_ylabel('Average Precision', fontsize=10)
ax3.plot(np.arange(len(ap_val_epoch)), ap_val_epoch, linewidth=1.0, color='orangered')
ax3.scatter(np.arange(len(ap_val_epoch)), ap_val_epoch, marker='o', edgecolor='orangered', facecolor='orangered', s=10, linewidth=0.5, alpha=0.5)

### plot accuracy per epoch
ax4.set_xlabel('Epoch', fontsize=10, labelpad=10)
ax4.set_ylabel('Accuracy', fontsize=10)
ax4.plot(np.arange(len(acc_val_epoch)), acc_val_epoch, color='orangered', linewidth=1.0)
ax4.scatter(np.arange(len(acc_val_epoch)), acc_val_epoch, marker='o', edgecolor='orangered', facecolor='orangered', s=10, linewidth=0.5, alpha=0.5)

### save plot
plot_fname = 'r' + str(run).zfill(2) + '-i' + str(n_epochs) + '-plot.pdf'
plt.savefig(os.path.join(args.m_out, plot_fname), bbox_inches='tight', dpi=200, rastersized=True, alpha=True)

