test.py - v1 - aine/test

test.py
65 строк2 KB
OlgaTop
fix: update plot save path to specified directory
23 июн 2026, 13:12
23 июн 2026, 13:126c0b8c8
#!/opt/miniconda/envs/fairy/bin/python
"""
Read EPV from NetCDF file and plot it.
"""

import sys
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

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
    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)

    # Create plot
    plt.figure(figsize=(10, 6))
    contour = plt.contourf(epv_slice, levels=20, cmap='viridis')
    plt.colorbar(contour, label='EPV')
    plt.title('Ertel Potential Vorticity (EPV)')
    plt.xlabel('Longitude index')
    plt.ylabel('Latitude index')
    plt.tight_layout()
    plt.savefig('/data0/home/141/test/v1/EPV_plot.png', dpi=150)
    print("Plot saved as /data0/home/141/test/v1/EPV_plot.png")
    # Optionally show
    # plt.show()

if __name__ == '__main__':
    main()