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