Plot import stats.
[lightning.git] / import.go
index e24b25dfac3ecc26a8aac0346835fb9b527b27f0..a43737232d06fb65816cb297f72e776b4d84de0f 100644 (file)
--- a/import.go
+++ b/import.go
@@ -4,9 +4,11 @@ import (
        "bufio"
        "compress/gzip"
        "encoding/gob"
+       "encoding/json"
        "errors"
        "flag"
        "fmt"
+       "image/color/palette"
        "io"
        "net/http"
        _ "net/http/pprof"
@@ -23,6 +25,10 @@ import (
 
        "git.arvados.org/arvados.git/sdk/go/arvados"
        log "github.com/sirupsen/logrus"
+       "gonum.org/v1/plot"
+       "gonum.org/v1/plot/plotter"
+       "gonum.org/v1/plot/vg"
+       "gonum.org/v1/plot/vg/draw"
 )
 
 type importer struct {
@@ -34,6 +40,7 @@ type importer struct {
        skipOOO        bool
        outputTiles    bool
        includeNoCalls bool
+       outputStats    string
        encoder        *gob.Encoder
 }
 
@@ -54,6 +61,7 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
        flags.BoolVar(&cmd.skipOOO, "skip-ooo", false, "skip out-of-order tags")
        flags.BoolVar(&cmd.outputTiles, "output-tiles", false, "include tile variant sequences in output file")
        flags.BoolVar(&cmd.includeNoCalls, "include-no-calls", false, "treat tiles with no-calls as regular tiles")
+       flags.StringVar(&cmd.outputStats, "output-stats", "", "output stats to `file` (json)")
        priority := flags.Int("priority", 500, "container request priority")
        pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
        loglevel := flags.String("loglevel", "info", "logging threshold (trace, debug, info, warn, error, fatal, or panic)")
@@ -118,6 +126,7 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
                        fmt.Sprintf("-skip-ooo=%v", cmd.skipOOO),
                        fmt.Sprintf("-output-tiles=%v", cmd.outputTiles),
                        fmt.Sprintf("-include-no-calls=%v", cmd.includeNoCalls),
+                       "-output-stats", "/mnt/output/stats.json",
                        "-tag-library", cmd.tagLibraryFile,
                        "-ref", cmd.refFile,
                        "-o", cmd.outputFile,
@@ -156,6 +165,7 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
 
        tilelib := &tileLibrary{taglib: taglib, includeNoCalls: cmd.includeNoCalls, skipOOO: cmd.skipOOO}
        if cmd.outputTiles {
+               cmd.encoder.Encode(LibraryEntry{TagSet: taglib.Tags()})
                tilelib.encoder = cmd.encoder
        }
        go func() {
@@ -179,17 +189,17 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
        return 0
 }
 
-func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, error) {
+func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, []importStats, error) {
        var input io.ReadCloser
        input, err := os.Open(infile)
        if err != nil {
-               return nil, err
+               return nil, nil, err
        }
        defer input.Close()
        if strings.HasSuffix(infile, ".gz") {
                input, err = gzip.NewReader(input)
                if err != nil {
-                       return nil, err
+                       return nil, nil, err
                }
                defer input.Close()
        }
@@ -223,12 +233,19 @@ func (cmd *importer) loadTagLibrary() (*tagLibrary, error) {
        return &taglib, nil
 }
 
+var (
+       vcfFilenameRe    = regexp.MustCompile(`\.vcf(\.gz)?$`)
+       fasta1FilenameRe = regexp.MustCompile(`\.1\.fa(sta)?(\.gz)?$`)
+       fasta2FilenameRe = regexp.MustCompile(`\.2\.fa(sta)?(\.gz)?$`)
+       fastaFilenameRe  = regexp.MustCompile(`\.fa(sta)?(\.gz)?$`)
+)
+
 func listInputFiles(paths []string) (files []string, err error) {
        for _, path := range paths {
                if fi, err := os.Stat(path); err != nil {
                        return nil, fmt.Errorf("%s: stat failed: %s", path, err)
                } else if !fi.IsDir() {
-                       if !strings.HasSuffix(path, ".2.fasta") || strings.HasSuffix(path, ".2.fasta.gz") {
+                       if !fasta2FilenameRe.MatchString(path) {
                                files = append(files, path)
                        }
                        continue
@@ -244,23 +261,27 @@ func listInputFiles(paths []string) (files []string, err error) {
                }
                sort.Strings(names)
                for _, name := range names {
-                       if strings.HasSuffix(name, ".vcf") || strings.HasSuffix(name, ".vcf.gz") {
+                       if vcfFilenameRe.MatchString(name) {
                                files = append(files, filepath.Join(path, name))
-                       } else if strings.HasSuffix(name, ".1.fasta") || strings.HasSuffix(name, ".1.fasta.gz") {
+                       } else if fastaFilenameRe.MatchString(name) && !fasta2FilenameRe.MatchString(name) {
                                files = append(files, filepath.Join(path, name))
                        }
                }
                d.Close()
        }
        for _, file := range files {
-               if strings.HasSuffix(file, ".1.fasta") || strings.HasSuffix(file, ".1.fasta.gz") {
-                       continue
-               } else if _, err := os.Stat(file + ".csi"); err == nil {
-                       continue
-               } else if _, err = os.Stat(file + ".tbi"); err == nil {
+               if fastaFilenameRe.MatchString(file) {
                        continue
+               } else if vcfFilenameRe.MatchString(file) {
+                       if _, err := os.Stat(file + ".csi"); err == nil {
+                               continue
+                       } else if _, err = os.Stat(file + ".tbi"); err == nil {
+                               continue
+                       } else {
+                               return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
+                       }
                } else {
-                       return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
+                       return nil, fmt.Errorf("don't know how to handle filename %s", file)
                }
        }
        return
@@ -270,48 +291,77 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
        starttime := time.Now()
        errs := make(chan error, 1)
        todo := make(chan func() error, len(infiles)*2)
+       allstats := make([][]importStats, len(infiles)*2)
        var encodeJobs sync.WaitGroup
-       for _, infile := range infiles {
-               infile := infile
+       for idx, infile := range infiles {
+               idx, infile := idx, infile
                var phases sync.WaitGroup
                phases.Add(2)
                variants := make([][]tileVariantID, 2)
-               if strings.HasSuffix(infile, ".1.fasta") || strings.HasSuffix(infile, ".1.fasta.gz") {
+               if fasta1FilenameRe.MatchString(infile) {
                        todo <- func() error {
                                defer phases.Done()
                                log.Printf("%s starting", infile)
                                defer log.Printf("%s done", infile)
-                               tseqs, err := cmd.tileFasta(tilelib, infile)
+                               tseqs, stats, err := cmd.tileFasta(tilelib, infile)
+                               allstats[idx*2] = stats
                                var kept, dropped int
                                variants[0], kept, dropped = tseqs.Variants()
                                log.Printf("%s found %d unique tags plus %d repeats", infile, kept, dropped)
                                return err
                        }
-                       infile2 := regexp.MustCompile(`\.1\.fasta(\.gz)?$`).ReplaceAllString(infile, `.2.fasta$1`)
+                       infile2 := fasta1FilenameRe.ReplaceAllString(infile, `.2.fa$1$2`)
                        todo <- func() error {
                                defer phases.Done()
                                log.Printf("%s starting", infile2)
                                defer log.Printf("%s done", infile2)
-                               tseqs, err := cmd.tileFasta(tilelib, infile2)
+                               tseqs, stats, err := cmd.tileFasta(tilelib, infile2)
+                               allstats[idx*2+1] = stats
                                var kept, dropped int
                                variants[1], kept, dropped = tseqs.Variants()
-                               log.Printf("%s found %d unique tags plus %d repeats", infile, kept, dropped)
+                               log.Printf("%s found %d unique tags plus %d repeats", infile2, kept, dropped)
+
                                return err
                        }
-               } else {
+               } else if fastaFilenameRe.MatchString(infile) {
+                       todo <- func() error {
+                               defer phases.Done()
+                               defer phases.Done()
+                               log.Printf("%s starting", infile)
+                               defer log.Printf("%s done", infile)
+                               tseqs, stats, err := cmd.tileFasta(tilelib, infile)
+                               allstats[idx*2] = stats
+                               if err != nil {
+                                       return err
+                               }
+                               totlen := 0
+                               for _, tseq := range tseqs {
+                                       totlen += len(tseq)
+                               }
+                               log.Printf("%s tiled %d seqs, total len %d", infile, len(tseqs), totlen)
+                               return cmd.encoder.Encode(LibraryEntry{
+                                       CompactSequences: []CompactSequence{{Name: infile, TileSequences: tseqs}},
+                               })
+                       }
+                       // Don't write out a CompactGenomes entry
+                       continue
+               } else if vcfFilenameRe.MatchString(infile) {
                        for phase := 0; phase < 2; phase++ {
                                phase := phase
                                todo <- func() error {
                                        defer phases.Done()
                                        log.Printf("%s phase %d starting", infile, phase+1)
                                        defer log.Printf("%s phase %d done", infile, phase+1)
-                                       tseqs, err := cmd.tileGVCF(tilelib, infile, phase)
+                                       tseqs, stats, err := cmd.tileGVCF(tilelib, infile, phase)
+                                       allstats[idx*2] = stats
                                        var kept, dropped int
                                        variants[phase], kept, dropped = tseqs.Variants()
                                        log.Printf("%s phase %d found %d unique tags plus %d repeats", infile, phase+1, kept, dropped)
                                        return err
                                }
                        }
+               } else {
+                       panic(fmt.Sprintf("bug: unhandled filename %q", infile))
                }
                encodeJobs.Add(1)
                go func() {
@@ -320,20 +370,8 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
                        if len(errs) > 0 {
                                return
                        }
-                       ntags := len(variants[0])
-                       if ntags < len(variants[1]) {
-                               ntags = len(variants[1])
-                       }
-                       flat := make([]tileVariantID, ntags*2)
-                       for i := 0; i < ntags; i++ {
-                               for hap := 0; hap < 2; hap++ {
-                                       if i < len(variants[hap]) {
-                                               flat[i*2+hap] = variants[hap][i]
-                                       }
-                               }
-                       }
                        err := cmd.encoder.Encode(LibraryEntry{
-                               CompactGenomes: []CompactGenome{{Name: infile, Variants: flat}},
+                               CompactGenomes: []CompactGenome{{Name: infile, Variants: flatten(variants)}},
                        })
                        if err != nil {
                                select {
@@ -364,19 +402,42 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
                                        }
                                }
                                remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
-                               ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
-                               eta := time.Now().Add(ttl)
-                               log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
+                               if remain < cap(todo) {
+                                       ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
+                                       eta := time.Now().Add(ttl)
+                                       log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
+                               }
                        }
                }()
        }
        tileJobs.Wait()
        encodeJobs.Wait()
+
        go close(errs)
-       return <-errs
+       err := <-errs
+       if err != nil {
+               return err
+       }
+
+       if cmd.outputStats != "" {
+               f, err := os.OpenFile(cmd.outputStats, os.O_CREATE|os.O_WRONLY, 0666)
+               if err != nil {
+                       return err
+               }
+               var flatstats []importStats
+               for _, stats := range allstats {
+                       flatstats = append(flatstats, stats...)
+               }
+               err = json.NewEncoder(f).Encode(flatstats)
+               if err != nil {
+                       return err
+               }
+       }
+
+       return nil
 }
 
-func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, err error) {
+func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, stats []importStats, err error) {
        if cmd.refFile == "" {
                err = errors.New("cannot import vcf: reference data (-ref) not specified")
                return
@@ -408,7 +469,7 @@ func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (t
                return
        }
        defer consensus.Wait()
-       tileseq, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
+       tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
        if err != nil {
                return
        }
@@ -423,3 +484,82 @@ func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (t
        }
        return
 }
+
+func flatten(variants [][]tileVariantID) []tileVariantID {
+       ntags := 0
+       for _, v := range variants {
+               if ntags < len(v) {
+                       ntags = len(v)
+               }
+       }
+       flat := make([]tileVariantID, ntags*2)
+       for i := 0; i < ntags; i++ {
+               for hap := 0; hap < 2; hap++ {
+                       if i < len(variants[hap]) {
+                               flat[i*2+hap] = variants[hap][i]
+                       }
+               }
+       }
+       return flat
+}
+
+type importstatsplot struct{}
+
+func (cmd *importstatsplot) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
+       err := cmd.Plot(stdin, stdout)
+       if err != nil {
+               log.Errorf("%s", err)
+               return 1
+       }
+       return 0
+}
+
+func (cmd *importstatsplot) Plot(stdin io.Reader, stdout io.Writer) error {
+       var stats []importStats
+       err := json.NewDecoder(stdin).Decode(&stats)
+       if err != nil {
+               return err
+       }
+
+       p, err := plot.New()
+       if err != nil {
+               return err
+       }
+       p.Title.Text = "coverage preserved by import (excl X<0.65)"
+       p.X.Label.Text = "input base calls ÷ sequence length"
+       p.Y.Label.Text = "output base calls ÷ input base calls"
+       p.Add(plotter.NewGrid())
+
+       data := map[string]plotter.XYs{}
+       for _, stat := range stats {
+               data[stat.InputLabel] = append(data[stat.InputLabel], plotter.XY{
+                       X: float64(stat.InputCoverage) / float64(stat.InputLength),
+                       Y: float64(stat.TileCoverage) / float64(stat.InputCoverage),
+               })
+       }
+
+       nextInPalette := 0
+       for label, xys := range data {
+               s, err := plotter.NewScatter(xys)
+               if err != nil {
+                       return err
+               }
+               s.GlyphStyle.Color = palette.Plan9[nextInPalette%len(palette.Plan9)]
+               s.GlyphStyle.Radius = vg.Millimeter / 2
+               s.GlyphStyle.Shape = draw.CrossGlyph{}
+               nextInPalette += 7
+               p.Add(s)
+               if false {
+                       p.Legend.Add(label, s)
+               }
+       }
+       p.X.Min = 0.65
+       p.X.Max = 1
+
+       w, err := p.WriterTo(8*vg.Inch, 6*vg.Inch, "svg")
+       if err != nil {
+               return err
+       }
+       _, err = w.WriteTo(stdout)
+       return err
+}