X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/bbae3808c62fe7b40893c1eb0910bd1757192802..f5d3fd181c83be196302b84ca32f0c1433778a4d:/export.go diff --git a/export.go b/export.go index 739f8efa12..710fd324eb 100644 --- a/export.go +++ b/export.go @@ -1,4 +1,4 @@ -package main +package lightning import ( "bufio" @@ -11,6 +11,7 @@ import ( "net/http" _ "net/http/pprof" "os" + "path" "sort" "strings" "sync" @@ -28,11 +29,13 @@ type outputFormat struct { var ( outputFormats = map[string]outputFormat{ - "hgvs": outputFormatHGVS, - "vcf": outputFormatVCF, + "hgvs-onehot": outputFormatHGVSOneHot, + "hgvs": outputFormatHGVS, + "vcf": outputFormatVCF, } - outputFormatHGVS = outputFormat{Print: printHGVS} - outputFormatVCF = outputFormat{Print: printVCF, PadLeft: true} + outputFormatHGVS = outputFormat{Print: printHGVS} + outputFormatHGVSOneHot = outputFormat{Print: printHGVSOneHot} + outputFormatVCF = outputFormat{Print: printVCF, PadLeft: true} ) type exporter struct { @@ -57,7 +60,7 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std 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") + labelsFilename := flags.String("output-labels", "", "also output genome labels csv `file`") err = flags.Parse(args) if err == flag.ErrHelp { err = nil @@ -88,8 +91,8 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std Name: "lightning export", Client: arvados.NewClientFromEnv(), ProjectUUID: *projectUUID, - RAM: 240000000000, - VCPUs: 32, + RAM: 1600000000000, + VCPUs: 64, Priority: *priority, } err = runner.TranslatePaths(inputFilename) @@ -103,7 +106,14 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std } *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"} + runner.Args = []string{"export", "-local=true", + "-ref", *refname, + "-output-format", *outputFormatStr, + "-output-bed", *outputBed, + "-output-labels", "/mnt/output/labels.csv", + "-i", *inputFilename, + "-o", "/mnt/output/export.csv", + } var output string output, err = runner.Run() if err != nil { @@ -113,11 +123,16 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std return 0 } - input, err := os.Open(*inputFilename) + in, err := open(*inputFilename) if err != nil { return 1 } - defer input.Close() + defer in.Close() + input, ok := in.(io.ReadSeeker) + if !ok { + err = fmt.Errorf("%s: %T cannot seek", *inputFilename, in) + return 1 + } // Error out early if seeking doesn't work on the input file. _, err = input.Seek(0, io.SeekEnd) @@ -129,25 +144,15 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std return 1 } - var mtx sync.Mutex var cgs []CompactGenome - tilelib := tileLibrary{ - includeNoCalls: true, + tilelib := &tileLibrary{ + retainNoCalls: true, + compactGenomes: map[string][]tileVariantID{}, } - err = tilelib.LoadGob(context.Background(), input, func(cg CompactGenome) { - if *pick != "" && *pick != cg.Name { - return - } - log.Debugf("export: pick %q", cg.Name) - mtx.Lock() - defer mtx.Unlock() - cgs = append(cgs, cg) - }) + err = tilelib.LoadGob(context.Background(), input, strings.HasSuffix(*inputFilename, ".gz"), nil) if err != nil { return 1 } - sort.Slice(cgs, func(i, j int) bool { return cgs[i].Name < cgs[j].Name }) - log.Printf("export: pick %q => %d genomes", *pick, len(cgs)) refseq, ok := tilelib.refseqs[*refname] if !ok { @@ -160,6 +165,33 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std return 1 } + names := cgnames(tilelib) + for _, name := range names { + cgs = append(cgs, CompactGenome{Name: name, Variants: tilelib.compactGenomes[name]}) + } + if *labelsFilename != "" { + log.Infof("writing labels to %s", *labelsFilename) + var f *os.File + f, err = os.OpenFile(*labelsFilename, os.O_CREATE|os.O_WRONLY, 0777) + if err != nil { + return 1 + } + defer f.Close() + _, outBasename := path.Split(*outputFilename) + for i, name := range names { + _, err = fmt.Fprintf(f, "%d,%q,%q\n", i, trimFilenameForLabel(name), outBasename) + if err != nil { + err = fmt.Errorf("write %s: %w", *labelsFilename, err) + return 1 + } + } + err = f.Close() + if err != nil { + err = fmt.Errorf("close %s: %w", *labelsFilename, err) + return 1 + } + } + _, err = input.Seek(0, io.SeekStart) if err != nil { return 1 @@ -188,7 +220,7 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std bedbufw = bufio.NewWriter(bedout) } - err = cmd.export(bufw, bedout, input, tilelib.taglib.keylen, refseq, cgs) + err = cmd.export(bufw, bedout, input, strings.HasSuffix(*inputFilename, ".gz"), tilelib, refseq, cgs) if err != nil { return 1 } @@ -210,14 +242,14 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std return 1 } } - err = input.Close() + err = in.Close() if err != nil { return 1 } return 0 } -func (cmd *exporter) export(out, bedout 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, tilelib *tileLibrary, refseq map[string][]tileLibRef, cgs []CompactGenome) error { need := map[tileLibRef]bool{} var seqnames []string for seqname, librefs := range refseq { @@ -242,9 +274,9 @@ func (cmd *exporter) export(out, bedout io.Writer, librdr io.Reader, taglen int, 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} + libref := tilelib.getRef(tv.Tag, tv.Sequence) if need[libref] { tileVariant[libref] = tv } @@ -271,46 +303,66 @@ func (cmd *exporter) export(out, bedout io.Writer, librdr io.Reader, taglen int, return fmt.Errorf("%d needed tiles are missing from library", len(missing)) } - log.Infof("assembling %d sequences concurrently", len(seqnames)) - var wg sync.WaitGroup + if true { + // low memory mode + for _, seqname := range seqnames { + log.Infof("assembling %q", seqname) + cmd.exportSeq(out, bedout, tilelib.taglib.keylen, seqname, refseq[seqname], tileVariant, cgs) + log.Infof("assembled %q", seqname) + } + return nil + } + 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()) - }() + ready := make([]chan struct{}, len(seqnames)) + for i := range ready { + ready[i] = make(chan struct{}) } - wg.Wait() - wg.Add(1) + var output sync.WaitGroup + output.Add(1) go func() { - defer wg.Done() + defer output.Done() for i, seqname := range seqnames { + <-ready[i] log.Infof("writing outbuf %s", seqname) io.Copy(out, &outbuf[i]) + log.Infof("writing outbuf %s done", seqname) } }() - if bedout != nil { - wg.Add(1) - go func() { - defer wg.Done() + output.Add(1) + go func() { + defer output.Done() + if bedout != nil { for i, seqname := range seqnames { + <-ready[i] log.Infof("writing bedbuf %s", seqname) io.Copy(bedout, &bedbuf[i]) + log.Infof("writing bedbuf %s done", seqname) } + } + }() + + throttle := throttle{Max: 8} + log.Infof("assembling %d sequences in %d goroutines", len(seqnames), throttle.Max) + for seqidx, seqname := range seqnames { + seqidx, seqname := seqidx, seqname + outbuf := &outbuf[seqidx] + bedbuf := &bedbuf[seqidx] + if bedout == nil { + bedbuf = nil + } + throttle.Acquire() + go func() { + defer throttle.Release() + defer close(ready[seqidx]) + cmd.exportSeq(outbuf, bedbuf, tilelib.taglib.keylen, seqname, refseq[seqname], tileVariant, cgs) + log.Infof("assembled %q to outbuf %d bedbuf %d", seqname, outbuf.Len(), bedbuf.Len()) }() } - wg.Wait() + + output.Wait() return nil } @@ -332,10 +384,16 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string, continue } tagcoverage++ - genometile := tileVariant[tileLibRef{Tag: libref.Tag, Variant: variant}] 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 @@ -355,7 +413,6 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string, 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) @@ -399,8 +456,13 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string, thickstart = 0 } thickend := refpos + // coverage score, 0 to 1000 - score := 1000 * tagcoverage / len(cgs) / 2 + score := 1000 + if len(cgs) > 0 { + score = 1000 * tagcoverage / len(cgs) / 2 + } + fmt.Fprintf(bedw, "%s %d %d %d %d . %d %d\n", seqname, tilestart, tileend, libref.Tag, @@ -466,3 +528,31 @@ func printHGVS(out io.Writer, seqname string, varslice []hgvs.Variant) { } out.Write([]byte{'\n'}) } + +func printHGVSOneHot(out io.Writer, seqname string, varslice []hgvs.Variant) { + vars := map[hgvs.Variant]bool{} + for _, v := range varslice { + if v.Ref != v.New { + vars[v] = true + } + } + + // sort variants to ensure output is deterministic + sorted := make([]hgvs.Variant, 0, len(vars)) + for v := range vars { + sorted = append(sorted, v) + } + sort.Slice(sorted, func(a, b int) bool { return hgvs.Less(sorted[a], sorted[b]) }) + + 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 { + out.Write([]byte("\t1")) + } else { + out.Write([]byte("\t0")) + } + } + out.Write([]byte{'\n'}) + } +}