PySDM_examples.Grabowski_and_Pawlowska_2023.simulation

  1import numpy as np
  2from PySDM_examples.utils import BasicSimulation
  3
  4from PySDM import Particulator
  5from PySDM.backends import CPU
  6from PySDM.backends.impl_numba.test_helpers import scipy_ode_condensation_solver
  7from PySDM.dynamics import AmbientThermodynamics, Condensation
  8from PySDM.environments import Parcel
  9from PySDM.initialisation.hygroscopic_equilibrium import equilibrate_wet_radii
 10from PySDM.initialisation.sampling.spectral_sampling import ConstantMultiplicity
 11from PySDM.physics import si
 12
 13
 14class Simulation(BasicSimulation):
 15    def __init__(
 16        self,
 17        settings,
 18        products=None,
 19        scipy_solver=False,
 20    ):
 21
 22        environment = Parcel(
 23            dt=settings.timestep,
 24            p0=settings.initial_pressure,
 25            initial_relative_humidity=settings.initial_relative_humidity,
 26            T0=settings.initial_temperature,
 27            w=settings.vertical_velocity,
 28            mass_of_dry_air=44 * si.kg,
 29            backend=CPU(
 30                formulae=settings.formulae, override_jit_flags={"parallel": False}
 31            ),
 32        )
 33        volume = environment.mass_of_dry_air / settings.initial_air_density
 34        attributes = {
 35            k: np.empty(0)
 36            for k in ("dry volume", "kappa times dry volume", "multiplicity")
 37        }
 38
 39        assert len(settings.aerosol_modes_by_kappa.keys()) == 1
 40        kappa = tuple(settings.aerosol_modes_by_kappa.keys())[0]
 41        spectrum = settings.aerosol_modes_by_kappa[kappa]
 42
 43        r_dry, n_per_volume = ConstantMultiplicity(spectrum).sample_deterministic(
 44            settings.n_sd
 45        )
 46        v_dry = settings.formulae.trivia.volume(radius=r_dry)
 47        attributes["multiplicity"] = np.append(
 48            attributes["multiplicity"], n_per_volume * volume
 49        )
 50        attributes["dry volume"] = np.append(attributes["dry volume"], v_dry)
 51        attributes["kappa times dry volume"] = np.append(
 52            attributes["kappa times dry volume"], v_dry * kappa
 53        )
 54        r_wet = equilibrate_wet_radii(
 55            r_dry=settings.formulae.trivia.radius(volume=attributes["dry volume"]),
 56            environment=environment,
 57            kappa_times_dry_volume=attributes["kappa times dry volume"],
 58        )
 59        attributes["volume"] = settings.formulae.trivia.volume(radius=r_wet)
 60
 61        particulator = Particulator(
 62            n_sd=settings.n_sd,
 63            environment=environment,
 64            dynamics=(
 65                AmbientThermodynamics(),
 66                Condensation(rtol_thd=settings.rtol_thd, rtol_x=settings.rtol_x),
 67            ),
 68            attributes=attributes,
 69            products=products,
 70            requested_attributes=(
 71                "critical saturation",
 72                "equilibrium saturation",
 73                "critical volume",
 74            ),
 75        )
 76
 77        super().__init__(particulator=particulator)
 78        if scipy_solver:
 79            scipy_ode_condensation_solver.patch_particulator(self.particulator)
 80
 81        self.output_attributes = {
 82            "volume": tuple([] for _ in range(self.particulator.n_sd)),
 83            "dry volume": tuple([] for _ in range(self.particulator.n_sd)),
 84            "critical saturation": tuple([] for _ in range(self.particulator.n_sd)),
 85            "equilibrium saturation": tuple([] for _ in range(self.particulator.n_sd)),
 86            "critical volume": tuple([] for _ in range(self.particulator.n_sd)),
 87            "multiplicity": tuple([] for _ in range(self.particulator.n_sd)),
 88        }
 89        self.settings = settings
 90
 91        self.__sanity_checks(attributes, volume)
 92
 93    def __sanity_checks(self, attributes, volume):
 94        for attribute in attributes.values():
 95            assert attribute.shape[0] == self.particulator.n_sd
 96        np.testing.assert_approx_equal(
 97            sum(attributes["multiplicity"]) / volume,
 98            sum(
 99                mode.norm_factor
100                for mode in self.settings.aerosol_modes_by_kappa.values()
101            ),
102            significant=4,
103        )
104
105    def _save(self, output):
106        for key, attr in self.output_attributes.items():
107            attr_data = self.particulator.attributes[key].to_ndarray()
108            for drop_id in range(self.particulator.n_sd):
109                attr[drop_id].append(attr_data[drop_id])
110        super()._save(output)
111
112    def run(self):
113        output_products = super()._run(
114            self.settings.nt, self.settings.steps_per_output_interval
115        )
116        return {"products": output_products, "attributes": self.output_attributes}