PySDM_examples.Bulenok_2023_MasterThesis.utils

  1import gc
  2import json
  3import os
  4import time
  5
  6import numba
  7import numpy as np
  8from matplotlib import pyplot
  9from open_atmos_jupyter_utils import show_plot
 10
 11from PySDM.backends import CPU, GPU
 12
 13
 14class ProductsNames:
 15    super_particle_count = "super_particle_count"
 16    total_volume = "total_volume"
 17    total_number = "total_number"
 18
 19
 20def print_all_products(particulator):
 21    print(
 22        ProductsNames.total_number,
 23        particulator.products[ProductsNames.total_number].get(),
 24    )
 25    print(
 26        ProductsNames.total_volume,
 27        particulator.products[ProductsNames.total_volume].get(),
 28    )
 29    print(
 30        ProductsNames.super_particle_count,
 31        particulator.products[ProductsNames.super_particle_count].get(),
 32    )
 33
 34
 35def get_prod_dict(particulator):
 36    d = {
 37        ProductsNames.total_number: list(
 38            particulator.products[ProductsNames.total_number].get()
 39        ),
 40        ProductsNames.total_volume: list(
 41            particulator.products[ProductsNames.total_volume].get()
 42        ),
 43        ProductsNames.super_particle_count: list(
 44            particulator.products[ProductsNames.super_particle_count].get()
 45        ),
 46    }
 47
 48    return d
 49
 50
 51def measure_time_for_each_step(particulator, n_steps):
 52    particulator.run(steps=1)
 53
 54    res = []
 55    for _ in range(n_steps):
 56        t0 = time.time()
 57        particulator.run(steps=1)
 58        t1 = time.time()
 59
 60        res.append(t1 - t0)
 61
 62    return res
 63
 64
 65def measure_time_per_timestep(particulator, n_steps):
 66    particulator.run(steps=1)
 67
 68    t0 = time.time()
 69    particulator.run(steps=n_steps)
 70    t1 = time.time()
 71
 72    return (t1 - t0) / n_steps
 73
 74
 75def go_benchmark(
 76    setup_sim,
 77    n_sds,
 78    n_steps,
 79    seeds,
 80    numba_n_threads=None,
 81    double_precision=True,
 82    sim_run_filename=None,
 83    total_number=None,
 84    dv=None,
 85    time_measurement_fun=measure_time_per_timestep,
 86    backends=(CPU, GPU),
 87):
 88    products = {}
 89    results = {}
 90
 91    backend_configs = []
 92    if CPU in backends:
 93        cpu_backends_configs = [(CPU, i) for i in numba_n_threads]
 94        backend_configs = [*backend_configs, *cpu_backends_configs]
 95    if GPU in backends:
 96        backend_configs.append((GPU, None))
 97
 98    for backend_class, n_threads in backend_configs:
 99        backend_name = backend_class.__name__
100        if n_threads:
101            numba.set_num_threads(n_threads)
102            backend_name += "_" + str(numba.get_num_threads())
103
104        results[backend_name] = {}
105        products[backend_name] = {}
106
107        print()
108        print("before")
109
110        for n_sd in n_sds:
111            print("\n")
112            print(backend_name, n_sd)
113
114            results[backend_name][n_sd] = {}
115            products[backend_name][n_sd] = {}
116
117            for seed in seeds:
118                gc.collect()
119
120                particulator = setup_sim(
121                    n_sd,
122                    backend_class,
123                    seed,
124                    double_precision=double_precision,
125                    total_number=total_number,
126                    dv=dv,
127                )
128
129                print()
130                print("products before simulation")
131                print_all_products(particulator)
132
133                print()
134                print("start simulation")
135
136                elapsed_time = time_measurement_fun(particulator, n_steps)
137
138                print()
139                print("products after simulation")
140                print_all_products(particulator)
141
142                results[backend_name][n_sd][seed] = elapsed_time
143                products[backend_name][n_sd][seed] = get_prod_dict(particulator)
144
145                gc.collect()
146                del particulator
147                gc.collect()
148
149    if sim_run_filename:
150        write_to_file(filename=f"{sim_run_filename}-products.txt", d=products)
151
152    return results
153
154
155def process_results(res_d, axis=None):
156    processed_d = {}
157    for backend in res_d.keys():
158        processed_d[backend] = {}
159
160        for n_sd in res_d[backend].keys():
161            processed_d[backend][n_sd] = {}
162
163            vals = res_d[backend][n_sd].values()
164            vals = np.array(list(vals))
165
166            processed_d[backend][n_sd]["mean"] = np.mean(vals, axis=axis)
167            processed_d[backend][n_sd]["std"] = np.std(vals, axis=axis)
168            processed_d[backend][n_sd]["max"] = np.amax(vals, axis=axis)
169            processed_d[backend][n_sd]["min"] = np.amin(vals, axis=axis)
170
171    return processed_d
172
173
174def write_to_file(filename, d):
175    assert not os.path.isfile(filename), filename
176
177    with open(filename, "w", encoding="utf-8") as fp:
178        json.dump(d, fp)
179
180
181class PlottingHelpers:
182    @staticmethod
183    def get_backend_markers(backends):
184        markers = {backend: "o" if "Numba" in backend else "x" for backend in backends}
185        return markers
186
187    @staticmethod
188    def get_sorted_backend_list(processed_d):
189        backends = list(processed_d.keys())
190
191        backends.sort()
192        backends.sort(key=lambda x: int(x[6:]) if "Numba_" in x else 100**10)
193
194        return backends
195
196    @staticmethod
197    def get_n_sd_list(backends, processed_d):
198        x = []
199
200        for backend in backends:
201            for n_sd in processed_d[backend].keys():
202                if n_sd not in x:
203                    x.append(n_sd)
204
205        x.sort()
206        return x
207
208
209def plot_processed_results(
210    processed_d,
211    show=True,
212    plot_label="",
213    plot_title=None,
214    metric="min",
215    plot_filename=None,
216    markers=None,
217    colors=None,
218):
219    backends = PlottingHelpers.get_sorted_backend_list(processed_d)
220
221    if markers is None:
222        markers = PlottingHelpers.get_backend_markers(backends)
223
224    x = PlottingHelpers.get_n_sd_list(backends, processed_d)
225
226    for backend in backends:
227        y = []
228        for n_sd in x:
229            v = processed_d[backend][n_sd][metric]
230            assert isinstance(v, (float, int)), "must be scalar"
231            y.append(v)
232
233        if colors:
234            pyplot.plot(
235                x,
236                y,
237                label=backend + plot_label,
238                marker=markers[backend],
239                color=colors[backend],
240                linewidth=2,
241            )
242        else:
243            pyplot.plot(
244                x, y, label=backend + plot_label, marker=markers[backend], linewidth=2
245            )
246
247    pyplot.legend()
248    pyplot.xscale("log", base=2)
249    pyplot.yscale("log", base=2)
250    pyplot.ylim(bottom=2**-15, top=2**3)
251
252    pyplot.grid()
253    pyplot.xticks(x)
254    pyplot.xlabel("number of super-droplets")
255    pyplot.ylabel("wall time per timestep [s]")
256
257    if plot_title:
258        pyplot.title(plot_title)
259
260    if show:
261        if plot_filename:
262            show_plot(filename=plot_filename)
263        else:
264            pyplot.show()
265
266
267def plot_processed_on_same_plot(coal_d, break_d, coal_break_d):
268    plot_processed_results(coal_d, plot_label="-c", show=False)
269    plot_processed_results(break_d, plot_label="-b", show=False)
270    plot_processed_results(coal_break_d, plot_label="-cb", show=False)
271
272    show_plot()
273
274
275def plot_time_per_step(
276    processed_d,
277    n_sd,
278    show=True,
279    plot_label="",
280    plot_title=None,
281    metric="mean",
282    plot_filename=None,
283    step_from_to=None,
284):
285    backends = PlottingHelpers.get_sorted_backend_list(processed_d)
286
287    markers = PlottingHelpers.get_backend_markers(backends)
288
289    for backend in backends:
290        y = processed_d[backend][n_sd][metric]
291        x = np.arange(len(y))
292
293        if step_from_to is not None:
294            x = x[step_from_to[0] : step_from_to[1]]
295            y = y[step_from_to[0] : step_from_to[1]]
296
297        pyplot.plot(x, y, label=backend + plot_label, marker=markers[backend])
298
299    pyplot.legend()
300    pyplot.grid()
301    pyplot.xticks(x)
302    pyplot.xlabel("number of super-droplets")
303    pyplot.ylabel("wall time per timestep [s]")
304
305    if plot_title:
306        pyplot.title(plot_title + f"(n_sd: {n_sd})")
307
308    if show:
309        if plot_filename:
310            show_plot(filename=plot_filename)
311        else:
312            pyplot.show()