Era5MBWindtoDeflt3D

"""
Created on Jun 19 2024

@author: Faezah Maghsoodifar

This script processes wind and pressure data from a netCDF file generated by the ERA5 dataset.
It reads mean sea level pressure, eastward wind component, and northward wind component values,
flipping the data along the latitude axis for proper alignment. The script preserves time units
from the input file and saves the processed data into a new netCDF file in the NETCDF4 format.
The output file includes dimensions for time, latitude, and longitude, along with variables for
pressure and wind components, each with appropriate units and standard names.
"""

import os
import numpy as np
from netCDF4 import Dataset


def process_netcdf(input_filepath, output_filepath):
    """
    Process wind and pressure data and save it to a new NetCDF file.

    Parameters:
        input_filepath (str): Path to the input netCDF file.
        output_filepath (str): Path to save the output netCDF file.
    """
    try:
        # Open the input netCDF file in read mode
        with Dataset(input_filepath, 'r') as ncid0:
            # Print available variables to detect the time variable dynamically
            print("Variables in the input file:")
            print(ncid0.variables.keys())

            # Detect the time variable name dynamically
            time_var_name = 'time'  # Default name
            if 'time' not in ncid0.variables:
                for var in ncid0.variables:
                    if 'time' in var.lower():
                        time_var_name = var
                        break
                else:
                    raise ValueError("No time variable found in the input file.")

            print(f"Using time variable: {time_var_name}")

            # Read and adjust longitude values to [-180, 180]
            lon = ncid0.variables['longitude'][:]
            lon[lon >= 180] -= 360  # Adjust longitudes
            lon = np.sort(lon)

            # Read latitude values
            lat = ncid0.variables['latitude'][:]

            # Read time values and their units
            time = ncid0.variables[time_var_name][:]
            time_units = ncid0.variables[time_var_name].units  # Fetch original time units

            # Read pressure and wind components
            Pmsl = ncid0.variables['msl'][:, :, :]  # Mean sea level pressure (Pa)
            U10 = ncid0.variables['u10'][:, :, :]  # Eastward wind component (m/s)
            V10 = ncid0.variables['v10'][:, :, :]  # Northward wind component (m/s)

            # Flip the data along the latitude axis
            Pmsl_flipped = np.flip(Pmsl, axis=1)
            U10_flipped = np.flip(U10, axis=1)
            V10_flipped = np.flip(V10, axis=1)

            # Create the output netCDF file
            with Dataset(output_filepath, 'w', format='NETCDF4') as ncid:
                # Define dimensions
                ncid.createDimension('time', len(time))
                ncid.createDimension('y', len(lat))
                ncid.createDimension('x', len(lon))

                # Define longitude
                var_lon = ncid.createVariable('X', 'f8', ('x',))
                var_lon.units = 'degrees_east'
                var_lon.standard_name = 'longitude'
                var_lon[:] = lon

                # Define latitude
                var_lat = ncid.createVariable('Y', 'f8', ('y',))
                var_lat.units = 'degrees_north'
                var_lat.standard_name = 'latitude'
                var_lat[:] = lat

                # Define time and preserve original units
                var_time = ncid.createVariable('time', 'f8', ('time',))
                var_time.units = time_units  # Use original time units
                var_time.calendar = 'gregorian'
                var_time.standard_name = 'time'
                var_time.long_name = 'Time'
                var_time[:] = time

                # Define pressure
                var_p = ncid.createVariable('air_pressure_fixed_height', 'f4', ('time', 'y', 'x'))
                var_p.units = 'Pa'
                var_p.long_name = 'Mean sea level pressure'
                var_p.standard_name = 'air_pressure'
                var_p[:, :, :] = Pmsl_flipped

                # Define U wind component
                var_u10 = ncid.createVariable('eastward_wind', 'f4', ('time', 'y', 'x'))
                var_u10.units = 'm s**-1'
                var_u10.long_name = '10 metre U wind component'
                var_u10.standard_name = 'eastward_wind'
                var_u10[:, :, :] = U10_flipped

                # Define V wind component
                var_v10 = ncid.createVariable('northward_wind', 'f4', ('time', 'y', 'x'))
                var_v10.units = 'm s**-1'
                var_v10.long_name = '10 metre V wind component'
                var_v10.standard_name = 'northward_wind'
                var_v10[:, :, :] = V10_flipped

                # Print confirmation message
                print(f"Processed data saved to: {output_filepath}")

    except Exception as e:
        print(f"An error occurred: {e}")


def main():
    """
    Main function to set paths and process the netCDF file.
    """
    # USER-DEFINED INPUT: Set input/output paths
    dinput = input("Enter the input directory path: ").strip()  # Input directory
    doutput = input("Enter the output directory path: ").strip()  # Output directory

    filename = 'data_stream-oper.nc'  # Input file name
    input_filepath = os.path.join(dinput, filename)
    output_filepath = os.path.join(doutput, filename.replace('.nc', '_dfm.nc'))

    # Verify paths before processing
    if not os.path.isfile(input_filepath):
        print(f"Error: Input file '{input_filepath}' does not exist.")
        return
    if not os.path.isdir(doutput):
        print(f"Error: Output directory '{doutput}' does not exist.")
        return

    # Process the netCDF file
    process_netcdf(input_filepath, output_filepath)


if __name__ == "__main__":
    main()