Preserving particle ID between loops using MPI

[julia-neme]

Question

Question

Hi! I am new to Parcels, so apologies if this question has been answered elsewhere - I've done a search but couldn't find an answer.

I am using version 3.1.4. to backtrack particles, releasing every 7 days during the first year. I am running this on an HPC using MPI, and due to walltime/efficiency limitations I am running loops of 7 days duration. I have found that particle ID is not preserved across loops. Is there a way to preserve IDs using MPI? My code is below.

Thank you very much in advance!!

Supporting code/error messages

# coding: utf-8

import numpy as np
import xarray as xr
import sys
import parcels
import warnings
from datetime import timedelta
from  glob import glob
from parcels import FieldSet, ParticleSet, Variable, JITParticle, AdvectionRK4_3D, StatusCode

# check we are using version 3.1.0
print('parcels', parcels.__version__)
print(parcels.__file__)

warnings.filterwarnings("ignore", category=FutureWarning, message=".*xarray will not decode timedelta values.*")
warnings.filterwarnings("ignore", category=UserWarning, message="The specified chunks separate the stored chunks")

# get run count argument and gdata_dir that was passed to python script:
run_count = int(sys.argv[1])

output_directory = '/scratch/v45/jn8053/CDW_pathways/sector1_missing_test/loop_'+str(run_count).zfill(5)+'/'

######## Parcels Code below ### ---------------------------------
data_path = '/g/data/cj50/access-om2/raw-output/access-om2-01/01deg_jra55v140_iaf/'
ufiles = sorted(glob(data_path+'output*/ocean/ocean-3d-u-1-daily-mean-ym_*.nc'))
vfiles = sorted(glob(data_path+'output*/ocean/ocean-3d-v-1-daily-mean-ym_*.nc'))
wfiles = sorted(glob(data_path+'output*/ocean/ocean-3d-wt-1-daily-mean-ym_*.nc'))
tfiles = sorted(glob(data_path+'output*/ocean/ocean-3d-temp-1-daily-mean-ym_*.nc'))
sfiles = sorted(glob(data_path+'output*/ocean/ocean-3d-salt-1-daily-mean-ym_*.nc'))
mfiles = sorted(glob(data_path+'output*/ocean/ocean-2d-mld-1-daily-mean-ym_*.nc'))

# Define mesh mask which contains the coordinates
mesh_mask = '/scratch/x77/jn8053/CDW_pathways/ocean.nc'

# Define a dictionary that tells Parcels where the data is,
# and where to find the coordinates. 

filenames = {'U': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': mesh_mask, 'data': ufiles},
             'V': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': mesh_mask, 'data': vfiles},
             'W': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': mesh_mask, 'data': wfiles},
             'T': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': mesh_mask, 'data': tfiles},
             'S': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': mesh_mask, 'data': sfiles},
             'M': {'lon': mesh_mask, 'lat': mesh_mask, 'data': mfiles}
             }
# Now define a dictionary that specifies the variable names as given in the netCDF files  
variables = {'U': 'u', 'V': 'v', 'W': 'wt', 'T': 'temp', 'S': 'salt', 'M': 'mld'}

dimensions = {'U': {'lon': 'xu_ocean', 'lat': 'yu_ocean', 'depth': 'sw_ocean', 'time': 'time'},
              'V': {'lon': 'xu_ocean', 'lat': 'yu_ocean', 'depth': 'sw_ocean', 'time': 'time'},
              'W': {'lon': 'xu_ocean', 'lat': 'yu_ocean', 'depth': 'sw_ocean', 'time': 'time'},
              'T': {'lon': 'xu_ocean', 'lat': 'yu_ocean', 'depth': 'sw_ocean', 'time': 'time'},
              'S': {'lon': 'xu_ocean', 'lat': 'yu_ocean', 'depth': 'sw_ocean', 'time': 'time'},
              'M': {'lon': 'xu_ocean', 'lat': 'yu_ocean', 'time': 'time'}
              }

cs = {"U": {"lon": ("xu_ocean", 400), "lat": ("yu_ocean", 300), "depth": ("st_ocean", 75), "time": ("time", 1)}, 
      "V": {"lon": ("xu_ocean", 400), "lat": ("yu_ocean", 300), "depth": ("st_ocean", 75), "time": ("time", 1)},
      "W": {"lon": ("xt_ocean", 400), "lat": ("yt_ocean", 300), "depth": ("sw_ocean", 75), "time": ("time", 1)},
      "T": {"lon": ("xt_ocean", 400), "lat": ("yt_ocean", 300), "depth": ("st_ocean", 75), "time": ("time", 1)},
      "S": {"lon": ("xt_ocean", 400), "lat": ("yt_ocean", 300), "depth": ("st_ocean", 75), "time": ("time", 1)},
      "M": {"lon": ("xt_ocean", 400), "lat": ("yt_ocean", 300), "time": ("time", 1)}
      }

# Read in the fieldset using the Parcels `FieldSet.from_mom5` function. 
fieldset = FieldSet.from_mom5(filenames, variables, dimensions,
                              mesh = 'spherical', 
                              chunksize = cs, 
                              tracer_interp_method = 'bgrid_tracer')
fieldset.add_constant('halo_west', fieldset.U.grid.lon[0])
fieldset.add_constant('halo_east', fieldset.U.grid.lon[-1])
fieldset.add_periodic_halo(zonal=True)

def periodicBC(particle, fieldset, time):
    if particle.lon < fieldset.halo_west:
        particle_dlon += fieldset.halo_east - fieldset.halo_west
    elif particle.lon > fieldset.halo_east:
        particle_dlon -= fieldset.halo_east - fieldset.halo_west

def checkOutOfBounds(particle, fieldset, time):
    if particle.state == StatusCode.ErrorOutOfBounds:
        particle.delete()

class SampleParticle(JITParticle):
    temp = Variable('temp', dtype=np.float32, initial = 0.0)
    salt = Variable('salt', dtype=np.float32, initial = 0.0)
    mld = Variable('mld', dtype=np.float32, initial = 0.0)
    inmld = Variable('inmld', dtype=np.float32, initial = 0.0)

# Kernel to sample the model temperature at the particle's position
def SampleTemp(particle, fieldset, time):
    particle.temp = fieldset.T[time, particle.depth, particle.lat, particle.lon]
    particle.salt = fieldset.S[time, particle.depth, particle.lat, particle.lon]
    particle.mld = fieldset.M[time, 0, particle.lat, particle.lon]

def mld_flag(particle, fieldset, time):
    if particle.mld > particle.depth:
        particle.inmld = 1

def DeleteParticle(particle, fieldset, time):
    if particle.inmld == 1:
        particle.delete()

#------------------------------------------------------------------------------------
# Define particle set here for first run:
# for first run locate release locations, else load from file:
print('Starting run '+str(run_count))

if run_count == 0:
    ####### Define the particle set for the first run
    ds = xr.open_dataarray('/scratch/x77/jn8053/CDW_pathways/release_1000m/sector1/release_locs.nc')
    lons = ds.sel(dim='lon')
    lats = ds.sel(dim='lat')
    depths = ds.sel(dim='depth')

    ######## Define the particle set ### -----------------------
    pset = ParticleSet.from_list(fieldset=fieldset, 
        pclass=SampleParticle, 
        lon=lons, 
        lat=lats, 
        depth=depths,
        time=fieldset.U.grid.time[-1],
        )
elif run_count > 0 and run_count < 52:
    restart_file = '/scratch/v45/jn8053/CDW_pathways/sector1_missing_test/loop_'+str(run_count-1).zfill(5)+'/CDWParticles_merged.zarr'
    # Define particle set from file:
    pset = ParticleSet.from_particlefile(fieldset=fieldset,
                                         pclass=SampleParticle,
                                         filename=restart_file,
                                         restart=True,
                                         restarttime=np.nanmin)
    # Define the particle set for the first run
    ds = xr.open_dataarray('/scratch/x77/jn8053/CDW_pathways/release_1000m/sector1/release_locs.nc')
    lons = ds.sel(dim='lon')
    lats = ds.sel(dim='lat')
    depths = ds.sel(dim='depth')
    time_restart = fieldset.U.grid.time[-7*run_count-1] 

    pset1 = ParticleSet.from_list(fieldset=fieldset,
            pclass=SampleParticle,  
            time = time_restart,
            lon = lons,      
            lat = lats,      
            depth = depths,    
            )

    pset.add(pset1)
                    
else:
    restart_file = '/scratch/v45/jn8053/CDW_pathways/sector1_missing_test/loop_'+str(run_count-1).zfill(5)+'/CDWParticles_merged.zarr'
    # Define particle set from file:
    pset = ParticleSet.from_particlefile(fieldset=fieldset,
                                         pclass=SampleParticle,
                                         filename=restart_file,
                                         restart=True,
                                         restarttime=np.nanmin)
    
# Set your output location: 
output_filename = 'CDWParticles.zarr'

output_file = pset.ParticleFile(name=output_directory + output_filename, 
                                outputdt=timedelta(days=7) 
                               ) 

kernels = [AdvectionRK4_3D, periodicBC, checkOutOfBounds, SampleTemp, mld_flag, DeleteParticle]

if run_count < 52:
    pset.execute(kernels, 
                 runtime=timedelta(days=8),   
                 dt=-timedelta(hours=1),        
                 output_file=output_file,     
                )
else:
    pset.execute(kernels, 
                 runtime=timedelta(days=50),   
                 dt=-timedelta(hours=1),        
                 output_file=output_file,     
                )
     
print('Run '+str(run_count)+' finished')

[erikvansebille]

Hi @julia-neme, thanks for reaching out - and for using Parcels!

I am not using v3 anymore so can't easily check, but I think the 'problem' is that you use MPI. Can you confirm that the particle ID is preserved if you don't use MPI with the code above?

An alternative solution may be to simply open the file and create a particles (including id) from it inside your script. Something like (not tested)

pfile = xr.open_zarr(str(filename))
pset = ParticleSet.from_list(
    fieldset=fieldset,
    pclass=SampleParticle,
    time=pfile["time].values[:, -1],
    lat=pfile["lat'].values[:, -1],
    lon=pfile["lon"].values[:, -1],
    depth=pfile["z"].values[:, -1],
    id=pfile["trajectory"].values[:],
)

See below for the code how the ParticleSet.from_particlefile() works under the hood - perhaps useful for inspiration/debugging?

Parcels/parcels/particleset.py

Lines 761 to 871 in e6b6c70

	def from_particlefile(
	cls, fieldset, pclass, filename, restart=True, restarttime=None, repeatdt=None, lonlatdepth_dtype=None, **kwargs
	):
	"""Initialise the ParticleSet from a zarr ParticleFile.
	This creates a new ParticleSet based on locations of all particles written
	in a zarr ParticleFile at a certain time. Particle IDs are preserved if restart=True
	Parameters
	----------
	fieldset : parcels.fieldset.FieldSet
	mod:`parcels.fieldset.FieldSet` object from which to sample velocity
	pclass : parcels.particle.JITParticle or parcels.particle.ScipyParticle
	Particle class. May be a particle class as defined in parcels, or a subclass defining a custom particle.
	filename : str
	Name of the particlefile from which to read initial conditions
	restart : bool
	BSignal if pset is used for a restart (default is True).
	In that case, Particle IDs are preserved.
	restarttime :
	time at which the Particles will be restarted. Default is the last time written.
	Alternatively, restarttime could be a time value (including np.datetime64) or
	a callable function such as np.nanmin. The last is useful when running with dt < 0.
	repeatdt : datetime.timedelta or float, optional
	Optional interval on which to repeat the release of the ParticleSet. Either timedelta object, or float in seconds.
	lonlatdepth_dtype :
	Floating precision for lon, lat, depth particle coordinates.
	It is either np.float32 or np.float64. Default is np.float32 if fieldset.U.interp_method is 'linear'
	and np.float64 if the interpolation method is 'cgrid_velocity'
	**kwargs :
	Keyword arguments passed to the particleset constructor.
	"""
	if repeatdt is not None:
	warnings.warn(
	f"Note that the `repeatdt` argument is not retained from {filename}, and that "
	"setting a new repeatdt will start particles from the _new_ particle "
	"locations.",
	ParticleSetWarning,
	stacklevel=2,
	)
	pfile = xr.open_zarr(str(filename))
	pfile_vars = [v for v in pfile.data_vars]
	vars = {}
	to_write = {}
	for v in pclass.getPType().variables:
	if v.name in pfile_vars:
	vars[v.name] = np.ma.filled(pfile.variables[v.name], np.nan)
	elif (
	v.name
	not in [
	"xi",
	"yi",
	"zi",
	"ti",
	"dt",
	"depth",
	"id",
	"obs_written",
	"state",
	"lon_nextloop",
	"lat_nextloop",
	"depth_nextloop",
	"time_nextloop",
	]
	and v.to_write
	):
	raise RuntimeError(f"Variable {v.name} is in pclass but not in the particlefile")
	to_write[v.name] = v.to_write
	vars["depth"] = np.ma.filled(pfile.variables["z"], np.nan)
	vars["id"] = np.ma.filled(pfile.variables["trajectory"], np.nan)
	for v in ["lon", "lat", "depth", "time"]:
	to_write[v] = True
	if isinstance(vars["time"][0, 0], np.timedelta64):
	vars["time"] = np.array([t / np.timedelta64(1, "s") for t in vars["time"]])
	if restarttime is None:
	restarttime = np.nanmax(vars["time"])
	elif callable(restarttime):
	restarttime = restarttime(vars["time"])
	else:
	restarttime = restarttime
	inds = np.where(vars["time"] == restarttime)
	for v in vars:
	if to_write[v] is True:
	vars[v] = vars[v][inds]
	elif to_write[v] == "once":
	vars[v] = vars[v][inds[0]]
	if v not in ["lon", "lat", "depth", "time", "id"]:
	kwargs[v] = vars[v]
	if restart:
	pclass.setLastID(0) # reset to zero offset
	else:
	vars["id"] = None
	return cls(
	fieldset=fieldset,
	pclass=pclass,
	lon=vars["lon"],
	lat=vars["lat"],
	depth=vars["depth"],
	time=vars["time"],
	pid_orig=vars["id"],
	lonlatdepth_dtype=lonlatdepth_dtype,
	repeatdt=repeatdt,
	**kwargs,
	)

[julia-neme]

Hi @erikvansebille thank you for the prompt reply! After a bit more testing, you are right the problem was using MPI. MPI stores the output in a separate .zarr file for each processor, which means that after a run is finished I postprocess it into a single .zarr to feed to the next loop. I think there was an issue with my postprocessing (possible not sorting along trajectory), but it is working now!

Sorry about posting, I should have given it a few more days of work. Thank you again!!