Note
Go to the end to download the full example code.
Testing the RAMSES solver self-gravity with the Jeans instability#
CI test for the self-gravity of the RAMSES solver, applied through the external forces with the
trapezoidal coupling (cfg.set_ext_force_coupling(ExtForceCoupling.Trapezoidal)).
A uniform periodic box with a small standing density perturbation \(\delta \rho = A \rho_0 \cos(k x)\) (gas at rest, isentropic) follows the linear dispersion relation \(\omega^2 = c_s^2 k^2 - 4 \pi G \rho_0\):
above the Jeans length (\(c_s k < \sqrt{4 \pi G \rho_0}\)) the perturbation grows as \(\cosh(\gamma t)\), with \(\gamma^2 = 4 \pi G \rho_0 - c_s^2 k^2\),
below it, it oscillates as \(\cos(\omega t)\), slower than a sound wave.
The growth rate (with both Poisson solvers) and the oscillation frequency are checked against the linear theory. The gravitational constant comes from the unit system, set so that \(G = 1\). The box is split into several patches, so that the ghost exchanges of the Poisson solvers are tested between patches (and between ranks when run with MPI).
22 import numpy as np
23
24 import shamrock
25
26 # If we use the shamrock executable to run this script instead of the python interpreter,
27 # we should not initialize the system as the shamrock executable needs to handle specific MPI logic
28 if not shamrock.sys.is_initialized():
29 shamrock.change_loglevel(1)
30 shamrock.sys.init("0:0")
31
32 from shamrock.model_ramses import HLLC_PVRS, ExtForceCoupling
33
34 shamrock.enable_experimental_features()
35
36 gamma = 5.0 / 3.0
37 rho0 = 1.0
38 amp = 1e-3
39 k = 2 * np.pi # one wavelength in the unit box
40
41 # G = 1 in the code units (unit_mass = 1 / G_SI kg)
42 G_SI = shamrock.Constants(shamrock.UnitSystem()).G()
43 units = shamrock.UnitSystem(unit_time=1.0, unit_length=1.0, unit_mass=1.0 / G_SI)
44 G = shamrock.Constants(units).G()
45 four_pi_G_rho0 = 4 * np.pi * G * rho0
46
47 sz = 1 << 1
48 base = 8 # 16 cells per wavelength
49 block_size = 8 # cells per AMR block
50
51
52 def run(cs, solver, dt, nsteps):
53 """Return the times and the amplitudes of the cos(k x) mode of the density"""
54
55 ctx = shamrock.Context()
56 ctx.pdata_layout_new()
57
58 model = shamrock.get_Model_Ramses(context=ctx, vector_type="f64_3", grid_repr="i64_3")
59
60 cfg = model.gen_default_config()
61 cfg.set_scale_factor(1 / (sz * base))
62 cfg.set_eos_gamma(gamma)
63 cfg.set_slope_lim_vanleer_sym()
64 cfg.set_riemann_solver(HLLC_PVRS)
65 cfg.set_face_time_interpolation(True)
66 cfg.set_units(units)
67 if solver == "cg":
68 cfg.set_gravity_mode_cg()
69 else:
70 cfg.set_gravity_mode_bicgstab()
71 # well above the round-off floor of the (singular, periodic) Poisson problem: when the
72 # perturbation is small (e.g. near a zero of cos(omega t)) a tolerance at that floor is not
73 # reached and the CG iterations diverge (see the TODO in NodeCGLoop)
74 cfg.set_self_gravity_tol(1e-8)
75 cfg.set_self_gravity_Niter_max(500)
76 cfg.set_ext_force_coupling(ExtForceCoupling.Trapezoidal)
77 model.set_solver_config(cfg)
78
79 # split the 512 blocks into several patches
80 model.init_scheduler(64, 1)
81 model.make_base_grid((0, 0, 0), (sz, sz, sz), (base, base, base))
82
83 P0 = cs * cs * rho0 / gamma
84
85 def rho_map(rmin, rmax):
86 x = 0.5 * (rmin[0] + rmax[0])
87 return rho0 * (1.0 + amp * np.cos(k * x))
88
89 def rhoetot_map(rmin, rmax):
90 # isentropic perturbation
91 return P0 * (rho_map(rmin, rmax) / rho0) ** gamma / (gamma - 1.0)
92
93 model.set_field_value_lambda_f64("rho", rho_map)
94 model.set_field_value_lambda_f64_3("rhovel", lambda rmin, rmax: (0.0, 0.0, 0.0))
95 model.set_field_value_lambda_f64("rhoetot", rhoetot_map)
96 model.set_field_value_lambda_f64("phi", lambda rmin, rmax: 0.0)
97 model.set_field_value_lambda_f64("phi_old", lambda rmin, rmax: 0.0)
98
99 def mode_amplitude():
100 data = ctx.collect_data()
101 if shamrock.sys.world_rank() != 0:
102 return 0.0, 0.0
103 rho = np.array(data["rho"])
104 x = np.array(
105 [
106 0.5 * (rmin[0] + rmax[0])
107 for cmin, cmax in zip(data["cell_min"], data["cell_max"])
108 for rmin, rmax in (
109 model.get_cell_coords((tuple(map(int, cmin)), tuple(map(int, cmax))), lid)
110 for lid in range(block_size)
111 )
112 ]
113 )
114
115 a = 2 * np.mean((rho - rho0) * np.cos(k * x)) / (rho0 * amp)
116 p_tot = np.sum(np.array(data["rhovel"]), axis=0)
117 return a, float(np.max(np.abs(p_tot)))
118
119 times = [0.0]
120 a0, _ = mode_amplitude()
121 amps = [a0]
122 p_max = 0.0
123 for i in range(nsteps):
124 model.evolve_once_override_time(dt * float(i), dt)
125 a, p = mode_amplitude()
126 times.append(dt * (i + 1))
127 amps.append(a)
128 p_max = max(p_max, p)
129
130 # every rank must own patches, otherwise the exchanges between ranks are not tested
131 owners = [p.node_owner_id for p in ctx.get_patch_list_global()]
132 patches_per_rank = np.bincount(owners, minlength=shamrock.sys.world_size())
133 if shamrock.sys.world_rank() == 0:
134 print(f"patches per rank: {patches_per_rank.tolist()}")
135 assert len(owners) > 1, "the box is not split into several patches"
136 assert np.all(patches_per_rank > 0), f"a rank owns no patch ({patches_per_rank.tolist()})"
137
138 return np.array(times), np.array(amps), p_max
139
140
141 def test_unstable(solver):
142 cs = 0.3
143 growth_th = np.sqrt(four_pi_G_rho0 - cs * cs * k * k)
144
145 dt, nsteps = 0.01, 50
146 t, a, p_max = run(cs, solver, dt, nsteps)
147
148 if shamrock.sys.world_rank() != 0:
149 return
150
151 # measured growth rate from a(t) = a(0) cosh(growth t) at the end of the run
152 growth = np.arccosh(a[-1] / a[0]) / t[-1]
153 rel_err = abs(growth - growth_th) / growth_th
154
155 print(
156 f"Jeans unstable ({solver}): a(t_end)/a(0) = {a[-1] / a[0]:.6f} "
157 f"(theory {np.cosh(growth_th * t[-1]):.6f}), growth rate = {growth:.6f} "
158 f"(theory {growth_th:.6f}), rel err = {rel_err:.3e}, max |total momentum| = {p_max:.3e}"
159 )
160
161 assert abs(a[0] - 1.0) < 1e-2, f"bad initial mode amplitude ({a[0]})"
162 # the first step only sees the pressure in the density update (the face predictor has no
163 # gravity source term, as in the reference scheme), so the growth starts at the second step
164 assert np.all(np.diff(a[1:]) > 0), f"{solver}: the perturbation does not grow monotonically"
165 assert rel_err < 0.05, f"{solver}: growth rate {growth} too far from theory {growth_th}"
166 assert p_max < 1e-12, f"{solver}: the total momentum is not conserved ({p_max})"
167
168 return growth
169
170
171 def test_stable():
172 cs = 1.0
173 omega_th = np.sqrt(cs * cs * k * k - four_pi_G_rho0)
174 omega_sound = cs * k
175
176 dt, nsteps = 0.01, 40
177 t, a, p_max = run(cs, "cg", dt, nsteps)
178
179 if shamrock.sys.world_rank() != 0:
180 return
181
182 # first zero of a(t) = cos(omega t), at t = pi / (2 omega)
183 i = int(np.argmax(a < 0))
184 assert i > 0, "the perturbation does not oscillate"
185 t_zero = t[i - 1] + (t[i] - t[i - 1]) * a[i - 1] / (a[i - 1] - a[i])
186 omega = np.pi / (2 * t_zero)
187 rel_err = abs(omega - omega_th) / omega_th
188
189 print(
190 f"Jeans stable: omega = {omega:.6f} (theory {omega_th:.6f}, sound wave {omega_sound:.6f}), "
191 f"rel err = {rel_err:.3e}, max |total momentum| = {p_max:.3e}"
192 )
193
194 assert rel_err < 0.05, f"oscillation frequency {omega} too far from theory {omega_th}"
195 assert p_max < 1e-12, f"the total momentum is not conserved ({p_max})"
196
197
198 growth_cg = test_unstable("cg")
199 growth_bicgstab = test_unstable("bicgstab")
200 test_stable()
201
202 if shamrock.sys.world_rank() == 0:
203 # both solvers converge to the same potential (tol = 1e-10)
204 rel_diff = abs(growth_cg - growth_bicgstab) / growth_cg
205 print(f"growth rate CG vs BiCGSTAB: rel diff = {rel_diff:.3e}")
206 assert rel_diff < 1e-6, f"CG and BiCGSTAB give different growth rates ({rel_diff})"
207
208 print("ramses_self_gravity_jeans: OK")
Estimated memory usage: 0 MB