Changeset 3783 for trunk/LMDZ.MARS/util


Ignore:
Timestamp:
May 28, 2025, 5:31:59 PM (3 weeks ago)
Author:
jbclement
Message:

Mars PCM:
Big improvement of Python scripts in util folder to analyse/display variables in NetCDF files.
JBC

Location:
trunk/LMDZ.MARS/util
Files:
2 edited

Legend:

Unmodified
Added
Removed

  • trunk/LMDZ.MARS/util/analyse_netcdf.py

    r3648 r3783  
     1#!/usr/bin/env python3
    12############################################################
    23### Python script to analyse a NetCDF file for debugging ###
    34############################################################
    45
    5 ### This script gives useful information about a NetCDF file
    6 ### to help for debugging. For each variable, it outputs the
    7 ### dimensions, the min & max values, the average value and
    8 ### warns the user in case of NaN or negative values.
    9 ### The file name is asked to the user in the terminal.
     6
     7"""
     8For each numeric variable, it outputs:
     9  - Dimensions and shape
     10  - Minimum & maximum values (ignoring NaNs)
     11  - Mean value (ignoring NaNs)
     12  - Warnings if the variable is entirely NaN or contains any NaNs/negative values
     13
     14Usage:
     15  1) Command-line mode:
     16       python analyze_netcdf.py /path/to/your_file.nc
     17 
     18  2) Interactive mode through the prompt:
     19       python analyze_netcdf.py
     20"""
     21
    1022
    1123import os
     24import sys
     25import glob
    1226import readline
    13 import glob
     27import argparse
     28import numpy as np
    1429from netCDF4 import Dataset
    15 import numpy as np
    1630
    17 ############################################################
    18 ### Setup readline for file name autocompletion
    19 def complete(text,state):
    20     line = readline.get_line_buffer().split()
    21     # Use glob to find all matching files/directories for the current text
    22     if '*' not in text:
    23         text += '*'
    24     matches = glob.glob(os.path.expanduser(text))
    25     # Add '/' if the match is a directory
    26     matches = [match + '/' if os.path.isdir(match) else match for match in matches]
    27    
     31
     32def complete_filename(text, state):
     33    """
     34    Tab-completion function for readline: completes filesystem paths.
     35    Appends '/' if the match is a directory.
     36    """
     37    # The text forms a partial path; glob for matching entries
     38    if "*" not in text:
     39        text_glob = text + "*"
     40    else:
     41        text_glob = text
     42    matches = glob.glob(os.path.expanduser(text_glob))
     43    # Add a trailing slash for directories
     44    matches = [m + "/" if os.path.isdir(m) else m for m in matches]
    2845    try:
    2946        return matches[state]
    … …  
    3148        return None
    3249
    33 ### Function to analyze a variable in a NetCDF file
     50
    3451def analyze_variable(variable):
    35     # Get the data for the variable
    36     data = variable[:]
     52    """
     53    Print summary statistics (min, max, mean) for a numeric NetCDF variable.
     54    Ignores NaNs when computing min/max/mean. Warns if any NaNs or negatives exist.
     55    """
     56    name = variable.name
     57    dims = variable.dimensions
     58    shape = variable.shape
    3759   
    38     # Calculate min, max and mean
    39     if np.isnan(data).all():
    40         min_val = np.nan
    41         max_val = np.nan
    42         mean_val = np.nan
    43     else:
    44         data_min = np.nanmin(data) # Min value ignoring NaN
    45         data_max = np.nanmax(data) # Max value ignoring NaN
    46         data_mean = np.nanmean(data) # Mean value ignoring NaN
    47    
    48     # Check if there are any NaN values
     60    try:
     61        # Read the entire array into memory; this may be large for huge datasets
     62        data = variable[:]
     63    except Exception as e:
     64        print(f"\nUnable to read variable '{name}': {e}")
     65        return
     66
     67    # If the array is a masked array, convert to a NumPy array with masked values as np.nan
     68    if hasattr(data, "mask"):
     69        # Fill masked entries with NaN so that np.nanmin / np.nanmax works correctly
     70        data = np.where(data.mask, np.nan, data.data)
     71
     72    # Determine if the variable has any valid (finite) data at all
     73    if np.all(np.isnan(data)):
     74        # Entirely NaN (or entirely masked)
     75        print(f"\nAnalysis of variable: {name}")
     76        print(f"  Dimensions: {dims}")
     77        print(f"  Shape     : {shape}")
     78        print("  Entire variable is NaN or masked.")
     79        return
     80
     81    # Compute min, max, mean ignoring NaNs
     82    data_min = np.nanmin(data)
     83    data_max = np.nanmax(data)
     84    data_mean = np.nanmean(data)
     85
     86    # Check for presence of NaNs and negative values
    4987    has_nan = np.isnan(data).any()
     88    has_negative = np.any(data < 0)
    5089
    51     # Check for negative values
    52     has_negative = (data < 0).any()
    53    
    54     # Print the results
    55     print(f"\nAnalysis of variable: {variable.name}")
    56     print(f"  Dimensions: {variable.dimensions}")
     90    # Output
     91    print(f"\nAnalysis of variable: {name}")
     92    print(f"  Dimensions: {dims}")
     93    print(f"  Shape     : {shape}")
    5794    print(f"  Min value : {data_min:>12.6e}")
    5895    print(f"  Max value : {data_max:>12.6e}")
    … …  
    63100        print(f"  \033[93mWarning: contains negative values!\033[0m")
    64101
    65 ### Main function
    66 def analyze_netcdf():
    67     # Ask for the file name
    68     readline.set_completer(complete)
    69     readline.parse_and_bind('tab: complete')
    70     file = input("Enter the name of the NetCDF file: ")
    71    
    72     # Open the NetCDF file
     102def analyze_netcdf_file(nc_path):
     103    """
     104    Open the NetCDF file at nc_path and analyze each numeric variable.
     105    """
     106    if not os.path.isfile(nc_path):
     107        print(f"Error: File '{nc_path}' not found.")
     108        return
     109
    73110    try:
    74         dataset = Dataset(file,mode='r')
    75     except FileNotFoundError:
    76         print(f"File '{file}' not found.")
     111        ds = Dataset(nc_path, mode='r')
     112    except Exception as e:
     113        print(f"Error: Unable to open '{nc_path}': {e}")
    77114        return
    78    
    79     # Iterate through all variables in the dataset to analyze them
    80     for variable_name in dataset.variables:
    81         variable = dataset.variables[variable_name]
    82         if np.issubdtype(variable[:].dtype,np.number):
     115
     116    print(f"\nOpened NetCDF file: {nc_path}")
     117    print(f"Number of variables: {len(ds.variables)}")
     118
     119    for var_name, variable in ds.variables.items():
     120        # Attempt to check if the dtype is numeric
     121        try:
     122            dtype = variable.dtype
     123        except Exception:
     124            # If reading dtype fails, skip it
     125            print(f"\nSkipping variable with unknown type: {var_name}")
     126            continue
     127
     128        if np.issubdtype(dtype, np.number) or hasattr(variable[:], "mask"):
    83129            analyze_variable(variable)
    84130        else:
    85             print(f"\nSkipping non-numeric variable: {variable.name}")
    86    
    87     # Close the NetCDF file
    88     dataset.close()
     131            print(f"\nSkipping non-numeric variable: {var_name}")
    89132
    90 ### Call the main function
    91 analyze_netcdf()
     133    ds.close()
     134    print("\nFinished analysis.\n")
     135
     136
     137def main():
     138    parser = argparse.ArgumentParser(
     139        description="Analyze a NetCDF file and report min/max/mean for each numeric variable."
     140    )
     141    parser.add_argument(
     142        "nc_file",
     143        nargs="?",
     144        help="Path to the NetCDF file (if omitted, you'll be prompted)."
     145    )
     146    args = parser.parse_args()
     147
     148    if args.nc_file:
     149        # Command-line mode: directly analyze the provided file path
     150        analyze_netcdf_file(args.nc_file)
     151    else:
     152        # Interactive mode: enable tab completion for filenames
     153        readline.set_completer(complete_filename)
     154        readline.parse_and_bind("tab: complete")
     155        try:
     156            user_input = input("Enter the path to the NetCDF file: ").strip()
     157        except (EOFError, KeyboardInterrupt):
     158            print("\nExiting.")
     159            return
     160
     161        if not user_input:
     162            print("No file specified. Exiting.")
     163            return
     164
     165        analyze_netcdf_file(user_input)
     166
     167
     168if __name__ == "__main__":
     169    main()
     170

  • trunk/LMDZ.MARS/util/display_netcdf.py

    r3681 r3783  
     1#!/usr/bin/env python3
    12##############################################################
    23### Python script to visualize a variable in a NetCDF file ###
    34##############################################################
    45
    5 ### This script can display any variable of a NetCDF file.
    6 ### The file name, the variable to display and eventually the
    7 ### dimension are asked to the user in the terminal.
     6"""
     7This script can display any numeric variable from a NetCDF file on a lat/lon map.
     8It supports variables of dimension:
     9  - (latitude, longitude)
     10  - (time, latitude, longitude)
     11  - (altitude, latitude, longitude)
     12  - (time, altitude, latitude, longitude)
     13
     14Usage:
     15  1) Command-line mode:
     16       python display_netcdf.py /path/to/your_file.nc --variable VAR_NAME [--time-index 0] [--alt-index 0] [--cmap viridis] [--output out.png]
     17
     18  2) Interactive mode through the prompt:
     19       python display_netcdf.py
     20
     21The script will:
     22  > Attempt to locate latitude and longitude variables (searching for
     23    names like 'latitude', 'lat', 'longitude', 'lon').
     24  > If there is exactly one variable in the dataset, select it automatically.
     25  > Prompt for time/altitude indices if needed (or accept via CLI).
     26  > Handle masked arrays, converting masked values to NaN.
     27  > Plot with a default colormap ('jet'), adjustable via --cmap.
     28  > Optionally save the figure instead of displaying it interactively.
     29"""
     30
    831
    932import os
     33import sys
     34import glob
    1035import readline
    11 import glob
     36import argparse
     37import numpy as np
     38import matplotlib.pyplot as plt
    1239from netCDF4 import Dataset
    13 import matplotlib.pyplot as plt
    14 import numpy as np
    15 
    16 ##############################################################
    17 ### Setup readline for file name autocompletion
    18 def complete(text,state):
    19     line = readline.get_line_buffer().split()
    20     # Use glob to find all matching files/directories for the current text
    21     if '*' not in text:
    22         text += '*'
    23     matches = glob.glob(os.path.expanduser(text))
    24     # Add '/' if the match is a directory
    25     matches = [match + '/' if os.path.isdir(match) else match for match in matches]
     40
     41
     42def complete_filename(text, state):
     43    """
     44    Readline tab-completion function for filesystem paths.
     45    """
     46    if "*" not in text:
     47        pattern = text + "*"
     48    else:
     49        pattern = text
     50    matches = glob.glob(os.path.expanduser(pattern))
     51    matches = [m + "/" if os.path.isdir(m) else m for m in matches]
    2652    try:
    2753        return matches[state]
    … …  
    2955        return None
    3056
    31 ### Function to handle autocomplete for variable names
    32 def complete_variable_names(variable_names):
     57
     58def make_varname_completer(varnames):
     59    """
     60    Returns a readline completer function for the given list of variable names.
     61    """
    3362    def completer(text, state):
    34         options = [name for name in variable_names if name.startswith(text)]
    35         if state < len(options):
     63        options = [name for name in varnames if name.startswith(text)]
     64        try:
    3665            return options[state]
    37         else:
     66        except IndexError:
    3867            return None
    3968    return completer
    4069
    41 ### Function to visualize a variable from a NetCDF file
    42 def visualize_variable():
    43     # Ask for the NetCDF file name
    44     readline.set_completer(complete)
    45     readline.parse_and_bind('tab: complete')
    46     file = input("Enter the name of the NetCDF file: ")
    47    
    48     # Open the NetCDF file
     70
     71# Helper functions to detect common dimension names
     72TIME_DIMS = ("Time", "time", "time_counter")
     73ALT_DIMS  = ("altitude",)
     74LAT_DIMS  = ("latitude", "lat")
     75LON_DIMS  = ("longitude", "lon")
     76
     77
     78def find_dim_index(dims, candidates):
     79    """
     80    Search through dims tuple for any name in candidates.
     81    Returns the index if found, else returns None.
     82    """
     83    for idx, dim in enumerate(dims):
     84        for cand in candidates:
     85            if cand.lower() == dim.lower():
     86                return idx
     87    return None
     88
     89
     90def find_coord_var(dataset, candidates):
     91    """
     92    Among dataset variables, return the first variable whose name matches any candidate.
     93    Returns None if none found.
     94    """
     95    for name in dataset.variables:
     96        for cand in candidates:
     97            if cand.lower() == name.lower():
     98                return name
     99    return None
     100
     101
     102# Core plotting helper
     103
     104def plot_variable(dataset, varname, time_index=None, alt_index=None, colormap="jet", output_path=None):
     105    """
     106    Extracts the requested slice from the variable and plots it on a lat/lon grid.
     107
     108    Parameters
     109    ----------
     110    dataset    : netCDF4.Dataset object (already open)
     111    varname    : string name of the variable to plot
     112    time_index : int or None (if variable has a time dimension)
     113    alt_index  : int or None (if variable has an altitude dimension)
     114    colormap   : string colormap name (passed to plt.contourf)
     115    output_path: string filepath to save figure, or None to display interactively
     116    """
     117    var = dataset.variables[varname]
     118    dims = var.dimensions
     119
    49120    try:
    50         dataset = Dataset(file,mode='r')
    51     except FileNotFoundError:
    52         print(f"File '{file}' not found.")
    53         return
    54 
    55     # Display available variables
    56     variable_names = list(dataset.variables.keys())
    57     print("Available variables:\n",variable_names)
    58    
    59     # Ask for the variable to display
    60     readline.set_completer(complete_variable_names(variable_names))
    61     variable_name = input("\nEnter the name of the variable you want to visualize: ")
    62    
    63     # Check if the variable exists
    64     if variable_name not in dataset.variables:
    65         print(f"Variable '{variable_name}' not found in the dataset.")
    66         dataset.close()
    67         return
    68    
    69     # Extract the selected variable
    70     variable = dataset.variables[variable_name][:]
    71    
    72     # Extract latitude, longitude and altitude
    73     latitude = dataset.variables['latitude'][:]
    74     longitude = dataset.variables['longitude'][:]
    75    
    76     # Check if the variable has altitude and time dimensions
    77     dimensions = dataset.variables[variable_name].dimensions
    78     print(f"\nDimensions of '{variable_name}': {dimensions}")
    79    
    80     # If the variable has a time dimension, ask for the time index
    81     if 'Time' in dimensions:
    82         if variable.shape[0] == 1:
    83             time_index = 0
     121        data_full = var[:]
     122    except Exception as e:
     123        print(f"Error: Cannot read data for variable '{varname}': {e}")
     124        return
     125
     126    # Convert masked array to NaN
     127    if hasattr(data_full, "mask"):
     128        data_full = np.where(data_full.mask, np.nan, data_full.data)
     129
     130    # Identify dimension indices
     131    t_idx = find_dim_index(dims, TIME_DIMS)
     132    a_idx = find_dim_index(dims, ALT_DIMS)
     133    lat_idx = find_dim_index(dims, LAT_DIMS)
     134    lon_idx = find_dim_index(dims, LON_DIMS)
     135
     136    # Check that lat and lon dims exist
     137    if lat_idx is None or lon_idx is None:
     138        print("Error: Could not find 'latitude'/'lat' and 'longitude'/'lon' dimensions for plotting.")
     139        return
     140
     141    # Build a slice with defaults
     142    slicer = [slice(None)] * len(dims)
     143    if t_idx is not None:
     144        if time_index is None:
     145            print("Error: Variable has a time dimension; please supply a time index.")
     146            return
     147        slicer[t_idx] = time_index
     148    if a_idx is not None:
     149        if alt_index is None:
     150            print("Error: Variable has an altitude dimension; please supply an altitude index.")
     151            return
     152        slicer[a_idx] = alt_index
     153
     154    # Extract the 2D slice
     155    try:
     156        data_slice = data_full[tuple(slicer)]
     157    except Exception as e:
     158        print(f"Error: Could not slice variable '{varname}': {e}")
     159        return
     160
     161    # After slicing, data_slice should be 2D
     162    if data_slice.ndim != 2:
     163        print(f"Error: After slicing, data for '{varname}' is not 2D (ndim={data_slice.ndim}).")
     164        return
     165
     166    nlat, nlon = data_slice.shape
     167    # Handle too-small grid (1x1, 1xN, Nx1)
     168    if nlat < 2 or nlon < 2:
     169        lat_varname = find_coord_var(dataset, LAT_DIMS)
     170        lon_varname = find_coord_var(dataset, LON_DIMS)
     171        if lat_varname and lon_varname:
     172            lat_vals = dataset.variables[lat_varname][:]
     173            lon_vals = dataset.variables[lon_varname][:]
     174            if hasattr(lat_vals, "mask"):
     175                lat_vals = np.where(lat_vals.mask, np.nan, lat_vals.data)
     176            if hasattr(lon_vals, "mask"):
     177                lon_vals = np.where(lon_vals.mask, np.nan, lon_vals.data)
     178            # Single point
     179            if nlat == 1 and nlon == 1:
     180                print(f"Single data point: value={data_slice[0,0]} at (lon={lon_vals[0]}, lat={lat_vals[0]})")
     181                return
     182            # 1 x N -> plot vs lon
     183            if nlat == 1 and nlon > 1:
     184                x = lon_vals
     185                y = data_slice[0, :]
     186                plt.figure()
     187                plt.plot(x, y, marker='o')
     188                plt.xlabel(f"Longitude ({getattr(dataset.variables[lon_varname], 'units', 'degrees')})")
     189                plt.ylabel(varname)
     190                plt.title(f"{varname} (lat={lat_vals[0]})")
     191            # N x 1 -> plot vs lat
     192            elif nlon == 1 and nlat > 1:
     193                x = lat_vals
     194                y = data_slice[:, 0]
     195                plt.figure()
     196                plt.plot(x, y, marker='o')
     197                plt.xlabel(f"Latitude ({getattr(dataset.variables[lat_varname], 'units', 'degrees')})")
     198                plt.ylabel(varname)
     199                plt.title(f"{varname} (lon={lon_vals[0]})")
     200            else:
     201                print("Unexpected slice shape.")
     202                return
     203            if output_path:
     204                plt.savefig(output_path, bbox_inches="tight")
     205                print(f"Figure saved to '{output_path}'")
     206            else:
     207                plt.show()
     208            plt.close()
     209            return
     210
     211    # Locate coordinate variables
     212    lat_varname = find_coord_var(dataset, LAT_DIMS)
     213    lon_varname = find_coord_var(dataset, LON_DIMS)
     214    if lat_varname is None or lon_varname is None:
     215        print("Error: Could not locate latitude/longitude variables in the dataset.")
     216        return
     217
     218    lat_var = dataset.variables[lat_varname][:]
     219    lon_var = dataset.variables[lon_varname][:]
     220    if hasattr(lat_var, "mask"):
     221        lat_var = np.where(lat_var.mask, np.nan, lat_var.data)
     222    if hasattr(lon_var, "mask"):
     223        lon_var = np.where(lon_var.mask, np.nan, lon_var.data)
     224
     225    # Build 2D coordinate arrays
     226    if lat_var.ndim == 1 and lon_var.ndim == 1:
     227        lon2d, lat2d = np.meshgrid(lon_var, lat_var)
     228    elif lat_var.ndim == 2 and lon_var.ndim == 2:
     229        lat2d, lon2d = lat_var, lon_var
     230    else:
     231        print("Error: Latitude and longitude variables must both be either 1D or 2D.")
     232        return
     233
     234    # Retrieve units if available
     235    var_units = getattr(var, "units", "")
     236    lat_units = getattr(dataset.variables[lat_varname], "units", "degrees")
     237    lon_units = getattr(dataset.variables[lon_varname], "units", "degrees")
     238
     239    # Plot
     240    plt.figure(figsize=(10, 6))
     241    cf = plt.contourf(lon2d, lat2d, data_slice, cmap=colormap)
     242    cbar = plt.colorbar(cf)
     243    if var_units:
     244        cbar.set_label(f"{varname} ({var_units})")
     245    else:
     246        cbar.set_label(varname)
     247
     248    plt.xlabel(f"Longitude ({lon_units})")
     249    plt.ylabel(f"Latitude ({lat_units})")
     250    plt.title(f"{varname} Visualization")
     251
     252    if output_path:
     253        try:
     254            plt.savefig(output_path, bbox_inches="tight")
     255            print(f"Figure saved to '{output_path}'")
     256        except Exception as e:
     257            print(f"Error saving figure: {e}")
     258    else:
     259        plt.show()
     260    plt.close()
     261
     262
     263def visualize_variable_interactive(nc_path=None):
     264    """
     265    Interactive mode: if nc_path is provided, skip prompting for filename.
     266    Otherwise, prompt for filename. Then select variable (automatically if only one), prompt for indices.
     267    """
     268    # Determine file path
     269    if nc_path:
     270        file_input = nc_path
     271    else:
     272        readline.set_completer(complete_filename)
     273        readline.parse_and_bind("tab: complete")
     274        file_input = input("Enter the path to the NetCDF file: ").strip()
     275
     276    if not file_input:
     277        print("No file specified. Exiting.")
     278        return
     279    if not os.path.isfile(file_input):
     280        print(f"Error: '{file_input}' not found.")
     281        return
     282
     283    try:
     284        ds = Dataset(file_input, mode="r")
     285    except Exception as e:
     286        print(f"Error: Unable to open '{file_input}': {e}")
     287        return
     288
     289    varnames = list(ds.variables.keys())
     290    if not varnames:
     291        print("Error: No variables found in the dataset.")
     292        ds.close()
     293        return
     294
     295    # Auto-select if only one variable
     296    if len(varnames) == 1:
     297        var_input = varnames[0]
     298        print(f"Automatically selected the only variable: '{var_input}'")
     299    else:
     300        print("\nAvailable variables:")
     301        for name in varnames:
     302            print(f"  - {name}")
     303        print()
     304        readline.set_completer(make_varname_completer(varnames))
     305        var_input = input("Enter the name of the variable to visualize: ").strip()
     306        if var_input not in ds.variables:
     307            print(f"Error: Variable '{var_input}' not found. Exiting.")
     308            ds.close()
     309            return
     310
     311    dims = ds.variables[var_input].dimensions
     312    time_idx = None
     313    alt_idx = None
     314
     315    t_idx = find_dim_index(dims, TIME_DIMS)
     316    if t_idx is not None:
     317        length = ds.variables[var_input].shape[t_idx]
     318        if length > 1:
     319            while True:
     320                try:
     321                    user_t = input(f"Enter time index [0..{length - 1}]: ").strip()
     322                    if user_t == "":
     323                        print("No time index entered. Exiting.")
     324                        ds.close()
     325                        return
     326                    time_idx = int(user_t)
     327                    if 0 <= time_idx < length:
     328                        break
     329                except ValueError:
     330                    pass
     331                print(f"Invalid index. Enter an integer between 0 and {length - 1}.")
    84332        else:
    85             time_index = int(input(f"Enter the time index (0 to {variable.shape[0] - 1}): "))
    86     else:
    87         time_index = None
    88    
    89     # If the variable has an altitude dimension, ask for the altitude index
    90     if 'altitude' in dimensions:
    91         altitude = dataset.variables['altitude'][:]
    92         altitude_index = int(input(f"Enter the altitude index (0 to {altitude.shape[0] - 1}): "))
    93     else:
    94         altitude_index = None
    95    
    96     # Prepare the 2D slice for plotting
    97     if time_index is not None and altitude_index is not None:
    98         data_slice = variable[time_index,altitude_index,:,:]
    99     elif time_index is not None:
    100         data_slice = variable[time_index,:,:]
    101     elif altitude_index is not None:
    102         data_slice = variable[altitude_index,:,:]
    103     else:
    104         data_slice = variable[:,:]
    105    
    106     # Plot the selected variable
    107     plt.figure(figsize = (10,6))
    108     plt.contourf(longitude,latitude,data_slice,cmap = 'jet')
    109     plt.colorbar(label=f"{variable_name.capitalize()} (units)") # Adjust units based on your data
    110     plt.xlabel('Longitude (degrees)')
    111     plt.ylabel('Latitude (degrees)')
    112     plt.title(f"{variable_name.capitalize()} visualization")
    113    
    114     # Show the plot
    115     plt.show()
    116    
    117     # Close the NetCDF file
    118     dataset.close()
    119 
    120 ### Call the main function
    121 visualize_variable()
     333            time_idx = 0
     334            print("Only one time step available; using index 0.")
     335
     336    a_idx = find_dim_index(dims, ALT_DIMS)
     337    if a_idx is not None:
     338        length = ds.variables[var_input].shape[a_idx]
     339        if length > 1:
     340            while True:
     341                try:
     342                    user_a = input(f"Enter altitude index [0..{length - 1}]: ").strip()
     343                    if user_a == "":
     344                        print("No altitude index entered. Exiting.")
     345                        ds.close()
     346                        return
     347                    alt_idx = int(user_a)
     348                    if 0 <= alt_idx < length:
     349                        break
     350                except ValueError:
     351                    pass
     352                print(f"Invalid index. Enter an integer between 0 and {length - 1}.")
     353        else:
     354            alt_idx = 0
     355            print("Only one altitude level available; using index 0.")
     356
     357    plot_variable(
     358        dataset=ds,
     359        varname=var_input,
     360        time_index=time_idx,
     361        alt_index=alt_idx,
     362        colormap="jet",
     363        output_path=None
     364    )
     365    ds.close()
     366
     367
     368def visualize_variable_cli(nc_path, varname, time_index, alt_index, colormap, output_path):
     369    """
     370    Command-line mode: directly visualize based on provided arguments.
     371    """
     372    if not os.path.isfile(nc_path):
     373        print(f"Error: '{nc_path}' not found.")
     374        return
     375    try:
     376        ds = Dataset(nc_path, mode="r")
     377    except Exception as e:
     378        print(f"Error: Unable to open '{nc_path}': {e}")
     379        return
     380
     381    if varname not in ds.variables:
     382        print(f"Error: Variable '{varname}' not found in '{nc_path}'.")
     383        ds.close()
     384        return
     385
     386    plot_variable(
     387        dataset=ds,
     388        varname=varname,
     389        time_index=time_index,
     390        alt_index=alt_index,
     391        colormap=colormap,
     392        output_path=output_path
     393    )
     394    ds.close()
     395
     396
     397def main():
     398    parser = argparse.ArgumentParser(
     399        description="Visualize a 2D slice of a NetCDF variable on a latitude-longitude map."
     400    )
     401    parser.add_argument(
     402        "nc_file",
     403        nargs="?",
     404        help="Path to the NetCDF file (interactive if omitted)."
     405    )
     406    parser.add_argument(
     407        "--variable", "-v",
     408        help="Name of the variable to visualize."
     409    )
     410    parser.add_argument(
     411        "--time-index", "-t",
     412        type=int,
     413        help="Index along the time dimension (if applicable)."
     414    )
     415    parser.add_argument(
     416        "--alt-index", "-a",
     417        type=int,
     418        help="Index along the altitude dimension (if applicable)."
     419    )
     420    parser.add_argument(
     421        "--cmap", "-c",
     422        default="jet",
     423        help="Matplotlib colormap (default: 'jet')."
     424    )
     425    parser.add_argument(
     426        "--output", "-o",
     427        help="If provided, save the plot to this file instead of displaying it."
     428    )
     429
     430    args = parser.parse_args()
     431
     432    # If nc_file is provided but variable is missing: ask only for variable
     433    if args.nc_file and not args.variable:
     434        visualize_variable_interactive(nc_path=args.nc_file)
     435    # If either nc_file or variable is missing, run fully interactive
     436    elif not args.nc_file or not args.variable:
     437        visualize_variable_interactive()
     438    else:
     439        visualize_variable_cli(
     440            nc_path=args.nc_file,
     441            varname=args.variable,
     442            time_index=args.time_index,
     443            alt_index=args.alt_index,
     444            colormap=args.cmap,
     445            output_path=args.output
     446        )
     447
     448
     449if __name__ == "__main__":
     450    main()
     451
Note: See TracChangeset for help on using the changeset viewer.