diffuse_filtering.py

#! /usr/bin/env python


"""
Python implementation of the diffuse image filtering technique described by Rudnick (2002):
https://ui.adsabs.harvard.edu/abs/2002NewAR..46..101R/abstract
"""

from argparse import ArgumentParser

from astropy.io import fits
import numpy as np
from scipy import ndimage


def get_args():
    ps = ArgumentParser()
    ps.add_argument("image")
    ps.add_argument("scale", type=float, help="Scale for diffuse filtering in arcsec.")
    args = ps.parse_args()
    return args


def strip_wcsaxes(hdr):
    """Strip extra axes in a header object."""

    remove_keys = [key+i for key in 
                   ["CRVAL", "CDELT", "CRPIX", "CTYPE", "CUNIT", "NAXIS"]
                   for i in ["3", "4", "5"]]

    for key in remove_keys:
        if key in hdr.keys():
            del hdr[key]

    return hdr


def get_last2d(array):
    """https://stackoverflow.com/a/27111239"""
    if array.ndim <= 2:
        return array
    else:
        slc = [0] * (array.ndim - 2)
        slc += [slice(None), slice(None)]
        return array[tuple(slc)]


def max_filter(array, size):
    output_array = ndimage.maximum_filter(array, size=size)
    return output_array
def min_filter(array, size):
    output_array = ndimage.minimum_filter(array, size=size)
    return output_array

            
def filter_image(image, scale):

    hdu = fits.open(image)
    pixel_scale = int((scale/3600.) / hdu[0].header["CDELT2"])
    print("Filter scale: {} pixels".format(pixel_scale))

    data = get_last2d(hdu[0].data)
    hdr = strip_wcsaxes(hdu[0].header)

    e_filtered = min_filter(data, pixel_scale)
    d_filtered = max_filter(e_filtered, pixel_scale)

    filtered_map = data - d_filtered
    fits.writeto(image.replace(".fits", "") + ".open.fits", d_filtered, hdr, overwrite=True)
    fits.writeto(image.replace(".fits", "") + ".filtered.fits", filtered_map, hdr, overwrite=True)


        
def cli(args):
    filter_image(
        image=args.image,
        scale=args.scale
    )

if __name__ == "__main__":
    cli(get_args())