plot.go

   1 package main
   2
   3 import (
   4         "flag"
   5         "fmt"
   6         "io"
   7         _ "net/http/pprof"
   8
   9         "git.arvados.org/arvados.git/sdk/go/arvados"
  10 )
  11
  12 type pythonPlot struct{}
  13
  14 func (cmd *pythonPlot) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  15         var err error
  16         defer func() {
  17                 if err != nil {
  18                         fmt.Fprintf(stderr, "%s\n", err)
  19                 }
  20         }()
  21         flags := flag.NewFlagSet("", flag.ContinueOnError)
  22         flags.SetOutput(stderr)
  23         projectUUID := flags.String("project", "", "project `UUID` for output data")
  24         inputFilename := flags.String("i", "-", "input `file`")
  25         sampleCSVFilename := flags.String("labels-csv", "", "use first two columns of `labels.csv` as id->color mapping")
  26         sampleFastaDirname := flags.String("sample-fasta-dir", "", "`directory` containing fasta input files")
  27         err = flags.Parse(args)
  28         if err == flag.ErrHelp {
  29                 err = nil
  30                 return 0
  31         } else if err != nil {
  32                 return 2
  33         }
  34
  35         runner := arvadosContainerRunner{
  36                 Name:        "lightning plot",
  37                 Client:      arvados.NewClientFromEnv(),
  38                 ProjectUUID: *projectUUID,
  39                 RAM:         1 << 30,
  40                 VCPUs:       1,
  41                 Mounts: map[string]map[string]interface{}{
  42                         "/plot.py": map[string]interface{}{
  43                                 "kind":    "text",
  44                                 "content": plotscript,
  45                         },
  46                 },
  47         }
  48         err = runner.TranslatePaths(inputFilename, sampleCSVFilename, sampleFastaDirname)
  49         if err != nil {
  50                 return 1
  51         }
  52         runner.Prog = "python3"
  53         runner.Args = []string{"/plot.py", *inputFilename, *sampleCSVFilename, *sampleFastaDirname, "/mnt/output/plot.png"}
  54         var output string
  55         output, err = runner.Run()
  56         if err != nil {
  57                 return 1
  58         }
  59         fmt.Fprintln(stdout, output+"/plot.png")
  60         return 0
  61 }
  62
  63 var plotscript = `
  64 import csv
  65 import os
  66 import scipy
  67 import sys
  68
  69 infile = sys.argv[1]
  70 X = scipy.load(infile)
  71
  72 colors = None
  73 if sys.argv[2]:
  74     labels = {}
  75     for fnm in os.listdir(sys.argv[3]):
  76         if '.2.fasta' not in fnm:
  77             labels[fnm] = '---'
  78     if len(labels) != len(X):
  79         raise "len(inputdir) != len(inputarray)"
  80     with open(sys.argv[2], 'rt') as csvfile:
  81         for row in csv.reader(csvfile):
  82             ident=row[0]
  83             label=row[1]
  84             for fnm in labels:
  85                 if row[0] in fnm:
  86                     labels[fnm] = row[1]
  87     colors = []
  88     labelcolors = {
  89         'PUR': 'firebrick',
  90         'CLM': 'firebrick',
  91         'MXL': 'firebrick',
  92         'PEL': 'firebrick',
  93         'TSI': 'green',
  94         'IBS': 'green',
  95         'CEU': 'green',
  96         'GBR': 'green',
  97         'FIN': 'green',
  98         'LWK': 'coral',
  99         'MSL': 'coral',
 100         'GWD': 'coral',
 101         'YRI': 'coral',
 102         'ESN': 'coral',
 103         'ACB': 'coral',
 104         'ASW': 'coral',
 105         'KHV': 'royalblue',
 106         'CDX': 'royalblue',
 107         'CHS': 'royalblue',
 108         'CHB': 'royalblue',
 109         'JPT': 'royalblue',
 110         'STU': 'blueviolet',
 111         'ITU': 'blueviolet',
 112         'BEB': 'blueviolet',
 113         'GIH': 'blueviolet',
 114         'PJL': 'blueviolet',
 115     }
 116     for fnm in sorted(labels.keys()):
 117         colors.append(labelcolors[labels[fnm]])
 118
 119 from matplotlib.figure import Figure
 120 from matplotlib.patches import Polygon
 121 from matplotlib.backends.backend_agg import FigureCanvasAgg
 122 fig = Figure()
 123 ax = fig.add_subplot(111)
 124 ax.scatter(X[:,0], X[:,1], c=colors, s=60, marker='o', alpha=0.5)
 125 canvas = FigureCanvasAgg(fig)
 126 canvas.print_figure(sys.argv[4], dpi=80)
 127 `