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