Even on powerful servers, the COSI Data Challenge imaging computations initially took 8+ hours to fully run. This slow processing would limit the user’s ability to effectively play around with and analyze the data.

The primary bottleneck in the pipeline was the Richardson Lucy Image Deconvolution Algorithm. 3 of the main issues were:

1. Single-threaded convolution kernels (140s per iteration)

2. Inefficient background fitting (60s per iteration)

3. Excessive memory usage (70-100 GB for response matrices)

My Role: To optimize Python code while keeping it accessible to researchers.

Investigated various optimization methods, including JIT compilation, parallelization, and memory reduction techniques.

Tested each solution on three different imaging tasks, which were the image reconstructions for 511keV, Al26, and Point Sources. I measured speedups while ensuring quality image output, post deconvolution. I also verified that the results from my accelerated code matched the original algorithm output.

Achieved 60.1x average speedup across all tasks through my acceleration strategies.

JAXopt also provided 7-14x further improvement through better convergence, and I achieved a 50% reduction in memory usage.

I compiled my findings into a scientific poster which I presented as first author at the 2024 High Energy Astrophysics Division (HEAD) conference in Texas. I also met with with the original COSI Data Challenge team to discuss implementing my optimizations into upcoming data challenge releases.