WIP: replace Prior and posterior samples with ParameterDistributions

src/CalibrateEmulateSample.jl

@@ -4,7 +4,6 @@ using Distributions, Statistics, LinearAlgebra, DocStringExtensions
 
 # No internal deps, light external deps
 include("Observations.jl")
-include("Priors.jl")
 include("ParameterDistribution.jl")
 
 # No internal deps, heavy external deps

src/EKP.jl

@@ -1,6 +1,8 @@
 module EKP
 
-using ..Priors
+#using ..Priors
+using ..ParameterDistributionStorage
+
 using Random
 using Statistics
 using Distributions
@@ -46,8 +48,6 @@ $(DocStringExtensions.FIELDS)
 struct EKObj{FT<:AbstractFloat, IT<:Int, P<:Process}
     "vector of arrays of size N_ensemble x N_parameters containing the parameters (in each EK iteration a new array of parameters is added)"
      u::Vector{Array{FT, 2}}
-     "vector of parameter names"
-     unames::Vector{String}
      "vector of observations (length: N_data); mean of all observation samples"
      obs_mean::Vector{FT}
      "covariance of the observational noise, which is assumed to be normally distributed"
@@ -66,7 +66,6 @@ end
 
 # outer constructors
 function EKObj(parameters::Array{FT, 2},
-               parameter_names::Vector{String},
                obs_mean,
                obs_noise_cov::Array{FT, 2},
                process::P;
@@ -85,31 +84,21 @@ function EKObj(parameters::Array{FT, 2},
     # timestep store
     Δt = Array([Δt])
 
-    EKObj{FT, IT, P}(u, parameter_names, obs_mean, obs_noise_cov, N_ens, g,
+    EKObj{FT, IT, P}(u, obs_mean, obs_noise_cov, N_ens, g,
                      err, Δt, process)
 end
 
 
 """
-construct_initial_ensemble(N_ens::IT, priors::Array{Prior, 1}; rng_seed=42) where {IT<:Int}
+construct_initial_ensemble(prior::ParameterDistribution, N_ens::IT; rng_seed=42) where {IT<:Int}
 
 Construct the initial parameters, by sampling N_ens samples from specified
-prior distributions.
+prior distribution. Returned with parameters as rows
 """
-function construct_initial_ensemble(N_ens::IT, priors::Array{Prior, 1}; rng_seed=42) where {IT<:Int}
-    N_params = length(priors)
-    params = zeros(N_ens, N_params)
+function construct_initial_ensemble(prior::ParameterDistribution, N_ens::IT; rng_seed=42) where {IT<:Int}
     # Ensuring reproducibility of the sampled parameter values
     Random.seed!(rng_seed)
-    for i in 1:N_params
-        prior_i = priors[i].dist
-        if !(typeof(priors[i].dist) == Deterministic{Float64})
-            params[:, i] = rand(prior_i, N_ens)
-        else
-            params[:, i] = prior_i.value * ones(N_ens)
-        end
-    end
-
+    params = permutedims(sample_distribution(prior, N_ens), (2,1)) #this transpose is [N_ens x dim(param space)]
     return params
 end
 

src/MCMC.jl

@@ -1,7 +1,8 @@
 module MCMC
 
 using ..GPEmulator
-using ..Priors
+#using ..Priors
+using ..ParameterDistributionStorage
 
 using Statistics
 using Distributions
@@ -34,7 +35,7 @@ struct MCMCObj{FT<:AbstractFloat, IT<:Int}
     "covariance of the observational noise"
     obs_noise_cov::Array{FT, 2}
     "array of length N_parameters with the parameters' prior distributions"
-    prior::Array{Prior, 1}
+    prior::ParameterDistribution
     "MCMC step size"
     step::Array{FT}
     "Number of MCMC steps that are considered burnin"
@@ -69,7 +70,7 @@ where max_iter is the number of MCMC steps to perform (e.g., 100_000)
 function MCMCObj(
     obs_sample::Vector{FT},
     obs_noise_cov::Array{FT, 2},
-    priors::Array{Prior, 1},
+    prior::ParameterDistribution,
     step::FT,
     param_init::Vector{FT},
     max_iter::IT,
@@ -100,7 +101,7 @@ function MCMCObj(
     end
     MCMCObj{FT,IT}(obs_sample,
                    obs_noise_cov,
-                   priors,
+                   prior,
                    [step],
                    burnin,
                    param,
@@ -124,7 +125,14 @@ end
 
 
 function get_posterior(mcmc::MCMCObj)
-    return mcmc.posterior[mcmc.burnin+1:end, :]
+    #Return a parameter distributions object
+    posterior_samples = Samples(mcmc.posterior[mcmc.burnin+1:end,:]; params_are_columns=false) 
+    parameter_constraints = get_all_constraints(mcmc.prior) #live in same space as prior
+    parameter_names = get_name(mcmc.prior) #the same parameters as in prior
+    posterior_distribution = ParameterDistribution(posterior_samples, parameter_constraints, parameter_names)
+    return posterior_distribution
+    #return mcmc.posterior[mcmc.burnin+1:end, :]
+
 end
 
 function mcmc_sample!(mcmc::MCMCObj{FT}, g::Vector{FT}, gvar::Vector{FT}) where {FT}
@@ -174,30 +182,17 @@ end
 
 
 function log_prior(mcmc::MCMCObj{FT}) where {FT}
-    log_rho = FT[0]
-    # Assume independent priors for each parameter
-    priors = [mcmc.prior[i].dist for i in 1:length(mcmc.prior)]
-    for (param, prior_dist) in zip(mcmc.param, priors)
-        # get density at current parameter value
-        log_rho[1] += logpdf(prior_dist, param)
-    end
-
-    return log_rho[1]
+    return get_logpdf(mcmc.prior,mcmc.param)
 end
 
 
 function proposal(mcmc::MCMCObj)
 
-    variances = zeros(length(mcmc.param))
-    priors = [mcmc.prior[i].dist for i in 1:length(mcmc.prior)]
-    param_names = [mcmc.prior[i].param_name for i in 1:length(mcmc.prior)]
-    for (idx, prior) in enumerate(priors)
-        variances[idx] = var(prior)
-    end
-
+    proposal_covariance = get_cov(mcmc.prior)
+
     if mcmc.algtype == "rwm"
         prop_dist = MvNormal(zeros(length(mcmc.param)), 
-                             (mcmc.step[1]^2) * Diagonal(variances))
+                             (mcmc.step[1]^2) * proposal_covariance)
     end
     sample = mcmc.posterior[1 + mcmc.iter[1], :] .+ rand(prop_dist)