Seeding example for 1D environment of S&H 2012

PySDM/backends/impl_numba/methods/seeding_methods.py

@@ -14,10 +14,18 @@ def _seeding(self):
         def body(  # pylint: disable=too-many-arguments
             idx,
             multiplicity,
+            cell_id,
+            cell_origin,
+            pos_cell,
+            volume,
             extensive_attributes,
             seeded_particle_index,
             seeded_particle_multiplicity,
+            seeded_particle_cell_id,
+            seeded_particle_cell_origin,
+            seeded_particle_pos_cell,
             seeded_particle_extensive_attributes,
+            seeded_particle_volume,
             number_of_super_particles_to_inject: int,
         ):
             number_of_super_particles_already_injected = 0
@@ -35,10 +43,14 @@ def body(  # pylint: disable=too-many-arguments
                     ]
                     number_of_super_particles_already_injected += 1
                     multiplicity[i] = seeded_particle_multiplicity[s]
+                    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]
                     for a in range(len(extensive_attributes)):
                         extensive_attributes[a, i] = (
                             seeded_particle_extensive_attributes[a, s]
                         )
+                    volume[i] = seeded_particle_volume[s]
             assert (
                 number_of_super_particles_to_inject
                 == number_of_super_particles_already_injected
@@ -51,18 +63,34 @@ def seeding(
         *,
         idx,
         multiplicity,
+        cell_id,
+        cell_origin,
+        pos_cell,
+        volume,
         extensive_attributes,
         seeded_particle_index,
         seeded_particle_multiplicity,
+        seeded_particle_cell_id,
+        seeded_particle_cell_origin,
+        seeded_particle_pos_cell,
         seeded_particle_extensive_attributes,
+        seeded_particle_volume,
         number_of_super_particles_to_inject: int,
     ):
         self._seeding(
             idx=idx.data,
             multiplicity=multiplicity.data,
+            cell_id=cell_id.data,
+            cell_origin=cell_origin.data,
+            pos_cell=pos_cell.data,
+            volume=volume.data,
             extensive_attributes=extensive_attributes.data,
             seeded_particle_index=seeded_particle_index.data,
             seeded_particle_multiplicity=seeded_particle_multiplicity.data,
+            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_extensive_attributes=seeded_particle_extensive_attributes.data,
+            seeded_particle_volume=seeded_particle_volume.data,
             number_of_super_particles_to_inject=number_of_super_particles_to_inject,
         )

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: Sized,
+        seeded_particle_cell_origin: float,
+        seeded_particle_pos_cell: float,
+        seeded_particle_volume: float,
     ):
         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
@@ -42,6 +50,12 @@ def post_register_setup_when_attributes_are_known(self):
                 f" ({self.particulator.attributes.get_extensive_attribute_keys()})"
             )
 
+        self.tmp_dry_vol = self.seeded_particle_extensive_attributes["dry volume"]
+        self.tmp_kappa_times_dry_vol = self.seeded_particle_extensive_attributes[
+            "kappa times dry volume"
+        ]
+        self.tmp_cell_id = self.seeded_particle_cell_id
+
         self.index = self.particulator.Index.identity_index(
             len(self.seeded_particle_multiplicity)
         )
@@ -60,6 +74,24 @@ def post_register_setup_when_attributes_are_known(self):
                 ),
             )
         )
+        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),
+        )
         self.seeded_particle_extensive_attributes = (
             self.particulator.IndexedStorage.from_ndarray(
                 self.index,
@@ -90,5 +122,9 @@ def __call__(self):
                 seeded_particle_index=self.index,
                 number_of_super_particles_to_inject=number_of_super_particles_to_inject,
                 seeded_particle_multiplicity=self.seeded_particle_multiplicity,
+                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,
                 seeded_particle_extensive_attributes=self.seeded_particle_extensive_attributes,
             )

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:
@@ -444,6 +447,10 @@ def seeding(
         *,
         seeded_particle_index,
         seeded_particle_multiplicity,
+        seeded_particle_cell_id,
+        seeded_particle_cell_origin,
+        seeded_particle_pos_cell,
+        seeded_particle_volume,
         seeded_particle_extensive_attributes,
         number_of_super_particles_to_inject,
     ):
@@ -467,10 +474,18 @@ def seeding(
         self.backend.seeding(
             idx=self.attributes._ParticleAttributes__idx,
             multiplicity=self.attributes["multiplicity"],
+            cell_id=self.attributes["cell id"],
+            cell_origin=self.attributes["cell origin"],
+            pos_cell=self.attributes["position in cell"],
+            volume=self.attributes["volume"],
             extensive_attributes=self.attributes.get_extensive_attribute_storage(),
             seeded_particle_index=seeded_particle_index,
             seeded_particle_multiplicity=seeded_particle_multiplicity,
+            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_extensive_attributes=seeded_particle_extensive_attributes,
+            seeded_particle_volume=seeded_particle_volume,
             number_of_super_particles_to_inject=number_of_super_particles_to_inject,
         )
         self.attributes.reset_idx()

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