PySDM_examples.Arabas_and_Pawlowska_2011.simulation

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