Merge pull request #1501 from OceanParcels/paserticleset_add_variable_method

Implementing Particleset.add_variable() method

Author: erikvansebille

docs/examples/documentation_stuck_particles.ipynb

@@ -82,7 +82,6 @@
     "    FieldSet,\n",
     "    JITParticle,\n",
     "    ParticleSet,\n",
-    "    Variable,\n",
     "    download_example_dataset,\n",
     ")"
    ]
@@ -1222,8 +1221,7 @@
    },
    "outputs": [],
    "source": [
-    "class LandParticle(JITParticle):\n",
-    "    on_land = Variable(\"on_land\")\n",
+    "LandParticle = JITParticle.add_variable(\"on_land\")\n",
     "\n",
     "\n",
     "def Sample_land(particle, fieldset, time):\n",

docs/examples/documentation_unstuck_Agrid.ipynb

@@ -809,10 +809,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "class DisplacementParticle(JITParticle):\n",
-    "    dU = Variable(\"dU\")\n",
-    "    dV = Variable(\"dV\")\n",
-    "    d2s = Variable(\"d2s\", initial=1e3)\n",
+    "DisplacementParticle = JITParticle.add_variables(\n",
+    "    [\n",
+    "        Variable(\"dU\"),\n",
+    "        Variable(\"dV\"),\n",
+    "        Variable(\"d2s\", initial=1e3),\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "\n",
     "def set_displacement(particle, fieldset, time):\n",

docs/examples/example_globcurrent.py

@@ -13,7 +13,6 @@
     ParticleSet,
     ScipyParticle,
     TimeExtrapolationError,
-    Variable,
     download_example_dataset,
 )
 
@@ -99,8 +98,7 @@ def test_globcurrent_time_periodic(mode, rundays):
     for deferred_load in [True, False]:
         fieldset = set_globcurrent_fieldset(time_periodic=delta(days=365), deferred_load=deferred_load)
 
-        class MyParticle(ptype[mode]):
-            sample_var = Variable('sample_var', initial=0.)
+        MyParticle = ptype[mode].add_variable('sample_var', initial=0.)
 
         pset = ParticleSet(fieldset, pclass=MyParticle, lon=25, lat=-35, time=fieldset.U.grid.time[0])
 
@@ -194,8 +192,7 @@ def test_globcurrent_startparticles_between_time_arrays(mode, dt, with_starttime
     fieldset.add_field(Field.from_netcdf(fnamesFeb, ('P', 'eastward_eulerian_current_velocity'),
                                          {'lat': 'lat', 'lon': 'lon', 'time': 'time'}))
 
-    class MyParticle(ptype[mode]):
-        sample_var = Variable('sample_var', initial=0.)
+    MyParticle = ptype[mode].add_variable('sample_var', initial=0.)
 
     def SampleP(particle, fieldset, time):
         particle.sample_var += fieldset.P[time, particle.depth, particle.lat, particle.lon]

docs/examples/example_peninsula.py

@@ -130,11 +130,8 @@ def peninsula_example(fieldset, outfile, npart, mode='jit', degree=1,
     # First, we define a custom Particle class to which we add a
     # custom variable, the initial stream function value p.
     # We determine the particle base class according to mode.
-    class MyParticle(ptype[mode]):
-        # JIT compilation requires a-priori knowledge of the particle
-        # data structure, so we define additional variables here.
-        p = Variable('p', dtype=np.float32, initial=0.)
-        p_start = Variable('p_start', dtype=np.float32, initial=0)
+    MyParticle = ptype[mode].add_variable([Variable('p', dtype=np.float32, initial=0.),
+                                           Variable('p_start', dtype=np.float32, initial=0)])
 
     # Initialise particles
     if fieldset.U.grid.mesh == 'flat':

docs/examples/example_stommel.py

@@ -105,11 +105,11 @@ def stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4, gri
     dt = delta(hours=1)
     outputdt = delta(days=5)
 
-    class MyParticle(ParticleClass):
-        p = Variable('p', dtype=np.float32, initial=0.)
-        p_start = Variable('p_start', dtype=np.float32, initial=0.)
-        next_dt = Variable('next_dt', dtype=np.float64, initial=dt.total_seconds())
-        age = Variable('age', dtype=np.float32, initial=0.)
+    extra_vars = [Variable('p', dtype=np.float32, initial=0.),
+                  Variable('p_start', dtype=np.float32, initial=0.),
+                  Variable('next_dt', dtype=np.float64, initial=dt.total_seconds()),
+                  Variable('age', dtype=np.float32, initial=0.)]
+    MyParticle = ParticleClass.add_variables(extra_vars)
 
     if custom_partition_function:
         pset = ParticleSet.from_line(fieldset, size=npart, pclass=MyParticle, repeatdt=repeatdt,

docs/examples/parcels_tutorial.ipynb

@@ -228,7 +228,7 @@
      "output_type": "stream",
      "text": [
       "INFO: Output files are stored in EddyParticles.zarr.\n",
-      "100%|██████████| 518400.0/518400.0 [00:02<00:00, 187445.76it/s]\n"
+      "100%|██████████| 518400.0/518400.0 [00:03<00:00, 172443.78it/s]\n"
      ]
     }
    ],
@@ -569,42 +569,42 @@
        "</style>\n",
        "\n",
        "<div class=\"animation\">\n",
-       "  <img id=\"_anim_img8bfbea7abfae4cbf885a957f3f62a93b\">\n",
+       "  <img id=\"_anim_img39fc9ac9ecc54c368c0d0e047c7673e6\">\n",
        "  <div class=\"anim-controls\">\n",
-       "    <input id=\"_anim_slider8bfbea7abfae4cbf885a957f3f62a93b\" type=\"range\" class=\"anim-slider\"\n",
+       "    <input id=\"_anim_slider39fc9ac9ecc54c368c0d0e047c7673e6\" type=\"range\" class=\"anim-slider\"\n",
        "           name=\"points\" min=\"0\" max=\"1\" step=\"1\" value=\"0\"\n",
-       "           oninput=\"anim8bfbea7abfae4cbf885a957f3f62a93b.set_frame(parseInt(this.value));\">\n",
+       "           oninput=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.set_frame(parseInt(this.value));\">\n",
        "    <div class=\"anim-buttons\">\n",
-       "      <button title=\"Decrease speed\" aria-label=\"Decrease speed\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.slower()\">\n",
+       "      <button title=\"Decrease speed\" aria-label=\"Decrease speed\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.slower()\">\n",
        "          <i class=\"fa fa-minus\"></i></button>\n",
-       "      <button title=\"First frame\" aria-label=\"First frame\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.first_frame()\">\n",
+       "      <button title=\"First frame\" aria-label=\"First frame\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.first_frame()\">\n",
        "        <i class=\"fa fa-fast-backward\"></i></button>\n",
-       "      <button title=\"Previous frame\" aria-label=\"Previous frame\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.previous_frame()\">\n",
+       "      <button title=\"Previous frame\" aria-label=\"Previous frame\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.previous_frame()\">\n",
        "          <i class=\"fa fa-step-backward\"></i></button>\n",
-       "      <button title=\"Play backwards\" aria-label=\"Play backwards\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.reverse_animation()\">\n",
+       "      <button title=\"Play backwards\" aria-label=\"Play backwards\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.reverse_animation()\">\n",
        "          <i class=\"fa fa-play fa-flip-horizontal\"></i></button>\n",
-       "      <button title=\"Pause\" aria-label=\"Pause\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.pause_animation()\">\n",
+       "      <button title=\"Pause\" aria-label=\"Pause\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.pause_animation()\">\n",
        "          <i class=\"fa fa-pause\"></i></button>\n",
-       "      <button title=\"Play\" aria-label=\"Play\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.play_animation()\">\n",
+       "      <button title=\"Play\" aria-label=\"Play\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.play_animation()\">\n",
        "          <i class=\"fa fa-play\"></i></button>\n",
-       "      <button title=\"Next frame\" aria-label=\"Next frame\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.next_frame()\">\n",
+       "      <button title=\"Next frame\" aria-label=\"Next frame\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.next_frame()\">\n",
        "          <i class=\"fa fa-step-forward\"></i></button>\n",
-       "      <button title=\"Last frame\" aria-label=\"Last frame\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.last_frame()\">\n",
+       "      <button title=\"Last frame\" aria-label=\"Last frame\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.last_frame()\">\n",
        "          <i class=\"fa fa-fast-forward\"></i></button>\n",
-       "      <button title=\"Increase speed\" aria-label=\"Increase speed\" onclick=\"anim8bfbea7abfae4cbf885a957f3f62a93b.faster()\">\n",
+       "      <button title=\"Increase speed\" aria-label=\"Increase speed\" onclick=\"anim39fc9ac9ecc54c368c0d0e047c7673e6.faster()\">\n",
        "          <i class=\"fa fa-plus\"></i></button>\n",
        "    </div>\n",
-       "    <form title=\"Repetition mode\" aria-label=\"Repetition mode\" action=\"#n\" name=\"_anim_loop_select8bfbea7abfae4cbf885a957f3f62a93b\"\n",
+       "    <form title=\"Repetition mode\" aria-label=\"Repetition mode\" action=\"#n\" name=\"_anim_loop_select39fc9ac9ecc54c368c0d0e047c7673e6\"\n",
        "          class=\"anim-state\">\n",
-       "      <input type=\"radio\" name=\"state\" value=\"once\" id=\"_anim_radio1_8bfbea7abfae4cbf885a957f3f62a93b\"\n",
+       "      <input type=\"radio\" name=\"state\" value=\"once\" id=\"_anim_radio1_39fc9ac9ecc54c368c0d0e047c7673e6\"\n",
        "             >\n",
-       "      <label for=\"_anim_radio1_8bfbea7abfae4cbf885a957f3f62a93b\">Once</label>\n",
-       "      <input type=\"radio\" name=\"state\" value=\"loop\" id=\"_anim_radio2_8bfbea7abfae4cbf885a957f3f62a93b\"\n",
+       "      <label for=\"_anim_radio1_39fc9ac9ecc54c368c0d0e047c7673e6\">Once</label>\n",
+       "      <input type=\"radio\" name=\"state\" value=\"loop\" id=\"_anim_radio2_39fc9ac9ecc54c368c0d0e047c7673e6\"\n",
        "             checked>\n",
-       "      <label for=\"_anim_radio2_8bfbea7abfae4cbf885a957f3f62a93b\">Loop</label>\n",
-       "      <input type=\"radio\" name=\"state\" value=\"reflect\" id=\"_anim_radio3_8bfbea7abfae4cbf885a957f3f62a93b\"\n",
+       "      <label for=\"_anim_radio2_39fc9ac9ecc54c368c0d0e047c7673e6\">Loop</label>\n",
+       "      <input type=\"radio\" name=\"state\" value=\"reflect\" id=\"_anim_radio3_39fc9ac9ecc54c368c0d0e047c7673e6\"\n",
        "             >\n",
-       "      <label for=\"_anim_radio3_8bfbea7abfae4cbf885a957f3f62a93b\">Reflect</label>\n",
+       "      <label for=\"_anim_radio3_39fc9ac9ecc54c368c0d0e047c7673e6\">Reflect</label>\n",
        "    </form>\n",
        "  </div>\n",
        "</div>\n",
@@ -614,9 +614,9 @@
        "  /* Instantiate the Animation class. */\n",
        "  /* The IDs given should match those used in the template above. */\n",
        "  (function() {\n",
-       "    var img_id = \"_anim_img8bfbea7abfae4cbf885a957f3f62a93b\";\n",
-       "    var slider_id = \"_anim_slider8bfbea7abfae4cbf885a957f3f62a93b\";\n",
-       "    var loop_select_id = \"_anim_loop_select8bfbea7abfae4cbf885a957f3f62a93b\";\n",
+       "    var img_id = \"_anim_img39fc9ac9ecc54c368c0d0e047c7673e6\";\n",
+       "    var slider_id = \"_anim_slider39fc9ac9ecc54c368c0d0e047c7673e6\";\n",
+       "    var loop_select_id = \"_anim_loop_select39fc9ac9ecc54c368c0d0e047c7673e6\";\n",
        "    var frames = new Array(29);\n",
        "    \n",
        "  frames[0] = \"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAArwAAAH0CAYAAADfWf7fAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90\\\n",
@@ -13156,7 +13156,7 @@
        "    /* set a timeout to make sure all the above elements are created before\n",
        "       the object is initialized. */\n",
        "    setTimeout(function() {\n",
-       "        anim8bfbea7abfae4cbf885a957f3f62a93b = new Animation(frames, img_id, slider_id, 100.0,\n",
+       "        anim39fc9ac9ecc54c368c0d0e047c7673e6 = new Animation(frames, img_id, slider_id, 100.0,\n",
        "                                 loop_select_id);\n",
        "    }, 0);\n",
        "  })()\n",
@@ -13201,7 +13201,7 @@
      "output_type": "stream",
      "text": [
       "INFO: Output files are stored in EddyParticles_Bwd.zarr.\n",
-      "100%|██████████| 518400.0/518400.0 [00:02<00:00, 188464.76it/s]\n"
+      "100%|██████████| 518400.0/518400.0 [00:02<00:00, 176426.86it/s]\n"
      ]
     }
    ],
@@ -13323,7 +13323,7 @@
      "output_type": "stream",
      "text": [
       "INFO: Output files are stored in EddyParticles_WestVel.zarr.\n",
-      "100%|██████████| 172800.0/172800.0 [00:00<00:00, 180626.58it/s]\n"
+      "100%|██████████| 172800.0/172800.0 [00:00<00:00, 179532.85it/s]\n"
      ]
     }
    ],
@@ -13503,7 +13503,7 @@
      "output_type": "stream",
      "text": [
       "INFO: Output files are stored in GlobCurrentParticles.zarr.\n",
-      "100%|██████████| 864000.0/864000.0 [00:00<00:00, 1104294.63it/s]\n"
+      "100%|██████████| 864000.0/864000.0 [00:00<00:00, 1072517.72it/s]\n"
      ]
     }
    ],
@@ -13592,7 +13592,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now define a new `Particle` class that has an extra `Variable`: the pressure. We initialise this by sampling the `fieldset.P` field.\n"
+    "Now define a new `Particle` class that has an extra `Variable`: the pressure. This `particle.p` can be used to store the values of the `fieldset.P` field at the particle locations.\n"
    ]
   },
   {
@@ -13601,11 +13601,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "class SampleParticle(JITParticle):\n",
-    "    \"\"\"Define a new particle class with variable 'p'\n",
-    "    initialised by sampling the pressure\"\"\"\n",
-    "\n",
-    "    p = Variable(\"p\")"
+    "SampleParticle = JITParticle.add_variable(\"p\")"
    ]
   },
   {
@@ -13688,7 +13684,7 @@
      "output_type": "stream",
      "text": [
       "INFO: Output files are stored in PeninsulaPressure.zarr.\n",
-      "100%|██████████| 72000.0/72000.0 [00:00<00:00, 157840.98it/s]\n"
+      "100%|██████████| 72000.0/72000.0 [00:00<00:00, 143717.52it/s]\n"
      ]
     }
    ],
@@ -13768,7 +13764,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "First, we need to create a new `Particle` class that includes three extra variables. The `distance` variable will be written to output, but the auxiliary variables `prev_lon` and `prev_lat` won't be written to output (can be controlled using the `to_write` keyword)\n"
+    "First, we need to add three extra variables to the Particle Class. The `distance` variable will be written to output, but the auxiliary variables `prev_lon` and `prev_lat` won't be written to output (can be controlled using the `to_write` keyword)."
    ]
   },
   {
@@ -13777,19 +13773,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "class DistParticle(JITParticle):\n",
-    "    \"\"\"Define a new particle class that contains three extra variables\"\"\"\n",
-    "\n",
-    "    # the distance travelled by the particle\n",
-    "    distance = Variable(\"distance\", initial=0.0, dtype=np.float32)\n",
+    "extra_vars = [\n",
+    "    Variable(\"distance\", initial=0.0, dtype=np.float32),\n",
+    "    Variable(\"prev_lon\", dtype=np.float32, to_write=False, initial=attrgetter(\"lon\")),\n",
+    "    Variable(\"prev_lat\", dtype=np.float32, to_write=False, initial=attrgetter(\"lat\")),\n",
+    "]\n",
     "\n",
-    "    # the previous longitude and latitude of the particle\n",
-    "    prev_lon = Variable(\n",
-    "        \"prev_lon\", dtype=np.float32, to_write=False, initial=attrgetter(\"lon\")\n",
-    "    )\n",
-    "    prev_lat = Variable(\n",
-    "        \"prev_lat\", dtype=np.float32, to_write=False, initial=attrgetter(\"lat\")\n",
-    "    )"
+    "DistParticle = JITParticle.add_variables(extra_vars)"
    ]
   },
   {
@@ -13879,7 +13869,7 @@
      "output_type": "stream",
      "text": [
       "INFO: Output files are stored in GlobCurrentParticles_Dist.zarr.\n",
-      "100%|██████████| 518400.0/518400.0 [00:03<00:00, 147724.13it/s]\n"
+      "100%|██████████| 518400.0/518400.0 [00:03<00:00, 136275.28it/s]\n"
      ]
     }
    ],
@@ -13936,7 +13926,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.11.6"
   }
  },
  "nbformat": 4,

docs/examples/tutorial_Argofloats.ipynb

@@ -128,19 +128,21 @@
     "\n",
     "\n",
     "# Define a new Particle type including extra Variables\n",
-    "class ArgoParticle(JITParticle):\n",
-    "    # Phase of cycle:\n",
-    "    # init_descend=0,\n",
-    "    # drift=1,\n",
-    "    # profile_descend=2,\n",
-    "    # profile_ascend=3,\n",
-    "    # transmit=4\n",
-    "    cycle_phase = Variable(\"cycle_phase\", dtype=np.int32, initial=0.0)\n",
-    "    cycle_age = Variable(\"cycle_age\", dtype=np.float32, initial=0.0)\n",
-    "    drift_age = Variable(\"drift_age\", dtype=np.float32, initial=0.0)\n",
-    "    # if fieldset has temperature\n",
-    "    # temp = Variable('temp', dtype=np.float32, initial=np.nan)\n",
-    "\n",
+    "ArgoParticle = JITParticle.add_variables(\n",
+    "    [\n",
+    "        # Phase of cycle:\n",
+    "        # init_descend=0,\n",
+    "        # drift=1,\n",
+    "        # profile_descend=2,\n",
+    "        # profile_ascend=3,\n",
+    "        # transmit=4\n",
+    "        Variable(\"cycle_phase\", dtype=np.int32, initial=0.0),\n",
+    "        Variable(\"cycle_age\", dtype=np.float32, initial=0.0),\n",
+    "        Variable(\"drift_age\", dtype=np.float32, initial=0.0),\n",
+    "        # if fieldset has temperature\n",
+    "        # Variable('temp', dtype=np.float32, initial=np.nan),\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "# Initiate one Argo float in the Agulhas Current\n",
     "pset = ParticleSet(\n",

docs/examples/tutorial_NestedFields.ipynb

@@ -226,16 +226,11 @@
     }
    ],
    "source": [
-    "from parcels import Variable\n",
-    "\n",
-    "\n",
     "def SampleNestedFieldIndex(particle, fieldset, time):\n",
     "    particle.f = fieldset.F[time, particle.depth, particle.lat, particle.lon]\n",
     "\n",
     "\n",
-    "class SampleParticle(JITParticle):\n",
-    "    f = Variable(\"f\", dtype=np.int32)\n",
-    "\n",
+    "SampleParticle = JITParticle.add_variable(\"f\", dtype=np.int32)\n",
     "\n",
     "pset = ParticleSet(fieldset, pclass=SampleParticle, lon=[1000], lat=[500])\n",
     "pset.execute(SampleNestedFieldIndex, runtime=1)\n",

docs/examples/tutorial_analyticaladvection.ipynb

@@ -154,9 +154,12 @@
     "    particle.radius = fieldset.R[time, particle.depth, particle.lat, particle.lon]\n",
     "\n",
     "\n",
-    "class MyParticle(ScipyParticle):\n",
-    "    radius = Variable(\"radius\", dtype=np.float32, initial=0.0)\n",
-    "    radius_start = Variable(\"radius_start\", dtype=np.float32, initial=0.0)\n",
+    "MyParticle = ScipyParticle.add_variables(\n",
+    "    [\n",
+    "        Variable(\"radius\", dtype=np.float32, initial=0.0),\n",
+    "        Variable(\"radius_start\", dtype=np.float32, initial=0.0),\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "\n",
     "pset = ParticleSet(fieldsetRR, pclass=MyParticle, lon=0, lat=4e3, time=0)\n",

docs/examples/tutorial_delaystart.ipynb

@@ -32416,10 +32416,13 @@
     }
    ],
    "source": [
-    "class GrowingParticle(JITParticle):\n",
-    "    mass = Variable(\"mass\", initial=0)\n",
-    "    splittime = Variable(\"splittime\", initial=-1)\n",
-    "    splitmass = Variable(\"splitmass\", initial=0)\n",
+    "GrowingParticle = JITParticle.add_variables(\n",
+    "    [\n",
+    "        Variable(\"mass\", initial=0),\n",
+    "        Variable(\"splittime\", initial=-1),\n",
+    "        Variable(\"splitmass\", initial=0),\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "\n",
     "def GrowParticles(particle, fieldset, time):\n",