PySDM_examples.Pyrcel.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        *,
 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        n_sd = sum(settings.n_sd_per_mode)
 27        builder = Builder(
 28            n_sd=n_sd,
 29            backend=CPU(
 30                formulae=settings.formulae, override_jit_flags={"parallel": False}
 31            ),
 32            environment=Parcel(
 33                dt=settings.timestep,
 34                p0=settings.initial_pressure,
 35                initial_water_vapour_mixing_ratio=settings.initial_vapour_mixing_ratio,
 36                T0=settings.initial_temperature,
 37                w=settings.vertical_velocity,
 38                mass_of_dry_air=mass_of_dry_air,
 39            ),
 40        )
 41        if additional_attributes is not None:
 42            for attribute in additional_attributes:
 43                builder.request_attribute(attribute)
 44        builder.add_dynamic(AmbientThermodynamics())
 45        builder.add_dynamic(Condensation(rtol_thd=rtol_thd, rtol_x=rtol_x))
 46
 47        volume = (
 48            builder.particulator.environment.mass_of_dry_air
 49            / settings.initial_air_density
 50        )
 51        attributes = {
 52            k: np.empty(0)
 53            for k in ("dry volume", "kappa times dry volume", "multiplicity")
 54        }
 55        for i, (kappa, spectrum) in enumerate(settings.aerosol_modes_by_kappa.items()):
 56            sampling = ConstantMultiplicity(spectrum)
 57            r_dry, n_per_volume = sampling.sample_deterministic(
 58                settings.n_sd_per_mode[i]
 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=builder.particulator.environment,
 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            attr: tuple([] for _ in range(self.particulator.n_sd))
 83            for attr in ["volume"]
 84            + (list(additional_attributes) if additional_attributes is not None else [])
 85        }
 86        self.settings = settings
 87
 88        self.__sanity_checks(attributes, volume)
 89
 90    def __sanity_checks(self, attributes, volume):
 91        for attribute in attributes.values():
 92            assert attribute.shape[0] == self.particulator.n_sd
 93        np.testing.assert_approx_equal(
 94            sum(attributes["multiplicity"]) / volume,
 95            sum(
 96                mode.norm_factor
 97                for mode in self.settings.aerosol_modes_by_kappa.values()
 98            ),
 99            significant=4,
100        )
101
102    def _save(self, output):
103        for key, attr in self.output_attributes.items():
104            attr_data = self.particulator.attributes[key].to_ndarray()
105            for drop_id in range(self.particulator.n_sd):
106                attr[drop_id].append(attr_data[drop_id])
107        super()._save(output)
108
109    def run(self, observers=()):
110        for observer in observers:
111            self.particulator.observers.append(observer)
112        output_products = super()._run(
113            self.settings.nt, self.settings.steps_per_output_interval
114        )
115        return {"products": output_products, "attributes": self.output_attributes}