From bd25176b4b6c80ea6b4c61750baf2eb2d96b4f12 Mon Sep 17 00:00:00 2001
From: Hugo Buddelmeijer <hugo@buddelmeijer.nl>
Date: Mon, 28 Oct 2024 20:28:31 +0100
Subject: [PATCH] Fix and test calibration module

---
 scopesim_templates/calibration/calibration.py | 44 ++++++++++++-------
 .../test_calibration/test_calibration.py      | 29 ++++++++++++
 2 files changed, 58 insertions(+), 15 deletions(-)

diff --git a/scopesim_templates/calibration/calibration.py b/scopesim_templates/calibration/calibration.py
index b200023..b84bf5a 100644
--- a/scopesim_templates/calibration/calibration.py
+++ b/scopesim_templates/calibration/calibration.py
@@ -6,15 +6,17 @@
 import numpy as np
 from astropy import units as u
 
-from spextra import Spextrum
+from spextra import Spextrum, Passband
 
 from ..rc import Source, __config__
 from ..utils.general_utils import add_function_call_str
 from ..misc.misc import uniform_source
 
-__all__ = ["lamp",
-           "flat_field",
-           "empty_sky"]
+__all__ = [
+    "lamp",
+    "flat_field",
+    "empty_sky",
+]
 
 
 @add_function_call_str
@@ -62,7 +64,10 @@ def lamp(waves, fwhm, fluxes) -> Source:
 
     w_minmax = [waves.min().value, waves.max().value] * u.AA
     # A very faint spextrum to add the lines
-    spec = Spextrum.flat_spectrum(amplitude=40, waves=w_minmax)
+    spec = Spextrum.flat_spectrum(
+        amplitude=40 * u.ABmag,
+        waves=w_minmax,
+    )
     spec = spec.add_emi_lines(center=waves, fwhm=fwhm, flux=fluxes)
     src = uniform_source(sed=spec)
     # src.meta.update({"object": "lamp"})
@@ -71,8 +76,12 @@ def lamp(waves, fwhm, fluxes) -> Source:
 
 
 @add_function_call_str
-def flat_field(temperature=5000, amplitude=0*u.ABmag, filter_curve="V",
-               extend=60) -> Source:
+def flat_field(
+    temperature: u.Quantity[u.K] = 5000 * u.K,
+    amplitude: u.Quantity[u.ABmag] = 0 * u.ABmag,
+    filter_curve: str | Passband = "V",
+    extend=60,
+) -> Source:
     """
     Create a flat-field source to be used in ``ScopeSim``.
 
@@ -87,9 +96,9 @@ def flat_field(temperature=5000, amplitude=0*u.ABmag, filter_curve="V",
 
     Parameters
     ----------
-    temperature: float
+    temperature: ``astropy.Quantity``
         [Kelvin] Temperature of the lamp
-    amplitude: u.Quantity
+    amplitude: ``astropy.Quantity``
         [u.Quantity] amplitude of the lamp in u.ABmag, u.mag (vega) or u. STmag
     filter_curve: str
         any filter curve available for ``spextra``, used for scaling
@@ -98,17 +107,22 @@ def flat_field(temperature=5000, amplitude=0*u.ABmag, filter_curve="V",
     Returns
     -------
     src: Source
-
     """
     # I don't understand this warning but it probably has a reason...
     warnings.warn("Make sure to turn off the appropriate ScopeSim effects "
                   "during the simulation", RuntimeWarning)
 
-    spec = Spextrum.black_body_spectrum(temperature=temperature,
-                                        amplitude=amplitude,
-                                        filter_curve=filter_curve)
-    src = uniform_source(sed=spec, amplitude=amplitude,
-                         filter_curve=filter_curve, extend=extend)
+    spec = Spextrum.black_body_spectrum(
+        temperature=temperature,
+        amplitude=amplitude,
+        filter_curve=filter_curve,
+    )
+    src = uniform_source(
+        sed=spec,
+        amplitude=amplitude,
+        filter_curve=filter_curve,
+        extend=extend,
+    )
     # src.meta.update({"object": "flat_field"})
 
     return src
diff --git a/scopesim_templates/tests/test_calibration/test_calibration.py b/scopesim_templates/tests/test_calibration/test_calibration.py
index 985855d..ade6f48 100644
--- a/scopesim_templates/tests/test_calibration/test_calibration.py
+++ b/scopesim_templates/tests/test_calibration/test_calibration.py
@@ -1,4 +1,7 @@
+from astropy.io.fits import ImageHDU
 from astropy.table import Table
+from astropy import units as u
+import numpy as np
 from synphot import SourceSpectrum
 
 from scopesim_templates import calibration
@@ -13,3 +16,29 @@ def test_empty_sky_returns_source_object(self):
         assert (isinstance(getattr(sky.fields[0], "field", sky.fields[0]), Table)
                 or isinstance(sky.fields[0], Table))
         assert sky.fields[0]["ref"][0] == 0
+
+
+class TestFlatField:
+    def test_flat_field_returns_source_object(self):
+        flatfield = calibration.calibration.flat_field()
+        assert isinstance(flatfield, Source)
+        assert isinstance(flatfield.spectra[0], SourceSpectrum)
+        assert (isinstance(getattr(flatfield.fields[0], "field", flatfield.fields[0]), ImageHDU)
+                or isinstance(flatfield.fields[0], ImageHDU))
+
+
+class TestLamp:
+    def test_lamp_returns_source_object(self):
+        res = 20000  # Resolution of the instrument
+        waves = np.arange(1e4, 2.7e4, 100)
+        fwhm = 2.6*(waves/res)*np.ones(shape=waves.shape)  # Nyquist sampled
+        flux = 1e-10*np.ones(waves.shape)
+        lamp = calibration.calibration.lamp(
+            waves=waves,
+            fwhm=fwhm,
+            fluxes=flux,
+        )
+        assert isinstance(lamp, Source)
+        assert isinstance(lamp.spectra[0], SourceSpectrum)
+        assert (isinstance(getattr(lamp.fields[0], "field", lamp.fields[0]), ImageHDU)
+                or isinstance(lamp.fields[0], ImageHDU))