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,
20        backend=backend(formulae=settings.formulae),
21        environment=env,
22    )
23    attributes = {}
24    sampling = ConstantMultiplicity(settings.spectrum)
25    attributes["volume"], attributes["multiplicity"] = sampling.sample(settings.n_sd)
26    coalescence = Coalescence(
27        collision_kernel=settings.kernel, adaptive=settings.adaptive
28    )
29    builder.add_dynamic(coalescence)
30    products = (
31        ParticleVolumeVersusRadiusLogarithmSpectrum(
32            settings.radius_bins_edges, name="dv/dlnr"
33        ),
34        WallTime(),
35    )
36    particulator = builder.build(attributes, products)
37
38    for observer in observers:
39        particulator.observers.append(observer)
40
41    vals = {}
42    particulator.products["wall time"].reset()
43    for step in settings.output_steps:
44        particulator.run(step - particulator.n_steps)
45        vals[step] = particulator.products["dv/dlnr"].get()[0]
46        vals[step][:] *= settings.rho
47
48    exec_time = particulator.products["wall time"].get()
49    return vals, exec_time
50
51
52def main(plot: bool, save: Optional[str]):
53    with np.errstate(all="raise"):
54        settings = Settings()
55
56        settings.n_sd = 2**15
57
58        states, _ = run(settings)
59
60    with np.errstate(invalid="ignore"):
61        plotter = SpectrumPlotter(settings)
62        plotter.smooth = True
63        for step, vals in states.items():
64            _ = plotter.plot(vals, step * settings.dt)
65            # assert _ < 200  # TODO #327
66        if save is not None:
67            n_sd = settings.n_sd
68            plotter.save(save + "/" + f"{n_sd}_shima_fig_2" + "." + plotter.format)
69        if plot:
70            plotter.show()
71
72
73if __name__ == "__main__":
74    main(plot="CI" not in os.environ, save=None)
def run( settings, backend=functools.partial(<function _cached_backend>, backend_class=<class 'PySDM.backends.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,
21        backend=backend(formulae=settings.formulae),
22        environment=env,
23    )
24    attributes = {}
25    sampling = ConstantMultiplicity(settings.spectrum)
26    attributes["volume"], attributes["multiplicity"] = sampling.sample(settings.n_sd)
27    coalescence = Coalescence(
28        collision_kernel=settings.kernel, adaptive=settings.adaptive
29    )
30    builder.add_dynamic(coalescence)
31    products = (
32        ParticleVolumeVersusRadiusLogarithmSpectrum(
33            settings.radius_bins_edges, name="dv/dlnr"
34        ),
35        WallTime(),
36    )
37    particulator = builder.build(attributes, products)
38
39    for observer in observers:
40        particulator.observers.append(observer)
41
42    vals = {}
43    particulator.products["wall time"].reset()
44    for step in settings.output_steps:
45        particulator.run(step - particulator.n_steps)
46        vals[step] = particulator.products["dv/dlnr"].get()[0]
47        vals[step][:] *= settings.rho
48
49    exec_time = particulator.products["wall time"].get()
50    return vals, exec_time
def main(plot: bool, save: Optional[str]):
53def main(plot: bool, save: Optional[str]):
54    with np.errstate(all="raise"):
55        settings = Settings()
56
57        settings.n_sd = 2**15
58
59        states, _ = run(settings)
60
61    with np.errstate(invalid="ignore"):
62        plotter = SpectrumPlotter(settings)
63        plotter.smooth = True
64        for step, vals in states.items():
65            _ = plotter.plot(vals, step * settings.dt)
66            # assert _ < 200  # TODO #327
67        if save is not None:
68            n_sd = settings.n_sd
69            plotter.save(save + "/" + f"{n_sd}_shima_fig_2" + "." + plotter.format)
70        if plot:
71            plotter.show()