19566: glm one column at a time.
[lightning.git] / slicenumpy.go
index bef164d72071092b0f07b3ec233e91806d3e59b5..25fc466c116c4b9cc0a9508378ff33db91247c0b 100644 (file)
@@ -39,16 +39,18 @@ import (
 const annotationMaxTileSpan = 100
 
 type sliceNumpy struct {
-       filter                filter
-       threads               int
-       chi2CaseControlColumn string
-       chi2CaseControlFile   string
-       chi2Cases             []bool
-       chi2PValue            float64
-       minCoverage           int
-       cgnames               []string
-       includeVariant1       bool
-       debugTag              tagID
+       filter          filter
+       threads         int
+       chi2Cases       []bool
+       chi2PValue      float64
+       minCoverage     int
+       includeVariant1 bool
+       debugTag        tagID
+
+       cgnames         []string
+       samples         []sampleInfo
+       trainingSet     []int // samples index => training set index, or -1 if not in training set
+       trainingSetSize int
 }
 
 func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
@@ -59,6 +61,7 @@ func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, s
        }
        return 0
 }
+
 func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) error {
        flags := flag.NewFlagSet("", flag.ContinueOnError)
        flags.SetOutput(stderr)
@@ -78,14 +81,14 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
        hgvsChunked := flags.Bool("chunked-hgvs-matrix", false, "also generate hgvs-based matrix per chromosome")
        onehotSingle := flags.Bool("single-onehot", false, "generate one-hot tile-based matrix")
        onehotChunked := flags.Bool("chunked-onehot", false, "generate one-hot tile-based matrix per input chunk")
-       onlyPCA := flags.Bool("pca", false, "generate pca matrix")
+       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, "run principal component analysis, write components to pca.npy and samples.csv")
        pcaComponents := flags.Int("pca-components", 4, "number of PCA components")
        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.IntVar(&cmd.threads, "threads", 16, "number of memory-hungry assembly threads, and number of VCPUs to request for arvados container")
-       flags.StringVar(&cmd.chi2CaseControlFile, "chi2-case-control-file", "", "tsv file or directory indicating cases and controls for Χ² test (if directory, all .tsv files will be read)")
-       flags.StringVar(&cmd.chi2CaseControlColumn, "chi2-case-control-column", "", "name of case/control column in case-control files for Χ² test (value must be 0 for control, 1 for case)")
-       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.BoolVar(&cmd.includeVariant1, "include-variant-1", false, "include most common variant when building one-hot matrix")
        cmd.filter.Flags(flags)
        err := flags.Parse(args)
@@ -93,6 +96,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 != "" {
@@ -101,8 +106,8 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                }()
        }
 
-       if cmd.chi2PValue != 1 && (cmd.chi2CaseControlFile == "" || cmd.chi2CaseControlColumn == "") {
-               return fmt.Errorf("cannot use provided -chi2-p-value=%f because -chi2-case-control-file= or -chi2-case-control-column= value is empty", cmd.chi2PValue)
+       if cmd.chi2PValue != 1 && *samplesFilename == "" {
+               return fmt.Errorf("cannot use provided -chi2-p-value=%f because -samples= value is empty", cmd.chi2PValue)
        }
 
        cmd.debugTag = tagID(*debugTag)
@@ -118,7 +123,7 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                        KeepCache:   2,
                        APIAccess:   true,
                }
-               err = runner.TranslatePaths(inputDir, regionsFilename, &cmd.chi2CaseControlFile)
+               err = runner.TranslatePaths(inputDir, regionsFilename, samplesFilename)
                if err != nil {
                        return err
                }
@@ -134,11 +139,11 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                        "-chunked-hgvs-matrix=" + fmt.Sprintf("%v", *hgvsChunked),
                        "-single-onehot=" + fmt.Sprintf("%v", *onehotSingle),
                        "-chunked-onehot=" + fmt.Sprintf("%v", *onehotChunked),
+                       "-samples=" + *samplesFilename,
+                       "-case-control-only=" + fmt.Sprintf("%v", *caseControlOnly),
                        "-pca=" + fmt.Sprintf("%v", *onlyPCA),
                        "-pca-components=" + fmt.Sprintf("%d", *pcaComponents),
                        "-max-pca-tiles=" + fmt.Sprintf("%d", *maxPCATiles),
-                       "-chi2-case-control-file=" + cmd.chi2CaseControlFile,
-                       "-chi2-case-control-column=" + cmd.chi2CaseControlColumn,
                        "-chi2-p-value=" + fmt.Sprintf("%f", cmd.chi2PValue),
                        "-include-variant-1=" + fmt.Sprintf("%v", cmd.includeVariant1),
                        "-debug-tag=" + fmt.Sprintf("%d", cmd.debugTag),
@@ -176,6 +181,15 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                return err
        }
 
+       if *samplesFilename != "" {
+               cmd.samples, err = cmd.loadSampleInfo(*samplesFilename)
+               if err != nil {
+                       return err
+               }
+       } else if *caseControlOnly {
+               return fmt.Errorf("-case-control-only does not make sense without -samples")
+       }
+
        cmd.cgnames = nil
        var tagset [][]byte
        err = DecodeLibrary(in0, strings.HasSuffix(infiles[0], ".gz"), func(ent *LibraryEntry) error {
@@ -218,19 +232,53 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                return err
        }
        taglen := taglib.TagLen()
+       sort.Strings(cmd.cgnames)
 
        if len(cmd.cgnames) == 0 {
-               err = fmt.Errorf("no genomes found matching regexp %q", cmd.filter.MatchGenome)
-               return err
-       }
-       sort.Strings(cmd.cgnames)
-       err = cmd.useCaseControlFiles()
-       if err != nil {
-               return err
+               return fmt.Errorf("fatal: 0 matching samples in library, nothing to do")
        }
-       if len(cmd.cgnames) == 0 {
-               err = fmt.Errorf("fatal: 0 cases, 0 controls, nothing to do")
-               return err
+       cmd.trainingSet = make([]int, len(cmd.cgnames))
+       if *samplesFilename == "" {
+               cmd.trainingSetSize = len(cmd.cgnames)
+               for i, name := range cmd.cgnames {
+                       cmd.samples = append(cmd.samples, sampleInfo{
+                               id:         trimFilenameForLabel(name),
+                               isTraining: true,
+                       })
+                       cmd.trainingSet[i] = i
+               }
+       } else if len(cmd.cgnames) != len(cmd.samples) {
+               return fmt.Errorf("mismatched sample list: %d samples in library, %d in %s", len(cmd.cgnames), len(cmd.samples), *samplesFilename)
+       } else {
+               for i, name := range cmd.cgnames {
+                       if s := trimFilenameForLabel(name); s != cmd.samples[i].id {
+                               return fmt.Errorf("mismatched sample list: sample %d is %q in library, %q in %s", i, s, cmd.samples[i].id, *samplesFilename)
+                       }
+               }
+               if *caseControlOnly {
+                       for i := 0; i < len(cmd.samples); i++ {
+                               if !cmd.samples[i].isTraining && !cmd.samples[i].isValidation {
+                                       if i+1 < len(cmd.samples) {
+                                               copy(cmd.samples[i:], cmd.samples[i+1:])
+                                               copy(cmd.cgnames[i:], cmd.cgnames[i+1:])
+                                       }
+                                       cmd.samples = cmd.samples[:len(cmd.samples)-1]
+                                       cmd.cgnames = cmd.cgnames[:len(cmd.cgnames)-1]
+                                       i--
+                               }
+                       }
+               }
+               cmd.chi2Cases = nil
+               cmd.trainingSetSize = 0
+               for i := range cmd.cgnames {
+                       if cmd.samples[i].isTraining {
+                               cmd.trainingSet[i] = cmd.trainingSetSize
+                               cmd.trainingSetSize++
+                               cmd.chi2Cases = append(cmd.chi2Cases, cmd.samples[i].isCase)
+                       } else {
+                               cmd.trainingSet[i] = -1
+                       }
+               }
        }
        if cmd.filter.MinCoverage == 1 {
                // In the generic formula below, floating point
@@ -243,30 +291,38 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
        }
 
        {
-               labelsFilename := *outputDir + "/samples.csv"
-               log.Infof("writing labels to %s", labelsFilename)
+               samplesOutFilename := *outputDir + "/samples.csv"
+               log.Infof("writing sample metadata to %s", samplesOutFilename)
                var f *os.File
-               f, err = os.Create(labelsFilename)
+               f, err = os.Create(samplesOutFilename)
                if err != nil {
                        return err
                }
                defer f.Close()
-               for i, name := range cmd.cgnames {
-                       cc := 0
-                       if cmd.chi2Cases != nil && cmd.chi2Cases[i] {
-                               cc = 1
+               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,%q,%d\n", i, trimFilenameForLabel(name), cc)
+                       _, err = fmt.Fprintf(f, "%d,%s,%s,%s\n", i, si.id, cc, tv)
                        if err != nil {
-                               err = fmt.Errorf("write %s: %w", labelsFilename, err)
+                               err = fmt.Errorf("write %s: %w", samplesOutFilename, err)
                                return err
                        }
                }
                err = f.Close()
                if err != nil {
-                       err = fmt.Errorf("close %s: %w", labelsFilename, err)
+                       err = fmt.Errorf("close %s: %w", samplesOutFilename, err)
                        return err
                }
+               log.Print("done")
        }
 
        log.Info("indexing reference tiles")
@@ -606,6 +662,12 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                                        break
                                }
                                remap := variantRemap[tag-tagstart]
+                               if remap == nil {
+                                       // was not assigned above,
+                                       // because minCoverage
+                                       outcol++
+                                       continue
+                               }
                                maxv := tileVariantID(0)
                                for _, v := range remap {
                                        if maxv < v {
@@ -1128,18 +1190,27 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                                cols = (cols + 1) / 2
                                stride = stride * 2
                        }
-                       log.Printf("creating matrix: %d rows, %d cols, stride %d", len(cmd.cgnames), cols, stride)
-                       mtx := mat.NewDense(len(cmd.cgnames), cols, nil)
+                       if cols%2 == 1 {
+                               // we work with pairs of columns
+                               cols++
+                       }
+                       log.Printf("creating full matrix (%d rows) and training matrix (%d rows) with %d cols, stride %d", len(cmd.cgnames), cmd.trainingSetSize, cols, stride)
+                       mtxFull := mat.NewDense(len(cmd.cgnames), cols, nil)
+                       mtxTrain := mat.NewDense(cmd.trainingSetSize, cols, nil)
                        for i, c := range onehot[nzCount:] {
                                if int(c/2)%stride == 0 {
-                                       mtx.Set(int(onehot[i]), int(c/2)/stride*2+int(c)%2, 1)
+                                       outcol := int(c/2)/stride*2 + int(c)%2
+                                       mtxFull.Set(int(onehot[i]), outcol, 1)
+                                       if trainRow := cmd.trainingSet[int(onehot[i])]; trainRow >= 0 {
+                                               mtxTrain.Set(trainRow, outcol, 1)
+                                       }
                                }
                        }
                        log.Print("fitting")
                        transformer := nlp.NewPCA(*pcaComponents)
-                       transformer.Fit(mtx.T())
+                       transformer.Fit(mtxTrain.T())
                        log.Printf("transforming")
-                       pca, err := transformer.Transform(mtx.T())
+                       pca, err := transformer.Transform(mtxFull.T())
                        if err != nil {
                                return err
                        }
@@ -1172,6 +1243,43 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                                return err
                        }
                        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
+                               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
+                               }
+                       }
+                       err = f.Close()
+                       if err != nil {
+                               err = fmt.Errorf("close %s: %w", samplesOutFilename, err)
+                               return err
+                       }
+                       log.Print("done")
                }
        }
        if !*mergeOutput && !*onehotChunked && !*onehotSingle && !*onlyPCA {
@@ -1196,90 +1304,75 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                        return err
                }
        }
+
        return nil
 }
 
-// Read case/control files, remove non-case/control entries from
-// cmd.cgnames, and build cmd.chi2Cases.
-func (cmd *sliceNumpy) useCaseControlFiles() error {
-       if cmd.chi2CaseControlFile == "" {
-               return nil
+type sampleInfo struct {
+       id            string
+       isCase        bool
+       isControl     bool
+       isTraining    bool
+       isValidation  bool
+       pcaComponents []float64
+}
+
+// Read samples.csv file with case/control and training/validation
+// flags.
+func (cmd *sliceNumpy) loadSampleInfo(samplesFilename string) ([]sampleInfo, error) {
+       var si []sampleInfo
+       f, err := open(samplesFilename)
+       if err != nil {
+               return nil, err
        }
-       infiles, err := allFiles(cmd.chi2CaseControlFile, nil)
+       buf, err := io.ReadAll(f)
+       f.Close()
        if err != nil {
-               return err
+               return nil, err
        }
-       // index in cmd.cgnames => case(true) / control(false)
-       cc := map[int]bool{}
-       for _, infile := range infiles {
-               f, err := open(infile)
-               if err != nil {
-                       return err
+       lineNum := 0
+       for _, csv := range bytes.Split(buf, []byte{'\n'}) {
+               lineNum++
+               if len(csv) == 0 {
+                       continue
                }
-               buf, err := io.ReadAll(f)
-               f.Close()
-               if err != nil {
-                       return err
+               split := strings.Split(string(csv), ",")
+               if len(split) < 4 {
+                       return nil, fmt.Errorf("%d fields < 4 in %s line %d: %q", len(split), samplesFilename, lineNum, csv)
                }
-               ccCol := -1
-               for _, tsv := range bytes.Split(buf, []byte{'\n'}) {
-                       if len(tsv) == 0 {
-                               continue
-                       }
-                       split := strings.Split(string(tsv), "\t")
-                       if ccCol < 0 {
-                               // header row
-                               for col, name := range split {
-                                       if name == cmd.chi2CaseControlColumn {
-                                               ccCol = col
-                                               break
-                                       }
-                               }
-                               if ccCol < 0 {
-                                       return fmt.Errorf("%s: no column named %q in header row %q", infile, cmd.chi2CaseControlColumn, tsv)
-                               }
-                               continue
-                       }
-                       if len(split) <= ccCol {
-                               continue
-                       }
-                       pattern := split[0]
-                       found := -1
-                       for i, name := range cmd.cgnames {
-                               if strings.Contains(name, pattern) {
-                                       if found >= 0 {
-                                               log.Warnf("pattern %q in %s matches multiple genome IDs (%qs, %q)", pattern, infile, cmd.cgnames[found], name)
-                                       }
-                                       found = i
-                               }
-                       }
-                       if found < 0 {
-                               log.Warnf("pattern %q in %s does not match any genome IDs", pattern, infile)
-                               continue
-                       }
-                       if split[ccCol] == "0" {
-                               cc[found] = false
-                       }
-                       if split[ccCol] == "1" {
-                               cc[found] = true
+               if split[0] == "Index" && split[1] == "SampleID" && split[2] == "CaseControl" && split[3] == "TrainingValidation" {
+                       continue
+               }
+               idx, err := strconv.Atoi(split[0])
+               if err != nil {
+                       if lineNum == 1 {
+                               return nil, fmt.Errorf("header does not look right: %q", csv)
                        }
+                       return nil, fmt.Errorf("%s line %d: index: %s", samplesFilename, lineNum, err)
                }
-       }
-       allnames := cmd.cgnames
-       cmd.cgnames = nil
-       cmd.chi2Cases = nil
-       ncases := 0
-       for i, name := range allnames {
-               if cc, ok := cc[i]; ok {
-                       cmd.cgnames = append(cmd.cgnames, name)
-                       cmd.chi2Cases = append(cmd.chi2Cases, cc)
-                       if cc {
-                               ncases++
+               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",
+                       pcaComponents: pcaComponents,
+               })
        }
-       log.Printf("%d cases, %d controls, %d neither (dropped)", ncases, len(cmd.cgnames)-ncases, len(allnames)-len(cmd.cgnames))
-       return nil
+       return si, nil
 }
 
 func (cmd *sliceNumpy) filterHGVScolpair(colpair [2][]int8) bool {
@@ -1472,18 +1565,30 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
        if coverage < cmd.minCoverage {
                return nil, nil
        }
+       // "observed" array for p-value calculation (training set
+       // only)
        obs := make([][]bool, (maxv+1)*2) // 2 slices (hom + het) for each variant#
+       // one-hot output (all samples)
+       outcols := make([][]int8, (maxv+1)*2)
        for i := range obs {
-               obs[i] = make([]bool, len(cmd.cgnames))
+               obs[i] = make([]bool, cmd.trainingSetSize)
+               outcols[i] = make([]int8, len(cmd.cgnames))
        }
        for cgid, name := range cmd.cgnames {
+               tsid := cmd.trainingSet[cgid]
                cgvars := cgs[name].Variants[tagoffset*2:]
                tv0, tv1 := remap[cgvars[0]], remap[cgvars[1]]
                for v := tileVariantID(1); v <= maxv; v++ {
                        if tv0 == v && tv1 == v {
-                               obs[v*2][cgid] = true
+                               if tsid >= 0 {
+                                       obs[v*2][tsid] = true
+                               }
+                               outcols[v*2][cgid] = 1
                        } else if tv0 == v || tv1 == v {
-                               obs[v*2+1][cgid] = true
+                               if tsid >= 0 {
+                                       obs[v*2+1][tsid] = true
+                               }
+                               outcols[v*2+1][cgid] = 1
                        }
                }
        }
@@ -1496,11 +1601,16 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
                if col < 4 && !cmd.includeVariant1 {
                        continue
                }
-               p := pvalue(obs[col], cmd.chi2Cases)
+               var p float64
+               if len(cmd.samples[0].pcaComponents) > 0 {
+                       p = pvalueGLM(cmd.samples, obs[col])
+               } else {
+                       p = pvalue(obs[col], cmd.chi2Cases)
+               }
                if cmd.chi2PValue < 1 && !(p < cmd.chi2PValue) {
                        continue
                }
-               onehot = append(onehot, bool2int8(obs[col]))
+               onehot = append(onehot, outcols[col])
                xref = append(xref, onehotXref{
                        tag:     tag,
                        variant: tileVariantID(col >> 1),
@@ -1511,16 +1621,6 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
        return onehot, xref
 }
 
-func bool2int8(in []bool) []int8 {
-       out := make([]int8, len(in))
-       for i, v := range in {
-               if v {
-                       out[i] = 1
-               }
-       }
-       return out
-}
-
 // convert a []onehotXref with length N to a numpy-style []int32
 // matrix with N columns, one row per field of onehotXref struct.
 //