北京大学Julia语言入门讲义第14章： 统计学习介绍

14.2.3.1 随机模拟方法

随机模拟方法包括MCMC，
重要抽样，
序贯重要抽样等方法。
模拟方法可以获得渐近精确的结果，
对模型结构的要求少，
适用于各种复杂的模型。
其中的难点是选择合适的试投分布(proposal distribution)，
这体现了对问题本身的认识，
可以由用户提供，
也可以用数据训练调节。

14.2.3.2 变分推断

变分推断方法用概率质量函数或者概率密度函数逼近目标密度，
一般不能渐近收敛，
但是比随机模拟方法更适合大规模问题。
也可以将随机模拟与变分推断结合起来使用。

14.3.3.1 交叉验证的Julia代码

Julia中许多统计学习（机器学习）方法已经自带了交叉验证功能。
这里摘录McNicholas and Tait(2019)的生成交叉验证数据结构的代码。

# 要设置随机数种子需要调用Random标准库
using Random
# 设置随机数种子使得随机模拟可重复
Random.seed!(35)

# 将1到N的下标，划分为大致相同大小的k个组
# 输出为K个下标数组组成的数组
function cvind(N, k)
    ## 每组的大致点数
    gs = N ÷ k
    ## 将1到N下标随机乱序
    index = shuffle(collect(1:N))
    ## 指定组规模分组，这可能会使得最后一个组仅有几个点
    folds = collect(Iterators.partition(index, gs))
    ## 如果gs*k不等于N，将最后一个点很少的组与倒数第二个组合并
    if length(folds) > k
        folds[k] = vcat(folds[k], folds[k+1])
        deleteat!(folds, k+1)
    end
    return folds
end # function cvind

## 输入数据框dat, cvind函数的输出结果ind, 
## 产生交叉验证的k个自数据框的数组
function kfolds(dat, ind)
    ## 折数
    k = length(ind)
    ## 所有数据的下标
    ind1 = collect(1:size(dat, 1))
    ## 用字典形式返回结果，每一项是编号到字典的映射，
    ## 每一折有留出的用来评估的折以及其余折合并的数据集
    res = Dict{Int, Dict}()
    
    for i = 1:k
        ## 对第i折的产生训练集和测试集
        res2 = Dict{String, DataFrame}()
        ## 去掉第i折作为测试集，其余k-1折皆为训练集
        tr = setdiff(ind1, ind[i])
        res2["training"] = dat[tr, :]
        res2["testing"] = dat[ind[i], :]
        res[i] = res2
    end
    return res
end

using DecisionTree, Random
Random.seed!(35)

## food training data from chapter 3
y = df_food[:gpa]
tmp = convert(Array, df_food[setdiff(names(df_food), [:gpa])] )
xmat = convert(Array{Float64}, collect(Missings.replace(tmp, 0)))
names_food = setdiff(names(df_food), [:gpa])

# defaults to regression tree if y is a Float array
model = build_tree(y, xmat)

# prune tree: merge leaves having >= 85% combined purity (default: 100%)
modelp = prune_tree(model, 0.85)

# tree depth is 9
depth(modelp)

# print the tree to depth 3
print_tree(modelp, 3)
#=
Feature 183, Threshold 0.5
L-> Feature 15, Threshold 0.5
    L-> Feature 209, Threshold 0.5
        L->
        R->
    R-> Feature 209, Threshold 0.5
        L->
        R->
R-> Feature 15, Threshold 0.5
    L-> Feature 8, Threshold 0.5
        L->
        R->
    R-> Feature 23, Threshold 0.5
        L->
        R->
=#

# print the variables in the tree
for i in [183, 15, 209, 8, 23]
  println("Variable number: ", i, " name: ", names_food[i])
end

#=
Variable number: 183 name: self_perception_wgt3
Variable number: 15 name: cal_day3
Variable number: 209 name: waffle_cal1315
Variable number: 8 name: vitamins
Variable number: 23 name: comfoodr_code11
=#

using Statistics, StatsBase

function tunegrid_xgb(eta, maxdepth)
    n_eta = size(eta)[1]
    md = collect(1:maxdepth)
    n_md = size(md)[1]

    ## 2 cols to store the CV results
    res = zeros(*(maxdepth, n_eta), 4)
    n_res = size(res)[1]

    ## populate the res matrix with the tuning parameters
    res[:,1] = repeat(md, inner = Int64(/(n_res, n_md)))
    res[:,2] = repeat(eta, outer = Int64(/(n_res, n_eta)))

    return(res)
end

function binary_mcr(yp, ya; trace = false)
    yl = Array{Int8}(yp .> 0.5)
    disagree = Array{Int8}(ya .!= yl)
    mcr = mean(disagree)
    if trace
        #print(countmap(yl))
        println("yl: ", yl[1:4])
        println("ya: ", ya[1:4])
        println("disagree: ", disagree[1:4])
        println("mcr: ", mcr)
    end
    return( mcr )
end

## 计算二值偏差的辅助函数
function bd_helper(yp, ya)
    e_const = 1e-16
    pneg_const = 1.0 - yp
    if yp < e_const
        res = +(*(ya, log(e_const)), *( -(1.0, ya), log(-(1.0, e_const))))
    elseif pneg_const < e_const
        res = +(*(ya, log(-(1.0, e_const))), *( -(1.0, ya), log(e_const)))
    else
        res = +(*(ya, log(yp)), *( -(1.0, ya), log(pneg_const)))
    end
    return res
end

## Binomial Deviance（二值偏差）
function binomial_dev(yp, ya) #::Vector{T}) where {T <: Number}
    res = map(bd_helper, yp, ya)
    dev = *(-2, sum(res))
    return(dev)
end

function cvERR_xgb(model, Xtst, Ytst, error_fun; trace = false)
    y_p = XGBoost.predict(model, Xtst)
    if(trace)
        println("cvERR: y_p[1:5] ", y_p[1:5])
        println("cvERR: Ytst[1:5] ", Ytst[1:5])
        println("cvERR: error_fun(y_p, Ytst) ", error_fun(y_p, Ytst))
    end
    return(error_fun(y_p, Ytst))
end

function cvResult_xgb(
        Xmat, 
        Yvec, 
        cvparam, 
        tunegrid; 
        k = 5, 
        n = 50,
        trace = false )

    ## tunegrid could have two columns on input
    result = deepcopy(tunegrid) 
    n_tg = size(result)[1]

    ## k-Fold CV for each combination of parameters
    for i = 1:n_tg
        scores = cross_validate(
            trind -> xgboost(
                Xmat[trind,:], 1000, 
                label = Yvec[trind],
                param = cvparam, 
                max_depth = Int(tunegrid[i,1]),
                eta = tunegrid[i,2]),
            (c, trind) -> cvERR_xgb(
                c, Xmat[trind,:], Yvec[trind],
                binomial_dev, trace = true),
            ## total number of samples
            n,
            ## Stratified CV on the outcome
            StratifiedKfold(Yvec, k)
        ) # cross_validate()
        if(trace)
            println("cvResult_xgb: scores: ", scores)
            println("size(scores): ", size(scores))
        end
        result[i, 3] = mean(scores)
        result[i, 4] = std(scores)
    end

    return(result)
end

## 80/20 split
splt = 0.8
y_1_ind = N_array[df_recipe[:y] .== 1]
y_0_ind = N_array[df_recipe[:y] .== 0]

tr_index = sort(
  vcat(
    sample(y_1_ind, Int64(floor(*(N_y_1, splt))), replace = false),
    sample(y_0_ind, Int64(floor(*(N_y_0, splt))), replace = false)
  )
)
tst_index = setdiff(N_array, tr_index)

df_train = df_recipe[tr_index, :]
df_test = df_recipe[tst_index, :]

## Check proportions
println("train: n y: $(sum(df_train[:y]))")
println("train: % y: $(mean(df_train[:y]))")
println("test: n y: $(sum(df_test[:y]))")
println("test: % y: $(mean(df_test[:y]))")

## Static boosting parameters

param_cv = [
  "silent" => 1,
  "verbose" => 0,
  "booster" => "gbtree",
  "objective" => "binary:logistic",
  "save_period" => 0,
  "subsample" => 0.75,
  "colsample_bytree" => 0.75,
  "alpha" => 0,
  "lambda" => 1,
  "gamma" => 0
 ]

## training data: predictors are sparse matrix
tmp = convert(Array, df_train[setdiff(names(df_train), [:y])] )
xmat_tr = convert(SparseMatrixCSC{Float64,Int64},
                      collect(Missings.replace(tmp, 0)))

## CV Results

## eta shrinks the feature weights to help prevent overfilling by making
## the boosting process more conservative
## max_depth is the maximum depth of a tree; increasing this value makes
## the boosting process more complex
tunegrid = tunegrid_xgb([0.1, 0.25, 0.5, 0.75, 1], 5)
N_tr = size(xmat_tr)[1]

cv_res = cvResult_xgb(xmat_tr, y_tr, param_cv, tunegrid,
  k = 5, n = N_tr, trace = true )

## dataframe for plotting
cv_df = DataFrame(tree_depth = cv_res[:, 1],
  eta = cv_res[:, 2],
  mean_error = cv_res[:, 3],
  sd_error = cv_res[:, 4],
)
cv_df[:se_error] = map(x -> x / sqrt(5), cv_df[:sd_error])
cv_df[:mean_min] = map(-, cv_df[:mean_error], cv_df[:se_error])
cv_df[:mean_max] = map(+, cv_df[:mean_error], cv_df[:se_error])

## configurations within 1 se
min_dev = minimum(cv_df[:mean_error])
min_err = filter(row -> row[:mean_error] == min_dev, cv_df)
one_se = min_err[1, :mean_max]
possible_models = filter(row -> row[:mean_error] <= one_se, cv_df)

#####
## Test Set Predictions

tst_results = DataFrame(
 eta = Float64[],
 tree_depth = Int64[],
 mcr = Float64[],
 dev = Float64[])

## using model configuration selected above
pred_tmp = XGBoost.predict(xgboost(xmat_tr, 1000, label = y_tr,
  param = param_cv, eta =0.1 , max_depth = 2), xmat_tst)
tmp = [0.1, 2, binary_mcr(pred_tmp, y_tst), binomial_dev(pred_tmp, y_tst)]
push!(tst_results, tmp)

#####
## Variable importance from the overall data

fit_oa = xgboost(xmat_oa, 1000, label = y_oa, param = param_cv, eta = 0.1,
  max_depth = 2)

## names of features
names_oa =  map(string, setdiff(names(df_recipe), [:y]))
fit_oa_imp = importance(fit_oa, names_oa)

## DF for plotting
N_fi = length(fit_oa_imp)

imp_df = DataFrame(fname = String[], gain = Float64[], cover = Float64[])

for i = 1:N_fi
    tmp = [fit_oa_imp[i].fname, fit_oa_imp[i].gain, fit_oa_imp[i].cover]
    push!(imp_df, tmp)
end

sort!(imp_df, :gain, rev=true)

function tunegrid_xgb_reg(eta, maxdepth, alpha)
    n_eta = size(eta)[1]
    md = collect(1:maxdepth)
    n_md = size(md)[1]
    n_a = size(alpha)[1]

    ## 2 cols to store the CV results
    res = zeros(*(maxdepth, n_eta, n_a), 5)
    n_res = size(res)[1]
    println("N:", n_res, " e: ", n_eta, " m: ", n_md, " a: ", n_a)
    ## populate the res matrix with the tuning parameters
    n_md_i = Int64(/(n_res, n_md))
    res[:,1] = repeat(md, outer=n_md_i)
    res[:,2] = repeat(eta, inner = Int64(/(n_res, n_eta)))
    res[:,3] = repeat(repeach(alpha, (n_md)), outer=n_a)

    return(res)
end

function medae(yp, ya)
    ## element wise operations
    res = map(yp, ya) do x,y
        dif = -(x, y)
        return(abs(dif))
    end
    return(median(res))
end

function rmse(yp, ya)
    ## element wise operations
    res = map(yp, ya) do x,y
        dif = -(x, y)
        return(dif^2)
    end
    return(sqrt(mean(res)))
end

function cvResult_xgb_reg(
        Xmat, Yvec, cvparam, tunegrid; 
        k = 5, n = 50, trace = false )

    result = deepcopy(tunegrid)
    n_tg = size(result)[1]

    ## k-Fold CV for each combination of parameters
    for i = 1:n_tg
        scores = cross_validate(
            ## num_round investigate
            trind -> xgboost(
                Xmat[trind,:], 1000, label = Yvec[trind],
                param = cvparam, max_depth = Int(tunegrid[i,1]),
                eta = tunegrid[i,2], alpha = tunegrid[i,3]),
            (c, trind) -> cvERR_xgb(
                c, Xmat[trind,:], Yvec[trind], medae,
                trace = true),
            n, ## total number of samples
            Kfold(n, k)
        )
        if(trace)
            println("cvResult_xgb_reg: scores: ", scores)
            println("size(scores): ", size(scores))
        end
        result[i, 4] = mean(scores)
        result[i, 5] = std(scores)
    end

    return(result)
end

## 80/20 train test split
splt = 0.8

tr_index = sample(N_array, Int64(floor(*(N, splt))), replace = false)
tst_index = setdiff(N_array, tr_index)

df_train = df_food[tr_index, :]
df_test = df_food[tst_index, :]

## train data: predictors are sparse matrix
y_tr = convert(Array{Float64}, df_train[:gpa])
tmp = convert(Array, df_train[setdiff(names(df_train), [:gpa])] )
xmat_tr = convert(SparseMatrixCSC{Float64,Int64}, 
 collect(Missings.replace(tmp, 0)))

## Static boosting parameters
param_cv_reg = [
 "silent" => 1,
 "verbose" => 0,
 "booster" => "gbtree",
 "objective" => "reg:linear",
 "save_period" => 0,
 "subsample" => 0.75,
 "colsample_bytree" => 0.75,
 "lambda" => 1,
 "gamma" => 0
]

## CV Results

## eta shrinks the feature weights to help prevent overfilling by making
## the boosting process more conservative
## max_depth is the maximum depth of a tree; increasing this value makes
## the boosting process more complex
## alpha is an L1 regularization term on the weights; increasing this value 
## makes the boosting process more conservative
tunegrid = tunegrid_xgb_reg([0.1, 0.3, 0.5], 5, [0.1, 0.3, 0.5])
N_tr = size(xmat_tr)[1]

cv_res_reg = cvResult_xgb_reg(xmat_tr, y_tr, param_cv_reg, tunegrid,
 k = 5, n = N_tr, trace = true )

## add the standard error
cv_res_reg = hcat(cv_res_reg, ./(cv_res_reg[:,5], sqrt(5)))

## dataframe for plotting
cv_df_r = DataFrame(tree_depth = cv_res_reg[:, 1],
  eta = cv_res_reg[:, 2],
  alpha = cv_res_reg[:, 3],
  medae = cv_res_reg[:, 4],
  medae_sd = cv_res_reg[:, 5]
)
cv_df_r[:medae_se] = map(x -> /(x , sqrt(5)), cv_df_r[:medae_sd])
cv_df_r[:medae_min] = map(-, cv_df_r[:medae], cv_df_r[:medae_se])
cv_df_r[:medae_max] = map(+, cv_df_r[:medae], cv_df_r[:medae_se])

min_medae = minimum(cv_df_r[:medae])
min_err = filter(row -> row[:medae] == min_medae, cv_df_r)
one_se = min_err[1, :medae_max]

#######
## Test Set Predictions

tst_results = DataFrame(
  eta = Float64[],
  alpha = Float64[],
  tree_depth = Int64[],
  medae = Float64[],
  rmse = Float64[])

## using configuration chosen above
pred_tmp = XGBoost.predict(xgboost(xmat_tr, 1000, label = y_tr,
  param = param_cv_reg, eta =0.1 , max_depth = 1, alpha = 0.1), xmat_tst)
tmp = [0.1, 0.1, 1, medae(pred_tmp, y_tst), rmse(pred_tmp, y_tst)]
push!(tst_results, tmp)

#####
## Variable Importance from the overall data

fit_oa = xgboost(xmat_oa, 1000, label = y_oa, param = param_cv_reg,
  eta = 0.1, max_depth = 1, alpha = 0.1)

## names of features
# names_oa = convert(Array{String,1}, names(df_food))
names_oa = map(string, setdiff(names(df_train), [:gpa]))
fit_oa_imp = importance(fit_oa, names_oa)

## DF for plotting
N_fi = length(fit_oa_imp)

imp_df = DataFrame(fname = String[], gain = Float64[], cover = Float64[])

for i = 1:N_fi
    tmp = [fit_oa_imp[i].fname, fit_oa_imp[i].gain, fit_oa_imp[i].cover]
    push!(imp_df, tmp)
end

sort!(imp_df, :gain, rev=true)

## EM Algorithm for PPCA -- Helper functions


## 用Woodbury公式求逆
function wbiinv(q, p, psi, L)
    ## Woodbury Indentity eq 6.13
    Lt = transpose(L)
    Iq = Matrix{Float64}(I, q, q)
    Ip = Matrix{Float64}(I, p, p)
    psi2 = /(1, psi)
    m1 = *(psi2, Ip)
    m2c = +( *(psi, Iq), *(Lt, L) )
    m2 = *(psi2, L, inv(m2c), Lt)
    return( -(m1, m2) )
end

## 用Woodbury公式求行列式
function wbidet(q, p, psi, L)
  ## Woodbury Indentity eq 6.14
  Iq = Matrix{Float64}(I, q, q)
  num = psi^p
  den1 = *(transpose(L), wbiinv(q, p, psi, L), L)
  den = det(-(Iq, den1))
  return( /(num, den) )
end

## 似然函数
function ppca_ll(n, q, p, psi, L, S)
    ## log likelihood eq 6.8
    n2 = /(n, 2)
    c = /( *(n, p, log(*(2, pi))), 2)
    l1 = *(n2, log(wbidet(q, p, psi, L)))
    l2 = *(n2, tr(*(S, wbiinv(q, p, psi, L))))
    return( *(-1, +(c, l1, l2)) )
end

## 计算协方差阵
function cov_mat(X)
    n,p = size(X)
    mux = mean(X, dims = 1)
    X_c = .-(X, mux)
    res =  *( /(1,n), *(transpose(X_c), X_c))
    return(res)
end

## 用Aitken加速判断是否收敛
function aitkenaccel(l1, l2, l3, tol)
    ## l1 = l(k+1), l2 = l(k), l3 = l(k-1)
    conv = false
    num = -(l1, l2)
    den = -(l2, l3)
    ## eq 6.16
    ak = /(num, den)
    if ak <= 1.0
        c1 = -(l1, l2)
        c2 = -(1.0, ak)
        ## eq 6.17
        l_inf = +(l2, /(c2, c1))
        c3 = -(l_inf, l2)
        if  0.0 < c3 < tol
            conv = true
        end
    end
    return conv
end

## 计算自由参数个数
function ppca_fparams(p, q)
    ## number of free parameters in the ppca model
    c1 = *(p, q)
    c2 = *(-1, q, -(q, 1), 0.5)
    return( +(c1, c2, 1) )
end

## 计算BIC准则值
function BiC(ll, n, q, ρ)
    ## name does not conflict with StatsBase
    ## eq 6.19
    c1 = *(2, ll)
    c2 = *(ρ, log(n))
    return( -(c1, c2) )
end

using LinearAlgebra, Random
Random.seed!(429)

## EM Function for PPCA
function ppca(X; q = 2, thresh = 1e-5, maxit = 1e5)

     Iq = Matrix{Float64}(I, q, q)
     qI = *(q, Iq)
     n,p = size(X)
     ## eigfact has eval/evec smallest to largest
     ind = p:-1:(p-q+1)
     Sx = cov_mat(X)
     D = eigen(Sx)
     d = diagm(0 => map(sqrt, D.values[ind]))
     P = D.vectors[:, ind]

     ## initialize parameters
     L = *(P, d) ## pxq
     psi = *( /(1, p), tr( -( Sx, *(L, transpose(L)) ) ) )
     B = zeros(q,p)
     T = zeros(q,q)
     conv = true
     iter = 1
     ll = zeros(100)
     ll[1] = 1e4

    ## while not converged
    while(conv)
      ## above eq 6.12
      B = *(transpose(L), wbiinv(q, p, psi, L))  ## qxp
      T = -(qI, +( *(B,L), *(B, Sx, transpose(B)) ) ) ## qxq
      ## sec 6.4.3 - update Lambda_new
      L_new = *(Sx, transpose(B), inv(T)) ## pxq
      ## update psi_new
      psi_new = *( /(1, p), tr( -(Sx, *(L_new, B, Sx) ) )) #num

      iter += 1

      if iter > maxit
          conv = false
          println("ppca() while loop went over maxit parameter. iter = $iter")
      else
        ## stores at most the 100 most recent ll values
        if iter <= 100
          ll[iter]=ppca_ll(n, q, p, psi_new, L_new, Sx)
        else
          ll = circshift(ll, -1)
          ll[100] = ppca_ll(n, q, p, psi_new, L_new, Sx)
        end
        if 2 < iter < 101
              ## scales the threshold by the latest ll value
              thresh2 = *(-1, ll[iter], thresh)
              if aitkenaccel(ll[iter], ll[(iter-1)], ll[(iter-2)], thresh2)
               conv = false
              end
        else
          thresh2 = *(-1, ll[100], thresh)
          if aitkenaccel(ll[100], ll[(99)], ll[(98)], thresh2)
            conv = false
          end
        end
      end ## if maxit
      L = L_new
      psi = psi_new
    end ## while

    ## orthonormal coefficients
    coef = svd(L).U
    ## projections
    proj = *(X, coef)

    if iter <= 100
      resize!(ll, iter)
    end
    fp = ppca_fparams(p, q)

    bic_res = BiC(ll[end], n, q, fp)

    return(Dict(
        "iter" => iter,
        "ll" => ll,
        "beta" => B,
        "theta" => T,
        "lambda" => L,
        "psi" => psi,
        "q" => q,
        "sd" => diag(d),
        "coef" => coef,
        "proj" => proj,
        "bic" => bic_res
      ))
end

ppca1 = ppca(crab_mat_c, q=3, maxit = 1e3)

## MPPCA helper functions
using LinearAlgebra

function sigmaG(mu, xmat, Z)
    res = Dict{Int, Array}()
    N,g = size(Z)
    c1 = ./(1, sum(Z, dims = 1))
    for i = 1:g
        xmu = .-(xmat, transpose(mu[:,i]))
        zxmu =  .*(Z[:,i], xmu)
        res_g = *(c1[i], *(transpose(zxmu), zxmu)) 
        push!(res, i=>res_g)
    end
    return res
end

function muG(g, xmat, Z)
    N, p = size(xmat)
    mu = zeros(p, g)
    for i in 1:g
        num = sum(.*(Z[:,i], xmat), dims = 1)
        den = sum(Z[:, i])
        mu[:, i] = /(num, den)
    end
    return mu
end

## initializations
function init_LambdaG(q, g, Sx)
    res = Dict{Int, Array}()
    for i = 1:g
        p = size(Sx[i], 1)
        ind = p:-1:(p-q+1)
        D = eigen(Sx[i])
        d = diagm(0 => map(sqrt, D.values[ind]))
        P = D.vectors[:, ind]
        L = *(P,d)
        push!(res, i => L)
    end
    return res
end

function init_psi(g, Sx, L)
    res = Dict{Int, Float64}()
    for i = 1:g
        p = size(Sx[i], 1)
        psi = *(/(1, p), tr(-(Sx[i], *(L[i], transpose(L[i])))))
        push!(res, i=>psi)
    end
    return res
end

function init_dict0(g, r, c)
    res = Dict{Int, Array}()
    for i = 1:g
        push!(res, i=> zeros(r, c))
    end
    return res
end

## Updates
function update_B(q, p, psig, Lg)
    res = Dict{Int, Array}()
    g = length(psig)
    for i=1:g
        B = *(transpose(Lg[i]), wbiinv(q, p, psig[i], Lg[i]))
        push!(res, i=>B)
    end
    return res
end

function update_T(q, Bg, Lg, Sg)
    res = Dict{Int, Array}()
    Iq = Matrix{Float64}(I, q, q)
    qI = *(q, Iq)
    g = length(Bg)
    for i =1:g
        T = -(qI, +(*(Bg[i], Lg[i]), *(Bg[i], Sg[i], transpose(Bg[i]))))
        push!(res, i=>T)
    end
    return res
end

function update_L(Sg, Bg, Tg)
    res = Dict{Int, Array}()
    g = length(Bg)
    for i = 1:g
        L = *(Sg[i], transpose(Bg[i]), inv(Tg[i]))
        push!(res, i=>L)
    end
    return res
end

function update_psi(p, Sg, Lg, Bg)
    res = Dict{Int, Float64}()
    g = length(Bg)
    for i = 1:g
        psi = *( /(1, p), tr( -(Sg[i], *(Lg[i], Bg[i], Sg[i]) ) ) )
        push!(res, i=>psi)
    end
    return res
end

function update_zmat(xmat, mug, Lg, psig, pig)
    N,p = size(xmat)
    g = length(Lg)
    res = Matrix{Float64}(undef, N, g)
    Ip = Matrix{Float64}(I, p, p)
    for i = 1:g
        pI = *(psig[i], Ip)
        mu = mug[:, i]
        cov = +( *( Lg[i], transpose(Lg[i]) ), pI)
        pi_den = *(pig[i], pdf(MvNormal(mu, cov), transpose(xmat)))
        res[:,i] = pi_den
    end
    return ./(res, sum(res, dims = 2))
end

function mapz(Z)
    N,g = size(Z)
    res = Vector{Int}(undef, N)
    for i = 1:N
        res[i] = findmax(Z[i,:])[2]
    end
    return res
end

function mppca_ll(N, p, pig, q, psig, Lg, Sg)
    g = length(Lg)
    l1,l2,l3 = (0,0,0)
    c1 = /(N,2)
    c2 = *(-1, c1, p, g, log( *(2,pi) ))
    for i = 1:g
        l1 += log(pig[i])
        l2 += log(wbidet(q, p, psig[i], Lg[i]))
        l3 += tr(*(Sg[i], wbiinv(q, p, psig[i], Lg[i])))
     end
     l1b = *(N, l1)
     l2b = *(-1, c1, l2)
     l3b = *(-1, c1, l3)
     return(+(c2, l1b, l2b, l3b))
end

function mppca_fparams(p, q, g)
    ## number of free parameters in the ppca model
    c1 = *(p, q)
    c2 = *(-1, q, -(q, 1), 0.5)
    return(+( *( +(c1, c2), g), g))
end

function mppca_proj(X, G, map, L)
    res = Dict{Int, Array}()
    for i in 1:G
        coef = svd(L[i]).U
        sel = map .== i
        proj = *(X[sel, :], coef)
        push!(res, i=>proj)
    end
    return(res)
end

using Clustering
include("chp6_ppca_functions.jl")
include("chp6_mixppca_functions.jl")

## MPPCA function
function mixppca(X; q = 2, G = 2, thresh = 1e-5, maxit = 1e5, init = 1))
   ## initializations
   N, p = size(X)
   ## z_ig
   if init == 1
     ## random
     zmat = rand(Uniform(0,1), N, G)
     ## row sum to 1
     zmat = ./(zmat, sum(zmat, dims = 2))
   elseif init == 2
     # k-means
     kma = kmeans(permutedims(X), G; init=:rand).assignments
     zmat = zeros(N,G)
     for i in 1:N
       zmat[i, kma[i]] = 1
     end
   end
   n_g = sum(zmat, dims = 1)
   pi_g = /(n_g, N)
   mu_g = muG(G, X, zmat)
   S_g = sigmaG(mu_g, X, zmat)
   L_g = init_LambdaG(q, G, S_g)
   psi_g = init_psi(G, S_g, L_g)
   B_g = init_dict0(G, q, p)
   T_g = init_dict0(G, q, q)

   conv = true
   iter = 1
   ll = zeros(100)
   ll[1] = 1e4

  # while not converged
  while(conv)

    ## update pi_g and mu_g
    n_g = sum(zmat, dims = 1)
    pi_g = /(n_g, N)
    mu_g = muG(G, X, zmat)
    if iter > 1
      ## update z_ig
      zmat = update_zmat(X, mu_g, L_g, psi_g, pi_g)
    end
    ## compute S_g, Beta_g, Theta_g
    S_g = sigmaG(mu_g, X, zmat)
    B_g = update_B(q, p, psi_g, L_g)
    T_g = update_T(q, B_g, L_g, S_g)

    ## update Lambda_g psi_g
    L_g_new = update_L(S_g, B_g, T_g)
    psi_g_new = update_psi(p, S_g, L_g, B_g)
    ## update z_ig
    zmat = update_zmat(X, mu_g, L_g_new, psi_g_new, pi_g)

    iter += 1

    if iter > maxit
          conv = false
          println("mixppca() while loop went past maxit parameter. iter = 
            $iter")
    else
      ## stores at most the 100 most recent ll values
      if iter <= 100
         ll[iter] = mppca_ll(N, p, n_g, pi_g, q, psi_g, L_g, S_g)
      else
         ll = circshift(ll, -1)
         ll[100] = mppca_ll(N, p, n_g, pi_g, q, psi_g, L_g, S_g)
      end
      if 2 < iter < 101
        ## scales the threshold by the latest ll value
        thresh2 = *(-1, ll[iter], thresh)
        if aitkenaccel(ll[iter], ll[(iter-1)], ll[(iter-2)], thresh2)
           conv = false
        end
      else
        thresh2 = *(-1, ll[100], thresh)
        if aitkenaccel(ll[100], ll[(99)], ll[(98)], thresh2)
           conv = false
        end
      end
    end ## if maxit
    L_g = L_g_new
    psi_g = psi_g_new
  end ## while

  map_res = mapz(zmat)
  proj_res = mppca_proj(X, G, map_res, L_g)

  if iter <= 100
    resize!(ll, iter)
  end
  fp = mppca_fparams(p, q, G)

  bic_res = BiC(ll[end], N, q, fp)

  return Dict(
    "iter" => iter,
    "ll" => ll,
    "beta" => B_g,
    "theta" => T_g,
    "lambda" => L_g,
    "psi" => psi_g,
    "q" => q,
    "G" => G,
    "map" => map_res,
    "zmat" => zmat,
    "proj" => proj_res,
    "bic" => bic_res
  )
end

mixppca_12k = mixppca(cof_mat_c, q=1, G=2, maxit = 1e6, thresh = 1e-3, 
  init=2)