_ahocorasick.pyx 45.7 KB
Newer Older
Eric Coissac committed
1 2 3
# cython: language_level=3

from ._ahocorasick cimport *
4 5
from orgasm.apps.progress cimport ProgressBar 

6
from orgasm.utils.dna import reverseComplement
Eric Coissac committed
7
from ._bitvector cimport BitVector
8 9

import time
10
import math
11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78

#
# This dict is a mix off all mitochondrial, chroroplast and the universal genetic code
# It serves for back translating purpose
#
 
cdef dict UniversalGeneticCode     = {b"A" : (b"GCT", b"GCC", b"GCA", b"GCG"),
                                      b"C" : (b"TGT", b"TGC"),
                                      b"D" : (b"GAT", b"GAC"),
                                      b"E" : (b"GAA", b"GAG"),
                                      b"F" : (b"TTT", b"TTC"),
                                      b"G" : (b"GGT", b"GGC", b"GGA", b"GGG"),
                                      b"H" : (b"CAT", b"CAC"),
                                      b"I" : (b"ATT", b"ATC", b"ATA"),
                                      b"K" : (b"AAA", b"AAG"),
                                      b"L" : (b"TTA", b"TTG", b"TAG", b"CTT", b"CTC", b"CTA", b"CTG"),
                                      b"M" : (b"ATG", b"ATA"),
                                      b"N" : (b"AAT", b"AAC", b"AAA"),
                                      b"P" : (b"CCT", b"CCC", b"CCA", b"CCG"),
                                      b"Q" : (b"CAA", b"CAG"),
                                      b"R" : (b"CGT", b"CGC", b"CGA", b"CGG", b"AGA", b"AGG"),
                                      b"S" : (b"TCT", b"TCC", b"TCA", b"TCG", b"AGT", b"AGC", b"AGA", b"AGG"),
                                      b"T" : (b"ACT", b"ACC", b"ACA", b"ACG", b"CTT", b"CTC", b"CTA", b"CTG"),
                                      b"V" : (b"GTT", b"GTC", b"GTA", b"GTG"),
                                      b"W" : (b"TGA", b"TGG"),
                                      b"Y" : (b"TAT", b"TAC", b"TAA"),
                                     }

cdef dict CompUniversalGeneticCode = {b"A" : (b"AGC", b"GGC", b"TGC", b"CGC"),
                                      b"C" : (b"ACA", b"GCA"),
                                      b"D" : (b"ATC", b"GTC"),
                                      b"E" : (b"TTC", b"CTC"),
                                      b"F" : (b"AAA", b"GAA"),
                                      b"G" : (b"ACC", b"GCC", b"TCC", b"CCC"),
                                      b"H" : (b"ATG", b"GTG"),
                                      b"I" : (b"AAT", b"GAT", b"TAT"),
                                      b"K" : (b"TTT", b"CTT"),
                                      b"L" : (b"TAA", b"CAA", b"CTA", b"AAG", b"GAG", b"TAG", b"CAG"),
                                      b"M" : (b"CAT", b"TAT"),
                                      b"N" : (b"ATT", b"GTT", b"TTT"),
                                      b"P" : (b"AGG", b"GGG", b"TGG", b"CGG"),
                                      b"Q" : (b"TTG", b"CTG"),
                                      b"R" : (b"ACG", b"GCG", b"TCG", b"CCG", b"TCT", b"CCT"),
                                      b"S" : (b"AGA", b"GGA", b"TGA", b"CGA", b"ACT", b"GCT", b"TCT", b"CCT"),
                                      b"T" : (b"AGT", b"GGT", b"TGT", b"CGT", b"AAG", b"GAG", b"TAG", b"CAG"),
                                      b"V" : (b"AAC", b"GAC", b"TAC", b"CAC"),
                                      b"W" : (b"TCA", b"CCA"),
                                      b"Y" : (b"ATA", b"GTA", b"TTA"),
                                     }


cdef dict listCodons2Dict(tuple codons):
    cdef dict d={}
    cdef dict s
    cdef dict new
    cdef char* cc
    cdef dict      value
    cdef int  l
    cdef char letter[2]
    cdef bytes c

    letter[1]=0
    
    for c in codons:
        s = d
        cc= c
        for l in range(3):
            letter[0]=cc[l]
Eric Coissac committed
79 80 81 82
            value = s.get(letter,None)
            if value is None:
                value = {}
                s[letter]=value
83
            s=value
Eric Coissac committed
84
                
85 86
    return d

Eric Coissac committed
87
cpdef dict bindCodons(dict codon1, dict codon2):
88 89
    cdef PyObject* pkey
    cdef bytes     key
90
    cdef PyObject* pvalue  # @DuplicatedSignature
91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109
    cdef dict      l2
    cdef dict      l3
    cdef Py_ssize_t ppos1=0
    cdef Py_ssize_t ppos2=0
    cdef Py_ssize_t ppos3=0
    
    while PyDict_Next(codon1, &ppos1, &pkey, &pvalue):
        l2=<object>pvalue
        ppos2=0
        while PyDict_Next(l2, &ppos2, &pkey, &pvalue):
            l3=<object>pvalue
            ppos3=0
            while PyDict_Next(l3, &ppos3, &pkey, &pvalue):
                Py_XDECREF(pvalue)
                key=<object>pkey
                PyDict_SetItem(l3,key,codon2)

    return codon1

Eric Coissac committed
110
cpdef dict codon(char aa, bint direct=True):
111 112 113 114 115 116
    cdef tuple value
    cdef char  letter[2]
    
    letter[0]=aa
    letter[1]=0
    
Eric Coissac committed
117
    
118
    if direct:
Eric Coissac committed
119
#        value = <object>PyDict_GetItemString(UniversalGeneticCode,letter)
Eric Coissac committed
120
        value = UniversalGeneticCode.get(letter,())
121
    else:
Eric Coissac committed
122
#        value = <object>PyDict_GetItemString(CompUniversalGeneticCode,letter)        
Eric Coissac committed
123
        value = CompUniversalGeneticCode.get(letter,())        
Eric Coissac committed
124
    
125 126
    return listCodons2Dict(value)

Eric Coissac committed
127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168
cpdef bint homopolymer(bytes s):
    """
    Returns True if the word s is an homopolymer, an homedimer 
    or an homotrimer
    
    @param s: the string to test
    @type s: bytes   
    
    @return True if s is an homopolymer, an homedimer 
            or an homotrimer, False otherwise
    @rtype: bool
    """
    s2 = s[2:] + s[0:2]
    s3 = s[3:] + s[0:3]
    return s==s2 or s==s3
                
def enumerateword(dict automata, bytes prefix=b""):
    cdef bytes k
    cdef bytes w 
    
    if not automata:
        yield prefix
    else:
        for k in automata:
            for w in enumerateword(automata[k],prefix+k):
                yield w
          
cpdef dict word2automata(wordlist):
    cdef dict automata={}
    cdef bytes w
    cdef dict  a 
    cdef bytes l 
    cdef dict  d
    
    for w in wordlist:
        a = automata
        for l in w:
            d = a.get(l,{})
            a[l]=d
            a=d
            
    return automata
169 170

cdef double* buildShanonTable(size_t readsize, size_t* offset):
Eric Coissac committed
171
    cdef size_t codonmax = readsize // 3 + (1 if (readsize % 3) else 0)
172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204
    cdef size_t i=codonmax+1
    cdef size_t j=0
    cdef size_t k=0
    cdef double* table
    cdef double* line
    cdef double log6=log(6)
    cdef double freq
    
    while i!=0:
        k|=(i & 1) & (i>1)
        i >>= 1
        j+=1
        
    j+=k
    
    offset[0]=j
    codonmax=1<<j
        
    table = <double*> malloc((readsize+1)*(codonmax)*sizeof(double))
    
    for i in range(readsize+1):
        line = table + (i << offset[0])
        for j in range(codonmax):
            if j <=i and j > 0:
                freq = <double> j / <double> i
                freq = -freq * log(freq)/log6
            else:
                freq=0
            line[j]=freq 
            
    return table
            
    
205 206 207 208 209 210 211 212 213 214 215 216
 
cdef bint samereads(unsigned char* p1, unsigned char* p2, size_t length):
    cdef uint64_t *w1 = <uint64_t*>p1
    cdef uint64_t *w2 = <uint64_t*>p2
    cdef size_t   i
    
    for i in range(length):
        if w1[i]!=w2[i]:
            return False
        
    return True   

217
cdef class AhoCorasick:
218
    
219
    def __init__(self,initsize=100000):
220 221 222 223
        self._finalized = False
        self._step = initsize
        self._depth= 0
        self._maxseqid=0
224
        self._maxpos=0
225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244
        
        # Allocate state arena
        
        self._states = <dnastate*> malloc(sizeof(dnastate)
                                          * initsize)
        
        if self._states==NULL:
            raise RuntimeError("No more memory for automata states")
        
        self._statesize = initsize
        
        # initialize initial state
        
        self._states.match = <protmatch*>-1
        self._states.a     = NULL
        self._states.c     = NULL
        self._states.t     = NULL
        self._states.g     = NULL
        
        self._seqid       = {}
245
        self._rseqid       = {}
246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344
        
        self._statecount = 1
        
        # allocate match arena
        
        self._matches = <protmatch*> malloc(sizeof(protmatch)
                                            * initsize)
        
        if self._matches==NULL:
            raise RuntimeError("No more memory for automata matches")
        
        self._matchsize = initsize
        self._matchcount= 0
        
        
        
    def __dealloc__(self):
        if (self._states != NULL):
            free(<void*>self._states)
        if (self._matches != NULL):
            free(<void*>self._matches)
            
            
    def __len__(self):
        return self._statecount
    
    def addWord(self,bytes word,int protid, int position):
        cdef char* cword = word 
        
        assert not self._finalized, "You cannot add words to a finilized automata"

        self.addCWord(cword,PyBytes_Size(word),protid,position)


    cdef dnastate *getNextState(self,dnastate *base, int letter):
        cdef dnastate** table=&(base.a)
        cdef dnastate*  state = table[letter]
        cdef dnastate*  arena
        cdef size_t     baseid    
    
        if state == NULL:
            if self._finalized:
                return NULL
            else:
                                
                if self._statecount == self._statesize:
                    #
                    # the state arena if full we need to enlarge it
                    #
                    
                    # we save baseid to convert it to the new arena adressing system 
                    baseid = base - self._states
                    
                    # new size of the arena and correponding allocation
                    self._statesize+=self._step
#                    print "realloc buffer to size %d" % self._statesize
                    arena = <dnastate*> realloc(self._states,
                                                sizeof(dnastate)
                                                * self._statesize)
                                         
                    if arena == NULL:
                        # We were not able to enlarge arena
                        return NULL
                    else:
                        self._states = arena
                       
                    #compute the new base and table address
                    base = arena + baseid
                    table=&(base.a)
                
                state = self._states + self._statecount
                table[letter] = <dnastate*>self._statecount
                self._statecount+=1
                
                state.match = <protmatch*>-1
                state.a     = NULL
                state.c     = NULL
                state.t     = NULL
                state.g     = NULL
                
                return state
                
        else:
            if self._finalized:
                return state 
            else:
                return self._states + <size_t> state
            
    cdef protmatch* setAsTerminal(self, dnastate* base, int protid, int position):
        """
        setAsTerminal function is an internal function that is called from the
        addWord method. It must be only used on none finalized automata, but 
        no checking for this property is done a precondition
        """
        cdef protmatch* curmatch = NULL
        cdef protmatch* nextmatch= NULL
        cdef size_t     curmatchid
        cdef protmatch* arena
        
345 346 347
        if position > self._maxpos:
            self._maxpos = position 
        
348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407
        if base.match!=<protmatch*>-1:
            nextmatch = self._matches + <size_t>(base.match)        
        
        while(nextmatch):
            curmatch = nextmatch
            if (    curmatch.protid == protid 
                and curmatch.position == position):
                return curmatch
            else:
                if curmatch.next:
                    nextmatch = self._matches + <size_t>curmatch.next
                else:
                    nextmatch = NULL
                    
        # There no match for position@protid register to this state
        # so we have to create a new one
            
        if self._matchcount == self._matchsize:
            # I save the curmatchid to restore curmatch later
            if curmatch :
                curmatchid = curmatch - self._matches
                        
            # the match arena if full we need to enlarge it
            
            self._matchsize+=self._step
#            print "reallocate matches to %d" % self._matchsize
            arena = <protmatch*> realloc(self._matches,
                                         sizeof(protmatch)
                                         * self._matchsize)
                                 
            if arena == NULL:
                # We were not able to enlarge arena
                return NULL
            else:
                self._matches = arena
                
            # I restore curmatch 
            if curmatch :
                curmatch = self._matches + curmatchid

                
        # compute the address 
        nextmatch = self._matches + self._matchcount
        
        if curmatch : 
            curmatch.next = <protmatch*>self._matchcount
        else:
            base.match = <protmatch*>self._matchcount
            
        self._matchcount+=1
        
        nextmatch.protid = protid
        nextmatch.position = position 
        nextmatch.next = NULL 
        
        if protid > self._maxseqid:
            self._maxseqid=protid
        
        return nextmatch 
    
408 409 410 411 412 413 414 415 416 417 418
    cdef dnastate* simpleMatch(self,char* cword, size_t lword):
        """
        Checks if cword a dna word of length lword is present in 
        the automata. The function returns the state at the end of
        the word in the automata or null if the word is not found

        This function is used for computing the error link during the
        finalization step
        
        warning : This is a private function that should not be used directly        
        """
419 420 421 422
        cdef char   letter
        cdef dnastate** table
        cdef dnastate*  curstate
        cdef dnastate*  nextstate
423 424 425 426
                
        # We cannot match a word of length null
        # neither longer than the longest pattern
        # in the automata
427
        
428
        if lword>self._depth:
429 430 431 432 433 434
            return NULL 
        
        nextstate = self._states 
        
        while (nextstate!=NULL and lword):
            curstate = nextstate
435 436

            letter = cword[0]
437 438
            table    = &(curstate.a)
            nextstate = table[letter]
439
            
440 441 442 443
            cword+=1
            lword-=1
            
        if lword==0 and nextstate!=NULL:
444
            return nextstate
445 446 447 448
        else:
            return NULL
        
        
449
    cdef dnastate* longestSuffix(self,char* cword, size_t lword):
450 451 452 453 454 455 456 457 458
        """
        identify the longest suffix of cword that is recognized by
        the Aho-Corasick automata.
        
        This function is used for computing the error link during the
        finalization step
        
        warning : This is a private function that should not be used directly        
        """
459 460
        cdef dnastate* suffix=NULL

461
        
462
        while lword >= 1 and suffix==NULL:
463 464
            cword+=1
            lword-=1
465 466

            suffix = self.simpleMatch(cword,lword)
467 468 469 470 471 472 473 474 475 476 477
            
        return suffix
            
        
    cdef setErrorLink(self):
        """
        Add error link to the Aho-Corasick automata during its finalization
        step.
        
        warning : This is a private function that should not be used directly
        """
478
        cdef char* buffer
479 480 481 482
        cdef size_t*     stack
        cdef dnastate*  curstate
        cdef dnastate*  nextstate
        cdef dnastate** table
483
        cdef dnastate* links
484 485 486
        cdef dnastate** errors
        cdef size_t i
        cdef size_t j
487
        cdef int pos=0
488 489 490 491
        cdef size_t lsuffix
        cdef size_t lword
         

492
#        print "depth : %d -> %d" % (self._depth,sizeof(size_t)*(self._depth+1)*4)
493 494 495 496 497

        stack = <size_t*>malloc(sizeof(size_t)*(self._depth+1)*4)

        errors = <dnastate**>malloc(sizeof(dnastate*) * self._statecount)
        bzero(<void*>errors,sizeof(dnastate*) * self._statecount)
498 499
        
        buffer=<char*>malloc(sizeof(char)*(self._depth+1))
500

501 502 503 504 505
        #
        # We explore trie with a deep first algorithm
        #
        
        # I put on the stack the root of the automata
506
        stack[0] = 0
507 508 509 510
        pos = 0

        j=0 # used to scan the four branches of a node     
       
511

512 513
        while pos >= 0:
#            print "Run on state : [%d]  %d-%d" % (pos,stack[pos],j)
514

515 516
            # I consider the last state on the stack 
            curstate = self._states + stack[pos] 
517 518 519 520 521 522 523 524 525 526 527
            table=&(curstate.a)          
            
            # We skip the empty link
            while j < 4 and table[j]==NULL:
                j+=1
                
            # on the loop exit or j==4 and we have explored all the branches
            # or j<4 and we found a non empty branch
                
            if j < 4:
                
528
                # We found a link we push it on the stack
529 530 531
                
                buffer[pos]=j
                pos+=1
532 533
                stack[pos]= (table[j] - self._states)
#                print "push %d [%d]" % (pos,j)
534
                j=0
535 536 537
                            
                # ===> Then we loop on the main while
                
538 539 540
            else:
                
                # There is no more link on this state
541
                # --> we cannot go deeper
542
                                
543 544 545 546
                links = self.longestSuffix(buffer,pos)
                errors[stack[pos]]=links
                                           
                # we pop out the state
547 548
                pos-=1
                j = buffer[pos]+1
549 550 551 552 553 554
         
        errors[0]=NULL       

#        for i in range(1,self._statecount):
#            if errors[i]==NULL:
#                print "no link on %d" % i
555 556 557 558
                
        for i in range(self._statecount):
            curstate=self._states + i 
            for j in range(4):
559 560 561 562 563 564 565 566 567 568 569 570 571 572
                nextstate=curstate
                table=&(nextstate.a)
                if table[j]==NULL:
                    while nextstate!=NULL and table[j]==NULL:
                        nextstate = errors[nextstate - self._states]
                        table=&(nextstate.a)
                    if nextstate==NULL:
#                        print "jump to %d[%d] -> root" % (i,j)
                        nextstate = self._states
                    else:
                        nextstate = table[j]
#                        print "jump to %d[%d] -> %d" % (i,j,nextstate-self._states)
                    table =&(curstate.a) 
                    table[j]=nextstate
573 574 575 576
                
                
        free(errors)
        free(stack)
577
        free(buffer)
578 579 580 581 582 583

    cpdef finalize(self):
        cdef size_t i 
        cdef size_t j 
        cdef size_t dest
        cdef dnastate* state
Eric Coissac committed
584
        cdef dnastate** table  # @DuplicatedSignature
585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608
        cdef protmatch* match
        
        if self._finalized:
            return None
        
        self._finalized = True

        for i in range(self._statecount):
            state = self._states + i
            
            if state.match!=<protmatch*>-1:
                state.match = self._matches + <size_t>state.match
            else:
                state.match = NULL
                
            table=&(state.a)
            for j in range(4):
                if table[j]!=NULL:
                    table[j]=self._states + <size_t>table[j]
                    
        for i in range(self._matchcount):
            match = self._matches + i
            if match.next!=NULL:
                match.next=self._matches + <size_t>match.next
609
                                                            
610 611 612 613 614 615 616
        self.setErrorLink()          
        
      
    cdef int addCWord(self,char* cword, size_t lword, int protid, int position):
        cdef int   letter
        cdef int   i
        cdef dnastate* state = self._states
617
        cdef protmatch* match  # @DuplicatedSignature
618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710
                        
        for i in range(lword):
            letter = cword[i] 
            letter >>=1
            letter &=3
            state = self.getNextState(state,letter)
            
            if state==NULL:
                return 0
            
        match = self.setAsTerminal(state,protid,position)
            
        if match == NULL:
            return 0
         
        if lword > self._depth:
            self._depth=lword  
             
        return self._statecount
    
    
    cdef int addAutomata(self,dict automata,int protid, int position):
        cdef list       stack=[(automata,0,0)]
        cdef tuple      stackitem
        cdef PyObject*  pkey
        cdef char*      ckey  
        cdef int        letter
        cdef PyObject*  pvalue
        cdef dict       value
        cdef Py_ssize_t ppos=0
        cdef dnastate*  state
        cdef dnastate*  nstate
        cdef protmatch* match
        cdef size_t     stateid
        cdef size_t     lword=0
       
        while PyList_GET_SIZE(stack)>0 :
            stackitem = stack.pop()
            automata  = <object>PyTuple_GET_ITEM(stackitem,0)
            stateid   = PyInt_AsSsize_t(<object>PyTuple_GET_ITEM(stackitem,1))
            lword     = PyInt_AsSsize_t(<object>PyTuple_GET_ITEM(stackitem,2))
            state     = self._states + stateid
            ppos      = 0
            
            lword+=1
            
            while PyDict_Next(automata, &ppos, &pkey, &pvalue):
                
                if lword > self._depth:
                    self._depth=lword
                #
                # Get the nucleotide code
                #
                
                ckey = PyBytes_AS_STRING(pkey)
                letter = ckey[0]
                letter>>=1
                letter&=3
                
                #
                # create or gate the corresponding node in the automata
                #
                
                nstate = self.getNextState(state,letter)
                
                # potentially recompute state address if reallocation occured 
                state     = self._states + stateid
                
                if state==NULL:
                    return 0
                
                # 
                # if value is an empty dict then this is a terminal state
                # else we have to push it on the stack
                #
                
                value=<object>pvalue
                
#                print stateid,"->",(nstate - self._states),value
                
                if PyDict_Size(value):
                    stack.append((value,(nstate - self._states),lword))
#                    print "on continue"
                else:
#                    print "Etat final"
                    match = self.setAsTerminal(nstate,protid,position)       
#                    print "Etat final ok"
                    if match == NULL:
                        return 0
                   
        return self._statecount
                
        
711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767
    
    
    cpdef dict match(self,bytes sequence):
        cdef char*  csequence = sequence 
        cdef size_t lseq      = len(sequence)
        cdef size_t pos      
        cdef dnastate* state 
        cdef dnastate* nstate 
        cdef dnastate** table 
        cdef protmatch* pmatch
        cdef int letter
        cdef dict results = {}
        cdef PyObject*  plpos
        cdef list lpos
        cdef int protid
        
                
        if not self._finalized:
            self.finalize()
        
        state = self._states
        
        for pos in range(lseq):
            table = &(state.a)
            letter = csequence[pos]
            letter >>=1
            letter &=3
            nstate = table[letter]
            pmatch = nstate.match
            while (pmatch!=NULL):
                plpos = PyDict_GetItem(results,PyInt_FromLong(pmatch.protid))
                if plpos==NULL:
                    lpos = []
                    PyDict_SetItem(results,PyInt_FromLong(pmatch.protid),lpos)
                else:
                    lpos = <object>plpos
                lpos.append((pos,pmatch.position))
                pmatch = pmatch.next
            state = nstate
                    
        return results 
                
                
           
            
            
        
    
    cpdef DiGraphMultiEdge asGraph(self):
        cdef DiGraphMultiEdge g 
        cdef size_t i 
        cdef size_t j 
        cdef size_t dest
        cdef dnastate* ndest
        cdef dict node
        cdef dict edge
        cdef dnastate** table
Eric Coissac committed
768
        cdef protmatch* match  # @DuplicatedSignature
769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837
        cdef size_t matchcount = 0
        cdef size_t pnode
        
        g=DiGraphMultiEdge(b"AhoCorasick")
        
        for i in range(self._statecount):
            g.addNode(i)
            node=g.getNodeAttr(i)
            node[b'label'] = b"%d" % i
            node[b'graphics'] = {}
            if i==0:
                node[b'graphics'][b'fill']=b'#008000'
                node[b'graphics'][b'type']=b'rectangle'  
                
            match = self._states[i].match    
            if ((not self._finalized and match!=<protmatch*>-1) or
                (self._finalized and match!=NULL)):
                node[b'graphics'][b'fill']=b'#800000'
                node[b'graphics'][b'type']=b'triangle'
                
                if not self._finalized:
                    match = self._matches + <size_t>match
                
                pnode = i
                while(match!=NULL):
                    #print "%d" % <size_t>match
                    matchcount+=1
                    g.addNode(self._statecount + matchcount)
                    node = g.getNodeAttr(self._statecount + matchcount)
                    node[b'graphics'] = {"type" : 'circle',
                                         'fill' : "#000080"
                                        }
                    node[b'label'] = "%d,%d" % (match.protid,match.position)
                    g.addEdge(pnode,self._statecount + matchcount)  
                    edge = g.getEdgeAttr(pnode,self._statecount + matchcount)
                    edge['graphics']={b'arrow':b'last',
                                      b'fill':b'#000040'
                                     }            
                    pnode = self._statecount + matchcount
                    match = match.next
                    if match!=NULL and not self._finalized:
                        match = self._matches + <size_t>match
                
        for i in range(self._statecount):
            table=&(self._states[i].a)
            for j in range(4):
                if self._finalized:
                    ndest = table[j]
                    if ndest:
                        dest  = ndest - self._states
                    else:
                        dest = 0
                else:
                    dest  = <size_t>table[j]
                    ndest = <dnastate*>-1
                    if dest==0:
                        ndest=NULL
                    
                if ndest != NULL:
                    g.addEdge(i,dest)
                    edge = g.getEdgeAttr(i,dest)
                    edge[b"label"]=b'%s' % (b"actg"[j])
                    edge['graphics']={b'arrow':b'last',
                                      b'fill':b'#080000'
                                     }            
        
        return g
    

838
cdef class ProtAhoCorasick(AhoCorasick):
839 840

    cpdef int addSequence(self,bytes sequence, object seqid=1, size_t kup=4):
Eric Coissac committed
841 842 843 844 845 846 847
        cdef int    position
        cdef int    k
        cdef char*  csequence
        cdef char   letter[2]
        cdef dict   c
        cdef dict   cn
        cdef int    osize
848
        cdef size_t erreur
Eric Coissac committed
849
        cdef int    id
850 851 852 853 854
        
        if self._finalized:
            return 0
        
        
Eric Coissac committed
855
        if seqid in self._seqid:
856
            id = self._seqid[seqid]
857 858
        else:
            id = self._maxseqid + 1
859 860
            self._seqid[seqid]=id
            self._rseqid[id]=seqid
861 862
            
        # register the size of the automata before add the protein
Eric Coissac committed
863
                
864 865 866 867
        osize = self._statecount
        
        letter[1]=0
        csequence = sequence
Eric Coissac committed
868
                
869
        # Backtranslate according to the forward strand
Eric Coissac committed
870
        
871 872 873 874 875 876
        for position in range(len(sequence)-kup+1):
            
            # for each peptide of size kup build a small automata with all
            # codon variants
            
            c = codon(csequence[position+kup-1])
Eric Coissac committed
877
            
Eric Coissac committed
878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894
            if c:
                erreur=0
                for k in range(position+kup-2,position-1,-1):
                    cn = codon(csequence[k])
                    if cn:
                        c = bindCodons(cn,c)
                                                    
                    else:
                        erreur=1
                        
                if erreur==0:
                    c = word2automata([w for w in enumerateword(c) if not homopolymer(w)])
                    erreur = self.addAutomata(c, id, (position+kup) * 3 - 1)
                    
                    if erreur==0:
                        return 0
                                
895 896 897 898 899 900 901 902 903
        # Backtranslate according to the reverse strand
                
        for position in range(len(sequence)-1,kup-2,-1):
            
            # for each peptide of size kup build a small automata with all
            # anticodon variants
            
            c = codon(csequence[position-kup+1],False)
            
Eric Coissac committed
904 905 906 907 908 909 910 911 912 913 914 915 916 917 918
            if c:
                erreur=0
                for k in range(position-kup+2,position+1):
                    cn = codon(csequence[k],False)
                    if cn:
                        c = bindCodons(cn,c)
                    else:
                        erreur=1
                        
                if erreur==0:
                    c = word2automata([w for w in enumerateword(c) if not homopolymer(w)])
        
                    erreur = self.addAutomata(c, -id, position * 3 + 2)
                    if erreur==0:
                        return 0
919 920

        return self._statecount - osize
921 922

    cpdef scanIndex(self,Index seqindex,int minmatch=15, int maxmatch=-1, int covmin=1, bint progress=True):
923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947
        """
        
        
            @param seqindex: A NGS read index
            @type seqindex: orgasm.indexer._orgasm.Index instance
            
            @param minmatch: minimum count of identified word to take into
                             account the match
            @type minimum: int
                             
            @param maxmatch: maximum count of identified word to take into
                             account the match. Repeated can leads to huge
                             wordcount without real similarity
                             By default (-1) this value is set to readsize/3
                             (I.e amino acid count).
            @type maximum: int
            
            @param covmin: minimum count observed for a read to be taken into 
                           account
            @type covmin: int
            
            @param progress: Boolean flg indicating if a progress bar must be
                             displayed
            @type progress: bool
        """
948
        cdef unsigned char* nucs
949 950
        cdef size_t i  # @DuplicatedSignature
        cdef size_t j  # @DuplicatedSignature
951 952
        cdef int k
        cdef list ks
953 954 955
        cdef size_t readcount  = seqindex._index.readCount
        cdef size_t recordsize = seqindex._index.recordSize
        cdef size_t readsize   = seqindex._index.readSize
956 957
        
        # the number of non-overlaping 4mer in a read
958
        cdef size_t readsize4  = (readsize >> 2) + (1 if (readsize & 3)  else 0)
959 960
        
        # the number of non-overlaping 32mer in a read
961
        cdef size_t readsize64 = (readsize >> 5) + (1 if (readsize & 31) else 0)
962 963
        
        # Pointer to the first read record
964 965 966
        cdef unsigned char*  records    = <unsigned char*>seqindex._index.records
        cdef unsigned char*  nuc4
        cdef unsigned char*  nuc4bis
967
        cdef BitVector   wordcount
968
        cdef size_t wordmax=0
969 970
        cdef double shamin
        cdef size_t wc[6]
971 972 973
        cdef size_t rp[6]
        cdef dict matchpos
        
974 975 976 977 978 979
        cdef size_t* count
        cdef int wct
        cdef int wcmax
        cdef int phase 
        cdef int protidmax=0
        cdef int framemax=0
980 981
        cdef int loc
        cdef int locmax
982
        
983 984 985 986 987 988 989 990 991
        cdef size_t nbreads
        cdef dnastate* state 
        cdef dnastate* nstate 
        cdef dnastate** table 
        cdef int letter
        cdef dict results = {}
        cdef int readid
        cdef PyObject*  plpos        
        cdef list lpos
992
        cdef int pid 
993 994 995 996 997
        cdef int p 
        cdef double  shanon
        cdef size_t  offset
        cdef double* shanonTable = buildShanonTable(readsize,&offset)
        cdef double* shanonLine
998
        cdef ProgressBar pb
999
        
1000 1001 1002 1003 1004
        # mid is the number of read frame globally available for each 
        # protein stored in the automata
        cdef int mid = self._maxseqid * 6 + 3
        
        # Finalize the automata if needed
1005 1006 1007
        if not self._finalized:
            self.finalize()
        
1008
        # Computes the default maxmatch according to the read size
1009
        if maxmatch < 0:
Eric Coissac committed
1010
            maxmatch = readsize // 3
1011
            
1012 1013
        i=0
        
1014
        wordcount = BitVector(mid,self._maxpos)
1015
        
1016
        nuc4 = records + i * (recordsize)
1017
        
1018
        if (progress):
1019 1020
            pb=ProgressBar(readcount)
#            progressBar(1,readcount,True)
1021

1022
        while i < readcount:
1023 1024
                
            # We start a new read so we set counter to 0
1025 1026
            wordcount.clear()     
            matchpos={}   
1027 1028
            
            readid = i
1029
#            print "@",i
1030
            
Eric Coissac committed
1031
            # setup automata to the initial state
1032 1033 1034 1035 1036 1037 1038 1039 1040
            state = self._states
            table = &(state.a)
            for j in range(readsize4):
                nucs = <unsigned char*>&(expanded8bitsnuc[nuc4[j]])
                for k in range(4):
                    letter = nucs[k]
                    nstate = table[letter]
                    pmatch = nstate.match
                    while (pmatch!=NULL):
1041
#                        print "coucou %d %d" %(pmatch.protid,pmatch.position)
1042
                        pid = pmatch.protid * 3 + ((j*4+k) % 3)
1043 1044 1045
                        if pid < 0 :
                            pid += mid
                        wordcount.set(pid,pmatch.position)
1046
                        matchpos[pid]=min(matchpos.get(pid,65535),pmatch.position)
1047 1048 1049 1050
                        pmatch = pmatch.next
                    state = nstate
                    table = &(state.a)
                    
1051
            shamin=1.0
1052 1053
            wordmax=0
            protidmax=0
1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064
            framemax=0
            
            count = wordcount.counter()
            
            for k in range(1,self._maxseqid+1):
                pid = k * 3
                
                wc[3]=count[pid]
                wc[4]=count[pid+1]
                wc[5]=count[pid+2]
                
1065 1066 1067 1068
                rp[3]=matchpos.get(pid,65535)
                rp[4]=matchpos.get(pid+1,65535)
                rp[5]=matchpos.get(pid+2,65535)
                
1069 1070 1071 1072 1073 1074 1075
                #memcpy(<void*>(wc+3),<void*>(count+pid),3*sizeof(size_t))
                
                pid = mid - pid 
                 
                wc[0]=count[pid]
                wc[1]=count[pid+1]
                wc[2]=count[pid+2]
1076 1077 1078 1079

                rp[0]=matchpos.get(pid,65535)
                rp[1]=matchpos.get(pid+1,65535)
                rp[2]=matchpos.get(pid+2,65535)
1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100
                
                # memcpy(<void*>(wc),<void*>(count+pid),3*sizeof(size_t))
                
                wct = wc[0] + wc[1] + wc[2] + \
                      wc[3] + wc[4] + wc[5]
                                            
                if wct > minmatch:
                    shanonLine = shanonTable + (wct << offset)
                    shanon=shanonLine[wc[0]] + \
                           shanonLine[wc[1]] + \
                           shanonLine[wc[2]] + \
                           shanonLine[wc[3]] + \
                           shanonLine[wc[4]] + \
                           shanonLine[wc[5]]
                
                    wcmax=0
                    phase=0
                    for p in range(6):
                        if wc[p]>wcmax:
                            wcmax=wc[p]
                            phase=p-3
1101
                            loc=rp[p]
1102 1103 1104 1105 1106
                else:
                    shanon=1.0
                    wcmax=0
                    phase=0

1107 1108 1109
                if ((shanon<0.1) and
                    (wcmax >=minmatch) and
                    (wcmax >wordmax)):
1110 1111
                    shamin = shanon
                    wordmax = wcmax
1112
                    protidmax=k
1113 1114 1115
                    framemax= phase
                    locmax=loc
                    
1116

1117 1118 1119 1120 1121
                
            nbreads=1
            nuc4bis=nuc4 + recordsize
            i+=1

1122
            if (progress):
1123
                pb(i)
1124

1125 1126 1127 1128 1129 1130
            while (i< readcount and 
                   samereads(nuc4, nuc4bis, readsize64)):
                nbreads+=1
                nuc4bis+= recordsize
                i+=1
                
1131 1132
#            print "Best match : %d -> %d (%d)" % (protidmax,wordmax,nbreads) 

1133 1134
            nuc4 = nuc4bis
            
1135 1136
            if protidmax>0 and shamin < 0.1 and nbreads >= covmin:
                if framemax < 0:
1137 1138 1139 1140 1141 1142 1143 1144 1145 1146
                    readid = -readid - 1
                else:
                    readid+=1 
                    
                plpos = PyDict_GetItem(results,PyInt_FromLong(protidmax))
                if plpos==NULL:
                    lpos = []
                    PyDict_SetItem(results,PyInt_FromLong(protidmax),lpos)
                else:
                    lpos = <object>plpos
1147
                lpos.append((readid,wordmax,nbreads,framemax,shamin,locmax))
1148 1149
        
        #<------------------------------ End of the while loop ------------------------>
1150
                
Eric Coissac committed
1151
        ks = list(results.keys())
1152 1153 1154 1155 1156 1157

        for k in ks:
            results[self._rseqid[k]]=results[k]
            del results[k]
        

1158
        free(shanonTable)          
1159 1160
        return results
                
1161
 
1162
cdef class NucAhoCorasick(AhoCorasick):
1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183

    cpdef int addSequence(self,bytes sequence, object seqid=1, size_t kup=12):
        cdef int   position
        cdef int   k
        cdef char* csequence
        cdef char  letter[2]
        cdef dict  c
        cdef int osize
        cdef size_t erreur
        cdef int   id
        
        if self._finalized:
            return 0
        
        
        if self._seqid.has_key(seqid):
            id = self._seqid[seqid]
        else:
            id = self._maxseqid + 1
            self._seqid[seqid]=id
            self._rseqid[id]=seqid
1184
            
1185 1186 1187
        # register the size of the automata before add the protein
        
        osize = self._statecount
1188
        
1189 1190 1191 1192 1193 1194 1195 1196 1197 1198
        # Backtranslate according to the forward strand

        for position in range(len(sequence)-kup+1):
            
            # for each peptide of size kup build a small automata with all
            # codon variants
            self.addWord(sequence[position:position+kup],id,position)
            self.addWord(reverseComplement(sequence[position:position+kup]),-id,position)

        return self._statecount - osize
1199
    
1200 1201
    cpdef scanIndex(self,Index seqindex,int minmatch=5, int maxmatch=-1, int covmin=1, bint progress=True):
        """
1202 1203
        
        
1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248
            @param seqindex: A NGS read index
            @type seqindex: orgasm.indexer._orgasm.Index instance
            
            @param minmatch: minimum count of identified word to take into
                             account the match
            @type minimum: int
                             
            @param maxmatch: maximum count of identified word to take into
                             account the match. Repeated can leads to huge
                             wordcount without real similarity
                             By default (-1) this value is set to readsize/3
                             (I.e amino acid count).
            @type maximum: int
            
            @param covmin: minimum count observed for a read to be taken into 
                           account
            @type covmin: int
            
            @param progress: Boolean flg indicating if a progress bar must be
                             displayed
            @type progress: bool
        """
        cdef unsigned char* nucs
        cdef size_t i  # @DuplicatedSignature
        cdef size_t j  # @DuplicatedSignature
        cdef int k
        cdef list ks
        cdef size_t readcount  = seqindex._index.readCount
        cdef size_t recordsize = seqindex._index.recordSize
        cdef size_t readsize   = seqindex._index.readSize
        
        # the number of non-overlaping 4mer in a read
        cdef size_t readsize4  = (readsize >> 2) + (1 if (readsize & 3)  else 0)
        
        # the number of non-overlaping 32mer in a read
        cdef size_t readsize64 = (readsize >> 5) + (1 if (readsize & 31) else 0)
        
        # Pointer to the first read record
        cdef unsigned char*  records    = <unsigned char*>seqindex._index.records
        cdef unsigned char*  nuc4
        cdef unsigned char*  nuc4bis
        cdef BitVector   wordcount
        cdef size_t wordmax=0
        cdef double shamin
        cdef size_t wc[6]
1249 1250 1251
        cdef size_t rp[6]
        cdef dict matchpos
        
1252 1253 1254 1255 1256 1257
        cdef size_t* count
        cdef int wct
        cdef int wcmax
        cdef int phase 
        cdef int protidmax=0
        cdef int framemax=0
1258 1259
        cdef int loc
        cdef int locmax
1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291
        
        cdef size_t nbreads
        cdef dnastate* state 
        cdef dnastate* nstate 
        cdef dnastate** table 
        cdef int letter
        cdef dict results = {}
        cdef int readid
        cdef PyObject*  plpos        
        cdef list lpos
        cdef int pid 
        cdef int p 
        cdef size_t  offset
        
        # mid is the number of read frame globally available for each 
        # protein stored in the automata
        cdef int mid = self._maxseqid * 2 + 1
        # Finalize the automata if needed
        if not self._finalized:
            self.finalize()
        
        # Computes the default maxmatch according to the read size
        if maxmatch < 0:
            maxmatch = readsize 
            
        i=0
        
        wordcount = BitVector(mid,self._maxpos)
        
        nuc4 = records + i * (recordsize)
        
        if (progress):
1292
            pb=ProgressBar(readcount)
1293 1294

        while i < readcount:
1295
                
1296
            # We start a new read so we set counter to 0
1297 1298
            wordcount.clear()     
            matchpos={}   
1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316
            
            readid = i
#            print "@",i
            
            # setup automata to the initial state
            state = self._states
            table = &(state.a)
            for j in range(readsize4):
                nucs = <unsigned char*>&(expanded8bitsnuc[nuc4[j]])
                for k in range(4):
                    letter = nucs[k]
                    nstate = table[letter]
                    pmatch = nstate.match
                    while (pmatch!=NULL):
                        pid = pmatch.protid
                        if pid < 0 :
                            pid += mid
                        wordcount.set(pid,pmatch.position)
1317
                        matchpos[pid]=min(matchpos.get(pid,65535),pmatch.position)
1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329
                        pmatch = pmatch.next
                    state = nstate
                    table = &(state.a)
                    
            wordmax=0
            protidmax=0
            framemax=0
            
            count = wordcount.counter()
            
            for k in range(1,self._maxseqid+1):
                pid = k 
1330
                
1331
                wc[1]=count[pid]
1332
                rp[1]=matchpos.get(pid,65535)                
1333
                #memcpy(<void*>(wc+3),<void*>(count+pid),3*sizeof(size_t))
1334
                
1335 1336 1337
                pid = mid - pid 
                 
                wc[0]=count[pid]
1338
                rp[0]=matchpos.get(pid,65535)
1339 1340 1341 1342
                
                # memcpy(<void*>(wc),<void*>(count+pid),3*sizeof(size_t))
                
                wct = wc[0] + wc[1]
1343 1344 1345 1346 1347 1348 1349 1350
                if wc[0] > wc[1]:
                    phase=-1
                    wcmax=wc[0]
                    loc=rp[0]
                else:
                    phase=1
                    wcmax=wc[1]
                    loc=rp[1]
1351 1352 1353 1354 1355 1356
                

                if (wcmax >=minmatch):
                    wordmax = wcmax
                    protidmax=k
                    framemax= phase 
1357
                    locmax=loc
1358 1359 1360 1361 1362 1363 1364

                
            nbreads=1
            nuc4bis=nuc4 + recordsize
            i+=1

            if (progress):
1365
                pb(i)
1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379

            while (i< readcount and 
                   samereads(nuc4, nuc4bis, readsize64)):
                nbreads+=1
                nuc4bis+= recordsize
                i+=1
                
#            print "Best match : %d -> %d (%d)" % (protidmax,wordmax,nbreads) 

            nuc4 = nuc4bis
            
            if protidmax>0 and nbreads >= covmin:
                if framemax < 0:
                    readid = -readid - 1
1380
                else:
1381
                    readid+=1 
1382
                    
1383 1384 1385 1386 1387 1388
                plpos = PyDict_GetItem(results,PyInt_FromLong(protidmax))
                if plpos==NULL:
                    lpos = []
                    PyDict_SetItem(results,PyInt_FromLong(protidmax),lpos)
                else:
                    lpos = <object>plpos
1389
                lpos.append((readid,wordmax,nbreads,framemax,0,locmax))
Eric Coissac committed
1390
                
1391
                
Eric Coissac committed
1392
        ks = list(results.keys())
1393

1394 1395 1396 1397 1398 1399
        for k in ks:
            results[self._rseqid[k]]=results[k]
            del results[k]
        

        return results
1400