Merge pull request #207 from OceanParcels/script_convert_indexed_outp…

…ut_to_array

Script to convert indexed output to array

.travis.yml

@@ -29,6 +29,11 @@ install:
     - pip install -r requirements.txt
     - git submodule update --init --recursive
 
+before_script: # configure a headless display to test plot generation
+    - "export DISPLAY=:99.0"
+    - "sh -e /etc/init.d/xvfb start"
+    - sleep 3 # give xvfb some time to start
+
 script:
     - export PYTHONPATH=`pwd`:$PYTHONPATH
     - flake8 parcels

parcels/scripts/__init__.py

@@ -1 +1,2 @@
 from .plotParticles import plotTrajectoriesFile  # NOQA get flake8 to ignore unused import.
+from .convert_IndexedOutputToArray import *  # NOQA

parcels/scripts/convert_IndexedOutputToArray.py

@@ -0,0 +1,72 @@
+from netCDF4 import Dataset
+import numpy as np
+from progressbar import ProgressBar
+from argparse import ArgumentParser
+
+
+def convert_IndexedOutputToArray(file_in, file_out):
+    """Script to convert a trajectory file as outputted by Parcels
+     in Indexed format to the easier-to-handle Array format
+
+     :param file_in: name of the input file in Indexed format
+     :param file_out: name of the output file"""
+
+    pfile_in = Dataset(file_in, 'r')
+    if 'trajectory' in pfile_in.dimensions:
+        print file_in+' appears to be in array format already. Doing nothing..'
+        return
+
+    class IndexedTrajectories(object):
+        """IndexedTrajectories class that holds info on the indices where the
+        individual particle trajectories start and end"""
+        def __init__(self, pfile):
+            self.ids = pfile.variables['trajectory'][:]
+            self.indices = np.argsort(self.ids)
+            self.ids = self.ids[self.indices]
+            self.starts = np.insert([x + 1 for x in np.where(np.diff(self.ids) > 0)], 0, 0)
+            self.ends = np.append(self.starts[1:], [len(self.ids)-1])
+            self.lengths = self.ends - self.starts
+            self.nid = len(self.starts)
+            self.nobs = np.max(self.lengths)
+            for i in range(self.nid):
+                # Test whether all ids in a trajectory are the same
+                assert all(self.ids[self.starts[i]:self.ends[i]] == self.ids[self.starts[i]])
+
+    trajs = IndexedTrajectories(pfile_in)
+
+    pfile_out = Dataset("%s" % file_out, "w", format="NETCDF4")
+    pfile_out.createDimension("obs", trajs.nobs)
+    pfile_out.createDimension("trajectory", trajs.nid)
+    coords = ("trajectory", "obs")
+
+    id = pfile_out.createVariable("trajectory", "i4", ("trajectory",))
+    id.long_name = "Unique identifier for each particle"
+    id.cf_role = "trajectory_id"
+    id[:] = np.array([trajs.ids[p] for p in trajs.starts])
+
+    # create dict of all variables, except 'trajectory' as that is already created above
+    var = {}
+    for v in pfile_in.variables:
+        if str(v) != 'trajectory':
+            varin = pfile_in.variables[v]
+            var[v] = pfile_out.createVariable(v, "f4", coords, fill_value=np.nan)
+            # copy all attributes, except Fill_Value which is set automatically
+            var[v].setncatts({k: varin.getncattr(k) for k in varin.ncattrs() if k != '_FillValue'})
+
+    pbar = ProgressBar()
+    for i in pbar(range(trajs.nid)):
+        ii = np.sort(trajs.indices[trajs.starts[i]:trajs.ends[i]])
+        for v in var:
+            var[v][i, 0:trajs.lengths[i]] = pfile_in.variables[v][ii]
+
+    pfile_out.sync()
+
+
+if __name__ == "__main__":
+    p = ArgumentParser(description="""Converting Indexed Parcels output to Array format""")
+    p.add_argument('-i', '--file_in', type=str,
+                   help='Name of input file in indexed form')
+    p.add_argument('-o', '--file_out', type=str,
+                   help='Name of output file in array form')
+    args = p.parse_args()
+    convert_IndexedOutputToArray(file_in=args.file_in, file_out=args.file_out)

parcels/scripts/plotParticles.py

@@ -11,7 +11,7 @@
 
 
 def plotTrajectoriesFile(filename, mode='2d', tracerfile=None, tracerfield='P',
-                         tracerlon='x', tracerlat='y', recordedvar=None):
+                         tracerlon='x', tracerlat='y', recordedvar=None, show_plt=True):
     """Quick and simple plotting of Parcels trajectories
 
     :param filename: Name of Parcels-generated NetCDF file with particle positions
@@ -24,6 +24,7 @@ def plotTrajectoriesFile(filename, mode='2d', tracerfile=None, tracerfield='P',
     :param tracerlat: Name of latitude dimension of variable to show as background
     :param recordedvar: Name of variable used to color particles in scatter-plot.
                 Only works in 'movie2d' or 'movie2d_notebook' mode.
+    :param show_plt: Boolean whether plot should directly be show (for py.test)
     """
 
     if plt is None:
@@ -34,10 +35,10 @@ def plotTrajectoriesFile(filename, mode='2d', tracerfile=None, tracerfield='P',
     lon = pfile.variables['lon']
     lat = pfile.variables['lat']
     z = pfile.variables['z']
+    time = pfile.variables['time'][:]
     if len(lon.shape) == 1:
         type = 'indexed'
         id = pfile.variables['trajectory'][:]
-        time = pfile.variables['time'][:]
     else:
         type = 'array'
 
@@ -59,7 +60,7 @@ def plotTrajectoriesFile(filename, mode='2d', tracerfile=None, tracerfield='P',
             for p in range(len(lon)):
                 ax.plot(lon[p, :], lat[p, :], z[p, :], '.-')
         elif type == 'indexed':
-            for t in range(max(id)+1):
+            for t in np.unique(id):
                 ax.plot(lon[id == t], lat[id == t],
                         z[id == t], '.-')
         ax.set_xlabel('Longitude')
@@ -69,11 +70,18 @@ def plotTrajectoriesFile(filename, mode='2d', tracerfile=None, tracerfield='P',
         if type == 'array':
             plt.plot(np.transpose(lon), np.transpose(lat), '.-')
         elif type == 'indexed':
-            for t in range(max(id)+1):
+            for t in np.unique(id):
                 plt.plot(lon[id == t], lat[id == t], '.-')
         plt.xlabel('Longitude')
         plt.ylabel('Latitude')
     elif mode == 'movie2d' or 'movie2d_notebook':
+        if type == 'array' and any(time[:, 0] != time[0, 0]):
+            # since particles don't start at the same time, treat as indexed
+            type = 'indexed'
+            id = pfile.variables['trajectory'][:].flatten()
+            lon = lon[:].flatten()
+            lat = lat[:].flatten()
+            time = time.flatten()
 
         fig = plt.figure()
         ax = plt.axes(xlim=(np.amin(lon), np.amax(lon)), ylim=(np.amin(lat), np.amax(lat)))
@@ -84,7 +92,7 @@ def plotTrajectoriesFile(filename, mode='2d', tracerfile=None, tracerfield='P',
             mintime = min(time)
             scat = ax.scatter(lon[time == mintime], lat[time == mintime],
                               s=60, cmap=plt.get_cmap('autumn'))
-            frames = np.unique(time)
+            frames = np.unique(time[~np.isnan(time)])
 
         def animate(t):
             if type == 'array':
@@ -102,7 +110,9 @@ def animate(t):
         plt.close()
         return anim
     else:
-        plt.show()
+        if show_plt:
+            plt.show()
+        return plt
 
 
 if __name__ == "__main__":
@@ -125,4 +135,5 @@ def animate(t):
 
     plotTrajectoriesFile(args.particlefile, mode=args.mode, tracerfile=args.tracerfile,
                          tracerfield=args.tracerfilefield, tracerlon=args.tracerfilelon,
-                         tracerlat=args.tracerfilelat, recordedvar=args.recordedvar)
+                         tracerlat=args.tracerfilelat, recordedvar=args.recordedvar,
+                         show_plt=True)

tests/test_scripts.py

@@ -0,0 +1,43 @@
+from parcels import (FieldSet, ParticleSet, JITParticle, AdvectionRK4,
+                     plotTrajectoriesFile, convert_IndexedOutputToArray)
+from datetime import timedelta as delta
+import numpy as np
+import pytest
+from os import path, pardir
+
+
+def create_outputfiles(dir):
+    datafile = path.join(path.dirname(__file__), pardir, 'examples',
+                         'Peninsula_data', 'peninsula')
+
+    fieldset = FieldSet.from_nemo(datafile, allow_time_extrapolation=True)
+    pset = ParticleSet(fieldset=fieldset, lon=[], lat=[], pclass=JITParticle)
+    npart = 10
+    delaytime = delta(hours=1)
+    endtime = delta(hours=24)
+    x = 3. * (1. / 1.852 / 60)
+    y = (fieldset.U.lat[0] + x, fieldset.U.lat[-1] - x)
+    lat = np.linspace(y[0], y[1], npart, dtype=np.float32)
+
+    fp_index = dir.join("DelayParticle")
+    output_file = pset.ParticleFile(name=fp_index, type="indexed")
+
+    for t in range(npart):
+        pset.add(JITParticle(lon=x, lat=lat[t], fieldset=fieldset))
+        pset.execute(AdvectionRK4, runtime=delaytime, dt=delta(minutes=5),
+                     interval=delaytime, starttime=delaytime*t, output_file=output_file)
+
+    pset.execute(AdvectionRK4, runtime=endtime-npart*delaytime, starttime=delaytime*npart,
+                 dt=delta(minutes=5), interval=delta(hours=1), output_file=output_file)
+
+    fp_array = dir.join("DelayParticle_array")
+    convert_IndexedOutputToArray(fp_index+'.nc', fp_array+'.nc')
+    return fp_index, fp_array
+
+
+@pytest.mark.parametrize('mode', ['2d', '3d', 'movie2d'])
+@pytest.mark.parametrize('fp_type', ['index', 'array'])
+def test_plotting(mode, tmpdir, fp_type):
+    fp_index, fp_array = create_outputfiles(tmpdir)
+    fp = fp_array if fp_type == 'array' else fp_index
+    plotTrajectoriesFile(fp+'.nc', mode=mode, show_plt=False)