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