从ALMA的data cube导出moment0

#MKmomentmap.py
#!/usr/bin/python
# -*- coding: utf-8 -*-
##  Zhang Chao 2020-11-27   initial version

from spectral_cube import SpectralCube
from astropy.io import fits
import astropy.units as u
import astropy.constants as c
import numpy as np
import os


def make_moment(fitsfile,window_type,window_range,order):
    '''
    input parameters:
    fitsfile: ALMA data cube file, i.e., fitsfile = "member_uid_*.fits" 
    window_type: A string. The window_type is only suported as 'velocity', 'frequency', and 'channel'. 

    output: The output hdu has a unit of 'K km s-1' or 'km s-1' for zero moment map or first moment map, respectively.
    '''

    f = fits.open(fitsfile)
    hdr = f[0].header
    if ('m' in hdr.cards['CUNIT3'][1]):
        cube = SpectralCube.read(fitsfile)
    else:
        hdu = change_xaxis_in_velocity(fitsfile)
        cube = SpectralCube.read(hdu)

    
    if window_type != 'velocity':
        window_range = change_to_velocity_range(hdu,window_type,window_range)

#----debug---------
#    print(window_range)
#    subcube = cube.subcube(xlo=140,xhi=235,ylo=165,yhi=255)
    slab = cube.spectral_slab(window_range[0]*u.km/u.s, window_range[1]*u.km/u.s)
    slab = slab.to(u.K) 
    if order==0:
        moment = slab.moment(order=order)
        moment = moment*1e-3
        hdr = moment.header
        hdr['BUNIT'] = 'K km s-1'
        hdu = fits.PrimaryHDU(moment.data,hdr)
#        if os.path.exists('%s_%s_m%i.fits' %(source,molecule,order)):
#            os.remove('%s_%s_m%i.fits' %(source,molecule,order))
#        moment.write('%s_%s_m%i.fits' %(source,molecule,order))
    elif order==1:

        m0 = slab.moment(order=0)
        moment_map = slab.moment(order=order)
        x,y =  moment_map.data.shape[0],moment_map.data.shape[1]
        vlsr = np.zeros((x,y))
        vlsr[vlsr==0]=(window_range[0]+window_range[1])/2

        moment_map.data = (moment_map.data-vlsr*1000)/1000
        mask = m0<10000*u.K*u.m/u.s
        moment_map[mask]=np.nan
    
        hdr = moment_map.header
        hdr['BUNIT'] = 'km s-1'
        hdu = fits.PrimaryHDU(moment_map.data,hdr)

    return (hdu)
#        if os.path.exists('%s_%s_m%i.fits' %(source,molecule,order)):
#            os.remove('%s_%s_m%i.fits' %(source,molecule,order))
#        hdu.write('%s_%s_m%i.fits' %(source,molecule,order))


def change_xaxis_in_velocity(fitsfile):
    f = fits.open(fitsfile)
    hdr = f[0].header
    delta_v = hdr['CDELT3']/hdr['CRVAL3']*c.c/1e3/u.m*u.s
    crvalue = hdr['NAXIS3']*delta_v.value
    hdr.set('CTYPE3','VELOCITY')
    hdr.set('CRVAL3',crvalue)
    hdr.set('CDELT3',-1*delta_v.value)
    hdr.set('CUNIT3', 'km/s')
    hdu = fits.PrimaryHDU(f[0].data,header=hdr)
    return (hdu)

def change_to_velocity_range(hdu,window_type,window_range):
    if window_type =='channel':
        window_range = [hdu.header['CRVAL3']+window_range[1]*hdu.header['CDELT3'],hdu.header['CRVAL3']+window_range[0]*hdu.header['CDELT3']]
    elif window_type == 'frequency':
        delta_frequency,reference_frequency = get_deltafrequency(hdu.header)

#--------debug------
#        print(delta_frequency,reference_frequency)
        window_range = [int((window_range[0]*1e6-reference_frequency)/delta_frequency),int((window_range[1]*1e6-reference_frequency)/delta_frequency)+1]
        window_range = [hdu.header['CRVAL3']+window_range[1]*hdu.header['CDELT3'],hdu.header['CRVAL3']+window_range[0]*hdu.header['CDELT3']]
    return (window_range)

def get_deltafrequency(hdr):
    if ('FILNAM01' and 'FILNAM03' and 'FILNAM04' and 'FILNAM05' and 'FILNAM06' and 'FILNAM09' in hdr)==False:
        print("This file is only used for ATOMS data reduction.")
        delta_frequency = None
        reference_frequency = None
    elif os.path.exists('member.'+hdr['FILNAM01']+'.'+hdr['FILNAM03']+'.'+hdr['FILNAM04']+'.'+hdr['FILNAM05']+'.'+hdr['FILNAM06']+'.'+hdr['FILNAM09']+'.fits'):
        original_file_header = fits.open('member.'+hdr['FILNAM01']+'.'+hdr['FILNAM03']+'.'+hdr['FILNAM04']+'.'+hdr['FILNAM05']+'.'+hdr['FILNAM06']+'.'+hdr['FILNAM09']+'.fits')[0].header
        delta_frequency = original_file_header['CDELT3']
        reference_frequency = original_file_header['CRVAL3']
    else:
        print("Please give the path of the original cube file.")
        delta_frequency = None
        reference_frequency = None
    return (delta_frequency,reference_frequency)