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