R/linear_stretch.R

In azvoleff/teamlucc: TEAM land use and cover change data processing toolkit

# Note only one of pct or n_sd should be specified
.do_stretch <- function(x, pct, n_sd, max_val) {
    stopifnot(missing(n_sd) | missing(pct))
    if (!missing(pct)) {
        lower <- quantile(x, prob=0 + pct/100, na.rm=TRUE)
        upper <- quantile(x, prob=1 - pct/100, na.rm=TRUE)
        x[x > upper] <- upper
        x[x < lower] <- lower
        x <- ((x - lower) / (upper-lower)) * max_val
    } else {
        x_m <- mean(x, na.rm=TRUE)
        x_sd <- sd(x, na.rm=TRUE)
        lower <- x_m - n_sd * x_sd
        upper <- x_m + n_sd * x_sd
        x[x > upper] <- upper
        x[x < lower] <- lower
        x <- ((x - lower) / (upper-lower)) * max_val
    }
    return(x)
}

#' Apply a linear stretch to an image
#'
#' Applies a linear stretch to an image (default linear 2% stretch), and 
#' returns the image with each band individually stretched and rescaled to 
#' range between zero and \code{max_val} (default of \code{max_val} is 1).
#'
#' Note only one of \code{pct} or \code{n_sd} can be specified.
#'
#' @export
#' @import raster
#' @param x image to stretch
#' @param pct percent stretch (defaults to 2 if not specified)
#' @param n_sd number of standard deviations for stretch
#' @param max_val maximum value of final output (image will be rescaled to 
#' range from 0 - \code{max_val})
#' @return image with stretch applied
linear_stretch <- function(x, pct, n_sd, max_val=1) {
    if (!missing(pct) & !missing(n_sd)) {
        error("Only one of pct and n_sd should be specified")
    }
    if (missing(pct) & missing(n_sd)) {
        warning("pct and n_sd not specified - defaulting to 2% linear stretch")
        pct <- 2
    }
    # Applies linear stretch (2 percent by default). Assumes image is arranged 
    # with bands in columns. Returns the image with stretch applied and bands 
    # rescaled to range from 0 - max_val.
    if (!missing(pct)) {
        if ((pct <= 0) | (pct >= 50)) {
            stop('pct must be > 0 and < 50')
        }
    }
    if (!missing(n_sd)) {
        if ((n_sd <= 0) | (n_sd >= 5)) {
            stop('n_sd must be > 0 and < 5')
        }
    }
    if (class(x) %in% c('RasterLayer')) {
        x <- setValues(x, .do_stretch(getValues(x), pct, n_sd, max_val))
        return(x)
    } else if (class(x) %in% c('RasterStack', 'RasterBrick')) {
        for (n in 1:nlayers(x)) {
            x <- setValues(x, .do_stretch(getValues(raster(x, layer=n)),
                                          pct, n_sd, max_val), layer=n)
        }
        return(x)
    } else if (is.null(dim(x)) || (length(dim(x)) == 2) || (dim(x)[3] == 1)) {
        # Handle a vector, matrix, or n x m x 1 array
        return(.do_stretch(x, pct, n_sd, max_val))
    }
}