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