Note
Go to the end to download the full example code.
Testing the external forces coupling modes in RAMSES solver#
CI test for the two time quadratures of the external forces source term
(cfg.set_ext_force_coupling), with a constant gravitational acceleration:
free fall of a uniform periodic box (gas + one dust species): the momentum must follow \(\rho \mathbf{g} t\) exactly in both modes. The trapezoidal rule must conserve the internal energy, while the midpoint rule has the known \(-\rho g^2 \Delta t^2 / 2\) drift per step (see the TODO in
ApplyExtForceGas.cpp). This is checked with and without the time interpolation of the face states.one step on a non uniform moving state: for a constant acceleration both modes give the same momentum update, and their energy differs by \(\Delta t^2 \rho^{n+1/2} g^2 / 2\).
17 import numpy as np
18
19 import shamrock
20
21 # If we use the shamrock executable to run this script instead of the python interpreter,
22 # we should not initialize the system as the shamrock executable needs to handle specific MPI logic
23 if not shamrock.sys.is_initialized():
24 shamrock.change_loglevel(1)
25 shamrock.sys.init("0:0")
26
27 from shamrock.model_ramses import ExtForceCoupling
28
29 gamma = 1.4
30 g = np.array([0.0, -1.0, 0.0])
31 g_norm2 = float(np.dot(g, g))
32
33 sz = 1 << 1
34 base = 4
35
36
37 def zero_vel(rmin, rmax):
38 return (0.0, 0.0, 0.0)
39
40
41 def run(
42 mode,
43 rho_map,
44 rho_dust_map,
45 dt,
46 nsteps,
47 vel_map=zero_vel,
48 vel_dust_map=zero_vel,
49 face_time_interpolation=True,
50 ):
51 ctx = shamrock.Context()
52 ctx.pdata_layout_new()
53
54 model = shamrock.get_Model_Ramses(context=ctx, vector_type="f64_3", grid_repr="i64_3")
55
56 cfg = model.gen_default_config()
57 cfg.set_scale_factor(1 / (sz * base))
58 cfg.set_eos_gamma(gamma)
59 cfg.set_dust_mode_dhll(1)
60 cfg.set_face_time_interpolation(face_time_interpolation)
61 cfg.add_ext_force_constant_gravity(g=tuple(g))
62 cfg.set_ext_force_coupling(mode)
63 model.set_solver_config(cfg)
64
65 model.init_scheduler(int(1e7), 1)
66 model.make_base_grid((0, 0, 0), (sz, sz, sz), (base, base, base))
67
68 model.set_field_value_lambda_f64("rho", rho_map)
69
70 def rhovel_map(rmin, rmax):
71 return tuple(rho_map(rmin, rmax) * np.array(vel_map(rmin, rmax)))
72
73 def rhoetot_map(rmin, rmax):
74 v = np.array(vel_map(rmin, rmax))
75 return 0.5 * rho_map(rmin, rmax) * float(np.dot(v, v)) + 1.0 / (gamma - 1.0)
76
77 def rhovel_dust_map(rmin, rmax):
78 return tuple(rho_dust_map(rmin, rmax) * np.array(vel_dust_map(rmin, rmax)))
79
80 model.set_field_value_lambda_f64_3("rhovel", rhovel_map)
81 model.set_field_value_lambda_f64("rhoetot", rhoetot_map)
82 model.set_field_value_lambda_f64("rho_dust", rho_dust_map, 0)
83 model.set_field_value_lambda_f64_3("rhovel_dust", rhovel_dust_map, 0)
84
85 data_init = ctx.collect_data()
86
87 for i in range(nsteps):
88 model.evolve_once_override_time(dt * float(i), dt)
89
90 return data_init, ctx.collect_data()
91
92
93 def internal_energy(data):
94 rho = np.array(data["rho"])
95 rhov = np.array(data["rhovel"])
96 return np.array(data["rhoetot"]) - 0.5 * np.sum(rhov * rhov, axis=1) / rho
97
98
99 def max_rel_err(a, b, scale):
100 return float(np.max(np.abs(np.array(a) - np.array(b)))) / scale
101
102
103 def test_free_fall():
104 rho0, rho_d0 = 1.5, 0.7
105 dt, nsteps = 0.01, 20
106 t = dt * nsteps
107
108 cases = [
109 (mode, face_interp)
110 for mode in [ExtForceCoupling.Midpoint, ExtForceCoupling.Trapezoidal]
111 for face_interp in [True, False]
112 ]
113 for mode, face_interp in cases:
114 data_init, data = run(
115 mode,
116 lambda rmin, rmax: rho0,
117 lambda rmin, rmax: rho_d0,
118 dt,
119 nsteps,
120 face_time_interpolation=face_interp,
121 )
122
123 # the collected data is only complete on rank 0
124 if shamrock.sys.world_rank() != 0:
125 continue
126
127 # no pressure gradient: the gas and the dust are in free fall
128 err_rhov = max_rel_err(data["rhovel"], [rho0 * g * t], rho0 * t)
129 err_rhov_d = max_rel_err(data["rhovel_dust"], [rho_d0 * g * t], rho_d0 * t)
130
131 de_int = internal_energy(data) - internal_energy(data_init)
132 if mode == ExtForceCoupling.Trapezoidal:
133 de_int_expected = 0.0
134 else:
135 de_int_expected = -nsteps * rho0 * g_norm2 * dt**2 / 2
136 err_e_int = float(np.max(np.abs(de_int - de_int_expected))) / (1.0 / (gamma - 1.0))
137
138 print(
139 f"free fall {mode}, face interp {face_interp}: err rhovel = {err_rhov:.3e}, err rhovel_dust = {err_rhov_d:.3e}, "
140 f"err e_int = {err_e_int:.3e}"
141 )
142
143 assert err_rhov < 1e-13, (
144 f"{mode}, face interp {face_interp}: gas momentum is not rho g t ({err_rhov})"
145 )
146 assert err_rhov_d < 1e-13, (
147 f"{mode}, face interp {face_interp}: dust momentum is not rho_d g t ({err_rhov_d})"
148 )
149 assert err_e_int < 1e-12, (
150 f"{mode}, face interp {face_interp}: unexpected internal energy change ({err_e_int})"
151 )
152
153
154 def test_modes_one_step():
155 def rho_map(rmin, rmax):
156 x, y, z = (np.array(rmin) + np.array(rmax)) / 2
157 return 1.0 + 0.3 * np.sin(2 * np.pi * x) * np.cos(2 * np.pi * y)
158
159 def rho_dust_map(rmin, rmax):
160 x, y, z = (np.array(rmin) + np.array(rmax)) / 2
161 return 0.5 + 0.2 * np.cos(2 * np.pi * z)
162
163 # a non uniform velocity, so that the flux derivatives (dt rho, dt rhov) are not zero
164 def vel_map(rmin, rmax):
165 x, y, z = (np.array(rmin) + np.array(rmax)) / 2
166 return (0.3 * np.sin(2 * np.pi * y), 0.2 * np.cos(2 * np.pi * x), 0.1)
167
168 def vel_dust_map(rmin, rmax):
169 x, y, z = (np.array(rmin) + np.array(rmax)) / 2
170 return (0.1, 0.2 * np.sin(2 * np.pi * z), 0.3 * np.cos(2 * np.pi * x))
171
172 dt = 0.01
173
174 args = (rho_map, rho_dust_map, dt, 1, vel_map, vel_dust_map)
175 data_init, data_mid = run(ExtForceCoupling.Midpoint, *args)
176 _, data_trap = run(ExtForceCoupling.Trapezoidal, *args)
177
178 # the collected data is only complete on rank 0
179 if shamrock.sys.world_rank() != 0:
180 return
181
182 # check that the setup does exercise the flux derivatives
183 assert not np.array_equal(np.array(data_init["rho"]), np.array(data_mid["rho"]))
184 assert not np.array_equal(np.array(data_init["rho_dust"]), np.array(data_mid["rho_dust"]))
185
186 # the density update does not depend on the external forces
187 assert np.array_equal(np.array(data_mid["rho"]), np.array(data_trap["rho"]))
188
189 rhov_scale = float(np.max(np.abs(np.array(data_mid["rhovel"]))))
190 rhov_d_scale = float(np.max(np.abs(np.array(data_mid["rhovel_dust"]))))
191 err_rhov = max_rel_err(data_mid["rhovel"], data_trap["rhovel"], rhov_scale)
192 err_rhov_d = max_rel_err(data_mid["rhovel_dust"], data_trap["rhovel_dust"], rhov_d_scale)
193
194 rho_half = 0.5 * (np.array(data_init["rho"]) + np.array(data_mid["rho"]))
195 de_expected = dt**2 * rho_half * g_norm2 / 2
196 de = np.array(data_trap["rhoetot"]) - np.array(data_mid["rhoetot"])
197 err_de = float(np.max(np.abs(de - de_expected) / de_expected))
198
199 print(
200 f"one step midpoint vs trapezoidal: err rhovel = {err_rhov:.3e}, "
201 f"err rhovel_dust = {err_rhov_d:.3e}, err energy difference = {err_de:.3e}"
202 )
203
204 assert err_rhov < 1e-13, f"modes give different gas momentum ({err_rhov})"
205 assert err_rhov_d < 1e-13, f"modes give different dust momentum ({err_rhov_d})"
206 assert err_de < 1e-8, f"unexpected energy difference between the modes ({err_de})"
207
208
209 test_free_fall()
210 test_modes_one_step()
211
212 if shamrock.sys.world_rank() == 0:
213 print("ramses_ext_force_coupling: OK")
Estimated memory usage: 0 MB