In [ ]:
from __future__ import print_function

import numpy as np
import glob


## Boilerplate path hack to give access to full clustered_SNe package
import sys, os
if __package__ is None:
    if os.pardir not in sys.path[0]:
        file_dir = os.getcwd()
        sys.path.insert(0, os.path.join(file_dir, 
                                        os.pardir, 
                                        os.pardir))
        
from clustered_SNe import data_dir_default
from clustered_SNe.analysis.constants import m_proton, pc, yr, M_solar, \
                                   metallicity_solar
from clustered_SNe.analysis.parse import Overview, RunSummary, \
                                         Inputs, \
                                         get_full_id_from_partial_id, \
                                         checkpoint_num_from_filename

from clustered_SNe.analysis.database_helpers import session, \
                                                Simulation, \
                                                Simulation_Inputs, \
                                                Simulation_Status

In [ ]:
data_dir = data_dir_default

In [ ]:
all_ids = [Overview(filename).id for filename in glob.glob(os.path.join(data_dir,
                                                                        "*_overview.dat"))]

In [ ]:
ids = []
for id in all_ids:
    filenames = glob.glob(os.path.join(data_dir, id + "_checkpoint_*.dat"))
    last_filename = sorted(filenames)[-1]
    last_file_number = checkpoint_num_from_filename(last_filename)
    if (last_file_number % 100) == 99:
        if session.query(Simulation_Status).get(id).status == "Ready":
            ids.append(id)
        
len(ids)

In [ ]:
data_dir = "../tmp/"

In [ ]:
ids = [
]

In [ ]:
restart_dir = "restart_incomplete"

qsub_systems = set(["SGE", "PBS"])
qsub_system = "PBS"
assert(qsub_system in qsub_systems)
    
f_submit = open("submit_script", mode="w")
f_submit.write("#!/bin/bash \n")

f_copy = open("copy_script", mode="w")
f_copy.write("#!/bin/bash \n")
f_copy.write("mkdir -p ../" + restart_dir + "\n")

for id in ids:

    print("id: ", id)
    
    overview = Overview(os.path.join(data_dir, id + "_overview.dat"))
    if overview.num_SNe == 0:
        continue
    
    filenames = glob.glob(os.path.join(data_dir, id + "_checkpoint_*.dat"))
    filenames = sorted(filenames)
    last_checkpoint_filename  = filenames[-1]
    last_checkpoint_number = checkpoint_num_from_filename(last_checkpoint_filename)
    
    with open(last_checkpoint_filename, 'r') as f:
        line = f.readline()
        last_checkpoint_time = float(line.split()[3])
    
    restart_checkpoints_needed = 100 - ((last_checkpoint_number+1) % 100)
    print(last_checkpoint_filename)
    print(((last_checkpoint_number+1) % 100))
    print("new checkpoints_needed: ", restart_checkpoints_needed)
    
    restart_time_needed = overview.SNe_times.max() * 3 - last_checkpoint_time
    if last_checkpoint_number >= 99:
        restart_time_needed = 3e15
    print("extra time needed: ", restart_time_needed, "[s]")
    
    try:
        cooling_redshift = overview.inputs.Cooling_Redshift
        print("redshift: ", cooling_redshift)
    except AttributeError:
        cooling_redshift = 0

    print("")

    CFL = .1
    
    batch_name = "restart.batch"
    if qsub_system == "PBS":
        batch_name = id + "_" + batch_name
        
    f_submit.write("qsub " +
                       "-v RESTART_ID=" + id + ","
                       + "RESTART_N_CHECKPOINTS=" + str(restart_checkpoints_needed) + ","
                       + "RESTART_DELTA_TIME=" + str(restart_time_needed) + ","
                       + "RESTART_CFL=" + str(CFL) + ","
                       + "RESTART_COOLING_REDSHIFT=" + str(cooling_redshift)
                    + " -N " + batch_name + " "
                    + " ../scripts/" + qsub_system + "/restart.batch.env_variables \n")
    f_copy.write("cp -p " + id + "* ../" + restart_dir + " \n")
        
f_submit.close()
f_copy.close()

In [ ]:
for id in ids:
    session.query(Simulation_Status).filter(Simulation_Sta?tus.id == id).update({"status":"Running"})
    sim_status = session.query(Simulation_Status).get(id)
    print(sim_status.status)
    print(id)

In [ ]:
session.commit()