Use tsv filename when using tab separator.
[lightning.git] / export.go
index 8d455b15ac82429dc7221fc20fc8846ce2f6a7eb..5e0ecf43e1889316e5ab3766667850bdaa34ce9e 100644 (file)
--- a/export.go
+++ b/export.go
@@ -14,18 +14,26 @@ import (
        "path/filepath"
        "runtime"
        "sort"
+       "strconv"
        "strings"
        "sync"
        "time"
 
        "git.arvados.org/arvados.git/sdk/go/arvados"
        "github.com/arvados/lightning/hgvs"
+       "github.com/klauspost/pgzip"
        log "github.com/sirupsen/logrus"
 )
 
+type tvVariant struct {
+       hgvs.Variant
+       librefs map[tileLibRef]bool
+}
+
 type outputFormat struct {
        Filename string
-       Print    func(out io.Writer, seqname string, varslice []hgvs.Variant)
+       Head     func(out io.Writer, cgs []CompactGenome)
+       Print    func(out io.Writer, seqname string, varslice []tvVariant)
        PadLeft  bool
 }
 
@@ -33,16 +41,20 @@ var (
        outputFormats = map[string]outputFormat{
                "hgvs-onehot": outputFormatHGVSOneHot,
                "hgvs":        outputFormatHGVS,
+               "pvcf":        outputFormatPVCF,
                "vcf":         outputFormatVCF,
        }
-       outputFormatHGVS       = outputFormat{Filename: "out.csv", Print: printHGVS}
-       outputFormatHGVSOneHot = outputFormat{Filename: "out.csv", Print: printHGVSOneHot}
-       outputFormatVCF        = outputFormat{Filename: "out.vcf", Print: printVCF, PadLeft: true}
+       outputFormatHGVS       = outputFormat{Filename: "out.tsv", Head: headNone, Print: printHGVS}
+       outputFormatHGVSOneHot = outputFormat{Filename: "out.tsv", Head: headNone, Print: printHGVSOneHot}
+       outputFormatPVCF       = outputFormat{Filename: "out.vcf", Head: headPVCF, Print: printPVCF, PadLeft: true}
+       outputFormatVCF        = outputFormat{Filename: "out.vcf", Head: headVCF, Print: printVCF, PadLeft: true}
+       headNone               = func(io.Writer, []CompactGenome) {}
 )
 
 type exporter struct {
        outputFormat   outputFormat
        outputPerChrom bool
+       compress       bool
        maxTileSize    int
 }
 
@@ -63,9 +75,10 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
        refname := flags.String("ref", "", "reference genome `name`")
        inputDir := flags.String("input-dir", ".", "input `directory`")
        outputDir := flags.String("output-dir", ".", "output `directory`")
-       outputFormatStr := flags.String("output-format", "hgvs", "output `format`: hgvs or vcf")
+       outputFormatStr := flags.String("output-format", "hgvs", "output `format`: hgvs, pvcf, or vcf")
        outputBed := flags.String("output-bed", "", "also output bed `file`")
        flags.BoolVar(&cmd.outputPerChrom, "output-per-chromosome", true, "output one file per chromosome")
+       flags.BoolVar(&cmd.compress, "z", false, "write gzip-compressed output files")
        labelsFilename := flags.String("output-labels", "", "also output genome labels csv `file`")
        flags.IntVar(&cmd.maxTileSize, "max-tile-size", 50000, "don't try to make annotations for tiles bigger than given `size`")
        err = flags.Parse(args)
@@ -132,6 +145,7 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
                        "-max-tile-size", fmt.Sprintf("%d", cmd.maxTileSize),
                        "-input-dir", *inputDir,
                        "-output-dir", "/mnt/output",
+                       "-z=" + fmt.Sprintf("%v", cmd.compress),
                }
                var output string
                output, err = runner.Run()
@@ -269,22 +283,42 @@ func (cmd *exporter) export(outdir string, bedout io.Writer, tilelib *tileLibrar
        }
        if cmd.outputPerChrom {
                for i, seqname := range seqnames {
-                       f, err := os.OpenFile(filepath.Join(outdir, strings.Replace(cmd.outputFormat.Filename, ".", "."+seqname+".", 1)), os.O_CREATE|os.O_WRONLY, 0666)
+                       fnm := filepath.Join(outdir, strings.Replace(cmd.outputFormat.Filename, ".", "."+seqname+".", 1))
+                       if cmd.compress {
+                               fnm += ".gz"
+                       }
+                       f, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
                        if err != nil {
                                return err
                        }
                        defer f.Close()
                        log.Infof("writing %q", f.Name())
                        outw[i] = f
+                       if cmd.compress {
+                               z := pgzip.NewWriter(f)
+                               defer z.Close()
+                               outw[i] = z
+                       }
+                       cmd.outputFormat.Head(outw[i], cgs)
                }
        } else {
                fnm := filepath.Join(outdir, cmd.outputFormat.Filename)
-               log.Infof("writing %q", fnm)
-               out, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY, 0666)
+               if cmd.compress {
+                       fnm += ".gz"
+               }
+               f, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
                if err != nil {
                        return err
                }
-               defer out.Close()
+               defer f.Close()
+               log.Infof("writing %q", fnm)
+               var out io.Writer = f
+               if cmd.compress {
+                       z := pgzip.NewWriter(out)
+                       defer z.Close()
+                       out = z
+               }
+               cmd.outputFormat.Head(out, cgs)
                merge(out, outw, "output")
        }
        if bedout != nil {
@@ -303,24 +337,44 @@ func (cmd *exporter) export(outdir string, bedout io.Writer, tilelib *tileLibrar
                        if bedw != nil {
                                defer bedw.Close()
                        }
-                       defer outw.Close()
                        outwb := bufio.NewWriterSize(outw, 8*1024*1024)
-                       defer outwb.Flush()
                        cmd.exportSeq(outwb, bedw, tilelib.taglib.keylen, seqname, refseq[seqname], tilelib, cgs)
+                       err := outwb.Flush()
+                       throttle.Report(err)
+                       err = outw.Close()
+                       throttle.Report(err)
                }()
        }
 
        merges.Wait()
        throttle.Wait()
-       return nil
+       return throttle.Err()
 }
 
 // Align genome tiles to reference tiles, write diffs to outw, and (if
 // bedw is not nil) write tile coverage to bedw.
 func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string, reftiles []tileLibRef, tilelib *tileLibrary, cgs []CompactGenome) {
+       t0 := time.Now()
+       progressbar := time.NewTicker(time.Minute)
+       defer progressbar.Stop()
+       var outmtx sync.Mutex
+       defer outmtx.Lock()
        refpos := 0
-       variantAt := map[int][]hgvs.Variant{} // variantAt[chromOffset][genomeIndex*2+phase]
+       variantAt := map[int][]tvVariant{} // variantAt[chromOffset][genomeIndex*2+phase]
        for refstep, libref := range reftiles {
+               select {
+               case <-progressbar.C:
+                       var eta interface{}
+                       if refstep > 0 {
+                               fin := t0.Add(time.Duration(float64(time.Now().Sub(t0)) * float64(len(reftiles)) / float64(refstep)))
+                               eta = fmt.Sprintf("%v (%v)", fin.Format(time.RFC3339), fin.Sub(time.Now()))
+                       } else {
+                               eta = "N/A"
+                       }
+                       log.Printf("exportSeq: %s: refstep %d of %d, %.0f/s, ETA %v", seqname, refstep, len(reftiles), float64(refstep)/time.Now().Sub(t0).Seconds(), eta)
+               default:
+               }
+               diffs := map[tileLibRef][]hgvs.Variant{}
                refseq := tilelib.TileVariantSequence(libref)
                tagcoverage := 0 // number of times the start tag was found in genomes -- max is len(cgs)*2
                for cgidx, cg := range cgs {
@@ -336,30 +390,40 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
                                if variant == libref.Variant {
                                        continue
                                }
-                               genomeseq := tilelib.TileVariantSequence(tileLibRef{Tag: libref.Tag, Variant: variant})
-                               if len(genomeseq) == 0 {
-                                       // Hash is known but sequence
-                                       // is not, e.g., retainNoCalls
-                                       // was false during import
-                                       continue
-                               }
-                               if len(genomeseq) > cmd.maxTileSize {
-                                       continue
-                               }
-                               refSequence := refseq
-                               // If needed, extend the reference
-                               // sequence up to the tag at the end
-                               // of the genomeseq sequence.
-                               refstepend := refstep + 1
-                               for refstepend < len(reftiles) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genomeseq[len(genomeseq)-taglen:]) && len(refSequence) <= cmd.maxTileSize {
-                                       if &refSequence[0] == &refseq[0] {
-                                               refSequence = append([]byte(nil), refSequence...)
+                               glibref := tileLibRef{Tag: libref.Tag, Variant: variant}
+                               vars, ok := diffs[glibref]
+                               if !ok {
+                                       genomeseq := tilelib.TileVariantSequence(glibref)
+                                       if len(genomeseq) == 0 {
+                                               // Hash is known but sequence
+                                               // is not, e.g., retainNoCalls
+                                               // was false during import
+                                               continue
                                        }
-                                       refSequence = append(refSequence, tilelib.TileVariantSequence(reftiles[refstepend])...)
-                                       refstepend++
+                                       if len(genomeseq) > cmd.maxTileSize {
+                                               continue
+                                       }
+                                       refSequence := refseq
+                                       // If needed, extend the
+                                       // reference sequence up to
+                                       // the tag at the end of the
+                                       // genomeseq sequence.
+                                       refstepend := refstep + 1
+                                       for refstepend < len(reftiles) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genomeseq[len(genomeseq)-taglen:]) && len(refSequence) <= cmd.maxTileSize {
+                                               if &refSequence[0] == &refseq[0] {
+                                                       refSequence = append([]byte(nil), refSequence...)
+                                               }
+                                               refSequence = append(refSequence, tilelib.TileVariantSequence(reftiles[refstepend])...)
+                                               refstepend++
+                                       }
+                                       // (TODO: handle no-calls)
+                                       if len(refSequence) <= cmd.maxTileSize {
+                                               refstr := strings.ToUpper(string(refSequence))
+                                               genomestr := strings.ToUpper(string(genomeseq))
+                                               vars, _ = hgvs.Diff(refstr, genomestr, time.Second)
+                                       }
+                                       diffs[glibref] = vars
                                }
-                               // (TODO: handle no-calls)
-                               vars, _ := hgvs.Diff(strings.ToUpper(string(refSequence)), strings.ToUpper(string(genomeseq)), time.Second)
                                for _, v := range vars {
                                        if cmd.outputFormat.PadLeft {
                                                v = v.PadLeft()
@@ -367,10 +431,15 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
                                        v.Position += refpos
                                        varslice := variantAt[v.Position]
                                        if varslice == nil {
-                                               varslice = make([]hgvs.Variant, len(cgs)*2)
+                                               varslice = make([]tvVariant, len(cgs)*2)
                                                variantAt[v.Position] = varslice
                                        }
-                                       varslice[cgidx*2+phase] = v
+                                       varslice[cgidx*2+phase].Variant = v
+                                       if varslice[cgidx*2+phase].librefs == nil {
+                                               varslice[cgidx*2+phase].librefs = map[tileLibRef]bool{glibref: true}
+                                       } else {
+                                               varslice[cgidx*2+phase].librefs[glibref] = true
+                                       }
                                }
                        }
                }
@@ -379,7 +448,7 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
                // Flush entries from variantAt that are behind
                // refpos. Flush all entries if this is the last
                // reftile of the path/chromosome.
-               var flushpos []int
+               flushpos := make([]int, 0, len(variantAt))
                lastrefstep := refstep == len(reftiles)-1
                for pos := range variantAt {
                        if lastrefstep || pos <= refpos {
@@ -387,7 +456,8 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
                        }
                }
                sort.Slice(flushpos, func(i, j int) bool { return flushpos[i] < flushpos[j] })
-               for _, pos := range flushpos {
+               flushvariants := make([][]tvVariant, len(flushpos))
+               for i, pos := range flushpos {
                        varslice := variantAt[pos]
                        delete(variantAt, pos)
                        for i := range varslice {
@@ -395,8 +465,15 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
                                        varslice[i].Position = pos
                                }
                        }
-                       cmd.outputFormat.Print(outw, seqname, varslice)
+                       flushvariants[i] = varslice
                }
+               outmtx.Lock()
+               go func() {
+                       defer outmtx.Unlock()
+                       for _, varslice := range flushvariants {
+                               cmd.outputFormat.Print(outw, seqname, varslice)
+                       }
+               }()
                if bedw != nil && len(refseq) > 0 {
                        tilestart := refpos - len(refseq) + taglen
                        tileend := refpos
@@ -424,21 +501,57 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
        }
 }
 
-func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
-       refs := map[string]map[string]int{}
+func bucketVarsliceByRef(varslice []tvVariant) map[string]map[string]int {
+       byref := map[string]map[string]int{}
        for _, v := range varslice {
                if v.Ref == "" && v.New == "" {
                        continue
                }
-               alts := refs[v.Ref]
+               alts := byref[v.Ref]
                if alts == nil {
                        alts = map[string]int{}
-                       refs[v.Ref] = alts
+                       byref[v.Ref] = alts
+               }
+               alts[v.New]++
+       }
+       return byref
+}
+
+func headVCF(out io.Writer, cgs []CompactGenome) {
+       fmt.Fprint(out, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n")
+}
+
+func printVCF(out io.Writer, seqname string, varslice []tvVariant) {
+       for ref, alts := range bucketVarsliceByRef(varslice) {
+               altslice := make([]string, 0, len(alts))
+               for alt := range alts {
+                       altslice = append(altslice, alt)
                }
-               alts[v.New] = 0
+               sort.Strings(altslice)
+
+               info := "AC="
+               for i, a := range altslice {
+                       if i > 0 {
+                               info += ","
+                       }
+                       info += strconv.Itoa(alts[a])
+               }
+               fmt.Fprintf(out, "%s\t%d\t.\t%s\t%s\t.\t.\t%s\n", seqname, varslice[0].Position, ref, strings.Join(altslice, ","), info)
        }
-       for ref, alts := range refs {
-               var altslice []string
+}
+
+func headPVCF(out io.Writer, cgs []CompactGenome) {
+       fmt.Fprintln(out, `##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">`)
+       fmt.Fprintf(out, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT")
+       for _, cg := range cgs {
+               fmt.Fprintf(out, "\t%s", cg.Name)
+       }
+       fmt.Fprintf(out, "\n")
+}
+
+func printPVCF(out io.Writer, seqname string, varslice []tvVariant) {
+       for ref, alts := range bucketVarsliceByRef(varslice) {
+               altslice := make([]string, 0, len(alts))
                for alt := range alts {
                        altslice = append(altslice, alt)
                }
@@ -446,11 +559,14 @@ func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
                for i, a := range altslice {
                        alts[a] = i + 1
                }
-               fmt.Fprintf(out, "%s\t%d\t%s\t%s", seqname, varslice[0].Position, ref, strings.Join(altslice, ","))
+               fmt.Fprintf(out, "%s\t%d\t.\t%s\t%s\t.\t.\t.\tGT", seqname, varslice[0].Position, ref, strings.Join(altslice, ","))
                for i := 0; i < len(varslice); i += 2 {
                        v1, v2 := varslice[i], varslice[i+1]
                        a1, a2 := alts[v1.New], alts[v2.New]
                        if v1.Ref != ref {
+                               // variant on allele 0 belongs on a
+                               // different output line -- same
+                               // chr,pos but different "ref" length
                                a1 = 0
                        }
                        if v2.Ref != ref {
@@ -462,13 +578,13 @@ func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
        }
 }
 
-func printHGVS(out io.Writer, seqname string, varslice []hgvs.Variant) {
+func printHGVS(out io.Writer, seqname string, varslice []tvVariant) {
        for i := 0; i < len(varslice)/2; i++ {
                if i > 0 {
                        out.Write([]byte{'\t'})
                }
                var1, var2 := varslice[i*2], varslice[i*2+1]
-               if var1 == var2 {
+               if var1.Variant == var2.Variant {
                        if var1.Ref == var1.New {
                                out.Write([]byte{'.'})
                        } else {
@@ -481,11 +597,11 @@ func printHGVS(out io.Writer, seqname string, varslice []hgvs.Variant) {
        out.Write([]byte{'\n'})
 }
 
-func printHGVSOneHot(out io.Writer, seqname string, varslice []hgvs.Variant) {
+func printHGVSOneHot(out io.Writer, seqname string, varslice []tvVariant) {
        vars := map[hgvs.Variant]bool{}
        for _, v := range varslice {
                if v.Ref != v.New {
-                       vars[v] = true
+                       vars[v.Variant] = true
                }
        }
 
@@ -499,7 +615,7 @@ func printHGVSOneHot(out io.Writer, seqname string, varslice []hgvs.Variant) {
        for _, v := range sorted {
                fmt.Fprintf(out, "%s.%s", seqname, v.String())
                for i := 0; i < len(varslice); i += 2 {
-                       if varslice[i] == v || varslice[i+1] == v {
+                       if varslice[i].Variant == v || varslice[i+1].Variant == v {
                                out.Write([]byte("\t1"))
                        } else {
                                out.Write([]byte("\t0"))