PySDM_examples.Bartman_2020_MasterThesis.fig_5_SCIPY_VS_ADAPTIVE

  1import os
  2
  3import matplotlib
  4import matplotlib.pyplot as plt
  5import numpy as np
  6from matplotlib.collections import LineCollection
  7from PySDM_examples.Arabas_and_Shima_2017.settings import setups
  8from PySDM_examples.Arabas_and_Shima_2017.simulation import Simulation
  9
 10from PySDM.backends import CPU, GPU
 11from PySDM.backends.impl_numba.test_helpers import scipy_ode_condensation_solver
 12
 13
 14def data(n_output, rtols, schemes, setups_num):
 15    resultant_data = {}
 16    for scheme in schemes:
 17        resultant_data[scheme] = {}
 18        if scheme == "SciPy":
 19            for rtol in rtols:
 20                resultant_data[scheme][rtol] = []
 21            for settings_idx in range(setups_num):
 22                settings = setups[settings_idx]
 23                settings.n_output = n_output
 24                simulation = Simulation(settings)
 25                scipy_ode_condensation_solver.patch_particulator(
 26                    simulation.particulator
 27                )
 28                results = simulation.run()
 29                for rtol in rtols:
 30                    resultant_data[scheme][rtol].append(results)
 31        else:
 32            for rtol in rtols:
 33                resultant_data[scheme][rtol] = []
 34                for settings_idx in range(setups_num):
 35                    settings = setups[settings_idx]
 36                    settings.rtol_x = rtol
 37                    settings.rtol_thd = rtol
 38                    settings.n_output = n_output
 39                    simulation = Simulation(
 40                        settings, backend=CPU if scheme == "CPU" else GPU
 41                    )
 42                    results = simulation.run()
 43                    resultant_data[scheme][rtol].append(results)
 44    return resultant_data
 45
 46
 47def add_color_line(fig, ax, x, y, z):
 48    points = np.array([x, y]).T.reshape(-1, 1, 2)
 49    segments = np.concatenate([points[:-1], points[1:]], axis=1)
 50    z = np.array(z)
 51    vmin = min(np.nanmin(z), np.nanmax(z) / 2)
 52    lc = LineCollection(
 53        segments,
 54        cmap=plt.get_cmap("plasma"),
 55        norm=matplotlib.colors.LogNorm(vmax=1, vmin=vmin),
 56    )
 57    lc.set_array(z)
 58    lc.set_linewidth(3)
 59
 60    ax.add_collection(lc)
 61    fig.colorbar(lc, ax=ax)
 62
 63
 64def plot(plot_data, rtols, schemes, setups_num, show_plot, path=None):
 65    _rtol = "$r_{tol}$"
 66
 67    plt.ioff()
 68    fig, axs = plt.subplots(
 69        setups_num, len(rtols), sharex=True, sharey=True, figsize=(10, 13)
 70    )
 71    for settings_idx in range(setups_num):
 72        SCIPY_S = None
 73        PySDM_S = None
 74        for rtol_idx, _rtol in enumerate(rtols):
 75            ax = axs[settings_idx, rtol_idx]
 76            for scheme in schemes:
 77                datum = plot_data[scheme][_rtol][settings_idx]
 78                S = datum["S"]
 79                z = datum["z"]
 80                dt = datum["dt_cond_min"]
 81                if scheme == "SciPy":
 82                    ax.plot(S, z, label=scheme, color="grey")
 83                    SCIPY_S = np.array(S)
 84                else:
 85                    add_color_line(fig, ax, S, z, dt)
 86                    PySDM_S = np.array(S)
 87            if SCIPY_S is not None and PySDM_S is not None:
 88                mae = np.mean(np.abs(SCIPY_S - PySDM_S))
 89                ax.set_title(f"MAE: {mae:.4E}")
 90            ax.set_xlim(-7.5e-3, 7.5e-3)
 91            ax.set_ylim(0, 180)
 92            ax.get_xaxis().set_minor_locator(matplotlib.ticker.AutoMinorLocator())
 93            ax.grid()
 94            plt.setp(ax.get_xticklabels(), rotation=30, horizontalalignment="right")
 95    for i, ax in enumerate(axs[:, 0]):
 96        ax.set(ylabel=r"$\bf{settings: " + str(i) + "}$\ndisplacement [m]")
 97    for i, ax in enumerate(axs[-1, :]):
 98        ax.set(xlabel="supersaturation\n" + r"$\bf{r_{tol}: " + str(rtols[i]) + "}$")
 99
100    plt.tight_layout()
101
102    if path is not None:
103        plt.savefig(path + ".pdf", format="pdf")
104    if show_plot:
105        plt.show()
106
107
108def main(save=None, show_plot=True):
109    rtols = [1e-7, 1e-11]
110    schemes = ["CPU", "SciPy"]
111    setups_num = len(setups)
112    input_data = data(80 if "CI" not in os.environ else 20, rtols, schemes, setups_num)
113    plot(input_data, rtols, schemes, setups_num, show_plot, save)
114
115
116if __name__ == "__main__":
117    main("SCIPY_VS_ADAPTIVE", show_plot="CI" not in os.environ)