Move from Parcels v3 to v4

environment.yml

@@ -1,20 +1,23 @@
-name: ship
+name: ship_parcelsv4 #TODO: revert back to 'ship' before proper release...
 channels:
   - conda-forge
+  - https://repo.prefix.dev/parcels
 dependencies:
   - click
-  - parcels >3.1.0
+  - parcels =4.0.0alpha0
   - pyproj >= 3, < 4
   - sortedcontainers == 2.4.0
   - opensimplex == 0.4.5
-  - numpy >=1, < 2
+  - numpy >=2.1
   - pydantic >=2, <3
   - pip
   - pyyaml
   - copernicusmarine >= 2.2.2
   - openpyxl
   - yaspin
   - textual
+  # - pip:
+  #     - git+https://github.com/OceanParcels/parcels.git@v4-dev
 
   # linting
   - pre-commit

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 ]
 dependencies = [
     "click",
-    "parcels >3.1.0",
+    "parcels @ git+https://github.com/OceanParcels/parcels.git@v4-dev",
     "pyproj >= 3, < 4",
     "sortedcontainers == 2.4.0",
     "opensimplex == 0.4.5",
@@ -68,7 +68,11 @@ filterwarnings = [
   "error",
   "default::DeprecationWarning",
   "error::DeprecationWarning:virtualship",
-  "ignore:ParticleSet is empty.*:RuntimeWarning" # TODO: Probably should be ignored in the source code
+  "ignore:ParticleSet is empty.*:RuntimeWarning", # TODO: Probably should be ignored in the source code
+  "ignore:divide by zero encountered *:RuntimeWarning",
+  "ignore:invalid value encountered *:RuntimeWarning",
+  "ignore:This is an alpha version of Parcels v4*:UserWarning",
+  "ignore:numpy.ndarray size changed*:RuntimeWarning",
 ]
 log_cli_level = "INFO"
 testpaths = [

src/virtualship/expedition/input_data.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+import xarray as xr
 from parcels import Field, FieldSet
 
 
@@ -95,40 +96,16 @@ def _load_ship_fieldset(cls, directory: Path) -> FieldSet:
             "V": directory.joinpath("ship_uv.nc"),
             "S": directory.joinpath("ship_s.nc"),
             "T": directory.joinpath("ship_t.nc"),
+            "bathymetry": directory.joinpath("bathymetry.nc"),
         }
-        variables = {"U": "uo", "V": "vo", "S": "so", "T": "thetao"}
-        dimensions = {
-            "lon": "longitude",
-            "lat": "latitude",
-            "time": "time",
-            "depth": "depth",
-        }
-
-        # create the fieldset and set interpolation methods
-        fieldset = FieldSet.from_netcdf(
-            filenames, variables, dimensions, allow_time_extrapolation=True
+        ds = xr.open_mfdataset(
+            [filenames["U"], filenames["T"], filenames["S"], filenames["bathymetry"]]
         )
-        fieldset.T.interp_method = "linear_invdist_land_tracer"
-        fieldset.S.interp_method = "linear_invdist_land_tracer"
-
-        # make depth negative
-        for g in fieldset.gridset.grids:
-            g.negate_depth()
-
-        # add bathymetry data
-        bathymetry_file = directory.joinpath("bathymetry.nc")
-        bathymetry_variables = ("bathymetry", "deptho")
-        bathymetry_dimensions = {"lon": "longitude", "lat": "latitude"}
-        bathymetry_field = Field.from_netcdf(
-            bathymetry_file, bathymetry_variables, bathymetry_dimensions
-        )
-        # make depth negative
-        bathymetry_field.data = -bathymetry_field.data
-        fieldset.add_field(bathymetry_field)
-
-        # read in data already
-        fieldset.computeTimeChunk(0, 1)
-
+        ds = ds.rename_vars({"deptho": "bathymetry"})
+        ds["bathymetry"] = -ds["bathymetry"]
+        ds["depth"] = -ds["depth"]
+        ds = ds.rename({"so": "S", "thetao": "T"})
+        fieldset = FieldSet.from_copernicusmarine(ds)
         return fieldset
 
     @classmethod
@@ -203,26 +180,9 @@ def _load_drifter_fieldset(cls, directory: Path) -> FieldSet:
             "V": directory.joinpath("drifter_uv.nc"),
             "T": directory.joinpath("drifter_t.nc"),
         }
-        variables = {"U": "uo", "V": "vo", "T": "thetao"}
-        dimensions = {
-            "lon": "longitude",
-            "lat": "latitude",
-            "time": "time",
-            "depth": "depth",
-        }
-
-        fieldset = FieldSet.from_netcdf(
-            filenames, variables, dimensions, allow_time_extrapolation=False
-        )
-        fieldset.T.interp_method = "linear_invdist_land_tracer"
-
-        # make depth negative
-        for g in fieldset.gridset.grids:
-            g.negate_depth()
-
-        # read in data already
-        fieldset.computeTimeChunk(0, 1)
-
+        ds = xr.open_mfdataset([filenames["U"], filenames["T"]])
+        ds = ds.rename({"thetao": "T"})
+        fieldset = FieldSet.from_copernicusmarine(ds)
         return fieldset
 
     @classmethod

src/virtualship/instruments/adcp.py

@@ -3,23 +3,25 @@
 from pathlib import Path
 
 import numpy as np
-from parcels import FieldSet, ParticleSet, ScipyParticle, Variable
+from parcels import FieldSet, Particle, ParticleSet, Variable
 
 from virtualship.models import Spacetime
 
-# we specifically use ScipyParticle because we have many small calls to execute
-# there is some overhead with JITParticle and this ends up being significantly faster
-_ADCPParticle = ScipyParticle.add_variables(
+_ADCPParticle = Particle.add_variable(
     [
         Variable("U", dtype=np.float32, initial=np.nan),
         Variable("V", dtype=np.float32, initial=np.nan),
     ]
 )
 
 
-def _sample_velocity(particle, fieldset, time):
-    particle.U, particle.V = fieldset.UV.eval(
-        time, particle.depth, particle.lat, particle.lon, applyConversion=False
+def _sample_velocity(particles, fieldset):
+    particles.U, particles.V = fieldset.UV.eval(
+        particles.time,
+        particles.z,
+        particles.lat,
+        particles.lon,
+        applyConversion=False,
     )
 
 

src/virtualship/instruments/argo_float.py

@@ -1,19 +1,18 @@
 """Argo float instrument."""
 
-import math
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from pathlib import Path
 
 import numpy as np
 from parcels import (
-    AdvectionRK4,
     FieldSet,
-    JITParticle,
+    Particle,
     ParticleSet,
     StatusCode,
     Variable,
 )
+from parcels.kernels import AdvectionRK4
 
 from virtualship.models import Spacetime
 
@@ -31,7 +30,7 @@ class ArgoFloat:
     drift_days: float
 
 
-_ArgoParticle = JITParticle.add_variables(
+_ArgoParticle = Particle.add_variable(
     [
         Variable("cycle_phase", dtype=np.int32, initial=0.0),
         Variable("cycle_age", dtype=np.float32, initial=0.0),
@@ -48,71 +47,86 @@ class ArgoFloat:
 )
 
 
-def _argo_float_vertical_movement(particle, fieldset, time):
-    if particle.cycle_phase == 0:
-        # Phase 0: Sinking with vertical_speed until depth is drift_depth
-        particle_ddepth += (  # noqa Parcels defines particle_* variables, which code checkers cannot know.
-            particle.vertical_speed * particle.dt
-        )
-        if particle.depth + particle_ddepth <= particle.drift_depth:
-            particle_ddepth = particle.drift_depth - particle.depth
-            particle.cycle_phase = 1
-
-    elif particle.cycle_phase == 1:
-        # Phase 1: Drifting at depth for drifttime seconds
-        particle.drift_age += particle.dt
-        if particle.drift_age >= particle.drift_days * 86400:
-            particle.drift_age = 0  # reset drift_age for next cycle
-            particle.cycle_phase = 2
-
-    elif particle.cycle_phase == 2:
-        # Phase 2: Sinking further to max_depth
-        particle_ddepth += particle.vertical_speed * particle.dt
-        if particle.depth + particle_ddepth <= particle.max_depth:
-            particle_ddepth = particle.max_depth - particle.depth
-            particle.cycle_phase = 3
-
-    elif particle.cycle_phase == 3:
-        # Phase 3: Rising with vertical_speed until at surface
-        particle_ddepth -= particle.vertical_speed * particle.dt
-        particle.cycle_age += (
-            particle.dt
-        )  # solve issue of not updating cycle_age during ascent
-        if particle.depth + particle_ddepth >= particle.min_depth:
-            particle_ddepth = particle.min_depth - particle.depth
-            particle.temperature = (
-                math.nan
-            )  # reset temperature to NaN at end of sampling cycle
-            particle.salinity = math.nan  # idem
-            particle.cycle_phase = 4
-        else:
-            particle.temperature = fieldset.T[
-                time, particle.depth, particle.lat, particle.lon
-            ]
-            particle.salinity = fieldset.S[
-                time, particle.depth, particle.lat, particle.lon
-            ]
-
-    elif particle.cycle_phase == 4:
-        # Phase 4: Transmitting at surface until cycletime is reached
-        if particle.cycle_age > particle.cycle_days * 86400:
-            particle.cycle_phase = 0
-            particle.cycle_age = 0
-
-    if particle.state == StatusCode.Evaluate:
-        particle.cycle_age += particle.dt  # update cycle_age
-
-
-def _keep_at_surface(particle, fieldset, time):
+def ArgoPhase1(particles, fieldset):
+    dt = particles.dt / np.timedelta64(1, "s")  # convert dt to seconds
+
+    def SinkingPhase(p):
+        """Phase 0: Sinking with p.vertical_speed until depth is driftdepth."""
+        p.dz += p.verticle_speed * dt
+        p.cycle_phase = np.where(p.z + p.dz >= p.drift_depth, 1, p.cycle_phase)
+        p.dz = np.where(p.z + p.dz >= p.drift_depth, p.drift_depth - p.z, p.dz)
+
+    SinkingPhase(particles[particles.cycle_phase == 0])
+
+
+def ArgoPhase2(particles, fieldset):
+    dt = particles.dt / np.timedelta64(1, "s")  # convert dt to seconds
+
+    def DriftingPhase(p):
+        """Phase 1: Drifting at depth for drift_time seconds."""
+        p.drift_age += dt
+        p.cycle_phase = np.where(p.drift_age >= p.drift_time, 2, p.cycle_phase)
+        p.drift_age = np.where(p.drift_age >= p.drift_time, 0, p.drift_age)
+
+    DriftingPhase(particles[particles.cycle_phase == 1])
+
+
+def ArgoPhase3(particles, fieldset):
+    dt = particles.dt / np.timedelta64(1, "s")  # convert dt to seconds
+
+    def SecondSinkingPhase(p):
+        """Phase 2: Sinking further to max_depth."""
+        p.dz += p.vertical_speed * dt
+        p.cycle_phase = np.where(p.z + p.dz >= p.max_depth, 3, p.cycle_phase)
+        p.dz = np.where(p.z + p.dz >= p.max_depth, p.max_depth - p.z, p.dz)
+
+    SecondSinkingPhase(particles[particles.cycle_phase == 2])
+
+
+def ArgoPhase4(particles, fieldset):
+    dt = particles.dt / np.timedelta64(1, "s")  # convert dt to seconds
+
+    def RisingPhase(p):
+        """Phase 3: Rising with p.vertical_speed until at surface."""
+        p.dz -= p.vertical_speed * dt
+        p.temp = fieldset.temp[p.time, p.z, p.lat, p.lon]
+        p.cycle_phase = np.where(p.z + p.dz <= fieldset.mindepth, 4, p.cycle_phase)
+
+    RisingPhase(particles[particles.cycle_phase == 3])
+
+
+def ArgoPhase5(particles, fieldset):
+    def TransmittingPhase(p):
+        """Phase 4: Transmitting at surface until cycletime (cycle_days * 86400 [seconds]) is reached."""
+        p.cycle_phase = np.where(p.cycle_age >= p.cycle_days * 86400, 0, p.cycle_phase)
+        p.cycle_age = np.where(p.cycle_age >= p.cycle_days * 86400, 0, p.cycle_age)
+
+    TransmittingPhase(particles[particles.cycle_phase == 4])
+
+
+def ArgoPhase6(particles, fieldset):
+    dt = particles.dt / np.timedelta64(1, "s")  # convert dt to seconds
+    particles.cycle_age += dt  # update cycle_age
+
+
+def _keep_at_surface(particles, fieldset):
     # Prevent error when float reaches surface
-    if particle.state == StatusCode.ErrorThroughSurface:
-        particle.depth = particle.min_depth
-        particle.state = StatusCode.Success
+    particles.z = np.where(
+        particles.state == StatusCode.ErrorThroughSurface,
+        particles.min_depth,
+        particles.z,
+    )
+    particles.state = np.where(
+        particles.state == StatusCode.ErrorThroughSurface,
+        StatusCode.Success,
+        particles.state,
+    )
 
 
-def _check_error(particle, fieldset, time):
-    if particle.state >= 50:  # This captures all Errors
-        particle.delete()
+def _check_error(particles, fieldset):
+    particles.state = np.where(
+        particles.state >= 50, StatusCode.Delete, particles.state
+    )  # captures all errors
 
 
 def simulate_argo_floats(
@@ -174,7 +188,12 @@ def simulate_argo_floats(
     # execute simulation
     argo_float_particleset.execute(
         [
-            _argo_float_vertical_movement,
+            ArgoPhase1,
+            ArgoPhase2,
+            ArgoPhase3,
+            ArgoPhase4,
+            ArgoPhase5,
+            ArgoPhase6,
             AdvectionRK4,
             _keep_at_surface,
             _check_error,