PySDM_examples.seeding.simulation

  1import numpy as np
  2
  3from PySDM_examples.seeding.settings import Settings
  4
  5from PySDM import Particulator
  6from PySDM.backends import CPU
  7from PySDM.environments import Parcel
  8from PySDM.dynamics import Condensation, AmbientThermodynamics, Coalescence, Seeding
  9from PySDM.dynamics.collisions.collision_kernels import Geometric
 10from PySDM.initialisation.sampling.spectral_sampling import ConstantMultiplicity
 11from PySDM import products
 12from PySDM.physics import si
 13
 14
 15class Simulation:
 16    def __init__(self, settings: Settings):
 17        environment = Parcel(
 18            dt=settings.timestep,
 19            mass_of_dry_air=settings.mass_of_dry_air,
 20            w=settings.updraft,
 21            initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
 22            p0=settings.initial_total_pressure,
 23            T0=settings.initial_temperature,
 24            backend=CPU(
 25                formulae=settings.formulae, override_jit_flags={"parallel": False}
 26            ),
 27        )
 28        r_dry, n_in_dv = ConstantMultiplicity(
 29            settings.initial_aerosol_dry_radii
 30        ).sample_deterministic(n_sd=settings.n_sd_initial, backend=environment.backend)
 31        attributes = environment.init_attributes(
 32            n_in_dv=n_in_dv, kappa=settings.initial_aerosol_kappa, r_dry=r_dry
 33        )
 34        self.particulator = Particulator(
 35            environment=environment,
 36            n_sd=settings.n_sd_seeding + settings.n_sd_initial,
 37            dynamics=[
 38                AmbientThermodynamics(),
 39                Condensation(),
 40            ]
 41            + (
 42                []
 43                if not settings.enable_collisions
 44                else [
 45                    Coalescence(collision_kernel=Geometric()),
 46                ]
 47            )
 48            + [
 49                Seeding(
 50                    **{
 51                        k: getattr(settings, k)
 52                        for k in (
 53                            "super_droplet_injection_rate",
 54                            "seeded_particle_multiplicity",
 55                            "seeded_particle_extensive_attributes",
 56                        )
 57                    }
 58                ),
 59            ],
 60            attributes={
 61                k: np.pad(
 62                    array=v,
 63                    pad_width=(0, settings.n_sd_seeding),
 64                    mode="constant",
 65                    constant_values=np.nan if k == "multiplicity" else 0,
 66                )
 67                for k, v in attributes.items()
 68            },
 69            products=(
 70                products.SuperDropletCountPerGridbox(name="sd_count"),
 71                products.Time(),
 72                products.WaterMixingRatio(
 73                    radius_range=(settings.rain_water_radius_threshold, np.inf),
 74                    name="rain water mixing ratio",
 75                ),
 76                products.EffectiveRadius(
 77                    name="r_eff",
 78                    unit="um",
 79                    radius_range=(0.5 * si.um, 25 * si.um),
 80                ),
 81                products.ParticleConcentration(
 82                    name="n_drop",
 83                    unit="cm^-3",
 84                    radius_range=(0.5 * si.um, 25 * si.um),
 85                ),
 86            ),
 87        )
 88        self.n_steps = int(settings.t_max // settings.timestep)
 89
 90    def run(self):
 91        output = {
 92            "attributes": {"water mass": []},
 93            "products": {key: [] for key in self.particulator.products},
 94        }
 95        for step in range(self.n_steps + 1):
 96            if step != 0:
 97                self.particulator.run(steps=1)
 98            for key, attr in output["attributes"].items():
 99                data = self.particulator.attributes[key].to_ndarray(raw=True)
100                data[data == 0] = np.nan
101                attr.append(data)
102            for key, prod in output["products"].items():
103                value = self.particulator.products[key].get()
104                if not isinstance(value, float):
105                    (value,) = value
106                prod.append(float(value))
107        for out in ("attributes", "products"):
108            for key, val in output[out].items():
109                output[out][key] = np.array(val)
110        return output