Merge pull request #1723 from cta-observatory/add_noise

Components to modify dl1 images

ctapipe/image/__init__.py

@@ -59,10 +59,12 @@
     ring_completeness,
     ring_containment,
 )
+from .modifications import ImageModifier
 from .image_processor import ImageProcessor
 
 
 __all__ = [
+    "ImageModifier",
     "ImageProcessor",
     "hillas_parameters",
     "HillasParameterizationError",

ctapipe/image/image_processor.py

@@ -16,6 +16,7 @@
 from ..instrument import SubarrayDescription
 from . import (
     ImageCleaner,
+    ImageModifier,
     concentration_parameters,
     descriptive_statistics,
     hillas_parameters,
@@ -89,6 +90,8 @@ def __init__(
         self.clean = ImageCleaner.from_name(
             self.image_cleaner_type, subarray=subarray, parent=self
         )
+        self.modify = ImageModifier(subarray=subarray, parent=self)
+
         self.check_image = ImageQualityQuery(parent=self)
         if self.use_telescope_frame:
             telescope_frame = TelescopeFrame()
@@ -194,6 +197,7 @@ def _process_telescope_event(self, event):
             else:
                 geometry = self.subarray.tel[tel_id].camera.geometry
             # compute image parameters only if requested to write them
+            dl1_camera.image = self.modify(tel_id=tel_id, image=dl1_camera.image)
             dl1_camera.image_mask = self.clean(
                 tel_id=tel_id,
                 image=dl1_camera.image,

ctapipe/image/modifications.py

@@ -0,0 +1,189 @@
+import numpy as np
+from numba import njit
+from ..core.component import TelescopeComponent
+from ..core.traits import FloatTelescopeParameter, BoolTelescopeParameter, Int
+from ..instrument import SubarrayDescription
+
+
+__all__ = [
+    "ImageModifier",
+]
+
+
+def _add_noise(image, noise_level, rng=None, correct_bias=True):
+    """
+    Create a new image with added poissonian noise
+    """
+    if not rng:
+        rng = np.random.default_rng()
+
+    noisy_image = image.copy()
+    noise = rng.poisson(noise_level, size=image.shape)
+    noisy_image += noise
+
+    if correct_bias:
+        noisy_image -= noise_level
+
+    return noisy_image
+
+
+@njit(cache=True)
+def _smear_psf_randomly(
+    image, fraction, indices, indptr, smear_probabilities, seed=None
+):
+    """
+    Create a new image with values smeared across the
+    neighbor pixels specified by `indices` and `indptr`.
+    These are what defines the sparse neighbor matrix
+    and are available as attributes from the neighbor matrix.
+    The amount of charge that is distributed away from a given
+    pixel is drawn from a poissonian distribution.
+    The distribution of this charge among the neighboring
+    pixels follows a multinomial.
+    Pixels at the camera edge lose charge this way.
+    No geometry is available in this function due to the numba optimization,
+    so the indices, indptr and smear_probabilities have to match
+    to get sensible results.
+
+    Parameters:
+    -----------
+    image: ndarray
+        1d array of the pixel charge values
+    fraction: float
+        fraction of charge that will be distributed among neighbors (modulo poissonian)
+    indices: ndarray[int]
+        CSR format index array of the neighbor matrix
+    indptr: ndarray[int]
+        CSR format index pointer array of the neighbor matrix
+    smear_probabilities: ndarray[float]
+        shape: (n_neighbors, )
+        A priori distribution of the charge amongst neighbors.
+        In most cases probably of the form np.full(n_neighbors, 1/n_neighbors)
+    seed: int
+        Random seed for the numpy rng.
+        Because this is numba optimized, a rng instance can not be used here
+
+    Returns:
+    --------
+    new_image: ndarray
+        1d array with smeared values
+    """
+    new_image = image.copy()
+    if seed is not None:
+        np.random.seed(seed)
+
+    for pixel in range(len(image)):
+
+        if image[pixel] <= 0:
+            continue
+
+        to_smear = np.random.poisson(image[pixel] * fraction)
+        if to_smear == 0:
+            continue
+
+        # remove light from current pixel
+        new_image[pixel] -= to_smear
+
+        # add light to neighbor pixels
+        neighbors = indices[indptr[pixel] : indptr[pixel + 1]]
+        n_neighbors = len(neighbors)
+
+        # we always distribute the charge as if the maximum number
+        # of neighbors of a geoemtry is present, so that charge
+        # on the edges of the camera is lost
+        neighbor_charges = np.random.multinomial(to_smear, smear_probabilities)
+
+        for n in range(n_neighbors):
+            neighbor = neighbors[n]
+            new_image[neighbor] += neighbor_charges[n]
+
+    return new_image
+
+
+class ImageModifier(TelescopeComponent):
+    """
+    Component to tune simulated background to
+    overserved NSB values.
+    A differentiation between bright and dim pixels is taking place
+    because this happens at DL1a level and in general the integration window
+    would change for peak-searching extraction algorithms with different background levels
+    introducing a bias to the charge in dim pixels.
+
+    The performed steps include:
+    - Smearing of the image (simulating a worse PSF)
+    - Adding poissonian noise (different for bright and dim pixels)
+    - Adding a bias to dim pixel charges (see above)
+    """
+
+    psf_smear_factor = FloatTelescopeParameter(
+        default_value=0.0, help="Fraction of light to move to each neighbor"
+    ).tag(config=True)
+    noise_transition_charge = FloatTelescopeParameter(
+        default_value=0.0, help="separation between dim and bright pixels"
+    ).tag(config=True)
+    noise_bias_dim_pixels = FloatTelescopeParameter(
+        default_value=0.0, help="extra bias to add in dim pixels"
+    ).tag(config=True)
+    noise_level_dim_pixels = FloatTelescopeParameter(
+        default_value=0.0, help="expected extra noise in dim pixels"
+    ).tag(config=True)
+    noise_level_bright_pixels = FloatTelescopeParameter(
+        default_value=0.0, help="expected extra noise in bright pixels"
+    ).tag(config=True)
+    noise_correct_bias = BoolTelescopeParameter(
+        default_value=True, help="If True subtract the expected noise from the image."
+    ).tag(config=True)
+    rng_seed = Int(default_value=1337, help="Seed for the random number generator").tag(
+        config=True
+    )
+
+    def __init__(
+        self, subarray: SubarrayDescription, config=None, parent=None, **kwargs
+    ):
+        """
+        Parameters
+        ----------
+        subarray: SubarrayDescription
+            Description of the subarray. Provides information about the
+            camera which are useful in calibration. Also required for
+            configuring the TelescopeParameter traitlets.
+        config: traitlets.loader.Config
+            Configuration specified by config file or cmdline arguments.
+            Used to set traitlet values.
+            This is mutually exclusive with passing a ``parent``.
+        parent: ctapipe.core.Component or ctapipe.core.Tool
+            Parent of this component in the configuration hierarchy,
+            this is mutually exclusive with passing ``config``
+        """
+
+        super().__init__(subarray=subarray, config=config, parent=parent, **kwargs)
+        self.rng = np.random.default_rng(self.rng_seed)
+
+    def __call__(self, tel_id, image, rng=None):
+        geom = self.subarray.tel[tel_id].camera.geometry
+        smeared_image = _smear_psf_randomly(
+            image=image,
+            fraction=self.psf_smear_factor.tel[tel_id],
+            indices=geom.neighbor_matrix_sparse.indices,
+            indptr=geom.neighbor_matrix_sparse.indptr,
+            smear_probabilities=np.full(geom.max_neighbors, 1 / geom.max_neighbors),
+            seed=self.rng.integers(0, np.iinfo(np.int64).max),
+        )
+
+        noise = np.where(
+            image > self.noise_transition_charge.tel[tel_id],
+            self.noise_level_bright_pixels.tel[tel_id],
+            self.noise_level_dim_pixels.tel[tel_id],
+        )
+        image_with_noise = _add_noise(
+            smeared_image,
+            noise,
+            rng=self.rng,
+            correct_bias=self.noise_correct_bias.tel[tel_id],
+        )
+        bias = np.where(
+            image > self.noise_transition_charge.tel[tel_id],
+            0,
+            self.noise_bias_dim_pixels.tel[tel_id],
+        )
+        return (image_with_noise + bias).astype(np.float32)

ctapipe/image/tests/test_modifications.py

@@ -0,0 +1,72 @@
+import numpy as np
+from ctapipe.instrument import CameraGeometry
+from ctapipe.image import modifications
+
+
+def test_add_noise():
+    """
+    Test that adding noise changes a dummy
+    image in the expected way.
+    """
+    image = np.array([0, 0, 5, 1, 0, 0])
+    rng = np.random.default_rng(42)
+    # test different noise per pixel:
+    noise = [6, 8, 0, 7, 9, 12]
+    noisy = modifications._add_noise(image, noise, rng, correct_bias=False)
+    assert image[2] == noisy[2]
+    # For other seeds there exists a probability > 0 for no noise added at all
+    assert noisy.sum() > image.sum()
+
+    # test scalar
+    noisy = modifications._add_noise(image, 20, rng, correct_bias=False)
+    diff_no_bias = noisy - image
+    assert (noisy > image).all()
+
+    # test bias
+    noisy = modifications._add_noise(image, 20, rng, correct_bias=True)
+    assert np.sum(diff_no_bias) > np.sum(noisy - image)
+
+
+def test_smear_image():
+    """
+    Test that smearing the image leads to the expected results.
+    For random smearing this will not work with all seeds.
+    With the selected seed it will not lose charge
+    (No pixel outside of the camera receives light) but
+    for positive smear factors the image will be different
+    from the input
+    (At least for one pixel a positive charge is selected to
+    be distributed among neighbors).
+    """
+    seed = 20
+
+    # Hexagonal geometry -> Thats why we divide by 6 below
+    geom = CameraGeometry.from_name("LSTCam")
+    image = np.zeros_like(geom.pix_id, dtype=np.float64)
+    # select two pixels, one at the edge with only 5 neighbors
+    signal_pixels = [1, 1853]
+    neighbors = geom.neighbor_matrix[signal_pixels]
+
+    for signal_value in [1, 5]:
+        image[signal_pixels] = signal_value
+        for fraction in [0, 0.2, 1]:
+            # random smearing
+            # The seed is important here (See below)
+            smeared = modifications._smear_psf_randomly(
+                image,
+                fraction=fraction,
+                indices=geom.neighbor_matrix_sparse.indices,
+                indptr=geom.neighbor_matrix_sparse.indptr,
+                smear_probabilities=np.full(6, 1 / 6),
+                seed=seed,
+            )
+            neighbors_1 = smeared[neighbors[0]]
+
+            # this can be False if the "pseudo neighbor" of pixel
+            # 1853 is selected (outside of the camera)
+            assert np.isclose(image.sum(), smeared.sum())
+            assert np.isclose(np.sum(neighbors_1) + smeared[1], image[1])
+            # this can be False if for both pixels a 0 is
+            # drawn from the poissonian (especially with signal value 1)
+            if fraction > 0:
+                assert not ((image > 0) == (smeared > 0)).all()

ctapipe/instrument/camera/geometry.py

@@ -408,16 +408,16 @@ def _all_pixel_areas_equal(self):
         return ~np.any(~np.isclose(self.pix_area.value, self.pix_area[0].value), axis=0)
 
     def image_index_to_cartesian_index(self, pixel_index):
-        '''
+        """
         Convert pixel index in the 1d image representation to row and col
-        '''
+        """
         rows, cols = self._pixel_positions_2d
         return rows[pixel_index], cols[pixel_index]
 
     def cartesian_index_to_image_index(self, row, col):
-        '''
+        """
         Convert cartesian index (row, col) to pixel index in 1d representation.
-        '''
+        """
         return self._pixel_indices_cartesian[row, col]
 
     @lazyproperty
@@ -665,6 +665,10 @@ def neighbors(self):
     def neighbor_matrix(self):
         return self.neighbor_matrix_sparse.A
 
+    @lazyproperty
+    def max_neighbors(self):
+        return self.neighbor_matrix_sparse.sum(axis=1).max()
+
     @lazyproperty
     def neighbor_matrix_sparse(self):
         if self._neighbors is not None:

ctapipe/io/hdf5tableio.py

@@ -35,7 +35,8 @@
     "uint16": tables.UInt16Col,
     "uint32": tables.UInt32Col,
     "uint64": tables.UInt64Col,
-    "bool": tables.BoolCol,
+    "bool": tables.BoolCol,     # python bool
+    "bool_": tables.BoolCol,    # np.bool_
 }
 
 
@@ -256,9 +257,10 @@ class Schema(tables.IsDescription):
                     pos += 1
                 except ValueError:
                     self.log.warning(
-                        f"Column {col_name} of "
-                        f"container {container.__class__.__name__} in "
-                        f"table {table_name} not writable, skipping"
+                        f"Column {col_name}"
+                        f" with value {value!r} of type {type(value)} "
+                        f" of container {container.__class__.__name__} in"
+                        f" table {table_name} not writable, skipping"
                     )
         self._schemas[table_name] = Schema
         meta["CTAPIPE_VERSION"] = ctapipe.__version__

ctapipe/tools/process.py

@@ -8,7 +8,7 @@
 from ..calib import CameraCalibrator, GainSelector
 from ..core import QualityQuery, Tool
 from ..core.traits import Bool, classes_with_traits, flag
-from ..image import ImageCleaner, ImageProcessor
+from ..image import ImageCleaner, ImageProcessor, ImageModifier
 from ..image.extractor import ImageExtractor
 from ..io import DataLevel, DataWriter, EventSource, SimTelEventSource, write_table
 from ..io.datawriter import DATA_MODEL_VERSION
@@ -140,6 +140,7 @@ class ProcessorTool(Tool):
         + classes_with_traits(ImageExtractor)
         + classes_with_traits(GainSelector)
         + classes_with_traits(QualityQuery)
+        + classes_with_traits(ImageModifier)
         + classes_with_traits(EventTypeFilter)
     )