Skip to content

Commit

Permalink
Vectorization of passband computations (#972)
Browse files Browse the repository at this point in the history 
* Optimizes blackbody extinction computation.
It still runs the loops, but it computes auxiliary quantities outside of the loop.

* Vectorizes blackbody extinction computation.

* Full vectorization of model atmosphere extinction.
This brings down per-passband computation from ~1 day to ~5 minutes.

* Removes obsolete, commented-out piece of code.
  • Loading branch information
@aprsa
aprsa authored Oct 25, 2024
1 parent 2092b6e commit 6decdf5
Showing 1 changed file with 28 additions and 44 deletions.
72 changes: 28 additions & 44 deletions phoebe/atmospheres/passbands.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -746,42 +746,17 @@ def compute_blackbody_intensities(self, teffs=None, include_extinction=False, rv
axbx = libphoebe.gordon_extinction(self.wl)
ax, bx = axbx[:,0], axbx[:,1]

pgrid = np.empty(shape=(len(teffs), len(ebvs), len(rvs), 1))
egrid = np.empty(shape=(len(teffs), len(ebvs), len(rvs), 1))
# The following code broadcasts arrays so that integration can be vectorized:
bb_sed = self._planck(self.wl[:, None], teffs[None, :]) # (54, 97)
ptf = self.ptf(self.wl)[:, None] # (54, 1)
Alam = 10**(-0.4 * ebvs[None, :, None] * (rvs[None, None, :] * ax[:, None, None] + bx[:, None, None])) # Shape (54, 30, 16)

for ti, teff in enumerate(teffs):
bb_sed = self._planck(self.wl, teff)
for ei, ebv in enumerate(ebvs):
for ri, rv in enumerate(rvs):
Alam = 10**(-0.4 * ebv * (rv * ax + bx))
egrid[ti, ei, ri, 0] = np.trapz(self.ptf(self.wl) * bb_sed * Alam, axis=0) / np.trapz(self.ptf(self.wl) * bb_sed, axis=0)
pgrid[ti, ei, ri, 0] = np.trapz(self.wl * self.ptf(self.wl) * bb_sed * Alam, axis=0) / np.trapz(self.wl * self.ptf(self.wl) * bb_sed, axis=0)
egrid = np.trapz(ptf[:, :, None, None] * bb_sed[:, :, None, None] * Alam[:, None, :, :], self.wl, axis=0) / np.trapz(ptf[:, :, None, None] * bb_sed[:, :, None, None], self.wl, axis=0)
pgrid = np.trapz(self.wl[:, None, None, None] * ptf[:, :, None, None] * bb_sed[:, :, None, None] * Alam[:, None, :, :], self.wl, axis=0) / np.trapz(self.wl[:, None, None, None] * ptf[:, :, None, None] * bb_sed[:, :, None, None], self.wl, axis=0)

self.ndp['blackbody'] = ndpolator.Ndpolator(basic_axes=(axes[0],))
self.ndp['blackbody'].register('ext@photon', associated_axes=(axes[1], axes[2]), grid=pgrid)
self.ndp['blackbody'].register('ext@energy', associated_axes=(axes[1], axes[2]), grid=egrid)

# ebv_column = np.tile(ebvs, len(rvs))
# rvebv_column = ebv_column*np.repeat(rvs, len(ebvs))
# ext_pars = np.vstack((rvebv_column, ebv_column))
# ext_func = libphoebe.gordon_extinction(wls) # (47, 2)
# ext = ext_func @ ext_pars # (47, 480)
# flux_fracs = 10**(-0.4*ext) # (47, 480)
# integrand_energy = (pbpfs_energy[:, None, :]*flux_fracs[:, :, None]).T # ~25ms (97, 480, 47)
# integrand_photon = (pbpfs_photon[:, None, :]*flux_fracs[:, :, None]).T # ~25ms (97, 480, 47)

# extincted_intensities_energy = np.trapz(integrand_energy, self.wl, axis=2) # (97, 480)
# extincted_intensities_photon = np.trapz(integrand_photon, self.wl, axis=2) # (97, 480)

# non_extincted_intensities_energy = np.trapz(pbpfs_energy, self.wl, axis=0)[:,None] # (97, 1)
# non_extincted_intensities_photon = np.trapz(pbpfs_photon, self.wl, axis=0)[:,None] # (97, 1)

# egrid = (extincted_intensities_energy/non_extincted_intensities_energy).reshape(len(teffs), len(ebvs), len(rvs), 1) # (97, 16, 30, 1)
# pgrid = (extincted_intensities_photon/non_extincted_intensities_photon).reshape(len(teffs), len(ebvs), len(rvs), 1) # (97, 16, 30, 1)

# self.ndp['blackbody'] = ndpolator.Ndpolator(basic_axes=(axes[0],))
# self.ndp['blackbody'].register('ext@photon', associated_axes=(axes[1], axes[2]), grid=pgrid)
# self.ndp['blackbody'].register('ext@energy', associated_axes=(axes[1], axes[2]), grid=egrid)
self.ndp['blackbody'].register('ext@photon', associated_axes=(axes[1], axes[2]), grid=pgrid[..., None])
self.ndp['blackbody'].register('ext@energy', associated_axes=(axes[1], axes[2]), grid=egrid[..., None])

if 'blackbody:ext' not in self.content:
self.content.append('blackbody:ext')
Expand Down Expand Up @@ -879,6 +854,10 @@ def compute_intensities(self, atm, path, include_extinction=False, rvs=None, ebv

ints_energy, ints_photon = np.empty(nmodels*len(mus)), np.empty(nmodels*len(mus))

keep = (wls >= self.ptf_table['wl'][0]) & (wls <= self.ptf_table['wl'][-1])
wls = wls[keep]
ptf = self.ptf(wls)

if include_extinction:
if ebvs is None:
ebvs = np.linspace(0., 3., 30)
Expand All @@ -889,8 +868,11 @@ def compute_intensities(self, atm, path, include_extinction=False, rvs=None, ebv
ext_photon_grid = np.empty(shape=(len(np.unique(teffs)), len(np.unique(loggs)), len(np.unique(abuns)), len(ebvs), len(rvs), 1))
ext_energy_grid = np.empty(shape=(len(np.unique(teffs)), len(np.unique(loggs)), len(np.unique(abuns)), len(ebvs), len(rvs), 1))

keep = (wls >= self.ptf_table['wl'][0]) & (wls <= self.ptf_table['wl'][-1])
wls = wls[keep]
axbx = libphoebe.gordon_extinction(wls)
ax, bx = axbx[:,0], axbx[:,1]

# The following code broadcasts arrays so that integration can be vectorized:
Alam = 10**(-0.4 * ebvs[None, :, None] * (rvs[None, None, :] * ax[:, None, None] + bx[:, None, None]))

for i, model in tqdm(enumerate(models), desc=atm, total=len(models), disable=not verbose, unit=' models'):
with fits.open(model) as hdu:
Expand All @@ -899,7 +881,7 @@ def compute_intensities(self, atm, path, include_extinction=False, rvs=None, ebv
# trim intensities to the passband limits:
seds = seds[:,keep]

pbints_energy = self.ptf(wls)*seds
pbints_energy = ptf*seds
fluxes_energy = np.trapz(pbints_energy, wls)

pbints_photon = wls*pbints_energy
Expand All @@ -912,14 +894,16 @@ def compute_intensities(self, atm, path, include_extinction=False, rvs=None, ebv
ints_photon[i*len(mus):(i+1)*len(mus)] = np.log10(fluxes_photon/self.ptf_photon_area) # photon-weighted intensity

if include_extinction:
axbx = libphoebe.gordon_extinction(wls)
ax, bx = axbx[:,0], axbx[:,1]
for ei, ebv in enumerate(ebvs):
for ri, rv in enumerate(rvs):
Alam = 10**(-0.4 * ebv * (rv * ax + bx))
t = (teffs[i] == ext_axes[0], loggs[i] == ext_axes[1], abuns[i] == ext_axes[2], ebvs[ei] == ext_axes[3], rvs[ri] == ext_axes[4],0)
ext_energy_grid[t] = np.trapz(self.ptf(wls) * seds * Alam, wls)[-1] / np.trapz(self.ptf(wls) * seds, wls)[-1]
ext_photon_grid[t] = np.trapz(wls * self.ptf(wls) * seds * Alam, wls)[-1] / np.trapz(wls * self.ptf(wls) * seds, wls)[-1]
# we only use normal emergent intensities here for simplicity:
epbints = pbints_energy[-1].reshape(-1, 1)
egrid = np.trapz(epbints[:, :, None, None] * Alam[:, None, :, :], wls, axis=0) / np.trapz(epbints[:, :, None, None], wls, axis=0)

ppbints = pbints_photon[-1].reshape(-1, 1)
pgrid = np.trapz(ppbints[:, :, None, None] * Alam[:, None, :, :], wls, axis=0) / np.trapz(ppbints[:, :, None, None], wls, axis=0)

t = (teffs[i] == ext_axes[0], loggs[i] == ext_axes[1], abuns[i] == ext_axes[2])
ext_energy_grid[t] = egrid.reshape(len(ebvs), len(rvs), 1)
ext_photon_grid[t] = pgrid.reshape(len(ebvs), len(rvs), 1)

basic_axes = (np.unique(teffs), np.unique(loggs), np.unique(abuns))
self.ndp[atm] = ndpolator.Ndpolator(basic_axes=basic_axes)
Expand Down

0 comments on commit 6decdf5

Please sign in to comment.