PySDM.environments.parcel

  1"""
  2Zero-dimensional adiabatic parcel framework
  3"""
  4
  5from typing import List, Optional
  6
  7import numpy as np
  8
  9from PySDM.environments.impl.moist import Moist
 10from PySDM.impl.mesh import Mesh
 11from PySDM.initialisation.equilibrate_wet_radii import (
 12    default_rtol,
 13    equilibrate_wet_radii,
 14)
 15from PySDM.environments.impl import register_environment
 16
 17
 18@register_environment()
 19class Parcel(Moist):  # pylint: disable=too-many-instance-attributes
 20    def __init__(
 21        self,
 22        *,
 23        dt,
 24        mass_of_dry_air: float,
 25        p0: float,
 26        initial_water_vapour_mixing_ratio: float,
 27        T0: float,
 28        w: [float, callable],
 29        z0: float = 0,
 30        mixed_phase=False,
 31        variables: Optional[List[str]] = None,
 32    ):
 33        variables = (variables or []) + ["rhod", "z"]
 34        super().__init__(dt, Mesh.mesh_0d(), variables, mixed_phase=mixed_phase)
 35
 36        self.p0 = p0
 37        self.initial_water_vapour_mixing_ratio = initial_water_vapour_mixing_ratio
 38        self.T0 = T0
 39        self.z0 = z0
 40        self.mass_of_dry_air = mass_of_dry_air
 41        self.w = w if callable(w) else lambda _: w
 42        self.delta_liquid_water_mixing_ratio = np.nan
 43
 44    @property
 45    def dv(self):
 46        rhod_mean = (self.get_predicted("rhod")[0] + self["rhod"][0]) / 2
 47        return self.particulator.formulae.trivia.volume_of_density_mass(
 48            rhod_mean, self.mass_of_dry_air
 49        )
 50
 51    def register(self, builder):
 52        formulae = builder.particulator.formulae
 53        pd0 = formulae.trivia.p_d(self.p0, self.initial_water_vapour_mixing_ratio)
 54        rhod0 = formulae.state_variable_triplet.rhod_of_pd_T(pd0, self.T0)
 55        self.mesh.dv = formulae.trivia.volume_of_density_mass(
 56            rhod0, self.mass_of_dry_air
 57        )
 58
 59        Moist.register(self, builder)
 60
 61        self["water_vapour_mixing_ratio"][:] = self.initial_water_vapour_mixing_ratio
 62        self["thd"][:] = formulae.trivia.th_std(pd0, self.T0)
 63        self["rhod"][:] = rhod0
 64        self["z"][:] = self.z0
 65
 66        self._tmp["water_vapour_mixing_ratio"][
 67            :
 68        ] = self.initial_water_vapour_mixing_ratio
 69        self.sync_parcel_vars()
 70        Moist.sync(self)
 71        self.notify()
 72
 73    def init_attributes(
 74        self,
 75        *,
 76        n_in_dv: [float, np.ndarray],
 77        kappa: float,
 78        r_dry: [float, np.ndarray],
 79        rtol=default_rtol,
 80        include_dry_volume_in_attribute: bool = True,
 81    ):
 82        if not isinstance(n_in_dv, np.ndarray):
 83            r_dry = np.array([r_dry])
 84            n_in_dv = np.array([n_in_dv])
 85
 86        attributes = {}
 87        dry_volume = self.particulator.formulae.trivia.volume(radius=r_dry)
 88        attributes["kappa times dry volume"] = dry_volume * kappa
 89        attributes["multiplicity"] = n_in_dv
 90        r_wet = equilibrate_wet_radii(
 91            r_dry=r_dry,
 92            environment=self,
 93            kappa_times_dry_volume=attributes["kappa times dry volume"],
 94            rtol=rtol,
 95        )
 96        attributes["volume"] = self.particulator.formulae.trivia.volume(radius=r_wet)
 97        if include_dry_volume_in_attribute:
 98            attributes["dry volume"] = dry_volume
 99        return attributes
100
101    def advance_parcel_vars(self):
102        """compute new values of displacement, dry-air density and volume,
103        and write them to self._tmp and self.mesh.dv"""
104        dt = self.particulator.dt
105        formulae = self.particulator.formulae
106        T = self["T"][0]
107        p = self["p"][0]
108
109        dz_dt = self.w((self.particulator.n_steps + 1 / 2) * dt)  # "mid-point"
110        water_vapour_mixing_ratio = (
111            self["water_vapour_mixing_ratio"][0]
112            - self.delta_liquid_water_mixing_ratio / 2
113        )
114
115        # derivative evaluated at p_old, T_old, mixrat_mid, w_mid
116        drho_dz = formulae.hydrostatics.drho_dz(
117            p=p,
118            T=T,
119            water_vapour_mixing_ratio=water_vapour_mixing_ratio,
120            lv=formulae.latent_heat_vapourisation.lv(T),
121            d_liquid_water_mixing_ratio__dz=(
122                self.delta_liquid_water_mixing_ratio / dz_dt / dt
123            ),
124        )
125        drhod_dz = drho_dz  # TODO #407
126
127        self.particulator.backend.explicit_euler(self._tmp["z"], dt, dz_dt)
128        self.particulator.backend.explicit_euler(
129            self._tmp["rhod"], dt, dz_dt * drhod_dz
130        )
131
132        self.mesh.dv = formulae.trivia.volume_of_density_mass(
133            (self._tmp["rhod"][0] + self["rhod"][0]) / 2, self.mass_of_dry_air
134        )
135
136    def get_thd(self):
137        return self["thd"]
138
139    def get_water_vapour_mixing_ratio(self):
140        return self["water_vapour_mixing_ratio"]
141
142    def sync_parcel_vars(self):
143        self.delta_liquid_water_mixing_ratio = (
144            self._tmp["water_vapour_mixing_ratio"][0]
145            - self["water_vapour_mixing_ratio"][0]
146        )
147        for var in self.variables:
148            self._tmp[var][:] = self[var][:]
149
150    def sync(self):
151        self.sync_parcel_vars()
152        self.advance_parcel_vars()
153        super().sync()