multi_figure.py 2.56 KB
Newer Older
1
2
3
# -*- coding: utf-8 -*-

import numpy as np
4
from nifty import dependency_injector as gdi
5
6
7
from figure_base import FigureBase
from figure_3D import Figure3D

8
9
plotly = gdi.get('plotly')

10

Theo Steininger's avatar
Theo Steininger committed
11
# TODO: add nice height and width defaults for multifigure
12
class MultiFigure(FigureBase):
Theo Steininger's avatar
Theo Steininger committed
13
14
    def __init__(self, rows, columns, title=None, width=None, height=None,
                 subfigures=None):
15
16
        if 'plotly' not in gdi:
            raise ImportError("The module plotly is needed but not available.")
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
        super(MultiFigure, self).__init__(title, width, height)
        self.subfigures = np.empty((rows, columns), dtype=np.object)
        self.subfigures[:] = subfigures

    @property
    def rows(self):
        return self.subfigures.shape[0]

    @property
    def columns(self):
        return self.subfigures.shape[1]

    def add_subfigure(self, figure, row, column):
        self.subfigures[row, column] = figure

    def to_plotly(self):
Theo Steininger's avatar
Theo Steininger committed
33
34
35
        title_extractor = lambda z: z.title if z else ""
        sub_titles = \
            tuple(np.vectorize(title_extractor)(self.subfigures.flatten()))
36
37
38

        specs_setter = \
            lambda z: {'is_3d': True} if isinstance(z, Figure3D) else {}
Theo Steininger's avatar
Theo Steininger committed
39
40
41
        sub_specs = \
            list(map(list, np.vectorize(specs_setter)(self.subfigures)))

42
43
        multi_figure_plotly_object = plotly.tools.make_subplots(
                                                   self.rows,
44
45
46
                                                   self.columns,
                                                   subplot_titles=sub_titles,
                                                   specs=sub_specs)
Theo Steininger's avatar
Theo Steininger committed
47
        # TODO resolve bug with titles and 3D subplots
48
        for index, fig in np.ndenumerate(self.subfigures):
Theo Steininger's avatar
Theo Steininger committed
49
50
51
52
53
            if fig:
                for plot in fig.plots:
                    multi_figure_plotly_object.append_trace(plot.to_plotly(),
                                                            index[0]+1,
                                                            index[1]+1)
54
55
56
57

        multi_figure_plotly_object['layout'].update(height=self.height,
                                                    width=self.width,
                                                    title=self.title)
Theo Steininger's avatar
Theo Steininger committed
58

59
60
61
62
        return multi_figure_plotly_object


    @staticmethod
Theo Steininger's avatar
Theo Steininger committed
63
64
65
66
    def from_figures_cols(figures, cols=2, title=None, width=None, height=None):
        multi_figure = MultiFigure((len(figures) + cols -1)/cols , cols, title, width, height)
        for i in range(len(figures)):
            multi_figure.add_subfigure(figures[i], i/cols, i%cols)
67
68
69
70

        return multi_figure