annotate.go

   1 // Copyright (C) The Lightning Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package lightning
   6
   7 import (
   8         "bufio"
   9         "context"
  10         "errors"
  11         "flag"
  12         "fmt"
  13         "io"
  14         "io/ioutil"
  15         "net/http"
  16         _ "net/http/pprof"
  17         "os"
  18         "runtime"
  19         "sort"
  20         "strconv"
  21         "strings"
  22         "sync"
  23
  24         "git.arvados.org/arvados.git/sdk/go/arvados"
  25         "github.com/arvados/lightning/hgvs"
  26         log "github.com/sirupsen/logrus"
  27 )
  28
  29 type annotatecmd struct {
  30         dropTiles        []bool
  31         variantHash      bool
  32         maxTileSize      int
  33         tag2tagid        map[string]tagID
  34         reportAnnotation func(tag tagID, outcol int, variant tileVariantID, refname string, seqname string, pdi hgvs.Variant)
  35 }
  36
  37 func (cmd *annotatecmd) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  38         var err error
  39         defer func() {
  40                 if err != nil {
  41                         fmt.Fprintf(stderr, "%s\n", err)
  42                 }
  43         }()
  44         flags := flag.NewFlagSet("", flag.ContinueOnError)
  45         flags.SetOutput(stderr)
  46         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  47         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  48         projectUUID := flags.String("project", "", "project `UUID` for output data")
  49         priority := flags.Int("priority", 500, "container request priority")
  50         inputFilename := flags.String("i", "-", "input `file` (library)")
  51         outputFilename := flags.String("o", "-", "output `file`")
  52         flags.BoolVar(&cmd.variantHash, "variant-hash", false, "output variant hash instead of index")
  53         flags.IntVar(&cmd.maxTileSize, "max-tile-size", 50000, "don't try to make annotations for tiles bigger than given `size`")
  54         err = flags.Parse(args)
  55         if err == flag.ErrHelp {
  56                 err = nil
  57                 return 0
  58         } else if err != nil {
  59                 return 2
  60         }
  61
  62         if *pprof != "" {
  63                 go func() {
  64                         log.Println(http.ListenAndServe(*pprof, nil))
  65                 }()
  66         }
  67         if !*runlocal {
  68                 if *outputFilename != "-" {
  69                         err = errors.New("cannot specify output file in container mode: not implemented")
  70                         return 1
  71                 }
  72                 runner := arvadosContainerRunner{
  73                         Name:        "lightning annotate",
  74                         Client:      arvados.NewClientFromEnv(),
  75                         ProjectUUID: *projectUUID,
  76                         RAM:         80000000000,
  77                         VCPUs:       16,
  78                         Priority:    *priority,
  79                 }
  80                 err = runner.TranslatePaths(inputFilename)
  81                 if err != nil {
  82                         return 1
  83                 }
  84                 runner.Args = []string{"annotate", "-local=true", fmt.Sprintf("-variant-hash=%v", cmd.variantHash), "-max-tile-size", strconv.Itoa(cmd.maxTileSize), "-i", *inputFilename, "-o", "/mnt/output/tilevariants.csv"}
  85                 var output string
  86                 output, err = runner.Run()
  87                 if err != nil {
  88                         return 1
  89                 }
  90                 fmt.Fprintln(stdout, output+"/tilevariants.csv")
  91                 return 0
  92         }
  93
  94         var input io.ReadCloser
  95         if *inputFilename == "-" {
  96                 input = ioutil.NopCloser(stdin)
  97         } else {
  98                 input, err = os.Open(*inputFilename)
  99                 if err != nil {
 100                         return 1
 101                 }
 102                 defer input.Close()
 103         }
 104
 105         var output io.WriteCloser
 106         if *outputFilename == "-" {
 107                 output = nopCloser{stdout}
 108         } else {
 109                 output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0666)
 110                 if err != nil {
 111                         return 1
 112                 }
 113                 defer output.Close()
 114         }
 115         bufw := bufio.NewWriterSize(output, 4*1024*1024)
 116
 117         tilelib := &tileLibrary{
 118                 retainNoCalls:       true,
 119                 retainTileSequences: true,
 120         }
 121         err = tilelib.LoadGob(context.Background(), input, strings.HasSuffix(*inputFilename, ".gz"))
 122         if err != nil {
 123                 return 1
 124         }
 125         err = cmd.exportTileDiffs(bufw, tilelib)
 126         if err != nil {
 127                 return 1
 128         }
 129         err = bufw.Flush()
 130         if err != nil {
 131                 return 1
 132         }
 133         err = output.Close()
 134         if err != nil {
 135                 return 1
 136         }
 137         err = input.Close()
 138         if err != nil {
 139                 return 1
 140         }
 141         return 0
 142 }
 143
 144 func (cmd *annotatecmd) exportTileDiffs(outw io.Writer, tilelib *tileLibrary) error {
 145         tagset := tilelib.taglib.Tags()
 146         if len(tagset) == 0 {
 147                 return errors.New("cannot annotate library without tags")
 148         }
 149         taglen := len(tagset[0])
 150         var refs []string
 151         for name := range tilelib.refseqs {
 152                 refs = append(refs, name)
 153         }
 154         cmd.tag2tagid = make(map[string]tagID, len(tagset))
 155         for tagid, tagseq := range tagset {
 156                 cmd.tag2tagid[string(tagseq)] = tagID(tagid)
 157         }
 158         sort.Strings(refs)
 159         log.Infof("len(refs) %d", len(refs))
 160
 161         outch := make(chan string, runtime.NumCPU()*2)
 162         var outwg sync.WaitGroup
 163         defer outwg.Wait()
 164         outwg.Add(1)
 165         go func() {
 166                 defer outwg.Done()
 167                 for s := range outch {
 168                         io.WriteString(outw, s)
 169                 }
 170         }()
 171         defer close(outch)
 172
 173         nseqs := 0
 174         for _, refcs := range tilelib.refseqs {
 175                 nseqs += len(refcs)
 176         }
 177
 178         throttle := &throttle{Max: runtime.NumCPU()*2 + nseqs*2 + 1}
 179         defer throttle.Wait()
 180
 181         for _, refname := range refs {
 182                 refname := refname
 183                 refcs := tilelib.refseqs[refname]
 184                 var seqnames []string
 185                 for seqname := range refcs {
 186                         seqnames = append(seqnames, seqname)
 187                 }
 188                 sort.Strings(seqnames)
 189                 for _, seqname := range seqnames {
 190                         seqname := seqname
 191                         throttle.Acquire()
 192                         if throttle.Err() != nil {
 193                                 break
 194                         }
 195                         go func() {
 196                                 defer throttle.Release()
 197                                 throttle.Report(cmd.annotateSequence(throttle, outch, tilelib, taglen, refname, seqname, refcs[seqname], len(refs) > 1))
 198                         }()
 199                 }
 200         }
 201         throttle.Wait()
 202         return throttle.Err()
 203 }
 204
 205 func (cmd *annotatecmd) annotateSequence(throttle *throttle, outch chan<- string, tilelib *tileLibrary, taglen int, refname, seqname string, reftiles []tileLibRef, refnamecol bool) error {
 206         refnamefield := ""
 207         if refnamecol {
 208                 refnamefield = "," + trimFilenameForLabel(refname)
 209         }
 210         var refseq []byte
 211         // tilestart[123] is the index into refseq
 212         // where the tile for tag 123 was placed.
 213         tilestart := map[tagID]int{}
 214         tileend := map[tagID]int{}
 215         for _, libref := range reftiles {
 216                 if libref.Variant < 1 {
 217                         return fmt.Errorf("reference %q seq %q uses variant zero at tag %d", refname, seqname, libref.Tag)
 218                 }
 219                 seq := tilelib.TileVariantSequence(libref)
 220                 if len(seq) < taglen {
 221                         return fmt.Errorf("reference %q seq %q uses tile %d variant %d with sequence len %d < taglen %d", refname, seqname, libref.Tag, libref.Variant, len(seq), taglen)
 222                 }
 223                 overlap := taglen
 224                 if len(refseq) == 0 {
 225                         overlap = 0
 226                 }
 227                 tilestart[libref.Tag] = len(refseq) - overlap
 228                 refseq = append(refseq, seq[overlap:]...)
 229                 tileend[libref.Tag] = len(refseq)
 230         }
 231         log.Infof("seq %s len(refseq) %d len(tilestart) %d", seqname, len(refseq), len(tilestart))
 232         // outtag is tag's index in the subset of tags that aren't
 233         // dropped. If there are 10M tags and half are dropped by
 234         // dropTiles, tag ranges from 0 to 10M-1 and outtag ranges
 235         // from 0 to 5M-1.
 236         //
 237         // IOW, in the matrix built by cgs2array(), {tag} is
 238         // represented by columns {outtag}*2 and {outtag}*2+1.
 239         outcol := -1
 240         for tag, tvs := range tilelib.variant {
 241                 if len(cmd.dropTiles) > tag && cmd.dropTiles[tag] {
 242                         continue
 243                 }
 244                 tag := tagID(tag)
 245                 outcol++
 246                 // Must shadow outcol var to use safely in goroutine below.
 247                 outcol := outcol
 248                 refstart, ok := tilestart[tag]
 249                 if !ok {
 250                         // Tag didn't place on this reference
 251                         // sequence. (It might place on the same
 252                         // chromosome in a genome anyway, but we don't
 253                         // output the annotations that would result.)
 254                         // outch <- fmt.Sprintf("%d,%d,-1%s\n", tag, outcol, refnamefield)
 255                         continue
 256                 }
 257                 for variant := 1; variant <= len(tvs); variant++ {
 258                         variant, hash := tileVariantID(variant), tvs[variant-1]
 259                         tileseq := tilelib.TileVariantSequence(tileLibRef{Tag: tag, Variant: variant})
 260                         if len(tileseq) == 0 {
 261                                 continue
 262                         } else if len(tileseq) < taglen {
 263                                 return fmt.Errorf("tilevar %d,%d has sequence len %d < taglen %d", tag, variant, len(tileseq), taglen)
 264                         }
 265                         var refpart []byte
 266                         endtag := string(tileseq[len(tileseq)-taglen:])
 267                         if endtagid, ok := cmd.tag2tagid[endtag]; !ok {
 268                                 // Tile variant doesn't end on a tag, so it can only place at the end of a chromosome.
 269                                 refpart = refseq[refstart:]
 270                                 log.Warnf("%x tilevar %d,%d endtag not in ref: %s", hash[:13], tag, variant, endtag)
 271                         } else if refendtagstart, ok := tilestart[endtagid]; !ok {
 272                                 // Ref ends a chromsome with a (possibly very large) variant of this tile, but genomes with this tile don't.
 273                                 // Give up. (TODO: something smarter)
 274                                 log.Debugf("%x not annotating tilevar %d,%d because end tag %d is not in ref", hash[:13], tag, variant, endtagid)
 275                                 continue
 276                         } else {
 277                                 // Non-terminal tile vs. non-terminal reference.
 278                                 refpart = refseq[refstart : refendtagstart+taglen]
 279                                 log.Tracef("\n%x tilevar %d,%d endtag %s endtagid %d refendtagstart %d", hash[:13], tag, variant, endtag, endtagid, refendtagstart)
 280                         }
 281                         if len(refpart) > cmd.maxTileSize {
 282                                 log.Warnf("%x tilevar %d,%d skipping long diff, ref %s seq %s pos %d ref len %d", hash[:13], tag, variant, refname, seqname, refstart, len(refpart))
 283                                 continue
 284                         }
 285                         if len(tileseq) > cmd.maxTileSize {
 286                                 log.Warnf("%x tilevar %d,%d skipping long diff, ref %s seq %s variant len %d", hash[:13], tag, variant, refname, seqname, len(tileseq))
 287                                 continue
 288                         }
 289                         // log.Printf("\n%x @ refstart %d \n< %s\n> %s\n", tv.Blake2b, refstart, refpart, tileseq)
 290
 291                         throttle.Acquire()
 292                         go func() {
 293                                 defer throttle.Release()
 294                                 diffs, _ := hgvs.Diff(strings.ToUpper(string(refpart)), strings.ToUpper(string(tileseq)), 0)
 295                                 for _, diff := range diffs {
 296                                         diff.Position += refstart
 297                                         var varid string
 298                                         if cmd.variantHash {
 299                                                 varid = fmt.Sprintf("%x", hash)[:13]
 300                                         } else {
 301                                                 varid = fmt.Sprintf("%d", variant)
 302                                         }
 303                                         outch <- fmt.Sprintf("%d,%d,%s%s,%s:g.%s\n", tag, outcol, varid, refnamefield, seqname, diff.String())
 304                                         if cmd.reportAnnotation != nil {
 305                                                 cmd.reportAnnotation(tag, outcol, variant, refname, seqname, diff)
 306                                         }
 307                                 }
 308                         }()
 309                 }
 310         }
 311         return nil
 312 }