heatmap.py 1.42 KB
Newer Older
1
# -*- coding: utf-8 -*-
2
import numpy as np
3

4
from nifty.plotting.colormap import Colormap
5
from nifty.plotting.plotly_wrapper import PlotlyWrapper
6

7

8
9
class Heatmap(PlotlyWrapper):
    def __init__(self, data, color_map=None, webgl=False,
10
                 smoothing=False):  # smoothing 'best', 'fast', False
Mihai Baltac's avatar
Mihai Baltac committed
11
12
13
14
15
16
        if isinstance(data, list):
            self.data = np.zeros((data[0].shape))
            for arr in data:
                self.data = np.add(self.data, arr)
        else:
            self.data = data
17
18
19
20
21

        if color_map is not None:
            if not isinstance(color_map, Colormap):
                raise TypeError("Provided color_map must be an instance of "
                                "the NIFTy Colormap class.")
22
        self.color_map = color_map
23
24
25
        self.webgl = webgl
        self.smoothing = smoothing

26
27
28
29
    @property
    def figure_dimension(self):
        return 2

30
    def to_plotly(self):
31
        plotly_object = dict()
32
        plotly_object['z'] = self.data
33
34
35
36
        plotly_object['showscale'] = False
        if self.color_map:
            plotly_object['colorscale'] = self.color_map.to_plotly()
            plotly_object['colorbar'] = dict(title=self.color_map.name, x=0.42)
37
38
39
40
41
42
43
        if self.webgl:
            plotly_object['type'] = 'heatmapgl'
        else:
            plotly_object['type'] = 'heatmap'
        if self.smoothing:
            plotly_object['zsmooth'] = self.smoothing
        return plotly_object