"""
Copyright (C) 2018, Riccardo Pavesi
E-mail: pavesiriccardo3 -at- gmail.com
Updated versions of the software are available through github:
https://github.com/pavesiriccardo/uvmodeldisk
If you have found this software useful for your research,
I would appreciate an acknowledgment to the use of the
"Disk modeling in the UV plane (UVmodeldisk) routines of Pavesi et al., (2018a)".
[https://arxiv.org/abs/1803.08048]
This software is provided as is without any warranty whatsoever.
For details of permissions granted please see LICENCE.md
"""
import numpy as np
from KinMS import KinMS
import uvutil
from astropy.io import fits
import pymultinest,emcee
from galario.double import sampleImage
class uvmodeldisk(object):
def __init__(self,xs,ys,vs,cellsize,dv,UVFITSfile,cube_template_file):
"""
This function initializes the modeling object.
Parameters
----------
xs: float
X-axis size of the model image (in arcsec)
ys: float
Y-axis size of the model image (in arcsec)
vs: float
Velocity-axis size of the model image (in km/s)
cellsize: float
Arcseconds in a pixel
dv: float
Km/s in a channel
UVFITSfile: string
Filename of the UVFITS visibility file
cube_template_file: string
Filename of the image cube to be used as a model template (the dimensions and coordinates are the only thing that's used).
It's best obtained by imaging the visibilities themselves, using mode='channel'. This ensures that the cell size is sufficient to suitable for the UV coverage.
Returns
-------
"""
self.xs=xs
self.ys=ys #y-axis size of the model image (in arcsec)
self.vs=vs
self.cellsize=cellsize
self.dv=dv
self.UVFITSfile=UVFITSfile #'.uvfits'
self.cube_template_file=cube_template_file #'.fits'
self.sbrad=np.arange(0,2,.01) #radius vector defining the light profile, sbprof is the light intensity evaluated at these radial values
self.velrad=self.sbrad #use the same radius vector for the velocity definition, used by velprof to define the velocity profile
self.nsamps=1e5 #number of clouds, should be sufficient
self.vis_complex_data, self.wgt_data = uvutil.visload(self.UVFITSfile) #load visibility and weights
self.vis_complex_model_template=np.copy(self.vis_complex_data)
self.wgt_data=self.wgt_data.flatten()
self.good_vis=self.wgt_data>0
self.wgt_data=self.wgt_data[self.good_vis] #get rid of flagged data
self.vis_complex_data=self.vis_complex_data.flatten()[self.good_vis]
self.modelheader = fits.getheader(self.cube_template_file)
self.uu_cube, self.vv_cube, self.ww_cube = uvutil.uvload(self.UVFITSfile)
self.pcd = uvutil.pcdload(self.UVFITSfile)
self.Nxpix=self.modelheader['NAXIS1']
self.Nypix=self.modelheader['NAXIS2']
self.Nchan=self.modelheader['NAXIS3']
self.Nxpix_small=int(self.xs/self.cellsize)
self.Nypix_small=int(self.ys/self.cellsize)
self.Nchan_small=int(self.vs/self.dv)
#In this code we assume the continuum emission can be modelled as a 2D Gaussian, with known parameters. This is often the case and the continuum parameters are often precisely known. If they are not, include the continuum parameters as fitting parameters and make sure the visibilities cover a good range of continuum-only channels.
self.y,self.x,self.zz=np.mgrid[:self.Nypix_small,:self.Nxpix_small,:self.Nchan_small] #define a coordinate cube of the same size as the KinMS model, to make the continuum model, to be added in.
self.offset_lnlike=0
self.xpos_center_padded,self.ypos_center_padded= int(self.Nxpix/2),int(self.Nypix/2)
def set_ellipt_gauss_continuum(self,cont_xcen,cont_ycen,cont_maj_FWHM,cont_min_FWHM,cont_posang,cont_flux):
"""
This function defines the continuum model.
Parameters
----------
cont_xcen,cont_ycen: (float,float)
X-axis and Y-axis offset of continuum center from image center (in arcsec)
cont_maj_FWHM,cont_min_FWHM: (float,float)
FWHM of the major and minor axes of the continuum 2D Gaussian (in arcsec)
cont_posang: float
Position angle of the ellipse, counterclockwise from the positive x axis (in degrees)
cont_flux: float
Integrated flux of the continuum emission (in mJy)
Returns
-------
"""
xcen=self.x-cont_xcen/self.cellsize-self.Nxpix_small/2.
ycen=self.y-cont_ycen/self.cellsize-self.Nypix_small/2.
cc=np.cos(cont_posang/180*np.pi)
ss=np.sin(cont_posang/180*np.pi)
r2=(xcen*cc+ycen*ss)**2/(cont_maj_FWHM/self.cellsize/2.355)**2+(xcen*ss-ycen*cc)**2/(cont_min_FWHM/self.cellsize/2.355)**2
self.model_cont=cont_flux*1e-3*np.exp(-r2/2)/(2*np.pi*(cont_maj_FWHM/self.cellsize/2.355)*(cont_min_FWHM/self.cellsize/2.355))
def loglike_multinest(self,cube,ndim,npar):
return self.loglike(cube)
def my_prior(self,cube):
return 0
def loglike(self,cube):
'''
This function calculates the model ln(likelihood).
Parameters
----------
gassigma: float
The gas dispersion in km/s
bright_std: float
The size as gaussian FWHM in arcsec
vmax: float
The max velocity in km/s, the asymptote of the arctan
vel_scale: float
The radial distance, in arcsec, where velocity goes to vmax/2
vel_cen: float
The velocity center in km/s away from central channel, increasing vel_cen moves it to later channels, i.e. same direction as the channel ordering
inc: float
Inclination, is 0 for face on and 90 for edge-on
posang: float
Position angle starting with red emission from horizontal to the right (toward decreasing RA), and increasing counterclockwise (when we have a - in front of vmax). posang near 0, and - in front of vmax, means the emission moves right to left as channels increase. positive posang rotates this pattern toward the north, in a counterclockwise manner
x_cen, y_cen: (float,float)
Center position for the disk, they are in arcsec. y_cen actually controls the x-axis and positive means increasing x of center. x_cen controls y axis, and positive means increasing y of center
intflux: float
Line integrated flux in Jy km/s
Returns
-------
ln_like+self.offset_lnlike: float
We offset the ln_like value by a constant to make the number be small enough for Multinest to be able to work with.
Definition of posang for the rotating disk:
If velprof=vmax*np.arctan(velrad/vel_scale)/np.pi*2
posang 0: to the right
posang 90: downwards
posang 180: to the left
posang 270 (=-90): upward
If velprof=-vmax*np.arctan(velrad/vel_scale)/np.pi*2
posang 0: to the left
posang 90: upward
posang 180: to the right
posang 270 (=-90): downward
For example: if the emission is moving from left to right, as channels increase (toward lower frequency and higher velocity, red). Then need posang=180 if minus sign in front of vmax.
'''
gassigma,bright_std,vmax,vel_scale,vel_cen,inc,posang,x_cen,y_cen,intflux=cube[0],cube[1],cube[2],cube[3],cube[4],cube[5],cube[6],cube[7],cube[8],cube[9]
model_cont=self.model_cont
sbprof=np.exp(-self.sbrad**2/2/(bright_std/2.355)**2)
velprof=vmax*np.arctan(self.velrad/vel_scale)/np.pi*2
self.modelimage=KinMS(self.xs,
self.ys,
self.vs,
cellSize=self.cellsize,
dv=self.dv,
beamSize=0,
inc=inc,
gasSigma=gassigma,
sbProf=sbprof,
sbRad=self.sbrad,
velRad=self.velrad,
velProf=velprof,
diskThick=0,
cleanOut=True,
ra=0,
dec=0,
nSamps=self.nsamps,
posAng=posang,
intFlux=intflux,
inClouds=[],
vLOS_clouds=[],
flux_clouds=0,
vSys=0,
restFreq=115.271e9,
phaseCen=np.array([x_cen,y_cen]),
vOffset=vel_cen,
fixSeed=False,
vRadial=0,
vPosAng=0,
#vPhaseCen=np.array([x_cen,y_cen]),
#That ^ was a workaround for a bug in KinMS that has now been fixed
returnClouds=False,
gasGrav=False,
fileName=False)
model = model_cont+self.modelimage
xpos,ypos=self.xpos_center_padded,self.ypos_center_padded
model_padded=np.transpose(np.pad(model,((ypos-self.Nypix_small/2,self.Nypix-ypos-self.Nypix_small/2),(xpos-self.Nxpix_small/2,self.Nxpix-xpos-self.Nxpix_small/2),(0,0)),mode='constant'),(2,0,1))
modelimage_cube = model_padded
vis_complex_model=np.copy(self.vis_complex_model_template)
for chan in range(modelimage_cube.shape[0]):
uu=np.ones((self.uu_cube.shape[0],self.uu_cube.shape[1],1,self.uu_cube.shape[3]))
vv=np.ones((self.vv_cube.shape[0],self.vv_cube.shape[1],1,self.vv_cube.shape[3]))
uu[:,:,0,:]=self.uu_cube[:,:,chan,:]
vv[:,:,0,:]=self.vv_cube[:,:,chan,:]
modelimage=modelimage_cube[chan]
uushape = uu.shape
uu = uu.flatten()
vv = vv.flatten()
uu=uu.copy(order='C') #This
vv=vv.copy(order='C') #this
modelimage=np.roll(np.flip(modelimage,axis=0),1,axis=0).copy(order='C')#.byteswap().newbyteorder() #This
model_complex = sampleImage(modelimage, np.absolute(self.modelheader['CDELT1'])/180*np.pi, uu, vv) #this uses galario
#model_complex = sample_vis.uvmodel(modelimage, modelheader, uu, vv, pcd)
vis_complex = model_complex.reshape(uushape)
vis_complex_model[:,:,chan,:]=vis_complex[:,:,0,:]
#replace_visibilities('HZ10_spw01_comb.uvfits','my_img_mod.fits','model_visib.uvfits')
#vis_complex_model,bb = uvutil.visload('model_visib.uvfits')
vis_complex_model=vis_complex_model.flatten()[self.good_vis]
def find_param(scale):
diff_all=np.abs(self.vis_complex_data-vis_complex_model*scale)
return np.sum(self.wgt_data*diff_all*diff_all)
ln_like=-0.5*find_param(1)
print(ln_like)
return ln_like+self.offset_lnlike
def run_Multinest(self,Nthreads,output_filename='temporary_model'):
'''
This function runs Multinest using the self.my_prior priors, which need to be appropriately set to a Multinest-type of priors.
Parameters
Nthreads: int
Number of threads per Multinest instance used by Galario to compute the model visibilities
Output_filename: string
Multinest output base file names. See pymultinest docs for how to read and analyze those.
'''
n_params = 10
from galario import double
double.threads(Nthreads)
pymultinest.run(self.loglike_multinest, self.my_prior, n_params, outputfiles_basename=output_filename, resume = True, verbose = True,sampling_efficiency='model')
def lnprob(self,cube):
'''
This function combines the emcee-type prior and the log likelihood to compute the log probability. Should be used with emcee.
'''
if self.my_prior(cube)>-np.inf:
return self.my_prior(cube)+self.loglike(cube)
else:
return -np.inf
def find_max_prob(self,starting_guess,Nthreads_galario):
'''
This function tries to find the maximum probability parameter values. This may be useful for ML fitting or, e.g., to start an MCMC sampler such as emcee.
Parameters
Nthreads_galario: int
Number of threads to be used by Galario
starting_guess: list of float
This should be a list of disk fitting parameters to be used as a starting point for the optimizer.
Returns
result['x']: list of float
The list of parameters which achieves the highest probability, as computed by the optimizer.
'''
from galario import double
double.threads(Nthreads_galario)
nll = lambda *args: -self.lnprob(*args)
from scipy.optimize import minimize
result=minimize(nll,starting_guess,method='Nelder-Mead')
print(result)
return result['x']
def run_emcee(self,Nthreads_galario,Nthreads_emcee,nwalkers,starting_point,Nsteps,output_filename='temporary_model'):
'''
This function runs emcee to produce MCMC samples from the probability distribution defined by my_prior and loglike.
Parameters
Nthreads_galario, Nthreads_emcee: (int,int)
Number of threads to be used by Galario and by emcee, respectively
nwalkers: int
Number of Emcee walkers, see emcee documentation
starting_point: list of float
List of parameters to be used as starting point for the MCMC chain
Nsteps: int
Number of MCMC steps, needs to be >100 because I burn 100 by default (easy to change this below). The functions saves the samples every 100 steps.
output_filename: string
Filename for the numpy save file storing the samples
'''
ndim=10
from galario import double
double.threads(Nthreads_galario)
pos = [np.array(starting_point)*(1+.1*np.random.randn(ndim)) for i in range(nwalkers)]
sampler = emcee.EnsembleSampler(nwalkers, ndim, picklable_boundmethod(self.lnprob),threads=Nthreads_emcee)
sampler.run_mcmc(pos, 100)
for idx in range(int(Nsteps/100)):
sampler.run_mcmc(None, 100)
samples = sampler.chain[:, 10:, :].reshape((-1, ndim))
np.save(output_filename,samples)
print(idx,' of ',int(Nsteps/100))
#Necessary in order to use the lnprob instance method within the emcee sampler, to make it picklable
class picklable_boundmethod(object):
def __init__(self, mt):
self.mt = mt
def __getstate__(self):
return self.mt.im_self, self.mt.im_func.__name__
def __setstate__(self, (s,fn)):
self.mt = getattr(s, fn)
def __call__(self, *a, **kw):
return self.mt(*a, **kw)
'''
#Example of priors for Multinest
def prior(cube,ndim,nparams):
from Priors_multinest import Priors
pri=Priors()
cube[0]=pri.LogPrior(cube[0],10.,700.)
cube[1]=pri.LogPrior(cube[1],.1,2.)
cube[2]=pri.LogPrior(cube[2],10.,2000.)
cube[3]=pri.LogPrior(cube[3],.01,.3)
cube[4]=pri.UniformPrior(cube[4],-400.,400.)
cube[5]=pri.SinPrior(cube[5],0.,90.)
cube[6]=pri.UniformPrior(cube[6],60.,120.)
cube[7]=pri.UniformPrior(cube[7],-.15,.15)
cube[8]=pri.UniformPrior(cube[8],-.15,.15)
cube[9]=pri.LogPrior(cube[9],1,5)
#Example of priors for emcee
def prior(cube):
from Priors_emcee import Priors
pri=Priors()
lnprior=0
lnprior+=pri.LogPrior(cube[0],10.,700.)
lnprior+=pri.LogPrior(cube[1],.1,2.)
lnprior+=pri.LogPrior(cube[2],10.,2000.)
lnprior+=pri.LogPrior(cube[3],.01,.3)
lnprior+=pri.UniformPrior(cube[4],-400.,400.)
lnprior+=pri.SinPrior(cube[5],0.,90.)
lnprior+=pri.UniformPrior(cube[6],60.,120.)
lnprior+=pri.UniformPrior(cube[7],-.15,.15)
lnprior+=pri.UniformPrior(cube[8],-.15,.15)
lnprior+=pri.LogPrior(cube[9],1.,5.)
return lnprior
'''