PySDM_examples.Srivastava_1982.example

  1from collections import namedtuple
  2
  3import numpy as np
  4from matplotlib import pyplot
  5from PySDM_examples.Srivastava_1982.equations import Equations, EquationsHelpers
  6from PySDM_examples.Srivastava_1982.settings import SimProducts
  7from PySDM_examples.Srivastava_1982.simulation import Simulation
  8
  9from PySDM.dynamics import Collision
 10from PySDM.dynamics.collisions.breakup_efficiencies import ConstEb
 11from PySDM.dynamics.collisions.breakup_fragmentations import ConstantMass
 12from PySDM.dynamics.collisions.coalescence_efficiencies import ConstEc
 13from PySDM.dynamics.collisions.collision_kernels import ConstantK
 14
 15NO_BOUNCE = ConstEb(1)
 16
 17
 18def coalescence_and_breakup_eq13(
 19    settings=None, n_steps=256, n_realisations=2, title=None, warn_overflows=True
 20):
 21    # arrange
 22    seeds = list(range(n_realisations))
 23
 24    collision_rate = settings.srivastava_c + settings.srivastava_beta
 25    simulation = Simulation(
 26        n_steps=n_steps,
 27        settings=settings,
 28        collision_dynamic=Collision(
 29            collision_kernel=ConstantK(a=collision_rate),
 30            coalescence_efficiency=ConstEc(settings.srivastava_c / collision_rate),
 31            breakup_efficiency=NO_BOUNCE,
 32            fragmentation_function=ConstantMass(c=settings.frag_mass),
 33            warn_overflows=warn_overflows,
 34        ),
 35    )
 36
 37    x = np.arange(n_steps + 1, dtype=float)
 38    sim_products = simulation.run_convergence_analysis(x, seeds=seeds)
 39
 40    secondary_products = get_pysdm_secondary_products(
 41        products=sim_products, total_volume=settings.total_volume
 42    )
 43
 44    pysdm_results = get_processed_results(secondary_products)
 45
 46    equations = Equations(
 47        M=settings.total_volume * settings.rho / settings.frag_mass,
 48        c=settings.srivastava_c,
 49        beta=settings.srivastava_beta,
 50    )
 51    equation_helper = EquationsHelpers(
 52        settings.total_volume,
 53        settings.total_number_0,
 54        settings.rho,
 55        frag_mass=settings.frag_mass,
 56    )
 57    m0 = equation_helper.m0()
 58
 59    x_log = compute_log_space(x)
 60    analytic_results = get_breakup_coalescence_analytic_results(
 61        equations, settings, m0, x, x_log
 62    )
 63
 64    prods = [
 65        k
 66        for k in list(pysdm_results.values())[0].keys()
 67        if k != SimProducts.PySDM.total_volume.name
 68    ]
 69
 70    add_to_plot_simulation_results(
 71        prods,
 72        settings.n_sds,
 73        x,
 74        pysdm_results=pysdm_results,
 75        analytic_results=analytic_results,
 76        title=title,
 77    )
 78
 79    results = namedtuple("_", "pysdm, analytic")(
 80        pysdm=pysdm_results, analytic=analytic_results
 81    )
 82    return results
 83
 84
 85def get_processed_results(res):
 86    processed = {}
 87    for n_sd in res.keys():
 88        processed[n_sd] = {}
 89        for prod in res[n_sd].keys():
 90            processed[n_sd][prod] = {}
 91            all_runs = np.asarray(list(res[n_sd][prod].values()))
 92
 93            processed[n_sd][prod]["avg"] = np.mean(all_runs, axis=0)
 94            processed[n_sd][prod]["max"] = np.max(all_runs, axis=0)
 95            processed[n_sd][prod]["min"] = np.min(all_runs, axis=0)
 96            processed[n_sd][prod]["std"] = np.std(all_runs, axis=0)
 97
 98    return processed
 99
100
101# TODO #1045 (not needed)
102def get_pysdm_secondary_products(products, total_volume):
103    pysdm_results = products
104    for n_sd in products.keys():
105        pysdm_results[n_sd][
106            SimProducts.Computed.mean_drop_volume_total_volume_ratio.name
107        ] = {}
108
109        for k in pysdm_results[n_sd][SimProducts.PySDM.total_numer.name].keys():
110            pysdm_results[n_sd][
111                SimProducts.Computed.mean_drop_volume_total_volume_ratio.name
112            ][k] = compute_drop_volume_total_volume_ratio(
113                mean_volume=total_volume
114                / products[n_sd][SimProducts.PySDM.total_numer.name][k],
115                total_volume=total_volume,
116            )
117    return pysdm_results
118
119
120def get_coalescence_analytic_results(equations, settings, m0, x, x_log):
121    mean_mass10 = (
122        equations.eq10(m0, equations.tau(x * settings.dt)) * settings.frag_mass
123    )
124    mean_mass_ratio_log = equations.eq10(m0, equations.tau(x_log * settings.dt))
125
126    return get_analytic_results(equations, settings, mean_mass10, mean_mass_ratio_log)
127
128
129def get_breakup_coalescence_analytic_results(equations, settings, m0, x, x_log):
130    mean_mass13 = (
131        equations.eq13(m0, equations.tau(x * settings.dt)) * settings.frag_mass
132    )
133    mean_mass_ratio_log = equations.eq13(m0, equations.tau(x_log * settings.dt))
134
135    return get_analytic_results(equations, settings, mean_mass13, mean_mass_ratio_log)
136
137
138def get_analytic_results(equations, settings, mean_mass, mean_mass_ratio):
139    res = {}
140    res[SimProducts.Computed.mean_drop_volume_total_volume_ratio.name] = (
141        compute_drop_volume_total_volume_ratio(
142            mean_volume=mean_mass / settings.rho, total_volume=settings.total_volume
143        )
144    )
145    res[SimProducts.PySDM.total_numer.name] = equations.M / mean_mass_ratio
146    return res
147
148
149def compute_log_space(x, shift=0, num_points=1000, eps=1e-1):
150    assert eps < x[1]
151    return (
152        np.logspace(np.log10(x[0] if x[0] != 0 else eps), np.log10(x[-1]), num_points)
153        + shift
154    )
155
156
157# TODO #1045 (not needed)
158def compute_drop_volume_total_volume_ratio(mean_volume, total_volume):
159    return mean_volume / total_volume * 100
160
161
162def add_to_plot_simulation_results(
163    prods,
164    n_sds,
165    x,
166    pysdm_results=None,
167    analytic_results=None,
168    title=None,
169):
170    fig = pyplot.figure(layout="constrained", figsize=(10, 4))
171    _wide = 14
172    _shrt = 8
173    _mrgn = 1
174    gs = fig.add_gridspec(nrows=20, ncols=2 * _wide + _shrt + 2 * _mrgn)
175    axs = (
176        fig.add_subplot(gs[2:, _shrt + _mrgn : _shrt + _mrgn + _wide]),
177        fig.add_subplot(gs[2:, 0:_shrt]),
178        fig.add_subplot(gs[2:, _shrt + 2 * _mrgn + _wide :]),
179    )
180    axs[1].set_ylim([6, 3500])
181
182    expons = [3, 5, 7, 9, 11]
183
184    axs[1].set_yscale(SimProducts.PySDM.super_particle_count.plot_yscale)
185    axs[1].set_yticks([2**e for e in expons], [f"$2^{{{e}}}$" for e in expons])
186
187    if title:
188        fig.suptitle(title)
189
190    ylims = {}
191    for prod in prods:
192        ylims[prod] = (np.inf, -np.inf)
193        for n_sd in n_sds:
194            ylims[prod] = (
195                min(ylims[prod][0], 0.75 * np.amin(pysdm_results[n_sd][prod]["avg"])),
196                max(ylims[prod][1], 1.25 * np.amax(pysdm_results[n_sd][prod]["avg"])),
197            )
198
199    for i, prod in enumerate(prods):
200        # plot numeric
201        if pysdm_results:
202            for n_sd in n_sds:
203                y_model = pysdm_results[n_sd][prod]
204
205                axs[i].step(
206                    x,
207                    y_model["avg"],
208                    where="mid",
209                    label=f"initial #SD: {n_sd}",
210                    linewidth=1 + np.log(n_sd) / 3,
211                )
212                axs[i].fill_between(
213                    x,
214                    y_model["avg"] - y_model["std"],
215                    y_model["avg"] + y_model["std"],
216                    alpha=0.2,
217                )
218
219        # plot analytic
220        if analytic_results:
221            add_analytic_result_to_axs(axs[i], prod, x, analytic_results)
222
223        # cosmetics
224        axs[i].set_ylabel(SimProducts.get_prod_by_name(prod).plot_title)
225
226        axs[i].grid()
227        axs[i].set_xlabel("step: t / dt")
228
229        if prod != SimProducts.PySDM.super_particle_count.name:
230            bottom = ylims[prod][0]
231            top = ylims[prod][1]
232
233            slope = (
234                pysdm_results[n_sds[-1]][prod]["avg"][-1]
235                - pysdm_results[n_sds[-1]][prod]["avg"][0]
236            )
237            if prod == SimProducts.PySDM.total_numer.name and slope < 0:
238                axs[i].set_ylim(0.2 * bottom, top)
239            else:
240                axs[i].set_ylim(bottom, top)
241
242    axs[0].legend()
243    return fig, axs
244
245
246def add_analytic_result_to_axs(axs_i, prod, x, res, key=""):
247    if prod != SimProducts.PySDM.super_particle_count.name:
248        x_theory = x
249        y_theory = res[prod]
250
251        if prod == SimProducts.PySDM.total_numer.name:
252            if y_theory.shape != x_theory.shape:
253                x_theory = compute_log_space(x)
254
255                axs_i.set_yscale(SimProducts.PySDM.total_numer.plot_yscale)
256                axs_i.set_xscale(SimProducts.PySDM.total_numer.plot_xscale)
257                axs_i.set_xlim(x_theory[0], None)
258
259        axs_i.plot(
260            x_theory,
261            y_theory,
262            label=f"analytic {key}",
263            linestyle=":",
264            linewidth=2,
265            color="black",
266        )