-
Notifications
You must be signed in to change notification settings - Fork 1
/
prep_snec.py
145 lines (128 loc) · 6.35 KB
/
prep_snec.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
import argparse
import glob
import os
import stat
import sys
import shutil
import numpy as np
import yaml
import profile_manipulation
import parfile_organizer
def get_input_parameters(filename='user_def.param'):
'''
Read user supplied file to get initial setup parameters
Input:
filename: name of user supplied file with inputs
Output:
parameters: dictionary of parameters
'''
ofile = open('user_def.param', 'r')
parameters = yaml.load(ofile)
for iparam in parameters['expand_params']:
parameters[iparam] = np.arange(*parameters[iparam])
return parameters
def create_directory(i_ni_mass, i_ni_mix, imass, ienergy, idensity, iradius, overwrite=False):
'''
Create directory structure, one directory for each model run
'''
path = os.path.join('snec_models',
'Ni_mass_{:1.4f}'.format(i_ni_mass),
'Ni_mixing_{:1.1f}'.format(i_ni_mix),
'M{:2.1f}'.format(imass), #TODO - figure out formating so if string is single digit, prepended with a 0
'E_{:1.3f}'.format(ienergy),
'K_{:2.1f}'.format(idensity), #TODO - figure out formating so if string is single digit, prepended with a 0
'R_{}'.format(int(iradius)))
if os.path.exists(path) and overwrite is False:
overwrite=input('overwrite existing path (y, n): {} '.format(path))
if overwrite == 'y':
overwrite = True
else:
sys.exit()
if (os.path.exists(path) and overwrite is True):
shutil.rmtree(path)
os.makedirs(path, exist_ok=True)
os.makedirs(os.path.join(path, 'profiles'), exist_ok=True)
os.makedirs(os.path.join(path, 'Data'), exist_ok=True)
return path
def copy_files(basepath, imass, idensity, iradius):
'''
Copy relevant profiles, executable, and opacity tables to each directory
'''
profile_basename = os.path.join('sukhbold_profiles_wind',
'M{:2.1f}'.format(imass), #TODO - figure out formating so if string is single digit, prepended with a 0
'K_{:2.1f}'.format(idensity), #TODO - figure out formating so if string is single digit, prepended with a 0
'R_{}'.format(int(iradius)))
profile_list = glob.glob(os.path.join(profile_basename, 's*')) #get the .iso and .dat files
for ifile in profile_list:
print('Copying {} to {}'.format(ifile, os.path.join(basepath, 'profiles', os.path.basename(ifile))))
shutil.copyfile(ifile, os.path.join(basepath, 'profiles', os.path.basename(ifile)))
shutil.copytree('tables', os.path.join(basepath, 'tables'))
shutil.copyfile('snec', os.path.join(basepath, 'snec'))
shutil.copymode('snec', os.path.join(basepath, 'snec'))
def write_directory_list(path_list):
'''
Create a file to execute snec for each array job
'''
ofile = open('input_dir_list.txt'.format(array_num), 'w')
for ipath in path_list:
ofile.write('{}\n'.format(ipath))
ofile.close()
os.chmod('input_dir_list.txt'.format(array_num), stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH | stat.S_IWUSR | stat.S_IXUSR)
def write_sbatch_job(timeout, array_num):
'''
Creates script to be called by sbatch
'''
ofile = open('snec_array_job.sh', 'w')
ofile.write(
'''#! /bin/bash -l
#SBATCH -t {}:00:00
#SBATCH --mail-type=begin
#SBATCH --mail-type=end
#SBATCH [email protected]
#SBATCH --job-name='M={}-{},E={}-{},K={}-{}, R={}-{}'
#SBATCH -o slurm_%A_%a.out # Standard output
#SBATCH -e slurm_%A_%a.err # Standard error
#SBATCH --array=1-{}
export SEEDFILE=input_dir_list.txt
export SEED=$(cat $SEEDFILE | head -n $SLURM_ARRAY_TASK_ID | tail -n 1)
mkdir -p /scratch/bostroem/job_$SLURM_ARRAY_TASK_ID
cp -r $SEED/* /scratch/bostroem/job_$SLURM_ARRAY_TASK_ID
cd /scratch/bostroem/job_$SLURM_ARRAY_TASK_ID
srun ./snec &>snec.out
#make sure you captured all your data
cp /scratch/bostroem/job_$SLURM_ARRAY_TASK_ID/snec.out $SEED
cp -r /scratch/bostroem/job_$SLURM_ARRAY_TASK_ID/Data $SEED/Data
rm -rf /scratch/bostroem/job_$SLURM_ARRAY_TASK_ID
'''.format(timeout, parameters['mass'][0], parameters['mass'][-1],
parameters['explosion_energy'][0], parameters['explosion_energy'][-1],
parameters['density_1D'][0], parameters['density_1D'][-1],
parameters['wind_extent'][0], parameters['wind_extent'][-1],
array_num) )
ofile.close()
os.chmod('snec_array_job.sh', stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH | stat.S_IWUSR | stat.S_IRUSR)
if __name__ == "__main__":
#If running on peloton, need to add the line 'module load bio3' to bash_profile
parser = argparse.ArgumentParser()
parser.add_argument('timeout', type=int,
help='the maximum time for a job to run, this should be twice what you expect it to take')
args = parser.parse_args()
snec_dir = os.getcwd()
parameters = get_input_parameters()
profile_manipulation.add_wind(parameters)
array_num=0
path_list = []
for indx_ni_mass, i_ni_mass in enumerate(parameters['ni_mass']):
for indx_ni_mix, i_ni_mix in enumerate(parameters['ni_mix']):
for indx_mass, imass in enumerate(parameters['mass']):
for indx_energy, ienergy in enumerate(parameters['explosion_energy']):
for indx_density, idensity in enumerate(parameters['density_1D']):
for indx_radius, iradius in enumerate(parameters['wind_extent']):
basepath=create_directory(i_ni_mass, i_ni_mix, imass, ienergy, idensity, iradius, overwrite=False)
copy_files(basepath, imass, idensity, iradius)
parfile_organizer.write_parfile(i_ni_mass, i_ni_mix, imass, ienergy, parameters, basepath)
array_num +=1
path_list.append(os.path.join(snec_dir, basepath))
write_directory_list(path_list)
write_sbatch_job(args.timeout, array_num)
print('Ready to run, start your jobs with the command:')
print('sbatch --partition low ./snec_array_job.sh'.format(array_num))