From 093ef6afa8eff8b34c01e8fe07f97a706fd9db20 Mon Sep 17 00:00:00 2001
From: Simon Garnier <garnier@njit.edu>
Date: Sat, 13 Jul 2024 17:35:56 +0200
Subject: [PATCH] Change cluster assignment.

---
 R/gravitree.R | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/R/gravitree.R b/R/gravitree.R
index ee29df0..acf44d7 100644
--- a/R/gravitree.R
+++ b/R/gravitree.R
@@ -49,12 +49,12 @@ gravitree <- function(x, m, k = NULL, sample = 1, na_rm = FALSE) {
       ix <- cluster == i
       .wcov(xx[ix, , drop = FALSE], mm[ix])
     })
-    tmp <- rep(NA, nr)
-    tmp[ix] <- cluster
-    tmp[notix] <- Rfast::rowMins(sapply(1:k, function(i) {
-      .Mahalanobis(x[notix, , drop = FALSE], centers[[i]]$center, centers[[i]]$cov)
-    }))
-    cluster <- tmp
+    # tmp <- rep(NA, nr)
+    # tmp[ix] <- cluster
+    # tmp[notix] <- Rfast::rowMins(sapply(1:k, function(i) {
+    #   .Mahalanobis(x[notix, , drop = FALSE], centers[[i]]$center, centers[[i]]$cov)
+    # }))
+    # cluster <- tmp
   } else {
     centers <- lapply(1:k, function(i) {
       ix <- cluster == i
@@ -62,6 +62,10 @@ gravitree <- function(x, m, k = NULL, sample = 1, na_rm = FALSE) {
     })
   }
 
+  cluster <- Rfast::rowMins(sapply(1:k, function(i) {
+    .Mahalanobis(x, centers[[i]]$center, centers[[i]]$cov)
+  }))
+
   list(
     tree = tree,
     cluster = cluster,