zblz / naima

Derivation of non-thermal particle distributions through MCMC spectral fitting
http://naima.readthedocs.io
BSD 3-Clause "New" or "Revised" License
45 stars 53 forks source link

Improve speed of SSC InverseCompton #182

Open kbruegge opened 5 years ago

kbruegge commented 5 years ago

I'm currently fitting SSC models. Its pretty slow though. You think it might be a good idea to try and speed up the radiative models with numba? See here for a quick overview: https://numba.pydata.org/numba-doc/latest/user/5minguide.html

I can make some quick speed tests next week.

kbruegge commented 5 years ago

So after some profiling it turns out that the numerical part might not be the problem. The first call of flux of the InverseCompton model is much slower than following ones. I ran some tests and got

▶ python speed_test.py           
First call took:  10.792217016220093
Second call took:  0.00581812858581543

I'm investigating further.

kbruegge commented 5 years ago

It only shows that strongly when using an SSC photonfield.

Here's my test code btw:

from naima.models import PowerLaw, InverseCompton, Synchrotron
from astropy.constants import c
import numpy as np
import astropy.units as u

def ssc_model():
    electron_population = PowerLaw(1E48/u.eV, alpha=2, e_0=1 * u.TeV)
    SYN = Synchrotron(electron_population, B=100 * u.uG, Eemax=50 * u.PeV, Eemin=0.1 * u.MeV)

    Rpwn = 2.0 * u.pc
    Esy = np.logspace(-7, 12, 80) * u.eV
    Lsy = SYN.flux(Esy, distance=0 * u.cm)
    phn_sy = Lsy / (4 * np.pi * Rpwn ** 2 * c) * 2.25

    IC = InverseCompton(
        electron_population,
        seed_photon_fields=[
            "CMB",
            ["FIR", 70 * u.K, 0.5 * u.eV / u.cm ** 3],
            ["NIR", 5000 * u.K, 1 * u.eV / u.cm ** 3],
            ["SSC", Esy, phn_sy],
        ],
        Eemax=50 * u.PeV,
        Eemin=0.1 * u.GeV,
    )

    return SYN, IC

SYN, IC = ssc_model()

energies = np.logspace(-8.1, 3, 300) * u.TeV
import time
start = time.time()
IC.flux(energies)
end = time.time()
print(f"First call took:  {end-start}")

start = time.time()
IC.flux(energies)
end = time.time()
print(f"Second call took:  {end-start}")
kbruegge commented 5 years ago

Turns out there's a memoize decorator on that. Stupid me.

kbruegge commented 5 years ago

Note to self: it seems like its the _iso_ic_on_monochromatic function.

zblz commented 5 years ago

Thanks for looking into this, @mackaiver. SSC is intrinsically much more compute intensive that BB or monochromatic origins as naima needs to iterate over the wavelengths of the Synchroton spectrum when computing the IC spectrum.

I experimented with numba in the past, but found that since most of the computations in naima are handled by numpy, the improvements where not significant enough for the increase in complexity to be worth it. However, I tried that quite a long time ago and things might have changed, so I'd encourage you to give it a go.

Turns out there's a memoize decorator on that.

You can disable memoization on a given model by setting IC._memoize = False.

For any speed improvement work, it might be worth also looking at the naima-benchmarks repo, where there are a set of airspeed velocity benchmarks that time a few of the naima models.