PySDM_examples.Yang_et_al_2018.simulation

  1import numpy as np
  2
  3import PySDM.products as PySDM_products
  4from PySDM import Formulae, Particulator
  5from PySDM.backends import CPU
  6from PySDM.dynamics import AmbientThermodynamics, Condensation
  7from PySDM.environments import Parcel
  8from PySDM.physics import si
  9
 10
 11class Simulation:
 12    def __init__(self, settings, backend=CPU):
 13        dt_output = (
 14            settings.total_time / settings.n_steps
 15        )  # TODO #334 overwritten in notebook
 16        self.n_substeps = 1  # TODO #334 use condensation substeps
 17        while dt_output / self.n_substeps >= settings.dt_max:
 18            self.n_substeps += 1
 19        self.formulae = Formulae(
 20            diffusion_coordinate=settings.coord,
 21            saturation_vapour_pressure="AugustRocheMagnus",
 22        )
 23        self.bins_edges = self.formulae.trivia.volume(settings.r_bins_edges)
 24
 25        env = Parcel(
 26            dt=dt_output / self.n_substeps,
 27            mass_of_dry_air=settings.mass_of_dry_air,
 28            p0=settings.p0,
 29            initial_relative_humidity=settings.RH0,
 30            T0=settings.T0,
 31            w=settings.w,
 32            z0=settings.z0,
 33            backend=backend(
 34                formulae=self.formulae, override_jit_flags={"parallel": False}
 35            ),
 36        )
 37
 38        condensation = Condensation(
 39            adaptive=settings.adaptive,
 40            rtol_x=settings.rtol_x,
 41            rtol_thd=settings.rtol_thd,
 42            dt_cond_range=settings.dt_cond_range,
 43        )
 44
 45        products = [
 46            PySDM_products.ParticleSizeSpectrumPerVolume(
 47                name="Particles Wet Size Spectrum",
 48                radius_bins_edges=settings.r_bins_edges,
 49            ),
 50            PySDM_products.CondensationTimestepMin(name="dt_cond_min"),
 51            PySDM_products.CondensationTimestepMax(name="dt_cond_max"),
 52            PySDM_products.RipeningRate(),
 53            PySDM_products.MeanRadius(
 54                name="r_mean_gt_1_um", radius_range=(1 * si.um, np.inf)
 55            ),
 56            PySDM_products.ActivatedMeanRadius(
 57                name="r_act", count_activated=True, count_unactivated=False
 58            ),
 59            PySDM_products.Time(name="t"),
 60        ]
 61
 62        attributes = env.init_attributes(
 63            n_in_dv=settings.n, kappa=settings.kappa, r_dry=settings.r_dry
 64        )
 65
 66        self.particulator = Particulator(
 67            attributes=attributes,
 68            products=products,
 69            n_sd=settings.n_sd,
 70            environment=env,
 71            dynamics=(
 72                AmbientThermodynamics(),
 73                condensation,
 74            ),
 75        )
 76
 77        self.n_steps = settings.n_steps
 78
 79    def save(self, output):
 80        _sp = self.particulator
 81        cell_id = 0
 82        output["r_bins_values"].append(
 83            (_sp.products["Particles Wet Size Spectrum"].get()).T
 84        )
 85        volume = _sp.attributes["volume"].to_ndarray()
 86        output["r"].append((self.formulae.trivia.radius(volume=volume)).T)
 87        output["S"].append(_sp.environment["RH"][cell_id] - 1)
 88        output["t"].append(_sp.products["t"].get())
 89        for key in ("water_vapour_mixing_ratio", "T", "z"):
 90            output[key].append(_sp.environment[key][cell_id])
 91        for key in (
 92            "dt_cond_max",
 93            "dt_cond_min",
 94            "ripening rate",
 95            "r_mean_gt_1_um",
 96            "r_act",
 97        ):
 98            output[key].append(_sp.products[key].get()[cell_id].copy())
 99
100    def run(self):
101        output = {
102            key: []
103            for key in (
104                "r",
105                "S",
106                "z",
107                "t",
108                "water_vapour_mixing_ratio",
109                "T",
110                "r_bins_values",
111                "dt_cond_max",
112                "dt_cond_min",
113                "ripening rate",
114                "r_mean_gt_1_um",
115                "r_act",
116            )
117        }
118
119        self.save(output)
120        for _ in range(self.n_steps):
121            self.particulator.run(self.n_substeps)
122            self.save(output)
123        for k, v in output.items():
124            output[k] = np.asarray(v)
125        return output