Monte Carlo Method for Displacement Dynamic

PySDM/backends/impl_numba/methods/displacement_methods.py

@@ -14,12 +14,15 @@
 @numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False}})
 # pylint: disable=too-many-arguments
 def calculate_displacement_body_common(
-    dim, droplet, scheme, _l, _r, displacement, courant, position_in_cell, n_substeps
+    dim, droplet, scheme, _l, _r, displacement, courant, position_in_cell, cell_id, n_substeps, enable_monte_carlo, u01
 ):
     displacement[dim, droplet] = scheme(
         position_in_cell[dim, droplet],
+        cell_id[droplet],
         courant[_l] / n_substeps,
         courant[_r] / n_substeps,
+        enable_monte_carlo,
+        u01
     )
 
 
@@ -28,13 +31,14 @@ class DisplacementMethods(BackendMethods):
     @numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False, "cache": False}})
     # pylint: disable=too-many-arguments
     def calculate_displacement_body_1d(
-        dim, scheme, displacement, courant, cell_origin, position_in_cell, n_substeps
+        dim, scheme, displacement, courant, cell_origin, position_in_cell, cell_id, n_substeps, enable_monte_carlo, rng
     ):
         length = displacement.shape[1]
         for droplet in numba.prange(length):  # pylint: disable=not-an-iterable
             # Arakawa-C grid
             _l = cell_origin[0, droplet]
             _r = cell_origin[0, droplet] + 1
+            u01 = rng[droplet]
             calculate_displacement_body_common(
                 dim,
                 droplet,
@@ -44,14 +48,17 @@ def calculate_displacement_body_1d(
                 displacement,
                 courant,
                 position_in_cell,
+                cell_id,
                 n_substeps,
+                enable_monte_carlo,
+                u01,
             )
 
     @staticmethod
     @numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False, "cache": False}})
     # pylint: disable=too-many-arguments
     def calculate_displacement_body_2d(
-        dim, scheme, displacement, courant, cell_origin, position_in_cell, n_substeps
+        dim, scheme, displacement, courant, cell_origin, position_in_cell, cell_id, n_substeps, enable_monte_carlo, rng
     ):
         length = displacement.shape[1]
         for droplet in numba.prange(length):  # pylint: disable=not-an-iterable
@@ -64,6 +71,7 @@ def calculate_displacement_body_2d(
                 cell_origin[0, droplet] + 1 * (dim == 0),
                 cell_origin[1, droplet] + 1 * (dim == 1),
             )
+            u01 = rng[droplet]
             calculate_displacement_body_common(
                 dim,
                 droplet,
@@ -73,14 +81,17 @@ def calculate_displacement_body_2d(
                 displacement,
                 courant,
                 position_in_cell,
+                cell_id,
                 n_substeps,
+                enable_monte_carlo,
+                u01,
             )
 
     @staticmethod
     @numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False, "cache": False}})
     # pylint: disable=too-many-arguments
     def calculate_displacement_body_3d(
-        dim, scheme, displacement, courant, cell_origin, position_in_cell, n_substeps
+        dim, scheme, displacement, courant, cell_origin, position_in_cell, cell_id, n_substeps, enable_monte_carlo, rng
     ):
         n_sd = displacement.shape[1]
         for droplet in numba.prange(n_sd):  # pylint: disable=not-an-iterable
@@ -95,6 +106,7 @@ def calculate_displacement_body_3d(
                 cell_origin[1, droplet] + 1 * (dim == 1),
                 cell_origin[2, droplet] + 1 * (dim == 2),
             )
+            u01 = rng[droplet]
             calculate_displacement_body_common(
                 dim,
                 droplet,
@@ -104,11 +116,14 @@ def calculate_displacement_body_3d(
                 displacement,
                 courant,
                 position_in_cell,
+                cell_id,
                 n_substeps,
+                enable_monte_carlo,
+                u01,
             )
 
     def calculate_displacement(
-        self, *, dim, displacement, courant, cell_origin, position_in_cell, n_substeps
+        self, *, dim, displacement, courant, cell_origin, position_in_cell, cell_id, n_substeps, enable_monte_carlo, rng
     ):
         n_dims = len(courant.shape)
         scheme = self.formulae.particle_advection.displacement
@@ -120,7 +135,10 @@ def calculate_displacement(
                 courant.data,
                 cell_origin.data,
                 position_in_cell.data,
+                cell_id.data,
                 n_substeps,
+                enable_monte_carlo,
+                rng.uniform(0., 1., displacement.data.shape[1]),
             )
         elif n_dims == 2:
             DisplacementMethods.calculate_displacement_body_2d(
@@ -130,7 +148,10 @@ def calculate_displacement(
                 courant.data,
                 cell_origin.data,
                 position_in_cell.data,
+                cell_id.data,
                 n_substeps,
+                enable_monte_carlo,
+                rng.uniform(0., 1., displacement.data.shape[1]),
             )
         elif n_dims == 3:
             DisplacementMethods.calculate_displacement_body_3d(
@@ -140,7 +161,10 @@ def calculate_displacement(
                 courant.data,
                 cell_origin.data,
                 position_in_cell.data,
+                cell_id.data,
                 n_substeps,
+                enable_monte_carlo,
+                rng.uniform(0., 1., displacement.data.shape[1]),
             )
         else:
             raise NotImplementedError()
@@ -173,7 +197,8 @@ def body(
                     # and crossed precip-counting level
                     position_within_column < precipitation_counting_level_index
                 ):
-                    rainfall_mass += abs(water_mass[idx[i]]) * multiplicity[idx[i]]
+                    rainfall_mass += abs(water_mass[idx[i]]) * \
+                        multiplicity[idx[i]]
                     idx[i] = flag
                     healthy[0] = 0
             return rainfall_mass

PySDM/backends/impl_thrust_rtc/methods/displacement_methods.py

@@ -24,8 +24,11 @@ def __calculate_displacement_body(self):
                     "courant_shape_0",
                     "courant_shape_1",
                     "cell_origin",
+                    "cell_id",
                     "position_in_cell",
                     "n_substeps",
+                    "enable_monte_carlo",
+                    "u01",
                 ),
                 name_iter="i",
                 body=f"""
@@ -52,8 +55,11 @@ def __calculate_displacement_body(self):
             displacement[i + n_sd * dim] = {
                 self.formulae.particle_advection.displacement.c_inline(
                     position_in_cell="position_in_cell[i + n_sd * dim]",
+                    cell_id="cell_id[i + n_sd * dim]",
                     c_l="courant[_l] / n_substeps",
-                    c_r="courant[_r] / n_substeps"
+                    c_r="courant[_r] / n_substeps",
+                    enable_monte_carlo="false",
+                    u01="u01"
                 )
             };
             """.replace(
@@ -96,7 +102,7 @@ def __flag_precipitated_body(self):
 
     @nice_thrust(**NICE_THRUST_FLAGS)
     def calculate_displacement(
-        self, *, dim, displacement, courant, cell_origin, position_in_cell, n_substeps
+        self, *, dim, displacement, courant, cell_origin, position_in_cell, cell_id, n_substeps, enable_monte_carlo, rng
     ):
         n_dim = len(courant.shape)
         n_sd = position_in_cell.shape[1]
@@ -110,8 +116,11 @@ def calculate_displacement(
                 trtc.DVInt64(courant.shape[0]),
                 trtc.DVInt64(courant.shape[1] if n_dim > 2 else -1),
                 cell_origin.data,
+                cell_id.data,
                 position_in_cell.data,
                 trtc.DVInt64(n_substeps),
+                trtc.DVBool(enable_monte_carlo),
+                trtc.DVDouble(rng.uniform(0., 1.))
             ),
         )
 

PySDM/backends/impl_thrust_rtc/test_helpers/cpp2python.py

@@ -199,7 +199,8 @@ def make(self):
         arcsinh as asinh,
         sinh,
         maximum, minimum,
-        where,  # TODO #1295
+        where, # TODO #1295
+        absolute, sign
     )
     import numba
 

PySDM/dynamics/displacement.py

@@ -21,12 +21,14 @@ class Displacement:  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
         enable_sedimentation=False,
+        enable_monte_carlo=False,
         precipitation_counting_level_index: int = 0,
         adaptive=DEFAULTS.adaptive,
         rtol=DEFAULTS.rtol,
     ):  # pylint: disable=too-many-arguments
         self.particulator = None
         self.enable_sedimentation = enable_sedimentation
+        self.enable_monte_carlo = enable_monte_carlo
         self.dimension = None
         self.grid = None
         self.courant = None
@@ -101,11 +103,12 @@ def __call__(self):
         # TIP: not need all array only [idx[:sd_num]]
         cell_origin = self.particulator.attributes["cell origin"]
         position_in_cell = self.particulator.attributes["position in cell"]
+        cell_id = self.particulator.attributes["cell id"]
 
         self.precipitation_mass_in_last_step = 0.0
         for _ in range(self._n_substeps):
             self.calculate_displacement(
-                self.displacement, self.courant, cell_origin, position_in_cell
+                self.displacement, self.courant, cell_origin, position_in_cell, cell_id
             )
             self.update_position(position_in_cell, self.displacement)
             if self.enable_sedimentation:
@@ -122,16 +125,23 @@ def __call__(self):
             self.particulator.attributes.mark_updated(key)
 
     def calculate_displacement(
-        self, displacement, courant, cell_origin, position_in_cell
+        self, displacement, courant, cell_origin, position_in_cell, cell_id
     ):
+        if self.enable_sedimentation and self.enable_monte_carlo:
+            dt = self.particulator.dt / self._n_substeps
+            dt_over_dz = dt / self.particulator.mesh.dz
+            courant += self.particulator.attributes["relative fall velocity"] / dt_over_dz
+
         self.particulator.calculate_displacement(
             displacement=displacement,
             courant=courant,
             cell_origin=cell_origin,
+            cell_id=cell_id,
             position_in_cell=position_in_cell,
             n_substeps=self._n_substeps,
+            enable_monte_carlo=self.enable_monte_carlo,
         )
-        if self.enable_sedimentation:
+        if self.enable_sedimentation and not self.enable_monte_carlo: 
             displacement_z = displacement[self.dimension - 1, :]
             dt = self.particulator.dt / self._n_substeps
             dt_over_dz = dt / self.particulator.mesh.dz

PySDM/particulator.py

@@ -427,16 +427,19 @@ def flag_out_of_column(self):
         self.attributes.sanitize()
 
     def calculate_displacement(
-        self, *, displacement, courant, cell_origin, position_in_cell, n_substeps
+        self, *, displacement, courant, cell_origin, cell_id, position_in_cell, n_substeps, enable_monte_carlo
     ):
         for dim in range(len(self.environment.mesh.grid)):
             self.backend.calculate_displacement(
                 dim=dim,
                 displacement=displacement,
                 courant=courant[dim],
                 cell_origin=cell_origin,
+                cell_id=cell_id,
                 position_in_cell=position_in_cell,
                 n_substeps=n_substeps,
+                enable_monte_carlo=enable_monte_carlo,
+                rng=self.Random(1,1).generator,
             )
 
     def isotopic_fractionation(self, heavy_isotopes: tuple):

PySDM/physics/particle_advection/explicit_in_space.py

@@ -1,12 +1,34 @@
 """
 basic explicit-in-space Euler scheme
 """
-
+import numpy as np
 
 class ExplicitInSpace:  # pylint: disable=too-few-public-methods
     def __init__(self, _):
         pass
 
     @staticmethod
-    def displacement(_, position_in_cell, c_l, c_r):
-        return c_l * (1 - position_in_cell) + c_r * position_in_cell
+    def displacement(_, position_in_cell, cell_id, c_l, c_r, enable_monte_carlo, u01):
+        return (
+            position_in_cell
+            + (
+                np.floor(
+                    np.abs(max(c_l, c_r))
+                )
+                * np.sign(
+                    np.abs(max(c_l, c_r))
+                    / max(c_l, c_r)
+                )
+            )
+            + (
+                np.abs(max(c_l, c_r)) > u01
+            )
+            * np.sign(
+                np.abs(max(c_l, c_r))
+                / max(c_l, c_r)
+            )
+        ) if enable_monte_carlo else (
+            c_l
+            * (1 - position_in_cell)
+            + c_r * position_in_cell
+        )

PySDM/physics/particle_advection/implicit_in_space.py

@@ -1,12 +1,42 @@
 """
 eqs. 14-16 in [Arabas et al. 2015](https://doi.org/10.5194/gmd-8-1677-2015)
 """
-
+import numpy as np
 
 class ImplicitInSpace:  # pylint: disable=too-few-public-methods
     def __init__(self, _):
         pass
 
     @staticmethod
-    def displacement(_, position_in_cell, c_l, c_r):
-        return (c_l * (1 - position_in_cell) + c_r * position_in_cell) / (1 - c_r + c_l)
+    def displacement(_, position_in_cell, cell_id, c_l, c_r, enable_monte_carlo, u01):
+        return (
+            position_in_cell
+            + (
+                np.floor(
+                    np.abs(max(c_l, c_r))
+                )
+                * np.sign(
+                    np.abs(max(c_l, c_r))
+                    / max(c_l, c_r)
+                )
+            )
+            + (
+                np.abs(max(c_l, c_r)) > u01
+            )
+            * np.sign(
+                np.abs(max(c_l, c_r))
+                / max(c_l, c_r)
+            )
+        ) if enable_monte_carlo else (
+            (
+                c_l
+                * (1 - position_in_cell)
+                + c_r
+                * position_in_cell
+            )
+            / (
+                1
+                - c_r
+                + c_l
+            )
+        )

tests/unit_tests/dynamics/displacement/displacement_settings.py

@@ -21,9 +21,10 @@ def __init__(self, n_sd=1, grid=None, positions=None, courant_field_data=None):
         self.sedimentation = False
         self.dt = None
 
-    def get_displacement(self, backend, scheme, adaptive=True):
+    def get_displacement(self, backend, scheme, adaptive=True, enable_monte_carlo=False):
         formulae = Formulae(particle_advection=scheme)
-        particulator = DummyParticulator(backend, n_sd=len(self.n), formulae=formulae)
+        particulator = DummyParticulator(
+            backend, n_sd=len(self.n), formulae=formulae)
         particulator.environment = DummyEnvironment(
             timestep=self.dt, grid=self.grid, courant_field_data=self.courant_field_data
         )
@@ -39,7 +40,8 @@ def get_displacement(self, backend, scheme, adaptive=True):
             "position in cell": position_in_cell,
         }
         particulator.build(attributes)
-        sut = Displacement(enable_sedimentation=self.sedimentation, adaptive=adaptive)
+        sut = Displacement(enable_sedimentation=self.sedimentation,
+                           adaptive=adaptive, enable_monte_carlo=enable_monte_carlo)
         sut.register(particulator)
         sut.upload_courant_field(self.courant_field_data)