Rariariari arbitrary transform and weights

R/aggregate.R

@@ -12,6 +12,8 @@
 #' @param left.open logical; used for time intervals, see \link{findInterval} and \link{cut.POSIXt}
 #' @param as_points see \link[stars]{st_as_sf}: shall raster pixels be taken as points, or small square polygons?
 #' @param exact logical; if \code{TRUE}, use \link[exactextractr]{coverage_fraction} to compute exact overlap fractions of polygons with raster cells
+#' @param weights single-layer \code{SpatRaster} or \code{stars} object on the same grid as \code{x} with secondary per-cell weights, e.g. population or cropland area. Only used when \code{exact = TRUE}, and cells with \code{NA} weight are treated as zero weight. Not supported for \code{stars_proxy} objects
+#' @param transform function or one-sided formula, evaluated by \code{rlang::as_function}, applied to cell values before aggregation; e.g. \code{~ .x^2} or \code{~ pmax(0, .x - 10) - pmax(0, .x - 30)}. Units on \code{x} are dropped when set. Note: any transform that derives parameters from the values it is given is applied separately to each attribute (and to each chunk for \code{stars_proxy} objects), so its results may not be comparable across attributes or chunks, e.g. \code{splines::bs(.x)} places knots at quantiles of its input, so fix such parameters explicitly, as in \code{splines::bs(.x, knots = ..., Boundary.knots = ...)}
 #' @seealso \link[sf]{aggregate}, \link[sf]{st_interpolate_aw}, \link{st_extract}, https://github.com/r-spatial/stars/issues/317
 #' @export
 #' @aliases aggregate
@@ -67,9 +69,19 @@
 #' }
 #' agg = aggregate(s, f, mean)
 #' plot(agg)
+#'
+#' # exact = TRUE with secondary weights: population-weighted mean
+#' # population density per municipality
+#' if (requireNamespace("exactextractr", quietly = TRUE) &&
+#'     requireNamespace("terra", quietly = TRUE)) {
+#'  dens = read_stars(system.file("sao_miguel/gpw_v411_2020_density_2020.tif", package = "exactextractr"))
+#'  pop = read_stars(system.file("sao_miguel/gpw_v411_2020_count_2020.tif", package = "exactextractr"))
+#'  conc = sf::read_sf(system.file("sao_miguel/concelhos.gpkg", package = "exactextractr"))
+#'  aggregate(dens, conc, mean, exact = TRUE, weights = pop, na.rm = TRUE)
+#' }
 aggregate.stars = function(x, by, FUN, ..., drop = FALSE, join = st_intersects, 
 		as_points = any(st_dimension(by) == 2, na.rm = TRUE), rightmost.closed = FALSE,
-		left.open = FALSE, exact = FALSE) {
+		left.open = FALSE, exact = FALSE, weights = NULL, transform = NULL) {
 
 	fn_name = substr(deparse1(substitute(FUN)), 1, 20)
 	classes = c("sf", "sfc", "POSIXct", "Date", "character", "function", "stars")
@@ -85,30 +97,79 @@ aggregate.stars = function(x, by, FUN, ..., drop = FALSE, join = st_intersects,
 	} else
 		geom = "geometry"
 	stopifnot(!missing(FUN), is.function(FUN))
+
+	if (!exact && !is.null(weights))
+		warning("for exact=FALSE, weights is ignored")
+	if (!is.null(transform)) {
+		tf = rlang::as_function(transform)
+		# tf may return a vector or a matrix, the matrix branch appends a `term`
+		# dimension to x, the vector branch keeps dim(y). for example splines::bs(...)
+		# or ~ cbind(.x, .x^2, .x^3) returns a matrix while x^2 returns a vector
+		tr = lapply(x, function(y) tf(as.vector(y)))
+		d = st_dimensions(x)
+		if (is.matrix(tr[[1]])) {
+			k = ncol(tr[[1]])
+			if (!all(sapply(tr, function(v) is.matrix(v) && ncol(v) == k)))
+				stop("transform must return a matrix with the same number of columns for every attribute")
+			if (!all(mapply(function(v, y) nrow(v) == length(y), tr, x)))
+				stop("transform must return one row per cell value")
+			cn = colnames(tr[[1]])
+			# any unnamed column makes all names default to t1..tk:
+			term_names = if (is.null(cn) || any(is.na(cn) | !nzchar(cn))) paste0("t", seq_len(k)) else cn
+			d = create_dimensions(append(d, list(term = create_dimension(values = term_names))), attr(d, "raster"))
+			x = st_as_stars(mapply(function(v, y) array(v, dim = c(dim(y), k)), tr, x, SIMPLIFY = FALSE),
+				dimensions = d)
+		} else {
+			if (!all(mapply(function(v, y) length(v) == length(y), tr, x)))
+				stop("transform must return one value per cell value")
+			x = st_as_stars(mapply(function(v, y) array(v, dim = dim(y)), tr, x, SIMPLIFY = FALSE),
+				dimensions = d)
+		}
+	}
 
 	if (exact && inherits(by, c("sfc_POLYGON", "sfc_MULTIPOLYGON")) && has_raster(x)) {
-    	if (!requireNamespace("raster", quietly = TRUE))
-        	stop("package raster required, please install it first") # nocov
+    	if (!requireNamespace("terra", quietly = TRUE))
+        	stop("package terra required, please install it first") # nocov
     	if (!requireNamespace("exactextractr", quietly = TRUE))
         	stop("package exactextractr required, please install it first") # nocov
 		x = st_upfront(x)
 		d = st_dimensions(x)[1:2]
 		r = st_as_stars(list(a = array(1, dim = dim(d))), dimensions = d)
-		e = exactextractr::coverage_fraction(as(r, "Raster"), by)
-		st = do.call(raster::stack, e)
-		m = raster::getValues(st)
-		if (!identical(FUN, sum)) { # see https://github.com/r-spatial/stars/issues/289
-			if (isTRUE(as.character(as.list(FUN)[[3]])[2] == "mean"))
-				m = sweep(m, 2, colSums(m), "/") # mean: divide weights by the sum of weights
-			else
-				stop("for exact=TRUE, FUN should either be mean or sum")
+		template = as(r, "SpatRaster")
+		e = exactextractr::coverage_fraction(template, by)
+		m = terra::values(do.call(c, e))
+		if (!is.null(weights)) {
+			if (inherits(weights, "stars"))
+				weights = as(weights, "SpatRaster")
+			if (!inherits(weights, "SpatRaster"))
+				stop("weights must be a SpatRaster or single-attribute stars object")
+			if (terra::nlyr(weights) > 1)
+				stop("weights must be a single-layer raster")
+			terra::compareGeom(weights, template) # errors if weights does not align with x
+			w = terra::values(weights)[, 1]
+			w[is.na(w)] = 0
+			if (all(w == 0))
+				stop("weights are all zero")
+			m = m * w
 		}
+		is_mean = !identical(FUN, sum) # see https://github.com/r-spatial/stars/issues/289
+		if (is_mean && !isTRUE(as.character(as.list(FUN)[[3]])[2] == "mean"))
+			stop("for exact=TRUE, FUN should either be mean or sum")
+		na.rm = isTRUE(list(...)$na.rm)
 		new_dim = c(prod(dim(x)[1:2]), prod(dim(x)[-(1:2)]))
 		out_dim = c(ncol(m), dim(x)[-(1:2)])
-		if (isTRUE(list(...)$na.rm))
-			x = st_as_stars(lapply(x, function(y) { y[is.na(y)] = 0.0; y }), dimensions = st_dimensions(x))
-		agg = lapply(x, function(a) array(t(m) %*% array(a, dim = new_dim), dim = out_dim))
-		# %*% dropped units, so to propagate units, if present we need to copy (mean/sum):
+		agg = lapply(x, function(a) {
+			v = array(a, dim = new_dim)
+			nas = is.na(v)
+			v[nas] = 0
+			num = crossprod(m, v)
+			if (is_mean) # na.rm also drops NA cells from the denominator:
+				num = num / (if (na.rm) crossprod(m, !nas) else colSums(m))
+			if (!na.rm && any(nas)) # a group is NA when it covers an NA cell:
+				num[crossprod(m, nas) > 0] = NA
+			array(num, dim = out_dim)
+		})
+		# crossprod dropped units, so to propagate units, if present we need to copy (mean/sum):
 		d = create_dimensions(append(setNames(list(create_dimension(values = by)), geom),
 			st_dimensions(x)[-(1:2)]))
 		for (i in seq_along(x)) {
@@ -118,6 +179,8 @@ aggregate.stars = function(x, by, FUN, ..., drop = FALSE, join = st_intersects,
 		}
 		return(st_as_stars(agg, dimensions = d))
 	}
+	if (exact && !is.null(weights))
+		warning("weights is ignored: exact aggregation requires a polygonal `by' and a raster x")
 
 	values = NULL
 	drop_y = FALSE
@@ -272,6 +335,8 @@ aggregate.stars_proxy = function(x, by, FUN, ...) {
 		if (inherits(by, "stars"))
 			by = st_as_sfc(by, as_points = FALSE)
 		by = st_geometry(by)
+		if (isTRUE(list(...)$exact) && !is.null(list(...)$weights))
+			stop("weights is not supported for stars_proxy objects")
 
 		# this assumes each result of a [ selection is small enough to hold in memory
 		l = lapply(seq_along(by), 

R/subset.R

@@ -16,7 +16,7 @@
 #' x[,1:100,100:200,] # select x and y by range
 #' x["L7_ETMs.tif"] # select attribute
 #' xy = structure(list(x = c(293253.999046018, 296400.196497684), y = c(9113801.64775462,
-#' 9111328.49619133)), .Names = c("x", "y"))
+#' 9111328.49619133)), names = c("x", "y"))
 #' pts = st_as_sf(data.frame(do.call(cbind, xy)), coords = c("x", "y"), crs = st_crs(x))
 #' image(x, axes = TRUE)
 #' plot(st_as_sfc(st_bbox(pts)), col = NA, add = TRUE)

tests/aggregate.R

@@ -31,12 +31,28 @@ write_stars(st, tmp)
 sfc = st_set_crs(st_as_sfc(red, as_points = FALSE), st_crs(st))
 (a = aggregate(st, st_sf(a = 1, geom = sfc), mean))
 (a = aggregate(st, sfc, mean))
-if (require(raster)) {
+if (requireNamespace("terra", quietly = TRUE) && requireNamespace("exactextractr", quietly = TRUE)) {
  print(a <- aggregate(st, sfc, mean, exact = TRUE))
  print(a[[1]])
  print(sum(a[[1]])*30 == sum(1:720))
+ # weights:
+ w = st_as_stars(list(w = array(rep(1:5, 24), dim = c(x = 10, y = 12))), dimensions = st_dimensions(st)[1:2])
+ print(all.equal(aggregate(st, sfc, mean, exact = TRUE, weights = w)[[1]][1,1], weighted.mean(st[[1]][1:5,1:6,1], rep(1:5, 6))))
+ print(sum(aggregate(st, sfc, sum, exact = TRUE, weights = w)[[1]][,1]) == sum(st[[1]][,,1] * rep(1:5, 24)))
+ # na.rm: only groups covering an NA cell become NA
+ na_st = st
+ na_st[[1]][1,1,1] = NA
+ a = aggregate(na_st, sfc, mean, exact = TRUE)[[1]]
+ print(is.na(a[1,1]) && !anyNA(a[-1,]) && !anyNA(a[,-1]))
+ print(all.equal(aggregate(na_st, sfc, mean, exact = TRUE, na.rm = TRUE)[[1]][1,1], mean(na_st[[1]][1:5,1:6,1], na.rm = TRUE)))
 }
 
+# transform:
+print(all.equal(aggregate(st, sfc, mean, transform = ~ .x^2)[[1]], aggregate(st^2, sfc, mean)[[1]]))
+a = aggregate(st, sfc, mean, transform = ~ cbind(lin = .x, sq = .x^2))
+print(st_get_dimension_values(a, "term"))
+print(all.equal(unname(a[[1]][,,2]), unname(aggregate(st^2, sfc, mean)[[1]][,])))
+
 tm0 = as.Date("2019-02-19") + -1:8
 (a = aggregate(st, tm0, mean, na.rm = TRUE))
 (a = aggregate(st, "days", mean, na.rm = TRUE))

tests/crop.R

@@ -2,7 +2,7 @@ suppressPackageStartupMessages(library(stars))
 tif = system.file("tif/L7_ETMs.tif", package = "stars")
 x = read_stars(tif)
 xy = structure(list(x = c(293253.999046018, 296400.196497684), y = c(9113801.64775462, 
-9111328.49619133)), .Names = c("x", "y"))
+9111328.49619133)), names = c("x", "y"))
 pts = st_as_sf(data.frame(do.call(cbind, xy)), coords = c("x", "y"), crs = st_crs(x))
 image(x, axes = TRUE)
 plot(st_as_sfc(st_bbox(pts)), col = NA, add = TRUE)