Add Quantization effect

scopesim/effects/electronic.py

@@ -386,7 +386,6 @@
             above = np.invert(below)
             data[below] = np.random.poisson(data[below]).astype(float)
             data[above] = np.random.normal(data[above], np.sqrt(data[above]))
-            data = np.floor(data)
             new_imagehdu = fits.ImageHDU(data=data, header=det._hdu.header)
             det._hdu = new_imagehdu
 
@@ -585,6 +584,35 @@
         return det
 
 
+class Quantization(Effect):
+    """Converts raw data to whole photons."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        params = {
+            "z_order": [825],
+            "dtype": "uint32",
+        }
+        self.meta.update(params)
+        self.meta.update(kwargs)
+
+    def apply_to(self, obj, **kwargs):
+        if not isinstance(obj, DetectorBase):
+            return obj
+
+        new_dtype = self.meta["dtype"]
+        if not np.issubdtype(new_dtype, np.integer):
+            logging.warning("Setting quantized data to dtype %s, which is not "
+                            "an integer subtype.", new_dtype)
+
+        # This used to create a new ImageHDU with the same header but the data
+        # set to the modified data. It should be fine to simply re-assign the
+        # data attribute, but just in case it's not...
+        obj._hdu.data = np.floor(obj._hdu.data).astype(new_dtype)
+
+        return obj
+
+
 def make_ron_frame(image_shape, noise_std, n_channels, channel_fraction,
                    line_fraction, pedestal_fraction, read_fraction):
     shape = image_shape

scopesim/tests/tests_effects/test_Quantization.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+"""Contains tests for Shutter class."""
+
+import pytest
+import numpy as np
+
+from scopesim.detector import Detector
+from scopesim.effects.electronic import Quantization
+
+from scopesim.tests.mocks.py_objects.header_objects import _implane_header
+
+
+@pytest.fixture
+def mock_detector():
+    det = Detector(_implane_header())
+    return det
+
+
+@pytest.fixture(scope="function", name="det")
+def detector_with_data(mock_detector):
+    det = mock_detector
+    width = det._hdu.data.shape[1]
+    det._hdu.data[:] = 1.2
+    det._hdu.data[:, width//2] = 1.99
+    return det
+
+
+class TestInit:
+    def test_initialised_with_nothing(self):
+        quant = Quantization()
+        assert isinstance(quant, Quantization)
+
+
+class TestApplyTo:
+    def test_floors_pixels_to_integer_with_default(self, det):
+        assert det.data.sum() > det.data.size * 1.2
+        quant = Quantization()
+        det = quant.apply_to(det)
+        assert det.data.sum() == det.data.size
+        assert det.data.dtype == np.uint32
+
+    def test_floors_pixels_to_integer_with_custom(self, det):
+        assert det.data.sum() > det.data.size * 1.2
+        quant = Quantization(dtype="int16")
+        det = quant.apply_to(det)
+        assert det.data.sum() == det.data.size
+        assert det.data.dtype == np.int16
+
+    def test_logs_warning_for_float_dtype(self, det, caplog):
+        assert det.data.sum() > det.data.size * 1.2
+        quant = Quantization(dtype=float)
+        det = quant.apply_to(det)
+        assert det.data.sum() == det.data.size
+        # Test sub-dtype because exact dtype for "float" may depend on OS
+        assert np.issubdtype(det.data.dtype, np.floating)
+        assert "Setting quantized data to dtype" in caplog.text