Option to export one vcf/csv file per chromosome.
[lightning.git] / tilelib.go
index 8c88956b97685df9a24176a812999103543d5051..8073eba18360733918b9fa705618e55e088581c0 100644 (file)
@@ -1,4 +1,4 @@
-package main
+package lightning
 
 import (
        "bufio"
@@ -7,6 +7,7 @@ import (
        "encoding/gob"
        "fmt"
        "io"
+       "os"
        "regexp"
        "runtime"
        "sort"
@@ -14,6 +15,7 @@ import (
        "sync"
        "sync/atomic"
 
+       "github.com/klauspost/pgzip"
        log "github.com/sirupsen/logrus"
        "golang.org/x/crypto/blake2b"
 )
@@ -60,9 +62,9 @@ type tileLibrary struct {
        variant        [][][blake2b.Size256]byte
        refseqs        map[string]map[string][]tileLibRef
        compactGenomes map[string][]tileVariantID
-       // count [][]int
-       seq      map[[blake2b.Size256]byte][]byte
-       variants int64
+       seq2           map[[2]byte]map[[blake2b.Size256]byte][]byte
+       seq2lock       map[[2]byte]sync.Locker
+       variants       int64
        // if non-nil, write out any tile variants added while tiling
        encoder *gob.Encoder
 
@@ -141,7 +143,7 @@ func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap m
                                        }
                                        return
                                }
-                               log.Tracef("loadCompactGenomes: cg %s tag %d variant %d => %d", cg.Name, tag, variant, newvariant)
+                               // log.Tracef("loadCompactGenomes: cg %s tag %d variant %d => %d", cg.Name, tag, variant, newvariant)
                                cg.Variants[i] = newvariant
                        }
                        if onLoadGenome != nil {
@@ -208,6 +210,238 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian
        return nil
 }
 
+func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGenome func(CompactGenome)) 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)
+                       }
+               }
+               return nil
+       }
+       log.Infof("LoadDir: walk dir %s", path)
+       err := walk(path)
+       if err != nil {
+               return err
+       }
+       ctx, cancel := context.WithCancel(ctx)
+       defer cancel()
+       var mtx sync.Mutex
+       allcgs := make([][]CompactGenome, len(files))
+       allcseqs := make([][]CompactSequence, len(files))
+       allvariantmap := map[tileLibRef]tileVariantID{}
+       errs := make(chan error, len(files))
+       log.Infof("LoadDir: read %d files", len(files))
+       for fileno, path := range files {
+               fileno, path := fileno, path
+               go func() {
+                       f, err := open(path)
+                       if err != nil {
+                               errs <- err
+                               return
+                       }
+                       defer f.Close()
+                       defer log.Infof("LoadDir: finished reading %s", path)
+
+                       var variantmap = map[tileLibRef]tileVariantID{}
+                       var cgs []CompactGenome
+                       var cseqs []CompactSequence
+                       err = DecodeLibrary(f, strings.HasSuffix(path, ".gz"), func(ent *LibraryEntry) error {
+                               if ctx.Err() != nil {
+                                       return ctx.Err()
+                               }
+                               if len(ent.TagSet) > 0 {
+                                       mtx.Lock()
+                                       if tilelib.taglib == nil || tilelib.taglib.Len() != len(ent.TagSet) {
+                                               // load first set of tags, or
+                                               // report mismatch if 2 sets
+                                               // have different #tags.
+                                               if err := tilelib.loadTagSet(ent.TagSet); err != nil {
+                                                       mtx.Unlock()
+                                                       return err
+                                               }
+                                       }
+                                       mtx.Unlock()
+                               }
+                               for _, tv := range ent.TileVariants {
+                                       variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence).Variant
+                               }
+                               cgs = append(cgs, ent.CompactGenomes...)
+                               cseqs = append(cseqs, ent.CompactSequences...)
+                               return nil
+                       })
+                       allcgs[fileno] = cgs
+                       allcseqs[fileno] = cseqs
+                       mtx.Lock()
+                       defer mtx.Unlock()
+                       for k, v := range variantmap {
+                               allvariantmap[k] = v
+                       }
+                       errs <- err
+               }()
+       }
+       for range files {
+               err := <-errs
+               if err != nil {
+                       return err
+               }
+       }
+
+       log.Info("LoadDir: loadCompactGenomes")
+       var flatcgs []CompactGenome
+       for _, cgs := range allcgs {
+               flatcgs = append(flatcgs, cgs...)
+       }
+       err = tilelib.loadCompactGenomes(flatcgs, allvariantmap, onLoadGenome)
+       if err != nil {
+               return err
+       }
+
+       log.Info("LoadDir: loadCompactSequences")
+       var flatcseqs []CompactSequence
+       for _, cseqs := range allcseqs {
+               flatcseqs = append(flatcseqs, cseqs...)
+       }
+       err = tilelib.loadCompactSequences(flatcseqs, allvariantmap)
+       if err != nil {
+               return err
+       }
+
+       log.Info("LoadDir done")
+       return nil
+}
+
+func (tilelib *tileLibrary) WriteDir(dir string) error {
+       nfiles := 128
+       files := make([]*os.File, nfiles)
+       for i := range files {
+               f, err := os.OpenFile(fmt.Sprintf("%s/library.%04d.gob.gz", dir, i), os.O_CREATE|os.O_WRONLY, 0666)
+               if err != nil {
+                       return err
+               }
+               defer f.Close()
+               files[i] = f
+       }
+       bufws := make([]*bufio.Writer, nfiles)
+       for i := range bufws {
+               bufws[i] = bufio.NewWriterSize(files[i], 1<<26)
+       }
+       zws := make([]*pgzip.Writer, nfiles)
+       for i := range zws {
+               zws[i] = pgzip.NewWriter(bufws[i])
+               defer zws[i].Close()
+       }
+       encoders := make([]*gob.Encoder, nfiles)
+       for i := range encoders {
+               encoders[i] = gob.NewEncoder(zws[i])
+       }
+
+       cgnames := make([]string, 0, len(tilelib.compactGenomes))
+       for name := range tilelib.compactGenomes {
+               cgnames = append(cgnames, name)
+       }
+       sort.Strings(cgnames)
+
+       log.Infof("WriteDir: writing %d files", nfiles)
+       ctx, cancel := context.WithCancel(context.Background())
+       defer cancel()
+       errs := make(chan error, nfiles)
+       for start := range files {
+               start := start
+               go func() {
+                       err := encoders[start].Encode(LibraryEntry{TagSet: tilelib.taglib.Tags()})
+                       if err != nil {
+                               errs <- err
+                               return
+                       }
+                       if start == 0 {
+                               // For now, just write all the refs to
+                               // the first file
+                               for name, tseqs := range tilelib.refseqs {
+                                       err := encoders[start].Encode(LibraryEntry{CompactSequences: []CompactSequence{{
+                                               Name:          name,
+                                               TileSequences: tseqs,
+                                       }}})
+                                       if err != nil {
+                                               errs <- err
+                                               return
+                                       }
+                               }
+                       }
+                       for i := start; i < len(cgnames); i += nfiles {
+                               err := encoders[start].Encode(LibraryEntry{CompactGenomes: []CompactGenome{{
+                                       Name:     cgnames[i],
+                                       Variants: tilelib.compactGenomes[cgnames[i]],
+                               }}})
+                               if err != nil {
+                                       errs <- err
+                                       return
+                               }
+                       }
+                       tvs := []TileVariant{}
+                       for tag := start; tag < len(tilelib.variant) && ctx.Err() == nil; tag += nfiles {
+                               tvs = tvs[:0]
+                               for idx, hash := range tilelib.variant[tag] {
+                                       tvs = append(tvs, TileVariant{
+                                               Tag:      tagID(tag),
+                                               Variant:  tileVariantID(idx + 1),
+                                               Blake2b:  hash,
+                                               Sequence: tilelib.hashSequence(hash),
+                                       })
+                               }
+                               err := encoders[start].Encode(LibraryEntry{TileVariants: tvs})
+                               if err != nil {
+                                       errs <- err
+                                       return
+                               }
+                       }
+                       errs <- nil
+               }()
+       }
+       for range files {
+               err := <-errs
+               if err != nil {
+                       return err
+               }
+       }
+       log.Info("WriteDir: flushing")
+       for i := range zws {
+               err := zws[i].Close()
+               if err != nil {
+                       return err
+               }
+               err = bufws[i].Flush()
+               if err != nil {
+                       return err
+               }
+               err = files[i].Close()
+               if err != nil {
+                       return err
+               }
+       }
+       log.Info("WriteDir: done")
+       return nil
+}
+
 // Load library data from rdr. Tile variants might be renumbered in
 // the process; in that case, genomes variants will be renumbered to
 // match.
@@ -248,6 +482,36 @@ func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool,
        return nil
 }
 
+func (tilelib *tileLibrary) dump(out io.Writer) {
+       printTV := func(tag int, variant tileVariantID) {
+               if variant < 1 {
+                       fmt.Fprintf(out, " -")
+               } else if tag >= len(tilelib.variant) {
+                       fmt.Fprintf(out, " (!tag=%d)", tag)
+               } else if int(variant) > len(tilelib.variant[tag]) {
+                       fmt.Fprintf(out, " (tag=%d,!variant=%d)", tag, variant)
+               } else {
+                       fmt.Fprintf(out, " %x", tilelib.variant[tag][variant-1][:8])
+               }
+       }
+       for refname, refseqs := range tilelib.refseqs {
+               for seqname, seq := range refseqs {
+                       fmt.Fprintf(out, "ref %s %s", refname, seqname)
+                       for _, libref := range seq {
+                               printTV(int(libref.Tag), libref.Variant)
+                       }
+                       fmt.Fprintf(out, "\n")
+               }
+       }
+       for name, cg := range tilelib.compactGenomes {
+               fmt.Fprintf(out, "cg %s", name)
+               for tag, variant := range cg {
+                       printTV(tag/2, variant)
+               }
+               fmt.Fprintf(out, "\n")
+       }
+}
+
 type importStats struct {
        InputFile              string
        InputLabel             string
@@ -463,12 +727,31 @@ func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
        vlock.Unlock()
 
        if tilelib.retainTileSequences && !dropSeq {
-               tilelib.mtx.Lock()
-               if tilelib.seq == nil {
-                       tilelib.seq = map[[blake2b.Size256]byte][]byte{}
+               seqCopy := append([]byte(nil), seq...)
+               if tilelib.seq2 == nil {
+                       tilelib.mtx.Lock()
+                       if tilelib.seq2 == nil {
+                               tilelib.seq2lock = map[[2]byte]sync.Locker{}
+                               m := map[[2]byte]map[[blake2b.Size256]byte][]byte{}
+                               var k [2]byte
+                               for i := 0; i < 256; i++ {
+                                       k[0] = byte(i)
+                                       for j := 0; j < 256; j++ {
+                                               k[1] = byte(j)
+                                               m[k] = map[[blake2b.Size256]byte][]byte{}
+                                               tilelib.seq2lock[k] = &sync.Mutex{}
+                                       }
+                               }
+                               tilelib.seq2 = m
+                       }
+                       tilelib.mtx.Unlock()
                }
-               tilelib.seq[seqhash] = append([]byte(nil), seq...)
-               tilelib.mtx.Unlock()
+               var k [2]byte
+               copy(k[:], seqhash[:])
+               locker := tilelib.seq2lock[k]
+               locker.Lock()
+               tilelib.seq2[k][seqhash] = seqCopy
+               locker.Unlock()
        }
 
        if tilelib.encoder != nil {
@@ -489,11 +772,17 @@ func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
        return tileLibRef{Tag: tag, Variant: variant}
 }
 
+func (tilelib *tileLibrary) hashSequence(hash [blake2b.Size256]byte) []byte {
+       var partition [2]byte
+       copy(partition[:], hash[:])
+       return tilelib.seq2[partition][hash]
+}
+
 func (tilelib *tileLibrary) TileVariantSequence(libref tileLibRef) []byte {
        if libref.Variant == 0 || len(tilelib.variant) <= int(libref.Tag) || len(tilelib.variant[libref.Tag]) < int(libref.Variant) {
                return nil
        }
-       return tilelib.seq[tilelib.variant[libref.Tag][libref.Variant-1]]
+       return tilelib.hashSequence(tilelib.variant[libref.Tag][libref.Variant-1])
 }
 
 // Tidy deletes unreferenced tile variants and renumbers variants so
@@ -564,18 +853,30 @@ func (tilelib *tileLibrary) Tidy() {
        // Apply remap to genomes and reference sequences, so they
        // refer to the same tile variants using the changed IDs.
        log.Print("Tidy: apply remap")
+       var wg sync.WaitGroup
        for _, cg := range tilelib.compactGenomes {
-               for idx, variant := range cg {
-                       cg[idx] = remap[tagID(idx/2)][variant]
-               }
+               cg := cg
+               wg.Add(1)
+               go func() {
+                       defer wg.Done()
+                       for idx, variant := range cg {
+                               cg[idx] = remap[tagID(idx/2)][variant]
+                       }
+               }()
        }
        for _, refcs := range tilelib.refseqs {
                for _, refseq := range refcs {
-                       for i, tv := range refseq {
-                               refseq[i].Variant = remap[tv.Tag][tv.Variant]
-                       }
+                       refseq := refseq
+                       wg.Add(1)
+                       go func() {
+                               defer wg.Done()
+                               for i, tv := range refseq {
+                                       refseq[i].Variant = remap[tv.Tag][tv.Variant]
+                               }
+                       }()
                }
        }
+       wg.Wait()
        log.Print("Tidy: done")
 }