X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/2c21a370ec77fa18ee73ea1a12ad8639025ca837..e4839e500e2da73b476f095425c0347ae7c5de97:/exportnumpy.go

diff --git a/exportnumpy.go b/exportnumpy.go
index 3736ef0184..27c86a3210 100644
--- a/exportnumpy.go
+++ b/exportnumpy.go
@@ -1,7 +1,12 @@
-package main
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package lightning
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"errors"
 	"flag"
@@ -12,10 +17,15 @@ import (
 	_ "net/http/pprof"
 	"os"
 	"sort"
+	"strconv"
 	"strings"
+	"sync"
+	"sync/atomic"
 
 	"git.arvados.org/arvados.git/sdk/go/arvados"
+	"github.com/arvados/lightning/hgvs"
 	"github.com/kshedden/gonpy"
+	"github.com/sirupsen/logrus"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -36,11 +46,15 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
 	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")
-	inputFilename := flags.String("i", "-", "input `file`")
-	outputFilename := flags.String("o", "-", "output `file`")
-	annotationsFilename := flags.String("output-annotations", "", "output `file` for tile variant annotations tsv")
-	librefsFilename := flags.String("output-onehot2tilevar", "", "when using -one-hot, create tsv `file` mapping column# to tag# and variant#")
+	inputDir := flags.String("input-dir", "./in", "input `directory`")
+	outputDir := flags.String("output-dir", "./out", "output `directory`")
+	annotationsFilename := flags.String("output-annotations", "", "output `file` for tile variant annotations csv")
+	librefsFilename := flags.String("output-onehot2tilevar", "", "when using -one-hot, create csv `file` mapping column# to tag# and variant#")
+	labelsFilename := flags.String("output-labels", "", "output `file` for genome labels csv")
+	regionsFilename := flags.String("regions", "", "only output columns/annotations that intersect regions in specified bed `file`")
+	expandRegions := flags.Int("expand-regions", 0, "expand specified regions by `N` base pairs on each side`")
 	onehot := flags.Bool("one-hot", false, "recode tile variants as one-hot")
+	chunks := flags.Int("chunks", 1, "split output into `N` numpy files")
 	cmd.filter.Flags(flags)
 	err = flags.Parse(args)
 	if err == flag.ErrHelp {
@@ -57,32 +71,33 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
 	}
 
 	if !*runlocal {
-		if *outputFilename != "-" {
-			err = errors.New("cannot specify output file in container mode: not implemented")
-			return 1
-		}
 		runner := arvadosContainerRunner{
 			Name:        "lightning export-numpy",
 			Client:      arvados.NewClientFromEnv(),
 			ProjectUUID: *projectUUID,
-			RAM:         128000000000,
-			VCPUs:       32,
+			RAM:         500000000000,
+			VCPUs:       96,
 			Priority:    *priority,
+			KeepCache:   1,
+			APIAccess:   true,
 		}
-		err = runner.TranslatePaths(inputFilename)
+		err = runner.TranslatePaths(inputDir, regionsFilename)
 		if err != nil {
 			return 1
 		}
 		runner.Args = []string{"export-numpy", "-local=true",
+			"-pprof", ":6060",
 			fmt.Sprintf("-one-hot=%v", *onehot),
-			"-i", *inputFilename,
-			"-o", "/mnt/output/matrix.npy",
-			"-output-annotations", "/mnt/output/annotations.tsv",
-			"-output-onehot2tilevar", "/mnt/output/onehot2tilevar.tsv",
-			"-max-variants", fmt.Sprintf("%d", cmd.filter.MaxVariants),
-			"-min-coverage", fmt.Sprintf("%f", cmd.filter.MinCoverage),
-			"-max-tag", fmt.Sprintf("%d", cmd.filter.MaxTag),
+			"-input-dir", *inputDir,
+			"-output-dir", "/mnt/output",
+			"-output-annotations", "/mnt/output/annotations.csv",
+			"-output-onehot2tilevar", "/mnt/output/onehot2tilevar.csv",
+			"-output-labels", "/mnt/output/labels.csv",
+			"-regions", *regionsFilename,
+			"-expand-regions", fmt.Sprintf("%d", *expandRegions),
+			"-chunks", fmt.Sprintf("%d", *chunks),
 		}
+		runner.Args = append(runner.Args, cmd.filter.Args()...)
 		var output string
 		output, err = runner.Run()
 		if err != nil {
@@ -92,26 +107,12 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
 		return 0
 	}
 
-	var input io.ReadCloser
-	if *inputFilename == "-" {
-		input = ioutil.NopCloser(stdin)
-	} else {
-		input, err = os.Open(*inputFilename)
-		if err != nil {
-			return 1
-		}
-		defer input.Close()
-	}
 	tilelib := &tileLibrary{
 		retainNoCalls:       true,
 		retainTileSequences: true,
 		compactGenomes:      map[string][]tileVariantID{},
 	}
-	err = tilelib.LoadGob(context.Background(), input, strings.HasSuffix(*inputFilename, ".gz"), nil)
-	if err != nil {
-		return 1
-	}
-	err = input.Close()
+	err = tilelib.LoadDir(context.Background(), *inputDir)
 	if err != nil {
 		return 1
 	}
@@ -121,15 +122,61 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
 	log.Info("tidying")
 	tilelib.Tidy()
 
+	log.Info("building lowqual map")
+	lowqual := lowqual(tilelib)
+	names := cgnames(tilelib)
+
+	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 := "matrix.npy"
+		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
+		}
+	}
+
+	log.Info("determining which tiles intersect given regions")
+	dropTiles, err := chooseTiles(tilelib, *regionsFilename, *expandRegions)
+	if err != nil {
+		return 1
+	}
+
+	annotation2tvs := map[string]map[hgvs.Variant][]tileLibRef{}
 	if *annotationsFilename != "" {
-		log.Infof("writing annotations")
+		log.Info("writing annotations")
 		var annow io.WriteCloser
 		annow, err = os.OpenFile(*annotationsFilename, os.O_CREATE|os.O_WRONLY, 0666)
 		if err != nil {
 			return 1
 		}
 		defer annow.Close()
-		err = (&annotatecmd{maxTileSize: 5000}).exportTileDiffs(annow, tilelib)
+		var mtx sync.Mutex
+		err = (&annotatecmd{
+			maxTileSize: 5000,
+			dropTiles:   dropTiles,
+			reportAnnotation: func(tag tagID, _ int, variant tileVariantID, refname string, seqname string, pdi hgvs.Variant) {
+				mtx.Lock()
+				defer mtx.Unlock()
+				if annotation2tvs[seqname] == nil {
+					annotation2tvs[seqname] = map[hgvs.Variant][]tileLibRef{}
+				}
+				annotation2tvs[seqname][pdi] = append(annotation2tvs[seqname][pdi], tileLibRef{Tag: tag, Variant: variant})
+			},
+		}).exportTileDiffs(annow, tilelib)
 		if err != nil {
 			return 1
 		}
@@ -139,47 +186,149 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
 		}
 	}
 
-	log.Info("building numpy array")
-	out, rows, cols := cgs2array(tilelib)
+	var lastErr atomic.Value
+	var wg sync.WaitGroup
+	for seqname, pdivars := range annotation2tvs {
+		seqname, pdivars := seqname, pdivars
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			log.Infof("choosing hgvs columns for seq %s", seqname)
+			var pdis []hgvs.Variant
+			for pdi, librefs := range pdivars {
+				// Include this HGVS column if it was
+				// seen in a variant of any
+				// non-dropped tile.
+				for _, libref := range librefs {
+					if int(libref.Tag) >= len(dropTiles) || !dropTiles[libref.Tag] {
+						pdis = append(pdis, pdi)
+						break
+					}
+				}
+			}
+			sort.Slice(pdis, func(i, j int) bool {
+				if cmp := pdis[i].Position - pdis[j].Position; cmp != 0 {
+					return cmp < 0
+				} else if pdis[i].Ref != pdis[j].Ref {
+					return pdis[i].Ref < pdis[j].Ref
+				} else {
+					return pdis[i].New < pdis[j].New
+				}
+			})
+			log.Infof("writing column labels for seq %s", seqname)
+			var buf bytes.Buffer
+			for _, pdi := range pdis {
+				fmt.Fprintf(&buf, "%s:g.%s\n", seqname, pdi.String())
+			}
+			err := ioutil.WriteFile(*outputDir+"/"+seqname+".columns.csv", buf.Bytes(), 0777)
+			if err != nil {
+				lastErr.Store(err)
+				return
+			}
+
+			log.Infof("building hgvs matrix for seq %s", seqname)
+			data := make([]int8, len(names)*len(pdis)*2)
+			for row, name := range names {
+				cg := tilelib.compactGenomes[name]
+				rowstart := row * len(pdis) * 2
+				for col, pdi := range pdis {
+					for _, libref := range pdivars[pdi] {
+						if len(cg) <= int(libref.Tag)*2+1 {
+							continue
+						}
+						for phase := 0; phase < 2; phase++ {
+							if cg[int(libref.Tag)*2+phase] == libref.Variant {
+								data[rowstart+col*2+phase] = 1
+							}
+						}
+					}
+				}
+			}
+			log.Infof("writing hgvs numpy for seq %s", seqname)
+			f, err := os.OpenFile(*outputDir+"/"+seqname+".npy", os.O_CREATE|os.O_WRONLY, 0777)
+			if err != nil {
+				lastErr.Store(err)
+				return
+			}
+			defer f.Close()
+			// gonpy closes our writer and ignores errors. Give it a nopCloser so we can close f properly.
+			npw, err := gonpy.NewWriter(nopCloser{f})
+			if err != nil {
+				lastErr.Store(err)
+				return
+			}
+			npw.Shape = []int{len(names), len(pdis) * 2}
+			err = npw.WriteInt8(data)
+			if err != nil {
+				lastErr.Store(err)
+				return
+			}
+			err = f.Close()
+			if err != nil {
+				lastErr.Store(err)
+				return
+			}
+		}()
+	}
+	wg.Wait()
+	if e, ok := lastErr.Load().(error); ok {
+		err = e
+		return 1
+	}
 
-	log.Info("writing numpy file")
-	var output io.WriteCloser
-	if *outputFilename == "-" {
-		output = nopCloser{stdout}
-	} else {
-		output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
+	chunksize := (len(tilelib.variant) + *chunks - 1) / *chunks
+	for chunk := 0; chunk < *chunks; chunk++ {
+		log.Infof("preparing chunk %d of %d", chunk+1, *chunks)
+		tagstart := chunk * chunksize
+		tagend := tagstart + chunksize
+		if tagend > len(tilelib.variant) {
+			tagend = len(tilelib.variant)
+		}
+		out, rows, cols := cgs2array(tilelib, names, lowqual, dropTiles, tagstart, tagend)
+
+		var npw *gonpy.NpyWriter
+		var output io.WriteCloser
+		fnm := *outputDir + "/matrix.npy"
+		if *chunks > 1 {
+			fnm = fmt.Sprintf("%s/matrix.%d.npy", *outputDir, chunk)
+		}
+		output, err = os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY, 0777)
 		if err != nil {
 			return 1
 		}
 		defer output.Close()
-	}
-	bufw := bufio.NewWriter(output)
-	npw, err := gonpy.NewWriter(nopCloser{bufw})
-	if err != nil {
-		return 1
-	}
-	if *onehot {
-		log.Info("recoding to onehot")
-		recoded, librefs, recodedcols := recodeOnehot(out, cols)
-		out, cols = recoded, recodedcols
-		if *librefsFilename != "" {
-			log.Infof("writing onehot column mapping")
-			err = cmd.writeLibRefs(*librefsFilename, tilelib, librefs)
-			if err != nil {
-				return 1
+		bufw := bufio.NewWriter(output)
+		npw, err = gonpy.NewWriter(nopCloser{bufw})
+		if err != nil {
+			return 1
+		}
+		if *onehot {
+			log.Info("recoding to onehot")
+			recoded, librefs, recodedcols := recodeOnehot(out, cols)
+			out, cols = recoded, recodedcols
+			if *librefsFilename != "" {
+				log.Infof("writing onehot column mapping")
+				err = cmd.writeLibRefs(*librefsFilename, tilelib, librefs)
+				if err != nil {
+					return 1
+				}
 			}
 		}
-	}
-	log.Info("writing numpy")
-	npw.Shape = []int{rows, cols}
-	npw.WriteInt16(out)
-	err = bufw.Flush()
-	if err != nil {
-		return 1
-	}
-	err = output.Close()
-	if err != nil {
-		return 1
+		log.WithFields(logrus.Fields{
+			"filename": fnm,
+			"rows":     rows,
+			"cols":     cols,
+		}).Info("writing numpy")
+		npw.Shape = []int{rows, cols}
+		npw.WriteInt16(out)
+		err = bufw.Flush()
+		if err != nil {
+			return 1
+		}
+		err = output.Close()
+		if err != nil {
+			return 1
+		}
 	}
 	return 0
 }
@@ -191,7 +340,7 @@ func (*exportNumpy) writeLibRefs(fnm string, tilelib *tileLibrary, librefs []til
 	}
 	defer f.Close()
 	for i, libref := range librefs {
-		_, err = fmt.Fprintf(f, "%d\t%d\t%d\n", i, libref.Tag, libref.Variant)
+		_, err = fmt.Fprintf(f, "%d,%d,%d\n", i, libref.Tag, libref.Variant)
 		if err != nil {
 			return err
 		}
@@ -199,26 +348,22 @@ func (*exportNumpy) writeLibRefs(fnm string, tilelib *tileLibrary, librefs []til
 	return f.Close()
 }
 
-func cgs2array(tilelib *tileLibrary) (data []int16, rows, cols int) {
-	var cgnames []string
+func cgnames(tilelib *tileLibrary) (cgnames []string) {
 	for name := range tilelib.compactGenomes {
 		cgnames = append(cgnames, name)
 	}
-	sort.Strings(cgnames)
-
-	rows = len(tilelib.compactGenomes)
-	for _, cg := range tilelib.compactGenomes {
-		if cols < len(cg) {
-			cols = len(cg)
-		}
-	}
+	sort.Slice(cgnames, func(i, j int) bool {
+		return trimFilenameForLabel(cgnames[i]) < trimFilenameForLabel(cgnames[j])
+	})
+	return
+}
 
-	// flag low-quality tile variants so we can change to -1 below
-	lowqual := make([]map[tileVariantID]bool, cols/2)
+func lowqual(tilelib *tileLibrary) (lowqual []map[tileVariantID]bool) {
+	lowqual = make([]map[tileVariantID]bool, len(tilelib.variant))
 	for tag, variants := range tilelib.variant {
 		lq := lowqual[tag]
 		for varidx, hash := range variants {
-			if len(tilelib.seq[hash]) == 0 {
+			if len(tilelib.hashSequence(hash)) == 0 {
 				if lq == nil {
 					lq = map[tileVariantID]bool{}
 					lowqual[tag] = lq
@@ -227,18 +372,145 @@ func cgs2array(tilelib *tileLibrary) (data []int16, rows, cols int) {
 			}
 		}
 	}
+	return
+}
 
+func cgs2array(tilelib *tileLibrary, names []string, lowqual []map[tileVariantID]bool, dropTiles []bool, tagstart, tagend int) (data []int16, rows, cols int) {
+	rows = len(tilelib.compactGenomes)
+	for tag := tagstart; tag < tagend; tag++ {
+		if len(dropTiles) <= tag || !dropTiles[tag] {
+			cols += 2
+		}
+	}
 	data = make([]int16, rows*cols)
-	for row, name := range cgnames {
-		for i, v := range tilelib.compactGenomes[name] {
-			if v > 0 && lowqual[i/2][v] {
-				data[row*cols+i] = -1
+	for row, name := range names {
+		cg := tilelib.compactGenomes[name]
+		outidx := 0
+		for tag := tagstart; tag < tagend && tag*2+1 < len(cg); tag++ {
+			if len(dropTiles) > tag && dropTiles[tag] {
+				continue
+			}
+			for phase := 0; phase < 2; phase++ {
+				v := cg[tag*2+phase]
+				if v > 0 && lowqual[tag][v] {
+					data[row*cols+outidx] = -1
+				} else {
+					data[row*cols+outidx] = int16(v)
+				}
+				outidx++
+			}
+		}
+	}
+	return
+}
+
+func makeMask(regionsFilename string, expandRegions int) (*mask, error) {
+	log.Printf("makeMask: reading %s", regionsFilename)
+	rfile, err := zopen(regionsFilename)
+	if err != nil {
+		return nil, err
+	}
+	defer rfile.Close()
+	regions, err := io.ReadAll(rfile)
+	if err != nil {
+		return nil, err
+	}
+
+	log.Print("makeMask: building mask")
+	var mask mask
+	for _, line := range bytes.Split(regions, []byte{'\n'}) {
+		if bytes.HasPrefix(line, []byte{'#'}) {
+			continue
+		}
+		fields := bytes.Split(line, []byte{'\t'})
+		if len(fields) < 3 {
+			continue
+		}
+		refseqname := string(fields[0])
+		if strings.HasPrefix(refseqname, "chr") {
+			refseqname = refseqname[3:]
+		}
+		start, err1 := strconv.Atoi(string(fields[1]))
+		end, err2 := strconv.Atoi(string(fields[2]))
+		if err1 == nil && err2 == nil {
+			// BED
+		} else {
+			start, err1 = strconv.Atoi(string(fields[3]))
+			end, err2 = strconv.Atoi(string(fields[4]))
+			if err1 == nil && err2 == nil {
+				// GFF/GTF
+				end++
 			} else {
-				data[row*cols+i] = int16(v)
+				return nil, fmt.Errorf("cannot parse input line as BED or GFF/GTF: %q", line)
+			}
+		}
+		mask.Add(refseqname, start-expandRegions, end+expandRegions)
+	}
+	log.Print("makeMask: mask.Freeze")
+	mask.Freeze()
+	return &mask, nil
+}
+
+func chooseTiles(tilelib *tileLibrary, regionsFilename string, expandRegions int) (drop []bool, err error) {
+	if regionsFilename == "" {
+		return
+	}
+	mask, err := makeMask(regionsFilename, expandRegions)
+	if err != nil {
+		return
+	}
+
+	tagset := tilelib.taglib.Tags()
+	if len(tagset) == 0 {
+		err = errors.New("cannot choose tiles by region in a library without tags")
+		return
+	}
+	taglen := len(tagset[0])
+
+	log.Print("chooseTiles: check ref tiles")
+	// Find position+size of each reference tile, and if it
+	// intersects any of the desired regions, set drop[tag]=false.
+	//
+	// (Note it doesn't quite work to do the more obvious thing --
+	// start with drop=false and change to true when ref tiles
+	// intersect target regions -- because that would give us
+	// drop=false for tiles that don't appear at all in the
+	// reference.)
+	//
+	// TODO: (optionally?) don't drop tags for which some tile
+	// variants are spanning tiles, i.e., where the reference tile
+	// does not intersect the desired regions, but a spanning tile
+	// from a genome does.
+	drop = make([]bool, len(tilelib.variant))
+	for i := range drop {
+		drop[i] = true
+	}
+	for refname, refseqs := range tilelib.refseqs {
+		for refseqname, reftiles := range refseqs {
+			if strings.HasPrefix(refseqname, "chr") {
+				refseqname = refseqname[3:]
+			}
+			tileend := 0
+			for _, libref := range reftiles {
+				if libref.Variant < 1 {
+					err = fmt.Errorf("reference %q seq %q uses variant zero at tag %d", refname, refseqname, libref.Tag)
+					return
+				}
+				seq := tilelib.TileVariantSequence(libref)
+				if len(seq) < taglen {
+					err = fmt.Errorf("reference %q seq %q uses tile %d variant %d with sequence len %d < taglen %d", refname, refseqname, libref.Tag, libref.Variant, len(seq), taglen)
+					return
+				}
+				tilestart := tileend
+				tileend = tilestart + len(seq) - taglen
+				if mask.Check(refseqname, tilestart, tileend) {
+					drop[libref.Tag] = false
+				}
 			}
 		}
 	}
 
+	log.Print("chooseTiles: done")
 	return
 }
 
@@ -284,3 +556,17 @@ type nopCloser struct {
 }
 
 func (nopCloser) Close() error { return nil }
+
+func trimFilenameForLabel(s string) string {
+	if i := strings.LastIndex(s, "/"); i >= 0 {
+		s = s[i+1:]
+	}
+	s = strings.TrimSuffix(s, ".gz")
+	s = strings.TrimSuffix(s, ".fa")
+	s = strings.TrimSuffix(s, ".fasta")
+	s = strings.TrimSuffix(s, ".1")
+	s = strings.TrimSuffix(s, ".2")
+	s = strings.TrimSuffix(s, ".gz")
+	s = strings.TrimSuffix(s, ".vcf")
+	return s
+}