Fix some tests.
[lightning.git] / slicenumpy.go
index e4dd766c84586ea691161cc57c0ea1498bf83592..34cd777458ab93d60505b5dfa19b3aaa0280dd32 100644 (file)
@@ -8,6 +8,7 @@ import (
        "bufio"
        "bytes"
        "encoding/gob"
+       "encoding/json"
        "errors"
        "flag"
        "fmt"
@@ -39,20 +40,24 @@ import (
 const annotationMaxTileSpan = 100
 
 type sliceNumpy struct {
-       filter          filter
-       threads         int
-       chi2Cases       []bool
-       chi2PValue      float64
-       pcaComponents   int
-       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 {
@@ -73,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)")
@@ -89,8 +95,11 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
        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 (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)
@@ -124,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 {
@@ -143,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", 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),
                }
@@ -188,9 +201,6 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                if err != nil {
                        return err
                }
-               if len(cmd.samples[0].pcaComponents) > 0 {
-                       cmd.pvalue = glmPvalueFunc(cmd.samples, cmd.pcaComponents)
-               }
        } else if *caseControlOnly {
                return fmt.Errorf("-case-control-only does not make sense without -samples")
        }
@@ -290,49 +300,41 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                        }
                }
        }
-       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 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 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 {
-                       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")
@@ -359,7 +361,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                                return err
                        }
                        foundthistag := false
-                       taglib.FindAll(tiledata[:len(tiledata)-1], func(tagid tagID, offset, _ int) {
+                       taglib.FindAll(bufio.NewReader(bytes.NewReader(tiledata[:len(tiledata)-1])), nil, func(tagid tagID, offset, _ int) {
                                if !foundthistag && tagid == libref.Tag {
                                        foundthistag = true
                                        return
@@ -516,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
@@ -532,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
@@ -554,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]
@@ -866,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 {
@@ -1183,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
@@ -1254,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 {
@@ -1378,13 +1372,60 @@ func loadSampleInfo(samplesFilename string) ([]sampleInfo, error) {
                        isCase:        split[2] == "1",
                        isControl:     split[2] == "0",
                        isTraining:    split[3] == "1",
-                       isValidation:  split[3] == "0",
+                       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
@@ -1533,6 +1574,7 @@ type onehotXref struct {
        variant tileVariantID
        hom     bool
        pvalue  float64
+       maf     float64
 }
 
 const onehotXrefSize = unsafe.Sizeof(onehotXref{})
@@ -1561,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 {
@@ -1604,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
@@ -1611,6 +1658,22 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
                if col < 4 && !cmd.includeVariant1 {
                        continue
                }
+               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
@@ -1621,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.
 //
@@ -1634,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)
@@ -1643,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
 }