import.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "compress/gzip"
   6         "encoding/gob"
   7         "encoding/json"
   8         "errors"
   9         "flag"
  10         "fmt"
  11         "image/color/palette"
  12         "io"
  13         "net/http"
  14         _ "net/http/pprof"
  15         "os"
  16         "os/exec"
  17         "path/filepath"
  18         "regexp"
  19         "runtime"
  20         "sort"
  21         "strings"
  22         "sync"
  23         "sync/atomic"
  24         "time"
  25
  26         "git.arvados.org/arvados.git/sdk/go/arvados"
  27         log "github.com/sirupsen/logrus"
  28         "gonum.org/v1/plot"
  29         "gonum.org/v1/plot/plotter"
  30         "gonum.org/v1/plot/vg"
  31         "gonum.org/v1/plot/vg/draw"
  32 )
  33
  34 type importer struct {
  35         tagLibraryFile string
  36         refFile        string
  37         outputFile     string
  38         projectUUID    string
  39         runLocal       bool
  40         skipOOO        bool
  41         outputTiles    bool
  42         includeNoCalls bool
  43         outputStats    string
  44         encoder        *gob.Encoder
  45 }
  46
  47 func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  48         var err error
  49         defer func() {
  50                 if err != nil {
  51                         fmt.Fprintf(stderr, "%s\n", err)
  52                 }
  53         }()
  54         flags := flag.NewFlagSet("", flag.ContinueOnError)
  55         flags.SetOutput(stderr)
  56         flags.StringVar(&cmd.tagLibraryFile, "tag-library", "", "tag library fasta `file`")
  57         flags.StringVar(&cmd.refFile, "ref", "", "reference fasta `file`")
  58         flags.StringVar(&cmd.outputFile, "o", "-", "output `file`")
  59         flags.StringVar(&cmd.projectUUID, "project", "", "project `UUID` for output data")
  60         flags.BoolVar(&cmd.runLocal, "local", false, "run on local host (default: run in an arvados container)")
  61         flags.BoolVar(&cmd.skipOOO, "skip-ooo", false, "skip out-of-order tags")
  62         flags.BoolVar(&cmd.outputTiles, "output-tiles", false, "include tile variant sequences in output file")
  63         flags.BoolVar(&cmd.includeNoCalls, "include-no-calls", false, "treat tiles with no-calls as regular tiles")
  64         flags.StringVar(&cmd.outputStats, "output-stats", "", "output stats to `file` (json)")
  65         priority := flags.Int("priority", 500, "container request priority")
  66         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  67         loglevel := flags.String("loglevel", "info", "logging threshold (trace, debug, info, warn, error, fatal, or panic)")
  68         err = flags.Parse(args)
  69         if err == flag.ErrHelp {
  70                 err = nil
  71                 return 0
  72         } else if err != nil {
  73                 return 2
  74         } else if cmd.tagLibraryFile == "" {
  75                 fmt.Fprintln(os.Stderr, "cannot import without -tag-library argument")
  76                 return 2
  77         } else if flags.NArg() == 0 {
  78                 flags.Usage()
  79                 return 2
  80         }
  81
  82         if *pprof != "" {
  83                 go func() {
  84                         log.Println(http.ListenAndServe(*pprof, nil))
  85                 }()
  86         }
  87
  88         lvl, err := log.ParseLevel(*loglevel)
  89         if err != nil {
  90                 return 2
  91         }
  92         log.SetLevel(lvl)
  93
  94         if !cmd.runLocal {
  95                 runner := arvadosContainerRunner{
  96                         Name:        "lightning import",
  97                         Client:      arvados.NewClientFromEnv(),
  98                         ProjectUUID: cmd.projectUUID,
  99                         RAM:         60000000000,
 100                         VCPUs:       16,
 101                         Priority:    *priority,
 102                 }
 103                 err = runner.TranslatePaths(&cmd.tagLibraryFile, &cmd.refFile, &cmd.outputFile)
 104                 if err != nil {
 105                         return 1
 106                 }
 107                 inputs := flags.Args()
 108                 for i := range inputs {
 109                         err = runner.TranslatePaths(&inputs[i])
 110                         if err != nil {
 111                                 return 1
 112                         }
 113                 }
 114                 if cmd.outputFile == "-" {
 115                         cmd.outputFile = "/mnt/output/library.gob"
 116                 } else {
 117                         // Not yet implemented, but this should write
 118                         // the collection to an existing collection,
 119                         // possibly even an in-place update.
 120                         err = errors.New("cannot specify output file in container mode: not implemented")
 121                         return 1
 122                 }
 123                 runner.Args = append([]string{"import",
 124                         "-local=true",
 125                         "-loglevel=" + *loglevel,
 126                         fmt.Sprintf("-skip-ooo=%v", cmd.skipOOO),
 127                         fmt.Sprintf("-output-tiles=%v", cmd.outputTiles),
 128                         fmt.Sprintf("-include-no-calls=%v", cmd.includeNoCalls),
 129                         "-output-stats", "/mnt/output/stats.json",
 130                         "-tag-library", cmd.tagLibraryFile,
 131                         "-ref", cmd.refFile,
 132                         "-o", cmd.outputFile,
 133                 }, inputs...)
 134                 var output string
 135                 output, err = runner.Run()
 136                 if err != nil {
 137                         return 1
 138                 }
 139                 fmt.Fprintln(stdout, output+"/library.gob")
 140                 return 0
 141         }
 142
 143         infiles, err := listInputFiles(flags.Args())
 144         if err != nil {
 145                 return 1
 146         }
 147
 148         taglib, err := cmd.loadTagLibrary()
 149         if err != nil {
 150                 return 1
 151         }
 152
 153         var output io.WriteCloser
 154         if cmd.outputFile == "-" {
 155                 output = nopCloser{stdout}
 156         } else {
 157                 output, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
 158                 if err != nil {
 159                         return 1
 160                 }
 161                 defer output.Close()
 162         }
 163         bufw := bufio.NewWriter(output)
 164         cmd.encoder = gob.NewEncoder(bufw)
 165
 166         tilelib := &tileLibrary{taglib: taglib, includeNoCalls: cmd.includeNoCalls, skipOOO: cmd.skipOOO}
 167         if cmd.outputTiles {
 168                 cmd.encoder.Encode(LibraryEntry{TagSet: taglib.Tags()})
 169                 tilelib.encoder = cmd.encoder
 170         }
 171         go func() {
 172                 for range time.Tick(10 * time.Minute) {
 173                         log.Printf("tilelib.Len() == %d", tilelib.Len())
 174                 }
 175         }()
 176
 177         err = cmd.tileInputs(tilelib, infiles)
 178         if err != nil {
 179                 return 1
 180         }
 181         err = bufw.Flush()
 182         if err != nil {
 183                 return 1
 184         }
 185         err = output.Close()
 186         if err != nil {
 187                 return 1
 188         }
 189         return 0
 190 }
 191
 192 func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, []importStats, error) {
 193         var input io.ReadCloser
 194         input, err := os.Open(infile)
 195         if err != nil {
 196                 return nil, nil, err
 197         }
 198         defer input.Close()
 199         if strings.HasSuffix(infile, ".gz") {
 200                 input, err = gzip.NewReader(input)
 201                 if err != nil {
 202                         return nil, nil, err
 203                 }
 204                 defer input.Close()
 205         }
 206         return tilelib.TileFasta(infile, input)
 207 }
 208
 209 func (cmd *importer) loadTagLibrary() (*tagLibrary, error) {
 210         log.Printf("tag library %s load starting", cmd.tagLibraryFile)
 211         f, err := os.Open(cmd.tagLibraryFile)
 212         if err != nil {
 213                 return nil, err
 214         }
 215         defer f.Close()
 216         var rdr io.ReadCloser = f
 217         if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
 218                 rdr, err = gzip.NewReader(f)
 219                 if err != nil {
 220                         return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
 221                 }
 222                 defer rdr.Close()
 223         }
 224         var taglib tagLibrary
 225         err = taglib.Load(rdr)
 226         if err != nil {
 227                 return nil, err
 228         }
 229         if taglib.Len() < 1 {
 230                 return nil, fmt.Errorf("cannot tile: tag library is empty")
 231         }
 232         log.Printf("tag library %s load done", cmd.tagLibraryFile)
 233         return &taglib, nil
 234 }
 235
 236 var (
 237         vcfFilenameRe    = regexp.MustCompile(`\.vcf(\.gz)?$`)
 238         fasta1FilenameRe = regexp.MustCompile(`\.1\.fa(sta)?(\.gz)?$`)
 239         fasta2FilenameRe = regexp.MustCompile(`\.2\.fa(sta)?(\.gz)?$`)
 240         fastaFilenameRe  = regexp.MustCompile(`\.fa(sta)?(\.gz)?$`)
 241 )
 242
 243 func listInputFiles(paths []string) (files []string, err error) {
 244         for _, path := range paths {
 245                 if fi, err := os.Stat(path); err != nil {
 246                         return nil, fmt.Errorf("%s: stat failed: %s", path, err)
 247                 } else if !fi.IsDir() {
 248                         if !fasta2FilenameRe.MatchString(path) {
 249                                 files = append(files, path)
 250                         }
 251                         continue
 252                 }
 253                 d, err := os.Open(path)
 254                 if err != nil {
 255                         return nil, fmt.Errorf("%s: open failed: %s", path, err)
 256                 }
 257                 defer d.Close()
 258                 names, err := d.Readdirnames(0)
 259                 if err != nil {
 260                         return nil, fmt.Errorf("%s: readdir failed: %s", path, err)
 261                 }
 262                 sort.Strings(names)
 263                 for _, name := range names {
 264                         if vcfFilenameRe.MatchString(name) {
 265                                 files = append(files, filepath.Join(path, name))
 266                         } else if fastaFilenameRe.MatchString(name) && !fasta2FilenameRe.MatchString(name) {
 267                                 files = append(files, filepath.Join(path, name))
 268                         }
 269                 }
 270                 d.Close()
 271         }
 272         for _, file := range files {
 273                 if fastaFilenameRe.MatchString(file) {
 274                         continue
 275                 } else if vcfFilenameRe.MatchString(file) {
 276                         if _, err := os.Stat(file + ".csi"); err == nil {
 277                                 continue
 278                         } else if _, err = os.Stat(file + ".tbi"); err == nil {
 279                                 continue
 280                         } else {
 281                                 return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
 282                         }
 283                 } else {
 284                         return nil, fmt.Errorf("don't know how to handle filename %s", file)
 285                 }
 286         }
 287         return
 288 }
 289
 290 func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
 291         starttime := time.Now()
 292         errs := make(chan error, 1)
 293         todo := make(chan func() error, len(infiles)*2)
 294         allstats := make([][]importStats, len(infiles)*2)
 295         var encodeJobs sync.WaitGroup
 296         for idx, infile := range infiles {
 297                 idx, infile := idx, infile
 298                 var phases sync.WaitGroup
 299                 phases.Add(2)
 300                 variants := make([][]tileVariantID, 2)
 301                 if fasta1FilenameRe.MatchString(infile) {
 302                         todo <- func() error {
 303                                 defer phases.Done()
 304                                 log.Printf("%s starting", infile)
 305                                 defer log.Printf("%s done", infile)
 306                                 tseqs, stats, err := cmd.tileFasta(tilelib, infile)
 307                                 allstats[idx*2] = stats
 308                                 var kept, dropped int
 309                                 variants[0], kept, dropped = tseqs.Variants()
 310                                 log.Printf("%s found %d unique tags plus %d repeats", infile, kept, dropped)
 311                                 return err
 312                         }
 313                         infile2 := fasta1FilenameRe.ReplaceAllString(infile, `.2.fa$1$2`)
 314                         todo <- func() error {
 315                                 defer phases.Done()
 316                                 log.Printf("%s starting", infile2)
 317                                 defer log.Printf("%s done", infile2)
 318                                 tseqs, stats, err := cmd.tileFasta(tilelib, infile2)
 319                                 allstats[idx*2+1] = stats
 320                                 var kept, dropped int
 321                                 variants[1], kept, dropped = tseqs.Variants()
 322                                 log.Printf("%s found %d unique tags plus %d repeats", infile2, kept, dropped)
 323
 324                                 return err
 325                         }
 326                 } else if fastaFilenameRe.MatchString(infile) {
 327                         todo <- func() error {
 328                                 defer phases.Done()
 329                                 defer phases.Done()
 330                                 log.Printf("%s starting", infile)
 331                                 defer log.Printf("%s done", infile)
 332                                 tseqs, stats, err := cmd.tileFasta(tilelib, infile)
 333                                 allstats[idx*2] = stats
 334                                 if err != nil {
 335                                         return err
 336                                 }
 337                                 totlen := 0
 338                                 for _, tseq := range tseqs {
 339                                         totlen += len(tseq)
 340                                 }
 341                                 log.Printf("%s tiled %d seqs, total len %d", infile, len(tseqs), totlen)
 342                                 return cmd.encoder.Encode(LibraryEntry{
 343                                         CompactSequences: []CompactSequence{{Name: infile, TileSequences: tseqs}},
 344                                 })
 345                         }
 346                         // Don't write out a CompactGenomes entry
 347                         continue
 348                 } else if vcfFilenameRe.MatchString(infile) {
 349                         for phase := 0; phase < 2; phase++ {
 350                                 phase := phase
 351                                 todo <- func() error {
 352                                         defer phases.Done()
 353                                         log.Printf("%s phase %d starting", infile, phase+1)
 354                                         defer log.Printf("%s phase %d done", infile, phase+1)
 355                                         tseqs, stats, err := cmd.tileGVCF(tilelib, infile, phase)
 356                                         allstats[idx*2] = stats
 357                                         var kept, dropped int
 358                                         variants[phase], kept, dropped = tseqs.Variants()
 359                                         log.Printf("%s phase %d found %d unique tags plus %d repeats", infile, phase+1, kept, dropped)
 360                                         return err
 361                                 }
 362                         }
 363                 } else {
 364                         panic(fmt.Sprintf("bug: unhandled filename %q", infile))
 365                 }
 366                 encodeJobs.Add(1)
 367                 go func() {
 368                         defer encodeJobs.Done()
 369                         phases.Wait()
 370                         if len(errs) > 0 {
 371                                 return
 372                         }
 373                         err := cmd.encoder.Encode(LibraryEntry{
 374                                 CompactGenomes: []CompactGenome{{Name: infile, Variants: flatten(variants)}},
 375                         })
 376                         if err != nil {
 377                                 select {
 378                                 case errs <- err:
 379                                 default:
 380                                 }
 381                         }
 382                 }()
 383         }
 384         go close(todo)
 385         var tileJobs sync.WaitGroup
 386         var running int64
 387         for i := 0; i < runtime.NumCPU()*9/8+1; i++ {
 388                 tileJobs.Add(1)
 389                 atomic.AddInt64(&running, 1)
 390                 go func() {
 391                         defer tileJobs.Done()
 392                         defer atomic.AddInt64(&running, -1)
 393                         for fn := range todo {
 394                                 if len(errs) > 0 {
 395                                         return
 396                                 }
 397                                 err := fn()
 398                                 if err != nil {
 399                                         select {
 400                                         case errs <- err:
 401                                         default:
 402                                         }
 403                                 }
 404                                 remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
 405                                 if remain < cap(todo) {
 406                                         ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
 407                                         eta := time.Now().Add(ttl)
 408                                         log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
 409                                 }
 410                         }
 411                 }()
 412         }
 413         tileJobs.Wait()
 414         encodeJobs.Wait()
 415
 416         go close(errs)
 417         err := <-errs
 418         if err != nil {
 419                 return err
 420         }
 421
 422         if cmd.outputStats != "" {
 423                 f, err := os.OpenFile(cmd.outputStats, os.O_CREATE|os.O_WRONLY, 0666)
 424                 if err != nil {
 425                         return err
 426                 }
 427                 var flatstats []importStats
 428                 for _, stats := range allstats {
 429                         flatstats = append(flatstats, stats...)
 430                 }
 431                 err = json.NewEncoder(f).Encode(flatstats)
 432                 if err != nil {
 433                         return err
 434                 }
 435         }
 436
 437         return nil
 438 }
 439
 440 func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, stats []importStats, err error) {
 441         if cmd.refFile == "" {
 442                 err = errors.New("cannot import vcf: reference data (-ref) not specified")
 443                 return
 444         }
 445         args := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase + 1), infile}
 446         indexsuffix := ".tbi"
 447         if _, err := os.Stat(infile + ".csi"); err == nil {
 448                 indexsuffix = ".csi"
 449         }
 450         if out, err := exec.Command("docker", "image", "ls", "-q", "lightning-runtime").Output(); err == nil && len(out) > 0 {
 451                 args = append([]string{
 452                         "docker", "run", "--rm",
 453                         "--log-driver=none",
 454                         "--volume=" + infile + ":" + infile + ":ro",
 455                         "--volume=" + infile + indexsuffix + ":" + infile + indexsuffix + ":ro",
 456                         "--volume=" + cmd.refFile + ":" + cmd.refFile + ":ro",
 457                         "lightning-runtime",
 458                 }, args...)
 459         }
 460         consensus := exec.Command(args[0], args[1:]...)
 461         consensus.Stderr = os.Stderr
 462         stdout, err := consensus.StdoutPipe()
 463         defer stdout.Close()
 464         if err != nil {
 465                 return
 466         }
 467         err = consensus.Start()
 468         if err != nil {
 469                 return
 470         }
 471         defer consensus.Wait()
 472         tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
 473         if err != nil {
 474                 return
 475         }
 476         err = stdout.Close()
 477         if err != nil {
 478                 return
 479         }
 480         err = consensus.Wait()
 481         if err != nil {
 482                 err = fmt.Errorf("%s phase %d: bcftools: %s", infile, phase, err)
 483                 return
 484         }
 485         return
 486 }
 487
 488 func flatten(variants [][]tileVariantID) []tileVariantID {
 489         ntags := 0
 490         for _, v := range variants {
 491                 if ntags < len(v) {
 492                         ntags = len(v)
 493                 }
 494         }
 495         flat := make([]tileVariantID, ntags*2)
 496         for i := 0; i < ntags; i++ {
 497                 for hap := 0; hap < 2; hap++ {
 498                         if i < len(variants[hap]) {
 499                                 flat[i*2+hap] = variants[hap][i]
 500                         }
 501                 }
 502         }
 503         return flat
 504 }
 505
 506 type importstatsplot struct{}
 507
 508 func (cmd *importstatsplot) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
 509         err := cmd.Plot(stdin, stdout)
 510         if err != nil {
 511                 log.Errorf("%s", err)
 512                 return 1
 513         }
 514         return 0
 515 }
 516
 517 func (cmd *importstatsplot) Plot(stdin io.Reader, stdout io.Writer) error {
 518         var stats []importStats
 519         err := json.NewDecoder(stdin).Decode(&stats)
 520         if err != nil {
 521                 return err
 522         }
 523
 524         p, err := plot.New()
 525         if err != nil {
 526                 return err
 527         }
 528         p.Title.Text = "coverage preserved by import (excl X<0.65)"
 529         p.X.Label.Text = "input base calls ÷ sequence length"
 530         p.Y.Label.Text = "output base calls ÷ input base calls"
 531         p.Add(plotter.NewGrid())
 532
 533         data := map[string]plotter.XYs{}
 534         for _, stat := range stats {
 535                 data[stat.InputLabel] = append(data[stat.InputLabel], plotter.XY{
 536                         X: float64(stat.InputCoverage) / float64(stat.InputLength),
 537                         Y: float64(stat.TileCoverage) / float64(stat.InputCoverage),
 538                 })
 539         }
 540
 541         nextInPalette := 0
 542         for label, xys := range data {
 543                 s, err := plotter.NewScatter(xys)
 544                 if err != nil {
 545                         return err
 546                 }
 547                 s.GlyphStyle.Color = palette.Plan9[nextInPalette%len(palette.Plan9)]
 548                 s.GlyphStyle.Radius = vg.Millimeter / 2
 549                 s.GlyphStyle.Shape = draw.CrossGlyph{}
 550                 nextInPalette += 7
 551                 p.Add(s)
 552                 if false {
 553                         p.Legend.Add(label, s)
 554                 }
 555         }
 556         p.X.Min = 0.65
 557         p.X.Max = 1
 558
 559         w, err := p.WriterTo(8*vg.Inch, 6*vg.Inch, "svg")
 560         if err != nil {
 561                 return err
 562         }
 563         _, err = w.WriteTo(stdout)
 564         return err
 565 }