blpd.model

  1"""
  2`Model` class for running the LPD model.
  3"""
  4
  5import importlib
  6import math
  7import warnings
  8from copy import deepcopy as copy
  9
 10import numpy as np
 11
 12from . import lpd
 13from .utils import disable_numba
 14from .utils import enable_numba
 15from .utils import numbify
 16from .utils import unnumbify
 17
 18__all__ = (
 19    "Model",
 20    "INPUT_PARAM_DEFAULTS",
 21    "calc_MW_derived_params",
 22    "compare_params",
 23)
 24
 25
 26INPUT_PARAM_DEFAULTS = {
 27    #
 28    # particle sources
 29    "release_height": 1.05,  # m; really should draw from a distribution of release heights for each particle!!
 30    "source_positions": [(0, 0)],  # point locs of the particle sources; must be iterable
 31    "dNp_per_dt_per_source": 2,  # should be int
 32    #
 33    # canopy
 34    "canopy_height": 1.1,  # m; h / h_c
 35    "total_LAI": 2.0,  # (total) leaf area index
 36    "foliage_drag_coeff": 0.2,  # C_d
 37    #
 38    # turbulence
 39    "ustar": 0.25,  # m; u_*; friction velocity above canopy (assuming const shear layer)
 40    "von_Karman_constant": 0.4,  # k
 41    "Kolmogorov_C0": 5.5,
 42    #
 43    # run options
 44    "dt": 0.25,  # s; time step for the 1-O Newton FT scheme; this is what Pratt used
 45    "t_tot": 100.0,  # s; total time of the run
 46    "dt_out": 1.0,
 47    "continuous_release": True,
 48    "use_numba": True,
 49    "chemistry_on": False,
 50    #
 51    # chemistry
 52    "fv_0": {},  # fv: floral volatiles. initial values (mol, ug, or somesuch) can be provided here.
 53    "n_air_cm3": 2.62e19,  # (dry) air number density (molec cm^-3)
 54    "oxidants_ppbv": {"O3": 40.0, "OH": 1.0e-4, "NO3": 1.0e-5},
 55    #
 56    # Massman and Weil (MW) canopy wind model parameters (could be dict)
 57    "MW_c1": 0.28,  # c's for above-canopy wind profile
 58    "MW_c2": 0.37,
 59    "MW_c3": 15.1,
 60    "MW_gam1": 2.40,  # gam_i = sig_i/u_star (velocity std's above canopy, in sfc layer)
 61    "MW_gam2": 1.90,
 62    "MW_gam3": 1.25,
 63    "MW_alpha": 0.05,  # parameter that controls in-canopy sigma_w and sigma_u
 64    "MW_A2": 0.6,  # appears to be unused, but is part of their model
 65    # only alpha and A_1 are empirical constants (MW p. 89)
 66}
 67"""
 68Input parameter defaults.
 69"""
 70# could do more dict nesting like in pyAPES...
 71
 72
 73# TODO: opposite of this -- calculate above MW params from normal wind params
 74def calc_MW_derived_params(p):
 75    r"""Calculate Massman and Weil (1999) parameters
 76    from the base wind profile parameters
 77    $c_i$, $\gamma_i$, $\alpha$ ($i = 1\colon3$)
 78    and other model input parameters.
 79
 80    References
 81    ----------
 82    * [Massman and Weil (1999)](https://doi.org/10.1023/A:1001810204560) [MW]
 83    """
 84    cd = p["foliage_drag_coeff"]
 85    ustar = p["ustar"]
 86    LAI = p["total_LAI"]
 87    h = p["canopy_height"]
 88    kconstant = p["von_Karman_constant"]
 89
 90    c1 = p["MW_c1"]
 91    c2 = p["MW_c2"]
 92    c3 = p["MW_c3"]
 93    gam1 = p["MW_gam1"]
 94    gam2 = p["MW_gam2"]
 95    gam3 = p["MW_gam3"]
 96    alpha = p["MW_alpha"]
 97
 98    # derived params
 99    nu1 = (gam1**2 + gam2**2 + gam3**2) ** (-0.5)  # MW p. 86
100    nu3 = (gam1**2 + gam2**2 + gam3**2) ** (1.5)
101    nu2 = nu3 / 6 - gam3**2 / (2 * nu1)
102    # Lam2 = 7/(3*alpha**2*nu1*nu3) + [1/3 - gam3**2*nu1**2]/(3*alpha**2*nu1*nu2)  # the first Lambda^2
103    Lam2 = 3 * nu1**2 / alpha**2  # the simplified Lambda^2 expr; MW p. 87
104    Lam = math.sqrt(Lam2)
105    uh = ustar / (c1 - c2 * math.exp(-c3 * cd * LAI))  # u(h); MW Eq. 5
106    n = cd * LAI / (2 * ustar**2 / uh**2)  # MW Eq. 4, definitely here "n" not "nu"
107    B1 = -(9 * ustar / uh) / (2 * alpha * nu1 * (9 / 4 - Lam**2 * ustar**4 / uh**4))
108
109    d = h * (1 - (1 / (2 * n)) * (1 - math.exp(-2 * n)))  # displacement height
110
111    z0 = (h - d) * math.exp(-kconstant * uh / ustar)  # roughness length
112
113    # Calculate dissipation at canopy top to choose matching approach (Massman and Weil)
114    epsilon_ah = (ustar**3) / (kconstant * (h - d))
115    sig_eh = ustar * (nu3) ** (1 / 3)  # not used elsewhere
116    epsilon_ch = sig_eh**3 * (cd * LAI / h) / (nu3 * alpha)
117    if epsilon_ah >= epsilon_ch:
118        epflag = True
119    else:  # eps_a(h) < eps_c(h)  => usually indicates a relatively dense canopy
120        epflag = False
121
122    r = {
123        "MW_nu1": nu1,
124        "MW_nu2": nu2,
125        "MW_nu3": nu3,
126        "MW_Lam": Lam,
127        "MW_n": n,  # should be eta? (think not)
128        "MW_B1": B1,
129        "U_h": uh,  # mean wind speed at canopy height U(h)
130        "displacement_height": d,
131        "roughness_length": z0,
132        "MW_epsilon_a_h": epsilon_ah,  # above-canopy TKE dissipation rate at h
133        "MW_epsilon_c_h": epsilon_ch,  # in-canopy " "
134        "MW_epsilon_ah_gt_ch": epflag,  # name needs to be more descriptive
135    }
136
137    return r
138
139
140def compare_params(p, p0=None, input_params_only=False, print_message=True):
141    """Compare parameters dict `p` to reference `p0`.
142    `p0` is assumed to be the default parameters dict if not provided.
143    Use `input_params_only=True` to only see those differences in the message.
144    """
145    if p0 is None:
146        p0 = Model().p
147        p0_name = "default"
148    else:
149        p0_name = "reference"
150
151    same = True
152    if any(p[k] != p0[k] for k in p0):
153        same = False
154
155        if input_params_only:
156            p0 = {k: p0[k] for k in p0 if k in INPUT_PARAM_DEFAULTS}
157
158        if print_message:
159            t = f"parameter: {p0_name} --> current"
160            print(t)
161            print("-" * len(t))
162        for k, v0 in sorted(p0.items(), key=lambda x: x[0].lower()):  # don't separate uppercase
163            v = p[k]
164            if v0 != v:
165                if print_message:
166                    print(f"'{k}': {v0} --> {v}")
167
168    if same:
169        if print_message:
170            print(f"all params same as {p0_name}")
171
172    return same
173
174
175class Model:
176    """The LPD model."""
177
178    # class variables (as opposed to instance)
179    _p_user_input_default = INPUT_PARAM_DEFAULTS
180    _p_default_MW = calc_MW_derived_params(_p_user_input_default)
181    _p_input_default = {**_p_user_input_default, **_p_default_MW}
182
183    def __init__(self, p=None):
184        """
185        Parameters
186        ----------
187        p : dict
188            User-supplied input parameters (to update the defaults (`INPUT_PARAM_DEFAULTS`)).
189            `Model.update_p` can also be used to update parameters
190            after creating the `Model` instance.
191        """
192        self.p = copy(Model._p_input_default)  # start with defaults
193        """Model parameters dict. This includes both *input* and *derived*
194        parameters and so should not be modified directly
195        (instead use `Model.update_p`).
196        """
197        if p is None:
198            p = {}
199        self.update_p(**p)  # calculate derived params
200
201        # checks (could move to separate `check_p` method or to `update_p`)
202        assert (
203            self.p["release_height"] <= self.p["canopy_height"]
204        ), "particles must be released within canopy"
205        assert (
206            np.modf(self.p["dt_out"] / self.p["dt"])[0] == 0
207        ), "output interval must be a multiple of dt"
208
209        self._init_state()
210        self._init_hist()
211
212    def update_p(self, **kwargs):
213        """Use `**kwargs` (allowed user input parameters only) to check/update all model parameters
214        (in place).
215        """
216        allowed_keys = self._p_user_input_default.keys()
217        for k, v in kwargs.items():
218            if k not in allowed_keys:
219                msg = f"key '{k}' is not in the default parameter list. ignoring it."
220                warnings.warn(msg)
221            else:
222                if isinstance(v, dict):
223                    self.p[k].update(v)
224                else:
225                    self.p[k] = v
226
227        # calculated parameters (probably should be in setter methods or something to prevent inconsistencies)
228        # i.e., user changing them without using `update_p`
229        self.p["N_sources"] = len(self.p["source_positions"])
230
231        # calculate oxidant concentrations from ppbv values and air number density
232        n_a = self.p["n_air_cm3"]
233        conc_ox = {}
234        for ox_name, ox_ppbv in self.p["oxidants_ppbv"].items():
235            conc_ox[ox_name] = n_a * ox_ppbv * 1e-9
236        self.p.update({"conc_oxidants": conc_ox})
237
238        # calculate number of time steps: N_t from t_tot
239        t_tot = self.p["t_tot"]
240        dt = self.p["dt"]
241        N_t = math.floor(t_tot / dt)  # number of time steps
242        if abs(N_t - t_tot / dt) > 0.01:
243            msg = f"N was rounded down from {t_tot/dt:.4f} to {N_t}"
244            warnings.warn(msg)
245        self.p["N_t"] = N_t  # TODO: consistentify style between N_p and N_t
246
247        # calculate total run number of particles: Np_tot
248        dNp_dt_ds = self.p["dNp_per_dt_per_source"]
249        N_s = self.p["N_sources"]
250        if self.p["continuous_release"]:
251            Np_tot = N_t * dNp_dt_ds * N_s
252            Np_tot_per_source = int(Np_tot / N_s)
253
254        else:
255            Np_tot = dNp_dt_ds * N_s
256            Np_tot_per_source = dNp_dt_ds
257        self.p["Np_tot"] = Np_tot
258        self.p["Np_tot_per_source"] = Np_tot_per_source
259
260        # some variables change the lengths of the state and hist arrays
261        if any(
262            k in kwargs
263            for k in ["t_tot", "dNp_per_dt_per_source", "N_sources", "continuous_release"]
264        ):
265            self._init_state()
266            self._init_hist()
267
268        # some variables affect the derived MW variables
269        #
270        # these are the non-model-parameter inputs:
271        MW_inputs = [
272            "foliage_drag_coeff",
273            "ustar",
274            "total_LAI",
275            "canopy_height",
276            "von_Karman_constant",
277        ]
278        # check for these and also the MW model parameters
279        if any(k in kwargs for k in MW_inputs) or any(k[:2] == "MW" for k in kwargs):
280            self.p.update(calc_MW_derived_params(self.p))
281
282        return self
283
284    # TODO: could change self.p to self._p, but have self.p return a view,
285    #       but give error if user tries to set items
286
287    # TODO: init conc at this time too?
288    def _init_state(self):
289        Np_tot = self.p["Np_tot"]
290        # also could do as 3-D? (x,y,z coords)
291
292        # particle positions
293        xp = np.empty((Np_tot,))
294        yp = np.empty((Np_tot,))
295        zp = np.empty((Np_tot,))
296
297        # local wind speed (perturbation?) at particle positions
298        up = np.zeros((Np_tot,))  # should initial u be = horiz wind speed? or + a perturbation?
299        vp = np.zeros((Np_tot,))
300        wp = np.zeros((Np_tot,))
301
302        # seems ideal to generate sources and initial pos for each before going into time loop
303        # can't be generator because we go over them multiple times
304        # but could have a generator that generatoes that sources for each time step?
305        N_sources = self.p["N_sources"]
306        Np_tot_per_source = self.p["Np_tot_per_source"]
307        release_height = self.p["release_height"]
308        source_positions = self.p["source_positions"]
309        for isource in range(N_sources):
310            ib = isource * Np_tot_per_source
311            ie = (isource + 1) * Np_tot_per_source
312            xp[ib:ie] = source_positions[isource][0]
313            yp[ib:ie] = source_positions[isource][1]
314            zp[ib:ie] = release_height
315
316        # rearrange by time (instead of by source)
317        for p_ in (xp, yp, zp):
318            groups = []
319            for isource in range(N_sources):
320                ib = isource * Np_tot_per_source
321                ie = (isource + 1) * Np_tot_per_source
322                groups.append(p_[ib:ie])
323            p_[:] = np.column_stack(groups).flatten()
324        # assuming sources same strength everywhere for now
325
326        self.state = {
327            # 'k': 0,
328            # 't': 0,
329            # 'Np_k': 0,
330            "xp": xp,
331            "yp": yp,
332            "zp": zp,
333            "up": up,
334            "vp": vp,
335            "wp": wp,
336        }
337
338    def _init_hist(self):
339        if self.p["continuous_release"]:
340            hist = False
341        else:
342            # set up hist if dt_out
343            # should use xarray for this (and for other stuff too; like sympl)
344            if self.p["dt_out"] <= 0:
345                raise ValueError("dt_out must be pos. to use single-release mode")
346                # TODO: should be ok to do continuous_release without hist if want
347
348            t_tot = self.p["t_tot"]
349            dt_out = self.p["dt_out"]
350            Np_tot = self.p["Np_tot"]
351            hist = dict()
352            N_t_hist = int(t_tot / dt_out) + 1
353            hist["pos"] = np.empty((Np_tot, N_t_hist, 3))  # particle, x, y, z
354            hist["ws"] = np.empty((Np_tot, N_t_hist, 3))
355            # could use list instead of particle dim
356            # to allow for those with different record lengths
357
358            xp = self.state["xp"]
359            yp = self.state["yp"]
360            zp = self.state["zp"]
361            up = self.state["up"]
362            vp = self.state["vp"]
363            wp = self.state["wp"]
364            hist["pos"][:, 0, :] = np.column_stack((xp, yp, zp))  # initial positions
365            hist["ws"][:, 0, :] = np.column_stack((up, vp, wp))
366
367        self.hist = hist
368
369    def run(self):
370        """Run the LPD model, integrating the particles using
371        `blpd.lpd.integrate_particles_one_timestep`.
372        """
373        import datetime
374
375        # TODO: change to `time.perf_counter` for this?
376        self._clock_time_run_start = datetime.datetime.now()
377
378        Np_k = 0  # initially tracking 0 particles
379        # Np_tot = self.p['Np_tot']
380        dt = self.p["dt"]
381        dt_out = self.p["dt_out"]
382        N_t = self.p["N_t"]  # number of time steps
383        # t_tot = self.p['t_tot']
384        dNp_dt_ds = self.p["dNp_per_dt_per_source"]
385        N_s = self.p["N_sources"]
386        # outer loop could be particles instead of time. might make some parallelization easier
387
388        # init of hist and state could go here
389
390        if self.p["use_numba"]:
391            enable_numba()  # ensure numba compilation is not disabled
392
393            # > prepare p for numba
394            p_for_nb = {k: v for k, v in self.p.items() if not isinstance(v, (str, list, dict))}
395            # p_for_nb = {k: v for k, v in self.p.items() if isinstance(v, (int, float, np.ndarray))}
396            # p_nb = numbify(p_for_nb)
397            p_nb = numbify(p_for_nb, zerod_only=True)  # only floats/ints
398
399            # > prepare state for numba
400            #  leaving out t and k for now, pass as arguments instead
401            # state_for_nb = {k: v for k, v in self.state.items() if k in ('xp', 'yp', 'zp', 'up', 'vp', 'wp')}
402            state_for_nb = self.state
403            state_nb = numbify(state_for_nb)
404
405            # for debug
406            self.p_nb = p_nb
407            self.state_nb = state_nb
408
409            state_run = state_nb
410            p_run = p_nb
411
412        else:  # model is set not to use numba (for checking the performance advantage of using numba)
413            disable_numba()  # disable numba compilation
414
415            state_run = self.state
416            p_run = self.p
417
418        # changing numba config in lpd redefines the njit decorated functions
419        # but that isn't recognized right away unless we do this re-import
420        # reloading numba in lpd doesn't seem to work
421        # - at least not the first time trying to run after changing use_numba True->False
422        importlib.reload(lpd)
423
424        # print(state_run['xp'].shape)
425
426        for k in range(1, N_t + 1):
427            if self.p["continuous_release"]:
428                Np_k += dNp_dt_ds * N_s
429            else:  # only release at k=1 (at k=0 the particles are inside their release point)
430                if k == 1:
431                    Np_k += dNp_dt_ds * N_s
432                else:
433                    pass
434
435            t = k * dt  # current (elapsed) time
436
437            if self.p["use_numba"]:
438                state_run.update(numbify({"k": k, "t": t, "Np_k": Np_k}))
439            else:
440                state_run.update({"k": [k], "t": [t], "Np_k": [Np_k]})
441            # print(state_run['xp'].shape)
442
443            # pass numba-ified dicts here
444
445            lpd.integrate_particles_one_timestep(state_run, p_run)
446            # integrate_particles_one_timestep(state_run, p_run)
447
448            # TODO: option to save avg / other stats in addition to instantaneous? or specify avg vs instant?
449            if self.hist is not False:
450                if t % dt_out == 0:
451                    o = int(t // dt_out)  # note that `int()` floors anyway
452                    xp = state_run["xp"]
453                    yp = state_run["yp"]
454                    zp = state_run["zp"]
455                    up = state_run["up"]
456                    vp = state_run["vp"]
457                    wp = state_run["wp"]
458                    self.hist["pos"][:, o, :] = np.column_stack((xp, yp, zp))
459                    self.hist["ws"][:, o, :] = np.column_stack((up, vp, wp))
460
461        if self.p["use_numba"]:
462            self.state = unnumbify(state_run)
463        else:
464            self.state = state_run
465
466        self._clock_time_run_end = datetime.datetime.now()
467
468        # self._maybe_run_chem()
469
470        return self
471
472    # def _maybe_run_chem(self):
473    #     # note: adds conc dataset to self.state
474
475    #     # this check/correction logic could be somewhere else
476    #     if self.p['chemistry_on']:
477    #         if not self.p['continuous_release']:
478    #             warnings.warn(
479    #                 'chemistry is calculated only for the continuous release option (`continuous_release=True`). not calculating chemistry',
480    #                 stacklevel=2,
481    #             )
482    #             self.p['chemistry_on'] = False
483
484    #     if self.p["chemistry_on"]:
485    #         conc = chem_calc_options["fixed_oxidants"](self.to_xarray())
486
487    #     else:
488    #         conc = False
489
490    #     self.state.update({
491    #         'conc': conc
492    #     })
493    #     # maybe should instead remove 'conc' from state if not doing chem
494
495    def to_xarray(self):
496        """Create and return an `xr.Dataset` of the LPD run."""
497        # TODO: smoothing/skipping options to reduce storage needed?
498        import json
499
500        import xarray as xr
501
502        ip_coord_tup = (
503            "ip",
504            np.arange(self.p["Np_tot"]),
505            {"long_name": "Lagrangian particle index"},
506        )
507        if self.hist:  # continuous release run
508            t = np.arange(0, self.p["t_tot"] + self.p["dt_out"], self.p["dt_out"])
509            # ^ note can use `pd.to_timedelta(t, unit="s")`
510            dims = ("ip", "t")
511            coords = {
512                "ip": ip_coord_tup,
513                "t": ("t", t, {"long_name": "Simulation elapsed time", "units": "s"}),
514            }
515            x = self.hist["pos"][..., 0]
516            y = self.hist["pos"][..., 1]
517            z = self.hist["pos"][..., 2]
518            u = self.hist["ws"][..., 0]
519            v = self.hist["ws"][..., 1]
520            w = self.hist["ws"][..., 2]
521        else:  # no hist, only current state
522            dims = ("ip",)
523            coords = {"ip": ip_coord_tup}
524            x = self.state["xp"]
525            y = self.state["yp"]
526            z = self.state["zp"]
527            u = self.state["up"]
528            v = self.state["vp"]
529            w = self.state["wp"]
530
531        data_vars = {
532            "x": (dims, x, {"long_name": "$x$", "units": "m"}),
533            "y": (dims, y, {"long_name": "$y$", "units": "m"}),
534            "z": (dims, z, {"long_name": "$z$", "units": "m"}),
535            "u": (dims, u, {"long_name": "$u$", "units": "m s$^{-1}$"}),
536            "v": (dims, v, {"long_name": "$v$", "units": "m s$^{-1}$"}),
537            "w": (dims, w, {"long_name": "$w$", "units": "m s$^{-1}$"}),
538        }
539
540        # Serialize model parameters in JSON to allow saving in netCDF and loading later
541        attrs = {
542            "run_completed": self._clock_time_run_end,
543            "run_runtime": self._clock_time_run_end - self._clock_time_run_start,
544            # TODO: package version once the packaging is better...
545            "p_json": json.dumps(self.p),
546        }
547
548        ds = xr.Dataset(
549            coords=coords,
550            data_vars=data_vars,
551            attrs=attrs,
552        )
553
554        # TODO: Add extra useful coordinates: t_out for non-hist and t as Timedelta for hist
555
556        return ds
557
558    def to_xr(self):
559        """Alias of `Model.to_xarray`."""
560
561    def plot(self, **kwargs):
562        """Make a plot of the results (type based on run type).
563        `**kwargs` are passed through to the relevant plotting function.
564        * single release: `blpd.plot.trajectories`
565        * continuous release: `blpd.plot.final_pos_scatter`
566        """
567        # first check if model has been run
568        p = self.p
569        state = self.state
570        hist = self.hist
571        from . import plot
572
573        if np.all(state["up"] == 0):  # model probably hasn't been run
574            pass  # silently do nothing for now
575        else:
576            if (p["continuous_release"] is False) and hist:
577                plot.trajectories(self.to_xarray(), **kwargs)
578            else:
579                plot.final_pos_scatter(self.to_xarray(), **kwargs)