PySDM.environments.impl.moist

  1"""
  2common logic for environments featuring moist-air thermodynamics
  3"""
  4
  5from abc import abstractmethod
  6
  7import numpy as np
  8
  9
 10class Moist:
 11    def __init__(self, dt, mesh, variables, mixed_phase=False):
 12        variables += ["water_vapour_mixing_ratio", "thd", "T", "p", "RH"]
 13        if mixed_phase:
 14            variables += ["a_w_ice", "RH_ice"]
 15        all_vars_unique = len(variables) == len(set(variables))
 16        assert all_vars_unique
 17
 18        self.particulator = None
 19        self.dt = dt
 20        self.mesh = mesh
 21        self.variables = variables
 22        self._values = None
 23        self._tmp = None
 24        self._nan_field = None
 25
 26    def register(self, builder):
 27        self.particulator = builder.particulator
 28        self.particulator.observers.append(self)
 29
 30        if self.particulator.formulae.ventilation.__name__ != "Neglect":
 31            for var in ("air density", "air dynamic viscosity"):
 32                if var not in self.variables:
 33                    self.variables += [var]
 34
 35        self._values = {"predicted": None, "current": self._allocate(self.variables)}
 36        self._tmp = self._allocate(self.variables)
 37
 38        self._nan_field = self._allocate(("_",))["_"]
 39        self._nan_field.fill(np.nan)
 40
 41    def _allocate(self, variables):
 42        result = {}
 43        for var in variables:
 44            result[var] = self.particulator.Storage.empty((self.mesh.n_cell,), float)
 45        return result
 46
 47    def __getitem__(self, key: str):
 48        """returns a Storage representing the variable (field) at a given key or
 49        otherwise a NaN-filled Storage if the key is not found (in order to simplify
 50        generic code which uses optional variables, e.g. air viscosity, etc.)"""
 51        if key in self._values["current"]:
 52            return self._values["current"][key]
 53        return self._nan_field
 54
 55    def get_predicted(self, key: str):
 56        if self._values["predicted"] is None:
 57            raise AssertionError(
 58                "It seems the AmbientThermodynamics dynamic was not added"
 59                " when building particulator"
 60            )
 61        return self._values["predicted"][key]
 62
 63    def _recalculate_temperature_pressure_relative_humidity(self, target):
 64        self.particulator.backend.temperature_pressure_rh(
 65            rhod=target["rhod"],
 66            thd=target["thd"],
 67            water_vapour_mixing_ratio=target["water_vapour_mixing_ratio"],
 68            T=target["T"],
 69            p=target["p"],
 70            RH=target["RH"],
 71        )
 72
 73    def sync(self):
 74        target = self._tmp
 75        target["water_vapour_mixing_ratio"].ravel(self.get_water_vapour_mixing_ratio())
 76        target["thd"].ravel(self.get_thd())
 77
 78        self._recalculate_temperature_pressure_relative_humidity(target)
 79
 80        if "a_w_ice" in self.variables:
 81            self.particulator.backend.a_w_ice(
 82                T=target["T"],
 83                p=target["p"],
 84                RH=target["RH"],
 85                water_vapour_mixing_ratio=target["water_vapour_mixing_ratio"],
 86                a_w_ice=target["a_w_ice"],
 87                RH_ice=target["RH_ice"],
 88            )
 89        if "air density" in self.variables:
 90            self.particulator.backend.air_density(
 91                water_vapour_mixing_ratio=target["water_vapour_mixing_ratio"],
 92                rhod=target["rhod"],
 93                output=target["air density"],
 94            )
 95        if "air dynamic viscosity" in self.variables:
 96            self.particulator.backend.air_dynamic_viscosity(
 97                temperature=target["T"],
 98                output=target["air dynamic viscosity"],
 99            )
100        self._values["predicted"] = target
101
102    @abstractmethod
103    def get_water_vapour_mixing_ratio(self) -> np.ndarray:
104        raise NotImplementedError()
105
106    @abstractmethod
107    def get_thd(self) -> np.ndarray:
108        raise NotImplementedError()
109
110    def notify(self):
111        if self._values["predicted"] is None:
112            return
113
114        self._tmp = self._values["current"]
115        self._values["current"] = self._values["predicted"]
116        self._values["predicted"] = None