slicenumpy.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         "runtime"
  17         "sort"
  18         "strings"
  19
  20         "git.arvados.org/arvados.git/sdk/go/arvados"
  21         "github.com/arvados/lightning/hgvs"
  22         "github.com/kshedden/gonpy"
  23         "github.com/sirupsen/logrus"
  24         log "github.com/sirupsen/logrus"
  25 )
  26
  27 type sliceNumpy struct {
  28         filter filter
  29 }
  30
  31 func (cmd *sliceNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  32         var err error
  33         defer func() {
  34                 if err != nil {
  35                         fmt.Fprintf(stderr, "%s\n", err)
  36                 }
  37         }()
  38         flags := flag.NewFlagSet("", flag.ContinueOnError)
  39         flags.SetOutput(stderr)
  40         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  41         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  42         projectUUID := flags.String("project", "", "project `UUID` for output data")
  43         priority := flags.Int("priority", 500, "container request priority")
  44         inputDir := flags.String("input-dir", "./in", "input `directory`")
  45         outputDir := flags.String("output-dir", "./out", "output `directory`")
  46         ref := flags.String("ref", "", "reference name (if blank, choose last one that appears in input)")
  47         regionsFilename := flags.String("regions", "", "only output columns/annotations that intersect regions in specified bed `file`")
  48         expandRegions := flags.Int("expand-regions", 0, "expand specified regions by `N` base pairs on each side`")
  49         cmd.filter.Flags(flags)
  50         err = flags.Parse(args)
  51         if err == flag.ErrHelp {
  52                 err = nil
  53                 return 0
  54         } else if err != nil {
  55                 return 2
  56         }
  57
  58         if *pprof != "" {
  59                 go func() {
  60                         log.Println(http.ListenAndServe(*pprof, nil))
  61                 }()
  62         }
  63
  64         if !*runlocal {
  65                 runner := arvadosContainerRunner{
  66                         Name:        "lightning slice-numpy",
  67                         Client:      arvados.NewClientFromEnv(),
  68                         ProjectUUID: *projectUUID,
  69                         RAM:         250000000000,
  70                         VCPUs:       32,
  71                         Priority:    *priority,
  72                         KeepCache:   1,
  73                         APIAccess:   true,
  74                 }
  75                 err = runner.TranslatePaths(inputDir, regionsFilename)
  76                 if err != nil {
  77                         return 1
  78                 }
  79                 runner.Args = []string{"slice-numpy", "-local=true",
  80                         "-pprof", ":6060",
  81                         "-input-dir", *inputDir,
  82                         "-output-dir", "/mnt/output",
  83                         "-regions", *regionsFilename,
  84                         "-expand-regions", fmt.Sprintf("%d", *expandRegions),
  85                 }
  86                 runner.Args = append(runner.Args, cmd.filter.Args()...)
  87                 var output string
  88                 output, err = runner.Run()
  89                 if err != nil {
  90                         return 1
  91                 }
  92                 fmt.Fprintln(stdout, output)
  93                 return 0
  94         }
  95
  96         infiles, err := allGobFiles(*inputDir)
  97         if err != nil {
  98                 return 1
  99         }
 100         if len(infiles) == 0 {
 101                 err = fmt.Errorf("no input files found in %s", *inputDir)
 102                 return 1
 103         }
 104         sort.Strings(infiles)
 105
 106         var cgnames []string
 107         var refseq map[string][]tileLibRef
 108         in0, err := open(infiles[0])
 109         if err != nil {
 110                 return 1
 111         }
 112         taglen := -1
 113         DecodeLibrary(in0, strings.HasSuffix(infiles[0], ".gz"), func(ent *LibraryEntry) error {
 114                 if len(ent.TagSet) > 0 {
 115                         taglen = len(ent.TagSet[0])
 116                 }
 117                 for _, cseq := range ent.CompactSequences {
 118                         if cseq.Name == *ref || *ref == "" {
 119                                 refseq = cseq.TileSequences
 120                         }
 121                 }
 122                 for _, cg := range ent.CompactGenomes {
 123                         cgnames = append(cgnames, cg.Name)
 124                 }
 125                 return nil
 126         })
 127         if err != nil {
 128                 return 1
 129         }
 130         in0.Close()
 131         if refseq == nil {
 132                 err = fmt.Errorf("%s: reference sequence not found", infiles[0])
 133                 return 1
 134         }
 135         if taglen < 0 {
 136                 err = fmt.Errorf("tagset not found")
 137                 return 1
 138         }
 139         sort.Strings(cgnames)
 140
 141         {
 142                 labelsFilename := *outputDir + "/labels.csv"
 143                 log.Infof("writing labels to %s", labelsFilename)
 144                 var f *os.File
 145                 f, err = os.Create(labelsFilename)
 146                 if err != nil {
 147                         return 1
 148                 }
 149                 defer f.Close()
 150                 for i, name := range cgnames {
 151                         _, err = fmt.Fprintf(f, "%d,%q\n", i, trimFilenameForLabel(name))
 152                         if err != nil {
 153                                 err = fmt.Errorf("write %s: %w", labelsFilename, err)
 154                                 return 1
 155                         }
 156                 }
 157                 err = f.Close()
 158                 if err != nil {
 159                         err = fmt.Errorf("close %s: %w", labelsFilename, err)
 160                         return 1
 161                 }
 162         }
 163
 164         log.Info("building list of reference tiles to load") // TODO: more efficient if we had saved all ref tiles in slice0
 165         type reftileinfo struct {
 166                 seqname  string // chr1
 167                 pos      int    // distance from start of chr1 to start of tile
 168                 tiledata []byte // acgtggcaa...
 169         }
 170         reftile := map[tagID]*reftileinfo{}
 171         for seqname, cseq := range refseq {
 172                 for _, libref := range cseq {
 173                         reftile[libref.Tag] = &reftileinfo{seqname: seqname}
 174                 }
 175         }
 176         log.Info("loading reference tiles from all slices")
 177         throttle := throttle{Max: runtime.NumCPU()}
 178         for _, infile := range infiles {
 179                 infile := infile
 180                 throttle.Go(func() error {
 181                         defer log.Infof("%s: done", infile)
 182                         f, err := open(infile)
 183                         if err != nil {
 184                                 return err
 185                         }
 186                         defer f.Close()
 187                         return DecodeLibrary(f, strings.HasSuffix(infile, ".gz"), func(ent *LibraryEntry) error {
 188                                 for _, tv := range ent.TileVariants {
 189                                         if dst, ok := reftile[tv.Tag]; ok {
 190                                                 dst.tiledata = tv.Sequence
 191                                         }
 192                                 }
 193                                 return nil
 194                         })
 195                 })
 196         }
 197         throttle.Wait()
 198
 199         log.Info("reconstructing reference sequences")
 200         for seqname, cseq := range refseq {
 201                 seqname, cseq := seqname, cseq
 202                 throttle.Go(func() error {
 203                         defer log.Printf("... %s done", seqname)
 204                         pos := 0
 205                         for _, libref := range cseq {
 206                                 rt := reftile[libref.Tag]
 207                                 rt.pos = pos
 208                                 pos += len(rt.tiledata) - taglen
 209                         }
 210                         return nil
 211                 })
 212         }
 213         throttle.Wait()
 214
 215         log.Info("TODO: determining which tiles intersect given regions")
 216
 217         log.Info("generating annotations and numpy matrix for each slice")
 218         for infileIdx, infile := range infiles {
 219                 infileIdx, infile := infileIdx, infile
 220                 throttle.Go(func() error {
 221                         defer log.Infof("%s: done", infile)
 222                         seq := map[tagID][][]byte{}
 223                         cgs := make(map[string]CompactGenome, len(cgnames))
 224                         f, err := open(infile)
 225                         if err != nil {
 226                                 return err
 227                         }
 228                         defer f.Close()
 229                         err = DecodeLibrary(f, strings.HasSuffix(infile, ".gz"), func(ent *LibraryEntry) error {
 230                                 for _, tv := range ent.TileVariants {
 231                                         variants := seq[tv.Tag]
 232                                         for len(variants) <= int(tv.Variant) {
 233                                                 variants = append(variants, nil)
 234                                         }
 235                                         variants[int(tv.Variant)] = tv.Sequence
 236                                         seq[tv.Tag] = variants
 237                                 }
 238                                 for _, cg := range ent.CompactGenomes {
 239                                         cgs[cg.Name] = cg
 240                                 }
 241                                 return nil
 242                         })
 243                         if err != nil {
 244                                 return err
 245                         }
 246                         tagstart := cgs[cgnames[0]].StartTag
 247                         tagend := cgs[cgnames[0]].EndTag
 248
 249                         log.Infof("TODO: %s: filtering", infile)
 250                         log.Infof("TODO: %s: tidying", infile)
 251                         log.Infof("TODO: %s: lowqual to -1", infile)
 252
 253                         annotationsFilename := fmt.Sprintf("%s/matrix.%04d.annotations.csv", *outputDir, infileIdx)
 254                         log.Infof("%s: writing annotations to %s", infile, annotationsFilename)
 255                         annof, err := os.Create(annotationsFilename)
 256                         if err != nil {
 257                                 return err
 258                         }
 259                         annow := bufio.NewWriterSize(annof, 1<<20)
 260                         for tag, variants := range seq {
 261                                 rt, ok := reftile[tag]
 262                                 if !ok {
 263                                         // Reference does not use any
 264                                         // variant of this tile.
 265                                         // TODO: log this? mention it
 266                                         // in annotations?
 267                                         continue
 268                                 }
 269                                 outcol := tag - tagID(tagstart)
 270                                 reftilestr := strings.ToUpper(string(rt.tiledata))
 271                                 for v, seq := range variants {
 272                                         if len(seq) == 0 || !bytes.HasSuffix(rt.tiledata, seq[len(seq)-taglen:]) {
 273                                                 continue
 274                                         }
 275                                         if lendiff := len(rt.tiledata) - len(seq); lendiff < -1000 || lendiff > 1000 {
 276                                                 continue
 277                                         }
 278                                         diffs, _ := hgvs.Diff(reftilestr, strings.ToUpper(string(seq)), 0)
 279                                         for _, diff := range diffs {
 280                                                 diff.Position += rt.pos
 281                                                 fmt.Fprintf(annow, "%d,%d,%d,%s:g.%s\n", tag, outcol, v, rt.seqname, diff.String())
 282                                         }
 283                                 }
 284                         }
 285                         err = annow.Flush()
 286                         if err != nil {
 287                                 return err
 288                         }
 289                         err = annof.Close()
 290                         if err != nil {
 291                                 return err
 292                         }
 293
 294                         log.Infof("%s: preparing numpy", infile)
 295                         rows := len(cgnames)
 296                         cols := 2 * int(tagend-tagstart)
 297                         out := make([]int16, rows*cols)
 298                         for row, name := range cgnames {
 299                                 out := out[row*cols:]
 300                                 for col, v := range cgs[name].Variants {
 301                                         if variants, ok := seq[tagstart+tagID(col/2)]; ok && len(variants) > int(v) && len(variants[v]) > 0 {
 302                                                 out[col] = int16(v)
 303                                         } else {
 304                                                 out[col] = -1
 305                                         }
 306                                 }
 307                         }
 308
 309                         fnm := fmt.Sprintf("%s/matrix.%04d.npy", *outputDir, infileIdx)
 310                         output, err := os.Create(fnm)
 311                         if err != nil {
 312                                 return err
 313                         }
 314                         defer output.Close()
 315                         bufw := bufio.NewWriter(output)
 316                         npw, err := gonpy.NewWriter(nopCloser{bufw})
 317                         if err != nil {
 318                                 return err
 319                         }
 320                         log.WithFields(logrus.Fields{
 321                                 "filename": fnm,
 322                                 "rows":     rows,
 323                                 "cols":     cols,
 324                         }).Info("writing numpy")
 325                         npw.Shape = []int{rows, cols}
 326                         npw.WriteInt16(out)
 327                         err = bufw.Flush()
 328                         if err != nil {
 329                                 return err
 330                         }
 331                         err = output.Close()
 332                         if err != nil {
 333                                 return err
 334                         }
 335                         return nil
 336                 })
 337         }
 338         if err = throttle.Wait(); err != nil {
 339                 return 1
 340         }
 341         return 0
 342 }
 343
 344 func (*sliceNumpy) writeLibRefs(fnm string, tilelib *tileLibrary, librefs []tileLibRef) error {
 345         f, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY, 0666)
 346         if err != nil {
 347                 return err
 348         }
 349         defer f.Close()
 350         for i, libref := range librefs {
 351                 _, err = fmt.Fprintf(f, "%d,%d,%d\n", i, libref.Tag, libref.Variant)
 352                 if err != nil {
 353                         return err
 354                 }
 355         }
 356         return f.Close()
 357 }