R/maximize.R

In acabassi/variational-ogc: Variational Outcome-Guided Clustering

#' Maximization step for Gaussian response variable
#'
#' @param y Vector of responses of length N
#' @param model Model parameters
#' @param prior Prior parameters
#' @param ms Model selection (boolean)
#' @return Updated model parameters
#' @export
#'

maximizeGaussResp = function(y, model, prior, ms){

    beta0 = prior$beta_y
    m0 = prior$m_y
    v0 = prior$v_y
    M0 = prior$M_y

    Resp = model$Resp

    K = dim(Resp)[2]

    xbar = rep(0,K)
    M = rep(NA,K)
    m = rep(NA,K)

    Nk = colSums(Resp) # (10.51)
    beta = beta0 + Nk # (10.60)
    v = v0 + Nk  # (10.63)

    for(k in 1:K){
        xbar[k] = update_xbar(y, Resp[,k]) # (10.52)
        M[k] = update_invW(y, Resp[,k], M0, beta0, m0, xbar[k], indep = T) # (10.62)
        m[k] = update_m(beta0, m0, Nk[k], xbar[k], beta[k]) # (10.61)
    }

    model$m_y = m
    model$M_y = M
    model$v_y = v
    model$beta_y = beta
    model

}

#' Maximization step for multivariate Gaussian
#'
#' @param X NxD data matrix
#' @param model Model parameters
#' @param prior Prior parameters
#' @param ms Model selection (boolean)
#' @param indep Independent covariates (boolean)
#' @return Updated model parameters
#' @export
#'

maximizeGauss = function(X, model, prior, ms, indep){

    alpha0 = prior$alpha
    beta0 = prior$beta
    m0 = prior$m
    v0 = prior$v
    M0 = prior$M
    logW0 = prior$logW
    M0 = prior$M
    Resp = model$Resp

    if(is.vector(X)){
        N = length(X)
        D = 1
    }else{
        N = dim(X)[1]
        D = dim(X)[2]
    }
    K = dim(Resp)[2]

    Nk = colSums(Resp) # (10.51)
    beta = beta0 + Nk # (10.60)
    v = v0 + Nk  # (10.63)

    xbar = matrix(0,D,K)
    m = matrix(NA,D,K)

    if(indep){
        M = matrix(NA,D,K)
        for(k in 1:K){
            xbar[,k] = update_xbar(X, Resp[,k]) # (10.52)
            m[,k] = update_m(beta0, m0, Nk[k], xbar[,k], beta[k]) # (10.61)
        }
    }else{
        M = array(NA, c(D,D,K))
        for(k in 1:K){
            xbar[,k] = update_xbar(X, Resp[,k]) # (10.52)
            M[,,k] = update_invW(X, Resp[,k], M0, beta0, m0, xbar[,k], indep) # (10.62)
            m[,k] = update_m(beta0, m0, Nk[k], xbar[,k], beta[k]) # (10.61)
        }
    }

    if(!ms){
        alpha0 = prior$alpha
        alpha = alpha0 + Nk # (10.58)
        model$alpha = alpha
    }else{
        piw = colSums(Resp)/N
        model$piw = piw
    }

    model$m = m
    model$M = M
    model$v = v
    model$beta = beta
    model
}

#' Maximization step for categorical response
#'
#' @param X NxD data matrix
#' @param model Model parameters
#' @param prior Prior parameters
#' @return Updated model parameters
#' @export
#'

maximizeCatResp = function(X, model, prior){

    eps0 = prior$eps
    Resp = model$Resp

    R = dim(eps0)[1]
    K = dim(eps0)[2]

    eps = matrix(NA,R,K)

    for(k in 1:K){
        eps[,k] = update_eps(eps0[,k],Resp[,k], X, R)
    }

    model$eps = eps
    model
}

#' Update xbar
#'
#' @param X NxD data matrix.
#' @param Resp Vector of responsibilities for the mixture component of interest.
#' @export
#'
update_xbar = function (X, Resp){

    Nk = sum(Resp)
    if(is.vector(X)){
        xbarNk = sum(Resp*X)
        xbar = xbarNk/Nk
    }else{
        xbarNk = colSums(sweep(X, MARGIN=1, Resp, `*`))
        xbar = xbarNk/Nk
    }
    if(is.nan(sum(xbar))) stop()
    xbar
}

#' Update M = inv(W)
#'
#' @param X NxD data matrix.
#' @param Resp_k Vector of responsibilities for mixture component k.
#' @param M0_k Prior precision for mixture component k.
#' @param beta0_k Prior beta for mixture component k.
#' @param m0_k Prior m for mixture component k.
#' @param xbar_k Weighted mean for mixture component k.
#' @param indep Independent covariates (boolean).
#' @export
#'
update_invW = function(X, Resp_k, M0_k, beta0_k, m0_k, xbar_k, indep){ ### Questo dovrebbe essere a posto 05/02

    # Resp_k = Resp[,k]; M0_k = M0; beta0_k = beta0; m0_k = m0; xbar_k = xbar[,k]

    Nk = sum(Resp_k)

    if(is.vector(X)){ # Univariate Gaussian
        N = length(X)
        SkNk = 0
        for(n in 1:N) SkNk = SkNk + Resp_k[n]*(X[n]-xbar_k)^2 # (10.53) univariate
        M_k = M0_k + SkNk + (beta0_k*Nk)/(beta0_k+Nk)*(xbar_k-m0_k)^2
    }else if(indep){ # Product of univariate Gaussians
        N = dim(X)[1]
        D = dim(X)[2]
        SkNk = rep(0,D)
        for(n in 1:N) SkNk = SkNk + Resp_k[n]*(X[n,]-xbar_k)^2 # (10.53) univariate
        M_k = M0_k + SkNk + (beta0_k*Nk)/(beta0_k+Nk)*(xbar_k-m0_k)^2
    }else{ # Multivariate Gaussian
        N = dim(X)[1]
        D = dim(X)[2]
        SkNk = matrix(0,D,D)
        for(n in 1:N) SkNk = SkNk + Resp_k[n]*tcrossprod(X[n,]-xbar_k) # (10.53)
        M_k = M0_k + SkNk + (beta0_k*Nk)/(beta0_k+Nk)*tcrossprod(xbar_k-m0_k)

    }
    M_k
}

#' Update mean vector
#'
#' @param beta0 Prior beta.
#' @param m0 Prior m.
#' @param Nk Number of observations in component k.
#' @param xbar Weighted mean of component k.
#' @param beta Parameter for precision of mu.
#' @export
#'
update_m = function(beta0, m0, Nk, xbar, beta){
    m = (beta0*m0 + Nk*xbar)/beta
    if(is.nan(sum(m))) stop()
    m
}

#' Update espilon
#'
#' @param eps0 Prior eps.
#' @param Resp Responsibilities.
#' @param y Vector of responses.
#' @param R number of categories.
#' @export
#'
update_eps = function(eps0, Resp, y, R){
    # TO DO check if y contains R consecutive integers starting from 1
    ntilde = rep(NA, R)
    for(r in 1:R) ntilde[r] = sum(Resp*(y==r))
    eps = eps0 + ntilde
}