torch smoothing

Author: alexcarssuncao

DESCRIPTION

@@ -315,3 +315,4 @@ Collate:
     'sits_xlsx.R'
     'zzz.R'
 Config/roxygen2/version: 8.1.0
+RoxygenNote: 7.3.2

NAMESPACE

@@ -650,9 +650,7 @@ export(sits_view)
 export(sits_xgboost)
 importFrom(Rcpp,sourceCpp)
 importFrom(dplyr,.data)
-importFrom(lubridate,
-  "%m+%",
-  "%within%"
-)
+importFrom(lubridate,"%m+%")
+importFrom(lubridate,"%within%")
 importFrom(utils,read.csv)
 useDynLib(sits, .registration = TRUE)

R/api_bstorch.R

@@ -2,7 +2,7 @@
 #' @name .torch_smooth_bayes_fraction
 #' @keywords internal
 #' @noRd
-#' @author Gilberto Camara, \email{gilberto.camara@@inpe.br}
+#' @author Alexandre Assuncao, \email{alexcarssuncao@@gmail.com}
 #'
 #' @description
 #' Drop-in replacement for the C++ \code{bayes_smoother_fraction()} function.
@@ -34,7 +34,7 @@
                                          window_size,
                                          smoothness,
                                          neigh_fraction) {
-    # Select compute device
+    # Select device
     device <- if (torch::cuda_is_available()) {
         "cuda"
     } else if (torch::backends_mps_is_available()) {
@@ -43,93 +43,75 @@
         "cpu"
     }
 
-    nbands  <- ncol(logits)
-    npix    <- nrows * ncols
-    leg     <- window_size %/% 2L
-    win_sq  <- as.integer(window_size * window_size)
+    nbands  <- ncol(logits)                          # Number of bands in data
+    npix    <- nrows * ncols                         # Number of pixels in data
+    leg     <- window_size %/% 2L                    # Floor of window size over two
+    win_sq  <- as.integer(window_size * window_size) # Number of pixels in window
 
-    # ---- Build input tensor [1, nbands, nrows, ncols] ----
+    # Build input tensor [1, nbands, nrows, ncols]
     x <- torch::torch_tensor(
         t(logits),
         dtype  = torch::torch_float32(),
         device = device
     )$view(c(nbands, nrows, ncols))$unsqueeze(1L)
 
-    # ---- Reflect-pad spatial dimensions (all bands at once — cheap) ----
+    # Padding spatial dimensions with reflect [1, nbands, nrows+2*leg, ncols+2*leg]
+    # NOTE: Due to overlapping tiles, the padded values are later discarded
     x_pad <- torch::nnf_pad(
         x,
         pad  = c(leg, leg, leg, leg),
         mode = "reflect"
     )
-    # x_pad: [1, nbands, nrows+2*leg, ncols+2*leg]
 
-    # ---- Reusable scalar tensors and position index ----
-    zero_t   <- torch::torch_tensor(
-        0.0, dtype = torch::torch_float32(), device = device
-    )
-    neginf_t <- torch::torch_tensor(
-        -Inf, dtype = torch::torch_float32(), device = device
-    )
-    # Position index for top-fraction selection [1, win_sq]
+    # Aux zero tensor
+    zero_t   <- torch::torch_tensor(0.0, dtype = torch::torch_float32(), device = device)
+    # Aux minus inf tensor
+    neginf_t <- torch::torch_tensor(-Inf, dtype = torch::torch_float32(), device = device)
+    # Aux index tensor for neighbor selection [1, win_sq]
     pos <- torch::torch_tensor(
         seq_len(win_sq),
         dtype  = torch::torch_float32(),
         device = device
     )$view(c(1L, win_sq))
 
-    # ---- Original pixel values for all bands [npix, nbands] ----
-    x0_all <- torch::torch_tensor(
-        logits, dtype = torch::torch_float32(), device = device
-    )
+    # Original pixel values for all bands [npix, nbands]
+    x0_all <- torch::torch_tensor(logits, dtype = torch::torch_float32(), device = device)
 
-    # ---- Process one band at a time to limit GPU memory ----
-    # Without this, nnf_unfold on all bands produces a
-    # [nbands, npix, win_sq] tensor that can easily exceed GPU memory.
+    # Storage for results per-band (Avoid blowing up GPU memory)
     band_results <- vector("list", nbands)
+
+    # --- MAIN LOOP
     for (b in seq_len(nbands)) {
-        # Extract single band from padded tensor: [1, 1, H_pad, W_pad]
+        # Get band b from padded tensor, i.e. [1, 1, H_pad, W_pad]
         x_b <- x_pad[, b, , , drop = FALSE]
-
-        # Unfold windows: [1, win_sq, npix] → [npix, win_sq]
+        # Unfold windows, i.e. [1, win_sq, npix] to [npix, win_sq]
         wins_b <- torch::nnf_unfold(
             x_b,
             kernel_size = window_size,
             stride      = 1L,
             padding     = 0L
         )$squeeze(1L)$t()
-        # wins_b: [npix, win_sq]
-
-        # ---- NaN handling ----
+        # Dealing with NA
         nan_mask_b <- torch::torch_isnan(wins_b)
-
-        # ---- Select neighbourhood values ----
+        # Select neighbourhood values
         if (neigh_fraction == 1.0) {
             selected_b <- torch::torch_where(nan_mask_b, zero_t, wins_b)
             sel_mask_b <- !nan_mask_b
-            n_b        <- (!nan_mask_b)$sum(dim = -1L)$to(
-                dtype = torch::torch_float32()
-            )
+            n_b        <- (!nan_mask_b)$sum(dim = -1L)$to(dtype = torch::torch_float32())
         } else {
             wins_sort_b <- torch::torch_where(nan_mask_b, neginf_t, wins_b)
-            sorted_b    <- torch::torch_sort(
-                wins_sort_b, dim = -1L, descending = TRUE
-            )[[1L]]
+            sorted_b    <- torch::torch_sort(wins_sort_b, dim = -1L, descending = TRUE)[[1L]]
 
-            valid_b    <- (!nan_mask_b)$sum(dim = -1L)$to(
-                dtype = torch::torch_float32()
-            )
-            neigh_hi_b <- torch::torch_ceil(
-                neigh_fraction * valid_b
-            )$clamp_min(1L)
+            valid_b    <- (!nan_mask_b)$sum(dim = -1L)$to(dtype = torch::torch_float32())
+            neigh_hi_b <- torch::torch_ceil(neigh_fraction * valid_b)$clamp_min(1L)
 
             sel_mask_b <- pos$le(neigh_hi_b$unsqueeze(-1L))
             selected_b <- torch::torch_where(sel_mask_b, sorted_b, zero_t)
             n_b        <- neigh_hi_b
         }
-
-        # ---- Unbiased mean and variance ----
+        # Calculating empirical mean
         m0_b <- selected_b$sum(dim = -1L) / n_b
-
+        # Calculating empirical variance
         diff_sq_b <- torch::torch_where(
             sel_mask_b,
             (selected_b - m0_b$unsqueeze(-1L))$pow(2L),
@@ -140,25 +122,24 @@
             diff_sq_b$sum(dim = -1L) / (n_b - 1.0),
             zero_t
         )
-
-        # ---- Bayesian update for this band ----
-        x0_b   <- x0_all[, b]
-        w_b    <- s0_b / (s0_b + smoothness[b])
+        # Bayesian update for band b
+        x0_b    <- x0_all[, b]
+        w_b     <- s0_b / (s0_b + smoothness[b])
         bayes_b <- w_b * x0_b + (1.0 - w_b) * m0_b
         use_m0  <- torch::torch_isnan(x0_b) | s0_b$lt(1e-4)
-
+        # Store results
         band_results[[b]] <- torch::torch_where(use_m0, m0_b, bayes_b)
     }
 
-    # ---- Stack bands and convert to R matrix [npix, nbands] ----
+    # Stack all bands and convert to matrix with dim [npix, nbands]
     result <- torch::torch_stack(band_results, dim = 2L)
     as.matrix(
         result$cpu()$to(dtype = torch::torch_float64())
     )
 }
 
 #' @title Torch Bayesian smoother closure factory
-#' @name .smooth_fn_bayes_torch
+#' @name  .smooth_fn_bayes_torch
 #' @keywords internal
 #' @noRd
 #'
@@ -240,12 +221,16 @@
                           output_dir,
                           version,
                           progress) {
+    # Define Bayesian Smoothing function
     smooth_fn <- .smooth_fn_bayes_torch(
         window_size    = window_size,
         neigh_fraction = neigh_fraction,
         smoothness     = smoothness
     )
+    # The overlap makes sure that the padded values added for torch
+    # computations are discarded
     overlap <- ceiling(window_size / 2L) - 1L
+    # Use sits cube api to iterate over tiles
     .cube_foreach_tile(cube, function(tile) {
         .smooth_tile(
             tile           = tile,

R/sits_bstorch.R

@@ -2,8 +2,7 @@
 #'
 #' @name sits_smooth_torch
 #'
-#' @author Gilberto Camara, \email{gilberto.camara@@inpe.br}
-#' @author Rolf Simoes, \email{rolfsimoes@@gmail.com}
+#' @author Alexandre Assuncao, \email{alexcarssuncao@@gmail.com}
 #'
 #' @description
 #' Torch-backed alternative to \code{\link[sits]{sits_smooth}}.  Applies the
@@ -131,8 +130,8 @@ sits_smooth_torch.probs_cube <- function(cube,
         memsize           = memsize,
         multicores        = multicores
     )
-    if (.parallel_start(workers = multicores))
-        on.exit(.parallel_stop(), add = TRUE)
+    started <- .parallel_start(workers = multicores)
+    on.exit(.parallel_stop(started), add = TRUE)
 
     # Call torch-backed orchestrator
     .smooth_torch(