eloceanografo / marginallogdensities.jl Goto Github PK

using Plots
using MarginalLogDensities
using Distributions
using Random
using Optim

function urchin_V(ω, p, g, a)
    ω = exp(ω)
    p = exp(p)
    g = exp(g)
    am = log(p/(g*ω))/g
    if a < am
        return ω*exp(g*a)
    else
        return p/g + p*(a-am)
    end
end

# number of data
const n = 100

# true parameters

# fixed effects
θ_true = zeros(206)
θ_true[1] = -4.0
θ_true[2] = -0.5
θ_true[3] = log(0.2)
θ_true[4] = 1
θ_true[5] = log(0.1)
θ_true[6] = log(0.1)

# random effects
θ_true[7:106] = rand(Normal(θ_true[2], exp(θ_true[3])), n) #g
θ_true[107:206] = rand(Normal(θ_true[4], exp(θ_true[5])), n) #p

# simulate true data

# urchin ages
a = Float64.(sample(1:30, n))

# urchin volumes
v = urchin_V.(repeat([θ_true[1]],n), θ_true[107:206], θ_true[7:106], a)
v_samp = rand.(Normal.(sqrt.(v), exp(θ_true[6]))) # data + measurement error for likelihood computation

# plot the "data"
scatter(a,v,legend=false)
scatter(a,v_samp,legend=false)

# likelihood function for model with fixed and random effects
function loglik_urchin(θ::Vector{T}) where {T<:Real}

    # parameters
    log_ω = θ[1]
    log_g = θ[7:106]
    log_p = θ[107:206]
    μ_g = θ[2]
    log_σ_g = θ[3]
    μ_p = θ[4]
    log_σ_p = θ[5]
    log_σ = θ[6]

    # estimated volumes of urchins conditional on current θ
    v_est = urchin_V.(repeat([log_ω],n), log_p, log_g, a)

    # data likelihood
    data_lik = sum(logpdf.(Normal.(sqrt.(v_est), exp(log_σ)), v_samp))

    # random effect (growth rates) likelihood
    g_lik = loglikelihood(Normal(μ_g, exp(log_σ_g)), log_g)
    p_lik = loglikelihood(Normal(μ_p, exp(log_σ_p)), log_p)

    return data_lik + g_lik + p_lik
end

loglik_urchin(θ_true)

n_θ = length(θ_true)
marginal_ix = collect(7:206)
marginalloglik_urchin = MarginalLogDensity(loglik_urchin, n_θ, marginal_ix)

# check it works at the true values
marginalloglik_urchin(θ_true[7:end], θ_true[1:6])

# just fixed effects
marginalloglik_urchin(θ_true[1:6])