X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/3144aedf02bc1c241ac0666bf01905814eeed71a..3364ff0d5ecd6a0f820281625d16ef13f8521160:/export.go

diff --git a/export.go b/export.go
index a37ec725bf..b8590b973f 100644
--- a/export.go
+++ b/export.go
@@ -1,4 +1,4 @@
-package main
+package lightning
 
 import (
 	"bufio"
@@ -11,6 +11,8 @@ import (
 	"net/http"
 	_ "net/http/pprof"
 	"os"
+	"path/filepath"
+	"runtime"
 	"sort"
 	"strings"
 	"sync"
@@ -22,21 +24,26 @@ import (
 )
 
 type outputFormat struct {
-	Print   func(out io.Writer, seqname string, varslice []hgvs.Variant)
-	PadLeft bool
+	Filename string
+	Print    func(out io.Writer, seqname string, varslice []hgvs.Variant)
+	PadLeft  bool
 }
 
 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{Filename: "out.csv", Print: printHGVS}
+	outputFormatHGVSOneHot = outputFormat{Filename: "out.csv", Print: printHGVSOneHot}
+	outputFormatVCF        = outputFormat{Filename: "out.vcf", Print: printVCF, PadLeft: true}
 )
 
 type exporter struct {
-	outputFormat outputFormat
+	outputFormat   outputFormat
+	outputPerChrom bool
+	maxTileSize    int
 }
 
 func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
@@ -49,15 +56,18 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
 	flags := flag.NewFlagSet("", flag.ContinueOnError)
 	flags.SetOutput(stderr)
 	pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
+	pprofdir := flags.String("pprof-dir", "", "write Go profile data to `directory` periodically")
 	runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
 	projectUUID := flags.String("project", "", "project `UUID` for output data")
 	priority := flags.Int("priority", 500, "container request priority")
 	refname := flags.String("ref", "", "reference genome `name`")
-	inputFilename := flags.String("i", "-", "input `file` (library)")
-	outputFilename := flags.String("o", "-", "output `file`")
+	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")
 	outputBed := flags.String("output-bed", "", "also output bed `file`")
-	pick := flags.String("pick", "", "`name` of single genome to export")
+	flags.BoolVar(&cmd.outputPerChrom, "output-per-chromosome", true, "output one file per chromosome")
+	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)
 	if err == flag.ErrHelp {
 		err = nil
@@ -65,6 +75,10 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
 	} else if err != nil {
 		return 2
 	}
+	if flag.NArg() > 0 {
+		err = fmt.Errorf("extra unparsed command line arguments: %q", flag.Args())
+		return 2
+	}
 
 	if f, ok := outputFormats[*outputFormatStr]; !ok {
 		err = fmt.Errorf("invalid output format %q", *outputFormatStr)
@@ -78,21 +92,24 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
 			log.Println(http.ListenAndServe(*pprof, nil))
 		}()
 	}
+	if *pprofdir != "" {
+		go writeProfilesPeriodically(*pprofdir)
+	}
 
 	if !*runlocal {
-		if *outputFilename != "-" {
-			err = errors.New("cannot specify output file in container mode: not implemented")
+		if *outputDir != "." {
+			err = errors.New("cannot specify output directory in container mode: not implemented")
 			return 1
 		}
 		runner := arvadosContainerRunner{
 			Name:        "lightning export",
 			Client:      arvados.NewClientFromEnv(),
 			ProjectUUID: *projectUUID,
-			RAM:         240000000000,
-			VCPUs:       32,
+			RAM:         700000000000,
+			VCPUs:       64,
 			Priority:    *priority,
 		}
-		err = runner.TranslatePaths(inputFilename)
+		err = runner.TranslatePaths(inputDir)
 		if err != nil {
 			return 1
 		}
@@ -103,7 +120,18 @@ 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",
+			"-pprof", ":6000",
+			"-pprof-dir", "/mnt/output",
+			"-ref", *refname,
+			"-output-format", *outputFormatStr,
+			"-output-bed", *outputBed,
+			"-output-labels", "/mnt/output/labels.csv",
+			"-output-per-chromosome=" + fmt.Sprintf("%v", cmd.outputPerChrom),
+			"-max-tile-size", fmt.Sprintf("%d", cmd.maxTileSize),
+			"-input-dir", *inputDir,
+			"-output-dir", "/mnt/output",
+		}
 		var output string
 		output, err = runner.Run()
 		if err != nil {
@@ -113,41 +141,16 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
 		return 0
 	}
 
-	input, err := os.Open(*inputFilename)
-	if err != nil {
-		return 1
-	}
-	defer input.Close()
-
-	// Error out early if seeking doesn't work on the input file.
-	_, err = input.Seek(0, io.SeekEnd)
-	if err != nil {
-		return 1
-	}
-	_, err = input.Seek(0, io.SeekStart)
-	if err != nil {
-		return 1
-	}
-
-	var mtx sync.Mutex
 	var cgs []CompactGenome
-	tilelib := tileLibrary{
-		includeNoCalls: true,
+	tilelib := &tileLibrary{
+		retainNoCalls:       true,
+		retainTileSequences: 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.LoadDir(context.Background(), *inputDir, 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,22 +163,31 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
 		return 1
 	}
 
-	_, err = input.Seek(0, io.SeekStart)
-	if err != nil {
-		return 1
+	names := cgnames(tilelib)
+	for _, name := range names {
+		cgs = append(cgs, CompactGenome{Name: name, Variants: tilelib.compactGenomes[name]})
 	}
-
-	var output io.WriteCloser
-	if *outputFilename == "-" {
-		output = nopCloser{stdout}
-	} else {
-		output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0666)
+	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()
+		for i, name := range names {
+			_, err = fmt.Fprintf(f, "%d,%q,%q\n", i, trimFilenameForLabel(name), cmd.outputFormat.Filename)
+			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
 		}
-		defer output.Close()
 	}
-	bufw := bufio.NewWriter(output)
 
 	var bedout *os.File
 	var bedbufw *bufio.Writer
@@ -188,15 +200,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)
-	if err != nil {
-		return 1
-	}
-	err = bufw.Flush()
-	if err != nil {
-		return 1
-	}
-	err = output.Close()
+	err = cmd.export(*outputDir, bedout, tilelib, refseq, cgs)
 	if err != nil {
 		return 1
 	}
@@ -210,58 +214,22 @@ func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, std
 			return 1
 		}
 	}
-	err = input.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 {
-	need := map[tileLibRef]bool{}
+func (cmd *exporter) export(outdir string, bedout io.Writer, tilelib *tileLibrary, refseq map[string][]tileLibRef, cgs []CompactGenome) error {
 	var seqnames []string
+	var missing []tileLibRef
 	for seqname, librefs := range refseq {
 		seqnames = append(seqnames, seqname)
 		for _, libref := range librefs {
-			if libref.Variant != 0 {
-				need[libref] = true
+			if libref.Variant != 0 && tilelib.TileVariantSequence(libref) == nil {
+				missing = append(missing, libref)
 			}
 		}
 	}
 	sort.Strings(seqnames)
 
-	for _, cg := range cgs {
-		for i, variant := range cg.Variants {
-			if variant == 0 {
-				continue
-			}
-			libref := tileLibRef{Tag: tagID(i / 2), Variant: variant}
-			need[libref] = true
-		}
-	}
-
-	log.Infof("export: loading %d tile variants", len(need))
-	tileVariant := map[tileLibRef]TileVariant{}
-	err := DecodeLibrary(librdr, func(ent *LibraryEntry) error {
-		for _, tv := range ent.TileVariants {
-			libref := tileLibRef{Tag: tv.Tag, Variant: tv.Variant}
-			if need[libref] {
-				tileVariant[libref] = tv
-			}
-		}
-		return nil
-	})
-	if err != nil {
-		return err
-	}
-
-	log.Infof("export: loaded %d tile variants", len(tileVariant))
-	var missing []tileLibRef
-	for libref := range need {
-		if _, ok := tileVariant[libref]; !ok {
-			missing = append(missing, libref)
-		}
-	}
 	if len(missing) > 0 {
 		if limit := 100; len(missing) > limit {
 			log.Warnf("first %d missing tiles: %v", limit, missing[:limit])
@@ -271,56 +239,85 @@ 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
-	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
+	outw := make([]io.WriteCloser, len(seqnames))
+	bedw := make([]io.WriteCloser, len(seqnames))
+
+	var merges sync.WaitGroup
+	merge := func(dst io.Writer, src []io.WriteCloser, label string) {
+		var mtx sync.Mutex
+		for i, seqname := range seqnames {
+			pr, pw := io.Pipe()
+			src[i] = pw
+			merges.Add(1)
+			seqname := seqname
+			go func() {
+				defer merges.Done()
+				log.Infof("writing %s %s", seqname, label)
+				scanner := bufio.NewScanner(pr)
+				for scanner.Scan() {
+					mtx.Lock()
+					dst.Write(scanner.Bytes())
+					dst.Write([]byte{'\n'})
+					mtx.Unlock()
+				}
+				log.Infof("writing %s %s done", seqname, label)
+			}()
 		}
-		// 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()
+	if cmd.outputPerChrom {
 		for i, seqname := range seqnames {
-			log.Infof("writing outbuf %s", seqname)
-			io.Copy(out, &outbuf[i])
+			f, err := os.OpenFile(filepath.Join(outdir, strings.Replace(cmd.outputFormat.Filename, ".", "."+seqname+".", 1)), os.O_CREATE|os.O_WRONLY, 0666)
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			log.Infof("writing %q", f.Name())
+			outw[i] = f
 		}
-	}()
+	} 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 err != nil {
+			return err
+		}
+		defer out.Close()
+		merge(out, outw, "output")
+	}
 	if bedout != nil {
-		wg.Add(1)
+		merge(bedout, bedw, "bed")
+	}
+
+	throttle := throttle{Max: runtime.NumCPU()}
+	log.Infof("assembling %d sequences in %d goroutines", len(seqnames), throttle.Max)
+	for seqidx, seqname := range seqnames {
+		seqidx, seqname := seqidx, seqname
+		outw := outw[seqidx]
+		bedw := bedw[seqidx]
+		throttle.Acquire()
 		go func() {
-			defer wg.Done()
-			for i, seqname := range seqnames {
-				log.Infof("writing bedbuf %s", seqname)
-				io.Copy(bedout, &bedbuf[i])
+			defer throttle.Release()
+			if bedw != nil {
+				defer bedw.Close()
 			}
+			defer outw.Close()
+			outwb := bufio.NewWriter(outw)
+			defer outwb.Flush()
+			cmd.exportSeq(outwb, bedw, tilelib.taglib.keylen, seqname, refseq[seqname], tilelib, cgs)
 		}()
 	}
-	wg.Wait()
+
+	merges.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) {
+func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string, reftiles []tileLibRef, tilelib *tileLibrary, cgs []CompactGenome) {
 	refpos := 0
 	variantAt := map[int][]hgvs.Variant{} // variantAt[chromOffset][genomeIndex*2+phase]
 	for refstep, libref := range reftiles {
-		reftile := tileVariant[libref]
+		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 {
 			for phase := 0; phase < 2; phase++ {
@@ -335,27 +332,35 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
 				if variant == libref.Variant {
 					continue
 				}
-				genometile := tileVariant[tileLibRef{Tag: libref.Tag, Variant: variant}]
-				refSequence := reftile.Sequence
+				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 genometile sequence.
+				// of the genomeseq 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] {
+				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, tileVariant[reftiles[refstepend]].Sequence...)
+					refSequence = append(refSequence, tilelib.TileVariantSequence(reftiles[refstepend])...)
 					refstepend++
 				}
 				// (TODO: handle no-calls)
-				vars, _ := hgvs.Diff(strings.ToUpper(string(refSequence)), strings.ToUpper(string(genometile.Sequence)), time.Second)
+				vars, _ := hgvs.Diff(strings.ToUpper(string(refSequence)), strings.ToUpper(string(genomeseq)), 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)
@@ -365,7 +370,7 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
 				}
 			}
 		}
-		refpos += len(reftile.Sequence) - taglen
+		refpos += len(refseq) - taglen
 
 		// Flush entries from variantAt that are behind
 		// refpos. Flush all entries if this is the last
@@ -388,8 +393,8 @@ func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string,
 			}
 			cmd.outputFormat.Print(outw, seqname, varslice)
 		}
-		if bedw != nil && len(reftile.Sequence) > 0 {
-			tilestart := refpos - len(reftile.Sequence) + taglen
+		if bedw != nil && len(refseq) > 0 {
+			tilestart := refpos - len(refseq) + taglen
 			tileend := refpos
 			if !lastrefstep {
 				tileend += taglen
@@ -399,8 +404,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 +476,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'})
+	}
+}