Skip to content

Commit

Permalink
Merge pull request #14 from ecmwf/develop
Browse files Browse the repository at this point in the history
Adds initial implementation of visualising unstructured grids
  • Loading branch information
JamesVarndell authored Sep 24, 2024
2 parents b37cece + 0ec5a35 commit d88da6f
Show file tree
Hide file tree
Showing 5 changed files with 337 additions and 3 deletions.
3 changes: 2 additions & 1 deletion docs/examples/gallery/gallery.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,8 @@
"* [Model orography](gridded-data/model-orography.ipynb)\n",
"* [Temperature and pressure](gridded-data/temperature-and-pressure.ipynb)\n",
"* [Time zones](gridded-data/time-zones.ipynb)\n",
"* [Weather forecast steps](gridded-data/weather-forecast-steps.ipynb)"
"* [Weather forecast steps](gridded-data/weather-forecast-steps.ipynb)\n",
"* [Unstructured grids](gridded-data/unstructured-grids.ipynb)"
]
},
{
Expand Down
220 changes: 220 additions & 0 deletions docs/examples/gallery/gridded-data/unstructured-grids.ipynb

Large diffs are not rendered by default.

23 changes: 22 additions & 1 deletion src/earthkit/plots/components/subplots.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import warnings
from itertools import cycle

import earthkit.data
Expand All @@ -21,6 +22,7 @@

from earthkit.plots import identifiers
from earthkit.plots.components.layers import Layer
from earthkit.plots.geo import grids
from earthkit.plots.metadata.formatters import (
LayerFormatter,
SourceFormatter,
Expand Down Expand Up @@ -338,7 +340,8 @@ def _extract_plottables(
style = auto.guess_style(
source, units=units or source.units, **override_kwargs
)
if (data is None and z is None) or (z is not None and not z):

if data is None and z is None:
z_values = None
else:
z_values = style.convert_units(source.z_values, source.units)
Expand Down Expand Up @@ -371,6 +374,19 @@ def _extract_plottables(
x_values = source.x_values
y_values = source.y_values

if method_name in (
"contour",
"contourf",
"pcolormesh",
) and not grids.is_structured(x_values, y_values):
x_values, y_values, z_values = grids.interpolate_unstructured(
x_values,
y_values,
z_values,
method=kwargs.pop("interpolation_method", "linear"),
)
extract_domain = False

if every is not None:
x_values = x_values[::every]
y_values = y_values[::every]
Expand All @@ -384,6 +400,11 @@ def _extract_plottables(
z_values,
source_crs=source.crs,
)
if "interpolation_method" in kwargs:
kwargs.pop("interpolation_method")
warnings.warn(
"The 'interpolation_method' argument is only valid for unstructured data."
)
mappable = getattr(style, method_name)(
self.ax, x_values, y_values, z_values, **kwargs
)
Expand Down
92 changes: 92 additions & 0 deletions src/earthkit/plots/geo/grids.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
# Copyright 2024, European Centre for Medium Range Weather Forecasts.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from scipy.interpolate import griddata


def is_structured(lat, lon, tol=1e-5):
"""
Determines whether the latitude and longitude points form a structured grid.
Parameters:
- lat: A 1D or 2D array of latitude points.
- lon: A 1D or 2D array of longitude points.
- tol: Tolerance for floating-point comparison (default 1e-5).
Returns:
- True if the data is structured (grid), False if it's unstructured.
"""

lat = np.asarray(lat)
lon = np.asarray(lon)

# Check if there are consistent spacing in latitudes and longitudes
unique_lat = np.unique(lat)
unique_lon = np.unique(lon)

# Structured grid condition: the number of unique lat/lon values should multiply to the number of total points
if len(unique_lat) * len(unique_lon) == len(lat) * len(lon):
# Now check if the spacing is consistent
lat_diff = np.diff(unique_lat)
lon_diff = np.diff(unique_lon)

# Check if lat/lon differences are consistent
lat_spacing_consistent = np.all(np.abs(lat_diff - lat_diff[0]) < tol)
lon_spacing_consistent = np.all(np.abs(lon_diff - lon_diff[0]) < tol)

return lat_spacing_consistent and lon_spacing_consistent

# If the product of unique lat/lon values doesn't match total points, it's unstructured
return False


def interpolate_unstructured(x, y, z, resolution=1000, method="linear"):
"""
Interpolates unstructured data to a structured grid, handling NaNs in z-values
and preventing interpolation across large gaps.
Parameters:
- x: 1D array of x-coordinates.
- y: 1D array of y-coordinates.
- z: 1D array of z values.
- resolution: The number of points along each axis for the structured grid.
- method: Interpolation method ('linear', 'nearest', 'cubic').
- gap_threshold: The distance threshold beyond which interpolation is not performed (set to NaN).
Returns:
- grid_x: 2D grid of x-coordinates.
- grid_y: 2D grid of y-coordinates.
- grid_z: 2D grid of interpolated z-values, with NaNs in large gap regions.
"""
# Filter out NaN values from z and corresponding x, y
mask = ~np.isnan(z)
x_filtered = x[mask]
y_filtered = y[mask]
z_filtered = z[mask]

# Create a structured grid
grid_x, grid_y = np.mgrid[
x.min() : x.max() : resolution * 1j, y.min() : y.max() : resolution * 1j
]

# Interpolate the filtered data onto the structured grid
grid_z = griddata(
np.column_stack((x_filtered, y_filtered)),
z_filtered,
(grid_x, grid_y),
method=method,
)

return grid_x, grid_y, grid_z
2 changes: 1 addition & 1 deletion src/earthkit/plots/styles/legends.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# limitations under the License.


DEFAULT_LEGEND_LABEL = "{variable_name} ({units})"
DEFAULT_LEGEND_LABEL = "" # {variable_name} ({units})"

_DISJOINT_LEGEND_LOCATIONS = {
"bottom": {
Expand Down

0 comments on commit d88da6f

Please sign in to comment.