import json import numpy as np import pandas as pd import torch import dgl import gradio as gr import plotly.graph_objects as go import plotly.express as px from math import radians, cos, sin from sklearn.neighbors import KDTree import torch.nn as nn # Define normalization parameters NORM_PARAMS = { 'v_min': 0.0508155134766646, 'v_max': 23.99703788008771, 'x_min': -97.40959, 'x_max': -96.55169890366584, 'y_min': 32.587689999999995, 'y_max': 33.067024421728206 } # Define the model class class WindGNN(nn.Module): def __init__(self, in_feats=5, hidden_size=128, out_feats=3, num_layers=3): super(WindGNN, self).__init__() self.layers = nn.ModuleList() self.input_proj = nn.Linear(in_feats, hidden_size) for _ in range(num_layers): self.layers.append(dgl.nn.GraphConv(hidden_size, hidden_size)) self.velocity_head = nn.Sequential( nn.Linear(hidden_size, hidden_size // 4), nn.LayerNorm(hidden_size // 4), nn.ReLU(), nn.Dropout(0.2), nn.Linear(hidden_size // 4, 1), nn.Sigmoid() ) self.direction_head = nn.Sequential( nn.Linear(hidden_size, hidden_size // 4), nn.LayerNorm(hidden_size // 4), nn.ReLU(), nn.Dropout(0.2), nn.Linear(hidden_size // 4, 2) ) self.layer_norms = nn.ModuleList([ nn.LayerNorm(hidden_size) for _ in range(num_layers) ]) self.dropout = nn.Dropout(0.2) def forward(self, g, features): h = self.input_proj(features) for i, (conv, norm) in enumerate(zip(self.layers, self.layer_norms)): h_new = conv(g, h) h_new = norm(h_new) h_new = nn.functional.relu(h_new) h_new = self.dropout(h_new) h = h + h_new if i > 0 else h_new velocity = self.velocity_head(h) direction = self.direction_head(h) direction = nn.functional.normalize(direction, dim=1) return torch.cat([velocity, direction], dim=1) # Denormalize predictions def denormalize_predictions(pred): velocity = pred[:, 0] * (NORM_PARAMS['v_max'] - NORM_PARAMS['v_min']) + NORM_PARAMS['v_min'] direction = np.rad2deg(np.arctan2(pred[:, 1], pred[:, 2])) % 360 return velocity, direction # Function to create plotly visualization with error annotations def plot_errors(box_errors): fig = go.Figure() colors = px.colors.qualitative.Set3 for i, box_error in enumerate(box_errors): coords = box_error['coords'] lon_min, lon_max = np.min(coords[:, 0]), np.max(coords[:, 0]) lat_min, lat_max = np.min(coords[:, 1]), np.max(coords[:, 1]) center_lon = np.mean([lon_min, lon_max]) center_lat = np.mean([lat_min, lat_max]) - (lat_max - lat_min) * 0.35 # Add box boundaries fig.add_trace(go.Scatter( x=[lon_min, lon_max, lon_max, lon_min, lon_min], y=[lat_min, lat_min, lat_max, lat_max, lat_min], mode='lines', line=dict(color=colors[i % len(colors)]) )) # Add annotations for velocity and direction errors fig.add_annotation( x=center_lon, y=center_lat + 0.02, # Adjust position for velocity error at the top text=f"{box_error['velocity_error']:.2f} m/s", showarrow=False, font=dict(size=10) # Adjust font size if needed ) fig.add_annotation( x=center_lon, y=center_lat - 0.01, # Slightly closer to the center for direction error text=f"{box_error['direction_error']:.2f}°", showarrow=False, font=dict(size=10) ) fig.update_layout( title='Wind Prediction Error After 15 Minutes', xaxis_title='Longitude', yaxis_title='Latitude', showlegend=False, hovermode='closest', width=700, # Reduced width height=700 # Reduced height ) fig.update_yaxes(scaleanchor="x", scaleratio=1) return fig # Inference function for Gradio interface def inference(current_csv, future_csv): current_data = pd.read_csv(current_csv) future_data = pd.read_csv(future_csv) # Create graph from current data coords = current_data[['x', 'y']].to_numpy() speeds = current_data['v'].to_numpy() directions = current_data['d'].to_numpy() # Normalize data norm_x = (coords[:, 0] - NORM_PARAMS['x_min']) / (NORM_PARAMS['x_max'] - NORM_PARAMS['x_min']) norm_y = (coords[:, 1] - NORM_PARAMS['y_min']) / (NORM_PARAMS['y_max'] - NORM_PARAMS['y_min']) norm_speeds = (speeds - NORM_PARAMS['v_min']) / (NORM_PARAMS['v_max'] - NORM_PARAMS['v_min']) directions_rad = np.deg2rad(directions) sin_dir = np.sin(directions_rad) cos_dir = np.cos(directions_rad) features = np.column_stack([norm_x, norm_y, norm_speeds, sin_dir, cos_dir]) # Load and run model device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = WindGNN().to(device) checkpoint = torch.load('5.pth', map_location=device) model.load_state_dict(checkpoint['model_state_dict']) model.eval() # Create DGL graph tree = KDTree(coords) distances, indices = tree.query(coords, k=9) src_nodes = [] dst_nodes = [] for i in range(len(coords)): for j in range(1, 9): neighbor_idx = indices[i][j] src_nodes.append(i) dst_nodes.append(neighbor_idx) g = dgl.graph((torch.tensor(src_nodes), torch.tensor(dst_nodes))) g.ndata['feat'] = torch.FloatTensor(features).to(device) # Predict with model with torch.no_grad(): predictions = model(g, g.ndata['feat']).cpu().numpy() # Denormalize predictions pred_velocity, pred_direction = denormalize_predictions(predictions) # Calculate errors true_velocity = future_data['v'].to_numpy() true_direction = future_data['d'].to_numpy() velocity_errors = np.abs(pred_velocity - true_velocity) direction_errors = np.abs(pred_direction - true_direction) mean_velocity_error = np.mean(velocity_errors) mean_direction_error = np.mean(direction_errors) max_velocity_error = np.max(velocity_errors) min_velocity_error = np.min(velocity_errors) max_direction_error = np.max(direction_errors) min_direction_error = np.min(direction_errors) # Prepare box errors box_errors = [] points_per_box = len(coords) // 24 for i in range(24): start_idx = i * points_per_box end_idx = (i + 1) * points_per_box box_error = { 'coords': coords[start_idx:end_idx], 'velocity_error': np.mean(velocity_errors[start_idx:end_idx]), 'direction_error': np.mean(direction_errors[start_idx:end_idx]) } box_errors.append(box_error) fig = plot_errors(box_errors) # Prepare detailed error information error_info = ( f"Mean Velocity Error: {mean_velocity_error:.3f} m/s, " f"Mean Direction Error: {mean_direction_error:.3f}°\n" f"Max Velocity Error: {max_velocity_error:.3f} m/s, " f"Min Velocity Error: {min_velocity_error:.3f} m/s\n" f"Max Direction Error: {max_direction_error:.3f}°, " f"Min Direction Error: {min_direction_error:.3f}°" ) # Combine original and predicted data into a DataFrame for display result_df = pd.DataFrame({ 'x': current_data['x'], 'y': current_data['y'], 'Original Velocity': true_velocity, 'Predicted Velocity': pred_velocity, 'Original Direction': true_direction, 'Predicted Direction': pred_direction }) return fig, error_info, result_df # Paths to example CSV files example_csv_files = [ ["2021-05-03_0500.csv", "2021-05-03_0515.csv"], ["2021-05-03_0515.csv", "2021-05-03_0530.csv"], ["2021-07-01_0700.csv", "2021-07-01_0715.csv"], ["2021-07-01_0715.csv", "2021-07-01_0730.csv"] ] # Gradio Interface iface = gr.Interface( fn=inference, inputs=[gr.File(file_types=['.csv'], label="Current Wind Data CSV"), gr.File(file_types=['.csv'], label="Future Wind Data CSV")], outputs=["plot", "text", gr.DataFrame(label="Prediction Results")], title="Wind Prediction Model", description="Upload CSV files containing current and 15-minute future wind data to see detailed prediction errors per box.", examples=example_csv_files ) iface.launch()