Possible route to investigate: Pipeline for creating out-of-memory AIA data cubes with NDCube and Dask

[wtbarnes]

So this is very late in the game I realize to be adding another project idea, but we have quite a few people on the team so I figured it couldn't hurt. This is a bit of a mess and I will try to come back later and clean it up and maybe some more concrete steps will become more clear once we talk things through a bit.

One very common task that people want to do with AIA images is to analyze them in time by stacking many sequential images. Few people care about looking at just 1 AIA image. Especially at the full 4K resolution, this can be extremely challenging, computationally and just logistically, for a significant chunk of time. Even for smaller cutouts this can be an issue as accessing and aligning these cutouts is often difficult.

This typical workflow would look like,

1. Query and download AIA data over some interval and for a selected number of wavelengths.
   (optionally also deconvolve with PSF though this isn't often done)
2. Apply prep operations to every image such that each wavelength, at a given timestep, is on the same pixel grid
3. Pick a reference time and align all images to this reference time to remove the effect of the Sun's rotation such that a given pixel at each time corresponds to the same physical location on the Sun. Alternatively, use some type of cross-correlation technique to align the images
4. Stack the aligned images into a cube of dimensions nx-by-ny-by-nt, where nt corresponds to the number of images in the selected interval. There will be as many separate cubes as there are wavelengths.
5. Perform scientific analysis on these aligned data cubes (e.g. extracting light curves from specific regions, computing time lags in each pixel, computing DEMs in each pixel as a function of time)

The steps through step 3 all have to occur on a by-image basis.

In an ideal world, I would be able to pass in a stack of AIA FITS files and process it, in parallel, into an out-of-core aligned AIA data cube.

I'm not exactly sure how this would fit into the aiapy package, but I think showing a prototype of this kind of workflow, especially accelerated with things like numba or cuda, would be really beneficial to the community.

I see this less about creating any custom tools for AIA analysis, but more about demonstrating how fast and powerful this workflow can be in Python.

Tools we would probably need:

• ndcube
• sunpy
• reproject (for image alignment in time)
• Dask

Main challenges:

• Parallelizing the prep step (we are already planning to address this)
• Loading the FITS files lazily
• Transitioning from per-file operations to operations on the whole cube
• Making these workflows the same, regardless of whether Dask is being used or not, i.e. all of this should be able to be done in memory and in serial too

As a possible starting point, see this notebook: https://gitlab.com/wtbarnes/aia-on-pleiades/-/blob/master/notebooks/tidy/data-cubes.ipynb

[cbard]

This is a really good idea! I was just thinking about how our proposed acceleration pathways seem slightly wasted on individual images, so this fits well with trying to expand the reduction pipeline to multiple images.

For step 3 (removing Sun's rotation), what would the data array look like? Would data[i,j,k,t1] be the same solar lat. and long. as data[i,j,k,t2]? Or would it be the same relative location within the aligned and centered solar disk? The latter would be easy to stack, but the former would have to be careful about points appearing/disappearing off the edge. This might require a creative way of storing the data such that we don't have unnecessary NULL storage for those points not on disk over the full time range. Of course, we could also make the assumption that the cutouts are fully on-disk over the time range, but we would need a separate function to verify this assumption.

[wtbarnes]

For step 3 (removing Sun's rotation), what would the data array look like? Would data[i,j,k,t1] be the same solar lat. and long. as data[i,j,k,t2]?

Yes that's correct. So for example if I want to look at how a particular bit of plasma varies as a function of time, I can take a slice in time at a particular pixel (i,j) (I'm assuming k in your example is wavelength?) by doing data[i,j,:].

You're right about points disappearing and reappearing. Usually one takes care to choose observations where the active region (or whatever ROI you have) doesn't move too far from disk center such that no points are lost in the part of the image that you actually care about. This also means of course you can only choose a time interval that's less than half a solar rotation as well, but in general this is not a huge problem. For my research problems, I'm usually interested in timescales significantly less than that.

[wtbarnes]

To add to this, I'm also planning on investigating whether cupy can be used to speedup the calculation of time lags over AIA data cubes.

Currently, I've done this with Dask (see https://gitlab.com/wtbarnes/aiacube/-/blob/master/aiacube/time_lag.py) but I'm interested to see whether it can be done in pure cupy or cupy + Dask.

[wtbarnes]

Potentially plan: Benchmark the following three cases for a time lag calculation

• Pure Numpy
• Dask
• Cupy
• Dask + Cupy

There is a question of whether option 1 and 3 will even work because of the memory requirements.

[cbard]

The GPUs on the NVIDIA cluster have 16 GB on board memory, is your stack larger than that? Alternatively, it sounds like splitting up the data using Dask may be the way to go.

[wtbarnes]

For "realistic" workflows, it is bigger. I may have to experiment with smaller stacks to begin with. I have access to a multi-node GPU cluster so I may experiment there with some larger stacks as well.