PySDM_examples.Shipway_and_Hill_2012.simulation

  1from collections import namedtuple
  2from typing import List, Any
  3
  4import numpy as np
  5from PySDM_examples.Shipway_and_Hill_2012.mpdata_1d import MPDATA_1D
  6
  7import PySDM.products as PySDM_products
  8from PySDM import Particulator
  9from PySDM.backends import CPU
 10from PySDM.dynamics import (
 11    AmbientThermodynamics,
 12    Coalescence,
 13    Condensation,
 14    Displacement,
 15    EulerianAdvection,
 16)
 17from PySDM.environments.kinematic_1d import Kinematic1D
 18from PySDM.impl.mesh import Mesh
 19from PySDM.initialisation.sampling import spatial_sampling, spectral_sampling
 20
 21
 22class Simulation:
 23    def __init__(self, settings, backend=CPU):
 24        self.nt = settings.nt
 25        self.z0 = -settings.particle_reservoir_depth
 26        self.save_spec_and_attr_times = settings.save_spec_and_attr_times
 27        self.number_of_bins = settings.number_of_bins
 28
 29        self.particulator = None
 30        self.output_attributes = None
 31        self.output_products = None
 32
 33        self.mesh = Mesh(
 34            grid=(settings.nz,),
 35            size=(settings.z_max + settings.particle_reservoir_depth,),
 36        )
 37
 38        def zZ_to_z_above_reservoir(zZ):
 39            z_above_reservoir = zZ * (settings.nz * settings.dz) + self.z0
 40            return z_above_reservoir
 41
 42        mpdata = MPDATA_1D(
 43            nz=settings.nz,
 44            dt=settings.dt,
 45            mpdata_settings=settings.mpdata_settings,
 46            advector_of_t=lambda t: settings.rho_times_w(t) * settings.dt / settings.dz,
 47            advectee_of_zZ_at_t0=lambda zZ: settings.water_vapour_mixing_ratio(
 48                zZ_to_z_above_reservoir(zZ)
 49            ),
 50            g_factor_of_zZ=lambda zZ: settings.rhod(zZ_to_z_above_reservoir(zZ)),
 51        )
 52
 53        env = Kinematic1D(
 54            dt=settings.dt,
 55            mesh=self.mesh,
 56            thd_of_z=settings.thd,
 57            rhod_of_z=settings.rhod,
 58            z0=-settings.particle_reservoir_depth,
 59            backend=backend(formulae=settings.formulae, n_dims=1),
 60            solvers=mpdata,
 61        )
 62
 63        _extra_nz = settings.particle_reservoir_depth // settings.dz
 64        _z_vec = settings.dz * np.linspace(
 65            -_extra_nz, settings.nz - _extra_nz, settings.nz + 1
 66        )
 67        self.g_factor_vec = settings.rhod(_z_vec)
 68
 69        dynamics: List[Any] = [AmbientThermodynamics()]
 70
 71        if settings.enable_condensation:
 72            dynamics.append(
 73                Condensation(
 74                    adaptive=settings.condensation_adaptive,
 75                    rtol_thd=settings.condensation_rtol_thd,
 76                    rtol_x=settings.condensation_rtol_x,
 77                    update_thd=settings.condensation_update_thd,
 78                )
 79            )
 80        dynamics.append(EulerianAdvection())
 81
 82        self.products = []
 83        if settings.precip:
 84            self.add_collision_dynamic(dynamics, settings, self.products)
 85
 86        dynamics.append(
 87            Displacement(
 88                enable_sedimentation=settings.precip,
 89                precipitation_counting_level_index=int(
 90                    settings.particle_reservoir_depth / settings.dz
 91                ),
 92            )
 93        )
 94        self.attributes = env.init_attributes(
 95            spatial_discretisation=spatial_sampling.Pseudorandom(),
 96            spectral_discretisation=spectral_sampling.ConstantMultiplicity(
 97                spectrum=settings.wet_radius_spectrum_per_mass_of_dry_air
 98            ),
 99            kappa=settings.kappa,
100            collisions_only=not settings.enable_condensation,
101            z_part=settings.z_part,
102            n_sd=settings.n_sd,
103        )
104        self.products += [
105            PySDM_products.WaterMixingRatio(
106                name="cloud water mixing ratio",
107                unit="g/kg",
108                radius_range=settings.cloud_water_radius_range,
109            ),
110            PySDM_products.WaterMixingRatio(
111                name="rain water mixing ratio",
112                unit="g/kg",
113                radius_range=settings.rain_water_radius_range,
114            ),
115            PySDM_products.AmbientDryAirDensity(name="rhod"),
116            PySDM_products.AmbientDryAirPotentialTemperature(name="thd"),
117            PySDM_products.ParticleSizeSpectrumPerVolume(
118                name="wet spectrum", radius_bins_edges=settings.r_bins_edges
119            ),
120            PySDM_products.ParticleConcentration(
121                name="nc", radius_range=settings.cloud_water_radius_range
122            ),
123            PySDM_products.ParticleConcentration(
124                name="nr", radius_range=settings.rain_water_radius_range
125            ),
126            PySDM_products.ParticleConcentration(
127                name="na", radius_range=(0, settings.cloud_water_radius_range[0])
128            ),
129            PySDM_products.MeanRadius(),
130            PySDM_products.EffectiveRadius(
131                radius_range=settings.cloud_water_radius_range
132            ),
133            PySDM_products.SuperDropletCountPerGridbox(),
134            PySDM_products.AveragedTerminalVelocity(
135                name="rain averaged terminal velocity",
136                radius_range=settings.rain_water_radius_range,
137            ),
138            PySDM_products.AmbientRelativeHumidity(name="RH", unit="%"),
139            PySDM_products.AmbientPressure(name="p"),
140            PySDM_products.AmbientTemperature(name="T"),
141            PySDM_products.AmbientWaterVapourMixingRatio(
142                name="water_vapour_mixing_ratio"
143            ),
144        ]
145        if settings.enable_condensation:
146            self.products.extend(
147                [
148                    PySDM_products.RipeningRate(name="ripening"),
149                    PySDM_products.ActivatingRate(name="activating"),
150                    PySDM_products.DeactivatingRate(name="deactivating"),
151                    PySDM_products.PeakSaturation(),
152                    PySDM_products.ParticleSizeSpectrumPerVolume(
153                        name="dry spectrum",
154                        radius_bins_edges=settings.r_bins_edges_dry,
155                        dry=True,
156                    ),
157                ]
158            )
159        if settings.precip:
160            self.products.extend(
161                [
162                    PySDM_products.CollisionRatePerGridbox(
163                        name="collision_rate",
164                    ),
165                    PySDM_products.CollisionRateDeficitPerGridbox(
166                        name="collision_deficit",
167                    ),
168                    PySDM_products.CoalescenceRatePerGridbox(
169                        name="coalescence_rate",
170                    ),
171                    PySDM_products.SurfacePrecipitation(),
172                ]
173            )
174        self.particulator = Particulator(
175            n_sd=settings.n_sd,
176            environment=env,
177            dynamics=dynamics,
178            attributes=self.attributes,
179            products=tuple(self.products),
180        )
181
182        self.output_attributes = {
183            "cell origin": [],
184            "position in cell": [],
185            "radius": [],
186            "multiplicity": [],
187        }
188        self.output_products = {}
189        for k, v in self.particulator.products.items():
190            if len(v.shape) == 0:
191                self.output_products[k] = np.zeros(self.nt + 1)
192            elif len(v.shape) == 1:
193                self.output_products[k] = np.zeros((self.mesh.grid[-1], self.nt + 1))
194            elif len(v.shape) == 2:
195                number_of_time_sections = len(self.save_spec_and_attr_times)
196                self.output_products[k] = np.zeros(
197                    (self.mesh.grid[-1], self.number_of_bins, number_of_time_sections)
198                )
199
200    @staticmethod
201    def add_collision_dynamic(dynamics, settings, _):
202        dynamics.append(
203            Coalescence(
204                collision_kernel=settings.collision_kernel,
205                adaptive=settings.coalescence_adaptive,
206            )
207        )
208
209    def save_scalar(self, step):
210        for k, v in self.particulator.products.items():
211            if len(v.shape) > 1:
212                continue
213            if len(v.shape) == 1:
214                self.output_products[k][:, step] = v.get()
215            else:
216                self.output_products[k][step] = v.get()
217
218    def save_spectrum(self, index):
219        for k, v in self.particulator.products.items():
220            if len(v.shape) == 2:
221                self.output_products[k][:, :, index] = v.get()
222
223    def save_attributes(self):
224        for k, v in self.output_attributes.items():
225            v.append(self.particulator.attributes[k].to_ndarray())
226
227    def save(self, step):
228        self.save_scalar(step)
229        time = step * self.particulator.dt
230        if len(self.save_spec_and_attr_times) > 0 and (
231            np.min(
232                np.abs(
233                    np.ones_like(self.save_spec_and_attr_times) * time
234                    - np.array(self.save_spec_and_attr_times)
235                )
236            )
237            < 0.1
238        ):
239            save_index = np.argmin(
240                np.abs(
241                    np.ones_like(self.save_spec_and_attr_times) * time
242                    - np.array(self.save_spec_and_attr_times)
243                )
244            )
245            self.save_spectrum(save_index)
246            self.save_attributes()
247
248    def run(self):
249        mesh = self.particulator.mesh
250
251        assert "t" not in self.output_products and "z" not in self.output_products
252        self.output_products["t"] = np.linspace(
253            0, self.nt * self.particulator.dt, self.nt + 1, endpoint=True
254        )
255        self.output_products["z"] = np.linspace(
256            self.z0 + mesh.dz / 2,
257            self.z0 + (mesh.grid[-1] - 1 / 2) * mesh.dz,
258            mesh.grid[-1],
259            endpoint=True,
260        )
261
262        self.save(0)
263        for step in range(self.nt):
264            mpdata = self.particulator.environment.solvers
265            mpdata.update_advector_field()
266            if "Displacement" in self.particulator.dynamics:
267                self.particulator.dynamics["Displacement"].upload_courant_field(
268                    (mpdata.advector / self.g_factor_vec,)
269                )
270            self.particulator.run(steps=1)
271            self.save(step + 1)
272
273        Outputs = namedtuple("Outputs", "products attributes")
274        output_results = Outputs(self.output_products, self.output_attributes)
275        return output_results