blpd.lpd

  1"""
  2Lagrangian stochastic particle dispersion.
  3Using [Numba](https://numba.pydata.org/), particles are integrated in parallel.
  4
  5Although it is possible to use these functions directly,
  6the intention is for `blpd.model.Model` to be used.
  7"""
  8
  9import numpy as np
 10from numba import njit
 11from numba import prange
 12
 13
 14@njit
 15def _calc_fd_params_above_canopy(x, p):
 16    """Calculate the dispersion parameters above the canopy."""
 17    z = x[2]
 18
 19    ustar = p["ustar"]
 20    kconstant = p["von_Karman_constant"]
 21    d = p["displacement_height"]
 22
 23    z0 = p["roughness_length"]
 24    gam1 = p["MW_gam1"]
 25    gam2 = p["MW_gam2"]
 26    gam3 = p["MW_gam3"]
 27
 28    umean = (ustar / kconstant) * np.log((z - d) / z0)
 29    dumeandz = (ustar / (kconstant)) * (1 / (z - d))
 30
 31    tau11 = (ustar * gam1) ** 2
 32    tau22 = (ustar * gam2) ** 2
 33    tau33 = (ustar * gam3) ** 2
 34    tau13 = -(ustar**2)
 35
 36    dtau11dz = 0
 37    dtau22dz = 0
 38    dtau33dz = 0
 39    dtau13dz = 0
 40
 41    lam11, lam22, lam33, lam13 = _calc_Rodean_lambdas(tau11, tau22, tau33, tau13)
 42
 43    # Dissipation - above canopy (log wind profile)
 44    # ref: e.g., MW Eq. 12
 45    epsilon = (ustar**3) / (kconstant * (z - d))
 46
 47    return (
 48        umean,
 49        dumeandz,
 50        dtau11dz,
 51        dtau22dz,
 52        dtau33dz,
 53        dtau13dz,
 54        lam11,
 55        lam22,
 56        lam33,
 57        lam13,
 58        epsilon,
 59    )
 60
 61
 62@njit
 63def _calc_fd_params_in_canopy(x, p):
 64    """Calculate the dispersion parameters within the canopy."""
 65    z = x[2]
 66
 67    ustar = p["ustar"]  # at h?
 68    kconstant = p["von_Karman_constant"]
 69    d = p["displacement_height"]  # for log wind profile; d < h
 70    h = p["canopy_height"]
 71    uh = p["U_h"]  # wind speed at canopy height
 72    cd = p["foliage_drag_coeff"]
 73    LAI = p["total_LAI"]
 74
 75    # z0 = p['roughness_length']
 76    gam1 = p["MW_gam1"]
 77    gam2 = p["MW_gam2"]
 78    gam3 = p["MW_gam3"]
 79
 80    nu1 = p["MW_nu1"]
 81    # nu2 = p['MW_nu2']  # unused?
 82    nu3 = p["MW_nu3"]
 83    Lam = p["MW_Lam"]
 84    n = p["MW_n"]
 85    B1 = p["MW_B1"]
 86    alpha = p["MW_alpha"]
 87
 88    epflag = p["MW_epsilon_ah_gt_ch"]
 89    epsilon_ah = p["MW_epsilon_a_h"]
 90    epsilon_ch = p["MW_epsilon_c_h"]
 91
 92    # zeta(z): "cumulative leaf drag area per unit planform area" (MW p. 82) (ζ)
 93    # The code we started from replace terms zeta(z)/zeta(h) by z/h.
 94    # This assumes a vertically homogeneous LAI dist (LAD const with z in canopy)
 95    # that accumulates from z=0, not z=h_c like LAI.
 96    # zeta(h) is replaced by cd*LAI
 97
 98    umean = uh * np.exp(-n * (1 - (z / h)))
 99    dumeandz = uh * (n / h) * np.exp(-n * (1 - (z / h)))
100
101    zet_h = cd * LAI  # zeta(h) total leaf drag area
102
103    # ref: MW eqn. 10, p. 87
104    sig_e = ustar * (
105        nu3 * np.exp(-Lam * zet_h * (1 - z / h))
106        + B1 * (np.exp(-3 * n * (1 - z / h)) - np.exp(-Lam * zet_h * (1 - z / h)))
107    ) ** (1.0 / 3)
108
109    # I can't find a source for this eqn
110    dsig_e2dz = (
111        (2 / 3)
112        * ustar**2
113        * (
114            (
115                nu3 * np.exp(-Lam * zet_h * (1 - z / h))
116                + B1 * (np.exp(-3 * n * (1 - z / h)) - np.exp(-Lam * zet_h * (1 - z / h)))
117            )
118            ** (-1.0 / 3)
119        )
120        * (
121            (nu3 * Lam * zet_h / h) * np.exp(-Lam * zet_h * (1 - z / h))
122            + B1
123            * (
124                (3 * n / h) * np.exp(-3 * n * (1 - z / h))
125                - (Lam * zet_h / h) * np.exp(-Lam * zet_h * (1 - z / h))
126            )
127        )
128    )
129
130    # From MW eq. 11, in the canopy:
131    #   gam_i * nu_1 * sig_e = sig_i
132    # and above the canopy:
133    #   sig_i/u_star = gam_i    (p. 86)
134
135    tau11 = (gam1 * nu1 * sig_e) ** 2
136    dtau11dz = ((gam1 * nu1) ** 2) * dsig_e2dz
137
138    tau22 = (gam2 * nu1 * sig_e) ** 2
139    dtau22dz = ((gam2 * nu1) ** 2) * dsig_e2dz
140
141    tau33 = (gam3 * nu1 * sig_e) ** 2
142    dtau33dz = ((gam3 * nu1) ** 2) * dsig_e2dz
143
144    tau13 = -(ustar**2) * np.exp(-2 * n * (1 - (z / h)))
145    dtau13dz = -(ustar**2) * (2 * n / h) * np.exp(-2 * n * (1 - (z / h)))
146
147    # Dissipation
148    # ref: MW p. 88
149    if epflag:
150        epsilon = (sig_e**3) * zet_h / (h * nu3 * alpha) * (epsilon_ah / epsilon_ch)
151    else:
152        if z <= d:  # 0 < z <= d
153            epsilon = (sig_e**3) * zet_h / (h * nu3 * alpha)
154        else:  # d < z <= h_c
155            scale1 = zet_h * sig_e**3 / (nu3 * alpha * ustar**3)
156            scale2 = h / (kconstant * (z - d))
157            # ^ this is a potential source of div by 0 (if z v close to d)
158            epsilon = (ustar**3 / h) * min(scale1, scale2)
159
160    lam11, lam22, lam33, lam13 = _calc_Rodean_lambdas(tau11, tau22, tau33, tau13)
161
162    return (
163        umean,
164        dumeandz,
165        dtau11dz,
166        dtau22dz,
167        dtau33dz,
168        dtau13dz,
169        lam11,
170        lam22,
171        lam33,
172        lam13,
173        epsilon,
174    )
175
176
177@njit
178def _calc_Rodean_lambdas(tau11, tau22, tau33, tau13):
179    # ref: Pratt thesis Eq. 2.14, p. 15
180    lam11 = 1.0 / (tau11 - ((tau13**2) / tau33))
181    lam22 = 1.0 / tau22
182    lam33 = 1.0 / (tau33 - ((tau13**2) / tau11))
183    lam13 = 1.0 / (tau13 - ((tau11 * tau33) / tau13))
184
185    return lam11, lam22, lam33, lam13
186
187
188@njit
189def calc_xtend(x, u, p):
190    r"""Calculate the position tendencies
191    (i.e., the velocity components $u_i + \delta u_i$)
192    for one particle with position vector `x` and velocity vector `u`.
193
194    `p` (parameters) must be a Numba typed dict (:class:`numba.typed.Dict`).
195    """
196    z = x[2]  # x, y, z
197    u1, u2, u3 = u  # u, v, w
198
199    h_c = p["canopy_height"]
200    C0 = p["Kolmogorov_C0"]  # a Kolmogorov constant (3--10)
201    dt = p["dt"]
202
203    # Calculate fd (fluid dynamics) params
204    if z >= h_c:
205        (
206            U1,
207            dU1dx3,
208            dtau11dx3,
209            dtau22dx3,
210            dtau33dx3,
211            dtau13dx3,
212            lam11,
213            lam22,
214            lam33,
215            lam13,
216            eps,
217        ) = _calc_fd_params_above_canopy(x, p)
218    else:
219        (
220            U1,
221            dU1dx3,
222            dtau11dx3,
223            dtau22dx3,
224            dtau33dx3,
225            dtau13dx3,
226            lam11,
227            lam22,
228            lam33,
229            lam13,
230            eps,
231        ) = _calc_fd_params_in_canopy(x, p)
232
233    # Calculate new positions
234    # ref: Pratt thesis Eqs. 2.12-13, p. 15
235    # dW_j is an incremental Wiener process with a mean of 0, and a variance of dt
236    # simulated as sqrt(dt)*N(0,1)  (p. 12)
237
238    sqrt_C0eps = np.sqrt(C0 * eps)
239    sqrt_dt = np.sqrt(dt)
240
241    # TODO: if wanted to be able to use other sorts of integrators,
242    # would have to freeze the random seed or something?
243    # (otherwise the tends will be different every time)
244    # or separate the random part from the rest of the tend calculation.
245    # The dist to draw from could be an input.
246
247    # x-1 component
248    randn = np.random.standard_normal()  # different for each component
249    dW1 = sqrt_dt * randn
250    du1 = (
251        (-C0 * eps / 2 * (lam11 * (u1 - U1) + lam13 * u3) + dU1dx3 * u3 + dtau13dx3 / 2) * dt
252        + (
253            dtau11dx3 * (lam11 * (u1 - U1) + lam13 * u3)
254            + dtau13dx3 * (lam13 * (u1 - U1) + lam33 * u3)
255        )
256        * (u3 / 2 * dt)
257        + sqrt_C0eps * dW1
258    )
259
260    # x-2 component
261    randn = np.random.standard_normal()
262    dW2 = sqrt_dt * randn
263    du2 = (-C0 * eps / 2 * lam22 * u2 + dtau22dx3 * lam22 * u2 * u3 / 2) * dt + sqrt_C0eps * dW2
264
265    # x-3 component
266    randn = np.random.standard_normal()
267    dW3 = sqrt_dt * randn
268    du3 = (
269        (-C0 * eps / 2 * (lam13 * (u1 - U1) + lam33 * u3) + dtau33dx3 / 2) * dt
270        + (
271            dtau13dx3 * (lam11 * (u1 - U1) + lam13 * u3)
272            + dtau33dx3 * (lam13 * (u1 - U1) + lam33 * u3)
273        )
274        * (u3 / 2 * dt)
275        + sqrt_C0eps * dW3
276    )
277
278    # Check for too-high values
279    if np.any(np.abs(np.array([du1, du2, du3])) > 50):
280        print("one du is too high:", du1, du2, du3)
281        print("umean, dumeandz, epsilon", U1, dU1dx3, eps)
282        if z >= h_c:
283            print("  z >= h_c")
284        else:
285            print("  z < h_c, z =", z)
286        print("resetting new ws to 0")
287        du1 = 0.0
288        du2 = 0.0
289        du3 = 0.0
290        u1 = 0.0
291        u2 = 0.0
292        u3 = 0.0
293
294    return [u1 + du1, u2 + du2, u3 + du3]
295
296
297@njit
298def _integrate_particle_one_timestep(x, u, p):
299    """Take the state for one particle and integrate."""
300    dt = p["dt"]
301
302    dxdt = calc_xtend(x, u, p)  # u, v, w
303
304    # Adjustments for undesired new z position
305    # (below z = 0 or above PBL inversion)
306    # TODO: original code noted that this needs work
307    # This current method seems to correspond to completely elastic collision,
308    # but maybe inelastic or some probability of deposition could be another option.
309    z = x[2]  # before integration
310    znew = z + dxdt[2] * dt
311    if znew <= 0.1:
312        dxdt = [dxdt[0], dxdt[1], -dxdt[2]]
313        xnew = [x[0] + dxdt[0] * dt, x[1] + dxdt[1] * dt, 0.1]
314
315    elif znew >= 800.0:  # reflect at z_i too
316        dxdt = [dxdt[0], dxdt[1], -dxdt[2]]
317        xnew = [x[0] + dxdt[0] * dt, x[1] + dxdt[1] * dt, 800.0]
318
319    else:
320        xnew = [xi + dxidt * dt for (xi, dxidt) in zip(x, dxdt)]
321
322    # The new position (x) tendencies are the (new) velocity components
323    unew = dxdt
324
325    return xnew, unew
326
327
328# @njit
329@njit(parallel=True)
330def integrate_particles_one_timestep(state, p):
331    """Integrate all particles one time step, modifying the current `state` (in place).
332
333    `state` and `p` (parameters) must be Numba typed dicts (:class:`numba.typed.Dict`)
334    in order for `@njit` to work if Numba is not disabled.
335    """
336    Np_k = int(state["Np_k"][0])
337    xp = state["xp"]
338    yp = state["yp"]
339    zp = state["zp"]
340    up = state["up"]
341    vp = state["vp"]
342    wp = state["wp"]
343
344    # for i in range(Np_k):
345    for i in prange(Np_k):  # pylint: disable=not-an-iterable
346        # State for particle i: position and (wind) velocity
347        x = [xp[i], yp[i], zp[i]]  # position of particle i at current time step
348        u = [up[i], vp[i], wp[i]]  # local wind speed for particle
349
350        # Calculate new state
351        x, u = _integrate_particle_one_timestep(x, u, p)
352
353        # Update state arrays in place
354        xp[i] = x[0]
355        yp[i] = x[1]
356        zp[i] = x[2]
357        up[i] = u[0]
358        vp[i] = u[1]
359        wp[i] = u[2]
360
361        # Note that we are not really time-integrating the wind speed.
362        # We are making assumptions about the distribution of perturbations
363        # about a mean x-direction wind!
364
365        # del x, u
366        # ^ seems like Numba might require this? (gave errors otherwise)
367        # Now, using numba v0.49.1 this raises error:
368        #   `CompilerError: Illegal IR, del found at: del pos`
369        # If trying to use sufficiently older versions of numba, may need to put it back.
370        # TODO: maybe change numba version spec in setup.cfg to >= 0.49