dump.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         "bytes"
  10         "flag"
  11         "fmt"
  12         "io"
  13         "net/http"
  14         _ "net/http/pprof"
  15         "os"
  16         "regexp"
  17         "runtime"
  18         "sort"
  19         "strconv"
  20         "strings"
  21         "sync"
  22         "sync/atomic"
  23
  24         "git.arvados.org/arvados.git/sdk/go/arvados"
  25         log "github.com/sirupsen/logrus"
  26         "golang.org/x/crypto/blake2b"
  27 )
  28
  29 type dump struct {
  30         filter       filter
  31         cgnames      []string
  32         selectedTags map[tagID]bool
  33 }
  34
  35 func (cmd *dump) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  36         err := cmd.run(prog, args, stdin, stdout, stderr)
  37         if err != nil {
  38                 return 1
  39         }
  40         return 0
  41 }
  42
  43 func (cmd *dump) run(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) error {
  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         inputDir := flags.String("input-dir", "./in", "input `directory`")
  51         outputDir := flags.String("output-dir", "./out", "output `directory`")
  52         ref := flags.String("ref", "", "reference name (if blank, choose last one that appears in input)")
  53         regionsFilename := flags.String("regions", "", "only output columns/annotations that intersect regions in specified bed `file`")
  54         expandRegions := flags.Int("expand-regions", 0, "expand specified regions by `N` base pairs on each side`")
  55         selectedTags := flags.String("tags", "", "tag numbers to dump")
  56         cmd.filter.Flags(flags)
  57         err := flags.Parse(args)
  58         if err == flag.ErrHelp {
  59                 return nil
  60         } else if err != nil {
  61                 return err
  62         }
  63
  64         if *pprof != "" {
  65                 go func() {
  66                         log.Println(http.ListenAndServe(*pprof, nil))
  67                 }()
  68         }
  69
  70         if !*runlocal {
  71                 runner := arvadosContainerRunner{
  72                         Name:        "lightning slice-numpy",
  73                         Client:      arvados.NewClientFromEnv(),
  74                         ProjectUUID: *projectUUID,
  75                         RAM:         750000000000,
  76                         VCPUs:       96,
  77                         Priority:    *priority,
  78                         KeepCache:   2,
  79                         APIAccess:   true,
  80                 }
  81                 err = runner.TranslatePaths(inputDir, regionsFilename)
  82                 if err != nil {
  83                         return err
  84                 }
  85                 runner.Args = []string{"dump", "-local=true",
  86                         "-pprof=:6060",
  87                         "-input-dir=" + *inputDir,
  88                         "-output-dir=/mnt/output",
  89                         "-regions=" + *regionsFilename,
  90                         "-expand-regions=" + fmt.Sprintf("%d", *expandRegions),
  91                         "-tags=" + *selectedTags,
  92                 }
  93                 runner.Args = append(runner.Args, cmd.filter.Args()...)
  94                 output, err := runner.Run()
  95                 if err != nil {
  96                         return err
  97                 }
  98                 fmt.Fprintln(stdout, output)
  99                 return nil
 100         }
 101
 102         if *selectedTags != "" {
 103                 cmd.selectedTags = map[tagID]bool{}
 104                 for _, tagstr := range strings.Split(*selectedTags, ",") {
 105                         tag, err := strconv.ParseInt(tagstr, 10, 64)
 106                         if err != nil {
 107                                 return err
 108                         }
 109                         cmd.selectedTags[tagID(tag)] = true
 110                 }
 111         }
 112
 113         infiles, err := allFiles(*inputDir, matchGobFile)
 114         if err != nil {
 115                 return err
 116         }
 117         if len(infiles) == 0 {
 118                 return fmt.Errorf("no input files found in %s", *inputDir)
 119         }
 120         sort.Strings(infiles)
 121
 122         var refseq map[string][]tileLibRef
 123         var reftiledata = make(map[tileLibRef][]byte, 11000000)
 124         in0, err := open(infiles[0])
 125         if err != nil {
 126                 return err
 127         }
 128
 129         matchGenome, err := regexp.Compile(cmd.filter.MatchGenome)
 130         if err != nil {
 131                 err = fmt.Errorf("-match-genome: invalid regexp: %q", cmd.filter.MatchGenome)
 132                 return err
 133         }
 134
 135         cmd.cgnames = nil
 136         taglen := -1
 137         DecodeLibrary(in0, strings.HasSuffix(infiles[0], ".gz"), func(ent *LibraryEntry) error {
 138                 if len(ent.TagSet) > 0 {
 139                         taglen = len(ent.TagSet[0])
 140                 }
 141                 for _, cseq := range ent.CompactSequences {
 142                         if cseq.Name == *ref || *ref == "" {
 143                                 refseq = cseq.TileSequences
 144                         }
 145                 }
 146                 for _, cg := range ent.CompactGenomes {
 147                         if matchGenome.MatchString(cg.Name) {
 148                                 cmd.cgnames = append(cmd.cgnames, cg.Name)
 149                         }
 150                 }
 151                 for _, tv := range ent.TileVariants {
 152                         if tv.Ref {
 153                                 reftiledata[tileLibRef{tv.Tag, tv.Variant}] = tv.Sequence
 154                         }
 155                 }
 156                 return nil
 157         })
 158         if err != nil {
 159                 return err
 160         }
 161         in0.Close()
 162         if refseq == nil {
 163                 return fmt.Errorf("%s: reference sequence not found", infiles[0])
 164         }
 165         if taglen < 0 {
 166                 return fmt.Errorf("tagset not found")
 167         }
 168         if len(cmd.cgnames) == 0 {
 169                 return fmt.Errorf("no genomes found matching regexp %q", cmd.filter.MatchGenome)
 170         }
 171         sort.Strings(cmd.cgnames)
 172
 173         log.Info("indexing reference tiles")
 174         type reftileinfo struct {
 175                 variant  tileVariantID
 176                 seqname  string // chr1
 177                 pos      int    // distance from start of chromosome to starttag
 178                 tiledata []byte // acgtggcaa...
 179         }
 180         isdup := map[tagID]bool{}
 181         reftile := map[tagID]*reftileinfo{}
 182         for seqname, cseq := range refseq {
 183                 pos := 0
 184                 for _, libref := range cseq {
 185                         tiledata := reftiledata[libref]
 186                         if len(tiledata) == 0 {
 187                                 return fmt.Errorf("missing tiledata for tag %d variant %d in %s in ref", libref.Tag, libref.Variant, seqname)
 188                         }
 189                         if isdup[libref.Tag] {
 190                                 log.Printf("dropping reference tile %+v from %s @ %d, tag not unique", libref, seqname, pos)
 191                         } else if reftile[libref.Tag] != nil {
 192                                 log.Printf("dropping reference tile %+v from %s @ %d, tag not unique", tileLibRef{Tag: libref.Tag, Variant: reftile[libref.Tag].variant}, reftile[libref.Tag].seqname, reftile[libref.Tag].pos)
 193                                 delete(reftile, libref.Tag)
 194                                 log.Printf("dropping reference tile %+v from %s @ %d, tag not unique", libref, seqname, pos)
 195                                 isdup[libref.Tag] = true
 196                         } else {
 197                                 reftile[libref.Tag] = &reftileinfo{
 198                                         seqname:  seqname,
 199                                         variant:  libref.Variant,
 200                                         tiledata: tiledata,
 201                                         pos:      pos,
 202                                 }
 203                         }
 204                         pos += len(tiledata) - taglen
 205                 }
 206                 log.Printf("... %s done, len %d", seqname, pos+taglen)
 207         }
 208
 209         var mask *mask
 210         if *regionsFilename != "" {
 211                 log.Printf("loading regions from %s", *regionsFilename)
 212                 mask, err = makeMask(*regionsFilename, *expandRegions)
 213                 if err != nil {
 214                         return err
 215                 }
 216                 log.Printf("before applying mask, len(reftile) == %d", len(reftile))
 217                 log.Printf("deleting reftile entries for regions outside %d intervals", mask.Len())
 218                 for tag, rt := range reftile {
 219                         if !mask.Check(strings.TrimPrefix(rt.seqname, "chr"), rt.pos, rt.pos+len(rt.tiledata)) {
 220                                 delete(reftile, tag)
 221                         }
 222                 }
 223                 log.Printf("after applying mask, len(reftile) == %d", len(reftile))
 224         }
 225
 226         if cmd.selectedTags != nil {
 227                 log.Printf("deleting reftile entries other than %d selected tags", len(cmd.selectedTags))
 228                 for tag := range reftile {
 229                         if !cmd.selectedTags[tag] {
 230                                 delete(reftile, tag)
 231                         }
 232                 }
 233                 log.Printf("after applying selected tags, len(reftile) == %d", len(reftile))
 234         }
 235
 236         dumpVariantsName := fmt.Sprintf("%s/variants.csv", *outputDir)
 237         log.Infof("writing %s", dumpVariantsName)
 238         dumpVariantsF, err := os.Create(dumpVariantsName)
 239         if err != nil {
 240                 return err
 241         }
 242         dumpVariantsW := bufio.NewWriterSize(dumpVariantsF, 1<<20)
 243         mtx := sync.Mutex{}
 244
 245         throttleMem := throttle{Max: runtime.GOMAXPROCS(0)}
 246         log.Infof("reading %d slices with max concurrency %d", len(infiles), throttleMem.Max)
 247         var done int64
 248         for infileIdx, infile := range infiles {
 249                 infileIdx, infile := infileIdx, infile
 250                 throttleMem.Go(func() error {
 251                         seq := make(map[tagID][]TileVariant, 50000)
 252                         cgs := make(map[string]CompactGenome, len(cmd.cgnames))
 253                         f, err := open(infile)
 254                         if err != nil {
 255                                 return err
 256                         }
 257                         defer f.Close()
 258                         log.Infof("%04d: reading %s", infileIdx, infile)
 259                         err = DecodeLibrary(f, strings.HasSuffix(infile, ".gz"), func(ent *LibraryEntry) error {
 260                                 for _, tv := range ent.TileVariants {
 261                                         if tv.Ref {
 262                                                 continue
 263                                         }
 264                                         if mask != nil && reftile[tv.Tag] == nil {
 265                                                 // Don't waste
 266                                                 // time/memory on
 267                                                 // masked-out tiles.
 268                                                 continue
 269                                         }
 270                                         variants := seq[tv.Tag]
 271                                         if len(variants) == 0 {
 272                                                 variants = make([]TileVariant, 100)
 273                                         }
 274                                         for len(variants) <= int(tv.Variant) {
 275                                                 variants = append(variants, TileVariant{})
 276                                         }
 277                                         variants[int(tv.Variant)] = tv
 278                                         seq[tv.Tag] = variants
 279                                 }
 280                                 for _, cg := range ent.CompactGenomes {
 281                                         if !matchGenome.MatchString(cg.Name) {
 282                                                 continue
 283                                         }
 284                                         // pad to full slice size
 285                                         // to avoid out-of-bounds
 286                                         // checks later
 287                                         if sliceSize := 2 * int(cg.EndTag-cg.StartTag); len(cg.Variants) < sliceSize {
 288                                                 cg.Variants = append(cg.Variants, make([]tileVariantID, sliceSize-len(cg.Variants))...)
 289                                         }
 290                                         cgs[cg.Name] = cg
 291                                 }
 292                                 return nil
 293                         })
 294                         if err != nil {
 295                                 return err
 296                         }
 297                         tagstart := cgs[cmd.cgnames[0]].StartTag
 298                         tagend := cgs[cmd.cgnames[0]].EndTag
 299
 300                         // TODO: filters
 301
 302                         log.Infof("%04d: renumber/dedup variants for tags %d-%d", infileIdx, tagstart, tagend)
 303                         variantRemap := make([][]tileVariantID, tagend-tagstart)
 304                         throttleCPU := throttle{Max: runtime.GOMAXPROCS(0)}
 305                         for tag, variants := range seq {
 306                                 tag, variants := tag, variants
 307                                 throttleCPU.Go(func() error {
 308                                         count := make(map[[blake2b.Size256]byte]int, len(variants))
 309
 310                                         rt := reftile[tag]
 311                                         var rthash [blake2b.Size256]byte
 312                                         if rt != nil {
 313                                                 rthash = blake2b.Sum256(rt.tiledata)
 314                                                 count[rthash] = 0
 315                                         }
 316
 317                                         for _, cg := range cgs {
 318                                                 idx := int(tag-tagstart) * 2
 319                                                 for allele := 0; allele < 2; allele++ {
 320                                                         v := cg.Variants[idx+allele]
 321                                                         if v > 0 && len(variants[v].Sequence) > 0 {
 322                                                                 count[variants[v].Blake2b]++
 323                                                         }
 324                                                 }
 325                                         }
 326                                         // hash[i] will be the hash of
 327                                         // the variant(s) that should
 328                                         // be at rank i (0-based).
 329                                         hash := make([][blake2b.Size256]byte, 0, len(count))
 330                                         for b := range count {
 331                                                 hash = append(hash, b)
 332                                         }
 333                                         sort.Slice(hash, func(i, j int) bool {
 334                                                 bi, bj := &hash[i], &hash[j]
 335                                                 if ci, cj := count[*bi], count[*bj]; ci != cj {
 336                                                         return ci > cj
 337                                                 } else {
 338                                                         return bytes.Compare((*bi)[:], (*bj)[:]) < 0
 339                                                 }
 340                                         })
 341                                         // rank[b] will be the 1-based
 342                                         // new variant number for
 343                                         // variants whose hash is b.
 344                                         rank := make(map[[blake2b.Size256]byte]tileVariantID, len(hash))
 345                                         for i, h := range hash {
 346                                                 rank[h] = tileVariantID(i + 1)
 347                                         }
 348                                         // remap[v] will be the new
 349                                         // variant number for original
 350                                         // variant number v.
 351                                         remap := make([]tileVariantID, len(variants))
 352                                         for i, tv := range variants {
 353                                                 remap[i] = rank[tv.Blake2b]
 354                                         }
 355                                         variantRemap[tag-tagstart] = remap
 356                                         if rt != nil {
 357                                                 rt.variant = rank[rthash]
 358                                         }
 359                                         return nil
 360                                 })
 361                         }
 362                         throttleCPU.Wait()
 363
 364                         for tag := tagstart; tag < tagend; tag++ {
 365                                 rt, ok := reftile[tag]
 366                                 if !ok {
 367                                         // Excluded by specified
 368                                         // regions, or reference does
 369                                         // not use any variant of this
 370                                         // tile. (TODO: log this?
 371                                         // mention it in annotations?)
 372                                         continue
 373                                 }
 374                                 remap := variantRemap[tag-tagstart]
 375                                 maxv := tileVariantID(0)
 376                                 for _, v := range remap {
 377                                         if maxv < v {
 378                                                 maxv = v
 379                                         }
 380                                 }
 381                                 variants := seq[tag]
 382
 383                                 mtx.Lock()
 384                                 fmt.Fprintf(dumpVariantsW, "%d,%d,1,%s,%d,%s\n", tag, rt.variant, rt.seqname, rt.pos+1, bytes.ToUpper(rt.tiledata))
 385                                 mtx.Unlock()
 386
 387                                 done := make([]bool, maxv+1)
 388                                 for v, tv := range variants {
 389                                         v := remap[v]
 390                                         if v == 0 || v == rt.variant || done[v] {
 391                                                 continue
 392                                         } else {
 393                                                 done[v] = true
 394                                         }
 395                                         mtx.Lock()
 396                                         fmt.Fprintf(dumpVariantsW, "%d,%d,0,%s,%d,%s\n", tag, v, rt.seqname, rt.pos+1, bytes.ToUpper(tv.Sequence))
 397                                         mtx.Unlock()
 398                                 }
 399                         }
 400                         log.Infof("%s: done (%d/%d)", infile, int(atomic.AddInt64(&done, 1)), len(infiles))
 401                         return nil
 402                 })
 403         }
 404         if err = throttleMem.Wait(); err != nil {
 405                 return err
 406         }
 407         err = dumpVariantsW.Flush()
 408         if err != nil {
 409                 return err
 410         }
 411         err = dumpVariantsF.Close()
 412         if err != nil {
 413                 return err
 414         }
 415         return nil
 416 }