|
|
|
|
|
import os |
|
|
import numpy as np |
|
|
import matplotlib.pyplot as plt |
|
|
import matplotlib.ticker as mticker |
|
|
import cartopy.crs as ccrs |
|
|
import cartopy.feature as cfeature |
|
|
from .interpolation import interpolate_grid |
|
|
from .basemaps import draw_etopo_basemap |
|
|
|
|
|
def export_frames_as_jpgs(animator, output_folder: str, include_metadata: bool = True): |
|
|
os.makedirs(output_folder, exist_ok=True) |
|
|
|
|
|
meta = animator.datasets[0].attrs |
|
|
legend_text = ( |
|
|
f"Run name: {meta.get('run_name', 'N/A')}\n" |
|
|
f"Run time: {meta.get('run_time', 'N/A')}\n" |
|
|
f"Met data: {meta.get('met_data', 'N/A')}\n" |
|
|
f"Start of release: {meta.get('start_of_release', 'N/A')}\n" |
|
|
f"End of release: {meta.get('end_of_release', 'N/A')}\n" |
|
|
f"Source strength: {meta.get('source_strength', 'N/A')} g / s\n" |
|
|
f"Release location: {meta.get('release_location', 'N/A')}\n" |
|
|
f"Release height: {meta.get('release_height', 'N/A')} m asl\n" |
|
|
f"Run duration: {meta.get('run_duration', 'N/A')}" |
|
|
) |
|
|
|
|
|
for z_index, z_val in enumerate(animator.levels): |
|
|
z_dir = os.path.join(output_folder, f"ash_T1-Tn_Z{z_index+1}") |
|
|
os.makedirs(z_dir, exist_ok=True) |
|
|
|
|
|
valid_mask = np.stack([ |
|
|
ds['ash_concentration'].values[z_index] for ds in animator.datasets |
|
|
]).max(axis=0) > 0 |
|
|
y_idx, x_idx = np.where(valid_mask) |
|
|
|
|
|
if y_idx.size == 0 or x_idx.size == 0: |
|
|
print(f"Z level {z_val} km has no valid data. Skipping...") |
|
|
continue |
|
|
|
|
|
y_min, y_max = y_idx.min(), y_idx.max() |
|
|
x_min, x_max = x_idx.min(), x_idx.max() |
|
|
|
|
|
for t in range(len(animator.datasets)): |
|
|
data = animator.datasets[t]['ash_concentration'].values[z_index] |
|
|
interp = interpolate_grid(data, animator.lon_grid, animator.lat_grid) |
|
|
if np.isfinite(interp).sum() == 0: |
|
|
continue |
|
|
|
|
|
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 8), subplot_kw={'projection': ccrs.PlateCarree()}) |
|
|
valid_vals = interp[np.isfinite(interp)] |
|
|
min_val = np.nanmin(valid_vals) |
|
|
max_val = np.nanmax(valid_vals) |
|
|
log_cutoff = 1e-3 |
|
|
log_ratio = max_val / (min_val + 1e-6) |
|
|
use_log = min_val > log_cutoff and log_ratio > 100 |
|
|
|
|
|
levels = np.logspace(np.log10(log_cutoff), np.log10(max_val), 20) if use_log else np.linspace(0, max_val, 20) |
|
|
data_for_plot = np.where(interp > log_cutoff, interp, np.nan) if use_log else interp |
|
|
scale_label = "Hybrid Log" if use_log else "Linear" |
|
|
|
|
|
|
|
|
c1 = ax1.contourf(animator.lons, animator.lats, data_for_plot, levels=levels, |
|
|
cmap="rainbow", alpha=0.6, transform=ccrs.PlateCarree()) |
|
|
draw_etopo_basemap(ax1, mode='stock') |
|
|
ax1.set_extent([animator.lons.min(), animator.lons.max(), animator.lats.min(), animator.lats.max()]) |
|
|
ax1.set_title(f"T{t+1} | Alt: {z_val} km (Full - {scale_label})") |
|
|
ax1.coastlines(); ax1.add_feature(cfeature.BORDERS) |
|
|
ax1.add_feature(cfeature.LAND); ax1.add_feature(cfeature.OCEAN) |
|
|
|
|
|
|
|
|
buffer_y = int((y_max - y_min) * 0.5) |
|
|
buffer_x = int((x_max - x_min) * 0.5) |
|
|
|
|
|
y_start = max(0, y_min - buffer_y) |
|
|
y_end = min(data_for_plot.shape[0], y_max + buffer_y + 1) |
|
|
x_start = max(0, x_min - buffer_x) |
|
|
x_end = min(data_for_plot.shape[1], x_max + buffer_x + 1) |
|
|
|
|
|
zoom = data_for_plot[y_start:y_end, x_start:x_end] |
|
|
lon_zoom = animator.lons[x_start:x_end] |
|
|
lat_zoom = animator.lats[y_start:y_end] |
|
|
|
|
|
c2 = ax2.contourf(lon_zoom, lat_zoom, zoom, levels=levels, |
|
|
cmap="rainbow", alpha=0.6, transform=ccrs.PlateCarree()) |
|
|
draw_etopo_basemap(ax2, mode='stock') |
|
|
ax2.set_extent([lon_zoom.min(), lon_zoom.max(), lat_zoom.min(), lat_zoom.max()]) |
|
|
ax2.set_title(f"T{t+1} | Alt: {z_val} km (Zoom - {scale_label})") |
|
|
ax2.coastlines(); ax2.add_feature(cfeature.BORDERS) |
|
|
ax2.add_feature(cfeature.LAND); ax2.add_feature(cfeature.OCEAN) |
|
|
|
|
|
for ax in [ax1, ax2]: |
|
|
ax.text(0.01, 0.98, f"Time step T{t+1}", transform=ax.transAxes, |
|
|
fontsize=9, color='white', va='top', ha='left', |
|
|
bbox=dict(facecolor='black', alpha=0.4, boxstyle='round')) |
|
|
|
|
|
if include_metadata: |
|
|
for ax in [ax1, ax2]: |
|
|
ax.text(0.01, 0.01, |
|
|
f"Source: NAME model\nRes: {animator.x_res:.2f}°\n{meta.get('run_name', 'N/A')}", |
|
|
transform=ax.transAxes, fontsize=8, color='white', |
|
|
bbox=dict(facecolor='black', alpha=0.5)) |
|
|
ax1.annotate(legend_text, xy=(0.75, 0.99), xycoords='axes fraction', |
|
|
fontsize=8, ha='left', va='top', |
|
|
bbox=dict(boxstyle="round", facecolor="white", edgecolor="gray"), |
|
|
annotation_clip=False) |
|
|
|
|
|
cbar = fig.colorbar(c1, ax=[ax1, ax2], orientation='vertical', shrink=0.75, pad=0.03) |
|
|
cbar.set_label("Ash concentration (g/m³)") |
|
|
formatter = mticker.FuncFormatter(lambda x, _: f'{x:.2g}') |
|
|
cbar.ax.yaxis.set_major_formatter(formatter) |
|
|
|
|
|
if use_log: |
|
|
cbar.ax.text(1.1, 1.02, "log scale", transform=cbar.ax.transAxes, |
|
|
fontsize=9, color='gray', rotation=90, ha='left', va='bottom') |
|
|
|
|
|
frame_path = os.path.join(z_dir, f"frame_{t+1:04d}.jpg") |
|
|
plt.savefig(frame_path, dpi=150, bbox_inches='tight') |
|
|
plt.close(fig) |
|
|
print(f"Saved {frame_path}") |
|
|
|
