Herschel-sed-fitter代码

翻到去年写的sed点到点拟合的代码。

######! /usr/bin python3
#This progress is wrote by Lixia Yuan
#Version v2.0 modified by Chao Zhang at Dec. 26, 2019
#Version v2.0.1 modified by Chao Zhang at Jun. 30, 2020
#This script requires dust_emissivity which can be installed by "pip install https://github.com/keflavich/dust_emissivity/archive/master.zip"
#This script requires higal_sedfitter which can be installed by "pip install https://github.com/keflavich/Hi-Gal_SEDfitter/archive/master.zip"
#This script requires fits_tools which can be installed by "pip install https://github.com/keflavich/FITS_tools/archive/master.zip"

import os
import sys
import numpy as np
from astropy.coordinates import SkyCoord
from scipy.optimize import curve_fit
import astropy.units as u
from astropy.io import fits
from astropy.convolution import convolve, Gaussian2DKernel, Gaussian1DKernel
import math
import astropy.constants as c
from dust_emissivity.dust import kappa
#from higal_sedfitter.higal_beams import beams
from astropy.wcs import WCS

rad = 1/np.pi*180*3600

#Crop image. This is the extra-module. 
def go_crop(source_coord=None,map=None,lambdas=None):
	'''
	Parameters
	---
	source_coord : 2-dimentional sky coordinates.
		>> source_coord = SkyCoord('04h21m10.5s','27d02m05.9s',frame='fk5')
	map : list or array
		The default map is [[-5,5],[-5,5]], where unit is arcmin. You can give map as [[map_RA1,map_RA2],[map_DEC1,map_DEC2]].
	lambdas : list or array
		The default lambdas is [160,250,350,500], where unit is micrometre.
	'''
	if lambdas is None:
		lambdas=[160,250,350,500]
	if map is None:
		map = np.array([[-5,5],[-5,5]])
	else:
		map = np.array(map)
	if source_coord is None:
		source_coord = SkyCoord('04h21m10.5s','27d02m05.9s',frame='fk5')
	for i in lambdas:
		crop_image(source_coord,str(i),map[0]*60,map[1]*60) 
		
#crop image function
def crop_image(source_coord,lamb,RA,DEC):
	if os.path.exists(lamb+'_old.fits') is True:
		os.remove(lamb+'.fits')
	if os.path.exists(lamb+'.fits') is True:
		f = fits.open(lamb+'.fits')
		f.writeto(lamb+'_old.fits')
	else:
		f = fits.open(lamb+'_old.fits')

	hdr = f[1].header
	w = WCS(hdr)
	pixel1,pixel2=w.wcs_world2pix(source_coord.ra.deg,source_coord.dec.deg,0)
	pixel1 = int(pixel1)+1
	pixel2 = int(pixel2)+1
	crdelt1 = abs(hdr['CDELT1'])*3600
	crdelt2 = abs(hdr['CDELT2'])*3600
	f[1].data = f[1].data[pixel2+int(DEC[0]/crdelt2)-1:pixel2+int(DEC[1]/crdelt2)+1,pixel1+int(RA[0]/crdelt1)-1:pixel1+int(RA[1]/crdelt1)+1]
	hdr.set('CRPIX1',1+abs(int(RA[0]/crdelt1)))
	hdr.set('CRPIX2',1+abs(int(DEC[0]/crdelt2)))
	hdr.set('NAXIS1',int((RA[1]-RA[0])/crdelt1)+2)
	hdr.set('NAXIS2',int((DEC[1]-DEC[0])/crdelt2)+2)
	hdr.set('CRVAL1',source_coord.ra.deg)
	hdr.set('CRVAL2',source_coord.dec.deg)
	hdu = fits.PrimaryHDU(data=f[1].data,header=hdr)

	if os.path.exists(lamb+'_crop.fits') is True:
		os.remove(lamb+'_crop.fits')
	hdu.writeto(lamb+'_crop.fits')
	f = fits.open(lamb+'_crop.fits')
	if os.path.exists(lamb+'.fits') is True:
		os.remove(lamb+'.fits')
	f.writeto(lamb+'.fits')
	os.remove(lamb+'_crop.fits')
#foreground/background removed, or namely large scale structure removed.
#undo
def go_removeback(box=None,lambdas=None):
	'''
	Statements: The foreground/background removed algorithm is according to Berry et al. 2007 (2007ASPC..376..425B).
	Parameters
	---
	box : source size in 250 um with 2-dimension list or array. The default box is 5 arcmin*5 arcmin.
	lambdas : list or array
		The default lambdas is [160,250,350,500], where unit is micrometre.
	'''
	if lambdas is None:
		lambdas=[160,250,350,500]

	if box is None:
		box = np.array([5,5])
	else:
		box = np.array(box)

	for i in lambdas:
		removed_background(box,str(i))

def removed_background(box,lamb):
	if os.path.exists(lamb+'_unremoveback.fits') is True:
		os.remove(lamb+'.fits')
		f = fits.open(lamb+'_unremoveback.fits')
		f.writeto(lamb+'.fits')
		f = fits.open(lamb+'.fits')
		os.remove(lamb+'_unremoveback.fits')
	else:
		f = fits.open(lamb+'.fits')

	f.writeto(lamb+'_unremoveback.fits')
	data = f[0].data
	hdr = f[0].header
	box = [int(box[0]/abs(hdr['CDELT1'])/60),int(box[1]/abs(hdr['CDELT2'])/60)]
	print(box)
	min_data=np.zeros((data.shape))
	for j in range(data.shape[0]):
		for k in range(data.shape[1]):
			if j-int(box[0]/2)-1<0:
				pixel_x_min=0
			else:
				pixel_x_min=j-int(box[0]/2)-1
			if j+int(box[0]/2)+1>data.shape[0]:
				pixel_x_max=data.shape[0]
			else:
				pixel_x_max=j+int(box[0]/2)+1
			if k-int(box[1]/2)-1<0:
				pixel_y_min=0
			else:
				pixel_y_min=k-int(box[1]/2)-1
			if k+int(box[1]/2)+1>data.shape[1]:
				pixel_y_max=data.shape[1]
			else:
				pixel_y_max=k+int(box[1]/2)+1

			min_data[j,k] = np.nanmin(data[pixel_x_min:pixel_x_max,pixel_y_min:pixel_y_max])

	max_data = np.zeros((data.shape))
	for j in range(data.shape[0]):
		for k in range(data.shape[1]):
			if j-int(box[0]/2)-1<0:
				pixel_x_min=0
			else:
				pixel_x_min=j-int(box[0]/2)-1
			if j+int(box[0]/2)+1>data.shape[0]:
				pixel_x_max=data.shape[0]
			else:
				pixel_x_max=j+int(box[0]/2)+1
			if k-int(box[1]/2)-1<0:
				pixel_y_min=0
			else:
				pixel_y_min=k-int(box[1]/2)-1
			if k+int(box[1]/2)+1>data.shape[1]:
				pixel_y_max=data.shape[1]
			else:
				pixel_y_max=k+int(box[1]/2)+1

			max_data[j,k] = np.nanmax(min_data[pixel_x_min:pixel_x_max,pixel_y_min:pixel_y_max])

	mean_data = np.zeros((data.shape))
	for j in range(data.shape[0]):
		for k in range(data.shape[1]):
			if j-int(box[0]/2)-1<0:
				pixel_x_min=0
			else:
				pixel_x_min=j-int(box[0]/2)-1
			if j+int(box[0]/2)+1>data.shape[0]:
				pixel_x_max=data.shape[0]
			else:
				pixel_x_max=j+int(box[0]/2)+1
			if k-int(box[1]/2)-1<0:
				pixel_y_min=0
			else:
				pixel_y_min=k-int(box[1]/2)-1
			if k+int(box[1]/2)+1>data.shape[1]:
				pixel_y_max=data.shape[1]
			else:
				pixel_y_max=k+int(box[1]/2)+1

			mean_data[j,k] = np.nanmean(max_data[pixel_x_min:pixel_x_max,pixel_y_min:pixel_y_max])

	new_data = data-mean_data
	hdu0 = fits.PrimaryHDU(data=new_data,header=hdr)
	hdu1 = fits.PrimaryHDU(data=mean_data,header=hdr)
	os.remove(lamb+'.fits')
	hdu0.writeto(lamb+'.fits')
	if os.path.exists(lamb+'_filter.fits'):
		os.remove(lamb+'_filter.fits')
	hdu1.writeto(lamb+'_filter.fits')



class SedFitter():
	'''
	Sed_Fitter can be used to fit SED by point to point and output the column density and dust temperature distribution images.
	Attention: The given fits file must be 'lambda.fits', such as '160.fits', '250.fits', '350.fits' and '500.fits'.
	Parameters
	----------
	lambdas : list or array
		The default lambdas is [160,250,350,500], where unit is micrometre.
	unit : str
		The Given unit is the converse_unit. Default is 'MJy/sr'. You can change the unit to 'Jy/pixel'.
	go_converse_unit : Default is True. If you do not need to converse unit, this should be change to None.
	go_convolve : Default is True. If do this, 'remove' must be True. If you do not need to convolve data, this should be change to None.
	go_sed : SED fitting process. Default is True.
	remove : Default is True. To delete the results previously output.
	'''
	def __init__(self, lambdas=None, unit=None, go_converse_unit=True, go_convolve=True, go_sed=True, remove=True):
		#modify the following parameters to aply your process
		pixels = [3.2, 6, 10,14]
		rms_sr = [0.0023/3.2/3.2/1e6*rad*rad, 5.4,3.4,1.9]
		rms_pixel = [0.0023,3.8*1e6/rad/rad*10*10,3.4*1e6/rad/rad*10*10,1.9*1e6/rad/rad*14*14]
		

		if lambdas is None:
			lambdas = [160,250,350,500]
			
		if unit is None:
			unit = 'MJy/sr'
			
		if remove is True:
			remove_fits_file(lambdas,go_converse_unit,go_convolve)
			
		fwhms = np.array([13.5,18.1,24.9,36.4])
		
		s = 0
		for i in lambdas:
			fwhm = fwhms[s]
			pixel = pixels[s]
			s = s+1
			if go_converse_unit is True:       #Do unit conversion proess
				hdu = fits.open(str(i)+'.fits')
				conv_u(hdu,unit,str(i),fwhm,pixel)
		
		s = 0
		for i in lambdas:
			fwhm = fwhms[s]
			pixel = pixels[s]
			s = s+1
			if go_convolve is True:            #Do convolve process
				convos(unit,str(i),fwhm)
				
		if go_sed is True:                    #Do SED fitting process
			im_160 = fits.open(str(lambdas[0])+'_convolve.fits')
			if im_160[0].header['BUNIT']=='MJy/sr':
				rms = rms_sr
			else:
				rms = rms_pixel
			sed_fitter(lambdas,rms)

		

#the unit conversion function
def conv_u(hdu,unit,lamb,fwhm,pixel):
	hdr=hdu[0].header
	if unit=='MJy/sr':
		if hdr['BUNIT']=='Jy/beam':
			conv_paras = rad**2/1e6/fwhm/fwhm
		
		elif hdr['BUNIT']=='Jy/pixel':
			conv_paras = rad**2/1e6/pixel/pixel
		
		elif hdr['BUNIT']=='MJy/sr':
			conv_paras = 1
		
		hdr.set('BUNIT','MJy/sr')
		hdr.set('QTTY____','MJy/sr')
		
	elif unit=='Jy/pixel':
		if hdr['BUNIT']=='Jy/beam':
			conv_paras = pixel*pixel/fwhm/fwhm
				
		elif hdr['BUNIT']=='Jy/pixel':
			conv_paras = 1
			
		elif hdr['BUNIT']=='MJy/sr':
			conv_paras = pixel*pixel*1e6/rad**2

			
		hdr.set('BUNIT','Jy/pixel')
		hdr.set('QTTY____','Jy/pixel')
	
	x,y = hdu[0].data.shape[0],hdu[0].data.shape[1]
	data_arr=np.zeros((x,y))
	for i in range(x):
		for j in range(y):
			data_arr[i,j]=hdu[0].data[i,j]*conv_paras
		
	hdu=fits.PrimaryHDU(data_arr,hdr)
	hdu.writeto(lamb+'_conv_unit.fits',overwrite=True)

#the convolve function
def convos(unit,lamb,fwhm):
	from FITS_tools.hcongrid import hcongrid
	hdu = fits.open(lamb+'_conv_unit.fits')
	if lamb!=str(500):
		gauss=np.sqrt(36.4**2-fwhm**2)
		pixel_size=gauss / 3600 / abs(hdu[0].header['CDELT2']) / np.sqrt(8*np.log(2))
		k1 = Gaussian1DKernel(pixel_size,mode='integrate').array    #It can be avoided that kernel become very large in the case gauss_pixel is much less than 1 when using gaussian2dkernel. 
		k2 = np.array([k1]*len(k1))
		gauss_kernel = k2*k2.T
		hdu[0].data=convolve(hdu[0].data,gauss_kernel,boundary='extend')
		target_header = fits.getheader('500_conv_unit.fits')
		new_data = hcongrid(hdu[0].data,hdu[0].header,target_header)
		hdu = fits.PrimaryHDU(new_data,target_header)
		hdu.writeto(lamb+'_convolve.fits', overwrite=True)
	else:
		hdu.writeto(lamb+'_convolve.fits', overwrite=True)


#define flux function, input are lamda(in micronmetre), T and N, output is spectral density(in MJy/sr)
def flux(lamb,T,N):
	freq=c.c.cgs/(lamb*u.cm*1e-4)
	kappa_v = kappa(lamb*u.um)
	tau = kappa_v*c.m_p.cgs*2.33*(N/u.cm**2)
	B_T = np.exp(c.h.cgs*freq/c.k_B.cgs/(T*u.K))-1
	B_T = 2*c.h.cgs*freq**3/B_T/c.c.cgs**2
	f = B_T*(1-np.exp(-tau))*1e17/u.erg*u.cm**2          #1 erg/cm2 = 1e17 MJy
	return f

def sed_fitter(lambdas,rms):
	im_160 = fits.open(str(lambdas[0])+'_convolve.fits')
	im_250 = fits.open(str(lambdas[1])+'_convolve.fits')
	im_350 = fits.open(str(lambdas[2])+'_convolve.fits')
	im_500 = fits.open(str(lambdas[3])+'_convolve.fits')
	pixel_x,pixel_y = im_160[0].shape[0],im_160[0].shape[1]
	
	T_data=np.zeros((pixel_x,pixel_y))
	N_data=np.zeros((pixel_x,pixel_y))
	T_hdr=im_160[0].header
	N_hdr=im_160[0].header
	T_hdr['BUNIT']='K'
	T_hdr['BTYPE']='Temperature'
	T_hdr['QTTY____']='K'
	N_hdr['BUNIT']='cm^-2'
	N_hdr['BTYPE']='Column Density'
	N_hdr['QTTY____']='cm^-2'
	err_160=13.1
	err_250=21.6
	err_350=18.5
	err_500=10.1
	for i in range(pixel_x):
		for j in range(pixel_y):
			xdata=[]
			ydata=[]
			errdata=[]
			#only use data that lagrer than a*rms, a can be specified
			if im_160[0].data[i,j]> 3*rms[0]: #3*sigma
				xdata.append(160)
				ydata.append(im_160[0].data[i,j])
				errdata.append(err_160)
			if im_250[0].data[i,j]> 3*rms[1]:
				xdata.append(250)
				ydata.append(im_250[0].data[i,j])
				errdata.append(err_250)
			if im_350[0].data[i,j]> 3*rms[2]:
				xdata.append(350)
				ydata.append(im_350[0].data[i,j])
				errdata.append(err_350)
			if im_500[0].data[i,j]> 3*rms[3]:
				xdata.append(500)
				ydata.append(im_500[0].data[i,j])
				errdata.append(err_500)
			if len(xdata)< 3:
				T_data[i,j]=np.nan
				N_data[i,j]=np.nan
			#fit the SED of available data
			else:
				#fit the SED, popt is the result and pcov is the error
				popt,pcov=curve_fit(flux, xdata, ydata, p0=[15,1e22],method='lm')
				#write T and N(in cm^-2) to data array
				if popt[0]>5: #fit the curve successfully
					T_data[i,j]=popt[0]
					N_data[i,j]=popt[1]
				else:
					T_data[i,j]=np.nan
					N_data[i,j]=np.nan

	a,b = int(N_data.shape[1]/2),int(N_data.shape[0]/2)
	import matplotlib.pyplot as plt
	x = np.arange(900)+101
	y = flux(x,T_data[a,b],N_data[a,b])
	plt.cla()
	plt.plot(x,y,c='k')
	x1 = [160,250,350,500]
	y1 = [im_160[0].data[a,b],im_250[0].data[a,b],im_350[0].data[a,b],im_500[0].data[a,b]]
	y_err = [rms[0],rms[1],rms[2],rms[3]]
	plt.errorbar(x1,y1,yerr=rms,color='blue',ecolor='red',fmt='o',capsize=4)
	plt.xlabel("$\\lambda$ ($\mu$m)")
	plt.ylabel("Flux (MJy)")
	plt.text(700,0.95*np.max(y1),"\ ")
	plt.text(700,0.85*np.max(y1),"T$_{dust}$=%.2f (K)" %(T_data[a,b]))
	plt.text(700,0.75*np.max(y1),"N=%.2e (cm$^{-2}$)" %(N_data[a,b]))
	plt.savefig('./sedfitting.pdf')

	hdu_T=fits.PrimaryHDU(T_data,T_hdr)
	hdu_N=fits.PrimaryHDU(N_data,N_hdr)
	hdu_T.writeto('T.fits',clobber=True)
	hdu_N.writeto('N.fits',clobber=True)

def remove_fits_file(lambdas,go_convolve,go_converse_unit):
	for i in lambdas:
		files = os.listdir(os.getcwd())
		for file in files:
			if file=='N.fits':
				os.remove(file)
				print(file+" deleted.")
			if file=='T.fits':
				os.remove(file)
				print(file+" deleted.")
			if go_convolve is True:
				if file==str(i)+'_convolve.fits':
					os.remove(file)
					print(file+" deleted.")
			if go_converse_unit is True:
				if file==str(i)+'_conv_unit.fits':
					os.remove(file)
					print(file+" deleted.")

def reload_hdu(lamb):
	files = os.listdir(os.getcwd())
	have_converse_unit = None
	have_convolve_data = None
	for file in files:
		if file==str(lamb)+'_convolve.fits':
			have_convolve_data = True
		elif file==str(lamb)+'_conv_unit.fits':
			have_converse_unit = True
	
	if have_convolve_data is True:
		hdu = fits.open(str(lamb)+'_convolve.fits')
	elif have_converse_unit is True:
		hdu = fits.open(str(lamb)+'_conv_unit.fits')
	else:
		for file in files:
			if file==str(lamb)+'.fits':
				hdu = fits.open(str(lamb)+'.fits')
				
	if hdu is None:
		raise ValueError("The name of fits file are not in regular. Please change the name to '160.fits', '250.fits', '350.fits' or '500.fits'.")
	else:
		return hdu