Note

Go to the end to download the full example code.

SPH benchmark for an accretion disc#

This example tests the performance of the SPH solver for a locally isothermal accretion disc orbiting a central sink particle. It is the disc counterpart of the homogeneous density box benchmark: the resolution is automatically adapted to the available memory and number of processes, and the same step is replayed several times to measure the solver rate.

11 import datetime
12 import json
13 import math
14 import time
15 from statistics import mean, stdev
16
17 import shamrock
18
19 device_properties = shamrock.sys.get_compute_device_properties()
20
21 microbench_results = shamrock.sys.get_microbench_results(allow_run=True)
22 if len(microbench_results) == 0:
23     print("no microbench results, please run with --benchmark-mpi")
24     raise ValueError("no microbench results")
25
26 memory_gb = device_properties["global_mem_size"] / (1e9)
27
28 N_target_base = 2 ** int(math.log2(memory_gb * 1e6 / 1.5))
29 print(f"N_target_base = {N_target_base}")
30 print(f"memory_gb = {memory_gb}")
31 print(f"device_properties = {device_properties}")
32
33 N_target_base = min(N_target_base, 2**25)
34
35 if device_properties["type"] == "CPU":
36     N_target_base = min(N_target_base, 2**23)
37
38 shamrock.backends.reset_mem_info_max()
39
40 compute_multiplier = shamrock.sys.world_size()
41 # compute_multiplier = 12
42 scheduler_split_val = int(2e7)
43 scheduler_merge_val = 1
44
45 N_target = N_target_base * compute_multiplier

Disc parameters (same disc as the dusty disc perf test, without the dust)

 49 si = shamrock.UnitSystem()
 50 sicte = shamrock.Constants(si)
 51 codeu = shamrock.UnitSystem(
 52     unit_time=sicte.year(),  # year
 53     unit_length=sicte.au(),  # astro unit
 54     unit_mass=sicte.sol_mass(),
 55 )
 56
 57 center_mass = 1.0  # sol mass
 58 center_racc = 8.0  # au
 59
 60 disc = shamrock.utils.disc_setup.StandardDisc(
 61     units=codeu,
 62     center_mass=center_mass,
 63     disc_mass=0.05,  # sol mass
 64     rin=10.0,  # au
 65     rout=150.0,  # au
 66     H_r_0=0.1,
 67     q=0.5,
 68     p=3.0 / 2.0,
 69     r0=10.0,
 70     rotation="subkeplerian_3d",
 71     inner_tapering=True,
 72 )
 73
 74 bsize = disc.rout * 2
 75 bmin = (-bsize, -bsize, -bsize)
 76 bmax = (bsize, bsize, bsize)
 77
 78 if shamrock.sys.world_rank() == 0:
 79     print("N_target_base", N_target_base)
 80     print("compute_multiplier", compute_multiplier)
 81     print("scheduler_split_val", scheduler_split_val)
 82     print("scheduler_merge_val", scheduler_merge_val)
 83     print("N_target", N_target)
 84
 85 ctx = shamrock.Context()
 86 ctx.pdata_layout_new()
 87
 88 model = shamrock.get_Model_SPH(context=ctx, vector_type="f64_3", sph_kernel="M4")
 89
 90 cfg = model.gen_default_config()
 91 cfg.set_artif_viscosity_VaryingCD10(
 92     alpha_min=0.0, alpha_max=1, sigma_decay=0.1, alpha_u=1, beta_AV=2
 93 )
 94 cfg.set_eos_locally_isothermalLP07(cs0=disc.cs0(), q=disc.q, r0=disc.r0)
 95 cfg.add_kill_sphere(center=(0, 0, 0), radius=bsize)  # kill particles outside the simulation box
 96 cfg.set_units(codeu)
 97 cfg.set_particle_mass(disc.part_mass(N_target))
 98 cfg.set_cfl_cour(0.1)
 99 cfg.set_cfl_force(0.1)
100 cfg.print_status()
101 model.set_solver_config(cfg)
102 model.init_scheduler(scheduler_split_val, scheduler_merge_val)
103
104 model.resize_simulation_box(bmin, bmax)
105
106 setup = model.get_setup()
107 gen = disc.make_generator(setup, N_target, random_seed=666)
108
109 # Kind of optimized for Aurora
110 setup.apply_setup(
111     gen,
112     gen_step=int(scheduler_split_val / 8),
113     insert_step=int(scheduler_split_val * 2),
114     msg_count_limit=1024,
115     rank_comm_size_limit=int(scheduler_split_val) * 2,
116     max_msg_size=int(scheduler_split_val / 8),
117     do_setup_log=False,
118 )
119
120 model.add_sink(center_mass, (0, 0, 0), (0, 0, 0), center_racc)
121
122 shamrock.backends.reset_mem_info_max()
123
124 # converge smoothing length and compute initial dt
125 # (the disc is generated with a rough guess of h, hence the larger coarse tolerance)
126 model.change_htolerances(coarse=1.3, fine=1.1)
127 model.timestep()
128 model.change_htolerances(coarse=1.1, fine=1.1)
129
130 # Now run the actual benchmark for 10 steps
131 res_rates = []
132 res_cnts = []
133 res_system_metrics = []
134 res_mpi_timers = []
135
136 """
137 Here we insert callbacks to measure solver MPI usage by fetching the timers twice at the begining and end of the step
138 """
139 before_mpi_timers, after_mpi_timers = None, None
140
141
142 def callback_before_mpi_timer():
143     global before_mpi_timers
144     before_mpi_timers = shamrock.comm.get_timers()
145
146
147 def callback_after_mpi_timer():
148     global after_mpi_timers
149     after_mpi_timers = shamrock.comm.get_timers()
150
151
152 model.add_timestep_callback(step_begin=callback_before_mpi_timer, step_end=callback_after_mpi_timer)
153
154 for i in range(10):
155     if shamrock.sys.world_rank() == 0:
156         print("running step ", i + 1, "/", 10, " ...")
157
158     shamrock.sys.mpi_barrier()
159
160     # To replay the same step
161     model.set_next_dt(0.0)
162     model.timestep()
163
164     if shamrock.sys.world_rank() == 0:
165         print("collecting results ...")
166
167     tmp_res_rate, tmp_res_cnt, tmp_system_metrics = (
168         model.solver_logs_last_rate(),
169         model.solver_logs_last_obj_count(),
170         model.solver_logs_last_system_metrics(),
171     )
172     res_rates.append(tmp_res_rate)
173     res_cnts.append(tmp_res_cnt)
174     res_system_metrics.append(tmp_system_metrics)
175     res_mpi_timers.append(shamrock.comm.mpi_timers_delta(before_mpi_timers, after_mpi_timers))
176
177     if shamrock.sys.world_rank() == 0:
178         print("sleeping 1 second ...")
179
180     time.sleep(1)
181
182     if shamrock.sys.world_rank() == 0:
183         print("done sleeping 1 second ...")
184
185 # result is the best rate of the 10 steps
186 res_rate, res_cnt = max(res_rates), res_cnts[0]
187
188 # index of the max rate
189 max_rate_index = res_rates.index(max(res_rates))
190 max_rate_system_metrics = res_system_metrics[max_rate_index]
191 max_mpi_timers = res_mpi_timers[max_rate_index]
192 step_time = res_cnt / res_rate
193
194 if shamrock.sys.world_rank() == 0:
195     result_text = ""
196     result_text += f"--- final score for N_target_base={N_target_base} ---\n"
197     result_text += f"world size  : {shamrock.sys.world_size()}\n"
198     result_text += f"result rate : {res_rate}\n"
199     result_text += f"result cnt  : {res_cnt}\n"
200     result_text += f"cnt/rank    : {res_cnt / shamrock.sys.world_size()}\n"
201     result_text += f"result rate per rank : {res_rate / shamrock.sys.world_size()}\n"
202     result_text += f"rates infos : max={max(res_rates)}, min={min(res_rates)}, mean={mean(res_rates)}, stddev={stdev(res_rates)}\n"
203     result_text += f"res_rates = {res_rates}\n"
204     result_text += f"res_cnts = {res_cnts}\n"
205     result_text += f"step time = {step_time}\n"
206
207     dic_out = {
208         "device_properties": device_properties,
209         "microbench_results": shamrock.sys.get_microbench_results(),
210         "shamrock_version": shamrock.version_string(),
211         "shamrock_compiler_id_string": shamrock.get_compiler_id_string(),
212         "shamrock_compile_flags": shamrock.get_compile_arg(),
213         "date": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
214         "world_size": shamrock.sys.world_size(),
215         "rate": res_rate,
216         "cnt": res_cnt,
217         "step_time": step_time,
218         "mpi_timers": max_mpi_timers,
219     }
220
221     # print the system metrics
222     metrics_duration = max_rate_system_metrics["duration"]
223     result_text += "system metrics:\n"
224     for key, value in max_rate_system_metrics.items():
225         if not key == "duration":
226             result_text += f"{key}: {value} J\n"
227             dic_out[key] = value
228
229     for key, value in max_rate_system_metrics.items():
230         if not key == "duration":
231             result_text += f"avg power {key} / step time : {value / metrics_duration} W\n"
232             dic_out[f"power_{key}"] = value / metrics_duration
233
234     dic_out["system_metric_duration"] = metrics_duration
235
236     result_text += "---------submit this result--------\n"
237     result_text += f"{json.dumps(dic_out, indent=4)}\n"
238     result_text += "-----------------------------------\n"
239
240     print("current results:")
241     print(result_text)

Estimated memory usage: 0 MB

Gallery generated by Sphinx-Gallery