Merge branch '19524-pca'
[lightning.git] / plot.go
1 // Copyright (C) The Lightning Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package lightning
6
7 import (
8         _ "embed"
9         "flag"
10         "fmt"
11         "io"
12         _ "net/http/pprof"
13         "os/exec"
14         "strings"
15
16         "git.arvados.org/arvados.git/sdk/go/arvados"
17 )
18
19 type pythonPlot struct{}
20
21 //go:embed plot.py
22 var plotscript string
23
24 func (cmd *pythonPlot) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
25         var err error
26         defer func() {
27                 if err != nil {
28                         fmt.Fprintf(stderr, "%s\n", err)
29                 }
30         }()
31         flags := flag.NewFlagSet("", flag.ContinueOnError)
32         flags.SetOutput(stderr)
33         projectUUID := flags.String("project", "", "project `UUID` for output data")
34         inputFilename := flags.String("i", "-", "input `file`")
35         outputFilename := flags.String("o", "", "output `filename` (e.g., './plot.png')")
36         sampleListFilename := flags.String("samples", "", "use second column of `samples.csv` as complete list of sample IDs")
37         phenotypeFilename := flags.String("phenotype", "", "use `phenotype.csv` as id->phenotype mapping (column 0 is sample id)")
38         cat1Column := flags.Int("phenotype-cat1-column", 1, "0-based column `index` of 1st category in phenotype.csv file")
39         cat2Column := flags.Int("phenotype-cat2-column", -1, "0-based column `index` of 2nd category in phenotype.csv file")
40         xComponent := flags.Int("x", 1, "1-based PCA component to plot on x axis")
41         yComponent := flags.Int("y", 2, "1-based PCA component to plot on y axis")
42         priority := flags.Int("priority", 500, "container request priority")
43         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
44         err = flags.Parse(args)
45         if err == flag.ErrHelp {
46                 err = nil
47                 return 0
48         } else if err != nil {
49                 return 2
50         }
51
52         runner := arvadosContainerRunner{
53                 Name:        "lightning plot",
54                 Client:      arvados.NewClientFromEnv(),
55                 ProjectUUID: *projectUUID,
56                 RAM:         4 << 30,
57                 VCPUs:       1,
58                 Priority:    *priority,
59                 Mounts: map[string]map[string]interface{}{
60                         "/plot.py": map[string]interface{}{
61                                 "kind":    "text",
62                                 "content": plotscript,
63                         },
64                 },
65         }
66         if !*runlocal {
67                 err = runner.TranslatePaths(inputFilename, sampleListFilename, phenotypeFilename)
68                 if err != nil {
69                         return 1
70                 }
71                 *outputFilename = "/mnt/output/plot.png"
72         }
73         args = []string{
74                 *inputFilename,
75                 fmt.Sprintf("%d", *xComponent),
76                 fmt.Sprintf("%d", *yComponent),
77                 *sampleListFilename,
78                 *phenotypeFilename,
79                 fmt.Sprintf("%d", *cat1Column),
80                 fmt.Sprintf("%d", *cat2Column),
81                 *outputFilename,
82         }
83         if *runlocal {
84                 if *outputFilename == "" {
85                         fmt.Fprintln(stderr, "error: must specify -o filename.png in local mode (or try -help)")
86                         return 1
87                 }
88                 cmd := exec.Command("python3", append([]string{"-"}, args...)...)
89                 cmd.Stdin = strings.NewReader(plotscript)
90                 cmd.Stdout = stdout
91                 cmd.Stderr = stderr
92                 err = cmd.Run()
93                 if err != nil {
94                         return 1
95                 }
96                 return 0
97         }
98         runner.Prog = "python3"
99         runner.Args = append([]string{"/plot.py"}, args...)
100         var output string
101         output, err = runner.Run()
102         if err != nil {
103                 return 1
104         }
105         fmt.Fprintln(stdout, output+"/plot.png")
106         return 0
107 }