Converts Pixie-16 list-mode data files to a variety of other filetypes

pixie16_list_mode_converter.py

""" SPDX-License-Identifier: Apache-2.0 """
import argparse
import json
from io import BytesIO
import logging
import os
import sys

from dolosse.hardware.xia.pixie16.list_mode_data_mask import ListModeDataMask
from dolosse.hardware.xia.pixie16.list_mode_data_decoder import decode_listmode_data
import pandas as pd
import matplotlib.pyplot as plt

"""
Copyright 2022 XIA LLC, All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

""" lmd_analysis.py

"""

logging.basicConfig(stream=sys.stdout, datefmt="%Y-%m-%dT%H:%M:%S.000Z", level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')


def convert_to_jsonl(args, filename):
    out = open(filename, "w")
    for file in args.file:
        logging.info(f"Processing {file}")

        if os.path.exists(file):
            if os.path.getsize(file):
                logging.info(
                    f"Sending {file} with size of {os.path.getsize(file)} bytes for processing.")
                try:
                    logging.info(f"Started to read {file}")
                    raw = open(file, 'rb').read()
                    logging.info("Finished reading file.")
                    logging.info(f"Started to decode data")
                    results = decode_listmode_data(BytesIO(raw),
                                                   ListModeDataMask(args.freq, args.rev))
                    logging.info(f"Finished decoding data")
                    while results:
                        out.write(json.dumps(results.pop(0)))
                        out.write("\n")
                except Exception as ex:
                    logging.error(ex)
                    continue
            else:
                logging.warning(f'{file} is empty!!')
                continue
        else:
            logging.warning(f'{file} does not exist!!')
            continue
    out.close()


def analyze_files(filename):
    if not os.path.exists(filename):
        raise RuntimeError(f"{filename} does not exist!!")
    df = pd.read_json(filename, lines=True)
    logging.info(f'Total events processed: {len(df)}')
    df.to_parquet(f'{os.path.splitext(filename)[0]}.parquet', compression='snappy')
    df.to_csv(f'{os.path.splitext(filename)[0]}.csv')

    df.hist(column='energy', by='channel', bins=51)
    plt.suptitle("List-Mode Energy Histograms")
    plt.show()

    df.hist(column='event_length', by='channel', bins=51)
    plt.show()

    for chan in df['channel'].unique():
        # timediff = df[df['channel'] == chan].sort_values(
        #     'event_time_low').rolling(window=2).apply(lambda x: x.iloc[1] - x.iloc[0])
        # timediff.hist(column=['event_time_low'])
        # plt.suptitle(f"Channel {chan} diffs:")
        # plt.show()
        logging.info(
            f"Total events for channel {chan}: {df[df['channel'] == chan]['energy'].count()}")


def lmd_analysis(args):
    base_filename, ext = os.path.splitext(args.file[0])

    if len(args.file) > 1 and ext in ('.jsonl' or '.parquet'):
        raise RuntimeError("We will not aggregate jsonl files at this time.")

    if len(args.file) == 1:
        filename = f'{base_filename}.jsonl'
    else:
        filename = f'{base_filename}-agg.jsonl'

    if len(args.file) > 1 or ext == '.bin':
        convert_to_jsonl(args, filename)

    analyze_files(filename)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Test to generate a full suite of List-mode data.')
    parser.add_argument('-f', '--file', type=str, nargs='+', dest="file",
                        required=True, help="The list of filenames that we want to analyze.")
    parser.add_argument('--freq', type=int, dest='freq', required=True,
                        help="The sampling frequency used to collect list-mode data. Ex. 250")
    parser.add_argument('--rev', type=int, dest='rev', required=True,
                        help="The firmware used to collect list-mode data. Ex. 30474")

    try:
        lmd_analysis(parser.parse_args())
    except KeyboardInterrupt:
        logging.info("Received keyboard interrupt. Stopping execution.")
    except Exception as e:
        logging.error(repr(e))