choosesamples.go

   1 // Copyright (C) The Lightning Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package lightning
   6
   7 import (
   8         "bytes"
   9         "errors"
  10         "flag"
  11         "fmt"
  12         "io"
  13         "math/rand"
  14         "net/http"
  15         _ "net/http/pprof"
  16         "os"
  17         "regexp"
  18         "sort"
  19         "strings"
  20
  21         "git.arvados.org/arvados.git/sdk/go/arvados"
  22         log "github.com/sirupsen/logrus"
  23 )
  24
  25 type chooseSamples struct {
  26         filter filter
  27 }
  28
  29 func (cmd *chooseSamples) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  30         err := cmd.run(prog, args, stdin, stdout, stderr)
  31         if err != nil {
  32                 fmt.Fprintf(stderr, "%s\n", err)
  33                 return 1
  34         }
  35         return 0
  36 }
  37
  38 func (cmd *chooseSamples) run(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) error {
  39         flags := flag.NewFlagSet("", flag.ContinueOnError)
  40         flags.SetOutput(stderr)
  41         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  42         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  43         projectUUID := flags.String("project", "", "project `UUID` for output data")
  44         priority := flags.Int("priority", 500, "container request priority")
  45         inputDir := flags.String("input-dir", "./in", "input `directory`")
  46         outputDir := flags.String("output-dir", "./out", "output `directory`")
  47         trainingSetSize := flags.Float64("training-set-size", 0.8, "number (or proportion, if <=1) of eligible samples to assign to the training set")
  48         caseControlFilename := flags.String("case-control-file", "", "tsv file or directory indicating cases and controls (if directory, all .tsv files will be read)")
  49         caseControlColumn := flags.String("case-control-column", "", "name of case/control column in case-control files (value must be 0 for control, 1 for case)")
  50         randSeed := flags.Int64("random-seed", 0, "PRNG seed")
  51         cmd.filter.Flags(flags)
  52         err := flags.Parse(args)
  53         if err == flag.ErrHelp {
  54                 return nil
  55         } else if err != nil {
  56                 return err
  57         }
  58         if *caseControlFilename == "" {
  59                 return errors.New("must provide -case-control-file")
  60         }
  61         if *caseControlColumn == "" {
  62                 return errors.New("must provide -case-control-column")
  63         }
  64
  65         if *pprof != "" {
  66                 go func() {
  67                         log.Println(http.ListenAndServe(*pprof, nil))
  68                 }()
  69         }
  70
  71         if !*runlocal {
  72                 runner := arvadosContainerRunner{
  73                         Name:        "lightning choose-samples",
  74                         Client:      arvados.NewClientFromEnv(),
  75                         ProjectUUID: *projectUUID,
  76                         RAM:         16000000000,
  77                         VCPUs:       4,
  78                         Priority:    *priority,
  79                         KeepCache:   2,
  80                         APIAccess:   true,
  81                 }
  82                 err = runner.TranslatePaths(inputDir, caseControlFilename)
  83                 if err != nil {
  84                         return err
  85                 }
  86                 runner.Args = []string{"choose-samples", "-local=true",
  87                         "-pprof=:6060",
  88                         "-input-dir=" + *inputDir,
  89                         "-output-dir=/mnt/output",
  90                         "-case-control-file=" + *caseControlFilename,
  91                         "-case-control-column=" + *caseControlColumn,
  92                         "-training-set-size=" + fmt.Sprintf("%f", *trainingSetSize),
  93                         "-random-seed=" + fmt.Sprintf("%d", *randSeed),
  94                 }
  95                 runner.Args = append(runner.Args, cmd.filter.Args()...)
  96                 var output string
  97                 output, err = runner.Run()
  98                 if err != nil {
  99                         return err
 100                 }
 101                 fmt.Fprintln(stdout, output)
 102                 return nil
 103         }
 104
 105         infiles, err := allFiles(*inputDir, matchGobFile)
 106         if err != nil {
 107                 return err
 108         }
 109         if len(infiles) == 0 {
 110                 err = fmt.Errorf("no input files found in %s", *inputDir)
 111                 return err
 112         }
 113         sort.Strings(infiles)
 114
 115         in0, err := open(infiles[0])
 116         if err != nil {
 117                 return err
 118         }
 119
 120         matchGenome, err := regexp.Compile(cmd.filter.MatchGenome)
 121         if err != nil {
 122                 err = fmt.Errorf("-match-genome: invalid regexp: %q", cmd.filter.MatchGenome)
 123                 return err
 124         }
 125
 126         var sampleIDs []string
 127         err = DecodeLibrary(in0, strings.HasSuffix(infiles[0], ".gz"), func(ent *LibraryEntry) error {
 128                 for _, cg := range ent.CompactGenomes {
 129                         if matchGenome.MatchString(cg.Name) {
 130                                 sampleIDs = append(sampleIDs, cg.Name)
 131                         }
 132                 }
 133                 return nil
 134         })
 135         if err != nil {
 136                 return err
 137         }
 138         in0.Close()
 139
 140         if len(sampleIDs) == 0 {
 141                 err = fmt.Errorf("no genomes found matching regexp %q", cmd.filter.MatchGenome)
 142                 return err
 143         }
 144         sort.Strings(sampleIDs)
 145         caseControl, err := cmd.loadCaseControlFiles(*caseControlFilename, *caseControlColumn, sampleIDs)
 146         if err != nil {
 147                 return err
 148         }
 149         if len(caseControl) == 0 {
 150                 err = fmt.Errorf("fatal: 0 cases, 0 controls, nothing to do")
 151                 return err
 152         }
 153
 154         var trainingSet, validationSet []int
 155         for i := range caseControl {
 156                 trainingSet = append(trainingSet, i)
 157         }
 158         sort.Ints(trainingSet)
 159         wantlen := int(*trainingSetSize)
 160         if *trainingSetSize <= 1 {
 161                 wantlen = int(*trainingSetSize * float64(len(trainingSet)))
 162         }
 163         randsrc := rand.NewSource(*randSeed)
 164         for tslen := len(trainingSet); tslen > wantlen; {
 165                 i := int(randsrc.Int63()) % tslen
 166                 validationSet = append(validationSet, trainingSet[i])
 167                 tslen--
 168                 trainingSet[i] = trainingSet[tslen]
 169                 trainingSet = trainingSet[:tslen]
 170         }
 171         sort.Ints(trainingSet)
 172         sort.Ints(validationSet)
 173
 174         samplesFilename := *outputDir + "/samples.csv"
 175         log.Infof("writing sample metadata to %s", samplesFilename)
 176         var f *os.File
 177         f, err = os.Create(samplesFilename)
 178         if err != nil {
 179                 return err
 180         }
 181         defer f.Close()
 182         _, err = fmt.Fprint(f, "Index,SampleID,CaseControl,TrainingValidation\n")
 183         if err != nil {
 184                 return err
 185         }
 186         tsi := 0 // next idx in training set
 187         vsi := 0 // next idx in validation set
 188         for i, name := range sampleIDs {
 189                 var cc, tv string
 190                 if len(trainingSet) > tsi && trainingSet[tsi] == i {
 191                         tv = "1"
 192                         tsi++
 193                         if caseControl[i] {
 194                                 cc = "1"
 195                         } else {
 196                                 cc = "0"
 197                         }
 198                 } else if len(validationSet) > vsi && validationSet[vsi] == i {
 199                         tv = "0"
 200                         vsi++
 201                         if caseControl[i] {
 202                                 cc = "1"
 203                         } else {
 204                                 cc = "0"
 205                         }
 206                 }
 207                 _, err = fmt.Fprintf(f, "%d,%s,%s,%s\n", i, trimFilenameForLabel(name), cc, tv)
 208                 if err != nil {
 209                         err = fmt.Errorf("write %s: %w", samplesFilename, err)
 210                         return err
 211                 }
 212         }
 213         err = f.Close()
 214         if err != nil {
 215                 err = fmt.Errorf("close %s: %w", samplesFilename, err)
 216                 return err
 217         }
 218         return nil
 219 }
 220
 221 // Read case/control file(s). Returned map m has m[i]==true if
 222 // sampleIDs[i] is case, m[i]==false if sampleIDs[i] is control.
 223 func (cmd *chooseSamples) loadCaseControlFiles(path, colname string, sampleIDs []string) (map[int]bool, error) {
 224         infiles, err := allFiles(path, nil)
 225         if err != nil {
 226                 return nil, err
 227         }
 228         // index in sampleIDs => case(true) / control(false)
 229         cc := map[int]bool{}
 230         // index in sampleIDs => true if matched by multiple patterns in case/control files
 231         dup := map[int]bool{}
 232         for _, infile := range infiles {
 233                 f, err := open(infile)
 234                 if err != nil {
 235                         return nil, err
 236                 }
 237                 buf, err := io.ReadAll(f)
 238                 f.Close()
 239                 if err != nil {
 240                         return nil, err
 241                 }
 242                 ccCol := -1
 243                 for _, tsv := range bytes.Split(buf, []byte{'\n'}) {
 244                         if len(tsv) == 0 {
 245                                 continue
 246                         }
 247                         split := strings.Split(string(tsv), "\t")
 248                         if ccCol < 0 {
 249                                 // header row
 250                                 for col, name := range split {
 251                                         if name == colname {
 252                                                 ccCol = col
 253                                                 break
 254                                         }
 255                                 }
 256                                 if ccCol < 0 {
 257                                         return nil, fmt.Errorf("%s: no column named %q in header row %q", infile, colname, tsv)
 258                                 }
 259                                 continue
 260                         }
 261                         if len(split) <= ccCol {
 262                                 continue
 263                         }
 264                         pattern := split[0]
 265                         found := -1
 266                         for i, name := range sampleIDs {
 267                                 if strings.Contains(name, pattern) {
 268                                         if found >= 0 {
 269                                                 log.Warnf("pattern %q in %s matches multiple sample IDs (%q, %q)", pattern, infile, sampleIDs[found], name)
 270                                         }
 271                                         if dup[i] {
 272                                                 continue
 273                                         } else if _, ok := cc[i]; ok {
 274                                                 log.Warnf("multiple patterns match sample ID %q, omitting from cases/controls", name)
 275                                                 dup[i] = true
 276                                                 delete(cc, i)
 277                                                 continue
 278                                         }
 279                                         found = i
 280                                         if split[ccCol] == "0" {
 281                                                 cc[found] = false
 282                                         }
 283                                         if split[ccCol] == "1" {
 284                                                 cc[found] = true
 285                                         }
 286                                 }
 287                         }
 288                         if found < 0 {
 289                                 log.Warnf("pattern %q in %s does not match any genome IDs", pattern, infile)
 290                                 continue
 291                         }
 292                 }
 293         }
 294         return cc, nil
 295 }