|
| 1 | +# era5_to_rocketpy_patch.py |
| 2 | +# Author: Amanda Wech |
| 3 | +# Mission: Make ERA5 NetCDF work with RocketPy EnvironmentAnalysis |
| 4 | +# Status: LIFTOFF ACHIEVED |
| 5 | + |
| 6 | +import xarray as xr |
| 7 | +import numpy as np |
| 8 | +import os |
| 9 | + |
| 10 | +def create_era5_compatible_surface(input_path, output_path, target_lat, target_lon): |
| 11 | + """Patch ERA5 surface data to be RocketPy-compatible.""" |
| 12 | + print(f"Loading ERA5 surface: {input_path}") |
| 13 | + ds = xr.open_dataset(input_path) |
| 14 | + |
| 15 | + # Rename time |
| 16 | + if 'valid_time' in ds: |
| 17 | + ds = ds.rename({'valid_time': 'time'}) |
| 18 | + print(" → Renamed 'valid_time' to 'time'") |
| 19 | + |
| 20 | + # Add v10, u100, v100 |
| 21 | + if 'v10' not in ds: |
| 22 | + ds['v10'] = xr.zeros_like(ds['u10']) |
| 23 | + print(" → Added 'v10' (zeroed)") |
| 24 | + ds['u100'] = ds['u10'] |
| 25 | + ds['v100'] = ds['v10'] |
| 26 | + print(" → Added u100/v100") |
| 27 | + |
| 28 | + # Required variables |
| 29 | + for var, val in {'t2m': 293.15, 'sp': 95000.0, 'tp': 0.0, 'z': 0.0}.items(): |
| 30 | + if var not in ds: |
| 31 | + ds[var] = ds['u10'] * 0 + val |
| 32 | + print(f" → Added '{var}' = {val}") |
| 33 | + |
| 34 | + # Optional variables |
| 35 | + for var, val in {'cbh': 1000.0, 'i10fg': 10.0, 'tcc': 0.0, 'blh': 500.0, |
| 36 | + 'sshf': 0.0, 'slhf': 0.0, 'ssrd': 0.0}.items(): |
| 37 | + if var not in ds: |
| 38 | + ds[var] = ds['u10'] * 0 + val |
| 39 | + print(f" → Added '{var}' = {val}") |
| 40 | + |
| 41 | + # Snap to grid |
| 42 | + lats = ds['latitude'].values |
| 43 | + lons = ds['longitude'].values |
| 44 | + best_lat = float(lats[np.argmin(np.abs(lats - target_lat))]) |
| 45 | + best_lon = float(lons[np.argmin(np.abs(lons - target_lon))]) |
| 46 | + dist_km = ((target_lat - best_lat)**2 + (target_lon - best_lon)**2)**0.5 * 111 |
| 47 | + print(f" → Grid point: ({best_lat}, {best_lon}) | ~{dist_km:.1f} km") |
| 48 | + |
| 49 | + ds.to_netcdf(output_path) |
| 50 | + print(f"ERA5 SURFACE PATCHED → {output_path}") |
| 51 | + ds.close() |
| 52 | + return best_lat, best_lon |
| 53 | + |
| 54 | + |
| 55 | +def create_dummy_pressure(output_path, best_lat, best_lon): |
| 56 | + """Create 2x2 dummy pressure file to avoid RocketPy index bugs.""" |
| 57 | + if os.path.exists(output_path): |
| 58 | + os.remove(output_path) |
| 59 | + print(f"Deleted old: {output_path}") |
| 60 | + |
| 61 | + lat_vals = np.array([best_lat - 0.1, best_lat + 0.1]) |
| 62 | + lon_vals = np.array([best_lon - 0.1, best_lon + 0.1]) |
| 63 | + level_vals = np.array([1000, 850, 700, 500, 300, 200, 100]) |
| 64 | + time_vals = np.array([], dtype='datetime64[h]') |
| 65 | + shape = (0, len(level_vals), 2, 2) |
| 66 | + empty_4d = np.full(shape, np.nan) |
| 67 | + |
| 68 | + ds = xr.Dataset( |
| 69 | + { |
| 70 | + 'u': (['time', 'level', 'latitude', 'longitude'], empty_4d), |
| 71 | + 'v': (['time', 'level', 'latitude', 'longitude'], empty_4d), |
| 72 | + 't': (['time', 'level', 'latitude', 'longitude'], empty_4d), |
| 73 | + 'z': (['time', 'level', 'latitude', 'longitude'], empty_4d), |
| 74 | + }, |
| 75 | + coords={ |
| 76 | + 'time': ('time', time_vals), |
| 77 | + 'level': ('level', level_vals, {'units': 'hPa'}), |
| 78 | + 'latitude': ('latitude', lat_vals, {'units': 'degrees_north'}), |
| 79 | + 'longitude': ('longitude', lon_vals, {'units': 'degrees_east'}) |
| 80 | + } |
| 81 | + ) |
| 82 | + ds.to_netcdf(output_path) |
| 83 | + print(f"DUMMY PRESSURE CREATED → {output_path}") |
0 commit comments