mantaray 0.2.0

Ray-tracing solver for ocean surface gravity waves that integrates the wave ray equations over spatially varying currents and bathymetry.
Documentation
''' Plot single ray output on top of depth map

A script that lists the .nc and .txt files in the current directory and asks the
user to select two. The netcdf file is the depth map which will be the
background for the plotting. Then the text file has the output from ray tracing
a single wave, and that data is plotted on top of the depth map. Then the plot
is displayed.

This script currently only works and was tested for space separated files, where
each row is the state of the wave at the given timestep. The columns  of the
file are t, x, y, kx, ky.

Generated by phind ai
'''

import glob
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

def get_input_files(file_extension):
    """
    Lists all files in the current directory with the specified extension and asks the user to select one.
    """
    files = glob.glob(f"*.{file_extension}")
    
    for i, filename in enumerate(files):
        print(f"{i+1}: {filename}")
        
    file_num = int(input(f"Enter the number of the file you want to use: ")) - 1
    
    if file_num >= 0 and file_num < len(files):
        print(f"You selected {files[file_num]}")
        return files[file_num]
    else:
        print("Invalid choice. Exiting.")
        exit(1)
        

def read_nc_file(file_path):
    """
    Reads the file and extracts the depth, x and y values.
    """
    try:
        ds = xr.open_dataset(file_path)
        return ds.depth.values, ds.x.values, ds.y.values
    except Exception as e:
        print(f"Error reading file {file_path}: {str(e)}")
        

def read_txt_file(file_path):
    """
    Reads the file and extracts the t, x, y, kx and ky values.
    """
    try:
        with open(file_path, 'r') as file:
            data = np.genfromtxt(file_path, skip_header=1).T
            return data[0], data[1], data[2], data[3], data[4]
    except Exception as e:
        print(f"Error reading file {file_path}: {str(e)}")


def create_plot(depth_arr, x_arr, y_arr, x, y):
    """
    Creates the plot.
    """
    fig, ax = plt.subplots()

    # Set the background
    im = ax.imshow(depth_arr, extent=[x_arr[0], x_arr[-1], y_arr[0], y_arr[-1]], origin='lower')
    
    # Color bar
    cbar = ax.figure.colorbar(im, ax=ax)
    cbar.ax.set_ylabel("Depth", rotation=-90, va="bottom")
    
    # Set number of ticks
    num_ticks = 10
    x_ticks = np.linspace(x_arr[0], x_arr[-1], num_ticks, dtype=int)
    y_ticks = np.linspace(y_arr[0], y_arr[-1], num_ticks, dtype=int)
    
    # Display ticks
    ax.set_xticks(x_ticks)
    ax.set_yticks(y_ticks)
    ax.set_xticklabels(x_ticks)
    ax.set_yticklabels(y_ticks)
    
    # Set axes labels and title
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_title("Position of wave at x(t), y(t)")
    
    # Plot the resulting line
    ax.plot(x, y, linewidth=2, color='k')
    
    plt.show()


if __name__ == "__main__":
    nc_file = get_input_files('nc')
    depth_arr, x_arr, y_arr = read_nc_file(nc_file)
    
    txt_file = get_input_files('txt')
    t, x, y, kx, ky = read_txt_file(txt_file)
    
    create_plot(depth_arr, x_arr, y_arr, x, y)