introducing FlagAttribute and FlagCoalescence to mark particles that took part in coalescence within last timestep

PySDM/attributes/impl/__init__.py

@@ -11,3 +11,4 @@
 from .maximum_attribute import MaximumAttribute
 from .attribute_registry import register_attribute, get_attribute_class
 from .intensive_attribute import IntensiveAttribute
+from .flag_attribute import FlagAttribute

PySDM/attributes/impl/flag_attribute.py

@@ -0,0 +1,21 @@
+""" a Boolean attribute that gets reset for False after copying out the recorded values """
+
+from .attribute import Attribute
+
+
+class FlagAttribute(Attribute):
+    def __init__(self, builder, name):
+        super().__init__(builder, name, dtype=bool)
+        builder.particulator.observers.append(self)
+
+    def notify(self):
+        source = self.particulator.dynamics["Collision"].flag_coalescence
+        self.data[:] = source
+        source[:] = False
+
+    def allocate(self, idx):
+        super().allocate(idx)
+        self.data[:] = False
+
+    def get(self):
+        return self.data

PySDM/attributes/physics/__init__.py

@@ -18,3 +18,4 @@
 from .terminal_velocity import TerminalVelocity
 from .volume import Volume
 from .reynolds_number import ReynoldsNumber
+from .flags import FlagCoalescence

PySDM/attributes/physics/flags.py

@@ -0,0 +1,7 @@
+from ..impl import register_attribute, FlagAttribute
+
+
+@register_attribute()
+class FlagCoalescence(FlagAttribute):
+    def __init__(self, builder):
+        super().__init__(builder, name="coalescence flag")

PySDM/backends/impl_numba/methods/collisions_methods.py

@@ -43,9 +43,10 @@ def flag_zero_multiplicity(j, k, multiplicity, healthy):
 
 @numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False}})
 def coalesce(  # pylint: disable=too-many-arguments
-    i, j, k, cid, multiplicity, gamma, attributes, coalescence_rate
+    i, j, k, cid, multiplicity, gamma, attributes, coalescence_rate, flag_coalescence
 ):
     atomic_add(coalescence_rate, cid, gamma[i] * multiplicity[k])
+    flag_coalescence[j] = flag_coalescence[k] = True
     new_n = multiplicity[j] - gamma[i] * multiplicity[k]
     if new_n > 0:
         multiplicity[j] = new_n
@@ -269,6 +270,7 @@ def body(
             max_multiplicity,
             warn_overflows,
             particle_mass,
+            flag_coalescence,
         ):
             # pylint: disable=not-an-iterable,too-many-nested-blocks,too-many-locals
             for i in numba.prange(length // 2):
@@ -289,6 +291,7 @@ def body(
                         gamma,
                         attributes,
                         coalescence_rate,
+                        flag_coalescence,
                     )
                 else:
                     _break_up(
@@ -421,6 +424,7 @@ def body(
             cell_id,
             coalescence_rate,
             is_first_in_pair,
+            flag_coalescence,
         ):
             for (
                 i
@@ -439,6 +443,7 @@ def body(
                     gamma,
                     attributes,
                     coalescence_rate,
+                    flag_coalescence,
                 )
                 flag_zero_multiplicity(j, k, multiplicity, healthy)
 
@@ -455,6 +460,7 @@ def collision_coalescence(
         cell_id,
         coalescence_rate,
         is_first_in_pair,
+        flag_coalescence,
     ):
         self._collision_coalescence_body(
             multiplicity=multiplicity.data,
@@ -466,6 +472,7 @@ def collision_coalescence(
             cell_id=cell_id.data,
             coalescence_rate=coalescence_rate.data,
             is_first_in_pair=is_first_in_pair.indicator.data,
+            flag_coalescence=flag_coalescence.data,
         )
 
     def collision_coalescence_breakup(
@@ -488,6 +495,7 @@ def collision_coalescence_breakup(
         warn_overflows,
         particle_mass,
         max_multiplicity,
+        flag_coalescence,
     ):
         # pylint: disable=too-many-locals
         self._collision_coalescence_breakup_body(
@@ -509,6 +517,7 @@ def collision_coalescence_breakup(
             max_multiplicity=max_multiplicity,
             warn_overflows=warn_overflows,
             particle_mass=particle_mass.data,
+            flag_coalescence=flag_coalescence.data,
         )
 
     @cached_property

PySDM/backends/impl_thrust_rtc/methods/collisions_methods.py

@@ -34,11 +34,13 @@
     VectorView<real_type> gamma,
     VectorView<real_type> attributes,
     VectorView<int64_t> coalescence_rate,
+    VectorView<bool_t> flag_coalescence,
     int64_t n_attr,
     int64_t n_sd
   ) {
     auto cid = cell_id[j];
     atomicAdd((unsigned long long int*)&coalescence_rate[cid], (unsigned long long int)(gamma[i] * multiplicity[k]));
+    flag_coalescence[j] = flag_coalescence[k] = true;
     auto new_n = multiplicity[j] - gamma[i] * multiplicity[k];
     if (new_n > 0) {
         multiplicity[j] = new_n;
@@ -328,6 +330,7 @@ def __collision_coalescence_body(self):
                 "cell_id",
                 "coalescence_rate",
                 "is_first_in_pair",
+                "flag_coalescence",
             ),
             name_iter="i",
             body=f"""
@@ -341,7 +344,7 @@ def __collision_coalescence_body(self):
                 return;
             }}
 
-            Commons::coalesce(i, j, k, cell_id, multiplicity, gamma, attributes, coalescence_rate, n_attr, n_sd);
+            Commons::coalesce(i, j, k, cell_id, multiplicity, gamma, attributes, coalescence_rate, flag_coalescence, n_attr, n_sd);
 
             Commons::flag_zero_multiplicity(j, k, multiplicity, healthy);
             """.replace(
@@ -369,6 +372,7 @@ def __collision_coalescence_breakup_body(self):
                 "is_first_in_pair",
                 "max_multiplicity",
                 "particle_mass",
+                "flag_coalescence",
                 "n_sd",
                 "n_attr",
             ),
@@ -398,6 +402,7 @@ def __collision_coalescence_breakup_body(self):
                     gamma,
                     attributes,
                     coalescence_rate,
+                    flag_coalescence,
                     n_attr,
                     n_sd
                 );
@@ -793,6 +798,7 @@ def collision_coalescence(
         cell_id,
         coalescence_rate,
         is_first_in_pair,
+        flag_coalescence,
     ):
         if len(idx) < 2:
             return
@@ -811,6 +817,7 @@ def collision_coalescence(
                 cell_id.data,
                 coalescence_rate.data,
                 is_first_in_pair.indicator.data,
+                flag_coalescence.data,
             ),
         )
 
@@ -835,6 +842,7 @@ def collision_coalescence_breakup(  # pylint: disable=unused-argument,too-many-l
         warn_overflows,
         particle_mass,
         max_multiplicity,
+        flag_coalescence,
     ):
         if warn_overflows:
             raise NotImplementedError()
@@ -861,6 +869,7 @@ def collision_coalescence_breakup(  # pylint: disable=unused-argument,too-many-l
                 is_first_in_pair.indicator.data,
                 trtc.DVInt64(max_multiplicity),
                 particle_mass.data,
+                flag_coalescence.data,
                 n_sd,
                 n_attr,
             ),

PySDM/backends/impl_thrust_rtc/test_helpers/cpp2python.py

@@ -35,6 +35,7 @@
     "(float)": "np.float32",
     "return;": "continue",
     "void*": "",
+    "VectorView<bool_t> ": "",
     "VectorView<int64_t> ": "",
     "VectorView<float> ": "",
     "VectorView<double> ": "",

PySDM/dynamics/collisions/collision.py

@@ -98,6 +98,8 @@ def __init__(
         self.breakup_rate = None
         self.breakup_rate_deficit = None
 
+        self.flag_coalescence = None
+
     def register(self, builder):
         self.particulator = builder.particulator
         rnd_args = {
@@ -148,6 +150,10 @@ def register(self, builder):
         )
         self.coalescence_rate = self.particulator.Storage.from_ndarray(*counter_args)
 
+        self.flag_coalescence = self.particulator.Storage.from_ndarray(
+            np.full(self.particulator.n_sd, fill_value=False, dtype=bool),
+        )
+
         if self.enable_breakup:
             self.n_fragment = self.particulator.PairwiseStorage.empty(
                 **empty_args_pairwise
@@ -231,6 +237,7 @@ def step(self):
             is_first_in_pair=self.is_first_in_pair,
             warn_overflows=self.warn_overflows,
             max_multiplicity=self.max_multiplicity,
+            flag_coalescence=self.flag_coalescence,
         )
 
     def toss_candidate_pairs_and_sort_within_pair_by_multiplicity(

PySDM/impl/particle_attributes_factory.py

@@ -10,6 +10,7 @@
     DummyAttribute,
     ExtensiveAttribute,
     MaximumAttribute,
+    FlagAttribute,
 )
 from PySDM.impl.particle_attributes import ParticleAttributes
 from PySDM.attributes.impl import get_attribute_class
@@ -35,6 +36,7 @@ def attributes(particulator, req_attr, attributes):
                     get_attribute_class("multiplicity"),
                     CellAttribute,
                     DummyAttribute,
+                    FlagAttribute,
                 ),
             ):
                 raise AssertionError()
@@ -47,10 +49,10 @@ def attributes(particulator, req_attr, attributes):
         )
 
         for attr in req_attr.values():
-            if isinstance(attr, (DerivedAttribute, DummyAttribute)):
+            if isinstance(attr, (DerivedAttribute, DummyAttribute, FlagAttribute)):
                 if attr.name in attributes:
                     raise ValueError(
-                        f"attribute '{attr.name}' is a dummy/derived one,"
+                        f"attribute '{attr.name}' is a dummy/derived/flag attribute,"
                         f" but values were provided"
                     )
                 attr.allocate(idx)

PySDM/particulator.py

@@ -164,6 +164,7 @@ def collision_coalescence_breakup(
         is_first_in_pair,
         warn_overflows,
         max_multiplicity,
+        flag_coalescence,
     ):
         # pylint: disable=too-many-locals
         idx = self.attributes._ParticleAttributes__idx
@@ -190,6 +191,7 @@ def collision_coalescence_breakup(
                 warn_overflows=warn_overflows,
                 particle_mass=self.attributes["water mass"],
                 max_multiplicity=max_multiplicity,
+                flag_coalescence=flag_coalescence,
             )
         else:
             self.backend.collision_coalescence(
@@ -201,6 +203,7 @@ def collision_coalescence_breakup(
                 cell_id=cell_id,
                 coalescence_rate=coalescence_rate,
                 is_first_in_pair=is_first_in_pair,
+                flag_coalescence=flag_coalescence,
             )
         self.attributes.sanitize()
         self.attributes.mark_updated("multiplicity")

tests/unit_tests/attributes/test_flags.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+from PySDM.builder import Builder
+from PySDM.dynamics import Coalescence
+from PySDM.dynamics.collisions.collision_kernels import Golovin
+from PySDM.environments import Box
+from PySDM.physics import si
+
+
+def make_one_step_and_record_flag(particulator, output, enable):
+    particulator.dynamics["Collision"].enable = enable
+    particulator.run(steps=1)
+    output += [particulator.attributes["flag coalescence"].to_ndarray().copy()]
+
+
+def test_flag_coalescence(backend_class, n_sd=1000, golovin_b=1e15):
+    # arrange
+    builder = Builder(
+        backend=backend_class(), n_sd=n_sd, environment=Box(dv=1 * si.m**3, dt=1 * si.s)
+    )
+    builder.add_dynamic(Coalescence(collision_kernel=Golovin(b=golovin_b)))
+    builder.request_attribute("flag coalescence")
+    particulator = builder.build(
+        attributes={
+            "multiplicity": np.full(n_sd, fill_value=1e5),
+            "water mass": np.full(n_sd, fill_value=1 * si.ug),
+        }
+    )
+
+    # act
+    output = []
+    make_one_step_and_record_flag(particulator, output, enable=False)
+    make_one_step_and_record_flag(particulator, output, enable=True)
+    make_one_step_and_record_flag(particulator, output, enable=False)
+    make_one_step_and_record_flag(particulator, output, enable=True)
+
+    # assert
+    assert (output[0] == False).all()
+    assert (output[1] == True).any()
+    assert (output[2] == False).all()
+    assert (output[3] != output[2]).any()