-
Notifications
You must be signed in to change notification settings - Fork 0
/
colorgrid.py
36 lines (25 loc) · 1.24 KB
/
colorgrid.py
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
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
class Colorgrid:
def __init__(self, X, figure = None, colormap = 'jet', ylim = (-1,1), xlim = (-1,1), norm = None, origin='lower', interpolation='nearest', color_label='colorscale'):
# "figure" should be a matplotlib figure
# initialize figure and axes
if figure is None:
self.fig = plt.figure()
else:
self.fig = plt.figure(figure)
self.ax0 = self.fig.add_axes([0.1,0.1, 0.75, 0.75])
self.ax1 = self.fig.add_axes([0.85,0.1,0.05,0.75])
self.cmap = plt.get_cmap(colormap)
# make the heatmap
im = self.ax0.imshow( X,
cmap=self.cmap,
extent=(xlim[0], xlim[1], ylim[0], ylim[1]),
origin=origin,
interpolation=interpolation)
#im.set_interpolation(interpolation)
# make the colorbar
norm = plt.Normalize(X.min(), X.max())
cb = matplotlib.colorbar.ColorbarBase(self.ax1, cmap=self.cmap, norm=norm, orientation='vertical', boundaries=None)
self.ax1.set_ylabel(color_label)