Skip to content

Commit e8199d9

Browse files
aaronjridleyaaronjridley
committed
FEAT: create series of restarts
1 parent 4a9c804 commit e8199d9

File tree

1 file changed

+276
-0
lines changed

1 file changed

+276
-0
lines changed

‎srcPython/run_restarts.py

Lines changed: 276 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,276 @@
1+
#!/usr/bin/env python3
2+
3+
import argparse
4+
import datetime as dt
5+
import json
6+
import os
7+
8+
# ----------------------------------------------------------------------
9+
# Function to parse input arguments
10+
# ----------------------------------------------------------------------
11+
12+
def get_args_restart():
13+
14+
parser = argparse.ArgumentParser(description =
15+
'Run series of restarts for Aether')
16+
parser.add_argument('-input',
17+
help = 'aether json file for whole run',
18+
default = 'aether.json.whole')
19+
20+
parser.add_argument('-mpi',
21+
help = 'mpirun command',
22+
default = 'mpirun')
23+
24+
parser.add_argument('-post',
25+
help = 'post processing command',
26+
default = '~/bin/postAether.py')
27+
28+
parser.add_argument('-aether',
29+
help = 'aether command',
30+
default = './aether')
31+
32+
parser.add_argument('-rundir',
33+
help = 'path to the default run directory',
34+
default = './share/run')
35+
36+
parser.add_argument('-restarts',
37+
help = 'number of restarts',
38+
default = 1, type = int)
39+
40+
parser.add_argument('-ensembles',
41+
help = 'number of ensemble members',
42+
default = 3, type = int)
43+
44+
parser.add_argument('-blocks',
45+
help = 'number of blocks in each ensemble member',
46+
default = 1, type = int)
47+
48+
parser.add_argument('-totaltime',
49+
help = 'total time (in seconds) of run (-1 = ignore)',
50+
default = -1, type = int)
51+
52+
parser.add_argument('-test', \
53+
help='just output the files and dont run', \
54+
action="store_true")
55+
56+
parser.add_argument('-dowhole', \
57+
help='include the whole run for comparison', \
58+
action="store_true")
59+
60+
args = parser.parse_args()
61+
62+
return args
63+
64+
# ----------------------------------------------------------------------
65+
# Get the start time, end time, and delta-t from the whole run dict:
66+
# ----------------------------------------------------------------------
67+
68+
def get_times(wholeDict):
69+
70+
if ('StartTime' in wholeDict):
71+
st = wholeDict['StartTime']
72+
startTime = dt.datetime(st[0], st[1], st[2], st[3], st[4], st[5])
73+
else:
74+
startTime = None
75+
if ('EndTime' in wholeDict):
76+
et = wholeDict['EndTime']
77+
endTime = dt.datetime(et[0], et[1], et[2], et[3], et[4], et[5])
78+
else:
79+
endTime = None
80+
if (startTime and endTime):
81+
deltaTime = (endTime - startTime).total_seconds()
82+
else:
83+
deltaTime = -1.0
84+
85+
return startTime, endTime, deltaTime
86+
87+
# ----------------------------------------------------------------------
88+
# run os command
89+
# ----------------------------------------------------------------------
90+
91+
def run(command, isTest, isVerbose):
92+
if (isVerbose):
93+
print(' -> Running : ', command)
94+
if (not isTest):
95+
os.system(command)
96+
return
97+
98+
# ----------------------------------------------------------------------
99+
# Main code:
100+
# ----------------------------------------------------------------------
101+
102+
if __name__ == '__main__':
103+
104+
args = get_args_restart()
105+
106+
isVerbose = True
107+
108+
inFile = args.input
109+
nRestarts = args.restarts
110+
nMembers = args.ensembles
111+
nBlocks = args.blocks
112+
113+
nProcs = nBlocks * nMembers
114+
115+
aetherCommand = args.aether
116+
117+
runDir = args.rundir
118+
notInRunDir = False
119+
120+
if (not os.path.exists(aetherCommand)):
121+
print('-> Can not find aether command : ', aetherCommand)
122+
notInRunDir = True
123+
124+
if (not os.path.exists('./UA')):
125+
print('-> Can not find UA directory!')
126+
notInRunDir = True
127+
128+
if (notInRunDir):
129+
if (os.path.exists(runDir)):
130+
print('Found rundir... copying')
131+
dirPre = './run.restarts'
132+
cdPre = 'cd ' + dirPre + ' ; '
133+
cdPost = ' ; cd ..'
134+
command = 'rm -rf ' + dirPre
135+
command = 'cp -r ' + runDir + ' ' + dirPre
136+
run(command, args.test, isVerbose)
137+
else:
138+
print('-> Can not find rundir to copy!')
139+
print('use -rundir to tell where to find the default run directory')
140+
exit()
141+
else:
142+
cdPre = ''
143+
cdPost = ''
144+
dirPre = '.'
145+
146+
runCommand = cdPre + args.mpi + ' -np %d ' % nProcs
147+
runCommand = runCommand + aetherCommand + cdPost
148+
149+
postCommand = args.post + ' -rm'
150+
151+
print('Number of Ensemble Members : ', nMembers)
152+
print('Number of blocks in each member : ', nBlocks)
153+
print(' --> nProcessors : ', nProcs)
154+
155+
with open(inFile, 'r') as filePointer:
156+
# Reading from json file
157+
wholeDict = json.load(filePointer)
158+
filePointer.close()
159+
160+
sTime, eTime, dTime = get_times(wholeDict)
161+
162+
if (args.totaltime > 0):
163+
dTime = args.totaltime
164+
165+
if (dTime < 0.0):
166+
print('Some sort of error determining time interval for run')
167+
print('please check StartTime and EndTime in whole run file')
168+
exit()
169+
170+
if ('Perturb' in wholeDict):
171+
perturb = wholeDict['Perturb']
172+
else:
173+
perturb = {'f107': {'Mean' : 1.0,
174+
'Std' : 0.0,
175+
'Add': False,
176+
'Constant': True}}
177+
wholeDict['Perturb'] = perturb
178+
179+
if ('Ensembles' in wholeDict):
180+
wholeDict['Ensembles']['nMembers'] = nMembers
181+
else:
182+
wholeDict['Ensembles']['nMembers'] = nMembers
183+
184+
interval = dTime / (nRestarts + 1)
185+
186+
if ('Restart' in wholeDict):
187+
wholeDict['Restart']['do'] = False
188+
wholeDict['Restart']['dt'] = interval
189+
else:
190+
wholeDict['Restart'] = {'do' : False,
191+
'dt' : interval}
192+
193+
subDict = {'Ensembles' : {'nMembers': nMembers},
194+
'Restart': {'do' : True, 'dt': interval},
195+
'Perturb': perturb}
196+
197+
print(' --> StartTime set to : ', sTime)
198+
199+
for iRun in range(nRestarts + 2):
200+
201+
isTest = args.test
202+
if (iRun == 0):
203+
if (not args.dowhole):
204+
isTest = True
205+
206+
cRun = '%04d' % iRun
207+
print('Iteration: ', iRun)
208+
if (iRun > 0):
209+
endtime = sTime + dt.timedelta(seconds = interval * iRun)
210+
else:
211+
endtime = sTime + dt.timedelta(seconds = dTime)
212+
endtimeAsArray = [endtime.year,
213+
endtime.month,
214+
endtime.day,
215+
endtime.hour,
216+
endtime.minute,
217+
endtime.second]
218+
print(' --> EndTime set to : ', endtime)
219+
220+
if (args.test):
221+
outFile = dirPre + '/aether_' + cRun + '.json'
222+
else:
223+
outFile = dirPre + '/aether.json'
224+
225+
if (iRun <= 1):
226+
wholeDict['EndTime'] = endtimeAsArray
227+
jsonObject = json.dumps(wholeDict)
228+
else:
229+
subDict['EndTime'] = endtimeAsArray
230+
jsonObject = json.dumps(subDict)
231+
232+
print(' --> Writing File : ', outFile)
233+
with open(outFile, "w") as outfile:
234+
outfile.write(jsonObject)
235+
236+
# ------------------------------------------------------------------
237+
# run Aether
238+
command = runCommand + ' 2>&1 | tee output_' + cRun + '.log'
239+
run(command, isTest, isVerbose)
240+
241+
# ----------------------
242+
# remove target directories
243+
command = \
244+
'cd ' + dirPre + '/UA ; ' + \
245+
'rm -rf output_' + cRun + \
246+
' ; cd ..' + cdPost
247+
run(command, isTest, isVerbose)
248+
249+
command = \
250+
'cd ' + dirPre + '/UA ; ' + \
251+
'rm -rf restartOut_' + cRun + \
252+
' ; cd ..' + cdPost
253+
run(command, isTest, isVerbose)
254+
255+
# ----------------------
256+
# Post Processing
257+
command = \
258+
'cd ' + dirPre + '/UA/output ; ' + \
259+
postCommand + \
260+
' ; cd .. ; ' + \
261+
'mv output output_' + cRun + \
262+
' ; mkdir output ' + \
263+
' ; cd ..' + cdPost
264+
run(command, isTest, isVerbose)
265+
266+
# ----------------------
267+
# Set up restart directories
268+
command = \
269+
'cd ' + dirPre + '/UA ; ' + \
270+
'mv restartOut restartOut_' + cRun + \
271+
' ; rm -f ./restartIn ' + \
272+
' ; ln -s restartOut_' + cRun + ' ./restartIn ' + \
273+
' ; mkdir restartOut ' + \
274+
' ; cd ..' + cdPost
275+
run(command, isTest, isVerbose)
276+

0 commit comments

Comments
 (0)