R/expect.R

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

#' Expectation step for mixture of multivariate Gaussians
#'
#' @param X NxD data matrix.
#' @param model Model parameters.
#' @param ms Model selection (boolean).
#' @param indep Independent covariates (boolean).
#' @return Updated model parameters.
#' @export
#'
expectGauss = function(X, model, ms, indep){

    if(!ms){ alpha = model$alpha
    }else{   alpha = model$piw
    }

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

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

    K = length(v)

    logRho = matrix(NA, N, K)

    if(indep){
        for (k in 1:K){
            logRho[,k] = expeCat(alpha, k, ms) + expeGauss(X, D, N, v[k], m[,k], M[,k], beta[k], ms, indep)
        }
    }else{
        logW = rep(NA,K)
        for (k in 1:K){
            ElnPi = expeCat(alpha, k, ms)
            logW[k] = update_logW(M[,,k])
            logRho[,k] = expeCat(alpha, k, ms) + expeGauss(X, D, N, v[k], m[,k], M[,,k], beta[k], ms, indep, logW[k])
        }
        model$logW = logW
    }

    Resp = logRho_to_Resp(logRho)

    model$Resp = Resp

    model
}

#' Expectation step for profile regression with categorical response
#'
#' @param X NxD data matrix.
#' @param y Response vector of length N.
#' @param model Model parameters.
#' @param ms Model selection (boolean).
#' @param indep Independent covariates (boolean).
#' @return Updated model parameters.
#' @export
#'
expectGauss_catResp = function(X, y, model, ms, indep){

    if(!ms){ alpha = model$alpha
    }else{   alpha = model$piw
    }

    v = model$v
    beta = model$beta
    m = model$m
    M = model$M
    eps = model$eps

    N = dim(X)[1]
    D = dim(X)[2]
    K = dim(eps)[2]
    R = dim(eps)[1] # number of categories in each covariate

    logRho = matrix(NA, N, K)
    digammaeps = matrix(NA, N, K)

    if(indep){
        for (k in 1:K){
            ElnPhi =  digamma(eps[y,k]) - digamma(sum(eps[,k]))
            logRho[,k] = expeCat(alpha, k, ms) + expeGauss(X, D, N, v[k], m[,k], M[,k], beta[k], ms, indep) + ElnPhi
        }
    }else{
        logW = rep(NA, K)
        for (k in 1:K){
            ElnPhi =  digamma(eps[y,k]) - digamma(sum(eps[,k]))
            logW[k] = update_logW(M[,,k])
            logRho[,k] = expeCat(alpha, k, ms) + expeGauss(X, D, N, v[k], m[,k], M[,,k], beta[k], ms, indep,  logW[k]) + ElnPhi
        }
        model$logW = logW
    }

    Resp = logRho_to_Resp(logRho)

    model$Resp = Resp
    model
}

#' Expectation step for profile regression with Gaussian response
#'
#'@param X NxD data matrix.
#'@param y Response vector of length N.
#'@param model Model parameters.
#'@param ms Boolean flag that indicates if model selection is required.
#'@param indep Independent covariates (boolean).
#'@export
#'
expectGauss_gaussResp = function(X, y, model, ms, indep){

    if(!ms){ alpha = model$alpha
    }else{   alpha = model$piw
    }

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

       v_y = model$v_y
    beta_y = model$beta_y
       m_y = model$m_y
       M_y = model$M_y

    N = dim(X)[1]
    D = dim(X)[2]
    K = length(v)

    logRho = matrix(NA, N, K)
    logW = rep(NA, K)
    digammaeps = matrix(NA, N, K)

    if(indep){
        for (k in 1:K){
            logRho[,k] =
                        expeCat(alpha, k, ms)
                        + expeGauss(X, D, N, v[k],   m[,k],  M[,k],  beta[k],   ms, indep)
                        + expeGauss(y, 1, N, v_y[k], m_y[k], M_y[k], beta_y[k], ms, indep)
        }
    }else{
        for (k in 1:K){
            logW = rep(NA, K)
            logW[k] = update_logW(M[,,k])
            logRho[,k] =
                        expeCat(alpha, k, ms)
                        + expeGauss(X, D, N, v[k],   m[,k],  M[,k],  beta[k],   ms, indep, logW[,k])
                        + expeGauss(y, 1, N, v_y[k], m_y[k], M_y[k], beta_y[k], ms, indep)
        }
    }

    Resp = logRho_to_Resp(logRho)

    model$Resp = Resp
    model$logW = logW
    model
}

#' Update log(W)
#' @param M = 1/W
#' @export
#'
update_logW = function(M) {
    if(is.matrix(M)){
        logW = log(1/det(M))
    }else{
        logW = log(1/M)
    }
    return(logW)
}

#' Update log(Rho) for independent Gaussians
#' @param X NxD data matrix.
#' @param D Number of covariates.
#' @param N Number of observations.
#' @param v_k Precision hyperparameter.
#' @param m_k Mean hyperparameter.
#' @param M_k Precision hyperparameter.
#' @param beta_k Mean hyperparameter.
#' @param ms Boolean flag that indicates if model selection is required.
#' @param indep Boolean flag that indicates if the covariates are considered to be independent.
#' @param logW_k log(1/det(M)), not needed if indep is TRUE).
#' @export
#'
expeGauss = function(X, D, N, v_k, m_k, M_k, beta_k, ms, indep, logW_k){

    if(D==1){
        # v = model$v_y[k]; m = model$m_y[k]; M = model$M_y[k]; beta = model$beta_y[k]; ms = FALSE
        ElnLa = digamma(0.5*v_k) + log(0.5*M_k) # (10.65) univariate case
        diff = X - m_k
        ExmuLaxmu = 1/beta_k + v_k*(diff^2)/M_k
        logRho = 0.5*ElnLa - 0.5*ExmuLaxmu
    }else if(indep){
        # v = model$v[k]; m = model$m[,k]; M = model$M[,k]; beta = model$beta[k]; ms = FALSE
        ElnLa = D*digamma(0.5*v_k) + sum(log(0.5*M_k)) # (10.65) univariate case ### non sono sicura
        diff = sweep(X, 2, m_k, FUN = "-") ### non sono sicura
        ExmuLaxmu = D/beta_k + v_k*rowSums(sweep(diff^2,MARGIN = 2,1/M_k,`*`)) ### non sono sicura
        logRho = 0.5*ElnLa - 0.5*ExmuLaxmu
    }else{
        # logW_k = logW[k]; v_k = v[k]; m_k = m[,k]; M_k = M[,,k]; beta_k = beta[k]
        ElnLa = D*log(2) + logW_k; for(d in 1:D) ElnLa = ElnLa + digamma(0.5*(v_k+1-d)) # (10.65)
        diff = sweep(X, 2, m_k, FUN="-")
        xmuLaxmu = rep(NA, N); for(n in 1:N) xmuLaxmu[n] = t(diff[n,])%*%solve(M_k,diff[n,]) # TO DO sistemare
        ExmuLaxmu = D/beta_k + v_k*xmuLaxmu # (10.64)
        logRho = 0.5*ElnLa - 0.5*ExmuLaxmu # (10.46)
    }
    return(logRho)
}

#' Expectation of log(pi)
#' @param alpha Dirichlet hyperparameter.
#' @param k Mixture component of interest.
#' @param ms  Boolean flag that indicates if model selection is required.
#' @export
expeCat = function(alpha, k, ms){
    if(!ms){ ElnPi = digamma(alpha[k]) - digamma(sum(alpha)) # (10.66)
    }else{   ElnPi = log(alpha[k])
    }
    ElnPi
}

#' Transform log(Rho) matrix to responsibilities matrix
#' @param logRho log(Rho).
#' @return Matrix of responsibilities.
#' @export
logRho_to_Resp = function(logRho){

    maxLogRho = apply(logRho,1,max)
    logRho = logRho - maxLogRho
    logSumExpLogRho =  log(apply(exp(logRho),1,sum))
    logResp = logRho - logSumExpLogRho
    Resp = exp(logResp)
}