Add function for analysing tree structure of dendrograms

scripts/astrodendro_example_comparisions.py

@@ -3,7 +3,8 @@
 
 from astrodendro.dendrogram import Dendrogram
 from dendro import DistributedDendrogram
-from dendro.utils import compare_dendrograms, get_deviation_level
+from dendro.utils import get_deviation_level
+from dendro.analysis import compare_dendrograms
 
 
 def get_data_and_params():
@@ -37,13 +38,17 @@ def compare_examples(ntasks=4):
         ntasks=ntasks, **params, min_npix_loc=params["min_npix"]
     )
 
-    compare_dendrograms(d_astrodendro, d_exact)
+    overlap_fig, overlap_ax = plt.subplots()
 
     titles = ["astrodrendro", "exact", "intermediate", "fast"]
     for j in [0, 4, 8]:
         for i, dendrogram in enumerate(
             [d_astrodendro, d_exact, d_intermediate, d_fast]
         ):
+            if j == 0 and dendrogram is not d_astrodendro:  # plot overlap
+                overlap = compare_dendrograms(d_astrodendro, dendrogram)
+                overlap_ax.plot(overlap.keys(), overlap.values(), label=f"{titles[i]}")
+
             ax = axs[i]
             ax.imshow(
                 params["data"],
@@ -73,6 +78,13 @@ def compare_examples(ntasks=4):
     fig.tight_layout()
     fig.savefig("compare_astrodendro_example_v3.pdf", dpi=300, bbox_inches="tight")
 
+    overlap_ax.set_xlabel("level")
+    overlap_ax.set_ylabel("overlap")
+    overlap_ax.legend(frameon=False)
+    overlap_fig.savefig(
+        "compare_astrodendro_example_v3_overlap.pdf", dpi=300, bbox_inches="tight"
+    )
+
 
 def iterations_heatmap(ntasks=4, steps_min_delta=11, steps_min_npix=11):
     params = get_data_and_params()

src/dendro/analysis.py

@@ -0,0 +1,72 @@
+import numpy as np
+
+
+def get_overlap_between_index_arrays(indices_a, indices_b):
+    a = np.ascontiguousarray(indices_a)
+    b = np.ascontiguousarray(indices_b)
+    a_view = a.view([("", a.dtype)] * a.shape[1]).ravel()
+    b_view = b.view([("", b.dtype)] * b.shape[1]).ravel()
+
+    mask = np.isin(a_view, b_view)
+    assert np.allclose(mask.shape, (a.shape[0],))  # sanity check
+
+    return mask.sum(), mask.shape[0]
+
+
+def get_overlap_between_structures(a, b):
+    indices_a = np.ascontiguousarray(np.vstack(a.indices(subtree=True)).T)
+    indices_b = np.ascontiguousarray(np.vstack(b.indices(subtree=True)).T)
+    return get_overlap_between_index_arrays(indices_a, indices_b)
+
+
+def find_structure_with_most_overlap_in_layer(structure_a, layer_b):
+    overlaps = [
+        get_overlap_between_structures(structure_a, structure_b)
+        for structure_b in layer_b
+    ]
+    if len(overlaps) == 0:
+        return None, (0, len(structure_a.values(subtree=True)))
+    most_overlap_with = np.argmax([me[0] for me in overlaps])
+    return layer_b[most_overlap_with], overlaps[most_overlap_with]
+
+
+def _recursive_comparison_function(layer_a, layer_b, res=None):
+    if res is None:
+        res = {}
+
+    for structure_a in layer_a:
+        structure_b, overlap = find_structure_with_most_overlap_in_layer(
+            structure_a, layer_b
+        )
+
+        level = structure_a.level
+        if structure_b is not None:
+            assert level == structure_b.level  # sanity check
+        if level not in res.keys():
+            res[level] = []
+        res[level] += [
+            overlap,
+        ]
+
+        if structure_a.children is not None:
+            structure_b_children = (
+                structure_b._children if structure_b is not None else []
+            )
+            _recursive_comparison_function(
+                structure_a.children, structure_b_children, res=res
+            )
+
+    return res
+
+
+def compare_dendrograms(referenence, approximation):
+    all_overlaps = _recursive_comparison_function(
+        referenence.trunk, approximation.trunk
+    )
+    res = {}
+    for level, overlaps in all_overlaps.items():
+        _overlapping = np.array([me[0] for me in overlaps])
+        _num = np.array([me[1] for me in overlaps])
+        mean = _overlapping.sum() / _num.sum()
+        res[level] = mean
+    return res