PySDM_examples.Shima_et_al_2009.example

 1import os
 2from typing import Optional
 3
 4import numpy as np
 5from PySDM_examples.Shima_et_al_2009.settings import Settings
 6from PySDM_examples.Shima_et_al_2009.spectrum_plotter import SpectrumPlotter
 7
 8from PySDM.backends import CPU
 9from PySDM.builder import Builder
10from PySDM.dynamics import Coalescence
11from PySDM.environments import Box
12from PySDM.initialisation.sampling.spectral_sampling import ConstantMultiplicity
13from PySDM.products import ParticleVolumeVersusRadiusLogarithmSpectrum, WallTime
14
15
16def run(settings, backend=CPU, observers=()):
17    env = Box(dv=settings.dv, dt=settings.dt)
18    builder = Builder(
19        n_sd=settings.n_sd, backend=backend(formulae=settings.formulae), environment=env
20    )
21    attributes = {}
22    sampling = ConstantMultiplicity(settings.spectrum)
23    attributes["volume"], attributes["multiplicity"] = sampling.sample(settings.n_sd)
24    coalescence = Coalescence(
25        collision_kernel=settings.kernel, adaptive=settings.adaptive
26    )
27    builder.add_dynamic(coalescence)
28    products = (
29        ParticleVolumeVersusRadiusLogarithmSpectrum(
30            settings.radius_bins_edges, name="dv/dlnr"
31        ),
32        WallTime(),
33    )
34    particulator = builder.build(attributes, products)
35
36    for observer in observers:
37        particulator.observers.append(observer)
38
39    vals = {}
40    particulator.products["wall time"].reset()
41    for step in settings.output_steps:
42        particulator.run(step - particulator.n_steps)
43        vals[step] = particulator.products["dv/dlnr"].get()[0]
44        vals[step][:] *= settings.rho
45
46    exec_time = particulator.products["wall time"].get()
47    return vals, exec_time
48
49
50def main(plot: bool, save: Optional[str]):
51    with np.errstate(all="raise"):
52        settings = Settings()
53
54        settings.n_sd = 2**15
55
56        states, _ = run(settings)
57
58    with np.errstate(invalid="ignore"):
59        plotter = SpectrumPlotter(settings)
60        plotter.smooth = True
61        for step, vals in states.items():
62            _ = plotter.plot(vals, step * settings.dt)
63            # assert _ < 200  # TODO #327
64        if save is not None:
65            n_sd = settings.n_sd
66            plotter.save(save + "/" + f"{n_sd}_shima_fig_2" + "." + plotter.format)
67        if plot:
68            plotter.show()
69
70
71if __name__ == "__main__":
72    main(plot="CI" not in os.environ, save=None)
def run(settings, backend=<class 'PySDM.backends.numba.Numba'>, observers=()):
17def run(settings, backend=CPU, observers=()):
18    env = Box(dv=settings.dv, dt=settings.dt)
19    builder = Builder(
20        n_sd=settings.n_sd, backend=backend(formulae=settings.formulae), environment=env
21    )
22    attributes = {}
23    sampling = ConstantMultiplicity(settings.spectrum)
24    attributes["volume"], attributes["multiplicity"] = sampling.sample(settings.n_sd)
25    coalescence = Coalescence(
26        collision_kernel=settings.kernel, adaptive=settings.adaptive
27    )
28    builder.add_dynamic(coalescence)
29    products = (
30        ParticleVolumeVersusRadiusLogarithmSpectrum(
31            settings.radius_bins_edges, name="dv/dlnr"
32        ),
33        WallTime(),
34    )
35    particulator = builder.build(attributes, products)
36
37    for observer in observers:
38        particulator.observers.append(observer)
39
40    vals = {}
41    particulator.products["wall time"].reset()
42    for step in settings.output_steps:
43        particulator.run(step - particulator.n_steps)
44        vals[step] = particulator.products["dv/dlnr"].get()[0]
45        vals[step][:] *= settings.rho
46
47    exec_time = particulator.products["wall time"].get()
48    return vals, exec_time
def main(plot: bool, save: Optional[str]):
51def main(plot: bool, save: Optional[str]):
52    with np.errstate(all="raise"):
53        settings = Settings()
54
55        settings.n_sd = 2**15
56
57        states, _ = run(settings)
58
59    with np.errstate(invalid="ignore"):
60        plotter = SpectrumPlotter(settings)
61        plotter.smooth = True
62        for step, vals in states.items():
63            _ = plotter.plot(vals, step * settings.dt)
64            # assert _ < 200  # TODO #327
65        if save is not None:
66            n_sd = settings.n_sd
67            plotter.save(save + "/" + f"{n_sd}_shima_fig_2" + "." + plotter.format)
68        if plot:
69            plotter.show()