Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

start of plotting "help" module #258

Merged
merged 23 commits into from
Nov 27, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ jobs:
fail-fast: false
matrix:
os: ["ubuntu-latest", "macos-latest", "windows-latest"]
python-version: ["3.9", "3.10", "3.11"]
python-version: ["3.9", "3.10", "3.11", "3.12"]
steps:
- uses: actions/checkout@v3
- uses: mamba-org/setup-micromamba@v1
Expand All @@ -29,7 +29,8 @@ jobs:
- name: Testing
shell: bash -l {0}
run: |
pip install coverage
pip install coverage
pip install matplotlib-base cartopy
coverage run -m unittest discover -s tests -p "*.py"
coverage xml
- name: Upload coverage reports to Codecov
Expand Down
16 changes: 16 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,13 @@ In your virtual environment, type:
pip install clouddrift
```

To install optional dependencies needed by the `clouddrift.plotting` module,
type:

```
pip install matplotlib-base cartopy
```

##### Conda:

First add `conda-forge` to your channels in your Conda configuration (`~/.condarc`):
Expand All @@ -84,6 +91,13 @@ then install CloudDrift:
conda install clouddrift
```

To install optional dependencies needed by the `clouddrift.plotting` module,
type:

```
conda install matplotlib-base cartopy
```

#### Development branch:

If you need the latest development version, you can install it directly from this GitHub repository.
Expand Down Expand Up @@ -120,13 +134,15 @@ With pip:
python3 -m venv .venv
source .venv/bin/activate
pip install .
pip install matplotlib-base cartopy
```

With Conda:

```
conda env create -f environment.yml
conda activate clouddrift
conda install matplotlib-base cartopy
```

Then, run the tests like this:
Expand Down
1 change: 1 addition & 0 deletions clouddrift/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import clouddrift.adapters
import clouddrift.datasets
import clouddrift.kinematics
import clouddrift.plotting
import clouddrift.ragged
import clouddrift.signal
import clouddrift.sphere
Expand Down
215 changes: 215 additions & 0 deletions clouddrift/plotting.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,215 @@
"""
This module provides a function to easily and efficiently plot trajectories stored in a ragged array.
"""

from clouddrift.ragged import segment, rowsize_to_index
import numpy as np
import pandas as pd
from typing import Optional, Union
import xarray as xr
import pandas as pd
from typing import Optional, Union
from clouddrift.ragged import segment, rowsize_to_index


def plot_ragged(
ax,
longitude: Union[list, np.ndarray, pd.Series, xr.DataArray],
latitude: Union[list, np.ndarray, pd.Series, xr.DataArray],
rowsize: Union[list, np.ndarray, pd.Series, xr.DataArray],
*args,
colors: Optional[Union[list, np.ndarray, pd.Series, xr.DataArray]] = None,
tolerance: Optional[Union[float, int]] = 180,
**kwargs,
):
"""Plot trajectories from a ragged array dataset on a Matplotlib Axes
or a Cartopy GeoAxes object ``ax``.

This function wraps Matplotlib's ``plot`` function (``plt.plot``) and
``LineCollection`` (``matplotlib.collections``) to efficiently plot
trajectories from a ragged array dataset.

Parameters
----------
ax: matplotlib.axes.Axes or cartopy.mpl.geoaxes.GeoAxes
Axis to plot on.
longitude : array-like
Longitude sequence. Unidimensional array input.
latitude : array-like
Latitude sequence. Unidimensional array input.
rowsize : list
List of integers specifying the number of data points in each row.
*args : tuple
Additional arguments to pass to ``ax.plot``.
colors : array-like
Colors to use for plotting. If colors is the same shape as longitude and latitude,
the trajectories are splitted into segments and each segment is colored according
to the corresponding color value. If colors is the same shape as rowsize, the
trajectories are uniformly colored according to the corresponding color value.
tolerance : float
Longitude tolerance gap between data points (in degrees) for segmenting trajectories.
For periodic domains, the tolerance parameter should be set to the maximum allowed gap
between data points. Defaults to 180.
**kwargs : dict
Additional keyword arguments to pass to ``ax.plot``.

Returns
-------
list of matplotlib.lines.Line2D or matplotlib.collections.LineCollection
The plotted lines or line collection. Can be used to set a colorbar
after plotting or extract information from the lines.

Examples
--------

Plot the first 100 trajectories from the gdp1h dataset, assigning
a different color to each trajectory:

>>> from clouddrift import datasets
>>> import matplotlib.pyplot as plt
>>> ds = datasets.gdp1h()
>>> ds = subset(ds, {"ID": ds.ID[:100].values}).load()
>>> fig = plt.figure()
>>> ax = fig.add_subplot(1, 1, 1)

>>> plot_ragged(
>>> ax,
>>> ds.lon,
>>> ds.lat,
>>> ds.rowsize,
>>> colors=np.arange(len(ds.rowsize))
>>> )

To plot the same trajectories, but assigning a different color to each
observation and specifying a colormap:

>>> fig = plt.figure()
>>> ax = fig.add_subplot(1, 1, 1)
>>> time = [v.astype(np.int64) / 86400 / 1e9 for v in ds.time.values]
>>> lc = plot_ragged(
>>> ax,
>>> ds.lon,
>>> ds.lat,
>>> ds.rowsize,
>>> colors=np.floor(time),
philippemiron marked this conversation as resolved.
Show resolved Hide resolved
>>> cmap="inferno"
>>> )
>>> fig.colorbar(lc[0])
>>> ax.set_xlim([-180, 180])
>>> ax.set_ylim([-90, 90])

Finally, to plot the same trajectories, but using a cartopy
projection:

>>> import cartopy.crs as ccrs
>>> fig = plt.figure()
>>> ax = fig.add_subplot(1, 1, 1, projection=ccrs.Mollweide())
philippemiron marked this conversation as resolved.
Show resolved Hide resolved
>>> time = [v.astype(np.int64) / 86400 / 1e9 for v in ds.time.values]
>>> lc = plot_ragged(
>>> ax,
>>> ds.lon,
>>> ds.lat,
>>> ds.rowsize,
>>> colors=np.arange(len(ds.rowsize)),
>>> transform=ccrs.PlateCarree(),
>>> cmap=cmocean.cm.ice,
>>> )

Raises
------
ValueError
If longitude and latitude arrays do not have the same shape.
If colors do not have the same shape as longitude and latitude arrays or rowsize.
If ax is not a matplotlib Axes or GeoAxes object.
If ax is a GeoAxes object and the transform keyword argument is not provided.

ImportError
If matplotlib is not installed.
If the axis is a GeoAxes object and cartopy is not installed.
"""

# optional dependency
try:
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.collections import LineCollection
from matplotlib import cm
except ImportError:
raise ImportError("missing optional dependency 'matplotlib'")

Check warning on line 138 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L132-L138

Added lines #L132 - L138 were not covered by tests

if hasattr(ax, "coastlines"): # check if GeoAxes without cartopy
try:
from cartopy.mpl.geoaxes import GeoAxes

Check warning on line 142 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L140-L142

Added lines #L140 - L142 were not covered by tests

if isinstance(ax, GeoAxes) and not kwargs.get("transform"):
raise ValueError(

Check warning on line 145 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L144-L145

Added lines #L144 - L145 were not covered by tests
"For GeoAxes, the transform keyword argument must be provided."
)
except ImportError:
raise ImportError("missing optional dependency 'cartopy'")
elif not isinstance(ax, plt.Axes):
raise ValueError("ax must be either: plt.Axes or GeoAxes.")

Check warning on line 151 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L148-L151

Added lines #L148 - L151 were not covered by tests

if np.sum(rowsize) != len(longitude):
raise ValueError("The sum of rowsize must equal the length of lon and lat.")

Check warning on line 154 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L153-L154

Added lines #L153 - L154 were not covered by tests

if len(longitude) != len(latitude):
raise ValueError("lon and lat must have the same length.")

Check warning on line 157 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L156-L157

Added lines #L156 - L157 were not covered by tests

if colors is None:
colors = np.arange(len(rowsize))
elif colors is not None and (len(colors) not in [len(longitude), len(rowsize)]):
raise ValueError("shape colors must match the shape of lon/lat or rowsize.")

Check warning on line 162 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L159-L162

Added lines #L159 - L162 were not covered by tests

# define a colormap
cmap = kwargs.pop("cmap", cm.viridis)

Check warning on line 165 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L165

Added line #L165 was not covered by tests

# define a normalization obtain uniform colors
# for the sequence of lines or LineCollection
norm = kwargs.pop(

Check warning on line 169 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L169

Added line #L169 was not covered by tests
"norm", mcolors.Normalize(vmin=np.nanmin(colors), vmax=np.nanmax(colors))
)

mpl_plot = True if colors is None or len(colors) == len(rowsize) else False
traj_idx = rowsize_to_index(rowsize)

Check warning on line 174 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L173-L174

Added lines #L173 - L174 were not covered by tests

lines = []
for i in range(len(rowsize)):
lon_i, lat_i = (

Check warning on line 178 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L176-L178

Added lines #L176 - L178 were not covered by tests
longitude[traj_idx[i] : traj_idx[i + 1]],
latitude[traj_idx[i] : traj_idx[i + 1]],
)

start = 0
for length in segment(lon_i, tolerance, rowsize=segment(lon_i, -tolerance)):
end = start + length

Check warning on line 185 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L183-L185

Added lines #L183 - L185 were not covered by tests

if mpl_plot:
line = ax.plot(

Check warning on line 188 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L187-L188

Added lines #L187 - L188 were not covered by tests
lon_i[start:end],
lat_i[start:end],
c=cmap(norm(colors[i])) if colors is not None else None,
*args,
**kwargs,
)
else:
colors_i = colors[traj_idx[i] : traj_idx[i + 1]]
segments = np.column_stack(

Check warning on line 197 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L196-L197

Added lines #L196 - L197 were not covered by tests
[
lon_i[start : end - 1],
lat_i[start : end - 1],
lon_i[start + 1 : end],
lat_i[start + 1 : end],
]
).reshape(-1, 2, 2)
line = LineCollection(segments, cmap=cmap, norm=norm, *args, **kwargs)
line.set_array(

Check warning on line 206 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L205-L206

Added lines #L205 - L206 were not covered by tests
# color of a segment is the average of its two data points
np.convolve(colors_i[start:end], [0.5, 0.5], mode="valid")
)
ax.add_collection(line)

Check warning on line 210 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L210

Added line #L210 was not covered by tests

start = end
lines.append(line)

Check warning on line 213 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L212-L213

Added lines #L212 - L213 were not covered by tests

return lines

Check warning on line 215 in clouddrift/plotting.py

View check run for this annotation

Codecov / codecov/patch

clouddrift/plotting.py#L215

Added line #L215 was not covered by tests
1 change: 1 addition & 0 deletions docs/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ Auto-generated summary of CloudDrift's API. For more details and examples, refer
adapters.subsurface_floats
datasets
kinematics
plotting
ragged
raggedarray
signal
Expand Down
16 changes: 15 additions & 1 deletion docs/install.rst
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,13 @@ In your virtual environment, type:

pip install clouddrift

To install optional dependencies needed by the ``clouddrift.plotting`` module,
type:

.. code-block:: text

pip install matplotlib-base cartopy

Conda
-----

Expand All @@ -31,6 +38,13 @@ then install CloudDrift:

conda install clouddrift

To install optional dependencies needed by the ``clouddrift.plotting`` module,
type:

.. code-block:: text

conda install matplotlib-base cartopy

Developers
----------

Expand Down Expand Up @@ -61,4 +75,4 @@ Then, run the tests like this:

python -m unittest tests/*.py

A quick how-to guide is provided on the `Usage <https://cloud-drift.github.io/clouddrift/usage.html>`_ page.
A quick how-to guide is provided on the `Usage <https://cloud-drift.github.io/clouddrift/usage.html>`_ page.
Loading