anchor_base.py

import numpy as np
import copy
import collections
import logging
from typing import Callable, Tuple, Set, Dict, Sequence

logger = logging.getLogger(__name__)


def matrix_subset(matrix: np.ndarray, n_samples: int) -> np.ndarray:
    """
    Parameters
    ----------
    matrix
        Matrix to sample from
    n_samples
        Number of samples returned

    Returns
    -------
    Sample of the input matrix.
    """
    if matrix.shape[0] == 0:
        return matrix
    n_samples = min(matrix.shape[0], n_samples)
    return matrix[np.random.choice(matrix.shape[0], n_samples, replace=False)]


class AnchorBaseBeam(object):

    def __init__(self) -> None:
        """
        Initialize the anchor beam search class.
        """
        pass

    @staticmethod
    def kl_bernoulli(p: float, q: float) -> float:
        """
        Compute KL-divergence between 2 probabilities p and q.

        Parameters
        ----------
        p
            Probability
        q
            Probability

        Returns
        -------
        KL-divergence
        """
        p = min(0.9999999999999999, max(0.0000001, p))
        q = min(0.9999999999999999, max(0.0000001, q))
        return (p * np.log(float(p) / q) + (1 - p) *
                np.log(float(1 - p) / (1 - q)))

    @staticmethod
    def dup_bernoulli(p: float, level: float, n_iter: int = 17) -> float:
        """
        Update upper precision bound for a candidate anchor dependent on the KL-divergence.

        Parameters
        ----------
        p
            Precision of candidate anchor
        level
            beta / nb of samples
        n_iter
            Number of iterations during lower bound update

        Returns
        -------
        Updated upper precision bound
        """
        # TODO: where does 17x sampling come from?
        lm = p
        um = min(min(1, p + np.sqrt(level / 2.)), 1)
        for j in range(1, n_iter):
            qm = (um + lm) / 2.
            if AnchorBaseBeam.kl_bernoulli(p, qm) > level:
                um = qm
            else:
                lm = qm
        return um

    @staticmethod
    def dlow_bernoulli(p: float, level: float, n_iter: int = 17) -> float:
        """
        Update lower precision bound for a candidate anchor dependent on the KL-divergence.

        Parameters
        ----------
        p
            Precision of candidate anchor
        level
            beta / nb of samples
        n_iter
            Number of iterations during lower bound update

        Returns
        -------
        Updated lower precision bound
        """
        um = p
        lm = max(min(1, p - np.sqrt(level / 2.)), 0)  # lower bound
        for j in range(1, n_iter):
            qm = (um + lm) / 2.
            if AnchorBaseBeam.kl_bernoulli(p, qm) > level:  # KL-divergence > threshold level
                lm = qm
            else:
                um = qm
        return lm

    @staticmethod
    def compute_beta(n_features: int, t: int, delta: float) -> float:
        """
        Parameters
        ----------
        n_features
            Number of candidate anchors
        t
            Iteration number
        delta

        Returns
        -------
        Level used to update upper and lower precision bounds.
        """
        # TODO: where do magic numbers come from?
        alpha = 1.1
        k = 405.5
        temp = np.log(k * n_features * (t ** alpha) / delta)
        return temp + np.log(temp)

    @staticmethod
    def lucb(sample_fns: list, initial_stats: dict, epsilon: float, delta: float, batch_size: int, top_n: int,
             verbose: bool = False, verbose_every: int = 1) -> Sequence:
        """
        Parameters
        ----------
        sample_fns
            List with sample functions for each candidate anchor
        initial_stats
            Dictionary with lists containing nb of samples used and where sample predictions equal the desired label
        epsilon
            Precision bound tolerance for convergence
        delta
            Used to compute beta
        batch_size
            Number of samples
        top_n
            Min of beam width size or number of candidate anchors
        verbose
            Whether to print intermediate output
        verbose_every
            Whether to print intermediate output every verbose_every steps

        Returns
        -------
        Indices of best anchor options. Number of indices equals min of beam width or nb of candidate anchors.
        """
        # n_features equals to the nb of candidate anchors
        n_features = len(sample_fns)

        # initiate arrays for number of samples, positives (sample prediction equals desired label), ...
        # ... upper and lower precision bounds for each anchor candidate
        n_samples = np.array(initial_stats['n_samples'])
        positives = np.array(initial_stats['positives'])
        ub = np.zeros(n_samples.shape)
        lb = np.zeros(n_samples.shape)
        for f in np.where(n_samples == 0)[0]:
            n_samples[f] += 1  # set min samples for each anchor candidate to 1
            positives[f] += sample_fns[f](1)  # add labels.sum() for the anchor candidate

        if n_features == top_n:  # return all options b/c of beam search width
            return range(n_features)

        means = positives / n_samples  # fraction sample predictions equal to desired label
        t = 1

        def update_bounds(t: int) -> Tuple[np.ndarray, np.ndarray]:
            """
            Parameters
            ----------
            t
                Iteration number

            Returns
            -------
            Upper and lower precision bound indices.
            """
            sorted_means = np.argsort(means)  # ascending sort of anchor candidates by precision

            beta = AnchorBaseBeam.compute_beta(n_features, t, delta)

            # J = the beam width top anchor candidates with highest precision
            # not_J = the rest
            J = sorted_means[-top_n:]
            not_J = sorted_means[:-top_n]

            for f in not_J:  # update upper bound for lowest precision anchor candidates
                ub[f] = AnchorBaseBeam.dup_bernoulli(means[f], beta / n_samples[f])

            for f in J:  # update lower bound for highest precision anchor candidates
                lb[f] = AnchorBaseBeam.dlow_bernoulli(means[f], beta / n_samples[f])

            # for the low precision anchor candidates, compute the upper precision bound and keep the index ...
            # ... of the anchor candidate with the highest upper precision value -> ut
            # for the high precision anchor candidates, compute the lower precision bound and keep the index ...
            # ... of the anchor candidate with the lowest lower precision value -> lt
            ut = not_J[np.argmax(ub[not_J])]
            lt = J[np.argmin(lb[J])]
            return ut, lt

        # keep updating the upper and lower precision bounds until the difference between the best upper ...
        # ... precision bound of the low precision anchors and the worst lower precision bound of the high ...
        # ... precision anchors is smaller than eps
        ut, lt = update_bounds(t)
        B = ub[ut] - lb[lt]
        verbose_count = 0
        while B > epsilon:
            verbose_count += 1
            if verbose and verbose_count % verbose_every == 0:
                print('Best: %d (mean:%.10f, n: %d, lb:%.4f)' %
                      (lt, means[lt], n_samples[lt], lb[lt]), end=' ')
                print('Worst: %d (mean:%.4f, n: %d, ub:%.4f)' %
                      (ut, means[ut], n_samples[ut], ub[ut]), end=' ')
                print('B = %.2f' % B)
            n_samples[ut] += batch_size
            positives[ut] += sample_fns[ut](batch_size)  # sample new batch of data
            means[ut] = positives[ut] / n_samples[ut]
            n_samples[lt] += batch_size
            positives[lt] += sample_fns[lt](batch_size)  # sample new batch of data
            means[lt] = positives[lt] / n_samples[lt]
            t += 1
            ut, lt = update_bounds(t)
            B = ub[ut] - lb[lt]
        sorted_means = np.argsort(means)
        return sorted_means[-top_n:]

    @staticmethod
    def make_tuples(previous_best: list, state: dict) -> list:
        """
        Parameters
        ----------
        previous_best
            List with tuples of anchor candidates
        state
            Dictionary with the relevant metrics like coverage and samples for candidate anchors

        Returns
        -------
        List with tuples of candidate anchors with additional metadata.
        """
        # compute some variables used later on
        normalize_tuple = lambda x: tuple(sorted(set(x)))  # noqa E731
        all_features = range(state['n_features'])
        coverage_data = state['coverage_data']
        current_idx = state['current_idx']
        data = state['data'][:current_idx]
        labels = state['labels'][:current_idx]

        # initially, every feature separately is an anchor
        if len(previous_best) == 0:
            tuples = [(x,) for x in all_features]
            for x in tuples:
                pres = data[:, x[0]].nonzero()[0]
                state['t_idx'][x] = set(pres)
                state['t_nsamples'][x] = float(len(pres))
                state['t_positives'][x] = float(labels[pres].sum())
                state['t_order'][x].append(x[0])
                state['t_coverage_idx'][x] = set(coverage_data[:, x[0]].nonzero()[0])
                state['t_coverage'][x] = (float(len(state['t_coverage_idx'][x])) / coverage_data.shape[0])
            return tuples

        # create new anchors: add a feature to every anchor in current best
        new_tuples = set()  # type: Set[tuple]
        for f in all_features:
            for t in previous_best:
                new_t = normalize_tuple(t + (f,))
                if len(new_t) != len(t) + 1:
                    continue
                if new_t not in new_tuples:
                    new_tuples.add(new_t)
                    state['t_order'][new_t] = copy.deepcopy(state['t_order'][t])
                    state['t_order'][new_t].append(f)
                    state['t_coverage_idx'][new_t] = (state['t_coverage_idx'][t].intersection(
                        state['t_coverage_idx'][(f,)]))
                    state['t_coverage'][new_t] = (float(len(state['t_coverage_idx'][new_t])) / coverage_data.shape[0])
                    t_idx = np.array(list(state['t_idx'][t]))
                    t_data = state['data'][t_idx]
                    present = np.where(t_data[:, f] == 1)[0]
                    state['t_idx'][new_t] = set(t_idx[present])
                    idx_list = list(state['t_idx'][new_t])
                    state['t_nsamples'][new_t] = float(len(idx_list))
                    state['t_positives'][new_t] = np.sum(state['labels'][idx_list])
        return list(new_tuples)

    @staticmethod
    def get_sample_fns(sample_fn: Callable, tuples: list, state: dict, data_type: str = None) -> list:
        """
        Parameters
        ----------
        sample_fn
            Sample function, returns both raw and categorized data as well as labels
        tuples
            List of anchor candidates
        state
            Dictionary with the relevant metrics like coverage and samples for candidate anchors
        data_type
            Data type for raw data

        Returns
        -------
        List with sample functions for each candidate anchor.
        """

        def complete_sample_fn(t: tuple, n: int) -> int:
            """
            Parameters
            ----------
            t
                Anchor candidate
            n
                Number of samples

            Returns
            -------
            Sum of where sampled data equals desired label of observation to be explained.
            """
            raw_data, data, labels = sample_fn(list(t), n)
            current_idx = state['current_idx']
            idxs = range(current_idx, current_idx + n)
            state['t_idx'][t].update(idxs)
            state['t_nsamples'][t] += n
            state['t_positives'][t] += labels.sum()
            state['data'][idxs] = data
            state['raw_data'][idxs] = raw_data
            state['labels'][idxs] = labels
            state['current_idx'] += n
            if state['current_idx'] >= state['data'].shape[0] - max(1000, n):
                prealloc_size = state['prealloc_size']
                current_idx = data.shape[0]
                state['data'] = np.vstack((state['data'], np.zeros((prealloc_size, data.shape[1]), data.dtype)))
                dtype = data_type if data_type is not None else raw_data.dtype
                state['raw_data'] = np.vstack((state['raw_data'], np.zeros((prealloc_size, raw_data.shape[1]),
                                                                           dtype=dtype)))
                state['labels'] = np.hstack((state['labels'], np.zeros(prealloc_size, labels.dtype)))
            return labels.sum()

        sample_fns = []
        for t in tuples:
            sample_fns.append(lambda n, t=t: complete_sample_fn(t, n))

        return sample_fns

    @staticmethod
    def get_initial_statistics(tuples: list, state: dict) -> dict:
        """
        Parameters
        ----------
        tuples
            Candidate anchors
        state
            Dictionary with the relevant metrics like coverage and samples for candidate anchors

        Returns
        -------
        Dictionary with lists containing nb of samples used and where sample predictions equal the desired label.
        """
        stats = {
            'n_samples': [],
            'positives': []
        }  # type: Dict[str, list]
        for t in tuples:
            stats['n_samples'].append(state['t_nsamples'][t])
            stats['positives'].append(state['t_positives'][t])
        return stats

    @staticmethod
    def get_anchor_from_tuple(t: tuple, state: dict) -> dict:
        """
        Parameters
        ----------
        t
            Anchor
        state
            Dictionary with the relevant metrics like coverage and samples for candidate anchors

        Returns
        -------
        Anchor dictionary with anchor features and additional metadata.
        """
        # TODO: This is wrong, some of the intermediate anchors may not exist.
        anchor = {'feature': [], 'mean': [], 'precision': [],
                  'coverage': [], 'examples': [], 'all_precision': 0}  # type: dict
        anchor['num_preds'] = state['data'].shape[0]
        normalize_tuple = lambda x: tuple(sorted(set(x)))  # noqa E731
        current_t = tuple()  # type: tuple
        for f in state['t_order'][t]:
            current_t = normalize_tuple(current_t + (f,))
            mean = (state['t_positives'][current_t] / state['t_nsamples'][current_t])
            anchor['feature'].append(f)
            anchor['mean'].append(mean)
            anchor['precision'].append(mean)
            anchor['coverage'].append(state['t_coverage'][current_t])

            # add examples where anchor does or does not hold
            raw_idx = list(state['t_idx'][current_t])
            raw_data = state['raw_data'][raw_idx]
            covered_true = (state['raw_data'][raw_idx][state['labels'][raw_idx] == 1])
            covered_false = (state['raw_data'][raw_idx][state['labels'][raw_idx] == 0])
            exs = {}
            exs['covered'] = matrix_subset(raw_data, 10)
            exs['covered_true'] = matrix_subset(covered_true, 10)
            exs['covered_false'] = matrix_subset(covered_false, 10)
            exs['uncovered_true'] = np.array([])
            exs['uncovered_false'] = np.array([])
            anchor['examples'].append(exs)
        return anchor

    @staticmethod
    def anchor_beam(sample_fn: Callable, delta: float = 0.05, epsilon: float = 0.1, batch_size: int = 10,
                    desired_confidence: float = 1, beam_size: int = 1, verbose: bool = False,
                    epsilon_stop: float = 0.05, min_samples_start: int = 0, max_anchor_size: int = None,
                    verbose_every: int = 1, stop_on_first: bool = False, coverage_samples: int = 10000,
                    data_type: str = None, **kwargs) -> dict:
        """
        Parameters
        ----------
        sample_fn
            Function used to sample from training set which returns (raw) data and labels
        delta
            Used to compute beta
        epsilon
            Precision bound tolerance for convergence
        batch_size
            Number of samples
        desired_confidence
            Desired level of precision, tau in paper
        beam_size
            Beam width
        verbose
            Whether to print intermediate output
        epsilon_stop
            Confidence bound margin around desired precision
        min_samples_start
            Min number of initial samples
        max_anchor_size
            Max number of features in anchor
        verbose_every
            Whether to print intermediate output every verbose_every steps
        stop_on_first
            Stop on first valid anchor found
        coverage_samples
            Number of samples used to compute coverage
        data_type
            Data type for raw data

        Returns
        -------
        Explanation dictionary containing anchors with metadata like coverage and precision.
        """
        # initiate empty anchor dict
        anchor = {'feature': [], 'mean': [], 'precision': [],
                  'coverage': [], 'examples': [], 'all_precision': 0}

        # random (b/c first argument is empty) sample nb of coverage_samples from training data
        _, coverage_data, _ = sample_fn([], coverage_samples, compute_labels=False)

        # sample by default 1 or min_samples_start more random value(s)
        raw_data, data, labels = sample_fn([], max(1, min_samples_start))

        # mean = fraction of labels sampled data that equals the label of the instance to be explained ...
        # ... and is equivalent to prec(A) in paper (eq.2)
        # get lower precision bound lb
        mean = labels.mean()
        beta = np.log(1. / delta)
        lb = AnchorBaseBeam.dlow_bernoulli(mean, beta / data.shape[0])

        # while prec(A) > tau (precision constraint) for A=[] and prec_lb(A) < tau - eps ...
        # ... (lower precision bound below tau with margin eps), keep sampling data until lb is high enough
        while mean > desired_confidence and lb < desired_confidence - epsilon:
            nraw_data, ndata, nlabels = sample_fn([], batch_size)
            data = np.vstack((data, ndata))
            raw_data = np.vstack((raw_data, nraw_data))
            labels = np.hstack((labels, nlabels))
            mean = labels.mean()
            lb = AnchorBaseBeam.dlow_bernoulli(mean, beta / data.shape[0])

        # if prec_lb(A) > tau for A=[] then the empty anchor already satisfies the constraints ...
        # ... and an empty anchor is returned
        if lb > desired_confidence:
            anchor['num_preds'] = data.shape[0]
            anchor['all_precision'] = mean
            return anchor

        # initialize variables
        prealloc_size = batch_size * 10000
        current_idx = data.shape[0]
        data = np.vstack((data, np.zeros((prealloc_size, data.shape[1]), data.dtype)))
        dtype = data_type if data_type is not None else raw_data.dtype
        raw_data = np.vstack((raw_data, np.zeros((prealloc_size, raw_data.shape[1]), dtype=dtype)))
        labels = np.hstack((labels, np.zeros(prealloc_size, labels.dtype)))
        n_features = data.shape[1]
        state = {'t_idx': collections.defaultdict(lambda: set()),
                 't_nsamples': collections.defaultdict(lambda: 0.),
                 't_positives': collections.defaultdict(lambda: 0.),
                 'data': data,
                 'prealloc_size': prealloc_size,
                 'raw_data': raw_data,
                 'labels': labels,
                 'current_idx': current_idx,
                 'n_features': n_features,
                 't_coverage_idx': collections.defaultdict(lambda: set()),
                 't_coverage': collections.defaultdict(lambda: 0.),
                 'coverage_data': coverage_data,
                 't_order': collections.defaultdict(lambda: list())
                 }
        current_size = 1
        best_of_size = {0: []}  # type: Dict[int, list]
        best_coverage = -1
        best_tuple = ()
        if max_anchor_size is None:
            max_anchor_size = n_features

        # find best anchor using beam search until max anchor size
        while current_size <= max_anchor_size:

            # create new candidate anchors by adding features to current best anchors
            tuples = AnchorBaseBeam.make_tuples(best_of_size[current_size - 1], state)

            # goal is to max coverage given precision constraint P(prec(A) > tau) > 1 - delta (eq.4)
            # so keep tuples with higher coverage than current best coverage
            tuples = [x for x in tuples if state['t_coverage'][x] > best_coverage]

            # if no better coverage found with added features -> break
            if len(tuples) == 0:
                break

            # build sample functions for each tuple in tuples list
            # these functions sample randomly for all features except for the ones in the candidate anchors
            # for the features in the anchor it uses the same category (categorical features) or samples from ...
            # ... the same bin (discretized numerical features) as the feature in the observation that is explained
            sample_fns = AnchorBaseBeam.get_sample_fns(sample_fn, tuples, state, data_type=dtype)

            # for each tuple, get initial nb of samples used and prec(A)
            initial_stats = AnchorBaseBeam.get_initial_statistics(tuples, state)

            # apply KL-LUCB and return anchor options (nb of options = beam width)
            # anchor options are in the form of indices of candidate anchors
            chosen_tuples = AnchorBaseBeam.lucb(sample_fns,
                                                initial_stats,
                                                epsilon, delta,
                                                batch_size,
                                                min(beam_size, len(tuples)),
                                                verbose=verbose,
                                                verbose_every=verbose_every)

            # store best anchors for the given anchor size (nb of features in the anchor)
            best_of_size[current_size] = [tuples[x] for x in chosen_tuples]
            if verbose:
                print('Best of size ', current_size, ':')

            # for each candidate anchor:
            # update precision, lower and upper bounds until precision constraints are met
            # update best anchor if coverage is larger than current best coverage
            stop_this = False
            for i, t in zip(chosen_tuples, best_of_size[current_size]):

                # choose at most (beam_size - 1) tuples at each step with at most n_feature steps
                beta = np.log(1. / (delta / (1 + (beam_size - 1) * n_features)))

                # get precision, lower and upper bounds, and coverage for candidate anchor
                mean = state['t_positives'][t] / state['t_nsamples'][t]
                lb = AnchorBaseBeam.dlow_bernoulli(mean, beta / state['t_nsamples'][t])
                ub = AnchorBaseBeam.dup_bernoulli(mean, beta / state['t_nsamples'][t])
                coverage = state['t_coverage'][t]

                if verbose:
                    print(i, mean, lb, ub)

                # while prec(A) >= tau and prec_lb(A) < tau - eps or prec(A) < tau and prec_ub(A) > tau + eps
                # sample more data and update lower and upper precision bounds ...
                # ... b/c respectively either prec_lb(A) or prec(A) needs to improve
                while ((mean >= desired_confidence and lb < desired_confidence - epsilon_stop) or
                       (mean < desired_confidence and ub >= desired_confidence + epsilon_stop)):
                    # sample a batch of data, get new precision, lb and ub values
                    sample_fns[i](batch_size)
                    mean = state['t_positives'][t] / state['t_nsamples'][t]
                    lb = AnchorBaseBeam.dlow_bernoulli(mean, beta / state['t_nsamples'][t])
                    ub = AnchorBaseBeam.dup_bernoulli(mean, beta / state['t_nsamples'][t])

                if verbose:
                    print('%s mean = %.2f lb = %.2f ub = %.2f coverage: %.2f n: %d' %
                          (t, mean, lb, ub, coverage, state['t_nsamples'][t]))

                # if prec(A) > tau and prec_lb(A) > tau - eps then we found an eligible anchor
                if mean >= desired_confidence and lb > desired_confidence - epsilon_stop:

                    if verbose:
                        print('Found eligible anchor ', t, 'Coverage:',
                              coverage, 'Is best?', coverage > best_coverage)

                    # coverage eligible anchor needs to be bigger than current best coverage
                    if coverage > best_coverage:
                        best_coverage = coverage
                        best_tuple = t
                        if best_coverage == 1 or stop_on_first:
                            stop_this = True
            if stop_this:
                break
            current_size += 1

        # if no anchor is found, choose highest precision of best anchor candidate from every round
        if best_tuple == ():
            logger.warning('Could not find an anchor satisfying the {} precision constraint. Now returning '
                           'the best non-eligible anchor.'.format(desired_confidence))
            tuples = []
            for i in range(0, current_size):
                tuples.extend(best_of_size[i])
            sample_fns = AnchorBaseBeam.get_sample_fns(sample_fn, tuples, state, data_type=dtype)
            initial_stats = AnchorBaseBeam.get_initial_statistics(tuples, state)
            chosen_tuples = AnchorBaseBeam.lucb(sample_fns, initial_stats, epsilon,
                                                delta, batch_size, 1, verbose=verbose)
            best_tuple = tuples[chosen_tuples[0]]

        # return explanation dictionary
        return AnchorBaseBeam.get_anchor_from_tuple(best_tuple, state)