Note

Go to the end to download the full example code.

Dusty SPH disc#

A disc but with dust

 8 import os
 9 from enum import Enum
10
11 import matplotlib.pyplot as plt
12 import numpy as np
13 from shamrock.external import coala
14 from shamrock.utils.analysis import StandardPlotHelper
15 from shamrock.utils.DustMRNDistribution import DustMRNDistribution
16 from shamrock.utils.SimulationRunner import SimulationRunner, callback, simulation_setup
17
18 import shamrock
19
20 shamrock.enable_experimental_features()
21
22 # If we use the shamrock executable to run this script instead of the python interpreter,
23 # we should not initialize the system as the shamrock executable needs to handle specific MPI logic
24 if not shamrock.sys.is_initialized():
25     shamrock.change_loglevel(1)
26     shamrock.sys.init("0:0")

Use shamrock documentation style for matplotlib

30 shamrock.matplotlib.set_shamrock_mpl_style()

COALA only acts on dust, so ndust = 0 runs without it

 35 configs = [
 36     {"kernel": "M4", "ndust": 0, "coala": False},
 37     {"kernel": "M6", "ndust": 0, "coala": False},
 38 ]
 39
 40 for n in range(1, 100):
 41     for use_coala in (False, True):
 42         configs.append({"kernel": "M6", "ndust": n, "coala": use_coala})
 43
 44
 45 for c in configs:
 46     # %%
 47     # Sim parameters
 48     si = shamrock.UnitSystem()
 49     sicte = shamrock.Constants(si)
 50     codeu = shamrock.UnitSystem(
 51         unit_time=sicte.year(),  # year
 52         unit_length=sicte.au(),  # astro unit
 53         unit_mass=sicte.sol_mass(),
 54     )
 55     ucte = shamrock.Constants(codeu)
 56     codeu_kg_m3 = codeu.get("kg") * codeu.get("m", power=-3)
 57
 58     # Resolution
 59     Npart = int(os.environ.get("NPART", "100000"))
 60     print(f"Npart = {Npart} (NPART={os.environ.get('NPART', 'not set')!r})")
 61
 62     # Domain decomposition parameters
 63     scheduler_split_val = int(1.0e7)  # split patches with more than 1e7 particles
 64     scheduler_merge_val = scheduler_split_val // 16
 65
 66     # Dump and plot frequency and duration of the simulation
 67     analysis_frequency = 2.5  # years
 68     high_freq_analysis_frequency = 0.1  # years
 69     t_end = 1000.0  # years
 70     t_inject = 0.0  # years
 71
 72     # Sink parameters
 73     center_mass = 1.0  # sol mass
 74     center_racc = 8.0  # au
 75
 76     # Disc parameters
 77     disc = shamrock.utils.disc_setup.StandardDisc(
 78         units=codeu,
 79         center_mass=center_mass,
 80         disc_mass=0.05,  # sol mass
 81         rin=10.0,  # au
 82         rout=150.0,  # au
 83         H_r_0=0.1,
 84         q=0.5,
 85         p=3.0 / 2.0,
 86         r0=10.0,
 87         rotation="subkeplerian_3d",
 88         inner_tapering=True,
 89     )
 90
 91     # Solver parameters
 92     kernel = c["kernel"]
 93
 94     # Hydro parameters
 95     alpha_min = 0.0
 96     alpha_max = 1
 97     sigma_decay = 0.1
 98     alpha_u = 1
 99     beta_AV = 2
100
101     # Dust parameters
102     ndust = c["ndust"]
103     gamma = 1.4
104
105     use_coala = c["coala"] and ndust > 0
106
107     print(f"ndust = {ndust}, use_coala = {use_coala}")
108
109     mrn_pow = 3.5
110     mrn_cutoff_si = np.inf  # would be 250e-9 normally
111
112     epsilon_base = 0.01
113
114     rho_grains_si_edges = np.array([2.3 * 1000 for _ in range(ndust + 1)])  # 2.3 g.cm^-3
115     grain_size_si_edges = np.logspace(-6.5, -1.5, ndust + 1)  # 10um -> 1mm
116
117     mrn_distribution = DustMRNDistribution(
118         codeu, mrn_pow, mrn_cutoff_si, grain_size_si_edges, rho_grains_si_edges
119     )
120
121     class DustLimiter(Enum):
122         NONE = "none"
123         SMOOTH = "smooth"
124         BALLABIO = "ballabio"
125         HARD = "hard"
126
127     _dust_limiter_env = os.environ.get("DUST_LIMITER", "none")
128     limiter = DustLimiter(_dust_limiter_env)
129     print(f"limiter = {limiter} (DUST_LIMITER={_dust_limiter_env!r})")
130
131     if use_coala:
132         vfrag_threshold = 1000000 * codeu.get("m") / codeu.get("s")
133         Q = 5
134         rhodust_eps = 1e-23 * codeu_kg_m3
135         K0_multiplier = 1
136
137     # Integrator parameters
138     C_cour = 0.1
139     C_force = 0.1
140
141     sim_folder = (
142         f"_to_trash/circular_dustydisc_{ndust}_{Npart}_{kernel}_{limiter.value}_coala_{use_coala}/"
143     )
144
145     dump_folder = sim_folder + "dump/"
146     analysis_folder = sim_folder + "analysis/"
147     plot_folder = analysis_folder + "plots/"
148
149     dump_prefix = dump_folder + "dump_"
150
151     # %%
152     # Create the dump directory if it does not exist
153     if shamrock.sys.world_rank() == 0:
154         os.makedirs(sim_folder, exist_ok=True)
155         os.makedirs(dump_folder, exist_ok=True)
156         os.makedirs(analysis_folder, exist_ok=True)
157         os.makedirs(plot_folder, exist_ok=True)
158
159     # %%
160     # Utility functions and quantities deduced from the base one
161
162     # Deduced quantities
163
164     bsize = disc.rout * 2
165     bmin = (-bsize, -bsize, -bsize)
166     bmax = (bsize, bsize, bsize)
167     profiles = disc.get_profiles()
168
169     grain_size = mrn_distribution.grain_size
170     grain_size_si = mrn_distribution.grain_size_si
171     rho_grains = mrn_distribution.rho_grains
172     massgrid_edges = mrn_distribution.massgrid_edges
173     mrn_weight = mrn_distribution.mrn_weight
174
175     if use_coala:
176         K0 = np.pi * ((4.0 / 3.0) * np.pi * rho_grains[0]) ** (-2.0 / 3.0)
177         K0 *= K0_multiplier
178         print(f"K0 = {K0}")
179
180         tabflux_coag = coala.coala_precalc_tabflux_coag(K0, ndust, Q, massgrid_edges)
181
182     # %%
183     # Start the context
184     # The context holds the data of the code
185     # We then init the layout of the field (e.g. the list of fields used by the solver)
186
187     ctx = shamrock.Context()
188     ctx.pdata_layout_new()
189
190     # %%
191     # Attach a SPH model to the context
192
193     model = shamrock.get_Model_SPH(context=ctx, vector_type="f64_3", sph_kernel=kernel)
194
195     cfg = model.gen_default_config()
196
197     cfg.set_artif_viscosity_VaryingCD10(
198         alpha_min=alpha_min,
199         alpha_max=alpha_max,
200         sigma_decay=sigma_decay,
201         alpha_u=alpha_u,
202         beta_AV=beta_AV,
203     )
204
205     cfg.set_eos_locally_isothermalLP07(cs0=disc.cs0(), q=disc.q, r0=disc.r0)
206
207     if ndust > 0:
208         cfg.set_dust_drag_epstein(gamma, grain_size, rho_grains)
209
210         if limiter == DustLimiter.NONE:
211             cfg.set_dust_mode_monofluid_tva(
212                 nvar=ndust, ensure_s_j_positivity=False, smooth_s_positivity_limiter=False
213             )
214         elif limiter == DustLimiter.SMOOTH:
215             cfg.set_dust_mode_monofluid_tva(
216                 nvar=ndust, ensure_s_j_positivity=False, smooth_s_positivity_limiter=True
217             )
218         elif limiter == DustLimiter.BALLABIO:
219             cfg.set_dust_mode_monofluid_tva(
220                 nvar=ndust, ensure_s_j_positivity=False, smooth_s_positivity_limiter=False
221             )
222             cfg.set_dust_ballabio_ts_limiter(True)
223         elif limiter == DustLimiter.HARD:
224             cfg.set_dust_mode_monofluid_tva(
225                 nvar=ndust, ensure_s_j_positivity=True, smooth_s_positivity_limiter=False
226             )
227
228         if use_coala:
229             cfg.set_dust_evol_coala_coag(rhodust_eps, vfrag_threshold, massgrid_edges, tabflux_coag)
230
231     cfg.add_kill_sphere(center=(0, 0, 0), radius=bsize)  # kill particles outside the simulation box
232
233     cfg.set_units(codeu)
234     cfg.set_particle_mass(disc.part_mass(Npart))
235
236     print(C_cour, C_force)
237
238     # Set the CFL
239     cfg.set_cfl_cour(C_cour)
240     cfg.set_cfl_force(C_force)
241     cfg.set_show_cfl_detail(True)
242
243     cfg.set_smoothing_length_density_based_neigh_lim(500)
244
245     # On a chaotic disc, we disable to two stage search to avoid giant leaves
246     cfg.set_tree_reduction_level(6)
247     cfg.set_two_stage_search(False)
248
249     # Set the solver config to be the one stored in cfg
250     model.set_solver_config(cfg)
251
252     # Print the solver config
253     model.get_current_config().print_status()
254
255     # Init the scheduler & fields
256     model.init_scheduler(scheduler_split_val, scheduler_merge_val)
257
258     # Set the simulation box size
259     model.resize_simulation_box(bmin, bmax)
260
261     # Create the setup
262
263     setup = model.get_setup()
264     gen_disc = disc.make_generator(setup, Npart, random_seed=666)
265
266     # Print the dot graph of the setup
267     if shamrock.sys.world_rank() == 0:
268         print(gen_disc.get_dot())
269
270     # Apply the setup
271     setup.apply_setup(gen_disc)
272
273     # correct the momentum and barycenter of the disc to 0
274     analysis_momentum = shamrock.model_sph.analysisTotalMomentum(model=model)
275     total_momentum = analysis_momentum.get_total_momentum()
276
277     if shamrock.sys.world_rank() == 0:
278         print(f"disc momentum = {total_momentum}")
279
280     model.apply_momentum_offset((-total_momentum[0], -total_momentum[1], -total_momentum[2]))
281
282     # Correct the barycenter before adding the sink
283     analysis_barycenter = shamrock.model_sph.analysisBarycenter(model=model)
284     barycenter, disc_mass = analysis_barycenter.get_barycenter()
285
286     if shamrock.sys.world_rank() == 0:
287         print(f"disc barycenter = {barycenter}")
288
289     model.apply_position_offset((-barycenter[0], -barycenter[1], -barycenter[2]))
290
291     total_momentum = shamrock.model_sph.analysisTotalMomentum(model=model).get_total_momentum()
292
293     if shamrock.sys.world_rank() == 0:
294         print(f"disc momentum after correction = {total_momentum}")
295
296     barycenter, disc_mass = shamrock.model_sph.analysisBarycenter(model=model).get_barycenter()
297
298     if shamrock.sys.world_rank() == 0:
299         print(f"disc barycenter after correction = {barycenter}")
300
301     if not np.allclose(total_momentum, 0.0):
302         raise RuntimeError("disc momentum is not 0")
303     if not np.allclose(barycenter, 0.0):
304         raise RuntimeError("disc barycenter is not 0")
305
306     # now that the barycenter & momentum are 0, we can add the sink
307     model.add_sink(center_mass, (0, 0, 0), (0, 0, 0), center_racc)
308
309     # Run a single step to init the integrator and smoothing length of the particles
310     # Here the htolerance is the maximum factor of evolution of the smoothing length in each
311     # Smoothing length iterations, increasing it affect the performance negatively but increase the
312     # convergence rate of the smoothing length
313     # this is why we increase it temporely to 1.3 before lowering it back to 1.1 (default value)
314     # Note that both ``change_htolerances`` can be removed and it will work the same but would converge
315     # more slowly at the first timestep
316
317     model.change_htolerances(coarse=1.3, fine=1.1)
318     model.timestep()
319     model.change_htolerances(coarse=1.1, fine=1.1)
320
321     def compute_sj_new_j(patchdata, j):
322         pmass = model.get_particle_mass()
323
324         hpart = patchdata["hpart"]
325         rho = pmass * (model.get_hfact() / np.array(hpart)) ** 3
326
327         epsilon_target = epsilon_base * mrn_weight[j]
328         print(f"epsilon_target = {epsilon_target} {j}")
329         s = np.sqrt(rho * epsilon_target)
330
331         print(
332             f"s = {s} {np.isnan(s).any()} epsilon_target = {epsilon_target} mrn_weight = {mrn_weight[j]}, rho = {rho}"
333         )
334
335         return s
336
337         # Add the dust
338
339     for k in range(ndust):
340
341         def compute_sj_new(patchdata):
342             return compute_sj_new_j(patchdata, k)
343
344         model.overwrite_field_value_f64("s_j", compute_sj_new, k)
345
346     model.set_dt(0.0)  # to help the corrector on next step after adding dust
347
348     # %%
349     # Evolve the simulation
350     model.solver_logs_reset_cumulated_step_time()
351     model.solver_logs_reset_step_count()
352
353     res_rates = []
354     for i in range(10):
355         model.timestep()
356         model.set_dt(0.0)
357         res_rates.append(model.solver_logs_last_rate())
358
359     c2 = c.copy()
360     c2["rate"] = np.max(res_rates)
361
362     print(c2)

Estimated memory usage: 0 MB

Gallery generated by Sphinx-Gallery