Snapshot from this documentation build. SHA-256: 3f2bc81de967d1a213d0da8ab39f4ee5e3f680147498a054bca81e05cbf6285d
1 # Time-center a gas density source without retaining the previous full grid.
2 # From continuity, rho(t-dt) = rho(t) + dt div(m)(t) + O(dt^2).
3
4 export predict_density_backward!, predict_density_perturbation_backward!,
5 apply_cell_center_gravity!
6 export initialize_lattice_displacements!
7 export deposit_lattice_displacements!, deposit_lattice_displacement_delta!
8 export drift_lattice_displacements!
9 export deposit_particle_velocity!
10
11 @kernel function _deposit_particle_velocity_k!(density,momentum,
12 @Const(px),@Const(py),@Const(pz),@Const(v),N::Int,
13 lattice::Val{LATTICE}) where {LATTICE}
14 p=@index(Global)
15 @inbounds begin
16 T=eltype(density)
17 q=p-1
18 i=LATTICE ? q%N : 0
19 j=LATTICE ? (q÷N)%N : 0
20 k=LATTICE ? q÷(N*N) : 0
21 shift=LATTICE ? zero(T) : -T(0.5)
22 x=T(N)*T(px[p])+shift
23 y=T(N)*T(py[p])+shift
24 z=T(N)*T(pz[p])+shift
25 ix=unsafe_trunc(Int,floor(x)); iy=unsafe_trunc(Int,floor(y))
26 iz=unsafe_trunc(Int,floor(z))
27 fx=x-T(ix); fy=y-T(iy); fz=z-T(iz)
28 for dz in 0:1,dy in 0:1,dx in 0:1
29 w=(dx==0 ? one(T)-fx : fx)*
30 (dy==0 ? one(T)-fy : fy)*(dz==0 ? one(T)-fz : fz)
31 c=mod(i+ix+dx,N)+N*mod(j+iy+dy,N)+N*N*mod(k+iz+dz,N)+1
32 KA.@atomic density[c] += w
33 KA.@atomic momentum[c] += w*T(v[p])
34 end
35 end
36 end
37
38 @kernel function _normalize_particle_velocity_k!(velocity,@Const(density))
39 c=@index(Global)
40 @inbounds velocity[c]=density[c]>0 ? velocity[c]/density[c] : zero(eltype(velocity))
41 end
42
43 """
44 deposit_particle_velocity!(velocity, density, px, py, pz, v; N,
45 lattice_displacements=false)
46
47 CIC-deposit particle mass and one momentum component at current positions,
48 then form the mass-weighted velocity within each Eulerian cell. Empty cells
49 have velocity zero and remain identifiable by `density==0`. Displacements
50 are relative to the implicit cell-centred lattice when requested. This is a
51 mesh velocity estimator, not a particle-label Fourier transform or an
52 unambiguous single-stream velocity after shell crossing.
53 """
54 function deposit_particle_velocity!(velocity,density,px,py,pz,v;
55 N::Integer,lattice_displacements::Bool=false)
56 N>=2 || throw(ArgumentError("particle velocity grid requires N>=2"))
57 length(velocity)==length(density)==N^3 || throw(DimensionMismatch("velocity grid size"))
58 all(length(a)==length(v) for a in (px,py,pz)) ||
59 throw(DimensionMismatch("particle array sizes"))
60 lattice_displacements && length(v)!=N^3 &&
61 throw(DimensionMismatch("implicit lattice needs N^3 particles"))
62 velocity===density && throw(ArgumentError("velocity and density must not alias"))
63 fill!(density,zero(eltype(density))); fill!(velocity,zero(eltype(velocity)))
64 be=KA.get_backend(density)
65 _deposit_particle_velocity_k!(be)(density,velocity,px,py,pz,v,Int(N),
66 Val(lattice_displacements);ndrange=length(v))
67 KA.synchronize(be)
68 _normalize_particle_velocity_k!(be)(velocity,density;ndrange=length(density))
69 KA.synchronize(be)
70 velocity
71 end
72
73 @kernel function _initialize_lattice_displacements_k!(dxp, dyp, dzp, vx, vy, vz)
74 p = @index(Global)
75 @inbounds begin
76 dxp[p] = zero(eltype(dxp)); dyp[p] = zero(eltype(dyp)); dzp[p] = zero(eltype(dzp))
77 vx[p] = zero(eltype(vx)); vy[p] = zero(eltype(vy)); vz[p] = zero(eltype(vz))
78 end
79 end
80
81 """
82 initialize_lattice_displacements!(dxp, dyp, dzp, vx, vy, vz)
83
84 Initialize particles represented by displacements from an implicit regular
85 lattice. Particle `p` has an unperturbed centre reconstructed from `p` and the
86 grid size, so no absolute-position arrays are required.
87 """
88 function initialize_lattice_displacements!(dxp, dyp, dzp, vx, vy, vz)
89 n = length(dxp)
90 all(length(a) == n for a in (dyp, dzp, vx, vy, vz)) ||
91 throw(DimensionMismatch("particle arrays must have equal lengths"))
92 be = KA.get_backend(dxp)
93 _initialize_lattice_displacements_k!(be)(dxp, dyp, dzp, vx, vy, vz; ndrange=n)
94 return dxp, dyp, dzp, vx, vy, vz
95 end
96
97 @kernel function _deposit_lattice_displacements_k!(rho,
98 @Const(dxp), @Const(dyp), @Const(dzp),
99 @Const(vx), @Const(vy), @Const(vz),
100 N::Int, disp, shift)
101 p = @index(Global)
102 @inbounds begin
103 T = eltype(rho)
104 q = p - 1
105 i = q % N
106 j = (q ÷ N) % N
107 k = q ÷ (N * N)
108 # Resolve the fractional coordinate before adding the integer cell.
109 # Adding N*dxp to i rounds away early-time drifts for large i in f32.
110 base = T(0.5) + shift
111 ux = base + T(N) * (T(dxp[p]) + disp * T(vx[p]))
112 uy = base + T(N) * (T(dyp[p]) + disp * T(vy[p]))
113 uz = base + T(N) * (T(dzp[p]) + disp * T(vz[p]))
114 ox = unsafe_trunc(Int, floor(ux))
115 oy = unsafe_trunc(Int, floor(uy))
116 oz = unsafe_trunc(Int, floor(uz))
117 # Compute both weights from the local coordinate. The small member of
118 # each pair then remains representable whether the drift is positive
119 # or negative.
120 wxa = T(ox + 1) - ux; wxb = ux - T(ox)
121 wya = T(oy + 1) - uy; wyb = uy - T(oy)
122 wza = T(oz + 1) - uz; wzb = uz - T(oz)
123 ia = mod(i + ox, N); ib = mod(i + ox + 1, N)
124 ja = mod(j + oy, N); jb = mod(j + oy + 1, N)
125 ka = mod(k + oz, N); kb = mod(k + oz + 1, N)
126 Nj = N; Nk = N * N
127 KA.@atomic rho[ia + Nj*ja + Nk*ka + 1] += wxa*wya*wza
128 KA.@atomic rho[ib + Nj*ja + Nk*ka + 1] += wxb*wya*wza
129 KA.@atomic rho[ia + Nj*jb + Nk*ka + 1] += wxa*wyb*wza
130 KA.@atomic rho[ib + Nj*jb + Nk*ka + 1] += wxb*wyb*wza
131 KA.@atomic rho[ia + Nj*ja + Nk*kb + 1] += wxa*wya*wzb
132 KA.@atomic rho[ib + Nj*ja + Nk*kb + 1] += wxb*wya*wzb
133 KA.@atomic rho[ia + Nj*jb + Nk*kb + 1] += wxa*wyb*wzb
134 KA.@atomic rho[ib + Nj*jb + Nk*kb + 1] += wxb*wyb*wzb
135 end
136 end
137
138 @inline function _unit_product_delta(lx, ly, lz)
139 # (1-lx)*(1-ly)*(1-lz) - 1, without subtracting two values near one.
140 return -(lx + ly + lz) + lx*ly + lx*lz + ly*lz - lx*ly*lz
141 end
142
143 @kernel function _deposit_lattice_displacement_delta_k!(delta,
144 @Const(dxp), @Const(dyp),
145 @Const(dzp), @Const(vx),
146 @Const(vy), @Const(vz),
147 N::Int, disp, shift)
148 p = @index(Global)
149 @inbounds begin
150 T = eltype(delta)
151 q = p - 1
152 i = q % N
153 j = (q ÷ N) % N
154 k = q ÷ (N * N)
155 base = T(0.5) + shift
156 ux = base + T(N) * (T(dxp[p]) + disp * T(vx[p]))
157 uy = base + T(N) * (T(dyp[p]) + disp * T(vy[p]))
158 uz = base + T(N) * (T(dzp[p]) + disp * T(vz[p]))
159 ox = unsafe_trunc(Int, floor(ux))
160 oy = unsafe_trunc(Int, floor(uy))
161 oz = unsafe_trunc(Int, floor(uz))
162 wxa = T(ox + 1) - ux; wxb = ux - T(ox)
163 wya = T(oy + 1) - uy; wyb = uy - T(oy)
164 wza = T(oz + 1) - uz; wzb = uz - T(oz)
165 ia = mod(i + ox, N); ib = mod(i + ox + 1, N)
166 ja = mod(j + oy, N); jb = mod(j + oy + 1, N)
167 ka = mod(k + oz, N); kb = mod(k + oz + 1, N)
168 Nj = N; Nk = N * N
169
170 h000 = ia == i && ja == j && ka == k
171 h100 = ib == i && ja == j && ka == k
172 h010 = ia == i && jb == j && ka == k
173 h110 = ib == i && jb == j && ka == k
174 h001 = ia == i && ja == j && kb == k
175 h101 = ib == i && ja == j && kb == k
176 h011 = ia == i && jb == j && kb == k
177 h111 = ib == i && jb == j && kb == k
178
179 v000 = h000 ? _unit_product_delta(wxb, wyb, wzb) : wxa*wya*wza
180 v100 = h100 ? _unit_product_delta(wxa, wyb, wzb) : wxb*wya*wza
181 v010 = h010 ? _unit_product_delta(wxb, wya, wzb) : wxa*wyb*wza
182 v110 = h110 ? _unit_product_delta(wxa, wya, wzb) : wxb*wyb*wza
183 v001 = h001 ? _unit_product_delta(wxb, wyb, wza) : wxa*wya*wzb
184 v101 = h101 ? _unit_product_delta(wxa, wyb, wza) : wxb*wya*wzb
185 v011 = h011 ? _unit_product_delta(wxb, wya, wza) : wxa*wyb*wzb
186 v111 = h111 ? _unit_product_delta(wxa, wya, wza) : wxb*wyb*wzb
187
188 KA.@atomic delta[ia + Nj*ja + Nk*ka + 1] += v000
189 KA.@atomic delta[ib + Nj*ja + Nk*ka + 1] += v100
190 KA.@atomic delta[ia + Nj*jb + Nk*ka + 1] += v010
191 KA.@atomic delta[ib + Nj*jb + Nk*ka + 1] += v110
192 KA.@atomic delta[ia + Nj*ja + Nk*kb + 1] += v001
193 KA.@atomic delta[ib + Nj*ja + Nk*kb + 1] += v101
194 KA.@atomic delta[ia + Nj*jb + Nk*kb + 1] += v011
195 KA.@atomic delta[ib + Nj*jb + Nk*kb + 1] += v111
196 if !(h000 || h100 || h010 || h110 || h001 || h101 || h011 || h111)
197 KA.@atomic delta[i + Nj*j + Nk*k + 1] -= one(T)
198 end
199 end
200 end
201
202 """
203 deposit_lattice_displacements!(rho, dxp, dyp, dzp, vx, vy, vz; N, disp=0, shift=-0.5)
204
205 CIC-deposit one unit-mass particle per cell from displacement coordinates. The
206 lattice centre and displacement are combined in cell coordinates, preserving
207 sub-ULP box displacements in Float32. `rho` is cleared before deposit.
208 """
209 function deposit_lattice_displacements!(rho::AbstractArray{T}, dxp, dyp, dzp,
210 vx, vy, vz; N::Integer,
211 disp::Real=0, shift::Real=-0.5) where {T}
212 length(rho) == N^3 || throw(DimensionMismatch("rho must contain N^3 cells"))
213 length(dxp) == N^3 || throw(DimensionMismatch("lattice must contain N^3 particles"))
214 fill!(rho, zero(T))
215 be = KA.get_backend(rho)
216 _deposit_lattice_displacements_k!(be)(rho, dxp, dyp, dzp, vx, vy, vz,
217 Int(N), T(disp), T(shift);
218 ndrange=length(dxp))
219 return rho
220 end
221
222 """
223 deposit_lattice_displacement_delta!(delta, dxp, dyp, dzp, vx, vy, vz;
224 N, disp=0, shift=-0.5)
225
226 Deposit the density contrast of one particle per implicit lattice cell. The
227 kernel accumulates the displaced CIC weights minus the undisplaced lattice
228 weight directly, so perturbations much smaller than `eps(Float32)` are not
229 lost by first constructing a density near one. No additional grid is needed.
230 """
231 function deposit_lattice_displacement_delta!(delta::AbstractArray{T}, dxp, dyp,
232 dzp, vx, vy, vz; N::Integer,
233 disp::Real=0,
234 shift::Real=-0.5) where {T}
235 N >= 2 || throw(ArgumentError("lattice density-contrast deposit requires N >= 2"))
236 length(delta) == N^3 ||
237 throw(DimensionMismatch("delta must contain N^3 cells"))
238 length(dxp) == N^3 ||
239 throw(DimensionMismatch("lattice must contain N^3 particles"))
240 fill!(delta, zero(T))
241 be = KA.get_backend(delta)
242 _deposit_lattice_displacement_delta_k!(be)(
243 delta, dxp, dyp, dzp, vx, vy, vz, Int(N), T(disp), T(shift);
244 ndrange=length(dxp),
245 )
246 return delta
247 end
248
249 @kernel function _drift_lattice_displacements_k!(dxp, dyp, dzp,
250 @Const(vx), @Const(vy), @Const(vz),
251 coef, dowrap::Int)
252 p = @index(Global)
253 @inbounds begin
254 x = dxp[p] + coef * vx[p]
255 y = dyp[p] + coef * vy[p]
256 z = dzp[p] + coef * vz[p]
257 if dowrap == 1
258 x -= floor(x + oftype(x, 0.5))
259 y -= floor(y + oftype(y, 0.5))
260 z -= floor(z + oftype(z, 0.5))
261 end
262 dxp[p] = x; dyp[p] = y; dzp[p] = z
263 end
264 end
265
266 "Accumulate a drift into zero-centred periodic lattice displacements."
267 function drift_lattice_displacements!(dxp::AbstractVector{T}, dyp, dzp, vx, vy, vz;
268 coef::Real, wrap::Bool=true) where {T}
269 be = KA.get_backend(dxp)
270 _drift_lattice_displacements_k!(be)(dxp, dyp, dzp, vx, vy, vz,
271 T(coef), wrap ? 1 : 0;
272 ndrange=length(dxp))
273 return dxp, dyp, dzp
274 end
275
276 @kernel function _predict_density_backward_k!(dst,
277 @Const(rho), @Const(mx),
278 @Const(my), @Const(mz),
279 dt_back, inv2dx,
280 nx::Int, ny::Int, nz::Int)
281 c = @index(Global)
282 @inbounds begin
283 i = (c - 1) % nx + 1
284 q = (c - 1) ÷ nx
285 j = q % ny + 1
286 k = q ÷ ny + 1
287 im = i == 1 ? nx : i - 1
288 ip = i == nx ? 1 : i + 1
289 jm = j == 1 ? ny : j - 1
290 jp = j == ny ? 1 : j + 1
291 km = k == 1 ? nz : k - 1
292 kp = k == nz ? 1 : k + 1
293
294 cim = ((k - 1) * ny + (j - 1)) * nx + im
295 cip = ((k - 1) * ny + (j - 1)) * nx + ip
296 cjm = ((k - 1) * ny + (jm - 1)) * nx + i
297 cjp = ((k - 1) * ny + (jp - 1)) * nx + i
298 ckm = ((km - 1) * ny + (j - 1)) * nx + i
299 ckp = ((kp - 1) * ny + (j - 1)) * nx + i
300 divm = ((mx[cip] - mx[cim]) +
301 (my[cjp] - my[cjm]) +
302 (mz[ckp] - mz[ckm])) * inv2dx
303 dst[c] = rho[c] + dt_back * divm
304 end
305 end
306
307 """
308 predict_density_backward!(dst, rho, mx, my, mz, dims; dx, dt_back)
309
310 Predict density at `t - dt_back` from conserved state at `t` using the periodic
311 continuity equation. This second-order midpoint source reconstruction avoids a
312 previous-density grid in operator-split particle gravity.
313 """
314 function predict_density_backward!(dst::AbstractArray{T}, rho, mx, my, mz,
315 dims::NTuple{3,Int}; dx::Real,
316 dt_back::Real) where {T}
317 length(dst) == prod(dims) || throw(DimensionMismatch("dst does not match dims"))
318 be = KA.get_backend(dst)
319 _predict_density_backward_k!(be)(dst, rho, mx, my, mz,
320 T(dt_back), T(0.5) / T(dx), dims...;
321 ndrange=length(dst))
322 return dst
323 end
324
325 function predict_density_backward!(dst::AbstractArray, s::MHDState,
326 dt_back::Real)
327 predict_density_backward!(dst, s.U[1], s.U[2], s.U[3], s.U[4], s.dims;
328 dx=s.dx, dt_back=dt_back)
329 end
330
331 """
332 predict_density_perturbation_backward!(dst, delta_rho, mx, my, mz, dims;
333 dx, dt_back)
334
335 Time-center an authoritative density perturbation without adding and subtracting
336 the homogeneous background in Float32.
337 """
338 function predict_density_perturbation_backward!(dst::AbstractArray{T},delta_rho,
339 mx,my,mz,dims::NTuple{3,Int};
340 dx::Real,dt_back::Real) where {T}
341 length(dst)==prod(dims) || throw(DimensionMismatch("dst does not match dims"))
342 be=KA.get_backend(dst)
343 _predict_density_backward_k!(be)(dst,delta_rho,mx,my,mz,
344 T(dt_back),T(0.5)/T(dx),dims...;ndrange=length(dst))
345 dst
346 end
347
348 @kernel function _apply_cell_center_gravity_k!(rho, mx, my, mz, E,
349 @Const(phi), energy_residual,
350 dt, inv2dx,
351 nx::Int, ny::Int, nz::Int,
352 track_energy::Val{TRACK}) where {TRACK}
353 c = @index(Global)
354 @inbounds begin
355 i = (c - 1) % nx + 1
356 q = (c - 1) ÷ nx
357 j = q % ny + 1
358 k = q ÷ ny + 1
359 im = i == 1 ? nx : i - 1
360 ip = i == nx ? 1 : i + 1
361 jm = j == 1 ? ny : j - 1
362 jp = j == ny ? 1 : j + 1
363 km = k == 1 ? nz : k - 1
364 kp = k == nz ? 1 : k + 1
365
366 cim = ((k - 1) * ny + (j - 1)) * nx + im
367 cip = ((k - 1) * ny + (j - 1)) * nx + ip
368 cjm = ((k - 1) * ny + (jm - 1)) * nx + i
369 cjp = ((k - 1) * ny + (jp - 1)) * nx + i
370 ckm = ((km - 1) * ny + (j - 1)) * nx + i
371 ckp = ((kp - 1) * ny + (j - 1)) * nx + i
372
373 gx = -(phi[cip] - phi[cim]) * inv2dx
374 gy = -(phi[cjp] - phi[cjm]) * inv2dx
375 gz = -(phi[ckp] - phi[ckm]) * inv2dx
376 r = rho[c]
377 mx0 = mx[c]; my0 = my[c]; mz0 = mz[c]
378 dmx = r * gx * dt
379 dmy = r * gy * dt
380 dmz = r * gz * dt
381 mx[c] = mx0 + dmx
382 my[c] = my0 + dmy
383 mz[c] = mz0 + dmz
384 # Exact kinetic-energy change for a constant acceleration over the kick.
385 de = dt * (mx0 * gx + my0 * gy + mz0 * gz) +
386 (dt * dt * r) * (gx * gx + gy * gy + gz * gz) / oftype(r, 2)
387 E[c] = _compensated_update!(energy_residual,c,E[c],de,track_energy)
388 end
389 end
390
391 """
392 apply_cell_center_gravity!(state, phi, dt)
393
394 Apply the periodic cell-centred acceleration `-grad(phi)` to gas momentum and
395 total energy. The energy update is the exact kinetic-energy change of the kick,
396 so internal and magnetic energy are unchanged. The operation is allocation-free
397 on CPU, CUDA, and Metal backends.
398 """
399 function apply_cell_center_gravity!(s::MHDState, phi::AbstractArray, dt::Real;
400 energy_compensation=nothing)
401 length(phi) == prod(s.dims) || throw(DimensionMismatch("phi does not match state dimensions"))
402 be = KA.get_backend(s.U[1])
403 T = eltype(s.U[1])
404 residual=energy_compensation===nothing ? s.U[5] : energy_compensation
405 length(residual)==length(s.U[5]) ||
406 throw(DimensionMismatch("energy compensation does not match state dimensions"))
407 _apply_cell_center_gravity_k!(be)(s.U[1], s.U[2], s.U[3], s.U[4], s.U[5], phi,
408 residual,T(dt), T(0.5) / T(s.dx), s.dims...,
409 Val(energy_compensation!==nothing);
410 ndrange=length(phi))
411 return s
412 end