Fix some tests.
[lightning.git] / tilelib.go
index a3b63ec78f9bbf2481301eb0f689c7eac2028edd..7d4599d4e254bda48a7c0b6436e6ffcaf1d44f30 100644 (file)
@@ -11,6 +11,7 @@ import (
        "encoding/gob"
        "fmt"
        "io"
+       "math/rand"
        "os"
        "regexp"
        "runtime"
@@ -61,6 +62,7 @@ type tileLibrary struct {
        retainNoCalls       bool
        skipOOO             bool
        retainTileSequences bool
+       useDups             bool
 
        taglib         *tagLibrary
        variant        [][][blake2b.Size256]byte
@@ -71,6 +73,12 @@ type tileLibrary struct {
        variants       int64
        // if non-nil, write out any tile variants added while tiling
        encoder *gob.Encoder
+       // set Ref flag when writing new variants to encoder
+       encodeRef bool
+
+       onAddTileVariant func(libref tileLibRef, hash [blake2b.Size256]byte, seq []byte) error
+       onAddGenome      func(CompactGenome) error
+       onAddRefseq      func(CompactSequence) error
 
        mtx   sync.RWMutex
        vlock []sync.Locker
@@ -116,12 +124,12 @@ func (tilelib *tileLibrary) loadTileVariants(tvs []TileVariant, variantmap map[t
        for _, tv := range tvs {
                // Assign a new variant ID (unique across all inputs)
                // for each input variant.
-               variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence).Variant
+               variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence, tv.Ref).Variant
        }
        return nil
 }
 
-func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap map[tileLibRef]tileVariantID, onLoadGenome func(CompactGenome)) error {
+func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap map[tileLibRef]tileVariantID) error {
        log.Debugf("loadCompactGenomes: %d", len(cgs))
        var wg sync.WaitGroup
        errs := make(chan error, 1)
@@ -150,8 +158,15 @@ func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap m
                                // log.Tracef("loadCompactGenomes: cg %s tag %d variant %d => %d", cg.Name, tag, variant, newvariant)
                                cg.Variants[i] = newvariant
                        }
-                       if onLoadGenome != nil {
-                               onLoadGenome(cg)
+                       if tilelib.onAddGenome != nil {
+                               err := tilelib.onAddGenome(cg)
+                               if err != nil {
+                                       select {
+                                       case errs <- err:
+                                       default:
+                                       }
+                                       return
+                               }
                        }
                        if tilelib.encoder != nil {
                                err := tilelib.encoder.Encode(LibraryEntry{
@@ -203,6 +218,12 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian
                                return err
                        }
                }
+               if tilelib.onAddRefseq != nil {
+                       err := tilelib.onAddRefseq(cseq)
+                       if err != nil {
+                               return err
+                       }
+               }
                log.Infof("loadCompactSequences: checking %s done", cseq.Name)
        }
        tilelib.mtx.Lock()
@@ -217,36 +238,39 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian
        return nil
 }
 
-func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGenome func(CompactGenome)) error {
+func allFiles(path string, re *regexp.Regexp) ([]string, error) {
        var files []string
-       var walk func(string) error
-       walk = func(path string) error {
-               f, err := open(path)
-               if err != nil {
-                       return err
-               }
-               defer f.Close()
-               fis, err := f.Readdir(-1)
-               if err != nil {
-                       files = append(files, path)
-                       return nil
-               }
-               for _, fi := range fis {
-                       if fi.Name() == "." || fi.Name() == ".." {
-                               continue
-                       } else if child := path + "/" + fi.Name(); fi.IsDir() {
-                               err = walk(child)
-                               if err != nil {
-                                       return err
-                               }
-                       } else if strings.HasSuffix(child, ".gob") || strings.HasSuffix(child, ".gob.gz") {
-                               files = append(files, child)
+       f, err := open(path)
+       if err != nil {
+               return nil, err
+       }
+       defer f.Close()
+       fis, err := f.Readdir(-1)
+       if err != nil {
+               return []string{path}, nil
+       }
+       for _, fi := range fis {
+               if fi.Name() == "." || fi.Name() == ".." {
+                       continue
+               } else if child := path + "/" + fi.Name(); fi.IsDir() {
+                       add, err := allFiles(child, re)
+                       if err != nil {
+                               return nil, err
                        }
+                       files = append(files, add...)
+               } else if re == nil || re.MatchString(child) {
+                       files = append(files, child)
                }
-               return nil
        }
+       sort.Strings(files)
+       return files, nil
+}
+
+var matchGobFile = regexp.MustCompile(`\.gob(\.gz)?$`)
+
+func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string) error {
        log.Infof("LoadDir: walk dir %s", path)
-       err := walk(path)
+       files, err := allFiles(path, matchGobFile)
        if err != nil {
                return err
        }
@@ -290,7 +314,7 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
                                        mtx.Unlock()
                                }
                                for _, tv := range ent.TileVariants {
-                                       variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence).Variant
+                                       variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence, tv.Ref).Variant
                                }
                                cgs = append(cgs, ent.CompactGenomes...)
                                cseqs = append(cseqs, ent.CompactSequences...)
@@ -318,7 +342,7 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
        for _, cgs := range allcgs {
                flatcgs = append(flatcgs, cgs...)
        }
-       err = tilelib.loadCompactGenomes(flatcgs, allvariantmap, onLoadGenome)
+       err = tilelib.loadCompactGenomes(flatcgs, allvariantmap)
        if err != nil {
                return err
        }
@@ -455,9 +479,7 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
 // Load library data from rdr. Tile variants might be renumbered in
 // the process; in that case, genomes variants will be renumbered to
 // match.
-//
-// If onLoadGenome is non-nil, call it on each CompactGenome entry.
-func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool, onLoadGenome func(CompactGenome)) error {
+func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool) error {
        cgs := []CompactGenome{}
        cseqs := []CompactSequence{}
        variantmap := map[tileLibRef]tileVariantID{}
@@ -481,7 +503,7 @@ func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool,
        if ctx.Err() != nil {
                return ctx.Err()
        }
-       err = tilelib.loadCompactGenomes(cgs, variantmap, onLoadGenome)
+       err = tilelib.loadCompactGenomes(cgs, variantmap)
        if err != nil {
                return err
        }
@@ -523,38 +545,17 @@ func (tilelib *tileLibrary) dump(out io.Writer) {
 }
 
 type importStats struct {
-       InputFile              string
-       InputLabel             string
-       InputLength            int
-       InputCoverage          int
-       PathLength             int
-       DroppedOutOfOrderTiles int
+       InputFile             string
+       InputLabel            string
+       InputLength           int
+       InputCoverage         int
+       PathLength            int
+       DroppedRepeatedTags   int
+       DroppedOutOfOrderTags int
 }
 
-func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChromosome *regexp.Regexp) (tileSeq, []importStats, error) {
+func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChromosome *regexp.Regexp, isRef bool) (tileSeq, []importStats, error) {
        ret := tileSeq{}
-       type jobT struct {
-               label string
-               fasta []byte
-       }
-       todo := make(chan jobT, 1)
-       scanner := bufio.NewScanner(rdr)
-       go func() {
-               defer close(todo)
-               var fasta []byte
-               var seqlabel string
-               for scanner.Scan() {
-                       buf := scanner.Bytes()
-                       if len(buf) > 0 && buf[0] == '>' {
-                               todo <- jobT{seqlabel, append([]byte(nil), fasta...)}
-                               seqlabel, fasta = strings.SplitN(string(buf[1:]), " ", 2)[0], fasta[:0]
-                               log.Debugf("%s %s reading fasta", filelabel, seqlabel)
-                       } else {
-                               fasta = append(fasta, bytes.ToLower(buf)...)
-                       }
-               }
-               todo <- jobT{seqlabel, fasta}
-       }()
        type foundtag struct {
                pos   int
                tagid tagID
@@ -566,84 +567,134 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro
        skippedSequences := 0
        taglen := tilelib.taglib.TagLen()
        var stats []importStats
-       for job := range todo {
-               if len(job.fasta) == 0 {
-                       continue
-               } else if !matchChromosome.MatchString(job.label) {
+
+       in := bufio.NewReader(rdr)
+readall:
+       for {
+               var seqlabel string
+               // Advance to next '>', then
+               // read seqlabel up to \r?\n
+       readseqlabel:
+               for seqlabelStarted := false; ; {
+                       rune, _, err := in.ReadRune()
+                       if err == io.EOF {
+                               break readall
+                       } else if err != nil {
+                               return nil, nil, err
+                       }
+                       switch {
+                       case rune == '\r':
+                       case seqlabelStarted && rune == '\n':
+                               break readseqlabel
+                       case seqlabelStarted:
+                               seqlabel += string(rune)
+                       case rune == '>':
+                               seqlabelStarted = true
+                       default:
+                       }
+               }
+               log.Debugf("%s %s reading fasta", filelabel, seqlabel)
+               if !matchChromosome.MatchString(seqlabel) {
                        skippedSequences++
                        continue
                }
-               log.Debugf("%s %s tiling", filelabel, job.label)
+               log.Debugf("%s %s tiling", filelabel, seqlabel)
 
+               fasta := bytes.NewBuffer(nil)
                found = found[:0]
-               tilelib.taglib.FindAll(job.fasta, func(tagid tagID, pos, taglen int) {
+               err := tilelib.taglib.FindAll(in, fasta, func(tagid tagID, pos, taglen int) {
                        found = append(found, foundtag{pos: pos, tagid: tagid})
                })
+               if err != nil {
+                       return nil, nil, err
+               }
                totalFoundTags += len(found)
                if len(found) == 0 {
-                       log.Warnf("%s %s no tags found", filelabel, job.label)
+                       log.Warnf("%s %s no tags found", filelabel, seqlabel)
                }
 
-               skipped := 0
+               droppedDup := 0
+               if !tilelib.useDups {
+                       // Remove any tags that appeared more than once
+                       dup := map[tagID]bool{}
+                       for _, ft := range found {
+                               _, dup[ft.tagid] = dup[ft.tagid]
+                       }
+                       dst := 0
+                       for _, ft := range found {
+                               if !dup[ft.tagid] {
+                                       found[dst] = ft
+                                       dst++
+                               }
+                       }
+                       droppedDup = len(found) - dst
+                       log.Infof("%s %s dropping %d non-unique tags", filelabel, seqlabel, droppedDup)
+                       found = found[:dst]
+               }
+
+               droppedOOO := 0
                if tilelib.skipOOO {
-                       log.Infof("%s %s keeping longest increasing subsequence", filelabel, job.label)
                        keep := longestIncreasingSubsequence(len(found), func(i int) int { return int(found[i].tagid) })
                        for i, x := range keep {
                                found[i] = found[x]
                        }
-                       skipped = len(found) - len(keep)
+                       droppedOOO = len(found) - len(keep)
+                       log.Infof("%s %s dropping %d out-of-order tags", filelabel, seqlabel, droppedOOO)
                        found = found[:len(keep)]
                }
 
-               log.Infof("%s %s getting %d librefs", filelabel, job.label, len(found))
-               throttle := &throttle{Max: runtime.NumCPU()}
+               log.Infof("%s %s getting %d librefs", filelabel, seqlabel, len(found))
                path = path[:len(found)]
                var lowquality int64
-               for i, f := range found {
-                       i, f := i, f
-                       throttle.Acquire()
-                       go func() {
-                               defer throttle.Release()
-                               var startpos, endpos int
-                               if i == 0 {
-                                       startpos = 0
-                               } else {
-                                       startpos = f.pos
-                               }
-                               if i == len(found)-1 {
-                                       endpos = len(job.fasta)
-                               } else {
-                                       endpos = found[i+1].pos + taglen
-                               }
-                               path[i] = tilelib.getRef(f.tagid, job.fasta[startpos:endpos])
-                               if countBases(job.fasta[startpos:endpos]) != endpos-startpos {
-                                       atomic.AddInt64(&lowquality, 1)
-                               }
-                       }()
+               // Visit each element of found, but start at a random
+               // index, to reduce the likelihood of lock contention
+               // when importing many samples concurrently.
+               startpoint := rand.Int() % len(found)
+               for offset := range found {
+                       i := startpoint + offset
+                       if i >= len(found) {
+                               i -= len(found)
+                       }
+                       f := found[i]
+                       var startpos, endpos int
+                       if i == 0 {
+                               startpos = 0
+                       } else {
+                               startpos = f.pos
+                       }
+                       if i == len(found)-1 {
+                               endpos = fasta.Len()
+                       } else {
+                               endpos = found[i+1].pos + taglen
+                       }
+                       path[i] = tilelib.getRef(f.tagid, fasta.Bytes()[startpos:endpos], isRef)
+                       if countBases(fasta.Bytes()[startpos:endpos]) != endpos-startpos {
+                               lowquality++
+                       }
                }
-               throttle.Wait()
 
-               log.Infof("%s %s copying path", filelabel, job.label)
+               log.Infof("%s %s copying path", filelabel, seqlabel)
 
                pathcopy := make([]tileLibRef, len(path))
                copy(pathcopy, path)
-               ret[job.label] = pathcopy
+               ret[seqlabel] = pathcopy
 
-               basesIn := countBases(job.fasta)
-               log.Infof("%s %s fasta in %d coverage in %d path len %d low-quality %d skipped-out-of-order %d", filelabel, job.label, len(job.fasta), basesIn, len(path), lowquality, skipped)
+               basesIn := countBases(fasta.Bytes())
+               log.Infof("%s %s fasta in %d coverage in %d path len %d low-quality %d", filelabel, seqlabel, fasta.Len(), basesIn, len(path), lowquality)
                stats = append(stats, importStats{
-                       InputFile:              filelabel,
-                       InputLabel:             job.label,
-                       InputLength:            len(job.fasta),
-                       InputCoverage:          basesIn,
-                       PathLength:             len(path),
-                       DroppedOutOfOrderTiles: skipped,
+                       InputFile:             filelabel,
+                       InputLabel:            seqlabel,
+                       InputLength:           fasta.Len(),
+                       InputCoverage:         basesIn,
+                       PathLength:            len(path),
+                       DroppedOutOfOrderTags: droppedOOO,
+                       DroppedRepeatedTags:   droppedDup,
                })
 
                totalPathLen += len(path)
        }
        log.Printf("%s tiled with total path len %d in %d sequences (skipped %d sequences that did not match chromosome regexp, skipped %d out-of-order tags)", filelabel, totalPathLen, len(ret), skippedSequences, totalFoundTags-totalPathLen)
-       return ret, stats, scanner.Err()
+       return ret, stats, nil
 }
 
 func (tilelib *tileLibrary) Len() int64 {
@@ -652,7 +703,7 @@ func (tilelib *tileLibrary) Len() int64 {
 
 // Return a tileLibRef for a tile with the given tag and sequence,
 // adding the sequence to the library if needed.
-func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
+func (tilelib *tileLibrary) getRef(tag tagID, seq []byte, usedByRef bool) tileLibRef {
        dropSeq := false
        if !tilelib.retainNoCalls {
                for _, b := range seq {
@@ -764,21 +815,25 @@ func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
                locker.Unlock()
        }
 
+       saveSeq := seq
+       if dropSeq {
+               // Save the hash, but not the sequence
+               saveSeq = nil
+       }
        if tilelib.encoder != nil {
-               saveSeq := seq
-               if dropSeq {
-                       // Save the hash, but not the sequence
-                       saveSeq = nil
-               }
                tilelib.encoder.Encode(LibraryEntry{
                        TileVariants: []TileVariant{{
                                Tag:      tag,
+                               Ref:      usedByRef,
                                Variant:  variant,
                                Blake2b:  seqhash,
                                Sequence: saveSeq,
                        }},
                })
        }
+       if tilelib.onAddTileVariant != nil {
+               tilelib.onAddTileVariant(tileLibRef{tag, variant}, seqhash, saveSeq)
+       }
        return tileLibRef{Tag: tag, Variant: variant}
 }