Check writing of initial values of sampled Variables

docs/user_guide/examples/tutorial_sampling.ipynb

@@ -11,8 +11,8 @@
     "In this tutorial we will go through how particles can sample `Fields`, using temperature as an example. Along the way we will get to know the parcels class `Variable` (see [here](https://parcels.readthedocs.io/en/latest/reference/particles.html#parcels.particle.Variable) for the documentation) and some of its methods. This tutorial covers several applications of a sampling setup:\n",
     "\n",
     "- [**Basic along trajectory sampling**](#basic-sampling)\n",
-    "- [**Sampling velocity fields**](#sampling-velocity-fields)\n",
-    "- [**Sampling initial conditions**](#sampling-initial-values)"
+    "- [**Sampling initial Field values**](#sampling-initial-field-values)\n",
+    "- [**Sampling velocity fields**](#sampling-velocity-fields)"
    ]
   },
   {
@@ -200,6 +200,78 @@
     "plt.show()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Sampling initial Field values\n",
+    "\n",
+    "When creating a ParticleSet, the initial values of `Variables` are set as a constant (default `0`, but in the example above set to `np.nan`). Those will also be the values that appear as first entry in the ParticleFile."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\n",
+    "    [\n",
+    "        traj[\"temperature\"][0]\n",
+    "        for traj in df.partition_by(\"particle_id\", maintain_order=True)\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If you want to sample the initial value of a field, you can do so by calling the field with the ParticleSet as an argument. For example, in, the code above you would add one line (`pset.temperature = fieldset.thetao[pset]`) to sample the initial temperature values, beofe calling `pset.execute(...)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-output"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "pset = parcels.ParticleSet(\n",
+    "    fieldset=fieldset, pclass=SampleParticle, x=lon, y=lat, z=z, t=time\n",
+    ")\n",
+    "pset.temperature = fieldset.thetao[pset]  # <-- sample initial temperature values\n",
+    "\n",
+    "output_file = parcels.ParticleFile(\n",
+    "    \"sampletemp_withinit.parquet\", outputdt=timedelta(hours=1)\n",
+    ")\n",
+    "\n",
+    "pset.execute(\n",
+    "    [parcels.kernels.AdvectionRK2, SampleT],\n",
+    "    runtime=timedelta(hours=30),\n",
+    "    dt=timedelta(minutes=5),\n",
+    "    output_file=output_file,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = parcels.read_particlefile(\"sampletemp_withinit.parquet\")\n",
+    "print(\n",
+    "    [\n",
+    "        round(traj[\"temperature\"][0], 2)\n",
+    "        for traj in df.partition_by(\"particle_id\", maintain_order=True)\n",
+    "    ]\n",
+    ")"
+   ]
+  },
   {
    "attachments": {},
    "cell_type": "markdown",

docs/user_guide/getting_started/explanation_concepts.md

@@ -94,6 +94,7 @@ pset = parcels.ParticleSet(fieldset=fieldset, pclass=parcels.Particle, t=t, z=z,
 ```{admonition} 🖥️ Learn more about how to create ParticleSets
 :class: seealso
 - [Release particles at different times](../examples/tutorial_delaystart.ipynb)
+- [Sampling initial Field values](../examples/tutorial_sampling.ipynb#sampling-initial-field-values)
 ```
 
 ## 3. Kernels

docs/user_guide/getting_started/tutorial_quickstart.md

@@ -112,6 +112,14 @@ ax = temperature.axes
 ax.scatter(lon, lat, s=40, c='w', edgecolors='r');
 ```
 
+If you also want to sample the initial value of a field, you can do so by calling the field with the `ParticleSet` as an argument. For example, in the code above you would add one line (assuming that the ParticleSet has a variable `temperature`) to sample the initial temperature values, before calling `pset.execute(...)`.
+
+```python
+pset.temperature = fieldset.thetao[pset]
+```
+
+See the [sampling tutorial](../examples/tutorial_sampling.ipynb#sampling-initial-field-values) for more information on initial sampling of fields.
+
 ## Compute: `Kernel`
 
 After setting up the input data and particle start locations and times, we need to specify what calculations to do with

src/parcels/_core/field.py

@@ -8,6 +8,7 @@
 import numpy as np
 
 from parcels._core.index_search import GRID_SEARCH_ERROR, LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS, _search_time_index
+from parcels._core.particleset import ParticleSet
 from parcels._core.particlesetview import ParticleSetView
 from parcels._core.statuscodes import (
     AllParcelsErrorCodes,
@@ -186,7 +187,7 @@ def eval(self, t: datetime, z, y, x, particles=None):
     def __getitem__(self, key):
         self._check_velocitysampling()
         try:
-            if isinstance(key, ParticleSetView):
+            if isinstance(key, (ParticleSetView, ParticleSet)):
                 return self.eval(key.t, key.z, key.y, key.x, key)
             else:
                 return self.eval(*key)
@@ -295,7 +296,7 @@ def eval(self, t: datetime, z, y, x, particles=None):
 
     def __getitem__(self, key):
         try:
-            if isinstance(key, ParticleSetView):
+            if isinstance(key, (ParticleSetView, ParticleSet)):
                 return self.eval(key.t, key.z, key.y, key.x, key)
             else:
                 return self.eval(*key)

tests/test_particlefile.py

@@ -304,6 +304,29 @@ def IncreaseAge(particles, fieldset):  # pragma: no cover
         assert (df_traj["age"] == traj_time).all()
 
 
+@pytest.mark.parametrize("npart", [1, 10])
+def test_sampling_initial_value(fieldset, npart, tmp_parquet):
+    # Test that inital value of a field gets sampled
+
+    SampleParticle = get_default_particle(np.float64).add_variable(Variable("sample", initial=np.nan))
+
+    def SampleKernel(particles, fieldset):  # pragma: no cover
+        particles.sample = fieldset.U[particles]
+
+    x = np.zeros(npart)
+    y = np.zeros(npart)
+    t = np.zeros(npart, dtype="timedelta64[s]")
+
+    pset = ParticleSet(fieldset, pclass=SampleParticle, x=x, y=y, t=t)
+    pset.sample = fieldset.U[pset]  # Sample initial value
+
+    ofile = ParticleFile(tmp_parquet, outputdt=np.timedelta64(1, "s"))
+    pset.execute(SampleKernel, runtime=np.timedelta64(2, "s"), dt=np.timedelta64(1, "s"), output_file=ofile)
+
+    df = parcels.read_particlefile(tmp_parquet)
+    np.testing.assert_allclose(df["sample"].is_finite().all(), True)
+
+
 def test_reset_dt(fieldset, tmp_parquet):
     # Assert that p.dt gets reset when a write_time is not a multiple of dt
     # for p.dt=0.02 to reach outputdt=0.05 and endtime=0.1, the steps should be [0.2, 0.2, 0.1, 0.2, 0.2, 0.1], resulting in 6 kernel executions