FEM_python/visualization.py at master · DimaAndreev30/FEM_python · 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
77
78
79
80
81
82
"""
Пример использования классов Normalize и LinearSegmentedColormap
из модуля colors библиотеки matplotlib:

from matplotlib.colors import Normalize, LinearSegmentedColormap

normalizer = Normalize()
colormap = LinearSegmentedColormap(
    'map',
    {'red':   [(0.0,  0.0, 0.0),
               (0.5,  1.0, 1.0),
               (1.0,  1.0, 1.0)],

     'green': [(0.0,  0.0, 0.0),
               (0.25, 0.0, 0.0),
               (0.75, 1.0, 1.0),
               (1.0,  1.0, 1.0)],

     'blue':  [(0.0,  0.0, 0.0),
               (0.5,  0.0, 0.0),
               (1.0,  1.0, 1.0)]})

colormap(normalizer(data).data) # -> выдает цвета в формате RGBA
"""

"""
Пример использования модуля cm библиотеки matplotlib:

from matplotlib import cm

cm.cmap_d.keys() для вывода названий предустановленных карт
'viridis' - для приятной картинки
'coolwarm' - для симметрии отрицательных/положительных значений
'gray' - для черно-белой картинки

cmap = cm.get_cmap('viridis')
mappable = cm.ScalarMappable(cmap=cmap)
mappable.to_rgba(data)
plt.colorbar(mappable)
"""


import matplotlib.pyplot as plt
import matplotlib.cm as cm


def visualizeCellsValues(mesh, values,
                         contour=False, contour_color=[[0.0, 0.0, 0.0, 0.4]],
                         mappable=None, cmap=cm.get_cmap('viridis'),
                         bar_orientation='vertical', bar_shrink=1.0,
                         **kwargs):
    """
    Визуализация значений values на ячейках сетки mesh
    """

    if mappable is None:
        mappable = cm.ScalarMappable(cmap=cmap)

    plt.gca().set_aspect('equal')
    mesh.visualize(color=mappable.to_rgba(values), fill=True, **kwargs)
    if contour is True:
        mesh.visualize(color=contour_color, fill=False, **kwargs)
    plt.colorbar(mappable, orientation=bar_orientation, shrink=bar_shrink)


def visualizeNodesValues(mesh, values, levels=50,
                         contour=False, contour_color=[[0.0, 0.0, 0.0, 0.4]],
                         mappable=None, cmap=cm.get_cmap('viridis'),
                         bar_orientation='vertical', bar_shrink=1.0,
                         **kwargs):
    """
    Визуализация значений values на ячейках сетки mesh
    """

    if mappable is None:
        mappable = cm.ScalarMappable(cmap=cmap)

    plt.gca().set_aspect('equal')
    mesh.visualizeNodes(values=mappable.norm(values), cmap=mappable.cmap, levels=levels, **kwargs)
    if contour is True:
        mesh.visualize(color=contour_color, fill=False, **kwargs)
    plt.colorbar(mappable, orientation=bar_orientation, shrink=bar_shrink)