-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathutils.py
127 lines (101 loc) · 4.01 KB
/
utils.py
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
116
117
118
119
120
121
122
123
124
125
126
127
import sys
# import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from solver import solve_poisson
# ----------------------------------------------------------------------------
###########################
######### GENERAL #########
###########################
def progbar(iterable, length=30, symbol='='):
"""Wrapper generator function for an iterable.
Prints a progressbar when yielding an item. \\
Args:
iterable: an object supporting iteration
length: length of the progressbar
"""
n = len(iterable)
for i, item in enumerate(iterable):
steps = int(length/n*(i+1))
sys.stdout.write('\r')
sys.stdout.write(f"[{symbol*steps:{length}}] {(100/n*(i+1)):.1f}%")
sys.stdout.flush()
yield item
# ----------------------------------------------------------------------------
###########################
###### VISUALIZATION ######
###########################
def plot_images(images, labels=[], nrow=3, cmap='viridis'):
n_images = len(images)
for i, image in enumerate(images):
plt.subplot(-(n_images//-nrow), nrow, i+1)
plt.imshow(image, cmap=cmap)
if labels:
plt.title(labels[i], fontdict={'fontsize': 22})
plt.xlabel('x', fontdict={'fontsize': 18})
plt.ylabel('y', fontdict={'fontsize': 18})
plt.colorbar(shrink=0.35)
# ----------------------------------------------------------------------------
def make_colored_plot(A, num_levels=None, s=None, mul=1., cmap='afmhot'):
s = A.max()-A.min() if s is None else s
cms = cm.get_cmap(cmap)
if num_levels is None:
return cms(mul*(A - A.min()) / s)
Ac = np.floor((A - A.min())/s*num_levels) / num_levels
return cms(mul*Ac)
# ----------------------------------------------------------------------------
def make_frames(values,
U, configU=None, configU_t=None,
Rho=None, configRho=None, configRho_t=None,
Eps=None, configEps=None, configEps_t=None,
out=['phi', 'E']
):
out_dict = {k: [] for k in out}
if configU is not None:
configU_t = lambda _: configU
if configRho is not None:
configRho_t = lambda _: configRho
if configEps is not None:
configEps_t = lambda _: configEps
Uc, RhoC, EpsC = U, Rho, Eps
for val in progbar(values):
if configU_t is not None:
Uc = configU_t(val)(U.copy())
if configRho_t is not None:
RhoC = configRho_t(val)(Rho.copy())
if configEps_t is not None:
EpsC = configEps_t(val)(Eps.copy())
if 'U' in out:
out_dict['U'].append(Uc)
if 'phi' in out:
phi = solve_poisson(Uc, RhoC, EpsC)
out_dict['phi'].append(phi)
if 'E' in out:
E = get_E_abs(phi)
out_dict['E'].append(E)
return out_dict
# ----------------------------------------------------------------------------
def colorize_frames(frames, num_levels=None, mul=1.):
max_val = np.max([frame - np.min(frame) for frame in frames])
frames_c = [make_colored_plot(frame,
num_levels=num_levels,
s=max_val,
mul=mul,
cmap='afmhot')[:,:,:3]
for frame in frames]
return frames_c
# ----------------------------------------------------------------------------
def plot_frames(frames, num=20, nrow=5):
for i, frame in enumerate(frames[::len(frames)//num]):
plt.subplot(-(num//-nrow), nrow, i+1)
plt.imshow(frame)
# ----------------------------------------------------------------------------
###########################
######### PHYSICS #########
###########################
def get_E_abs(phi):
Ey, Ex = np.gradient(phi)
E = np.sqrt(Ex**2 + Ey**2)
return E
# ----------------------------------------------------------------------------