Unzip speed up #20
Description
Running line_profiler and profiling every function in evolve.py has yeilded the following:
Timer unit: 1e-06 s
Total time: 0.059587 s
File: evolve.py
Function: load_sfh at line 96
Line # Hits Time Per Hit % Time Line Contents
==============================================================
96 @profile
97 def load_sfh(self):
98 '''
99 takes in input SFH file and extend backwards to start from 1e-3 Gyr
100 '''
101 2 6 3.0 0.0 try:
102 2 48635 24317.5 81.6 vals = np.loadtxt(self.SFH_file)
103 2 6 3.0 0.0 scale = [1e-9,1e9] # Gyr conversions for time, SFR
104 2 74 37.0 0.1 sfh = vals*scale # converts time in Gyr and SFR in Msun/Gyr
105 # extrapolates SFH back to 0.001Gyr using SFH file and power law (gamma)
106 2 9273 4636.5 15.6 final_sfh = f.extra_sfh(sfh, self.gamma)
107 2 1593 796.5 2.7 self.sfh = np.array(final_sfh)
108 except:
109 logger.error("File '%s' will not parse" % self.SFH_file)
110 self.sfh = None
Total time: 0.559438 s
File: evolve.py
Function: sfr at line 112
Line # Hits Time Per Hit % Time Line Contents
==============================================================
112 @profile
113 def sfr(self, t):
114 '''
115 define sfr as function to look up nearest sfr value at any specified time
116 '''
117 42898 20311 0.5 3.6 try:
118 42898 509412 11.9 91.1 vals = find_nearest(self.sfh,t)
119 42898 29715 0.7 5.3 return vals[1]
120 except:
121 logger.error("No SFH yet")
Total time: 95.0901 s
File: evolve.py
Function: gas_metal_dust_mass at line 123
Line # Hits Time Per Hit % Time Line Contents
==============================================================
123 @profile
124 def gas_metal_dust_mass(self, sn_rate):
125 '''
126 Calculates the gas, metal and dust mass from stars
127 note mass is only ejected after stars die ie when
128 t - taum (where taum is lifetime of star) > 0
129 '''
130 # initialize
131 2 2 1.0 0.0 mg = self.gasmass_init
132 2 2 1.0 0.0 mstars = 0
133 2 2 1.0 0.0 md = 0
134 2 2 1.0 0.0 md_all = 0
135 2 2 1.0 0.0 md_stars = 0
136 2 2 1.0 0.0 md_gg = 0
137 2 2 1.0 0.0 metals = 0
138 2 1 0.5 0.0 prev_t = 1e-3
139 2 1 0.5 0.0 metals_pre = 0
140 2 2 1.0 0.0 mstars_list = []
141 2 2 1.0 0.0 z = []
142 2 2 1.0 0.0 z_lookup = []
143 2 2 1.0 0.0 sfr_list = []
144 2 2 1.0 0.0 sfr_lookup = []
145 2 2 1.0 0.0 all_results = []
146 # Limit time to less than tend
147 2 4 2.0 0.0 time = self.sfh[:,0]
148 2 33 16.5 0.0 time = time[time < self.tend]
149 2 21 10.5 0.0 now = datetime.now()
150 # TIME integral
151 3263 6748 2.1 0.0 for item, t in enumerate(time):
152 3262 5910 1.8 0.0 r_sn = sn_rate [item]
153 3262 4627 1.4 0.0 metallicity = metals/mg
154
155 # start appending arrays for needing later
156 3262 6703 2.1 0.0 z.append([t,metallicity])
157 3262 1423655 436.4 1.5 z_lookup = np.array(z)
158 3262 90205 27.7 0.1 sfr_list.append([t,self.sfr(t)])
159 3262 1225041 375.5 1.3 sfr_lookup = np.array(sfr_list)
160
161 '''
162 STARS: dM_stars = sfr(t) * dt
163 '''
164 3262 77753 23.8 0.1 dmstars = self.sfr(t)
165
166 '''
167 GAS: dMg = (-sfr(t) + e(t) + inflows(t) - outflows(t)) * dt
168 set up astration, inflow, outflow components
169 '''
170 3262 56955 17.5 0.1 gas_ast = self.sfr(t)
171 3262 62668 19.2 0.1 gas_inf = f.inflows(self.sfr(t), self.inflows['xSFR'])
172 3262 61147 18.7 0.1 gas_out = f.outflows(self.sfr(t), self.outflows['xSFR'])
173
174 '''
175 METALS: dMz = (-Z*sfr(t) + ez(t) + Z*inflows(t) - Z*outflows(t)) * dt
176 set up astration, inflow and outflow components
177 '''
178 3262 59347 18.2 0.1 metals_ast = f.astration(metals,mg,self.sfr(t))
179 3262 4713 1.4 0.0 if self.outflows['metals']:
180 1864 33706 18.1 0.0 metals_out = metallicity*f.outflows(self.sfr(t), self.outflows['xSFR'])
181 else:
182 1398 1638 1.2 0.0 metals_out = 0.
183 3262 59859 18.4 0.1 metals_inf = self.inflows['metals']*f.inflows(self.sfr(t), self.inflows['xSFR'])
184
185 '''
186 DUST: dMd = (-Md/Mg*sfr(t) + ed(t) + Md/Mg*inflows(t) - Md/Mg*outflows(t)
187 - (1-f)*Md/t_destroy + f(1-Md/Mg)*Md/t_graingrowth) * dt
188 set up astration, inflows, outflows, destruction, grain growth components
189 '''
190 3262 4479 1.4 0.0 if self.outflows['dust']:
191 1864 33644 18.0 0.0 mdust_out = (md/mg)*f.outflows(self.sfr(t), self.outflows['xSFR'])
192 else:
193 1398 1647 1.2 0.0 mdust_out = 0.
194 3262 59203 18.1 0.1 mdust_inf = self.inflows['dust']*f.inflows(self.sfr(t), self.inflows['xSFR'])
195 3262 55883 17.1 0.1 mdust_ast = f.astration(md,mg,self.sfr(t))
196
197 3262 81253 24.9 0.1 mdust_gg, t_gg = f.graingrowth(self.choice_dust[2], self.epsilon,mg, self.sfr(t), metallicity, md, self.coldfraction)
198 3262 22796 7.0 0.0 mdust_des, t_des = f.destroy_dust(self.choice_des, self.destroy_ism, mg, r_sn, md, self.coldfraction)
199
200 '''
201 Get ejected masses from stars when they die
202 gas_ej = e(t): ejected gas mass from stars of mass m at t = taum
203 metals_stars = ez(t): ejected metal mass from stars of mass m at t = taum (fresh + recycled)
204 mdust_stars = ed(t): ejected dust mass from stars of mass m at t = taum (fresh + recycled)
205 '''
206 gas_ej, metals_stars, mdust_stars = \
207 3262 88276663 27062.1 92.8 f.mass_integral(self.choice_dust, self.reduce_sn, t, metallicity, sfr_lookup, z_lookup, self.imf)
208
209 '''
210 integrate over time for gas, metals and stars (mg, metals, md)
211 '''
212 3262 12345 3.8 0.0 dmg = -gas_ast + gas_ej + gas_inf - gas_out
213 3262 7334 2.2 0.0 dmetals = -metals_ast + metals_stars + metals_pre + metals_inf - metals_out
214 3262 7084 2.2 0.0 ddust = -mdust_ast + mdust_stars + mdust_inf - mdust_out + mdust_gg - mdust_des
215 # dust_source_all separates out the dust sources (Md vs t) wihtout including sinks (Astration etc)
216 # and grain growth separately (this is the Md vs time contributed by dust sources)
217 3262 4919 1.5 0.0 dust_source_all = mdust_stars + mdust_gg
218 3262 4751 1.5 0.0 dt = t - prev_t # calculate next time step
219 3262 4068 1.2 0.0 prev_t = t
220 3262 5564 1.7 0.0 mstars += dmstars*dt
221 3262 4789 1.5 0.0 mg += dmg*dt # gas mass integral
222 3262 5506 1.7 0.0 if mg <= 0:
223 # exit program if all ISM removed
224 1 18 18.0 0.0 print ('Oops you have no interstellar medium left')
225 1 3 3.0 0.0 break
226 3261 4682 1.4 0.0 metals += dmetals*dt # metal mass integral
227 3261 4636 1.4 0.0 md += ddust*dt # dust mass integral
228 3261 4712 1.4 0.0 md_all += dust_source_all*dt # dust mass sources integral
229 3261 4751 1.5 0.0 md_gg += mdust_gg*dt # dust source from grain growth only
230 3261 4771 1.5 0.0 md_stars += mdust_stars*dt # dust source from stars only
231 3261 1657357 508.2 1.7 Z = zip(*z_lookup) # write metallicity to an array
232 3261 1487612 456.2 1.6 s_f_r = zip(*sfr_lookup) # write SFR lookup array
233 3261 10769 3.3 0.0 if mg <= 0. or metals <=0: # write dust/metals ratio
234 1353 1740 1.3 0.0 dust_to_metals = 0.
235 else:
236 1908 3620 1.9 0.0 dust_to_metals = md/metals
237 3261 5994 1.8 0.0 all_results.append((t, mg, mstars, metals, metallicity, \
238 3261 114921 35.2 0.1 md, dust_to_metals, self.sfr(t)*1e-9, \
239 3261 6936 2.1 0.0 md_all, md_stars, md_gg, t_des, t_gg))
240 # to test code kinks
241 2 108 54.0 0.0 print("Gas, metal and dust mass exterior loop %s" % str(datetime.now()-now))
242 2 4818 2409.0 0.0 return np.array(all_results)
Total time: 2.26193 s
File: evolve.py
Function: supernova_rate at line 244
Line # Hits Time Per Hit % Time Line Contents
==============================================================
244 @profile
245 def supernova_rate(self):
246 '''
247 Calculates the SN rate at time t by integrating over mass m
248 '''
249 # initialize
250 2 4 2.0 0.0 sn_rate_list = []
251 2 2 1.0 0.0 dm = 0.01
252 2 2 1.0 0.0 prev_t = 1e-3
253 # define time array
254 2 8 4.0 0.0 time = self.sfh[:,0]
255 2 55 27.5 0.0 time = time[time < self.tend]
256 3291 3391 1.0 0.1 for t in time:
257 # need to clear the sn_rates as we don't want them adding up
258 3289 2447 0.7 0.1 sn_rate = 0.
259 3289 2478 0.8 0.1 dsn_rate = 0.
260 3289 2951 0.9 0.1 if t < 0.049:
261 2180 1374063 630.3 60.7 m = lookup_fn(t_lifetime,'lifetime_high_metals',t)[0]
262 else:
263 1109 805 0.7 0.0 m = 9.
264 119825 95893 0.8 4.2 while m < 40.:
265 116536 88941 0.8 3.9 if m > 10.:
266 111744 82717 0.7 3.7 dm = 0.5
267 116536 428777 3.7 19.0 sn_rate += f.initial_mass_function(m, self.imf_type)*dm
268 116536 95611 0.8 4.2 m += dm
269 3289 73646 22.4 3.3 r_sn = self.sfr(t)*sn_rate # units in N per Gyr
270 3289 3470 1.1 0.2 dt = t - prev_t
271 3289 2478 0.8 0.1 prev_t = t
272 3289 3895 1.2 0.2 sn_rate_list.append(r_sn)
273 2 300 150.0 0.0 return np.array(sn_rate_list)
It looks like the unzipping at lines 231 and 232 of gas_metal_dust_mass() are taking most of the time.
I'll have a look at speeding that up.