source: orange-bioinformatics/obiGO.py @ 1609:6b0bb985598f

Revision 1609:6b0bb985598f, 61.7 KB checked in by Ales Erjavec <ales.erjavec@…>, 2 years ago (diff)

Can now specify specific annotation and ontology revision.

Line 
1"""obiGO is a Gene Ontology (GO) Handling Library.
2
3"""
4
5from urllib import urlretrieve
6from gzip import GzipFile
7import tarfile
8import gzip
9import shutil
10import sys, os
11import re, cPickle
12from datetime import datetime
13from collections import defaultdict
14
15import obiProb, obiGene
16import orngEnviron
17
18try:
19    import orngServerFiles
20    default_database_path = os.path.join(orngServerFiles.localpath(), "GO")
21except Exception:
22    default_database_path = os.curdir
23
24_CVS_REVISION_RE = re.compile(r"^(rev)?\d+\.\d+$")
25
26evidenceTypes = {
27##Experimental
28    'EXP': 'Inferred from Experiment',
29    'IDA': 'Inferred from Direct Assay',
30    'IPI': 'Inferred from Physical Interaction', ## [with <database:protein_name>]',
31    'IMP': 'Inferred from Mutant Phenotype',
32    'IGI': 'Inferred from Genetic Interaction', ## [with <database:gene_symbol[allele_symbol]>]',
33    'IEP': 'Inferred from Expression Pattern',
34##Computational Analysis Evidence Codes
35    'ISS': 'Inferred from Sequence Similarity', ## [with <database:sequence_id>] ',
36    'ISA': 'Inferred from Sequence Alignment',
37    'ISO': 'Inferred from Sequence Orthology',
38    'ISM': 'Inferred from Sequence Model',
39    'IGC': 'Inferred from Genomic Context',
40    'RCA': 'Inferred from Reviewed Computational Analysis',
41##Author Statement Evidence Codes
42    'TAS': 'Traceable author statement',
43    'NAS': 'Non-traceable author statement',
44##Curatorial Statement Evidence Codes
45    'IC': 'Inferred by curator',
46    'ND': 'No biological data available',
47##Computationally-assigned Evidence Codes
48    'IEA': 'Inferred from electronic annotation', ## [to <database:id>]',
49##Obsolete Evidence Codes
50    'NR': 'Not Recorded(Obsolete)'
51}
52##evidenceDict={"IMP":1, "IGI":2, "IPI":4, "ISS":8, "IDA":16, "IEP":32, "IEA":64,
53##              "TAS":128, "NAS":256, "ND":512, "IC":1024, "RCA":2048, "IGC":4096, "RCA":8192, "NR":16384}
54
55evidenceDict=defaultdict(int, [(e, 2**i) for i, e in enumerate(evidenceTypes.keys())])
56
57evidenceTypesOrdered = [
58'EXP',
59'IDA',
60'IPI',
61'IMP',
62'IGI',
63'IEP',
64##Computational Analysis Evidence Codes
65'ISS',
66'ISA',
67'ISO',
68'ISM',
69'IGC',
70'RCA',
71##Author Statement Evidence Codes
72'TAS',
73'NAS',
74##Curatorial Statement Evidence Codes
75'IC',
76'ND',
77##Computationally-assigned Evidence Codes
78'IEA',
79##Obsolete Evidence Codes
80'NR'
81]
82
83multiplicitySet=set(["alt_id","is_a","subset","synonym","related_synonym",
84                     "exact_synonym","broad_synonym","narrow_synonym",
85                     "xref_analog","xref_unknown","relationship"])
86
87multipleTagSet = multiplicitySet
88
89annotationFields=["DB","DB_Object_ID","DB_Object_Symbol","Qualifier","GO_ID",
90                  "DB_Reference","Evidence_Code","With_From","Aspect",
91                  "DB_Object_Name","DB_Object_Synonym","DB_Object_Type",
92                  "Taxon","Date","Assigned_by"]
93
94annotationFieldsDict={"DB":0,
95                      "DB_Object_ID":1,
96                      "DB_Object_Symbol":2,
97                      "Qualifier":3,
98                      "GO_ID":4,
99                      "GO ID":4,
100                      "GOID":4,
101                      "DB_Reference":5,
102                      "DB:Reference":5,
103                      "Evidence_Code":6,
104                      "Evidence Code":6,
105                      "Evidence_code":6, #compatibility with  older revisions
106                      "With_or_From":7,
107                      "With (or) From":7,
108                      "Aspect":8,
109                      "DB_Object_Name":9,
110                      "DB_Object_Synonym":10,
111                      "DB_Object_Type":11,
112                      "taxon":12,
113                      "Date":13,
114                      "Assigned_by":14}   
115
116builtinOBOObjects = ["""
117[Typedef]
118id: is_a
119name: is_a
120range: OBO:TERM_OR_TYPE
121domain: OBO:TERM_OR_TYPE
122definition: The basic subclassing relationship [OBO:defs]"""
123,
124"""[Typedef]
125id: disjoint_from
126name: disjoint_from
127range: OBO:TERM
128domain: OBO:TERM
129definition: Indicates that two classes are disjoint [OBO:defs]"""
130,
131"""[Typedef]
132id: instance_of
133name: instance_of
134range: OBO:TERM
135domain: OBO:INSTANCE
136definition: Indicates the type of an instance [OBO:defs]"""
137,
138"""[Typedef]
139id: inverse_of
140name: inverse_of
141range: OBO:TYPE
142domain: OBO:TYPE
143definition: Indicates that one relationship type is the inverse of another [OBO:defs]"""
144,
145"""[Typedef]
146id: union_of
147name: union_of
148range: OBO:TERM
149domain: OBO:TERM
150definition: Indicates that a term is the union of several others [OBO:defs]"""
151,
152"""[Typedef]
153id: intersection_of
154name: intersection_of
155range: OBO:TERM
156domain: OBO:TERM
157definition: Indicates that a term is the intersection of several others [OBO:defs]"""]
158
159class OBOObject(object):
160    """ Represents a generic OBO object (e.g. Term, Typedef, Instance, ...)
161    Example:
162    >>> OBOObject(r"[Term]\nid: FOO:001\nname: bar", ontology)
163    """
164    _INTERN_TAGS = ["id", "name", "namespace", "alt_id", "is_a"]
165    def __init__(self, stanza=None, ontology=None):
166        self.ontology = ontology
167        self._lines = []
168        self.values = {}
169        self.related = set()
170        self.relatedTo = set()
171        if stanza:
172            self.ParseStanza(stanza)
173
174    def ParseStanza(self, stanza):
175        intern_tags = set(self._INTERN_TAGS)
176        for line in stanza.split("\n"):
177            if ":" not in line:
178                continue
179            tag, rest = line.split(":", 1)
180            value, modifiers, comment = "", "", ""
181            if "!" in rest:
182                rest, comment = rest.split("!")
183            if "{" in rest:
184                value, modifiers = rest.split("{", 1)
185                modifiers = modifiers.strip("}")
186            else:
187                value = rest
188            tag = intern(tag)
189            value = value.strip()
190            comment = comment.strip()
191            if tag in intern_tags:
192                value, comment = intern(value), intern(comment)
193            self._lines.append((tag, value, modifiers, comment))
194            if tag in multipleTagSet:
195                self.values.setdefault(tag, []).append(value)
196            else:
197                self.values[tag] = value
198        self.related = set(self.GetRelatedObjects())
199        self.__dict__.update(self.values)
200        if "def" in self.__dict__:
201            self.__dict__["def_"] = self.__dict__["def"]
202       
203
204    def GetRelatedObjects(self):
205        """ Return a list of tuple pairs where the first element is relationship
206        typeId and the second id of object to whom the relationship applys to.
207        """
208        ##TODO: add other defined Typedef ids
209        typeIds = [intern("is_a")]
210        result = [(typeId, id) for typeId in typeIds for id in self.values.get(typeId, [])] 
211        result = result + [tuple(map(intern, r.split(None, 1))) for r in self.values.get("relationship", [])]
212        return result
213
214    def __repr__(self):
215        """ Return a string representation of the object in OBO format
216        """
217        repr = "[%s]\n" % type(self).__name__
218        for tag, value, modifiers, comment in self._lines:
219            repr = repr + tag + ": " + value
220            if modifiers:
221                repr = repr + "{ " + modifiers + " }"
222            if comment:
223                repr = repr + " ! " + comment
224            repr = repr + "\n"
225        return repr
226
227    def __str__(self):
228        """ Return the OBO object id entry
229        """
230        return "%s: %s" % (self.id, self.name)
231
232    def __iter__(self):
233        """ Iterates over sub terms
234        """
235        for typeId, id in self.relatedTo:
236            yield (typeId, self.ontology[id])
237       
238class Term(OBOObject):
239    pass
240
241class Typedef(OBOObject):
242    pass
243
244class Instance(OBOObject):
245    pass
246       
247class Ontology(object):
248    """ Ontology is the main class representing a gene ontology.
249   
250    Example::
251        >>> ontology = Ontology("my_gene_ontology.obo")
252       
253    """
254    version = 1
255    def __init__(self, file=None, progressCallback=None, rev=None):
256        """ Initialize the ontology from file (if `None` the default gene
257        ontology will be loaded). The optional `progressCallback` will be
258        called with a single argument to report on the progress.
259       
260        """
261        self.terms = {}
262        self.typedefs = {}
263        self.instances = {}
264        self.slimsSubset = set()
265        if isinstance(file, basestring):
266            self.ParseFile(file, progressCallback)
267           
268        elif rev is not None:
269            if not _CVS_REVISION_RE.match(rev):
270                raise ValueError("Invalid revision format.")
271            if rev.startswith("rev"):
272                rev = rev[3:]
273            pc = lambda v: progressCallback(v/2.0) \
274                 if progressCallback else None
275            filename = os.path.join(default_database_path, "gene_ontology_edit@rev%s.obo" % rev)
276            if not os.path.exists(filename):
277                self.DownloadOntologyAtRev(rev, filename, pc)
278            self.ParseFile(filename, lambda v: progressCallback(v/2.0 + 50) \
279                                     if progressCallback else None)
280        else:
281            fool = self.Load(progressCallback)
282            self.__dict__ = fool.__dict__ ## A fool and his attributes are soon parted
283
284    @classmethod
285    def Load(cls, progressCallback=None):
286        """ A class method that tries to load the ontology file from
287        default_database_path. It looks for a filename starting with
288        'gene_ontology'. If not found it will download it.
289        """
290        filename = os.path.join(default_database_path, "gene_ontology_edit.obo.tar.gz")
291        if not os.path.isfile(filename) and not os.path.isdir(filename):
292            import orngServerFiles
293            orngServerFiles.download("GO", "gene_ontology_edit.obo.tar.gz")
294        try:
295            return cls(filename, progressCallback=progressCallback)
296        except (IOError, OSError), ex:
297            print ex
298            raise Exception, "Could not locate ontology file"
299       
300    def ParseFile(self, file, progressCallback=None):
301        """ Parse the file. file can be a filename string or an open filelike
302        object. The optional progressCallback will be called with a single
303        argument to report on the progress.
304        """
305        if type(file) == str:
306            if os.path.isfile(file) and tarfile.is_tarfile(file):
307                f = tarfile.open(file).extractfile("gene_ontology_edit.obo")
308            elif os.path.isfile(file):
309                f = open(file)
310            else:
311                f = open(os.path.join(file, "gene_ontology_edit.obo"))
312        else:
313            f = file
314       
315        data = f.readlines()
316        data = "".join([line for line in data if not line.startswith("!")])
317        self.header = data[: data.index("[Term]")]
318        c=re.compile("\[.+?\].*?\n\n", re.DOTALL)
319        data=c.findall(data)
320
321        import orngMisc
322        milestones = orngMisc.progressBarMilestones(len(data), 90)
323        for i, block in enumerate(builtinOBOObjects + data):
324            if block.startswith("[Term]"):
325                term = Term(block, self)
326                self.terms[term.id] = term
327            elif block.startswith("[Typedef]"):
328                typedef = Typedef(block, self)
329                self.typedefs[typedef.id] = typedef
330            elif block.startswith("[Instance]"):
331                instance = Instance(block, self)
332                self.instances[instance.id] = instance
333            if progressCallback and i in milestones:
334                progressCallback(90.0*i/len(data))
335       
336        self.aliasMapper = {}
337        self.reverseAliasMapper = defaultdict(set)
338        milestones = orngMisc.progressBarMilestones(len(self.terms), 10)
339        for i, (id, term) in enumerate(self.terms.iteritems()):
340            for typeId, parent in term.related:
341                self.terms[parent].relatedTo.add((typeId, id))
342            try:
343                self.aliasMapper.update([(alt_id, id) for alt_id in term.alt_id])
344                self.reverseAliasMapper[id].union_update(term.alt_id)
345            except AttributeError:
346                pass
347            if progressCallback and i in milestones:
348                progressCallback(90.0 + 10.0*i/len(self.terms))
349
350    def GetDefinedSlimsSubsets(self):
351        """ Return a list of defined subsets.
352        """
353        return [line.split()[1] for line in self.header.split("\n") if line.startswith("subsetdef:")]
354
355    def SetSlimsSubset(self, subset):
356        """ Set the slims term subset to subset. If subset is a string it
357        must equal one of the defined subsetdef.
358        """
359        if type(subset) == str:
360            self.slimsSubset = [id for id, term in self.terms.items() if subset in getattr(term, "subset", set())]
361        else:
362            self.slimsSubset = set(subset)
363
364    def GetSlimsSubset(self, subset):
365        return [id for id, term in self.terms.items() if subset in getattr(term, "subset", set())]
366
367    def GetSlimTerms(self, termId):
368        """ Return a list of slim terms for termId.
369        """
370        queue = set([termId])
371        visited = set()
372        slims = set()
373        while queue:
374            term = queue.pop()
375            visited.add(term)
376            if term in self.slimsSubset:
377                slims.add(term)
378            else:
379                queue.update(set(id for typeId, id in self[term].related) - visited)
380        return slims
381
382    def ExtractSuperGraph(self, terms):
383        """ Return all super terms of terms up to the most general one.
384        """
385        terms = [terms] if type(terms) == str else terms
386        visited = set()
387        queue = set(terms)
388        while queue:
389            term = queue.pop()
390            visited.add(term)
391            queue.update(set(id for typeId, id in self[term].related) - visited)
392        return visited
393
394    def ExtractSubGraph(self, terms):
395        """ Return all sub terms of terms.
396        """
397        terms = [terms] if type(terms) == str else terms
398        visited = set()
399        queue = set(terms)
400        while queue:
401            term = queue.pop()
402            visited.add(term)
403            queue.update(set(id for typeId, id in self[term].relatedTo) - visited)
404        return visited
405
406    def GetTermDepth(self, term, cache_={}):
407        """ Return the minimum depth of a term (length of the shortest
408        path to this term from the top level term).
409        """
410        if term not in cache_:
411            cache_[term] = min([self.GetTermDepth(parent) + 1 for typeId, parent in self[term].related] or [1])
412        return cache_[term]
413
414    def __getitem__(self, id):
415        """ Return object with id (same as ontology.terms[id]
416        """
417        return self.terms.get(id, self.terms.get(self.aliasMapper.get(id, id)))
418
419    def __iter__(self):
420        """ Iterate over all ids in ontology
421        """
422        return iter(self.terms)
423
424    def __len__(self):
425        """ Return number of objects in ontology
426        """
427        return len(self.terms)
428
429    def __contains__(self, id):
430        return id in self.terms or id in self.aliasMapper
431
432    @staticmethod
433    def DownloadOntology(file, progressCallback=None):
434        tFile = tarfile.open(file, "w:gz") if isinstance(file, basestring) else file
435        tmpDir = os.path.join(orngEnviron.bufferDir, "tmp_go/")
436        try:
437            os.mkdir(tmpDir)
438        except Exception:
439            pass
440        urlretrieve("http://www.geneontology.org/ontology/gene_ontology_edit.obo", 
441                    os.path.join(tmpDir, "gene_ontology_edit.obo"),
442                    progressCallback and __progressCallbackWrapper(progressCallback))
443        tFile.add(os.path.join(tmpDir, "gene_ontology_edit.obo"),
444                  "gene_ontology_edit.obo")
445        tFile.close()
446        os.remove(os.path.join(tmpDir, "gene_ontology_edit.obo"))
447
448    @staticmethod
449    def DownloadOntologyAtRev(rev, filename=None, progressCallback=None):
450        import shutil
451        import urllib2
452        url = "http://cvsweb.geneontology.org/cgi-bin/cvsweb.cgi/~checkout~/go/ontology/gene_ontology_edit.obo?rev=%s" % rev
453        url += ";content-type=text%2Fplain"
454        if filename is None:
455            filename = os.path.join(default_database_path, "gene_ontology_edit@rev%s.obo" % rev)
456        r = urllib2.urlopen(url)
457        shutil.copyfileobj(r, open(filename, "wb"))
458       
459       
460_re_obj_name_ = re.compile("([a-zA-z0-9-_]+)")
461
462class AnnotationRecord(object):
463    """Holds the data for an annotation record read from the annotation file.
464    Fields can be accessed with the names: DB, DB_Object_ID, DB_Object_Symbol,
465    Qualifier, GO_ID, DB_Reference, Evidence_code, With_or_From, Aspect,
466    DB_Object_Name, DB_Object_Synonym, DB_Object_Type, taxon, Date,
467    Assigned_by (e.g. rec.GO_ID) or by supplying the original name of the
468    field (see http://geneontology.org/GO.annotation.shtml#file) to the get
469    method (e.g. rec.get("GO ID")) The object also provides the following
470    data members for quicker access: geneName, GOId, evidence, aspect and
471    alias(a list of aliases)
472   
473    """
474    __slots__ = annotationFields + ["geneName", "GOId", "evidence",
475                                    "aspect", "alias", "additionalAliases"]
476    def __init__(self, fullText):
477        """\
478        :param fulText: A single line from the annotation file.
479       
480        """
481        for slot, val in zip(annotationFields, fullText.split("\t")):
482            setattr(self, slot, intern(val))
483        self.geneName = self.DB_Object_Symbol
484        self.GOId = self.GO_ID
485        self.evidence = self.Evidence_Code
486        self.aspect = self.Aspect
487        self.alias = list(map(intern, self.DB_Object_Synonym.split("|")))
488
489        self.additionalAliases = []
490        if ":" in self.DB_Object_Name:
491            self.additionalAliases = [] #_re_obj_name_.findall(self.DB_Object_Name.split(":")[0])
492
493    def __getattr__(self, name):
494        if name in annotationFieldsDict:
495            return getattr(self, self.__slots__[annotationFieldsDict[name]])
496        else:
497            raise AttributeError(name)
498
499
500class Annotations(object):
501    """Annotations object holds the annotations.
502    """
503    version = 2
504    def __init__(self, file=None, ontology=None, genematcher=None, progressCallback=None, rev=None):
505        """Initialize the annotations from file by calling ParseFile on it.
506        The ontology must be an instance of Ontology class.
507        The optional progressCallback will be called with a single argument
508        to report on the progress.
509       
510        """
511        self.file = file
512        self.ontology = ontology
513        self.allAnnotations = defaultdict(list)
514        self.geneAnnotations = defaultdict(list)
515        self.termAnnotations = defaultdict(list)
516        self._geneNames = None
517        self._geneNamesDict = None
518        self._aliasMapper = None
519        self.additionalAliases = {}
520        self.annotations = []
521        self.header = ""
522        self.genematcher = genematcher
523        self.taxid = None
524        if type(file) in [list, set, dict, Annotations]:
525            for ann in file:
526                self.AddAnnotation(ann)
527            if type(file, Annotations):
528                taxid = file.taxid
529        elif isinstance(file, basestring) and os.path.exists(file):
530            self.ParseFile(file, progressCallback)
531            try:
532                self.taxid = to_taxid(os.path.basename(file).split(".")[1]).pop()
533            except IOError:
534                pass
535        elif file is not None:
536            # Organism code
537            if rev is not None:
538                if not _CVS_REVISION_RE.match(rev):
539                    raise ValueError("Invalid revision format")
540               
541                if rev.startswith("rev"):
542                    rev = rev[3:]
543                code = self.organism_name_search(file)
544                filename = os.path.join(default_database_path,
545                                        "gene_association.%s@rev%s.tar.gz" % (code, rev))
546               
547                if not os.path.exists(filename):
548                    self.DownloadAnnotationsAtRev(code, rev, filename,
549                                                  progressCallback
550                                                  )
551                   
552                self.ParseFile(filename, progressCallback)
553                self.taxid = to_taxid(code).pop()
554            else:
555                a = self.Load(file, ontology, genematcher, progressCallback)
556                self.__dict__ = a.__dict__
557                self.taxid = to_taxid(organism_name_search(file)).pop()
558        if not self.genematcher and self.taxid:
559            import obiGene
560            self.genematcher = obiGene.matcher([obiGene.GMGO(self.taxid)] + \
561                                               ([obiGene.GMDicty(), [obiGene.GMGO(self.taxid),
562                                                                     obiGene.GMDicty()]] \
563                                                if self.taxid == "352472"  else []))
564        if self.genematcher:
565            self.genematcher.set_targets(self.geneNames)
566       
567    @classmethod
568    def organism_name_search(cls, org):
569        ids = to_taxid(org) 
570        import obiTaxonomy as tax
571        if not ids:
572            ids = [org] if org in Taxonomy().common_org_map.keys() + Taxonomy().code_map.keys() else []
573        if not ids:
574            ids = tax.to_taxid(org, mapTo=Taxonomy().keys())
575        if not ids:
576            ids = tax.search(org, exact=True)
577            ids = set(ids).intersection(Taxonomy().keys())
578        if not ids:
579            ids = tax.search(org)
580            ids = set(ids).intersection(Taxonomy().keys())
581        codes = set([from_taxid(id) for id in ids])
582        if len(codes) > 1:
583            raise tax.MultipleSpeciesException, ", ".join(["%s: %s" % \
584                            (str(from_taxid(id)), tax.name(id)) for id in ids])
585        elif len(codes) == 0:
586            raise tax.UnknownSpeciesIdentifier, org
587        return codes.pop()
588
589    @classmethod
590    def organism_version(cls, name):
591        name = organism_name_search(name)
592        orngServerFiles.localpath_download("GO", "gene_association.%s.tar.gz" % name)
593        return ("v%i." % cls.version) + orngServerFiles.info("GO",
594                        "gene_association.%s.tar.gz" % name)["datetime"]
595
596    def SetOntology(self, ontology):
597        """ Set the ontology to use in the annotations mapping.
598        """
599        self.allAnnotations = defaultdict(list)
600        self._ontology = ontology
601
602    def GetOntology(self):
603        return self._ontology
604
605    ontology = property(GetOntology, SetOntology,
606                        doc="Ontology object for annotations")
607   
608    @classmethod
609    def Load(cls, org, ontology=None, genematcher=None, progressCallback=None):
610        """A class method that tries to load the association file for the
611        given organism from default_database_path.
612        """
613        import orngServerFiles
614        code = organism_name_search(org)
615       
616        file = "gene_association.%s.tar.gz" % code
617
618        path = os.path.join(orngServerFiles.localpath("GO"), file)
619        if not os.path.exists(path):
620            orngServerFiles.download("GO", file)
621        return cls(path, ontology=ontology, genematcher=genematcher, progressCallback=progressCallback)
622   
623    def ParseFile(self, file, progressCallback=None):
624        """ Parse and load the annotations from file. Report progress
625        with progressCallback.
626        File can be:
627            - a tarball containing the association file named gene_association
628            - a directory name containing the association file named gene_association
629            - a path to the actual association file
630            - an open file-like object of the association file
631        """
632        if type(file) == str:
633            if os.path.isfile(file) and tarfile.is_tarfile(file):
634                f = tarfile.open(file).extractfile("gene_association")
635            elif os.path.isfile(file) and file.endswith(".gz"):
636                f = gzip.open(file) 
637            elif os.path.isfile(file):
638                f = open(file)
639            else:
640                f = open(os.path.join(file, "gene_association"))
641        else:
642            f = file
643        lines = [line for line in f.read().split("\n") if line.strip()]
644        import orngMisc
645        milestones = orngMisc.progressBarMilestones(len(lines), 100)
646        for i,line in enumerate(lines):
647            if line.startswith("!"):
648                self.header = self.header + line + "\n"
649                continue
650           
651            a=AnnotationRecord(line)
652            self.AddAnnotation(a)
653#            self.annotations.append(a)
654            if progressCallback and i in milestones:
655                progressCallback(100.0*i/len(lines))
656
657    def AddAnnotation(self, a):
658        """ Add a single `AnotationRecord` instance to this `Annotations`
659        object.
660        """
661        if not isinstance(a, AnnotationRecord):
662            a = AnnotationRecord(a)
663        if not a.geneName or not a.GOId or a.Qualifier == "NOT":
664            return
665       
666        self.geneAnnotations[a.geneName].append(a)
667        self.annotations.append(a)
668        self.termAnnotations[a.GOId].append(a)
669        self.allAnnotations = defaultdict(list)
670       
671        self._geneNames = None
672        self._geneNamesDict = None
673        self._aliasMapper = None
674
675    @property
676    def geneNamesDict(self):
677        if getattr(self, "_geneNamesDict", None) is None:
678            self._geneNamesDict = defaultdict(set)
679            for alias, name in self.aliasMapper.iteritems():
680                self._geneNamesDict[name].add(alias)
681        return self._geneNamesDict
682
683    @property
684    def geneNames(self):
685        if getattr(self, "_geneNames", None) is None:
686            self._geneNames = set([ann.geneName for ann in self.annotations])
687        return self._geneNames
688
689    @property
690    def aliasMapper(self):
691        if getattr(self, "_aliasMapper", None) is None:
692            self._aliasMapper = {}
693            for ann in self.annotations:
694                self._aliasMapper.update([(alias, ann.geneName) for alias in ann.alias + \
695                                          [ann.geneName, ann.DB_Object_ID]])
696        return self._aliasMapper
697
698    def GetGeneNamesTranslator(self, genes):
699        """ Return a dictionary mapping canonical names (DB_Object_Symbol)
700        to `genes`.
701         
702        """
703        def alias(gene):
704            if self.genematcher:
705                return self.genematcher.umatch(gene)
706            else:
707                return gene if gene in self.geneNames else \
708                        self.aliasMapper.get(gene,
709                             self.additionalAliases.get(gene, None))
710                       
711        return dict([(alias(gene), gene) for gene in genes if alias(gene)])
712
713    def _CollectAnnotations(self, id, visited):
714        """ Recursive function collects and caches all annotations for id
715        """
716        if id not in self.allAnnotations and id not in visited:
717            if id in self.ontology.reverseAliasMapper:
718                annotations = [self.termAnnotations.get(alt_id, []) \
719                               for alt_id in self.ontology.reverseAliasMapper[id]] + \
720                                             [self.termAnnotations[id]]
721            else:
722                ## annotations for this term alone
723                annotations = [self.termAnnotations[id]] 
724            visited.add(id)
725            for typeId, child in self.ontology[id].relatedTo:
726                aa = self._CollectAnnotations(child, visited)
727                if type(aa) == set:
728                    ## if it was already reduced in GetAllAnnotations
729                    annotations.append(aa)
730                else:
731                    annotations.extend(aa)
732            self.allAnnotations[id] = annotations
733        return self.allAnnotations[id]
734
735    def GetAllAnnotations(self, id):
736        """ Return a set of all annotations (instances if `AnnotationRectord`)
737        for GO term `id` and all it's subterms.
738       
739        :param id: GO term id
740        :type id: str
741       
742        """
743        visited = set()
744        id = self.ontology.aliasMapper.get(id, id)
745        if id not in self.allAnnotations or type(self.allAnnotations[id]) == list:
746            annot_set = set()
747            for annots in self._CollectAnnotations(id, set()):
748                annot_set.update(annots)
749            self.allAnnotations[id] = annot_set
750        return self.allAnnotations[id]
751
752    def GetAllGenes(self, id, evidenceCodes = None):
753        """ Return a list of genes annotated by specified `evidenceCodes`
754        to GO term 'id' and all it's subterms."
755       
756        :param id: GO term id
757        :type id: str
758       
759        :param evidneceCodes: List of evidence codes to consider when
760                              matching annotations to terms.
761        :type evidenceCodes: list-of-strings
762        """
763        evidenceCodes = set(evidenceCodes or evidenceDict.keys())
764        annotations = self.GetAllAnnotations(id)
765        return list(set([ann.geneName for ann in annotations if ann.Evidence_Code in evidenceCodes]))
766
767    def GetEnrichedTerms(self, genes, reference=None, evidenceCodes=None, 
768                         slimsOnly=False, aspect=None, prob=obiProb.Binomial(),
769                         useFDR=True, progressCallback=None):
770        """ Return a dictionary of enriched terms, with tuples of
771        (list_of_genes, p_value, reference_count) for items and term
772        ids as keys. P-Values are FDR adjusted if useFDR is True (default).
773       
774        :param genes: List of genes
775        :param reference: list of genes (if None all genes included in the
776                          annotations will be used).
777        :param evidenceCodes: List of evidence codes to consider.
778        :param slimsOnly: If `True` return only slim terms
779        :param aspect: Which aspects to use. Use all by default. "P", "F", "C"
780            or a set containing these elements.
781        """
782        revGenesDict = self.GetGeneNamesTranslator(genes)
783        genes = set(revGenesDict.keys())
784        if reference:
785            refGenesDict = self.GetGeneNamesTranslator(reference)
786            reference = set(refGenesDict.keys())
787        else:
788            reference = self.geneNames
789
790        if aspect == None:
791            aspects_set = set(["P", "C", "F"])
792        else:
793            aspects_set = set([aspect]) if isinstance(aspect, basestring) else aspect
794
795        evidenceCodes = set(evidenceCodes or evidenceDict.keys())
796        annotations = [ann for gene in genes for ann in self.geneAnnotations[gene] \
797                       if ann.Evidence_Code in evidenceCodes and ann.Aspect in aspects_set]
798        refAnnotations = set([ann for gene in reference for ann in self.geneAnnotations[gene] \
799                              if ann.Evidence_Code in evidenceCodes and ann.Aspect in aspects_set])
800        annotationsDict = defaultdict(set)
801        for ann in annotations:
802            annotationsDict[ann.GO_ID].add(ann)
803           
804        if slimsOnly and not self.ontology.slimsSubset:
805            import warnings
806            warnings.warn("Unspecified slims subset in the ontology! Using 'goslim_generic' subset", UserWarning)
807            self.ontology.SetSlimsSubset("goslim_generic")
808           
809        terms = self.ontology.ExtractSuperGraph(annotationsDict.keys())
810        res = {}
811        import orngMisc
812        milestones = orngMisc.progressBarMilestones(len(terms), 100)
813        for i, term in enumerate(terms):
814            if slimsOnly and term not in self.ontology.slimsSubset:
815                continue
816            allAnnotations = self.GetAllAnnotations(term).intersection(refAnnotations)
817##            allAnnotations.intersection_update(refAnnotations)
818            allAnnotatedGenes = set([ann.geneName for ann in allAnnotations])
819            mappedGenes = genes.intersection(allAnnotatedGenes)
820##            if not mappedGenes:
821##                print >> sys.stderr, term, sorted(genes)
822##                print >> sys.stderr, sorted(allAnnotatedGenes)
823##                return
824            if len(reference) > len(allAnnotatedGenes):
825                mappedReferenceGenes = reference.intersection(allAnnotatedGenes)
826            else:
827                mappedReferenceGenes = allAnnotatedGenes.intersection(reference)
828            res[term] = ([revGenesDict[g] for g in mappedGenes],
829                         prob.p_value(len(mappedGenes), len(reference),
830                                      len(mappedReferenceGenes), len(genes)),
831                         len(mappedReferenceGenes))
832            if progressCallback and i in milestones:
833                progressCallback(100.0 * i / len(terms))
834        if useFDR:
835            res = sorted(res.items(), key = lambda (_1, (_2, p, _3)): p)
836            res = dict([(id, (genes, p, ref)) \
837                        for (id, (genes, _, ref)), p in zip(res, obiProb.FDR([p for _, (_, p, _) in res]))])
838        return res
839
840    def GetAnnotatedTerms(self, genes, directAnnotationOnly=False, evidenceCodes=None, progressCallback=None):
841        """ Return all terms that are annotated by genes with evidenceCodes.
842        """
843        genes = [genes] if type(genes) == str else genes
844        revGenesDict = self.GetGeneNamesTranslator(genes)
845        genes = set(revGenesDict.keys())
846        evidenceCodes = set(evidenceCodes or evidenceDict.keys())
847        annotations = [ann for gene in genes for ann in self.geneAnnotations[gene] \
848                       if ann.Evidence_Code in evidenceCodes]
849        dd = defaultdict(set)
850        for ann in annotations:
851            dd[ann.GO_ID].add(revGenesDict.get(ann.geneName, ann.geneName))
852        if not directAnnotationOnly:
853            terms = self.ontology.ExtractSuperGraph(dd.keys())
854            for i, term in enumerate(terms):
855                termAnnots = self.GetAllAnnotations(term).intersection(annotations)
856##                termAnnots.intersection_update(annotations)
857                dd[term].update([revGenesDict.get(ann.geneName, ann.geneName) for ann in termAnnots])
858        return dict(dd)
859
860    def DrawEnrichmentGraph(self, terms, clusterSize, refSize=None, file="graph.png", width=None, height=None, precison=3):
861        refSize = len(self.geneNames) if refSize == None else refSize
862        sortedterms = sorted(terms.items(), key=lambda term:term[1][1])
863        fdr = dict(zip([t[0] for t in sortedterms], obiProb.FDR([t[1][1] for t in sortedterms])))
864        termsList = [(term, (float(len(terms[term][0]))/clusterSize) / (float(terms[term][2])/refSize),
865                      len(terms[term][0]), terms[term][2], terms[term][1],
866                      fdr[term], terms[term][0]) for term in terms]
867                         
868        drawEnrichmentGraph(termsList, file, width, height, ontology=self.ontology, precison=precison)
869
870    def __add__(self, iterable):
871        """ Return a new Annotations object with combined annotations
872        """
873        return Annotations([a for a in self] + [a for a in iterable], ontology=self.ontology)
874
875    def __iadd__(self, iterable):
876        """ Add annotations to this instance
877        """
878        self.extend(iterable)
879        return self
880
881    def __contains__(self, item):
882        return item in self.annotations
883           
884    def __iter__(self):
885        """ Iterate over all AnnotationRecord objects in annotations
886        """
887        return iter(self.annotations)
888
889    def __len__(self):
890        """ Return the number of annotations
891        """
892        return len(self.annotations)
893
894    def __getitem__(self, index):
895        """ Return the i-th annotation record
896        """
897        return self.annotations[index]
898
899    def __getslice__(self, *args):
900        return self.annotations.__getslice__(*args)
901
902    def add(self, line):
903        """ Add one annotation
904        """
905        self.AddAnnotation(line)
906
907    def append(self, line):
908        """ Add one annotation
909        """
910        self.AddAnnotation(line)
911
912    def extend(self, lines):
913        """ Add multiple annotations
914        """
915        for line in lines:
916            self.AddAnnotation(line)
917
918    def RemapGenes(self, map):
919        """
920        """
921        from copy import copy
922        for gene in map:
923            annotations = self.geneAnnotations[gene]
924            for ann in annotations:
925                for name in map[gene]:
926                    ann1 = copy(ann)
927                    ann1.geneName = name
928                    self.add(ann1)
929        self.genematcher = obiGene.GMDirect()
930        try:
931            del self._geneNames
932        except Exception:
933            pass
934        self.genematcher.set_targets(self.geneNames)
935   
936    @staticmethod
937    def DownloadAnnotations(org, file, progressCallback=None):
938        tFile = tarfile.open(file, "w:gz") if type(file) == str else file
939        tmpDir = os.path.join(orngEnviron.bufferDir, "tmp_go/")
940        try:
941            os.mkdir(tmpDir)
942        except Exception:
943            pass
944        fileName = "gene_association." + org + ".gz"
945        urlretrieve("http://www.geneontology.org/gene-associations/" + fileName,
946                    os.path.join(tmpDir, fileName),
947                    progressCallback and __progressCallbackWraper(progressCallback))
948        gzFile = GzipFile(os.path.join(tmpDir, fileName), "r")
949        file = open(os.path.join(tmpDir, "gene_association." + org), "w")
950        file.writelines(gzFile.readlines())
951        file.flush()
952        file.close()
953##        tFile = tarfile.open(os.path.join(tmpDir, "gene_association." + org + ".tar.gz"), "w:gz")
954        tFile.add(os.path.join(tmpDir, "gene_association." + org), "gene_association")
955        annotation = Annotations(os.path.join(tmpDir, "gene_association." + org),
956                    genematcher=obiGene.GMDirect(), progressCallback=progressCallback)
957        cPickle.dump(annotation.geneNames, open(os.path.join(tmpDir, "gene_names.pickle"), "wb"))
958        tFile.add(os.path.join(tmpDir, "gene_names.pickle"), "gene_names.pickle")
959        tFile.close()
960        os.remove(os.path.join(tmpDir, "gene_association." + org))
961        os.remove(os.path.join(tmpDir, "gene_names.pickle"))
962       
963    @staticmethod
964    def DownloadAnnotationsAtRev(org, rev, filename=None, progressCallback=None):
965        import urllib2
966        import shutil
967        if filename is None:
968            filename = os.path.join(default_database_path,
969                                    "gene_association.%s@rev%s.tar.gz" % (code, rev))
970        url = "http://cvsweb.geneontology.org/cgi-bin/cvsweb.cgi/~checkout~/go/gene-associations/gene_association.%s.gz?rev=%s" % (org, rev)
971        url += ";content-type=application%2Fx-gzip"
972        r = urllib2.urlopen(url)
973        shutil.copyfileobj(r, open(filename, "wb"))
974
975from obiTaxonomy import pickled_cache
976
977@pickled_cache(None, [("GO", "taxonomy.pickle"), ("Taxonomy", "ncbi_taxonomy.tar.gz")])
978def organism_name_search(name):
979    return Annotations.organism_name_search(name)
980
981def filterByPValue(terms, maxPValue=0.1):
982    """ Filters the terms by the p-value. Asumes terms is a dict with
983    the same structure as returned from GetEnrichedTerms.
984   
985    """
986    return dict(filter(lambda (k,e): e[1]<=maxPValue, terms.items()))
987
988def filterByFrequency(terms, minF=2):
989    """ Filters the terms by the cluster frequency. Asumes terms is
990    a dict with the same structure as returned from GetEnrichedTerms.
991   
992    """
993    return dict(filter(lambda (k,e): len(e[0])>=minF, terms.items()))
994
995def filterByRefFrequency(terms, minF=4):
996    """ Filters the terms by the reference frequency. Asumes terms is
997    a dict with the same structure as returned from GetEnrichedTerms.
998   
999    """
1000    return dict(filter(lambda (k,e): e[2]>=minF, terms.items()))
1001
1002##def drawEnrichmentGraph(termsList, clusterSize, refSize, filename="graph.png", width=None, height=None):
1003##    if type(termsList) == dict:
1004##        termsList = [(term, (float(len(termsList[term][0]))/clusterSize) / (float(termsList[term][2])/refSize),
1005##                      len(termsList[term][0]), termsList[term][2], termsList[term][1], 1.0, termsList[term][0]) for term in termsList]
1006##                     
1007##                     
1008##                             
1009##    drawEnrichmentGraph_tostreamMk2(termsList, open(filename, "wb"), width, height)
1010
1011def drawEnrichmentGraph_tostream(GOTerms, clusterSize, refSize, fh, width=None, height=None):
1012    def getParents(term):
1013        parents = extractGODAG([term])
1014        parents = filter(lambda t: t.id in GOTerms and t.id!=term, parents)
1015        c = []
1016        map(c.extend, [getParents(t.id) for t in parents])
1017        parents = filter(lambda t: t not in c, parents)
1018        return parents
1019    parents = dict([(term, getParents(term)) for term in GOTerms])
1020    #print "Parentes", parents
1021    def getChildren(term):
1022        return filter(lambda t: term in [p.id for p in parents[t]], GOTerms.keys())
1023    topLevelTerms = filter(lambda t: not parents[t], parents.keys())
1024    #print "Top level terms", topLevelTerms
1025    termsList=[]
1026    def collect(term, parent):
1027        termsList.append(
1028            ((float(len(GOTerms[term][0]))/clusterSize) / (float(GOTerms[term][2])/refSize),
1029            len(GOTerms[term][0]),
1030            GOTerms[term][2],
1031            "%.4f" % GOTerms[term][1],
1032            loadedGO.termDict[term].name,
1033            loadedGO.termDict[term].id,
1034            ", ".join(GOTerms[term][0]),
1035            parent)
1036            )
1037##        print float(len(GOTerms[term][0])), float(GOTerms[term][2]), clusterSize, refSize
1038        parent = len(termsList)-1
1039        for c in getChildren(term):
1040            collect(c, parent)
1041                         
1042    for topTerm in topLevelTerms:
1043        collect(topTerm, None)
1044
1045    drawEnrichmentGraphPIL_tostream(termsList, fh, width, height)
1046
1047def drawEnrichmentGraph(enriched, file="graph.png", width=None, height=None, header=None, ontology = None, precison=3):
1048    file = open(file, "wb") if type(file) == str else file
1049    drawEnrichmentGraph_tostreamMk2(enriched, file,  width, height, header, ontology, precison)
1050   
1051def drawEnrichmentGraph_tostreamMk2(enriched, fh, width, height, header=None, ontology = None, precison=4):
1052    ontology = ontology if ontology else Ontology()
1053    header = header if header else ["List", "Total", "p-value", "FDR", "Names", "Genes"]
1054    GOTerms = dict([(t[0], t) for t in enriched if t[0] in ontology])
1055    def getParents(term):
1056        parents = ontology.ExtractSuperGraph([term])
1057        parents = [id for id in parents if id in GOTerms and id != term]
1058        c = reduce(set.union, [set(ontology.ExtractSuperGraph([id])) - set([id]) for id in parents], set())
1059        parents = [t for t in parents if t not in c]
1060        return parents
1061    parents = dict([(term, getParents(term)) for term in GOTerms])
1062    #print "Parentes", parents
1063    def getChildren(term):
1064        return [id for id in GOTerms if term in parents[id]]
1065    topLevelTerms = [id for id in parents if not parents[id]]
1066    #print "Top level terms", topLevelTerms
1067    termsList=[]
1068    fmt = "%" + ".%if" % precison
1069    def collect(term, parent):
1070##        termsList.append(
1071##            ((float(len(GOTerms[term][0]))/clusterSize) / (float(GOTerms[term][2])/refSize),
1072##            len(GOTerms[term][0]),
1073##            GOTerms[term][2],
1074##            "%.4f" % GOTerms[term][1],
1075##            loadedGO.termDict[term].name,
1076##            loadedGO.termDict[term].id,
1077##            ", ".join(GOTerms[term][0]),
1078##            parent)
1079##            )
1080        termsList.append(GOTerms[term][1:4] + \
1081                         (fmt % GOTerms[term][4],
1082                          fmt % GOTerms[term][5],
1083                          ontology[term].name,
1084                          ", ".join(GOTerms[term][6])) + (parent,))
1085##        print float(len(GOTerms[term][0])), float(GOTerms[term][2]), clusterSize, refSize
1086        parent = len(termsList)-1
1087        for c in getChildren(term):
1088            collect(c, parent)
1089                         
1090    for topTerm in topLevelTerms:
1091        collect(topTerm, None)
1092    for entry in enriched:
1093        if entry[0] not in ontology:
1094            termsList.append(entry[1:4] + \
1095                             (fmt % entry[4],
1096                              fmt % entry[5],
1097                              entry[0],
1098                              ", ".join(entry[6])) + (None,))
1099
1100    drawEnrichmentGraphPIL_tostream(termsList, header, fh, width, height)
1101##    drawEnrichmentGraphPylab_tostream(termsList, header, fh, width, height)
1102   
1103def drawEnrichmentGraphPIL_tostream(termsList, headers, fh, width=None, height=None):
1104    from PIL import Image, ImageDraw, ImageFont
1105    backgroundColor = (255, 255, 255)
1106    textColor = (0, 0, 0)
1107    graphColor = (0, 0, 255)
1108    fontSize = height==None and 12 or (height-60)/len(termsList)
1109    font = ImageFont.load_default()
1110    try:
1111        font = ImageFont.truetype("arial.ttf", fontSize)
1112    except:
1113        pass
1114    getMaxTextHeightHint = lambda l: max([font.getsize(t)[1] for t in l])
1115    getMaxTextWidthHint = lambda l: max([font.getsize(t)[0] for t in l])
1116    maxFoldWidth = width!=None and min(150, width/6) or 150
1117    maxFoldEnrichment = max([t[0] for t in termsList])
1118    foldNormalizationFactor = float(maxFoldWidth)/maxFoldEnrichment
1119    foldWidths = [int(foldNormalizationFactor*term[0]) for term in termsList]
1120    treeStep = 10
1121    treeWidth = {}
1122    for i, term in enumerate(termsList):
1123        treeWidth[i] = (term[-1]==None and 1 or treeWidth[term[-1]]+1)
1124    treeStep = width!=None and min(treeStep, width/(6*max(treeWidth.values())) or 2) or treeStep
1125    treeWidth = [w*treeStep + foldWidths[i] for i, w in treeWidth.items()]
1126    treeWidth = max(treeWidth) - maxFoldWidth
1127    verticalMargin = 10
1128    horizontalMargin = 10
1129##    print verticalMargin, maxFoldWidth, treeWidth
1130##    treeWidth = 100
1131    firstColumnStart = verticalMargin + maxFoldWidth + treeWidth + 10
1132    secondColumnStart = firstColumnStart + getMaxTextWidthHint([str(t[1]) for t in termsList]+[headers[0]]) + 2
1133    thirdColumnStart = secondColumnStart + getMaxTextWidthHint([str(t[2]) for t in termsList]+[headers[1]]) + 2
1134    fourthColumnStart = thirdColumnStart + getMaxTextWidthHint([str(t[3]) for t in termsList]+[headers[2]]) + 2
1135    fifthColumnStart = fourthColumnStart + getMaxTextWidthHint([str(t[4]) for t in termsList]+[headers[3]]) + 4
1136##    maxAnnotationTextWidth = width==None and getMaxTextWidthHint([str(t[4]) for t in termsList]+["Annotation"]) or (width - fourthColumnStart - verticalMargin) * 2 / 3
1137    maxAnnotationTextWidth = width==None and getMaxTextWidthHint([str(t[5]) for t in termsList]+[headers[4]]) or max((width - fifthColumnStart - verticalMargin) * 2 / 3, getMaxTextWidthHint([t[5] for t in termsList]+[headers[4]]))
1138    sixthColumnStart  = fifthColumnStart + maxAnnotationTextWidth + 4
1139    maxGenesTextWidth = width==None and getMaxTextWidthHint([str(t[6]) for t in termsList]+[headers[5]]) or (width - fifthColumnStart - verticalMargin) / 3
1140   
1141    legendHeight = font.getsize("1234567890")[1]*2
1142    termHeight = font.getsize("A")[1]
1143##    print fourthColumnStart, maxAnnotationTextWidth, verticalMargin
1144    width = sixthColumnStart + maxGenesTextWidth + verticalMargin
1145    height = len(termsList)*termHeight+2*(legendHeight+horizontalMargin)
1146
1147    image = Image.new("RGB", (width, height), backgroundColor)
1148    draw = ImageDraw.Draw(image)
1149
1150    def truncText(text, maxWidth, append=""):
1151        #print getMaxTextWidthHint([text]), maxAnnotationTextWidth
1152        if getMaxTextWidthHint([text])>maxWidth:
1153            while getMaxTextWidthHint([text+"..."+append])>maxWidth and text:
1154                text = text[:-1]
1155            if text:
1156                text = text+"..."+append
1157            else:
1158                text = append
1159        return text
1160    currentY = horizontalMargin + legendHeight
1161    connectAtX = {}
1162    for i, term in enumerate(termsList):
1163        draw.line([(verticalMargin, currentY+termHeight/2), (verticalMargin + foldWidths[i], currentY+termHeight/2)], width=termHeight-2, fill=graphColor)
1164        draw.text((firstColumnStart, currentY), str(term[1]), font=font, fill=textColor)
1165        draw.text((secondColumnStart, currentY), str(term[2]), font=font, fill=textColor)
1166        draw.text((thirdColumnStart, currentY), str(term[3]), font=font, fill=textColor)
1167        draw.text((fourthColumnStart, currentY), str(term[4]), font=font, fill=textColor)
1168##        annotText = width!=None and truncText(str(term[5]), maxAnnotationTextWidth, str(term[5])) or str(term[4])
1169        annotText = width!=None and truncText(str(term[5]), maxAnnotationTextWidth)
1170        draw.text((fifthColumnStart, currentY), annotText, font=font, fill=textColor)
1171        genesText = width!=None and truncText(str(term[6]), maxGenesTextWidth) or str(term[6])
1172        draw.text((sixthColumnStart, currentY), genesText, font=font, fill=textColor)
1173        lineEnd = term[-1]==None and firstColumnStart-10 or connectAtX[term[-1]]
1174        draw.line([(verticalMargin+foldWidths[i]+1, currentY+termHeight/2), (lineEnd, currentY+termHeight/2)], width=1, fill=textColor)
1175        if term[-1]!=None:
1176            draw.line([(lineEnd, currentY+termHeight/2), (lineEnd, currentY+termHeight/2 - termHeight*(i-term[-1]))], width=1, fill=textColor)
1177        connectAtX[i] = lineEnd - treeStep
1178        currentY+=termHeight
1179
1180    currentY = horizontalMargin
1181    draw.text((firstColumnStart, currentY), headers[0], font=font, fill=textColor)
1182    draw.text((secondColumnStart, currentY), headers[1], font=font, fill=textColor)
1183    draw.text((thirdColumnStart, currentY), headers[2], font=font, fill=textColor)
1184    draw.text((fourthColumnStart, currentY), headers[3], font=font, fill=textColor)
1185    draw.text((fifthColumnStart, currentY), headers[4], font=font, fill=textColor)
1186    draw.text((sixthColumnStart, currentY), headers[5], font=font, fill=textColor)
1187
1188    horizontalMargin = 0
1189    #draw.line([(verticalMargin, height - horizontalMargin - legendHeight), (verticalMargin + maxFoldWidth, height - horizontalMargin - legendHeight)], width=1, fill=textColor)
1190    draw.line([(verticalMargin, horizontalMargin + legendHeight), (verticalMargin + maxFoldWidth, horizontalMargin + legendHeight)], width=1, fill=textColor)
1191    maxLabelWidth = getMaxTextWidthHint([" "+str(i) for i in range(int(maxFoldEnrichment+1))])
1192    numOfLegendLabels = max(int(maxFoldWidth/maxLabelWidth), 2)
1193    for i in range(numOfLegendLabels+1):
1194        #draw.line([(verticalMargin + i*maxFoldWidth/10, height - horizontalMargin - legendHeight/2), (verticalMargin + i*maxFoldWidth/10, height - horizontalMargin - legendHeight)], width=1, fill=textColor)
1195        #draw.text((verticalMargin + i*maxFoldWidth/10 - font.getsize(str(i))[0]/2, height - horizontalMargin - legendHeight/2), str(i), font=font, fill=textColor)
1196
1197        label = str(int(i*maxFoldEnrichment/numOfLegendLabels))
1198        draw.line([(verticalMargin + i*maxFoldWidth/numOfLegendLabels, horizontalMargin + legendHeight/2), (verticalMargin + i*maxFoldWidth/numOfLegendLabels, horizontalMargin + legendHeight)], width=1, fill=textColor)
1199        draw.text((verticalMargin + i*maxFoldWidth/numOfLegendLabels - font.getsize(label)[0]/2, horizontalMargin), label, font=font, fill=textColor)
1200       
1201    image.save(fh)
1202
1203def drawEnrichmentGraphPylab_tostream(termsList, headers, fh, width=None, height=None, show=True):
1204    from matplotlib import pyplot as plt
1205    from matplotlib.patches import Rectangle
1206   
1207    maxFoldWidth = width!=None and min(150, width/6) or 150
1208    maxFoldEnrichment = max([t[0] for t in termsList])
1209    foldNormalizationFactor = float(maxFoldWidth)/maxFoldEnrichment
1210##    foldWidths = [int(foldNormalizationFactor*term[0]) for term in termsList]
1211    foldWidths = [term[0] for term in termsList]
1212    treeStep = maxFoldEnrichment*0.05
1213    treeWidth = {}
1214
1215    for i, term in enumerate(termsList):
1216        treeWidth[i] = (term[-1]==None and treeStep or treeWidth[term[-1]] + treeStep)
1217    maxTreeWidth = max(treeWidth)
1218
1219    connectAt = {}
1220    cellText = []
1221    axes1 = plt.axes([0.1, 0.1, 0.2, 0.8])
1222    for i, line in enumerate(termsList):
1223        enrichment, n, m, p_val, fdr_val, name, genes, parent = line
1224        r = Rectangle((0, len(termsList) - i - 0.4), enrichment, 0.8)
1225        plt.gca().add_patch(r)
1226        plt.plot([enrichment, connectAt.get(parent, maxFoldEnrichment + maxTreeWidth)], [len(termsList) - i, len(termsList) - i], color="black")
1227        connectAt[i] = connectAt.get(parent, maxFoldEnrichment + maxTreeWidth) - treeStep
1228        if parent != None:
1229            plt.plot([connectAt.get(parent)]*2, [len(termsList) - i, len(termsList) - parent], color="black")
1230        cellText.append((str(n), str(m), p_val, fdr_val, name, genes))
1231
1232##    from orngClustering import TableTextLayout
1233##    text = TableTextLayout((maxFoldEnrichment*1.1, len(termsList)), cellText)
1234    from orngClustering import TablePlot
1235    if True:
1236        axes2 = plt.axes([0.3, 0.1, 0.6, 0.8], sharey=axes1)
1237        axes2.set_axis_off()
1238        table = TablePlot((0, len(termsList)), axes=plt.gca())
1239        for i, line in enumerate(cellText):
1240            for j, text in enumerate(line):
1241                table.add_cell(i, j,width=len(text), height=1, text=text, loc="left", edgecolor="w", facecolor="w")
1242
1243        table.set_figure(plt.gcf())
1244        plt.gca().add_artist(table)
1245        plt.gca()._set_artist_props(table)
1246##    plt.text(3, 3, "\n".join(["\t".join(text) for text in cellText]))
1247
1248##    table = plt.table(cellText=cellText, colLabels=headers, loc="right")
1249##    table.set_transform(plt.gca().transData)
1250##   
1251##    table.set_xy(20,20)
1252    plt.show()
1253   
1254class Taxonomy(object):
1255    """Maps NCBI taxonomy ids to coresponding GO organism codes
1256    """
1257    common_org_map = {"297284":"9913", "30523":"9913", # Bos taurus
1258                      "5782":"352472", "44689":"352472", "366501":"352472", # Dictyostelium discoideum
1259                      "83333": "562", # Escherichia coli
1260                      "52545":"4530", "4532":"4530", "65489":"4530", "4533":"4530", "77588":"4530", "29689":"4530",
1261                      "4538":"4530", "40148":"4530", "29690":"4530", "110450":"4530", "4534":"4530", "83309":"4530",
1262                      "4528":"4530", "127571":"4530", "40149":"4530", "83307":"4530", "63629":"4530", "4536": "4530",
1263                      "4535":"4530", "4537":"4530", "65491":"4530", "83308":"4530", "4529":"4530", "4530":"4530",
1264                      "39946":"4530", "39947":"4530", "110451":"4530", "364100":"4530", "364099":"4530", "4539":"4530",
1265                      }
1266    code_map = {"3702":"tair",  # Arabidopsis thaliana
1267                "9913":"goa_cow", # Bos taurus
1268                "6239":"wb", # Caenorhabditis elegans
1269                "3055":None, # Chlamydomonas reinhardtii
1270                "7955":"zfin", # Danio rerio (zebrafish)
1271                "352472":"dictyBase", # Dictyostelium discoideum
1272                "7227": "fb", # Drosophila melanogaster
1273                "562": "ecocyc", # Escherichia coli
1274                "11103": None, # Hepatitis C virus
1275                "9606": "goa_human",  # Homo sapiens
1276                "10090": "mgi", # Mus musculus
1277                "2104": None,  # Mycoplasma pneumoniae
1278                "4530": "gramene_oryza",  # Oryza sativa
1279                "5833": "GeneDB_Pfalciparum",  # Plasmodium falciparum
1280                "4754": None,  # Pneumocystis carinii
1281                "10116": "rgd", # Rattus norvegicus
1282                "4932": "sgd",  # Saccharomyces cerevisiae
1283                "4896": "GeneDB_Spombe", # Schizosaccharomyces pombe
1284                "31033": None, # Takifugu rubripes
1285                "8355": None,  # Xenopus laevis
1286                "4577": None # Zea mays
1287                }
1288    version = 1
1289    __shared_state = {"tax": None}
1290    def __init__(self):
1291        self.__dict__ = self.__shared_state
1292        if not self.tax:
1293            import orngServerFiles
1294            path = orngServerFiles.localpath_download("GO", "taxonomy.pickle")
1295            if os.path.isfile(path):
1296                self.tax = cPickle.load(open(path, "rb"))
1297            else:
1298                orngServerFiles.download("GO", "taxonomy.pickle")
1299                self.tax = cPickle.load(open(path, "rb"))
1300               
1301    def __getitem__(self, key):
1302        key = self.common_org_map.get(key, key)
1303        return self.code_map[key]
1304   
1305    def keys(self):
1306        return list(set(self.common_org_map.keys() + self.code_map.keys()))
1307   
1308#    @classmethod
1309#    def get_taxonomy(cls):
1310#        import urllib2 as url
1311#        import sgmllib
1312#        organisms
1313#        class MyParser(sgmllib.SGMLParser):
1314#            inTable = False
1315#            def start_table(self, attributes):
1316#                self.inTable = dict(attributes).get("summary", False)
1317#            def end_table(self):
1318#                self.inTable = False
1319#            def start
1320   
1321def from_taxid(id):
1322    """ Return a set of GO organism codes that correspond to NCBI taxonomy id
1323    """
1324    return Taxonomy()[id]
1325
1326def to_taxid(db_code):
1327    """ Return a set of NCBI taxonomy ids from db_code GO organism annotations
1328    """
1329    r = [key for key, val in Taxonomy().code_map.items() if db_code == val]
1330    return set(r)
1331   
1332
1333class __progressCallbackWrapper:
1334    def __init__(self, callback):
1335        self.callback = callback
1336    def __call__(self, bCount, bSize, fSize):
1337        fSize = 10000000 if fSize == -1 else fSize
1338        self.callback(100*bCount*bSize/fSize)
1339       
1340from obiGenomicsUpdate import Update as UpdateBase
1341
1342import urllib2
1343
1344class Update(UpdateBase):
1345    def __init__(self, local_database_path=None, progressCallback=None):
1346        UpdateBase.__init__(self, local_database_path or getDataDir(), progressCallback)
1347    def CheckModified(self, addr, date=None):
1348        return date < self.GetLastModified(addr) if date else True
1349       
1350    def CheckModifiedOrg(self, org):
1351        return self.CheckModified("http://www.geneontology.org/gene-associations/gene_association." + org + ".gz", self.LastModifiedOrg(org))
1352   
1353    def LastModifiedOrg(self, org):
1354        return self.shelve.get((Update.UpdateAnnotation, (org,)), None)
1355
1356    def GetLastModified(self, addr):
1357        stream = urllib2.urlopen(addr)
1358        return datetime.strptime(stream.headers.get("Last-Modified"), "%a, %d %b %Y %H:%M:%S %Z")
1359##        return stream.headers.get("Last-Modified")
1360
1361    def GetAvailableOrganisms(self):
1362        source = urllib2.urlopen("http://www.geneontology.org/gene-associations/").read()
1363        return [s.split(".")[1] for s in sorted(set(re.findall("gene_association\.[a-zA-z0-9_]+?\.gz", source)))]
1364
1365    def GetDownloadedOrganisms(self):
1366        return [name.split(".")[1] for name in os.listdir(self.local_database_path) if name.startswith("gene_association")]
1367
1368    def IsUpdatable(self, func, args):
1369        if func == Update.UpdateOntology:
1370            return self.CheckModified("http://www.geneontology.org/ontology/gene_ontology.obo", self.shelve.get((Update.UpdateOntology, ()), None))
1371        elif func == Update.UpdateAnnotation:
1372            return self.CheckModifiedOrg(args[0])
1373           
1374    def GetDownloadable(self):
1375        orgs = set(self.GetAvailableOrganisms()) - set(self.GetDownloadedOrganisms())
1376        ret = []
1377        if (Update.UpdateOntology, ()) not in self.shelve:
1378            ret.append((Update.UpdateOntology, ()))
1379        if orgs:
1380            ret.extend([(Update.UpdateAnnotation, (org,)) for org in orgs])
1381        return ret
1382
1383    def UpdateOntology(self):
1384        Ontology.DownloadOntology(os.path.join(self.local_database_path, "gene_ontology_edit.obo.tar.gz"), self.progressCallback)
1385        self._update(Update.UpdateOntology, (), self.GetLastModified("http://www.geneontology.org/ontology/gene_ontology.obo"))
1386
1387    def UpdateAnnotation(self, org):
1388        Annotations.DownloadAnnotations(org, os.path.join(self.local_database_path, "gene_association." + org + ".tar.gz"), self.progressCallback)
1389        self._update(Update.UpdateAnnotation, (org,), self.GetLastModified("http://www.geneontology.org/gene-associations/gene_association." + org + ".gz"))
1390       
1391    def UpdateTaxonomy(self, org):
1392        exclude = ["goa_uniprot", "goa_pdb", "GeneDB_tsetse", "reactome", "goa_zebrafish", "goa_rat", "goa_mouse"]
1393
1394        orgs = self.GetAvailableOrganisms()
1395        tax = defaultdict(set)
1396
1397        for org in orgs:
1398            if org in exclude:
1399                continue
1400            try:
1401                a = obiGO.Annotations(os.path.join(self.local_database_path, "gene_association." + org + ".tar.gz"))
1402                taxons = set(ann.taxon for ann in a.annotations)
1403                for taxId in [t.split(":")[-1] for t in taxons if "|" not in t]: ## exclude taxons with cardinality 2
1404                    tax[taxId].add(org)
1405            except Exception, ex:
1406                print ex
1407               
1408        cPickle.dump(dict(tax), open(os.path.join(path, "taxonomy.pickle"), "wb"))
1409           
1410
1411def _test1():
1412##    Ontology.DownloadOntology("ontology_arch.tar.gz")
1413##    Annotations.DownloadAnnotations("sgd", "annotations_arch.tar.gz")
1414    def _print(f):
1415        print f
1416    o = Ontology("ontology_arch.tar.gz")
1417    a = Annotations("annotations_arch.tar.gz", ontology=o)
1418   
1419    a.GetEnrichedTerms(sorted(a.geneNames)[:100])#, progressCallback=_print)
1420##    profile.runctx("a.GetEnrichedTerms(sorted(a.geneNames)[:100])", {"a":a}, {})
1421    a.GetEnrichedTerms(sorted(a.geneNames)[:100])#, progressCallback=_print)
1422    d1 = a.GetEnrichedTerms(sorted(a.geneNames)[:1000])#, progressCallback=_print)
1423   
1424##    print a.GetEnrichedTerms(sorted(a.geneNames)[:100])#, progressCallback=_print)
1425
1426def _test2():
1427    o = Ontology()
1428    a = Annotations("human", ontology=o)
1429    clusterGenes = sorted(a.geneNames)[:100]
1430    for i in range(10):
1431        terms = a.GetEnrichedTerms(sorted(a.geneNames)[:200], aspect=["P"])
1432        a.GetEnrichedTerms(sorted(a.geneNames)[:200], aspect=["C"])
1433        a.GetEnrichedTerms(sorted(a.geneNames)[:200], aspect=["F"])
1434        print i
1435#    a.DrawEnrichmentGraph(filterByPValue(terms), len(clusterGenes), len(a.geneNames))
1436             
1437##    drawEnrichmentGraph([("bal", 1.0, 5, 6, 0.1, 0.4, ["vv"]),
1438##                        ("GO:0019079", 0.5, 5, 6, 0.1, 0.4, ["cc", "bb"]),
1439##                        ("GO:0022415", 0.4, 5, 7, 0.11, 0.4, ["cc1", "bb"])], open("graph.png", "wb"), None, None)
1440
1441def _test3():
1442    o = Ontology()
1443    a = Annotations("sgd", ontology=o)
1444##    a = Annotations(list(a)[3:len(a)/3], ontology=o)
1445    clusterGenes = sorted(a.geneNames)[:1] + sorted(a.geneNames)[-1:]
1446##    clusterGenes = [g + "_" + str(i%5) for g in sorted(a.geneNames)[:2]]
1447    exonMap = dict([(gene, [gene+"_E%i" %i for i in range(10)]) for gene in a.geneNames])
1448    a.RemapGenes(exonMap)
1449##    o.reverseAliasMapper = o.aliasMapper = {}
1450    terms = a.GetEnrichedTerms(exonMap.values()[0][:2] + exonMap.values()[-1][2:])
1451##    terms = a.GetEnrichedTerms(clusterGenes)
1452    print terms
1453##    a.DrawEnrichmentGraph(filterByPValue(terms), len(clusterGenes), len(a.geneNames))
1454    a.DrawEnrichmentGraph(filterByPValue(terms, maxPValue=0.1), len(clusterGenes), len(a.geneNames))
1455   
1456if __name__ == "__main__":
1457    _test2()
Note: See TracBrowser for help on using the repository browser.