PySDM_examples.featured.pyrcel_docs_testcase.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        *,
 19        products=None,
 20        scipy_solver=False,
 21        rtol_thd=1e-10,
 22        rtol_x=1e-10,
 23        mass_of_dry_air=44 * si.kg,
 24        additional_attributes=None,
 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=mass_of_dry_air,
 33            backend=CPU(formulae=settings.formulae),
 34        )
 35        n_sd = sum(settings.n_sd_per_mode)
 36        volume = environment.mass_of_dry_air / settings.initial_air_density
 37        attributes = {
 38            k: np.empty(0)
 39            for k in ("dry volume", "kappa times dry volume", "multiplicity")
 40        }
 41        for i, (kappa, spectrum) in enumerate(settings.aerosol_modes_by_kappa.items()):
 42            sampling = ConstantMultiplicity(spectrum)
 43            r_dry, n_per_volume = sampling.sample_deterministic(
 44                settings.n_sd_per_mode[i]
 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_kwargs = {
 62            "n_sd": n_sd,
 63            "environment": environment,
 64            "dynamics": (
 65                AmbientThermodynamics(),
 66                Condensation(rtol_thd=rtol_thd, rtol_x=rtol_x),
 67            ),
 68            "attributes": attributes,
 69            "products": products,
 70        }
 71
 72        if additional_attributes is not None:
 73            particulator_kwargs["requested_attributes"] = additional_attributes
 74
 75        super().__init__(particulator=Particulator(**particulator_kwargs))
 76
 77        if scipy_solver:
 78            scipy_ode_condensation_solver.patch_particulator(self.particulator)
 79
 80        self.output_attributes = {
 81            attr: tuple([] for _ in range(self.particulator.n_sd))
 82            for attr in ["volume"]
 83            + (list(additional_attributes) if additional_attributes is not None else [])
 84        }
 85        self.settings = settings
 86
 87        self.__sanity_checks(attributes, volume)
 88
 89    def __sanity_checks(self, attributes, volume):
 90        for attribute in attributes.values():
 91            assert attribute.shape[0] == self.particulator.n_sd
 92        np.testing.assert_approx_equal(
 93            sum(attributes["multiplicity"]) / volume,
 94            sum(
 95                mode.norm_factor
 96                for mode in self.settings.aerosol_modes_by_kappa.values()
 97            ),
 98            significant=4,
 99        )
100
101    def _save(self, output):
102        for key, attr in self.output_attributes.items():
103            attr_data = self.particulator.attributes[key].to_ndarray()
104            for drop_id in range(self.particulator.n_sd):
105                attr[drop_id].append(attr_data[drop_id])
106        super()._save(output)
107
108    def run(self, observers=()):
109        for observer in observers:
110            self.particulator.observers.append(observer)
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}