PySDM_examples.utils.kinematic_2d.simulation

  1import numpy as np
  2from PySDM_examples.utils.kinematic_2d.make_default_product_collection import (
  3    make_default_product_collection,
  4)
  5from PySDM_examples.utils.kinematic_2d.mpdata_2d import MPDATA_2D
  6from PySDM_examples.utils import DummyController
  7
  8from PySDM.backends import CPU
  9from PySDM.builder import Builder
 10from PySDM.dynamics import (
 11    AmbientThermodynamics,
 12    Coalescence,
 13    Collision,
 14    Condensation,
 15    Displacement,
 16    EulerianAdvection,
 17    Freezing,
 18)
 19from PySDM.environments import Kinematic2D
 20from PySDM.initialisation.sampling import spatial_sampling
 21
 22
 23_BACKEND_CACHE = {}
 24
 25
 26def backend_factory(backend_class, formulae):
 27    key = backend_class.__name__ + ":" + str(formulae)
 28    if key not in _BACKEND_CACHE:
 29        _BACKEND_CACHE[key] = backend_class(formulae=formulae)
 30    return _BACKEND_CACHE[key]
 31
 32
 33class Simulation:
 34    def __init__(self, settings, storage, SpinUp, backend_class=CPU):
 35        self.settings = settings
 36        self.storage = storage
 37        self.particulator = None
 38        self.backend_class = backend_class
 39        self.SpinUp = SpinUp
 40
 41    @property
 42    def products(self):
 43        return self.particulator.products
 44
 45    def reinit(self, products=None):
 46        formulae = self.settings.formulae
 47        backend = backend_factory(self.backend_class, formulae=formulae)
 48        environment = Kinematic2D(
 49            dt=self.settings.dt,
 50            grid=self.settings.grid,
 51            size=self.settings.size,
 52            rhod_of=self.settings.rhod_of_zZ,
 53            mixed_phase=self.settings.processes["freezing"],
 54        )
 55        builder = Builder(
 56            n_sd=self.settings.n_sd, backend=backend, environment=environment
 57        )
 58
 59        if products is not None:
 60            products = list(products)
 61        else:
 62            products = make_default_product_collection(self.settings)
 63
 64        if self.settings.processes["fluid advection"]:
 65            builder.add_dynamic(AmbientThermodynamics())
 66        if self.settings.processes["condensation"]:
 67            kwargs = {}
 68            if not self.settings.condensation_adaptive:
 69                kwargs["substeps"] = (self.settings.condensation_substeps,)
 70            condensation = Condensation(
 71                rtol_x=self.settings.condensation_rtol_x,
 72                rtol_thd=self.settings.condensation_rtol_thd,
 73                adaptive=self.settings.condensation_adaptive,
 74                dt_cond_range=self.settings.condensation_dt_cond_range,
 75                schedule=self.settings.condensation_schedule,
 76                **kwargs,
 77            )
 78            builder.add_dynamic(condensation)
 79        if self.settings.processes["fluid advection"]:
 80            initial_profiles = {
 81                "th": self.settings.initial_dry_potential_temperature_profile,
 82                "water_vapour_mixing_ratio": self.settings.initial_vapour_mixing_ratio_profile,
 83            }
 84            advectees = dict(
 85                (
 86                    key,
 87                    np.repeat(
 88                        profile.reshape(1, -1),
 89                        builder.particulator.environment.mesh.grid[0],
 90                        axis=0,
 91                    ),
 92                )
 93                for key, profile in initial_profiles.items()
 94            )
 95            solver = MPDATA_2D(
 96                advectees=advectees,
 97                stream_function=self.settings.stream_function,
 98                rhod_of_zZ=self.settings.rhod_of_zZ,
 99                dt=self.settings.dt,
100                grid=self.settings.grid,
101                size=self.settings.size,
102                n_iters=self.settings.mpdata_iters,
103                infinite_gauge=self.settings.mpdata_iga,
104                nonoscillatory=self.settings.mpdata_fct,
105                third_order_terms=self.settings.mpdata_tot,
106            )
107            builder.add_dynamic(EulerianAdvection(solver))
108        if self.settings.processes["particle advection"]:
109            builder.add_dynamic(
110                Displacement(
111                    enable_sedimentation=self.settings.processes["sedimentation"],
112                    adaptive=self.settings.displacement_adaptive,
113                    rtol=self.settings.displacement_rtol,
114                )
115            )
116        if (
117            self.settings.processes["coalescence"]
118            and self.settings.processes["breakup"]
119        ):
120            builder.add_dynamic(
121                Collision(
122                    collision_kernel=self.settings.kernel,
123                    enable_breakup=self.settings.processes["breakup"],
124                    coalescence_efficiency=self.settings.coalescence_efficiency,
125                    breakup_efficiency=self.settings.breakup_efficiency,
126                    fragmentation_function=self.settings.breakup_fragmentation,
127                    adaptive=self.settings.coalescence_adaptive,
128                    dt_coal_range=self.settings.coalescence_dt_coal_range,
129                    substeps=self.settings.coalescence_substeps,
130                    optimized_random=self.settings.coalescence_optimized_random,
131                )
132            )
133        elif (
134            self.settings.processes["coalescence"]
135            and not self.settings.processes["breakup"]
136        ):
137            builder.add_dynamic(
138                Coalescence(
139                    collision_kernel=self.settings.kernel,
140                    adaptive=self.settings.coalescence_adaptive,
141                    dt_coal_range=self.settings.coalescence_dt_coal_range,
142                    substeps=self.settings.coalescence_substeps,
143                    optimized_random=self.settings.coalescence_optimized_random,
144                )
145            )
146        assert not (
147            self.settings.processes["breakup"]
148            and not self.settings.processes["coalescence"]
149        )
150        if self.settings.processes["freezing"]:
151            builder.add_dynamic(
152                Freezing(
153                    immersion_freezing=self.settings.freezing_immersion,
154                    thaw=self.settings.freezing_thaw,
155                )
156            )
157
158        attributes = builder.particulator.environment.init_attributes(
159            spatial_discretisation=spatial_sampling.Pseudorandom(),
160            dry_radius_spectrum=self.settings.spectrum_per_mass_of_dry_air,
161            kappa=self.settings.kappa,
162            n_sd=self.settings.n_sd
163            // (2 if self.settings.freezing_inp_frac != 1 else 1),
164        )
165
166        if self.settings.processes["freezing"]:
167            attributes["signed water mass"] = attributes.pop("water mass")
168
169            if self.settings.freezing_inp_spec is None:
170                immersed_surface_area = formulae.trivia.sphere_surface(
171                    diameter=2 * formulae.trivia.radius(volume=attributes["dry volume"])
172                )
173            else:
174                immersed_surface_area = self.settings.freezing_inp_spec.percentiles(
175                    np.random.random(attributes["dry volume"].size),  # TODO #599: seed
176                )
177
178            if self.settings.freezing_immersion == "singular":
179                attributes["freezing temperature"] = (
180                    formulae.freezing_temperature_spectrum.invcdf(
181                        np.random.random(immersed_surface_area.size),  # TODO #599: seed
182                        immersed_surface_area,
183                    )
184                )
185            else:
186                attributes["immersed surface area"] = immersed_surface_area
187
188            if self.settings.freezing_inp_frac != 1:
189                assert self.settings.n_sd % 2 == 0
190                assert 0 < self.settings.freezing_inp_frac < 1
191                freezing_attribute = {
192                    "singular": "freezing temperature",
193                    "time-dependent": "immersed surface area",
194                }[self.settings.freezing_immersion]
195                for name, array in attributes.items():
196                    if array.shape[-1] != self.settings.n_sd // 2:
197                        raise AssertionError(f"attribute >>{name}<< has wrong size")
198                    array = array.copy()
199                    full_shape = list(array.shape)
200                    orig = slice(None, full_shape[-1])
201                    copy = slice(orig.stop, None)
202                    full_shape[-1] *= 2
203                    attributes[name] = np.empty(full_shape, dtype=array.dtype)
204                    if name == freezing_attribute:
205                        attributes[name][orig] = array
206                        attributes[name][copy] = 0
207                    elif name == "multiplicity":
208                        attributes[name][orig] = array * self.settings.freezing_inp_frac
209                        attributes[name][copy] = array * (
210                            1 - self.settings.freezing_inp_frac
211                        )
212                    elif len(array.shape) > 1:  # particle positions
213                        # TODO #599: seed
214                        for dim, _ in enumerate(array.shape):
215                            # only to make particles not shadow each other in visualisations
216                            attributes[name][dim, orig] = array[dim, :]
217                            attributes[name][dim, copy] = np.random.permutation(
218                                array[dim, :]
219                            )
220                    else:
221                        attributes[name][orig] = array
222                        attributes[name][copy] = array
223
224                non_zero_per_gridbox = np.count_nonzero(
225                    attributes[freezing_attribute]
226                ) / np.prod(self.settings.grid)
227                assert non_zero_per_gridbox == self.settings.n_sd_per_gridbox / 2
228
229        self.particulator = builder.build(attributes, tuple(products))
230
231        if self.SpinUp is not None:
232            self.SpinUp(self.particulator, self.settings.n_spin_up)
233        if self.storage is not None:
234            self.storage.init(self.settings)
235
236    def run(self, controller=DummyController(), vtk_exporter=None):
237        with controller:
238            for step in self.settings.output_steps:
239                if controller.panic:
240                    break
241
242                self.particulator.run(step - self.particulator.n_steps)
243
244                self.store(step)
245
246                if vtk_exporter is not None:
247                    vtk_exporter.export_attributes(self.particulator)
248                    vtk_exporter.export_products(self.particulator)
249
250                controller.set_percent(step / self.settings.output_steps[-1])
251
252    def store(self, step):
253        for name, product in self.particulator.products.items():
254            self.storage.save(product.get(), step, name)