climplot 0.1.0__py3-none-any.whl

climplot/metrics.py

@@ -0,0 +1,404 @@
+"""
+Area-weighted statistical metrics for climate data.
+
+This module provides functions for calculating area-weighted statistics
+that are essential for accurate global and regional analyses of gridded
+climate data.
+
+Why Area Weighting Matters
+--------------------------
+Ocean/atmosphere grid cells near the poles are much smaller than equatorial
+cells. Simple averaging treats all cells equally, over-weighting high
+latitudes. This can produce errors of ~1 cm in global sea level means
+or ~2 cm in RMSE calculations.
+
+Examples
+--------
+>>> import climplot
+>>> # Area-weighted global mean
+>>> gmsl = climplot.area_weighted_mean(ssh, areacello, dim=['yh', 'xh'])
+
+>>> # Area-weighted RMSE
+>>> rmse = climplot.area_weighted_rmse(model, obs, areacello, dim=['yh', 'xh'])
+"""
+
+import numpy as np
+import xarray as xr
+from typing import Union, List, Optional
+
+
+def area_weighted_mean(
+    data: xr.DataArray,
+    weights: xr.DataArray,
+    dim: Optional[Union[str, List[str]]] = None,
+) -> Union[xr.DataArray, float]:
+    """
+    Calculate area-weighted mean.
+
+    Parameters
+    ----------
+    data : DataArray
+        Data to average
+    weights : DataArray
+        Area weights (e.g., areacello grid cell areas)
+    dim : str or list of str, optional
+        Dimensions to average over. If None, averages over all.
+        Common: dim=['yh', 'xh'] for spatial mean.
+
+    Returns
+    -------
+    DataArray or float
+        Area-weighted mean
+
+    Examples
+    --------
+    >>> # Global mean SSH
+    >>> gmsl = climplot.area_weighted_mean(ssh, areacello, dim=['yh', 'xh'])
+
+    >>> # Time series of global mean
+    >>> gmsl_ts = climplot.area_weighted_mean(ssh_3d, areacello, dim=['yh', 'xh'])
+
+    Notes
+    -----
+    - NaN values are automatically excluded
+    - Weights are normalized internally
+    """
+    # Mask where data or weights are NaN
+    valid_mask = ~(np.isnan(data) | np.isnan(weights))
+
+    data_masked = data.where(valid_mask)
+    weights_masked = weights.where(valid_mask)
+
+    weighted_sum = (data_masked * weights_masked).sum(dim=dim, skipna=True)
+    total_weight = weights_masked.sum(dim=dim, skipna=True)
+
+    return weighted_sum / total_weight
+
+
+def area_weighted_bias(
+    model: xr.DataArray,
+    obs: xr.DataArray,
+    weights: xr.DataArray,
+    dim: Optional[Union[str, List[str]]] = None,
+) -> Union[xr.DataArray, float]:
+    """
+    Calculate area-weighted mean bias (model - obs).
+
+    Parameters
+    ----------
+    model : DataArray
+        Model data
+    obs : DataArray
+        Observational data
+    weights : DataArray
+        Area weights
+    dim : str or list of str, optional
+        Dimensions to average over
+
+    Returns
+    -------
+    DataArray or float
+        Mean bias (positive = model too high)
+
+    Examples
+    --------
+    >>> bias = climplot.area_weighted_bias(ssh_model, ssh_obs, areacello, dim=['yh', 'xh'])
+    """
+    diff = model - obs
+    return area_weighted_mean(diff, weights, dim=dim)
+
+
+def area_weighted_rmse(
+    model: xr.DataArray,
+    obs: xr.DataArray,
+    weights: xr.DataArray,
+    dim: Optional[Union[str, List[str]]] = None,
+) -> Union[xr.DataArray, float]:
+    """
+    Calculate area-weighted root-mean-square error.
+
+    Parameters
+    ----------
+    model : DataArray
+        Model data
+    obs : DataArray
+        Observational data
+    weights : DataArray
+        Area weights
+    dim : str or list of str, optional
+        Dimensions to average over
+
+    Returns
+    -------
+    DataArray or float
+        RMSE (always positive)
+
+    Examples
+    --------
+    >>> rmse = climplot.area_weighted_rmse(ssh_model, ssh_obs, areacello, dim=['yh', 'xh'])
+    """
+    squared_diff = (model - obs) ** 2
+    mse = area_weighted_mean(squared_diff, weights, dim=dim)
+    return np.sqrt(mse)
+
+
+def area_weighted_std(
+    data: xr.DataArray,
+    weights: xr.DataArray,
+    dim: Optional[Union[str, List[str]]] = None,
+) -> Union[xr.DataArray, float]:
+    """
+    Calculate area-weighted standard deviation.
+
+    Parameters
+    ----------
+    data : DataArray
+        Data to compute std for
+    weights : DataArray
+        Area weights
+    dim : str or list of str, optional
+        Dimensions to compute over
+
+    Returns
+    -------
+    DataArray or float
+        Standard deviation
+
+    Examples
+    --------
+    >>> spatial_std = climplot.area_weighted_std(ssh, areacello, dim=['yh', 'xh'])
+    """
+    mean = area_weighted_mean(data, weights, dim=dim)
+    squared_dev = (data - mean) ** 2
+    variance = area_weighted_mean(squared_dev, weights, dim=dim)
+    return np.sqrt(variance)
+
+
+def area_weighted_corr(
+    x: xr.DataArray,
+    y: xr.DataArray,
+    weights: xr.DataArray,
+    dim: Optional[Union[str, List[str]]] = None,
+) -> Union[xr.DataArray, float]:
+    """
+    Calculate area-weighted pattern correlation coefficient.
+
+    Parameters
+    ----------
+    x : DataArray
+        First field (e.g., model)
+    y : DataArray
+        Second field (e.g., observations)
+    weights : DataArray
+        Area weights
+    dim : str or list of str, optional
+        Dimensions to compute over
+
+    Returns
+    -------
+    DataArray or float
+        Correlation coefficient (-1 to 1)
+
+    Examples
+    --------
+    >>> r = climplot.area_weighted_corr(ssh_model, ssh_obs, areacello, dim=['yh', 'xh'])
+    """
+    x_mean = area_weighted_mean(x, weights, dim=dim)
+    y_mean = area_weighted_mean(y, weights, dim=dim)
+
+    x_anom = x - x_mean
+    y_anom = y - y_mean
+
+    covariance = area_weighted_mean(x_anom * y_anom, weights, dim=dim)
+    x_var = area_weighted_mean(x_anom**2, weights, dim=dim)
+    y_var = area_weighted_mean(y_anom**2, weights, dim=dim)
+
+    return covariance / np.sqrt(x_var * y_var)
+
+
+# ============================================================================
+# Time Series Metrics (1D arrays, no area weighting)
+# ============================================================================
+
+
+def timeseries_bias(
+    model: Union[np.ndarray, xr.DataArray],
+    obs: Union[np.ndarray, xr.DataArray],
+) -> float:
+    """
+    Calculate mean bias for time series.
+
+    Parameters
+    ----------
+    model : array-like
+        Model time series
+    obs : array-like
+        Observed time series
+
+    Returns
+    -------
+    float
+        Mean bias (positive = model too high)
+
+    Examples
+    --------
+    >>> bias = climplot.timeseries_bias(gmsl_model, gmsl_obs)
+    """
+    diff = np.array(model) - np.array(obs)
+    return float(np.nanmean(diff))
+
+
+def timeseries_rmse(
+    model: Union[np.ndarray, xr.DataArray],
+    obs: Union[np.ndarray, xr.DataArray],
+) -> float:
+    """
+    Calculate RMSE for time series.
+
+    Parameters
+    ----------
+    model : array-like
+        Model time series
+    obs : array-like
+        Observed time series
+
+    Returns
+    -------
+    float
+        RMSE (always positive)
+
+    Examples
+    --------
+    >>> rmse = climplot.timeseries_rmse(gmsl_model, gmsl_obs)
+    """
+    diff = np.array(model) - np.array(obs)
+    return float(np.sqrt(np.nanmean(diff**2)))
+
+
+def timeseries_corr(
+    x: Union[np.ndarray, xr.DataArray],
+    y: Union[np.ndarray, xr.DataArray],
+) -> float:
+    """
+    Calculate correlation coefficient for time series.
+
+    Parameters
+    ----------
+    x : array-like
+        First time series
+    y : array-like
+        Second time series
+
+    Returns
+    -------
+    float
+        Correlation coefficient (-1 to 1)
+
+    Examples
+    --------
+    >>> r = climplot.timeseries_corr(gmsl_model, gmsl_obs)
+    """
+    x_arr = np.array(x)
+    y_arr = np.array(y)
+
+    valid = ~(np.isnan(x_arr) | np.isnan(y_arr))
+    x_valid = x_arr[valid]
+    y_valid = y_arr[valid]
+
+    return float(np.corrcoef(x_valid, y_valid)[0, 1])
+
+
+def timeseries_std(data: Union[np.ndarray, xr.DataArray]) -> float:
+    """
+    Calculate standard deviation for time series.
+
+    Parameters
+    ----------
+    data : array-like
+        Time series data
+
+    Returns
+    -------
+    float
+        Standard deviation
+
+    Examples
+    --------
+    >>> std = climplot.timeseries_std(gmsl)
+    """
+    return float(np.nanstd(data))
+
+
+# ============================================================================
+# Summary Functions
+# ============================================================================
+
+
+def metrics_summary(
+    model: xr.DataArray,
+    obs: xr.DataArray,
+    weights: xr.DataArray,
+    dim: Optional[Union[str, List[str]]] = None,
+    name: str = "Model",
+) -> dict:
+    """
+    Calculate comprehensive validation metrics.
+
+    Parameters
+    ----------
+    model : DataArray
+        Model data
+    obs : DataArray
+        Observational data
+    weights : DataArray
+        Area weights
+    dim : str or list of str, optional
+        Dimensions to compute over
+    name : str, optional
+        Model name for output
+
+    Returns
+    -------
+    dict
+        Dictionary with bias, rmse, correlation, means, and stds
+
+    Examples
+    --------
+    >>> metrics = climplot.metrics_summary(ssh_model, ssh_obs, areacello, dim=['yh', 'xh'])
+    >>> print(f"RMSE: {metrics['rmse']:.3f}")
+    """
+    return {
+        "bias": float(area_weighted_bias(model, obs, weights, dim=dim)),
+        "rmse": float(area_weighted_rmse(model, obs, weights, dim=dim)),
+        "correlation": float(area_weighted_corr(model, obs, weights, dim=dim)),
+        "model_mean": float(area_weighted_mean(model, weights, dim=dim)),
+        "obs_mean": float(area_weighted_mean(obs, weights, dim=dim)),
+        "model_std": float(area_weighted_std(model, weights, dim=dim)),
+        "obs_std": float(area_weighted_std(obs, weights, dim=dim)),
+    }
+
+
+def print_metrics_summary(metrics: dict, name: str = "Model"):
+    """
+    Print formatted metrics summary.
+
+    Parameters
+    ----------
+    metrics : dict
+        Dictionary from metrics_summary()
+    name : str
+        Model name for display
+
+    Examples
+    --------
+    >>> metrics = climplot.metrics_summary(model, obs, weights, dim=['yh', 'xh'])
+    >>> climplot.print_metrics_summary(metrics, name='OM5 b01')
+    """
+    print(f"\n{name} Validation Metrics:")
+    print(f"  Bias (model - obs):     {metrics['bias']:+.4f}")
+    print(f"  RMSE:                   {metrics['rmse']:.4f}")
+    print(f"  Correlation:            {metrics['correlation']:.4f}")
+    print(f"  Model mean:             {metrics['model_mean']:.4f}")
+    print(f"  Obs mean:               {metrics['obs_mean']:.4f}")
+    print(f"  Model std:              {metrics['model_std']:.4f}")
+    print(f"  Obs std:                {metrics['obs_std']:.4f}")

climplot/panels.py

@@ -0,0 +1,295 @@
+"""
+Multi-panel figure utilities.
+
+This module provides functions for creating and labeling multi-panel
+figures with consistent styling and colorbars.
+
+Examples
+--------
+>>> import climplot
+>>> climplot.publication()
+>>> fig, axes = climplot.panel_figure(2, 3)  # 2 rows, 3 columns
+"""
+
+import matplotlib.pyplot as plt
+from pathlib import Path
+from typing import Tuple, Optional, List, Union
+import numpy as np
+
+
+def panel_figure(
+    nrows: int,
+    ncols: int,
+    figsize: Optional[Tuple[float, float]] = None,
+    projection=None,
+    constrained_layout: bool = True,
+    **kwargs,
+) -> Tuple[plt.Figure, np.ndarray]:
+    """
+    Create a multi-panel figure.
+
+    Parameters
+    ----------
+    nrows : int
+        Number of rows
+    ncols : int
+        Number of columns
+    figsize : tuple, optional
+        Figure size. If None, auto-calculated based on panels.
+    projection : cartopy.crs.Projection, optional
+        Map projection for all panels. If None, creates regular axes.
+    constrained_layout : bool, optional
+        Use constrained layout for better spacing. Default is True.
+    **kwargs
+        Additional arguments passed to plt.subplots()
+
+    Returns
+    -------
+    fig : Figure
+        The figure object
+    axes : ndarray
+        Array of axes objects (shape: nrows x ncols)
+
+    Examples
+    --------
+    >>> # 2x3 panel figure
+    >>> fig, axes = climplot.panel_figure(2, 3)
+    >>> for ax in axes.flat:
+    ...     ax.plot([1, 2, 3], [1, 4, 9])
+
+    >>> # Panel figure with map projections
+    >>> import cartopy.crs as ccrs
+    >>> fig, axes = climplot.panel_figure(2, 2, projection=ccrs.Robinson())
+    """
+    if figsize is None:
+        # Auto-calculate size
+        panel_width = 3.5 if projection is None else 3.0
+        panel_height = 2.5 if projection is None else 2.0
+        figsize = (ncols * panel_width, nrows * panel_height)
+
+    # Handle map projections
+    if projection is not None:
+        kwargs["subplot_kw"] = {"projection": projection}
+
+    fig, axes = plt.subplots(
+        nrows, ncols, figsize=figsize, constrained_layout=constrained_layout, **kwargs
+    )
+
+    # Ensure axes is always 2D array for consistency
+    if nrows == 1 and ncols == 1:
+        axes = np.array([[axes]])
+    elif nrows == 1:
+        axes = axes.reshape(1, -1)
+    elif ncols == 1:
+        axes = axes.reshape(-1, 1)
+
+    return fig, axes
+
+
+def add_panel_labels(
+    axes: Union[np.ndarray, List[plt.Axes]],
+    labels: Optional[List[str]] = None,
+    fontsize: int = 10,
+    fontweight: str = "bold",
+    x: float = 0.02,
+    y: float = 0.98,
+    ha: str = "left",
+    va: str = "top",
+):
+    """
+    Add panel labels (a. b. c. d.) to multi-panel figures.
+
+    Parameters
+    ----------
+    axes : array-like
+        List or array of axes objects
+    labels : list of str, optional
+        Custom labels. If None, uses lowercase letters with periods.
+    fontsize : int, optional
+        Font size. Default is 10.
+    fontweight : str, optional
+        Font weight. Default is 'bold'.
+    x, y : float, optional
+        Position in axes coordinates. Default is (0.02, 0.98).
+    ha, va : str, optional
+        Horizontal/vertical alignment. Default is 'left', 'top'.
+
+    Examples
+    --------
+    >>> fig, axes = climplot.panel_figure(2, 2)
+    >>> climplot.add_panel_labels(axes.flatten())
+    """
+    # Flatten if needed
+    if hasattr(axes, "flatten"):
+        axes = axes.flatten()
+
+    if labels is None:
+        labels = [f"{chr(97 + i)}." for i in range(len(axes))]
+
+    for ax, label in zip(axes, labels):
+        ax.text(
+            x,
+            y,
+            label,
+            transform=ax.transAxes,
+            fontsize=fontsize,
+            fontweight=fontweight,
+            va=va,
+            ha=ha,
+        )
+
+
+def add_colorbar(
+    mappable,
+    ax: plt.Axes,
+    label: str,
+    orientation: str = "horizontal",
+    extend: str = "both",
+    **kwargs,
+) -> plt.colorbar:
+    """
+    Add a colorbar with consistent formatting.
+
+    Parameters
+    ----------
+    mappable : ScalarMappable
+        The plot object (e.g., from pcolormesh)
+    ax : Axes
+        The axes to attach colorbar to
+    label : str
+        Colorbar label with units
+    orientation : str, optional
+        'horizontal' or 'vertical'. Default is 'horizontal'.
+    extend : str, optional
+        Out-of-range handling. Default is 'both'.
+    **kwargs
+        Additional arguments passed to plt.colorbar()
+
+    Returns
+    -------
+    Colorbar
+        The colorbar object
+
+    Examples
+    --------
+    >>> cs = ax.pcolormesh(lon, lat, data)
+    >>> cbar = climplot.add_colorbar(cs, ax, 'Temperature (K)')
+    """
+    if orientation == "horizontal":
+        cbar_kwargs = {"orientation": "horizontal", "pad": 0.05, "fraction": 0.046}
+    else:
+        cbar_kwargs = {"orientation": "vertical", "pad": 0.05, "fraction": 0.046}
+
+    cbar_kwargs.update(kwargs)
+    cbar = plt.colorbar(mappable, ax=ax, extend=extend, **cbar_kwargs)
+    cbar.set_label(label, fontsize=8)
+    cbar.ax.tick_params(labelsize=8)
+
+    return cbar
+
+
+def bottom_colorbar(
+    mappable,
+    fig: plt.Figure,
+    axes: np.ndarray,
+    label: str,
+    extend: str = "both",
+    **kwargs,
+) -> plt.colorbar:
+    """
+    Add a single colorbar below all panels.
+
+    Parameters
+    ----------
+    mappable : ScalarMappable
+        The plot object (e.g., from pcolormesh)
+    fig : Figure
+        The figure object
+    axes : ndarray
+        Array of all axes
+    label : str
+        Colorbar label with units
+    extend : str, optional
+        Out-of-range handling. Default is 'both'.
+    **kwargs
+        Additional arguments passed to fig.colorbar()
+
+    Returns
+    -------
+    Colorbar
+        The colorbar object
+
+    Examples
+    --------
+    >>> fig, axes = climplot.panel_figure(2, 2)
+    >>> for ax in axes.flat:
+    ...     cs = ax.pcolormesh(lon, lat, data)
+    >>> cbar = climplot.bottom_colorbar(cs, fig, axes, 'SSH (m)')
+    """
+    # Flatten axes
+    if hasattr(axes, "flatten"):
+        axes_flat = axes.flatten()
+    else:
+        axes_flat = axes
+
+    cbar = fig.colorbar(
+        mappable,
+        ax=list(axes_flat),
+        orientation="horizontal",
+        pad=0.05,
+        fraction=0.03,
+        extend=extend,
+        **kwargs,
+    )
+    cbar.set_label(label, fontsize=8)
+    cbar.ax.tick_params(labelsize=8)
+
+    return cbar
+
+
+def save_figure(
+    filename: str,
+    output_dir: str = "figures",
+    dpi: int = 300,
+    **kwargs,
+) -> Path:
+    """
+    Save figure with consistent settings.
+
+    Parameters
+    ----------
+    filename : str
+        Filename (should include extension)
+    output_dir : str or Path, optional
+        Directory to save to. Default is 'figures'.
+    dpi : int, optional
+        Resolution. Default is 300.
+    **kwargs
+        Additional arguments passed to plt.savefig()
+
+    Returns
+    -------
+    Path
+        Full path where figure was saved
+
+    Examples
+    --------
+    >>> fig, ax = plt.subplots()
+    >>> ax.plot([1, 2, 3], [1, 4, 9])
+    >>> climplot.save_figure('test_plot.png')
+    """
+    output_dir = Path(output_dir)
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    save_path = output_dir / filename
+
+    save_kwargs = {
+        "dpi": dpi,
+        "bbox_inches": "tight",
+        "facecolor": "white",
+    }
+    save_kwargs.update(kwargs)
+
+    plt.savefig(save_path, **save_kwargs)
+
+    return save_path