test.py - v2 - aine/test

test.py
144 строки5 KB
OlgaTop
feat: use filled contours with color scale based on actual data min/max
23 июн 2026, 14:40
23 июн 2026, 14:4064817ea
#!/opt/miniconda/envs/fairy/bin/python
"""
Read EPV from NetCDF file and plot it with geographic basemap in polar projection.
"""

import sys
import os
try:
    import xarray as xr
except ImportError:
    print("Error: xarray module not found. Please install it (e.g., pip install xarray).", file=sys.stderr)
    sys.exit(1)
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

# Cartopy for geographic projections and features
try:
    import cartopy.crs as ccrs
    import cartopy.feature as cfeature
except ImportError:
    print("Error: cartopy module not found. Please install it (e.g., pip install cartopy).", file=sys.stderr)
    sys.exit(1)

def main():
    filepath = "/data2/filesystem2/141f/all_winters_isentropic_data/nc_isentropic/2026/01/ncdir_840_day28/EPVday28.nc"
    try:
        ds = xr.open_dataset(filepath)
    except FileNotFoundError:
        print(f"Error: file not found at {filepath}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error opening file: {e}", file=sys.stderr)
        sys.exit(1)

    # Print variable names for debugging
    print("Variables in file:", list(ds.data_vars))

    # Assume EPV variable exists
    if 'EPV' not in ds.data_vars:
        print("Variable 'EPV' not found in file.", file=sys.stderr)
        ds.close()
        sys.exit(1)

    epv = ds['EPV'].values

    # Try to read latitude and longitude coordinates
    # Common names: lat, latitude, Lat, LAT, lon, longitude, Lon, LON
    lat_var = None
    lon_var = None
    for candidate in ['lat', 'latitude', 'Lat', 'LAT']:
        if candidate in ds.coords:
            lat_var = candidate
            break
    for candidate in ['lon', 'longitude', 'Lon', 'LON']:
        if candidate in ds.coords:
            lon_var = candidate
            break

    if lat_var is None or lon_var is None:
        print("Warning: latitude/longitude coordinates not found in dataset. Using index-based coordinates.", file=sys.stderr)
        # Create dummy coordinates based on indices
        lat = np.arange(epv.shape[-2])
        lon = np.arange(epv.shape[-1])
    else:
        lat = ds[lat_var].values
        lon = ds[lon_var].values

    ds.close()

    # If 4D (time, level, lat, lon) take first time and first level
    if epv.ndim == 4:
        epv_slice = epv[0, 0, :, :]
    elif epv.ndim == 3:
        epv_slice = epv[0, :, :]
    elif epv.ndim == 2:
        epv_slice = epv
    else:
        print(f"Unexpected dimensions: {epv.shape}", file=sys.stderr)
        sys.exit(1)

    # Ensure lat/lon are 2D for pcolormesh (meshgrid)
    if lat.ndim == 1 and lon.ndim == 1:
        lon2d, lat2d = np.meshgrid(lon, lat)
    else:
        lon2d, lat2d = lon, lat

    # Output directory and file
    outdir = "/data0/home/141/test/pic"
    outfile = os.path.join(outdir, "EPV_plot.png")
    os.makedirs(outdir, exist_ok=True)

    # Create polar stereographic projection
    proj = ccrs.NorthPolarStereo()

    # Create figure with projection
    fig, ax = plt.subplots(figsize=(10, 10),
                           subplot_kw={'projection': proj})

    # Add geographic features
    ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
    ax.add_feature(cfeature.BORDERS, linewidth=0.3, alpha=0.5)
    # Optionally add land/ocean
    # ax.add_feature(cfeature.LAND, color='lightgray')
    # ax.add_feature(cfeature.OCEAN, color='lightblue')

    # Set map extent (adjust as needed)
    ax.set_extent([-180, 180, 30, 90], crs=ccrs.PlateCarree())

    # Convert to PVU (multiply by 1e6)
    epv_pvu = epv_slice * 1e6

    # Plot EPV data using filled contours (contourf)
    # Data coordinates are in PlateCarree (lat/lon)
    # Determine levels based on actual min/max of data
    vmin = np.nanmin(epv_pvu)
    vmax = np.nanmax(epv_pvu)
    # Add a small margin to avoid clipping
    margin = 0.01 * (vmax - vmin) if vmax != vmin else 1.0
    levels = np.linspace(vmin - margin, vmax + margin, 20)

    contour = ax.contourf(lon2d, lat2d, epv_pvu,
                          levels=levels,
                          transform=ccrs.PlateCarree(),
                          cmap='viridis',
                          extend='both')

    # Add colorbar for filled contours
    cbar = plt.colorbar(contour, ax=ax, orientation='horizontal',
                        pad=0.05, shrink=0.8, label='EPV [PVU]')

    # Title
    ax.set_title('Ertel Potential Vorticity (EPV) [PVU] – Polar Stereographic')

    # Save figure
    plt.savefig(outfile, dpi=150, bbox_inches='tight')
    print(f"Plot saved as {outfile}")

    # Optionally show
    # plt.show()

if __name__ == '__main__':
    main()