Allow importing all-hom (reference) data from single fasta file.
[lightning.git] / vcf2fasta.go
index 383ef85ca76b5a6ad15c83778082549d995edbe9..2f02e44a28b8b9188ba8700c4ce15c88a08647c5 100644 (file)
@@ -1,8 +1,10 @@
 package main
 
 import (
+       "bufio"
        "bytes"
        "compress/gzip"
+       "context"
        "errors"
        "flag"
        "fmt"
@@ -13,9 +15,12 @@ import (
        "os"
        "os/exec"
        "path/filepath"
+       "regexp"
        "runtime"
+       "strconv"
        "strings"
        "sync"
+       "syscall"
 
        "git.arvados.org/arvados.git/sdk/go/arvados"
        log "github.com/sirupsen/logrus"
@@ -111,7 +116,10 @@ func (cmd *vcf2fasta) RunCommand(prog string, args []string, stdin io.Reader, st
                                },
                        },
                }
-               err = runner.TranslatePaths(&cmd.refFile)
+               if cmd.mask {
+                       runner.RAM += int64(cmd.vcpus) << 31
+               }
+               err = runner.TranslatePaths(&cmd.refFile, &cmd.genomeFile)
                if err != nil {
                        return 1
                }
@@ -122,7 +130,12 @@ func (cmd *vcf2fasta) RunCommand(prog string, args []string, stdin io.Reader, st
                                return 1
                        }
                }
-               runner.Args = append([]string{"vcf2fasta", "-local=true", "-ref", cmd.refFile, fmt.Sprintf("-mask=%v", cmd.mask), "-gvcf-regions.py", "/gvcf_regions.py", "-output-dir", "/mnt/output"}, inputs...)
+               runner.Args = append([]string{"vcf2fasta",
+                       "-local=true",
+                       "-ref", cmd.refFile, fmt.Sprintf("-mask=%v", cmd.mask),
+                       "-genome", cmd.genomeFile,
+                       "-gvcf-regions.py", "/gvcf_regions.py",
+                       "-output-dir", "/mnt/output"}, inputs...)
                var output string
                output, err = runner.Run()
                if err != nil {
@@ -187,7 +200,7 @@ func maybeInDocker(args, mountfiles []string) []string {
                return args
        }
        dockerrun := []string{
-               "docker", "run", "--rm",
+               "docker", "run", "--rm", "-i",
                "--log-driver=none",
        }
        for _, f := range mountfiles {
@@ -199,9 +212,12 @@ func maybeInDocker(args, mountfiles []string) []string {
 }
 
 func (cmd *vcf2fasta) vcf2fasta(infile string, phase int) error {
+       ctx, cancel := context.WithCancel(context.Background())
+       defer cancel()
+
        _, basename := filepath.Split(infile)
        outfile := filepath.Join(cmd.outputDir, fmt.Sprintf("%s.%d.fasta.gz", basename, phase))
-       outf, err := os.OpenFile(outfile, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0777)
+       outf, err := os.OpenFile(outfile, os.O_CREATE|os.O_WRONLY, 0777)
        if err != nil {
                return fmt.Errorf("error opening output file: %s", err)
        }
@@ -209,74 +225,184 @@ func (cmd *vcf2fasta) vcf2fasta(infile string, phase int) error {
        gzipw := gzip.NewWriter(outf)
        defer gzipw.Close()
 
-       var maskfile *os.File // reading side of a pipe if we're running bedtools, otherwise nil
+       var maskfifo string // filename of mask fifo if we're running bedtools, otherwise ""
 
        var wg sync.WaitGroup
-       errs := make(chan error, 3)
+       errs := make(chan error, 2)
        if cmd.mask {
-               if cmd.genomeFile == "" {
-                       return errors.New("cannot apply mask without -genome argument")
-               }
-               bedr, bedw, err := os.Pipe()
+               chrSize := map[string]int{}
+
+               vcffile, err := os.Open(infile)
                if err != nil {
                        return err
                }
-               bedargs := []string{"python", "-", "--gvcf_type", "gatk", infile}
-               bed := exec.Command(bedargs[0], bedargs[1:]...)
+               defer vcffile.Close()
+               var rdr io.Reader = vcffile
+               if strings.HasSuffix(infile, ".gz") {
+                       rdr, err = gzip.NewReader(vcffile)
+                       if err != nil {
+                               return err
+                       }
+               }
+               contigre := regexp.MustCompile(`([^=,]*)=([^>,]*)`)
+               scanner := bufio.NewScanner(rdr)
+               for scanner.Scan() {
+                       if s := scanner.Text(); !strings.HasPrefix(s, "##") {
+                               break
+                       } else if !strings.HasPrefix(s, "##contig=<") {
+                               continue
+                       } else {
+                               kv := map[string]string{}
+                               for _, m := range contigre.FindAllStringSubmatch(s[10:], -1) {
+                                       kv[m[1]] = m[2]
+                               }
+                               if kv["ID"] != "" && kv["length"] != "" {
+                                       chrSize[kv["ID"]], _ = strconv.Atoi(kv["length"])
+                               }
+                       }
+               }
+               if err = scanner.Err(); err != nil {
+                       return fmt.Errorf("error scanning input file %q: %s", infile, err)
+               }
+               var regions bytes.Buffer
+               bedargs := []string{"python2", "-", "--gvcf_type", "gatk", infile}
+               bed := exec.CommandContext(ctx, bedargs[0], bedargs[1:]...)
                bed.Stdin = bytes.NewBuffer(cmd.gvcfRegionsPyData)
-               bed.Stdout = bedw
+               bed.Stdout = &regions
                bed.Stderr = cmd.stderr
-               wg.Add(1)
-               go func() {
-                       defer wg.Done()
-                       log.Printf("running %v", bed.Args)
-                       errs <- bed.Run()
-               }()
+               log.Printf("running %v", bed.Args)
+               err = bed.Run()
+               log.Printf("exited %v", bed.Args)
+               if err != nil {
+                       return fmt.Errorf("gvcf_regions: %s", err)
+               }
 
-               bedcompr, bedcompw, err := os.Pipe()
+               if cmd.genomeFile != "" {
+                       // Read chromosome sizes from genome file in
+                       // case any weren't specified in the VCF
+                       // header.
+                       genomeFile, err := os.Open(cmd.genomeFile)
+                       if err != nil {
+                               return fmt.Errorf("error opening genome file %q: %s", cmd.genomeFile, err)
+                       }
+                       scanner := bufio.NewScanner(genomeFile)
+                       for scanner.Scan() {
+                               var chr string
+                               var size int
+                               _, err := fmt.Sscanf(scanner.Text(), "%s\t%d", &chr, &size)
+                               if err != nil {
+                                       return fmt.Errorf("error parsing genome file %q: %s", cmd.genomeFile, err)
+                               }
+                               if chrSize[chr] == 0 {
+                                       chrSize[chr] = size
+                               }
+                       }
+                       if err = scanner.Err(); err != nil {
+                               return fmt.Errorf("error scanning genome file %q: %s", cmd.genomeFile, err)
+                       }
+               }
+
+               // "bedtools complement" expects the chromosome sizes
+               // ("genome file") to appear in the same order as the
+               // chromosomes in the input vcf, so we need to sort
+               // them.
+               scanner = bufio.NewScanner(bytes.NewBuffer(append([]byte(nil), regions.Bytes()...)))
+               var sortedGenomeFile bytes.Buffer
+               for scanner.Scan() {
+                       var chr string
+                       var size int
+                       _, err := fmt.Sscanf(scanner.Text(), "%s\t%d", &chr, &size)
+                       if err != nil {
+                               return fmt.Errorf("error parsing gvcf_regions output: %s", err)
+                       }
+                       if size, ok := chrSize[chr]; ok {
+                               fmt.Fprintf(&sortedGenomeFile, "%s\t%d\n", chr, size)
+                               delete(chrSize, chr)
+                       }
+               }
+
+               // The bcftools --mask argument needs to end in ".bed"
+               // in order to be parsed as a BED file, so we need to
+               // use a named pipe instead of stdin.
+               tempdir, err := ioutil.TempDir("", "")
+               if err != nil {
+                       return fmt.Errorf("TempDir: %s", err)
+               }
+               defer os.RemoveAll(tempdir)
+               maskfifo = filepath.Join(tempdir, "fifo.bed")
+               err = syscall.Mkfifo(maskfifo, 0600)
+               if err != nil {
+                       return fmt.Errorf("mkfifo: %s", err)
+               }
+
+               // bedtools complement can't seem to read from a pipe
+               // reliably -- "Error: line number 1 of file
+               // /dev/stdin has 1 fields, but 3 were expected." --
+               // so we stage to a temp file.
+               regionsFile := filepath.Join(tempdir, "gvcf_regions.bed")
+               err = ioutil.WriteFile(regionsFile, regions.Bytes(), 0644)
                if err != nil {
                        return err
                }
-               bedcompargs := []string{"bedtools", "complement", "-i", "/dev/stdin", "-g", cmd.genomeFile}
-               bedcompargs = maybeInDocker(bedcompargs, []string{cmd.genomeFile})
-               bedcomp := exec.Command(bedcompargs[0], bedcompargs[1:]...)
-               bedcomp.Stdin = bedr
-               bedcomp.Stdout = bedcompw
-               bedcomp.Stderr = cmd.stderr
+
                wg.Add(1)
                go func() {
                        defer wg.Done()
+
+                       maskfifow, err := os.OpenFile(maskfifo, os.O_WRONLY, 0)
+                       if err != nil {
+                               errs <- err
+                               return
+                       }
+                       defer maskfifow.Close()
+
+                       bedcompargs := []string{"bedtools", "complement", "-i", regionsFile, "-g", "/dev/stdin"}
+                       bedcompargs = maybeInDocker(bedcompargs, []string{cmd.genomeFile})
+                       bedcomp := exec.CommandContext(ctx, bedcompargs[0], bedcompargs[1:]...)
+                       bedcomp.Stdin = &sortedGenomeFile
+                       bedcomp.Stdout = maskfifow
+                       bedcomp.Stderr = cmd.stderr
                        log.Printf("running %v", bedcomp.Args)
-                       errs <- bedcomp.Run()
+                       err = bedcomp.Run()
+                       log.Printf("exited %v", bedcomp.Args)
+                       if err != nil {
+                               errs <- fmt.Errorf("bedtools complement: %s", err)
+                               return
+                       }
+                       err = maskfifow.Close()
+                       if err != nil {
+                               errs <- err
+                               return
+                       }
                }()
-               maskfile = bedcompr
        }
 
        wg.Add(1)
        go func() {
                defer wg.Done()
                consargs := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase)}
-               if maskfile != nil {
-                       consargs = append(consargs, "--mask", "/dev/stdin")
+               if maskfifo != "" {
+                       consargs = append(consargs, "--mask", maskfifo)
                }
                consargs = append(consargs, infile)
                indexsuffix := ".tbi"
                if _, err := os.Stat(infile + ".csi"); err == nil {
                        indexsuffix = ".csi"
                }
-               consargs = maybeInDocker(consargs, []string{infile, infile + indexsuffix, cmd.refFile})
+               mounts := []string{infile, infile + indexsuffix, cmd.refFile}
+               if maskfifo != "" {
+                       mounts = append(mounts, maskfifo)
+               }
+               consargs = maybeInDocker(consargs, mounts)
 
-               consensus := exec.Command(consargs[0], consargs[1:]...)
+               consensus := exec.CommandContext(ctx, consargs[0], consargs[1:]...)
                consensus.Stderr = os.Stderr
                consensus.Stdout = gzipw
                consensus.Stderr = cmd.stderr
-               if maskfile != nil {
-                       consensus.Stdin = maskfile
-               }
                log.Printf("running %v", consensus.Args)
                err = consensus.Run()
                if err != nil {
-                       errs <- err
+                       errs <- fmt.Errorf("bcftools consensus: %s", err)
                        return
                }
                err = gzipw.Close()
@@ -294,6 +420,7 @@ func (cmd *vcf2fasta) vcf2fasta(infile string, phase int) error {
 
        for err := range errs {
                if err != nil {
+                       cancel()
                        wg.Wait()
                        return err
                }