exportnumpy.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "errors"
   6         "flag"
   7         "fmt"
   8         "io"
   9         "log"
  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 )
  17
  18 type exportNumpy struct {
  19         output io.Writer
  20 }
  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         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         cmd.output = stdout
  44
  45         if *pprof != "" {
  46                 go func() {
  47                         log.Println(http.ListenAndServe(*pprof, nil))
  48                 }()
  49         }
  50
  51         if !*runlocal {
  52                 if *outputFilename != "" {
  53                         err = errors.New("cannot specify output file in container mode: not implemented")
  54                         return 1
  55                 }
  56                 runner := arvadosContainerRunner{
  57                         Name:        "lightning export-numpy",
  58                         Client:      arvados.NewClientFromEnv(),
  59                         ProjectUUID: *projectUUID,
  60                         RAM:         64000000000,
  61                         VCPUs:       2,
  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                 err = runner.Run()
  69                 if err != nil {
  70                         return 1
  71                 }
  72                 return 0
  73         }
  74
  75         cgs, err := ReadCompactGenomes(stdin)
  76         if err != nil {
  77                 return 1
  78         }
  79         cols := 0
  80         for _, cg := range cgs {
  81                 if cols < len(cg.Variants) {
  82                         cols = len(cg.Variants)
  83                 }
  84         }
  85         rows := len(cgs)
  86         out := make([]uint16, rows*cols)
  87         for row, cg := range cgs {
  88                 for i, v := range cg.Variants {
  89                         out[row*cols+i] = uint16(v)
  90                 }
  91         }
  92
  93         var output io.WriteCloser
  94         if *outputFilename == "" {
  95                 output = nopCloser{cmd.output}
  96         } else {
  97                 output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
  98                 if err != nil {
  99                         return 1
 100                 }
 101                 defer output.Close()
 102         }
 103         bufw := bufio.NewWriter(output)
 104         npw, err := gonpy.NewWriter(nopCloser{bufw})
 105         if err != nil {
 106                 return 1
 107         }
 108         npw.Shape = []int{rows, cols}
 109         npw.WriteUint16(out)
 110         err = bufw.Flush()
 111         if err != nil {
 112                 return 1
 113         }
 114         err = output.Close()
 115         if err != nil {
 116                 return 1
 117         }
 118         return 0
 119 }
 120
 121 type nopCloser struct {
 122         io.Writer
 123 }
 124
 125 func (nopCloser) Close() error { return nil }