PySDM_examples.Shima_et_al_2009.example_timing

 1import os
 2
 3from matplotlib import pyplot as plt
 4from PySDM_examples.Shima_et_al_2009.settings import Settings
 5
 6from PySDM.backends import Numba, ThrustRTC
 7from PySDM.builder import Builder
 8from PySDM.dynamics import Coalescence
 9from PySDM.environments import Box
10from PySDM.initialisation.sampling.spectral_sampling import ConstantMultiplicity
11from PySDM.products import WallTime
12
13
14def run(settings, backend):
15    env = Box(dv=settings.dv, dt=settings.dt)
16    builder = Builder(n_sd=settings.n_sd, backend=backend, environment=env)
17    attributes = {}
18    sampling = ConstantMultiplicity(settings.spectrum)
19    attributes["volume"], attributes["multiplicity"] = sampling.sample_deterministic(
20        settings.n_sd
21    )
22    builder.add_dynamic(Coalescence(collision_kernel=settings.kernel))
23    particles = builder.build(attributes, products=(WallTime(),))
24
25    states = {}
26    last_wall_time = None
27    for step in settings.output_steps:
28        particles.run(step - particles.n_steps)
29        last_wall_time = particles.products["wall time"].get()
30
31    return states, last_wall_time
32
33
34def main(plot: bool):
35    settings = Settings()
36    settings.steps = [100, 3600] if "CI" not in os.environ else [1, 2]
37
38    times = {}
39    for backend in (ThrustRTC, Numba):
40        nsds = [2**n for n in range(12, 19, 3)]
41        key = backend.__name__
42        times[key] = []
43        for sd in nsds:
44            settings.n_sd = sd
45            _, wall_time = run(settings, backend())
46            times[key].append(wall_time)
47
48    for backend, t in times.items():
49        plt.plot(nsds, t, label=backend, linestyle="--", marker="o")
50    plt.ylabel("wall time [s]")
51    plt.xlabel("number of particles")
52    plt.grid()
53    plt.legend()
54    plt.loglog(base=2)
55    if plot:
56        plt.show()
57
58
59if __name__ == "__main__":
60    main(plot="CI" not in os.environ)
def run(settings, backend):
15def run(settings, backend):
16    env = Box(dv=settings.dv, dt=settings.dt)
17    builder = Builder(n_sd=settings.n_sd, backend=backend, environment=env)
18    attributes = {}
19    sampling = ConstantMultiplicity(settings.spectrum)
20    attributes["volume"], attributes["multiplicity"] = sampling.sample_deterministic(
21        settings.n_sd
22    )
23    builder.add_dynamic(Coalescence(collision_kernel=settings.kernel))
24    particles = builder.build(attributes, products=(WallTime(),))
25
26    states = {}
27    last_wall_time = None
28    for step in settings.output_steps:
29        particles.run(step - particles.n_steps)
30        last_wall_time = particles.products["wall time"].get()
31
32    return states, last_wall_time
def main(plot: bool):
35def main(plot: bool):
36    settings = Settings()
37    settings.steps = [100, 3600] if "CI" not in os.environ else [1, 2]
38
39    times = {}
40    for backend in (ThrustRTC, Numba):
41        nsds = [2**n for n in range(12, 19, 3)]
42        key = backend.__name__
43        times[key] = []
44        for sd in nsds:
45            settings.n_sd = sd
46            _, wall_time = run(settings, backend())
47            times[key].append(wall_time)
48
49    for backend, t in times.items():
50        plt.plot(nsds, t, label=backend, linestyle="--", marker="o")
51    plt.ylabel("wall time [s]")
52    plt.xlabel("number of particles")
53    plt.grid()
54    plt.legend()
55    plt.loglog(base=2)
56    if plot:
57        plt.show()