PySDM_examples.Arabas_and_Pawlowska_2011.simulation

  1import numpy as np
  2
  3from PySDM_examples.utils.basic_simulation import BasicSimulation
  4
  5from PySDM import products, Particulator
  6from PySDM.backends import CPU
  7from PySDM.dynamics import AmbientThermodynamics, Condensation
  8from PySDM.environments import Parcel
  9from PySDM.initialisation import discretise_multiplicities
 10from PySDM.initialisation.hygroscopic_equilibrium import equilibrate_wet_radii
 11from PySDM.initialisation.sampling import spectral_sampling
 12
 13
 14class Simulation(BasicSimulation):
 15    def __init__(
 16        self,
 17        settings,
 18        product_list=None,
 19    ):
 20        self.settings = settings
 21
 22        self.environment = Parcel(
 23            dt=settings.dt,
 24            mass_of_dry_air=settings.mass_of_dry_air,
 25            p0=settings.p0,
 26            initial_relative_humidity=settings.RH0,
 27            T0=settings.T0,
 28            w=settings.w,
 29            backend=CPU(settings.formulae),
 30        )
 31
 32        attributes, self.mode_id = self._make_attributes()
 33
 34        if product_list is None:
 35            product_list = (
 36                products.AmbientRelativeHumidity(name="RH"),
 37                products.Time(name="time"),
 38                products.AmbientTemperature(name="T"),
 39            )
 40
 41        particulator = Particulator(
 42            n_sd=settings.n_sd,
 43            environment=self.environment,
 44            dynamics=(
 45                AmbientThermodynamics(),
 46                Condensation(),
 47            ),
 48            attributes=attributes,
 49            products=product_list,
 50            requested_attributes=("radius",),
 51        )
 52
 53        super().__init__(
 54            particulator=particulator,
 55            output_attributes=["radius"],
 56        )
 57
 58    def _make_attributes(self):
 59        settings = self.settings
 60        formulae = settings.formulae
 61        environment = self.environment
 62
 63        parcel_volume = environment.mass_of_dry_air / settings.initial_air_density
 64
 65        n_components = len(settings.aerosol_modes_by_kappa)
 66
 67        if settings.n_sd % n_components != 0:
 68            raise ValueError(
 69                f"settings.n_sd={settings.n_sd} must be divisible by "
 70                f"n_components={n_components}"
 71            )
 72
 73        n_sd_per_component = settings.n_sd // n_components
 74
 75        dry_volume_parts = []
 76        kappa_vdry_parts = []
 77        multiplicity_parts = []
 78        mode_id_parts = []
 79
 80        for component_id, (kappa, spectrum) in enumerate(
 81            settings.aerosol_modes_by_kappa.items()
 82        ):
 83            r_dry, concentration = spectral_sampling.Logarithmic(
 84                spectrum=spectrum
 85            ).sample_deterministic(n_sd_per_component)
 86
 87            v_dry = formulae.trivia.volume(radius=r_dry)
 88
 89            dry_volume_parts.append(v_dry)
 90            kappa_vdry_parts.append(kappa * v_dry)
 91
 92            multiplicity_parts.append(
 93                discretise_multiplicities(concentration * parcel_volume)
 94            )
 95
 96            mode_id_parts.append(
 97                np.full(n_sd_per_component, component_id, dtype=np.int64)
 98            )
 99
100        dry_volume = np.concatenate(dry_volume_parts)
101        kappa_times_dry_volume = np.concatenate(kappa_vdry_parts)
102        multiplicity = np.concatenate(multiplicity_parts)
103        mode_id = np.concatenate(mode_id_parts)
104
105        r_wet = equilibrate_wet_radii(
106            r_dry=formulae.trivia.radius(volume=dry_volume),
107            environment=environment,
108            kappa_times_dry_volume=kappa_times_dry_volume,
109        )
110
111        attributes = {
112            "multiplicity": multiplicity,
113            "dry volume": dry_volume,
114            "kappa times dry volume": kappa_times_dry_volume,
115            "volume": formulae.trivia.volume(radius=r_wet),
116        }
117
118        self._sanity_check_attributes(attributes, mode_id, parcel_volume)
119
120        return attributes, mode_id
121
122    def _sanity_check_attributes(self, attributes, mode_id, volume):
123        for attribute in attributes.values():
124            assert attribute.shape[0] == self.settings.n_sd
125
126        assert mode_id.shape[0] == self.settings.n_sd
127
128        assert np.all(attributes["multiplicity"] > 0)
129        assert np.all(attributes["dry volume"] > 0)
130        assert np.all(attributes["volume"] >= attributes["dry volume"])
131
132        kappa_eff = attributes["kappa times dry volume"] / attributes["dry volume"]
133
134        for component_id, kappa in enumerate(
135            self.settings.aerosol_modes_by_kappa.keys()
136        ):
137            mask = mode_id == component_id
138            assert np.any(mask)
139            assert np.allclose(kappa_eff[mask], kappa)
140
141        np.testing.assert_allclose(
142            np.sum(attributes["multiplicity"]) / volume,
143            self.settings.total_aerosol_concentration,
144            rtol=1e-2,
145        )
146
147    def run(self):
148        return self._run(
149            nt=self.settings.output_interval * self.settings.output_points,
150            steps_per_output_interval=self.settings.output_interval,
151        )