data-visualizer/api.py at master · AtLomDan/data-visualizer · 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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
from typing import Literal

import numpy as np

from normalizers import (
    MagnitudeNormalizer,
    RationalNormalizer,
    LinearNormalizer,
    LogNormalizer,
)

from strategies import (
    ColorCombineStrategy,
    ValueStrategy,
    SaturationStrategy,
    BothStrategy,
    PhaseOnlyStrategy,
    MagnitudeGrayStrategy,
)

from matrices import (
    RandomMatrixBuilder,
    DiagonalMatrixBuilder,
    UnitaryMatrixBuilder,
)

from mapper import ComplexColorMapper
from plotter import ComplexArrayPlotter

__version__ = "0.2.0"

__all__ = [
    "__version__",
    "MagnitudeNormalizer",
    "RationalNormalizer",
    "LinearNormalizer",
    "LogNormalizer",
    "ColorCombineStrategy",
    "ValueStrategy",
    "SaturationStrategy",
    "BothStrategy",
    "PhaseOnlyStrategy",
    "MagnitudeGrayStrategy",
    "RandomMatrixBuilder",
    "DiagonalMatrixBuilder",
    "UnitaryMatrixBuilder",
    "ComplexColorMapper",
    "ComplexArrayPlotter",
    "plot_array",
    "plot_arrays",
    "get_array",
    "get_diagonal_array",
    "get_unitary_array",
]


def plot_array(
    arr: np.ndarray,
    normalizer: MagnitudeNormalizer | None = None,
    combine_strategy: ColorCombineStrategy | None = None,
    show_values: bool = False,
    bar_orientation: Literal["vertical", "horizontal"] = "horizontal",
    show_individual_bars: bool = True,
    single_color_bar: bool = False,
) -> None:
    plot_arrays(
        arr,
        normalizer=normalizer,
        combine_strategy=combine_strategy,
        show_values=show_values,
        bar_orientation=bar_orientation,
        show_individual_bars=show_individual_bars,
        single_color_bar=single_color_bar,
    )


def plot_arrays(
    *arrays: np.ndarray,
    normalizer: MagnitudeNormalizer | None = None,
    combine_strategy: ColorCombineStrategy | None = None,
    show_values: bool = False,
    bar_orientation: Literal["vertical", "horizontal"] = "horizontal",
    show_individual_bars: bool = True,
    single_color_bar: bool = True,
) -> None:
    if not arrays:
        raise ValueError("plot_arrays requires at least one array")

    mapper = ComplexColorMapper(
        normalizer=normalizer,
        combine_strategy=combine_strategy,
    )
    plotter = ComplexArrayPlotter(mapper)
    plotter.plot_arrays(
        *arrays,
        show_values=show_values,
        bar_orientation=bar_orientation,
        show_individual_bars=show_individual_bars,
        single_color_bar=single_color_bar,
    )


def get_array(n: int) -> np.ndarray:
    return RandomMatrixBuilder.build(n)


def get_diagonal_array(
    n: int,
    diag_values: np.ndarray | None = None,
) -> np.ndarray:
    return DiagonalMatrixBuilder.build(n, diag_values)


def get_unitary_array(n: int) -> np.ndarray:
    return UnitaryMatrixBuilder.build(n)