compositional-explanations 0.0.0.1__py3-none-any.whl

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