exportnumpy.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "context"
   6         "errors"
   7         "flag"
   8         "fmt"
   9         "io"
  10         "io/ioutil"
  11         "net/http"
  12         _ "net/http/pprof"
  13         "os"
  14         "sort"
  15         "strings"
  16
  17         "git.arvados.org/arvados.git/sdk/go/arvados"
  18         "github.com/kshedden/gonpy"
  19         log "github.com/sirupsen/logrus"
  20 )
  21
  22 type exportNumpy struct {
  23         filter filter
  24 }
  25
  26 func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  27         var err error
  28         defer func() {
  29                 if err != nil {
  30                         fmt.Fprintf(stderr, "%s\n", err)
  31                 }
  32         }()
  33         flags := flag.NewFlagSet("", flag.ContinueOnError)
  34         flags.SetOutput(stderr)
  35         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  36         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  37         projectUUID := flags.String("project", "", "project `UUID` for output data")
  38         priority := flags.Int("priority", 500, "container request priority")
  39         inputFilename := flags.String("i", "-", "input `file`")
  40         outputFilename := flags.String("o", "-", "output `file`")
  41         annotationsFilename := flags.String("output-annotations", "", "output `file` for tile variant annotations tsv")
  42         librefsFilename := flags.String("output-onehot2tilevar", "", "when using -one-hot, create tsv `file` mapping column# to tag# and variant#")
  43         onehot := flags.Bool("one-hot", false, "recode tile variants as one-hot")
  44         cmd.filter.Flags(flags)
  45         err = flags.Parse(args)
  46         if err == flag.ErrHelp {
  47                 err = nil
  48                 return 0
  49         } else if err != nil {
  50                 return 2
  51         }
  52
  53         if *pprof != "" {
  54                 go func() {
  55                         log.Println(http.ListenAndServe(*pprof, nil))
  56                 }()
  57         }
  58
  59         if !*runlocal {
  60                 if *outputFilename != "-" {
  61                         err = errors.New("cannot specify output file in container mode: not implemented")
  62                         return 1
  63                 }
  64                 runner := arvadosContainerRunner{
  65                         Name:        "lightning export-numpy",
  66                         Client:      arvados.NewClientFromEnv(),
  67                         ProjectUUID: *projectUUID,
  68                         RAM:         128000000000,
  69                         VCPUs:       32,
  70                         Priority:    *priority,
  71                 }
  72                 err = runner.TranslatePaths(inputFilename)
  73                 if err != nil {
  74                         return 1
  75                 }
  76                 runner.Args = []string{"export-numpy", "-local=true",
  77                         fmt.Sprintf("-one-hot=%v", *onehot),
  78                         "-i", *inputFilename,
  79                         "-o", "/mnt/output/matrix.npy",
  80                         "-output-annotations", "/mnt/output/annotations.tsv",
  81                         "-output-onehot2tilevar", "/mnt/output/onehot2tilevar.tsv",
  82                         "-max-variants", fmt.Sprintf("%d", cmd.filter.MaxVariants),
  83                         "-min-coverage", fmt.Sprintf("%f", cmd.filter.MinCoverage),
  84                         "-max-tag", fmt.Sprintf("%d", cmd.filter.MaxTag),
  85                 }
  86                 var output string
  87                 output, err = runner.Run()
  88                 if err != nil {
  89                         return 1
  90                 }
  91                 fmt.Fprintln(stdout, output+"/matrix.npy")
  92                 return 0
  93         }
  94
  95         var input io.ReadCloser
  96         if *inputFilename == "-" {
  97                 input = ioutil.NopCloser(stdin)
  98         } else {
  99                 input, err = os.Open(*inputFilename)
 100                 if err != nil {
 101                         return 1
 102                 }
 103                 defer input.Close()
 104         }
 105         tilelib := &tileLibrary{
 106                 retainNoCalls:       true,
 107                 retainTileSequences: true,
 108                 compactGenomes:      map[string][]tileVariantID{},
 109         }
 110         err = tilelib.LoadGob(context.Background(), input, strings.HasSuffix(*inputFilename, ".gz"), nil)
 111         if err != nil {
 112                 return 1
 113         }
 114         err = input.Close()
 115         if err != nil {
 116                 return 1
 117         }
 118
 119         log.Info("filtering")
 120         cmd.filter.Apply(tilelib)
 121         log.Info("tidying")
 122         tilelib.Tidy()
 123
 124         if *annotationsFilename != "" {
 125                 log.Infof("writing annotations")
 126                 var annow io.WriteCloser
 127                 annow, err = os.OpenFile(*annotationsFilename, os.O_CREATE|os.O_WRONLY, 0666)
 128                 if err != nil {
 129                         return 1
 130                 }
 131                 defer annow.Close()
 132                 err = (&annotatecmd{maxTileSize: 5000}).exportTileDiffs(annow, tilelib)
 133                 if err != nil {
 134                         return 1
 135                 }
 136                 err = annow.Close()
 137                 if err != nil {
 138                         return 1
 139                 }
 140         }
 141
 142         log.Info("building numpy array")
 143         out, rows, cols := cgs2array(tilelib)
 144
 145         log.Info("writing numpy file")
 146         var output io.WriteCloser
 147         if *outputFilename == "-" {
 148                 output = nopCloser{stdout}
 149         } else {
 150                 output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
 151                 if err != nil {
 152                         return 1
 153                 }
 154                 defer output.Close()
 155         }
 156         bufw := bufio.NewWriter(output)
 157         npw, err := gonpy.NewWriter(nopCloser{bufw})
 158         if err != nil {
 159                 return 1
 160         }
 161         if *onehot {
 162                 log.Info("recoding to onehot")
 163                 recoded, librefs, recodedcols := recodeOnehot(out, cols)
 164                 out, cols = recoded, recodedcols
 165                 if *librefsFilename != "" {
 166                         log.Infof("writing onehot column mapping")
 167                         err = cmd.writeLibRefs(*librefsFilename, tilelib, librefs)
 168                         if err != nil {
 169                                 return 1
 170                         }
 171                 }
 172         }
 173         log.Info("writing numpy")
 174         npw.Shape = []int{rows, cols}
 175         npw.WriteInt16(out)
 176         err = bufw.Flush()
 177         if err != nil {
 178                 return 1
 179         }
 180         err = output.Close()
 181         if err != nil {
 182                 return 1
 183         }
 184         return 0
 185 }
 186
 187 func (*exportNumpy) writeLibRefs(fnm string, tilelib *tileLibrary, librefs []tileLibRef) error {
 188         f, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY, 0666)
 189         if err != nil {
 190                 return err
 191         }
 192         defer f.Close()
 193         for i, libref := range librefs {
 194                 _, err = fmt.Fprintf(f, "%d\t%d\t%d\n", i, libref.Tag, libref.Variant)
 195                 if err != nil {
 196                         return err
 197                 }
 198         }
 199         return f.Close()
 200 }
 201
 202 func cgs2array(tilelib *tileLibrary) (data []int16, rows, cols int) {
 203         var cgnames []string
 204         for name := range tilelib.compactGenomes {
 205                 cgnames = append(cgnames, name)
 206         }
 207         sort.Strings(cgnames)
 208
 209         rows = len(tilelib.compactGenomes)
 210         for _, cg := range tilelib.compactGenomes {
 211                 if cols < len(cg) {
 212                         cols = len(cg)
 213                 }
 214         }
 215
 216         // flag low-quality tile variants so we can change to -1 below
 217         lowqual := make([]map[tileVariantID]bool, cols/2)
 218         for tag, variants := range tilelib.variant {
 219                 lq := lowqual[tag]
 220                 for varidx, hash := range variants {
 221                         if len(tilelib.seq[hash]) == 0 {
 222                                 if lq == nil {
 223                                         lq = map[tileVariantID]bool{}
 224                                         lowqual[tag] = lq
 225                                 }
 226                                 lq[tileVariantID(varidx+1)] = true
 227                         }
 228                 }
 229         }
 230
 231         data = make([]int16, rows*cols)
 232         for row, name := range cgnames {
 233                 for i, v := range tilelib.compactGenomes[name] {
 234                         if v > 0 && lowqual[i/2][v] {
 235                                 data[row*cols+i] = -1
 236                         } else {
 237                                 data[row*cols+i] = int16(v)
 238                         }
 239                 }
 240         }
 241
 242         return
 243 }
 244
 245 func recodeOnehot(in []int16, incols int) (out []int16, librefs []tileLibRef, outcols int) {
 246         rows := len(in) / incols
 247         maxvalue := make([]int16, incols)
 248         for row := 0; row < rows; row++ {
 249                 for col := 0; col < incols; col++ {
 250                         if v := in[row*incols+col]; maxvalue[col] < v {
 251                                 maxvalue[col] = v
 252                         }
 253                 }
 254         }
 255         outcol := make([]int, incols)
 256         dropped := 0
 257         for incol, maxv := range maxvalue {
 258                 outcol[incol] = outcols
 259                 if maxv == 0 {
 260                         dropped++
 261                 }
 262                 for v := 1; v <= int(maxv); v++ {
 263                         librefs = append(librefs, tileLibRef{Tag: tagID(incol), Variant: tileVariantID(v)})
 264                         outcols++
 265                 }
 266         }
 267         log.Printf("recodeOnehot: dropped %d input cols with zero maxvalue", dropped)
 268
 269         out = make([]int16, rows*outcols)
 270         for inidx, row := 0, 0; row < rows; row++ {
 271                 outrow := out[row*outcols:]
 272                 for col := 0; col < incols; col++ {
 273                         if v := in[inidx]; v > 0 {
 274                                 outrow[outcol[col]+int(v)-1] = 1
 275                         }
 276                         inidx++
 277                 }
 278         }
 279         return
 280 }
 281
 282 type nopCloser struct {
 283         io.Writer
 284 }
 285
 286 func (nopCloser) Close() error { return nil }