19566: Option to limit pca components used in glm. Fix onehot use.

[lightning.git] / plot.go

diff --git a/plot.go b/plot.go
index 51f5c8236c27e62a2134f706999ee8752a4ce731..6deaff6fc7d01057b0b4cfd1c49537cbacc7edb0 100644 (file)
--- a/plot.go
+++ b/plot.go
@@ -10,6 +10,8 @@ import (
        "fmt"
        "io"
        _ "net/http/pprof"
+       "os/exec"
+       "strings"
 
        "git.arvados.org/arvados.git/sdk/go/arvados"
 )
@@ -30,16 +32,24 @@ func (cmd *pythonPlot) RunCommand(prog string, args []string, stdin io.Reader, s
        flags.SetOutput(stderr)
        projectUUID := flags.String("project", "", "project `UUID` for output data")
        inputFilename := flags.String("i", "-", "input `file`")
+       outputFilename := flags.String("o", "", "output `filename` (e.g., './plot.png')")
        sampleListFilename := flags.String("samples", "", "use second column of `samples.csv` as complete list of sample IDs")
        phenotypeFilename := flags.String("phenotype", "", "use `phenotype.csv` as id->phenotype mapping (column 0 is sample id)")
-       phenotypeColumn := flags.Int("phenotype-column", 1, "0-based column `index` of phenotype in phenotype.csv file")
+       cat1Column := flags.Int("phenotype-cat1-column", 1, "0-based column `index` of 1st category in phenotype.csv file")
+       cat2Column := flags.Int("phenotype-cat2-column", -1, "0-based column `index` of 2nd category in phenotype.csv file")
+       xComponent := flags.Int("x", 1, "1-based PCA component to plot on x axis")
+       yComponent := flags.Int("y", 2, "1-based PCA component to plot on y axis")
        priority := flags.Int("priority", 500, "container request priority")
+       runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
        err = flags.Parse(args)
        if err == flag.ErrHelp {
                err = nil
                return 0
        } else if err != nil {
                return 2
+       } else if flags.NArg() > 0 {
+               err = fmt.Errorf("errant command line arguments after parsed flags: %v", flags.Args())
+               return 2
        }
 
        runner := arvadosContainerRunner{
@@ -56,12 +66,40 @@ func (cmd *pythonPlot) RunCommand(prog string, args []string, stdin io.Reader, s
                        },
                },
        }
-       err = runner.TranslatePaths(inputFilename, sampleListFilename, phenotypeFilename)
-       if err != nil {
-               return 1
+       if !*runlocal {
+               err = runner.TranslatePaths(inputFilename, sampleListFilename, phenotypeFilename)
+               if err != nil {
+                       return 1
+               }
+               *outputFilename = "/mnt/output/plot.png"
+       }
+       args = []string{
+               *inputFilename,
+               fmt.Sprintf("%d", *xComponent),
+               fmt.Sprintf("%d", *yComponent),
+               *sampleListFilename,
+               *phenotypeFilename,
+               fmt.Sprintf("%d", *cat1Column),
+               fmt.Sprintf("%d", *cat2Column),
+               *outputFilename,
+       }
+       if *runlocal {
+               if *outputFilename == "" {
+                       fmt.Fprintln(stderr, "error: must specify -o filename.png in local mode (or try -help)")
+                       return 1
+               }
+               cmd := exec.Command("python3", append([]string{"-"}, args...)...)
+               cmd.Stdin = strings.NewReader(plotscript)
+               cmd.Stdout = stdout
+               cmd.Stderr = stderr
+               err = cmd.Run()
+               if err != nil {
+                       return 1
+               }
+               return 0
        }
        runner.Prog = "python3"
-       runner.Args = []string{"/plot.py", *inputFilename, *sampleListFilename, *phenotypeFilename, fmt.Sprintf("%d", *phenotypeColumn), "/mnt/output/plot.png"}
+       runner.Args = append([]string{"/plot.py"}, args...)
        var output string
        output, err = runner.Run()
        if err != nil {