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 Builder
  6from PySDM.physics import si
  7from PySDM.backends import CPU
  8from PySDM.products import (
  9    PeakSupersaturation,
 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        const = settings.formulae.constants
 32        pvs_water = settings.formulae.saturation_vapour_pressure.pvs_water
 33        initial_water_vapour_mixing_ratio = const.eps / (
 34            settings.p0 / settings.RH0 / pvs_water(settings.T0) - 1
 35        )
 36
 37        n_gccn = np.count_nonzero(table_3.NA) if gccn else 0
 38
 39        builder = Builder(
 40            n_sd=N_SD_NON_GCCN + n_gccn,
 41            backend=CPU(
 42                formulae=settings.formulae, override_jit_flags={"parallel": False}
 43            ),
 44            environment=Parcel(
 45                dt=settings.dt,
 46                mass_of_dry_air=666 * si.kg,
 47                p0=settings.p0,
 48                initial_water_vapour_mixing_ratio=initial_water_vapour_mixing_ratio,
 49                T0=settings.T0,
 50                w=settings.vertical_velocity,
 51                z0=settings.z0,
 52            ),
 53        )
 54
 55        additional_derived_attributes = ("radius", "equilibrium supersaturation")
 56        for additional_derived_attribute in additional_derived_attributes:
 57            builder.request_attribute(additional_derived_attribute)
 58
 59        builder.add_dynamic(
 60            AmbientThermodynamics()
 61        )  # TODO #1266: order matters here, but error message is not saying it!
 62        builder.add_dynamic(Condensation())
 63        if gravitational_coalsecence:
 64            builder.add_dynamic(Coalescence(collision_kernel=Geometric()))
 65
 66        self.r_dry, n_in_unit_volume = Logarithmic(
 67            spectrum=settings.dry_radii_spectrum,
 68        ).sample(builder.particulator.n_sd - n_gccn)
 69
 70        if gccn:
 71            nonzero_concentration_mask = np.nonzero(table_3.NA)
 72            self.r_dry = np.concatenate(
 73                [self.r_dry, table_3.RD[nonzero_concentration_mask]]
 74            )
 75            n_in_unit_volume = np.concatenate(
 76                [n_in_unit_volume, table_3.NA[nonzero_concentration_mask]]
 77            )  # TODO #1266: check which temp, pres, RH assumed in the paper for NA???
 78
 79        pd0 = settings.formulae.trivia.p_d(
 80            settings.p0, initial_water_vapour_mixing_ratio
 81        )
 82        rhod0 = settings.formulae.state_variable_triplet.rhod_of_pd_T(pd0, settings.T0)
 83
 84        attributes = builder.particulator.environment.init_attributes(
 85            n_in_dv=n_in_unit_volume
 86            * builder.particulator.environment.mass_of_dry_air
 87            / rhod0,
 88            kappa=settings.kappa,
 89            r_dry=self.r_dry,
 90        )
 91
 92        super().__init__(
 93            builder.build(
 94                attributes=attributes,
 95                products=(
 96                    PeakSupersaturation(name="S_max"),
 97                    ParcelDisplacement(name="z"),
 98                    Time(name="t"),
 99                    ActivatedMeanRadius(
100                        name="r_mean_act", count_activated=True, count_unactivated=False
101                    ),
102                    RadiusStandardDeviation(
103                        name="r_std_act", count_activated=True, count_unactivated=False
104                    ),
105                ),
106            )
107        )
108
109        # TODO #1266: copied from G & P 2023
110        self.output_attributes = {
111            attr: tuple([] for _ in range(self.particulator.n_sd))
112            for attr in additional_derived_attributes
113        }
114
115    def run(
116        self, *, n_steps: int = 2250, steps_per_output_interval: int = 10
117    ):  # TODO #1266: essentially copied from G & P 2023
118        output_products = super()._run(
119            nt=n_steps, steps_per_output_interval=steps_per_output_interval
120        )
121        return {"products": output_products, "attributes": self.output_attributes}
122
123    def _save(self, output):  # TODO #1266: copied from G&P 2023
124        for key, attr in self.output_attributes.items():
125            attr_data = self.particulator.attributes[key].to_ndarray()
126            for drop_id in range(self.particulator.n_sd):
127                attr[drop_id].append(attr_data[drop_id])
128        super()._save(output)