Added MCMC parameter estimation for massless Two-Body problem

code/SubTasks/TwoBody/Cornerplot.jl

+__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

code/SubTasks/TwoBody/DataSets.jl

+__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

code/SubTasks/TwoBody/EoMs.jl

@@ -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

code/SubTasks/TwoBody/Models.jl

@@ -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
 

code/SubTasks/TwoBody/ParEst.jl

+__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