Adding WW3 mesh writer

[SorooshMani-NOAA]

Based on the work done by @yunfangsun, add the writer for WW3 mesh file format.

[SorooshMani-NOAA]

From @yunfangsun:

[...] scripts for preparing the WWIII mesh input file.

1. break the quads into triangular:

from __future__ import annotations
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path

import f90nml
import numpy as np
from pyschism.mesh.base import Elements
from pyschism.mesh.base import Gr3
from pyschism.mesh.gridgr3 import Gr3Field
from pyschism.param.param import Param


def break_quads(pyschism_mesh: Gr3) -> Gr3 | Gr3Field:
    # Create new Elements and set it for the Gr3.elements
    quads = pyschism_mesh.quads
    if len(quads) == 0:
        new_mesh = deepcopy(pyschism_mesh)

    else:
        tmp = quads[:,2:]
        tmp = np.insert(tmp, 0, quads[:, 0], axis=1)
        broken = np.vstack((quads[:, :3], tmp))
        trias = pyschism_mesh.triangles
        final_trias = np.vstack((trias, broken))
        # NOTE: Node IDs and indexs are the same as before
        elements = {
            idx+1: list(map(pyschism_mesh.nodes.get_id_by_index, tri))
            for idx, tri in enumerate(final_trias)
        }

        new_mesh = deepcopy(pyschism_mesh)
        new_mesh.elements = Elements(pyschism_mesh.nodes, elements)


    return new_mesh
############################
mesh_obj4 = Gr3.open('new_mesh_1.gr3', crs=4326)
triangular_mesh4 = break_quads(mesh_obj4)

2. I have converted the following MATLAB function to write the WWIII mesh (https://github.com/NOAA-EMC/WW3-tools/blob/develop/matlab_tools/WW3_mesh_write.m)

import numpy as np
import matplotlib.pyplot as plt

def WW3_mesh_write_from_object(mesh, filename, plott=False):
    # Extract the node (vertex) data from the mesh object
    x = mesh.x
    y = mesh.y
    z = mesh.values  # Assuming 'values' stores the depth. Adjust this if it's different.
   
    # Extract the element (triangle) data from the mesh object
    triangles = mesh.triangles  # Assuming triangles provides (n x 3) array of vertex indices

    # It's not clear how open boundary nodes are represented in `triangular_mesh4`
    # For now, I'll omit that from the function. You may need to add that if it's provided in the mesh object.

    with open(filename, 'w') as f:
        f.write('$MeshFormat\n')
        f.write('2 0 8\n')
        f.write('$EndMeshFormat\n')
        f.write('$Nodes\n')
        f.write(f'{len(x)}\n')

        for i in range(len(x)):
            f.write(f'{i+1} {x[i]:.5f} {y[i]:.5f} {z[i]:.5f}\n')
       
        f.write('$EndNodes\n')
        f.write('$Elements\n')
        f.write(f'{len(triangles)}\n')

        for i, tri in enumerate(triangles):
            f.write(f'{i+1} 2 3 0 {i+1} 0 {tri[0]+1} {tri[1]+1} {tri[2]+1}\n')

        f.write('$EndElements\n')

    if plott:
        plt.figure(figsize=(11, 8))
        avg_z_values = np.mean(z[triangles], axis=1)

        plt.tripcolor(x, y, triangles, facecolors=avg_z_values, edgecolors='k')

        plt.colorbar(label='Depth (m)')
        plt.xlabel('Longitude')
        plt.ylabel('Latitude')
        plt.title('WW3 Grid - Depth [m]')
        plt.tight_layout()
        plt.savefig(filename + '.png', dpi=300)
        plt.show()
#############################
# Example Usage:
WW3_mesh_write_from_object(triangular_mesh4, 'output_filename.txt', plott=True)

3. The WWIII read function (https://github.com/NOAA-EMC/WW3-tools/blob/develop/matlab_tools/WW3_mesh_read.m) is also converted to test if the mesh is correctly written.

import numpy as np
import matplotlib.pyplot as plt

def WW3_mesh_read(filename, plott=False):
    # Read file content
    with open(filename, 'r') as f:
        content = f.readlines()
   
    # Extract node and element count
    N_n = int(content[4].split()[0])
    N_e = int(content[7 + N_n].split()[0])

    # Extract node and element data
    nodes = np.array([list(map(float, line.split()[1:4])) for line in content[5:5+N_n]])
#    TRI = np.array([list(map(int, line.split()[6:9])) for line in content[8+N_n:8+N_n+N_e]])
    TRI = np.array([list(map(int, line.split()[6:9])) for line in content[8+N_n:8+N_n+N_e]]) - 1

    x = nodes[:, 0]
    y = nodes[:, 1]
    z = nodes[:, 2]
    tri = TRI

    # Plot
    if plott:
        fig, ax = plt.subplots(figsize=(11, 8))
        vmin_value = 0  # example minimum value
        vmax_value = 20   # example maximum value

       
       
        avg_z_values = np.mean(z[tri], axis=1)
        ax.tripcolor(x, y, tri, facecolors=avg_z_values, edgecolors='k', cmap='jet')
        mappable = ax.tripcolor(x, y, tri, facecolors=avg_z_values, edgecolors='k', cmap='jet', vmin=vmin_value, vmax=vmax_value)

       
        ax_pos = ax.get_position()

    # Define colorbar axes, [left, bottom, width, height]
#        cbar_ax = fig.add_axes([ax_pos.x1 + 0.02, ax_pos.y0, 0.02, ax_pos.height])
       
        ax.set_position([ax_pos.x0, ax_pos.y0, ax_pos.width * 0.9, ax_pos.height])

    # Define colorbar axes, [left, bottom, width, height]
        cbar_ax = fig.add_axes([ax_pos.x0 + ax_pos.width * 0.9 + 0.02, ax_pos.y0, 0.02, ax_pos.height])

        cbar = plt.colorbar(mappable=mappable, cax=cbar_ax, orientation='vertical')

#        cbar = plt.colorbar(mappable=mappable, ax=ax, orientation='vertical')        
       
       
#        ax.tripcolor(x, y, tri, facecolors=z, edgecolors='k', cmap='jet')
#        cbar = plt.colorbar(ax=ax, orientation='vertical')
        cbar.set_label('Depth (m)', fontsize=14)
       
        ax.set_xlabel('Longitude °')
        ax.set_ylabel('Latitude °')
        ax.set_title('Depth (m)', fontsize=15)

#        ax.set_xlim([min(x) - 0.1 * (max(x) - min(x)), max(x) + 0.1 * (max(x) - min(x))])
#        ax.set_ylim([min(y) - 0.1 * (max(y) - min(y)), max(y) + 0.1 * (max(y) - min(y))])

 
        lon_min, lon_max = -82.8, -82.65  # example values
        lat_min, lat_max = 27.6, 27.7  # example values
        ax.set_xlim(lon_min, lon_max)
        ax.set_ylim(lat_min, lat_max)
       
        average_latitude = np.mean([lat_min, lat_max])
        aspect_ratio = 1 / np.cos(average_latitude * np.pi / 180)
       
        ax.set_aspect(aspect_ratio)      


        ax.grid(False)
        plt.tight_layout()
        plt.savefig(filename + '.png', dpi=300)
        plt.show()

    return tri, x, y, z
################################
# Example Usage:
tri, x, y, z = WW3_mesh_read('output_filename.txt', plott=True)

The function "WW3_mesh_write_from_object(mesh, filename, plott=False) " and "WW3_mesh_read" could be added to OCSmesh to write the WWIII mesh input file.

[yunfangsun]

@SorooshMani-NOAA Hi Soroosh,

I have made some modifications for WW3_mesh_write_from_object, could you please make a correction for it:

import numpy as np
import matplotlib.pyplot as plt

def WW3_mesh_write_from_object(mesh, filename, plott=False):
    # Extract the node (vertex) data from the mesh object
    x = mesh.x
    y = mesh.y
    z = mesh.values  # Assuming 'values' stores the depth. Adjust this if it's different.
    
    # Extract the element (triangle) data from the mesh object
    triangles = mesh.triangles   # Assuming triangles provides (n x 3) array of vertex indices

    # It's not clear how open boundary nodes are represented in `triangular_mesh4`
    # For now, I'll omit that from the function. You may need to add that if it's provided in the mesh object.

    OB_ID=mesh.boundaries.open.index_id
    OB_ID_flattened = np.unique(np.concatenate(OB_ID.values))
    
    with open(filename, 'w') as f:
        f.write('$MeshFormat\n')
        f.write('2 0 8\n')
        f.write('$EndMeshFormat\n')
        f.write('$Nodes\n')
        f.write(f'{len(x)}\n')

        for i in range(len(x)):
            f.write(f'{i+1} {x[i]:.5f} {y[i]:.5f} {z[i]:.5f}\n')
        
        f.write('$EndNodes\n')
        f.write('$Elements\n')
        f.write(f'{len(triangles)+ len(OB_ID_flattened)}\n')

#        for i, tri in enumerate(triangles):
#            f.write(f'{i+1} 2 3 0 {i+1} 0 {tri[0]+1} {tri[1]+1} {tri[2]+1}\n')



        m = 0
        for ob_id in OB_ID_flattened:
            m += 1
            f.write(f'{m} 15 2 0 0 {ob_id}\n')

        for i, elem in enumerate(triangles, start=m+1):
            f.write(f'{i} 2 3 0 {i-m} 0 {elem[0]+1} {elem[1]+1} {elem[2]+1}\n')
        
        f.write('$EndElements\n')       
        
        
    if plott:
        plt.figure(figsize=(11, 8))
        avg_z_values = np.mean(z[triangles], axis=1)

        plt.tripcolor(x, y, triangles, facecolors=avg_z_values, edgecolors='k')

        plt.colorbar(label='Depth (m)')
        plt.xlabel('Longitude')
        plt.ylabel('Latitude')
        plt.title('WW3 Grid - Depth [m]')
        plt.tight_layout()
        plt.savefig(filename + '.png', dpi=300)
        plt.show()

# Example Usage:
WW3_mesh_write_from_object(hgrid, 'atlantic1.msh', plott=True)

Thank you very much!

[pvelissariou1]

@yunfangsun , @SorooshMani-NOAA I have the complete matlab script: adcirc2ww3_mesh.m on hera at: /scratch2/STI/coastal_temp/save/Panagiotis.Velissariou/data_FlSaLaMa.
This generates the WW3 "msh" file using Ali's WW3_mesh_write.m function and the readfort14.m function from OceanMesh2D. It might be helpful to both of you to check for missing functionalities in WW3_mesh_write_from_object. It works for me for all the meshes I tested.