Dev

examples/create_test_dataset.py

@@ -14,10 +14,12 @@
     '16Nov2015_NorKyst_z_surface/norkyst800_subset_16Nov2015.nc')
 o.add_reader([reader_norkyst, reader_arome])
 
-# Seeding some particles
-o.seed_elements(lon=4.4, lat=60.1, radius=1000, number=1000,
-                time=reader_arome.start_time)
+# Seeding some particles, continuous spill
+o.seed_elements(lon=4.2, lat=60.1, radius=1000, number=1000,
+                time=[reader_arome.start_time, reader_arome.end_time])
 
 # Running model
 o.run(end_time=reader_norkyst.end_time,
-      export_variables=['mass_oil'], outfile='openoil.nc')
+      export_variables=['viscosity', 'z'], outfile='openoil.nc')
+
+o.plot(fast=True, linecolor='viscosity')

examples/example_opendrift.py

@@ -13,21 +13,17 @@
 
 #%%
 # Importing a trajectory dataset from a simulation with OpenDrift.
-# decode_coords is needed so that lon and lat are not interpreted as coordinate variables.
-with lzma.open('openoil.nc.xz') as oil:
-    ds = xr.open_dataset(oil, decode_coords=False)
-    ds.load()
-    # Requirement that status>=0 is needed since non-valid points are not masked in OpenDrift output
-    ds = ds.where(ds.status>=0)  # only active particles
+ds = xr.open_dataset('openoil.nc')
 
 #%%
 # Displaying some basic information about this dataset
 print(ds.traj)
 
 #%%
-# Making a basic plot of trajectories
-ds.traj.plot()
-plt.title('Basic trajectory plot')
+# Plotting trajectories, colored by oil viscosity
+mappable = ds.traj.plot(color=ds.viscosity, alpha=1)
+plt.title('Oil trajectories')
+plt.colorbar(mappable, orientation='horizontal', label=f'Oil viscosity  [{ds.viscosity.units}]')
 plt.show()
 
 #%%

examples/example_parcels.py

@@ -19,7 +19,25 @@
 #%%
 # Basic plot
 ds.traj.plot(land='mask', margin=1)
+
 # TODO: we must allow no time dimension for the below to work
 #ds.mean('trajectory', skipna=True).traj.plot(color='r', label='Mean trajectory')
+# In the meantime, we regrid to a regular 1D dataset to allow plotting a mean trajectory
+ds = ds.traj.gridtime('1h')
+ds.mean('trajectory').traj.plot(color='r', label='Mean trajectory')
+plt.legend()
+plt.show()
 
+#%%
+# Calculating skillscore
+# Defining the first trajectory to be the "true"
+ds_true = ds.isel(trajectory=0)
+skillscore = ds.traj.skill(expected=ds_true, method='liu-weissberg', tolerance_threshold=1)
+
+#%%
+# Plotting trajectories, colored by their skillscore, compared to the "true" trajectory (black)
+mappable = ds.traj.plot(land='mask', color=skillscore)
+ds_true.traj.plot(color='k', linewidth=3, label='"True" trajectory')
+plt.colorbar(mappable=mappable, orientation='horizontal', label='Skillscore per trajectory')
+plt.legend()
 plt.show()

tests/fixtures.py

@@ -9,11 +9,9 @@ def plot(pytestconfig):
 
 @pytest.fixture
 def opendrift_sim():
-    oil = Path(__file__).parent.parent / 'examples' / 'openoil.nc.xz'
-    with lzma.open(oil) as oil:
-        ds = xr.open_dataset(oil)
-        ds.load()
-        return ds
+    oil = Path(__file__).parent.parent / 'examples' / 'openoil.nc'
+    ds = xr.open_dataset(oil)
+    return ds
 
 @pytest.fixture
 def barents():

trajan/plot/__init__.py

@@ -168,6 +168,12 @@ def lines(self, *args, **kwargs):
             y = self.ds.traj.tlat.values.T
 
         if hasattr(kwargs['color'], 'shape'):
+            if isinstance(kwargs['color'], xr.DataArray):
+                kwargs['color'] = kwargs['color'].values
+            # TODO: it should be possible to make this much faster, especially for fixed color per trajectory
+            if kwargs['color'].ndim == 1:  # Fixed color per line
+                assert len(kwargs['color']) == x.shape[1]
+                kwargs['color'] = kwargs['color'][:, np.newaxis] * np.ones(x.shape).T
             from matplotlib.collections import LineCollection
             c = kwargs.pop('color').T
             if hasattr(c, 'values'):