text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Feb  7 10:53:03 2017

@author: jerrywzy, Pablo Carbonell
"""
import re
import os, subprocess
import json
import csv
import argparse
import quickRsim
import numpy as np
import pandas as pd
from rdkit.Chem import AllChem, Draw
from time import time
from brs_utils import print_OK, print_FAILED

class preLoad(object):
    """ Container of precomputed data """
    def __init__(self, datadir):
        self.datadir = datadir
        print ("Acquiring databases...")
        self.readData()

    def readData(self):
        """ Read all data into memory """
        print('   + Get fasta...', end = '', flush=True)
        start = time()
        self.fasta('seqs.fasta')
        print_OK(time()-start)

        print('   + Get fpData...', end = '', flush=True)
        start = time()
        self.fpData()
        print_OK(time()-start)

        print('   + Get seqData...', end = '', flush=True)
        start = time()
        self.seqData(['reac_seqs.tsv', 'upclst.json', 'clstrep.json', "seq_org.tsv", "org_lineage.csv"])
        print_OK(time()-start)

        print('   + Get reacData...', end = '', flush=True)
        start = time()
        self.reacData(['reac_smi.csv','reac_xref.tsv',"rxn_consensus_20160612.txt",
                        'reac_prop.tsv', 'brenda-mnxref2.tsv', 'sabiork-mnxref2.tsv', 'reac_xref_v3.tsv'])
        print_OK(time()-start)

    def fasta(self, ffile="seqs.fasta"):
        (sequence, names, descriptions, fulldescriptions, osource, pexistence) = readFasta(self.datadir, ffile)
        self.sequence = sequence
        self.names = names
        self.descriptions = descriptions
        self.fulldescriptions = fulldescriptions
        self.osource = osource
        self.pexistence = pexistence

    def fpData(self):
        self.fp = {}
        for fpid in quickRsim.fingerprint():
            self.fp[fpid] = quickRsim.loadFingerprint(self.datadir, fpid)

    def seqData(self, fl):
        with open(os.path.join(self.datadir, fl[0])) as f:
            self.MnxToUprot = {}
            self.UprotToMnx = {}
            for row in f:
                mnxr, db, seqid, source, ec = row.split('\t')
                if db == 'uniprot':
                    if mnxr not in self.MnxToUprot:
                        self.MnxToUprot[mnxr] = set()
                    self.MnxToUprot[mnxr].add(seqid)
                    if seqid not in self.UprotToMnx:
                        self.UprotToMnx[seqid] = set()
                    self.UprotToMnx[seqid].add(mnxr)

        with open(os.path.join(self.datadir, fl[1])) as f2:
            self.upclst = json.load(f2)

        with open(os.path.join(self.datadir, fl[2])) as f3:
            self.clstrep= json.load(f3)
        self.seqorg = seqOrganism(self.datadir, fl[3])
        self.tax = readTaxonomy(self.datadir, fl[4])

    def reacData(self, smf):
        """ Transitional downgrade mapping for mnx v3 to v2.
        In the future everything will be based on v3.
        """
        smiFile = os.path.join(self.datadir, smf[0])
        rxnRefFile = os.path.join(self.datadir, smf[1])
        rxnConsensus = os.path.join(self.datadir, smf[2])
        rxnProp = os.path.join(self.datadir, smf[3])
        rxnBrenda = os.path.join(self.datadir, smf[4])
        rxnSabiork = os.path.join(self.datadir, smf[5])
        rxnv3 = None
        if len(smf) > 6:
            rxnv3 = os.path.join(self.datadir, smf[6])
        self.smir = {}
        if os.path.exists(smiFile):
            self.smir = reactionSmiles(smiFile)
        self.rxnref = {}
        if os.path.exists(rxnRefFile):
            self.rxnref = reactionXref(rxnRefFile, rxnBrenda, rxnSabiork,rxnv3)
        self.rxndir = {}
        if os.path.exists(rxnConsensus):
            self.rxndir = readRxnCons(rxnConsensus)
        self.ecrxn = {}
        self.rxnec = {}
        if os.path.exists(rxnProp):
            self.ecrxn, self.rxnec = readRxnProp(rxnProp)
        self.ecsmi = ecSmiles(self.ecrxn, self.smir, self.rxnref)

def availableFingerprints():
    fp = quickRsim.fingerprint()
    return sorted(fp, key=lambda x: fp[x][4])

def sanitizeRxn(rxninfo, outrxn):
    """ It works both with the smiles string or a rxn file """
    try:
        if os.path.exists(rxninfo):
            if open(rxninfo).readline().startswith('$RXN'):
                rxn = AllChem.ReactionFromRxnFile(rxninfo)
            else:
                smarts =  open(rxninfo).readline().rstrip()
                try:
                    rxn = AllChem.ReactionFromSmarts(smarts, useSmiles=True)
                except:
                    rxn = AllChem.ReactionFromSmarts(smarts)
        smi = AllChem.ReactionToSmiles(rxn)
        with open(outrxn+'.smi', 'w') as handler:
            handler.write(smi)
        mdl = AllChem.ReactionToRxnBlock(rxn)
        with open(outrxn+'.rxn', 'w') as handler:
            handler.write(mdl)
        return smi
    except:
        return ''

def sanitizeSmarts(smarts, outrxn):
    """ It works both with the smiles string or a rxn file """
    try:
        try:
            rxn = AllChem.ReactionFromSmarts(smarts, useSmiles=True)
        except:
            rxn = AllChem.ReactionFromSmarts(smarts)
        smi = AllChem.ReactionToSmiles(rxn)
        with open(outrxn+'.smi', 'w') as handler:
            handler.write(smi)
        mdl = AllChem.ReactionToRxnBlock(rxn)
        with open(outrxn+'.rxn', 'w') as handler:
            handler.write(mdl)
        return smi
    except:
        return ''

def display_reaction(rxninfo, outfolder, outname, marvin=False):
    """ It works both with the smiles string or a rxn file """
    if marvin:
        try:
            cmd = ['molconvert', 'svg:w500', rxninfo]
            job = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            out, err = job.communicate()
            outimage = os.path.join(outfolder, outname+'.svg')
            return outimage, (600, 400)
        except:
            return '', (0,0)
    else:
        try:
            if os.path.exists(rxninfo):
                if open(rxninfo).readline().startswith('$RXN'):
                    rxn = AllChem.ReactionFromRxnFile(rxninfo)
                else:
                    smarts =  open(rxninfo).readline()
                    rxn = AllChem.ReactionFromSmarts(smarts)
                outimage = os.path.join(outfolder, outname+'.png')
                im = Draw.ReactionToImage(rxn)
                size = im.size
                im.save(outimage)
                return outimage, size
            else:
                rxn = AllChem.ReactionFromSmarts(rxninfo)
                outimage = os.path.join(outfolder, outname+'.png')
                im = Draw.ReactionToImage(rxn)
                size = im.size
                im.save(outimage)
                return outimage, size
        except:
            return '', (0,0)

def seqOrganism(datadir, fileSeqOrg):
    seqorg = {}
    for line in open(os.path.join(datadir, fileSeqOrg)):
        row = line.rstrip().split('\t')
        seqorg[ row[0] ] = (row[1], row[2])
    return seqorg

def readTaxonomy(datadir, fileLineage):
    tax = {}
    with open(os.path.join(datadir, fileLineage)) as handler:
        for row in csv.reader(handler):
            tax[row[0]] = row
    return tax

def taxDistance(tax, host, target):
    if host in tax and target in tax:
        hostLineage = set(tax[host])
        targetLineage = set(tax[target])
        distance = 1 + len(hostLineage ^ targetLineage)
        return distance
    else:
        return '-'

def seqScore(newscore=None):
    import string
    # Initial score
    vdict = {
        string.ascii_uppercase[9]: ('Reaction similarity:', 100.0, True),
        string.ascii_uppercase[8]: ('Sequence conservation:', 1.0, True),
        string.ascii_uppercase[4]: ('Sequence taxonomic distance:', -1.0, True),
        string.ascii_uppercase[7]: ('Uniprot protein evidence:', -0.1, True),
        string.ascii_uppercase[12]: ('Percentage helices:', 0.0, False),
        string.ascii_uppercase[13]: ('Percentage sheets:', 0.0, False),
        string.ascii_uppercase[14]: ('Percentage turns:', 0.0, False),
        string.ascii_uppercase[15]: ('Molecular weight:', 0.0, False),
        string.ascii_uppercase[16]: ('Isoelectric point:', 0.0, False),
        string.ascii_uppercase[17]: ('Percentage polar amino acids:', 0.0, False)
    }
    nvdict = vdict.copy()
    # Reference order (alternatively, perhaps easier to keep table order)
    clist = [string.ascii_uppercase[x] for x in [9,8,4,7,12,13,14,15,16,17]]
    # Update score if given and well-formed
    update = False
    if newscore is not None:
        nscore = [x[0] for x in newscore]
        if len( set(nscore) & set(clist) ) > 0:
            for val in vdict:
                nvdict[val] = (vdict[val][0], vdict[val][1], False)
            for val in newscore:
                if val[0] in nvdict:
                    try:
                        newval = float(val[1])
                        nvdict[val[0]] = (nvdict[val[0]][0], newval, True)
                        update = True
                    except:
                        continue
    if not update:
        nvdict = vdict
    score = []
    for x in clist:
        score.append( (x,) + nvdict[x] )
    return score

def updateScore(csvfile, score):
    """ Add or update score column and reorder """
    import string
    head, rows = read_csv(csvfile)
    data = pd.read_csv(csvfile)
    data.index = data.index + 1
    cols = data.columns.tolist()
    sco = pd.Series(np.zeros(len(data[cols[0]])), index=data.index)
    if 'Score' not in cols:
        data['Score'] = sco
        cols = ['Score'] + cols
        data = data[cols]
    colk = list(string.ascii_uppercase)
    for sc in score:
        try:
            coln = colk.index(sc[0])
            val = sc[2]
            checked = sc[3]
            if checked:
                sco += val * data.iloc[:,coln]
        except:
            continue
    data['Score'] = sco
    data = data.sort_values('Score', ascending=False)
    updateMSA(os.path.dirname(csvfile), [[v] for v in data['Seq. ID']])
    data = data.reset_index(drop=True)
    data.index = data.index + 1
    data.rename_axis('Select', axis="columns")
    data.to_csv(csvfile, quoting=csv.QUOTE_ALL, index=False)
    return data

def readFasta(datadir, fileFasta, limit=None):

    from Bio import SeqIO

    sequence = {}
    descriptions = {}
    fulldescriptions = {}
    osource = {}
    pexistence = {}
    names = []
    seen = set()
    seen_add = seen.add

    for seq_record in SeqIO.parse(os.path.join(datadir, fileFasta), "fasta"):
        ming = seq_record.id
        try:
            idonly = re.search(r'\|(.*?)\|',ming)
            x = idonly.group(1)
        except:
            x = ming
        fulldesc = seq_record.description
        desc = fulldesc.rsplit('OS=')[0]
        try:
            shortdesc = " ".join(desc.split()[1:])
        except:
            shortdesc = desc
        try:
            orgsource = fulldesc.rsplit('OS=')[1]
        except:
            orgsource = '-'
        shortos = orgsource.rsplit('GN=')[0]
        if ',' in shortdesc:
            y = shortdesc.replace(",", ";")
        else:
            y = shortdesc
        pe = orgsource.rsplit('PE=')[1].split(' ')[0]
        if x not in seen:
            names.append(x)
            seen_add(x)
        myseq = seq_record.seq
        sequence[x]=(myseq)
        descriptions[x]= y
        osource[x]=shortos
        fulldescriptions[x] = fulldesc
        pexistence[x] = pe
        # A practical limit hardcoded
        if limit is not None:
            try:
                if len(sequence) > limit:
                    break
            except:
                continue

    return (sequence, names, descriptions, fulldescriptions, osource, pexistence)

def readRxnCons(consensus):

    f = open(consensus, 'r')

    MnxDir = {}

    for line in f:
        splitdata = line.split()
        Mnx = splitdata[1]
        dirxn = splitdata[2]
        MnxDir[Mnx] = dirxn

    return (MnxDir)

def getMnxSim(rxnInput, datadir, outdir, drxn=0, fp='RDK', pc=None):
    """ Commmand line arguments of quickRsim """

    args = [datadir, fp] + rxnInput + ['-out', os.path.join(outdir,'results_quickRsim.txt')]
    quickRsim.run( quickRsim.arguments(args), pc )
    MnxSim = {}
    MnxDirPref = readRxnCons(os.path.join(datadir, "rxn_consensus_20160612.txt"))
    MnxDirUsed = {}
    EcNumber = {}

    if drxn==1:
        fileout = open(os.path.join(outdir, "results_quickRsim.txt"), 'r')
        for line in fileout:
            splitdata = line.split()
            S1 = splitdata[2]
            S2 = splitdata[3]
            SMILES = splitdata[4]
            EC = ''
            if len(splitdata) > 5:
                EC = splitdata[5]
            Mnx = splitdata[1]
            EcNumber[Mnx] = EC
            try:
                direction = MnxDirPref[Mnx]
                if direction == '1':
                    MnxSim[Mnx]=S1
                    MnxDirUsed[Mnx]= '1'
                elif direction == '-1':
                    MnxSim[Mnx]=S2
                    MnxDirUsed[Mnx]= '-1'
                else:
                    if S1 >= S2:
                        MnxSim[Mnx]=S1
                        MnxDirUsed[Mnx]= '1'
                    elif S2 > S1:
                        MnxSim[Mnx]=S2
                        MnxDirUsed[Mnx]= '-1'
            except KeyError:
                MnxDirPref[Mnx] = "N/A"   # for missing Keys?
                if S1 >= S2:
                    MnxSim[Mnx]=S1
                else:
                    MnxSim[Mnx]=S2
        fileout.close()

        return (MnxSim, MnxDirPref, MnxDirUsed, SMILES, EcNumber)

    else:

        fileout = open(os.path.join(outdir, "results_quickRsim.txt"), 'r')
        for line in fileout:
            splitdata = line.split()
            S1 = splitdata[2]
            S2 = splitdata[3]
            SMILES = splitdata[4]
            EC = ''
            if len(splitdata) > 5:
                EC = splitdata[5]
            Mnx = splitdata[1]
            EcNumber[Mnx] = EC
            if S1 > S2:
                MnxDirUsed[Mnx]='1'
                MnxSim[Mnx]=S1
            elif S2 > S1:
                MnxDirUsed[Mnx]='-1'
                MnxSim[Mnx]=S2
            elif S2 == S1:
                MnxDirUsed[Mnx]='0'
                MnxSim[Mnx]=S2
        fileout.close()

        return (MnxSim, MnxDirPref, MnxDirUsed, SMILES, EcNumber)

def readRxnProp(rxnprop):
    ecrxn = {}
    rxnec = {}
    with open(rxnprop) as handler:
        for line in handler:
            if line.startswith('#'):
                continue
            row = line.rstrip().split('\t')
            rxnid = row[0]
            ec = row[4].split(';')
            if len(ec) > 0:
                for e in ec:
                    if e == '':
                        continue
                    if e not in ecrxn:
                        ecrxn[e] = set()
                    ecrxn[e].add(rxnid)
                    if rxnid not in rxnec:
                        rxnec[rxnid] = set()
                    rxnec[rxnid].add(e)
    return ecrxn, rxnec

def reactionXref(rxnRefFile, rxnBrenda, rxnSabiork, rxnv3=None):
    rxnref = {}
    for xref in (rxnRefFile, rxnBrenda, rxnSabiork):
        with open(xref) as handler:
            for line in handler:
                if line.startswith('#'):
                    continue
                row = line.rstrip().split('\t')
                rxnref[row[0]] = row[1]
                rxnref['mnx:'+row[1]] = row[1]
    # Downgrade v3 to v2
    if rxnv3 is not None:
        for line in open(rxnv3):
            if line.startswith('#'):
                continue
            row = line.rstrip().split('\t')
            try:
                db, rid = row[0].split(':')
                if db == 'deprecated':
                    rxnref['mnx:'+row[1]] = rid
            except:
                continue
    return rxnref

def ecSmiles(ecrxn, rsmi, rxnref):
    smiec = {}
    for ec in ecrxn:
        for r in set(ecrxn[ec]) & set(rsmi):
            s = rsmi[r][0]
            if s not in smiec:
                smiec[s] = set()
            smiec[s].add(ec)
    for ec in ecrxn:
        """ Take the shortest SMARTS being unique for the EC class if possible """
        hits = sorted(set(ecrxn[ec]) & set(rsmi), key=lambda x: len(rsmi[x][0]))
        for h in hits:
            if len(rsmi[h]) > 0:
                rxnid = 'ec:'+ec
                rxnref[rxnid] = h
                if len(smiec[rsmi[h][0]]) == 1:
                       break

def reactionSmiles(rxnSmilesFile):
    rsmi = {}
    with open(rxnSmilesFile) as handler:
        header = handler.readline()
        for line in handler:
            row = line.rstrip().split(',')
            if len(row) > 1:
                rid = row[0]
                rs1 = row[1]
                lr = rs1.split('>>')
                rs2 = lr[1]+'>>'+lr[0]
                rsmi[rid] = (rs1, rs2)

    return rsmi

def pepstats(file, outdir):
    outfile = os.path.join(outdir, "results.pepstats")
    args = ("pepstats -sequence {0} -outfile ".format(file) + outfile)
    os.system(args)

    f = open(outfile, "r")

    hydrop = {}
    weight = {}
    isoelec = {}
    polar = {}

    for line in f:
        if "PEPSTATS of" in line:
            splitdata = line.split()
            seq = splitdata[2].split('_')[0]
        elif "Molecular weight = " in line:
            splitdata = line.split()
            w = splitdata[3]
            weight[seq] = w
        elif "Isoelectric Point = " in line:
            splitdata = line.split()
            i = splitdata[3]
            isoelec[seq] = i
        elif "Polar	" in line:
            splitdata = line.split()
            percent = splitdata[3]
            polar[seq] = percent
            seq
    return (hydrop, weight, isoelec, polar)

def noAmbiguousSeqs(infile, outfile):
    """ Remove ambigous amino acid codes """
    from Bio.Data.IUPACData import protein_letters_1to3, extended_protein_values
    from Bio import SeqIO
    from Bio.Seq import Seq
    from Bio.SeqRecord import SeqRecord
    from Bio.Alphabet import generic_protein

    newrecords = []
    for record in SeqIO.parse(infile, "fasta"):
        newseq = ''
        for aa in record.seq:
            if aa not in protein_letters_1to3:
                newseq += extended_protein_values[aa][0]
            else:
                newseq += aa
        newrecords.append( SeqRecord(Seq(newseq), id = record.id, description = record.description) )
    SeqIO.write(newrecords, outfile, "fasta")

def garnier(file, outdir):
    fixfile = file+'.fix.fasta'
    noAmbiguousSeqs(file, fixfile)
    outfile = os.path.join(outdir, "garnier.txt")

    args = ("garnier -sequence {0} -outfile ".format(fixfile) + outfile)
    os.system(args)

    f = open(outfile, "r")

    helices = {}
    sheets = {}
    turns = {}
    coils = {}

    for line in f:
        if "Sequence:" in line:
            splitdata = line.split()
            seq = splitdata[2].split('_')[0]
        elif "percent:" in line:
            percents = line.split()
            h = percents[3]
            e = percents[5]
            t = percents[7]
            c = percents[9]
            helices[seq] = h
            sheets[seq] = e
            turns[seq] = t
            coils[seq] = c

    return (helices, sheets, turns, coils)

def updateMSA(outdir, sortrows):
    """ Update the fasta MSA file in the order of the scores """
    align_fasta = os.path.join(outdir, "sequences_aln.fasta")
    if not os.path.exists(align_fasta):
        return
    fasta = {}
    with open(align_fasta) as handler:
        for line in handler:
            if line.startswith('>'):
                splitdata = line.split()
                upid = splitdata[0][1:]
                if len(upid.split('|')) > 1:
                    upid = upid.split('|')[1]
                elif len(upid.split('_')) > 1:
                    upid = upid.split('_')[0]
            if upid not in fasta:
                fasta[upid] = []
            fasta[upid].append(line)

    if len(fasta) == 0:
        return
    outfile = align_fasta
    with open(outfile, 'w') as hw:
        for row in sortrows:
            try:
                pid = row[0]
                if pid in fasta:
                    for seq in fasta[pid]:
                        hw.write(seq)
            except:
                continue

def doMSA(finallistfile, outdir):
    outfile = os.path.join(outdir, "sequences.score_ascii")
    outfile_html = os.path.join(outdir, "sequences.score_ascii.score_html")
    outfile_aln = os.path.join(outdir, "sequences.score_ascii.fasta_aln")
    align_html = os.path.join(outdir, "sequences_score.html")
    align_fasta = os.path.join(outdir, "sequences_aln.fasta")
    treefile = os.path.join(outdir, "sequences.dnd")
    args = ("t_coffee -in {0} -mode quickaln -output=score_ascii,fasta_aln,score_html -outfile ".format(finallistfile) +outfile+ " -newtree "+treefile)
    os.system(args)
    if os.path.exists(outfile_html):
        os.rename(outfile_html, align_html)
    if os.path.exists(outfile_aln):
        os.rename(outfile_aln, align_fasta)

    f = open(outfile, "r")

    cons = {}
    for line in f:
        if "   :  " in line:
            splitdata = line.split()
            upid = splitdata[0]
            if len(upid.split('|')) > 1:
                   upid = upid.split('|')[1]
            elif len(upid.split('_')) > 1:
                   upid = upid.split('_')[0]
            score = splitdata[2]
            cons[upid] = score

    return cons

def read_csv(csvfile):
    rows = []
    if os.path.exists(csvfile):
        with open(csvfile) as handler:
            cv = csv.reader(handler)
            head = next(cv)
            for row in cv:
                rows.append(row)
    return head, rows

def write_csv(csvfilepath, head, rows):
    with open (csvfilepath, 'w') as csvfile:
        writer = csv.writer(csvfile, delimiter=',', quotechar='"', quoting=csv.QUOTE_ALL)
        writer.writerow(head)
        for r in rows:
            writer.writerow(r)

def sort_rows(rows, columns):
    for i in range(len(columns), 0, -1):
        key = columns[i-1]
        if key != 0:
            if key < 0:
                try:
                    rows.sort(key = lambda x: -float(x[abs(key)-1]))
                except:
                    rows.sort(key = lambda x: x[abs(key)-1], reverse=True)
            else:
                try:
                    rows.sort(key = lambda x: float(x[key-1]))
                except:
                    rows.sort(key = lambda x: x[key-1])
    return rows

def write_fasta(fastaFile, targets, pc, short=False, info=False, maxlength=25):
    def shorten(s, maxlength=maxlength):
        ns = re.sub('\s+', '_', s)
        ns = re.sub('\|','_', ns)
        if len(ns) > maxlength:
            ns = ''.join(list(ns)[0:maxlength])+'...'
        return ns
    with open(fastaFile, "w") as f:
        for t in targets:
            try:
                seq = pc.sequence[t]
                if info:
                    desc = ''
                    try:
                        desc = pc.descriptions[t]
                    except:
                        pass
                    org = ''
                    try:
                        org = pc.seqorg[t][1]
                    except:
                        pass
                    ecl = set()
                    if t in pc.UprotToMnx:
                        for r in pc.UprotToMnx[t]:
                            if r in pc.rxnec:
                                ecl |= pc.rxnec[r]
                    ec = ';'.join( sorted(ecl) )
                    st = '_'.join( [t, shorten(desc,30), shorten(org), ec] )
                    print ('>{0} \n{1}'.format(st, seq), file=f)
                elif short:
                    print ('>{0} \n{1}'.format(t, seq), file=f)
                else:
                    fdesc = pc.fulldescriptions[t]
                    print ('>{0} \n{1}'.format(fdesc, seq), file=f)
            except KeyError:
                pass

def short_fasta(fastafile):
    dirname = os.path.dirname(fastafile)
    basename = os.path.basename(fastafile)
    shortfile = 'short_'+basename
    shortname = os.path.join(dirname, shortfile)
    with open(fastafile) as fasta, open(shortname, 'w') as fastaw:
        for line in fasta:
            if line.startswith('>'):
                head = line.split('OS=')
                if len(head[0].split('|')) >= 2:
                    seqid = head[0].split('|')[1]
                    line = '>'+seqid+'\n'
            fastaw.write(line)
    return shortfile

def extend_sequences(initialfastafile, fastafile, workfolder, noMSA):
    """ Extend the fasta file """
    csvfile = os.path.join(workfolder, 'selenzy_results.csv')
    shortfile = short_fasta(os.path.join(workfolder, fastafile))
    try:
        # TO do: Check that does not assume Uniprot format...
        (sequence, names, descriptions,
         fulldescriptions, osource, pexistence) = readFasta(workfolder, fastafile, limit=1000)

        if len(sequence) == 0:
            return csvfile

        (hydrop, weight, isoelec, polar, helices,
         sheets, turns, coils) = sequence_properties(os.path.join(workfolder, shortfile))
        fasta = open(os.path.join(workfolder, initialfastafile)).readlines()
        fasta.extend(open(os.path.join(workfolder, fastafile)).readlines())
        with open(os.path.join(workfolder, initialfastafile), 'w') as handler:
            for line in fasta:
                handler.write(line)
        shortfile = short_fasta(os.path.join(workfolder, initialfastafile))
        if not noMSA:
            cons = conservation_properties(os.path.join(workfolder, shortfile))
        else:
            cons = ({},{})

    except:
        return csvfile
    # Extend the csvfile
    head, rows = read_csv(csvfile)
    # Update the conservation scores
    if not noMSA:
        for r in rows:
            seqid = r[head.index('Seq. ID')]
            if seqid in cons:
                r[head.index('Consv. Score')] = cons[seqid]
    for k in range(0, len(fulldescriptions)):
        n = names[k]
        try:
            conservation = cons[n]
        except:
            conservation = 0
        try:
            h = helices[n]
            e = sheets[n]
            t = turns[n]
            c = coils[n]
        except:
            h = '-'
            e = '-'
            t = '-'
            c = '-'
        try:
            w = weight[n]
            i = isoelec[n]
            pol = polar[n]
        except:
            w = '-'
            i = '-'
            pol = '-'
        try:
            description = descriptions[n]
            source = osource[n]
            ext = pexistence[n]
        except:
            description = '-'
            source = '-'
        row = [0.0, n, description, source, -1,
               '-', '-', ext, conservation, 1, 1, 1,
               h, e, t, c, w, i, pol, '-', '-']
        rows.append( row )
    write_csv(csvfile, head, rows)
    return csvfile

def sequence_properties(fastaShortNameFile):
    #analysis of final list of sequences
    (hydrop, weight, isoelec, polar) = pepstats(fastaShortNameFile, os.path.dirname(fastaShortNameFile))
    (helices, sheets, turns, coils) = garnier(fastaShortNameFile,  os.path.dirname(fastaShortNameFile))
    return hydrop, weight, isoelec, polar, helices, sheets, turns, coils

def conservation_properties(fastaFile):
    cons = doMSA(fastaFile,  os.path.dirname(fastaFile))
    return cons

def analyse(rxnInput, targ, datadir, outdir, csvfilename, pdir=0, host='83333', NoMSA=False, pc=None, fp='RDK'):

    datadir = os.path.join(datadir)
    outdir = os.path.join(outdir)

#    rxnname = os.path.splitext(rxn)[0]
#    csvname = rxnname.rsplit('/', 1)[-1]

    if csvfilename:
        csvfilename = csvfilename
    else:
        csvfilename = "results_selenzy.csv"
    if pc is None:
        pc = readData(datadir)

    print ("Running quickRsim...")
    try:
        (MnxSim, MnxDirPref, MnxDirUsed, Smiles, EcNumber) = getMnxSim(rxnInput, datadir, outdir, pdir, fp, pc)
    except:
        return False, pc

    sequence = pc.sequence
    names = pc.names
    descriptions = pc.descriptions
    fulldescriptions = pc.fulldescriptions
    osource = pc.osource
    pexistence = pc.pexistence
    seqorg = pc.seqorg
    tax = pc.tax
    MnxToUprot = pc.MnxToUprot
    upclst = pc.upclst
    clstrep = pc.clstrep
    smir = pc.smir

    list_mnx = sorted(MnxSim, key=MnxSim.__getitem__, reverse=True)  #allow user to manipulate window of initial rxn id list
    print ("Creating initial MNX list...")
    targets = set()
    UprotToMnx = {}

    # first creating fasta file, f, for further data extraction
    for x in list_mnx:
        up = MnxToUprot.get(x)
        if up is not None:
            for y in up:
                if len(targets) >= int(targ):    # allow user to manipulate desired number of entries for resulting table
                    break
                else:
                    UprotToMnx[y] = x
                    try:
                        targets.add(y)
                    except KeyError:
                        pass
            else:
                continue
            break
    fastaFile = os.path.join(outdir, "sequences.fasta")
    write_fasta(fastaFile, targets, pc)
    # Avoid issues with sequence ids
    fastaShortNameFile = os.path.join(outdir, "seqids.fasta")
    fastaInfoNameFile = os.path.join(outdir, "seqinfo.fasta")
    write_fasta(fastaShortNameFile, targets, pc, short=True)
    write_fasta(fastaInfoNameFile, targets, pc, info=True)

    (hydrop, weight, isoelec, polar, helices, sheets, turns, coils) = sequence_properties(fastaShortNameFile)
    if not NoMSA:
        cons = conservation_properties(fastaInfoNameFile)
    else:
        cons = {}

    print ("Acquiring sequence properties...")
    # final table, do all data and value storing before this!
    tdist = {}
    rows = []

    for y in targets:
        try:
            # Essential sequence information
            desc = descriptions[y]
            fdesc = fulldescriptions[y]
            org = osource[y]
            ext = pexistence[y]
            mnx = UprotToMnx[y]
            rxnsimpre = float(MnxSim[mnx])
            rxnsim = float("{0:.5f}".format(rxnsimpre))

            # Non-essential sequence information
            try:
                cn = upclst.get(y)
                repid = clstrep[cn]
            except:
                cn = 0
                repid = y
            try:
                conservation = float(cons[y])
            except:
                conservation = 0.0
            try:
                ecid = EcNumber[mnx]
            except:
                ecid = ''
            try:
                h = helices[y]
                e = sheets[y]
                t = turns[y]
                c = coils[y]
            except:
                h = '-'
                e = '-'
                t = '-'
                c = '-'
            try:
                w = weight[y]
                i = isoelec[y]
                pol = polar[y]
            except:
                w = '-'
                i = '-'
                pol = '-'
            try:
                rxndirused = MnxDirUsed[mnx]
                rxndirpref = MnxDirPref[mnx]
            except:
                rxndirused = 1
                rxndirpred = 1

            mnxSmiles = ''
            if mnx in smir:
                if rxndirused == 1:
                    mnxSmiles = smir[mnx][0]
                else:
                    mnxSmiles = smir[mnx][1]
            if org not in tdist:
                if y in seqorg:
                    tdist[org] = taxDistance(tax, host, seqorg[y][0])
                else:
                    tdist[org] = '-'

            rows.append( (y, desc, org, tdist[org], mnx, ecid, ext, conservation, rxnsim, rxndirused, rxndirpref, h, e, t, c, w, i, pol, Smiles, mnxSmiles) )

        except KeyError:
            pass
    sortrows = sort_rows(rows, (-10, -9, 4) )
    updateMSA(outdir, sortrows)

    head = ('Seq. ID','Description', 'Organism Source', 'Tax. distance', 'Rxn. ID', 'EC Number', 'Uniprot protein evidence', 'Consv. Score',
            'Rxn Sim.', "Direction Used", "Direction Preferred",
            '% helices', '% sheets', '% turns', '% coils', 'Mol. Weight', 'Isoelec. Point', 'Polar %','Query', 'Hit')

    write_csv(os.path.join(outdir, csvfilename), head, sortrows)

    print ("CSV file created.")
    return True, pc

max_stars_repo_name: brsynth/rpSelenzyme

max_issues_repo_name: brsynth/rpSelenzyme

max_forks_repo_name: brsynth/rpSelenzyme


text:
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Core TensorFlow types."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import textwrap

from typing import Union
import numpy as np

from tensorflow.python.types import doc_typealias
from tensorflow.python.util.tf_export import tf_export

# TODO(mdan): Consider adding ABC once the dependence on isinstance is reduced.
# TODO(mdan): Add type annotations.

class Tensor(object):
  """The base class of all dense Tensor objects.

  A dense tensor has a static data type (dtype), and may have a static rank and
  shape. Tensor objects are immutable. Mutable objects may be backed by a Tensor
  which holds the unique handle that identifies the mutable object.
  """

  @property
  def dtype(self):
    pass

  @property
  def shape(self):
    pass

class Symbol(Tensor):
  """Symbolic "graph" Tensor.

  These objects represent the output of an op definition and do not carry a
  value.
  """
  pass

class Value(Tensor):
  """Tensor that can be associated with a value (aka "eager tensor").

  These objects represent the (usually future) output of executing an op
  immediately.
  """

  def numpy(self):
    pass

# TODO(rahulkamat): Add missing types that are convertible to Tensor.
TensorLike = Union[Tensor, int, float, bool, str, complex, tuple, list,
                   np.ndarray]
doc_typealias.document(
    obj=TensorLike,
    doc=textwrap.dedent("""\
      Union of all types that can be converted to a `tf.Tensor` by `tf.convert_to_tensor`.

      This definition may be used in user code. Additional types may be added
      in the future as more input types are supported.

      Example:

      ```
      def foo(x: TensorLike):
        pass
      ```

      This definition passes static type verification for:

      ```
      foo(tf.constant([1, 2, 3]))
      foo([1, 2, 3])
      foo(np.array([1, 2, 3]))
      ```
      """),
)
tf_export("types.experimental.TensorLike").export_constant(
    __name__, "TensorLike")

ext: py

reason: import numpy


text:
# -*- coding: utf-8 -*-
"""Type of the elements in each table of the database.

The database is formed by numpy arrays and the type of the elements of each one
is specified in this module as a module level variable.

Attributes:
    departement_dtype (str): The definition of the numpy dtype for the
        elements of the departement table.
    postal_dtype (str): The definition of the numpy dtype for the elements of
        the postal table.
    commune_dtype (str): The definition of the numpy dtype for the elements of
        the commune table.
    voie_dtype (str): The definition of the numpy dtype for the elements of
        the voie table.
    localisation_dtype (str): The definition of the numpy dtype for the
        elements of the localisation table.
    kdtree_dtype (str): The definition of the numpy dtype for the elements of
        the kdtree table.
    dtypes (:obj:`dict` of :obj:`str`): A python dictionary to easily access
        the dtypes definitions.

"""

import numpy as np

departement_dtype = np.dtype([
    ('code', 'U3'),
    ('start', 'int32'),
    ('end', 'int32'),
])

postal_dtype = np.dtype([
    ('code', 'int32'),
    ('start', 'int32'),
    ('end', 'int32'),
    ('ref_id', 'int32'),
])

commune_dtype = np.dtype([
    ('normalise', 'U32'),
    ('nom', 'U32'),
    ('code_insee', 'U5'),
    ('longitude', 'int32'),
    ('latitude', 'int32'),
    ('start', 'int32'),
    ('end', 'int32'),
    ('ref_id', 'int32'),
])

voie_dtype = np.dtype([
    ('normalise', 'U47'),
    ('nom', 'U65'),
    ('longitude', 'int32'),
    ('latitude', 'int32'),
    ('start', 'int32'),
    ('end', 'int32'),
    ('ref_id', 'int32'),
])

localisation_dtype = np.dtype([
    ('numero', 'int16'),
    ('repetition', 'U3'),
    ('longitude', 'int32'),
    ('latitude', 'int32'),
    ('ref_id', 'int32'),
])

kdtree_dtype = np.dtype([
    ('longitude', 'int32'),
    ('latitude', 'int32'),
    ('limit_left', 'int32'),
    ('limit_right', 'int32'),
    ('limit_bottom', 'int32'),
    ('limit_top', 'int32'),
    ('dimension', 'int8'),
    ('left', 'int32'),
    ('right', 'int32'),
    ('ref_id', 'int32'),
])

dtypes = {
    'departement': departement_dtype,
    'postal': postal_dtype,
    'commune': commune_dtype,
    'voie': voie_dtype,
    'localisation': localisation_dtype,
    'commune_index': 'int32',
    'postal_index': 'int32',
    'voie_index': 'int32',
    'kdtree': kdtree_dtype
}

max_stars_repo_name: Mega-Touk/Geocoding

max_issues_repo_name: Mega-Touk/Geocoding

max_forks_repo_name: Mega-Touk/Geocoding


text:
import gzip
import pickle
import numpy as np
import torch
from torch import nn
import random
from tqdm import tqdm
import os
import subprocess
import collections
import igraph
import argparse
import pdb
import pygraphviz as pgv
import sys
from PIL import Image

# create a parser to save graph arguments
cmd_opt = argparse.ArgumentParser()
graph_args, _ = cmd_opt.parse_known_args()

'''gobal variables'''
CIRCUIT_NODE_NUM = 5
CIRCUIT_NODE_TYPE = 26

'''dataset generation'''

def clamp(minn, maxn, n):
    return max(min(maxn, n), minn)

def sample_topo_vector(edge_num):
    topo_vector = []
    for e in range(edge_num):
        if e <= 1:
            topo_vector.append(np.random.randint(0, 11))
        elif e == 2:
            topo_vector.append(np.random.randint(0, 25))
        else:
            topo_vector.append(np.random.randint(0, 5))
    return topo_vector

def sample_full_random(edge_num):
    topo_vector = []
    for e in range(edge_num):
        topo_vector.append(np.random.randint(0, 25))
    return topo_vector

def vector2row(topo_vector):
    CIRCUIT_DAG = []
    for i in range(CIRCUIT_NODE_NUM):
        if i == 0:
            CIRCUIT_DAG.append([0])  # node Vin
        elif i == 1:
            CIRCUIT_DAG.append([1, 5])  # node 1
        elif i == 2:
            CIRCUIT_DAG.append([1, topo_vector[0], 6])  # node 2
        elif i == 3:
            CIRCUIT_DAG.append([2, topo_vector[1], topo_vector[2], 5])  # node Vout
        elif i == 4:
            CIRCUIT_DAG.append([3, 0, topo_vector[3], topo_vector[4], 0])  # node GND
        else:
            pass
    return CIRCUIT_DAG

def row2vector(row):
    topo = []
    for i in range(len(row)):
        if i == 2:
            topo.append(row[i][1])
        elif i == 3:
            topo.append(row[i][1])
            topo.append(row[i][2])
        elif i == 4:
            topo.append(row[i][2])
            topo.append(row[i][3])
        else:
            pass
    return np.array(topo)

def g2topo(g):
    # plot_DAG(g, './', 'fuck')
    row_ = []
    for vo in range(2, g.vcount() - 1):
        for vi in range(1, vo):
            _eid = g.get_eid(vi, vo)
            row_.append(int(g.es[_eid]['weight']))
    topo_id = [1, 3, 4, 7, 8]
    topo = []
    for i in topo_id:
        if i == 1 or i == 3:
            topo.append(clamp(0, 11, row_[i]))
        elif i == 4:
            topo.append(clamp(0, 25, row_[i]))
        else:
            topo.append(clamp(0, 5, row_[i]))
    return np.array(topo)

'''Data preprocessing'''

def one_hot(idx, length):
    idx = torch.LongTensor([idx]).unsqueeze(0)
    x = torch.zeros((1, length)).scatter_(1, idx, 1)
    return x

def reorder_data2(value):
    # value = np.load("%s.npy" % datafile)
    value_new = np.zeros_like(value)
    value_new[0] = value[0]
    value_new[1] = value[3]
    value_new[2] = value[1]
    value_new[3] = value[2]
    value_new[4] = value[4]
    value_new[5] = value[5]
    return value_new

def load_CIRCUIT_graphs(rand_seed=0):
    # load DAG format CIRCUITs to igraphs
    g_list = []
    max_n = 10  # no connection to gnd, delete node gnd
    num_edge_feature = 32
    # max_n = 5  # maximum number of nodes
    value1 = np.load("../circuit1/value1000.npy")
    # print(value[0])
    value2 = np.load("../circuit2/value1000.npy")
    # len = np.shape(value[0])
    # print(value[0])
    value1 = value1.reshape((1000, 5, 32))
    value2 = value2.reshape((1000, 6, 32))
    value2 = reorder_data2(value2)
    # value = value[:, :, :num_edge_feature]
    performance1 = np.load("../circuit1/performance1000.npy")
    performance2 = np.load("../circuit2/performance1000.npy")
    for (value_, performance_) in zip(value1, performance1):
        g = decode_CIRCUIT_to_igraph1(value_)
        g_list.append((g, performance_))
    for (value_, performance_) in zip(value2, performance2):
        g = decode_CIRCUIT_to_igraph2(value_)
        g_list.append((g, performance_))
    graph_args.nvt = CIRCUIT_NODE_TYPE  # original types + virtual start/end types
    graph_args.max_n = max_n  # maximum number of nodes
    graph_args.edge_feature = num_edge_feature
    ng = len(g_list)
    print('# node types: %d' % graph_args.nvt)
    random.Random(rand_seed).shuffle(g_list)
    return g_list[:int(ng * 0.9)], g_list[int(ng * 0.9):], graph_args

def decode_CIRCUIT_to_igraph1(value):
    g = igraph.Graph(directed=True)
    n = 7
    g.add_vertices(n)
    g.vs[0]['type'] = 0  # virtual start
    #g.vs[0]['color'] = 'red'
    #g.vs[0]['label'] = 'start'
    g.vs[1]['type'] = 6  # gm+
    #g.vs[1]['color'] = 'blue'
    #g.vs[1]['label'] = 'gm+'
    g.vs[2]['type'] = 7  # gm-
    #g.vs[2]['color'] = 'green'
    #g.vs[2]['label'] = 'gm-'
    g.vs[3]['type'] = 6  # gm+
    #g.vs[3]['color'] = 'blue'
    #g.vs[3]['label'] = 'gm+'
    g.vs[4]['type'] = 7  # gm-
    #g.vs[4]['color'] = 'green'
    #g.vs[4]['label'] = 'gm-'
    g.vs[5]['type'] = 11  # gm-_c_series
    #g.vs[5]['color'] = 'brown'
    #g.vs[5]['label'] = 'gm-_c_series'
    g.vs[6]['type'] = 1  # virtual end
    #g.vs[6]['color'] = 'yellow'
    #g.vs[6]['label'] = 'end'

    for i in range(n):
        if i == 0 or i == n - 1:
            g.vs[i]['param'] = np.zeros(32)
        else:
            g.vs[i]['param'] = value[i - 1]

    # for i in range(n):
    #    print(g.vs[i]['param'])
    g.add_edge(0, 1)
    g.add_edge(1, 2)
    g.add_edge(2, 3)
    g.add_edge(3, 6)
    g.add_edge(1, 4)
    g.add_edge(4, 6)
    g.add_edge(0, 5)
    g.add_edge(5, 6)

    return g

def decode_CIRCUIT_to_igraph2(value):
    g = igraph.Graph(directed=True)
    n = 8
    g.add_vertices(n)
    g.vs[0]['type'] = 0  # virtual start
    #g.vs[0]['color'] = 'red'
    #g.vs[0]['label'] = 'start'
    g.vs[1]['type'] = 6  # gm+
    #g.vs[1]['color'] = 'blue'
    #g.vs[1]['label'] = 'gm+'
    g.vs[2]['type'] = 6  # gm+
    #g.vs[2]['color'] = 'blue'
    #g.vs[2]['label'] = 'gm+'
    g.vs[3]['type'] = 7  # gm-
    #g.vs[3]['color'] = 'green'
    #g.vs[3]['label'] = 'gm-'
    g.vs[4]['type'] = 8  # gm+_r_series
    #g.vs[4]['color'] = 'grey'
    #g.vs[4]['label'] = 'gm+_r_series'
    g.vs[5]['type'] = 6  # gm+
    #g.vs[5]['color'] = 'blue'
    #g.vs[5]['label'] = 'gm+'
    g.vs[6]['type'] = 9  # gm+_c_series
    #g.vs[6]['color'] = 'purple'
    #g.vs[6]['label'] = 'gm+_c_series'
    g.vs[7]['type'] = 1  # virtual end
    #g.vs[7]['color'] = 'yellow'
    #g.vs[7]['label'] = 'end'

    for i in range(n):
        if i == 0 or i == n - 1:
            g.vs[i]['param'] = np.zeros(32)
        else:
            g.vs[i]['param'] = value[i - 1]

    g.add_edge(0, 1)
    g.add_edge(1, 2)
    g.add_edge(2, 7)
    g.add_edge(1, 3)
    g.add_edge(3, 5)
    g.add_edge(5, 7)
    g.add_edge(0, 4)
    g.add_edge(4, 5)
    g.add_edge(0, 6)
    g.add_edge(6, 7)

    return g

'''Network visualization'''

def plot_DAG(g, res_dir, name, pdf=False):
    # backbone: puts all nodes in a straight line
    file_name = os.path.join(res_dir, name + '.png')
    if pdf:
        file_name = os.path.join(res_dir, name + '.pdf')
    draw_network(g, file_name)
    return file_name

def draw_network(g, path):
    graph = pgv.AGraph(directed=True, strict=True, fontname='Helvetica', arrowtype='open')
    # add vertex
    for idx in range(0, g.vcount()):
        add_node(graph, g, idx)
    # add edge
    for idx in range(1, g.vcount()):
        for node in g.get_adjlist(igraph.IN)[idx]:
            graph.add_edge(node, idx)
    graph.layout(prog='dot')
    graph.draw(path)

def add_node(graph, g, node_id, shape='box', style='filled'):
    # if g.vs[node_id]['type'] == 0 and node_id == 0:
    if g.vs[node_id]['type'] == 0:
        g.vs[node_id]['color'] = 'red'
        g.vs[node_id]['label'] = 'start'
    elif g.vs[node_id]['type'] == 1:
        g.vs[node_id]['color'] = 'yellow'
        g.vs[node_id]['label'] = 'end'
    elif g.vs[node_id]['type'] == 6:
        g.vs[node_id]['color'] = 'blue'
        g.vs[node_id]['label'] = 'gm+'
    elif g.vs[node_id]['type'] == 7:
        g.vs[node_id]['color'] = 'green'
        g.vs[node_id]['label'] = 'gm-'
    elif g.vs[node_id]['type'] == 8:
        g.vs[node_id]['color'] = 'grey'
        g.vs[node_id]['label'] = 'gm+_r_series'
    elif g.vs[node_id]['type'] == 9:
        g.vs[node_id]['color'] = 'purple'
        g.vs[node_id]['label'] = 'gm+_c_series'
    elif g.vs[node_id]['type'] == 11:
        g.vs[node_id]['color'] = 'brown'
        g.vs[node_id]['label'] = 'gm-_c_series'
    else:
        g.vs[node_id]['color'] = 'white'
        g.vs[node_id]['label'] = 'new'
    graph.add_node(node_id, label=g.vs[node_id]['label'], color='black', fillcolor=g.vs[node_id]['color'],
                   shape=shape, style=style, fontsize=24)
    #if not (g.vs[node_id]['type'] == 0 and node_id != 0):
    #    graph.add_node(node_id, label=g.vs[node_id]['label'], color='black', fillcolor=g.vs[node_id]['color'],
    #                   shape=shape, style=style, fontsize=24)

'''Validity and novelty functions'''

def is_same_DAG(g0, g1):
    # Correct rate of edge type prediction
    all_e = 0
    right_e = 0
    # fix_e = [[1,2],[2,3],[3,4],[1,5],[4,5]]
    for vo in range(2, g0.vcount() - 1):
        for vi in g0.get_adjlist(igraph.IN)[vo]:
            eid0 = g0.get_eid(vi, vo)
            eid1 = g1.get_eid(vi, vo)
            g0_edge_type = g0.es['weight'][eid0]
            g1_edge_type = g1.es['weight'][eid1]
            if g0_edge_type == g1_edge_type:
                right_e = right_e + 1
            all_e = all_e + 1
    return right_e / all_e

def load_module_state(model, state_name):
    pretrained_dict = torch.load(state_name)
    model_dict = model.state_dict()

    # to delete, to correct grud names
    '''
    new_dict = {}
    for k, v in pretrained_dict.items():
        if k.startswith('grud_forward'):
            new_dict['grud'+k[12:]] = v
        else:
            new_dict[k] = v
    pretrained_dict = new_dict
    '''

    # 1. filter out unnecessary keys
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}

    # 2. overwrite entries in the existing state dict
    model_dict.update(pretrained_dict)
    # 3. load the new state dict
    model.load_state_dict(pretrained_dict)
    return

def decode_igraph_to_CIRCUIT(g):
    # decode an igraph to a flattend CIRCUIT string
    n = g.vcount()
    res = []
    adjlist = g.get_adjlist(igraph.IN)
    for i in range(1, n - 1):
        res.append(int(g.vs[i]['type']) - 2)
        row = [0] * (i - 1)
        for j in adjlist[i]:
            if j < i - 1:
                row[j] = 1
        res += row
    return ' '.join(str(x) for x in res)
    # return res

def decode_from_latent_space(
        latent_points, model, decode_attempts=500, n_nodes='variable', return_igraph=False):
    # decode points from the VAE model's latent space multiple attempts
    # and return the most common decoded graphs
    if n_nodes != 'variable':
        check_n_nodes = True  # check whether the decoded graphs have exactly n nodes
    else:
        check_n_nodes = False
    decoded_arcs = []  # a list of lists of igraphs
    pbar = tqdm(range(decode_attempts))
    for i in pbar:
        current_decoded_arcs = model.decode(latent_points)
        decoded_arcs.append(current_decoded_arcs)
        pbar.set_description("Decoding attempts {}/{}".format(i, decode_attempts))

    # We see which ones are decoded to be valid architectures
    valid_arcs = []  # a list of lists of strings
    if return_igraph:
        str2igraph = {}  # map strings to igraphs
    pbar = tqdm(range(latent_points.shape[0]))
    for i in pbar:
        valid_arcs.append([])
        for j in range(decode_attempts):
            arc = decoded_arcs[j][i]  # arc is an igraph
            if not check_n_nodes or check_n_nodes and arc.vcount() == n_nodes:
                cur = decode_igraph_to_CIRCUIT(arc)  # a flat circuit igraph string
                if return_igraph:
                    str2igraph[cur] = arc
                valid_arcs[i].append(cur)
        pbar.set_description("Check validity for {}/{}".format(i, latent_points.shape[0]))

    # select the most common decoding as the final architecture
    final_arcs = []  # a list of lists of strings
    pbar = tqdm(range(latent_points.shape[0]))
    for i in pbar:
        valid_curs = valid_arcs[i]
        aux = collections.Counter(valid_curs)
        if len(aux) > 0:
            arc, num_arc = list(aux.items())[np.argmax(aux.values())]
        else:
            arc = None
            num_arc = 0
        final_arcs.append(arc)
        pbar.set_description("Latent point {}'s most common decoding ratio: {}/{}".format(
            i, num_arc, len(valid_curs)))

    if return_igraph:
        final_arcs_igraph = [str2igraph[x] if x is not None else None for x in final_arcs]
        return final_arcs_igraph, final_arcs
    return final_arcs


text:
##################################################
# MIT License
#
# Copyright (c) 2019 Learning Equality
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
##################################################

from config import BULLET_THRESHOLD
from functools import reduce
from PIL import Image, ImageDraw
import cv2
import numpy as np
import re

class BoundingBox(object):
    def __init__(self, x1, y1, x2, y2):
        if x1 >= x2:
            print(x1, y1, x2, y2)
        assert x1 < x2
        assert y1 < y2
        self.x1 = x1
        self.y1 = y1
        self.x2 = x2
        self.y2 = y2

    def center(self):
        return (self.x1 + self.x2) / 2, (self.y1 + self.y2) / 2

    def area(self):
        return self.width() * self.height()

    def width(self):
        return self.x2 - self.x1

    def height(self):
        return self.y2 - self.y1

    def expanded(self, factor, axis="both"):
        assert axis in ["both", "x", "y"], "`axis` must be one of 'both', 'x', or 'y'"
        x1, y1, x2, y2 = self.x1, self.y1, self.x2, self.y2
        if axis in ["x", "both"]:
            width = self.width()
            x1 -= factor * width
            x2 += factor * width
        elif axis in ["y", "both"]:
            height = self.height()
            y1 -= factor * height
            y2 += factor * height
        return BoundingBox(int(x1), int(y1), int(x2), int(y2))

    def shrunk(self, factor, axis="both"):
        return self.expanded(-factor, axis=axis)

    def shifted(self, x, y):
        return BoundingBox(self.x1 + x, self.y1 + y, self.x2 + x, self.y2 + y)

    def __and__(self, other):
        # returns the intersection of the bounding boxes
        x1 = max(self.x1, other.x1)
        y1 = max(self.y1, other.y1)
        x2 = min(self.x2, other.x2)
        y2 = min(self.y2, other.y2)
        if (x2 <= x1) or (y2 <= y1):
            return None
        return BoundingBox(x1, y1, x2, y2)

    def __or__(self, other):
        # returns the union of the bounding boxes (box that contains both)
        x1 = min(self.x1, other.x1)
        y1 = min(self.y1, other.y1)
        x2 = max(self.x2, other.x2)
        y2 = max(self.y2, other.y2)
        return BoundingBox(x1, y1, x2, y2)

    def overlap(self, other, axis="both"):
        """
        Calculate the Intersection over Union (IoU) of two bounding boxes.
        Adapted from: https://stackoverflow.com/questions/25349178/calculating-percentage-of-bounding-box-overlap-for-image-detector-evaluation
        """
        assert axis in ["both", "x", "y"], "`axis` must be one of 'both', 'x', or 'y'"
        if axis == "x":
            self = BoundingBox(self.x1, 0, self.x2, 1)
            other = BoundingBox(other.x1, 0, other.x2, 1)
        elif axis == "y":
            self = BoundingBox(0, self.y1, 1, self.y2)
            other = BoundingBox(0, other.y1, 1, other.y2)

        intersection = self & other
        if intersection is None:
            return 0.0
        intersection_area = intersection.area()

        # compute the intersection over union by taking the intersection area and dividing it
        # by the sum of the two areas minus the intersection area
        return intersection_area / float(self.area() + other.area() - intersection_area)

    def get_subimage(self, img):
        return img[self.y1 : self.y2, self.x1 : self.x2, :].copy()

    def __contains__(self, item):
        intersection = self & item
        if intersection is None:
            return False
        return item.overlap(intersection) > 0.8

    def __str__(self):
        return "({}, {})/({}, {})".format(self.x1, self.y1, self.x2, self.y2)

    def __repr__(self):
        return "<BoundingBox: {}>".format(str(self))

class BoundingBoxSet(list):
    def __init__(self, *args, overlap_threshold=0.4, **kwargs):
        self.overlap_threshold = overlap_threshold
        super().__init__(*args, **kwargs)

    def get_outer_box(self):
        return reduce(lambda a, b: a | b, self)

    def __and__(self, other):
        # returns the bounding boxes that are in both sets (intersection)
        results = BoundingBoxSet(overlap_threshold=self.overlap_threshold)
        for box_a in self:
            for box_b in other:
                if box_a.overlap(box_b) > self.overlap_threshold:
                    results.append(box_a & box_b)
        return results

    def __or__(self, other):
        # returns the bounding boxes that are in either set (union)
        results = BoundingBoxSet(overlap_threshold=self.overlap_threshold)
        for box_a in self:
            for box_b in other:
                if box_a.overlap(box_b) > self.overlap_threshold:
                    results.append(box_a & box_b)
        return results

    def __add__(self, other):
        return self | other

    def __sub__(self, other):
        # returns the set of boxes in self but not in other
        results = BoundingBoxSet(overlap_threshold=self.overlap_threshold)
        for box_a in self:
            found = False
            for box_b in other:
                if box_a.overlap(box_b) > self.overlap_threshold:
                    found = True
                    break
            if not found:
                results.append(box_a)
        return results

    def __contains__(self, item):
        for box in self:
            if box in item and item in box:
                return True
        return False

    def deduplicate(self):
        unique = BoundingBoxSet([], overlap_threshold=self.overlap_threshold)
        for box in self:
            if box not in unique:
                unique.append(box)
        return unique

class Word(object):
    """
        Attributes:
            text: string of characters
            bounding_box: BoundingBox object with coordinates for word
    """

    bounding_box = None
    text = ""

    def __init__(self, text, bounding_box):
        self.bounding_box = bounding_box
        self.text = text

    def __repr__(self):
        return '<Word: "{}" @ {}>'.format(self.text, str(self.bounding_box))

class Line(object):
    """
        Attributes:
            words: a list of Words
            fontweight: float representing how bold the line of text is
    """

    words = None
    fontweight = None
    column_box = None

    def __init__(self, words, fontweight=None, column_box=None):
        self.fontweight = fontweight
        self.words = words or []
        self.column_box = column_box

    def add_word(self, word):
        self.words.append(word)

    def get_box(self):
        if not self.words:
            return None
        return BoundingBox(
            min(word.bounding_box.x1 for word in self.words),
            min(word.bounding_box.y1 for word in self.words),
            max(word.bounding_box.x2 for word in self.words),
            max(word.bounding_box.y2 for word in self.words),
        )

    def get_text(self):
        return " ".join([word.text for word in self.words])

    def get_indentation(self, word=None, units="col_width"):
        assert self.column_box, "line must have column_box set to calculate indentation"
        assert units in ["col_width", "pixels", "line_height"]
        if word is None:
            word = self.words[0]
        elif isinstance(word, int):
            word = self.words[word]
        indent = word.bounding_box.x1 - self.column_box.x1
        if units == "col_width":
            return indent / self.column_box.width()
        elif units == "line_height":
            return indent / np.mean([word.bounding_box.height() for word in self.words])
        elif units == "pixels":
            return indent

    def __repr__(self):
        return "<Line: {} @ {}>".format(
            [word.text for word in self.words], str(self.get_box())
        )

    def extract_bullet_by_space(self):
        # Get average character width and use that to detect wider spaces
        character_sizes = []
        for word in self.words:
            width = word.bounding_box.x2 - word.bounding_box.x1
            char_size = width / len(word.text)
            character_sizes.extend([char_size for _ in range(len(word.text))])

        threshold = np.mean(character_sizes) * BULLET_THRESHOLD
        for index, word in enumerate(self.words[:-1]):
            space = self.words[index + 1].bounding_box.x1 - word.bounding_box.x2
            if space > threshold:
                bullet_words = self.words[: index + 1]
                self.words = self.words[index + 1 :]
                return Word(
                    " ".join([w.text for w in bullet_words]),
                    BoundingBox(
                        min(word.bounding_box.x1 for word in bullet_words),
                        min(word.bounding_box.y1 for word in bullet_words),
                        max(word.bounding_box.x2 for word in bullet_words),
                        max(word.bounding_box.y2 for word in bullet_words),
                    ),
                )

    def extract_bullet_by_pattern(
        self, bullet_patterns=[r"•", r"-", r"\d+\.\d+\.\d+", r"[a-zA-Z0-9]{1,3}\)"]
    ):
        text = ""
        for index, word in enumerate(self.words):
            text += word.text.strip()
            for pattern in bullet_patterns:
                if re.match(pattern, text):
                    bullet_words = self.words[: index + 1]
                    self.words = self.words[index + 1 :]
                    return Word(
                        text.replace(" ", ""),
                        BoundingBoxSet(
                            word.bounding_box for word in bullet_words
                        ).get_outer_box(),
                    )

class Item(object):
    lines = None
    bullet = None
    tabs = None

    def __init__(self, lines, bullet=None):
        self.bullet = bullet
        self.lines = lines or []

    def set_bullet(self, bullet):
        self.bullet = bullet

    def add_lines(self, lines):
        self.lines.extend(lines)

    def get_box(self, include_bullet=False):
        boxes = [line.get_box() for line in self.lines]
        if include_bullet and self.bullet:
            boxes = [self.bullet.bounding_box] + boxes
        boxes = BoundingBoxSet([box for box in boxes if box])
        return boxes.get_outer_box()

    def get_indentation(self, include_bullet=False, units="col_width"):
        assert len(self.lines) > 0, "can't get indentation for item with no lines"
        word = self.bullet if self.bullet and include_bullet else self.lines[0].words[0]
        return self.lines[0].get_indentation(word=word, units=units)

    def average_fontweight(self):
        fontweights = [
            line.fontweight for line in self.lines if line.fontweight is not None
        ]
        if fontweights:
            return np.mean(fontweights)
        else:
            return None

    def get_text(self, separator=" "):
        text = separator.join([line.get_text().strip() for line in self.lines])
        text = (
            text.replace(" ,", ",")
            .replace(" .", ".")
            .replace(" :", ":")
            .replace(" ;", ";")
            .replace("( ", "(")
            .replace(" )", ")")
            .lstrip(".")
            .strip()
        )
        for match in re.findall(r"\d+\. \d+\. \d+", text):
            text = text.replace(match, match.replace(" ", ""))
        return text

    def __str__(self):
        text = "'{}'".format(self.get_text(separator=" "))
        if self.bullet:
            text = "[{}] {}".format(self.bullet.text, text)
        return text

    def __repr__(self):
        return "<Item: {} @ {}>".format(str(self), self.get_box())

class ItemList(list):
    def get_box(self, include_bullet=False):
        items = BoundingBoxSet(
            [item.get_box(include_bullet=include_bullet) for item in self]
        )
        return items.get_outer_box()

    def add_item(self, item):
        if item.lines:
            self.append(item)

    def combine_lines(
        self,
        header_text=[
            "Content",
            "Specific Objectives",
            "Suggested Resources",
            "Suggested Further Assessment",
            "Notes",
        ],
        factor_in_fontweight=False,
    ):
        new_items = ItemList([])
        current_item = Item([])
        prev_bold = False
        for item in self:
            # check whether there's a bullet, or the boldness of the text changed
            bold = (
                (item.lines[0].fontweight > 1)
                if factor_in_fontweight and item.lines[0].fontweight
                else False
            )
            has_bullet = bool(item.bullet)
            boldness_changed = bold != prev_bold
            if has_bullet or boldness_changed:
                new_items.add_item(current_item)
                current_item = Item([], bullet=item.bullet)
            current_item.add_lines(item.lines)
            prev_bold = bold
        new_items.add_item(current_item)

        itemlist = ItemList([item for item in new_items if item.get_text().strip()])

        # split on header text
        for i, item in list(reversed(list(enumerate(itemlist)))):
            if len(item.lines) == 1:
                continue
            found = False
            for header in header_text:
                if header in item.get_text():
                    section_starts = set([0, len(item.lines)])
                    found = ""
                    started_at = None
                    for j, line in enumerate(item.lines):
                        seeking = header[len(found) :].strip()
                        if seeking.startswith(line.get_text()):
                            if started_at is None:  # found start
                                started_at = j
                            found += " " + line.get_text()
                            if found.strip() == header:  # found end
                                section_starts.add(started_at)
                                section_starts.add(j + 1)
                                found = ""
                                started_at = None
                    section_starts = list(sorted(section_starts))
                    new_items = []
                    for ind in range(len(section_starts) - 1):
                        lines = item.lines[
                            section_starts[ind] : section_starts[ind + 1]
                        ]
                        if ind == 0:
                            new_items.append(Item(lines, bullet=item.bullet))
                        else:
                            new_items.append(Item(lines))
                    itemlist[i : i + 1] = new_items

        return itemlist

class PageImage(np.ndarray):

    _annotated_array = None

    box = None

    def __new__(
        subtype,
        source,
        dtype=float,
        buffer=None,
        offset=0,
        strides=None,
        order=None,
        box=None,
    ):
        if isinstance(source, str):
            source = cv2.imread(source)
        elif isinstance(source, Image.Image):
            source = np.array(source)

        shape = source.shape
        dtype = source.dtype
        buffer = source.copy().data

        obj = super(PageImage, subtype).__new__(
            subtype, shape, dtype, buffer, offset, strides, order
        )
        obj.box = box

        return obj

    def __array_finalize__(self, obj):
        if obj is None:
            return
        # Note that it is here, rather than in the __new__ method, that we set default values,
        # because this method sees all creation of default objects - with the __new__ constructor,
        # but also with arr.view(PageImage).
        self.box = getattr(obj, "box", None)

    def _initialize_annotated_array(self):
        if not self._annotated_array is not None:
            self._annotated_array = self.copy()

    def _repr_png_(self):
        self._initialize_annotated_array()
        return self.as_pil_image(annotated=True)._repr_png_()

    def clear(self):
        self._initialize_annotated_array()
        self._annotated_array.data = self.copy().data

    def draw_box(self, box, color=(255, 0, 0), width=2):
        self._initialize_annotated_array()
        if isinstance(color[0], float):
            color = tuple(map(int, np.array(color[:3]) * 255))
        if hasattr(box, "bounding_box"):
            box = box.bounding_box
        elif hasattr(box, "get_box"):
            box = box.get_box()
        elif hasattr(box, "get_outer_box"):
            box = box.get_outer_box()
        if isinstance(box, tuple):
            box = BoundingBox(*box)
        cv2.rectangle(
            self._annotated_array, (box.x1, box.y1), (box.x2, box.y2), color, width
        )

    def as_pil_image(self, full=False, annotated=False):
        image = Image.fromarray(self._annotated_array if annotated else self)
        if self.box and not full:
            image = image.crop((self.box.x1, self.box.y1, self.box.x2, self.box.y2))
        return image

    def subimage(self, box, relative_to_page=False):
        if relative_to_page or not self.box:
            offset_x = 0
            offset_y = 0
        else:
            offset_x = self.box.x1
            offset_y = self.box.y1
        return self[
            box.x1 + offset_x : box.x2 + offset_x,
            box.y1 + offset_y : box.y2 + offset_y,
            :,
        ]

max_stars_repo_name: learningequality/curriculum-scanner

max_issues_repo_name: learningequality/curriculum-scanner

max_forks_repo_name: learningequality/curriculum-scanner


text:
from collections import namedtuple
import sklearn.linear_model as lm
import numpy as np
from houghvst.estimation import gat, regions

def compute_score(sigmas, tol):
    x = (sigmas - 1) / tol
    weights = np.exp(-(x ** 2))
    score = np.sum(weights)
    if score > 0:
        sigma_est = np.sum(sigmas * weights) / score
    else:
        sigma_est = np.nan
    return sigma_est, score

AccumulatorSpace = namedtuple('AccumulatorSpace', ['score', 'sigma_sq_range',
                                                   'alpha_range', 'sigma_sq',
                                                   'alpha'])

def hough_estimation(blocks, sigma_sq_range, alpha_range, tol=1e-2):
    score = np.zeros((len(sigma_sq_range), len(alpha_range)))

    for i_a, alpha in enumerate(alpha_range):
        for i_s, sigma_sq in enumerate(sigma_sq_range):
            # print('{} / {}'.format(i_a, len(alpha_range)), '--',
            #       '{} / {}'.format(i_s, len(sigma_range)))

            blocks_gat = gat.compute_gat(blocks, sigma_sq, alpha=alpha)
            sigmas = np.std(blocks_gat, axis=(1, 2), ddof=1)
            score[i_s, i_a] = compute_score(sigmas, tol)[1]

    max_score_idx = np.argmax(score)
    best_params = np.unravel_index(max_score_idx, score.shape)
    sigma_sq_est = sigma_sq_range[best_params[0]]
    alpha_est = alpha_range[best_params[1]]

    print('\tHighest score=', score[best_params[0], best_params[1]])

    acc = AccumulatorSpace(score, sigma_sq_range, alpha_range, sigma_sq_est,
                           alpha_est)
    return sigma_sq_est, alpha_est, acc

def hough_estimation_point(img, sigma, alpha, tol=2e-3):
    return hough_estimation(img, [sigma], [alpha], tol=tol)

def compute_mean_var(blocks):
    means = np.mean(blocks, axis=(1, 2))
    variances = np.var(blocks, axis=(1, 2), ddof=1)
    return means, variances

def regress_sigma_alpha(means, variances, verbose=True):
    # reg = lm.LinearRegression(fit_intercept=True)
    reg = lm.HuberRegressor(alpha=0, fit_intercept=True)
    reg.fit(means[:, np.newaxis], variances)

    alpha = reg.coef_[0]
    sigma_sq = reg.intercept_
    return sigma_sq, alpha

def initial_estimate_sigma_alpha(blocks):
    means, variances = compute_mean_var(blocks)
    sigma_est, alpha_est = regress_sigma_alpha(means, variances, verbose=False)
    return sigma_est, alpha_est

EstimationResult = namedtuple('EstimationResult',
                              ['alpha_init', 'sigma_sq_init',
                               'alpha', 'sigma_sq',
                               'acc_space_init', 'acc_space'])

def estimate_vst_movie(movie, block_size=8, stride=8):
    blocks = []
    for img in movie:
        blocks.append(regions.im2col(img, block_size, stride))
    blocks = np.vstack(blocks)
    return estimate_vst_blocks(blocks)

def estimate_vst_image(img, block_size=8, stride=8):
    blocks = regions.im2col(img, block_size, stride)
    return estimate_vst_blocks(blocks)

def estimate_vst_blocks(blocks):
    sigma_sq_init, alpha_init = initial_estimate_sigma_alpha(blocks)
    print('\tinitial alpha = {}; sigma^2 = {}'.format(alpha_init,
                                                      sigma_sq_init))

    diff_s = np.maximum(2e3, np.abs(sigma_sq_init))
    diff_a = alpha_init * 0.9

    sigma_sq_range = np.linspace(sigma_sq_init - diff_s, sigma_sq_init + diff_s,
                                 num=100)
    alpha_range = np.linspace(alpha_init - diff_a, alpha_init + diff_a,
                              num=100)
    sigma_sq_mid, alpha_mid, acc_init = hough_estimation(blocks, sigma_sq_range,
                                                        alpha_range)
    print('\tmid alpha = {}; sigma^2 = {}'.format(alpha_mid, sigma_sq_mid))

    diff_s /= 10
    diff_a /= 4
    sigma_sq_range = np.linspace(sigma_sq_mid - diff_s, sigma_sq_mid + diff_s,
                                 num=100)
    alpha_range = np.linspace(alpha_mid - diff_a, alpha_mid + diff_a,
                              num=100)
    sigma_sq_final, alpha_final, acc = hough_estimation(blocks, sigma_sq_range,
                                                        alpha_range)
    print('\talpha = {}; sigma^2 = {}'.format(alpha_final, sigma_sq_final))

    return EstimationResult(alpha_init, sigma_sq_init, alpha_final,
                            sigma_sq_final, acc_init, acc)

max_stars_repo_name: marianotepper/hough-anscombe

max_issues_repo_name: marianotepper/hough-anscombe

max_forks_repo_name: marianotepper/hough-anscombe


text:
#! /usr/bin/env python3
# coding: utf-8

import numpy as np
import sys, os
import librosa
import soundfile as sf
import time
import pickle as pic

from configure_FastModel import *
from FastFCA import FastFCA

class FastMNMF(FastFCA):
    """ Blind Source Separation Using Fast Multichannel Nonnegative Matrix Factorization (FastMNMF)

    X_FTM: the observed complex spectrogram
    Q_FMM: diagonalizer that converts a spatial covariance matrix (SCM) to a diagonal matrix
    G_NFM: diagonal elements of the diagonalized SCMs
    W_NFK: basis vectors for each source
    H_NKT: activations for each source
    lambda_NFT: power spectral densities of each source (W_NFK @ H_NKT)
    Qx_power_FTM: power spectra of Qx
    Y_FTM: \sum_n lambda_NFT G_NM
    """

    def __init__(self, n_source=2, n_basis=8, xp=np, init_SCM="unit", n_bit=64, seed=0):
        """ initialize FastMNMF

        Parameters:
        -----------
            n_source: int
                The number of sources
            n_iteration: int
                The number of iteration to update all variables
            n_basis: int
                The number of bases of each source
            init_SCM: str
                How to initialize covariance matrix {circular, gradual, obs, ILRMA}
                About circular and gradual initialization, please check my paper:
                    Kouhei Sekiguchi, Yoshiaki Bando, Aditya Arie Nugraha, Kazuyoshi Yoshii, Tatsuya Kawahara:
                    Fast Multichannel Nonnegative Matrix Factorization with Directivity-Aware
                        Jointly-Diagonalizable Spatial Covariance Matrices for Blind Source Separation,
                    IEEE/ACM Transactions on Audio, Speech, and Language Processing, accepted, 2020.
            n_bit:int (32 or 64)
                The number of bits for floating point number.
                '32' may degrade the peformance in exchange for lower computational cost.
                32 -> float32 and complex64
                64 -> float64 and complex128
        """
        super(FastMNMF, self).__init__(n_source=n_source, xp=xp, init_SCM=init_SCM, n_bit=n_bit, seed=seed)
        self.method_name = "FastMNMF"
        self.n_basis = n_basis

    def initialize_PSD(self):
        self.W_NFK = self.xp.random.rand(self.n_source, self.n_freq, self.n_basis).astype(self.TYPE_FLOAT)
        self.H_NKT = self.xp.random.rand(self.n_source, self.n_basis, self.n_time).astype(self.TYPE_FLOAT)
        self.lambda_NFT = self.W_NFK @ self.H_NKT

    def make_filename_suffix(self):
        self.filename_suffix = f"S={self.n_source}-it={self.n_iteration}-L={self.n_basis}-init={self.init_SCM}"

        if self.n_bit != 64:
            self.filename_suffix += f"-bit={self.n_bit}"
        if hasattr(self, "file_id"):
            self.filename_suffix += "-ID={}".format(self.file_id)
        print("param:", self.filename_suffix)

    def update(self):
        self.update_WH()
        self.update_CovarianceDiagElement()
        self.udpate_Diagonalizer()
        self.normalize()

    def update_WH(self):
        tmp1_NFT = (self.G_NFM[:, :, None] * (self.Qx_power_FTM / (self.Y_FTM ** 2))[None]).sum(axis=3)
        tmp2_NFT = (self.G_NFM[:, :, None] / self.Y_FTM[None]).sum(axis=3)
        a_W = (self.H_NKT[:, None] * tmp1_NFT[:, :, None]).sum(axis=3)  # N F K T M
        b_W = (self.H_NKT[:, None] * tmp2_NFT[:, :, None]).sum(axis=3)
        a_H = (self.W_NFK[..., None] * tmp1_NFT[:, :, None] ).sum(axis=1) # N F K T M
        b_H = (self.W_NFK[..., None] * tmp2_NFT[:, :, None]).sum(axis=1) # N F K T M
        self.W_NFK *= self.xp.sqrt(a_W / b_W)
        self.H_NKT *= self.xp.sqrt(a_H / b_H)

        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS
        self.Y_FTM = (self.lambda_NFT[..., None] * self.G_NFM[:, :, None]).sum(axis=0)

    def normalize(self):
        phi_F = self.xp.sum(self.Q_FMM * self.Q_FMM.conj(), axis=(1, 2)).real / self.n_mic
        self.Q_FMM /= self.xp.sqrt(phi_F)[:, None, None]
        self.G_NFM /= phi_F[None, :, None]

        mu_NF = (self.G_NFM).sum(axis=2).real
        self.G_NFM /= mu_NF[:, :, None]
        self.W_NFK *= mu_NF[:, :, None]

        nu_NK = self.W_NFK.sum(axis=1)
        self.W_NFK /= nu_NK[:, None]
        self.H_NKT *= nu_NK[:, :, None]
        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS

        self.reset_variable()

    def save_parameter(self, filename):
        param_list = [self.lambda_NFT, self.G_NFM, self.Q_FMM, self.W_NFK, self.H_NKT]
        if self.xp != np:
            param_list = [self.convert_to_NumpyArray(param) for param in param_list]

        pic.dump(param_list, open(filename, "wb"))

    def load_parameter(self, filename):
        param_list = pic.load(open(filename, "rb"))
        if self.xp != np:
            param_list = [cuda.to_gpu(param) for param in param_list]

        self.lambda_NFT, self.G_NFM, self.Q_FMM, self.W_NFK, self.H_NKT = param_list

if __name__ == "__main__":
    import argparse
    import pickle as pic
    import sys, os

    parser = argparse.ArgumentParser()
    parser.add_argument('input_filename', type= str, help='filename of the multichannel observed signals')
    parser.add_argument(     '--file_id', type= str, default=    "None", help='file id')
    parser.add_argument(         '--gpu', type= int, default=         0, help='GPU ID')
    parser.add_argument(       '--n_fft', type= int, default=      1024, help='number of frequencies')
    parser.add_argument(    '--n_source', type= int, default=         4, help='number of noise sources')
    parser.add_argument(     '--n_basis', type= int, default=         8, help='number of bases')
    parser.add_argument(       '--n_mic', type= int, default=         8, help='number of microphones')
    parser.add_argument(    '--init_SCM', type= str, default= "gradual", help='circular, gradual, obs, ILRMA')
    parser.add_argument( '--n_iteration', type= int, default=       100, help='number of iteration')
    parser.add_argument(       '--n_bit', type= int, default=        64, help='number of bits for floating point number')
    args = parser.parse_args()

    if args.gpu < 0:
        import numpy as xp
    else:
        import cupy as xp
        print("Use GPU " + str(args.gpu))
        xp.cuda.Device(args.gpu).use()

    wav, fs = sf.read(args.input_filename)
    wav = wav.T
    M = min(args.n_mic, len(wav))
    for m in range(M):
        tmp = librosa.core.stft(wav[m], n_fft=args.n_fft, hop_length=int(args.n_fft/4))
        if m == 0:
            spec = np.zeros([tmp.shape[0], tmp.shape[1], M], dtype=np.complex)
        spec[:, :, m] = tmp

    spec += (np.random.rand(spec.shape[0], spec.shape[1], M) + np.random.rand(spec.shape[0], spec.shape[1], M) * 1j) * np.abs(spec).max() * 1e-5

    separater = FastMNMF(n_source=args.n_source, n_basis=args.n_basis, xp=xp, init_SCM=args.init_SCM, n_bit=args.n_bit, seed=0)
    separater.load_spectrogram(spec)
    separater.file_id = args.file_id
    separater.solve(n_iteration=args.n_iteration, save_likelihood=False, save_parameter=False, save_wav=False, save_path="./", interval_save_parameter=25)

max_stars_repo_name: oucxlw/SoundSourceSeparation

max_issues_repo_name: oucxlw/SoundSourceSeparation

max_forks_repo_name: oucxlw/SoundSourceSeparation

reason: import numpy,import cupy


text:
import math
import numpy as np
import matplotlib.pyplot as plt
import dynamic_window_approach as dwa

class Robot:
    def __init__(
        self,
        variables,
        x=0.0,
        y=0.0,
        yaw=math.radians(0),
        v=0.0,
        omega=0.0,
    ):
        self.setState([x, y, yaw, v, omega])
        self.setGoal(0, 0)

        self.var = variables

        self.__currentState = [x, y, yaw, v, omega]
        self.__targetState = np.array([0, 0, 0, 0, 0])

        self.config = dwa.Config()
        self.config.robot_type = dwa.RobotType.rectangle

        self.__obstacles = []
        self.__trajectory = []
        self.__predicted_trajectory = []

        self.isNear = False
        self.isSameYaw = False
        self.isSameSpeed = False

        self.pred_yaw = 0
        self.missedCounter = 0

    @property
    def state(self) -> list:
        return self.__state

    @property
    def goal(self) -> list:
        return self.__goal

    @property
    def trajectory(self) -> list:
        return self.__trajectory

    @property
    def obstacles(self) -> list:
        return self.__obstacles

    @property
    def targetState(self) -> list:
        return self.__targetState

    @property
    def currentState(self) -> list:
        return self.__currentState

    def setState(self, arr: list) -> None:
        self.__state = np.array(arr)

    def setPredictedTrajectory(self, arr: list) -> None:
        self.__predicted_trajectory = arr

    def setGoal(self, x: float, y: float) -> None:
        self.__goal = np.array([x, y])

    def updateTrajectory(self, state: list) -> None:
        if len(self.__trajectory) > 0:
            self.__trajectory = np.vstack((self.__trajectory, state))
        else:
            self.__trajectory = np.array(self.__state)

    def updateObstacle(self, x: float, y: float) -> None:
        if len(self.__obstacles) > 0:
            self.__obstacles = np.concatenate((self.__obstacles, np.array([[x, y]])))
        else:
            self.__obstacles = np.array([[x, y]])

    def updateCurrentState(self, val, what: int) -> None:
        oldState = self.__currentState

        if what == "position":
            [x, y] = val
            self.__currentState = [x, y, oldState[2], oldState[3], oldState[4]]
        elif what == "rotation":
            self.__currentState = [
                oldState[0],
                oldState[1],
                val,
                oldState[3],
                oldState[4],
            ]
        elif what == "velocity":
            self.__currentState = [
                oldState[0],
                oldState[1],
                oldState[2],
                val,
                oldState[4],
            ]
        elif what == "omega":
            self.__currentState = [
                oldState[0],
                oldState[1],
                oldState[2],
                oldState[3],
                val,
            ]

    def calcPath(self, cur_state=None) -> list:
        [x, y, yaw, v, omega] = cur_state if cur_state else self.__state

        u, predicted_trajectory = dwa.dwa_control(
            np.array([x, y, yaw, v, omega]),
            self.config,
            self.__goal,
            self.__obstacles,
        )

        new_state = dwa.motion(
            np.array([x, y, yaw, v, omega]), u, self.config.dt
        )

        if cur_state == None:
            self.__targetState = new_state
            self.setPredictedTrajectory(predicted_trajectory)

        return new_state

    def isCoordNear(self, current, goal, thresh) -> bool:
        dist_to_goal = math.hypot(current[0] - goal[0], current[1] - goal[1])
        return dist_to_goal <= thresh

    def isFinished(self) -> bool:
        self.isCoordNear(self.__currentState, self.goal, self.config.robot_radius)

    def isStateClose(self) -> bool:
        [x, y, yaw, v, omega] = self.__currentState
        [curX, curY, curYaw, curV, curOmega] = self.__targetState

        i = self.isCoordNear([curX, curY], [x, y], 0.05)
        j = abs(curYaw - yaw) <= 0.1
        k = abs(curV - v) <= 0.1

        self.isNear = self.isNear or i
        self.isSameYaw = self.isSameYaw or j
        self.isSameSpeed = self.isSameSpeed or k

        if self.isNear and self.isSameYaw and self.isSameSpeed:
            self.isNear = False
            self.isSameYaw = False
            self.isSameSpeed = False
            self.missedCounter = 0

            if i and j and k:
                return [True, True]

            return [True, False]
        else:
            self.missedCounter += 1
            return [False, False]

    def predictRotation(self) -> None:
        oldState = self.currentState

        print(f"current rotation = {math.degrees(oldState[2])}")

        # Keep calculating until the robot starts
        # to move from its current location
        state = oldState
        while self.isCoordNear(oldState[:2], state[:2], 0.05):
            state = self.calcPath(state).tolist()

            if abs(math.degrees(oldState[2]) - math.degrees(state[2])) >= 20:
                break

        self.pred_yaw = state[2]
        print(f"NEW rotation = {math.degrees(self.pred_yaw)}")

    def plotAllStuff(self) -> None:
        plt.cla()

        # For stopping simulation with the esc key.
        plt.gcf().canvas.mpl_connect(
            "key_release_event",
            lambda event: [exit(0) if event.key == "escape" else None],
        )

        if len(self.__predicted_trajectory) > 0:
            plt.plot(
                self.__predicted_trajectory[:, 0],
                self.__predicted_trajectory[:, 1],
                "-g",
            )

        if len(self.__obstacles) > 0:
            plt.plot(self.__obstacles[:, 0], self.__obstacles[:, 1], "ok")

        plt.plot(self.__state[0], self.__state[1], "xr")
        plt.plot(self.__goal[0], self.__goal[1], "xb")

        dwa.plot_robot(self.__state[0], self.__state[1], self.__state[2], self.config)
        dwa.plot_arrow(self.__state[0], self.__state[1], self.__state[2], length=0.3)

        plt.axis("equal")
        plt.grid(True)
        plt.pause(0.1)

max_stars_repo_name: fakhrip/Robotics-Research

max_issues_repo_name: fakhrip/Robotics-Research

max_forks_repo_name: fakhrip/Robotics-Research