forked from ibackus/ICgen
/
ICgen_utils.py
511 lines (346 loc) · 13.5 KB
/
ICgen_utils.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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
# -*- coding: utf-8 -*-
"""
Created on Tue Aug 5 17:37:03 2014
@author: ibackus
"""
__version__ = "$Revision: 1 $"
# $Source$
# External modules
import subprocess
import numpy as np
import pynbody
SimArray = pynbody.array.SimArray
import os
import re
# ICgen modules
from ICglobal_settings import global_settings
import isaac
def Qeff(ICobj, bins=None):
if bins is None:
bins = ICobj.sigma.r_bins
# Constants
G = SimArray([1.0],'G')
kB = SimArray([1.0], 'k')
if not hasattr(ICobj, 'snapshot'):
raise ValueError('Could not find snapshot. Must generate ICs first')
snap = ICobj.snapshot
T = snap.g['temp']
r = snap.g['rxy']
M = snap.s['mass']
m = ICobj.settings.physical.m
cs = np.sqrt(kB*T/m)
omega = snap.g['vt']/r
sigma = ICobj.sigma(r)
r_edges, h_binned = isaac.height(snap, bins=bins)
r_cent = (r_edges[1:] + r_edges[0:-1])/2
h_spl = isaac.extrap1d(r_cent, h_binned)
h = SimArray(h_spl(r), h_binned.units)
Q = (cs*omega/(np.pi*G*sigma)).in_units('1')
Q1 = Q * ((h/r).in_units('1'))**0.192
dummy, Q1_binned, dummy2 = isaac.binned_mean(r, Q1, bins=r_edges)
return r_edges, Q1_binned
class larray(list):
"""
A simple subclass of list for containing listified arrays
(see ICgen_utils.listify)
Should be instantiated by ICgen_utils.listify
USAGE:
Creating an larray object:
a = larray() # blank larray
a = larray(shape) # sets a.shape = shape
Then one can append to an larray just as with a list
a.append(stuff)
To return return a normal array:
array = a.delistify()
"""
def __init__(self, shape = None):
list.__init__(self)
self.shape = shape
def delistify(self):
"""
Return an array made from self. See ICgen_utils.delistify
"""
return delistify(self)
def listify(array, max_element=10**7):
"""
Breaks up an array or SimArray into chunks and saves as an larray object
(essentially, a list). Useful for pickling very large arrays which otherwise
may throw an error with pickle
Whenever possible the array is NOT copied, rather a view is returned. This
depends on the functionality of array.ravel() (see numpy.ravel)
**ARGUMENTS**
array : array_like or SimArray
Input array to listify
max_element : int
Maximimum number of elements per chunk (i.e., per list item)
**RETURNS**
list_array : larray
A listified array.
See Also : ICgen_utils.delistify
"""
# Initialize
shape = array.shape
n_elements = np.prod(shape)
if n_elements <= max_element:
return array
else:
out_list = larray(shape)
array = array.ravel()
# Number per slice
N = int(max_element)
counter = 0
i = 0
while counter < n_elements:
out_list.append(array[counter:counter+N])
i += 1
counter += N
return out_list
def delistify(in_list):
"""
Reverses listify.
Takes an larray object (a listified array) and returns the original array
**ARGUMENTS**
in_list : larray
Listified array
**RETURNS**
array : array_like or SimArray
array from the original data
See Also : ICgen_utils.listify
"""
if isinstance(in_list, larray):
if isinstance(in_list[0], SimArray):
array = SimArray(np.concatenate(in_list), in_list[0].units)
else:
array = np.concatenate(in_list)
return array.reshape(in_list.shape)
else:
return in_list
def est_eps(smoothlength_file, nstar=1):
"""
Estimate gravitational softening length (eps) from a ChaNGa output .smoothlength
file. eps is estimated as 1/2 the mean smoothing length
**ARGUENTS**
smoothlength_file : str
Filename of the .smoothlength file
nstar : int
Number of star particles present
**RETURNS**
eps : float
Estimate of the gravitational softening length in simulation units
"""
# Open ChaNGa output file containing smoothing lengths for all particles
f = open(smoothlength_file, 'r')
# Total number of particles (incl. star)
nParticles = int(f.readline().strip())
# Allocate smoothing length array
smoothlength = np.zeros(nParticles, dtype=np.float32)
# Read smoothing lengths
for i, line in enumerate(f):
smoothlength[i] = float(line.strip())
# Calculate eps, ignoring star particle
mean_smooth = (smoothlength.sum() - smoothlength[-nstar])/(nParticles-nstar)
eps = mean_smooth/2
return eps
def est_time_step(param_name, preset='default', dDelta0=100, changa_args='', runner_args=''):
"""
A routine to automatically estimate a reasonable time-step size for ChaNGa.
The idea is to have about half the particles fall into the lowest rung (ie
the big time step). This is done by calculating the rung distribution for
a large time step by running ChaNGa and killing ChaNGa once it has output
rung distribution.
NOTE: this is still fairly alpha. A better version would probably not
just run ChaNGa and then try to kill it. To be safe, a local ChaNGa preset
should probably be used.
**ARGUMENTS**
param_name : str
Filename for a ChaNGa .param file which defines parameters for the
snapshot. The snapshot must already be saved to disk
preset : str
changa_runner preset to use. See ICglobal_settings.global_settings
dDelta0 : int or float
Some large time step that should place all the particles at higher
rungs.
changa_args : str
Additional command line arguments to pass to changa. CANNOT include
-n (number of time steps) or -dt (timestep size)
runner_args : str
Additional command line arguments to pass to the runner, ie to
charmrun or mpirun
**RETURNS**
dDelta : float
Estimated reasonable time step that places half the particles on the
lowest rung (ie the big time step)
"""
settings = global_settings['changa_presets'][preset]
changa_name = settings[2]
runner_name = settings[0]
changa_args += ' -n 1 -dt {0}'.format(dDelta0)
command = changa_command(param_name, preset, changa_args=changa_args, runner_args=runner_args)
rung_line = ''
p = changa_run(command, verbose=False)
for line in iter(p.stdout.readline, ''):
if 'rung distribution' in line.lower():
# Kill the runner
kill_command = 'pkill -9 ' + runner_name
pkill = subprocess.Popen(kill_command.split(), \
stdout=subprocess.PIPE)
pkill.wait()
# Kill ChaNGa
kill_command = 'pkill -9 ' + changa_name
pkill = subprocess.Popen(kill_command.split(), \
stdout=subprocess.PIPE)
pkill.wait()
rung_line = line.strip()
break
if rung_line == '':
raise RuntimeError('ChaNGa failed to output rung distribution')
rung_list = re.findall('\d+', rung_line)
rung_hist = np.array(rung_list).astype(float)
rung_edges = np.arange(len(rung_hist) + 1, dtype=float)
s = np.cumsum(rung_hist)
Ntot = s[-1]
# Find first bin which gives us more than half the total number
for i, n in enumerate(s):
if n >= 0.5*Ntot:
ind = i
break
# Calculate the median rung
rung_med = rung_edges[ind] + (0.5*Ntot - s[ind-1])/rung_hist[ind]
# Now estimate a time step that will fit about half the particles on the
# lowest rung (ie the big time step)
dDelta = dDelta0 * 2.0**(-rung_med+1)
return dDelta
def changa_run(command, verbose = True, logfile_name=None, force_wait=False):
"""
A wrapper for running ChaNGa
**ARGUMENTS**
command : str
A full command line command for running ChaNGa. Can be produced from
defaults using ICgen_utils.changa_command
verbose : bool
(optional) Flag for printing ChaNGa output to stdout.
If True - stdout is printed. This will effectively makes changa_run
wait on ChaNGa completion
logfile_name : str
(optional) If set, saves ChaNGa output to file
force_wait : bool
(optional) Default = False
If set, forces wait on ChaNGa before completion
**RETURNS**
p : subprocess.Popen
A process object created by subprocess.Popen for the ChaNGa command
"""
if logfile_name is not None:
logfile = open(logfile_name, 'w')
logfile.close()
logfile = open(logfile_name, 'a')
if verbose:
output = subprocess.PIPE
p = subprocess.Popen(command.split(), stderr=output, stdout=output)
for line in iter(p.stdout.readline, ''):
print line,
if logfile_name is not None:
logfile.write(line)
p.wait()
else:
if logfile_name is not None:
output = logfile
else:
output = subprocess.PIPE
p = subprocess.Popen(command.split(), stderr=output, stdout=output)
if force_wait:
p.wait()
return p
def changa_command(param_name, preset=None, changa_bin=None, changa_args='', runner_args=''):
"""
A utility for created command line commands for running ChaNGa
**ARGUMENTS**
param_name : str
Filename of the .param file used for ChaNGa
preset : str
if None, the default preset is used
Presets are defined in global_settings
changa_bin : str
Path to the ChaNGa binary to use. If None, defaults are used
Overrides preset binary
changa_args : str
Additional user supplied arguments for ChaNGa
runner_args : str
Additional user supplied arguments for the runner (ie charmrun or mpirun)
**RETURNS**
command : str
A command line command for running ChaNGa
"""
# Contains all the presets
preset_dict = global_settings['changa_presets']
# Load the preset
if preset is None:
preset = preset_dict['default']
preset_list = preset_dict[preset]
# Get full path to ChaNGa binary
if changa_bin is None:
changa_bin = preset_list[2]
changa_bin = os.popen('which ' + changa_bin).read().strip()
if '' == changa_bin:
raise RuntimeError, 'Could not find ChaNGa. Try different preset'
# Merge user defined extra arguments
runner_args = ' '.join([preset_list[1], runner_args])
changa_args = ' '.join([preset_list[3], changa_args])
runner = preset_list[0]
command = ' '.join([runner, runner_args, changa_bin, changa_args, param_name])
command = ' '.join(command.split())
return command
def arg_cat(arg_list):
"""
STILL ALPHA!!!
arg_str = arg_cat([args1, args2, ...])
Concatenates a list of various command line arguments. arg_list should
be a list containing command line argument strings.
Priority is given to later arguments. So arg_list[2] overwrites arg_list[1]
if they share any flags
**EXAMPLES**
args1 = '-n 20 +cd 13 ++b fire testit'
args2 = '-n 20 +cd 300'
print arg_cat([args1, args2])
returns:
+cd 300 -n 20 testit ++b fire
"""
args_dict = {}
# Loop through all sets of arguments
for args in arg_list:
# Split args if it's not already a list
if isinstance((args), str):
args = args.split()
# Parse arguments
counter = 0
while counter < len(args):
key = args[counter]
if (key[0] == '-') or (key[0] == '+'):
# We have a flag followed by its value
val = args[counter+1]
counter += 2
else:
# We just have an argument
val = ''
counter += 1
args_dict[key] = val
args_str = ''
for key, val in args_dict.iteritems():
if val == '':
# Tack on to the end
args_str = ' '.join([args_str, key])
else:
# Place at beginning
args_str = ' '.join([key, val, args_str])
args_str = ' '.join(args_str.split())
return args_str
def checkversion(obj):
"""
Checks the version of an object. Returns -1 if there is no version
"""
if hasattr(obj, '__version__'):
return obj.__version__
elif ('__version__' in obj) and isinstance(obj, dict):
return obj['__version__']
else:
return -1