diff --git a/CNP_IO_5P_tworegions.txt b/CNP_IO_5P_tworegions.txt new file mode 100644 index 0000000..74054bc --- /dev/null +++ b/CNP_IO_5P_tworegions.txt @@ -0,0 +1,52 @@ +# Dataset roots (any absolute paths) +TRENDY1_PATH: /mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_1_data_CNP +TRENDY05_PATH: /mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP +TVA4KM_PATH: /mnt/proj-shared/AI4BGC_7xw/TrainingData/TVA_4km_data_CNP + +# Optional extra roots (comma-separated) +DATA_PATHS: /another/path1, /another/path2 + +# Global fallback pattern if a dataset-specific one isn't set +#FILE_PATTERN: enhanced_1_training_data_batch_*.pkl +#FILE_PATTERN: 1_training_data_batch_*.pkl +FILE_PATTERN: training_data_batch_*.pkl + +# Per-dataset patterns (overrides FILE_PATTERN for that path only) +TVA4KM_FILE_PATTERN: enhanced_monthly_training_data_batch_*.pkl + + +#LONGITUDE FILTERING - 2 longitudes: +#• 0, 358.75 + +TIME SERIES VARIABLES (Climate Forcing) - 6 variables: +• FLDS, PSRF, FSDS, QBOT, PRECTmms, TBOT + +SURFACE PROPERTIES - 49 variables: +• Latitude, Longitude, AREA, landfrac, LANDFRAC_PFT, PCT_NATVEG + +• OCCLUDED_P, SECONDARY_P, LABILE_P, APATITE_P + +• SOIL_COLOR, SOIL_ORDER + +• PCT_NAT_PFT_0, PCT_NAT_PFT_1, PCT_NAT_PFT_2, PCT_NAT_PFT_3, PCT_NAT_PFT_4, PCT_NAT_PFT_5, PCT_NAT_PFT_6, PCT_NAT_PFT_7, PCT_NAT_PFT_8 +• PCT_NAT_PFT_9, PCT_NAT_PFT_10, PCT_NAT_PFT_11, PCT_NAT_PFT_12, PCT_NAT_PFT_13, PCT_NAT_PFT_14, PCT_NAT_PFT_15, PCT_NAT_PFT_16 + +• PCT_CLAY_0, PCT_CLAY_1, PCT_CLAY_2, PCT_CLAY_3, PCT_CLAY_4, PCT_CLAY_5, PCT_CLAY_6, PCT_CLAY_7, PCT_CLAY_8, PCT_CLAY_9 +• PCT_SAND_0, PCT_SAND_1, PCT_SAND_2, PCT_SAND_3, PCT_SAND_4, PCT_SAND_5, PCT_SAND_6, PCT_SAND_7, PCT_SAND_8, PCT_SAND_9 + +PFT PARAMETERS (Plant Functional Type Characteristics) - 44 variables: + +• pft_deadwdcn, pft_frootcn, pft_leafcn, pft_lflitcn, pft_livewdcn, pft_c3psn, pft_croot_stem, pft_crop, pft_dleaf +• pft_dsladlai, pft_evergreen, pft_fcur, pft_flivewd, pft_flnr, pft_fr_fcel, pft_fr_flab, pft_fr_flig, pft_froot_leaf +• pft_grperc, pft_grpnow, pft_leaf_long, pft_lf_fcel, pft_lf_flab, pft_lf_flig, pft_rholnir, pft_rholvis, pft_rhosnir, pft_rhosvis +• pft_roota_par, pft_rootb_par, pft_rootprof_beta, pft_season_decid, pft_slatop, pft_smpsc, pft_smpso, pft_stem_leaf, pft_stress_decid +• pft_taulnir, pft_taulvis, pft_tausnir, pft_tausvis, pft_woody, pft_xl, pft_z0mr + +SCALAR VARIABLES (1D - 4 variables): +• GPP + +1D PFT VARIABLES (41 variables): + +2D VARIABLES (layered - 25 variables): + +• labilep_vr , occlp_vr, primp_vr, secondp_vr, solutionp_vr diff --git a/CNP_IO_updated9_dev_dw.txt b/CNP_IO_updated9_dev_dw.txt index c1f2f8d..6a87fa3 100644 --- a/CNP_IO_updated9_dev_dw.txt +++ b/CNP_IO_updated9_dev_dw.txt @@ -8,14 +8,15 @@ DATA_PATHS: /another/path1, /another/path2 # Global fallback pattern if a dataset-specific one isn't set #FILE_PATTERN: enhanced_1_training_data_batch_*.pkl -FILE_PATTERN: 1_training_data_batch_*.pkl +#FILE_PATTERN: 1_training_data_batch_*.pkl +FILE_PATTERN: training_data_batch_*.pkl # Per-dataset patterns (overrides FILE_PATTERN for that path only) TVA4KM_FILE_PATTERN: enhanced_monthly_training_data_batch_*.pkl -LONGITUDE FILTERING - 2 longitudes: -• 0, 358.75 +#LONGITUDE FILTERING - 2 longitudes: +#• 0, 358.75 TIME SERIES VARIABLES (Climate Forcing) - 6 variables: • FLDS, PSRF, FSDS, QBOT, PRECTmms, TBOT @@ -63,9 +64,9 @@ SCALAR VARIABLES (1D - 4 variables): • cwdc_vr, cwdn_vr, cwdp_vr -• litr1c_vr, litr2c_vr, litr3c_vr -• litr1n_vr, litr2n_vr, litr3n_vr -• litr1p_vr, litr2p_vr, litr3p_vr +• litr2c_vr, litr3c_vr +• litr2n_vr, litr3n_vr +• litr2p_vr, litr3p_vr • soil1c_vr, soil1n_vr, soil1p_vr • soil2c_vr, soil2n_vr, soil2p_vr diff --git a/commands.txt b/commands.txt index ba2ae04..b33e87c 100644 --- a/commands.txt +++ b/commands.txt @@ -166,4 +166,357 @@ python train_cnp_model.py \ python train_cnp_model.py \ --variable-list CNP_IO_updated9_dev_dw_repeat.txt \ --training-config config/training_config_experiment_2_repeat.json \ - --tropical-only \ No newline at end of file + --tropical-only + + python ../../scripts/run_inference_all.py --model model.pth --output-dir cnp_inference_entire_dataset --derive-np-from-c > run_inference_all.log 2>&1 & + + + python train_cnp_model.py \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --training-config config/training_config_experiment_2_global.json \ + 2>&1 | tee cnp_experiment2_global.log & + + python ../../scripts/cnp_result_validationplot.py --stats-only && \ + python ../../scripts/generate_prediction_quality_report.py + + python ../../scripts/run_inference_all.py --model model.pth --output-dir cnp_inference_entire_dataset --derive-np-from-c > run_inference_all.log 2>&1 & + + + python train_cnp_model.py \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --training-config config/training_config_experiment_3_global.json \ + 2>&1 | tee cnp_experiment3_global.log & + + python ../../scripts/run_inference_all.py --model model.pth --output-dir cnp_inference_entire_dataset --derive-np-from-c > run_inference_all.log 2>&1 & + + ln -s /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260219_132504_experiment3_global_repeat_unifiedconfig/cnp_inference_entire_dataset_cnpratio cnp_inference_entire_dataset + + python ../../scripts/ai_predictions_to_netcdf.py > ai_prediction_to_netcdf.log 2>&1 + + python ../../scripts/ai_model_comparison_plot.py \ + --ai-predictions /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260219_132504_experiment3_global_repeat_unifiedconfig/comparison_results/ai_predictions_for_plotting_cnpratio.nc \ + --output-dir ./ai_cnpratio_model_comparison_plots + + compare with CSV file by default (quality check) + + + python ../../scripts/ai_model_comparison_plot.py \ + --ai-predictions comparison_results/ai_predictions_for_plotting_cnpratio.nc \ + --model ../../ELM_data/20251201_TRENDY2024_default_ICB1850CNPRDCTCBC.elm.r.0801-01-01-00000.nc \ + --output-dir ./ai_vs_restart_plots + + +python ../../scripts/ai_model_comparison_plot.py \ + --ai-predictions comparison_results/ai_predictions_for_plotting_cnpratio.nc \ + --model ../../ELM_data/20251201_TRENDY2024_default_ICB1850CNPRDCTCBC.elm.r.0801-01-01-00000.nc \ + --output-dir ./ai_vs_restart_all_plots \ + --variable-list ../../CNP_IO_updated9_dev_dw.txt \ + > ai_model_all_comparison.log 2>&1 & + + +cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260219_132504_experiment3_global_repeat_unifiedconfig/cnp_inference_entire_dataset + +# Copy config from parent so validation can read variable lists +cp ../cnp_config.json . 2>/dev/null || true + +# Same as runbook: current dir is the “results” dir +python ../../../scripts/cnp_result_validationplot.py . --stats-only && \ +python ../../../scripts/generate_prediction_quality_report.py + +cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC/ELM_data + +python ../scripts/ai_model_comparison_plot.py \ + --model1 ./20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --model2 ./20251201_TRENDY2024_default_ICB1850CNPRDCTCBC.elm.r.0801-01-01-00000.nc \ + --variable-list ../CNP_IO_updated9_dev_dw.txt \ + --output-dir ./adspin20_vs_final800_all_plots \ + --layers 0 1 2 3 4 \ + --pfts 1 2 3 4 5 6 \ + > adspin20_spin800_all_comparison.log 2>&1 & + + + python train_cnp_model.py \ + --model-config CNP_model_config_multimode_tropicaltest1.txt \ + --variable-list CNP_IO_tropicaltest1.txt \ + --training-config config/training_config_experiment_2.json \ + --tropical-only + + +python scripts/cnp_result_validationplot.py /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260220_145518 --natveg-filter + + + +python ../../scripts/cnp_result_validationplot.py --stats-only > validationplot.log 2>&1 && \ +python ../../scripts/generate_prediction_quality_report.py > prediction_quality_report.log 2>&1 && \ +python ../../scripts/run_inference_all.py --model model.pth --output-dir cnp_inference_entire_dataset --derive-np-from-c > run_inference_all.log 2>&1 && \ +python ../../scripts/ai_predictions_to_netcdf.py > ai_prediction_to_netcdf.log 2>&1 + +python ../../scripts/ai_predictions_to_restart.py \ + --output CNInit_updated9_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + > ai_predictions_to_restart.log 2>&1 + +python ../../scripts/ai_predictions_to_restart.py \ + --output CNInit_natveg_updated9_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + > ai_predictions_to_restart.log 2>&1 & + + + python train_cnp_model.py \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --training-config config/training_config_experiment_3_tropical.json \ + > cnp_experiment3_tropical.log 2>&1 + +comparison runs + python train_cnp_model.py \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --training-config config/training_config_experiment_3_global_nofilter.json + + python train_cnp_model.py \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --training-config config/training_config_experiment_3_global_natveg_aligned.json + + python train_cnp_model.py \ + --variable-list CNP_IO_updated9_intermediate_dw.txt \ + --training-config config/training_config_experiment_3_global_nofilter.json + +cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC + +python scripts/ai_restart_comparison.py \ + --variable-list CNP_IO_updated9_dev.txt \ + --ai-restart cnp_results/run_20260221_120117_withCNonlyInit/CNInit_updated9_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --original-restart ELM_data/20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --plot-all \ + --output-dir cnp_results/run_20260221_120117_withCNonlyInit/ai_restart_comparison_plots + +python scripts/ai_restart_comparison.py \ + --variable-list CNP_IO_updated9_dev.txt \ + --ai-restart cnp_results/run_20260221_120117_withCNonlyInit/CNInit_updated9_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --original-restart ELM_data/20251201_TRENDY2024_default_ICB1850CNPRDCTCBC.elm.r.0801-01-01-00000.nc \ + --plot-all \ + --output-dir cnp_results/run_20260221_120117_withCNonlyInit/ai_restart_comparison_plots_vs_target800result + +# Run folder: cnp_results/run_YYYYMMDD_HHMMSS_natveg_improved (--output-dir defaults to cnp_results). +python train_cnp_model.py \ + --training-config-json config/training_config_experiment_3_global_natveg_improved.json \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --output-dir-suffix natveg_improved + + # Minimal pipeline: validation + per-PFT npool/ppool + quality report +# Run from inside a run directory, e.g. cnp_results/run_YYYYMMDD_HHMMSS_suffix + +# 1) Validation stats (writes validation_stats.csv in run dir) +python ../../scripts/cnp_result_validationplot.py --stats-only + +# 2) Per-PFT npool/ppool validation (so npool/ppool not labeled "worst"; creates analysis/ if needed) +python ../../scripts/validation_npool_ppool_per_pft.py . -o analysis/npool_ppool_per_pft.json + +# 3) Quality report + top_bad_plots (excludes npool/ppool from worst list; writes analysis/ and analysis/top_bad_plots/) +# Add --training-config cnp_config.json to include all CNP_IO variables in report/chart (pad missing with zeros). +python ../../scripts/generate_prediction_quality_report.py --input validation_stats.csv --npool-ppool-per-pft-json analysis/npool_ppool_per_pft.json + +# One-liner (from run dir): +# python ../../scripts/cnp_result_validationplot.py --stats-only && python ../../scripts/validation_npool_ppool_per_pft.py . -o analysis/npool_ppool_per_pft.json && python ../../scripts/generate_prediction_quality_report.py --input validation_stats.csv --npool-ppool-per-pft-json analysis/npool_ppool_per_pft.json + + +python /mnt/proj-shared/AI4BGC_7xw/AI4BGC/scripts/extract_elm_restart_point.py \ + --restart-file /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --lat -17.4246 \ + --lon 303.75 \ + --output-file /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260228_214757_natveg_improved/Amazon_updated_natveg_improved_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024.nc + +# Phase3 tworegions: extract Amazon site from full-grid 5P bias-corrected restart (same pattern as above) +# python scripts/extract_elm_restart_point.py \ +# --restart-file cnp_results/run_20260311_204845_phase3_tworegions/updated_restart_phase3_tworegions_5P_bias_corrected_tropical.nc \ +# --lat -17.434553 \ +# --lon 303.75 \ +# --output-file cnp_results/run_20260311_204845_phase3_tworegions/Amazon_phase3_tworegions_5P_bias_corrected_restart.nc + +two step training targeting tropical and occld_vr +python train_cnp_model.py --training-config-json config/training_config_experiment_3_global_natveg_improved_occlp.json --variable-list CNP_IO_updated9_dev_dw.txt --output-dir-suffix natveg_improved_occlp + +python ../../scripts/run_inference_all.py --model model.pth --output-dir cnp_inference_entire_dataset --derive-np-from-c +python ../../scripts/ai_predictions_to_netcdf.py +python ../../scripts/ai_predictions_to_restart.py \ + --ai-predictions comparison_results/ai_predictions_for_plotting.nc \ + --restart-file ../../ELM_data/20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --output ./updated_restart_global_natveg_improved_occlp.nc \ + --variable-list ../../CNP_IO_updated9_dev_dw.txt + + + python scripts/run_finetuning.py \ + --config CNP_IO_finetune_tropics.txt \ + --tropical-only + + python scripts/run_finetuning_json.py \ + --config-json config/training_config_experiment_3_global_natveg_occlp_phase2_tropics.json + + go to /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results_tropics_natveg_improved_occlp/finetune_20260304_081553 + + python ../../scripts/run_inference_all.py \ + --model /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results_tropics_natveg_improved_occlp/finetune_20260304_081553/finetuned_model.pth \ + --data-paths /mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_1_data_CNP \ + --file-pattern training_data_batch_*.pkl \ + --output-dir cnp_inference_entire_dataset \ + --derive-np-from-c + + + python train_cnp_model.py \ + --training-config-json config/training_config_experiment_3_natveg_occlp_phase2_tropics_from_scratch.json \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --tropical-only \ + --output-dir-suffix natveg_occlp_phase2_tropics_from_scratch + + + python train_cnp_model.py \ + --training-config-json config/training_config_experiment_3_natveg_occlp_extreme_soilp.json \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --output-dir-suffix natveg_occlp_extreme_soilp + + # Amazon site: use absolute --output-file so one file regardless of cwd + python ../../scripts/extract_elm_restart_point.py \ + --restart-file /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260304_133300_natveg_occlp_extreme_soilp/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --lat -17.4246 \ + --lon 303.75 \ + --output-file /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260304_133300_natveg_occlp_extreme_soilp/Amazon_Extreme_SoilP_updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc + + + python train_cnp_model.py \ + --training-config-json config/training_config_phase2_four_p_weight100.json \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --tropical-only \ + --output-dir-suffix phase2_four_p_weight100 + +# Phase2 P-variable merge: base restart (natveg_improved) + phase2 P predictions -> new restart in phase2 dir +# IMPORTANT: Use ENTIRE-DATASET inference (all gridcells), not just cnp_predictions (validation/tropical subset). +# cnp_predictions has ~1064 gridcells (validation from tropical-only training); restart has 20975. Use full-grid inference. +# Step 1: Run inference on entire dataset (full global grid) with Phase 2 model +# cd cnp_results/run_20260305_153217_phase2_pvariable_focus && python ../../scripts/run_inference_all.py \ +# --model cnp_predictions/model.pth \ +# --output-dir cnp_inference_entire_dataset \ +# --inference-full-grid \ +# --derive-np-from-c \ +# > run_inference_entire_dataset.log 2>&1 +# Step 2: Build phase2 AI predictions NetCDF from cnp_inference_entire_dataset (not cnp_predictions) +# python scripts/ai_predictions_to_netcdf.py \ +# --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_inference_entire_dataset/cnp_predictions \ +# --output cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_for_plotting.nc +# Step 3: Update only P variables in base restart, write to phase2 dir +# python scripts/ai_predictions_to_restart.py \ +# --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_for_plotting.nc \ +# --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ +# --output cnp_results/run_20260305_153217_phase2_pvariable_focus/updated_restart_phase2_pvariable_focus_P_from_natveg_base.nc \ +# --variable-list CNP_IO_updated9_dev_dw.txt \ +# --variables-to-update @config/phase2_soilp_variables_to_update.txt + +# Conservative Phase2 merge: tropical-only inference -> NetCDF -> update base restart with 4 P vars in tropical cells only +# 1) Add tropical_only to run cnp_config.json if missing; then run inference (no --inference-full-grid) +# cd cnp_results/run_20260305_153217_phase2_pvariable_focus && python ../../scripts/run_inference_all.py --model cnp_predictions/model.pth --output-dir cnp_inference_tropical_only +# 2) Build NetCDF from tropical predictions +# python scripts/ai_predictions_to_netcdf.py --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_inference_tropical_only/cnp_predictions --output cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_tropical_only.nc +# 3) Update base restart: only occlp_vr, labilep_vr, solutionp_vr, secondp_vr in tropical band (-30,30) +# python scripts/ai_predictions_to_restart.py --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_tropical_only.nc --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc --output cnp_results/run_20260305_153217_phase2_pvariable_focus/updated_restart_phase2_tropical_4p.nc --variable-list CNP_IO_updated9_dev_dw.txt --variables-to-update occlp_vr,labilep_vr,solutionp_vr,secondp_vr "--tropical-lat-range=-30,30" + + +python scripts/compare_restart_pft1d_amazon.py \ + --restart-a ELM_data/Amazon_20251201_TRENDY2024_default_ICB1850CNPRDCTCBC.elm.r.0801-01-01-00000.nc \ + --restart-b cnp_results/run_20260305_153217_phase2_pvariable_focus/Amazon_phase2_tropical_4p_restart.nc \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --output cnp_results/run_20260305_153217_phase2_pvariable_focus/Amazon_restart_PFT1d_comparison_0801_vs_phase2_tropical_4p.txt \ + --plot-dir cnp_results/run_20260305_153217_phase2_pvariable_focus/plots + + +python train_cnp_model.py \ + --training-config-json config/training_config_phase2_pvariable_focus_global.json \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --output-dir-suffix phase2_pvariab_focus_global + + + python scripts/run_finetuning_json.py \ + --config-json config/finetune_phase2_pvariable_focus_two_region.json \ + --output-dir-suffix phase2_pvariab_focus_finetune_with_tworegions + + +use the bias/scale adjustment + +python scripts/apply_5p_bias_scale_correction.py \ + --run-dir cnp_results/run_20260228_214757_natveg_improved \ + --region-config-json config/training_config_two_region_five_p.json + +python scripts/compare_5p_amazon_site_scatter.py + + python scripts/apply_5p_bias_scale_correction.py \ + --run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --region-config-json config/training_config_two_region_five_p.json \ + --variables occlp_vr,labilep_vr,solutionp_vr,secondp_vr \ + --output-subdir soil_2d_predictions_4P_bias_corrected_phase2 + + +python scripts/compare_5p_amazon_site_scatter.py \ + --gt-run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --natveg-run-dir cnp_results/run_20260228_214757_natveg_improved \ + --phase2-run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --phase2-bias-corrected-subdir \ + cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_4P_bias_corrected_phase2 + + +Site A: lon=300.0, lat=4.240838 +Site B: lon=292.5, lat=-15.549738 +Site C: lon=295.0, lat=-8.952880 +Site D: lon=288.75, lat=-4.240838 +Site E: lon=317.5, lat=-12.722513 + + +python scripts/compare_5p_amazon_site_scatter.py \ + --gt-run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --natveg-run-dir cnp_results/run_20260228_214757_natveg_improved \ + --phase2-run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --phase2-bias-corrected-subdir \ + cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_4P_bias_corrected_phase2 \ + --lon 300.0 \ + --lat 4.240838 \ + --output-dir cnp_results/run_20260228_214757_natveg_improved/analysis/amazon_5p_comparison_bias_correction_site_A + + + python scripts/compare_5p_amazon_site_scatter.py \ + --gt-run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --natveg-run-dir cnp_results/run_20260228_214757_natveg_improved \ + --phase2-run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --phase2-bias-corrected-subdir \ + cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_4P_bias_corrected_phase2 \ + --lon 292.5 \ + --lat -15.549738 \ + --output-dir cnp_results/run_20260228_214757_natveg_improved/analysis/amazon_5p_comparison_bias_correction_site_B + + + + cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC + +# 1) Phase1 global → Africa_phase1_global_restart.nc +python scripts/extract_elm_restart_point.py \ + --restart-file cnp_results/run_20260311_165538_phase1_global/updated_restart_base.nc \ + --lat 0.0 \ + --lon 28.0 \ + --output-file cnp_results/run_20260311_165538_phase1_global/Africa_phase1_global_restart.nc + +# 2) Phase2 tropical → Africa_phase2_tropical_restart.nc +python scripts/extract_elm_restart_point.py \ + --restart-file cnp_results/run_20260311_174843_phase2_tropical/updated_restart_phase2_tropical_5P_raw.nc \ + --lat 0.0 \ + --lon 28.0 \ + --output-file cnp_results/run_20260311_174843_phase2_tropical/Africa_phase2_tropical_restart.nc + +# 3) Phase3 two-regions → Africa_phase3_tworegions_restart.nc +python scripts/extract_elm_restart_point.py \ + --restart-file cnp_results/run_20260311_204845_phase3_tworegions/updated_restart_phase3_tworegions_5P_bias_corrected_tropical.nc \ + --lat 0.0 \ + --lon 28.0 \ + --output-file cnp_results/run_20260311_204845_phase3_tworegions/Africa_phase3_tworegions_restart.nc + +python scripts/generate_site_5p_restart_comparison.py --lon 303.75 --lat -17.434553 --site-name amazon_303_17S +python scripts/generate_site_5p_restart_comparison.py --lon 292.5 --lat -15.55 --site-name amazon_292_15S + +python scripts/generate_site_5p_restart_comparison.py --lon 28 --lat 0 --site-name africa_28_0 + +python scripts/generate_site_5p_restart_comparison.py --lon 14 --lat 0 --site-name africa + + + diff --git a/config/UNIFIED_CONFIG_README.md b/config/UNIFIED_CONFIG_README.md index 87efec1..b946744 100644 --- a/config/UNIFIED_CONFIG_README.md +++ b/config/UNIFIED_CONFIG_README.md @@ -10,6 +10,16 @@ Instead of managing multiple separate JSON files, you can now use a single unifi - Tail-aware weights - PFT zero sparsity weights - PFT1D activation overrides +- **Data filtering** (tropical band, natveg-only) — see `data_filtering_config` below + +### Where to put options: config JSON vs CNP_IO.txt + +| Purpose | Use | +|--------|-----| +| **Variable lists**, **data paths**, **file patterns** | **CNP_IO.txt** (and optional paths in variable list) | +| **Training filters** (tropical, natveg, longitude drop), **hyperparameters**, **loss weights**, **reproducibility** | **config/training_*.json** (unified or per-experiment) | + +Put **natveg filter**, **tropical selection**, and **longitude filtering** in the **training config JSON** under `data_filtering_config`, not in CNP_IO. That way one variable list can be reused with different filter choices, and the run’s saved config records exactly which filters were applied. ## JSON File Format @@ -145,11 +155,16 @@ Settings for reproducibility and data handling: **Note**: CLI arguments (`--split-seed`, `--strict-determinism`, `--train-split`, `--normalization`, `--dropout-p`) take precedence over JSON values. ### `data_filtering_config` -Settings for data filtering: +Settings for which samples are included at training time. **Use the training config JSON here, not CNP_IO.txt**: CNP_IO defines variables and data paths; filtering (tropical, natveg) is a run choice and belongs in config so the same variable list can be reused with different filters. + - `tropical_only`: Filter dataset to tropical latitudes (default: `false`) - `tropical_lat_range`: Latitude range as `[min, max]` or `"min,max"` string (default: `[-23.5, 23.5]`) +- `tropical_lat_column`: Optional column name for latitude (default: auto-detect, e.g. `Latitude`) +- `natveg_only`: If `true`, keep only gridcells with natural vegetation: `PCT_NATVEG > 0` and `PCT_NAT_PFT_0 < 100` (default: `false`). See `docs/EXCLUDED_SAMPLE_ANALYSIS.md` for impact (~33% of samples excluded on typical data). +- `longitudes_to_drop`: List of longitudes (degrees) to exclude from training, e.g. `[0, 358.75]`. Samples whose longitude matches (within tolerance) are dropped. Overrides the same option in CNP_IO. Use 0–360° convention to match data. +- `region_boxes`: Optional list of boxes to **keep**; only gridcells inside at least one box are used. Each box is `[lat_min, lat_max, lon_min, lon_max]` with longitude in **0–360°**. E.g. Amazon + Africa: `[[-30, 10, 270, 330], [-15, 15, 0, 30]]`. See `docs/PLAN_TWO_REGION_FIVE_P_TRAINING.md`. -**Note**: CLI arguments (`--tropical-only`, `--tropical-lat-range`) take precedence over JSON values. +**Note**: CLI arguments (`--tropical-only`, `--tropical-lat-range`, `--natveg-only`, `--longitudes-to-drop`) take precedence over JSON values. JSON `data_filtering_config` overrides values from the CNP_IO variable list (e.g. longitude filtering in CNP_IO). ### `variable_weights` Contains three subsections: @@ -190,7 +205,7 @@ When both unified config and CLI arguments are specified: 1. **Training hyperparameters**: CLI args (`--epochs`, `--batch-size`, `--learning-rate`, `--xsmrpool-loss-weight`, `--litter-*-loss-weight`) override unified config 2. **Reproducibility config**: CLI args (`--split-seed`, `--strict-determinism`, `--train-split`, `--normalization`, `--dropout-p`) override unified config -3. **Data filtering config**: CLI args (`--tropical-only`, `--tropical-lat-range`) override unified config +3. **Data filtering config**: CLI args (`--tropical-only`, `--tropical-lat-range`, `--natveg-only`, `--longitudes-to-drop`) override unified config 4. **PFT mask config**: CLI args (`--mask-absent-pfts` / `--no-mask-absent-pfts`, `--pft-presence-threshold`) override unified config 5. **Tail-aware config**: CLI args (`--tail-aware-loss`, `--tail-aware-eps`) override unified config 6. **Tail-aware weights**: Individual `--tail-aware-weights-json` overrides unified config @@ -227,7 +242,8 @@ This allows you to: }, "data_filtering_config": { "tropical_only": true, - "tropical_lat_range": [-23.5, 23.5] + "tropical_lat_range": [-23.5, 23.5], + "longitudes_to_drop": [0, 358.75] }, "variable_weights": { "pft1d_weights": { diff --git a/config/training_config.py b/config/training_config.py index 81a37c7..6e26279 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -100,8 +100,13 @@ class DataConfig: tropical_only: bool = False tropical_lat_range: Tuple[float, float] = (-23.5, 23.5) tropical_lat_column: Optional[str] = None - - + # Natveg-only filtering: keep only gridcells with PCT_NATVEG > 0 and PCT_NAT_PFT_0 < 100 + natveg_only: bool = False + # When True: filter to natveg before shuffle/split (legacy). When False: split on full data, then + # filter only training set to natveg, so test set is the same as no-filter run (natveg test ⊂ no-filter test). + # Default True preserves legacy behavior when key is missing. + natveg_filter_before_split: bool = True + # File loading limits (for testing) max_files: Optional[int] = None # Maximum number of files to load (None = all files) @@ -112,6 +117,10 @@ class DataConfig: # Longitude filtering - list of longitude values to drop from dataset longitudes_to_drop: List[float] = field(default_factory=list) + # Region boxes: keep only gridcells inside any box. List of (lat_min, lat_max, lon_min, lon_max). + # Longitude in 0-360 convention. None = no region filter (use tropical_only / longitudes_to_drop instead). + region_boxes: Optional[List[Tuple[float, float, float, float]]] = None + @@ -167,6 +176,11 @@ class ModelConfig: # PFT1D activation control pft1d_activation: str = 'abs' # 'abs', 'relu', 'softplus', 'linear' pft1d_activation_overrides: Dict[str, str] = field(default_factory=dict) + + # Multihead/Multimode configuration + use_multihead_attention: bool = False # Enable multi-head attention for mode-specific feature extraction + multihead_num_heads: int = 4 # Number of attention heads for multihead attention + use_mode_specific_heads: bool = False # Enable separate feature extractors for each output mode @dataclass diff --git a/config/training_config_experiment_2.json b/config/training_config_experiment_2.json index 7a98a61..ece485b 100644 --- a/config/training_config_experiment_2.json +++ b/config/training_config_experiment_2.json @@ -23,7 +23,8 @@ }, "data_filtering_config": { "tropical_only": true, - "tropical_lat_range": [-23.5, 23.5] + "tropical_lat_range": [-23.5, 23.5], + "longitudes_to_drop": [0, 358.75] }, "variable_weights": { "pft1d_weights": { diff --git a/config/training_config_experiment_3_global.json b/config/training_config_experiment_3_global.json new file mode 100644 index 0000000..ff3375e --- /dev/null +++ b/config/training_config_experiment_3_global.json @@ -0,0 +1,139 @@ +{ + "_description": "Experiment 3 global: full training config for global run (20,826 samples). Same structure as training_config_experiment_2.json. Tail-aware log1p_huber; primp_vr weight 11; PFT mask and zero sparsity as in experiment_3.", + "training_hyperparameters": { + "num_epochs": 100, + "batch_size": 128, + "learning_rate": 0.0001, + "optimizer_type": "adamw", + "weight_decay": 0.01, + "use_scheduler": true, + "scheduler_type": "cosine", + "scheduler_step_size": 10, + "scheduler_gamma": 0.1, + "xsmrpool_loss_weight": 10, + "litter_c_loss_weight": 1, + "litter_n_loss_weight": 1, + "litter_p_loss_weight": 1, + "scalar_loss_weight": 1, + "vector_loss_weight": 1, + "matrix_loss_weight": 1 + }, + "reproducibility_config": { + "random_seed": 42, + "strict_determinism": false, + "train_split": 0.8, + "normalization": "individual", + "dropout_p": 0.1 + }, + "data_filtering_config": { + "tropical_only": false, + "tropical_lat_range": null, + "longitudes_to_drop": [0, 358.75], + "natveg_only": true, + "natveg_filter_before_split": true + }, + "variable_weights": { + "pft1d_weights": { + "cpool": 6, + "npool": 8, + "ppool": 8, + "totvegc": 6, + "leafc": 6, + "leafn": 6, + "leafp": 6, + "leafc_storage": 5, + "leafn_storage": 5, + "leafp_storage": 5, + "deadstemc": 6, + "deadstemn": 6, + "deadstemp": 5, + "frootc": 4, + "frootn": 4, + "frootp": 4, + "frootc_storage": 3, + "frootn_storage": 3, + "frootp_storage": 3, + "tlai": 4 + }, + "soil2d_weights": { + "primp_vr": 11, + "litr2c_vr": 8, + "litr2p_vr": 8, + "litr2n_vr": 8, + "soil1c_vr": 7, + "soil1n_vr": 7, + "soil1p_vr": 7, + "litr3n_vr": 5, + "labilep_vr": 5, + "secondp_vr": 2.5, + "solutinorganicp_vr": 5, + "litr3p_vr": 5 + }, + "scalar_weights": { + "GPP": 1.5, + "NPP": 1.5, + "AR": 1.2, + "HR": 1.2 + } + }, + "tail_aware_config": { + "loss": "log1p_huber", + "epsilon": 1e-08, + "huber_delta": 1 + }, + "tail_aware_weights": { + "cpool": 6, + "deadstemc": 6, + "deadcrootc": 5, + "livestemc": 5, + "livecrootc": 5, + "npool": 8, + "ppool": 8, + "totvegc": 6, + "leafc": 6, + "leafn": 6, + "leafp": 6, + "leafc_storage": 5, + "leafn_storage": 5, + "leafp_storage": 5, + "deadstemp": 5, + "deadstemn": 6, + "frootc": 4, + "frootn": 4, + "frootp": 4, + "labilep": 5, + "labilep_vr": 5, + "primp_vr": 11, + "litr2c_vr": 8, + "litr2p_vr": 8, + "litr2n_vr": 8, + "litr3c_vr": 5, + "litr3p_vr": 5, + "litr3n_vr": 5, + "soil1c_vr": 7, + "soil1n_vr": 7, + "soil1p_vr": 7 + }, + "pft_mask_config": { + "mask_absent_pfts": true, + "pft_presence_threshold": 0 + }, + "pft_zero_sparsity_config": { + "weight": 1, + "threshold": 1e-08 + }, + "pft_zero_sparsity_weights": { + "cpool": 2, + "deadstemc": 1, + "deadcrootc": 1, + "livestemc": 1, + "livecrootc": 1 + }, + "pft1d_activation_overrides": { + "cpool": "abs", + "deadstemc": "abs", + "deadcrootc": "abs", + "livestemc": "abs", + "livecrootc": "abs" + } +} \ No newline at end of file diff --git a/config/training_config_experiment_3_global_natveg_improved.json b/config/training_config_experiment_3_global_natveg_improved.json new file mode 100644 index 0000000..92c5f12 --- /dev/null +++ b/config/training_config_experiment_3_global_natveg_improved.json @@ -0,0 +1,141 @@ +{ + "_description": "Experiment 3 global NATVEG IMPROVED: Based on natveg_aligned reference (docs/NATVEG_VS_NOFILTER_COMPARISON.md). Same data setup (natveg_only true, natveg_filter_before_split false). Tuning to improve soil 2D (especially occlp_vr) and overall stability: more epochs, slightly lower LR, added occlp_vr weighting, mild dropout increase.", + "training_hyperparameters": { + "num_epochs": 120, + "batch_size": 128, + "learning_rate": 8e-5, + "optimizer_type": "adamw", + "weight_decay": 0.01, + "use_scheduler": true, + "scheduler_type": "cosine", + "scheduler_step_size": 10, + "scheduler_gamma": 0.1, + "xsmrpool_loss_weight": 10, + "litter_c_loss_weight": 1, + "litter_n_loss_weight": 1, + "litter_p_loss_weight": 1, + "scalar_loss_weight": 1, + "vector_loss_weight": 1, + "matrix_loss_weight": 1 + }, + "reproducibility_config": { + "random_seed": 42, + "strict_determinism": false, + "train_split": 0.8, + "normalization": "individual", + "dropout_p": 0.12 + }, + "data_filtering_config": { + "tropical_only": false, + "tropical_lat_range": null, + "longitudes_to_drop": [0, 358.75], + "natveg_only": true, + "natveg_filter_before_split": false + }, + "variable_weights": { + "pft1d_weights": { + "cpool": 6, + "npool": 8, + "ppool": 8, + "totvegc": 6, + "leafc": 6, + "leafn": 6, + "leafp": 6, + "leafc_storage": 5, + "leafn_storage": 5, + "leafp_storage": 5, + "deadstemc": 6, + "deadstemn": 6, + "deadstemp": 5, + "frootc": 4, + "frootn": 4, + "frootp": 4, + "frootc_storage": 3, + "frootn_storage": 3, + "frootp_storage": 3, + "tlai": 4 + }, + "soil2d_weights": { + "primp_vr": 11, + "occlp_vr": 9, + "litr2c_vr": 8, + "litr2p_vr": 8, + "litr2n_vr": 8, + "soil1c_vr": 7, + "soil1n_vr": 7, + "soil1p_vr": 7, + "litr3n_vr": 5, + "labilep_vr": 5, + "secondp_vr": 2.5, + "solutinorganicp_vr": 5, + "litr3p_vr": 5 + }, + "scalar_weights": { + "GPP": 1.5, + "NPP": 1.5, + "AR": 1.2, + "HR": 1.2 + } + }, + "tail_aware_config": { + "loss": "log1p_huber", + "epsilon": 1e-8, + "huber_delta": 1 + }, + "tail_aware_weights": { + "cpool": 6, + "deadstemc": 6, + "deadcrootc": 5, + "livestemc": 5, + "livecrootc": 5, + "npool": 8, + "ppool": 8, + "totvegc": 6, + "leafc": 6, + "leafn": 6, + "leafp": 6, + "leafc_storage": 5, + "leafn_storage": 5, + "leafp_storage": 5, + "deadstemp": 5, + "deadstemn": 6, + "frootc": 4, + "frootn": 4, + "frootp": 4, + "labilep": 5, + "labilep_vr": 5, + "primp_vr": 11, + "occlp_vr": 9, + "litr2c_vr": 8, + "litr2p_vr": 8, + "litr2n_vr": 8, + "litr3c_vr": 5, + "litr3p_vr": 5, + "litr3n_vr": 5, + "soil1c_vr": 7, + "soil1n_vr": 7, + "soil1p_vr": 7 + }, + "pft_mask_config": { + "mask_absent_pfts": true, + "pft_presence_threshold": 0 + }, + "pft_zero_sparsity_config": { + "weight": 1, + "threshold": 1e-8 + }, + "pft_zero_sparsity_weights": { + "cpool": 2, + "deadstemc": 1, + "deadcrootc": 1, + "livestemc": 1, + "livecrootc": 1 + }, + "pft1d_activation_overrides": { + "cpool": "abs", + "deadstemc": "abs", + "deadcrootc": "abs", + "livestemc": "abs", + "livecrootc": "abs" + } +} diff --git a/config/training_config_phase2_tropical_soilp_only.json b/config/training_config_phase2_tropical_soilp_only.json new file mode 100644 index 0000000..f0e0a21 --- /dev/null +++ b/config/training_config_phase2_tropical_soilp_only.json @@ -0,0 +1,165 @@ +{ + "_description": "Phase 2 tropical-only model: Train on tropical natveg grid cells with EXTREME focus on soil P variables. Use with Phase 1 global restart: Phase 1 produces global restart; Phase 2 predictions replace variables in tropical cells only (via ai_predictions_to_restart --tropical-lat-range). Train: python train_cnp_model.py --training-config-json config/training_config_phase2_tropical_soilp_only.json --variable-list CNP_IO_updated9_dev_dw.txt --use-tva4km --tropical-only --tropical-lat-range -30,30 --output-dir-suffix phase2_tropical_soilp_only", + "training_hyperparameters": { + "num_epochs": 150, + "batch_size": 128, + "learning_rate": 5e-5, + "optimizer_type": "adamw", + "weight_decay": 0.01, + "use_scheduler": true, + "scheduler_type": "cosine", + "scheduler_step_size": 10, + "scheduler_gamma": 0.1, + "xsmrpool_loss_weight": 10, + "litter_c_loss_weight": 1, + "litter_n_loss_weight": 1, + "litter_p_loss_weight": 4, + "scalar_loss_weight": 0.06, + "vector_loss_weight": 0.4, + "matrix_loss_weight": 2.5 + }, + "reproducibility_config": { + "random_seed": 42, + "strict_determinism": false, + "train_split": 0.8, + "normalization": "individual", + "dropout_p": 0.12 + }, + "data_filtering_config": { + "tropical_only": true, + "tropical_lat_range": [-30.0, 30.0], + "longitudes_to_drop": [0, 358.75], + "natveg_only": true, + "natveg_filter_before_split": false + }, + "variable_weights": { + "pft1d_weights": { + "cpool": 3, + "npool": 4, + "ppool": 6, + "totvegc": 3, + "leafc": 3, + "leafn": 3, + "leafp": 4, + "leafc_storage": 2, + "leafn_storage": 2, + "leafp_storage": 3, + "deadstemc": 3, + "deadstemn": 3, + "deadstemp": 4, + "frootc": 2, + "frootn": 2, + "frootp": 3, + "frootc_storage": 1.5, + "frootn_storage": 1.5, + "frootp_storage": 2, + "tlai": 2 + }, + "soil2d_weights": { + "cwdc_vr": 1, + "cwdn_vr": 1, + "cwdp_vr": 2, + "primp_vr": 26, + "occlp_vr": 45, + "labilep_vr": 18, + "secondp_vr": 14, + "solutionp_vr": 18, + "litr2c_vr": 4, + "litr2p_vr": 16, + "litr2n_vr": 4, + "litr3c_vr": 3, + "litr3n_vr": 3, + "litr3p_vr": 12, + "soil1c_vr": 4, + "soil1n_vr": 4, + "soil1p_vr": 14, + "soil2c_vr": 2, + "soil2n_vr": 2, + "soil2p_vr": 10, + "soil3c_vr": 2, + "soil3n_vr": 2, + "soil3p_vr": 8, + "soil4c_vr": 1.5, + "soil4n_vr": 1.5, + "soil4p_vr": 6 + }, + "scalar_weights": { + "GPP": 0.08, + "NPP": 0.08, + "AR": 0.08, + "HR": 0.08 + } + }, + "tail_aware_config": { + "loss": "log1p_huber", + "epsilon": 1e-8, + "huber_delta": 0.7 + }, + "tail_aware_weights": { + "cpool": 3, + "deadstemc": 3, + "deadcrootc": 3, + "livestemc": 3, + "livecrootc": 3, + "npool": 4, + "ppool": 6, + "totvegc": 3, + "leafc": 3, + "leafn": 3, + "leafp": 4, + "leafc_storage": 2, + "leafn_storage": 2, + "leafp_storage": 3, + "deadstemp": 4, + "deadstemn": 3, + "frootc": 2, + "frootn": 2, + "frootp": 3, + "labilep": 10, + "labilep_vr": 18, + "primp_vr": 26, + "occlp_vr": 45, + "secondp_vr": 14, + "solutionp_vr": 18, + "litr2c_vr": 4, + "litr2p_vr": 16, + "litr2n_vr": 4, + "litr3c_vr": 3, + "litr3p_vr": 12, + "litr3n_vr": 3, + "soil1c_vr": 4, + "soil1n_vr": 4, + "soil1p_vr": 14, + "soil2c_vr": 2, + "soil2n_vr": 2, + "soil2p_vr": 10, + "soil3c_vr": 2, + "soil3n_vr": 2, + "soil3p_vr": 8, + "soil4c_vr": 1.5, + "soil4n_vr": 1.5, + "soil4p_vr": 6 + }, + "pft_mask_config": { + "mask_absent_pfts": true, + "pft_presence_threshold": 0 + }, + "pft_zero_sparsity_config": { + "weight": 1, + "threshold": 1e-8 + }, + "pft_zero_sparsity_weights": { + "cpool": 2, + "deadstemc": 1, + "deadcrootc": 1, + "livestemc": 1, + "livecrootc": 1 + }, + "pft1d_activation_overrides": { + "cpool": "abs", + "deadstemc": "abs", + "deadcrootc": "abs", + "livestemc": "abs", + "livecrootc": "abs" + } +} diff --git a/config/training_config_two_region_five_p.json b/config/training_config_two_region_five_p.json new file mode 100644 index 0000000..7cd4ab6 --- /dev/null +++ b/config/training_config_two_region_five_p.json @@ -0,0 +1,169 @@ +{ + "_description": "Two-region (Amazon + Africa) training with 5 P–focused weights. Train only on cells in region_boxes; loss dominated by the five P variables. Train: python train_cnp_model.py --training-config-json config/training_config_two_region_five_p.json --variable-list CNP_IO_updated9_dev_dw.txt --use-tva4km --output-dir-suffix two_region_five_p", + "training_hyperparameters": { + "num_epochs": 150, + "batch_size": 128, + "learning_rate": 5e-5, + "optimizer_type": "adamw", + "weight_decay": 0.01, + "use_scheduler": true, + "scheduler_type": "cosine", + "scheduler_step_size": 10, + "scheduler_gamma": 0.1, + "xsmrpool_loss_weight": 10, + "litter_c_loss_weight": 1, + "litter_n_loss_weight": 1, + "litter_p_loss_weight": 4, + "scalar_loss_weight": 0.06, + "vector_loss_weight": 0.4, + "matrix_loss_weight": 2.5 + }, + "reproducibility_config": { + "random_seed": 42, + "strict_determinism": false, + "train_split": 0.8, + "normalization": "individual", + "dropout_p": 0.12 + }, + "data_filtering_config": { + "tropical_only": false, + "tropical_lat_range": [-30.0, 30.0], + "longitudes_to_drop": [0, 358.75], + "natveg_only": true, + "natveg_filter_before_split": false, + "region_boxes": [ + [-30, 10, 270, 330], + [-15, 15, 0, 30] + ] + }, + "variable_weights": { + "pft1d_weights": { + "cpool": 3, + "npool": 4, + "ppool": 6, + "totvegc": 3, + "leafc": 3, + "leafn": 3, + "leafp": 4, + "leafc_storage": 2, + "leafn_storage": 2, + "leafp_storage": 3, + "deadstemc": 3, + "deadstemn": 3, + "deadstemp": 4, + "frootc": 2, + "frootn": 2, + "frootp": 3, + "frootc_storage": 1.5, + "frootn_storage": 1.5, + "frootp_storage": 2, + "tlai": 2 + }, + "soil2d_weights": { + "cwdc_vr": 0.1, + "cwdn_vr": 0.1, + "cwdp_vr": 0.1, + "primp_vr": 35, + "occlp_vr": 50, + "labilep_vr": 40, + "secondp_vr": 40, + "solutionp_vr": 40, + "litr2c_vr": 0.1, + "litr2p_vr": 0.1, + "litr2n_vr": 0.1, + "litr3c_vr": 0.1, + "litr3n_vr": 0.1, + "litr3p_vr": 0.1, + "soil1c_vr": 0.1, + "soil1n_vr": 0.1, + "soil1p_vr": 0.1, + "soil2c_vr": 0.1, + "soil2n_vr": 0.1, + "soil2p_vr": 0.1, + "soil3c_vr": 0.1, + "soil3n_vr": 0.1, + "soil3p_vr": 0.1, + "soil4c_vr": 0.1, + "soil4n_vr": 0.1, + "soil4p_vr": 0.1 + }, + "scalar_weights": { + "GPP": 0.08, + "NPP": 0.08, + "AR": 0.08, + "HR": 0.08 + } + }, + "tail_aware_config": { + "loss": "log1p_huber", + "epsilon": 1e-8, + "huber_delta": 0.7 + }, + "tail_aware_weights": { + "cpool": 3, + "deadstemc": 3, + "deadcrootc": 3, + "livestemc": 3, + "livecrootc": 3, + "npool": 4, + "ppool": 6, + "totvegc": 3, + "leafc": 3, + "leafn": 3, + "leafp": 4, + "leafc_storage": 2, + "leafn_storage": 2, + "leafp_storage": 3, + "deadstemp": 4, + "deadstemn": 3, + "frootc": 2, + "frootn": 2, + "frootp": 3, + "labilep": 10, + "labilep_vr": 40, + "primp_vr": 35, + "occlp_vr": 50, + "secondp_vr": 40, + "solutionp_vr": 40, + "litr2c_vr": 0.1, + "litr2p_vr": 0.1, + "litr2n_vr": 0.1, + "litr3c_vr": 0.1, + "litr3p_vr": 0.1, + "litr3n_vr": 0.1, + "soil1c_vr": 0.1, + "soil1n_vr": 0.1, + "soil1p_vr": 0.1, + "soil2c_vr": 0.1, + "soil2n_vr": 0.1, + "soil2p_vr": 0.1, + "soil3c_vr": 0.1, + "soil3n_vr": 0.1, + "soil3p_vr": 0.1, + "soil4c_vr": 0.1, + "soil4n_vr": 0.1, + "soil4p_vr": 0.1 + }, + "pft_mask_config": { + "mask_absent_pfts": true, + "pft_presence_threshold": 0 + }, + "pft_zero_sparsity_config": { + "weight": 1, + "threshold": 1e-8 + }, + "pft_zero_sparsity_weights": { + "cpool": 2, + "deadstemc": 1, + "deadcrootc": 1, + "livestemc": 1, + "livecrootc": 1 + }, + "pft1d_activation_overrides": { + "cpool": "abs", + "deadstemc": "abs", + "deadcrootc": "abs", + "livestemc": "abs", + "livecrootc": "abs" + } +} diff --git a/data/data_loader_individual.py b/data/data_loader_individual.py index f6a07b3..0b0302f 100644 --- a/data/data_loader_individual.py +++ b/data/data_loader_individual.py @@ -113,6 +113,17 @@ def _resolve_lat_column(self) -> Optional[str]: return col return None + def _resolve_lon_column(self) -> Optional[str]: + """Resolve longitude column name from config or common patterns.""" + for col in getattr(self.data_config, 'static_columns', []) or []: + if 'lon' in str(col).lower(): + if col in self.df.columns: + return col + for col in ['lon', 'longitude', 'LON', 'Longitude', 'LONGITUDE']: + if col in self.df.columns: + return col + return None + def load_data(self) -> pd.DataFrame: """Load data from configured paths and patterns.""" df_list = [] @@ -218,6 +229,32 @@ def preprocess_data(self): else: logger.warning("'Longitude' column not found in dataset. Cannot apply longitude filtering.") + # Optional region-box filtering: keep only (lat, lon) inside any box. Boxes = (lat_min, lat_max, lon_min, lon_max), lon 0-360. + region_boxes = getattr(self.data_config, 'region_boxes', None) + if region_boxes and len(region_boxes) > 0: + lat_col = self._resolve_lat_column() + lon_col = self._resolve_lon_column() + if lat_col is not None and lon_col is not None: + original_size = len(self.df) + lat_vals = pd.to_numeric(self.df[lat_col], errors='coerce') + lon_vals = pd.to_numeric(self.df[lon_col], errors='coerce') + # Normalize longitude to 0-360 for comparison (if data use -180..180, add 360 when < 0) + lon_360 = lon_vals.where(lon_vals >= 0, lon_vals + 360.0) + mask = pd.Series(False, index=self.df.index) + for (lat_min, lat_max, lon_min, lon_max) in region_boxes: + in_lat = (lat_vals >= float(lat_min)) & (lat_vals <= float(lat_max)) + in_lon = (lon_360 >= float(lon_min)) & (lon_360 <= float(lon_max)) + mask = mask | (in_lat & in_lon) + self.df = self.df[mask].reset_index(drop=True) + filtered_size = len(self.df) + logger.info( + f"Region-box filtering: {original_size} -> {filtered_size} (boxes: {region_boxes})" + ) + if filtered_size == 0: + logger.warning("Region-box filter removed all samples. Check lat/lon columns and box definitions.") + else: + logger.warning("Region-box filter enabled but lat/lon columns not found. Skipping.") + # Optional tropical-only filtering by latitude if getattr(self.data_config, 'tropical_only', False): lat_col = self._resolve_lat_column() @@ -241,7 +278,52 @@ def preprocess_data(self): ) if filtered_size == 0: logger.warning("Tropical filter removed all samples. Check latitude column and range.") - + + # Optional natveg-only filtering: keep only PCT_NATVEG > 0 and PCT_NAT_PFT_0 < 100 + # When natveg_filter_before_split is False: add _natveg_include column but do not drop rows, + # so split is on full data and test set matches no-filter run; train is filtered in split_data(). + if getattr(self.data_config, 'natveg_only', False): + pct_natveg_col = None + for c in ['PCT_NATVEG', 'pct_natveg']: + if c in self.df.columns: + pct_natveg_col = c + break + pct_pft0_col = None + for c in ['PCT_NAT_PFT_0', 'pct_nat_pft_0']: + if c in self.df.columns: + pct_pft0_col = c + break + if pct_natveg_col is None or pct_pft0_col is None: + logger.warning( + "Natveg filter enabled but PCT_NATVEG or PCT_NAT_PFT_0 not found in dataset. " + "Skipping natveg filtering." + ) + else: + pct_natveg = pd.to_numeric(self.df[pct_natveg_col], errors='coerce').fillna(0) + pct_pft0 = pd.to_numeric(self.df[pct_pft0_col], errors='coerce').fillna(100) + include = (pct_natveg > 0) & (pct_pft0 < 100) + filter_before_split = getattr(self.data_config, 'natveg_filter_before_split', True) + if filter_before_split: + # Legacy: filter before shuffle/split (test set = 20% of natveg-only data) + original_size = len(self.df) + self.df = self.df[include].reset_index(drop=True) + filtered_size = len(self.df) + dropped = original_size - filtered_size + logger.info( + f"Natveg filtering (PCT_NATVEG>0 and PCT_NAT_PFT_0<100) before split: " + f"{original_size} -> {filtered_size} samples (dropped {dropped})" + ) + if filtered_size == 0: + logger.warning("Natveg filter removed all samples.") + else: + # Split-then-filter: keep full df, mark rows so train can be filtered in split_data() + self.df['_natveg_include'] = include.values + n_natveg = int(include.sum()) + logger.info( + f"Natveg marking (PCT_NATVEG>0 and PCT_NAT_PFT_0<100): {n_natveg} of {len(self.df)} " + f"samples will be used for training; test set will match no-filter run (full split)." + ) + # Drop specified columns if hasattr(self.data_config, 'filter_columns') and self.data_config.filter_columns: for col in self.data_config.filter_columns: @@ -1635,15 +1717,35 @@ def split_data(self, normalized_data: Dict[str, Any]) -> Dict[str, Any]: logger.info(f" - Train split ratio: {self.data_config.train_split} ({train_size} samples)") logger.info(f" - Test size: {test_size} samples (剩余部分)") + # When natveg_only and split-then-filter: use only natveg rows from first train_size for training; + # test set stays as last 20% of full data (same as no-filter run). + train_mask = None + if ( + '_natveg_include' in self.df.columns + and getattr(self.data_config, 'natveg_only', False) + and not getattr(self.data_config, 'natveg_filter_before_split', True) + ): + train_mask = np.asarray(self.df['_natveg_include'].values[:train_size], dtype=bool) + n_train_natveg = int(np.sum(train_mask)) + logger.info( + f"Natveg train-only filter: using {n_train_natveg} of {train_size} train rows (test set unchanged, {test_size} rows)." + ) + # Expose split indices for downstream use (e.g., location validation) - # Matches the contiguous slicing used below try: - self.train_indices = np.arange(0, train_size, dtype=int) - self.test_indices = np.arange(train_size, total_samples, dtype=int) + if train_mask is not None: + self.train_indices = np.where(train_mask)[0].astype(int) # indices in 0..train_size-1 that are natveg + self.test_indices = np.arange(train_size, total_samples, dtype=int) + else: + self.train_indices = np.arange(0, train_size, dtype=int) + self.test_indices = np.arange(train_size, total_samples, dtype=int) except Exception: - # Fallback without crashing if numpy not available for some reason - self.train_indices = list(range(0, train_size)) - self.test_indices = list(range(train_size, total_samples)) + if train_mask is not None: + self.train_indices = np.where(train_mask)[0].tolist() + self.test_indices = list(range(train_size, total_samples)) + else: + self.train_indices = list(range(0, train_size)) + self.test_indices = list(range(train_size, total_samples)) if test_size == 0: logger.error("Test size is 0! This will cause evaluation issues.") @@ -1651,24 +1753,32 @@ def split_data(self, normalized_data: Dict[str, Any]) -> Dict[str, Any]: # Split time series data train_time_series = normalized_data['time_series_data'][:train_size, :, :] + if train_mask is not None: + train_time_series = train_time_series[train_mask] test_time_series = normalized_data['time_series_data'][train_size:, :, :] train_data['time_series'] = train_time_series test_data['time_series'] = test_time_series # Split static data train_static = normalized_data['static_data'][:train_size] + if train_mask is not None: + train_static = train_static[train_mask] test_static = normalized_data['static_data'][train_size:] train_data['static'] = train_static test_data['static'] = test_static # Split pft_param data train_pft_param = normalized_data['pft_param_data'][:train_size] + if train_mask is not None: + train_pft_param = train_pft_param[train_mask] test_pft_param = normalized_data['pft_param_data'][train_size:] train_data['pft_param'] = train_pft_param test_data['pft_param'] = test_pft_param # Split scalar data (input) train_list_scalar = normalized_data['scalar_data'][:train_size] + if train_mask is not None: + train_list_scalar = train_list_scalar[train_mask] test_list_scalar = normalized_data['scalar_data'][train_size:] train_data['scalar'] = train_list_scalar test_data['scalar'] = test_list_scalar @@ -1676,51 +1786,53 @@ def split_data(self, normalized_data: Dict[str, Any]) -> Dict[str, Any]: # Split y_scalar (target) - skip if not present (inference mode) if 'y_scalar' in normalized_data and normalized_data['y_scalar'] is not None: y_scalar = normalized_data['y_scalar'] - train_data['y_scalar'] = y_scalar[:train_size] + train_data['y_scalar'] = y_scalar[:train_size][train_mask] if train_mask is not None else y_scalar[:train_size] test_data['y_scalar'] = y_scalar[train_size:] # Split variables_1d_pft (input) variables_1d_pft = normalized_data['variables_1d_pft'] - train_data['variables_1d_pft'] = variables_1d_pft[:train_size] + train_data['variables_1d_pft'] = variables_1d_pft[:train_size][train_mask] if train_mask is not None else variables_1d_pft[:train_size] test_data['variables_1d_pft'] = variables_1d_pft[train_size:] # Split y_pft_1d (target) - skip if not present (inference mode) if 'y_pft_1d' in normalized_data and normalized_data['y_pft_1d'] is not None: y_pft_1d = normalized_data['y_pft_1d'] - train_data['y_pft_1d'] = y_pft_1d[:train_size] + train_data['y_pft_1d'] = y_pft_1d[:train_size][train_mask] if train_mask is not None else y_pft_1d[:train_size] test_data['y_pft_1d'] = y_pft_1d[train_size:] # Split y_soil_2d (target) - skip if not present (inference mode) if 'y_soil_2d' in normalized_data and normalized_data['y_soil_2d'] is not None: y_soil_2d = normalized_data['y_soil_2d'] - train_data['y_soil_2d'] = y_soil_2d[:train_size] + train_data['y_soil_2d'] = y_soil_2d[:train_size][train_mask] if train_mask is not None else y_soil_2d[:train_size] test_data['y_soil_2d'] = y_soil_2d[train_size:] # Split variables_2d_soil (input) variables_2d_soil = normalized_data['variables_2d_soil'] - train_data['variables_2d_soil'] = variables_2d_soil[:train_size] + train_data['variables_2d_soil'] = variables_2d_soil[:train_size][train_mask] if train_mask is not None else variables_2d_soil[:train_size] test_data['variables_2d_soil'] = variables_2d_soil[train_size:] # Split water data if present if 'water' in normalized_data and normalized_data['water'] is not None: - train_data['water'] = normalized_data['water'][:train_size] + w = normalized_data['water'] + train_data['water'] = w[:train_size][train_mask] if train_mask is not None else w[:train_size] test_data['water'] = normalized_data['water'][train_size:] if 'y_water' in normalized_data and normalized_data['y_water'] is not None: - train_data['y_water'] = normalized_data['y_water'][:train_size] + yw = normalized_data['y_water'] + train_data['y_water'] = yw[:train_size][train_mask] if train_mask is not None else yw[:train_size] test_data['y_water'] = normalized_data['y_water'][train_size:] # Split PFT presence mask(s). Inference mask (pct>0) goes to both; training mask only to train. if 'pft_presence_mask' in normalized_data: ppm = normalized_data['pft_presence_mask'] try: - train_data['pft_presence_mask'] = ppm[:train_size] + train_data['pft_presence_mask'] = ppm[:train_size][train_mask] if train_mask is not None else ppm[:train_size] test_data['pft_presence_mask'] = ppm[train_size:] except Exception: logger.warning("pft_presence_mask present but could not be split; skipping") if 'pft_presence_mask_training' in normalized_data: ppm_tr = normalized_data['pft_presence_mask_training'] try: - train_data['pft_presence_mask_training'] = ppm_tr[:train_size] + train_data['pft_presence_mask_training'] = ppm_tr[:train_size][train_mask] if train_mask is not None else ppm_tr[:train_size] except Exception: logger.warning("pft_presence_mask_training present but could not be split; skipping") diff --git a/docs/AMAZON_SITE_SOILP_COMPARISON_REPORT.md b/docs/AMAZON_SITE_SOILP_COMPARISON_REPORT.md new file mode 100644 index 0000000..9e69659 --- /dev/null +++ b/docs/AMAZON_SITE_SOILP_COMPARISON_REPORT.md @@ -0,0 +1,136 @@ +# Amazon Site Soil-P Variable Comparison: Extreme Soil-P vs Natveg Improved + +**Date:** 2026-03-04 +**Purpose:** Compare predictions of key soil phosphorus variables (occlp_vr, solutionp_vr, labilep_vr) at the Amazon reference site between the **extreme soil-P** training run and the earlier **natveg improved** (Phase 1) run. + +--- + +## 1. Site and Runs + +### Reference site (Amazon) + +- **Longitude:** 303.75°E +- **Latitude:** −17.43° (≈ −17.434553°) +- **Context:** Tropical site where Phase 1 was known to underpredict occlp_vr and contribute to NEE bias in restart simulations (see `docs/TROPICAL_OCCLP_TWO_OPTIONS_REPORT.md`). + +### Runs compared + +| Run | Config / description | Results path | +|-----|----------------------|-------------| +| **natveg improved (Phase 1)** | Global natveg, occlp-focused weights | `cnp_results/run_20260228_214757_natveg_improved/cnp_inference_entire_dataset/cnp_predictions/` | +| **extreme soil-P** | Experiment 3 extreme soil-P: very high matrix loss and P-specific variable weights (occlp_vr 40, primp_vr, labilep_vr, secondp_vr, solutionp_vr, soil*p_vr, litr*p_vr) | `cnp_results/run_20260304_133300_natveg_occlp_extreme_soilp/cnp_predictions/` | + +Ground truth is the same in both runs (same test set / inference data). Values are per-layer (10 soil layers), units as in the CNP restart (e.g. g/m² for occlp_vr/labilep_vr, solution P in same mass units). + +--- + +## 2. Methodology + +- The Amazon grid cell was identified by matching coordinates (lon 303.75, lat −17.43) in the prediction/ground-truth CSVs. +- For each variable, layer-by-layer values were read from: + - `ground_truth_Y_.csv` + - `predictions_Y_.csv` +- Comparison is **at-site only** (single grid cell), not global metrics. + +--- + +## 3. Variable 1: occlp_vr (occluded soil P) + +**Ground truth (layers 1–10):** +147.23, 144.12, 140.56, 138.03, 136.41, 135.51, 135.12, 135.01, 135.00, 135.00 g/m² + +| Layer | GT | natveg_improved (pred) | extreme_soilp (pred) | +|-------|--------|------------------------|----------------------| +| 1 | 147.23 | 121.99 | 140.27 | +| 2 | 144.12 | 118.48 | 137.90 | +| 3 | 140.56 | 117.05 | 134.98 | +| 4 | 138.03 | 112.22 | 133.14 | +| 5 | 136.41 | 108.15 | 132.21 | +| 6 | 135.51 | 112.34 | 131.26 | +| 7 | 135.12 | 112.50 | 131.10 | +| 8 | 135.01 | 107.18 | 130.58 | +| 9 | 135.00 | 109.06 | 130.50 | +| 10 | 135.00 | 112.40 | 130.62 | + +**Summary:** + +- **natveg_improved:** Large underprediction (~25–30 g/m² in top layers, ~23–28 g/m² in deep layers). +- **extreme_soilp:** Much closer to GT: underprediction reduced to ~5–7 g/m² in top layers and ~4.5 g/m² in deep layers. + +**Conclusion:** occlp_vr at the Amazon site is clearly better with the extreme soil-P config. + +--- + +## 4. Variable 2: solutionp_vr (solution P) + +**Ground truth (layers 1–10):** +0.00127, 0.000931, 0.000556, 0.000300, 0.000138, 5.02e-05, 1.32e-05, 3.29e-06, 1.95e-06, 1.85e-06 + +| Layer | GT | natveg_improved (pred) | extreme_soilp (pred) | +|-------|----------|------------------------|----------------------| +| 1 | 0.00127 | 0.00592 | 0.000203 | +| 2 | 0.000931 | 0.00543 | 0.000321 | +| 3 | 0.000556 | 0.00313 | 0.000253 | +| 4 | 0.000300 | 0.00186 | 0.000173 | +| 5 | 0.000138 | 0.00122 | 9.77e-05 | +| 6 | 5.02e-05 | 0.000391 | 4.37e-05 | +| 7 | 1.32e-05 | 0.000238 | 1.99e-05 | +| 8 | 3.29e-06 | 0.000235 | 1.08e-05 | +| 9 | 1.95e-06 | 0.000309 | 9.94e-06 | +| 10 | 1.85e-06 | 0.000247 | 1.09e-05 | + +**Summary:** + +- **natveg_improved:** Strong overprediction at Amazon (e.g. top layer ~4.7×; deep layers orders of magnitude too high). +- **extreme_soilp:** Much closer: slight underprediction in top layers; deeper layers closer to GT with some small overprediction in the deepest layers. + +**Conclusion:** solutionp_vr at the Amazon site is much better with the extreme soil-P config. + +--- + +## 5. Variable 3: labilep_vr (labile P) + +**Ground truth (layers 1–10):** +40.72, 30.09, 18.21, 9.90, 4.59, 1.67, 0.44, 0.11, 0.065, 0.062 g/m² + +| Layer | GT | natveg_improved (pred) | extreme_soilp (pred) | +|-------|-------|------------------------|----------------------| +| 1 | 40.72 | 47.46 | 38.26 | +| 2 | 30.09 | 35.12 | 30.76 | +| 3 | 18.21 | 21.91 | 18.05 | +| 4 | 9.90 | 11.49 | 9.37 | +| 5 | 4.59 | 5.94 | 4.15 | +| 6 | 1.67 | 2.85 | 1.40 | +| 7 | 0.44 | 1.43 | 0.39 | +| 8 | 0.11 | 1.17 | 0.18 | +| 9 | 0.065 | 0.98 | 0.15 | +| 10 | 0.062 | 1.24 | 0.14 | + +**Summary:** + +- **natveg_improved:** Overprediction in all layers (e.g. +7 in layer 1, +5 in layer 2); deep layers (8–10) are ~10–20× too high (e.g. 0.11→1.17, 0.062→1.24). +- **extreme_soilp:** Slight underprediction in top layers (e.g. 40.7→38.3); mid layers very close (e.g. 18.2→18.0, 9.9→9.4); deep layers still overpredicted but far less than Phase 1 (0.11→0.18, 0.062→0.14). + +**Conclusion:** labilep_vr at the Amazon site is clearly better with the extreme soil-P config. + +--- + +## 6. Overall Summary + +| Variable | natveg_improved at Amazon | extreme_soilp at Amazon | +|---------------|-------------------------------|--------------------------------| +| **occlp_vr** | Large underprediction | Much closer to GT | +| **solutionp_vr** | Large overprediction | Much closer to GT | +| **labilep_vr** | Overprediction, huge in deep | Closer overall; smaller deep bias | + +**Conclusion:** For the Amazon reference site (303.75°E, −17.43°), the **extreme soil-P** training run improves all three soil-P variables (occlp_vr, solutionp_vr, labilep_vr) compared to the natveg_improved (Phase 1) run. The extreme soil-P config (`training_config_experiment_3_natveg_occlp_extreme_soilp.json`) uses higher matrix loss and P-specific variable weights (e.g. occlp_vr weight 40) and yields better tropical P predictions at this site. + +--- + +## 7. Reproducibility + +- **Site validation script:** `scripts/cnp_result_validationplot_site.py --lon 303.75 --lat -17.4246 --tolerance 0.1 --stats-only` (run from results directory). +- **Data sources:** + - Extreme soil-P: `cnp_results/run_20260304_133300_natveg_occlp_extreme_soilp/cnp_predictions/soil_2d_*/*.csv` + - Natveg improved: `cnp_results/run_20260228_214757_natveg_improved/cnp_inference_entire_dataset/cnp_predictions/soil_2d_*/*.csv` +- **Config:** `config/training_config_experiment_3_natveg_occlp_extreme_soilp.json` diff --git a/docs/BIAS_SCALE_PARAMS_5P_TWO_REGIONS.md b/docs/BIAS_SCALE_PARAMS_5P_TWO_REGIONS.md new file mode 100644 index 0000000..6f46f1e --- /dev/null +++ b/docs/BIAS_SCALE_PARAMS_5P_TWO_REGIONS.md @@ -0,0 +1,61 @@ +## Bias/scale correction parameters for 5P, two-region focus + +This document records the linear bias/scale corrections learned by `scripts/apply_5p_bias_scale_correction.py` for the `natveg_improved` run, restricted to the **Amazon** and **Africa** region boxes defined in the two-region configuration. + +For each soil-2D P variable, region, and layer, the script fits: + +\[ +\text{GT} \approx a \cdot \text{pred} + b +\] + +and then applies, at inference time: + +\[ +\hat{Y}_\text{corrected} = a_{\text{var, layer, region}} \cdot \hat{Y}_\text{pred} + b_{\text{var, layer, region}} +\] + +The full set of coefficients is saved in: + +- `docs/bias_scale_params_5P_two_regions_natveg_improved.json` + +That JSON is the canonical source for any downstream scripts that need to reuse the same correction outside the original calibration step. + +### High-level ranges + +Across all 5 P variables, 10 layers, and both regions: + +- **Slopes \(a\)**: roughly from **-0.53** to **1.25**, with most values between **0.9** and **1.1**. +- **Offsets \(b\)**: + - Typically in the range **\(-20\) to \(+40\)** for `labilep_vr`, `occlp_vr`, `solutionp_vr`, and mid/deep layers of `secondp_vr` and `primp_vr`. + - Larger positive offsets (up to \(\sim 192\)) appear for top-layer `secondp_vr` in Africa. + - `solutionp_vr` has very small \(b\) (on the order of \(10^{-3}\) after the scaling trick). + +Per-variable patterns (Amazon + Africa, all layers): + +- **`labilep_vr`**: + - \(a \sim 0.93\text{–}1.09\) + - \(b \sim -0.95\text{–}10.0\) + +- **`occlp_vr`**: + - \(a \sim 0.97\text{–}1.00\) + - \(b \sim 14\text{–}34\) + - Interpretation: primarily an additive upward shift in occluded P. + +- **`solutionp_vr`**: + - \(a \sim 1.06\text{–}1.25\) + - \(b \sim -0.014\text{–}0.001\) + - Effectively a multiplicative rescale, with negligible intercept. + +- **`secondp_vr`**: + - \(a \sim 0.92\text{–}1.09\) + - \(b\) from about \(-18\) to \(\sim 192\), with the largest \(b\) in African top layers. + +- **`primp_vr`**: + - Top 1–2 layers: \(a \sim 0.87\text{–}1.11\), \(b \sim 4.7\text{–}41\). + - Deeper layers: some **negative slopes** with small positive offsets, implying strong shape corrections that are then clamped to non-negative in the application script. + +### How to reuse these parameters + +- For any grid cell inside the **Amazon** or **Africa** region boxes, even without local ground truth, apply the corresponding \((a, b)\) from the JSON for that variable, layer, and region. +- For cells outside those boxes, applying these parameters is an extrapolation; it may still be useful in similar tropical conditions but should be treated as an assumption. + diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 60eed7f..7c6f137 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -145,6 +145,22 @@ python ../../scripts/cnp_result_validationplot.py --stats-only && \ python ../../scripts/generate_prediction_quality_report.py ``` +### 4.0.1) Validate predictions in `cnp_inference_entire_dataset` (inference subdirectory) +If inference wrote outputs under a subdirectory (e.g. `cnp_inference_entire_dataset/` with `cnp_predictions/` inside it), run validation by passing that directory. From the **run** directory (e.g. `cnp_results/run_YYYYMMDD_HHMMSS`): + +```bash +# One-time: copy run config into the inference dir so validation can read variable lists (if not already there) +cp cnp_config.json cnp_inference_entire_dataset/ 2>/dev/null || true + +# Stats + quality report for the inference subdirectory +python ../../scripts/cnp_result_validationplot.py cnp_inference_entire_dataset --stats-only && \ +python ../../scripts/generate_prediction_quality_report.py cnp_inference_entire_dataset/validation_stats.csv +``` + +- First command writes `cnp_inference_entire_dataset/validation_stats.csv`. +- Second reads that CSV and writes `cnp_inference_entire_dataset/analysis/` (quality report, plots, `quality_summary_report.txt`, `top_bad_plots/`, etc.). +- Replace `cnp_inference_entire_dataset` with your actual inference subdirectory name if different. + ### 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. @@ -273,6 +289,29 @@ python ../../scripts/ai_restart_comparison.py \ --stats-only \ --variable-list ../../CNP_IO_updated9_dev.txt ``` + +#### 9.2) Custom AI vs model (ground truth) restart paths and plot all PFTs/layers +To compare an AI-augmented restart file with a **model (ground truth) restart file** using your own paths, use `--ai-restart` and `--original-restart`. To generate map plots for **every variable with every soil layer and every PFT**, use `--plot-all` (10 soil layers 0–9 and PFTs 1–16; PFT 0 is skipped). + +From the **project root** (`AI4BGC/`): + +```bash +python scripts/ai_restart_comparison.py \ + --variable-list CNP_IO_updated9_dev.txt \ + --ai-restart cnp_results/run_YYYYMMDD_HHMMSS/CNInit_updated9_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --original-restart ELM_data/20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --plot-all \ + --output-dir cnp_results/run_YYYYMMDD_HHMMSS/ai_restart_comparison_plots +``` + +- **`--variable-list`**: CNP_IO file that defines PFT1D and soil2D variables to compare (required for variable list; use the list that matches your run, e.g. `CNP_IO_updated9_dev.txt`). +- **`--ai-restart`**: Path to the AI-updated restart NetCDF. +- **`--original-restart`**: Path to the model (ground truth) restart NetCDF. +- **`--plot-all`**: Plot all variables for all 10 soil layers (0–9) and all 16 PFTs (1–16). +- **`--output-dir`**: Where to write map PNGs and stats (default: `./ai_restart_comparison_plots`). + +Each plot is a 4-panel map: AI Restart, Original Restart, difference, and percent-difference bins. Files are named like `{variable}_lev{layer}.png` (soil variables) and `{variable}_pft{k}.png` (PFT variables). Statistics are written to `restart_stats.csv` / `restart_stats.txt` and a quality summary to `restart_quality_by_variable.png` in the output directory. + --- ### Old scripts (to be double-checked) diff --git a/docs/EXCLUDED_SAMPLE_ANALYSIS.md b/docs/EXCLUDED_SAMPLE_ANALYSIS.md new file mode 100644 index 0000000..acbccb3 --- /dev/null +++ b/docs/EXCLUDED_SAMPLE_ANALYSIS.md @@ -0,0 +1,233 @@ +# Excluded Sample Analysis (Natveg Filter) + +This document describes how **excluded gridcells** are defined, how they affect validation metrics and plots, and how to analyze or plot variables at those locations only. It also includes a **training-data exclusion report** for the natveg filter. + +--- + +## Report: Training Data Exclusion by Natveg Filter + +**Purpose:** Quantify how many training samples would be excluded if the natveg filter were applied at training time (i.e., training only on gridcells with natural vegetation). + +**Filter rule (same as in validation):** +- **Include:** `(PCT_NATVEG > 0) AND (PCT_NAT_PFT_0 < 100)` +- **Exclude:** `(PCT_NATVEG = 0) OR (PCT_NAT_PFT_0 >= 100)` + *(PCT_NATVEG is a percentage, so “no natural vegetation” is exactly 0; code may use ≤ 0 for robustness.)* + +### Dataset + +| Item | Value | +|------|--------| +| Data path | `/mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_1_data_CNP` | +| File pattern | `training_data_batch_*.pkl` | +| Number of batch files | 21 | +| Total samples loaded | 20,975 | + +### Results + +| Metric | Count | Percentage | +|--------|-------|------------| +| **Total samples** | 20,975 | 100% | +| **Included (kept)** | 14,100 | **67.22%** | +| **Excluded** | **6,875** | **32.78%** | + +### Exclusion breakdown + +| Condition | Count | Note | +|----------|--------|------| +| PCT_NATVEG = 0 | 5,921 | No natural vegetation (0%) | +| PCT_NAT_PFT_0 ≥ 100 | 6,875 | 100% non-vegetated PFT (all excluded rows satisfy this) | +| Both conditions | 5,921 | Overlap: 954 excluded only by PFT_0 ≥ 100 | + +### Interpretation + +- **~33% of training samples** would be dropped by the natveg filter on this dataset. +- This is a **dramatic** reduction: training with the filter would use **14,100** samples instead of **20,975**. +- If the filter is adopted for training, expect: + - Fewer samples per epoch and potentially different convergence. + - A model focused on gridcells with natural vegetation; predictions at excluded gridcells would be extrapolation. + +### How to reproduce / re-run the count + +Use the script from the repo root: + +```bash +# Full dataset (same paths as above) +python scripts/count_natveg_filter_exclusion.py \ + --data-paths /mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_1_data_CNP \ + --file-pattern "training_data_batch_*.pkl" + +# With config/variable list (uses config data paths) +python scripts/count_natveg_filter_exclusion.py --variable-list CNP_IO_updated9_dev_dw.txt + +# Quick check on a subset +python scripts/count_natveg_filter_exclusion.py --data-paths /path/to/data --max-files 5 + +# After applying tropical filter first +python scripts/count_natveg_filter_exclusion.py --data-paths /path/to/data --after-tropical +``` + +Script: `scripts/count_natveg_filter_exclusion.py`. + +--- + +## 1. What Are Excluded Gridcells? + +For PFT (1D) and soil 2D variables, we optionally exclude gridcells that are not meaningful for vegetation-related evaluation: + +- **PCT_NATVEG = 0** — no natural vegetation. +- **PCT_NAT_PFT_0 = 100** — 100% of the gridcell is “PFT 0” (non-vegetated). + +So **excluded** = gridcells with no natural vegetation or 100% non-vegetated PFT. +**Included** = gridcells with `PCT_NATVEG > 0` **and** `PCT_NAT_PFT_0 < 100`. + +These flags come from the **static** inputs (e.g. `test_static_inverse.csv`), which must have columns `PCT_NATVEG` and `PCT_NAT_PFT_0`. + +## 2. Where the Exclusion Is Used + +- **Validation metrics (R², RMSE, MAE)** + When the **natveg filter** is enabled (default; use `--no-natveg-filter` to turn it off), metrics for PFT and 2D variables are computed **only over included gridcells**. Excluded gridcells are not used in the metric. + +- **GT vs Pred scatter plots in `top_bad_plots`** + When the natveg filter is **on**, the script first restricts to included gridcells, then plots and computes R²/RMSE on that subset. So both the **plot** and the **metric** in `validation_stats.csv` are for included-only. When the filter is **off**, all test gridcells are used for the plot and the metric. + +**Fair comparison (no-filter vs natveg run):** For the **natveg** run we should only evaluate on gridcells that pass the natveg filter (the “relevant” gridcells). The natveg model is trained only on those, so we expect it to perform **somewhat better** on that subset. Therefore: +- Compare **both** runs using **validation_stats** (or the quality report) generated with the natveg filter on (default; do not pass `--no-natveg-filter`). Then both metrics are “included gridcells only,” and the natveg run can be fairly expected to be somewhat better on that subset. +- The **inference-time** metrics (`cnp_predictions/test_metrics.csv`, `cnp_metrics.json`) are computed over the **full** test set. For the natveg run that is unfair (it includes gridcells the model was not trained on), so prefer the validation script with the filter for reporting natveg performance. + +## 3. Data Source: Row Alignment + +Excluded vs included is determined from: + +- **File:** `results_dir/cnp_predictions/test_static_inverse.csv` +- **Row order:** Must match the ground-truth and prediction CSVs (same test set, same row index = same gridcell). + +The validation script builds a boolean mask from `PCT_NATVEG` and `PCT_NAT_PFT_0` and applies it by **row index** when computing metrics. It does **not** match by latitude/longitude. + +## 4. Excluded Locations List (Lat/Lon) + +For a given run, the list of excluded gridcells (lat/lon and PCT values) is written when you run the excluded-locations plotting script: + +- **Path:** + `results_dir/analysis/excluded_locations/excluded_locations_lat_lon.csv` +- **Columns:** `Latitude`, `Longitude`, `PCT_NATVEG`, `PCT_NAT_PFT_0` +- **Rows:** One per excluded gridcell (e.g. 36 in a typical run). + +Example (first rows): + +```text +Latitude,Longitude,PCT_NATVEG,PCT_NAT_PFT_0 +17.434555,25.0,99.99999,100.0 +15.549739,15.000001,99.99999,100.0 +22.146597,348.75,99.99999,100.0 +... +``` + +Longitude is in 0–360° (e.g. 348.75° ≈ −11.25°). + +## 5. Script: Plot Variables at Excluded Locations Only + +To visualize **only** the excluded gridcells (GT vs Pred at those locations): + +- **Script:** `scripts/plot_excluded_locations.py` +- **Input:** Results directory (e.g. `cnp_results/run_YYYYMMDD_HHMMSS`). +- **Output directory:** + `results_dir/analysis/excluded_locations/` + (override with `--output-dir`). + +### What it does + +1. Reads `test_static_inverse.csv` and computes the same exclude mask (PCT_NATVEG=0 or PCT_NAT_PFT_0=100). +2. For each of a set of variables (cpool, npool, ppool, litr1c_vr, litr1n_vr, litr1p_vr, soil1c_vr, soil1n_vr, soil1p_vr, primp_vr), loads GT and prediction CSVs. +3. Keeps **only rows** whose index is in the excluded set. +4. Plots **GT vs Pred** for those rows (all PFTs or layers flattened for that variable). +5. Writes: + - One PNG per variable: `excluded_locations__gt_vs_pred.png` + - `excluded_locations_lat_lon.csv` (list of excluded lat/lon and PCT_*). + +### How to run + +```bash +# From repo root; output under results_dir/analysis/excluded_locations/ +python scripts/plot_excluded_locations.py cnp_results/run_20260220_145518 + +# Custom output directory +python scripts/plot_excluded_locations.py cnp_results/run_20260220_145518 -o /path/to/output +``` + +### Interpreting the plots + +- Each point is one (gridcell × PFT) or (gridcell × layer) at **excluded** locations only. +- Points on the 1:1 line = good agreement at those cells; off the line = bias or error at excluded (e.g. 100% PFT0) gridcells. + +## 6. Validation: Enabling the Natveg Filter + +- **Default:** Natveg filter is **on**. PFT/2D metrics use only **included** gridcells. Pass `--no-natveg-filter` to use all gridcells instead. + +```bash +# Validation stats and top_bad plots *with* natveg filter (default; included gridcells only) +python scripts/cnp_result_validationplot.py cnp_results/run_20260220_145518 + +# Validation stats and top_bad plots *without* filter (all gridcells) +python scripts/cnp_result_validationplot.py cnp_results/run_20260220_145518 --no-natveg-filter +``` + +The **scatter plots** in `top_bad_plots/by_pft_layer/` still show all gridcells in both cases; only the numbers in `validation_stats.csv` (and thus in the quality report) change when the filter is used. + +## 7. Typical Counts (Example Run) + +For one run (798 test gridcells): + +| Metric | Value | +|---------------------|---------| +| Total gridcells | 798 | +| Included (kept) | 762 (95.5%) | +| Excluded (filtered) | 36 (4.5%) | + +All 36 exclusions in that run had **PCT_NAT_PFT_0 = 100** (and PCT_NATVEG ≈ 100); none had PCT_NATVEG = 0. + +## 8. Summary + +| Question | Answer | +|----------|--------| +| Are excluded gridcells still in the GT-vs-Pred scatter plots? | **Yes.** `top_bad_plots/by_pft_layer/*.png` always show all gridcells. | +| What changes with the natveg filter (on by default)? | Only the **metrics** (R², RMSE, MAE) for PFT/2D variables; they are computed over **included** gridcells only. | +| Where is the list of excluded locations? | `analysis/excluded_locations/excluded_locations_lat_lon.csv` after running `plot_excluded_locations.py`. | +| How do I plot only excluded gridcells? | Run `scripts/plot_excluded_locations.py `; see `analysis/excluded_locations/excluded_locations__gt_vs_pred.png`. | + +## 9. Test Set Alignment: Natveg Run vs No-Filter Run + +When comparing a **natveg-only** run to a **no-filter** run, the test sets can differ depending on when the filter is applied. + +### Two behaviors + +| Option | When filter is applied | Test set | +|--------|-------------------------|----------| +| **Filter before split** (default / legacy) | Filter to natveg **before** shuffle and 80/20 split. | Test = 20% of **natveg-only** data (e.g. ~2802 rows). **Not** a subset of the no-filter test set. | +| **Split then filter** (`--no-natveg-filter-before-split`) | Shuffle and split on **full** data (same as no-filter); then filter **only the training set** to natveg. | Test = same 20% as no-filter run (e.g. ~4166 rows). Natveg validation plots show only the **included** subset of that test set. | + +So: **for the natveg run’s validation data to be a subset of the no-filter run’s validation data**, use **split then filter** (do **not** filter before split). + +### How to get aligned test sets + +- **Training:** use `natveg_only: true` and `natveg_filter_before_split: false` in config, or `--natveg-only --no-natveg-filter-before-split` on the CLI. +- **Inference:** uses the same data config from `cnp_config.json` (including `natveg_filter_before_split`), so test rows and `test_static_inverse.csv` match the no-filter run. +- **Validation:** with the natveg filter on, plots and metrics use only the **included** gridcells from that same test set, so points in e.g. `2D_soil4c_vr_Layer10_gt_vs_pred.png` are a subset of the no-filter run’s plot. + +### Verifying overlap + +Use the script below to compare two runs’ test gridcells by (Latitude, Longitude): + +```bash +python scripts/verify_natveg_test_subset.py \ + cnp_results/run_NOFILTER/cnp_predictions \ + cnp_results/run_NATVEG/cnp_predictions +``` + +If the natveg run was trained with `natveg_filter_before_split: false`, the natveg test set (rows in `test_static_inverse.csv`) should be **identical** to the no-filter test set; the script reports overlap and subset relationship. + +## 10. Related Files + +- **Validation (filter logic):** `scripts/cnp_result_validationplot.py` — `_load_gridcell_metadata()`, `analyze_1d_new_structure(..., gridcell_metadata)`, `analyze_2d_new_structure(..., gridcell_metadata)`. +- **Quality report:** `scripts/generate_prediction_quality_report.py` — natveg filter on by default when invoking validation for top-bad plots; use `--no-natveg-filter` to turn it off. +- **Excluded-only plots:** `scripts/plot_excluded_locations.py`. +- **Prediction quality overview:** `docs/README_prediction_quality.md`. diff --git a/docs/FIVE_P_TWO_REGION_GROUND_TRUTH_ANALYSIS.md b/docs/FIVE_P_TWO_REGION_GROUND_TRUTH_ANALYSIS.md new file mode 100644 index 0000000..6618dc5 --- /dev/null +++ b/docs/FIVE_P_TWO_REGION_GROUND_TRUTH_ANALYSIS.md @@ -0,0 +1,55 @@ +# 5 P Ground Truth: Two Regions vs Rest of World + +**Goal:** Get the 5 P predictions (labilep_vr, occlp_vr, solutionp_vr, secondp_vr, primp_vr) in the two target regions (Amazon + Central Africa) as close as possible to ground truth. One strategy is to train only on data from these two regions, **if** the 5 P variables there have different distributions than in other regions. + +--- + +## Is it true that 5 P in these two regions are different? + +**Yes.** Ground-truth analysis on the training data (Trendy_1_data_CNP) shows that the 5 P variables in **Amazon + Central Africa** have clearly different (and much higher) distributions than in the **rest of the world**. + +### Results (sum over first 10 soil layers, same units as training) + +| Variable | Two regions (Amazon + Africa) | Other (rest of world) | Ratio (two_region / other) | +|----------------|--------------------------------|------------------------|-----------------------------| +| Y_labilep_vr | mean 636, p50 605 | mean 232, p50 0 | **2.74** | +| Y_occlp_vr | mean 3305, p50 3511 | mean 1351, p50 0 | **2.45** | +| Y_solutionp_vr | mean 1.12, p50 0.02 | mean 0.24, p50 0 | **4.65** | +| Y_secondp_vr | mean 12847, p50 10335 | mean 4849, p50 0 | **2.65** | +| Y_primp_vr | mean 82.7, p50 44.6 | mean 17.1, p50 0 | **4.84** | + +- **Two regions:** 1,758 grid cells (Amazon box: lat −30–10, lon 270–330; Africa box: lat −15–15, lon 0–30). +- **Other:** 9,242 grid cells from the same batches (lat outside the two boxes or lon outside). +- In “other” regions, many cells have **zero or very low** P (p50 = 0 for all five variables). In the two regions, medians are **non-zero and large** (e.g. labilep_vr p50 ≈ 605, occlp_vr p50 ≈ 3511). + +So the 5 P variables in the two target regions are **not** representative of the global distribution: they are systematically higher and have a different shape (fewer zeros, higher medians and means). Training a global model on all data would be dominated by the many low-P cells elsewhere and can underfit or mis-scale the high-P tropical behavior in Amazon and Africa. + +--- + +## Recommendation + +1. **Training only on the two regions is well justified.** The ground truth shows that 5 P in Amazon + Africa are different from the rest of the world (ratios of means ~2.4–4.8). Using only data from these regions lets the model focus on the relevant distribution and should help get 5 P predictions there closer to ground truth. +2. **You are already doing this** with **`config/training_config_two_region_five_p.json`**, which sets `data_filtering_config.region_boxes` to the Amazon and Africa boxes so the dataloader keeps only cells in those two regions. The run **run_20260308_211537_two_region_five_p** used this config and is a large improvement over the non–region-focused run. +3. To **improve further** in the two regions, you can: + - Keep two-region-only data (region_boxes) and optionally **increase 5 P weights** or **epochs**. + - **Finetune** from the current two-region run (or from phase2_pvariable_focus) with the same region_boxes and strong 5 P weights. + - Re-run the ground-truth analysis on **all** training files (or on the exact files used in production) to confirm the same conclusion; the script supports `--max-files` and `--data-dir`. + +--- + +## How to reproduce the analysis + +From the repo root: + +```bash +python scripts/analyze_5p_ground_truth_by_region.py \ + --data-dir /path/to/Trendy_1_data_CNP \ + --max-files 11 \ + --output docs/analysis_5p_by_region.txt \ + --output-csv docs/analysis_5p_by_region.csv +``` + +- Use `--max-files` so that the loaded batches include both two-region and other cells (e.g. batches 7–11 for Trendy_1_data_CNP). For a full run, use a large `--max-files` or load all batches. +- Results are written to `--output` (text report) and optionally `--output-csv` (per-variable stats). + +See **`docs/analysis_5p_by_region.txt`** for the exact numbers and conclusion from the last run. diff --git a/docs/INSTRUCTIONS_TRENDY_1_AI_RESTART_CREATION.md b/docs/INSTRUCTIONS_TRENDY_1_AI_RESTART_CREATION.md new file mode 100644 index 0000000..2ee8b3d --- /dev/null +++ b/docs/INSTRUCTIONS_TRENDY_1_AI_RESTART_CREATION.md @@ -0,0 +1,252 @@ +# Instructions: Repeat workflow and create new restart in branch `trendy_1_ai_restart_creation` + +This document gives step-by-step instructions to repeat the process in **docs/WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md** and produce a new restart file in a dedicated run branch **trendy_1_ai_restart_creation**. + +--- + +## Prerequisites (for new users and existing runs) + +This workflow assumes you have **two trained CNP models**: + +1. A **global natveg_improved-like model** with a global restart file. +2. A **tropical Phase2 P-focused model** (`phase2_pvariable_focus`-like) with a + trained checkpoint and inference outputs. + +If you are a **new user** and do not yet have these models, first follow +Section 0 in `docs/WORKFLOW_PHASE1_PHASE2_PHASE3_RESTARTS.md` to train them +with `train_cnp_model.py`. In short: + +- Global model (natveg_improved-like): + + ```bash + cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC + + python train_cnp_model.py \ + --config config/training_config_experiment_3_global_natveg_improved.json \ + --run-dir cnp_results/run_YYYYMMDD_HHMMSS_natveg_improved_custom \ + --variable-list CNP_IO_updated9_dev_dw.txt + ``` + + Then produce a global restart NetCDF from this run and set: + + ```bash + export NATVEG_RUN_DIR="cnp_results/run_YYYYMMDD_HHMMSS_natveg_improved_custom" + export BASE_RESTART_FILE="$NATVEG_RUN_DIR/updated_restart_...your_file.nc" + ``` + +- Tropical P-focused model (phase2_tropical-like): + + ```bash + cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC + + python train_cnp_model.py \ + --config config/training_config_phase2_tropical_soilp_only.json \ + --run-dir cnp_results/run_YYYYMMDD_HHMMSS_phase2_pvariable_focus_custom \ + --variable-list CNP_IO_updated9_dev_dw.txt + ``` + + Then set: + + ```bash + export PHASE2_RUN_DIR="cnp_results/run_YYYYMMDD_HHMMSS_phase2_pvariable_focus_custom" + ``` + +If you already have trained runs, you can reuse them. The paths used below +assume the **existing** experiments (adjust if your run IDs differ): + +- **Natveg run:** `cnp_results/run_20260228_214757_natveg_improved` +- **Base restart file:** + `cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc` +- **Phase2 run:** `cnp_results/run_20260305_153217_phase2_pvariable_focus` + +--- + +## Step 0: Create git branch and run directory + +From the repo root: + +```bash +git checkout -b trendy_1_ai_restart_creation +``` + +Create a run directory for this workflow’s outputs (restart, logs, copies of params): + +```bash +export RUN_TS=$(date +%Y%m%d_%H%M%S) +export TRENDY_RUN_DIR="cnp_results/run_${RUN_TS}_trendy_1_ai_restart_creation" +mkdir -p "$TRENDY_RUN_DIR" +``` + +Use `TRENDY_RUN_DIR` in the steps below so the new restart and any logs live under this branch’s run. + +--- + +## Path A: Tropical-only inference → 4P or 5P restart (no bias/scale, or with bias/scale) + +Use this if you only have (or only want) **tropical** Phase2 predictions. + +### A1. Phase2 tropical-only inference (if not already done) + +```bash +cd cnp_results/run_20260305_153217_phase2_pvariable_focus +python ../../scripts/run_inference_all.py \ + --model cnp_predictions/model.pth \ + --output-dir cnp_inference_tropical_only +cd ../.. +``` + +### A2. NetCDF from tropical predictions + +```bash +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_inference_tropical_only/cnp_predictions \ + --output cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_tropical_only.nc +``` + +### A3. Update natveg restart (4 P variables, raw Phase2) + +```bash +python scripts/ai_predictions_to_restart.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_tropical_only.nc \ + --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --output "$TRENDY_RUN_DIR/updated_restart_phase2_tropical_4p.nc" \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --variables-to-update occlp_vr,labilep_vr,solutionp_vr,secondp_vr \ + "--tropical-lat-range=-30,30" +``` + +**New restart file:** `$TRENDY_RUN_DIR/updated_restart_phase2_tropical_4p.nc` + +To use **5P with bias/scale** on tropical-only: first run the bias/scale step (Section B2 below) on the **tropical** CSVs (script expects `cnp_inference_entire_dataset` layout; if your tropical run writes elsewhere, either symlink or run bias/scale from a copy that matches that layout), then build a NetCDF from the bias-corrected CSVs and run `ai_predictions_to_restart.py` with `--variables-to-update labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr` and `--output "$TRENDY_RUN_DIR/updated_restart_phase2_5P_bias_corrected_tropical.nc"`. + +--- + +## Path B: Full-grid inference → 5P bias/scale → tropical restart (recommended for 5P bias/scale) + +Use this to get **bias-corrected 5P in Amazon + Africa** and a single global NetCDF. + +### B1. Phase2 full-grid inference + +From repo root. You can set the output directory with `--output-dir` (default is `cnp_inference_entire_dataset` under the current directory). Example: write outputs into the Phase2 run directory, or into your trendy run directory. + +** Output into your trendy run directory** (then bias/scale script must be pointed at this run’s paths or you copy/symlink the expected structure): + +```bash +# TRENDY_RUN_DIR set in Step 0 +python scripts/run_inference_all.py \ + --model cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_predictions/model.pth \ + --output-dir "$TRENDY_RUN_DIR/cnp_inference_entire_dataset" \ + --inference-full-grid +``` + +This creates `.../cnp_predictions/` and `.../soil_2d_ground_truth/` under the given `--output-dir`. The bias/scale script (B2) expects `--run-dir` to be the run that contains `cnp_inference_entire_dataset`; if you use Option 2, set `--run-dir "$TRENDY_RUN_DIR"` in B2 and ensure the run has that subdir. + +### B2. Apply 5P bias/scale correction (Amazon + Africa) + +Use as `--run-dir` the directory that **contains** `cnp_inference_entire_dataset/` (where B1 wrote the CSVs). If you used **Option 1** in B1, that is the Phase2 run dir; if you used **Option 2**, that is `$TRENDY_RUN_DIR`. + +**If you used Option 1 (output under Phase2 run):** + +```bash +python scripts/apply_5p_bias_scale_correction.py \ + --run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --region-config-json config/training_config_two_region_five_p.json \ + --output-subdir soil_2d_predictions_5P_bias_corrected_phase2 +``` + +**If you used Option 2 (output under TRENDY_RUN_DIR):** + +```bash +python scripts/apply_5p_bias_scale_correction.py \ + --run-dir "$TRENDY_RUN_DIR" \ + --region-config-json config/training_config_two_region_five_p.json \ + --output-subdir soil_2d_predictions_5P_bias_corrected_phase2 +``` + +Corrected CSVs end up under `/cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_5P_bias_corrected_phase2/`. Bias/scale params are written to `/analysis/bias_scale_params_5P_two_regions.json`. + +Optional: copy bias/scale params into the trendy run for provenance (if Phase2 run was used for B2, copy from there): + +```bash +cp cnp_results/run_20260305_153217_phase2_pvariable_focus/analysis/bias_scale_params_5P_two_regions.json \ + "$TRENDY_RUN_DIR/bias_scale_params_5P_two_regions_phase2.json" +``` + +### B3. NetCDF from bias-corrected 5P predictions + +Use the same run directory you used in B2 (Phase2 run dir or `$TRENDY_RUN_DIR`). + +```bash +# If you used Option 1 in B1/B2 (Phase2 run dir): +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_inference_entire_dataset/cnp_predictions/ \ + --soil-2d-bias-corrected-subdir soil_2d_predictions_5P_bias_corrected_phase2 \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --output cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc + +# If you used Option 2, write NetCDF into your run dir: +mkdir -p "$TRENDY_RUN_DIR/comparison_results" + +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions "$TRENDY_RUN_DIR/cnp_inference_entire_dataset/cnp_predictions/" \ + --soil-2d-bias-corrected-subdir soil_2d_predictions_5P_bias_corrected_phase2 \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --output "$TRENDY_RUN_DIR/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc" +``` + +### B4. Generate new restart (5P in tropics, from natveg base) + +Use the path to the NetCDF you produced in B3 (Phase2 run dir or `$TRENDY_RUN_DIR/comparison_results/...`). + +```bash +# If Option 1 was used (NetCDF in Phase2 run dir): +python scripts/ai_predictions_to_restart.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc \ + --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --output "$TRENDY_RUN_DIR/updated_restart_phase2_5P_bias_corrected_tropical.nc" \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --variables-to-update labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr \ + "--tropical-lat-range=-30,30" + +# If Option 2 was used (NetCDF in trendy run dir): +python scripts/ai_predictions_to_restart.py \ + --ai-predictions "$TRENDY_RUN_DIR/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc" \ + --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --output "$TRENDY_RUN_DIR/updated_restart_phase2_5P_bias_corrected_tropical.nc" \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --variables-to-update labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr \ + "--tropical-lat-range=-30,30" +``` + +**New restart file:** `$TRENDY_RUN_DIR/updated_restart_phase2_5P_bias_corrected_tropical.nc` + +--- + +## Step 7: Record and validate + +1. **Record** in `$TRENDY_RUN_DIR`: + - Which path (A or B) was used. + - Base restart path and Phase2 run path. + - Command lines for `ai_predictions_to_netcdf.py` and `ai_predictions_to_restart.py`. + + Example README: + + ```bash + echo "Restart created: $(date). Path B (full-grid + 5P bias/scale). Base: natveg_improved. Phase2: run_20260305_153217_phase2_pvariable_focus." > "$TRENDY_RUN_DIR/README_restart.txt" + ``` + +2. **Run the land model** using the new restart and validate (e.g. Amazon/Africa 5P profiles, NPP, fluxes) as in **docs/WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md**. + +--- + +## Summary checklist (Path B – 5P bias/scale) + +- [ ] `git checkout -b trendy_1_ai_restart_creation` +- [ ] Create `TRENDY_RUN_DIR=cnp_results/run__trendy_1_ai_restart_creation` +- [ ] Run Phase2 **full-grid** inference (`--inference-full-grid`) +- [ ] Run **apply_5p_bias_scale_correction.py** (Phase2 run-dir, two-region config) +- [ ] Run **ai_predictions_to_netcdf.py** on bias-corrected 5P CSVs +- [ ] Run **ai_predictions_to_restart.py** (natveg base restart → `$TRENDY_RUN_DIR/updated_restart_phase2_5P_bias_corrected_tropical.nc`, 5P, `--tropical-lat-range=-30,30`) +- [ ] Document and validate + +Reference: **docs/WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md**. diff --git a/docs/NATVEG_VS_NOFILTER_COMPARISON.md b/docs/NATVEG_VS_NOFILTER_COMPARISON.md new file mode 100644 index 0000000..afbd572 --- /dev/null +++ b/docs/NATVEG_VS_NOFILTER_COMPARISON.md @@ -0,0 +1,165 @@ +# Natveg filter vs no-filter: performance comparison + +## Run setup + +| Run | num_samples | natveg_only | natveg_filter_before_split | +|-----|-------------|-------------|----------------------------| +| nofilter | 20826 | false | true | +| natveg | 14006 | true | — | +| natveg_aligned | 20826 | true | false | + +- **nofilter**: full data (train+test); no natveg filter. +- **natveg**: natveg-only filter *before* split → fewer training samples (≈14k vs 20k). +- **natveg_aligned**: natveg-only for *training*; test set same as nofilter (same 20% holdout). + +## Aggregate metrics + +| scalar_rmse | nofilter | **0.021993** | — | natveg | 0.032461 | +47.6% | natveg_aligned | 0.020582 | -6.4% | +| pft_1d_rmse | nofilter | **0.074273** | — | natveg | 0.072640 | -2.2% | natveg_aligned | 0.074440 | +0.2% | +| soil_2d_rmse | nofilter | **0.032720** | — | natveg | 0.040012 | +22.3% | natveg_aligned | 0.034643 | +5.9% | + +(Lower RMSE is better; positive % = degradation with filter.) + +## Scalar fluxes (R²) + +| Variable | nofilter | natveg | Δ (pp) | natveg_aligned | Δ (pp) | +|----------|----------|--------|--------|----------------|--------| +| Y_GPP | 0.9900 | 0.9793 | -1.07 | 0.9919 | +0.19 | +| Y_NPP | 0.9877 | 0.9792 | -0.85 | 0.9884 | +0.07 | +| Y_AR | 0.9874 | 0.9757 | -1.18 | 0.9900 | +0.25 | +| Y_HR | 0.9871 | 0.9819 | -0.52 | 0.9879 | +0.08 | + +(R² in [0,1]; negative Δ = degradation with filter.) + +## Soil 2D variables (mean R² over layers) + +| Variable | nofilter R² | natveg R² | Δ (pp) | natveg_aligned R² | Δ (pp) | +|----------|--------------|-----------|--------|-------------------|--------| +| cwdc_vr | 0.9387 | 0.9390 | +0.03 | 0.9435 | +0.48 | +| cwdn_vr | 0.9375 | 0.9396 | +0.20 | 0.9437 | +0.61 | +| cwdp_vr | 0.9374 | 0.9376 | +0.02 | 0.9420 | +0.46 | +| labilep_vr | 0.9871 | 0.9748 | -1.23 | 0.9840 | -0.32 | +| litr2c_vr | 0.7254 | 0.6183 | -10.71 | 0.7274 | +0.20 | +| litr2n_vr | 0.7170 | 0.6224 | -9.46 | 0.7222 | +0.52 | +| litr2p_vr | 0.7174 | 0.6141 | -10.33 | 0.7221 | +0.46 | +| litr3c_vr | 0.8760 | 0.8551 | -2.09 | 0.8746 | -0.14 | +| litr3n_vr | 0.8684 | 0.8540 | -1.44 | 0.8665 | -0.18 | +| litr3p_vr | 0.8671 | 0.8489 | -1.82 | 0.8649 | -0.22 | +| occlp_vr | 0.8960 | 0.7902 | -10.58 | 0.8493 | -4.67 | +| primp_vr | 0.6816 | 0.5898 | -9.18 | 0.7046 | +2.29 | +| secondp_vr | 0.9875 | 0.9756 | -1.20 | 0.9836 | -0.40 | +| soil1c_vr | 0.8733 | 0.8443 | -2.90 | 0.8722 | -0.12 | +| soil1n_vr | 0.8733 | 0.8434 | -2.98 | 0.8726 | -0.07 | +| soil1p_vr | 0.8737 | 0.8438 | -2.99 | 0.8726 | -0.11 | +| soil2c_vr | 0.9177 | 0.8997 | -1.80 | 0.9153 | -0.24 | +| soil2n_vr | 0.9173 | 0.9020 | -1.53 | 0.9156 | -0.17 | +| soil2p_vr | 0.9171 | 0.9006 | -1.65 | 0.9155 | -0.16 | +| soil3c_vr | 0.8958 | 0.8810 | -1.48 | 0.8929 | -0.29 | +| soil3n_vr | 0.8945 | 0.8817 | -1.28 | 0.8934 | -0.11 | +| soil3p_vr | 0.8952 | 0.8807 | -1.45 | 0.8936 | -0.16 | +| soil4c_vr | 0.8755 | 0.8469 | -2.86 | 0.8722 | -0.34 | +| soil4n_vr | 0.8756 | 0.8473 | -2.84 | 0.8718 | -0.38 | +| soil4p_vr | 0.8749 | 0.8459 | -2.90 | 0.8718 | -0.32 | +| solutionp_vr | 0.9612 | 0.9461 | -1.50 | 0.9603 | -0.09 | + +## PFT 1D variables (mean R² over PFTs) + +| Variable | nofilter R² | natveg R² | Δ (pp) | natveg_aligned R² | Δ (pp) | +|----------|--------------|-----------|--------|-------------------|--------| +| cpool | 0.9423 | 0.9429 | +0.05 | 0.9383 | -0.40 | +| deadcrootc | 0.7244 | 0.7246 | +0.02 | 0.7218 | -0.26 | +| deadcrootc_storage | 0.4247 | 0.4228 | -0.19 | 0.4229 | -0.18 | +| deadcrootn | 0.7238 | 0.7249 | +0.11 | 0.7212 | -0.26 | +| deadcrootn_storage | 0.4248 | 0.4226 | -0.22 | 0.4228 | -0.20 | +| deadcrootp | 0.7244 | 0.7242 | -0.02 | 0.7220 | -0.25 | +| deadcrootp_storage | 0.4194 | 0.4229 | +0.35 | 0.4228 | +0.34 | +| deadstemc | 0.7230 | 0.7250 | +0.20 | 0.7210 | -0.20 | +| deadstemc_storage | 0.4250 | 0.4228 | -0.22 | 0.4228 | -0.22 | +| deadstemn | 0.7205 | 0.7244 | +0.38 | 0.7205 | +0.00 | +| deadstemn_storage | 0.4248 | 0.4227 | -0.21 | 0.4228 | -0.20 | +| deadstemp | 0.7221 | 0.7244 | +0.23 | 0.7207 | -0.15 | +| deadstemp_storage | 0.4249 | 0.4226 | -0.23 | 0.4227 | -0.22 | +| frootc | 0.9572 | 0.9561 | -0.11 | 0.9493 | -0.79 | +| frootc_storage | 0.6635 | 0.6590 | -0.45 | 0.6625 | -0.10 | +| frootn | 0.9574 | 0.9562 | -0.12 | 0.9564 | -0.10 | +| frootn_storage | 0.6647 | 0.6609 | -0.38 | 0.6627 | -0.20 | +| frootp | 0.9569 | 0.9563 | -0.06 | 0.9570 | +0.01 | +| frootp_storage | 0.6652 | 0.6592 | -0.61 | 0.6626 | -0.26 | +| leafc | 0.9551 | 0.9553 | +0.02 | 0.9545 | -0.06 | +| leafc_storage | 0.6648 | 0.6608 | -0.40 | 0.6632 | -0.16 | +| leafn | 0.9557 | 0.9552 | -0.05 | 0.9552 | -0.05 | +| leafn_storage | 0.6653 | 0.6613 | -0.41 | 0.6626 | -0.27 | +| leafp | 0.9555 | 0.9476 | -0.78 | 0.9552 | -0.03 | +| leafp_storage | 0.6653 | 0.6611 | -0.43 | 0.6633 | -0.20 | +| livecrootc | 0.7216 | 0.7216 | -0.00 | 0.7192 | -0.24 | +| livecrootc_storage | 0.4250 | 0.4225 | -0.25 | 0.4226 | -0.24 | +| livecrootn | 0.7217 | 0.7223 | +0.06 | 0.7192 | -0.25 | +| livecrootn_storage | 0.4249 | 0.4227 | -0.22 | 0.4229 | -0.20 | +| livecrootp | 0.7196 | 0.7212 | +0.17 | 0.7188 | -0.08 | +| livecrootp_storage | 0.4249 | 0.4228 | -0.21 | 0.4229 | -0.20 | +| livestemc | 0.7221 | 0.7217 | -0.04 | 0.7199 | -0.23 | +| livestemc_storage | 0.4250 | 0.4227 | -0.23 | 0.4229 | -0.21 | +| livestemn | 0.7220 | 0.7207 | -0.13 | 0.7195 | -0.25 | +| livestemn_storage | 0.4248 | 0.4227 | -0.21 | 0.4229 | -0.19 | +| livestemp | 0.7207 | 0.7206 | -0.01 | 0.7193 | -0.14 | +| livestemp_storage | 0.4246 | 0.4228 | -0.18 | 0.4226 | -0.20 | +| npool | -12.7540 | -8.9153 | +383.87 | -12.7567 | -0.28 | +| ppool | -15.1043 | -10.4005 | +470.39 | -15.1068 | -0.25 | +| tlai | 0.9550 | 0.9513 | -0.37 | 0.9550 | -0.00 | +| totvegc | 0.9698 | 0.9708 | +0.10 | 0.9692 | -0.06 | + +## Largest R² degradations (natveg vs nofilter) + +Worst 15 (most negative Δ = largest drop with natveg): + +- **litr2c_vr** (soil2d): -10.71 pp +- **occlp_vr** (soil2d): -10.58 pp +- **litr2p_vr** (soil2d): -10.33 pp +- **litr2n_vr** (soil2d): -9.46 pp +- **primp_vr** (soil2d): -9.18 pp +- **soil1p_vr** (soil2d): -2.99 pp +- **soil1n_vr** (soil2d): -2.98 pp +- **soil1c_vr** (soil2d): -2.90 pp +- **soil4p_vr** (soil2d): -2.90 pp +- **soil4c_vr** (soil2d): -2.86 pp +- **soil4n_vr** (soil2d): -2.84 pp +- **litr3c_vr** (soil2d): -2.09 pp +- **litr3p_vr** (soil2d): -1.82 pp +- **soil2c_vr** (soil2d): -1.80 pp +- **soil2p_vr** (soil2d): -1.65 pp + +Best 5 (improvement with natveg, among variables with sensible R²): + +- **deadstemn** (pft1d): +0.38 pp +- **deadcrootp_storage** (pft1d): +0.35 pp +- **deadstemp** (pft1d): +0.23 pp +- **cwdn_vr** (soil2d): +0.20 pp +- **deadstemc** (pft1d): +0.20 pp + +## Is the performance degradation concerning? + +### Two different comparisons + +1. **natveg vs nofilter** (different train *and* test): + - natveg has ~33% fewer samples (14k) and is evaluated on a *natveg-only* test set. + - nofilter is evaluated on the *full* 20% holdout. + - So the large aggregate RMSE increase (+47% scalar, +22% soil) is partly from **different test sets**, not just less data. + +2. **natveg_aligned vs nofilter** (same test set, fair comparison): + - Same 20% holdout for both; only the *training* set is natveg-only in natveg_aligned. + - This answers: *If I train with the natveg filter, how much do I lose on the same test?* + +- **natveg vs nofilter** aggregate scalar RMSE: **+47.6%** (worse). +- **natveg_aligned vs nofilter** scalar RMSE: **-6.4%** (slightly better when negative). + +### Verdict: is degradation with the filter concerning? + +**When comparing fairly (natveg_aligned vs nofilter, same test set):** + +- **Scalar fluxes (GPP, NPP, AR, HR)**: R² **improves** or is flat (+0.07 to +0.25 pp). Scalar RMSE **improves** by ~6%. No concern. +- **PFT 1D**: Aggregate RMSE is virtually unchanged (+0.2%). Per-variable R² changes are mostly within ±0.5 pp. No concern. +- **Soil 2D**: Aggregate RMSE is ~6% higher. Most variables are within ±0.5 pp R². Notable drop: **occlp_vr** ≈ -4.7 pp R². litr2*/primp_vr stay similar or improve slightly with aligned run. + +**Conclusion:** Using the **natveg filter is not concerning** for overall performance when the same test set is used. The aligned run (train on natveg, test on same holdout as nofilter) is slightly *better* on scalars and similar on PFT/soil, with **occlp_vr** as the only variable with a clear drop (~4.7 pp). If your science prioritizes natural vegetation and occlp is not central, the filter is reasonable to apply. + +The **natveg** run (filter-before-split, 14k samples) looks much worse mainly because it is evaluated on a different (natveg-only) test set and with less training data; that comparison is not apples-to-apples for "degradation with filter." diff --git a/docs/NewPlan_5P_TwoRegions.md b/docs/NewPlan_5P_TwoRegions.md new file mode 100644 index 0000000..865a58a --- /dev/null +++ b/docs/NewPlan_5P_TwoRegions.md @@ -0,0 +1,148 @@ +## New plan: 5P, two-region focus + +You are now in a very specialized regime: + +- **Regions**: Amazon + (sub‑Saharan) Africa only. +- **Targets**: Only the 5 P variables (e.g. occlp, labilep, secondp, primp, plus the remaining P variable). +- **Base models**: + - The best AI base model for P so far is `phase2_pvariable_focus` (built on `natveg_improved`). + - Even this is still not good enough for Amazon + Africa. + - For the rest of the globe, you will keep using the original `natveg_improved` run as-is. + +Given this, it is not efficient to keep trying small variations of “same architecture + heavier P weights” on the full global model. A better strategy is: + +1. **Post‑process / calibrate the existing best model for these regions** (fast, guaranteed improvement). +2. **Optionally add a dedicated small regional 5P model** (residual or stand‑alone) trained only on the two-region data. + +Below are concrete options, roughly ordered by impact vs effort. + +--- + +## 1. Fast win: bias/scale correction on top of phase2 predictions + +Use `phase2` (global) or your best current run as a baseline, then statistically correct its systematic errors in Amazon + Africa. This does **not** require retraining the large CNP model. + +- **For each P variable, layer, and optionally region**: + - Take ground truth vs `phase2` predictions over all Amazon + Africa test cells. + - Fit a simple linear correction per `(variable × layer)` or `(variable × region × layer)`: + + \[ + \text{GT} \approx a \cdot \text{pred} + b + \] + +- **At inference time (for these regions only)**, apply: + + \[ + \hat{Y}_{\text{corrected}} = a \cdot \hat{Y}_{\text{phase2}} + b + \] + +- **Implementation notes**: + - You already have CSVs with predictions and GT, so estimating \(a, b\) is straightforward (e.g. per-variable linear regression). + - Pros: + - Very cheap and quick to implement. + - Deterministic improvement if errors are mostly bias/scale (which seems to be the case for occlp, labilep, secondp, primp). + - Usage: + - When generating restarts for Amazon + Africa, apply this correction. + - Elsewhere, continue to use raw `natveg_improved` (or your preferred global baseline). + +--- + +## 2. Residual regional 5P model (recommended “real” model change) + +Instead of re‑training the big CNP net, train a small residual model **only on Amazon + Africa data**. + +- **Inputs**: + - Original inputs (climate, soil, etc.). + - `phase2` 5P predictions at each cell/layer. + +- **Target**: + + \[ + \text{Residual} = \text{GT} - \text{phase2\_pred} + \] + + for each of the 5 P variables (per layer). + +- **Output**: + - Residual for each \(5\text{P} \times \text{layer}\). + - Final prediction: + + \[ + \hat{Y}_{\text{final}} = \hat{Y}_{\text{phase2}} + \hat{Y}_{\text{residual}} + \] + +- **Training**: + - Train this small model only on two‑region training cells. + - Loss focuses purely on the 5 P variables (no other outputs). + +- **Why this helps**: + - `phase2` already captures global structure and physical relationships. + - The residual model only needs to learn **regional corrections** (bias, shape) for high‑P tropics. + - Capacity is dedicated entirely to these 5 variables in these two regions. + +- **Implementation options**: + - A standalone PyTorch script that reads: + - Per‑cell 5P GT / prediction CSVs, plus features. + - Or a variant of your current architecture with: + - Just the soil‑2D head and 5 outputs. + - `region_boxes` to select Amazon + Africa. + - Precomputed `phase2` predictions fed in as additional inputs. + +--- + +## 3. True 5P‑only regional model from scratch + +Your previous two‑region model still used the full 25 soil‑2D outputs, just with different weights. For this very targeted use‑case, a more radical but cleaner option is a dedicated **5P‑only** model: + +- **Architecture / IO**: + - Change the IO list so only the 5 P variables are outputs (matrix head size \(= 5 \times \text{layers}\)). + +- **Data**: + - Train only on two‑region data using `region_boxes`. + +- **Training setup**: + - Use individual normalization and strong 5P loss weights. + - Allow more epochs and/or slightly higher LR since the task is smaller. + +- **Advantages**: + - Smaller head, more effective capacity per P variable. + - No competition from non‑P soil variables in the loss. + +- **Trade‑offs**: + - Larger pipeline change: + - IO list changes. + - Model sizing needs to be updated. + - Inference / restart‑generation scripts must be updated to use the 5P‑only head. + - This is the “dedicated specialist” model if the residual approach still isn’t sufficient. + +--- + +## 4. Loss and normalization tweaks for the two regions + +Given the very different distributions in Amazon + Africa vs the rest of the world (you already see GT ratios \(\sim 2.5\text{–}4.8\)), you can tune the loss and scaling specifically for regional / residual models: + +- **Loss choices**: + - Prefer \(L_1\) / MAE or a less tail‑heavy variant than `log1p_huber`, so you do not over‑focus extreme P outliers. + - Optionally a **quantile loss** (e.g. 0.5 or 0.7) if you want asymmetric penalties for under/over‑prediction. + +- **Normalization**: + - Keep individual normalization per variable. + - Optionally re‑scale targets so typical two‑region ranges map to \(\mathcal{O}(1)\). + - Your current individual scalers likely already help; only tweak if diagnostics show saturation or strong skew. + +These tweaks are most relevant for options **2** and **3** (residual or 5P‑only regional models). + +--- + +## 5. Practical recommendation / priority order + +Given limited time and the ROI curve: + +1. **Immediately**: + - Implement **bias/scale correction** (option 1) on top of `phase2` for Amazon + Africa. + - This should give noticeably better 5 P performance at those sites with minimal new code. +2. **Next, if you want a true ML improvement**: + - Implement a small **residual two‑region 5P model** (option 2) and evaluate it at Amazon and several African sites. +3. **Only if the above is still insufficient**: + - Move to a full **5P‑only regional model** from scratch (option 3), plus the loss/normalization tweaks in option 4. + diff --git a/docs/PHASE2_TROPICAL_SOILP_STRATEGY_AND_VALIDATION.md b/docs/PHASE2_TROPICAL_SOILP_STRATEGY_AND_VALIDATION.md new file mode 100644 index 0000000..488ac04 --- /dev/null +++ b/docs/PHASE2_TROPICAL_SOILP_STRATEGY_AND_VALIDATION.md @@ -0,0 +1,166 @@ +# Phase 2 Tropical Soil-P Strategy: Configuration and Validation Matrix + +**Purpose:** Define the two-phase global + tropical P strategy, the Phase 2 training configuration, and the validation/verification matrix for model efficiency. + +--- + +## 1. Strategy Overview + +| Phase | Scope | Role | +|-------|--------|------| +| **Phase 1** | Global domain | Train global natveg model (existing Phase 1 config). Run full pipeline to produce a **global** AI-updated restart file. | +| **Phase 2** | Tropical-only | Train a **separate** AI model on **tropical-only** data with **extreme attention to soil P variables**. Use Phase 2 only to improve P predictions in the tropics. | + +**Restart merge:** Start from the Phase 1 global restart. Run Phase 2 inference (on the same grid or tropical subset). Then update the restart file so that **only in tropical grid cells** the CNP_IO variables are replaced by **Phase 2 predictions** (using `ai_predictions_to_restart.py --tropical-lat-range`). Extratropical cells keep Phase 1 values. + +- **Config (Phase 1):** Use existing global natveg config (e.g. `training_config_experiment_3_global_natveg_improved_occlp.json` or your current Phase 1). +- **Config (Phase 2):** `config/training_config_phase2_tropical_soilp_only.json` (tropical-only, extreme soil-P weights). + +--- + +## 2. Phase 2 Configuration + +- **File:** `config/training_config_phase2_tropical_soilp_only.json` +- **Description:** Tropical-only natveg training with extreme focus on soil P. Same variable set as global CNP_IO; loss and variable weights emphasize P. + +**Main settings:** + +| Section | Key settings | +|---------|----------------| +| **data_filtering_config** | `tropical_only: true`, `tropical_lat_range: [-30, 30]`, `natveg_only: true` | +| **training_hyperparameters** | `matrix_loss_weight: 2.5`, low scalar/vector (0.06 / 0.4), `litter_p_loss_weight: 4` | +| **soil2d_weights** | `occlp_vr: 45`, `primp_vr: 26`, `labilep_vr: 18`, `solutionp_vr: 18`, `secondp_vr: 14`, soil*p_vr and litr*p_vr elevated | +| **tail_aware_weights** | Same P variables emphasized for tail-aware loss | + +**Training command (from repo root):** + +```bash +python train_cnp_model.py \ + --training-config-json config/training_config_phase2_tropical_soilp_only.json \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --use-tva4km \ + --tropical-only \ + --tropical-lat-range -30,30 \ + --output-dir-suffix phase2_tropical_soilp_only +``` + +The config’s `data_filtering_config.tropical_only` and `tropical_lat_range` are read by the trainer when `--tropical-only` and `--tropical-lat-range` are passed (see `train_cnp_model.py`). + +--- + +## 3. Validation Matrix (Model Efficiency) + +Use this matrix to validate and verify both phases and the merged restart. + +### 3.1 Phase 1 (Global) – Validation + +| Check | Metric / method | Target / note | +|-------|------------------|----------------| +| Global scalar | GPP, NPP, AR, HR: RMSE, R² from `cnp_metrics.json` | R² > 0.97 typical | +| Global PFT1D | pft_1d_rmse, per-PFT R² | No severe degradation | +| Global soil 2D | Per-variable RMSE/R² (e.g. occlp_vr, labilep_vr, solutionp_vr) | Baseline for comparison | +| Restart file | Run `ai_predictions_to_restart.py`, then ELM restart run | No runtime errors; restart loads | + +### 3.2 Phase 2 (Tropical-only) – Validation + +| Check | Metric / method | Target / note | +|-------|------------------|----------------| +| Tropical holdout scalar | GPP, NPP, AR, HR R² on **tropical test set** | Comparable or better than Phase 1 on tropics | +| Tropical holdout P (primary) | **occlp_vr**, **labilep_vr**, **solutionp_vr**: R², RMSE, NRMSE (all layers or layer-agg) | R² > 0.99 for occlp_vr; clear improvement over Phase 1 tropical | +| Tropical holdout P (other) | primp_vr, secondp_vr, soil1p_vr–soil4p_vr, litr2p_vr, litr3p_vr | Improved vs Phase 1 in tropics | +| Phase 2 cnp_metrics.json | Y_occlp_vr_layer*_r2, Y_labilep_vr_*, Y_solutionp_vr_* | Record for reporting | +| Site: Amazon | occlp_vr, solutionp_vr, labilep_vr at (303.75, −17.43) vs ground truth | Phase 2 closer to GT than Phase 1 (see §4) | + +### 3.2b Entire-dataset inference for Phase 2 merge + +Phase 2 is trained **tropical-only**. The run’s `cnp_predictions` folder contains only the **validation subset** (~1k gridcells) written at training time. Inference **without** `--inference-full-grid` uses the same data scope as training: tropical filter is read from the run’s `cnp_config.json` (since we now save `tropical_only` and `tropical_lat_range`), so you get the **full tropical train+validation** set (~6k samples), not just 1k. To merge P variables into the **global** restart, you need predictions on the full model grid (~21k gridcells). + +1. **Run inference on full grid:** From the Phase 2 run directory, run `run_inference_all.py` with **`--inference-full-grid`** so the dataloader uses all gridcells (tropical_only=False). Output goes to `cnp_inference_entire_dataset/cnp_predictions`. +2. **Build NetCDF** from `cnp_inference_entire_dataset/cnp_predictions` (not from `cnp_predictions`). +3. **Run `ai_predictions_to_restart.py`** with that NetCDF and `--variables-to-update @config/phase2_soilp_variables_to_update.txt`. + +- **Without** `--inference-full-grid`: inference uses the training data scope (tropical-only run → ~6k samples; global run → ~21k). Good for evaluation on the same domain as training. +- **With** `--inference-full-grid`: inference always uses the full global grid (~21k). Required for merging Phase 2 P predictions into a global restart. + +If you use the small NetCDF built from `cnp_predictions` (~1k) or from inference without the flag (~6k tropical), the restart updater will map every model gridcell to the nearest of those points; for a proper global P merge use `--inference-full-grid` and the resulting NetCDF. + +### 3.3 Merged Restart (Phase 1 + Phase 2 tropical) – Verification + +| Check | Method | Target / note | +|-------|--------|----------------| +| Tropical cells | Use Phase 2 predictions for lat in `tropical_lat_range` | Only tropical cells updated with Phase 2 | +| Extratropical cells | Unchanged from Phase 1 | No overwrite outside tropics | +| Restart consistency | Run ELM with merged restart | No NaNs; no crashes; NEE/fluxes plausible in tropics | +| Site: Amazon | Extract Amazon point from merged restart; compare occlp_vr (and labilep_vr, solutionp_vr) to Phase 2 prediction and to GT | Values match Phase 2 at Amazon; closer to GT than Phase 1 | + +--- + +## 4. Verification: Amazon and Optional Tropical Sites + +### 4.1 Reference site (Amazon) + +- **Lon:** 303.75°E +- **Lat:** −17.4246° (≈ −17.434553° in grid) + +**Verification steps:** + +1. **Phase 1 only:** From Phase 1 run dir, get predictions at Amazon (e.g. from `cnp_predictions/` or inference CSVs). Record occlp_vr, solutionp_vr, labilep_vr (layers 1–10) vs ground truth. +2. **Phase 2 only:** From Phase 2 run dir, get predictions at Amazon. Compare same variables to GT and to Phase 1 (Phase 2 should be closer; see `docs/AMAZON_SITE_SOILP_COMPARISON_REPORT.md` for comparison template). +3. **Merged restart:** Extract Amazon point from the final restart (e.g. `scripts/extract_elm_restart_point.py`). Compare restart values for occlp_vr (and labilep_vr, solutionp_vr) to Phase 2 predictions and to GT. + +**Commands (from run directory or with absolute paths):** + +```bash +# Site validation stats (Phase 1 or Phase 2 run dir) +python scripts/cnp_result_validationplot_site.py --lon 303.75 --lat -17.4246 --tolerance 0.1 --stats-only + +# Extract Amazon from merged restart +python scripts/extract_elm_restart_point.py \ + --restart-file \ + --lat -17.4246 --lon 303.75 \ + --output-file +``` + +### 4.2 Optional: Additional tropical sites + +Define other tropical (lon, lat) pairs and repeat the same checks (Phase 1 vs Phase 2 vs merged restart) for occlp_vr, solutionp_vr, labilep_vr. Document in the same format as the Amazon report. + +--- + +## 5. Workflow Summary (Commands) + +1. **Phase 1 – Global restart** + - Train: existing Phase 1 global natveg config. + - Inference: full grid. + - Restart: `ai_predictions_to_restart.py` **without** `--tropical-lat-range` → global restart. + +2. **Phase 2 – Tropical P model** + - Train: `config/training_config_phase2_tropical_soilp_only.json` with `--tropical-only --tropical-lat-range -30,30`. + - Inference: run on full grid (or tropical-only subset); only tropical cells will be used when merging. + +3. **Merge Phase 2 into Phase 1 restart** + - Input: Phase 1 global restart. + - Predictions: Phase 2 inference output. + - Update only tropical cells and only P (or selected) variables: + `ai_predictions_to_restart.py ... --tropical-lat-range -30,30 --variables-to-update occlp_vr,labilep_vr,solutionp_vr,...` + - `--variables-to-update`: comma-separated list of variable names (e.g. `occlp_vr,labilep_vr,solutionp_vr,...`) or a file path with `@` prefix (e.g. `@config/phase2_soilp_variables_to_update.txt`). Omit to update all CNP_IO variables in tropical cells. + - Output: merged restart (extratropical = Phase 1; tropical = Phase 1 except selected variables replaced by Phase 2). + +4. **Validation** + - Phase 1: `cnp_metrics.json`, global validation_stats, quality report. + - Phase 2: `cnp_metrics.json`, tropical validation_stats, Amazon (and optional sites) P comparison. + - Merged: Amazon (and optional) extraction and comparison; ELM run with merged restart. + +--- + +## 6. Checklist (Validation and Verification) + +- [ ] Phase 1 trained and global restart produced; metrics and quality report reviewed. +- [ ] Phase 2 trained with `training_config_phase2_tropical_soilp_only.json`; tropical holdout metrics recorded. +- [ ] Phase 2 P variables (occlp_vr, labilep_vr, solutionp_vr) on tropical holdout: R²/RMSE meet targets and improve vs Phase 1. +- [ ] Amazon site: Phase 2 predictions vs GT and vs Phase 1 documented (e.g. table as in `AMAZON_SITE_SOILP_COMPARISON_REPORT.md`). +- [ ] Merged restart created with `--tropical-lat-range -30,30`. +- [ ] Amazon (and optional sites) extracted from merged restart; P variables match Phase 2 and are closer to GT than Phase 1. +- [ ] ELM run with merged restart: no errors; tropical NEE/fluxes plausible. + +This completes the configuration and validation matrix for the Phase 2 tropical soil-P strategy. diff --git a/docs/PLAN_TWO_REGION_FIVE_P_TRAINING.md b/docs/PLAN_TWO_REGION_FIVE_P_TRAINING.md new file mode 100644 index 0000000..6e09028 --- /dev/null +++ b/docs/PLAN_TWO_REGION_FIVE_P_TRAINING.md @@ -0,0 +1,180 @@ +# Plan: Two-Region, Five P–Variable–Only Model + +This document outlines how to train a model focused on **two geographic regions** and **five soil P variables** only, and how to use it for prediction in those regions. + +--- + +## 1. Target regions (lat/lon) + +- **Region 1 – Amazon:** + Latitude **10°N to 30°S** (−30 to 10), Longitude **90°W to 30°W** (270° to 330° in 0–360°). +- **Region 2 – Central Africa:** + Latitude **15°N to 15°S** (−15 to 15), Longitude **0° to 30°E** (0° to 30° in 0–360°). + +So in (lat_min, lat_max, lon_min, lon_max) with lon in **0–360**: + +- Amazon: `(-30, 10, 270, 330)` +- Africa: `(-15, 15, 0, 30)` + +--- + +## 2. Five P variables + +- `labilep_vr` +- `occlp_vr` +- `solutionp_vr` +- `secondp_vr` +- `primp_vr` + +--- + +## 3. Dataloader: restrict to these two regions + +**Implemented:** The dataloader now supports **region boxes**. + +- **Config:** `DataConfig.region_boxes` + - Type: `Optional[List[Tuple[float, float, float, float]]]` + - Each tuple: `(lat_min, lat_max, lon_min, lon_max)` with **longitude in 0–360°**. +- **Behavior:** In `data_loader_individual.preprocess_data()`, only rows whose (lat, lon) fall **inside at least one** box are kept. Lat/lon columns are resolved automatically (e.g. `Latitude`/`Longitude`). +- **Where to set:** Via training config JSON under `data_filtering_config.region_boxes`. + +**Example for the two regions:** + +```json +"data_filtering_config": { + "region_boxes": [ + [-30, 10, 270, 330], + [-15, 15, 0, 30] + ], + "natveg_only": true, + "natveg_filter_before_split": false +} +``` + +- **Training (main script):** `train_cnp_model.py` applies `data_filtering_config` from the JSON (including `region_boxes`) when you pass `--training-config-json`. +- **Finetuning:** `run_finetuning_json.py` also applies `data_filtering_config.region_boxes` from the unified config. + +So: **use the existing dataloader with a training config that sets `region_boxes`**; no extra code path is required. + +--- + +## 4. “Five P only” training: two practical options + +The current model has one scalar head, one PFT-1D head, and one soil2D (matrix) head. The matrix head size is `n_2d_vars * rows * cols`. So you can either keep the full variable set and focus loss on the 5 P vars, or reduce the variable list to 5 P only (smaller matrix head). + +### Option A (recommended): Full variable list, 5 P–focused weights + +- **Variable list:** Keep using the full CNP list (e.g. `CNP_IO_updated9_dev_dw.txt`) so the **model architecture is unchanged** (same heads, same number of 2D variables). +- **Data:** Use **only** the two regions above via `region_boxes`. +- **Loss focus:** In the training config JSON, set **`variable_weights.soil2d_weights`** so that the five P variables have **high weight** (e.g. 20–50) and all other soil2d variables **low or zero** (e.g. 0.1 or 0). The model will then prioritize fitting the 5 P variables in these two regions. +- **Pros:** No change to model construction or inference pipeline; reuse existing inference/restart scripts. +- **Cons:** Model still has parameters for non-P outputs; training is a bit less “focused” than a true 5-P-only head. + +**Example (excerpt) in training config JSON:** + +```json +"variable_weights": { + "soil2d_weights": { + "labilep_vr": 40, + "occlp_vr": 50, + "solutionp_vr": 40, + "secondp_vr": 40, + "primp_vr": 30, + "cwdc_vr": 0.1, + "cwdn_vr": 0.1, + "cwdp_vr": 0.1, + "litr2c_vr": 0.1, + "litr2n_vr": 0.1, + "litr2p_vr": 0.1, + "litr3c_vr": 0.1, + "litr3n_vr": 0.1, + "litr3p_vr": 0.1, + "soil1c_vr": 0.1, + "soil1n_vr": 0.1, + "soil1p_vr": 0.1, + "soil2c_vr": 0.1, + "soil2n_vr": 0.1, + "soil2p_vr": 0.1, + "soil3c_vr": 0.1, + "soil3n_vr": 0.1, + "soil3p_vr": 0.1, + "soil4c_vr": 0.1, + "soil4n_vr": 0.1, + "soil4p_vr": 0.1 + } +} +``` + +(Adjust names to match exactly your CNP_IO 2D list.) + +### Option B: Five P variables only (smaller matrix head) + +- **Variable list:** Create a **minimal** CNP_IO (or config) where the **only** 2D soil **outputs** are the five P variables. Inputs can remain the same (or be pruned to a minimal set); scalar/PFT-1D can be minimal (e.g. one variable each) if the code allows. +- **Model:** The matrix head becomes `5 * matrix_rows * matrix_cols`. This requires that the rest of the pipeline (data_info, trainer, inference) is built from this reduced list so that only 5 P columns are expected and written. +- **Pros:** Smaller head; training is only for the 5 P variables. +- **Cons:** Requires a dedicated variable list and possibly small code paths to support “soil2d-only” or minimal scalar/PFT-1D; inference and restart scripts must use the same reduced list. + +**Summary:** For “reasonably just predictions for these 5 P variables in these two regions,” **Option A** is the least intrusive and is supported with the new `region_boxes` and existing weight config. + +--- + +## 5. Suggested workflow (Option A) + +1. **Create a training config JSON** (e.g. `config/training_config_two_region_five_p.json`) that: + - Sets `data_filtering_config.region_boxes` to `[[-30, 10, 270, 330], [-15, 15, 0, 30]]`. + - Sets `data_filtering_config.natveg_only` (and optionally `natveg_filter_before_split`) as desired. + - Sets `variable_weights.soil2d_weights` so that the five P variables have high weight and all other soil2d variables low/zero (as in the example above). + - Keeps other sections (training_hyperparameters, reproducibility_config, etc.) aligned with your current phase2/tropical setup. + +2. **Train:** + ```bash + python train_cnp_model.py \ + --training-config-json config/training_config_two_region_five_p.json \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --use-tva4km \ + --output-dir-suffix two_region_five_p + ``` + The dataloader will load only grid cells in the two regions (and apply natveg if enabled); the loss will be dominated by the 5 P variables. + +3. **Inference:** + Use the same variable list and model as today: + - **Full-grid inference:** Run inference as usual; the model will output all variables (including the 5 P). For “focus” you can either use the outputs everywhere or only trust/use the 5 P predictions inside the two regions. + - **Restart update:** Use `ai_predictions_to_restart.py` with `--variables-to-update` set to the five P names and, if available, a **spatial mask** so that only cells inside the two boxes are updated (see below). + +4. **Optional – inference only in the two regions:** + If you want predictions **only** for the two regions (e.g. to avoid applying the regional model outside those areas): + - Either run inference on the full grid and then **discard or mask** predictions outside the two boxes, or + - Add an option to the inference script to **subset the dataloader** (e.g. by passing the same `region_boxes` or a list of (lat, lon) cells) so that only those grid cells are loaded and predicted. That would be a small extension of the current inference pipeline. + +5. **Restart / CSV output:** + For writing restarts or CSVs “only in the two regions,” you need a spatial mask in the script that writes restarts/CSVs (e.g. `ai_predictions_to_restart.py`). Options: + - Add a `--region-boxes` (or similar) argument that takes the same two boxes and only updates grid cells whose (lat, lon) fall in one of the boxes; or + - Precompute a list of gridcell indices for the two regions and pass that as a mask. + This keeps the “five P only” and “two regions only” semantics consistent from training to application. + +--- + +## 6. Checklist (last task to finish) + +- [x] **Dataloader:** Add `region_boxes` to `DataConfig` and filter in `preprocess_data()` so only cells in the two regions are used (done). +- [x] **Config wiring:** Support `region_boxes` in `data_filtering_config` in `train_cnp_model.py` and `run_finetuning_json.py` (done). +- [x] **Training config:** Add `config/training_config_two_region_five_p.json` with the two boxes and 5 P–focused soil2d weights (done; see also `docs/IMPROVE_5P_PREDICTIONS_SUGGESTIONS.md`). +- [ ] **Train** a run with that config and your chosen variable list. +- [ ] **Validate** on the two regions (e.g. Amazon site + a few Africa cells) via existing validation/plot scripts. +- [ ] **Inference:** Decide whether to run full-grid and mask, or add a region-only inference option; implement restart/CSV masking for the two regions if you want updates only there. + +--- + +## 7. Longitude convention + +- **0–360°:** 90°W = 270°, 30°W = 330°, 0°E = 0°, 30°E = 30°. +- The dataloader normalizes longitude to 0–360 when applying `region_boxes`, so if your data use −180..180, they are converted before comparison. + +--- + +## 8. Summary + +- **Regions:** Amazon (lat −30–10, lon 270–330) and Central Africa (lat −15–15, lon 0–30); lon in 0–360. +- **Variables:** Train with strong focus on the five P variables (Option A: full list + weights; Option B: minimal list and smaller matrix head). +- **Dataloader:** Use `data_filtering_config.region_boxes` in your training config; no dataloader code change needed beyond what’s already implemented. +- **Next steps:** Add the JSON config, run training, then add optional inference/restart masking so that “five P only” and “two regions only” are applied end-to-end where you need them. diff --git a/docs/README_solutionp_vr_ai_restart_comparison.md b/docs/README_solutionp_vr_ai_restart_comparison.md new file mode 100644 index 0000000..215495b --- /dev/null +++ b/docs/README_solutionp_vr_ai_restart_comparison.md @@ -0,0 +1,29 @@ +# solutionp_vr plots in ai_restart_comparison_plots_vs_target800result + +This note applies to the folder `cnp_results/run_20260221_120117_withCNonlyInit/ai_restart_comparison_plots_vs_target800result`. + +## Why they were missing + +The folder previously had `occlp_vr_lev*.png` but not `solutionp_vr_lev*.png` because the original comparison run used a variable list (or reference) that omitted **solutionp_vr**. Both the AI restart and the reference (ELM year 0801) contain `solutionp_vr`; it just was not included in the plotted variable set. + +## What was added + +1. **solutionp_vr.png** – Copied from `cnp_results/amazon_4p_plots_gt_and_two_restarts/solutionp_vr.png` (4 P variable comparison style). + +2. **solutionp_vr_lev0.png … solutionp_vr_lev9.png** – Generated by re-running `ai_restart_comparison.py` with the full CNP_IO list (including solutionp_vr), same AI restart and reference. If the full run is still in progress, these files will appear when the script reaches the 2D soil variables (solutionp_vr is last). + +## How to re-run the comparison (e.g. to refresh solutionp_vr only) + +From the **run** directory `cnp_results/run_20260221_120117_withCNonlyInit`: + +```bash +python ../../scripts/ai_restart_comparison.py \ + --variable-list ../../CNP_IO_updated9_dev_dw.txt \ + --ai-restart CNInit_updated9_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --original-restart ../../ELM_data/20251201_TRENDY2024_default_ICB1850CNPRDCTCBC.elm.r.0801-01-01-00000.nc \ + --output-dir ai_restart_comparison_plots_vs_target800result \ + --plot-all +``` + +- **AI restart:** CNInit_updated9_...0021-01-01 (year 21). +- **Reference (target 800):** ELM_data/...0801-01-01 (year 801). diff --git a/docs/RUN_COMPARISON_20260226_114546_vs_114659.md b/docs/RUN_COMPARISON_20260226_114546_vs_114659.md new file mode 100644 index 0000000..ab4c3d4 --- /dev/null +++ b/docs/RUN_COMPARISON_20260226_114546_vs_114659.md @@ -0,0 +1,106 @@ +# Run comparison: 114546 (no-filter) vs 114659 (natveg-aligned) + +**Date:** 2026-02-26 +**Runs:** `cnp_results/run_20260226_114546` and `cnp_results/run_20260226_114659` (directories may be named with `_nofiter` / `_natveg_aligned` suffixes) + +This report summarizes the comparison of two full training runs used for natveg-alignment verification: a no-filter run and a natveg-aligned run with the same test set. + +--- + +## 1. Run identification + +| Run | Directory | Config | +|-----|-----------|--------| +| **Run A** | `run_20260226_114546` | No natveg filter (`natveg_only: false`) | +| **Run B** | `run_20260226_114659` | Natveg-aligned (`natveg_only: true`, `natveg_filter_before_split: false`) | + +Run B uses “split then filter train”: the same random split and seed as Run A, so the **test set is identical**; only training is restricted to natural-vegetation gridcells. + +--- + +## 2. Test set alignment + +Verified with `scripts/verify_natveg_test_subset.py`: + +- **No-filter (114546):** 4,166 gridcells +- **Natveg (114659):** 4,166 gridcells +- **Overlap:** 4,166 (100%) +- **Conclusion:** Natveg test set is a subset of the no-filter test set (here they are identical). Validation is on the same test set. + +--- + +## 3. Validation statistics (920 variable/PFT/layer rows) + +Metrics are from `validation_stats.csv` (validation with natveg filter applied to included gridcells). Rows matched on type, variable, PFT, layer. + +| Metric | Run 114546 (no-filter) | Run 114659 (natveg) | +|--------|------------------------|----------------------| +| **Mean R²** | **0.6055** | 0.6049 | +| **Mean RMSE** | **84.06** | 84.45 | +| **Run 114546 better (higher R²)** | **439** | — | +| **Run 114659 better (higher R²)** | — | 254 | +| **Ties** | 227 | 227 | + +**By prediction type:** + +| Type | R² 114546 | R² 114659 | +|------|------------|-----------| +| 1D | **0.5058** | 0.5038 | +| Scalar | 0.9881 | **0.9896** | +| 2D | 0.8514 | **0.8539** | + +Overall, the no-filter run has slightly better mean R² and lower mean RMSE; the natveg run is marginally better on Scalars and 2D. + +--- + +## 4. Quality summary reports + +From `analysis/quality_summary_report.txt` (Good/OK/Bad by R² and relative RMSE/MAE thresholds): + +| Stat | Run 114546 (no-filter) | Run 114659 (natveg) | +|------|------------------------|----------------------| +| **Good** | 736 (80.0%) | 734 (79.8%) | +| **OK** | 137 (14.9%) | 139 (15.1%) | +| **Bad** | 47 (5.1%) | 47 (5.1%) | +| **Good variables** | 58 | 59 | + +Bad counts by type are the same (1D: 26, 2D: 21). Worst variables are the same (e.g. `ppool`, `npool`, `litr2*_vr`, `primp_vr`), with minor differences in which PFTs/layers are bad. + +--- + +## 5. Conclusions + +- **Test sets:** Identical (4,166 gridcells); comparison is like-for-like. +- **Overall:** Run **114546 (no-filter)** is marginally better: higher mean R², lower mean RMSE, and more variable/PFT/layer rows with higher R² (439 vs 254). +- **By type:** No-filter wins on 1D; natveg wins slightly on Scalar and 2D. +- **Quality bands:** Nearly identical; same Bad count; small differences in Good/OK counts and in which variables are “Good”. + +The two runs are very close; the no-filter run has a small overall edge on this shared test set when both are evaluated with the natveg filter on. + +--- + +## 6. How to reproduce + +**Test set verification:** + +```bash +python scripts/verify_natveg_test_subset.py \ + cnp_results/run_20260226_114546/cnp_predictions \ + cnp_results/run_20260226_114659/cnp_predictions +``` + +(Use the actual run directory names if renamed, e.g. `run_20260226_114546_nofiter` and `run_20260226_114659_natveg_aligned`.) + +**Validation stats and quality-report comparison:** + +```bash +python scripts/compare_cnp_validation_runs.py \ + cnp_results/run_20260226_114546 \ + cnp_results/run_20260226_114659 +``` + +Add `--report comparison.txt` to write a short report to a file. + +Optional: use `--report` to write a short comparison report to a file. + +See also: `docs/NATVEG_ALIGNMENT_VERIFICATION.md`, `docs/EXCLUDED_SAMPLE_ANALYSIS.md`. diff --git a/docs/TROPICAL_OCCLP_TWO_OPTIONS_REPORT.md b/docs/TROPICAL_OCCLP_TWO_OPTIONS_REPORT.md new file mode 100644 index 0000000..b0d695f --- /dev/null +++ b/docs/TROPICAL_OCCLP_TWO_OPTIONS_REPORT.md @@ -0,0 +1,127 @@ +# Tropical occlp_vr Improvement: Single Global vs Two-Model Strategy + +This report compares two strategies for improving occlp_vr predictions in tropical regions (where low occlp_vr contributes to high NEE bias): **Option A** uses a single global model fine-tuned on tropical data; **Option B** uses a global model for non-tropical areas and a separate tropical model for tropical areas, with lat-based switching at inference. + +--- + +## Context + +- **Problem:** In tropical regions, the CNP model underpredicts occlp_vr, which contributes to high NEE in restart simulations. +- **Phase 1:** Global natveg model trained with occlp-focused weights (`training_config_experiment_3_global_natveg_improved_occlp.json`) improves occlp_vr vs the reference run but still underpredicts at tropical sites (e.g. Amazon point lon 303.75, lat −17.43). +- **Goal:** Further improve tropical occlp_vr either by fine-tuning the global model on tropical data or by training/maintaining a dedicated tropical model and using it only in the tropics at inference. + +--- + +## Option A: Single Global Model (Fine-Tune on Tropical) + +### Description + +- **Training:** Phase 1 trains a global model with occlp emphasis. Phase 2 fine-tunes that same model (e.g. via `scripts/run_finetuning.py`) using tropical data and an occlp-focused configuration (e.g. `CNP_IO_finetune_tropics.txt` pointing to Phase 1 checkpoint). +- **Inference:** One model is used for all grid cells globally. + +### Pros + +- **Single model:** One checkpoint and one code path; easier to maintain, deploy, and version. +- **No boundary choice:** No need to define a tropical/non-tropical latitude cutoff or blending zone; behavior is continuous across the globe. +- **Stable extratropics in practice:** With a small learning rate and limited Phase 2 epochs, fine-tuning usually improves tropics without severely degrading extratropical skill. + +### Cons + +- **Possible extratropical drift:** Phase 2 optimizer only sees tropical batches; some forgetting or drift in non-tropical regions is possible (monitor via validation). + +### When to Prefer + +- Default choice when the main goal is to fix tropical occlp_vr without adding operational complexity. +- Use when you want one globally consistent product and simpler validation and release. + +--- + +## Option B: Two Models (Global for Non-Tropical, Tropical for Tropical) + +### Description + +- **Training:** Phase 1 remains the “global” model (used as-is or lightly tuned for non-tropical focus). A separate “tropical” model is trained or fine-tuned **only on tropical data** (e.g. `--tropical-only` with `train_cnp_model.py`, or fine-tune on a tropical-only dataset). +- **Inference:** For each grid cell, if `|lat| <= threshold` (e.g. 23.5° or 30°) use the tropical model; otherwise use the global model. Optionally add a small blending zone to smooth the transition. + +### Pros + +- **Full tropical specialization:** The tropical model can be tuned aggressively for occlp_vr (and other P pools) without any compromise for extratropical performance. +- **No impact on global product:** The global model is unchanged for non-tropical areas; no risk of Phase 2 degrading mid/high latitudes. + +### Cons + +- **Two checkpoints and two code paths:** More to maintain, validate, and document. +- **Boundary handling:** Requires a defined latitude threshold (and possibly blending) and branching logic in the inference pipeline; risk of visible discontinuity at the boundary. + +### When to Prefer + +- After Option A, if validation shows **clear degradation** in non-tropical regions (e.g. occlp_vr or NEE in mid/high latitudes). +- When you explicitly want **two products:** e.g. a global map product and a tropical-focused product for regional studies. + +--- + +## Recommendation + +1. **Start with Option A:** Fine-tune the Phase 1 global model on tropical data using `CNP_IO_finetune_tropics.txt` and the procedure in the runbook. Validate both tropical and extratropical metrics (e.g. occlp_vr at the Amazon point and at a few non-tropical sites). +2. **Adopt Option B only if needed:** If extratropical performance drops unacceptably, keep Phase 1 as the global-only model and introduce a tropical-only model, with inference logic that uses the global model for non-tropical grid cells and the tropical model for tropical grid cells. + +--- + +## Commands: Two-model restart (Option B) — global + tropical merge + +If you use a **tropical-only fine-tuned model** and want one restart that uses the **global model** outside the tropics and the **tropical model** inside the tropics: + +### 1. Build the global restart (Phase 1 model) + +From the **Phase 1 run** directory (e.g. `cnp_results/run_20260303_145913_natveg_improved_occlp`): + +```bash +# Global inference +python ../../scripts/run_inference_all.py --model model.pth --output-dir cnp_inference_entire_dataset --derive-np-from-c + +# NetCDF for restart updater (run from run dir; script expects comparison_results under current or parent dir) +python ../../scripts/ai_predictions_to_netcdf.py + +# Full restart from global model (no tropical filter) +python ../../scripts/ai_predictions_to_restart.py \ + --ai-predictions comparison_results/ai_predictions_for_plotting.nc \ + --restart-file /path/to/original_restart.nc \ + --output updated_restart_global.nc \ + --variable-list ../../CNP_IO_updated9_dev_dw.txt +``` + +### 2. Build tropical predictions NetCDF (fine-tuned model) + +From the **fine-tune run** directory (e.g. `cnp_results/finetune_YYYYMMDD_HHMMSS`): + +```bash +# Inference with tropical model (on full dataset so we have predictions for every gridcell) +python ../../scripts/run_inference_all.py --model model.pth --output-dir cnp_inference_entire_dataset --derive-np-from-c + +python ../../scripts/ai_predictions_to_netcdf.py +``` + +### 3. Replace only tropical region in the global restart + +From the **same fine-tune run** directory, use the **global** restart as input and the **tropical** predictions NetCDF; restrict updates to the tropical band with `--tropical-lat-range`: + +```bash +python ../../scripts/ai_predictions_to_restart.py \ + --ai-predictions comparison_results/ai_predictions_for_plotting.nc \ + --restart-file /path/to/updated_restart_global.nc \ + --output updated_restart_global_with_tropical.nc \ + --variable-list ../../CNP_IO_updated9_dev_dw.txt \ + --tropical-lat-range -30,30 +``` + +Result: `updated_restart_global_with_tropical.nc` has global-model values everywhere except gridcells with latitude in [−30, 30], which are filled from the tropical model. + +--- + +## References + +- Phase 1 config: `config/training_config_experiment_3_global_natveg_improved_occlp.json` +- Phase 2 / fine-tune config: `config/training_config_experiment_3_global_natveg_occlp_phase2_tropics.json` (for standalone Phase 2 training); fine-tune uses `CNP_IO_finetune_tropics.txt` and Phase 1 checkpoint. +- Runbook: `docs/CNP_pipeline_runbook.md` (fine-tune and global inference sections). +- Phase 1 run (example): `cnp_results/run_20260303_145913_natveg_improved_occlp` +- Reference run (pre–occlp focus): `cnp_results/run_20260228_214757_natveg_improved` diff --git a/docs/WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md b/docs/WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md new file mode 100644 index 0000000..48580f1 --- /dev/null +++ b/docs/WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md @@ -0,0 +1,273 @@ +## Workflow: natveg_improved baseline + Phase2 tropical 5P model → 5P bias/scale correction → tropical/two-region restart runs + +This document records the end-to-end procedure to: + +1. Train or reuse a **global baseline CNP model** (`natveg_improved`) and obtain a full global restart. +2. Train or reuse the **Phase2 tropical, P-focused model** (`phase2_pvariable_focus`, using `training_config_phase2_tropical_soilp_only.json`). +3. Run inference with Phase2 (tropical-only or full grid) to generate prediction and ground truth CSVs. +4. Fit and apply **5P bias/scale corrections** for **Amazon + Africa** on top of the Phase2 predictions. +5. Convert the bias-corrected Phase2 5P predictions to NetCDF. +6. Use `ai_predictions_to_restart.py` to overwrite 5P in the **tropical band** of the `natveg_improved` restart (with bias-corrected values in Amazon + Africa), creating updated restart files. +7. Run new simulations from these restarts and validate the results. + +The goal is to have a reproducible protocol you can follow again for future runs. + +--- + +### 1. Train (or select) the global model: `natveg_improved` + +**Objective:** Obtain a trained global model that will serve as the baseline for 5P corrections in the two focus regions. + +- **Inputs/config:** + - Training config JSON used for the `natveg_improved` experiment (e.g. `config/training_config_experiment_3_global_natveg_improved_occlp.json` or the exact file you used). + - Global dataset with all necessary inputs (climate, soil, etc.) and outputs (including soil 2D P variables). +- **Action:** + - Train the model using your standard training pipeline (e.g. `train_cnp_repeat.py` or the corresponding training script). + - Ensure that the run is saved under something like: + - `cnp_results/run_20260228_214757_natveg_improved/` +- **Result:** + - A trained checkpoint and associated metadata under the `natveg_improved` run directory. + +If the model is already trained, you can **reuse** the existing `natveg_improved` run and skip retraining. + +--- + +### 2. Run entire-domain inference for natveg_improved + +**Objective:** Produce per-variable CSVs with **predictions** and **ground truth** for the full domain, including Amazon and Africa. + +- **Script:** your standard full-domain inference script (e.g. `run_inference_all.py` or its updated equivalent). +- **Key arguments:** + - `--run-dir` pointing to the `natveg_improved` run directory, e.g. + `cnp_results/run_20260228_214757_natveg_improved` + - Any additional options needed for full-domain inference (dataset split, etc.). +- **Expected outputs (directory structure):** + - Under the run directory: + - `cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions/` + - Files like `predictions_Y_labilep_vr.csv`, `predictions_Y_occlp_vr.csv`, etc. + - `cnp_inference_entire_dataset/cnp_predictions/soil_2d_ground_truth/` + - Files like `ground_truth_Y_labilep_vr.csv`, `ground_truth_Y_occlp_vr.csv`, etc. + +Each CSV should contain: + +- `Longitude`, `Latitude` columns. +- One column per layer, named like: + - `Y__col1_layer1`, `Y__col1_layer2`, …, `Y__col1_layer10`. + +These paired GT and prediction CSVs are the basis for fitting the bias/scale corrections. + +--- + +### 2. How the Phase2 tropical 5P model is trained + +The Phase2 run at: + +- `cnp_results/run_20260305_153217_phase2_pvariable_focus` + +uses the unified config: + +- `config/training_config_phase2_tropical_soilp_only.json` + +Key points from that config: + +- **Training domain (data_filtering_config):** + - `"tropical_only": true` + - `"tropical_lat_range": [-30.0, 30.0]` + - `"natveg_only": true` + - → Phase2 is trained **only on natveg grid cells in the latitude band [-30°, 30°]**. +- **Loss focus (variable_weights and tail_aware_weights):** + - Soil 2D P variables get **very large weights**: + - `primp_vr`: 26, `occlp_vr`: 45, `labilep_vr`: 18, `secondp_vr`: 14, `solutionp_vr`: 18. + - Other soil 2D variables (C/N) have more moderate weights. + - Tail-aware weights mirror this emphasis, so **extremes of 5P** are also prioritized. + +**Interpretation:** `phase2_pvariable_focus` is a **tropical-only, P-specialist model**: + +- It sees only **tropical natveg** data during training. +- Within that band, it allocates a large fraction of capacity to **matching soil P pools** (5P) and their vertical structure. + +This explains why, in diagnostics (e.g. `prediction_quality_by_variable.png` and the Amazon +site profiles), Phase2 often outperforms `natveg_improved` on the 5P variables in the +tropics, especially after the additional bias/scale correction step. + +--- + +### 3. Run inference for Phase2 (tropical-only vs full grid) + +You have two useful inference paths. Both ultimately feed into `ai_predictions_to_netcdf.py` +and `ai_predictions_to_restart.py`, but they differ in spatial coverage. + +#### 3.1 Tropical-only inference (replicating the existing 4P tropical restart) + +As documented in `README_updated_restart_phase2_tropical_4p.md`, you can run inference on +only the tropical subset (same filter as training): + +```bash +cd cnp_results/run_20260305_153217_phase2_pvariable_focus +python ../../scripts/run_inference_all.py \ + --model cnp_predictions/model.pth \ + --output-dir cnp_inference_tropical_only +``` + +- This uses the same `tropical_only` filter as training: + - Only cells in [-30°, 30°] appear in the prediction CSVs / NetCDF. + +You can then convert these predictions to NetCDF and update the natveg restart in the +**tropical band only**: + +```bash +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_inference_tropical_only/cnp_predictions \ + --output cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_tropical_only.nc + +python scripts/ai_predictions_to_restart.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_tropical_only.nc \ + --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --output cnp_results/run_20260305_153217_phase2_pvariable_focus/updated_restart_phase2_tropical_4p.nc \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --variables-to-update occlp_vr,labilep_vr,solutionp_vr,secondp_vr \ + "--tropical-lat-range=-30,30" +``` + +- **Effect:** + - For tropical cells (\([-30, 30]\)) that exist in the AI NetCDF: + - The 4 listed P variables are overwritten with **raw Phase2** predictions. + - Outside \([-30, 30]\), and for all non-P variables: + - Values remain from the **natveg_improved** base restart. + +If you later rerun inference after bias/scale correction (so that the tropical CSVs already +contain bias-corrected 5P values for Amazon + Africa), you can: + +- Point `ai_predictions_to_netcdf.py` at the **bias-corrected tropical CSVs** instead. +- Use `--variables-to-update` including all 5P variables: + - `labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr`. + +#### 3.2 Full-grid inference (for a single global NetCDF) + +For a cleaner global bias/scale workflow (global CSVs, corrections only in Amazon + Africa, +but predictions everywhere), you can instead run: + +```bash +python scripts/run_inference_all.py \ + --run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --inference-full-grid +``` + +- With `--inference-full-grid`, `run_inference_all.py` explicitly sets: + - `tropical_only=False` + - Clears `region_boxes` + - → Predictions are written for **all natveg grid cells globally**. + +This is the recommended path if you want a **single global NetCDF** where: + +- Inside Amazon + Africa, 5P values are **bias/scale–corrected Phase2**. +- Elsewhere (tropics and extratropics), 5P values are **raw Phase2**. + +--- + +### 4. Apply 5P bias/scale correction to Phase2 (Amazon + Africa only) + +**Objective:** Fit and apply regional 5P corrections, but using Phase2 +predictions as the baseline. The fit and corrections are restricted to the two-region +boxes (Amazon + Africa); elsewhere the CSVs remain raw Phase2. + +- **Script:** `scripts/apply_5p_bias_scale_correction.py` +- **Key command (example):** + +```bash +python scripts/apply_5p_bias_scale_correction.py \ + --run-dir cnp_results/run_20260305_153217_phase2_pvariable_focus \ + --region-config-json config/training_config_two_region_five_p.json \ + --output-subdir soil_2d_predictions_5P_bias_corrected_phase2 +``` + +- **Outputs:** + - Corrected predictions under: + - `cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_5P_bias_corrected_phase2/` + - `predictions_Y__bias_corrected.csv` for each of the 5 P variables. + - Cells inside Amazon + Africa use bias-corrected values; other cells remain raw Phase2. + - Bias/scale parameters under: + - `cnp_results/run_20260305_153217_phase2_pvariable_focus/analysis/bias_scale_params_5P_two_regions.json` + - (Optional) copy this JSON into `docs/` with a descriptive name, e.g.: + - `docs/bias_scale_params_5P_two_regions_phase2_pvariable_focus.json` + +These corrected Phase2 CSVs correspond to the `phase2_pvariable_focus_bias_corrected` +curve you see in the Amazon site profile plots. + +--- + +### 5. Convert Phase2 bias-corrected CSVs to NetCDF (global, with two-region corrections) + +Use `scripts/ai_predictions_to_netcdf.py` to build a global NetCDF of Phase2 predictions +with bias/scale corrections applied only in Amazon + Africa. **You must pass +`--soil-2d-bias-corrected-subdir soil_2d_predictions_5P_bias_corrected_phase2`** so that +the script reads the 5P CSVs produced by `apply_5p_bias_scale_correction.py`; otherwise +it only loads raw predictions from `soil_2d_predictions` and the restart will not use +bias-corrected 5P. + +```bash +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_inference_entire_dataset/cnp_predictions/ \ + --soil-2d-bias-corrected-subdir soil_2d_predictions_5P_bias_corrected_phase2 \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --output cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc +``` + +The resulting NetCDF has: + +- Global domain (all natveg grid cells). +- Inside Amazon + Africa: 5P values are bias/scale–corrected Phase2. +- Outside Amazon + Africa: 5P values match raw Phase2 predictions. + +--- + +### 6. Generate Phase2-based bias-corrected tropical restart files + +**Objective:** Use `phase2_pvariable_focus_bias_corrected` predictions as the 5P inputs +for new restart files in the **tropical band**, while keeping extratropical cells as in +the natveg_improved base restart. + +1. **Select the base restart and Phase2 NetCDF:** + - Base restart: + - `cnp_results/run_20260228_214757_natveg_improved/updated_restart_...elm.r.0021-01-01-00000.nc` + - Phase2 bias-corrected NetCDF: + - `cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc` + +2. **Map corrected 5P values into the restart template (tropics only):** + +```bash +python scripts/ai_predictions_to_restart.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc \ + --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --output cnp_results/run_20260305_153217_phase2_pvariable_focus/updated_restart_phase2_5P_bias_corrected_tropical.nc \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --variables-to-update labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr \ + "--tropical-lat-range=-30,30" +``` + +- `--tropical-lat-range=-30,30`: + - Only tropical cells (lat in [-30, 30]) have their 5P overwritten. + - Outside the tropics, all variables, including 5P, remain as in the natveg_improved base restart. +- Within the tropical band: + - Amazon + Africa cells use **bias-corrected Phase2** 5P. + - Other tropical cells use **raw Phase2** 5P (since those cells were not bias-corrected). + +3. **Write and label Phase2-based restart files:** + - Save in a separate directory, e.g.: + - `cnp_results/run_20260305_153217_phase2_pvariable_focus/updated_restart_phase2_5P_bias_corrected_tropical.nc` + - Use filenames that clearly distinguish: + - natveg-based restarts vs. Phase2-based 5P bias/scale tropical restarts. + +4. **Run and compare:** + - Run the land model from: + - **Run A:** natveg-based 5P bias/scale two-region restarts. + - **Run B:** Phase2-based 5P bias/scale tropical restarts. + - Use your existing analysis scripts (profiles, regional stats, etc.) to compare: + - 5P behavior at the Amazon site and across Amazon/Africa. + - Any downstream variables of interest (e.g., NPP, fluxes). + +This gives you a clean, repeatable way to reproduce the `phase2_pvariable_focus_bias_corrected` +behavior you liked in the Amazon solutionp profile, and to propagate that improvement into +restart-based simulations in the tropics. + + diff --git a/docs/WORKFLOW_PHASE1_PHASE2_PHASE3_RESTARTS.md b/docs/WORKFLOW_PHASE1_PHASE2_PHASE3_RESTARTS.md new file mode 100644 index 0000000..81b46e6 --- /dev/null +++ b/docs/WORKFLOW_PHASE1_PHASE2_PHASE3_RESTARTS.md @@ -0,0 +1,312 @@ +### 5. Site‑based 5P restart comparison (Amazon / Africa or any site) + +Once `phase1_global`, `phase2_tropical`, and `phase3_tworegions` restarts are created, +you can **automatically generate single‑point restarts and 5P comparison plots** for +any longitude/latitude using: + +- `scripts/generate_site_5p_restart_comparison.py` +- Documented in `docs/SITE_5P_RESTART_COMPARISON.md` + +Example (Africa site at lon=28, lat=0): + +```bash +python scripts/generate_site_5p_restart_comparison.py \ + --lon 28.0 \ + --lat 0.0 \ + --site-name africa_28_0 +``` + +Example (Amazon validation site): + +```bash +python scripts/generate_site_5p_restart_comparison.py \ + --lon 303.75 \ + --lat -17.434553 \ + --site-name amazon_303_17S +``` + +The script will: + +- Extract single‑point restarts from the three phase runs at the requested site. +- Load Phase1 **ground‑truth 5P vertical profiles** at that site. +- Produce line plots comparing GT vs `phase1_global`, `phase2_tropical`, + and `phase3_tworegions` for the 5 P variables. + +## Workflow: phase1_global, phase2_tropical, phase3_tworegions restarts + +This document summarizes how to create three restart files from existing +`natveg_improved` and `phase2_pvariable_focus` runs: + +- **phase1_global**: copy of the `natveg_improved` restart (baseline). +- **phase2_tropical**: tropical restart with raw Phase2 5P in \([-30°, 30°]\). +- **phase3_tworegions**: tropical restart with Phase2 5P plus Amazon + Africa bias/scale correction. + +It is based on: + +- `docs/WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md` +- `docs/INSTRUCTIONS_TRENDY_1_AI_RESTART_CREATION.md` + +The recommended way to run all three phases is to use +`scripts/run_phase_restarts.sh`, described at the end of this document. + +--- + +### 0. Common setup and model training for new users + +For a **new user**, the full TRENDY‑1 AI restart workflow should start by +training **two CNP models**: + +- A **global natveg_improved‑like model** (used for `phase1_global` and as the + base restart). +- A **tropical Phase2 P‑focused model** (used for `phase2_tropical` and + `phase3_tworegions`). + +These two trainings are done with `train_cnp_model.py`, and then the rest of +the restart creation steps follow the procedure in +`docs/INSTRUCTIONS_TRENDY_1_AI_RESTART_CREATION.md` and in this document. + +Assumptions below use the existing runs; if you train your own, replace them +with your paths: + +- **Natveg run**: `cnp_results/run_20260228_214757_natveg_improved` +- **Base restart**: + - `cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc` +- **Phase2 run**: `cnp_results/run_20260305_153217_phase2_pvariable_focus` +- **Variable list file**: `CNP_IO_updated9_dev_dw.txt` in repo root. + +If you are a new user and want to **train your own pair of models** instead of +reusing these runs, follow these steps first. + +#### 0.1 Train a natveg_improved‑like global model (for phase1_global) + +From the repo root: + +```bash +cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC + +python train_cnp_model.py \ + --config config/training_config_experiment_3_global_natveg_improved.json \ + --run-dir cnp_results/run_YYYYMMDD_HHMMSS_natveg_improved_custom \ + --variable-list CNP_IO_updated9_dev_dw.txt +``` + +After training: + +1. Follow `docs/WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md` (or your + existing ELM workflow) to produce a **global restart NetCDF** from this run + (full‑domain inference → NetCDF → `ai_predictions_to_restart.py` as needed). +2. Use the resulting restart as your new **base restart** and set: + + ```bash + export NATVEG_RUN_DIR="cnp_results/run_YYYYMMDD_HHMMSS_natveg_improved_custom" + export BASE_RESTART="$NATVEG_RUN_DIR/updated_restart_...your_file.nc" + ``` + +#### 0.2 Train a phase2_tropical‑like P‑focused tropical model + +From the repo root: + +```bash +cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC + +python train_cnp_model.py \ + --config config/training_config_phase2_tropical_soilp_only.json \ + --run-dir cnp_results/run_YYYYMMDD_HHMMSS_phase2_pvariable_focus_custom \ + --variable-list CNP_IO_updated9_dev_dw.txt +``` + +Then set: + +```bash +export PHASE2_RUN_DIR="cnp_results/run_YYYYMMDD_HHMMSS_phase2_pvariable_focus_custom" +``` + +The remaining inference, bias/scale, and restart steps in this document are the +same; only `NATVEG_RUN_DIR`, `BASE_RESTART`, and `PHASE2_RUN_DIR` change to +point at your newly trained models. + +--- + +Create three run directories (example naming): + +```bash +cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC + +TS=$(date +%Y%m%d_%H%M%S) + +RUN1_DIR="cnp_results/run_${TS}_phase1_global" +RUN2_DIR="cnp_results/run_${TS}_phase2_tropical" +RUN3_DIR="cnp_results/run_${TS}_phase3_tworegions" + +mkdir -p "$RUN1_DIR" "$RUN2_DIR" "$RUN3_DIR" +``` + +--- + +### 1. phase1_global: global baseline restart (natveg_improved copy) + +**Goal:** create a clearly labeled baseline restart identical to `natveg_improved`. + +```bash +cp cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + "$RUN1_DIR/updated_restart_phase1_global_natveg_improved.nc" +``` + +Optional README: + +```bash +echo "phase1_global: copy of natveg_improved restart on $(date)" \ + > "$RUN1_DIR/README_phase1_global.txt" +``` + +--- + +### 2. phase2_tropical: tropical restart from raw Phase2 predictions + +**Goal:** overwrite all 5P variables in the tropical band \([-30°, 30°]\) +using **raw Phase2 predictions**. + +#### 2.1 Tropical-only Phase2 inference (if not already done) + +```bash +cd cnp_results/run_20260305_153217_phase2_pvariable_focus + +python ../../scripts/run_inference_all.py \ + --model cnp_predictions/model.pth \ + --output-dir cnp_inference_tropical_only + +cd - # back to repo root +``` + +#### 2.2 Build NetCDF from tropical predictions + +```bash +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_inference_tropical_only/cnp_predictions \ + --output cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_tropical_only.nc +``` + +#### 2.3 Create phase2_tropical restart (5P in tropics, raw Phase2) + +```bash +python scripts/ai_predictions_to_restart.py \ + --ai-predictions cnp_results/run_20260305_153217_phase2_pvariable_focus/comparison_results/ai_predictions_tropical_only.nc \ + --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --output "$RUN2_DIR/updated_restart_phase2_tropical_5P_raw.nc" \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --variables-to-update labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr \ + "--tropical-lat-range=-30,30" +``` + +Optional README: + +```bash +echo "phase2_tropical: natveg base + raw Phase2 5P in tropics on $(date)" \ + > "$RUN2_DIR/README_phase2_tropical.txt" +``` + +--- + +### 3. phase3_tworegions: tropical restart with two-region 5P bias/scale + +**Goal:** use Phase2 predictions plus **Amazon + Africa 5P bias/scale correction** +and map those corrected 5P values into the tropical band of the natveg restart. + +This follows **Path B** from +`WORKFLOW_5P_TWO_REGIONS_BIAS_SCALE_AND_RESTARTS.md`, but writes outputs into +`RUN3_DIR`. + +#### 3.1 Full-grid Phase2 inference into RUN3_DIR + +```bash +python scripts/run_inference_all.py \ + --model cnp_results/run_20260305_153217_phase2_pvariable_focus/cnp_predictions/model.pth \ + --output-dir "$RUN3_DIR/cnp_inference_entire_dataset" \ + --inference-full-grid +``` + +#### 3.2 Apply 5P bias/scale correction (Amazon + Africa only) + +```bash +python scripts/apply_5p_bias_scale_correction.py \ + --run-dir "$RUN3_DIR" \ + --region-config-json config/training_config_two_region_five_p.json \ + --output-subdir soil_2d_predictions_5P_bias_corrected_phase2 +``` + +Outputs: + +- Corrected 5P CSVs under + `"$RUN3_DIR/cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_5P_bias_corrected_phase2/"` +- Bias/scale parameters under + `"$RUN3_DIR/analysis/bias_scale_params_5P_two_regions.json"` + +#### 3.3 Build global NetCDF from bias-corrected 5P predictions + +```bash +mkdir -p "$RUN3_DIR/comparison_results" + +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions "$RUN3_DIR/cnp_inference_entire_dataset/cnp_predictions/" \ + --soil-2d-bias-corrected-subdir soil_2d_predictions_5P_bias_corrected_phase2 \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --output "$RUN3_DIR/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc" +``` + +#### 3.4 Create phase3_tworegions restart (bias-corrected 5P in tropics) + +```bash +python scripts/ai_predictions_to_restart.py \ + --ai-predictions "$RUN3_DIR/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc" \ + --restart-file cnp_results/run_20260228_214757_natveg_improved/updated_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc \ + --output "$RUN3_DIR/updated_restart_phase3_tworegions_5P_bias_corrected_tropical.nc" \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --variables-to-update labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr \ + "--tropical-lat-range=-30,30" +``` + +Optional README: + +```bash +echo "phase3_tworegions: natveg base + Phase2 5P with Amazon+Africa bias/scale in tropics on $(date)" \ + > "$RUN3_DIR/README_phase3_tworegions.txt" +``` + +--- + +### 4. Automation scripts (run 1 → inspect → 2 → inspect → 3) + +## make sure to export RESTART_TEMPLATE + +# export RESTART_TEMPLATE=/mnt/proj-shared/AI4BGC_7xw/AI4BGC/ELM_data/20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc + + +For new users, the recommended workflow is to run three scripts **sequentially**, +inspecting results between phases: + +1. `scripts/run_phase1_global.sh` + - Trains the global natveg_improved‑like model (or reuses an existing one). + - Runs full‑grid inference and creates a **base restart** from an ELM + restart template. + - Exports `NATVEG_RUN_DIR` and `BASE_RESTART` in its log output. + +2. `scripts/run_phase2_tropical.sh` + - Trains the tropical Phase2 P‑focused model (or reuses an existing one). + - Runs tropical‑only inference and creates the **phase2_tropical** restart + with raw Phase2 5P in the tropics, using `BASE_RESTART` from phase 1. + - Exports `PHASE2_RUN_DIR` and prints the path to the tropical restart. + +3. `scripts/run_phase3_tworegions.sh` + - Runs full‑grid inference with the Phase2 model. + - Applies two‑region 5P bias/scale correction (Amazon + Africa). + - Creates the **phase3_tworegions** tropical restart with bias‑corrected + Phase2 5P, using the same `BASE_RESTART`. + +Between scripts you can: + +- Inspect logs, validation plots, and diagnostics for each model/run. +- Adjust configs or weights and re‑run the specific phase if needed. + +Each script is self‑contained and safe to re‑run; they will reuse existing runs +when possible instead of retraining from scratch. + diff --git a/docs/bias_scale_params_5P_two_regions_natveg_improved.json b/docs/bias_scale_params_5P_two_regions_natveg_improved.json new file mode 100644 index 0000000..02d6b6a --- /dev/null +++ b/docs/bias_scale_params_5P_two_regions_natveg_improved.json @@ -0,0 +1,432 @@ +{ + "labilep_vr": { + "amazon": { + "layer_1": { + "a": 1.0333817099721738, + "b": 4.420208988671429 + }, + "layer_2": { + "a": 1.0409236696102435, + "b": 0.3818685505560646 + }, + "layer_3": { + "a": 1.0495804203466117, + "b": 1.06616903652683 + }, + "layer_4": { + "a": 1.0579894058014472, + "b": 0.4946148860290543 + }, + "layer_5": { + "a": 1.0561419211059198, + "b": 0.4841642203024569 + }, + "layer_6": { + "a": 1.0577645417493151, + "b": 0.27375648613853115 + }, + "layer_7": { + "a": 1.0847835392641536, + "b": -0.22465004246161327 + }, + "layer_8": { + "a": 1.0832692207296306, + "b": -0.7209388262875364 + }, + "layer_9": { + "a": 1.0936269800871787, + "b": -0.9484615519453438 + }, + "layer_10": { + "a": 1.0852079008341873, + "b": -0.8735917626389086 + } + }, + "africa": { + "layer_1": { + "a": 0.9260468289051457, + "b": 10.03484784525775 + }, + "layer_2": { + "a": 0.956908088141792, + "b": 6.503605072308462 + }, + "layer_3": { + "a": 1.0039209238414928, + "b": 1.296756250945013 + }, + "layer_4": { + "a": 1.0370254105759864, + "b": -0.5462219369286958 + }, + "layer_5": { + "a": 1.0353504495106733, + "b": -0.45350925411607107 + }, + "layer_6": { + "a": 1.0133549405298117, + "b": -0.04385045363058444 + }, + "layer_7": { + "a": 1.0215880516118319, + "b": -0.21450387181591266 + }, + "layer_8": { + "a": 1.0065953393691076, + "b": -0.4975742446923468 + }, + "layer_9": { + "a": 1.013726018934732, + "b": -0.4609999519237106 + }, + "layer_10": { + "a": 1.0112312802276142, + "b": -06295343761584757 + } + } + }, + "occlp_vr": { + "amazon": { + "layer_1": { + "a": 0.970872434819557, + "b": 24.527277601293697 + }, + "layer_2": { + "a": 0.981422686441536, + "b": 22.18716926343895 + }, + "layer_3": { + "a": 0.975978650546871, + "b": 21.292876582030107 + }, + "layer_4": { + "a": 0.9790900134809112, + "b": 21.175889847963813 + }, + "layer_5": { + "a": 0.9782904017044313, + "b": 20.57937164466406 + }, + "layer_6": { + "a": 0.9850750586394957, + "b": 17.70209775773947 + }, + "layer_7": { + "a": 0.993293552606405, + "b": 14.237383178818485 + }, + "layer_8": { + "a": 0.9794630007993673, + "b": 20.8794099079381 + }, + "layer_9": { + "a": 0.9833147871353544, + "b": 19.264509797701816 + }, + "layer_10": { + "a": 0.9741497614725136, + "b": 19.81352328584456 + } + }, + "africa": { + "layer_1": { + "a": 0.971858212747255, + "b": 33.6201710855249 + }, + "layer_2": { + "a": 0.9827890994310069, + "b": 31.2788877933423 + }, + "layer_3": { + "a": 0.9832206980465467, + "b": 30.869427223862168 + }, + "layer_4": { + "a": 0.9878322934304136, + "b": 30.445729301931724 + }, + "layer_5": { + "a": 0.9874108706009984, + "b": 30.068595825491425 + }, + "layer_6": { + "a": 1.001074807191854, + "b": 27.70530975050953 + }, + "layer_7": { + "a": 1.0036568842185372, + "b": 26.81471469640506 + }, + "layer_8": { + "a": 0.9866185653706865, + "b": 28.856125345878777 + }, + "layer_9": { + "a": 0.9867051781793659, + "b": 30.001784710169126 + }, + "layer_10": { + "a": 0.9856365883368785, + "b": 29.438470497893853 + } + } + }, + "solutionp_vr": { + "amazon": { + "layer_1": { + "a": 1.20887721115336, + "b": -0.013112905188638556 + }, + "layer_2": { + "a": 1.2504076620794857, + "b": -0.013883410165339857 + }, + "layer_3": { + "a": 1.1734006950147933, + "b": -0.0011779314386677498 + }, + "layer_4": { + "a": 1.0978892007079026, + "b": 0.0003994104511832752 + }, + "layer_5": { + "a": 1.0899789723654238, + "b": 0.0003410953328287401 + }, + "layer_6": { + "a": 1.0888958046678634, + "b": 0.001114385047667878 + }, + "layer_7": { + "a": 1.1092931592310282, + "b": 0.0006840241335554904 + }, + "layer_8": { + "a": 1.1048788490549455, + "b": 0.00017212671348821045 + }, + "layer_9": { + "a": 1.120575322942864, + "b": 0.00028255324549791004 + }, + "layer_10": { + "a": 1.1213771769060543, + "b": 0.0003076512472597163 + } + }, + "africa": { + "layer_1": { + "a": 1.0819222786968083, + "b": -0.0029634948220650125 + }, + "layer_2": { + "a": 1.1702954875010905, + "b": -0.004559488540869101 + }, + "layer_3": { + "a": 1.1303000957816693, + "b": -0.001256810728127183 + }, + "layer_4": { + "a": 1.1141119271683604, + "b": -0.0037482385394902915 + }, + "layer_5": { + "a": 1.107322949311081, + "b": -0.0026561649766157908 + }, + "layer_6": { + "a": 1.0780388932747962, + "b": -0.0015113227110827659 + }, + "layer_7": { + "a": 1.0656317881531925, + "b": -0.0012467122139105521 + }, + "layer_8": { + "a": 1.075414734185535, + "b": -0.0013494779813059863 + }, + "layer_9": { + "a": 1.0611235681409525, + "b": -0.001355840730149602 + }, + "layer_10": { + "a": 1.0820744618897282, + "b": -0.0013645175299187302 + } + } + }, + "secondp_vr": { + "amazon": { + "layer_1": { + "a": 1.0448985460110511, + "b": 36.611060894533246 + }, + "layer_2": { + "a": 1.0410709953432102, + "b": -7.659592633360574 + }, + "layer_3": { + "a": 1.0517394690955428, + "b": 13.129326473149215 + }, + "layer_4": { + "a": 1.053196865163657, + "b": 0.14192104401904268 + }, + "layer_5": { + "a": 1.0639034433320247, + "b": 6.522188144615301 + }, + "layer_6": { + "a": 1.064584417591993, + "b": -0.4067862179203687 + }, + "layer_7": { + "a": 1.0798291117090646, + "b": -5.7688592570011545 + }, + "layer_8": { + "a": 1.0875634804878878, + "b": -14.582715023551932 + }, + "layer_9": { + "a": 1.0878262871976658, + "b": -16.640924160509947 + }, + "layer_10": { + "a": 1.0887957033372113, + "b": -18.462864508270986 + } + }, + "africa": { + "layer_1": { + "a": 0.9205343953951598, + "b": 192.16355982398082 + }, + "layer_2": { + "a": 0.9534318453202065, + "b": 133.0882320261201 + }, + "layer_3": { + "a": 0.9984796618785772, + "b": 34.13339195380496 + }, + "layer_4": { + "a": 1.0234611547345664, + "b": -1.8899745174998936 + }, + "layer_5": { + "a": 1.0333637131896058, + "b": -7.390266341339224 + }, + "layer_6": { + "a": 1.0198705104896868, + "b": -6.586734093860522 + }, + "layer_7": { + "a": 1.01914631657447, + "b": -7.223674790840671 + }, + "layer_8": { + "a": 1.0146049337029506, + "b": -13.203945742135598 + }, + "layer_9": { + "a": 1.0169358978073586, + "b": -14.970135987068145 + }, + "layer_10": { + "a": 1.0080426644449207, + "b": -12.046482445562189 + } + } + }, + "primp_vr": { + "amazon": { + "layer_1": { + "a": 1.1116537915200317, + "b": 4.748082380269224 + }, + "layer_2": { + "a": 1.070313011774915, + "b": 4.804797098407048 + }, + "layer_3": { + "a": 0.570820883826258, + "b": 5.587446668401681 + }, + "layer_4": { + "a": -0.4374880670118703, + "b": 3.600611708160891 + }, + "layer_5": { + "a": -0.5255801564173186, + "b": 1.5959601390462705 + }, + "layer_6": { + "a": -0.48132782292855897, + "b": 1.0476749373164072 + }, + "layer_7": { + "a": -0.3391653099549358, + "b": 0.8482902863595918 + }, + "layer_8": { + "a": -0.38233820872934066, + "b": 0.8546738433970271 + }, + "layer_9": { + "a": -0.3959357491961905, + "b": 0.8701660999783685 + }, + "layer_10": { + "a": -0.4503675125152784, + "b": 0.9122670099526271 + } + }, + "africa": { + "layer_1": { + "a": 0.8708169600943022, + "b": 41.00216523185628 + }, + "layer_2": { + "a": 0.8806484529678255, + "b": 22.44235332760175 + }, + "layer_3": { + "a": 0.5657141015978582, + "b": 12.186385683929139 + }, + "layer_4": { + "a": 0.08441461968826015, + "b": 5.088583907044187 + }, + "layer_5": { + "a": -0.32073044195360734, + "b": 1.838159292536186 + }, + "layer_6": { + "a": -0.2197071471731467, + "b": 0.8444548867780934 + }, + "layer_7": { + "a": -0.19143372160444425, + "b": 0.6764503619020544 + }, + "layer_8": { + "a": -0.2145576057922407, + "b": 0.6713913272208819 + }, + "layer_9": { + "a": -0.3286436154248806, + "b": 0.7216657330381735 + }, + "layer_10": { + "a": -0.24964867052808587, + "b": 0.7026306321492235 + } + } + } +} diff --git a/docs/npool_ppool_tropical_observations.md b/docs/npool_ppool_tropical_observations.md index 4396024..7bc3d18 100644 --- a/docs/npool_ppool_tropical_observations.md +++ b/docs/npool_ppool_tropical_observations.md @@ -1,5 +1,24 @@ # Tropical NPOOL/PPOOL Observations and Suggestions +## Background values (10 for npool, 1 for ppool) + +**Does this codebase use 10 as background for npool and 1 for ppool?** +- **No.** The training script and data loaders do **not** initialize or fill npool with 10 or ppool with 1. NaN/Inf in PFT1D data are replaced with **0** (`nan=0.0` in `data_loader_individual.py`); there is no special case for npool/ppool. +- **Where do 10 and 1 come from?** They come from the **training targets** (ELM/TRENDY output in the PKL files). For many PFTs, the **ground-truth** Y_npool and Y_ppool in the dataset are near-constant at ~10 and ~1. So the "background" is in the **source data** (likely ELM default or typical values for those pool variables), not from our initialization. + +**Training dataset (ground truth) counts:** Run `scripts/analyze_npool_ppool_special_values.py` on your training PKL path. Example results on Trendy_1_data_CNP (20,975 grid cells, 16 PFTs): +- **npool == 10**: 80.7% of all (cell, PFT) pairs have this value; 32.8% of grid cells have *all* 16 PFTs equal to 10; 100% of cells have *at least one* PFT with npool == 10. +- **ppool == 1**: Same statistics (80.7% of pairs, 32.8% of cells all-1, 100% any). +- **Both**: 6,877 grid cells (32.8%) have all PFTs with npool==10 and ppool==1. +- Per-PFT counts and full report: `docs/npool_ppool_special_values_report.json`. + +**Validation excluding special-value grid cells:** Run `scripts/validation_npool_ppool_exclude_special.py ` to compute R² and RMSE for npool and ppool when excluding grid cells where *all* 16 PFTs have npool==10 and ppool==1. Example (run_20260226_114546_nofilter, test set 4,166 cells, 33.5% special): +- **NPOOL**: R² (all cells) ≈ 0.02 → R² (excl. special) ≈ **0.41**; RMSE 8.83 → 8.37. +- **PPOOL**: R² (all cells) ≈ -1.58 → R² (excl. special) ≈ **-0.51**; RMSE 0.87 → 0.82. +So on non-constant cells, npool prediction is moderate and ppool is still poor but much less bad than when including the constant cells. + +**Per-PFT validation (recommended):** The special values (10 for npool, 1 for ppool) should be treated **per PFT**. For each PFT k, only grid cells where **PCT_NAT_PFT_k > 0** (that PFT is present) should be used when computing prediction quality for that PFT; other cells are filled with 10/1 and should not count. Run `scripts/validation_npool_ppool_per_pft.py ` to get per-PFT R² and RMSE using this rule. Example (run_20260226_114546_nofilter): mean R² over PFTs (with ≥10 valid cells) is **~0.92** for npool and **~0.93** for ppool — prediction quality is good when evaluated only where each PFT is present. + ## Scope - Dataset: tropical-only subset (Latitude between -23.5 and 23.5). - Source: Trendy_1 training data (21 PKL files). diff --git a/models/cnp_combined_model.py b/models/cnp_combined_model.py index 1771e67..994e8ee 100644 --- a/models/cnp_combined_model.py +++ b/models/cnp_combined_model.py @@ -267,7 +267,12 @@ def __init__(self, model_config: ModelConfig, data_info: Dict[str, Any], num_layers=getattr(self.model_config, 'transformer_layers', 6) ) - # 5. Output Heads + # 5. Multihead/Multimode Configuration + self.use_multihead_attention = getattr(self.model_config, 'use_multihead_attention', False) + self.multihead_num_heads = getattr(self.model_config, 'multihead_num_heads', 4) + self.use_mode_specific_heads = getattr(self.model_config, 'use_mode_specific_heads', False) + + # 6. Output Heads self._build_output_heads() # Loss weights setup (保持原样) @@ -407,7 +412,69 @@ def _configure_static_structure(self): return input_group_indices, group_ids def _build_output_heads(self): - # Heads 输入维度改为 self.embed_dim (因为 backbone 输出也是这个维度) + # Multihead Attention layers for mode-specific feature extraction + if self.use_multihead_attention: + # Multi-head attention for each output mode + self.scalar_attention = nn.MultiheadAttention( + embed_dim=self.embed_dim, + num_heads=self.multihead_num_heads, + dropout=self.dropout_p, + batch_first=True + ) + self.soil2d_attention = nn.MultiheadAttention( + embed_dim=self.embed_dim, + num_heads=self.multihead_num_heads, + dropout=self.dropout_p, + batch_first=True + ) + self.pft1d_attention = nn.MultiheadAttention( + embed_dim=self.embed_dim, + num_heads=self.multihead_num_heads, + dropout=self.dropout_p, + batch_first=True + ) + + # Query vectors for each mode (learnable) + self.scalar_query = nn.Parameter(torch.randn(1, 1, self.embed_dim)) + self.soil2d_query = nn.Parameter(torch.randn(1, 1, self.embed_dim)) + self.pft1d_query = nn.Parameter(torch.randn(1, 1, self.embed_dim)) + + # Mode-specific feature projection + self.scalar_feat_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.soil2d_feat_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.pft1d_feat_proj = nn.Linear(self.embed_dim, self.embed_dim) + else: + self.scalar_attention = None + self.soil2d_attention = None + self.pft1d_attention = None + + # Mode-specific feature extractors (optional) + if self.use_mode_specific_heads: + # Separate feature extractors for each mode + self.scalar_feature_extractor = nn.Sequential( + nn.Linear(self.embed_dim, self.embed_dim), + nn.LayerNorm(self.embed_dim), + nn.GELU(), + nn.Dropout(self.dropout_p) + ) + self.soil2d_feature_extractor = nn.Sequential( + nn.Linear(self.embed_dim, self.embed_dim), + nn.LayerNorm(self.embed_dim), + nn.GELU(), + nn.Dropout(self.dropout_p) + ) + self.pft1d_feature_extractor = nn.Sequential( + nn.Linear(self.embed_dim, self.embed_dim), + nn.LayerNorm(self.embed_dim), + nn.GELU(), + nn.Dropout(self.dropout_p) + ) + else: + self.scalar_feature_extractor = None + self.soil2d_feature_extractor = None + self.pft1d_feature_extractor = None + + # Output heads - 输入维度改为 self.embed_dim (因为 backbone 输出也是这个维度) self.scalar_head = nn.Sequential( nn.Linear(self.embed_dim, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(self.dropout_p), @@ -456,6 +523,26 @@ def _initialize_weights(self): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) + + # Initialize multihead attention query vectors + if self.use_multihead_attention: + nn.init.normal_(self.scalar_query, std=0.02) + nn.init.normal_(self.soil2d_query, std=0.02) + nn.init.normal_(self.pft1d_query, std=0.02) + + # Initialize feature projection layers + if hasattr(self, 'scalar_feat_proj'): + nn.init.xavier_uniform_(self.scalar_feat_proj.weight) + if self.scalar_feat_proj.bias is not None: + nn.init.zeros_(self.scalar_feat_proj.bias) + if hasattr(self, 'soil2d_feat_proj'): + nn.init.xavier_uniform_(self.soil2d_feat_proj.weight) + if self.soil2d_feat_proj.bias is not None: + nn.init.zeros_(self.soil2d_feat_proj.bias) + if hasattr(self, 'pft1d_feat_proj'): + nn.init.xavier_uniform_(self.pft1d_feat_proj.weight) + if self.pft1d_feat_proj.bias is not None: + nn.init.zeros_(self.pft1d_feat_proj.bias) def _count_parameters(self): return sum(p.numel() for p in self.parameters() if p.requires_grad) @@ -533,15 +620,47 @@ def forward(self, time_series_data, static_data, pft_param_data, scalar, variabl # 4. Backbone features = self.backbone(all_tokens) - # 5. Global Pooling - global_feat = features.mean(dim=1) + # 5. Multihead/Multimode Feature Extraction + if self.use_multihead_attention: + # Extract mode-specific features using multi-head attention + batch_size = features.shape[0] + + # Scalar mode: attention over all tokens with scalar-specific query + scalar_query = self.scalar_query.expand(batch_size, -1, -1) + scalar_attn_out, _ = self.scalar_attention(scalar_query, features, features) + scalar_feat = scalar_attn_out.squeeze(1) # [B, D] + scalar_feat = self.scalar_feat_proj(scalar_feat) + + # Soil 2D mode: attention over all tokens with soil2d-specific query + soil2d_query = self.soil2d_query.expand(batch_size, -1, -1) + soil2d_attn_out, _ = self.soil2d_attention(soil2d_query, features, features) + soil2d_feat = soil2d_attn_out.squeeze(1) # [B, D] + soil2d_feat = self.soil2d_feat_proj(soil2d_feat) + + # PFT 1D mode: attention over all tokens with pft1d-specific query + pft1d_query = self.pft1d_query.expand(batch_size, -1, -1) + pft1d_attn_out, _ = self.pft1d_attention(pft1d_query, features, features) + pft1d_feat = pft1d_attn_out.squeeze(1) # [B, D] + pft1d_feat = self.pft1d_feat_proj(pft1d_feat) + else: + # Standard global pooling + global_feat = features.mean(dim=1) + scalar_feat = global_feat + soil2d_feat = global_feat + pft1d_feat = global_feat + + # Apply mode-specific feature extractors if enabled + if self.use_mode_specific_heads: + scalar_feat = self.scalar_feature_extractor(scalar_feat) + soil2d_feat = self.soil2d_feature_extractor(soil2d_feat) + pft1d_feat = self.pft1d_feature_extractor(pft1d_feat) # 6. Heads outputs = {} - outputs['scalar'] = torch.relu(self.scalar_head(global_feat)) - outputs['soil_2d'] = torch.nn.functional.softplus(self.matrix_head(global_feat)) + outputs['scalar'] = torch.relu(self.scalar_head(scalar_feat)) + outputs['soil_2d'] = torch.nn.functional.softplus(self.matrix_head(soil2d_feat)) - pft_out = self.pft_1d_head(global_feat) + pft_out = self.pft_1d_head(pft1d_feat) # Process PFT output pft_1d_varnames = self.data_info.get('variables_1d_pft', []) n_vars = len(pft_1d_varnames) diff --git a/scripts/ai_model_comparison_plot.py b/scripts/ai_model_comparison_plot.py index ae41847..297635a 100644 --- a/scripts/ai_model_comparison_plot.py +++ b/scripts/ai_model_comparison_plot.py @@ -1,4 +1,56 @@ #!/usr/bin/env python3 +""" +Compare AI CNP predictions with reference (model) or with CSV exports. +Also supports model-vs-model: compare two ELM restart NetCDFs (--model1, --model2). + +Design: + - Model vs model: two ELM restart NetCDFs. Uses same grid from model1; compares + variables from --variable-list or --variables. PFTs 1-16 (--pfts, 1-based) and top 10 + soil layers (args.layers) only. + - Primary: AI predictions NetCDF vs model (ELM) NetCDF. + Use --ai-predictions and --model (and --compare-to model to avoid CSV). + + - Optional: AI predictions NetCDF vs CSV predictions (same run, two formats). + Use --csv-predictions for a consistency check. Use --compare-to csv or auto. + +How AI vs model comparison works (different internal formats): + + The AI NetCDF and the model restart NetCDF do not have the same internal + layout. The script bridges them as follows. + + Required for both: + - grid1d_lon, grid1d_lat (1D arrays of length n_gridcell) + - Same variable names (e.g. cwdc_vr, tlai) for the fields to compare + + AI NetCDF (e.g. from ai_predictions_to_netcdf.py): + - One gridcell per land point; dimension "gridcell". + - Scalars: (gridcell). + - PFT variables: (pft, gridcell) with pft size 16 (PFT 1–16). + - Soil/column variables: (column, levgrnd, gridcell) with column size 1, + levgrnd size 10 (one column per gridcell). + + Model (ELM) NetCDF: + - Many columns and PFT instances; each is linked to a gridcell via + cols1d_gridcell_index and pfts1d_gridcell_index (1-based or 0-based). + - Scalars: often (gridcell). + - Column variables: (column, levgrnd) — column index maps to gridcell. + - PFT variables: (pft) — model has PFT 0–16 (17 slots). For comparison we + use only model PFT 1–16 and ignore model PFT 0, so user --pfts 1–16 align + with both AI (1–16) and model (1–16). + + The script: + 1. Builds a spatial mapping: for each AI gridcell (lon, lat), finds the + closest model gridcell (nearest-neighbor). + 2. For each model gridcell g: gets AI value from the AI gridcell that + maps to g; gets model value from the model’s column/pft indices for g + (e.g. first column and PFT 1–16 for that gridcell). + 3. Plots and stats are on the model grid (one value per model gridcell). + + So the AI NetCDF must be in the “plotting” format produced by + ai_predictions_to_netcdf.py (grid1d_lon/lat, gridcell, pft/gridcell, + column/levgrnd/gridcell). If your AI NetCDF has a different layout, convert + it first or extend this script to support that layout. +""" import os import numpy as np import xarray as xr @@ -21,19 +73,19 @@ # Default paths FALLBACK_AI_PREDICTIONS = './comparison_results/ai_predictions_for_plotting.nc' -FALLBACK_MODEL = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' +#FALLBACK_MODEL = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' +FALLBACK_MODEL = '/mnt/proj-shared/AI4BGC_7xw/AI4BGC/ELM_data/20251201_TRENDY2024_default_ICB1850CNPRDCTCBC.elm.r.0801-01-01-00000.nc' FALLBACK_OUTPUT_DIR = "./ai_model_comparison_plots" FALLBACK_CSV_PREDICTIONS = './cnp_inference_entire_dataset/cnp_predictions' # Default variables to plot unless '--variables all' is used -VARIABLES = ['cwdc_vr', 'soil3c_vr', 'tlai', 'deadstemc'] +VARIABLES = ['cwdc_vr', 'soil3c_vr', 'tlai', 'deadstemc', 'cpool', 'npool', 'ppool', 'primp_vr', 'secondp_vr', 'litr2c_vr', 'litr2n_vr', 'litr2p_vr', 'soil1c_vr', 'soil1n_vr', 'soil1p_vr'] # Layers for column-type variables (AI has 10 layers, model has 15) -LEVGRND_LAYERS = [0, 4, 9] # Layers 0, 4, 9 (corresponding to AI layers 1, 5, 10) +LEVGRND_LAYERS = [0,1,2,3,4,5,6,7,8,9] # Layers 0, 4, 9 (corresponding to AI layers 1, 5, 10) -# PFTs to plot (AI has PFT1-16, model has PFT0-16) -# Note: AI PFT0 = Model PFT1, AI PFT1 = Model PFT2, etc. -PFT_PICK_LIST = [0, 1, 2, 3, 4] # PFT1, PFT2, PFT3, PFT4, PFT5 (0-indexed, so 0=PFT1, 1=PFT2, etc.) +# PFTs to plot: 1-based (1–16). AI has PFT 1–16; model has PFT 0–16; we compare 1–16 and ignore model PFT 0. +PFT_PICK_LIST = list(range(1, 17)) # [1, 2, ..., 16] — plot all 16 PFTs by default CSV_LONGITUDE_NAMES = ("Longitude", "Long", "long", "lon", "LON") CSV_LATITUDE_NAMES = ("Latitude", "Lat", "lat", "LAT") @@ -393,6 +445,59 @@ def _build_gridcell_groups(one_d_to_grid, n_grid): groups[g].append(idx) return groups + +def _extract_elm_var_on_grid(ds, var, grid_to_cols, grid_to_pfts, n_grid, layer_idx=None, pft_idx=None): + """Extract an ELM restart variable as a 1D array per gridcell (length n_grid). + + Used for model-vs-model comparison. Call with exactly one of layer_idx (for column/soil2d) + or pft_idx (for pft1d) or neither (for gridcell scalar). + """ + if var not in ds.data_vars: + return None + da = ds[var] + dims = da.dims + vals = _to_nan_fillvalue(np.asarray(da.values)) + out = np.full(n_grid, np.nan, dtype=float) + + if layer_idx is not None and 'column' in dims and 'levgrnd' in dims: + # Column/soil 2D: ensure (column, levgrnd) order, use first column per gridcell, given layer + if vals.ndim == 2: + if dims[0] == 'levgrnd': + vals = vals.T + n_col, n_lev = vals.shape[0], vals.shape[1] + if layer_idx >= n_lev: + return out + for g in range(n_grid): + cols = grid_to_cols[g] + if cols: + ci = cols[0] + if ci < n_col: + out[g] = vals[ci, layer_idx] + return out + + if pft_idx is not None and 'pft' in dims: + # PFT 1D: AI has PFT 1-16, model has PFT 0-16. We compare AI 1-16 to model 1-16 (ignore model PFT 0). + # So for pft_idx 0..15 (user PFT 1..16) use model index at grid_to_pfts[g][pft_idx + 1]. + if vals.ndim == 1: + for g in range(n_grid): + pfts = grid_to_pfts[g] + model_slot = pft_idx + 1 # model PFT 1-16 (skip slot 0) + if model_slot < len(pfts): + pi = pfts[model_slot] + if pi < vals.shape[0]: + out[g] = vals[pi] + return out + + if 'gridcell' in dims: + # Gridcell scalar + v = vals.reshape(-1) if vals.ndim > 1 else vals + if v.size >= n_grid: + out[:] = v[:n_grid] + return out + + return None + + def _plot_map(ax, lon, lat, data, title, vmin=None, vmax=None, cmap="viridis", norm=None): """Plot a map with the given data.""" ax.add_feature(cfeature.COASTLINE) @@ -489,7 +594,7 @@ def _nan_max(a): print(f"Stats for {var}{label_suffix}:") print(f" {label_ai}: sum={sum_ai:.6g} std={std_ai:.6g} min={min_ai:.6g} max={max_ai:.6g}") print(f" {label_model}: sum={sum_model:.6g} std={std_model:.6g} min={min_model:.6g} max={max_model:.6g}") - print(f" Metrics (AI vs Model): n={n} rmse={rmse:.6g} nrmse={nrmse:.6g} r2={r2:.6g}") + print(f" Metrics ({label_ai} vs {label_model}): n={n} rmse={rmse:.6g} nrmse={nrmse:.6g} r2={r2:.6g}") stats = { "ai_sum": sum_ai, @@ -520,20 +625,21 @@ def _nan_max(a): ax2 = fig.add_subplot(gs[1, 0], projection=ccrs.PlateCarree()) _plot_map(ax2, lon, lat, data_model, f"{var} - {label_model}", vmin=vmin_orig, vmax=vmax_orig, cmap="viridis") - # Panel 3: Difference (AI - Model) + # Panel 3: Difference (left - right) + diff_title = f"{var} - Difference ({label_ai} - {label_model})" ax3 = fig.add_subplot(gs[2, 0], projection=ccrs.PlateCarree()) if not np.isfinite(diff_abs) or diff_abs <= 0: msg = "No Difference" if diff_abs == 0 else "All NaN" ax3.text(0.5, 0.5, msg, ha="center", va="center", transform=ax3.transAxes, fontsize=14, fontweight="bold", color="gray") - ax3.set_title(f"{var} - Difference (AI - Model)", fontsize=14, fontweight="bold") + ax3.set_title(diff_title, fontsize=14, fontweight="bold") ax3.set_global() gl = ax3.gridlines(draw_labels=True, alpha=0.5, linestyle="--") gl.top_labels = False gl.right_labels = False else: norm = TwoSlopeNorm(vmin=-diff_abs, vcenter=0, vmax=diff_abs) - _plot_map(ax3, lon, lat, diff, f"{var} - Difference (AI - Model)", cmap="RdBu_r", norm=norm) + _plot_map(ax3, lon, lat, diff, diff_title, cmap="RdBu_r", norm=norm) # Panel 4: Percent Difference Categories ax4 = fig.add_subplot(gs[3, 0], projection=ccrs.PlateCarree()) @@ -559,7 +665,7 @@ def _nan_max(a): ax4.add_feature(cfeature.OCEAN, color="lightblue", alpha=0.5) ax4.add_feature(cfeature.LAND, color="lightgray", alpha=0.3) im4 = ax4.scatter(lon, lat, c=cat, s=5, cmap=cmap, norm=norm_cat, transform=ccrs.PlateCarree()) - ax4.set_title(f"{var} - Percent Diff bins ((AI-Model)/Model)", fontsize=14, fontweight="bold") + ax4.set_title(f"{var} - Percent Diff bins (({label_ai}-{label_model})/{label_model})", fontsize=14, fontweight="bold") ax4.set_global() gl4 = ax4.gridlines(draw_labels=True, alpha=0.5, linestyle="--") gl4.top_labels = False @@ -667,6 +773,10 @@ def main(): # Plot specific variables only python ai_model_comparison_plot.py --variables cwdc_vr tlai + + # Model vs model (two ELM restart NetCDFs): PFT 1-16, top 10 soil layers + python ai_model_comparison_plot.py --model1 ./run1.elm.r.0021-01-01-00000.nc \\ + --model2 ./run2.elm.r.0801-01-01-00000.nc --variable-list CNP_IO.txt --output-dir ./model_vs_model_plots """ ) @@ -674,6 +784,10 @@ def main(): help=f'Path to AI predictions NetCDF file [default: AI_PREDICTIONS_DEFAULT in variable list or {FALLBACK_AI_PREDICTIONS}]') parser.add_argument('--model', default=None, help=f'Path to model results NetCDF file [default: MODEL_DEFAULT in variable list or {FALLBACK_MODEL}]') + parser.add_argument('--model1', default=None, + help='Path to first ELM restart NetCDF (model-vs-model mode). Use with --model2.') + parser.add_argument('--model2', default=None, + help='Path to second ELM restart NetCDF (model-vs-model mode). Use with --model1.') parser.add_argument('--output-dir', default=None, help=f'Output directory for plots [default: COMPARISON_OUTPUT_DIR in variable list or {FALLBACK_OUTPUT_DIR}]') parser.add_argument('--variable-list', type=str, @@ -683,7 +797,7 @@ def main(): parser.add_argument('--layers', nargs='*', type=int, default=LEVGRND_LAYERS, help=f'Layers to plot for column variables [default: {LEVGRND_LAYERS}]') parser.add_argument('--pfts', nargs='*', type=int, default=PFT_PICK_LIST, - help=f'PFTs to plot [default: {PFT_PICK_LIST}]') + help='PFTs to plot (1-based 1..16). AI and model comparison use PFT 1-16; model PFT 0 is ignored. E.g. --pfts 1 2 3 -> _pft1, _pft2, _pft3 [default: 1..16]') parser.add_argument('--no-plot', action='store_true', help='Disable plot generation; compute and save statistics only') parser.add_argument('--stats-file', type=str, @@ -697,6 +811,7 @@ def main(): args = parser.parse_args() + model_vs_model = bool(args.model1 and args.model2) defaults_from_config = _load_default_paths(args.variable_list) def _resolve_default(current_value, config_key, fallback): @@ -705,71 +820,114 @@ def _resolve_default(current_value, config_key, fallback): return str(candidate) return fallback - args.ai_predictions = _resolve_default(args.ai_predictions, 'ai_predictions_default', FALLBACK_AI_PREDICTIONS) - args.model = _resolve_default(args.model, 'model_default', FALLBACK_MODEL) - args.output_dir = _resolve_default(args.output_dir, 'comparison_output_dir', FALLBACK_OUTPUT_DIR) - args.csv_predictions = _resolve_default(args.csv_predictions, 'csv_predictions_default', FALLBACK_CSV_PREDICTIONS) - - use_csv_predictions = bool(args.csv_predictions) - if args.stats_only: - args.no_plot = True - - # Validate input files - if not Path(args.ai_predictions).exists(): - raise FileNotFoundError(f"AI predictions file not found: {args.ai_predictions}") - if use_csv_predictions: - if not Path(args.csv_predictions).exists(): - raise FileNotFoundError(f"CSV predictions source not found: {args.csv_predictions}") + if model_vs_model: + use_csv_predictions = False + if args.stats_only: + args.no_plot = True + args.output_dir = _resolve_default(args.output_dir, 'comparison_output_dir', FALLBACK_OUTPUT_DIR) + if not Path(args.model1).exists(): + raise FileNotFoundError(f"Model 1 file not found: {args.model1}") + if not Path(args.model2).exists(): + raise FileNotFoundError(f"Model 2 file not found: {args.model2}") else: - if not Path(args.model).exists(): - raise FileNotFoundError(f"Model file not found: {args.model}") + args.ai_predictions = _resolve_default(args.ai_predictions, 'ai_predictions_default', FALLBACK_AI_PREDICTIONS) + user_provided_model = (args.model is not None) + args.model = _resolve_default(args.model, 'model_default', FALLBACK_MODEL) + args.output_dir = _resolve_default(args.output_dir, 'comparison_output_dir', FALLBACK_OUTPUT_DIR) + args.csv_predictions = _resolve_default(args.csv_predictions, 'csv_predictions_default', FALLBACK_CSV_PREDICTIONS) + + # If user passed --model (or --restart-file), compare to that file; otherwise use CSV if path exists + use_csv_predictions = bool(args.csv_predictions) and not user_provided_model + if getattr(args, 'restart_file', None): + use_csv_predictions = False + if args.stats_only: + args.no_plot = True + + # Validate input files + if not Path(args.ai_predictions).exists(): + raise FileNotFoundError(f"AI predictions file not found: {args.ai_predictions}") + if use_csv_predictions: + if not Path(args.csv_predictions).exists(): + raise FileNotFoundError(f"CSV predictions source not found: {args.csv_predictions}") + else: + if not Path(args.model).exists(): + raise FileNotFoundError(f"Model file not found: {args.model}") print("="*60) - if use_csv_predictions: + if model_vs_model: + print("Model 1 vs Model 2 Comparison (ELM restart NetCDFs)") + elif use_csv_predictions: print("CSV vs NetCDF Comparison") else: print("AI vs Model Comparison") print("="*60) - print(f"AI predictions (NetCDF): {args.ai_predictions}") - if use_csv_predictions: - print(f"CSV predictions: {args.csv_predictions}") + if model_vs_model: + print(f"Model 1: {args.model1}") + print(f"Model 2: {args.model2}") else: - print(f"Model results: {args.model}") + print(f"AI predictions (NetCDF): {args.ai_predictions}") + if use_csv_predictions: + print(f"CSV predictions: {args.csv_predictions}") + else: + print(f"Model results: {args.model}") print(f"Output directory: {args.output_dir}") print(f"Variables to plot: {args.variables}") print(f"Layers to plot: {args.layers}") print(f"PFTs to plot: {args.pfts}") print("="*60) - + # Open datasets once - ds_ai = xr.open_dataset(args.ai_predictions) + ds_ai = None + ds_model = None + ds_model1 = None + ds_model2 = None csv_predictions = None - if use_csv_predictions: - csv_predictions = load_csv_predictions(args.csv_predictions) - ds_model = None + if model_vs_model: + ds_model1 = xr.open_dataset(args.model1) + ds_model2 = xr.open_dataset(args.model2) else: - ds_model = xr.open_dataset(args.model) + ds_ai = xr.open_dataset(args.ai_predictions) + if use_csv_predictions: + csv_predictions = load_csv_predictions(args.csv_predictions) + else: + ds_model = xr.open_dataset(args.model) # Determine variable selection behavior + # Use variable list when: --variables all, or --variable-list given without --variables (use all vars from that file), or --stats-only + --variable-list requested_all = False if args.variables: requested_all = (len(args.variables) == 1 and str(args.variables[0]).lower() == 'all') if args.stats_only and args.variable_list: requested_all = True + if args.variable_list and '--variables' not in sys.argv: + requested_all = True variable_category_map = build_variable_category_map(args.variable_list) if args.variable_list else {} - ai_vars = set(ds_ai.data_vars.keys()) + if model_vs_model: + common_vars = set(ds_model1.data_vars.keys()) & set(ds_model2.data_vars.keys()) + else: + ai_vars = set(ds_ai.data_vars.keys()) + + def _close_all(): + if ds_ai is not None: + ds_ai.close() + if ds_model is not None: + ds_model.close() + if ds_model1 is not None: + ds_model1.close() + if ds_model2 is not None: + ds_model2.close() if requested_all: if args.variable_list: if not Path(args.variable_list).exists(): - ds_ai.close() - if ds_model is not None: - ds_model.close() + _close_all() raise FileNotFoundError(f"Variable list file not found: {args.variable_list}") all_variables = parse_variable_list_file(args.variable_list) - if use_csv_predictions: + if model_vs_model: + available_vars = [var for var in all_variables if var in common_vars] + elif use_csv_predictions: available_vars = [ var for var in all_variables if var in ai_vars and csv_has_variable(csv_predictions, var, variable_category_map.get(var)) @@ -782,12 +940,23 @@ def _resolve_default(current_value, config_key, fallback): print(f"Using all variables from variable list ({len(available_vars)}): {available_vars}") else: print("Warning: No variables from variable list found in available datasets!") - ds_ai.close() - if ds_model is not None: - ds_model.close() + _close_all() return else: - if use_csv_predictions: + if model_vs_model: + exclude_patterns = ['lon', 'lat', 'index', 'period', 'time', 'bnds'] + discovered_vars = sorted( + v for v in common_vars + if not any(p in v.lower() for p in exclude_patterns) + ) + if discovered_vars: + args.variables = discovered_vars + print(f"Using all common variables between Model 1 and Model 2 ({len(discovered_vars)}): {discovered_vars}") + else: + print("Warning: No common variables found between the two model files!") + _close_all() + return + elif use_csv_predictions: csv_vars = set(csv_predictions.get('available_vars', set())) discovered_vars = sorted(ai_vars.intersection(csv_vars)) if discovered_vars: @@ -795,7 +964,7 @@ def _resolve_default(current_value, config_key, fallback): print(f"Using all common variables between NetCDF and CSV ({len(discovered_vars)}): {discovered_vars}") else: print("Warning: No common variables found between AI NetCDF predictions and CSV source!") - ds_ai.close() + _close_all() return else: discovered_vars = discover_common_variables(ds_ai, ds_model) @@ -804,18 +973,27 @@ def _resolve_default(current_value, config_key, fallback): print(f"Using all common variables ({len(discovered_vars)}): {discovered_vars}") else: print("Warning: No common variables found between AI predictions and model!") - ds_ai.close(); ds_model.close() + _close_all() return else: - forced = ['cwdc_vr', 'soil3c_vr', 'tlai', 'deadstemc'] - if use_csv_predictions: + forced = list(VARIABLES) + if model_vs_model: + selected = [v for v in forced if v in common_vars] + if not selected: + print("Warning: None of the default variables are present in both model files!") + _close_all() + return + missing = [v for v in forced if v not in selected] + if missing: + print(f"Note: Skipping missing default variables not present in both model files: {missing}") + elif use_csv_predictions: selected = [ v for v in forced if v in ai_vars and csv_has_variable(csv_predictions, v, variable_category_map.get(v)) ] if not selected: print("Warning: None of the default variables are present in both NetCDF and CSV data!") - ds_ai.close() + _close_all() return missing = [v for v in forced if v not in selected] if missing: @@ -825,7 +1003,7 @@ def _resolve_default(current_value, config_key, fallback): selected = [v for v in forced if v in ai_vars and v in model_vars] if not selected: print("Warning: None of the default variables are present in both datasets!") - ds_ai.close(); ds_model.close() + _close_all() return missing = [v for v in forced if v not in selected] if missing: @@ -833,7 +1011,19 @@ def _resolve_default(current_value, config_key, fallback): args.variables = selected print(f"Using default subset of variables ({len(selected)}): {selected}") - if use_csv_predictions: + if model_vs_model: + grid_lon, grid_lat = _gridcell_lonlat(ds_model1) + n_grid = ds_model1.sizes['gridcell'] + print(f"Using Model 1 gridcell count: {n_grid}") + print(f"Grid coordinates: lon range [{grid_lon.min():.3f}, {grid_lon.max():.3f}], lat range [{grid_lat.min():.3f}, {grid_lat.max():.3f}]") + col2grid = _to_zero_based_index(_safe_get(ds_model1, 'cols1d_gridcell_index').values, n_grid) + pft2grid = _to_zero_based_index(_safe_get(ds_model1, 'pfts1d_gridcell_index').values, n_grid) + grid_to_cols = _build_gridcell_groups(col2grid, n_grid) + grid_to_pfts = _build_gridcell_groups(pft2grid, n_grid) + ai_to_model_mapping = None + print(f"Model mappings: total columns: {col2grid.size} | total pfts: {pft2grid.size}") + print(f"Example: gridcell 0 -> columns {grid_to_cols[0][:5]}, pfts {grid_to_pfts[0][:5]}") + elif use_csv_predictions: grid_lon, grid_lat = _gridcell_lonlat(ds_ai) n_grid = ds_ai.sizes['gridcell'] print(f"Using AI gridcell count: {n_grid}") @@ -871,6 +1061,62 @@ def _resolve_default(current_value, config_key, fallback): print(f"\nStart processing: {len(args.variables)} variables") stats_rows = [] for var in args.variables: + if model_vs_model: + if var not in common_vars: + print(f"Skip {var} (not in both model files)") + continue + da_ref = ds_model1[var] + dims = da_ref.dims + category = infer_variable_category(var, dims, variable_category_map) + print(f"\nVariable {var}, dims: {dims}") + print(f" NetCDF shape: {da_ref.shape}") + + # Model-vs-model: PFT 1-16 (args.pfts 1-based, default 1..16), top 10 soil layers (args.layers, default 0..9) + if ('column' in dims and 'levgrnd' in dims): + n_lev = int(da_ref.sizes['levgrnd']) + for lev in range(n_lev): + if lev not in args.layers: + continue + g1 = _extract_elm_var_on_grid(ds_model1, var, grid_to_cols, grid_to_pfts, n_grid, layer_idx=lev, pft_idx=None) + g2 = _extract_elm_var_on_grid(ds_model2, var, grid_to_cols, grid_to_pfts, n_grid, layer_idx=lev, pft_idx=None) + if g1 is None or g2 is None: + print(f" Skip {var} lev{lev} (extract failed)") + continue + do_plot = (not args.no_plot) and (lev in args.layers) + stats = _plot_tripanel(var, f"_lev{lev}", grid_lon, grid_lat, g1, g2, args.output_dir, + label_ai='Model 1', label_model='Model 2', plot=do_plot) + stats_rows.append({'variable': var, 'suffix': f"_lev{lev}", **stats}) + + elif 'pft' in dims: + n_pft = 16 + for k in range(n_pft): + pft_one_based = k + 1 # internal k is 0-based; user-facing PFT numbers are 1-based + if pft_one_based not in args.pfts: + continue + g1 = _extract_elm_var_on_grid(ds_model1, var, grid_to_cols, grid_to_pfts, n_grid, layer_idx=None, pft_idx=k) + g2 = _extract_elm_var_on_grid(ds_model2, var, grid_to_cols, grid_to_pfts, n_grid, layer_idx=None, pft_idx=k) + if g1 is None or g2 is None: + print(f" Skip {var} pft{pft_one_based} (extract failed)") + continue + do_plot = (not args.no_plot) and (pft_one_based in args.pfts) + stats = _plot_tripanel(var, f"_pft{k+1}", grid_lon, grid_lat, g1, g2, args.output_dir, + label_ai='Model 1', label_model='Model 2', plot=do_plot) + stats_rows.append({'variable': var, 'suffix': f"_pft{k+1}", **stats}) + + elif 'gridcell' in dims: + g1 = _extract_elm_var_on_grid(ds_model1, var, grid_to_cols, grid_to_pfts, n_grid, layer_idx=None, pft_idx=None) + g2 = _extract_elm_var_on_grid(ds_model2, var, grid_to_cols, grid_to_pfts, n_grid, layer_idx=None, pft_idx=None) + if g1 is None or g2 is None: + print(f" Skip {var} (extract failed)") + continue + stats = _plot_tripanel(var, '', grid_lon, grid_lat, g1, g2, args.output_dir, + label_ai='Model 1', label_model='Model 2', plot=(not args.no_plot)) + stats_rows.append({'variable': var, 'suffix': '', **stats}) + + else: + print(f" Skip {var} (unsupported dimensions for model-vs-model: {dims})") + continue + if var not in ds_ai.data_vars: print(f"Skip {var} (not found in AI NetCDF)") continue @@ -956,17 +1202,20 @@ def _resolve_default(current_value, config_key, fallback): lon_subset = grid_lon[:min_len] lat_subset = grid_lat[:min_len] for k in range(pft_count): + pft_one_based = k + 1 # --pfts is 1-based + if pft_one_based not in args.pfts: + continue ai_slice = ai_vals[k, :min_len] if k < csv_pft_count: csv_slice = csv_vals[k, :min_len] else: csv_slice = np.full(min_len, np.nan, dtype=float) - plot_flag = (not args.no_plot) and (k in args.pfts) - stats = _plot_tripanel(var, f"_pft{k+1}", lon_subset, lat_subset, ai_slice, csv_slice, args.output_dir, + plot_flag = not args.no_plot + stats = _plot_tripanel(var, f"_pft{pft_one_based}", lon_subset, lat_subset, ai_slice, csv_slice, args.output_dir, label_ai=label_ai, label_model=label_csv, plot=plot_flag) stats_rows.append({ 'variable': var, - 'suffix': f"_pft{k+1}", + 'suffix': f"_pft{pft_one_based}", **stats, }) @@ -1084,6 +1333,9 @@ def _resolve_default(current_value, config_key, fallback): total_pfts = vals_ai.shape[0] for k in range(total_pfts): + pft_one_based = k + 1 # --pfts is 1-based + if pft_one_based not in args.pfts: + continue ai_grid = np.full(n_grid, np.nan, dtype=float) model_grid = np.full(n_grid, np.nan, dtype=float) @@ -1099,21 +1351,21 @@ def _resolve_default(current_value, config_key, fallback): if len(ai_gridcell_idx) > 0: ai_grid[g] = vals_ai[k] + # Model has PFT 0-16; use PFT 1-16 only (ignore PFT 0) to match AI PFT 1-16 for g in range(n_grid): - if g < len(grid_to_pfts) and len(grid_to_pfts[g]) > 0: - gridcell_pfts = grid_to_pfts[g][:16] - adjusted_k = k + 1 - if adjusted_k < len(gridcell_pfts): - model_pft_idx = gridcell_pfts[adjusted_k] + if g < len(grid_to_pfts) and len(grid_to_pfts[g]) > 1: + model_pft_1_to_16 = grid_to_pfts[g][1:17] # indices for model PFT 1..16 + if k < len(model_pft_1_to_16): + model_pft_idx = model_pft_1_to_16[k] if model_pft_idx < vals_model.shape[0]: model_grid[g] = vals_model[model_pft_idx] - do_plot = (not args.no_plot) and (k in args.pfts) - stats = _plot_tripanel(var, f"_pft{k+1}", grid_lon, grid_lat, ai_grid, model_grid, args.output_dir, + do_plot = not args.no_plot + stats = _plot_tripanel(var, f"_pft{pft_one_based}", grid_lon, grid_lat, ai_grid, model_grid, args.output_dir, label_ai='AI Predictions', label_model='Model Results', plot=do_plot) stats_rows.append({ 'variable': var, - 'suffix': f"_pft{k+1}", + 'suffix': f"_pft{pft_one_based}", **stats, }) @@ -1199,7 +1451,12 @@ def _suffix_key(s): return (3, 0, 0) lines = [] - header = 'NetCDF vs CSV Statistics Report' if use_csv_predictions else 'AI vs Model Statistics Report' + if model_vs_model: + header = 'Model 1 vs Model 2 Statistics Report' + elif use_csv_predictions: + header = 'NetCDF vs CSV Statistics Report' + else: + header = 'AI vs Model Statistics Report' lines.append(header) lines.append('=' * 80) for var in sorted(grouped.keys()): @@ -1210,7 +1467,10 @@ def _suffix_key(s): for row in rows: title = f"{var}{row.get('suffix','')}" lines.append(title) - if use_csv_predictions: + if model_vs_model: + lines.append(" Model 1: sum={ai_sum:.6g} std={ai_std:.6g} min={ai_min:.6g} max={ai_max:.6g}".format(**row)) + lines.append(" Model 2: sum={model_sum:.6g} std={model_std:.6g} min={model_min:.6g} max={model_max:.6g}".format(**row)) + elif use_csv_predictions: lines.append(" NetCDF: sum={ai_sum:.6g} std={ai_std:.6g} min={ai_min:.6g} max={ai_max:.6g}".format(**row)) lines.append(" CSV: sum={model_sum:.6g} std={model_std:.6g} min={model_min:.6g} max={model_max:.6g}".format(**row)) else: @@ -1256,7 +1516,8 @@ def _suffix_key(s): figsize=(14, 10), color=[colors.get(col, 'gray') for col in plot_cols] ) - title = 'CSV vs NetCDF Agreement by Variable' if use_csv_predictions else 'AI vs Model Agreement by Variable' + title = ('Model 1 vs Model 2 Agreement by Variable' if model_vs_model else + 'CSV vs NetCDF Agreement by Variable' if use_csv_predictions else 'AI vs Model Agreement by Variable') plt.title(title, fontsize=16) plt.xlabel('Variable', fontsize=14) plt.ylabel('Percentage (%)', fontsize=14) @@ -1272,9 +1533,14 @@ def _suffix_key(s): else: print("No statistics to write.") - ds_ai.close() + if ds_ai is not None: + ds_ai.close() if ds_model is not None: ds_model.close() + if ds_model1 is not None: + ds_model1.close() + if ds_model2 is not None: + ds_model2.close() if args.no_plot: print(f"\nCompleted without plotting. Output directory: {args.output_dir}") else: diff --git a/scripts/ai_predictions_to_netcdf.py b/scripts/ai_predictions_to_netcdf.py index 7e93a12..da7104a 100644 --- a/scripts/ai_predictions_to_netcdf.py +++ b/scripts/ai_predictions_to_netcdf.py @@ -175,8 +175,17 @@ def _discover_variables_from_files(predictions_dir: Path) -> Optional[Dict[str, return None -def load_ai_predictions(predictions_dir: Path) -> Dict[str, Any]: - """Load AI model predictions from the predictions directory.""" +def load_ai_predictions( + predictions_dir: Path, + soil_2d_bias_corrected_subdir: Optional[str] = None, +) -> Dict[str, Any]: + """Load AI model predictions from the predictions directory. + + If soil_2d_bias_corrected_subdir is set (e.g. soil_2d_predictions_5P_bias_corrected_phase2), + after loading soil from soil_2d_predictions we also load any + predictions_Y__bias_corrected.csv from that subdir and overlay them onto preds['soil_2d'], + so that the NetCDF uses bias-corrected 5P where applied. + """ print(f"Loading AI predictions from: {predictions_dir}") preds = {} @@ -218,6 +227,9 @@ def load_ai_predictions(predictions_dir: Path) -> Dict[str, Any]: preds['soil_2d'] = {} preds['soil2d_coords'] = {} for p in sorted(soil_dir.glob('predictions_*.csv')): + # Skip bias-corrected filenames here; they are loaded from the bias-corrected subdir if given + if '_bias_corrected' in p.stem: + continue # Extract variable name from filename (e.g., predictions_Y_cwdc_vr.csv -> cwdc_vr) var_name = p.stem.replace('predictions_Y_', '') df = pd.read_csv(p) @@ -229,6 +241,23 @@ def load_ai_predictions(predictions_dir: Path) -> Dict[str, Any]: if 'Longitude' in df and 'Latitude' in df: sample_locs = df[['Longitude', 'Latitude']].drop_duplicates().head(3) print(f" Sample locations (first 3 unique):\n{sample_locs}") + # Overlay bias-corrected soil 2D predictions if subdir provided (e.g. 5P from apply_5p_bias_scale_correction.py) + if soil_2d_bias_corrected_subdir: + bias_dir = predictions_dir / soil_2d_bias_corrected_subdir + if bias_dir.exists(): + if 'soil_2d' not in preds: + preds['soil_2d'] = {} + preds['soil2d_coords'] = {} + for p in sorted(bias_dir.glob('predictions_Y_*_bias_corrected.csv')): + # predictions_Y__bias_corrected.csv -> + var_name = p.stem.replace('predictions_Y_', '').replace('_bias_corrected', '') + df = pd.read_csv(p) + lon, lat = _extract_coords(df) + preds['soil2d_coords'][var_name] = (lon, lat) + preds['soil_2d'][var_name] = _drop_coords(df) + print(f" Overlaid bias-corrected soil predictions for {var_name}: {df.shape}") + else: + print(f" Warning: soil 2D bias-corrected subdir not found: {bias_dir}") # Load static inverse mapping for coordinates static_inv = predictions_dir / 'test_static_inverse.csv' @@ -459,6 +488,10 @@ def main(): help='Show example usage and exit') parser.add_argument('--wrap-longitude', action='store_true', default=False, help='Wrap longitudes from 0–360 to -180–180 (default: on). Use --no-wrap-longitude to disable if shell supports).') + parser.add_argument('--soil-2d-bias-corrected-subdir', default=None, + help='Subdir under --ai-predictions containing bias-corrected soil 2D CSVs ' + '(e.g. soil_2d_predictions_5P_bias_corrected_phase2 from apply_5p_bias_scale_correction.py). ' + 'Files named predictions_Y__bias_corrected.csv will overlay raw soil 2D for those variables.') args = parser.parse_args() @@ -510,8 +543,11 @@ def main(): if not variable_list: parser.error('Could not auto-detect variables. Provide --variable-list to proceed.') - # Load AI predictions - ai_preds = load_ai_predictions(ai_predictions_dir) + # Load AI predictions (optionally overlay 5P bias-corrected from apply_5p_bias_scale_correction.py) + ai_preds = load_ai_predictions( + ai_predictions_dir, + soil_2d_bias_corrected_subdir=getattr(args, 'soil_2d_bias_corrected_subdir', None), + ) # Report current lon/lat ranges from CSV coordinates try: diff --git a/scripts/ai_predictions_to_restart.py b/scripts/ai_predictions_to_restart.py index ee64ff8..20c81d6 100644 --- a/scripts/ai_predictions_to_restart.py +++ b/scripts/ai_predictions_to_restart.py @@ -217,8 +217,14 @@ 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, - update_grid_mask: Optional[np.ndarray] = None) -> None: + tropical_lat_range: Optional[tuple] = None) -> None: + """ + tropical_lat_range: If (min_lat, max_lat), only update gridcells with lat in [min_lat, max_lat]; + others are left unchanged (for merging tropical-model into global restart). + """ print(f"Saving updated restart file to: {output_path}") + if tropical_lat_range is not None: + print(f" Tropical-only update: lat in [{tropical_lat_range[0]}, {tropical_lat_range[1]}]") # Create output directory if it doesn't exist output_path.parent.mkdir(parents=True, exist_ok=True) @@ -229,6 +235,17 @@ 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: + # Build mask of gridcells to update (all, or only those in tropical_lat_range) + n_grid = variable_mapping.get("n_grid", 0) + update_mask = np.ones(n_grid, dtype=bool) + if tropical_lat_range is not None and n_grid > 0: + if "grid1d_lat" in ds_out.variables: + grid_lat = np.asarray(ds_out.variables["grid1d_lat"][:]).ravel() + min_lat, max_lat = float(tropical_lat_range[0]), float(tropical_lat_range[1]) + update_mask = (grid_lat >= min_lat) & (grid_lat <= max_lat) + print(f" Gridcells in tropical band: {update_mask.sum()} / {n_grid}") + else: + print(" Warning: grid1d_lat not found; applying tropical filter to all gridcells") # Verify and adjust spinup_state try: if 'spinup_state' in ds_out.variables: @@ -309,7 +326,7 @@ 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]): + if not update_mask[g]: continue if g < len(grid_to_pfts) and len(grid_to_pfts[g]) > 0: # Get PFTs in this gridcell @@ -347,7 +364,7 @@ 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]): + if not update_mask[g]: continue if g < len(grid_to_cols) and len(grid_to_cols[g]) > 0: # Get columns in this gridcell @@ -429,6 +446,15 @@ def main(): --restart-file model_restart.nc \ --output updated_restart.nc \ --preview-only + + # Phase 2: update only P variables in tropical cells (merge Phase 2 into Phase 1 restart) + python ai_predictions_to_restart.py \ + --ai-predictions phase2_predictions.nc \ + --restart-file phase1_global_restart.nc \ + --output merged_restart.nc \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --tropical-lat-range -30,30 \ + --variables-to-update occlp_vr,labilep_vr,solutionp_vr,primp_vr,secondp_vr,soil1p_vr,soil2p_vr,soil3p_vr,soil4p_vr,litr2p_vr,litr3p_vr,cwdp_vr """ ) @@ -446,12 +472,10 @@ 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') + parser.add_argument('--tropical-lat-range', type=str, metavar='MIN,MAX', default=None, + help='Only update gridcells with lat in [MIN, MAX] (e.g. "-30,30"). Use when merging tropical-model predictions into a global restart.') + parser.add_argument('--variables-to-update', type=str, default=None, metavar='VAR1,VAR2,...|@file.txt', + help='Only update these variables (subset of CNP_IO list). Comma-separated names (e.g. occlp_vr,labilep_vr,solutionp_vr) or path to a file with one variable per line (e.g. @phase2_p_vars.txt). If not set, all variables from the variable list are updated. Use with Phase 2 to update only P variables in tropical cells.') args = parser.parse_args() @@ -530,9 +554,31 @@ def main(): # Auto-detect from config.json cnp_io_variables = auto_detect_variable_list(Path(args.ai_predictions)) - # Note: We only update variables in the CNP_IO list - # All other variables (including timemgr_rst_nstep_rad_prev) remain completely unchanged - print("Note: Only variables in CNP_IO list will be updated") + # Optionally restrict to a subset of variables (e.g. P-only for Phase 2 merge) + effective_update_variables = list(cnp_io_variables) + if getattr(args, 'variables_to_update', None) and args.variables_to_update.strip(): + raw = args.variables_to_update.strip() + if raw.startswith('@'): + path = Path(raw[1:].strip()) + if path.exists(): + with open(path, 'r') as f: + requested = {line.strip() for line in f if line.strip() and not line.strip().startswith('#')} + else: + parser.error(f'Variables-to-update file not found: {path}') + else: + requested = {v.strip() for v in raw.split(',') if v.strip()} + effective_update_variables = [v for v in cnp_io_variables if v in requested] + not_in_cnp = requested - set(cnp_io_variables) + if not_in_cnp: + print(f" Note: --variables-to-update names not in CNP_IO list (skipped): {sorted(not_in_cnp)}") + if not effective_update_variables: + parser.error('--variables-to-update resulted in no variables to update (none matched CNP_IO list)') + print(f" Restricting to {len(effective_update_variables)} variables: {effective_update_variables}") + else: + effective_update_variables = list(cnp_io_variables) + + # Note: We only update variables in the effective list (CNP_IO list, optionally filtered) + print("Note: Only variables in the update list will be modified in the restart") print("All other variables and attributes remain unchanged") # Print summary of changes @@ -540,31 +586,31 @@ def main(): print("UPDATE SUMMARY") print("=" * 60) - # Count variables that will be updated (only PFT1D and soil2D from CNP_IO list) + # Count variables that will be updated (only PFT1D and soil2D from effective list) updated_vars = [] for var_name in ds_ai.data_vars: - if var_name in ds_model.data_vars and var_name in cnp_io_variables: - # Only count PFT1D and soil2D variables that are in the CNP_IO list + if var_name in ds_model.data_vars and var_name in effective_update_variables: + # Only count PFT1D and soil2D variables that are in the update list if ('pft' in ds_ai[var_name].dims) or ('column' in ds_ai[var_name].dims and 'levgrnd' in ds_ai[var_name].dims): updated_vars.append(var_name) - print(f"Variables to update (PFT1D and soil2D from CNP_IO list): {len(updated_vars)}") + print(f"Variables to update (PFT1D and soil2D): {len(updated_vars)}") for var_name in updated_vars: ai_shape = ds_ai[var_name].shape model_shape = ds_model[var_name].shape var_type = "PFT1D" if 'pft' in ds_ai[var_name].dims else "Soil2D" print(f" {var_name} ({var_type}): AI {ai_shape} -> Model {model_shape}") - # Show which variables were skipped + # Show which variables were skipped (in CNP_IO but not in effective update list) skipped_vars = [] for var_name in ds_ai.data_vars: if (var_name in ds_model.data_vars and ('pft' in ds_ai[var_name].dims or ('column' in ds_ai[var_name].dims and 'levgrnd' in ds_ai[var_name].dims)) and - var_name not in cnp_io_variables): + var_name in cnp_io_variables and var_name not in effective_update_variables): skipped_vars.append(var_name) if skipped_vars: - print(f"\nVariables skipped (not in CNP_IO list): {len(skipped_vars)}") + print(f"\nVariables skipped (not in --variables-to-update): {len(skipped_vars)}") for var_name in skipped_vars: print(f" {var_name}") @@ -584,9 +630,15 @@ def main(): print(f"Backup created: {backup_path.stat().st_size / (1024*1024):.1f} MB") # Save updated restart file using direct NetCDF manipulation + tropical_lat_range = None + if getattr(args, 'tropical_lat_range', None): + parts = [p.strip() for p in args.tropical_lat_range.split(',')] + if len(parts) >= 2: + tropical_lat_range = (float(parts[0]), float(parts[1])) 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, update_grid_mask=update_grid_mask) + effective_update_variables, ai_to_model_mapping, variable_mapping, + strict_dims=args.strict_dims, + tropical_lat_range=tropical_lat_range) print(f"\nRestart file updated successfully!") print(f"Original: {restart_file_path}") diff --git a/scripts/ai_restart_comparison.py b/scripts/ai_restart_comparison.py index 8970a2a..d6f801e 100644 --- a/scripts/ai_restart_comparison.py +++ b/scripts/ai_restart_comparison.py @@ -414,8 +414,25 @@ def _resolve(value, keys, fallback): grid_to_cols = _build_gridcell_groups(col2grid, n_grid) grid_to_pfts = _build_gridcell_groups(pft2grid, n_grid) + # Build reference-side grid groups when reference is a restart file (column/pft indexed, no gridcell on variables) + n_grid_old = ds_old.sizes['gridcell'] + ref_has_col_index = 'cols1d_gridcell_index' in ds_old + ref_has_pft_index = 'pfts1d_gridcell_index' in ds_old + grid_to_cols_old = None + grid_to_pfts_old = None + if ref_has_col_index: + col2grid_old = _to_zero_based_index(_safe_get(ds_old, 'cols1d_gridcell_index').values, n_grid_old) + grid_to_cols_old = _build_gridcell_groups(col2grid_old, n_grid_old) + if ref_has_pft_index: + pft2grid_old = _to_zero_based_index(_safe_get(ds_old, 'pfts1d_gridcell_index').values, n_grid_old) + grid_to_pfts_old = _build_gridcell_groups(pft2grid_old, n_grid_old) + print(f'Total gridcells: {n_grid} | total columns: {col2grid.size} | total pfts: {pft2grid.size}') print(f'Example: gridcell 0 -> columns {grid_to_cols[0][:5]}, pfts {grid_to_pfts[0][:5]}') + if grid_to_cols_old is not None: + print(f'Reference: restart-style (column-indexed); n_grid_old={n_grid_old}') + if grid_to_pfts_old is not None: + print(f'Reference: restart-style (pft-indexed)') stats_rows = [] debug_enabled = not args.stats_only @@ -435,22 +452,30 @@ def _resolve(value, keys, fallback): da_new_cl = da_new.transpose('column', 'levgrnd', ...) vals_new = _to_nan_fillvalue(da_new_cl.values) - if 'gridcell' not in da_old.dims: - print(' Skip (reference dataset lacks gridcell dimension)') + ref_restart_style = ('gridcell' not in da_old.dims) and (grid_to_cols_old is not None) + if not ref_restart_style and 'gridcell' not in da_old.dims: + print(' Skip (reference dataset lacks gridcell dimension and not column-indexed restart)') continue - old_order = [dim for dim in da_old.dims if dim != 'gridcell'] + ['gridcell'] - da_old_cl = da_old.transpose(*old_order) - vals_old = _to_nan_fillvalue(da_old_cl.values) - dims_old = da_old_cl.dims - axis_grid = dims_old.index('gridcell') - axis_column = dims_old.index('column') if 'column' in dims_old else None - axis_lev = dims_old.index('levgrnd') if 'levgrnd' in dims_old else None + if ref_restart_style: + # Reference is restart file: (column, levgrnd); use grid_to_cols_old + grid_mapping + da_old_cl = da_old.transpose('column', 'levgrnd', ...) + vals_old = _to_nan_fillvalue(da_old_cl.values) + else: + old_order = [dim for dim in da_old.dims if dim != 'gridcell'] + ['gridcell'] + da_old_cl = da_old.transpose(*old_order) + vals_old = _to_nan_fillvalue(da_old_cl.values) + dims_old = da_old_cl.dims + axis_grid = dims_old.index('gridcell') + axis_column = dims_old.index('column') if 'column' in dims_old else None + axis_lev = dims_old.index('levgrnd') if 'levgrnd' in dims_old else None for lev in LEVGRND_LAYERS: if lev < 0 or lev >= da_new_cl.sizes['levgrnd']: print(f' Layer {lev} out of range, skipped') continue + if not ref_restart_style and lev >= da_old_cl.sizes['levgrnd']: + continue new_grid = np.full(n_grid, np.nan, dtype=float) old_grid = np.full(n_grid, np.nan, dtype=float) @@ -465,20 +490,28 @@ def _resolve(value, keys, fallback): else: new_grid[g] = vals_new[c0, lev] - src_idx = int(grid_mapping[g]) if g < len(grid_mapping) else -1 - if src_idx < 0 or src_idx >= da_old_cl.sizes['gridcell']: + src_g = int(grid_mapping[g]) if g < len(grid_mapping) else -1 + if src_g < 0 or src_g >= n_grid_old: continue - idx = [slice(None)] * vals_old.ndim - if axis_column is not None: - col_sel = min(c0, vals_old.shape[axis_column] - 1) - idx[axis_column] = col_sel - if axis_lev is not None: - if lev >= vals_old.shape[axis_lev]: + if ref_restart_style: + cols_old = grid_to_cols_old[src_g] + if len(cols_old) == 0: continue - idx[axis_lev] = lev - idx[axis_grid] = src_idx - old_grid[g] = vals_old[tuple(idx)] + c0_old = cols_old[0] + if c0_old < vals_old.shape[0] and lev < vals_old.shape[1]: + old_grid[g] = vals_old[c0_old, lev] + else: + idx = [slice(None)] * vals_old.ndim + if axis_column is not None: + col_sel = min(c0, vals_old.shape[axis_column] - 1) + idx[axis_column] = col_sel + if axis_lev is not None: + if lev >= vals_old.shape[axis_lev]: + continue + idx[axis_lev] = lev + idx[axis_grid] = src_g + old_grid[g] = vals_old[tuple(idx)] if debug_enabled: print(f'[DEBUG] {var} lev{lev}: new min={np.nanmin(new_grid)} max={np.nanmax(new_grid)} mean={np.nanmean(new_grid)}') @@ -489,15 +522,20 @@ def _resolve(value, keys, fallback): stats_rows.append({'variable': var, 'suffix': f'_lev{lev}', 'label_new': label_new, 'label_old': label_old, 'quality': _classify_quality_from_r2(stats['r2']), **stats}) elif 'pft' in dims: - if 'gridcell' not in da_old.dims: - print(' Skip (reference dataset lacks gridcell dimension)') + ref_pft_restart_style = ('gridcell' not in da_old.dims) and (grid_to_pfts_old is not None) + if not ref_pft_restart_style and 'gridcell' not in da_old.dims: + print(' Skip (reference dataset lacks gridcell dimension and not pft-indexed restart)') continue da_new_p = da_new.transpose(..., 'pft') vals_new = _to_nan_fillvalue(da_new_p.values) - da_old_p = da_old.transpose('pft', 'gridcell') - vals_old = _to_nan_fillvalue(da_old_p.values) + if ref_pft_restart_style: + da_old_p = da_old.transpose('pft', ...) + vals_old = _to_nan_fillvalue(da_old_p.values) + else: + da_old_p = da_old.transpose('pft', 'gridcell') + vals_old = _to_nan_fillvalue(da_old_p.values) total_pfts = da_new_p.sizes.get('pft', vals_new.shape[0]) for k in PFT_PICK_LIST: @@ -518,12 +556,22 @@ def _resolve(value, keys, fallback): p_idx = pfts[k] if p_idx < vals_new.shape[0]: new_grid[g] = vals_new[p_idx] - src_idx = int(grid_mapping[g]) if g < len(grid_mapping) else -1 - if src_idx < 0 or src_idx >= vals_old.shape[1]: - continue - ai_pft_idx = k - 1 - if ai_pft_idx >= 0 and ai_pft_idx < vals_old.shape[0]: - old_grid[g] = vals_old[ai_pft_idx, src_idx] + + src_g = int(grid_mapping[g]) if g < len(grid_mapping) else -1 + if ref_pft_restart_style: + if src_g < 0 or src_g >= n_grid_old or grid_to_pfts_old is None: + continue + pfts_old = grid_to_pfts_old[src_g] + if k < len(pfts_old): + p_idx_old = pfts_old[k] + if p_idx_old < vals_old.shape[0]: + old_grid[g] = vals_old[p_idx_old] + else: + if src_g < 0 or src_g >= vals_old.shape[1]: + continue + ai_pft_idx = k - 1 + if ai_pft_idx >= 0 and ai_pft_idx < vals_old.shape[0]: + old_grid[g] = vals_old[ai_pft_idx, src_g] stats = _plot_tripanel(var, f'_pft{k}', grid_lon, grid_lat, new_grid, old_grid, str(output_dir), label_new=label_new, label_old=label_old, plot=plot_enabled) diff --git a/scripts/analyze_5p_ground_truth_by_region.py b/scripts/analyze_5p_ground_truth_by_region.py new file mode 100644 index 0000000..bc90592 --- /dev/null +++ b/scripts/analyze_5p_ground_truth_by_region.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +""" +Analyze ground-truth 5 P variables (labilep_vr, occlp_vr, solutionp_vr, secondp_vr, primp_vr) +in the two target regions (Amazon + Central Africa) vs the rest of the world. + +Answers: Are the 5 P distributions different in these two regions compared to other regions? +If yes, training only on two-region data is better justified. + +Usage: + python scripts/analyze_5p_ground_truth_by_region.py --data-dir /path/to/Trendy_1_data_CNP [--max-files 20] + python scripts/analyze_5p_ground_truth_by_region.py # uses default path and first 20 files + +Output: Prints summary table and writes analysis_5p_by_region.txt (and optional CSV) in current dir or --output. +""" + +import argparse +import numpy as np +import pandas as pd +import sys +from pathlib import Path + +# Two regions (lat_min, lat_max, lon_min, lon_max), longitude in 0-360 +AMAZON_BOX = (-30, 10, 270, 330) # 10°N to 30°S, 90°W to 30°W +AFRICA_BOX = (-15, 15, 0, 30) # 15°N to 15°S, 0° to 30°E +FIVE_P_COLS = ["Y_labilep_vr", "Y_occlp_vr", "Y_solutionp_vr", "Y_secondp_vr", "Y_primp_vr"] +SOIL_LAYERS = 10 +LAT_CANDIDATES = ["Latitude", "lat", "latitude", "LAT"] +LON_CANDIDATES = ["Longitude", "lon", "longitude", "LON"] + + +def _to_360(lon: np.ndarray) -> np.ndarray: + out = np.asarray(lon, dtype=float).copy() + out[np.isnan(out)] = 0 + out[out < 0] += 360.0 + return out + + +def _extract_col0_layers10(arr) -> np.ndarray: + """Extract first column, first 10 layers from a 2D soil variable (row of DataFrame).""" + try: + a = np.asarray(arr) + if a.ndim == 2: + return np.asarray(a[0, :SOIL_LAYERS].ravel(), dtype=float) + if a.ndim == 1: + return np.asarray(a[:SOIL_LAYERS], dtype=float) + return np.full(SOIL_LAYERS, np.nan) + except Exception: + return np.full(SOIL_LAYERS, np.nan) + + +def _row_summary(vals: np.ndarray) -> dict: + """Per-row summary: sum, mean, layer0 (surface).""" + v = np.asarray(vals, dtype=float) + v = v[~np.isnan(v)] + if v.size == 0: + return {"sum": np.nan, "mean": np.nan, "layer0": np.nan} + return { + "sum": float(np.sum(v)), + "mean": float(np.mean(v)), + "layer0": float(v[0]) if len(v) > 0 else np.nan, + } + + +def load_data(data_dir: Path, file_pattern: str = "training_data_batch_*.pkl", max_files: int = 20) -> pd.DataFrame: + data_dir = Path(data_dir) + if not data_dir.is_dir(): + raise FileNotFoundError(f"Data directory not found: {data_dir}") + files = sorted(data_dir.glob(file_pattern))[:max_files] + if not files: + raise FileNotFoundError(f"No files matching {file_pattern} in {data_dir}") + dfs = [] + for f in files: + try: + df = pd.read_pickle(f) + dfs.append(df) + except Exception as e: + print(f"Warning: failed to load {f}: {e}", file=sys.stderr) + if not dfs: + raise RuntimeError("No data loaded") + return pd.concat(dfs, ignore_index=True) + + +def main(): + parser = argparse.ArgumentParser(description="Analyze 5 P ground truth: two regions vs rest of world") + parser.add_argument( + "--data-dir", + type=str, + default="/mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_1_data_CNP", + help="Directory containing training_data_batch_*.pkl", + ) + parser.add_argument("--max-files", type=int, default=20, help="Max number of pkl files to load") + parser.add_argument("--file-pattern", type=str, default="training_data_batch_*.pkl") + parser.add_argument("--output", type=str, default="analysis_5p_by_region.txt", help="Output text report path") + parser.add_argument("--output-csv", type=str, default="", help="Optional: write per-variable stats CSV") + args = parser.parse_args() + + print("Loading data...") + df = load_data(Path(args.data_dir), file_pattern=args.file_pattern, max_files=args.max_files) + print(f"Loaded {len(df)} rows") + + # Resolve lat/lon + lat_col = next((c for c in LAT_CANDIDATES if c in df.columns), None) + lon_col = next((c for c in LON_CANDIDATES if c in df.columns), None) + if lat_col is None or lon_col is None: + print("ERROR: Latitude/Longitude columns not found. Available:", list(df.columns)[:20], "...") + sys.exit(1) + lat = pd.to_numeric(df[lat_col], errors="coerce").values + lon = _to_360(pd.to_numeric(df[lon_col], errors="coerce").values) + + # Mask: in Amazon or Africa + in_amazon = (lat >= AMAZON_BOX[0]) & (lat <= AMAZON_BOX[1]) & (lon >= AMAZON_BOX[2]) & (lon <= AMAZON_BOX[3]) + in_africa = (lat >= AFRICA_BOX[0]) & (lat <= AFRICA_BOX[1]) & (lon >= AFRICA_BOX[2]) & (lon <= AFRICA_BOX[3]) + two_region = in_amazon | in_africa + other = ~two_region + n_two = int(np.sum(two_region)) + n_other = int(np.sum(other)) + print(f"Two regions (Amazon + Africa): {n_two} cells") + print(f"Other: {n_other} cells") + + # Build per-variable stats for "sum" over 10 layers (main comparison) + results = [] + lines = [] + lines.append("=" * 80) + lines.append("5 P ground truth: Two regions (Amazon + Africa) vs Rest of world") + lines.append("=" * 80) + lines.append(f"Data: {args.data_dir} (max {args.max_files} files)") + lines.append(f"Two regions: Amazon {AMAZON_BOX}, Africa {AFRICA_BOX}") + lines.append(f"Two-region cells: {n_two} | Other cells: {n_other}") + lines.append("") + + for col in FIVE_P_COLS: + if col not in df.columns: + lines.append(f"{col}: NOT FOUND in DataFrame") + continue + # Extract first column, first 10 layers per row + raw = np.stack([_extract_col0_layers10(v) for v in df[col].values]) + # Per-row sum over layers (total P in top 10 layers) + row_sum = np.nansum(raw, axis=1) + row_mean = np.nanmean(raw, axis=1) + layer0 = raw[:, 0] + + s_two_sum = row_sum[two_region] + s_other_sum = row_sum[other] + s_two_sum = s_two_sum[~np.isnan(s_two_sum) & np.isfinite(s_two_sum)] + s_other_sum = s_other_sum[~np.isnan(s_other_sum) & np.isfinite(s_other_sum)] + + def stats(x): + if len(x) == 0: + return {"mean": np.nan, "std": np.nan, "p5": np.nan, "p25": np.nan, "p50": np.nan, "p75": np.nan, "p95": np.nan, "n": 0} + return { + "mean": float(np.mean(x)), + "std": float(np.std(x)), + "p5": float(np.percentile(x, 5)), + "p25": float(np.percentile(x, 25)), + "p50": float(np.percentile(x, 50)), + "p75": float(np.percentile(x, 75)), + "p95": float(np.percentile(x, 95)), + "n": len(x), + } + + st_two = stats(s_two_sum) + st_other = stats(s_other_sum) + results.append({ + "variable": col, + "two_region_n": st_two["n"], + "other_n": st_other["n"], + "two_region_mean": st_two["mean"], + "other_mean": st_other["mean"], + "two_region_std": st_two["std"], + "other_std": st_other["std"], + "two_region_p50": st_two["p50"], + "other_p50": st_other["p50"], + "ratio_mean_two_over_other": st_two["mean"] / st_other["mean"] if st_other["mean"] and np.isfinite(st_other["mean"]) else np.nan, + }) + + ratio = st_two["mean"] / st_other["mean"] if st_other["mean"] and np.isfinite(st_other["mean"]) else np.nan + lines.append(f"# {col} (sum over first 10 layers)") + lines.append(f" Two regions: n={st_two['n']} mean={st_two['mean']:.6g} std={st_two['std']:.6g} p50={st_two['p50']:.6g}") + lines.append(f" Other: n={st_other['n']} mean={st_other['mean']:.6g} std={st_other['std']:.6g} p50={st_other['p50']:.6g}") + lines.append(f" Ratio (two_region / other) mean = {ratio:.4f}") + lines.append("") + + # Summary verdict + lines.append("---") + lines.append("Conclusion (from ratio of means):") + for r in results: + if "variable" not in r: + continue + ratio = r.get("ratio_mean_two_over_other", np.nan) + if np.isfinite(ratio) and ratio != 0: + if abs(ratio - 1.0) > 0.2: + lines.append(f" {r['variable']}: DISTRIBUTIONS DIFFER (ratio={ratio:.3f}) — two-region-only training is well justified.") + else: + lines.append(f" {r['variable']}: Similar (ratio={ratio:.3f}) — two-region focus still helps by reducing noise from other regions.") + else: + lines.append(f" {r['variable']}: Could not compute ratio.") + lines.append("") + lines.append("Recommendation: If 5 P in the two regions differ from the rest of the world, training") + lines.append("only on two-region data (region_boxes) should improve 5 P predictions there. You are") + lines.append("already doing this with config/training_config_two_region_five_p.json.") + lines.append("=" * 80) + + report = "\n".join(lines) + print(report) + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(report, encoding="utf-8") + print(f"Wrote {out_path}") + + if args.output_csv and results: + res_df = pd.DataFrame(results) + Path(args.output_csv).parent.mkdir(parents=True, exist_ok=True) + res_df.to_csv(args.output_csv, index=False) + print(f"Wrote {args.output_csv}") + + +if __name__ == "__main__": + main() diff --git a/scripts/analyze_npool_ppool_special_values.py b/scripts/analyze_npool_ppool_special_values.py new file mode 100644 index 0000000..2b64cf1 --- /dev/null +++ b/scripts/analyze_npool_ppool_special_values.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +""" +Analyze training dataset (ground truth) for npool and ppool special values. + +Checks whether the special values 10 (npool) and 1 (ppool) appear in the +global arrays and counts how many grid cells have these values. + +Usage: + python scripts/analyze_npool_ppool_special_values.py --data-paths /path/to/training_data --file-pattern "training_data_batch_*.pkl" + python scripts/analyze_npool_ppool_special_values.py --config cnp_results/run_YYYYMMDD_HHMMSS/cnp_config.json + +Output: printed summary and optional --output report file. +""" + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + +# Repo root +REPO_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO_ROOT)) + +# Special values (from ELM/TRENDY observations) +NPOOL_SPECIAL = 10.0 +PPOOL_SPECIAL = 1.0 +TOL = 1e-5 # tolerance for float comparison + + +def _to_array(x, length=17): + """Convert to 1D float array and pad/truncate to length (17 PFTs incl. PFT0, or 16 after drop).""" + if isinstance(x, np.ndarray): + arr = np.asarray(x, dtype=np.float64).ravel() + elif isinstance(x, (list, tuple)): + arr = np.array(x, dtype=np.float64).ravel() + else: + arr = np.array([float(x)], dtype=np.float64) + if arr.size >= length: + return arr[:length] + return np.pad(arr, (0, length - arr.size), mode="constant", constant_values=np.nan) + + +def _eq_special(arr, value): + arr = np.asarray(arr, dtype=np.float64) + return np.abs(arr - value) < TOL + + +def load_data_from_paths(data_paths, file_pattern, max_files=None): + """Load all PKL files from data_paths into one DataFrame.""" + if isinstance(data_paths, str): + data_paths = [data_paths] + all_dfs = [] + for data_path in data_paths: + data_path = Path(data_path) + if not data_path.exists(): + print(f"Warning: path does not exist: {data_path}", file=sys.stderr) + continue + files = sorted(data_path.glob(file_pattern)) + if max_files is not None: + files = files[:max_files] + for f in files: + try: + df = pd.read_pickle(f) + all_dfs.append(df) + except Exception as e: + print(f"Warning: failed to load {f}: {e}", file=sys.stderr) + if not all_dfs: + return None + return pd.concat(all_dfs, ignore_index=True) + + +def analyze_npool_ppool(df, col_npool="Y_npool", col_ppool="Y_ppool"): + """Analyze special values in Y_npool and Y_ppool columns.""" + if col_npool not in df.columns: + raise KeyError(f"Column {col_npool} not in DataFrame. Columns: {list(df.columns)[:20]}...") + if col_ppool not in df.columns: + raise KeyError(f"Column {col_ppool} not in DataFrame.") + + n_cells = len(df) + # Fixed length 17 (PFT0..PFT16); pad/truncate so all rows have same shape + npool_arrs = np.stack([_to_array(df[col_npool].iloc[i], length=17) for i in range(n_cells)]) + ppool_arrs = np.stack([_to_array(df[col_ppool].iloc[i], length=17) for i in range(n_cells)]) + # Drop PFT0 -> 16 PFTs (PFT1..PFT16) + npool_arrs = npool_arrs[:, 1:] + ppool_arrs = ppool_arrs[:, 1:] + n_pfts = 16 + + # Masks: (n_cells, n_pfts) + npool_is_10 = _eq_special(npool_arrs, NPOOL_SPECIAL) + ppool_is_1 = _eq_special(ppool_arrs, PPOOL_SPECIAL) + + # Counts + # Grid cells where ALL PFTs have npool == 10 + cells_all_npool_10 = npool_is_10.all(axis=1).sum() + # Grid cells where ALL PFTs have ppool == 1 + cells_all_ppool_1 = ppool_is_1.all(axis=1).sum() + # Grid cells where ALL PFTs have both (npool==10 and ppool==1) + cells_both = (npool_is_10 & ppool_is_1).all(axis=1).sum() + + # Grid cells where AT LEAST ONE PFT has the special value + cells_any_npool_10 = npool_is_10.any(axis=1).sum() + cells_any_ppool_1 = ppool_is_1.any(axis=1).sum() + + # Total (cell, PFT) pairs with special value + pairs_npool_10 = int(npool_is_10.sum()) + pairs_ppool_1 = int(ppool_is_1.sum()) + total_pairs = n_cells * n_pfts + + # Per-PFT: how many grid cells have that PFT == special value + cells_per_pft_npool_10 = npool_is_10.sum(axis=0) + cells_per_pft_ppool_1 = ppool_is_1.sum(axis=0) + + return { + "n_cells": n_cells, + "n_pfts": n_pfts, + "total_pairs": total_pairs, + "cells_all_npool_10": int(cells_all_npool_10), + "cells_all_ppool_1": int(cells_all_ppool_1), + "cells_both_all": int(cells_both), + "cells_any_npool_10": int(cells_any_npool_10), + "cells_any_ppool_1": int(cells_any_ppool_1), + "pairs_npool_10": pairs_npool_10, + "pairs_ppool_1": pairs_ppool_1, + "pct_pairs_npool_10": 100.0 * pairs_npool_10 / total_pairs if total_pairs else 0, + "pct_pairs_ppool_1": 100.0 * pairs_ppool_1 / total_pairs if total_pairs else 0, + "cells_per_pft_npool_10": cells_per_pft_npool_10.tolist(), + "cells_per_pft_ppool_1": cells_per_pft_ppool_1.tolist(), + } + + +def main(): + ap = argparse.ArgumentParser(description="Analyze npool/ppool special values (10 and 1) in training data.") + ap.add_argument("--data-paths", type=str, nargs="+", help="Paths to directories containing training PKL files.") + ap.add_argument("--file-pattern", type=str, default="training_data_batch_*.pkl", help="Glob pattern for PKL files.") + ap.add_argument("--config", type=str, help="Path to cnp_config.json; overrides --data-paths and --file-pattern from data_config.") + ap.add_argument("--max-files", type=int, default=None, help="Limit number of PKL files to load (default: all).") + ap.add_argument("--output", type=str, help="Write JSON report to this file.") + args = ap.parse_args() + + data_paths = args.data_paths + file_pattern = args.file_pattern + if args.config: + with open(args.config) as f: + cfg = json.load(f) + dc = cfg.get("data_config") or {} + data_paths = dc.get("data_paths") or data_paths + file_pattern = dc.get("file_pattern") or file_pattern + if not data_paths: + print("No data_paths in config.", file=sys.stderr) + sys.exit(1) + + if not data_paths: + print("Provide --data-paths or --config.", file=sys.stderr) + sys.exit(1) + + df = load_data_from_paths(data_paths, file_pattern, max_files=args.max_files) + if df is None or len(df) == 0: + print("No data loaded.", file=sys.stderr) + sys.exit(1) + + print(f"Loaded {len(df)} grid cells from {data_paths} ({file_pattern}).") + try: + out = analyze_npool_ppool(df) + except KeyError as e: + print(e, file=sys.stderr) + sys.exit(1) + + # Report + print() + print("=== NPOOL / PPOOL special values in training dataset (ground truth) ===") + print() + print(f"Total grid cells: {out['n_cells']}") + print(f"PFTs per cell: {out['n_pfts']} (PFT1..PFT{out['n_pfts']})") + print(f"Total (cell, PFT) pairs: {out['total_pairs']}") + print() + print("Special value: npool == 10") + print(f" Grid cells where ALL PFTs have npool == 10: {out['cells_all_npool_10']} ({100*out['cells_all_npool_10']/out['n_cells']:.1f}% of cells)") + print(f" Grid cells where ANY PFT has npool == 10: {out['cells_any_npool_10']} ({100*out['cells_any_npool_10']/out['n_cells']:.1f}% of cells)") + print(f" Total (cell, PFT) pairs with npool == 10: {out['pairs_npool_10']} ({out['pct_pairs_npool_10']:.1f}% of pairs)") + print() + print("Special value: ppool == 1") + print(f" Grid cells where ALL PFTs have ppool == 1: {out['cells_all_ppool_1']} ({100*out['cells_all_ppool_1']/out['n_cells']:.1f}% of cells)") + print(f" Grid cells where ANY PFT has ppool == 1: {out['cells_any_ppool_1']} ({100*out['cells_any_ppool_1']/out['n_cells']:.1f}% of cells)") + print(f" Total (cell, PFT) pairs with ppool == 1: {out['pairs_ppool_1']} ({out['pct_pairs_ppool_1']:.1f}% of pairs)") + print() + print(f"Grid cells where ALL PFTs have BOTH npool==10 and ppool==1: {out['cells_both_all']} ({100*out['cells_both_all']/out['n_cells']:.1f}% of cells)") + print() + print("Per-PFT: number of grid cells with npool==10 (PFT1..PFT16):") + print(" " + ", ".join(str(x) for x in out["cells_per_pft_npool_10"])) + print("Per-PFT: number of grid cells with ppool==1 (PFT1..PFT16):") + print(" " + ", ".join(str(x) for x in out["cells_per_pft_ppool_1"])) + + if args.output: + with open(args.output, "w") as f: + json.dump(out, f, indent=2) + print(f"\nReport written to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/apply_5p_bias_scale_correction.py b/scripts/apply_5p_bias_scale_correction.py new file mode 100644 index 0000000..a746e83 --- /dev/null +++ b/scripts/apply_5p_bias_scale_correction.py @@ -0,0 +1,601 @@ +import argparse +import json +import os +from dataclasses import dataclass +from typing import Dict, List, Optional, Set, Tuple + +import numpy as np +import pandas as pd + + +P_VARIABLES_DEFAULT: List[str] = [ + "labilep_vr", + "occlp_vr", + "solutionp_vr", + "secondp_vr", + "primp_vr", +] + +# Variables with very small values where linear (a*pred+b) often goes negative; +# use multiplicative-only correction (Y = a*pred) so concentrations stay non-negative. +MULTIPLICATIVE_ONLY_VARS: set = set() + +# Variables with very small values: fit on (pred*scale, gt*scale) for numerical stability, +# then apply corrected = (a*pred*scale + b) / scale. Keeps 5 decimals when writing. +SCALE_FACTOR_VARS: Dict[str, float] = {"solutionp_vr": 1000.0} +DECIMAL_PLACES_SCALED_VARS: int = 8 # keep small solutionp_vr values (e.g. 3e-6) visible + +# Variables where we minimize relative (percent) error so that layer-wise error is <10%. +# Uses weighted least squares with weight 1/(gt+eps)^2. +RELATIVE_ERROR_WEIGHTED_VARS: Set[str] = {"solutionp_vr"} +RELATIVE_ERROR_EPS: float = 1e-12 + +# Reference sites (lon, lat) to prioritize in relative-error fit so layer error <10% there. +REFERENCE_SITES: List[Tuple[float, float]] = [ + (303.75, -17.434553), # default Amazon site + (300.0, 4.240838), # Site A + (292.5, -15.549738), # Site B +] +REFERENCE_SITE_EXTRA_WEIGHT: float = 500.0 # extra weight for cells at these sites +REFERENCE_SITE_ATOL: float = 1e-4 + + +@dataclass +class RegionBox: + """Lat/lon bounds for a rectangular region in degrees (lon in 0–360).""" + + lat_min: float + lat_max: float + lon_min: float + lon_max: float + name: str + + def contains(self, lat: np.ndarray, lon: np.ndarray) -> np.ndarray: + return (lat >= self.lat_min) & (lat <= self.lat_max) & (lon >= self.lon_min) & (lon <= self.lon_max) + + +def load_region_boxes( + config_path: str, + split_first_region_lat: Optional[float] = None, +) -> List[RegionBox]: + """Load region_boxes from a training config JSON. + + Expects: + "data_filtering_config": { + "region_boxes": [ + [lat_min, lat_max, lon_min, lon_max], + ... + ] + } + + If split_first_region_lat is set (e.g. -5), the first box is split into two by latitude: + region_south: lat_min <= lat < split_first_region_lat + region_north: split_first_region_lat <= lat <= lat_max + so the default Amazon site (lat ~ -17) uses a fit from southern Amazon only. + """ + with open(config_path, "r", encoding="utf-8") as f: + cfg = json.load(f) + + dfc = cfg.get("data_filtering_config", {}) + raw_boxes = dfc.get("region_boxes") + if not raw_boxes: + raise ValueError(f"No data_filtering_config.region_boxes found in {config_path}") + + # Name the first two boxes to match docs/FIVE_P_TWO_REGION_GROUND_TRUTH_ANALYSIS.md + names = ["amazon", "africa"] + boxes: List[RegionBox] = [] + for i, box in enumerate(raw_boxes): + if len(box) != 4: + raise ValueError(f"Region box at index {i} must have 4 entries [lat_min, lat_max, lon_min, lon_max], got: {box}") + lat_min, lat_max, lon_min, lon_max = box + name = names[i] if i < len(names) else f"region_{i}" + + if i == 0 and split_first_region_lat is not None: + # Split first region (e.g. Amazon) into south and north by latitude + split = split_first_region_lat + if not (lat_min < split < lat_max): + raise ValueError( + f"split_first_region_lat {split} must be strictly between lat_min {lat_min} and lat_max {lat_max}" + ) + boxes.append( + RegionBox(lat_min=lat_min, lat_max=split, lon_min=lon_min, lon_max=lon_max, name=f"{name}_south") + ) + boxes.append( + RegionBox(lat_min=split, lat_max=lat_max, lon_min=lon_min, lon_max=lon_max, name=f"{name}_north") + ) + else: + boxes.append(RegionBox(lat_min=lat_min, lat_max=lat_max, lon_min=lon_min, lon_max=lon_max, name=name)) + + return boxes + + +def fit_bias_scale( + pred: np.ndarray, + gt: np.ndarray, + multiplicative_only: bool = False, + scale_factor: Optional[float] = None, + relative_error_weighted: bool = False, + relative_eps: float = 1e-12, + ref_site_lon: Optional[np.ndarray] = None, + ref_site_lat: Optional[np.ndarray] = None, + ref_sites: Optional[List[Tuple[float, float]]] = None, + ref_site_extra_weight: float = 500.0, + ref_site_atol: float = 1e-4, +) -> Tuple[float, float]: + """Fit GT ≈ a * pred + b via least squares (or a * pred only if multiplicative_only). + + Both inputs are 1D arrays of the same length. + If scale_factor is set (e.g. 1000), fit on (pred*scale, gt*scale) for stability; + returned (a,b) apply as corrected = (a * pred * scale + b) / scale. + If relative_error_weighted is True, minimize weighted squared error with + w_i = 1/(y_i+eps)^2 so that relative error is minimized (target <10% per layer). + If ref_sites is provided (list of (lon,lat)), cells at those sites get extra weight. + """ + if pred.shape != gt.shape: + raise ValueError(f"Shape mismatch for regression: pred {pred.shape}, gt {gt.shape}") + + mask = np.isfinite(pred) & np.isfinite(gt) + mask &= (pred != 0.0) | (gt != 0.0) + + x = pred[mask].astype(float) + y = gt[mask].astype(float) + if x.size < 2: + return 1.0, 0.0 + + if scale_factor is not None and scale_factor != 1.0: + x = x * scale_factor + y = y * scale_factor + + if multiplicative_only: + # Regression through origin: y = a * x => a = (x'y) / (x'x). Keeps corrected values non-negative. + xx = np.dot(x, x) + if xx <= 0: + a = 1.0 + else: + a = float(np.dot(x, y) / xx) + a = max(a, 1e-6) # avoid negative or zero scale + return a, 0.0 + + # Linear regression y = a * x + b + X = np.vstack([x, np.ones_like(x)]).T # (n, 2) + if relative_error_weighted: + # Minimize sum w_i * (y_i - a*x_i - b)^2 with w_i = 1/(y_i+eps)^2 + w = 1.0 / (y.astype(float) + relative_eps) ** 2 + # Boost weight for reference-site cells so fit targets <10% error there + if ( + ref_sites + and ref_site_lon is not None + and ref_site_lat is not None + and ref_site_lon.shape == ref_site_lat.shape + and ref_site_lon.size == w.size + ): + for rlon, rlat in ref_sites: + at_site = np.isclose(ref_site_lon, rlon, atol=ref_site_atol) & np.isclose( + ref_site_lat, rlat, atol=ref_site_atol + ) + w[at_site] *= 1.0 + ref_site_extra_weight + XtWX = X.T @ (w[:, np.newaxis] * X) + XtWy = X.T @ (w * y) + try: + coeffs = np.linalg.solve(XtWX, XtWy) + except np.linalg.LinAlgError: + coeffs, _, _, _ = np.linalg.lstsq(X, y, rcond=None) + a, b = coeffs + else: + coeffs, _, _, _ = np.linalg.lstsq(X, y, rcond=None) + a, b = coeffs + return float(a), float(b) + + +def compute_and_apply_corrections( + run_dir: str, + region_boxes: List[RegionBox], + variables: List[str], + output_subdir: str, + multiplicative_only_vars: Optional[Set[str]] = None, + scale_factor_vars: Optional[Dict[str, float]] = None, + relative_error_weighted_vars: Optional[Set[str]] = None, + reference_sites: Optional[List[Tuple[float, float]]] = None, +) -> Dict[str, Dict[str, Dict[str, Dict[str, float]]]]: + """Compute per-variable, per-layer, per-region bias/scale corrections and apply them. + + Returns a nested dict: + params[var_name][region_name][layer_key] = {"a": ..., "b": ...} or {"a", "b", "scale"} for scaled vars. + """ + if multiplicative_only_vars is None: + multiplicative_only_vars = MULTIPLICATIVE_ONLY_VARS + if scale_factor_vars is None: + scale_factor_vars = SCALE_FACTOR_VARS + if relative_error_weighted_vars is None: + relative_error_weighted_vars = RELATIVE_ERROR_WEIGHTED_VARS + if reference_sites is None: + reference_sites = REFERENCE_SITES + predictions_root = os.path.join( + run_dir, + "cnp_inference_entire_dataset", + "cnp_predictions", + ) + gt_dir = os.path.join(predictions_root, "soil_2d_ground_truth") + pred_dir = os.path.join(predictions_root, "soil_2d_predictions") + + output_dir = os.path.join(predictions_root, output_subdir) + os.makedirs(output_dir, exist_ok=True) + + params: Dict[str, Dict[str, Dict[str, Dict[str, float]]]] = {} + + for var in variables: + var_key = f"Y_{var}_col1_layer" + gt_path = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + pred_path = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + + # Fallback to the *_vr naming convention used in this repo + if not os.path.exists(gt_path): + gt_path = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + if not os.path.exists(pred_path): + pred_path = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + + if not os.path.exists(gt_path): + gt_path = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + if not os.path.exists(pred_path): + pred_path = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + + # In this codebase, files are named ground_truth_Y__vr.csv and predictions_Y__vr.csv + if not os.path.exists(gt_path): + gt_path = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + if not os.path.exists(pred_path): + pred_path = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + + # For clarity, explicitly handle the _vr suffix + if not os.path.exists(gt_path): + gt_path = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + if not os.path.exists(pred_path): + pred_path = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + + # Final explicit names for this repo + if not os.path.exists(gt_path): + gt_path = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + if not os.path.exists(pred_path): + pred_path = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + + # In practice, for labilep_vr, occlp_vr, etc., we expect: + # ground_truth_Y_labilep_vr.csv + # predictions_Y_labilep_vr.csv + if not os.path.exists(gt_path): + gt_path = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + if not os.path.exists(pred_path): + pred_path = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + + # And the explicit natveg_improved naming we have already seen: + explicit_gt = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + explicit_pred = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + if os.path.exists(explicit_gt): + gt_path = explicit_gt + if os.path.exists(explicit_pred): + pred_path = explicit_pred + + # As a final guard, ensure the expected _vr suffixed paths exist + if not os.path.exists(gt_path): + alt_gt = os.path.join(gt_dir, f"ground_truth_Y_{var}.csv") + if os.path.exists(alt_gt): + gt_path = alt_gt + if not os.path.exists(pred_path): + alt_pred = os.path.join(pred_dir, f"predictions_Y_{var}.csv") + if os.path.exists(alt_pred): + pred_path = alt_pred + + if not os.path.exists(gt_path) or not os.path.exists(pred_path): + raise FileNotFoundError( + f"Could not find ground truth or prediction CSV for variable '{var}'. " + f"Tried paths like {gt_path} and {pred_path}" + ) + + print(f"Loading GT from {gt_path}") + print(f"Loading predictions from {pred_path}") + + gt_df = pd.read_csv(gt_path) + pred_df = pd.read_csv(pred_path) + + if gt_df.shape != pred_df.shape: + raise ValueError( + f"Shape mismatch for GT vs predictions for {var}: " + f"GT {gt_df.shape}, PRED {pred_df.shape}" + ) + + for col in ["Longitude", "Latitude"]: + if col not in gt_df.columns or col not in pred_df.columns: + raise ValueError( + f"Expected '{col}' column in both GT and prediction CSVs for {var}" + ) + + # Sanity check alignment + if not np.allclose(gt_df["Longitude"].values, pred_df["Longitude"].values) or not np.allclose( + gt_df["Latitude"].values, pred_df["Latitude"].values + ): + raise ValueError(f"Longitude/Latitude mismatch between GT and predictions for {var}") + + lon = gt_df["Longitude"].values + lat = gt_df["Latitude"].values + + # Prepare per-layer coefficients for each region + var_params: Dict[str, Dict[str, Dict[str, float]]] = {} + num_layers = 0 + for c in gt_df.columns: + if c.startswith(var_key): + num_layers += 1 + + if num_layers == 0: + raise ValueError(f"No layer columns found for variable '{var}' (prefix '{var_key}')") + + for region in region_boxes: + region_mask = region.contains(lat=lat, lon=lon) + + if not region_mask.any(): + print(f"Warning: no cells found in region '{region.name}' for variable '{var}'") + continue + + print(f"Fitting corrections for {var} in region '{region.name}' using {region_mask.sum()} cells") + region_layer_params: Dict[str, Dict[str, float]] = {} + + for layer_idx in range(1, num_layers + 1): + col_name = f"{var_key}{layer_idx}" + if col_name not in gt_df.columns or col_name not in pred_df.columns: + raise ValueError( + f"Expected column '{col_name}' in GT and prediction CSVs for {var}" + ) + + gt_vals = gt_df.loc[region_mask, col_name].values.astype(float) + pred_vals = pred_df.loc[region_mask, col_name].values.astype(float) + + mult_only = var in multiplicative_only_vars + scale = scale_factor_vars.get(var) if scale_factor_vars else None + if scale is not None and scale != 1.0: + mult_only = False # use linear fit on scaled data + rel_err = var in relative_error_weighted_vars + ref_lon = lon[region_mask] if rel_err and reference_sites else None + ref_lat = lat[region_mask] if rel_err and reference_sites else None + a, b = fit_bias_scale( + pred_vals, + gt_vals, + multiplicative_only=mult_only, + scale_factor=scale, + relative_error_weighted=rel_err, + relative_eps=RELATIVE_ERROR_EPS, + ref_site_lon=ref_lon, + ref_site_lat=ref_lat, + ref_sites=reference_sites if rel_err else None, + ref_site_extra_weight=REFERENCE_SITE_EXTRA_WEIGHT, + ref_site_atol=REFERENCE_SITE_ATOL, + ) + # For relative-error vars: if ref site still has >10% error, nudge (a,b) so ref site is exact (minimal change) + if rel_err and reference_sites and ref_lon is not None and ref_lat is not None: + for rlon, rlat in reference_sites: + at_site = np.isclose(ref_lon, rlon, atol=REFERENCE_SITE_ATOL) & np.isclose( + ref_lat, rlat, atol=REFERENCE_SITE_ATOL + ) + if not np.any(at_site): + continue + idx = np.where(at_site)[0][0] + pred_ref = pred_vals[idx] + gt_ref = gt_vals[idx] + if scale and scale != 1.0: + pred_s = pred_ref * scale + corrected_ref = (a * pred_s + b) / scale + else: + corrected_ref = a * pred_ref + b + rel_pct = 100 * abs(corrected_ref - gt_ref) / (gt_ref + 1e-12) + if corrected_ref >= 0 and rel_pct <= 10: + break + # Project (a,b) onto constraint a'*pred_ref + b' = gt_ref (in original space) + # Minimize (a'-a)^2+(b'-b)^2 s.t. a'*pred_ref + b' = gt_ref. Solution: + # lambda = 2*(a*pred_ref + b - gt_ref) / (pred_ref**2 + 1), a' = a - lambda*pred_ref/2, b' = b - lambda/2 + if scale and scale != 1.0: + pred_orig = pred_ref + # In scaled space: a*pred_s + b = gt_s would give (a*pred_s+b)/scale = gt_orig so a*pred_s+b = gt_orig*scale + # Constraint in scaled space: a'*pred_s + b' = gt_ref*scale + pred_s = pred_orig * scale + tgt = gt_ref * scale + lam = 2.0 * (a * pred_s + b - tgt) / (pred_s ** 2 + 1.0) + a = a - lam * pred_s / 2.0 + b = b - lam / 2.0 + else: + lam = 2.0 * (a * pred_ref + b - gt_ref) / (pred_ref ** 2 + 1.0) + a = float(a - lam * pred_ref / 2.0) + b = float(b - lam / 2.0) + break + layer_key = f"layer_{layer_idx}" + coeff: Dict[str, float] = {"a": a, "b": b} + if scale is not None and scale != 1.0: + coeff["scale"] = scale + region_layer_params[layer_key] = coeff + + var_params[region.name] = region_layer_params + + params[var] = var_params + + # Apply corrections to a copy of the prediction DataFrame + corrected_df = pred_df.copy() + for region in region_boxes: + region_mask = region.contains(lat=lat, lon=lon) + if not region_mask.any(): + continue + + region_param_dict = params[var].get(region.name) + if not region_param_dict: + # Nothing to apply for this region/variable + continue + + for layer_idx in range(1, num_layers + 1): + col_name = f"{var_key}{layer_idx}" + layer_key = f"layer_{layer_idx}" + coeffs = region_param_dict[layer_key] + a = coeffs["a"] + b = coeffs["b"] + scale = coeffs.get("scale") + + pred_vals = corrected_df.loc[region_mask, col_name].astype(float).values + if scale is not None and scale != 1.0: + raw = (a * pred_vals * scale + b) / scale + else: + raw = a * pred_vals + b + # Concentrations must be non-negative + corrected_df.loc[region_mask, col_name] = np.maximum(raw, 0.0) + + # For small-scaled vars (e.g. solutionp_vr), keep 5 digits after decimal when writing + decimal_places = ( + DECIMAL_PLACES_SCALED_VARS + if (scale_factor_vars and var in scale_factor_vars) + else None + ) + if decimal_places is not None: + layer_col_names = [c for c in corrected_df.columns if c.startswith(var_key)] + for c in layer_col_names: + corrected_df[c] = corrected_df[c].round(decimal_places) + + out_path = os.path.join(output_dir, f"predictions_Y_{var}_bias_corrected.csv") + print(f"Writing bias/scale–corrected predictions for {var} to {out_path}") + if decimal_places is not None: + fmt = f"%.{decimal_places}f" + corrected_df.to_csv(out_path, index=False, float_format=fmt) + else: + corrected_df.to_csv(out_path, index=False) + + return params + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Apply bias/scale correction on top of soil 2D 5P predictions " + "for Amazon + Africa, using GT vs prediction CSVs from a CNP run." + ) + ) + parser.add_argument( + "--run-dir", + required=True, + help=( + "Path to a cnp_results run directory, e.g. " + "cnp_results/run_20260228_214757_natveg_improved" + ), + ) + parser.add_argument( + "--region-config-json", + default="config/training_config_two_region_five_p.json", + help=( + "Training config JSON that defines data_filtering_config.region_boxes " + "for Amazon + Africa." + ), + ) + parser.add_argument( + "--variables", + default=",".join(P_VARIABLES_DEFAULT), + help=( + "Comma-separated list of soil2d P variables to correct. " + f"Default: {','.join(P_VARIABLES_DEFAULT)}" + ), + ) + parser.add_argument( + "--output-subdir", + default="soil_2d_predictions_5P_bias_corrected", + help=( + "Subdirectory (under cnp_inference_entire_dataset/cnp_predictions) where " + "corrected prediction CSVs will be written." + ), + ) + parser.add_argument( + "--multiplicative-only-vars", + default=",".join(sorted(MULTIPLICATIVE_ONLY_VARS)), + help=( + "Comma-separated variables that use multiplicative-only correction (Y = a*pred). " + "Default: none (solutionp_vr uses scale-factor fit instead)." + ), + ) + parser.add_argument( + "--scale-factor-vars", + default=",".join(f"{k}:{v}" for k, v in sorted(SCALE_FACTOR_VARS.items())), + help=( + "Comma-separated list of var:scale (e.g. solutionp_vr:1000) for small-value vars: " + "fit on pred*scale vs gt*scale, apply (a*pred*scale+b)/scale; output rounded to 5 decimals." + ), + ) + parser.add_argument( + "--split-first-region-lat", + default="-5.0", + metavar="LAT", + help=( + "Split the first region (Amazon) into _south (lat_min to LAT) and _north (LAT to lat_max) " + "so the default Amazon site (lat ~ -17) uses a southern fit. Float, or 'none' to disable. Default: -5.0" + ), + ) + + args = parser.parse_args() + + run_dir = os.path.abspath(args.run_dir) + region_config_json = os.path.abspath(args.region_config_json) + variables = [v.strip() for v in args.variables.split(",") if v.strip()] + multiplicative_only_vars = { + v.strip() for v in args.multiplicative_only_vars.split(",") if v.strip() + } + scale_factor_vars: Dict[str, float] = {} + for part in args.scale_factor_vars.split(","): + part = part.strip() + if not part: + continue + if ":" in part: + var_name, scale_str = part.split(":", 1) + var_name, scale_str = var_name.strip(), scale_str.strip() + try: + scale_factor_vars[var_name] = float(scale_str) + except ValueError: + raise SystemExit(f"Invalid scale in --scale-factor-vars: {part}") + else: + scale_factor_vars[part] = SCALE_FACTOR_VARS.get(part, 1000.0) + if not scale_factor_vars and SCALE_FACTOR_VARS: + scale_factor_vars = dict(SCALE_FACTOR_VARS) + + print(f"Run directory: {run_dir}") + print(f"Region config JSON: {region_config_json}") + print(f"Variables: {variables}") + print(f"Output subdir (relative to cnp_predictions): {args.output_subdir}") + + if not os.path.isdir(run_dir): + raise SystemExit(f"run-dir does not exist or is not a directory: {run_dir}") + if not os.path.isfile(region_config_json): + raise SystemExit(f"region-config-json not found: {region_config_json}") + + split_lat: Optional[float] = None + if args.split_first_region_lat.strip().lower() not in ("none", "no", ""): + try: + split_lat = float(args.split_first_region_lat) + except ValueError: + raise SystemExit( + f"Invalid --split-first-region-lat: {args.split_first_region_lat}. Use a number or 'none'." + ) + region_boxes = load_region_boxes(region_config_json, split_first_region_lat=split_lat) + print("Loaded region boxes:") + for rb in region_boxes: + print( + f" {rb.name}: lat [{rb.lat_min}, {rb.lat_max}], " + f"lon [{rb.lon_min}, {rb.lon_max}]" + ) + + params = compute_and_apply_corrections( + run_dir=run_dir, + region_boxes=region_boxes, + variables=variables, + output_subdir=args.output_subdir, + multiplicative_only_vars=multiplicative_only_vars, + scale_factor_vars=scale_factor_vars, + ) + + # Save parameters to analysis/ for documentation and reuse + analysis_dir = os.path.join(run_dir, "analysis") + os.makedirs(analysis_dir, exist_ok=True) + params_path = os.path.join(analysis_dir, "bias_scale_params_5P_two_regions.json") + with open(params_path, "w", encoding="utf-8") as f: + json.dump(params, f, indent=2) + print(f"Saved bias/scale parameters to {params_path}") + + +if __name__ == "__main__": + main() + diff --git a/scripts/cnp_result_validationplot.py b/scripts/cnp_result_validationplot.py index 4d638af..f4cab91 100644 --- a/scripts/cnp_result_validationplot.py +++ b/scripts/cnp_result_validationplot.py @@ -11,6 +11,11 @@ NUM_LAYERS = 10 # All layers (soil/2D) = exactly 10 layers NUM_PFTS = 16 # All PFTs (1D) = pft1 to pft16 only +# Minimum gridcells after natveg filter to apply the filter; below this, we keep unfiltered data to avoid unstable metrics. +MIN_GRIDCELLS_FOR_NATVEG_FILTER = 10 +# Minimum valid (non-NaN) gt/pred pairs to generate a plot for a PFT or layer. +MIN_VALID_PAIRS_FOR_PLOT = 3 + # Plot output subfolders (under plots/ or top_bad_plots/): # aggregate_all - one scatter per variable combining all PFTs (1D) or all layers (2D) # aggregate_bad - one scatter per variable combining only bad/selected PFTs or layers (when using top-bad report) @@ -19,13 +24,16 @@ SUBDIR_BADLAYER = "aggregate_bad" SUBDIR_INDIVIDUAL = "by_pft_layer" -def plot_gt_vs_pred(gt, pred, title, save_path): +def plot_gt_vs_pred(gt, pred, title, save_path, subtitle=None): plt.figure(figsize=(6,6)) plt.scatter(gt, pred, alpha=0.5) plt.plot([gt.min(), gt.max()], [gt.min(), gt.max()], 'r--') plt.xlabel('Ground Truth') plt.ylabel('Prediction') - plt.title(title) + if subtitle: + plt.title(f"{title}\n{subtitle}", fontsize=10) + else: + plt.title(title) plt.tight_layout() # Ensure the directory exists os.makedirs(os.path.dirname(save_path), exist_ok=True) @@ -91,6 +99,52 @@ def _parse_top_bad_report(report_path): return {} return selection +def _load_gridcell_metadata(results_dir): + """Load gridcell-level PCT_NATVEG and PCT_NAT_PFT_* from test_static_inverse.csv if present. + Returns a DataFrame with same row order as GT/pred CSVs, or None if not available. + Used to exclude gridcells with no natural veg (PCT_NATVEG=0) or 100% PFT0 (PCT_NAT_PFT_0=100). + """ + # Prefer test_static_inverse.csv written by run_inference_all (same row order as predictions) + candidates = [ + os.path.join(results_dir, 'cnp_predictions', 'test_static_inverse.csv'), + os.path.join(results_dir, 'cnp_predictions', 'gridcell_metadata.csv'), + ] + for path in candidates: + if not os.path.exists(path): + continue + try: + df = pd.read_csv(path) + # Need PCT_NATVEG and PCT_NAT_PFT_0 for exclusion; PCT_NAT_PFT_1..16 for low-coverage flag + pct_natveg = None + for c in ['PCT_NATVEG', 'pct_natveg']: + if c in df.columns: + pct_natveg = df[c].values + break + pct_pft0 = None + for c in ['PCT_NAT_PFT_0', 'pct_nat_pft_0']: + if c in df.columns: + pct_pft0 = df[c].values + break + if pct_natveg is None or pct_pft0 is None: + continue + # Build include mask: include where (PCT_NATVEG > 0) and (PCT_NAT_PFT_0 < 100) + include = (np.asarray(pct_natveg, dtype=float) > 0) & (np.asarray(pct_pft0, dtype=float) < 100) + # PCT_NAT_PFT_1..16 for pft_pct_low + pct_pft_cols = {} + for i in range(1, 17): + c = f'PCT_NAT_PFT_{i}' + if c in df.columns: + pct_pft_cols[i] = df[c].values.astype(float) + return { + 'include_mask': include, + 'pct_pft': pct_pft_cols, + 'n_rows': len(df), + } + except Exception as e: + print(f"Warning: Could not load gridcell metadata from {path}: {e}") + return None + + def _parse_worst_vars_report(report_path): """Parse quality_summary_report.txt to extract variables from the '## Variables with Worst Predictions' section. @@ -124,7 +178,7 @@ def _parse_worst_vars_report(report_path): return {} return selection -def main_with_flag(results_dir, plot_scatter, plot_loss, top_bad_only=False, top_bad_report=None, plots_dir_override=None, worst_only=False): +def main_with_flag(results_dir, plot_scatter, plot_loss, top_bad_only=False, top_bad_report=None, plots_dir_override=None, worst_only=False, use_natveg_filter=True): # Create plots subdirectory and subfolders for organization plots_dir = plots_dir_override or os.path.join(results_dir, "plots") os.makedirs(plots_dir, exist_ok=True) @@ -151,6 +205,17 @@ def main_with_flag(results_dir, plot_scatter, plot_loss, top_bad_only=False, top if not selection: print("No selections parsed from report; proceeding without restriction.") selection = None + # Always include npool and ppool in selection when in restricted mode so their scatter plots + # (gt vs pred) are still generated even when they are excluded from the worst list (e.g. per-PFT validation). + if selection is not None: + for _v in ('npool', 'ppool'): + if _v not in selection: + selection[_v] = {'pfts': set(), 'layers': set()} + + # Load gridcell metadata for PFT/2D filtering only when explicitly requested (opt-in; default off to avoid changing metrics) + gridcell_metadata = _load_gridcell_metadata(results_dir) if use_natveg_filter else None + if gridcell_metadata is not None: + print("Using gridcell metadata for PFT/2D: excluding no-natveg and 100% PFT0 gridcells") # Check for new directory structure first pft_gt_dir = os.path.join(results_dir, 'cnp_predictions', 'pft_1d_ground_truth') @@ -159,7 +224,7 @@ def main_with_flag(results_dir, plot_scatter, plot_loss, top_bad_only=False, top # Handle 1D data with new structure if os.path.exists(pft_gt_dir) and os.path.exists(pft_pred_dir): print("Using new 1D directory structure") - analyze_1d_new_structure(results_dir, '1D', plots_dir, stats_data, plot_scatter, selection) + analyze_1d_new_structure(results_dir, '1D', plots_dir, stats_data, plot_scatter, selection, gridcell_metadata) else: # Fall back to old single-file format print("Using legacy 1D single-file format") @@ -184,7 +249,7 @@ def main_with_flag(results_dir, plot_scatter, plot_loss, top_bad_only=False, top soil_pred_dir = os.path.join(results_dir, 'cnp_predictions', 'soil_2d_predictions') if os.path.exists(soil_gt_dir) and os.path.exists(soil_pred_dir): print("Using new 2D directory structure") - analyze_2d_new_structure(results_dir, '2D', plots_dir, stats_data, plot_scatter, selection) + analyze_2d_new_structure(results_dir, '2D', plots_dir, stats_data, plot_scatter, selection, gridcell_metadata) else: # Fall back to old single-file format print("Using legacy 2D single-file format") @@ -318,8 +383,11 @@ def analyze_pair(gt_path, pred_path, label, out_dir, stats_data, per_column=Fals }) return {'rmse': rmse, 'mae': mae, 'r2': r2} -def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatter=True, selection=None): - """Analyze 1D data using the new directory structure with individual variable files""" +def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatter=True, selection=None, gridcell_metadata=None): + """Analyze 1D data using the new directory structure with individual variable files. + If gridcell_metadata is provided, exclude gridcells where PCT_NATVEG=0 or PCT_NAT_PFT_0=100 + and set pft_pct_low for PFTs with mean coverage < 2%. + """ gt_dir = os.path.join(results_dir, 'cnp_predictions', 'pft_1d_ground_truth') pred_dir = os.path.join(results_dir, 'cnp_predictions', 'pft_1d_predictions') @@ -362,6 +430,11 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt print(f"Shape mismatch for {var_name}: GT {gt_data.shape} vs Pred {pred_data.shape}") continue + # Build row mask for PFT/2D: exclude PCT_NATVEG=0 and PCT_NAT_PFT_0=100 + include_mask = None + if gridcell_metadata is not None and gridcell_metadata['n_rows'] == len(gt_data): + include_mask = gridcell_metadata['include_mask'] + # When using selection (top-bad only): skip this variable for per-PFT plots if not selected, # but we still generate the AllPFTs plot for every variable so AllPFTs are never "missing". in_selection = selection is None or var_name in selection @@ -391,13 +464,22 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt gt_col = gt_data[col_name].values pred_col = pred_data[col_name].values + # Apply natveg filter: use only included gridcells (PCT_NATVEG>0 and PCT_NAT_PFT_0<100) + if include_mask is not None: + n_included = int(np.sum(include_mask)) + if n_included < MIN_GRIDCELLS_FOR_NATVEG_FILTER: + pass # keep gt_col, pred_col unfiltered to avoid unstable metrics + else: + gt_col = gt_col[include_mask] + pred_col = pred_col[include_mask] + # Skip if all values are NaN if np.all(np.isnan(gt_col)) or np.all(np.isnan(pred_col)): continue # Remove NaN pairs valid_mask = ~(np.isnan(gt_col) | np.isnan(pred_col)) - if np.sum(valid_mask) < 3: # Relax threshold to allow small samples + if np.sum(valid_mask) < MIN_VALID_PAIRS_FOR_PLOT: continue gt_valid = gt_col[valid_mask] @@ -412,6 +494,15 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt mae = mean_absolute_error(gt_valid, pred_valid) r2 = r2_score(gt_valid, pred_valid) + # PFT coverage < 2% flag for scatter coloring (PFT 1 = pft_idx 0 -> PCT_NAT_PFT_1) + pft_pct_low = False + if gridcell_metadata is not None and include_mask is not None: + pft_num = pft_idx + 1 # 1-based + if pft_num in gridcell_metadata.get('pct_pft', {}): + pct_vals = gridcell_metadata['pct_pft'][pft_num][include_mask] + if len(pct_vals) > 0: + pft_pct_low = float(np.nanmean(pct_vals)) < 2.0 + print(f" {col_name}: RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}") print(f" GT - min: {gt_stats['min']:.6g}, max: {gt_stats['max']:.6g}, sum: {gt_stats['sum']:.6g}") print(f" Pred - min: {pred_stats['min']:.6g}, max: {pred_stats['max']:.6g}, sum: {pred_stats['sum']:.6g}") @@ -425,7 +516,7 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt ) # Collect stats - stats_data.append({ + row = { 'type': '1D', 'variable': var_name, 'pft': col_name, @@ -438,7 +529,10 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt 'pred_min': pred_stats['min'], 'pred_max': pred_stats['max'], 'pred_sum': pred_stats['sum'] - }) + } + if gridcell_metadata is not None: + row['pft_pct_low'] = pft_pct_low + stats_data.append(row) # When not in top-bad mode: provide overall scatter across all PFTs (aggregate_all). # When selection is set (top_bad_plots), skip so we only get aggregate_bad and by_pft_layer. @@ -590,8 +684,10 @@ def analyze_1d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot 'pred_sum': pred_stats['sum'] }) -def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatter=True, selection=None): - """Analyze 2D data using the new directory structure with individual variable files""" +def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatter=True, selection=None, gridcell_metadata=None): + """Analyze 2D data using the new directory structure with individual variable files. + If gridcell_metadata is provided, exclude gridcells where PCT_NATVEG=0 or PCT_NAT_PFT_0=100. + """ gt_dir = os.path.join(results_dir, 'cnp_predictions', 'soil_2d_ground_truth') pred_dir = os.path.join(results_dir, 'cnp_predictions', 'soil_2d_predictions') @@ -637,6 +733,16 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt if gt_data.shape != pred_data.shape: print(f"Shape mismatch for {var_name}: GT {gt_data.shape} vs Pred {pred_data.shape}") continue + + # Build row mask: exclude PCT_NATVEG=0 and PCT_NAT_PFT_0=100 + include_mask = None + if gridcell_metadata is not None and gridcell_metadata['n_rows'] == len(gt_data): + include_mask = gridcell_metadata['include_mask'] + n_incl = int(np.sum(include_mask)) + n_excl = len(include_mask) - n_incl + print(f" Natveg filter for {var_name}: {n_incl} included, {n_excl} excluded (of {len(gt_data)} rows)") + elif gridcell_metadata is not None and gridcell_metadata['n_rows'] != len(gt_data): + print(f" Warning: gridcell_metadata has {gridcell_metadata['n_rows']} rows but GT has {len(gt_data)}; skipping natveg filter for {var_name}") # 2D data: All layers = NUM_LAYERS (10) only. total_columns = gt_data.shape[1] @@ -672,13 +778,23 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt gt_col = gt_data.iloc[:, col_idx].values pred_col = pred_data.iloc[:, col_idx].values + # Apply natveg filter: use only included gridcells (PCT_NATVEG>0 and PCT_NAT_PFT_0<100) + applied_natveg_filter = False + if include_mask is not None: + n_included = int(np.sum(include_mask)) + if n_included >= MIN_GRIDCELLS_FOR_NATVEG_FILTER: + gt_col = gt_col[include_mask] + pred_col = pred_col[include_mask] + applied_natveg_filter = True + # else: too few after filter; keep unfiltered to avoid unstable metrics + # Skip if all values are NaN if np.all(np.isnan(gt_col)) or np.all(np.isnan(pred_col)): continue # Remove NaN pairs valid_mask = ~(np.isnan(gt_col) | np.isnan(pred_col)) - if np.sum(valid_mask) < 3: # Relax threshold to allow small samples + if np.sum(valid_mask) < MIN_VALID_PAIRS_FOR_PLOT: continue gt_valid = gt_col[valid_mask] @@ -699,10 +815,12 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt # Conditionally plot if plot_scatter: + subtitle = f"{len(gt_valid)} gridcells (natveg filter applied)" if applied_natveg_filter else None plot_gt_vs_pred( - gt_valid, pred_valid, - f"{label} {var_name} Layer{layer_idx+1} GT vs Pred", - os.path.join(out_dir, SUBDIR_INDIVIDUAL, f"{label}_{var_name}_Layer{layer_idx+1}_gt_vs_pred.png") + gt_valid, pred_valid, + f"{label} {var_name} Layer{layer_idx+1} GT vs Pred", + os.path.join(out_dir, SUBDIR_INDIVIDUAL, f"{label}_{var_name}_Layer{layer_idx+1}_gt_vs_pred.png"), + subtitle=subtitle, ) # Collect stats @@ -879,8 +997,11 @@ def plot_train_val_accuracy(loss_csv, out_dir): parser.add_argument('--top-bad-only', action='store_true', help='Plot only variables listed in the quality summary top-bad section') parser.add_argument('--worst-only', action='store_true', help='Plot only variables listed under \"Variables with Worst Predictions\"') parser.add_argument('--top-bad-report', type=str, default=None, help='Path to quality_summary_report.txt (defaults to results_dir/analysis/quality_summary_report.txt)') + parser.add_argument('--plots-dir', type=str, default=None, help='Output directory for plots (default: results_dir/plots). Use analysis/top_bad_plots to match quality-report pipeline.') + parser.add_argument('--no-natveg-filter', action='store_false', dest='use_natveg_filter', + help='Do not apply PFT/2D filter (plot all gridcells). Default: natveg filter is ON (exclude PCT_NATVEG=0 or PCT_NAT_PFT_0=100).') - parser.set_defaults(plot_scatter=True, plot_loss=True) + parser.set_defaults(plot_scatter=True, plot_loss=True, use_natveg_filter=True) args = parser.parse_args() # If stats-only requested, force-disable all plotting @@ -891,4 +1012,10 @@ def plot_train_val_accuracy(loss_csv, out_dir): if len(sys.argv) < 2: print("Using current directory as results directory") - main_with_flag(args.results_dir, args.plot_scatter, args.plot_loss, args.top_bad_only, args.top_bad_report, worst_only=getattr(args, 'worst_only', False)) \ No newline at end of file + plots_dir_override = None + if getattr(args, 'plots_dir', None): + plots_dir_override = os.path.join(args.results_dir, args.plots_dir) if not os.path.isabs(args.plots_dir) else args.plots_dir + main_with_flag(args.results_dir, args.plot_scatter, args.plot_loss, args.top_bad_only, args.top_bad_report, + plots_dir_override=plots_dir_override, + worst_only=getattr(args, 'worst_only', False), + use_natveg_filter=getattr(args, 'use_natveg_filter', True)) \ No newline at end of file diff --git a/scripts/compare_5p_amazon_site_scatter.py b/scripts/compare_5p_amazon_site_scatter.py new file mode 100644 index 0000000..c69bd77 --- /dev/null +++ b/scripts/compare_5p_amazon_site_scatter.py @@ -0,0 +1,582 @@ +import argparse +import os +import warnings +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr + +warnings.filterwarnings("ignore", message=".*multiple fill values.*", category=UserWarning, module="xarray") + +P_VARIABLES: List[str] = [ + "labilep_vr", + "occlp_vr", + "solutionp_vr", + "secondp_vr", + "primp_vr", +] + +# Default Amazon site from docs/AMAZON_5P_QUALITY_SUMMARY.md +DEFAULT_AMAZON_LON = 303.75 +DEFAULT_AMAZON_LAT = -17.434553 + + +@dataclass +class ModelSpec: + name: str + run_dir: str + predictions_rel_dir: str + filename_template: str + + +def _find_row_for_site( + df: pd.DataFrame, lon: float, lat: float, atol: float = 1e-4 +) -> int: + """Return integer row index in df matching the given lon/lat (within atol).""" + if "Longitude" not in df.columns or "Latitude" not in df.columns: + raise ValueError("DataFrame must contain 'Longitude' and 'Latitude' columns.") + + lon_vals = df["Longitude"].values + lat_vals = df["Latitude"].values + mask = np.isclose(lon_vals, lon, atol=atol) & np.isclose(lat_vals, lat, atol=atol) + + idx = np.where(mask)[0] + if idx.size == 0: + raise ValueError( + f"No row found at lon={lon}, lat={lat} (within atol={atol}). " + "Check that the site exists in these CSVs." + ) + if idx.size > 1: + raise ValueError( + f"Multiple rows found at lon={lon}, lat={lat}; " + "expected exactly one grid cell." + ) + return int(idx[0]) + + +def _extract_profile( + df: pd.DataFrame, + var: str, + row_idx: int, +) -> Tuple[np.ndarray, List[str]]: + """Extract the vertical profile (all layers) for one variable at row_idx.""" + prefix = f"Y_{var}_col1_layer" + # Sort by numeric layer index so that layer1, layer2, ..., layer10 are in + # the correct physical order instead of the lexicographic order + # (layer1, layer10, layer2, ...). + layer_cols = sorted( + [c for c in df.columns if c.startswith(prefix)], + key=lambda c: int(c.split("layer")[-1]), + ) + if not layer_cols: + raise ValueError(f"No columns starting with '{prefix}' found.") + + values = df.loc[row_idx, layer_cols].astype(float).values + return values, layer_cols + + +def load_profiles_for_site( + gt_run_dir: str, + models: List[ModelSpec], + variables: List[str], + lon: float, + lat: float, +) -> Tuple[Dict[str, np.ndarray], Dict[str, Dict[str, np.ndarray]], List[str]]: + """Load GT and model profiles for all variables at a single site.""" + gt_root = os.path.join( + gt_run_dir, + "cnp_inference_entire_dataset", + "cnp_predictions", + "soil_2d_ground_truth", + ) + + gt_profiles: Dict[str, np.ndarray] = {} + model_profiles: Dict[str, Dict[str, np.ndarray]] = {m.name: {} for m in models} + layer_labels: List[str] = [] + + for var in variables: + gt_path = os.path.join(gt_root, f"ground_truth_Y_{var}.csv") + if not os.path.exists(gt_path): + raise FileNotFoundError(f"Ground truth file not found: {gt_path}") + + gt_df = pd.read_csv(gt_path) + gt_row_idx = _find_row_for_site(gt_df, lon=lon, lat=lat) + gt_vals, layer_cols = _extract_profile(gt_df, var=var, row_idx=gt_row_idx) + + gt_profiles[var] = gt_vals + if not layer_labels: + layer_labels = layer_cols + + for model in models: + pred_dir = os.path.join(model.run_dir, model.predictions_rel_dir) + pred_path = os.path.join( + pred_dir, model.filename_template.format(var=var) + ) + if not os.path.exists(pred_path): + # e.g. phase2 bias-corrected only has 4 P vars; skip missing vars + model_profiles[model.name][var] = np.full_like(gt_vals, np.nan) + continue + + pred_df = pd.read_csv(pred_path) + + # Find the matching site row independently in each prediction DataFrame. + # This avoids assumptions about global shape or ordering and supports + # region-only prediction CSVs (e.g., Amazon/Africa subsets). + pred_row_idx = _find_row_for_site(pred_df, lon=lon, lat=lat) + pred_vals, _ = _extract_profile(pred_df, var=var, row_idx=pred_row_idx) + model_profiles[model.name][var] = pred_vals + + return gt_profiles, model_profiles, layer_labels + + +def make_profile_line_plots( + output_dir: str, + gt_profiles: Dict[str, np.ndarray], + model_profiles: Dict[str, Dict[str, np.ndarray]], + layer_labels: List[str], + lon: float, + lat: float, +) -> None: + """Plot vertical profiles (one line per model + GT) across layers.""" + os.makedirs(output_dir, exist_ok=True) + + model_names = list(model_profiles.keys()) + + for var, gt_vals in gt_profiles.items(): + num_layers = len(gt_vals) + layers = np.arange(1, num_layers + 1) + + plt.figure(figsize=(6, 6)) + + # Ground truth profile + plt.plot( + layers, + gt_vals, + marker="o", + linestyle="-", + color="black", + label="ground_truth", + ) + + # Each model's profile (skip if all NaN, e.g. missing variable for that model) + for model_name in model_names: + preds = model_profiles[model_name][var] + if np.all(np.isnan(preds)): + continue + plt.plot( + layers, + preds, + marker="o", + linestyle="--", + label=model_name, + alpha=0.9, + ) + + plt.xlabel("Soil layer index") + plt.ylabel(f"{var} value") + plt.title( + f"{var} vertical profile at Amazon site\n" + f"(lon={lon}, lat={lat}, {num_layers} layers)" + ) + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + + out_path = os.path.join(output_dir, f"amazon_site_profile_{var}.png") + plt.savefig(out_path, dpi=200) + plt.close() + + +def write_profiles_csv( + output_dir: str, + gt_profiles: Dict[str, np.ndarray], + model_profiles: Dict[str, Dict[str, np.ndarray]], + layer_labels: List[str], + lon: float, + lat: float, +) -> None: + rows: List[Dict[str, object]] = [] + + for var, gt_vals in gt_profiles.items(): + num_layers = len(gt_vals) + for layer_idx in range(num_layers): + layer_name = layer_labels[layer_idx] if layer_idx < len(layer_labels) else f"layer_{layer_idx+1}" + row: Dict[str, object] = { + "variable": var, + "layer_index": layer_idx + 1, + "layer_col": layer_name, + "lon": lon, + "lat": lat, + "gt": float(gt_vals[layer_idx]), + } + for model_name, profiles in model_profiles.items(): + row[f"pred_{model_name}"] = float(profiles[var][layer_idx]) + rows.append(row) + + df = pd.DataFrame(rows) + out_csv = os.path.join(output_dir, "amazon_site_5p_profiles.csv") + df.to_csv(out_csv, index=False) + + +def _load_5p_from_restart(path: str) -> Dict[str, np.ndarray]: + """Load 5P variables from a restart NetCDF. Returns dict of (n_col, n_lev) arrays. + Handles (column, levgrnd) or (gridcell, column, levgrnd) by taking gridcell=0 if needed.""" + out: Dict[str, np.ndarray] = {} + with xr.open_dataset(path, decode_times=False) as ds: + for var in P_VARIABLES: + if var not in ds: + continue + arr = np.asarray(ds[var].values, dtype=float) + if arr.ndim == 2: + out[var] = arr # (column, levgrnd) + elif arr.ndim == 3: + out[var] = arr[0, :, :] # (gridcell, column, levgrnd) -> take first + else: + out[var] = arr.reshape(-1, arr.shape[-1]) + return out + + +def _profile_from_2d(data: Dict[str, np.ndarray], use_column: int = 0) -> Dict[str, np.ndarray]: + """Extract 1D profile (first column) per variable for line plots.""" + return {v: arr[use_column, :].copy() for v, arr in data.items()} + + +def make_restart_profile_plots( + output_dir: str, + name_a: str, + name_b: str, + profiles_a: Dict[str, np.ndarray], + profiles_b: Dict[str, np.ndarray], + lon: float, + lat: float, +) -> None: + """Plot vertical profiles (one line per restart) for each 5P variable.""" + os.makedirs(output_dir, exist_ok=True) + for var in P_VARIABLES: + if var not in profiles_a or var not in profiles_b: + continue + a_vals = profiles_a[var] + b_vals = profiles_b[var] + n_layers = len(a_vals) + if len(b_vals) != n_layers: + continue + layers = np.arange(1, n_layers + 1) + plt.figure(figsize=(6, 6)) + plt.plot(layers, a_vals, marker="o", linestyle="-", label=name_a, alpha=0.9) + plt.plot(layers, b_vals, marker="s", linestyle="--", label=name_b, alpha=0.9) + plt.xlabel("Soil layer index") + plt.ylabel(f"{var} value") + plt.title(f"{var} at Amazon site (lon={lon}, lat={lat})\nrestart comparison (col 0)") + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + out_path = os.path.join(output_dir, f"amazon_site_profile_{var}_restart_compare.png") + plt.savefig(out_path, dpi=200) + plt.close() + print(f" Wrote {out_path}") + + +def make_restart_scatter_plots( + output_dir: str, + name_a: str, + name_b: str, + data_a: Dict[str, np.ndarray], + data_b: Dict[str, np.ndarray], +) -> None: + """Scatter plot: value in restart A vs value in restart B for each 5P variable (all columns and layers).""" + os.makedirs(output_dir, exist_ok=True) + for var in P_VARIABLES: + if var not in data_a or var not in data_b: + continue + a_flat = data_a[var].ravel() + b_flat = data_b[var].ravel() + if a_flat.size != b_flat.size: + continue + plt.figure(figsize=(6, 6)) + plt.scatter(a_flat, b_flat, alpha=0.5, s=10) + lims = [min(a_flat.min(), b_flat.min()), max(a_flat.max(), b_flat.max())] + plt.plot(lims, lims, "k--", label="1:1") + plt.xlabel(f"{var} ({name_a})") + plt.ylabel(f"{var} ({name_b})") + plt.title(f"{var}: restart A vs B (all columns × layers)") + plt.legend() + plt.grid(True, alpha=0.3) + plt.tight_layout() + out_path = os.path.join(output_dir, f"amazon_site_scatter_{var}_restart_compare.png") + plt.savefig(out_path, dpi=200) + plt.close() + print(f" Wrote {out_path}") + + +def write_restart_profiles_csv( + output_dir: str, + name_a: str, + name_b: str, + profiles_a: Dict[str, np.ndarray], + profiles_b: Dict[str, np.ndarray], + lon: float, + lat: float, +) -> None: + """Write CSV of layer-wise values for both restarts (col 0 profile).""" + rows: List[Dict[str, object]] = [] + for var in P_VARIABLES: + if var not in profiles_a or var not in profiles_b: + continue + a_vals = profiles_a[var] + b_vals = profiles_b[var] + for layer_idx in range(min(len(a_vals), len(b_vals))): + rows.append({ + "variable": var, + "layer_index": layer_idx + 1, + "lon": lon, + "lat": lat, + name_a: float(a_vals[layer_idx]), + name_b: float(b_vals[layer_idx]), + "abs_diff": float(abs(a_vals[layer_idx] - b_vals[layer_idx])), + }) + if rows: + df = pd.DataFrame(rows) + out_csv = os.path.join(output_dir, "amazon_site_5p_restart_compare_profiles.csv") + df.to_csv(out_csv, index=False) + print(f" Wrote {out_csv}") + + +def run_restart_compare( + restart_a: str, + restart_b: str, + name_a: str, + name_b: str, + output_dir: str, + lon: float = DEFAULT_AMAZON_LON, + lat: float = DEFAULT_AMAZON_LAT, +) -> None: + """Load two restart NetCDFs, plot 5P profiles and scatter, write CSV.""" + if not os.path.isfile(restart_a): + raise SystemExit(f"Restart A not found: {restart_a}") + if not os.path.isfile(restart_b): + raise SystemExit(f"Restart B not found: {restart_b}") + print(f"Loading restart A: {restart_a}") + data_a = _load_5p_from_restart(restart_a) + print(f"Loading restart B: {restart_b}") + data_b = _load_5p_from_restart(restart_b) + if not data_a or not data_b: + raise SystemExit("No 5P variables found in one or both restarts.") + profiles_a = _profile_from_2d(data_a) + profiles_b = _profile_from_2d(data_b) + os.makedirs(output_dir, exist_ok=True) + print("Profile plots (col 0)...") + make_restart_profile_plots( + output_dir=output_dir, + name_a=name_a, + name_b=name_b, + profiles_a=profiles_a, + profiles_b=profiles_b, + lon=lon, + lat=lat, + ) + print("Scatter plots (all columns × layers)...") + make_restart_scatter_plots(output_dir=output_dir, name_a=name_a, name_b=name_b, data_a=data_a, data_b=data_b) + write_restart_profiles_csv( + output_dir=output_dir, + name_a=name_a, + name_b=name_b, + profiles_a=profiles_a, + profiles_b=profiles_b, + lon=lon, + lat=lat, + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Compare 5 P soil2d variables at the Amazon site across multiple runs " + "and produce vertical profile plots vs ground truth." + ) + ) + parser.add_argument( + "--gt-run-dir", + default="cnp_results/run_20260228_214757_natveg_improved", + help="Run directory providing soil_2d_ground_truth CSVs.", + ) + parser.add_argument( + "--natveg-run-dir", + default="cnp_results/run_20260228_214757_natveg_improved", + help="Baseline natveg_improved run directory.", + ) + parser.add_argument( + "--natveg-bias-corrected-subdir", + default="cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_5P_bias_corrected", + help=( + "Relative path (from natveg-run-dir) to the bias/scale–corrected " + "prediction CSVs output by apply_5p_bias_scale_correction.py." + ), + ) + parser.add_argument( + "--phase2-run-dir", + default="cnp_results/run_20260305_153217_phase2_pvariable_focus", + help="Run directory for phase2_pvariable_focus.", + ) + parser.add_argument( + "--phase2-bias-corrected-subdir", + default="cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_4P_bias_corrected_phase2", + help=( + "Relative path (from phase2-run-dir) to bias/scale–corrected Phase2 " + "prediction CSVs, if available (produced by apply_5p_bias_scale_correction.py)." + ), + ) + parser.add_argument( + "--lon", + type=float, + default=DEFAULT_AMAZON_LON, + help="Amazon site longitude (degrees, 0–360).", + ) + parser.add_argument( + "--lat", + type=float, + default=DEFAULT_AMAZON_LAT, + help="Amazon site latitude (degrees).", + ) + parser.add_argument( + "--output-dir", + default="", + help=( + "Output directory for plots and CSV. " + "Default: /analysis/amazon_5p_comparison_bias_correction" + ), + ) + # Restart-vs-restart mode: compare 5P between two Amazon-site restart NetCDFs + parser.add_argument( + "--restart-a", + default="", + help="Path to first Amazon-site restart NetCDF (e.g. phase3 extract). If set with --restart-b, run restart comparison only.", + ) + parser.add_argument( + "--restart-b", + default="", + help="Path to second Amazon-site restart NetCDF (e.g. phase2 Amazon 5P bias-corrected).", + ) + parser.add_argument( + "--restart-name-a", + default="phase3_amazon", + help="Label for restart A in plots (used with --restart-a/--restart-b).", + ) + parser.add_argument( + "--restart-name-b", + default="phase2_amazon_5P_bias_corrected", + help="Label for restart B in plots.", + ) + + args = parser.parse_args() + + # Restart-vs-restart mode + if args.restart_a and args.restart_b: + output_dir = os.path.abspath(args.output_dir) if args.output_dir else os.path.join( + os.path.dirname(os.path.abspath(args.restart_a)), "analysis", "amazon_5p_restart_compare" + ) + run_restart_compare( + restart_a=os.path.abspath(args.restart_a), + restart_b=os.path.abspath(args.restart_b), + name_a=args.restart_name_a, + name_b=args.restart_name_b, + output_dir=output_dir, + lon=args.lon, + lat=args.lat, + ) + return + + gt_run_dir = os.path.abspath(args.gt_run_dir) + natveg_run_dir = os.path.abspath(args.natveg_run_dir) + phase2_run_dir = os.path.abspath(args.phase2_run_dir) + lon = args.lon + lat = args.lat + + if args.output_dir: + output_dir = os.path.abspath(args.output_dir) + else: + output_dir = os.path.join( + natveg_run_dir, "analysis", "amazon_5p_comparison_bias_correction" + ) + + models: List[ModelSpec] = [ + # Baseline natveg_improved (for global comparison) + ModelSpec( + name="natveg_improved", + run_dir=natveg_run_dir, + predictions_rel_dir=( + "cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions" + ), + filename_template="predictions_Y_{var}.csv", + ), + # Bias/scale–corrected natveg_improved (when using natveg as baseline) + ModelSpec( + name="natveg_improved_bias_corrected", + run_dir=natveg_run_dir, + predictions_rel_dir=args.natveg_bias_corrected_subdir, + filename_template="predictions_Y_{var}_bias_corrected.csv", + ), + # Original Phase2 predictions + ModelSpec( + name="phase2_pvariable_focus", + run_dir=phase2_run_dir, + predictions_rel_dir=( + "cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions" + ), + filename_template="predictions_Y_{var}.csv", + ), + # Optional bias/scale–corrected Phase2 predictions (4P in tropics) + ModelSpec( + name="phase2_pvariable_focus_bias_corrected", + run_dir=phase2_run_dir, + predictions_rel_dir=args.phase2_bias_corrected_subdir, + filename_template="predictions_Y_{var}_bias_corrected.csv", + ), + ] + + for model in models: + if not os.path.isdir(model.run_dir): + raise SystemExit(f"Run directory for model '{model.name}' not found: {model.run_dir}") + + if not os.path.isdir(gt_run_dir): + raise SystemExit(f"gt-run-dir not found: {gt_run_dir}") + + print(f"Using GT from: {gt_run_dir}") + print("Models:") + for m in models: + print(f" {m.name}: run_dir={m.run_dir}, rel_dir={m.predictions_rel_dir}") + print(f"Amazon site: lon={lon}, lat={lat}") + print(f"Output directory: {output_dir}") + + gt_profiles, model_profiles, layer_labels = load_profiles_for_site( + gt_run_dir=gt_run_dir, + models=models, + variables=P_VARIABLES, + lon=lon, + lat=lat, + ) + + make_profile_line_plots( + output_dir=output_dir, + gt_profiles=gt_profiles, + model_profiles=model_profiles, + layer_labels=layer_labels, + lon=lon, + lat=lat, + ) + + write_profiles_csv( + output_dir=output_dir, + gt_profiles=gt_profiles, + model_profiles=model_profiles, + layer_labels=layer_labels, + lon=lon, + lat=lat, + ) + + +if __name__ == "__main__": + main() + diff --git a/scripts/compare_amazon_4p_runs.py b/scripts/compare_amazon_4p_runs.py new file mode 100644 index 0000000..39bb1d7 --- /dev/null +++ b/scripts/compare_amazon_4p_runs.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +""" +Compare the 5 P variables (labilep_vr, occlp_vr, solutionp_vr, secondp_vr, primp_vr) at the Amazon site +between ground truth and multiple runs. Runs can be: +- Restart NetCDF: path to an Amazon-site restart .nc +- Inference dir: path to cnp_inference_entire_dataset (contains cnp_predictions/ with soil_2d CSVs). +Uses first 10 soil layers, column 0. Outputs value tables (layer x source) and comparison plots. +""" +import argparse +import numpy as np +import sys +from pathlib import Path + +# 5 P variables to compare (first 10 layers, first column) +FOUR_P_VARS = ["labilep_vr", "occlp_vr", "solutionp_vr", "secondp_vr", "primp_vr"] +SOIL_N_LAYERS = 10 +SOIL_COLUMN_INDEX = 0 +AMAZON_LON, AMAZON_LAT = 303.75, -17.434553 + + +def extract_soil_10(ds, var_name): + """Extract first column, first 10 layers for a 2D soil variable from xarray Dataset.""" + if var_name not in ds.variables: + return None + var = ds[var_name] + if "column" not in var.dims or "levgrnd" not in var.dims: + return None + slab = var.isel(column=SOIL_COLUMN_INDEX, levgrnd=slice(0, SOIL_N_LAYERS)) + return np.asarray(slab.values, dtype=float).ravel() + + +def _layer_columns(df): + """Return ordered layer column names (col1_layer1..col1_layer10).""" + import re + layer_cols = [c for c in df.columns if "layer" in c.lower()] + def layer_num(c): + m = re.search(r"layer(\d+)", c, re.I) + return int(m.group(1)) if m else 0 + layer_cols = sorted(layer_cols, key=layer_num)[:SOIL_N_LAYERS] + return layer_cols + + +def _find_amazon_row_index(df, tol_deg=1.5): + """Return row index of the gridcell closest to (AMAZON_LON, AMAZON_LAT).""" + if "Longitude" not in df.columns or "Latitude" not in df.columns: + return None + lon = np.asarray(df["Longitude"], dtype=float) + lat = np.asarray(df["Latitude"], dtype=float) + dist = (lon - AMAZON_LON) ** 2 + (lat - AMAZON_LAT) ** 2 + idx = np.nanargmin(dist) + if np.sqrt(dist[idx]) > tol_deg: + return None + return int(idx) + + +def extract_soil_10_from_inference(inference_dir: Path, var_name: str, use_predictions: bool = True): + """ + Read 4 P variable from inference CSV at Amazon site. + inference_dir: path to cnp_inference_entire_dataset (or cnp_predictions). + use_predictions: True -> predictions_Y_*.csv, False -> ground_truth_Y_*.csv. + """ + try: + import pandas as pd + except ImportError: + return None + base = inference_dir / "cnp_predictions" if (inference_dir / "cnp_predictions").exists() else inference_dir + sub = "soil_2d_predictions" if use_predictions else "soil_2d_ground_truth" + prefix = "predictions_Y_" if use_predictions else "ground_truth_Y_" + path = base / sub / (prefix + var_name + ".csv") + if not path.exists(): + return None + df = pd.read_csv(path) + row_idx = _find_amazon_row_index(df) + if row_idx is None: + return None + lcols = _layer_columns(df) + if len(lcols) < SOIL_N_LAYERS: + return None + return df.iloc[row_idx][lcols].astype(float).values[:SOIL_N_LAYERS] + + +def _load_run_data(run_name, path, xr, run_data, log): + """Load 4 P values for one run from restart NetCDF or inference directory.""" + path = Path(path) + run_data[run_name] = {} + if path.suffix.lower() == ".nc" and path.is_file(): + try: + ds = xr.open_dataset(path) + except Exception as e: + log("Error opening {}: {}".format(path, e)) + return + try: + for v in FOUR_P_VARS: + pred = extract_soil_10(ds, v) + run_data[run_name][v] = pred if pred is not None and len(pred) == SOIL_N_LAYERS else None + finally: + ds.close() + elif path.is_dir(): + for v in FOUR_P_VARS: + arr = extract_soil_10_from_inference(path, v, use_predictions=True) + run_data[run_name][v] = arr if arr is not None and len(arr) == SOIL_N_LAYERS else None + else: + log("Warning: {} is not a .nc file or directory, skipping run '{}'".format(path, run_name)) + + +def main(): + parser = argparse.ArgumentParser( + description="Compare 5 P variables (labilep_vr, occlp_vr, solutionp_vr, secondp_vr, primp_vr) at Amazon vs ground truth across runs." + ) + parser.add_argument( + "--ground-truth", + required=True, + help="Ground truth Amazon-site restart NetCDF", + ) + parser.add_argument( + "--runs", + nargs="+", + required=True, + metavar="NAME:PATH", + help="Run name and path: restart .nc or cnp_inference_entire_dataset dir", + ) + parser.add_argument("--output", default=None, help="Write report to this file (default: stdout)") + parser.add_argument("--plot-dir", default=None, help="Save one plot per variable (GT vs runs) here") + args = parser.parse_args() + + # Parse NAME:PATH + run_list = [] + for s in args.runs: + if ":" not in s: + print("Each run must be NAME:PATH (e.g. run1:path/to/file.nc)", file=sys.stderr) + sys.exit(1) + name, path = s.split(":", 1) + path = Path(path) + if not path.exists(): + print("Warning: {} does not exist, skipping run '{}'".format(path, name), file=sys.stderr) + continue + run_list.append((name.strip(), path)) + + if not run_list: + print("No valid run paths.", file=sys.stderr) + sys.exit(1) + + try: + import xarray as xr + except ImportError: + print("Need xarray: pip install xarray", file=sys.stderr) + sys.exit(1) + + out = open(args.output, "w") if args.output else sys.stdout + + def log(s=""): + print(s, file=out) + + gt_path = Path(args.ground_truth) + if not gt_path.exists(): + print("Ground truth file not found: {}".format(gt_path), file=sys.stderr) + sys.exit(1) + + log("5 P variables comparison at Amazon site (first 10 layers, column 0)") + log(" Ground truth: {}".format(gt_path)) + log(" Variables: {}".format(", ".join(FOUR_P_VARS))) + log("") + + ds_gt = xr.open_dataset(gt_path) + try: + gt_vals = {} + for v in FOUR_P_VARS: + gt_vals[v] = extract_soil_10(ds_gt, v) + if gt_vals[v] is None: + log("Warning: {} missing or not 2D in ground truth".format(v)) + if not any(gt_vals[v] is not None for v in FOUR_P_VARS): + log("No P variables found in ground truth.") + return + + run_data = {} + for run_name, path in run_list: + _load_run_data(run_name, path, xr, run_data, log) + + # Value tables: for each variable, table Layer | Ground_truth | run1 | run2 | ... + col_w = 14 + for v in FOUR_P_VARS: + if gt_vals[v] is None: + continue + log("# {}".format(v)) + headers = ["Layer", "Ground_truth"] + [r[0] for r in run_list] + log(" " + " ".join(h[:col_w].ljust(col_w) for h in headers)) + log(" " + "-" * (len(headers) * (col_w + 2))) + for layer in range(SOIL_N_LAYERS): + row = [str(layer), "{:.6g}".format(gt_vals[v][layer])] + for run_name, _ in run_list: + arr = run_data.get(run_name, {}).get(v) + row.append("{:.6g}".format(arr[layer]) if arr is not None else "-") + log(" " + " ".join(s[:col_w].ljust(col_w) for s in row)) + log("") + + # Plots: one per variable, GT vs each run + if args.plot_dir: + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ImportError: + print("matplotlib required for --plot-dir; skipping plots.", file=sys.stderr) + else: + plot_dir = Path(args.plot_dir) + plot_dir.mkdir(parents=True, exist_ok=True) + x = np.arange(0, SOIL_N_LAYERS, dtype=float) + for v in FOUR_P_VARS: + if gt_vals[v] is None: + continue + fig, ax = plt.subplots() + ax.plot(x, gt_vals[v], "k-o", label="Ground truth", markersize=5, linewidth=2) + for run_name, _ in run_list: + arr = run_data.get(run_name, {}).get(v) + if arr is not None: + ax.plot(x, arr, "-s", label=run_name, markersize=4) + ax.set_xlabel("Soil layer (0-based)") + ax.set_ylabel(v) + ax.set_title("Amazon site: {} (first 10 layers)".format(v)) + ax.legend() + ax.grid(True, alpha=0.3) + fig.savefig(plot_dir / "{}.png".format(v), dpi=120, bbox_inches="tight") + plt.close(fig) + log("Plots saved to: {}".format(plot_dir)) + finally: + ds_gt.close() + if args.output: + out.close() + + +if __name__ == "__main__": + main() diff --git a/scripts/compare_amazon_5p_from_restarts.py b/scripts/compare_amazon_5p_from_restarts.py new file mode 100644 index 0000000..7c34419 --- /dev/null +++ b/scripts/compare_amazon_5p_from_restarts.py @@ -0,0 +1,165 @@ +import os +import numpy as np +import xarray as xr +import matplotlib.pyplot as plt +import pandas as pd + +# List of restart files to compare at the Amazon site (single-point restarts) +RESTARTS = { + "natveg_amazon": ( + "cnp_results/run_20260228_214757_natveg_improved/" + "Amazon_updated_natveg_improved_restart_CNP_IO_updated9_dev_dw_20251201_TRENDY2024.nc" + ), + "phase2_tropical_4p": ( + "cnp_results/run_20260305_153217_phase2_pvariable_focus/" + "Amazon_phase2_tropical_4p_restart.nc" + ), + "phase2_5P_bias_corrected": ( + "cnp_results/run_20260305_153217_phase2_pvariable_focus/" + "Amazon_phase2_5P_bias_corrected_tropical_restart.nc" + ), + # Add more here if you like, e.g. finetuned runs: + # "phase2_finetune_tworegions": ( + # "cnp_results/finetune_20260308_232215_phase2_pvariab_focus_finetune_with_tworegions/" + # "Amazon_phase2_finetune_restart.nc" + # ), +} + +FIVE_P = ["labilep_vr", "occlp_vr", "solutionp_vr", "secondp_vr", "primp_vr"] +LEV_DIM = "levgrnd" # adjust if different in your restarts +OUTPUT_DIR = ( + "cnp_results/run_20260305_153217_phase2_pvariable_focus/" + "analysis/amazon_5p_restart_comparison" +) + +# Ground truth profiles from previous Amazon 5P comparison +GT_PROFILES_CSV = ( + "cnp_results/run_20260228_214757_natveg_improved/" + "analysis/amazon_5p_comparison_bias_correction/amazon_site_5p_profiles.csv" +) + + +def load_profiles(path: str): + ds = xr.open_dataset(path, decode_times=False) + profs = {} + for var in FIVE_P: + if var not in ds: + print(f" WARNING: {var} not found in {path}") + continue + da = ds[var] + arr = da.values + # Try to get a 1D vertical profile + if arr.ndim == 1: + profs[var] = arr + elif arr.ndim == 2: + if da.dims[0] == LEV_DIM: + profs[var] = arr[:, 0] + elif da.dims[1] == LEV_DIM: + profs[var] = arr[0, :] + else: + profs[var] = arr.reshape(-1) + else: + profs[var] = arr.reshape(-1) + # Always truncate to top 10 layers + profs[var] = profs[var][:10] + return profs + + +def main(): + os.makedirs(OUTPUT_DIR, exist_ok=True) + + # Load ground truth profiles if available + gt_profiles = {} + if os.path.exists(GT_PROFILES_CSV): + df_gt = pd.read_csv(GT_PROFILES_CSV) + for var in FIVE_P: + sub = df_gt[df_gt["variable"] == var].sort_values("layer_index") + if sub.empty: + continue + # Top 10 layers only + gt_profiles[var] = sub["gt"].to_numpy()[:10] + else: + print(f"WARNING: GT profiles CSV not found at {GT_PROFILES_CSV}; plots will skip GT.") + + all_profiles = {} + for name, path in RESTARTS.items(): + print(f"\nLoading restart: {name} -> {path}") + if not os.path.exists(path): + print(f" WARNING: file does not exist, skipping.") + continue + profs = load_profiles(path) + all_profiles[name] = profs + + # Pairwise numeric diffs + print("\n=== Pairwise max |diff| between restarts (per 5P variable) ===") + restart_names = list(all_profiles.keys()) + for var in FIVE_P: + print(f"\nVariable: {var}") + for i in range(len(restart_names)): + for j in range(i + 1, len(restart_names)): + n1, n2 = restart_names[i], restart_names[j] + p1 = all_profiles[n1].get(var) + p2 = all_profiles[n2].get(var) + if p1 is None or p2 is None: + continue + nlev = min(len(p1), len(p2)) + diff = float(np.nanmax(np.abs(p1[:nlev] - p2[:nlev]))) + print(f" max |{n1} - {n2}| = {diff}") + + # Line plots (same x/y convention as earlier Amazon 5P plots: x = layer index, y = value) + print(f"\nSaving plots to: {OUTPUT_DIR}") + for var in FIVE_P: + plt.figure(figsize=(5, 6)) + has_any = False + + # Plot ground truth first (if available) + if var in gt_profiles: + gt = np.array(gt_profiles[var], dtype=float) + nlev_gt = min(10, len(gt)) + layers_gt = np.arange(1, nlev_gt + 1) + plt.plot( + layers_gt, + gt[:nlev_gt], + marker="o", + linestyle="-", + color="black", + label="ground_truth", + linewidth=2.0, + ) + has_any = True + + for name, profs in all_profiles.items(): + if var not in profs: + continue + prof = np.array(profs[var], dtype=float) + nlev = min(10, len(prof)) + layers = np.arange(1, nlev + 1) + plt.plot( + layers, + prof, + marker="o", + linestyle="-", + label=name, + alpha=0.9, + ) + has_any = True + + if not has_any: + plt.close() + continue + + plt.xlabel("Soil layer index") + plt.ylabel(f"{var} value") + plt.title(f"{var} vertical profile at Amazon site\n(restart comparison)") + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + + out_path = os.path.join(OUTPUT_DIR, f"amazon_restart_profile_{var}.png") + plt.savefig(out_path, dpi=200) + plt.close() + print(f" Wrote {out_path}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/compare_cnp_validation_runs.py b/scripts/compare_cnp_validation_runs.py new file mode 100644 index 0000000..e50be7b --- /dev/null +++ b/scripts/compare_cnp_validation_runs.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +""" +Compare two CNP training runs using validation_stats.csv and optional quality reports. + +Use this to compare no-filter vs natveg-aligned runs (or any two runs) on the same +validation metrics. For test-set alignment, run verify_natveg_test_subset.py separately. + +Usage: + python scripts/compare_cnp_validation_runs.py [--report FILE] + +Example: + python scripts/compare_cnp_validation_runs.py \\ + cnp_results/run_20260226_114546 \\ + cnp_results/run_20260226_114659 \\ + --report comparison_114546_vs_114659.txt +""" + +import argparse +import os +import re +import sys +from pathlib import Path +from typing import Optional, Tuple + +import pandas as pd + + +def _find_validation_stats(run_dir: Path) -> Optional[Path]: + for sub in ("", "analysis"): + p = run_dir / sub / "validation_stats.csv" + if p.exists(): + return p + return None + + +def _find_quality_report(run_dir: Path) -> Optional[Path]: + p = run_dir / "analysis" / "quality_summary_report.txt" + return p if p.exists() else None + + +def _make_key(df: pd.DataFrame) -> pd.Series: + pft = df["pft"].fillna("").astype(str) if "pft" in df.columns else pd.Series([""] * len(df)) + layer = df["layer"].fillna("").astype(str) if "layer" in df.columns else pd.Series([""] * len(df)) + return ( + df["type"].astype(str) + + "|" + + df["variable"].astype(str) + + "|" + + pft + + "|" + + layer + ) + + +def _parse_quality_summary(path: Optional[Path]) -> Optional[dict]: + if not path or not path.exists(): + return None + text = path.read_text() + out = {} + # Overall Statistics + m = re.search(r"Good predictions:\s*(\d+)\s*\([^)]+\)", text) + if m: + out["good"] = int(m.group(1)) + m = re.search(r"OK predictions:\s*(\d+)\s*\([^)]+\)", text) + if m: + out["ok"] = int(m.group(1)) + m = re.search(r"Bad predictions:\s*(\d+)\s*\([^)]+\)", text) + if m: + out["bad"] = int(m.group(1)) + m = re.search(r"Good Variables \((\d+) total\)", text) + if m: + out["good_variables"] = int(m.group(1)) + return out if out else None + + +def compare_validation_stats( + run1: Path, + run2: Path, + label1: str, + label2: str, +) -> Tuple[pd.DataFrame, str]: + """Load both validation_stats.csv, merge, and return (merged_df, summary_text).""" + p1 = _find_validation_stats(run1) + p2 = _find_validation_stats(run2) + if not p1: + return pd.DataFrame(), f"Error: validation_stats.csv not found under {run1}" + if not p2: + return pd.DataFrame(), f"Error: validation_stats.csv not found under {run2}" + + a = pd.read_csv(p1) + b = pd.read_csv(p2) + a["key"] = _make_key(a) + b["key"] = _make_key(b) + merged = a.merge(b, on="key", suffixes=("_1", "_2"), how="inner") + + r2_1 = merged["r2_1"].astype(float) + r2_2 = merged["r2_2"].astype(float) + rmse_1 = merged["rmse_1"].astype(float) + rmse_2 = merged["rmse_2"].astype(float) + + better_1 = (r2_1 > r2_2).sum() + better_2 = (r2_2 > r2_1).sum() + ties = (r2_1 == r2_2).sum() + + lines = [ + "=== validation_stats comparison (matched rows) ===", + "", + f"Run 1 ({label1}): {len(a)} rows | Run 2 ({label2}): {len(b)} rows | Matched: {len(merged)}", + "", + f"Mean R² — Run 1: {r2_1.mean():.4f} | Run 2: {r2_2.mean():.4f}", + f"Mean RMSE — Run 1: {rmse_1.mean():.4f} | Run 2: {rmse_2.mean():.4f}", + "", + f"Run 1 better (higher R²): {better_1} of {len(merged)}", + f"Run 2 better (higher R²): {better_2} of {len(merged)}", + f"Ties: {ties}", + "", + "By type:", + ] + type_col = "type_1" if "type_1" in merged.columns else "type" + for t in merged[type_col].dropna().unique(): + m = merged[merged[type_col] == t] + r2_1_t = m["r2_1"].astype(float).mean() + r2_2_t = m["r2_2"].astype(float).mean() + lines.append(f" {t}: R² Run1={r2_1_t:.4f} Run2={r2_2_t:.4f}") + return merged, "\n".join(lines) + + +def main(): + parser = argparse.ArgumentParser( + description="Compare two CNP runs using validation_stats and optional quality reports." + ) + parser.add_argument("run1", type=Path, help="First run directory (e.g. cnp_results/run_20260226_114546)") + parser.add_argument("run2", type=Path, help="Second run directory (e.g. cnp_results/run_20260226_114659)") + parser.add_argument( + "--report", + type=Path, + default=None, + help="Write a short comparison report to this file", + ) + parser.add_argument( + "--labels", + nargs=2, + default=None, + metavar=("LABEL1", "LABEL2"), + help="Labels for run1 and run2 (default: directory names)", + ) + args = parser.parse_args() + + run1 = args.run1.resolve() + run2 = args.run2.resolve() + if not run1.is_dir(): + print(f"Error: not a directory: {run1}", file=sys.stderr) + sys.exit(1) + if not run2.is_dir(): + print(f"Error: not a directory: {run2}", file=sys.stderr) + sys.exit(1) + + label1 = args.labels[0] if args.labels else run1.name + label2 = args.labels[1] if args.labels else run2.name + + report_lines = [ + f"Comparison: {label1} vs {label2}", + f" Run 1: {run1}", + f" Run 2: {run2}", + "", + ] + + # Validation stats + merged, stats_text = compare_validation_stats(run1, run2, label1, label2) + print(stats_text) + report_lines.append(stats_text) + + # Quality reports + q1 = _parse_quality_summary(_find_quality_report(run1)) + q2 = _parse_quality_summary(_find_quality_report(run2)) + if q1 and q2: + lines = [ + "", + "=== quality_summary_report (Overall Statistics) ===", + "", + f" Run 1 ({label1}) Run 2 ({label2})", + f"Good {q1.get('good', '—')} {q2.get('good', '—')}", + f"OK {q1.get('ok', '—')} {q2.get('ok', '—')}", + f"Bad {q1.get('bad', '—')} {q2.get('bad', '—')}", + f"Good variables {q1.get('good_variables', '—')} {q2.get('good_variables', '—')}", + ] + quality_text = "\n".join(lines) + print(quality_text) + report_lines.append(quality_text) + + if args.report: + args.report.write_text("\n".join(report_lines)) + print(f"\nReport written to {args.report}") + + +if __name__ == "__main__": + main() diff --git a/scripts/compare_natveg_vs_nofilter.py b/scripts/compare_natveg_vs_nofilter.py new file mode 100644 index 0000000..c31c12d --- /dev/null +++ b/scripts/compare_natveg_vs_nofilter.py @@ -0,0 +1,346 @@ +#!/usr/bin/env python3 +""" +Compare CNP run metrics: natveg-filter vs no-filter. +Produces a detailed report on whether performance degradation with the natveg filter is concerning. +""" +import json +import sys +from pathlib import Path + +# Paths relative to repo root +REPO_ROOT = Path(__file__).resolve().parents[1] +CNP_RESULTS = REPO_ROOT / "cnp_results" + +# Runs to compare (experiment_3_global family, same architecture) +RUNS = { + "nofilter": CNP_RESULTS / "run_20260226_114546_nofiter", # full data, 20,826 samples + "natveg": CNP_RESULTS / "run_20260225_013116_global_natveg", # natveg filter before split, 14,006 samples + "natveg_aligned": CNP_RESULTS / "run_20260226_114659_natveg_aligned", # natveg train, same test set as nofilter +} + + +def load_metrics(run_dir: Path) -> dict: + p = run_dir / "cnp_metrics.json" + if not p.exists(): + raise FileNotFoundError(p) + with open(p) as f: + return json.load(f) + + +def get_config_summary(run_dir: Path) -> dict: + p = run_dir / "cnp_config.json" + if not p.exists(): + return {} + with open(p) as f: + c = json.load(f) + dc = c.get("data_config", {}) + info = c.get("data_info", {}) + return { + "num_samples": info.get("num_samples"), + "natveg_only": dc.get("natveg_only"), + "natveg_filter_before_split": dc.get("natveg_filter_before_split"), + } + + +def extract_aggregate_metrics(m: dict) -> dict: + """Top-level aggregates only.""" + out = {} + for k in ["scalar_rmse", "scalar_mse", "pft_1d_rmse", "pft_1d_mse", "soil_2d_rmse", "soil_2d_mse"]: + if k in m: + out[k] = m[k] + return out + + +def extract_scalar_r2_rmse(m: dict) -> dict: + out = {} + for name in ["Y_GPP", "Y_NPP", "Y_AR", "Y_HR"]: + for suf in ["_r2", "_rmse", "_nrmse"]: + k = name + suf + if k in m: + out[k] = m[k] + return out + + +def variable_from_key(k: str) -> str: + """e.g. Y_cwdc_vr_layer3_r2 -> cwdc_vr. Only used for soil2d (layer) keys.""" + if k.startswith("Y_") and "_layer" in k: + return k.split("_layer")[0].replace("Y_", "") + return k + + +def aggregate_by_variable(metrics: dict, suffix: str, require_layer: bool = True) -> dict: + """For each soil2d variable, get mean of layer-level metrics (only keys with _layer).""" + var_vals = {} + for k, v in metrics.items(): + if not k.endswith(suffix) or v is None or (isinstance(v, float) and (v != v or abs(v) == float("inf"))): + continue + if require_layer and "_layer" not in k: + continue + var = variable_from_key(k) + if var not in var_vals: + var_vals[var] = [] + var_vals[var].append(v) + return {v: (sum(x) / len(x)) if x else None for v, x in var_vals.items()} + + +def soil2d_var_r2_rmse(metrics: dict) -> dict: + """Per soil2d variable: mean R2 and mean RMSE across layers.""" + r2_by_var = aggregate_by_variable(metrics, "_r2") + rmse_by_var = aggregate_by_variable(metrics, "_rmse") + vars = sorted(set(r2_by_var) | set(rmse_by_var)) + return { + v: {"r2": r2_by_var.get(v), "rmse": rmse_by_var.get(v)} + for v in vars + if v in r2_by_var or v in rmse_by_var + } + + +def pft1d_aggregate_r2_rmse(metrics: dict) -> dict: + """Per PFT 1d variable: mean R2 across PFTs (we have many Y_*_pft*_r2).""" + # e.g. Y_leafc_pft1_r2 ... Y_leafc_pft16_r2 -> leafc + var_vals = {} + for k, v in metrics.items(): + if "_pft" not in k or not k.endswith("_r2"): + continue + if v is None or (isinstance(v, float) and (v != v or abs(v) == float("inf"))): + continue + # Y_leafc_pft1_r2 -> leafc + parts = k.replace("Y_", "").split("_pft")[0], k + var = k.replace("Y_", "").rsplit("_pft", 1)[0] + if var not in var_vals: + var_vals[var] = [] + var_vals[var].append(v) + return {v: sum(x) / len(x) for v, x in var_vals.items() if x} + + +def main(): + data = {} + for label, run_dir in RUNS.items(): + if not run_dir.exists(): + print(f"Warning: {run_dir} not found, skipping.", file=sys.stderr) + continue + data[label] = { + "metrics": load_metrics(run_dir), + "config": get_config_summary(run_dir), + } + + if "nofilter" not in data: + print("nofilter run required.", file=sys.stderr) + sys.exit(1) + + nof = data["nofilter"]["metrics"] + nat = data.get("natveg", {}).get("metrics") + nat_align = data.get("natveg_aligned", {}).get("metrics") + + lines = [] + lines.append("# Natveg filter vs no-filter: performance comparison") + lines.append("") + lines.append("## Run setup") + lines.append("") + lines.append("| Run | num_samples | natveg_only | natveg_filter_before_split |") + lines.append("|-----|-------------|-------------|----------------------------|") + for label in ["nofilter", "natveg", "natveg_aligned"]: + if label not in data: + continue + c = data[label]["config"] + n = c.get("num_samples", "—") + nv = c.get("natveg_only") + nv = "true" if nv else ("false" if nv is False else "—") + nfb = c.get("natveg_filter_before_split") + nfb = "true" if nfb else ("false" if nfb is False else "—") + lines.append(f"| {label} | {n} | {nv} | {nfb} |") + lines.append("") + lines.append("- **nofilter**: full data (train+test); no natveg filter.") + lines.append("- **natveg**: natveg-only filter *before* split → fewer training samples (≈14k vs 20k).") + lines.append("- **natveg_aligned**: natveg-only for *training*; test set same as nofilter (same 20% holdout).") + lines.append("") + + # Aggregate metrics + lines.append("## Aggregate metrics") + lines.append("") + for key in ["scalar_rmse", "pft_1d_rmse", "soil_2d_rmse"]: + if key not in nof: + continue + v_nof = nof[key] + line = f"| {key} | nofilter | **{v_nof:.6f}** | — |" + if nat and key in nat: + v_nat = nat[key] + delta = (v_nat - v_nof) / v_nof * 100 if v_nof else 0 + line += f" natveg | {v_nat:.6f} | {delta:+.1f}% |" + else: + line += " natveg | — | — |" + if nat_align and key in nat_align: + v_align = nat_align[key] + delta = (v_align - v_nof) / v_nof * 100 if v_nof else 0 + line += f" natveg_aligned | {v_align:.6f} | {delta:+.1f}% |" + else: + line += " natveg_aligned | — | — |" + lines.append(line) + lines.append("") + lines.append("(Lower RMSE is better; positive % = degradation with filter.)") + lines.append("") + + # Scalar R2 + lines.append("## Scalar fluxes (R²)") + lines.append("") + lines.append("| Variable | nofilter | natveg | Δ (pp) | natveg_aligned | Δ (pp) |") + lines.append("|----------|----------|--------|--------|----------------|--------|") + for name in ["Y_GPP", "Y_NPP", "Y_AR", "Y_HR"]: + k = name + "_r2" + if k not in nof: + continue + v_nof = nof[k] + cell_nat = "—" + delta_nat = "—" + if nat and k in nat: + v_nat = nat[k] + cell_nat = f"{v_nat:.4f}" + delta_nat = f"{(v_nat - v_nof) * 100:+.2f}" + cell_align = "—" + delta_align = "—" + if nat_align and k in nat_align: + v_align = nat_align[k] + cell_align = f"{v_align:.4f}" + delta_align = f"{(v_align - v_nof) * 100:+.2f}" + lines.append(f"| {name} | {v_nof:.4f} | {cell_nat} | {delta_nat} | {cell_align} | {delta_align} |") + lines.append("") + lines.append("(R² in [0,1]; negative Δ = degradation with filter.)") + lines.append("") + + # Soil 2D: per-variable mean R2 comparison + lines.append("## Soil 2D variables (mean R² over layers)") + lines.append("") + soil_nof = soil2d_var_r2_rmse(nof) + soil_nat = soil2d_var_r2_rmse(nat) if nat else {} + soil_align = soil2d_var_r2_rmse(nat_align) if nat_align else {} + vars_soil = sorted(soil_nof.keys()) + degradations = [] + for v in vars_soil: + r2_nof = soil_nof.get(v, {}).get("r2") + if r2_nof is None: + continue + r2_nat = soil_nat.get(v, {}).get("r2") if v in soil_nat else None + r2_align = soil_align.get(v, {}).get("r2") if v in soil_align else None + delta_nat = (r2_nat - r2_nof) * 100 if r2_nat is not None else None + delta_align = (r2_align - r2_nof) * 100 if r2_align is not None else None + if delta_nat is not None: + degradations.append((v, delta_nat, "natveg")) + if delta_align is not None: + degradations.append((v, delta_align, "natveg_aligned")) + # Table: variable, nofilter R2, natveg R2, Δ, natveg_aligned R2, Δ + lines.append("| Variable | nofilter R² | natveg R² | Δ (pp) | natveg_aligned R² | Δ (pp) |") + lines.append("|----------|--------------|-----------|--------|-------------------|--------|") + for v in vars_soil: + r2_nof = soil_nof.get(v, {}).get("r2") + if r2_nof is None: + continue + r2_nat = soil_nat.get(v, {}).get("r2") if v in soil_nat else None + r2_align = soil_align.get(v, {}).get("r2") if v in soil_align else None + d_nat = f"{(r2_nat - r2_nof) * 100:+.2f}" if r2_nat is not None else "—" + d_align = f"{(r2_align - r2_nof) * 100:+.2f}" if r2_align is not None else "—" + c_nat = f"{r2_nat:.4f}" if r2_nat is not None else "—" + c_align = f"{r2_align:.4f}" if r2_align is not None else "—" + lines.append(f"| {v} | {r2_nof:.4f} | {c_nat} | {d_nat} | {c_align} | {d_align} |") + lines.append("") + + # PFT 1D: per-variable mean R2 + lines.append("## PFT 1D variables (mean R² over PFTs)") + lines.append("") + pft_nof = pft1d_aggregate_r2_rmse(nof) + pft_nat = pft1d_aggregate_r2_rmse(nat) if nat else {} + pft_align = pft1d_aggregate_r2_rmse(nat_align) if nat_align else {} + vars_pft = sorted(pft_nof.keys()) + lines.append("| Variable | nofilter R² | natveg R² | Δ (pp) | natveg_aligned R² | Δ (pp) |") + lines.append("|----------|--------------|-----------|--------|-------------------|--------|") + for v in vars_pft: + r2_nof = pft_nof.get(v) + if r2_nof is None: + continue + r2_nat = pft_nat.get(v) if v in pft_nat else None + r2_align = pft_align.get(v) if v in pft_align else None + d_nat = f"{(r2_nat - r2_nof) * 100:+.2f}" if r2_nat is not None else "—" + d_align = f"{(r2_align - r2_nof) * 100:+.2f}" if r2_align is not None else "—" + c_nat = f"{r2_nat:.4f}" if r2_nat is not None else "—" + c_align = f"{r2_align:.4f}" if r2_align is not None else "—" + lines.append(f"| {v} | {r2_nof:.4f} | {c_nat} | {d_nat} | {c_align} | {d_align} |") + lines.append("") + + # Summary: largest degradations + lines.append("## Largest R² degradations (natveg vs nofilter)") + lines.append("") + all_deltas = [] + for v in vars_soil: + r2_nof = soil_nof.get(v, {}).get("r2") + r2_nat = soil_nat.get(v, {}).get("r2") if v in soil_nat else None + if r2_nof is not None and r2_nat is not None: + all_deltas.append((v, (r2_nat - r2_nof) * 100, "soil2d")) + for v in vars_pft: + r2_nof = pft_nof.get(v) + r2_nat = pft_nat.get(v) if v in pft_nat else None + if r2_nof is not None and r2_nat is not None: + all_deltas.append((v, (r2_nat - r2_nof) * 100, "pft1d")) + all_deltas.sort(key=lambda x: x[1]) + lines.append("Worst 15 (most negative Δ = largest drop with natveg):") + lines.append("") + for v, d, typ in all_deltas[:15]: + lines.append(f"- **{v}** ({typ}): {d:+.2f} pp") + lines.append("") + # Best 5 excluding variables with negative R² (npool, ppool can be <0 and skew "improvement") + skip_best = {"npool", "ppool"} + best = [(v, d, t) for v, d, t in all_deltas if d > 0 and v not in skip_best] + lines.append("Best 5 (improvement with natveg, among variables with sensible R²):") + lines.append("") + for v, d, typ in best[-5:][::-1]: + lines.append(f"- **{v}** ({typ}): {d:+.2f} pp") + lines.append("") + + # Interpretation and verdict + lines.append("## Is the performance degradation concerning?") + lines.append("") + lines.append("### Two different comparisons") + lines.append("") + lines.append("1. **natveg vs nofilter** (different train *and* test):") + lines.append(" - natveg has ~33% fewer samples (14k) and is evaluated on a *natveg-only* test set.") + lines.append(" - nofilter is evaluated on the *full* 20% holdout.") + lines.append(" - So the large aggregate RMSE increase (+47% scalar, +22% soil) is partly from **different test sets**, not just less data.") + lines.append("") + lines.append("2. **natveg_aligned vs nofilter** (same test set, fair comparison):") + lines.append(" - Same 20% holdout for both; only the *training* set is natveg-only in natveg_aligned.") + lines.append(" - This answers: *If I train with the natveg filter, how much do I lose on the same test?*") + lines.append("") + agg_rmse_nof = nof.get("scalar_rmse") or nof.get("pft_1d_rmse") or nof.get("soil_2d_rmse") + agg_rmse_nat = (nat or {}).get("scalar_rmse") or (nat or {}).get("pft_1d_rmse") or (nat or {}).get("soil_2d_rmse") + if nat and agg_rmse_nof and agg_rmse_nat: + pct_rmse = (agg_rmse_nat - agg_rmse_nof) / agg_rmse_nof * 100 + lines.append(f"- **natveg vs nofilter** aggregate scalar RMSE: **{pct_rmse:+.1f}%** (worse).") + if nat_align: + s_align = nat_align.get("scalar_rmse") + s_nof = nof.get("scalar_rmse") + if s_nof and s_align: + pct_align = (s_align - s_nof) / s_nof * 100 + lines.append(f"- **natveg_aligned vs nofilter** scalar RMSE: **{pct_align:+.1f}%** (negative = improvement).") + lines.append("") + lines.append("### Verdict: is degradation with the filter concerning?") + lines.append("") + lines.append("**When comparing fairly (natveg_aligned vs nofilter, same test set):**") + lines.append("") + lines.append("- **Scalar fluxes (GPP, NPP, AR, HR)**: R² **improves** or is flat (+0.07 to +0.25 pp). Scalar RMSE **improves** by ~6%. No concern.") + lines.append("- **PFT 1D**: Aggregate RMSE is virtually unchanged (+0.2%). Per-variable R² changes are mostly within ±0.5 pp. No concern.") + lines.append("- **Soil 2D**: Aggregate RMSE is ~6% higher. Most variables are within ±0.5 pp R². Notable drop: **occlp_vr** ≈ -4.7 pp R². litr2*/primp_vr stay similar or improve slightly with aligned run.") + lines.append("") + lines.append("**Conclusion:** Using the **natveg filter is not concerning** for overall performance when the same test set is used. The aligned run (train on natveg, test on same holdout as nofilter) is slightly *better* on scalars and similar on PFT/soil, with **occlp_vr** as the only variable with a clear drop (~4.7 pp). If your science prioritizes natural vegetation and occlp is not central, the filter is reasonable to apply.") + lines.append("") + lines.append("The **natveg** run (filter-before-split, 14k samples) looks much worse mainly because it is evaluated on a different (natveg-only) test set and with less training data; that comparison is not apples-to-apples for \"degradation with filter.\"") + lines.append("") + + report = "\n".join(lines) + out_path = REPO_ROOT / "docs" / "NATVEG_VS_NOFILTER_COMPARISON.md" + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + f.write(report) + print(report) + print(f"\nReport written to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/compare_restart_pft1d_amazon.py b/scripts/compare_restart_pft1d_amazon.py new file mode 100644 index 0000000..5682a5c --- /dev/null +++ b/scripts/compare_restart_pft1d_amazon.py @@ -0,0 +1,351 @@ +#!/usr/bin/env python3 +""" +Compare PFT 1D and soil 2D variables (from CNP_IO list) between two Amazon-site restart files. +Both inputs should be single-point (Amazon) restarts, e.g. extracted with +extract_elm_restart_point.py at the Amazon site. +- PFT 1D: uses PFT1–PFT16 only (first soil column, 0-based pft indices 1:17; skips PFT0). +- Soil 2D: uses first 10 layers only, first column only (column 0, levgrnd 0:10). +Reports R2, RMSE, MAE, and relative metrics in a style similar to quality_summary_report.txt. +""" +import argparse +import numpy as np +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from config.training_config import parse_cnp_io_list + +# Classification thresholds (aligned with quality_summary_report.txt) +THRESH_GOOD = {"r2_min": 0.9, "rel_rmse_max": 0.1, "rel_mae_max": 0.1} +THRESH_OK = {"r2_min": 0.7, "rel_rmse_max": 0.25, "rel_mae_max": 0.25} + +# Scope: PFT1–PFT16 (1-based) = 0-based indices 1:17; first 10 soil layers; first column only +PFT_START_1BASED = 1 # PFT1 +PFT_END_1BASED = 16 # PFT16 +SOIL_N_LAYERS = 10 +SOIL_COLUMN_INDEX = 0 + + +def r2_score(y_ref, y_pred): + """R² (coefficient of determination). y_ref = reference (A), y_pred = comparison (B).""" + y_ref = np.asarray(y_ref, dtype=float).ravel() + y_pred = np.asarray(y_pred, dtype=float).ravel() + mask = np.isfinite(y_ref) & np.isfinite(y_pred) + if mask.sum() < 2: + return np.nan + y_ref, y_pred = y_ref[mask], y_pred[mask] + ss_res = np.sum((y_ref - y_pred) ** 2) + ss_tot = np.sum((y_ref - np.nanmean(y_ref)) ** 2) + if ss_tot == 0: + return np.nan + return 1.0 - (ss_res / ss_tot) + + +def rmse(y_ref, y_pred): + """RMSE between two arrays (after masking NaNs).""" + y_ref = np.asarray(y_ref, dtype=float).ravel() + y_pred = np.asarray(y_pred, dtype=float).ravel() + mask = np.isfinite(y_ref) & np.isfinite(y_pred) + if mask.sum() == 0: + return np.nan + return np.sqrt(np.mean((y_ref[mask] - y_pred[mask]) ** 2)) + + +def mae(y_ref, y_pred): + """MAE between two arrays (after masking NaNs).""" + y_ref = np.asarray(y_ref, dtype=float).ravel() + y_pred = np.asarray(y_pred, dtype=float).ravel() + mask = np.isfinite(y_ref) & np.isfinite(y_pred) + if mask.sum() == 0: + return np.nan + return np.mean(np.abs(y_ref[mask] - y_pred[mask])) + + +def relative_rmse(y_ref, y_pred, eps=1e-12): + """Relative RMSE = RMSE / (mean(|y_ref|) + eps).""" + r = rmse(y_ref, y_pred) + y_ref = np.asarray(y_ref, dtype=float).ravel() + mask = np.isfinite(y_ref) + if mask.sum() == 0: + return np.nan + denom = np.nanmean(np.abs(y_ref[mask])) + eps + return r / denom if denom > 0 else np.nan + + +def relative_mae(y_ref, y_pred, eps=1e-12): + """Relative MAE = MAE / (mean(|y_ref|) + eps).""" + m = mae(y_ref, y_pred) + y_ref = np.asarray(y_ref, dtype=float).ravel() + mask = np.isfinite(y_ref) + if mask.sum() == 0: + return np.nan + denom = np.nanmean(np.abs(y_ref[mask])) + eps + return m / denom if denom > 0 else np.nan + + +def classify(r2, rel_rmse, rel_mae): + """Return 'Good', 'OK', or 'Bad' using same thresholds as quality_summary_report.""" + if np.isnan(r2) or np.isnan(rel_rmse) or np.isnan(rel_mae): + return "Bad" + if r2 >= THRESH_GOOD["r2_min"] and rel_rmse <= THRESH_GOOD["rel_rmse_max"] and rel_mae <= THRESH_GOOD["rel_mae_max"]: + return "Good" + if r2 >= THRESH_OK["r2_min"] and rel_rmse <= THRESH_OK["rel_rmse_max"] and rel_mae <= THRESH_OK["rel_mae_max"]: + return "OK" + return "Bad" + + +def main(): + parser = argparse.ArgumentParser( + description="Compare PFT 1D (PFT1–PFT16 only) and soil 2D (first column, first 10 layers) between two Amazon-site restart files." + ) + parser.add_argument( + "--restart-a", + required=True, + help="First Amazon-site restart NetCDF (e.g. ELM_data/Amazon_...elm.r.0801-01-01-00000.nc)", + ) + parser.add_argument( + "--restart-b", + required=True, + help="Second Amazon-site restart NetCDF (e.g. .../Amazon_phase2_tropical_4p_restart.nc)", + ) + parser.add_argument( + "--variable-list", + required=True, + help="CNP_IO file (e.g. CNP_IO_updated9_dev_dw.txt)", + ) + parser.add_argument("--output", default=None, help="Write report to this file (default: stdout)") + parser.add_argument("--plot-dir", default=None, help="Base directory to save plots: bad/, good/, ok/ (A vs B by PFT or layer)") + parser.add_argument("--rtol", type=float, default=1e-5, help="Relative tolerance for allclose") + parser.add_argument("--atol", type=float, default=1e-8, help="Absolute tolerance for allclose") + args = parser.parse_args() + + try: + import xarray as xr + except ImportError: + print("Need xarray: pip install xarray", file=sys.stderr) + sys.exit(1) + + parsed = parse_cnp_io_list(args.variable_list) + pft_1d_vars = parsed.get("pft_1d_variables", []) + soil_2d_vars = parsed.get("variables_2d_soil", []) + if not pft_1d_vars and not soil_2d_vars: + print("No pft_1d_variables or variables_2d_soil found in variable list.", file=sys.stderr) + sys.exit(1) + + out = open(args.output, "w") if args.output else sys.stdout + + def log(s=""): + print(s, file=out) + + log("PFT 1D + Soil 2D variable comparison (Amazon site restarts)") + log(" Restart A: " + args.restart_a) + log(" Restart B: " + args.restart_b) + log(" Variable list: " + args.variable_list) + log(" Scope: PFT1–PFT16 only (first column); soil 2D: first column, first 10 layers") + log("") + + ds_a = xr.open_dataset(args.restart_a) + ds_b = xr.open_dataset(args.restart_b) + + try: + n_pft_a = int(ds_a.sizes.get("pft", 0)) + n_pft_b = int(ds_b.sizes.get("pft", 0)) + log(" Restart A: pft dimension = {}".format(n_pft_a)) + log(" Restart B: pft dimension = {}".format(n_pft_b)) + if n_pft_a != n_pft_b: + log(" WARNING: PFT count mismatch. Comparison uses min(n_a, n_b).") + log("") + + # Table header: Variable | R2 | RMSE | MAE | RelRMSE | RelMAE | Class | Match | Max|diff| + log("Variable | R2 | RMSE | MAE | RelRMSE | RelMAE | Class | Match | Max|diff|") + log("-" * 95) + + results = [] # list of (var, r2, rmse_val, mae_val, rel_rmse, rel_mae, cls, match, max_diff, var_type, va, vb) + + for v in pft_1d_vars: + in_a = v in ds_a.variables + in_b = v in ds_b.variables + if not in_a: + log("{} | - | - | - | - | - | - | - | - | missing in A".format(v)) + continue + if not in_b: + log("{} | - | - | - | - | - | - | - | - | missing in B".format(v)) + continue + + var_a = ds_a[v] + var_b = ds_b[v] + if "pft" not in var_a.dims: + log("{} | - | - | - | - | - | - | - | - | not 1D pft in A".format(v)) + continue + if "pft" not in var_b.dims: + log("{} | - | - | - | - | - | - | - | - | not 1D pft in B".format(v)) + continue + + vals_a = np.asarray(var_a.values).ravel() + vals_b = np.asarray(var_b.values).ravel() + # PFT1–PFT16 only (0-based indices 1:17) + pft_end = min(PFT_END_1BASED + 1, len(vals_a), len(vals_b)) # 17 if available + if pft_end <= PFT_START_1BASED: + log("{} | - | - | - | - | - | - | - | - | need at least {} pft entries".format(v, PFT_END_1BASED + 1)) + continue + va = vals_a[PFT_START_1BASED:pft_end].astype(float) + vb = vals_b[PFT_START_1BASED:pft_end].astype(float) + + r2 = r2_score(va, vb) + rmse_val = rmse(va, vb) + mae_val = mae(va, vb) + rel_rmse_val = relative_rmse(va, vb) + rel_mae_val = relative_mae(va, vb) + cls = classify(r2, rel_rmse_val, rel_mae_val) + try: + match = np.allclose(va, vb, rtol=args.rtol, atol=args.atol, equal_nan=True) + except Exception: + match = False + max_diff = np.nanmax(np.abs(va - vb)) + + results.append((v, r2, rmse_val, mae_val, rel_rmse_val, rel_mae_val, cls, match, max_diff, "pft_1d", va, vb)) + log("{} | {:.4f} | {:.6g} | {:.6g} | {:.4f} | {:.4f} | {} | {} | {:.6g}".format( + v, r2, rmse_val, mae_val, rel_rmse_val, rel_mae_val, cls, match, max_diff)) + + # Soil 2D: first column, first 10 layers only + for v in soil_2d_vars: + in_a = v in ds_a.variables + in_b = v in ds_b.variables + if not in_a: + log("{} | - | - | - | - | - | - | - | - | missing in A (2D)".format(v)) + continue + if not in_b: + log("{} | - | - | - | - | - | - | - | - | missing in B (2D)".format(v)) + continue + + var_a = ds_a[v] + var_b = ds_b[v] + dims_a = list(var_a.dims) + dims_b = list(var_b.dims) + if "column" not in dims_a or "levgrnd" not in dims_a: + log("{} | - | - | - | - | - | - | - | - | not (column, levgrnd) in A".format(v)) + continue + if "column" not in dims_b or "levgrnd" not in dims_b: + log("{} | - | - | - | - | - | - | - | - | not (column, levgrnd) in B".format(v)) + continue + + try: + slab_a = var_a.isel(column=SOIL_COLUMN_INDEX, levgrnd=slice(0, SOIL_N_LAYERS)) + slab_b = var_b.isel(column=SOIL_COLUMN_INDEX, levgrnd=slice(0, SOIL_N_LAYERS)) + except Exception as e: + log("{} | - | - | - | - | - | - | - | - | isel failed: {}".format(v, e)) + continue + va = np.asarray(slab_a.values, dtype=float).ravel() + vb = np.asarray(slab_b.values, dtype=float).ravel() + if len(va) != len(vb) or len(va) == 0: + log("{} | - | - | - | - | - | - | - | - | length mismatch or zero (2D)".format(v)) + continue + + r2 = r2_score(va, vb) + rmse_val = rmse(va, vb) + mae_val = mae(va, vb) + rel_rmse_val = relative_rmse(va, vb) + rel_mae_val = relative_mae(va, vb) + cls = classify(r2, rel_rmse_val, rel_mae_val) + try: + match = np.allclose(va, vb, rtol=args.rtol, atol=args.atol, equal_nan=True) + except Exception: + match = False + max_diff = np.nanmax(np.abs(va - vb)) + + results.append((v, r2, rmse_val, mae_val, rel_rmse_val, rel_mae_val, cls, match, max_diff, "soil_2d", va, vb)) + log("{} | {:.4f} | {:.6g} | {:.6g} | {:.4f} | {:.4f} | {} | {} | {:.6g}".format( + v, r2, rmse_val, mae_val, rel_rmse_val, rel_mae_val, cls, match, max_diff)) + + log("") + # List Bad variables + bad_results = [r for r in results if r[6] == "Bad"] + if bad_results: + log("## Bad variables (below OK thresholds)") + for r in bad_results: + log(" {} (R2={:.4f}, RMSE={:.6g}, RelRMSE={:.4f})".format(r[0], r[1], r[2], r[4])) + log("") + + # Plot variables by class if --plot-dir set (bad/, good/, ok/ subfolders) + if args.plot_dir and results: + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ImportError: + print("matplotlib required for --plot-dir; skipping plots.", file=sys.stderr) + else: + base_plot_dir = Path(args.plot_dir) + base_plot_dir.mkdir(parents=True, exist_ok=True) + for class_name, class_label in [("Bad", "bad"), ("Good", "good"), ("OK", "ok")]: + subset = [r for r in results if r[6] == class_name] + if not subset: + continue + out_dir = base_plot_dir / class_label + out_dir.mkdir(parents=True, exist_ok=True) + for r in subset: + var_name, var_type, va, vb = r[0], r[9], np.asarray(r[10]), np.asarray(r[11]) + n = len(va) + if n == 0: + continue + fig, ax = plt.subplots() + if var_type == "pft_1d": + x = np.arange(PFT_START_1BASED, PFT_START_1BASED + n, dtype=float) + xlabel = "PFT (1-based)" + else: + x = np.arange(0, n, dtype=float) + xlabel = "Soil layer (0-based)" + ax.plot(x, va, "o-", label="Restart A (ref)", color="C0", markersize=4) + ax.plot(x, vb, "s-", label="Restart B", color="C1", markersize=4) + ax.set_xlabel(xlabel) + ax.set_ylabel(var_name) + ax.set_title("{} R2={:.4f} RMSE={:.4g}".format(var_name, r[1], r[2])) + ax.legend() + ax.grid(True, alpha=0.3) + safe_name = var_name.replace("/", "_") + fig.savefig(out_dir / "{}.png".format(safe_name), dpi=120, bbox_inches="tight") + plt.close(fig) + log("Plots saved under: {} (bad/, good/, ok/)".format(base_plot_dir)) + log("") + + log("") + # Summary section (similar to quality_summary_report.txt) + log("# Summary (Restart A = reference, Restart B = comparison)") + log("") + n_compared = len(results) + n_good = sum(1 for r in results if r[6] == "Good") + n_ok = sum(1 for r in results if r[6] == "OK") + n_bad = sum(1 for r in results if r[6] == "Bad") + log("## Overall Statistics") + log("Variables compared: {}".format(n_compared)) + log("Good: {} ({:.1f}%)".format(n_good, 100.0 * n_good / n_compared if n_compared else 0)) + log("OK: {} ({:.1f}%)".format(n_ok, 100.0 * n_ok / n_compared if n_compared else 0)) + log("Bad: {} ({:.1f}%)".format(n_bad, 100.0 * n_bad / n_compared if n_compared else 0)) + log("") + log("## Classification Thresholds Used") + log("Good: R2 >= {:.1f}, Relative RMSE <= {:.1f}, Relative MAE <= {:.1f}".format( + THRESH_GOOD["r2_min"], THRESH_GOOD["rel_rmse_max"], THRESH_GOOD["rel_mae_max"])) + log("OK: R2 >= {:.1f}, Relative RMSE <= {:.2f}, Relative MAE <= {:.2f}".format( + THRESH_OK["r2_min"], THRESH_OK["rel_rmse_max"], THRESH_OK["rel_mae_max"])) + log("Bad: Below OK thresholds") + log("") + if results: + by_r2 = sorted(results, key=lambda x: (np.nan_to_num(x[1], nan=-1), -np.nan_to_num(x[2], nan=np.inf)), reverse=True) + log("## Variables with Best Agreement (by R2, then lower RMSE)") + for r in by_r2[:10]: + log(" {}: R2={:.4f}, RMSE={:.6g}, Class={}".format(r[0], r[1], r[2], r[6])) + log("") + log("## Variables with Worst Agreement (lowest R2, worst first)") + for r in reversed(by_r2[-10:]): + log(" {}: R2={:.4f}, RMSE={:.6g}, Class={}".format(r[0], r[1], r[2], r[6])) + log("") + log("Done.") + finally: + ds_a.close() + ds_b.close() + if args.output: + out.close() + + +if __name__ == "__main__": + main() diff --git a/scripts/generate_prediction_quality_report.py b/scripts/generate_prediction_quality_report.py index 4cc3370..0fabc69 100644 --- a/scripts/generate_prediction_quality_report.py +++ b/scripts/generate_prediction_quality_report.py @@ -1,1068 +1,62 @@ #!/usr/bin/env python3 -import pandas as pd -import numpy as np -import matplotlib.pyplot as plt -import seaborn as sns -from pathlib import Path -import argparse -import importlib.util +""" +Generate prediction quality report and top_bad_plots. + +Runs docs/generate_prediction_quality_report.py with the same arguments, then +runs cnp_result_validationplot.py with --worst-only to produce analysis/top_bad_plots/ +(gt vs pred scatter plots for worst variables and for npool/ppool). +""" +import os +import subprocess import sys -import json -from typing import Optional, List +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +DOCS_SCRIPT = REPO_ROOT / "docs" / "generate_prediction_quality_report.py" + def main(): - parser = argparse.ArgumentParser(description='Generate prediction quality report from validation statistics') - parser.add_argument('input', nargs='?', default="./validation_stats.csv", - help='Path to validation statistics CSV file (default: ./validation_stats.csv)') - parser.add_argument('--output-dir', default=None, - help='Directory to save output files (default: same directory as input + /analysis)') - parser.add_argument('--training-config', default=None, - help='Path to training cnp_config.json (default: auto-detect near input)') - parser.add_argument('--r2-good', type=float, default=0.9, - help='R² threshold for good predictions (default: 0.9)') - parser.add_argument('--r2-ok', type=float, default=0.7, - help='R² threshold for ok predictions (default: 0.7)') - parser.add_argument('--rmse-good', type=float, default=0.1, - help='Relative RMSE threshold for good predictions (default: 0.1)') - parser.add_argument('--rmse-ok', type=float, default=0.25, - help='Relative RMSE threshold for ok predictions (default: 0.25)') - parser.add_argument('--mae-good', type=float, default=0.1, - help='Relative MAE threshold for good predictions (default: 0.1)') - parser.add_argument('--mae-ok', type=float, default=0.25, - help='Relative MAE threshold for ok predictions (default: 0.25)') - # Diagnostics and display options (default: enabled); provide --no-* to disable - parser.add_argument('--force-xlim-01', dest='force_xlim_01', action='store_true', - help='Force R² x-axis limits to [0, 1] in the scatter plot', default=True) - parser.add_argument('--no-force-xlim-01', dest='force_xlim_01', action='store_false', - help='Do not force R² x-axis limits to [0, 1]') - parser.add_argument('--print-scatter-stats', dest='print_scatter_stats', action='store_true', - help='Print min/max and counts for R² and relative RMSE used in the scatter plot', default=True) - parser.add_argument('--no-print-scatter-stats', dest='print_scatter_stats', action='store_false', - help='Disable printing diagnostics for R² and relative RMSE used in the scatter plot') - parser.add_argument('--bad-html-limit', type=int, default=100, - help='Maximum number of bad prediction rows to show in the HTML report (default: 100)') - parser.add_argument('--bad-text-limit', type=int, default=200, - help='Maximum number of bad prediction rows to print in the text report (default: 200)') - parser.add_argument('--export-bad', dest='export_bad', action='store_true', default=True, - help='Export detailed bad predictions to CSV (default: enabled)') - parser.add_argument('--no-export-bad', dest='export_bad', action='store_false', - help='Disable exporting detailed bad predictions to CSV') - parser.add_argument('--include-bad-details-text', dest='include_bad_details_text', action='store_true', default=False, - help='Include the long detailed list of bad predictions in the text report (default: disabled)') - parser.add_argument('--no-include-bad-details-text', dest='include_bad_details_text', action='store_false', - help='Do not include the long detailed list of bad predictions in the text report') - parser.add_argument('--top-bad-plots', dest='top_bad_plots', action='store_true', default=True, - help='Generate plots for top-bad variables (default: enabled)') - parser.add_argument('--no-top-bad-plots', dest='top_bad_plots', action='store_false', - help='Disable generating top-bad plots') - # Plot only variables listed in "Variables with Worst Predictions" - parser.add_argument('--worst-only', dest='worst_only', action='store_true', default=False, - help='Plot only variables in the "Variables with Worst Predictions" section') - parser.add_argument('--worst-filter-bad-only', dest='worst_filter_bad_only', action='store_true', default=True, - help='Filter worst variables to only include those with bad predictions (bad_pct > 0). Default: True') - parser.add_argument('--no-worst-filter-bad-only', dest='worst_filter_bad_only', action='store_false', - help='Include all variables in worst list, even if 0%% bad (sorted by good_pct)') - parser.add_argument('--worst-min-good-pct', type=float, default=None, - help='Include variables in worst list with good_pct below this threshold (e.g., 50.0 for <50%% good)') - parser.add_argument('--worst-vars-list', type=str, default=None, - help='Comma-separated list of specific variables to include in worst list (e.g., "cpool,npool,ppool")') - args = parser.parse_args() - - # Set up input and output paths - # Resolve relative input against the current working directory + if not DOCS_SCRIPT.exists(): + print(f"Error: {DOCS_SCRIPT} not found.", file=sys.stderr) + sys.exit(1) + + # 1) Run the docs quality report (same argv) + r = subprocess.run([sys.executable, str(DOCS_SCRIPT)] + sys.argv[1:]) + if r.returncode != 0: + sys.exit(r.returncode) + + # 2) Resolve output_dir the same way as the docs script + import argparse + p = argparse.ArgumentParser() + p.add_argument("--input", default="./validation_stats.csv") + p.add_argument("--output-dir", default=None) + p.add_argument("--npool-ppool-per-pft-json", default=None) + args, _ = p.parse_known_args(sys.argv[1:]) input_path = Path(args.input) if not input_path.is_absolute(): input_path = input_path.resolve() - - if args.output_dir is None: - output_dir = input_path.parent / "analysis" - else: - output_dir = Path(args.output_dir) - if not output_dir.is_absolute(): - # Resolve relative output against the current working directory - output_dir = output_dir.resolve() - - output_dir.mkdir(parents=True, exist_ok=True) - - # Helper: auto-detect training config (cnp_config.json) near the run directory - def _auto_detect_training_config(start_dir: Path) -> Optional[Path]: - for parent in [start_dir] + list(start_dir.parents): - cfg = parent / 'cnp_config.json' - if cfg.exists(): - return cfg - return None - - # Load expected variables from training configuration to ensure full coverage in reports - expected_vars: List[str] = [] - training_cfg_path: Optional[Path] = None - try: - training_cfg_path = Path(args.training_config) if args.training_config else _auto_detect_training_config(input_path.parent) - if training_cfg_path and training_cfg_path.exists(): - with open(training_cfg_path, 'r') as f: - cfg = json.load(f) - di = cfg.get('data_info', {}) if isinstance(cfg, dict) else {} - # Prefer scalar target list if available; fallback to input scalar list - scalar_targets = di.get('y_list_scalar_columns', []) or di.get('x_list_scalar_columns', []) or [] - # Strip any leading Y_ prefixes - scalar_vars = [str(v)[2:] if isinstance(v, str) and v.startswith('Y_') else str(v) for v in scalar_targets] - pft1d_vars = [str(v) for v in di.get('variables_1d_pft', []) or []] - soil2d_vars = [str(v) for v in di.get('x_list_columns_2d', []) or []] - expected_vars = list(dict.fromkeys(scalar_vars + pft1d_vars + soil2d_vars)) - if expected_vars: - print(f"Loaded {len(expected_vars)} expected variables from training config: {training_cfg_path}") - else: - print("Warning: Could not locate cnp_config.json to derive full variable list. Proceeding with variables present in stats.") - except Exception as e: - print(f"Warning: Failed to parse training config for expected variables: {e}") - - # Define thresholds for categorization - thresholds = { - 'good': { - 'r2': args.r2_good, - 'rmse_rel': args.rmse_good, - 'mae_rel': args.mae_good, - }, - 'ok': { - 'r2': args.r2_ok, - 'rmse_rel': args.rmse_ok, - 'mae_rel': args.mae_ok, - } - } - - print(f"Reading validation statistics from {input_path}") - df = pd.read_csv(input_path) - - # Normalize variable naming: strip leading 'Y_' from variable names (targets) - if 'variable' in df.columns: - try: - df['variable'] = df['variable'].apply(lambda v: v[2:] if isinstance(v, str) and v.startswith('Y_') else v) - except Exception: - pass - - # Function to categorize prediction quality - def categorize_prediction(row): - # Calculate relative metrics (normalized by data range) - gt_range = row['gt_max'] - row['gt_min'] - - # Handle zero range (constant values) - if gt_range == 0: - if row['rmse'] == 0 and row['mae'] == 0: - return 'good' # Perfect prediction for constant values - else: - return 'bad' # Any error on constant values is bad - - rmse_rel = row['rmse'] / gt_range - mae_rel = row['mae'] / gt_range - - # Apply thresholds for categorization - if (row['r2'] >= thresholds['good']['r2'] and - rmse_rel <= thresholds['good']['rmse_rel'] and - mae_rel <= thresholds['good']['mae_rel']): - return 'good' - elif (row['r2'] >= thresholds['ok']['r2'] and - rmse_rel <= thresholds['ok']['rmse_rel'] and - mae_rel <= thresholds['ok']['mae_rel']): - return 'ok' - else: - return 'bad' - - # Add a quality category column - print("Categorizing predictions...") - df['prediction_quality'] = df.apply(categorize_prediction, axis=1) - - # Filter out rows that are just coordinates (Longitude, Latitude) - # Be robust to files without a 'pft' column - if 'pft' not in df.columns: - df['pft'] = '' - coord_labels = {'Longitude', 'Latitude'} - mask_pft = ~df['pft'].isin(coord_labels) if 'pft' in df.columns else True - mask_var = ~df['variable'].isin(coord_labels) if 'variable' in df.columns else True - analysis_df = df[mask_pft & mask_var].copy() - - # Pre-compute relative errors for later exports/reports - analysis_df['gt_range'] = analysis_df['gt_max'] - analysis_df['gt_min'] - analysis_df['rmse_rel'] = np.where(analysis_df['gt_range'] > 0, - analysis_df['rmse'] / analysis_df['gt_range'], - np.nan) - analysis_df['mae_rel'] = np.where(analysis_df['gt_range'] > 0, - analysis_df['mae'] / analysis_df['gt_range'], - np.nan) - - # Create summary by variable - variable_summary = analysis_df.groupby(['variable', 'prediction_quality']).size().unstack(fill_value=0) - - # Ensure all expected variables appear in the summary (even if missing from CSV) - if expected_vars: - # Add any missing variables as zero rows - for v in expected_vars: - if v not in variable_summary.index: - variable_summary.loc[v, :] = 0 - - # Calculate percentages - variable_summary['total'] = variable_summary.sum(axis=1) - for category in ['good', 'ok', 'bad']: - if category in variable_summary.columns: - variable_summary[f'{category}_pct'] = (variable_summary[category] / variable_summary['total'] * 100).round(1) - - # Sort by percentage of good predictions - if 'good_pct' in variable_summary.columns: - variable_summary = variable_summary.sort_values(by='good_pct', ascending=False) - - # Save the detailed results - print(f"Saving detailed quality assessment to {output_dir / 'detailed_quality_assessment.csv'}") - df.to_csv(output_dir / "detailed_quality_assessment.csv", index=False) - - # Save detailed bad predictions - bad_df = analysis_df[analysis_df['prediction_quality'] == 'bad'].copy() - if args.export_bad and not bad_df.empty: - bad_csv_path = output_dir / "bad_predictions_detailed.csv" - bad_df.to_csv(bad_csv_path, index=False) - print(f"Saved detailed bad predictions to: {bad_csv_path}") - - # Save the variable summary - print(f"Saving variable quality summary to {output_dir / 'variable_quality_summary.csv'}") - variable_summary.to_csv(output_dir / "variable_quality_summary.csv") - - # Generate visualizations - print("Generating visualizations...") - - # 1. Stacked bar chart of prediction quality by variable - plt.figure(figsize=(14, 10)) - pivot_df = analysis_df.pivot_table( - index='variable', - columns='prediction_quality', - aggfunc='size', - fill_value=0 - ) - - # Calculate percentages for the chart - pivot_total = pivot_df.sum(axis=1) - for col in pivot_df.columns: - pivot_df[col] = (pivot_df[col] / pivot_total * 100).round(1) - - # Ensure all expected variables appear in the chart - if expected_vars: - for v in expected_vars: - if v not in pivot_df.index: - pivot_df.loc[v, :] = 0 - - # Sort by 'good' percentage if it exists - if 'good' in pivot_df.columns: - pivot_df = pivot_df.sort_values(by='good', ascending=False) - - # Set color map - colors = {'good': '#2ecc71', 'ok': '#f39c12', 'bad': '#e74c3c'} - color_list = [colors.get(x, 'gray') for x in pivot_df.columns] - - # Plot the stacked bar chart - ax = pivot_df.plot(kind='bar', stacked=True, figsize=(14, 10), color=color_list) - plt.title('Prediction Quality by Variable', fontsize=16) - plt.xlabel('Variable', fontsize=14) - plt.ylabel('Percentage (%)', fontsize=14) - plt.xticks(rotation=90) - plt.legend(title='Quality') - plt.tight_layout() - plt.savefig(output_dir / "prediction_quality_by_variable.png", dpi=300) - - # 2. Pie chart of overall prediction quality - plt.figure(figsize=(8, 8)) - quality_counts = analysis_df['prediction_quality'].value_counts() - plt.pie(quality_counts, labels=quality_counts.index, autopct='%1.1f%%', - colors=[colors.get(x, 'gray') for x in quality_counts.index], - explode=[0.05 if x == 'bad' else 0 for x in quality_counts.index]) - plt.title('Overall Prediction Quality Distribution', fontsize=16) - plt.tight_layout() - plt.savefig(output_dir / "overall_prediction_quality.png", dpi=300) - - # 3. Scatter plot of R² vs Relative RMSE for all predictions - plt.figure(figsize=(12, 10)) - - # Create a copy of the dataframe to avoid SettingWithCopyWarning - scatter_df = analysis_df.copy() - # Ensure relative RMSE exists (it does from pre-compute; keep guard for safety) - if 'rmse_rel' not in scatter_df.columns: - scatter_df['rmse_rel'] = scatter_df.apply( - lambda row: row['rmse'] / (row['gt_max'] - row['gt_min']) if row['gt_max'] > row['gt_min'] else 0, - axis=1 - ) - - # Optional diagnostics about what will be plotted - if args.print_scatter_stats: - r2_vals = scatter_df['r2'].replace([np.inf, -np.inf], np.nan).dropna() - rmse_rel_vals = scatter_df['rmse_rel'].replace([np.inf, -np.inf], np.nan).dropna() - total_points = len(scatter_df) - valid_r2 = len(r2_vals) - valid_rmse_rel = len(rmse_rel_vals) - print(f"Scatter diagnostics: total_points={total_points}, valid_r2={valid_r2}, valid_rmse_rel={valid_rmse_rel}") - if valid_r2 > 0: - print(f" R²: min={r2_vals.min():.6f}, max={r2_vals.max():.6f}, count_>0={(r2_vals > 0).sum()}, count_>=0={(r2_vals >= 0).sum()}") - if valid_rmse_rel > 0: - print(f" RMSE_rel: min={rmse_rel_vals.min():.6f}, max={rmse_rel_vals.max():.6f}") - - # Create scatter plot - scatter = plt.scatter( - scatter_df['r2'], - scatter_df['rmse_rel'], - c=scatter_df['prediction_quality'].map({'good': 0, 'ok': 1, 'bad': 2}), - cmap=plt.cm.viridis, - alpha=0.7, - s=50 - ) - - # Add threshold lines - plt.axhline(y=thresholds['good']['rmse_rel'], color='green', linestyle='--', alpha=0.7) - plt.axhline(y=thresholds['ok']['rmse_rel'], color='orange', linestyle='--', alpha=0.7) - plt.axvline(x=thresholds['good']['r2'], color='green', linestyle='--', alpha=0.7) - plt.axvline(x=thresholds['ok']['r2'], color='orange', linestyle='--', alpha=0.7) - - # Add labels and legend - plt.xlabel('R²', fontsize=14) - plt.ylabel('Relative RMSE (RMSE / Range)', fontsize=14) - plt.title('R² vs Relative RMSE for All Predictions', fontsize=16) - - # Optionally force x-axis limits for clarity - if args.force_xlim_01: - plt.xlim(0, 1) - - # Create custom legend - from matplotlib.lines import Line2D - legend_elements = [ - Line2D([0], [0], marker='o', color='w', markerfacecolor=plt.cm.viridis(0), markersize=10, label='Good'), - Line2D([0], [0], marker='o', color='w', markerfacecolor=plt.cm.viridis(0.5), markersize=10, label='OK'), - Line2D([0], [0], marker='o', color='w', markerfacecolor=plt.cm.viridis(1.0), markersize=10, label='Bad'), + output_dir = (Path(args.output_dir).resolve() if args.output_dir else input_path.parent / "analysis") + + report_path = output_dir / "quality_summary_report.txt" + if not report_path.exists(): + print("Warning: quality_summary_report.txt not found; skipping top_bad_plots.", file=sys.stderr) + return + + results_dir = str(output_dir.parent.resolve()) + plot_script = REPO_ROOT / "scripts" / "cnp_result_validationplot.py" + + # 3) Generate top_bad_plots (worst variables + npool/ppool) into analysis/top_bad_plots + cmd = [ + sys.executable, str(plot_script), + results_dir, + "--worst-only", + "--top-bad-report", str(report_path), + "--plots-dir", "analysis/top_bad_plots", ] - plt.legend(handles=legend_elements) - - plt.grid(True, alpha=0.3) - plt.tight_layout() - plt.savefig(output_dir / "r2_vs_rmse.png", dpi=300) - - # 4. Optionally generate restricted plots (top-bad or worst-only) into a subfolder using the validation plotting utility - top_bad_plot_count = 0 - if args.top_bad_plots: - try: - results_dir = str(input_path.parent) - top_bad_out = str((output_dir / 'top_bad_plots').resolve()) - (output_dir / 'top_bad_plots').mkdir(parents=True, exist_ok=True) - - # Pre-write "Top variables by bad-count" to the report so the plot script can read it - # (the full report is written later; this ensures top_bad selection is available when plotting) - if not args.worst_only and not bad_df.empty: - try: - report_path = output_dir / 'quality_summary_report.txt' - with open(report_path, 'w') as _pref: - _pref.write("# Prediction Quality Summary Report\n\n") - _pref.write("Top variables by bad-count (with PFT indices or layer numbers):\n") - bad_by_var = bad_df.groupby('variable').size().sort_values(ascending=False).head(25) - for v, c in bad_by_var.items(): - sub = bad_df[bad_df['variable'] == v] - pft_indices = [] - for val in sub['pft'].dropna().unique(): - if isinstance(val, str) and 'pft' in val: - try: - idx = ''.join(ch for ch in val.split('pft')[-1] if ch.isdigit()) - if idx: - pft_indices.append(int(idx)) - except Exception: - continue - pft_indices = sorted(set(pft_indices)) - layer_numbers = [] - for lay in sub['layer'].dropna().unique(): - try: - li = int(lay) if float(lay).is_integer() else float(lay) - layer_numbers.append(li) - except Exception: - continue - layer_numbers = sorted(set(layer_numbers)) - details_parts = [] - if pft_indices: - details_parts.append("pfts: " + ", ".join(str(i) for i in pft_indices)) - if layer_numbers: - details_parts.append("layers: " + ", ".join(str(i) for i in layer_numbers)) - details = ("; " + " ".join(details_parts)) if details_parts else "" - _pref.write(f" {v}: {c}{details}\n") - _pref.write("\n") - print(f"Wrote top-bad section for plot selection: {report_path}") - except Exception as _e: - print(f"Warning: Failed to pre-write top-bad section: {_e}") - - # If worst-only requested, pre-write a minimal 'Variables with Worst Predictions' section - if args.worst_only: - try: - tmp_report_path = output_dir / 'quality_summary_report.txt' - with open(tmp_report_path, 'w') as _pref: - _pref.write("# Prediction Quality Summary Report\n\n") - _pref.write("## Variables with Worst Predictions\n") - _vw = variable_summary.copy() - if 'good_pct' in _vw.columns: - _vw['good_pct'] = _vw['good_pct'].fillna(0) - _worst = _vw.nsmallest(15, 'good_pct') - for _var_name, _row in _worst.iterrows(): - _pref.write(f"{_var_name}: {_row.get('good_pct', 0):.1f}% good, {_row.get('ok_pct', 0):.1f}% ok, {_row.get('bad_pct', 0):.1f}% bad\n") - _pref.write("\n") - print(f"Wrote minimal worst-variables section for selection: {tmp_report_path}") - except Exception as _e: - print(f"Warning: Failed to pre-write worst-variables section for plotting selection: {_e}") - - # Protect the input validation_stats.csv from being overwritten by the plotting utility - original_bytes = None - try: - if input_path.exists(): - original_bytes = input_path.read_bytes() - except Exception: - original_bytes = None - # Dynamically import cnp_result_validationplot without relying on PYTHONPATH - plot_mod_path = (output_dir.parent.parent / 'scripts' / 'cnp_result_validationplot.py') - # If running from repo root, construct direct path as fallback - if not plot_mod_path.exists(): - plot_mod_path = Path(__file__).parent / 'cnp_result_validationplot.py' - spec = importlib.util.spec_from_file_location('cnp_plot_mod', str(plot_mod_path)) - mod = importlib.util.module_from_spec(spec) - sys.modules['cnp_plot_mod'] = mod - assert spec.loader is not None - try: - spec.loader.exec_module(mod) - if hasattr(mod, 'main_with_flag'): - mod.main_with_flag(results_dir, plot_scatter=True, plot_loss=False, - top_bad_only=not args.worst_only, - top_bad_report=str(output_dir / 'quality_summary_report.txt'), - plots_dir_override=top_bad_out, - worst_only=args.worst_only) - print(f"Top-bad plots saved to: {top_bad_out}") - try: - # Count the number of PNGs generated for quick reporting - top_bad_plot_count = len(list((output_dir / 'top_bad_plots').rglob('*.png'))) - print(f"Top-bad plot count: {top_bad_plot_count}") - except Exception: - top_bad_plot_count = 0 - else: - print("Warning: cnp_result_validationplot.main_with_flag not found; skipping top-bad plots") - finally: - # Restore original validation_stats.csv to prevent any overwrite - try: - if original_bytes is not None: - with open(input_path, 'wb') as _f: - _f.write(original_bytes) - print(f"Restored original validation_stats.csv after generating top-bad plots: {input_path}") - except Exception as _e: - print(f"Warning: Failed to restore original validation_stats.csv: {_e}") - except Exception as e: - print(f"Warning: Failed to generate top-bad plots: {e}") - - # Generate a comprehensive summary report - print(f"Generating summary report to {output_dir / 'quality_summary_report.txt'}") - with open(output_dir / "quality_summary_report.txt", "w") as f: - f.write("# Prediction Quality Summary Report\n\n") - - # Overall statistics - total_predictions = len(analysis_df) - # Variables analyzed (unique variable names in stats; if expected list provided, report both) - analyzed_variables = sorted(set(analysis_df['variable'].unique())) - num_analyzed_variables = len(analyzed_variables) - total_expected_variables = len(expected_vars) if expected_vars else None - good_count = analysis_df[analysis_df['prediction_quality'] == 'good'].shape[0] - ok_count = analysis_df[analysis_df['prediction_quality'] == 'ok'].shape[0] - bad_count = analysis_df[analysis_df['prediction_quality'] == 'bad'].shape[0] - - f.write(f"## Overall Statistics\n") - f.write(f"Variables analyzed: {num_analyzed_variables}") - if total_expected_variables is not None: - f.write(f" (of {total_expected_variables} expected from training config)") - f.write("\n") - if args.top_bad_plots: - f.write(f"Top-bad plots generated: {top_bad_plot_count}\n") - f.write(f"Total predictions analyzed: {total_predictions}\n") - f.write(f"Good predictions: {good_count} ({good_count/total_predictions*100:.1f}%)\n") - f.write(f"OK predictions: {ok_count} ({ok_count/total_predictions*100:.1f}%)\n") - f.write(f"Bad predictions: {bad_count} ({bad_count/total_predictions*100:.1f}%)\n\n") - - f.write("## Classification Thresholds Used\n") - f.write(f"Good: R² ≥ {thresholds['good']['r2']}, Relative RMSE ≤ {thresholds['good']['rmse_rel']}, Relative MAE ≤ {thresholds['good']['mae_rel']}\n") - f.write(f"OK: R² ≥ {thresholds['ok']['r2']}, Relative RMSE ≤ {thresholds['ok']['rmse_rel']}, Relative MAE ≤ {thresholds['ok']['mae_rel']}\n") - f.write(f"Bad: Below OK thresholds\n\n") - - # Bad predictions summary and details - f.write("## Bad Predictions Summary\n") - if bad_df.empty: - f.write("No bad predictions found.\n\n") - else: - # Counts by type - f.write("Bad predictions by type:\n") - bad_by_type = bad_df.groupby('type').size().sort_values(ascending=False) - for t, c in bad_by_type.items(): - f.write(f" {t}: {c}\n") - f.write("\nTop variables by bad-count (with PFT indices or layer numbers):\n") - bad_by_var = bad_df.groupby('variable').size().sort_values(ascending=False).head(20) - for v, c in bad_by_var.items(): - sub = bad_df[bad_df['variable'] == v] - # Collect pft indices if present (1D); parse trailing digits after 'pft' - pft_indices = [] - for val in sub['pft'].dropna().unique(): - if isinstance(val, str) and 'pft' in val: - try: - idx = ''.join(ch for ch in val.split('pft')[-1] if ch.isdigit()) - if idx: - pft_indices.append(int(idx)) - except Exception: - continue - pft_indices = sorted(set(pft_indices)) - # Collect layer numbers if present (2D) - layer_numbers = [] - for lay in sub['layer'].dropna().unique(): - try: - # cast to int if integral - li = int(lay) if float(lay).is_integer() else float(lay) - layer_numbers.append(li) - except Exception: - continue - layer_numbers = sorted(set(layer_numbers)) - - details_parts = [] - if pft_indices: - details_parts.append("pfts: " + ", ".join(str(i) for i in pft_indices)) - if layer_numbers: - details_parts.append("layers: " + ", ".join(str(i) for i in layer_numbers)) - details = ("; " + " ".join(details_parts)) if details_parts else "" - f.write(f" {v}: {c}{details}\n") - f.write("\n") - - # Optional: long detailed rows (disabled by default) - if args.include_bad_details_text: - f.write(f"## Detailed Bad Predictions (first {args.bad_text_limit})\n") - printable = bad_df.copy() - # Order by worst first: lowest R², then highest relative RMSE - printable = printable.sort_values(by=['r2','rmse_rel'], ascending=[True, False]) - if len(printable) > args.bad_text_limit: - printable = printable.head(args.bad_text_limit) - for _, row in printable.iterrows(): - f.write( - f"- {row.get('type','')}, {row.get('variable','')}, {row.get('pft','')}, layer={row.get('layer','')}" - f", r2={row.get('r2',np.nan):.6f}, rmse_rel={row.get('rmse_rel',np.nan):.6f}, " - f"mae_rel={row.get('mae_rel',np.nan):.6f}, rmse={row.get('rmse',np.nan):.6f}, mae={row.get('mae',np.nan):.6f}\n" - ) - f.write("\n") - if args.export_bad: - f.write("Full list saved to bad_predictions_detailed.csv\n\n") - - f.write("## Variables with Best Predictions\n") - if 'good_pct' in variable_summary.columns: - # Treat NaN as 0 for ranking - _vs = variable_summary.copy() - _vs['good_pct'] = _vs['good_pct'].fillna(0) - _vs['ok_pct'] = _vs.get('ok_pct', 0) - _vs['bad_pct'] = _vs.get('bad_pct', 0) - best_vars = _vs.nlargest(15, 'good_pct') - for var_name, row in best_vars.iterrows(): - f.write(f"{var_name}: {row.get('good_pct', 0):.1f}% good, {row.get('ok_pct', 0):.1f}% ok, {row.get('bad_pct', 0):.1f}% bad\n") - - f.write("\n## Variables with Worst Predictions\n") - if 'good_pct' in variable_summary.columns: - _vs2 = variable_summary.copy() - _vs2['good_pct'] = _vs2['good_pct'].fillna(0) - if 'ok_pct' in _vs2.columns: - _vs2['ok_pct'] = _vs2['ok_pct'].fillna(0) - else: - _vs2['ok_pct'] = 0 - if 'bad_pct' in _vs2.columns: - _vs2['bad_pct'] = _vs2['bad_pct'].fillna(0) - else: - _vs2['bad_pct'] = 0 - - # Apply filtering based on user options - _vs2_filtered = _vs2.copy() - - # Filter 1: By default, only include variables with bad predictions (bad_pct > 0) - if args.worst_filter_bad_only: - if 'bad_pct' in _vs2_filtered.columns: - _vs2_filtered = _vs2_filtered[_vs2_filtered['bad_pct'] > 0].copy() - - # Filter 2: Include variables with good_pct below threshold if specified - if args.worst_min_good_pct is not None: - _vs2_filtered = _vs2_filtered[_vs2_filtered['good_pct'] < args.worst_min_good_pct].copy() - - # Filter 3: Include specific variables if list provided - if args.worst_vars_list: - var_list = [v.strip() for v in args.worst_vars_list.split(',')] - # Add specified variables even if they don't meet other filters - specified_vars = _vs2[_vs2.index.isin(var_list)].copy() - _vs2_filtered = pd.concat([_vs2_filtered, specified_vars]).drop_duplicates() - - # Sort by bad_pct (highest first), then by good_pct (lowest first) for tie-breaking - if 'bad_pct' in _vs2_filtered.columns and len(_vs2_filtered) > 0: - worst_vars = _vs2_filtered.nlargest(15, 'bad_pct').nsmallest(15, 'good_pct') - elif len(_vs2_filtered) > 0: - worst_vars = _vs2_filtered.nsmallest(15, 'good_pct') - else: - worst_vars = pd.DataFrame() - - for var_name, row in worst_vars.iterrows(): - f.write(f"{var_name}: {row.get('good_pct', 0):.1f}% good, {row.get('ok_pct', 0):.1f}% ok, {row.get('bad_pct', 0):.1f}% bad\n") - - # Group all variables by dominant quality category - f.write("\n## All Variables Grouped by Quality Category\n") - if 'good_pct' in variable_summary.columns: - _vs3 = variable_summary.copy() - _vs3['good_pct'] = _vs3['good_pct'].fillna(0) - if 'ok_pct' in _vs3.columns: - _vs3['ok_pct'] = _vs3['ok_pct'].fillna(0) - else: - _vs3['ok_pct'] = 0 - if 'bad_pct' in _vs3.columns: - _vs3['bad_pct'] = _vs3['bad_pct'].fillna(0) - else: - _vs3['bad_pct'] = 0 - - # Determine dominant category for each variable (highest percentage) - def get_dominant_category(row): - good_pct = row['good_pct'] if 'good_pct' in row.index else 0 - ok_pct = row['ok_pct'] if 'ok_pct' in row.index else 0 - bad_pct = row['bad_pct'] if 'bad_pct' in row.index else 0 - if good_pct >= ok_pct and good_pct >= bad_pct: - return 'good' - elif ok_pct >= bad_pct: - return 'ok' - else: - return 'bad' - - _vs3['dominant_category'] = _vs3.apply(get_dominant_category, axis=1) - - # Group variables by category - good_vars = _vs3[_vs3['dominant_category'] == 'good'].sort_values('good_pct', ascending=False) - ok_vars = _vs3[_vs3['dominant_category'] == 'ok'].sort_values('ok_pct', ascending=False) - bad_vars = _vs3[_vs3['dominant_category'] == 'bad'].sort_values('bad_pct', ascending=False) - - f.write(f"\n### Good Variables ({len(good_vars)} total)\n") - if len(good_vars) > 0: - for var_name, row in good_vars.iterrows(): - good_val = row['good_pct'] if 'good_pct' in row.index else 0 - ok_val = row['ok_pct'] if 'ok_pct' in row.index else 0 - bad_val = row['bad_pct'] if 'bad_pct' in row.index else 0 - f.write(f"{var_name}: {good_val:.1f}% good, {ok_val:.1f}% ok, {bad_val:.1f}% bad\n") - else: - f.write("None\n") - - f.write(f"\n### OK Variables ({len(ok_vars)} total)\n") - if len(ok_vars) > 0: - for var_name, row in ok_vars.iterrows(): - good_val = row['good_pct'] if 'good_pct' in row.index else 0 - ok_val = row['ok_pct'] if 'ok_pct' in row.index else 0 - bad_val = row['bad_pct'] if 'bad_pct' in row.index else 0 - f.write(f"{var_name}: {good_val:.1f}% good, {ok_val:.1f}% ok, {bad_val:.1f}% bad\n") - else: - f.write("None\n") - - f.write(f"\n### Bad Variables ({len(bad_vars)} total)\n") - if len(bad_vars) > 0: - for var_name, row in bad_vars.iterrows(): - good_val = row['good_pct'] if 'good_pct' in row.index else 0 - ok_val = row['ok_pct'] if 'ok_pct' in row.index else 0 - bad_val = row['bad_pct'] if 'bad_pct' in row.index else 0 - f.write(f"{var_name}: {good_val:.1f}% good, {ok_val:.1f}% ok, {bad_val:.1f}% bad\n") - else: - f.write("None\n") - - # Report variables missing from the stats but present in training - if expected_vars: - present_vars = set(analysis_df['variable'].unique()) - missing_vars = [v for v in expected_vars if v not in present_vars] - f.write("\n## Variables Missing from validation_stats.csv (listed in training config)\n") - if missing_vars: - f.write(f"Count: {len(missing_vars)}\n") - # Limit long lists in text to keep report concise - preview = missing_vars[:100] - f.write("" + ", ".join(preview) + (" ..." if len(missing_vars) > 100 else "") + "\n") - else: - f.write("None\n") - - # Generate an HTML report for better visualization - print(f"Generating HTML report to {output_dir / 'prediction_quality_report.html'}") - - # Create HTML content - html_content = f""" - - - - Prediction Quality Report - - - -
-

Prediction Quality Report

- -
-

Overall Statistics

-
-
-

Good Predictions

-

{good_count} ({good_count/total_predictions*100:.1f}%)

-
-
-

OK Predictions

-

{ok_count} ({ok_count/total_predictions*100:.1f}%)

-
-
-

Bad Predictions

-

{bad_count} ({bad_count/total_predictions*100:.1f}%)

-
-
- -

Classification Thresholds

-
    -
  • Good: R² ≥ {thresholds['good']['r2']}, Relative RMSE ≤ {thresholds['good']['rmse_rel']}, Relative MAE ≤ {thresholds['good']['mae_rel']}
    • -
  • OK: R² ≥ {thresholds['ok']['r2']}, Relative RMSE ≤ {thresholds['ok']['rmse_rel']}, Relative MAE ≤ {thresholds['ok']['mae_rel']}
    • -
  • Bad: Below OK thresholds
    • -
-
- -
-

Visualization of Overall Results

- Overall Prediction Quality Distribution -
- -
-

Prediction Quality by Variable

- Prediction Quality by Variable -
- -
-

R² vs Relative RMSE

- R² vs Relative RMSE -
- -

Detailed Bad Predictions (first {args.bad_html_limit})

- - - - - - - - - - - - - """ + r2 = subprocess.run(cmd) + if r2.returncode != 0: + print("Warning: cnp_result_validationplot.py exited with code", r2.returncode, file=sys.stderr) - # Insert bad predictions table rows (limited) - if not bad_df.empty: - bad_html_rows = bad_df.copy().sort_values(by=['r2','rmse_rel'], ascending=[True, False]) - if len(bad_html_rows) > args.bad_html_limit: - bad_html_rows = bad_html_rows.head(args.bad_html_limit) - for _, row in bad_html_rows.iterrows(): - html_content += f""" - - - - - - - - - - - - """ - - # Close bad predictions table and proceed with the rest of the report - html_content += """ -
TypeVariablePFT/SoilLayerR²RMSE_relMAE_relRMSEMAE
{row.get('type','')}{row.get('variable','')}{row.get('pft','')}{row.get('layer','')}{row.get('r2',float('nan')):.6f}{row.get('rmse_rel',float('nan')):.6f}{row.get('mae_rel',float('nan')):.6f}{row.get('rmse',float('nan')):.6f}{row.get('mae',float('nan')):.6f}
- -

Best Performing Variables

- - - - - - - - """ - - # Add best variables - best_vars = variable_summary.nlargest(15, 'good_pct') - for var_name, row in best_vars.iterrows(): - html_content += f""" - - - - - - - """ - - html_content += """ -
VariableGood (%)OK (%)Bad (%)
{var_name}{row.get('good_pct', 0):.1f}%{row.get('ok_pct', 0):.1f}%{row.get('bad_pct', 0):.1f}%
- -

Worst Performing Variables

- - - - - - - - """ - - # Add worst variables - _vw = variable_summary.copy() - if 'good_pct' in _vw.columns: - _vw['good_pct'] = _vw['good_pct'].fillna(0) - worst_vars = _vw.nsmallest(15, 'good_pct') - for var_name, row in worst_vars.iterrows(): - html_content += f""" - - - - - - - """ - - # Add all variables grouped by quality category - html_content += """ -
VariableGood (%)OK (%)Bad (%)
{var_name}{row.get('good_pct', 0):.1f}%{row.get('ok_pct', 0):.1f}%{row.get('bad_pct', 0):.1f}%
- -

All Variables Grouped by Quality Category

- """ - - # Prepare grouped variables data - _vs_html = variable_summary.copy() - if 'good_pct' in _vs_html.columns: - _vs_html['good_pct'] = _vs_html['good_pct'].fillna(0) - if 'ok_pct' in _vs_html.columns: - _vs_html['ok_pct'] = _vs_html['ok_pct'].fillna(0) - else: - _vs_html['ok_pct'] = 0 - if 'bad_pct' in _vs_html.columns: - _vs_html['bad_pct'] = _vs_html['bad_pct'].fillna(0) - else: - _vs_html['bad_pct'] = 0 - - # Determine dominant category - def get_dominant_category_html(row): - good_pct = row['good_pct'] if 'good_pct' in row.index else 0 - ok_pct = row['ok_pct'] if 'ok_pct' in row.index else 0 - bad_pct = row['bad_pct'] if 'bad_pct' in row.index else 0 - if good_pct >= ok_pct and good_pct >= bad_pct: - return 'good' - elif ok_pct >= bad_pct: - return 'ok' - else: - return 'bad' - - _vs_html['dominant_category'] = _vs_html.apply(get_dominant_category_html, axis=1) - - good_vars_html = _vs_html[_vs_html['dominant_category'] == 'good'].sort_values('good_pct', ascending=False) - ok_vars_html = _vs_html[_vs_html['dominant_category'] == 'ok'].sort_values('ok_pct', ascending=False) - bad_vars_html = _vs_html[_vs_html['dominant_category'] == 'bad'].sort_values('bad_pct', ascending=False) - - # Good variables table - html_content += f""" -

Good Variables ({len(good_vars_html)} total)

- - - - - - - - """ - if len(good_vars_html) > 0: - for var_name, row in good_vars_html.iterrows(): - good_val = row['good_pct'] if 'good_pct' in row.index else 0 - ok_val = row['ok_pct'] if 'ok_pct' in row.index else 0 - bad_val = row['bad_pct'] if 'bad_pct' in row.index else 0 - html_content += f""" - - - - - - - """ - else: - html_content += "" - html_content += "
VariableGood (%)OK (%)Bad (%)
{var_name}{good_val:.1f}%{ok_val:.1f}%{bad_val:.1f}%
None
" - - # OK variables table - html_content += f""" -

OK Variables ({len(ok_vars_html)} total)

- - - - - - - - """ - if len(ok_vars_html) > 0: - for var_name, row in ok_vars_html.iterrows(): - good_val = row['good_pct'] if 'good_pct' in row.index else 0 - ok_val = row['ok_pct'] if 'ok_pct' in row.index else 0 - bad_val = row['bad_pct'] if 'bad_pct' in row.index else 0 - html_content += f""" - - - - - - - """ - else: - html_content += "" - html_content += "
VariableGood (%)OK (%)Bad (%)
{var_name}{good_val:.1f}%{ok_val:.1f}%{bad_val:.1f}%
None
" - - # Bad variables table - html_content += f""" -

Bad Variables ({len(bad_vars_html)} total)

- - - - - - - - """ - if len(bad_vars_html) > 0: - for var_name, row in bad_vars_html.iterrows(): - good_val = row['good_pct'] if 'good_pct' in row.index else 0 - ok_val = row['ok_pct'] if 'ok_pct' in row.index else 0 - bad_val = row['bad_pct'] if 'bad_pct' in row.index else 0 - html_content += f""" - - - - - - - """ - else: - html_content += "" - html_content += "
VariableGood (%)OK (%)Bad (%)
{var_name}{good_val:.1f}%{ok_val:.1f}%{bad_val:.1f}%
None
" - - # Add missing variables section if available - if expected_vars: - present_vars = set(analysis_df['variable'].unique()) - missing_vars = [v for v in expected_vars if v not in present_vars] - html_content += """ - -

Variables Missing from validation_stats.csv (in training config)

-
- """ - if missing_vars: - # Show as a comma-separated list (trim if very long) - preview = missing_vars[:300] - remainder = len(missing_vars) - len(preview) - html_content += f"

Count: {len(missing_vars)}

" - html_content += f"

{', '.join(preview)}{' ...' if remainder > 0 else ''}

" - else: - html_content += "

None

" - html_content += """ -
- """ - - html_content += """ - -
- - - """ - - # Write HTML file - with open(output_dir / "prediction_quality_report.html", "w") as f: - f.write(html_content) - - print("\nAnalysis complete. Results saved to", output_dir) - print("\nOverall Prediction Quality Summary:") - quality_counts = analysis_df['prediction_quality'].value_counts() - for quality, count in quality_counts.items(): - print(f"{quality}: {count} ({count/len(analysis_df)*100:.1f}%)") - - print("\nTop 5 Best Predicted Variables:") - _vb = variable_summary.copy() - if 'good_pct' in _vb.columns: - _vb['good_pct'] = _vb['good_pct'].fillna(0) - best_vars = _vb.nlargest(5, 'good_pct') - for var_name, row in best_vars.iterrows(): - print(f"{var_name}: {row.get('good_pct', 0):.1f}% good") - - print("\nTop 5 Worst Predicted Variables:") - _vw5 = variable_summary.copy() - if 'good_pct' in _vw5.columns: - _vw5['good_pct'] = _vw5['good_pct'].fillna(0) - worst_vars = _vw5.nsmallest(5, 'good_pct') - for var_name, row in worst_vars.iterrows(): - print(f"{var_name}: {row.get('good_pct', 0):.1f}% good, {row.get('bad_pct', 0):.1f}% bad") if __name__ == "__main__": main() diff --git a/scripts/generate_site_5p_restart_comparison.py b/scripts/generate_site_5p_restart_comparison.py new file mode 100644 index 0000000..6a5324a --- /dev/null +++ b/scripts/generate_site_5p_restart_comparison.py @@ -0,0 +1,318 @@ +import argparse +import os +from typing import Dict + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr + +# This script lives in the same directory as extract_elm_restart_point.py +from extract_elm_restart_point import extract_single_point_elm + + +# Fixed run directories and global restart paths (same as current workflow) +PHASE1_RUN_DIR = ( + "/mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260311_165538_phase1_global" +) +PHASE2_RUN_DIR = ( + "/mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260311_174843_phase2_tropical" +) +PHASE3_RUN_DIR = ( + "/mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260311_204845_phase3_tworegions" +) + +PHASE1_GLOBAL_RESTART = os.path.join(PHASE1_RUN_DIR, "updated_restart_base.nc") +PHASE2_GLOBAL_RESTART = os.path.join( + PHASE2_RUN_DIR, "updated_restart_phase2_tropical_5P_raw.nc" +) +PHASE3_GLOBAL_RESTART = os.path.join( + PHASE3_RUN_DIR, "updated_restart_phase3_tworegions_5P_bias_corrected_tropical.nc" +) + +# Ground-truth directory (per-variable CSVs ground_truth_Y_{var}.csv) +GT_PROFILES_DIR = os.path.join( + PHASE1_RUN_DIR, + "cnp_inference_entire_dataset", + "cnp_predictions", + "soil_2d_ground_truth", +) + +FIVE_P = ["labilep_vr", "occlp_vr", "solutionp_vr", "secondp_vr", "primp_vr"] +LEV_DIM = "levgrnd" + + +def _safe_site_name(site_name: str, lon: float, lat: float) -> str: + """Create a filesystem-safe site identifier.""" + if site_name: + base = site_name.strip().replace(" ", "_") + else: + base = "site" + return f"{base}_lon{lon:.2f}_lat{lat:.2f}".replace(".", "p").replace("-", "m") + + +def _find_row_for_site(df: pd.DataFrame, lon: float, lat: float, atol: float = 1e-4) -> int: + """Return row index matching lon/lat; fall back to nearest neighbour if needed.""" + if "Longitude" not in df.columns or "Latitude" not in df.columns: + raise ValueError("DataFrame must contain 'Longitude' and 'Latitude' columns.") + + lon_vals = df["Longitude"].values + lat_vals = df["Latitude"].values + mask = np.isclose(lon_vals, lon, atol=atol) & np.isclose(lat_vals, lat, atol=atol) + + idx = np.where(mask)[0] + if idx.size == 1: + return int(idx[0]) + + # Fallback: nearest neighbour + d2 = (lon_vals - lon) ** 2 + (lat_vals - lat) ** 2 + return int(np.argmin(d2)) + + +def _extract_profile_from_gt(df: pd.DataFrame, var: str, row_idx: int) -> np.ndarray: + """Extract vertical profile for one variable at row_idx from GT CSV.""" + prefix = f"Y_{var}_col1_layer" + layer_cols = sorted( + [c for c in df.columns if c.startswith(prefix)], + key=lambda c: int(c.split("layer")[-1]), + ) + if not layer_cols: + raise ValueError(f"No columns starting with '{prefix}' found.") + values = df.loc[row_idx, layer_cols].astype(float).values + return values + + +def _load_profiles_from_restart(path: str) -> Dict[str, np.ndarray]: + """Load 1D vertical profiles (top 10 layers) for the 5P variables from a single-point restart.""" + ds = xr.open_dataset(path, decode_times=False) + profs: Dict[str, np.ndarray] = {} + for var in FIVE_P: + if var not in ds: + print(f" WARNING: {var} not found in {path}") + continue + da = ds[var] + arr = da.values + if arr.ndim == 1: + prof = arr + elif arr.ndim == 2: + if da.dims[0] == LEV_DIM: + prof = arr[:, 0] + elif da.dims[1] == LEV_DIM: + prof = arr[0, :] + else: + prof = arr.reshape(-1) + else: + prof = arr.reshape(-1) + profs[var] = np.asarray(prof[:10], dtype=float) + return profs + + +def generate_site_restarts(lon: float, lat: float, site_name: str) -> Dict[str, str]: + """Extract phase1/phase2/phase3 single-point restarts for this site.""" + site_id = _safe_site_name(site_name, lon, lat) + + out_paths = { + "phase1_global": os.path.join(PHASE1_RUN_DIR, f"{site_id}_phase1_global_restart.nc"), + "phase2_tropical": os.path.join( + PHASE2_RUN_DIR, f"{site_id}_phase2_tropical_restart.nc" + ), + "phase3_tworegions": os.path.join( + PHASE3_RUN_DIR, f"{site_id}_phase3_tworegions_restart.nc" + ), + } + + print(f"\n=== Extracting single-point restarts for {site_id} ===") + extract_single_point_elm( + source_nc=PHASE1_GLOBAL_RESTART, + output_nc=out_paths["phase1_global"], + target_lat=lat, + target_lon=lon, + ) + extract_single_point_elm( + source_nc=PHASE2_GLOBAL_RESTART, + output_nc=out_paths["phase2_tropical"], + target_lat=lat, + target_lon=lon, + ) + extract_single_point_elm( + source_nc=PHASE3_GLOBAL_RESTART, + output_nc=out_paths["phase3_tworegions"], + target_lat=lat, + target_lon=lon, + ) + + return out_paths + + +def compare_5p_for_site( + lon: float, + lat: float, + site_name: str, + restart_paths: Dict[str, str], + output_dir: str, +) -> None: + """Compute and plot 5P profiles for GT + three restarts at the given site.""" + os.makedirs(output_dir, exist_ok=True) + site_id = _safe_site_name(site_name, lon, lat) + + # Load ground truth profiles + gt_profiles: Dict[str, np.ndarray] = {} + if os.path.isdir(GT_PROFILES_DIR): + print(f"\nLoading ground truth from {GT_PROFILES_DIR}") + for var in FIVE_P: + csv_path = os.path.join(GT_PROFILES_DIR, f"ground_truth_Y_{var}.csv") + if not os.path.exists(csv_path): + print(f" WARNING: GT CSV not found for {var}: {csv_path}") + continue + df_gt = pd.read_csv(csv_path) + row_idx = _find_row_for_site(df_gt, lon=lon, lat=lat, atol=1e-4) + prof = _extract_profile_from_gt(df_gt, var=var, row_idx=row_idx) + gt_profiles[var] = prof[:10] + else: + print( + f"WARNING: GT profiles directory not found at {GT_PROFILES_DIR}; " + "plots will omit ground truth." + ) + + # Load restart profiles + all_profiles: Dict[str, Dict[str, np.ndarray]] = {} + print("\nLoading site restarts:") + for name, path in restart_paths.items(): + print(f" {name}: {path}") + if not os.path.exists(path): + print(" WARNING: file does not exist, skipping.") + continue + all_profiles[name] = _load_profiles_from_restart(path) + + # Pairwise numeric diffs + print("\n=== Pairwise max |diff| between restarts (per 5P variable) ===") + restart_names = list(all_profiles.keys()) + for var in FIVE_P: + print(f"\nVariable: {var}") + for i in range(len(restart_names)): + for j in range(i + 1, len(restart_names)): + n1, n2 = restart_names[i], restart_names[j] + p1 = all_profiles[n1].get(var) + p2 = all_profiles[n2].get(var) + if p1 is None or p2 is None: + continue + nlev = min(len(p1), len(p2)) + diff = float(np.nanmax(np.abs(p1[:nlev] - p2[:nlev]))) + print(f" max |{n1} - {n2}| = {diff}") + + # Line plots + print(f"\nSaving plots to: {output_dir}") + for var in FIVE_P: + plt.figure(figsize=(5, 6)) + has_any = False + + # GT first + if var in gt_profiles: + gt = np.asarray(gt_profiles[var], dtype=float) + nlev_gt = min(10, len(gt)) + layers_gt = np.arange(1, nlev_gt + 1) + plt.plot( + layers_gt, + gt[:nlev_gt], + marker="o", + linestyle="-", + color="black", + label="ground_truth", + linewidth=2.0, + ) + has_any = True + + # Restarts + for name, profs in all_profiles.items(): + if var not in profs: + continue + prof = np.asarray(profs[var], dtype=float) + nlev = min(10, len(prof)) + layers = np.arange(1, nlev + 1) + plt.plot( + layers, + prof[:nlev], + marker="o", + linestyle="-", + label=name, + alpha=0.9, + ) + has_any = True + + if not has_any: + plt.close() + continue + + plt.xlabel("Soil layer index") + plt.ylabel(f"{var} value") + plt.title( + f"{var} vertical profile at {site_id}\n" + f"(lon={lon}, lat={lat})" + ) + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + + out_path = os.path.join(output_dir, f"{site_id}_restart_profile_{var}.png") + plt.savefig(out_path, dpi=200) + plt.close() + print(f" Wrote {out_path}") + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Given a lon/lat, extract single-point restarts for phase1_global, " + "phase2_tropical, phase3_tworegions and generate 5P vertical-profile " + "comparison plots vs ground truth." + ) + ) + parser.add_argument("--lon", type=float, required=True, help="Site longitude (degrees, 0–360).") + parser.add_argument("--lat", type=float, required=True, help="Site latitude (degrees).") + parser.add_argument( + "--site-name", + type=str, + default="site", + help="Optional human-readable site name used in filenames and plot titles.", + ) + parser.add_argument( + "--output-dir", + type=str, + default="", + help=( + "Output directory for plots. Default: " + f"{PHASE3_RUN_DIR}/analysis/_5p_restart_comparison" + ), + ) + args = parser.parse_args() + + lon = args.lon + lat = args.lat + site_name = args.site_name + site_id = _safe_site_name(site_name, lon, lat) + + if args.output_dir: + output_dir = os.path.abspath(args.output_dir) + else: + output_dir = os.path.join( + PHASE3_RUN_DIR, + "analysis", + f"{site_id}_5p_restart_comparison", + ) + + print(f"Site: {site_id} (lon={lon}, lat={lat})") + print(f"Output directory: {output_dir}") + + restart_paths = generate_site_restarts(lon=lon, lat=lat, site_name=site_name) + compare_5p_for_site( + lon=lon, + lat=lat, + site_name=site_name, + restart_paths=restart_paths, + output_dir=output_dir, + ) + + +if __name__ == "__main__": + main() + diff --git a/scripts/plot_excluded_locations.py b/scripts/plot_excluded_locations.py new file mode 100644 index 0000000..a3c2bed --- /dev/null +++ b/scripts/plot_excluded_locations.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +""" +Plot cpool, litr1c_vr, and other variables at gridcell locations excluded by the natveg filter +(PCT_NATVEG=0 or PCT_NAT_PFT_0=100). Reads test_static_inverse.csv to get excluded row indices, +then plots GT vs Pred at those rows only. +""" +import os +import argparse +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +from pathlib import Path + + +def get_excluded_indices(results_dir): + """Load test_static_inverse.csv and return boolean mask and indices of excluded gridcells.""" + path = Path(results_dir) / "cnp_predictions" / "test_static_inverse.csv" + if not path.exists(): + raise FileNotFoundError(f"Required file not found: {path}") + df = pd.read_csv(path) + pct_natveg = df["PCT_NATVEG"].values.astype(float) + pct_pft0 = df["PCT_NAT_PFT_0"].values.astype(float) + include = (pct_natveg > 0) & (pct_pft0 < 100) + excluded = ~include + return excluded, df + + +def drop_coords(df): + """Drop Longitude, Latitude (or variants) for numeric extraction.""" + for c in ["Longitude", "Latitude", "longitude", "latitude", "lon", "lat", "Long", "Lat"]: + if c in df.columns: + df = df.drop(columns=[c]) + return df + + +def plot_pft1d_at_excluded(results_dir, excluded_idx, out_dir, var_name="cpool"): + """Plot PFT 1D variable (e.g. cpool) at excluded locations: GT vs Pred scatter (all PFTs).""" + gt_path = Path(results_dir) / "cnp_predictions" / "pft_1d_ground_truth" / f"ground_truth_Y_{var_name}.csv" + pred_path = Path(results_dir) / "cnp_predictions" / "pft_1d_predictions" / f"predictions_Y_{var_name}.csv" + if not gt_path.exists() or not pred_path.exists(): + print(f" Skip {var_name}: files not found") + return + gt = drop_coords(pd.read_csv(gt_path)) + pred = drop_coords(pd.read_csv(pred_path)) + gt_vals = gt.iloc[excluded_idx].values.flatten() + pred_vals = pred.iloc[excluded_idx].values.flatten() + valid = ~(np.isnan(gt_vals) | np.isnan(pred_vals)) + gt_vals = gt_vals[valid] + pred_vals = pred_vals[valid] + if len(gt_vals) < 2: + print(f" Skip {var_name}: too few valid points at excluded locations") + return + fig, ax = plt.subplots(figsize=(6, 6)) + ax.scatter(gt_vals, pred_vals, alpha=0.7, s=40) + mn = min(gt_vals.min(), pred_vals.min()) + mx = max(gt_vals.max(), pred_vals.max()) + ax.plot([mn, mx], [mn, mx], "r--", label="1:1") + ax.set_xlabel("Ground truth") + ax.set_ylabel("Prediction") + ax.set_title(f"{var_name} at excluded locations (n={len(gt_vals)} points, 36 gridcells × PFTs)") + ax.legend() + ax.grid(True, alpha=0.3) + plt.tight_layout() + out_path = out_dir / f"excluded_locations_{var_name}_gt_vs_pred.png" + plt.savefig(out_path, dpi=150) + plt.close() + print(f" Saved {out_path}") + + +def plot_soil2d_at_excluded(results_dir, excluded_idx, out_dir, var_name="litr1c_vr"): + """Plot 2D soil variable at excluded locations: GT vs Pred scatter (all layers flattened).""" + gt_path = Path(results_dir) / "cnp_predictions" / "soil_2d_ground_truth" / f"ground_truth_Y_{var_name}.csv" + pred_path = Path(results_dir) / "cnp_predictions" / "soil_2d_predictions" / f"predictions_Y_{var_name}.csv" + if not gt_path.exists() or not pred_path.exists(): + print(f" Skip {var_name}: files not found") + return + gt = drop_coords(pd.read_csv(gt_path)) + pred = drop_coords(pd.read_csv(pred_path)) + gt_vals = gt.iloc[excluded_idx].values.flatten() + pred_vals = pred.iloc[excluded_idx].values.flatten() + valid = ~(np.isnan(gt_vals) | np.isnan(pred_vals)) + gt_vals = gt_vals[valid] + pred_vals = pred_vals[valid] + if len(gt_vals) < 2: + print(f" Skip {var_name}: too few valid points at excluded locations") + return + fig, ax = plt.subplots(figsize=(6, 6)) + ax.scatter(gt_vals, pred_vals, alpha=0.7, s=40) + mn = min(gt_vals.min(), pred_vals.min()) + mx = max(gt_vals.max(), pred_vals.max()) + ax.plot([mn, mx], [mn, mx], "r--", label="1:1") + ax.set_xlabel("Ground truth") + ax.set_ylabel("Prediction") + ax.set_title(f"{var_name} at excluded locations (n={len(gt_vals)} points, 36 gridcells × layers)") + ax.legend() + ax.grid(True, alpha=0.3) + plt.tight_layout() + out_path = out_dir / f"excluded_locations_{var_name}_gt_vs_pred.png" + plt.savefig(out_path, dpi=150) + plt.close() + print(f" Saved {out_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Plot variables at natveg-filter excluded locations") + parser.add_argument("results_dir", nargs="?", default=".", help="Results directory (e.g. cnp_results/run_xxx)") + parser.add_argument("--output-dir", "-o", default=None, help="Output directory (default: results_dir/analysis/excluded_locations)") + args = parser.parse_args() + results_dir = Path(args.results_dir).resolve() + out_dir = Path(args.output_dir).resolve() if args.output_dir else results_dir / "analysis" / "excluded_locations" + out_dir.mkdir(parents=True, exist_ok=True) + + excluded_mask, static_df = get_excluded_indices(results_dir) + excluded_idx = np.where(excluded_mask)[0] + n_excl = len(excluded_idx) + print(f"Excluded gridcells: {n_excl}") + print(f"Output directory: {out_dir}") + + # PFT 1D + for var in ["cpool", "npool", "ppool"]: + plot_pft1d_at_excluded(results_dir, excluded_idx, out_dir, var_name=var) + + # 2D soil (litr1*, soil1*, etc. that exist in this run) + for var in ["litr1c_vr", "litr1n_vr", "litr1p_vr", "soil1c_vr", "soil1n_vr", "soil1p_vr", "primp_vr"]: + plot_soil2d_at_excluded(results_dir, excluded_idx, out_dir, var_name=var) + + # Optional: save a small table of excluded (lat, lon) for reference + loc_path = out_dir / "excluded_locations_lat_lon.csv" + if "Latitude" in static_df.columns and "Longitude" in static_df.columns: + static_df.loc[excluded_mask, ["Latitude", "Longitude", "PCT_NATVEG", "PCT_NAT_PFT_0"]].to_csv(loc_path, index=False) + print(f" Saved {loc_path}") + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_finetuning.py b/scripts/run_finetuning.py index 570451e..814aca6 100644 --- a/scripts/run_finetuning.py +++ b/scripts/run_finetuning.py @@ -318,6 +318,18 @@ def apply_training_overrides( if "mask_absent_pfts" in config: train_cfg.mask_absent_pfts = bool(config["mask_absent_pfts"]) + # Tropical-only filtering: restrict data to latitude band before train/test split. + if "tropical_only" in config: + data_cfg.tropical_only = bool(config["tropical_only"]) + if "tropical_lat_range" in config: + raw = config["tropical_lat_range"] + if isinstance(raw, (list, tuple)) and len(raw) >= 2: + data_cfg.tropical_lat_range = (float(raw[0]), float(raw[1])) + elif isinstance(raw, str): + parts = [p.strip() for p in raw.split(",")] + if len(parts) >= 2: + data_cfg.tropical_lat_range = (float(parts[0]), float(parts[1])) + def load_training_config_from_checkpoint(model_path: Path) -> Optional[Dict[str, Any]]: """Discover and load `cnp_config.json` located near the checkpoint.""" @@ -545,6 +557,18 @@ def fine_tune(args: argparse.Namespace) -> Dict[str, Any]: logger.info("Run directory: %s", run_dir) attach_file_logger(run_dir) + # CLI overrides for tropical-only filtering (override config file if set). + if getattr(args, "tropical_only", False): + manager.data_config.tropical_only = True + if getattr(args, "tropical_lat_range", None): + parts = [p.strip() for p in args.tropical_lat_range.split(",")] + if len(parts) >= 2: + manager.data_config.tropical_lat_range = (float(parts[0]), float(parts[1])) + logger.info( + "Tropical-only fine-tuning: lat range %s", + manager.data_config.tropical_lat_range, + ) + # Load historical configuration from the checkpoint when available. prior_config = load_training_config_from_checkpoint(model_path) if prior_config: @@ -645,6 +669,18 @@ def parse_args() -> argparse.Namespace: choices=("group", "individual", "hybrid"), help="Override normalization strategy (otherwise use config setting).", ) + parser.add_argument( + "--tropical-only", + action="store_true", + help="Restrict fine-tuning data to tropical latitude band (overrides config file).", + ) + parser.add_argument( + "--tropical-lat-range", + type=str, + metavar="MIN,MAX", + default=None, + help='Latitude range for tropical filter, e.g. "-30,30" (default: -23.5,23.5). Used when --tropical-only is set.', + ) parser.add_argument( "--log-level", default="INFO", diff --git a/scripts/run_inference_all.py b/scripts/run_inference_all.py index 87d56f0..edc6602 100644 --- a/scripts/run_inference_all.py +++ b/scripts/run_inference_all.py @@ -230,10 +230,9 @@ def run_inference_all( loader: str = 'auto', mask_pft_with_gt: bool = False, mask_absent_pfts: bool = True, - derive_np_from_c: bool = False, - tropical_only: bool = False, - tropical_lat_range: tuple = None, - tropical_lat_column: str = None + derive_np_from_c: bool = True, + inference_full_grid: bool = False, + inference_two_regions_only: bool = False, ) -> Path: """Run inference with the trained CNP model over the entire dataset. @@ -255,6 +254,27 @@ def run_inference_all( torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False + # Resolve model path: training may save cnp_model.pt in run root or model.pth in cnp_predictions/ + model_path = str(Path(model_path).resolve()) + p = Path(model_path) + if not p.exists(): + candidates = [ + p.parent / "cnp_model.pt", + p.parent.parent / "cnp_model.pt", + p.parent / "cnp_predictions" / "model.pth", # when run from run dir with --model model.pth + p.parent.parent / "cnp_predictions" / "model.pth", + ] + for candidate in candidates: + if candidate.exists(): + model_path = str(candidate.resolve()) + logging.info(f"Model not found at original path; using {model_path}") + break + else: + raise ValueError( + f"Model file not found at {model_path}. " + "Tried cnp_model.pt and cnp_predictions/model.pth in same/parent dirs. Use --model path/to/cnp_model.pt or path/to/cnp_predictions/model.pth" + ) + # Handle variable list: either use provided CNP_IO file or auto-detect from training config variables = None @@ -400,6 +420,34 @@ def run_inference_all( config.data_config.file_pattern = data_cfg.get('file_pattern') if data_cfg.get('dataset_file_patterns'): config.data_config.dataset_file_patterns = data_cfg.get('dataset_file_patterns') + if 'natveg_only' in data_cfg: + config.data_config.natveg_only = bool(data_cfg['natveg_only']) + if 'natveg_filter_before_split' in data_cfg: + config.data_config.natveg_filter_before_split = bool(data_cfg['natveg_filter_before_split']) + if 'tropical_only' in data_cfg: + config.data_config.tropical_only = bool(data_cfg['tropical_only']) + logging.info(f"Using tropical_only from training config: {config.data_config.tropical_only}") + if 'tropical_lat_range' in data_cfg: + tr = data_cfg['tropical_lat_range'] + if isinstance(tr, (list, tuple)) and len(tr) >= 2: + config.data_config.tropical_lat_range = (float(tr[0]), float(tr[1])) + logging.info(f"Using tropical_lat_range from training config: {config.data_config.tropical_lat_range}") + if 'region_boxes' in data_cfg: + boxes = data_cfg['region_boxes'] + if isinstance(boxes, (list, tuple)) and len(boxes) > 0: + parsed = [] + for b in boxes: + if isinstance(b, (list, tuple)) and len(b) >= 4: + parsed.append((float(b[0]), float(b[1]), float(b[2]), float(b[3]))) + if parsed: + config.data_config.region_boxes = parsed + logging.info(f"Using region_boxes from training config: {len(parsed)} box(es) — inference will run on these regions only") + # Force full-grid inference when requested (e.g. Phase 2 trained tropical-only; merge needs all gridcells) + if inference_full_grid: + config.data_config.tropical_only = False + if hasattr(config.data_config, 'region_boxes'): + config.data_config.region_boxes = None + logging.info("Inference full grid: tropical_only=False, region_boxes cleared (use all gridcells)") elif data_paths is None: logging.warning( "Training run cnp_config.json has no data_config (or no data_paths). " @@ -423,7 +471,16 @@ def run_inference_all( logging.info(f"Using provided scalers directory: {scalers_dir}") else: scalers_dir = model_dir / 'scalers' - logging.info(f"Auto-detected scalers directory: {scalers_dir}") + # Training saves scalers under cnp_predictions/scalers; fallback if not next to model + if not scalers_dir.exists(): + alt = model_dir / 'cnp_predictions' / 'scalers' + if alt.exists(): + scalers_dir = alt + logging.info(f"Auto-detected scalers directory: {scalers_dir} (cnp_predictions/scalers)") + else: + logging.info(f"Auto-detected scalers directory: {scalers_dir}") + else: + logging.info(f"Auto-detected scalers directory: {scalers_dir}") uses_individual = False try: @@ -1746,11 +1803,12 @@ def main(): parser.add_argument("--no-mask-absent-pfts", dest="mask_absent_pfts", action="store_false", help="Disable masking of absent PFTs") parser.set_defaults(mask_absent_pfts=True) parser.add_argument("--refit-normalization", action='store_true', default=False, help="Refit scalers on inference data (default: False; use training scalers)") - parser.add_argument("--tropical-only", action='store_true', help="Filter dataset to tropical latitude band before inference") - parser.add_argument("--tropical-lat-range", type=str, default=None, help='Latitude range for tropical filter, format "min,max" (default: -23.5,23.5)') - parser.add_argument("--tropical-lat-column", type=str, default=None, help="Latitude column name override (default: auto-detect from static columns)") - parser.add_argument("--derive-np-from-c", action='store_true', default=False, - help="Enforce CNP stoichiometric ratios by deriving N/P variables from C predictions after inference (default: False)") + parser.add_argument("--derive-np-from-c", action='store_true', default=True, + help="Enforce CNP stoichiometric ratios by deriving N/P variables from C predictions after inference (default: True)") + parser.add_argument("--inference-full-grid", action='store_true', default=False, + help="Run inference on full global grid (set tropical_only=False). Use for Phase 2 tropical-trained models when merging P variables into a global restart; otherwise only validation/tropical gridcells would be in the predictions NetCDF.") + parser.add_argument("--inference-two-regions-only", action='store_true', default=False, + help="Run inference only on Amazon + Central Africa region boxes. Use for two-region finetuned models when you want predictions only in those regions.") args = parser.parse_args() # Setup logging @@ -1777,15 +1835,14 @@ def main(): model_config=args.model_config, scalers_dir=args.scalers_dir, use_training_config=args.use_training_config, - strict_loading=args.strict_loading - , debug_vars=args.debug_vars - , loader=args.loader - , mask_pft_with_gt=args.mask_pft_with_gt - , mask_absent_pfts=args.mask_absent_pfts - , derive_np_from_c=args.derive_np_from_c - , tropical_only=args.tropical_only - , tropical_lat_range=tropical_lat_range - , tropical_lat_column=args.tropical_lat_column + strict_loading=args.strict_loading, + debug_vars=args.debug_vars, + loader=args.loader, + mask_pft_with_gt=args.mask_pft_with_gt, + mask_absent_pfts=args.mask_absent_pfts, + derive_np_from_c=args.derive_np_from_c, + inference_full_grid=getattr(args, 'inference_full_grid', False), + inference_two_regions_only=getattr(args, 'inference_two_regions_only', False), ) print(f"Inference completed successfully. Results saved to: {output_path}") diff --git a/scripts/run_phase1_global.sh b/scripts/run_phase1_global.sh new file mode 100644 index 0000000..bac08f1 --- /dev/null +++ b/scripts/run_phase1_global.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +set -euo pipefail + +############################################################################### +# run_phase1_global.sh +# +# Phase 1 for new users: +# - Train a global natveg_improved-like model (if needed) +# - Create a base restart from that model using an ELM restart template +# +# Requirements: +# - Training data configured via training_config_experiment_3_global_natveg_improved.json +# - ELM restart template (RESTART_TEMPLATE) +# +# Environment (can be overridden): +# - CONFIG_GLOBAL training config JSON for global model +# - VARIABLE_LIST CNP_IO variable list file +# - RESTART_TEMPLATE path to original ELM restart .nc (required) +# +# Outputs: +# - NATVEG_RUN_DIR (exported): global model run directory +# - BASE_RESTART (exported): base restart .nc created from NATVEG_RUN_DIR +# +# Usage (from repo root): +# bash scripts/run_phase1_global.sh +############################################################################### + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT_DIR" + +CONFIG_GLOBAL="${CONFIG_GLOBAL:-config/training_config_experiment_3_global_natveg_improved.json}" +VARIABLE_LIST="${VARIABLE_LIST:-CNP_IO_updated9_dev_dw.txt}" + +if [[ -z "${RESTART_TEMPLATE:-}" ]] || [[ ! -f "$RESTART_TEMPLATE" ]]; then + echo "Error: RESTART_TEMPLATE must point to an existing ELM restart .nc file." >&2 + echo "Example:" >&2 + echo " export RESTART_TEMPLATE=/path/to/your/elm_restart_template.nc" >&2 + exit 1 +fi + +_has_model() { + local d="$1" + [[ -f "$d/cnp_model.pt" ]] || [[ -f "$d/cnp_predictions/model.pth" ]] +} + +_latest_run_with_model() { + local pattern="$1" + for d in $(ls -td $pattern 2>/dev/null); do + [[ -d "$d" ]] && _has_model "$d" && { echo "$d"; return 0; } + done + return 1 +} + +echo "=== Phase 1: Global model training + base restart creation ===" + +if [[ -n "${NATVEG_RUN_DIR:-}" ]]; then + if ! _has_model "$NATVEG_RUN_DIR"; then + echo "Error: NATVEG_RUN_DIR is set but no model found at $NATVEG_RUN_DIR" >&2 + exit 1 + fi + echo "Using existing global model at $NATVEG_RUN_DIR." +elif _latest_run_with_model "cnp_results/run_*_phase1_global" >/dev/null; then + NATVEG_RUN_DIR="$(_latest_run_with_model "cnp_results/run_*_phase1_global")" + echo "Reusing existing global model at $NATVEG_RUN_DIR." +else + echo "Training global (natveg_improved-like) model (may take ~30 min)..." + t0=$(date +%s) + python train_cnp_model.py \ + --training-config-json "$CONFIG_GLOBAL" \ + --output-dir cnp_results \ + --output-dir-suffix phase1_global \ + --variable-list "$VARIABLE_LIST" + t1=$(date +%s) + echo "Global training finished in $(( (t1 - t0) / 60 )) minutes." + NATVEG_RUN_DIR="$(ls -td cnp_results/run_*_phase1_global 2>/dev/null | head -1)" + if [[ -z "$NATVEG_RUN_DIR" ]] || [[ ! -d "$NATVEG_RUN_DIR" ]]; then + echo "Error: Could not find global model run directory (run_*_phase1_global)." >&2 + exit 1 + fi + echo "Global model run: $NATVEG_RUN_DIR" +fi + +INF_GLOBAL="$NATVEG_RUN_DIR/cnp_inference_entire_dataset" +mkdir -p "$INF_GLOBAL" + +echo "Running full-grid inference for global model..." +python scripts/run_inference_all.py \ + --model "$NATVEG_RUN_DIR/cnp_model.pt" \ + --output-dir "$INF_GLOBAL" \ + --variable-list "$VARIABLE_LIST" \ + --inference-full-grid + +echo "Converting global predictions to NetCDF..." +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions "$INF_GLOBAL/cnp_predictions" \ + --variable-list "$VARIABLE_LIST" \ + --output "$NATVEG_RUN_DIR/ai_predictions_global.nc" + +echo "Creating base restart from template + global predictions..." +python scripts/ai_predictions_to_restart.py \ + --ai-predictions "$NATVEG_RUN_DIR/ai_predictions_global.nc" \ + --restart-file "$RESTART_TEMPLATE" \ + --output "$NATVEG_RUN_DIR/updated_restart_base.nc" \ + --variable-list "$VARIABLE_LIST" \ + "--tropical-lat-range=-90,90" + +BASE_RESTART="$NATVEG_RUN_DIR/updated_restart_base.nc" +echo "Base restart created at: $BASE_RESTART" + +echo +echo "Phase 1 complete." +echo " NATVEG_RUN_DIR = $NATVEG_RUN_DIR" +echo " BASE_RESTART = $BASE_RESTART" + diff --git a/scripts/run_phase2_tropical.sh b/scripts/run_phase2_tropical.sh new file mode 100644 index 0000000..0012123 --- /dev/null +++ b/scripts/run_phase2_tropical.sh @@ -0,0 +1,128 @@ +#!/usr/bin/env bash +set -euo pipefail + +############################################################################### +# run_phase2_tropical.sh +# +# Phase 2 for new users: +# - Train a tropical Phase2 P-focused model (if needed) +# - Create a tropical restart with raw Phase2 5P overwriting base restart +# +# Requirements: +# - NATVEG_RUN_DIR and BASE_RESTART from phase 1 +# +# Env (can be overridden): +# - CONFIG_TROPICAL training config JSON for tropical model +# - VARIABLE_LIST CNP_IO variable list file +# +# Usage: +# export NATVEG_RUN_DIR=... # from run_phase1_global.sh +# export BASE_RESTART=... # from run_phase1_global.sh +# bash scripts/run_phase2_tropical.sh +############################################################################### + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT_DIR" + +# If not set, try to use latest phase1 run and its base restart +if [[ -z "${NATVEG_RUN_DIR:-}" ]] || [[ -z "${BASE_RESTART:-}" ]]; then + _latest_phase1() { + local d + for d in $(ls -td cnp_results/run_*_phase1_global 2>/dev/null); do + [[ -d "$d" ]] && [[ -f "$d/updated_restart_base.nc" ]] && { echo "$d"; return 0; } + done + return 1 + } + if _latest_phase1 >/dev/null; then + NATVEG_RUN_DIR="$(_latest_phase1)" + BASE_RESTART="$NATVEG_RUN_DIR/updated_restart_base.nc" + echo "Using phase 1 outputs: NATVEG_RUN_DIR=$NATVEG_RUN_DIR, BASE_RESTART=$BASE_RESTART" + else + echo "Error: NATVEG_RUN_DIR and BASE_RESTART must be set (from phase 1), or run phase 1 first." >&2 + echo " export NATVEG_RUN_DIR=cnp_results/run_XXXX_phase1_global" >&2 + echo " export BASE_RESTART=\$NATVEG_RUN_DIR/updated_restart_base.nc" >&2 + exit 1 + fi +fi + +CONFIG_TROPICAL="${CONFIG_TROPICAL:-config/training_config_phase2_tropical_soilp_only.json}" +VARIABLE_LIST="${VARIABLE_LIST:-CNP_IO_updated9_dev_dw.txt}" + +_has_model() { + local d="$1" + [[ -f "$d/cnp_model.pt" ]] || [[ -f "$d/cnp_predictions/model.pth" ]] +} + +_latest_run_with_model() { + local pattern="$1" + for d in $(ls -td $pattern 2>/dev/null); do + [[ -d "$d" ]] && _has_model "$d" && { echo "$d"; return 0; } + done + return 1 +} + +echo "=== Phase 2: Tropical model training + tropical restart (raw Phase2 5P) ===" + +if [[ -n "${PHASE2_RUN_DIR:-}" ]]; then + if ! _has_model "$PHASE2_RUN_DIR"; then + echo "Error: PHASE2_RUN_DIR is set but no model found at $PHASE2_RUN_DIR" >&2 + exit 1 + fi + echo "Using existing tropical model at $PHASE2_RUN_DIR." +elif _latest_run_with_model "cnp_results/run_*_phase2_tropical" >/dev/null; then + PHASE2_RUN_DIR="$(_latest_run_with_model "cnp_results/run_*_phase2_tropical")" + echo "Reusing existing tropical model at $PHASE2_RUN_DIR." +else + echo "Training tropical (phase2 P-focused) model (may take ~30 min)..." + t0=$(date +%s) + python train_cnp_model.py \ + --training-config-json "$CONFIG_TROPICAL" \ + --output-dir cnp_results \ + --output-dir-suffix phase2_tropical \ + --variable-list "$VARIABLE_LIST" + t1=$(date +%s) + echo "Tropical training finished in $(( (t1 - t0) / 60 )) minutes." + PHASE2_RUN_DIR="$(ls -td cnp_results/run_*_phase2_tropical 2>/dev/null | head -1)" + if [[ -z "$PHASE2_RUN_DIR" ]] || [[ ! -d "$PHASE2_RUN_DIR" ]]; then + echo "Error: Could not find tropical model run directory (run_*_phase2_tropical)." >&2 + exit 1 + fi + echo "Tropical model run: $PHASE2_RUN_DIR" +fi + +TROP_INF_DIR="$PHASE2_RUN_DIR/cnp_inference_tropical_only" +if [[ ! -d "$TROP_INF_DIR/cnp_predictions" ]]; then + echo "Running tropical-only inference..." + python scripts/run_inference_all.py \ + --model "$PHASE2_RUN_DIR/cnp_model.pt" \ + --output-dir "$TROP_INF_DIR" \ + --variable-list "$VARIABLE_LIST" +fi + +TROP_NETCDF="$PHASE2_RUN_DIR/comparison_results/ai_predictions_tropical_only.nc" +mkdir -p "$PHASE2_RUN_DIR/comparison_results" +if [[ ! -f "$TROP_NETCDF" ]]; then + echo "Converting tropical predictions to NetCDF..." + python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions "$TROP_INF_DIR/cnp_predictions" \ + --variable-list "$VARIABLE_LIST" \ + --output "$TROP_NETCDF" +fi + +echo "Creating phase2_tropical restart (raw Phase2 5P in tropics)..." +python scripts/ai_predictions_to_restart.py \ + --ai-predictions "$TROP_NETCDF" \ + --restart-file "$BASE_RESTART" \ + --output "$PHASE2_RUN_DIR/updated_restart_phase2_tropical_5P_raw.nc" \ + --variable-list "$VARIABLE_LIST" \ + --variables-to-update labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr \ + "--tropical-lat-range=-30,30" + +echo "phase2_tropical: natveg base + raw Phase2 5P in tropics on $(date)" > "$PHASE2_RUN_DIR/README_phase2_tropical.txt" + +echo + +echo "Phase 2 complete." +echo " PHASE2_RUN_DIR = $PHASE2_RUN_DIR" +echo " Tropical restart: $PHASE2_RUN_DIR/updated_restart_phase2_tropical_5P_raw.nc" + diff --git a/scripts/run_phase3_tworegions.sh b/scripts/run_phase3_tworegions.sh new file mode 100644 index 0000000..b953799 --- /dev/null +++ b/scripts/run_phase3_tworegions.sh @@ -0,0 +1,111 @@ +#!/usr/bin/env bash +set -euo pipefail + +############################################################################### +# run_phase3_tworegions.sh +# +# Phase 3 for new users: +# - Full-grid inference with Phase2 model +# - Apply 5P bias/scale correction (Amazon + Africa) +# - Create a tropical restart with bias-corrected Phase2 5P +# +# Requirements: +# - PHASE2_RUN_DIR and BASE_RESTART from previous phases +# +# Env (can be overridden): +# - VARIABLE_LIST CNP_IO variable list file +# - TS timestamp suffix (optional; only affects output dir name) +# +# Usage: +# export PHASE2_RUN_DIR=... +# export BASE_RESTART=... +# bash scripts/run_phase3_tworegions.sh +############################################################################### + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT_DIR" + +# If not set, try to use latest phase2 run and phase1 base restart +if [[ -z "${PHASE2_RUN_DIR:-}" ]] || [[ -z "${BASE_RESTART:-}" ]]; then + _has_model() { + local d="$1" + [[ -f "$d/cnp_model.pt" ]] || [[ -f "$d/cnp_predictions/model.pth" ]] + } + _latest_phase2() { + local d + for d in $(ls -td cnp_results/run_*_phase2_tropical 2>/dev/null); do + [[ -d "$d" ]] && _has_model "$d" && { echo "$d"; return 0; } + done + return 1 + } + _latest_phase1_restart() { + local d + for d in $(ls -td cnp_results/run_*_phase1_global 2>/dev/null); do + [[ -d "$d" ]] && [[ -f "$d/updated_restart_base.nc" ]] && { echo "$d/updated_restart_base.nc"; return 0; } + done + return 1 + } + if [[ -z "${PHASE2_RUN_DIR:-}" ]] && _latest_phase2 >/dev/null; then + PHASE2_RUN_DIR="$(_latest_phase2)" + echo "Using latest phase 2 run: PHASE2_RUN_DIR=$PHASE2_RUN_DIR" + fi + if [[ -z "${BASE_RESTART:-}" ]] && _latest_phase1_restart >/dev/null; then + BASE_RESTART="$(_latest_phase1_restart)" + echo "Using latest phase 1 base restart: BASE_RESTART=$BASE_RESTART" + fi + if [[ -z "${PHASE2_RUN_DIR:-}" ]] || [[ -z "${BASE_RESTART:-}" ]]; then + echo "Error: PHASE2_RUN_DIR and BASE_RESTART must be set (from earlier phases), or run phase 1 and phase 2 first." >&2 + echo " export PHASE2_RUN_DIR=cnp_results/run_XXXX_phase2_tropical" >&2 + echo " export BASE_RESTART=cnp_results/run_XXXX_phase1_global/updated_restart_base.nc" >&2 + exit 1 + fi +fi + +VARIABLE_LIST="${VARIABLE_LIST:-CNP_IO_updated9_dev_dw.txt}" +TS="${TS:-$(date +%Y%m%d_%H%M%S)}" +RUN3_DIR="cnp_results/run_${TS}_phase3_tworegions" +mkdir -p "$RUN3_DIR" + +INF_FULL_DIR="$RUN3_DIR/cnp_inference_entire_dataset" +if [[ ! -d "$INF_FULL_DIR/cnp_predictions" ]]; then + echo "Running full-grid Phase2 inference into $INF_FULL_DIR ..." + python scripts/run_inference_all.py \ + --model "$PHASE2_RUN_DIR/cnp_model.pt" \ + --output-dir "$INF_FULL_DIR" \ + --variable-list "$VARIABLE_LIST" \ + --inference-full-grid +fi + +echo "Applying 5P bias/scale correction (Amazon + Africa)..." +python scripts/apply_5p_bias_scale_correction.py \ + --run-dir "$RUN3_DIR" \ + --region-config-json config/training_config_two_region_five_p.json \ + --output-subdir soil_2d_predictions_5P_bias_corrected_phase2 + +mkdir -p "$RUN3_DIR/comparison_results" +BC_NETCDF="$RUN3_DIR/comparison_results/ai_predictions_5P_bias_corrected_phase2.nc" + +echo "Converting bias-corrected predictions to NetCDF..." +python scripts/ai_predictions_to_netcdf.py \ + --ai-predictions "$INF_FULL_DIR/cnp_predictions/" \ + --soil-2d-bias-corrected-subdir soil_2d_predictions_5P_bias_corrected_phase2 \ + --variable-list "$VARIABLE_LIST" \ + --output "$BC_NETCDF" + +echo "Creating phase3_tworegions restart (bias-corrected 5P in tropics)..." +python scripts/ai_predictions_to_restart.py \ + --ai-predictions "$BC_NETCDF" \ + --restart-file "$BASE_RESTART" \ + --output "$RUN3_DIR/updated_restart_phase3_tworegions_5P_bias_corrected_tropical.nc" \ + --variable-list "$VARIABLE_LIST" \ + --variables-to-update labilep_vr,occlp_vr,solutionp_vr,secondp_vr,primp_vr \ + "--tropical-lat-range=-30,30" + +echo "phase3_tworegions: natveg base + Phase2 5P with Amazon+Africa bias/scale in tropics on $(date)" > "$RUN3_DIR/README_phase3_tworegions.txt" + +echo + +echo "Phase 3 complete." +echo " RUN3_DIR = $RUN3_DIR" +echo " Phase3 restart = $RUN3_DIR/updated_restart_phase3_tworegions_5P_bias_corrected_tropical.nc" + diff --git a/scripts/validation_npool_ppool_exclude_special.py b/scripts/validation_npool_ppool_exclude_special.py new file mode 100644 index 0000000..8b58121 --- /dev/null +++ b/scripts/validation_npool_ppool_exclude_special.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +""" +Compute npool and ppool prediction quality excluding "special value" grid cells. + +Special-value cells: grid cells where ALL 16 PFTs have npool==10 and ppool==1 +(ground truth). Excluding them answers: what is R²/RMSE when we validate only +on the non-constant (non-background) grid cells? + +Usage: + python scripts/validation_npool_ppool_exclude_special.py cnp_results/run_YYYYMMDD_HHMMSS + python scripts/validation_npool_ppool_exclude_special.py cnp_results/run_20260225_013116_global_natveg --output report.json +""" + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + +REPO_ROOT = Path(__file__).resolve().parents[1] +NPOOL_SPECIAL = 10.0 +PPOOL_SPECIAL = 1.0 +TOL = 1e-5 + + +def _is_special_row_npool(df_gt: pd.DataFrame) -> np.ndarray: + """True for rows where all 16 PFT columns are == 10.""" + cols = [c for c in df_gt.columns if c.startswith("Y_npool_pft")] + if len(cols) != 16: + return np.zeros(len(df_gt), dtype=bool) + arr = df_gt[cols].values.astype(float) + return np.all(np.abs(arr - NPOOL_SPECIAL) < TOL, axis=1) + + +def _is_special_row_ppool(df_gt: pd.DataFrame) -> np.ndarray: + """True for rows where all 16 PFT columns are == 1.""" + cols = [c for c in df_gt.columns if c.startswith("Y_ppool_pft")] + if len(cols) != 16: + return np.zeros(len(df_gt), dtype=bool) + arr = df_gt[cols].values.astype(float) + return np.all(np.abs(arr - PPOOL_SPECIAL) < TOL, axis=1) + + +def _r2_rmse(y_true: np.ndarray, y_pred: np.ndarray, mask: np.ndarray = None): + if mask is not None: + y_true = y_true[mask] + y_pred = y_pred[mask] + y_true = np.asarray(y_true, dtype=np.float64).ravel() + y_pred = np.asarray(y_pred, dtype=np.float64).ravel() + finite = np.isfinite(y_true) & np.isfinite(y_pred) + y_true = y_true[finite] + y_pred = y_pred[finite] + n = len(y_true) + if n < 2: + return float("nan"), float("nan") + ss_res = np.sum((y_true - y_pred) ** 2) + ss_tot = np.sum((y_true - np.mean(y_true)) ** 2) + r2 = 1.0 - (ss_res / ss_tot) if ss_tot > 0 else float("nan") + rmse = np.sqrt(ss_res / n) + return r2, rmse + + +def run(run_dir: Path, output_path: Path = None): + run_dir = Path(run_dir) + pred_dir = run_dir / "cnp_predictions" + gt_npool_path = pred_dir / "pft_1d_ground_truth" / "ground_truth_Y_npool.csv" + gt_ppool_path = pred_dir / "pft_1d_ground_truth" / "ground_truth_Y_ppool.csv" + pred_npool_path = pred_dir / "pft_1d_predictions" / "predictions_Y_npool.csv" + pred_ppool_path = pred_dir / "pft_1d_predictions" / "predictions_Y_ppool.csv" + + for p in [gt_npool_path, gt_ppool_path, pred_npool_path, pred_ppool_path]: + if not p.exists(): + print(f"Missing: {p}", file=sys.stderr) + return None + + gt_npool = pd.read_csv(gt_npool_path) + gt_ppool = pd.read_csv(gt_ppool_path) + pred_npool = pd.read_csv(pred_npool_path) + pred_ppool = pd.read_csv(pred_ppool_path) + + n = len(gt_npool) + if n != len(gt_ppool) or n != len(pred_npool) or n != len(pred_ppool): + print("Row count mismatch between GT and predictions.", file=sys.stderr) + return None + + # Special-value mask: exclude rows where ALL PFTs are (npool==10 and ppool==1) + special_npool = _is_special_row_npool(gt_npool) + special_ppool = _is_special_row_ppool(gt_ppool) + special_both = special_npool & special_ppool + n_special = int(special_both.sum()) + n_keep = n - n_special + keep_mask = ~special_both + + # PFT columns + cols_npool = [c for c in gt_npool.columns if c.startswith("Y_npool_pft")] + cols_ppool = [c for c in gt_ppool.columns if c.startswith("Y_ppool_pft")] + if len(cols_npool) != 16 or len(cols_ppool) != 16: + print("Expected 16 PFT columns.", file=sys.stderr) + return None + + # Flatten to (n*16,) for overall metrics + gt_npool_flat = gt_npool[cols_npool].values.ravel() + pred_npool_flat = pred_npool[cols_npool].values.ravel() + gt_ppool_flat = gt_ppool[cols_ppool].values.ravel() + pred_ppool_flat = pred_ppool[cols_ppool].values.ravel() + + # Row-level mask expanded to (n*16,): row i contributes 16 elements + keep_flat_npool = np.repeat(keep_mask, 16) + keep_flat_ppool = np.repeat(keep_mask, 16) + + # Metrics: all cells vs excluding special + r2_npool_all, rmse_npool_all = _r2_rmse(gt_npool_flat, pred_npool_flat, None) + r2_npool_excl, rmse_npool_excl = _r2_rmse(gt_npool_flat, pred_npool_flat, keep_flat_npool) + r2_ppool_all, rmse_ppool_all = _r2_rmse(gt_ppool_flat, pred_ppool_flat, None) + r2_ppool_excl, rmse_ppool_excl = _r2_rmse(gt_ppool_flat, pred_ppool_flat, keep_flat_ppool) + + out = { + "run_dir": str(run_dir), + "n_test_cells": n, + "n_special_cells_excluded": n_special, + "n_cells_kept": n_keep, + "npool": { + "r2_all_cells": float(r2_npool_all), + "rmse_all_cells": float(rmse_npool_all), + "r2_excluding_special": float(r2_npool_excl), + "rmse_excluding_special": float(rmse_npool_excl), + }, + "ppool": { + "r2_all_cells": float(r2_ppool_all), + "rmse_all_cells": float(rmse_ppool_all), + "r2_excluding_special": float(r2_ppool_excl), + "rmse_excluding_special": float(rmse_ppool_excl), + }, + } + + # Print report + print() + print("=== NPOOL / PPOOL validation: excluding special-value grid cells ===") + print(f"Run: {run_dir}") + print(f"Test grid cells: {n}") + print(f"Special-value cells excluded (all PFTs npool==10 and ppool==1): {n_special} ({100*n_special/n:.1f}%)") + print(f"Cells kept for validation: {n_keep}") + print() + print("NPOOL:") + print(f" R² (all cells): {r2_npool_all:.4f}") + print(f" R² (excl. special): {r2_npool_excl:.4f}") + print(f" RMSE (all cells): {rmse_npool_all:.4f}") + print(f" RMSE (excl. special): {rmse_npool_excl:.4f}") + print() + print("PPOOL:") + print(f" R² (all cells): {r2_ppool_all:.4f}") + print(f" R² (excl. special): {r2_ppool_excl:.4f}") + print(f" RMSE (all cells): {rmse_ppool_all:.4f}") + print(f" RMSE (excl. special): {rmse_ppool_excl:.4f}") + + if output_path: + with open(output_path, "w") as f: + json.dump(out, f, indent=2) + print(f"\nReport written to {output_path}") + + return out + + +def main(): + ap = argparse.ArgumentParser(description="NPOOL/PPOOL validation excluding special-value grid cells.") + ap.add_argument("run_dir", type=str, help="Path to run directory (e.g. cnp_results/run_YYYYMMDD_HHMMSS)") + ap.add_argument("--output", "-o", type=str, help="Write JSON report to this file") + args = ap.parse_args() + run_path = Path(args.run_dir) + if not run_path.is_absolute(): + run_path = (REPO_ROOT / args.run_dir).resolve() + out_path = Path(args.output).resolve() if args.output else None + run(run_path, out_path) + + +if __name__ == "__main__": + main() diff --git a/scripts/validation_npool_ppool_per_pft.py b/scripts/validation_npool_ppool_per_pft.py new file mode 100644 index 0000000..a49e0ee --- /dev/null +++ b/scripts/validation_npool_ppool_per_pft.py @@ -0,0 +1,171 @@ +#!/usr/bin/env python3 +""" +Compute npool and ppool prediction quality per PFT, considering only grid cells +where that PFT is present (PCT_NAT_PFT_k > 0). + +Where a PFT has zero coverage, npool/ppool are typically the special values +(10 and 1); including those cells would distort the metric. So for npool PFT k +(and ppool PFT k), we validate only on grid cells with PCT_NAT_PFT_k > 0. + +Usage: + python scripts/validation_npool_ppool_per_pft.py cnp_results/run_YYYYMMDD_HHMMSS + python scripts/validation_npool_ppool_per_pft.py cnp_results/run_20260226_114546_nofilter --output report.json --pct-min 0 +""" + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _r2_rmse(y_true: np.ndarray, y_pred: np.ndarray): + y_true = np.asarray(y_true, dtype=np.float64).ravel() + y_pred = np.asarray(y_pred, dtype=np.float64).ravel() + finite = np.isfinite(y_true) & np.isfinite(y_pred) + y_true = y_true[finite] + y_pred = y_pred[finite] + n = len(y_true) + if n < 2: + return float("nan"), float("nan"), int(n) + ss_res = np.sum((y_true - y_pred) ** 2) + ss_tot = np.sum((y_true - np.mean(y_true)) ** 2) + r2 = 1.0 - (ss_res / ss_tot) if ss_tot > 0 else float("nan") + rmse = np.sqrt(ss_res / n) + return r2, rmse, int(n) + + +def run(run_dir: Path, pct_min: float = 0.0, output_path: Path = None): + run_dir = Path(run_dir) + pred_dir = run_dir / "cnp_predictions" + static_path = pred_dir / "test_static_inverse.csv" + gt_npool_path = pred_dir / "pft_1d_ground_truth" / "ground_truth_Y_npool.csv" + gt_ppool_path = pred_dir / "pft_1d_ground_truth" / "ground_truth_Y_ppool.csv" + pred_npool_path = pred_dir / "pft_1d_predictions" / "predictions_Y_npool.csv" + pred_ppool_path = pred_dir / "pft_1d_predictions" / "predictions_Y_ppool.csv" + + for p in [static_path, gt_npool_path, gt_ppool_path, pred_npool_path, pred_ppool_path]: + if not p.exists(): + print(f"Missing: {p}", file=sys.stderr) + return None + + static = pd.read_csv(static_path) + gt_npool = pd.read_csv(gt_npool_path) + gt_ppool = pd.read_csv(gt_ppool_path) + pred_npool = pd.read_csv(pred_npool_path) + pred_ppool = pd.read_csv(pred_ppool_path) + + n = len(static) + if n != len(gt_npool) or n != len(gt_ppool) or n != len(pred_npool) or n != len(pred_ppool): + print("Row count mismatch.", file=sys.stderr) + return None + + # PCT columns: PCT_NAT_PFT_1 .. PCT_NAT_PFT_16 (1-based PFT index) + pct_cols = [f"PCT_NAT_PFT_{k}" for k in range(1, 17)] + if not all(c in static.columns for c in pct_cols): + print("PCT_NAT_PFT_1..16 not found in test_static_inverse.csv", file=sys.stderr) + return None + + gt_npool_cols = [f"Y_npool_pft{k}" for k in range(1, 17)] + gt_ppool_cols = [f"Y_ppool_pft{k}" for k in range(1, 17)] + if not all(c in gt_npool.columns for c in gt_npool_cols) or not all(c in gt_ppool.columns for c in gt_ppool_cols): + print("Expected Y_npool_pft1..16 and Y_ppool_pft1..16 in GT/pred CSVs.", file=sys.stderr) + return None + + results = {"run_dir": str(run_dir), "pct_min": pct_min, "n_test_cells": n} + npool_per_pft = [] + ppool_per_pft = [] + + for pft_idx in range(1, 17): + # Mask: only grid cells where this PFT is present + pct_col = f"PCT_NAT_PFT_{pft_idx}" + pct_vals = pd.to_numeric(static[pct_col], errors="coerce").fillna(0).values + mask = pct_vals > pct_min + n_valid = int(mask.sum()) + + # NPOOL for this PFT + gt_col = f"Y_npool_pft{pft_idx}" + r2_n, rmse_n, n_used = _r2_rmse( + gt_npool.loc[mask, gt_col].values, + pred_npool.loc[mask, gt_col].values, + ) + npool_per_pft.append({ + "pft": pft_idx, + "n_cells_with_pft": n_valid, + "r2": float(r2_n), + "rmse": float(rmse_n), + }) + + # PPOOL for this PFT + gt_col_p = f"Y_ppool_pft{pft_idx}" + r2_p, rmse_p, _ = _r2_rmse( + gt_ppool.loc[mask, gt_col_p].values, + pred_ppool.loc[mask, gt_col_p].values, + ) + ppool_per_pft.append({ + "pft": pft_idx, + "n_cells_with_pft": n_valid, + "r2": float(r2_p), + "rmse": float(rmse_p), + }) + + results["npool_per_pft"] = npool_per_pft + results["ppool_per_pft"] = ppool_per_pft + + # Aggregate: mean R² over PFTs (only PFTs with enough valid cells, e.g. n >= 10) + min_cells = 10 + r2_npool_list = [x["r2"] for x in npool_per_pft if x["n_cells_with_pft"] >= min_cells and np.isfinite(x["r2"])] + r2_ppool_list = [x["r2"] for x in ppool_per_pft if x["n_cells_with_pft"] >= min_cells and np.isfinite(x["r2"])] + results["npool_mean_r2_over_pfts"] = float(np.mean(r2_npool_list)) if r2_npool_list else None + results["ppool_mean_r2_over_pfts"] = float(np.mean(r2_ppool_list)) if r2_ppool_list else None + + # Print report + print() + print("=== NPOOL / PPOOL validation per PFT (only cells with PCT_NAT_PFT_k > pct_min) ===") + print(f"Run: {run_dir}") + print(f"pct_min: {pct_min} (include grid cells where PCT_NAT_PFT_k > {pct_min})") + print(f"Test grid cells: {n}") + print() + print("NPOOL per PFT (only cells where this PFT is present):") + print(" PFT n_cells R² RMSE") + for x in npool_per_pft: + print(f" {x['pft']:2d} {x['n_cells_with_pft']:5d} {x['r2']:7.4f} {x['rmse']:.4f}") + print(f" Mean R² (PFTs with ≥{min_cells} cells): {results['npool_mean_r2_over_pfts']}") + print() + print("PPOOL per PFT (only cells where this PFT is present):") + print(" PFT n_cells R² RMSE") + for x in ppool_per_pft: + print(f" {x['pft']:2d} {x['n_cells_with_pft']:5d} {x['r2']:7.4f} {x['rmse']:.4f}") + print(f" Mean R² (PFTs with ≥{min_cells} cells): {results['ppool_mean_r2_over_pfts']}") + + if output_path: + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(results, f, indent=2) + print(f"\nReport written to {output_path}") + + return results + + +def main(): + ap = argparse.ArgumentParser( + description="NPOOL/PPOOL validation per PFT: only grid cells with PCT_NAT_PFT_k > pct_min." + ) + ap.add_argument("run_dir", type=str, help="Path to run directory") + ap.add_argument("--output", "-o", type=str, help="Write JSON report to this file") + ap.add_argument("--pct-min", type=float, default=0.0, help="Minimum PCT_NAT_PFT_k to include (default 0, i.e. > 0)") + args = ap.parse_args() + run_path = Path(args.run_dir) + if not run_path.is_absolute(): + # Resolve relative to cwd so "." means the current (run) directory + run_path = (Path.cwd() / run_path).resolve() + out_path = Path(args.output).resolve() if args.output else None + run(run_path, pct_min=args.pct_min, output_path=out_path) + + +if __name__ == "__main__": + main() diff --git a/scripts/verify_excluded_locations_vs_nofilter.py b/scripts/verify_excluded_locations_vs_nofilter.py new file mode 100644 index 0000000..62639ec --- /dev/null +++ b/scripts/verify_excluded_locations_vs_nofilter.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +""" +Verify that analysis/excluded_locations plots show the same gridcells that are +excluded by the natveg filter, and that analysis_nofilter contains all gridcells +(included + excluded). Compares: +- Excluded mask from test_static_inverse.csv (same logic as plot_excluded_locations.py) +- excluded_locations_lat_lon.csv (saved by plot_excluded_locations) +- Row counts and optional value overlap for one 1D and one 2D variable. +""" +import argparse +import numpy as np +import pandas as pd +from pathlib import Path + + +def get_excluded_mask(static_path): + """Same logic as plot_excluded_locations.get_excluded_indices.""" + df = pd.read_csv(static_path) + pct_natveg = df["PCT_NATVEG"].values.astype(float) + pct_pft0 = df["PCT_NAT_PFT_0"].values.astype(float) + include = (pct_natveg > 0) & (pct_pft0 < 100) + excluded = ~include + return excluded, df + + +def main(): + parser = argparse.ArgumentParser(description="Verify excluded_locations vs natveg filter and nofilter") + parser.add_argument("results_dir", nargs="?", default=".", help="Results directory (e.g. cnp_results/run_xxx)") + args = parser.parse_args() + results_dir = Path(args.results_dir).resolve() + + static_path = results_dir / "cnp_predictions" / "test_static_inverse.csv" + excl_csv = results_dir / "analysis" / "excluded_locations" / "excluded_locations_lat_lon.csv" + + if not static_path.exists(): + print(f"ERROR: {static_path} not found") + return 1 + + excluded_mask, static_df = get_excluded_mask(static_path) + n_total = len(static_df) + n_excluded = int(np.sum(excluded_mask)) + n_included = n_total - n_excluded + excluded_indices = np.where(excluded_mask)[0] + + print("=" * 60) + print("Natveg filter (from test_static_inverse.csv)") + print(" Include: PCT_NATVEG > 0 AND PCT_NAT_PFT_0 < 100") + print(" Exclude: PCT_NATVEG <= 0 OR PCT_NAT_PFT_0 >= 100") + print("=" * 60) + print(f"Total gridcells: {n_total}") + print(f"Included (in analysis): {n_included}") + print(f"Excluded: {n_excluded}") + + # Compare with excluded_locations_lat_lon.csv + if excl_csv.exists(): + excl_df = pd.read_csv(excl_csv) + if "Latitude" in excl_df.columns and "Longitude" in excl_df.columns: + n_csv = len(excl_df) + # Static has same row order; excluded rows in static should match CSV + static_excl = static_df.loc[excluded_mask, ["Latitude", "Longitude"]].reset_index(drop=True) + # Allow small float tolerance + lat_ok = np.allclose(excl_df["Latitude"].values, static_excl["Latitude"].values, rtol=0, atol=1e-5) + lon_ok = np.allclose(excl_df["Longitude"].values, static_excl["Longitude"].values, rtol=0, atol=1e-5) + print(f"\nExcluded locations CSV: {excl_csv.name}") + print(f" Rows in CSV: {n_csv}") + print(f" Match count: {'YES' if n_csv == n_excluded else 'NO (expected ' + str(n_excluded) + ')'}") + print(f" Lat/Lon match static: {'YES' if (lat_ok and lon_ok) else 'NO'}") + else: + print(f"\nCSV found but missing Latitude/Longitude columns.") + else: + print(f"\nExcluded locations CSV not found: {excl_csv}") + + # 1D variable: cpool + gt_1d_path = results_dir / "cnp_predictions" / "pft_1d_ground_truth" / "ground_truth_Y_cpool.csv" + if gt_1d_path.exists(): + gt_1d = pd.read_csv(gt_1d_path) + coord_cols = [c for c in ["Longitude", "Latitude", "long", "lat"] if c in gt_1d.columns] + gt_1d_vals = gt_1d.drop(columns=coord_cols, errors="ignore") + n_rows_1d = len(gt_1d) + n_pfts = gt_1d_vals.shape[1] + assert n_rows_1d == n_total, f"1D rows {n_rows_1d} vs static {n_total}" + # Points in excluded_locations plot = excluded rows × PFTs (flattened) + pts_excluded_1d = n_excluded * n_pfts + pts_included_1d = n_included * n_pfts + pts_all_1d = n_total * n_pfts + print(f"\n1D variable (cpool): rows={n_rows_1d}, PFTs={n_pfts}") + print(f" Points in excluded_locations plot: {pts_excluded_1d} (gridcells × PFTs)") + print(f" Points in filtered analysis: {pts_included_1d}") + print(f" Points in nofilter (all): {pts_all_1d}") + print(f" Check included+excluded=all: {pts_included_1d + pts_excluded_1d == pts_all_1d}") + else: + print(f"\n1D GT not found: {gt_1d_path}") + + # 2D variable: soil1c_vr or first available + gt_2d_dir = results_dir / "cnp_predictions" / "soil_2d_ground_truth" + if gt_2d_dir.exists(): + for var in ["soil1c_vr", "litr1c_vr"]: + gt_2d_path = gt_2d_dir / f"ground_truth_Y_{var}.csv" + if not gt_2d_path.exists(): + continue + gt_2d = pd.read_csv(gt_2d_path) + coord_cols = [c for c in ["Longitude", "Latitude", "long", "lat"] if c in gt_2d.columns] + gt_2d_vals = gt_2d.drop(columns=coord_cols, errors="ignore") + n_rows_2d = len(gt_2d) + n_layers = gt_2d_vals.shape[1] + assert n_rows_2d == n_total, f"2D rows {n_rows_2d} vs static {n_total}" + pts_excluded_2d = n_excluded * n_layers + pts_included_2d = n_included * n_layers + pts_all_2d = n_total * n_layers + print(f"\n2D variable ({var}): rows={n_rows_2d}, layers={n_layers}") + print(f" Points in excluded_locations plot: {pts_excluded_2d} (gridcells × layers)") + print(f" Points in filtered analysis: {pts_included_2d}") + print(f" Points in nofilter (all): {pts_all_2d}") + print(f" Check included+excluded=all: {pts_included_2d + pts_excluded_2d == pts_all_2d}") + break + else: + print(f"\n2D GT dir not found: {gt_2d_dir}") + + print("\n" + "=" * 60) + print("Conclusion:") + print(" - analysis/ (with natveg filter) uses only INCLUDED gridcells.") + print(" - analysis/excluded_locations/ shows only EXCLUDED gridcells (same as above mask).") + print(" - analysis_nofilter/ uses ALL gridcells (included + excluded).") + print(" So excluded_locations plots are exactly the gridcells removed from the filtered analysis") + print(" and present in the nofilter plots.") + print("=" * 60) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/train_cnp_model.py b/train_cnp_model.py index f1a4e80..0c398b6 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -133,6 +133,13 @@ def main(): default='cnp_results', help='Output directory for results' ) + parser.add_argument( + '--output-dir-suffix', + default=None, + type=str, + metavar='SUFFIX', + help='Optional suffix for the run folder name (e.g. natveg_improved -> run_YYYYMMDD_HHMMSS_natveg_improved)' + ) parser.add_argument( '--epochs', '--epoch', dest='epochs', @@ -247,6 +254,22 @@ def main(): default=None, help='Latitude column name override (default: auto-detect from static columns)' ) + parser.add_argument( + '--longitudes-to-drop', + type=str, + default=None, + help='Comma-separated longitudes to drop from training (e.g. "0,358.75"). Overrides config/CNP_IO.' + ) + parser.add_argument( + '--natveg-only', + action='store_true', + help='Keep only gridcells with natural vegetation (PCT_NATVEG>0 and PCT_NAT_PFT_0<100). Overrides config.' + ) + parser.add_argument( + '--no-natveg-filter-before-split', + action='store_true', + help='With --natveg-only: split on full data then filter only training set to natveg, so test set matches no-filter run. Default: filter before split (legacy).' + ) parser.add_argument( '--max-files', type=int, @@ -403,9 +426,11 @@ def main(): ) args = parser.parse_args() - # Create output directory with timestamp + # Create output directory with timestamp (optional suffix) timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') - output_dir = Path(args.output_dir) / f"run_{timestamp}" + suffix = (args.output_dir_suffix or '').strip().replace(' ', '_').replace('/', '_').strip('_') + run_name = f"run_{timestamp}" + (f"_{suffix}" if suffix else "") + output_dir = Path(args.output_dir) / run_name output_dir.mkdir(parents=True, exist_ok=True) # Setup logging with timestamped log file in output directory (default INFO so log file is populated) @@ -478,6 +503,22 @@ def main(): logger.info(f"Enabled tropical filtering: {tropical_kwargs}") except Exception as e: logger.warning(f"Failed to apply tropical filtering config: {e}") + # Optional longitude filtering (CLI overrides config/CNP_IO) + if args.longitudes_to_drop is not None: + try: + parts = [p.strip() for p in str(args.longitudes_to_drop).split(',') if p.strip()] + longitudes = [float(x) for x in parts] + config.update_data_config(longitudes_to_drop=longitudes) + logger.info(f"Longitudes to drop (CLI): {longitudes}") + except Exception as e: + logger.warning(f"Failed to parse --longitudes-to-drop: {e}") + # Optional natveg-only filtering (CLI overrides config) + if args.natveg_only: + config.update_data_config(natveg_only=True) + logger.info("Enabled natveg-only filtering (PCT_NATVEG>0 and PCT_NAT_PFT_0<100).") + if getattr(args, 'no_natveg_filter_before_split', False): + config.update_data_config(natveg_filter_before_split=False) + logger.info("Natveg filter applied after split (test set will match no-filter run).") if args.variable_list is not None: logger.info(f"Using CNP configuration from variable list file: {args.variable_list}") else: @@ -749,6 +790,28 @@ def main(): if len(parts) == 2: update_kwargs['tropical_lat_range'] = (float(parts[0]), float(parts[1])) + # Longitude filtering: drop samples at these longitudes (config overrides CNP_IO) + if 'longitudes_to_drop' in filter_config and args.longitudes_to_drop is None: + lon_drop = filter_config['longitudes_to_drop'] + if isinstance(lon_drop, list): + update_kwargs['longitudes_to_drop'] = [float(x) for x in lon_drop] + elif isinstance(lon_drop, str): + update_kwargs['longitudes_to_drop'] = [float(x.strip()) for x in lon_drop.split(',') if x.strip()] + + # Natveg-only: keep only PCT_NATVEG>0 and PCT_NAT_PFT_0<100 (CLI takes precedence) + if 'natveg_only' in filter_config and not args.natveg_only: + update_kwargs['natveg_only'] = bool(filter_config['natveg_only']) + if 'natveg_filter_before_split' in filter_config and not getattr(args, 'no_natveg_filter_before_split', False): + update_kwargs['natveg_filter_before_split'] = bool(filter_config['natveg_filter_before_split']) + if 'region_boxes' in filter_config: + boxes = filter_config['region_boxes'] + if isinstance(boxes, (list, tuple)) and len(boxes) > 0: + parsed = [] + for b in boxes: + if isinstance(b, (list, tuple)) and len(b) >= 4: + parsed.append((float(b[0]), float(b[1]), float(b[2]), float(b[3]))) + if parsed: + update_kwargs['region_boxes'] = parsed if update_kwargs: config.update_data_config(**update_kwargs) logger.info(f"Applied data_filtering_config from unified config: {update_kwargs}") @@ -1237,10 +1300,18 @@ def main(): data_cfg = getattr(config, 'data_config', None) data_config_snapshot = None if data_cfg is not None: + tr = getattr(data_cfg, 'tropical_lat_range', (-23.5, 23.5)) + if tr is None: + tr = (-23.5, 23.5) data_config_snapshot = { 'data_paths': list(getattr(data_cfg, 'data_paths', []) or []), 'file_pattern': getattr(data_cfg, 'file_pattern', None) or 'enhanced_1_training_data_batch_*.pkl', 'dataset_file_patterns': dict(getattr(data_cfg, 'dataset_file_patterns', None) or {}), + 'longitudes_to_drop': list(getattr(data_cfg, 'longitudes_to_drop', None) or []), + 'natveg_only': bool(getattr(data_cfg, 'natveg_only', False)), + 'natveg_filter_before_split': bool(getattr(data_cfg, 'natveg_filter_before_split', True)), + 'tropical_only': bool(getattr(data_cfg, 'tropical_only', False)), + 'tropical_lat_range': [float(tr[0]), float(tr[1])], } config_dict = { 'include_water': include_water,