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
  14         "git.arvados.org/arvados.git/sdk/go/arvados"
  15         "github.com/kshedden/gonpy"
  16         log "github.com/sirupsen/logrus"
  17 )
  18
  19 type exportNumpy struct{}
  20
  21 func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  22         var err error
  23         defer func() {
  24                 if err != nil {
  25                         fmt.Fprintf(stderr, "%s\n", err)
  26                 }
  27         }()
  28         flags := flag.NewFlagSet("", flag.ContinueOnError)
  29         flags.SetOutput(stderr)
  30         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  31         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  32         projectUUID := flags.String("project", "", "project `UUID` for output data")
  33         priority := flags.Int("priority", 500, "container request priority")
  34         inputFilename := flags.String("i", "-", "input `file`")
  35         outputFilename := flags.String("o", "-", "output `file`")
  36         err = flags.Parse(args)
  37         if err == flag.ErrHelp {
  38                 err = nil
  39                 return 0
  40         } else if err != nil {
  41                 return 2
  42         }
  43
  44         if *pprof != "" {
  45                 go func() {
  46                         log.Println(http.ListenAndServe(*pprof, nil))
  47                 }()
  48         }
  49
  50         if !*runlocal {
  51                 if *outputFilename != "-" {
  52                         err = errors.New("cannot specify output file in container mode: not implemented")
  53                         return 1
  54                 }
  55                 runner := arvadosContainerRunner{
  56                         Name:        "lightning export-numpy",
  57                         Client:      arvados.NewClientFromEnv(),
  58                         ProjectUUID: *projectUUID,
  59                         RAM:         64000000000,
  60                         VCPUs:       2,
  61                         Priority:    *priority,
  62                 }
  63                 err = runner.TranslatePaths(inputFilename)
  64                 if err != nil {
  65                         return 1
  66                 }
  67                 runner.Args = []string{"export-numpy", "-local=true", "-i", *inputFilename, "-o", "/mnt/output/library.npy"}
  68                 var output string
  69                 output, err = runner.Run()
  70                 if err != nil {
  71                         return 1
  72                 }
  73                 fmt.Fprintln(stdout, output+"/library.npy")
  74                 return 0
  75         }
  76
  77         var input io.ReadCloser
  78         if *inputFilename == "-" {
  79                 input = ioutil.NopCloser(stdin)
  80         } else {
  81                 input, err = os.Open(*inputFilename)
  82                 if err != nil {
  83                         return 1
  84                 }
  85                 defer input.Close()
  86         }
  87         cgs, err := ReadCompactGenomes(input)
  88         if err != nil {
  89                 return 1
  90         }
  91         err = input.Close()
  92         if err != nil {
  93                 return 1
  94         }
  95         cols := 0
  96         for _, cg := range cgs {
  97                 if cols < len(cg.Variants) {
  98                         cols = len(cg.Variants)
  99                 }
 100         }
 101         rows := len(cgs)
 102         out := make([]uint16, rows*cols)
 103         for row, cg := range cgs {
 104                 for i, v := range cg.Variants {
 105                         out[row*cols+i] = uint16(v)
 106                 }
 107         }
 108
 109         var output io.WriteCloser
 110         if *outputFilename == "-" {
 111                 output = nopCloser{stdout}
 112         } else {
 113                 output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
 114                 if err != nil {
 115                         return 1
 116                 }
 117                 defer output.Close()
 118         }
 119         bufw := bufio.NewWriter(output)
 120         npw, err := gonpy.NewWriter(nopCloser{bufw})
 121         if err != nil {
 122                 return 1
 123         }
 124         npw.Shape = []int{rows, cols}
 125         npw.WriteUint16(out)
 126         err = bufw.Flush()
 127         if err != nil {
 128                 return 1
 129         }
 130         err = output.Close()
 131         if err != nil {
 132                 return 1
 133         }
 134         return 0
 135 }
 136
 137 type nopCloser struct {
 138         io.Writer
 139 }
 140
 141 func (nopCloser) Close() error { return nil }