Gzip gob files.
[lightning.git] / export.go
index 6fde64f50800ecf67fb0a39a70ff1dd29163c010..b6b160ddb79311fe3a5964b25e2ca0efe482abdc 100644 (file)
--- a/export.go
+++ b/export.go
@@ -56,6 +56,7 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
        inputFilename := flags.String("i", "-", "input `file` (library)")
        outputFilename := flags.String("o", "-", "output `file`")
        outputFormatStr := flags.String("output-format", "hgvs", "output `format`: hgvs or vcf")
+       outputBed := flags.String("output-bed", "", "also output bed `file`")
        pick := flags.String("pick", "", "`name` of single genome to export")
        err = flags.Parse(args)
        if err == flag.ErrHelp {
@@ -87,15 +88,22 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
                        Name:        "lightning export",
                        Client:      arvados.NewClientFromEnv(),
                        ProjectUUID: *projectUUID,
-                       RAM:         128000000000,
-                       VCPUs:       2,
+                       RAM:         240000000000,
+                       VCPUs:       32,
                        Priority:    *priority,
                }
                err = runner.TranslatePaths(inputFilename)
                if err != nil {
                        return 1
                }
-               runner.Args = []string{"export", "-local=true", "-pick", *pick, "-ref", *refname, "-output-format", *outputFormatStr, "-i", *inputFilename, "-o", "/mnt/output/export.csv"}
+               if *outputBed != "" {
+                       if strings.Contains(*outputBed, "/") {
+                               err = fmt.Errorf("cannot use -output-bed filename %q containing '/' char", *outputBed)
+                               return 1
+                       }
+                       *outputBed = "/mnt/output/" + *outputBed
+               }
+               runner.Args = []string{"export", "-local=true", "-pick", *pick, "-ref", *refname, "-output-format", *outputFormatStr, "-output-bed", *outputBed, "-i", *inputFilename, "-o", "/mnt/output/export.csv"}
                var output string
                output, err = runner.Run()
                if err != nil {
@@ -123,8 +131,10 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
 
        var mtx sync.Mutex
        var cgs []CompactGenome
-       var tilelib tileLibrary
-       err = tilelib.LoadGob(context.Background(), input, func(cg CompactGenome) {
+       tilelib := tileLibrary{
+               retainNoCalls: true,
+       }
+       err = tilelib.LoadGob(context.Background(), input, strings.HasSuffix(*inputFilename, ".gz"), func(cg CompactGenome) {
                if *pick != "" && *pick != cg.Name {
                        return
                }
@@ -159,14 +169,26 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
        if *outputFilename == "-" {
                output = nopCloser{stdout}
        } else {
-               output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
+               output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0666)
                if err != nil {
                        return 1
                }
                defer output.Close()
        }
        bufw := bufio.NewWriter(output)
-       err = cmd.export(bufw, input, tilelib.taglib.keylen, refseq, cgs)
+
+       var bedout *os.File
+       var bedbufw *bufio.Writer
+       if *outputBed != "" {
+               bedout, err = os.OpenFile(*outputBed, os.O_CREATE|os.O_WRONLY, 0666)
+               if err != nil {
+                       return 1
+               }
+               defer bedout.Close()
+               bedbufw = bufio.NewWriter(bedout)
+       }
+
+       err = cmd.export(bufw, bedout, input, strings.HasSuffix(*inputFilename, ".gz"), tilelib.taglib.keylen, refseq, cgs)
        if err != nil {
                return 1
        }
@@ -178,6 +200,16 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
        if err != nil {
                return 1
        }
+       if bedout != nil {
+               err = bedbufw.Flush()
+               if err != nil {
+                       return 1
+               }
+               err = bedout.Close()
+               if err != nil {
+                       return 1
+               }
+       }
        err = input.Close()
        if err != nil {
                return 1
@@ -185,7 +217,7 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
        return 0
 }
 
-func (cmd *exporter) export(out io.Writer, librdr io.Reader, taglen int, refseq map[string][]tileLibRef, cgs []CompactGenome) error {
+func (cmd *exporter) export(out, bedout io.Writer, librdr io.Reader, gz bool, taglen int, refseq map[string][]tileLibRef, cgs []CompactGenome) error {
        need := map[tileLibRef]bool{}
        var seqnames []string
        for seqname, librefs := range refseq {
@@ -210,7 +242,7 @@ func (cmd *exporter) export(out io.Writer, librdr io.Reader, taglen int, refseq
 
        log.Infof("export: loading %d tile variants", len(need))
        tileVariant := map[tileLibRef]TileVariant{}
-       err := DecodeLibrary(librdr, func(ent *LibraryEntry) error {
+       err := DecodeLibrary(librdr, gz, func(ent *LibraryEntry) error {
                for _, tv := range ent.TileVariants {
                        libref := tileLibRef{Tag: tv.Tag, Variant: tv.Variant}
                        if need[libref] {
@@ -239,79 +271,149 @@ func (cmd *exporter) export(out io.Writer, librdr io.Reader, taglen int, refseq
                return fmt.Errorf("%d needed tiles are missing from library", len(missing))
        }
 
+       log.Infof("assembling %d sequences concurrently", len(seqnames))
+       var wg sync.WaitGroup
+       outbuf := make([]bytes.Buffer, len(seqnames))
+       bedbuf := make([]bytes.Buffer, len(seqnames))
+       for seqidx, seqname := range seqnames {
+               seqname := seqname
+               outbuf := &outbuf[seqidx]
+               bedbuf := &bedbuf[seqidx]
+               if bedout == nil {
+                       bedbuf = nil
+               }
+               // TODO: limit number of goroutines and unflushed bufs to ncpus
+               wg.Add(1)
+               go func() {
+                       defer wg.Done()
+                       cmd.exportSeq(outbuf, bedbuf, taglen, seqname, refseq[seqname], tileVariant, cgs)
+                       log.Infof("assembled %q to outbuf %d bedbuf %d", seqname, outbuf.Len(), bedbuf.Len())
+               }()
+       }
+       wg.Wait()
+
+       wg.Add(1)
+       go func() {
+               defer wg.Done()
+               for i, seqname := range seqnames {
+                       log.Infof("writing outbuf %s", seqname)
+                       io.Copy(out, &outbuf[i])
+               }
+       }()
+       if bedout != nil {
+               wg.Add(1)
+               go func() {
+                       defer wg.Done()
+                       for i, seqname := range seqnames {
+                               log.Infof("writing bedbuf %s", seqname)
+                               io.Copy(bedout, &bedbuf[i])
+                       }
+               }()
+       }
+       wg.Wait()
+       return nil
+}
+
+// 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, tileVariant map[tileLibRef]TileVariant, cgs []CompactGenome) {
        refpos := 0
-       for _, seqname := range seqnames {
-               variantAt := map[int][]hgvs.Variant{} // variantAt[chromOffset][genomeIndex*2+phase]
-               for refstep, libref := range refseq[seqname] {
-                       reftile := tileVariant[libref]
-                       for cgidx, cg := range cgs {
-                               for phase := 0; phase < 2; phase++ {
-                                       if len(cg.Variants) <= int(libref.Tag)*2+phase {
-                                               continue
-                                       }
-                                       variant := cg.Variants[int(libref.Tag)*2+phase]
-                                       if variant == 0 {
-                                               continue
-                                       }
-                                       genometile := tileVariant[tileLibRef{Tag: libref.Tag, Variant: variant}]
-                                       if variant == libref.Variant {
-                                               continue
+       variantAt := map[int][]hgvs.Variant{} // variantAt[chromOffset][genomeIndex*2+phase]
+       for refstep, libref := range reftiles {
+               reftile := tileVariant[libref]
+               tagcoverage := 0 // number of times the start tag was found in genomes -- max is len(cgs)*2
+               for cgidx, cg := range cgs {
+                       for phase := 0; phase < 2; phase++ {
+                               if len(cg.Variants) <= int(libref.Tag)*2+phase {
+                                       continue
+                               }
+                               variant := cg.Variants[int(libref.Tag)*2+phase]
+                               if variant == 0 {
+                                       continue
+                               }
+                               tagcoverage++
+                               if variant == libref.Variant {
+                                       continue
+                               }
+                               genometile := tileVariant[tileLibRef{Tag: libref.Tag, Variant: variant}]
+                               if len(genometile.Sequence) == 0 {
+                                       // Hash is known but sequence
+                                       // is not, e.g., retainNoCalls
+                                       // was false during import
+                                       continue
+                               }
+                               refSequence := reftile.Sequence
+                               // If needed, extend the reference
+                               // sequence up to the tag at the end
+                               // of the genometile sequence.
+                               refstepend := refstep + 1
+                               for refstepend < len(reftiles) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genometile.Sequence[len(genometile.Sequence)-taglen:]) {
+                                       if &refSequence[0] == &reftile.Sequence[0] {
+                                               refSequence = append([]byte(nil), refSequence...)
                                        }
-                                       refSequence := reftile.Sequence
-                                       // If needed, extend the
-                                       // reference sequence up to
-                                       // the tag at the end of the
-                                       // genometile sequence.
-                                       refstepend := refstep + 1
-                                       for refstepend < len(refseq[seqname]) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genometile.Sequence[len(genometile.Sequence)-taglen:]) {
-                                               if &refSequence[0] == &reftile.Sequence[0] {
-                                                       refSequence = append([]byte(nil), refSequence...)
-                                               }
-                                               refSequence = append(refSequence, tileVariant[refseq[seqname][refstepend]].Sequence...)
-                                               refstepend++
+                                       refSequence = append(refSequence, tileVariant[reftiles[refstepend]].Sequence...)
+                                       refstepend++
+                               }
+                               // (TODO: handle no-calls)
+                               vars, _ := hgvs.Diff(strings.ToUpper(string(refSequence)), strings.ToUpper(string(genometile.Sequence)), time.Second)
+                               for _, v := range vars {
+                                       if cmd.outputFormat.PadLeft {
+                                               v = v.PadLeft()
                                        }
-                                       vars, _ := hgvs.Diff(strings.ToUpper(string(refSequence)), strings.ToUpper(string(genometile.Sequence)), time.Second)
-                                       for _, v := range vars {
-                                               if cmd.outputFormat.PadLeft {
-                                                       v = v.PadLeft()
-                                               }
-                                               v.Position += refpos
-                                               log.Debugf("%s seq %s phase %d tag %d tile diff %s\n", cg.Name, seqname, phase, libref.Tag, v.String())
-                                               varslice := variantAt[v.Position]
-                                               if varslice == nil {
-                                                       varslice = make([]hgvs.Variant, len(cgs)*2)
-                                                       variantAt[v.Position] = varslice
-                                               }
-                                               varslice[cgidx*2+phase] = v
+                                       v.Position += refpos
+                                       log.Debugf("%s seq %s phase %d tag %d tile diff %s\n", cg.Name, seqname, phase, libref.Tag, v.String())
+                                       varslice := variantAt[v.Position]
+                                       if varslice == nil {
+                                               varslice = make([]hgvs.Variant, len(cgs)*2)
+                                               variantAt[v.Position] = varslice
                                        }
+                                       varslice[cgidx*2+phase] = v
                                }
                        }
-                       refpos += len(reftile.Sequence) - taglen
+               }
+               refpos += len(reftile.Sequence) - taglen
 
-                       // Flush entries from variantAt that are
-                       // behind refpos. Flush all entries if this is
-                       // the last reftile of the path/chromosome.
-                       var flushpos []int
-                       lastrefstep := refstep == len(refseq[seqname])-1
-                       for pos := range variantAt {
-                               if lastrefstep || pos <= refpos {
-                                       flushpos = append(flushpos, pos)
-                               }
+               // Flush entries from variantAt that are behind
+               // refpos. Flush all entries if this is the last
+               // reftile of the path/chromosome.
+               var flushpos []int
+               lastrefstep := refstep == len(reftiles)-1
+               for pos := range variantAt {
+                       if lastrefstep || pos <= refpos {
+                               flushpos = append(flushpos, pos)
                        }
-                       sort.Slice(flushpos, func(i, j int) bool { return flushpos[i] < flushpos[j] })
-                       for _, pos := range flushpos {
-                               varslice := variantAt[pos]
-                               delete(variantAt, pos)
-                               for i := range varslice {
-                                       if varslice[i].Position == 0 {
-                                               varslice[i].Position = pos
-                                       }
+               }
+               sort.Slice(flushpos, func(i, j int) bool { return flushpos[i] < flushpos[j] })
+               for _, pos := range flushpos {
+                       varslice := variantAt[pos]
+                       delete(variantAt, pos)
+                       for i := range varslice {
+                               if varslice[i].Position == 0 {
+                                       varslice[i].Position = pos
                                }
-                               cmd.outputFormat.Print(out, seqname, varslice)
                        }
+                       cmd.outputFormat.Print(outw, seqname, varslice)
+               }
+               if bedw != nil && len(reftile.Sequence) > 0 {
+                       tilestart := refpos - len(reftile.Sequence) + taglen
+                       tileend := refpos
+                       if !lastrefstep {
+                               tileend += taglen
+                       }
+                       thickstart := tilestart + taglen
+                       if refstep == 0 {
+                               thickstart = 0
+                       }
+                       thickend := refpos
+                       // coverage score, 0 to 1000
+                       score := 1000 * tagcoverage / len(cgs) / 2
+                       fmt.Fprintf(bedw, "%s %d %d %d %d . %d %d\n",
+                               seqname, tilestart, tileend,
+                               libref.Tag,
+                               score,
+                               thickstart, thickend)
                }
        }
-       return nil
 }
 
 func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
@@ -337,8 +439,8 @@ func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
                        alts[a] = i + 1
                }
                fmt.Fprintf(out, "%s\t%d\t%s\t%s", seqname, varslice[0].Position, ref, strings.Join(altslice, ","))
-               for i := 0; i < len(varslice)/2; i++ {
-                       v1, v2 := varslice[i*2], varslice[i*2+1]
+               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 {
                                a1 = 0