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