-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathbolshoireader.py
151 lines (141 loc) · 7.03 KB
/
bolshoireader.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
import numpy as np
import pdb
class bolshoireader:
def __init__(self,fn, minmass, maxmass, bolshoidir):
self.filename = fn
self.bolshoidir = bolshoidir
trees=[]
with open(bolshoidir+fn,'r') as f:
line = f.readline()
while line!='':
sline = line.split()
treenum = int(sline[0])
mh0 = float(sline[1])
weightedmdot = float(sline[2])
if minmass<mh0 and mh0<maxmass:
trees.append([treenum, mh0, weightedmdot])
line=f.readline()
self.trees = sorted(trees, key=lambda tree: tree[2])
self.keys = [tree[2] for tree in self.trees]
self.masses = np.log10([tree[1] for tree in self.trees]) # log10(M_h)
self.minimumMass = np.min(self.masses)
self.maximumMass = np.max(self.masses)
self.stored = False
def storeAll(self):
self.xs = np.zeros( (len(self.trees), 1000) )
with open(self.bolshoidir+'rf_data4.txt', 'r') as f:
all_lines = f.readlines()
treecounter = 0
bad_rows = []
for treei, tree in enumerate(self.trees):
treecounter+=1
if treecounter % 100 == 0:
print ("Loading up tree ", treecounter, ' of ', len(self.trees))
treestring = 'tree '+repr(tree[0])+'\n'
treeind = all_lines.index(treestring)
try:
self.xs[treei, : ] = [float( all_lines[lineind].split()[0] ) for lineind in range(treeind+1,treeind+1001) ]
except:
print ("Ran into trouble!")
bad_rows.append(treei)
for k, row in enumerate(bad_rows):
random_ind = int((len(self.trees)-1)*np.random.random())
if not random_ind in bad_rows:
self.xs[row, :] = self.xs[random_ind, :]
self.stored = True
def returnnearest(self, dirname, filename, value, logMh0):
''' copy a file with the closest value of tree[2]'''
#ind = np.searchsorted(self.keys, value)
assert value>=0 and value<=1
if logMh0 < self.minimumMass+0.25:
acceptable = self.masses<self.minimumMass+0.5
elif logMh0 > self.maximumMass-0.25:
acceptable = self.masses>self.maximumMass-0.5
else:
acceptable = np.logical_and( self.masses>logMh0-0.25, self.masses<logMh0+0.25 )
nAcceptable = np.sum(np.ones(len(self.keys))[acceptable])
indAccept = int(nAcceptable*value) # pick one of the acceptable trees
ind = np.array(range(len(self.keys)))[acceptable][indAccept]
if ind == len(self.keys):
ind=ind-1
assert ind<len(self.keys) and ind>-1
#rffile = bolshoidir+ repr(self.trees[ind][0])+'_rf.txt'
#shutil.copy(rffile, dirname)
# This chunk of code searches the file rf_data.txt for
# the appropriate tree number from the register as identified by ind above.
# Once found, the subsequent lines (containing the values of the random factors
# for the main progenitor accretion history of that tree) are copied to a file
# in the working directory of the gidget run.
# If we've already cached the x values, use that!
if self.stored:
return self.xs[ind,:]
else:
xs = []
with open(self.bolshoidir+'rf_data4.txt','r') as f:
line = f.readline()
while line!='':
if 'tree '+repr(self.trees[ind][0]) in line:
print ("Copying tree with index ",ind," and Bolshoi ID ",self.trees[ind][0])
break # we found the tree we're looking for
line = f.readline()
if line=='':
print ("Failed to find line tree "+repr(self.trees[ind][0]))
line = f.readline()
collectlines=[]
while not 'tree' in line and line!='':
xs.append( float(line.split()[0]) )
collectlines.append(line)
line = f.readline()
# with open(dirname+'/'+filename,'w') as ff:
# assert len(collectlines)>0
# ### Major change! Add in a line at the beginning which tells us the z=0 halo mass of the bolshoi tree from which these random factors are drawn
# ff.write(str(self.trees[ind][1])+'\n')
# for line in collectlines:
# ff.write(line)
return xs[:-1]
def copynearest(self, dirname, filename, value, logMh0):
''' copy a file with the closest value of tree[2]'''
#ind = np.searchsorted(self.keys, value)
assert value>=0 and value<=1
if logMh0 < self.minimumMass+0.25:
acceptable = self.masses<self.minimumMass+0.5
elif logMh0 > self.maximumMass-0.25:
acceptable = self.masses>self.maximumMass-0.5
else:
acceptable = np.logical_and( self.masses>logMh0-0.25, self.masses<logMh0+0.25 )
nAcceptable = np.sum(np.ones(len(self.keys))[acceptable])
indAccept = int(nAcceptable*value) # pick one of the acceptable trees
ind = np.array(range(len(self.keys)))[acceptable][indAccept]
if ind == len(self.keys):
ind=ind-1
assert ind<len(self.keys) and ind>-1
#rffile = bolshoidir+ repr(self.trees[ind][0])+'_rf.txt'
#shutil.copy(rffile, dirname)
# This chunk of code searches the file rf_data.txt for
# the appropriate tree number from the register as identified by ind above.
# Once found, the subsequent lines (containing the values of the random factors
# for the main progenitor accretion history of that tree) are copied to a file
# in the working directory of the gidget run.
with open(self.bolshoidir+'rf_data4.txt','r') as f:
line = f.readline()
while line!='':
if 'tree '+repr(self.trees[ind][0]) in line:
print ("Copying tree with index ",ind," and Bolshoi ID ",self.trees[ind][0])
break # we found the tree we're looking for
line = f.readline()
if line=='':
print ("Failed to find line tree "+repr(self.trees[ind][0]))
line = f.readline()
collectlines=[]
while not 'tree' in line and line!='':
collectlines.append(line)
line = f.readline()
with open(dirname+'/'+filename,'w') as ff:
assert len(collectlines)>0
### Major change! Add in a line at the beginning which tells us the z=0 halo mass of the bolshoi tree from which these random factors are drawn
ff.write(str(self.trees[ind][1])+'\n')
for line in collectlines:
ff.write(line)
def test(self):
for i in range(len(self.keys)):
globalBolshoiReader.copynearest('~/bolshoitest/', 'test_'+str(i).zfill(5)+'_inputRandomFactors.txt', i)