Log dimensions.

[lightning.git] / exportnumpy.go

diff --git a/exportnumpy.go b/exportnumpy.go
index 8ccf11e707b3f17b930323a877f9f20e106836b8..fd198116b1512c01f6fb51db67bc82354edd40d5 100644 (file)
--- a/exportnumpy.go
+++ b/exportnumpy.go
@@ -7,13 +7,14 @@ import (
        "fmt"
        "io"
        "io/ioutil"
-       "log"
        "net/http"
        _ "net/http/pprof"
        "os"
+       "sort"
 
        "git.arvados.org/arvados.git/sdk/go/arvados"
        "github.com/kshedden/gonpy"
+       log "github.com/sirupsen/logrus"
 )
 
 type exportNumpy struct{}
@@ -30,8 +31,10 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
        pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
        runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
        projectUUID := flags.String("project", "", "project `UUID` for output data")
+       priority := flags.Int("priority", 500, "container request priority")
        inputFilename := flags.String("i", "-", "input `file`")
        outputFilename := flags.String("o", "-", "output `file`")
+       onehot := flags.Bool("one-hot", false, "recode tile variants as one-hot")
        err = flags.Parse(args)
        if err == flag.ErrHelp {
                err = nil
@@ -55,14 +58,15 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
                        Name:        "lightning export-numpy",
                        Client:      arvados.NewClientFromEnv(),
                        ProjectUUID: *projectUUID,
-                       RAM:         64000000000,
+                       RAM:         128000000000,
                        VCPUs:       2,
+                       Priority:    *priority,
                }
                err = runner.TranslatePaths(inputFilename)
                if err != nil {
                        return 1
                }
-               runner.Args = []string{"export-numpy", "-local=true", "-i", *inputFilename, "-o", "/mnt/output/library.npy"}
+               runner.Args = []string{"export-numpy", "-local=true", fmt.Sprintf("-one-hot=%v", *onehot), "-i", *inputFilename, "-o", "/mnt/output/library.npy"}
                var output string
                output, err = runner.Run()
                if err != nil {
@@ -90,19 +94,9 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
        if err != nil {
                return 1
        }
-       cols := 0
-       for _, cg := range cgs {
-               if cols < len(cg.Variants) {
-                       cols = len(cg.Variants)
-               }
-       }
-       rows := len(cgs)
-       out := make([]uint16, rows*cols)
-       for row, cg := range cgs {
-               for i, v := range cg.Variants {
-                       out[row*cols+i] = uint16(v)
-               }
-       }
+       sort.Slice(cgs, func(i, j int) bool { return cgs[i].Name < cgs[j].Name })
+
+       out, rows, cols := cgs2array(cgs)
 
        var output io.WriteCloser
        if *outputFilename == "-" {
@@ -119,6 +113,9 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
        if err != nil {
                return 1
        }
+       if *onehot {
+               out, cols = recodeOnehot(out, cols)
+       }
        npw.Shape = []int{rows, cols}
        npw.WriteUint16(out)
        err = bufw.Flush()
@@ -132,6 +129,49 @@ func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader,
        return 0
 }
 
+func cgs2array(cgs []CompactGenome) (data []uint16, rows, cols int) {
+       rows = len(cgs)
+       for _, cg := range cgs {
+               if cols < len(cg.Variants) {
+                       cols = len(cg.Variants)
+               }
+       }
+       data = make([]uint16, rows*cols)
+       for row, cg := range cgs {
+               for i, v := range cg.Variants {
+                       data[row*cols+i] = uint16(v)
+               }
+       }
+       return
+}
+
+func recodeOnehot(in []uint16, incols int) ([]uint16, int) {
+       rows := len(in) / incols
+       maxvalue := make([]uint16, incols)
+       for row := 0; row < rows; row++ {
+               for col := 0; col < incols; col++ {
+                       if v := in[row*incols+col]; maxvalue[col] < v {
+                               maxvalue[col] = v
+                       }
+               }
+       }
+       outcol := make([]int, incols)
+       outcols := 0
+       for incol, v := range maxvalue {
+               outcol[incol] = outcols
+               outcols += int(v)
+       }
+       out := make([]uint16, rows*outcols)
+       for row := 0; row < rows; row++ {
+               for col := 0; col < incols; col++ {
+                       if v := in[row*incols+col]; v > 0 {
+                               out[row*outcols+outcol[col]+int(v)-1] = 1
+                       }
+               }
+       }
+       return out, outcols
+}
+
 type nopCloser struct {
        io.Writer
 }