PySDM_examples.Pyrcel.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(
 34                formulae=settings.formulae, override_jit_flags={"parallel": False}
 35            ),
 36        )
 37        n_sd = sum(settings.n_sd_per_mode)
 38        volume = environment.mass_of_dry_air / settings.initial_air_density
 39        attributes = {
 40            k: np.empty(0)
 41            for k in ("dry volume", "kappa times dry volume", "multiplicity")
 42        }
 43        for i, (kappa, spectrum) in enumerate(settings.aerosol_modes_by_kappa.items()):
 44            sampling = ConstantMultiplicity(spectrum)
 45            r_dry, n_per_volume = sampling.sample_deterministic(
 46                settings.n_sd_per_mode[i]
 47            )
 48            v_dry = settings.formulae.trivia.volume(radius=r_dry)
 49            attributes["multiplicity"] = np.append(
 50                attributes["multiplicity"], n_per_volume * volume
 51            )
 52            attributes["dry volume"] = np.append(attributes["dry volume"], v_dry)
 53            attributes["kappa times dry volume"] = np.append(
 54                attributes["kappa times dry volume"], v_dry * kappa
 55            )
 56        r_wet = equilibrate_wet_radii(
 57            r_dry=settings.formulae.trivia.radius(volume=attributes["dry volume"]),
 58            environment=environment,
 59            kappa_times_dry_volume=attributes["kappa times dry volume"],
 60        )
 61        attributes["volume"] = settings.formulae.trivia.volume(radius=r_wet)
 62
 63        particulator_kwargs = {
 64            "n_sd": n_sd,
 65            "environment": environment,
 66            "dynamics": (
 67                AmbientThermodynamics(),
 68                Condensation(rtol_thd=rtol_thd, rtol_x=rtol_x),
 69            ),
 70            "attributes": attributes,
 71            "products": products,
 72        }
 73
 74        if additional_attributes is not None:
 75            particulator_kwargs["requested_attributes"] = additional_attributes
 76
 77        super().__init__(particulator=Particulator(**particulator_kwargs))
 78
 79        if scipy_solver:
 80            scipy_ode_condensation_solver.patch_particulator(self.particulator)
 81
 82        self.output_attributes = {
 83            attr: tuple([] for _ in range(self.particulator.n_sd))
 84            for attr in ["volume"]
 85            + (list(additional_attributes) if additional_attributes is not None else [])
 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, observers=()):
111        for observer in observers:
112            self.particulator.observers.append(observer)
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}