PySDM_examples.featured.Srivastava_1982.example

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