Make an AIA movie out of a bunch of FITS files using sunpy

make_aia_movie.py

from datetime import datetime
import functools
import multiprocessing as mp
import os
import pathlib
import sys

import imageio
import numpy as np 
import sunpy.map
import matplotlib.pyplot as plt
# plt.style.use(os.getenv('MPL_INTERACTIVE_STYLE'))

'''
Call `default_movie` with a directory full of FITS files to make an AIA movie.

You can customize how it is made with the more granular functions `make_images`
and `combine_images`.

W Setterberg 2023
'''

def default_movie(in_dir: str) -> None:
    out_dir = f'{in_dir}/aia-movie'

    fits_in = [f'{in_dir}/{fn}' for fn in os.listdir(in_dir)]
    ncpu = mp.cpu_count()
    split_files = np.array_split(fits_in, ncpu)
    to_map = functools.partial(make_images, out_dir=out_dir)

    with mp.Pool(ncpu) as p:
        files = p.map(to_map, split_files)

    files = [e for sublist in files for e in sublist]
    gif_out = out_dir + '/aia-movie.gif'
    combine_images(files, gif_out)


def make_images(fns: list[str], out_dir: str) -> list[str]:
    os.makedirs(out_dir, exist_ok=True)
    out_path = pathlib.Path(out_dir)
    files_out = []
    fig = plt.figure(figsize=(8, 8), layout='constrained')
    for fn in fns:
        fig.clear()
        if not fn.endswith('.fits'): continue
        print('start img', fn)
        p = pathlib.Path(fn)
        m = sunpy.map.Map(p)
        m.plot()

        outf = (out_path / p.stem).with_suffix('.jpg')
        files_out.append(outf)
        fig.savefig(fname=outf, facecolor='white', dpi=100)
        print('done img', fn)

    return [str(f) for f in files_out]


def sort_image_fns(fns: list[str]) -> list[str]:
    ''' assumes default AIA file naming '''
    def key(s: str) -> datetime:
        stem = str(pathlib.Path(s).stem)
        utc_str = stem.split('_')[1]
        return datetime.fromisoformat(utc_str)

    return list(sorted(fns, key=key))


def combine_images(image_fns: list[str], out_fn: str, frame_duration: float=0.5) -> None:
    sorted_fns = sort_image_fns(image_fns)
    with imageio.get_writer(out_fn, mode='I', duration=frame_duration) as wr:
        for fn in sorted_fns:
            print('start vid', fn)
            img = imageio.imread(fn)
            wr.append_data(img)
            print('end vid', fn)


if __name__ == '__main__':
    default_movie(sys.argv[1])