PySDM_examples.Jensen_and_Nugent_2017.simulation

  1import numpy as np
  2from PySDM_examples.utils import BasicSimulation
  3from PySDM_examples.Jensen_and_Nugent_2017.settings import Settings
  4from PySDM_examples.Jensen_and_Nugent_2017 import table_3
  5from PySDM import Particulator
  6from PySDM.physics import si
  7from PySDM.backends import CPU
  8from PySDM.products import (
  9    PeakSaturation,
 10    ParcelDisplacement,
 11    Time,
 12    ActivatedMeanRadius,
 13    RadiusStandardDeviation,
 14)
 15from PySDM.environments import Parcel
 16from PySDM.dynamics import Condensation, AmbientThermodynamics, Coalescence
 17from PySDM.dynamics.collisions.collision_kernels import Geometric
 18from PySDM.initialisation.sampling.spectral_sampling import Logarithmic
 19
 20# note: 100 in caption of Table 1
 21N_SD_NON_GCCN = 100
 22
 23
 24class Simulation(BasicSimulation):
 25    def __init__(
 26        self,
 27        settings: Settings,
 28        gccn: bool = False,
 29        gravitational_coalsecence: bool = False,
 30    ):
 31
 32        n_gccn = np.count_nonzero(table_3.NA) if gccn else 0
 33        environment = Parcel(
 34            dt=settings.dt,
 35            mass_of_dry_air=666 * si.kg,
 36            p0=settings.p0,
 37            initial_relative_humidity=settings.RH0,
 38            T0=settings.T0,
 39            w=settings.vertical_velocity,
 40            z0=settings.z0,
 41            backend=CPU(
 42                formulae=settings.formulae,
 43                override_jit_flags={"parallel": False},
 44            ),
 45        )
 46
 47        self.r_dry, n_in_unit_volume = Logarithmic(
 48            spectrum=settings.dry_radii_spectrum,
 49        ).sample_deterministic(N_SD_NON_GCCN)
 50
 51        if gccn:
 52            nonzero_concentration_mask = np.nonzero(table_3.NA)
 53            self.r_dry = np.concatenate(
 54                [self.r_dry, table_3.RD[nonzero_concentration_mask]]
 55            )
 56            n_in_unit_volume = np.concatenate(
 57                [n_in_unit_volume, table_3.NA[nonzero_concentration_mask]]
 58            )  # TODO #1266: check which temp, pres, RH assumed in the paper for NA???
 59
 60        pd0 = settings.formulae.trivia.p_d(
 61            settings.p0,
 62            settings.formulae.trivia.water_vapour_mixing_ratio(
 63                settings.p0,
 64                settings.RH0,
 65                settings.formulae.saturation_vapour_pressure.pvs_water(settings.T0),
 66            ),
 67        )
 68        rhod0 = settings.formulae.state_variable_triplet.rhod_of_pd_T(pd0, settings.T0)
 69
 70        attributes = environment.init_attributes(
 71            n_in_dv=n_in_unit_volume * environment.mass_of_dry_air / rhod0,
 72            kappa=settings.kappa,
 73            r_dry=self.r_dry,
 74        )
 75
 76        super().__init__(
 77            Particulator(
 78                n_sd=N_SD_NON_GCCN + n_gccn,
 79                environment=environment,
 80                dynamics=[
 81                    # TODO #1266: order matters here, but error message is not saying it!
 82                    AmbientThermodynamics(),
 83                    Condensation(),
 84                ]
 85                + (
 86                    []
 87                    if not gravitational_coalsecence
 88                    else [Coalescence(collision_kernel=Geometric())]
 89                ),
 90                attributes=attributes,
 91                products=(
 92                    PeakSaturation(name="S_max"),
 93                    ParcelDisplacement(name="z"),
 94                    Time(name="t"),
 95                    ActivatedMeanRadius(
 96                        name="r_mean_act", count_activated=True, count_unactivated=False
 97                    ),
 98                    RadiusStandardDeviation(
 99                        name="r_std_act", count_activated=True, count_unactivated=False
100                    ),
101                ),
102                requested_attributes=(
103                    additional_derived_attributes := (
104                        "radius",
105                        "equilibrium saturation",
106                    )
107                ),
108            )
109        )
110
111        # TODO #1266: copied from G & P 2023
112        self.output_attributes = {
113            attr: tuple([] for _ in range(self.particulator.n_sd))
114            for attr in additional_derived_attributes
115        }
116
117    def run(
118        self, *, n_steps: int = 2250, steps_per_output_interval: int = 10
119    ):  # TODO #1266: essentially copied from G & P 2023
120        output_products = super()._run(
121            nt=n_steps, steps_per_output_interval=steps_per_output_interval
122        )
123        return {"products": output_products, "attributes": self.output_attributes}
124
125    def _save(self, output):  # TODO #1266: copied from G&P 2023
126        for key, attr in self.output_attributes.items():
127            attr_data = self.particulator.attributes[key].to_ndarray()
128            for drop_id in range(self.particulator.n_sd):
129                attr[drop_id].append(attr_data[drop_id])
130        super()._save(output)