complex_field_visualization.py 768 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import numpy as np
from matplotlib.colors import hsv_to_rgb


def complex_to_hsv(arr, background='dark'):
    """
    2d complex field to colorfull field
    Args:
        arr: numpy array with complex dtype
        background: 'dark' or 'light'

    Returns:

    """
Matthijs's avatar
Matthijs committed
15
    norm = np.max([np.max(abs(arr)), 1e-15])
16
17
18
19
    objecthsv = np.zeros((3,) + arr.shape)
    objecthsv[0] = np.angle(-arr) / 2 / np.pi + 0.5
    if background == 'dark':
        objecthsv[1] = np.ones(arr.shape)
Matthijs's avatar
Matthijs committed
20
        objecthsv[2] = abs(arr) / norm
21
22
    else:
        objecthsv[2] = np.ones(arr.shape)
Matthijs's avatar
Matthijs committed
23
        objecthsv[1] = abs(arr) / norm
24
25
26
27
28
    return np.transpose(objecthsv, (1, 2, 0))


def complex_to_rgb(arr, background='dark'):
    return hsv_to_rgb(complex_to_hsv(arr, background=background))