source: orange-bioinformatics/obiGO.py @ 1351:c9470ca4ca88

Revision 1351:c9470ca4ca88, 58.6 KB checked in by ales_erjavec <ales.erjavec@…>, 3 years ago (diff)

Using intern function to reduce the memory consumption of loaded annotations terms.

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