make_upset_plot.py

from upsetplot import from_memberships
import sys

def membership_line_to_array(line: str, base_name: str):
    """
    Fromm a line membership to an array membership: 
    eg : 1010 -> ['geno1', 'geno3'] (if base_name == 'geno')
    """
    res = []
    i=0
    for char in line: 
        i+=1
        if char == '1': res.append(f"{base_name}{i}")
    return res

def load_venn_from_fredy(venn_file_name: str, base_name: str='geno'):
    """
    From a venn file generated by fredy as : 
    #Venn:  
    0000 111875
    1000 346
    0100 357
    1100 272
    0010 398
    1010 0
    0110 31
    1110 107
    0001 362
    1001 1
    0101 9
    1101 89
    0011 199
    1011 76
    0111 54
    1111 2057

    generate a data usable by upsetplot:
    venn_data = from_memberships(
        [[],
        ['geno1'],
        ['geno2'],
        ['geno1', 'geno2'],
        ['geno3'],
        ['geno1', 'geno3'],
        ...
        ],
        data=[111875,346,357,272,398,272,...]
    )
    """
    with open(venn_file_name) as venn_file:
        members = []     # first array = arrays of membership, second array: corresponding counts
        abundances = []
        for line in venn_file.readlines():
            if line[0]=='#': continue
            sline = line.strip().split()
            members.append(membership_line_to_array(sline[0], base_name))
            abundances.append(int(sline[1]))
    return from_memberships(members,data=abundances)



membership = load_venn_from_fredy(sys.argv[1])
from upsetplot import plot
plot(membership, show_percentages=True)  
from matplotlib import pyplot
pyplot.show() 

## TODO: deal with input files
## TODO: propose interactive way to select members