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:       64,
 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+"/export.csv")
 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 *os.File
 194         var bedbufw *bufio.Writer
 195         if *outputBed != "" {
 196                 bedout, err = os.OpenFile(*outputBed, os.O_CREATE|os.O_WRONLY, 0666)
 197                 if err != nil {
 198                         return 1
 199                 }
 200                 defer bedout.Close()
 201                 bedbufw = bufio.NewWriter(bedout)
 202         }
 203
 204         err = cmd.export(*outputDir, bedout, tilelib, refseq, cgs)
 205         if err != nil {
 206                 return 1
 207         }
 208         if bedout != nil {
 209                 err = bedbufw.Flush()
 210                 if err != nil {
 211                         return 1
 212                 }
 213                 err = bedout.Close()
 214                 if err != nil {
 215                         return 1
 216                 }
 217         }
 218         return 0
 219 }
 220
 221 func (cmd *exporter) export(outdir string, bedout io.Writer, tilelib *tileLibrary, refseq map[string][]tileLibRef, cgs []CompactGenome) error {
 222         var seqnames []string
 223         var missing []tileLibRef
 224         for seqname, librefs := range refseq {
 225                 seqnames = append(seqnames, seqname)
 226                 for _, libref := range librefs {
 227                         if libref.Variant != 0 && tilelib.TileVariantSequence(libref) == nil {
 228                                 missing = append(missing, libref)
 229                         }
 230                 }
 231         }
 232         sort.Strings(seqnames)
 233
 234         if len(missing) > 0 {
 235                 if limit := 100; len(missing) > limit {
 236                         log.Warnf("first %d missing tiles: %v", limit, missing[:limit])
 237                 } else {
 238                         log.Warnf("missing tiles: %v", missing)
 239                 }
 240                 return fmt.Errorf("%d needed tiles are missing from library", len(missing))
 241         }
 242
 243         outw := make([]io.WriteCloser, len(seqnames))
 244         bedw := make([]io.WriteCloser, len(seqnames))
 245
 246         var merges sync.WaitGroup
 247         merge := func(dst io.Writer, src []io.WriteCloser, label string) {
 248                 var mtx sync.Mutex
 249                 for i, seqname := range seqnames {
 250                         pr, pw := io.Pipe()
 251                         src[i] = pw
 252                         merges.Add(1)
 253                         seqname := seqname
 254                         go func() {
 255                                 defer merges.Done()
 256                                 log.Infof("writing %s %s", seqname, label)
 257                                 scanner := bufio.NewScanner(pr)
 258                                 for scanner.Scan() {
 259                                         mtx.Lock()
 260                                         dst.Write(scanner.Bytes())
 261                                         dst.Write([]byte{'\n'})
 262                                         mtx.Unlock()
 263                                 }
 264                                 log.Infof("writing %s %s done", seqname, label)
 265                         }()
 266                 }
 267         }
 268         if cmd.outputPerChrom {
 269                 for i, seqname := range seqnames {
 270                         f, err := os.OpenFile(filepath.Join(outdir, strings.Replace(cmd.outputFormat.Filename, ".", "."+seqname+".", 1)), os.O_CREATE|os.O_WRONLY, 0666)
 271                         if err != nil {
 272                                 return err
 273                         }
 274                         defer f.Close()
 275                         log.Infof("writing %q", f.Name())
 276                         outw[i] = f
 277                 }
 278         } else {
 279                 fnm := filepath.Join(outdir, cmd.outputFormat.Filename)
 280                 log.Infof("writing %q", fnm)
 281                 out, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY, 0666)
 282                 if err != nil {
 283                         return err
 284                 }
 285                 defer out.Close()
 286                 merge(out, outw, "output")
 287         }
 288         if bedout != nil {
 289                 merge(bedout, bedw, "bed")
 290         }
 291
 292         throttle := throttle{Max: runtime.NumCPU()}
 293         log.Infof("assembling %d sequences in %d goroutines", len(seqnames), throttle.Max)
 294         for seqidx, seqname := range seqnames {
 295                 seqidx, seqname := seqidx, seqname
 296                 outw := outw[seqidx]
 297                 bedw := bedw[seqidx]
 298                 throttle.Acquire()
 299                 go func() {
 300                         defer throttle.Release()
 301                         if bedw != nil {
 302                                 defer bedw.Close()
 303                         }
 304                         defer outw.Close()
 305                         outwb := bufio.NewWriter(outw)
 306                         defer outwb.Flush()
 307                         cmd.exportSeq(outwb, bedw, tilelib.taglib.keylen, seqname, refseq[seqname], tilelib, cgs)
 308                 }()
 309         }
 310
 311         merges.Wait()
 312         return nil
 313 }
 314
 315 // Align genome tiles to reference tiles, write diffs to outw, and (if
 316 // bedw is not nil) write tile coverage to bedw.
 317 func (cmd *exporter) exportSeq(outw, bedw io.Writer, taglen int, seqname string, reftiles []tileLibRef, tilelib *tileLibrary, cgs []CompactGenome) {
 318         refpos := 0
 319         variantAt := map[int][]hgvs.Variant{} // variantAt[chromOffset][genomeIndex*2+phase]
 320         for refstep, libref := range reftiles {
 321                 refseq := tilelib.TileVariantSequence(libref)
 322                 tagcoverage := 0 // number of times the start tag was found in genomes -- max is len(cgs)*2
 323                 for cgidx, cg := range cgs {
 324                         for phase := 0; phase < 2; phase++ {
 325                                 if len(cg.Variants) <= int(libref.Tag)*2+phase {
 326                                         continue
 327                                 }
 328                                 variant := cg.Variants[int(libref.Tag)*2+phase]
 329                                 if variant == 0 {
 330                                         continue
 331                                 }
 332                                 tagcoverage++
 333                                 if variant == libref.Variant {
 334                                         continue
 335                                 }
 336                                 genomeseq := tilelib.TileVariantSequence(tileLibRef{Tag: libref.Tag, Variant: variant})
 337                                 if len(genomeseq) == 0 {
 338                                         // Hash is known but sequence
 339                                         // is not, e.g., retainNoCalls
 340                                         // was false during import
 341                                         continue
 342                                 }
 343                                 if len(genomeseq) > cmd.maxTileSize {
 344                                         continue
 345                                 }
 346                                 refSequence := refseq
 347                                 // If needed, extend the reference
 348                                 // sequence up to the tag at the end
 349                                 // of the genomeseq sequence.
 350                                 refstepend := refstep + 1
 351                                 for refstepend < len(reftiles) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genomeseq[len(genomeseq)-taglen:]) && len(refSequence) <= cmd.maxTileSize {
 352                                         if &refSequence[0] == &refseq[0] {
 353                                                 refSequence = append([]byte(nil), refSequence...)
 354                                         }
 355                                         refSequence = append(refSequence, tilelib.TileVariantSequence(reftiles[refstepend])...)
 356                                         refstepend++
 357                                 }
 358                                 // (TODO: handle no-calls)
 359                                 vars, _ := hgvs.Diff(strings.ToUpper(string(refSequence)), strings.ToUpper(string(genomeseq)), time.Second)
 360                                 for _, v := range vars {
 361                                         if cmd.outputFormat.PadLeft {
 362                                                 v = v.PadLeft()
 363                                         }
 364                                         v.Position += refpos
 365                                         varslice := variantAt[v.Position]
 366                                         if varslice == nil {
 367                                                 varslice = make([]hgvs.Variant, len(cgs)*2)
 368                                                 variantAt[v.Position] = varslice
 369                                         }
 370                                         varslice[cgidx*2+phase] = v
 371                                 }
 372                         }
 373                 }
 374                 refpos += len(refseq) - taglen
 375
 376                 // Flush entries from variantAt that are behind
 377                 // refpos. Flush all entries if this is the last
 378                 // reftile of the path/chromosome.
 379                 var flushpos []int
 380                 lastrefstep := refstep == len(reftiles)-1
 381                 for pos := range variantAt {
 382                         if lastrefstep || pos <= refpos {
 383                                 flushpos = append(flushpos, pos)
 384                         }
 385                 }
 386                 sort.Slice(flushpos, func(i, j int) bool { return flushpos[i] < flushpos[j] })
 387                 for _, pos := range flushpos {
 388                         varslice := variantAt[pos]
 389                         delete(variantAt, pos)
 390                         for i := range varslice {
 391                                 if varslice[i].Position == 0 {
 392                                         varslice[i].Position = pos
 393                                 }
 394                         }
 395                         cmd.outputFormat.Print(outw, seqname, varslice)
 396                 }
 397                 if bedw != nil && len(refseq) > 0 {
 398                         tilestart := refpos - len(refseq) + taglen
 399                         tileend := refpos
 400                         if !lastrefstep {
 401                                 tileend += taglen
 402                         }
 403                         thickstart := tilestart + taglen
 404                         if refstep == 0 {
 405                                 thickstart = 0
 406                         }
 407                         thickend := refpos
 408
 409                         // coverage score, 0 to 1000
 410                         score := 1000
 411                         if len(cgs) > 0 {
 412                                 score = 1000 * tagcoverage / len(cgs) / 2
 413                         }
 414
 415                         fmt.Fprintf(bedw, "%s %d %d %d %d . %d %d\n",
 416                                 seqname, tilestart, tileend,
 417                                 libref.Tag,
 418                                 score,
 419                                 thickstart, thickend)
 420                 }
 421         }
 422 }
 423
 424 func printVCF(out io.Writer, seqname string, varslice []hgvs.Variant) {
 425         refs := map[string]map[string]int{}
 426         for _, v := range varslice {
 427                 if v.Ref == "" && v.New == "" {
 428                         continue
 429                 }
 430                 alts := refs[v.Ref]
 431                 if alts == nil {
 432                         alts = map[string]int{}
 433                         refs[v.Ref] = alts
 434                 }
 435                 alts[v.New] = 0
 436         }
 437         for ref, alts := range refs {
 438                 var altslice []string
 439                 for alt := range alts {
 440                         altslice = append(altslice, alt)
 441                 }
 442                 sort.Strings(altslice)
 443                 for i, a := range altslice {
 444                         alts[a] = i + 1
 445                 }
 446                 fmt.Fprintf(out, "%s\t%d\t%s\t%s", seqname, varslice[0].Position, ref, strings.Join(altslice, ","))
 447                 for i := 0; i < len(varslice); i += 2 {
 448                         v1, v2 := varslice[i], varslice[i+1]
 449                         a1, a2 := alts[v1.New], alts[v2.New]
 450                         if v1.Ref != ref {
 451                                 a1 = 0
 452                         }
 453                         if v2.Ref != ref {
 454                                 a2 = 0
 455                         }
 456                         fmt.Fprintf(out, "\t%d/%d", a1, a2)
 457                 }
 458                 out.Write([]byte{'\n'})
 459         }
 460 }
 461
 462 func printHGVS(out io.Writer, seqname string, varslice []hgvs.Variant) {
 463         for i := 0; i < len(varslice)/2; i++ {
 464                 if i > 0 {
 465                         out.Write([]byte{'\t'})
 466                 }
 467                 var1, var2 := varslice[i*2], varslice[i*2+1]
 468                 if var1 == var2 {
 469                         if var1.Ref == var1.New {
 470                                 out.Write([]byte{'.'})
 471                         } else {
 472                                 fmt.Fprintf(out, "%s:g.%s", seqname, var1.String())
 473                         }
 474                 } else {
 475                         fmt.Fprintf(out, "%s:g.[%s];[%s]", seqname, var1.String(), var2.String())
 476                 }
 477         }
 478         out.Write([]byte{'\n'})
 479 }
 480
 481 func printHGVSOneHot(out io.Writer, seqname string, varslice []hgvs.Variant) {
 482         vars := map[hgvs.Variant]bool{}
 483         for _, v := range varslice {
 484                 if v.Ref != v.New {
 485                         vars[v] = true
 486                 }
 487         }
 488
 489         // sort variants to ensure output is deterministic
 490         sorted := make([]hgvs.Variant, 0, len(vars))
 491         for v := range vars {
 492                 sorted = append(sorted, v)
 493         }
 494         sort.Slice(sorted, func(a, b int) bool { return hgvs.Less(sorted[a], sorted[b]) })
 495
 496         for _, v := range sorted {
 497                 fmt.Fprintf(out, "%s.%s", seqname, v.String())
 498                 for i := 0; i < len(varslice); i += 2 {
 499                         if varslice[i] == v || varslice[i+1] == v {
 500                                 out.Write([]byte("\t1"))
 501                         } else {
 502                                 out.Write([]byte("\t0"))
 503                         }
 504                 }
 505                 out.Write([]byte{'\n'})
 506         }
 507 }