compositional-explanations 0.0.0.1__tar.gz

compositional_explanations-0.0.0.1/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2018 The Python Packaging Authority
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

compositional_explanations-0.0.0.1/PKG-INFO

@@ -0,0 +1,86 @@
+Metadata-Version: 2.1
+Name: compositional_explanations
+Version: 0.0.0.1
+Summary: A small example package
+Author-email: Biagio La Rosa <bilarosa@ucsc.edu>
+License: MIT
+Project-URL: Homepage, https://github.com/aiea-lab/optimal-compositional-explanations
+Project-URL: Issues, https://github.com/aiea-lab/optimal-compositional-explanations/issues
+Platform: UNKNOWN
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.22.2
+Requires-Dist: scipy>=1.10.1
+Requires-Dist: torch>=1.14.0
+
+
+# Compositional Explanations Package
+This package provides functions to compute compositional explanations for a given bitmap and a set of boolean masks. The package includes two main methods for computing explanations: optimal and beam search. Additionally, it provides metrics to evaluate the quality of the explanations.
+
+The package assumes you are able to provide the following inputs:
+- `bitmaps`: A boolean tensor representing the bitmap to be explained for a given neuron. The bitmap represents the activation of the neuron for a set of inputs. A cell is 1 if the neuron is activated for that feature or whole input, and 0 otherwise.
+- `masks`: A list of boolean tensors representing the masks for each feature. Each mask represents whether a given concept/feature is annotated in the given position or samples. A cell is 1 if the feature is present in that position or sample, and 0 otherwise.
+- `disjoint_info`: A boolean tensor representing the disjointness of the masks. The tensor is of shape (num_masks, num_masks) and indicates whether two masks are disjoint (i.e., they do not overlap). A cell is 1 if the two masks are disjoint, and 0 otherwise.
+- `length`: An integer representing the maximum length of the explanation formula. The formula is a combination of masks that best explains the bitmap.
+- `beam_size`: An integer representing the beam size for the beam search method. This parameter controls the number of candidate explanations to consider at each step of the search.
+- `device`: A string representing the device to run the computations on. Highly suggested to use a GPU for faster computations.
+- `cache_dir`: A string representing the directory to cache heuristic informaation for the given masks. This is useful for large datasets where computing the heuristics can be time-consuming. Default is None, which means no caching will be used.
+
+Useful functions include:
+- `optimal.compute_optimal_explanations(bitmaps, masks, disjoint_info, length, device, cache_dir)`: Computes the optimal explanation formula for the given bitmap and masks. This method can be slow (~1 hour per neuron) for high complexity scenarios.
+- `beam.compute_beam_explanations(bitmaps, masks, disjoint_info, length, beam_size, device, cache_dir)`: Computes the explanation formula for the given bitmap and masks using a beam search method. This method is faster than the optimal method but cannot guarantee the optimality of the explanation. It is recommended to use this method for large datasets or when a faster computation is needed.
+- `metrics.compute_iou_from_masks(formula, masks, bitmaps)`: Computes the Intersection over Union (IoU) metric for the given explanation formula, masks, and bitmap. The IoU metric measures the overlap between the explanation masks and the bitmap, providing a measure of how well the explanation captures the bitmap's activation.
+
+## Example usage of the compositional_explanations package
+
+```python
+from compositional_explanations import beam, optimal, metrics
+import torch
+
+# Create 5 boolean random masks of shape (4, 4)
+masks = [torch.randint(0, 2, (4, 4), dtype=torch.bool) for _ in range(5)]
+
+concept_index = 0
+for mask in masks:
+    print(f"Mask {concept_index}:")
+    print(mask)
+    concept_index += 1
+
+# Create a random bitmap of shape (4, 4)
+bitmap = torch.randint(0, 2, (4, 4), dtype=torch.bool)
+
+print(f"Bitmap:")
+print(bitmap)
+
+# Compute disjoint matrix (inefficent, we suggest to compute them directly from the annotations)
+disjoint_matrix = torch.ones((len(masks), len(masks)), dtype=torch.bool)
+for i in range(len(masks)):
+    for j in range(i + 1, len(masks)):
+        disjoint_matrix[i, j] = not torch.any(masks[i] & masks[j])
+        disjoint_matrix[j, i] = disjoint_matrix[i, j]
+
+print("Disjoint Matrix:")
+print(disjoint_matrix)
+# Compute optimal explanations
+best_formula_optimal = optimal.compute_optimal_explanations(
+    bitmaps=bitmap, masks=masks, disjoint_info=disjoint_matrix, length=3, device=torch.device("cuda"), cache_dir=None
+)
+optimal_iou = metrics.compute_iou_from_masks(formula=best_formula_optimal, masks=masks, bitmaps=bitmap)
+
+print("Optimal Explanation:")
+print("Best formula:", best_formula_optimal)
+print("Best IoU:", optimal_iou)
+
+best_formula_beam = beam.compute_beam_explanations(
+    bitmaps=bitmap, masks=masks, disjoint_info=disjoint_matrix, length=3, beam_size=5, device=torch.device("cuda"), cache_dir=None
+)   
+beam_iou = metrics.compute_iou_from_masks(formula=best_formula_beam, masks=masks, bitmaps=bitmap)
+print("Beam Explanation:")
+print("Best formula:", best_formula_beam)
+print("Best IoU:", beam_iou)
+print()
+print("Beam Explanation is optimal:", beam_iou == optimal_iou)
+```

compositional_explanations-0.0.0.1/README.md

@@ -0,0 +1,68 @@
+
+# Compositional Explanations Package
+This package provides functions to compute compositional explanations for a given bitmap and a set of boolean masks. The package includes two main methods for computing explanations: optimal and beam search. Additionally, it provides metrics to evaluate the quality of the explanations.
+
+The package assumes you are able to provide the following inputs:
+- `bitmaps`: A boolean tensor representing the bitmap to be explained for a given neuron. The bitmap represents the activation of the neuron for a set of inputs. A cell is 1 if the neuron is activated for that feature or whole input, and 0 otherwise.
+- `masks`: A list of boolean tensors representing the masks for each feature. Each mask represents whether a given concept/feature is annotated in the given position or samples. A cell is 1 if the feature is present in that position or sample, and 0 otherwise.
+- `disjoint_info`: A boolean tensor representing the disjointness of the masks. The tensor is of shape (num_masks, num_masks) and indicates whether two masks are disjoint (i.e., they do not overlap). A cell is 1 if the two masks are disjoint, and 0 otherwise.
+- `length`: An integer representing the maximum length of the explanation formula. The formula is a combination of masks that best explains the bitmap.
+- `beam_size`: An integer representing the beam size for the beam search method. This parameter controls the number of candidate explanations to consider at each step of the search.
+- `device`: A string representing the device to run the computations on. Highly suggested to use a GPU for faster computations.
+- `cache_dir`: A string representing the directory to cache heuristic informaation for the given masks. This is useful for large datasets where computing the heuristics can be time-consuming. Default is None, which means no caching will be used.
+
+Useful functions include:
+- `optimal.compute_optimal_explanations(bitmaps, masks, disjoint_info, length, device, cache_dir)`: Computes the optimal explanation formula for the given bitmap and masks. This method can be slow (~1 hour per neuron) for high complexity scenarios.
+- `beam.compute_beam_explanations(bitmaps, masks, disjoint_info, length, beam_size, device, cache_dir)`: Computes the explanation formula for the given bitmap and masks using a beam search method. This method is faster than the optimal method but cannot guarantee the optimality of the explanation. It is recommended to use this method for large datasets or when a faster computation is needed.
+- `metrics.compute_iou_from_masks(formula, masks, bitmaps)`: Computes the Intersection over Union (IoU) metric for the given explanation formula, masks, and bitmap. The IoU metric measures the overlap between the explanation masks and the bitmap, providing a measure of how well the explanation captures the bitmap's activation.
+
+## Example usage of the compositional_explanations package
+
+```python
+from compositional_explanations import beam, optimal, metrics
+import torch
+
+# Create 5 boolean random masks of shape (4, 4)
+masks = [torch.randint(0, 2, (4, 4), dtype=torch.bool) for _ in range(5)]
+
+concept_index = 0
+for mask in masks:
+    print(f"Mask {concept_index}:")
+    print(mask)
+    concept_index += 1
+
+# Create a random bitmap of shape (4, 4)
+bitmap = torch.randint(0, 2, (4, 4), dtype=torch.bool)
+
+print(f"Bitmap:")
+print(bitmap)
+
+# Compute disjoint matrix (inefficent, we suggest to compute them directly from the annotations)
+disjoint_matrix = torch.ones((len(masks), len(masks)), dtype=torch.bool)
+for i in range(len(masks)):
+    for j in range(i + 1, len(masks)):
+        disjoint_matrix[i, j] = not torch.any(masks[i] & masks[j])
+        disjoint_matrix[j, i] = disjoint_matrix[i, j]
+
+print("Disjoint Matrix:")
+print(disjoint_matrix)
+# Compute optimal explanations
+best_formula_optimal = optimal.compute_optimal_explanations(
+    bitmaps=bitmap, masks=masks, disjoint_info=disjoint_matrix, length=3, device=torch.device("cuda"), cache_dir=None
+)
+optimal_iou = metrics.compute_iou_from_masks(formula=best_formula_optimal, masks=masks, bitmaps=bitmap)
+
+print("Optimal Explanation:")
+print("Best formula:", best_formula_optimal)
+print("Best IoU:", optimal_iou)
+
+best_formula_beam = beam.compute_beam_explanations(
+    bitmaps=bitmap, masks=masks, disjoint_info=disjoint_matrix, length=3, beam_size=5, device=torch.device("cuda"), cache_dir=None
+)   
+beam_iou = metrics.compute_iou_from_masks(formula=best_formula_beam, masks=masks, bitmaps=bitmap)
+print("Beam Explanation:")
+print("Best formula:", best_formula_beam)
+print("Best IoU:", beam_iou)
+print()
+print("Beam Explanation is optimal:", beam_iou == optimal_iou)
+```

compositional_explanations-0.0.0.1/pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "compositional_explanations"
+version = "0.0.0.1"
+authors = [
+  { name="Biagio La Rosa", email="bilarosa@ucsc.edu" },
+]
+description = "A small example package"
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+
+dependencies = [
+    "numpy>=1.22.2",
+    "scipy>=1.10.1",
+    "torch>=1.14.0",
+]
+
+license = {text = "MIT"}
+
+[project.urls]
+Homepage = "https://github.com/aiea-lab/optimal-compositional-explanations"
+Issues = "https://github.com/aiea-lab/optimal-compositional-explanations/issues"

compositional_explanations-0.0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

compositional_explanations-0.0.0.1/src/compositional_explanations/beam.py

@@ -0,0 +1,405 @@
+"""Beam-search variant guided by the optimal heuristic family."""
+
+from collections import Counter
+import queue as Q
+
+import torch
+
+from compositional_explanations import formula as F
+from compositional_explanations import optimal_sample_heuristic, metrics
+from utils import optimal_utils, mask_utils
+
+def beam_search(
+    search_space,
+    *,
+    masks,
+    beam_masks,
+    bitmaps,
+    beam_limit,
+    previous_beam=None,
+):
+    """Perform the beam search on the search space.
+
+    Args:
+        search_space (list): A list of formulas.
+        masks (dict): A dictionary of concept masks. Each mask is a tensor of
+            shape (N, H, W).
+        beam_masks (dict): A dictionary of cached formula masks for formulas in
+            the current beam.
+        bitmaps (torch.Tensor): A tensor of shape (N, H, W) where N is the
+            number of samples.
+        beam_limit (int): The beam size.
+        previous_beam (dict): A dictionary mapping beam formulas to their IoU.
+
+    Returns:
+        current_beam_formulas (list): A list of formulas.
+        current_beam_iou (list): A list of IoU values corresponding to
+            `current_beam_formulas`.
+        visited_indices (int): The number of formulas whose exact IoU was evaluated.
+    """
+    if previous_beam is None:
+        previous_beam = {}
+    current_beam = Q.PriorityQueue(beam_limit)
+    current_beam_iou = []
+    current_beam_formulas = []
+    minimum = 0
+    visited_indices = 0
+    best_formula = None
+    # Init beam with previous best
+    for label, iou in previous_beam.items():
+        if not current_beam.full():
+            current_beam.put((iou, label))
+            minimum = current_beam.queue[0][0]
+        elif iou > minimum:
+            current_beam.get()
+            current_beam.put((iou, label))
+            minimum = current_beam.queue[0][0]
+
+    # Set minimum
+    if current_beam.empty():
+        minimum = 0
+    else:
+        minimum = current_beam.queue[0][0]
+
+    # Iterate over the search space
+    for candidate_formula in search_space:
+        e_iou = candidate_formula.iou
+
+        # If the estimated IoU is less than the minimum, we can stop the search if the beam is full
+        if current_beam.full() and e_iou < minimum:
+            break
+
+        # skip equivalent formulas of the current beam
+        if best_formula and hash(candidate_formula) == hash(best_formula):
+            continue
+
+        # Compute IoU
+        masks_formula = mask_utils.get_formula_mask(
+            candidate_formula, masks, beam_masks, device=bitmaps.device
+        )
+        iou = metrics.iou(masks_formula, bitmaps)
+
+        # Update visited nodes
+        visited_indices += 1
+
+        # Update beam
+        if not current_beam.full():
+            candidate_formula.iou = iou
+            current_beam.put((iou, candidate_formula))
+            minimum = current_beam.queue[0][0]
+        elif iou > minimum:
+            candidate_formula.iou = iou
+            current_beam.get()
+            current_beam.put((iou, candidate_formula))
+            minimum = current_beam.queue[0][0]
+
+    # Extract formulas and iou from the beam
+    for _ in range(current_beam.qsize()):
+        iou, candidate = current_beam.get()
+        current_beam_formulas.append(candidate)
+        current_beam_iou.append(iou)
+    return current_beam_formulas, current_beam_iou, visited_indices
+
+
+def compute_next_search_space(formulas, candidate_labels):
+    """Compute the next search space starting from the current beam
+    of formulas.
+
+    Args:
+        formulas (list): A list of formulas.
+        candidate_labels (list): A list of candidate labels.
+
+    Returns:
+        search_space (list): A list of formulas.
+    """
+    search_space = []
+
+    for formula in formulas:
+        vals_formula = set(formula.get_vals())
+        for candidate_term in candidate_labels:
+            # remove dummy cases with void masks or equivalent formulas
+            if candidate_term.val in vals_formula:
+                continue
+            for op, negate in [(F.Or, False), (F.And, False), (F.And, True)]:
+                candidate_to_attach = candidate_term
+                if negate:
+                    candidate_to_attach = F.Not(candidate_to_attach)
+                candidate_formula = op(formula, candidate_to_attach)
+                candidate_formula.iou = 1.0
+
+                search_space.append(candidate_formula)
+    search_space = list(set(search_space))
+    return search_space
+
+
+def compute_beam_quantities(beam, masks, beam_masks, heuristic_info, bitmaps):
+    """Computes the quantities for each formula in the beam.
+    Args:
+        beam (dict): A dictionary containing the formulas in the beam and their corresponding iou.
+        masks (list): A list of concept masks. Each mask is a tensor of shape (H*W).
+        beam_masks (dict): A dictionary containing the masks for each formula in the beam.
+        heuristic_info (tuple): A tuple containing the masks_info, neuron_quantities, and concept quantities for the optimal heuristic.
+        bitmaps (torch.Tensor): A tensor of shape (N, H*W) where N is the number of sample.
+    Returns:
+        tuple: A tuple containing the beam info and the updated beam masks.
+    """
+    seg_quantities, neuron_quantities, _ = heuristic_info
+    (neuron_unique, _), (neuron_common, _), _, _, _, _ = neuron_quantities
+    common_elements, unique_elements, _ = seg_quantities
+
+    # Elements re-used by all the nodes, we pre-move them to gpu if available
+    unique_elements = unique_elements.to(bitmaps.device)
+    common_elements = common_elements.to(bitmaps.device)
+    beam_info = {}
+    for label in beam:
+        if label in beam_masks or isinstance(label, F.Leaf):
+            # We already have the information for this label, we skip it
+            continue
+
+        # Compute the mask for the label
+        label_mask = mask_utils.get_formula_mask(label, masks, beam_masks, device=bitmaps.device)
+
+        # Compute the quantities from the mask
+        label_quantities = optimal_utils.compute_quantities_vector(
+            label_mask=label_mask,
+            bitmaps=bitmaps,
+            common_elements=common_elements,
+            unique_elements=unique_elements,
+            neuron_common=neuron_common,
+            neuron_unique=neuron_unique,
+        )
+
+        # Compute the info from the quantities
+        label_info = optimal_utils.get_concept_info(label_quantities)
+
+        beam_info[label] = label_info
+        label_mask = label_mask.cpu()
+        beam_masks[label] = label_mask
+    return beam_info, beam_masks
+
+
+def get_beam_info(
+    *, beam, masks, beam_masks, heuristic_info, label_mapping, bitmaps, length
+):
+    """Gets the information for each formula in the beam and updates the heuristic information and the label mapping.
+    Args:
+        beam (dict): A dictionary containing the formulas in the beam and their corresponding iou.
+        masks (list): A list of concept masks. Each mask is a tensor of shape (H*W).
+        beam_masks (dict): A dictionary containing the masks for each formula in the beam.
+        heuristic_info (tuple): A tuple containing the masks_info, neuron_quantities, and concept quantities for the optimal heuristic.
+        label_mapping (dict): A dictionary mapping labels to the indices of their corresponding masks.
+        bitmaps (torch.Tensor): A tensor of shape (N, H*W) where N is the number of sample.
+        length (int): The maximum length of the formulas to search.
+    Returns:
+        tuple: A tuple containing the updated beam masks and the updated heuristic information and label mapping.
+    """
+    beam_info, beam_masks = compute_beam_quantities(
+        beam, masks, beam_masks, heuristic_info, bitmaps
+    )
+    new_heuristic_info, new_label_mapping = optimal_utils.update_heuristic_info(
+        nodes_info=beam_info,
+        heuristic_info=heuristic_info,
+        label_mapping=label_mapping,
+        max_length=length,
+    )
+    return beam_masks, (new_heuristic_info, new_label_mapping)
+
+
+def sort_search_space_by(
+    *,
+    search_space,
+    label_mapping,
+    heuristic_info,
+    disjoint_info,
+    num_hits,
+    max_size_mask
+):
+    """Sorts the search space based on the estimated iou for each formula in the search space.
+    Args:
+        search_space (list): A list of formulas to sort.
+        label_mapping (dict): A dictionary mapping labels to the indices of their corresponding masks.
+        heuristic_info (tuple): A tuple containing the masks_info, neuron_quantities, and concept quantities for the optimal heuristic.
+        disjoint_info (dict): A dictionary containing the disjoint information for the concepts.
+        num_hits (int): The number of hits in the bitmaps.
+        max_size_mask (int): The maximum size of the mask.
+    Returns:
+        list: A sorted list of formulas based on the estimated iou.
+    """
+    _, neuron_quantities, _ = heuristic_info
+    for index_formula, candidate_formula in enumerate(search_space):
+        label_quantities = optimal_utils.estimate_label_quantities(
+            heuristic=optimal_sample_heuristic,
+            label=candidate_formula,
+            label_mapping=label_mapping,
+            heuristic_info=heuristic_info,
+            max_size_mask=max_size_mask,
+            disjoint_info=disjoint_info,
+        )
+        if label_quantities is None:
+            # Label discarded at the previous step
+            esti_iou = 0.0
+        else:
+
+            esti_iou = optimal_utils.compute_max_iou_from_label_info(
+                label_quantities=label_quantities,
+                num_hits=num_hits,
+                neuron_quantities=neuron_quantities,
+            )
+        search_space[index_formula].iou = esti_iou
+
+    # Sort the search space based on the estimated iou in descending order
+    search_space = sorted(search_space, key=lambda x: x.iou, reverse=True)
+    return search_space
+
+
+def perform_search(
+    *,
+    masks,
+    bitmaps,
+    heuristic_info,
+    disjoint_info,
+    max_size_mask,
+    beam_size=5,
+    length=3
+):
+    """Performs the beam optimal search to find the best formula that explains the bitmaps given the masks and the heuristic information.
+    Args:
+        masks (list): A list of concept masks. Each mask is a tensor of shape (H*W).
+        bitmaps (torch.Tensor): A tensor of shape (N, H*W) where N is the number of sample.
+        heuristic_info (tuple): A tuple containing the masks_info, neuron_quantities, and concept quantities for the optimal heuristic.
+        disjoint_info (dict): A dictionary containing the disjoint information for the concepts.
+        max_size_mask (int): The maximum size of the mask.
+        beam_size (int): The beam size for the search.
+        length (int): The maximum length of the formulas to search.
+    Returns:
+        tuple: A tuple containing the best formula, its iou, the total number of visited nodes, the number of expanded nodes, and the number of estimated nodes.
+    """
+
+    # Number of hits
+    num_hits = bitmaps.sum()
+
+    # Utilities
+    label_mapping = {}
+    leaf_mapping = {F.Leaf(c): c for c in range(len(masks))}
+    label_mapping.update(leaf_mapping)
+
+    # Extract first beam and candidate concepts
+    candidate_labels = [F.Leaf(c) for c in range(len(masks))]
+    
+    _, neuron_quantities, concept_quantities = heuristic_info
+    iou_atoms = {
+        k: optimal_utils.compute_max_iou_from_label_info(
+            concept_quantities[label_mapping[k]], neuron_quantities, num_hits
+        )
+        for k in candidate_labels
+    }
+    iou_atoms = Counter(iou_atoms)
+    first_beam_num = min(len(iou_atoms), beam_size * 2)
+    beam_atoms = {
+        lab: iou for lab, iou in iou_atoms.most_common(first_beam_num) if iou > 0
+    }
+
+    if len(beam_atoms) == 0:
+        return None
+
+    # Beam Search
+    beam_masks = {}
+    beam = beam_atoms.copy()
+    for previous_beam_length in range(1, length):
+        # Only expand formulas of the previous beam length to avoid regenerating beam tree already explored
+        to_expand = [lab for lab in beam.keys() if len(lab) == previous_beam_length]
+        search_space = compute_next_search_space(
+            to_expand,
+            candidate_labels,
+        )
+
+        # Compute the estimation for the next frontier
+        sorted_search_space = sort_search_space_by(
+            search_space=search_space,
+            label_mapping=label_mapping,
+            heuristic_info=heuristic_info,
+            num_hits=num_hits,
+            max_size_mask=max_size_mask,
+            disjoint_info=disjoint_info,
+        )
+
+        # # If we are in the last iteration, we set the beam size to 1 to get only the best formula
+        if previous_beam_length == length - 1:
+            beam_size = 1
+
+        # Perform beam search
+        next_beam_formulas, next_beam_iou, beam_visited = beam_search(
+            sorted_search_space,
+            masks=masks,
+            previous_beam=beam,
+            beam_masks=beam_masks,
+            bitmaps=bitmaps,
+            beam_limit=beam_size,
+        )
+
+        # Update top formulas
+        for index_beam in range(len(next_beam_formulas)):
+            beam.update({next_beam_formulas[index_beam]: next_beam_iou[index_beam]})
+
+        # Trim the beam
+        beam = dict(Counter(beam).most_common(beam_size))
+
+        # Update infos if there are step left
+        if previous_beam_length < length - 1:
+            beam_masks, (heuristic_info, label_mapping) = get_beam_info(
+                beam=beam.keys(),
+                masks=masks,
+                beam_masks=beam_masks,
+                heuristic_info=heuristic_info,
+                label_mapping=label_mapping,
+                bitmaps=bitmaps,
+                length=length,
+            )
+
+    top_result = Counter(beam).most_common(1)[0]
+
+    best_label = top_result[0]
+    return best_label
+
+
+def compute_beam_explanations(
+    *, bitmaps, masks, disjoint_info, length, beam_size, device, cache_dir=None
+):
+    """Computes the beam explanations for the given bitmaps and masks.
+    Args:
+        bitmaps (torch.Tensor): A tensor of shape (N, H*W) where N is the number of sample.
+        masks (list): A list of concept masks. Each mask is a tensor of shape (H*W).
+        disjoint_info (dict): A dictionary containing the disjoint information for the concepts.
+        length (int): The maximum length of the formulas to search.
+        device (torch.device): The device to perform the computations on.
+        cache_dir (str, optional): The directory to store the cached information for the masks. If None, the information is not cached and is recomputed every time. Defaults to None.
+    Returns:
+        tuple: A tuple containing the best formula, its iou, the total number of visited nodes, the number of expanded nodes, and the number of estimated nodes.
+    """
+
+    max_size_mask = torch.numel(masks[0])
+    bitmaps = bitmaps.to(device)
+    
+    # Get masks info
+    masks_info = mask_utils.get_dataset_info(masks, cache_dir)
+    
+    # Get quantities
+    neuron_quantities, concept_quantities = optimal_utils.get_optimal_heuristic_info(
+        masks=masks, bitmaps=bitmaps, masks_quantities=masks_info
+    )
+
+    heuristic_info = (
+        masks_info,
+        neuron_quantities,
+        concept_quantities,
+    )
+    best_label = perform_search(
+        masks=masks,
+        bitmaps=bitmaps,
+        heuristic_info=heuristic_info,
+        disjoint_info=disjoint_info,
+        max_size_mask=max_size_mask,
+        beam_size=beam_size,
+        length=length,
+    )
+    return best_label