From 38a02de24e9139093bd03036f128b36e99fa67a1 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Fri, 13 Feb 2026 17:46:55 -0800 Subject: [PATCH] Add tropical-only region option for restart generation and update runbook --- docs/CNP_pipeline_runbook.md | 151 +++++++++++++++-- scripts/ai_predictions_to_restart.py | 244 +++++++++++++++++++++------ 2 files changed, 333 insertions(+), 62 deletions(-) diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 8778d1c..e1f83d0 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -3,7 +3,7 @@ This guide walks through the end-to-end workflow using a user-defined CNP_IO list. ### 1) Create your CNP_IO list -Use the CNP_IO template to create a user-defined list, e.g. `CNP_IO_demo1.txt`). +Use the CNP_IO template to create a user-defined list, e.g. `CNP_IO_demo.txt`). ### 2) Train the AI model Run training with your variable list (edit the filename as needed): @@ -15,6 +15,40 @@ python train_cnp_model.py --variable-list CNP_IO_demo.txt --epoch 100 2>&1 & Notes: - This launches training in the background and redirects logs to stdout/stderr. - The run output directory will be created under `cnp_results/run_YYYYMMDD_HHMMSS`. +- Optional: restrict training/test split to tropical latitudes only: +```bash +python train_cnp_model.py --variable-list CNP_IO_demo.txt --epoch 100 \ + --tropical-only +``` +- Optional: customize the latitude range or column name: +```bash +python train_cnp_model.py --variable-list CNP_IO_demo.txt --epoch 100 \ + --tropical-only --tropical-lat-range -23.5,23.5 --tropical-lat-column Latitude +``` + +### 2a) Fine-tune a pretrained model (optional) +If you already have a trained checkpoint and want to continue training on a TVA-style dataset, use the fine-tuning helper. Populate the necessary paths in your CNP_IO file (e.g. `CNP_IO_updated9_dev_gao.txt`): + +``` +fine_tuning_dataset_path: /path/to/TVA_dataset_root +pretrained_model_path: /path/to/base_run/cnp_predictions/model.pth +output_finetuned_model_dir: ./cnp_results/ +finetuned_model_filename: model.pth +num_epochs: 150 +``` + +Then launch fine-tuning: + +```bash +python scripts/run_finetuning.py +``` + +What happens: +- The script reuses the original training pipeline (same preprocessing, normalization, ModelTrainer). +- `fine_tuning_dataset_path` + `FILE_PATTERN` point to the TVA batches to load. +- `pretrained_model_path` is loaded before training begins so weights pick up where the base run ended. +- Outputs are written to `cnp_results/run_YYYYMMDD_HHMMSS/` (same layout as full training). + Check `cnp_predictions/model.pth`, `cnp_config.json`, `cnp_metrics.json`, and the refreshed prediction CSVs. ### 3) Navigate to the run directory @@ -22,25 +56,41 @@ Notes: cd cnp_results/run_YYYYMMDD_HHMMSS # e.g., cnp_results/run_20250815_205419 ``` -### 4) Validate predictions vs ground truth (test split) using the current run directory as default -Generates scatter plots and statistics using individually normalized results. +### 4) Validate predictions vs ground truth (test split) +Generates quick statistics and a prediction quality report. The report now also creates filtered plots for the “top variables by bad-count”. ```bash -python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & +python ../../scripts/cnp_result_validationplot.py --stats-only +python ../../scripts/generate_prediction_quality_report.py ``` -### 4.5) Generate comprehensive prediction quality report -Creates detailed quality analysis categorizing predictions as "good", "ok", or "bad" based on statistical thresholds. +### 4.1) Highlight a site in all scatter plots (optional) +Highlight a given site (lon/lat) in every GT vs Pred scatter plot. The site will be marked as a red star. ```bash -python ../../scripts/generate_prediction_quality_report.py > prediction_quality_report.log 2>&1 & +python ../../scripts/cnp_result_validationplot_site.py . \ + --lon 303.75 --lat -17.4246 ``` -**Output**: Creates `analysis/` directory with: -- Quality assessment CSV files -- Interactive HTML report -- Visualization charts (bar charts, pie charts, scatter plots) -- Text summary report +Options and behavior: +- The script matches samples by coordinate with a tolerance (default: 0.01 degrees). +- Use `--tolerance N` to relax the matching if no point is found. +- Use `--coord-csv /path/to/test_static_inverse.csv` to explicitly provide a coordinate source. +- Supports `--top-bad-only` / `--worst-only` for restricted plotting. + +**Output**: scatter plots are saved under `plots_site/`. + +Options and behavior: +- The quality report saves outputs under `analysis/`: + - `detailed_quality_assessment.csv`, `variable_quality_summary.csv`, `overall_prediction_quality.png`, `prediction_quality_by_variable.png`, `r2_vs_rmse.png`, `prediction_quality_report.html`, `quality_summary_report.txt`. + - `bad_predictions_detailed.csv` (full list of rows classified as bad; can be disabled). + - `top_bad_plots/` folder with filtered scatter plots for variables listed as “top variables by bad-count”. +- Useful flags: + - `--no-top-bad-plots`: skip generating `analysis/top_bad_plots/`. + - `--bad-html-limit N`, `--bad-text-limit N`, `--no-export-bad`. + - `--force-xlim-01/--no-force-xlim-01`, `--print-scatter-stats/--no-print-scatter-stats`. + +**Output**: `analysis/` contains CSV/PNG/HTML reports plus `top_bad_plots/`. **Quality Classification**: - **Good**: R² ≥ 0.9, Relative RMSE ≤ 0.1, Relative MAE ≤ 0.1 @@ -63,14 +113,25 @@ Creates a NetCDF file containing all AI predictions for plotting and comparison: ```bash python ../../scripts/ai_predictions_to_netcdf.py > ai_prediction_to_netcdf.log 2>&1 & ``` +To verify the CSV predictions match the generated NetCDF: +```bash +python ../../scripts/verify_predictions_in_netcdf.py +``` ### 7) Generate comparison plots Creates map plots for the variables in your CNP_IO list under -`./ai_model_comparison_plots/comparison_CNP_IO_demo1`. +`./ai_model_comparison_plots/comparison_CNP_IO_demo`. ```bash python ../../scripts/ai_model_comparison_plot.py > ai_model_comparison.log 2>&1 & ``` +#### 7.1) Stats-only mode (no plotting) +To validate that the variable values from the CSV file are consistent with those in the NetCDF file, you can compare their summary statistics (sum, standard deviation, minimum, and maximum) for all variables and their corresponding layers/PFTs using the --stats-only and --variable-list options: +```bash +python ../../scripts/ai_model_comparison_plot.py \ + --stats-only \ + --variable-list ../../CNP_IO_updated9_dev.txt +``` ### 8) Create a new ELM restart file using AI predictions @@ -81,6 +142,31 @@ python ../../scripts/ai_predictions_to_restart.py > ai_predictions_to_restart.lo Outputs a new restart file derived from `original_20250408_trendytest_ICB1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc`. +#### 8.1) Tropical-only merge into a global restart file +If your inference was generated with `--tropical-only`, you can overwrite only tropical gridcells in a global restart file while keeping non-tropical regions unchanged. + +Recommended: run preview first (no file write): +```bash +python ../../scripts/ai_predictions_to_restart.py \ + --restart-file /path/to/global_restart.nc \ + --output ./updated_restart_tropical_merge.nc \ + --merge-scope tropical-only \ + --preview-only +``` + +Then run without preview to write output: +```bash +python ../../scripts/ai_predictions_to_restart.py \ + --restart-file /path/to/global_restart.nc \ + --output ./updated_restart_tropical_merge.nc \ + --merge-scope tropical-only +``` + +Notes: +- Default tropical range is `-23.5,23.5` (`--tropical-lat-range` to override). +- Coordinate matching tolerance defaults to `1e-4` (`--coord-tol` to override). +- Confirm log line `Overwrite scope: tropical-only (eligible gridcells: ...)` before production runs. + ### 9) Compare restart files Compares selected layers and PFTs; optionally verify with `restart_variable_plot.py`. @@ -88,7 +174,13 @@ Compares selected layers and PFTs; optionally verify with `restart_variable_plot python ../../scripts/ai_restart_comparison.py --variable-list ../../CNP_IO_demo1.txt --layers 0,5,9 --pfts 0,1,2,3,4,5 # Optionally use ../../scripts/restart_variable_plot.py for manual verification ``` - +#### 9.1) Stats-only mode (no plotting) +To validate that the variable values in the predicted NetCDF file are consistent with those in the updated restart file, you can compare their summary statistics (sum, standard deviation, minimum, and maximum) for all variables and their corresponding layers/PFTs using the --stats-only and --variable-list options: +```bash +python ../../scripts/ai_restart_comparison.py \ + --stats-only \ + --variable-list ../../CNP_IO_updated9_dev.txt +``` --- ### Old scripts (to be double-checked) @@ -104,3 +196,34 @@ python ../../scripts/update_restart_with_aipredictions.py \ Optional: `inference_map.py` can create statistics and interactive maps (currently not working well). +### Run Inference on the TVA Dataset +The TVA workflow allows you to perform site-specific inference using a trained model and generate updated restart files for targeted locations. Given one or more geographic coordinates, the script extracts all required variables from the TVA dataset, runs model inference to obtain predicted values, and produces a new restart file reflecting the AI-updated state. + +All codes are organized under the ./TVA_1_Sample folder. +Geographic coordinates are defined in locations.csv. +```bash +python ./TVA_1_Sample/run_workflow.py \ + --restart-file /path/to/20year_restart_file.nc \ + --model-path /path/to/trained_model_TVA.pt \ + --dataset-root /path/to/TVA_dataset \ + --variable-list ../CNP_IO_updated9_dev.txt +``` + +### Extract Single Point Data from ELM Restart File +The `extract_elm_restart_point.py` script extracts all data for specified geographic coordinates from a global ELM restart NetCDF file. It correctly handles ELM's multi-level hierarchical structure (gridcell → topounit → landunit → column → pft) and creates a subset NetCDF file containing only the data for the target location. + +**Key Features**: +- Extracts data at all hierarchical levels for a single geographic point +- Automatically finds the nearest neighbor gridcell to target coordinates +- Preserves all metadata and encoding information from the original file +- Supports both 0-360° and -180 to 180° longitude formats +- Can be used via command-line arguments or with hardcoded defaults + +**Usage**: +```bash +python scripts/extract_elm_restart_point.py \ + --restart-file /path/to/input_restart_file.nc \ + --lat 35.833332(Target latitude coordinate) \ + --lon -84.208336(Target longitude coordinate) \ + --output-file single_point_extracted.nc +``` \ No newline at end of file diff --git a/scripts/ai_predictions_to_restart.py b/scripts/ai_predictions_to_restart.py index e9b64d9..038dc6c 100644 --- a/scripts/ai_predictions_to_restart.py +++ b/scripts/ai_predictions_to_restart.py @@ -25,6 +25,7 @@ import sys import json import re +import math # Project imports sys.path.append(str(Path(__file__).resolve().parents[1])) @@ -57,6 +58,72 @@ def _build_gridcell_groups(one_d_to_grid, n_grid): return groups +def _parse_lat_range(lat_range_text: str) -> tuple[float, float]: + """Parse latitude range string in 'min,max' format.""" + parts = [p.strip() for p in str(lat_range_text).split(",")] + if len(parts) != 2: + raise ValueError(f"Invalid latitude range '{lat_range_text}'. Expected format: min,max") + lat_min, lat_max = float(parts[0]), float(parts[1]) + if lat_min > lat_max: + lat_min, lat_max = lat_max, lat_min + return lat_min, lat_max + + +def _wrap_lon(lon: float) -> float: + """Wrap longitude to [-180, 180) for stable coordinate matching.""" + return ((float(lon) + 180.0) % 360.0) - 180.0 + + +def _coord_key(lon: float, lat: float, decimals: int) -> tuple[float, float]: + return (round(_wrap_lon(lon), decimals), round(float(lat), decimals)) + + +def build_update_grid_mask( + variable_mapping: Dict[str, Any], + merge_scope: str, + tropical_lat_range: tuple[float, float], + coord_tol: float +) -> np.ndarray: + """Build per-gridcell update mask for restart overwrite.""" + n_grid = int(variable_mapping["n_grid"]) + mask_all = np.ones(n_grid, dtype=bool) + if merge_scope == "all": + print("Merge scope: all model gridcells (backward-compatible behavior)") + return mask_all + + model_lon = np.asarray(variable_mapping["model_lon"], dtype=float) + model_lat = np.asarray(variable_mapping["model_lat"], dtype=float) + ai_lon = np.asarray(variable_mapping["ai_lon"], dtype=float) + ai_lat = np.asarray(variable_mapping["ai_lat"], dtype=float) + + lat_min, lat_max = tropical_lat_range + in_tropical = np.isfinite(model_lat) & (model_lat >= lat_min) & (model_lat <= lat_max) + + tol = float(coord_tol) + if not np.isfinite(tol) or tol <= 0: + tol = 1e-6 + decimals = max(0, int(math.ceil(-math.log10(tol)))) + + ai_coord_keys = set() + for lon, lat in zip(ai_lon, ai_lat): + if np.isfinite(lon) and np.isfinite(lat): + ai_coord_keys.add(_coord_key(lon, lat, decimals)) + + has_ai_match = np.zeros(n_grid, dtype=bool) + for g in range(n_grid): + if not (np.isfinite(model_lon[g]) and np.isfinite(model_lat[g])): + continue + has_ai_match[g] = _coord_key(model_lon[g], model_lat[g], decimals) in ai_coord_keys + + update_mask = in_tropical & has_ai_match + print(f"Merge scope: tropical-only (lat in [{lat_min}, {lat_max}])") + print(f"Coordinate tolerance for matching: {tol} (rounded decimals: {decimals})") + print(f" Tropical model gridcells: {int(in_tropical.sum())}/{n_grid}") + print(f" Model gridcells matched in AI coords: {int(has_ai_match.sum())}/{n_grid}") + print(f" Gridcells eligible for overwrite: {int(update_mask.sum())}/{n_grid}") + return update_mask + + def load_datasets(ai_predictions_path: Path, restart_file_path: Path) -> tuple[xr.Dataset, xr.Dataset]: """Load AI predictions and model restart datasets.""" print("Loading datasets...") @@ -91,13 +158,9 @@ def create_spatial_mapping(ds_ai: xr.Dataset, ds_model: xr.Dataset) -> tuple[np. from scipy.spatial.distance import cdist ai_coords = np.column_stack([ai_lon, ai_lat]) model_coords = np.column_stack([model_lon, model_lat]) - distances = cdist(ai_coords, model_coords) - ai_to_model_mapping = np.argmin(distances, axis=1) - - print(f" Spatial mapping created: {len(ai_to_model_mapping)} AI -> {len(set(ai_to_model_mapping))} Model") - print(f" Mapping range: AI gridcell 0->model gridcell {ai_to_model_mapping[0]}") - print(f" Mapping range: AI gridcell {len(ai_lon)-1}->model gridcell {ai_to_model_mapping[-1]}") - + distances = cdist(model_coords, ai_coords) + model_to_ai_mapping = np.argmin(distances, axis=1) # 索引是模型格点, 值是最近的 AI 格点 + print(f" Spatial mapping created: {len(model_to_ai_mapping)} Model -> {len(set(model_to_ai_mapping))} AI") # Get grid information from the MODEL file as the master coordinate system n_grid = ds_model.sizes["gridcell"] print(f" Using MODEL gridcell count: {n_grid}") @@ -116,10 +179,14 @@ def create_spatial_mapping(ds_ai: xr.Dataset, ds_model: xr.Dataset) -> tuple[np. variable_mapping = { 'grid_to_cols': grid_to_cols, 'grid_to_pfts': grid_to_pfts, - 'n_grid': n_grid + 'n_grid': n_grid, + 'ai_lon': ai_lon, + 'ai_lat': ai_lat, + 'model_lon': model_lon, + 'model_lat': model_lat } - return ai_to_model_mapping, variable_mapping + return model_to_ai_mapping, variable_mapping def auto_detect_variable_list(ai_predictions_path: Path) -> list: @@ -148,8 +215,9 @@ def auto_detect_variable_list(ai_predictions_path: Path) -> list: def create_updated_restart_file(restart_file_path: Path, output_path: Path, ai_predictions_path: Path, cnp_io_variables: List[str], - ai_to_model_mapping: np.ndarray, variable_mapping: Dict[str, Any]) -> None: - """Directly update the restart file using netCDF4 without xarray encoding issues.""" + model_to_ai_mapping: np.ndarray, variable_mapping: Dict[str, Any], + strict_dims: bool = False, + update_grid_mask: Optional[np.ndarray] = None) -> None: print(f"Saving updated restart file to: {output_path}") # Create output directory if it doesn't exist @@ -161,8 +229,65 @@ def create_updated_restart_file(restart_file_path: Path, output_path: Path, # Open the output file for direct modification with nc.Dataset(output_path, 'r+') as ds_out: + # Verify and adjust spinup_state + try: + if 'spinup_state' in ds_out.variables: + spin_var = ds_out.variables['spinup_state'] + try: + orig_val = np.array(spin_var[:]).item() if spin_var.size == 1 else None + except Exception: + orig_val = None + if orig_val is not None: + print(f"spinup_state in original restart (copied): {orig_val}") + if orig_val != 1: + print("Warning: Expected spinup_state==1 for adspinup; proceeding to set final_spinup (0) anyway") + else: + print("Warning: Could not read scalar value of spinup_state; proceeding to set to 0") + # Set to final_spinup mode (0) + try: + spin_var[...] = 0 + print("Set spinup_state to 0 (final_spinup) in updated restart") + except Exception as e: + print(f"Warning: Failed to set spinup_state to 0: {e}") + else: + print("Warning: 'spinup_state' variable not found in restart; skipping spinup flag update") + except Exception as e: + print(f"Warning: spinup_state check/update failed: {e}") # Load AI predictions with nc.Dataset(ai_predictions_path, 'r') as ds_ai: + # Helpers for shape/dimension checks + def _fail_or_warn(msg: str) -> bool: + if strict_dims: + raise ValueError(msg) + print(f"Warning: {msg} — skipping this variable") + return False + + def _check_pft_compat(ai_var: nc.Variable, model_var: nc.Variable) -> bool: + # Expect AI dims to include pft and gridcell + ai_dims = list(ai_var.dimensions) + if not ('pft' in ai_dims and 'gridcell' in ai_dims): + return _fail_or_warn(f"PFT var '{ai_var.name}' missing required dims (has {ai_dims}, need ['pft','gridcell'])") + # Model var should be 1D over pfts1d (or equivalent) + if len(model_var.shape) < 1: + return _fail_or_warn(f"Model PFT var '{model_var.name}' has invalid shape {model_var.shape}") + # Require at least 16 PFT slots (PFT1..PFT16). We skip PFT0 by design. + if model_var.shape[0] < 16: + return _fail_or_warn(f"Model PFT var '{model_var.name}' has insufficient length {model_var.shape[0]} (<16)") + return True + + def _check_soil_compat(ai_var: nc.Variable, model_var: nc.Variable) -> bool: + # Expect AI dims: (column, levgrnd, gridcell) + ai_dims = list(ai_var.dimensions) + required = {'column','levgrnd','gridcell'} + if not required.issubset(set(ai_dims)): + return _fail_or_warn(f"Soil var '{ai_var.name}' missing required dims (has {ai_dims}, need {sorted(required)})") + if len(model_var.shape) < 2: + return _fail_or_warn(f"Model soil var '{model_var.name}' has invalid shape {model_var.shape}") + # Need at least 10 layers in model to write top 10 + if model_var.shape[1] < 10: + return _fail_or_warn(f"Model soil var '{model_var.name}' has insufficient levgrnd={model_var.shape[1]} (<10)") + return True + # Update PFT variables for var_name in ds_ai.variables: if (var_name in ds_out.variables and @@ -174,33 +299,31 @@ def create_updated_restart_file(restart_file_path: Path, output_path: Path, print(f" AI data shape: {ai_data.shape}") print(f" Model variable shape: {model_var.shape}") + # Dimension compatibility check + if not _check_pft_compat(ds_ai.variables[var_name], model_var): + continue + # Get the grid-to-pfts mapping if 'grid_to_pfts' in variable_mapping: grid_to_pfts = variable_mapping['grid_to_pfts'] # For each model gridcell, update PFT data for g in range(variable_mapping['n_grid']): + if update_grid_mask is not None and not bool(update_grid_mask[g]): + continue if g < len(grid_to_pfts) and len(grid_to_pfts[g]) > 0: # Get PFTs in this gridcell gridcell_pfts = grid_to_pfts[g] - # Find corresponding AI gridcell using spatial mapping - ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] - if len(ai_gridcell_idx) > 0: - ai_gridcell_idx = ai_gridcell_idx[0] - - # Update each PFT instance in this gridcell - # Only update PFT1-PFT16 (skip PFT0), and only in the first column - # Get the first 16 PFTs in this gridcell (exact same as working script) - gridcell_pfts = gridcell_pfts[:16] # First 16 PFTs - for pft_idx, model_pft_idx in enumerate(gridcell_pfts): - # Skip PFT0 (index 0), start from PFT1 (index 1) - if 1 <= pft_idx <= 16 and model_pft_idx < len(model_var): - # AI PFT0 -> Model PFT1, AI PFT1 -> Model PFT2, etc. - # Adjust index: AI PFT k corresponds to Model PFT (k+1) in the first 16 - adjusted_k = pft_idx - 1 # AI PFT0 -> Model PFT1, AI PFT1 -> Model PFT2 - if adjusted_k < ai_data.shape[0]: - model_var[model_pft_idx] = ai_data[adjusted_k, ai_gridcell_idx] + + ai_gridcell_idx = model_to_ai_mapping[g] + gridcell_pfts = gridcell_pfts[:16] + for pft_idx, model_pft_idx in enumerate(gridcell_pfts): + # Skip PFT0 (index 0), start from PFT1 (index 1) + if 1 <= pft_idx <= 16 and model_pft_idx < len(model_var): + adjusted_k = pft_idx - 1 + if adjusted_k < ai_data.shape[0]: + model_var[model_pft_idx] = ai_data[adjusted_k, ai_gridcell_idx] # Update soil variables for var_name in ds_ai.variables: @@ -214,38 +337,38 @@ def create_updated_restart_file(restart_file_path: Path, output_path: Path, print(f" AI data shape: {ai_data.shape}") print(f" Model variable shape: {model_var.shape}") + # Dimension compatibility check + if not _check_soil_compat(ds_ai.variables[var_name], model_var): + continue + # Get the grid-to-cols mapping if 'grid_to_cols' in variable_mapping: grid_to_cols = variable_mapping['grid_to_cols'] # For each model gridcell, update column data for g in range(variable_mapping['n_grid']): + if update_grid_mask is not None and not bool(update_grid_mask[g]): + continue if g < len(grid_to_cols) and len(grid_to_cols[g]) > 0: # Get columns in this gridcell gridcell_cols = grid_to_cols[g] - # Find corresponding AI gridcell using spatial mapping - ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] - if len(ai_gridcell_idx) > 0: - ai_gridcell_idx = ai_gridcell_idx[0] - - # Update first column in this gridcell (use AI column 0) - if len(gridcell_cols) > 0: - model_col_idx = gridcell_cols[0] # First column of this gridcell - if model_col_idx < model_var.shape[0]: - # Update only first 10 layers for this column (even if model has 15 layers) - layers_to_update = min(10, ai_data.shape[1]) - for layer_idx in range(layers_to_update): - # Handle AI data indexing - shape is (column, levgrnd, gridcell) - if ai_data.ndim == 3: - model_var[model_col_idx, layer_idx] = ai_data[0, layer_idx, ai_gridcell_idx] - else: - model_var[model_col_idx, layer_idx] = ai_data[0, layer_idx] + + ai_gridcell_idx = model_to_ai_mapping[g] + + if len(gridcell_cols) > 0: + model_col_idx = gridcell_cols[0] + if model_col_idx < model_var.shape[0]: + layers_to_update = min(10, ai_data.shape[1]) + for layer_idx in range(layers_to_update): + if ai_data.ndim == 3: + model_var[model_col_idx, layer_idx] = ai_data[0, layer_idx, ai_gridcell_idx] + else: + model_var[model_col_idx, layer_idx] = ai_data[0, layer_idx] print(f"Updated restart file saved successfully!") print(f"File size: {output_path.stat().st_size / (1024*1024):.1f} MB") - def get_varlist_name_from_config(ai_predictions_path): for parent in [ai_predictions_path.parent] + list(ai_predictions_path.parents): config_path = parent / 'cnp_config.json' @@ -321,6 +444,14 @@ def main(): help='Preview changes without saving updated restart file') parser.add_argument('--backup', action='store_true', help='Create backup of original restart file before updating') + parser.add_argument('--strict-dims', action='store_true', + help='Abort on any dimension mismatch instead of skipping') + parser.add_argument('--merge-scope', choices=['all', 'tropical-only'], default='all', + help='Overwrite scope: all model gridcells (default) or tropical-only') + parser.add_argument('--tropical-lat-range', type=str, default='-23.5,23.5', + help='Latitude range used when --merge-scope tropical-only, format "min,max"') + parser.add_argument('--coord-tol', type=float, default=1e-4, + help='Coordinate match tolerance for tropical-only merge') args = parser.parse_args() @@ -350,6 +481,7 @@ def main(): print(f"Output: {output_path}") print(f"Preview only: {args.preview_only}") print(f"Create backup: {args.backup}") + print(f"Merge scope: {args.merge_scope}") print("=" * 60) # Load datasets @@ -358,6 +490,19 @@ def main(): # Create spatial mapping ai_to_model_mapping, variable_mapping = create_spatial_mapping(ds_ai, ds_model) + tropical_lat_range = (-23.5, 23.5) + if args.merge_scope == 'tropical-only': + try: + tropical_lat_range = _parse_lat_range(args.tropical_lat_range) + except Exception as e: + parser.error(f"Invalid --tropical-lat-range: {e}") + update_grid_mask = build_update_grid_mask( + variable_mapping=variable_mapping, + merge_scope=args.merge_scope, + tropical_lat_range=tropical_lat_range, + coord_tol=args.coord_tol + ) + # Parse CNP_IO list if provided, else auto-detect cnp_io_variables = [] if args.variable_list: @@ -428,6 +573,7 @@ def main(): print(f" AI gridcells: {ds_ai.sizes.get('gridcell', 'N/A')}") print(f" Model gridcells: {ds_model.sizes.get('gridcell', 'N/A')}") print(f" Spatial mapping: {len(ai_to_model_mapping)} AI -> {len(set(ai_to_model_mapping))} Model") + print(f" Effective overwrite gridcells: {int(np.count_nonzero(update_grid_mask))}") if not args.preview_only: # Create backup if requested @@ -439,7 +585,8 @@ def main(): # Save updated restart file using direct NetCDF manipulation create_updated_restart_file(restart_file_path, output_path, ai_predictions_path, - cnp_io_variables, ai_to_model_mapping, variable_mapping) + cnp_io_variables, ai_to_model_mapping, variable_mapping, + strict_dims=args.strict_dims, update_grid_mask=update_grid_mask) print(f"\nRestart file updated successfully!") print(f"Original: {restart_file_path}") @@ -454,6 +601,7 @@ def main(): print(f" Other columns: Preserved (not modified)") print(f" Spatial mapping: Used geographic coordinates to map AI gridcells to model gridcells") print(f" Coordinate system: Model coordinates used as master reference for alignment") + print(f" Overwrite scope: {args.merge_scope} (eligible gridcells: {int(np.count_nonzero(update_grid_mask))})") print(f" Important: Only CNP_IO variables were modified - all other variables and attributes unchanged") print(f"\nYou can now use the updated restart file for model simulations!") @@ -467,4 +615,4 @@ def main(): if __name__ == '__main__': - main() + main() \ No newline at end of file