diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 138aa5a..e1f83d0 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -142,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`. diff --git a/scripts/ai_predictions_to_restart.py b/scripts/ai_predictions_to_restart.py index 320f2f6..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...") @@ -112,7 +179,11 @@ 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 model_to_ai_mapping, variable_mapping @@ -145,7 +216,8 @@ 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], model_to_ai_mapping: np.ndarray, variable_mapping: Dict[str, Any], - strict_dims: bool = False) -> None: + 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 @@ -237,6 +309,8 @@ def _check_soil_compat(ai_var: nc.Variable, model_var: nc.Variable) -> bool: # 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] @@ -273,6 +347,8 @@ def _check_soil_compat(ai_var: nc.Variable, model_var: nc.Variable) -> bool: # 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] @@ -370,6 +446,12 @@ def main(): 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() @@ -399,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 @@ -407,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: @@ -477,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 @@ -489,7 +586,7 @@ 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, - strict_dims=args.strict_dims) + strict_dims=args.strict_dims, update_grid_mask=update_grid_mask) print(f"\nRestart file updated successfully!") print(f"Original: {restart_file_path}") @@ -504,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!") @@ -517,4 +615,4 @@ def main(): if __name__ == '__main__': - main() + main() \ No newline at end of file