Skip to content

Commit 5a80172

Browse files
reinhold-willcoxreinhold-willcox
committed
added compasUtils file
1 parent 20146a4 commit 5a80172

File tree

1 file changed

+355
-0
lines changed

1 file changed

+355
-0
lines changed

‎compas_python_utils/compasUtils.py‎

Lines changed: 355 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,355 @@
1+
import h5py as h5
2+
import numpy as np
3+
import pandas as pd
4+
from numpy.dtypes import StringDType
5+
6+
7+
########################################################################
8+
# ## Function to print the data from a given COMPAS HDF5 group in a readable pandas template
9+
10+
def printCompasDetails(data, *seeds, mask=()):
11+
"""
12+
Function to print the full Compas output for given seeds, optionally with an additional mask
13+
"""
14+
list_of_keys = list(data.keys())
15+
16+
# Check if seed parameter exists - if not, just print without (e.g RunDetails)
17+
if ('SEED' in list_of_keys) | ('SEED>MT' in list_of_keys): # Most output files
18+
#SEED>MT is a relic from older versions, but we leave this in for backwards compatibility
19+
20+
# Set the seed name parameter, mask on seeds as needed, and set the index
21+
seedVariableName='SEED' if ('SEED' in list_of_keys) else 'SEED>MT'
22+
list_of_keys.remove(seedVariableName) # this is the index above, don't want to include it
23+
24+
allSeeds = data[seedVariableName][()]
25+
seedsMask = np.isin(allSeeds, seeds)
26+
if len(seeds) == 0: # if any seed is included, do not reset the mask
27+
seedsMask = np.ones_like(allSeeds).astype(bool)
28+
if mask == ():
29+
mask = np.ones_like(allSeeds).astype(bool)
30+
mask &= seedsMask
31+
32+
df = pd.DataFrame.from_dict({param: data[param][()][mask] for param in list(data.keys())}).set_index(seedVariableName).T
33+
34+
else: # No seed parameter, so do custom print for Run Details
35+
36+
# Get just the keys without the -Derivation suffix - those will be a second column
37+
keys_not_derivations = []
38+
for key in list_of_keys:
39+
if '-Derivation' not in key:
40+
keys_not_derivations.append(key)
41+
42+
# Some parameter values are string types, formatted as np.bytes_, need to convert back
43+
def convert_strings(param_array):
44+
if isinstance(param_array[0], np.bytes_):
45+
return param_array.astype(str)
46+
else:
47+
return param_array
48+
49+
df_keys = pd.DataFrame.from_dict({param: convert_strings(data[param][()]) for param in keys_not_derivations }).T
50+
nCols = df_keys.shape[1] # Required only because if we combine RDs, we get many columns (should fix later)
51+
df_keys.columns = ['Parameter']*nCols
52+
df_drvs = pd.DataFrame.from_dict({param: convert_strings(data[param+'-Derivation'][()]) for param in keys_not_derivations }).T
53+
df_drvs.columns = ['Derivation']*nCols
54+
df = pd.concat([df_keys, df_drvs], axis=1)
55+
56+
# Add units as first col
57+
units_dict = {key:data[key].attrs['units'].astype(str) for key in list_of_keys}
58+
df.insert(loc=0, column='(units)', value=pd.Series(units_dict))
59+
return df
60+
61+
62+
63+
########################################################################
64+
# ## Get event histories of MT data, SN data, and combined MT, SN data
65+
66+
def getMtEvents(MT):
67+
"""
68+
This function takes in the `BSE_RLOF` output category from COMPAS, and returns the information
69+
on the Mass Transfer (MT) events that happen for each seed. The events do not have to be in order,
70+
either chronologically or by seed, this function will reorder them as required.
71+
72+
OUT:
73+
returnedSeeds (list): ordered list of the unique seeds in the MT file
74+
returnedEvents (list): list of sublists, where each sublist contains all the MT events for a given seed.
75+
MT event tuples take the form :
76+
(stellarTypePrimary, stellarTypeSecondary, isRlof1, isRlof2, isCEE, isMrg)
77+
returnTimes (list): is a list of sublists of times of each of the MT events
78+
"""
79+
mtSeeds = MT['SEED'][()]
80+
mtTimes = MT['Time<MT'][()]
81+
mtPrimaryStype = MT['Stellar_Type(1)<MT'][()]
82+
mtSecondaryStype = MT['Stellar_Type(2)<MT'][()]
83+
mtIsRlof1 = MT['RLOF(1)>MT'][()] == 1
84+
mtIsRlof2 = MT['RLOF(2)>MT'][()] == 1
85+
mtIsCEE = MT['CEE>MT'][()] == 1
86+
mtIsMrg = MT['Merger'][()] == 1
87+
88+
# We want the return arrays sorted by seed, so sort here.
89+
mtSeedsInds = np.lexsort((mtTimes, mtSeeds)) # sort by seeds then times - lexsort sorts by the last column first...
90+
mtSeeds = mtSeeds[mtSeedsInds]
91+
mtTimes = mtTimes[mtSeedsInds]
92+
mtPrimaryStype = mtPrimaryStype[mtSeedsInds]
93+
mtSecondaryStype = mtSecondaryStype[mtSeedsInds]
94+
mtIsRlof1 = mtIsRlof1[mtSeedsInds]
95+
mtIsRlof2 = mtIsRlof2[mtSeedsInds]
96+
mtIsCEE = mtIsCEE[mtSeedsInds]
97+
mtIsMrg = mtIsMrg[mtSeedsInds]
98+
99+
# Process the MT events
100+
returnedSeeds = [] # array of seeds - will only contain seeds that have MT events
101+
returnedEvents = [] # array of MT events for each seed in returnedSeeds
102+
returnedTimes = [] # array of times for each event in returnedEvents (for each seed in returnedSeeds)
103+
104+
lastSeed = -1 # initialize most recently processed seed
105+
for seedIndex, thisSeed in enumerate(mtSeeds): # iterate over all RLOF file entries
106+
thisTime = mtTimes[seedIndex] # time for this RLOF file entry
107+
thisEvent = (mtPrimaryStype[seedIndex], mtSecondaryStype[seedIndex],
108+
mtIsRlof1[seedIndex], mtIsRlof2[seedIndex], mtIsCEE[seedIndex], mtIsMrg[seedIndex]) # construct event tuple
109+
110+
# If this is an entirely new seed:
111+
if thisSeed != lastSeed: # same seed as last seed processed?
112+
returnedSeeds.append(thisSeed) # no - new seed, record it
113+
returnedTimes.append([thisTime]) # initialize the list of event times for this seed
114+
returnedEvents.append([thisEvent]) # initialize the list of events for this seed
115+
lastSeed = thisSeed # update the latest seed
116+
117+
# Add event, if it is not a duplicate
118+
try:
119+
eventIndex = returnedEvents[-1].index(thisEvent) # find eventIndex of this particular event tuple in the array of events for this seed
120+
if thisTime > returnedTimes[-1][eventIndex]: # ^ if event is not a duplicate, this will throw a ValueError
121+
returnedTimes[-1][eventIndex] = thisTime # if event is duplicate, update time to the later of the duplicates
122+
except ValueError: # event is not a duplicate:
123+
returnedEvents[-1].append(thisEvent) # record new event tuple for this seed
124+
returnedTimes[-1].append(thisTime) # record new event time for this seed
125+
126+
return returnedSeeds, returnedEvents, returnedTimes # see above for description
127+
128+
129+
def getSnEvents(SN):
130+
"""
131+
This function takes in the `BSE_Supernovae` output category from COMPAS, and returns the information
132+
on the Supernova (SN) events that happen for each seed. The events do not have to be in order chronologically,
133+
this function will reorder them as required.
134+
135+
OUT:
136+
returnedSeeds (list): ordered list of all the unique seeds in the SN file
137+
returnedEvents (list): list of sublists, where each sublist contains all the SN events for a given seed.
138+
SN event tuples take the form :
139+
(stellarTypeProgenitor, stellarTypeRemnant, whichStarIsProgenitor, isBinaryUnbound)
140+
returnedTimes (list): is a list of sublists of times of each of the SN events
141+
"""
142+
snSeeds = SN['SEED'][()]
143+
snTimes = SN['Time'][()]
144+
snProgStype = SN['Stellar_Type_Prev(SN)'][()]
145+
snRemnStype = SN['Stellar_Type(SN)'][()]
146+
snWhichProg = SN['Supernova_State'][()]
147+
snIsUnbound = SN['Unbound'][()] == 1
148+
149+
# We want the return arrays sorted by seed, so sort here.
150+
snSeedsInds = np.lexsort((snTimes, snSeeds)) # sort by seeds then times - lexsort sorts by the last column first...
151+
snSeeds = snSeeds[snSeedsInds]
152+
snTimes = snTimes[snSeedsInds]
153+
snProgStype = snProgStype[snSeedsInds]
154+
snRemnStype = snRemnStype[snSeedsInds]
155+
snWhichProg = snWhichProg[snSeedsInds]
156+
snIsUnbound = snIsUnbound[snSeedsInds]
157+
158+
# Process the SN events
159+
returnedSeeds = [] # array of seeds - will only contain seeds that have SN events
160+
returnedEvents = [] # array of SN events for each seed in returnedSeeds
161+
returnedTimes = [] # array of times for each event in returnedEvents (for each seed in returnedSeeds)
162+
163+
lastSeed = -1 # initialize most recently processed seed
164+
for seedIndex, thisSeed in enumerate(snSeeds): # iterate over all SN file entries
165+
thisTime = snTimes[seedIndex] # time for this SN file entry
166+
thisEvent = (snProgStype[seedIndex], snRemnStype[seedIndex],
167+
snWhichProg[seedIndex], snIsUnbound[seedIndex]) # construct event tuple
168+
169+
# If this is an entirely new seed:
170+
if thisSeed != lastSeed: # same seed as last seed processed?
171+
returnedSeeds.append(thisSeed) # no - new seed, record it
172+
returnedTimes.append([thisTime]) # initialize the list of event times for this seed
173+
returnedEvents.append([thisEvent]) # initialize the list of events for this seed
174+
lastSeed = thisSeed # update the latest seed
175+
else: # yes - second SN event for this seed
176+
returnedTimes[-1].append(thisTime) # append time at end of array
177+
returnedEvents[-1].append(thisEvent) # append event at end of array
178+
179+
return returnedSeeds, returnedEvents, returnedTimes # see above for description
180+
181+
182+
def getEventHistory(h5file, exclude_null=False):
183+
"""
184+
Get the event history for all seeds, including both RLOF and SN events, in chronological order.
185+
IN:
186+
h5file (h5.File() type): COMPAS HDF5 output file
187+
exclude_null (bool): whether or not to exclude seeds which undergo no RLOF or SN events
188+
OUT:
189+
returnedSeeds (list): ordered list of all seeds in the output
190+
returnedEvents (list): a list (corresponding to the ordered seeds above) of the
191+
collected SN and MT events from the getMtEvents and getSnEvents functions above,
192+
themselves ordered chronologically
193+
"""
194+
SP = h5file['BSE_System_Parameters']
195+
MT = h5file['BSE_RLOF']
196+
SN = h5file['BSE_Supernovae']
197+
allSeeds = SP['SEED'][()] # get all seeds
198+
mtSeeds, mtEvents, mtTimes = getMtEvents(MT) # get MT events
199+
snSeeds, snEvents, snTimes = getSnEvents(SN) # get SN events
200+
201+
numMtSeeds = len(mtSeeds) # number of MT events
202+
numSnSeeds = len(snSeeds) # number of SN events
203+
204+
if numMtSeeds < 1 and numSnSeeds < 1: return [] # no events - return empty history
205+
206+
returnedSeeds = [] # array of seeds - will only contain seeds that have events (of any type)
207+
returnedEvents = [] # array of events - same size as returnedSeeds (includes event times)
208+
209+
eventOrdering = ['RL', 'CE', 'SN', 'MG'] # order of preference for simultaneous events
210+
211+
mtIndex = 0 # index into MT events arrays
212+
snIndex = 0 # index into SN events arrays
213+
214+
if exclude_null:
215+
seedsToIterate = np.sort(np.unique(np.append(mtSeeds, snSeeds))) # iterate over all the seeds that have either MT or SN events
216+
else:
217+
seedsToIterate = allSeeds
218+
219+
idxOrdered = np.argsort(seedsToIterate)
220+
returnedSeeds = [None] * np.size(seedsToIterate) # array of seeds - will only contain seeds that have events (of any type)
221+
returnedEvents = [None] * np.size(seedsToIterate) # array of events - same size as returnedSeeds (includes event times)
222+
223+
for idx in idxOrdered:
224+
seed = seedsToIterate[idx]
225+
seedEvents = [] # initialise the events for the seed being processed
226+
227+
# Collect any MT events for this seed, add the time of the event and the event type
228+
while mtIndex < numMtSeeds and mtSeeds[mtIndex] == seed:
229+
for eventIndex, event in enumerate(mtEvents[mtIndex]):
230+
_, _, isRL1, isRL2, isCEE, isMrg = event
231+
eventKey = 'MG' if isMrg else 'CE' if isCEE else 'RL' # event type: Mrg, CEE, RLOF 1->2, RLOF 2->1
232+
seedEvents.append((eventKey, mtTimes[mtIndex][eventIndex], *mtEvents[mtIndex][eventIndex]))
233+
mtIndex += 1
234+
235+
# Collect any SN events for this seed, add the time of the event and the event type
236+
while snIndex < numSnSeeds and snSeeds[snIndex] == seed:
237+
for eventIndex, event in enumerate(snEvents[snIndex]):
238+
eventKey = 'SN'
239+
seedEvents.append((eventKey, snTimes[snIndex][eventIndex], *snEvents[snIndex][eventIndex]))
240+
snIndex += 1
241+
242+
seedEvents.sort(key=lambda ev:(ev[1], eventOrdering.index(ev[0]))) # sort the events by time and event type (MT before SN if at the same time)
243+
244+
returnedSeeds[idx] = seed # record the seed in the seeds array being returned
245+
returnedEvents[idx] = seedEvents # record the events for this seed in the events array being returned
246+
247+
return returnedSeeds, returnedEvents # see above for details
248+
249+
250+
251+
###########################################
252+
# ## Produce strings of the event histories
253+
254+
stellarTypeDict = {
255+
0: 'MS',
256+
1: 'MS',
257+
2: 'HG',
258+
3: 'GB',
259+
4: 'GB',
260+
5: 'GB',
261+
6: 'GB',
262+
7: 'HE',
263+
8: 'HE',
264+
9: 'HE',
265+
10: 'WD',
266+
11: 'WD',
267+
12: 'WD',
268+
13: 'NS',
269+
14: 'BH',
270+
15: 'MR',
271+
16: 'MS',
272+
}
273+
274+
def buildEventString(events, useIntStypes=False):
275+
"""
276+
Function to produce a string representing the event history of a single binary for quick readability.
277+
NOTE: unvectorized, do the vectorization later for a speed up
278+
IN:
279+
events (list of tuples): events output from getEventHistory()
280+
OUT:
281+
eventString (string): string representing the event history of the binary
282+
283+
MT strings look like:
284+
P>S, P<S, or P=S where P is primary type, S is secondary type,
285+
and >, < is RLOF (1->2 or 1<-2) or otherwise = for CEE or & for Merger
286+
287+
SN strings look like:
288+
P*SR for star1 the SN progenitor,or
289+
R*SP for star2 the SN progenitor,
290+
where P is progenitor type, R is remnant type,
291+
S is state (I for intact, U for unbound)
292+
293+
Event strings for the same seed are separated by the undesrcore character ('_')
294+
"""
295+
def _remap_stype(int_stype):
296+
if useIntStypes:
297+
return str(int_stype)
298+
else:
299+
return stellarTypeDict[int_stype]
300+
301+
# Empty event
302+
if len(events) == 0:
303+
return 'NA'
304+
305+
eventStr = ''
306+
for event in events:
307+
if event[0] == 'SN': # SN event
308+
_, time, stypeP, stypeC, whichSN, isUnbound = event
309+
if whichSN == 1: # Progenitor or Remnant depending upon which star is the SN
310+
charL = _remap_stype(stypeP)
311+
charR = _remap_stype(stypeC)
312+
else:
313+
charL = _remap_stype(stypeC)
314+
charR = _remap_stype(stypeP)
315+
charM = '*u' if isUnbound else '*i' # unbound or intact
316+
else: # Any of the MT events
317+
_, time, stype1, stype2, isRL1, isRL2, isCEE, isMrg = event
318+
charL = _remap_stype(stype1) # primary stellar type
319+
charR = _remap_stype(stype2) # secondary stellar type
320+
charM = '&' if isMrg \
321+
else '=' if isCEE \
322+
else '<' if isRL2 \
323+
else '>' # event type: CEE, RLOF 2->1, RLOF 1->2
324+
eventStr += "{}{}{}_".format(charL, charM, charR) # event string for this star, _ is event separator
325+
326+
return np.array(eventStr[:-1], dtype=np.str_) # return event string for this star (pop the last underscore first)
327+
328+
def getEventStrings(h5file=None, allEvents=None, useIntStypes=False):
329+
"""
330+
Function to calculate the event history strings for either the entire Compas output, or some list of events
331+
IN: One of
332+
h5file (h5.File() type): COMPAS HDF5 output file
333+
allEvents (list of tuples)
334+
OUT:
335+
eventStrings (list): list of strings of the event history of each seed
336+
"""
337+
338+
# If output is
339+
if (h5file == None) & (allEvents == None):
340+
return
341+
elif (allEvents == None):
342+
_, allEvents = getEventHistory(h5file)
343+
344+
eventStrings = np.zeros(len(allEvents), dtype=StringDType())
345+
for ii, eventsForGivenSeed in enumerate(allEvents):
346+
eventString = buildEventString(eventsForGivenSeed, useIntStypes=useIntStypes)
347+
eventStrings[ii] = eventString # append event string for this star (pop the last underscore first)
348+
349+
return eventStrings
350+
351+
""
352+
353+
354+
""
355+

0 commit comments

Comments
 (0)