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