R/plotting-ellipse.R

In lgatto/pRoloc: A unifying bioinformatics framework for spatial proteomics

##' A function to plot probabiltiy ellipses on marker PCA plots to
##' visualise and assess TAGM models.
##'
##' Note that when running PCA, this function does not scale the data
##' (centring is performed), as opposed to [plot2D()]. Only marker
##' proteins are displayed; the protein of unknown location, that are
##' not used to estimate the MAP parameters, are filtered out.
##'
##' @param object An [`MSnbase::MSnset`] containing quantitative
##'     spatial proteomics data.
##' @param params An [`MAPParams`] with the TAGM-MAP parameters, as
##'     generated by `tagmMapTrain`.
##' @param dims A `numeric(2)` with the principal components along
##'     which to project the data. Default is `c(1, 2)`.
##' @param method The method used. Currently `"MAP"` only.
##' @param ... Additional parameters passed to [plot2D()].
##' @return A PCA plot of the marker data with probability
##'     ellipises. The outer ellipse contains 99% of the total
##'     probability whilst the middle and inner ellipses contain 95%
##'     and 90% of the probability respectively. The centres of the
##'     clusters are represented by black circumpunct (circled dot).
##'
##' @seealso [plot2D()] to visualise spatial proteomics data using
##'     various dimensionality reduction methods. For details about
##'     TAGM models, see [tagmPredict()] and the *pRoloc-bayesian*
##'     vignette.
plotEllipse <- function(object,
                         params,
                         dims = c(1,2),
                         method = "MAP",
                         ...) {
    ## make marker MSnSet
    markersubset <- markerMSnSet(object)
    markers <- getMarkerClasses(object)
    mydata <- exprs(markersubset)
    K <- length(markers)

    D <- ncol(mydata)
    N <- nrow(mydata)

    muk <- params@posteriors$mu
    sigmak <- params@posteriors$sigma

    ## scale data correctly and format
    sigmaROT <- array(0, c(K, D, D))
    pca <- prcomp(mydata, center = TRUE, scale = FALSE)
    plot2D(markersubset, dims = dims, methargs = list(scale = FALSE, center= TRUE), ...)
    meanROT <- scale(muk, center = pca$center, scale = pca$scale) %*% pca$rotation
    points(meanROT[, dims], pch = 1, cex = 1.4, col = 'black')
    points(meanROT[, dims], pch = 19, cex = 0.8, col = 'black')
    for (j in 1:K) {
        sigmaROT[j, , ] <- scale(sigmak[j, , ], scale=pca$scale, center = FALSE)
    }

    for (j in 1:K) {
        myCol <-  getStockcol()[j]
        mixtools::ellipse(meanROT[j, dims], sigmaROT[j, dims, dims], alpha = .1,  npoints = 250, newplot = FALSE, col = myCol)
        mixtools::ellipse(meanROT[j, dims], sigmaROT[j, dims, dims], alpha = .01, npoints = 250, newplot = FALSE, col = myCol)
        mixtools::ellipse(meanROT[j, dims], sigmaROT[j, dims, dims], alpha = .05, npoints = 250, newplot = FALSE, col = myCol)
    }

}



## This function is not exported and kept for backwards compatibility
## (for instance for the code in the tagm paper). It is not meant to
## be used anymore and will be removed at some point.
plotEllipse_old <- function(object,
                            components = c(1,2),
                            mu0 = NULL,
                            lambda0 = 0.01,
                            nu0 = NULL,
                            S0 = NULL,
                            method = c("MAP", "MCMC"),
                            ...) {

    method <- match.arg(method)

    ## make marker MSnSet
    markersubset <- markerMSnSet(object)
    marker.names <- getMarkerClasses(object)
    mydata <- exprs(markersubset)
    K <- length(marker.names)

    ## create dataset
    D <- ncol(mydata)
    N <- nrow(mydata)

    if (is.null(nu0)) {
        nu0 <- D + 2
    }
    if (is.null(S0)) {
        S0 <- diag( colSums(( mydata - mean( mydata)) ^ 2) / N)/( K ^ (1/D))
    }
    if(is.null(mu0)){
        mu0 <- colMeans( mydata)
    }


    ## create storage for posterior parameters
    mk <- matrix(0, nrow = K, ncol = D)
    lambdak <- matrix(0, nrow = K, ncol = 1)
    nuk <- matrix(0, nrow = K, ncol = 1)
    sk <- array(0, dim = c(K, D, D))

    ## create storage for cluster parameters
    muk <- matrix(0, nrow = K, ncol = D)
    sigmak <- array(0, dim = c(K, D, D))
    xk <- matrix(0, nrow = K, ncol = D)

    ## update prior with training data
    nk <- tabulate(fData(markersubset)[, "markers"])
    for(j in 1:K)
        xk[j, ] <- colMeans(mydata[fData(markersubset)[, "markers"] == getMarkerClasses(markersubset)[j], ])

    lambdak <- lambda0 + nk
    nuk <- nu0 + nk
    mk <- (nk * xk + lambda0 * mu0) / lambdak

    for(j in 1:K) {
        sk[j, , ] <- S0 + t(mydata[fData(markersubset)[,"markers"] == getMarkerClasses(markersubset)[j], ]) %*%
            mydata[fData(markersubset)[, "markers"] == getMarkerClasses(markersubset)[j],] +
            lambda0 * mu0 %*% t(mu0) - lambdak[j] * mk[j, ] %*% t(mk[j, ])
    }
    ## initial posterior mode
    muk <- mk
    if (method == "MAP") {
        for (j in 1:K) sigmak[j, , ] <- sk[j, , ] / (nuk[j] + D + 1)
    } else { ## method == "MCMC"
        for (j in 1:K) sigmak[j, , ] <- sk[j, , ] / (nuk[j] - D - 1)
    }


    ## scale data correctly and format
    sigmaROT <- array(0, c(K, D, D))
    par(mfrow = c(1,1))
    pca <- prcomp(mydata, center = TRUE, scale = FALSE)
    plot2D(markersubset, dims = components, methargs = list(scale = FALSE, center= TRUE),...)
    meanROT <- scale(muk, center = pca$center, scale = pca$scale) %*% pca$rotation
    points(meanROT[, components], pch=19, col='black')
    for (j in 1:K)
        sigmaROT[j, , ] <- scale(sigmak[j, , ], scale=pca$scale, center = FALSE)


    ## ellipses from Mixtools package
    for(j in 1:K){
        mixtools::ellipse(meanROT[j,components],sigmaROT[j,components,components], alpha = .1, npoints=250, newplot = FALSE)
        mixtools::ellipse(meanROT[j,components],sigmaROT[j,components,components], alpha = .01, npoints=250, newplot = FALSE)
        mixtools::ellipse(meanROT[j,components],sigmaROT[j,components,components], alpha = .05, npoints=250, newplot = FALSE)
    }

}