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