diff --git a/src/xdust/halos/galhalo.py b/src/xdust/halos/galhalo.py index a091d02..88a0284 100644 --- a/src/xdust/halos/galhalo.py +++ b/src/xdust/halos/galhalo.py @@ -153,7 +153,7 @@ def variable_profile(self, time, lc, dist=8.0, tnow=None): Returns ------- - numpy.ndarray (NE x NTH) [fabs units / arcsec**2] + astropy.units.Quantity (NE x NTH) [fabs units / arcsec**2] Scattering halo intensity as a function of energy and observation angle. """ @@ -171,11 +171,11 @@ def variable_profile(self, time, lc, dist=8.0, tnow=None): theta_rad = self.theta/3600./180.*np.pi ne, ntheta = len(self.lam), len(self.theta) - inten = np.zeros(shape=(ne, ntheta)) + inten = np.zeros(shape=(ne, ntheta)) * self.norm_int.unit lctm = (time-tzero) for i in range(len(self.lam)): - deltat = time_delay(self.theta, self.x, dist) * u.second.to(u.day) + deltat = time_delay(self.theta.to(u.arcsec).value, self.x, dist) * u.second.to(u.day) t = tnow - deltat for j in range(ntheta): inten[i,j] += np.interp(t[j], time, lc * self.norm_int[i,j] * self.fabs[i]) @@ -237,16 +237,17 @@ def fake_variable_image(self, time, lc, arf, to simulate the number of counts in each pixel. """ # intensity cube (NE x NTH), phot/cm^2/s/arcsec^2 + var_profile = self.variable_profile(time, lc, dist=dist, tnow=tnow) # Decide which energy indexes to use if lmin is None: imin = 0 else: - imin = min(np.arange(len(self.lam))[self.lam >= lmin]) + imin = min(np.arange(len(self.lam))[self.lam >= lmin * self.lam.unit]) if lmax is None: iend = len(self.lam) else: - iend = max(np.arange(len(self.lam))[self.lam <= lmax]) + iend = max(np.arange(len(self.lam))[self.lam <= lmax * self.lam.unit]) # set up image grid xlen, ylen = num_pix @@ -261,13 +262,7 @@ def fake_variable_image(self, time, lc, arf, arf = InterpolatedUnivariateSpline(arf_x, arf_y, k=1) # Conversion erg -> ct for each energy bin - if self.lam_unit in ['Angs', 'Angstrom', 'angs', 'angstrom']: - ener = self.lam * u.angstrom - elif self.lam_unit in ['kev', 'keV']: - ener = self.lam * u.keV - else: - ener = self.lam * u.Unit(self.lam_unit) - int_conv = arf(ener.to(u.keV, equivalencies=u.spectral())) + int_conv = arf(self.lam.to(u.keV, equivalencies=u.spectral()).value) # cm^2 ct/phot r_asec = radius * pix_scale @@ -601,5 +596,4 @@ def time_delay(alpha, x, dkpc): Assumes small angle scattering. """ delta_x = path_diff(alpha, x) - d_cm = dkpc * 1.e3 * u.pc.to('cm') # cm - return delta_x * (d_cm / c.c).to('s').value # seconds + return delta_x * (dkpc * u.kpc / c.c).to('s').value diff --git a/src/xdust/halos/halo.py b/src/xdust/halos/halo.py index 11bbe74..8bc910d 100644 --- a/src/xdust/halos/halo.py +++ b/src/xdust/halos/halo.py @@ -403,23 +403,18 @@ def fake_image(self, arf, src_flux, exposure, arf = InterpolatedUnivariateSpline(arf_x, arf_y, k=1) # Source counts to use for each energy bin - if self.lam_unit == 'angs': - ltemp = self.lam * u.angstrom - ltemp_kev = ltemp.to(u.keV, equivalencies=u.spectral()).value - arf_temp = arf(ltemp_kev)[::-1] - src_counts = src_flux * arf_temp * exposure - else: - src_counts = src_flux * arf(self.lam) * exposure + lam_kev = self.lam.to(u.keV, equivalencies=u.spectral()).value + src_counts = src_flux * arf(lam_kev) * exposure # Decide which energy indexes to use if lmin is None: imin = 0 else: - imin = min(np.arange(len(self.lam))[self.lam >= lmin]) + imin = min(np.arange(len(self.lam))[self.lam >= lmin * self.lam.unit]) if lmax is None: iend = len(self.lam) else: - iend = max(np.arange(len(self.lam))[self.lam <= lmax]) + iend = max(np.arange(len(self.lam))[self.lam <= lmax * self.lam.unit]) #iend = imax #if imax < 0: diff --git a/tests/test_galhalo.py b/tests/test_galhalo.py index 45858bd..f8a30e6 100644 --- a/tests/test_galhalo.py +++ b/tests/test_galhalo.py @@ -2,6 +2,8 @@ import numpy as np from scipy.integrate import trapezoid as trapz import astropy.units as u +from astropy.io import fits +from astropy.table import Table from xdust.halos import * from xdust import grainpop @@ -167,3 +169,84 @@ def test_halo_io(test): assert np.all(percent_diff(new_halo.taux,test.taux) < 0.01) assert np.all(percent_diff(new_halo.norm_int.flatten(),test.norm_int.flatten()) < 0.01) assert new_halo.lam.unit == 'keV' + +def create_arf(tmp_path, outfile, evals, effarea) -> str: + """ + Create a dummy ARF file for testing + """ + t = Table() + t['ENERG_LO'] = evals[:-1] + t['ENERG_HI'] = evals[1:] + t['SPECRESP'] = np.full(len(evals) - 1, effarea) + + hdu_table = fits.BinTableHDU(t.as_array(), name='SPECRESP') + hdu_primary = fits.PrimaryHDU() + hdul = fits.HDUList([hdu_primary, hdu_table]) + hdul.writeto(tmp_path / outfile, overwrite=True) + return str(tmp_path / outfile) + +@pytest.fixture +def halo(): + new_halo = galhalo.UniformGalHalo(EVALS, THVALS) + new_halo.calculate(GPOP) + return new_halo + +@pytest.fixture +def screen_halo(): + new_halo = galhalo.ScreenGalHalo(EVALS, THVALS) + new_halo.calculate(GPOP, x=0.5) + new_halo.calculate_intensity(FABS) + return new_halo + +@pytest.fixture +def arf(tmp_path): + arf = create_arf(tmp_path, 'test_arf.fits', EVALS, 100.0) + return arf + +def test_fake_image(halo, arf): + image = halo.fake_image(arf, src_flux=FABS, exposure=1e4, pix_scale=1.0, num_pix=[8, 16]) + assert image.shape == (16, 8) + assert np.all(image >= 0.0) + assert np.all(image == np.floor(image)) + assert image.sum() > 0 + +def test_fake_image_lmin_lmax(halo, arf): + restricted = halo.fake_image(arf, src_flux=FABS, exposure=1e4, pix_scale=1.0, num_pix=[8, 16], lmin=0.5, lmax=5.0) + unrestricted = halo.fake_image(arf, src_flux=FABS, exposure=1e4, pix_scale=1.0, num_pix=[8, 16]) + assert 0 < restricted.sum() < unrestricted.sum() + +def test_time_delay(): + result1 = galhalo.time_delay(100.0, 0.5, 8.0) + assert result1 == pytest.approx(96770, rel=1e-3) + + result2 = galhalo.time_delay(200.0, 0.5, 8.0) + assert result2 == pytest.approx(4 * result1, rel=1e-3) + +def test_variable_profile(screen_halo): + time = np.linspace(0, 100, 10) + lc = np.ones_like(time) + dist = 8.0 + var_profile1 = screen_halo.variable_profile(time, lc, dist=dist) + assert var_profile1.shape == (len(EVALS), len(THVALS)) + assert np.sum(var_profile1) > 0 + + var_profile2 = screen_halo.variable_profile(time, 2 * lc, dist=dist) + assert var_profile2.shape == (len(EVALS), len(THVALS)) + assert np.sum(var_profile2) > 0 + assert np.sum(var_profile2).value == pytest.approx(2 * np.sum(var_profile1).value) + +def test_fake_variable_image(screen_halo, arf): + time = np.linspace(0, 100, 10) + lc = np.ones_like(time) + image = screen_halo.fake_variable_image(time, lc=lc, arf=arf, exposure=1e4, pix_scale=1.0, num_pix=[8, 16]) + assert image.shape == (16, 8) + assert np.all(image >= 0.0) + assert np.all(image == np.floor(image)) + assert image.sum() > 0 + +def test_fake_variable_image_lmin_lmax(screen_halo, arf): + time = np.linspace(0, 100, 10) + lc = np.ones_like(time) + restricted = screen_halo.fake_variable_image(time, lc=lc, arf=arf, exposure=1e4, num_pix=[8, 16], lmin=0.5, lmax=5.0) + unrestricted = screen_halo.fake_variable_image(time, lc=lc, arf=arf, exposure=1e4, num_pix=[8, 16]) + assert 0 < restricted.sum() < unrestricted.sum()