PCSPSF/plotter.py at master · oxford-pcs/PCSPSF · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
from copy import deepcopy

import pylab as plt
import numpy as np

class plotter():
  def __init__(self):
    self.PLOTDATA = []		# plot information

  def addImagePlot(self, title, data, cb=True, extent=None, xl=None, yl=None,
    overplot=False):
    '''
      Add image data to plot.
    '''
    self.PLOTDATA.append({"title": title, "data": deepcopy(data), "cb": cb,
      "type": "im", "extent": extent, "xl": xl, "yl": yl, "overplot": overplot})

  def addScatterPlot(self, title, y, x=None, color='w', ls='--', cb=True,
    xl=None, yl=None, xr=None, yr=None, overplot=False):
    '''
      Add xy scatter data to plot.
    '''
    self.PLOTDATA.append({"title": title, "x": deepcopy(x), "y": deepcopy(y),
      "color": color, "ls": ls, "cb": cb, "type": "scatter", "xl": xl, "yl": yl,
      "xr": xr, "yr": yr, "overplot": overplot})

  def addTextToPlot(self, x, y, s, color='w', fontsize=12):
    '''
      Add text data to plot.
    '''
    self.PLOTDATA.append({"x": x, "y": y, "text": s, "color": color,
      "fontsize": fontsize, "type": "text", "overplot": True})

  def _reset(self):
    self.__init__()

  def draw(self, nrows, ncols):
    '''
      Sequentially draws data from [self.PLOTDATA].
    '''
    fig = plt.figure(figsize=(14, 10))
    nplot = 1
    for d in self.PLOTDATA:
      if not d['overplot']:
        plt.subplot(ncols,nrows,nplot)
        plt.title(d['title'], fontsize=12)
        nplot = nplot+1
      if d['type'] == 'im':
        plt.imshow(d['data'], extent=d['extent'], interpolation="none")
      if d['xl'] is not None:
        plt.xlabel(d['xl'], fontsize=12)
      if d['yl'] is not None:
        plt.ylabel(d['yl'], fontsize=12)
      if d['cb']:
        plt.colorbar()
      if d['type'] == 'scatter':
        if d['x'] == None:
          plt.plot(d['y'], color=d['color'], linestyle=d['ls'])
        else:
          plt.plot(d['x'], d['y'], color=d['color'], linestyle=d['ls'])
        if d['xl'] is not None:
          plt.xlabel(d['xl'], fontsize=12)
        if d['yl'] is not None:
          plt.ylabel(d['yl'], fontsize=12)
        if d['xr'] is not None:
          plt.xlim(d['xr'])
        if d['yr'] is not None:
          plt.ylim(d['yr'])
      if d['type'] == 'text':
        plt.text(d['x'], d['y'], d['text'], color=d['color'],
          fontsize=d['fontsize'])

    plt.tight_layout()
    plt.show()
    self._reset()