import os import math import subprocess from pathlib import Path import matplotlib.pyplot as plt import numpy as np import xarray as xr from matplotlib.backends.backend_pdf import PdfPages from matplotlib.lines import Line2D from matplotlib.patches import Patch # ----------------------------------------------------------------------------- # Configuration # ----------------------------------------------------------------------------- training_data_path = Path( "/global/cfs/projectdirs/m4359/huiwan/m4792/huiwan/mlaer/" "e3sm_output/interstitial_r4_output_3hourly/hist/" ) # Select multiple NetCDF files from training_data_path. training_data_file_prefix = "interstitial_r4_output_3hourly.eam.h0." # Marker strings may be full timestamps, e.g. "2010-01-01-00000", # or shorter substrings, e.g. "01-01". # # If a marker matches multiple files: # - the FIRST match is used for the first-file marker; # - the LAST match is used for the last-file marker. # # Set a marker to None or "" to use the beginning/end of the file list. training_data_first_marker = "2010-01-01-00000" training_data_last_marker = "2010-01-10-00000" # Select every Nth file from the inclusive first/last range. training_data_file_stride = 1 # Nested variable list: # - each inner list defines one variable group # - each group gets its own multi-page PDF for PDF/CDF analyses # - the final ruler PDF uses the same row/column organization variables = [ [ "d_vmr_num_a1_AerMic", "d_vmr_num_a2_AerMic", "d_vmr_num_a4_AerMic", ], [ "d_vmr_so4_a1_AerMic", "d_vmr_so4_a2_AerMic", "d_vmr_so4_a3_AerMic", "d_vmr_H2SO4_AerMic", ], [ "d_vmr_soa_a1_AerMic", "d_vmr_soa_a2_AerMic", "d_vmr_soa_a3_AerMic", "d_vmr_SOAG_AerMic", ], [ "d_vmr_bc_a1_AerMic", "d_vmr_bc_a4_AerMic", "d_vmr_pom_a1_AerMic", "d_vmr_pom_a4_AerMic", ], [ "d_vmr_ncl_a1_AerMic", "d_vmr_ncl_a2_AerMic", "d_vmr_mom_a1_AerMic", "d_vmr_mom_a2_AerMic", "d_vmr_mom_a4_AerMic", ], ] cloud_fraction_variable = "i_cldfr" # If True: # - read/subset/mask the requested variables # - create only the combined transformation-ruler PDF # - skip transformation PDF/CDF calculations # - skip scale-only and z-score calculations # - skip all per-group comparison PDF files # # If False: # - create the full PDF/CDF comparison files # - then create the combined ruler PDF ruler_only = True ruler_only = False time_start = 0 time_end = 0 time_stride = 1 lev_start = 27 lev_end = 71 lev_stride = 1 # Use 500 bins for all PDF histograms. n_bins = 100 overlay_pdf_bins = 200 # Layout for PDF/CDF pages. panel_width = 5.0 panel_height = 5.0 * 0.618 panel_wspace = 0.45 panel_hspace = 0.65 # Important: comparison pages are saved WITHOUT bbox_inches="tight". # This preserves the full fixed-width subplot grid, including intentionally # blank columns for shorter variable groups, and keeps page titles centered # relative to the common page size. # Transformations used in the comparison pages. transformations = [ "symlog", "power_1_7", "power_1_5", "power_1_3", ] transform_colors = { "symlog": "tab:green", "power_1_7": "tab:purple", "power_1_5": "tab:red", "power_1_3": "tab:orange", } # Ruler-only no-transformation reference. no_transform_color = "tab:blue" # Baseline page PDF color. baseline_pdf_color = "tab:blue" # Histogram style. pdf_alpha = 0.35 overlay_pdf_alpha = 0.25 overlay_pdf_linewidth = 0.8 # CDF style. cdf_color = "black" cdf_linewidth = 1.8 cdf_ymin = -0.02 # Reference-line styles. median_line_color = "0.5" median_linewidth = 1.0 zero_line_color = "black" zero_linewidth = 0.8 # Percentiles used on transformed PDF/CDF pages and ruler page. percentiles_to_mark = [5, 35, 65, 95] top_percentile_fontsize = 6 percentile_box_fontsize = 6 # Ruler layout. ruler_box_aspect = 0.5 ruler_panel_width = 4.0 ruler_panel_height = 2.8 ruler_wspace = 0.35 ruler_hspace = 0.55 ruler_legend_y = -0.33 ruler_legend_fontsize = 7.0 ruler_legend_columnspacing = 0.65 ruler_legend_handlelength = 1.5 ruler_legend_handletextpad = 0.30 # ----------------------------------------------------------------------------- # Output # ----------------------------------------------------------------------------- plot_path = "/global/cfs/projectdirs/m4359/www/huiwan/aims-atm/" plot_url = "https://portal.nersc.gov/project/m4359/huiwan/aims-atm/" # Output prefixes are derived automatically from the common variable-name # prefix inside main(), e.g. "d_vmr" or "i_vmr". comparison_file_prefix = None ruler_plot_file = None plot_fmt = "pdf" # ----------------------------------------------------------------------------- # General helpers # ----------------------------------------------------------------------------- def flatten_variable_groups(variable_groups): """Flatten nested variable groups while preserving order.""" return [ variable_name for group in variable_groups for variable_name in group ] def common_variable_prefix(variable_groups): """ Derive a filename-friendly common prefix from all variable names. Examples -------- d_vmr_num_a1_AerMic, d_vmr_so4_a1_AerMic -> d_vmr i_vmr_num_a1_preAerMic, i_vmr_so4_a1_preAerMic -> i_vmr """ names = flatten_variable_groups(variable_groups) if not names: raise ValueError("The nested `variables` list is empty.") prefix = os.path.commonprefix(names) # Prefer a token boundary rather than a partial variable-name token. if "_" in prefix: prefix = prefix[:prefix.rfind("_")] prefix = prefix.rstrip("_") if not prefix: prefix = "variables" return prefix def discover_training_data_files( data_path, filename_prefix, first_marker=None, last_marker=None, file_stride=1, ): """ Use `ls -1` to discover and select a sequence of NetCDF files. Only filenames that start with `filename_prefix` and end in ".nc" are considered. `first_marker` selects the first matching filename in the ls list. `last_marker` selects the last matching filename in the ls list. The first/last range is inclusive, and `file_stride` is then applied to that range. """ data_path = Path(data_path) if not data_path.is_dir(): raise NotADirectoryError( f"Training-data path does not exist or is not a directory: " f"{data_path}" ) if not isinstance(file_stride, int) or file_stride < 1: raise ValueError( "training_data_file_stride must be an integer >= 1." ) result = subprocess.run( ["ls", "-1", str(data_path)], check=True, capture_output=True, text=True, ) filenames = [ line.strip() for line in result.stdout.splitlines() if line.strip() ] candidate_names = [ name for name in filenames if ( name.startswith(filename_prefix) and name.endswith(".nc") ) ] if not candidate_names: raise FileNotFoundError( "No NetCDF files were found that start with " f"{filename_prefix!r} in {data_path}" ) def matching_indices(marker): if marker is None or marker == "": return [] return [ i for i, name in enumerate(candidate_names) if marker in name ] first_matches = matching_indices(first_marker) if first_marker not in (None, "") and not first_matches: raise FileNotFoundError( "No candidate filename contains the first-file marker " f"{first_marker!r}." ) first_index = first_matches[0] if first_matches else 0 last_matches = matching_indices(last_marker) if last_marker not in (None, "") and not last_matches: raise FileNotFoundError( "No candidate filename contains the last-file marker " f"{last_marker!r}." ) last_index = ( last_matches[-1] if last_matches else len(candidate_names) - 1 ) if first_index > last_index: raise ValueError( "The selected first file occurs after the selected last file:\n" f" first = {candidate_names[first_index]}\n" f" last = {candidate_names[last_index]}" ) selected_names = candidate_names[ first_index:last_index + 1:file_stride ] if not selected_names: raise RuntimeError( "The file-selection settings produced an empty file list." ) return [ data_path / name for name in selected_names ] def make_index_slice(start, end, stride): """Treat start == end as selection of one index.""" if start is not None and end is not None and start == end: return slice(start, start + 1, stride) return slice(start, end, stride) def calculate_symlog_scale(values): """Use the smallest finite nonzero magnitude as symlog_scale.""" values = np.asarray(values) values = values[np.isfinite(values)] abs_nonzero = np.abs(values[values != 0.0]) if abs_nonzero.size == 0: raise ValueError( "Cannot determine symlog scale: no finite nonzero values." ) return np.min(abs_nonzero) def transform_values(values, transformation): """Apply a sign-preserving nonlinear transformation.""" values = np.asarray(values) values = values[np.isfinite(values)] if transformation == "none": return values if transformation == "symlog": symlog_scale = calculate_symlog_scale(values) return ( np.sign(values) * np.log10(1.0 + np.abs(values) / symlog_scale) ) power_map = { "power_1_3": 1.0 / 3.0, "power_1_5": 1.0 / 5.0, "power_1_7": 1.0 / 7.0, } if transformation in power_map: power = power_map[transformation] return np.sign(values) * np.abs(values) ** power raise ValueError(f"Unknown transformation {transformation!r}.") def transform_scalar_values(values, transformation, symlog_scale=None): """Transform arbitrary scalar values with the same transform definition.""" values = np.asarray(values, dtype=float) if transformation == "none": return values if transformation == "symlog": if symlog_scale is None: raise ValueError( "symlog_scale must be supplied when mapping scalar values." ) return ( np.sign(values) * np.log10(1.0 + np.abs(values) / symlog_scale) ) power_map = { "power_1_3": 1.0 / 3.0, "power_1_5": 1.0 / 5.0, "power_1_7": 1.0 / 7.0, } if transformation in power_map: power = power_map[transformation] return np.sign(values) * np.abs(values) ** power raise ValueError(f"Unknown transformation {transformation!r}.") def transform_abs_magnitudes(abs_values, transformation, raw_values): """Transform nonnegative magnitudes for the ruler analysis.""" abs_values = np.asarray(abs_values, dtype=float) if transformation == "none": return abs_values if transformation == "symlog": symlog_scale = calculate_symlog_scale(raw_values) return np.log10(1.0 + abs_values / symlog_scale) power_map = { "power_1_3": 1.0 / 3.0, "power_1_5": 1.0 / 5.0, "power_1_7": 1.0 / 7.0, } if transformation in power_map: return abs_values ** power_map[transformation] raise ValueError(f"Unknown transformation {transformation!r}.") def transformation_label(transformation): labels = { "none": "No transformation", "symlog": "Symmetric log", "power_1_3": "Power 1/3", "power_1_5": "Power 1/5", "power_1_7": "Power 1/7", } return labels[transformation] def scale_by_std(values): """Scale without centering: scaled = values / std(values).""" values = np.asarray(values) values = values[np.isfinite(values)] if values.size == 0: return values, np.nan std = np.std(values, ddof=0) if not np.isfinite(std) or std == 0.0: raise ValueError( "Cannot scale by std because the standard deviation is zero " "or invalid." ) return values / std, std def zscore_normalize(values): """Z-score normalize after transformation.""" values = np.asarray(values) values = values[np.isfinite(values)] if values.size == 0: return values, np.nan, np.nan mean = np.mean(values) std = np.std(values, ddof=0) if not np.isfinite(std) or std == 0.0: raise ValueError( "Cannot z-score normalize because the standard deviation is " "zero or invalid." ) return (values - mean) / std, mean, std def format_original_value(value): """Format original percentile values compactly.""" if value == 0: return "0" abs_value = abs(value) if abs_value < 1.0e-3 or abs_value >= 1.0e4: return f"{value:.2e}" return f"{value:.4g}" def get_original_percentiles(raw_values): """Return configured original-value percentiles.""" raw_values = np.asarray(raw_values) raw_values = raw_values[np.isfinite(raw_values)] return np.percentile( raw_values, percentiles_to_mark, ) # ----------------------------------------------------------------------------- # Panel annotations # ----------------------------------------------------------------------------- def set_panel_title( ax, variable_name, transformation=None, transformation_color="black", ): """Show variable name on line 1 and transformation on line 2.""" ax.text( 0.5, 1.19, variable_name, transform=ax.transAxes, ha="center", va="bottom", fontsize=10, fontweight="bold", color="black", clip_on=False, ) if transformation is not None: ax.text( 0.5, 1.11, transformation_label(transformation), transform=ax.transAxes, ha="center", va="bottom", fontsize=9, fontweight="bold", color=transformation_color, clip_on=False, ) def add_original_percentile_top_axis( ax_hist, raw_values, transformation, std_scale=None, ): """Add outward top ticks at P5/P35/P65/P95.""" percentile_values = get_original_percentiles(raw_values) if transformation == "symlog": symlog_scale = calculate_symlog_scale(raw_values) else: symlog_scale = None transformed_positions = transform_scalar_values( percentile_values, transformation, symlog_scale=symlog_scale, ) if std_scale is not None: transformed_positions = transformed_positions / std_scale ax_top = ax_hist.secondary_xaxis("top") ax_top.set_xticks(transformed_positions) ax_top.set_xticklabels([]) ax_top.tick_params( axis="x", direction="out", length=4, width=0.8, pad=1, ) for percentile, xpos in zip( percentiles_to_mark, transformed_positions, ): ax_hist.annotate( f"P{percentile}", xy=(xpos, 1.0), xycoords=("data", "axes fraction"), xytext=(0, -3), textcoords="offset points", ha="center", va="top", fontsize=top_percentile_fontsize, color="black", clip_on=True, ) def add_percentile_value_box(ax_hist, raw_values): """Write original percentile values inside the panel.""" percentile_values = get_original_percentiles(raw_values) lines = ["Original-value percentiles"] lines.extend( f"P{p}: {format_original_value(value)}" for p, value in zip( percentiles_to_mark, percentile_values, ) ) xmin, xmax = ax_hist.get_xlim() xmid = 0.5 * (xmin + xmax) if xmin <= 0.0 <= xmax and 0.0 > xmid: x_box = 0.03 ha = "left" else: x_box = 0.97 ha = "right" ax_hist.text( x_box, 0.88, "\n".join(lines), transform=ax_hist.transAxes, ha=ha, va="top", fontsize=percentile_box_fontsize, color="black", linespacing=1.05, bbox={ "facecolor": "white", "edgecolor": "0.75", "alpha": 0.78, "pad": 1.5, }, zorder=10, ) # ----------------------------------------------------------------------------- # Shared PDF + CDF overlay panel # ----------------------------------------------------------------------------- def plot_pdf_cdf_overlay( ax_hist, values, variable_name, pdf_color, xlabel, raw_values_for_percentiles=None, transformation=None, std_scale=None, ): """Plot filled PDF and black CDF in one panel.""" values = np.asarray(values) values = values[np.isfinite(values)] if values.size == 0: ax_hist.text( 0.5, 0.5, "No finite values", ha="center", va="center", transform=ax_hist.transAxes, ) return # PDF ax_hist.hist( values, bins=n_bins, density=True, histtype="stepfilled", color=pdf_color, edgecolor=pdf_color, linewidth=0.8, alpha=pdf_alpha, ) ax_hist.set_xlabel(xlabel) ax_hist.set_ylabel("Probability density", color=pdf_color) ax_hist.tick_params(axis="y", colors=pdf_color) ax_hist.spines["left"].set_color(pdf_color) pdf_ymax = ax_hist.get_ylim()[1] ax_hist.set_ylim( -0.02 * pdf_ymax, 1.10 * pdf_ymax, ) ax_hist.axvline( 0.0, color=zero_line_color, linestyle="-", linewidth=zero_linewidth, zorder=4, ) median_value = np.median(values) ax_hist.axvline( median_value, color=median_line_color, linestyle="--", linewidth=median_linewidth, zorder=4, ) # CDF sorted_values = np.sort(values) cdf = np.arange(1, sorted_values.size + 1) / sorted_values.size ax_cdf = ax_hist.twinx() ax_cdf.plot( sorted_values, cdf, color=cdf_color, linewidth=cdf_linewidth, ) ax_cdf.set_ylabel("Cumulative probability", color=cdf_color) ax_cdf.tick_params(axis="y", colors=cdf_color) ax_cdf.spines["right"].set_color(cdf_color) ax_cdf.set_ylim(cdf_ymin, 1.10) ax_cdf.set_yticks(np.arange(0.0, 1.01, 0.1)) ax_cdf.grid( True, axis="y", which="major", alpha=0.25, linestyle="--", ) ax_cdf.axhline( 0.5, color=median_line_color, linestyle="--", linewidth=median_linewidth, zorder=2, ) set_panel_title( ax_hist, variable_name, transformation=transformation, transformation_color=pdf_color, ) if ( raw_values_for_percentiles is not None and transformation is not None ): add_original_percentile_top_axis( ax_hist, raw_values_for_percentiles, transformation, std_scale=std_scale, ) add_percentile_value_box( ax_hist, raw_values_for_percentiles, ) # ----------------------------------------------------------------------------- # Data preparation # ----------------------------------------------------------------------------- def prepare_scaled_data(values): """Prepare scale-only and z-score versions for every transformation.""" scale_only_data = {} zscore_data = {} diagnostics = {} for transformation in transformations: transformed = transform_values(values, transformation) scaled, std = scale_by_std(transformed) zscored, mean, zstd = zscore_normalize(transformed) scale_only_data[transformation] = scaled zscore_data[transformation] = zscored info = { "std_before_scaling": std, "mean_before_zscore": mean, "std_before_zscore": zstd, } if transformation == "symlog": info["symlog_scale"] = calculate_symlog_scale(values) diagnostics[transformation] = info return scale_only_data, zscore_data, diagnostics # ----------------------------------------------------------------------------- # Page helpers for one variable group # ----------------------------------------------------------------------------- def make_original_overlay_page( pdf, variable_names, raw_values_by_variable, n_plot_cols, ): """Page 1: original values, PDF + CDF overlaid.""" ncols = n_plot_cols fig, axes = plt.subplots( nrows=1, ncols=ncols, figsize=( panel_width * ncols, panel_height, ), squeeze=False, ) axes = axes.ravel() for i, variable_name in enumerate(variable_names): plot_pdf_cdf_overlay( axes[i], raw_values_by_variable[variable_name], variable_name, pdf_color=baseline_pdf_color, xlabel="Original value", ) for i in range(len(variable_names), ncols): axes[i].set_visible(False) fig.suptitle( "Original distributions: no transformation, no scaling", fontsize=14, fontweight="bold", ) fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.90]) fig.subplots_adjust( wspace=panel_wspace, ) pdf.savefig(fig) plt.close(fig) def make_unscaled_transform_page( pdf, variable_names, raw_values_by_variable, n_plot_cols, ): """Page 2: one transformation per row, no std scaling.""" nrows = len(transformations) ncols = n_plot_cols fig, axes = plt.subplots( nrows=nrows, ncols=ncols, figsize=( panel_width * ncols, panel_height * nrows, ), squeeze=False, ) for row, transformation in enumerate(transformations): pdf_color = transform_colors[transformation] for col, variable_name in enumerate(variable_names): raw_values = raw_values_by_variable[variable_name] transformed = transform_values( raw_values, transformation, ) plot_pdf_cdf_overlay( axes[row, col], transformed, variable_name, pdf_color=pdf_color, xlabel="Transformed value", raw_values_for_percentiles=raw_values, transformation=transformation, std_scale=None, ) for col in range(len(variable_names), ncols): axes[row, col].set_visible(False) fig.suptitle( "Sign-preserving transformations: no std scaling", fontsize=14, fontweight="bold", ) fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.97]) fig.subplots_adjust( wspace=panel_wspace, hspace=panel_hspace, ) pdf.savefig(fig) plt.close(fig) def make_scale_only_page( pdf, variable_names, raw_values_by_variable, scaled_by_variable, diagnostics_by_variable, n_plot_cols, ): """Page 3: transformation + scale-only normalization.""" nrows = len(transformations) ncols = n_plot_cols fig, axes = plt.subplots( nrows=nrows, ncols=ncols, figsize=( panel_width * ncols, panel_height * nrows, ), squeeze=False, ) for row, transformation in enumerate(transformations): pdf_color = transform_colors[transformation] for col, variable_name in enumerate(variable_names): raw_values = raw_values_by_variable[variable_name] scaled_values = scaled_by_variable[variable_name][ transformation ] std_scale = diagnostics_by_variable[ variable_name ][transformation]["std_before_scaling"] plot_pdf_cdf_overlay( axes[row, col], scaled_values, variable_name, pdf_color=pdf_color, xlabel="Transformed value / std", raw_values_for_percentiles=raw_values, transformation=transformation, std_scale=std_scale, ) for col in range(len(variable_names), ncols): axes[row, col].set_visible(False) fig.suptitle( "Sign-preserving transformations + scale-only normalization", fontsize=14, fontweight="bold", ) fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.97]) fig.subplots_adjust( wspace=panel_wspace, hspace=panel_hspace, ) pdf.savefig(fig) plt.close(fig) def make_zscore_page( pdf, variable_names, zscore_by_variable, n_plot_cols, ): """Page 4: transformation + z-score normalization.""" nrows = len(transformations) ncols = n_plot_cols fig, axes = plt.subplots( nrows=nrows, ncols=ncols, figsize=( panel_width * ncols, panel_height * nrows, ), squeeze=False, ) for row, transformation in enumerate(transformations): pdf_color = transform_colors[transformation] for col, variable_name in enumerate(variable_names): zvalues = zscore_by_variable[variable_name][ transformation ] plot_pdf_cdf_overlay( axes[row, col], zvalues, variable_name, pdf_color=pdf_color, xlabel="Z-score after transformation", raw_values_for_percentiles=None, transformation=transformation, std_scale=None, ) for col in range(len(variable_names), ncols): axes[row, col].set_visible(False) fig.suptitle( "Sign-preserving transformations + z-score normalization", fontsize=14, fontweight="bold", ) fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.97]) fig.subplots_adjust( wspace=panel_wspace, hspace=panel_hspace, ) pdf.savefig(fig) plt.close(fig) def make_overlaid_transform_page( pdf, variable_names, data_by_variable, normalization_label, xlabel, n_plot_cols, ): """ Two-row comparison page: top row = PDFs for all transformations overlaid bottom row = CDFs for all transformations overlaid Used for scale-only and z-score comparisons. """ ncols = n_plot_cols fig, axes = plt.subplots( nrows=2, ncols=ncols, figsize=( panel_width * ncols, panel_height * 2, ), squeeze=False, ) legend_handles = [ Line2D( [0], [0], color=transform_colors[t], linewidth=2.0, label=transformation_label(t), ) for t in transformations ] for col, variable_name in enumerate(variable_names): # PDFs ax_pdf = axes[0, col] for i, transformation in enumerate(transformations): values = data_by_variable[variable_name][transformation] color = transform_colors[transformation] if i == 0: ax_pdf.hist( values, bins=overlay_pdf_bins, density=True, histtype="stepfilled", color=color, edgecolor=color, alpha=overlay_pdf_alpha, linewidth=overlay_pdf_linewidth, ) else: ax_pdf.hist( values, bins=overlay_pdf_bins, density=True, histtype="step", color=color, alpha=1.0, linewidth=overlay_pdf_linewidth, ) pdf_ymax = ax_pdf.get_ylim()[1] ax_pdf.set_ylim( -0.02 * pdf_ymax, 1.10 * pdf_ymax, ) ax_pdf.set_xlabel(xlabel) ax_pdf.set_ylabel("Probability density") ax_pdf.set_title( variable_name, fontsize=10, fontweight="bold", ) if col == 0: ax_pdf.legend( handles=legend_handles, title="Transformation", loc="best", fontsize=8, ) # CDFs ax_cdf = axes[1, col] for transformation in transformations: values = data_by_variable[variable_name][transformation] color = transform_colors[transformation] sorted_values = np.sort(values) cdf = ( np.arange(1, sorted_values.size + 1) / sorted_values.size ) ax_cdf.plot( sorted_values, cdf, color=color, linewidth=1.8, ) ax_cdf.axhline( 0.5, color=median_line_color, linestyle="--", linewidth=median_linewidth, ) ax_cdf.set_ylim(cdf_ymin, 1.10) ax_cdf.set_yticks(np.arange(0.0, 1.01, 0.1)) ax_cdf.grid( True, axis="y", which="major", alpha=0.25, linestyle="--", ) ax_cdf.set_xlabel(xlabel) ax_cdf.set_ylabel("Cumulative probability") ax_cdf.set_title( variable_name, fontsize=10, fontweight="bold", ) for col in range(len(variable_names), ncols): axes[0, col].set_visible(False) axes[1, col].set_visible(False) fig.suptitle( f"Overlaid transformed distributions: {normalization_label}", fontsize=14, fontweight="bold", ) fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.95]) fig.subplots_adjust( wspace=panel_wspace, hspace=0.45, ) pdf.savefig(fig) plt.close(fig) def write_group_comparison_pdf( group_index, variable_names, raw_values_by_variable, scaled_by_variable, zscore_by_variable, diagnostics_by_variable, n_plot_cols, comparison_file_prefix, ): """Create one multi-page PDF for one inner list of variables.""" plot_file = f"{comparison_file_prefix}_{group_index}" output_file = os.path.join( plot_path, f"{plot_file}.{plot_fmt}", ) with PdfPages(output_file) as pdf: # Page 1 make_original_overlay_page( pdf, variable_names, raw_values_by_variable, n_plot_cols, ) # Page 2 make_unscaled_transform_page( pdf, variable_names, raw_values_by_variable, n_plot_cols, ) # Page 3 make_scale_only_page( pdf, variable_names, raw_values_by_variable, scaled_by_variable, diagnostics_by_variable, n_plot_cols, ) # Page 4 make_zscore_page( pdf, variable_names, zscore_by_variable, n_plot_cols, ) # Page 5 make_overlaid_transform_page( pdf, variable_names, scaled_by_variable, normalization_label="scale-only normalization", xlabel="Transformed value / std", n_plot_cols=n_plot_cols, ) # Page 6 make_overlaid_transform_page( pdf, variable_names, zscore_by_variable, normalization_label="z-score normalization", xlabel="Z-score after transformation", n_plot_cols=n_plot_cols, ) os.chmod(output_file, 0o644) print("") print(f"Created comparison PDF for variable-group index {group_index}:") print(f" {output_file}") print(f" {plot_url}{plot_file}.{plot_fmt}") # ----------------------------------------------------------------------------- # Combined ruler page # ----------------------------------------------------------------------------- def make_combined_ruler_pdf( raw_values_by_variable, ruler_plot_file, ): """ Create one PDF containing one ruler page. The panel layout matches the nested `variables` list: - each inner list is one row - variable position within the inner list determines the column - unused cells are hidden """ ruler_transformations = [ "none", "symlog", "power_1_7", "power_1_5", "power_1_3", ] ruler_colors = { "none": no_transform_color, "symlog": transform_colors["symlog"], "power_1_7": transform_colors["power_1_7"], "power_1_5": transform_colors["power_1_5"], "power_1_3": transform_colors["power_1_3"], } n_group_rows = len(variables) n_group_cols = max(len(group) for group in variables) fig, axes = plt.subplots( nrows=n_group_rows, ncols=n_group_cols, figsize=( ruler_panel_width * n_group_cols, ruler_panel_height * n_group_rows, ), squeeze=False, ) legend_handles = [ Patch( facecolor="none", edgecolor="0.3", hatch="......", label="< P5", ), Patch( facecolor="0.5", edgecolor="0.3", alpha=0.4, label="P5-P35", ), Patch( facecolor="0.5", edgecolor="0.3", alpha=0.8, label="P35-P65", ), Patch( facecolor="0.5", edgecolor="0.3", alpha=0.4, label="P65-P95", ), Patch( facecolor="none", edgecolor="0.3", hatch="//////", label="> P95", ), ] fontsize_base = 9 for group_row, variable_group in enumerate(variables): for group_col in range(n_group_cols): ax = axes[group_row, group_col] if group_col >= len(variable_group): ax.set_visible(False) continue variable_name = variable_group[group_col] raw_values = np.asarray( raw_values_by_variable[variable_name] ) raw_values = raw_values[np.isfinite(raw_values)] abs_nonzero = np.abs( raw_values[raw_values != 0.0] ) if abs_nonzero.size == 0: ax.text( 0.5, 0.5, "No finite nonzero values", ha="center", va="center", transform=ax.transAxes, ) ax.set_title( variable_name, fontsize=fontsize_base+1.5, fontweight="bold", ) continue percentile_values = np.percentile( abs_nonzero, percentiles_to_mark, ) max_abs = np.max(abs_nonzero) y_positions = np.arange( len(ruler_transformations) ) for ruler_row, transformation in enumerate( ruler_transformations ): color = ruler_colors[transformation] landmarks_original = np.concatenate( ( [0.0], percentile_values, [max_abs], ) ) landmarks_transformed = ( transform_abs_magnitudes( landmarks_original, transformation, raw_values, ) ) transformed_max = landmarks_transformed[-1] if ( not np.isfinite(transformed_max) or transformed_max <= 0.0 ): continue normalized_landmarks = ( landmarks_transformed / transformed_max ) widths = np.diff(normalized_landmarks) left = 0.0 # < P5 ax.barh( ruler_row, widths[0], left=left, height=0.60, facecolor="none", edgecolor=color, linewidth=1.0, hatch="....", ) left += widths[0] # P5-P35 ax.barh( ruler_row, widths[1], left=left, height=0.60, color=color, alpha=0.4, edgecolor=color, linewidth=0.7, ) left += widths[1] # P35-P65 ax.barh( ruler_row, widths[2], left=left, height=0.60, color=color, alpha=0.8, edgecolor=color, linewidth=0.7, ) left += widths[2] # P65-P95 ax.barh( ruler_row, widths[3], left=left, height=0.60, color=color, alpha=0.4, edgecolor=color, linewidth=0.7, ) left += widths[3] # > P95 ax.barh( ruler_row, widths[4], left=left, height=0.60, facecolor="none", edgecolor=color, linewidth=1.0, hatch="/////", ) for boundary in normalized_landmarks[1:-1]: ax.plot( [boundary, boundary], [ ruler_row - 0.30, ruler_row + 0.30, ], color=color, linewidth=0.65, ) ax.set_xlim(0.0, 1.0) ax.set_xlabel( "Fraction of transformed range", fontsize=fontsize_base, ) ax.tick_params( axis="x", labelsize=fontsize_base-1, ) ax.set_yticks(y_positions) ax.set_yticklabels( [ "No transfm.", "Sym. log", "Power 1/7", "Power 1/5", "Power 1/3", ], fontsize=fontsize_base-1, ) for tick_label, transformation in zip( ax.get_yticklabels(), ruler_transformations, ): tick_label.set_color( ruler_colors[transformation] ) ax.invert_yaxis() ax.set_title( variable_name, fontsize=fontsize_base-1, fontweight="bold", ) ax.set_box_aspect(ruler_box_aspect) ax.grid( True, axis="x", linestyle="--", alpha=0.25, ) landmark_lines = [ "Original nonzero |q|", ] landmark_lines.extend( f"P{p}: {format_original_value(value)}" for p, value in zip( percentiles_to_mark, percentile_values, ) ) landmark_lines.append( f"max: {format_original_value(max_abs)}" ) ax.text( 0.96, 0.09, "\n".join(landmark_lines), transform=ax.transAxes, ha="right", va="bottom", fontsize=fontsize_base-2.5, bbox={ "facecolor": "white", "edgecolor": "0.75", "alpha": 0.90, "pad": 2.5, }, ) ax.legend( handles=legend_handles, loc="upper center", bbox_to_anchor=( 0.5, ruler_legend_y, ), ncol=5, fontsize=ruler_legend_fontsize, frameon=False, borderaxespad=0.0, columnspacing=( ruler_legend_columnspacing ), handlelength=( ruler_legend_handlelength ), handletextpad=( ruler_legend_handletextpad ), ) fig.suptitle( ( "How each transformation allocates numerical range\n" "Boundaries: P5, P35, P65, P95 of original nonzero |q|" ), fontsize=14, fontweight="bold", ) fig.tight_layout( rect=[0.0, 0.03, 1.0, 0.96] ) fig.subplots_adjust( left=0.045, right=0.985, bottom=0.05, top=0.92, wspace=ruler_wspace, hspace=ruler_hspace, ) ruler_output_file = os.path.join( plot_path, f"{ruler_plot_file}.{plot_fmt}", ) fig.savefig( ruler_output_file, format=plot_fmt, bbox_inches="tight", ) plt.close(fig) os.chmod(ruler_output_file, 0o644) print("") print("Created combined ruler PDF:") print(f" {ruler_output_file}") print( f" {plot_url}{ruler_plot_file}.{plot_fmt}" ) # ----------------------------------------------------------------------------- # Main # ----------------------------------------------------------------------------- def main(): selection = { "time": make_index_slice( time_start, time_end, time_stride, ), "lev": make_index_slice( lev_start, lev_end, lev_stride, ), } os.makedirs( plot_path, exist_ok=True, ) all_variables = flatten_variable_groups( variables ) variable_prefix = common_variable_prefix( variables ) comparison_file_prefix = ( f"{variable_prefix}_transformation_scaling_comparison" ) ruler_plot_file = ( f"{variable_prefix}_transformation_ruler" ) n_plot_cols = max( len(group) for group in variables ) print( f"Common variable prefix: {variable_prefix}" ) print( f"Using {n_plot_cols} plot columns for every comparison PDF." ) if len(all_variables) != len(set(all_variables)): raise ValueError( "The nested `variables` list contains duplicate variable names." ) # ----------------------------------------------------------------- # Discover the NetCDF files to use. # ----------------------------------------------------------------- training_data_files = discover_training_data_files( training_data_path, training_data_file_prefix, first_marker=training_data_first_marker, last_marker=training_data_last_marker, file_stride=training_data_file_stride, ) print("") print("Selected training-data files:") print(f" path: {training_data_path}") print(f" prefix: {training_data_file_prefix}") print(f" first marker: {training_data_first_marker!r}") print(f" last marker: {training_data_last_marker!r}") print(f" file stride: {training_data_file_stride}") print( f" number of selected files: {len(training_data_files)}" ) for i, training_data_file in enumerate(training_data_files): print( f" [{i:4d}] {training_data_file.name}" ) print("") print( f"Selected time indices in EVERY file: " f"start={time_start}, end={time_end}, stride={time_stride}" ) print( f"Selected lev indices in EVERY file: " f"start={lev_start}, end={lev_end}, stride={lev_stride}" ) # Accumulate one list of arrays per variable, one array per selected file. value_chunks_by_variable = { variable_name: [] for variable_name in all_variables } total_clear_sky_points = 0 required_variables = ( all_variables + [cloud_fraction_variable] ) for file_index, training_data_file in enumerate( training_data_files ): print("") print( f"Reading file {file_index + 1}/{len(training_data_files)}: " f"{training_data_file}" ) with xr.open_dataset(training_data_file) as ds: missing_variables = [ variable_name for variable_name in required_variables if variable_name not in ds.variables ] if missing_variables: raise KeyError( f"File {training_data_file.name} is missing requested " "variables: " + ", ".join(missing_variables) ) cloud_fraction = ds[ cloud_fraction_variable ].isel(selection) clear_sky_mask = ( cloud_fraction == 0 ) file_clear_sky_points = int( clear_sky_mask.sum().values ) total_clear_sky_points += file_clear_sky_points print( f" Clear-sky points " f"({cloud_fraction_variable} == 0): " f"{file_clear_sky_points:,}" ) for variable_name in all_variables: data = ds[ variable_name ].isel(selection) clear_sky_data = data.where( clear_sky_mask ) values = np.asarray( clear_sky_data.values ).ravel() values = values[ np.isfinite(values) ] value_chunks_by_variable[ variable_name ].append(values) print("") print( f"Total clear-sky points across selected files: " f"{total_clear_sky_points:,}" ) # ----------------------------------------------------------------- # Concatenate values across selected files. # ----------------------------------------------------------------- raw_values_by_variable = {} scaled_by_variable = {} zscore_by_variable = {} diagnostics_by_variable = {} for variable_name in all_variables: chunks = value_chunks_by_variable[ variable_name ] values = ( np.concatenate(chunks) if chunks else np.array([], dtype=float) ) raw_values_by_variable[ variable_name ] = values print("") print( f"{variable_name}: " f"{values.size:,} finite clear-sky values " f"from {len(training_data_files)} selected files" ) if values.size == 0: raise ValueError( f"{variable_name}: no finite clear-sky values remain " "after the requested file/time/lev selections." ) percentile_values = np.percentile( values, percentiles_to_mark, ) print( " original percentiles: " + ", ".join( f"P{p}={value:.6e}" for p, value in zip( percentiles_to_mark, percentile_values, ) ) ) if not ruler_only: ( scaled_data, zscore_data, diagnostics, ) = prepare_scaled_data(values) scaled_by_variable[ variable_name ] = scaled_data zscore_by_variable[ variable_name ] = zscore_data diagnostics_by_variable[ variable_name ] = diagnostics for transformation in transformations: info = diagnostics[ transformation ] line = ( f" {transformation:10s} | " f"std_before_scaling=" f"{info['std_before_scaling']:.6e}" ) if "symlog_scale" in info: line += ( f" | symlog_scale=" f"{info['symlog_scale']:.6e}" ) print(line) # ----------------------------------------------------------------- # Create the PDF/CDF comparison files unless ruler_only=True. # ----------------------------------------------------------------- if not ruler_only: for group_index, variable_names in enumerate( variables ): write_group_comparison_pdf( group_index, variable_names, raw_values_by_variable, scaled_by_variable, zscore_by_variable, diagnostics_by_variable, n_plot_cols, comparison_file_prefix, ) else: print("") print( "ruler_only = True: " "skipping all PDF/CDF comparison files." ) # The ruler uses only raw_values_by_variable. make_combined_ruler_pdf( raw_values_by_variable, ruler_plot_file, ) print("") print("**********") print("Done.") if not ruler_only: print("") print("Comparison files:") for group_index in range(len(variables)): plot_file = ( f"{comparison_file_prefix}_" f"{group_index}" ) print( f" {plot_url}" f"{plot_file}.{plot_fmt}" ) print("") print("Ruler file:") print( f" {plot_url}" f"{ruler_plot_file}.{plot_fmt}" ) print("") print("**********") if __name__ == "__main__": main()