Source code for core.interpolation

"""
Interpolation utilities for geophysical data processing.
"""
import numpy as np
from scipy.interpolate import griddata
from typing import Tuple, List, Optional, Union


[docs]def interpolate_to_profile(data: np.ndarray, X_grid: np.ndarray, Y_grid: np.ndarray, X_pro: np.ndarray, Y_pro: np.ndarray, method: str = 'linear') -> np.ndarray: """ Interpolate 2D data onto a profile line Args: data: 2D array of values to interpolate X_grid: X coordinates of original grid (meshgrid) Y_grid: Y coordinates of original grid (meshgrid) X_pro: X coordinates of profile points Y_pro: Y coordinates of profile points method: Interpolation method ('linear' or 'nearest') Returns: Interpolated values along profile """ X_new = X_grid.ravel() Y_new = Y_grid.ravel() return griddata((X_new, Y_new), np.array(data).ravel(), (np.array(X_pro).ravel(), np.array(Y_pro).ravel()), method=method)
[docs]def setup_profile_coordinates(point1: List[int], point2: List[int], surface_data: np.ndarray, origin_x: float = 0.0, origin_y: float = 0.0, pixel_width: float = 1.0, pixel_height: float = -1.0, num_points: int = 200) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Set up profile coordinates based on surface elevation data between two points Args: point1: Starting point indices [col, row] point2: Ending point indices [col, row] surface_data: 2D array of surface elevation data origin_x: X coordinate of origin origin_y: Y coordinate of origin pixel_width: Width of each pixel pixel_height: Height of each pixel (negative for top-down) num_points: Number of points along profile Returns: X_pro: X coordinates along profile Y_pro: Y coordinates along profile L_profile: Distances along profile XX: X coordinate grid YY: Y coordinate grid """ # Create coordinate grids x = origin_x + pixel_width * np.arange(surface_data.shape[1]) y = origin_y + pixel_height * np.arange(surface_data.shape[0]) XX, YY = np.meshgrid(x, y) # Handle no-data values surface_data = surface_data.copy() surface_data[surface_data == 0] = np.nan # Calculate start and end positions P1_pos = np.array([x[point1[0]], y[point1[1]]]) P2_pos = np.array([x[point2[0]], y[point2[1]]]) # Calculate total distance dis = np.sqrt(np.sum((P1_pos - P2_pos)**2)) # Generate profile coordinates X_pro = (x[point1[0]] - x[point2[0]])/dis * np.linspace(0, dis, num_points)[:-1] + x[point2[0]] Y_pro = (y[point1[1]] - y[point2[1]])/dis * np.linspace(0, dis, num_points)[:-1] + y[point2[1]] # Calculate profile distances L_profile = np.sqrt((X_pro - X_pro[0])**2 + (Y_pro - Y_pro[0])**2) return X_pro, Y_pro, L_profile, XX, YY
[docs]def interpolate_structure_to_profile(structure_data: List[np.ndarray], X_grid: np.ndarray, Y_grid: np.ndarray, X_pro: np.ndarray, Y_pro: np.ndarray) -> np.ndarray: """ Interpolate multiple structure layers onto profile Args: structure_data: List of 2D arrays for each layer X_grid: X coordinates of original grid Y_grid: Y coordinates of original grid X_pro: X coordinates of profile points Y_pro: Y coordinates of profile points Returns: Array of interpolated values with shape (n_layers, n_points) """ structure = [] for layer in structure_data: interpolated = interpolate_to_profile(layer, X_grid, Y_grid, X_pro, Y_pro) structure.append(interpolated) return np.array(structure)
[docs]def prepare_2D_profile_data(data: np.ndarray, XX: np.ndarray, YY: np.ndarray, X_pro: np.ndarray, Y_pro: np.ndarray) -> np.ndarray: """ Interpolate multiple 2D gridded data layers onto a profile line. Args: data: 3D array of gridded data (n_layers, ny, nx) XX, YY: Coordinate grids from meshgrid X_pro, Y_pro: Profile line coordinates Returns: Interpolated values along profile (n_layers, n_profile_points) """ n_layers = data.shape[0] profile_values = [] X_new = XX.ravel() Y_new = YY.ravel() for i in range(n_layers): layer_values = griddata((X_new, Y_new), data[i].ravel(), (X_pro.ravel(), Y_pro.ravel()), method='linear') profile_values.append(layer_values) return np.array(profile_values)
[docs]def interpolate_to_mesh(property_values: np.ndarray, profile_distance: np.ndarray, depth_values: np.ndarray, mesh_x: np.ndarray, mesh_y: np.ndarray, mesh_markers: np.ndarray, ID, layer_markers: list = [3, 0, 2]) -> np.ndarray: """ Interpolate property values from profile to mesh with layer-specific handling. Args: property_values: Property values array (n_points) profile_distance: Distance along profile (n_points) depth_values: Depth values array (n_layers, n_points) mesh_x: X coordinates of mesh cells mesh_y: Y coordinates of mesh cells mesh_markers: Markers indicating different layers in mesh layer_markers: List of marker values for each layer Returns: Interpolated values for mesh cells """ # Initialize output array result = np.zeros_like(mesh_markers, dtype=float) # print(profile_distance.shape) # print(depth_values.shape) # print(property_values.shape) L_profile_new = np.repeat(profile_distance.reshape(1,-1),property_values.shape[0],axis=0) Depth = depth_values[:14] maxele = 0 # set 0 here for marker in layer_markers: # For each layer marker, interpolate property values to mesh grid # Note: ID is used to identify which layer to interpolate for # Interpolate property values to mesh grid for each layer grid_z1 = griddata((L_profile_new[ID==marker].ravel(),Depth[ID==marker].ravel()- maxele), property_values[ID==marker].ravel(), (mesh_x[mesh_markers==marker], mesh_y[mesh_markers==marker]), method='linear') temp_ID = np.isnan(grid_z1) grid_z2 = griddata((L_profile_new[ID==marker].ravel(),Depth[ID==marker].ravel()- maxele), property_values[ID==marker].ravel(), (mesh_x[mesh_markers==marker], mesh_y[mesh_markers==marker]), method='nearest') grid_z1[temp_ID] = grid_z2[temp_ID] result[mesh_markers==marker] = grid_z1.copy() # # Interpolate property values to mesh grid for each layer # grid_z1 = griddata((L_profile_new[ID==0].ravel(),Depth[ID==0].ravel()- maxele), property_values[ID==0].ravel(), (mesh_x[mesh_markers==0], mesh_y[mesh_markers==0]), method='linear') # temp_ID = np.isnan(grid_z1) # grid_z2 = griddata((L_profile_new[ID==0].ravel(),Depth[ID==0].ravel()- maxele), property_values[ID==0].ravel(), (mesh_x[mesh_markers==0], mesh_y[mesh_markers==0]), method='nearest') # grid_z1[temp_ID] = grid_z2[temp_ID] # result[mesh_markers==0] = grid_z1.copy() #result = griddata((L_profile_new.ravel(),depth_values[:14].ravel()), property_values.ravel(), (mesh_x, mesh_y), method='nearest') return result
[docs]class ProfileInterpolator: """Class for handling interpolation of data to/from profiles."""
[docs] def __init__(self, point1: List[int], point2: List[int], surface_data: np.ndarray, origin_x: float = 0.0, origin_y: float = 0.0, pixel_width: float = 1.0, pixel_height: float = -1.0, num_points: int = 200): """ Initialize profile interpolator with reference points and surface data. Args: point1: Starting point indices [col, row] point2: Ending point indices [col, row] surface_data: 2D array of surface elevation data origin_x, origin_y: Coordinates of origin pixel_width, pixel_height: Pixel dimensions num_points: Number of points along profile """ self.point1 = point1 self.point2 = point2 self.surface_data = surface_data self.origin_x = origin_x self.origin_y = origin_y self.pixel_width = pixel_width self.pixel_height = pixel_height self.num_points = num_points # Set up profile coordinates self.X_pro, self.Y_pro, self.L_profile, self.XX, self.YY = setup_profile_coordinates( point1, point2, surface_data, origin_x, origin_y, pixel_width, pixel_height, num_points ) # Get surface profile self.surface_profile = interpolate_to_profile( surface_data, self.XX, self.YY, self.X_pro, self.Y_pro )
[docs] def interpolate_layer_data(self, layer_data: List[np.ndarray]) -> np.ndarray: """ Interpolate multiple layer data to profile. Args: layer_data: List of 2D arrays for each layer Returns: Array of interpolated values (n_layers, n_profile_points) """ return interpolate_structure_to_profile( layer_data, self.XX, self.YY, self.X_pro, self.Y_pro )
[docs] def interpolate_3d_data(self, data: np.ndarray) -> np.ndarray: """ Interpolate 3D data (n_layers, ny, nx) to profile. Args: data: 3D array of values Returns: Array of interpolated values (n_layers, n_profile_points) """ return prepare_2D_profile_data( data, self.XX, self.YY, self.X_pro, self.Y_pro )
[docs] def interpolate_to_mesh(self, property_values: np.ndarray, depth_values: np.ndarray, mesh_x: np.ndarray, mesh_y: np.ndarray, mesh_markers: np.ndarray, ID: np.ndarray, layer_markers: list = [3, 0, 2]) -> np.ndarray: """ Interpolate property values from profile to mesh with layer-specific handling. Args: property_values: Property values array (n_points or n_layers, n_points) depth_values: Depth values array (n_layers, n_points) mesh_x, mesh_y: Coordinates of mesh cells mesh_markers: Markers indicating different layers in mesh layer_markers: List of marker values for each layer Returns: Interpolated values for mesh cells """ return interpolate_to_mesh( property_values, self.L_profile, depth_values, mesh_x, mesh_y, mesh_markers, ID,layer_markers )
[docs]def create_surface_lines(L_profile: np.ndarray, structure: np.ndarray, top_idx: int = 0, mid_idx: int = 4, bot_idx: int = 12) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ Create surface and boundary lines from structure data Args: L_profile: Distance along profile structure: Interpolated structure data top_idx: Index for top surface mid_idx: Index for middle boundary bot_idx: Index for bottom boundary Returns: surface: Surface coordinates line1: First boundary coordinates line2: Second boundary coordinates """ # Extract and reshape structure layers S1 = structure[top_idx,:].reshape(-1,1) S2 = structure[mid_idx,:].reshape(-1,1) S3 = structure[bot_idx,:].reshape(-1,1) # Create coordinate arrays surface = np.hstack((L_profile.reshape(-1,1), S1)) line1 = np.hstack((L_profile.reshape(-1,1), S2)) line2 = np.hstack((L_profile.reshape(-1,1), S3)) # Normalize by maximum elevation #maxele = np.nanmax(surface[:,1]) surface[:,1] = surface[:,1] #- maxele line1[:,1] = line1[:,1] #- maxele line2[:,1] = line2[:,1] #- maxele return surface, line1, line2