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