X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/08aa5f2c415150de2738378b6d4fbce88a9e6a59..a0a52e9e50ff032df3605a0fbadb2725b67010f7:/slicenumpy.go

diff --git a/slicenumpy.go b/slicenumpy.go
index bed143e149..e60c50adb1 100644
--- a/slicenumpy.go
+++ b/slicenumpy.go
@@ -8,6 +8,7 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/gob"
+	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
@@ -39,18 +40,24 @@ import (
 const annotationMaxTileSpan = 100
 
 type sliceNumpy struct {
-	filter          filter
-	threads         int
-	chi2Cases       []bool
-	chi2PValue      float64
-	minCoverage     int
-	includeVariant1 bool
-	debugTag        tagID
+	filter             filter
+	threads            int
+	chi2Cases          []bool
+	chi2PValue         float64
+	pvalueMinFrequency float64
+	maxFrequency       float64
+	pcaComponents      int
+	minCoverage        int
+	minCoverageAll     bool
+	includeVariant1    bool
+	debugTag           tagID
 
 	cgnames         []string
 	samples         []sampleInfo
 	trainingSet     []int // samples index => training set index, or -1 if not in training set
 	trainingSetSize int
+	pvalue          func(onehot []bool) float64
+	pvalueCallCount int64
 }
 
 func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
@@ -71,6 +78,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 	arvadosVCPUs := flags.Int("arvados-vcpus", 96, "number of VCPUs to request for arvados container")
 	projectUUID := flags.String("project", "", "project `UUID` for output data")
 	priority := flags.Int("priority", 500, "container request priority")
+	preemptible := flags.Bool("preemptible", true, "request preemptible instance")
 	inputDir := flags.String("input-dir", "./in", "input `directory`")
 	outputDir := flags.String("output-dir", "./out", "output `directory`")
 	ref := flags.String("ref", "", "reference name (if blank, choose last one that appears in input)")
@@ -83,12 +91,15 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 	onehotChunked := flags.Bool("chunked-onehot", false, "generate one-hot tile-based matrix per input chunk")
 	samplesFilename := flags.String("samples", "", "`samples.csv` file with training/validation and case/control groups (see 'lightning choose-samples')")
 	caseControlOnly := flags.Bool("case-control-only", false, "drop samples that are not in case/control groups")
-	onlyPCA := flags.Bool("pca", false, "generate pca matrix")
-	pcaComponents := flags.Int("pca-components", 4, "number of PCA components")
+	onlyPCA := flags.Bool("pca", false, "run principal component analysis, write components to pca.npy and samples.csv")
+	flags.IntVar(&cmd.pcaComponents, "pca-components", 4, "number of PCA components to compute / use in logistic regression")
 	maxPCATiles := flags.Int("max-pca-tiles", 0, "maximum tiles to use as PCA input (filter, then drop every 2nd colum pair until below max)")
 	debugTag := flags.Int("debug-tag", -1, "log debugging details about specified tag")
+	flags.BoolVar(&cmd.minCoverageAll, "min-coverage-all", false, "apply -min-coverage filter based on all samples, not just training set")
 	flags.IntVar(&cmd.threads, "threads", 16, "number of memory-hungry assembly threads, and number of VCPUs to request for arvados container")
-	flags.Float64Var(&cmd.chi2PValue, "chi2-p-value", 1, "do Χ² test and omit columns with p-value above this threshold")
+	flags.Float64Var(&cmd.chi2PValue, "chi2-p-value", 1, "do Χ² test (or logistic regression if -samples file has PCA components) and omit columns with p-value above this threshold")
+	flags.Float64Var(&cmd.pvalueMinFrequency, "pvalue-min-frequency", 0.01, "skip p-value calculation on tile variants below this frequency in the training set")
+	flags.Float64Var(&cmd.maxFrequency, "max-frequency", 1, "do not output variants above this frequency in the training set")
 	flags.BoolVar(&cmd.includeVariant1, "include-variant-1", false, "include most common variant when building one-hot matrix")
 	cmd.filter.Flags(flags)
 	err := flags.Parse(args)
@@ -96,6 +107,8 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 		return nil
 	} else if err != nil {
 		return err
+	} else if flags.NArg() > 0 {
+		return fmt.Errorf("errant command line arguments after parsed flags: %v", flags.Args())
 	}
 
 	if *pprof != "" {
@@ -120,6 +133,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 			Priority:    *priority,
 			KeepCache:   2,
 			APIAccess:   true,
+			Preemptible: *preemptible,
 		}
 		err = runner.TranslatePaths(inputDir, regionsFilename, samplesFilename)
 		if err != nil {
@@ -139,10 +153,13 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 			"-chunked-onehot=" + fmt.Sprintf("%v", *onehotChunked),
 			"-samples=" + *samplesFilename,
 			"-case-control-only=" + fmt.Sprintf("%v", *caseControlOnly),
+			"-min-coverage-all=" + fmt.Sprintf("%v", cmd.minCoverageAll),
 			"-pca=" + fmt.Sprintf("%v", *onlyPCA),
-			"-pca-components=" + fmt.Sprintf("%d", *pcaComponents),
+			"-pca-components=" + fmt.Sprintf("%d", cmd.pcaComponents),
 			"-max-pca-tiles=" + fmt.Sprintf("%d", *maxPCATiles),
 			"-chi2-p-value=" + fmt.Sprintf("%f", cmd.chi2PValue),
+			"-pvalue-min-frequency=" + fmt.Sprintf("%f", cmd.pvalueMinFrequency),
+			"-max-frequency=" + fmt.Sprintf("%f", cmd.maxFrequency),
 			"-include-variant-1=" + fmt.Sprintf("%v", cmd.includeVariant1),
 			"-debug-tag=" + fmt.Sprintf("%d", cmd.debugTag),
 		}
@@ -180,7 +197,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 	}
 
 	if *samplesFilename != "" {
-		cmd.samples, err = cmd.loadSampleInfo(*samplesFilename)
+		cmd.samples, err = loadSampleInfo(*samplesFilename)
 		if err != nil {
 			return err
 		}
@@ -277,50 +294,47 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 				cmd.trainingSet[i] = -1
 			}
 		}
+		if cmd.pvalue == nil {
+			cmd.pvalue = func(onehot []bool) float64 {
+				return pvalue(onehot, cmd.chi2Cases)
+			}
+		}
 	}
-	if cmd.filter.MinCoverage == 1 {
-		// In the generic formula below, floating point
-		// arithmetic can effectively push the coverage
-		// threshold above 1.0, which is impossible/useless.
-		// 1.0 needs to mean exactly 100% coverage.
+
+	if cmd.minCoverageAll {
 		cmd.minCoverage = len(cmd.cgnames)
 	} else {
-		cmd.minCoverage = int(math.Ceil(cmd.filter.MinCoverage * float64(len(cmd.cgnames))))
+		cmd.minCoverage = cmd.trainingSetSize
+	}
+	if cmd.filter.MinCoverage < 1 {
+		cmd.minCoverage = int(math.Ceil(cmd.filter.MinCoverage * float64(cmd.minCoverage)))
 	}
 
-	{
-		samplesOutFilename := *outputDir + "/samples.csv"
-		log.Infof("writing sample metadata to %s", samplesOutFilename)
-		var f *os.File
-		f, err = os.Create(samplesOutFilename)
+	if len(cmd.samples[0].pcaComponents) > 0 {
+		cmd.pvalue = glmPvalueFunc(cmd.samples, cmd.pcaComponents)
+		// Unfortunately, statsmodel/glm lib logs stuff to
+		// os.Stdout when it panics on an unsolvable
+		// problem. We recover() from the panic in glm.go, but
+		// we also need to commandeer os.Stdout to avoid
+		// producing large quantities of logs.
+		stdoutWas := os.Stdout
+		defer func() { os.Stdout = stdoutWas }()
+		os.Stdout, err = os.Open(os.DevNull)
 		if err != nil {
 			return err
 		}
-		defer f.Close()
-		for i, si := range cmd.samples {
-			var cc, tv string
-			if si.isCase {
-				cc = "1"
-			} else if si.isControl {
-				cc = "0"
-			}
-			if si.isTraining {
-				tv = "1"
-			} else {
-				tv = "0"
-			}
-			_, err = fmt.Fprintf(f, "%d,%s,%s,%s\n", i, si.id, cc, tv)
-			if err != nil {
-				err = fmt.Errorf("write %s: %w", samplesOutFilename, err)
-				return err
-			}
-		}
-		err = f.Close()
-		if err != nil {
-			err = fmt.Errorf("close %s: %w", samplesOutFilename, err)
-			return err
-		}
-		log.Print("done")
+	}
+
+	// cgnamemap[name]==true for samples that we are including in
+	// output
+	cgnamemap := map[string]bool{}
+	for _, name := range cmd.cgnames {
+		cgnamemap[name] = true
+	}
+
+	err = writeSampleInfo(cmd.samples, *outputDir)
+	if err != nil {
+		return err
 	}
 
 	log.Info("indexing reference tiles")
@@ -504,7 +518,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 					if cmd.filter.MaxTag >= 0 && cg.StartTag > tagID(cmd.filter.MaxTag) {
 						return errSkip
 					}
-					if !matchGenome.MatchString(cg.Name) {
+					if !cgnamemap[cg.Name] {
 						continue
 					}
 					// pad to full slice size
@@ -520,7 +534,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 			if err == errSkip {
 				return nil
 			} else if err != nil {
-				return fmt.Errorf("%04d: DecodeLibrary(%s): err", infileIdx, infile)
+				return fmt.Errorf("%04d: DecodeLibrary(%s): %w", infileIdx, infile, err)
 			}
 			tagstart := cgs[cmd.cgnames[0]].StartTag
 			tagend := cgs[cmd.cgnames[0]].EndTag
@@ -542,7 +556,11 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 						count[blake2b.Sum256(rt.tiledata)] = 0
 					}
 
-					for cgname, cg := range cgs {
+					for cgidx, cgname := range cmd.cgnames {
+						if !cmd.minCoverageAll && !cmd.samples[cgidx].isTraining {
+							continue
+						}
+						cg := cgs[cgname]
 						idx := int(tag-tagstart) * 2
 						for allele := 0; allele < 2; allele++ {
 							v := cg.Variants[idx+allele]
@@ -854,7 +872,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 						if cmd.filter.MaxTag >= 0 && tag > tagID(cmd.filter.MaxTag) {
 							break
 						}
-						if rt := reftile[tag]; rt == nil || rt.excluded {
+						if rt := reftile[tag]; mask != nil && (rt == nil || rt.excluded) {
 							continue
 						}
 						if v == 0 {
@@ -1171,6 +1189,17 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 			if err != nil {
 				return err
 			}
+			fnm = fmt.Sprintf("%s/stats.json", *outputDir)
+			j, err := json.Marshal(map[string]interface{}{
+				"pvalueCallCount": cmd.pvalueCallCount,
+			})
+			if err != nil {
+				return err
+			}
+			err = os.WriteFile(fnm, j, 0777)
+			if err != nil {
+				return err
+			}
 		}
 		if *onlyPCA {
 			cols := 0
@@ -1205,7 +1234,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 				}
 			}
 			log.Print("fitting")
-			transformer := nlp.NewPCA(*pcaComponents)
+			transformer := nlp.NewPCA(cmd.pcaComponents)
 			transformer.Fit(mtxTrain.T())
 			log.Printf("transforming")
 			pca, err := transformer.Transform(mtxFull.T())
@@ -1242,42 +1271,19 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
 			}
 			log.Print("done")
 
-			samplesOutFilename := *outputDir + "/samples.csv"
-			log.Infof("writing sample metadata to %s", samplesOutFilename)
-			var f *os.File
-			f, err = os.Create(samplesOutFilename)
-			if err != nil {
-				return err
-			}
-			defer f.Close()
-			for i, si := range cmd.samples {
-				var cc, tv string
-				if si.isCase {
-					cc = "1"
-				} else if si.isControl {
-					cc = "0"
-				}
-				if si.isTraining {
-					tv = "1"
-				} else {
-					tv = "0"
-				}
-				var pcavals string
+			log.Print("copying pca components to sampleInfo")
+			for i := range cmd.samples {
+				cmd.samples[i].pcaComponents = make([]float64, outcols)
 				for c := 0; c < outcols; c++ {
-					pcavals += fmt.Sprintf(",%f", pca.At(i, c))
-				}
-				_, err = fmt.Fprintf(f, "%d,%s,%s,%s%s\n", i, si.id, cc, tv, pcavals)
-				if err != nil {
-					err = fmt.Errorf("write %s: %w", samplesOutFilename, err)
-					return err
+					cmd.samples[i].pcaComponents[c] = pca.At(i, c)
 				}
 			}
-			err = f.Close()
+			log.Print("done")
+
+			err = writeSampleInfo(cmd.samples, *outputDir)
 			if err != nil {
-				err = fmt.Errorf("close %s: %w", samplesOutFilename, err)
 				return err
 			}
-			log.Print("done")
 		}
 	}
 	if !*mergeOutput && !*onehotChunked && !*onehotSingle && !*onlyPCA {
@@ -1317,7 +1323,7 @@ type sampleInfo struct {
 
 // Read samples.csv file with case/control and training/validation
 // flags.
-func (cmd *sliceNumpy) loadSampleInfo(samplesFilename string) ([]sampleInfo, error) {
+func loadSampleInfo(samplesFilename string) ([]sampleInfo, error) {
 	var si []sampleInfo
 	f, err := open(samplesFilename)
 	if err != nil {
@@ -1351,17 +1357,75 @@ func (cmd *sliceNumpy) loadSampleInfo(samplesFilename string) ([]sampleInfo, err
 		if idx != len(si) {
 			return nil, fmt.Errorf("%s line %d: index %d out of order", samplesFilename, lineNum, idx)
 		}
+		var pcaComponents []float64
+		if len(split) > 4 {
+			for _, s := range split[4:] {
+				f, err := strconv.ParseFloat(s, 64)
+				if err != nil {
+					return nil, fmt.Errorf("%s line %d: cannot parse float %q: %s", samplesFilename, lineNum, s, err)
+				}
+				pcaComponents = append(pcaComponents, f)
+			}
+		}
 		si = append(si, sampleInfo{
-			id:           split[1],
-			isCase:       split[2] == "1",
-			isControl:    split[2] == "0",
-			isTraining:   split[3] == "1",
-			isValidation: split[3] == "0",
+			id:            split[1],
+			isCase:        split[2] == "1",
+			isControl:     split[2] == "0",
+			isTraining:    split[3] == "1",
+			isValidation:  split[3] == "0" && len(split[2]) > 0, // fix errant 0s in input
+			pcaComponents: pcaComponents,
 		})
 	}
 	return si, nil
 }
 
+func writeSampleInfo(samples []sampleInfo, outputDir string) error {
+	fnm := outputDir + "/samples.csv"
+	log.Infof("writing sample metadata to %s", fnm)
+	f, err := os.Create(fnm)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	pcaLabels := ""
+	if len(samples) > 0 {
+		for i := range samples[0].pcaComponents {
+			pcaLabels += fmt.Sprintf(",PCA%d", i)
+		}
+	}
+	_, err = fmt.Fprintf(f, "Index,SampleID,CaseControl,TrainingValidation%s\n", pcaLabels)
+	if err != nil {
+		return err
+	}
+	for i, si := range samples {
+		var cc, tv string
+		if si.isCase {
+			cc = "1"
+		} else if si.isControl {
+			cc = "0"
+		}
+		if si.isTraining {
+			tv = "1"
+		} else if si.isValidation {
+			tv = "0"
+		}
+		var pcavals string
+		for _, pcaval := range si.pcaComponents {
+			pcavals += fmt.Sprintf(",%f", pcaval)
+		}
+		_, err = fmt.Fprintf(f, "%d,%s,%s,%s%s\n", i, si.id, cc, tv, pcavals)
+		if err != nil {
+			return fmt.Errorf("write %s: %w", fnm, err)
+		}
+	}
+	err = f.Close()
+	if err != nil {
+		return fmt.Errorf("close %s: %w", fnm, err)
+	}
+	log.Print("done")
+	return nil
+}
+
 func (cmd *sliceNumpy) filterHGVScolpair(colpair [2][]int8) bool {
 	if cmd.chi2PValue >= 1 {
 		return true
@@ -1510,6 +1574,7 @@ type onehotXref struct {
 	variant tileVariantID
 	hom     bool
 	pvalue  float64
+	maf     float64
 }
 
 const onehotXrefSize = unsafe.Sizeof(onehotXref{})
@@ -1538,7 +1603,11 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
 	}
 	tagoffset := tag - chunkstarttag
 	coverage := 0
-	for _, cg := range cgs {
+	for cgidx, cgname := range cmd.cgnames {
+		if !cmd.minCoverageAll && !cmd.samples[cgidx].isTraining {
+			continue
+		}
+		cg := cgs[cgname]
 		alleles := 0
 		for _, v := range cg.Variants[tagoffset*2 : tagoffset*2+2] {
 			if v > 0 && int(v) < len(seq[tag]) && len(seq[tag][v].Sequence) > 0 {
@@ -1581,6 +1650,7 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
 	}
 	var onehot [][]int8
 	var xref []onehotXref
+	var maf float64
 	for col := 2; col < len(obs); col++ {
 		// col 0,1 correspond to tile variant 0, i.e.,
 		// no-call; col 2,3 correspond to the most common
@@ -1588,7 +1658,23 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
 		if col < 4 && !cmd.includeVariant1 {
 			continue
 		}
-		p := pvalue(obs[col], cmd.chi2Cases)
+		if col&1 == 0 {
+			maf = homhet2maf(obs[col : col+2])
+			if maf < cmd.pvalueMinFrequency {
+				// Skip both columns (hom and het) if
+				// allele frequency is below threshold
+				col++
+				continue
+			}
+			if maf > cmd.maxFrequency {
+				// Skip both columns if allele
+				// frequency is above threshold
+				col++
+				continue
+			}
+		}
+		atomic.AddInt64(&cmd.pvalueCallCount, 1)
+		p := cmd.pvalue(obs[col])
 		if cmd.chi2PValue < 1 && !(p < cmd.chi2PValue) {
 			continue
 		}
@@ -1598,11 +1684,29 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
 			variant: tileVariantID(col >> 1),
 			hom:     col&1 == 0,
 			pvalue:  p,
+			maf:     maf,
 		})
 	}
 	return onehot, xref
 }
 
+func homhet2maf(onehot [][]bool) float64 {
+	if len(onehot[0]) == 0 {
+		return 0
+	}
+	n := 0
+	for i := range onehot[0] {
+		if onehot[0][i] {
+			// hom
+			n += 2
+		} else if onehot[1][i] {
+			// het
+			n += 1
+		}
+	}
+	return float64(n) / float64(len(onehot[0])*2)
+}
+
 // convert a []onehotXref with length N to a numpy-style []int32
 // matrix with N columns, one row per field of onehotXref struct.
 //
@@ -1611,7 +1715,7 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
 // P-value row contains 1000000x actual p-value.
 func onehotXref2int32(xrefs []onehotXref) []int32 {
 	xcols := len(xrefs)
-	xdata := make([]int32, 5*xcols)
+	xdata := make([]int32, 6*xcols)
 	for i, xref := range xrefs {
 		xdata[i] = int32(xref.tag)
 		xdata[xcols+i] = int32(xref.variant)
@@ -1620,6 +1724,7 @@ func onehotXref2int32(xrefs []onehotXref) []int32 {
 		}
 		xdata[xcols*3+i] = int32(xref.pvalue * 1000000)
 		xdata[xcols*4+i] = int32(-math.Log10(xref.pvalue) * 1000000)
+		xdata[xcols*5+i] = int32(xref.maf * 1000000)
 	}
 	return xdata
 }