Add test_create_particle_data

Author: VeckoTheGecko

parcels/particle.py

@@ -194,7 +194,7 @@ def create_particle_data(
     nparticles: int,
     ngrids: int,
     time_interval: TimeInterval,
-    initial: dict[str, np.array] | None,
+    initial: dict[str, np.array] | None = None,
 ):
     if initial is None:
         initial = {}
@@ -231,18 +231,20 @@ def create_particle_data(
             name_to_copy = var.initial(_attrgetter_helper)
             data[var.name] = data[name_to_copy].copy()
         else:
-            data[var.name] = _create_array_for_variable(var, nparticles)
+            data[var.name] = _create_array_for_variable(var, nparticles, time_interval)
     return data
 
 
-def _create_array_for_variable(variable: Variable, nparticles: int):
+def _create_array_for_variable(variable: Variable, nparticles: int, time_interval: TimeInterval):
     assert not isinstance(variable.initial, operator.attrgetter), (
         "This function cannot handle attrgetter initial values."
     )
+    if (dtype := variable.dtype) is _SAME_AS_FIELDSET_TIME_INTERVAL.VALUE:
+        dtype = type(time_interval.left)
     return np.full(
         shape=(nparticles,),
         fill_value=variable.initial,
-        dtype=variable.dtype,
+        dtype=dtype,
     )
 
 

tests/v4/test_particle.py

@@ -1,7 +1,9 @@
 import numpy as np
 import pytest
 
-from parcels.particle import ParticleClass, Variable
+from parcels._core.utils.time import TimeInterval
+from parcels._datasets.structured.generic import TIME
+from parcels.particle import _SAME_AS_FIELDSET_TIME_INTERVAL, Particle, ParticleClass, Variable, create_particle_data
 
 
 def test_variable_init():
@@ -104,3 +106,39 @@ def test_particleclass_add_variable_collision():
 
     with pytest.raises(ValueError, match="Variable name already exists: "):
         p_initial.add_variable([Variable("vara", dtype=np.float32, to_write=True)])
+
+
+@pytest.mark.parametrize(
+    "particle",
+    [
+        ParticleClass(
+            variables=[
+                Variable("vara", dtype=np.float32, initial=1.0),
+                Variable("varb", dtype=np.float32, initial=2.0),
+            ]
+        ),
+        Particle,
+    ],
+)
+@pytest.mark.parametrize("nparticles", [5, 10])
+def test_create_particle_data(particle, nparticles):
+    time_interval = TimeInterval(TIME[0], TIME[-1])
+    ngrids = 4
+    data = create_particle_data(pclass=particle, nparticles=nparticles, ngrids=ngrids, time_interval=time_interval)
+
+    assert isinstance(data, dict)
+    assert len(data) == len(particle.variables) + 1  # ei variable is separate
+
+    variables = {var.name: var for var in particle.variables}
+
+    for variable_name in variables.keys():
+        variable = variables[variable_name]
+        variable_array = data[variable_name]
+
+        assert variable_array.shape[0] == nparticles
+
+        dtype = variable.dtype
+        if dtype is _SAME_AS_FIELDSET_TIME_INTERVAL.VALUE:
+            dtype = type(time_interval.left)
+
+        assert variable_array.dtype == dtype