Feature/density tracks

CHANGELOG.md

@@ -12,7 +12,9 @@ Code freeze date: YYYY-MM-DD
 
 ### Added
 
-- `climada.hazard.tc_tracks.TCTracks.from_FAST` function, add Australia basin (AU) [#993](https://github.com/CLIMADA-project/climada_python/pull/993)
+-`climada.hazard.tc_tracks.compute_track_density` function, `climada.util.coordinates.compute_grid_cell_area`
+function, `climada.hazard.plot.plot_track_density` function  [#1003](https://github.com/CLIMADA-project/climada_python/pull/1003)
+-`climada.hazard.tc_tracks.TCTracks.from_FAST` function, add Australia basin (AU) [#993](https://github.com/CLIMADA-project/climada_python/pull/993)
 - Add `osm-flex` package to CLIMADA core [#981](https://github.com/CLIMADA-project/climada_python/pull/981)
 - `doc.tutorial.climada_entity_Exposures_osm.ipynb` tutorial explaining how to use `osm-flex`with CLIMADA
 - `climada.util.coordinates.bounding_box_global` function [#980](https://github.com/CLIMADA-project/climada_python/pull/980)

climada/hazard/plot.py

@@ -21,6 +21,8 @@
 
 import logging
 
+import cartopy.crs as ccrs
+import cartopy.feature as cfeature
 import matplotlib.pyplot as plt
 import numpy as np
 from deprecation import deprecated
@@ -158,6 +160,119 @@ def plot_intensity(
 
         raise ValueError("Provide one event id or one centroid id.")
 
+    @staticmethod
+    def plot_track_density(
+        hist: np.ndarray,
+        axis=None,
+        projection=ccrs.Mollweide(),
+        add_features: dict = None,
+        title: str = None,
+        figsize=(12, 6),
+        cbar_kwargs: dict = {
+            "orientation": "horizontal",
+            "pad": 0.05,
+            "shrink": 0.8,
+            "label": "Track Density [n° tracks / km²]",
+        },
+        **kwargs,
+    ):
+        """
+        Plot the track density of tropical cyclone tracks on a customizable world map.
+
+        Parameters:
+        ----------
+        hist: np.ndarray
+            2D histogram of track density.
+        axis: GeoAxes, optional
+            Existing Cartopy axis.
+        projection: cartopy.crs, optional
+            Projection for the map.
+        add_features: dict
+            Dictionary of map features to add. Keys can be 'land', 'coastline', 'borders', and
+            'lakes'. Values are Booleans indicating whether to include each feature.
+        title: str
+            Title of the plot.
+        figsize: tuple
+            Figure size when creating a new figure.
+        cbar_kwargs: dict
+            dictionary containing keyword arguments passed to cbar
+        kwargs:
+            Additional keyword arguments passed to `ax.contourf`.
+
+        Returns:
+        -------
+        axis: GeoAxes
+            The plot axis.
+
+
+        Example:
+        --------
+        >>> axis = plot_track_density(
+        ...     hist=hist,
+        ...     cmap='Spectral_r',
+        ...     cbar_kwargs={'shrink': 0.8,
+                            'label': 'Cyclone Density [n° tracks / km²]',
+                            'pad': 0.1},
+        ...     add_features={
+        ...         'land': True,
+        ...         'coastline': True,
+        ...         'borders': False,
+        ...         'lakes': False
+        ...     },
+        ...     title='My Tropical Cyclone Track Density Map',
+        ...     figsize=(10, 5),
+        ...     levels=20
+        ... )
+
+        """
+
+        # Default features
+        default_features = {
+            "land": True,
+            "coastline": True,
+            "borders": False,
+            "lakes": False,
+        }
+        add_features = add_features or default_features
+
+        # Sample data
+        lon = np.linspace(-180, 180, hist.shape[1])
+        lat = np.linspace(-90, 90, hist.shape[0])
+
+        # Create figure and axis if not provided
+        if axis is None:
+            _, axis = plt.subplots(
+                figsize=figsize, subplot_kw={"projection": projection}
+            )
+
+        # Add requested features
+        if add_features.get("land", False):
+            land = cfeature.NaturalEarthFeature(
+                category="physical",
+                name="land",
+                scale="50m",
+                facecolor="lightgrey",
+                alpha=0.6,
+            )
+            axis.add_feature(land)
+        if add_features.get("coastline", False):
+            axis.add_feature(cfeature.COASTLINE, linewidth=0.5)
+        if add_features.get("borders", False):
+            axis.add_feature(cfeature.BORDERS, linestyle=":")
+        if add_features.get("lakes", False):
+            axis.add_feature(cfeature.LAKES, alpha=0.4, edgecolor="black")
+
+        # Plot density with contourf
+        contourf = axis.contourf(lon, lat, hist, transform=ccrs.PlateCarree(), **kwargs)
+
+        # Add colorbar
+        plt.colorbar(contourf, ax=axis, **cbar_kwargs)
+        # Title setup
+        if title:
+            axis.set_title(title, fontsize=16)
+
+        return axis
+
     def plot_fraction(self, event=None, centr=None, smooth=True, axis=None, **kwargs):
         """Plot fraction values for a selected event or centroid.
 

climada/hazard/tc_tracks.py

@@ -38,7 +38,6 @@
 import matplotlib.cm as cm_mp
 import matplotlib.pyplot as plt
 import netCDF4 as nc
-import numba
 import numpy as np
 import pandas as pd
 import pathos
@@ -51,6 +50,7 @@
 from matplotlib.lines import Line2D
 from shapely.geometry import LineString, MultiLineString, Point
 from sklearn.metrics import DistanceMetric
+from tqdm import tqdm
 
 import climada.hazard.tc_tracks_synth
 import climada.util.coordinates as u_coord
@@ -2980,3 +2980,178 @@ def _zlib_from_dataarray(data_var: xr.DataArray) -> bool:
     if np.issubdtype(data_var.dtype, float) or np.issubdtype(data_var.dtype, int):
         return True
     return False
+
+
+def compute_track_density(
+    tc_track: TCTracks,
+    res: int = 5,
+    genesis: bool = False,
+    norm: str = None,
+    filter_tracks: bool = True,
+    wind_min: float = None,
+    wind_max: float = None,
+) -> tuple[np.ndarray, tuple]:
+    """Compute absolute and normalized tropical cyclone track density. Before using this function,
+    apply the same temporal resolution to all tracks by calling :py:meth:`equal_timestep` on the
+    TCTrack object. Due to the computational cost of the this function, it is not recommended to
+    use a grid resolution higher tha 0.1°. First, this function creates 2D bins of the specified
+    resolution (e.g. 1° x 1°). Second, since tracks are not lines but a series of points, it counts
+    the number of points per bin. Lastly, it returns the absolute or normalized count per bin.
+    To plot the output of this function use :py:meth:`plot_track_density`.
+
+    Parameters:
+    ----------
+    tc_track: TCT track object
+        track object containing a list of all tracks
+    res: int (optional), default: 5°
+        resolution in degrees of the grid bins in which the density will be computed
+    genesis: bool, Default = False
+        If true the function computes the track density of only the genesis location of tracks
+    norm: bool (optional), default = False
+        If False it returns the number of samples in each bin. If True, returns the
+        specified normalization function at each bin computed as count_bin / grid_area.
+    filter_tracks: bool (optional) default: True
+        If True the track density is computed as the number of different tracks crossing a grid
+        cell. If False, the track density takes into account how long the track stayed in each
+        grid cell. Hence slower tracks increase the density if the parameter is set to False.
+    wind_min: float (optional), default: None
+        Minimum wind speed above which to select tracks.
+    wind_max: float (optional), default: None
+        Maximal wind speed below which to select tracks.
+    Returns:
+    -------
+    hist_count: np.ndarray
+        2D matrix containing the the absolute count per gridd cell of track point or the normalized
+        number of track points, depending on the norm parameter.
+
+    Example:
+    --------
+    >>> tc_tracks = TCTrack.from_ibtracs_netcdf("path_to_file")
+    >>> tc_tracks.equal_timestep(time_steph_h = 1)
+    >>> hist_count, *_ = compute_track_density(res = 1)
+    >>> plot_track_density(hist_count)
+
+    """
+
+    limit_ratio: float = 1.12 * 1.1  # record tc speed 112km/h -> 1.12°/h + 10% margin
+    time_value: float = (
+        tc_track.data[0].time_step[0].values.astype(float)
+    )  # Type hint for jenkins
+
+    if time_value > (res / limit_ratio):
+        warnings.warn(
+            "The time step is too big for the current resolution. For the desired resolution, \n"
+            f"apply a time step of {res/limit_ratio}h."
+        )
+    elif res < 0.1:
+        warnings.warn(
+            "The resolution is too high. The computation might take several minutes \n"
+            "to hours. Consider using a resolution below 0.1°."
+        )
+
+    # define grid resolution and bounds for density computation
+    lat_bins: np.ndarray = np.linspace(-90, 90, int(180 / res))
+    lon_bins: np.ndarray = np.linspace(-180, 180, int(360 / res))
+    # compute 2D density
+    if genesis:
+        hist_count = compute_genesis_density(
+            tc_track=tc_track, lat_bins=lat_bins, lon_bins=lon_bins
+        )
+    else:
+        hist_count = np.zeros((len(lat_bins) - 1, len(lon_bins) - 1))
+        for track in tqdm(tc_track.data, desc="Processing Tracks"):
+
+            # select according to wind speed
+            wind_speed = track.max_sustained_wind.values
+            if wind_min and wind_max:
+                index = np.where((wind_speed >= wind_min) & (wind_speed <= wind_max))[0]
+            elif wind_min and not wind_max:
+                index = np.where(wind_speed >= wind_min)[0]
+            elif wind_max and not wind_min:
+                index = np.where(wind_speed <= wind_max)[0]
+            else:
+                index = slice(None)  # select all the track
+
+            # compute 2D density
+            hist_new, _, _ = np.histogram2d(
+                track.lat.values[index],
+                track.lon.values[index],
+                bins=[lat_bins, lon_bins],
+                density=False,
+            )
+            if filter_tracks:
+                hist_new[hist_new > 1] = 1
+
+            hist_count += hist_new
+
+    if norm:
+        hist_count = normalize_hist(res=res, hist_count=hist_count, norm=norm)
+
+    return hist_count, lat_bins, lon_bins
+
+
+def compute_genesis_density(
+    tc_track: TCTracks, lat_bins: np.ndarray, lon_bins: np.ndarray
+) -> np.ndarray:
+    """Compute the density of track genesis locations. This function works under the hood
+    of :py:meth:`compute_track_density`. If it is called with the parameter genesis = True,
+    the function return the number of genesis points per grid cell.
+
+    Parameters:
+    -----------
+
+    tc_track: TCT track object
+        track object containing a list of all tracks
+    lat_bins: 1D np.array
+        array containg the latitude bins
+    lon_bins: 1D np.array
+        array containg the longitude bins
+
+    Returns:
+    --------
+    hist_count: 2D np.array
+        array containing the number of genesis points per grid cell
+    """
+    # Extract the first lat and lon from each dataset
+    first_lats = np.array([ds.lat.values[0] for ds in tc_track.data])
+    first_lons = np.array([ds.lon.values[0] for ds in tc_track.data])
+
+    # compute 2D density of genesis points
+    hist_count, _, _ = np.histogram2d(
+        first_lats,
+        first_lons,
+        bins=[lat_bins, lon_bins],
+        density=False,
+    )
+
+    return hist_count
+
+
+def normalize_hist(
+    res: int,
+    hist_count: np.ndarray,
+    norm: str,
+) -> np.ndarray:
+    """Normalize the number of points per grid cell by the grid cell area or by the total sum of
+    the histogram count.
+
+    Parameters:
+    -----------
+    res: float
+        resolution of grid cells
+    hist_count: 2D np.array
+        Array containing the count of tracks or genesis locations
+    norm: str
+        if norm = area normalize by gird cell area, if norm = sum normalize by the total sum
+    Returns:
+    ---------
+
+    """
+
+    if norm == "area":
+        grid_area, _ = u_coord.compute_grid_cell_area(res=res)
+        norm_hist: np.ndarray = hist_count / grid_area
+    elif norm == "sum":
+        norm_hist: np.ndarray = hist_count / hist_count.sum()
+
+    return norm_hist