Prune set of GWAS based on LD and significance. A graph of all SNVs is
    constructed with edges for LD >= 0.8 and the most significant SNV per
    connected component is kept. 
    
    Parameters
    ----------
    df : pandas.DataFrame
        Pandas dataframe with unique SNVs. The index is of the form chrom:pos
        where pos is the one-based position of the SNV. The columns must include
        chrom, start, end, and pvalue. chrom, start, end make a zero-based bed
        file with the SNV coordinates.

    ld_beds : dict
        Dict whose keys are chromosomes and whose values are filenames of
        tabixed LD bed files. An LD bed file looks like "chr1    11007   11008
        11008:11012:1" where the first three columns are the zero-based
        half-open coordinate of the SNV and the fourth column has the one-based
        coordinate followed of the SNV followed by the one-based coordinate of a
        different SNV and the LD between them. In this example, the variants are
        in perfect LD. The bed file should also contain the reciprocal line for
        this LD relationship: "chr1    11011   11012   11012:11008:1".

    snvs : list
        List of SNVs to filter against. If a SNV is not in this list, it will
        not be included. If you are working with GWAS SNPs, this is useful for
        filtering out SNVs that aren't in the SNPsnap database for instance.

    Returns
    -------
    out : pandas.DataFrame
        Pandas dataframe in the same format as the input dataframe but with only
        independent SNVs.

def ld_prune(df, ld_beds, snvs=None):
    """
    Prune set of GWAS based on LD and significance. A graph of all SNVs is
    constructed with edges for LD >= 0.8 and the most significant SNV per
    connected component is kept. 
    
    Parameters
    ----------
    df : pandas.DataFrame
        Pandas dataframe with unique SNVs. The index is of the form chrom:pos
        where pos is the one-based position of the SNV. The columns must include
        chrom, start, end, and pvalue. chrom, start, end make a zero-based bed
        file with the SNV coordinates.

    ld_beds : dict
        Dict whose keys are chromosomes and whose values are filenames of
        tabixed LD bed files. An LD bed file looks like "chr1    11007   11008
        11008:11012:1" where the first three columns are the zero-based
        half-open coordinate of the SNV and the fourth column has the one-based
        coordinate followed of the SNV followed by the one-based coordinate of a
        different SNV and the LD between them. In this example, the variants are
        in perfect LD. The bed file should also contain the reciprocal line for
        this LD relationship: "chr1    11011   11012   11012:11008:1".

    snvs : list
        List of SNVs to filter against. If a SNV is not in this list, it will
        not be included. If you are working with GWAS SNPs, this is useful for
        filtering out SNVs that aren't in the SNPsnap database for instance.

    Returns
    -------
    out : pandas.DataFrame
        Pandas dataframe in the same format as the input dataframe but with only
        independent SNVs.
    """
    import networkx as nx
    import tabix
    if snvs:
        df = df.ix[set(df.index) & set(snvs)]
    keep = set()
    for chrom in ld_beds.keys():
        tdf = df[df['chrom'].astype(str) == chrom]
        if tdf.shape[0] > 0:
            f = tabix.open(ld_beds[chrom])
            # Make a dict where each key is a SNP and the values are all of the
            # other SNPs in LD with the key.
            ld_d = {}
            for j in tdf.index:
                p = tdf.ix[j, 'end']
                ld_d[p] = []
                try:
                    r = f.query(chrom, p - 1, p)
                    while True:
                        try:
                            n = r.next()
                            p1, p2, r2 = n[-1].split(':')
                            if float(r2) >= 0.8:
                                ld_d[p].append(int(p2))
                        except StopIteration:
                            break
                except TabixError:
                    continue
            # Make adjacency matrix for LD.
            cols = sorted(list(set(
                [item for sublist in ld_d.values() for item in sublist])))
            t = pd.DataFrame(0, index=ld_d.keys(), columns=cols)
            for k in ld_d.keys():
                t.ix[k, ld_d[k]] = 1
            t.index = ['{}:{}'.format(chrom, x) for x in t.index]
            t.columns = ['{}:{}'.format(chrom, x) for x in t.columns]
            # Keep all SNPs not in LD with any others. These will be in the index
            # but not in the columns.
            keep |= set(t.index) - set(t.columns)
            # Filter so we only have SNPs that are in LD with at least one other
            # SNP.
            ind = list(set(t.columns) & set(t.index))
            # Keep one most sig. SNP per connected subgraph.
            t = t.ix[ind, ind]
            g = nx.Graph(t.values)
            c = nx.connected_components(g)
            while True:
                try:
                    sg = c.next()
                    s = tdf.ix[t.index[list(sg)]]
                    keep.add(s[s.pvalue == s.pvalue.min()].index[0])
                except StopIteration:
                    break
    out = df.ix[keep]
    return out