Commit 1a4b1652 authored by Rustam Gainutdinov's avatar Rustam Gainutdinov
Browse files

Added MCMC parameter estimation for massless Two-Body problem

parent 04e1013a
__precompile__()
module Cornerplot
#Pretty corner plot for MCMC chains
using PyPlot
using Statistics
export chain_estimates, cornerplot
function chain_estimates(fchain)
n = size(fchain)[1]
fchmeds = [quantile!(fchain[i,:], 0.5) for i in 1:n]
f_alims = [quantile!(fchain[i,:], 0.16) for i in 1:n]
f_blims = [quantile!(fchain[i,:], 0.84) for i in 1:n]
f_lerr = fchmeds - f_alims
f_rerr = f_blims - fchmeds
return fchmeds, f_lerr, f_rerr
end
function cornerplot(fchain, lnames, lnames_ed, hbinsize, cmap, errsize)
println("Printing...")
n = size(fchain)[1]
figure(figsize = (25,25))
fchmeds, f_lerr, f_rerr = chain_estimates(fchain)
fchmedsstr = string.(round.(fchmeds; sigdigits=3))
f_lerrstr = string.(round.(f_lerr; sigdigits=3))
f_rerrstr = string.(round.(f_rerr; sigdigits=3))
lborders = fchmeds .- f_lerr .* errsize
rborders = fchmeds .+ f_rerr .* errsize
flims = [(lborders[i], rborders[i]) for i in 1:n]
for i in 1:n
for j in 1:i
ax = subplot2grid((n, n), (i-1, j-1))
if i == j
mesh = fchain[i,:]
mesh = mesh[flims[i][1] .< mesh .< flims[i][2]]
#xlim(flims[i])
hist(mesh, bins=hbinsize)
title("\$"*lnames[i]*"="*fchmedsstr[i]*"^{+"*f_rerrstr[i]*"}_{-"*f_lerrstr[i]*"}\$", fontsize=9)
else
#xlim(flims[j])
#ylim(flims[i])
mesh1 = fchain[j,:]
mesh2 = fchain[i,:]
mesh2 = mesh2[flims[j][1] .< mesh1 .< flims[j][2]]
mesh1 = mesh1[flims[j][1] .< mesh1 .< flims[j][2]]
mesh1 = mesh1[flims[i][1] .< mesh2 .< flims[i][2]]
mesh2 = mesh2[flims[i][1] .< mesh2 .< flims[i][2]]
hexbin(mesh1, mesh2, gridsize=hbinsize, cmap=cmap)
end
if i < n
setp(ax.get_xticklabels(), visible=false)
else
xlabel(lnames_ed[j], fontsize=11)
end
if j > 1 || j == i && j == 1
setp(ax.get_yticklabels(), visible=false)
elseif i > 1
ylabel(lnames_ed[i], fontsize=11)
end
end
end
tight_layout()
end
end
__precompile__()
module DataSets
using CSV
export sstars_data, θ₀
sstars_df = CSV.read("./data/sstarsdata.csv", header=false)
sstars_t = sstars_df[!, 1]
S2_RA = sstars_df[!, 2]
S2_ΔRA = sstars_df[!, 3]
S2_Dec = sstars_df[!, 4]
S2_ΔDec = sstars_df[!, 5]
S2_RV = sstars_df[!, 6]
S2_ΔRV = sstars_df[!, 7]
S38_RA = sstars_df[!, 8]
S38_ΔRA = sstars_df[!, 9]
S38_Dec = sstars_df[!, 10]
S38_ΔDec = sstars_df[!, 11]
S38_RV = sstars_df[!, 12]
S38_ΔRV = sstars_df[!, 13]
S55_RA = sstars_df[!, 14]
S55_ΔRA = sstars_df[!, 15]
S55_Dec = sstars_df[!, 16]
S55_ΔDec = sstars_df[!, 17]
S55_RV = sstars_df[!, 18]
S55_ΔRV = sstars_df[!, 19]
sstars_art = sstars_df[!, 20]
sstars_data = [sstars_t, [S2_RA, S2_Dec, S2_RV, S38_RA, S38_Dec, S38_RV, S55_RA, S55_Dec, S55_RV], [S2_ΔRA, S2_ΔDec, S2_ΔRV, S38_ΔRA, S38_ΔDec, S38_ΔRV, S55_ΔRA, S55_ΔDec, S55_ΔRV]]
θ₀ = [98947.158163, 161544.112937, -2.19582537e+02, 6.04500473e+02, 1.33902287e+02, 2.26126780e+02, 215932.83605999997, 29862.5963615, 6.71724897e+01, 5.87446040e+02, 1.70792316e+02, 9.34728478e+01, 139381.778738, -37591.510541999996, -3.45931227e+02, 1.00503663e+03, 1.53023909e+02, 3.23354746e+02, 4.31, 8.3, 0.0, 0.0, -3.49118611e-06, 1.47495749e-05, 1.0, 1.0]
end
......@@ -34,7 +34,4 @@ module EoMs
end
#S2 params: 98947.158163, 161544.112937, -0.0007324487518145469, 0.0020163968040454333, 4.31, 1.0, 1.0
#θ₀ = [(98947.158163, 161544.112937, -2.19582537e+02, 6.04500473e+02, 1.33902287e+02, 2.26126780e+02), (215932.83605999997, 29862.5963615, 6.71724897e+01, 5.87446040e+02, 1.70792316e+02, 9.34728478e+01), (139381.778738, -37591.510541999996, -3.45931227e+02, 1.00503663e+03, 1.53023909e+02, 3.23354746e+02), 4.31, 8.3, 0.0, 0.0, -3.49118611e-06, 1.47495749e-05, 1.0, 1.0]
end
......@@ -64,8 +64,8 @@ module Models
function Rømer_delay(sol, Ω, i, t_obs)
z = t -> turnorbit(sol(t)[1], sol(t)[2], 0.0, Ω, i)
f = t_em_arg -> t_em_arg - z(t_em_arg) - t_obs
z = t -> turnorbit(sol(t)[1], sol(t)[2], 0.0, Ω, i)[3]
f = t_em_arg -> (t_em_arg - z(t_em_arg) - t_obs)
t_em = find_zero(f, t_obs)
return t_em
......@@ -74,23 +74,19 @@ module Models
function model_massless(θ, 𝐭_input)
S2_orb, S38_orb, S55_orb, M, R0, x_SgrA, y_SgrA, vx_SgrA, vy_SgrA, β, γ = θ
S2_x0, S2_y0, S2_vx0, S2_vy0, S2_i, S2_Ω, S38_x0, S38_y0, S38_vx0, S38_vy0, S38_i, S38_Ω, S55_x0, S55_y0, S55_vx0, S55_vy0, S55_i, S55_Ω, M, R0, x_SgrA, y_SgrA, vx_SgrA, vy_SgrA, β, γ = θ
S2_x0, S2_y0, S2_vx0, S2_vy0, S2_i, S2_Ω = S2_orb
S38_x0, S38_y0, S38_vx0, S38_vy0, S38_i, S38_Ω = S38_orb
S55_x0, S55_y0, S55_vx0, S55_vy0, S55_i, S55_Ω = S55_orb
S2_sol = solve_EoM_massless(S2_x0, S2_y0, S2_vx0 / 299792.458, S2_vy0 / 299792.458, M, β, γ)
S38_sol = solve_EoM_massless(S38_x0, S38_y0, S38_vx0 / 299792.458, S38_vy0 / 299792.458, M, β, γ)
S55_sol = solve_EoM_massless(S55_x0, S55_y0, S55_vx0 / 299792.458, S55_vy0 / 299792.458, M, β, γ)
S2_sol = solve_EoM_massless(S2_x0, S2_y0, S2_vx0 / c_to_kms, S2_vy0 / c_to_kms, M, β, γ)
S38_sol = solve_EoM_massless(S38_x0, S38_y0, S38_vx0 / c_to_kms, S38_vy0 / c_to_kms, M, β, γ)
S55_sol = solve_EoM_massless(S55_x0, S55_y0, S55_vx0 / c_to_kms, S55_vy0 / c_to_kms, M, β, γ)
𝐭_em_S2 = Rømer_delay.(Ref(S2_sol), Ref(S2_Ω), Ref(S2_i), 𝐭_input)
𝐭_em_S38 = Rømer_delay.(Ref(S38_sol), Ref(S38_Ω), Ref(S38_i), 𝐭_input)
𝐭_em_S55 = Rømer_delay.(Ref(S55_sol), Ref(S55_Ω), Ref(S55_i), 𝐭_input)
S2_𝐱, S2_𝐲, S2_𝐯𝐱, S2_𝐯𝐲 = S2_sol.(𝐭_em_S2)
S38_𝐱, S38_𝐲, S38_𝐯𝐱, S38_𝐯𝐲 = S38_sol.(𝐭_em_S38)
S55_𝐱, S55_𝐲, S55_𝐯𝐱, S55_𝐯𝐲 = S55_sol.(𝐭_em_S55)
S2_𝐱, S2_𝐲, S2_𝐯𝐱, S2_𝐯𝐲 = [hcat(S2_sol.(𝐭_em_S2)...)[i,:] for i in 1:4]
S38_𝐱, S38_𝐲, S38_𝐯𝐱, S38_𝐯𝐲 = [hcat(S38_sol.(𝐭_em_S38)...)[i,:] for i in 1:4]
S55_𝐱, S55_𝐲, S55_𝐯𝐱, S55_𝐯𝐲 = [hcat(S55_sol.(𝐭_em_S55)...)[i,:] for i in 1:4]
S2_𝐫 = turnorbit.(S2_𝐱, S2_𝐲, Ref(0.0), Ref(S2_Ω), Ref(S2_i))
S2_𝐯 = turnorbit.(S2_𝐯𝐱, S2_𝐯𝐲, Ref(0.0), Ref(S2_Ω), Ref(S2_i))
......@@ -101,19 +97,19 @@ module Models
Dist = R0 / r̂_to_kpc
S2_RA = -((S2_𝐫[:,1] + vx_SgrA*(𝐭_em_S2 .- t_beg)) ./ Dist) .* rad_to_arcsec
S2_Dec = ((S2_𝐫[:,2] + vy_SgrA*(𝐭_em_S2 .- t_beg)) ./ Dist) .* rad_to_arcsec
S2_RV = (-S2_𝐯[:,3] .- M ./ .(S2_𝐱.^2 + S2_𝐲.^2) .+ (S2_𝐯𝐱.^2 .+ S2_𝐯𝐲.^2) ./ 2.0) .* c_to_kms
S2_RA = -((hcat(S2_𝐫...)[1,:] .+ vx_SgrA .* (𝐭_em_S2 .- t_beg)) ./ Dist) .* rad_to_arcsec
S2_Dec = ((hcat(S2_𝐫...)[2,:] .+ vy_SgrA .* (𝐭_em_S2 .- t_beg)) ./ Dist) .* rad_to_arcsec
S2_RV = (-hcat(S2_𝐯...)[3,:] .- M ./ .(S2_𝐱.^2 .+ S2_𝐲.^2) .+ (S2_𝐯𝐱.^2 .+ S2_𝐯𝐲.^2) ./ 2.0) .* c_to_kms
S38_RA = -((S38_𝐫[:,1] + vx_SgrA*(𝐭_em_S38 .- t_beg)) ./ Dist) .* rad_to_arcsec
S38_Dec = ((S38_𝐫[:,2] + vy_SgrA*(𝐭_em_S38 .- t_beg)) ./ Dist) .* rad_to_arcsec
S38_RV = (-S38_𝐯[:,3] .- M / .(S38_𝐱.^2 + S38_𝐲.^2) .+ (S38_𝐯𝐱.^2 .+ S38_𝐯𝐲.^2) ./ 2.0) .* c_to_kms
S38_RA = -((hcat(S38_𝐫...)[1,:] .+ vx_SgrA .* (𝐭_em_S38 .- t_beg)) ./ Dist) .* rad_to_arcsec
S38_Dec = ((hcat(S38_𝐫...)[2,:] .+ vy_SgrA .* (𝐭_em_S38 .- t_beg)) ./ Dist) .* rad_to_arcsec
S38_RV = (-hcat(S38_𝐯...)[3,:] .- M ./ .(S38_𝐱.^2 .+ S38_𝐲.^2) .+ (S38_𝐯𝐱.^2 .+ S38_𝐯𝐲.^2) ./ 2.0) .* c_to_kms
S55_RA = -((S55_𝐫[:, 1] + vx_SgrA*(𝐭_em_S55 .- t_beg)) ./ Dist) .* rad_to_arcsec
S55_Dec = ((S55_𝐫[:, 2] + vy_SgrA*(𝐭_em_S55 .- t_beg)) ./ Dist) .* rad_to_arcsec
S55_RV = (-S55_𝐯[:, 3] .- M ./ .(S55_𝐱.^2 .+ S55_𝐲.^2) .+ (S55_𝐯𝐱.^2 .+ S55_𝐯𝐲.^2) ./ 2.0) .* c_to_kms
S55_RA = -((hcat(S55_𝐫...)[1,:] .+ vx_SgrA .* (𝐭_em_S55 .- t_beg)) ./ Dist) .* rad_to_arcsec
S55_Dec = ((hcat(S55_𝐫...)[2,:] .+ vy_SgrA .* (𝐭_em_S55 .- t_beg)) ./ Dist) .* rad_to_arcsec
S55_RV = (-hcat(S55_𝐯...)[3,:] .- M ./ .(S55_𝐱.^2 .+ S55_𝐲.^2) .+ (S55_𝐯𝐱.^2 .+ S55_𝐯𝐲.^2) ./ 2.0) .* c_to_kms
return (S2_RA, S2_Dec, S2_RV, S38_RA, S38_Dec, S38_RV, S55_RA, S55_Dec, S55_RV)
return [S2_RA, S2_Dec, S2_RV, S38_RA, S38_Dec, S38_RV, S55_RA, S55_Dec, S55_RV]
end
......
__precompile__()
module ParEst
#MCMC sampling module
include("EoMs.jl")
include("Models.jl")
include("Cornerplot.jl")
using .EoMs
using .Models
using .Cornerplot
using AffineInvariantMCMC
export lnprob, MCMCrun
function lnlike(θ, input_data)
t, x, Δx = input_data
χ² = ((x .- model_massless(θ, t)) ./ Δx) .^ 2
return -0.5 * sum(χ²)
end
function lnprior(θ)
S2_x0, S2_y0, S2_vx0, S2_vy0, S2_i, S2_Ω, S38_x0, S38_y0, S38_vx0, S38_vy0, S38_i, S38_Ω, S55_x0, S55_y0, S55_vx0, S55_vy0, S55_i, S55_Ω, M, R0, x_SgrA, y_SgrA, vx_SgrA, vy_SgrA, β, γ = θ
bcond = 0.0 S2_i 180.0 && 0.0 S38_i 180.0 && 0.0 S55_i 180.0 && 0.0 S2_Ω 360.0 && 0.0 S38_Ω 360.0 && 0.0 S55_Ω 360.0 && √sum(S2_vx0^2 + S2_vy0^2) < 1.0 && √sum(S38_vx0^2 + S38_vy0^2) < 1.0 && √sum(S55_vx0^2 + S55_vy0^2) < 1.0 && M > 0 && 0.0 < β < 2.0 && 0.0 < γ < 2.0
if bcond
return 0.0
else
return -Inf
end
end
function lnprob(θ, input_data)
lp = lnprior(θ)
if lp == -Inf
return -Inf
else
return lnlike(θ, input_data)
end
end
function MCMCrun(llhood_par, input_data, nwalkers, niter, initial)
llhood = θ -> llhood_par(θ, input_data)
ndim = length(initial)
θ00 = [initial + initial * 1e-6 .* randn(ndim) for i=1:nwalkers]
θ0 = reshape(hcat(θ00...), (ndim, nwalkers))
println("Running burn-in...")
chain, loglikevals = AffineInvariantMCMC.sample(llhood, nwalkers, θ0, 100, 1)
println("Running production...")
chain, loglikevals = AffineInvariantMCMC.sample(llhood, nwalkers, chain[:, :, end], niter, 1)
flatchain, flatloglikevals = AffineInvariantMCMC.flattenmcmcarray(chain, loglikevals)
θ_max = flatchain[:,argmax(flatloglikevals)]
return flatchain, flatloglikevals, θ_max
end
end
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment