From 770f014ee681a3b84c7cad54094aa3bebdff08d2 Mon Sep 17 00:00:00 2001 From: George Wong Date: Thu, 17 Dec 2020 18:30:49 -0600 Subject: [PATCH 01/28] add optional multiprocessing to movie scattering method --- ehtim/scattering/stochastic_optics.py | 33 +++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/ehtim/scattering/stochastic_optics.py b/ehtim/scattering/stochastic_optics.py index c3e7a413..d336cf7a 100644 --- a/ehtim/scattering/stochastic_optics.py +++ b/ehtim/scattering/stochastic_optics.py @@ -18,6 +18,9 @@ from ehtim.observing.obs_helpers import * from ehtim.const_def import * #Note: C is m/s rather than cm/s. +from multiprocessing import cpu_count +from multiprocessing import Pool + import math import cmath @@ -404,6 +407,10 @@ def MakePhaseScreen(self, EpsilonScreen, Reference_Image, obs_frequency_Hz=0.0, return phi_Image + def Scatter2(self, args, kwargs): + """Call self.Scatter with expanded args and kwargs.""" + return self.Scatter(*args, **kwargs) + def Scatter(self, Unscattered_Image, Epsilon_Screen=np.array([]), obs_frequency_Hz=0.0, Vx_km_per_s=50.0, Vy_km_per_s=0.0, t_hr=0.0, ea_ker=None, sqrtQ=None, Linearized_Approximation=False, DisplayImage=False, Force_Positivity=False, use_approximate_form=True): """Scatter an image using the specified epsilon screen. All lengths should be specified in centimeters @@ -522,7 +529,7 @@ def Scatter(self, Unscattered_Image, Epsilon_Screen=np.array([]), obs_frequency_ return AI_Image - def Scatter_Movie(self, Unscattered_Movie, Epsilon_Screen=np.array([]), obs_frequency_Hz=0.0, Vx_km_per_s=50.0, Vy_km_per_s=0.0, framedur_sec=None, N_frames = None, sqrtQ=None, Linearized_Approximation=False, Force_Positivity=False,Return_Image_List=False): + def Scatter_Movie(self, Unscattered_Movie, Epsilon_Screen=np.array([]), obs_frequency_Hz=0.0, Vx_km_per_s=50.0, Vy_km_per_s=0.0, framedur_sec=None, N_frames = None, sqrtQ=None, Linearized_Approximation=False, Force_Positivity=False, Return_Image_List=False, processes=0): """Scatter a movie using the specified epsilon screen. The movie can either be a movie object, an image list, or a static image If scattering a list of images or static image, the frame duration in seconds (framedur_sec) must be specified If scattering a static image, the total number of frames must be specified (N_frames) @@ -542,6 +549,7 @@ def Scatter_Movie(self, Unscattered_Movie, Epsilon_Screen=np.array([]), obs_freq Linearized_Approximation (bool): If True, uses a linearized approximation for the scattering (Eq. 10 of Johnson & Narayan 2016). If False, uses Eq. 9 of that paper. Force_Positivity (bool): If True, eliminates negative flux from the scattered image from the linearized approximation. Return_Image_List (bool): If True, returns a list of the scattered frames. If False, returns a movie object. + processes (int): Number of cores to use in multiprocessing. Default value (0) means no multiprocessing. Uses all available cores if processes < 0. Returns: Scattered_Movie: Either a movie object or a list of images, depending on the flag Return_Image_List. @@ -623,9 +631,26 @@ def get_frame(j): if Epsilon_Screen.shape[0] == 0: Epsilon_Screen = MakeEpsilonScreen(N, N) - scattered_im_List = [ self.Scatter(get_frame(j), Epsilon_Screen, obs_frequency_Hz = obs_frequency_Hz, Vx_km_per_s = Vx_km_per_s, Vy_km_per_s = Vy_km_per_s, - t_hr=tlist_hr[j], sqrtQ=sqrtQ, Linearized_Approximation=Linearized_Approximation, Force_Positivity=Force_Positivity) for j in range(N_frames)] - + # possibly parallelize + if processes < 0: + processes = cpu_count() + processes = min(processes, N_frames) + + # generate scattered images + if processes > 0: + pool = Pool(processes=processes) + args = [ + ( + [get_frame(j), Epsilon_Screen], + dict(obs_frequency_Hz = obs_frequency_Hz, Vx_km_per_s = Vx_km_per_s, Vy_km_per_s = Vy_km_per_s, t_hr=tlist_hr[j], sqrtQ=sqrtQ, Linearized_Approximation=Linearized_Approximation, Force_Positivity=Force_Positivity) + ) for j in range(N_frames) + ] + scattered_im_List = pool.starmap(self.Scatter2, args) + pool.close() + pool.join() + else: + scattered_im_List = [self.Scatter(get_frame(j), Epsilon_Screen, obs_frequency_Hz = obs_frequency_Hz, Vx_km_per_s = Vx_km_per_s, Vy_km_per_s = Vy_km_per_s, t_hr=tlist_hr[j], sqrtQ=sqrtQ, Linearized_Approximation=Linearized_Approximation, Force_Positivity=Force_Positivity) for j in range(N_frams)] + if Return_Image_List == True: return scattered_im_List From 64796c77c72a2fe99b3c07fe073cb7c3a5532923 Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Tue, 12 Jan 2021 21:48:40 -0500 Subject: [PATCH 02/28] fixed bug in blur_gauss() with latest version of scipy.signal.fftconvolve --- ehtim/image.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ehtim/image.py b/ehtim/image.py index 720905dd..2ebacd5f 100644 --- a/ehtim/image.py +++ b/ehtim/image.py @@ -1395,7 +1395,7 @@ def blur(imarr, gauss): return imarr_blur # Convolve the primary image - imarr = (self.imvec).reshape(self.ydim, self.xdim) + imarr = (self.imvec).reshape(self.ydim, self.xdim).astype('float64') imarr_blur = blur(imarr, gauss) # Make new image object @@ -1409,7 +1409,7 @@ def blur(imarr, gauss): continue polvec = self._imdict[pol] if len(polvec): - polarr = polvec.reshape(self.ydim, self.xdim) + polarr = polvec.reshape(self.ydim, self.xdim).astype('float64') if frac_pol: polarr = blur(polarr, gausspol) outim.add_pol_image(polarr, pol) @@ -1418,7 +1418,7 @@ def blur(imarr, gauss): mflist_out = [] for mfvec in self._mflist: if len(mfvec): - mfarr = mfvec.reshape(self.ydim, self.xdim) + mfarr = mfvec.reshape(self.ydim, self.xdim).astype('float64') mfarr = blur(mfarr, gauss) mfvec_out = mfarr.flatten() else: From 0cdbf19252bcb6fca6c82e18e11d28a6815c7342 Mon Sep 17 00:00:00 2001 From: Katie Bouman Date: Tue, 2 Feb 2021 09:46:47 -0800 Subject: [PATCH 03/28] updated starwarps to have a lightcurve flux constraint --- ehtim/imaging/starwarps.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/ehtim/imaging/starwarps.py b/ehtim/imaging/starwarps.py index caee8f17..2bd4c95b 100644 --- a/ehtim/imaging/starwarps.py +++ b/ehtim/imaging/starwarps.py @@ -155,7 +155,7 @@ def forwardUpdates_apxImgs(mu, Lambda_orig, obs_List, A_orig, Q_orig, init_image for k in range(0,numLinIters): # F is the derivative of the Forward model with respect to the unknown parameters - meas, idealmeas, F, measCov, valid = getMeasurementTerms(obs_List[t], z_List_lin[t], measurement=measurement, mask=mask, normalize=normalize) + meas, idealmeas, F, measCov, valid = getMeasurementTerms(obs_List[t], z_List_lin[t], measurement=measurement, tot_flux=lightcurve[t], mask=mask, normalize=normalize) if valid: z_List_t_t[t].imvec[mask], P_List_t_t[t] = prodGaussiansLem2(F, measCov, meas, z_star_List_t_tm1[t].imvec[mask], P_star_List_t_tm1[t]) @@ -239,7 +239,7 @@ def backwardUpdates(mu, Lambda_orig, obs_List, A_orig, Q_orig, measurement={'vis # update - meas, idealmeas, F, measCov, valid = getMeasurementTerms(obs_List[t], apxImgs[t], measurement=measurement, mask=mask, normalize=normalize) + meas, idealmeas, F, measCov, valid = getMeasurementTerms(obs_List[t], apxImgs[t], measurement=measurement, tot_flux=lightcurve[t], mask=mask, normalize=normalize) if valid: z_t_t[t].imvec[mask], P_t_t[t] = prodGaussiansLem2(F, measCov, meas, z_star_t_tp1[t].imvec[mask], P_star_t_tp1[t]) @@ -284,6 +284,9 @@ def computeSuffStatistics(mu, Lambda, obs_List, Upsilon, theta, init_x, init_y, if mu[0].xdim!=mu[0].ydim: error('Error: This has only been checked thus far on square images!') + + if lightcurve == None and 'flux' in measurement.keys(): #KATIE ADDED FEB 1 2021 + error('Error: if you are using a flux constraint you must specify a lightcurve') if list(measurement.keys())==1 and measurement.keys()[0]=='vis': numLinIters = 1 @@ -611,7 +614,7 @@ def prodGaussiansLem2(A, Sigma, y, mu, Q): return (mean, covariance) -def getMeasurementTerms(obs, im, measurement={'vis': 1}, mask=[], normalize=False): +def getMeasurementTerms(obs, im, measurement={'vis': 1}, tot_flux=None, mask=[], normalize=False): if not np.sum(mask)==len(mask): raise ValueError('The code doenst currently work with a mask!') @@ -632,7 +635,13 @@ def getMeasurementTerms(obs, im, measurement={'vis': 1}, mask=[], normalize=Fals # check to see if you have data in the current obs try: - data, sigma, A = chisqdata(obs, im, mask, dtype=dname, ttype='direct') + if dname=='flux': + if tot_flux == None: + error('Error: if you are using a flux constraint you must specify a total flux (via the lightcurve)') + data = np.array([tot_flux]) + sigma = np.array([1]) + else: + data, sigma, A = chisqdata(obs, im, mask, dtype=dname, ttype='direct') count = count + 1 except: continue @@ -653,6 +662,9 @@ def getMeasurementTerms(obs, im, measurement={'vis': 1}, mask=[], normalize=Fals elif dname == 'logcamp': F = grad_logcamp(im.imvec, A) ideal = logcamp(im.imvec,A) + elif dname == 'flux': + F = grad_flux(im.imvec) + ideal = flux(im.imvec) #turn complex matrices to real if not np.allclose(data.imag,0): @@ -693,6 +705,12 @@ def grad_bs(imvec, Amatrices): out = pt1[:,None] * Amatrices[0] + pt2[:,None] * Amatrices[1] + pt3[:,None] * Amatrices[2] return out +def flux(imvec): + return np.sum(imvec) + +def grad_flux(imvec): + return np.ones((1, len(imvec))) + def cphase(imvec, Amatrices): """the closure phase""" i1 = np.dot(Amatrices[0], imvec) From 65fca10d29187f31d1f8d74a54b5101743b5236b Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Fri, 12 Feb 2021 17:59:50 -0500 Subject: [PATCH 04/28] added neggains to synthetic data with sigmat --- ehtim/features/rex.py | 7 +++++-- ehtim/image.py | 6 ++++-- ehtim/movie.py | 1 - ehtim/observing/obs_simulate.py | 13 +++++++++---- 4 files changed, 18 insertions(+), 9 deletions(-) diff --git a/ehtim/features/rex.py b/ehtim/features/rex.py index 5fdfcc3e..fb824750 100644 --- a/ehtim/features/rex.py +++ b/ehtim/features/rex.py @@ -797,7 +797,7 @@ def FindProfileSingle(imname, postprocdir, rmin_search=RPRIOR_MIN, rmax_search=RPRIOR_MAX, nrays_search=NRAYS_SEARCH, nrs_search=NRS_SEARCH, thresh_search=THRESH, fov_search=FOVP_SEARCH, n_search=NSEARCH, - flux_norm=NORMFLUX): + flux_norm=NORMFLUX,center=False): """find the best ring profile for an image and save results """ @@ -817,7 +817,10 @@ def FindProfileSingle(imname, postprocdir, im_raw = im_raw.blur_circ(blur*ehc.RADPERUAS, blur*ehc.RADPERUAS) # center image and regrid to uniform pixel size and fox - im = di.center_core(im_raw) + if center: + im = di.center_core(im_raw) # TODO -- why isn't this working? + else: + im = im_raw im_search = im.regrid_image(imsize, npix) im = im.regrid_image(imsize, npix) diff --git a/ehtim/image.py b/ehtim/image.py index 2ebacd5f..01d61c39 100644 --- a/ehtim/image.py +++ b/ehtim/image.py @@ -1503,8 +1503,10 @@ def gradim(imvec): imarr = imvec.reshape(self.ydim, self.xdim) - sx = ndi.sobel(imarr, axis=0, mode='constant') - sy = ndi.sobel(imarr, axis=1, mode='constant') + #sx = ndi.sobel(imarr, axis=0, mode='constant') + #sy = ndi.sobel(imarr, axis=1, mode='constant') + sx = ndi.sobel(imarr, axis=0, mode='nearest') + sy = ndi.sobel(imarr, axis=1, mode='nearest') # TODO: are these in the right order?? if gradtype == 'x': diff --git a/ehtim/movie.py b/ehtim/movie.py index 65a03591..f70c9d22 100644 --- a/ehtim/movie.py +++ b/ehtim/movie.py @@ -1162,7 +1162,6 @@ def observe_same(self, obs_in, repeat=False, # Jones Matrix Corruption & Calibration if jones: - print("Applying Jones Matrices to data . . . ") obsdata = simobs.add_jones_and_noise(obs, add_th_noise=add_th_noise, opacitycal=opacitycal, ampcal=ampcal, phasecal=phasecal, frcal=frcal, dcal=dcal, diff --git a/ehtim/observing/obs_simulate.py b/ehtim/observing/obs_simulate.py index e2348154..e99f6d4a 100644 --- a/ehtim/observing/obs_simulate.py +++ b/ehtim/observing/obs_simulate.py @@ -538,8 +538,6 @@ def make_jones(obs, opacitycal=True, ampcal=True, phasecal=True, dcal=True, for time in times_stable_amp ), float))) - gainR = np.abs(gainR) - gainL = np.abs(gainL) if neggains: gainR = np.exp(-np.abs(np.log(gainR))) gainL = np.exp(-np.abs(np.log(gainL))) @@ -553,14 +551,21 @@ def make_jones(obs, opacitycal=True, ampcal=True, phasecal=True, dcal=True, (1 + gain_mult * obsh.hashmultivariaterandn( len(scan_start_times), cov, site, 'gain', str(time), str(gain_mult), seed)) )) - gainR_interpolateor = interp1d(scan_start_times, gainR, kind='zero') - gainR = gainR_interpolateor(times_stable_amp) + gainL = np.sqrt(np.abs( (1 + gainL_constant) * (1 + gain_mult * obsh.hashmultivariaterandn( len(scan_start_times), cov, site, 'gain', str(time), str(gain_mult), seed)) )) + + if neggains: + gainR = np.exp(-np.abs(np.log(gainR))) + gainL = np.exp(-np.abs(np.log(gainL))) + + gainR_interpolateor = interp1d(scan_start_times, gainR, kind='zero') + gainR = gainR_interpolateor(times_stable_amp) + gainL_interpolateor = interp1d(scan_start_times, gainL, kind='zero') gainL = gainL_interpolateor(times_stable_amp) From e38d78eaaeb1ab467072256e2d23f461445c1df2 Mon Sep 17 00:00:00 2001 From: Chi-kwan Chan Date: Sat, 13 Feb 2021 02:54:52 -0700 Subject: [PATCH 05/28] Optimize network_cal_scan(), specifically, errfunc() --- ehtim/calibrating/network_cal.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/ehtim/calibrating/network_cal.py b/ehtim/calibrating/network_cal.py index e8e29052..c98178b3 100644 --- a/ehtim/calibrating/network_cal.py +++ b/ehtim/calibrating/network_cal.py @@ -328,9 +328,9 @@ def errfunc(gvpar): # choose to only scale ampliltudes or phases if method == "phase": - g = g / np.abs(g) # TODO: use exp(i*np.arg())? - if method == "amp": - g = np.abs(np.real(g)) + g = g / np.abs(g) + elif method == "amp": + g = np.abs(g) # append the default values to g for missing points # and to v for the zero baseline points @@ -351,17 +351,19 @@ def errfunc(gvpar): else: verr = vis - g1 * g2.conj() * v_scan - nan_mask = np.array([not np.isnan(viter) for viter in verr]) * \ - np.array([not np.isnan(viter) for viter in sigma_inv]) - verr = verr[nan_mask * vis_mask] - - chisq = np.sum((verr.real * sigma_inv[nan_mask * vis_mask])**2) + \ - np.sum((verr.imag * sigma_inv[nan_mask * vis_mask])**2) + chi = np.abs(verr) * sigma_inv + chisq = np.sum((chi * chi)[np.isfinite(chi) * vis_mask]) # prior on the gains g_fracerr = gain_tol - chisq_g = np.sum((np.log(np.abs(g))**2 / g_fracerr**2)) - chisq_v = np.sum((np.abs(v) / zbl_scan)**4) + if method == "phase": + chisq_g = 0 # because |g| == 1 so log(|g|) = 0 + elif method == "amp": + chisq_g = np.sum(np.log(g)**2) / g_fracerr**2 + else: + chisq_g = np.sum(np.log(np.abs(g))**2) / g_fracerr**2 + + chisq_v = np.sum(np.abs(v)**4) / zbl_scan**4 return chisq + chisq_g + chisq_v if np.max(g1_keys) > -1 or np.max(g2_keys) > -1: From 881a98a4d6f9058bd54c46f06896899122cb77a7 Mon Sep 17 00:00:00 2001 From: Chi-kwan Chan Date: Sun, 14 Feb 2021 22:42:55 -0700 Subject: [PATCH 06/28] Restore g = np.abs(np.real(g)) to push the gain to have zero phase --- ehtim/calibrating/network_cal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ehtim/calibrating/network_cal.py b/ehtim/calibrating/network_cal.py index c98178b3..a6ac3582 100644 --- a/ehtim/calibrating/network_cal.py +++ b/ehtim/calibrating/network_cal.py @@ -330,7 +330,7 @@ def errfunc(gvpar): if method == "phase": g = g / np.abs(g) elif method == "amp": - g = np.abs(g) + g = np.abs(np.real(g)) # append the default values to g for missing points # and to v for the zero baseline points From 8d9dcb8903a0a7b679a9d8ac443ea9bfa0545396 Mon Sep 17 00:00:00 2001 From: Chi-kwan Chan Date: Sun, 14 Feb 2021 22:44:00 -0700 Subject: [PATCH 07/28] Optimize further by replacing power `**` by multiplications --- ehtim/calibrating/network_cal.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/ehtim/calibrating/network_cal.py b/ehtim/calibrating/network_cal.py index a6ac3582..814253e9 100644 --- a/ehtim/calibrating/network_cal.py +++ b/ehtim/calibrating/network_cal.py @@ -359,11 +359,15 @@ def errfunc(gvpar): if method == "phase": chisq_g = 0 # because |g| == 1 so log(|g|) = 0 elif method == "amp": - chisq_g = np.sum(np.log(g)**2) / g_fracerr**2 + logg = np.log(g) + chisq_g = np.sum(logg * logg) / (g_fracerr * g_fracerr) else: - chisq_g = np.sum(np.log(np.abs(g))**2) / g_fracerr**2 + logabsg = np.log(np.abs(g)) + chisq_g = np.sum(logabsg * logabsg) / (g_fracerr * g_fracerr) - chisq_v = np.sum(np.abs(v)**4) / zbl_scan**4 + absv = np.abs(v) + vv = absv * absv + chisq_v = np.sum(vv * vv) / zbl_scan**4 return chisq + chisq_g + chisq_v if np.max(g1_keys) > -1 or np.max(g2_keys) > -1: From 60130c87cc2917c101701a22e811fc3b655bed61 Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Tue, 16 Feb 2021 00:20:33 -0500 Subject: [PATCH 08/28] removed debias=True default from some plotting scripts. Still need to figure out debiasing overall. --- ehtim/obsdata.py | 4 ++-- ehtim/plotting/comp_plots.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/ehtim/obsdata.py b/ehtim/obsdata.py index 041614f4..47892eac 100644 --- a/ehtim/obsdata.py +++ b/ehtim/obsdata.py @@ -3843,7 +3843,7 @@ def camp_quad(self, site1, site2, site3, site4, return outdata def plotall(self, field1, field2, - conj=False, debias=True, tag_bl=False, ang_unit='deg', timetype=False, + conj=False, debias=False, tag_bl=False, ang_unit='deg', timetype=False, axis=False, rangex=False, rangey=False, snrcut=0., color=ehc.SCOLORS[0], marker='o', markersize=ehc.MARKERSIZE, label=None, grid=True, ebar=True, axislabels=True, legend=False, @@ -4071,7 +4071,7 @@ def plotall(self, field1, field2, return x def plot_bl(self, site1, site2, field, - debias=True, ang_unit='deg', timetype=False, + debias=False, ang_unit='deg', timetype=False, axis=False, rangex=False, rangey=False, snrcut=0., color=ehc.SCOLORS[0], marker='o', markersize=ehc.MARKERSIZE, label=None, grid=True, ebar=True, axislabels=True, legend=False, diff --git a/ehtim/plotting/comp_plots.py b/ehtim/plotting/comp_plots.py index 98e91559..c337049b 100644 --- a/ehtim/plotting/comp_plots.py +++ b/ehtim/plotting/comp_plots.py @@ -40,7 +40,7 @@ def plotall_compare(obslist, imlist, field1, field2, - conj=False, debias=True, sgrscat=False, + conj=False, debias=False, sgrscat=False, ang_unit='deg', timetype='UTC', ttype='nfft', axis=False, rangex=False, rangey=False, snrcut=0., clist=COLORLIST, legendlabels=None, markersize=ehc.MARKERSIZE, @@ -106,7 +106,7 @@ def plotall_compare(obslist, imlist, field1, field2, def plot_bl_compare(obslist, imlist, site1, site2, field, - debias=True, sgrscat=False, + debias=False, sgrscat=False, ang_unit='deg', timetype='UTC', ttype='nfft', axis=False, rangex=False, rangey=False, snrcut=0., clist=COLORLIST, legendlabels=None, markersize=ehc.MARKERSIZE, @@ -261,7 +261,7 @@ def plot_cphase_compare(obslist, imlist, site1, site2, site3, def plot_camp_compare(obslist, imlist, site1, site2, site3, site4, vtype='vis', ctype='camp', camps=[], force_recompute=False, - debias=True, sgrscat=False, timetype='UTC', ttype='nfft', + debias=False, sgrscat=False, timetype='UTC', ttype='nfft', axis=False, rangex=False, rangey=False, snrcut=0., clist=COLORLIST, legendlabels=None, markersize=ehc.MARKERSIZE, export_pdf="", grid=False, ebar=True, From 15defd7c8646efbb7258b69fee007abb7ca44285 Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Tue, 16 Feb 2021 00:41:57 -0500 Subject: [PATCH 09/28] modified load_fits to read 3 column aipscc table --- ehtim/io/load.py | 46 +++++++++++++++++++++++++++++----------------- 1 file changed, 29 insertions(+), 17 deletions(-) diff --git a/ehtim/io/load.py b/ehtim/io/load.py index f24e76d2..d3e25191 100644 --- a/ehtim/io/load.py +++ b/ehtim/io/load.py @@ -314,27 +314,35 @@ def load_im_fits(filename, aipscc=False, pulse=ehc.PULSE_DEFAULT, deltay = aipscctab.data["DELTAY"] # check to make sure all the source types are point sources and gaussian components - checkmtype = np.abs(np.unique(aipscctab.data["TYPE OBJ"])) < 2.0 - if False in checkmtype.tolist(): - errmsg = "The primary AIPS CC table in the input FITS file has non point-source" - errmsg += " or Gaussian Source CC components, which are not currently supported." - raise ValueError(errmsg) + try: + checkmtype = np.abs(np.unique(aipscctab.data["TYPE OBJ"])) < 2.0 + if False in checkmtype.tolist(): + errmsg = "The primary AIPS CC table in the input FITS file has non point-source" + errmsg += " or Gaussian Source CC components, which are not currently supported." + raise ValueError(errmsg) + point_src = aipscctab.data["TYPE OBJ"] == 0 + gaussian_src = aipscctab.data["TYPE OBJ"] == 1 + except(KeyError): + print("Cannot load AIPS CC Table OBJ data -- assuming all CC components are point sources!") + point_src = np.ones(aipscctab.data.shape).astype(bool) + gaussian_src = np.zeros(aipscctab.data.shape).astype(bool) print("%d CC components are loaded." % (len(flux))) # compile the point source aipscc info - point_src = aipscctab.data["TYPE OBJ"] == 0 flux_ps = flux[point_src] deltax_ps = deltax[point_src] deltay_ps = deltay[point_src] - # compile the gaussian aipscc info - gaussian_src = aipscctab.data["TYPE OBJ"] == 1 - flux_gs = flux[gaussian_src] - deltax_gs = deltax[gaussian_src] - deltay_gs = deltay[gaussian_src] - maj_gs = aipscctab.data["MAJOR AX"][gaussian_src] - min_gs = aipscctab.data["MINOR AX"][gaussian_src] - pa_gs = aipscctab.data["POSANGLE"][gaussian_src] + # compile the gaussian aipscc info, if any + if np.any(gaussian_src): + flux_gs = flux[gaussian_src] + deltax_gs = deltax[gaussian_src] + deltay_gs = deltay[gaussian_src] + maj_gs = aipscctab.data["MAJOR AX"][gaussian_src] + min_gs = aipscctab.data["MINOR AX"][gaussian_src] + pa_gs = aipscctab.data["POSANGLE"][gaussian_src] + else: + flux_gs = [] # the map_coordinates delta x / delta y of each delta CC component are # relative to the reference pixel which is defined by CRPIX1 and CRPIX2. @@ -372,16 +380,20 @@ def load_im_fits(filename, aipscc=False, pulse=ehc.PULSE_DEFAULT, if header['BUNIT'].lower() == 'JY/BEAM'.lower(): print("converting Jy/Beam --> Jy/pixel") + bmaj = bmin = 1.0 # default values + if 'BMAJ' in list(header.keys()): bmaj = header['BMAJ'] bmin = header['BMIN'] + elif 'HISTORY' in list(header.keys()): # Alternate option, to read AIPS fits images print("No beam info in header; reading from AIPS HISTORY instead...") for line in header['HISTORY']: - if 'BMAJ' in line: + if 'BMAJ' in line and len(line.split())>6: bmaj = float(line.split()[3]) bmin = float(line.split()[5]) - else: + + if bmaj==1.0 and bmin==1.0: print("No beam info found! Assuming nominal values for conversion.") bmaj = bmin = 1.0 @@ -389,7 +401,7 @@ def load_im_fits(filename, aipscc=False, pulse=ehc.PULSE_DEFAULT, normalizer = (header['CDELT2'])**2 / beamarea if aipscc: - print("the computed normalizer will not be applied since loading the AIPS CC table") + print("the computed normalizer will not be applied since we are loading the AIPS CC table") else: image *= normalizer From 8973d2b96513cbf82db8c7a81d21a73d7bc5450f Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Fri, 19 Feb 2021 06:28:29 -0500 Subject: [PATCH 10/28] fixed bug in movie where copy() did not copy time or frame lists --- ehtim/movie.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ehtim/movie.py b/ehtim/movie.py index f70c9d22..e75859ae 100644 --- a/ehtim/movie.py +++ b/ehtim/movie.py @@ -396,7 +396,8 @@ def movie_args(self): """ frames2D = self.frames.reshape((self.nframes, self.ydim, self.xdim)) - arglist = [frames2D, self.times, self.psize, self.ra, self.dec] + arglist = [frames2D.copy(), self.times.copy(), self.psize, self.ra, self.dec] + #arglist = [frames2D, self.times, self.psize, self.ra, self.dec] argdict = {'rf': self.rf, 'polrep': self.polrep, 'pol_prim': self.pol_prim, 'pulse': self.pulse, 'source': self.source, 'mjd': self.mjd, 'interp': self.interp, 'bounds_error': self.bounds_error} From b4fa84b29314b84472a94329b0200a61e13e8ad9 Mon Sep 17 00:00:00 2001 From: Kazu Akiyama Date: Tue, 16 Mar 2021 10:44:23 -0400 Subject: [PATCH 11/28] fix uvw scaling The current function scales uvw's unit by setting up PSCAL/PZERO. This apparently works fine when we save uvfits and load it again because of a bug in astropy.io.fits. This won't work if we want to directly work with hdulist rather than save/load uvfits files, since obsdata.load_uvfits don't attempt to load PSCAL/PZERO for uvw coordinates. This fix modifies the location of scaling. PSCAL/PZERO for uv coordinates will be 1/0 respecitvely. Instead u, v will be directly normalized by rf. --- ehtim/io/save.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ehtim/io/save.py b/ehtim/io/save.py index ddd257e1..31dd81a7 100644 --- a/ehtim/io/save.py +++ b/ehtim/io/save.py @@ -472,13 +472,13 @@ def save_obs_uvfits(obs, fname, force_singlepol=None, polrep_out='circ'): header['CRPIX7'] = 1.e0 header['CROTA7'] = 0.e0 header['PTYPE1'] = 'UU---SIN' - header['PSCAL1'] = 1.0/obs.rf + header['PSCAL1'] = 1.0 header['PZERO1'] = 0.e0 header['PTYPE2'] = 'VV---SIN' - header['PSCAL2'] = 1.0/obs.rf + header['PSCAL2'] = 1.0 header['PZERO2'] = 0.e0 header['PTYPE3'] = 'WW---SIN' - header['PSCAL3'] = 1.0/obs.rf + header['PSCAL3'] = 1.0 header['PZERO3'] = 0.e0 header['PTYPE4'] = 'BASELINE' header['PSCAL4'] = 1.e0 @@ -528,8 +528,8 @@ def save_obs_uvfits(obs, fname, force_singlepol=None, polrep_out='circ'): tau2 = obsdata['tau2'] # uv are in lightseconds - u = obsdata['u'] - v = obsdata['v'] + u = obsdata['u']/obs.rf + v = obsdata['v']/obs.rf # rr, ll, lr, rl, weights From ed4d95d214c7359539c6455c622250ecd5283b3c Mon Sep 17 00:00:00 2001 From: Kazu Akiyama Date: Tue, 16 Mar 2021 11:33:09 -0400 Subject: [PATCH 12/28] allow directly load/save HDUList object for load/save_uvfits functions --- ehtim/io/load.py | 34 +++++++++++++++++++++++++++++----- ehtim/io/save.py | 10 ++++++---- ehtim/obsdata.py | 17 ++++++++++------- 3 files changed, 45 insertions(+), 16 deletions(-) diff --git a/ehtim/io/load.py b/ehtim/io/load.py index f24e76d2..0f74e941 100644 --- a/ehtim/io/load.py +++ b/ehtim/io/load.py @@ -1060,7 +1060,7 @@ def load_obs_uvfits(filename, polrep='stokes', flipbl=False, allow_singlepol=Tru force_singlepol=None, channel=all, IF=all, remove_nan=False): """Load observation data from a uvfits file. Args: - fname (str): path to input text file + fname (str or HDUList): path to input text file or HDUList object polrep (str): load data as either 'stokes' or 'circ' flipbl (bool): flip baseline phases if True. allow_singlepol (bool): If True and polrep='stokes', @@ -1075,11 +1075,15 @@ def load_obs_uvfits(filename, polrep='stokes', flipbl=False, allow_singlepol=Tru if not(polrep in ['stokes', 'circ']): raise Exception("polrep should be 'stokes' or 'circ' in load_uvfits") if not(force_singlepol is None or force_singlepol is False) and polrep != 'stokes': - raise Exception("force_singlepol is incompatible with polrep!='stokes' in load_uvfits") + raise Exception( + "force_singlepol is incompatible with polrep!='stokes' in load_uvfits") # Load the uvfits file - print("Loading uvfits: ", filename) - hdulist = fits.open(filename) + if isinstance(filename, fits.HDUList): + hdulist = filename.copy() + else: + print("Loading uvfits: ", filename) + hdulist = fits.open(filename) header = hdulist[0].header data = hdulist[0].data @@ -1279,7 +1283,27 @@ def load_obs_uvfits(filename, polrep='stokes', flipbl=False, allow_singlepol=Tru print("Warning: removing flagged data present!") # Obs Times - jds = data['DATE'][mask].astype('d') + data['_DATE'][mask].astype('d') + paridx = data.parnames.index("DATE")+1 + if "PSCAL%d"%(paridx) in header.keys(): + jd1scal = header["PSCAL%d"%(paridx)] + else: + jd1scal = 1.0 + if "PZERO%d"%(paridx) in header.keys(): + jd1zero = header["PZERO%d"%(paridx)] + else: + jd1zero = 0.0 + if "PSCAL%d"%(paridx) in header.keys(): + jd2scal = header["PSCAL%d"%(paridx+1)] + else: + jd2scal = 1.0 + if "PZERO%d"%(paridx+1) in header.keys(): + jd2zero = header["PZERO%d"%(paridx+1)] + else: + jd2zero = 0.0 + + jds = jd1scal * data['DATE'][mask].astype('d') + jd1zero + jds += jd2scal * data['_DATE'][mask].astype('d') + jd2zero + mjd = int(np.min(jds) - 2400000.5) times = (jds - 2400000.5 - mjd) * 24.0 diff --git a/ehtim/io/save.py b/ehtim/io/save.py index 31dd81a7..324a448a 100644 --- a/ehtim/io/save.py +++ b/ehtim/io/save.py @@ -401,7 +401,7 @@ def save_obs_txt(obs, fname): return -def save_obs_uvfits(obs, fname, force_singlepol=None, polrep_out='circ'): +def save_obs_uvfits(obs, fname=None, force_singlepol=None, polrep_out='circ'): """Save observation data to uvfits. To save Stokes I as a single polarization (e.g., only RR) set force_singlepol='R' or 'L' """ @@ -815,9 +815,11 @@ def save_obs_uvfits(obs, fname, force_singlepol=None, polrep_out='circ'): print("No NX table in saved uvfits") # Write final HDUList to file - hdulist_new.writeto(fname, overwrite=True) - - return + if fname is None: + return hdulist_new.copy() + else: + hdulist_new.writeto(fname, overwrite=True) + return None def save_obs_oifits(obs, fname, flux=1.0): diff --git a/ehtim/obsdata.py b/ehtim/obsdata.py index 041614f4..0cbb95e5 100644 --- a/ehtim/obsdata.py +++ b/ehtim/obsdata.py @@ -4480,22 +4480,25 @@ def save_txt(self, fname): return - def save_uvfits(self, fname, force_singlepol=False, polrep_out='circ'): + def save_uvfits(self, fname=None, force_singlepol=False, polrep_out='circ'): """Save visibility data to uvfits file. - Args: - fname (str): path to output text file + fname (str): path to output text file. If not specified, return HDUList. force_singlepol (str): if 'R' or 'L', will interpret stokes I field as 'RR' or 'LL' polrep_out (str): 'circ' or 'stokes': how data should be stored in the uvfits file """ if (force_singlepol is not False) and (self.polrep != 'stokes'): - raise Exception("force_singlepol is incompatible with polrep!='stokes'") + raise Exception( + "force_singlepol is incompatible with polrep!='stokes'") - ehtim.io.save.save_obs_uvfits(self, fname, - force_singlepol=force_singlepol, polrep_out=polrep_out) + output = ehtim.io.save.save_obs_uvfits( + self, fname, force_singlepol=force_singlepol, polrep_out=polrep_out) - return + if fname is None: + return output + else: + return def save_oifits(self, fname, flux=1.0): """ Save visibility data to oifits. Polarization data is NOT saved. From 0d6b2ac70cfe261f4c39b4b6dc4b371b883d6e7a Mon Sep 17 00:00:00 2001 From: Kazu Akiyama Date: Tue, 16 Mar 2021 11:33:54 -0400 Subject: [PATCH 13/28] fix a minor issue around the header ["DATE-OBS"] --- ehtim/io/save.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ehtim/io/save.py b/ehtim/io/save.py index 324a448a..3c1351ec 100644 --- a/ehtim/io/save.py +++ b/ehtim/io/save.py @@ -430,7 +430,7 @@ def save_obs_uvfits(obs, fname=None, force_singlepol=None, polrep_out='circ'): header['OBSDEC'] = obs.dec header['OBJECT'] = obs.source header['MJD'] = float(obs.mjd) - header['DATE-OBS'] = Time(obs.mjd + MJD_0, format='jd', scale='utc', out_subfmt='date').iso + header['DATE-OBS'] = Time(obs.mjd + MJD_0, format='jd', scale='utc').iso[0:10] header['BSCALE'] = 1.0 header['BZERO'] = 0.0 header['BUNIT'] = 'JY' From d11e76ecf530bc52f4c6232c542b82d10ce0af94 Mon Sep 17 00:00:00 2001 From: Kazu Akiyama Date: Tue, 16 Mar 2021 11:37:43 -0400 Subject: [PATCH 14/28] minor update in the pydoc of load_uvfits --- ehtim/obsdata.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ehtim/obsdata.py b/ehtim/obsdata.py index 0cbb95e5..9352fa6c 100644 --- a/ehtim/obsdata.py +++ b/ehtim/obsdata.py @@ -4614,7 +4614,7 @@ def load_uvfits(fname, flipbl=False, remove_nan=False, force_singlepol=None, """Load observation data from a uvfits file. Args: - fname (str): path to input text file + fname (str or HDUList): path to input text file or HDUList object flipbl (bool): flip baseline phases if True. remove_nan (bool): True to remove nans from missing polarizations polrep (str): load data as either 'stokes' or 'circ' From d9f19967d39e0ad1f3473af9d2e9887e0df650f9 Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Tue, 16 Mar 2021 13:41:44 -0400 Subject: [PATCH 15/28] updated Imager reporting of solver res.message for python2/3 compatibility --- ehtim/imager.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ehtim/imager.py b/ehtim/imager.py index 14c1222c..05f48665 100644 --- a/ehtim/imager.py +++ b/ehtim/imager.py @@ -316,10 +316,15 @@ def callback_func(xcur): dname_key = dname + ('_%i' % i) outstr += "chi2_%s : %0.2f " % (dname_key, chi2_term_dict[dname_key]) - print("time: %f s" % (tstop - tstart)) - print("J: %f" % res.fun) - print(outstr) - print(res.message.decode()) + try: + print("time: %f s" % (tstop - tstart)) + print("J: %f" % res.fun) + print(outstr) + if isinstance(res.message,str): print(res.message) + else: print(res.message.decode()) + except: # TODO -- issues for some users with res.message + pass + print("==============================") # Embed image From dbdc66e4831e9a471eb2e700bd6a452851201677 Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Tue, 16 Mar 2021 14:58:35 -0400 Subject: [PATCH 16/28] added obsdata.make_hdulist to get hdulist data instead of saving uvfits / cleaned up save.py docstrings --- ehtim/io/save.py | 71 ++++++++++++++++++++++++++++++------------------ ehtim/obsdata.py | 31 +++++++++++++++------ 2 files changed, 67 insertions(+), 35 deletions(-) diff --git a/ehtim/io/save.py b/ehtim/io/save.py index 3c1351ec..a7903e9f 100644 --- a/ehtim/io/save.py +++ b/ehtim/io/save.py @@ -44,6 +44,7 @@ def save_im_txt(im, fname, mjd=False, time=False): """Save image data to text file. Args: + im (Image): image object fname (str): path to output text file mjd (int): MJD of saved image time (float): UTC time of saved image @@ -110,12 +111,11 @@ def save_im_txt(im, fname, mjd=False, time=False): return # TODO save image in circular basis? - - def save_im_fits(im, fname, mjd=False, time=False): """Save image data to a fits file. Args: + im (Image): image object fname (str): path to output fits file mjd (int): MJD of saved image time (float): UTC time of saved image @@ -188,6 +188,7 @@ def save_mov_hdf5(mov, fname, mjd=False): """Save movie data to an hdf5 file. Args: + mov (Movie): movie object fname (str): basename of output fits file mjd (int): MJD of saved movie @@ -197,26 +198,11 @@ def save_mov_hdf5(mov, fname, mjd=False): # TODO: Currently only supports one polarization! with h5py.File(fname, 'w') as file: - # if sys.version_info > (3,0): - # dt = h5py.special_dtype(vlen=str) - # else: - # dt = dtype=h5py.special_dtype(vlen=unicode)) - - # dt = dtype=h5py.special_dtype(vlen=bytes)) head = file.create_dataset('header', (0,), dtype="S10") if mjd is False: mjd = mov.mjd -# head.attrs['mjd'] = str(mjd) -# head.attrs['psize'] = str(mov.psize) -# head.attrs['source'] = str(mov.source) -# head.attrs['ra'] = str(mov.ra) -# head.attrs['dec'] = str(mov.dec) -# head.attrs['rf'] = str(mov.rf) -# head.attrs['polrep'] = str(mov.polrep) -# head.attrs['pol_prim'] = str(mov.pol_prim) - head.attrs['mjd'] = np.string_(str(mjd)) head.attrs['psize'] = np.string_(str(mov.psize)) head.attrs['source'] = np.string_(str(mov.source)) @@ -248,6 +234,7 @@ def save_mov_fits(mov, fname, mjd=False): """Save movie data to series of fits files. Args: + mov (Movie): movie object fname (str): basename of output fits file mjd (int): MJD of saved movie @@ -271,6 +258,7 @@ def save_mov_txt(mov, fname, mjd=False): """Save movie data to series of text files. Args: + mov (Movie): movie object fname (str): basename of output text file mjd (int): MJD of saved movie @@ -296,6 +284,12 @@ def save_mov_txt(mov, fname, mjd=False): def save_array_txt(arr, fname): """Save the array data in a text file. + + Args: + arr (Array): array object + fname (str): name of output text file + + Returns: """ if type(arr) == np.ndarray: @@ -329,6 +323,12 @@ def save_array_txt(arr, fname): ################################################################################################## def save_obs_txt(obs, fname): """Save the observation data in a text file. + + Args: + obs (Obsdata): obsdata object + fname (str): name of output text file + + Returns: """ # output times must be in utc @@ -403,7 +403,15 @@ def save_obs_txt(obs, fname): def save_obs_uvfits(obs, fname=None, force_singlepol=None, polrep_out='circ'): """Save observation data to uvfits. - To save Stokes I as a single polarization (e.g., only RR) set force_singlepol='R' or 'L' + + Args: + obs (Obsdata): obsdata object + fname (str): path to output fits file, or None to return HDUList only + force_singlepol (str): if 'R' or 'L', will interpret stokes I field as 'RR' or 'LL' + polrep_out (str): 'circ' or 'stokes': how data should be stored in the uvfits file + Returns: + hdulist (astropy.io.fits.HDUList) + """ # output times must be in utc @@ -815,18 +823,22 @@ def save_obs_uvfits(obs, fname=None, force_singlepol=None, polrep_out='circ'): print("No NX table in saved uvfits") # Write final HDUList to file - if fname is None: - return hdulist_new.copy() - else: + if fname is not None: hdulist_new.writeto(fname, overwrite=True) - return None + + return hdulist_new.copy() + def save_obs_oifits(obs, fname, flux=1.0): - """ Save visibility data to oifits - Polarization data is NOT saved - Antenna diameter and exact times are currently incorrectt - Please contact Katie Bouman (klbouman@mit.edu) for any questions on this function + """Save visibility data to oifits file. + Polarization data is NOT saved + NOTE: as of 2021, this function is very out-of-date and should be updated + Args: + obs (Obsdata): obsdata object + fname (str): path to output uvfits file. + flux (float): Flux density normalization + Returns: """ # TODO: Add polarization to oifits?? @@ -919,7 +931,12 @@ def save_obs_oifits(obs, fname, flux=1.0): def save_dtype_txt(obs, fname, dtype='cphase'): - """Save the dtype data in a text file. + """Save the data product of type 'dtype' in a text file. + Args: + obs (Obsdata): obsdata object + fname (str): path to output text file + dtype (str): desired data type + Returns: """ head = ("SRC: %s \n" % obs.source + diff --git a/ehtim/obsdata.py b/ehtim/obsdata.py index a5eae423..0118141e 100644 --- a/ehtim/obsdata.py +++ b/ehtim/obsdata.py @@ -4480,10 +4480,10 @@ def save_txt(self, fname): return - def save_uvfits(self, fname=None, force_singlepol=False, polrep_out='circ'): + def save_uvfits(self, fname, force_singlepol=False, polrep_out='circ'): """Save visibility data to uvfits file. Args: - fname (str): path to output text file. If not specified, return HDUList. + fname (str): path to output uvfits file. force_singlepol (str): if 'R' or 'L', will interpret stokes I field as 'RR' or 'LL' polrep_out (str): 'circ' or 'stokes': how data should be stored in the uvfits file """ @@ -4492,13 +4492,28 @@ def save_uvfits(self, fname=None, force_singlepol=False, polrep_out='circ'): raise Exception( "force_singlepol is incompatible with polrep!='stokes'") - output = ehtim.io.save.save_obs_uvfits( - self, fname, force_singlepol=force_singlepol, polrep_out=polrep_out) + output = ehtim.io.save.save_obs_uvfits(self, fname, + force_singlepol=force_singlepol, polrep_out=polrep_out) + + return + + def make_hdulist(self, force_singlepol=False, polrep_out='circ'): + """Returns an hdulist in the same format as in a saved .uvfits file. + Args: + force_singlepol (str): if 'R' or 'L', will interpret stokes I field as 'RR' or 'LL' + polrep_out (str): 'circ' or 'stokes': how data should be stored in the uvfits file + Returns: + hdulist (astropy.io.fits.HDUList) + """ + + if (force_singlepol is not False) and (self.polrep != 'stokes'): + raise Exception( + "force_singlepol is incompatible with polrep!='stokes'") + + hdulist = ehtim.io.save.save_obs_uvfits(self, None, + force_singlepol=force_singlepol, polrep_out=polrep_out) + return hdulist - if fname is None: - return output - else: - return def save_oifits(self, fname, flux=1.0): """ Save visibility data to oifits. Polarization data is NOT saved. From 0cd10e7c77b1e32f77f05ee427d2c41578319f0f Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Tue, 16 Mar 2021 15:06:52 -0400 Subject: [PATCH 17/28] cleaned up docstrings in load.py --- ehtim/io/load.py | 252 +++++++++++++++++++++++++++-------------------- 1 file changed, 143 insertions(+), 109 deletions(-) diff --git a/ehtim/io/load.py b/ehtim/io/load.py index f29ed39e..a7e2896f 100644 --- a/ehtim/io/load.py +++ b/ehtim/io/load.py @@ -52,9 +52,14 @@ def load_vex(fname): - """Read in .vex files. and function to observe them + """Read in .vex files. Assumes there is only 1 MODE in vex file Hotaka Shiokawa - 2017 + + Args: + fname (str): path to input .vex file + Returns: + vex (Vex): Vex file object """ print("Loading vexfile: ", fname) return ehtim.vex.Vex(fname) @@ -780,9 +785,17 @@ def load_movie_dat(basename, nframes, startframe=0, framedur_sec=1, psize=-1, ################################################################################################### def load_array_txt(filename, ephemdir='ephemeris'): """Read an array from a text file and return an Array object - Sites with x=y=z=0 are spacecraft - 2TLE ephemeris loaded from ephemdir + Sites with x=y=z=0 are spacecraft - TLE ephemeris loaded from ephemdir + + Args: + filename (str): path to input text file + ephemdir (str): directory with TLE files for spacecraft + + Returns: + arr (Array): Array object loaded from file """ + tdata = np.loadtxt(filename, dtype=bytes, comments='#').astype(str) if tdata[0][0].lower() == 'site': tdata = tdata[1:] @@ -965,121 +978,22 @@ def load_obs_txt(filename, polrep='stokes'): out = out.switch_polrep(polrep_out=polrep) return out - -def load_obs_maps(arrfile, obsspec, ifile, qfile=0, ufile=0, vfile=0, - src=ehc.SOURCE_DEFAULT, mjd=ehc.MJD_DEFAULT, ampcal=False, phasecal=False): - """Read an observation from a maps text file and return an Obsdata object - """ - # Read telescope parameters from the array file - tdata = np.loadtxt(arrfile, dtype=bytes).astype(str) - tdata = [np.array((x[0], float(x[1]), float(x[2]), float(x[3]), - float(x[-1]), float(x[-1]), 0., 0., 0., 0., 0.), - dtype=ehc.DTARR) for x in tdata] - tdata = np.array(tdata) - - # Read parameters from the obs_spec - f = open(obsspec) - stop = False - while not stop: - line = f.readline().split() - if line == [] or line[0] == '\\': - continue - elif line[0] == 'FOV_center_RA': - x = line[2].split(':') - ra = float(x[0]) + float(x[1]) / 60.0 + float(x[2]) / 3600.0 - elif line[0] == 'FOV_center_Dec': - x = line[2].split(':') - dec = np.sign(float(x[0])) * (abs(float(x[0])) + - float(x[1]) / 60.0 + float(x[2]) / 3600.0) - elif line[0] == 'Corr_int_time': - tint = float(line[2]) - elif line[0] == 'Corr_chan_bw': # TODO what if multiple channels? - bw = float(line[2]) * 1e6 # in MHz - elif line[0] == 'Channel': # TODO what if multiple scans with different params? - rf = float(line[2].split(':')[0]) * 1e6 - elif line[0] == 'Scan_start': - x = line[2].split(':') # TODO properly compute MJD! - elif line[0] == 'Endscan': - stop = True - f.close() - - # Load the data, convert to list format, return object - datatable = [] - f = open(ifile) - - for line in f: - line = line.split() - if not (line[0] in ['UV', 'Scan', '\n']): - time = line[0].split(':') - time = float(time[2]) + float(time[3]) / 60.0 + float(time[4]) / 3600.0 - u = float(line[1]) * 1000 - v = float(line[2]) * 1000 - bl = line[4].split('-') - t1 = tdata[int(bl[0]) - 1]['site'] - t2 = tdata[int(bl[1]) - 1]['site'] - tau1 = 0. - tau2 = 0. - vis = float(line[7][:-1]) * np.exp(1j * float(line[8][:-1]) * ehc.DEGREE) - sigma = float(line[10]) - datatable.append(np.array((time, tint, t1, t2, tau1, tau2, - u, v, vis, 0.0, 0.0, 0.0, - sigma, 0.0, 0.0, 0.0), dtype=ehc.DTPOL_STOKES)) - - datatable = np.array(datatable) - - # TODO qfile ufile and vfile must have exactly the same format as ifile! - # add some consistency check - if not qfile == 0: - f = open(qfile) - i = 0 - for line in f: - line = line.split() - if not (line[0] in ['UV', 'Scan', '\n']): - datatable[i]['qvis'] = float(line[7][:-1]) * \ - np.exp(1j * float(line[8][:-1]) * ehc.DEGREE) - datatable[i]['qsigma'] = float(line[10]) - i += 1 - - if not ufile == 0: - f = open(ufile) - i = 0 - for line in f: - line = line.split() - if not (line[0] in ['UV', 'Scan', '\n']): - datatable[i]['uvis'] = float(line[7][:-1]) * \ - np.exp(1j * float(line[8][:-1]) * ehc.DEGREE) - datatable[i]['usigma'] = float(line[10]) - i += 1 - - if not vfile == 0: - f = open(vfile) - i = 0 - for line in f: - line = line.split() - if not (line[0] in ['UV', 'Scan', '\n']): - datatable[i]['vvis'] = float(line[7][:-1]) * \ - np.exp(1j * float(line[8][:-1]) * ehc.DEGREE) - datatable[i]['vsigma'] = float(line[10]) - i += 1 - - # Return the data object - return ehtim.obsdata.Obsdata(ra, dec, rf, bw, datatable, tdata, - source=src, mjd=mjd, polrep='stokes') - - # TODO can we save new telescope array terms and flags to uvfits and load them? -def load_obs_uvfits(filename, polrep='stokes', flipbl=False, allow_singlepol=True, - force_singlepol=None, channel=all, IF=all, remove_nan=False): +def load_obs_uvfits(filename, polrep='stokes', flipbl=False, + allow_singlepol=True, force_singlepol=None, + channel=all, IF=all, remove_nan=False): """Load observation data from a uvfits file. + Args: - fname (str or HDUList): path to input text file or HDUList object + filename (str or HDUList): path to either an input text file or an HDUList object polrep (str): load data as either 'stokes' or 'circ' flipbl (bool): flip baseline phases if True. allow_singlepol (bool): If True and polrep='stokes', treat single-polarization data as Stokes I force_singlepol (str): 'R' or 'L' to load only 1 polarization and treat as Stokes I channel (list): list of channels to average in the import. channel=all averages all - IF (list): list of IFs to average in the import. IF=all averages all IFS + IF (list): list of IFs to average in the import. IF=all averages all + remove_nan (bool): whether or not to remove entries with nan data Returns: obs (Obsdata): Obsdata object loaded from file """ @@ -1609,9 +1523,129 @@ def load_obs_oifits(filename, flux=1.0): return ehtim.obsdata.Obsdata(ra, dec, rf, bw, datatable, tarr, polrep='stokes', source=src, mjd=time[0]) +def load_obs_maps(arrfile, obsspec, ifile, qfile=0, ufile=0, vfile=0, + src=ehc.SOURCE_DEFAULT, mjd=ehc.MJD_DEFAULT, ampcal=False, phasecal=False): + """Read an observation from a maps text file and return an Obsdata object. + + Args: + arrfile (str): path to input array file + obsspec (str): path to input obs spec file + ifile (str): path to input Stokes I data file + qfile (str): path to input Stokes Q data file + ufile (str): path to input Stokes U data file + vfile (str): path to input Stokes V data file + src (str): source name + mjd (int): integer observation MJD + ampcal (bool): True if amplitude calibrated + phasecal (bool): True if phase calibrated + + Returns: + obs (Obsdata): Obsdata object loaded from file + """ + # Read telescope parameters from the array file + tdata = np.loadtxt(arrfile, dtype=bytes).astype(str) + tdata = [np.array((x[0], float(x[1]), float(x[2]), float(x[3]), + float(x[-1]), float(x[-1]), 0., 0., 0., 0., 0.), + dtype=ehc.DTARR) for x in tdata] + tdata = np.array(tdata) + + # Read parameters from the obs_spec + f = open(obsspec) + stop = False + while not stop: + line = f.readline().split() + if line == [] or line[0] == '\\': + continue + elif line[0] == 'FOV_center_RA': + x = line[2].split(':') + ra = float(x[0]) + float(x[1]) / 60.0 + float(x[2]) / 3600.0 + elif line[0] == 'FOV_center_Dec': + x = line[2].split(':') + dec = np.sign(float(x[0])) * (abs(float(x[0])) + + float(x[1]) / 60.0 + float(x[2]) / 3600.0) + elif line[0] == 'Corr_int_time': + tint = float(line[2]) + elif line[0] == 'Corr_chan_bw': # TODO what if multiple channels? + bw = float(line[2]) * 1e6 # in MHz + elif line[0] == 'Channel': # TODO what if multiple scans with different params? + rf = float(line[2].split(':')[0]) * 1e6 + elif line[0] == 'Scan_start': + x = line[2].split(':') # TODO properly compute MJD! + elif line[0] == 'Endscan': + stop = True + f.close() + + # Load the data, convert to list format, return object + datatable = [] + f = open(ifile) + + for line in f: + line = line.split() + if not (line[0] in ['UV', 'Scan', '\n']): + time = line[0].split(':') + time = float(time[2]) + float(time[3]) / 60.0 + float(time[4]) / 3600.0 + u = float(line[1]) * 1000 + v = float(line[2]) * 1000 + bl = line[4].split('-') + t1 = tdata[int(bl[0]) - 1]['site'] + t2 = tdata[int(bl[1]) - 1]['site'] + tau1 = 0. + tau2 = 0. + vis = float(line[7][:-1]) * np.exp(1j * float(line[8][:-1]) * ehc.DEGREE) + sigma = float(line[10]) + datatable.append(np.array((time, tint, t1, t2, tau1, tau2, + u, v, vis, 0.0, 0.0, 0.0, + sigma, 0.0, 0.0, 0.0), dtype=ehc.DTPOL_STOKES)) + + datatable = np.array(datatable) + + # TODO qfile ufile and vfile must have exactly the same format as ifile! + # add some consistency check + if not qfile == 0: + f = open(qfile) + i = 0 + for line in f: + line = line.split() + if not (line[0] in ['UV', 'Scan', '\n']): + datatable[i]['qvis'] = float(line[7][:-1]) * \ + np.exp(1j * float(line[8][:-1]) * ehc.DEGREE) + datatable[i]['qsigma'] = float(line[10]) + i += 1 + + if not ufile == 0: + f = open(ufile) + i = 0 + for line in f: + line = line.split() + if not (line[0] in ['UV', 'Scan', '\n']): + datatable[i]['uvis'] = float(line[7][:-1]) * \ + np.exp(1j * float(line[8][:-1]) * ehc.DEGREE) + datatable[i]['usigma'] = float(line[10]) + i += 1 + + if not vfile == 0: + f = open(vfile) + i = 0 + for line in f: + line = line.split() + if not (line[0] in ['UV', 'Scan', '\n']): + datatable[i]['vvis'] = float(line[7][:-1]) * \ + np.exp(1j * float(line[8][:-1]) * ehc.DEGREE) + datatable[i]['vsigma'] = float(line[10]) + i += 1 + + # Return the data object + return ehtim.obsdata.Obsdata(ra, dec, rf, bw, datatable, tdata, + source=src, mjd=mjd, polrep='stokes') def load_dtype_txt(obs, filename, dtype='cphase'): - """Load the dtype data in a text file and put it in the obs + + """Load the dtype data in a text file and put it in the already-created obs object + Args: + obs (Obsdata): obsdata object + filename (str): path to output text file + dtype (str): desired data type + Returns: """ print("Loading text observation: ", filename) From 47aa3c0d395140727dcda7c62d6c349140bed950 Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Tue, 16 Mar 2021 21:48:33 -0400 Subject: [PATCH 18/28] update save.py fix second instance of astropy.Time(out_subformat='date') bug --- ehtim/io/save.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ehtim/io/save.py b/ehtim/io/save.py index a7903e9f..9f9dd2a3 100644 --- a/ehtim/io/save.py +++ b/ehtim/io/save.py @@ -680,8 +680,12 @@ def save_obs_uvfits(obs, fname=None, force_singlepol=None, polrep_out='circ'): head['ARRAYZ'] = 0.e0 # TODO change the reference date - rdate_tt_new = Time(obs.mjd + MJD_0, format='jd', scale='utc', out_subfmt='date') - rdate_out = rdate_tt_new.iso + #rdate_tt_new = Time(obs.mjd + MJD_0, format='jd', scale='utc', out_subfmt='date') + #rdate_out = rdate_tt_new.iso + + rdate_tt_new = Time(obs.mjd + MJD_0, format='jd', scale='utc') + rdate_out = rdate_tt_new.iso[0:10] + rdate_tt_new.out_subfmt = 'float' # TODO -- needed to fix subformat issue in astropy 4.0 rdate_jd_out = rdate_tt_new.jd rdate_gstiao_out = rdate_tt_new.sidereal_time('apparent', 'greenwich').degree From e8c37ee5313f5999ce190ef0b969d60482902e1f Mon Sep 17 00:00:00 2001 From: Aviad Date: Thu, 3 Jun 2021 19:06:05 -0700 Subject: [PATCH 19/28] adding user defined ea_ker to Scatter_Movie Extending the support for user defined scattering kernel that exists for images to movies. --- ehtim/scattering/stochastic_optics.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ehtim/scattering/stochastic_optics.py b/ehtim/scattering/stochastic_optics.py index d336cf7a..b2520474 100644 --- a/ehtim/scattering/stochastic_optics.py +++ b/ehtim/scattering/stochastic_optics.py @@ -529,7 +529,7 @@ def Scatter(self, Unscattered_Image, Epsilon_Screen=np.array([]), obs_frequency_ return AI_Image - def Scatter_Movie(self, Unscattered_Movie, Epsilon_Screen=np.array([]), obs_frequency_Hz=0.0, Vx_km_per_s=50.0, Vy_km_per_s=0.0, framedur_sec=None, N_frames = None, sqrtQ=None, Linearized_Approximation=False, Force_Positivity=False, Return_Image_List=False, processes=0): + def Scatter_Movie(self, Unscattered_Movie, Epsilon_Screen=np.array([]), obs_frequency_Hz=0.0, Vx_km_per_s=50.0, Vy_km_per_s=0.0, framedur_sec=None, N_frames = None, ea_ker=None, sqrtQ=None, Linearized_Approximation=False, Force_Positivity=False, Return_Image_List=False, processes=0): """Scatter a movie using the specified epsilon screen. The movie can either be a movie object, an image list, or a static image If scattering a list of images or static image, the frame duration in seconds (framedur_sec) must be specified If scattering a static image, the total number of frames must be specified (N_frames) @@ -545,6 +545,7 @@ def Scatter_Movie(self, Unscattered_Movie, Epsilon_Screen=np.array([]), obs_freq Vy_km_per_s (float): Velocity of the scattering screen in the y direction (toward North) in km/s. framedur_sec (float): Duration of each frame, in seconds. Only needed if Unscattered_Movie is not a movie object. N_frames (int): Total number of frames. Only needed if Unscattered_Movie is a static image object. + ea_ker (2D ndarray): The used can optionally pass a precomputed array of the ensemble-average blurring kernel. sqrtQ (2D ndarray): The used can optionally pass a precomputed array of the square root of the power spectrum. Linearized_Approximation (bool): If True, uses a linearized approximation for the scattering (Eq. 10 of Johnson & Narayan 2016). If False, uses Eq. 9 of that paper. Force_Positivity (bool): If True, eliminates negative flux from the scattered image from the linearized approximation. @@ -649,7 +650,7 @@ def get_frame(j): pool.close() pool.join() else: - scattered_im_List = [self.Scatter(get_frame(j), Epsilon_Screen, obs_frequency_Hz = obs_frequency_Hz, Vx_km_per_s = Vx_km_per_s, Vy_km_per_s = Vy_km_per_s, t_hr=tlist_hr[j], sqrtQ=sqrtQ, Linearized_Approximation=Linearized_Approximation, Force_Positivity=Force_Positivity) for j in range(N_frams)] + scattered_im_List = [self.Scatter(get_frame(j), Epsilon_Screen, obs_frequency_Hz = obs_frequency_Hz, Vx_km_per_s = Vx_km_per_s, Vy_km_per_s = Vy_km_per_s, t_hr=tlist_hr[j], ea_ker=ea_ker, sqrtQ=sqrtQ, Linearized_Approximation=Linearized_Approximation, Force_Positivity=Force_Positivity) for j in range(N_frams)] if Return_Image_List == True: return scattered_im_List From 86f31a5a01eaf2e24c6ebbadfd61ad71759ee3e1 Mon Sep 17 00:00:00 2001 From: michaeldjohnson Date: Fri, 18 Jun 2021 13:33:35 -0400 Subject: [PATCH 20/28] Added model object, model fitting, and example script. --- ehtim/__init__.py | 2 + ehtim/model.py | 2063 +++++++++++++++++++++++ ehtim/modeling/__init__.py | 9 + ehtim/modeling/modeling_utils.py | 2660 ++++++++++++++++++++++++++++++ examples/example_modeling.py | 94 ++ setup.py | 1 + 6 files changed, 4829 insertions(+) create mode 100644 ehtim/model.py create mode 100644 ehtim/modeling/__init__.py create mode 100644 ehtim/modeling/modeling_utils.py create mode 100644 examples/example_modeling.py diff --git a/ehtim/__init__.py b/ehtim/__init__.py index dd963114..b60f0b8e 100644 --- a/ehtim/__init__.py +++ b/ehtim/__init__.py @@ -16,6 +16,7 @@ import ehtim.observing from ehtim.const_def import * from ehtim.imaging.imager_utils import imager_func +from ehtim.modeling.modeling_utils import modeler_func import ehtim.imaging from ehtim.features import rex import ehtim.features @@ -40,6 +41,7 @@ import ehtim.array import ehtim.movie import ehtim.image +import ehtim.model import warnings diff --git a/ehtim/model.py b/ehtim/model.py new file mode 100644 index 00000000..5f983095 --- /dev/null +++ b/ehtim/model.py @@ -0,0 +1,2063 @@ +# model.py +# an interferometric model class + +from __future__ import division +from __future__ import print_function +from builtins import str +from builtins import range +from builtins import object + +import numpy as np +import scipy.special as sps +import scipy.integrate as integrate +import scipy.interpolate as interpolate +import copy + +import ehtim.observing.obs_simulate as simobs +import ehtim.observing.pulses + +from ehtim.const_def import * +from ehtim.observing.obs_helpers import * +#from ehtim.modeling.modeling_utils import * + +import ehtim.image as image + +from ehtim.const_def import * + +LINE_THICKNESS = 2 # Thickness of 1D models on the image, in pixels +FOV_DEFAULT = 100.*RADPERUAS +NPIX_DEFAULT = 256 +COMPLEX_BASIS = 'abs-arg' # Basis for representing (most) complex quantities: 'abs-arg' or 're-im' + +########################################################################################################################################### +#Model object +########################################################################################################################################### + +def model_params(model_type, model_params=None, fit_pol=False, fit_cpol=False): + """Return the ordered list of model parameters for a specified model type. This order must match that of the gradient function, sample_1model_grad_uv. + """ + + if COMPLEX_BASIS == 're-im': + complex_labels = ['_re','_im'] + elif COMPLEX_BASIS == 'abs-arg': + complex_labels = ['_abs','_arg'] + else: + raise Exception('COMPLEX_BASIS ' + COMPLEX_BASIS + ' not recognized!') + + params = [] + + # Function to add polarimetric parameters; these must be added before stretch parameters + def add_pol(): + if fit_pol: + if model_type.find('mring') == -1: + params.append('pol_frac') + params.append('pol_evpa') + else: + for j in range(-(len(model_params['beta_list_pol'])-1)//2,(len(model_params['beta_list_pol'])+1)//2): + params.append('betapol' + str(j) + complex_labels[0]) + params.append('betapol' + str(j) + complex_labels[1]) + if fit_cpol: + if model_type.find('mring') == -1: + params.append('cpol_frac') + else: + for j in range(len(model_params['beta_list_cpol'])): + params.append('betacpol' + str(j) + complex_labels[0]) + params.append('betacpol' + str(j) + complex_labels[1]) + + if model_type == 'point': + params = ['F0','x0','y0'] + add_pol() + elif model_type == 'circ_gauss': + params = ['F0','FWHM','x0','y0'] + add_pol() + elif model_type == 'gauss': + params = ['F0','FWHM_maj','FWHM_min','PA','x0','y0'] + add_pol() + elif model_type == 'disk': + params = ['F0','d','x0','y0'] + add_pol() + elif model_type == 'blurred_disk': + params = ['F0','d','alpha','x0','y0'] + add_pol() + elif model_type == 'ring': + params = ['F0','d','x0','y0'] + add_pol() + elif model_type == 'stretched_ring': + params = ['F0','d','x0','y0','stretch','stretch_PA'] + add_pol() + elif model_type == 'thick_ring': + params = ['F0','d','alpha','x0','y0'] + add_pol() + elif model_type == 'stretched_thick_ring': + params = ['F0','d','alpha','x0','y0','stretch','stretch_PA'] + add_pol() + elif model_type == 'mring': + params = ['F0','d','x0','y0'] + for j in range(len(model_params['beta_list'])): + params.append('beta' + str(j+1) + complex_labels[0]) + params.append('beta' + str(j+1) + complex_labels[1]) + add_pol() + elif model_type == 'stretched_mring': + params = ['F0','d','x0','y0'] + for j in range(len(model_params['beta_list'])): + params.append('beta' + str(j+1) + complex_labels[0]) + params.append('beta' + str(j+1) + complex_labels[1]) + add_pol() + params.append('stretch') + params.append('stretch_PA') + elif model_type == 'thick_mring': + params = ['F0','d','alpha','x0','y0'] + for j in range(len(model_params['beta_list'])): + params.append('beta' + str(j+1) + complex_labels[0]) + params.append('beta' + str(j+1) + complex_labels[1]) + add_pol() + elif model_type == 'thick_mring_floor': + params = ['F0','d','alpha','ff','x0','y0'] + for j in range(len(model_params['beta_list'])): + params.append('beta' + str(j+1) + complex_labels[0]) + params.append('beta' + str(j+1) + complex_labels[1]) + add_pol() + elif model_type == 'stretched_thick_mring': + params = ['F0','d','alpha', 'x0','y0'] + for j in range(len(model_params['beta_list'])): + params.append('beta' + str(j+1) + complex_labels[0]) + params.append('beta' + str(j+1) + complex_labels[1]) + add_pol() + params.append('stretch') + params.append('stretch_PA') + elif model_type == 'stretched_thick_mring_floor': + params = ['F0','d','alpha','ff', 'x0','y0'] + for j in range(len(model_params['beta_list'])): + params.append('beta' + str(j+1) + complex_labels[0]) + params.append('beta' + str(j+1) + complex_labels[1]) + add_pol() + params.append('stretch') + params.append('stretch_PA') + else: + print('Model ' + model_init.models[j] + ' not recognized.') + params = [] + + return params + +def default_prior(model_type,model_params=None,fit_pol=False,fit_cpol=False): + """Return the default model prior and transformation for a specified model type + """ + + if COMPLEX_BASIS == 're-im': + complex_labels = ['_re','_im'] + complex_priors = [{'prior_type':'flat','min':-0.5,'max':0.5}, {'prior_type':'flat','min':-0.5,'max':0.5}] + complex_priors2 = [{'prior_type':'flat','min':-1,'max':1}, {'prior_type':'flat','min':-1,'max':1}] + elif COMPLEX_BASIS == 'abs-arg': + complex_labels = ['_abs','_arg'] + # Note: angle range here must match np.angle(). Need to properly define wrapped distributions + complex_priors = [{'prior_type':'flat','min':0.0,'max':0.5}, {'prior_type':'flat','min':-np.pi, 'max':np.pi}] + complex_priors2 = [{'prior_type':'flat','min':0.0,'max':1.0}, {'prior_type':'flat','min':-np.pi, 'max':np.pi}] + else: + raise Exception('COMPLEX_BASIS ' + COMPLEX_BASIS + ' not recognized!') + + prior = {'F0':{'prior_type':'none','transform':'log'}, + 'x0':{'prior_type':'none'}, + 'y0':{'prior_type':'none'}} + if model_type == 'point': + pass + elif model_type == 'circ_gauss': + prior['FWHM'] = {'prior_type':'none','transform':'log'} + elif model_type == 'gauss': + prior['FWHM_maj'] = {'prior_type':'positive','transform':'log'} + prior['FWHM_min'] = {'prior_type':'positive','transform':'log'} + prior['PA'] = {'prior_type':'none'} + elif model_type == 'disk': + prior['d'] = {'prior_type':'positive','transform':'log'} + elif model_type == 'blurred_disk': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['alpha'] = {'prior_type':'positive','transform':'log'} + elif model_type == 'ring': + prior['d'] = {'prior_type':'positive','transform':'log'} + elif model_type == 'stretched_ring': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['stretch'] = {'prior_type':'positive','transform':'log'} + prior['stretch_PA'] = {'prior_type':'none'} + elif model_type == 'thick_ring': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['alpha'] = {'prior_type':'positive','transform':'log'} + elif model_type == 'stretched_thick_ring': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['alpha'] = {'prior_type':'positive','transform':'log'} + prior['stretch'] = {'prior_type':'positive','transform':'log'} + prior['stretch_PA'] = {'prior_type':'none'} + elif model_type == 'mring': + prior['d'] = {'prior_type':'positive','transform':'log'} + for j in range(len(model_params['beta_list'])): + prior['beta' + str(j+1) + complex_labels[0]] = complex_priors[0] + prior['beta' + str(j+1) + complex_labels[1]] = complex_priors[1] + elif model_type == 'stretched_mring': + prior['d'] = {'prior_type':'positive','transform':'log'} + for j in range(len(model_params['beta_list'])): + prior['beta' + str(j+1) + complex_labels[0]] = complex_priors[0] + prior['beta' + str(j+1) + complex_labels[1]] = complex_priors[1] + prior['stretch'] = {'prior_type':'positive','transform':'log'} + prior['stretch_PA'] = {'prior_type':'none'} + elif model_type == 'thick_mring': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['alpha'] = {'prior_type':'positive','transform':'log'} + for j in range(len(model_params['beta_list'])): + prior['beta' + str(j+1) + complex_labels[0]] = complex_priors[0] + prior['beta' + str(j+1) + complex_labels[1]] = complex_priors[1] + elif model_type == 'thick_mring_floor': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['alpha'] = {'prior_type':'positive','transform':'log'} + prior['ff'] = {'prior_type':'flat','min':0,'max':1} + for j in range(len(model_params['beta_list'])): + prior['beta' + str(j+1) + complex_labels[0]] = complex_priors[0] + prior['beta' + str(j+1) + complex_labels[1]] = complex_priors[1] + elif model_type == 'stretched_thick_mring': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['alpha'] = {'prior_type':'positive','transform':'log'} + for j in range(len(model_params['beta_list'])): + prior['beta' + str(j+1) + complex_labels[0]] = complex_priors[0] + prior['beta' + str(j+1) + complex_labels[1]] = complex_priors[1] + prior['stretch'] = {'prior_type':'positive','transform':'log'} + prior['stretch_PA'] = {'prior_type':'none'} + elif model_type == 'stretched_thick_mring_floor': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['alpha'] = {'prior_type':'positive','transform':'log'} + prior['ff'] = {'prior_type':'flat','min':0,'max':1} + for j in range(len(model_params['beta_list'])): + prior['beta' + str(j+1) + complex_labels[0]] = complex_priors[0] + prior['beta' + str(j+1) + complex_labels[1]] = complex_priors[1] + prior['stretch'] = {'prior_type':'positive','transform':'log'} + prior['stretch_PA'] = {'prior_type':'none'} + else: + print('Model not recognized!') + + if fit_pol: + if model_type.find('mring') == -1: + prior['pol_frac'] = {'prior_type':'flat','min':0.0,'max':1.0} + prior['pol_evpa'] = {'prior_type':'flat','min':0.0,'max':np.pi} + else: + for j in range(-(len(model_params['beta_list_pol'])-1)//2,(len(model_params['beta_list_pol'])+1)//2): + prior['betapol' + str(j) + complex_labels[0]] = complex_priors2[0] + prior['betapol' + str(j) + complex_labels[1]] = complex_priors2[1] + + if fit_cpol: + if model_type.find('mring') == -1: + prior['cpol_frac'] = {'prior_type':'flat','min':-1.0,'max':1.0} + else: + for j in range(len(model_params['beta_list_cpol'])): + prior['betacpol' + str(j) + complex_labels[0]] = complex_priors2[0] + prior['betacpol' + str(j) + complex_labels[1]] = complex_priors2[1] + + return prior + +def stretch_xy(x, y, params): + x_stretch = ((x - params['x0']) * (np.cos(params['stretch_PA'])**2 + np.sin(params['stretch_PA'])**2 / params['stretch']) + + (y - params['y0']) * np.cos(params['stretch_PA']) * np.sin(params['stretch_PA']) * (1.0/params['stretch'] - 1.0)) + y_stretch = ((y - params['y0']) * (np.cos(params['stretch_PA'])**2 / params['stretch'] + np.sin(params['stretch_PA'])**2) + + (x - params['x0']) * np.cos(params['stretch_PA']) * np.sin(params['stretch_PA']) * (1.0/params['stretch'] - 1.0)) + return (params['x0'] + x_stretch,params['y0'] + y_stretch) + +def stretch_uv(u, v, params): + u_stretch = (u * (np.cos(params['stretch_PA'])**2 + np.sin(params['stretch_PA'])**2 * params['stretch']) + + v * np.cos(params['stretch_PA']) * np.sin(params['stretch_PA']) * (params['stretch'] - 1.0)) + v_stretch = (v * (np.cos(params['stretch_PA'])**2 * params['stretch'] + np.sin(params['stretch_PA'])**2) + + u * np.cos(params['stretch_PA']) * np.sin(params['stretch_PA']) * (params['stretch'] - 1.0)) + return (u_stretch,v_stretch) + +def get_const_polfac(model_type, params, pol): + # Return the scaling factor for models with constant fractional polarization + + if model_type.find('mring') != -1: + # mring models have polarization information specified differently than a constant scaling factor + return 1.0 + + try: + if pol == 'I': + return 1.0 + elif pol == 'Q': + return params['pol_frac'] * np.cos(2.0 * params['pol_evpa']) + elif pol == 'U': + return params['pol_frac'] * np.sin(2.0 * params['pol_evpa']) + elif pol == 'V': + return params['cpol_frac'] + elif pol == 'P': + return params['pol_frac'] * np.exp(1j * 2.0 * params['pol_evpa']) + elif pol == 'RR': + return get_const_polfac(model_type, params, 'I') + get_const_polfac(model_type, params, 'V') + elif pol == 'RL': + return get_const_polfac(model_type, params, 'Q') + 1j*get_const_polfac(model_type, params, 'U') + elif pol == 'LR': + return get_const_polfac(model_type, params, 'Q') - 1j*get_const_polfac(model_type, params, 'U') + elif pol == 'LL': + return get_const_polfac(model_type, params, 'I') - get_const_polfac(model_type, params, 'V') + except Exception: + pass + + return 0.0 + +def sample_1model_xy(x, y, model_type, params, psize=1.*RADPERUAS, pol='I'): + if pol == 'Q': + return np.real(sample_1model_xy(x, y, model_type, params, psize=psize, pol='P')) + elif pol == 'U': + return np.imag(sample_1model_xy(x, y, model_type, params, psize=psize, pol='P')) + elif pol in ['I','V','P']: + pass + else: + raise Exception('Polarization ' + pol + ' not implemented!') + + if model_type == 'point': + val = params['F0'] * (np.abs( x - params['x0']) < psize/2.0) * (np.abs( y - params['y0']) < psize/2.0) + elif model_type == 'circ_gauss': + sigma = params['FWHM'] / (2. * np.sqrt(2. * np.log(2.))) + val = (params['F0']*psize**2 * 4.0 * np.log(2.)/(np.pi * params['FWHM']**2) * + np.exp(-((x - params['x0'])**2 + (y - params['y0'])**2)/(2*sigma**2))) + elif model_type == 'gauss': + sigma_maj = params['FWHM_maj'] / (2. * np.sqrt(2. * np.log(2.))) + sigma_min = params['FWHM_min'] / (2. * np.sqrt(2. * np.log(2.))) + cth = np.cos(params['PA']) + sth = np.sin(params['PA']) + val = (params['F0']*psize**2 * 4.0 * np.log(2.)/(np.pi * params['FWHM_maj'] * params['FWHM_min']) * + np.exp(-((y - params['y0'])*np.cos(params['PA']) + (x - params['x0'])*np.sin(params['PA']))**2/(2*sigma_maj**2) + + -((x - params['x0'])*np.cos(params['PA']) - (y - params['y0'])*np.sin(params['PA']))**2/(2*sigma_min**2))) + elif model_type == 'disk': + val = params['F0']*psize**2/(np.pi*params['d']**2/4.) * (np.sqrt((x - params['x0'])**2 + (y - params['y0'])**2) < params['d']/2.0) + elif model_type == 'blurred_disk': + # Note: the exact form of a blurred disk requires numeric integration + + # This is the peak brightness of the blurred disk + I_peak = 4.0/(np.pi*params['d']**2) * (1.0 - 2.0**(-params['d']**2/params['alpha']**2)) + + # Constant prefactor + prefac = 32.0 * np.log(2.0)/(np.pi * params['alpha']**2 * params['d']**2) + + def f(r): + return integrate.quad(lambda rp: + prefac * rp * np.exp( -4.0 * np.log(2.0)/params['alpha']**2 * (r**2 + rp**2 - 2.0*r * rp) ) + * sps.ive(0, 8.0*np.log(2.0) * r * rp/params['alpha']**2), + 0, params['d']/2.0, limit=1000, epsabs=I_peak/1e9, epsrel=1.0e-6)[0] + f=np.vectorize(f) + r = np.sqrt((x - params['x0'])**2 + (y - params['y0'])**2) + + # For images, it's much quicker to do the 1-D problem and interpolate + if np.ndim(r) > 0: + r_min = np.min(r) + r_max = np.max(r) + r_list = np.linspace(r_min, r_max, int((r_max-r_min)/(params['alpha']) * 20)) + if len(r_list) < len(np.ravel(r))/2 and len(r) > 100: + f = interpolate.interp1d(r_list, f(r_list), kind='cubic') + val = params['F0'] * psize**2 * f(r) + elif model_type == 'ring': + val = (params['F0']*psize**2/(np.pi*params['d']*psize*LINE_THICKNESS) + * (params['d']/2.0 - psize*LINE_THICKNESS/2 < np.sqrt((x - params['x0'])**2 + (y - params['y0'])**2)) + * (params['d']/2.0 + psize*LINE_THICKNESS/2 > np.sqrt((x - params['x0'])**2 + (y - params['y0'])**2))) + elif model_type == 'thick_ring': + r = np.sqrt((x - params['x0'])**2 + (y - params['y0'])**2) + z = 4.*np.log(2.) * r * params['d']/params['alpha']**2 + val = (params['F0']*psize**2 * 4.0 * np.log(2.)/(np.pi * params['alpha']**2) + * np.exp(-4.*np.log(2.)/params['alpha']**2*(r**2 + params['d']**2/4.) + z) + * sps.ive(0, z)) + elif model_type == 'mring': + phi = np.angle((y - params['y0']) + 1j*(x - params['x0'])) + if pol == 'I': + beta_factor = (1.0 + np.sum([2.*np.real(params['beta_list'][m-1] * np.exp(1j * m * phi)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + elif pol == 'V' and len(params['beta_list_cpol']) > 0: + beta_factor = params['beta_list_cpol'][0] + np.sum([2.*np.real(params['beta_list_cpol'][m] * np.exp(1j * m * phi)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + elif pol == 'P' and len(params['beta_list_pol']) > 0: + num_coeff = len(params['beta_list_pol']) + beta_factor = np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * np.exp(1j * m * phi) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + else: + beta_factor = 0.0 + + val = (params['F0']*psize**2/(np.pi*params['d']*psize*LINE_THICKNESS) + * beta_factor + * (params['d']/2.0 - psize*LINE_THICKNESS/2 < np.sqrt((x - params['x0'])**2 + (y - params['y0'])**2)) + * (params['d']/2.0 + psize*LINE_THICKNESS/2 > np.sqrt((x - params['x0'])**2 + (y - params['y0'])**2))) + elif model_type == 'thick_mring': + phi = np.angle((y - params['y0']) + 1j*(x - params['x0'])) + r = np.sqrt((x - params['x0'])**2 + (y - params['y0'])**2) + z = 4.*np.log(2.) * r * params['d']/params['alpha']**2 + if pol == 'I': + beta_factor = (sps.ive(0, z) + np.sum([2.*np.real(sps.ive(m, z) * params['beta_list'][m-1] * np.exp(1j * m * phi)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + elif pol == 'V' and len(params['beta_list_cpol']) > 0: + beta_factor = (sps.ive(0, z) * params['beta_list_cpol'][0] + np.sum([2.*np.real(sps.ive(m, z) * params['beta_list_cpol'][m] * np.exp(1j * m * phi)) for m in range(1,len(params['beta_list_cpol']))],axis=0)) + elif pol == 'P' and len(params['beta_list_pol']) > 0: + num_coeff = len(params['beta_list_pol']) + beta_factor = np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * sps.ive(m, z) * np.exp(1j * m * phi) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + else: + # Note: not all polarizations accounted for yet (need RR, RL, LR, LL; do these by calling for linear combinations of I, Q, U, V)! + beta_factor = 0.0 + + val = (params['F0']*psize**2 * 4.0 * np.log(2.)/(np.pi * params['alpha']**2) + * np.exp(-4.*np.log(2.)/params['alpha']**2*(r**2 + params['d']**2/4.) + z) + * beta_factor) + elif model_type == 'thick_mring_floor': + val = (1.0 - params['ff']) * sample_1model_xy(x, y, 'thick_mring', params, psize=psize, pol=pol) + val += params['ff'] * sample_1model_xy(x, y, 'blurred_disk', params, psize=psize, pol=pol) + elif model_type[:9] == 'stretched': + params_stretch = params.copy() + params_stretch['F0'] /= params['stretch'] + val = sample_1model_xy(*stretch_xy(x, y, params), model_type[10:], params_stretch, psize, pol=pol) + else: + print('Model ' + model_type + ' not recognized!') + val = 0.0 + return val * get_const_polfac(model_type, params, pol) + +def sample_1model_uv(u, v, model_type, params, pol='I', jonesdict=None): + if jonesdict is not None: + # Define the various lists + fr1 = jonesdict['fr1'] # Field rotation of site 1 + fr2 = jonesdict['fr2'] # Field rotation of site 2 + DR1 = jonesdict['DR1'] # Right leakage term of site 1 + DL1 = jonesdict['DL1'] # Left leakage term of site 1 + DR2 = np.conj(jonesdict['DR2']) # Right leakage term of site 2 + DL2 = np.conj(jonesdict['DL2']) # Left leakage term of site 2 + # Sample the model without leakage + RR = sample_1model_uv(u, v, model_type, params, pol='RR') + RL = sample_1model_uv(u, v, model_type, params, pol='RL') + LR = sample_1model_uv(u, v, model_type, params, pol='LR') + LL = sample_1model_uv(u, v, model_type, params, pol='LL') + # Apply the Jones matrices + RRp = RR + LR * DR1 * np.exp( 2j*fr1) + RL * DR2 * np.exp(-2j*fr2) + LL * DR1 * DR2 * np.exp( 2j*(fr1-fr2)) + RLp = RL + LL * DR1 * np.exp( 2j*fr1) + RR * DL2 * np.exp( 2j*fr2) + LR * DR1 * DL2 * np.exp( 2j*(fr1+fr2)) + LRp = LR + RR * DL1 * np.exp(-2j*fr1) + LL * DR2 * np.exp(-2j*fr2) + RL * DL1 * DR2 * np.exp(-2j*(fr1+fr2)) + LLp = LL + LR * DL2 * np.exp( 2j*fr2) + RL * DL1 * np.exp(-2j*fr1) + RR * DL1 * DL2 * np.exp(-2j*(fr1-fr2)) + # Return the specified polarization + if pol == 'RR': return RRp + elif pol == 'RL': return RLp + elif pol == 'LR': return LRp + elif pol == 'LL': return LLp + elif pol == 'I': return 0.5 * (RRp + LLp) + elif pol == 'Q': return 0.5 * (LRp + RLp) + elif pol == 'U': return 0.5j* (LRp - RLp) + elif pol == 'V': return 0.5 * (RRp - LLp) + elif pol == 'P': return RLp + else: + raise Exception('Polarization ' + pol + ' not recognized!') + + if pol == 'Q': + return 0.5 * (sample_1model_uv(u, v, model_type, params, pol='P') + np.conj(sample_1model_uv(-u, -v, model_type, params, pol='P'))) + elif pol == 'U': + return -0.5j * (sample_1model_uv(u, v, model_type, params, pol='P') - np.conj(sample_1model_uv(-u, -v, model_type, params, pol='P'))) + elif pol in ['I','V','P']: + pass + elif pol == 'RR': + return sample_1model_uv(u, v, model_type, params, pol='I') + sample_1model_uv(u, v, model_type, params, pol='V') + elif pol == 'LL': + return sample_1model_uv(u, v, model_type, params, pol='I') - sample_1model_uv(u, v, model_type, params, pol='V') + elif pol == 'RL': + return sample_1model_uv(u, v, model_type, params, pol='Q') + 1j*sample_1model_uv(u, v, model_type, params, pol='U') + elif pol == 'LR': + return sample_1model_uv(u, v, model_type, params, pol='Q') - 1j*sample_1model_uv(u, v, model_type, params, pol='U') + else: + raise Exception('Polarization ' + pol + ' not implemented!') + + if model_type == 'point': + val = params['F0'] * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + elif model_type == 'circ_gauss': + val = (params['F0'] + * np.exp(-np.pi**2/(4.*np.log(2.)) * (u**2 + v**2) * params['FWHM']**2) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + elif model_type == 'gauss': + u_maj = u*np.sin(params['PA']) + v*np.cos(params['PA']) + u_min = u*np.cos(params['PA']) - v*np.sin(params['PA']) + val = (params['F0'] + * np.exp(-np.pi**2/(4.*np.log(2.)) * ((u_maj * params['FWHM_maj'])**2 + (u_min * params['FWHM_min'])**2)) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + elif model_type == 'disk': + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + #Add a small offset to avoid issues with division by zero + z += (z == 0.0) * 1e-10 + val = (params['F0'] * 2.0/z * sps.jv(1, z) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + elif model_type == 'blurred_disk': + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + #Add a small offset to avoid issues with division by zero + z += (z == 0.0) * 1e-10 + val = (params['F0'] * 2.0/z * sps.jv(1, z) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.)))) + elif model_type == 'ring': + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + val = (params['F0'] * sps.jv(0, z) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + elif model_type == 'thick_ring': + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + val = (params['F0'] * sps.jv(0, z) + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + elif model_type == 'mring': + phi = np.angle(v + 1j*u) + # Flip the baseline sign to match eht-imaging conventions + phi += np.pi + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + + if pol == 'I': + beta_factor = (sps.jv(0, z) + + np.sum([params['beta_list'][m-1] * sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * sps.jv(-m, z) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + elif pol == 'V' and len(params['beta_list_cpol']) > 0: + beta_factor = (params['beta_list_cpol'][0] * sps.jv(0, z) + + np.sum([params['beta_list_cpol'][m] * sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * sps.jv(-m, z) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0)) + elif pol == 'P' and len(params['beta_list_pol']) > 0: + num_coeff = len(params['beta_list_pol']) + beta_factor = np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + else: + beta_factor = 0.0 + + val = params['F0'] * beta_factor * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + elif model_type == 'thick_mring': + phi = np.angle(v + 1j*u) + # Flip the baseline sign to match eht-imaging conventions + phi += np.pi + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + + if pol == 'I': + beta_factor = (sps.jv(0, z) + + np.sum([params['beta_list'][m-1] * sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * sps.jv(-m, z) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + elif pol == 'V' and len(params['beta_list_cpol']) > 0: + beta_factor = (params['beta_list_cpol'][0] * sps.jv(0, z) + + np.sum([params['beta_list_cpol'][m] * sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * sps.jv(-m, z) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0)) + elif pol == 'P' and len(params['beta_list_pol']) > 0: + num_coeff = len(params['beta_list_pol']) + beta_factor = np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + else: + beta_factor = 0.0 + + val = (params['F0'] * beta_factor + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + elif model_type == 'thick_mring_floor': + val = (1.0 - params['ff']) * sample_1model_uv(u, v, 'thick_mring', params, pol=pol, jonesdict=jonesdict) + val += params['ff'] * sample_1model_uv(u, v, 'blurred_disk', params, pol=pol, jonesdict=jonesdict) + elif model_type[:9] == 'stretched': + params_stretch = params.copy() + params_stretch['x0'] = 0.0 + params_stretch['y0'] = 0.0 + val = sample_1model_uv(*stretch_uv(u,v,params), model_type[10:], params_stretch, pol=pol) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + else: + print('Model ' + model_type + ' not recognized!') + val = 0.0 + return val * get_const_polfac(model_type, params, pol) + +def sample_1model_graduv_uv(u, v, model_type, params, pol='I', jonesdict=None): + # Gradient of the visibility function, (dV/du, dV/dv) + # This function makes it convenient to, e.g., compute gradients of stretched images and to compute the model centroid + + if jonesdict is not None: + # Define the various lists + fr1 = jonesdict['fr1'] # Field rotation of site 1 + fr2 = jonesdict['fr2'] # Field rotation of site 2 + DR1 = jonesdict['DR1'] # Right leakage term of site 1 + DL1 = jonesdict['DL1'] # Left leakage term of site 1 + DR2 = np.conj(jonesdict['DR2']) # Right leakage term of site 2 + DL2 = np.conj(jonesdict['DL2']) # Left leakage term of site 2 + # Sample the model without leakage + RR = sample_1model_graduv_uv(u, v, model_type, params, pol='RR').reshape(2,len(u)) + RL = sample_1model_graduv_uv(u, v, model_type, params, pol='RL').reshape(2,len(u)) + LR = sample_1model_graduv_uv(u, v, model_type, params, pol='LR').reshape(2,len(u)) + LL = sample_1model_graduv_uv(u, v, model_type, params, pol='LL').reshape(2,len(u)) + # Apply the Jones matrices + RRp = (RR + LR * DR1 * np.exp( 2j*fr1) + RL * DR2 * np.exp(-2j*fr2) + LL * DR1 * DR2 * np.exp( 2j*(fr1-fr2))) + RLp = (RL + LL * DR1 * np.exp( 2j*fr1) + RR * DL2 * np.exp( 2j*fr2) + LR * DR1 * DL2 * np.exp( 2j*(fr1+fr2))) + LRp = (LR + RR * DL1 * np.exp(-2j*fr1) + LL * DR2 * np.exp(-2j*fr2) + RL * DL1 * DR2 * np.exp(-2j*(fr1+fr2))) + LLp = (LL + LR * DL2 * np.exp( 2j*fr2) + RL * DL1 * np.exp(-2j*fr1) + RR * DL1 * DL2 * np.exp(-2j*(fr1-fr2))) + # Return the specified polarization + if pol == 'RR': return RRp + elif pol == 'RL': return RLp + elif pol == 'LR': return LRp + elif pol == 'LL': return LLp + elif pol == 'I': return 0.5 * (RRp + LLp) + elif pol == 'Q': return 0.5 * (LRp + RLp) + elif pol == 'U': return 0.5j* (LRp - RLp) + elif pol == 'V': return 0.5 * (RRp - LLp) + elif pol == 'P': return RLp + else: + raise Exception('Polarization ' + pol + ' not recognized!') + + if pol == 'Q': + return 0.5 * (sample_1model_graduv_uv(u, v, model_type, params, pol='P') + np.conj(sample_1model_graduv_uv(-u, -v, model_type, params, pol='P'))) + elif pol == 'U': + return -0.5j * (sample_1model_graduv_uv(u, v, model_type, params, pol='P') - np.conj(sample_1model_graduv_uv(-u, -v, model_type, params, pol='P'))) + elif pol in ['I','V','P']: + pass + elif pol == 'RR': + return sample_1model_graduv_uv(u, v, model_type, params, pol='I') + sample_1model_graduv_uv(u, v, model_type, params, pol='V') + elif pol == 'LL': + return sample_1model_graduv_uv(u, v, model_type, params, pol='I') - sample_1model_graduv_uv(u, v, model_type, params, pol='V') + elif pol == 'RL': + return sample_1model_graduv_uv(u, v, model_type, params, pol='Q') + 1j*sample_1model_graduv_uv(u, v, model_type, params, pol='U') + elif pol == 'LR': + return sample_1model_graduv_uv(u, v, model_type, params, pol='Q') - 1j*sample_1model_graduv_uv(u, v, model_type, params, pol='U') + else: + raise Exception('Polarization ' + pol + ' not implemented!') + + vis = sample_1model_uv(u, v, model_type, params, jonesdict=jonesdict) + if model_type == 'point': + val = np.array([ 1j * 2.0 * np.pi * params['x0'] * vis, + 1j * 2.0 * np.pi * params['y0'] * vis]) + elif model_type == 'circ_gauss': + val = np.array([ (1j * 2.0 * np.pi * params['x0'] - params['FWHM']**2 * np.pi**2 * u/(2. * np.log(2.))) * vis, + (1j * 2.0 * np.pi * params['y0'] - params['FWHM']**2 * np.pi**2 * v/(2. * np.log(2.))) * vis]) + elif model_type == 'gauss': + u_maj = u*np.sin(params['PA']) + v*np.cos(params['PA']) + u_min = u*np.cos(params['PA']) - v*np.sin(params['PA']) + val = np.array([ (1j * 2.0 * np.pi * params['x0'] - params['FWHM_maj']**2 * np.pi**2 * u_maj/(2. * np.log(2.)) * np.sin(params['PA']) - params['FWHM_min']**2 * np.pi**2 * u_min/(2. * np.log(2.)) * np.cos(params['PA'])) * vis, + (1j * 2.0 * np.pi * params['y0'] - params['FWHM_maj']**2 * np.pi**2 * u_maj/(2. * np.log(2.)) * np.cos(params['PA']) + params['FWHM_min']**2 * np.pi**2 * u_min/(2. * np.log(2.)) * np.sin(params['PA'])) * vis]) + elif model_type == 'disk': + # Take care of the degenerate origin point by a small offset + #v += (u==0.)*(v==0.)*1e-10 + uvdist = (u**2 + v**2 + (u==0.)*(v==0.)*1e-10)**0.5 + z = np.pi * params['d'] * uvdist + bessel_deriv = 0.5 * (sps.jv( 0, z) - sps.jv( 2, z)) + val = np.array([ (1j * 2.0 * np.pi * params['x0'] - u/uvdist**2) * vis + + params['F0'] * 2./z * np.pi * params['d'] * u/uvdist * bessel_deriv * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + (1j * 2.0 * np.pi * params['y0'] - v/uvdist**2) * vis + + params['F0'] * 2./z * np.pi * params['d'] * v/uvdist * bessel_deriv * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) ]) + elif model_type == 'blurred_disk': + # Take care of the degenerate origin point by a small offset + #u += (u==0.)*(v==0.)*1e-10 + uvdist = (u**2 + v**2 + (u==0.)*(v==0.)*1e-10)**0.5 + z = np.pi * params['d'] * uvdist + blur = np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + bessel_deriv = 0.5 * (sps.jv( 0, z) - sps.jv( 2, z)) + val = np.array([ (1j * 2.0 * np.pi * params['x0'] - u/uvdist**2 - params['alpha']**2 * np.pi**2 * u/(2. * np.log(2.))) * vis + + params['F0'] * 2./z * np.pi * params['d'] * u/uvdist * bessel_deriv * blur * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + (1j * 2.0 * np.pi * params['y0'] - v/uvdist**2 - params['alpha']**2 * np.pi**2 * v/(2. * np.log(2.))) * vis + + params['F0'] * 2./z * np.pi * params['d'] * v/uvdist * bessel_deriv * blur * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) ]) + elif model_type == 'ring': + # Take care of the degenerate origin point by a small offset + u += (u==0.)*(v==0.)*1e-10 + uvdist = (u**2 + v**2 + (u==0.)*(v==0.)*1e-10)**0.5 + z = np.pi * params['d'] * uvdist + val = np.array([ 1j * 2.0 * np.pi * params['x0'] * vis + - params['F0'] * np.pi*params['d']*u/uvdist * sps.jv(1, z) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + 1j * 2.0 * np.pi * params['y0'] * vis + - params['F0'] * np.pi*params['d']*v/uvdist * sps.jv(1, z) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) ]) + elif model_type == 'thick_ring': + uvdist = (u**2 + v**2)**0.5 + #Add a small offset to avoid issues with division by zero + uvdist += (uvdist == 0.0) * 1e-10 + z = np.pi * params['d'] * uvdist + val = np.array([ (1j * 2.0 * np.pi * params['x0'] - params['alpha']**2 * np.pi**2 * u/(2. * np.log(2.))) * vis + - params['F0'] * np.pi*params['d']*u/uvdist * sps.jv(1, z) + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + (1j * 2.0 * np.pi * params['y0'] - params['alpha']**2 * np.pi**2 * v/(2. * np.log(2.))) * vis + - params['F0'] * np.pi*params['d']*v/uvdist * sps.jv(1, z) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) ]) + elif model_type == 'mring': + # Take care of the degenerate origin point by a small offset + u += (u==0.)*(v==0.)*1e-10 + phi = np.angle(v + 1j*u) + # Flip the baseline sign to match eht-imaging conventions + phi += np.pi + uvdist = (u**2 + v**2 + (u==0.)*(v==0.)*1e-10)**0.5 + dphidu = v/uvdist**2 + dphidv = -u/uvdist**2 + z = np.pi * params['d'] * uvdist + + if pol == 'I': + beta_factor_u = (-np.pi * params['d'] * u/uvdist * sps.jv(1, z) + + np.sum([params['beta_list'][m-1] * sps.jv( m, z) * ( 1j * m * dphidu) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * sps.jv(-m, z) * (-1j * m * dphidu) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([params['beta_list'][m-1] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * 0.5 * (sps.jv(-m-1, z) - sps.jv(-m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + beta_factor_v = (-np.pi * params['d'] * v/uvdist * sps.jv(1, z) + + np.sum([params['beta_list'][m-1] * sps.jv( m, z) * ( 1j * m * dphidv) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * sps.jv(-m, z) * (-1j * m * dphidv) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([params['beta_list'][m-1] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * 0.5 * (sps.jv(-m-1, z) - sps.jv(-m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + + elif pol == 'V' and len(params['beta_list_cpol']) > 0: + beta_factor_u = (-np.pi * params['d'] * u/uvdist * sps.jv(1, z) * params['beta_list_cpol'][0] + + np.sum([params['beta_list_cpol'][m] * sps.jv( m, z) * ( 1j * m * dphidu) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * sps.jv(-m, z) * (-1j * m * dphidu) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([params['beta_list_cpol'][m] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * 0.5 * (sps.jv(-m-1, z) - sps.jv(-m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0)) + beta_factor_v = (-np.pi * params['d'] * v/uvdist * sps.jv(1, z) + + np.sum([params['beta_list_cpol'][m] * sps.jv( m, z) * ( 1j * m * dphidv) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * sps.jv(-m, z) * (-1j * m * dphidv) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([params['beta_list_cpol'][m] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * 0.5 * (sps.jv(-m-1, z) - sps.jv(-m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0)) + elif pol == 'P' and len(params['beta_list_pol']) > 0: + num_coeff = len(params['beta_list_pol']) + beta_factor_u = ( + np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * sps.jv( m, z) * ( 1j * m * dphidu) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + + np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0)) + beta_factor_v = ( + np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * sps.jv( m, z) * ( 1j * m * dphidv) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + + np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0)) + else: + beta_factor_u = beta_factor_v = 0.0 + + val = np.array([ + 1j * 2.0 * np.pi * params['x0'] * vis + + params['F0'] * beta_factor_u + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + 1j * 2.0 * np.pi * params['y0'] * vis + + params['F0'] * beta_factor_v + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) ]) + elif model_type == 'thick_mring': + # Take care of the degenerate origin point by a small offset + u += (u==0.)*(v==0.)*1e-10 + phi = np.angle(v + 1j*u) + # Flip the baseline sign to match eht-imaging conventions + phi += np.pi + uvdist = (u**2 + v**2 + (u==0.)*(v==0.)*1e-10)**0.5 + dphidu = v/uvdist**2 + dphidv = -u/uvdist**2 + z = np.pi * params['d'] * uvdist + + if pol == 'I': + beta_factor_u = (-np.pi * params['d'] * u/uvdist * sps.jv(1, z) + + np.sum([params['beta_list'][m-1] * sps.jv( m, z) * ( 1j * m * dphidu) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * sps.jv(-m, z) * (-1j * m * dphidu) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([params['beta_list'][m-1] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * 0.5 * (sps.jv(-m-1, z) - sps.jv(-m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + beta_factor_v = (-np.pi * params['d'] * v/uvdist * sps.jv(1, z) + + np.sum([params['beta_list'][m-1] * sps.jv( m, z) * ( 1j * m * dphidv) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * sps.jv(-m, z) * (-1j * m * dphidv) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([params['beta_list'][m-1] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * 0.5 * (sps.jv(-m-1, z) - sps.jv(-m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + + elif pol == 'V' and len(params['beta_list_cpol']) > 0: + beta_factor_u = (-np.pi * params['d'] * u/uvdist * sps.jv(1, z) * params['beta_list_cpol'][0] + + np.sum([params['beta_list_cpol'][m] * sps.jv( m, z) * ( 1j * m * dphidu) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * sps.jv(-m, z) * (-1j * m * dphidu) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([params['beta_list_cpol'][m] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * 0.5 * (sps.jv(-m-1, z) - sps.jv(-m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0)) + beta_factor_v = (-np.pi * params['d'] * v/uvdist * sps.jv(1, z) + + np.sum([params['beta_list_cpol'][m] * sps.jv( m, z) * ( 1j * m * dphidv) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * sps.jv(-m, z) * (-1j * m * dphidv) * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([params['beta_list_cpol'][m] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * 0.5 * (sps.jv(-m-1, z) - sps.jv(-m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0)) + elif pol == 'P' and len(params['beta_list_pol']) > 0: + num_coeff = len(params['beta_list_pol']) + beta_factor_u = (0.0 + + np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * sps.jv( m, z) * ( 1j * m * dphidu) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + + np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * u/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0)) + beta_factor_v = (0.0 + + np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * sps.jv( m, z) * ( 1j * m * dphidv) * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + + np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * params['d'] * v/uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0)) + else: + beta_factor_u = beta_factor_v = 0.0 + + val = np.array([ + (1j * 2.0 * np.pi * params['x0'] - params['alpha']**2 * np.pi**2 * u/(2. * np.log(2.))) * vis + + params['F0'] * beta_factor_u + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + (1j * 2.0 * np.pi * params['y0'] - params['alpha']**2 * np.pi**2 * v/(2. * np.log(2.))) * vis + + params['F0'] * beta_factor_v + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) ]) + elif model_type == 'thick_mring_floor': + val = (1.0 - params['ff']) * sample_1model_graduv_uv(u, v, 'thick_mring', params, pol=pol, jonesdict=jonesdict) + val += params['ff'] * sample_1model_graduv_uv(u, v, 'blurred_disk', params, pol=pol, jonesdict=jonesdict) + elif model_type[:9] == 'stretched': + # Take care of the degenerate origin point by a small offset + u += (u==0.)*(v==0.)*1e-10 + params_stretch = params.copy() + params_stretch['x0'] = 0.0 + params_stretch['y0'] = 0.0 + (u_stretch, v_stretch) = stretch_uv(u,v,params) + + # First calculate the gradient of the unshifted but stretched image + grad0 = sample_1model_graduv_uv(u_stretch, v_stretch, model_type[10:], params_stretch, pol=pol) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + grad = grad0.copy() * 0.0 + grad[0] = ( grad0[0] * (np.cos(params['stretch_PA'])**2 + np.sin(params['stretch_PA'])**2*params['stretch']) + + grad0[1] * ((params['stretch'] - 1.0) * np.cos(params['stretch_PA']) * np.sin(params['stretch_PA']))) + grad[1] = ( grad0[1] * (np.cos(params['stretch_PA'])**2*params['stretch'] + np.sin(params['stretch_PA'])**2) + + grad0[0] * ((params['stretch'] - 1.0) * np.cos(params['stretch_PA']) * np.sin(params['stretch_PA']))) + + # Add the gradient term from the shift + vis = sample_1model_uv(u_stretch, v_stretch, model_type[10:], params_stretch, jonesdict=jonesdict) + grad[0] += vis * 1j * 2.0 * np.pi * params['x0'] * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + grad[1] += vis * 1j * 2.0 * np.pi * params['y0'] * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + + val = grad + else: + print('Model ' + model_type + ' not recognized!') + val = 0.0 + return val * get_const_polfac(model_type, params, pol) + +def sample_1model_grad_leakage_uv_re(u, v, model_type, params, pol, site, hand, jonesdict): + # Convenience function to calculate the gradient with respect to the real part of a specified site/hand leakage + + # Define the various lists + fr1 = jonesdict['fr1'] # Field rotation of site 1 + fr2 = jonesdict['fr2'] # Field rotation of site 2 + DR1 = jonesdict['DR1'] # Right leakage term of site 1 + DL1 = jonesdict['DL1'] # Left leakage term of site 1 + DR2 = np.conj(jonesdict['DR2']) # Right leakage term of site 2 + DL2 = np.conj(jonesdict['DL2']) # Left leakage term of site 2 + # Sample the model without leakage + RR = sample_1model_uv(u, v, model_type, params, pol='RR') + RL = sample_1model_uv(u, v, model_type, params, pol='RL') + LR = sample_1model_uv(u, v, model_type, params, pol='LR') + LL = sample_1model_uv(u, v, model_type, params, pol='LL') + + # Figure out which terms to include in the gradient + DR1mask = 0.0 + (hand == 'R') * (jonesdict['t1'] == site) + DR2mask = 0.0 + (hand == 'R') * (jonesdict['t2'] == site) + DL1mask = 0.0 + (hand == 'L') * (jonesdict['t1'] == site) + DL2mask = 0.0 + (hand == 'L') * (jonesdict['t2'] == site) + + # These are the leakage gradient terms + RRp = LR * DR1mask * np.exp( 2j*fr1) + RL * DR2mask * np.exp(-2j*fr2) + LL * DR1mask * DR2 * np.exp( 2j*(fr1-fr2)) + LL * DR1 * DR2mask * np.exp( 2j*(fr1-fr2)) + RLp = LL * DR1mask * np.exp( 2j*fr1) + RR * DL2mask * np.exp( 2j*fr2) + LR * DR1mask * DL2 * np.exp( 2j*(fr1+fr2)) + LR * DR1 * DL2mask * np.exp( 2j*(fr1+fr2)) + LRp = RR * DL1mask * np.exp(-2j*fr1) + LL * DR2mask * np.exp(-2j*fr2) + RL * DL1mask * DR2 * np.exp(-2j*(fr1+fr2)) + RL * DL1 * DR2mask * np.exp(-2j*(fr1+fr2)) + LLp = LR * DL2mask * np.exp( 2j*fr2) + RL * DL1mask * np.exp(-2j*fr1) + RR * DL1mask * DL2 * np.exp(-2j*(fr1-fr2)) + RR * DL1 * DL2mask * np.exp(-2j*(fr1-fr2)) + + # Return the specified polarization + if pol == 'RR': return RRp + elif pol == 'RL': return RLp + elif pol == 'LR': return LRp + elif pol == 'LL': return LLp + elif pol == 'I': return 0.5 * (RRp + LLp) + elif pol == 'Q': return 0.5 * (LRp + RLp) + elif pol == 'U': return 0.5j* (LRp - RLp) + elif pol == 'V': return 0.5 * (RRp - LLp) + elif pol == 'P': return RLp + else: + raise Exception('Polarization ' + pol + ' not recognized!') + +def sample_1model_grad_leakage_uv_im(u, v, model_type, params, pol, site, hand, jonesdict): + # Convenience function to calculate the gradient with respect to the imaginary part of a specified site/hand leakage + # The tricky thing here is the conjugation of the second leakage site, flipping the sign of the gradient + + # Define the various lists + fr1 = jonesdict['fr1'] # Field rotation of site 1 + fr2 = jonesdict['fr2'] # Field rotation of site 2 + DR1 = jonesdict['DR1'] # Right leakage term of site 1 + DL1 = jonesdict['DL1'] # Left leakage term of site 1 + DR2 = np.conj(jonesdict['DR2']) # Right leakage term of site 2 + DL2 = np.conj(jonesdict['DL2']) # Left leakage term of site 2 + # Sample the model without leakage + RR = sample_1model_uv(u, v, model_type, params, pol='RR') + RL = sample_1model_uv(u, v, model_type, params, pol='RL') + LR = sample_1model_uv(u, v, model_type, params, pol='LR') + LL = sample_1model_uv(u, v, model_type, params, pol='LL') + + # Figure out which terms to include in the gradient + DR1mask = 0.0 + (hand == 'R') * (jonesdict['t1'] == site) + DR2mask = 0.0 + (hand == 'R') * (jonesdict['t2'] == site) + DL1mask = 0.0 + (hand == 'L') * (jonesdict['t1'] == site) + DL2mask = 0.0 + (hand == 'L') * (jonesdict['t2'] == site) + + # These are the leakage gradient terms + RRp = 1j*( LR * DR1mask * np.exp( 2j*fr1) - RL * DR2mask * np.exp(-2j*fr2) + LL * DR1mask * DR2 * np.exp( 2j*(fr1-fr2)) - LL * DR1 * DR2mask * np.exp( 2j*(fr1-fr2))) + RLp = 1j*( LL * DR1mask * np.exp( 2j*fr1) - RR * DL2mask * np.exp( 2j*fr2) + LR * DR1mask * DL2 * np.exp( 2j*(fr1+fr2)) - LR * DR1 * DL2mask * np.exp( 2j*(fr1+fr2))) + LRp = 1j*( RR * DL1mask * np.exp(-2j*fr1) - LL * DR2mask * np.exp(-2j*fr2) + RL * DL1mask * DR2 * np.exp(-2j*(fr1+fr2)) - RL * DL1 * DR2mask * np.exp(-2j*(fr1+fr2))) + LLp = 1j*(-LR * DL2mask * np.exp( 2j*fr2) + RL * DL1mask * np.exp(-2j*fr1) + RR * DL1mask * DL2 * np.exp(-2j*(fr1-fr2)) - RR * DL1 * DL2mask * np.exp(-2j*(fr1-fr2))) + + # Return the specified polarization + if pol == 'RR': return RRp + elif pol == 'RL': return RLp + elif pol == 'LR': return LRp + elif pol == 'LL': return LLp + elif pol == 'I': return 0.5 * (RRp + LLp) + elif pol == 'Q': return 0.5 * (LRp + RLp) + elif pol == 'U': return 0.5j* (LRp - RLp) + elif pol == 'V': return 0.5 * (RRp - LLp) + elif pol == 'P': return RLp + else: + raise Exception('Polarization ' + pol + ' not recognized!') + +def sample_1model_grad_uv(u, v, model_type, params, pol='I', fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + # Gradient of the model for each model parameter + + if jonesdict is not None: + # Define the various lists + fr1 = jonesdict['fr1'] # Field rotation of site 1 + fr2 = jonesdict['fr2'] # Field rotation of site 2 + DR1 = jonesdict['DR1'] # Right leakage term of site 1 + DL1 = jonesdict['DL1'] # Left leakage term of site 1 + DR2 = np.conj(jonesdict['DR2']) # Right leakage term of site 2 + DL2 = np.conj(jonesdict['DL2']) # Left leakage term of site 2 + # Sample the gradients without leakage + RR = sample_1model_grad_uv(u, v, model_type, params, pol='RR', fit_pol=fit_pol, fit_cpol=fit_cpol) + RL = sample_1model_grad_uv(u, v, model_type, params, pol='RL', fit_pol=fit_pol, fit_cpol=fit_cpol) + LR = sample_1model_grad_uv(u, v, model_type, params, pol='LR', fit_pol=fit_pol, fit_cpol=fit_cpol) + LL = sample_1model_grad_uv(u, v, model_type, params, pol='LL', fit_pol=fit_pol, fit_cpol=fit_cpol) + # Apply the Jones matrices + RRp = (RR + LR * DR1 * np.exp( 2j*fr1) + RL * DR2 * np.exp(-2j*fr2) + LL * DR1 * DR2 * np.exp( 2j*(fr1-fr2))) + RLp = (RL + LL * DR1 * np.exp( 2j*fr1) + RR * DL2 * np.exp( 2j*fr2) + LR * DR1 * DL2 * np.exp( 2j*(fr1+fr2))) + LRp = (LR + RR * DL1 * np.exp(-2j*fr1) + LL * DR2 * np.exp(-2j*fr2) + RL * DL1 * DR2 * np.exp(-2j*(fr1+fr2))) + LLp = (LL + LR * DL2 * np.exp( 2j*fr2) + RL * DL1 * np.exp(-2j*fr1) + RR * DL1 * DL2 * np.exp(-2j*(fr1-fr2))) + # Return the specified polarization + if pol == 'RR': grad = RRp + elif pol == 'RL': grad = RLp + elif pol == 'LR': grad = LRp + elif pol == 'LL': grad = LLp + elif pol == 'I': grad = 0.5 * (RRp + LLp) + elif pol == 'Q': grad = 0.5 * (LRp + RLp) + elif pol == 'U': grad = 0.5j* (LRp - RLp) + elif pol == 'V': grad = 0.5 * (RRp - LLp) + elif pol == 'P': grad = RLp + else: + raise Exception('Polarization ' + pol + ' not recognized!') + # If necessary, add the gradient components from the leakage terms + # Each leakage term has two corresponding gradient terms: d/dRe and d/dIm. + if fit_leakage: + # 'leakage_fit' is a list of tuples [site, 'R' or 'L'] denoting the fitted leakage terms + for (site, hand) in jonesdict['leakage_fit']: + grad = np.vstack([grad, sample_1model_grad_leakage_uv_re(u, v, model_type, params, pol, site, hand, jonesdict), sample_1model_grad_leakage_uv_im(u, v, model_type, params, pol, site, hand, jonesdict)]) + + return grad + + if pol == 'Q': + return 0.5 * (sample_1model_grad_uv(u, v, model_type, params, pol='P', fit_pol=fit_pol, fit_cpol=fit_cpol) + np.conj(sample_1model_grad_uv(-u, -v, model_type, params, pol='P', fit_pol=fit_pol, fit_cpol=fit_cpol))) + elif pol == 'U': + return -0.5j * (sample_1model_grad_uv(u, v, model_type, params, pol='P', fit_pol=fit_pol, fit_cpol=fit_cpol) - np.conj(sample_1model_grad_uv(-u, -v, model_type, params, pol='P', fit_pol=fit_pol, fit_cpol=fit_cpol))) + elif pol in ['I','V','P']: + pass + elif pol == 'RR': + return sample_1model_grad_uv(u, v, model_type, params, pol='I', fit_pol=fit_pol, fit_cpol=fit_cpol) + sample_1model_grad_uv(u, v, model_type, params, pol='V', fit_pol=fit_pol, fit_cpol=fit_cpol) + elif pol == 'LL': + return sample_1model_grad_uv(u, v, model_type, params, pol='I', fit_pol=fit_pol, fit_cpol=fit_cpol) - sample_1model_grad_uv(u, v, model_type, params, pol='V', fit_pol=fit_pol, fit_cpol=fit_cpol) + elif pol == 'RL': + return sample_1model_grad_uv(u, v, model_type, params, pol='Q', fit_pol=fit_pol, fit_cpol=fit_cpol) + 1j*sample_1model_grad_uv(u, v, model_type, params, pol='U', fit_pol=fit_pol, fit_cpol=fit_cpol) + elif pol == 'LR': + return sample_1model_grad_uv(u, v, model_type, params, pol='Q', fit_pol=fit_pol, fit_cpol=fit_cpol) - 1j*sample_1model_grad_uv(u, v, model_type, params, pol='U', fit_pol=fit_pol, fit_cpol=fit_cpol) + else: + raise Exception('Polarization ' + pol + ' not implemented!') + + if model_type == 'point': # F0, x0, y0 + val = np.array([ np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + 1j * 2.0 * np.pi * u * params['F0'] * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + 1j * 2.0 * np.pi * v * params['F0'] * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))]) + elif model_type == 'circ_gauss': # F0, FWHM, x0, y0 + gauss = (params['F0'] * np.exp(-np.pi**2/(4.*np.log(2.)) * (u**2 + v**2) * params['FWHM']**2) + *np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + val = np.array([ 1.0/params['F0'] * gauss, + -np.pi**2/(2.*np.log(2.)) * (u**2 + v**2) * params['FWHM'] * gauss, + 1j * 2.0 * np.pi * u * gauss, + 1j * 2.0 * np.pi * v * gauss]) + elif model_type == 'gauss': # F0, FWHM_maj, FWHM_min, PA, x0, y0 + u_maj = u*np.sin(params['PA']) + v*np.cos(params['PA']) + u_min = u*np.cos(params['PA']) - v*np.sin(params['PA']) + vis = (params['F0'] + * np.exp(-np.pi**2/(4.*np.log(2.)) * ((u_maj * params['FWHM_maj'])**2 + (u_min * params['FWHM_min'])**2)) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + val = np.array([ 1.0/params['F0'] * vis, + -np.pi**2/(2.*np.log(2.)) * params['FWHM_maj'] * u_maj**2 * vis, + -np.pi**2/(2.*np.log(2.)) * params['FWHM_min'] * u_min**2 * vis, + -np.pi**2/(2.*np.log(2.)) * (params['FWHM_maj']**2 - params['FWHM_min']**2) * u_maj * u_min * vis, + 1j * 2.0 * np.pi * u * vis, + 1j * 2.0 * np.pi * v * vis]) + elif model_type == 'disk': # F0, d, x0, y0 + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + #Add a small offset to avoid issues with division by zero + z += (z == 0.0) * 1e-10 + vis = (params['F0'] * 2.0/z * sps.jv(1, z) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + val = np.array([ 1.0/params['F0'] * vis, + -(params['F0'] * 2.0/z * sps.jv(2, z) * np.pi * (u**2 + v**2)**0.5 * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) , + 1j * 2.0 * np.pi * u * vis, + 1j * 2.0 * np.pi * v * vis]) + elif model_type == 'blurred_disk': # F0, d, alpha, x0, y0 + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + #Add a small offset to avoid issues with division by zero + z += (z == 0.0) * 1e-10 + blur = np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + vis = (params['F0'] * 2.0/z * sps.jv(1, z) + * blur + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + val = np.array([ 1.0/params['F0'] * vis, + -params['F0'] * 2.0/z * sps.jv(2, z) * np.pi * (u**2 + v**2)**0.5 * blur * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + -np.pi**2 * (u**2 + v**2) * params['alpha']/(2.*np.log(2.)) * vis, + 1j * 2.0 * np.pi * u * vis, + 1j * 2.0 * np.pi * v * vis]) + elif model_type == 'ring': # F0, d, x0, y0 + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + val = np.array([ sps.jv(0, z) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + -np.pi * (u**2 + v**2)**0.5 * params['F0'] * sps.jv(1, z) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + 2.0 * np.pi * 1j * u * params['F0'] * sps.jv(0, z) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])), + 2.0 * np.pi * 1j * v * params['F0'] * sps.jv(0, z) * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))]) + elif model_type == 'thick_ring': # F0, d, alpha, x0, y0 + z = np.pi * params['d'] * (u**2 + v**2)**0.5 + vis = (params['F0'] * sps.jv(0, z) + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + val = np.array([ 1.0/params['F0'] * vis, + -(params['F0'] * np.pi * (u**2 + v**2)**0.5 * sps.jv(1, z) + * np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))), + -np.pi**2 * (u**2 + v**2) * params['alpha']/(2.*np.log(2.)) * vis, + 1j * 2.0 * np.pi * u * vis, + 1j * 2.0 * np.pi * v * vis]) + elif model_type in ['mring','thick_mring']: # F0, d, [alpha], x0, y0, beta1_re/abs, beta1_im/arg, beta2_re/abs, beta2_im/arg, ... + phi = np.angle(v + 1j*u) + # Flip the baseline sign to match eht-imaging conventions + phi += np.pi + uvdist = (u**2 + v**2)**0.5 + z = np.pi * params['d'] * uvdist + if model_type == 'thick_mring': + alpha_factor = np.exp(-(np.pi * params['alpha'] * (u**2 + v**2)**0.5)**2/(4. * np.log(2.))) + else: + alpha_factor = 1 + + # Only one of the beta_lists will affect the measurement and have non-zero gradients. Figure out which: + # These are for the derivatives wrt diameter + if pol == 'I': + beta_factor = (-np.pi * uvdist * sps.jv(1, z) + + np.sum([params['beta_list'][m-1] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0) + + np.sum([np.conj(params['beta_list'][m-1]) * 0.5 * (sps.jv( -m-1, z) - sps.jv( -m+1, z)) * np.pi * uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list'])+1)],axis=0)) + elif pol == 'P' and len(params['beta_list_pol']) > 0: + num_coeff = len(params['beta_list_pol']) + beta_factor = np.sum([params['beta_list_pol'][m + (num_coeff-1)//2] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(-(num_coeff-1)//2,(num_coeff+1)//2)],axis=0) + elif pol == 'V' and len(params['beta_list_cpol']) > 0: + beta_factor = (-np.pi * uvdist * sps.jv(1, z) * params['beta_list_cpol'][0] + + np.sum([params['beta_list_cpol'][m] * 0.5 * (sps.jv( m-1, z) - sps.jv( m+1, z)) * np.pi * uvdist * np.exp( 1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0) + + np.sum([np.conj(params['beta_list_cpol'][m]) * 0.5 * (sps.jv( -m-1, z) - sps.jv( -m+1, z)) * np.pi * uvdist * np.exp(-1j * m * (phi - np.pi/2.)) for m in range(1,len(params['beta_list_cpol']))],axis=0)) + else: + beta_factor = 0.0 + + vis = sample_1model_uv(u, v, model_type, params, pol=pol, jonesdict=jonesdict) + grad = [ 1.0/params['F0'] * vis, + (params['F0'] * alpha_factor * beta_factor * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])))] + if model_type == 'thick_mring': + grad.append(-np.pi**2/(2.*np.log(2)) * uvdist**2 * params['alpha'] * vis) + grad.append(1j * 2.0 * np.pi * u * vis) + grad.append(1j * 2.0 * np.pi * v * vis) + + if pol=='I': + # Add derivatives of the beta terms + for m in range(1,len(params['beta_list'])+1): + beta_grad_re = params['F0'] * alpha_factor * ( + sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) + sps.jv(-m, z) * np.exp(-1j * m * (phi - np.pi/2.)) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + beta_grad_im = 1j * params['F0'] * alpha_factor * ( + sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) - sps.jv(-m, z) * np.exp(-1j * m * (phi - np.pi/2.)) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + if COMPLEX_BASIS == 're-im': + grad.append(beta_grad_re) + grad.append(beta_grad_im) + elif COMPLEX_BASIS == 'abs-arg': + beta_abs = np.abs(params['beta_list'][m-1]) + beta_arg = np.angle(params['beta_list'][m-1]) + grad.append(beta_grad_re * np.cos(beta_arg) + beta_grad_im * np.sin(beta_arg)) + grad.append(-beta_abs * np.sin(beta_arg) * beta_grad_re + beta_abs * np.cos(beta_arg) * beta_grad_im) + else: + raise Exception('COMPLEX_BASIS ' + COMPLEX_BASIS + ' not recognized!') + else: + [grad.append(np.zeros_like(grad[0])) for _ in range(2*len(params['beta_list']))] + + if pol=='P' and fit_pol: + # Add derivatives of the beta terms + num_coeff = len(params['beta_list_pol']) + for m in range(-(num_coeff-1)//2,(num_coeff+1)//2): + beta_grad_re = params['F0'] * alpha_factor * sps.jv( m, z) * np.exp( 1j * m * (phi - np.pi/2.)) + beta_grad_im = 1j * beta_grad_re + if COMPLEX_BASIS == 're-im': + grad.append(beta_grad_re) + grad.append(beta_grad_im) + elif COMPLEX_BASIS == 'abs-arg': + beta_abs = np.abs(params['beta_list_pol'][m+(num_coeff-1)//2]) + beta_arg = np.angle(params['beta_list_pol'][m+(num_coeff-1)//2]) + grad.append(beta_grad_re * np.cos(beta_arg) + beta_grad_im * np.sin(beta_arg)) + grad.append(-beta_abs * np.sin(beta_arg) * beta_grad_re + beta_abs * np.cos(beta_arg) * beta_grad_im) + else: + raise Exception('COMPLEX_BASIS ' + COMPLEX_BASIS + ' not recognized!') + elif pol!='P' and fit_pol: + [grad.append(np.zeros_like(grad[0])) for _ in range(2*len(params['beta_list_pol']))] + + val = np.array(grad) + elif model_type == 'thick_mring_floor': # F0, d, [alpha], ff, x0, y0, beta1_re/abs, beta1_im/arg, beta2_re/abs, beta2_im/arg, ... + # We need to stich together the two gradients for the mring and the disk; we also need to add the gradient for the floor fraction ff + grad_mring = (1.0 - params['ff']) * sample_1model_grad_uv(u, v, 'thick_mring', params, pol=pol, fit_pol=fit_pol, fit_cpol=fit_cpol) + grad_disk = params['ff'] * sample_1model_grad_uv(u, v, 'blurred_disk', params, pol=pol, fit_pol=fit_pol, fit_cpol=fit_cpol) + + # mring: F0, d, alpha, x0, y0, beta1_re/abs, beta1_im/arg, beta2_re/abs, beta2_im/arg, ... + # disk: F0, d, alpha, x0, y0 + + # Here are derivatives for F0, d, and alpha + grad = [] + for j in range(3): + grad.append(grad_mring[j] + grad_disk[j]) + + # Here is the derivative for ff + grad.append( params['F0'] * (grad_disk[0]/params['ff'] - grad_mring[0]/(1.0 - params['ff'])) ) + + # Now the derivatives for x0 and y0 + grad.append(grad_mring[3] + grad_disk[3]) + grad.append(grad_mring[4] + grad_disk[4]) + + # Add remaining gradients for the mring + for j in range(5,len(grad_mring)): + grad.append(grad_mring[j]) + + val = np.array(grad) + elif model_type[:9] == 'stretched': + # Start with the model visibility + vis = sample_1model_uv(u, v, model_type, params, pol=pol, jonesdict=jonesdict) + + # Next, calculate the gradient wrt model parameters other than stretch and stretch_PA + # These are the same as the gradient of the unstretched model on stretched baselines + params_stretch = params.copy() + params_stretch['x0'] = 0.0 + params_stretch['y0'] = 0.0 + (u_stretch, v_stretch) = stretch_uv(u,v,params) + grad = (sample_1model_grad_uv(u_stretch, v_stretch, model_type[10:], params_stretch, pol=pol, fit_pol=fit_pol, fit_cpol=fit_cpol) + * np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0']))) + + # Add the gradient terms for the centroid + grad[model_params(model_type, params).index('x0')] = 1j * 2.0 * np.pi * u * vis + grad[model_params(model_type, params).index('y0')] = 1j * 2.0 * np.pi * v * vis + + # Now calculate the gradient with respect to stretch and stretch PA + grad_uv = sample_1model_graduv_uv(u_stretch, v_stretch, model_type[10:], params_stretch, pol=pol) + grad_stretch = grad_uv.copy() * 0.0 + grad_stretch[0] = ( grad_uv[0] * (u * np.sin(params['stretch_PA'])**2 + v * np.sin(params['stretch_PA']) * np.cos(params['stretch_PA'])) + + grad_uv[1] * (v * np.cos(params['stretch_PA'])**2 + u * np.sin(params['stretch_PA']) * np.cos(params['stretch_PA']))) + grad_stretch[0] *= np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + + grad_stretch[1] = ( grad_uv[0] * (params['stretch'] - 1.0) * ( u * np.sin(2.0 * params['stretch_PA']) + v * np.cos(2.0 * params['stretch_PA'])) + + grad_uv[1] * (params['stretch'] - 1.0) * (-v * np.sin(2.0 * params['stretch_PA']) + u * np.cos(2.0 * params['stretch_PA']))) + grad_stretch[1] *= np.exp(1j * 2.0 * np.pi * (u * params['x0'] + v * params['y0'])) + + val = np.concatenate([grad, grad_stretch]) + else: + print('Model ' + model_type + ' not recognized!') + val = 0.0 + + grad = val * get_const_polfac(model_type, params, pol) + + if (fit_pol or fit_cpol) and model_type.find('mring') == -1: + # Add gradient contributions for models that have constant polarization + if fit_pol: + # Add gradient wrt pol_frac if the polarization is P, otherwise ignore + grad_params = copy.deepcopy(params) + grad_params['pol_frac'] = 1.0 + grad = np.vstack([grad, (pol == 'P') * sample_1model_uv(u, v, model_type, grad_params, pol=pol, jonesdict=jonesdict)]) + + # Add gradient wrt pol_evpa if the polarization is P, otherwise ignore + grad_params = copy.deepcopy(params) + grad_params['pol_frac'] *= 2j + grad = np.vstack([grad, (pol == 'P') * sample_1model_uv(u, v, model_type, grad_params, pol=pol, jonesdict=jonesdict)]) + if fit_cpol: + # Add gradient wrt cpol_frac + grad_params = copy.deepcopy(params) + grad_params['cpol_frac'] = 1.0 + grad = np.vstack([grad, (pol == 'V') * sample_1model_uv(u, v, model_type, grad_params, pol=pol, jonesdict=jonesdict)]) + + return grad + +def sample_model_xy(models, params, x, y, psize=1.*RADPERUAS, pol='I'): + return np.sum(sample_1model_xy(x, y, models[j], params[j], psize=psize,pol=pol) for j in range(len(models))) + +def sample_model_uv(models, params, u, v, pol='I', jonesdict=None): + return np.sum(sample_1model_uv(u, v, models[j], params[j], pol=pol, jonesdict=jonesdict) for j in range(len(models))) + +def sample_model_graduv_uv(models, params, u, v, pol='I', jonesdict=None): + # Gradient of a sum of models wrt (u,v) + return np.sum([sample_1model_graduv_uv(u, v, models[j], params[j], pol=pol, jonesdict=jonesdict) for j in range(len(models))],axis=0) + +def sample_model_grad_uv(models, params, u, v, pol='I', fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + # Gradient of a sum of models for each parameter + if fit_leakage == False: + return np.concatenate([sample_1model_grad_uv(u, v, models[j], params[j], pol=pol, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) for j in range(len(models))]) + else: + # Need to sum the leakage contributions + allgrad = [sample_1model_grad_uv(u, v, models[j], params[j], pol=pol, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) for j in range(len(models))] + n_leakage = len(jonesdict['leakage_fit'])*2 + grad = np.concatenate([allgrad[j][:-n_leakage] for j in range(len(models))]) + grad_leakage = np.sum([allgrad[j][-n_leakage:] for j in range(len(models))],axis=0) + return np.concatenate([grad, grad_leakage]) + +def blur_circ_1model(model_type, params, fwhm): + """Blur a single model, returning new model type and associated parameters + + Args: + fwhm (float) : Full width at half maximum of the kernel (radians) + + Returns: + (dict) : Dictionary with new 'model_type' and new 'params' + """ + + model_type_blur = model_type + params_blur = params.copy() + + if model_type == 'point': + model_type_blur = 'circ_gauss' + params_blur['FWHM'] = fwhm + elif model_type == 'circ_gauss': + params_blur['FWHM'] = (params_blur['FWHM']**2 + fwhm**2)**0.5 + elif model_type == 'gauss': + params_blur['FWHM_maj'] = (params_blur['FWHM_maj']**2 + fwhm**2)**0.5 + params_blur['FWHM_min'] = (params_blur['FWHM_min']**2 + fwhm**2)**0.5 + elif 'thick' in model_type or 'blurred' in model_type: + params_blur['alpha'] = (params_blur['alpha']**2 + fwhm**2)**0.5 + elif model_type == 'disk': + model_type_blur = 'blurred_' + model_type + params_blur['alpha'] = fwhm + elif model_type == 'ring' or model_type == 'mring': + model_type_blur = 'thick_' + model_type + params_blur['alpha'] = fwhm + elif model_type == 'stretched_ring' or model_type == 'stretched_mring': + model_type_blur = 'stretched_thick_' + model_type[10:] + params_blur['alpha'] = fwhm + else: + raise Exception("A blurred " + model_type + " is not yet a supported model!") + + return {'model_type':model_type_blur, 'params':params_blur} + +class Model(object): + """A model with analytic representations in the image and visibility domains. + + Attributes: + """ + + def __init__(self, ra=RA_DEFAULT, dec=DEC_DEFAULT, pa=0.0, + polrep='stokes', pol_prim=None, + rf=RF_DEFAULT, source=SOURCE_DEFAULT, + mjd=MJD_DEFAULT, time=0.): + + """A model with analytic representations in the image and visibility domains. + + Args: + + Returns: + """ + + # The model is a sum of component models, each defined by a tag and associated parameters + self.pol_prim = 'I' + self.polrep = 'stokes' + self._imdict = {'I':{'models':[],'params':[]},'Q':{'models':[],'params':[]},'U':{'models':[],'params':[]},'V':{'models':[],'params':[]}} + + # Save the image metadata + self.ra = float(ra) + self.dec = float(dec) + self.pa = float(pa) + self.rf = float(rf) + self.source = str(source) + self.mjd = int(mjd) + if time > 24: + self.mjd += int((time - time % 24)/24) + self.time = float(time % 24) + else: + self.time = time + + @property + def models(self): + return self._imdict[self.pol_prim]['models'] + + @models.setter + def models(self, model_list): + self._imdict[self.pol_prim]['models'] = model_list + + @property + def params(self): + return self._imdict[self.pol_prim]['params'] + + @params.setter + def params(self, param_list): + self._imdict[self.pol_prim]['params'] = param_list + + def copy(self): + """Return a copy of the Model object. + + Args: + + Returns: + (Model): copy of the Model. + """ + out = Model(ra=self.ra, dec=self.dec, pa=self.pa, polrep=self.polrep, pol_prim=self.pol_prim,rf=self.rf,source=self.source,mjd=self.mjd,time=self.time) + out.models = copy.deepcopy(self.models) + out.params = copy.deepcopy(self.params.copy()) + return out + + def switch_polrep(self, polrep_out='stokes', pol_prim_out=None): + + """Return a new model with the polarization representation changed + Args: + polrep_out (str): the polrep of the output data + pol_prim_out (str): The default image: I,Q,U or V for Stokes, RR,LL,LR,RL for circ + + Returns: + (Model): new Model object with potentially different polrep + """ + + # Note: this currently does nothing, but it is put here for compatibility with functions such as selfcal + if polrep_out not in ['stokes','circ']: + raise Exception("polrep_out must be either 'stokes' or 'circ'") + if pol_prim_out is None: + if polrep_out=='stokes': pol_prim_out = 'I' + elif polrep_out=='circ': pol_prim_out = 'RR' + + return self.copy() + + def N_models(self): + """Return the number of model components + + Args: + + Returns: + (int): number of model components + """ + return len(self.models) + + def total_flux(self): + """Return the total flux of the model in Jy. + + Args: + + Returns: + (float) : model total flux (Jy) + """ + return np.real(self.sample_uv(0,0)) + + def blur_circ(self, fwhm): + """Return a new model, equal to the current one convolved with a circular Gaussian kernel + + Args: + fwhm (float) : Full width at half maximum of the kernel (radians) + + Returns: + (Model) : Blurred model + """ + + out = self.copy() + + for j in range(len(out.models)): + blur_model = blur_circ_1model(out.models[j], out.params[j], fwhm) + out.models[j] = blur_model['model_type'] + out.params[j] = blur_model['params'] + + return out + + def add_point(self, F0 = 1.0, x0 = 0.0, y0 = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add a point source model. + + Args: + F0 (float): The total flux of the point source (Jy) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + + Returns: + (Model): Updated Model + """ + + out = self.copy() + out.models.append('point') + out.params.append({'F0':F0,'x0':x0,'y0':y0,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_circ_gauss(self, F0 = 1.0, FWHM = 50.*RADPERUAS, x0 = 0.0, y0 = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add a circular Gaussian model. + + Args: + F0 (float): The total flux of the Gaussian (Jy) + FWHM (float): The FWHM of the Gaussian (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + + Returns: + (Model): Updated Model + """ + out = self.copy() + out.models.append('circ_gauss') + out.params.append({'F0':F0,'FWHM':FWHM,'x0':x0,'y0':y0,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_gauss(self, F0 = 1.0, FWHM_maj = 50.*RADPERUAS, FWHM_min = 50.*RADPERUAS, PA = 0.0, x0 = 0.0, y0 = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add an anisotropic Gaussian model. + + Args: + F0 (float): The total flux of the Gaussian (Jy) + FWHM_maj (float): The FWHM of the Gaussian major axis (radians) + FWHM_min (float): The FWHM of the Gaussian minor axis (radians) + PA (float): Position angle of the major axis, east of north (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + + Returns: + (Model): Updated Model + """ + out = self.copy() + out.models.append('gauss') + out.params.append({'F0':F0,'FWHM_maj':FWHM_maj,'FWHM_min':FWHM_min,'PA':PA,'x0':x0,'y0':y0,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_disk(self, F0 = 1.0, d = 50.*RADPERUAS, x0 = 0.0, y0 = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add a circular disk model. + + Args: + F0 (float): The total flux of the disk (Jy) + d (float): The diameter (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + + Returns: + (Model): Updated Model + """ + out = self.copy() + out.models.append('disk') + out.params.append({'F0':F0,'d':d,'x0':x0,'y0':y0,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_blurred_disk(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, x0 = 0.0, y0 = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add a circular disk model that is blurred with a circular Gaussian kernel. + + Args: + F0 (float): The total flux of the disk (Jy) + d (float): The diameter (radians) + alpha (float): The blurring (FWHM of Gaussian convolution) (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + + Returns: + (Model): Updated Model + """ + out = self.copy() + out.models.append('blurred_disk') + out.params.append({'F0':F0,'d':d,'alpha':alpha,'x0':x0,'y0':y0,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_ring(self, F0 = 1.0, d = 50.*RADPERUAS, x0 = 0.0, y0 = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add a ring model with infinitesimal thickness. + + Args: + F0 (float): The total flux of the ring (Jy) + d (float): The diameter (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + + Returns: + (Model): Updated Model + """ + out = self.copy() + out.models.append('ring') + out.params.append({'F0':F0,'d':d,'x0':x0,'y0':y0,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_stretched_ring(self, F0 = 1.0, d = 50.*RADPERUAS, x0 = 0.0, y0 = 0.0, stretch = 1.0, stretch_PA = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add a stretched ring model with infinitesimal thickness. + + Args: + F0 (float): The total flux of the ring (Jy) + d (float): The diameter (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + stretch (float): The stretch to apply (1.0 = no stretch) + stretch_PA (float): Position angle of the stretch, east of north (radians) + + Returns: + (Model): Updated Model + """ + out = self.copy() + out.models.append('stretched_ring') + out.params.append({'F0':F0,'d':d,'x0':x0,'y0':y0,'stretch':stretch,'stretch_PA':stretch_PA,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_thick_ring(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, x0 = 0.0, y0 = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add a ring model with finite thickness, determined by circular Gaussian convolution of a thin ring. + For details, see Appendix G of https://iopscience.iop.org/article/10.3847/2041-8213/ab0e85/pdf + + Args: + F0 (float): The total flux of the ring (Jy) + d (float): The ring diameter (radians) + alpha (float): The ring thickness (FWHM of Gaussian convolution) (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + + Returns: + (Model): Updated Model + """ + out = self.copy() + out.models.append('thick_ring') + out.params.append({'F0':F0,'d':d,'alpha':alpha,'x0':x0,'y0':y0,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_stretched_thick_ring(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, x0 = 0.0, y0 = 0.0, stretch = 1.0, stretch_PA = 0.0, pol_frac = 0.0, pol_evpa = 0.0, cpol_frac = 0.0): + """Add a ring model with finite thickness, determined by circular Gaussian convolution of a thin ring. + For details, see Appendix G of https://iopscience.iop.org/article/10.3847/2041-8213/ab0e85/pdf + + Args: + F0 (float): The total flux of the ring (Jy) + d (float): The ring diameter (radians) + alpha (float): The ring thickness (FWHM of Gaussian convolution) (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + stretch (float): The stretch to apply (1.0 = no stretch) + stretch_PA (float): Position angle of the stretch, east of north (radians) + + Returns: + (Model): Updated Model + """ + out = self.copy() + out.models.append('stretched_thick_ring') + out.params.append({'F0':F0,'d':d,'alpha':alpha,'x0':x0,'y0':y0,'stretch':stretch,'stretch_PA':stretch_PA,'pol_frac':pol_frac,'pol_evpa':pol_evpa,'cpol_frac':cpol_frac}) + return out + + def add_mring(self, F0 = 1.0, d = 50.*RADPERUAS, x0 = 0.0, y0 = 0.0, beta_list = None, beta_list_pol = None, beta_list_cpol = None): + """Add a ring model with azimuthal brightness variations determined by a Fourier mode expansion. + For details, see Eq. 18-20 of https://arxiv.org/abs/1907.04329 + + Args: + F0 (float): The total flux of the ring (Jy), which is also beta_0. + d (float): The diameter (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + beta_list (list): List of complex Fourier coefficients, [beta_1, beta_2, ...]. + Negative indices are determined by the condition beta_{-m} = beta_m*. + Indices are all scaled by F0 = beta_0, so they are dimensionless. + Returns: + (Model): Updated Model + """ + if beta_list is None: beta_list = [] + if beta_list_pol is None: beta_list_pol = [] + if beta_list_cpol is None: beta_list_cpol = [] + + out = self.copy() + if beta_list is None: + beta_list = [0.0] + out.models.append('mring') + out.params.append({'F0':F0,'d':d,'beta_list':np.array(beta_list, dtype=np.complex_),'beta_list_pol':np.array(beta_list_pol, dtype=np.complex_),'beta_list_cpol':np.array(beta_list_cpol, dtype=np.complex_),'x0':x0,'y0':y0}) + return out + + def add_stretched_mring(self, F0 = 1.0, d = 50.*RADPERUAS, x0 = 0.0, y0 = 0.0, beta_list = None, beta_list_pol = None, beta_list_cpol = None, stretch = 1.0, stretch_PA = 0.0): + """Add a stretched ring model with azimuthal brightness variations determined by a Fourier mode expansion. + For details, see Eq. 18-20 of https://arxiv.org/abs/1907.04329 + + Args: + F0 (float): The total flux of the ring (Jy), which is also beta_0. + d (float): The diameter (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + beta_list (list): List of complex Fourier coefficients, [beta_1, beta_2, ...]. + Negative indices are determined by the condition beta_{-m} = beta_m*. + Indices are all scaled by F0 = beta_0, so they are dimensionless. + stretch (float): The stretch to apply (1.0 = no stretch) + stretch_PA (float): Position angle of the stretch, east of north (radians) + Returns: + (Model): Updated Model + """ + if beta_list is None: beta_list = [] + if beta_list_pol is None: beta_list_pol = [] + if beta_list_cpol is None: beta_list_cpol = [] + + out = self.copy() + if beta_list is None: + beta_list = [0.0] + out.models.append('stretched_mring') + out.params.append({'F0':F0,'d':d,'beta_list':np.array(beta_list, dtype=np.complex_),'beta_list_pol':np.array(beta_list_pol, dtype=np.complex_),'beta_list_cpol':np.array(beta_list_cpol, dtype=np.complex_),'x0':x0,'y0':y0,'stretch':stretch,'stretch_PA':stretch_PA}) + return out + + def add_thick_mring(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, x0 = 0.0, y0 = 0.0, beta_list = None, beta_list_pol = None, beta_list_cpol = None): + """Add a ring model with azimuthal brightness variations determined by a Fourier mode expansion and thickness determined by circular Gaussian convolution. + For details, see Eq. 18-20 of https://arxiv.org/abs/1907.04329 + The Gaussian convolution calculation is a trivial generalization of Appendix G of https://iopscience.iop.org/article/10.3847/2041-8213/ab0e85/pdf + + Args: + F0 (float): The total flux of the ring (Jy), which is also beta_0. + d (float): The ring diameter (radians) + alpha (float): The ring thickness (FWHM of Gaussian convolution) (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + beta_list (list): List of complex Fourier coefficients, [beta_1, beta_2, ...]. + Negative indices are determined by the condition beta_{-m} = beta_m*. + Indices are all scaled by F0 = beta_0, so they are dimensionless. + Returns: + (Model): Updated Model + """ + if beta_list is None: beta_list = [] + if beta_list_pol is None: beta_list_pol = [] + if beta_list_cpol is None: beta_list_cpol = [] + + out = self.copy() + if beta_list is None: + beta_list = [0.0] + out.models.append('thick_mring') + out.params.append({'F0':F0,'d':d,'beta_list':np.array(beta_list, dtype=np.complex_),'beta_list_pol':np.array(beta_list_pol, dtype=np.complex_),'beta_list_cpol':np.array(beta_list_cpol, dtype=np.complex_),'alpha':alpha,'x0':x0,'y0':y0}) + return out + + def add_thick_mring_floor(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, ff=0.0, x0 = 0.0, y0 = 0.0, beta_list = None, beta_list_pol = None, beta_list_cpol = None): + """Add a ring model with azimuthal brightness variations determined by a Fourier mode expansion, thickness determined by circular Gaussian convolution, and a floor + The floor is a blurred disk, with diameter d and blurred FWHM alpha + For details, see Eq. 18-20 of https://arxiv.org/abs/1907.04329 + The Gaussian convolution calculation is a trivial generalization of Appendix G of https://iopscience.iop.org/article/10.3847/2041-8213/ab0e85/pdf + + Args: + F0 (float): The total flux of the ring (Jy), which is also beta_0. + d (float): The ring diameter (radians) + alpha (float): The ring thickness (FWHM of Gaussian convolution) (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + ff (float): The fraction of the total flux in the floor + beta_list (list): List of complex Fourier coefficients, [beta_1, beta_2, ...]. + Negative indices are determined by the condition beta_{-m} = beta_m*. + Indices are all scaled by F0 = beta_0, so they are dimensionless. + Returns: + (Model): Updated Model + """ + if beta_list is None: beta_list = [] + if beta_list_pol is None: beta_list_pol = [] + if beta_list_cpol is None: beta_list_cpol = [] + + out = self.copy() + if beta_list is None: + beta_list = [0.0] + out.models.append('thick_mring_floor') + out.params.append({'F0':F0,'d':d,'beta_list':np.array(beta_list, dtype=np.complex_),'beta_list_pol':np.array(beta_list_pol, dtype=np.complex_),'beta_list_cpol':np.array(beta_list_cpol, dtype=np.complex_),'alpha':alpha,'x0':x0,'y0':y0,'ff':ff}) + return out + + def add_stretched_thick_mring(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, x0 = 0.0, y0 = 0.0, beta_list = None, beta_list_pol = None, beta_list_cpol = None, stretch = 1.0, stretch_PA = 0.0): + """Add a ring model with azimuthal brightness variations determined by a Fourier mode expansion and thickness determined by circular Gaussian convolution. + For details, see Eq. 18-20 of https://arxiv.org/abs/1907.04329 + The Gaussian convolution calculation is a trivial generalization of Appendix G of https://iopscience.iop.org/article/10.3847/2041-8213/ab0e85/pdf + + Args: + F0 (float): The total flux of the ring (Jy), which is also beta_0. + d (float): The ring diameter (radians) + alpha (float): The ring thickness (FWHM of Gaussian convolution) (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + beta_list (list): List of complex Fourier coefficients, [beta_1, beta_2, ...]. + Negative indices are determined by the condition beta_{-m} = beta_m*. + Indices are all scaled by F0 = beta_0, so they are dimensionless. + stretch (float): The stretch to apply (1.0 = no stretch) + stretch_PA (float): Position angle of the stretch, east of north (radians) + Returns: + (Model): Updated Model + """ + if beta_list is None: beta_list = [] + if beta_list_pol is None: beta_list_pol = [] + if beta_list_cpol is None: beta_list_cpol = [] + + out = self.copy() + if beta_list is None: + beta_list = [0.0] + out.models.append('stretched_thick_mring') + out.params.append({'F0':F0,'d':d,'beta_list':np.array(beta_list, dtype=np.complex_),'beta_list_pol':np.array(beta_list_pol, dtype=np.complex_),'beta_list_cpol':np.array(beta_list_cpol, dtype=np.complex_),'alpha':alpha,'x0':x0,'y0':y0,'stretch':stretch,'stretch_PA':stretch_PA}) + return out + + def add_stretched_thick_mring_floor(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, ff=0.0, x0 = 0.0, y0 = 0.0, beta_list = None, beta_list_pol = None, beta_list_cpol = None, stretch = 1.0, stretch_PA = 0.0): + """Add a ring model with azimuthal brightness variations determined by a Fourier mode expansion and thickness determined by circular Gaussian convolution. + For details, see Eq. 18-20 of https://arxiv.org/abs/1907.04329 + The Gaussian convolution calculation is a trivial generalization of Appendix G of https://iopscience.iop.org/article/10.3847/2041-8213/ab0e85/pdf + + Args: + F0 (float): The total flux of the ring (Jy), which is also beta_0. + d (float): The ring diameter (radians) + alpha (float): The ring thickness (FWHM of Gaussian convolution) (radians) + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + beta_list (list): List of complex Fourier coefficients, [beta_1, beta_2, ...]. + Negative indices are determined by the condition beta_{-m} = beta_m*. + Indices are all scaled by F0 = beta_0, so they are dimensionless. + stretch (float): The stretch to apply (1.0 = no stretch) + stretch_PA (float): Position angle of the stretch, east of north (radians) + Returns: + (Model): Updated Model + """ + if beta_list is None: beta_list = [] + if beta_list_pol is None: beta_list_pol = [] + if beta_list_cpol is None: beta_list_cpol = [] + + out = self.copy() + if beta_list is None: + beta_list = [0.0] + out.models.append('stretched_thick_mring_floor') + out.params.append({'F0':F0,'d':d,'beta_list':np.array(beta_list, dtype=np.complex_),'beta_list_pol':np.array(beta_list_pol, dtype=np.complex_),'beta_list_cpol':np.array(beta_list_cpol, dtype=np.complex_),'alpha':alpha,'x0':x0,'y0':y0,'stretch':stretch,'stretch_PA':stretch_PA,'ff':ff}) + return out + + def sample_xy(self, x, y, psize=1.*RADPERUAS, pol='I'): + """Sample model image on the specified x and y coordinates + + Args: + x (float): x coordinate (dimensionless) + y (float): y coordinate (dimensionless) + + Returns: + (float): Image brightness (Jy/radian^2) + """ + return sample_model_xy(self.models, self.params, x, y, psize=psize, pol=pol) + + def sample_uv(self, u, v, polrep_obs='Stokes', pol='I', jonesdict=None): + """Sample model visibility on the specified u and v coordinates + + Args: + u (float): u coordinate (dimensionless) + v (float): v coordinate (dimensionless) + + Returns: + (complex): complex visibility (Jy) + """ + return sample_model_uv(self.models, self.params, u, v, pol=pol, jonesdict=jonesdict) + + def sample_graduv_uv(self, u, v, pol='I', jonesdict=None): + """Sample model visibility gradient on the specified u and v coordinates wrt (u,v) + + Args: + u (float): u coordinate (dimensionless) + v (float): v coordinate (dimensionless) + + Returns: + (complex): complex visibility (Jy) + """ + return sample_model_graduv_uv(self.models, self.params, u, v, pol=pol, jonesdict=jonesdict) + + def sample_grad_uv(self, u, v, pol='I', fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """Sample model visibility gradient on the specified u and v coordinates wrt all model parameters + + Args: + u (float): u coordinate (dimensionless) + v (float): v coordinate (dimensionless) + + Returns: + (complex): complex visibility (Jy) + """ + return sample_model_grad_uv(self.models, self.params, u, v, pol=pol, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + + def centroid(self, pol=None): + """Compute the location of the image centroid (corresponding to the polarization pol) + Note: This quantity is only well defined for total intensity + + Args: + pol (str): The polarization for which to find the image centroid + + Returns: + (np.array): centroid positions (x0,y0) in radians + """ + + if pol is None: pol=self.pol_prim + if not (pol in list(self._imdict.keys())): + raise Exception("for polrep==%s, pol must be in " % + self.polrep + ",".join(list(self._imdict.keys()))) + + return np.real(self.sample_graduv_uv(0,0)/(2.*np.pi*1j))/self.total_flux() + + def default_prior(self,fit_pol=False,fit_cpol=False): + return [default_prior(self.models[j],self.params[j],fit_pol=fit_pol,fit_cpol=fit_cpol) for j in range(self.N_models())] + + def display(self, fov=FOV_DEFAULT, npix=NPIX_DEFAULT, polrep='stokes', pol_prim=None, pulse=PULSE_DEFAULT, time=0., **kwargs): + return self.make_image(fov, npix, polrep, pol_prim, pulse, time).display(**kwargs) + + def make_image(self, fov, npix, polrep='stokes', pol_prim=None, pulse=PULSE_DEFAULT, time=0.): + """Sample the model onto a square image. + + Args: + fov (float): the field of view of each axis in radians + npix (int): the number of pixels on each axis + ra (float): The source Right Ascension in fractional hours + dec (float): The source declination in fractional degrees + rf (float): The image frequency in Hz + + source (str): The source name + polrep (str): polarization representation, either 'stokes' or 'circ' + pol_prim (str): The default image: I,Q,U or V for Stokes, RR,LL,LR,RL for Circular + pulse (function): The function convolved with the pixel values for continuous image. + + mjd (int): The integer MJD of the image + time (float): The observing time of the image (UTC hours) + + Returns: + (Image): an image object + """ + + pdim = fov/float(npix) + npix = int(npix) + imarr = np.zeros((npix,npix)) + outim = image.Image(imarr, pdim, self.ra, self.dec, + polrep=polrep, pol_prim=pol_prim, + rf=self.rf, source=self.source, mjd=self.mjd, time=time, pulse=pulse) + + return self.image_same(outim) + + def image_same(self, im): + """Create an image of the model with parameters equal to a reference image. + + Args: + im (Image): the reference image + + Returns: + (Image): image of the model + """ + out = im.copy() + xlist = np.arange(0,-im.xdim,-1)*im.psize + (im.psize*im.xdim)/2.0 - im.psize/2.0 + ylist = np.arange(0,-im.ydim,-1)*im.psize + (im.psize*im.ydim)/2.0 - im.psize/2.0 + + x_grid, y_grid = np.meshgrid(xlist, ylist) + imarr = self.sample_xy(x_grid, y_grid, im.psize) + out.imvec = imarr.flatten() # Change this to init with image_args + + # Add the remaining polarizations + for pol in ['Q','U','V']: + out.add_pol_image(self.sample_xy(x_grid, y_grid, im.psize, pol=pol), pol) + + return out + + def observe_same_nonoise(self, obs, **kwargs): + """Observe the model on the same baselines as an existing observation, without noise. + + Args: + obs (Obsdata): the existing observation + + Returns: + (Obsdata): an observation object with no noise + """ + + # Copy data to be safe + obsdata = copy.deepcopy(obs.data) + + # Load optional parameters + jonesdict = kwargs.get('jonesdict',None) + + # Compute visibilities and put them into the obsdata + if obs.polrep=='stokes': + obsdata['vis'] = self.sample_uv(obs.data['u'], obs.data['v'], pol='I', jonesdict=jonesdict) + obsdata['qvis'] = self.sample_uv(obs.data['u'], obs.data['v'], pol='Q', jonesdict=jonesdict) + obsdata['uvis'] = self.sample_uv(obs.data['u'], obs.data['v'], pol='U', jonesdict=jonesdict) + obsdata['vvis'] = self.sample_uv(obs.data['u'], obs.data['v'], pol='V', jonesdict=jonesdict) + elif obs.polrep=='circ': + obsdata['rrvis'] = self.sample_uv(obs.data['u'], obs.data['v'], pol='RR', jonesdict=jonesdict) + obsdata['rlvis'] = self.sample_uv(obs.data['u'], obs.data['v'], pol='RL', jonesdict=jonesdict) + obsdata['lrvis'] = self.sample_uv(obs.data['u'], obs.data['v'], pol='LR', jonesdict=jonesdict) + obsdata['llvis'] = self.sample_uv(obs.data['u'], obs.data['v'], pol='LL', jonesdict=jonesdict) + + obs_no_noise = ehtim.obsdata.Obsdata(obs.ra, obs.dec, obs.rf, obs.bw, obsdata, obs.tarr, + source=obs.source, mjd=obs.mjd, polrep=obs.polrep, + ampcal=True, phasecal=True, opacitycal=True, + dcal=True, frcal=True, + timetype=obs.timetype, scantable=obs.scans) + + return obs_no_noise + + def observe_same(self, obs_in, add_th_noise=True, sgrscat=False, ttype=False, # Note: sgrscat and ttype are kept for consistency with comp_plots + opacitycal=True, ampcal=True, phasecal=True, + dcal=True, frcal=True, rlgaincal=True, + stabilize_scan_phase=False, stabilize_scan_amp=False, neggains=False, + jones=False, inv_jones=False, + tau=TAUDEF, taup=GAINPDEF, + gain_offset=GAINPDEF, gainp=GAINPDEF, + dterm_offset=DTERMPDEF, caltable_path=None, seed=False, **kwargs): + + """Observe the image on the same baselines as an existing observation object and add noise. + + Args: + obs_in (Obsdata): the existing observation + + add_th_noise (bool): if True, baseline-dependent thermal noise is added + opacitycal (bool): if False, time-dependent gaussian errors are added to opacities + ampcal (bool): if False, time-dependent gaussian errors are added to station gains + phasecal (bool): if False, time-dependent station-based random phases are added + frcal (bool): if False, feed rotation angle terms are added to Jones matrices. + dcal (bool): if False, time-dependent gaussian errors added to D-terms. + + stabilize_scan_phase (bool): if True, random phase errors are constant over scans + stabilize_scan_amp (bool): if True, random amplitude errors are constant over scans + neggains (bool): if True, force the applied gains to be <1 + meaning that you have overestimated your telescope's performance + + + jones (bool): if True, uses Jones matrix to apply mis-calibration effects + inv_jones (bool): if True, applies estimated inverse Jones matrix + (not including random terms) to a priori calibrate data + + tau (float): the base opacity at all sites, + or a dict giving one opacity per site + taup (float): the fractional std. dev. of the random error on the opacities + gainp (float): the fractional std. dev. of the random error on the gains + gain_offset (float): the base gain offset at all sites, + or a dict giving one offset per site + dterm_offset (float): the base dterm offset at all sites, + or a dict giving one dterm offset per site + + caltable_path (string): The path and prefix of a saved caltable + + seed (int): seeds the random component of the noise terms. DO NOT set to 0! + + Returns: + (Obsdata): an observation object + """ + + if seed!=False: + np.random.seed(seed=seed) + + obs = self.observe_same_nonoise(obs_in, **kwargs) + + # Jones Matrix Corruption & Calibration + if jones: + obsdata = simobs.add_jones_and_noise(obs, add_th_noise=add_th_noise, + opacitycal=opacitycal, ampcal=ampcal, + phasecal=phasecal, dcal=dcal, frcal=frcal, + rlgaincal=rlgaincal, + stabilize_scan_phase=stabilize_scan_phase, + stabilize_scan_amp=stabilize_scan_amp, + neggains=neggains, + gainp=gainp, taup=taup, gain_offset=gain_offset, + dterm_offset=dterm_offset, + caltable_path=caltable_path, seed=seed) + + obs = ehtim.obsdata.Obsdata(obs.ra, obs.dec, obs.rf, obs.bw, obsdata, obs.tarr, + source=obs.source, mjd=obs.mjd, polrep=obs_in.polrep, + ampcal=ampcal, phasecal=phasecal, opacitycal=opacitycal, + dcal=dcal, frcal=frcal, + timetype=obs.timetype, scantable=obs.scans) + + if inv_jones: + obsdata = simobs.apply_jones_inverse(obs, + opacitycal=opacitycal, dcal=dcal, frcal=frcal) + + obs = ehtim.obsdata.Obsdata(obs.ra, obs.dec, obs.rf, obs.bw, obsdata, obs.tarr, + source=obs.source, mjd=obs.mjd, polrep=obs_in.polrep, + ampcal=ampcal, phasecal=phasecal, + opacitycal=True, dcal=True, frcal=True, + timetype=obs.timetype, scantable=obs.scans) + #these are always set to True after inverse jones call + + + # No Jones Matrices, Add noise the old way + # NOTE There is an asymmetry here - in the old way, we don't offer the ability to *not* + # unscale estimated noise. + else: + + if caltable_path: + print('WARNING: the caltable is only saved if you apply noise with a Jones Matrix') + + obsdata = simobs.add_noise(obs, add_th_noise=add_th_noise, + ampcal=ampcal, phasecal=phasecal, opacitycal=opacitycal, + stabilize_scan_phase=stabilize_scan_phase, + stabilize_scan_amp=stabilize_scan_amp, + gainp=gainp, taup=taup, gain_offset=gain_offset, seed=seed) + + obs = ehtim.obsdata.Obsdata(obs.ra, obs.dec, obs.rf, obs.bw, obsdata, obs.tarr, + source=obs.source, mjd=obs.mjd, polrep=obs_in.polrep, + ampcal=ampcal, phasecal=phasecal, + opacitycal=True, dcal=True, frcal=True, + timetype=obs.timetype, scantable=obs.scans) + #these are always set to True after inverse jones call + + return obs + + def observe(self, array, tint, tadv, tstart, tstop, bw, + mjd=None, timetype='UTC', polrep_obs=None, + elevmin=ELEV_LOW, elevmax=ELEV_HIGH, + fix_theta_GMST=False, add_th_noise=True, + opacitycal=True, ampcal=True, phasecal=True, + dcal=True, frcal=True, rlgaincal=True, + stabilize_scan_phase=False, stabilize_scan_amp=False, + jones=False, inv_jones=False, + tau=TAUDEF, taup=GAINPDEF, + gainp=GAINPDEF, gain_offset=GAINPDEF, + dterm_offset=DTERMPDEF, seed=False, **kwargs): + + """Generate baselines from an array object and observe the image. + + Args: + array (Array): an array object containing sites with which to generate baselines + tint (float): the scan integration time in seconds + tadv (float): the uniform cadence between scans in seconds + tstart (float): the start time of the observation in hours + tstop (float): the end time of the observation in hours + bw (float): the observing bandwidth in Hz + + mjd (int): the mjd of the observation, if set as different from the image mjd + timetype (str): how to interpret tstart and tstop; either 'GMST' or 'UTC' + elevmin (float): station minimum elevation in degrees + elevmax (float): station maximum elevation in degrees + + polrep_obs (str): 'stokes' or 'circ' sets the data polarimetric representation + + fix_theta_GMST (bool): if True, stops earth rotation to sample fixed u,v + sgrscat (bool): if True, the visibilites will be blurred by the Sgr A* kernel + add_th_noise (bool): if True, baseline-dependent thermal noise is added + opacitycal (bool): if False, time-dependent gaussian errors are added to opacities + ampcal (bool): if False, time-dependent gaussian errors are added to station gains + phasecal (bool): if False, time-dependent station-based random phases are added + frcal (bool): if False, feed rotation angle terms are added to Jones matrices. + dcal (bool): if False, time-dependent gaussian errors added to Jones matrices D-terms. + + stabilize_scan_phase (bool): if True, random phase errors are constant over scans + stabilize_scan_amp (bool): if True, random amplitude errors are constant over scans + jones (bool): if True, uses Jones matrix to apply mis-calibration effects + otherwise uses old formalism without D-terms + inv_jones (bool): if True, applies estimated inverse Jones matrix + (not including random terms) to calibrate data + + tau (float): the base opacity at all sites, + or a dict giving one opacity per site + taup (float): the fractional std. dev. of the random error on the opacities + gain_offset (float): the base gain offset at all sites, + or a dict giving one gain offset per site + gainp (float): the fractional std. dev. of the random error on the gains + + dterm_offset (float): the base dterm offset at all sites, + or a dict giving one dterm offset per site + + seed (int): seeds the random component of noise added. DO NOT set to 0! + + Returns: + (Obsdata): an observation object + """ + + # Generate empty observation + print("Generating empty observation file . . . ") + + if mjd == None: + mjd = self.mjd + if polrep_obs is None: + polrep_obs=self.polrep + + obs = array.obsdata(self.ra, self.dec, self.rf, bw, tint, tadv, tstart, tstop, mjd=mjd, + polrep=polrep_obs, tau=tau, timetype=timetype, + elevmin=elevmin, elevmax=elevmax, fix_theta_GMST=fix_theta_GMST) + + # Observe on the same baselines as the empty observation and add noise + obs = self.observe_same(obs, add_th_noise=add_th_noise, + opacitycal=opacitycal,ampcal=ampcal, + phasecal=phasecal,dcal=dcal, + frcal=frcal, rlgaincal=rlgaincal, + stabilize_scan_phase=stabilize_scan_phase, + stabilize_scan_amp=stabilize_scan_amp, + gainp=gainp,gain_offset=gain_offset, + tau=tau, taup=taup, + dterm_offset=dterm_offset, + jones=jones, inv_jones=inv_jones, seed=seed, **kwargs) + + obs.mjd = mjd + + return obs + + def save_txt(self,filename): + # Header + import ehtim.observing.obs_helpers as obshelp + mjd = float(self.mjd) + time = self.time + mjd += (time/24.) + + head = ("SRC: %s \n" % self.source + + "RA: " + obshelp.rastring(self.ra) + "\n" + + "DEC: " + obshelp.decstring(self.dec) + "\n" + + "MJD: %.6f \n" % (float(mjd)) + + "RF: %.4f GHz" % (self.rf/1e9)) + # Models + out = [] + for j in range(self.N_models()): + out.append(self.models[j]) + out.append(str(self.params[j]).replace('\n','').replace('complex128','np.complex128').replace('array','np.array')) + np.savetxt(filename, out, header=head, fmt="%s") + + def load_txt(self,filename): + lines = open(filename).read().splitlines() + + src = ' '.join(lines[0].split()[2:]) + ra = lines[1].split() + self.ra = float(ra[2]) + float(ra[4])/60.0 + float(ra[6])/3600.0 + dec = lines[2].split() + self.dec = np.sign(float(dec[2])) * (abs(float(dec[2])) + float(dec[4])/60.0 + float(dec[6])/3600.0) + mjd_float = float(lines[3].split()[2]) + self.mjd = int(mjd_float) + self.time = (mjd_float - self.mjd) * 24 + self.rf = float(lines[4].split()[2]) * 1e9 + + self.models = lines[5::2] + self.params = [eval(x) for x in lines[6::2]] + +def load_txt(filename): + out = Model() + out.load_txt(filename) + return out diff --git a/ehtim/modeling/__init__.py b/ehtim/modeling/__init__.py new file mode 100644 index 00000000..2d00f4e5 --- /dev/null +++ b/ehtim/modeling/__init__.py @@ -0,0 +1,9 @@ +""" +.. module:: ehtim.modeling + :platform: Unix + :synopsis: EHT Modeling Utilities: modeling functions + +.. moduleauthor:: Michael Johnson (mjohnson@cfa.harvard.edu) + +""" +from ..const_def import * diff --git a/ehtim/modeling/modeling_utils.py b/ehtim/modeling/modeling_utils.py new file mode 100644 index 00000000..9c445bcd --- /dev/null +++ b/ehtim/modeling/modeling_utils.py @@ -0,0 +1,2660 @@ +# modeling_utils.py +# General modeling functions for total intensity VLBI data +# +# Copyright (C) 2020 Michael Johnson +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +## TODO ## +# >> return jonesdict for all data types <- requires significant modification to eht-imaging +# >> Deal with nans in fitting (mask chisqdata) <- mostly done +# >> Add optional transform for leakage and gains + +from __future__ import division +from __future__ import print_function +from __future__ import absolute_import +from builtins import range + +import string +import time +import numpy as np +import scipy.optimize as opt +import scipy.ndimage as nd +import scipy.ndimage.filters as filt +import matplotlib.pyplot as plt +import scipy.special as sps +import scipy.stats as stats +import copy + +import ehtim.obsdata as obsdata +import ehtim.image as image +import ehtim.model as model +import ehtim.caltable as caltable + +from ehtim.const_def import * +from ehtim.observing.obs_helpers import * +from ehtim.statistics.dataframes import * + +from IPython import display + +################################################################################################## +# Constants & Definitions +################################################################################################## + +MAXLS = 100 # maximum number of line searches in L-BFGS-B +NHIST = 100 # number of steps to store for hessian approx +MAXIT = 100 # maximum number of iterations +STOP = 1.e-8 # convergence criterion + +BOUNDS_MIN = -1e4 +BOUNDS_MAX = 1e4 +BOUNDS_GAUSS_NSIGMA = 10. +BOUNDS_EXP_NSIGMA = 10. + +PRIOR_MIN = 1e-200 # to avoid problems with log-prior + +DATATERMS = ['vis', 'bs', 'amp', 'cphase', 'cphase_diag', 'camp', 'logcamp', 'logcamp_diag', 'logamp', 'pvis', 'm', 'rlrr', 'rlll', 'lrrr', 'lrll','rrll','llrr','polclosure'] + +nit = 0 # global variable to track the iteration number in the plotting callback +globdict = {} # global dictionary with all parameters related to the model fitting (mainly for efficient parallelization, but also very useful for debugging) + +# Details on each fitted parameter (convenience rescaling factor and associated unit) +PARAM_DETAILS = {'F0':[1.,'Jy'], 'FWHM':[RADPERUAS,'uas'], 'FWHM_maj':[RADPERUAS,'uas'], 'FWHM_min':[RADPERUAS,'uas'], + 'd':[RADPERUAS,'uas'], 'PA':[np.pi/180.,'deg'], 'alpha':[RADPERUAS,'uas'], 'ff':[1.,''], + 'x0':[RADPERUAS,'uas'], 'y0':[RADPERUAS,'uas'], 'stretch':[1.,''], 'stretch_PA':[np.pi/180.,'deg'], + 'arg':[np.pi/180.,'deg'], 'evpa':[np.pi/180.,'deg']} + +GAIN_PRIOR_DEFAULT = {'prior_type':'lognormal','sigma':0.1,'mu':0.0,'shift':-1.0} +LEAKAGE_PRIOR_DEFAULT = {'prior_type':'flat','min':-0.5,'max':0.5} +N_POSTERIOR_SAMPLES = 100 + +################################################################################################## +# Priors +################################################################################################## + +def cdf(x, prior_params): + """Compute the cumulative distribution function CDF(x) of a given prior at a given point x + + Args: + x (float): Value at which to compute the CDF + prior_params (dict): Dictionary with information about the prior + + Returns: + float: CDF(x) + """ + if prior_params['prior_type'] == 'flat': + return ( (x > prior_params['max']) * 1.0 + + (x > prior_params['min']) * (x < prior_params['max']) * (x - prior_params['min'])/(prior_params['max'] - prior_params['min'])) + elif prior_params['prior_type'] == 'gauss': + return 0.5 * (1.0 + sps.erf( (x - prior_params['mean'])/(prior_params['std'] * np.sqrt(2.0)) )) + elif prior_params['prior_type'] == 'exponential': + return (1.0 - np.exp(-x/prior_params['std'])) * (x >= 0.0) + elif prior_params['prior_type'] == 'lognormal': + return (x > prior_params['shift']) * (0.5 * sps.erfc( (prior_params['mu'] - np.log(x - prior_params['shift']))/(np.sqrt(2.0) * prior_params['sigma']))) + elif prior_params['prior_type'] == 'positive': + raise Exception('CDF is not defined for prior type "positive"') + elif prior_params['prior_type'] == 'none': + raise Exception('CDF is not defined for prior type "none"') + elif prior_params['prior_type'] == 'fixed': + raise Exception('CDF is not defined for prior type "fixed"') + else: + raise Exception('Prior type ' + prior_params['prior_type'] + ' not recognized!') + +def cdf_inverse(x, prior_params): + """Compute the inverse cumulative distribution function of a given prior at a given point 0 <= x <= 1 + + Args: + x (float): Value at which to compute the inverse CDF + prior_params (dict): Dictionary with information about the prior + + Returns: + float: Inverse CDF at x + """ + if prior_params['prior_type'] == 'flat': + return prior_params['min'] * (1.0 - x) + prior_params['max'] * x + elif prior_params['prior_type'] == 'gauss': + return prior_params['mean'] - np.sqrt(2.0) * prior_params['std'] * sps.erfcinv(2.0 * x) + elif prior_params['prior_type'] == 'exponential': + return prior_params['std'] * np.log(1.0/(1.0 - x)) + elif prior_params['prior_type'] == 'lognormal': + return np.exp( prior_params['mu'] - np.sqrt(2.0) * prior_params['sigma'] * sps.erfcinv(2.0 * x)) + prior_params['shift'] + elif prior_params['prior_type'] == 'positive': + raise Exception('CDF is not defined for prior type "positive"') + elif prior_params['prior_type'] == 'none': + raise Exception('CDF is not defined for prior type "none"') + elif prior_params['prior_type'] == 'fixed': + raise Exception('CDF is not defined for prior type "fixed"') + else: + raise Exception('Prior type ' + prior_params['prior_type'] + ' not recognized!') + +def param_bounds(prior_params): + """Compute the parameter boundaries associated with a given prior + + Args: + prior_params (dict): Dictionary with information about the prior + + Returns: + list: 2-element list specifying the allowed parameter range: [min,max] + """ + if prior_params.get('transform','') == 'cdf': + bounds = [0.0, 1.0] + elif prior_params['prior_type'] == 'flat': + bounds = [prior_params['min'],prior_params['max']] + elif prior_params['prior_type'] == 'gauss': + bounds = [prior_params['mean'] - prior_params['std'] * BOUNDS_GAUSS_NSIGMA, prior_params['mean'] + prior_params['std'] * BOUNDS_GAUSS_NSIGMA] + elif prior_params['prior_type'] == 'exponential': + bounds = [PRIOR_MIN, BOUNDS_EXP_NSIGMA * prior_params['std']] + elif prior_params['prior_type'] == 'lognormal': + bounds = [prior_params['shift'], prior_params['shift'] + np.exp(prior_params['mu'] + BOUNDS_GAUSS_NSIGMA * prior_params['sigma'])] + elif prior_params['prior_type'] == 'positive': + bounds = [PRIOR_MIN, BOUNDS_MAX] + elif prior_params['prior_type'] == 'none': + bounds = [BOUNDS_MIN,BOUNDS_MAX] + elif prior_params['prior_type'] == 'fixed': + bounds = [1.0, 1.0] + else: + print('Prior type not recognized!') + bounds = [BOUNDS_MIN,BOUNDS_MAX] + + return bounds + +def prior_func(x, prior_params): + """Compute the value of a 1-D prior P(x) at a specified value x. + + Args: + x (float): Value at which to compute the prior + prior_params (dict): Dictionary with information about the prior + + Returns: + float: Prior value P(x) + """ + + if prior_params['prior_type'] == 'flat': + return (x >= prior_params['min']) * (x <= prior_params['max']) * 1.0/(prior_params['max'] - prior_params['min']) + PRIOR_MIN + elif prior_params['prior_type'] == 'gauss': + return 1./((2.*np.pi)**0.5 * prior_params['std']) * np.exp(-(x - prior_params['mean'])**2/(2.*prior_params['std']**2)) + elif prior_params['prior_type'] == 'exponential': + return (1./prior_params['std'] * np.exp(-x/prior_params['std'])) * (x >= 0.0) + PRIOR_MIN + elif prior_params['prior_type'] == 'lognormal': + return (x > prior_params['shift']) * ( + 1.0/((2.0*np.pi)**0.5 * prior_params['sigma'] * (x - prior_params['shift'])) + * np.exp( -(np.log(x - prior_params['shift']) - prior_params['mu'])**2/(2.0 * prior_params['sigma']**2) ) ) + elif prior_params['prior_type'] == 'positive': + return (x >= 0.0) * 1.0 + PRIOR_MIN + elif prior_params['prior_type'] == 'none': + return 1.0 + elif prior_params['prior_type'] == 'fixed': + return 1.0 + else: + print('Prior not recognized!') + return 1.0 + +def prior_grad_func(x, prior_params): + """Compute the value of the derivative of a 1-D prior, dP/dx at a specified value x. + + Args: + x (float): Value at which to compute the prior derivative + prior_params (dict): Dictionary with information about the prior + + Returns: + float: Prior derivative value dP/dx(x) + """ + + if prior_params['prior_type'] == 'flat': + return 0.0 + elif prior_params['prior_type'] == 'gauss': + return -(x - prior_params['mean'])/((2.*np.pi)**0.5 * prior_params['std']**3) * np.exp(-(x - prior_params['mean'])**2/(2.*prior_params['std']**2)) + elif prior_params['prior_type'] == 'exponential': + return (-1./prior_params['std']**2 * np.exp(-x/prior_params['std'])) * (x >= 0.0) + elif prior_params['prior_type'] == 'lognormal': + return (x > prior_params['shift']) * ( + (prior_params['mu'] - prior_params['sigma']**2 - np.log(x - prior_params['shift'])) + / ((2.0*np.pi)**0.5 * prior_params['sigma']**3 * (x - prior_params['shift'])**2) + * np.exp( -(np.log(x - prior_params['shift']) - prior_params['mu'])**2/(2.0 * prior_params['sigma']**2) ) ) + elif prior_params['prior_type'] == 'positive': + return 0.0 + elif prior_params['prior_type'] == 'none': + return 0.0 + elif prior_params['prior_type'] == 'fixed': + return 0.0 + else: + print('Prior not recognized!') + return 0.0 + +def transform_param(x, x_prior, inverse=True): + """Compute a specified coordinate transformation T(x) of a parameter value x + + Args: + x (float): Untransformed value + x_prior (dict): Dictionary with information about the transformation + inverse (bool): Whether to compute the forward or inverse transform. + + Returns: + float: Transformed parameter value + """ + + try: + transform = x_prior['transform'] + except: + transform = 'none' + pass + + if transform == 'log': + if inverse: + return np.exp(x) + else: + return np.log(x) + elif transform == 'cdf': + if inverse: + return cdf_inverse(x, x_prior) + else: + return cdf(x, x_prior) + else: + return x + +def transform_grad_param(x, x_prior): + """Compute the gradient of a specified coordinate transformation T(x) of a parameter value x + + Args: + x (float): Untransformed value + x_prior (dict): Dictionary with information about the transformation + + Returns: + float: Gradient of transformation, dT/dx(x) + """ + + try: + transform = x_prior['transform'] + except: + transform = 'none' + pass + + if transform == 'log': + return np.exp(x) + elif transform == 'cdf': + return 1.0/prior_func(transform_param(x,x_prior),x_prior) + else: + return 1.0 + +################################################################################################## +# Helper functions +################################################################################################## +def shrink_prior(prior, model, shrink=0.1): + """Shrink a specified prior volume by centering on a specified fitted model + + Args: + prior (list): Model prior (list of dictionaries, one per model component) + model (Model): Model to draw central values from + shrink (float): Factor to shrink each prior width by + + Returns: + prior (list): Model prior with restricted volume + """ + + prior_shrunk = copy.deepcopy(prior) + f = 1.0 + + #TODO: this doesn't work for beta lists yet! + + for j in range(len(prior_shrunk)): + for key in prior_shrunk[j].keys(): + if prior_shrunk[j][key]['prior_type'] == 'flat': + x = model.params[j][key] + w = prior_shrunk[j][key]['max'] - prior_shrunk[j][key]['min'] + prior_shrunk[j][key]['min'] = x - w/2 + prior_shrunk[j][key]['max'] = x + w/2 + if prior_shrunk[j][key]['min'] < prior[j][key]['min']: prior_shrunk[j][key]['min'] = prior[j][key]['min'] + if prior_shrunk[j][key]['max'] > prior[j][key]['max']: prior_shrunk[j][key]['max'] = prior[j][key]['max'] + f *= (prior_shrunk[j][key]['max'] - prior_shrunk[j][key]['min'])/w + else: + pass + + print('(New Prior Volume)/(Original Prior Volume:',f) + + return prior_shrunk + +def selfcal(Obsdata, model, + gain_init=None, gain_prior=None, + minimizer_func='scipy.optimize.minimize', minimizer_kwargs=None, + bounds=None, use_bounds=True, + processes=-1, msgtype='bar', quiet=True, **kwargs): + """Self-calibrate a specified observation to a given model, accounting for gain priors + + Args: + + Returns: + """ + + # This is just a convenience function. It will call modeler_func() scan-by-scan fitting only gains. + # This function differs from ehtim.calibrating.self_cal in the inclusion of gain priors + tlist = Obsdata.tlist() + res_list = [] + for j in range(len(tlist)): + if msgtype not in ['none','']: + prog_msg(j, len(tlist), msgtype, j-1) + obs = Obsdata.copy() + obs.data = tlist[j] + res_list.append(modeler_func(obs, model, model_prior=None, d1='amp', + fit_model=False, fit_gains=True,gain_init=gain_init,gain_prior=gain_prior, + minimizer_func=minimizer_func, minimizer_kwargs=minimizer_kwargs, + bounds=bounds, use_bounds=use_bounds, processes=-1, quiet=quiet, **kwargs)) + + # Assemble a single caltable to return + allsites = Obsdata.tarr['site'] + caldict = res_list[0]['caltable'].data + for j in range(1,len(tlist)): + row = res_list[j]['caltable'].data + for site in allsites: + try: dat = row[site] + except KeyError: continue + + try: caldict[site] = np.append(caldict[site], row[site]) + except KeyError: caldict[site] = dat + + caltable = ehtim.caltable.Caltable(obs.ra, obs.dec, obs.rf, obs.bw, caldict, obs.tarr, + source=obs.source, mjd=obs.mjd, timetype=obs.timetype) + + return caltable + +def make_param_map(model_init, model_prior, minimizer_func, fit_model, fit_pol=False, fit_cpol=False): + # Define the mapping between solved parameters and the model + # Each fitted model parameter can be rescaled to give values closer to order unity + param_map = [] # Define mapping for every fitted parameter: model component #, parameter name, rescale multiplier internal, unit, rescale multiplier external + param_mask = [] # True or False for whether to fit each model parameter (because the gradient is computed for all model parameters) + for j in range(model_init.N_models()): + params = model.model_params(model_init.models[j],model_init.params[j], fit_pol=fit_pol, fit_cpol=fit_cpol) + for param in params: + if fit_model == False: + param_mask.append(False) + elif model_prior[j][param]['prior_type'] != 'fixed': + param_mask.append(True) + param_type = param + if len(param_type.split('_')) == 2 and param_type not in PARAM_DETAILS: + param_type = param_type.split('_')[1] + try: + if model_prior[j][param].get('transform','') == 'cdf' or minimizer_func in ['dynesty_static','dynesty_dynamic','pymc3']: + param_map.append([j,param,1,PARAM_DETAILS[param_type][1],PARAM_DETAILS[param_type][0]]) + else: + param_map.append([j,param,PARAM_DETAILS[param_type][0],PARAM_DETAILS[param_type][1],PARAM_DETAILS[param_type][0]]) + except: + param_map.append([j,param,1,'',1]) + pass + else: + param_mask.append(False) + return (param_map, param_mask) + +def compute_likelihood_constants(d1, d2, d3, sigma1, sigma2, sigma3): + # Compute the correct data weights (hyperparameters) and the correct extra constant for the log-likelihood + alpha_d1 = alpha_d2 = alpha_d3 = ln_norm1 = ln_norm2 = ln_norm3 = 0.0 + + try: + alpha_d1 = 0.5 * len(sigma1) + ln_norm1 = -np.sum(np.log((2.0*np.pi)**0.5 * sigma1)) + except: pass + try: + alpha_d2 = 0.5 * len(sigma2) + ln_norm2 = -np.sum(np.log((2.0*np.pi)**0.5 * sigma2)) + except: pass + try: + alpha_d3 = 0.5 * len(sigma3) + ln_norm3 = -np.sum(np.log((2.0*np.pi)**0.5 * sigma3)) + except: pass + + if d1 in ['vis','bs','m','pvis','rrll','llrr','lrll','rlll','lrrr','rlrr','polclosure']: + alpha_d1 *= 2 + ln_norm1 *= 2 + if d2 in ['vis','bs','m','pvis','rrll','llrr','lrll','rlll','lrrr','rlrr','polclosure']: + alpha_d2 *= 2 + ln_norm2 *= 2 + if d3 in ['vis','bs','m','pvis','rrll','llrr','lrll','rlll','lrrr','rlrr','polclosure']: + alpha_d3 *= 2 + ln_norm3 *= 2 + ln_norm = ln_norm1 + ln_norm2 + ln_norm3 + + return (alpha_d1, alpha_d2, alpha_d3, ln_norm) + +def default_gain_prior(sites): + print('No gain prior specified. Defaulting to ' + str(GAIN_PRIOR_DEFAULT) + ' for all sites.') + gain_prior = {} + for site in sites: + gain_prior[site] = GAIN_PRIOR_DEFAULT + return gain_prior + +def caltable_to_gains(caltab, gain_list): + # Generate an ordered list of gains from a caltable + # gain_list is a set of tuples (time, site) + gains = [np.abs(caltab.data[site]['rscale'][caltab.data[site]['time'] == time][0]) - 1.0 for (time, site) in gain_list] + return gains + +def make_gain_map(Obsdata, gain_prior): + # gain_list gives all unique (time,site) pairs + # gains_t1 gives the gain index for the first site in each measurement + # gains_t2 gives the gain index for the second site in each measurement + gain_list = [] + for j in range(len(Obsdata.data)): + if ([Obsdata.data[j]['time'],Obsdata.data[j]['t1']] not in gain_list) and (gain_prior[Obsdata.data[j]['t1']]['prior_type'] != 'fixed'): + gain_list.append([Obsdata.data[j]['time'],Obsdata.data[j]['t1']]) + if ([Obsdata.data[j]['time'],Obsdata.data[j]['t2']] not in gain_list) and (gain_prior[Obsdata.data[j]['t2']]['prior_type'] != 'fixed'): + gain_list.append([Obsdata.data[j]['time'],Obsdata.data[j]['t2']]) + + # Now determine the appropriate mapping; use the final index for all ignored gains, which default to 1 + def gain_index(j, tnum): + try: + return gain_list.index([Obsdata.data[j]['time'],Obsdata.data[j][tnum]]) + except: + return len(gain_list) + + gains_t1 = [gain_index(j, 't1') for j in range(len(Obsdata.data))] + gains_t2 = [gain_index(j, 't2') for j in range(len(Obsdata.data))] + + return (gain_list, gains_t1, gains_t2) + +def make_bounds(model_prior, param_map, gain_prior, gain_list, n_gains, leakage_fit, leakage_prior): + bounds = [] + for j in range(len(param_map)): + pm = param_map[j] + pb = param_bounds(model_prior[pm[0]][pm[1]]) + if (model_prior[pm[0]][pm[1]]['prior_type'] not in ['positive','none','fixed']) and (model_prior[pm[0]][pm[1]].get('transform','') != 'cdf'): + pb[0] = transform_param(pb[0]/pm[2], model_prior[pm[0]][pm[1]], inverse=False) + pb[1] = transform_param(pb[1]/pm[2], model_prior[pm[0]][pm[1]], inverse=False) + bounds.append(pb) + for j in range(n_gains): + pb = param_bounds(gain_prior[gain_list[j][1]]) + if (gain_prior[gain_list[j][1]]['prior_type'] not in ['positive','none','fixed']) and (gain_prior[gain_list[j][1]].get('transform','') != 'cdf'): + pb[0] = transform_param(pb[0], gain_prior[gain_list[j][1]], inverse=False) + pb[1] = transform_param(pb[1], gain_prior[gain_list[j][1]], inverse=False) + bounds.append(pb) + for j in range(len(leakage_fit)): + for cpart in ['re','im']: + prior = leakage_prior[leakage_fit[j][0]][leakage_fit[j][1]][cpart] + pb = param_bounds(prior) + if (prior['prior_type'] not in ['positive','none','fixed']) and (prior.get('transform','') != 'cdf'): + pb[0] = transform_param(pb[0], prior, inverse=False) + pb[1] = transform_param(pb[1], prior, inverse=False) + bounds.append(pb) + + return np.array(bounds) + +# Determine multiplicative factor for the gains (amplitude only) +def gain_factor(dtype,gains,gains_t1,gains_t2, fit_or_marginalize_gains): + global globdict + + if not fit_or_marginalize_gains: + if globdict['gain_init'] == None: + return 1 + else: + gains = globdict['gain_init'] + + if globdict['marginalize_gains']: + gains = globdict['gain_init'] + + if dtype in ['amp','vis']: + gains_wzero = np.append(gains,0.0) + return (1.0 + gains_wzero[gains_t1])*(1.0 + gains_wzero[gains_t2]) + else: + return 1 + +def gain_factor_separate(dtype,gains,gains_t1,gains_t2, fit_or_marginalize_gains): + # Determine the pair of multiplicative factors for the gains (amplitude only) + # Note: these are not displaced by unity! + global globdict + + if not fit_or_marginalize_gains: + if globdict['gain_init'] == None: + return (0., 0.) + else: + gains = globdict['gain_init'] + + if globdict['marginalize_gains']: + gains = globdict['gain_init'] + + if dtype in ['amp','vis']: + gains_wzero = np.append(gains,0.0) + return (gains_wzero[gains_t1], gains_wzero[gains_t2]) + else: + return (0, 0) + +def prior_leakage(leakage, leakage_fit, leakage_prior, fit_leakage): + # Compute the log-prior contribution from the fitted leakage terms + if fit_leakage: + cparts = ['re','im'] + return np.sum([np.log(prior_func(leakage[j], leakage_prior[leakage_fit[j//2][0]][leakage_fit[j//2][1]][cparts[j%2]])) for j in range(len(leakage))]) + else: + return 0.0 + +def prior_leakage_grad(leakage, leakage_fit, leakage_prior, fit_leakage): + # Compute the log-prior contribution to the gradient from the leakages + if fit_leakage: + cparts = ['re','im'] + f = np.array([prior_func(leakage[j], leakage_prior[leakage_fit[j//2][0]][leakage_fit[j//2][1]][cparts[j%2]]) for j in range(len(leakage))]) + df = np.array([prior_grad_func(leakage[j], leakage_prior[leakage_fit[j//2][0]][leakage_fit[j//2][1]][cparts[j%2]]) for j in range(len(leakage))]) + return df/f + else: + return [] + +def prior_gain(gains, gain_list, gain_prior, fit_gains): + # Compute the log-prior contribution from the gains + if fit_gains: + return np.sum([np.log(prior_func(gains[j], gain_prior[gain_list[j][1]])) for j in range(len(gains))]) + else: + return 0.0 + +def prior_gain_grad(gains, gain_list, gain_prior, fit_gains): + # Compute the log-prior contribution to the gradient from the gains + if fit_gains: + f = np.array([prior_func(gains[j], gain_prior[gain_list[j][1]]) for j in range(len(gains))]) + df = np.array([prior_grad_func(gains[j], gain_prior[gain_list[j][1]]) for j in range(len(gains))]) + return df/f + else: + return [] + +def transform_params(params, param_map, minimizer_func, model_prior, inverse=True): + if minimizer_func not in ['dynesty_static','dynesty_dynamic','pymc3']: + return [transform_param(params[j], model_prior[param_map[j][0]][param_map[j][1]], inverse=inverse) for j in range(len(params))] + else: + # For dynesty or pymc3, over-ride all specified parameter transformations to assume CDF mapping to the hypercube + # However, the passed parameters to the objective function and gradient are *not* transformed (i.e., they are not in the hypercube), thus the transformation does not need to be inverted + return params + +def set_params(params, trial_model, param_map, minimizer_func, model_prior): + tparams = transform_params(params, param_map, minimizer_func, model_prior) + + for j in range(len(params)): + if param_map[j][1] in trial_model.params[param_map[j][0]].keys(): + trial_model.params[param_map[j][0]][param_map[j][1]] = tparams[j] * param_map[j][2] + else: # In this case, the parameter is a list of complex numbers, so the real/imaginary or abs/arg components need to be assigned + if param_map[j][1].find('cpol') != -1: + param_type = 'beta_list_cpol' + idx = int(param_map[j][1].split('_')[0][8:]) + elif param_map[j][1].find('pol') != -1: + param_type = 'beta_list_pol' + idx = int(param_map[j][1].split('_')[0][7:]) + (len(trial_model.params[param_map[j][0]][param_type])-1)//2 + elif param_map[j][1].find('beta') != -1: + param_type = 'beta_list' + idx = int(param_map[j][1].split('_')[0][4:]) - 1 + else: + raise Exception('Unsure how to interpret ' + param_map[j][1]) + + curval = trial_model.params[param_map[j][0]][param_type][idx] + if param_map[j][1][-2:] == 're': + trial_model.params[param_map[j][0]][param_type][idx] = tparams[j] * param_map[j][2] + np.imag(curval)*1j + elif param_map[j][1][-2:] == 'im': + trial_model.params[param_map[j][0]][param_type][idx] = tparams[j] * param_map[j][2] * 1j + np.real(curval) + elif param_map[j][1][-3:] == 'abs': + trial_model.params[param_map[j][0]][param_type][idx] = tparams[j] * param_map[j][2] * np.exp(1j * np.angle(curval)) + elif param_map[j][1][-3:] == 'arg': + trial_model.params[param_map[j][0]][param_type][idx] = np.abs(curval) * np.exp(1j * tparams[j] * param_map[j][2]) + else: + print('Parameter ' + param_map[j][1] + ' not understood!') + +# Define prior +def prior(params, param_map, model_prior, minimizer_func): + tparams = transform_params(params, param_map, minimizer_func, model_prior) + return np.sum([np.log(prior_func(tparams[j]*param_map[j][2], model_prior[param_map[j][0]][param_map[j][1]])) for j in range(len(params))]) + +def prior_grad(params, param_map, model_prior, minimizer_func): + tparams = transform_params(params, param_map, minimizer_func, model_prior) + f = np.array([prior_func(tparams[j]*param_map[j][2], model_prior[param_map[j][0]][param_map[j][1]]) for j in range(len(params))]) + df = np.array([prior_grad_func(tparams[j]*param_map[j][2], model_prior[param_map[j][0]][param_map[j][1]]) for j in range(len(params))]) + return df/f + +# Define constraint functions +def flux_constraint(trial_model, alpha_flux, flux): + if alpha_flux == 0.0: + return 0.0 + + return ((trial_model.total_flux() - flux)/flux)**2 + +def flux_constraint_grad(trial_model, alpha_flux, flux, params, param_map): + if alpha_flux == 0.0: + return 0.0 + + fluxmask = np.zeros_like(params) + for j in range(len(param_map)): + if param_map[j][1] == 'F0': + fluxmask[j] = 1.0 + + return 2.0 * (trial_model.total_flux() - flux)/flux * fluxmask + +################################################################################################## +# Define the chi^2 and chi^2 gradient functions +################################################################################################## +def laplace_approximation(trial_model, dtype, data, uv, sigma, gains_t1, gains_t2): + # Compute the approximate contribution to the log-likelihood by marginalizing over gains + global globdict + + if globdict['marginalize_gains'] == True and dtype == 'amp': + # Add the log-likelihood term from analytic gain marginalization + # Create the Hessian matrix for the argument of the exponential + gain_hess = np.zeros((len(globdict['gain_list']), len(globdict['gain_list']))) + + # Add the terms from the likelihood + gain = gain_factor(dtype,None,gains_t1,gains_t2,True) + amp_model = np.abs(trial_model.sample_uv(uv[:,0],uv[:,1])) + amp_bar = gain*data + sigma_bar = gain*sigma + (g1, g2) = gain_factor_separate(dtype,None,gains_t1,gains_t2,True) + + # Each amplitude *measurement* (not fitted gain parameter!) contributes to the hessian in four places; two diagonal and two off-diagonal + for j in range(len(gain)): + gain_hess[gains_t1[j],gains_t1[j]] += amp_model[j] * (3.0 * amp_model[j] - 2.0 * amp_bar[j])/((1.0 + g1[j])**2 * sigma_bar[j]**2) + gain_hess[gains_t2[j],gains_t2[j]] += amp_model[j] * (3.0 * amp_model[j] - 2.0 * amp_bar[j])/((1.0 + g2[j])**2 * sigma_bar[j]**2) + gain_hess[gains_t1[j],gains_t2[j]] += amp_model[j] * (2.0 * amp_model[j] - amp_bar[j])/((1.0 + g1[j])*(1.0 + g2[j]) * sigma_bar[j]**2) + gain_hess[gains_t2[j],gains_t1[j]] += amp_model[j] * (2.0 * amp_model[j] - amp_bar[j])/((1.0 + g1[j])*(1.0 + g2[j]) * sigma_bar[j]**2) + + # Add contributions from the prior to the diagonal. This ranges over the fitted gain parameters. + # Note: for the Laplace approximation, only Gaussian gain priors have any effect! + for j in range(len(globdict['gain_list'])): + t = globdict['gain_list'][j][1] + if globdict['gain_prior'][t]['prior_type'] == 'gauss': + gain_hess[j,j] += 1.0/globdict['gain_prior'][t]['std'] + elif globdict['gain_prior'][t]['prior_type'] == 'flat': + gain_hess[j,j] += 0.0 + elif globdict['gain_prior'][t]['prior_type'] == 'exponential': + gain_hess[j,j] += 0.0 + elif globdict['gain_prior'][t]['prior_type'] == 'fixed': + gain_hess[j,j] += 0.0 + else: + raise Exception('Gain prior not implemented!') + return np.log((2.0 * np.pi)**(len(gain)/2.0) * np.abs(np.linalg.det(gain_hess))**-0.5) + else: + return 0.0 + +def laplace_list(): + global globdict + l1 = laplace_approximation(globdict['trial_model'], globdict['d1'], globdict['data1'], globdict['uv1'], globdict['sigma1'], globdict['gains_t1'], globdict['gains_t2']) + l2 = laplace_approximation(globdict['trial_model'], globdict['d2'], globdict['data2'], globdict['uv2'], globdict['sigma2'], globdict['gains_t1'], globdict['gains_t2']) + l3 = laplace_approximation(globdict['trial_model'], globdict['d3'], globdict['data3'], globdict['uv3'], globdict['sigma3'], globdict['gains_t1'], globdict['gains_t2']) + return (l1, l2, l3) + +def chisq_wgain(trial_model, dtype, data, uv, sigma, jonesdict, gains, gains_t1, gains_t2, fit_or_marginalize_gains): + global globdict + gain = gain_factor(dtype,gains,gains_t1,gains_t2,fit_or_marginalize_gains) + log_likelihood = chisq(trial_model, uv, gain*data, gain*sigma, dtype, jonesdict) + return log_likelihood + +def chisqgrad_wgain(trial_model, dtype, data, uv, sigma, jonesdict, gains, gains_t1, gains_t2, fit_or_marginalize_gains, param_mask, fit_pol=False, fit_cpol=False, fit_leakage=False): + gain = gain_factor(dtype,gains,gains_t1,gains_t2,fit_or_marginalize_gains) + return chisqgrad(trial_model, uv, gain*data, gain*sigma, jonesdict, dtype, param_mask, fit_or_marginalize_gains, gains, gains_t1, gains_t2, fit_pol, fit_cpol, fit_leakage) + +def chisq_list(gains): + global globdict + chi2_1 = chisq_wgain(globdict['trial_model'], globdict['d1'], globdict['data1'], globdict['uv1'], globdict['sigma1'], globdict['jonesdict1'], gains, globdict['gains_t1'], globdict['gains_t2'], globdict['fit_gains'] + globdict['marginalize_gains']) + chi2_2 = chisq_wgain(globdict['trial_model'], globdict['d2'], globdict['data2'], globdict['uv2'], globdict['sigma2'], globdict['jonesdict2'], gains, globdict['gains_t1'], globdict['gains_t2'], globdict['fit_gains'] + globdict['marginalize_gains']) + chi2_3 = chisq_wgain(globdict['trial_model'], globdict['d3'], globdict['data3'], globdict['uv3'], globdict['sigma3'], globdict['jonesdict3'], gains, globdict['gains_t1'], globdict['gains_t2'], globdict['fit_gains'] + globdict['marginalize_gains']) + return (chi2_1, chi2_2, chi2_3) + +def update_leakage(leakage): + # This function updates the 'jonesdict' entries based on current leakage estimates + # leakage is list of the fitted parameters (re and im are separate) + # station_leakages is the dictionary containing all station leakages, some of which may be fixed + global globdict + if len(leakage) == 0: return + + station_leakages = globdict['station_leakages'] + leakage_fit = globdict['leakage_fit'] + # First, update the entries in the leakage dictionary + for j in range(len(leakage)//2): + station_leakages[leakage_fit[j][0]][leakage_fit[j][1]] = leakage[2*j] + 1j * leakage[2*j + 1] + + # Now, recompute the jonesdict objects + for j in range(1,4): + jonesdict = globdict['jonesdict' + str(j)] + if jonesdict is not None: + if type(jonesdict) is dict: + jonesdict['DR1'] = np.array([station_leakages[jonesdict['t1'][_]]['R'] for _ in range(len(jonesdict['t1']))]) + jonesdict['DR2'] = np.array([station_leakages[jonesdict['t2'][_]]['R'] for _ in range(len(jonesdict['t1']))]) + jonesdict['DL1'] = np.array([station_leakages[jonesdict['t1'][_]]['L'] for _ in range(len(jonesdict['t1']))]) + jonesdict['DL2'] = np.array([station_leakages[jonesdict['t2'][_]]['L'] for _ in range(len(jonesdict['t1']))]) + jonesdict['leakage_fit'] = globdict['leakage_fit'] + else: + # In this case, the data product requires a list of jonesdicts + for jonesdict2 in jonesdict: + jonesdict2['DR1'] = np.array([station_leakages[jonesdict2['t1'][_]]['R'] for _ in range(len(jonesdict2['t1']))]) + jonesdict2['DR2'] = np.array([station_leakages[jonesdict2['t2'][_]]['R'] for _ in range(len(jonesdict2['t1']))]) + jonesdict2['DL1'] = np.array([station_leakages[jonesdict2['t1'][_]]['L'] for _ in range(len(jonesdict2['t1']))]) + jonesdict2['DL2'] = np.array([station_leakages[jonesdict2['t2'][_]]['L'] for _ in range(len(jonesdict2['t1']))]) + jonesdict2['leakage_fit'] = globdict['leakage_fit'] + +################################################################################################## +# Define the objective function and gradient +################################################################################################## +def objfunc(params, force_posterior=False): + global globdict + # Note: model parameters can have transformations applied; gains and leakage do not + set_params(params[:globdict['n_params']], globdict['trial_model'], globdict['param_map'], globdict['minimizer_func'], globdict['model_prior']) + gains = params[globdict['n_params']:(globdict['n_params'] + globdict['n_gains'])] + leakage = params[(globdict['n_params'] + globdict['n_gains']):] + update_leakage(leakage) + + if globdict['marginalize_gains']: + # Ugh, the use of global variables totally messes this up + _globdict = globdict + # This doesn't handle the passed gain_init properly because the dimensions are incorrect + _globdict['gain_init'] = caltable_to_gains(selfcal(globdict['Obsdata'], globdict['trial_model'], gain_init=None, gain_prior=globdict['gain_prior'], msgtype='none'),globdict['gain_list']) + globdict = _globdict + + (chi2_1, chi2_2, chi2_3) = chisq_list(gains) + datterm = ( globdict['alpha_d1'] * chi2_1 + + globdict['alpha_d2'] * chi2_2 + + globdict['alpha_d3'] * chi2_3) + + if globdict['marginalize_gains']: + (l1, l2, l3) = laplace_list() + datterm += l1 + l2 + l3 + + if (globdict['minimizer_func'] not in ['dynesty_static','dynesty_dynamic','pymc3']) or force_posterior: + priterm = prior(params[:globdict['n_params']], globdict['param_map'], globdict['model_prior'], globdict['minimizer_func']) + priterm += prior_gain(params[globdict['n_params']:(globdict['n_params'] + globdict['n_gains'])], globdict['gain_list'], globdict['gain_prior'], globdict['fit_gains']) + priterm += prior_leakage(params[(globdict['n_params'] + globdict['n_gains']):], globdict['leakage_fit'], globdict['leakage_prior'], globdict['fit_leakage']) + else: + priterm = 0.0 + fluxterm = globdict['alpha_flux'] * flux_constraint(globdict['trial_model'], globdict['alpha_flux'], globdict['flux']) + + return datterm - priterm + fluxterm - globdict['ln_norm'] + +def objgrad(params): + global globdict + set_params(params[:globdict['n_params']], globdict['trial_model'], globdict['param_map'], globdict['minimizer_func'], globdict['model_prior']) + gains = params[globdict['n_params']:(globdict['n_params'] + globdict['n_gains'])] + leakage = params[(globdict['n_params'] + globdict['n_gains']):] + update_leakage(leakage) + + datterm = ( globdict['alpha_d1'] * chisqgrad_wgain(globdict['trial_model'], globdict['d1'], globdict['data1'], globdict['uv1'], globdict['sigma1'], globdict['jonesdict1'], gains, globdict['gains_t1'], globdict['gains_t2'], globdict['fit_gains'] + globdict['marginalize_gains'], globdict['param_mask'], globdict['fit_pol'], globdict['fit_cpol'], globdict['fit_leakage']) + + globdict['alpha_d2'] * chisqgrad_wgain(globdict['trial_model'], globdict['d2'], globdict['data2'], globdict['uv2'], globdict['sigma2'], globdict['jonesdict2'], gains, globdict['gains_t1'], globdict['gains_t2'], globdict['fit_gains'] + globdict['marginalize_gains'], globdict['param_mask'], globdict['fit_pol'], globdict['fit_cpol'], globdict['fit_leakage']) + + globdict['alpha_d3'] * chisqgrad_wgain(globdict['trial_model'], globdict['d3'], globdict['data3'], globdict['uv3'], globdict['sigma3'], globdict['jonesdict3'], gains, globdict['gains_t1'], globdict['gains_t2'], globdict['fit_gains'] + globdict['marginalize_gains'], globdict['param_mask'], globdict['fit_pol'], globdict['fit_cpol'], globdict['fit_leakage'])) + + if globdict['minimizer_func'] not in ['dynesty_static','dynesty_dynamic','pymc3']: + priterm = np.concatenate([prior_grad(params[:globdict['n_params']], globdict['param_map'], + globdict['model_prior'], globdict['minimizer_func']), + prior_gain_grad(params[globdict['n_params']:(globdict['n_params'] + globdict['n_gains'])], + globdict['gain_list'], globdict['gain_prior'], globdict['fit_gains']), + prior_leakage_grad(params[(globdict['n_params'] + globdict['n_gains']):], globdict['leakage_fit'], + globdict['leakage_prior'], globdict['fit_leakage'])]) + else: + priterm = 0.0 + fluxterm = globdict['alpha_flux'] * flux_constraint_grad(params, globdict['alpha_flux'], globdict['flux'], params, globdict['param_map']) + + grad = datterm - priterm + fluxterm + + if globdict['minimizer_func'] not in ['dynesty_static','dynesty_dynamic','pymc3']: + for j in range(globdict['n_params']): + grad[j] *= globdict['param_map'][j][2] * transform_grad_param(params[j], globdict['model_prior'][globdict['param_map'][j][0]][globdict['param_map'][j][1]]) + else: + # For dynesty or pymc3, over-ride all specified parameter transformations to assume CDF + # However, the passed parameters are *not* transformed (i.e., they are not in the hypercube) + # The Jacobian still needs to account for the parameter transformation + for j in range(len(params)): + if j < globdict['n_params']: + j2 = j + x = params[j2] + prior_params = globdict['model_prior'][globdict['param_map'][j2][0]][globdict['param_map'][j2][1]] + grad[j] /= prior_func(x,prior_params) + elif j < globdict['n_params'] + globdict['n_gains']: + j2 = j-globdict['n_params'] + x = gains[j2] + prior_params = globdict['gain_prior'][globdict['gain_list'][j2][1]] + grad[j] /= prior_func(x, prior_params) + else: + cparts = ['re','im'] + j2 = j-globdict['n_params']-globdict['n_gains'] + x = leakage[j2] + prior_params = globdict['leakage_prior'][globdict['leakage_fit'][j2//2][0]][globdict['leakage_fit'][j2//2][1]][cparts[j2%2]] + grad[j] /= prior_func(x, prior_params) + + if globdict['test_gradient']: + print('Testing the gradient at ',params) + import copy + dx = 1e-5 + grad_numeric = np.zeros(len(grad)) + f1 = objfunc(params) + print('Objective Function:',f1) + print('\nNumeric Gradient Check: Analytic Numeric') + for j in range(len(grad)): + if globdict['minimizer_func'] in ['dynesty_static','dynesty_dynamic','pymc3']: + dx = np.abs(params[j]) * 1e-6 + + params2 = copy.deepcopy(params) + params2[j] += dx + f2 = objfunc(params2) + grad_numeric[j] = (f2 - f1)/dx + + if globdict['minimizer_func'] in ['dynesty_static','dynesty_dynamic','pymc3']: + if j < globdict['n_params']: + j2 = j + x = params[j2] + prior_params = globdict['model_prior'][globdict['param_map'][j2][0]][globdict['param_map'][j2][1]] + grad_numeric[j] /= prior_func(x,prior_params) + elif j < globdict['n_params'] + globdict['n_gains']: + j2 = j-globdict['n_params'] + x = gains[j2] + prior_params = globdict['gain_prior'][globdict['gain_list'][j2][1]] + grad_numeric[j] /= prior_func(x, prior_params) + else: + cparts = ['re','im'] + j2 = j-globdict['n_params']-globdict['n_gains'] + x = leakage[j2] + prior_params = globdict['leakage_prior'][globdict['leakage_fit'][j2//2][0]][globdict['leakage_fit'][j2//2][1]][cparts[j2%2]] + grad_numeric[j] /= prior_func(x, prior_params) + + if j < globdict['n_params']: + print('\nNumeric Gradient Check:',globdict['param_map'][j][0],globdict['param_map'][j][1],grad[j],grad_numeric[j]) + else: + print('\nNumeric Gradient Check:',grad[j],grad_numeric[j]) + + return grad + +################################################################################################## +# Modeler +################################################################################################## +def modeler_func(Obsdata, model_init, model_prior, + d1='vis', d2=False, d3=False, + normchisq = False, alpha_d1=0, alpha_d2=0, alpha_d3=0, + flux=1.0, alpha_flux=0, + fit_model=True, fit_pol=False, fit_cpol=False, + fit_gains=False,marginalize_gains=False,gain_init=None,gain_prior=None, + fit_leakage=False, leakage_init=None, leakage_prior=None, + minimizer_func='scipy.optimize.minimize', + minimizer_kwargs=None, + bounds=None, use_bounds=False, + processes=-1, + test_gradient=False, quiet=False, **kwargs): + + """Fit a specified model. + + Args: + Obsdata (Obsdata): The Obsdata object with VLBI data + model_init (Model): The Model object to fit + model_prior (dict): Priors for each model parameter + + d1 (str): The first data term; options are 'vis', 'bs', 'amp', 'cphase', 'cphase_diag' 'camp', 'logcamp', 'logcamp_diag', 'm' + d2 (str): The second data term; options are 'vis', 'bs', 'amp', 'cphase', 'cphase_diag' 'camp', 'logcamp', 'logcamp_diag', 'm' + d3 (str): The third data term; options are 'vis', 'bs', 'amp', 'cphase', 'cphase_diag' 'camp', 'logcamp', 'logcamp_diag', 'm' + + normchisq (bool): If False (default), automatically assign weights alpha_d1-3 to match the true log-likelihood. + alpha_d1 (float): The first data term weighting. + alpha_d2 (float): The second data term weighting. Default value of zero will automatically assign weights to match the true log-likelihood. + alpha_d2 (float): The third data term weighting. Default value of zero will automatically assign weights to match the true log-likelihood. + + flux (float): Total flux of the fitted model + alpha_flux (float): Hyperparameter controlling how strongly to constrain that the total flux matches the specified flux. + + fit_model (bool): Whether or not to fit the model parameters + fit_pol (bool): Whether or not to fit linear polarization parameters + fit_cpol (bool): Whether or not to fit circular polarization parameters + fit_gains (bool): Whether or not to fit time-dependent amplitude gains for each station + marginalize_gains (bool): Whether or not to perform analytic gain marginalization (via the Laplace approximation to the posterior) + + gain_init (list or caltable): Initial gain amplitudes to apply; these can be specified even if gains aren't fitted + gain_prior (dict): Dictionary with the gain prior for each site. + + minimizer_func (str): Minimizer function to use. Current options are: + 'scipy.optimize.minimize' + 'scipy.optimize.dual_annealing' + 'scipy.optimize.basinhopping' + 'dynesty_static' + 'dynesty_dynamic' + 'pymc3' + minimizer_kwargs (dict): kwargs passed to the minimizer. + + bounds (list): List of parameter bounds for the fitted parameters (will automatically compute if needed) + use_bounds (bool): Whether or not to use bounds when fitting (required for some minimizers) + + processes (int): Number of processes to use for a multiprocessing pool. -1 disables multiprocessing; 0 uses all that are available. Only used for dynesty. + + Returns: + dict: Dictionary with fitted model ('model') and other diagnostics that are minimizer-dependent + """ + + global nit, globdict + nit = n_params = 0 + ln_norm = 0.0 + + if fit_model == False and fit_gains == False and fit_leakage == False: + raise Exception('Both fit_model, fit_gains, and fit_leakage are False. Must fit something!') + + if fit_gains == True and marginalize_gains == True: + raise Exception('Both fit_gains and marginalize_gains are True. Cannot do both!') + + if fit_gains == False and marginalize_gains == False and gain_init is not None: + if not quiet: print('Both fit_gains and marginalize_gains are False but gain_init was passed. Applying these gains as a fixed correction!') + + if minimizer_kwargs is None: + minimizer_kwargs = {} + + # Specifications for verbosity during fits + show_updates = kwargs.get('show_updates',True) + update_interval = kwargs.get('update_interval',1) + run_nested_kwargs = kwargs.get('run_nested_kwargs',{}) + + # Make sure data and regularizer options are ok + if not d1 and not d2 and not d3: + raise Exception("Must have at least one data term!") + if (not ((d1 in DATATERMS) or d1==False)) or (not ((d2 in DATATERMS) or d2==False)): + raise Exception("Invalid data term: valid data terms are: " + ' '.join(DATATERMS)) + + # Create the trial model + trial_model = model_init.copy() + + # Define mapping for every fitted parameter: model component index, parameter name, rescale multiplier, unit + (param_map, param_mask) = make_param_map(model_init, model_prior, minimizer_func, fit_model, fit_pol, fit_cpol) + + # Get data and info for the data terms + if type(Obsdata) is obsdata.Obsdata: + (data1, sigma1, uv1, jonesdict1) = chisqdata(Obsdata, d1) + (data2, sigma2, uv2, jonesdict2) = chisqdata(Obsdata, d2) + (data3, sigma3, uv3, jonesdict3) = chisqdata(Obsdata, d3) + elif type(Obsdata) is list: + # Combine a list of observations into one. + # Allow these to be from multiple sources for polarimetric zero-baseline purposes. + # Main thing for different sources is to compute field rotation before combining + def combine_data(d1,s1,u1,j1,d2,s2,u2,j2): + d = np.concatenate([d1,d2]) + s = np.concatenate([s1,s2]) + u = np.concatenate([u1,u2]) + j = j1.copy() + for key in ['fr1', 'fr2', 't1', 't2', 'DR1', 'DR2', 'DL1', 'DL2']: + j[key] = np.concatenate([j1[key],j2[key]]) + return (d, s, u, j) + + (data1, sigma1, uv1, jonesdict1) = chisqdata(Obsdata[0], d1) + (data2, sigma2, uv2, jonesdict2) = chisqdata(Obsdata[0], d2) + (data3, sigma3, uv3, jonesdict3) = chisqdata(Obsdata[0], d3) + for j in range(1,len(Obsdata)): + (data1b, sigma1b, uv1b, jonesdict1b) = chisqdata(Obsdata[j], d1) + (data2b, sigma2b, uv2b, jonesdict2b) = chisqdata(Obsdata[j], d2) + (data3b, sigma3b, uv3b, jonesdict3b) = chisqdata(Obsdata[j], d3) + + if data1b is not False: + (data1, sigma1, uv1, jonesdict1) = combine_data(data1,sigma1,uv1,jonesdict1,data1b,sigma1b,uv1b,jonesdict1b) + if data2b is not False: + (data2, sigma2, uv2, jonesdict2) = combine_data(data2,sigma2,uv2,jonesdict2,data2b,sigma2b,uv2b,jonesdict2b) + if data3b is not False: + (data3, sigma3, uv3, jonesdict3) = combine_data(data3,sigma3,uv3,jonesdict3,data3b,sigma3b,uv3b,jonesdict3b) + + alldata = np.concatenate([_.data for _ in Obsdata]) + Obsdata = Obsdata[0] + Obsdata.data = alldata + else: + raise Exception("Observation format not recognized!") + + if fit_leakage or leakage_init is not None: + # Determine what leakage terms must be fitted. At most, this would be L & R complex leakages terms for every site + # leakage_fit is a list of tuples [site, hand] that will be fitted + leakage_fit = [] + if fit_leakage: + import copy # Error on the next line if this isn't done again. Why python, why?!? + # Start with the list of all sites + sites = list(set(np.concatenate(Obsdata.unpack(['t1','t2']).tolist()))) + + # Add missing entries to leakage_prior + # leakage_prior is a nested dictionary with keys of station, hand, re/im + leakage_prior_init = copy.deepcopy(leakage_prior) + if leakage_prior_init is None: leakage_prior_init = {} + leakage_prior = {} + for s in sites: + leakage_prior[s] = {} + for pol in ['R','L']: + leakage_prior[s][pol] = {} + for cpart in ['re','im']: + # check to see if a prior is specified for the complex part, the pol, or the site (in that order) + if leakage_prior_init.get(s,{}).get(pol,{}).get(cpart,{}).get('prior_type','') != '': + leakage_prior[s][pol][cpart] = leakage_prior_init[s][pol][cpart] + elif leakage_prior_init.get(s,{}).get(pol,{}).get('prior_type','') != '': + leakage_prior[s][pol][cpart] = copy.deepcopy(leakage_prior_init[s][pol]) + elif leakage_prior_init.get(s,{}).get('prior_type','') != '': + leakage_prior[s][pol][cpart] = copy.deepcopy(leakage_prior_init[s]) + else: + leakage_prior[s][pol][cpart] = copy.deepcopy(LEAKAGE_PRIOR_DEFAULT) + + if Obsdata.polrep == 'stokes': + for s in sites: + for pol in ['R','L']: + if leakage_prior[s][pol]['re']['prior_type'] == 'fixed': continue + leakage_fit.append([s,pol]) + else: + vislist = Obsdata.unpack(['t1','t2','rlvis','lrvis']) + # Only fit leakage for sites that include cross hand visibilities + DR = list(set(np.concatenate([vislist[~np.isnan(vislist['rlvis'])]['t1'], vislist[~np.isnan(vislist['lrvis'])]['t2']]))) + DL = list(set(np.concatenate([vislist[~np.isnan(vislist['lrvis'])]['t1'], vislist[~np.isnan(vislist['rlvis'])]['t2']]))) + [leakage_fit.append([s,'R']) for s in DR if leakage_prior[s]['R']['re']['prior_type'] != 'fixed'] + [leakage_fit.append([s,'L']) for s in DL if leakage_prior[s]['L']['re']['prior_type'] != 'fixed'] + sites = list(set(np.concatenate([DR,DL]))) + + if type(leakage_init) is dict: + station_leakages = copy.deepcopy(leakage_init) + else: + station_leakages = {} + + # Add missing entries to station_leakages + for s in sites: + for pol in ['R','L']: + if s not in station_leakages.keys(): + station_leakages[s] = {} + if 'R' not in station_leakages[s].keys(): + station_leakages[s]['R'] = 0.0 + if 'L' not in station_leakages[s].keys(): + station_leakages[s]['L'] = 0.0 + else: + # Disable leakage computations + jonesdict1 = jonesdict2 = jonesdict3 = None + leakage_fit = [] + station_leakages = None + + if normchisq == False: + if not quiet: print('Assigning data weights to give the correct log-likelihood...') + (alpha_d1, alpha_d2, alpha_d3, ln_norm) = compute_likelihood_constants(d1, d2, d3, sigma1, sigma2, sigma3) + else: + ln_norm = 0.0 + + # Determine the mapping between solution gains and the input visibilities + # Use passed gains even if fit_gains=False and marginalize_gains=False + # NOTE: THERE IS A PROBLEM IN THIS IMPLEMENTATION. A fixed gain prior is ignored. However, gain_init may still want to apply a constant correction, especially when passing a caltable. + # We should maybe have two gain lists: one for constant gains and one for fitted gains + mean_g1 = mean_g2 = 0.0 + if fit_gains or marginalize_gains: + if gain_prior is None: + gain_prior = default_gain_prior(Obsdata.tarr['site']) + (gain_list, gains_t1, gains_t2) = make_gain_map(Obsdata, gain_prior) + if type(gain_init) == caltable.Caltable: + if not quiet: print('Converting gain_init from caltable to a list') + gain_init = caltable_to_gains(gain_init, gain_list) + if gain_init is None: + if not quiet: print('Initializing all gain corrections to be zero') + gain_init = np.zeros(len(gain_list)) + else: + if len(gain_init) != len(gain_list): + raise Exception('Gain initialization has incorrect dimensions! %d %d' % (len(gain_init), len(gain_list))) + if fit_gains: + n_gains = len(gain_list) + elif marginalize_gains: + n_gains = 0 + else: + if gain_init is None: + n_gains = 0 + gain_list = [] + gains_t1 = gains_t2 = None + else: + if gain_prior is None: + gain_prior = default_gain_prior(Obsdata.tarr['site']) + (gain_list, gains_t1, gains_t2) = make_gain_map(Obsdata, gain_prior) + if type(gain_init) == caltable.Caltable: + if not quiet: print('Converting gain_init from caltable to a list') + gain_init = caltable_to_gains(gain_init, gain_list) + + if fit_leakage: + leakage_init = np.zeros(len(leakage_fit) * 2) + for j in range(len(leakage_init)//2): + leakage_init[2*j] = np.real(station_leakages[leakage_fit[j][0]][leakage_fit[j][1]]) + leakage_init[2*j + 1] = np.imag(station_leakages[leakage_fit[j][0]][leakage_fit[j][1]]) + else: + leakage_init = [] + + # Initial parameters + param_init = [] + for j in range(len(param_map)): + pm = param_map[j] + if param_map[j][1] in trial_model.params[param_map[j][0]].keys(): + param_init.append(transform_param(model_init.params[pm[0]][pm[1]]/pm[2], model_prior[pm[0]][pm[1]],inverse=False)) + else: # In this case, the parameter is a list of complex numbers, so the real/imaginary or abs/arg components need to be assigned + if param_map[j][1].find('cpol') != -1: + param_type = 'beta_list_cpol' + idx = int(param_map[j][1].split('_')[0][8:]) + elif param_map[j][1].find('pol') != -1: + param_type = 'beta_list_pol' + idx = int(param_map[j][1].split('_')[0][7:]) + (len(trial_model.params[param_map[j][0]][param_type])-1)//2 + elif param_map[j][1].find('beta') != -1: + param_type = 'beta_list' + idx = int(param_map[j][1].split('_')[0][4:]) - 1 + else: + raise Exception('Unsure how to interpret ' + param_map[j][1]) + + curval = model_init.params[param_map[j][0]][param_type][idx] + if param_map[j][1][-2:] == 're': + param_init.append(transform_param(np.real( model_init.params[pm[0]][param_type][idx]/pm[2]), model_prior[pm[0]][pm[1]],inverse=False)) + elif param_map[j][1][-2:] == 'im': + param_init.append(transform_param(np.imag( model_init.params[pm[0]][param_type][idx]/pm[2]), model_prior[pm[0]][pm[1]],inverse=False)) + elif param_map[j][1][-3:] == 'abs': + param_init.append(transform_param(np.abs( model_init.params[pm[0]][param_type][idx]/pm[2]), model_prior[pm[0]][pm[1]],inverse=False)) + elif param_map[j][1][-3:] == 'arg': + param_init.append(transform_param(np.angle(model_init.params[pm[0]][param_type][idx])/pm[2], model_prior[pm[0]][pm[1]],inverse=False)) + else: + if not quiet: print('Parameter ' + param_map[j][1] + ' not understood!') + n_params = len(param_init) + + # Note: model parameters can have transformations applied; gains and leakage do not + if fit_gains: # Do not add these if marginalize_gains == True + param_init += list(gain_init) + if fit_leakage: + param_init += list(leakage_init) + + if minimizer_func not in ['dynesty_static','dynesty_dynamic','pymc3']: + # Define bounds (irrelevant for dynesty or pymc3) + if use_bounds == False and minimizer_func in ['scipy.optimize.dual_annealing']: + if not quiet: print('Bounds are required for ' + minimizer_func + '! Setting use_bounds=True.') + use_bounds = True + if use_bounds == False and bounds is not None: + if not quiet: print('Bounds passed but use_bounds=False; setting use_bounds=True.') + use_bounds = True + if bounds is None and use_bounds: + if not quiet: print('No bounds passed. Setting nominal bounds.') + bounds = make_bounds(model_prior, param_map, gain_prior, gain_list, n_gains, leakage_fit, leakage_prior) + if use_bounds == False: + bounds = None + + # Gather global variables into a dictionary + globdict = {'trial_model':trial_model, + 'd1':d1, 'd2':d2, 'd3':d3, + 'data1':data1, 'sigma1':sigma1, 'uv1':uv1, 'jonesdict1':jonesdict1, + 'data2':data2, 'sigma2':sigma2, 'uv2':uv2, 'jonesdict2':jonesdict2, + 'data3':data3, 'sigma3':sigma3, 'uv3':uv3, 'jonesdict3':jonesdict3, + 'alpha_d1':alpha_d1, 'alpha_d2':alpha_d2, 'alpha_d3':alpha_d3, + 'n_params': n_params, 'n_gains':n_gains, 'n_leakage':len(leakage_init), + 'model_prior':model_prior, 'param_map':param_map, 'param_mask':param_mask, + 'gain_prior':gain_prior, 'gain_list':gain_list, 'gain_init':gain_init, + 'fit_leakage':fit_leakage, 'leakage_init':leakage_init, 'leakage_fit':leakage_fit, 'station_leakages':station_leakages, 'leakage_prior':leakage_prior, + 'show_updates':show_updates, 'update_interval':update_interval, 'gains_t1':gains_t1, 'gains_t2':gains_t2, + 'minimizer_func':minimizer_func,'Obsdata':Obsdata, + 'fit_pol':fit_pol, 'fit_cpol':fit_cpol, + 'flux':flux, 'alpha_flux':alpha_flux, 'fit_gains':fit_gains, 'marginalize_gains':marginalize_gains, 'ln_norm':ln_norm, 'param_init':param_init, 'test_gradient':test_gradient} + if fit_leakage: + update_leakage(leakage_init) + + + # Define the function that reports progress + def plotcur(params_step, *args): + global nit, globdict + if globdict['show_updates'] and (nit % globdict['update_interval'] == 0) and (quiet == False): + if globdict['n_params'] > 0: + print('Params:',params_step[:globdict['n_params']]) + print('Transformed Params:',transform_params(params_step[:globdict['n_params']], globdict['param_map'], globdict['minimizer_func'], globdict['model_prior'])) + gains = params_step[globdict['n_params']:(globdict['n_params'] + globdict['n_gains'])] + leakage = params_step[(globdict['n_params'] + globdict['n_gains']):] + if len(leakage): + print('leakage:',leakage) + update_leakage(leakage) + (chi2_1, chi2_2, chi2_3) = chisq_list(gains) + print("i: %d chi2_1: %0.2f chi2_2: %0.2f chi2_3: %0.2f prior: %0.2f" % (nit, chi2_1, chi2_2, chi2_3, prior(params_step[:globdict['n_params']], globdict['param_map'], globdict['model_prior'], globdict['minimizer_func']))) + nit += 1 + + # Print initial statistics + if not quiet: + print("Initial Objective Function: %f" % (objfunc(param_init))) + if d1 in DATATERMS: + print("Total Data 1: ", (len(data1))) + if d2 in DATATERMS: + print("Total Data 2: ", (len(data2))) + if d3 in DATATERMS: + print("Total Data 3: ", (len(data3))) + print("Total Fitted Real Parameters #: ",(len(param_init))) + print("Fitted Model Parameters: ",[_[1] for _ in param_map]) + print('Fitting Leakage Terms for:',leakage_fit) + plotcur(param_init) + + # Run the minimization + tstart = time.time() + ret = {} + if minimizer_func == 'scipy.optimize.minimize': + min_kwargs = {'method':minimizer_kwargs.get('method','L-BFGS-B'), + 'options':{'maxiter':MAXIT, 'ftol':STOP, 'maxcor':NHIST,'gtol':STOP,'maxls':MAXLS}} + + if 'options' in minimizer_kwargs.keys(): + for key in minimizer_kwargs['options'].keys(): + min_kwargs['options'][key] = minimizer_kwargs['options'][key] + + for key in minimizer_kwargs.keys(): + if key in ['options','method']: + continue + else: + min_kwargs[key] = minimizer_kwargs[key] + + res = opt.minimize(objfunc, param_init, jac=objgrad, callback=plotcur, bounds=bounds, **min_kwargs) + elif minimizer_func == 'scipy.optimize.dual_annealing': + min_kwargs = {} + min_kwargs['local_search_options'] = {'jac':objgrad, + 'method':'L-BFGS-B','options':{'maxiter':MAXIT, 'ftol':STOP, 'maxcor':NHIST,'gtol':STOP,'maxls':MAXLS}} + if 'local_search_options' in minimizer_kwargs.keys(): + for key in minimizer_kwargs['local_search_options'].keys(): + min_kwargs['local_search_options'][key] = minimizer_kwargs['local_search_options'][key] + + for key in minimizer_kwargs.keys(): + if key in ['local_search_options']: + continue + min_kwargs[key] = minimizer_kwargs[key] + + res = opt.dual_annealing(objfunc, x0=param_init, bounds=bounds, callback=plotcur, **min_kwargs) + elif minimizer_func == 'scipy.optimize.basinhopping': + min_kwargs = {} + for key in minimizer_kwargs.keys(): + min_kwargs[key] = minimizer_kwargs[key] + + res = opt.basinhopping(objfunc, param_init, **min_kwargs) + elif minimizer_func == 'pymc3': + ######################## + ## Sample using pymc3 ## + ######################## + import pymc3 as pm + import theano + import theano.tensor as tt + + # To simplfy things, we'll use cdf transforms to map everything to a hypercube, as in dynesty + + # First, define a theano Op for our likelihood function + # This is based on the example here: https://docs.pymc.io/notebooks/blackbox_external_likelihood.html + class LogLike(tt.Op): + itypes = [tt.dvector] # expects a vector of parameter values when called + otypes = [tt.dscalar] # outputs a single scalar value (the log likelihood) + + def __init__(self, objfunc, objgrad): + # add inputs as class attributes + self.objfunc = objfunc + self.objgrad = objgrad + self.logpgrad = LogLikeGrad(objfunc, objgrad) + + def prior_transform(self, u): + # This function transforms samples from the unit hypercube (u) to the target prior (x) + global globdict + cparts = ['re','im'] + model_params_u = u[:n_params] + gain_params_u = u[n_params:(n_params+n_gains)] + leakage_params_u = u[(n_params+n_gains):] + model_params_x = [cdf_inverse( model_params_u[j], globdict['model_prior'][globdict['param_map'][j][0]][globdict['param_map'][j][1]]) for j in range(len(model_params_u))] + gain_params_x = [cdf_inverse( gain_params_u[j], globdict['gain_prior'][globdict['gain_list'][j][1]]) for j in range(len(gain_params_u))] + leakage_params_x = [cdf_inverse(leakage_params_u[j], globdict['leakage_prior'][globdict['leakage_fit'][j//2][0]][leakage_fit[j//2][1]][cparts[j%2]]) for j in range(len(leakage_params_u))] + return np.concatenate([model_params_x, gain_params_x, leakage_params_x]) + + def perform(self, node, inputs, outputs): + # the method that is used when calling the Op + theta, = inputs # this will contain my variables + # Transform from the hypercube to the prior + x = self.prior_transform(theta) + + # call the log-likelihood function + logl = -self.objfunc(x) + + outputs[0][0] = np.array(logl) # output the log-likelihood + + def grad(self, inputs, g): + # the method that calculates the vector-Jacobian product + # http://deeplearning.net/software/theano_versions/dev/extending/op.html#grad + theta, = inputs + return [g[0]*self.logpgrad(theta)] + + class LogLikeGrad(tt.Op): + """ + This Op will be called with a vector of values and also return a vector of + values - the gradients in each dimension. + """ + itypes = [tt.dvector] + otypes = [tt.dvector] + + def __init__(self, objfunc, objgrad): + self.objfunc = objfunc + self.objgrad = objgrad + + def prior_transform(self, u): + # This function transforms samples from the unit hypercube (u) to the target prior (x) + global globdict + cparts = ['re','im'] + model_params_u = u[:n_params] + gain_params_u = u[n_params:(n_params+n_gains)] + leakage_params_u = u[(n_params+n_gains):] + model_params_x = [cdf_inverse( model_params_u[j], globdict['model_prior'][globdict['param_map'][j][0]][globdict['param_map'][j][1]]) for j in range(len(model_params_u))] + gain_params_x = [cdf_inverse( gain_params_u[j], globdict['gain_prior'][globdict['gain_list'][j][1]]) for j in range(len(gain_params_u))] + leakage_params_x = [cdf_inverse(leakage_params_u[j], globdict['leakage_prior'][globdict['leakage_fit'][j//2][0]][leakage_fit[j//2][1]][cparts[j%2]]) for j in range(len(leakage_params_u))] + return np.concatenate([model_params_x, gain_params_x, leakage_params_x]) + + def perform(self, node, inputs, outputs): + theta, = inputs + x = self.prior_transform(theta) + outputs[0][0] = -self.objgrad(x) + + # create the log-likelihood Op + logl = LogLike(objfunc, objgrad) + + # Define the sampler keywords + min_kwargs = {} + for key in minimizer_kwargs.keys(): + min_kwargs[key] = minimizer_kwargs[key] + + # Define the initial value if not passed + if 'start' not in min_kwargs.keys(): + cparts = ['re','im'] + model_params_x = param_init[:n_params] + gain_params_x = param_init[n_params:(n_params+n_gains)] + leakage_params_x = param_init[(n_params+n_gains):] + model_params_u = [cdf( model_params_x[j], globdict['model_prior'][globdict['param_map'][j][0]][globdict['param_map'][j][1]]) for j in range(len(model_params_x))] + gain_params_u = [cdf( gain_params_x[j], globdict['gain_prior'][globdict['gain_list'][j][1]]) for j in range(len(gain_params_x))] + leakage_params_u = [cdf(leakage_params_x[j], globdict['leakage_prior'][globdict['leakage_fit'][j//2][0]][leakage_fit[j//2][1]][cparts[j%2]]) for j in range(len(leakage_params_x))] + param_init_u = np.concatenate([model_params_u, gain_params_u, leakage_params_u]) + min_kwargs['start'] = {} + for j in range(len(param_init)): + min_kwargs['start']['var' + str(j)] = param_init_u[j] + + # Setup the sampler + with pm.Model() as model: + theta = tt.as_tensor_variable([ pm.Uniform('var' + str(j), lower=0., upper=1.) for j in range(len(param_init)) ]) + pm.DensityDist('likelihood', lambda v: logl(v), observed={'v': theta}) + trace = pm.sample(**min_kwargs) + + # Extract useful sampling diagnostics. + samples_u = np.vstack([trace['var' + str(j)] for j in range(len(param_init))]).T # samples in the hypercube + samples = np.array([logl.prior_transform(u) for u in samples_u]) # samples + mean = np.mean(samples,axis=0) + var = np.var(samples,axis=0) + + # Compute the log-posterior + if not quiet: print('Calculating the posterior values for the samples...') + logposterior = np.array([-objfunc(x, force_posterior=True) for x in samples]) + + # Select the MAP + j_MAP = np.argmax(logposterior) + MAP = samples[j_MAP] + + # Return a model determined by the MAP + set_params(MAP[:n_params], trial_model, param_map, minimizer_func, model_prior) + gains = MAP[n_params:(n_params+n_gains)] + leakage = MAP[(n_params+n_gains):] + update_leakage(leakage) + + # Return the sampler + ret['trace'] = trace + ret['mean'] = mean + ret['map'] = MAP + ret['std'] = var**0.5 + ret['samples'] = samples + ret['logposterior'] = logposterior + + # Return a set of models from the posterior + posterior_models = [] + for j in range(N_POSTERIOR_SAMPLES): + posterior_model = trial_model.copy() + set_params(samples[-j][:n_params], posterior_model, param_map, minimizer_func, model_prior) + posterior_models.append(posterior_model) + ret['posterior_models'] = posterior_models + + # Return data that has been rescaled based on 'natural' units for each parameter + import copy + samples_natural = copy.deepcopy(samples) + samples_natural[:,:n_params] /= np.array([_[4] for _ in param_map]) + ret['samples_natural'] = samples_natural + + # Return the names of the fitted parameters + labels = [] + labels_natural = [] + for _ in param_map: + labels.append(_[1].replace('_','-')) + labels_natural.append(_[1].replace('_','-')) + if _[3] != '': + labels_natural[-1] += ' (' + _[3] + ')' + for _ in gain_list: + labels.append(str(_[0]) + ' ' + _[1]) + labels_natural.append(str(_[0]) + ' ' + _[1]) + for _ in leakage_fit: + for coord in ['re','im']: + labels.append(_[0] + ',' + _[1] + ',' + coord) + labels_natural.append(_[0] + ',' + _[1] + ',' + coord) + + ret['labels'] = labels + ret['labels_natural'] = labels_natural + elif minimizer_func in ['dynesty_static','dynesty_dynamic']: + ########################## + ## Sample using dynesty ## + ########################## + import dynesty + from dynesty import utils as dyfunc + # Define the functions that dynesty requires + def prior_transform(u): + # This function transforms samples from the unit hypercube (u) to the target prior (x) + global globdict + cparts = ['re','im'] + model_params_u = u[:n_params] + gain_params_u = u[n_params:(n_params+n_gains)] + leakage_params_u = u[(n_params+n_gains):] + model_params_x = [cdf_inverse( model_params_u[j], globdict['model_prior'][globdict['param_map'][j][0]][globdict['param_map'][j][1]]) for j in range(len(model_params_u))] + gain_params_x = [cdf_inverse( gain_params_u[j], globdict['gain_prior'][globdict['gain_list'][j][1]]) for j in range(len(gain_params_u))] + leakage_params_x = [cdf_inverse(leakage_params_u[j], globdict['leakage_prior'][globdict['leakage_fit'][j//2][0]][leakage_fit[j//2][1]][cparts[j%2]]) for j in range(len(leakage_params_u))] + return np.concatenate([model_params_x, gain_params_x, leakage_params_x]) + + def loglike(x): + return -objfunc(x) + + def grad(x): + return -objgrad(x) + + # Setup a multiprocessing pool if needed + if processes >= 0: + import pathos.multiprocessing as mp + from multiprocessing import cpu_count + if processes == 0: processes = int(cpu_count()) + + # Ensure efficient memory allocation among the processes and separate trial models for each + def init(_globdict): + global globdict + globdict = _globdict + if processes >= 0: + globdict['trial_model'] = globdict['trial_model'].copy() + + return + + pool = mp.Pool(processes=processes, initializer=init, initargs=(globdict,)) + if not quiet: print('Using a pool with %d processes' % processes) + else: + pool = processes = None + + # Setup the sampler + if minimizer_func == 'dynesty_static': + sampler = dynesty.NestedSampler(loglike, prior_transform, ndim=len(param_init), gradient=grad, pool=pool, queue_size=processes, **minimizer_kwargs) + else: + sampler = dynesty.DynamicNestedSampler(loglike, prior_transform, ndim=len(param_init), gradient=grad, pool=pool, queue_size=processes, **minimizer_kwargs) + + # Run the sampler + sampler.run_nested(**run_nested_kwargs) + + # Print the sampler summary + res = sampler.results + if not quiet: + try: res.summary() + except: pass + + # Extract useful sampling diagnostics. + samples = res.samples # samples + weights = np.exp(res.logwt - res.logz[-1]) # normalized weights + mean, cov = dyfunc.mean_and_cov(samples, weights) + + # Compute the log-posterior + if not quiet: print('Calculating the posterior values for the samples...') + if pool is not None: + from functools import partial + def logpost(j): + return -objfunc(samples[j], force_posterior=True) + + logposterior = pool.map(logpost, range(len(samples))) + else: + logposterior = np.array([-objfunc(x, force_posterior=True) for x in samples]) + + # Close the pool (this may not be the desired behavior if the sampling is to be iterative!) + if pool is not None: + pool.close() + + # Select the MAP + j_MAP = np.argmax(logposterior) + MAP = samples[j_MAP] + + # Resample from the posterior + samples = dyfunc.resample_equal(samples, weights) + + # Return a model determined by the MAP + set_params(MAP[:n_params], trial_model, param_map, minimizer_func, model_prior) + gains = MAP[n_params:(n_params+n_gains)] + leakage = MAP[(n_params+n_gains):] + update_leakage(leakage) + + # Return the sampler + ret['sampler'] = sampler + ret['mean'] = mean + ret['map'] = MAP + ret['std'] = cov.diagonal()**0.5 + ret['samples'] = samples + ret['logposterior'] = logposterior + + # Return a set of models from the posterior + posterior_models = [] + for j in range(N_POSTERIOR_SAMPLES): + posterior_model = trial_model.copy() + set_params(samples[j][:n_params], posterior_model, param_map, minimizer_func, model_prior) + posterior_models.append(posterior_model) + ret['posterior_models'] = posterior_models + + # Return data that has been rescaled based on 'natural' units for each parameter + import copy + res_natural = copy.deepcopy(res) + res_natural.samples[:,:n_params] /= np.array([_[4] for _ in param_map]) + samples_natural = samples[:,:n_params]/np.array([_[4] for _ in param_map]) + ret['res_natural'] = res_natural + ret['samples_natural'] = samples_natural + + # Return the names of the fitted parameters + labels = [] + labels_natural = [] + for _ in param_map: + labels.append(_[1].replace('_','-')) + labels_natural.append(_[1].replace('_','-')) + if _[3] != '': + labels_natural[-1] += ' (' + _[3] + ')' + for _ in gain_list: + labels.append(str(_[0]) + ' ' + _[1]) + labels_natural.append(str(_[0]) + ' ' + _[1]) + for _ in leakage_fit: + for coord in ['re','im']: + labels.append(_[0] + ',' + _[1] + ',' + coord) + labels_natural.append(_[0] + ',' + _[1] + ',' + coord) + + ret['labels'] = labels + ret['labels_natural'] = labels_natural + else: + raise Exception('Minimizer function ' + minimizer_func + ' is not recognized!') + + # Format and print summary and fitted parameters + tstop = time.time() + trial_model = globdict['trial_model'] + + if not quiet: + print("\ntime: %f s" % (tstop - tstart)) + print("\nFitted Parameters:") + if minimizer_func not in ['dynesty_static','dynesty_dynamic','pymc3']: + out = res.x + set_params(out[:n_params], trial_model, param_map, minimizer_func, model_prior) + gains = out[n_params:(n_params + n_gains)] + leakage = out[(n_params + n_gains):] + update_leakage(leakage) + tparams = transform_params(out[:n_params], param_map, minimizer_func, model_prior) + if not quiet: + cur_idx = -1 + if len(param_map): + print('Model Parameters:') + for j in range(len(param_map)): + if param_map[j][0] != cur_idx: + cur_idx = param_map[j][0] + print(model_init.models[cur_idx] + ' (component %d/%d):' % (cur_idx+1,model_init.N_models())) + print(('\t' + param_map[j][1] + ': %f ' + param_map[j][3]) % (tparams[j] * param_map[j][2]/param_map[j][4])) + print('\n') + + if len(leakage_fit): + print('Leakage (%; re, im):') + for j in range(len(leakage_fit)): + print('\t' + leakage_fit[j][0] + ', ' + leakage_fit[j][1] + ': %2.2f %2.2f' % (leakage[2*j]*100,leakage[2*j + 1]*100)) + print('\n') + + print("Final Chi^2_1: %f Chi^2_2: %f Chi^2_3: %f" % chisq_list(gains)) + print("J: %f" % res.fun) + print(res.message) + else: + if not quiet: + cur_idx = -1 + if len(param_map): + print('Model Parameters (mean and std):') + for j in range(len(param_map)): + if param_map[j][0] != cur_idx: + cur_idx = param_map[j][0] + print(model_init.models[cur_idx] + ' (component %d/%d):' % (cur_idx+1,model_init.N_models())) + print(('\t' + param_map[j][1] + ': %f +/- %f ' + param_map[j][3]) % (ret['mean'][j] * param_map[j][2]/param_map[j][4], ret['std'][j] * param_map[j][2]/param_map[j][4])) + print('\n') + + if len(leakage_fit): + print('Leakage (%; re, im):') + for j in range(len(leakage_fit)): + j2 = 2*j + n_params + n_gains + print(('\t' + leakage_fit[j][0] + ', ' + leakage_fit[j][1] + + ': %2.2f +/- %2.2f, %2.2f +/- %2.2f') + % (mean[j2]*100,cov[j2,j2]**0.5 * 100,mean[j2+1]*100,cov[j2+1,j2+1]**0.5 * 100)) + print('\n') + + # Return fitted model + ret['model'] = trial_model + ret['param_map'] = param_map + ret['chisq_list'] = chisq_list(gains) + try: ret['res'] = res + except: pass + + if fit_gains: + ret['gains'] = gains + + # Create and return a caltable + caldict = {} + for site in set(np.array(gain_list)[:,1]): + caldict[site] = [] + + for j in range(len(gains)): + caldict[gain_list[j][1]].append((gain_list[j][0], (1.0 + gains[j]), (1.0 + gains[j]))) + + for site in caldict.keys(): + caldict[site] = np.array(caldict[site], dtype=DTCAL) + + caltable = ehtim.caltable.Caltable(Obsdata.ra, Obsdata.dec, Obsdata.rf, Obsdata.bw, caldict, Obsdata.tarr, + source=Obsdata.source, mjd=Obsdata.mjd, timetype=Obsdata.timetype) + ret['caltable'] = caltable + + # If relevant, return useful quantities associated with the leakage + if station_leakages is not None: + ret['station_leakages'] = station_leakages + tarr = Obsdata.tarr.copy() + for s in station_leakages.keys(): + if 'R' in station_leakages[s].keys(): tarr[Obsdata.tkey[s]]['dr'] = station_leakages[s]['R'] + if 'L' in station_leakages[s].keys(): tarr[Obsdata.tkey[s]]['dl'] = station_leakages[s]['L'] + ret['tarr'] = tarr + + return ret + +################################################################################################## +# Wrapper Functions +################################################################################################## + +def chisq(model, uv, data, sigma, dtype, jonesdict=None): + """return the chi^2 for the appropriate dtype + """ + + chisq = 1 + if not dtype in DATATERMS: + return chisq + + if dtype == 'vis': + chisq = chisq_vis(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'amp': + chisq = chisq_amp(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'logamp': + chisq = chisq_logamp(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'bs': + chisq = chisq_bs(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'cphase': + chisq = chisq_cphase(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'cphase_diag': + chisq = chisq_cphase_diag(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'camp': + chisq = chisq_camp(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'logcamp': + chisq = chisq_logcamp(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'logcamp_diag': + chisq = chisq_logcamp_diag(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'pvis': + chisq = chisq_pvis(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'm': + chisq = chisq_m(model, uv, data, sigma, jonesdict=jonesdict) + elif dtype in ['rrll','llrr','rlrr','rlll','lrrr','lrll']: + chisq = chisq_fracpol(dtype[:2],dtype[2:],model, uv, data, sigma, jonesdict=jonesdict) + elif dtype == 'polclosure': + chisq = chisq_polclosure(model, uv, data, sigma, jonesdict=jonesdict) + + return chisq + +def chisqgrad(model, uv, data, sigma, jonesdict, dtype, param_mask, fit_gains=False, gains=None, gains_t1=None, gains_t2=None, fit_pol=False, fit_cpol=False, fit_leakage=False): + """return the chi^2 gradient for the appropriate dtype + """ + global globdict + + n_chisqgrad = len(param_mask) + if fit_leakage: + n_chisqgrad += 2*len(globdict['leakage_fit']) + + chisqgrad = np.zeros(n_chisqgrad) + if fit_gains: + gaingrad = np.zeros_like(gains) + else: + gaingrad = np.array([]) + + # Now we need to be sure to put the gradient in the correct order: model parameters, then gains, then leakage + param_mask_full = np.zeros(len(chisqgrad), dtype=bool) + leakage_mask_full = np.zeros(len(chisqgrad), dtype=bool) + param_mask_full[:len(param_mask)] = param_mask + leakage_mask_full[len(param_mask):] = ~leakage_mask_full[len(param_mask):] + + if not dtype in DATATERMS: + return np.concatenate([chisqgrad[param_mask_full],gaingrad,chisqgrad[leakage_mask_full]]) + + if dtype == 'vis': + chisqgrad = chisqgrad_vis(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'amp': + chisqgrad = chisqgrad_amp(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + + if fit_gains: + i1 = model.sample_uv(uv[:,0],uv[:,1], jonesdict=jonesdict) + amp_samples = np.abs(i1) + amp = data + pp = ((amp - amp_samples) * amp_samples) / (sigma**2) + gaingrad = 2.0/(1.0 + np.array(gains)) * np.array([np.sum(pp[(np.array(gains_t1) == j) + (np.array(gains_t2) == j)]) for j in range(len(gains))])/len(data) + elif dtype == 'logamp': + chisqgrad = chisqgrad_logamp(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'bs': + chisqgrad = chisqgrad_bs(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'cphase': + chisqgrad = chisqgrad_cphase(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'cphase_diag': + chisqgrad = chisqgrad_cphase_diag(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'camp': + chisqgrad = chisqgrad_camp(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'logcamp': + chisqgrad = chisqgrad_logcamp(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'logcamp_diag': + chisqgrad = chisqgrad_logcamp_diag(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'pvis': + chisqgrad = chisqgrad_pvis(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype == 'm': + chisqgrad = chisqgrad_m(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + elif dtype in ['rrll','llrr','rlrr','rlll','lrrr','lrll']: + chisqgrad = chisqgrad_fracpol(dtype[:2],dtype[2:],model, uv, data, sigma, jonesdict=jonesdict, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage) + elif dtype == 'polclosure': + chisqgrad = chisqgrad_polclosure(model, uv, data, sigma, fit_pol=fit_pol, fit_cpol=fit_cpol, fit_leakage=fit_leakage, jonesdict=jonesdict) + + return np.concatenate([chisqgrad[param_mask_full],gaingrad,chisqgrad[leakage_mask_full]]) + +def chisqdata(Obsdata, dtype, pol='I', **kwargs): + + """Return the data, sigma, and matrices for the appropriate dtype + """ + + (data, sigma, uv, jonesdict) = (False, False, False, None) + + if dtype == 'vis': + (data, sigma, uv, jonesdict) = chisqdata_vis(Obsdata, pol=pol, **kwargs) + elif dtype == 'amp' or dtype == 'logamp': + (data, sigma, uv, jonesdict) = chisqdata_amp(Obsdata, pol=pol,**kwargs) + elif dtype == 'bs': + (data, sigma, uv, jonesdict) = chisqdata_bs(Obsdata, pol=pol,**kwargs) + elif dtype == 'cphase': + (data, sigma, uv, jonesdict) = chisqdata_cphase(Obsdata, pol=pol,**kwargs) + elif dtype == 'cphase_diag': + (data, sigma, uv, jonesdict) = chisqdata_cphase_diag(Obsdata, pol=pol,**kwargs) + elif dtype == 'camp': + (data, sigma, uv, jonesdict) = chisqdata_camp(Obsdata, pol=pol,**kwargs) + elif dtype == 'logcamp': + (data, sigma, uv, jonesdict) = chisqdata_logcamp(Obsdata, pol=pol,**kwargs) + elif dtype == 'logcamp_diag': + (data, sigma, uv, jonesdict) = chisqdata_logcamp_diag(Obsdata, pol=pol,**kwargs) + elif dtype == 'pvis': + (data, sigma, uv, jonesdict) = chisqdata_pvis(Obsdata, pol=pol,**kwargs) + elif dtype == 'm': + (data, sigma, uv, jonesdict) = chisqdata_m(Obsdata, pol=pol,**kwargs) + elif dtype in ['rrll','llrr','rlrr','rlll','lrrr','lrll']: + (data, sigma, uv, jonesdict) = chisqdata_fracpol(Obsdata,dtype[:2],dtype[2:],jonesdict=jonesdict) + elif dtype == 'polclosure': + (data, sigma, uv, jonesdict) = chisqdata_polclosure(Obsdata,jonesdict=jonesdict) + + return (data, sigma, uv, jonesdict) + + +################################################################################################## +# Chi-squared and Gradient Functions +################################################################################################## + +def chisq_vis(model, uv, vis, sigma, jonesdict=None): + """Visibility chi-squared""" + + samples = model.sample_uv(uv[:,0],uv[:,1], jonesdict=jonesdict) + return np.sum(np.abs((samples-vis)/sigma)**2)/(2*len(vis)) + +def chisqgrad_vis(model, uv, vis, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the visibility chi-squared""" + + samples = model.sample_uv(uv[:,0],uv[:,1], jonesdict=jonesdict) + wdiff = (vis - samples)/(sigma**2) + grad = model.sample_grad_uv(uv[:,0],uv[:,1],fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + + out = -np.real(np.dot(grad.conj(), wdiff))/len(vis) + return out + +def chisq_amp(model, uv, amp, sigma, jonesdict=None): + """Visibility Amplitudes (normalized) chi-squared""" + + amp_samples = np.abs(model.sample_uv(uv[:,0],uv[:,1], jonesdict=jonesdict)) + return np.sum(np.abs((amp - amp_samples)/sigma)**2)/len(amp) + +def chisqgrad_amp(model, uv, amp, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the amplitude chi-squared""" + + i1 = model.sample_uv(uv[:,0],uv[:,1], jonesdict=jonesdict) + amp_samples = np.abs(i1) + + pp = ((amp - amp_samples) * amp_samples) / (sigma**2) / i1 + grad = model.sample_grad_uv(uv[:,0],uv[:,1],fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + out = (-2.0/len(amp)) * np.real(np.dot(grad, pp)) + return out + +def chisq_bs(model, uv, bis, sigma, jonesdict=None): + """Bispectrum chi-squared""" + + bisamples = model.sample_uv(uv[0][:,0],uv[0][:,1]) * model.sample_uv(uv[1][:,0],uv[1][:,1]) * model.sample_uv(uv[2][:,0],uv[2][:,1]) + chisq= np.sum(np.abs(((bis - bisamples)/sigma))**2)/(2.*len(bis)) + return chisq + +def chisqgrad_bs(model, uv, bis, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the bispectrum chi-squared""" + + V1 = model.sample_uv(uv[0][:,0],uv[0][:,1]) + V2 = model.sample_uv(uv[1][:,0],uv[1][:,1]) + V3 = model.sample_uv(uv[2][:,0],uv[2][:,1]) + V1_grad = model.sample_grad_uv(uv[0][:,0],uv[0][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V2_grad = model.sample_grad_uv(uv[1][:,0],uv[1][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V3_grad = model.sample_grad_uv(uv[2][:,0],uv[2][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + bisamples = V1 * V2 * V3 + wdiff = ((bis - bisamples).conj())/(sigma**2) + pt1 = wdiff * V2 * V3 + pt2 = wdiff * V1 * V3 + pt3 = wdiff * V1 * V2 + out = -np.real(np.dot(pt1, V1_grad.T) + np.dot(pt2, V2_grad.T) + np.dot(pt3, V3_grad.T))/len(bis) + return out + +def chisq_cphase(model, uv, clphase, sigma, jonesdict=None): + """Closure Phases (normalized) chi-squared""" + clphase = clphase * DEGREE + sigma = sigma * DEGREE + + V1 = model.sample_uv(uv[0][:,0],uv[0][:,1]) + V2 = model.sample_uv(uv[1][:,0],uv[1][:,1]) + V3 = model.sample_uv(uv[2][:,0],uv[2][:,1]) + + clphase_samples = np.angle(V1 * V2 * V3) + chisq= (2.0/len(clphase)) * np.sum((1.0 - np.cos(clphase-clphase_samples))/(sigma**2)) + return chisq + +def chisqgrad_cphase(model, uv, clphase, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the closure phase chi-squared""" + clphase = clphase * DEGREE + sigma = sigma * DEGREE + + V1 = model.sample_uv(uv[0][:,0],uv[0][:,1]) + V2 = model.sample_uv(uv[1][:,0],uv[1][:,1]) + V3 = model.sample_uv(uv[2][:,0],uv[2][:,1]) + V1_grad = model.sample_grad_uv(uv[0][:,0],uv[0][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V2_grad = model.sample_grad_uv(uv[1][:,0],uv[1][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V3_grad = model.sample_grad_uv(uv[2][:,0],uv[2][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + + clphase_samples = np.angle(V1 * V2 * V3) + + pref = np.sin(clphase - clphase_samples)/(sigma**2) + pt1 = pref/V1 + pt2 = pref/V2 + pt3 = pref/V3 + out = -(2.0/len(clphase)) * np.imag(np.dot(pt1, V1_grad.T) + np.dot(pt2, V2_grad.T) + np.dot(pt3, V3_grad.T)) + return out + +def chisq_cphase_diag(model, uv, clphase_diag, sigma, jonesdict=None): + """Diagonalized closure phases (normalized) chi-squared""" + clphase_diag = np.concatenate(clphase_diag) * DEGREE + sigma = np.concatenate(sigma) * DEGREE + + uv_diag = uv[0] + tform_mats = uv[1] + + clphase_diag_samples = [] + for iA, uv3 in enumerate(uv_diag): + i1 = model.sample_uv(uv3[0][:,0],uv3[0][:,1]) + i2 = model.sample_uv(uv3[1][:,0],uv3[1][:,1]) + i3 = model.sample_uv(uv3[2][:,0],uv3[2][:,1]) + + clphase_samples = np.angle(i1 * i2 * i3) + clphase_diag_samples.append(np.dot(tform_mats[iA],clphase_samples)) + clphase_diag_samples = np.concatenate(clphase_diag_samples) + + chisq = (2.0/len(clphase_diag)) * np.sum((1.0 - np.cos(clphase_diag-clphase_diag_samples))/(sigma**2)) + return chisq + +def chisqgrad_cphase_diag(model, uv, clphase_diag, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the diagonalized closure phase chi-squared""" + clphase_diag = clphase_diag * DEGREE + sigma = sigma * DEGREE + + uv_diag = uv[0] + tform_mats = uv[1] + + deriv = np.zeros(len(model.sample_grad_uv(0,0,fit_pol=fit_pol,fit_cpol=fit_cpol))) + for iA, uv3 in enumerate(uv_diag): + + i1 = model.sample_uv(uv3[0][:,0],uv3[0][:,1]) + i2 = model.sample_uv(uv3[1][:,0],uv3[1][:,1]) + i3 = model.sample_uv(uv3[2][:,0],uv3[2][:,1]) + + i1_grad = model.sample_grad_uv(uv3[0][:,0],uv3[0][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + i2_grad = model.sample_grad_uv(uv3[1][:,0],uv3[1][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + i3_grad = model.sample_grad_uv(uv3[2][:,0],uv3[2][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + + clphase_samples = np.angle(i1 * i2 * i3) + clphase_diag_samples = np.dot(tform_mats[iA],clphase_samples) + + clphase_diag_measured = clphase_diag[iA] + clphase_diag_sigma = sigma[iA] + + term1 = np.dot(np.dot((np.sin(clphase_diag_measured-clphase_diag_samples)/(clphase_diag_sigma**2.0)),(tform_mats[iA]/i1)),i1_grad.T) + term2 = np.dot(np.dot((np.sin(clphase_diag_measured-clphase_diag_samples)/(clphase_diag_sigma**2.0)),(tform_mats[iA]/i2)),i2_grad.T) + term3 = np.dot(np.dot((np.sin(clphase_diag_measured-clphase_diag_samples)/(clphase_diag_sigma**2.0)),(tform_mats[iA]/i3)),i3_grad.T) + deriv += -2.0*np.imag(term1 + term2 + term3) + + deriv *= 1.0/np.float(len(np.concatenate(clphase_diag))) + + return deriv + +def chisq_camp(model, uv, clamp, sigma, jonesdict=None): + """Closure Amplitudes (normalized) chi-squared""" + + V1 = model.sample_uv(uv[0][:,0],uv[0][:,1]) + V2 = model.sample_uv(uv[1][:,0],uv[1][:,1]) + V3 = model.sample_uv(uv[2][:,0],uv[2][:,1]) + V4 = model.sample_uv(uv[3][:,0],uv[3][:,1]) + + clamp_samples = np.abs(V1 * V2 / (V3 * V4)) + chisq = np.sum(np.abs((clamp - clamp_samples)/sigma)**2)/len(clamp) + return chisq + +def chisqgrad_camp(model, uv, clamp, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the closure amplitude chi-squared""" + + V1 = model.sample_uv(uv[0][:,0],uv[0][:,1]) + V2 = model.sample_uv(uv[1][:,0],uv[1][:,1]) + V3 = model.sample_uv(uv[2][:,0],uv[2][:,1]) + V4 = model.sample_uv(uv[3][:,0],uv[3][:,1]) + V1_grad = model.sample_grad_uv(uv[0][:,0],uv[0][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V2_grad = model.sample_grad_uv(uv[1][:,0],uv[1][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V3_grad = model.sample_grad_uv(uv[2][:,0],uv[2][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V4_grad = model.sample_grad_uv(uv[3][:,0],uv[3][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + + clamp_samples = np.abs((V1 * V2)/(V3 * V4)) + + pp = ((clamp - clamp_samples) * clamp_samples)/(sigma**2) + pt1 = pp/V1 + pt2 = pp/V2 + pt3 = -pp/V3 + pt4 = -pp/V4 + out = (-2.0/len(clamp)) * np.real(np.dot(pt1, V1_grad.T) + np.dot(pt2, V2_grad.T) + np.dot(pt3, V3_grad.T) + np.dot(pt4, V4_grad.T)) + return out + +def chisq_logcamp(model, uv, log_clamp, sigma, jonesdict=None): + """Log Closure Amplitudes (normalized) chi-squared""" + + a1 = np.abs(model.sample_uv(uv[0][:,0],uv[0][:,1])) + a2 = np.abs(model.sample_uv(uv[1][:,0],uv[1][:,1])) + a3 = np.abs(model.sample_uv(uv[2][:,0],uv[2][:,1])) + a4 = np.abs(model.sample_uv(uv[3][:,0],uv[3][:,1])) + + samples = np.log(a1) + np.log(a2) - np.log(a3) - np.log(a4) + chisq = np.sum(np.abs((log_clamp - samples)/sigma)**2) / (len(log_clamp)) + return chisq + +def chisqgrad_logcamp(model, uv, log_clamp, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the Log closure amplitude chi-squared""" + + V1 = model.sample_uv(uv[0][:,0],uv[0][:,1]) + V2 = model.sample_uv(uv[1][:,0],uv[1][:,1]) + V3 = model.sample_uv(uv[2][:,0],uv[2][:,1]) + V4 = model.sample_uv(uv[3][:,0],uv[3][:,1]) + V1_grad = model.sample_grad_uv(uv[0][:,0],uv[0][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V2_grad = model.sample_grad_uv(uv[1][:,0],uv[1][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V3_grad = model.sample_grad_uv(uv[2][:,0],uv[2][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + V4_grad = model.sample_grad_uv(uv[3][:,0],uv[3][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + + log_clamp_samples = np.log(np.abs(V1)) + np.log(np.abs(V2)) - np.log(np.abs(V3)) - np.log(np.abs(V4)) + + pp = (log_clamp - log_clamp_samples) / (sigma**2) + pt1 = pp / V1 + pt2 = pp / V2 + pt3 = -pp / V3 + pt4 = -pp / V4 + out = (-2.0/len(log_clamp)) * np.real(np.dot(pt1, V1_grad.T) + np.dot(pt2, V2_grad.T) + np.dot(pt3, V3_grad.T) + np.dot(pt4, V4_grad.T)) + return out + +def chisq_logcamp_diag(model, uv, log_clamp_diag, sigma, jonesdict=None): + """Diagonalized log closure amplitudes (normalized) chi-squared""" + + log_clamp_diag = np.concatenate(log_clamp_diag) + sigma = np.concatenate(sigma) + + uv_diag = uv[0] + tform_mats = uv[1] + + log_clamp_diag_samples = [] + for iA, uv4 in enumerate(uv_diag): + + a1 = np.abs(model.sample_uv(uv4[0][:,0],uv4[0][:,1])) + a2 = np.abs(model.sample_uv(uv4[1][:,0],uv4[1][:,1])) + a3 = np.abs(model.sample_uv(uv4[2][:,0],uv4[2][:,1])) + a4 = np.abs(model.sample_uv(uv4[3][:,0],uv4[3][:,1])) + + log_clamp_samples = np.log(a1) + np.log(a2) - np.log(a3) - np.log(a4) + log_clamp_diag_samples.append(np.dot(tform_mats[iA],log_clamp_samples)) + + log_clamp_diag_samples = np.concatenate(log_clamp_diag_samples) + + chisq = np.sum(np.abs((log_clamp_diag - log_clamp_diag_samples)/sigma)**2) / (len(log_clamp_diag)) + return chisq + +def chisqgrad_logcamp_diag(model, uv, log_clamp_diag, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the diagonalized log closure amplitude chi-squared""" + + uv_diag = uv[0] + tform_mats = uv[1] + + deriv = np.zeros(len(model.sample_grad_uv(0,0,fit_pol=fit_pol,fit_cpol=fit_cpol))) + for iA, uv4 in enumerate(uv_diag): + + i1 = model.sample_uv(uv4[0][:,0],uv4[0][:,1]) + i2 = model.sample_uv(uv4[1][:,0],uv4[1][:,1]) + i3 = model.sample_uv(uv4[2][:,0],uv4[2][:,1]) + i4 = model.sample_uv(uv4[3][:,0],uv4[3][:,1]) + + i1_grad = model.sample_grad_uv(uv4[0][:,0],uv4[0][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + i2_grad = model.sample_grad_uv(uv4[1][:,0],uv4[1][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + i3_grad = model.sample_grad_uv(uv4[2][:,0],uv4[2][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + i4_grad = model.sample_grad_uv(uv4[3][:,0],uv4[3][:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + + log_clamp_samples = np.log(np.abs(i1)) + np.log(np.abs(i2)) - np.log(np.abs(i3)) - np.log(np.abs(i4)) + log_clamp_diag_samples = np.dot(tform_mats[iA],log_clamp_samples) + + log_clamp_diag_measured = log_clamp_diag[iA] + log_clamp_diag_sigma = sigma[iA] + + term1 = np.dot(np.dot(((log_clamp_diag_measured-log_clamp_diag_samples)/(log_clamp_diag_sigma**2.0)),(tform_mats[iA]/i1)),i1_grad.T) + term2 = np.dot(np.dot(((log_clamp_diag_measured-log_clamp_diag_samples)/(log_clamp_diag_sigma**2.0)),(tform_mats[iA]/i2)),i2_grad.T) + term3 = np.dot(np.dot(((log_clamp_diag_measured-log_clamp_diag_samples)/(log_clamp_diag_sigma**2.0)),(tform_mats[iA]/i3)),i3_grad.T) + term4 = np.dot(np.dot(((log_clamp_diag_measured-log_clamp_diag_samples)/(log_clamp_diag_sigma**2.0)),(tform_mats[iA]/i4)),i4_grad.T) + deriv += -2.0*np.real(term1 + term2 - term3 - term4) + + deriv *= 1.0/np.float(len(np.concatenate(log_clamp_diag))) + + return deriv + +def chisq_logamp(model, uv, amp, sigma, jonesdict=None): + """Log Visibility Amplitudes (normalized) chi-squared""" + + # to lowest order the variance on the logarithm of a quantity x is + # sigma^2_log(x) = sigma^2/x^2 + logsigma = sigma / amp + + amp_samples = np.abs(model.sample_uv(uv[:,0],uv[:,1])) + return np.sum(np.abs((np.log(amp) - np.log(amp_samples))/logsigma)**2)/len(amp) + +def chisqgrad_logamp(model, uv, amp, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the Log amplitude chi-squared""" + + # to lowest order the variance on the logarithm of a quantity x is + # sigma^2_log(x) = sigma^2/x^2 + logsigma = sigma / amp + + i1 = model.sample_uv(uv[:,0],uv[:,1]) + amp_samples = np.abs(i1) + + V_grad = model.sample_grad_uv(uv[:,0],uv[:,1],fit_pol=fit_pol,fit_cpol=fit_cpol) + + pp = ((np.log(amp) - np.log(amp_samples))) / (logsigma**2) / i1 + out = (-2.0/len(amp)) * np.real(np.dot(pp, V_grad.T)) + return out + + +def chisq_pvis(model, uv, pvis, psigma, jonesdict=None): + """Polarimetric visibility chi-squared + """ + + psamples = model.sample_uv(uv[:,0],uv[:,1],pol='P',jonesdict=jonesdict) + return np.sum(np.abs((psamples-pvis)/psigma)**2)/(2*len(pvis)) + +def chisqgrad_pvis(model, uv, pvis, psigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """Polarimetric visibility chi-squared gradient + """ + samples = model.sample_uv(uv[:,0],uv[:,1],pol='P',jonesdict=jonesdict) + wdiff = (pvis - samples)/(psigma**2) + grad = model.sample_grad_uv(uv[:,0],uv[:,1],pol='P',fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + + out = -np.real(np.dot(grad.conj(), wdiff))/len(pvis) + return out + +def chisq_m(model, uv, m, msigma, jonesdict=None): + """Polarimetric ratio chi-squared + """ + + msamples = model.sample_uv(uv[:,0],uv[:,1],pol='P',jonesdict=jonesdict)/model.sample_uv(uv[:,0],uv[:,1],pol='I',jonesdict=jonesdict) + + return np.sum(np.abs((m - msamples))**2/(msigma**2)) / (2*len(m)) + +def chisqgrad_m(model, uv, mvis, msigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the polarimetric ratio chisq + """ + + samp_P = model.sample_uv(uv[:,0],uv[:,1],pol='P',jonesdict=jonesdict) + samp_I = model.sample_uv(uv[:,0],uv[:,1],pol='I',jonesdict=jonesdict) + grad_P = model.sample_grad_uv(uv[:,0],uv[:,1],pol='P',fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + grad_I = model.sample_grad_uv(uv[:,0],uv[:,1],pol='I',fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + + msamples = samp_P/samp_I + wdiff = (mvis - msamples)/(msigma**2) + # Get the gradient from the quotient rule + grad = ( grad_P * samp_I - grad_I * samp_P)/samp_I**2 + + return -np.real(np.dot(grad.conj(), wdiff))/len(mvis) + +def chisq_fracpol(upper, lower, model, uv, m, msigma, jonesdict=None): + """Polarimetric ratio chi-squared + """ + + msamples = model.sample_uv(uv[:,0],uv[:,1],pol=upper.upper(),jonesdict=jonesdict)/model.sample_uv(uv[:,0],uv[:,1],pol=lower.upper(),jonesdict=jonesdict) + + return np.sum(np.abs((m - msamples))**2/(msigma**2)) / (2*len(m)) + +def chisqgrad_fracpol(upper, lower, model, uv, mvis, msigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the polarimetric ratio chisq + """ + + samp_upper = model.sample_uv(uv[:,0],uv[:,1],pol=upper.upper(),jonesdict=jonesdict) + samp_lower = model.sample_uv(uv[:,0],uv[:,1],pol=lower.upper(),jonesdict=jonesdict) + grad_upper = model.sample_grad_uv(uv[:,0],uv[:,1],pol=upper.upper(),fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + grad_lower = model.sample_grad_uv(uv[:,0],uv[:,1],pol=lower.upper(),fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + + msamples = samp_upper/samp_lower + wdiff = (mvis - msamples)/(msigma**2) + # Get the gradient from the quotient rule + grad = ( grad_upper * samp_lower - grad_lower * samp_upper)/samp_lower**2 + + return -np.real(np.dot(grad.conj(), wdiff))/len(mvis) + +def chisq_polclosure(model, uv, vis, sigma, jonesdict=None): + """Polarimetric ratio chi-squared + """ + + RL = model.sample_uv(uv[:,0],uv[:,1],pol='RL',jonesdict=jonesdict) + LR = model.sample_uv(uv[:,0],uv[:,1],pol='LR',jonesdict=jonesdict) + RR = model.sample_uv(uv[:,0],uv[:,1],pol='RR',jonesdict=jonesdict) + LL = model.sample_uv(uv[:,0],uv[:,1],pol='LL',jonesdict=jonesdict) + samples = (RL * LR)/(RR * LL) + + return np.sum(np.abs((vis - samples))**2/(sigma**2)) / (2*len(vis)) + +def chisqgrad_polclosure(model, uv, vis, sigma, fit_pol=False, fit_cpol=False, fit_leakage=False, jonesdict=None): + """The gradient of the polarimetric ratio chisq + """ + + RL = model.sample_uv(uv[:,0],uv[:,1],pol='RL',jonesdict=jonesdict) + LR = model.sample_uv(uv[:,0],uv[:,1],pol='LR',jonesdict=jonesdict) + RR = model.sample_uv(uv[:,0],uv[:,1],pol='RR',jonesdict=jonesdict) + LL = model.sample_uv(uv[:,0],uv[:,1],pol='LL',jonesdict=jonesdict) + + dRL = model.sample_grad_uv(uv[:,0],uv[:,1],pol='RL',fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + dLR = model.sample_grad_uv(uv[:,0],uv[:,1],pol='LR',fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + dRR = model.sample_grad_uv(uv[:,0],uv[:,1],pol='RR',fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + dLL = model.sample_grad_uv(uv[:,0],uv[:,1],pol='LL',fit_pol=fit_pol,fit_cpol=fit_cpol,fit_leakage=fit_leakage,jonesdict=jonesdict) + + samples = (RL * LR)/(RR * LL) + wdiff = (vis - samples)/(sigma**2) + + # Get the gradient from the quotient rule + samp_upper = RL * LR + samp_lower = RR * LL + grad_upper = RL * dLR + dRL * LR + grad_lower = RR * dLL + dRR * LL + grad = ( grad_upper * samp_lower - grad_lower * samp_upper)/samp_lower**2 + + return -np.real(np.dot(grad.conj(), wdiff))/len(vis) + + +################################################################################################## +# Chi^2 Data functions +################################################################################################## +def apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol): + """apply systematic noise to VISIBILITIES or AMPLITUDES + data_arr should have fields 't1','t2','u','v','vis','amp','sigma' + + returns: (uv, vis, amp, sigma) + """ + + vtype=vis_poldict[pol] + atype=amp_poldict[pol] + etype=sig_poldict[pol] + + t1 = data_arr['t1'] + t2 = data_arr['t2'] + + sigma = data_arr[etype] + amp = data_arr[atype] + try: + vis = data_arr[vtype] + except ValueError: + vis = amp.astype('c16') + + snrmask = np.abs(amp/sigma) >= snrcut + + if type(systematic_noise) is dict: + sys_level = np.zeros(len(t1)) + for i in range(len(t1)): + if t1[i] in systematic_noise.keys(): + t1sys = systematic_noise[t1[i]] + else: + t1sys = 0. + if t2[i] in systematic_noise.keys(): + t2sys = systematic_noise[t2[i]] + else: + t2sys = 0. + + if t1sys<0 or t2sys<0: + sys_level[i] = -1 + else: + sys_level[i] = np.sqrt(t1sys**2 + t2sys**2) + else: + sys_level = np.sqrt(2)*systematic_noise*np.ones(len(t1)) + + mask = sys_level>=0. + mask = snrmask * mask + + sigma = np.linalg.norm([sigma, sys_level*np.abs(amp)], axis=0)[mask] + vis = vis[mask] + amp = amp[mask] + uv = np.hstack((data_arr['u'].reshape(-1,1), data_arr['v'].reshape(-1,1)))[mask] + return (uv, vis, amp, sigma) + +def make_jonesdict(Obsdata, data_arr): + # Make a dictionary with entries needed to form the Jones matrices + # Currently, this only works for data types on a single baseline (e.g., closure quantities aren't supported yet) + + # Get the names of each station for every measurement + t1 = data_arr['t1'] + t2 = data_arr['t2'] + + # Get the elevation of each station + el1 = data_arr['el1']*np.pi/180. + el2 = data_arr['el2']*np.pi/180. + + # Get the parallactic angle of each station + par1 = data_arr['par_ang1']*np.pi/180. + par2 = data_arr['par_ang2']*np.pi/180. + + # Compute the full field rotation angle for each site, based information in the Obsdata Array + fr_elev1 = np.array([Obsdata.tarr[Obsdata.tkey[o['t1']]]['fr_elev'] for o in data_arr]) + fr_elev2 = np.array([Obsdata.tarr[Obsdata.tkey[o['t2']]]['fr_elev'] for o in data_arr]) + fr_par1 = np.array([Obsdata.tarr[Obsdata.tkey[o['t1']]]['fr_par'] for o in data_arr]) + fr_par2 = np.array([Obsdata.tarr[Obsdata.tkey[o['t2']]]['fr_par'] for o in data_arr]) + fr_off1 = np.array([Obsdata.tarr[Obsdata.tkey[o['t1']]]['fr_off'] for o in data_arr]) + fr_off2 = np.array([Obsdata.tarr[Obsdata.tkey[o['t2']]]['fr_off'] for o in data_arr]) + fr1 = fr_elev1*el1 + fr_par1*par1 + fr_off1*np.pi/180. + fr2 = fr_elev2*el2 + fr_par2*par2 + fr_off2*np.pi/180. + + # Now populate the left and right D-term entries based on the Obsdata Array + DR1 = np.array([Obsdata.tarr[Obsdata.tkey[o['t1']]]['dr'] for o in data_arr]) + DL1 = np.array([Obsdata.tarr[Obsdata.tkey[o['t1']]]['dl'] for o in data_arr]) + DR2 = np.array([Obsdata.tarr[Obsdata.tkey[o['t2']]]['dr'] for o in data_arr]) + DL2 = np.array([Obsdata.tarr[Obsdata.tkey[o['t2']]]['dl'] for o in data_arr]) + + return {'fr1':fr1,'fr2':fr2,'t1':t1,'t2':t2, + 'DR1':DR1, 'DR2':DR2, 'DL1':DL1, 'DL2':DL2} + +def chisqdata_vis(Obsdata, pol='I', **kwargs): + """Return the data, sigmas, and fourier matrix for visibilities + """ + + # unpack keyword args + systematic_noise = kwargs.get('systematic_noise',0.) + snrcut = kwargs.get('snrcut',0.) + debias = kwargs.get('debias',True) + weighting = kwargs.get('weighting','natural') + + # unpack data + vtype=vis_poldict[pol] + atype=amp_poldict[pol] + etype=sig_poldict[pol] + data_arr = Obsdata.unpack(['t1','t2','u','v',vtype,atype,etype,'el1','el2','par_ang1','par_ang2'], debias=debias) + (uv, vis, amp, sigma) = apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol) + + jonesdict = make_jonesdict(Obsdata, data_arr) + + return (vis, sigma, uv, jonesdict) + +def chisqdata_amp(Obsdata, pol='I',**kwargs): + """Return the data, sigmas, and fourier matrix for visibility amplitudes + """ + + # unpack keyword args + systematic_noise = kwargs.get('systematic_noise',0.) + snrcut = kwargs.get('snrcut',0.) + debias = kwargs.get('debias',True) + weighting = kwargs.get('weighting','natural') + + # unpack data + vtype=vis_poldict[pol] + atype=amp_poldict[pol] + etype=sig_poldict[pol] + if (Obsdata.amp is None) or (len(Obsdata.amp)==0) or pol!='I': + data_arr = Obsdata.unpack(['time','t1','t2','u','v',vtype,atype,etype,'el1','el2','par_ang1','par_ang2'], debias=debias) + + else: # TODO -- pre-computed with not stokes I? + print("Using pre-computed amplitude table in amplitude chi^2!") + if not type(Obsdata.amp) in [np.ndarray, np.recarray]: + raise Exception("pre-computed amplitude table is not a numpy rec array!") + data_arr = Obsdata.amp + + + # apply systematic noise and SNR cut + # TODO -- after pre-computed?? + (uv, vis, amp, sigma) = apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol) + + # data weighting + if weighting=='uniform': + sigma = np.median(sigma) * np.ones(len(sigma)) + + jonesdict = make_jonesdict(Obsdata, data_arr) + + return (amp, sigma, uv, jonesdict) + +def chisqdata_bs(Obsdata, pol='I',**kwargs): + """return the data, sigmas, and fourier matrices for bispectra + """ + + # unpack keyword args + #systematic_noise = kwargs.get('systematic_noise',0.) #this will break with a systematic noise dict + maxset = kwargs.get('maxset',False) + if maxset: count='max' + else: count='min' + + snrcut = kwargs.get('snrcut',0.) + debias = kwargs.get('debias',True) + weighting = kwargs.get('weighting','natural') + + # unpack data + vtype=vis_poldict[pol] + if (Obsdata.bispec is None) or (len(Obsdata.bispec)==0) or pol!='I': + biarr = Obsdata.bispectra(mode="all", vtype=vtype, count=count,snrcut=snrcut) + + else: # TODO -- pre-computed with not stokes I? + print("Using pre-computed bispectrum table in cphase chi^2!") + if not type(Obsdata.bispec) in [np.ndarray, np.recarray]: + raise Exception("pre-computed bispectrum table is not a numpy rec array!") + biarr = Obsdata.bispec + # reduce to a minimal set + if count!='max': + biarr = reduce_tri_minimal(Obsdata, biarr) + + uv1 = np.hstack((biarr['u1'].reshape(-1,1), biarr['v1'].reshape(-1,1))) + uv2 = np.hstack((biarr['u2'].reshape(-1,1), biarr['v2'].reshape(-1,1))) + uv3 = np.hstack((biarr['u3'].reshape(-1,1), biarr['v3'].reshape(-1,1))) + bi = biarr['bispec'] + sigma = biarr['sigmab'] + + #add systematic noise + #sigma = np.linalg.norm([biarr['sigmab'], systematic_noise*np.abs(biarr['bispec'])], axis=0) + + # data weighting + if weighting=='uniform': + sigma = np.median(sigma) * np.ones(len(sigma)) + + return (bi, sigma, (uv1, uv2, uv3), None) + +def chisqdata_cphase(Obsdata, pol='I',**kwargs): + """Return the data, sigmas, and fourier matrices for closure phases + """ + + # unpack keyword args + maxset = kwargs.get('maxset',False) + uv_min = kwargs.get('cp_uv_min', False) + if maxset: count='max' + else: count='min' + + snrcut = kwargs.get('snrcut',0.) + systematic_cphase_noise = kwargs.get('systematic_cphase_noise',0.) + weighting = kwargs.get('weighting','natural') + + # unpack data + vtype=vis_poldict[pol] + if (Obsdata.cphase is None) or (len(Obsdata.cphase)==0) or pol!='I': + clphasearr = Obsdata.c_phases(mode="all", vtype=vtype, count=count, uv_min=uv_min, snrcut=snrcut) + else: #TODO precomputed with not Stokes I + print("Using pre-computed cphase table in cphase chi^2!") + if not type(Obsdata.cphase) in [np.ndarray, np.recarray]: + raise Exception("pre-computed closure phase table is not a numpy rec array!") + clphasearr = Obsdata.cphase + # reduce to a minimal set + if count!='max': + clphasearr = reduce_tri_minimal(Obsdata, clphasearr) + + uv1 = np.hstack((clphasearr['u1'].reshape(-1,1), clphasearr['v1'].reshape(-1,1))) + uv2 = np.hstack((clphasearr['u2'].reshape(-1,1), clphasearr['v2'].reshape(-1,1))) + uv3 = np.hstack((clphasearr['u3'].reshape(-1,1), clphasearr['v3'].reshape(-1,1))) + clphase = clphasearr['cphase'] + sigma = clphasearr['sigmacp'] + + #add systematic cphase noise (in DEGREES) + sigma = np.linalg.norm([sigma, systematic_cphase_noise*np.ones(len(sigma))], axis=0) + + # data weighting + if weighting=='uniform': + sigma = np.median(sigma) * np.ones(len(sigma)) + + return (clphase, sigma, (uv1, uv2, uv3), None) + +def chisqdata_cphase_diag(Obsdata, pol='I',**kwargs): + """Return the data, sigmas, and fourier matrices for diagonalized closure phases + """ + + # unpack keyword args + maxset = kwargs.get('maxset',False) + uv_min = kwargs.get('cp_uv_min', False) + if maxset: count='max' + else: count='min' + + snrcut = kwargs.get('snrcut',0.) + + # unpack data + vtype=vis_poldict[pol] + clphasearr = Obsdata.c_phases_diag(vtype=vtype,count=count,snrcut=snrcut,uv_min=uv_min) + + # loop over timestamps + clphase_diag = [] + sigma_diag = [] + uv_diag = [] + tform_mats = [] + for ic, cl in enumerate(clphasearr): + + # get diagonalized closure phases and errors + clphase_diag.append(cl[0]['cphase']) + sigma_diag.append(cl[0]['sigmacp']) + + # get uv arrays + u1 = cl[2][:,0].astype('float') + v1 = cl[3][:,0].astype('float') + uv1 = np.hstack((u1.reshape(-1,1), v1.reshape(-1,1))) + + u2 = cl[2][:,1].astype('float') + v2 = cl[3][:,1].astype('float') + uv2 = np.hstack((u2.reshape(-1,1), v2.reshape(-1,1))) + + u3 = cl[2][:,2].astype('float') + v3 = cl[3][:,2].astype('float') + uv3 = np.hstack((u3.reshape(-1,1), v3.reshape(-1,1))) + + # compute Fourier matrices + uv = (uv1, + uv2, + uv3 + ) + uv_diag.append(uv) + + # get transformation matrix for this timestamp + tform_mats.append(cl[4].astype('float')) + + # combine Fourier and transformation matrices into tuple for outputting + uvmatrices = (np.array(uv_diag),np.array(tform_mats)) + + return (np.array(clphase_diag), np.array(sigma_diag), uvmatrices, None) + +def chisqdata_camp(Obsdata, pol='I',**kwargs): + """Return the data, sigmas, and fourier matrices for closure amplitudes + """ + # unpack keyword args + maxset = kwargs.get('maxset',False) + if maxset: count='max' + else: count='min' + + snrcut = kwargs.get('snrcut',0.) + debias = kwargs.get('debias',True) + weighting = kwargs.get('weighting','natural') + + # unpack data & mask low snr points + vtype=vis_poldict[pol] + if (Obsdata.camp is None) or (len(Obsdata.camp)==0) or pol!='I': + clamparr = Obsdata.c_amplitudes(mode='all', count=count, ctype='camp', debias=debias, snrcut=snrcut) + else: # TODO -- pre-computed with not stokes I? + print("Using pre-computed closure amplitude table in closure amplitude chi^2!") + if not type(Obsdata.camp) in [np.ndarray, np.recarray]: + raise Exception("pre-computed closure amplitude table is not a numpy rec array!") + clamparr = Obsdata.camp + # reduce to a minimal set + if count!='max': + clamparr = reduce_quad_minimal(Obsdata, clamparr, ctype='camp') + + uv1 = np.hstack((clamparr['u1'].reshape(-1,1), clamparr['v1'].reshape(-1,1))) + uv2 = np.hstack((clamparr['u2'].reshape(-1,1), clamparr['v2'].reshape(-1,1))) + uv3 = np.hstack((clamparr['u3'].reshape(-1,1), clamparr['v3'].reshape(-1,1))) + uv4 = np.hstack((clamparr['u4'].reshape(-1,1), clamparr['v4'].reshape(-1,1))) + clamp = clamparr['camp'] + sigma = clamparr['sigmaca'] + + # data weighting + if weighting=='uniform': + sigma = np.median(sigma) * np.ones(len(sigma)) + + return (clamp, sigma, (uv1, uv2, uv3, uv4), None) + +def chisqdata_logcamp(Obsdata, pol='I', **kwargs): + """Return the data, sigmas, and fourier matrices for log closure amplitudes + """ + # unpack keyword args + maxset = kwargs.get('maxset',False) + if maxset: count='max' + else: count='min' + + snrcut = kwargs.get('snrcut',0.) + debias = kwargs.get('debias',True) + weighting = kwargs.get('weighting','natural') + + # unpack data & mask low snr points + vtype=vis_poldict[pol] + if (Obsdata.logcamp is None) or (len(Obsdata.logcamp)==0) or pol!='I': + clamparr = Obsdata.c_amplitudes(mode='all', count=count, vtype=vtype, ctype='logcamp', debias=debias, snrcut=snrcut) + else: # TODO -- pre-computed with not stokes I? + print("Using pre-computed log closure amplitude table in log closure amplitude chi^2!") + if not type(Obsdata.logcamp) in [np.ndarray, np.recarray]: + raise Exception("pre-computed log closure amplitude table is not a numpy rec array!") + clamparr = Obsdata.logcamp + # reduce to a minimal set + if count!='max': + clamparr = reduce_quad_minimal(Obsdata, clamparr, ctype='logcamp') + + uv1 = np.hstack((clamparr['u1'].reshape(-1,1), clamparr['v1'].reshape(-1,1))) + uv2 = np.hstack((clamparr['u2'].reshape(-1,1), clamparr['v2'].reshape(-1,1))) + uv3 = np.hstack((clamparr['u3'].reshape(-1,1), clamparr['v3'].reshape(-1,1))) + uv4 = np.hstack((clamparr['u4'].reshape(-1,1), clamparr['v4'].reshape(-1,1))) + clamp = clamparr['camp'] + sigma = clamparr['sigmaca'] + + # data weighting + if weighting=='uniform': + sigma = np.median(sigma) * np.ones(len(sigma)) + + return (clamp, sigma, (uv1, uv2, uv3, uv4), None) + +def chisqdata_logcamp_diag(Obsdata, pol='I', **kwargs): + """Return the data, sigmas, and fourier matrices for diagonalized log closure amplitudes + """ + # unpack keyword args + maxset = kwargs.get('maxset',False) + if maxset: count='max' + else: count='min' + + snrcut = kwargs.get('snrcut',0.) + debias = kwargs.get('debias',True) + + # unpack data & mask low snr points + vtype=vis_poldict[pol] + clamparr = Obsdata.c_log_amplitudes_diag(vtype=vtype,count=count,debias=debias,snrcut=snrcut) + + # loop over timestamps + clamp_diag = [] + sigma_diag = [] + uv_diag = [] + tform_mats = [] + for ic, cl in enumerate(clamparr): + + # get diagonalized log closure amplitudes and errors + clamp_diag.append(cl[0]['camp']) + sigma_diag.append(cl[0]['sigmaca']) + + # get uv arrays + u1 = cl[2][:,0].astype('float') + v1 = cl[3][:,0].astype('float') + uv1 = np.hstack((u1.reshape(-1,1), v1.reshape(-1,1))) + + u2 = cl[2][:,1].astype('float') + v2 = cl[3][:,1].astype('float') + uv2 = np.hstack((u2.reshape(-1,1), v2.reshape(-1,1))) + + u3 = cl[2][:,2].astype('float') + v3 = cl[3][:,2].astype('float') + uv3 = np.hstack((u3.reshape(-1,1), v3.reshape(-1,1))) + + u4 = cl[2][:,3].astype('float') + v4 = cl[3][:,3].astype('float') + uv4 = np.hstack((u4.reshape(-1,1), v4.reshape(-1,1))) + + # compute Fourier matrices + uv = (uv1, + uv2, + uv3, + uv4 + ) + uv_diag.append(uv) + + # get transformation matrix for this timestamp + tform_mats.append(cl[4].astype('float')) + + # combine Fourier and transformation matrices into tuple for outputting + uvmatrices = (np.array(uv_diag),np.array(tform_mats)) + + return (np.array(clamp_diag), np.array(sigma_diag), uvmatrices, None) + +def chisqdata_pvis(Obsdata, pol='I', **kwargs): + data_arr = Obsdata.unpack(['t1','t2','u','v','pvis','psigma','el1','el2','par_ang1','par_ang2'], conj=True) + uv = np.hstack((data_arr['u'].reshape(-1,1), data_arr['v'].reshape(-1,1))) + mask = np.isfinite(data_arr['pvis'] + data_arr['psigma']) # don't include nan (missing data) or inf (division by zero) + jonesdict = make_jonesdict(Obsdata, data_arr[mask]) + return (data_arr['pvis'][mask], data_arr['psigma'][mask], uv[mask], jonesdict) + +def chisqdata_m(Obsdata, pol='I',**kwargs): + debias = kwargs.get('debias',True) + data_arr = Obsdata.unpack(['t1','t2','u','v','m','msigma','el1','el2','par_ang1','par_ang2'], conj=True, debias=False) + uv = np.hstack((data_arr['u'].reshape(-1,1), data_arr['v'].reshape(-1,1))) + mask = np.isfinite(data_arr['m'] + data_arr['msigma']) # don't include nan (missing data) or inf (division by zero) + jonesdict = make_jonesdict(Obsdata, data_arr[mask]) + return (data_arr['m'][mask], data_arr['msigma'][mask], uv[mask], jonesdict) + +def chisqdata_fracpol(Obsdata, pol_upper,pol_lower,**kwargs): + debias = kwargs.get('debias',True) + data_arr = Obsdata.unpack(['t1','t2','u','v','m','msigma','el1','el2','par_ang1','par_ang2','rrvis','rlvis','lrvis','llvis','rramp','rlamp','lramp','llamp','rrsigma','rlsigma','lrsigma','llsigma'], conj=False, debias=True) + uv = np.hstack((data_arr['u'].reshape(-1,1), data_arr['v'].reshape(-1,1))) + + upper = data_arr[pol_upper + 'vis'] + lower = data_arr[pol_lower + 'vis'] + upper_amp = data_arr[pol_upper + 'amp'] + lower_amp = data_arr[pol_lower + 'amp'] + upper_sig = data_arr[pol_upper + 'sigma'] + lower_sig = data_arr[pol_lower + 'sigma'] + + sig = ((upper_sig/lower_amp)**2 + (lower_sig*upper_amp/lower_amp**2)**2)**0.5 + + # Mask bad data + mask = np.isfinite(upper + lower + sig) # don't include nan (missing data) or inf (division by zero) + jonesdict = make_jonesdict(Obsdata, data_arr[mask]) + + return ((upper/lower)[mask], sig[mask], uv[mask], jonesdict) + +def chisqdata_polclosure(Obsdata, **kwargs): + debias = kwargs.get('debias',True) + data_arr = Obsdata.unpack(['t1','t2','u','v','m','msigma','el1','el2','par_ang1','par_ang2','rrvis','rlvis','lrvis','llvis','rramp','rlamp','lramp','llamp','rrsigma','rlsigma','lrsigma','llsigma'], conj=False, debias=True) + uv = np.hstack((data_arr['u'].reshape(-1,1), data_arr['v'].reshape(-1,1))) + + RL = data_arr['rlvis'] + LR = data_arr['lrvis'] + RR = data_arr['rrvis'] + LL = data_arr['llvis'] + vis = (RL * LR)/(RR * LL) + sig = (np.abs(LR/(LL*RR) * data_arr['rlsigma'])**2 + +np.abs(RL/(LL*RR) * data_arr['lrsigma'])**2 + +np.abs(LR*RL/(RR**2*LL) * data_arr['rrsigma'])**2 + +np.abs(RL*LR/(LL**2*RR) * data_arr['llsigma'])**2)**0.5 + + # Mask bad data + mask = np.isfinite(vis + sig) # don't include nan (missing data) or inf (division by zero) + jonesdict = make_jonesdict(Obsdata, data_arr[mask]) + + return (vis[mask], sig[mask], uv[mask], jonesdict) diff --git a/examples/example_modeling.py b/examples/example_modeling.py new file mode 100644 index 00000000..56f2d833 --- /dev/null +++ b/examples/example_modeling.py @@ -0,0 +1,94 @@ +# Note: this is an example sequence of commands to run in ipython +# The matplotlib windows may not open/close properly if you run this directly as a script + +from __future__ import division +from __future__ import print_function + +import matplotlib.pyplot as plt +plt.ion + +import numpy as np +import ehtim as eh + +# Load a sample array +eht = eh.array.load_txt('../arrays/EHT2019.txt') + +### Make a simple model ### + +# Start with an empty model object +mod = eh.model.Model() + +# Add a ring model +mod = mod.add_ring(F0=1.5, d=40.*eh.RADPERUAS) + +# View the model +mod.display() + +# Add another model component +mod = mod.add_circ_gauss(1., 20.*eh.RADPERUAS, x0=-15.*eh.RADPERUAS, y0=20.*eh.RADPERUAS) + +# View the model after blurring with a circular Gaussian +mod.blur_circ(20.*eh.RADPERUAS).display() + +# Make an image of the model +im = mod.make_image(200.*eh.RADPERUAS, 1024) + +# Observe the model +tint_sec = 5 +tadv_sec = 3600 +tstart_hr = 0 +tstop_hr = 24 +bw_hz = 1e9 +obs = mod.observe(eht, tint_sec, tadv_sec, tstart_hr, tstop_hr, bw_hz, ampcal=True, phasecal=False) + +# Compare the observation to the model +eh.comp_plots.plotall_obs_im_compare(obs, mod, 'uvdist', 'amp') +eh.comp_plots.plotall_obs_im_compare(obs, mod, 'uvdist', 'phase') + +### Next, we'll try some model fitting ### + +# First, we need to define the model that we are using to fit an observation +# Some algorithms (e.g., gradient descent) also use this as the initialization +mod_init = eh.model.Model() +mod_init = mod_init.add_ring(1.5, 40.*eh.RADPERUAS) +mod_init = mod_init.add_circ_gauss(1., 20.*eh.RADPERUAS) + +# Next, define the prior for each fitted parameter +# Because we don't have absolute phase information, we'll force the ring to be centered on the origin +mod_prior = mod_init.default_prior() +# Ring: +mod_prior[0]['F0'] = {'prior_type':'flat', 'min':0.0, 'max':2.0} +mod_prior[0]['d'] = {'prior_type':'flat', 'min':1.0*eh.RADPERUAS, 'max':50.0*eh.RADPERUAS} +mod_prior[0]['x0'] = {'prior_type':'fixed'} +mod_prior[0]['y0'] = {'prior_type':'fixed'} +# Gaussian +mod_prior[1]['F0'] = {'prior_type':'flat', 'min':0.0, 'max':2.0} +mod_prior[1]['FWHM'] = {'prior_type':'flat', 'min':1.0*eh.RADPERUAS, 'max':50.0*eh.RADPERUAS} +mod_prior[1]['x0'] = {'prior_type':'gauss','mean':0.0,'std':20.*eh.RADPERUAS} +mod_prior[1]['y0'] = {'prior_type':'gauss','mean':0.0,'std':20.*eh.RADPERUAS} + +### Fit the model using dynamic nested sampling; this estimates the full posterior ### +# Fit using amplitudes and closure phases +mod_fit = eh.modeler_func(obs, mod_init, mod_prior, d1='amp', d2='cphase', minimizer_func='dynesty_dynamic') + +# View the fitted model (the MAP) +mod_fit['model'].display() +mod_fit['model'].blur_circ(20.*eh.RADPERUAS).display() + +# View a few random samples from the posterior +eh.imaging.dynamical_imaging.plot_im_List([_.make_image(100.*eh.RADPERUAS, 128) for _ in mod_fit['posterior_models'][:10]]) + +# Compare the fitted model to the data +eh.comp_plots.plotall_obs_im_compare(obs,mod_fit['model'],'uvdist','amp') + +# Compare a sample of fitted models drawn from the posterior to the data +eh.comp_plots.plotall_obs_im_compare(obs,mod_fit['posterior_models'],'uvdist','amp') + +# Compare posteriors from with the true model values ('natural' uses natural units and rescalings) +from dynesty import plotting as dyplot +cfig, caxes = dyplot.cornerplot(mod_fit['res_natural'], labels=mod_fit['labels_natural'], truths=[1.5, 40, 1.0, 20, -15, 20]) +cfig.set_size_inches((2.5*len(mod_fit['labels']),2.5*len(mod_fit['labels']))) +cfig.show() + +# Save the MAP +mod_fit['model'].save_txt('sample_modelfit.txt') diff --git a/setup.py b/setup.py index 08048abf..b4137f73 100755 --- a/setup.py +++ b/setup.py @@ -22,6 +22,7 @@ def read(fname): "ehtim.calibrating", "ehtim.imaging", "ehtim.io", + "ehtim.modeling", "ehtim.observing", "ehtim.plotting", "ehtim.features", From 3df6f2b0a27d26ae4645f1d0af3c728f0bd4ee61 Mon Sep 17 00:00:00 2001 From: michaeldjohnson Date: Sat, 19 Jun 2021 09:43:07 -0400 Subject: [PATCH 21/28] Fixed merge_obs to track polrep. --- ehtim/imaging/dynamical_imaging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ehtim/imaging/dynamical_imaging.py b/ehtim/imaging/dynamical_imaging.py index bc7a2777..a19d917f 100644 --- a/ehtim/imaging/dynamical_imaging.py +++ b/ehtim/imaging/dynamical_imaging.py @@ -305,7 +305,7 @@ def merge_obs(obs_List): #The important things to merge are the mjd and the data data_merge = np.hstack([obs.data for obs in obs_List]) - return obsdata.Obsdata(obs_List[0].ra, obs_List[0].dec, obs_List[0].rf, obs_List[0].bw, data_merge, obs_List[0].tarr, source=obs_List[0].source, mjd=obs_List[0].mjd, ampcal=obs_List[0].ampcal, phasecal=obs_List[0].phasecal) + return obsdata.Obsdata(obs_List[0].ra, obs_List[0].dec, obs_List[0].rf, obs_List[0].bw, data_merge, obs_List[0].tarr, polrep=obs_List[0].polrep, source=obs_List[0].source, mjd=obs_List[0].mjd, ampcal=obs_List[0].ampcal, phasecal=obs_List[0].phasecal) def average_im_list(im_List): """Return the average of a list of images From d677f1388c92318221d2fb419c022b1b5dc5f023 Mon Sep 17 00:00:00 2001 From: michaeldjohnson Date: Sat, 19 Jun 2021 10:09:51 -0400 Subject: [PATCH 22/28] Minor changes to imgsum to allow compatibility with a model. --- ehtim/plotting/summary_plots.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ehtim/plotting/summary_plots.py b/ehtim/plotting/summary_plots.py index 82645949..f7de5af3 100644 --- a/ehtim/plotting/summary_plots.py +++ b/ehtim/plotting/summary_plots.py @@ -152,6 +152,8 @@ def imgsum(im_or_mov, obs, obs_uncal, outname, outdir='.', title='imgsum', comme # TODO --- ok to always extrapolate? if force_extrapolate: im_or_mov.reset_interp(bounds_error=False) + elif hasattr(im_or_mov, 'make_image'): + im_display = im_or_mov.make_image(obs.res() * 10., 512) else: im_display = im_or_mov.copy() @@ -239,7 +241,7 @@ def imgsum(im_or_mov, obs, obs_uncal, outname, outdir='.', title='imgsum', comme ha='left', va='center', transform=ax.transAxes) ax.text(.23, .5, "%0.0f GHz" % (im_or_mov.rf/1.e9), fontsize=fs, ha='left', va='center', transform=ax.transAxes) - ax.text(.23, .3, "%0.1f $\mu$as" % (im_or_mov.fovx()/ehc.RADPERUAS), fontsize=fs, + ax.text(.23, .3, "%0.1f $\mu$as" % (im_display.fovx()/ehc.RADPERUAS), fontsize=fs, ha='left', va='center', transform=ax.transAxes) ax.text(.23, .1, "%0.2f Jy" % flux, fontsize=fs, ha='left', va='center', transform=ax.transAxes) From 88615d913524acf1afdd631f3c0cfb7a57cc789d Mon Sep 17 00:00:00 2001 From: michaeldjohnson Date: Sat, 19 Jun 2021 10:22:31 -0400 Subject: [PATCH 23/28] Added model class functionality to chisq --- ehtim/obsdata.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ehtim/obsdata.py b/ehtim/obsdata.py index 0118141e..35ecd6d4 100644 --- a/ehtim/obsdata.py +++ b/ehtim/obsdata.py @@ -959,6 +959,7 @@ def chisq(self, im_or_mov, dtype='vis', pol='I', ttype='nfft', mask=[], **kwargs # TODO -- should import this at top, but the circular dependencies create a mess... import ehtim.imaging.imager_utils as iu + import ehtim.modeling.modeling_utils as mu # Movie -- weighted sum of all frame chi^2 values if hasattr(im_or_mov, 'get_image'): @@ -991,6 +992,11 @@ def chisq(self, im_or_mov, dtype='vis', pol='I', ttype='nfft', mask=[], **kwargs chisq = np.sum(np.array(num_list) * np.array(chisq_list)) / np.sum(num_list) + # Model -- single chi^2 + elif hasattr(im_or_mov,'N_models'): + (data, sigma, uv, jonesdict) = mu.chisqdata(self, dtype, pol, **kwargs) + chisq = mu.chisq(im_or_mov, uv, data, sigma, dtype, jonesdict) + # Image -- single chi^2 else: im = im_or_mov From 437a86851ce0c096fbd6f38a9d477d2ac2a4c5ec Mon Sep 17 00:00:00 2001 From: michaeldjohnson Date: Wed, 23 Jun 2021 12:36:44 -0400 Subject: [PATCH 24/28] Fixed how caltable library is loaded. --- ehtim/modeling/modeling_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ehtim/modeling/modeling_utils.py b/ehtim/modeling/modeling_utils.py index 9c445bcd..b2fbc05b 100644 --- a/ehtim/modeling/modeling_utils.py +++ b/ehtim/modeling/modeling_utils.py @@ -362,10 +362,10 @@ def selfcal(Obsdata, model, try: caldict[site] = np.append(caldict[site], row[site]) except KeyError: caldict[site] = dat - caltable = ehtim.caltable.Caltable(obs.ra, obs.dec, obs.rf, obs.bw, caldict, obs.tarr, + ct = caltable.Caltable(obs.ra, obs.dec, obs.rf, obs.bw, caldict, obs.tarr, source=obs.source, mjd=obs.mjd, timetype=obs.timetype) - return caltable + return ct def make_param_map(model_init, model_prior, minimizer_func, fit_model, fit_pol=False, fit_cpol=False): # Define the mapping between solved parameters and the model @@ -1622,9 +1622,9 @@ def logpost(j): for site in caldict.keys(): caldict[site] = np.array(caldict[site], dtype=DTCAL) - caltable = ehtim.caltable.Caltable(Obsdata.ra, Obsdata.dec, Obsdata.rf, Obsdata.bw, caldict, Obsdata.tarr, + ct = caltable.Caltable(Obsdata.ra, Obsdata.dec, Obsdata.rf, Obsdata.bw, caldict, Obsdata.tarr, source=Obsdata.source, mjd=Obsdata.mjd, timetype=Obsdata.timetype) - ret['caltable'] = caltable + ret['caltable'] = ct # If relevant, return useful quantities associated with the leakage if station_leakages is not None: From 6a9482f7f45455928d05c566cb35b4b1b4d62de2 Mon Sep 17 00:00:00 2001 From: michaeldjohnson Date: Fri, 9 Jul 2021 14:31:50 -0400 Subject: [PATCH 25/28] Changed closure phase normalization term from Gaussian approximation to (exact) von Mises. --- ehtim/modeling/modeling_utils.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/ehtim/modeling/modeling_utils.py b/ehtim/modeling/modeling_utils.py index b2fbc05b..0df088f1 100644 --- a/ehtim/modeling/modeling_utils.py +++ b/ehtim/modeling/modeling_utils.py @@ -411,6 +411,14 @@ def compute_likelihood_constants(d1, d2, d3, sigma1, sigma2, sigma3): ln_norm3 = -np.sum(np.log((2.0*np.pi)**0.5 * sigma3)) except: pass + # Use the correct von Mises normalization if using closure phase + if d1 in ['cphase','cphase_diag']: + ln_norm1 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/sigma1**2))) + if d2 in ['cphase','cphase_diag']: + ln_norm2 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/sigma2**2))) + if d3 in ['cphase','cphase_diag']: + ln_norm2 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/sigma3**2))) + if d1 in ['vis','bs','m','pvis','rrll','llrr','lrll','rlll','lrrr','rlrr','polclosure']: alpha_d1 *= 2 ln_norm1 *= 2 From e4b246dc0686c945e2efe7e988c0be1cd98a4fa1 Mon Sep 17 00:00:00 2001 From: michaeldjohnson Date: Fri, 10 Sep 2021 23:32:38 -0400 Subject: [PATCH 26/28] Fixed likelihood prefactor for closure phases (was missing degrees-to-radians conversion for sigma) --- ehtim/modeling/modeling_utils.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ehtim/modeling/modeling_utils.py b/ehtim/modeling/modeling_utils.py index 0df088f1..7fb6d9ad 100644 --- a/ehtim/modeling/modeling_utils.py +++ b/ehtim/modeling/modeling_utils.py @@ -411,13 +411,14 @@ def compute_likelihood_constants(d1, d2, d3, sigma1, sigma2, sigma3): ln_norm3 = -np.sum(np.log((2.0*np.pi)**0.5 * sigma3)) except: pass + # If using closure phase, the sigma is given in degrees, not radians! # Use the correct von Mises normalization if using closure phase if d1 in ['cphase','cphase_diag']: - ln_norm1 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/sigma1**2))) + ln_norm1 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/(sigma1 * DEGREE)**2))) if d2 in ['cphase','cphase_diag']: - ln_norm2 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/sigma2**2))) + ln_norm2 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/(sigma2 * DEGREE)**2))) if d3 in ['cphase','cphase_diag']: - ln_norm2 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/sigma3**2))) + ln_norm3 = -np.sum(np.log(2.0*np.pi*sps.ive(0, 1.0/(sigma3 * DEGREE)**2))) if d1 in ['vis','bs','m','pvis','rrll','llrr','lrll','rlll','lrrr','rlrr','polclosure']: alpha_d1 *= 2 @@ -873,6 +874,7 @@ def modeler_func(Obsdata, model_init, model_prior, fit_model=True, fit_pol=False, fit_cpol=False, fit_gains=False,marginalize_gains=False,gain_init=None,gain_prior=None, fit_leakage=False, leakage_init=None, leakage_prior=None, + fit_noise_model=False, minimizer_func='scipy.optimize.minimize', minimizer_kwargs=None, bounds=None, use_bounds=False, From 37e986030b1e10f6e3b89a189ccad4a219ab891e Mon Sep 17 00:00:00 2001 From: michaeldjohnson Date: Mon, 4 Oct 2021 17:53:59 -0400 Subject: [PATCH 27/28] Added mring model with a Gaussian floor. --- ehtim/model.py | 85 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/ehtim/model.py b/ehtim/model.py index 5f983095..f20f4cab 100644 --- a/ehtim/model.py +++ b/ehtim/model.py @@ -117,6 +117,13 @@ def add_pol(): params.append('beta' + str(j+1) + complex_labels[0]) params.append('beta' + str(j+1) + complex_labels[1]) add_pol() + elif model_type == 'thick_mring_Gfloor': + params = ['F0','d','alpha','ff','FWHM','x0','y0'] + for j in range(len(model_params['beta_list'])): + params.append('beta' + str(j+1) + complex_labels[0]) + params.append('beta' + str(j+1) + complex_labels[1]) + add_pol() + elif model_type == 'stretched_thick_mring': params = ['F0','d','alpha', 'x0','y0'] for j in range(len(model_params['beta_list'])): @@ -210,6 +217,14 @@ def default_prior(model_type,model_params=None,fit_pol=False,fit_cpol=False): for j in range(len(model_params['beta_list'])): prior['beta' + str(j+1) + complex_labels[0]] = complex_priors[0] prior['beta' + str(j+1) + complex_labels[1]] = complex_priors[1] + elif model_type == 'thick_mring_Gfloor': + prior['d'] = {'prior_type':'positive','transform':'log'} + prior['alpha'] = {'prior_type':'positive','transform':'log'} + prior['ff'] = {'prior_type':'flat','min':0,'max':1} + prior['FWHM'] = {'prior_type':'positive','transform':'log'} + for j in range(len(model_params['beta_list'])): + prior['beta' + str(j+1) + complex_labels[0]] = complex_priors[0] + prior['beta' + str(j+1) + complex_labels[1]] = complex_priors[1] elif model_type == 'stretched_thick_mring': prior['d'] = {'prior_type':'positive','transform':'log'} prior['alpha'] = {'prior_type':'positive','transform':'log'} @@ -392,6 +407,9 @@ def f(r): elif model_type == 'thick_mring_floor': val = (1.0 - params['ff']) * sample_1model_xy(x, y, 'thick_mring', params, psize=psize, pol=pol) val += params['ff'] * sample_1model_xy(x, y, 'blurred_disk', params, psize=psize, pol=pol) + elif model_type == 'thick_mring_Gfloor': + val = (1.0 - params['ff']) * sample_1model_xy(x, y, 'thick_mring', params, psize=psize, pol=pol) + val += params['ff'] * sample_1model_xy(x, y, 'circ_gauss', params, psize=psize, pol=pol) elif model_type[:9] == 'stretched': params_stretch = params.copy() params_stretch['F0'] /= params['stretch'] @@ -531,6 +549,9 @@ def sample_1model_uv(u, v, model_type, params, pol='I', jonesdict=None): elif model_type == 'thick_mring_floor': val = (1.0 - params['ff']) * sample_1model_uv(u, v, 'thick_mring', params, pol=pol, jonesdict=jonesdict) val += params['ff'] * sample_1model_uv(u, v, 'blurred_disk', params, pol=pol, jonesdict=jonesdict) + elif model_type == 'thick_mring_Gfloor': + val = (1.0 - params['ff']) * sample_1model_uv(u, v, 'thick_mring', params, pol=pol, jonesdict=jonesdict) + val += params['ff'] * sample_1model_uv(u, v, 'circ_gauss', params, pol=pol, jonesdict=jonesdict) elif model_type[:9] == 'stretched': params_stretch = params.copy() params_stretch['x0'] = 0.0 @@ -755,6 +776,9 @@ def sample_1model_graduv_uv(u, v, model_type, params, pol='I', jonesdict=None): elif model_type == 'thick_mring_floor': val = (1.0 - params['ff']) * sample_1model_graduv_uv(u, v, 'thick_mring', params, pol=pol, jonesdict=jonesdict) val += params['ff'] * sample_1model_graduv_uv(u, v, 'blurred_disk', params, pol=pol, jonesdict=jonesdict) + elif model_type == 'thick_mring_Gfloor': + val = (1.0 - params['ff']) * sample_1model_graduv_uv(u, v, 'thick_mring', params, pol=pol, jonesdict=jonesdict) + val += params['ff'] * sample_1model_graduv_uv(u, v, 'circ_gauss', params, pol=pol, jonesdict=jonesdict) elif model_type[:9] == 'stretched': # Take care of the degenerate origin point by a small offset u += (u==0.)*(v==0.)*1e-10 @@ -1085,6 +1109,36 @@ def sample_1model_grad_uv(u, v, model_type, params, pol='I', fit_pol=False, fit_ grad.append(grad_mring[3] + grad_disk[3]) grad.append(grad_mring[4] + grad_disk[4]) + # Add remaining gradients for the mring + for j in range(5,len(grad_mring)): + grad.append(grad_mring[j]) + + val = np.array(grad) + elif model_type == 'thick_mring_Gfloor': # F0, d, [alpha], ff, FWHM, x0, y0, beta1_re/abs, beta1_im/arg, beta2_re/abs, beta2_im/arg, ... + # We need to stich together the two gradients for the mring and the gaussian; we also need to add the gradient for the floor fraction ff + grad_mring = (1.0 - params['ff']) * sample_1model_grad_uv(u, v, 'thick_mring', params, pol=pol, fit_pol=fit_pol, fit_cpol=fit_cpol) + grad_gauss = params['ff'] * sample_1model_grad_uv(u, v, 'circ_gauss', params, pol=pol, fit_pol=fit_pol, fit_cpol=fit_cpol) + + # mring: F0, d, alpha, x0, y0, beta1_re/abs, beta1_im/arg, beta2_re/abs, beta2_im/arg, ... + # gauss: F0, [d, alpha] FWHM, x0, y0 + + grad = [] + grad.append(grad_mring[0] + grad_gauss[0]) # Here are derivatives for F0 + + # Here are derivatives for d, and alpha + grad.append(grad_mring[1]) + grad.append(grad_mring[2]) + + # Here is the derivative for ff + grad.append( params['F0'] * (grad_gauss[0]/params['ff'] - grad_mring[0]/(1.0 - params['ff'])) ) + + # Now the derivatives for FWHM + grad.append(grad_gauss[1]) + + # Now the derivatives for x0 and y0 + grad.append(grad_mring[3] + grad_gauss[2]) + grad.append(grad_mring[4] + grad_gauss[3]) + # Add remaining gradients for the mring for j in range(5,len(grad_mring)): grad.append(grad_mring[j]) @@ -1604,6 +1658,37 @@ def add_thick_mring_floor(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERU out.params.append({'F0':F0,'d':d,'beta_list':np.array(beta_list, dtype=np.complex_),'beta_list_pol':np.array(beta_list_pol, dtype=np.complex_),'beta_list_cpol':np.array(beta_list_cpol, dtype=np.complex_),'alpha':alpha,'x0':x0,'y0':y0,'ff':ff}) return out + def add_thick_mring_Gfloor(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, ff=0.0, FWHM = 50.*RADPERUAS, x0 = 0.0, y0 = 0.0, beta_list = None, beta_list_pol = None, beta_list_cpol = None): + """Add a ring model with azimuthal brightness variations determined by a Fourier mode expansion, thickness determined by circular Gaussian convolution, and a floor + The floor is a circular Gaussian, with size FWHM + For details, see Eq. 18-20 of https://arxiv.org/abs/1907.04329 + The Gaussian convolution calculation is a trivial generalization of Appendix G of https://iopscience.iop.org/article/10.3847/2041-8213/ab0e85/pdf + + Args: + F0 (float): The total flux of the model + d (float): The ring diameter (radians) + alpha (float): The ring thickness (FWHM of Gaussian convolution) (radians) + FWHM (float): The Gaussian FWHM + x0 (float): The x-coordinate (radians) + y0 (float): The y-coordinate (radians) + ff (float): The fraction of the total flux in the floor + beta_list (list): List of complex Fourier coefficients, [beta_1, beta_2, ...]. + Negative indices are determined by the condition beta_{-m} = beta_m*. + Indices are all scaled by F0 = beta_0, so they are dimensionless. + Returns: + (Model): Updated Model + """ + if beta_list is None: beta_list = [] + if beta_list_pol is None: beta_list_pol = [] + if beta_list_cpol is None: beta_list_cpol = [] + + out = self.copy() + if beta_list is None: + beta_list = [0.0] + out.models.append('thick_mring_Gfloor') + out.params.append({'F0':F0,'d':d,'beta_list':np.array(beta_list, dtype=np.complex_),'beta_list_pol':np.array(beta_list_pol, dtype=np.complex_),'beta_list_cpol':np.array(beta_list_cpol, dtype=np.complex_),'alpha':alpha,'x0':x0,'y0':y0,'ff':ff,'FWHM':FWHM}) + return out + def add_stretched_thick_mring(self, F0 = 1.0, d = 50.*RADPERUAS, alpha = 10.*RADPERUAS, x0 = 0.0, y0 = 0.0, beta_list = None, beta_list_pol = None, beta_list_cpol = None, stretch = 1.0, stretch_PA = 0.0): """Add a ring model with azimuthal brightness variations determined by a Fourier mode expansion and thickness determined by circular Gaussian convolution. For details, see Eq. 18-20 of https://arxiv.org/abs/1907.04329 From 0b04dfaaf8288ec7c73984a22a0cf4b78d26c623 Mon Sep 17 00:00:00 2001 From: Andrew Chael Date: Fri, 29 Oct 2021 13:52:03 -0400 Subject: [PATCH 28/28] update setup.py to version 1.2.3 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index b4137f73..31fe89ae 100755 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ def read(fname): if __name__ == "__main__": setup(name="ehtim", - version = "1.2.2", + version = "1.2.3", author = "Andrew Chael", author_email = "achael@princeton.edu",