text:
#!/usr/bin/env python2.7
""" Read and Parse PDB files

Author: Hongjun Bai  (bbhhjj@gmail.com)

Last modified: 06/03/2014, by HB
Fixed all known problems before 06/03/2014

TODO: update to python3
"""

import sys
import math
import argparse
import gzip
import itertools as it
import copy

""" 20 amino acids:
    ALA     A    alanine
    CYS     C    cysteine
    ASP     D    aspartic acid
    GLU     E    glutamic acid
    PHE     F    phenylalanine
    GLY     G    glycine
    HIS     H    histidine
    ILE     I    isoleucine
    LYS     K    lysine
    LEU     L    leucine
    MET     M    methionine
    ASN     N    asparagine
    PRO     P    proline
    GLN     Q    glutamine
    ARG     R    arginine
    SER     S    serine
    THR     T    threonine
    VAL     V    valine
    TRP     W    tryptophan
    TYR     Y    tyrosine
    MSE     M    Se-Met
"""

AA1to3 = {'A': 'ALA', 'C': 'CYS', 'D': 'ASP', 'E': 'GLU', 'F': 'PHE',
          'G': 'GLY', 'H': 'HIS', 'I': 'ILE', 'K': 'LYS', 'L': 'LEU',
          'M': 'MET', 'N': 'ASN', 'P': 'PRO', 'Q': 'GLN', 'R': 'ARG',
          'S': 'SER', 'T': 'THR', 'V': 'VAL', 'W': 'TRP', 'Y': 'TYR'}
AA3to1 = {'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',
          'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',
          'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',
          'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',
          'MSE': 'M', 'UNK': 'X'}
BackBoneAtom = set(['C', 'O', 'N', 'CA'])

# Useful functions for structure process
def map_ab(a_prot, b_prot, atom_name):
    """ Get index map {a_prot:b_prot}

    Keyword arguments:
    a_prot -- pdbps.Protein instance
    b_prot -- pdbps.Protein instance
    atom_name -- atom_name of selected atom in PDB

    Output:
    {a_prot_resid: b_prot_resid, ...}

    """
    shift = 0  # assuming SeqID is identical in a_prot and b_prot
    aidx_2_seqid = dict((i, res.name+str(res.id)) for i, res in enumerate(a_prot.residues))
    seqid_2_bidx = dict((res.name+str(res.id+shift), i) for i, res in enumerate(b_prot.residues))
    map = {}
    for a_idx, seqid in aidx_2_seqid.items():
        try:
            b_idx = seqid_2_bidx[seqid]
            atom = b_prot.residues[b_idx].atom[atom_name]
        except KeyError:  # Either residue not in b_prot or selected atom not in atomDict
            continue
        else:
            map[a_idx] = b_idx
    return map

def sqdist(res0, res1, atom_name0, atom_name1):
    try:
        coor0 = res0.atom[atom_name0].coor
        coor1 = res1.atom[atom_name1].coor
    except ValueError:
        sys.stderr.write("%s of %s and/or %s of %s is absent" %(atom_name0, res0.name+res0.chain_id+str(res0.id), atom_name1, res1.name+res1.chain_id+str(res1.id)))
    else:
        dist2 = sum((a-b)*(a-b) for a, b in it.izip (coor0, coor1))
    return dist2

try:
    import numpy
except ImportError:
    print("NumPy is used for distance calcualtion. Please install NumPy first")
    print("if you plan to use sqdist_mat() and sele_by_dist().")
else:
    def sqdist_mat(xyza, xyzb):
        ''' Get the distance matrix between coords array xyza and xyzb.

        Input:
            xyza: [[xa1, ya1, za1], [xa2, ya2, za2], ...]
            xyzb: [[xb1, yb1, zb1], [xb2, yb2, zb2], ...]

        Output:
            distmatrix: (an x bn)
            [[D_a1_b1, D_a1_b2, D_a1_b3, ..., D_a1_bn],
             [D_a2_b1, D_a2_b2, D_a2_b3, ..., D_a2_bn],
             .
             .
             .
             [D_an_b1, D_an_b2, D_an_b3, ..., D_an_bn],
        '''
        sizea = xyza.shape[0]
        sizeb = xyzb.shape[0]
        mat_a = xyza.reshape(sizea, 1, 3)
        mat_a = mat_a.repeat(sizeb, axis=1)
        # mat_a:
        # [[[xa1, ya1, za1], [[xa1, ya1, za1], ...],
        #  [[xa2, ya2, za2], [[xa2, ya2, za2], ...],
        #  .
        #  .
        #  .
        #  [[xan, yan, zan], [[xan, yan, zan], ...]]
        mat_b = xyzb.reshape(1, sizeb, 3)
        mat_b = mat_b.repeat(sizea, axis=0)
        # mat_b:
        # [[[xb1, yb1, zb1], [xb2, yb2, zb2], ...],
        #  [[xb1, yb1, zb1], [xb2, yb2, zb2], ...],
        #  .
        #  .
        #  .
        #  [[xb1, yb1, zb1], [xb2, yb2, zb2], ...]]
        dist = mat_a - mat_b
        dist = numpy.sum(dist * dist, axis=2)
        return dist

    def sele_by_dist(atoms0, atoms1, dist_cutoff):
        coords0 = numpy.array([a.coor for a in atoms0])
        coords1 = numpy.array([a.coor for a in atoms1])
        dist = sqdist_mat(coords0, coords1)
        min_dist = numpy.min(dist, axis=1)
        selected = (min_dist < dist_cutoff*dist_cutoff)
        selected_atoms = [atom for inside, atom in it.izip(selected, atoms0) if inside]
        return selected_atoms
#

if sys.version_info < (3,):
    range = xrange

AllowedHetRes = set('MSE NAG MAN BMA'.split())

class Protein:
    """  Parse PDB file and provide a basic framework to work with protein structure.

    Example:
    prot = Protein('mypdbfile.pdb')  # or prot = Protein('mypdbfile.pdb.gz')
    xyz = [a.coor for a in prot.atoms()]  # or xyz = [a.coor for a in prot.atoms('CA')]
    print('Contained chains: %s' %(''.join(prot.chains())))
    print('%s' % (prot.seq()))
    newBFs = [1]*prot.size()
    prot.assign_bf(newBFs)
    prot.resort()
    prot.fwrite('processed.pdb')  # or prot.fwrite('processed.pdb.gz')

    """
    def atoms(self, atom_name=None, pos0=None, pos1=None, chains=None):
        """ Shorthand to access atoms (a generator) """
        if pos0 is None or pos0 < 0 or pos0 > pos1:
            i = 0
        else:
            i = pos0
        if pos1 is None or pos1 > self.size() or i >= pos1:
            pos1 = self.size()
        j = 0
        if atom_name is None:
            while i < pos1:
                res = self.residues[i]
                if chains is None or res.chain_id in chains:
                    yield res.atoms[j]
                    if j < res.size() - 1:
                        j += 1
                    else:
                        i += 1; j = 0
                else:
                    i += 1
        else:  # Limite atom_name in ['C', 'N', 'O', 'CA', 'CB']
            assert((atom_name in BackBoneAtom) or atom_name == 'CB')
            while i < pos1:
                res = self.residues[i]
                if chains is None or res.chain_id in chains:
                    try:
                        if atom_name == 'CB' and res.name == 'GLY':
                            yield res.atom['CA']
                        else:
                            yield res.atom[atom_name]
                    except KeyError:
                        print('(%s) residue %s:%s %3d has no %s atom.'
                              %(self.name, res.chain_id, res.name, res.id, atom_name))
                i += 1

    def size(self):
        return len(self.residues)

    def seq(self, chains=None):
        myseq = []
        for res in self.residues:
            if ((chains is None or res.chain_id in chains)):
                if res.name in AA3to1:
                    myseq.append(AA3to1[res.name])
                else:
                    myseq.append('X')
        return ''.join(myseq)

    def chains(self):
        return set(res.chain_id for res in self.residues)

    def assign_bf(self, bf, scale=True):
        assert(len(self.residues) == len(bf))
        if scale:
            upper = max(bf)
            lower = min(bf)
            if lower == upper: lower -= 1
            assigned = [(x-lower)*98.0/(upper-lower)+1.0  for x in bf]
        else:
            assigned = bf
        for res, property in it.izip(self.residues, assigned):
            for atom in res.atoms:
                atom.bf = property

    def fwrite(self, name, msg=None):
        with open_loudly(name, 'w', msg) as outfile:
            outfile.write(str(self))

    def resort(self, first_rid=1, first_aid=1):
        for i, residue in enumerate(self.residues):
            residue.id = first_rid + i
        for i, atom in enumerate(self.atoms()):
            atom.serial = first_aid + i

    def __init__(self, source=None, atom_name=None, chains=None):
        pdblines = []
        if isinstance(source, str):     #
            with open_loudly(source) as infile:
                pdblines = infile.readlines()
            self.name = source
        elif isinstance(source, file):  #
        #elif has_method(source, "readlines"):  #
            source.seek(0)
            pdblines = source.readlines()
            self.name = source.name
        elif isinstance(source, list) and isinstance(source[0], Residue): #
            self.name = 'None'
            self.residues = []
            for res in source:
                if (((chains is not None) and (res.chain_id not in chains)) or
                    ((atom_name is not None) and (atom_name not in res.atom))):
                    continue
                if atom_name is not None:
                    temp_res.atoms = [res.atom[atom_name]]
                    temp_res.atom = {atom_name: res.atom[atom_name]}
                else:
                    temp_res = copy.deepcopy(res)
                self.residues.append(temp_res)
        else:
            pass  # Not decided whether I need other type of constructor yet
        if pdblines:
            pdbatoms = []
            for line in pdblines:
                # Atom record of protein (including MSE, which may start with HETATOM)
                if ((line.startswith('ATOM') or  # atom line
                    (line.startswith('HETATM') and line[17:20] in AllowedHetRes)) and  # MSE or Glycans
                    (atom_name is None or line[12:16].strip() == atom_name) and  # atom name
                    (chains is None or line[21] in chains)):  # chain or chains
                    pdbatoms.append(parse_atom(line))
                if line.startswith('END'): break  # Keep only the first model
            assert(pdbatoms)
            self.residues = []
            raw_res = []
            for i, atom in enumerate(pdbatoms[:-1]):
                raw_res.append(atom)
                if atom.resid != pdbatoms[i+1].resid or atom.icode != pdbatoms[i+1].icode:
                    self.residues.append(Residue(raw_res))
                    raw_res = []
            raw_res.append(pdbatoms[-1])
            self.residues.append(Residue(raw_res))
        if self.residues:
            self.residue = dict((res.uniq_id(), res) for res in self.residues)

    def __str__(self):
        header = "REMARK This pdb file is generated by pdblib.py\n"
        header += "REMARK Input file: "+self.name+'\n'
        header += "REMARK Included chains: "+' '.join(self.chains())+'\n'
        if self.residues:
            ch0 = self.residues[0].chain_id
            body = []
            for i, res in enumerate(self.residues):
                if res.chain_id != ch0:  # Chain TER
                    last_res = self.residues[i-1]
                    last_atom = last_res.atoms[-1]
                    body.append('%-6s%5d      %3s%2s%4d\n' %('TER', last_atom.serial+1, last_res.name, last_res.chain_id, last_res.id))
                    ch0 = res.chain_id
                body.append(str(res))
            # TER of last chain
            last_res = self.residues[-1]
            last_atom = last_res.atoms[-1]
            body.append('%-6s%5d      %3s%2s%4d\n' %('TER', last_atom.serial+1, last_res.name, last_res.chain_id, last_res.id))
            body = ''.join(record for record in body)
        else:
            body = "REMARK No residues found!\n"
        tail = 'END\n'
        return header + body + tail

class Residue:
    def uniq_id(self):
        id = self.chain_id + self.name + str(self.id) + self.icode
        return id.strip()

    def size(self):
        return len(self.atoms)

    def to_ala(self):
        if self.name is not 'G':  # Gly is kept as is.
            ala_atoms = set(['C', 'O', 'N', 'CA', 'CB'])
            self.atoms = [a for a in self.atoms if a.name in ala_atoms]
            self.name = 'A'
            self.atom = dict((x.name, x) for x in self.atoms)

    def __init__(self, pdbatoms):
        try:
            atom0 = pdbatoms[0]
        except KeyError:
            sys.stderr.write("Empty pdbatom list used for residue initialization")
            sys.exit()
        else:
            self.name = atom0.resname
            self.chain_id = atom0.chain_id
            self.id = atom0.resid
            self.icode = atom0.icode
            self.seg = atom0.seg
            #
            self.atoms = [Atom(pdbatom, self) for pdbatom in pdbatoms]
            self.atom = dict((x.name, x) for x in self.atoms)

    def __str__(self):
        return ''.join(str(atom) for atom in self.atoms)

class Atom:
    def __init__(self, pdbatom, parent=None):
        self.type = pdbatom.type
        self.serial = pdbatom.serial
        self.name = pdbatom.name
        self.altloc = pdbatom.altloc
        self.coor = pdbatom.coor
        self.occupancy = pdbatom.occupancy
        self.bf = pdbatom.bf
        self.element = pdbatom.element
        self.charge = pdbatom.charge
        #
        self.parent = parent

    def __str__(self):
        res = self.parent
        if self.name[0] in '0123456789':
            name_format = '%-4s'
        elif len(self.name) > 3:
            name_format  = '%4s'
        else:
            name_format = ' %-3s'
        atom_format = '%-6s%5d '+name_format+'%1s%3s%2s%4d%1s   %8.3f%8.3f%8.3f%6.2f%6.2f      %-4s%2s%2s\n'
        return atom_format %(self.type, self.serial, self.name, self.altloc, res.name,
                res.chain_id, res.id, res.icode, self.coor[0], self.coor[1], self.coor[2],
                self.occupancy, self.bf, res.seg, self.element, self.charge)

def parse_atom(line):
    """
    #  Atom record format in pdbfile (downloaded from rcsb.org, 08-22, 2006)
    #+      COLUMNS   DATA TYPE       FIELD         DEFINITION
    #+ -----------------------------------------------------------------------------
    #+ 0    1 -  6    Record name     ATOM
    #+ 1    7 - 11    Integer         serial        Atom serial number.
    #+ 2    13 - 16   Atom            name          Atom name.
    #+ 3    17        Character       altLoc        Alternate location indicator.
    #+ 4    18 - 20   Residue name    resName       Residue name.
    #+ 5    22        Character       chainID       Chain identifier.
    #+ 6    23 - 26   Integer         resSeq        Residue sequence number.
    #+ 7    27        AChar           iCode         Code for insertion of residues.
    #+ 8    31 - 38   Real(8.3)       x             Orthogonal coordinates for X in Angstroms
    #+ 9    39 - 46   Real(8.3)       y             Orthogonal coordinates for Y in Angstroms
    #+ 0    47 - 54   Real(8.3)       z             Orthogonal coordinates for Z in Angstroms
    #+ 1    55 - 60   Real(6.2)       occupancy     Occupancy.
    #+ 2    61 - 66   Real(6.2)       tempFactor    Temperature factor.
    #+ 3    73 - 76   LString(4)      segID         Segment identifier, left-justified.
    #+ 4    77 - 78   LString(2)      element       Element symbol, right-justified.
    #+ 5    79 - 80   LString(2)      charge        Charge on the atom.
    """
    a = line.rstrip()
    assert (len(a) >= 54)
    class PdbAtomEntry: pass
    pdbatom = PdbAtomEntry()
    pdbatom.type = a[0:6].rstrip()
    pdbatom.serial = int(a[6:11])
    pdbatom.name = a[11:16].strip()
    pdbatom.altloc = a[16]
    pdbatom.resname = a[17:20]
    pdbatom.chain_id = a[21]
    pdbatom.resid = int(a[22:26])
    pdbatom.icode = a[26]
    try:
        coor = [float(a[30:38]), float(a[38:46]), float(a[46:54])]
    except ValueError:  # Handle it if coords is missing
        coor = [0.0, 0.0, 0.0]
    # "try ... except ..." approach is faster than following way:
    # coor = [0.0, 0.0, 0.0]
    # if a[27:38].strip() != '': coor[0] = float(a[27:38])
    # if a[38:46].strip() != '': coor[1] = float(a[38:46])
    # if a[46:54].strip() != '': coor[2] = float(a[46:54])
    pdbatom.coor = coor
    #In some CASP models, the format of following parts is wrong or missing.
    try:
        pdbatom.occupancy = float(a[54:60])
    except ValueError:
        pdbatom.occupancy = 1.0
    try:
        pdbatom.bf = float(a[60:66])
    except ValueError:
        pdbatom.bf = 0.0
    try:
        pdbatom.seg = a[72:76].strip()
    except ValueError:
        pdbatom.seg = ' '
    try:
        pdbatom.element = a[76:78].strip()
    except ValueError:
        pdbatom.element = ' '
    try:
        pdbatom.charge = a[78:80].strip()
    except ValueError:
        pdbatom.charge = ' '
    return pdbatom

# Detect the defects in structure
#Auxiliary functions of check_stru():
def find_altloc(protein):
    altloc_res = []
    for res in protein.residues:
        altloc_tags = list(set(a.altloc for a in res.atoms if a.altloc != ' '))
        altloc_tags.sort()  # Ordered like: ['A', 'B', 'C']
        if len(altloc_tags) > 1:
            altloc_res.append([res, altloc_tags])
    return altloc_res

def find_gap(protein, cutoff=3.0):
    gaps = []
    for i, res in enumerate(protein.residues[:-1]):
        next_res = protein.residues[i+1]
        if (res.chain_id == next_res.chain_id and
            (res.id + 1 != next_res.id or res.icode != next_res.icode)):
            try:
                distij = sqdist(res, next_res, 'C', 'N')
            except KeyError:
                distij = 1.0E8
            if distij > cutoff * cutoff:
                gaps.append([res, next_res, math.sqrt(distij)])
    return gaps

def clean_hydrogen(protein, loudly=False):
    def hydrogen(a):
        return a.name.startswith('H') or ((a.name[0] in '0123456789') and a.name[1] == 'H')
    if loudly: cnt = 0
    for res in protein.residues:
        if loudly: cnt += len(res.atoms)
        res.atoms = [atom for atom in it.ifilterfalse(hydrogen, res.atoms)]
        if loudly: cnt -= len(res.atoms)
    if loudly and cnt > 0:
        print("%d hydrogen atoms have been removed." %(cnt))

def clean_missing_ca(protein, loudly=False):
    if loudly:
        for res in protein.residues:
            if 'CA' not in res.atom and res.name in AA3to1:
                print("Residue %s:%s%c %3d has no CA atom, deleted."
                        %(res.chain_id, res.name, res.icode, res.id))
    protein.residues = filter(lambda res: 'CA' in res.atom or res.name in AllowedHetRes, protein.residues)
    if protein.residues:
        protein.residue = dict((res.uniq_id(), res) for res in protein.residues)

def clean_altloc(protein, loudly=False):
    altlocs = find_altloc(protein)
    for res, tags in altlocs:
        if loudly:
            altloc = '|'.join(tags)
            print("altloc_tags(%s) of residue %s:%s%c %3d has been cleaned."
                    %(altloc, res.chain_id, res.name, res.icode, res.id))
        preserved_tags = set([' ', tags[0]])
        res.atoms = filter(lambda atom: atom.altloc in preserved_tags, res.atoms)
        res.atom = dict((x.name, x) for x in res.atoms)  #rebuild atomDict
# End of auxiliary functions

def check_stru(prot, logfile=None):
    """ Handle defects in structure (for MD or other purpose) """
    if logfile is not None:
        log = open_loudly(logfile, 'a+')
    else:
        log = sys.stdout
    # 0, Size of protein
    log.write('%s: %d AA\n' %(prot.name, prot.size()))
    # 1, missing backbone atoms
    missing_bbatoms = []
    for res in prot.residues:
        if res.name not in AA1to3: continue
        missing_atoms = [name for name in BackBoneAtom if name not in res.atom]
        if missing_atoms:
            missing_bbatoms.append((res, missing_atoms))
    if missing_bbatoms: log.write("Missing backbone atom(s) detected:\n")
    for res, atom_names in missing_bbatoms:
        log.write("Residue %s:%s %3d missing [ %s ].\n"
                %(res.chain_id, res.name, res.id, ', '.join(atom_names)))
    # 2, has altloc_tags
    altloc_tags = find_altloc(prot)
    if altloc_tags: log.write("Alternative locations detected:\n")
    for res, tags in altloc_tags:
        log.write("Residue %s:%s %3d has %d alternative locations (%s).\n"
                %(res.chain_id, res.name, res.id, len(tags), '|'.join(tags)))
    # 3, gaps in chain
    gaps = find_gap(prot)
    if gaps: log.write("Gap(s) detected:\n")
    for resi, resj, dist in gaps:
        log.write("Gap bettween %s %3d and %s %3d of chain %s: %6.2f.\n"
                %(resi.name, resi.id, resj.name, resj.id, resi.chain_id, dist))
    # 4, insertion
    insList = [res for res in prot.residues if res.icode != ' ']
    if insList: log.write("Insertion(s) detected:\n")
    for res in insList:
        log.write("Inserted residue: %s:%s %3d%s\n" %(res.chain_id, res.name, res.id, res.icode))
    # 5, unknown residues
    unknownList = [res for res in prot.residues if res.name not in AA3to1]
    if unknownList: log.write("Unknow residue(s) detected:")
    for res in unknownList:
        log.write("Unknown residue: %s:%s %3d\n" %(res.chain_id, res.name, res.id))
    if len(missing_bbatoms) + len(altloc_tags) + len(gaps) + len(insList) + len(unknownList) < 1:
        log.write("No known issue is detected.\n")
    return len(missing_bbatoms) + len(gaps) + len(unknownList) < 1

# Auxiliary functions
# Enable gzip file handle be used by 'with' statement
class GzipFile(gzip.GzipFile):
    def __enter__(cls):
        if cls.fileobj is None:
            raise ValueError("I/O operation on closed GzipFile object")
        return cls

    def __exit__(cls, *args):
        cls.close()
# End

def has_method(o, name):
    return callable(getattr(o, name, None))

def open_loudly(file, specifications='r', msg=None):
    if file.endswith('.gz'):
        open_fun = GzipFile
    else:
        open_fun = open
    try:
        return open_fun(file, specifications)
    except IOError:
        print('File open error. %s cannot be opened' %(file))
        if msg is not None: print(msg)
        sys.exit()

def parse():
    parser = argparse.ArgumentParser(description='Check if there is any problem about the protein structure')
    parser.add_argument('inputPDB', help='PDB structure file to be checked', type=argparse.FileType('rt'))
    parser.add_argument('-c', '--chain', help='specify which chain should be readed')
    parser.add_argument('-o', '--output', help='outfile name of selected chain')
    return parser.parse_args()

def main():
    para = parse()
    prot = Protein(para.inputPDB, chains=para.chain)
    print("#> %s" %(' '.join(sys.argv)))
    clean_hydrogen(prot, loudly=True)
    clean_altloc(prot, loudly=True)
    clean_missing_ca(prot, loudly=True)
    check_stru(prot)
    for ch in prot.chains():
        current_seq = prot.seq(chains=ch)
        #print '[%s:%s] %4d %s' %(para.inputPDB.name[:-4], ch, len(current_seq), current_seq)
        print('>[%s:%s] %4d' %(para.inputPDB.name[:-4], ch, len(current_seq)))
        print(current_seq)
    if para.output:
        #prot.resort()
        prot.fwrite(para.output)
    ## Test memory Usage
    #from pympler.asizeof import asizeof
    #print('Estimated memory usage of storing %s is %d byte.' %(sys.argv[1], asizeof(prot)))
    #prot = Protein(para.inputPDB, atom_name='CA', chains=para.chain)
    #print('If only CA atoms stored:')
    #print('Estimated memory usage of storing %s is %d byte.' %(sys.argv[1], asizeof(prot)))

if __name__ == "__main__":
    #import cProfile
    #cProfile.run('main()')
    main()


text:
import numpy as np
import pytest
import xarray as xr
from vcm.cubedsphere.constants import COORD_X_CENTER, COORD_Y_CENTER
from vcm.select import (
    mask_to_surface_type,
    get_latlon_grid_coords,
    RegionOfInterest,
    zonal_average_approximate,
)

def test_zonal_average_approximate():
    lat = xr.DataArray(np.linspace(0, 9, 10), dims=["x"]).rename("lat")
    data = xr.DataArray(
        [[i * j for i in range(10)] for j in [0, 1]], dims=["z", "x"]
    ).rename("data")
    zonal_avg = zonal_average_approximate(lat, data, bins=np.arange(0, 10, 2),)
    # bins are (low, high]
    np.testing.assert_allclose(zonal_avg.isel(z=0), np.zeros(4))
    np.testing.assert_allclose(zonal_avg.isel(z=1), [1.5, 3.5, 5.5, 7.5])

@pytest.fixture()
def test_surface_type_grid():
    centered_coords = {COORD_Y_CENTER: [0], COORD_X_CENTER: [0, 1, 2]}
    slmsk_grid = xr.DataArray(
        [[0, 1, 2]], dims=[COORD_Y_CENTER, COORD_X_CENTER], coords=centered_coords
    )
    ds_slmsk = xr.Dataset({"land_sea_mask": slmsk_grid, "checkvar": slmsk_grid})

    return ds_slmsk

@pytest.fixture()
def test_latlon_grid():
    corner_coords = {"tile": [1], COORD_Y_CENTER: range(10), COORD_X_CENTER: range(10)}
    lat = xr.DataArray(
        [[range(10) for i in range(10)]],
        dims=["tile", COORD_Y_CENTER, COORD_X_CENTER],
        coords=corner_coords,
    )
    lon = xr.DataArray(
        [[[i for j in range(10)] for i in range(10)]],
        dims=["tile", COORD_Y_CENTER, COORD_X_CENTER],
        coords=corner_coords,
    )
    grid = xr.Dataset({"lon": lon, "lat": lat})
    return grid

def test_mask_to_surface_type(test_surface_type_grid):
    sea = mask_to_surface_type(test_surface_type_grid, "sea").checkvar.values
    np.testing.assert_equal(sea, np.array([[0, np.nan, np.nan]]))
    land = mask_to_surface_type(test_surface_type_grid, "land").checkvar.values
    np.testing.assert_equal(land, np.array([[np.nan, 1, np.nan]]))
    seaice = mask_to_surface_type(test_surface_type_grid, "seaice").checkvar.values
    np.testing.assert_equal(seaice, np.array([[np.nan, np.nan, 2]]))
    no_mask = mask_to_surface_type(test_surface_type_grid, "global").checkvar.values
    np.testing.assert_equal(no_mask, np.array([[0, 1, 2]]))

def test_get_latlon_grid_coords(test_latlon_grid):
    test_exact_pt = get_latlon_grid_coords(
        test_latlon_grid,
        lat=3,
        lon=2,
        init_search_width=0.5,
        search_width_increment=0.5,
        max_search_width=3,
    )
    assert test_exact_pt[COORD_X_CENTER] == 3
    assert test_exact_pt[COORD_Y_CENTER] == 2
    test_within_radius = get_latlon_grid_coords(
        test_latlon_grid,
        lat=3.25,
        lon=2.25,
        init_search_width=0.45,
        search_width_increment=0.5,
        max_search_width=3,
    )
    assert test_within_radius[COORD_X_CENTER] == 3
    assert test_within_radius[COORD_Y_CENTER] == 2
    test_multiple_pts = get_latlon_grid_coords(
        test_latlon_grid,
        lat=9,
        lon=9,
        init_search_width=2.5,
        search_width_increment=0.5,
        max_search_width=3,
    )
    assert test_multiple_pts[COORD_X_CENTER] in [9, 8, 7]
    assert test_multiple_pts[COORD_Y_CENTER] in [9, 8, 7]
    with pytest.raises(ValueError):
        test_no_match = get_latlon_grid_coords(
            test_latlon_grid,
            lat=13.5,
            lon=12.5,
            init_search_width=0.45,
            search_width_increment=0.5,
            max_search_width=3,
        )
        del test_no_match  # to placate linter

@pytest.mark.parametrize(
    "lat_bounds, lon_bounds, expected_avg",
    [
        pytest.param((1.0, 4.0), (1.0, 4.0), 2.0, id="test area averaging"),
        pytest.param((0.0, 3.0), (0.0, 3.0), 1.0, id="test lat/lon selection"),
    ],
)
def test_roi_average(lat_bounds, lon_bounds, expected_avg):
    roi = RegionOfInterest(lat_bounds, lon_bounds)
    xrange, yrange = 5, 5
    spike_value = 3
    area = [
        [3.0 if x == 3 and y == 3 else 1.0 for x in range(xrange)]
        for y in range(yrange)
    ]
    var = [
        [spike_value if x == 3 and y == 3 else 1.0 for x in range(xrange)]
        for y in range(yrange)
    ]
    ds = xr.Dataset(
        {
            "lat": (
                ["x", "y"],
                np.array([[x for x in range(xrange)] for y in range(yrange)]),
            ),
            "lon": (
                ["x", "y"],
                np.array([[y for x in range(xrange)] for y in range(yrange)]),
            ),
            "var": (["x", "y"], var),
            "area": (["x", "y"], area),
        }
    )
    assert roi.average(ds)["var"].values == expected_avg

ext: py

reason: import numpy


text:
import jieba
import pypinyin
import numpy as np
import os
import copy

#sent = '知情人士称，特朗普已表示不会?参加:当选,总统拜登的就职典礼。'
sp_tokens = '。？！?!.;；'
with open(os.path.join('./vocab','vocab_hanzi.txt')) as F:
    cn_vocab = F.read().split('\n')

cn2idx = dict([(c,i) for i,c in enumerate(cn_vocab)])
idx2cn = dict([(i,c) for i,c in enumerate(cn_vocab)])

def split2sent(text):
    new_sub = [text]

    while True:
        sub = copy.deepcopy(new_sub)
        new_sub = []
        for s in sub:
            sp = False
            for t in sp_tokens:
                if t in s:
                    new_sub += s.split(t)
                    sp = True
                    break

            if not sp:
                new_sub += [s]
        if len(new_sub) == len(sub):
            break
    tokens = [a for a in text if a in sp_tokens]
    return new_sub,tokens

def correct_tone3(pys):

    for i in range(2,len(pys)):
        if pys[i][-1] == '3' and  pys[i-1][-1] == '3' and  pys[i-2][-1] == '3':
            pys[i-1] = pys[i-1][:-1]+'2'#change the middle one

    for i in range(1,len(pys)):
        if pys[i][-1] == '3':
            if pys[i-1][-1]=='3':
                pys[i-1] = pys[i-1][:-1]+'2'
    return pys

def correct_tone4(pys):
    for i in range(len(pys)-1):
        if pys[i] == 'bu4':
            if pys[i+1][-1]=='4':
                pys[i] = 'bu2'

    return pys

def replace_with_sp(pys):
    for i,p in enumerate(pys):
        if p in ',，、':
            pys[i] = 'sp1'
    return pys
def correct_tone5(pys):
    for i in range(len(pys)):
            if pys[i][-1] not in '1234':
                pys[i] += '5'
    return pys

def convert(sent):
    pys = []
    for words in list(jieba.cut(sent)):
        py = pypinyin.pinyin(words,pypinyin.TONE3)
        py = [p[0] for p in py]
        pys += py
    pys = replace_with_sp(pys)
    pys = correct_tone3(pys)
    pys = correct_tone4(pys)
    pys = correct_tone5(pys)

    #for py in all_pys:
    all_pys =['sil']+ pys+['sil']

    return ' '.join(all_pys)

def convert_cn(cn_sentence):
    cn_sentence = list(cn_sentence)
    cn_sentence = ['sp1' if c in ',，：:、' else c for c in cn_sentence]
    cn_sentence = ['sil']+cn_sentence+['sil']
    cn_text = np.array([[cn2idx[c] for c in cn_sentence]])

    return cn_text


text:
import sys
sys.setrecursionlimit(int(1e4))
sys.path.insert(0, "./semantics/dreamcoder")

import random
from collections import defaultdict, Counter
import json
import math
import os
import datetime
import numpy as np
import torch

from dreamcoder.dreamcoder import commandlineArguments, explorationCompression
from dreamcoder.utilities import eprint, flatten, testTrainSplit, numberOfCPUs
from dreamcoder.grammar import Grammar
from dreamcoder.task import Task
from dreamcoder.type import Context, arrow, tbool, tlist, tint, t0, UnificationFailure
from dreamcoder.recognition import RecurrentFeatureExtractor
from dreamcoder.program import Program, Invented
from dreamcoder.frontier import Frontier, FrontierEntry

from dreamcoder.domains.hint.hintPrimitives import McCarthyPrimitives
from dreamcoder.domains.hint.main import main, list_options, LearnedFeatureExtractor

from utils import SYMBOLS

class ProgramWrapper(object):
    def __init__(self, prog):
        try:
            self.fn = prog.evaluate([])
        except RecursionError as e:
            self.fn = None
        self.prog = prog
        self.arity = len(prog.infer().functionArguments())
        self._name = None
        self.cache = {} # used for fast computation

    def __call__(self, *inputs):
        if len(inputs) != self.arity or None in inputs:
            raise TypeError
        if inputs in self.cache:
            return self.cache[inputs]
        fn = self.fn
        for x in inputs:
            fn = fn(x)
        self.cache[inputs] = fn
        return fn

    def __eq__(self, prog): # only used for removing equivalent semantics
        if self.arity != prog.arity:
            return False
        if isinstance(self.fn, int) and isinstance(prog.fn, int):
            return self.fn == prog.fn
        # if self.y is not None and prog.y is not None:
        #     assert len(self.y) == len(prog.y) # the program should be evaluated on same examples
        #     return np.mean(self.y[self.y!=None] == prog.y[self.y!=None]) > 0.95
        # return self.prog == prog.prog
        return False

    def __str__(self):
        return "%s %s"%(self.name, self.prog)

    @property
    def name(self):
        if self._name is not None: return self._name
        if isinstance(self.fn, int):
            self._name = str(self.fn)
        else:
            self._name = "fn"
            pass # TODO: assign name based on the function
        return self._name

    def evaluate(self, examples, store_y=True):
        ys = []
        for exp in examples:
            try:
                y = self(*exp)
            except (TypeError, RecursionError) as e:
                y = None
            ys.append(y)
        return ys

def compute_likelihood(program=None, examples=None):
    if examples is None:
        return 0., None
    elif program is None:
        res = [True if len(xs) == 0 and y is None else False for xs, y in examples ]
        return np.mean(res), np.array(res)
    else:
        pred = program.evaluate([e[0] for e in examples], store_y=False)
        gt = np.array([e[1] for e in examples])
        res = pred == gt
        return np.mean(res), np.array(res)

class Semantics(object):
    def __init__(self, idx, program=None, fewshot=False, learnable=True):
        self.idx = idx
        self.examples = []
        self.program = program
        self.arity = None
        self.solved = False
        self.likelihood = 0.
        self.fewshot = fewshot
        self.learnable = learnable

    def update_examples(self, examples):
        if len(examples) < 10 and not self.fewshot:
            self.clear()
            return

        if None in [x[1] for x in examples]:
            counts = Counter([x[1] for x in examples])
            if counts[None] / len(examples) >= 0.8:
                self.program = None
            else:
                examples = [x for x in examples if x[1] is not None]

        arity = Counter([len(x[0]) for x in examples]).most_common(1)[0][0]
        examples = [x[:2] for x in examples if len(x[0]) == arity]

        self.arity = arity
        self.examples = examples
        self.likelihood, self.res = compute_likelihood(self.program, self.examples)
        self.check_solved()

    def update_program(self, entry):
        program = ProgramWrapper(entry.program)
        likelihood = compute_likelihood(program, self.examples)[0]
        if (likelihood > self.likelihood) or \
            (likelihood == self.likelihood and len(str(program)) < len(str(self.program))):
            self.program = program
            self.likelihood = likelihood
            self.check_solved()

    def check_solved(self):
        if self.arity == 0 and self.likelihood > 0. and self.program is not None:
            self.solved = True
        elif self.arity > 0 and self.likelihood >= 0.9 and len(set(self.examples)) >= 80 and '#' not in str(self.program): # for + -
            self.solved = True
        elif self.arity > 0 and self.likelihood >= 0.95 and len(set(self.examples)) >= 80 and '#' in str(self.program):
            self.solved = True
        elif self.fewshot and self.likelihood >= 0.95 and len(set(self.examples)) >= 10:
            self.solved = True
        else:
            self.solved = False

    def __call__(self, *inputs):
        if self.program is None and len(inputs) == 0:
            return None
        return self.program(*inputs)

    def make_task(self):
        min_examples = 30 if self.arity is not None and self.arity > 0 else 10
        min_examples = min_examples if not self.fewshot else 0
        max_examples = 100
        examples = self.examples
        if len(examples) < min_examples or self.solved or None in [x[1] for x in examples]:
            return None
        task_type = arrow(*([tint]*(self.arity + 1)))
        if len(examples) > max_examples:
            wrong_examples = [e for e, r in zip(examples, self.res) if not r]
            right_examples = [e for e, r in zip(examples, self.res) if r]
            right_examples = random.choices(right_examples, k=max_examples-len(wrong_examples))
            examples = wrong_examples + right_examples
            examples = random.sample(examples, k=max_examples)
        return Task(str(self.idx), task_type, examples)

    def clear(self):
        self.examples = []
        self.program = None
        self.arity = None
        self.solved = False
        self.likelihood = 0.

    def save(self):
        model = {'idx': self.idx, 'solved': self.solved, 'likelihood': self.likelihood, 'arity': self.arity}
        model['program'] = None if self.program is None else self.program.prog
        return model

    def load(self, model):
        self.idx = model['idx']
        self.solved = model['solved']
        self.likelihood = model['likelihood']
        self.arity = model['arity']
        self.program = None if model['program'] is None else ProgramWrapper(model['program'])

class DreamCoder(object):
    def __init__(self):
        args = commandlineArguments(
            enumerationTimeout=200, activation='tanh', iterations=1, recognitionTimeout=3600,
            a=3, maximumFrontier=5, topK=2, pseudoCounts=30.0,
            helmholtzRatio=0.5, structurePenalty=1.,
            CPUs=min(numberOfCPUs(), 8),
            extras=list_options)

        args['noConsolidation'] = True
        args.pop("random_seed")
        args['contextual'] = True
        args['biasOptimal'] = True
        args['auxiliaryLoss'] = True
        args['activation'] = "relu"
        args['useDSL'] = False

        extractor = {
            "learned": LearnedFeatureExtractor,
        }[args.pop("extractor")]
        extractor.H = args.pop("hidden")

        timestamp = datetime.datetime.now().isoformat()
        outputDirectory = "tmp/%s"%timestamp
        os.system("mkdir -p %s"%outputDirectory)

        args.update({
            "featureExtractor": extractor,
            "outputPrefix": "%s/hint"%outputDirectory,
            "evaluationTimeout": 0.0005,
        })
        args.pop("maxTasks")
        args.pop("split")

        self.primitives = McCarthyPrimitives()
        baseGrammar = Grammar.uniform(self.primitives)
        self.grammar = baseGrammar
        self.train_args = args
        self.semantics = [Semantics(i) for i in range(len(SYMBOLS))]
        self.allFrontiers = None
        self.helmholtzFrontiers = None

    def __call__(self):
        return self.semantics

    def save(self):
        model = [smt.save() for smt in self.semantics]
        return model

    def load(self, model):
        if model is None:
            return
        assert len(self.semantics) == len(model)
        for i in range(len(self.semantics)):
            self.semantics[i].load(model[i])

    def extend(self, n):
        for smt in self.semantics:
            smt.learnable = False
        idx = len(SYMBOLS) - 1
        self.semantics.append(Semantics(idx, fewshot=True))
        self.primitives.extend([Invented(smt.program.prog) for smt in self.semantics if not smt.learnable and smt.arity > 0])

    def rescore_frontiers(self, tasks):
        if self.allFrontiers is None:
            return
        print('Rescoring %d frontiers...'%len(self.allFrontiers))
        id2task = {t.name: t for t in tasks}
        id2frontier = {f.task.name: f for f in self.allFrontiers}
        allFrontiers = {}
        for name in id2task.keys():
            task = id2task[name]
            examples = task.examples
            if name not in id2frontier:
                frontier = Frontier([], task=task)
            else:
                frontier = id2frontier[name]
                frontier.task = task
                for entry in frontier.entries:
                    program = ProgramWrapper(entry.program)
                    entry.logLikelihood = float(np.log(compute_likelihood(program=program, examples=examples)[0]))
                    entry.logPosterior = entry.logLikelihood + entry.logPrior
                frontier.removeLowLikelihood(low=0.1)

            allFrontiers[task] = frontier
        self.allFrontiers = allFrontiers

    def learn(self, dataset):
        tasks = []
        max_arity = 0
        for smt, exps in zip(self.semantics, dataset):
            if not smt.learnable:
                continue
            smt.update_examples(exps)
            t = smt.make_task()
            if t is not None:
                tasks.append(t)
                max_arity = max(smt.arity, max_arity)
        self.train_args['enumerationTimeout'] = 5 if max_arity == 0 else 300
        # self.train_args['iterations'] = 1 if max_arity == 0 else 3
        n_solved = len(['' for t in self.semantics if t.solved])
        print("Semantics: %d/%d/%d (total/solved/learn)."%(len(self.semantics), n_solved, len(tasks)))
        if len(tasks) == 0:
            self._print_semantics()
            return
        self._print_tasks(tasks)
        self.update_grammar()
        print(self.grammar)
        # print(self.allFrontiers)
        self.rescore_frontiers(tasks)
        # if self.allFrontiers is not None:
        #     print(self.allFrontiers.values())

        if self.helmholtzFrontiers is not None:
            requests_old ={x.task.request for x in self.helmholtzFrontiers()}
            requests = {t.request for t in tasks}
            # if new requests, discard old helmholtz frontiers
            if requests != requests_old:
                self.helmholtzFrontiers = None

        result = explorationCompression(self.grammar, tasks, allFrontiers=self.allFrontiers, helmholtzFrontiers=self.helmholtzFrontiers, **self.train_args)
        self.allFrontiers = list(result.allFrontiers.values())
        self.helmholtzFrontiers = result.helmholtzFrontiers

        for frontier in result.taskSolutions.values():
            if not frontier.entries: continue
            symbol_idx = int(frontier.task.name)
            # print(frontier)
            self.semantics[symbol_idx].update_program(frontier.bestPosterior)
        # examples = [xs for t in tasks for xs, y in t.examples]
        # self._removeEquivalentSemantics(examples)
        self._removeEquivalentSemantics()
        self._print_semantics()
        # self.grammar = result.grammars[-1]

    def update_grammar(self):
        programs = [Invented(smt.program.prog) for smt in self.semantics
            if smt.learnable and smt.solved and smt.program is not None and smt.program.arity > 0 and '#' not in str(smt.program)]
            # if '#' in the program, the program uses a invented primitive, it is very likely to have a high computation cost.
            # Therefore we don't add this program into primitives, since it might slow the enumeration a lot.
            # it might be resolved by increasing the enumeration time
        new_grammar = Grammar.uniform(self.primitives + programs)
        # self.train_args['enumerationTimeout'] += 100 * len(programs)
        if new_grammar != self.grammar:
            self.grammar = new_grammar
            self.helmholtzFrontiers = None
            self.allFrontiers = None
            print("Update grammar with invented programs and set frontiers to none.")

    def _print_semantics(self):
        for smt in sorted(self.semantics, key=lambda x: x.idx):
            print("Symbol-%02d: %s %.2f"%(smt.idx, smt.program, smt.likelihood))
            # print("Solved!" if smt.solved else "")

    def _print_tasks(self, tasks):
        for task in tasks:
            print("Symbol-%02d (%s), Samples: %3d, "%(int(task.name), task.request, len(task.examples)), Counter(task.examples).most_common(10))

        json.dump([t.examples for t in tasks], open('outputs/tasks.json', 'w'))

    def _removeEquivalentSemantics(self, examples=None):
        if examples is not None:
            examples = list(set(examples))
            for smt in self.semantics:
                if smt.program is not None:
                    smt.program.evaluate(examples)

        for i in range(len(self.semantics) - 1):
            smt_i = self.semantics[i]
            if smt_i.program is None:
                continue
            for j in range(i+1, len(self.semantics)):
                smt_j = self.semantics[j]
                if smt_j.program is None:
                    continue
                if smt_i.program == smt_j.program:
                    if len(smt_i.examples) >= len(smt_j.examples):
                        smt_j.clear()
                    else:
                        smt_i.clear()
                        break

max_stars_repo_name: liqing-ustc/ANS

max_issues_repo_name: liqing-ustc/CLL-NeSy

max_forks_repo_name: liqing-ustc/CLL-NeSy


text:
import numpy as np
from skimage.measure import label

from imlib.image import objects

many_objects = np.array(
    [
        [1, 1, 0, 1, 1, 1, 1, 0, 0],
        [1, 1, 0, 1, 1, 1, 1, 0, 0],
        [0, 0, 0, 1, 1, 1, 1, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 1],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 1],
        [0, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 0, 0, 1],
        [0, 1, 1, 1, 1, 1, 0, 0, 1],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
    ],
    dtype=bool,
)

largest_2 = np.array(
    [
        [0, 0, 0, 1, 1, 1, 1, 0, 0],
        [0, 0, 0, 1, 1, 1, 1, 0, 0],
        [0, 0, 0, 1, 1, 1, 1, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
    ],
    dtype=bool,
)

largest = np.array(
    [
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
    ],
    dtype=bool,
)

def test_get_largest_non_zero_object():
    assert (
        largest == objects.get_largest_non_zero_object(label(many_objects))
    ).all()

def test_keep_n_largest_objects():
    assert (largest == objects.keep_n_largest_objects(many_objects)).all()
    assert (
        largest_2 == objects.keep_n_largest_objects(many_objects, n=2)
    ).all()

max_stars_repo_name: noisysky/imlib

max_issues_repo_name: noisysky/imlib

max_forks_repo_name: noisysky/imlib


text:
from Auxiliar_Functions import *
from scipy.integrate import odeint
R_sol = 6.96340*10**8 #m
R_sol= R_sol*(8*10**5) #eV⁻¹

def N_e_sol(r):
    r_0=R_sol/10.54
    N_0=245 #N_A cm⁻³
    N_0=N_A*N_0*(1.973*10**(-5))**3#eV³
    N = N_0*math.exp(-r/r_0)

    return N

#EDO
def f_solar(y, t, params):
    S1, S2, S3 = y      # unpack current values of y
    E= params  # unpack parameters

    N_e=N_e_sol(t)
    #Mixing parameters in matter
    theta_M=theta_eff(delta_m2_21,theta_21,Acc(N_e,E))
    delta_m2_M=delta_m2_eff(delta_m2_21,theta_21,Acc(N_e,E))
    omega_M=delta_m2_M/(2*E)

    B1,B2,B3 = omega_M*np.array(B_vec(3,theta_M))
    derivs = [S2*B3-S3*B2,      # list of dy/dt=f functions
              S3*B1-S1*B3,
              S1*B2-S2*B1]
    return derivs

def solver_solar(E,r_i,r_f):
    #E [MeV], r_i [R_solar],  r_i [R_solar]
    E=E*10**6 #eV
    #omega_0=delta_m2_21/(2*E*10**6)
    r_step=0.000001 # units of R_solar
    r_per_R_sol=np.arange(r_i,r_f,r_step) # units of R_solar
    r=R_sol*r_per_R_sol #eV⁻¹

    # Initial values
    S1_0 = 0
    S2_0 = 0
    S3_0 = 1.0  #Electronic Neutrino

    # Initial conditions for ODE solver
    y0 = [S1_0, S2_0, S3_0]

    # Call the ODE solver
    psoln=odeint(f_solar, y0, r, args=(E,))
    psoln_trans=np.transpose(psoln,(1,0))

    S1=psoln_trans[0]
    S2=psoln_trans[1]
    S3=psoln_trans[2]
    S=np.array([S1,S2,S3])

    B=[]
    for r_i in r:
        N_e=N_e_sol(r_i)
        #Mixing parameters in matter
        theta_M=theta_eff(delta_m2_21,theta_21,Acc(N_e,E))
        delta_m2_M=delta_m2_eff(delta_m2_21,theta_21,Acc(N_e,E))
        omega_M=delta_m2_M/(2*E)
        B.append(omega_M*np.array(B_vec(3,theta_M)))
    B=np.transpose(B,(1,0))

    return S,B,r_per_R_sol

max_stars_repo_name: pedrodedin/Neutrino-Collective-Effects

max_issues_repo_name: pedrodedin/Neutrino-Collective-Effects

max_forks_repo_name: pedrodedin/Neutrino-Collective-Effects

reason: from scipy


text:
"""
Class for interfacing with the camera.
"""
import time
from io import BytesIO
from PIL import Image
import numpy as np

'''
Not sure why this isn't working, but it didn't find the picamera library or raise an error
Setting this manually for now

import importlib
picam_spec = importlib.util.find_spec("picamera")
picam_found = picam_spec is not None
if picam_found:
    from picamera import PiCamera

'''
from picamera import PiCamera
picam_found = True

RESOLUTION = (1280,720)
ZOOM = (320/1920, 180/1080, 1280/1920, 720/1080)
AWB_MODE = "off"
AWB_GAINS = 1.6
FRAMERATE = 30
SHUTTER_SPEED = 5000
EXPOSURE_MODE = "off"

class Camera(object):
    """
    Class that reads from the camera.
    """

    def __init__(self):
        """
        Initialization of the camera.
        """
        if picam_found:
            print ("Starting Camera")
            self.camera = PiCamera(sensor_mode = 1)
            self.camera.resolution = RESOLUTION
            self.camera.zoom = ZOOM
            self.camera.awb_mode = AWB_MODE
            self.camera.awb_gains = AWB_GAINS
            self.camera.framerate = FRAMERATE
            #time.sleep(3)
            #self.camera.shutter_speed = SHUTTER_SPEED
            #self.camera.exposure_mode = EXPOSURE_MODE

    def capture_stream(self)->BytesIO:
        """
        Return a PIL image.
        """
        stream = BytesIO()
        if picam_found:
            self.camera.capture(stream, format="jpeg", quality=65, resize = (int(RESOLUTION[0]/2), int(RESOLUTION[1]/2))) #switch to rgb later
        stream.seek(0)
        return stream

    def capture_image(self)->Image:
        """
        Return a PIL image.
        """
        return Image.open(self.capture_stream())

    def capture_array(self)->np.ndarray:
        """
        Return a 3 channel np.array (RGB)
        """
        output = np.empty((RESOLUTION[1], RESOLUTION[0], 3), dtype=np.uint8)
        if picam_found:
            self.camera.capture(output, format="rgb")
        return output

    def open(self):
        """
        Start the camera preview
        """
        #self.camera.start_preview()
        #time.sleep(2)

    def __enter__(self):
        """
        Context manager enter
        """
        self.open()
        return self

    def close(self):
        """
        Called to clean up camera context.
        """
        if picam_found:
            self.camera.close()

    def __exit__(self, *args):
        """
        Context manager exit
        """
        self.close()

def preview_camera():
    """
    Method to preview the camera (requires a GUI).
    """
    import timeit
    num_images=10
    with Camera() as cam:
        start = timeit.default_timer()
        for i in range(num_images):
            print("Capturing raw image " + str(i) )
            _ = cam.capture_array()
        diff = timeit.default_timer() - start
        print ("Captured " + str(num_images) + " raw images in " + str(diff) + "s.")
        start = timeit.default_timer()
        for i in range(num_images):
            print("Capturing jpeg image " + str(i) )
            _ = cam.capture_image()
        diff = timeit.default_timer() - start
        print ("Captured " + str(num_images) + " jpeg images in " + str(diff) + "s.")

max_stars_repo_name: npretor/circuitprinter

max_issues_repo_name: npretor/circuitprinter

max_forks_repo_name: npretor/circuitprinter


text:
import mdtraj as md
import numpy as np
import parmed as pmd
import pytest
from openff.toolkit.topology import Molecule
from openff.toolkit.typing.engines.smirnoff import ForceField
from openff.units import unit

from openff.interchange.components.interchange import Interchange
from openff.interchange.drivers import get_amber_energies, get_openmm_energies
from openff.interchange.testing import _BaseTest

kj_mol = unit.kilojoule / unit.mol

class TestAmber(_BaseTest):
    def test_inpcrd(self, parsley):
        mol = Molecule.from_smiles(10 * "C")
        mol.name = "HPER"
        mol.generate_conformers(n_conformers=1)

        out = Interchange.from_smirnoff(force_field=parsley, topology=mol.to_topology())
        out.box = [4, 4, 4]
        out.positions = mol.conformers[0]
        out.positions = unit.nanometer * np.round(out.positions.m_as(unit.nanometer), 5)

        out.to_inpcrd("internal.inpcrd")
        out._to_parmed().save("parmed.inpcrd")

        coords1 = pmd.load_file("internal.inpcrd").coordinates
        coords2 = pmd.load_file("parmed.inpcrd").coordinates

        np.testing.assert_equal(coords1, coords2)

    @pytest.mark.parametrize(
        "smiles",
        [
            "C",
            "CC",  # Adds a proper torsion term(s)
            "C=O",  # Simplest molecule with any improper torsion
            "OC=O",  # Simplest molecule with a multi-term torsion
            "CCOC",  # This hits t86, which has a non-1.0 idivf
            "C1COC(=O)O1",  # This adds an improper, i2
        ],
    )
    @pytest.mark.parametrize("constrained", [True, False])
    def test_amber_energy(self, smiles, constrained):
        """Basic test to see if the amber energy driver is functional"""
        mol = Molecule.from_smiles(smiles)
        mol.generate_conformers(n_conformers=1)
        top = mol.to_topology()
        top.mdtop = md.Topology.from_openmm(top.to_openmm())

        if constrained:
            sage = ForceField("openff-2.0.0.offxml")
        else:
            sage = ForceField("openff_unconstrained-2.0.0.offxml")

        off_sys = Interchange.from_smirnoff(sage, top)

        off_sys.box = [4, 4, 4]
        off_sys.positions = mol.conformers[0]

        omm_energies = get_openmm_energies(off_sys)
        amb_energies = get_amber_energies(off_sys)

        # MT: I _think_ some of these errors are the result of Amber reporting energies
        # to 0.001 kcal/mol, which introduces error on the order of ~0.002 kJ/mol
        # TODO: Figure out why bond and angle energies are reported differently
        #       in constrained systems
        #       https://github.com/openforcefield/openff-interchange/issues/323
        omm_energies.compare(
            amb_energies,
            custom_tolerances={
                "Bond": (0.1 if constrained else 0.001) * kj_mol,
                "Angle": (0.05 if constrained else 0.001) * kj_mol,
                "Torsion": (0.005 if constrained else 0.001) * kj_mol,
                "vdW": 0.02 * kj_mol,
                "Electrostatics": (0.5 if constrained else 0.05) * kj_mol,
            },
        )

max_stars_repo_name: mikemhenry/openff-interchange

max_issues_repo_name: mikemhenry/openff-interchange

max_forks_repo_name: mikemhenry/openff-interchange


text:
import flymc as ff
import numpy as np
import numpy.random as npr

def logistic_regression_chain(x, t, N_iter=100, stepsize=1, th0=1, q=0.1, y0=1, seed=None):

    # Set seed
    npr.seed(seed)

    # Obtain joint distributions over z and th and set step functions
    model = ff.LogisticModel(x, t, th0=th0, y0=y0)
    z__stepper = ff.zStepMH(model.log_pseudo_lik , q)
    th_stepper = ff.ThetaStepMH(model.log_p_joint, stepsize)

    # Initialize
    N, D = x.shape
    th = np.random.randn(D)*th0
    z = ff.BrightnessVars(N)

    # run chain
    th_chain = np.zeros((N_iter,  ) +  th.shape)
    z_chain  = np.zeros((N_iter, N), dtype=bool)

    for i in range(N_iter):
        th = th_stepper.step(th, z)
        z  = z__stepper.step(th ,z)
        # Record the intermediate results
        th_chain[i,:] = th.copy()
        z_chain[i,z.bright] = 1

    print "th0 = ", th0, "frac accepted is", th_stepper.frac_accepted,  "bright frac is:", np.mean(z_chain)
    return th_chain, z_chain

max_stars_repo_name: feynmanliang/firefly-monte-carlo

max_issues_repo_name: feynmanliang/firefly-monte-carlo

max_forks_repo_name: HIPS/firefly-monte-carlo


text:
#!/usr/bin/python
"""
Plot improvement in Omega_K, gamma, or FOM as a function of z (Fig. 14),
(Fig. 15), (Fig. 18).
"""
import numpy as np
import pylab as P
from rfwrapper import rf
import matplotlib.patches
import matplotlib.cm
import matplotlib.ticker
import os
from radiofisher import euclid

cosmo = rf.experiments.cosmo

nsig = 1.5 # No. of sigma to plot out to
aspect = 1. #1.7 # Aspect ratio of range (w = aspect * h)

#TYPE = 'omegak'
#TYPE = 'gamma'
TYPE = 'FOM'

if TYPE == 'omegak':
    param1 = "omegak"
    label1 = "\Omega_K"
    fid1 = 0.
    fig_name = "fig15-ok-improvement.pdf"
elif TYPE == 'gamma':
    param1 = "gamma"
    label1 = "\gamma"
    fid1 = cosmo['gamma']
    fig_name = "fig18-gamma-improvement.pdf"
else:
    param1 = "fom"
    fig_name = "fig14-fom-improvement.pdf"

USE_DETF_PLANCK_PRIOR = True
MARGINALISE_GAMMA = True # Marginalise over gamma
MARGINALISE_CURVATURE = True # Marginalise over Omega_K
MARGINALISE_INITIAL_PK = True # Marginalise over n_s, sigma_8
MARGINALISE_OMEGAB = True # Marginalise over Omega_baryons
MARGINALISE_W0WA = True # Marginalise over (w0, wa)

names = ['EuclidRef_paper', 'exptL_paper', 'aexptM_paper'] #, 'exptS']
labels = ['DETF IV', 'Facility', 'Stage II'] #, 'Stage I']
colours = ['#CC0000', '#1619A1', '#5B9C0A', '#FFB928']
linestyle = [[2, 4, 6, 4], [], [8, 4], [3, 4]]

# Fiducial value and plotting
fig = P.figure()
ax = fig.add_subplot(111)

_k = range(len(names))
for k in _k:
    root = "output/" + names[k]

    # Load cosmo fns.
    dat = np.atleast_2d( np.genfromtxt(root+"-cosmofns-zc.dat") ).T
    zc, Hc, dAc, Dc, fc = dat
    zs, Hs, dAs, Ds, fs = np.genfromtxt(root+"-cosmofns-smooth.dat").T
    kc = np.genfromtxt(root+"-fisher-kc.dat").T

    # Load Fisher matrices as fn. of z
    Nbins = zc.size
    F_list = [np.genfromtxt(root+"-fisher-full-%d.dat" % i) for i in range(Nbins)]

    # Add each redshift bin in turn, and recalculate FOM.
    zmax = []; fom = []
    #for l in range(1, len(F_list)):
    for l in range(1, len(F_list)+1):

        # EOS FISHER MATRIX
        # Actually, (aperp, apar) are (D_A, H)
        pnames = rf.load_param_names(root+"-fisher-full-0.dat")
        zfns = ['b_HI',]
        excl = ['Tb', 'f', 'H', 'DA', 'apar', 'aperp', 'pk*', 'N_eff', 'fs8', 'bs8']
        F, lbls = rf.combined_fisher_matrix( F_list[:l],
                                                    expand=zfns, names=pnames,
                                                    exclude=excl )
        # DETF Planck prior
        print "*** Using DETF Planck prior ***"
        l2 = ['n_s', 'w0', 'wa', 'omega_b', 'omegak', 'omegaDE', 'h', 'sigma8']
        F_detf = euclid.detf_to_rf("DETF_PLANCK_FISHER.txt", cosmo, omegab=False)
        Fpl, lbls = rf.add_fisher_matrices(F, F_detf, lbls, l2, expand=True)

        # Decide whether to fix various parameters
        fixed_params = []
        if not MARGINALISE_GAMMA: fixed_params += ['gamma',]
        if not MARGINALISE_CURVATURE: fixed_params += ['omegak',]
        if not MARGINALISE_INITIAL_PK: fixed_params += ['n_s', 'sigma8']
        if not MARGINALISE_OMEGAB: fixed_params += ['omega_b',]
        if not MARGINALISE_W0WA: fixed_params += ['w0', 'wa']

        if len(fixed_params) > 0:
            Fpl, lbls = rf.combined_fisher_matrix( [Fpl,], expand=[],
                         names=lbls, exclude=fixed_params )

        # Really hopeful H0 prior
        #ph = lbls.index('h')
        #Fpl[ph, ph] += 1./(0.012)**2.

        # Invert matrix
        cov_pl = np.linalg.inv(Fpl)

        # Plot improvement in chosen parameter
        if param1 == 'fom':
            p1 = lbls.index('w0'); p2 = lbls.index('wa')
            _fom = rf.figure_of_merit(p1, p2, None, cov=cov_pl)
        else:
            pp = lbls.index(param1)
            _fom = 1. / np.sqrt(cov_pl[pp, pp])
        zmax.append(zc[:l][-1])
        fom.append(_fom)

    # Plot curve for this experiment
    line = ax.plot(zmax, fom, color=colours[k], ls='solid', lw=1.8,
                   marker='o', label=labels[k])
    line[0].set_dashes(linestyle[k])

# Report on what options were used
print "-"*50
s1 = "Marginalised over Omega_K" if MARGINALISE_CURVATURE else "Fixed Omega_K"
s2 = "Marginalised over ns, sigma8" if MARGINALISE_INITIAL_PK else "Fixed ns, sigma8"
s3 = "Marginalised over Omega_b" if MARGINALISE_OMEGAB else "Fixed Omega_b"
s4 = "Marginalised over w0, wa" if MARGINALISE_W0WA else "Fixed w0, wa"
print "NOTE:", s1
print "NOTE:", s2
print "NOTE:", s3
print "NOTE:", s4

# Axis ticks and labels
ax.legend(prop={'size':'x-large'}, bbox_to_anchor=[0.96, 0.30], frameon=False)
ax.tick_params(axis='both', which='major', labelsize=20, size=8., width=1.5, pad=8.)
ax.set_xlim((0.0, 2.6))
ax.set_xlabel("$z_\mathrm{max}$", fontdict={'fontsize':'xx-large'}, labelpad=15.)
if param1 == 'fom':
    ax.set_ylabel("$\mathrm{FOM}$", fontdict={'fontsize':'xx-large'}, labelpad=15.)
    ax.set_ylim((0., 430.))
elif param1=='omegak':
    ax.set_ylabel("$[\sigma({%s})]^{-1}$" % label1, fontdict={'fontsize':'xx-large'},
                  labelpad=15.)
    ax.set_ylim((0., 1250.))
else:
    ax.set_ylabel("$[\sigma({%s})]^{-1}$" % label1, fontdict={'fontsize':'xx-large'},
                  labelpad=15.)
    ax.set_ylim((0., 60.))

# Set size and save
P.tight_layout()
P.gcf().set_size_inches(8.,6.)
P.savefig(fig_name, transparent=True)
P.show()

max_stars_repo_name: sjforeman/RadioFisher

max_issues_repo_name: sjforeman/RadioFisher

max_forks_repo_name: sjforeman/RadioFisher


text:
import parl
import paddle
import numpy as np

__all__ = ['PaddleAgent']

class PaddleAgent(parl.Agent):
    def __init__(self, algorithm):
        super(PaddleAgent, self).__init__(algorithm)

        self.alg.sync_target(decay=0)

    def predict(self, obs):
        obs = paddle.to_tensor(obs.reshape(1, -1).astype(np.float32))
        action = self.alg.predict(obs)
        action_numpy = action.cpu().numpy()[0]
        return action_numpy

    def sample(self, obs):
        obs = paddle.to_tensor(obs.reshape(1, -1).astype(np.float32))
        action, _ = self.alg.sample(obs)
        action_numpy = action.cpu().numpy()[0]
        return action_numpy

    def learn(self, obs, action, reward, next_obs, terminal):
        terminal = np.expand_dims(terminal, -1)
        reward = np.expand_dims(reward, -1)

        obs = paddle.to_tensor(obs)
        action = paddle.to_tensor(action)
        reward = paddle.to_tensor(reward)
        next_obs = paddle.to_tensor(next_obs)
        terminal = paddle.to_tensor(terminal)
        critic_loss, actor_loss = self.alg.learn(obs, action, reward, next_obs,
                                                 terminal)
        return critic_loss, actor_loss


text:
# 用于编程测试一些函数
import numpy as np
from keras.models import load_model
model = load_model('model_200_10.h5')
x_input = np.zeros((1,361),dtype=int)
x_input[0,180] = 1
x_input[0,181] = 2
print(x_input)
print(x_input.shape)
score = model.predict(x_input) * 20000 - 10000
print(score)
score = score[0,0]
print(score)
print("------------")


text:
import random

import numpy as np
from PIL import Image, ImageEnhance
from skimage.color import rgb2lab

import tensorflow as tf

backbone = tf.keras.applications.ResNet50(include_top=False, pooling='avg')
preprocess = tf.keras.applications.resnet50.preprocess_input

class Env(object):
    """Training env wrapper of image processing RL problem"""
    def __init__(self, src, max_episode_length=20, reward_scale=0.1):
        """
        Args:
            src (list[str, str]): list of raw and retouched path, initial
                                  state will sample from it uniformly
            max_episode_length (int): max number of actions can be taken
        """
        self._src = src
        self._backbone = backbone
        self._preprocess = preprocess
        self._rgb_state = None
        self._lab_state = None
        self._target_lab = None
        self._current_diff = None
        self._count = 0
        self._max_episode_length = max_episode_length
        self._reward_scale = reward_scale
        self._info = dict()

    def _state_feature(self):
        s = self._preprocess(self._rgb_state)
        s = tf.expand_dims(s, axis=0)
        context = self._backbone(s).numpy().astype('float32')
        hist_rgb = get_rgb_hist(self._rgb_state).astype('float32')
        hist_lab = get_lab_hist(self._lab_state).astype('float32')
        return np.concatenate([context, hist_rgb, hist_lab], 1)

    def step(self, action):
        """One step"""
        self._count += 1
        self._rgb_state = self._transit(action)
        self._lab_state = rgb2lab(self._rgb_state)
        reward = self._reward()
        done = self._count >= self._max_episode_length or action == 0
        return self._state_feature(), reward, done, self._info

    def reset(self):
        """Reset"""
        self._count = 0
        raw, retouched = map(Image.open, random.choice(self._src))
        self._rgb_state = np.asarray(raw)
        self._lab_state = rgb2lab(self._rgb_state)
        self._target_lab = rgb2lab(np.asarray(retouched))
        self._current_diff = self._diff(self._lab_state)
        self._info['max_reward'] = self._current_diff
        return self._state_feature()

    def _diff(self, lab):
        reward = np.sqrt(((self._target_lab - lab) ** 2).sum(2)).mean()
        return reward * self._reward_scale

    def _reward(self):
        diff = self._current_diff
        self._current_diff = self._diff(self._lab_state)
        return min(max(diff - self._current_diff, -1), 1)  # reward clip

    def _transit(self, a):
        s = self._rgb_state.copy()
        if a == 0:
            return s
        elif a == 1:
            return np.asarray(adjust_contrast(Image.fromarray(s), 0.95))
        elif a == 2:
            return np.asarray(adjust_contrast(Image.fromarray(s), 1.05))
        elif a == 3:
            return np.asarray(adjust_saturation(Image.fromarray(s), 0.95))
        elif a == 4:
            return np.asarray(adjust_saturation(Image.fromarray(s), 1.05))
        elif a == 5:
            return np.asarray(adjust_brightness(Image.fromarray(s), 0.95))
        elif a == 6:
            return np.asarray(adjust_brightness(Image.fromarray(s), 1.05))
        elif a == 7:
            return adjust_channels(s, 0.95, [0, 1])
        elif a == 8:
            return adjust_channels(s, 1.05, [0, 1])
        elif a == 9:
            return adjust_channels(s, 0.95, [2, 1])
        elif a == 10:
            return adjust_channels(s, 1.05, [2, 1])
        elif a == 11:
            return adjust_channels(s, 0.95, [0, 2])
        elif a == 12:
            return adjust_channels(s, 1.05, [0, 2])
        else:
            raise NotImplementedError

def get_lab_hist(lab):
    """Get hist of lab image"""
    lab = lab.reshape(-1, 3)
    hist, _ = np.histogramdd(lab, bins=(10, 10, 10),
                             range=((0, 100), (-60, 60), (-60, 60)))
    return hist.reshape(1, 1000) / 1000.0

def get_rgb_hist(lab):
    """Get hist of lab image"""
    lab = lab.reshape(-1, 3)
    hist, _ = np.histogramdd(lab, bins=(10, 10, 10),
                             range=((0, 255), (0, 255), (0, 255)))
    return hist.reshape(1, 1000) / 1000.0

def adjust_contrast(image_rgb, contrast_factor):
    """Adjust contrast"""
    enhancer = ImageEnhance.Contrast(image_rgb)
    img = enhancer.enhance(contrast_factor)
    return img

def adjust_saturation(image_rgb, saturation_factor):
    """Adjust saturation"""
    enhancer = ImageEnhance.Color(image_rgb)
    img = enhancer.enhance(saturation_factor)
    return img

def adjust_brightness(image_rgb, brightness_factor):
    """Adjust brightness"""
    enhancer = ImageEnhance.Brightness(image_rgb)
    img = enhancer.enhance(brightness_factor)
    return img

def adjust_channels(array, factor, channels):
    """Adjust channel pixel value"""
    for c in channels:
        array[:, :, c] = (array[:, :, c] * factor).clip(0, 255).astype('uint8')
    return array


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""
import numpy as np
from scipy.signal import detrend
import matplotlib.pyplot as plt
from joblib import Parallel, delayed
from mne.filter import filter_data, notch_filter
from scikits.samplerate import resample
from biosppy.signals import ecg

seg_mask_explanation = [
    'normal',
    'around sleep stage change point',
    'NaN in sleep stage',
    'NaN in signal',
    'overly high/low amplitude',
    'flat signal',
    'NaN in feature',
    'NaN in spectrum',
    'muscle artifact',
    'spurious spectrum',
    'missing or additional R peak']

def fast_resample(A, r, n_jobs=-1):
    """
    A: time points on the last axis
    r: newFs / oldFs
    """
    if A.ndim<=0:
        raise TypeError('ndim must >0')
    elif A.ndim==1:
        y = resample(A, r, 'sinc_best')
    else:
        shape = A.shape
        A = A.reshape(-1,shape[-1])
        y = Parallel(n_jobs=n_jobs, verbose=False)(delayed(resample)(A[i], r, 'sinc_best') for i in range(len(A)))
        y = np.array(y).astype(float)
        ylen = int(round(A.shape[-1]*r))
        if y.shape[1] < ylen:
            y = np.c_[y, np.repeat(y[:,[-1]], ylen-y.shape[1], axis=1)]
        elif y.shape[1] > ylen:
            y = y[:,:ylen]
        y = y.reshape(shape[:-1]+(y.shape[-1],))
    return y

def segment_ecg_signal(signal, labels, window_time, step_time, Fs, newFs=None, start_end_remove_window_num=1, n_jobs=-1, to_remove_mean=False, amplitude_thres=500):#, BW=4,
    """Segment ECG signals.

    Arguments:
    signal -- np.ndarray, size=(channel_num, sample_num)
    labels -- np.ndarray, size=(sample_num,), if none, pass in None
    window_time -- in seconds
    Fz -- in Hz

    Keyword arguments:
    notch_freq
    bandpass_freq
    start_end_remove_window_num -- default 0, number of windows removed at the beginning and the end of the signal
    amplitude_thres -- default 1000, mark all segments with np.any(signal_seg>=amplitude_thres)=True
    to_remove_mean -- default False, whether to remove the mean of signal from each channel

    Outputs:
    signal segments -- a list of np.ndarray, each has size=(window_size, channel_num)
    labels --  a list of labels of each window, if none, returns None
    segment start ids -- a list of starting ids of each window in (sample_num,)
    segment masks --
    """
    std_thres = 2
    std_thres2 = 5
    flat_seconds = 30

    if newFs is not None and Fs!=newFs:
        signal = fast_resample(signal, newFs*1./Fs)
        if labels is not None:
            labels = np.repeat(labels[::int(30*Fs)], int(30*newFs))
            minlen = min(signal.shape[-1], len(labels))
            signal = signal[...,:minlen]
            labels = labels[:minlen]
        Fs = newFs

    window_size = int(round(window_time*Fs))
    step_size = int(round(step_time*Fs))

    start_ids = np.arange(0, signal.shape[1]-window_size+1, step_size)
    if labels is not None:
        labels = labels[start_ids+window_size//2]
        if start_end_remove_window_num>0:
            start_ids = start_ids[start_end_remove_window_num:-start_end_remove_window_num]
            labels = labels[start_end_remove_window_num:-start_end_remove_window_num]
        assert len(start_ids)==len(labels)

    seg_masks = [seg_mask_explanation[0]]*len(start_ids)

    if labels is not None:
        if np.any(np.isnan(labels)):
            ids = np.where(np.isnan(labels))[0]
            for i in ids:
                seg_masks[i] = seg_mask_explanation[2]

    assert signal.shape[0]==1
    ecg_analysis_res = ecg.ecg(signal=signal.flatten(),sampling_rate=Fs,show=False)
    rpeaks = ecg_analysis_res['rpeaks']

    signal_segs = signal[:,list(map(lambda x:np.arange(x,x+window_size), start_ids))].transpose(1,0,2)  # (#window, #ch, window_size+2padding)

    ## find nan in signal

    nan2d = np.any(np.isnan(signal_segs), axis=2)
    nan1d = np.where(np.any(nan2d, axis=1))[0]
    for i in nan1d:
        seg_masks[i] = seg_mask_explanation[3]

    amplitude_large2d = np.any(np.abs(signal_segs)>amplitude_thres, axis=2)
    amplitude_large1d = np.where(np.any(amplitude_large2d, axis=1))[0]
    for i in amplitude_large1d:
        seg_masks[i] = seg_mask_explanation[4]

    flat_length = int(round(flat_seconds*Fs))
    # if there is any flat signal with flat_length
    short_segs = signal_segs.reshape(signal_segs.shape[0], signal_segs.shape[1], signal_segs.shape[2]//flat_length, flat_length)
    flat2d = np.any(detrend(short_segs, axis=3).std(axis=3)<=std_thres, axis=2)
    flat2d = np.logical_or(flat2d, np.std(signal_segs,axis=2)<=std_thres2)
    flat1d = np.where(np.any(flat2d, axis=1))[0]
    for i in flat1d:
        seg_masks[i] = seg_mask_explanation[5]

    # artifact removal for the rr peaks using ADARRI
    adarri = np.log1p(np.abs(np.diff(np.diff(rpeaks))))
    artifact_pos = np.where(adarri>=5)[0]+2
    rpeak_q10, rpeak_q90 = np.percentile(signal[0,rpeaks],(10,90))
    artifact_pos = artifact_pos.tolist() + np.where((signal[0,rpeaks]<rpeak_q10-300) | (signal[0,rpeaks]>rpeak_q90+300))[0].tolist()
    artifact_pos = np.sort(np.unique(artifact_pos))

    rr_peak_artifact1d = []
    for ap in artifact_pos:
        rr_peak_artifact1d.extend(np.where((start_ids<=rpeaks[ap]) & (rpeaks[ap]<start_ids+window_size))[0])
    rr_peak_artifact1d = np.sort(np.unique(rr_peak_artifact1d))
    for i in rr_peak_artifact1d:
        seg_masks[i] = seg_mask_explanation[10]

    # create binary signal
    signal = np.zeros_like(signal, dtype='double')
    signal[0,rpeaks] = 1.
    signal_segs = signal[:,list(map(lambda x:np.arange(x,x+window_size), start_ids))].transpose(1,0,2)  # (#window, #ch, window_size+2padding)

    # mark binary signals with all 0's
    fewbeats2d = signal_segs.sum(axis=2)<4.5*60/3 # can't have less than 20 beats/min
    fewbeats1d = np.where(np.any(fewbeats2d, axis=1))[0]
    for i in fewbeats1d:
        seg_masks[i] = seg_mask_explanation[10]

    return signal_segs, labels, start_ids, seg_masks#, specs, freq

def segment_chest_signal(signal, labels, window_time, step_time, Fs, newFs=None, notch_freq=None, bandpass_freq=None, start_end_remove_window_num=1, amplitude_thres=None, n_jobs=-1, to_remove_mean=False):#, BW=4
    """Segment CHEST signals.

    Arguments:
    signal -- np.ndarray, size=(channel_num, sample_num)
    labels -- np.ndarray, size=(sample_num,)
    window_time -- in seconds
    Fz -- in Hz

    Keyword arguments:
    notch_freq
    bandpass_freq
    start_end_remove_window_num -- default 0, number of windows removed at the beginning and the end of the signal
    amplitude_thres -- default 1000, mark all segments with np.any(signal_seg>=amplitude_thres)=True
    to_remove_mean -- default False, whether to remove the mean of signal from each channel

    Outputs:
    signal segments -- a list of np.ndarray, each has size=(window_size, channel_num)
    labels --  a list of labels of each window
    segment start ids -- a list of starting ids of each window in (sample_num,)
    segment masks --
    """
    std_thres1 = 5
    std_thres2 = 10

    window_size = int(round(window_time*Fs))
    step_size = int(round(step_time*Fs))

    start_ids = np.arange(0, signal.shape[1]-window_size+1, step_size)
    if len(start_ids)!=len(labels):
        labels = labels[start_ids+window_size//2]

    if start_end_remove_window_num>0:
        start_ids = start_ids[start_end_remove_window_num:-start_end_remove_window_num]
        labels = labels[start_end_remove_window_num:-start_end_remove_window_num]
    assert len(start_ids)==len(labels)

    seg_masks = [seg_mask_explanation[0]]*len(start_ids)

    if np.any(np.isnan(labels)):
        ids = np.where(np.isnan(labels))[0]
        for i in ids:
            seg_masks[i] = seg_mask_explanation[2]

    assert signal.shape[0]==1

    ## filter signal
    if notch_freq is not None and np.max(bandpass_freq)>=notch_freq:
        signal = notch_filter(signal, Fs, notch_freq, n_jobs=n_jobs, verbose='ERROR')  # (#window, #ch, window_size+2padding)
    signal = filter_data(detrend(signal, axis=1), Fs, bandpass_freq[0], bandpass_freq[1], n_jobs=n_jobs, verbose='ERROR')  # take the value starting from *padding*, (#window, #ch, window_size+2padding)
    signal_segs = signal[:,list(map(lambda x:np.arange(x,x+window_size), start_ids))].transpose(1,0,2)  # (#window, #ch, window_size+2padding)

    if newFs is not None and Fs!=newFs:
        signal_segs = fast_resample(signal_segs, newFs*1./Fs)
        Fs = newFs

    ## find nan in signal

    nan2d = np.any(np.isnan(signal_segs), axis=2)
    nan1d = np.where(np.any(nan2d, axis=1))[0]
    for i in nan1d:
        seg_masks[i] = '%s_%s'%(seg_mask_explanation[3], np.where(nan2d[i])[0])

    if amplitude_thres is not None:
        ## find large amplitude in signal

        amplitude_large2d = np.any(np.abs(signal_segs)>amplitude_thres, axis=2)
        amplitude_large1d = np.where(np.any(amplitude_large2d, axis=1))[0]
        for i in amplitude_large1d:
            seg_masks[i] = '%s_%s'%(seg_mask_explanation[4], np.where(amplitude_large2d[i])[0])

    ## find flat signal

    flat_seconds = 6
    flat_length = int(round(flat_seconds*Fs))
    assert signal_segs.shape[2]//flat_length*flat_length==signal_segs.shape[2]
    short_segs = signal_segs.reshape(signal_segs.shape[0], signal_segs.shape[1], signal_segs.shape[2]//flat_length, flat_length)
    flat2d = np.any(detrend(short_segs, axis=3).std(axis=3)<=std_thres1, axis=2)
    flat2d = np.logical_or(flat2d, np.std(signal_segs,axis=2)<=std_thres2)
    flat1d = np.where(np.any(flat2d, axis=1))[0]
    for i in flat1d:
        seg_masks[i] = '%s_%s'%(seg_mask_explanation[5], np.where(flat2d[i])[0])

    return signal_segs, labels, start_ids, seg_masks#, specs, freq

reason: import numpy,from scipy


text:
import functions
from scipy import optimize

def exponentialFit(xData, yData): #bounds will be hardcoded for now
    popt, pcov = optimize.curve_fit(functions.exponential,
    xData, yData, bounds = ((1e-05, 0, -15), (1, 5e-01, 15)), sigma=xData ** 2, absolute_sigma=True)

    return popt, pcov

def logisticfit(xData, yData):
    popt, pcov = optimize.curve_fit(functions.logistic,
    xData, yData, bounds = ((0, 0, 0), (1000000, 500, 1)), sigma = xData **2, absolute_sigma=True)

    return popt, pcov

def logisticDistributionFit(xData, yData):
    popt, pcov = optimize.curve_fit(functions.logisticDistribution,
    xData, yData) #TODO: implement bounds based on real world observations

    return popt, pcov

max_stars_repo_name: siddharthbharthulwar/nCov-Tracker

max_issues_repo_name: siddharthbharthulwar/nCov-Tracker

max_forks_repo_name: siddharthbharthulwar/nCov-Tracker


text:
#!/usr/bin/env python
## Author: Narmada M. Balasooriya   ##
##         University of Peradeniya ##
##         Sri Lanka                ##
## Editor: Leontios Mavropalias     ##
##         UoC                      ##

#####################################
## Import the necessary libraries ###
#####################################

from __future__ import division, print_function, absolute_import

from skimage import color, io
import numpy as np
import os
from glob import glob
import cv2

from six.moves import cPickle
import pickle

np.set_printoptions(suppress=True)
########################################
### Imports picture files
########################################

# TumorA = 0
# TumorB = 1
# TumorC = 2

###vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
### Change below for your dataset
files_path_tumorA = '/mnt/share/mri_dataset/cheng/lbl_orig/1/'
files_path_tumorB = '/mnt/share/mri_dataset/cheng/lbl_orig/2/'
files_path_tumorC = '/mnt/share/mri_dataset/cheng/lbl_orig/3/'

tumorA_path = os.path.join(files_path_tumorA, '*.tiff')
tumorB_path = os.path.join(files_path_tumorB, '*.tiff')
tumorC_path = os.path.join(files_path_tumorC, '*.tiff')
###^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

print("tumor A path")

# BUG: only tumor A is fuled, the rest empty
tumorA = sorted(glob(tumorA_path))
tumorB = sorted(glob(tumorB_path))
tumorC = sorted(glob(tumorC_path))

n_files = len(tumorA) + len(tumorB) + len(tumorC)
print("######print here")
print(n_files)
print("##########")

size_image = 64

allX = np.zeros((n_files, size_image, size_image, 3), dtype='uint8')
ally = np.zeros((n_files), dtype='int32')
count = 0
y_count = 0
for f in tumorA:
    try:
        #img = io.imread(f)
        img = cv2.imread(f)
        new_img = cv2.resize(img,(size_image,size_image))
        allX[count] = np.array(new_img)
        ally[y_count] = 0
        count += 1
        y_count += 1
    except exception as e:
        print("Tumour A not set")
        raise e

print("tumorA done")
for f in tumorB:
    try:
        #img = io.imread(f)
        img = cv2.imread(f)
        new_img = cv2.resize(img,(size_image,size_image))
        allX[count] = np.array(new_img)
        ally[y_count] = 1
        count += 1
        y_count += 1
    except exception as e:
        print("Tumour B not set")
        raise e
print("tumorB done")

for f in tumorC:
    try:
        #img = io.imread(f)
        img = cv2.imread(f)
        new_img = cv2.resize(img,(size_image,size_image))
        allX[count] = np.array(new_img)
        ally[y_count] = 2
        count += 1
        y_count += 1
    except exception as e:
        print("Tumour C not set")
        raise e

print("tumorC done")
print("data are split")

pkl_fname = 'dataset_cheng_full_nonseg_lbl.pkl'
f = open(pkl_fname, 'wb')

print("pickle file open")
cPickle.dump((allX, ally), f, protocol=cPickle.HIGHEST_PROTOCOL)
print("pickle dumped at:\n%s" % pkl_fname)
f.close()

print("finished")