Merge branch '19997-glm'
[lightning.git] / slicenumpy.go
index a9fee14c8d692430370b94f8ea0f3b35ca62fa43..895c3c15fa2d91cfd7b153f65d6b78a9c59ad806 100644 (file)
@@ -8,6 +8,7 @@ import (
        "bufio"
        "bytes"
        "encoding/gob"
+       "encoding/json"
        "errors"
        "flag"
        "fmt"
@@ -39,16 +40,21 @@ 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
+       pcaComponents   int
+       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
+       pvalue          func(onehot []bool) float64
+       pvalueCallCount int64
 }
 
 func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
@@ -59,6 +65,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)
@@ -68,6 +75,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)")
@@ -78,14 +86,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")
-       pcaComponents := flags.Int("pca-components", 4, "number of PCA components")
+       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")
+       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.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 +101,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 +111,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)
@@ -117,8 +127,9 @@ 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, &cmd.chi2CaseControlFile)
+               err = runner.TranslatePaths(inputDir, regionsFilename, samplesFilename)
                if err != nil {
                        return err
                }
@@ -134,11 +145,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),
+                       "-pca-components=" + fmt.Sprintf("%d", cmd.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 +187,15 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                return err
        }
 
+       if *samplesFilename != "" {
+               cmd.samples, err = 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 +238,58 @@ 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.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
@@ -242,31 +301,31 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                cmd.minCoverage = int(math.Ceil(cmd.filter.MinCoverage * float64(len(cmd.cgnames))))
        }
 
-       {
-               labelsFilename := *outputDir + "/samples.csv"
-               log.Infof("writing labels to %s", labelsFilename)
-               var f *os.File
-               f, err = os.Create(labelsFilename)
+       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, name := range cmd.cgnames {
-                       cc := 0
-                       if cmd.chi2Cases != nil && cmd.chi2Cases[i] {
-                               cc = 1
-                       }
-                       _, err = fmt.Fprintf(f, "%d,%q,%d\n", i, trimFilenameForLabel(name), cc)
-                       if err != nil {
-                               err = fmt.Errorf("write %s: %w", labelsFilename, err)
-                               return err
-                       }
-               }
-               err = f.Close()
-               if err != nil {
-                       err = fmt.Errorf("close %s: %w", labelsFilename, err)
-                       return err
-               }
+       }
+
+       // 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")
@@ -450,7 +509,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
@@ -606,6 +665,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 {
@@ -1111,6 +1176,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
@@ -1125,21 +1201,30 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                        log.Printf("have %d one-hot cols", cols)
                        stride := 1
                        for *maxPCATiles > 0 && cols > *maxPCATiles*2 {
-                               cols = cols / 2
+                               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 := nlp.NewPCA(cmd.pcaComponents)
+                       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 +1257,20 @@ func (cmd *sliceNumpy) run(prog string, args []string, stdin io.Reader, stdout,
                                return err
                        }
                        log.Print("done")
+
+                       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++ {
+                                       cmd.samples[i].pcaComponents[i] = pca.At(i, c)
+                               }
+                       }
+                       log.Print("done")
+
+                       err = writeSampleInfo(cmd.samples, *outputDir)
+                       if err != nil {
+                               return err
+                       }
                }
        }
        if !*mergeOutput && !*onehotChunked && !*onehotSingle && !*onlyPCA {
@@ -1196,89 +1295,121 @@ 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 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
+               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)
                        }
-                       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
+                       return nil, fmt.Errorf("%s line %d: index: %s", samplesFilename, lineNum, 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)
                                }
-                       }
-                       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
+                               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" && len(split[2]) > 0, // fix errant 0s in input
+                       pcaComponents: pcaComponents,
+               })
        }
-       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++
-                       }
+       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)
                }
        }
-       log.Printf("%d cases, %d controls, %d neither (dropped)", ncases, len(cmd.cgnames)-ncases, len(allnames)-len(cmd.cgnames))
+       _, 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
 }
 
@@ -1472,18 +1603,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 +1639,12 @@ func (cmd *sliceNumpy) tv2homhet(cgs map[string]CompactGenome, maxv tileVariantI
                if col < 4 && !cmd.includeVariant1 {
                        continue
                }
-               p := pvalue(obs[col], cmd.chi2Cases)
+               atomic.AddInt64(&cmd.pvalueCallCount, 1)
+               p := cmd.pvalue(obs[col])
                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 +1655,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.
 //