Absolute Flux Calculation Fix (#120)

* Fixing absolute flux calculation which was missing a squared. * Adding tests for flux calibration. * Adding tests for flux calibration but no check yet * flux calibrate tests * add more tests * fix warning * units change * plot=False * un-add tests --------- Co-authored-by: kelle <[email protected]>
BDNYC · Dec 21, 2024 · aa04f48 · aa04f48
1 parent f1b7430
commit aa04f48
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 20 deletions.
diff --git a/sedkit/spectrum.py b/sedkit/spectrum.py
@@ -23,7 +23,7 @@
 import numpy as np
 from pandas import DataFrame
 from svo_filters import Filter
-
+import warnings
 
 from . import mcmc as mc
 from . import utilities as u
@@ -528,7 +528,7 @@ def flux_calibrate(self, distance, target_distance=10 * q.pc, flux_units=None):
         if self.unc is None:
             unc = None
         else:
-            term1 = (self.spectrum[2] * distance[0] / target_distance).to(flux_units)
+            term1 = (self.spectrum[2] * (distance[0] / target_distance) ** 2).to(flux_units)
             term2 = (2 * self.spectrum[1] * (distance[1] * distance[0] / target_distance**2)).to(flux_units)
             unc = np.sqrt(term1**2 + term2**2)
 
@@ -1494,7 +1494,11 @@ def __init__(self, file, wave_units=None, flux_units=None, ext=0,
 
         # Sanity check for wave_units
         if data[0].min() > 100 and wave_units == q.um:
-            print("WARNING: Your wavelength range ({} - {}) looks like Angstroms. Are you sure it's {}?".format(data[0].min(), data[0].max(), wave_units))
+            msg = (
+                f"Your wavelength range ({data[0].min()} - {data[0].max()})"
+                f"looks like Angstroms. Are you sure it's {wave_units}?"
+            )
+            warnings.warn(msg)
 
         # Apply units
         wave = data[0] * wave_units

diff --git a/tests/test_sed.py b/tests/test_sed.py
@@ -217,10 +217,10 @@ def test_edit_spectrum(self):
         s.add_spectrum(self.spec1)
 
         # Smooth it
-        s.edit_spectrum(0, smooth={'beta': 5}, plot=True)
+        s.edit_spectrum(0, smooth={'beta': 5}, plot=False)
 
         # Unsmooth it
-        s.edit_spectrum(0, restore=True)
+        s.edit_spectrum(0, restore=True, plot=False)
 
         # Bad beta
         self.assertRaises((ValueError, TypeError), s.edit_spectrum, idx=0, smooth={'beta': 'foo'})

diff --git a/tests/test_spectrum.py b/tests/test_spectrum.py
@@ -149,6 +149,20 @@ def test_integrate(self):
         self.assertAlmostEqual(fbol[0].value, 4000, places=1)
         self.assertNotEqual(str(fbol[1].value), 'nan')
 
+    def test_flux_calibrate(self):
+        """ Test that flux calibrate is working properly"""
+        # 1 parsec Distance with unc
+        close_distance = 1 * q.pc, 0.05 * q.pc
+        abs_sed_close = sp.Spectrum.flux_calibrate(self.flat1, close_distance)
+        assert np.isclose(np.mean(abs_sed_close.flux),0.01, rtol=0.005)
+        assert np.isclose(np.mean(abs_sed_close.unc),0.001, rtol=0.005)
+
+        # 15 parsec Distance with unc
+        far_distance = 15 * q.pc, 2 * q.pc
+        abs_sed_far = sp.Spectrum.flux_calibrate(self.flat1, far_distance)
+        assert np.isclose(np.mean(abs_sed_far.flux), 2.25, rtol=0.005)
+        assert np.isclose(np.mean(abs_sed_far.unc), 0.6, rtol=0.005)
+
     def test_interpolate(self):
         """Test interpolate method"""
         spec1 = self.flat1
@@ -226,7 +240,7 @@ def test_norm_to_spec(self):
         s2 = self.flat2
 
         # Normalize 1 to 2 and check that they are close
-        s3 = s1.norm_to_spec(s2, plot=True)
+        s3 = s1.norm_to_spec(s2, plot=False)
         self.assertAlmostEqual(np.nanmean(s2.flux), np.nanmean(s3.flux), places=4)
         self.assertNotEqual(s2.size, s3.size)
         self.assertEqual(s1.size, s3.size)
@@ -236,24 +250,19 @@ def test_trim(self):
         # Test include
         s1 = copy.copy(self.flat1)
         trimmed = s1.trim(include=[(0.8 * q.um, 2 * q.um)], concat=False)
-        # self.assertTrue(len(trimmed) == 1)
-        # self.assertNotEqual(self.flat1.size, trimmed[0].size)
+        assert len(trimmed[0].wave) == 115
 
         # Test exclude
-        s1 = copy.copy(self.flat1)
-        trimmed = s1.trim(exclude=[(0.8 * q.um, 3 * q.um)], concat=False)
-        # self.assertNotEqual(self.flat1.size, trimmed[0].size)
+        trimmed2 = s1.trim(exclude=[(0.8 * q.um, 3 * q.um)], concat=False)
 
         # Test split
-        s1 = copy.copy(self.flat1)
-        trimmed = s1.trim(exclude=[(0.8 * q.um, 0.9 * q.um)], concat=False)
-        # self.assertTrue(len(trimmed) == 2)
-        # self.assertNotEqual(self.flat1.size, trimmed[0].size)
+        trimmed3 = s1.trim(exclude=[(0.8 * q.um, 0.9 * q.um)], concat=False)
+        # assert len(trimmed3[0].wave) == 115
+
 
         # Test concat
-        s1 = copy.copy(self.flat1)
-        trimmed = s1.trim(exclude=[(0.8 * q.um, 0.9 * q.um)], concat=True)
-        # self.assertNotEqual(self.flat1.size, trimmed.size)
+        trimmed4 = s1.trim(exclude=[(0.8 * q.um, 0.9 * q.um)], concat=True)
+        # assert len(trimmed4.wave) == 173
 
 
 class TestFileSpectrum(unittest.TestCase):
@@ -266,11 +275,13 @@ def setUp(self):
 
     def test_fits(self):
         """Test that a fits file can be loaded"""
-        spec = sp.FileSpectrum(self.fitsfile, wave_units='um', flux_units='erg/s/cm2/AA')
+        spec = sp.FileSpectrum(self.fitsfile, wave_units='um', flux_units='erg s-1 cm-2 AA-1')
+        assert spec
 
     def test_txt(self):
         """Test that a txt file can be loaded"""
-        spec = sp.FileSpectrum(self.txtfile, wave_units='um', flux_units='erg/s/cm2/AA')
+        spec = sp.FileSpectrum(self.txtfile, wave_units='um', flux_units='erg s-1 cm-2 AA-1')
+        assert spec
 
 
 class TestVega(unittest.TestCase):