from copy import deepcopy
from typing import List, Tuple
import astropy.units as u
import numpy as np
from astropy.table import Table
from exosim.models.signal import Signal
from exosim.output import Output
from exosim.tasks.task import Task
from exosim.utils import RunConfig
from exosim.utils.checks import check_units
from exosim.utils.iterators import iterate_over_chunks
from exosim.utils.types import ArrayType
[docs]class AddCosmicRays(Task):
"""
Task to simulate the addition of cosmic rays to sub-exposures in a detector.
This task models the impact of cosmic rays on a detector during the exposure time. The model assumes
that the cosmic rays can interact with the detector pixels in various predefined shapes like a cross,
horizontal rectangle, and vertical rectangle. The number of cosmic ray events and their impact on the
detector pixels are calculated based on the given cosmic ray flux, detector characteristics, and
integration times for the sub-exposures.
Notes
-----
- This task assumes that the cosmic ray events saturate the affected pixels, setting their value to the full well depth.
- The probabilities for the interaction shapes are configurable. The task issues a warning if the sum of provided probabilities is not 1.
- The cosmic ray flux is specified in ct/s/cm^2 and is scaled based on the pixel size and detector dimensions.
This is a default class with standardised inputs and outputs.
The user can load this class and overwrite the "model" method
to implement a custom Task to replace this.
"""
def __init__(self):
"""
Parameters
----------
subexposures: :class:`~exosim.models.signal.Counts`
sub-exposures cached signal
parameters: dict
channel parameters dictionary
integration_times: :class:`~astropy.units.Quantity`
sub-exposures integration times
outputs: :class:`~exosim.output.output.Output` (optional)
output file
"""
self.add_task_param("subexposures", "sub-exposures cached signal")
self.add_task_param("parameters", "channel parameters dictionary")
self.add_task_param(
"integration_times",
"subexposures integration times",
)
self.add_task_param("output", "output file", None)
[docs] def execute(self):
self.info("adding cosmic rays")
subexposures = self.get_task_param("subexposures")
parameters = self.get_task_param("parameters")
integration_times = self.get_task_param("integration_times")
output = self.get_task_param("output")
self.model(subexposures, parameters, integration_times, output)
[docs] def model(
self,
subexposures: Signal,
parameters: dict,
integration_times: ArrayType,
output=None,
) -> None:
"""
Default model to simulate the addition of cosmic rays to the sub-exposures.
This method saturates the hit pixels in the sub-exposure data based on the cosmic ray rate,
detector properties, and the given integration times. The cosmic ray interactions could be in various shapes
like single pixel, lines, squares, etc., which are defined in the configuration.
Parameters
----------
subexposures : Signal
The sub-exposures' cached signal.
parameters : dict
The channel parameters dictionary containing detector information.
integration_times : ArrayType
The integration times for the sub-exposures.
output : Output, optional
The output file where the cosmic ray events will be stored.
Notes
-----
The method assumes multiple possible shapes for the interaction:
- Single pixel
- Vertical line: Saturates two pixels vertically aligned.
- Horizontal line: Saturates two pixels horizontally aligned.
- Square: Saturates four pixels in a square.
- Cross: Saturates five pixels in a cross shape.
- Horizontal rectangle: Saturates six pixels in a horizontal rectangle shape.
- Vertical rectangle: Saturates six pixels in a vertical rectangle shape.
"""
# Extract relevant parameters from the provided dictionary
rate = parameters["detector"]["cosmic_rays_rate"].astype(np.float64)
rate = check_units(rate, "ct/s/cm^2")
saturation_rate = parameters["detector"]["saturation_rate"]
spatial_pix = parameters["detector"]["spatial_pix"]
spectral_pix = parameters["detector"]["spectral_pix"]
well_depth = parameters["detector"]["well_depth"]
well_depth = check_units(well_depth, "ct", force=True).value
pixel_size = parameters["detector"]["delta_pix"].astype(np.float64)
pixel_size = check_units(pixel_size, "cm", force=False)
randomise = parameters["detector"].get("cosmic_rays_randomise", False)
# Calculate the number of events to be added to each sub-exposure
events_counter = self.count_events(
rate,
pixel_size,
spectral_pix,
spatial_pix,
integration_times,
saturation_rate,
randomise,
)
self.info(f"cosmic events_counter: {events_counter}")
# Get interaction shapes and their probabilities
shapes, probs = self.shapes_and_probs(parameters["detector"])
# Initialize a table to store information about the cosmic ray events
storing_info = Table(
names=(
"Sub-exposure",
"spectral_pix",
"spatial_pix",
"saturated_pix_rand",
"saturated_pix_seed",
"saturated_pix",
),
dtype=("i4", "i4", "i4", "f8", "i4", "U16"),
)
# Iterate over the sub-exposures and add cosmic ray events
for chunk in iterate_over_chunks(
subexposures.dataset, desc="adding cosmic rays"
):
data = deepcopy(subexposures.dataset[chunk])
events_counter_chunk = events_counter[chunk[0]]
for t in range(data.shape[0]):
# Randomly select pixels for the cosmic ray events
x_pix = RunConfig.random_generator.choice(
spectral_pix, events_counter_chunk[t]
)
y_pix = RunConfig.random_generator.choice(
spatial_pix, events_counter_chunk[t]
)
for x, y in zip(x_pix, y_pix):
info_row = [chunk[0].start + t, x, y]
# Randomly select a shape for the cosmic ray event
rand = RunConfig.random_generator.uniform(0, 1)
info_row.append(rand)
info_row.append(RunConfig.random_seed)
r_ = 0
for s, r in zip(shapes, probs):
r_ += r
if rand <= r_:
shape = s
break
shape_str = ""
for i, j in zip(shape[0], shape[1]):
shape_str += f"({i + x}, {j + y})"
try:
data[t, j + y, i + x] = well_depth
except IndexError:
continue
info_row.append(shape_str)
storing_info.add_row(info_row)
# Update the sub-exposure data
subexposures.dataset[chunk] = data
subexposures.output.flush()
# Write the cosmic ray events information to the output file if provided
if output and issubclass(output.__class__, Output):
out_grp = output.create_group("cosmic rays")
out_grp.write_table("cosmic rays events", storing_info)
out_grp.write_array("cosmic events counter", events_counter)
if randomise:
out_grp.write_scalar("random seed", RunConfig.random_seed)
[docs] def count_events(
self,
rate: float,
pixel_size: u.Quantity,
spatial_pix: int,
spectral_pix: int,
integration_times: ArrayType,
saturation_rate: float,
randomise: bool,
) -> np.ndarray:
"""
Calculate the number of cosmic ray events in each sub-exposure, randomly distributing them
while keeping the probability proportional to the exposure time.
Parameters
----------
rate : float
Cosmic rays flux rate in ct/s/cm^2.
pixel_size : float
Size of a detector pixel in cm^2.
spatial_pix : int
Number of spatial pixels in the detector.
spectral_pix : int
Number of spectral pixels in the detector.
integration_times : :class:`~astropy.units.Quantity`
Sub-exposure integration times.
saturation_rate : float
Saturation rate for the detector.
randomise : bool
A flag to randomize the total number of events
Returns
-------
events_counter_i : np.ndarray
An array containing the number of events for each sub-exposure, rounded to the nearest integer.
"""
# Scale the rate based on detector characteristics
scaled_rate = (
rate * pixel_size * pixel_size * spatial_pix * spectral_pix
)
scaled_rate = check_units(scaled_rate, "ct/s")
# Log the calculated rate
self.info(f"cosmic rays rate: {scaled_rate}")
# Compute the total expected number of events
total_integration_time = np.sum(integration_times)
self.debug(f"integration tiems: {integration_times}")
total_events = (
scaled_rate * total_integration_time * saturation_rate
).value
if randomise:
self.debug("Randomizing total number of events.")
# Randomize total event
total_events = RunConfig.random_generator.poisson(total_events)
self.info(f"cosmic total events: {total_events}")
# Normalize integration times to create a probability distribution
integration_weights = integration_times / total_integration_time
probabilities = integration_weights # Convert to simple float array
# Generate a random multinomial distribution of events across sub-exposures
events_counter = np.random.multinomial(
int(total_events), probabilities
)
self.debug(f"Probabilities: {probabilities}")
self.debug(f"Probabilities sum: {np.sum(probabilities)}")
self.debug(f"Events counter before multinomial: {int(total_events)}")
return events_counter
[docs] def shapes_and_probs(self, parameters: dict) -> Tuple[List, List]:
"""
Generate cosmic ray interaction shapes and their corresponding probabilities.
Parameters
----------
parameters : dict
A dictionary containing optional 'interaction_shapes' which is a sub-dictionary
specifying shapes and their probabilities.
Returns
-------
shapes : list of tuples
A list of tuples representing shapes of cosmic ray interactions.
probs : list of float
A list of probabilities corresponding to the shapes.
Raises
------
ValueError
If the sum of provided probabilities for all shapes is greater than 1.
Examples
--------
>>> shapes, probs = AddCosmicRays.shapes_and_probs({"interaction_shapes": {"single": 0.5, "line_h": 0.3, "line_v": 0.2}})
>>> print(shapes)
[([0], [0]), ([0, 1], [0, 0]), ([0, 0], [0, 1]), ([0, 0, 1, 1], [0, 1, 0, 1]), ([0, 0, 0, -1, 1], [0, 1, -1, 0, 0])]
>>> print(probs)
[0.5, 0.3, 0.2, 0.1, 0.1]
>>> shapes, probs = AddCosmicRays.shapes_and_probs({"interaction_shapes": {"single": 0.5, "line_h": 0.3, "line_v": 0.2}})
>>> print(shapes)
Warnings
--------
A warning is issued if the sum of provided probabilities for all shapes is not 1.
"""
# Predefined cosmic ray interaction shapes
predefined_shapes = {
"single": ([0], [0]),
"line_h": ([0, 1], [0, 0]),
"line_v": ([0, 0], [0, 1]),
"square": ([0, 0, 1, 1], [0, 1, 0, 1]),
"cross": ([0, 0, 0, -1, 1], [0, 1, -1, 0, 0]),
"rect_h": ([0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1]),
"rect_v": ([0, 0, 0, 1, 1, 1], [0, 1, 2, 0, 1, 2]),
}
# Initialize lists to store shapes and probabilities
shapes, probs = [], []
# Check if interaction shapes are defined in the parameters
if "interaction_shapes" in parameters:
for shape_key, prob in parameters["interaction_shapes"].items():
if shape_key in predefined_shapes:
shapes.append(predefined_shapes[shape_key])
probs.append(prob)
# Check if probabilities sum up to 1
remaining_prob = 1.0 - sum(probs)
if remaining_prob < 0:
self.error("Total probability greater than 1.")
raise ValueError("Total probability greater than 1.")
# Adjust the probability for the 'single' shape if needed
if remaining_prob != parameters["interaction_shapes"].get(
"single", 0
):
self.warning(
"Total probability of all interaction shapes should sum to 1. "
"Adjusting 'single' interaction shape to make the sum equal to 1."
)
shapes.append(predefined_shapes["single"])
probs.append(remaining_prob)
return shapes, probs
# Default to 'single' shape with probability 1 if no shapes are defined
return [predefined_shapes["single"]], [1.0]
