PySDM_examples.Grabowski_and_Pawlowska_2023.simulation

  1import numpy as np
  2from PySDM_examples.utils import BasicSimulation
  3
  4from PySDM import Builder
  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        builder = Builder(
 22            n_sd=settings.n_sd,
 23            backend=CPU(
 24                formulae=settings.formulae, override_jit_flags={"parallel": False}
 25            ),
 26            environment=Parcel(
 27                dt=settings.timestep,
 28                p0=settings.initial_pressure,
 29                initial_water_vapour_mixing_ratio=settings.initial_vapour_mixing_ratio,
 30                T0=settings.initial_temperature,
 31                w=settings.vertical_velocity,
 32                mass_of_dry_air=44 * si.kg,
 33            ),
 34        )
 35        builder.add_dynamic(AmbientThermodynamics())
 36        builder.add_dynamic(
 37            Condensation(rtol_thd=settings.rtol_thd, rtol_x=settings.rtol_x)
 38        )
 39        for attribute in (
 40            "critical saturation",
 41            "equilibrium saturation",
 42            "critical volume",
 43        ):
 44            builder.request_attribute(attribute)
 45
 46        env = builder.particulator.environment
 47        volume = env.mass_of_dry_air / settings.initial_air_density
 48        attributes = {
 49            k: np.empty(0)
 50            for k in ("dry volume", "kappa times dry volume", "multiplicity")
 51        }
 52
 53        assert len(settings.aerosol_modes_by_kappa.keys()) == 1
 54        kappa = tuple(settings.aerosol_modes_by_kappa.keys())[0]
 55        spectrum = settings.aerosol_modes_by_kappa[kappa]
 56
 57        r_dry, n_per_volume = ConstantMultiplicity(spectrum).sample_deterministic(
 58            settings.n_sd
 59        )
 60        v_dry = settings.formulae.trivia.volume(radius=r_dry)
 61        attributes["multiplicity"] = np.append(
 62            attributes["multiplicity"], n_per_volume * volume
 63        )
 64        attributes["dry volume"] = np.append(attributes["dry volume"], v_dry)
 65        attributes["kappa times dry volume"] = np.append(
 66            attributes["kappa times dry volume"], v_dry * kappa
 67        )
 68        r_wet = equilibrate_wet_radii(
 69            r_dry=settings.formulae.trivia.radius(volume=attributes["dry volume"]),
 70            environment=env,
 71            kappa_times_dry_volume=attributes["kappa times dry volume"],
 72        )
 73        attributes["volume"] = settings.formulae.trivia.volume(radius=r_wet)
 74
 75        super().__init__(
 76            particulator=builder.build(attributes=attributes, products=products)
 77        )
 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}