BHC speedup (memory optimization and some batching)

Author: martinalex000

bhc/core/bhc.py

@@ -48,7 +48,11 @@ def build(self):
         ij = n_objects - 1
 
         # for every pair of data points
+        pair_count = n_objects * (n_objects - 1) // 2
+        tmp_merge = np.empty((pair_count, 5), dtype=float)
+        row = 0
         for i in range(n_objects):
+            log_p_k_row = self.model.row_of_log_likelihood_for_pairs(self.data, i)
             for j in range(i + 1, n_objects):
                 # compute log(d_k)
                 n_ch = n[i] + n[j]
@@ -59,28 +63,22 @@ def build(self):
                 # compute log(pi_k)
                 log_pik = np.log(self.alpha) + gammaln(n_ch) - log_dk
                 # compute log(p_k)
-                data_merged = np.vstack((self.data[i], self.data[j]))
-                log_p_k = self.model.calc_log_mlh(data_merged)
+                log_p_k = log_p_k_row[j - i - 1]  # since j starts at i + 1
                 # compute log(r_k)
                 log_p_ch = log_p[i] + log_p[j]
                 r1 = log_pik + log_p_k
                 r2 = log_d_ch - log_dk + log_p_ch
                 log_r = r1 - r2
                 # store results
-                merge_info = [i, j, log_r, r1, r2]
-                tmp_merge = (
-                    merge_info
-                    if tmp_merge is None
-                    else np.vstack((tmp_merge, merge_info))
-                )
+                tmp_merge[row] = [i, j, log_r, r1, r2]
+                row += 1
 
         # find clusters to merge
         arc_list = np.empty(0, dtype=api.Arc)
+        data_per_cluster = [np.array([self.data[i]]) for i in range(n_objects)]
         while active_nodes.size > 1:
             # find i, j with the highest probability of the merged hypothesis
-            max_log_rk = np.max(tmp_merge[:, 2])
-            ids_matched = np.argwhere(tmp_merge[:, 2] == max_log_rk)
-            position = np.min(ids_matched)
+            position = np.argmax(tmp_merge[:, 2])  # returns the first occurrence
             i, j, log_r, r1, r2 = tmp_merge[position]
             i = int(i)
             j = int(j)
@@ -91,12 +89,6 @@ def build(self):
                 hierarchy_cut = True
                 break
 
-            # turn nodes i,j off
-            tmp_merge[np.argwhere(tmp_merge[:, 0] == i).flatten(), 2] = -np.inf
-            tmp_merge[np.argwhere(tmp_merge[:, 1] == i).flatten(), 2] = -np.inf
-            tmp_merge[np.argwhere(tmp_merge[:, 0] == j).flatten(), 2] = -np.inf
-            tmp_merge[np.argwhere(tmp_merge[:, 1] == j).flatten(), 2] = -np.inf
-
             # new node ij
             ij = n.size
             n_ch = n[i] + n[j]
@@ -107,7 +99,12 @@ def build(self):
                 self.alpha, n[ij], log_d_ch
             )
             log_d = np.append(log_d, log_d_ij)
-            # update assignments
+            # update cluster assignments
+            data_per_cluster.append(
+                np.vstack((data_per_cluster[i], data_per_cluster[j]))
+            )
+            data_per_cluster[i] = None
+            data_per_cluster[j] = None
             assignments[np.argwhere(assignments == i)] = ij
             assignments[np.argwhere(assignments == j)] = ij
 
@@ -121,14 +118,20 @@ def build(self):
             j_idx = np.argwhere(active_nodes == j).flatten()
             active_nodes = np.delete(active_nodes, [i_idx, j_idx])
             active_nodes = np.append(active_nodes, ij)
+
+            # clean up tmp_merge
+            # keep rows where neither column 0 nor column 1 equals i or j
+            mask = ~np.isin(tmp_merge[:, :2], [i, j]).any(axis=1)
+            tmp_merge = tmp_merge[mask]
+
             # compute log(p_ij)
             t1 = np.maximum(r1, r2)
             t2 = np.minimum(r1, r2)
             log_p_ij = t1 + np.log(1 + np.exp(t2 - t1))
             log_p = np.append(log_p, log_p_ij)
 
             # for every pair ij x active
-            x_mat_ij = self.data[np.argwhere(assignments == ij).flatten()]
+            collected_merge_info = np.empty((len(active_nodes) - 1, 5), dtype=float)
             for k in range(active_nodes.size - 1):
                 # compute log(d_k)
                 n_ch = n[k] + n[ij]
@@ -139,18 +142,19 @@ def build(self):
                 # compute log(pi_k)
                 log_pik = np.log(self.alpha) + gammaln(n_ch) - log_dij
                 # compute log(p_k)
-                data_merged = self.data[
-                    np.argwhere(assignments == active_nodes[k]).flatten()
-                ]
-                log_p_ij = self.model.calc_log_mlh(np.vstack((x_mat_ij, data_merged)))
+                data_merged = np.vstack(
+                    (data_per_cluster[ij], data_per_cluster[active_nodes[k]])
+                )
+                log_p_ij = self.model.calc_log_mlh(data_merged)
                 # compute log(r_k)
                 log_p_ch = log_p[ij] + log_p[active_nodes[k]]
                 r1 = log_pik + log_p_ij
                 r2 = log_d_ch - log_dij + log_p_ch
                 log_r = r1 - r2
                 # store results
-                merge_info = [ij, active_nodes[k], log_r, r1, r2]
-                tmp_merge = np.vstack((tmp_merge, merge_info))
+                collected_merge_info[k] = [ij, active_nodes[k], log_r, r1, r2]
+
+            tmp_merge = np.vstack((tmp_merge, collected_merge_info))
 
         return api.Result(
             arc_list,

bhc/core/prior.py

@@ -36,6 +36,53 @@ def calc_log_mlh(self, x_mat):
         log_prior = NormalInverseWishart.__calc_log_prior(s_mat_p, rp, vp)
         return log_prior - self.log_prior0 - LOG2PI * (n * d / 2.0)
 
+    def row_of_log_likelihood_for_pairs(
+        self,
+        X,  # (N, d) data matrix
+        i,  # index of the row you want (int)
+    ):
+        """
+        Returns 1D array containing the log-likelihoods for pairs of points needed for the
+        initialization of bhc. This function combines i with all other points j > i and returns
+        the log-likelihood of those clusters (containing two points each).
+        """
+        N, d = X.shape
+        if d != self.s_mat.shape[0]:
+            raise ValueError("data dimension and prior scale matrix do not match")
+
+        # ------------------------------------------------------------------
+        # Pairwise sufficient statistics – only for j > i (batched)
+        # ------------------------------------------------------------------
+        # slice of points that matter
+        Xj = X[i + 1 :]  # shape (N-i-1, d)
+        diff = X[i] - Xj  # broadcasted automatically
+        x_bar = 0.5 * (X[i] + Xj)  # (N-i-1, d)
+
+        # Scatter matrix S = ½ diff·diffᵀ  → (N-i-1, d, d)
+        S = 0.5 * np.einsum("...i,...j->...ij", diff, diff)
+        # Term (r·2/(r+2))·(x̄‑m)(x̄‑m)ᵀ
+        dt = x_bar - self.m  # (N-i-1, d)
+        outer_dt = np.einsum("...i,...j->...ij", dt, dt)  # (N-i-1, d, d)
+        term = (self.r * 2.0 / (self.r + 2.0)) * outer_dt
+        # Posterior scale matrix for each pair
+        s_mat_p = self.s_mat[None, :, :] + S + term  # (N-i-1, d, d)
+
+        # ------------------------------------------------------------------
+        # Log‑posterior for each pair (batched)
+        # ------------------------------------------------------------------
+        rp = self.r + 2.0  # each cluster has two points
+        vp = self.v + 2.0
+        sign, logdet = slogdet(s_mat_p)  # (N-i-1,)
+        log_prior_post = (
+            LOG2 * (vp * d / 2.0)
+            + (d / 2.0) * np.log(2.0 * np.pi / rp)
+            + multigammaln(vp / 2.0, d)
+            - (vp / 2.0) * logdet
+        )  # (N-i-1,)
+
+        # Final log-likelihood for each pair
+        return log_prior_post - self.log_prior0 - LOG2PI * d  # (N-i-1,)
+
     @staticmethod
     def __calc_log_prior(s_mat, r, v):
         d = s_mat.shape[0]