Merge pull request #31 from stscl/dev

update `hclustgeo_disc()`

R/stratification.R

@@ -59,10 +59,12 @@ discretize_vector = \(x, n, method = 'natural',
 #' please see `stats::hclust()`.
 #' @param scale (optional) Whether to scaled the dissimilarities matrix, default is `TRUE`.
 #' @param wt (optional) Vector with the weights of the observations. By default, `wt` is `NULL`.
+#' @param cut (optional) Whether to cut the `hclust` tree, default is `TRUE`.
 #' @param ... (optional) Other arguments passed to `stats::dist()`.
 #'
-#' @return A `vector` with grouped memberships if `n` are `scalar`, otherwise a `matrix` with grouped
-#' memberships is returned where each column corresponds to the elements of `n`, respectively.
+#' @return When `cut` is `TRUE`, returns a grouped membership: a `vector` if `n` is a scalar,
+#' a `matrix` (columns correspond to elements of `n`) if not; otherwise, returns a `vector`
+#' of the permuted original observations.
 #' @export
 #'
 #' @examples
@@ -72,7 +74,7 @@ discretize_vector = \(x, n, method = 'natural',
 #'
 hclustgeo_disc = \(data, n, alpha = 0.5, D1 = NULL,
                    hclustm = "ward.D2", scale = TRUE,
-                   wt = NULL, ...){
+                   wt = NULL, cut = TRUE, ...){
   if (inherits(data,"sf")) {
     if (alpha != 0 & is.null(D1)) {
       D1 = sdsfun::sf_distance_matrix(data)
@@ -85,6 +87,11 @@ hclustgeo_disc = \(data, n, alpha = 0.5, D1 = NULL,
   }
   D0 = as.matrix(stats::dist(data,...))
   deltadist = stats::as.dist(RcppHClustGeoMat(D0,D1,alpha,scale,wt))
-  resh = stats::hclust(deltadist,method=hclustm,members=wt)
-  return(stats::cutree(resh,k = n))
+  resh = stats::hclust(deltadist,method = hclustm,members = wt)
+
+  if (cut) {
+    return(stats::cutree(resh,k = n))
+  } else {
+    return(resh$order)
+  }
 }