gvcf2numpy.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "compress/gzip"
   6         "flag"
   7         "fmt"
   8         "io"
   9         "log"
  10         "net/http"
  11         _ "net/http/pprof"
  12         "os"
  13         "os/exec"
  14         "path/filepath"
  15         "regexp"
  16         "runtime"
  17         "sort"
  18         "strings"
  19         "sync"
  20         "sync/atomic"
  21         "time"
  22
  23         "github.com/kshedden/gonpy"
  24 )
  25
  26 type gvcf2numpy struct {
  27         tagLibraryFile string
  28         refFile        string
  29         output         io.Writer
  30 }
  31
  32 func (cmd *gvcf2numpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  33         var err error
  34         defer func() {
  35                 if err != nil {
  36                         fmt.Fprintf(stderr, "%s\n", err)
  37                 }
  38         }()
  39         flags := flag.NewFlagSet("", flag.ContinueOnError)
  40         flags.SetOutput(stderr)
  41         flags.StringVar(&cmd.tagLibraryFile, "tag-library", "", "tag library fasta `file`")
  42         flags.StringVar(&cmd.refFile, "ref", "", "reference fasta `file`")
  43         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  44         err = flags.Parse(args)
  45         if err == flag.ErrHelp {
  46                 err = nil
  47                 return 0
  48         } else if err != nil {
  49                 return 2
  50         } else if cmd.refFile == "" || cmd.tagLibraryFile == "" {
  51                 fmt.Fprintln(os.Stderr, "cannot run without -tag-library and -ref arguments")
  52                 return 2
  53         } else if flags.NArg() == 0 {
  54                 flags.Usage()
  55                 return 2
  56         }
  57         cmd.output = stdout
  58
  59         if *pprof != "" {
  60                 go func() {
  61                         log.Println(http.ListenAndServe(*pprof, nil))
  62                 }()
  63         }
  64
  65         infiles, err := listInputFiles(flags.Args())
  66         if err != nil {
  67                 return 1
  68         }
  69
  70         tilelib, err := cmd.loadTileLibrary()
  71         if err != nil {
  72                 return 1
  73         }
  74         go func() {
  75                 for range time.Tick(10 * time.Second) {
  76                         log.Printf("tilelib.Len() == %d", tilelib.Len())
  77                 }
  78         }()
  79         tseqs, err := cmd.tileGVCFs(tilelib, infiles)
  80         if err != nil {
  81                 return 1
  82         }
  83         err = cmd.printVariants(tseqs)
  84         if err != nil {
  85                 return 1
  86         }
  87         return 0
  88 }
  89
  90 func (cmd *gvcf2numpy) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, error) {
  91         var input io.ReadCloser
  92         input, err := os.Open(infile)
  93         if err != nil {
  94                 return nil, err
  95         }
  96         defer input.Close()
  97         if strings.HasSuffix(infile, ".gz") {
  98                 input, err = gzip.NewReader(input)
  99                 if err != nil {
 100                         return nil, err
 101                 }
 102                 defer input.Close()
 103         }
 104         return tilelib.TileFasta(infile, input)
 105 }
 106
 107 func (cmd *gvcf2numpy) loadTileLibrary() (*tileLibrary, error) {
 108         log.Printf("tag library %s load starting", cmd.tagLibraryFile)
 109         f, err := os.Open(cmd.tagLibraryFile)
 110         if err != nil {
 111                 return nil, err
 112         }
 113         defer f.Close()
 114         var rdr io.ReadCloser = f
 115         if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
 116                 rdr, err = gzip.NewReader(f)
 117                 if err != nil {
 118                         return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
 119                 }
 120                 defer rdr.Close()
 121         }
 122         var taglib tagLibrary
 123         err = taglib.Load(rdr)
 124         if err != nil {
 125                 return nil, err
 126         }
 127         if taglib.Len() < 1 {
 128                 return nil, fmt.Errorf("cannot tile: tag library is empty")
 129         }
 130         log.Printf("tag library %s load done", cmd.tagLibraryFile)
 131         return &tileLibrary{taglib: &taglib}, nil
 132 }
 133
 134 func listInputFiles(paths []string) (files []string, err error) {
 135         for _, path := range paths {
 136                 if fi, err := os.Stat(path); err != nil {
 137                         return nil, fmt.Errorf("%s: stat failed: %s", path, err)
 138                 } else if !fi.IsDir() {
 139                         if !strings.HasSuffix(path, ".2.fasta") || strings.HasSuffix(path, ".2.fasta.gz") {
 140                                 files = append(files, path)
 141                         }
 142                         continue
 143                 }
 144                 d, err := os.Open(path)
 145                 if err != nil {
 146                         return nil, fmt.Errorf("%s: open failed: %s", path, err)
 147                 }
 148                 defer d.Close()
 149                 names, err := d.Readdirnames(0)
 150                 if err != nil {
 151                         return nil, fmt.Errorf("%s: readdir failed: %s", path, err)
 152                 }
 153                 sort.Strings(names)
 154                 for _, name := range names {
 155                         if strings.HasSuffix(name, ".vcf") || strings.HasSuffix(name, ".vcf.gz") {
 156                                 files = append(files, filepath.Join(path, name))
 157                         } else if strings.HasSuffix(name, ".1.fasta") || strings.HasSuffix(name, ".1.fasta.gz") {
 158                                 files = append(files, filepath.Join(path, name))
 159                         }
 160                 }
 161                 d.Close()
 162         }
 163         for _, file := range files {
 164                 if strings.HasSuffix(file, ".1.fasta") || strings.HasSuffix(file, ".1.fasta.gz") {
 165                         continue
 166                 } else if _, err := os.Stat(file + ".csi"); err == nil {
 167                         continue
 168                 } else if _, err = os.Stat(file + ".tbi"); err == nil {
 169                         continue
 170                 } else {
 171                         return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
 172                 }
 173         }
 174         return
 175 }
 176
 177 func (cmd *gvcf2numpy) tileGVCFs(tilelib *tileLibrary, infiles []string) ([]tileSeq, error) {
 178         starttime := time.Now()
 179         errs := make(chan error, 1)
 180         tseqs := make([]tileSeq, len(infiles)*2)
 181         todo := make(chan func() error, len(infiles)*2)
 182         var wg sync.WaitGroup
 183         for i, infile := range infiles {
 184                 i, infile := i, infile
 185                 if strings.HasSuffix(infile, ".1.fasta") || strings.HasSuffix(infile, ".1.fasta.gz") {
 186                         todo <- func() (err error) {
 187                                 log.Printf("%s starting", infile)
 188                                 defer log.Printf("%s done", infile)
 189                                 tseqs[i*2], err = cmd.tileFasta(tilelib, infile)
 190                                 return
 191                         }
 192                         infile2 := regexp.MustCompile(`\.1\.fasta(\.gz)?$`).ReplaceAllString(infile, `.2.fasta$1`)
 193                         todo <- func() (err error) {
 194                                 log.Printf("%s starting", infile2)
 195                                 defer log.Printf("%s done", infile2)
 196                                 tseqs[i*2+1], err = cmd.tileFasta(tilelib, infile2)
 197                                 return
 198                         }
 199                 } else {
 200                         for phase := 0; phase < 2; phase++ {
 201                                 phase := phase
 202                                 todo <- func() (err error) {
 203                                         log.Printf("%s phase %d starting", infile, phase+1)
 204                                         defer log.Printf("%s phase %d done", infile, phase+1)
 205                                         tseqs[i*2+phase], err = cmd.tileGVCF(tilelib, infile, phase)
 206                                         return
 207                                 }
 208                         }
 209                 }
 210         }
 211         go close(todo)
 212         running := int64(runtime.NumCPU())
 213         for i := 0; i < runtime.NumCPU(); i++ {
 214                 wg.Add(1)
 215                 go func() {
 216                         defer wg.Done()
 217                         defer atomic.AddInt64(&running, -1)
 218                         for fn := range todo {
 219                                 if len(errs) > 0 {
 220                                         return
 221                                 }
 222                                 err := fn()
 223                                 if err != nil {
 224                                         select {
 225                                         case errs <- err:
 226                                         default:
 227                                         }
 228                                 }
 229                                 remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
 230                                 ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
 231                                 eta := time.Now().Add(ttl)
 232                                 log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
 233                         }
 234                 }()
 235         }
 236         wg.Wait()
 237         go close(errs)
 238         return tseqs, <-errs
 239 }
 240
 241 func (cmd *gvcf2numpy) printVariants(tseqs []tileSeq) error {
 242         maxtag := tagID(-1)
 243         for _, tseq := range tseqs {
 244                 for _, path := range tseq {
 245                         for _, tvar := range path {
 246                                 if maxtag < tvar.tag {
 247                                         maxtag = tvar.tag
 248                                 }
 249                         }
 250                 }
 251         }
 252         rows := len(tseqs) / 2
 253         cols := 2 * (int(maxtag) + 1)
 254         out := make([]uint16, rows*cols)
 255         for row := 0; row < len(tseqs)/2; row++ {
 256                 for phase := 0; phase < 2; phase++ {
 257                         for _, path := range tseqs[row*2+phase] {
 258                                 for _, tvar := range path {
 259                                         out[row*cols+2*int(tvar.tag)+phase] = uint16(tvar.variant)
 260                                 }
 261                         }
 262                 }
 263         }
 264         w := bufio.NewWriter(cmd.output)
 265         npw, err := gonpy.NewWriter(nopCloser{w})
 266         if err != nil {
 267                 return err
 268         }
 269         npw.Shape = []int{rows, cols}
 270         npw.WriteUint16(out)
 271         return w.Flush()
 272 }
 273
 274 type nopCloser struct {
 275         io.Writer
 276 }
 277
 278 func (nopCloser) Close() error { return nil }
 279
 280 func (cmd *gvcf2numpy) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, err error) {
 281         args := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase + 1), infile}
 282         indexsuffix := ".tbi"
 283         if _, err := os.Stat(infile + ".csi"); err == nil {
 284                 indexsuffix = ".csi"
 285         }
 286         if out, err := exec.Command("docker", "image", "ls", "-q", "lightning-runtime").Output(); err == nil && len(out) > 0 {
 287                 args = append([]string{
 288                         "docker", "run", "--rm",
 289                         "--log-driver=none",
 290                         "--volume=" + infile + ":" + infile + ":ro",
 291                         "--volume=" + infile + indexsuffix + ":" + infile + indexsuffix + ":ro",
 292                         "--volume=" + cmd.refFile + ":" + cmd.refFile + ":ro",
 293                         "lightning-runtime",
 294                 }, args...)
 295         }
 296         consensus := exec.Command(args[0], args[1:]...)
 297         consensus.Stderr = os.Stderr
 298         stdout, err := consensus.StdoutPipe()
 299         defer stdout.Close()
 300         if err != nil {
 301                 return
 302         }
 303         err = consensus.Start()
 304         if err != nil {
 305                 return
 306         }
 307         defer consensus.Wait()
 308         tileseq, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
 309         if err != nil {
 310                 return
 311         }
 312         err = stdout.Close()
 313         if err != nil {
 314                 return
 315         }
 316         err = consensus.Wait()
 317         if err != nil {
 318                 err = fmt.Errorf("%s phase %d: bcftools: %s", infile, phase, err)
 319                 return
 320         }
 321         return
 322 }