'''

#########################################################

This file containts classes to handle data processing

Author: Julian Kates-Harbeck, jkatesharbeck@g.harvard.edu

This work was supported by the DOE CSGF program.

#########################################################

'''

from __future__ import print_function

from os import listdir,remove

import time

import sys

import os

import numpy as np

import pathos.multiprocessing as mp

from plasma.utils.processing import *

from plasma.primitives.shots import ShotList

from plasma.utils.downloading import mkdirdepth

class Preprocessor(object):

def __init__(self,conf):

self.conf = conf

def clean_shot_lists(self):

shot_list_dir = self.conf['paths']['shot_list_dir']

paths = [os.path.join(shot_list_dir, f) for f in listdir(shot_list_dir) if os.path.isfile(os.path.join(shot_list_dir, f))]

for path in paths:

self.clean_shot_list(path)

def clean_shot_list(self,path):

data = np.loadtxt(path)

ending_idx = path.rfind('.')

new_path = append_to_filename(path,'_clear')

if len(np.shape(data)) < 2:

#nondisruptive

nd_times = -1.0*np.ones_like(data)

data_two_column = np.vstack((data,nd_times)).transpose()

np.savetxt(new_path,data_two_column,fmt = '%d %f')

print('created new file: {}'.format(new_path))

print('deleting old file: {}'.format(path))

os.remove(path)

def all_are_preprocessed(self):

return os.path.isfile(self.get_shot_list_path())

def preprocess_all(self):

conf = self.conf

shot_files_all = conf['paths']['shot_files_all']

# shot_files_train = conf['paths']['shot_files']

# shot_files_test = conf['paths']['shot_files_test']

# shot_list_dir = conf['paths']['shot_list_dir']

use_shots = conf['data']['use_shots']

train_frac = conf['training']['train_frac']

use_shots_train = int(round(train_frac*use_shots))

use_shots_test = int(round((1-train_frac)*use_shots))

# print(use_shots_train)

# print(use_shots_test) #each print out 100,000

# if len(shot_files_test) > 0:

# return self.preprocess_from_files(shot_list_dir,shot_files_train,machines_train,use_shots_train) + \

# self.preprocess_from_files(shot_list_dir,shot_files_test,machines_train,use_shots_test)

# else:

return self.preprocess_from_files(shot_files_all,use_shots)

def preprocess_from_files(self,shot_files,use_shots):

#all shots, including invalid ones

all_signals = self.conf['paths']['all_signals']

shot_list = ShotList()

shot_list.load_from_shot_list_files_objects(shot_files,all_signals)

shot_list_picked = shot_list.random_sublist(use_shots)

#empty

used_shots = ShotList()

use_cores = max(1,mp.cpu_count()-2)

pool = mp.Pool(use_cores)

print('running in parallel on {} processes'.format(pool._processes))

start_time = time.time()

for (i,shot) in enumerate(pool.imap_unordered(self.preprocess_single_file,shot_list_picked)):

#for (i,shot) in enumerate(map(self.preprocess_single_file,shot_list_picked)):

sys.stdout.write('\r{}/{}'.format(i,len(shot_list_picked)))

used_shots.append_if_valid(shot)

pool.close()

pool.join()

print('Finished Preprocessing {} files in {} seconds'.format(len(shot_list_picked),time.time()-start_time))

print('Omitted {} shots of {} total.'.format(len(shot_list_picked) - len(used_shots),len(shot_list_picked)))

print('{}/{} disruptive shots'.format(used_shots.num_disruptive(),len(used_shots)))

if len(used_shots) == 0:

print("WARNING: All shots were omitted, please ensure raw data is complete and available at {}.".format(self.conf['paths']['signal_prepath']))

return used_shots

def preprocess_single_file(self,shot):

processed_prepath = self.conf['paths']['processed_prepath']

recompute = self.conf['data']['recompute']

# print('({}/{}): '.format(num_processed,use_shots))

if recompute or not shot.previously_saved(processed_prepath):

shot.preprocess(self.conf)

shot.save(processed_prepath)

else:

try:

shot.restore(processed_prepath,light=True)

sys.stdout.write('\r{} exists.'.format(shot.number))

except:

shot.preprocess(self.conf)

shot.save(processed_prepath)

sys.stdout.write('\r{} exists but corrupted, resaved.'.format(shot.number))

shot.make_light()

return shot

def get_individual_channel_dirs(self):

signals_dirs = self.conf['paths']['signals_dirs']

def get_shot_list_path(self):

return self.conf['paths']['saved_shotlist_path']

def load_shotlists(self):

path = self.get_shot_list_path()

data = np.load(path,encoding="latin1")

shot_list_train = data['shot_list_train'][()]

shot_list_validate = data['shot_list_validate'][()]

shot_list_test = data['shot_list_test'][()]

if isinstance(shot_list_train,ShotList):

return shot_list_train,shot_list_validate,shot_list_test

else:

return ShotList(shot_list_train),ShotList(shot_list_validate),ShotList(shot_list_test)

def save_shotlists(self,shot_list_train,shot_list_validate,shot_list_test):

path = self.get_shot_list_path()

mkdirdepth(path)

np.savez(path,shot_list_train=shot_list_train,shot_list_validate=shot_list_validate,shot_list_test=shot_list_test)

def apply_bleed_in(conf,shot_list_train,shot_list_validate,shot_list_test):

np.random.seed(2)

num = conf['data']['bleed_in']

new_shots = []

if num > 0:

shot_list_bleed = ShotList()

print('applying bleed in with {} disruptive shots\n'.format(num))

num_total = len(shot_list_test)

num_d = shot_list_test.num_disruptive()

num_nd = num_total - num_d

assert(num_d >= num), "Not enough disruptive shots {} to cover bleed in {}".format(num_d,num)

num_sampled_d = 0

num_sampled_nd = 0

while num_sampled_d < num:

s = shot_list_test.sample_shot()

shot_list_bleed.append(s)

if conf['data']['bleed_in_remove_from_test']:

shot_list_test.remove(s)

if s.is_disruptive:

num_sampled_d += 1

else:

num_sampled_nd += 1

print("Sampled {} shots, {} disruptive, {} nondisruptive".format(num_sampled_nd+num_sampled_d,num_sampled_d,num_sampled_nd))

print("Before adding: training shots: {} validation shots: {}".format(len(shot_list_train),len(shot_list_validate)))

assert(num_sampled_d == num)

if conf['data']['bleed_in_equalize_sets']:#add bleed-in shots to training and validation set repeatedly

print("Applying equalized bleed in")

for shot_list_curr in [shot_list_train,shot_list_validate]:

for i in range(len(shot_list_curr)):

s = shot_list_bleed.sample_shot()

shot_list_curr.append(s)

elif conf['data']['bleed_in_repeat_fac'] > 1:

repeat_fac = conf['data']['bleed_in_repeat_fac']

print("Applying bleed in with repeat factor {}".format(repeat_fac))

num_to_sample = int(round(repeat_fac*len(shot_list_bleed)))

for i in range(num_to_sample):

s = shot_list_bleed.sample_shot()

shot_list_train.append(s)

shot_list_validate.append(s)

else: #add each shot only once

print("Applying bleed in without repetition")

for s in shot_list_bleed:

shot_list_train.append(s)

shot_list_validate.append(s)

print("After adding: training shots: {} validation shots: {}".format(len(shot_list_train),len(shot_list_validate)))

print("Added bleed in shots to training and validation sets")

# if num_d > 0:

# for i in range(num):

# s = shot_list_test.sample_single_class(True)

# shot_list_train.append(s)

# shot_list_validate.append(s)

# if conf['data']['bleed_in_remove_from_test']:

# shot_list_test.remove(s)

# else:

# print('No disruptive shots in test set, omitting bleed in')

# if num_nd > 0:

# for i in range(num):

# s = shot_list_test.sample_single_class(False)

# shot_list_train.append(s)

# shot_list_validate.append(s)

# if conf['data']['bleed_in_remove_from_test']:

# shot_list_test.remove(s)

# else:

# print('No nondisruptive shots in test set, omitting bleed in')

return shot_list_train,shot_list_validate,shot_list_test

def guarantee_preprocessed(conf):

pp = Preprocessor(conf)

if pp.all_are_preprocessed():

print("shots already processed.")

shot_list_train,shot_list_validate,shot_list_test = pp.load_shotlists()

else:

print("preprocessing all shots",end='')

pp.clean_shot_lists()

shot_list = pp.preprocess_all()

shot_list.sort()

shot_list_train,shot_list_test = shot_list.split_train_test(conf)

num_shots = len(shot_list_train) + len(shot_list_test)

validation_frac = conf['training']['validation_frac']

if validation_frac <= 0.05:

print('Setting validation to a minimum of 0.05')

validation_frac = 0.05

shot_list_train,shot_list_validate = shot_list_train.split_direct(1.0-validation_frac,do_shuffle=True)

pp.save_shotlists(shot_list_train,shot_list_validate,shot_list_test)

shot_list_train,shot_list_validate,shot_list_test = apply_bleed_in(conf,shot_list_train,shot_list_validate,shot_list_test)

print('validate: {} shots, {} disruptive'.format(len(shot_list_validate),shot_list_validate.num_disruptive()))

print('training: {} shots, {} disruptive'.format(len(shot_list_train),shot_list_train.num_disruptive()))

print('testing: {} shots, {} disruptive'.format(len(shot_list_test),shot_list_test.num_disruptive()))

print("...done")

return shot_list_train,shot_list_validate,shot_list_test