export.go

   1 package lightning
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "context"
   7         "errors"
   8         "flag"
   9         "fmt"
  10         "io"
  11         "net/http"
  12         _ "net/http/pprof"
  13         "os"
  14         "path/filepath"
  15         "runtime"
  16         "sort"
  17         "strconv"
  18         "strings"
  19         "sync"
  20         "time"
  21
  22         "git.arvados.org/arvados.git/sdk/go/arvados"
  23         "github.com/arvados/lightning/hgvs"
  24         log "github.com/sirupsen/logrus"
  25 )
  26
  27 type outputFormat struct {
  28         Filename string
  29         Head     func(out io.Writer, cgs []CompactGenome)
  30         Print    func(out io.Writer, seqname string, varslice []hgvs.Variant)
  31         PadLeft  bool
  32 }
  33
  34 var (
  35         outputFormats = map[string]outputFormat{
  36                 "hgvs-onehot": outputFormatHGVSOneHot,
  37                 "hgvs":        outputFormatHGVS,
  38                 "pvcf":        outputFormatPVCF,
  39                 "vcf":         outputFormatVCF,
  40         }
  41         outputFormatHGVS       = outputFormat{Filename: "out.csv", Head: headNone, Print: printHGVS}
  42         outputFormatHGVSOneHot = outputFormat{Filename: "out.csv", Head: headNone, Print: printHGVSOneHot}
  43         outputFormatPVCF       = outputFormat{Filename: "out.vcf", Head: headPVCF, Print: printPVCF, PadLeft: true}
  44         outputFormatVCF        = outputFormat{Filename: "out.vcf", Head: headVCF, Print: printVCF, PadLeft: true}
  45         headNone               = func(io.Writer, []CompactGenome) {}
  46 )
  47
  48 type exporter struct {
  49         outputFormat   outputFormat
  50         outputPerChrom bool
  51         maxTileSize    int
  52 }
  53
  54 func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  55         var err error
  56         defer func() {
  57                 if err != nil {
  58                         fmt.Fprintf(stderr, "%s\n", err)
  59                 }
  60         }()
  61         flags := flag.NewFlagSet("", flag.ContinueOnError)
  62         flags.SetOutput(stderr)
  63         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  64         pprofdir := flags.String("pprof-dir", "", "write Go profile data to `directory` periodically")
  65         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  66         projectUUID := flags.String("project", "", "project `UUID` for output data")
  67         priority := flags.Int("priority", 500, "container request priority")
  68         refname := flags.String("ref", "", "reference genome `name`")
  69         inputDir := flags.String("input-dir", ".", "input `directory`")
  70         outputDir := flags.String("output-dir", ".", "output `directory`")
  71         outputFormatStr := flags.String("output-format", "hgvs", "output `format`: hgvs, pvcf, or vcf")
  72         outputBed := flags.String("output-bed", "", "also output bed `file`")
  73         flags.BoolVar(&cmd.outputPerChrom, "output-per-chromosome", true, "output one file per chromosome")
  74         labelsFilename := flags.String("output-labels", "", "also output genome labels csv `file`")
  75         flags.IntVar(&cmd.maxTileSize, "max-tile-size", 50000, "don't try to make annotations for tiles bigger than given `size`")
  76         err = flags.Parse(args)
  77         if err == flag.ErrHelp {
  78                 err = nil
  79                 return 0
  80         } else if err != nil {
  81                 return 2
  82         }
  83         if flag.NArg() > 0 {
  84                 err = fmt.Errorf("extra unparsed command line arguments: %q", flag.Args())
  85                 return 2
  86         }
  87
  88         if f, ok := outputFormats[*outputFormatStr]; !ok {
  89                 err = fmt.Errorf("invalid output format %q", *outputFormatStr)
  90                 return 2
  91         } else {
  92                 cmd.outputFormat = f
  93         }
  94
  95         if *pprof != "" {
  96                 go func() {
  97                         log.Println(http.ListenAndServe(*pprof, nil))
  98                 }()
  99         }
 100         if *pprofdir != "" {
 101                 go writeProfilesPeriodically(*pprofdir)
 102         }
 103
 104         if !*runlocal {
 105                 if *outputDir != "." {
 106                         err = errors.New("cannot specify output directory in container mode: not implemented")
 107                         return 1
 108                 }
 109                 runner := arvadosContainerRunner{
 110                         Name:        "lightning export",
 111                         Client:      arvados.NewClientFromEnv(),
 112                         ProjectUUID: *projectUUID,
 113                         RAM:         700000000000,
 114                         VCPUs:       96,
 115                         Priority:    *priority,
 116                         APIAccess:   true,
 117                 }
 118                 err = runner.TranslatePaths(inputDir)
 119                 if err != nil {
 120                         return 1
 121                 }
 122                 if *outputBed != "" {
 123                         if strings.Contains(*outputBed, "/") {
 124                                 err = fmt.Errorf("cannot use -output-bed filename %q containing '/' char", *outputBed)
 125                                 return 1
 126                         }
 127                         *outputBed = "/mnt/output/" + *outputBed
 128                 }
 129                 runner.Args = []string{"export", "-local=true",
 130                         "-pprof", ":6000",
 131                         "-pprof-dir", "/mnt/output",
 132                         "-ref", *refname,
 133                         "-output-format", *outputFormatStr,
 134                         "-output-bed", *outputBed,
 135                         "-output-labels", "/mnt/output/labels.csv",
 136                         "-output-per-chromosome=" + fmt.Sprintf("%v", cmd.outputPerChrom),
 137                         "-max-tile-size", fmt.Sprintf("%d", cmd.maxTileSize),
 138                         "-input-dir", *inputDir,
 139                         "-output-dir", "/mnt/output",
 140                 }
 141                 var output string
 142                 output, err = runner.Run()
 143                 if err != nil {
 144                         return 1
 145                 }
 146                 fmt.Fprintln(stdout, output)
 147                 return 0
 148         }
 149
 150         var cgs []CompactGenome
 151         tilelib := &tileLibrary{
 152                 retainNoCalls:       true,
 153                 retainTileSequences: true,
 154                 compactGenomes:      map[string][]tileVariantID{},
 155         }
 156         err = tilelib.LoadDir(context.Background(), *inputDir, nil)
 157         if err != nil {
 158                 return 1
 159         }
 160
 161         refseq, ok := tilelib.refseqs[*refname]
 162         if !ok {
 163                 err = fmt.Errorf("reference name %q not found in input; have %v", *refname, func() (names []string) {
 164                         for name := range tilelib.refseqs {
 165                                 names = append(names, name)
 166                         }
 167                         return
 168                 }())
 169                 return 1
 170         }
 171
 172         names := cgnames(tilelib)
 173         for _, name := range names {
 174                 cgs = append(cgs, CompactGenome{Name: name, Variants: tilelib.compactGenomes[name]})
 175         }
 176         if *labelsFilename != "" {
 177                 log.Infof("writing labels to %s", *labelsFilename)
 178                 var f *os.File
 179                 f, err = os.OpenFile(*labelsFilename, os.O_CREATE|os.O_WRONLY, 0777)
 180                 if err != nil {
 181                         return 1
 182                 }
 183                 defer f.Close()
 184                 for i, name := range names {
 185                         _, err = fmt.Fprintf(f, "%d,%q,%q\n", i, trimFilenameForLabel(name), cmd.outputFormat.Filename)
 186                         if err != nil {
 187                                 err = fmt.Errorf("write %s: %w", *labelsFilename, err)
 188                                 return 1
 189                         }
 190                 }
 191                 err = f.Close()
 192                 if err != nil {
 193                         err = fmt.Errorf("close %s: %w", *labelsFilename, err)
 194                         return 1
 195                 }
 196         }
 197
 198         var bedout io.Writer
 199         var bedfile *os.File
 200         var bedbufw *bufio.Writer
 201         if *outputBed != "" {
 202                 bedfile, err = os.OpenFile(*outputBed, os.O_CREATE|os.O_WRONLY, 0666)
 203                 if err != nil {
 204                         return 1
 205                 }
 206                 defer bedfile.Close()
 207                 bedbufw = bufio.NewWriterSize(bedfile, 16*1024*1024)
 208                 bedout = bedbufw
 209         }
 210
 211         err = cmd.export(*outputDir, bedout, tilelib, refseq, cgs)
 212         if err != nil {
 213                 return 1
 214         }
 215         if bedout != nil {
 216                 err = bedbufw.Flush()
 217                 if err != nil {
 218                         return 1
 219                 }
 220                 err = bedfile.Close()
 221                 if err != nil {
 222                         return 1
 223                 }
 224         }
 225         return 0
 226 }
 227
 228 func (cmd *exporter) export(outdir string, bedout io.Writer, tilelib *tileLibrary, refseq map[string][]tileLibRef, cgs []CompactGenome) error {
 229         var seqnames []string
 230         var missing []tileLibRef
 231         for seqname, librefs := range refseq {
 232                 seqnames = append(seqnames, seqname)
 233                 for _, libref := range librefs {
 234                         if libref.Variant != 0 && tilelib.TileVariantSequence(libref) == nil {
 235                                 missing = append(missing, libref)
 236                         }
 237                 }
 238         }
 239         sort.Strings(seqnames)
 240
 241         if len(missing) > 0 {
 242                 if limit := 100; len(missing) > limit {
 243                         log.Warnf("first %d missing tiles: %v", limit, missing[:limit])
 244                 } else {
 245                         log.Warnf("missing tiles: %v", missing)
 246                 }
 247                 return fmt.Errorf("%d needed tiles are missing from library", len(missing))
 248         }
 249
 250         outw := make([]io.WriteCloser, len(seqnames))
 251         bedw := make([]io.WriteCloser, len(seqnames))
 252
 253         var merges sync.WaitGroup
 254         merge := func(dst io.Writer, src []io.WriteCloser, label string) {
 255                 var mtx sync.Mutex
 256                 for i, seqname := range seqnames {
 257                         pr, pw := io.Pipe()
 258                         src[i] = pw
 259                         merges.Add(1)
 260                         seqname := seqname
 261                         go func() {
 262                                 defer merges.Done()
 263                                 log.Infof("writing %s %s", seqname, label)
 264                                 scanner := bufio.NewScanner(pr)
 265                                 for scanner.Scan() {
 266                                         mtx.Lock()
 267                                         dst.Write(scanner.Bytes())
 268                                         dst.Write([]byte{'\n'})
 269                                         mtx.Unlock()
 270                                 }
 271                                 log.Infof("writing %s %s done", seqname, label)
 272                         }()
 273                 }
 274         }
 275         if cmd.outputPerChrom {
 276                 for i, seqname := range seqnames {
 277                         f, err := os.OpenFile(filepath.Join(outdir, strings.Replace(cmd.outputFormat.Filename, ".", "."+seqname+".", 1)), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
 278                         if err != nil {
 279                                 return err
 280                         }
 281                         defer f.Close()
 282                         log.Infof("writing %q", f.Name())
 283                         cmd.outputFormat.Head(f, cgs)
 284                         outw[i] = f
 285                 }
 286         } else {
 287                 fnm := filepath.Join(outdir, cmd.outputFormat.Filename)
 288                 log.Infof("writing %q", fnm)
 289                 out, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
 290                 if err != nil {
 291                         return err
 292                 }
 293                 defer out.Close()
 294                 cmd.outputFormat.Head(out, cgs)
 295                 merge(out, outw, "output")
 296         }
 297         if bedout != nil {
 298                 merge(bedout, bedw, "bed")
 299         }
 300
 301         throttle := throttle{Max: runtime.NumCPU()}
 302         log.Infof("assembling %d sequences in %d goroutines", len(seqnames), throttle.Max)
 303         for seqidx, seqname := range seqnames {
 304                 seqidx, seqname := seqidx, seqname
 305                 outw := outw[seqidx]
 306                 bedw := bedw[seqidx]
 307                 throttle.Acquire()
 308                 go func() {
 309                         defer throttle.Release()
 310                         if bedw != nil {
 311                                 defer bedw.Close()
 312                         }
 313                         defer outw.Close()
 314                         outwb := bufio.NewWriterSize(outw, 8*1024*1024)
 315                         defer outwb.Flush()
 316                         cmd.exportSeq(outwb, bedw, tilelib.taglib.keylen, seqname, refseq[seqname], tilelib, cgs)
 317                 }()
 318         }
 319
 320         merges.Wait()
 321         throttle.Wait()
 322         return nil
 323 }
 324
 325 // Align genome tiles to reference tiles, write diffs to outw, and (if
 326 // bedw is not nil) write tile coverage to bedw.
 327 func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string, reftiles []tileLibRef, tilelib *tileLibrary, cgs []CompactGenome) {
 328         t0 := time.Now()
 329         progressbar := time.NewTicker(time.Minute)
 330         defer progressbar.Stop()
 331         var outmtx sync.Mutex
 332         defer outmtx.Lock()
 333         refpos := 0
 334         variantAt := map[int][]hgvs.Variant{} // variantAt[chromOffset][genomeIndex*2+phase]
 335         for refstep, libref := range reftiles {
 336                 select {
 337                 case <-progressbar.C:
 338                         var eta interface{}
 339                         if refstep > 0 {
 340                                 fin := t0.Add(time.Duration(float64(time.Now().Sub(t0)) * float64(len(reftiles)) / float64(refstep)))
 341                                 eta = fmt.Sprintf("%v (%v)", fin.Format(time.RFC3339), fin.Sub(time.Now()))
 342                         } else {
 343                                 eta = "N/A"
 344                         }
 345                         log.Printf("exportSeq: %s: refstep %d of %d, %.0f/s, ETA %v", seqname, refstep, len(reftiles), float64(refstep)/time.Now().Sub(t0).Seconds(), eta)
 346                 default:
 347                 }
 348                 diffs := map[tileLibRef][]hgvs.Variant{}
 349                 refseq := tilelib.TileVariantSequence(libref)
 350                 tagcoverage := 0 // number of times the start tag was found in genomes -- max is len(cgs)*2
 351                 for cgidx, cg := range cgs {
 352                         for phase := 0; phase < 2; phase++ {
 353                                 if len(cg.Variants) <= int(libref.Tag)*2+phase {
 354                                         continue
 355                                 }
 356                                 variant := cg.Variants[int(libref.Tag)*2+phase]
 357                                 if variant == 0 {
 358                                         continue
 359                                 }
 360                                 tagcoverage++
 361                                 if variant == libref.Variant {
 362                                         continue
 363                                 }
 364                                 glibref := tileLibRef{Tag: libref.Tag, Variant: variant}
 365                                 vars, ok := diffs[glibref]
 366                                 if !ok {
 367                                         genomeseq := tilelib.TileVariantSequence(glibref)
 368                                         if len(genomeseq) == 0 {
 369                                                 // Hash is known but sequence
 370                                                 // is not, e.g., retainNoCalls
 371                                                 // was false during import
 372                                                 continue
 373                                         }
 374                                         if len(genomeseq) > cmd.maxTileSize {
 375                                                 continue
 376                                         }
 377                                         refSequence := refseq
 378                                         // If needed, extend the
 379                                         // reference sequence up to
 380                                         // the tag at the end of the
 381                                         // genomeseq sequence.
 382                                         refstepend := refstep + 1
 383                                         for refstepend < len(reftiles) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genomeseq[len(genomeseq)-taglen:]) && len(refSequence) <= cmd.maxTileSize {
 384                                                 if &refSequence[0] == &refseq[0] {
 385                                                         refSequence = append([]byte(nil), refSequence...)
 386                                                 }
 387                                                 refSequence = append(refSequence, tilelib.TileVariantSequence(reftiles[refstepend])...)
 388                                                 refstepend++
 389                                         }
 390                                         // (TODO: handle no-calls)
 391                                         refstr := strings.ToUpper(string(refSequence))
 392                                         genomestr := strings.ToUpper(string(genomeseq))
 393                                         vars, _ = hgvs.Diff(refstr, genomestr, time.Second)
 394                                         diffs[glibref] = vars
 395                                 }
 396                                 for _, v := range vars {
 397                                         if cmd.outputFormat.PadLeft {
 398                                                 v = v.PadLeft()
 399                                         }
 400                                         v.Position += refpos
 401                                         varslice := variantAt[v.Position]
 402                                         if varslice == nil {
 403                                                 varslice = make([]hgvs.Variant, len(cgs)*2)
 404                                                 variantAt[v.Position] = varslice
 405                                         }
 406                                         varslice[cgidx*2+phase] = v
 407                                 }
 408                         }
 409                 }
 410                 refpos += len(refseq) - taglen
 411
 412                 // Flush entries from variantAt that are behind
 413                 // refpos. Flush all entries if this is the last
 414                 // reftile of the path/chromosome.
 415                 flushpos := make([]int, 0, len(variantAt))
 416                 lastrefstep := refstep == len(reftiles)-1
 417                 for pos := range variantAt {
 418                         if lastrefstep || pos <= refpos {
 419                                 flushpos = append(flushpos, pos)
 420                         }
 421                 }
 422                 sort.Slice(flushpos, func(i, j int) bool { return flushpos[i] < flushpos[j] })
 423                 flushvariants := make([][]hgvs.Variant, len(flushpos))
 424                 for i, pos := range flushpos {
 425                         varslice := variantAt[pos]
 426                         delete(variantAt, pos)
 427                         for i := range varslice {
 428                                 if varslice[i].Position == 0 {
 429                                         varslice[i].Position = pos
 430                                 }
 431                         }
 432                         flushvariants[i] = varslice
 433                 }
 434                 outmtx.Lock()
 435                 go func() {
 436                         defer outmtx.Unlock()
 437                         for _, varslice := range flushvariants {
 438                                 cmd.outputFormat.Print(outw, seqname, varslice)
 439                         }
 440                 }()
 441                 if bedw != nil && len(refseq) > 0 {
 442                         tilestart := refpos - len(refseq) + taglen
 443                         tileend := refpos
 444                         if !lastrefstep {
 445                                 tileend += taglen
 446                         }
 447                         thickstart := tilestart + taglen
 448                         if refstep == 0 {
 449                                 thickstart = 0
 450                         }
 451                         thickend := refpos
 452
 453                         // coverage score, 0 to 1000
 454                         score := 1000
 455                         if len(cgs) > 0 {
 456                                 score = 1000 * tagcoverage / len(cgs) / 2
 457                         }
 458
 459                         fmt.Fprintf(bedw, "%s %d %d %d %d . %d %d\n",
 460                                 seqname, tilestart, tileend,
 461                                 libref.Tag,
 462                                 score,
 463                                 thickstart, thickend)
 464                 }
 465         }
 466 }
 467
 468 func bucketVarsliceByRef(varslice []hgvs.Variant) map[string]map[string]int {
 469         byref := map[string]map[string]int{}
 470         for _, v := range varslice {
 471                 if v.Ref == "" && v.New == "" {
 472                         continue
 473                 }
 474                 alts := byref[v.Ref]
 475                 if alts == nil {
 476                         alts = map[string]int{}
 477                         byref[v.Ref] = alts
 478                 }
 479                 alts[v.New]++
 480         }
 481         return byref
 482 }
 483
 484 func headVCF(out io.Writer, cgs []CompactGenome) {
 485         fmt.Fprint(out, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n")
 486 }
 487
 488 func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
 489         for ref, alts := range bucketVarsliceByRef(varslice) {
 490                 altslice := make([]string, 0, len(alts))
 491                 for alt := range alts {
 492                         altslice = append(altslice, alt)
 493                 }
 494                 sort.Strings(altslice)
 495
 496                 info := "AC="
 497                 for i, a := range altslice {
 498                         if i > 0 {
 499                                 info += ","
 500                         }
 501                         info += strconv.Itoa(alts[a])
 502                 }
 503                 fmt.Fprintf(out, "%s\t%d\t.\t%s\t%s\t.\t.\t%s\n", seqname, varslice[0].Position, ref, strings.Join(altslice, ","), info)
 504         }
 505 }
 506
 507 func headPVCF(out io.Writer, cgs []CompactGenome) {
 508         fmt.Fprintln(out, `##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">`)
 509         fmt.Fprintf(out, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT")
 510         for _, cg := range cgs {
 511                 fmt.Fprintf(out, "\t%s", cg.Name)
 512         }
 513         fmt.Fprintf(out, "\n")
 514 }
 515
 516 func printPVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
 517         for ref, alts := range bucketVarsliceByRef(varslice) {
 518                 altslice := make([]string, 0, len(alts))
 519                 for alt := range alts {
 520                         altslice = append(altslice, alt)
 521                 }
 522                 sort.Strings(altslice)
 523                 for i, a := range altslice {
 524                         alts[a] = i + 1
 525                 }
 526                 fmt.Fprintf(out, "%s\t%d\t.\t%s\t%s\t.\t.\t.\tGT", seqname, varslice[0].Position, ref, strings.Join(altslice, ","))
 527                 for i := 0; i < len(varslice); i += 2 {
 528                         v1, v2 := varslice[i], varslice[i+1]
 529                         a1, a2 := alts[v1.New], alts[v2.New]
 530                         if v1.Ref != ref {
 531                                 // variant on allele 0 belongs on a
 532                                 // different output line -- same
 533                                 // chr,pos but different "ref" length
 534                                 a1 = 0
 535                         }
 536                         if v2.Ref != ref {
 537                                 a2 = 0
 538                         }
 539                         fmt.Fprintf(out, "\t%d/%d", a1, a2)
 540                 }
 541                 out.Write([]byte{'\n'})
 542         }
 543 }
 544
 545 func printHGVS(out io.Writer, seqname string, varslice []hgvs.Variant) {
 546         for i := 0; i < len(varslice)/2; i++ {
 547                 if i > 0 {
 548                         out.Write([]byte{'\t'})
 549                 }
 550                 var1, var2 := varslice[i*2], varslice[i*2+1]
 551                 if var1 == var2 {
 552                         if var1.Ref == var1.New {
 553                                 out.Write([]byte{'.'})
 554                         } else {
 555                                 fmt.Fprintf(out, "%s:g.%s", seqname, var1.String())
 556                         }
 557                 } else {
 558                         fmt.Fprintf(out, "%s:g.[%s];[%s]", seqname, var1.String(), var2.String())
 559                 }
 560         }
 561         out.Write([]byte{'\n'})
 562 }
 563
 564 func printHGVSOneHot(out io.Writer, seqname string, varslice []hgvs.Variant) {
 565         vars := map[hgvs.Variant]bool{}
 566         for _, v := range varslice {
 567                 if v.Ref != v.New {
 568                         vars[v] = true
 569                 }
 570         }
 571
 572         // sort variants to ensure output is deterministic
 573         sorted := make([]hgvs.Variant, 0, len(vars))
 574         for v := range vars {
 575                 sorted = append(sorted, v)
 576         }
 577         sort.Slice(sorted, func(a, b int) bool { return hgvs.Less(sorted[a], sorted[b]) })
 578
 579         for _, v := range sorted {
 580                 fmt.Fprintf(out, "%s.%s", seqname, v.String())
 581                 for i := 0; i < len(varslice); i += 2 {
 582                         if varslice[i] == v || varslice[i+1] == v {
 583                                 out.Write([]byte("\t1"))
 584                         } else {
 585                                 out.Write([]byte("\t0"))
 586                         }
 587                 }
 588                 out.Write([]byte{'\n'})
 589         }
 590 }