exportnumpy.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "errors"
   6         "flag"
   7         "fmt"
   8         "io"
   9         "io/ioutil"
  10         "log"
  11         "net/http"
  12         _ "net/http/pprof"
  13         "os"
  14
  15         "git.arvados.org/arvados.git/sdk/go/arvados"
  16         "github.com/kshedden/gonpy"
  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                 err = runner.Run()
  67                 if err != nil {
  68                         return 1
  69                 }
  70                 return 0
  71         }
  72
  73         var input io.ReadCloser
  74         if *inputFilename == "-" {
  75                 input = ioutil.NopCloser(stdin)
  76         } else {
  77                 input, err = os.Open(*inputFilename)
  78                 if err != nil {
  79                         return 1
  80                 }
  81                 defer input.Close()
  82         }
  83         cgs, err := ReadCompactGenomes(input)
  84         if err != nil {
  85                 return 1
  86         }
  87         err = input.Close()
  88         if err != nil {
  89                 return 1
  90         }
  91         cols := 0
  92         for _, cg := range cgs {
  93                 if cols < len(cg.Variants) {
  94                         cols = len(cg.Variants)
  95                 }
  96         }
  97         rows := len(cgs)
  98         out := make([]uint16, rows*cols)
  99         for row, cg := range cgs {
 100                 for i, v := range cg.Variants {
 101                         out[row*cols+i] = uint16(v)
 102                 }
 103         }
 104
 105         var output io.WriteCloser
 106         if *outputFilename == "-" {
 107                 output = nopCloser{stdout}
 108         } else {
 109                 output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
 110                 if err != nil {
 111                         return 1
 112                 }
 113                 defer output.Close()
 114         }
 115         bufw := bufio.NewWriter(output)
 116         npw, err := gonpy.NewWriter(nopCloser{bufw})
 117         if err != nil {
 118                 return 1
 119         }
 120         npw.Shape = []int{rows, cols}
 121         npw.WriteUint16(out)
 122         err = bufw.Flush()
 123         if err != nil {
 124                 return 1
 125         }
 126         err = output.Close()
 127         if err != nil {
 128                 return 1
 129         }
 130         return 0
 131 }
 132
 133 type nopCloser struct {
 134         io.Writer
 135 }
 136
 137 func (nopCloser) Close() error { return nil }