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 == "") != (*caseControlColumn == "") {
  59                 return errors.New("must provide both -case-control-file and -case-control-column, or neither")
  60         }
  61
  62         if *pprof != "" {
  63                 go func() {
  64                         log.Println(http.ListenAndServe(*pprof, nil))
  65                 }()
  66         }
  67
  68         if !*runlocal {
  69                 runner := arvadosContainerRunner{
  70                         Name:        "lightning choose-samples",
  71                         Client:      arvados.NewClientFromEnv(),
  72                         ProjectUUID: *projectUUID,
  73                         RAM:         16000000000,
  74                         VCPUs:       4,
  75                         Priority:    *priority,
  76                         KeepCache:   2,
  77                         APIAccess:   true,
  78                 }
  79                 err = runner.TranslatePaths(inputDir, caseControlFilename)
  80                 if err != nil {
  81                         return err
  82                 }
  83                 runner.Args = []string{"choose-samples", "-local=true",
  84                         "-pprof=:6060",
  85                         "-input-dir=" + *inputDir,
  86                         "-output-dir=/mnt/output",
  87                         "-case-control-file=" + *caseControlFilename,
  88                         "-case-control-column=" + *caseControlColumn,
  89                         "-training-set-size=" + fmt.Sprintf("%f", *trainingSetSize),
  90                         "-random-seed=" + fmt.Sprintf("%d", *randSeed),
  91                 }
  92                 runner.Args = append(runner.Args, cmd.filter.Args()...)
  93                 var output string
  94                 output, err = runner.Run()
  95                 if err != nil {
  96                         return err
  97                 }
  98                 fmt.Fprintln(stdout, output)
  99                 return nil
 100         }
 101
 102         infiles, err := allFiles(*inputDir, matchGobFile)
 103         if err != nil {
 104                 return err
 105         }
 106         if len(infiles) == 0 {
 107                 err = fmt.Errorf("no input files found in %s", *inputDir)
 108                 return err
 109         }
 110         sort.Strings(infiles)
 111
 112         in0, err := open(infiles[0])
 113         if err != nil {
 114                 return err
 115         }
 116
 117         matchGenome, err := regexp.Compile(cmd.filter.MatchGenome)
 118         if err != nil {
 119                 err = fmt.Errorf("-match-genome: invalid regexp: %q", cmd.filter.MatchGenome)
 120                 return err
 121         }
 122
 123         var sampleIDs []string
 124         err = DecodeLibrary(in0, strings.HasSuffix(infiles[0], ".gz"), func(ent *LibraryEntry) error {
 125                 for _, cg := range ent.CompactGenomes {
 126                         if matchGenome.MatchString(cg.Name) {
 127                                 sampleIDs = append(sampleIDs, cg.Name)
 128                         }
 129                 }
 130                 return nil
 131         })
 132         if err != nil {
 133                 return err
 134         }
 135         in0.Close()
 136
 137         if len(sampleIDs) == 0 {
 138                 err = fmt.Errorf("no genomes found matching regexp %q", cmd.filter.MatchGenome)
 139                 return err
 140         }
 141         sort.Strings(sampleIDs)
 142         caseControl, err := cmd.loadCaseControlFiles(*caseControlFilename, *caseControlColumn, sampleIDs)
 143         if err != nil {
 144                 return err
 145         }
 146         if len(caseControl) == 0 {
 147                 err = fmt.Errorf("fatal: 0 cases, 0 controls, nothing to do")
 148                 return err
 149         }
 150
 151         var trainingSet, validationSet []int
 152         for i := range caseControl {
 153                 trainingSet = append(trainingSet, i)
 154         }
 155         sort.Ints(trainingSet)
 156         wantlen := int(*trainingSetSize)
 157         if *trainingSetSize <= 1 {
 158                 wantlen = int(*trainingSetSize * float64(len(trainingSet)))
 159         }
 160         randsrc := rand.NewSource(*randSeed)
 161         for tslen := len(trainingSet); tslen > wantlen; {
 162                 i := int(randsrc.Int63()) % tslen
 163                 validationSet = append(validationSet, trainingSet[i])
 164                 tslen--
 165                 trainingSet[i] = trainingSet[tslen]
 166                 trainingSet = trainingSet[:tslen]
 167         }
 168         sort.Ints(trainingSet)
 169         sort.Ints(validationSet)
 170
 171         samplesFilename := *outputDir + "/samples.csv"
 172         log.Infof("writing sample metadata to %s", samplesFilename)
 173         var f *os.File
 174         f, err = os.Create(samplesFilename)
 175         if err != nil {
 176                 return err
 177         }
 178         defer f.Close()
 179         _, err = fmt.Fprint(f, "Index,SampleID,CaseControl,TrainingValidation\n")
 180         if err != nil {
 181                 return err
 182         }
 183         tsi := 0 // next idx in training set
 184         vsi := 0 // next idx in validation set
 185         for i, name := range sampleIDs {
 186                 var cc, tv string
 187                 if len(trainingSet) > tsi && trainingSet[tsi] == i {
 188                         tv = "1"
 189                         tsi++
 190                         if caseControl[i] {
 191                                 cc = "1"
 192                         } else {
 193                                 cc = "0"
 194                         }
 195                 } else if len(validationSet) > vsi && validationSet[vsi] == i {
 196                         tv = "0"
 197                         vsi++
 198                         if caseControl[i] {
 199                                 cc = "1"
 200                         } else {
 201                                 cc = "0"
 202                         }
 203                 }
 204                 _, err = fmt.Fprintf(f, "%d,%s,%s,%s\n", i, trimFilenameForLabel(name), cc, tv)
 205                 if err != nil {
 206                         err = fmt.Errorf("write %s: %w", samplesFilename, err)
 207                         return err
 208                 }
 209         }
 210         err = f.Close()
 211         if err != nil {
 212                 err = fmt.Errorf("close %s: %w", samplesFilename, err)
 213                 return err
 214         }
 215         return nil
 216 }
 217
 218 // Read case/control file(s). Returned map m has m[i]==true if
 219 // sampleIDs[i] is case, m[i]==false if sampleIDs[i] is control.
 220 func (cmd *chooseSamples) loadCaseControlFiles(path, colname string, sampleIDs []string) (map[int]bool, error) {
 221         if path == "" {
 222                 // all samples are control group
 223                 cc := make(map[int]bool, len(sampleIDs))
 224                 for i := range sampleIDs {
 225                         cc[i] = false
 226                 }
 227                 return cc, nil
 228         }
 229         infiles, err := allFiles(path, nil)
 230         if err != nil {
 231                 return nil, err
 232         }
 233         // index in sampleIDs => case(true) / control(false)
 234         cc := map[int]bool{}
 235         // index in sampleIDs => true if matched by multiple patterns in case/control files
 236         dup := map[int]bool{}
 237         for _, infile := range infiles {
 238                 f, err := open(infile)
 239                 if err != nil {
 240                         return nil, err
 241                 }
 242                 buf, err := io.ReadAll(f)
 243                 f.Close()
 244                 if err != nil {
 245                         return nil, err
 246                 }
 247                 ccCol := -1
 248                 for _, tsv := range bytes.Split(buf, []byte{'\n'}) {
 249                         if len(tsv) == 0 {
 250                                 continue
 251                         }
 252                         split := strings.Split(string(tsv), "\t")
 253                         if ccCol < 0 {
 254                                 // header row
 255                                 for col, name := range split {
 256                                         if name == colname {
 257                                                 ccCol = col
 258                                                 break
 259                                         }
 260                                 }
 261                                 if ccCol < 0 {
 262                                         return nil, fmt.Errorf("%s: no column named %q in header row %q", infile, colname, tsv)
 263                                 }
 264                                 continue
 265                         }
 266                         if len(split) <= ccCol {
 267                                 continue
 268                         }
 269                         pattern := split[0]
 270                         found := -1
 271                         for i, name := range sampleIDs {
 272                                 if strings.Contains(name, pattern) {
 273                                         if found >= 0 {
 274                                                 log.Warnf("pattern %q in %s matches multiple sample IDs (%q, %q)", pattern, infile, sampleIDs[found], name)
 275                                         }
 276                                         if dup[i] {
 277                                                 continue
 278                                         } else if _, ok := cc[i]; ok {
 279                                                 log.Warnf("multiple patterns match sample ID %q, omitting from cases/controls", name)
 280                                                 dup[i] = true
 281                                                 delete(cc, i)
 282                                                 continue
 283                                         }
 284                                         found = i
 285                                         if split[ccCol] == "0" {
 286                                                 cc[found] = false
 287                                         }
 288                                         if split[ccCol] == "1" {
 289                                                 cc[found] = true
 290                                         }
 291                                 }
 292                         }
 293                         if found < 0 {
 294                                 log.Warnf("pattern %q in %s does not match any genome IDs", pattern, infile)
 295                                 continue
 296                         }
 297                 }
 298         }
 299         return cc, nil
 300 }