Seeding examples for 1D environment of S&H 2012 as well as 2D environment of M&G2007

PySDM/backends/impl_numba/methods/seeding_methods.py

@@ -19,6 +19,14 @@ def body(  # pylint: disable=too-many-arguments
             seeded_particle_multiplicity,
             seeded_particle_extensive_attributes,
             number_of_super_particles_to_inject: int,
+            seeded_particle_cell_id,
+            seeded_particle_cell_origin,
+            seeded_particle_pos_cell,
+            seeded_particle_volume,
+            cell_id,
+            cell_origin,
+            pos_cell,
+            volume,
         ):
             number_of_super_particles_already_injected = 0
             # TODO #1387 start enumerating from the end of valid particle set
@@ -39,6 +47,13 @@ def body(  # pylint: disable=too-many-arguments
                         extensive_attributes[a, i] = (
                             seeded_particle_extensive_attributes[a, s]
                         )
+
+                    if cell_id is not None:
+                        cell_id[i] = seeded_particle_cell_id[s]
+                        cell_origin[0][i] = seeded_particle_cell_origin[0][s]
+                        pos_cell[0][i] = seeded_particle_pos_cell[0][s]
+                        volume[i] = seeded_particle_volume[s]
+
             assert (
                 number_of_super_particles_to_inject
                 == number_of_super_particles_already_injected
@@ -56,6 +71,14 @@ def seeding(
         seeded_particle_multiplicity,
         seeded_particle_extensive_attributes,
         number_of_super_particles_to_inject: int,
+        seeded_particle_cell_id,
+        seeded_particle_cell_origin,
+        seeded_particle_pos_cell,
+        seeded_particle_volume,
+        cell_id,
+        cell_origin,
+        pos_cell,
+        volume,
     ):
         self._seeding(
             idx=idx.data,
@@ -65,4 +88,12 @@ def seeding(
             seeded_particle_multiplicity=seeded_particle_multiplicity.data,
             seeded_particle_extensive_attributes=seeded_particle_extensive_attributes.data,
             number_of_super_particles_to_inject=number_of_super_particles_to_inject,
+            seeded_particle_cell_id=seeded_particle_cell_id,
+            seeded_particle_cell_origin=seeded_particle_cell_origin,
+            seeded_particle_pos_cell=seeded_particle_pos_cell,
+            seeded_particle_volume=seeded_particle_volume,
+            cell_id=cell_id,
+            cell_origin=cell_origin,
+            pos_cell=pos_cell,
+            volume=volume,
         )

PySDM/dynamics/seeding.py

@@ -18,13 +18,21 @@ def __init__(
         super_droplet_injection_rate: callable,
         seeded_particle_extensive_attributes: dict,
         seeded_particle_multiplicity: Sized,
+        seeded_particle_cell_id=None,
+        seeded_particle_cell_origin=None,
+        seeded_particle_pos_cell=None,
+        seeded_particle_volume=None,
     ):
         for attr in seeded_particle_extensive_attributes.values():
             assert len(seeded_particle_multiplicity) == len(attr)
         self.particulator = None
         self.super_droplet_injection_rate = super_droplet_injection_rate
         self.seeded_particle_extensive_attributes = seeded_particle_extensive_attributes
         self.seeded_particle_multiplicity = seeded_particle_multiplicity
+        self.seeded_particle_cell_id = seeded_particle_cell_id
+        self.seeded_particle_cell_origin = seeded_particle_cell_origin
+        self.seeded_particle_pos_cell = seeded_particle_pos_cell
+        self.seeded_particle_volume = seeded_particle_volume
         self.rnd = None
         self.u01 = None
         self.index = None
@@ -67,6 +75,36 @@ def post_register_setup_when_attributes_are_known(self):
             )
         )
 
+        if self.particulator.environment.__class__.__name__ == "Parcel":
+            self.seeded_particle_cell_id = None
+            self.seeded_particle_cell_origin = None
+            self.seeded_particle_pos_cell = None
+            self.seeded_particle_volume = None
+
+        else:
+            self.seeded_particle_cell_id = (
+                self.particulator.IndexedStorage.from_ndarray(
+                    self.index,
+                    np.asarray(self.seeded_particle_cell_id),
+                )
+            )
+            self.seeded_particle_cell_origin = (
+                self.particulator.IndexedStorage.from_ndarray(
+                    self.index,
+                    np.asarray(self.seeded_particle_cell_origin),
+                )
+            )
+            self.seeded_particle_pos_cell = (
+                self.particulator.IndexedStorage.from_ndarray(
+                    self.index,
+                    np.asarray(self.seeded_particle_pos_cell),
+                )
+            )
+            self.seeded_particle_volume = self.particulator.IndexedStorage.from_ndarray(
+                self.index,
+                np.asarray(self.seeded_particle_volume),
+            )
+
     def __call__(self):
         if self.particulator.n_steps == 0:
             self.post_register_setup_when_attributes_are_known()
@@ -91,4 +129,8 @@ def __call__(self):
                 number_of_super_particles_to_inject=number_of_super_particles_to_inject,
                 seeded_particle_multiplicity=self.seeded_particle_multiplicity,
                 seeded_particle_extensive_attributes=self.seeded_particle_extensive_attributes,
+                seeded_particle_cell_id=self.seeded_particle_cell_id,
+                seeded_particle_cell_origin=self.seeded_particle_cell_origin,
+                seeded_particle_pos_cell=self.seeded_particle_pos_cell,
+                seeded_particle_volume=self.seeded_particle_volume,
             )

PySDM/particulator.py

@@ -67,6 +67,9 @@ def Random(self):
     def n_sd(self) -> int:
         return self.__n_sd
 
+    def n_sd_setter(self, value):
+        self.__n_sd += value
+
     @property
     def dt(self) -> float:
         if self.environment is not None:
@@ -446,6 +449,10 @@ def seeding(
         seeded_particle_multiplicity,
         seeded_particle_extensive_attributes,
         number_of_super_particles_to_inject,
+        seeded_particle_cell_id=None,
+        seeded_particle_cell_origin=None,
+        seeded_particle_pos_cell=None,
+        seeded_particle_volume=None,
     ):
         n_null = self.n_sd - self.attributes.super_droplet_count
         if n_null == 0:
@@ -464,15 +471,44 @@ def seeding(
                 Instead increase multiplicity of injected particles."
             )
 
-        self.backend.seeding(
-            idx=self.attributes._ParticleAttributes__idx,
-            multiplicity=self.attributes["multiplicity"],
-            extensive_attributes=self.attributes.get_extensive_attribute_storage(),
-            seeded_particle_index=seeded_particle_index,
-            seeded_particle_multiplicity=seeded_particle_multiplicity,
-            seeded_particle_extensive_attributes=seeded_particle_extensive_attributes,
-            number_of_super_particles_to_inject=number_of_super_particles_to_inject,
-        )
+        if self.environment.__class__.__name__ == "Parcel":
+            self.backend.seeding(
+                idx=self.attributes._ParticleAttributes__idx,
+                multiplicity=self.attributes["multiplicity"],
+                extensive_attributes=self.attributes.get_extensive_attribute_storage(),
+                seeded_particle_index=seeded_particle_index,
+                seeded_particle_multiplicity=seeded_particle_multiplicity,
+                seeded_particle_extensive_attributes=seeded_particle_extensive_attributes,
+                number_of_super_particles_to_inject=number_of_super_particles_to_inject,
+                seeded_particle_cell_id=seeded_particle_cell_id,
+                seeded_particle_cell_origin=seeded_particle_cell_origin,
+                seeded_particle_pos_cell=seeded_particle_pos_cell,
+                seeded_particle_volume=seeded_particle_volume,
+                cell_id=None,
+                cell_origin=None,
+                pos_cell=None,
+                volume=None,
+            )
+
+        else:
+            self.backend.seeding(
+                idx=self.attributes._ParticleAttributes__idx,
+                multiplicity=self.attributes["multiplicity"],
+                extensive_attributes=self.attributes.get_extensive_attribute_storage(),
+                seeded_particle_index=seeded_particle_index,
+                seeded_particle_multiplicity=seeded_particle_multiplicity,
+                seeded_particle_extensive_attributes=seeded_particle_extensive_attributes,
+                number_of_super_particles_to_inject=number_of_super_particles_to_inject,
+                seeded_particle_cell_id=seeded_particle_cell_id.data,
+                seeded_particle_cell_origin=seeded_particle_cell_origin.data,
+                seeded_particle_pos_cell=seeded_particle_pos_cell.data,
+                seeded_particle_volume=seeded_particle_volume.data,
+                cell_id=self.attributes["cell id"].data,
+                cell_origin=self.attributes["cell origin"].data,
+                pos_cell=self.attributes["position in cell"].data,
+                volume=self.attributes["volume"].data,
+            )
+
         self.attributes.reset_idx()
         self.attributes.sanitize()
 

examples/PySDM_examples/seeding/kinematic_1d_seeding.py

@@ -0,0 +1,92 @@
+"""
+Single-column time-varying-updraft framework with moisture advection handled by
+[PyMPDATA](http://github.com/open-atmos/PyMPDATA/)
+"""
+
+import numpy as np
+
+from PySDM.environments.impl.moist import Moist
+
+from PySDM.impl import arakawa_c
+from PySDM.initialisation.equilibrate_wet_radii import equilibrate_wet_radii
+from PySDM.environments.impl import register_environment
+
+
+@register_environment()
+class Kinematic1D(Moist):
+    def __init__(self, *, dt, mesh, thd_of_z, rhod_of_z, z0=0):
+        super().__init__(dt, mesh, [])
+        self.thd0 = thd_of_z(z0 + mesh.dz * arakawa_c.z_scalar_coord(mesh.grid))
+        self.rhod = rhod_of_z(z0 + mesh.dz * arakawa_c.z_scalar_coord(mesh.grid))
+        self.formulae = None
+
+    def register(self, builder):
+        super().register(builder)
+        self.formulae = builder.particulator.formulae
+        rhod = builder.particulator.Storage.from_ndarray(self.rhod)
+        self._values["current"]["rhod"] = rhod
+        self._tmp["rhod"] = rhod
+
+    def get_water_vapour_mixing_ratio(self) -> np.ndarray:
+        return self.particulator.dynamics["EulerianAdvection"].solvers.advectee
+
+    def get_thd(self) -> np.ndarray:
+        return self.thd0
+
+    def init_attributes(
+        self,
+        *,
+        spatial_discretisation,
+        spectral_discretisation,
+        kappa,
+        n_sd,
+        z_part=None,
+        collisions_only=False
+    ):
+        super().sync()
+        self.notify()
+
+        attributes = {}
+        with np.errstate(all="raise"):
+            positions = spatial_discretisation.sample(
+                backend=self.particulator.backend,
+                grid=self.mesh.grid,
+                n_sd=n_sd,
+                z_part=z_part,
+            )
+            (
+                attributes["cell id"],
+                attributes["cell origin"],
+                attributes["position in cell"],
+            ) = self.mesh.cellular_attributes(positions)
+
+            if collisions_only:
+                v_wet, n_per_kg = spectral_discretisation.sample(
+                    backend=self.particulator.backend, n_sd=n_sd
+                )
+                attributes["volume"] = v_wet
+            else:
+                r_dry, n_per_kg = spectral_discretisation.sample(
+                    backend=self.particulator.backend, n_sd=n_sd
+                )
+                attributes["dry volume"] = self.formulae.trivia.volume(radius=r_dry)
+                attributes["kappa times dry volume"] = attributes["dry volume"] * kappa
+                r_wet = equilibrate_wet_radii(
+                    r_dry=r_dry,
+                    environment=self,
+                    cell_id=attributes["cell id"],
+                    kappa_times_dry_volume=attributes["kappa times dry volume"],
+                )
+                attributes["volume"] = self.formulae.trivia.volume(radius=r_wet)
+
+            rhod = self["rhod"].to_ndarray()
+            cell_id = attributes["cell id"]
+            domain_volume = np.prod(np.array(self.mesh.size))
+
+        attributes["multiplicity"] = n_per_kg * rhod[cell_id] * domain_volume
+
+        return attributes
+
+    @property
+    def dv(self):
+        return self.mesh.dv