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(settings.n_sd)
20    builder.add_dynamic(Coalescence(collision_kernel=settings.kernel))
21    particles = builder.build(attributes, products=(WallTime(),))
22
23    states = {}
24    last_wall_time = None
25    for step in settings.output_steps:
26        particles.run(step - particles.n_steps)
27        last_wall_time = particles.products["wall time"].get()
28
29    return states, last_wall_time
30
31
32def main(plot: bool):
33    settings = Settings()
34    settings.steps = [100, 3600] if "CI" not in os.environ else [1, 2]
35
36    times = {}
37    for backend in (ThrustRTC, Numba):
38        nsds = [2**n for n in range(12, 19, 3)]
39        key = backend.__name__
40        times[key] = []
41        for sd in nsds:
42            settings.n_sd = sd
43            _, wall_time = run(settings, backend())
44            times[key].append(wall_time)
45
46    for backend, t in times.items():
47        plt.plot(nsds, t, label=backend, linestyle="--", marker="o")
48    plt.ylabel("wall time [s]")
49    plt.xlabel("number of particles")
50    plt.grid()
51    plt.legend()
52    plt.loglog(base=2)
53    if plot:
54        plt.show()
55
56
57if __name__ == "__main__":
58    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(settings.n_sd)
21    builder.add_dynamic(Coalescence(collision_kernel=settings.kernel))
22    particles = builder.build(attributes, products=(WallTime(),))
23
24    states = {}
25    last_wall_time = None
26    for step in settings.output_steps:
27        particles.run(step - particles.n_steps)
28        last_wall_time = particles.products["wall time"].get()
29
30    return states, last_wall_time
def main(plot: bool):
33def main(plot: bool):
34    settings = Settings()
35    settings.steps = [100, 3600] if "CI" not in os.environ else [1, 2]
36
37    times = {}
38    for backend in (ThrustRTC, Numba):
39        nsds = [2**n for n in range(12, 19, 3)]
40        key = backend.__name__
41        times[key] = []
42        for sd in nsds:
43            settings.n_sd = sd
44            _, wall_time = run(settings, backend())
45            times[key].append(wall_time)
46
47    for backend, t in times.items():
48        plt.plot(nsds, t, label=backend, linestyle="--", marker="o")
49    plt.ylabel("wall time [s]")
50    plt.xlabel("number of particles")
51    plt.grid()
52    plt.legend()
53    plt.loglog(base=2)
54    if plot:
55        plt.show()