An expectation maximization technique that produces pathway definitions
    devoid of crosstalk. That is, each gene is mapped to the pathway in
    which it has the greatest predicted impact; this removes any overlap
    between pathway definitions.

    Parameters
    -----------
    membership_matrix : numpy.array(float), shape = [n, k]
        The observed gene-to-pathway membership matrix, where n is the number
        of genes and k is the number of pathways we are interested in.
    max_iters : int (default=1000)
        The maximum number of expectation-maximization steps to take.

    Returns
    -----------
    dict(int -> set(int)), a dictionary mapping a pathway to a set of genes.
    These are the pathway definitions after the maximum impact estimation
    procedure has been applied to remove crosstalk.
      - The keys are ints corresponding to the pathway column indices in the
        membership matrix.
      - The values are sets of ints corresponding to gene row indices in the
        membership matrix.

def maximum_impact_estimation(membership_matrix, max_iters=1000):
    """An expectation maximization technique that produces pathway definitions
    devoid of crosstalk. That is, each gene is mapped to the pathway in
    which it has the greatest predicted impact; this removes any overlap
    between pathway definitions.

    Parameters
    -----------
    membership_matrix : numpy.array(float), shape = [n, k]
        The observed gene-to-pathway membership matrix, where n is the number
        of genes and k is the number of pathways we are interested in.
    max_iters : int (default=1000)
        The maximum number of expectation-maximization steps to take.

    Returns
    -----------
    dict(int -> set(int)), a dictionary mapping a pathway to a set of genes.
    These are the pathway definitions after the maximum impact estimation
    procedure has been applied to remove crosstalk.
      - The keys are ints corresponding to the pathway column indices in the
        membership matrix.
      - The values are sets of ints corresponding to gene row indices in the
        membership matrix.
    """
    # Initialize the probability vector as the sum of each column in the
    # membership matrix normalized by the sum of the entire membership matrix.
    # The probability at some index j in the vector represents the likelihood
    # that a pathway (column) j is defined by the current set of genes (rows)
    # in the membership matrix.
    pr_0 = np.sum(membership_matrix, axis=0) / np.sum(membership_matrix)
    pr_1 = _update_probabilities(pr_0, membership_matrix)
    epsilon = np.linalg.norm(pr_1 - pr_0)/100.

    pr_old = pr_1
    check_for_convergence = epsilon
    count = 0
    while epsilon > NEAR_ZERO and check_for_convergence >= epsilon:
        count += 1
        if count > max_iters:
            print("Reached the maximum number of iterations {0}".format(
                  max_iters))
            break
        pr_new = _update_probabilities(pr_old, membership_matrix)
        check_for_convergence = np.linalg.norm(pr_new - pr_old)
        pr_old = pr_new

    pr_final = pr_old  # renaming for readability

    corrected_pathway_definitions = {}
    n, k = membership_matrix.shape
    for gene_index in range(n):
        gene_membership = membership_matrix[gene_index]
        denominator = np.dot(gene_membership, pr_final)
        # Approximation is used to prevent divide by zero warning.
        # Since we are only looking for the _most_ probable pathway in which a
        # gene contributes its maximum impact, precision is not as important
        # as maintaining the relative differences between each
        # pathway's probability.
        if denominator < NEAR_ZERO:
            denominator = NEAR_ZERO
        # This is equivalent to one row in what Donato et al. (2013) refer
        # to as the underlying (latent) Z matrix.
        conditional_pathway_pr = (np.multiply(gene_membership, pr_final) /
                                  denominator)
        all_pathways_at_max = np.where(
            conditional_pathway_pr == conditional_pathway_pr.max())[0]
        gene_in_pathways = np.where(gene_membership == 1)[0]
        all_pathways_at_max = np.intersect1d(
            all_pathways_at_max, gene_in_pathways)
        pathway_index = np.random.choice(all_pathways_at_max)
        if pathway_index not in corrected_pathway_definitions:
            corrected_pathway_definitions[pathway_index] = set()
        corrected_pathway_definitions[pathway_index].add(gene_index)
    return corrected_pathway_definitions