PyMC Models to Turing.jl

Specific Domains Probabilistic programming
, ,
5

In general, you can just write Julia code that uses the probabilistic variables as though they were ordinary variables. So mean in the model above could be written similarly, minus the pm.Deterministic and T.as_tensor boilerplate.

I started translating the dark world example, but ended up rewriting it in a way that made more sense to me. Have not checked it for correctness or optimized it, but this at least shows conceptually one way to do it:

using Turing
using LinearAlgebra: norm
using Plots, StatsPlots


# Copied and pasted from the dark sky example linked above.  Columns are
# galaxy x and y positions, and the two components of their eccentricity.
# Transposed to get things column-major.
data = [1.62690000e+02   1.60006000e+03   1.14664000e-01  -1.90326000e-01;
    2.27228000e+03   5.40040000e+02   6.23555000e-01   2.14979000e-01;
    3.55364000e+03   2.69771000e+03   2.83527000e-01  -3.01870000e-01]'
galaxy_positions = data[1:2, :]
galaxy_eccs = data[3:4, :]

f_distance(galaxy_pos, halo_pos, c) = 1 / min(norm(galaxy_pos .- halo_pos), c)

function tangential_distance(galaxy_pos, halo_pos)
    Δ = galaxy_pos .- halo_pos
    t = 2 * atan(Δ[2] / Δ[1])
    return [-cos(t), -sin(t)]
end

function predict_ecc(galaxy_pos, halo_positions, halo_masses, c)
    d = [tangential_distance(galaxy_pos, hp) for hp in eachcol(halo_positions)]
    f = [f_distance(galaxy_pos, hp, c) for hp in eachcol(halo_positions)]
    return(reduce(+, d .* halo_masses .* f))
end
    
@model function DarkSky(galaxy_positions, galaxy_eccs, nhalos=3, c=240)
    n = size(galaxy_positions, 2)
    halo_positions ~ filldist(Uniform(0, 4200), 2, nhalos)
    halo_masses ~ filldist(Uniform(40, 180), nhalos)
    σ2 = 0.05

    # for each galaxy, predict its observed eccentricity and compare it with the data,
    # assuming m.v. normal errors with diagonal covariance
    for i in 1:n
        μ_ecc = predict_ecc(galaxy_positions[:, i], halo_positions, halo_masses, 240)
        galaxy_eccs[:, i] ~ MvNormal(μ_ecc, sqrt(σ2))
    end
end

model = DarkSky(galaxy_positions, galaxy_eccs, 3, 240)
chain = sample(model, NUTS(), 1000)
2 Likes