PySDM_examples.featured.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(formulae=settings.formulae),
 30        )
 31        volume = environment.mass_of_dry_air / settings.initial_air_density
 32        attributes = {
 33            k: np.empty(0)
 34            for k in ("dry volume", "kappa times dry volume", "multiplicity")
 35        }
 36
 37        assert len(settings.aerosol_modes_by_kappa.keys()) == 1
 38        kappa = tuple(settings.aerosol_modes_by_kappa.keys())[0]
 39        spectrum = settings.aerosol_modes_by_kappa[kappa]
 40
 41        r_dry, n_per_volume = ConstantMultiplicity(spectrum).sample_deterministic(
 42            settings.n_sd
 43        )
 44        v_dry = settings.formulae.trivia.volume(radius=r_dry)
 45        attributes["multiplicity"] = np.append(
 46            attributes["multiplicity"], n_per_volume * volume
 47        )
 48        attributes["dry volume"] = np.append(attributes["dry volume"], v_dry)
 49        attributes["kappa times dry volume"] = np.append(
 50            attributes["kappa times dry volume"], v_dry * kappa
 51        )
 52        r_wet = equilibrate_wet_radii(
 53            r_dry=settings.formulae.trivia.radius(volume=attributes["dry volume"]),
 54            environment=environment,
 55            kappa_times_dry_volume=attributes["kappa times dry volume"],
 56        )
 57        attributes["volume"] = settings.formulae.trivia.volume(radius=r_wet)
 58
 59        particulator = Particulator(
 60            n_sd=settings.n_sd,
 61            environment=environment,
 62            dynamics=(
 63                AmbientThermodynamics(),
 64                Condensation(rtol_thd=settings.rtol_thd, rtol_x=settings.rtol_x),
 65            ),
 66            attributes=attributes,
 67            products=products,
 68            requested_attributes=(
 69                "critical saturation",
 70                "equilibrium saturation",
 71                "critical volume",
 72            ),
 73        )
 74
 75        super().__init__(particulator=particulator)
 76        if scipy_solver:
 77            scipy_ode_condensation_solver.patch_particulator(self.particulator)
 78
 79        self.output_attributes = {
 80            "volume": tuple([] for _ in range(self.particulator.n_sd)),
 81            "dry volume": tuple([] for _ in range(self.particulator.n_sd)),
 82            "critical saturation": tuple([] for _ in range(self.particulator.n_sd)),
 83            "equilibrium saturation": tuple([] for _ in range(self.particulator.n_sd)),
 84            "critical volume": tuple([] for _ in range(self.particulator.n_sd)),
 85            "multiplicity": tuple([] for _ in range(self.particulator.n_sd)),
 86        }
 87        self.settings = settings
 88
 89        self.__sanity_checks(attributes, volume)
 90
 91    def __sanity_checks(self, attributes, volume):
 92        for attribute in attributes.values():
 93            assert attribute.shape[0] == self.particulator.n_sd
 94        np.testing.assert_approx_equal(
 95            sum(attributes["multiplicity"]) / volume,
 96            sum(
 97                mode.norm_factor
 98                for mode in self.settings.aerosol_modes_by_kappa.values()
 99            ),
100            significant=4,
101        )
102
103    def _save(self, output):
104        for key, attr in self.output_attributes.items():
105            attr_data = self.particulator.attributes[key].to_ndarray()
106            for drop_id in range(self.particulator.n_sd):
107                attr[drop_id].append(attr_data[drop_id])
108        super()._save(output)
109
110    def run(self):
111        output_products = super()._run(
112            self.settings.nt, self.settings.steps_per_output_interval
113        )
114        return {"products": output_products, "attributes": self.output_attributes}