X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/52c08db40bdfc267f83dadc9df78d1f77a427ccd..f5d3fd181c83be196302b84ca32f0c1433778a4d:/export.go

diff --git a/export.go b/export.go
index 2c51ff3eb9..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:         128000000000,
-			VCPUs:       2,
+			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,101 +303,173 @@ 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))
 	}
 
+	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))
+	ready := make([]chan struct{}, len(seqnames))
+	for i := range ready {
+		ready[i] = make(chan struct{})
+	}
+
+	var output sync.WaitGroup
+	output.Add(1)
+	go func() {
+		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)
+		}
+	}()
+	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())
+		}()
+	}
+
+	output.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]
-			coverage := int64(0) // number of ref bases that are called in genomes -- max is len(reftile.Sequence)*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
-					}
-					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()
 					}
-					// (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()
-						}
-						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
+					varslice := variantAt[v.Position]
+					if varslice == nil {
+						varslice = make([]hgvs.Variant, len(cgs)*2)
+						variantAt[v.Position] = varslice
 					}
-					coverage += int64(len(reftile.Sequence))
+					varslice[cgidx*2+phase] = v
 				}
 			}
-			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)
-				}
+		}
+		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(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)
 			}
-			if bedout != 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 * coverage / int64(len(reftile.Sequence)) / int64(len(cgs)) / 2
-				fmt.Fprintf(bedout, "%s %d %d %d %d . %d %d\n",
-					seqname, tilestart, tileend,
-					libref.Tag,
-					score,
-					thickstart, thickend)
+			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
+			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,
+				score,
+				thickstart, thickend)
 		}
 	}
-	return nil
 }
 
 func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
@@ -424,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'})
+	}
+}