Add functions to display seismic data in PyVista

[AlexanderJuestel]

Is your feature request related to a problem? Please describe.
Add functions to convert seismic data loaded with segysak to PyVista meshes

Describe the solution you'd like
Functionality wrapped in functions from https://gemgis.readthedocs.io/en/latest/getting_started/tutorial/56_Displaying_Seismic_Data_in_Pyvista.html

[AlexanderJuestel]

import xarray
from typing import Union


def clip_seismic_data(seismic_data: xarray.core.dataset.Dataset,
                     cdp_start: Union[int, type(None)] = None, 
                     cdp_end: Union[int, type(None)] = None) -> pd.DataFrame:
    
    # Checking that the seismic data is provided as xarray Dataset
    if not isinstance(seismic_data, xarray.core.dataset.Dataset):
        raise TypeError('The seismic data must be provided as xarray Dataset loaded ideally with segysak and its segy_loader')
    
    # Checking that cdp_start ist of type int or None
    if not isinstance(cdp_start, (int, type(None))):
        raise TypeError('The start CDP must be provided as int')
        
    # Checking that cdp_end ist of type int or None
    if not isinstance(cdp_end, (int, type(None))):
        raise TypeError('The end CDP must be provided as int')
        
    # Converting xarray DataSet to DataFrame
    df_seismic_data = seismic_data.to_dataframe()
    
    # Getting the start CDP number if it is not provided
    if not cdp_start:
        cdp_start = int(df_seismic_data.index[0][0])
    
    # Getting the end CDP number if it is not provided
    if not cdp_end:
        cdp_end = int(df_seismic_data.index[-1][0])
        
    # Select seismic data
    df_seismic_data_selection = df_seismic_data.loc[cdp_start:cdp_end]
    
    return df_seismic_data_selection
    
    

def seismic_to_array(seismic_data: xarray.core.dataset.Dataset,
                     cdp_start: Union[int, type(None)] = None, 
                     cdp_end: Union[int, type(None)] = None,
                     max_depth: Union[int, float, type(None)] = None) -> np.ndarray:
    
    # Checking that the seismic data is provided as xarray Dataset
    if not isinstance(seismic_data, xarray.core.dataset.Dataset):
        raise TypeError('The seismic data must be provided as xarray Dataset loaded ideally with segysak and its segy_loader')
    
    # Checking that cdp_start ist of type int or None
    if not isinstance(cdp_start, (int, type(None))):
        raise TypeError('The start CDP must be provided as int')
        
    # Checking that cdp_end ist of type int or None
    if not isinstance(cdp_end, (int, type(None))):
        raise TypeError('The end CDP must be provided as int')
        
        
    # Checking that the max_depth is of type int or float
    if not isinstance(max_depth, (int, float, type(None))):
        raise TypeError('The maximum depth in m or TWT must be provided as int or float')
        
    # Converting xarray DataSet to DataFrame
    df_seismic_data = seismic_data.to_dataframe()
        
    # Getting the start CDP number if it is not provided
    if not cdp_start:
        cdp_start = int(df_seismic_data.index[0][0])
    
    # Getting the end CDP number if it is not provided
    if not cdp_end:
        cdp_end = int(df_seismic_data.index[-1][0])
    
    # Clipping the seismic data
    df_seismic_data_selection = clip_seismic_data(seismic_data=seismic_data,
                                                  cdp_start=cdp_start,
                                                  cdp_end=cdp_end)
    
    # Getting the number of rows per CDP and number of cdps
    len_cdp = int(len(df_seismic_data_selection.loc[cdp_start]))
    num_cdp = int(len(df_seismic_data_selection)/len_cdp)
    
    # Getting the seismic data
    df_seismic_data_values = df_seismic_data_selection['data'].values
    
    # Reshaping the array
    df_seismic_data_values_reshaped = df_seismic_data_values.reshape(num_cdp, 
                                                                     len_cdp)
    
    # Getting the max_depth if it is not provided
    if not max_depth:
        max_depth = df_seismic_data_selection.loc[cdp_start].index[-1]
                    
    # Getting the number of samples based on max_depth
    num_indices = int((len_cdp-1)/(df_seismic_data_selection.loc[cdp_start].index.max()-df_seismic_data_selection.loc[cdp_start].index.min())*max_depth+1)
    
    # Selecting samples
    df_seismic_data_values_reshaped_selected = df_seismic_data_values_reshaped[:,:num_indices]
    
    
    return df_seismic_data_values_reshaped_selected

import pyproj

def seismic_to_mesh(seismic_data: xarray.core.dataset.Dataset,
                    cdp_start: Union[int, type(None)] = None, 
                    cdp_end: Union[int, type(None)] = None,
                    max_depth: Union[int, float] = None, 
                    sampling_rate: Union[int, type(None)] = None,
                    shift: int = 0,
                    source_crs: Union[str, pyproj.crs.crs.CRS] = None,
                    target_crs: Union[str, pyproj.crs.crs.CRS] = None, 
                    cdp_coords = None,
                    )-> pv.core.pointset.StructuredGrid:
    
    # Checking that the sampling_rate is provided
    if not isinstance(sampling_rate, (int, type(None))):
        raise TypeError('The sampling rate must be provided as integer')
        
    # Checking that the shift is of type int
    if not isinstance(shift, int):
        raise TypeError('The shift must be provided as integer')
        
     # Checking that the target_crs is of type string
    if not isinstance(source_crs, (str, type(None), pyproj.crs.crs.CRS)):
        raise TypeError('source_crs must be of type string or a pyproj object')
        
     # Checking that the target_crs is of type string
    if not isinstance(target_crs, (str, type(None), pyproj.crs.crs.CRS)):
        raise TypeError('target_crs must be of type string or a pyproj object')
                     
    # Getting the sampling rate if it is not provided
    if not sampling_rate:
        sampling_rate = seismic_data.to_dataframe().reset_index()['twt'][1] - seismic_data.to_dataframe().reset_index()['twt'][0]
    
    # Getting the seismic data as array
    seismic_data_array = seismic_to_array(seismic_data=seismic_data, 
                                          cdp_start=cdp_start,
                                          cdp_end=cdp_end,
                                          max_depth=max_depth)
    
    # Getting the number of traces and samples (columns and rows)
    ntraces, nsamples = seismic_data_array.shape 
    
    # Clipping the seismic data
    seismic_data = clip_seismic_data(seismic_data=seismic_data,
                                     cdp_start=cdp_start,
                                     cdp_end=cdp_end)
    
    # Getting the CDP coordinates
    try:
        cdp_coordinates = seismic_data[['cdp_x', 'cdp_y']].drop_duplicates().values
        cdp_x = cdp_coordinates[:,0]
        cdp_y = cdp_coordinates[:,1]
        
    # TODO: Implement that the seismic section can also be clipped. Currently, only the entire section can be converted to a mesh when the cdp coordinates are not present in the xarray DataSet
    except KeyError:
        cdp_x = cdp_coords[:,0]
        cdp_y = cdp_coords[:,1]
        cdp_coordinates = cdp_coords
    # Converting the coordinates
    if target_crs and source_crs:
        gdf_coords = gpd.GeoDataFrame(geometry=gpd.points_from_xy(x=cdp_x,
                                                                  y=cdp_y),
                                      crs=source_crs).to_crs(target_crs)   
        
        cdp_coordinates= np.array([gdf_coords.geometry.x.values, 
                                   gdf_coords.geometry.y.values]).T
    
    # Creating the path
    seismic_path = np.c_[cdp_coordinates, np.zeros(len(cdp_coordinates))]
    
    # Creating the points
    seismic_points = np.repeat(seismic_path, nsamples, axis=0)
    
    # Repeating the Z locations
    tp = np.arange(0, sampling_rate * nsamples, sampling_rate)
    tp = seismic_path[:, 2][:, None] - tp + shift
    seismic_points[:, -1] = tp.ravel()
    
    # Creating StrcuturedGrid
    grid = pv.StructuredGrid()
    grid.points = seismic_points
    grid.dimensions = nsamples, ntraces, 1
    
    # Adding the amplitude values 
    grid["values"] = seismic_data_array.ravel(order="C")
    
    return grid