Merge pull request #80 from p-slash/fix-flux-calibration

Fix flux calibration and add tests

py/qsonic/calibration.py

@@ -145,7 +145,7 @@ def _read(self, fname):
         return FastLinear1DInterp(waves_0, dwave, stacked_flux)
 
     def __init__(self, fname, comm=None, mpi_rank=0):
-        self.flux_interp = mpi_fnc_bcast(
+        self.stacked_flux_interp = mpi_fnc_bcast(
             self._read, comm, mpi_rank,
             f"Error loading FluxCalibrator from file {fname}.",
             fname)

py/qsonic/scripts/qsonic_calib.py

@@ -119,7 +119,7 @@ def mpi_set_targetid_list_to_remove(args, comm=None, mpi_rank=0):
             args.remove_targetid_list, dtype=int)
 
     if args.catalog:
-        catalog = qsonic.io.mpi_read_quasar_catalog(
+        catalog = qsonic.catalog.mpi_read_quasar_catalog(
             args.catalog, comm, mpi_rank, args.mock_analysis)
     else:
         catalog = None

tests/test_calibration.py

@@ -0,0 +1,51 @@
+import pytest
+
+import fitsio
+import numpy as np
+import numpy.testing as npt
+
+import qsonic.spectrum
+import qsonic.calibration
+
+
+@pytest.fixture()
+def my_setup_attributes(tmp_path):
+    fname = tmp_path / "attributes.fits"
+
+    with fitsio.FITS(fname, 'rw', clobber=True) as fts:
+        wavelength = np.linspace(3600, 9800, 775)
+        stacked_flux = 5. * np.ones(wavelength.size)
+        fts.write(
+            [wavelength, stacked_flux],
+            names=['lambda', 'stacked_flux'],
+            extname='STACKED_FLUX')
+
+        wavelength = np.linspace(3600, 9800, 50)
+        ones_array = np.ones(wavelength.size)
+        fts.write(
+            [wavelength, ones_array, ones_array, ones_array * 2, ones_array],
+            names=['lambda', 'var_lss', 'e_var_lss', 'eta', 'e_eta'],
+            extname='VAR_FUNC')
+
+    yield fname
+
+
+def test_calibrations(setup_data, my_setup_attributes):
+    fname = my_setup_attributes
+    ncal = qsonic.calibration.NoiseCalibrator(fname)
+    fcal = qsonic.calibration.FluxCalibrator(fname)
+
+    cat_by_survey, _, data = setup_data(1)
+    spectra_list = qsonic.spectrum.generate_spectra_list_from_data(
+        cat_by_survey, data)
+    spec = spectra_list[0]
+    spec.set_forest_region(3600., 6000., 1050., 1300.)
+    f0 = spec.flux['B'][0]
+
+    ncal.apply(spectra_list)
+    npt.assert_allclose(spec.forestivar['B'], 0.5)
+    npt.assert_allclose(spec.forestivar['R'], 0.5)
+
+    fcal.apply(spectra_list)
+    npt.assert_allclose(spec.forestflux['B'], f0 / 5)
+    npt.assert_allclose(spec.forestflux['R'], f0 / 5)