Skip to content

Commit f5c2f57

Browse files
sdat2sdat2
committed
swegnn
1 parent 5104e4c commit f5c2f57

File tree

3 files changed

+259
-2
lines changed

3 files changed

+259
-2
lines changed

‎adforce/dual_graph.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -652,6 +652,7 @@ def plot_swegnn_ghost_cells(path_in: str, output_path: str) -> None:
652652
ax.legend(ordered_handles, ordered_labels, fontsize="small", loc="best")
653653

654654
plt.savefig(output_path, bbox_inches="tight", dpi=300)
655+
plt.savefig(output_path.replace(".pdf", ".png"), bbox_inches="tight", dpi=400)
655656
plt.close()
656657
print(f"Ghost cell plot saved to {output_path}")
657658

‎adforce/mesh.py‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -610,6 +610,7 @@ def dual_graph_ds_from_mesh_ds_from_path(
610610

611611
def swegnn_dg_from_mesh_ds_from_path(
612612
path: str = os.path.join(DATA_PATH, "exp_0049"),
613+
time_downsampling: int = 36,
613614
use_dask: bool = True,
614615
) -> xr.Dataset:
615616
"""
@@ -620,6 +621,7 @@ def swegnn_dg_from_mesh_ds_from_path(
620621
621622
Args:
622623
path (str, optional): Path to ADCIRC fort.*.nc files. Defaults to DATA_PATH/exp_0049.
624+
time_downsampling (int, optional): Factor to downsample time dimension. Defaults to 36.
623625
use_dask (bool, optional): Whether to use dask for loading. Defaults to True.
624626
625627
Returns:
@@ -633,6 +635,9 @@ def swegnn_dg_from_mesh_ds_from_path(
633635
use_dask=use_dask,
634636
take_grad="static" # Only calculate static depth gradient
635637
)
638+
# Downsample time dimension, getting every nth timestep
639+
if time_downsampling > 1:
640+
dg = dg.isel(time=slice(None, None, time_downsampling))
636641

637642
# --- Step 2: Create a new Dataset for SWE-GNN format ---
638643
swegnn_ds = xr.Dataset()
@@ -752,18 +757,25 @@ def swegnn_dg_from_mesh_ds_from_path(
752757

753758

754759
@timeit
755-
def swegnn_netcdf_creation(path_in: str, path_out: str, use_dask: bool = True) -> None:
760+
def swegnn_netcdf_creation(path_in: str,
761+
path_out: str,
762+
time_downsampling: int = 36,
763+
use_dask: bool = True) -> None:
756764
"""
757765
Create a netCDF file formatted for mSWE-GNN from ADCIRC fort.*.nc files.
758766
759767
Args:
760768
path_in (str): Path to the directory containing fort.*.nc files.
761769
path_out (str): Path to save the output netCDF file. E.g. "swegnn_data.nc".
770+
time_downsampling (int, optional): Factor to downsample time dimension. Defaults to 36.
762771
use_dask (bool, optional): Whether to use dask for loading. Defaults to True.
763772
"""
764773
# Define encoding settings
765774

766-
swegnn_ds = swegnn_dg_from_mesh_ds_from_path(path=path_in, use_dask=use_dask)
775+
swegnn_ds = swegnn_dg_from_mesh_ds_from_path(path=path_in,
776+
time_downsampling=time_downsampling,
777+
use_dask=use_dask)
778+
767779

768780
encoding_dict = {}
769781
float_vars = [
@@ -2025,6 +2037,7 @@ def plot_contour(
20252037
# path=os.path.join(DATA_PATH, "exp_0049"), # Adjust path as needed
20262038
# # use_dask=False # Using use_dask=False might be simpler initially for debugging
20272039
# )
2040+
# python -m adforce.mesh
20282041
swegnn_netcdf_creation(os.path.join(DATA_PATH, "exp_0049"),
20292042
os.path.join(DATA_PATH, "exp_0049", "swegnn.nc")
20302043
)

‎adforce/tc_plots.py‎

Lines changed: 243 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,243 @@
1+
"""
2+
Module for plotting TC tracks and simulation summary data.
3+
4+
This script generates a summary plot of all U.S. landfalling TC tracks
5+
on the ADCIRC mesh, colored by intensity.
6+
7+
Usage:
8+
python -m adforce.plotting [--test-single] [--output-name "track_summary.pdf"]
9+
"""
10+
11+
import os
12+
import numpy as np
13+
import netCDF4 as nc
14+
import xarray as xr
15+
import pandas as pd
16+
from adcircpy import AdcircMesh
17+
import matplotlib.pyplot as plt
18+
from matplotlib.collections import LineCollection
19+
from matplotlib.colors import ListedColormap, BoundaryNorm
20+
import argparse
21+
from sithom.plot import plot_defaults, get_dim
22+
23+
# --- Assuming constants.py and ibtracs.py are accessible ---
24+
# Use .constants and .ibtracs if running as a module,
25+
# or adjust paths if running as a standalone script.
26+
try:
27+
from .constants import FIGURE_PATH, SETUP_PATH, PROJ_PATH
28+
from tcpips.ibtracs import na_landing_tcs
29+
except ImportError:
30+
print("Warning: Running as standalone. Assuming relative paths.")
31+
# Define fallback paths if run directly from project root
32+
PROJ_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
33+
SETUP_PATH = os.path.join(PROJ_PATH, 'setup')
34+
FIGURE_PATH = os.path.join(PROJ_PATH, 'figures')
35+
# This is complex, assumes tcpips is a sibling directory
36+
import sys
37+
sys.path.append(PROJ_PATH)
38+
from tcpips.ibtracs import na_landing_tcs
39+
40+
from .generate_training_data import _decode_char_array
41+
from .constants import DATA_PATH
42+
43+
plot_defaults()
44+
45+
# --- Saffir-Simpson Scale Definitions (in m/s) ---
46+
KNOTS_TO_MS = 0.514444
47+
48+
# Original wind speed bins in knots
49+
SS_BINS_KNOTS = [0, 34, 64, 83, 96, 113, 137, 500] # Bin edges
50+
51+
# Convert bins to m/s and round to one decimal
52+
SS_BINS = [np.round(kt * KNOTS_TO_MS, 1) for kt in SS_BINS_KNOTS]
53+
# Result: [0.0, 17.5, 32.9, 42.7, 49.4, 58.1, 70.5, 257.2]
54+
55+
SS_COLORS = ['blue', 'green', 'yellow', 'orange', 'red', 'darkred', 'magenta']
56+
57+
# Labels updated to reflect m/s
58+
SS_LABELS = [
59+
'TD (<17.5 m/s)', 'TS (17.5-32.8 m/s)', 'Cat 1 (32.9-42.6 m/s)',
60+
'Cat 2 (42.7-49.3 m/s)', 'Cat 3 (49.4-58.0 m/s)',
61+
'Cat 4 (58.1-70.4 m/s)', 'Cat 5 (70.5+ m/s)'
62+
]
63+
SS_CMAP = ListedColormap(SS_COLORS)
64+
SS_NORM = BoundaryNorm(SS_BINS, SS_CMAP.N) # This now uses the m/s bins
65+
66+
67+
# --- Plotting Function ---
68+
69+
def plot_all_tc_tracks_on_mesh(
70+
all_storms_ds: xr.Dataset,
71+
mesh_path: str,
72+
output_path: str,
73+
test_single: bool = False
74+
) -> None:
75+
"""
76+
Plots TC tracks from IBTrACS on the ADCIRC mesh.
77+
78+
Tracks are colored by wind intensity using the Saffir-Simpson scale.
79+
80+
Args:
81+
all_storms_ds (xr.Dataset): Dataset containing all storms from
82+
na_landing_tcs().
83+
mesh_path (str): Path to the fort.14 mesh file.
84+
output_path (str): Path to save the output .pdf plot.
85+
test_single (bool, optional): If True, plots only Katrina 2005.
86+
Defaults to False.
87+
"""
88+
print(f"Loading mesh from {mesh_path}...")
89+
try:
90+
with nc.Dataset(mesh_path, 'r') as ds:
91+
x_nodes = ds.variables['x'][:]
92+
y_nodes = ds.variables['y'][:]
93+
triangles = ds.variables['element'][:] -1 # adcircpy elements are 0-based
94+
except Exception as e:
95+
print(f"Error loading mesh: {e}. Cannot plot mesh background.")
96+
return
97+
98+
print("Setting up plot...")
99+
fig, ax = plt.subplots() #figsize=(15, 12))
100+
101+
# Plot the mesh (faintly)
102+
print("Plotting mesh background...")
103+
ax.triplot(
104+
x_nodes, y_nodes, triangles,
105+
color='grey', alpha=0.2, linewidth=0.1, label='ADCIRC Mesh'
106+
)
107+
108+
# Determine storms to plot
109+
if test_single:
110+
# Find Katrina 2005, just as in drive_all_adcirc
111+
i_ran = np.where([x == b"KATRINA" for x in all_storms_ds.name.values])[0]
112+
if len(i_ran) == 0:
113+
indices = [0] # Fallback to first storm
114+
print("Warning: Katrina 2005 not found. Plotting first storm.")
115+
else:
116+
indices = [i_ran[-1]] # Use last Katrina entry
117+
print("Plotting in single mode (Katrina 2005)...")
118+
else:
119+
indices = range(len(all_storms_ds.storm))
120+
print(f"Plotting all {len(indices)} storms...")
121+
122+
# Loop and plot each track
123+
tracks_plotted = 0
124+
for i in indices:
125+
storm_ds = all_storms_ds.isel(storm=i)
126+
storm_name = _decode_char_array(storm_ds['name'])
127+
128+
# Extract and clean data
129+
lat = storm_ds['usa_lat'].values
130+
lon = storm_ds['usa_lon'].values
131+
wind = storm_ds['usa_wind'].values * KNOTS_TO_MS # Convert to m/s
132+
133+
valid_mask = ~np.isnan(lat) & ~np.isnan(lon) & ~np.isnan(wind)
134+
lat, lon, wind = lat[valid_mask], lon[valid_mask], wind[valid_mask]
135+
136+
# --- CONVERT TO M/S ---
137+
# Data from IBTrACS (via na_landing_tcs) is in knots, convert to m/s
138+
wind = wind * KNOTS_TO_MS
139+
# ----------------------
140+
141+
if len(lat) < 2:
142+
# print(f"Skipping {storm_name} (insufficient data).")
143+
continue
144+
145+
# Create segments [[(x1, y1), (x2, y2)], [(x2, y2), (x3, y3)], ...]
146+
points = np.array([lon, lat]).T.reshape(-1, 1, 2)
147+
segments = np.concatenate([points[:-1], points[1:]], axis=1)
148+
149+
# Use wind at the start of each segment for coloring
150+
# This is now in m/s
151+
segment_winds = wind[:-1]
152+
153+
# Create a LineCollection
154+
lc = LineCollection(
155+
segments, cmap=SS_CMAP, norm=SS_NORM,
156+
linewidths=0.5, alpha=0.7 # Thinner, semi-transparent lines
157+
)
158+
lc.set_array(segment_winds) # Set array with m/s values
159+
ax.add_collection(lc)
160+
tracks_plotted += 1
161+
162+
print(f"Plotted {tracks_plotted} tracks.")
163+
164+
# --- Final plot formatting ---
165+
ax.set_aspect('equal')
166+
ax.set_xlabel("Longitude [$^{\circ}$E]")
167+
ax.set_ylabel("Latitude [$^{\circ}$N]")
168+
# title = "Historical U.S. Landfalling TC Tracks on ADCIRC Mesh"
169+
#if test_single:
170+
# title = "TC Track for Katrina (2005) on ADCIRC Mesh"
171+
#ax.set_title(title)
172+
173+
# Set plot limits to the mesh extent
174+
ax.set_xlim(x_nodes.min(), x_nodes.max())
175+
ax.set_ylim(y_nodes.min(), y_nodes.max())
176+
177+
# Add a custom colorbar
178+
# We create a dummy ScalarMappable to link the cmap and norm
179+
sm = plt.cm.ScalarMappable(cmap=SS_CMAP, norm=SS_NORM)
180+
sm.set_array([]) # Dummy array
181+
cbar = fig.colorbar(
182+
sm,
183+
ax=ax,
184+
boundaries=SS_BINS,
185+
# Center ticks in each color block
186+
ticks=[b + (SS_BINS[i+1]-b)/2 for i, b in enumerate(SS_BINS[:-1])],
187+
# spacing='proportional'
188+
)
189+
cbar.ax.set_yticklabels(SS_LABELS, fontsize='small')
190+
cbar.set_label("Wind Speed (m/s) - Saffir-Simpson Scale") # Updated label
191+
192+
# Save the figure
193+
os.makedirs(os.path.dirname(output_path), exist_ok=True)
194+
plt.savefig(output_path, bbox_inches='tight', dpi=300)
195+
plt.close(fig)
196+
print(f"✅ Successfully plotted tracks to {output_path}")
197+
198+
199+
# --- Main execution block to make this script runnable ---
200+
201+
if __name__ == "__main__":
202+
# python -m adforce.plotting --test-single
203+
parser = argparse.ArgumentParser(
204+
description="Plot TC tracks on the ADCIRC mesh."
205+
)
206+
parser.add_argument(
207+
"--test-single",
208+
action="store_true",
209+
help="If set, only process Katrina 2005 for testing."
210+
)
211+
parser.add_argument(
212+
"--output-name",
213+
type=str,
214+
default="all_tc_tracks.pdf",
215+
help="Name for the output plot file in the figures/tc_tracks directory."
216+
)
217+
args = parser.parse_args()
218+
219+
print("Loading all storm data from IBTrACS...")
220+
try:
221+
# This data is still loaded with wind speed in knots
222+
all_storms_ds = na_landing_tcs()
223+
except Exception as e:
224+
print(f"Error loading IBTrACS data: {e}")
225+
print("Please ensure 'tcpips' package is installed and accessible.")
226+
sys.exit(1)
227+
228+
output_file = args.output_name
229+
if args.test_single and "all_" in output_file:
230+
output_file = output_file.replace("all_", "test_single_")
231+
232+
# Define a dedicated subdir for these plots
233+
output_folder = os.path.join(FIGURE_PATH, "tc_tracks")
234+
os.makedirs(output_folder, exist_ok=True)
235+
output_plot_path = os.path.join(output_folder, output_file)
236+
237+
# The plotting function will now handle the conversion to m/s
238+
plot_all_tc_tracks_on_mesh(
239+
all_storms_ds=all_storms_ds,
240+
mesh_path=os.path.join(DATA_PATH, "exp_0049", "fort.63.nc"),
241+
output_path=output_plot_path,
242+
test_single=args.test_single
243+
)

0 commit comments

Comments
 (0)