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