PySDM_examples.Luettmer_et_al_2026.simulation

  1import numpy as np
  2import PySDM.products as PySDM_products
  3from PySDM import Particulator
  4from PySDM.dynamics import (
  5    AmbientThermodynamics,
  6    Condensation,
  7    Freezing,
  8    VapourDepositionOnIce,
  9)
 10from PySDM.environments import Parcel
 11from PySDM.initialisation import discretise_multiplicities
 12from PySDM.initialisation.hygroscopic_equilibrium import equilibrate_dry_radii
 13
 14
 15class Simulation:
 16    def __init__(self, settings):
 17
 18        self.dt = settings.dt
 19
 20        formulae = settings.formulae
 21
 22        self.silent = settings.silent
 23
 24        env = Parcel(
 25            mixed_phase=True,
 26            dt=self.dt,
 27            mass_of_dry_air=settings.mass_of_dry_air,
 28            p0=settings.initial_pressure,
 29            initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
 30            T0=settings.initial_temperature,
 31            w=settings.w_updraft,
 32            backend=settings.backend,
 33        )
 34
 35        self.n_sd = settings.n_sd
 36        self.multiplicities = discretise_multiplicities(
 37            settings.specific_concentration * env.mass_of_dry_air
 38        )
 39        self.r_wet = settings.r_wet
 40
 41        kappa = np.full_like(settings.r_wet, settings.kappa)
 42
 43        self.r_dry = equilibrate_dry_radii(
 44            r_wet=self.r_wet,
 45            environment=env,
 46            kappa=kappa,
 47        )
 48        v_dry = settings.formulae.trivia.volume(radius=self.r_dry)
 49        self.initial_mass = formulae.particle_shape_and_density.radius_to_mass(
 50            self.r_wet
 51        )
 52
 53        attributes = {
 54            "multiplicity": self.multiplicities,
 55            "dry volume": v_dry,
 56            "kappa times dry volume": kappa * v_dry,
 57            "signed water mass": self.initial_mass,
 58        }
 59
 60        products = (
 61            PySDM_products.ParcelDisplacement(name="z"),
 62            PySDM_products.Time(name="t"),
 63            PySDM_products.AmbientRelativeHumidity(name="RH", unit="%"),
 64            PySDM_products.AmbientRelativeHumidity(name="RH_ice", unit="%"),
 65            PySDM_products.AmbientTemperature(name="T"),
 66            PySDM_products.AmbientPressure(name="p", unit="hPa"),
 67            PySDM_products.WaterMixingRatio(name="LWC", radius_range=(0, np.inf)),
 68            PySDM_products.WaterMixingRatio(name="IWC", radius_range=(-np.inf, 0)),
 69            PySDM_products.AmbientWaterVapourMixingRatio(
 70                name="qv", var="water_vapour_mixing_ratio"
 71            ),
 72            PySDM_products.ParticleSpecificConcentration(
 73                name="ns", radius_range=(0, np.inf), unit="kg^-1"
 74            ),
 75            PySDM_products.ParticleSpecificConcentration(
 76                name="ni", radius_range=(-np.inf, 0), unit="kg^-1"
 77            ),
 78            PySDM_products.MeanRadius(name="rs", radius_range=(0, np.inf)),
 79            PySDM_products.MeanRadius(name="ri", radius_range=(-np.inf, 0)),
 80        )
 81        self.output_product_list = [
 82            "z",
 83            "t",
 84            "T",
 85            "p",
 86            "RH",
 87            "RH_ice",
 88            "LWC",
 89            "IWC",
 90            "qv",
 91            "ns",
 92            "ni",
 93            "rs",
 94            "ri",
 95        ]
 96
 97        self.particulator = Particulator(
 98            n_sd=settings.n_sd,
 99            environment=env,
100            dynamics=(
101                [
102                    AmbientThermodynamics(),
103                    Condensation(adaptive=True),
104                ]
105                + (
106                    [VapourDepositionOnIce(adaptive=True)]
107                    if settings.deposition_enable
108                    else []
109                )
110                + [
111                    Freezing(
112                        homogeneous_freezing=settings.hom_freezing_type,
113                        immersion_freezing=None,
114                    )
115                ]
116            ),
117            attributes=attributes,
118            products=products,
119            requested_attributes=(
120                "temperature of last freezing",
121                "supersaturation of last freezing",
122                "radius",
123                "wet to critical volume ratio",
124            ),
125        )
126
127        self.n_output = settings.n_output
128        if settings.n_output == 1:
129            self.n_substeps = 1
130        else:
131            self.n_substeps = int(self.n_output / self.dt)
132        self.t_max_duration = settings.t_max_duration
133
134    def save(self, output):
135        cell_id = 0
136
137        for key in self.output_product_list:
138            if key == "t":
139                output[key].append(self.particulator.products[key].get())
140            else:
141                output[key].append(self.particulator.products[key].get()[cell_id])
142
143        output["T_frz"] = self.particulator.attributes[
144            "temperature of last freezing"
145        ].data.tolist()
146        output["RHi_frz"] = self.particulator.attributes[
147            "supersaturation of last freezing"
148        ].data.tolist()
149        if not output["radius"]:
150            output["radius"] = self.particulator.attributes["radius"].data.tolist()
151            output["multiplicity"] = self.particulator.attributes[
152                "multiplicity"
153            ].data.tolist()
154
155    def run(self):
156
157        if not self.silent:
158            print("Starting simulation...")
159
160        output = {
161            "T_frz": [],
162            "RHi_frz": [],
163            "radius": [],
164            "multiplicity": [],
165        }
166        for key in self.output_product_list:
167            output[key] = []
168
169        self.save(output)
170
171        while True:
172
173            self.particulator.run(self.n_substeps)
174            self.save(output)
175
176            w_cr_v_ratio = self.particulator.attributes[
177                "wet to critical volume ratio"
178            ].data
179            sig_mass = self.particulator.attributes["signed water mass"].data
180            frozen = sig_mass < 0
181            unactivated = w_cr_v_ratio < 1
182            if any(frozen) and all(np.logical_or(frozen, unactivated)):
183                if not self.silent:
184                    print("all particles frozen or evaporated")
185                # Assert for water saturation
186                test_water_saturation = np.asarray(output["RH"])
187                # Sort out times before CCN activation & after first occurence of ice
188                test_water_saturation = np.where(
189                    np.asarray(output["rs"]) < 1e-6, 100.0, test_water_saturation
190                )
191                test_water_saturation = np.where(
192                    np.asarray(output["IWC"]) > 0.0, 100.0, test_water_saturation
193                )
194                if np.allclose(test_water_saturation, 100.0, rtol=5e-2) is False:
195                    print(
196                        "Warning: water saturation is too high outside "
197                        "of activation and mixed-phase environment"
198                    )
199
200                break
201            if output["t"][-1] >= self.t_max_duration * self.dt:
202                print("time exceeded")
203                break
204
205        return output