exportnumpy.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "flag"
   6         "fmt"
   7         "io"
   8         "log"
   9         "net/http"
  10         _ "net/http/pprof"
  11
  12         "github.com/kshedden/gonpy"
  13 )
  14
  15 type exportNumpy struct {
  16         output io.Writer
  17 }
  18
  19 func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  20         var err error
  21         defer func() {
  22                 if err != nil {
  23                         fmt.Fprintf(stderr, "%s\n", err)
  24                 }
  25         }()
  26         flags := flag.NewFlagSet("", flag.ContinueOnError)
  27         flags.SetOutput(stderr)
  28         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  29         err = flags.Parse(args)
  30         if err == flag.ErrHelp {
  31                 err = nil
  32                 return 0
  33         } else if err != nil {
  34                 return 2
  35         }
  36         cmd.output = stdout
  37
  38         if *pprof != "" {
  39                 go func() {
  40                         log.Println(http.ListenAndServe(*pprof, nil))
  41                 }()
  42         }
  43
  44         cgs, err := ReadCompactGenomes(stdin)
  45         if err != nil {
  46                 return 1
  47         }
  48         cols := 0
  49         for _, cg := range cgs {
  50                 if cols < len(cg.Variants) {
  51                         cols = len(cg.Variants)
  52                 }
  53         }
  54         rows := len(cgs)
  55         out := make([]uint16, rows*cols)
  56         for row, cg := range cgs {
  57                 for i, v := range cg.Variants {
  58                         out[row*cols+i] = uint16(v)
  59                 }
  60         }
  61         w := bufio.NewWriter(cmd.output)
  62         npw, err := gonpy.NewWriter(nopCloser{w})
  63         if err != nil {
  64                 return 1
  65         }
  66         npw.Shape = []int{rows, cols}
  67         npw.WriteUint16(out)
  68         err = w.Flush()
  69         if err != nil {
  70                 return 1
  71         }
  72         return 0
  73 }
  74
  75 type nopCloser struct {
  76         io.Writer
  77 }
  78
  79 func (nopCloser) Close() error { return nil }