|
| 1 | +#' @title Check if the torch smoother is available |
| 2 | +#' @name .torch_smooth_available |
| 3 | +#' @keywords internal |
| 4 | +#' @noRd |
| 5 | +#' @description Use torch when it is functional unless |
| 6 | +#' SITS_SMOOTH_FORCE_CPP is set to TRUE. |
| 7 | +#' |
| 8 | +#' @return A logical value |
| 9 | +.torch_smooth_available <- function() { |
| 10 | + force_cpp <- Sys.getenv("SITS_SMOOTH_FORCE_CPP", unset = "FALSE") |
| 11 | + if (toupper(force_cpp) == "TRUE") { |
| 12 | + return(FALSE) |
| 13 | + } |
| 14 | + .torch_is_functional() |
| 15 | +} |
| 16 | + |
| 17 | +#' @title Estimate memory used by torch smoothing windows |
| 18 | +#' @name .torch_smooth_block_memsize |
| 19 | +#' @keywords internal |
| 20 | +#' @noRd |
| 21 | +#' @description Estimate the memory used to materialize the neighborhood |
| 22 | +#' tensors for one raster block. Bands are processed sequentially and their |
| 23 | +#' intermediate tensors are released between iterations, so the dominant |
| 24 | +#' window allocation does not grow with the number of bands. |
| 25 | +#' |
| 26 | +#' @param block_size Number of pixels in the block, including overlap. |
| 27 | +#' @param window_size Side length of the square neighborhood. |
| 28 | +#' |
| 29 | +#' @return Estimated memory in GB. |
| 30 | +.torch_smooth_block_memsize <- function(block_size, window_size) { |
| 31 | + .jobs_block_memsize( |
| 32 | + block_size = block_size, |
| 33 | + npaths = window_size^2, |
| 34 | + nbytes = 4L, |
| 35 | + proc_bloat = .conf("processing_bloat_smooth_torch") |
| 36 | + ) |
| 37 | +} |
| 38 | + |
| 39 | +#' @title Torch-backed Bayesian smoother for probability cubes |
| 40 | +#' @name .torch_smooth_bayes_fraction |
| 41 | +#' @keywords internal |
| 42 | +#' @noRd |
| 43 | +#' @author Alexandre Assuncao, \email{alexcarssuncao@@gmail.com} |
| 44 | +#' |
| 45 | +#' @description |
| 46 | +#' Drop-in replacement for the C++ \code{bayes_smoother_fraction()} function. |
| 47 | +#' This is the default raster smoothing backend whenever torch and its native |
| 48 | +#' dependencies are functional. It uses torch tensor operations on CPU and |
| 49 | +#' runs them on a GPU when CUDA or MPS is available. |
| 50 | +#' |
| 51 | +#' Padding note: PyTorch \code{"reflect"} mode mirrors from the element |
| 52 | +#' adjacent to the boundary (boundary pixel is not repeated), whereas the |
| 53 | +#' C++ implementation repeats the boundary pixel once. The mismatch only |
| 54 | +#' affects the outermost \code{leg = floor(window_size / 2)} rows/columns of |
| 55 | +#' each processed block; those pixels are discarded by the overlap/crop_block |
| 56 | +#' mechanism in \code{.smooth_tile()}, so the final output is unaffected for |
| 57 | +#' all interior blocks. |
| 58 | +#' |
| 59 | +#' @param logits Numeric matrix \code{[nrows*ncols, nbands]} of |
| 60 | +#' logit-transformed probabilities. |
| 61 | +#' @param nrows Number of rows in the block. |
| 62 | +#' @param ncols Number of columns in the block. |
| 63 | +#' @param window_size Side length of the square neighbourhood |
| 64 | +#' (odd integer). |
| 65 | +#' @param smoothness Numeric vector of length \code{nbands}: the prior |
| 66 | +#' variance parameter per class. |
| 67 | +#' @param neigh_fraction Fraction of neighbourhood values (highest logits) |
| 68 | +#' used to estimate the prior. |
| 69 | +#' @return Numeric matrix \code{[nrows*ncols, nbands]} of smoothed logits. |
| 70 | +#' |
| 71 | +.torch_smooth_bayes_fraction <- function(logits, |
| 72 | + nrows, |
| 73 | + ncols, |
| 74 | + window_size, |
| 75 | + smoothness, |
| 76 | + neigh_fraction) { |
| 77 | + # Select device |
| 78 | + device <- if (.torch_gpu_available()) { |
| 79 | + if (torch::cuda_is_available()) { |
| 80 | + "cuda" |
| 81 | + } else { |
| 82 | + "mps" |
| 83 | + } |
| 84 | + } else { |
| 85 | + "cpu" |
| 86 | + } |
| 87 | + |
| 88 | + # Number of bands in data |
| 89 | + nbands <- ncol(logits) |
| 90 | + # Floor of window size over two |
| 91 | + leg <- window_size %/% 2L |
| 92 | + # Number of pixels in window |
| 93 | + win_sq <- as.integer(window_size * window_size) |
| 94 | + |
| 95 | + # Build input tensor [1, nbands, nrows, ncols] |
| 96 | + x <- torch::torch_tensor( |
| 97 | + t(logits), |
| 98 | + dtype = torch::torch_float32(), |
| 99 | + device = device |
| 100 | + )$view(c(nbands, nrows, ncols))$unsqueeze(1L) |
| 101 | + |
| 102 | + # Pad spatial dimensions with reflect mode. The resulting shape is |
| 103 | + # [1, nbands, nrows + 2 * leg, ncols + 2 * leg]. |
| 104 | + # NOTE: Due to overlapping tiles, the padded values are later discarded |
| 105 | + x_pad <- torch::nnf_pad( |
| 106 | + x, |
| 107 | + pad = c(leg, leg, leg, leg), |
| 108 | + mode = "reflect" |
| 109 | + ) |
| 110 | + |
| 111 | + # Aux zero tensor |
| 112 | + zero_t <- torch::torch_tensor( |
| 113 | + 0.0, |
| 114 | + dtype = torch::torch_float32(), |
| 115 | + device = device |
| 116 | + ) |
| 117 | + # Aux minus inf tensor |
| 118 | + neginf_t <- torch::torch_tensor( |
| 119 | + -Inf, |
| 120 | + dtype = torch::torch_float32(), |
| 121 | + device = device |
| 122 | + ) |
| 123 | + # Aux index tensor for neighbor selection [1, win_sq] |
| 124 | + pos <- torch::torch_tensor( |
| 125 | + seq_len(win_sq), |
| 126 | + dtype = torch::torch_float32(), |
| 127 | + device = device |
| 128 | + )$view(c(1L, win_sq)) |
| 129 | + |
| 130 | + # Original pixel values for all bands [npix, nbands] |
| 131 | + x0_all <- torch::torch_tensor( |
| 132 | + logits, |
| 133 | + dtype = torch::torch_float32(), |
| 134 | + device = device |
| 135 | + ) |
| 136 | + |
| 137 | + # Process each band separately to avoid exhausting GPU memory |
| 138 | + band_results <- purrr::map(seq_len(nbands), function(b) { |
| 139 | + # Release intermediate tensors from the previous band |
| 140 | + invisible(gc(full = TRUE)) |
| 141 | + # Get band b from padded tensor, i.e. [1, 1, H_pad, W_pad] |
| 142 | + x_b <- x_pad[, b, , , drop = FALSE] |
| 143 | + # Unfold windows, i.e. [1, win_sq, npix] to [npix, win_sq] |
| 144 | + wins_b <- torch::nnf_unfold( |
| 145 | + x_b, |
| 146 | + kernel_size = window_size, |
| 147 | + stride = 1L, |
| 148 | + padding = 0L |
| 149 | + )$squeeze(1L)$t() |
| 150 | + # Dealing with NA |
| 151 | + nan_mask_b <- torch::torch_isnan(wins_b) |
| 152 | + # Select neighbourhood values |
| 153 | + if (neigh_fraction == 1.0) { |
| 154 | + selected_b <- torch::torch_where(nan_mask_b, zero_t, wins_b) |
| 155 | + sel_mask_b <- !nan_mask_b |
| 156 | + n_b <- (!nan_mask_b)$sum(dim = -1L)$to( |
| 157 | + dtype = torch::torch_float32() |
| 158 | + ) |
| 159 | + } else { |
| 160 | + wins_sort_b <- torch::torch_where(nan_mask_b, neginf_t, wins_b) |
| 161 | + sorted_b <- torch::torch_sort( |
| 162 | + wins_sort_b, |
| 163 | + dim = -1L, |
| 164 | + descending = TRUE |
| 165 | + )[[1L]] |
| 166 | + |
| 167 | + valid_b <- (!nan_mask_b)$sum(dim = -1L)$to( |
| 168 | + dtype = torch::torch_float32() |
| 169 | + ) |
| 170 | + neigh_hi_b <- torch::torch_ceil( |
| 171 | + neigh_fraction * valid_b |
| 172 | + )$clamp_min(1L) |
| 173 | + |
| 174 | + sel_mask_b <- pos$le(neigh_hi_b$unsqueeze(-1L)) |
| 175 | + selected_b <- torch::torch_where(sel_mask_b, sorted_b, zero_t) |
| 176 | + n_b <- neigh_hi_b |
| 177 | + } |
| 178 | + # Calculating empirical mean |
| 179 | + m0_b <- selected_b$sum(dim = -1L) / n_b |
| 180 | + # Calculating empirical variance |
| 181 | + diff_sq_b <- torch::torch_where( |
| 182 | + sel_mask_b, |
| 183 | + (selected_b - m0_b$unsqueeze(-1L))$pow(2L), |
| 184 | + zero_t |
| 185 | + ) |
| 186 | + s0_b <- torch::torch_where( |
| 187 | + n_b$gt(1.0), |
| 188 | + diff_sq_b$sum(dim = -1L) / (n_b - 1.0), |
| 189 | + zero_t |
| 190 | + ) |
| 191 | + # Bayesian update for band b |
| 192 | + x0_b <- x0_all[, b] |
| 193 | + w_b <- s0_b / (s0_b + smoothness[b]) |
| 194 | + bayes_b <- w_b * x0_b + (1.0 - w_b) * m0_b |
| 195 | + use_m0 <- torch::torch_isnan(x0_b) | s0_b$lt(1e-4) |
| 196 | + torch::torch_where(use_m0, m0_b, bayes_b) |
| 197 | + }) |
| 198 | + # Release intermediate tensors from the last band |
| 199 | + invisible(gc(full = TRUE)) |
| 200 | + |
| 201 | + # Stack all bands and convert to matrix with dim [npix, nbands] |
| 202 | + result <- torch::torch_stack(band_results, dim = 2L) |
| 203 | + as.matrix( |
| 204 | + result$cpu()$to(dtype = torch::torch_float64()) |
| 205 | + ) |
| 206 | +} |
| 207 | + |
| 208 | +#' @title Torch Bayesian smoother closure factory |
| 209 | +#' @name .smooth_fn_bayes_torch |
| 210 | +#' @keywords internal |
| 211 | +#' @noRd |
| 212 | +#' @author Alexandre Assuncao, \email{alexcarssuncao@@gmail.com} |
| 213 | +#' |
| 214 | +#' @description |
| 215 | +#' Mirrors \code{.smooth_fn_bayes()} from \file{api_smooth.R} but calls the |
| 216 | +#' torch-backed \code{.torch_smooth_bayes_fraction()} instead of the C++ |
| 217 | +#' function. It applies the same logit and inverse-logit transforms on CPU, |
| 218 | +#' CUDA, or MPS. |
| 219 | +#' |
| 220 | +#' @param window_size Size of the neighbourhood |
| 221 | +#' (odd integer, min 5, max 21). |
| 222 | +#' @param neigh_fraction Fraction of highest-valued neighbours to use. |
| 223 | +#' @param smoothness Numeric vector (length = number of classes). |
| 224 | +#' @return A closure \code{function(values, block)} compatible with |
| 225 | +#' \code{.smooth_tile()}. |
| 226 | +#' |
| 227 | +.smooth_fn_bayes_torch <- function(window_size, |
| 228 | + neigh_fraction, |
| 229 | + smoothness) { |
| 230 | + # Check window size |
| 231 | + .check_int_parameter(window_size, min = 5L, max = 21L, is_odd = TRUE) |
| 232 | + # Check neigh_fraction |
| 233 | + .check_num_parameter(neigh_fraction, min = 0.1, max = 1.0) |
| 234 | + |
| 235 | + smooth_fn <- function(values, block) { |
| 236 | + # Record expected number of pixels |
| 237 | + input_pixels <- nrow(values) |
| 238 | + # Clamp to avoid ±Inf in logit |
| 239 | + values[values == 1.0] <- 0.999999 |
| 240 | + values[values == 0.0] <- 0.000001 |
| 241 | + # Transform to logits |
| 242 | + values <- log(values / (rowSums(values) - values)) |
| 243 | + # Apply torch Bayesian smoother |
| 244 | + values <- .torch_smooth_bayes_fraction( |
| 245 | + logits = values, |
| 246 | + nrows = .nrows(block), |
| 247 | + ncols = .ncols(block), |
| 248 | + window_size = window_size, |
| 249 | + smoothness = smoothness, |
| 250 | + neigh_fraction = neigh_fraction |
| 251 | + ) |
| 252 | + # Inverse logit |
| 253 | + values <- exp(values) / (exp(values) + 1.0) |
| 254 | + # Sanity check |
| 255 | + .check_processed_values(values, input_pixels) |
| 256 | + values |
| 257 | + } |
| 258 | + smooth_fn |
| 259 | +} |
0 commit comments