PySDM_examples.Pyrcel.tutorial_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 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, settings, products=None, scipy_solver=False, rtol_thd=1e-10, rtol_x=1e-10
17    ):
18        env = Parcel(
19            dt=settings.timestep,
20            p0=settings.initial_pressure,
21            initial_water_vapour_mixing_ratio=settings.initial_vapour_mixing_ratio,
22            T0=settings.initial_temperature,
23            w=settings.vertical_velocity,
24            mass_of_dry_air=44 * si.kg,
25        )
26        n_sd = sum(settings.n_sd_per_mode)
27        builder = Builder(n_sd=n_sd, backend=CPU(formulae=settings.formulae))
28        builder.set_environment(env)
29        builder.add_dynamic(AmbientThermodynamics())
30        builder.add_dynamic(Condensation(rtol_thd=rtol_thd, rtol_x=rtol_x))
31
32        volume = env.mass_of_dry_air / settings.initial_air_density
33        attributes = {
34            k: np.empty(0) for k in ("dry volume", "kappa times dry volume", "n")
35        }
36        for i, (kappa, spectrum) in enumerate(settings.aerosol_modes_by_kappa.items()):
37            sampling = ConstantMultiplicity(spectrum)
38            r_dry, n_per_volume = sampling.sample(settings.n_sd_per_mode[i])
39            v_dry = settings.formulae.trivia.volume(radius=r_dry)
40            attributes["n"] = np.append(attributes["n"], n_per_volume * volume)
41            attributes["dry volume"] = np.append(attributes["dry volume"], v_dry)
42            attributes["kappa times dry volume"] = np.append(
43                attributes["kappa times dry volume"], v_dry * kappa
44            )
45        r_wet = equilibrate_wet_radii(
46            r_dry=settings.formulae.trivia.radius(volume=attributes["dry volume"]),
47            environment=env,
48            kappa_times_dry_volume=attributes["kappa times dry volume"],
49        )
50        attributes["volume"] = settings.formulae.trivia.volume(radius=r_wet)
51
52        super().__init__(
53            particulator=builder.build(attributes=attributes, products=products)
54        )
55        if scipy_solver:
56            scipy_ode_condensation_solver.patch_particulator(self.particulator)
57
58        self.output_attributes = {
59            "volume": tuple([] for _ in range(self.particulator.n_sd))
60        }
61        self.settings = settings
62
63        self.__sanity_checks(attributes, volume)
64
65    def __sanity_checks(self, attributes, volume):
66        for attribute in attributes.values():
67            assert attribute.shape[0] == self.particulator.n_sd
68        np.testing.assert_approx_equal(
69            sum(attributes["multiplicity"]) / volume,
70            sum(
71                mode.norm_factor
72                for mode in self.settings.aerosol_modes_by_kappa.values()
73            ),
74            significant=4,
75        )
76
77    def _save(self, output):
78        for key, attr in self.output_attributes.items():
79            attr_data = self.particulator.attributes[key].to_ndarray()
80            for drop_id in range(self.particulator.n_sd):
81                attr[drop_id].append(attr_data[drop_id])
82        super()._save(output)
83
84    def run(self, observers=()):
85        for observer in observers:
86            self.particulator.observers.append(observer)
87        output_products = super()._run(
88            self.settings.nt, self.settings.steps_per_output_interval
89        )
90        return {"products": output_products, "attributes": self.output_attributes}