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         inputFilename := flags.String("i", "-", "input `file`")
  34         outputFilename := flags.String("o", "-", "output `file`")
  35         err = flags.Parse(args)
  36         if err == flag.ErrHelp {
  37                 err = nil
  38                 return 0
  39         } else if err != nil {
  40                 return 2
  41         }
  42
  43         if *pprof != "" {
  44                 go func() {
  45                         log.Println(http.ListenAndServe(*pprof, nil))
  46                 }()
  47         }
  48
  49         if !*runlocal {
  50                 if *outputFilename != "-" {
  51                         err = errors.New("cannot specify output file in container mode: not implemented")
  52                         return 1
  53                 }
  54                 runner := arvadosContainerRunner{
  55                         Name:        "lightning export-numpy",
  56                         Client:      arvados.NewClientFromEnv(),
  57                         ProjectUUID: *projectUUID,
  58                         RAM:         64000000000,
  59                         VCPUs:       2,
  60                 }
  61                 err = runner.TranslatePaths(inputFilename)
  62                 if err != nil {
  63                         return 1
  64                 }
  65                 runner.Args = []string{"export-numpy", "-local=true", "-i", *inputFilename, "-o", "/mnt/output/library.npy"}
  66                 var output string
  67                 output, err = runner.Run()
  68                 if err != nil {
  69                         return 1
  70                 }
  71                 fmt.Fprintln(stdout, output+"/library.npy")
  72                 return 0
  73         }
  74
  75         var input io.ReadCloser
  76         if *inputFilename == "-" {
  77                 input = ioutil.NopCloser(stdin)
  78         } else {
  79                 input, err = os.Open(*inputFilename)
  80                 if err != nil {
  81                         return 1
  82                 }
  83                 defer input.Close()
  84         }
  85         cgs, err := ReadCompactGenomes(input)
  86         if err != nil {
  87                 return 1
  88         }
  89         err = input.Close()
  90         if err != nil {
  91                 return 1
  92         }
  93         cols := 0
  94         for _, cg := range cgs {
  95                 if cols < len(cg.Variants) {
  96                         cols = len(cg.Variants)
  97                 }
  98         }
  99         rows := len(cgs)
 100         out := make([]uint16, rows*cols)
 101         for row, cg := range cgs {
 102                 for i, v := range cg.Variants {
 103                         out[row*cols+i] = uint16(v)
 104                 }
 105         }
 106
 107         var output io.WriteCloser
 108         if *outputFilename == "-" {
 109                 output = nopCloser{stdout}
 110         } else {
 111                 output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
 112                 if err != nil {
 113                         return 1
 114                 }
 115                 defer output.Close()
 116         }
 117         bufw := bufio.NewWriter(output)
 118         npw, err := gonpy.NewWriter(nopCloser{bufw})
 119         if err != nil {
 120                 return 1
 121         }
 122         npw.Shape = []int{rows, cols}
 123         npw.WriteUint16(out)
 124         err = bufw.Flush()
 125         if err != nil {
 126                 return 1
 127         }
 128         err = output.Close()
 129         if err != nil {
 130                 return 1
 131         }
 132         return 0
 133 }
 134
 135 type nopCloser struct {
 136         io.Writer
 137 }
 138
 139 func (nopCloser) Close() error { return nil }