import argparse from pathlib import Path import matplotlib.pyplot as plt import numpy as np from config import load_config, project_path DEFAULT_VARIABLES = ("t2m", "msl", "tp", "tcwv") VARIABLE_STYLES = { "t2m": ("2 m temperature", "Temperature (K)", "coolwarm"), "d2m": ("2 m dew point", "Temperature (K)", "coolwarm"), "sst": ("Sea-surface temperature", "Temperature (K)", "coolwarm"), "msl": ("Mean sea-level pressure", "Pressure (model units)", "viridis"), "tp": ("Total precipitation", "Precipitation (model units)", "Blues"), "tcwv": ("Total column water vapour", "Water vapour (model units)", "YlGnBu"), } def find_input(output_dir): candidates = sorted(Path(output_dir).glob("*.npy"), key=lambda path: path.stat().st_mtime) if not candidates: raise FileNotFoundError( f"No inference result (*.npy) was found in {output_dir}. " "Run scripts/inference.py first or pass --input explicitly." ) return candidates[-1] def select_forecast(array, step, ensemble): if array.ndim == 5: if ensemble >= array.shape[0]: raise IndexError(f"Ensemble index {ensemble} is outside array shape {array.shape}") array = array[ensemble] if array.ndim == 4: if step >= array.shape[0]: raise IndexError(f"Forecast step {step} is outside array shape {array.shape}") array = array[step] if array.ndim != 3: raise ValueError( "Expected an inference result shaped [time, channel, lat, lon], " f"[channel, lat, lon], or [ensemble, time, channel, lat, lon], got {array.shape}" ) return np.asarray(array, dtype=np.float32) def parse_variables(value, variables): names = [name.strip() for name in value.split(",") if name.strip()] if not names: raise ValueError("At least one variable must be selected with --variables") unknown = [name for name in names if name not in variables] if unknown: raise ValueError(f"Unknown variables: {unknown}. Available variables are configured in config.yaml") return names def plot_field(axis, colorbar_axis, field, variable, latitudes, longitudes): finite = field[np.isfinite(field)] if finite.size == 0: axis.set_title(f"{VARIABLE_STYLES.get(variable, (variable,))[0]} (no finite data)") axis.set_facecolor("#d9dde3") axis.text( 0.5, 0.5, "No finite values", transform=axis.transAxes, ha="center", va="center", fontsize=12, color="#263238", ) axis.set_xlabel("Longitude (degrees)") axis.set_ylabel("Latitude (degrees)") axis.set_xticks(np.arange(0, 361, 60)) axis.set_yticks(np.arange(-90, 91, 30)) colorbar_axis.set_visible(False) return low, high = np.percentile(finite, [2, 98]) if low == high: low, high = float(finite.min()), float(finite.max() + 1.0) label, colorbar_label, colormap = VARIABLE_STYLES.get( variable, (variable, "Model output", "viridis") ) image = axis.imshow( field, cmap=colormap, vmin=low, vmax=high, extent=(longitudes[0], longitudes[-1], latitudes[-1], latitudes[0]), interpolation="nearest", aspect="auto", ) axis.set_title(label) axis.set_xlabel("Longitude (degrees)") axis.set_ylabel("Latitude (degrees)") axis.set_xticks(np.arange(0, 361, 60)) axis.set_yticks(np.arange(-90, 91, 30)) axis.grid(color="white", alpha=0.3, linewidth=0.6) axis.text( 0.02, 0.03, f"range {finite.min():.3g} to {finite.max():.3g}\n" f"mean {finite.mean():.3g} | p02-p98 {low:.3g}-{high:.3g}", transform=axis.transAxes, color="white", fontsize=8, va="bottom", bbox={"facecolor": "black", "alpha": 0.55, "pad": 3, "edgecolor": "none"}, ) colorbar = plt.colorbar(image, cax=colorbar_axis, orientation="horizontal") colorbar.set_label(colorbar_label) def main(): config = load_config() data_config = config["data"] inference_config = config["inference"] variables = data_config["variables"] parser = argparse.ArgumentParser( description="Render FuXi-S2S forecast fields as a geographic multi-panel image." ) parser.add_argument( "--input", default=None, help="Inference .npy file; defaults to the newest file in inference.output_dir", ) parser.add_argument( "--output", default=None, help="PNG output path; defaults to inference.visualization_dir/.png", ) parser.add_argument("--title", default=inference_config["visualization_title"]) parser.add_argument("--step", type=int, default=0, help="Forecast time index to display") parser.add_argument("--ensemble", type=int, default=0, help="Ensemble member to display") parser.add_argument( "--variables", default=",".join(name for name in DEFAULT_VARIABLES if name in variables), help="Comma-separated variable names to display", ) args = parser.parse_args() input_path = Path(args.input) if args.input else find_input(project_path(inference_config["output_dir"])) array = select_forecast(np.load(input_path), args.step, args.ensemble) selected_variables = parse_variables(args.variables, variables) if array.shape[0] != len(variables): raise ValueError( f"Inference result has {array.shape[0]} channels, but config.yaml defines {len(variables)} variables" ) output = Path(args.output) if args.output else project_path(inference_config["visualization_dir"]) / f"{input_path.stem}.png" output.parent.mkdir(parents=True, exist_ok=True) height, width = array.shape[-2:] latitudes = np.linspace(90.0, -90.0, height) longitudes = np.linspace(0.0, 360.0, width, endpoint=False) figure = plt.figure(figsize=(16, 9), constrained_layout=True) layout = figure.add_gridspec( 2, len(selected_variables), height_ratios=[12, 1], width_ratios=[1] * len(selected_variables), ) figure.suptitle( f"{args.title}\n{input_path.name} | forecast step {args.step} | ensemble {args.ensemble}", fontsize=15, fontweight="bold", ) for index, variable in enumerate(selected_variables): axis = figure.add_subplot(layout[0, index]) colorbar_axis = figure.add_subplot(layout[1, index]) variable_index = variables.index(variable) plot_field(axis, colorbar_axis, array[variable_index], variable, latitudes, longitudes) colorbar_axis.set_xticks([]) colorbar_axis.set_yticks([]) figure.savefig(output, dpi=180, facecolor="white") plt.close(figure) print(f"Input: {input_path}") print(f"Output: {output}") if __name__ == "__main__": main()