Code to load STIX pixel data into some arrays/dicts and then sum it into a spectrogram

stix_pixel_loader.py

from astropy.io import fits
import astropy.table as atab
import astropy.time as atime
import astropy.units as u
import numpy as np
from stixdcpy import instrument as inst


def load_stix_pixel_data_to_spectrogram(fn: str) -> dict:
    pdat = load_stix_pixel_data(fn)

    # sum out the detectors and pixels
    cts = pdat['counts'].sum(axis=(1, 2))
    cts_err = np.sqrt((pdat['counts_error']**2).sum(axis=(1, 2)))

    pdat.update({'counts': cts, 'counts_error': cts_err})
    return pdat


def load_stix_pixel_data(fn: str) -> dict:
    with fits.open(fn) as f:
        start_date = atime.Time(f['primary'].header['date-beg'])
        earth_timedelta = f['primary'].header['ear_tdel'] << u.s

        data_tab = atab.QTable.read(f, format='fits', hdu='data')

        time_bins = start_date + data_tab['time']
        dt = data_tab['timedel']
        time_bins = time_bins - dt/2
        time_bins = atime.Time(
            np.concatenate((
                time_bins, [time_bins[-1] + dt[-1]]
            ))
        )

        # index ordering: (time, detector, pixel, energy bin)
        counts = data_tab['counts']
        counts_cmp_err = data_tab['counts_comp_err']
        counts_err = np.sqrt(counts.to_value(u.ct) + counts_cmp_err.to_value(u.ct)**2) << u.ct

        triggers = data_tab['triggers']
        trig_cmp_err = data_tab['triggers_comp_err']
        trig_cmp_err.shape
        triggers_err = np.sqrt(triggers + trig_cmp_err**2)

        nom_lt = livetime_correct(triggers, counts, dt.to_value(u.s))
        up_lt = livetime_correct(triggers + triggers_err, counts, dt.to_value(u.s))
        lo_lt = livetime_correct(triggers - triggers_err, counts, dt.to_value(u.s))

        prop_live_err = np.abs(up_lt['live_ratio'] - lo_lt['live_ratio']) / 2 / nom_lt['live_ratio']

        lt_corr_counts = counts / nom_lt['live_ratio'][:, :, None, None]
        lt_corr_counts_err = np.sqrt(counts_err**2 + (prop_live_err[:, :, None, None] * counts)**2)

        energy_tab = atab.QTable.read(f, format='fits', hdu='energies')
        ebins = np.unique(
            np.column_stack((energy_tab['e_low'], energy_tab['e_high'])).flatten()
        )

        AVOID_DETECTORS = [9, 10]
        AVOID_MASK = np.array([i not in AVOID_DETECTORS for i in range(1, 33)])
        detector_mask = (data_tab['detector_masks'][0] & AVOID_MASK).astype(bool)

    return {
        'energy_bin_edges': ebins,
        'time_bin_edges': time_bins,
        'livetime': nom_lt['live_ratio'],
        'counts': lt_corr_counts[:, detector_mask, :, :],
        'counts_error': lt_corr_counts_err[:, detector_mask, :, :],
        'earth_spacecraft_dt': earth_timedelta
    }


# Taken directly from stixdcpy
BETA = 0.94
FPGA_TAU = 10.1e-6
ASIC_TAU = 2.63e-6
TRIG_TAU = FPGA_TAU + ASIC_TAU


def livetime_correct(triggers, counts_arr, time_bins):
    """ Live time correction
    Args
        triggers: ndarray
            triggers in the spectrogram
        counts_arr:ndarray
            counts in the spectrogram
        time_bins: ndarray
            time_bins in the spectrogram
    Returns
    live_time_ratio: ndarray
        live time ratio of detectors
    count_rate:
        corrected count rate
    photons_in:
        rate of photons illuminating the detector group
    """
    time_bins = time_bins[:, None]
    photons_in = triggers / (time_bins - TRIG_TAU * triggers)
    #photon rate calculated using triggers
    live_ratio = np.zeros((time_bins.size, 32))
    time_bins = time_bins[:, :, None, None]
    count_rate = counts_arr / time_bins
    # print(counts_arr.shape)
    for det in range(32):
        trig_idx = inst.detector_id_to_trigger_index(det)
        nin = photons_in[:, trig_idx]
        live_ratio[:, det] = np.exp(
            -BETA * nin * ASIC_TAU * 1e-6) / (1 + nin * TRIG_TAU)
    corrected_rate = count_rate / live_ratio[:, :, None, None]
    return {
        'corrected_rates': corrected_rate,
        'count_rate': count_rate,
        'photons_in': photons_in,
        'live_ratio': live_ratio
    }