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 many 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
many waves, 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. The files are delimited by a line with "END".

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 for each ray.
    """
    try:
        with open(file_path, 'r') as file:
            lines = file.readlines()

        rays = []
        ray_data = []
        start_reading = False
        for line in lines:
            if "t x y kx ky" in line:  # The header of a new block of data
                start_reading = True
                continue  # skip the header line
            if start_reading and "END" in line:  # The end of current block of data
                if ray_data:  # If there is data from the current block convert to numpy array and store
                    ray = np.array(ray_data)
                    rays.append(ray)
                    ray_data = []  # Reset for new block
                start_reading = False
                continue 
            if start_reading:
                 values = list(map(float, line.strip().split()))
                 ray_data.append(values)

        # After going through all the lines, there may be a last block still not processed
        if ray_data:
            ray = np.array(ray_data)
            rays.append(ray)

        return rays
    except Exception as e:
        print(f"Error reading file {file_path}: {str(e)}")


def create_plot(depth_arr, x_arr, y_arr, rays):
    """
    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 [m]", 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 [m]")
    ax.set_ylabel("y [m]")
    ax.set_title("ray paths")
    
    # Plot each ray
    for ray in rays:
        t, x, y, kx, ky = ray.T  # transpose the ray data to get the rows as t, x, y, kx, ky
        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')
    rays = read_txt_file(txt_file)
    create_plot(depth_arr, x_arr, y_arr, rays)