X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/ce7cc440993b56c3f87b040b33d002074f304feb..dfc54ac31cf542a4c53aeaa02dabddb17767513b:/slicenumpy.go

diff --git a/slicenumpy.go b/slicenumpy.go
index 4d3bd2745d..2cb67960b6 100644
--- a/slicenumpy.go
+++ b/slicenumpy.go
@@ -10,12 +10,14 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
 	"regexp"
 	"runtime"
 	"sort"
+	"strconv"
 	"strings"
 	"sync/atomic"
 
@@ -49,6 +51,7 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 	ref := flags.String("ref", "", "reference name (if blank, choose last one that appears in input)")
 	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`")
+	mergeOutput := flags.Bool("merge-output", false, "merge output into one matrix.npy and one matrix.annotations.csv")
 	flags.IntVar(&cmd.threads, "threads", 16, "number of memory-hungry assembly threads")
 	cmd.filter.Flags(flags)
 	err = flags.Parse(args)
@@ -70,7 +73,7 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 			Name:        "lightning slice-numpy",
 			Client:      arvados.NewClientFromEnv(),
 			ProjectUUID: *projectUUID,
-			RAM:         650000000000,
+			RAM:         750000000000,
 			VCPUs:       96,
 			Priority:    *priority,
 			KeepCache:   2,
@@ -81,12 +84,13 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 			return 1
 		}
 		runner.Args = []string{"slice-numpy", "-local=true",
-			"-pprof", ":6060",
-			"-input-dir", *inputDir,
-			"-output-dir", "/mnt/output",
-			"-threads", fmt.Sprintf("%d", cmd.threads),
-			"-regions", *regionsFilename,
-			"-expand-regions", fmt.Sprintf("%d", *expandRegions),
+			"-pprof=:6060",
+			"-input-dir=" + *inputDir,
+			"-output-dir=/mnt/output",
+			"-threads=" + fmt.Sprintf("%d", cmd.threads),
+			"-regions=" + *regionsFilename,
+			"-expand-regions=" + fmt.Sprintf("%d", *expandRegions),
+			"-merge-output=" + fmt.Sprintf("%v", *mergeOutput),
 		}
 		runner.Args = append(runner.Args, cmd.filter.Args()...)
 		var output string
@@ -189,41 +193,60 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 	type reftileinfo struct {
 		variant  tileVariantID
 		seqname  string // chr1
-		pos      int    // distance from start of chr1 to start of tile
+		pos      int    // distance from start of chromosome to starttag
 		tiledata []byte // acgtggcaa...
 	}
+	isdup := map[tagID]bool{}
 	reftile := map[tagID]*reftileinfo{}
 	for seqname, cseq := range refseq {
+		pos := 0
 		for _, libref := range cseq {
-			reftile[libref.Tag] = &reftileinfo{
-				seqname:  seqname,
-				variant:  libref.Variant,
-				tiledata: reftiledata[libref],
+			tiledata := reftiledata[libref]
+			if len(tiledata) == 0 {
+				err = fmt.Errorf("missing tiledata for tag %d variant %d in %s in ref", libref.Tag, libref.Variant, seqname)
+				return 1
+			}
+			if isdup[libref.Tag] {
+				log.Printf("dropping reference tile %+v from %s @ %d, tag not unique", libref, seqname, pos)
+			} else if reftile[libref.Tag] != nil {
+				log.Printf("dropping reference tile %+v from %s @ %d, tag not unique", tileLibRef{Tag: libref.Tag, Variant: reftile[libref.Tag].variant}, reftile[libref.Tag].seqname, reftile[libref.Tag].pos)
+				delete(reftile, libref.Tag)
+				log.Printf("dropping reference tile %+v from %s @ %d, tag not unique", libref, seqname, pos)
+				isdup[libref.Tag] = true
+			} else {
+				reftile[libref.Tag] = &reftileinfo{
+					seqname:  seqname,
+					variant:  libref.Variant,
+					tiledata: tiledata,
+					pos:      pos,
+				}
 			}
+			pos += len(tiledata) - taglen
 		}
+		log.Printf("... %s done, len %d", seqname, pos+taglen)
 	}
 
-	throttleCPU := throttle{Max: runtime.GOMAXPROCS(0)}
-	log.Info("reconstructing reference sequences")
-	for seqname, cseq := range refseq {
-		seqname, cseq := seqname, cseq
-		throttleCPU.Go(func() error {
-			defer log.Printf("... %s done", seqname)
-			pos := 0
-			for _, libref := range cseq {
-				rt := reftile[libref.Tag]
-				rt.pos = pos
-				if len(rt.tiledata) == 0 {
-					return fmt.Errorf("missing tiledata for tag %d variant %d in %s in ref", libref.Tag, libref.Variant, seqname)
-				}
-				pos += len(rt.tiledata) - taglen
+	var mask *mask
+	if *regionsFilename != "" {
+		log.Printf("loading regions from %s", *regionsFilename)
+		mask, err = makeMask(*regionsFilename, *expandRegions)
+		if err != nil {
+			return 1
+		}
+		log.Printf("before applying mask, len(reftile) == %d", len(reftile))
+		log.Printf("deleting reftile entries for regions outside %d intervals", mask.Len())
+		for tag, rt := range reftile {
+			if !mask.Check(strings.TrimPrefix(rt.seqname, "chr"), rt.pos, rt.pos+len(rt.tiledata)) {
+				delete(reftile, tag)
 			}
-			return nil
-		})
+		}
+		log.Printf("after applying mask, len(reftile) == %d", len(reftile))
 	}
-	throttleCPU.Wait()
 
-	log.Info("TODO: determining which tiles intersect given regions")
+	var toMerge [][]int16
+	if *mergeOutput {
+		toMerge = make([][]int16, len(infiles))
+	}
 
 	throttleMem := throttle{Max: cmd.threads} // TODO: estimate using mem and data size
 	throttleNumpyMem := throttle{Max: cmd.threads/2 + 1}
@@ -245,6 +268,12 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 					if tv.Ref {
 						continue
 					}
+					if mask != nil && reftile[tv.Tag] == nil {
+						// Don't waste
+						// time/memory on
+						// masked-out tiles.
+						continue
+					}
 					variants := seq[tv.Tag]
 					if len(variants) == 0 {
 						variants = make([]TileVariant, 100)
@@ -279,6 +308,12 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 				go func() {
 					defer throttleCPU.Release()
 					count := make(map[[blake2b.Size256]byte]int, len(variants))
+
+					rt := reftile[tag]
+					if rt != nil {
+						count[blake2b.Sum256(rt.tiledata)] = 0
+					}
+
 					for _, cg := range cgs {
 						idx := int(tag-tagstart) * 2
 						if idx < len(cg.Variants) {
@@ -320,6 +355,9 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 						remap[i] = rank[tv.Blake2b]
 					}
 					variantRemap[tag-tagstart] = remap
+					if rt != nil {
+						rt.variant = rank[blake2b.Sum256(rt.tiledata)]
+					}
 				}()
 			}
 			throttleCPU.Wait()
@@ -331,30 +369,38 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 				return err
 			}
 			annow := bufio.NewWriterSize(annof, 1<<20)
-			for tag, variants := range seq {
+			outcol := 0
+			for tag := tagstart; tag < tagend; tag++ {
 				rt, ok := reftile[tag]
 				if !ok {
-					// Reference does not use any
-					// variant of this tile.
-					// TODO: log this? mention it
-					// in annotations?
+					if mask == nil {
+						outcol++
+					}
+					// Excluded by specified
+					// regions, or reference does
+					// not use any variant of this
+					// tile. (TODO: log this?
+					// mention it in annotations?)
 					continue
 				}
-				outcol := tag - tagID(tagstart)
+				fmt.Fprintf(annow, "%d,%d,%d,=,%s,%d,,,\n", tag, outcol, rt.variant, rt.seqname, rt.pos)
+				variants := seq[tag]
 				reftilestr := strings.ToUpper(string(rt.tiledata))
 				remap := variantRemap[tag-tagstart]
 				done := make([]bool, len(variants))
 				for v, tv := range variants {
 					v := remap[v]
-					if done[v] {
+					if v == rt.variant || done[v] {
 						continue
 					} else {
 						done[v] = true
 					}
 					if len(tv.Sequence) < taglen || !bytes.HasSuffix(rt.tiledata, tv.Sequence[len(tv.Sequence)-taglen:]) {
+						fmt.Fprintf(annow, "%d,%d,%d,,%s,%d,,,\n", tag, outcol, v, rt.seqname, rt.pos)
 						continue
 					}
 					if lendiff := len(rt.tiledata) - len(tv.Sequence); lendiff < -1000 || lendiff > 1000 {
+						fmt.Fprintf(annow, "%d,%d,%d,,%s,%d,,,\n", tag, outcol, v, rt.seqname, rt.pos)
 						continue
 					}
 					diffs, _ := hgvs.Diff(reftilestr, strings.ToUpper(string(tv.Sequence)), 0)
@@ -363,6 +409,7 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 						fmt.Fprintf(annow, "%d,%d,%d,%s:g.%s,%s,%d,%s,%s,%s\n", tag, outcol, v, rt.seqname, diff.String(), rt.seqname, diff.Position, diff.Ref, diff.New, diff.Left)
 					}
 				}
+				outcol++
 			}
 			err = annow.Flush()
 			if err != nil {
@@ -373,49 +420,39 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 				return err
 			}
 
-			throttleNumpyMem.Acquire()
 			log.Infof("%04d: preparing numpy", infileIdx)
+			throttleNumpyMem.Acquire()
 			rows := len(cgnames)
-			cols := 2 * int(tagend-tagstart)
+			cols := 2 * outcol
 			out := make([]int16, rows*cols)
 			for row, name := range cgnames {
 				out := out[row*cols:]
+				outcol := 0
 				for col, v := range cgs[name].Variants {
-					if variants, ok := seq[tagstart+tagID(col/2)]; ok && len(variants) > int(v) && len(variants[v].Sequence) > 0 {
-						out[col] = int16(variantRemap[col/2][v])
+					tag := tagstart + tagID(col/2)
+					if mask != nil && reftile[tag] == nil {
+						continue
+					}
+					if variants, ok := seq[tag]; ok && len(variants) > int(v) && len(variants[v].Sequence) > 0 {
+						out[outcol] = int16(variantRemap[tag-tagstart][v])
 					} else {
-						out[col] = -1
+						out[outcol] = -1
 					}
+					outcol++
 				}
 			}
 			seq = nil
 			throttleNumpyMem.Release()
 
-			fnm := fmt.Sprintf("%s/matrix.%04d.npy", *outputDir, infileIdx)
-			output, err := os.Create(fnm)
-			if err != nil {
-				return err
-			}
-			defer output.Close()
-			bufw := bufio.NewWriterSize(output, 1<<26)
-			npw, err := gonpy.NewWriter(nopCloser{bufw})
-			if err != nil {
-				return err
-			}
-			log.WithFields(log.Fields{
-				"filename": fnm,
-				"rows":     rows,
-				"cols":     cols,
-			}).Infof("%04d: writing numpy", infileIdx)
-			npw.Shape = []int{rows, cols}
-			npw.WriteInt16(out)
-			err = bufw.Flush()
-			if err != nil {
-				return err
-			}
-			err = output.Close()
-			if err != nil {
-				return err
+			if *mergeOutput {
+				log.Infof("%04d: matrix fragment %d rows x %d cols", infileIdx, rows, cols)
+				toMerge[infileIdx] = out
+			} else {
+				fnm := fmt.Sprintf("%s/matrix.%04d.npy", *outputDir, infileIdx)
+				err = writeNumpyInt16(fnm, out, rows, cols)
+				if err != nil {
+					return err
+				}
 			}
 			log.Infof("%s: done (%d/%d)", infile, int(atomic.AddInt64(&done, 1)), len(infiles))
 			return nil
@@ -424,5 +461,182 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
 	if err = throttleMem.Wait(); err != nil {
 		return 1
 	}
+	if *mergeOutput {
+		log.Info("merging output matrix and annotations")
+
+		annoFilename := fmt.Sprintf("%s/matrix.annotations.csv", *outputDir)
+		annof, err := os.Create(annoFilename)
+		if err != nil {
+			return 1
+		}
+		annow := bufio.NewWriterSize(annof, 1<<20)
+
+		rows := len(cgnames)
+		cols := 0
+		for _, chunk := range toMerge {
+			cols += len(chunk) / rows
+		}
+		out := make([]int16, rows*cols)
+		hgvsCols := map[string][2][]int16{} // hgvs -> [[g0,g1,g2,...], [g0,g1,g2,...]] (slice of genomes for each phase)
+		startcol := 0
+		for outIdx, chunk := range toMerge {
+			chunkcols := len(chunk) / rows
+			for row := 0; row < rows; row++ {
+				copy(out[row*cols+startcol:], chunk[row*chunkcols:(row+1)*chunkcols])
+			}
+			toMerge[outIdx] = nil
+
+			annotationsFilename := fmt.Sprintf("%s/matrix.%04d.annotations.csv", *outputDir, outIdx)
+			log.Infof("reading %s", annotationsFilename)
+			buf, err := os.ReadFile(annotationsFilename)
+			if err != nil {
+				return 1
+			}
+			err = os.Remove(annotationsFilename)
+			if err != nil {
+				return 1
+			}
+			for _, line := range bytes.Split(buf, []byte{'\n'}) {
+				if len(line) == 0 {
+					continue
+				}
+				fields := bytes.SplitN(line, []byte{','}, 9)
+				tag, _ := strconv.Atoi(string(fields[0]))
+				incol, _ := strconv.Atoi(string(fields[1]))
+				tileVariant, _ := strconv.Atoi(string(fields[2]))
+				hgvsID := string(fields[3])
+				seqname := string(fields[4])
+				pos, _ := strconv.Atoi(string(fields[5]))
+				refseq := fields[6]
+				if hgvsID == "" {
+					// Null entry for un-diffable
+					// tile variant
+					continue
+				}
+				if hgvsID == "=" {
+					// Null entry for ref tile
+					continue
+				}
+				if mask != nil && !mask.Check(strings.TrimPrefix(seqname, "chr"), pos, pos+len(refseq)) {
+					// The tile intersects one of
+					// the selected regions, but
+					// this particular HGVS
+					// variant does not.
+					continue
+				}
+				hgvsColPair := hgvsCols[hgvsID]
+				if hgvsColPair[0] == nil {
+					// values in new columns start
+					// out as -1 ("no data yet")
+					// or 0 ("=ref") here, may
+					// change to 1 ("hgvs variant
+					// present") below, either on
+					// this line or a future line.
+					hgvsColPair = [2][]int16{make([]int16, len(cgnames)), make([]int16, len(cgnames))}
+					rt, ok := reftile[tagID(tag)]
+					if !ok {
+						err = fmt.Errorf("bug: seeing annotations for tag %d, but it has no reftile entry", tag)
+						return 1
+					}
+					for ph := 0; ph < 2; ph++ {
+						for row := 0; row < rows; row++ {
+							v := chunk[row*chunkcols+incol*2+ph]
+							if tileVariantID(v) == rt.variant {
+								hgvsColPair[ph][row] = 0
+							} else {
+								hgvsColPair[ph][row] = -1
+							}
+						}
+					}
+					hgvsCols[hgvsID] = hgvsColPair
+					hgvsref := hgvs.Variant{
+						Position: pos,
+						Ref:      string(refseq),
+						New:      string(refseq),
+					}
+					fmt.Fprintf(annow, "%d,%d,%d,%s:g.%s,%s,%d,%s,%s,%s\n", tag, incol+startcol/2, rt.variant, seqname, hgvsref.String(), seqname, pos, refseq, refseq, fields[8])
+				}
+				fmt.Fprintf(annow, "%d,%d,%d,%s,%s,%d,%s,%s,%s\n", tag, incol+startcol/2, tileVariant, hgvsID, seqname, pos, refseq, fields[7], fields[8])
+				for ph := 0; ph < 2; ph++ {
+					for row := 0; row < rows; row++ {
+						v := chunk[row*chunkcols+incol*2+ph]
+						if int(v) == tileVariant {
+							hgvsColPair[ph][row] = 1
+						}
+					}
+				}
+			}
+
+			startcol += chunkcols
+		}
+		err = annow.Flush()
+		if err != nil {
+			return 1
+		}
+		err = annof.Close()
+		if err != nil {
+			return 1
+		}
+		err = writeNumpyInt16(fmt.Sprintf("%s/matrix.npy", *outputDir), out, rows, cols)
+		if err != nil {
+			return 1
+		}
+		out = nil
+
+		cols = len(hgvsCols) * 2
+		log.Printf("building hgvs-based matrix: %d rows x %d cols", rows, cols)
+		out = make([]int16, rows*cols)
+		hgvsIDs := make([]string, 0, len(hgvsCols))
+		for hgvsID := range hgvsCols {
+			hgvsIDs = append(hgvsIDs, hgvsID)
+		}
+		sort.Strings(hgvsIDs)
+		var hgvsLabels bytes.Buffer
+		for idx, hgvsID := range hgvsIDs {
+			fmt.Fprintf(&hgvsLabels, "%d,%s\n", idx, hgvsID)
+			for ph := 0; ph < 2; ph++ {
+				hgvscol := hgvsCols[hgvsID][ph]
+				for row, val := range hgvscol {
+					out[row*cols+idx*2+ph] = val
+				}
+			}
+		}
+		err = writeNumpyInt16(fmt.Sprintf("%s/hgvs.npy", *outputDir), out, rows, cols)
+		if err != nil {
+			return 1
+		}
+
+		fnm := fmt.Sprintf("%s/hgvs.annotations.csv", *outputDir)
+		log.Printf("writing hgvs labels: %s", fnm)
+		err = ioutil.WriteFile(fnm, hgvsLabels.Bytes(), 0777)
+		if err != nil {
+			return 1
+		}
+	}
 	return 0
 }
+
+func writeNumpyInt16(fnm string, out []int16, rows, cols int) error {
+	output, err := os.Create(fnm)
+	if err != nil {
+		return err
+	}
+	defer output.Close()
+	bufw := bufio.NewWriterSize(output, 1<<26)
+	npw, err := gonpy.NewWriter(nopCloser{bufw})
+	if err != nil {
+		return err
+	}
+	log.WithFields(log.Fields{
+		"filename": fnm,
+		"rows":     rows,
+		"cols":     cols,
+	}).Infof("writing numpy: %s", fnm)
+	npw.Shape = []int{rows, cols}
+	npw.WriteInt16(out)
+	err = bufw.Flush()
+	if err != nil {
+		return err
+	}
+	return output.Close()
+}