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