exportnumpy.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "errors"
   6         "flag"
   7         "fmt"
   8         "io"
   9         "io/ioutil"
  10         "net/http"
  11         _ "net/http/pprof"
  12         "os"
  13         "sort"
  14
  15         "git.arvados.org/arvados.git/sdk/go/arvados"
  16         "github.com/kshedden/gonpy"
  17         log "github.com/sirupsen/logrus"
  18 )
  19
  20 type exportNumpy struct{}
  21
  22 func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  23         var err error
  24         defer func() {
  25                 if err != nil {
  26                         fmt.Fprintf(stderr, "%s\n", err)
  27                 }
  28         }()
  29         flags := flag.NewFlagSet("", flag.ContinueOnError)
  30         flags.SetOutput(stderr)
  31         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  32         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  33         projectUUID := flags.String("project", "", "project `UUID` for output data")
  34         priority := flags.Int("priority", 500, "container request priority")
  35         inputFilename := flags.String("i", "-", "input `file`")
  36         outputFilename := flags.String("o", "-", "output `file`")
  37         onehot := flags.Bool("one-hot", false, "recode tile variants as one-hot")
  38         err = flags.Parse(args)
  39         if err == flag.ErrHelp {
  40                 err = nil
  41                 return 0
  42         } else if err != nil {
  43                 return 2
  44         }
  45
  46         if *pprof != "" {
  47                 go func() {
  48                         log.Println(http.ListenAndServe(*pprof, nil))
  49                 }()
  50         }
  51
  52         if !*runlocal {
  53                 if *outputFilename != "-" {
  54                         err = errors.New("cannot specify output file in container mode: not implemented")
  55                         return 1
  56                 }
  57                 runner := arvadosContainerRunner{
  58                         Name:        "lightning export-numpy",
  59                         Client:      arvados.NewClientFromEnv(),
  60                         ProjectUUID: *projectUUID,
  61                         RAM:         128000000000,
  62                         VCPUs:       2,
  63                         Priority:    *priority,
  64                 }
  65                 err = runner.TranslatePaths(inputFilename)
  66                 if err != nil {
  67                         return 1
  68                 }
  69                 runner.Args = []string{"export-numpy", "-local=true", fmt.Sprintf("-one-hot=%v", *onehot), "-i", *inputFilename, "-o", "/mnt/output/library.npy"}
  70                 var output string
  71                 output, err = runner.Run()
  72                 if err != nil {
  73                         return 1
  74                 }
  75                 fmt.Fprintln(stdout, output+"/library.npy")
  76                 return 0
  77         }
  78
  79         var input io.ReadCloser
  80         if *inputFilename == "-" {
  81                 input = ioutil.NopCloser(stdin)
  82         } else {
  83                 input, err = os.Open(*inputFilename)
  84                 if err != nil {
  85                         return 1
  86                 }
  87                 defer input.Close()
  88         }
  89         cgs, err := ReadCompactGenomes(input)
  90         if err != nil {
  91                 return 1
  92         }
  93         err = input.Close()
  94         if err != nil {
  95                 return 1
  96         }
  97         sort.Slice(cgs, func(i, j int) bool { return cgs[i].Name < cgs[j].Name })
  98
  99         out, rows, cols := cgs2array(cgs)
 100
 101         var output io.WriteCloser
 102         if *outputFilename == "-" {
 103                 output = nopCloser{stdout}
 104         } else {
 105                 output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
 106                 if err != nil {
 107                         return 1
 108                 }
 109                 defer output.Close()
 110         }
 111         bufw := bufio.NewWriter(output)
 112         npw, err := gonpy.NewWriter(nopCloser{bufw})
 113         if err != nil {
 114                 return 1
 115         }
 116         if *onehot {
 117                 out, cols = recodeOnehot(out, cols)
 118         }
 119         npw.Shape = []int{rows, cols}
 120         npw.WriteUint16(out)
 121         err = bufw.Flush()
 122         if err != nil {
 123                 return 1
 124         }
 125         err = output.Close()
 126         if err != nil {
 127                 return 1
 128         }
 129         return 0
 130 }
 131
 132 func cgs2array(cgs []CompactGenome) (data []uint16, rows, cols int) {
 133         rows = len(cgs)
 134         for _, cg := range cgs {
 135                 if cols < len(cg.Variants) {
 136                         cols = len(cg.Variants)
 137                 }
 138         }
 139         data = make([]uint16, rows*cols)
 140         for row, cg := range cgs {
 141                 for i, v := range cg.Variants {
 142                         data[row*cols+i] = uint16(v)
 143                 }
 144         }
 145         return
 146 }
 147
 148 func recodeOnehot(in []uint16, incols int) ([]uint16, int) {
 149         rows := len(in) / incols
 150         maxvalue := make([]uint16, incols)
 151         for row := 0; row < rows; row++ {
 152                 for col := 0; col < incols; col++ {
 153                         if v := in[row*incols+col]; maxvalue[col] < v {
 154                                 maxvalue[col] = v
 155                         }
 156                 }
 157         }
 158         outcol := make([]int, incols)
 159         outcols := 0
 160         dropped := 0
 161         for incol, v := range maxvalue {
 162                 outcol[incol] = outcols
 163                 if v == 0 {
 164                         dropped++
 165                 } else {
 166                         outcols += int(v)
 167                 }
 168         }
 169         log.Printf("recodeOnehot: dropped %d input cols with zero maxvalue", dropped)
 170
 171         out := make([]uint16, rows*outcols)
 172         for inidx, row := 0, 0; row < rows; row++ {
 173                 outrow := out[row*outcols:]
 174                 for col := 0; col < incols; col++ {
 175                         if v := in[inidx]; v > 0 {
 176                                 outrow[outcol[col]+int(v)-1] = 1
 177                         }
 178                         inidx++
 179                 }
 180         }
 181         return out, outcols
 182 }
 183
 184 type nopCloser struct {
 185         io.Writer
 186 }
 187
 188 func (nopCloser) Close() error { return nil }