make attribute saving logic compatible with variable state-vector size

open-atmos · Jan 31, 2025 · 3bb377b · 3bb377b
1 parent 59a0fa4
commit 3bb377b
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 9 deletions.
diff --git a/PySDM/dynamics/homogeneous_liquid_nucleation.py b/PySDM/dynamics/homogeneous_liquid_nucleation.py
@@ -1,3 +1,4 @@
+import numpy as np
 from PySDM.dynamics.impl import register_dynamic
 from PySDM.dynamics.impl import SuperParticleSpawningDynamic
 
@@ -41,7 +42,7 @@ def __call__(self):
             self.particulator.spawn(
                 spawned_particle_index=self.index,
                 number_of_super_particles_to_spawn=1,
-                spawned_particle_multiplicity=new_sd_multiplicity,
+                spawned_particle_multiplicity=np.asarray((new_sd_multiplicity,)),
                 spawned_particle_extensive_attributes=new_sd_extensive_attributes,
             )
             # TODO: subtract the water mass from ambient vapour
diff --git a/PySDM/particulator.py b/PySDM/particulator.py
@@ -464,7 +464,7 @@ def spawn(
                 Instead increase multiplicity of spawned particles."
             )
 
-        self.backend.spawn(
+        self.backend.spawning(
             idx=self.attributes._ParticleAttributes__idx,
             multiplicity=self.attributes["multiplicity"],
             extensive_attributes=self.attributes.get_extensive_attribute_storage(),

diff --git a/examples/PySDM_examples/Erinin_et_al_2025/expansion_simulation.py b/examples/PySDM_examples/Erinin_et_al_2025/expansion_simulation.py
@@ -119,10 +119,13 @@ def run_expansion(
 
     output = {product.name: [] for product in particulator.products.values()}
     output_attributes = {
-        "multiplicity": tuple([] for _ in range(particulator.n_sd)),
-        "volume": tuple([] for _ in range(particulator.n_sd)),
-        "critical volume": tuple([] for _ in range(particulator.n_sd)),
-        "critical supersaturation": tuple([] for _ in range(particulator.n_sd)),
+        k: []
+        for k in (
+            "multiplicity",
+            "volume",
+            "critical volume",
+            "critical supersaturation",
+        )
     }
 
     for _ in range(n_steps):
@@ -135,10 +138,13 @@ def run_expansion(
                 output[product.name].append(value)
             else:
                 output[product.name].append(value[0])
+        mult = particulator.attributes["multiplicity"].to_ndarray(raw=True)
         for key, attr in output_attributes.items():
-            attr_data = particulator.attributes[key].to_ndarray()
-            for drop_id in range(particulator.n_sd):
-                attr[drop_id].append(attr_data[drop_id])
+            if key == "multiplicity":
+                continue
+            data = particulator.attributes[key].to_ndarray(raw=True)
+            data[mult == 0] = np.nan
+            attr.append(data)
 
     dry_spectrum = particulator.products["dry:dN/dR"].get()
     wet_spectrum = particulator.products["wet:dN/dR"].get()