PySDM_examples.Lowe_et_al_2019.simulation

  1import numpy as np
  2from PySDM_examples.utils import BasicSimulation
  3
  4import PySDM.products as PySDM_products
  5from PySDM import Particulator
  6from PySDM.dynamics import AmbientThermodynamics, Condensation
  7from PySDM.environments import Parcel
  8from PySDM.initialisation.hygroscopic_equilibrium import equilibrate_wet_radii
  9from PySDM.initialisation.spectra import Sum
 10from PySDM.initialisation.sampling.spectral_sampling import ConstantMultiplicity
 11
 12
 13class Simulation(BasicSimulation):
 14    def __init__(self, settings, products=None):
 15        n_sd = settings.n_sd_per_mode * len(settings.aerosol.modes)
 16        environment = Parcel(
 17            dt=settings.dt,
 18            mass_of_dry_air=settings.mass_of_dry_air,
 19            p0=settings.p0,
 20            initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
 21            T0=settings.T0,
 22            w=settings.w,
 23            backend=settings.backend,
 24        )
 25
 26        attributes = {
 27            "dry volume": np.empty(0),
 28            "dry volume organic": np.empty(0),
 29            "kappa times dry volume": np.empty(0),
 30            "multiplicity": np.ndarray(0),
 31        }
 32        initial_volume = settings.mass_of_dry_air / settings.rho0
 33        for mode in settings.aerosol.modes:
 34            r_dry, n_in_dv = ConstantMultiplicity(
 35                spectrum=mode["spectrum"]
 36            ).sample_deterministic(settings.n_sd_per_mode)
 37            v_dry = settings.formulae.trivia.volume(radius=r_dry)
 38            attributes["multiplicity"] = np.append(
 39                attributes["multiplicity"], n_in_dv * initial_volume
 40            )
 41            attributes["dry volume"] = np.append(attributes["dry volume"], v_dry)
 42            attributes["dry volume organic"] = np.append(
 43                attributes["dry volume organic"], mode["f_org"] * v_dry
 44            )
 45            attributes["kappa times dry volume"] = np.append(
 46                attributes["kappa times dry volume"],
 47                v_dry * mode["kappa"][settings.model],
 48            )
 49        for attribute in attributes.values():
 50            assert attribute.shape[0] == n_sd
 51
 52        np.testing.assert_approx_equal(
 53            np.sum(attributes["multiplicity"]) / initial_volume,
 54            Sum(
 55                tuple(
 56                    settings.aerosol.modes[i]["spectrum"]
 57                    for i in range(len(settings.aerosol.modes))
 58                )
 59            ).norm_factor,
 60            significant=5,
 61        )
 62        r_wet = equilibrate_wet_radii(
 63            r_dry=settings.formulae.trivia.radius(volume=attributes["dry volume"]),
 64            environment=environment,
 65            kappa_times_dry_volume=attributes["kappa times dry volume"],
 66            f_org=attributes["dry volume organic"] / attributes["dry volume"],
 67        )
 68        attributes["volume"] = settings.formulae.trivia.volume(radius=r_wet)
 69
 70        if settings.model == "Constant":
 71            del attributes["dry volume organic"]
 72
 73        products = products or (
 74            PySDM_products.ParcelDisplacement(name="z"),
 75            PySDM_products.Time(name="t"),
 76            PySDM_products.PeakSaturation(name="S_max"),
 77            PySDM_products.AmbientRelativeHumidity(name="RH"),
 78            PySDM_products.ActivatedParticleConcentration(
 79                name="CDNC_cm3",
 80                unit="cm^-3",
 81                count_activated=True,
 82                count_unactivated=False,
 83            ),
 84            PySDM_products.ParticleSizeSpectrumPerVolume(
 85                radius_bins_edges=settings.wet_radius_bins_edges
 86            ),
 87            PySDM_products.ActivableFraction(name="Activated Fraction"),
 88            PySDM_products.WaterMixingRatio(),
 89            PySDM_products.AmbientDryAirDensity(name="rhod"),
 90            PySDM_products.ActivatedEffectiveRadius(
 91                name="reff", count_activated=True, count_unactivated=False
 92            ),
 93            PySDM_products.ParcelLiquidWaterPath(
 94                name="lwp", count_activated=True, count_unactivated=False
 95            ),
 96            PySDM_products.CloudOpticalDepth(name="tau"),
 97            PySDM_products.CloudAlbedo(name="albedo"),
 98        )
 99
100        particulator = Particulator(
101            n_sd=n_sd,
102            dynamics=(AmbientThermodynamics(), Condensation()),
103            environment=environment,
104            attributes=attributes,
105            products=products,
106        )
107        self.settings = settings
108        super().__init__(particulator=particulator)
109
110    def _save_scalars(self, output):
111        for k, v in self.particulator.products.items():
112            if len(v.shape) > 1 or k in ("lwp", "Activated Fraction", "tau", "albedo"):
113                continue
114            value = v.get()
115            if isinstance(value, np.ndarray) and value.size == 1:
116                value = value[0]
117            output[k].append(value)
118
119    def _save_final_timestep_products(self, output):
120        output["spectrum"] = self.particulator.products[
121            "particle size spectrum per volume"
122        ].get()
123
124        for name, args_call in {
125            "Activated Fraction": lambda: {"S_max": np.nanmax(output["S_max"])},
126            "lwp": lambda: {},
127            "tau": lambda: {
128                "effective_radius": output["reff"][-1],
129                "liquid_water_path": output["lwp"][0],
130            },
131            "albedo": lambda: {"optical_depth": output["tau"]},
132        }.items():
133            output[name] = self.particulator.products[name].get(**args_call())
134
135    def run(self):
136        output = {k: [] for k in self.particulator.products}
137        for step in self.settings.output_steps:
138            self.particulator.run(step - self.particulator.n_steps)
139            self._save_scalars(output)
140        self._save_final_timestep_products(output)
141        return output