PySDM.dynamics.condensation

  1"""
  2bespoke condensational growth solver
  3with implicit-in-particle-size integration and adaptive timestepping
  4"""
  5
  6from collections import namedtuple
  7
  8import numpy as np
  9
 10from PySDM.physics import si
 11from PySDM.dynamics.impl import register_dynamic
 12
 13DEFAULTS = namedtuple("_", ("rtol_x", "rtol_thd", "cond_range", "schedule"))(
 14    rtol_x=1e-6,
 15    rtol_thd=1e-6,
 16    cond_range=(1e-4 * si.second, 1 * si.second),
 17    schedule="dynamic",
 18)
 19
 20
 21@register_dynamic()
 22class Condensation:  # pylint: disable=too-many-instance-attributes
 23    def __init__(
 24        self,
 25        *,
 26        rtol_x=DEFAULTS.rtol_x,
 27        rtol_thd=DEFAULTS.rtol_thd,
 28        substeps: int = 1,
 29        adaptive: bool = True,
 30        dt_cond_range: tuple = DEFAULTS.cond_range,
 31        schedule: str = DEFAULTS.schedule,
 32        max_iters: int = 16,
 33        update_thd: bool = True,
 34    ):
 35        if adaptive and substeps != 1:
 36            raise ValueError(
 37                "if specifying substeps count manually, adaptivity must be disabled"
 38            )
 39
 40        self.particulator = None
 41        self.enable = True
 42
 43        self.rtol_x = rtol_x
 44        self.rtol_thd = rtol_thd
 45
 46        self.rh_max = None
 47        self.success = None
 48
 49        self.__substeps = substeps
 50        self.adaptive = adaptive
 51        self.counters = {}
 52        self.dt_cond_range = dt_cond_range
 53        self.schedule = schedule
 54        self.max_iters = max_iters
 55
 56        self.cell_order = None
 57
 58        self.update_thd = update_thd
 59
 60    def register(self, builder):
 61        self.particulator = builder.particulator
 62
 63        builder._set_condensation_parameters(
 64            dt_range=self.dt_cond_range,
 65            adaptive=self.adaptive,
 66            fuse=32,
 67            multiplier=2,
 68            RH_rtol=1e-7,
 69            max_iters=self.max_iters,
 70        )
 71        builder.request_attribute("critical volume")
 72        builder.request_attribute("kappa")
 73        builder.request_attribute("dry volume organic fraction")
 74        builder.request_attribute("Reynolds number")
 75
 76        for counter in ("n_substeps", "n_activating", "n_deactivating", "n_ripening"):
 77            self.counters[counter] = self.particulator.Storage.empty(
 78                self.particulator.mesh.n_cell, dtype=int
 79            )
 80            if counter == "n_substeps":
 81                self.counters[counter][:] = self.__substeps if not self.adaptive else -1
 82            else:
 83                self.counters[counter][:] = -1
 84
 85        self.rh_max = self.particulator.Storage.empty(
 86            self.particulator.mesh.n_cell, dtype=float
 87        )
 88        self.rh_max[:] = np.nan
 89        self.success = self.particulator.Storage.empty(
 90            self.particulator.mesh.n_cell, dtype=bool
 91        )
 92        self.success[:] = False
 93        self.cell_order = np.arange(self.particulator.mesh.n_cell)
 94
 95    def __call__(self):
 96        if self.enable:
 97            if self.schedule == "dynamic":
 98                self.cell_order = np.argsort(self.counters["n_substeps"])
 99            elif self.schedule == "static":
100                pass
101            else:
102                raise NotImplementedError()
103
104            self.particulator.condensation(
105                rtol_x=self.rtol_x,
106                rtol_thd=self.rtol_thd,
107                counters=self.counters,
108                RH_max=self.rh_max,
109                success=self.success,
110                cell_order=self.cell_order,
111            )
112            if not self.success.all():
113                raise RuntimeError("Condensation failed")
114            if not self.update_thd:
115                self.particulator.environment.get_predicted("thd").ravel(
116                    self.particulator.environment.get_thd()
117                )
118            # note: this makes order of dynamics matter
119            #       (e.g., condensation after chemistry or before)
120            self.particulator.update_TpRH()
121
122            if self.adaptive:
123                self.counters["n_substeps"][:] = np.maximum(
124                    self.counters["n_substeps"][:],
125                    int(self.particulator.dt / self.dt_cond_range[1]),
126                )
127                if self.dt_cond_range[0] != 0:
128                    self.counters["n_substeps"][:] = np.minimum(
129                        self.counters["n_substeps"][:],
130                        int(self.particulator.dt / self.dt_cond_range[0]),
131                    )