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         err = runner.Run()
  55         if err != nil {
  56                 return 1
  57         }
  58         return 0
  59 }
  60
  61 var plotscript = `
  62 import csv
  63 import os
  64 import scipy
  65 import sys
  66
  67 infile = sys.argv[1]
  68 X = scipy.load(infile)
  69
  70 colors = None
  71 if sys.argv[2]:
  72     labels = {}
  73     for fnm in os.listdir(sys.argv[3]):
  74         if '.2.fasta' not in fnm:
  75             labels[fnm] = '---'
  76     if len(labels) != len(X):
  77         raise "len(inputdir) != len(inputarray)"
  78     with open(sys.argv[2], 'rt') as csvfile:
  79         for row in csv.reader(csvfile):
  80             ident=row[0]
  81             label=row[1]
  82             for fnm in labels:
  83                 if row[0] in fnm:
  84                     labels[fnm] = row[1]
  85     colors = []
  86     labelcolors = {
  87         'PUR': 'firebrick',
  88         'CLM': 'firebrick',
  89         'MXL': 'firebrick',
  90         'PEL': 'firebrick',
  91         'TSI': 'green',
  92         'IBS': 'green',
  93         'CEU': 'green',
  94         'GBR': 'green',
  95         'FIN': 'green',
  96         'LWK': 'coral',
  97         'MSL': 'coral',
  98         'GWD': 'coral',
  99         'YRI': 'coral',
 100         'ESN': 'coral',
 101         'ACB': 'coral',
 102         'ASW': 'coral',
 103         'KHV': 'royalblue',
 104         'CDX': 'royalblue',
 105         'CHS': 'royalblue',
 106         'CHB': 'royalblue',
 107         'JPT': 'royalblue',
 108         'STU': 'blueviolet',
 109         'ITU': 'blueviolet',
 110         'BEB': 'blueviolet',
 111         'GIH': 'blueviolet',
 112         'PJL': 'blueviolet',
 113     }
 114     for fnm in sorted(labels.keys()):
 115         colors.append(labelcolors[labels[fnm]])
 116
 117 from matplotlib.figure import Figure
 118 from matplotlib.patches import Polygon
 119 from matplotlib.backends.backend_agg import FigureCanvasAgg
 120 fig = Figure()
 121 ax = fig.add_subplot(111)
 122 ax.scatter(X[:,0], X[:,1], c=colors, s=60, marker='o', alpha=0.5)
 123 canvas = FigureCanvasAgg(fig)
 124 canvas.print_figure(sys.argv[4], dpi=80)
 125 `