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         refpos := 0
 322         variantAt := map[int][]hgvs.Variant{} // variantAt[chromOffset][genomeIndex*2+phase]
 323         for refstep, libref := range reftiles {
 324                 refseq := tilelib.TileVariantSequence(libref)
 325                 tagcoverage := 0 // number of times the start tag was found in genomes -- max is len(cgs)*2
 326                 for cgidx, cg := range cgs {
 327                         for phase := 0; phase < 2; phase++ {
 328                                 if len(cg.Variants) <= int(libref.Tag)*2+phase {
 329                                         continue
 330                                 }
 331                                 variant := cg.Variants[int(libref.Tag)*2+phase]
 332                                 if variant == 0 {
 333                                         continue
 334                                 }
 335                                 tagcoverage++
 336                                 if variant == libref.Variant {
 337                                         continue
 338                                 }
 339                                 genomeseq := tilelib.TileVariantSequence(tileLibRef{Tag: libref.Tag, Variant: variant})
 340                                 if len(genomeseq) == 0 {
 341                                         // Hash is known but sequence
 342                                         // is not, e.g., retainNoCalls
 343                                         // was false during import
 344                                         continue
 345                                 }
 346                                 if len(genomeseq) > cmd.maxTileSize {
 347                                         continue
 348                                 }
 349                                 refSequence := refseq
 350                                 // If needed, extend the reference
 351                                 // sequence up to the tag at the end
 352                                 // of the genomeseq sequence.
 353                                 refstepend := refstep + 1
 354                                 for refstepend < len(reftiles) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genomeseq[len(genomeseq)-taglen:]) && len(refSequence) <= cmd.maxTileSize {
 355                                         if &refSequence[0] == &refseq[0] {
 356                                                 refSequence = append([]byte(nil), refSequence...)
 357                                         }
 358                                         refSequence = append(refSequence, tilelib.TileVariantSequence(reftiles[refstepend])...)
 359                                         refstepend++
 360                                 }
 361                                 // (TODO: handle no-calls)
 362                                 vars, _ := hgvs.Diff(strings.ToUpper(string(refSequence)), strings.ToUpper(string(genomeseq)), time.Second)
 363                                 for _, v := range vars {
 364                                         if cmd.outputFormat.PadLeft {
 365                                                 v = v.PadLeft()
 366                                         }
 367                                         v.Position += refpos
 368                                         varslice := variantAt[v.Position]
 369                                         if varslice == nil {
 370                                                 varslice = make([]hgvs.Variant, len(cgs)*2)
 371                                                 variantAt[v.Position] = varslice
 372                                         }
 373                                         varslice[cgidx*2+phase] = v
 374                                 }
 375                         }
 376                 }
 377                 refpos += len(refseq) - taglen
 378
 379                 // Flush entries from variantAt that are behind
 380                 // refpos. Flush all entries if this is the last
 381                 // reftile of the path/chromosome.
 382                 var flushpos []int
 383                 lastrefstep := refstep == len(reftiles)-1
 384                 for pos := range variantAt {
 385                         if lastrefstep || pos <= refpos {
 386                                 flushpos = append(flushpos, pos)
 387                         }
 388                 }
 389                 sort.Slice(flushpos, func(i, j int) bool { return flushpos[i] < flushpos[j] })
 390                 for _, pos := range flushpos {
 391                         varslice := variantAt[pos]
 392                         delete(variantAt, pos)
 393                         for i := range varslice {
 394                                 if varslice[i].Position == 0 {
 395                                         varslice[i].Position = pos
 396                                 }
 397                         }
 398                         cmd.outputFormat.Print(outw, seqname, varslice)
 399                 }
 400                 if bedw != nil && len(refseq) > 0 {
 401                         tilestart := refpos - len(refseq) + taglen
 402                         tileend := refpos
 403                         if !lastrefstep {
 404                                 tileend += taglen
 405                         }
 406                         thickstart := tilestart + taglen
 407                         if refstep == 0 {
 408                                 thickstart = 0
 409                         }
 410                         thickend := refpos
 411
 412                         // coverage score, 0 to 1000
 413                         score := 1000
 414                         if len(cgs) > 0 {
 415                                 score = 1000 * tagcoverage / len(cgs) / 2
 416                         }
 417
 418                         fmt.Fprintf(bedw, "%s %d %d %d %d . %d %d\n",
 419                                 seqname, tilestart, tileend,
 420                                 libref.Tag,
 421                                 score,
 422                                 thickstart, thickend)
 423                 }
 424         }
 425 }
 426
 427 func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
 428         refs := map[string]map[string]int{}
 429         for _, v := range varslice {
 430                 if v.Ref == "" && v.New == "" {
 431                         continue
 432                 }
 433                 alts := refs[v.Ref]
 434                 if alts == nil {
 435                         alts = map[string]int{}
 436                         refs[v.Ref] = alts
 437                 }
 438                 alts[v.New] = 0
 439         }
 440         for ref, alts := range refs {
 441                 var altslice []string
 442                 for alt := range alts {
 443                         altslice = append(altslice, alt)
 444                 }
 445                 sort.Strings(altslice)
 446                 for i, a := range altslice {
 447                         alts[a] = i + 1
 448                 }
 449                 fmt.Fprintf(out, "%s\t%d\t%s\t%s", seqname, varslice[0].Position, ref, strings.Join(altslice, ","))
 450                 for i := 0; i < len(varslice); i += 2 {
 451                         v1, v2 := varslice[i], varslice[i+1]
 452                         a1, a2 := alts[v1.New], alts[v2.New]
 453                         if v1.Ref != ref {
 454                                 a1 = 0
 455                         }
 456                         if v2.Ref != ref {
 457                                 a2 = 0
 458                         }
 459                         fmt.Fprintf(out, "\t%d/%d", a1, a2)
 460                 }
 461                 out.Write([]byte{'\n'})
 462         }
 463 }
 464
 465 func printHGVS(out io.Writer, seqname string, varslice []hgvs.Variant) {
 466         for i := 0; i < len(varslice)/2; i++ {
 467                 if i > 0 {
 468                         out.Write([]byte{'\t'})
 469                 }
 470                 var1, var2 := varslice[i*2], varslice[i*2+1]
 471                 if var1 == var2 {
 472                         if var1.Ref == var1.New {
 473                                 out.Write([]byte{'.'})
 474                         } else {
 475                                 fmt.Fprintf(out, "%s:g.%s", seqname, var1.String())
 476                         }
 477                 } else {
 478                         fmt.Fprintf(out, "%s:g.[%s];[%s]", seqname, var1.String(), var2.String())
 479                 }
 480         }
 481         out.Write([]byte{'\n'})
 482 }
 483
 484 func printHGVSOneHot(out io.Writer, seqname string, varslice []hgvs.Variant) {
 485         vars := map[hgvs.Variant]bool{}
 486         for _, v := range varslice {
 487                 if v.Ref != v.New {
 488                         vars[v] = true
 489                 }
 490         }
 491
 492         // sort variants to ensure output is deterministic
 493         sorted := make([]hgvs.Variant, 0, len(vars))
 494         for v := range vars {
 495                 sorted = append(sorted, v)
 496         }
 497         sort.Slice(sorted, func(a, b int) bool { return hgvs.Less(sorted[a], sorted[b]) })
 498
 499         for _, v := range sorted {
 500                 fmt.Fprintf(out, "%s.%s", seqname, v.String())
 501                 for i := 0; i < len(varslice); i += 2 {
 502                         if varslice[i] == v || varslice[i+1] == v {
 503                                 out.Write([]byte("\t1"))
 504                         } else {
 505                                 out.Write([]byte("\t0"))
 506                         }
 507                 }
 508                 out.Write([]byte{'\n'})
 509         }
 510 }