1import matplotlib.pyplot as plt
2import numpy as np
3
4from PySDM.physics import constants as const
5
6
7class _Plot:
8 def __init__(self, fig, ax):
9 self.fig, self.ax = fig, ax
10 self.ax.set_title(" ")
11
12
13class _ImagePlot(_Plot):
14 line_args = {"color": "red", "alpha": 0.666, "linestyle": ":", "linewidth": 3}
15
16 def __init__(
17 self, fig, ax, grid, size, product, *, show=False, lines=False, cmap="YlGnBu"
18 ):
19 super().__init__(fig, ax)
20 self.nans = np.full(grid, np.nan)
21
22 self.dx = size[0] / grid[0]
23 self.dz = size[1] / grid[1]
24
25 xlim = (0, size[0])
26 zlim = (0, size[1])
27
28 self.ax.set_xlim(xlim)
29 self.ax.set_ylim(zlim)
30
31 if lines:
32 self.lines = {"X": [None] * 2, "Z": [None] * 2}
33 self.lines["X"][0] = plt.plot([-1] * 2, zlim, **self.line_args)[0]
34 self.lines["Z"][0] = plt.plot(xlim, [-1] * 2, **self.line_args)[0]
35 self.lines["X"][1] = plt.plot([-1] * 2, zlim, **self.line_args)[0]
36 self.lines["Z"][1] = plt.plot(xlim, [-1] * 2, **self.line_args)[0]
37
38 data = np.full_like(self.nans, np.nan)
39 label = f"{product.name} [{product.unit}]"
40
41 self.ax.set_xlabel("X [m]")
42 self.ax.set_ylabel("Z [m]")
43
44 self.im = self.ax.imshow(
45 self._transpose(data), origin="lower", extent=(*xlim, *zlim), cmap=cmap
46 )
47 plt.colorbar(self.im, ax=self.ax).set_label(label)
48 if show:
49 plt.show()
50
51 @staticmethod
52 def _transpose(data):
53 if data is not None:
54 return data.T
55 return None
56
57 def update(self, data, step, data_range):
58 data = self._transpose(data)
59 if data is not None:
60 self.im.set_data(data)
61 if data_range is not None:
62 self.im.set_clim(vmin=data_range[0], vmax=data_range[1])
63 nanmin = np.nan
64 nanmax = np.nan
65 if np.isfinite(data).any():
66 nanmin = np.nanmin(data)
67 nanmax = np.nanmax(data)
68 self.ax.set_title(
69 f"min:{nanmin: .3g} max:{nanmax: .3g} t/dt_out:{step: >6}"
70 )
71
72 def update_lines(self, focus_x, focus_z):
73 self.lines["X"][0].set_xdata(x=(focus_x[0] + 0.15) * self.dx)
74 self.lines["Z"][0].set_ydata(y=(focus_z[0] + 0.15) * self.dz)
75 self.lines["X"][1].set_xdata(x=(focus_x[1] - 0.25) * self.dx)
76 self.lines["Z"][1].set_ydata(y=(focus_z[1] - 0.25) * self.dz)
77
78
79class _SpectrumPlot(_Plot):
80 def __init__(self, r_bins, initial_spectrum_per_mass_of_dry_air, show=True):
81 super().__init__(*plt.subplots(1, 1))
82 self.ax.set_xlim(np.amin(r_bins), np.amax(r_bins))
83 self.ax.set_xlabel("particle radius [μm]")
84 self.ax.set_ylabel("specific concentration density [mg$^{-1}$ μm$^{-1}$]")
85 self.ax.set_xscale("log")
86 self.ax.set_yscale("log")
87 self.ax.set_ylim(1, 5e3)
88 self.ax.grid(True)
89 vals = initial_spectrum_per_mass_of_dry_air.size_distribution(
90 r_bins * const.si.um
91 )
92 const.convert_to(vals, const.si.mg**-1 / const.si.um)
93 self.ax.plot(r_bins, vals, label="spectrum sampled at t=0")
94 self.spec_wet = self.ax.step(
95 r_bins,
96 np.full_like(r_bins, np.nan),
97 label="binned super-particle wet sizes",
98 )[0]
99 self.spec_dry = self.ax.step(
100 r_bins,
101 np.full_like(r_bins, np.nan),
102 label="binned super-particle dry sizes",
103 )[0]
104 self.ax.legend()
105 if show:
106 plt.show()
107
108 def update_wet(self, data, step):
109 self.spec_wet.set_ydata(data)
110 self.ax.set_title(f"t/dt_out:{step}")
111
112 def update_dry(self, dry):
113 self.spec_dry.set_ydata(dry)
114
115
116class _TimeseriesPlot(_Plot):
117 def __init__(self, fig, ax, times, show=True):
118 super().__init__(fig, ax)
119 self.ax.set_xlim(0, times[-1])
120 self.ax.set_xlabel("time [s]")
121 self.ax.set_ylabel("rainfall [mm/day]")
122 self.ax.grid(True)
123 self.ydata = np.full_like(times, np.nan, dtype=float)
124 self.timeseries = self.ax.step(times, self.ydata, where="pre")[0]
125 if show:
126 plt.show()
127
128 def update(self, data, data_range):
129 if data is not None:
130 self.ydata[0 : len(data)] = data[:]
131 if data_range[0] != data_range[1]:
132 self.ax.set_ylim(data_range[0], 1.1 * data_range[1])
133 else:
134 self.ydata[:] = np.nan
135 self.timeseries.set_ydata(self.ydata)
136
137
138class _TemperaturePlot(_Plot):
139 def __init__(self, T_bins, formulae, show=True):
140 super().__init__(*plt.subplots(1, 1))
141 self.formulae = formulae
142 self.ax.set_xlim(np.amax(T_bins), np.amin(T_bins))
143 self.ax.set_xlabel("temperature [K]")
144 self.ax.set_ylabel("freezable fraction / cdf [1]")
145 self.ax.set_ylim(-0.05, 1.05)
146 self.ax.grid(True)
147 self.spec = self.ax.step(
148 T_bins,
149 np.full_like(T_bins, np.nan),
150 label="binned super-particle attributes",
151 where="mid",
152 )[0]
153 self.ax.legend()
154 if show:
155 plt.show()
156
157 def update(self, data, step):
158 self.ax.set_title(f"t/dt_out:{step}")
159 self.spec.set_ydata(data)
160
161
162class _TerminalVelocityPlot(_Plot):
163 def __init__(self, radius_bins, formulae, show=True):
164 self.formulae = formulae
165 super().__init__(*plt.subplots(1, 1))
166
167 self.ax.set_xlim(
168 np.amin(radius_bins) / const.si.um, np.amax(radius_bins) / const.si.um
169 )
170 self.ax.set_xlabel("radius [μm]")
171 self.ax.set_ylabel("mean terminal velocity [m/s]")
172 self.ax.set_ylim(0, 0.1)
173 self.ax.grid(True)
174
175 self.radius_bins = radius_bins
176
177 if show:
178 plt.show()
179
180 def update(self, data_min, data_max, step):
181 self.ax.set_title(f"t/dt_out:{step}")
182 self.ax.collections.clear()
183 self.ax.fill_between(
184 (self.radius_bins[:-1] + np.diff(self.radius_bins) / 2) / const.si.um,
185 data_min,
186 data_max,
187 color="gray",
188 )