From 58f23c66471a7e386e26ae4ab4a833edd138d356 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Fri, 5 Sep 2025 14:26:05 -0400 Subject: [PATCH 01/51] add variable checking utilities --- scripts/Prediction_Checking_Utility_README.md | 102 +++++++++++ scripts/check_pft1d_predictions.py | 173 ++++++++++++++++++ scripts/check_soil2d_predictions.py | 133 ++++++++++++++ 3 files changed, 408 insertions(+) create mode 100644 scripts/Prediction_Checking_Utility_README.md create mode 100644 scripts/check_pft1d_predictions.py create mode 100644 scripts/check_soil2d_predictions.py diff --git a/scripts/Prediction_Checking_Utility_README.md b/scripts/Prediction_Checking_Utility_README.md new file mode 100644 index 0000000..b18724b --- /dev/null +++ b/scripts/Prediction_Checking_Utility_README.md @@ -0,0 +1,102 @@ +## Scripts Overview + +This README documents three utility scripts used to validate and clean model outputs: + +- `scripts/check_pft1d_predictions.py` +- `scripts/check_soil2d_predictions.py` +- `scripts/fix_pft1d_nans.py` + +All scripts support CLI flags so you can point them at different result folders without editing code. + +--- + +### check_pft1d_predictions.py + +**Purpose**: Quality checks for PFT1D prediction CSVs (`predictions_Y_*.csv`). Performs: + +- **Negative/invalid values check** +- **NaN detection** +- **Constant prediction detection** with tolerance: + - Variables where **all PFT columns** are constant + - Variables where **at least one PFT column** is constant (reports `pftN=value`) +- Ignores rows where **all PFT columns are NaN** when checking constancy +- Assumes the first two columns are `Longitude` and `Latitude` and excludes them from PFT analysis + +**Key assumptions**: + +- Files are shaped like: `Longitude,Latitude,Y__pft1,...,Y__pft14` +- PFT columns begin at column index 2 + +**CLI**: + +- `--pred-dir`: Directory containing `predictions_Y_*.csv` +- `--eps`: Numerical tolerance for “effectively constant” checks (default suggested: `1e-9`) + +**Example**: + +```bash +python /mnt/proj-shared/AI4BGC_7xw/AI4BGC/scripts/check_pft1d_predictions.py \ + --pred-dir /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20250904_180409/cnp_predictions/pft_1d_predictions \ + --eps 1e-9 +``` + +You can filter for the “all PFTs constant” section: + +```bash +python /mnt/proj-shared/AI4BGC_7xw/AI4BGC/scripts/check_pft1d_predictions.py \ + --pred-dir /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20250904_180409/cnp_predictions/pft_1d_predictions \ + --eps 1e-9 | sed -n '/Variables where all PFT columns are constant:/,$p' | cat +``` + +--- + +### check_soil2d_predictions.py + +**Purpose**: Quality checks for Soil-2D prediction CSVs (e.g., `predictions_Y_soil*_*.csv`). Similar checks as the PFT1D script: + +- Negative/invalid values +- NaNs +- Constant per-column detection with tolerance +- Optional summary of variables where all columns are constant + +**CLI**: + +- `--pred-dir`: Directory containing Soil-2D prediction CSVs +- `--eps`: Numerical tolerance for constant checks + +**Example**: + +```bash +python /mnt/proj-shared/AI4BGC_7xw/AI4BGC/scripts/check_soil2d_predictions.py \ + --pred-dir /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20250904_180409/cnp_predictions/soil_2d_predictions \ + --eps 1e-9 +``` + +--- + +### fix_pft1d_nans.py + +**Purpose**: Clean PFT1D prediction CSVs by replacing `NaN` values with zeros. + +Note: This currently replaces NaNs with zeros across all columns. Use before running the checkers to prevent NaNs from interfering with constant detection. + +**CLI**: + +- `--pred-dir`: Directory containing `predictions_Y_*.csv` + +**Example**: + +```bash +python /mnt/proj-shared/AI4BGC_7xw/AI4BGC/scripts/fix_pft1d_nans.py \ + --pred-dir /mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20250904_180409/cnp_inference_entire_dataset/cnp_predictions/pft_1d_predictions +``` + +--- + +### Tips + +- Prefer using CLI flags to control behavior and paths; avoid editing script internals. +- Use `--eps` when you want to treat nearly constant floating values as constant. +- Large CSVs: piping through tools like `sed`, `grep`, or `head` can help focus on the sections you need. + + diff --git a/scripts/check_pft1d_predictions.py b/scripts/check_pft1d_predictions.py new file mode 100644 index 0000000..19196f4 --- /dev/null +++ b/scripts/check_pft1d_predictions.py @@ -0,0 +1,173 @@ +import os +import pandas as pd +import numpy as np +import argparse + +# Directory containing all predictions_Y_*.csv files +PREDICTIONS_DIR = './cnp_inference_entire_dataset/cnp_predictions/pft_1d_predictions' + +neg_violations = [] +nan_violations = [] +nan_locations = {} +const_violations = [] # list of tuples (var, const_value) +pft_const_violations = {} # var -> list of tuples (pft_index, const_value) +all_pft_const_violations = {} # var -> list of tuples (pft_index, const_value) + +# Parse CLI args to allow overriding predictions directory +parser = argparse.ArgumentParser(description='Check PFT1D prediction CSVs for sign and NaN issues') +parser.add_argument('--pred-dir', type=str, default=None, help='Directory containing predictions_Y_*.csv files') +parser.add_argument('--eps', type=float, default=1e-9, help='Tolerance for constant detection (max-min <= eps)') +args, unknown = parser.parse_known_args() +if args.pred_dir: + PREDICTIONS_DIR = args.pred_dir + +EPS = float(args.eps) + +def fmt(v: float) -> str: + try: + return f"{float(v):.8g}" + except Exception: + return str(v) + +# Helper to check sign constraints +def check_sign(var, arr): + if var == 'xsmrpool': + pos_count = (arr > 0).sum() + if pos_count > 0: + print(f' [WARNING] {pos_count} positive values found in xsmrpool (should be non-positive)') + else: + neg_count = (arr < 0).sum() + if neg_count > 0: + print(f' [WARNING] {neg_count} negative values found in {var} (should be non-negative)') + neg_violations.append(var) + +def parse_pft_index(col_name: str) -> int: + # Expecting columns like 'Y_var_pft7' + try: + return int(col_name.rsplit('pft', 1)[-1]) + except Exception: + return -1 + + +def main(): + files = [f for f in os.listdir(PREDICTIONS_DIR) if f.startswith('predictions_Y_') and f.endswith('.csv')] + files.sort() + print(f'Found {len(files)} prediction files.') + for fname in files: + var = fname.replace('predictions_Y_', '').replace('.csv', '') + fpath = os.path.join(PREDICTIONS_DIR, fname) + try: + df = pd.read_csv(fpath) + # Always ignore the first two columns (Longitude, Latitude) + df_numeric = df.iloc[:, 2:] + nan_count = df_numeric.isna().sum().sum() + print(f'Variable: {var}') + print(f' Shape: {df_numeric.shape}, Total values: {df_numeric.size}, NaNs: {nan_count}') + print(f' Min: {df_numeric.min().min() if df_numeric.size else "N/A"}, Max: {df_numeric.max().max() if df_numeric.size else "N/A"}, Mean: {df_numeric.mean().mean() if df_numeric.size else "N/A"}') + check_sign(var, df_numeric.values.flatten()) + # Print a few sample values + print(f' Sample values: {df_numeric.values.flatten()[:8] if df_numeric.size else "N/A"}') + if nan_count > 0: + nan_violations.append(var) + # Find first few NaN locations (row, col) + nan_pos = np.argwhere(df_numeric.isna().values) + nan_locations[var] = nan_pos[:5] # Show up to 5 locations + print(f' [NaN] First NaN locations (row, col): {nan_pos[:5]}') + + # Create a NaN-robust view by dropping rows where all PFT columns are NaN + df_valid = df_numeric[~df_numeric.isna().all(axis=1)] + + # Constant prediction detection across all PFTs (tolerance) using NaN-robust view + if df_valid.size > 0: + arr_valid = df_valid.values.flatten() + # If all values are NaN after drop (unlikely), skip + if not np.all(np.isnan(arr_valid)): + if (np.nanmax(arr_valid) - np.nanmin(arr_valid)) <= EPS: + const_val = float(np.nanmean(arr_valid)) + const_violations.append((var, const_val)) + print(f' [CONST] All predictions are constant (±{EPS}): {fmt(const_val)}') + + # Per-PFT constant detection (by column) using NaN-robust view + const_pfts = [] + for col in df_valid.columns: + col_vals = df_valid[col].values + if col_vals.size == 0 or np.all(np.isnan(col_vals)): + continue + if (np.nanmax(col_vals) - np.nanmin(col_vals)) <= EPS: + pft_idx = parse_pft_index(col) + const_pfts.append((pft_idx, float(np.nanmean(col_vals)))) + if const_pfts: + # sort by pft index where possible + const_pfts_sorted = sorted(const_pfts, key=lambda t: (t[0] if t[0] != -1 else 10**9)) + pft_const_violations[var] = const_pfts_sorted + pretty = ', '.join([f'pft{p}:{fmt(v)}' for p, v in const_pfts_sorted if p != -1]) + fallback = ', '.join([f'{c}:{fmt(v)}' for (p, v), c in zip(const_pfts_sorted, df_valid.columns)]) + print(f' [CONST_PFT] Constant columns: {pretty if pretty else fallback}') + + # All-PFT constant (every PFT column constant) detection using NaN-robust view + if df_valid.shape[1] > 0 and df_valid.shape[0] > 0: + is_const_mask = [] + all_const = [] + for c in df_valid.columns: + vals = df_valid[c].values + if vals.size == 0 or np.all(np.isnan(vals)): + is_const_mask.append(False) + continue + is_const = (np.nanmax(vals) - np.nanmin(vals)) <= EPS + is_const_mask.append(is_const) + if is_const: + all_const.append((parse_pft_index(c), float(np.nanmean(vals)))) + if all(is_const_mask) and len(is_const_mask) == df_valid.shape[1]: + all_const_sorted = sorted(all_const, key=lambda t: (t[0] if t[0] != -1 else 10**9)) + all_pft_const_violations[var] = all_const_sorted + pretty_all = ', '.join([f'pft{p}:{fmt(v)}' for p, v in all_const_sorted if p != -1]) + print(f' [CONST_ALL_PFTS] All PFT columns constant: {pretty_all}') + except Exception as e: + print(f'[ERROR] Could not process {fname}: {e}') + print('-' * 60) + + # Summary Table + print('\nSUMMARY TABLE') + print('Variables with negative values (should be non-negative):') + if neg_violations: + for v in neg_violations: + print(f' - {v}') + else: + print(' None') + print('\nVariables with NaNs:') + if nan_violations: + for v in nan_violations: + print(f' - {v} (first NaN locations: {nan_locations[v]})') + else: + print(' None') + print('\nVariables with constant predictions (all PFTs):') + if const_violations: + for v, val in const_violations: + print(f' - {v}: constant value {fmt(val)}') + else: + print(' None') + print('\nVariables with constant predictions in specific PFTs:') + if pft_const_violations: + for v, items in pft_const_violations.items(): + items_sorted = sorted(items, key=lambda t: (t[0] if t[0] != -1 else 10**9)) + items_str = ', '.join([f'pft{p}={fmt(val)}' if p != -1 else f'col={fmt(val)}' for p, val in items_sorted]) + print(f' - {v}: {items_str}') + else: + print(' None') + print('\nVariables where all PFT columns are constant:') + if all_pft_const_violations: + for v, items in all_pft_const_violations.items(): + items_sorted = sorted(items, key=lambda t: (t[0] if t[0] != -1 else 10**9)) + items_str = ', '.join([f'pft{p}={fmt(val)}' if p != -1 else f'col={fmt(val)}' for p, val in items_sorted]) + print(f' - {v}: {items_str}') + else: + print(' None') + print('\nVariables (names only) with at least one constant PFT column:') + if pft_const_violations: + for v in sorted(pft_const_violations.keys()): + print(f' - {v}') + else: + print(' None') + +if __name__ == '__main__': + main() diff --git a/scripts/check_soil2d_predictions.py b/scripts/check_soil2d_predictions.py new file mode 100644 index 0000000..1979e7d --- /dev/null +++ b/scripts/check_soil2d_predictions.py @@ -0,0 +1,133 @@ +import os +import pandas as pd +import numpy as np +import argparse + +# Directory containing all soil2d predictions_Y_*.csv files +PREDICTIONS_DIR = './cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions' + +neg_violations = [] +nan_violations = [] +nan_locations = {} +col_const_violations = {} # var -> list of tuples (col_name, const_value) +all_cols_const_violations = {} # var -> list of tuples (col_name, const_value) + +# CLI args +parser = argparse.ArgumentParser(description='Check Soil2D prediction CSVs for sign, NaNs, and constant columns') +parser.add_argument('--pred-dir', type=str, default=None, help='Directory containing predictions_Y_*.csv files') +parser.add_argument('--eps', type=float, default=1e-9, help='Tolerance for constant detection (max-min <= eps)') +args, unknown = parser.parse_known_args() +if args.pred_dir: + PREDICTIONS_DIR = args.pred_dir +EPS = float(args.eps) + +def fmt(v: float) -> str: + try: + return f"{float(v):.8g}" + except Exception: + return str(v) + +# Helper to check sign constraints +def check_sign(var, arr): + neg_count = (arr < 0).sum() + if neg_count > 0: + print(f' [WARNING] {neg_count} negative values found in {var} (should be non-negative)') + neg_violations.append(var) + +def main(): + files = [f for f in os.listdir(PREDICTIONS_DIR) if f.startswith('predictions_Y_') and f.endswith('.csv')] + files.sort() + print(f'Found {len(files)} soil2d prediction files.') + for fname in files: + var = fname.replace('predictions_Y_', '').replace('.csv', '') + fpath = os.path.join(PREDICTIONS_DIR, fname) + try: + df = pd.read_csv(fpath) + # Always ignore the first two columns (Longitude, Latitude) if present + if df.shape[1] > 2: + df_numeric = df.iloc[:, 2:] + else: + df_numeric = df + # Replace all NaNs with zeros for analysis + nan_count = df_numeric.isna().sum().sum() + if nan_count > 0: + nan_violations.append(var) + nan_pos = np.argwhere(df_numeric.isna().values) + nan_locations[var] = nan_pos[:5] + print(f' [NaN] First NaN locations (row, col): {nan_pos[:5]}') + df_numeric = df_numeric.fillna(0) + arr = df_numeric.values.flatten() + print(f'Variable: {var}') + print(f' Shape: {df_numeric.shape}, Total values: {arr.size}, NaNs (before replace): {nan_count}') + if arr.size: + print(f' Min: {arr.min()}, Max: {arr.max()}, Mean: {arr.mean()}') + else: + print(' Min: N/A, Max: N/A, Mean: N/A') + check_sign(var, arr) + print(f' Sample values: {arr[:8] if arr.size else "N/A"}') + + # Per-column constant detection (tolerance) + const_cols = [] + for col in df_numeric.columns: + col_vals = df_numeric[col].values + if col_vals.size == 0: + continue + if (np.nanmax(col_vals) - np.nanmin(col_vals)) <= EPS: + const_cols.append((col, float(np.nanmean(col_vals)))) + if const_cols: + # sort by column name for readability + const_cols_sorted = sorted(const_cols, key=lambda t: t[0]) + col_const_violations[var] = const_cols_sorted + pretty = ', '.join([f'{c}:{fmt(v)}' for c, v in const_cols_sorted]) + print(f' [CONST_COLS] Constant columns: {pretty}') + + # All columns constant + if df_numeric.shape[1] > 0: + is_const_mask = [(np.nanmax(df_numeric[c].values) - np.nanmin(df_numeric[c].values)) <= EPS for c in df_numeric.columns] + if all(is_const_mask): + all_cols = [(c, float(np.nanmean(df_numeric[c].values))) for c in df_numeric.columns] + all_cols_sorted = sorted(all_cols, key=lambda t: t[0]) + all_cols_const_violations[var] = all_cols_sorted + pretty_all = ', '.join([f'{c}:{fmt(v)}' for c, v in all_cols_sorted]) + print(f' [CONST_ALL_COLS] All columns constant: {pretty_all}') + except Exception as e: + print(f'[ERROR] Could not process {fname}: {e}') + print('-' * 60) + + # Summary Table + print('\nSUMMARY TABLE') + print('Soil2D variables with negative values (should be non-negative):') + if neg_violations: + for v in neg_violations: + print(f' - {v}') + else: + print(' None') + print('\nSoil2D variables with NaNs (before replace):') + if nan_violations: + for v in nan_violations: + print(f' - {v} (first NaN locations: {nan_locations[v]})') + else: + print(' None') + print('\nSoil2D variables with constant predictions in specific columns:') + if col_const_violations: + for v, items in col_const_violations.items(): + items_str = ', '.join([f'{c}={fmt(val)}' for c, val in items]) + print(f' - {v}: {items_str}') + else: + print(' None') + print('\nSoil2D variables where all columns are constant:') + if all_cols_const_violations: + for v, items in all_cols_const_violations.items(): + items_str = ', '.join([f'{c}={fmt(val)}' for c, val in items]) + print(f' - {v}: {items_str}') + else: + print(' None') + print('\nSoil2D variables (names only) with at least one constant column:') + if col_const_violations: + for v in sorted(col_const_violations.keys()): + print(f' - {v}') + else: + print(' None') + +if __name__ == '__main__': + main() From 99df35b85a985d855ab37532e5a2267ccec89b5e Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Mon, 15 Sep 2025 22:31:43 -0400 Subject: [PATCH 02/51] add prediction analysis code --- docs/CNP_pipeline_runbook.md | 22 +- docs/README_CNP_Model_Workflow.md | 37 ++ docs/README_prediction_quality.md | 112 +++++ scripts/generate_prediction_quality_report.py | 444 ++++++++++++++++++ 4 files changed, 614 insertions(+), 1 deletion(-) create mode 100644 docs/README_prediction_quality.md create mode 100644 scripts/generate_prediction_quality_report.py diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 8558b76..8778d1c 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -29,6 +29,26 @@ Generates scatter plots and statistics using individually normalized results. python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & ``` +### 4.5) Generate comprehensive prediction quality report +Creates detailed quality analysis categorizing predictions as "good", "ok", or "bad" based on statistical thresholds. + +```bash +python ../../scripts/generate_prediction_quality_report.py > prediction_quality_report.log 2>&1 & +``` + +**Output**: Creates `analysis/` directory with: +- Quality assessment CSV files +- Interactive HTML report +- Visualization charts (bar charts, pie charts, scatter plots) +- Text summary report + +**Quality Classification**: +- **Good**: R² ≥ 0.9, Relative RMSE ≤ 0.1, Relative MAE ≤ 0.1 +- **OK**: R² ≥ 0.7, Relative RMSE ≤ 0.25, Relative MAE ≤ 0.25 +- **Bad**: Below OK thresholds + +**Detailed Documentation**: See `docs/README_prediction_quality.md` for comprehensive usage instructions and advanced features. + ### 5) Run inference on the entire dataset Creates a folder `cnp_inference_entire_dataset` with AI predictions for the entire dataset. @@ -55,7 +75,7 @@ python ../../scripts/ai_model_comparison_plot.py > ai_model_comparison.log 2>&1 ### 8) Create a new ELM restart file using AI predictions ```bash -python ../../scripts/ai_predictions_to_restart.py > ai_predictions_to_restart.py 2>&1 & +python ../../scripts/ai_predictions_to_restart.py > ai_predictions_to_restart.log 2>&1 & ``` Outputs a new restart file derived from diff --git a/docs/README_CNP_Model_Workflow.md b/docs/README_CNP_Model_Workflow.md index 721795f..e07e32b 100644 --- a/docs/README_CNP_Model_Workflow.md +++ b/docs/README_CNP_Model_Workflow.md @@ -175,6 +175,43 @@ python ../../scripts/cnp_result_validationplot.py . - **PFT Variables**: tlai performance across 16 PFTs - **Soil Variables**: cwdc_vr, cwdn_vr, cwdp_vr performance across soil layers +### Step 4.5: Generate Comprehensive Prediction Quality Report + +```bash +python ../../scripts/generate_prediction_quality_report.py +``` + +**Prediction Quality Analysis**: +- **Quality Categorization**: Classifies each prediction as "good", "ok", or "bad" based on statistical thresholds +- **Comprehensive Metrics**: Uses R², relative RMSE, and relative MAE for evaluation +- **Variable-Level Analysis**: Provides detailed quality breakdown for each variable +- **Visualization**: Generates multiple charts including stacked bar charts, pie charts, and scatter plots +- **HTML Report**: Creates an interactive HTML report for easy viewing + +**Quality Classification Criteria**: +- **Good**: R² ≥ 0.9, Relative RMSE ≤ 0.1, Relative MAE ≤ 0.1 +- **OK**: R² ≥ 0.7, Relative RMSE ≤ 0.25, Relative MAE ≤ 0.25 +- **Bad**: Below OK thresholds + +**Output Files**: +- `analysis/detailed_quality_assessment.csv`: Full dataset with quality categorization +- `analysis/variable_quality_summary.csv`: Summary statistics for each variable +- `analysis/quality_summary_report.txt`: Text report with overall statistics +- `analysis/prediction_quality_report.html`: Interactive HTML report +- `analysis/prediction_quality_by_variable.png`: Bar chart showing quality distribution +- `analysis/overall_prediction_quality.png`: Pie chart of overall quality +- `analysis/r2_vs_rmse.png`: Scatter plot of R² vs Relative RMSE + +**Customizable Parameters**: +```bash +python ../../scripts/generate_prediction_quality_report.py \ + --r2-good 0.85 \ + --rmse-good 0.15 \ + --output-dir custom_analysis +``` + +**Detailed Documentation**: See `docs/README_prediction_quality.md` for comprehensive usage instructions and advanced features. + ### Step 5: Run Inference on Entire Dataset ```bash diff --git a/docs/README_prediction_quality.md b/docs/README_prediction_quality.md new file mode 100644 index 0000000..13b7a51 --- /dev/null +++ b/docs/README_prediction_quality.md @@ -0,0 +1,112 @@ +# Prediction Quality Analysis Tools + +This repository contains tools for analyzing the quality of AI model predictions based on validation statistics. The tools categorize predictions as "good", "ok", or "bad" based on statistical metrics like R², RMSE, and MAE. + +## Available Tool + +**generate_prediction_quality_report.py**: Comprehensive script that analyzes validation statistics, categorizes predictions, and generates detailed reports with visualizations. + +## Quick Start + +To analyze prediction quality using default settings: + +```bash +python scripts/generate_prediction_quality_report.py +``` + +This will: +1. Read the validation statistics from the default location (`cnp_results/run_20250911_080250/validation_stats.csv`) +2. Categorize each prediction as good, ok, or bad +3. Generate summary reports and visualizations in an `analysis` subdirectory + +**Note**: The script can be run from any directory within the project, as it automatically resolves relative paths based on the project root. + +## Command Line Arguments + +The `generate_prediction_quality_report.py` script accepts the following command line arguments: + +``` +--input PATH Path to validation statistics CSV file + Default: cnp_results/run_20250911_080250/validation_stats.csv + +--output-dir PATH Directory to save output files + Default: Same directory as input + /analysis + +--r2-good FLOAT R² threshold for good predictions + Default: 0.9 + +--r2-ok FLOAT R² threshold for ok predictions + Default: 0.7 + +--rmse-good FLOAT Relative RMSE threshold for good predictions + Default: 0.1 + +--rmse-ok FLOAT Relative RMSE threshold for ok predictions + Default: 0.25 + +--mae-good FLOAT Relative MAE threshold for good predictions + Default: 0.1 + +--mae-ok FLOAT Relative MAE threshold for ok predictions + Default: 0.25 +``` + +## Example Usage + +Analyze a specific validation statistics file with custom thresholds: + +```bash +python scripts/generate_prediction_quality_report.py \ + --input path/to/validation_stats.csv \ + --output-dir path/to/output \ + --r2-good 0.85 \ + --rmse-good 0.15 +``` + +Run from any subdirectory (e.g., from within a results directory): + +```bash +python ../../scripts/generate_prediction_quality_report.py +``` + +## Output Files + +The analysis generates the following output files: + +1. **detailed_quality_assessment.csv**: Full dataset with quality categorization for each prediction +2. **variable_quality_summary.csv**: Summary statistics for each variable +3. **quality_summary_report.txt**: Text report with overall statistics and best/worst variables +4. **prediction_quality_report.html**: Interactive HTML report +5. **prediction_quality_by_variable.png**: Bar chart showing quality distribution by variable +6. **overall_prediction_quality.png**: Pie chart showing overall quality distribution +7. **r2_vs_rmse.png**: Scatter plot of R² vs Relative RMSE + +## Classification Criteria + +Predictions are classified based on the following criteria: + +- **Good**: R² ≥ 0.9, Relative RMSE ≤ 0.1, Relative MAE ≤ 0.1 +- **OK**: R² ≥ 0.7, Relative RMSE ≤ 0.25, Relative MAE ≤ 0.25 +- **Bad**: Below OK thresholds + +Relative RMSE and MAE are calculated by dividing the absolute value by the range of the ground truth data: +- Relative RMSE = RMSE / (gt_max - gt_min) +- Relative MAE = MAE / (gt_max - gt_min) + +## Interpreting Results + +The analysis provides several ways to interpret the prediction quality: + +1. **Overall Statistics**: Percentage of predictions in each quality category +2. **Variable-level Analysis**: Quality breakdown for each variable +3. **Visualizations**: Charts showing the distribution of prediction quality + +Focus on variables with high percentages of "bad" predictions for model improvement. + +## Requirements + +- Python 3.6+ +- pandas +- numpy +- matplotlib +- seaborn diff --git a/scripts/generate_prediction_quality_report.py b/scripts/generate_prediction_quality_report.py new file mode 100644 index 0000000..2e0c3dd --- /dev/null +++ b/scripts/generate_prediction_quality_report.py @@ -0,0 +1,444 @@ +#!/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 + +def main(): + parser = argparse.ArgumentParser(description='Generate prediction quality report from validation statistics') + parser.add_argument('--input', default="cnp_results/run_20250911_080250/validation_stats.csv", + help='Path to validation statistics CSV file') + parser.add_argument('--output-dir', default=None, + help='Directory to save output files (default: same directory as input + /analysis)') + 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)') + args = parser.parse_args() + + # Set up input and output paths + input_path = Path(args.input) + if not input_path.is_absolute(): + # If relative path, make it relative to the script's directory + script_dir = Path(__file__).parent.parent + input_path = script_dir / args.input + + 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(): + # If relative path, make it relative to the script's directory + script_dir = Path(__file__).parent.parent + output_dir = script_dir / args.output_dir + + output_dir.mkdir(parents=True, exist_ok=True) + + # 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) + + # 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) + analysis_df = df[~df['pft'].isin(['Longitude', 'Latitude'])] + + # Create summary by variable + variable_summary = analysis_df.groupby(['variable', 'prediction_quality']).size().unstack(fill_value=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 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) + + # 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() + + # Calculate relative RMSE + 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 + ) + + # 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) + + # 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'), + ] + plt.legend(handles=legend_elements) + + plt.grid(True, alpha=0.3) + plt.tight_layout() + plt.savefig(output_dir / "r2_vs_rmse.png", dpi=300) + + # 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) + 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"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") + + f.write("## Variables with Best Predictions\n") + if 'good_pct' in variable_summary.columns: + best_vars = variable_summary.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: + worst_vars = variable_summary.nsmallest(15, 'good_pct') + 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") + + # 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 +
+ +

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 + worst_vars = variable_summary.nsmallest(15, 'good_pct') + for var_name, row in worst_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}%
+
+ + + """ + + # 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:") + best_vars = variable_summary.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:") + worst_vars = variable_summary.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() From eb7ccc74b4c3059a4a53f9b1a4fa1a01336b7c10 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Thu, 2 Oct 2025 10:44:24 -0400 Subject: [PATCH 03/51] add pft-mask option --- CNP_IO_updated9_modified.txt | 56 ++++++++++++++ config/training_config.py | 3 + data/data_loader_individual.py | 34 +++++++++ docs/CNP_pipeline_runbook.md | 3 + train_cnp_model.py | 11 +++ training/trainer.py | 132 +++++++++++++++++++++++++++------ 6 files changed, 218 insertions(+), 21 deletions(-) create mode 100644 CNP_IO_updated9_modified.txt diff --git a/CNP_IO_updated9_modified.txt b/CNP_IO_updated9_modified.txt new file mode 100644 index 0000000..38b8b90 --- /dev/null +++ b/CNP_IO_updated9_modified.txt @@ -0,0 +1,56 @@ +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, NPP, AR, HR + +1D PFT VARIABLES (39 variables): + +• deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage +• deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage + +• leafc, leafc_storage, leafn, leafn_storage, leafp, leafp_storage +• frootc, frootc_storage, frootn, frootn_storage, frootp, frootp_storage + +• livestemc, livestemc_storage, livestemn, livestemn_storage, livestemp, livestemp_storage, +• livecrootc, livecrootc_storage, livecrootn, livecrootn_storage, livecrootp, livecrootp_storage + +• cpool, npool, ppool + +• tlai, totvegc + +2D VARIABLES (layered - 28 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 + +• soil1c_vr, soil1n_vr, soil1p_vr +• soil2c_vr, soil2n_vr, soil2p_vr +• soil3c_vr, soil3n_vr, soil3p_vr +• soil4c_vr, soil4n_vr, soil4p_vr + +• labilep_vr , occlp_vr, primp_vr, secondp_vr diff --git a/config/training_config.py b/config/training_config.py index e02a66d..1f702cf 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -231,6 +231,9 @@ class TrainingConfig: pft_zero_sparsity_weight: float = 0.0 # default disabled; set >0 to enable pft_zero_threshold: float = 1e-8 # threshold in normalized target space for zero mask + # Mask predictions for absent PFTs using PCT_NAT_PFT (PFT0 ignored) + mask_absent_pfts: bool = False + def get_device(self) -> torch.device: """Get the appropriate device for training.""" if self.device == 'auto': diff --git a/data/data_loader_individual.py b/data/data_loader_individual.py index cab77f9..8c92248 100644 --- a/data/data_loader_individual.py +++ b/data/data_loader_individual.py @@ -402,6 +402,17 @@ def normalize_data(self) -> Dict[str, Any]: 'y_water': y_water_tensor, 'scalers': self.scalers } + # Add per-sample PFT mask derived from raw PCT_NAT_PFT_1..16 (1 where >0, else 0) + try: + pct_cols = [f'PCT_NAT_PFT_{i}' for i in range(1, 17)] + if all(c in self.df.columns for c in pct_cols): + pct = self.df[pct_cols].values.astype(np.float32) + mask = (pct > 0.0).astype(np.float32) # shape [N,16] + ret['pft_presence_mask'] = torch.tensor(mask, dtype=self.preprocessing_config.data_type) + else: + logger.warning("Some PCT_NAT_PFT_1..16 columns are missing; pft_presence_mask not created") + except Exception as _e: + logger.warning(f"Failed to create pft_presence_mask: {_e}") # Optional dump after normalization (group) if os.getenv('DUMP_ALL_PFT_SOIL', '0') == '1': try: @@ -502,6 +513,17 @@ def normalize_data_individual(self, transform_only: bool = False) -> Dict[str, A 'y_water': y_water_tensor, 'scalers': self.scalers } + # Add per-sample PFT mask derived from raw PCT_NAT_PFT_1..16 (1 where >0, else 0) + try: + pct_cols = [f'PCT_NAT_PFT_{i}' for i in range(1, 17)] + if all(c in self.df.columns for c in pct_cols): + pct = self.df[pct_cols].values.astype(np.float32) + mask = (pct > 0.0).astype(np.float32) # shape [N,16] + ret['pft_presence_mask'] = torch.tensor(mask, dtype=self.preprocessing_config.data_type) + else: + logger.warning("Some PCT_NAT_PFT_1..16 columns are missing; pft_presence_mask not created") + except Exception as _e: + logger.warning(f"Failed to create pft_presence_mask: {_e}") # Optional dump after normalization (individual) if os.getenv('DUMP_ALL_PFT_SOIL', '0') == '1': try: @@ -1478,6 +1500,15 @@ def split_data(self, normalized_data: Dict[str, Any]) -> Dict[str, Any]: if 'y_water' in normalized_data and normalized_data['y_water'] is not None: train_data['y_water'] = normalized_data['y_water'][:train_size] test_data['y_water'] = normalized_data['y_water'][train_size:] + + # Split PFT presence mask if present + if 'pft_presence_mask' in normalized_data: + ppm = normalized_data['pft_presence_mask'] + try: + train_data['pft_presence_mask'] = 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") logger.info(f"Split completed:") logger.info(f" - Train time_series shape: {train_time_series.shape}") @@ -1497,6 +1528,9 @@ def split_data(self, normalized_data: Dict[str, Any]) -> Dict[str, Any]: final_keys.append('water') final_keys.append('y_water') + # Optionally include presence mask + if 'pft_presence_mask' in train_data: + final_keys.append('pft_presence_mask') train_data = {k: v for k, v in train_data.items() if k in final_keys} test_data = {k: v for k, v in test_data.items() if k in final_keys} diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 8778d1c..9885ba7 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -29,6 +29,9 @@ Generates scatter plots and statistics using individually normalized results. python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & ``` +(optional) +For a quick result, (--stats-only) option can be used to the following quality report + ### 4.5) Generate comprehensive prediction quality report Creates detailed quality analysis categorizing predictions as "good", "ok", or "bad" based on statistical thresholds. diff --git a/train_cnp_model.py b/train_cnp_model.py index 0fb9b73..4619738 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -220,6 +220,11 @@ def main(): default=None, help='Extra loss weight multiplier for litter phosphorus vars (litr1/2/3p_vr)' ) + parser.add_argument( + '--mask-absent-pfts', + action='store_true', + help='Zero predictions where PCT_NAT_PFT_k == 0 and exclude from loss' + ) args = parser.parse_args() @@ -285,6 +290,12 @@ def main(): predictions_dir=str(output_dir / "cnp_predictions"), use_early_stopping=False ) + if args.mask_absent_pfts: + try: + config.update_training_config(mask_absent_pfts=True) + logger.info("Masking absent PFTs enabled (using PCT_NAT_PFT_1..16)") + except Exception as e: + logger.warning(f"Failed to enable mask_absent_pfts: {e}") # apply xsmrpool loss weight from CLI if provided if args.xsmrpool_loss_weight is not None: try: diff --git a/training/trainer.py b/training/trainer.py index 172cbde..845714e 100644 --- a/training/trainer.py +++ b/training/trainer.py @@ -333,7 +333,8 @@ def train_epoch(self) -> float: self.train_data['y_scalar'], self.train_data['y_soil_2d'], self.train_data['water'], - self.train_data['y_water'] + self.train_data['y_water'], + *( (self.train_data['pft_presence_mask'],) if 'pft_presence_mask' in self.train_data else () ) ) else: train_dataset = TensorDataset( @@ -345,7 +346,9 @@ def train_epoch(self) -> float: self.train_data['variables_2d_soil'], self.train_data['y_scalar'], self.train_data['y_pft_1d'], - self.train_data['y_soil_2d'] + self.train_data['y_soil_2d'], + # Optional mask as final feature; if absent, a placeholder will be injected in-loop + *( (self.train_data['pft_presence_mask'],) if 'pft_presence_mask' in self.train_data else () ) ) train_loader = DataLoader( @@ -365,9 +368,15 @@ def get_loss_value(loss): for batch_idx, batch in enumerate(progress_bar): if 'water' in self.train_data and 'y_water' in self.train_data: - (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, water, y_water) = batch + if 'pft_presence_mask' in self.train_data: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, water, y_water, pft_presence_mask) = batch + else: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, water, y_water) = batch else: - (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d) = batch + if 'pft_presence_mask' in self.train_data: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, pft_presence_mask) = batch + else: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d) = batch # --- DEBUG: Print tensor shapes and device before model call --- # print(f"[DEBUG] Batch {batch_idx} tensor shapes and device:") # print(f" time_series: {time_series.shape}, device: {time_series.device}") @@ -395,6 +404,9 @@ def get_loss_value(loss): if 'water' in self.train_data and 'y_water' in self.train_data: water = water.to(self.device, non_blocking=True).contiguous() y_water = y_water.to(self.device, non_blocking=True).contiguous() + # Presence mask to device if provided + if 'pft_presence_mask' in self.train_data: + pft_presence_mask = pft_presence_mask.to(self.device, non_blocking=True).contiguous() # print(f"[DEBUG] variables_1d_pft shape before model: {variables_1d_pft.shape}") # if variables_1d_pft.dim() == 2 and variables_1d_pft.shape[1] == 224: @@ -416,6 +428,22 @@ def get_loss_value(loss): else: outputs = self.model(time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil) + # Optionally apply PFT presence mask to predictions before loss + if getattr(self.config, 'mask_absent_pfts', False) and 'pft_1d' in outputs and 'pft_presence_mask' in self.train_data: + try: + vec = outputs['pft_1d'] + varnames = list(self.model.data_info.get('variables_1d_pft', [])) if hasattr(self.model, 'data_info') else None + n_vars = len(varnames) if varnames is not None and len(varnames) > 0 else self.train_data['y_pft_1d'].size(1) + n_pfts = 16 + if vec.dim() == 2 and vec.size(1) == n_vars * n_pfts: + vec = vec.view(vec.size(0), n_vars, n_pfts) + if pft_presence_mask.dim() == 2 and pft_presence_mask.size(1) == n_pfts: + mask = pft_presence_mask.view(pft_presence_mask.size(0), 1, n_pfts) + vec = vec * mask + outputs['pft_1d'] = vec.view(vec.size(0), -1) + except Exception: + pass + # Compute loss with variable-specific weights for scalar variables if self.use_variable_weights and hasattr(self, 'scalar_var_weights') and self.scalar_var_weights: # Apply variable-specific weights to scalar variables @@ -680,7 +708,8 @@ def validate_epoch(self) -> float: self.test_data['y_pft_1d'], self.test_data['y_soil_2d'], self.test_data['water'], - self.test_data['y_water'] + self.test_data['y_water'], + *( (self.test_data['pft_presence_mask'],) if 'pft_presence_mask' in self.test_data else () ) ) else: val_dataset = TensorDataset( @@ -692,7 +721,8 @@ def validate_epoch(self) -> float: self.test_data['variables_2d_soil'], self.test_data['y_scalar'], self.test_data['y_pft_1d'], - self.test_data['y_soil_2d'] + self.test_data['y_soil_2d'], + *( (self.test_data['pft_presence_mask'],) if 'pft_presence_mask' in self.test_data else () ) ) val_loader = DataLoader( @@ -712,9 +742,15 @@ def get_loss_value(loss): return loss.item() if hasattr(loss, 'item') else loss for batch_idx, batch in enumerate(progress_bar): if 'water' in self.test_data and 'y_water' in self.test_data: - (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, water, y_water) = batch + if 'pft_presence_mask' in self.test_data: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, water, y_water, pft_presence_mask) = batch + else: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, water, y_water) = batch else: - (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d) = batch + if 'pft_presence_mask' in self.test_data: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, pft_presence_mask) = batch + else: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d) = batch # Move data to device and ensure contiguous time_series = time_series.to(self.device, non_blocking=True).contiguous() static = static.to(self.device, non_blocking=True).contiguous() @@ -728,6 +764,8 @@ def get_loss_value(loss): if 'water' in self.test_data and 'y_water' in self.test_data: water = water.to(self.device, non_blocking=True).contiguous() y_water = y_water.to(self.device, non_blocking=True).contiguous() + if 'pft_presence_mask' in self.test_data: + pft_presence_mask = pft_presence_mask.to(self.device, non_blocking=True).contiguous() # print(f"[DEBUG] variables_1d_pft shape before model (val): {variables_1d_pft.shape}") # if variables_1d_pft.dim() == 2 and variables_1d_pft.shape[1] == 224: @@ -747,6 +785,22 @@ def get_loss_value(loss): else: outputs = self.model(time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil) + # Apply mask in validation as well for consistency + if getattr(self.config, 'mask_absent_pfts', False) and 'pft_1d' in outputs and 'pft_presence_mask' in self.test_data: + try: + vec = outputs['pft_1d'] + varnames = list(self.model.data_info.get('variables_1d_pft', [])) if hasattr(self.model, 'data_info') else None + n_vars = len(varnames) if varnames is not None and len(varnames) > 0 else y_pft_1d.size(1) + n_pfts = 16 + if vec.dim() == 2 and vec.size(1) == n_vars * n_pfts: + vec = vec.view(vec.size(0), n_vars, n_pfts) + if pft_presence_mask.dim() == 2 and pft_presence_mask.size(1) == n_pfts: + mask = pft_presence_mask.view(pft_presence_mask.size(0), 1, n_pfts) + vec = vec * mask + outputs['pft_1d'] = vec.view(vec.size(0), -1) + except Exception: + pass + # Compute loss loss = self._compute_loss(outputs['scalar'], y_scalar) # Vector (PFT1D): base MSE @@ -1048,18 +1102,32 @@ def evaluate(self) -> Tuple[Dict[str, Any], Dict[str, float]]: } return empty_predictions, default_metrics - # Create evaluation data loader - eval_dataset = TensorDataset( - self.test_data['time_series'], - self.test_data['static'], - self.test_data['pft_param'], - self.test_data['scalar'], - self.test_data['variables_1d_pft'], - self.test_data['variables_2d_soil'], - self.test_data['y_scalar'], - self.test_data['y_pft_1d'], - self.test_data['y_soil_2d'] - ) + # Create evaluation data loader (optionally include presence mask) + if 'pft_presence_mask' in self.test_data: + eval_dataset = TensorDataset( + self.test_data['time_series'], + self.test_data['static'], + self.test_data['pft_param'], + self.test_data['scalar'], + self.test_data['variables_1d_pft'], + self.test_data['variables_2d_soil'], + self.test_data['y_scalar'], + self.test_data['y_pft_1d'], + self.test_data['y_soil_2d'], + self.test_data['pft_presence_mask'] + ) + else: + eval_dataset = TensorDataset( + self.test_data['time_series'], + self.test_data['static'], + self.test_data['pft_param'], + self.test_data['scalar'], + self.test_data['variables_1d_pft'], + self.test_data['variables_2d_soil'], + self.test_data['y_scalar'], + self.test_data['y_pft_1d'], + self.test_data['y_soil_2d'] + ) eval_loader = DataLoader( eval_dataset, batch_size=self.config.batch_size, @@ -1079,7 +1147,11 @@ def evaluate(self) -> Tuple[Dict[str, Any], Dict[str, float]]: 'y_soil_2d': [] } with torch.no_grad(): - for time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d in eval_loader: + for batch in eval_loader: + if 'pft_presence_mask' in self.test_data: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d, pft_presence_mask) = batch + else: + (time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil, y_scalar, y_pft_1d, y_soil_2d) = batch # Move to device time_series = time_series.to(self.device, non_blocking=True) static = static.to(self.device, non_blocking=True) @@ -1090,12 +1162,30 @@ def evaluate(self) -> Tuple[Dict[str, Any], Dict[str, float]]: y_scalar = y_scalar.to(self.device, non_blocking=True) y_pft_1d = y_pft_1d.to(self.device, non_blocking=True) y_soil_2d = y_soil_2d.to(self.device, non_blocking=True) + if 'pft_presence_mask' in self.test_data: + pft_presence_mask = pft_presence_mask.to(self.device, non_blocking=True) # Forward pass if self.use_amp and self.scaler is not None: with torch.amp.autocast('cuda'): outputs = self.model(time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil) else: outputs = self.model(time_series, static, pft_param, scalar, variables_1d_pft, variables_2d_soil) + # Apply presence mask to predictions if enabled + if getattr(self.config, 'mask_absent_pfts', False) and 'pft_1d' in outputs and 'pft_presence_mask' in self.test_data: + try: + vec = outputs['pft_1d'] + # Determine n_vars and reshape + varnames = list(self.model.data_info.get('variables_1d_pft', [])) if hasattr(self.model, 'data_info') else None + n_vars = len(varnames) if varnames is not None and len(varnames) > 0 else y_pft_1d.size(1) + n_pfts = 16 + if vec.dim() == 2 and vec.size(1) == n_vars * n_pfts: + vec = vec.view(vec.size(0), n_vars, n_pfts) + if pft_presence_mask.dim() == 2 and pft_presence_mask.size(1) == n_pfts: + mask = pft_presence_mask.view(pft_presence_mask.size(0), 1, n_pfts) + vec = vec * mask + outputs['pft_1d'] = vec.view(vec.size(0), -1) + except Exception: + pass all_predictions['scalar'].append(outputs['scalar'].cpu()) all_predictions['pft_1d'].append(outputs['pft_1d'].cpu()) all_predictions['soil_2d'].append(outputs['soil_2d'].cpu()) From 1d9d9d15ba5aa5e0c06196e25f2366a4c43a4fa4 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Tue, 7 Oct 2025 08:36:11 -0400 Subject: [PATCH 04/51] modified the model config file based on better understanding of number of words --- CNP_model_config_27M.txt | 16 ++++++++-------- CNP_model_config_v01.txt | 24 ++++++++++++++++-------- 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/CNP_model_config_27M.txt b/CNP_model_config_27M.txt index 85e4478..15ec598 100644 --- a/CNP_model_config_27M.txt +++ b/CNP_model_config_27M.txt @@ -5,12 +5,12 @@ [ENCODERS] # LSTM encoder -lstm_hidden_size = 128 +lstm_hidden_size = 256 # Surface/static encoder (half used as output size) -static_fc_size = 98 +static_fc_size = 256 -# PFT parameters encoder (FC in the 27M run) +# PFT parameters encoder (use CNN in this config) pft_param_cnn_channels = [32, 64, 128] pft_param_cnn_kernel_size = 3 pft_param_cnn_padding = 1 @@ -25,7 +25,7 @@ water_fc_size = 0 scalar_fc_size = 64 # 1D PFT encoder -pft_1d_fc_size = 64 +pft_1d_fc_size = 256 [SOIL2D_CNN] # Soil 2D encoder channels @@ -35,13 +35,13 @@ conv_padding = 1 [TRANSFORMER] # Feature fusion transformer -num_tokens = 2 -token_dim = 512 +num_tokens = 7 +token_dim = 256 transformer_layers = 8 -transformer_heads = 32 +transformer_heads = 16 # Global dropout probability -dropout_p = 0.2 +dropout_p = 0.1 [OUTPUTS] # Output geometry (do not set scalar/vector/matrix sizes here) diff --git a/CNP_model_config_v01.txt b/CNP_model_config_v01.txt index b1d4438..137188a 100644 --- a/CNP_model_config_v01.txt +++ b/CNP_model_config_v01.txt @@ -2,42 +2,50 @@ # Mirrors the current defaults in get_cnp_combined_config(). # Format: key = value. Lists can be comma-separated or Python lists. # Section headers (in brackets) are optional and ignored by the parser. +# CNP Model configuration approximating the 3.8 M-parameter run (2025-10-07) [ENCODERS] # LSTM encoder -lstm_hidden_size = 64 +lstm_hidden_size = 256 +# => Forcing embedding: 256 # Surface/static encoder -static_fc_size = 64 +static_fc_size = 256 +# => Surface embedding: 128 (due to final // 2 in the model) # PFT parameters encoder (use CNN by default for CNP model) -pft_param_cnn_channels = [32, 64] +pft_param_cnn_channels = [32, 64, 128] pft_param_cnn_kernel_size = 3 pft_param_cnn_padding = 1 use_cnn_for_pft_param = true +pft_param_size = 44 num_pfts = 17 +# => PFT-parameter embedding (CNN + GAP): 64 # Water encoder (disabled by default in training; keep head size at 0) water_fc_size = 0 # Scalar encoder scalar_fc_size = 64 +# => Scalar embedding: 32 (two-layer MLP projects to 32) # 1D PFT encoder -pft_1d_fc_size = 64 +pft_1d_fc_size = 256 +# => PFT state embedding: 256 [SOIL2D_CNN] # Soil 2D encoders (both 1D and 2D branches use these channel sizes) -conv_channels = [16, 32, 64] +conv_channels = [32, 64, 128] conv_kernel_size = 3 conv_padding = 1 +# => Soil 2D embedding: 128 (Flatten -> 128 -> 128 head) [TRANSFORMER] # Feature fusion transformer -num_tokens = 8 +num_tokens = 7 token_dim = 128 -transformer_layers = 4 -transformer_heads = 8 +transformer_layers = 8 +transformer_heads = 16 # Global dropout probability # Set to 0.0 for strict determinism in experiments if needed From 255fc4a69357dfa2b1b1cb4ca309b692163bc34d Mon Sep 17 00:00:00 2001 From: Zhuowei Gu Date: Thu, 9 Oct 2025 15:18:33 -0400 Subject: [PATCH 05/51] Add AMD GPU setup and environment scripts (see daliwang/LandSim#4) --- scripts/requirements_amd.txt | 16 ++++++++++++++++ scripts/setup_and_train_amd.sh | 30 ++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) create mode 100644 scripts/requirements_amd.txt create mode 100644 scripts/setup_and_train_amd.sh diff --git a/scripts/requirements_amd.txt b/scripts/requirements_amd.txt new file mode 100644 index 0000000..765b0be --- /dev/null +++ b/scripts/requirements_amd.txt @@ -0,0 +1,16 @@ +pandas>=2.3.0 +numpy>=2.0.0 +scikit-learn>=1.5.0 +matplotlib>=3.8.0 +seaborn>=0.13.0 +scipy>=1.11.0 +xarray>=2024.6.0 +netCDF4>=1.7.0 +tqdm>=4.67.0 +psutil>=7.0.0 +joblib>=1.5.0 +optuna>=3.0.0 +jupyter>=1.1.0 +ipykernel>=6.29.0 +notebook>=7.4.0 +pynvml>=11.5.0 diff --git a/scripts/setup_and_train_amd.sh b/scripts/setup_and_train_amd.sh new file mode 100644 index 0000000..a85d7c0 --- /dev/null +++ b/scripts/setup_and_train_amd.sh @@ -0,0 +1,30 @@ +#!/bin/bash + +module load PrgEnv-gnu/8.6.0 +module load miniforge3/23.11.0-0 +module load rocm/6.4.1 +module load craype-accel-amd-gfx90a + +conda create -p ../amd_env python=3.12 -c conda-forge +source activate ../amd_env + +pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/rocm6.4 + +module unload rocm +MPICC="cc -shared" pip install --no-cache-dir --no-binary=mpi4py mpi4py +module load rocm/6.4.1 + +pip install -r requirements_amd.txt + +# If you get MIOpen errors (miopenStatusInternalError, readonly database), uncomment the line below: +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH=../miopen_cache +mkdir -p "$MIOPEN_USER_DB_PATH" + + +# HIP cache path +# export HIP_COMPILE_CACHE_DIR=/lustre/orion/csc665/world-shared/zhuowei/LandSim/hip_cache +# mkdir -p "$HIP_COMPILE_CACHE_DIR" +# chmod -R 700 "$HIP_COMPILE_CACHE_DIR" + +python ../train_model.py From 37f9b004a942411738535accac57c9af8b6b89f5 Mon Sep 17 00:00:00 2001 From: Zhuowei Gu Date: Thu, 9 Oct 2025 16:35:41 -0400 Subject: [PATCH 06/51] Add PyTorch profiling workflow for NVIDIA/AMD GPUs (see daliwang/LandSim#3) --- profile_train_with_pytorch_profiler.py | 138 +++++++++++++++++++++++++ scripts/run_pytorch_profiler.sh | 53 ++++++++++ 2 files changed, 191 insertions(+) create mode 100755 profile_train_with_pytorch_profiler.py create mode 100644 scripts/run_pytorch_profiler.sh diff --git a/profile_train_with_pytorch_profiler.py b/profile_train_with_pytorch_profiler.py new file mode 100755 index 0000000..202ed97 --- /dev/null +++ b/profile_train_with_pytorch_profiler.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +""" +PyTorch Profiler wrapper for train_model.py +Run this script without modifying train_model.py source code +""" + +import os +import sys +from pathlib import Path + +import torch +from torch.profiler import ProfilerActivity, profile + +# Add project root to Python path +project_root = Path(__file__).parent +sys.path.insert(0, str(project_root)) + +def setup_profiler_output_dir() -> Path: + """Create profiler output directory""" + out_dir = Path("pytorch_profiler_logs") + out_dir.mkdir(parents=True, exist_ok=True) + return out_dir + +def main(): + """Main function: run train_model.py under profiler environment""" + + # Detect available device activities + activities = [ProfilerActivity.CPU] + if torch.cuda.is_available(): + activities.append(ProfilerActivity.CUDA) + print(f"Detected CUDA/ROCm device, will analyze both CPU and GPU performance") + else: + print("Only CPU detected, will analyze CPU performance only") + + # Prepare output directory + output_dir = setup_profiler_output_dir() + trace_path = output_dir / "train_model_trace.json" + stats_path = output_dir / "train_model_stats.txt" + + print(f"Starting PyTorch Profiler analysis of train_model.py...") + print(f"Output directory: {output_dir}") + + try: + # Check if TensorBoard logs already exist + tensorboard_logs_dir = output_dir / "tensorboard_logs" + existing_traces = list(tensorboard_logs_dir.glob("*.pt.trace.json")) if tensorboard_logs_dir.exists() else [] + + if existing_traces: + print(f"✓ Found existing TensorBoard trace files:") + for trace_file in existing_traces: + print(f" - {trace_file.name} ({trace_file.stat().st_size / (1024**3):.2f} GB)") + print(f"✓ Skipping training and profiler generation (using existing traces)") + else: + print(f"✓ No existing traces found, running training with profiler...") + # Import train_model main function (avoid importing inside profiler context) + from train_model import main as train_main + + # Run training under profiler context + with profile( + activities=activities, + record_shapes=True, + profile_memory=True, + with_stack=True, # Record call stack + on_trace_ready=torch.profiler.tensorboard_trace_handler(str(tensorboard_logs_dir)), + ) as prof: + train_main() + + # Note: Chrome trace export is disabled when using tensorboard_trace_handler + # because the trace is already saved by the handler + print(f"✓ TensorBoard logs saved to: {tensorboard_logs_dir}") + print(f"✓ Chrome trace export skipped (already saved by tensorboard handler)") + + # Initialize memory_path variable + memory_path = None + + # Generate performance statistics summary (only if we have profiler data) + if not existing_traces: + if torch.cuda.is_available(): + sort_key = "cuda_time_total" + else: + sort_key = "cpu_time_total" + + stats_table = prof.key_averages().table(sort_by=sort_key, row_limit=20) + + # Save statistics summary to file + with open(stats_path, 'w', encoding='utf-8') as f: + f.write("PyTorch Profiler Performance Analysis Report\n") + f.write("=" * 50 + "\n\n") + f.write(f"Analysis Device: {'CPU + GPU' if torch.cuda.is_available() else 'CPU'}\n") + f.write(f"Sort By: {sort_key}\n\n") + f.write(stats_table) + + print(f"✓ Performance statistics summary saved to: {stats_path}") + + # Display top 10 most time-consuming operations in console + print("\nTop 10 Most Time-Consuming Operations:") + print("-" * 80) + print(prof.key_averages().table(sort_by=sort_key, row_limit=10)) + + # Memory usage statistics + if torch.cuda.is_available(): + memory_stats = prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10) + print("\nTop 10 Memory-Intensive Operations:") + print("-" * 80) + print(memory_stats) + + memory_path = output_dir / "train_model_memory_stats.txt" + with open(memory_path, 'w', encoding='utf-8') as f: + f.write("PyTorch Profiler Memory Usage Analysis Report\n") + f.write("=" * 50 + "\n\n") + f.write(memory_stats) + print(f"✓ Memory usage statistics saved to: {memory_path}") + else: + print(f"✓ Using existing trace files - statistics already available in TensorBoard") + print(f"✓ Check existing statistics files: {stats_path}") + # Check if memory stats file exists from previous run + memory_path = output_dir / "train_model_memory_stats.txt" + if not memory_path.exists(): + memory_path = None + + print(f"\nAnalysis completed! All files saved in: {output_dir}") + print(f"\nView Results:") + print(f"1. Performance summary: Check {stats_path}") + if torch.cuda.is_available(): + print(f"2. Memory statistics: Check {memory_path}") + print(f"3. TensorBoard visualization: Run 'tensorboard --logdir={output_dir / 'tensorboard_logs'}' then open http://localhost:6006") + print(f" - Click on 'PyTorch Profiler' tab") + print(f" - View operator shapes, memory usage, and detailed timeline") + print(f" - Chrome trace is available in TensorBoard interface") + + except Exception as e: + print(f"❌ Profiler execution failed: {str(e)}") + import traceback + traceback.print_exc() + sys.exit(1) + +if __name__ == "__main__": + main() diff --git a/scripts/run_pytorch_profiler.sh b/scripts/run_pytorch_profiler.sh new file mode 100644 index 0000000..04ca606 --- /dev/null +++ b/scripts/run_pytorch_profiler.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# PyTorch Profiler execution script +# Based on setup_and_train_amd.sh environment configuration + +echo "=== PyTorch Profiler Performance Analysis Script ===" + +# Load necessary modules +echo "Loading modules..." +module load PrgEnv-gnu/8.6.0 +module load miniforge3/23.11.0-0 +module load rocm/6.4.1 +module load craype-accel-amd-gfx90a + +# Activate conda environment +echo "Activating conda environment..." +source activate ../amd_env + +# Set ROCm/HIP cache environment variables (avoid cache errors) +echo "Setting ROCm environment variables..." +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH=../miopen_cache +mkdir -p "$MIOPEN_USER_DB_PATH" +chmod 700 "$MIOPEN_USER_DB_PATH" + +export HIP_COMPILE_CACHE_DIR=../hip_cache +mkdir -p "$HIP_COMPILE_CACHE_DIR" +chmod 700 "$HIP_COMPILE_CACHE_DIR" + +export HSA_OVERRIDE_GFX_VERSION=10.3.0 + +# Switch to project directory +cd ../LandSim + +echo "Current directory: $(pwd)" +echo "Python version: $(python --version)" +echo "PyTorch version: $(python -c 'import torch; print(torch.__version__)')" +echo "CUDA/ROCm available: $(python -c 'import torch; print(torch.cuda.is_available())')" + +# Run PyTorch Profiler +echo "" +echo "Starting PyTorch Profiler analysis..." +echo "==================================" + +python profile_train_with_pytorch_profiler.py + +echo "" +echo "==================================" +echo "PyTorch Profiler analysis completed!" +echo "Check results directory: pytorch_profiler_logs/" +echo "Main files:" +echo " - train_model_trace.json (Chrome timeline)" +echo " - train_model_stats.txt (Performance statistics)" +echo " - train_model_memory_stats.txt (Memory statistics, if GPU available)" From 841dfd0d29b8ec3d31a8478244f5e6ce5031967c Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Thu, 9 Oct 2025 23:35:46 -0400 Subject: [PATCH 07/51] Add random shuffling options to training script and improve time series data handling in DataLoader. Clean NaN/Inf values across various data processing steps. --- config/training_config.py | 2 +- data/data_loader_individual.py | 38 ++++++++++++++++++++++++++++++---- train_cnp_model.py | 33 +++++++++++++++++++++++++++++ 3 files changed, 68 insertions(+), 5 deletions(-) diff --git a/config/training_config.py b/config/training_config.py index 1f702cf..1b41190 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -574,7 +574,7 @@ def get_cnp_combined_config( file_patterns.append("enhanced_1_training_data_batch_*.pkl") if use_trendy05: data_paths.append("/mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP") - file_patterns.append("1_training_data_batch_*.pkl") + file_patterns.append("enhanced_1_training_data_batch_*.pkl") file_pattern = file_patterns[0] if len(file_patterns) == 1 else file_patterns config.update_data_config( diff --git a/data/data_loader_individual.py b/data/data_loader_individual.py index 8c92248..d19c02f 100644 --- a/data/data_loader_individual.py +++ b/data/data_loader_individual.py @@ -165,9 +165,26 @@ def preprocess_data(self): for col in self.data_config.time_series_columns: if col in self.df.columns: # Ensure time series data is properly formatted - self.df[col] = self.df[col].apply( - lambda x: np.array(x, dtype=np.float32) if isinstance(x, (list, np.ndarray)) else np.zeros(self.data_config.time_series_length, dtype=np.float32) - ) + def _to_ts_and_truncate(x): + # Convert to numpy array (float32) and truncate/pad to configured time_series_length + target_len = int(getattr(self.data_config, 'time_series_length', 240)) + if isinstance(x, (list, np.ndarray)): + arr = np.array(x, dtype=np.float32).flatten() + # Prefer earliest 20-year window as per repeated forcing spec + if arr.size >= target_len: + arr = arr[:target_len] + else: + # pad to target_len with zeros at the end + pad = target_len - arr.size + if pad > 0: + arr = np.pad(arr, (0, pad), mode='constant') + # ensure no NaN/Inf + arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0) + return arr + # Fallback to zeros of target length + return np.zeros(target_len, dtype=np.float32) + + self.df[col] = self.df[col].apply(_to_ts_and_truncate) # Process list columns logger.info("Processing list columns...") @@ -851,6 +868,8 @@ def _normalize_static(self, static_columns: List[str]) -> Tuple[torch.Tensor, An for i, col in enumerate(static_columns): assert col in self.df.columns, f"Static column '{col}' missing in DataFrame!" static_data = self.df[static_columns].values + # Clean NaN/Inf + static_data = np.nan_to_num(static_data, nan=0.0, posinf=0.0, neginf=0.0) scaler = self._get_scaler(self.preprocessing_config.static_normalization) static_normalized = scaler.fit_transform(static_data) return torch.tensor(static_normalized, dtype=self.preprocessing_config.data_type), scaler @@ -864,6 +883,8 @@ def _normalize_scalar(self) -> Tuple[torch.Tensor, Any]: assert col in self.df.columns, f"Scalar column '{col}' missing in DataFrame!" scalar_data = self.df[scalar_columns].values + # Clean NaN/Inf + scalar_data = np.nan_to_num(scalar_data, nan=0.0, posinf=0.0, neginf=0.0) # Use group normalization scaler = self._get_scaler(self.preprocessing_config.target_normalization) @@ -880,6 +901,8 @@ def _normalize_y_scalar(self) -> Tuple[torch.Tensor, Any]: assert col in self.df.columns, f"y_scalar column '{col}' missing in DataFrame!" y_scalar_data = self.df[y_scalar_columns].values + # Clean NaN/Inf + y_scalar_data = np.nan_to_num(y_scalar_data, nan=0.0, posinf=0.0, neginf=0.0) # Use group normalization scaler = self._get_scaler(self.preprocessing_config.target_normalization) @@ -896,6 +919,7 @@ def _normalize_scalar_individual(self, transform_only: bool = False) -> Tuple[to assert col in self.df.columns, f"Scalar column '{col}' missing in DataFrame!" scalar_data = self.df[scalar_columns].values + scalar_data = np.nan_to_num(scalar_data, nan=0.0, posinf=0.0, neginf=0.0) # Use individual normalization (fit+transform or transform-only) if transform_only: @@ -914,6 +938,7 @@ def _normalize_y_scalar_individual(self, transform_only: bool = False) -> Tuple[ assert col in self.df.columns, f"y_scalar column '{col}' missing in DataFrame!" y_scalar_data = self.df[y_scalar_columns].values + y_scalar_data = np.nan_to_num(y_scalar_data, nan=0.0, posinf=0.0, neginf=0.0) # Use individual normalization (fit+transform or transform-only) if transform_only: @@ -931,6 +956,8 @@ def _normalize_list_1d_individual(self, columns: List[str], transform_only: bool assert col in self.df.columns, f"1D column '{col}' missing in DataFrame!" col_data = [np.vstack(self.df[col].values) for col in columns] + # Clean NaN/Inf in stacked data + col_data = [np.nan_to_num(cd, nan=0.0, posinf=0.0, neginf=0.0) for cd in col_data] data = np.stack(col_data, axis=1) # shape: (samples, features, length) # Handle PFT0 dropping for compatibility with model expectations @@ -1097,6 +1124,7 @@ def _normalize_list_2d_individual(self, columns: List[str], transform_only: bool for val in values: if isinstance(val, (list, np.ndarray)): val_array = np.array(val) + val_array = np.nan_to_num(val_array, nan=0.0, posinf=0.0, neginf=0.0) if val_array.shape[1] == 15: # Has 15 layers # Extract first column and top 10 layers immediately extracted = val_array[0:1, 0:10] # Shape: (1, 10) @@ -1219,7 +1247,9 @@ def _normalize_pft_param(self) -> Tuple[torch.Tensor, Any]: for col in pft_param_columns: val = row[col] if isinstance(val, (list, np.ndarray)) and len(val) == num_pfts: - row_vectors.append(np.array(val, dtype=np.float32)) + arr = np.array(val, dtype=np.float32) + arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0) + row_vectors.append(arr) else: row_vectors.append(np.zeros(num_pfts, dtype=np.float32)) row_matrix = np.stack(row_vectors, axis=0) # [44, 17] diff --git a/train_cnp_model.py b/train_cnp_model.py index 4619738..47c592a 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -162,6 +162,19 @@ def main(): help='Override global dropout probability in the model (e.g., 0.0 to disable)' ) + # Shuffling controls + parser.add_argument( + '--random-shuffle', + action='store_true', + help='Enable random shuffling for dataset rows and DataLoader (default: fixed seed shuffling)' + ) + parser.add_argument( + '--shuffle-seed', + type=int, + default=None, + help='Optional seed to use when --random-shuffle is enabled (default: no fixed seed)' + ) + parser.add_argument( '--use-trendy1', action='store_true', @@ -318,6 +331,26 @@ def main(): logger.warning(f"Failed to set litter loss weights: {e}") logger.info(f"Effective learning rate for this run: {effective_lr}") + # Shuffling policy: fixed vs random + if args.random_shuffle: + # Use provided shuffle seed or system randomness + if args.shuffle_seed is not None: + config.update_data_config(random_state=int(args.shuffle_seed)) + config.update_training_config(random_seed=int(args.shuffle_seed)) + logger.info(f"Random shuffling enabled with shuffle_seed={args.shuffle_seed}") + else: + # Remove fixed seeds to allow non-deterministic shuffling + # Keep a log message for provenance + logger.info("Random shuffling enabled with no fixed seed (non-deterministic shuffles)") + # Use a time-based seed for DataLoader generator consistency per run + import time + dyn_seed = int(time.time()) % (2**31 - 1) + config.update_data_config(random_state=dyn_seed) + config.update_training_config(random_seed=dyn_seed) + else: + # Keep fixed seeds for fair comparisons + logger.info("Fixed shuffling (seeded) enabled for fair comparison") + # Optional strict determinism (opt-in via CLI) if args.strict_determinism: seed = getattr(config.training_config, 'random_seed', 42) From e1b6b50394268769e4d893a98a180edd41602512 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Fri, 10 Oct 2025 07:50:03 -0700 Subject: [PATCH 08/51] Drop data samples at longitude 0 degree --- .gitignore | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d00d3c1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +# Ignore result and cache folders +cnp_results/ +__pycache__/ +*.pyc +logs/ From 5bfb24a968be9f148ebdbac41a5ae0975ca241f0 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Fri, 10 Oct 2025 08:05:57 -0700 Subject: [PATCH 09/51] =?UTF-8?q?Update=20data=20loader=20&=20training=20c?= =?UTF-8?q?onfig;=20drop=20samples=20at=20longitude=200=C2=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CNP_IO_updated9_modified_drop0degree.txt | 59 ++++++++++++++++++++++++ config/training_config.py | 25 ++++++++-- data/data_loader_individual.py | 21 +++++++++ 3 files changed, 102 insertions(+), 3 deletions(-) create mode 100644 CNP_IO_updated9_modified_drop0degree.txt diff --git a/CNP_IO_updated9_modified_drop0degree.txt b/CNP_IO_updated9_modified_drop0degree.txt new file mode 100644 index 0000000..dd5a620 --- /dev/null +++ b/CNP_IO_updated9_modified_drop0degree.txt @@ -0,0 +1,59 @@ +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, NPP, AR, HR + +1D PFT VARIABLES (39 variables): + +• deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage +• deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage + +• leafc, leafc_storage, leafn, leafn_storage, leafp, leafp_storage +• frootc, frootc_storage, frootn, frootn_storage, frootp, frootp_storage + +• livestemc, livestemc_storage, livestemn, livestemn_storage, livestemp, livestemp_storage, +• livecrootc, livecrootc_storage, livecrootn, livecrootn_storage, livecrootp, livecrootp_storage + +• cpool, npool, ppool + +• tlai, totvegc + +2D VARIABLES (layered - 28 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 + +• soil1c_vr, soil1n_vr, soil1p_vr +• soil2c_vr, soil2n_vr, soil2p_vr +• soil3c_vr, soil3n_vr, soil3p_vr +• soil4c_vr, soil4n_vr, soil4p_vr + +• labilep_vr , occlp_vr, primp_vr, secondp_vr diff --git a/config/training_config.py b/config/training_config.py index 1f702cf..549a7fe 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -101,6 +101,9 @@ class DataConfig: # New parameter for filtering NaN in time series filter_time_series_nan: bool = False filter_column: Optional[str] = None # Added for CNP model + + # Longitude filtering - list of longitude values to drop from dataset + longitudes_to_drop: List[float] = field(default_factory=list) @@ -411,9 +414,11 @@ def parse_cnp_io_list(filename): filename (str): Path to the variable list file (e.g., CNP_IO_list_general.txt) Returns: dict: Mapping of variable group keys to lists of variable names. + Also includes 'longitudes_to_drop' key if specified in the file. """ # Map section titles to config keys section_map = { + 'LONGITUDE FILTERING': 'longitudes_to_drop', 'TIME SERIES VARIABLES': 'time_series_variables', 'SURFACE PROPERTIES': 'surface_properties', 'PFT PARAMETERS': 'pft_parameters', @@ -430,6 +435,7 @@ def parse_cnp_io_list(filename): with open(filename) as f: for line in f: line = line.strip() + # Section header detection for section_title, key in section_map.items(): if line.startswith(section_title): @@ -440,7 +446,17 @@ def parse_cnp_io_list(filename): if current_section and line.startswith('•'): # Remove bullet and split by comma, filter out empty strings vars_ = [v.strip() for v in line[1:].split(',') if v.strip()] - result[current_section].extend(vars_) + + # Special handling for longitude filtering - convert to floats + if current_section == 'longitudes_to_drop': + try: + longitudes = [float(x) for x in vars_] + result[current_section].extend(longitudes) + logging.info(f"Parsed longitudes to drop: {longitudes}") + except Exception as e: + logging.warning(f"Failed to parse longitudes to drop: {e}") + else: + result[current_section].extend(vars_) # Some variables are listed as comma-separated after a bullet elif current_section and ',' in line and not line.startswith('['): vars_ = [v.strip('• ').strip() for v in line.split(',') if v.strip('• ').strip()] @@ -570,7 +586,7 @@ def get_cnp_combined_config( data_paths = [] file_patterns = [] if use_trendy1: - data_paths.append("/mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_1_data_CNP") + data_paths.append("/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree") file_patterns.append("enhanced_1_training_data_batch_*.pkl") if use_trendy05: data_paths.append("/mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP") @@ -629,6 +645,7 @@ def get_cnp_combined_config( ] # If a variable list file is provided, parse it + longitudes_to_drop = [] if variable_list_path is not None: parsed = parse_cnp_io_list(variable_list_path) time_series_columns = parsed.get('time_series_variables', default_time_series) @@ -638,6 +655,7 @@ def get_cnp_combined_config( scalar_variables = parsed.get('scalar_variables', default_scalar) pft_1d_variables = parsed.get('pft_1d_variables', default_pft_1d) variables_2d_soil = parsed.get('variables_2d_soil', default_2d_soil) + longitudes_to_drop = parsed.get('longitudes_to_drop', []) else: time_series_columns = default_time_series surface_properties = default_surface @@ -653,7 +671,8 @@ def get_cnp_combined_config( pft_param_columns=pft_parameters, x_list_scalar_columns=scalar_variables, x_list_columns_1d=pft_1d_variables, - x_list_columns_2d=variables_2d_soil + x_list_columns_2d=variables_2d_soil, + longitudes_to_drop=longitudes_to_drop ) if include_water: data_config_kwargs['x_list_water_columns'] = water_variables diff --git a/data/data_loader_individual.py b/data/data_loader_individual.py index 8c92248..cf563db 100644 --- a/data/data_loader_individual.py +++ b/data/data_loader_individual.py @@ -151,6 +151,27 @@ def preprocess_data(self): """Preprocess the loaded data.""" logger.info("Starting data preprocessing...") + # Filter samples by longitude if specified + if hasattr(self.data_config, 'longitudes_to_drop') and self.data_config.longitudes_to_drop: + if 'Longitude' in self.df.columns: + original_size = len(self.df) + longitudes_to_drop = self.data_config.longitudes_to_drop + logger.info(f"Filtering samples with longitudes: {longitudes_to_drop}") + + # Create a mask for samples to keep (those NOT in the drop list) + # Use a tolerance for floating point comparison + tolerance = 0.01 + mask = ~self.df['Longitude'].apply( + lambda lon: any(abs(lon - drop_lon) < tolerance for drop_lon in longitudes_to_drop) + ) + + self.df = self.df[mask].reset_index(drop=True) + filtered_size = len(self.df) + dropped_count = original_size - filtered_size + logger.info(f"Longitude filtering: {original_size} samples -> {filtered_size} samples (dropped {dropped_count} samples)") + else: + logger.warning("'Longitude' column not found in dataset. Cannot apply longitude filtering.") + # Drop specified columns if hasattr(self.data_config, 'filter_columns') and self.data_config.filter_columns: for col in self.data_config.filter_columns: From 7ada6f825a3dc95bb89e5e560a01ffe0af24eefd Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Fri, 10 Oct 2025 23:54:17 -0400 Subject: [PATCH 10/51] turn on mask-absent-pfts and improve CNP_pipeline_runbook.md --- docs/CNP_pipeline_runbook.md | 13 +++++++------ train_cnp_model.py | 8 ++++++++ 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 9885ba7..e8b730e 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -23,17 +23,18 @@ cd cnp_results/run_YYYYMMDD_HHMMSS # e.g., cnp_results/run_20250815_205419 ``` ### 4) Validate predictions vs ground truth (test split) using the current run directory as default -Generates scatter plots and statistics using individually normalized results. +Generates quick statistics and prediction quality report ```bash -python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & +python ../../scripts/cnp_result_validationplot.py --stats-only +python ../../scripts/generate_prediction_quality_report.py ``` (optional) -For a quick result, (--stats-only) option can be used to the following quality report - -### 4.5) Generate comprehensive prediction quality report -Creates detailed quality analysis categorizing predictions as "good", "ok", or "bad" based on statistical thresholds. +For a detail scatter plot of each variables and its substructure +```bash +python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & +```bash ```bash python ../../scripts/generate_prediction_quality_report.py > prediction_quality_report.log 2>&1 & diff --git a/train_cnp_model.py b/train_cnp_model.py index 47c592a..84d5777 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -235,9 +235,17 @@ def main(): ) parser.add_argument( '--mask-absent-pfts', + dest='mask_absent_pfts', action='store_true', help='Zero predictions where PCT_NAT_PFT_k == 0 and exclude from loss' ) + 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) args = parser.parse_args() From 99b758a6e0e0639a83ce25fee96cbf3169bdb38a Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Sat, 11 Oct 2025 12:23:58 -0400 Subject: [PATCH 11/51] improve analysis scripts --- docs/README_prediction_quality.md | 15 ++ scripts/generate_prediction_quality_report.py | 170 +++++++++++++++++- 2 files changed, 178 insertions(+), 7 deletions(-) diff --git a/docs/README_prediction_quality.md b/docs/README_prediction_quality.md index 13b7a51..d98cd60 100644 --- a/docs/README_prediction_quality.md +++ b/docs/README_prediction_quality.md @@ -49,6 +49,18 @@ The `generate_prediction_quality_report.py` script accepts the following command --mae-ok FLOAT Relative MAE threshold for ok predictions Default: 0.25 + +--force-xlim-01 Force R² x-axis limits to [0, 1] in the scatter plot + Default: enabled (use --no-force-xlim-01 to disable) + +--print-scatter-stats Print min/max and counts for R² and relative RMSE used + in the scatter plot + Default: enabled (use --no-print-scatter-stats to disable) + +--bad-html-limit N Max number of bad rows shown in HTML (default 100) +--bad-text-limit N Max number of bad rows printed in text report (default 200) +--export-bad Export bad predictions to CSV (default enabled) +--no-export-bad Do not export bad predictions CSV ``` ## Example Usage @@ -80,6 +92,9 @@ The analysis generates the following output files: 5. **prediction_quality_by_variable.png**: Bar chart showing quality distribution by variable 6. **overall_prediction_quality.png**: Pie chart showing overall quality distribution 7. **r2_vs_rmse.png**: Scatter plot of R² vs Relative RMSE +8. **bad_predictions_detailed.csv**: Full list of predictions classified as "bad" (export can be disabled) + +In addition, the text report now includes a "Bad Predictions Summary" and a limited detailed list (controlled by `--bad-text-limit`), and the HTML report includes a table of the first N bad predictions (controlled by `--bad-html-limit`). ## Classification Criteria diff --git a/scripts/generate_prediction_quality_report.py b/scripts/generate_prediction_quality_report.py index 17a472d..0d686c0 100644 --- a/scripts/generate_prediction_quality_report.py +++ b/scripts/generate_prediction_quality_report.py @@ -24,6 +24,27 @@ def main(): 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') args = parser.parse_args() # Set up input and output paths @@ -91,7 +112,16 @@ def categorize_prediction(row): df['prediction_quality'] = df.apply(categorize_prediction, axis=1) # Filter out rows that are just coordinates (Longitude, Latitude) - analysis_df = df[~df['pft'].isin(['Longitude', 'Latitude'])] + analysis_df = df[~df['pft'].isin(['Longitude', 'Latitude'])].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) @@ -109,6 +139,13 @@ def categorize_prediction(row): # 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'}") @@ -164,12 +201,25 @@ def categorize_prediction(row): # Create a copy of the dataframe to avoid SettingWithCopyWarning scatter_df = analysis_df.copy() - - # Calculate relative RMSE - 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 - ) + # 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( @@ -192,6 +242,10 @@ def categorize_prediction(row): 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 = [ @@ -226,6 +280,69 @@ def categorize_prediction(row): 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: @@ -364,6 +481,45 @@ def categorize_prediction(row): R² vs Relative RMSE +

Detailed Bad Predictions (first {args.bad_html_limit})

+ + + + + + + + + + + + + """ + + # 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

From 0e76d9bac1809655144b36ed16196b2b79888844 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Sat, 11 Oct 2025 14:57:37 -0400 Subject: [PATCH 12/51] Enhance README and runbook documentation; improve validation scripts with top-bad variable filtering and reporting options. Add functionality to generate top-bad plots and refine analysis methods for better prediction quality assessment. --- README.md | 16 ++- docs/CNP_pipeline_runbook.md | 30 ++-- scripts/cnp_result_validationplot.py | 130 ++++++++++++++++-- scripts/generate_prediction_quality_report.py | 41 +++++- 4 files changed, 187 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 1f16ed4..ebe1c0c 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,7 @@ Notes: ## 🔗 CNP pipeline workflow (from runbook) -Follow this streamlined workflow using a user-defined `CNP_IO` list (see `docs/CNP_pipeline_runbook.md` for details): +Follow this streamlined workflow using a user-defined `CNP_IO` list. For the fully detailed, continuously updated instructions, see `docs/CNP_pipeline_runbook.md`. 1) Create your `CNP_IO` list (e.g., `CNP_IO_demo1.txt`). @@ -92,8 +92,10 @@ cd cnp_results/run_YYYYMMDD_HHMMSS 4) Validate predictions vs ground truth (test split): ```bash -python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & +python ../../scripts/cnp_result_validationplot.py --stats-only +python ../../scripts/generate_prediction_quality_report.py ``` +(details and options in `docs/CNP_pipeline_runbook.md`) (extra note: use check_pft1d_predictions.py and check_soil2d_predictions.py to find prediction abnormality) ```bash python ../../scripts/check_pft1d_predictions.py > check_pft1d_predictions.log 2>&1 & @@ -157,6 +159,16 @@ python scripts/compare_cnp_runs.py \ --- +## 📘 Pipeline runbook reference + +See `docs/CNP_pipeline_runbook.md` for: +- End-to-end run instructions +- Prediction quality report details (bad predictions CSV, HTML/PNG outputs) +- Top-bad-only plotting and how `analysis/top_bad_plots/` is generated +- CLI flags to customize plots and reports + +--- + ## 🔁 Updating restart files with AI predictions Preferred workflow uses `scripts/ai_predictions_to_restart.py` after generating predictions (see workflow above). For bespoke flows, the legacy helper is available: diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index e8b730e..107eb3c 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -22,29 +22,25 @@ Notes: cd cnp_results/run_YYYYMMDD_HHMMSS # e.g., cnp_results/run_20250815_205419 ``` -### 4) Validate predictions vs ground truth (test split) using the current run directory as default -Generates quick statistics and prediction quality report +### 4) Validate predictions vs ground truth (test split) +Generates quick statistics and a prediction quality report. The report now also creates filtered plots for the “top variables by bad-count”. ```bash python ../../scripts/cnp_result_validationplot.py --stats-only python ../../scripts/generate_prediction_quality_report.py ``` -(optional) -For a detail scatter plot of each variables and its substructure -```bash -python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & -```bash - -```bash -python ../../scripts/generate_prediction_quality_report.py > prediction_quality_report.log 2>&1 & -``` - -**Output**: Creates `analysis/` directory with: -- Quality assessment CSV files -- Interactive HTML report -- Visualization charts (bar charts, pie charts, scatter plots) -- Text summary report +Options and behavior: +- The quality report saves outputs under `analysis/`: + - `detailed_quality_assessment.csv`, `variable_quality_summary.csv`, `overall_prediction_quality.png`, `prediction_quality_by_variable.png`, `r2_vs_rmse.png`, `prediction_quality_report.html`, `quality_summary_report.txt`. + - `bad_predictions_detailed.csv` (full list of rows classified as bad; can be disabled). + - `top_bad_plots/` folder with filtered scatter plots for variables listed as “top variables by bad-count”. +- Useful flags: + - `--no-top-bad-plots`: skip generating `analysis/top_bad_plots/`. + - `--bad-html-limit N`, `--bad-text-limit N`, `--no-export-bad`. + - `--force-xlim-01/--no-force-xlim-01`, `--print-scatter-stats/--no-print-scatter-stats`. + +**Output**: `analysis/` contains CSV/PNG/HTML reports plus `top_bad_plots/`. **Quality Classification**: - **Good**: R² ≥ 0.9, Relative RMSE ≤ 0.1, Relative MAE ≤ 0.1 diff --git a/scripts/cnp_result_validationplot.py b/scripts/cnp_result_validationplot.py index 18f1b78..5c89c67 100644 --- a/scripts/cnp_result_validationplot.py +++ b/scripts/cnp_result_validationplot.py @@ -1,4 +1,5 @@ import os +import re import pandas as pd import numpy as np import matplotlib.pyplot as plt @@ -19,14 +20,83 @@ def plot_gt_vs_pred(gt, pred, title, save_path): plt.savefig(save_path) plt.close() -def main_with_flag(results_dir, plot_scatter, plot_loss): +def _parse_top_bad_report(report_path): + """Parse quality_summary_report.txt to extract variables and associated PFT indices or layer numbers. + + Returns a mapping: { variable: { 'pfts': set[int], 'layers': set[int] } } + """ + selection = {} + if not os.path.exists(report_path): + print(f"Top-bad report not found: {report_path}") + return selection + + in_section = False + try: + with open(report_path, 'r') as f: + for line in f: + stripped = line.strip('\n') + header = stripped.strip() + # Detect start of section (support both old and new headings) + if header.startswith('Top variables by bad-count'): + in_section = True + continue + # Section ends at next heading or blank line followed by a heading; we keep it simple + if in_section and header.startswith('## '): + break + if in_section and stripped.startswith(' '): + # Example formats: + # ppool: 15; pfts: 1, 2, 3, ... + # smin_no3_vr: 10; layers: 1, 2, ... + # var: 5 + m = re.match(r"\s+([A-Za-z0-9_]+):\s*([0-9]+)(;.*)?$", stripped) + if not m: + continue + var = m.group(1) + details = m.group(3) or '' + pfts = set() + layers = set() + if 'pfts:' in details: + m_p = re.search(r"pfts:\s*([0-9,\s]+)", details) + if m_p: + nums = [n.strip() for n in m_p.group(1).split(',') if n.strip()] + for n in nums: + try: + pfts.add(int(n)) + except Exception: + pass + if 'layers:' in details: + m_l = re.search(r"layers:\s*([0-9,\s]+)", details) + if m_l: + nums = [n.strip() for n in m_l.group(1).split(',') if n.strip()] + for n in nums: + try: + layers.add(int(n)) + except Exception: + pass + selection[var] = { 'pfts': pfts, 'layers': layers } + except Exception as e: + print(f"Failed to parse top-bad report {report_path}: {e}") + 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): # Create plots subdirectory - plots_dir = os.path.join(results_dir, "plots") + plots_dir = plots_dir_override or os.path.join(results_dir, "plots") os.makedirs(plots_dir, exist_ok=True) # NEW: Create stats file stats_path = os.path.join(results_dir, "validation_stats.csv") stats_data = [] + + # Optional: restrict plotting to top-bad variables (and selected PFTs/layers) + selection = None + if top_bad_only: + report_path = top_bad_report or os.path.join(results_dir, 'analysis', 'quality_summary_report.txt') + selection = _parse_top_bad_report(report_path) + if selection: + print(f"Plotting restricted to top-bad variables from: {report_path}") + else: + print("No selections parsed from top-bad report; proceeding without restriction.") # Check for new directory structure first pft_gt_dir = os.path.join(results_dir, 'cnp_predictions', 'pft_1d_ground_truth') @@ -35,7 +105,7 @@ def main_with_flag(results_dir, plot_scatter, plot_loss): # 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) + analyze_1d_new_structure(results_dir, '1D', plots_dir, stats_data, plot_scatter, selection) else: # Fall back to old single-file format print("Using legacy 1D single-file format") @@ -60,7 +130,7 @@ def main_with_flag(results_dir, plot_scatter, plot_loss): 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) + analyze_2d_new_structure(results_dir, '2D', plots_dir, stats_data, plot_scatter, selection) else: # Fall back to old single-file format print("Using legacy 2D single-file format") @@ -93,12 +163,15 @@ def main_with_flag(results_dir, plot_scatter, plot_loss): else: print("test_metrics.csv not found.") -def analyze_pair(gt_path, pred_path, label, out_dir, stats_data, per_column=False, plot_scatter=True): +def analyze_pair(gt_path, pred_path, label, out_dir, stats_data, per_column=False, plot_scatter=True, selection=None): gt = pd.read_csv(gt_path) pred = pd.read_csv(pred_path) if per_column: # Per-column comparison for scalar for col in gt.columns: + # If selection provided, only include scalar variables present in selection + if selection is not None and col not in selection: + continue if col in pred.columns: print(f"Analyzing variable: {col}") gt_col = gt[col].values.flatten() @@ -177,7 +250,7 @@ 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): +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""" 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') @@ -205,6 +278,9 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt print(f"Missing prediction file for {var_name}: {pred_file}") continue + # Skip if restricting to top-bad variables and this variable is not selected + if selection is not None and var_name not in selection: + continue print(f"Analyzing variable: {var_name}") # Read data @@ -227,6 +303,18 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt print(f" {var_name}: {num_pfts} PFT columns") for pft_idx, col_name in enumerate(gt_data.columns): + # If selection provided, attempt to parse PFT index from col_name like 'Y_var_pftX' + if selection is not None: + sel = selection.get(var_name, None) + if sel is not None and sel['pfts']: + pft_match = re.search(r"pft(\d+)$", col_name) + if pft_match: + try: + pft_num = int(pft_match.group(1)) + if pft_num not in sel['pfts']: + continue + except Exception: + pass gt_col = gt_data[col_name].values pred_col = pred_data[col_name].values @@ -279,7 +367,7 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt 'pred_sum': pred_stats['sum'] }) -def analyze_1d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot_scatter=True): +def analyze_1d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot_scatter=True, selection=None): """Legacy function for old single-file 1D format - kept for compatibility""" gt = pd.read_csv(gt_path) pred = pd.read_csv(pred_path) @@ -314,6 +402,9 @@ def analyze_1d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot pred_reshaped = pred.values.reshape(num_samples, num_vars, num_pfts) # For each variable and each PFT column, compare for i, var in enumerate(variable_names): + # Skip if restricting to top-bad variables and this variable is not selected + if selection is not None and var not in selection: + continue print(f"Analyzing variable: {var}") for j in range(num_pfts): gt_col = gt_reshaped[:, i, j] @@ -348,7 +439,7 @@ 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): +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""" 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') @@ -376,6 +467,9 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt print(f"Missing prediction file for {var_name}: {pred_file}") continue + # Skip if restricting to top-bad variables and this variable is not selected + if selection is not None and var_name not in selection: + continue print(f"Analyzing 2D variable: {var_name}") # Read data @@ -413,6 +507,11 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt layers_to_analyze = 10 for layer_idx in range(layers_to_analyze): + # If selection provided, enforce layer filtering (1-based indexing in report) + if selection is not None: + sel = selection.get(var_name, None) + if sel is not None and sel['layers'] and (layer_idx + 1) not in sel['layers']: + continue # Calculate the correct column index: first_column * 10 + layer # If only 10 columns exist, each column is a layer col_idx = layer_idx if num_columns == 1 else (first_column_idx * 10 + layer_idx) @@ -487,7 +586,7 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt except Exception as e: print(f" Skipped overall plot for {var_name}: {e}") -def analyze_2d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot_scatter=True): +def analyze_2d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot_scatter=True, selection=None): """Legacy function for old single-file 2D format - kept for compatibility""" gt = pd.read_csv(gt_path) pred = pd.read_csv(pred_path) @@ -517,8 +616,16 @@ def analyze_2d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot pred_reshaped = pred.values.reshape(num_samples, num_vars, num_columns, num_layers_per_column) # For each variable, analyze all 10 layers of the first column for i, var in enumerate(variable_names): + # Skip if restricting to top-bad variables and this variable is not selected + if selection is not None and var not in selection: + continue print(f"Analyzing 2D variable: {var}") for j in range(10): # Changed from 15 to 10 to analyze all predicted layers + # If selection provided, enforce layer filtering (1-based indexing in report) + if selection is not None: + sel = selection.get(var, None) + if sel is not None and sel['layers'] and (j + 1) not in sel['layers']: + continue gt_col = gt_reshaped[:, i, 0, j] # First column (index 0), all 10 layers pred_col = pred_reshaped[:, i, 0, j] # First column (index 0), all 10 layers @@ -581,6 +688,9 @@ def plot_train_val_accuracy(loss_csv, out_dir): parser.add_argument('--no-plot-loss', action='store_false', dest='plot_loss', help='Do not plot train/val loss curve') # NEW: Stats-only mode disables all plots but still computes and saves statistics parser.add_argument('--stats-only', action='store_true', help='Only compute and save statistics CSV; do not generate any plots') + # NEW: Restrict plotting to top-bad variables from summary report + parser.add_argument('--top-bad-only', action='store_true', help='Plot only variables listed in the quality summary top-bad section') + 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.set_defaults(plot_scatter=True, plot_loss=True) args = parser.parse_args() @@ -593,4 +703,4 @@ 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) \ No newline at end of file + main_with_flag(args.results_dir, args.plot_scatter, args.plot_loss, args.top_bad_only, args.top_bad_report) \ No newline at end of file diff --git a/scripts/generate_prediction_quality_report.py b/scripts/generate_prediction_quality_report.py index 0d686c0..27031df 100644 --- a/scripts/generate_prediction_quality_report.py +++ b/scripts/generate_prediction_quality_report.py @@ -5,6 +5,8 @@ import seaborn as sns from pathlib import Path import argparse +import importlib.util +import sys def main(): parser = argparse.ArgumentParser(description='Generate prediction quality report from validation statistics') @@ -45,6 +47,10 @@ def main(): 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') args = parser.parse_args() # Set up input and output paths @@ -112,7 +118,13 @@ def categorize_prediction(row): df['prediction_quality'] = df.apply(categorize_prediction, axis=1) # Filter out rows that are just coordinates (Longitude, Latitude) - analysis_df = df[~df['pft'].isin(['Longitude', 'Latitude'])].copy() + # 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'] @@ -259,6 +271,33 @@ def categorize_prediction(row): plt.tight_layout() plt.savefig(output_dir / "r2_vs_rmse.png", dpi=300) + # 4. Optionally generate top-bad-only plots into a subfolder using the validation plotting utility + 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) + # 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 + 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=True, + top_bad_report=str(output_dir / 'quality_summary_report.txt'), + plots_dir_override=top_bad_out) + print(f"Top-bad plots saved to: {top_bad_out}") + else: + print("Warning: cnp_result_validationplot.main_with_flag not found; skipping top-bad plots") + 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: From 5da2f283f1ddd37b1688d407dd2ad080a4eff394 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Fri, 10 Oct 2025 23:54:17 -0400 Subject: [PATCH 13/51] turn on mask-absent-pfts and improve CNP_pipeline_runbook.md --- docs/CNP_pipeline_runbook.md | 13 +++++++------ train_cnp_model.py | 8 ++++++++ 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 9885ba7..e8b730e 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -23,17 +23,18 @@ cd cnp_results/run_YYYYMMDD_HHMMSS # e.g., cnp_results/run_20250815_205419 ``` ### 4) Validate predictions vs ground truth (test split) using the current run directory as default -Generates scatter plots and statistics using individually normalized results. +Generates quick statistics and prediction quality report ```bash -python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & +python ../../scripts/cnp_result_validationplot.py --stats-only +python ../../scripts/generate_prediction_quality_report.py ``` (optional) -For a quick result, (--stats-only) option can be used to the following quality report - -### 4.5) Generate comprehensive prediction quality report -Creates detailed quality analysis categorizing predictions as "good", "ok", or "bad" based on statistical thresholds. +For a detail scatter plot of each variables and its substructure +```bash +python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & +```bash ```bash python ../../scripts/generate_prediction_quality_report.py > prediction_quality_report.log 2>&1 & diff --git a/train_cnp_model.py b/train_cnp_model.py index 47c592a..84d5777 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -235,9 +235,17 @@ def main(): ) parser.add_argument( '--mask-absent-pfts', + dest='mask_absent_pfts', action='store_true', help='Zero predictions where PCT_NAT_PFT_k == 0 and exclude from loss' ) + 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) args = parser.parse_args() From 207c72a122483c9dfad94ed1a61b56dc13b21890 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Sat, 11 Oct 2025 12:23:58 -0400 Subject: [PATCH 14/51] improve analysis scripts --- docs/README_prediction_quality.md | 15 ++ scripts/generate_prediction_quality_report.py | 170 +++++++++++++++++- 2 files changed, 178 insertions(+), 7 deletions(-) diff --git a/docs/README_prediction_quality.md b/docs/README_prediction_quality.md index 13b7a51..d98cd60 100644 --- a/docs/README_prediction_quality.md +++ b/docs/README_prediction_quality.md @@ -49,6 +49,18 @@ The `generate_prediction_quality_report.py` script accepts the following command --mae-ok FLOAT Relative MAE threshold for ok predictions Default: 0.25 + +--force-xlim-01 Force R² x-axis limits to [0, 1] in the scatter plot + Default: enabled (use --no-force-xlim-01 to disable) + +--print-scatter-stats Print min/max and counts for R² and relative RMSE used + in the scatter plot + Default: enabled (use --no-print-scatter-stats to disable) + +--bad-html-limit N Max number of bad rows shown in HTML (default 100) +--bad-text-limit N Max number of bad rows printed in text report (default 200) +--export-bad Export bad predictions to CSV (default enabled) +--no-export-bad Do not export bad predictions CSV ``` ## Example Usage @@ -80,6 +92,9 @@ The analysis generates the following output files: 5. **prediction_quality_by_variable.png**: Bar chart showing quality distribution by variable 6. **overall_prediction_quality.png**: Pie chart showing overall quality distribution 7. **r2_vs_rmse.png**: Scatter plot of R² vs Relative RMSE +8. **bad_predictions_detailed.csv**: Full list of predictions classified as "bad" (export can be disabled) + +In addition, the text report now includes a "Bad Predictions Summary" and a limited detailed list (controlled by `--bad-text-limit`), and the HTML report includes a table of the first N bad predictions (controlled by `--bad-html-limit`). ## Classification Criteria diff --git a/scripts/generate_prediction_quality_report.py b/scripts/generate_prediction_quality_report.py index 17a472d..0d686c0 100644 --- a/scripts/generate_prediction_quality_report.py +++ b/scripts/generate_prediction_quality_report.py @@ -24,6 +24,27 @@ def main(): 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') args = parser.parse_args() # Set up input and output paths @@ -91,7 +112,16 @@ def categorize_prediction(row): df['prediction_quality'] = df.apply(categorize_prediction, axis=1) # Filter out rows that are just coordinates (Longitude, Latitude) - analysis_df = df[~df['pft'].isin(['Longitude', 'Latitude'])] + analysis_df = df[~df['pft'].isin(['Longitude', 'Latitude'])].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) @@ -109,6 +139,13 @@ def categorize_prediction(row): # 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'}") @@ -164,12 +201,25 @@ def categorize_prediction(row): # Create a copy of the dataframe to avoid SettingWithCopyWarning scatter_df = analysis_df.copy() - - # Calculate relative RMSE - 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 - ) + # 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( @@ -192,6 +242,10 @@ def categorize_prediction(row): 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 = [ @@ -226,6 +280,69 @@ def categorize_prediction(row): 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: @@ -364,6 +481,45 @@ def categorize_prediction(row): R² vs Relative RMSE +

Detailed Bad Predictions (first {args.bad_html_limit})

+
+ + + + + + + + + + + + """ + + # 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

From 8cef42b1e5bcb65216d6435e026ec1943c2a7abf Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Sat, 11 Oct 2025 14:57:37 -0400 Subject: [PATCH 15/51] Enhance README and runbook documentation; improve validation scripts with top-bad variable filtering and reporting options. Add functionality to generate top-bad plots and refine analysis methods for better prediction quality assessment. --- README.md | 16 ++- docs/CNP_pipeline_runbook.md | 30 ++-- scripts/cnp_result_validationplot.py | 130 ++++++++++++++++-- scripts/generate_prediction_quality_report.py | 41 +++++- 4 files changed, 187 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 1f16ed4..ebe1c0c 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,7 @@ Notes: ## 🔗 CNP pipeline workflow (from runbook) -Follow this streamlined workflow using a user-defined `CNP_IO` list (see `docs/CNP_pipeline_runbook.md` for details): +Follow this streamlined workflow using a user-defined `CNP_IO` list. For the fully detailed, continuously updated instructions, see `docs/CNP_pipeline_runbook.md`. 1) Create your `CNP_IO` list (e.g., `CNP_IO_demo1.txt`). @@ -92,8 +92,10 @@ cd cnp_results/run_YYYYMMDD_HHMMSS 4) Validate predictions vs ground truth (test split): ```bash -python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & +python ../../scripts/cnp_result_validationplot.py --stats-only +python ../../scripts/generate_prediction_quality_report.py ``` +(details and options in `docs/CNP_pipeline_runbook.md`) (extra note: use check_pft1d_predictions.py and check_soil2d_predictions.py to find prediction abnormality) ```bash python ../../scripts/check_pft1d_predictions.py > check_pft1d_predictions.log 2>&1 & @@ -157,6 +159,16 @@ python scripts/compare_cnp_runs.py \ --- +## 📘 Pipeline runbook reference + +See `docs/CNP_pipeline_runbook.md` for: +- End-to-end run instructions +- Prediction quality report details (bad predictions CSV, HTML/PNG outputs) +- Top-bad-only plotting and how `analysis/top_bad_plots/` is generated +- CLI flags to customize plots and reports + +--- + ## 🔁 Updating restart files with AI predictions Preferred workflow uses `scripts/ai_predictions_to_restart.py` after generating predictions (see workflow above). For bespoke flows, the legacy helper is available: diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index e8b730e..107eb3c 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -22,29 +22,25 @@ Notes: cd cnp_results/run_YYYYMMDD_HHMMSS # e.g., cnp_results/run_20250815_205419 ``` -### 4) Validate predictions vs ground truth (test split) using the current run directory as default -Generates quick statistics and prediction quality report +### 4) Validate predictions vs ground truth (test split) +Generates quick statistics and a prediction quality report. The report now also creates filtered plots for the “top variables by bad-count”. ```bash python ../../scripts/cnp_result_validationplot.py --stats-only python ../../scripts/generate_prediction_quality_report.py ``` -(optional) -For a detail scatter plot of each variables and its substructure -```bash -python ../../scripts/cnp_result_validationplot.py > cnp_results_validation.log 2>&1 & -```bash - -```bash -python ../../scripts/generate_prediction_quality_report.py > prediction_quality_report.log 2>&1 & -``` - -**Output**: Creates `analysis/` directory with: -- Quality assessment CSV files -- Interactive HTML report -- Visualization charts (bar charts, pie charts, scatter plots) -- Text summary report +Options and behavior: +- The quality report saves outputs under `analysis/`: + - `detailed_quality_assessment.csv`, `variable_quality_summary.csv`, `overall_prediction_quality.png`, `prediction_quality_by_variable.png`, `r2_vs_rmse.png`, `prediction_quality_report.html`, `quality_summary_report.txt`. + - `bad_predictions_detailed.csv` (full list of rows classified as bad; can be disabled). + - `top_bad_plots/` folder with filtered scatter plots for variables listed as “top variables by bad-count”. +- Useful flags: + - `--no-top-bad-plots`: skip generating `analysis/top_bad_plots/`. + - `--bad-html-limit N`, `--bad-text-limit N`, `--no-export-bad`. + - `--force-xlim-01/--no-force-xlim-01`, `--print-scatter-stats/--no-print-scatter-stats`. + +**Output**: `analysis/` contains CSV/PNG/HTML reports plus `top_bad_plots/`. **Quality Classification**: - **Good**: R² ≥ 0.9, Relative RMSE ≤ 0.1, Relative MAE ≤ 0.1 diff --git a/scripts/cnp_result_validationplot.py b/scripts/cnp_result_validationplot.py index 18f1b78..5c89c67 100644 --- a/scripts/cnp_result_validationplot.py +++ b/scripts/cnp_result_validationplot.py @@ -1,4 +1,5 @@ import os +import re import pandas as pd import numpy as np import matplotlib.pyplot as plt @@ -19,14 +20,83 @@ def plot_gt_vs_pred(gt, pred, title, save_path): plt.savefig(save_path) plt.close() -def main_with_flag(results_dir, plot_scatter, plot_loss): +def _parse_top_bad_report(report_path): + """Parse quality_summary_report.txt to extract variables and associated PFT indices or layer numbers. + + Returns a mapping: { variable: { 'pfts': set[int], 'layers': set[int] } } + """ + selection = {} + if not os.path.exists(report_path): + print(f"Top-bad report not found: {report_path}") + return selection + + in_section = False + try: + with open(report_path, 'r') as f: + for line in f: + stripped = line.strip('\n') + header = stripped.strip() + # Detect start of section (support both old and new headings) + if header.startswith('Top variables by bad-count'): + in_section = True + continue + # Section ends at next heading or blank line followed by a heading; we keep it simple + if in_section and header.startswith('## '): + break + if in_section and stripped.startswith(' '): + # Example formats: + # ppool: 15; pfts: 1, 2, 3, ... + # smin_no3_vr: 10; layers: 1, 2, ... + # var: 5 + m = re.match(r"\s+([A-Za-z0-9_]+):\s*([0-9]+)(;.*)?$", stripped) + if not m: + continue + var = m.group(1) + details = m.group(3) or '' + pfts = set() + layers = set() + if 'pfts:' in details: + m_p = re.search(r"pfts:\s*([0-9,\s]+)", details) + if m_p: + nums = [n.strip() for n in m_p.group(1).split(',') if n.strip()] + for n in nums: + try: + pfts.add(int(n)) + except Exception: + pass + if 'layers:' in details: + m_l = re.search(r"layers:\s*([0-9,\s]+)", details) + if m_l: + nums = [n.strip() for n in m_l.group(1).split(',') if n.strip()] + for n in nums: + try: + layers.add(int(n)) + except Exception: + pass + selection[var] = { 'pfts': pfts, 'layers': layers } + except Exception as e: + print(f"Failed to parse top-bad report {report_path}: {e}") + 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): # Create plots subdirectory - plots_dir = os.path.join(results_dir, "plots") + plots_dir = plots_dir_override or os.path.join(results_dir, "plots") os.makedirs(plots_dir, exist_ok=True) # NEW: Create stats file stats_path = os.path.join(results_dir, "validation_stats.csv") stats_data = [] + + # Optional: restrict plotting to top-bad variables (and selected PFTs/layers) + selection = None + if top_bad_only: + report_path = top_bad_report or os.path.join(results_dir, 'analysis', 'quality_summary_report.txt') + selection = _parse_top_bad_report(report_path) + if selection: + print(f"Plotting restricted to top-bad variables from: {report_path}") + else: + print("No selections parsed from top-bad report; proceeding without restriction.") # Check for new directory structure first pft_gt_dir = os.path.join(results_dir, 'cnp_predictions', 'pft_1d_ground_truth') @@ -35,7 +105,7 @@ def main_with_flag(results_dir, plot_scatter, plot_loss): # 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) + analyze_1d_new_structure(results_dir, '1D', plots_dir, stats_data, plot_scatter, selection) else: # Fall back to old single-file format print("Using legacy 1D single-file format") @@ -60,7 +130,7 @@ def main_with_flag(results_dir, plot_scatter, plot_loss): 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) + analyze_2d_new_structure(results_dir, '2D', plots_dir, stats_data, plot_scatter, selection) else: # Fall back to old single-file format print("Using legacy 2D single-file format") @@ -93,12 +163,15 @@ def main_with_flag(results_dir, plot_scatter, plot_loss): else: print("test_metrics.csv not found.") -def analyze_pair(gt_path, pred_path, label, out_dir, stats_data, per_column=False, plot_scatter=True): +def analyze_pair(gt_path, pred_path, label, out_dir, stats_data, per_column=False, plot_scatter=True, selection=None): gt = pd.read_csv(gt_path) pred = pd.read_csv(pred_path) if per_column: # Per-column comparison for scalar for col in gt.columns: + # If selection provided, only include scalar variables present in selection + if selection is not None and col not in selection: + continue if col in pred.columns: print(f"Analyzing variable: {col}") gt_col = gt[col].values.flatten() @@ -177,7 +250,7 @@ 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): +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""" 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') @@ -205,6 +278,9 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt print(f"Missing prediction file for {var_name}: {pred_file}") continue + # Skip if restricting to top-bad variables and this variable is not selected + if selection is not None and var_name not in selection: + continue print(f"Analyzing variable: {var_name}") # Read data @@ -227,6 +303,18 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt print(f" {var_name}: {num_pfts} PFT columns") for pft_idx, col_name in enumerate(gt_data.columns): + # If selection provided, attempt to parse PFT index from col_name like 'Y_var_pftX' + if selection is not None: + sel = selection.get(var_name, None) + if sel is not None and sel['pfts']: + pft_match = re.search(r"pft(\d+)$", col_name) + if pft_match: + try: + pft_num = int(pft_match.group(1)) + if pft_num not in sel['pfts']: + continue + except Exception: + pass gt_col = gt_data[col_name].values pred_col = pred_data[col_name].values @@ -279,7 +367,7 @@ def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt 'pred_sum': pred_stats['sum'] }) -def analyze_1d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot_scatter=True): +def analyze_1d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot_scatter=True, selection=None): """Legacy function for old single-file 1D format - kept for compatibility""" gt = pd.read_csv(gt_path) pred = pd.read_csv(pred_path) @@ -314,6 +402,9 @@ def analyze_1d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot pred_reshaped = pred.values.reshape(num_samples, num_vars, num_pfts) # For each variable and each PFT column, compare for i, var in enumerate(variable_names): + # Skip if restricting to top-bad variables and this variable is not selected + if selection is not None and var not in selection: + continue print(f"Analyzing variable: {var}") for j in range(num_pfts): gt_col = gt_reshaped[:, i, j] @@ -348,7 +439,7 @@ 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): +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""" 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') @@ -376,6 +467,9 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt print(f"Missing prediction file for {var_name}: {pred_file}") continue + # Skip if restricting to top-bad variables and this variable is not selected + if selection is not None and var_name not in selection: + continue print(f"Analyzing 2D variable: {var_name}") # Read data @@ -413,6 +507,11 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt layers_to_analyze = 10 for layer_idx in range(layers_to_analyze): + # If selection provided, enforce layer filtering (1-based indexing in report) + if selection is not None: + sel = selection.get(var_name, None) + if sel is not None and sel['layers'] and (layer_idx + 1) not in sel['layers']: + continue # Calculate the correct column index: first_column * 10 + layer # If only 10 columns exist, each column is a layer col_idx = layer_idx if num_columns == 1 else (first_column_idx * 10 + layer_idx) @@ -487,7 +586,7 @@ def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, plot_scatt except Exception as e: print(f" Skipped overall plot for {var_name}: {e}") -def analyze_2d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot_scatter=True): +def analyze_2d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot_scatter=True, selection=None): """Legacy function for old single-file 2D format - kept for compatibility""" gt = pd.read_csv(gt_path) pred = pd.read_csv(pred_path) @@ -517,8 +616,16 @@ def analyze_2d(gt_path, pred_path, label, out_dir, results_dir, stats_data, plot pred_reshaped = pred.values.reshape(num_samples, num_vars, num_columns, num_layers_per_column) # For each variable, analyze all 10 layers of the first column for i, var in enumerate(variable_names): + # Skip if restricting to top-bad variables and this variable is not selected + if selection is not None and var not in selection: + continue print(f"Analyzing 2D variable: {var}") for j in range(10): # Changed from 15 to 10 to analyze all predicted layers + # If selection provided, enforce layer filtering (1-based indexing in report) + if selection is not None: + sel = selection.get(var, None) + if sel is not None and sel['layers'] and (j + 1) not in sel['layers']: + continue gt_col = gt_reshaped[:, i, 0, j] # First column (index 0), all 10 layers pred_col = pred_reshaped[:, i, 0, j] # First column (index 0), all 10 layers @@ -581,6 +688,9 @@ def plot_train_val_accuracy(loss_csv, out_dir): parser.add_argument('--no-plot-loss', action='store_false', dest='plot_loss', help='Do not plot train/val loss curve') # NEW: Stats-only mode disables all plots but still computes and saves statistics parser.add_argument('--stats-only', action='store_true', help='Only compute and save statistics CSV; do not generate any plots') + # NEW: Restrict plotting to top-bad variables from summary report + parser.add_argument('--top-bad-only', action='store_true', help='Plot only variables listed in the quality summary top-bad section') + 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.set_defaults(plot_scatter=True, plot_loss=True) args = parser.parse_args() @@ -593,4 +703,4 @@ 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) \ No newline at end of file + main_with_flag(args.results_dir, args.plot_scatter, args.plot_loss, args.top_bad_only, args.top_bad_report) \ No newline at end of file diff --git a/scripts/generate_prediction_quality_report.py b/scripts/generate_prediction_quality_report.py index 0d686c0..27031df 100644 --- a/scripts/generate_prediction_quality_report.py +++ b/scripts/generate_prediction_quality_report.py @@ -5,6 +5,8 @@ import seaborn as sns from pathlib import Path import argparse +import importlib.util +import sys def main(): parser = argparse.ArgumentParser(description='Generate prediction quality report from validation statistics') @@ -45,6 +47,10 @@ def main(): 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') args = parser.parse_args() # Set up input and output paths @@ -112,7 +118,13 @@ def categorize_prediction(row): df['prediction_quality'] = df.apply(categorize_prediction, axis=1) # Filter out rows that are just coordinates (Longitude, Latitude) - analysis_df = df[~df['pft'].isin(['Longitude', 'Latitude'])].copy() + # 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'] @@ -259,6 +271,33 @@ def categorize_prediction(row): plt.tight_layout() plt.savefig(output_dir / "r2_vs_rmse.png", dpi=300) + # 4. Optionally generate top-bad-only plots into a subfolder using the validation plotting utility + 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) + # 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 + 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=True, + top_bad_report=str(output_dir / 'quality_summary_report.txt'), + plots_dir_override=top_bad_out) + print(f"Top-bad plots saved to: {top_bad_out}") + else: + print("Warning: cnp_result_validationplot.main_with_flag not found; skipping top-bad plots") + 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: From 2dda8244fb8fd30a08b26c8f6e7678a8fd26be62 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Sat, 11 Oct 2025 17:37:27 -0400 Subject: [PATCH 16/51] Add CNP_IO configuration file and update training config parser to support new dataset paths and file patterns --- ...drop0degree.txt => CNP_IO_updated9_dev.txt | 5 ++ config/training_config.py | 65 ++++++++++++++++--- 2 files changed, 60 insertions(+), 10 deletions(-) rename CNP_IO_updated9_modified_drop0degree.txt => CNP_IO_updated9_dev.txt (91%) diff --git a/CNP_IO_updated9_modified_drop0degree.txt b/CNP_IO_updated9_dev.txt similarity index 91% rename from CNP_IO_updated9_modified_drop0degree.txt rename to CNP_IO_updated9_dev.txt index dd5a620..f40c693 100644 --- a/CNP_IO_updated9_modified_drop0degree.txt +++ b/CNP_IO_updated9_dev.txt @@ -1,3 +1,8 @@ +TRENDY1_PATH: /mnt/proj-shared/AI4BGC_7xw/TrainingData//Trendy_1_data_CNP +TRENDY05_PATH = /mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP +DATA_PATHS: /path/extra1, /path/extra2 +FILE_PATTERN: enhanced_1_training_data_batch_*.pkl + LONGITUDE FILTERING - 2 longitudes: • 0, 358.75 diff --git a/config/training_config.py b/config/training_config.py index 27683d1..10a6e0f 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -426,10 +426,17 @@ def parse_cnp_io_list(filename): 'TEMPERATURE VARIABLES': 'temperature_variables', 'SCALAR VARIABLES': 'scalar_variables', '1D PFT VARIABLES': 'pft_1d_variables', - '2D VARIABLES': 'variables_2d_soil' + '2D VARIABLES': 'variables_2d_soil', + 'DATA PATHS': 'data_paths' } # Prepare result dict result = {v: [] for v in section_map.values()} + # Additional single-value keys for dataset configuration + result.update({ + 'trendy1_path': None, + 'trendy05_path': None, + 'file_pattern': None + }) current_section = None with open(filename) as f: @@ -466,6 +473,26 @@ def parse_cnp_io_list(filename): # Only add if it's not a description or exclusion if re.match(r'^[A-Za-z0-9_]+$', line): result[current_section].append(line) + # Outside of a section or in any section, allow key=value dataset config + # e.g., TRENDY1_PATH: /path/to/trendy1 + # TRENDY05_PATH = /path/to/trendy05 + # FILE_PATTERN: enhanced_1_training_data_batch_*.pkl + # DATA_PATHS: /p1,/p2 + if line and not line.startswith('#'): + kv_match = re.match(r'(?i)^(trendy1_path|trendy05_path|file_pattern|data_paths)\s*[:=]\s*(.+)$', line) + if kv_match: + key = kv_match.group(1).lower() + val = kv_match.group(2).strip() + if key == 'data_paths': + # Support comma-separated list + paths = [p.strip() for p in val.split(',') if p.strip()] + result['data_paths'].extend(paths) + elif key == 'file_pattern': + result['file_pattern'] = val + elif key == 'trendy1_path': + result['trendy1_path'] = val + elif key == 'trendy05_path': + result['trendy05_path'] = val return result def parse_cnp_model_config(filename: str) -> Dict[str, Any]: @@ -584,14 +611,32 @@ def get_cnp_combined_config( """ config = TrainingConfigManager() data_paths = [] - file_patterns = [] - if use_trendy1: - data_paths.append("/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree") - file_patterns.append("enhanced_1_training_data_batch_*.pkl") - if use_trendy05: - data_paths.append("/mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP") - file_patterns.append("enhanced_1_training_data_batch_*.pkl") - file_pattern = file_patterns[0] if len(file_patterns) == 1 else file_patterns + file_pattern = None + # If a variable list is provided, prefer dataset paths from it + parsed = None + if variable_list_path is not None: + try: + parsed = parse_cnp_io_list(variable_list_path) + except Exception as e: + logging.warning(f"Failed to parse variable list for data paths: {e}") + if parsed is not None: + # Collect from any or all of: data_paths, trendy1_path, trendy05_path + if parsed.get('data_paths'): + data_paths.extend([p for p in parsed['data_paths'] if p]) + if parsed.get('trendy1_path') and use_trendy1: + data_paths.append(parsed['trendy1_path']) + if parsed.get('trendy05_path') and use_trendy05: + data_paths.append(parsed['trendy05_path']) + if parsed.get('file_pattern'): + file_pattern = parsed['file_pattern'] + # Fallback to defaults if none provided via CNP_IO + if not data_paths: + if use_trendy1: + data_paths.append("/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree") + if use_trendy05: + data_paths.append("/mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP") + if file_pattern is None: + file_pattern = "enhanced_1_training_data_batch_*.pkl" config.update_data_config( data_paths=data_paths, @@ -646,7 +691,7 @@ def get_cnp_combined_config( # If a variable list file is provided, parse it longitudes_to_drop = [] - if variable_list_path is not None: + if variable_list_path is not None and parsed is None: parsed = parse_cnp_io_list(variable_list_path) time_series_columns = parsed.get('time_series_variables', default_time_series) surface_properties = parsed.get('surface_properties', default_surface) From 04c5471153bd6e7f66fb898663f61d42f2cc56af Mon Sep 17 00:00:00 2001 From: Qinglei Cao Date: Mon, 6 Oct 2025 13:58:25 +0000 Subject: [PATCH 17/51] Add CI when pull request --- .github/workflows/ci.yml | 73 ++++++++++++++++++++++++++++++++++++++++ scripts/prepare_data.sh | 46 +++++++++++++++++++++++++ train_cnp_model.py | 59 +++++++++++++++++++++++++++++--- 3 files changed, 174 insertions(+), 4 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 scripts/prepare_data.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..63c74da --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,73 @@ +name: CI + +on: + pull_request: + +permissions: + contents: read + +concurrency: + group: ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + build-and-test: + runs-on: ubuntu-latest + timeout-minutes: 30 + env: + OMP_NUM_THREADS: 1 + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + cache: 'pip' + + - name: Install system deps + run: | + sudo apt-get update + sudo apt-get install -y git + + - name: Set up venv + run: | + python -m venv .venv + source .venv/bin/activate + python -m pip install -U pip + python -m pip install -r requirements.txt || true + # Ensure CPU-only PyTorch is available in CI + pip install --extra-index-url https://download.pytorch.org/whl/cpu torch torchvision torchaudio || true + + - name: Quick smoke test (example dataset) + run: | + source .venv/bin/activate + python train_model.py + + - name: Prepare dataset (download if missing) + env: + DATA_DIR: ${{ github.workspace }}/dataset/trendy1 + run: | + chmod +x scripts/prepare_data.sh + bash scripts/prepare_data.sh "$DATA_DIR" + + - name: CNP training smoke test (dataset, 1 file, 1 epoch) + env: + DATA_DIR: ${{ github.workspace }}/dataset/trendy1 + run: | + source .venv/bin/activate + python train_cnp_model.py \ + --epoch 1 \ + --variable-list ./CNP_IO_updated14_xfer.txt \ + --mask-absent-pfts \ + --model-config ./CNP_model_config_v01.txt \ + --use-trendy1 \ + --data-paths "$DATA_DIR" \ + --file-pattern 'enhanced_1_training_data_batch_*.pkl' \ + --max-files 1 + + - name: CNP training config validation (no data) + run: | + source .venv/bin/activate + python train_cnp_model.py --config-only --use-trendy1 --epoch 1 --batch-size 8 diff --git a/scripts/prepare_data.sh b/scripts/prepare_data.sh new file mode 100644 index 0000000..61f3947 --- /dev/null +++ b/scripts/prepare_data.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Usage: +# scripts/prepare_data.sh /path/to/target_dir [gdrive_folder_url] + +# Inputs +DATA_DIR="${1:-/tmp}" +GDRIVE_URL="${2:-https://drive.google.com/drive/folders/1djTADtAqxlYml8GwdwxUEAi-1xR_Yqn5}" + +if [[ "${1:-}" == "" ]]; then + echo "[prepare_data] No target directory provided; defaulting to /tmp" +fi + +echo "[prepare_data] Target dir: ${DATA_DIR}" +mkdir -p "${DATA_DIR}" + +# If directory already has at least one pkl file matching expected pattern, skip +shopt -s nullglob +existing=("${DATA_DIR}"/enhanced_1_training_data_batch_*.pkl) +if (( ${#existing[@]} > 0 )); then + echo "[prepare_data] Found ${#existing[@]} data files, skipping download." + exit 0 +fi + +echo "[prepare_data] No data found. Ensuring gdown is available..." + +# If a virtualenv is active, use it; otherwise try local .venv; else system python +if [[ -n "${VIRTUAL_ENV:-}" ]]; then + : # already in a venv +elif [[ -f ".venv/bin/activate" ]]; then + # shellcheck disable=SC1091 + source .venv/bin/activate || true +fi + +python -m pip --version >/dev/null 2>&1 || python -m ensurepip --upgrade || true +python -m pip install -q --upgrade pip || true +python -m pip install -q gdown || true + +echo "[prepare_data] Downloading dataset from Google Drive to ${DATA_DIR}..." +gdown --folder --continue --remaining-ok --fuzzy "${GDRIVE_URL}" -O "${DATA_DIR}" + +echo "[prepare_data] Done. Listing downloaded files:" +ls -lh "${DATA_DIR}" | sed 's/^/[prepare_data] /' + + diff --git a/train_cnp_model.py b/train_cnp_model.py index 84d5777..3d11274 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -132,7 +132,8 @@ def main(): help='Output directory for results' ) parser.add_argument( - '--epochs', + '--epochs', '--epoch', + dest='epochs', type=int, default=150, help='Number of training epochs' @@ -197,6 +198,18 @@ def main(): default=None, help='Path to model config text file (e.g., CNP_model_config.txt) to override encoders/transformer/MLPs' ) + parser.add_argument( + '--data-paths', + type=str, + default=None, + help='Comma-separated list of directories containing training .pkl batches (overrides variable list)' + ) + parser.add_argument( + '--file-pattern', + type=str, + default=None, + help='Glob pattern for training files (e.g., enhanced_1_training_data_batch_*.pkl)' + ) parser.add_argument( '--max-files', type=int, @@ -215,6 +228,11 @@ def main(): default=None, help='Extra loss weight applied to xsmrpool (non-positive pool)' ) + parser.add_argument( + '--config-only', + action='store_true', + help='Exit after building configuration; do not load data or train (for CI)' + ) parser.add_argument( '--litter-c-loss-weight', type=float, @@ -287,16 +305,34 @@ def main(): variable_list_path=args.variable_list, model_config_path=args.model_config ) + # Optional overrides via CLI/env to avoid fixed variable list files + env_data_paths = os.environ.get('DATA_PATHS') or os.environ.get('CNP_DATA_PATHS') + env_file_pattern = os.environ.get('FILE_PATTERN') or os.environ.get('CNP_FILE_PATTERN') + final_data_paths = args.data_paths if args.data_paths else env_data_paths + final_file_pattern = args.file_pattern if args.file_pattern else env_file_pattern + if final_data_paths or final_file_pattern: + update_kwargs = {} + if final_data_paths: + # Support comma-separated paths + update_kwargs['data_paths'] = [p.strip() for p in str(final_data_paths).split(',') if p.strip()] + if final_file_pattern: + update_kwargs['file_pattern'] = str(final_file_pattern).strip() + try: + config.update_data_config(**update_kwargs) + logger.info(f"Applied data overrides: {update_kwargs}") + except Exception as e: + logger.warning(f"Failed to apply data overrides: {e}") if args.variable_list is not None: logger.info(f"Using CNP configuration from variable list file: {args.variable_list}") else: logger.info(f"Using default CNP variable configuration{' with water' if include_water else ' without water'}") if args.model_config is not None: logger.info(f"Applied model architecture overrides from: {args.model_config}") - # Set train/validation split to 50/50 + # Set train/validation split config.update_data_config(train_split=0.8) - # Ensure GPU and all files - config.update_training_config(device='cuda') + # Prefer GPU when available, otherwise CPU + device_str = 'cuda' if torch.cuda.is_available() else 'cpu' + config.update_training_config(device=device_str) config.update_data_config(max_files=args.max_files) # Turn off GPU monitoring and debug logging config.update_training_config(log_gpu_memory=False, log_gpu_utilization=False) @@ -359,6 +395,21 @@ def main(): # Keep fixed seeds for fair comparisons logger.info("Fixed shuffling (seeded) enabled for fair comparison") + # If only validating configuration, exit early before heavy work (used in CI) + if args.config_only: + # Validate 2D output alignment invariant + assert config.data_config.y_list_columns_2d == ['Y_' + v for v in config.data_config.x_list_columns_2d], \ + f"2D columns not aligned!\nX: {config.data_config.x_list_columns_2d}\nY: {config.data_config.y_list_columns_2d}" + # Log a brief summary and exit + logger.info("Configuration-only mode: built training/model/data configs successfully.") + logger.info(f"Data paths: {config.data_config.data_paths}") + logger.info(f"File pattern: {config.data_config.file_pattern}") + logger.info(f"Epochs: {config.training_config.num_epochs}, Batch size: {config.training_config.batch_size}") + return { + 'status': 'ok', + 'config_only': True + } + # Optional strict determinism (opt-in via CLI) if args.strict_determinism: seed = getattr(config.training_config, 'random_seed', 42) From 8b10e84b094510bb0f0fa6e0f8935dd023f33273 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Mon, 13 Oct 2025 17:44:17 -0400 Subject: [PATCH 18/51] remove old CNP_IO lists --- CNP_IO_default.txt | 41 -------- CNP_IO_list1.txt | 183 -------------------------------- CNP_IO_list2.txt | 257 --------------------------------------------- CNP_IO_mini.txt | 17 --- 4 files changed, 498 deletions(-) delete mode 100644 CNP_IO_default.txt delete mode 100644 CNP_IO_list1.txt delete mode 100644 CNP_IO_list2.txt delete mode 100644 CNP_IO_mini.txt diff --git a/CNP_IO_default.txt b/CNP_IO_default.txt deleted file mode 100644 index e2f5def..0000000 --- a/CNP_IO_default.txt +++ /dev/null @@ -1,41 +0,0 @@ -TIME SERIES VARIABLES (Climate Forcing) - 6 variables: -• FLDS, PSRF, FSDS, QBOT, PRECTmms, TBOT - -SURFACE PROPERTIES - 27 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, NPP, AR, HR - -1D PFT VARIABLES (12 variables): -• deadcrootc, deadcrootn, deadcrootp, deadstemc, deadstemn, deadstemp -• frootc, frootc_storage, leafc, leafc_storage, totvegc, tlai - -2D VARIABLES (layered - 27 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 -• sminn_vr, smin_no3_vr, smin_nh4_vr -• soil1c_vr, soil1n_vr, soil1p_vr -• soil2c_vr, soil2n_vr, soil2p_vr -• soil3c_vr, soil3n_vr, soil3p_vr -• soil4c_vr, soil4n_vr, soil4p_vr diff --git a/CNP_IO_list1.txt b/CNP_IO_list1.txt deleted file mode 100644 index 89bedc6..0000000 --- a/CNP_IO_list1.txt +++ /dev/null @@ -1,183 +0,0 @@ -AI MODEL VARIABLE LISTS FOR CONSTRUCTION - FINALIZED TEMPLATE -============================================================ - -DATASET OVERVIEW ----------------- -Dataset: Dataset 3 (Trend_1_data_CNP) -Total variables: 229 -Input variables: 119 (with water), 113 (without water) -Output variables: 52 (with water), 46 (without water) - -INPUT VARIABLES ORGANIZATION (FINALIZED) -======================================== - -1. FIRST 3 GROUPS (Time Series, Surface Properties, PFT Parameters) -------------------------------------------------------------------- -These groups contain the core forcing and parameter variables: - -TIME SERIES VARIABLES (Climate Forcing) - 6 variables: -• FLDS (Longwave downwelling solar flux) -• PSRF (Surface pressure) -• FSDS (Shortwave downwelling solar flux) -• QBOT (Specific humidity at lowest model level) -• PRECTmms (Total precipitation rate) -• TBOT (Temperature at lowest model level) - -SURFACE PROPERTIES - 27 variables: -Geographic: -• Latitude, Longitude, AREA, landfrac - -Soil Phosphorus Forms: -• OCCLUDED_P, SECONDARY_P, LABILE_P, APATITE_P - -PFT Coverage: -• 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 -• LANDFRAC_PFT -• PCT_NATVEG, SNOWDP - -Soil Texture: -• 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 Characteristics: -• 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 - -2. REMAINING 5 GROUPS (Water, Temperature, Scalar, 1D PFT, 2D SOIL) ----------------------------------------------------------- -These groups contain the state variables and fluxes: - -WATER VARIABLES (6 variables, optional): -• H2OCAN, H2OSFC, H2OSNO, TH2OSFC, H2OSOI_LIQ, H2OSOI_ICE - -TEMPERATURE VARIABLES (0 variables): -[EXCLUDED: T_GRND_R, T_GRND_U, T_LAKE, T_SOISNO] - -SCALAR VARIABLES (1D - 4 variables): -• GPP, NPP, AR, HR - -1D PFT VARIABLES (14 variables): -• deadcrootc, deadcrootn, deadcrootp, deadstemc, deadstemn, deadstemp -• frootc, frootc_storage, leafc, leafc_storage, totcolp, totlitc, totvegc, tlai -[EXCLUDED: H2OSOI_10CM, TS_TOPO, T_GRND_R, T_GRND_U, T_LAKE, T_SOISNO, abm, peatf, taf, cwdp] - -2D VARIABLES (layered - 27 variables): -Coarse Woody Debris and Secondary Phosphorus: -• cwdc_vr, cwdn_vr, cwdp_vr -[EXCLUDED: secondp_vr) - -Litter Variables: -• litr1c_vr, litr2c_vr, litr3c_vr -• litr1n_vr, litr2n_vr, litr3n_vr -• litr1p_vr, litr2p_vr, litr3p_vr - -Soil Mineral Nitrogen: -• sminn_vr, smin_no3_vr, smin_nh4_vr - -Soil Variables: -• soil1c_vr, soil1n_vr, soil1p_vr, soil2c_vr, soil2n_vr, soil2p_vr, soil3c_vr, soil3n_vr, soil3p_vr, soil4c_vr, soil4n_vr, soil4p_vr - -OUTPUT VARIABLES ORGANIZATION (FINALIZED) -======================================== - -OUTPUT VARIABLES (46 with water, 40 without water, all starting with Y_): - -WATER VARIABLES (6 variables, optional): -• Y_H2OCAN, Y_H2OSFC, Y_H2OSNO, Y_TH2OSFC, Y_H2OSOI_LIQ, Y_H2OSOI_ICE -[EXCLUDED FOR FIRST EXPERIMENTS] - -TEMPERATURE VARIABLES (7 variables): -• Y_T_GRND_R, Y_T_GRND_U, Y_T_LAKE, Y_T_SOISNO, Y_T_GRND_1_, Y_T_GRND_2_, Y_T_GRND_3_ -[EXCLUDED FOR FIRST EXPERIMENTS] - -SCALAR VARIABLES (4 variables): -• Y_GPP, Y_NPP, Y_AR, Y_HR - -1D PFT VARIABLES (14 variables): -• Y_deadcrootc, Y_deadcrootn, Y_deadcrootp -• Y_deadstemc, Y_deadstemn, Y_deadstemp, Y_frootc, Y_frootc_storage -• Y_leafc, Y_leafc_storage, Y_totcolp, Y_totlitc, Y_totvegc, Y_tlai -[NOTE: These PFT variables are 1D, not layered/2D.] - -2D VARIABLES (layered - 27 variables): - -Coarse Woody Debris and Secondary Phosphorus: -• Y_cwdc_vr, Y_cwdn_vr, Y_cwdp_vr -[EXCLUDED Y_secondp_vr] - -Litter Variables: -• Y_litr1c_vr, Y_litr2c_vr, Y_litr3c_vr -• Y_litr1n_vr, Y_litr2n_vr, Y_litr3n_vr -• Y_litr1p_vr, Y_litr2p_vr, Y_litr3p_vr - -Soil Mineral Nitrogen: -• Y_sminn_vr, Y_smin_no3_vr, Y_smin_nh4_vr - -Soil Variables: -• Y_soil3c_vr, Y_soil3n_vr, Y_soil3p_vr -• Y_soil1c_vr, Y_soil1n_vr, Y_soil1p_vr -• Y_soil2c_vr, Y_soil2n_vr, Y_soil2p_vr, -• Y_soil4c_vr, Y_soil4n_vr, Y_soil4p_vr - -EXCLUDED VARIABLES SUMMARY -========================== -INPUT EXCLUSIONS (7 variables): -• H2OSOI_10CM (water variable) -• T_GRND_R, T_GRND_U, T_LAKE, T_SOISNO (temperature variables) -• abm, peatf, taf, cwdp (other variables) -• secondp_vr - -OUTPUT EXCLUSIONS (1 variable): -• Y_LAKE_SOILC (lake soil carbon) -• Y_secondp_vr - -AI MODEL CONSTRUCTION TEMPLATE -============================== - -1. INPUT VARIABLE SELECTION (113 (119 with water)): - - Time Series: 6 variables (climate forcing) - - Surface Properties: 27 variables (geographic, soil phosphorus forms, PFT coverage, soil texture) - - PFT Parameters: 44 variables (plant functional type characteristics) - - Water: 6 variables (optional, H2OCAN, H2OSFC, H2OSNO, TH2OSFC, H2OSOI_LIQ, H2OSOI_ICE) - - Temperature: 0 variables (all excluded) - - Scalar: 5 variables (GPP, NPP, AR, HR) - - 2D PFT Variables: 14 variables (CNP pools, TLAI) - - 2D Variables: 27 variables (soil and litter properties, including soil variables) - -2. OUTPUT VARIABLE SELECTION (45 without water): - - Scalar: 4 variables (Y_GPP, Y_NPP, Y_AR, Y_HR) - - 2D PFT Variables: 14 variables (Y_CNP pools, Y_TLAI) - - 2D Variables: 27 variables (Y_soil and Y_litter properties) - - Water: 6 variables (optional, Y_H2OCAN, Y_H2OSFC, Y_H2OSNO, Y_TH2OSFC, Y_H2OSOI_LIQ, Y_H2OSOI_ICE) - - Temperature: 7 variables (Y_T variables) - EXCLUDED FOR FIRST EXPERIMENTS - -3. NEURAL NETWORK ARCHITECTURE CONSIDERATIONS: - - Time Series Group: 6 variables (LSTM/Transformer) - - Static Surface Group: 27 variables (Dense layers) - - PFT Parameters Group: 44 variables (Dense layers) - - Water Group: 6 variables (1D CNN or Dense, optional) - - Scalar Group: 4 variables (Dense layers) - - 1D PFT CNP Pools Group: 14 variables (Dense layers) - - 2D CNP Pools Group: 27 variables (2D CNN or specialized layers) - -4. DATA PREPROCESSING REQUIREMENTS: - - Normalize/standardize all variables appropriately - - Handle missing values - - Consider temporal dependencies for time series - - Separate processing for different variable dimensions - - Ensure proper alignment between input and output variables - -5. MODEL TRAINING CONSIDERATIONS: - - Use perfect correspondence variables for validation - - Consider multi-task learning for different variable types - - Implement appropriate loss functions for each variable group - - Monitor performance on excluded variables separately - -FINAL COUNTS: -============= -Input variables (first 3 groups): 77 (6 time series + 27 surface properties + 44 PFT parameters) -Input variables (remaining 5 groups): 45 without water (4 scalar + 14 1D PFT + 27 2D) -Total input variables: 122 without water -Total output variables: 45 without water -Total variables: 167 without water \ No newline at end of file diff --git a/CNP_IO_list2.txt b/CNP_IO_list2.txt deleted file mode 100644 index 6ddbb34..0000000 --- a/CNP_IO_list2.txt +++ /dev/null @@ -1,257 +0,0 @@ -VARIABLES IN NEW DATASET (from 1_training_data_batch_01.pkl) -================================================== -landfrac -Latitude -Longitude -FLDS -PSRF -FSDS -QBOT -PRECTmms -TBOT -LANDFRAC_PFT -PCT_NATVEG -AREA -peatf -abm -SOIL_COLOR -SOIL_ORDER -soil3c_vr -soil4c_vr -cwdc_vr -deadcrootc -deadstemc -tlai -GPP -SNOWDP -H2OSOI_10CM -HR -AR -NPP -OCCLUDED_P -SECONDARY_P -LABILE_P -APATITE_P -cwdn_vr -secondp_vr -cwdp_vr -cwdp -totcolp -totvegc -deadstemn -deadcrootn -deadstemp -deadcrootp -leafc -leafc_storage -frootc -frootc_storage -totlitc -leafn -leafn_storage -frootn -frootn_storage -leafp -leafp_storage -frootp -frootp_storage -livestemc -livestemc_storage -livestemn -livestemn_storage -livestemp -livestemp_storage -labilep_vr -occlp_vr -primp_vr -deadcrootc_storage -deadstemc_storage -livecrootc -livecrootc_storage -deadcrootn_storage -deadstemn_storage -livecrootn -livecrootn_storage -deadcrootp_storage -deadstemp_storage -livecrootp -livecrootp_storage -soil1c_vr -soil1n_vr -soil1p_vr -soil2c_vr -soil2n_vr -soil2p_vr -soil3n_vr -soil3p_vr -soil4n_vr -soil4p_vr -litr1c_vr -litr2c_vr -litr3c_vr -litr1n_vr -litr2n_vr -litr3n_vr -litr1p_vr -litr2p_vr -litr3p_vr -sminn_vr -smin_no3_vr -smin_nh4_vr -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_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 -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 -Y_soil3c_vr -Y_soil4c_vr -Y_cwdc_vr -Y_deadcrootc -Y_deadstemc -Y_tlai -Y_GPP -Y_HR -Y_AR -Y_NPP -Y_cwdn_vr -Y_secondp_vr -Y_cwdp_vr -Y_cwdp -Y_totcolp -Y_totvegc -Y_deadstemn -Y_deadcrootn -Y_deadstemp -Y_deadcrootp -Y_leafc -Y_leafc_storage -Y_frootc -Y_frootc_storage -Y_leafn -Y_leafn_storage -Y_frootn -Y_frootn_storage -Y_leafp -Y_leafp_storage -Y_frootp -Y_frootp_storage -Y_livestemc -Y_livestemc_storage -Y_livestemn -Y_livestemn_storage -Y_livestemp -Y_livestemp_storage -Y_labilep_vr -Y_occlp_vr -Y_primp_vr -Y_deadcrootc_storage -Y_deadstemc_storage -Y_livecrootc -Y_livecrootc_storage -Y_deadcrootn_storage -Y_deadstemn_storage -Y_livecrootn -Y_livecrootn_storage -Y_deadcrootp_storage -Y_deadstemp_storage -Y_livecrootp -Y_livecrootp_storage -Y_totlitc -Y_soil1c_vr -Y_soil1n_vr -Y_soil1p_vr -Y_soil2c_vr -Y_soil2n_vr -Y_soil2p_vr -Y_soil3n_vr -Y_soil3p_vr -Y_soil4n_vr -Y_soil4p_vr -Y_litr1c_vr -Y_litr2c_vr -Y_litr3c_vr -Y_litr1n_vr -Y_litr2n_vr -Y_litr3n_vr -Y_litr1p_vr -Y_litr2p_vr -Y_litr3p_vr -Y_sminn_vr -Y_smin_no3_vr -Y_smin_nh4_vr -pft_c3psn -pft_croot_stem -pft_crop -pft_deadwdcn -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_frootcn -pft_grperc -pft_grpnow -pft_leaf_long -pft_leafcn -pft_lf_fcel -pft_lf_flab -pft_lf_flig -pft_lflitcn -pft_livewdcn -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 diff --git a/CNP_IO_mini.txt b/CNP_IO_mini.txt deleted file mode 100644 index 4c4b724..0000000 --- a/CNP_IO_mini.txt +++ /dev/null @@ -1,17 +0,0 @@ -TIME SERIES VARIABLES (Climate Forcing) - 2 variables: -• FLDS, PSRF - -SURFACE PROPERTIES - 5 variables: -• Latitude, Longitude, AREA, landfrac, PCT_NAT_PFT_0, PCT_NAT_PFT_1 - -PFT PARAMETERS (Plant Functional Type Characteristics) - 5 variables: -• pft_deadwdcn, pft_frootcn, pft_leafcn, pft_lflitcn, pft_livewdcn - -SCALAR VARIABLES (1D - 4 variables): -• GPP, NPP, AR, HR - -1D PFT VARIABLES (2 variables): -• deadcrootc, deadcrootn - -2D VARIABLES (layered - 3 variables): -• cwdc_vr, cwdn_vr, cwdp_vr \ No newline at end of file From cae5024aa15a8a83807d57ef9d14a47db39d04c8 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Mon, 13 Oct 2025 18:09:23 -0400 Subject: [PATCH 19/51] read model config for inference --- scripts/run_inference_all.py | 54 +++++++++++++++++++++++++++++++++--- 1 file changed, 50 insertions(+), 4 deletions(-) diff --git a/scripts/run_inference_all.py b/scripts/run_inference_all.py index 30ff5eb..574cc9a 100644 --- a/scripts/run_inference_all.py +++ b/scripts/run_inference_all.py @@ -111,6 +111,31 @@ def _extract_variables_from_config(config_path: Path) -> dict: return None +# New: load model_config from cnp_config.json for exact architecture reuse +def _load_training_model_config_from_config(model_path: Path) -> dict: + """Load model_config dict from cnp_config.json in the training run directory.""" + import json + try: + model_dir = Path(model_path).parent + # Search current and parent directories for cnp_config.json + candidate_paths = [model_dir / 'cnp_config.json'] + [p / 'cnp_config.json' for p in model_dir.parents] + for config_path in candidate_paths: + if config_path.exists(): + try: + with open(config_path, 'r') as f: + cfg = json.load(f) + if isinstance(cfg, dict) and 'model_config' in cfg and isinstance(cfg['model_config'], dict): + logging.info(f"Loaded model_config from {config_path}") + return cfg['model_config'] + except Exception as e: + logging.warning(f"Failed reading model_config from {config_path}: {e}") + break + logging.warning("No model_config found in cnp_config.json near model path") + except Exception as e: + logging.warning(f"Error discovering model_config: {e}") + return None + + @@ -190,6 +215,24 @@ def run_inference_all( if model_config is not None and use_training_config: logging.warning("--model-config provided along with --use-training-config; training config will still govern variables and scalers. Model overrides only affect architecture sizing.") + # Apply model_config from training cnp_config.json if requested + json_model_config = None + if use_training_config: + json_model_config = _load_training_model_config_from_config(Path(model_path)) + if isinstance(json_model_config, dict): + applied = 0 + for k, v in json_model_config.items(): + if hasattr(config.model_config, k): + try: + setattr(config.model_config, k, v) + applied += 1 + except Exception as e: + logging.warning(f"Failed applying model_config.{k} from JSON: {e}") + else: + # Some fields may be added dynamically later; log and skip + logging.info(f"Ignoring unknown ModelConfig key in JSON: {k}") + logging.info(f"Applied {applied} model_config fields from training JSON") + # CRITICAL FIX: Apply the loaded variable configuration to ensure model compatibility if variables is not None: logging.info("Applying variable configuration to model config...") @@ -215,10 +258,13 @@ def run_inference_all( logging.info(f" PFT 1D variables: {len(config.data_config.x_list_columns_1d)} variables") logging.info(f" 2D soil variables: {len(config.data_config.x_list_columns_2d)} variables") - # Update model configuration for output dimensions - config.model_config.scalar_output_size = len(config.data_config.x_list_scalar_columns) - config.model_config.vector_output_size = len(config.data_config.x_list_columns_1d) - config.model_config.matrix_output_size = len(config.data_config.x_list_columns_2d) + # Update model configuration for output dimensions only if not provided by JSON + if not (isinstance(json_model_config, dict) and 'scalar_output_size' in json_model_config): + config.model_config.scalar_output_size = len(config.data_config.x_list_scalar_columns) + if not (isinstance(json_model_config, dict) and 'vector_output_size' in json_model_config): + config.model_config.vector_output_size = len(config.data_config.x_list_columns_1d) + if not (isinstance(json_model_config, dict) and 'matrix_output_size' in json_model_config): + config.model_config.matrix_output_size = len(config.data_config.x_list_columns_2d) logging.info(f"Updated model output dimensions:") logging.info(f" Scalar output size: {config.model_config.scalar_output_size}") From 9ca2b075e2a15d43379dd1f94bd0a49f614c1507 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Mon, 13 Oct 2025 18:29:57 -0400 Subject: [PATCH 20/51] remove unused test scripts --- CNP_IO_test1.txt | 32 -- scripts/test_individual_denorm.py | 145 ------- ...up_normalization_backward_compatibility.py | 407 ------------------ test_individual_scaler.py | 273 ------------ test_normalization_approaches.py | 370 ---------------- 5 files changed, 1227 deletions(-) delete mode 100644 CNP_IO_test1.txt delete mode 100644 scripts/test_individual_denorm.py delete mode 100644 test_group_normalization_backward_compatibility.py delete mode 100644 test_individual_scaler.py delete mode 100644 test_normalization_approaches.py diff --git a/CNP_IO_test1.txt b/CNP_IO_test1.txt deleted file mode 100644 index 50d3bc8..0000000 --- a/CNP_IO_test1.txt +++ /dev/null @@ -1,32 +0,0 @@ -TIME SERIES VARIABLES (Climate Forcing) - 6 variables: -• FLDS, PSRF, FSDS, QBOT, PRECTmms, TBOT - -SURFACE PROPERTIES - 27 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, NPP, AR, HR - -1D PFT VARIABLES (1 variables): -• tlai - -2D VARIABLES (layered - 3 variables): -• cwdc_vr, cwdn_vr, cwdp_vr diff --git a/scripts/test_individual_denorm.py b/scripts/test_individual_denorm.py deleted file mode 100644 index 33e2858..0000000 --- a/scripts/test_individual_denorm.py +++ /dev/null @@ -1,145 +0,0 @@ -#!/usr/bin/env python3 -""" -Test if the individual scalers from the new run can actually denormalize data. -""" - -import pickle -import pandas as pd -import numpy as np -from pathlib import Path -import sys -sys.path.append('..') - -run_dir = Path('../cnp_results/run_20250818_125330') -scalers_dir = run_dir / 'cnp_predictions' / 'scalers' - -def test_individual_denorm(): - """Test if individual scalers can denormalize data.""" - - # Load individual scalers - try: - with open(scalers_dir / 'individual_y_pft_1d_scaler.pkl', 'rb') as f: - pft_scaler = pickle.load(f) - print("✅ Loaded individual_y_pft_1d_scaler.pkl") - except Exception as e: - print(f"❌ Failed to load PFT scaler: {e}") - return - - try: - with open(scalers_dir / 'individual_y_soil_2d_scaler.pkl', 'rb') as f: - soil_scaler = pickle.load(f) - print("✅ Loaded individual_y_soil_2d_scaler.pkl") - except Exception as e: - print(f"❌ Failed to load soil scaler: {e}") - return - - # Load some predictions to test - try: - pft_pred = pd.read_csv(run_dir / 'cnp_predictions' / 'pft_1d_predictions' / 'predictions_Y_tlai.csv') - print(f"✅ Loaded PFT predictions: shape={pft_pred.shape}") - - # Take a few non-zero rows - non_zero_rows = pft_pred[(pft_pred != 0).any(axis=1)] - if len(non_zero_rows) > 0: - test_data = non_zero_rows.iloc[0:2].values - print(f"✅ Found {len(non_zero_rows)} non-zero rows, testing with first 2") - print(f" Test data shape: {test_data.shape}") - print(f" Sample values: {test_data[0, :5]}") # First 5 values of first row - else: - print("❌ No non-zero rows found in PFT predictions") - return - except Exception as e: - print(f"❌ Failed to load PFT predictions: {e}") - return - - # Test PFT 1D denormalization - print("\n--- Testing PFT 1D Denormalization ---") - try: - # The data should be in shape (samples, features) where features are the PFTs - # We need to reshape to (samples, pfts, variables) for the individual scaler - test_data_reshaped = test_data.reshape(test_data.shape[0], test_data.shape[1], 1) - print(f" Reshaped data shape: {test_data_reshaped.shape}") - - # Try to denormalize using the individual scaler - # The method signature is: inverse_transform_pft_1d(data, pft_names, variable_names) - pft_names = [f'PFT{i}' for i in range(16)] - variable_names = ['Y_tlai'] - denorm_data = pft_scaler.inverse_transform_pft_1d(test_data_reshaped, pft_names, variable_names) - print(f"✅ Denormalization successful!") - print(f" Original shape: {test_data.shape}") - print(f" Denormalized shape: {denorm_data.shape}") - print(f" Original values: {test_data[0, :5]}") - print(f" Denormalized values: {denorm_data[0, :5]}") - - # Reshape denorm_data back to original shape for comparison - denorm_data_flat = denorm_data.reshape(denorm_data.shape[0], -1) - print(f" Denormalized data flattened shape: {denorm_data_flat.shape}") - - # Check if values changed significantly - if np.allclose(test_data, denorm_data_flat, atol=1e-6): - print("⚠️ WARNING: Denormalized values are very close to original - possible identity transformation") - else: - print("✅ Denormalization appears to be working - values changed significantly") - - except Exception as e: - print(f"❌ PFT 1D denormalization failed: {e}") - print(f" Error type: {type(e).__name__}") - - # Test Soil 2D denormalization - print("\n--- Testing Soil 2D Denormalization ---") - try: - soil_pred = pd.read_csv(run_dir / 'cnp_predictions' / 'soil_2d_predictions' / 'predictions_Y_cwdc_vr.csv') - print(f"✅ Loaded soil predictions: shape={soil_pred.shape}") - - # Take a few non-zero rows - non_zero_rows = soil_pred[(soil_pred != 0).any(axis=1)] - if len(non_zero_rows) > 0: - test_data = non_zero_rows.iloc[0:2].values - print(f"✅ Found {len(non_zero_rows)} non-zero rows, testing with first 2") - print(f" Original shape: {test_data.shape}") - print(f" Sample values: {test_data[0, :5]}") # First 5 values of first row - else: - print("❌ No non-zero rows found in soil predictions") - return - except Exception as e: - print(f"❌ Failed to load soil predictions: {e}") - return - - try: - # The data should be in shape (samples, features) where features are (columns, layers) - # We need to reshape to (samples, columns, layers) for the individual scaler - # The data has 180 features, which should be 18 columns × 10 layers - test_data_reshaped = test_data.reshape(test_data.shape[0], 18, 10) # 18 columns, 10 layers - print(f" Reshaped data shape: {test_data_reshaped.shape}") - - # Try to denormalize using the individual scaler - # The method signature is: inverse_transform_soil_2d(data, variable_names, num_layers) - # But the data needs to be in shape (samples, variables, columns, layers) - # So we need to add a variable dimension - test_data_final = test_data_reshaped.reshape(test_data.shape[0], 1, 18, 10) # (samples, variables, columns, layers) - print(f" Final reshaped data shape: {test_data_final.shape}") - - variable_names = ['Y_cwdc_vr'] - denorm_data = soil_scaler.inverse_transform_soil_2d(test_data_final, variable_names, 10) - print(f"✅ Denormalization successful!") - print(f" Original shape: {test_data.shape}") - print(f" Denormalized shape: {denorm_data.shape}") - print(f" Original values: {test_data[0, :5]}") - print(f" Denormalized values: {denorm_data[0, :5]}") - - # Reshape denorm_data back to original shape for comparison - denorm_data_flat = denorm_data.reshape(denorm_data.shape[0], -1) - print(f" Denormalized data flattened shape: {denorm_data_flat.shape}") - - # Check if values changed significantly - if np.allclose(test_data, denorm_data_flat, atol=1e-6): - print("⚠️ WARNING: Denormalized values are very close to original - possible identity transformation") - else: - print("✅ Denormalization appears to be working - values changed significantly") - - except Exception as e: - print(f"❌ Soil 2D denormalization failed: {e}") - print(f" Error type: {type(e).__name__}") - -if __name__ == "__main__": - test_individual_denorm() diff --git a/test_group_normalization_backward_compatibility.py b/test_group_normalization_backward_compatibility.py deleted file mode 100644 index 7b31055..0000000 --- a/test_group_normalization_backward_compatibility.py +++ /dev/null @@ -1,407 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script for Group Normalization Backward Compatibility - -This script verifies that the new DataLoaderIndividual with group normalization -produces exactly the same results as the original system. -""" - -import sys -import os -import numpy as np -import pandas as pd -import torch -from pathlib import Path - -# Add the project root to the path -sys.path.insert(0, os.path.dirname(__file__)) - -try: - from config.training_config import DataConfig, PreprocessingConfig - from data.data_loader_individual import DataLoaderIndividual - print("✅ Successfully imported required modules") -except ImportError as e: - print(f"❌ Import error: {e}") - print("Please ensure you're running this from the project root directory") - sys.exit(1) - - -def create_test_data(): - """Create deterministic test data for reproducible testing.""" - print("🧪 Creating deterministic test data...") - - # Set random seed for reproducibility - np.random.seed(42) - - # Create sample data - n_samples = 20 # Small dataset for quick testing - - # Mock DataFrame with required columns - mock_data = { - # Time series data - 'FLDS': [np.random.uniform(100, 500, 240) for _ in range(n_samples)], - 'PSRF': [np.random.uniform(80000, 120000, 240) for _ in range(n_samples)], - 'FSDS': [np.random.uniform(-50, 50, 240) for _ in range(n_samples)], - 'QBOT': [np.random.uniform(0, 0.01, 240) for _ in range(n_samples)], - 'PRECTmms': [np.random.uniform(0, 0.01, 240) for _ in range(n_samples)], - 'TBOT': [np.random.uniform(200, 350, 240) for _ in range(n_samples)], - - # Static data - 'Latitude': np.random.uniform(30, 60, n_samples), - 'Longitude': np.random.uniform(-120, -60, n_samples), - 'AREA': np.random.uniform(1e10, 1e12, n_samples), - - # Scalar data with specific ranges to test normalization - 'GPP': np.random.uniform(0, 20, n_samples), # 0-20 range - 'NPP': np.random.uniform(0, 15, n_samples), # 0-15 range - 'AR': np.random.uniform(0, 10, n_samples), # 0-10 range - 'HR': np.random.uniform(0, 12, n_samples), # 0-12 range - - # Y scalar data - 'Y_GPP': np.random.uniform(0, 20, n_samples), - 'Y_NPP': np.random.uniform(0, 15, n_samples), - 'Y_AR': np.random.uniform(0, 10, n_samples), - 'Y_HR': np.random.uniform(0, 12, n_samples), - - # PFT1D data - 'tlai': [np.random.uniform(0, 10, 17) for _ in range(n_samples)], - 'deadstemc': [np.random.uniform(0, 1000, 17) for _ in range(n_samples)], - - # Y PFT1D data - 'Y_tlai': [np.random.uniform(0, 10, 17) for _ in range(n_samples)], - 'Y_deadstemc': [np.random.uniform(0, 1000, 17) for _ in range(n_samples)], - - # Soil2D data - 'cwdc_vr': [np.random.uniform(0, 1000, (360, 10)) for _ in range(n_samples)], - 'sminn_vr': [np.random.uniform(0, 50, (360, 10)) for _ in range(n_samples)], - - # Y Soil2D data - 'Y_cwdc_vr': [np.random.uniform(0, 1000, (360, 10)) for _ in range(n_samples)], - 'Y_sminn_vr': [np.random.uniform(0, 50, (360, 10)) for _ in range(n_samples)], - - # PFT param data - 'pft_deadwdcn': [np.random.uniform(0, 100, 17) for _ in range(n_samples)], - 'pft_frootcn': [np.random.uniform(0, 100, 17) for _ in range(n_samples)], - } - - df = pd.DataFrame(mock_data) - print(f"✅ Created test DataFrame with {len(df)} samples and {len(df.columns)} columns") - - # Print some sample values for verification - print("\n📊 Sample data values (first 3 samples):") - print(f" GPP: {df['GPP'].iloc[:3].values}") - print(f" NPP: {df['NPP'].iloc[:3].values}") - print(f" AR: {df['AR'].iloc[:3].values}") - print(f" HR: {df['HR'].iloc[:3].values}") - - return df - - -def test_group_normalization_basic(data_loader): - """Test basic group normalization functionality.""" - print("\n🧪 Testing basic group normalization...") - print("=" * 50) - - try: - # Use group normalization (default) - normalized_data = data_loader.normalize_data() - - print("✅ Group normalization completed successfully") - - # Check data shapes - print(f"✅ Scalar data shape: {normalized_data['scalar_data'].shape}") - print(f"✅ Y scalar data shape: {normalized_data['y_scalar'].shape}") - print(f"✅ PFT1D data shape: {normalized_data['variables_1d_pft'].shape}") - print(f"✅ Soil2D data shape: {normalized_data['variables_2d_soil'].shape}") - print(f"✅ Time series data shape: {normalized_data['time_series_data'].shape}") - print(f"✅ Static data shape: {normalized_data['static_data'].shape}") - - # Check data ranges (should be 0-1 for MinMaxScaler) - print("\n📊 Normalized data ranges (should be 0-1):") - print(f" Scalar data: [{normalized_data['scalar_data'].min():.6f}, {normalized_data['scalar_data'].max():.6f}]") - print(f" Y scalar data: [{normalized_data['y_scalar'].min():.6f}, {normalized_data['y_scalar'].max():.6f}]") - - # Check scaler types - scalers = normalized_data['scalers'] - print("\n📊 Scaler types used:") - for name, scaler in scalers.items(): - if scaler is not None: - if hasattr(scaler, 'normalization_type'): - print(f" {name}: IndividualScalerManager ({scaler.normalization_type})") - else: - print(f" {name}: {type(scaler).__name__}") - - return normalized_data - - except Exception as e: - print(f"❌ Group normalization failed: {e}") - import traceback - traceback.print_exc() - return None - - -def test_data_consistency(data_loader, normalized_data): - """Test that data is consistent and properly normalized.""" - print("\n🧪 Testing data consistency...") - print("=" * 50) - - try: - # Check that normalized data is in expected range (0-1 for MinMaxScaler) - print("📊 Checking normalized data ranges:") - - # Scalar data - scalar_data = normalized_data['scalar_data'] - scalar_min = scalar_data.min() - scalar_max = scalar_data.max() - print(f" Scalar data range: [{scalar_min:.6f}, {scalar_max:.6f}]") - - if 0.0 <= scalar_min <= 0.1 and 0.9 <= scalar_max <= 1.0: - print(" ✅ Scalar data properly normalized to [0,1] range") - else: - print(" ⚠️ Scalar data may not be properly normalized") - - # Y scalar data - y_scalar_data = normalized_data['y_scalar'] - y_scalar_min = y_scalar_data.min() - y_scalar_max = y_scalar_data.max() - print(f" Y scalar data range: [{y_scalar_min:.6f}, {y_scalar_max:.6f}]") - - if 0.0 <= y_scalar_min <= 0.1 and 0.9 <= y_scalar_max <= 1.0: - print(" ✅ Y scalar data properly normalized to [0,1] range") - else: - print(" ⚠️ Y scalar data may not be properly normalized") - - # Check that no NaN or infinite values - print("\n📊 Checking for data quality issues:") - has_nan = torch.isnan(scalar_data).any() or torch.isnan(y_scalar_data).any() - has_inf = torch.isinf(scalar_data).any() or torch.isinf(y_scalar_data).any() - - if not has_nan: - print(" ✅ No NaN values detected") - else: - print(" ❌ NaN values detected!") - - if not has_inf: - print(" ✅ No infinite values detected") - else: - print(" ❌ Infinite values detected!") - - # Check data types - print("\n📊 Checking data types:") - print(f" Scalar data type: {scalar_data.dtype}") - print(f" Y scalar data type: {y_scalar_data.dtype}") - - expected_dtype = torch.float32 - if scalar_data.dtype == expected_dtype and y_scalar_data.dtype == expected_dtype: - print(f" ✅ Data types are correct ({expected_dtype})") - else: - print(f" ⚠️ Data types may be incorrect (expected {expected_dtype})") - - return True - - except Exception as e: - print(f"❌ Data consistency test failed: {e}") - import traceback - traceback.print_exc() - return False - - -def test_scaler_functionality(data_loader, normalized_data): - """Test that scalers work correctly and can be used for inverse transformation.""" - print("\n🧪 Testing scaler functionality...") - print("=" * 50) - - try: - scalers = normalized_data['scalers'] - - # Test scalar scaler - if 'scalar' in scalers and scalers['scalar'] is not None: - scalar_scaler = scalers['scalar'] - print(f"✅ Scalar scaler found: {type(scalar_scaler).__name__}") - - # Test inverse transformation - if hasattr(scalar_scaler, 'inverse_transform'): - # Get some normalized data - sample_normalized = normalized_data['scalar_data'][:5] # First 5 samples - - # Convert to numpy for sklearn scalers - if isinstance(scalar_scaler, torch.Tensor): - print(" ⚠️ Scaler is a tensor, cannot test inverse transform") - else: - try: - # Inverse transform - sample_denormalized = scalar_scaler.inverse_transform(sample_normalized.numpy()) - print(f" ✅ Inverse transformation successful") - print(f" ✅ Denormalized shape: {sample_denormalized.shape}") - - # Check that denormalized values are reasonable - if sample_denormalized.min() >= 0 and sample_denormalized.max() <= 25: - print(" ✅ Denormalized values are in reasonable range") - else: - print(" ⚠️ Denormalized values may be out of expected range") - - except Exception as e: - print(f" ❌ Inverse transformation failed: {e}") - else: - print(" ⚠️ Scaler does not have inverse_transform method") - else: - print("⚠️ Scalar scaler not found or is None") - - # Test y_scalar scaler - if 'y_scalar' in scalers and scalers['y_scalar'] is not None: - y_scalar_scaler = scalers['y_scalar'] - print(f"✅ Y scalar scaler found: {type(y_scalar_scaler).__name__}") - else: - print("⚠️ Y scalar scaler not found or is None") - - return True - - except Exception as e: - print(f"❌ Scaler functionality test failed: {e}") - import traceback - traceback.print_exc() - return False - - -def test_scaler_persistence(data_loader): - """Test that scalers can be saved and loaded.""" - print("\n🧪 Testing scaler persistence...") - print("=" * 50) - - try: - # Create temporary directory - import tempfile - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = Path(temp_dir) - - # Save scalers - data_loader.save_scalers(temp_path) - print(f"✅ Scalers saved to {temp_path}") - - # Check if files were created - scaler_dir = temp_path / "scalers" - if scaler_dir.exists(): - print(f"✅ Scaler directory created: {scaler_dir}") - - # List saved files - saved_files = list(scaler_dir.rglob("*")) - print(f"✅ Saved {len(saved_files)} scaler files") - - for file_path in saved_files: - if file_path.is_file(): - print(f" - {file_path.name}") - else: - print("⚠️ Scaler directory not created") - - # Test loading scalers (create new instance) - new_data_loader = DataLoaderIndividual( - data_loader.data_config, - data_loader.preprocessing_config - ) - - new_data_loader.load_scalers(temp_path) - print("✅ Scalers loaded successfully") - - return True - - except Exception as e: - print(f"❌ Failed to test scaler persistence: {e}") - import traceback - traceback.print_exc() - return False - - -def main(): - """Run group normalization backward compatibility tests.""" - print("🧪 Testing Group Normalization Backward Compatibility") - print("=" * 70) - - # Create data loader - try: - data_config = DataConfig( - data_paths=["/tmp"], # Dummy path - file_pattern="*.pkl", - time_series_columns=["FLDS", "PSRF", "FSDS", "QBOT", "PRECTmms", "TBOT"], - time_series_length=240, - static_columns=["Latitude", "Longitude", "AREA"], - x_list_scalar_columns=["GPP", "NPP", "AR", "HR"], - y_list_scalar_columns=["Y_GPP", "Y_NPP", "Y_AR", "Y_HR"], - x_list_columns_1d=["tlai", "deadstemc"], - y_list_columns_1d=["Y_tlai", "Y_deadstemc"], - x_list_columns_2d=["cwdc_vr", "sminn_vr"], - y_list_columns_2d=["Y_cwdc_vr", "Y_sminn_vr"], - pft_param_columns=["pft_deadwdcn", "pft_frootcn"], - max_1d_length=17, - max_2d_rows=360, - max_2d_cols=10, - random_state=42 - ) - - preprocessing_config = PreprocessingConfig( - time_series_normalization="minmax", - static_normalization="minmax", - list_1d_normalization="minmax", - list_2d_normalization="minmax", - data_type=torch.float32 - ) - - data_loader = DataLoaderIndividual(data_config, preprocessing_config) - print("✅ DataLoaderIndividual created successfully") - - except Exception as e: - print(f"❌ Failed to create DataLoaderIndividual: {e}") - import traceback - traceback.print_exc() - return False - - # Integrate test data - test_df = create_test_data() - data_loader.df = test_df - - # Preprocess data - try: - data_loader.preprocess_data() - print("✅ Data preprocessing completed") - except Exception as e: - print(f"❌ Data preprocessing failed: {e}") - import traceback - traceback.print_exc() - return False - - # Test 1: Basic group normalization - normalized_data = test_group_normalization_basic(data_loader) - if normalized_data is None: - print("❌ Cannot proceed without normalized data") - return False - - # Test 2: Data consistency - if not test_data_consistency(data_loader, normalized_data): - print("⚠️ Data consistency test failed") - - # Test 3: Scaler functionality - if not test_scaler_functionality(data_loader, normalized_data): - print("⚠️ Scaler functionality test failed") - - # Test 4: Scaler persistence - if not test_scaler_persistence(data_loader): - print("⚠️ Scaler persistence test failed") - - print("\n" + "=" * 70) - print("🎉 Group normalization backward compatibility tests completed!") - print("\n📋 Summary:") - print("✅ Group normalization is working correctly") - print("✅ Data shapes and types are consistent") - print("✅ Normalization produces expected [0,1] ranges") - print("✅ Scalers can be saved and loaded") - print("\n💡 Next steps:") - print("1. Test with your actual data to verify compatibility") - print("2. Compare results with your existing system") - print("3. Once satisfied, test individual normalization") - print("4. Use hybrid approach for selective optimization") - - return True - - -if __name__ == "__main__": - success = main() - sys.exit(0 if success else 1) diff --git a/test_individual_scaler.py b/test_individual_scaler.py deleted file mode 100644 index 7749540..0000000 --- a/test_individual_scaler.py +++ /dev/null @@ -1,273 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script for IndividualScalerManager - -This script tests the basic functionality of the IndividualScalerManager -to ensure it works correctly before integration. -""" - -import numpy as np -import tempfile -import shutil -from pathlib import Path -import sys -import os - -# Add the data directory to the path -sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'data')) - -from individual_scaler_manager import IndividualScalerManager - - -def test_scalar_normalization(): - """Test individual scalar variable normalization.""" - print("Testing scalar normalization...") - - # Create sample data with different ranges - np.random.seed(42) - n_samples = 1000 - - # GPP: 0-20 gC/m²/day - gpp = np.random.uniform(0, 20, n_samples) - # NPP: 0-15 gC/m²/day - npp = np.random.uniform(0, 15, n_samples) - # AR: 0-10 gC/m²/day - ar = np.random.uniform(0, 10, n_samples) - # HR: 0-12 gC/m²/day - hr = np.random.uniform(0, 12, n_samples) - - # Combine into array - data = np.column_stack([gpp, npp, ar, hr]) - variable_names = ['GPP', 'NPP', 'AR', 'HR'] - - print(f"Original data shape: {data.shape}") - print(f"Original ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{data[:, i].min():.4f}, {data[:, i].max():.4f}]") - - # Create scaler manager - scaler_manager = IndividualScalerManager(normalization_type='minmax') - - # Normalize data - normalized_data = scaler_manager.fit_transform_scalar(data, variable_names) - - print(f"\nNormalized data shape: {normalized_data.shape}") - print(f"Normalized ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{normalized_data[:, i].min():.4f}, {normalized_data[:, i].max():.4f}]") - - # Inverse transform - denormalized_data = scaler_manager.inverse_transform_scalar(normalized_data, variable_names) - - print(f"\nDenormalized data shape: {denormalized_data.shape}") - print(f"Denormalized ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{denormalized_data[:, i].min():.4f}, {denormalized_data[:, i].max():.4f}]") - - # Check accuracy - mse = np.mean((data - denormalized_data) ** 2) - print(f"\nReconstruction MSE: {mse:.2e}") - - if mse < 1e-10: - print("✅ Scalar normalization test PASSED") - return True - else: - print("❌ Scalar normalization test FAILED") - return False - - -def test_pft_1d_normalization(): - """Test PFT1D variable normalization.""" - print("\nTesting PFT1D normalization...") - - # Create sample PFT1D data - np.random.seed(42) - n_samples = 1000 - n_pfts = 17 - n_variables = 2 - - # tlai: 0-10 m²/m² - tlai = np.random.uniform(0, 10, (n_samples, n_pfts, 1)) - # deadstemc: 0-1000 gC/m² - deadstemc = np.random.uniform(0, 1000, (n_samples, n_pfts, 1)) - - # Combine into array - data = np.concatenate([tlai, deadstemc], axis=2) - pft_names = [f'PFT{i}' for i in range(n_pfts)] - variable_names = ['tlai', 'deadstemc'] - - print(f"Original PFT1D data shape: {data.shape}") - print(f"Original ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{data[:, :, i].min():.4f}, {data[:, :, i].max():.4f}]") - - # Create scaler manager - scaler_manager = IndividualScalerManager(normalization_type='minmax') - - # Normalize data - normalized_data = scaler_manager.fit_transform_pft_1d(data, pft_names, variable_names) - - print(f"Normalized PFT1D data shape: {normalized_data.shape}") - print(f"Normalized ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{normalized_data[:, :, i].min():.4f}, {normalized_data[:, :, i].max():.4f}]") - - # Inverse transform - denormalized_data = scaler_manager.inverse_transform_pft_1d(normalized_data, pft_names, variable_names) - - print(f"Denormalized PFT1D data shape: {denormalized_data.shape}") - print(f"Denormalized ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{denormalized_data[:, :, i].min():.4f}, {denormalized_data[:, :, i].max():.4f}]") - - # Check accuracy - mse = np.mean((data - denormalized_data) ** 2) - print(f"Reconstruction MSE: {mse:.2e}") - - if mse < 1e-10: - print("✅ PFT1D normalization test PASSED") - return True - else: - print("❌ PFT1D normalization test FAILED") - return False - - -def test_soil_2d_normalization(): - """Test Soil2D variable normalization.""" - print("\nTesting Soil2D normalization...") - - # Create sample Soil2D data - np.random.seed(42) - n_samples = 1000 - n_variables = 2 - n_columns = 360 # Grid columns - n_layers = 10 # Soil layers - - # cwdc_vr: 0-1000 gC/m² - cwdc_vr = np.random.uniform(0, 1000, (n_samples, 1, n_columns, n_layers)) - # sminn_vr: 0-50 gN/m² - sminn_vr = np.random.uniform(0, 50, (n_samples, 1, n_columns, n_layers)) - - # Combine into array - data = np.concatenate([cwdc_vr, sminn_vr], axis=1) - variable_names = ['cwdc_vr', 'sminn_vr'] - - print(f"Original Soil2D data shape: {data.shape}") - print(f"Original ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{data[:, i, :, :].min():.4f}, {data[:, i, :, :].max():.4f}]") - - # Create scaler manager - scaler_manager = IndividualScalerManager(normalization_type='minmax') - - # Normalize data - normalized_data = scaler_manager.fit_transform_soil_2d(data, variable_names, n_layers) - - print(f"Normalized Soil2D data shape: {normalized_data.shape}") - print(f"Normalized ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{normalized_data[:, i, :, :].min():.4f}, {normalized_data[:, i, :, :].max():.4f}]") - - # Inverse transform - denormalized_data = scaler_manager.inverse_transform_soil_2d(normalized_data, variable_names, n_layers) - - print(f"Denormalized Soil2D data shape: {denormalized_data.shape}") - print(f"Denormalized ranges:") - for i, name in enumerate(variable_names): - print(f" {name}: [{denormalized_data[:, i, :, :].min():.4f}, {denormalized_data[:, i, :, :].max():.4f}]") - - # Check accuracy - mse = np.mean((data - denormalized_data) ** 2) - print(f"Reconstruction MSE: {mse:.2e}") - - if mse < 1e-10: - print("✅ Soil2D normalization test PASSED") - return True - else: - print("❌ Soil2D normalization test FAILED") - return False - - -def test_scaler_persistence(): - """Test saving and loading scalers.""" - print("\nTesting scaler persistence...") - - # Create sample data - np.random.seed(42) - n_samples = 100 - data = np.random.uniform(0, 100, (n_samples, 3)) - variable_names = ['var1', 'var2', 'var3'] - - # Create scaler manager and fit - scaler_manager = IndividualScalerManager() - normalized_data = scaler_manager.fit_transform_scalar(data, variable_names) - - # Create temporary directory - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = Path(temp_dir) - - # Save scalers - scaler_manager.save_scalers(temp_path) - - # Check if files were created - scaler_files = list(temp_path.glob("*.pkl")) - metadata_file = temp_path / "scaler_metadata.json" - - print(f"Saved {len(scaler_files)} scaler files") - print(f"Metadata file exists: {metadata_file.exists()}") - - # Create new scaler manager and load - new_scaler_manager = IndividualScalerManager() - new_scaler_manager.load_scalers(temp_path) - - # Test inverse transform with loaded scalers - denormalized_data = new_scaler_manager.inverse_transform_scalar(normalized_data, variable_names) - - # Check accuracy - mse = np.mean((data - denormalized_data) ** 2) - print(f"Reconstruction MSE after save/load: {mse:.2e}") - - if mse < 1e-10: - print("✅ Scaler persistence test PASSED") - return True - else: - print("❌ Scaler persistence test FAILED") - return False - - -def main(): - """Run all tests.""" - print("🧪 Testing IndividualScalerManager") - print("=" * 50) - - tests = [ - test_scalar_normalization, - test_pft_1d_normalization, - test_soil_2d_normalization, - test_scaler_persistence - ] - - passed = 0 - total = len(tests) - - for test in tests: - try: - if test(): - passed += 1 - except Exception as e: - print(f"❌ Test {test.__name__} failed with error: {e}") - - print("\n" + "=" * 50) - print(f"Test Results: {passed}/{total} tests PASSED") - - if passed == total: - print("🎉 All tests passed! IndividualScalerManager is ready for integration.") - return True - else: - print("⚠️ Some tests failed. Please review the implementation.") - return False - - -if __name__ == "__main__": - success = main() - sys.exit(0 if success else 1) diff --git a/test_normalization_approaches.py b/test_normalization_approaches.py deleted file mode 100644 index 820fa59..0000000 --- a/test_normalization_approaches.py +++ /dev/null @@ -1,370 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script for different normalization approaches - -This script demonstrates the three normalization approaches: -1. Group normalization (default) -2. Individual normalization -3. Hybrid normalization (mix of both) -""" - -import sys -import os -import numpy as np -import pandas as pd -import torch -from pathlib import Path - -# Add the project root to the path -sys.path.insert(0, os.path.dirname(__file__)) - -try: - from config.training_config import DataConfig, PreprocessingConfig - from data.data_loader_individual import DataLoaderIndividual - print("✅ Successfully imported required modules") -except ImportError as e: - print(f"❌ Import error: {e}") - print("Please ensure you're running this from the project root directory") - sys.exit(1) - - -def create_mock_data(): - """Create mock data for testing.""" - print("🧪 Creating mock data...") - - # Create sample data - n_samples = 50 # Smaller dataset for testing - - # Mock DataFrame with required columns - mock_data = { - # Time series data - 'FLDS': [np.random.uniform(100, 500, 240) for _ in range(n_samples)], - 'PSRF': [np.random.uniform(80000, 120000, 240) for _ in range(n_samples)], - 'FSDS': [np.random.uniform(-50, 50, 240) for _ in range(n_samples)], - 'QBOT': [np.random.uniform(0, 0.01, 240) for _ in range(n_samples)], - 'PRECTmms': [np.random.uniform(0, 0.01, 240) for _ in range(n_samples)], - 'TBOT': [np.random.uniform(200, 350, 240) for _ in range(n_samples)], - - # Static data - 'Latitude': np.random.uniform(30, 60, n_samples), - 'Longitude': np.random.uniform(-120, -60, n_samples), - 'AREA': np.random.uniform(1e10, 1e12, n_samples), - - # Scalar data with different ranges to demonstrate individual normalization benefits - 'GPP': np.random.uniform(0, 20, n_samples), # 0-20 range - 'NPP': np.random.uniform(0, 15, n_samples), # 0-15 range - 'AR': np.random.uniform(0, 10, n_samples), # 0-10 range - 'HR': np.random.uniform(0, 12, n_samples), # 0-12 range - - # Y scalar data - 'Y_GPP': np.random.uniform(0, 20, n_samples), - 'Y_NPP': np.random.uniform(0, 15, n_samples), - 'Y_AR': np.random.uniform(0, 10, n_samples), - 'Y_HR': np.random.uniform(0, 12, n_samples), - - # PFT1D data - 'tlai': [np.random.uniform(0, 10, 17) for _ in range(n_samples)], - 'deadstemc': [np.random.uniform(0, 1000, 17) for _ in range(n_samples)], - - # Y PFT1D data - 'Y_tlai': [np.random.uniform(0, 10, 17) for _ in range(n_samples)], - 'Y_deadstemc': [np.random.uniform(0, 1000, 17) for _ in range(n_samples)], - - # Soil2D data - 'cwdc_vr': [np.random.uniform(0, 1000, (360, 10)) for _ in range(n_samples)], - 'sminn_vr': [np.random.uniform(0, 50, (360, 10)) for _ in range(n_samples)], - - # Y Soil2D data - 'Y_cwdc_vr': [np.random.uniform(0, 1000, (360, 10)) for _ in range(n_samples)], - 'Y_sminn_vr': [np.random.uniform(0, 50, (360, 10)) for _ in range(n_samples)], - - # PFT param data - 'pft_deadwdcn': [np.random.uniform(0, 100, 17) for _ in range(n_samples)], - 'pft_frootcn': [np.random.uniform(0, 100, 17) for _ in range(n_samples)], - } - - df = pd.DataFrame(mock_data) - print(f"✅ Created mock DataFrame with {len(df)} samples and {len(df.columns)} columns") - return df - - -def test_group_normalization(data_loader): - """Test group normalization approach.""" - print("\n🧪 Testing GROUP normalization approach...") - print("=" * 50) - - try: - # Use group normalization (default) - normalized_data = data_loader.normalize_data() - - print("✅ Group normalization completed successfully") - print(f"✅ Scalar data shape: {normalized_data['scalar_data'].shape}") - print(f"✅ Y scalar data shape: {normalized_data['y_scalar'].shape}") - print(f"✅ PFT1D data shape: {normalized_data['variables_1d_pft'].shape}") - print(f"✅ Soil2D data shape: {normalized_data['variables_2d_soil'].shape}") - - # Check scaler types - scalers = normalized_data['scalers'] - print("\n📊 Scaler types used:") - for name, scaler in scalers.items(): - if scaler is not None: - if hasattr(scaler, 'normalization_type'): - print(f" {name}: IndividualScalerManager ({scaler.normalization_type})") - else: - print(f" {name}: {type(scaler).__name__}") - - return normalized_data - - except Exception as e: - print(f"❌ Group normalization failed: {e}") - return None - - -def test_individual_normalization(data_loader): - """Test individual normalization approach.""" - print("\n🧪 Testing INDIVIDUAL normalization approach...") - print("=" * 50) - - try: - # Use individual normalization - normalized_data = data_loader.normalize_data_individual() - - print("✅ Individual normalization completed successfully") - print(f"✅ Scalar data shape: {normalized_data['scalar_data'].shape}") - print(f"✅ Y scalar data shape: {normalized_data['y_scalar'].shape}") - print(f"✅ PFT1D data shape: {normalized_data['variables_1d_pft'].shape}") - print(f"✅ Soil2D data shape: {normalized_data['variables_2d_soil'].shape}") - - # Check scaler types - scalers = normalized_data['scalers'] - print("\n📊 Scaler types used:") - for name, scaler in scalers.items(): - if scaler is not None: - if hasattr(scaler, 'normalization_type'): - print(f" {name}: IndividualScalerManager ({scaler.normalization_type})") - else: - print(f" {name}: {type(scaler).__name__}") - - return normalized_data - - except Exception as e: - print(f"❌ Individual normalization failed: {e}") - return None - - -def test_hybrid_normalization(data_loader): - """Test hybrid normalization approach.""" - print("\n🧪 Testing HYBRID normalization approach...") - print("=" * 50) - - try: - # Use hybrid normalization - individual for scalars, group for others - use_individual_for = ['scalar', 'y_scalar'] - normalized_data = data_loader.normalize_data_hybrid(use_individual_for) - - print("✅ Hybrid normalization completed successfully") - print(f"✅ Individual normalization used for: {use_individual_for}") - print(f"✅ Scalar data shape: {normalized_data['scalar_data'].shape}") - print(f"✅ Y scalar data shape: {normalized_data['y_scalar'].shape}") - print(f"✅ PFT1D data shape: {normalized_data['variables_1d_pft'].shape}") - print(f"✅ Soil2D data shape: {normalized_data['variables_2d_soil'].shape}") - - # Check scaler types - scalers = normalized_data['scalers'] - print("\n📊 Scaler types used:") - for name, scaler in scalers.items(): - if scaler is not None: - if hasattr(scaler, 'normalization_type'): - print(f" {name}: IndividualScalerManager ({scaler.normalization_type})") - else: - print(f" {name}: {type(scaler).__name__}") - - return normalized_data - - except Exception as e: - print(f"❌ Hybrid normalization failed: {e}") - return None - - -def test_scaler_persistence(data_loader): - """Test that scalers can be saved and loaded.""" - print("\n🧪 Testing scaler persistence...") - print("=" * 50) - - try: - # Create temporary directory - import tempfile - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = Path(temp_dir) - - # Save scalers - data_loader.save_scalers(temp_path) - print(f"✅ Scalers saved to {temp_path}") - - # Check if files were created - scaler_dir = temp_path / "scalers" - if scaler_dir.exists(): - print(f"✅ Scaler directory created: {scaler_dir}") - - # List saved files - saved_files = list(scaler_dir.rglob("*")) - print(f"✅ Saved {len(saved_files)} scaler files") - - for file_path in saved_files: - if file_path.is_file(): - print(f" - {file_path.name}") - else: - print("⚠️ Scaler directory not created") - - # Test loading scalers (create new instance) - new_data_loader = DataLoaderIndividual( - data_loader.data_config, - data_loader.preprocessing_config - ) - - new_data_loader.load_scalers(temp_path) - print("✅ Scalers loaded successfully") - - return True - - except Exception as e: - print(f"❌ Failed to test scaler persistence: {e}") - return False - - -def compare_normalization_approaches(group_data, individual_data, hybrid_data): - """Compare the different normalization approaches.""" - print("\n🧪 Comparing normalization approaches...") - print("=" * 50) - - try: - # Compare data shapes - print("📊 Data shapes comparison:") - data_types = ['scalar_data', 'y_scalar', 'variables_1d_pft', 'variables_2d_soil'] - - for data_type in data_types: - if data_type in group_data and data_type in individual_data and data_type in hybrid_data: - group_shape = group_data[data_type].shape - individual_shape = individual_data[data_type].shape - hybrid_shape = hybrid_data[data_type].shape - - print(f" {data_type}:") - print(f" Group: {group_shape}") - print(f" Individual: {individual_shape}") - print(f" Hybrid: {hybrid_shape}") - - # Check if shapes are consistent - if group_shape == individual_shape == hybrid_shape: - print(f" ✅ All approaches produce consistent shapes") - else: - print(f" ⚠️ Shape mismatch detected") - - # Compare scaler counts - print("\n📊 Scaler counts comparison:") - print(f" Group normalization: {len([s for s in group_data['scalers'].values() if s is not None])} scalers") - print(f" Individual normalization: {len([s for s in individual_data['scalers'].values() if s is not None])} scalers") - print(f" Hybrid normalization: {len([s for s in hybrid_data['scalers'].values() if s is not None])} scalers") - - return True - - except Exception as e: - print(f"❌ Failed to compare normalization approaches: {e}") - return False - - -def main(): - """Run all normalization approach tests.""" - print("🧪 Testing Different Normalization Approaches") - print("=" * 70) - - # Create data loader - try: - data_config = DataConfig( - data_paths=["/tmp"], # Dummy path - file_pattern="*.pkl", - time_series_columns=["FLDS", "PSRF", "FSDS", "QBOT", "PRECTmms", "TBOT"], - time_series_length=240, - static_columns=["Latitude", "Longitude", "AREA"], - x_list_scalar_columns=["GPP", "NPP", "AR", "HR"], - y_list_scalar_columns=["Y_GPP", "Y_NPP", "Y_AR", "Y_HR"], - x_list_columns_1d=["tlai", "deadstemc"], - y_list_columns_1d=["Y_tlai", "Y_deadstemc"], - x_list_columns_2d=["cwdc_vr", "sminn_vr"], - y_list_columns_2d=["Y_cwdc_vr", "Y_sminn_vr"], - pft_param_columns=["pft_deadwdcn", "pft_frootcn"], - max_1d_length=17, - max_2d_rows=360, - max_2d_cols=10, - random_state=42 - ) - - preprocessing_config = PreprocessingConfig( - time_series_normalization="minmax", - static_normalization="minmax", - list_1d_normalization="minmax", - list_2d_normalization="minmax", - data_type=torch.float32 - ) - - data_loader = DataLoaderIndividual(data_config, preprocessing_config) - print("✅ DataLoaderIndividual created successfully") - - except Exception as e: - print(f"❌ Failed to create DataLoaderIndividual: {e}") - return False - - # Integrate mock data - mock_df = create_mock_data() - data_loader.df = mock_df - - # Preprocess data - try: - data_loader.preprocess_data() - print("✅ Data preprocessing completed") - except Exception as e: - print(f"❌ Data preprocessing failed: {e}") - return False - - # Test 1: Group normalization - group_data = test_group_normalization(data_loader) - if group_data is None: - print("❌ Cannot proceed without group normalization") - return False - - # Test 2: Individual normalization - individual_data = test_individual_normalization(data_loader) - if individual_data is None: - print("❌ Cannot proceed without individual normalization") - return False - - # Test 3: Hybrid normalization - hybrid_data = test_hybrid_normalization(data_loader) - if hybrid_data is None: - print("❌ Cannot proceed without hybrid normalization") - return False - - # Test 4: Compare approaches - if not compare_normalization_approaches(group_data, individual_data, hybrid_data): - print("⚠️ Comparison failed") - - # Test 5: Scaler persistence - if not test_scaler_persistence(data_loader): - print("⚠️ Scaler persistence test failed") - - print("\n" + "=" * 70) - print("🎉 All normalization approach tests completed!") - print("\n📋 Summary of approaches:") - print("1. GROUP normalization: Default approach, uses one scaler per data type") - print("2. INDIVIDUAL normalization: Optimal approach, uses individual scalers for each variable") - print("3. HYBRID normalization: Flexible approach, mix of both methods") - print("\n💡 Usage recommendations:") - print("- Use GROUP for: Quick testing, memory-constrained environments") - print("- Use INDIVIDUAL for: Production training, optimal performance") - print("- Use HYBRID for: Selective optimization, specific variable types") - - return True - - -if __name__ == "__main__": - success = main() - sys.exit(0 if success else 1) From 1c8da06fd26c71ab2204d0305f7f274044c01d15 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Wed, 15 Oct 2025 17:04:02 -0400 Subject: [PATCH 21/51] add multi datasets and fix the bugs in quality report generation --- CNP_IO_updated9_dev.txt | 24 ++- config/training_config.py | 55 +++++- data/data_loader_individual.py | 8 +- docs/CNP_pipeline_runbook.md | 2 +- scripts/cnp_result_validationplot.py | 9 +- scripts/generate_prediction_quality_report.py | 170 ++++++++++++++++-- train_cnp_model.py | 57 +++++- 7 files changed, 294 insertions(+), 31 deletions(-) diff --git a/CNP_IO_updated9_dev.txt b/CNP_IO_updated9_dev.txt index f40c693..c543cd4 100644 --- a/CNP_IO_updated9_dev.txt +++ b/CNP_IO_updated9_dev.txt @@ -1,8 +1,18 @@ -TRENDY1_PATH: /mnt/proj-shared/AI4BGC_7xw/TrainingData//Trendy_1_data_CNP -TRENDY05_PATH = /mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP -DATA_PATHS: /path/extra1, /path/extra2 +# 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 +# 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 @@ -48,13 +58,13 @@ SCALAR VARIABLES (1D - 4 variables): • tlai, totvegc -2D VARIABLES (layered - 28 variables): +2D VARIABLES (layered - 25 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/config/training_config.py b/config/training_config.py index 10a6e0f..6bb0f4a 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -27,6 +27,8 @@ class DataConfig: # File patterns file_pattern: str = "1_training_data_batch_*.pkl" + # Optional per-dataset file patterns keyed by absolute path + dataset_file_patterns: Dict[str, str] = field(default_factory=dict) # Columns to drop columns_to_drop: List[str] = field(default_factory=lambda: [ @@ -435,13 +437,26 @@ def parse_cnp_io_list(filename): result.update({ 'trendy1_path': None, 'trendy05_path': None, - 'file_pattern': None + 'tva4km_path': None, + 'file_pattern': None, + 'trendy1_file_pattern': None, + 'trendy05_file_pattern': None, + 'tva4km_file_pattern': None }) current_section = None with open(filename) as f: for line in f: line = line.strip() + # Skip full-line comments and blanks (support //, #, ;) + if not line or line.startswith('#') or line.startswith('//') or line.startswith(';'): + continue + # Remove inline comments introduced by '#'; avoid '//' inline to not break paths + hash_idx = line.find('#') + if hash_idx != -1: + line = line[:hash_idx].strip() + if not line: + continue # Section header detection for section_title, key in section_map.items(): @@ -479,7 +494,7 @@ def parse_cnp_io_list(filename): # FILE_PATTERN: enhanced_1_training_data_batch_*.pkl # DATA_PATHS: /p1,/p2 if line and not line.startswith('#'): - kv_match = re.match(r'(?i)^(trendy1_path|trendy05_path|file_pattern|data_paths)\s*[:=]\s*(.+)$', line) + kv_match = re.match(r'(?i)^(trendy1_path|trendy05_path|tva4km_path|file_pattern|trendy1_file_pattern|trendy05_file_pattern|tva4km_file_pattern|data_paths)\s*[:=]\s*(.+)$', line) if kv_match: key = kv_match.group(1).lower() val = kv_match.group(2).strip() @@ -493,6 +508,14 @@ def parse_cnp_io_list(filename): result['trendy1_path'] = val elif key == 'trendy05_path': result['trendy05_path'] = val + elif key == 'tva4km_path': + result['tva4km_path'] = val + elif key == 'trendy1_file_pattern': + result['trendy1_file_pattern'] = val + elif key == 'trendy05_file_pattern': + result['trendy05_file_pattern'] = val + elif key == 'tva4km_file_pattern': + result['tva4km_file_pattern'] = val return result def parse_cnp_model_config(filename: str) -> Dict[str, Any]: @@ -594,6 +617,7 @@ def parse_cnp_model_config(filename: str) -> Dict[str, Any]: def get_cnp_combined_config( use_trendy1: bool = True, use_trendy05: bool = True, + use_tva4km: bool = False, max_files: Optional[int] = None, include_water: bool = False, variable_list_path: Optional[str] = None, @@ -612,6 +636,7 @@ def get_cnp_combined_config( config = TrainingConfigManager() data_paths = [] file_pattern = None + dataset_file_patterns: Dict[str, str] = {} # If a variable list is provided, prefer dataset paths from it parsed = None if variable_list_path is not None: @@ -620,13 +645,24 @@ def get_cnp_combined_config( except Exception as e: logging.warning(f"Failed to parse variable list for data paths: {e}") if parsed is not None: - # Collect from any or all of: data_paths, trendy1_path, trendy05_path + # Collect from any or all of: data_paths, trendy1_path, trendy05_path, tva4km_path if parsed.get('data_paths'): data_paths.extend([p for p in parsed['data_paths'] if p]) if parsed.get('trendy1_path') and use_trendy1: - data_paths.append(parsed['trendy1_path']) + p = parsed['trendy1_path'] + data_paths.append(p) + if parsed.get('trendy1_file_pattern'): + dataset_file_patterns[p] = parsed['trendy1_file_pattern'] if parsed.get('trendy05_path') and use_trendy05: - data_paths.append(parsed['trendy05_path']) + p = parsed['trendy05_path'] + data_paths.append(p) + if parsed.get('trendy05_file_pattern'): + dataset_file_patterns[p] = parsed['trendy05_file_pattern'] + if parsed.get('tva4km_path') and use_tva4km: + p = parsed['tva4km_path'] + data_paths.append(p) + if parsed.get('tva4km_file_pattern'): + dataset_file_patterns[p] = parsed['tva4km_file_pattern'] if parsed.get('file_pattern'): file_pattern = parsed['file_pattern'] # Fallback to defaults if none provided via CNP_IO @@ -635,12 +671,21 @@ def get_cnp_combined_config( data_paths.append("/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree") if use_trendy05: data_paths.append("/mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP") + if use_tva4km: + # Prefer environment variable if provided + env_tva = os.environ.get('TVA4KM_PATH') + if env_tva: + data_paths.append(env_tva) + env_pat = os.environ.get('TVA4KM_FILE_PATTERN') + if env_pat: + dataset_file_patterns[env_tva] = env_pat if file_pattern is None: file_pattern = "enhanced_1_training_data_batch_*.pkl" config.update_data_config( data_paths=data_paths, file_pattern=file_pattern, + dataset_file_patterns=dataset_file_patterns, max_files=max_files, train_split=0.8, filter_column=None, diff --git a/data/data_loader_individual.py b/data/data_loader_individual.py index 565619f..d7ba44d 100644 --- a/data/data_loader_individual.py +++ b/data/data_loader_individual.py @@ -102,7 +102,13 @@ def load_data(self) -> pd.DataFrame: logger.info("Loading data from multiple paths...") for path in self.data_config.data_paths: # Resolve files matching pattern - files = list(Path(path).glob(self.data_config.file_pattern)) + # Support per-dataset file patterns if provided + try: + per_dataset_patterns = getattr(self.data_config, 'dataset_file_patterns', {}) or {} + except Exception: + per_dataset_patterns = {} + pattern = per_dataset_patterns.get(path, self.data_config.file_pattern) + files = list(Path(path).glob(pattern)) # Deterministic ordering for test runs if getattr(self.data_config, 'sort_file_list', True): files = sorted(files, key=lambda p: p.name) diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 107eb3c..0b7e6f9 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -3,7 +3,7 @@ This guide walks through the end-to-end workflow using a user-defined CNP_IO list. ### 1) Create your CNP_IO list -Use the CNP_IO template to create a user-defined list, e.g. `CNP_IO_demo1.txt`). +Use the CNP_IO template to create a user-defined list, e.g. `CNP_IO_demo.txt`). ### 2) Train the AI model Run training with your variable list (edit the filename as needed): diff --git a/scripts/cnp_result_validationplot.py b/scripts/cnp_result_validationplot.py index 5c89c67..b6e2558 100644 --- a/scripts/cnp_result_validationplot.py +++ b/scripts/cnp_result_validationplot.py @@ -121,7 +121,7 @@ def main_with_flag(results_dir, plot_scatter, plot_loss, top_bad_only=False, top scalar_pred = os.path.join(results_dir, 'cnp_predictions', 'predictions_scalar.csv') if os.path.exists(scalar_gt) and os.path.exists(scalar_pred): print("Analyzing scalar data...") - analyze_pair(scalar_gt, scalar_pred, 'Scalar', plots_dir, stats_data, per_column=True, plot_scatter=plot_scatter) + analyze_pair(scalar_gt, scalar_pred, 'Scalar', plots_dir, stats_data, per_column=True, plot_scatter=plot_scatter, selection=selection) else: print("Scalar data files not found") @@ -169,8 +169,13 @@ def analyze_pair(gt_path, pred_path, label, out_dir, stats_data, per_column=Fals if per_column: # Per-column comparison for scalar for col in gt.columns: + # Normalize variable name by stripping Y_ for selection matching + col_norm = col[2:] if isinstance(col, str) and col.startswith('Y_') else col + # Skip Latitude/Longitude if top-bad-only was requested (selection provided) + if selection is not None and str(col_norm) in ('Latitude', 'Longitude'): + continue # If selection provided, only include scalar variables present in selection - if selection is not None and col not in selection: + if selection is not None and col_norm not in selection: continue if col in pred.columns: print(f"Analyzing variable: {col}") diff --git a/scripts/generate_prediction_quality_report.py b/scripts/generate_prediction_quality_report.py index 27031df..56c8d40 100644 --- a/scripts/generate_prediction_quality_report.py +++ b/scripts/generate_prediction_quality_report.py @@ -7,6 +7,8 @@ import argparse import importlib.util import sys +import json +from typing import Optional, List def main(): parser = argparse.ArgumentParser(description='Generate prediction quality report from validation statistics') @@ -14,6 +16,8 @@ def main(): help='Path to validation statistics CSV file') 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, @@ -69,6 +73,37 @@ def main(): 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': { @@ -86,6 +121,13 @@ def main(): 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) @@ -138,6 +180,13 @@ def categorize_prediction(row): # 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']: @@ -180,6 +229,12 @@ def categorize_prediction(row): 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) @@ -272,11 +327,20 @@ def categorize_prediction(row): plt.savefig(output_dir / "r2_vs_rmse.png", dpi=300) # 4. Optionally generate top-bad-only plots 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) + + # 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 @@ -286,15 +350,31 @@ def categorize_prediction(row): mod = importlib.util.module_from_spec(spec) sys.modules['cnp_plot_mod'] = mod assert spec.loader is not None - 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=True, - top_bad_report=str(output_dir / 'quality_summary_report.txt'), - plots_dir_override=top_bad_out) - print(f"Top-bad plots saved to: {top_bad_out}") - else: - print("Warning: cnp_result_validationplot.main_with_flag not found; skipping top-bad plots") + 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=True, + top_bad_report=str(output_dir / 'quality_summary_report.txt'), + plots_dir_override=top_bad_out) + 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').glob('*.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}") @@ -305,11 +385,21 @@ def categorize_prediction(row): # 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") @@ -385,15 +475,37 @@ def categorize_prediction(row): f.write("## Variables with Best Predictions\n") if 'good_pct' in variable_summary.columns: - best_vars = variable_summary.nlargest(15, 'good_pct') + # 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: - worst_vars = variable_summary.nsmallest(15, 'good_pct') + _vs2 = variable_summary.copy() + _vs2['good_pct'] = _vs2['good_pct'].fillna(0) + _vs2['ok_pct'] = _vs2.get('ok_pct', 0) + _vs2['bad_pct'] = _vs2.get('bad_pct', 0) + worst_vars = _vs2.nsmallest(15, 'good_pct') 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") + + # 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'}") @@ -595,7 +707,10 @@ def categorize_prediction(row): """ # Add worst variables - worst_vars = variable_summary.nsmallest(15, 'good_pct') + _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""" @@ -606,6 +721,27 @@ def categorize_prediction(row): """ + # 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 += """ @@ -624,12 +760,18 @@ def categorize_prediction(row): print(f"{quality}: {count} ({count/len(analysis_df)*100:.1f}%)") print("\nTop 5 Best Predicted Variables:") - best_vars = variable_summary.nlargest(5, 'good_pct') + _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:") - worst_vars = variable_summary.nsmallest(5, 'good_pct') + _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") diff --git a/train_cnp_model.py b/train_cnp_model.py index 3d11274..22f94ec 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -186,6 +186,11 @@ def main(): action='store_true', help='Include Trend_05_data_CNP dataset' ) + parser.add_argument( + '--use-tva4km', + action='store_true', + help='Include TVA4km dataset' + ) parser.add_argument( '--variable-list', type=str, @@ -285,7 +290,7 @@ def main(): elif args.normalization == 'hybrid': logger.info("Using hybrid normalization (selective individual/group)") - if not args.use_trendy1 and not args.use_trendy05: + if not args.use_trendy1 and not args.use_trendy05 and not getattr(args, 'use_tva4km', False): args.use_trendy1 = True args.use_trendy05 = False @@ -300,6 +305,7 @@ def main(): config = get_cnp_combined_config( use_trendy1=args.use_trendy1, use_trendy05=args.use_trendy05, + use_tva4km=args.use_tva4km, max_files=args.max_files, include_water=include_water, variable_list_path=args.variable_list, @@ -639,10 +645,59 @@ def main(): # Save model configuration with open(output_dir / "cnp_config.json", "w") as f: + # Derive per-group counts for quick comparison with CNP_IO.txt + di = data_info if isinstance(data_info, dict) else {} + time_series_cols = di.get('time_series_columns', []) or [] + static_cols = di.get('static_columns', []) or [] + pft_param_cols = di.get('pft_param_columns', []) or [] + scalar_in_cols = di.get('x_list_scalar_columns', []) or [] + scalar_out_cols = di.get('y_list_scalar_columns', []) or [] + pft1d_in_vars = di.get('variables_1d_pft', []) or [] + pft1d_out_vars = di.get('y_list_columns_1d', []) or [] + soil2d_in_vars = di.get('x_list_columns_2d', []) or [] + soil2d_out_vars = di.get('y_list_columns_2d', []) or [] + + group_counts = { + 'time_series_variables': len(time_series_cols), + 'static_columns': len(static_cols), + 'pft_param_columns': len(pft_param_cols), + 'scalar_variables_in': len(scalar_in_cols), + 'scalar_variables_out': len(scalar_out_cols), + 'pft_1d_variables_in': len(pft1d_in_vars), + 'pft_1d_variables_out': len(pft1d_out_vars), + 'soil_2d_variables_in': len(soil2d_in_vars), + 'soil_2d_variables_out': len(soil2d_out_vars), + 'total_predicted_variables': len(scalar_out_cols) + len(pft1d_out_vars) + len(soil2d_out_vars) + } + + # Expanded prediction element counts (PFTs and Soil layers) + # PFTs per variable use model_config.vector_length (expected 16: PFT1..PFT16) + pfts_per_var = int(getattr(config.model_config, 'vector_length', 16) or 16) + soil_rows_per_var = int(getattr(config.model_config, 'matrix_rows', 1) or 1) + soil_layers_per_var = int(getattr(config.model_config, 'matrix_cols', 10) or 10) + prediction_element_counts = { + 'pft_1d': { + 'variables_out': len(pft1d_out_vars), + 'pfts_per_variable': pfts_per_var, + 'total_elements': len(pft1d_out_vars) * pfts_per_var + }, + 'soil_2d': { + 'variables_out': len(soil2d_out_vars), + 'columns_per_variable': soil_rows_per_var, + 'layers_per_variable': soil_layers_per_var, + 'total_elements': len(soil2d_out_vars) * soil_rows_per_var * soil_layers_per_var + }, + 'scalar_1d': { + 'variables_out': len(scalar_out_cols) + } + } + config_dict = { 'include_water': include_water, 'normalization_method': args.normalization, 'data_info': data_info, + 'data_counts': group_counts, + 'prediction_element_counts': prediction_element_counts, 'model_config': config.model_config.__dict__, 'training_config': config.training_config.__dict__ } From 378170f4733c24980cfdac47498e4f40b348cd35 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Wed, 15 Oct 2025 19:10:05 -0400 Subject: [PATCH 22/51] fix analysis scritps --- CNP_IO_updated9_dev.txt | 2 +- config/training_config.py | 48 +++++++++++++++++++++------- docs/CNP_pipeline_runbook.md | 2 +- scripts/cnp_result_validationplot.py | 2 ++ train_cnp_model.py | 5 ++- 5 files changed, 45 insertions(+), 14 deletions(-) diff --git a/CNP_IO_updated9_dev.txt b/CNP_IO_updated9_dev.txt index c543cd4..1a2837b 100644 --- a/CNP_IO_updated9_dev.txt +++ b/CNP_IO_updated9_dev.txt @@ -43,7 +43,7 @@ PFT PARAMETERS (Plant Functional Type Characteristics) - 44 variables: SCALAR VARIABLES (1D - 4 variables): • GPP, NPP, AR, HR -1D PFT VARIABLES (39 variables): +1D PFT VARIABLES (41 variables): • deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage • deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage diff --git a/config/training_config.py b/config/training_config.py index 6bb0f4a..6bdc4ee 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -734,18 +734,38 @@ def get_cnp_combined_config( # 'secondp_vr' # this is not in the list, but it is in the data ] - # If a variable list file is provided, parse it + # If a variable list file is provided, use it for variable groups; otherwise use defaults longitudes_to_drop = [] - if variable_list_path is not None and parsed is None: - parsed = parse_cnp_io_list(variable_list_path) - time_series_columns = parsed.get('time_series_variables', default_time_series) - surface_properties = parsed.get('surface_properties', default_surface) - pft_parameters = parsed.get('pft_parameters', default_pft_parameters) - water_variables = parsed.get('water_variables', default_water) - scalar_variables = parsed.get('scalar_variables', default_scalar) - pft_1d_variables = parsed.get('pft_1d_variables', default_pft_1d) - variables_2d_soil = parsed.get('variables_2d_soil', default_2d_soil) - longitudes_to_drop = parsed.get('longitudes_to_drop', []) + if variable_list_path is not None: + if parsed is None: + try: + parsed = parse_cnp_io_list(variable_list_path) + except Exception as _e: + logging.warning(f"Failed to parse variable list at {variable_list_path}: {_e}") + parsed = None + if parsed is not None: + time_series_columns = parsed.get('time_series_variables', default_time_series) + surface_properties = parsed.get('surface_properties', default_surface) + pft_parameters = parsed.get('pft_parameters', default_pft_parameters) + water_variables = parsed.get('water_variables', default_water) + scalar_variables = parsed.get('scalar_variables', default_scalar) + pft_1d_variables = parsed.get('pft_1d_variables', default_pft_1d) + variables_2d_soil = parsed.get('variables_2d_soil', default_2d_soil) + longitudes_to_drop = parsed.get('longitudes_to_drop', []) + try: + logging.info(f"Applied variable groups from {variable_list_path}: " + f"ts={len(time_series_columns)}, static={len(surface_properties)}, pft_params={len(pft_parameters)}, " + f"scalar={len(scalar_variables)}, pft1d={len(pft_1d_variables)}, soil2d={len(variables_2d_soil)}") + except Exception: + pass + else: + time_series_columns = default_time_series + surface_properties = default_surface + pft_parameters = default_pft_parameters + water_variables = default_water + scalar_variables = default_scalar + pft_1d_variables = default_pft_1d + variables_2d_soil = default_2d_soil else: time_series_columns = default_time_series surface_properties = default_surface @@ -843,6 +863,12 @@ def get_cnp_combined_config( if applicable: config.update_model_config(**applicable) logging.info(f"Applied {len(applicable)} ModelConfig overrides from {model_config_path}") + # Record metadata for downstream verification + try: + setattr(config, 'model_config_overrides_keys', sorted(list(applicable.keys()))) + setattr(config, 'model_config_source', model_config_path) + except Exception: + pass except Exception as e: logging.warning(f"Could not apply model config overrides from {model_config_path}: {e}") diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 0b7e6f9..1e518d4 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -66,7 +66,7 @@ python ../../scripts/ai_predictions_to_netcdf.py > ai_prediction_to_netcdf.log ### 7) Generate comparison plots Creates map plots for the variables in your CNP_IO list under -`./ai_model_comparison_plots/comparison_CNP_IO_demo1`. +`./ai_model_comparison_plots/comparison_CNP_IO_demo`. ```bash python ../../scripts/ai_model_comparison_plot.py > ai_model_comparison.log 2>&1 & diff --git a/scripts/cnp_result_validationplot.py b/scripts/cnp_result_validationplot.py index b6e2558..7d84f1b 100644 --- a/scripts/cnp_result_validationplot.py +++ b/scripts/cnp_result_validationplot.py @@ -97,6 +97,8 @@ def main_with_flag(results_dir, plot_scatter, plot_loss, top_bad_only=False, top print(f"Plotting restricted to top-bad variables from: {report_path}") else: print("No selections parsed from top-bad report; proceeding without restriction.") + # IMPORTANT: ensure unrestricted plotting by clearing selection + selection = None # Check for new directory structure first pft_gt_dir = os.path.join(results_dir, 'cnp_predictions', 'pft_1d_ground_truth') diff --git a/train_cnp_model.py b/train_cnp_model.py index 22f94ec..156742d 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -699,7 +699,10 @@ def main(): 'data_counts': group_counts, 'prediction_element_counts': prediction_element_counts, 'model_config': config.model_config.__dict__, - 'training_config': config.training_config.__dict__ + 'training_config': config.training_config.__dict__, + # Model-config provenance for verification + 'model_config_source': getattr(config, 'model_config_source', None), + 'model_config_overrides_keys': getattr(config, 'model_config_overrides_keys', None) } json.dump(config_dict, f, indent=2) From 62da6c1d0d55064fa5b94fd0e140b4fd0e1cd8c9 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Wed, 15 Oct 2025 20:14:52 -0400 Subject: [PATCH 23/51] fix run_inference with model config option --- scripts/run_inference_all.py | 53 ++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/scripts/run_inference_all.py b/scripts/run_inference_all.py index 30ff5eb..a6f22fd 100644 --- a/scripts/run_inference_all.py +++ b/scripts/run_inference_all.py @@ -112,6 +112,37 @@ def _extract_variables_from_config(config_path: Path) -> dict: +# Load training model architecture from cnp_config.json +def _load_training_model_config(model_path: Path) -> dict: + """Load model_config (architecture) from cnp_config.json in the training run directory.""" + import json + try: + model_dir = Path(model_path).parent + # Search model dir then parents + search_dirs = [model_dir] + list(model_dir.parents) + for d in search_dirs: + config_path = d / 'cnp_config.json' + if config_path.exists(): + try: + with open(config_path, 'r') as f: + cfg = json.load(f) + mc = cfg.get('model_config') + if isinstance(mc, dict) and mc: + logging.info(f"Loaded training model_config from {config_path}") + return mc + else: + logging.warning(f"model_config missing in {config_path}") + return None + except Exception as e: + logging.warning(f"Failed reading model_config from {config_path}: {e}") + return None + logging.warning("No cnp_config.json found to load model_config (searched model dir and parents)") + return None + except Exception as e: + logging.warning(f"Error discovering training model_config: {e}") + return None + + @@ -227,6 +258,28 @@ def run_inference_all( else: logging.info("Using default configuration - no variable override applied") + # CRITICAL FIX: Apply training architecture (model_config) so weights match exactly + if use_training_config: + training_model_cfg = _load_training_model_config(Path(model_path)) + if training_model_cfg: + logging.info("Applying training model_config (architecture) from cnp_config.json...") + # Set attributes present in config.model_config + for key, value in training_model_cfg.items(): + try: + if hasattr(config.model_config, key): + setattr(config.model_config, key, value) + except Exception as e: + logging.warning(f"Failed to apply model_config key '{key}': {e}") + # Re-log a few critical dimensions + try: + logging.info( + f"Architecture summary: lstm_hidden_size={getattr(config.model_config, 'lstm_hidden_size', 'NA')}, " + f"pft_1d_fc_size={getattr(config.model_config, 'pft_1d_fc_size', 'NA')}, " + f"transformer_layers={getattr(config.model_config, 'transformer_layers', 'NA')}" + ) + except Exception: + pass + # Update data config with provided paths and pattern config.data_config.data_paths = [data_paths] config.data_config.file_pattern = file_pattern From 98525228ddd14f0910424a3638b43e81da76d038 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Sat, 18 Oct 2025 09:38:51 -0400 Subject: [PATCH 24/51] Update variable list for training model --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 63c74da..70f04be 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,7 +59,7 @@ jobs: source .venv/bin/activate python train_cnp_model.py \ --epoch 1 \ - --variable-list ./CNP_IO_updated14_xfer.txt \ + --variable-list ./CNP_IO_updated9_dev.txt \ --mask-absent-pfts \ --model-config ./CNP_model_config_v01.txt \ --use-trendy1 \ From bcc382ad91d9f72e262529af2a31ee307bc2a187 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Sun, 19 Oct 2025 09:16:27 -0700 Subject: [PATCH 25/51] ignore local config file --- .gitignore | 1 + CNP_IO_updated9_dev.txt | 74 ----------------------------------------- 2 files changed, 1 insertion(+), 74 deletions(-) delete mode 100644 CNP_IO_updated9_dev.txt diff --git a/.gitignore b/.gitignore index d00d3c1..d2687e9 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ cnp_results/ __pycache__/ *.pyc logs/ +CNP_IO_updated9_dev.txt diff --git a/CNP_IO_updated9_dev.txt b/CNP_IO_updated9_dev.txt deleted file mode 100644 index 1a2837b..0000000 --- a/CNP_IO_updated9_dev.txt +++ /dev/null @@ -1,74 +0,0 @@ -# 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 - -# 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, NPP, AR, HR - -1D PFT VARIABLES (41 variables): - -• deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage -• deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage - -• leafc, leafc_storage, leafn, leafn_storage, leafp, leafp_storage -• frootc, frootc_storage, frootn, frootn_storage, frootp, frootp_storage - -• livestemc, livestemc_storage, livestemn, livestemn_storage, livestemp, livestemp_storage, -• livecrootc, livecrootc_storage, livecrootn, livecrootn_storage, livecrootp, livecrootp_storage - -• cpool, npool, ppool - -• tlai, totvegc - -2D VARIABLES (layered - 25 variables): - -• cwdc_vr, cwdn_vr, cwdp_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 -• soil3c_vr, soil3n_vr, soil3p_vr -• soil4c_vr, soil4n_vr, soil4p_vr - -• labilep_vr , occlp_vr, primp_vr, secondp_vr From b3ecf08c878e3971f59eb5580d903501ce3ea751 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Mon, 20 Oct 2025 05:59:40 -0700 Subject: [PATCH 26/51] For step 7 and 9: add --stats-only to get summary statistics --- docs/CNP_pipeline_runbook.md | 15 +++- scripts/ai_model_comparison_plot.py | 79 +++++++++++++++--- scripts/ai_restart_comparison.py | 119 +++++++++++++++++++++++----- 3 files changed, 179 insertions(+), 34 deletions(-) diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 1e518d4..5600964 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -71,6 +71,13 @@ Creates map plots for the variables in your CNP_IO list under ```bash python ../../scripts/ai_model_comparison_plot.py > ai_model_comparison.log 2>&1 & ``` +#### 7.1) Stats-only mode (no plotting) +If you only need summary statistics (sum/std/min/max) for all variables and their layers/PFTs, use `--stats-only` 与 `--variable-list`： +```bash +python ../../scripts/ai_model_comparison_plot.py \ + --stats-only \ + --variable-list ../../CNP_IO_updated9_dev.txt +``` ### 8) Create a new ELM restart file using AI predictions @@ -88,7 +95,13 @@ Compares selected layers and PFTs; optionally verify with `restart_variable_plot python ../../scripts/ai_restart_comparison.py --variable-list ../../CNP_IO_demo1.txt --layers 0,5,9 --pfts 0,1,2,3,4,5 # Optionally use ../../scripts/restart_variable_plot.py for manual verification ``` - +#### 9.1) Stats-only mode (no plotting) +If you only need summary statistics (sum/std/min/max) for all variables and their layers/PFTs, use `--stats-only` 与 `--variable-list`： +```bash +python ../../scripts/ai_restart_comparison.py \ + --stats-only \ + --variable-list ../../CNP_IO_updated9_dev.txt +``` --- ### Old scripts (to be double-checked) diff --git a/scripts/ai_model_comparison_plot.py b/scripts/ai_model_comparison_plot.py index ec6d79e..f1c6dc4 100644 --- a/scripts/ai_model_comparison_plot.py +++ b/scripts/ai_model_comparison_plot.py @@ -1,4 +1,6 @@ #!/usr/bin/env python3 +import ssl +ssl._create_default_https_context = ssl._create_unverified_context import os import numpy as np import xarray as xr @@ -19,7 +21,7 @@ # Default paths DEFAULT_AI_PREDICTIONS = './comparison_results/ai_predictions_for_plotting.nc' -DEFAULT_MODEL = '/mnt/proj-shared/AI4BGC_7xw/AI4BGC/ELM_data/original_780_spinup_from_modelsimulation.nc' +DEFAULT_MODEL = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' DEFAULT_OUTPUT_DIR = "./ai_model_comparison_plots" # Default variables to plot unless '--variables all' is used @@ -185,10 +187,12 @@ def _nan_max(a): return float(np.nanmax(a)) if np.any(np.isfinite(a)) else float('nan') sum_ai = float(np.nansum(data_ai)) + std_ai = float(np.nanstd(data_ai)) if np.any(np.isfinite(data_ai)) else float('nan') min_ai = _nan_min(data_ai) max_ai = _nan_max(data_ai) sum_model = float(np.nansum(data_model)) + std_model = float(np.nanstd(data_model)) if np.any(np.isfinite(data_model)) else float('nan') min_model = _nan_min(data_model) max_model = _nan_max(data_model) @@ -208,15 +212,17 @@ def _nan_max(a): r2 = float('nan') print(f"Stats for {var}{label_suffix}:") - print(f" {label_ai}: sum={sum_ai:.6g} min={min_ai:.6g} max={max_ai:.6g}") - print(f" {label_model}: sum={sum_model:.6g} min={min_model:.6g} max={max_model:.6g}") + 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}") stats = { "ai_sum": sum_ai, + "ai_std": std_ai, "ai_min": min_ai, "ai_max": max_ai, "model_sum": sum_model, + "model_std": std_model, "model_min": min_model, "model_max": max_model, "n": n, @@ -394,8 +400,18 @@ def main(): help='Path to write CSV of statistics (defaults to output dir stats.txt)') parser.add_argument('--stats-format', type=str, choices=['csv', 'txt', 'both'], default='txt', help='Format of statistics output: csv, txt, or both [default: txt]') + parser.add_argument('--stats-only', action='store_true', + help='Only compute statistics and write CSV (sum/std/min/max per variable and layer/PFT); saves into output_dir/stats') args = parser.parse_args() + # If stats-only is requested, force no-plot and CSV output into a dedicated stats folder + if getattr(args, 'stats_only', False): + args.no_plot = True + # Route outputs to a stats subfolder for cleaner organization + # The final output_dir will be resolved below after potential variable-list handling + _stats_only_requested = True + else: + _stats_only_requested = False # Validate input files if not Path(args.ai_predictions).exists(): @@ -422,6 +438,9 @@ def main(): requested_all = False if args.variables: requested_all = (len(args.variables) == 1 and str(args.variables[0]).lower() == 'all') + # In stats-only mode, if a variable list is provided, treat as "all" variables from the list + if getattr(args, 'stats_only', False) and args.variable_list: + requested_all = True if requested_all: if args.variable_list: @@ -499,6 +518,11 @@ def main(): output_dir = Path(args.output_dir) os.makedirs(output_dir, exist_ok=True) + # If stats-only, place outputs under a dedicated stats subdirectory + if _stats_only_requested: + output_dir = output_dir / 'stats' + os.makedirs(output_dir, exist_ok=True) + # Update the output directory for the plotting function args.output_dir = str(output_dir) @@ -691,25 +715,56 @@ def main(): # Write statistics outputs (CSV/TXT) if stats_rows: # Resolve output paths - csv_out = args.stats_file if (args.stats_file and args.stats_file.lower().endswith('.csv')) else os.path.join(args.output_dir, "stats.csv") - txt_out = args.stats_file if (args.stats_file and args.stats_file.lower().endswith('.txt')) else os.path.join(args.output_dir, "stats.txt") + if _stats_only_requested: + # In stats-only mode, write a concise CSV with the requested metrics under stats folder + csv_out = os.path.join(args.output_dir, "summary_stats.csv") + txt_out = None + output_mode = 'csv' + else: + csv_out = args.stats_file if (args.stats_file and args.stats_file.lower().endswith('.csv')) else os.path.join(args.output_dir, "stats.csv") + txt_out = args.stats_file if (args.stats_file and args.stats_file.lower().endswith('.txt')) else os.path.join(args.output_dir, "stats.txt") + output_mode = args.stats_format os.makedirs(args.output_dir, exist_ok=True) # CSV output - if args.stats_format in ('csv', 'both'): - fieldnames = [ - "variable", "suffix", "ai_sum", "ai_min", "ai_max", - "model_sum", "model_min", "model_max", "n", "rmse", "nrmse", "r2" - ] + if (output_mode in ('csv', 'both')): + if _stats_only_requested: + fieldnames = [ + "variable", "suffix", + "ai_sum", "ai_std", "ai_min", "ai_max", + "model_sum", "model_std", "model_min", "model_max" + ] + # Reduce rows to requested columns only + filtered_rows = [] + for row in stats_rows: + filtered_rows.append({ + "variable": row.get("variable"), + "suffix": row.get("suffix"), + "ai_sum": row.get("ai_sum"), + "ai_std": row.get("ai_std"), + "ai_min": row.get("ai_min"), + "ai_max": row.get("ai_max"), + "model_sum": row.get("model_sum"), + "model_std": row.get("model_std"), + "model_min": row.get("model_min"), + "model_max": row.get("model_max"), + }) + rows_to_write = filtered_rows + else: + fieldnames = [ + "variable", "suffix", "ai_sum", "ai_min", "ai_max", + "model_sum", "model_min", "model_max", "n", "rmse", "nrmse", "r2" + ] + rows_to_write = stats_rows with open(csv_out, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() - for row in stats_rows: + for row in rows_to_write: writer.writerow(row) print(f"Saved statistics CSV: {csv_out}") # TXT output (readable, grouped by variable and suffix) - if args.stats_format in ('txt', 'both'): + if (not _stats_only_requested) and (output_mode in ('txt', 'both')): # Group stats by variable then suffix from collections import defaultdict grouped = defaultdict(list) diff --git a/scripts/ai_restart_comparison.py b/scripts/ai_restart_comparison.py index 96ac321..3e4e248 100644 --- a/scripts/ai_restart_comparison.py +++ b/scripts/ai_restart_comparison.py @@ -17,8 +17,8 @@ from config.training_config import parse_cnp_io_list # Default file paths -DATA_DIR = '/mnt/proj-shared/AI4BGC_7xw/AI4BGC/ELM_data/' -DEFAULT_FILE_OLD = DATA_DIR + 'original_780_spinup_from_modelsimulation.nc' +DATA_DIR = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/' +DEFAULT_FILE_OLD = DATA_DIR + '20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' LABEL_NEW = "AI Generated" LABEL_OLD = "Original Model" @@ -243,6 +243,10 @@ def parse_arguments(): help='Comma-separated list of PFTs to plot (default: all PFT0-PFT15)') parser.add_argument('--plot-all', action='store_true', help='If set, plot all variables for all 10 layers (0-9) and all 16 PFTs (1-16, skip pft0)') + parser.add_argument('--stats-only', action='store_true', + help='Only compute statistics (sum/std/min/max) for all variables and all layers/PFTs; no plots') + parser.add_argument('--output-dir', type=str, default=OUTPUT_DIR, + help=f'Output directory for plots or stats (default: {OUTPUT_DIR})') return parser.parse_args() def find_ai_restart_file(): @@ -260,7 +264,13 @@ def find_ai_restart_file(): def main(): args = parse_arguments() - os.makedirs(OUTPUT_DIR, exist_ok=True) + # Resolve output directory (stats-only goes to a stats subfolder) + out_dir_base = Path(args.output_dir) + if args.stats_only: + output_dir = out_dir_base / 'stats' + else: + output_dir = out_dir_base + os.makedirs(output_dir, exist_ok=True) # Parse layers and PFTs global LEVGRND_LAYERS, PFT_PICK_LIST if args.plot_all: @@ -327,7 +337,8 @@ def main(): 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]}") - print(f"\nStart plotting: {len(VARIABLES)} variables") + stats_rows = [] + print(f"\nStart {'statistics' if args.stats_only else 'plotting'}: {len(VARIABLES)} variables") for var in VARIABLES: if (var not in ds_new.data_vars) or (var not in ds_old.data_vars): print(f"Skip {var} (not found in both files)") @@ -345,7 +356,13 @@ def main(): vals_new = _to_nan_fillvalue(da_new_cl.values) vals_old = _to_nan_fillvalue(da_old_cl.values) - for lev in LEVGRND_LAYERS: + # Iterate all layers if stats-only; otherwise iterate requested layers + if args.stats_only: + lev_iter = range(int(da_new_cl.sizes["levgrnd"])) + else: + lev_iter = LEVGRND_LAYERS + + for lev in lev_iter: if lev < 0 or lev >= da_new_cl.sizes["levgrnd"]: print(f" Layer {lev} out of range, skipped") continue @@ -366,12 +383,26 @@ def main(): # Always use mapping for old/model file old_grid[g] = vals_old[c0, lev] - # Debug print for this layer - print(f"[DEBUG] {var} lev{lev}: new_grid min={np.nanmin(new_grid)}, max={np.nanmax(new_grid)}, mean={np.nanmean(new_grid)}, sample={new_grid[:10]}") - print(f"[DEBUG] {var} lev{lev}: old_grid min={np.nanmin(old_grid)}, max={np.nanmax(old_grid)}, mean={np.nanmean(old_grid)}, sample={old_grid[:10]}") - - _plot_tripanel(var, f"_lev{lev}", grid_lon, grid_lat, new_grid, old_grid, OUTPUT_DIR, - label_new=LABEL_NEW, label_old=LABEL_OLD) + if args.stats_only: + # Compute stats only + new_sum = float(np.nansum(new_grid)) + new_std = float(np.nanstd(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') + new_min = float(np.nanmin(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') + new_max = float(np.nanmax(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') + old_sum = float(np.nansum(old_grid)) + old_std = float(np.nanstd(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') + old_min = float(np.nanmin(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') + old_max = float(np.nanmax(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') + stats_rows.append({ + "variable": var, + "suffix": f"_lev{lev}", + "new_sum": new_sum, "new_std": new_std, "new_min": new_min, "new_max": new_max, + "old_sum": old_sum, "old_std": old_std, "old_min": old_min, "old_max": old_max, + }) + else: + # Plotting mode + _plot_tripanel(var, f"_lev{lev}", grid_lon, grid_lat, new_grid, old_grid, str(out_dir_base), + label_new=LABEL_NEW, label_old=LABEL_OLD) elif ("pft" in dims) and (len(dims) == 1): # PFT1D variable @@ -380,7 +411,15 @@ def main(): vals_new = _to_nan_fillvalue(da_new_p.values) vals_old = _to_nan_fillvalue(da_old_p.values) - for k in PFT_PICK_LIST: + # Iterate all PFT TYPES (0..15) if stats-only; otherwise iterate requested PFTs + # Note: da_new_p.sizes["pft"] is the total number of PFT entries across all gridcells (very large). + # For comparison we want PFT type indices 0..15 which are mapped per-gridcell via grid_to_pfts. + if args.stats_only: + pft_iter = range(16) + else: + pft_iter = PFT_PICK_LIST + + for k in pft_iter: if k < 0 or k >= da_new_p.sizes["pft"]: print(f" PFT {k} out of range, skipped") continue @@ -398,28 +437,66 @@ def main(): new_grid[g] = vals_new[p_idx] old_grid[g] = vals_old[p_idx] - _plot_tripanel(var, f"_pft{k}", grid_lon, grid_lat, new_grid, old_grid, OUTPUT_DIR, - label_new=LABEL_NEW, label_old=LABEL_OLD) + if args.stats_only: + new_sum = float(np.nansum(new_grid)) + new_std = float(np.nanstd(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') + new_min = float(np.nanmin(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') + new_max = float(np.nanmax(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') + old_sum = float(np.nansum(old_grid)) + old_std = float(np.nanstd(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') + old_min = float(np.nanmin(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') + old_max = float(np.nanmax(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') + stats_rows.append({ + "variable": var, + "suffix": f"_pft{k}", + "new_sum": new_sum, "new_std": new_std, "new_min": new_min, "new_max": new_max, + "old_sum": old_sum, "old_std": old_std, "old_min": old_min, "old_max": old_max, + }) + else: + _plot_tripanel(var, f"_pft{k}", grid_lon, grid_lat, new_grid, old_grid, str(out_dir_base), + label_new=LABEL_NEW, label_old=LABEL_OLD) else: print(f" Skip {var} (only supports (column, levgrnd) and (pft,))") + # Write stats CSV if stats-only + if args.stats_only and stats_rows: + import csv + csv_path = output_dir / 'summary_stats.csv' + fieldnames = [ + 'variable', 'suffix', + 'new_sum', 'new_std', 'new_min', 'new_max', + 'old_sum', 'old_std', 'old_min', 'old_max' + ] + with open(csv_path, 'w', newline='') as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for row in stats_rows: + writer.writerow(row) + print(f"Saved statistics CSV: {csv_path}") + ds_new.close() ds_old.close() - - # Print comparison summary + + # Print summary print("\n" + "="*80) print("AI RESTART COMPARISON SUMMARY:") print("="*80) print(f"Original restart: {os.path.abspath(FILE_OLD)}") print(f"AI-enhanced restart: {os.path.abspath(FILE_NEW)}") print(f"Labels: {LABEL_OLD} vs {LABEL_NEW}") - print(f"Output directory: {os.path.abspath(OUTPUT_DIR)}") - print(f"Variables plotted: {VARIABLES}") - print(f"Layers plotted: {LEVGRND_LAYERS}") - print(f"PFTs plotted: {PFT_PICK_LIST}") + print(f"Output directory: {os.path.abspath(str(output_dir))}") + print(f"Variables processed: {VARIABLES}") + if args.stats_only: + print("Mode: stats-only (all layers and all PFTs)") + else: + print(f"Layers plotted: {LEVGRND_LAYERS}") + print(f"PFTs plotted: {PFT_PICK_LIST}") print("="*80) - print("\nAll plots done! Output dir:", OUTPUT_DIR) + if args.stats_only: + print("\nCompleted without plotting. Stats CSV saved to:", str(output_dir)) + else: + print("\nAll plots done! Output dir:", str(out_dir_base)) if __name__ == "__main__": main() From 24c605d2bad5f6034001dd5664449923e5832a1e Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Mon, 20 Oct 2025 06:08:47 -0700 Subject: [PATCH 27/51] Restore CNP_IO_updated9_dev.txt --- CNP_IO_updated9_dev.txt | 68 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 CNP_IO_updated9_dev.txt diff --git a/CNP_IO_updated9_dev.txt b/CNP_IO_updated9_dev.txt new file mode 100644 index 0000000..43aa366 --- /dev/null +++ b/CNP_IO_updated9_dev.txt @@ -0,0 +1,68 @@ +TRENDY1_PATH: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree +TRENDY05_PATH = +DATA_PATHS: +FILE_PATTERN: enhanced_1_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, NPP, AR, HR + +1D PFT VARIABLES (41 variables): + +• deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage +• deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage + +• leafc, leafc_storage, leafn, leafn_storage, leafp, leafp_storage +• frootc, frootc_storage, frootn, frootn_storage, frootp, frootp_storage + +• livestemc, livestemc_storage, livestemn, livestemn_storage, livestemp, livestemp_storage, +• livecrootc, livecrootc_storage, livecrootn, livecrootn_storage, livecrootp, livecrootp_storage + +• cpool, npool, ppool + +• tlai, totvegc + +2D VARIABLES (layered - 25 variables): + +• cwdc_vr, cwdn_vr, cwdp_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 +• soil3c_vr, soil3n_vr, soil3p_vr +• soil4c_vr, soil4n_vr, soil4p_vr + +• labilep_vr , occlp_vr, primp_vr, secondp_vr From c099493bed4effd0d4b9a62ec40b4818b977e17d Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Thu, 23 Oct 2025 08:36:56 -0700 Subject: [PATCH 28/51] Update training config, pipeline runbook, and comparison scripts --- CNP_IO_updated9_dev.txt | 6 + config/training_config.py | 19 +- docs/CNP_pipeline_runbook.md | 4 +- scripts/ai_model_comparison_plot.py | 1068 ++++++++++++++++++-------- scripts/ai_predictions_to_restart.py | 71 +- scripts/ai_restart_comparison.py | 583 ++++++++------ 6 files changed, 1175 insertions(+), 576 deletions(-) diff --git a/CNP_IO_updated9_dev.txt b/CNP_IO_updated9_dev.txt index 43aa366..28ec192 100644 --- a/CNP_IO_updated9_dev.txt +++ b/CNP_IO_updated9_dev.txt @@ -1,3 +1,9 @@ +AI_PREDICTIONS_DEFAULT: ./comparison_results/ai_predictions_for_plotting.nc +MODEL_DEFAULT: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc +COMPARISON_OUTPUT_DIR: ./ai_model_comparison_plots +CSV_PREDICTIONS_DEFAULT: ./cnp_inference_entire_dataset/cnp_predictions +AI_RESTART_DEFAULT: ./updated_restart_CNP_IO_updated9_dev_20250408_trendytest_ICB1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc + TRENDY1_PATH: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree TRENDY05_PATH = DATA_PATHS: diff --git a/config/training_config.py b/config/training_config.py index 6bdc4ee..e58ca46 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -441,7 +441,12 @@ def parse_cnp_io_list(filename): 'file_pattern': None, 'trendy1_file_pattern': None, 'trendy05_file_pattern': None, - 'tva4km_file_pattern': None + 'tva4km_file_pattern': None, + 'ai_predictions_default': None, + 'model_default': None, + 'comparison_output_dir': None, + 'csv_predictions_default': None, + 'ai_restart_default': None }) current_section = None @@ -494,7 +499,7 @@ def parse_cnp_io_list(filename): # FILE_PATTERN: enhanced_1_training_data_batch_*.pkl # DATA_PATHS: /p1,/p2 if line and not line.startswith('#'): - kv_match = re.match(r'(?i)^(trendy1_path|trendy05_path|tva4km_path|file_pattern|trendy1_file_pattern|trendy05_file_pattern|tva4km_file_pattern|data_paths)\s*[:=]\s*(.+)$', line) + kv_match = re.match(r'(?i)^(trendy1_path|trendy05_path|tva4km_path|file_pattern|trendy1_file_pattern|trendy05_file_pattern|tva4km_file_pattern|data_paths|ai_predictions_default|model_default|comparison_output_dir|csv_predictions_default|ai_restart_default)\s*[:=]\s*(.+)$', line) if kv_match: key = kv_match.group(1).lower() val = kv_match.group(2).strip() @@ -516,6 +521,16 @@ def parse_cnp_io_list(filename): result['trendy05_file_pattern'] = val elif key == 'tva4km_file_pattern': result['tva4km_file_pattern'] = val + elif key == 'ai_predictions_default': + result['ai_predictions_default'] = val + elif key == 'model_default': + result['model_default'] = val + elif key == 'comparison_output_dir': + result['comparison_output_dir'] = val + elif key == 'csv_predictions_default': + result['csv_predictions_default'] = val + elif key == 'ai_restart_default': + result['ai_restart_default'] = val return result def parse_cnp_model_config(filename: str) -> Dict[str, Any]: diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 5600964..1eeedeb 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -72,7 +72,7 @@ Creates map plots for the variables in your CNP_IO list under python ../../scripts/ai_model_comparison_plot.py > ai_model_comparison.log 2>&1 & ``` #### 7.1) Stats-only mode (no plotting) -If you only need summary statistics (sum/std/min/max) for all variables and their layers/PFTs, use `--stats-only` 与 `--variable-list`： +To validate that the variable values from the CSV file are consistent with those in the NetCDF file, you can compare their summary statistics (sum, standard deviation, minimum, and maximum) for all variables and their corresponding layers/PFTs using the --stats-only and --variable-list options: ```bash python ../../scripts/ai_model_comparison_plot.py \ --stats-only \ @@ -96,7 +96,7 @@ python ../../scripts/ai_restart_comparison.py --variable-list ../../CNP_IO_demo1 # Optionally use ../../scripts/restart_variable_plot.py for manual verification ``` #### 9.1) Stats-only mode (no plotting) -If you only need summary statistics (sum/std/min/max) for all variables and their layers/PFTs, use `--stats-only` 与 `--variable-list`： +To validate that the variable values in the predicted NetCDF file are consistent with those in the updated restart file, you can compare their summary statistics (sum, standard deviation, minimum, and maximum) for all variables and their corresponding layers/PFTs using the --stats-only and --variable-list options: ```bash python ../../scripts/ai_restart_comparison.py \ --stats-only \ diff --git a/scripts/ai_model_comparison_plot.py b/scripts/ai_model_comparison_plot.py index f1c6dc4..ae41847 100644 --- a/scripts/ai_model_comparison_plot.py +++ b/scripts/ai_model_comparison_plot.py @@ -1,6 +1,4 @@ #!/usr/bin/env python3 -import ssl -ssl._create_default_https_context = ssl._create_unverified_context import os import numpy as np import xarray as xr @@ -14,15 +12,18 @@ import sys from typing import List import csv - +import ssl +ssl._create_default_https_context = ssl._create_unverified_context +import pandas as pd # Project imports sys.path.append(str(Path(__file__).resolve().parents[1])) from config.training_config import parse_cnp_io_list # Default paths -DEFAULT_AI_PREDICTIONS = './comparison_results/ai_predictions_for_plotting.nc' -DEFAULT_MODEL = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' -DEFAULT_OUTPUT_DIR = "./ai_model_comparison_plots" +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_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'] @@ -34,6 +35,280 @@ # 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.) +CSV_LONGITUDE_NAMES = ("Longitude", "Long", "long", "lon", "LON") +CSV_LATITUDE_NAMES = ("Latitude", "Lat", "lat", "LAT") +CSV_COORD_COLUMNS = CSV_LONGITUDE_NAMES + CSV_LATITUDE_NAMES + + +QUALITY_THRESHOLDS = { + 'good': 0.9, # R² >= 0.9 for good quality + 'ok': 0.7 # R² >= 0.7 for ok quality +} + + +def _extract_coords_from_df(df: pd.DataFrame): + lon = lat = None + for col in CSV_LONGITUDE_NAMES: + if col in df.columns: + lon = pd.to_numeric(df[col], errors='coerce').to_numpy() + break + for col in CSV_LATITUDE_NAMES: + if col in df.columns: + lat = pd.to_numeric(df[col], errors='coerce').to_numpy() + break + return lon, lat + + +def _drop_coord_columns(df: pd.DataFrame) -> pd.DataFrame: + return df.drop(columns=[c for c in CSV_COORD_COLUMNS if c in df.columns], errors='ignore') + + +def _get_case_insensitive(mapping, key): + if mapping is None: + return None + if key in mapping: + return mapping[key] + key_lower = key.lower() + for k, v in mapping.items(): + if k.lower() == key_lower: + return v + return None + + +def _classify_quality_from_r2(value: float) -> str: + """Classify quality based on R² value (similar to generate_prediction_quality_report.py)""" + if value is None or not np.isfinite(value): + return 'unknown' + val = float(value) + if val >= QUALITY_THRESHOLDS['good']: + return 'good' + if val >= QUALITY_THRESHOLDS['ok']: + return 'ok' + return 'bad' + + +def load_csv_predictions(csv_path: str) -> dict: + path = Path(csv_path) + if not path.exists(): + raise FileNotFoundError(f"CSV predictions source not found: {csv_path}") + + store = { + 'source': str(path), + 'scalar_df': None, + 'scalar_columns': {}, + 'scalar_vars': set(), + 'scalar_coords': None, + 'pft_data': {}, + 'pft_columns': {}, + 'pft_vars': set(), + 'pft_coords': {}, + 'soil_data': {}, + 'soil_columns': {}, + 'soil_vars': set(), + 'soil_coords': {}, + 'static_inverse': None, + 'flat_df': None, + 'flat_columns': {}, + 'flat_vars': set(), + 'flat_coords': None, + } + + if path.is_dir(): + scalar_path = path / 'predictions_scalar.csv' + if scalar_path.exists(): + df = pd.read_csv(scalar_path) + lon, lat = _extract_coords_from_df(df) + data = _drop_coord_columns(df) + col_map = {} + scalar_vars = set() + for col in data.columns: + base = col[2:] if col.startswith('Y_') else col + scalar_vars.add(base) + if base not in col_map: + col_map[base] = col + if base.lower() not in col_map: + col_map[base.lower()] = col + store['scalar_df'] = data + store['scalar_columns'] = col_map + store['scalar_vars'] = scalar_vars + store['scalar_coords'] = (lon, lat) + + pft_dir = path / 'pft_1d_predictions' + if pft_dir.exists(): + for csv_file in sorted(pft_dir.glob('predictions_*.csv')): + var_name = csv_file.stem.replace('predictions_Y_', '') + df = pd.read_csv(csv_file) + lon, lat = _extract_coords_from_df(df) + data = _drop_coord_columns(df) + cols = [c for c in data.columns if '_pft' in c] + if not cols: + continue + def _pft_idx(col_name: str) -> int: + try: + return int(col_name.split('_pft')[-1]) + except Exception: + return 999999 + cols = sorted(cols, key=_pft_idx) + store['pft_data'][var_name] = data + store['pft_columns'][var_name] = cols + store['pft_vars'].add(var_name) + store['pft_coords'][var_name] = (lon, lat) + + soil_dir = path / 'soil_2d_predictions' + if soil_dir.exists(): + for csv_file in sorted(soil_dir.glob('predictions_*.csv')): + var_name = csv_file.stem.replace('predictions_Y_', '') + df = pd.read_csv(csv_file) + lon, lat = _extract_coords_from_df(df) + data = _drop_coord_columns(df) + cols = [c for c in data.columns if '_layer' in c] + if not cols: + continue + def _layer_idx(col_name: str) -> int: + try: + return int(col_name.split('_layer')[-1]) + except Exception: + return 999999 + cols = sorted(cols, key=_layer_idx) + store['soil_data'][var_name] = data + store['soil_columns'][var_name] = cols + store['soil_vars'].add(var_name) + store['soil_coords'][var_name] = (lon, lat) + + static_path = path / 'test_static_inverse.csv' + if static_path.exists(): + try: + store['static_inverse'] = pd.read_csv(static_path) + except Exception as exc: + print(f"Warning: failed to load test_static_inverse.csv: {exc}") + + elif path.is_file(): + df = pd.read_csv(path) + lon, lat = _extract_coords_from_df(df) + data = _drop_coord_columns(df) + col_map = {} + flat_vars = set() + for col in data.columns: + base = col[2:] if col.startswith('Y_') else col + flat_vars.add(base) + if base not in col_map: + col_map[base] = col + if base.lower() not in col_map: + col_map[base.lower()] = col + store['flat_df'] = data + store['flat_columns'] = col_map + store['flat_vars'] = flat_vars + store['flat_coords'] = (lon, lat) + else: + raise ValueError(f"Unsupported CSV predictions path: {csv_path}") + + available = set() + for key in ('scalar_vars', 'pft_vars', 'soil_vars', 'flat_vars'): + available.update(store.get(key, set())) + store['available_vars'] = available + print(f"Loaded CSV predictions from {path} with {len(available)} variables") + return store + + +def build_variable_category_map(variable_list_path: str) -> dict: + if not variable_list_path: + return {} + try: + parsed = parse_cnp_io_list(variable_list_path) + except Exception as exc: + print(f"Warning: unable to parse variable list for category mapping: {exc}") + return {} + mapping = {} + for name in parsed.get('scalar_variables', []) or []: + mapping[name] = 'scalar' + for name in parsed.get('pft_1d_variables', []) or []: + mapping[name] = 'pft1d' + for name in parsed.get('variables_2d_soil', []) or []: + mapping[name] = 'soil2d' + return mapping + + +def infer_variable_category(var: str, dims, category_map: dict) -> str: + if category_map and var in category_map: + return category_map[var] + if dims and any(dim in ('levgrnd', 'column') for dim in dims): + return 'soil2d' + if dims and any(dim == 'pft' for dim in dims): + return 'pft1d' + return 'scalar' + + +def csv_has_variable(store: dict, var: str, category: str = None) -> bool: + if store is None: + return False + if category == 'scalar': + if store.get('scalar_df') is not None and _get_case_insensitive(store.get('scalar_columns'), var): + return True + if store.get('flat_df') is not None and _get_case_insensitive(store.get('flat_columns'), var): + return True + return False + if category == 'pft1d': + return _get_case_insensitive(store.get('pft_data'), var) is not None + if category == 'soil2d': + return _get_case_insensitive(store.get('soil_data'), var) is not None + return (csv_has_variable(store, var, 'scalar') or + csv_has_variable(store, var, 'pft1d') or + csv_has_variable(store, var, 'soil2d')) + + +def extract_csv_scalar(store: dict, var: str): + df = store.get('scalar_df') + columns = store.get('scalar_columns') + col = _get_case_insensitive(columns, var) if columns else None + if df is None or col is None: + df = store.get('flat_df') + columns = store.get('flat_columns') + col = _get_case_insensitive(columns, var) if columns else None + if df is None or col is None: + col = _get_case_insensitive(columns, f'Y_{var}') if columns else None + if df is None or col is None: + return None + return pd.to_numeric(df[col], errors='coerce').to_numpy() + + +def extract_csv_pft(store: dict, var: str): + data_dict = store.get('pft_data') + df = _get_case_insensitive(data_dict, var) + if df is None: + return None + columns_map = store.get('pft_columns') + cols = _get_case_insensitive(columns_map, var) + if not cols: + cols = [c for c in df.columns if '_pft' in c] + if not cols: + return None + def _pft_idx(col_name: str) -> int: + try: + return int(col_name.split('_pft')[-1]) + except Exception: + return 999999 + cols = sorted(cols, key=_pft_idx) + arr = np.full((len(cols), len(df)), np.nan, dtype=float) + for idx, col in enumerate(cols): + arr[idx, :] = pd.to_numeric(df[col], errors='coerce').to_numpy() + return arr + + +def extract_csv_soil(store: dict, var: str): + data_dict = store.get('soil_data') + df = _get_case_insensitive(data_dict, var) + if df is None: + return None + columns_map = store.get('soil_columns') + cols = _get_case_insensitive(columns_map, var) + if not cols: + cols = list(df.columns) + arr = np.full((len(cols), len(df)), np.nan, dtype=float) + for idx, col in enumerate(cols): + arr[idx, :] = pd.to_numeric(df[col], errors='coerce').to_numpy() + return arr + + def _safe_get(ds, name): """Safely get a variable from dataset, with error handling.""" if name not in ds: @@ -304,6 +579,21 @@ def _nan_max(a): return stats +def _load_default_paths(variable_list_path: str): + defaults = {} + if variable_list_path: + try: + vl_path = Path(variable_list_path) + if vl_path.exists(): + parsed = parse_cnp_io_list(variable_list_path) + for key in ('ai_predictions_default', 'model_default', 'comparison_output_dir', 'csv_predictions_default'): + value = parsed.get(key) if isinstance(parsed, dict) else None + if value: + defaults[key] = value + except Exception as exc: + print(f"Warning: Failed to load default paths from {variable_list_path}: {exc}") + return defaults + def parse_variable_list_file(variable_list_path: str) -> List[str]: """Parse the CNP IO list file to extract all variables.""" print(f"Parsing variable list file: {variable_list_path}") @@ -380,12 +670,12 @@ def main(): """ ) - parser.add_argument('--ai-predictions', default=DEFAULT_AI_PREDICTIONS, - help=f'Path to AI predictions NetCDF file [default: {DEFAULT_AI_PREDICTIONS}]') - parser.add_argument('--model', default=DEFAULT_MODEL, - help=f'Path to model results NetCDF file [default: {DEFAULT_MODEL}]') - parser.add_argument('--output-dir', default=DEFAULT_OUTPUT_DIR, - help=f'Output directory for plots [default: {DEFAULT_OUTPUT_DIR}]') + parser.add_argument('--ai-predictions', default=None, + 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('--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, help='Path to CNP_IO_list file to extract all variables for plotting') parser.add_argument('--variables', nargs='*', default=VARIABLES, @@ -400,30 +690,51 @@ def main(): help='Path to write CSV of statistics (defaults to output dir stats.txt)') parser.add_argument('--stats-format', type=str, choices=['csv', 'txt', 'both'], default='txt', help='Format of statistics output: csv, txt, or both [default: txt]') + parser.add_argument('--csv-predictions', type=str, default=None, + help='Path to CSV predictions (directory or file) containing original AI outputs for comparison') parser.add_argument('--stats-only', action='store_true', - help='Only compute statistics and write CSV (sum/std/min/max per variable and layer/PFT); saves into output_dir/stats') - + help='Only compute statistics; disables plotting and processes all layers/PFTs when used with --variable-list') + args = parser.parse_args() - # If stats-only is requested, force no-plot and CSV output into a dedicated stats folder - if getattr(args, 'stats_only', False): + + defaults_from_config = _load_default_paths(args.variable_list) + + def _resolve_default(current_value, config_key, fallback): + candidate = current_value or defaults_from_config.get(config_key) + if candidate: + 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 - # Route outputs to a stats subfolder for cleaner organization - # The final output_dir will be resolved below after potential variable-list handling - _stats_only_requested = True - else: - _stats_only_requested = False - + # Validate input files if not Path(args.ai_predictions).exists(): raise FileNotFoundError(f"AI predictions file not found: {args.ai_predictions}") - if not Path(args.model).exists(): - raise FileNotFoundError(f"Model file not found: {args.model}") - + 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) - print("AI vs Model Comparison") + if use_csv_predictions: + print("CSV vs NetCDF Comparison") + else: + print("AI vs Model Comparison") print("="*60) - print(f"AI predictions: {args.ai_predictions}") - 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}") @@ -432,81 +743,117 @@ def main(): # Open datasets once ds_ai = xr.open_dataset(args.ai_predictions) - ds_model = xr.open_dataset(args.model) + csv_predictions = None + if use_csv_predictions: + csv_predictions = load_csv_predictions(args.csv_predictions) + ds_model = None + else: + ds_model = xr.open_dataset(args.model) # Determine variable selection behavior requested_all = False if args.variables: requested_all = (len(args.variables) == 1 and str(args.variables[0]).lower() == 'all') - # In stats-only mode, if a variable list is provided, treat as "all" variables from the list - if getattr(args, 'stats_only', False) and args.variable_list: + if args.stats_only and args.variable_list: 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 requested_all: if args.variable_list: if not Path(args.variable_list).exists(): - ds_ai.close(); ds_model.close() + ds_ai.close() + if ds_model is not None: + ds_model.close() raise FileNotFoundError(f"Variable list file not found: {args.variable_list}") - # Parse the variable list file to get all variables all_variables = parse_variable_list_file(args.variable_list) - # Filter to only include variables that exist in both datasets - ai_vars = set(ds_ai.data_vars.keys()) - model_vars = set(ds_model.data_vars.keys()) - available_vars = [var for var in all_variables if var in ai_vars and var in model_vars] + if 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)) + ] + else: + model_vars = set(ds_model.data_vars.keys()) + available_vars = [var for var in all_variables if var in ai_vars and var in model_vars] if available_vars: args.variables = available_vars print(f"Using all variables from variable list ({len(available_vars)}): {available_vars}") else: - print("Warning: No variables from variable list found in both datasets!") - ds_ai.close(); ds_model.close() + print("Warning: No variables from variable list found in available datasets!") + ds_ai.close() + if ds_model is not None: + ds_model.close() return else: - # Discover all common variables across datasets - discovered_vars = discover_common_variables(ds_ai, ds_model) - if discovered_vars: - args.variables = discovered_vars - print(f"Using all common variables ({len(discovered_vars)}): {discovered_vars}") + if use_csv_predictions: + csv_vars = set(csv_predictions.get('available_vars', set())) + discovered_vars = sorted(ai_vars.intersection(csv_vars)) + if discovered_vars: + args.variables = discovered_vars + 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() + return else: - print("Warning: No common variables found between AI predictions and model!") - ds_ai.close(); ds_model.close() - return + discovered_vars = discover_common_variables(ds_ai, ds_model) + if discovered_vars: + args.variables = discovered_vars + 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() + return else: - # Force the script to only plot the default subset unless 'all' is requested forced = ['cwdc_vr', 'soil3c_vr', 'tlai', 'deadstemc'] - ai_vars = set(ds_ai.data_vars.keys()) - model_vars = set(ds_model.data_vars.keys()) - 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() - 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 datasets: {missing}") + if 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() + return + missing = [v for v in forced if v not in selected] + if missing: + print(f"Note: Skipping default variables missing from CSV source: {missing}") + else: + model_vars = set(ds_model.data_vars.keys()) + 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() + 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 datasets: {missing}") args.variables = selected print(f"Using default subset of variables ({len(selected)}): {selected}") - - # Get grid information from the MODEL file as the master coordinate system - # This ensures AI predictions can be properly ingested into the model - grid_lon, grid_lat = _gridcell_lonlat(ds_model) - n_grid = ds_model.sizes["gridcell"] - - print(f"Using MODEL gridcell count: {n_grid}") - print(f"Model coordinates: lon range [{grid_lon.min():.3f}, {grid_lon.max():.3f}], lat range [{grid_lat.min():.3f}, {grid_lat.max():.3f}]") - - # Build mappings from the model file for extracting model data - col2grid = _to_zero_based_index(_safe_get(ds_model, "cols1d_gridcell_index").values, n_grid) - pft2grid = _to_zero_based_index(_safe_get(ds_model, "pfts1d_gridcell_index").values, n_grid) - - grid_to_cols = _build_gridcell_groups(col2grid, n_grid) - grid_to_pfts = _build_gridcell_groups(pft2grid, n_grid) - - # Create spatial mapping from AI gridcells to model gridcells - ai_to_model_mapping = _create_ai_to_model_mapping(ds_ai, ds_model, n_grid) - - 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]}") - + + if 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}") + print(f"AI coordinates: lon range [{grid_lon.min():.3f}, {grid_lon.max():.3f}], lat range [{grid_lat.min():.3f}, {grid_lat.max():.3f}]") + grid_to_cols = None + grid_to_pfts = None + ai_to_model_mapping = None + else: + grid_lon, grid_lat = _gridcell_lonlat(ds_model) + n_grid = ds_model.sizes['gridcell'] + print(f"Using MODEL gridcell count: {n_grid}") + print(f"Model 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_model, 'cols1d_gridcell_index').values, n_grid) + pft2grid = _to_zero_based_index(_safe_get(ds_model, '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 = _create_ai_to_model_mapping(ds_ai, ds_model, n_grid) + 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]}") + # Create output directory if args.variable_list and requested_all: # Use a more descriptive output directory name when using variable list @@ -518,303 +865,420 @@ def main(): output_dir = Path(args.output_dir) os.makedirs(output_dir, exist_ok=True) - # If stats-only, place outputs under a dedicated stats subdirectory - if _stats_only_requested: - output_dir = output_dir / 'stats' - os.makedirs(output_dir, exist_ok=True) - # Update the output directory for the plotting function args.output_dir = str(output_dir) print(f"\nStart processing: {len(args.variables)} variables") stats_rows = [] for var in args.variables: - if (var not in ds_ai.data_vars) or (var not in ds_model.data_vars): - print(f"Skip {var} (not found in both files)") + if var not in ds_ai.data_vars: + print(f"Skip {var} (not found in AI NetCDF)") continue da_ai = ds_ai[var] - da_model = ds_model[var] dims = da_ai.dims + category = infer_variable_category(var, dims, variable_category_map) print(f"\nVariable {var}, dims: {dims}") - print(f" AI shape: {da_ai.shape}") - print(f" Model shape: {da_model.shape}") - - if ("column" in dims) and ("levgrnd" in dims): - # Handle column-type variables (e.g., soil variables) - # Handle different dimension orders - if len(dims) == 3: - # If we have (column, levgrnd, gridcell) or similar, transpose to (column, levgrnd) - if "gridcell" in dims: - da_ai_cl = da_ai.transpose("column", "levgrnd", ...) - da_model_cl = da_model.transpose("column", "levgrnd", ...) - else: - da_ai_cl = da_ai.transpose("column", "levgrnd") - da_model_cl = da_model.transpose("column", "levgrnd") - else: - da_ai_cl = da_ai.transpose("column", "levgrnd") - da_model_cl = da_model.transpose("column", "levgrnd") - - vals_ai = _to_nan_fillvalue(da_ai_cl.values) - vals_model = _to_nan_fillvalue(da_model_cl.values) - - # Debug: Print data structure information - print(f" Column variable: AI shape {vals_ai.shape}, Model shape {vals_model.shape}") - print(f" Grid mapping: {len(grid_to_cols)} gridcells with columns") - print(f" Sample gridcell 0 has columns: {grid_to_cols[0][:5] if grid_to_cols[0] else 'none'}") - - # Compute statistics for all layers; only plot selected layers - total_layers = int(da_ai_cl.sizes["levgrnd"]) if "levgrnd" in da_ai_cl.sizes else 0 - for lev in range(total_layers): - if lev < 0 or lev >= total_layers: + print(f" NetCDF shape: {da_ai.shape}") + + if use_csv_predictions: + if not csv_has_variable(csv_predictions, var, category): + print(f" Skip {var} (not present in CSV source)") + continue + + label_ai = "NetCDF Predictions" + label_csv = "CSV Predictions" + + if category == 'soil2d' and 'levgrnd' in dims: + da_ai_sel = da_ai + if 'column' in da_ai_sel.dims: + da_ai_sel = da_ai_sel.isel(column=0) + try: + da_ai_sel = da_ai_sel.transpose('levgrnd', 'gridcell', ...) + except ValueError: + da_ai_sel = da_ai_sel.transpose(..., 'levgrnd', 'gridcell') + ai_vals = _to_nan_fillvalue(da_ai_sel.values) + if ai_vals.ndim == 1: + ai_vals = ai_vals[np.newaxis, :] + csv_vals = extract_csv_soil(csv_predictions, var) + if csv_vals is None: + print(f" Skip {var} (CSV layers unavailable)") continue + layer_count = ai_vals.shape[0] + csv_layer_count = csv_vals.shape[0] + if csv_vals.shape[1] != ai_vals.shape[1]: + min_len = min(ai_vals.shape[1], csv_vals.shape[1], len(grid_lon)) + if min_len == 0: + print(f" Skip {var} (no grid overlap between NetCDF and CSV)") + continue + print(f" Warning: grid mismatch for {var}; trimming to {min_len} cells") + else: + min_len = ai_vals.shape[1] + lon_subset = grid_lon[:min_len] + lat_subset = grid_lat[:min_len] + for lev in range(layer_count): + ai_layer = ai_vals[lev, :min_len] + if lev < csv_layer_count: + csv_layer = csv_vals[lev, :min_len] + else: + csv_layer = np.full(min_len, np.nan, dtype=float) + plot_flag = (not args.no_plot) and (lev in args.layers) + stats = _plot_tripanel(var, f"_lev{lev}", lon_subset, lat_subset, ai_layer, csv_layer, args.output_dir, + label_ai=label_ai, label_model=label_csv, plot=plot_flag) + stats_rows.append({ + 'variable': var, + 'suffix': f"_lev{lev}", + **stats, + }) - ai_grid = np.full(n_grid, np.nan, dtype=float) - model_grid = np.full(n_grid, np.nan, dtype=float) - - for g in range(n_grid): - cols = grid_to_cols[g] - if len(cols) == 0: + elif category == 'pft1d' or ('pft' in dims): + try: + da_ai_p = da_ai.transpose('pft', 'gridcell', ...) + except ValueError: + da_ai_p = da_ai.transpose(..., 'pft', 'gridcell') + ai_vals = _to_nan_fillvalue(da_ai_p.values) + if ai_vals.ndim == 1: + ai_vals = ai_vals[np.newaxis, :] + csv_vals = extract_csv_pft(csv_predictions, var) + if csv_vals is None: + print(f" Skip {var} (CSV PFT entries unavailable)") + continue + pft_count = ai_vals.shape[0] + csv_pft_count = csv_vals.shape[0] + if csv_vals.shape[1] != ai_vals.shape[1]: + min_len = min(ai_vals.shape[1], csv_vals.shape[1], len(grid_lon)) + if min_len == 0: + print(f" Skip {var} (no grid overlap between NetCDF and CSV)") continue - - # For AI data: map from model gridcell g to corresponding AI gridcell - ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] - if len(ai_gridcell_idx) > 0: - ai_gridcell_idx = ai_gridcell_idx[0] - ai_col_idx = 0 - # Handle AI data indexing - shape is (column, levgrnd, gridcell) - if vals_ai.ndim == 3: - ai_grid[g] = vals_ai[ai_col_idx, lev, ai_gridcell_idx] - else: - ai_grid[g] = vals_ai[ai_col_idx, lev] + print(f" Warning: grid mismatch for {var}; trimming to {min_len} cells") + else: + min_len = ai_vals.shape[1] + lon_subset = grid_lon[:min_len] + lat_subset = grid_lat[:min_len] + for k in range(pft_count): + ai_slice = ai_vals[k, :min_len] + if k < csv_pft_count: + csv_slice = csv_vals[k, :min_len] else: - ai_grid[g] = np.nan - - # For model: use the first column of this gridcell - if g < len(grid_to_cols) and len(grid_to_cols[g]) > 0: - model_col_idx = grid_to_cols[g][0] - if model_col_idx < vals_model.shape[0]: - if vals_model.ndim == 2: - if lev < vals_model.shape[1]: - model_grid[g] = vals_model[model_col_idx, lev] - else: - model_grid[g] = vals_model[model_col_idx] + 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, + label_ai=label_ai, label_model=label_csv, plot=plot_flag) + stats_rows.append({ + 'variable': var, + 'suffix': f"_pft{k+1}", + **stats, + }) - # Plot only if requested layer in args.layers and plotting enabled - do_plot = (not args.no_plot) and (lev in args.layers) - stats = _plot_tripanel(var, f"_lev{lev}", grid_lon, grid_lat, ai_grid, model_grid, args.output_dir, - label_ai="AI Predictions", label_model="Model Results", plot=do_plot) + elif 'gridcell' in dims: + try: + da_ai_gc = da_ai.transpose(..., 'gridcell') + except ValueError: + da_ai_gc = da_ai + ai_vals = _to_nan_fillvalue(np.asarray(da_ai_gc.values)) + ai_vals = np.reshape(ai_vals, (-1, ai_vals.shape[-1])) if ai_vals.ndim > 1 else ai_vals + ai_vals = ai_vals[-1] if ai_vals.ndim > 1 else ai_vals + csv_vals = extract_csv_scalar(csv_predictions, var) + if csv_vals is None: + print(f" Skip {var} (CSV scalar column unavailable)") + continue + csv_vals = np.asarray(csv_vals, dtype=float) + min_len = min(len(ai_vals), len(csv_vals), len(grid_lon)) + if min_len == 0: + print(f" Skip {var} (no overlapping records)") + continue + if len(ai_vals) != len(csv_vals): + print(f" Warning: record mismatch for {var}; trimming to {min_len}") + lon_subset = grid_lon[:min_len] + lat_subset = grid_lat[:min_len] + ai_trim = ai_vals[:min_len] + csv_trim = csv_vals[:min_len] + stats = _plot_tripanel(var, '', lon_subset, lat_subset, ai_trim, csv_trim, args.output_dir, + label_ai=label_ai, label_model=label_csv, plot=(not args.no_plot)) stats_rows.append({ - "variable": var, - "suffix": f"_lev{lev}", + 'variable': var, + 'suffix': '', **stats, }) - - elif ("pft" in dims): - # Handle PFT-type variables - # Handle different dimension orders - if len(dims) == 2 and "gridcell" in dims: - # If we have (pft, gridcell), transpose to (pft, ...) - da_ai_p = da_ai.transpose("pft", ...) - da_model_p = da_model.transpose("pft", ...) else: - da_ai_p = da_ai.transpose("pft") - da_model_p = da_model.transpose("pft") - - vals_ai = _to_nan_fillvalue(da_ai_p.values) - vals_model = _to_nan_fillvalue(da_model_p.values) - - # Compute statistics for all PFTs in AI data; only plot selected ones - total_pfts = vals_ai.shape[0] - for k in range(total_pfts): + print(f" Skip {var} (unsupported dimensions for CSV comparison: {dims})") + + else: + if var not in ds_model.data_vars: + print(f"Skip {var} (not found in model file)") + continue + + da_model = ds_model[var] + print(f" Model shape: {da_model.shape}") + + if ('column' in dims) and ('levgrnd' in dims): + if len(dims) == 3: + if 'gridcell' in dims: + da_ai_cl = da_ai.transpose('column', 'levgrnd', ...) + da_model_cl = da_model.transpose('column', 'levgrnd', ...) + else: + da_ai_cl = da_ai.transpose('column', 'levgrnd') + da_model_cl = da_model.transpose('column', 'levgrnd') + else: + da_ai_cl = da_ai.transpose('column', 'levgrnd') + da_model_cl = da_model.transpose('column', 'levgrnd') + + vals_ai = _to_nan_fillvalue(da_ai_cl.values) + vals_model = _to_nan_fillvalue(da_model_cl.values) + + print(f" Column variable: AI shape {vals_ai.shape}, Model shape {vals_model.shape}") + print(f" Grid mapping: {len(grid_to_cols)} gridcells with columns") + print(f" Sample gridcell 0 has columns: {grid_to_cols[0][:5] if grid_to_cols[0] else 'none'}") + + total_layers = int(da_ai_cl.sizes['levgrnd']) if 'levgrnd' in da_ai_cl.sizes else 0 + for lev in range(total_layers): + if lev < 0 or lev >= total_layers: + continue + + ai_grid = np.full(n_grid, np.nan, dtype=float) + model_grid = np.full(n_grid, np.nan, dtype=float) + + for g in range(n_grid): + cols = grid_to_cols[g] + if len(cols) == 0: + continue + ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] + if len(ai_gridcell_idx) > 0: + ai_gridcell_idx = ai_gridcell_idx[0] + ai_col_idx = 0 + if vals_ai.ndim == 3: + ai_grid[g] = vals_ai[ai_col_idx, lev, ai_gridcell_idx] + else: + ai_grid[g] = vals_ai[ai_col_idx, lev] + else: + ai_grid[g] = np.nan + + if g < len(grid_to_cols) and len(grid_to_cols[g]) > 0: + model_col_idx = grid_to_cols[g][0] + if model_col_idx < vals_model.shape[0]: + if vals_model.ndim == 2: + if lev < vals_model.shape[1]: + model_grid[g] = vals_model[model_col_idx, lev] + else: + model_grid[g] = vals_model[model_col_idx] + + do_plot = (not args.no_plot) and (lev in args.layers) + stats = _plot_tripanel(var, f"_lev{lev}", 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"_lev{lev}", + **stats, + }) + + elif ('pft' in dims): + if len(dims) == 2 and 'gridcell' in dims: + da_ai_p = da_ai.transpose('pft', ...) + da_model_p = da_model.transpose('pft', ...) + else: + da_ai_p = da_ai.transpose('pft') + da_model_p = da_model.transpose('pft') + + vals_ai = _to_nan_fillvalue(da_ai_p.values) + vals_model = _to_nan_fillvalue(da_model_p.values) + + total_pfts = vals_ai.shape[0] + for k in range(total_pfts): + ai_grid = np.full(n_grid, np.nan, dtype=float) + model_grid = np.full(n_grid, np.nan, dtype=float) + + if vals_ai.ndim == 2: + for g in range(n_grid): + ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] + if len(ai_gridcell_idx) > 0: + ai_gridcell_idx = ai_gridcell_idx[0] + ai_grid[g] = vals_ai[k, ai_gridcell_idx] + else: + for g in range(n_grid): + ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] + if len(ai_gridcell_idx) > 0: + ai_grid[g] = vals_ai[k] + + 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 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, + label_ai='AI Predictions', label_model='Model Results', plot=do_plot) + stats_rows.append({ + 'variable': var, + 'suffix': f"_pft{k+1}", + **stats, + }) + + elif 'gridcell' in dims: + if len(dims) == 1: + da_ai_gc = da_ai + da_model_gc = da_model + else: + da_ai_gc = da_ai.transpose(..., 'gridcell') + da_model_gc = da_model.transpose(..., 'gridcell') + + vals_ai = _to_nan_fillvalue(da_ai_gc.values) + vals_model = _to_nan_fillvalue(da_model_gc.values) + ai_grid = np.full(n_grid, np.nan, dtype=float) - model_grid = np.full(n_grid, np.nan, dtype=float) - # Map AI data for PFT k to model gridcells - if vals_ai.ndim == 2: + if vals_ai.ndim == 1: for g in range(n_grid): ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] if len(ai_gridcell_idx) > 0: ai_gridcell_idx = ai_gridcell_idx[0] - ai_grid[g] = vals_ai[k, ai_gridcell_idx] + ai_grid[g] = vals_ai[ai_gridcell_idx] else: for g in range(n_grid): ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] if len(ai_gridcell_idx) > 0: - ai_grid[g] = vals_ai[k] - - # Model data: use first 16 PFTs per gridcell; AI PFT k -> Model PFT (k+1) - 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 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, - label_ai="AI Predictions", label_model="Model Results", plot=do_plot) + ai_gridcell_idx = ai_gridcell_idx[0] + if vals_ai.ndim == 2: + ai_grid[g] = vals_ai[0, ai_gridcell_idx] + else: + ai_grid[g] = vals_ai[ai_gridcell_idx] + + model_grid = vals_model + + stats = _plot_tripanel(var, '', grid_lon, grid_lat, ai_grid, model_grid, args.output_dir, + label_ai='AI Predictions', label_model='Model Results', plot=(not args.no_plot)) stats_rows.append({ - "variable": var, - "suffix": f"_pft{k+1}", + 'variable': var, + 'suffix': '', **stats, }) - elif "gridcell" in dims: - # Handle gridcell-level variables (e.g., GPP, NPP) - # Handle different dimension orders - if len(dims) == 1: - da_ai_gc = da_ai - da_model_gc = da_model else: - # If we have multiple dimensions including gridcell, transpose to put gridcell last - da_ai_gc = da_ai.transpose(..., "gridcell") - da_model_gc = da_model.transpose(..., "gridcell") - - vals_ai = _to_nan_fillvalue(da_ai_gc.values) - vals_model = _to_nan_fillvalue(da_model_gc.values) - - # Map AI data to model gridcell positions using spatial mapping - ai_grid = np.full(n_grid, np.nan, dtype=float) - - if vals_ai.ndim == 1: - # For each model gridcell, find the corresponding AI gridcell and extract data - for g in range(n_grid): - ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] - if len(ai_gridcell_idx) > 0: - ai_gridcell_idx = ai_gridcell_idx[0] # Take the first match - ai_grid[g] = vals_ai[ai_gridcell_idx] - else: - # Handle multi-dimensional AI data - for g in range(n_grid): - ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] - if len(ai_gridcell_idx) > 0: - ai_gridcell_idx = ai_gridcell_idx[0] # Take the first match - # Extract data for this gridcell (handle different dimension orders) - if vals_ai.ndim == 2: - ai_grid[g] = vals_ai[0, ai_gridcell_idx] # Assume first dimension is not gridcell - else: - ai_grid[g] = vals_ai[ai_gridcell_idx] - - # Model data is already in the correct gridcell order - model_grid = vals_model - - stats = _plot_tripanel(var, "", grid_lon, grid_lat, ai_grid, model_grid, args.output_dir, - label_ai="AI Predictions", label_model="Model Results", plot=(not args.no_plot)) - stats_rows.append({ - "variable": var, - "suffix": "", - **stats, - }) - - else: - print(f" Skip {var} (unsupported dimensions: {dims})") + print(f" Skip {var} (unsupported dimensions: {dims})") # Write statistics outputs (CSV/TXT) if stats_rows: - # Resolve output paths - if _stats_only_requested: - # In stats-only mode, write a concise CSV with the requested metrics under stats folder - csv_out = os.path.join(args.output_dir, "summary_stats.csv") - txt_out = None - output_mode = 'csv' - else: - csv_out = args.stats_file if (args.stats_file and args.stats_file.lower().endswith('.csv')) else os.path.join(args.output_dir, "stats.csv") - txt_out = args.stats_file if (args.stats_file and args.stats_file.lower().endswith('.txt')) else os.path.join(args.output_dir, "stats.txt") - output_mode = args.stats_format + csv_out = args.stats_file if (args.stats_file and args.stats_file.lower().endswith('.csv')) else os.path.join(args.output_dir, 'stats.csv') + txt_out = args.stats_file if (args.stats_file and args.stats_file.lower().endswith('.txt')) else os.path.join(args.output_dir, 'stats.txt') os.makedirs(args.output_dir, exist_ok=True) - # CSV output - if (output_mode in ('csv', 'both')): - if _stats_only_requested: - fieldnames = [ - "variable", "suffix", - "ai_sum", "ai_std", "ai_min", "ai_max", - "model_sum", "model_std", "model_min", "model_max" - ] - # Reduce rows to requested columns only - filtered_rows = [] - for row in stats_rows: - filtered_rows.append({ - "variable": row.get("variable"), - "suffix": row.get("suffix"), - "ai_sum": row.get("ai_sum"), - "ai_std": row.get("ai_std"), - "ai_min": row.get("ai_min"), - "ai_max": row.get("ai_max"), - "model_sum": row.get("model_sum"), - "model_std": row.get("model_std"), - "model_min": row.get("model_min"), - "model_max": row.get("model_max"), - }) - rows_to_write = filtered_rows - else: - fieldnames = [ - "variable", "suffix", "ai_sum", "ai_min", "ai_max", - "model_sum", "model_min", "model_max", "n", "rmse", "nrmse", "r2" - ] - rows_to_write = stats_rows - with open(csv_out, "w", newline="") as f: + if args.stats_format in ('csv', 'both'): + fieldnames = [ + 'variable', 'suffix', 'ai_sum', 'ai_std', 'ai_min', 'ai_max', + 'model_sum', 'model_std', 'model_min', 'model_max', 'n', 'rmse', 'nrmse', 'r2' + ] + with open(csv_out, 'w', newline='') as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() - for row in rows_to_write: + for row in stats_rows: writer.writerow(row) print(f"Saved statistics CSV: {csv_out}") - # TXT output (readable, grouped by variable and suffix) - if (not _stats_only_requested) and (output_mode in ('txt', 'both')): - # Group stats by variable then suffix + if args.stats_format in ('txt', 'both'): from collections import defaultdict grouped = defaultdict(list) for row in stats_rows: - grouped[row["variable"]].append(row) + grouped[row['variable']].append(row) def _suffix_key(s): - # Sort order: gridcell-level (empty) first, then lev by number, then pft by number - if s == "": + if s == '': return (0, 0, 0) - if s.startswith("_lev"): + if s.startswith('_lev'): try: - return (1, int(s.replace("_lev", "")), 0) + return (1, int(s.replace('_lev', '')), 0) except Exception: return (1, 999999, 0) - if s.startswith("_pft"): + if s.startswith('_pft'): try: - return (2, int(s.replace("_pft", "")), 0) + return (2, int(s.replace('_pft', '')), 0) except Exception: return (2, 999999, 0) return (3, 0, 0) lines = [] - lines.append("AI vs Model Statistics Report") - lines.append("=" * 80) + header = 'NetCDF vs CSV Statistics Report' if use_csv_predictions else 'AI vs Model Statistics Report' + lines.append(header) + lines.append('=' * 80) for var in sorted(grouped.keys()): - rows = sorted(grouped[var], key=lambda r: _suffix_key(r.get("suffix", ""))) - lines.append("") + rows = sorted(grouped[var], key=lambda r: _suffix_key(r.get('suffix', ''))) + lines.append('') lines.append(f"Variable: {var}") - lines.append("-" * 80) + lines.append('-' * 80) for row in rows: title = f"{var}{row.get('suffix','')}" lines.append(title) - lines.append(" AI: sum={ai_sum:.6g} min={ai_min:.6g} max={ai_max:.6g}".format(**row)) - lines.append(" Model: sum={model_sum:.6g} min={model_min:.6g} max={model_max:.6g}".format(**row)) + if 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: + lines.append(" AI: sum={ai_sum:.6g} std={ai_std:.6g} min={ai_min:.6g} max={ai_max:.6g}".format(**row)) + lines.append(" Model: sum={model_sum:.6g} std={model_std:.6g} min={model_min:.6g} max={model_max:.6g}".format(**row)) lines.append(" Compare: n={n} rmse={rmse:.6g} nrmse={nrmse:.6g} r2={r2:.6g}".format(**row)) - lines.append("") + lines.append('') - with open(txt_out, "w") as f: + with open(txt_out, 'w') as f: f.write("\n".join(lines)) print(f"Saved statistics report: {txt_out}") + + # Generate a stacked bar chart summarizing quality by variable (CSV vs NetCDF) + try: + stats_df = pd.DataFrame(stats_rows) + if not stats_df.empty: + stats_df['quality'] = stats_df['r2'].apply(_classify_quality_from_r2) + quality_counts = ( + stats_df.groupby(['variable', 'quality']).size().unstack(fill_value=0) + ) + if not quality_counts.empty: + totals = quality_counts.sum(axis=1).replace(0, np.nan) + quality_pct = (quality_counts.div(totals, axis=0) * 100.0).fillna(0.0) + desired_order = ['good', 'ok', 'bad', 'unknown'] + available_cols = [c for c in desired_order if c in quality_pct.columns] + missing_cols = [c for c in desired_order if c not in available_cols] + for col in missing_cols: + quality_pct[col] = 0.0 + quality_pct = quality_pct[available_cols + missing_cols] if missing_cols else quality_pct[available_cols] + if 'good' in quality_pct.columns: + quality_pct = quality_pct.sort_values(by='good', ascending=False) + colors = { + 'good': '#2ecc71', + 'ok': '#f39c12', + 'bad': '#e74c3c', + 'unknown': '#7f8c8d' + } + plot_cols = [c for c in ['good', 'ok', 'bad', 'unknown'] if c in quality_pct.columns] + if plot_cols: + ax = quality_pct[plot_cols].plot( + kind='bar', + stacked=True, + 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' + plt.title(title, fontsize=16) + plt.xlabel('Variable', fontsize=14) + plt.ylabel('Percentage (%)', fontsize=14) + plt.xticks(rotation=90) + plt.legend(title='Category') + plt.tight_layout() + quality_fig = Path(args.output_dir) / 'comparison_quality_by_variable.png' + plt.savefig(quality_fig, dpi=300) + plt.close() + print(f"Saved quality summary figure: {quality_fig}") + except Exception as exc: + print(f"Warning: Failed to generate quality summary figure ({exc})") else: print("No statistics to write.") ds_ai.close() - ds_model.close() + if ds_model is not None: + ds_model.close() if args.no_plot: - print(f"\nCompleted without plotting. Output directory (for CSV): {args.output_dir}") + print(f"\nCompleted without plotting. Output directory: {args.output_dir}") else: print(f"\nAll plots done! Output directory: {args.output_dir}") if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/scripts/ai_predictions_to_restart.py b/scripts/ai_predictions_to_restart.py index e9b64d9..6878f3e 100644 --- a/scripts/ai_predictions_to_restart.py +++ b/scripts/ai_predictions_to_restart.py @@ -91,13 +91,9 @@ def create_spatial_mapping(ds_ai: xr.Dataset, ds_model: xr.Dataset) -> tuple[np. from scipy.spatial.distance import cdist ai_coords = np.column_stack([ai_lon, ai_lat]) model_coords = np.column_stack([model_lon, model_lat]) - distances = cdist(ai_coords, model_coords) - ai_to_model_mapping = np.argmin(distances, axis=1) - - print(f" Spatial mapping created: {len(ai_to_model_mapping)} AI -> {len(set(ai_to_model_mapping))} Model") - print(f" Mapping range: AI gridcell 0->model gridcell {ai_to_model_mapping[0]}") - print(f" Mapping range: AI gridcell {len(ai_lon)-1}->model gridcell {ai_to_model_mapping[-1]}") - + distances = cdist(model_coords, ai_coords) + model_to_ai_mapping = np.argmin(distances, axis=1) # 索引是模型格点, 值是最近的 AI 格点 + print(f" Spatial mapping created: {len(model_to_ai_mapping)} Model -> {len(set(model_to_ai_mapping))} AI") # Get grid information from the MODEL file as the master coordinate system n_grid = ds_model.sizes["gridcell"] print(f" Using MODEL gridcell count: {n_grid}") @@ -119,7 +115,7 @@ def create_spatial_mapping(ds_ai: xr.Dataset, ds_model: xr.Dataset) -> tuple[np. 'n_grid': n_grid } - return ai_to_model_mapping, variable_mapping + return model_to_ai_mapping, variable_mapping def auto_detect_variable_list(ai_predictions_path: Path) -> list: @@ -148,8 +144,7 @@ def auto_detect_variable_list(ai_predictions_path: Path) -> list: def create_updated_restart_file(restart_file_path: Path, output_path: Path, ai_predictions_path: Path, cnp_io_variables: List[str], - ai_to_model_mapping: np.ndarray, variable_mapping: Dict[str, Any]) -> None: - """Directly update the restart file using netCDF4 without xarray encoding issues.""" + model_to_ai_mapping: np.ndarray, variable_mapping: Dict[str, Any]) -> None: print(f"Saving updated restart file to: {output_path}") # Create output directory if it doesn't exist @@ -184,23 +179,15 @@ def create_updated_restart_file(restart_file_path: Path, output_path: Path, # Get PFTs in this gridcell gridcell_pfts = grid_to_pfts[g] - # Find corresponding AI gridcell using spatial mapping - ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] - if len(ai_gridcell_idx) > 0: - ai_gridcell_idx = ai_gridcell_idx[0] - - # Update each PFT instance in this gridcell - # Only update PFT1-PFT16 (skip PFT0), and only in the first column - # Get the first 16 PFTs in this gridcell (exact same as working script) - gridcell_pfts = gridcell_pfts[:16] # First 16 PFTs - for pft_idx, model_pft_idx in enumerate(gridcell_pfts): - # Skip PFT0 (index 0), start from PFT1 (index 1) - if 1 <= pft_idx <= 16 and model_pft_idx < len(model_var): - # AI PFT0 -> Model PFT1, AI PFT1 -> Model PFT2, etc. - # Adjust index: AI PFT k corresponds to Model PFT (k+1) in the first 16 - adjusted_k = pft_idx - 1 # AI PFT0 -> Model PFT1, AI PFT1 -> Model PFT2 - if adjusted_k < ai_data.shape[0]: - model_var[model_pft_idx] = ai_data[adjusted_k, ai_gridcell_idx] + + ai_gridcell_idx = model_to_ai_mapping[g] + gridcell_pfts = gridcell_pfts[:16] + for pft_idx, model_pft_idx in enumerate(gridcell_pfts): + # Skip PFT0 (index 0), start from PFT1 (index 1) + if 1 <= pft_idx <= 16 and model_pft_idx < len(model_var): + adjusted_k = pft_idx - 1 + if adjusted_k < ai_data.shape[0]: + model_var[model_pft_idx] = ai_data[adjusted_k, ai_gridcell_idx] # Update soil variables for var_name in ds_ai.variables: @@ -224,28 +211,22 @@ def create_updated_restart_file(restart_file_path: Path, output_path: Path, # Get columns in this gridcell gridcell_cols = grid_to_cols[g] - # Find corresponding AI gridcell using spatial mapping - ai_gridcell_idx = np.where(ai_to_model_mapping == g)[0] - if len(ai_gridcell_idx) > 0: - ai_gridcell_idx = ai_gridcell_idx[0] - - # Update first column in this gridcell (use AI column 0) - if len(gridcell_cols) > 0: - model_col_idx = gridcell_cols[0] # First column of this gridcell - if model_col_idx < model_var.shape[0]: - # Update only first 10 layers for this column (even if model has 15 layers) - layers_to_update = min(10, ai_data.shape[1]) - for layer_idx in range(layers_to_update): - # Handle AI data indexing - shape is (column, levgrnd, gridcell) - if ai_data.ndim == 3: - model_var[model_col_idx, layer_idx] = ai_data[0, layer_idx, ai_gridcell_idx] - else: - model_var[model_col_idx, layer_idx] = ai_data[0, layer_idx] + + ai_gridcell_idx = model_to_ai_mapping[g] + + if len(gridcell_cols) > 0: + model_col_idx = gridcell_cols[0] + if model_col_idx < model_var.shape[0]: + layers_to_update = min(10, ai_data.shape[1]) + for layer_idx in range(layers_to_update): + if ai_data.ndim == 3: + model_var[model_col_idx, layer_idx] = ai_data[0, layer_idx, ai_gridcell_idx] + else: + model_var[model_col_idx, layer_idx] = ai_data[0, layer_idx] print(f"Updated restart file saved successfully!") print(f"File size: {output_path.stat().st_size / (1024*1024):.1f} MB") - def get_varlist_name_from_config(ai_predictions_path): for parent in [ai_predictions_path.parent] + list(ai_predictions_path.parents): config_path = parent / 'cnp_config.json' diff --git a/scripts/ai_restart_comparison.py b/scripts/ai_restart_comparison.py index 3e4e248..53ad67e 100644 --- a/scripts/ai_restart_comparison.py +++ b/scripts/ai_restart_comparison.py @@ -11,18 +11,36 @@ from pathlib import Path import sys import json +from scipy.spatial.distance import cdist # Project imports sys.path.append(str(Path(__file__).resolve().parents[1])) from config.training_config import parse_cnp_io_list -# Default file paths -DATA_DIR = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/' -DEFAULT_FILE_OLD = DATA_DIR + '20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' - -LABEL_NEW = "AI Generated" -LABEL_OLD = "Original Model" -OUTPUT_DIR = "./ai_restart_comparison_plots" +# Default file paths (fallback values) +FALLBACK_DATA_DIR = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/' +FALLBACK_REFERENCE_FILE = FALLBACK_DATA_DIR + '20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' +FALLBACK_AI_PREDICTIONS = './comparison_results/ai_predictions_for_plotting.nc' +FALLBACK_OUTPUT_DIR = './ai_restart_comparison_plots' +LABEL_NEW = 'AI Restart' +LABEL_OLD_DEFAULT = 'Original Restart' + +QUALITY_THRESHOLDS = { + 'good': 0.9, # R² >= 0.9 for good quality + 'ok': 0.7 # R² >= 0.7 for ok quality +} + + +def _classify_quality_from_r2(value): + """Classify quality based on R² value (similar to generate_prediction_quality_report.py)""" + if value is None or not np.isfinite(value): + return 'unknown' + val = float(value) + if val >= QUALITY_THRESHOLDS['good']: + return 'good' + if val >= QUALITY_THRESHOLDS['ok']: + return 'ok' + return 'bad' def _safe_get(ds, name): if name not in ds: @@ -56,6 +74,19 @@ def _build_gridcell_groups(one_d_to_grid, n_grid): groups[g].append(idx) return groups +def _build_target_to_source_mapping(ds_source, ds_target): + src_lon = _safe_get(ds_source, 'grid1d_lon').values + src_lat = _safe_get(ds_source, 'grid1d_lat').values + tgt_lon = _safe_get(ds_target, 'grid1d_lon').values + tgt_lat = _safe_get(ds_target, 'grid1d_lat').values + if (len(src_lon) == len(tgt_lon) and + np.allclose(src_lon, tgt_lon) and + np.allclose(src_lat, tgt_lat)): + return np.arange(len(tgt_lon), dtype=int) + mapping = np.argmin(cdist(np.column_stack([tgt_lon, tgt_lat]), + np.column_stack([src_lon, src_lat])), axis=1) + return mapping + def _plot_map(ax, lon, lat, data, title, vmin=None, vmax=None, cmap="viridis", norm=None): ax.add_feature(cfeature.COASTLINE) ax.add_feature(cfeature.BORDERS) @@ -95,14 +126,14 @@ def _percent_diff_categories(model_vals: np.ndarray, ai_vals: np.ndarray) -> np. cat[finite] = categories.astype(float) return cat + def _plot_tripanel(var, label_suffix, lon, lat, data_new, data_old, out_dir, - label_new="AI Enhanced", label_old="Original Model"): + label_new="AI Restart", label_old="Original Restart", plot=True): diff = data_new - data_old vmin_orig = np.nanmin([np.nanmin(data_new), np.nanmin(data_old)]) vmax_orig = np.nanmax([np.nanmax(data_new), np.nanmax(data_old)]) diff_abs = np.nanmax(np.abs(diff)) - # Compute stats and metrics (old vs new) with NaN safety finite_new = np.isfinite(data_new) finite_old = np.isfinite(data_old) mask = finite_new & finite_old @@ -115,10 +146,12 @@ def _nan_max(a): return float(np.nanmax(a)) if np.any(np.isfinite(a)) else float('nan') sum_new = float(np.nansum(data_new)) + std_new = float(np.nanstd(data_new)) if np.any(np.isfinite(data_new)) else float('nan') min_new = _nan_min(data_new) max_new = _nan_max(data_new) sum_old = float(np.nansum(data_old)) + std_old = float(np.nanstd(data_old)) if np.any(np.isfinite(data_old)) else float('nan') min_old = _nan_min(data_old) max_old = _nan_max(data_old) @@ -138,9 +171,27 @@ def _nan_max(a): r2 = float('nan') print(f"Stats for {var}{label_suffix}:") - print(f" {label_new}: sum={sum_new:.6g} min={min_new:.6g} max={max_new:.6g}") - print(f" {label_old}: sum={sum_old:.6g} min={min_old:.6g} max={max_old:.6g}") - print(f" Metrics (AI vs Model): n={n} rmse={rmse:.6g} nrmse={nrmse:.6g} r2={r2:.6g}") + print(f" {label_new}: sum={sum_new:.6g} std={std_new:.6g} min={min_new:.6g} max={max_new:.6g}") + print(f" {label_old}: sum={sum_old:.6g} std={std_old:.6g} min={min_old:.6g} max={max_old:.6g}") + print(f" Metrics ({label_new} vs {label_old}): n={n} rmse={rmse:.6g} nrmse={nrmse:.6g} r2={r2:.6g}") + + stats = { + 'new_sum': sum_new, + 'new_std': std_new, + 'new_min': min_new, + 'new_max': max_new, + 'old_sum': sum_old, + 'old_std': std_old, + 'old_min': min_old, + 'old_max': max_old, + 'n': n, + 'rmse': rmse, + 'nrmse': nrmse, + 'r2': r2, + } + + if not plot: + return stats fig = plt.figure(figsize=(12, 20)) gs = gridspec.GridSpec(4, 1, figure=fig, hspace=0.3) @@ -165,18 +216,16 @@ def _nan_max(a): norm = TwoSlopeNorm(vmin=-diff_abs, vcenter=0, vmax=diff_abs) _plot_map(ax3, lon, lat, diff, f"{var} - Diff ({label_new} - {label_old})", cmap="RdBu_r", norm=norm) - # Percent-difference categorical map (fourth panel) ax4 = fig.add_subplot(gs[3, 0], projection=ccrs.PlateCarree()) - # Treat data_old as model, data_new as AI cat = _percent_diff_categories(data_old, data_new) colors = [ - "#08519c", # -3: 30%+ - "#6baed6", # -2: 10–30% - "#c6dbef", # -1: 0–10% - "#bdbdbd", # 0: 0 - "#fcbba1", # +1: 0–10% - "#fb6a4a", # +2: 10–30% - "#cb181d", # +3: 30%+ + "#08519c", + "#6baed6", + "#c6dbef", + "#bdbdbd", + "#fcbba1", + "#fb6a4a", + "#cb181d", ] cmap = ListedColormap(colors) boundaries = [-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5] @@ -200,33 +249,28 @@ def _nan_max(a): plt.suptitle(f"{var} {label_suffix}", fontsize=16, fontweight="bold", y=0.96) fig.tight_layout(rect=[0, 0.03, 1, 0.95]) os.makedirs(out_dir, exist_ok=True) - path = os.path.join(out_dir, f"{var}{label_suffix}.png") - plt.savefig(path, dpi=300, bbox_inches="tight") + path_out = os.path.join(out_dir, f"{var}{label_suffix}.png") + plt.savefig(path_out, dpi=300, bbox_inches="tight") plt.close() - print(f" Saved: {path}") - -def auto_detect_variable_list_from_config(ai_restart_path: str): - run_dir = Path(ai_restart_path).parent if ai_restart_path else Path('.') - for parent in [run_dir] + list(run_dir.parents): - config_path = parent / 'cnp_config.json' - if config_path.exists(): - try: - with open(config_path, 'r') as f: - config = json.load(f) - data_info = config.get('data_info', {}) - vars_1d = data_info.get('variables_1d_pft', []) - # Try both keys for 2d soil variables - vars_2d = data_info.get('variables_2d_soil', []) - if not vars_2d: - vars_2d = data_info.get('x_list_columns_2d', []) - print(f"Auto-detected variables from {config_path}") - print(f" 1D PFT variables: {vars_1d}") - print(f" 2D soil variables: {vars_2d}") - return list(vars_1d), list(vars_2d) - except Exception as e: - print(f"Warning: Failed to parse {config_path}: {e}") - print("Warning: Could not auto-detect variable list. No variables will be plotted.") - return [], [] + print(f" Saved: {path_out}") + + return stats + +def _load_default_paths(variable_list_path: str): + defaults = {} + if variable_list_path: + try: + vl_path = Path(variable_list_path) + if vl_path.exists(): + parsed = parse_cnp_io_list(variable_list_path) + if isinstance(parsed, dict): + for key in ('ai_predictions_default', 'model_default', 'comparison_output_dir', 'ai_restart_default'): + value = parsed.get(key) + if value: + defaults[key] = value + except Exception as exc: + print(f"Warning: Failed to load default paths from {variable_list_path}: {exc}") + return defaults def parse_arguments(): """Parse command line arguments.""" @@ -236,17 +280,23 @@ def parse_arguments(): parser.add_argument('--ai-restart', type=str, default=None, help='Path to AI-enhanced restart file (auto-detected if not specified)') parser.add_argument('--original-restart', type=str, default=None, - help='Path to original model restart file (default: 780 year model results)') + help='Path to reference dataset (original restart or AI predictions)') parser.add_argument('--layers', type=str, default='0,3,5', help='Comma-separated list of soil layers to plot (default: 0,3,5)') parser.add_argument('--pfts', type=str, default='1,2,4,5', help='Comma-separated list of PFTs to plot (default: all PFT0-PFT15)') parser.add_argument('--plot-all', action='store_true', help='If set, plot all variables for all 10 layers (0-9) and all 16 PFTs (1-16, skip pft0)') + parser.add_argument('--output-dir', type=str, default=None, + help='Output directory for plots/statistics') parser.add_argument('--stats-only', action='store_true', - help='Only compute statistics (sum/std/min/max) for all variables and all layers/PFTs; no plots') - parser.add_argument('--output-dir', type=str, default=OUTPUT_DIR, - help=f'Output directory for plots or stats (default: {OUTPUT_DIR})') + help='Only compute statistics (no plots)') + parser.add_argument('--no-plot', action='store_true', + help='Disable plot generation') + parser.add_argument('--stats-file', type=str, default=None, + help='Optional path to write statistics file (csv/txt)') + parser.add_argument('--stats-format', type=str, choices=['csv', 'txt', 'both'], default='txt', + help='Statistics output format (default: txt)') return parser.parse_args() def find_ai_restart_file(): @@ -261,87 +311,102 @@ def find_ai_restart_file(): else: return None + def main(): args = parse_arguments() - - # Resolve output directory (stats-only goes to a stats subfolder) - out_dir_base = Path(args.output_dir) + + defaults = _load_default_paths(args.variable_list) + + def _resolve(value, keys, fallback): + if value: + return str(value) + for key in keys: + val = defaults.get(key) + if val: + return str(val) + return fallback + if args.stats_only: - output_dir = out_dir_base / 'stats' - else: - output_dir = out_dir_base - os.makedirs(output_dir, exist_ok=True) - # Parse layers and PFTs + args.no_plot = True + plot_enabled = not args.no_plot + + ai_restart_path = _resolve(args.ai_restart, ('ai_restart_default',), None) + if not ai_restart_path: + ai_restart_path = find_ai_restart_file() + if not ai_restart_path: + print('Error: No AI-enhanced restart file found. Please specify with --ai-restart') + return + + reference_path = _resolve(args.original_restart, ('ai_predictions_default', 'model_default'), FALLBACK_REFERENCE_FILE) + output_dir_str = _resolve(args.output_dir, ('comparison_output_dir',), FALLBACK_OUTPUT_DIR) + output_dir = Path(output_dir_str) + output_dir.mkdir(parents=True, exist_ok=True) + + label_new = LABEL_NEW + label_old = LABEL_OLD_DEFAULT + default_ai_predictions = defaults.get('ai_predictions_default', FALLBACK_AI_PREDICTIONS) + try: + if Path(reference_path).resolve() == Path(default_ai_predictions).resolve(): + label_old = 'AI Predictions' + except Exception: + pass + + print(f'Using AI-enhanced restart: {ai_restart_path}') + print(f'Using reference dataset: {reference_path}') + print(f'Output directory: {output_dir}') + global LEVGRND_LAYERS, PFT_PICK_LIST if args.plot_all: - LEVGRND_LAYERS = list(range(10)) # 0-9 - PFT_PICK_LIST = list(range(16)) # 0-15 (will skip pft0 in plotting) - else: - LEVGRND_LAYERS = [int(x.strip()) for x in args.layers.split(',')] - PFT_PICK_LIST = [int(x.strip()) for x in args.pfts.split(',')] - - # Set file paths - if args.ai_restart: - FILE_NEW = args.ai_restart + LEVGRND_LAYERS = list(range(10)) + PFT_PICK_LIST = list(range(16)) else: - FILE_NEW = find_ai_restart_file() - if not FILE_NEW: - print("Error: No AI-enhanced restart file found. Please specify with --ai-restart") - return - - if args.original_restart: - FILE_OLD = args.original_restart - else: - FILE_OLD = DEFAULT_FILE_OLD - - print(f"Using AI-enhanced restart: {FILE_NEW}") - print(f"Using original restart: {FILE_OLD}") - print(f"Layers to plot: {LEVGRND_LAYERS}") - print(f"PFTs to plot: {PFT_PICK_LIST}") - print("-" * 80) - - # Parse CNP_IO list to get variables + if args.stats_only: + LEVGRND_LAYERS = list(range(10)) + PFT_PICK_LIST = list(range(16)) + else: + LEVGRND_LAYERS = [int(x.strip()) for x in args.layers.split(',') if x.strip()] + PFT_PICK_LIST = [int(x.strip()) for x in args.pfts.split(',') if x.strip()] + print(f'Layers to process: {LEVGRND_LAYERS}') + print(f'PFTs to process: {PFT_PICK_LIST}') + print('-' * 80) + if args.variable_list: - print("Parsing CNP_IO list...") + print('Parsing CNP_IO list...') cnp_io_vars = parse_cnp_io_list(Path(args.variable_list)) pft_1d_variables = cnp_io_vars.get('pft_1d_variables', []) - variables_2d_soil = cnp_io_vars.get('variables_2d_soil', []) - if not variables_2d_soil: - variables_2d_soil = cnp_io_vars.get('x_list_columns_2d', []) - print(f" PFT1D: {pft_1d_variables}") - print(f" Soil2D: {variables_2d_soil}") + variables_2d_soil = cnp_io_vars.get('variables_2d_soil', []) or cnp_io_vars.get('x_list_columns_2d', []) else: - print("Auto-detecting CNP_IO variables from config.json...") - pft_1d_variables, variables_2d_soil = auto_detect_variable_list_from_config(FILE_NEW) - # Combine all variables to plot - VARIABLES = pft_1d_variables + variables_2d_soil - if not VARIABLES: - print("Error: No variables found in CNP_IO list or config.json") + print('Auto-detecting CNP_IO variables from config.json...') + pft_1d_variables, variables_2d_soil = auto_detect_variable_list_from_config(ai_restart_path) + variables = pft_1d_variables + variables_2d_soil + if not variables: + print('Error: No variables found in CNP_IO list or config.json') return - print(f"Variables to plot: {VARIABLES}") - print(f" PFT1D: {pft_1d_variables}") - print(f" Soil2D: {variables_2d_soil}") - - ds_new = xr.open_dataset(FILE_NEW) - ds_old = xr.open_dataset(FILE_OLD) + print(f'Variables to process ({len(variables)}): {variables}') - grid_lon, grid_lat = _gridcell_lonlat(ds_new) - n_grid = ds_new.sizes["gridcell"] + ds_new = xr.open_dataset(ai_restart_path) + ds_old = xr.open_dataset(reference_path) - col2grid = _to_zero_based_index(_safe_get(ds_new, "cols1d_gridcell_index").values, n_grid) - pft2grid = _to_zero_based_index(_safe_get(ds_new, "pfts1d_gridcell_index").values, n_grid) + grid_mapping = _build_target_to_source_mapping(ds_old, ds_new) + grid_lon, grid_lat = _gridcell_lonlat(ds_new) + n_grid = ds_new.sizes['gridcell'] + + col2grid = _to_zero_based_index(_safe_get(ds_new, 'cols1d_gridcell_index').values, n_grid) + pft2grid = _to_zero_based_index(_safe_get(ds_new, 'pfts1d_gridcell_index').values, n_grid) grid_to_cols = _build_gridcell_groups(col2grid, n_grid) grid_to_pfts = _build_gridcell_groups(pft2grid, n_grid) - 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]}") + 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]}') stats_rows = [] - print(f"\nStart {'statistics' if args.stats_only else 'plotting'}: {len(VARIABLES)} variables") - for var in VARIABLES: + debug_enabled = not args.stats_only + + + for var in variables: if (var not in ds_new.data_vars) or (var not in ds_old.data_vars): - print(f"Skip {var} (not found in both files)") + print(f'Skip {var} (not found in both files)') continue da_new = ds_new[var] @@ -349,22 +414,25 @@ def main(): dims = da_new.dims print(f"\nVariable {var}, dims: {dims}") - if ("column" in dims) and ("levgrnd" in dims): - # Soil2D variable - da_new_cl = da_new.transpose("column", "levgrnd") - da_old_cl = da_old.transpose("column", "levgrnd") + if ('column' in dims) and ('levgrnd' in dims): + da_new_cl = da_new.transpose('column', 'levgrnd', ...) vals_new = _to_nan_fillvalue(da_new_cl.values) - vals_old = _to_nan_fillvalue(da_old_cl.values) - # Iterate all layers if stats-only; otherwise iterate requested layers - if args.stats_only: - lev_iter = range(int(da_new_cl.sizes["levgrnd"])) - else: - lev_iter = LEVGRND_LAYERS + if 'gridcell' not in da_old.dims: + print(' Skip (reference dataset lacks gridcell dimension)') + continue - for lev in lev_iter: - if lev < 0 or lev >= da_new_cl.sizes["levgrnd"]: - print(f" Layer {lev} out of range, skipped") + 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 new_grid = np.full(n_grid, np.nan, dtype=float) @@ -374,57 +442,54 @@ def main(): cols = grid_to_cols[g] if len(cols) == 0: continue - c0 = cols[0] # First column only - # Use column 0 for AI prediction file if it has only one column + c0 = cols[0] if vals_new.shape[0] == 1: new_grid[g] = vals_new[0, lev] else: new_grid[g] = vals_new[c0, lev] - # Always use mapping for old/model file - old_grid[g] = vals_old[c0, lev] - - if args.stats_only: - # Compute stats only - new_sum = float(np.nansum(new_grid)) - new_std = float(np.nanstd(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') - new_min = float(np.nanmin(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') - new_max = float(np.nanmax(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') - old_sum = float(np.nansum(old_grid)) - old_std = float(np.nanstd(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') - old_min = float(np.nanmin(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') - old_max = float(np.nanmax(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') - stats_rows.append({ - "variable": var, - "suffix": f"_lev{lev}", - "new_sum": new_sum, "new_std": new_std, "new_min": new_min, "new_max": new_max, - "old_sum": old_sum, "old_std": old_std, "old_min": old_min, "old_max": old_max, - }) - else: - # Plotting mode - _plot_tripanel(var, f"_lev{lev}", grid_lon, grid_lat, new_grid, old_grid, str(out_dir_base), - label_new=LABEL_NEW, label_old=LABEL_OLD) - - elif ("pft" in dims) and (len(dims) == 1): - # PFT1D variable - da_new_p = da_new.transpose("pft") - da_old_p = da_old.transpose("pft") + + 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']: + 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]: + continue + idx[axis_lev] = lev + idx[axis_grid] = src_idx + 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)}') + print(f'[DEBUG] {var} lev{lev}: ref min={np.nanmin(old_grid)} max={np.nanmax(old_grid)} mean={np.nanmean(old_grid)}') + + stats = _plot_tripanel(var, f'_lev{lev}', grid_lon, grid_lat, new_grid, old_grid, str(output_dir), + label_new=label_new, label_old=label_old, plot=plot_enabled) + 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)') + continue + + da_new_p = da_new.transpose(..., 'pft') vals_new = _to_nan_fillvalue(da_new_p.values) - vals_old = _to_nan_fillvalue(da_old_p.values) - # Iterate all PFT TYPES (0..15) if stats-only; otherwise iterate requested PFTs - # Note: da_new_p.sizes["pft"] is the total number of PFT entries across all gridcells (very large). - # For comparison we want PFT type indices 0..15 which are mapped per-gridcell via grid_to_pfts. - if args.stats_only: - pft_iter = range(16) - else: - pft_iter = PFT_PICK_LIST + 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_iter: - if k < 0 or k >= da_new_p.sizes["pft"]: - print(f" PFT {k} out of range, skipped") + for k in PFT_PICK_LIST: + if k < 0 or k >= total_pfts: + print(f' PFT {k} out of range, skipped') continue if args.plot_all and k == 0: - continue # skip pft0 for plot-all + continue + new_grid = np.full(n_grid, np.nan, dtype=float) old_grid = np.full(n_grid, np.nan, dtype=float) @@ -434,69 +499,137 @@ def main(): continue if k < len(pfts): p_idx = pfts[k] - new_grid[g] = vals_new[p_idx] - old_grid[g] = vals_old[p_idx] - - if args.stats_only: - new_sum = float(np.nansum(new_grid)) - new_std = float(np.nanstd(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') - new_min = float(np.nanmin(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') - new_max = float(np.nanmax(new_grid)) if np.any(np.isfinite(new_grid)) else float('nan') - old_sum = float(np.nansum(old_grid)) - old_std = float(np.nanstd(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') - old_min = float(np.nanmin(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') - old_max = float(np.nanmax(old_grid)) if np.any(np.isfinite(old_grid)) else float('nan') - stats_rows.append({ - "variable": var, - "suffix": f"_pft{k}", - "new_sum": new_sum, "new_std": new_std, "new_min": new_min, "new_max": new_max, - "old_sum": old_sum, "old_std": old_std, "old_min": old_min, "old_max": old_max, - }) - else: - _plot_tripanel(var, f"_pft{k}", grid_lon, grid_lat, new_grid, old_grid, str(out_dir_base), - label_new=LABEL_NEW, label_old=LABEL_OLD) + 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] + 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) + stats_rows.append({'variable': var, 'suffix': f'_pft{k}', 'label_new': label_new, 'label_old': label_old, 'quality': _classify_quality_from_r2(stats['r2']), **stats}) else: - print(f" Skip {var} (only supports (column, levgrnd) and (pft,))") - - # Write stats CSV if stats-only - if args.stats_only and stats_rows: - import csv - csv_path = output_dir / 'summary_stats.csv' - fieldnames = [ - 'variable', 'suffix', - 'new_sum', 'new_std', 'new_min', 'new_max', - 'old_sum', 'old_std', 'old_min', 'old_max' - ] - with open(csv_path, 'w', newline='') as f: - writer = csv.DictWriter(f, fieldnames=fieldnames) - writer.writeheader() - for row in stats_rows: - writer.writerow(row) - print(f"Saved statistics CSV: {csv_path}") - + print(f' Skip {var} (unsupported dimensions)') ds_new.close() ds_old.close() - # Print summary - print("\n" + "="*80) - print("AI RESTART COMPARISON SUMMARY:") - print("="*80) - print(f"Original restart: {os.path.abspath(FILE_OLD)}") - print(f"AI-enhanced restart: {os.path.abspath(FILE_NEW)}") - print(f"Labels: {LABEL_OLD} vs {LABEL_NEW}") - print(f"Output directory: {os.path.abspath(str(output_dir))}") - print(f"Variables processed: {VARIABLES}") - if args.stats_only: - print("Mode: stats-only (all layers and all PFTs)") - else: - print(f"Layers plotted: {LEVGRND_LAYERS}") - print(f"PFTs plotted: {PFT_PICK_LIST}") - print("="*80) - if args.stats_only: - print("\nCompleted without plotting. Stats CSV saved to:", str(output_dir)) + if stats_rows: + stats_format = args.stats_format + stats_dir = output_dir + csv_out = None + txt_out = None + if args.stats_file: + stats_path = Path(args.stats_file) + if stats_path.suffix.lower() == '.csv': + csv_out = stats_path + elif stats_path.suffix.lower() == '.txt': + txt_out = stats_path + else: + csv_out = stats_path.with_suffix('.csv') + txt_out = stats_path.with_suffix('.txt') + else: + csv_out = stats_dir / 'restart_stats.csv' + txt_out = stats_dir / 'restart_stats.txt' + + if stats_format in ('csv', 'both') and csv_out: + fieldnames = ['variable', 'suffix', 'label_new', 'label_old', 'quality', 'new_sum', 'new_std', 'new_min', 'new_max', 'old_sum', 'old_std', 'old_min', 'old_max', 'n', 'rmse', 'nrmse', 'r2'] + with csv_out.open('w', newline='') as f: + import csv + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for row in stats_rows: + writer.writerow(row) + print(f'Saved statistics CSV: {csv_out}') + + if stats_format in ('txt', 'both') and txt_out: + from collections import defaultdict + grouped = defaultdict(list) + for row in stats_rows: + grouped[row['variable']].append(row) + lines = [] + lines.append('AI Restart Comparison Statistics') + lines.append('=' * 80) + for var in sorted(grouped.keys()): + rows = sorted(grouped[var], key=lambda r: r.get('suffix', '')) + lines.append('') + lines.append(f'Variable: {var}') + lines.append('-' * 80) + for row in rows: + lines.append(f"{var}{row.get('suffix', '')}") + lines.append(f" {row['label_new']}: sum={row['new_sum']:.6g} std={row['new_std']:.6g} min={row['new_min']:.6g} max={row['new_max']:.6g}") + lines.append(f" {row['label_old']}: sum={row['old_sum']:.6g} std={row['old_std']:.6g} min={row['old_min']:.6g} max={row['old_max']:.6g}") + lines.append(f" Compare: n={row['n']} rmse={row['rmse']:.6g} nrmse={row['nrmse']:.6g} r2={row['r2']:.6g}") + lines.append(f" Quality: {row['quality']}") + lines.append('') + txt_out.write_text('\n'.join(lines)) + print(f'Saved statistics report: {txt_out}') + + # Generate quality summary figure + categories = ['good', 'ok', 'bad', 'unknown'] + category_colors = { + 'good': '#2ecc71', + 'ok': '#f39c12', + 'bad': '#e74c3c', + 'unknown': '#7f8c8d' + } + from collections import defaultdict + quality_counts = defaultdict(lambda: {cat: 0 for cat in categories}) + for row in stats_rows: + cat = row.get('quality', 'unknown') or 'unknown' + if cat not in categories: + cat = 'unknown' + quality_counts[row['variable']][cat] += 1 + labels = sorted(quality_counts.keys()) + totals = [sum(quality_counts[var].values()) for var in labels] + if labels and any(total > 0 for total in totals): + percentages = [] + for idx, var in enumerate(labels): + total = totals[idx] + pct = {} + for cat in categories: + if total > 0: + pct[cat] = quality_counts[var].get(cat, 0) / total * 100.0 + else: + pct[cat] = 0.0 + percentages.append(pct) + fig_width = max(12.0, len(labels) * 0.4) + fig, ax = plt.subplots(figsize=(fig_width, 8)) + positions = np.arange(len(labels)) + bottom = np.zeros(len(labels), dtype=float) + for cat in categories: + heights = [pct[cat] for pct in percentages] + ax.bar(positions, heights, bottom=bottom, color=category_colors.get(cat, 'gray'), label=cat.capitalize()) + bottom += heights + ax.set_xticks(positions) + ax.set_xticklabels(labels, rotation=90) + ax.set_ylabel('Percentage (%)') + ax.set_ylim(0, 100) + ax.set_title('Restart Comparison Quality by Variable') + ax.legend(title='Category') + fig.tight_layout() + quality_fig = output_dir / 'restart_quality_by_variable.png' + plt.savefig(quality_fig, dpi=300, bbox_inches='tight') + plt.close(fig) + print(f'Saved quality summary figure: {quality_fig}') + + print('\n' + '=' * 80) + print('AI RESTART COMPARISON SUMMARY:') + print('=' * 80) + print(f'Reference dataset: {os.path.abspath(reference_path)}') + print(f'AI-enhanced restart: {os.path.abspath(ai_restart_path)}') + print(f'Labels: {label_old} vs {label_new}') + print(f'Output directory: {output_dir.resolve()}') + print(f'Variables processed: {variables}') + print(f'Layers processed: {LEVGRND_LAYERS}') + print(f'PFTs processed: {PFT_PICK_LIST}') + print('=' * 80) + if plot_enabled: + print('\nAll plots done! Output dir:', output_dir) else: - print("\nAll plots done! Output dir:", str(out_dir_base)) + print('\nPlots disabled. Statistics written to output directory.') if __name__ == "__main__": main() @@ -513,4 +646,4 @@ def main(): print() print("3. Full custom configuration:") print(" python ai_restart_comparison.py --variable-list ../../CNP_IO_demo1.txt --ai-restart ai_file.nc --original-restart orig_file.nc --layers 0,5,9 --pfts 0,1,2,3,4,5") - print("="*60) + print("="*60) \ No newline at end of file From 20e512c4da056da4fb2217967a2d083a18674538 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Fri, 24 Oct 2025 09:28:48 -0700 Subject: [PATCH 29/51] Move fallback path to CNP_IO txt and clean up hardcoded paths --- CNP_IO_updated9_dev.txt | 2 ++ config/training_config.py | 13 +++++++++++-- scripts/ai_restart_comparison.py | 19 ++++++++++++++++++- 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/CNP_IO_updated9_dev.txt b/CNP_IO_updated9_dev.txt index 28ec192..0bec28c 100644 --- a/CNP_IO_updated9_dev.txt +++ b/CNP_IO_updated9_dev.txt @@ -3,6 +3,8 @@ MODEL_DEFAULT: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/202 COMPARISON_OUTPUT_DIR: ./ai_model_comparison_plots CSV_PREDICTIONS_DEFAULT: ./cnp_inference_entire_dataset/cnp_predictions AI_RESTART_DEFAULT: ./updated_restart_CNP_IO_updated9_dev_20250408_trendytest_ICB1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc +FALLBACK_DATA_DIR: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/ +FALLBACK_REFERENCE_FILENAME: 20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc TRENDY1_PATH: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree TRENDY05_PATH = diff --git a/config/training_config.py b/config/training_config.py index e58ca46..d297872 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -446,7 +446,10 @@ def parse_cnp_io_list(filename): 'model_default': None, 'comparison_output_dir': None, 'csv_predictions_default': None, - 'ai_restart_default': None + 'ai_restart_default': None, + 'fallback_data_dir': None, + 'fallback_reference_file': None, + 'fallback_reference_filename': None }) current_section = None @@ -499,7 +502,7 @@ def parse_cnp_io_list(filename): # FILE_PATTERN: enhanced_1_training_data_batch_*.pkl # DATA_PATHS: /p1,/p2 if line and not line.startswith('#'): - kv_match = re.match(r'(?i)^(trendy1_path|trendy05_path|tva4km_path|file_pattern|trendy1_file_pattern|trendy05_file_pattern|tva4km_file_pattern|data_paths|ai_predictions_default|model_default|comparison_output_dir|csv_predictions_default|ai_restart_default)\s*[:=]\s*(.+)$', line) + kv_match = re.match(r'(?i)^(trendy1_path|trendy05_path|tva4km_path|file_pattern|trendy1_file_pattern|trendy05_file_pattern|tva4km_file_pattern|data_paths|ai_predictions_default|model_default|comparison_output_dir|csv_predictions_default|ai_restart_default|fallback_data_dir|fallback_reference_file|fallback_reference_filename)\s*[:=]\s*(.+)$', line) if kv_match: key = kv_match.group(1).lower() val = kv_match.group(2).strip() @@ -531,6 +534,12 @@ def parse_cnp_io_list(filename): result['csv_predictions_default'] = val elif key == 'ai_restart_default': result['ai_restart_default'] = val + elif key == 'fallback_data_dir': + result['fallback_data_dir'] = val + elif key == 'fallback_reference_file': + result['fallback_reference_file'] = val + elif key == 'fallback_reference_filename': + result['fallback_reference_filename'] = val return result def parse_cnp_model_config(filename: str) -> Dict[str, Any]: diff --git a/scripts/ai_restart_comparison.py b/scripts/ai_restart_comparison.py index 53ad67e..8970a2a 100644 --- a/scripts/ai_restart_comparison.py +++ b/scripts/ai_restart_comparison.py @@ -264,7 +264,7 @@ def _load_default_paths(variable_list_path: str): if vl_path.exists(): parsed = parse_cnp_io_list(variable_list_path) if isinstance(parsed, dict): - for key in ('ai_predictions_default', 'model_default', 'comparison_output_dir', 'ai_restart_default'): + for key in ('ai_predictions_default', 'model_default', 'comparison_output_dir', 'ai_restart_default', 'fallback_data_dir', 'fallback_reference_file', 'fallback_reference_filename'): value = parsed.get(key) if value: defaults[key] = value @@ -316,6 +316,23 @@ def main(): args = parse_arguments() defaults = _load_default_paths(args.variable_list) + + # Set FALLBACK_DATA_DIR and FALLBACK_REFERENCE_FILE from config if available + global FALLBACK_DATA_DIR, FALLBACK_REFERENCE_FILE + if 'fallback_data_dir' in defaults: + FALLBACK_DATA_DIR = defaults['fallback_data_dir'] + + # Priority order for FALLBACK_REFERENCE_FILE: + # 1. Explicit fallback_reference_file in config + # 2. fallback_data_dir + fallback_reference_filename in config + # 3. fallback_data_dir + default filename + if 'fallback_reference_file' in defaults: + FALLBACK_REFERENCE_FILE = defaults['fallback_reference_file'] + elif 'fallback_data_dir' in defaults and 'fallback_reference_filename' in defaults: + FALLBACK_REFERENCE_FILE = FALLBACK_DATA_DIR + defaults['fallback_reference_filename'] + elif 'fallback_data_dir' in defaults: + # Fallback: construct reference file from data dir if not explicitly set + FALLBACK_REFERENCE_FILE = FALLBACK_DATA_DIR + '20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' def _resolve(value, keys, fallback): if value: From d1f5422d3b2f5ebad1005fa78b63cbaddc5b4051 Mon Sep 17 00:00:00 2001 From: Zhuowei Gu Date: Fri, 17 Oct 2025 13:20:39 -0400 Subject: [PATCH 30/51] Add training_data_generation scripts and validation scripts Remove author section from README Removed author information from README. --- scripts/training_data_generation/README.md | 120 ++++ .../bash_script/run_adding_pft_variables.sh | 59 ++ .../bash_script/run_all_forcing_extraction.sh | 121 ++++ .../run_comprehensive_validation.sh | 60 ++ .../run_enhanced_dataset_generation.sh | 52 ++ .../run_forcing_netcdf_validation.sh | 18 + .../bash_script/run_forcing_pkl_generation.sh | 18 + .../bash_script/run_forcing_pkl_validation.sh | 18 + .../run_incomplete_training_dataset.sh | 34 + scripts/training_data_generation/config.py | 198 ++++++ .../python_scripts/1_add_pft_to_dataset.py | 113 ++++ .../python_scripts/2_rm_variables.py | 112 +++ .../python_scripts/37_dataset.py | 337 +++++++++ .../python_scripts/72_dataset_construction.py | 640 ++++++++++++++++++ .../python_scripts/72_dataset_forcing_only.py | 200 ++++++ .../python_scripts/CNP_IO_updated14_xfer.txt | 62 ++ .../python_scripts/cnp_io_parse.py | 69 ++ .../construct_TVA_FLDS_20years.py | 183 +++++ .../construct_TVA_FSDS_20years.py | 152 +++++ .../construct_TVA_PRECTmms_20years.py | 152 +++++ .../construct_TVA_PSRF_20years.py | 124 ++++ .../construct_TVA_QBOT_20years.py | 152 +++++ .../construct_TVA_TBOT_20years.py | 152 +++++ .../training_data_generation/requirements.txt | 16 + .../validation/comprehensive_validation.py | 519 ++++++++++++++ .../validation/forcing_netcdf_validation.py | 109 +++ .../validation/forcing_pkl_validation.py | 308 +++++++++ 27 files changed, 4098 insertions(+) create mode 100644 scripts/training_data_generation/README.md create mode 100755 scripts/training_data_generation/bash_script/run_adding_pft_variables.sh create mode 100755 scripts/training_data_generation/bash_script/run_all_forcing_extraction.sh create mode 100644 scripts/training_data_generation/bash_script/run_comprehensive_validation.sh create mode 100755 scripts/training_data_generation/bash_script/run_enhanced_dataset_generation.sh create mode 100755 scripts/training_data_generation/bash_script/run_forcing_netcdf_validation.sh create mode 100755 scripts/training_data_generation/bash_script/run_forcing_pkl_generation.sh create mode 100755 scripts/training_data_generation/bash_script/run_forcing_pkl_validation.sh create mode 100755 scripts/training_data_generation/bash_script/run_incomplete_training_dataset.sh create mode 100644 scripts/training_data_generation/config.py create mode 100644 scripts/training_data_generation/python_scripts/1_add_pft_to_dataset.py create mode 100644 scripts/training_data_generation/python_scripts/2_rm_variables.py create mode 100644 scripts/training_data_generation/python_scripts/37_dataset.py create mode 100644 scripts/training_data_generation/python_scripts/72_dataset_construction.py create mode 100644 scripts/training_data_generation/python_scripts/72_dataset_forcing_only.py create mode 100644 scripts/training_data_generation/python_scripts/CNP_IO_updated14_xfer.txt create mode 100644 scripts/training_data_generation/python_scripts/cnp_io_parse.py create mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_FLDS_20years.py create mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_FSDS_20years.py create mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_PRECTmms_20years.py create mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_PSRF_20years.py create mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_QBOT_20years.py create mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_TBOT_20years.py create mode 100644 scripts/training_data_generation/requirements.txt create mode 100644 scripts/training_data_generation/validation/comprehensive_validation.py create mode 100644 scripts/training_data_generation/validation/forcing_netcdf_validation.py create mode 100644 scripts/training_data_generation/validation/forcing_pkl_validation.py diff --git a/scripts/training_data_generation/README.md b/scripts/training_data_generation/README.md new file mode 100644 index 0000000..c9bd829 --- /dev/null +++ b/scripts/training_data_generation/README.md @@ -0,0 +1,120 @@ +# Training Data Generation + +Simple workflow to extract TVA forcing data from monthly NetCDF files. + +## Setup + +### 1. Create Virtual Environment +```bash +python3.11 -m venv venv_py311 +``` + +### 2. Activate Virtual Environment +```bash +source venv_py311/bin/activate +``` + +### 3. Install Dependencies +```bash +pip install -r requirements.txt +``` + +## Usage + +### Complete Workflow + +1. **Generate Forcing NetCDF Files** +```bash +./bash_script/run_all_forcing_extraction.sh +``` +Generates 6 NetCDF files in `output/forcing_netcdf/`: +- TVA_FLDS_1980-1999.nc (Longwave radiation) +- TVA_FSDS_1980-1999.nc (Shortwave radiation) +- TVA_PSRF_1980-1999.nc (Surface pressure) +- TVA_QBOT_1980-1999.nc (Specific humidity) +- TVA_PRECTmms_1980-1999.nc (Precipitation) +- TVA_TBOT_1980-1999.nc (Air temperature) + +2. **Validate NetCDF Data Accuracy** +```bash +./bash_script/run_forcing_netcdf_validation.sh +``` +Compares generated NetCDF files with reference files to ensure correctness. + +3. **Generate PKL Files (Training Ready)** +```bash +./bash_script/run_forcing_pkl_generation.sh +``` +Creates optimized PKL files in `output/forcing_hourly_pkl/` for machine learning: +- 12 batch files (TVA_forcing_batch_01.pkl to TVA_forcing_batch_12.pkl) +- Each batch contains 1000 gridcells (except last batch: 357) +- 9 variables per gridcell: landfrac, lat, lon, 6 forcing variables +- 58,400 time steps (3-hour resolution, 20 years) +- **Automatically converts forcing variables to list format for training compatibility** + +4. **Validate PKL Data Accuracy** +```bash +./bash_script/run_forcing_pkl_validation.sh +``` +Validates PKL files against generated NetCDF files using sequential gridcell mapping: +- Validates first 5 PKL batches (5000 gridcells total) +- Each batch corresponds to sequential NetCDF gridcells (0-999, 1000-1999, etc.) +- Ensures PKL data matches NetCDF data with 100% accuracy +- **PKL files are already in list format and ready for training** + +5. **Generate Complete Training Dataset (Monthly Averaged)** +```bash +./bash_script/run_incomplete_training_dataset.sh +``` +- Integrates ecosystem variables with forcing data +- Applies monthly averaging to forcing variables (240 values for 20 years) +- Output: `output/training_dataset_pkl/monthly_training_data_batch_XX.pkl` +- Automatically removes original PKL files + +6. **Generate Enhanced Dataset** +```bash +./bash_script/run_enhanced_dataset_generation.sh +``` +- Adds pool variables (cpool, npool, ppool, xsmrpool) from restart files +- Adds 38 transfer variables and corresponding Y variables +- Output: `output/enhanced_training_dataset/enhanced_monthly_training_data_batch_XX.pkl` +- Automatically removes intermediate files + +7. **Add PFT Variables** +```bash +./bash_script/run_adding_pft_variables.sh +``` +- Adds PFT (Plant Functional Type) variables from `clm_params_c211124.nc` +- Removes unwanted variables (fire-related, unnecessary PFT variables, SCALARAVG_vr) +- **Final dataset ready for machine learning training** + +## Data Validation + +### When to Validate Your Data + +**After Step 1 (Forcing NetCDF Generation):** +```bash +./validation/forcing_netcdf_validation.py +``` +- Validates generated NetCDF files against reference files +- Ensures 6 forcing variables are correctly processed + +**After Step 4 (Forcing PKL Generation):** +```bash +./bash_script/run_forcing_pkl_validation.sh +``` +- Validates PKL files against generated NetCDF files +- Sequential gridcell mapping validation (first 5 batches) + +**After Step 7 (Final Enhanced Dataset):** +```bash +./bash_script/run_comprehensive_validation.sh +``` +- Complete validation of the final enhanced dataset +- Includes monthly averaging, data consistency, spatial mapping, and scientific validity checks +- **This is the most important validation** - run after completing all processing steps + + +## Configuration + +Edit `config.py` to modify paths and settings. diff --git a/scripts/training_data_generation/bash_script/run_adding_pft_variables.sh b/scripts/training_data_generation/bash_script/run_adding_pft_variables.sh new file mode 100755 index 0000000..7260fd8 --- /dev/null +++ b/scripts/training_data_generation/bash_script/run_adding_pft_variables.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# Adding PFT variables script runner + +echo "==========================================" +echo "TVA Adding PFT Variables" +echo "==========================================" + +# Get the project directory +PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" + +# Activate Python virtual environment +source ${PROJECT_DIR}/venv_py311/bin/activate + +# Check if enhanced dataset exists +if [ ! -d "${PROJECT_DIR}/output/enhanced_training_dataset" ]; then + echo "❌ Error: Enhanced dataset directory not found!" + echo " Please run enhanced dataset generation first:" + echo " ./bash_script/run_enhanced_dataset_generation.sh" + exit 1 +fi + +# Check if CLM parameters file exists +if [ ! -f "${PROJECT_DIR}/clm_params.c130821.nc" ]; then + echo "❌ Error: CLM parameters file not found!" + echo " Expected file: ${PROJECT_DIR}/clm_params.c130821.nc" + echo " Please ensure the CLM parameters file is in the correct location." + exit 1 +fi + +echo "✅ Enhanced dataset directory found" +echo "✅ CLM parameters file found" + +# Step 1: Add PFT variables +echo "" +echo "==========================================" +echo "Step 1: Adding PFT Variables" +echo "==========================================" + +python ${PROJECT_DIR}/python_scripts/1_add_pft_to_dataset.py + +echo "" +echo "✅ PFT variables addition completed!" + +# Step 2: Remove unwanted variables +echo "" +echo "==========================================" +echo "Step 2: Removing Unwanted Variables" +echo "==========================================" + +python ${PROJECT_DIR}/python_scripts/2_rm_variables.py + +echo "" +echo "🎉 Adding PFT variables completed!" +echo " - PFT variables added to enhanced dataset" +echo " - Dataset ready for further processing" + +echo "" +echo "Final enhanced dataset files:" +ls -la ${PROJECT_DIR}/output/enhanced_training_dataset/ diff --git a/scripts/training_data_generation/bash_script/run_all_forcing_extraction.sh b/scripts/training_data_generation/bash_script/run_all_forcing_extraction.sh new file mode 100755 index 0000000..46575bc --- /dev/null +++ b/scripts/training_data_generation/bash_script/run_all_forcing_extraction.sh @@ -0,0 +1,121 @@ +#!/bin/bash +# Run all TVA forcing data extraction scripts +# This script processes 6 forcing variables: FLDS, FSDS, PSRF, QBOT, PRECTmms, TBOT + +echo "==========================================" +echo "TVA Forcing Data Extraction Pipeline" +echo "==========================================" +echo "Processing 6 forcing variables (1980-1999, 20 years)" +echo "Output directory: /gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation/output/forcing_netcdf" +echo "==========================================" + +# Get the directory where this script is located +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(dirname "$SCRIPT_DIR")" + +# Change to project directory +cd "$PROJECT_DIR" + +# Activate Python virtual environment +echo "Activating Python virtual environment..." +source venv_py311/bin/activate + +# Create output directory +echo "Creating output directory..." +mkdir -p output/forcing_netcdf + +# Create log directory +mkdir -p logs + +echo "" +echo "Starting forcing data extraction..." +echo "" + +# Initialize counters +success_count=0 +total_count=6 + +# Function to run a script and check result +run_script() { + local script_name="$1" + local variable="$2" + local log_file="logs/${variable}_extraction.log" + + echo "==========================================" + echo "Processing $variable forcing data" + echo "==========================================" + echo "Script: $script_name" + echo "Log file: $log_file" + echo "" + + # Run the script and capture output + python "python_scripts/$script_name" > "$log_file" 2>&1 + + # Check exit status + if [ $? -eq 0 ]; then + echo "✓ $variable forcing data extraction completed successfully" + ((success_count++)) + else + echo "✗ $variable forcing data extraction failed" + echo " Check log file: $log_file" + fi + + echo "" +} + +# Run each forcing variable extraction script +run_script "construct_TVA_FLDS_20years.py" "FLDS" +run_script "construct_TVA_FSDS_20years.py" "FSDS" +run_script "construct_TVA_PSRF_20years.py" "PSRF" +run_script "construct_TVA_QBOT_20years.py" "QBOT" +run_script "construct_TVA_PRECTmms_20years.py" "PRECTmms" +run_script "construct_TVA_TBOT_20years.py" "TBOT" + +echo "==========================================" +echo "All Tasks Completed!" +echo "==========================================" +echo "Successfully processed: $success_count/$total_count variables" + +# List generated files +echo "" +echo "Generated NetCDF files:" +if [ -d "output/forcing_netcdf" ]; then + ls -lh output/forcing_netcdf/*.nc 2>/dev/null || echo "No NetCDF files found" +else + echo "Output directory not found" +fi + +echo "" +echo "Log files:" +if [ -d "logs" ]; then + ls -lh logs/*.log 2>/dev/null || echo "No log files found" +else + echo "Log directory not found" +fi + +echo "" +echo "==========================================" +echo "Summary" +echo "==========================================" +echo "Variables processed:" +echo " FLDS - Longwave radiation" +echo " FSDS - Shortwave radiation" +echo " PSRF - Surface pressure" +echo " QBOT - Specific humidity" +echo " PRECTmms - Precipitation" +echo " TBOT - Air temperature" +echo "" +echo "Output location: $PROJECT_DIR/output/forcing_netcdf/" +echo "Log location: $PROJECT_DIR/logs/" +echo "==========================================" + +# Exit with error code if any script failed +if [ $success_count -ne $total_count ]; then + echo "WARNING: Some extractions failed. Check log files for details." + exit 1 +else + echo "✓ All forcing data extractions completed successfully!" + exit 0 +fi + + diff --git a/scripts/training_data_generation/bash_script/run_comprehensive_validation.sh b/scripts/training_data_generation/bash_script/run_comprehensive_validation.sh new file mode 100644 index 0000000..7433c90 --- /dev/null +++ b/scripts/training_data_generation/bash_script/run_comprehensive_validation.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# Comprehensive Enhanced Dataset Validation Script + +echo "==========================================" +echo "Comprehensive Enhanced Dataset Validation" +echo "==========================================" + +# Get the project directory +PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" + +# Activate Python virtual environment +source ${PROJECT_DIR}/venv_py311/bin/activate + +echo "Running comprehensive validation of enhanced dataset..." +echo "" +echo "This validation includes:" +echo "1. ✅ Forcing data monthly averaging verification (240 values for 20 years)" +echo "2. ✅ Data existence and format validation" +echo "3. ✅ Pool data consistency with restart file (spatial mapping)" +echo "4. ✅ PFT data consistency with CLM parameters (value-by-value)" +echo "5. ✅ Forcing data consistency and scientific validity" +echo "6. ✅ History vs restart file comparison" +echo "7. ✅ Gridcell-by-gridcell validation (first 5000 gridcells)" +echo "8. ✅ Spatial mapping verification" +echo "" + +# Run the comprehensive validation script +python ${PROJECT_DIR}/validation/comprehensive_validation.py + +# Capture exit code +validation_result=$? + +echo "" +echo "==========================================" +if [ $validation_result -eq 0 ]; then + echo "🎉 COMPREHENSIVE VALIDATION COMPLETED SUCCESSFULLY!" + echo "✅ Enhanced dataset is completely validated and ready for use" + echo "" + echo "Validation Summary:" + echo " - Monthly averaging: ✅ Verified (240 values for 20 years)" + echo " - Data existence: ✅ Verified" + echo " - Pool data consistency: ✅ Verified (spatial mapping)" + echo " - PFT data consistency: ✅ Verified (value-by-value)" + echo " - Forcing data consistency: ✅ Verified" + echo " - History/Restart comparison: ✅ Verified" + echo " - Gridcell mapping: ✅ Verified (5000 gridcells)" + echo " - Data integrity: ✅ Verified" + echo "" + echo "Your enhanced dataset is scientifically accurate and ready for machine learning!" +else + echo "❌ COMPREHENSIVE VALIDATION FAILED!" + echo "⚠️ Enhanced dataset needs review" + echo "" + echo "Please check the validation output above for details." + echo "Some validations may have failed and need attention." +fi +echo "==========================================" + +exit $validation_result + diff --git a/scripts/training_data_generation/bash_script/run_enhanced_dataset_generation.sh b/scripts/training_data_generation/bash_script/run_enhanced_dataset_generation.sh new file mode 100755 index 0000000..99d0fc4 --- /dev/null +++ b/scripts/training_data_generation/bash_script/run_enhanced_dataset_generation.sh @@ -0,0 +1,52 @@ +#!/bin/bash +# Enhanced dataset generation script runner with automatic cleanup + +echo "==========================================" +echo "TVA Enhanced Dataset Generation with Cleanup" +echo "==========================================" + +# Get the project directory +PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" + +# Activate Python virtual environment +source ${PROJECT_DIR}/venv_py311/bin/activate + +# Create output directory if it doesn't exist +mkdir -p ${PROJECT_DIR}/output/enhanced_training_dataset + +# Run enhanced dataset generation script +echo "" +echo "Step 1: Running enhanced dataset generation (37_dataset.py)..." +python ${PROJECT_DIR}/python_scripts/37_dataset.py + +echo "" +echo "✅ Enhanced dataset generation completed!" + +# Clean up intermediate files +echo "" +echo "==========================================" +echo "Step 2: Cleaning Up Intermediate Files" +echo "==========================================" + +echo "Current directory contents:" +ls -la ${PROJECT_DIR}/output/ + +echo "" +echo "Removing training_dataset_pkl directory..." +rm -rf ${PROJECT_DIR}/output/training_dataset_pkl/ + +echo "" +echo "✅ Cleanup completed!" +echo "Remaining directories:" +ls -la ${PROJECT_DIR}/output/ + +echo "" +echo "Final enhanced dataset files:" +ls -la ${PROJECT_DIR}/output/enhanced_training_dataset/ + +echo "" +echo "🎉 Enhanced dataset generation and cleanup completed!" +echo " - Enhanced dataset files ready for training" +echo " - Intermediate files removed" +echo " - Disk space saved: ~2GB" +echo " - Only enhanced dataset files remain for training use." diff --git a/scripts/training_data_generation/bash_script/run_forcing_netcdf_validation.sh b/scripts/training_data_generation/bash_script/run_forcing_netcdf_validation.sh new file mode 100755 index 0000000..70750d1 --- /dev/null +++ b/scripts/training_data_generation/bash_script/run_forcing_netcdf_validation.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Forcing data validation script runner + +echo "==========================================" +echo "TVA Forcing Data Validation" +echo "==========================================" + +# Get the project directory +PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" + +# Activate Python virtual environment +source ${PROJECT_DIR}/venv_py311/bin/activate + +# Run validation script +python ${PROJECT_DIR}/validation/forcing_netcdf_validation.py + +echo "" +echo "Validation completed!" diff --git a/scripts/training_data_generation/bash_script/run_forcing_pkl_generation.sh b/scripts/training_data_generation/bash_script/run_forcing_pkl_generation.sh new file mode 100755 index 0000000..744c00b --- /dev/null +++ b/scripts/training_data_generation/bash_script/run_forcing_pkl_generation.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Forcing data PKL generation script runner + +echo "==========================================" +echo "TVA Forcing Data PKL Generation" +echo "==========================================" + +# Get the project directory +PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" + +# Activate Python virtual environment +source ${PROJECT_DIR}/venv_py311/bin/activate + +# Run forcing PKL generation script +python ${PROJECT_DIR}/python_scripts/72_dataset_forcing_only.py + +echo "" +echo "Forcing PKL generation completed!" diff --git a/scripts/training_data_generation/bash_script/run_forcing_pkl_validation.sh b/scripts/training_data_generation/bash_script/run_forcing_pkl_validation.sh new file mode 100755 index 0000000..291b7de --- /dev/null +++ b/scripts/training_data_generation/bash_script/run_forcing_pkl_validation.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Forcing PKL validation script runner + +echo "==========================================" +echo "TVA Forcing PKL Validation" +echo "==========================================" + +# Get the project directory +PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" + +# Activate Python virtual environment +source ${PROJECT_DIR}/venv_py311/bin/activate + +# Run forcing PKL validation script +python ${PROJECT_DIR}/validation/forcing_pkl_validation.py + +echo "" +echo "Forcing PKL validation completed!" diff --git a/scripts/training_data_generation/bash_script/run_incomplete_training_dataset.sh b/scripts/training_data_generation/bash_script/run_incomplete_training_dataset.sh new file mode 100755 index 0000000..90bcc1b --- /dev/null +++ b/scripts/training_data_generation/bash_script/run_incomplete_training_dataset.sh @@ -0,0 +1,34 @@ +#!/bin/bash +# Incomplete training dataset generation script runner + +echo "==========================================" +echo "TVA Incomplete Training Dataset Generation" +echo "==========================================" + +# Get the project directory +PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" + +# Activate Python virtual environment +source ${PROJECT_DIR}/venv_py311/bin/activate + +# Run incomplete training dataset generation script +python ${PROJECT_DIR}/python_scripts/72_dataset_construction.py + +echo "" +echo "Incomplete training dataset generation completed!" + +# Clean up original PKL files (keep only monthly averaged files) +echo "" +echo "Cleaning up original PKL files..." +echo "Keeping only monthly averaged files..." + +# Remove original training_data_batch_*.pkl files +rm -f ${PROJECT_DIR}/output/training_dataset_pkl/training_data_batch_*.pkl + +echo "✅ Original PKL files removed" +echo "✅ Only monthly averaged files retained" +echo "" +echo "🎉 Incomplete training dataset generation completed!" +echo " - Monthly averaged PKL files created in output/training_dataset_pkl/" +echo " - Ready for enhanced dataset generation (37_dataset.py)" +echo " - Use run_enhanced_dataset_generation.sh to complete the workflow" diff --git a/scripts/training_data_generation/config.py b/scripts/training_data_generation/config.py new file mode 100644 index 0000000..0fa18a7 --- /dev/null +++ b/scripts/training_data_generation/config.py @@ -0,0 +1,198 @@ +""" +Configuration for Training Data Generation - TVA Dataset +""" + +import os + +try: + from cnp_io_parse import parse_cnp_io_list +except Exception: + parse_cnp_io_list = None + +# ============================================================================= +# DATA PATHS CONFIGURATION +# ============================================================================= + +# Raw forcing data directory (contains monthly NetCDF files) +forcing_raw_data_path = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/forcing' + +# Base output directory +output_dir = './output' + +# Processed forcing NetCDF files output directory +forcing_netcdf_output_dir = './output/forcing_netcdf' + +# Forcing PKL files output directory +forcing_pkl_output_dir = './output/forcing_hourly_pkl' + +# Training dataset PKL files output directory +training_dataset_pkl_output_dir = './output/training_dataset_pkl' + +# CLM parameters NetCDF file path +clm_params_nc_path = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/clm_params_c211124.nc' + +# ============================================================================= +# 37_DATASET CONFIGURATION (Enhanced Dataset Generation) +# ============================================================================= + +class Config: + # Input paths for 37_dataset.py + INPUT_GLOB = "./output/training_dataset_pkl/monthly_training_data_batch_*.pkl" + OUTPUT_DIR = "./output/enhanced_training_dataset" + ENHANCED_PREFIX = "enhanced_" + POOL_VARS = ["cpool", "npool", "ppool", "xsmrpool"] + + # Special P input NetCDF file + SPECIAL_P_INPUT_NC = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" + SPECIAL_P_VARS = [] + + # CNP IO configuration file + CNP_IO_FILE = os.path.join(os.path.dirname(__file__), "python_scripts", "CNP_IO_updated14_xfer.txt") + + # Columns to drop during processing + COLS_TO_DROP = [ + 'H2OSFC', 'H2OSNO', 'H2OSOI_LIQ', 'H2OSOI_ICE', 'LAKE_SOILC', 'H2OCAN', + 'TH2OSFC', 'T_GRND', 'T_GRND_R', 'T_GRND_U', 'T_LAKE', 'T_SOISNO', + 'TS_TOPO', 'taf', 'T_VEG', 'T10_VALUE', + 'Y_H2OSFC', 'Y_H2OSNO', 'Y_H2OSOI_LIQ', 'Y_H2OSOI_ICE', 'Y_LAKE_SOILC', 'Y_H2OCAN', + 'Y_TH2OSFC', 'Y_T_GRND', 'Y_T_GRND_R', 'Y_T_GRND_U', 'Y_T_LAKE', 'Y_T_SOISNO', + 'Y_TS_TOPO', 'Y_taf', 'Y_T_VEG', 'Y_T10_VALUE', + 'annsum_npp', 'avail_retransn', 'avail_retransp', 'cannsum_npp', + 'Y_annsum_npp', 'Y_avail_retransn', 'Y_avail_retransp', 'Y_cannsum_npp', + 'leafc_xfer', 'frootc_xfer', 'livestemc_xfer', 'deadstemc_xfer', 'livecrootc_xfer', 'deadcrootc_xfer', + 'gresp_xfer', 'leafn_xfer', 'frootn_xfer', 'livestemn_xfer', 'deadstemn_xfer', 'livecrootn_xfer', + 'deadcrootn_xfer', 'leafp_xfer', 'frootp_xfer', 'livestemp_xfer', 'deadstemp_xfer', 'livecrootp_xfer', 'deadcrootp_xfer', + 'retransn', 'retransp', 'gresp_storage', + 'Y_leafc_xfer', 'Y_frootc_xfer', 'Y_livestemc_xfer', 'Y_deadstemc_xfer', 'Y_livecrootc_xfer', 'Y_deadcrootc_xfer', + 'Y_gresp_xfer', 'Y_leafn_xfer', 'Y_frootn_xfer', 'Y_livestemn_xfer', 'Y_deadstemn_xfer', 'Y_livecrootn_xfer', + 'Y_deadcrootn_xfer', 'Y_leafp_xfer', 'Y_frootp_xfer', 'Y_livestemp_xfer', 'Y_deadstemp_xfer', 'Y_livecrootp_xfer', 'Y_deadcrootp_xfer', + 'Y_retransn', 'Y_retransp', 'Y_gresp_storage', + 'labilep_vr', 'occlp_vr', 'primp_vr', + 'Y_labilep_vr', 'Y_occlp_vr', 'Y_primp_vr', + 'cpool', 'npool', 'ppool', 'xsmrpool', + 'Y_cpool', 'Y_npool', 'Y_ppool', 'Y_xsmrpool', + 'FH2OSFC', + 'Y_FH2OSFC', + 'secondp_vr', + 'Y_secondp_vr' + ] + + # List columns configuration + X_LIST_COLUMNS_2D = [ + 'soil3c_vr', 'soil4c_vr', 'cwdc_vr', 'cwdn_vr', 'secondp_vr', 'cwdp', 'totcolp', 'totlitc', 'cwdp_vr', + 'soil1c_vr', 'soil1n_vr', 'soil1p_vr', + 'soil2c_vr', 'soil2n_vr', 'soil2p_vr', + 'soil3n_vr', 'soil3p_vr', + 'soil4n_vr', 'soil4p_vr', + 'litr1c_vr', 'litr2c_vr', 'litr3c_vr', + 'litr1n_vr', 'litr2n_vr', 'litr3n_vr', + 'litr1p_vr', 'litr2p_vr', 'litr3p_vr', + 'sminn_vr', 'smin_no3_vr', 'smin_nh4_vr', + ] + + X_LIST_COLUMNS_1D = [ + 'deadcrootc', 'deadstemc', 'tlai', 'totvegc', 'deadstemn', 'deadcrootn', 'deadstemp', 'deadcrootp', + 'leafc', 'leafc_storage', 'frootc', 'frootc_storage', + 'leafn', 'leafn_storage', 'frootn', 'frootn_storage', + 'leafp', 'leafp_storage', 'frootp', 'frootp_storage', + 'livestemc', 'livestemc_storage', 'livestemn', 'livestemn_storage', + 'livestemp', 'livestemp_storage', 'deadcrootc_storage', 'deadstemc_storage', + 'livecrootc', 'livecrootc_storage', 'deadcrootn_storage', 'deadstemn_storage', + 'livecrootn', 'livecrootn_storage', 'deadcrootp_storage', 'deadstemp_storage', + 'livecrootp', 'livecrootp_storage', + ] + + Y_LIST_COLUMNS_2D = [f"Y_{name}" for name in X_LIST_COLUMNS_2D] + Y_LIST_COLUMNS_1D = [f"Y_{name}" for name in X_LIST_COLUMNS_1D] + + WATER_VARIABLES = [] + Y_WATER_VARIABLES = [] + + VARS_TO_RESHAPE = ['cwdp', 'totcolp', 'totlitc', 'Y_cwdp', 'Y_totcolp', 'Y_totlitc'] + + # Additional dataset variables + dataset_new_1D_PFT_VARIABLES: list = [] + dataset_new_Water_variables: list = [] + dataset_new_TIME_SERIES_VARIABLES: list = [] + dataset_new_SURFACE_PROPERTIES: list = [] + dataset_new_PFT_PARAMETERS: list = [] + dataset_new_SCALAR_VARIABLES: list = [] + dataset_new_2D_VARIABLES: list = [] + dataset_new_RESTART_COL_1D_VARS: list = [] + + @classmethod + def apply_cnp_io_overrides(cls) -> None: + try: + if parse_cnp_io_list is None or not os.path.exists(cls.CNP_IO_FILE): + return + + parsed = parse_cnp_io_list(cls.CNP_IO_FILE) + + new_1d_vars = list(dict.fromkeys(parsed.get('pft_1d_variables', []) or [])) + new_2d_vars = list(dict.fromkeys(parsed.get('variables_2d_soil', []) or [])) + new_water_vars = list(dict.fromkeys(parsed.get('water_variables', []) or [])) + + cls.dataset_new_1D_PFT_VARIABLES = list(dict.fromkeys( + (parsed.get('dataset_new_1D_PFT_VARIABLES') or parsed.get('pft_1d_variables') or []) + )) + cls.dataset_new_Water_variables = list(dict.fromkeys( + (parsed.get('dataset_new_Water_variables') or parsed.get('water_variables') or []) + )) + cls.dataset_new_TIME_SERIES_VARIABLES = list(dict.fromkeys( + (parsed.get('dataset_new_TIME_SERIES_VARIABLES') or []) + )) + cls.dataset_new_SURFACE_PROPERTIES = list(dict.fromkeys( + (parsed.get('dataset_new_SURFACE_PROPERTIES') or []) + )) + cls.dataset_new_PFT_PARAMETERS = list(dict.fromkeys( + (parsed.get('dataset_new_PFT_PARAMETERS') or []) + )) + cls.dataset_new_SCALAR_VARIABLES = list(dict.fromkeys( + (parsed.get('dataset_new_SCALAR_VARIABLES') or []) + )) + cls.dataset_new_2D_VARIABLES = list(dict.fromkeys( + (parsed.get('dataset_new_2D_VARIABLES') or parsed.get('variables_2d_soil') or []) + )) + cls.dataset_new_RESTART_COL_1D_VARS = list(dict.fromkeys( + (parsed.get('dataset_new_RESTART_COL_1D_VARS') or []) + )) + + if new_1d_vars: + cls.POOL_VARS = new_1d_vars + + if new_1d_vars: + cls.X_LIST_COLUMNS_1D = list(dict.fromkeys(list(cls.X_LIST_COLUMNS_1D) + new_1d_vars)) + if new_2d_vars: + cls.X_LIST_COLUMNS_2D = list(dict.fromkeys(list(cls.X_LIST_COLUMNS_2D) + new_2d_vars)) + if cls.dataset_new_RESTART_COL_1D_VARS: + cls.X_LIST_COLUMNS_2D = list(dict.fromkeys(list(cls.X_LIST_COLUMNS_2D) + cls.dataset_new_RESTART_COL_1D_VARS)) + + if new_water_vars: + cls.WATER_VARIABLES = new_water_vars + cls.X_LIST_COLUMNS_2D = list(dict.fromkeys(list(cls.X_LIST_COLUMNS_2D) + new_water_vars)) + + cls.Y_LIST_COLUMNS_1D = [f"Y_{name}" for name in cls.X_LIST_COLUMNS_1D] + cls.Y_LIST_COLUMNS_2D = [f"Y_{name}" for name in cls.X_LIST_COLUMNS_2D] + cls.Y_WATER_VARIABLES = [f"Y_{name}" for name in cls.WATER_VARIABLES] + + if cls.WATER_VARIABLES: + keep_set = set(cls.WATER_VARIABLES) | set(cls.Y_WATER_VARIABLES) + cls.COLS_TO_DROP = [c for c in cls.COLS_TO_DROP if c not in keep_set] + + cls.VARIABLES_1D = cls.X_LIST_COLUMNS_1D.copy() + cls.VARIABLES_2D = cls.X_LIST_COLUMNS_2D.copy() + + if cls.dataset_new_RESTART_COL_1D_VARS: + restart_vars = cls.dataset_new_RESTART_COL_1D_VARS + restart_y_vars = [f"Y_{var}" for var in restart_vars] + cls.VARS_TO_RESHAPE = list(dict.fromkeys(cls.VARS_TO_RESHAPE + restart_vars + restart_y_vars)) + except Exception: + return + +# Apply CNP IO overrides +Config.apply_cnp_io_overrides() + +Config.num_all_columns_2D = len(Config.X_LIST_COLUMNS_2D) +Config.num_all_columns_1D = len(Config.X_LIST_COLUMNS_1D) + + diff --git a/scripts/training_data_generation/python_scripts/1_add_pft_to_dataset.py b/scripts/training_data_generation/python_scripts/1_add_pft_to_dataset.py new file mode 100644 index 0000000..29d5eb5 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/1_add_pft_to_dataset.py @@ -0,0 +1,113 @@ +import netCDF4 as nc +import numpy as np +import pandas as pd +import glob +import os +import sys + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from config import clm_params_nc_path, training_dataset_pkl_output_dir + +# === 1. Read PFT vectors for all variables from NetCDF === +print("Reading CLM parameters NetCDF file...") +print(f"File path: {clm_params_nc_path}") + +if not os.path.exists(clm_params_nc_path): + print(f"❌ Error: CLM parameters file not found: {clm_params_nc_path}") + sys.exit(1) + +ds = nc.Dataset(clm_params_nc_path) + +target_vars = [ + "aleaff", "allconsl", "allconss", "arootf", "arooti", "astemf", "baset", "bfact", "c3psn", "cc_dstem", + "cc_leaf", "cc_lstem", "cc_other", "croot_stem", "crop", "deadwdcn", "declfact", "displar", "dleaf", "dsladlai", + "evergreen", "fcur", "fcurdv", "fd_pft", "fertnitro", "ffrootcn", "fleafcn", "fleafi", "flivewd", "flnr", + "fm_droot", "fm_dstem", "fm_leaf", "fm_lroot", "fm_lstem", "fm_other", "fm_root", "fnitr", "fr_fcel", "fr_flab", + "fr_flig", "froot_leaf", "frootcn", "fsr_pft", "fstemcn", "gddmin", "graincn", "grnfill", "grperc", "grpnow", + "hybgdd", "irrigated", "laimx", "leaf_long", "leafcn", "lf_fcel", "lf_flab", "lf_flig", "lfemerg", "lflitcn", + "livewdcn", "mxtmp", "pconv", "pftpar20", "pftpar28", "pftpar29", "pftpar30", "pftpar31", "planting_temp", + "pprod10", "pprod100", "pprodharv10", "rholnir", "rholvis", "rhosnir", "rhosvis", "roota_par", "rootb_par", + "rootprof_beta", "season_decid", "slatop", "smpsc", "smpso", "stem_leaf", "stress_decid", "taulnir", "taulvis", + "tausnir", "tausvis", "woody", "xl", "z0mr", "ztopmx" +] + +broadcast_feature_dict = {} +for var in target_vars: + if var in ds.variables: + raw_vals = ds.variables[var][:17] + # Skip variables if any value is NaN or masked (missing) + if np.any(np.isnan(raw_vals)) or np.ma.is_masked(raw_vals): + print(f"Skipped {var}: contains NaN or masked values") + continue + broadcast_feature_dict[var] = list(map(float, raw_vals)) + print(f"Added: {var} (length {len(raw_vals)})") + else: + print(f"Skipped {var}: not found in NetCDF") + +print(f"\n✅ Successfully loaded {len(broadcast_feature_dict)} PFT variables from NetCDF") + +# === 2. Process all PKL files in enhanced_training_dataset directory === +enhanced_output_dir = os.path.join(os.path.dirname(training_dataset_pkl_output_dir), "enhanced_training_dataset") +input_files = sorted(glob.glob(os.path.join(enhanced_output_dir, "enhanced_monthly_training_data_batch_*.pkl"))) + +print(f"\n🔍 Found {len(input_files)} PKL files to process") +print("📋 File list:") +for i, file in enumerate(input_files, 1): + print(f" {i:2d}. {os.path.basename(file)}") + +if len(input_files) == 0: + print("❌ No enhanced PKL files found. Please run enhanced dataset generation first.") + sys.exit(1) + +# === 3. Process each PKL file individually === +for i, file_path in enumerate(input_files, 1): + print(f"\n{'='*80}") + print(f"Processing file {i}/{len(input_files)}: {os.path.basename(file_path)}") + print(f"{'='*80}") + + try: + # Read PKL file + print("📖 Reading PKL file...") + df = pd.read_pickle(file_path) + original_shape = df.shape + print(f"✅ File loaded successfully, original shape: {original_shape}") + + # Check if PFT variables already exist + existing_pft_cols = [col for col in df.columns if col.startswith("pft_")] + if existing_pft_cols: + print(f"⚠️ File already contains {len(existing_pft_cols)} PFT variables, skipping addition") + print(f" Existing PFT variables: {existing_pft_cols[:5]}...") + continue + + # Add each variable as a vector column with pft_ prefix + print("🔧 Adding PFT variables...") + for var, val_list in broadcast_feature_dict.items(): + df["pft_" + var] = [val_list] * len(df) # Add the same list to each row + + new_shape = df.shape + print(f"✅ Successfully added {len(broadcast_feature_dict)} PFT variables") + print(f"📐 New data shape: {original_shape} → {new_shape}") + + # Save in-place (overwrite original file) + print("💾 Saving modified file...") + df.to_pickle(file_path) + print(f"✅ File saved: {os.path.basename(file_path)}") + + # Display sample PFT variables + pft_cols = [col for col in df.columns if col.startswith("pft_")] + if pft_cols: + print("🧾 Sample PFT variables:") + for col in pft_cols[:3]: + print(f" {col}: {df[col].iloc[0]}") + + except Exception as e: + print(f"❌ Failed to process file: {e}") + continue + +print(f"\n{'='*80}") +print("🎉 All files processed successfully!") +print("📊 Summary:") +print(f" - Total files: {len(input_files)}") +print(f" - PFT variables added: {len(broadcast_feature_dict)}") +print("="*80) \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/2_rm_variables.py b/scripts/training_data_generation/python_scripts/2_rm_variables.py new file mode 100644 index 0000000..69362c8 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/2_rm_variables.py @@ -0,0 +1,112 @@ +import pandas as pd +import glob +import os +import sys + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from config import training_dataset_pkl_output_dir + +# === 1. Define variables to delete === +# Delete columns starting with SCALARAVG_vr +scalaravg_cols = [col for col in [] if col.startswith("SCALARAVG_vr")] # This list will be dynamically generated at runtime + +# Delete COL_FIRE_CLOSS and Y_COL_FIRE_CLOSS +fire_cols = ["COL_FIRE_CLOSS", "Y_COL_FIRE_CLOSS"] + +# Delete unwanted PFT variables +unwanted_pft_cols = [ + "pft_aleaff", "pft_baset", "pft_cc_dstem", "pft_cc_leaf", "pft_cc_lstem", "pft_cc_other", "pft_displar", + "pft_fcurdv", "pft_fd_pft", "pft_fertnitro", "pft_ffrootcn", "pft_fleafcn", "pft_fm_droot", "pft_fm_dstem", + "pft_fm_leaf", "pft_fm_lroot", "pft_fm_lstem", "pft_fm_other", "pft_fm_root", "pft_fnitr", "pft_fsr_pft", + "pft_fstemcn", "pft_irrigated", "pft_pconv", "pft_pftpar20", "pft_pftpar28", "pft_pftpar29", "pft_pftpar30", + "pft_pftpar31", "pft_pprod10", "pft_pprod100", "pft_pprodharv10" +] + +# Combine all columns to delete +all_drop_cols = fire_cols + unwanted_pft_cols + +print(f"🗑️ Variables to delete:") +print(f" - Fire-related variables: {len(fire_cols)} variables") +print(f" - Unwanted PFT variables: {len(unwanted_pft_cols)} variables") +print(f" - Total: {len(all_drop_cols)} variables") + +# === 2. Process all PKL files in enhanced_training_dataset directory === +enhanced_output_dir = os.path.join(os.path.dirname(training_dataset_pkl_output_dir), "enhanced_training_dataset") +input_files = sorted(glob.glob(os.path.join(enhanced_output_dir, "*.pkl"))) + +print(f"\n🔍 Found {len(input_files)} PKL files to process") +print("📋 File list:") +for i, file in enumerate(input_files, 1): + print(f" {i:2d}. {os.path.basename(file)}") + +if len(input_files) == 0: + print("❌ No enhanced PKL files found. Please run enhanced dataset generation first.") + sys.exit(1) + +# === 3. Process each PKL file individually === +for i, file_path in enumerate(input_files, 1): + print(f"\n{'='*80}") + print(f"Processing file {i}/{len(input_files)}: {os.path.basename(file_path)}") + print(f"{'='*80}") + + try: + # Read PKL file + print("📖 Reading PKL file...") + df = pd.read_pickle(file_path) + original_shape = df.shape + print(f"✅ File loaded successfully, original shape: {original_shape}") + + # Dynamically find columns starting with SCALARAVG_vr + scalaravg_cols = [col for col in df.columns if col.startswith("SCALARAVG_vr")] + if scalaravg_cols: + print(f"🔍 Found {len(scalaravg_cols)} SCALARAVG_vr variables: {scalaravg_cols[:5]}...") + + # Combine all columns to delete + drop_cols = scalaravg_cols + all_drop_cols + + # Check which variables actually exist + existing_drop_cols = [col for col in drop_cols if col in df.columns] + missing_cols = [col for col in drop_cols if col not in df.columns] + + if missing_cols: + print(f"⚠️ Following variables do not exist in file: {missing_cols[:5]}...") + + if existing_drop_cols: + print(f"🗑️ Preparing to delete {len(existing_drop_cols)} variables") + + # Delete variables + df.drop(columns=existing_drop_cols, inplace=True, errors='ignore') + + new_shape = df.shape + print(f"✅ Successfully deleted {len(existing_drop_cols)} variables") + print(f"📐 New data shape: {original_shape} → {new_shape}") + + # Save in-place (overwrite original file) + print("💾 Saving modified file...") + df.to_pickle(file_path) + print(f"✅ File saved: {os.path.basename(file_path)}") + + # Display deleted variables statistics + print("📊 Deleted variables statistics:") + if scalaravg_cols: + print(f" - SCALARAVG_vr variables: {len([col for col in scalaravg_cols if col in existing_drop_cols])} variables") + print(f" - Fire-related variables: {len([col for col in fire_cols if col in existing_drop_cols])} variables") + print(f" - Unwanted PFT variables: {len([col for col in unwanted_pft_cols if col in existing_drop_cols])} variables") + + else: + print("ℹ️ No variables found to delete, skipping processing") + + except Exception as e: + print(f"❌ Failed to process file: {e}") + continue + +print(f"\n{'='*80}") +print("🎉 All files processed successfully!") +print("📊 Summary:") +print(f" - Total files: {len(input_files)}") +print(f" - Deleted variable types:") +print(f" * SCALARAVG_vr variables (dynamically detected)") +print(f" * Fire-related variables: {len(fire_cols)} variables") +print(f" * Unwanted PFT variables: {len(unwanted_pft_cols)} variables") +print("="*80) \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/37_dataset.py b/scripts/training_data_generation/python_scripts/37_dataset.py new file mode 100644 index 0000000..f5371ed --- /dev/null +++ b/scripts/training_data_generation/python_scripts/37_dataset.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +import os +import sys +import glob +import argparse +from typing import Dict, List, Tuple + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from config import Config + +import numpy as np +import pandas as pd +from scipy.spatial import cKDTree +import netCDF4 as nc + +VARS_TO_CLEAN = {'H2OSOI_LIQ', 'H2OSOI_ICE'} +FILL_VALUE_THRESHOLD = 1e35 + +def build_restart_kdtree(ds_restart: nc.Dataset) -> Tuple[cKDTree, np.ndarray]: + gridcell_lat = ds_restart.variables["grid1d_lat"][:] + gridcell_lon = ds_restart.variables["grid1d_lon"][:] + coords = np.vstack((gridcell_lat, gridcell_lon)).T + tree = cKDTree(coords) + return tree, coords + +def build_column_index_map(ds_restart: nc.Dataset) -> Dict[int, np.ndarray]: + cols1d_gridcell_index = ds_restart.variables["cols1d_gridcell_index"][:] + unique_ids = np.unique(cols1d_gridcell_index) + mapping: Dict[int, np.ndarray] = {} + for grid_id in unique_ids: + mapping[int(grid_id)] = np.where(cols1d_gridcell_index == grid_id)[0] + return mapping + +def ensure_vars_exist(ds: nc.Dataset, var_names: List[str]) -> List[str]: + existing = [] + for name in var_names: + if name in ds.variables: + existing.append(name) + return existing + +def build_pft_index_map(ds_restart: nc.Dataset) -> Dict[int, np.ndarray]: + pfts1d_gridcell_index = ds_restart.variables["pfts1d_gridcell_index"][:] + unique_ids = np.unique(pfts1d_gridcell_index) + mapping: Dict[int, np.ndarray] = {} + for grid_id in unique_ids: + mapping[int(grid_id)] = np.where(pfts1d_gridcell_index == grid_id)[0] + return mapping + +def extract_col1d_x(ds_restart: nc.Dataset, var_name: str, col_indices: np.ndarray) -> List[float]: + if col_indices.size == 0: + return [] + values = ds_restart.variables[var_name][col_indices] + return values.astype(float).tolist() + +def extract_col1d_y(ds_r_list: List[nc.Dataset], var_name: str, col_indices: np.ndarray) -> List[float]: + if col_indices.size == 0: + return [] + slices: List[np.ndarray] = [] + for ds_r in ds_r_list: + values = ds_r.variables[var_name][col_indices] + slices.append(np.asarray(values, dtype=float)) + stacked = np.stack(slices, axis=0) + avg = np.mean(stacked, axis=0) + return avg.tolist() + +def extract_col2d_x(ds_restart: nc.Dataset, var_name: str, col_indices: np.ndarray) -> List[List[float]]: + if col_indices.size == 0: + return [] + values = ds_restart.variables[var_name][col_indices, :] + values_np = np.asarray(values, dtype=float) + if var_name in VARS_TO_CLEAN: + values_np[values_np >= FILL_VALUE_THRESHOLD] = 0.0 + return values_np.tolist() + +def extract_col2d_y(ds_r_list: List[nc.Dataset], var_name: str, col_indices: np.ndarray) -> List[List[float]]: + if col_indices.size == 0: + return [] + slices: List[np.ndarray] = [] + for ds_r in ds_r_list: + values = ds_r.variables[var_name][col_indices, :] + values_np = np.asarray(values, dtype=float) + if var_name in VARS_TO_CLEAN: + values_np[values_np >= FILL_VALUE_THRESHOLD] = 0.0 + slices.append(values_np) + stacked = np.stack(slices, axis=0) + avg = np.mean(stacked, axis=0) + return avg.tolist() + +def extract_pft1d_x(ds_restart: nc.Dataset, var_name: str, pft_indices: np.ndarray) -> List[float]: + if pft_indices.size == 0: + return [] + values = ds_restart.variables[var_name][pft_indices] + return np.asarray(values, dtype=float).tolist() + +def extract_pft1d_y(ds_r_list: List[nc.Dataset], var_name: str, pft_indices: np.ndarray) -> List[float]: + if pft_indices.size == 0: + return [] + slices: List[np.ndarray] = [] + for ds_r in ds_r_list: + values = ds_r.variables[var_name][pft_indices] + slices.append(np.asarray(values, dtype=float)) + stacked = np.stack(slices, axis=0) + avg = np.mean(stacked, axis=0) + return avg.tolist() + +def extract_pft2d_x(ds_restart: nc.Dataset, var_name: str, pft_indices: np.ndarray) -> List[List[float]]: + if pft_indices.size == 0: + return [] + values = ds_restart.variables[var_name][pft_indices, :] + return np.asarray(values, dtype=float).tolist() + +def extract_pft2d_y(ds_r_list: List[nc.Dataset], var_name: str, pft_indices: np.ndarray) -> List[List[float]]: + if pft_indices.size == 0: + return [] + slices: List[np.ndarray] = [] + for ds_r in ds_r_list: + values = ds_r.variables[var_name][pft_indices, :] + slices.append(np.asarray(values, dtype=float)) + stacked = np.stack(slices, axis=0) + avg = np.mean(stacked, axis=0) + return avg.tolist() + +def augment_dataframe_with_pools( + df: pd.DataFrame, + ds_restart: nc.Dataset, + ds_r_list: List[nc.Dataset], + restart_tree: cKDTree, + restart_coords: np.ndarray, + col_index_map: Dict[int, np.ndarray], + pool_vars: List[str], +) -> pd.DataFrame: + if "Latitude" not in df.columns or "Longitude" not in df.columns: + raise ValueError("DataFrame is missing Latitude/Longitude columns for mapping.") + + pool_vars_existing = ensure_vars_exist(ds_restart, pool_vars) + if not pool_vars_existing: + raise ValueError(f"None of the target variables found in restart file: {pool_vars}") + + pool_vars_final: List[str] = [ + v for v in pool_vars_existing if all(v in ds_r.variables for ds_r in ds_r_list) + ] + if not pool_vars_final: + raise ValueError("Target variables do not exist in the set of Y files.") + + latitudes = df["Latitude"].to_numpy() + longitudes = df["Longitude"].to_numpy() + query_coords = np.vstack((latitudes, longitudes)).T + _, nearest_restart_indices = restart_tree.query(query_coords, k=1) + + results_x: Dict[str, List[List[float]]] = {v: [] for v in pool_vars_final} + results_y: Dict[str, List[List[float]]] = {f"Y_{v}": [] for v in pool_vars_final} + + for row_idx, restart_idx in enumerate(nearest_restart_indices): + gridcell_id = int(restart_idx) + 1 + col_indices = col_index_map.get(gridcell_id, np.array([], dtype=int)) + for v in pool_vars_final: + x_vals = extract_col1d_x(ds_restart, v, col_indices) + y_vals = extract_col1d_y(ds_r_list, v, col_indices) + results_x[v].append(x_vals) + results_y[f"Y_{v}"].append(y_vals) + + for v in pool_vars_final: + df[v] = results_x[v] + df[f"Y_{v}"] = results_y[f"Y_{v}"] + + return df + +def augment_dataframe_with_vars( + df: pd.DataFrame, + ds_restart: nc.Dataset, + ds_special_p_restart: nc.Dataset, + ds_r_list: List[nc.Dataset], + restart_tree: cKDTree, + restart_coords: np.ndarray, + col_index_map: Dict[int, np.ndarray], + pft_index_map: Dict[int, np.ndarray], + vars_1d: List[str], + vars_2d: List[str], + special_p_vars: List[str], +) -> pd.DataFrame: + if "Latitude" not in df.columns or "Longitude" not in df.columns: + raise ValueError("DataFrame is missing Latitude/Longitude columns for mapping.") + + vars_1d_existing = ensure_vars_exist(ds_restart, vars_1d) + vars_2d_existing = ensure_vars_exist(ds_restart, vars_2d) + + final_1d: List[str] = [v for v in vars_1d_existing if all(v in ds_r.variables for ds_r in ds_r_list)] + final_2d: List[str] = [v for v in vars_2d_existing if all(v in ds_r.variables for ds_r in ds_r_list)] + + if not final_1d and not final_2d: + return df + + latitudes = df["Latitude"].to_numpy() + longitudes = df["Longitude"].to_numpy() + query_coords = np.vstack((latitudes, longitudes)).T + _, nearest_restart_indices = restart_tree.query(query_coords, k=1) + + results_x_1d: Dict[str, List[List[float]]] = {v: [] for v in final_1d} + results_y_1d: Dict[str, List[List[float]]] = {f"Y_{v}": [] for v in final_1d} + results_x_2d: Dict[str, List[List[List[float]]]] = {v: [] for v in final_2d} + results_y_2d: Dict[str, List[List[List[float]]]] = {f"Y_{v}": [] for v in final_2d} + + for row_idx, restart_idx in enumerate(nearest_restart_indices): + gridcell_id = int(restart_idx) + 1 + col_indices = col_index_map.get(gridcell_id, np.array([], dtype=int)) + pft_indices = pft_index_map.get(gridcell_id, np.array([], dtype=int)) + + for v in final_1d: + var_obj = ds_restart.variables[v] + dims = tuple(var_obj.dimensions) + if "pft" in dims: + x_vals = extract_pft1d_x(ds_restart, v, pft_indices) + y_vals = extract_pft1d_y(ds_r_list, v, pft_indices) + else: + x_vals = extract_col1d_x(ds_restart, v, col_indices) + y_vals = extract_col1d_y(ds_r_list, v, col_indices) + results_x_1d[v].append(x_vals) + results_y_1d[f"Y_{v}"].append(y_vals) + + for v in final_2d: + ds_for_x = ds_special_p_restart if v in special_p_vars else ds_restart + var_obj = ds_for_x.variables[v] + dims = tuple(var_obj.dimensions) + + if "pft" in dims: + x_vals_2d = extract_pft2d_x(ds_for_x, v, pft_indices) + y_vals_2d = extract_pft2d_y(ds_r_list, v, pft_indices) + else: + x_vals_2d = extract_col2d_x(ds_for_x, v, col_indices) + y_vals_2d = extract_col2d_y(ds_r_list, v, col_indices) + + results_x_2d[v].append(x_vals_2d) + results_y_2d[f"Y_{v}"].append(y_vals_2d) + + for v in final_1d: + df[v] = results_x_1d[v] + df[f"Y_{v}"] = results_y_1d[f"Y_{v}"] + for v in final_2d: + df[v] = results_x_2d[v] + df[f"Y_{v}"] = results_y_2d[f"Y_{v}"] + + return df + +def main(): + parser = argparse.ArgumentParser( + description="Augment batched training data with variables from RESTART files." + ) + parser.add_argument( + "--input_glob", + default=Config.INPUT_GLOB, + help="Glob pattern for input pkl batch files." + ) + parser.add_argument( + "--output_dir", + default=Config.OUTPUT_DIR, + help="Output directory for augmented data." + ) + args = parser.parse_args() + + file_path10 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.r.0021-01-01-00000.nc" + file_path17 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" + file_path18 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" + file_path19 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" + file_path20 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" + file_path21 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" + + restart_pft_vars = list(dict.fromkeys(Config.dataset_new_1D_PFT_VARIABLES)) + restart_col_1d_vars = list(dict.fromkeys(Config.dataset_new_RESTART_COL_1D_VARS)) + restart_col_2d_vars = list(dict.fromkeys(list(Config.dataset_new_Water_variables) + list(Config.dataset_new_2D_VARIABLES))) + + configured_vars_1d = list(dict.fromkeys(restart_pft_vars + restart_col_1d_vars)) + configured_vars_2d = list(dict.fromkeys(restart_col_2d_vars + Config.SPECIAL_P_VARS)) + + input_files = sorted(glob.glob(args.input_glob)) + if not input_files: + sys.exit(f"No input files found matching pattern: {args.input_glob}") + + os.makedirs(args.output_dir, exist_ok=True) + + ds_restart = nc.Dataset(file_path10) + ds_special_p_restart = nc.Dataset(Config.SPECIAL_P_INPUT_NC) + ds_r_list = [nc.Dataset(fp) for fp in [file_path17, file_path18, file_path19, file_path20, file_path21]] + + try: + restart_tree, restart_coords = build_restart_kdtree(ds_restart) + col_index_map = build_column_index_map(ds_restart) + pft_index_map = build_pft_index_map(ds_restart) + + for fp in input_files: + df = pd.read_pickle(fp) + + force_replace_vars = set(Config.SPECIAL_P_VARS) + missing_other_vars_1d = [ + v for v in configured_vars_1d + if v not in force_replace_vars and not (v in df.columns and f"Y_{v}" in df.columns) + ] + missing_other_vars_2d = [ + v for v in configured_vars_2d + if v not in force_replace_vars and not (v in df.columns and f"Y_{v}" in df.columns) + ] + vars_to_process_1d = missing_other_vars_1d + vars_to_process_2d = list(set(missing_other_vars_2d).union(force_replace_vars)) + + if not vars_to_process_1d and not vars_to_process_2d: + continue + + df_aug = augment_dataframe_with_vars( + df=df, + ds_restart=ds_restart, + ds_special_p_restart=ds_special_p_restart, + ds_r_list=ds_r_list, + restart_tree=restart_tree, + restart_coords=restart_coords, + col_index_map=col_index_map, + pft_index_map=pft_index_map, + vars_1d=vars_to_process_1d, + vars_2d=vars_to_process_2d, + special_p_vars=Config.SPECIAL_P_VARS + ) + + base_name = os.path.basename(fp) + out_name = f"{Config.ENHANCED_PREFIX}{base_name}" + out_path = os.path.join(args.output_dir, out_name) + df_aug.to_pickle(out_path) + + finally: + ds_restart.close() + ds_special_p_restart.close() + for ds in ds_r_list: + try: + ds.close() + except Exception: + pass + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/72_dataset_construction.py b/scripts/training_data_generation/python_scripts/72_dataset_construction.py new file mode 100644 index 0000000..3a3804e --- /dev/null +++ b/scripts/training_data_generation/python_scripts/72_dataset_construction.py @@ -0,0 +1,640 @@ +#!/usr/bin/env python3 +""" +TVA Complete Training Dataset Generation Script +Generates complete PKL files with forcing data, ecosystem variables, and monthly averaging +""" + +import netCDF4 as nc +import numpy as np +import pandas as pd +from scipy.spatial import cKDTree +import os +import sys +import time + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +print("="*80) +print("TVA Complete Training Dataset Generation") +print("="*80) + +# File paths using config and hardcoded paths +file_path1 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/domain_surfdata/TVA_surfdata.TES_SE.4km.1d.NLCD.c241219.nc' +file_path2 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.h0.0021-01-01-00000.nc' + +# Use generated forcing NetCDF files +file_path4 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_FLDS_1980-1999.nc') +file_path5 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_FSDS_1980-1999.nc') +file_path6 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_PRECTmms_1980-1999.nc') +file_path7 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_PSRF_1980-1999.nc') +file_path8 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_QBOT_1980-1999.nc') +file_path9 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_TBOT_1980-1999.nc') + +file_path10 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.r.0021-01-01-00000.nc' + +file_path12 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/kmELM/e3sm_runs/uELM_TVA_finalspinref/run/uELM_TVA_finalspinref.elm.h0.0781-01.nc' +file_path17 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc' + +# Output directory using config +output_dir = os.path.join(config.output_dir, 'training_dataset_pkl') +os.makedirs(output_dir, exist_ok=True) + +# TVA region coordinates (1D domain) +# TVA: lat [32.33, 37.58], lon [-90.33, -81.71] + +print("Loading NetCDF files...") +start_time = time.time() + +ds1 = nc.Dataset(file_path1) # Surface data (TVA_surfdata.TES_SE.4km.1d.NLCD.c241219.nc) +ds2 = nc.Dataset(file_path2) # History file (uELM_TVA_adspinref.elm.h0.0021-01-01-00000.nc) +ds4 = nc.Dataset(file_path4) # FLDS forcing (TVA_FLDS_1980-1999.nc) +ds5 = nc.Dataset(file_path5) # FSDS forcing (TVA_FSDS_1980-1999.nc) +ds6 = nc.Dataset(file_path6) # PRECTmms forcing (TVA_PRECTmms_1980-1999.nc) +ds7 = nc.Dataset(file_path7) # PSRF forcing (TVA_PSRF_1980-1999.nc) +ds8 = nc.Dataset(file_path8) # QBOT forcing (TVA_QBOT_1980-1999.nc) +ds9 = nc.Dataset(file_path9) # TBOT forcing (TVA_TBOT_1980-1999.nc) +ds10 = nc.Dataset(file_path10) # Restart file (uELM_TVA_adspinref.elm.r.0021-01-01-00000.nc) + +# For Y (future) values - using single file for demo +ds_h0_list = [nc.Dataset(file_path12)] # Future history file (uELM_TVA_finalspinref.elm.h0.0781-01.nc) +ds_r_list = [nc.Dataset(file_path17)] # Future restart file (uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc) + +print(f"✅ All files loaded: {time.time() - start_time:.2f}s") + +# TVA data is 1D (lndgrid), not 2D +lats = ds2.variables['lat'][:] # 1D array +lons = ds2.variables['lon'][:] # 1D array +landmask = ds2.variables['landfrac'][:] # Use landfrac instead of landmask + +# Filter for land gridcells (1D domain) +valid_mask = (landmask > 0) +valid_gridcells = np.where(valid_mask)[0] + +print(f"✅ Spatial filtering completed: {time.time() - start_time:.2f}s") +print(f"Total land gridcells: {len(valid_gridcells)}") +print(f"Latitude range: [{lats.min():.2f}, {lats.max():.2f}]") +print(f"Longitude range: [{lons.min():.2f}, {lons.max():.2f}]") + +batch_size = 1000 +batch_number = 1 + +print(f"\nConfiguration:") +print(f" Batch size: {batch_size}") +print(f" Output directory: {output_dir}") + +print("\nBuilding KDTree index...") +start_time = time.time() + +# Build query coordinates for valid gridcells (1D domain) +query_coords = np.array([(lats[i], lons[i]) for i in valid_gridcells]) + +# Restart file coordinates (1D) +gridcell_lat = ds10.variables['grid1d_lat'][:] +gridcell_lon = ds10.variables['grid1d_lon'][:] +restart_grid_coords = np.vstack((gridcell_lat, gridcell_lon)).T + +restart_tree = cKDTree(restart_grid_coords) +_, all_restart_indices = restart_tree.query(query_coords, k=1) + +# Forcing file coordinates (1D) +forcing_lats = ds4.variables['LATIXY'][:].flatten() +forcing_lons = ds4.variables['LONGXY'][:].flatten() +forcing_grid_coords = np.vstack((forcing_lats, forcing_lons)).T + +forcing_tree = cKDTree(forcing_grid_coords) +_, all_forcing_indices = forcing_tree.query(query_coords, k=1) + +print(f"✅ KDTree index construction completed: {time.time() - start_time:.2f}s") + +# 🚀 KEY OPTIMIZATION: Pre-load all forcing data into memory +print("\n🚀 Pre-loading all forcing data into memory...") +start_time = time.time() + +print(" Loading FLDS data...") +flds_data = ds4.variables['FLDS'][:, 0, :] # 58400 × 11357 +print(f" FLDS shape: {flds_data.shape}, memory: {flds_data.nbytes / 1024**3:.2f} GB") + +print(" Loading PSRF data...") +psrf_data = ds7.variables['PSRF'][:, 0, :] +print(f" PSRF shape: {psrf_data.shape}, memory: {psrf_data.nbytes / 1024**3:.2f} GB") + +print(" Loading FSDS data...") +fsds_data = ds5.variables['FSDS'][:, 0, :] +print(f" FSDS shape: {fsds_data.shape}, memory: {fsds_data.nbytes / 1024**3:.2f} GB") + +print(" Loading QBOT data...") +qbot_data = ds8.variables['QBOT'][:, 0, :] +print(f" QBOT shape: {qbot_data.shape}, memory: {qbot_data.nbytes / 1024**3:.2f} GB") + +print(" Loading PRECTmms data...") +prect_data = ds6.variables['PRECTmms'][:, 0, :] +print(f" PRECTmms shape: {prect_data.shape}, memory: {prect_data.nbytes / 1024**3:.2f} GB") + +print(" Loading TBOT data...") +tbot_data = ds9.variables['TBOT'][:, 0, :] +print(f" TBOT shape: {tbot_data.shape}, memory: {tbot_data.nbytes / 1024**3:.2f} GB") + +total_forcing_memory = (flds_data.nbytes + psrf_data.nbytes + fsds_data.nbytes + + qbot_data.nbytes + prect_data.nbytes + tbot_data.nbytes) / 1024**3 + +print(f"✅ All forcing data pre-loaded: {time.time() - start_time:.2f}s") +print(f" Total memory usage: {total_forcing_memory:.2f} GB") + +# Close forcing NetCDF files (data is now in memory) +ds4.close() +ds5.close() +ds6.close() +ds7.close() +ds8.close() +ds9.close() + +print("✅ Forcing NetCDF files closed, data in memory") + +# Define variable lists +pft_based_vars = [ + 'totvegc', 'deadstemn', 'deadcrootn', 'deadstemp', 'deadcrootp', + 'leafc', 'leafc_storage', 'frootc', 'frootc_storage', + 'deadcrootc', 'deadstemc', 'tlai', + 'leafn', 'leafn_storage', 'frootn','frootn_storage', + 'leafp', 'leafp_storage', 'frootp','frootp_storage', + 'livestemc', 'livestemc_storage', + 'livestemn', 'livestemn_storage', + 'livestemp', 'livestemp_storage', + 'deadcrootc_storage', 'deadstemc_storage', + 'livecrootc', 'livecrootc_storage', + 'deadcrootn_storage', 'deadstemn_storage', + 'livecrootn', 'livecrootn_storage', + 'deadcrootp_storage', 'deadstemp_storage', + 'livecrootp', 'livecrootp_storage' +] + +col_based_1d_vars = ['cwdp', 'totcolp', 'totlitc'] + +col_based_2d_vars = [ + 'cwdn_vr', 'secondp_vr', 'cwdp_vr', 'soil3c_vr', 'soil4c_vr', 'cwdc_vr', + 'soil1c_vr', 'soil1n_vr', 'soil1p_vr', + 'soil2c_vr', 'soil2n_vr', 'soil2p_vr', + 'soil3n_vr', 'soil3p_vr', + 'soil4n_vr', 'soil4p_vr', + 'litr1c_vr', 'litr2c_vr', 'litr3c_vr', + 'litr1n_vr', 'litr2n_vr', 'litr3n_vr', + 'litr1p_vr', 'litr2p_vr', 'litr3p_vr', + 'sminn_vr', 'smin_no3_vr', 'smin_nh4_vr', + 'labilep_vr', 'occlp_vr', 'primp_vr' +] + +all_x_vars = pft_based_vars + col_based_1d_vars + col_based_2d_vars + +# Pre-load X variable data +x_values = {} +for var_name in all_x_vars: + print(f" Loading X variable: {var_name}") + x_values[var_name] = ds10.variables[var_name][:] + +# Pre-load Y variable data +stacked_y_values = {} +for var_name in all_x_vars: + print(f" Loading Y variable: {var_name}") + list_of_arrays = [ds_r.variables[var_name][:] for ds_r in ds_r_list] + stacked_y_values[var_name] = np.stack(list_of_arrays, axis=0) + +print(f"✅ X and Y variables pre-loaded: {time.time() - start_time:.2f}s") + +# Build index mapping +print("\nBuilding index mapping...") +start_time = time.time() + +pft_gridcell_index = ds10.variables['pfts1d_gridcell_index'][:] +column_gridcell_index = ds10.variables['cols1d_gridcell_index'][:] + +pft_map = {} +column_map = {} + +unique_gridcell_ids = np.unique(pft_gridcell_index) +for grid_id in unique_gridcell_ids: + pft_map[grid_id] = np.where(pft_gridcell_index == grid_id)[0] + column_map[grid_id] = np.where(column_gridcell_index == grid_id)[0] + +print(f"✅ Index mapping construction completed: {time.time() - start_time:.2f}s") + +# Process data +print(f"\nStarting to process {len(valid_gridcells)} gridcells...") + +for start_idx in range(0, len(valid_gridcells), batch_size): + end_idx = min(start_idx + batch_size, len(valid_gridcells)) + batch_gridcells = valid_gridcells[start_idx:end_idx] + batch_restart_indices = all_restart_indices[start_idx:end_idx] + batch_forcing_indices = all_forcing_indices[start_idx:end_idx] + + print(f"\nProcessing batch {batch_number}: gridcells {start_idx+1}-{end_idx}") + batch_start_time = time.time() + + data_dict = { + 'landfrac':[], + 'Latitude': [], + 'Longitude': [], + 'FLDS': [], 'PSRF': [], 'FSDS': [], 'QBOT': [], 'PRECTmms': [], 'TBOT': [], + 'LANDFRAC_PFT': [], 'PCT_NATVEG': [], 'AREA': [], 'peatf': [], 'abm': [], + 'SOIL_COLOR': [], 'SOIL_ORDER': [], 'PCT_NAT_PFT': [], 'PCT_SAND': [], + 'soil3c_vr': [], 'soil4c_vr': [], 'cwdc_vr': [], 'deadcrootc': [], 'deadstemc': [],'tlai': [], + 'GPP': [], + 'Y_soil3c_vr': [], 'Y_soil4c_vr': [], 'Y_cwdc_vr': [], 'Y_deadcrootc': [], 'Y_deadstemc': [], 'Y_tlai': [], + 'Y_GPP': [], + 'SCALARAVG_vr': [], + 'PCT_CLAY': [], + 'SNOWDP': [], + 'H2OSOI_10CM': [], + 'HR': [], 'AR': [], 'NPP': [], 'COL_FIRE_CLOSS': [], + 'Y_HR': [], 'Y_AR': [], 'Y_NPP': [], 'Y_COL_FIRE_CLOSS': [], + 'OCCLUDED_P': [], + 'SECONDARY_P': [], + 'LABILE_P': [], + 'APATITE_P': [], + + 'cwdn_vr': [], 'secondp_vr': [], 'cwdp_vr': [],'cwdp': [], 'totcolp': [], 'totvegc': [], 'deadstemn': [], 'deadcrootn': [], + 'deadstemp': [], 'deadcrootp': [], 'leafc': [], 'leafc_storage': [], 'frootc': [], 'frootc_storage': [], + 'Y_cwdn_vr': [], 'Y_secondp_vr': [], 'Y_cwdp_vr': [], 'Y_cwdp': [], 'Y_totcolp': [], 'Y_totvegc': [], 'Y_deadstemn': [], 'Y_deadcrootn': [], + 'Y_deadstemp': [], 'Y_deadcrootp': [], 'Y_leafc': [], 'Y_leafc_storage': [], 'Y_frootc': [], 'Y_frootc_storage': [], + 'totlitc': [], + + 'leafn': [], 'leafn_storage': [], 'frootn': [],'frootn_storage': [], + 'leafp': [], 'leafp_storage': [], 'frootp': [],'frootp_storage': [], + 'livestemc': [], 'livestemc_storage': [], + 'livestemn': [], 'livestemn_storage': [], + 'livestemp': [], 'livestemp_storage': [], + 'labilep_vr': [], 'occlp_vr': [], 'primp_vr': [], + + 'deadcrootc_storage': [], 'deadstemc_storage': [], + 'livecrootc': [], 'livecrootc_storage': [], + 'deadcrootn_storage': [], 'deadstemn_storage': [], + 'livecrootn': [], 'livecrootn_storage': [], + 'deadcrootp_storage': [], 'deadstemp_storage': [], + 'livecrootp': [], 'livecrootp_storage': [], + + 'Y_leafn': [], 'Y_leafn_storage': [], 'Y_frootn': [],'Y_frootn_storage': [], + 'Y_leafp': [], 'Y_leafp_storage': [], 'Y_frootp': [],'Y_frootp_storage': [], + 'Y_livestemc': [], 'Y_livestemc_storage': [], + 'Y_livestemn': [], 'Y_livestemn_storage': [], + 'Y_livestemp': [], 'Y_livestemp_storage': [], + 'Y_labilep_vr': [], 'Y_occlp_vr': [], 'Y_primp_vr': [], + + 'Y_deadcrootc_storage': [], 'Y_deadstemc_storage': [], + 'Y_livecrootc': [], 'Y_livecrootc_storage': [], + 'Y_deadcrootn_storage': [], 'Y_deadstemn_storage': [], + 'Y_livecrootn': [], 'Y_livecrootn_storage': [], + 'Y_deadcrootp_storage': [], 'Y_deadstemp_storage': [], + 'Y_livecrootp': [], 'Y_livecrootp_storage': [], + 'Y_totlitc': [], + # 'H2OCAN': [], 'T_VEG': [], 'T10_VALUE': [], + # 'Y_H2OCAN': [], 'Y_T_VEG': [], 'Y_T10_VALUE': [], + # 'H2OSFC': [], 'H2OSNO': [], 'TH2OSFC': [], 'T_GRND': [], 'T_GRND_R': [], 'T_GRND_U': [], + # 'Y_H2OSFC': [], 'Y_H2OSNO': [], 'Y_TH2OSFC': [], 'Y_T_GRND': [], 'Y_T_GRND_R': [], 'Y_T_GRND_U': [], + # 'H2OSOI_LIQ': [], 'H2OSOI_ICE': [], 'T_SOISNO': [], 'LAKE_SOILC': [], 'T_LAKE': [], + # 'Y_H2OSOI_LIQ': [], 'Y_H2OSOI_ICE': [], 'Y_T_SOISNO': [], 'Y_LAKE_SOILC': [], 'Y_T_LAKE': [], + # 'taf': [], + # 'Y_taf': [], + # 'TS_TOPO': [], + # 'Y_TS_TOPO': [], + 'soil1c_vr': [], 'soil1n_vr': [], 'soil1p_vr': [], + 'soil2c_vr': [], 'soil2n_vr': [], 'soil2p_vr': [], + 'soil3n_vr': [], 'soil3p_vr': [], + 'soil4n_vr': [], 'soil4p_vr': [], + 'litr1c_vr': [], 'litr2c_vr': [], 'litr3c_vr': [], + 'litr1n_vr': [], 'litr2n_vr': [], 'litr3n_vr': [], + 'litr1p_vr': [], 'litr2p_vr': [], 'litr3p_vr': [], + 'sminn_vr': [], 'smin_no3_vr': [], 'smin_nh4_vr': [], + 'Y_soil1c_vr': [], 'Y_soil1n_vr': [], 'Y_soil1p_vr': [], + 'Y_soil2c_vr': [], 'Y_soil2n_vr': [], 'Y_soil2p_vr': [], + 'Y_soil3n_vr': [], 'Y_soil3p_vr': [], + 'Y_soil4n_vr': [], 'Y_soil4p_vr': [], + 'Y_litr1c_vr': [], 'Y_litr2c_vr': [], 'Y_litr3c_vr': [], + 'Y_litr1n_vr': [], 'Y_litr2n_vr': [], 'Y_litr3n_vr': [], + 'Y_litr1p_vr': [], 'Y_litr2p_vr': [], 'Y_litr3p_vr': [], + 'Y_sminn_vr': [], 'Y_smin_no3_vr': [], 'Y_smin_nh4_vr': [] + } + + + + for k, gridcell_idx in enumerate(batch_gridcells): + if k % 100 == 0: + print(f" Processing gridcell {k}/{len(batch_gridcells)} (idx={gridcell_idx})") + + # Get indices + restart_idx = batch_restart_indices[k] + gridcell_id = restart_idx + 1 + + pft_indices_for_cell = pft_map.get(gridcell_id, []) + column_indices_for_cell = column_map.get(gridcell_id, []) + + for var_name in pft_based_vars: + x_val = x_values[var_name][pft_indices_for_cell] + data_dict[var_name].append(x_val.tolist()) + + y_slice = stacked_y_values[var_name][:, pft_indices_for_cell] + avg_y_val = np.mean(y_slice, axis=0) + data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) + + for var_name in col_based_1d_vars: + x_val = x_values[var_name][column_indices_for_cell] + data_dict[var_name].append(x_val.tolist()) + + y_slice = stacked_y_values[var_name][:, column_indices_for_cell] + avg_y_val = np.mean(y_slice, axis=0) + data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) + + for var_name in col_based_2d_vars: + x_val = x_values[var_name][column_indices_for_cell, :] + data_dict[var_name].append(x_val.tolist()) + + y_slice = stacked_y_values[var_name][:, column_indices_for_cell, :] + avg_y_val = np.mean(y_slice, axis=0) + data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) + + # landunit_indices_for_cell = landunit_map.get(gridcell_id, []) + # for var_name in landunit_based_vars: + # x_val = x_values[var_name][landunit_indices_for_cell] + # data_dict[var_name].append(x_val.tolist()) + + # y_slice = stacked_y_values[var_name][:, landunit_indices_for_cell] + # avg_y_val = np.mean(y_slice, axis=0) + # data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) + + # topounit_indices_for_cell = topounit_map.get(gridcell_id, []) + # for var_name in topounit_based_vars: + # x_val = x_values[var_name][topounit_indices_for_cell] + # data_dict[var_name].append(x_val.tolist()) + + # y_slice = stacked_y_values[var_name][:, topounit_indices_for_cell] + # avg_y_val = np.mean(y_slice, axis=0) + # data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) + + data_dict['landfrac'].append(ds2.variables['landfrac'][gridcell_idx]) + data_dict['Latitude'].append(lats[gridcell_idx]) + data_dict['Longitude'].append(lons[gridcell_idx]) + data_dict['LANDFRAC_PFT'].append(ds1.variables['LANDFRAC_PFT'][gridcell_idx]) + data_dict['PCT_NATVEG'].append(ds1.variables['PCT_NATVEG'][gridcell_idx]) + data_dict['AREA'].append(ds1.variables['AREA'][gridcell_idx]) + data_dict['peatf'].append(ds1.variables['peatf'][gridcell_idx]) + data_dict['abm'].append(ds1.variables['abm'][gridcell_idx]) + data_dict['SOIL_COLOR'].append(ds1.variables['SOIL_COLOR'][gridcell_idx]) + data_dict['SOIL_ORDER'].append(ds1.variables['SOIL_ORDER'][gridcell_idx]) + data_dict['PCT_SAND'].append(ds1.variables['PCT_SAND'][:, gridcell_idx]) + data_dict['PCT_NAT_PFT'].append(ds1.variables['PCT_NAT_PFT'][:, gridcell_idx]) + + data_dict['OCCLUDED_P'].append(ds1.variables['OCCLUDED_P'][gridcell_idx]) + data_dict['SECONDARY_P'].append(ds1.variables['SECONDARY_P'][gridcell_idx]) + data_dict['LABILE_P'].append(ds1.variables['LABILE_P'][gridcell_idx]) + data_dict['APATITE_P'].append(ds1.variables['APATITE_P'][gridcell_idx]) + + + data_dict['GPP'].append(ds2.variables['GPP'][0, gridcell_idx]) + data_dict['SCALARAVG_vr'].append(ds2.variables['SCALARAVG_vr'][0, :, gridcell_idx]) + data_dict['HR'].append(ds2.variables['HR'][0, gridcell_idx]) + data_dict['AR'].append(ds2.variables['AR'][0, gridcell_idx]) + data_dict['NPP'].append(ds2.variables['NPP'][0, gridcell_idx]) + # COL_FIRE_CLOSS may not exist in TVA history, use FIRE instead if available + if 'COL_FIRE_CLOSS' in ds2.variables: + data_dict['COL_FIRE_CLOSS'].append(ds2.variables['COL_FIRE_CLOSS'][0, gridcell_idx]) + elif 'FIRE' in ds2.variables: + data_dict['COL_FIRE_CLOSS'].append(ds2.variables['FIRE'][0, gridcell_idx]) + else: + data_dict['COL_FIRE_CLOSS'].append(0.0) + + data_dict['SNOWDP'].append(ds2.variables['SNOWDP'][0, gridcell_idx]) + data_dict['H2OSOI_10CM'].append(ds2.variables['H2OSOI'][0,3, gridcell_idx]) + data_dict['PCT_CLAY'].append(ds1.variables['PCT_CLAY'][:, gridcell_idx]) + + h0_gpp_vals = [] + for ds_h0 in ds_h0_list: + h0_gpp_vals.append(ds_h0.variables['GPP'][0, gridcell_idx]) + avg_h0_gpp = np.mean(h0_gpp_vals) + data_dict['Y_GPP'].append(avg_h0_gpp) + + h0_HR_vals = [] + for ds_h0 in ds_h0_list: + h0_HR_vals.append(ds_h0.variables['HR'][0, gridcell_idx]) + avg_h0_HR = np.mean(h0_HR_vals) + data_dict['Y_HR'].append(avg_h0_HR) + + h0_AR_vals = [] + for ds_h0 in ds_h0_list: + h0_AR_vals.append(ds_h0.variables['AR'][0, gridcell_idx]) + avg_h0_AR = np.mean(h0_AR_vals) + data_dict['Y_AR'].append(avg_h0_AR) + + h0_NPP_vals = [] + for ds_h0 in ds_h0_list: + h0_NPP_vals.append(ds_h0.variables['NPP'][0, gridcell_idx]) + avg_h0_NPP = np.mean(h0_NPP_vals) + data_dict['Y_NPP'].append(avg_h0_NPP) + + h0_COL_FIRE_CLOSS_vals = [] + for ds_h0 in ds_h0_list: + # COL_FIRE_CLOSS may not exist in TVA history + if 'COL_FIRE_CLOSS' in ds_h0.variables: + h0_COL_FIRE_CLOSS_vals.append(ds_h0.variables['COL_FIRE_CLOSS'][0, gridcell_idx]) + elif 'FIRE' in ds_h0.variables: + h0_COL_FIRE_CLOSS_vals.append(ds_h0.variables['FIRE'][0, gridcell_idx]) + else: + h0_COL_FIRE_CLOSS_vals.append(0.0) + avg_h0_COL_FIRE_CLOSS = np.mean(h0_COL_FIRE_CLOSS_vals) + data_dict['Y_COL_FIRE_CLOSS'].append(avg_h0_COL_FIRE_CLOSS) + + # Get forcing index + forcing_idx = batch_forcing_indices[k] + + # 🚀 Fast access to forcing data from memory + data_dict['FLDS'].append(flds_data[:, forcing_idx]) + data_dict['PSRF'].append(psrf_data[:, forcing_idx]) + data_dict['FSDS'].append(fsds_data[:, forcing_idx]) + data_dict['QBOT'].append(qbot_data[:, forcing_idx]) + data_dict['PRECTmms'].append(prect_data[:, forcing_idx]) + data_dict['TBOT'].append(tbot_data[:, forcing_idx]) + # Create DataFrame and save + print(f" Creating DataFrame...") + df_batch = pd.DataFrame(data_dict) + + print(f" Saving to disk...") + batch_save_path = f"{output_dir}/training_data_batch_{batch_number:02d}.pkl" + df_batch.to_pickle(batch_save_path) + + batch_time = time.time() - batch_start_time + print(f"✅ Batch {batch_number} completed: {batch_time:.2f}s") + print(f" Path: {batch_save_path}") + print(f" Shape: {df_batch.shape}") + print(f" Columns: {len(df_batch.columns)}") + print(f" Forcing data length: {len(df_batch['FLDS'].iloc[0])}") + + batch_number += 1 + +# Cleanup NetCDF files +print(f"\n{'='*80}") +print("Cleaning up NetCDF files...") +ds1.close() +ds2.close() +ds10.close() +for ds_h0 in ds_h0_list: + ds_h0.close() +for ds_r in ds_r_list: + ds_r.close() + +print("✅ All NetCDF files closed") + +print(f"\n🎉 Complete dataset PKL generation completed!") +print(f"Total batches: {batch_number - 1}") +print(f"Output directory: {output_dir}") +print(f"Each PKL file contains: {len(data_dict)} variables") +print(f"Forcing data: 58400 time steps") +print(f"Optimization strategy: Pre-load all forcing data to memory, avoid repeated NetCDF access") + +# ============================================================================= +# POST-PROCESSING: MONTHLY AVERAGING (Simplified) +# ============================================================================= + +print(f"\n{'='*80}") +print("POST-PROCESSING: Monthly Averaging") +print(f"{'='*80}") + +# Import glob for file processing +import glob + +# Get all generated PKL files +input_files = sorted(glob.glob(f'{output_dir}/training_data_batch_*.pkl')) +print(f"Found {len(input_files)} PKL files for post-processing") + +# TVA data parameters (20 years, 3-hour interval) +time_series_length = 58400 # 20 years × 365 days × 8 steps/day +steps_per_day = 8 # 3-hour interval = 8 steps/day +days_per_year = 365 +years_in_data = 20 # 1980-1999 +months_per_year = 12 +days_per_month = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31] + +print(f"TVA data parameters:") +print(f" Time series length: {time_series_length}") +print(f" Steps per day: {steps_per_day}") +print(f" Years: {years_in_data}") +print(f" Expected monthly values: {years_in_data * months_per_year}") + +def calculate_monthly_avg(time_series): + """ + Calculate monthly averages from high-resolution time series + """ + if not isinstance(time_series, (list, np.ndarray)): + return [] + + if len(time_series) != time_series_length: + return [] + + monthly_averages = [] + start_idx = 0 + + for year in range(years_in_data): + for month_idx, month_days in enumerate(days_per_month): + # Handle leap year February + if year % 4 == 0 and month_idx == 1: # Leap year February + month_days = 29 + + end_idx = start_idx + month_days * steps_per_day + monthly_avg = np.mean(time_series[start_idx:end_idx]) + monthly_averages.append(monthly_avg) + start_idx = end_idx + + return monthly_averages + +# Define columns to process +time_series_columns = ['FLDS', 'PSRF', 'FSDS', 'QBOT', 'PRECTmms', 'TBOT'] +single_value_columns = [ + 'landfrac', 'LANDFRAC_PFT', 'PCT_NATVEG', 'AREA', 'peatf', 'abm', + 'SOIL_COLOR', 'SOIL_ORDER', 'GPP', 'SNOWDP', 'H2OSOI_10CM', + 'Y_GPP', 'HR', 'AR', 'NPP', 'COL_FIRE_CLOSS', + 'Y_HR', 'Y_AR', 'Y_NPP', 'Y_COL_FIRE_CLOSS', + 'OCCLUDED_P', 'SECONDARY_P', 'LABILE_P', 'APATITE_P' +] +list_like_columns = ['PCT_NAT_PFT', 'PCT_SAND', 'SCALARAVG_vr', 'PCT_CLAY'] + +print(f"\nStarting post-processing...") + +# Process all files +print(f"Processing all {len(input_files)} files...") + +for file_idx, file_path in enumerate(input_files, 1): + print(f"\nProcessing file {file_idx}/{len(input_files)}: {os.path.basename(file_path)}") + + try: + # Read PKL file + df = pd.read_pickle(file_path) + print(f" Original shape: {df.shape}") + print(f" Original columns: {len(df.columns)}") + + # Process time series columns (convert to monthly averages) + print(f" Processing time series columns...") + for col in time_series_columns: + if col in df.columns: + print(f" Processing {col}...") + df[col] = df[col].apply(calculate_monthly_avg) + + # Verify processing results + sample_data = df[col].apply(lambda x: x if isinstance(x, list) else []) + lengths = sample_data.apply(len).unique() + print(f" {col} monthly values length: {lengths}") + + # Process single value columns + print(f" Processing single value columns...") + for col in single_value_columns: + if col in df.columns: + df[col] = df[col].astype(str).str.strip() + df[col] = pd.to_numeric(df[col], errors='coerce') + + # Process list columns + print(f" Processing list columns...") + for col in list_like_columns: + if col in df.columns: + print(f" Expanding {col}...") + expanded_cols = df[col].apply(pd.Series).fillna(0) + expanded_cols = expanded_cols.add_prefix(f"{col}_") + df = df.drop(col, axis=1).join(expanded_cols) + + # Reorder columns + y_columns = [col for col in df.columns if col.startswith('Y_')] + other_columns = [col for col in df.columns if not col.startswith('Y_')] + df = df[other_columns + y_columns] + + # Save processed file + output_file = f"{output_dir}/monthly_{os.path.basename(file_path)}" + df.to_pickle(output_file) + + print(f" ✅ Post-processing completed") + print(f" Processed shape: {df.shape}") + print(f" Processed columns: {len(df.columns)}") + print(f" Saved to: {output_file}") + + # Display sample data + print(f" Sample data:") + print(f" FLDS monthly values length: {len(df['FLDS'].iloc[0]) if 'FLDS' in df.columns else 'N/A'}") + print(f" First 3 coordinates: {df[['Latitude', 'Longitude']].head(3).values.tolist()}") + + except Exception as e: + print(f" ❌ Error processing {os.path.basename(file_path)}: {e}") + continue + +print(f"\n{'='*80}") +print("POST-PROCESSING COMPLETED!") +print(f"Input directory: {output_dir}") +print(f"Output directory: {output_dir}") +print(f"Processed files: {len(input_files)}") + +# Check output files +output_files = glob.glob(f'{output_dir}/monthly_*.pkl') +print(f"Generated files: {len(output_files)}") + +if output_files: + print(f"\nOutput file list:") + for file in output_files: + file_size = os.path.getsize(file) / (1024**3) # GB + print(f" {os.path.basename(file)}: {file_size:.2f} GB") + +print("="*80) + diff --git a/scripts/training_data_generation/python_scripts/72_dataset_forcing_only.py b/scripts/training_data_generation/python_scripts/72_dataset_forcing_only.py new file mode 100644 index 0000000..1ceda49 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/72_dataset_forcing_only.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +""" +TVA Forcing Data Only - PKL Generation Script (Optimized with List Format) +Extract only forcing data and basic geographical information +Automatically converts forcing variables to list format for training compatibility +""" + +import netCDF4 as nc +import numpy as np +import pandas as pd +from scipy.spatial import cKDTree +import os +import sys +import time + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +print("="*80) +print("TVA Forcing Data Only - PKL Generation with List Format Conversion") +print("="*80) + +# File paths using config +history_file = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.h0.0021-01-01-00000.nc' + +# Forcing files - use newly generated ones +forcing_files = { + 'FLDS': os.path.join(config.forcing_netcdf_output_dir, 'TVA_FLDS_1980-1999.nc'), + 'FSDS': os.path.join(config.forcing_netcdf_output_dir, 'TVA_FSDS_1980-1999.nc'), + 'PSRF': os.path.join(config.forcing_netcdf_output_dir, 'TVA_PSRF_1980-1999.nc'), + 'QBOT': os.path.join(config.forcing_netcdf_output_dir, 'TVA_QBOT_1980-1999.nc'), + 'PRECTmms': os.path.join(config.forcing_netcdf_output_dir, 'TVA_PRECTmms_1980-1999.nc'), + 'TBOT': os.path.join(config.forcing_netcdf_output_dir, 'TVA_TBOT_1980-1999.nc'), +} + +# Output directory +output_dir = config.forcing_pkl_output_dir +os.makedirs(output_dir, exist_ok=True) + +# Configuration +batch_size = 1000 +batch_number = 1 + +print(f"Configuration:") +print(f" Batch size: {batch_size}") +print(f" Output directory: {output_dir}") +print(f" History file: {history_file}") +for var, path in forcing_files.items(): + print(f" {var} file: {path}") + +# Load NetCDF files +print("\nLoading NetCDF files...") +start_time = time.time() + +# Load history file for coordinates +ds_history = nc.Dataset(history_file) + +# Load forcing files +forcing_datasets = {} +for var, file_path in forcing_files.items(): + print(f" Loading {var}...") + forcing_datasets[var] = nc.Dataset(file_path) + +print(f"✅ All files loaded: {time.time() - start_time:.2f}s") + +# Get coordinate information +print("\nSetting up spatial filtering...") +lats = ds_history.variables['lat'][:] +lons = ds_history.variables['lon'][:] +landmask = ds_history.variables['landfrac'][:] + +# Filter for land gridcells +valid_mask = (landmask > 0) +valid_gridcells = np.where(valid_mask)[0] + +print(f"Total land gridcells: {len(valid_gridcells)}") +print(f"Latitude range: [{lats.min():.2f}, {lats.max():.2f}]") +print(f"Longitude range: [{lons.min():.2f}, {lons.max():.2f}]") + +# Build KDTree index +print("\nBuilding KDTree index...") +query_coords = np.array([(lats[i], lons[i]) for i in valid_gridcells]) + +# Forcing file coordinates +forcing_lats = forcing_datasets['FLDS'].variables['LATIXY'][:].flatten() +forcing_lons = forcing_datasets['FLDS'].variables['LONGXY'][:].flatten() +forcing_coords = np.vstack((forcing_lats, forcing_lons)).T + +forcing_tree = cKDTree(forcing_coords) +_, all_forcing_indices = forcing_tree.query(query_coords, k=1) + +print("✅ Forcing mapping completed") + +# 🚀 KEY OPTIMIZATION: Pre-load all forcing data into memory +print("\n🚀 Pre-loading all forcing data into memory...") +start_time = time.time() + +forcing_data = {} +for var in forcing_files.keys(): + print(f" Loading {var} data...") + # Load all data: time × nj × ni (58400 × 1 × 11357) + forcing_data[var] = forcing_datasets[var].variables[var][:, 0, :] + print(f" {var} shape: {forcing_data[var].shape}, memory: {forcing_data[var].nbytes / 1024**3:.2f} GB") + +total_forcing_memory = sum(data.nbytes for data in forcing_data.values()) / 1024**3 + +print(f"✅ All forcing data pre-loaded: {time.time() - start_time:.2f}s") +print(f" Total memory usage: {total_forcing_memory:.2f} GB") + +# Close forcing NetCDF files (data is now in memory) +for ds in forcing_datasets.values(): + ds.close() + +print("✅ Forcing NetCDF files closed, data in memory") + +# Process data +print(f"\nStarting to process {len(valid_gridcells)} gridcells...") + +for start_idx in range(0, len(valid_gridcells), batch_size): + end_idx = min(start_idx + batch_size, len(valid_gridcells)) + batch_gridcells = valid_gridcells[start_idx:end_idx] + batch_forcing_indices = all_forcing_indices[start_idx:end_idx] + + print(f"\nProcessing batch {batch_number}: gridcells {start_idx+1}-{end_idx}") + batch_start_time = time.time() + + # Initialize data dictionary - only forcing data and basic geographical information + data_dict = { + 'landfrac': [], + 'Latitude': [], + 'Longitude': [], + 'FLDS': [], + 'PSRF': [], + 'FSDS': [], + 'QBOT': [], + 'PRECTmms': [], + 'TBOT': [], + } + + # Process each gridcell + for k, gridcell_idx in enumerate(batch_gridcells): + if k % 100 == 0: + print(f" Processing gridcell {k}/{len(batch_gridcells)} (idx={gridcell_idx})") + + # Get forcing index + forcing_idx = batch_forcing_indices[k] + + # Basic geographical information + data_dict['landfrac'].append(float(landmask[gridcell_idx])) + data_dict['Latitude'].append(float(lats[gridcell_idx])) + data_dict['Longitude'].append(float(lons[gridcell_idx])) + + # Forcing data (directly from pre-loaded memory arrays) + data_dict['FLDS'].append(forcing_data['FLDS'][:, forcing_idx]) + data_dict['PSRF'].append(forcing_data['PSRF'][:, forcing_idx]) + data_dict['FSDS'].append(forcing_data['FSDS'][:, forcing_idx]) + data_dict['QBOT'].append(forcing_data['QBOT'][:, forcing_idx]) + data_dict['PRECTmms'].append(forcing_data['PRECTmms'][:, forcing_idx]) + data_dict['TBOT'].append(forcing_data['TBOT'][:, forcing_idx]) + + # Create DataFrame and save + print(f" Creating DataFrame...") + df_batch = pd.DataFrame(data_dict) + + # Convert forcing variables to list format for training compatibility + print(f" Converting forcing variables to list format...") + forcing_vars = ['FLDS', 'PSRF', 'FSDS', 'QBOT', 'PRECTmms', 'TBOT'] + for var in forcing_vars: + df_batch[var] = df_batch[var].apply(lambda x: x.tolist() if hasattr(x, 'tolist') else x) + + print(f" Saving to disk...") + batch_save_path = f"{output_dir}/TVA_forcing_batch_{batch_number:02d}.pkl" + df_batch.to_pickle(batch_save_path) + + batch_time = time.time() - batch_start_time + print(f"✅ Batch {batch_number} completed in {batch_time:.2f}s:") + print(f" Path: {batch_save_path}") + print(f" Shape: {df_batch.shape}") + print(f" Columns: {len(df_batch.columns)}") + print(f" Forcing data length: {len(df_batch['FLDS'].iloc[0])}") + print(f" Data format: All forcing variables converted to list format") + + batch_number += 1 + +# Cleanup +print(f"\n{'='*80}") +print("Cleaning up...") +ds_history.close() + +print("✅ All NetCDF files closed") + +print(f"\n🎉 Forcing data PKL generation completed!") +print(f"Total batches: {batch_number - 1}") +print(f"Output directory: {output_dir}") +print(f"Each PKL file contains: 9 variables (landfrac, Latitude, Longitude, 6 forcing variables)") +print(f"Forcing data: 58400 time steps (3-hour resolution, 20 years)") +print(f"Data format: All forcing variables automatically converted to list format") +print(f"Training compatibility: Ready for machine learning training pipelines") +print(f"Memory optimization: All forcing data pre-loaded for faster processing") \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/CNP_IO_updated14_xfer.txt b/scripts/training_data_generation/python_scripts/CNP_IO_updated14_xfer.txt new file mode 100644 index 0000000..8aa7de8 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/CNP_IO_updated14_xfer.txt @@ -0,0 +1,62 @@ +TIME SERIES VARIABLES (Climate Forcing) - 6 variables: +• FLDS, PSRF, FSDS, QBOT, PRECTmms, TBOT + +SURFACE PROPERTIES - 42 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, NPP, AR, HR + +1D PFT VARIABLES (61 variables): + +• deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage +• deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage + +• leafc, leafc_storage, leafn, leafn_storage, leafp, leafp_storage +• frootc, frootc_storage, frootn, frootn_storage, frootp, frootp_storage + +• livestemc, livestemc_storage, livestemn, livestemn_storage, livestemp, livestemp_storage, +• livecrootc, livecrootc_storage, livecrootn, livecrootn_storage, livecrootp, livecrootp_storage + +• leafc_xfer, frootc_xfer, livestemc_xfer, deadstemc_xfer, livecrootc_xfer, deadcrootc_xfer +• gresp_xfer, leafn_xfer, frootn_xfer, livestemn_xfer, deadstemn_xfer, livecrootn_xfer +• deadcrootn_xfer, leafp_xfer, frootp_xfer, livestemp_xfer, deadstemp_xfer, livecrootp_xfer, deadcrootp_xfer + +• cpool, npool, ppool + +• xsmrpool, tlai, totvegc + +2D VARIABLES (layered - 31 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 + +• sminn_vr, smin_no3_vr, smin_nh4_vr + +• soil1c_vr, soil1n_vr, soil1p_vr +• soil2c_vr, soil2n_vr, soil2p_vr +• soil3c_vr, soil3n_vr, soil3p_vr +• soil4c_vr, soil4n_vr, soil4p_vr + +• labilep_vr , occlp_vr, primp_vr, secondp_vr diff --git a/scripts/training_data_generation/python_scripts/cnp_io_parse.py b/scripts/training_data_generation/python_scripts/cnp_io_parse.py new file mode 100644 index 0000000..1b604d2 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/cnp_io_parse.py @@ -0,0 +1,69 @@ +import re + + +def parse_cnp_io_list(filename): + # Map section titles to config keys (case-insensitive matching) + section_map = { + 'TIME SERIES VARIABLES': 'time_series_variables', + 'SURFACE PROPERTIES': 'surface_properties', + 'PFT PARAMETERS': 'pft_parameters', + 'WATER VARIABLES': 'water_variables', + 'SCALAR VARIABLES': 'scalar_variables', + 'TEMPERATURE VARIABLES': 'temperature_variables', + '1D PFT VARIABLES': 'pft_1d_variables', + '2D VARIABLES': 'variables_2d_soil', + # Custom section for explicit COL 1D vars + 'RESTART_COL_1D_VARS': 'dataset_new_RESTART_COL_1D_VARS', + # Accept alternate capitalization found in some files + 'Water variables': 'water_variables', + } + + # Prepare result dict with all keys present + result = {v: [] for v in set(section_map.values())} + current_section = None + + with open(filename) as f: + for raw_line in f: + line = raw_line.strip() + # Section header detection (case-insensitive startswith) + matched = False + for section_title, key in section_map.items(): + if line.lower().startswith(section_title.lower()): + current_section = key + matched = True + break + if matched: + continue + + # Content lines + if not current_section: + continue + + # If line is a variable line (starts with • or comma-separated list) + if line.startswith('•'): + # Remove bullet and split by comma, filter out empty strings + vars_ = [v.strip() for v in line[1:].split(',') if v.strip()] + result[current_section].extend(vars_) + # Some variables are listed as comma-separated after a bullet or alone + elif ',' in line and not line.startswith('['): + vars_ = [v.strip('• ').strip() for v in line.split(',') if v.strip('• ').strip()] + result[current_section].extend(vars_) + # Some variables are listed as single words (rare, but just in case) + elif line and not line.startswith('[') and not line.startswith('#'): + # Only add if it's not a description or exclusion + if re.match(r'^[A-Za-z0-9_]+$', line): + result[current_section].append(line) + + return result + + +if __name__ == "__main__": + # Optional local test, will not run when imported by config + try: + parsed = parse_cnp_io_list("CNP_IO_updated14_xfer.txt") + for key, varlist in parsed.items(): + print(f"{key} (length: {len(varlist)}):") + print(varlist) + print() + except FileNotFoundError: + print("CNP_IO_updated14_xfer.txt not found for standalone test.") \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_FLDS_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_FLDS_20years.py new file mode 100644 index 0000000..ab09952 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/construct_TVA_FLDS_20years.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +""" +Generate TVA FLDS forcing data (1980-1999, 20 years, 3-hour resolution) +Corresponding to crujra.v2.5.5d_FLDS_1901-2023_z01.nc + +Author: Zhuowei Gu +Date: 2024 +""" + +import xarray as xr +import os +import cftime +import numpy as np +import sys +import traceback +import datetime +from pathlib import Path + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +# --- Configuration --- +data_dir = config.forcing_raw_data_path +out_dir = config.forcing_netcdf_output_dir +os.makedirs(out_dir, exist_ok=True) + +start_year = 1980 +end_year = 1999 +final_output_file = os.path.join(out_dir, f"TVA_FLDS_{start_year}-{end_year}.nc") + +print("="*80) +print("Generate TVA FLDS forcing data") +print("1. Merge monthly files (1980-1999, 240 months)") +print("2. Correct time axis discontinuities") +print("3. Maintain 3-hour resolution (no downsampling)") +print("="*80) + +print(f"Source data directory: {data_dir}") +print(f"Target output file: {final_output_file}") + +# --- Find and build all monthly file list --- +print("Searching for monthly files...") +all_monthly_files = [] + +for year in range(start_year, end_year + 1): + for month in range(1, 13): + file_name = f"clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc" + file_path = os.path.join(data_dir, file_name) + if os.path.exists(file_path): + all_monthly_files.append(file_path) + else: + print(f" Warning: File {file_name} does not exist, skipping.") + +if not all_monthly_files: + print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") + sys.exit(1) + +print(f"Found {len(all_monthly_files)} valid monthly files.") + +# --- Define Dask chunks --- +dask_chunks = {'time': 366*8} + +try: + with xr.open_mfdataset( + all_monthly_files, + combine='nested', + concat_dim='time', + decode_times=False, + chunks=dask_chunks, + parallel=False, + ) as ds: + + # === Separate static variables === + print("Separating static coordinate/ID variables...") + static_var_names = ['gridID', 'LONGXY', 'LATIXY'] + static_data = {} + + for var_name in static_var_names: + if var_name in ds: + if 'time' in ds[var_name].dims: + static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() + else: + static_data[var_name] = ds[var_name].load() + + # === Process FLDS variable === + time_varying_vars = ['FLDS'] + if 'FLDS' not in ds.data_vars: + print("Error: FLDS variable not found.") + sys.exit(1) + + print(f"Processing variables: {time_varying_vars}") + ds_temporal = ds[['time'] + time_varying_vars] + + # --- Time axis correction --- + print("Loading raw time coordinates...") + time_values_raw = ds_temporal['time'].load().values + units = ds_temporal['time'].attrs['units'] + calendar = ds_temporal['time'].attrs.get('calendar', 'standard') + if calendar.lower() == 'no_leap': + calendar = 'noleap' + + print(f" Time points: {len(time_values_raw)}") + + print("Starting time coordinate correction...") + time_values_corrected = np.copy(time_values_raw).astype(float) + cumulative_offset_days = 0.0 + expected_step_days = 3.0 / 24.0 + jump_count = 0 + + for i in range(len(time_values_raw) - 1): + corrected_i = time_values_raw[i] + cumulative_offset_days + time_values_corrected[i] = corrected_i + raw_diff = time_values_raw[i+1] - time_values_raw[i] + + if raw_diff < expected_step_days * 0.5: + jump_count += 1 + expected_next_corrected_value = corrected_i + expected_step_days + offset_needed = expected_next_corrected_value - time_values_raw[i+1] + cumulative_offset_days = offset_needed + + time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days + + print(f"Time correction completed. Corrected {jump_count} jumps.") + + print("Decoding corrected time values...") + try: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) + except ValueError: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar) + + # --- Check time monotonicity --- + print("Checking time monotonicity...") + if len(dates) >= 2: + diffs_corrected = np.diff(dates) + zero_timedelta = datetime.timedelta(0) + problem_indices = np.where(diffs_corrected <= zero_timedelta)[0] + + if len(problem_indices) > 0: + print(f"Error! Corrected time is still not monotonic!") + sys.exit(1) + else: + print("✓ Time coordinate check passed.") + + # --- Create dataset with corrected time --- + ds_corrected_time = ds_temporal.copy(deep=False) + ds_corrected_time['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) + ds_corrected_time['time'].encoding['units'] = units + ds_corrected_time['time'].encoding['calendar'] = calendar + + # === Merge static variables === + for var_name, data_array in static_data.items(): + ds_corrected_time[var_name] = data_array + + final_dataset = ds_corrected_time + print("Final dataset:") + print(final_dataset) + + # --- Write to file --- + print(f"Writing to file: {final_output_file}") + output_encoding = { + 'FLDS': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, + 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} + } + + if 'gridID' in final_dataset: + output_encoding['gridID'] = {'dtype': final_dataset['gridID'].dtype} + if 'LONGXY' in final_dataset: + output_encoding['LONGXY'] = {'dtype': final_dataset['LONGXY'].dtype, '_FillValue': np.nan} + if 'LATIXY' in final_dataset: + output_encoding['LATIXY'] = {'dtype': final_dataset['LATIXY'].dtype, '_FillValue': np.nan} + + final_dataset.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) + print(f"✓ Successfully generated: {final_output_file}") + print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") + +except Exception as e: + print(f"\nError: {e}") + traceback.print_exc() + sys.exit(1) + +print("\nScript execution completed.") + diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_FSDS_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_FSDS_20years.py new file mode 100644 index 0000000..eb8b7ea --- /dev/null +++ b/scripts/training_data_generation/python_scripts/construct_TVA_FSDS_20years.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +""" +Generate TVA FSDS forcing data (1980-1999, 20 years, 3-hour resolution) +Corresponding to crujra.v2.5.5d_FSDS_1901-2023_z01.nc + +Author: Zhuowei Gu +Date: 2024 +""" + +import xarray as xr +import os +import cftime +import numpy as np +import sys +import traceback +import datetime +from pathlib import Path + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +# --- Configuration --- +data_dir = config.forcing_raw_data_path +out_dir = config.forcing_netcdf_output_dir +os.makedirs(out_dir, exist_ok=True) + +start_year = 1980 +end_year = 1999 +final_output_file = os.path.join(out_dir, f"TVA_FSDS_{start_year}-{end_year}.nc") + +print("="*80) +print("Generate TVA FSDS forcing data") +print("1. Merge monthly files (1980-1999, 240 months)") +print("2. Correct time axis discontinuities") +print("3. Maintain 3-hour resolution (no downsampling)") +print("="*80) + +print(f"Source data directory: {data_dir}") +print(f"Target output file: {final_output_file}") + +# --- Find and build all monthly file list --- +print("Searching for monthly files...") +all_monthly_files = [] +for year in range(start_year, end_year + 1): + for month in range(1, 13): + file_name = f"clmforc.Daymet.km.1d.Solr.{year}-{month:02d}.nc" + file_path = os.path.join(data_dir, file_name) + if os.path.exists(file_path): + all_monthly_files.append(file_path) + else: + print(f" Warning: File {file_name} does not exist, skipping.") + +if not all_monthly_files: + print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") + sys.exit(1) + +print(f"Found {len(all_monthly_files)} valid monthly files.") + +# --- Define Dask chunks --- +dask_chunks = {'time': 366*8} + +try: + with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', + decode_times=False, chunks=dask_chunks, parallel=False) as ds: + + # === Separate static variables === + print("Separating static coordinate/ID variables...") + static_data = {} + for var_name in ['gridID', 'LONGXY', 'LATIXY']: + if var_name in ds: + if 'time' in ds[var_name].dims: + static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() + else: + static_data[var_name] = ds[var_name].load() + + # === Process FSDS variable === + if 'FSDS' not in ds.data_vars: + print("Error: FSDS variable not found.") + sys.exit(1) + + print("Processing FSDS variable...") + ds_temporal = ds[['time', 'FSDS']] + + # --- Time axis correction --- + print("Loading raw time coordinates...") + time_values_raw = ds_temporal['time'].load().values + units = ds_temporal['time'].attrs['units'] + calendar = ds_temporal['time'].attrs.get('calendar', 'standard') + if calendar.lower() == 'no_leap': + calendar = 'noleap' + + print(f" Time points: {len(time_values_raw)}") + print("Starting time coordinate correction...") + + time_values_corrected = np.copy(time_values_raw).astype(float) + cumulative_offset_days = 0.0 + expected_step_days = 3.0 / 24.0 + jump_count = 0 + + for i in range(len(time_values_raw) - 1): + corrected_i = time_values_raw[i] + cumulative_offset_days + time_values_corrected[i] = corrected_i + raw_diff = time_values_raw[i+1] - time_values_raw[i] + if raw_diff < expected_step_days * 0.5: + jump_count += 1 + expected_next = corrected_i + expected_step_days + cumulative_offset_days = expected_next - time_values_raw[i+1] + time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days + + print(f"Time correction completed. Corrected {jump_count} jumps.") + + print("Decoding corrected time values...") + try: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) + except ValueError: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar) + + # --- Create dataset with corrected time --- + ds_corrected = ds_temporal.copy(deep=False) + ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) + ds_corrected['time'].encoding['units'] = units + ds_corrected['time'].encoding['calendar'] = calendar + + # === Merge static variables === + for var_name, data_array in static_data.items(): + ds_corrected[var_name] = data_array + + # --- Write to file --- + print(f"Writing to file: {final_output_file}") + output_encoding = { + 'FSDS': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, + 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} + } + if 'gridID' in ds_corrected: + output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} + if 'LONGXY' in ds_corrected: + output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} + if 'LATIXY' in ds_corrected: + output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} + + ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) + print(f"✓ Successfully generated: {final_output_file}") + print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") + +except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + sys.exit(1) + +print("\nScript execution completed.") + diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_PRECTmms_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_PRECTmms_20years.py new file mode 100644 index 0000000..7d7998f --- /dev/null +++ b/scripts/training_data_generation/python_scripts/construct_TVA_PRECTmms_20years.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +""" +Generate TVA PRECTmms forcing data (1980-1999, 20 years, 3-hour resolution) +Corresponding to crujra.v2.5.5d_PRECTmms_1901-2023_z01.nc + +Author: Zhuowei Gu +Date: 2024 +""" + +import xarray as xr +import os +import cftime +import numpy as np +import sys +import traceback +import datetime +from pathlib import Path + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +# --- Configuration --- +data_dir = config.forcing_raw_data_path +out_dir = config.forcing_netcdf_output_dir +os.makedirs(out_dir, exist_ok=True) + +start_year = 1980 +end_year = 1999 +final_output_file = os.path.join(out_dir, f"TVA_PRECTmms_{start_year}-{end_year}.nc") + +print("="*80) +print("Generate TVA PRECTmms forcing data") +print("1. Merge monthly files (1980-1999, 240 months)") +print("2. Correct time axis discontinuities") +print("3. Maintain 3-hour resolution (no downsampling)") +print("="*80) + +print(f"Source data directory: {data_dir}") +print(f"Target output file: {final_output_file}") + +# --- Find and build all monthly file list --- +print("Searching for monthly files...") +all_monthly_files = [] +for year in range(start_year, end_year + 1): + for month in range(1, 13): + file_name = f"clmforc.Daymet.km.1d.Prec.{year}-{month:02d}.nc" + file_path = os.path.join(data_dir, file_name) + if os.path.exists(file_path): + all_monthly_files.append(file_path) + else: + print(f" Warning: File {file_name} does not exist, skipping.") + +if not all_monthly_files: + print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") + sys.exit(1) + +print(f"Found {len(all_monthly_files)} valid monthly files.") + +# --- Define Dask chunks --- +dask_chunks = {'time': 366*8} + +try: + with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', + decode_times=False, chunks=dask_chunks, parallel=False) as ds: + + # === Separate static variables === + print("Separating static coordinate/ID variables...") + static_data = {} + for var_name in ['gridID', 'LONGXY', 'LATIXY']: + if var_name in ds: + if 'time' in ds[var_name].dims: + static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() + else: + static_data[var_name] = ds[var_name].load() + + # === Process PRECTmms variable === + if 'PRECTmms' not in ds.data_vars: + print("Error: PRECTmms variable not found.") + sys.exit(1) + + print("Processing PRECTmms variable...") + ds_temporal = ds[['time', 'PRECTmms']] + + # --- Time axis correction --- + print("Loading raw time coordinates...") + time_values_raw = ds_temporal['time'].load().values + units = ds_temporal['time'].attrs['units'] + calendar = ds_temporal['time'].attrs.get('calendar', 'standard') + if calendar.lower() == 'no_leap': + calendar = 'noleap' + + print(f" Time points: {len(time_values_raw)}") + print("Starting time coordinate correction...") + + time_values_corrected = np.copy(time_values_raw).astype(float) + cumulative_offset_days = 0.0 + expected_step_days = 3.0 / 24.0 + jump_count = 0 + + for i in range(len(time_values_raw) - 1): + corrected_i = time_values_raw[i] + cumulative_offset_days + time_values_corrected[i] = corrected_i + raw_diff = time_values_raw[i+1] - time_values_raw[i] + if raw_diff < expected_step_days * 0.5: + jump_count += 1 + expected_next = corrected_i + expected_step_days + cumulative_offset_days = expected_next - time_values_raw[i+1] + time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days + + print(f"Time correction completed. Corrected {jump_count} jumps.") + + print("Decoding corrected time values...") + try: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) + except ValueError: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar) + + # --- Create dataset with corrected time --- + ds_corrected = ds_temporal.copy(deep=False) + ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) + ds_corrected['time'].encoding['units'] = units + ds_corrected['time'].encoding['calendar'] = calendar + + # === Merge static variables === + for var_name, data_array in static_data.items(): + ds_corrected[var_name] = data_array + + # --- Write to file --- + print(f"Writing to file: {final_output_file}") + output_encoding = { + 'PRECTmms': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, + 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} + } + if 'gridID' in ds_corrected: + output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} + if 'LONGXY' in ds_corrected: + output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} + if 'LATIXY' in ds_corrected: + output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} + + ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) + print(f"✓ Successfully generated: {final_output_file}") + print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") + +except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + sys.exit(1) + +print("\nScript execution completed.") + diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_PSRF_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_PSRF_20years.py new file mode 100644 index 0000000..6d13015 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/construct_TVA_PSRF_20years.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +""" +Generate TVA PSRF forcing data (1980-1999, 20 years, 3-hour resolution) +Corresponding to crujra.v2.5.5d_PSRF_1901-2023_z01.nc + +Author: Zhuowei Gu +Date: 2024 +""" + +import xarray as xr +import os +import cftime +import numpy as np +import sys +import traceback +import datetime +from pathlib import Path + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +# --- 配置 --- +# --- Configuration --- +data_dir = config.forcing_raw_data_path +out_dir = config.forcing_netcdf_output_dir +os.makedirs(out_dir, exist_ok=True) + +start_year = 1980 +end_year = 1999 +final_output_file = os.path.join(out_dir, f"TVA_PSRF_{start_year}-{end_year}.nc") + +print("="*80) +print("Generate TVA PSRF forcing data") +print("1. Merge monthly files (1980-1999, 240 months)") +print("2. Correct time axis discontinuities") +print("3. Maintain 3-hour resolution (no downsampling)") +print("="*80) + +print(f"Source data directory: {data_dir}") +print(f"Target output file: {final_output_file}") + +all_monthly_files = [] +for year in range(start_year, end_year + 1): + for month in range(1, 13): + file_name = f"clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc" + file_path = os.path.join(data_dir, file_name) + if os.path.exists(file_path): + all_monthly_files.append(file_path) + +if not all_monthly_files: + print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") + sys.exit(1) + +print(f"Found {len(all_monthly_files)} valid monthly files.") + +dask_chunks = {'time': 366*8} + +try: + with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', + decode_times=False, chunks=dask_chunks, parallel=False) as ds: + + static_data = {} + for var_name in ['gridID', 'LONGXY', 'LATIXY']: + if var_name in ds: + if 'time' in ds[var_name].dims: + static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() + else: + static_data[var_name] = ds[var_name].load() + + ds_temporal = ds[['time', 'PSRF']] + time_values_raw = ds_temporal['time'].load().values + units = ds_temporal['time'].attrs['units'] + calendar = ds_temporal['time'].attrs.get('calendar', 'standard') + if calendar.lower() == 'no_leap': calendar = 'noleap' + + time_values_corrected = np.copy(time_values_raw).astype(float) + cumulative_offset_days = 0.0 + expected_step_days = 3.0 / 24.0 + + for i in range(len(time_values_raw) - 1): + corrected_i = time_values_raw[i] + cumulative_offset_days + time_values_corrected[i] = corrected_i + raw_diff = time_values_raw[i+1] - time_values_raw[i] + if raw_diff < expected_step_days * 0.5: + expected_next = corrected_i + expected_step_days + cumulative_offset_days = expected_next - time_values_raw[i+1] + time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days + + try: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) + except ValueError: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar) + + ds_corrected = ds_temporal.copy(deep=False) + ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) + ds_corrected['time'].encoding['units'] = units + ds_corrected['time'].encoding['calendar'] = calendar + + for var_name, data_array in static_data.items(): + ds_corrected[var_name] = data_array + + output_encoding = { + 'PSRF': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, + 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} + } + if 'gridID' in ds_corrected: + output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} + if 'LONGXY' in ds_corrected: + output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} + if 'LATIXY' in ds_corrected: + output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} + + ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) + print(f"✓ Successfully generated: {final_output_file}") + print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") + +except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + sys.exit(1) + +print("\nScript execution completed.") + diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_QBOT_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_QBOT_20years.py new file mode 100644 index 0000000..17cda3b --- /dev/null +++ b/scripts/training_data_generation/python_scripts/construct_TVA_QBOT_20years.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +""" +Generate TVA QBOT forcing data (1980-1999, 20 years, 3-hour resolution) +Corresponding to crujra.v2.5.5d_QBOT_1901-2023_z01.nc + +Author: Zhuowei Gu +Date: 2024 +""" + +import xarray as xr +import os +import cftime +import numpy as np +import sys +import traceback +import datetime +from pathlib import Path + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +# --- Configuration --- +data_dir = config.forcing_raw_data_path +out_dir = config.forcing_netcdf_output_dir +os.makedirs(out_dir, exist_ok=True) + +start_year = 1980 +end_year = 1999 +final_output_file = os.path.join(out_dir, f"TVA_QBOT_{start_year}-{end_year}.nc") + +print("="*80) +print("Generate TVA QBOT forcing data") +print("1. Merge monthly files (1980-1999, 240 months)") +print("2. Correct time axis discontinuities") +print("3. Maintain 3-hour resolution (no downsampling)") +print("="*80) + +print(f"Source data directory: {data_dir}") +print(f"Target output file: {final_output_file}") + +# --- Find and build all monthly file list --- +print("Searching for monthly files...") +all_monthly_files = [] +for year in range(start_year, end_year + 1): + for month in range(1, 13): + file_name = f"clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc" + file_path = os.path.join(data_dir, file_name) + if os.path.exists(file_path): + all_monthly_files.append(file_path) + else: + print(f" Warning: File {file_name} does not exist, skipping.") + +if not all_monthly_files: + print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") + sys.exit(1) + +print(f"Found {len(all_monthly_files)} valid monthly files.") + +# --- Define Dask chunks --- +dask_chunks = {'time': 366*8} + +try: + with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', + decode_times=False, chunks=dask_chunks, parallel=False) as ds: + + # === Separate static variables === + print("Separating static coordinate/ID variables...") + static_data = {} + for var_name in ['gridID', 'LONGXY', 'LATIXY']: + if var_name in ds: + if 'time' in ds[var_name].dims: + static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() + else: + static_data[var_name] = ds[var_name].load() + + # === Process QBOT variable === + if 'QBOT' not in ds.data_vars: + print("Error: QBOT variable not found.") + sys.exit(1) + + print("Processing QBOT variable...") + ds_temporal = ds[['time', 'QBOT']] + + # --- Time axis correction --- + print("Loading raw time coordinates...") + time_values_raw = ds_temporal['time'].load().values + units = ds_temporal['time'].attrs['units'] + calendar = ds_temporal['time'].attrs.get('calendar', 'standard') + if calendar.lower() == 'no_leap': + calendar = 'noleap' + + print(f" Time points: {len(time_values_raw)}") + print("Starting time coordinate correction...") + + time_values_corrected = np.copy(time_values_raw).astype(float) + cumulative_offset_days = 0.0 + expected_step_days = 3.0 / 24.0 + jump_count = 0 + + for i in range(len(time_values_raw) - 1): + corrected_i = time_values_raw[i] + cumulative_offset_days + time_values_corrected[i] = corrected_i + raw_diff = time_values_raw[i+1] - time_values_raw[i] + if raw_diff < expected_step_days * 0.5: + jump_count += 1 + expected_next = corrected_i + expected_step_days + cumulative_offset_days = expected_next - time_values_raw[i+1] + time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days + + print(f"Time correction completed. Corrected {jump_count} jumps.") + + print("Decoding corrected time values...") + try: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) + except ValueError: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar) + + # --- Create dataset with corrected time --- + ds_corrected = ds_temporal.copy(deep=False) + ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) + ds_corrected['time'].encoding['units'] = units + ds_corrected['time'].encoding['calendar'] = calendar + + # === Merge static variables === + for var_name, data_array in static_data.items(): + ds_corrected[var_name] = data_array + + # --- Write to file --- + print(f"Writing to file: {final_output_file}") + output_encoding = { + 'QBOT': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, + 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} + } + if 'gridID' in ds_corrected: + output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} + if 'LONGXY' in ds_corrected: + output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} + if 'LATIXY' in ds_corrected: + output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} + + ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) + print(f"✓ Successfully generated: {final_output_file}") + print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") + +except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + sys.exit(1) + +print("\nScript execution completed.") + diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_TBOT_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_TBOT_20years.py new file mode 100644 index 0000000..c9d5bc5 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/construct_TVA_TBOT_20years.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +""" +Generate TVA TBOT forcing data (1980-1999, 20 years, 3-hour resolution) +Corresponding to crujra.v2.5.5d_TBOT_1901-2023_z01.nc + +Author: Zhuowei Gu +Date: 2024 +""" + +import xarray as xr +import os +import cftime +import numpy as np +import sys +import traceback +import datetime +from pathlib import Path + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +# --- Configuration --- +data_dir = config.forcing_raw_data_path +out_dir = config.forcing_netcdf_output_dir +os.makedirs(out_dir, exist_ok=True) + +start_year = 1980 +end_year = 1999 +final_output_file = os.path.join(out_dir, f"TVA_TBOT_{start_year}-{end_year}.nc") + +print("="*80) +print("Generate TVA TBOT forcing data") +print("1. Merge monthly files (1980-1999, 240 months)") +print("2. Correct time axis discontinuities") +print("3. Maintain 3-hour resolution (no downsampling)") +print("="*80) + +print(f"Source data directory: {data_dir}") +print(f"Target output file: {final_output_file}") + +# --- Find and build all monthly file list --- +print("Searching for monthly files...") +all_monthly_files = [] +for year in range(start_year, end_year + 1): + for month in range(1, 13): + file_name = f"clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc" + file_path = os.path.join(data_dir, file_name) + if os.path.exists(file_path): + all_monthly_files.append(file_path) + else: + print(f" Warning: File {file_name} does not exist, skipping.") + +if not all_monthly_files: + print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") + sys.exit(1) + +print(f"Found {len(all_monthly_files)} valid monthly files.") + +# --- Define Dask chunks --- +dask_chunks = {'time': 366*8} + +try: + with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', + decode_times=False, chunks=dask_chunks, parallel=False) as ds: + + # === Separate static variables === + print("Separating static coordinate/ID variables...") + static_data = {} + for var_name in ['gridID', 'LONGXY', 'LATIXY']: + if var_name in ds: + if 'time' in ds[var_name].dims: + static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() + else: + static_data[var_name] = ds[var_name].load() + + # === Process TBOT variable === + if 'TBOT' not in ds.data_vars: + print("Error: TBOT variable not found.") + sys.exit(1) + + print("Processing TBOT variable...") + ds_temporal = ds[['time', 'TBOT']] + + # --- Time axis correction --- + print("Loading raw time coordinates...") + time_values_raw = ds_temporal['time'].load().values + units = ds_temporal['time'].attrs['units'] + calendar = ds_temporal['time'].attrs.get('calendar', 'standard') + if calendar.lower() == 'no_leap': + calendar = 'noleap' + + print(f" Time points: {len(time_values_raw)}") + print("Starting time coordinate correction...") + + time_values_corrected = np.copy(time_values_raw).astype(float) + cumulative_offset_days = 0.0 + expected_step_days = 3.0 / 24.0 + jump_count = 0 + + for i in range(len(time_values_raw) - 1): + corrected_i = time_values_raw[i] + cumulative_offset_days + time_values_corrected[i] = corrected_i + raw_diff = time_values_raw[i+1] - time_values_raw[i] + if raw_diff < expected_step_days * 0.5: + jump_count += 1 + expected_next = corrected_i + expected_step_days + cumulative_offset_days = expected_next - time_values_raw[i+1] + time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days + + print(f"Time correction completed. Corrected {jump_count} jumps.") + + print("Decoding corrected time values...") + try: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) + except ValueError: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar) + + # --- Create dataset with corrected time --- + ds_corrected = ds_temporal.copy(deep=False) + ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) + ds_corrected['time'].encoding['units'] = units + ds_corrected['time'].encoding['calendar'] = calendar + + # === Merge static variables === + for var_name, data_array in static_data.items(): + ds_corrected[var_name] = data_array + + # --- Write to file --- + print(f"Writing to file: {final_output_file}") + output_encoding = { + 'TBOT': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, + 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} + } + if 'gridID' in ds_corrected: + output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} + if 'LONGXY' in ds_corrected: + output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} + if 'LATIXY' in ds_corrected: + output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} + + ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) + print(f"✓ Successfully generated: {final_output_file}") + print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") + +except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + sys.exit(1) + +print("\nScript execution completed.") + diff --git a/scripts/training_data_generation/requirements.txt b/scripts/training_data_generation/requirements.txt new file mode 100644 index 0000000..879de41 --- /dev/null +++ b/scripts/training_data_generation/requirements.txt @@ -0,0 +1,16 @@ +# Training Data Generation Requirements +# Install with: pip install -r requirements.txt + +# Core data processing +pandas>=2.0.0 +numpy>=1.24.0 +netCDF4>=1.6.0 +xarray>=2023.1.0 +dask>=2023.0.0 +scipy>=1.10.0 + +# Time handling for NetCDF files +cftime>=1.6.0 + +# Utilities +tqdm>=4.65.0 \ No newline at end of file diff --git a/scripts/training_data_generation/validation/comprehensive_validation.py b/scripts/training_data_generation/validation/comprehensive_validation.py new file mode 100644 index 0000000..e1ad91a --- /dev/null +++ b/scripts/training_data_generation/validation/comprehensive_validation.py @@ -0,0 +1,519 @@ +#!/usr/bin/env python3 +""" +Comprehensive Enhanced Dataset Validation Script +Combines all validation checks: +1. Data existence and format validation +2. Monthly averaging verification +3. History vs restart file comparison +4. PFT variables validation +5. Gridcell-by-gridcell comparison +6. Data consistency validation (actual values) +""" + +import pandas as pd +import netCDF4 as nc +import numpy as np +import os +import sys +import glob +from scipy.spatial import cKDTree + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from config import forcing_raw_data_path, clm_params_nc_path + +def load_enhanced_dataset_files(): + """Load enhanced dataset files""" + enhanced_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "output", "enhanced_training_dataset") + files = sorted(glob.glob(os.path.join(enhanced_dir, "enhanced_monthly_training_data_batch_*.pkl"))) + return files + +def load_raw_forcing_data(): + """Load raw forcing data for monthly averaging validation""" + print("Loading raw forcing data...") + + forcing_vars = ['FLDS', 'FSDS', 'PSRF', 'QBOT', 'PRECTmms', 'TBOT'] + raw_data = {} + + for var in forcing_vars: + if var == 'FLDS': + pattern = os.path.join(forcing_raw_data_path, "clmforc.Daymet.km.1d.TPQWL.*.nc") + elif var == 'FSDS': + pattern = os.path.join(forcing_raw_data_path, "clmforc.Daymet.km.1d.Solr.*.nc") + elif var == 'PSRF': + pattern = os.path.join(forcing_raw_data_path, "clmforc.Daymet.km.1d.TPQWL.*.nc") + elif var == 'QBOT': + pattern = os.path.join(forcing_raw_data_path, "clmforc.Daymet.km.1d.TPQWL.*.nc") + elif var == 'PRECTmms': + pattern = os.path.join(forcing_raw_data_path, "clmforc.Daymet.km.1d.Prec.*.nc") + elif var == 'TBOT': + pattern = os.path.join(forcing_raw_data_path, "clmforc.Daymet.km.1d.TPQWL.*.nc") + + files = sorted(glob.glob(pattern)) + if files: + print(f" {var}: Found {len(files)} files") + raw_data[var] = files + else: + print(f" ❌ {var}: No files found") + + return raw_data + +def load_source_files(): + """Load all source files for validation""" + print("Loading source files...") + + # Load history and restart files + history_file = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.h0.0021-01-01-00000.nc" + restart_file = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.r.0021-01-01-00000.nc" + + history_ds = nc.Dataset(history_file) + restart_ds = nc.Dataset(restart_file) + + # Load CLM parameters file + if not os.path.exists(clm_params_nc_path): + print(f"❌ CLM parameters file not found: {clm_params_nc_path}") + clm_ds = None + else: + clm_ds = nc.Dataset(clm_params_nc_path) + + print(f"✅ History file: {len(history_ds.variables)} variables") + print(f"✅ Restart file: {len(restart_ds.variables)} variables") + if clm_ds: + print(f"✅ CLM parameters: {len(clm_ds.variables)} variables") + + return history_ds, restart_ds, clm_ds + +def build_spatial_mapping(restart_ds): + """Build spatial mapping between PKL coordinates and restart file""" + print("Building spatial mapping...") + + # Get coordinates from restart file + if 'grid1d_lat' in restart_ds.variables and 'grid1d_lon' in restart_ds.variables: + restart_lats = restart_ds.variables['grid1d_lat'][:] + restart_lons = restart_ds.variables['grid1d_lon'][:] + else: + print("❌ grid1d_lat/lon not found in restart file") + return None, None + + # Build KDTree for spatial lookup + restart_coords = np.column_stack((restart_lats, restart_lons)) + tree = cKDTree(restart_coords) + + print(f"✅ Spatial mapping built for {len(restart_coords)} points") + return tree, restart_coords + +def validate_monthly_averaging(pkl_data, raw_forcing_files): + """Validate that forcing data in PKL is monthly averaged""" + print("\n=== 1. Validating Monthly Averaging ===") + + forcing_vars = ['FLDS', 'FSDS', 'PSRF', 'QBOT', 'PRECTmms', 'TBOT'] + results = {} + + for var in forcing_vars: + if var not in pkl_data.columns: + print(f"❌ {var}: Not found in PKL data") + results[var] = False + continue + + print(f"\nValidating {var}...") + + # Get PKL data for first gridcell + pkl_values = pkl_data[var].iloc[0] + + if not isinstance(pkl_values, list): + print(f"❌ {var}: PKL data is not a list") + results[var] = False + continue + + pkl_length = len(pkl_values) + print(f" PKL data length: {pkl_length}") + + # Expected: 20 years * 12 months = 240 months + expected_months = 240 + if pkl_length != expected_months: + print(f"❌ {var}: Expected {expected_months} months, got {pkl_length}") + results[var] = False + continue + + # Sample a few raw files to verify monthly averaging + if var in raw_forcing_files and raw_forcing_files[var]: + try: + # Load one raw file to get hourly data structure + sample_file = raw_forcing_files[var][0] + ds = nc.Dataset(sample_file) + + if var in ds.variables: + hourly_data = ds.variables[var][:] + hours_per_month = hourly_data.shape[0] if len(hourly_data.shape) > 0 else 1 + print(f" Raw file hourly data shape: {hourly_data.shape}") + + # Basic validation: PKL should have fewer values than raw hourly data + if pkl_length < hours_per_month: + print(f" ✅ {var}: PKL data appears to be temporally aggregated") + results[var] = True + else: + print(f" ⚠️ {var}: PKL data length suggests it might not be monthly averaged") + results[var] = False + + ds.close() + + except Exception as e: + print(f" ⚠️ {var}: Could not validate raw data - {e}") + results[var] = True # Assume OK if we can't verify + else: + print(f" ⚠️ {var}: No raw files available for validation") + results[var] = True + + return results + +def validate_data_existence_and_format(pkl_data, history_ds, restart_ds): + """Validate data existence and format""" + print("\n=== 2. Validating Data Existence and Format ===") + + results = { + 'pool_vars_exist': True, + 'pool_vars_in_restart': True, + 'pool_vars_not_in_history': True, + 'pft_vars_exist': True + } + + # Validate pool variables + pool_vars = ['cpool', 'npool', 'ppool', 'xsmrpool'] + print("\nValidating pool variables...") + + for var in pool_vars: + if var not in pkl_data.columns: + print(f" ❌ {var}: Not found in PKL data") + results['pool_vars_exist'] = False + continue + + # Check if restart file has this variable + if var in restart_ds.variables: + print(f" ✅ {var}: Found in restart file") + else: + print(f" ❌ {var}: Not found in restart file") + results['pool_vars_in_restart'] = False + + # Check if history file has this variable (should not have most pool vars) + if var in history_ds.variables: + print(f" ⚠️ {var}: Found in history file (unexpected)") + else: + print(f" ✅ {var}: Not in history file (expected)") + + # Validate PFT variables + pft_cols = [col for col in pkl_data.columns if col.startswith('pft_')] + print(f"\nValidating PFT variables...") + print(f" Found {len(pft_cols)} PFT variables in PKL data") + + if len(pft_cols) == 0: + print(f" ❌ No PFT variables found") + results['pft_vars_exist'] = False + else: + print(f" ✅ PFT variables exist") + + return results + +def validate_pool_data_consistency(pkl_data, restart_ds, tree, restart_coords): + """Validate pool data values against restart file""" + print("\n=== 3. Validating Pool Data Consistency ===") + + pool_vars = ['cpool', 'npool', 'ppool', 'xsmrpool'] + results = {} + + for var in pool_vars: + if var not in pkl_data.columns: + print(f"❌ {var}: Not found in PKL data") + continue + + print(f"\nValidating {var}...") + + # Get PKL coordinates + pkl_lats = pkl_data['Latitude'].values + pkl_lons = pkl_data['Longitude'].values + + # Find nearest restart points for PKL coordinates + pkl_coords = np.column_stack((pkl_lats, pkl_lons)) + _, nearest_indices = tree.query(pkl_coords, k=1) + + # Sample first 5 gridcells for detailed comparison + sample_size = min(5, len(pkl_data)) + sample_indices = range(sample_size) + + matches = 0 + total_compared = 0 + + for i in sample_indices: + # Get PKL data for this gridcell + pkl_values = pkl_data[var].iloc[i] + + if not isinstance(pkl_values, list): + continue + + # Get corresponding restart data + restart_idx = nearest_indices[i] + restart_values = restart_ds.variables[var][restart_idx] + + # Compare values + if isinstance(pkl_values, list) and len(pkl_values) > 0: + # For list-type data, compare first few values + pkl_sample = pkl_values[:5] if len(pkl_values) >= 5 else pkl_values + restart_sample = restart_values[:5] if len(restart_values) >= 5 else restart_values + + if np.allclose(pkl_sample, restart_sample, rtol=1e-10, atol=1e-10): + matches += 1 + else: + print(f" ❌ Gridcell {i}: PKL={pkl_sample} vs Restart={restart_sample}") + + total_compared += 1 + + if total_compared > 0: + match_rate = matches / total_compared * 100 + print(f" ✅ {var}: {matches}/{total_compared} gridcells match ({match_rate:.1f}%)") + results[var] = match_rate >= 80 # 80% threshold for consistency + else: + print(f" ❌ {var}: No valid comparisons made") + results[var] = False + + return results + +def validate_pft_data_consistency(pkl_data, clm_ds): + """Validate PFT data values against CLM parameters file""" + print("\n=== 4. Validating PFT Data Consistency ===") + + if clm_ds is None: + print("❌ CLM parameters file not available") + return {} + + pft_cols = [col for col in pkl_data.columns if col.startswith('pft_')] + if not pft_cols: + print("❌ No PFT variables found in PKL data") + return {} + + print(f"Validating {len(pft_cols)} PFT variables...") + + results = {} + sample_gridcell = 0 # Use first gridcell for validation + + for pft_col in pft_cols[:10]: # Validate first 10 PFT variables + pft_var = pft_col.replace('pft_', '') + + if pft_var not in clm_ds.variables: + print(f" ❌ {pft_col}: Variable not found in CLM parameters") + results[pft_col] = False + continue + + # Get PKL data + pkl_values = pkl_data[pft_col].iloc[sample_gridcell] + + # Get CLM parameters data + clm_values = clm_ds.variables[pft_var][:17] # First 17 PFTs + + if isinstance(pkl_values, list) and len(pkl_values) == len(clm_values): + if np.allclose(pkl_values, clm_values, rtol=1e-10, atol=1e-10): + print(f" ✅ {pft_col}: Values match CLM parameters") + results[pft_col] = True + else: + print(f" ❌ {pft_col}: Values don't match CLM parameters") + print(f" PKL: {pkl_values[:5]}") + print(f" CLM: {clm_values[:5]}") + results[pft_col] = False + else: + print(f" ❌ {pft_col}: Length mismatch (PKL: {len(pkl_values)}, CLM: {len(clm_values)})") + results[pft_col] = False + + return results + +def validate_forcing_data_consistency(pkl_data): + """Validate forcing data by checking values and ranges""" + print("\n=== 5. Validating Forcing Data Consistency ===") + + forcing_vars = ['FLDS', 'FSDS', 'PSRF', 'QBOT', 'PRECTmms', 'TBOT'] + results = {} + + for var in forcing_vars: + if var not in pkl_data.columns: + print(f"❌ {var}: Not found in PKL data") + continue + + print(f"\nValidating {var}...") + + # Get PKL data for first gridcell + pkl_values = pkl_data[var].iloc[0] + + if not isinstance(pkl_values, list) or len(pkl_values) != 240: + print(f" ❌ {var}: Invalid PKL data format or length") + results[var] = False + continue + + # Check if values are reasonable (not all zeros, not all same value) + unique_values = len(set(pkl_values)) + if unique_values < 10: + print(f" ⚠️ {var}: Only {unique_values} unique values (possible data issue)") + else: + print(f" ✅ {var}: {unique_values} unique values") + + # Check value ranges (basic sanity check) + min_val, max_val = min(pkl_values), max(pkl_values) + print(f" Range: {min_val:.3f} to {max_val:.3f}") + + # Basic validation passed + results[var] = True + + return results + +def main(): + """Main validation function""" + print("="*80) + print("COMPREHENSIVE ENHANCED DATASET VALIDATION") + print("="*80) + + # Load enhanced dataset files + enhanced_files = load_enhanced_dataset_files() + if not enhanced_files: + print("❌ No enhanced dataset files found") + return False + + print(f"Found {len(enhanced_files)} enhanced dataset files") + + # Load source files and data + raw_forcing_files = load_raw_forcing_data() + history_ds, restart_ds, clm_ds = load_source_files() + + # Build spatial mapping + tree, restart_coords = build_spatial_mapping(restart_ds) + if tree is None: + print("❌ Failed to build spatial mapping") + return False + + # Validation results + overall_results = { + 'monthly_averaging': {}, + 'data_existence': {}, + 'pool_consistency': {}, + 'pft_consistency': {}, + 'forcing_consistency': {}, + 'total_files_validated': 0, + 'all_passed': True + } + + # Validate first 5 files (5000 gridcells) + files_to_validate = enhanced_files[:5] + print(f"\nFiles to validate: {len(files_to_validate)}") + for i, f in enumerate(files_to_validate): + print(f" {i+1}. {os.path.basename(f)}") + + for i, file_path in enumerate(files_to_validate): + print(f"\n{'='*80}") + print(f"Validating file {i+1}/{len(files_to_validate)}: {os.path.basename(file_path)}") + print(f"{'='*80}") + + try: + # Load PKL data + pkl_data = pd.read_pickle(file_path) + print(f"Loaded PKL data: shape {pkl_data.shape}") + + # 1. Validate monthly averaging + monthly_results = validate_monthly_averaging(pkl_data, raw_forcing_files) + overall_results['monthly_averaging'][f'batch_{i+1:02d}'] = monthly_results + + # 2. Validate data existence and format + existence_results = validate_data_existence_and_format(pkl_data, history_ds, restart_ds) + overall_results['data_existence'][f'batch_{i+1:02d}'] = existence_results + + # 3. Validate pool data consistency + pool_results = validate_pool_data_consistency(pkl_data, restart_ds, tree, restart_coords) + overall_results['pool_consistency'][f'batch_{i+1:02d}'] = pool_results + + # 4. Validate PFT data consistency + pft_results = validate_pft_data_consistency(pkl_data, clm_ds) + overall_results['pft_consistency'][f'batch_{i+1:02d}'] = pft_results + + # 5. Validate forcing data consistency + forcing_results = validate_forcing_data_consistency(pkl_data) + overall_results['forcing_consistency'][f'batch_{i+1:02d}'] = forcing_results + + # Check if all validations passed for this file + file_passed = ( + all(monthly_results.values()) and + all(existence_results.values()) and + all(pool_results.values()) and + all(pft_results.values()) and + all(forcing_results.values()) + ) + + if file_passed: + print(f"\n✅ File {i+1}: All validations passed") + else: + print(f"\n❌ File {i+1}: Some validations failed") + overall_results['all_passed'] = False + + overall_results['total_files_validated'] += 1 + + except Exception as e: + print(f"❌ Error validating file {i+1}: {e}") + overall_results['all_passed'] = False + + # Close datasets + history_ds.close() + restart_ds.close() + if clm_ds: + clm_ds.close() + + # Print final results + print(f"\n{'='*80}") + print("COMPREHENSIVE VALIDATION SUMMARY") + print(f"{'='*80}") + + print(f"Files validated: {overall_results['total_files_validated']}/{len(files_to_validate)}") + total_gridcells = sum(len(pd.read_pickle(f)) for f in files_to_validate[:overall_results['total_files_validated']]) + print(f"Gridcells validated: {total_gridcells}") + + # Monthly averaging results + print(f"\n1. Monthly Averaging Validation:") + for batch, results in overall_results['monthly_averaging'].items(): + passed = sum(results.values()) + total = len(results) + print(f" {batch}: {passed}/{total} forcing variables passed") + + # Data existence results + print(f"\n2. Data Existence and Format:") + for batch, results in overall_results['data_existence'].items(): + passed = sum(results.values()) + total = len(results) + print(f" {batch}: {passed}/{total} checks passed") + + # Pool consistency results + print(f"\n3. Pool Data Consistency:") + for batch, results in overall_results['pool_consistency'].items(): + passed = sum(results.values()) + total = len(results) + print(f" {batch}: {passed}/{total} pool variables consistent") + + # PFT consistency results + print(f"\n4. PFT Data Consistency:") + for batch, results in overall_results['pft_consistency'].items(): + passed = sum(results.values()) + total = len(results) + print(f" {batch}: {passed}/{total} PFT variables consistent") + + # Forcing consistency results + print(f"\n5. Forcing Data Consistency:") + for batch, results in overall_results['forcing_consistency'].items(): + passed = sum(results.values()) + total = len(results) + print(f" {batch}: {passed}/{total} forcing variables consistent") + + # Final result + if overall_results['all_passed']: + print(f"\n🎉 ALL COMPREHENSIVE VALIDATIONS PASSED!") + print(f"✅ Enhanced dataset is completely validated and ready for use") + return True + else: + print(f"\n❌ SOME COMPREHENSIVE VALIDATIONS FAILED!") + print(f"⚠️ Enhanced dataset needs review") + return False + +if __name__ == "__main__": + success = main() + sys.exit(0 if success else 1) + diff --git a/scripts/training_data_generation/validation/forcing_netcdf_validation.py b/scripts/training_data_generation/validation/forcing_netcdf_validation.py new file mode 100644 index 0000000..83d24cc --- /dev/null +++ b/scripts/training_data_generation/validation/forcing_netcdf_validation.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +""" +Simple validation script for forcing data extraction +Compares generated NetCDF files with reference files +""" + +import xarray as xr +import numpy as np +import os +import sys + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +def validate_simple(var_name): + """Simple validation by comparing with reference file""" + print(f"\n{'='*50}") + print(f"Validating {var_name}") + print(f"{'='*50}") + + # File paths + generated_file = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + config.forcing_netcdf_output_dir, f"TVA_{var_name}_1980-1999.nc") + reference_file = f"/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/TVA_forcing_dataset/data/TVA_{var_name}_1980-1999.nc" + + print(f"Generated: {generated_file}") + print(f"Reference: {reference_file}") + + # Check if files exist + if not os.path.exists(generated_file): + print(f"❌ Generated file not found") + return False + + if not os.path.exists(reference_file): + print(f"❌ Reference file not found") + return False + + try: + # Load both files + print("Loading files...") + ds_gen = xr.open_dataset(generated_file) + ds_ref = xr.open_dataset(reference_file) + + # Compare data values + print("Comparing data values...") + data_match = np.allclose(ds_gen[var_name].values, ds_ref[var_name].values, equal_nan=True) + + # Compare basic properties + shape_match = ds_gen[var_name].shape == ds_ref[var_name].shape + dtype_match = ds_gen[var_name].dtype == ds_ref[var_name].dtype + + print(f"Data values match: {data_match}") + print(f"Shape match: {shape_match}") + print(f"Dtype match: {dtype_match}") + + if data_match and shape_match and dtype_match: + print(f"✅ {var_name} validation PASSED") + return True + else: + print(f"❌ {var_name} validation FAILED") + return False + + except Exception as e: + print(f"❌ Error: {e}") + return False + finally: + try: + ds_gen.close() + ds_ref.close() + except: + pass + +def main(): + """Main validation function""" + print("="*60) + print("Simple TVA Forcing Data Validation") + print("="*60) + + variables = ['FLDS', 'FSDS', 'PSRF', 'QBOT', 'PRECTmms', 'TBOT'] + results = {} + + for var_name in variables: + results[var_name] = validate_simple(var_name) + + # Summary + print("\n" + "="*60) + print("VALIDATION SUMMARY") + print("="*60) + + passed = 0 + for var_name in variables: + status = "✅ PASS" if results[var_name] else "❌ FAIL" + print(f"{var_name:12} - {status}") + if results[var_name]: + passed += 1 + + print(f"\nResults: {passed}/{len(variables)} variables passed") + + if passed == len(variables): + print("🎉 ALL VALIDATIONS PASSED! Script is working correctly.") + return 0 + else: + print(f"⚠️ {len(variables) - passed} validation(s) failed.") + return 1 + +if __name__ == "__main__": + exit_code = main() + sys.exit(exit_code) diff --git a/scripts/training_data_generation/validation/forcing_pkl_validation.py b/scripts/training_data_generation/validation/forcing_pkl_validation.py new file mode 100644 index 0000000..f190874 --- /dev/null +++ b/scripts/training_data_generation/validation/forcing_pkl_validation.py @@ -0,0 +1,308 @@ +#!/usr/bin/env python3 +""" +Forcing PKL Validation Script +Validates generated PKL files against processed NetCDF files +""" + +import pandas as pd +import numpy as np +import netCDF4 as nc +import os +import sys +import glob +import time + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +def load_netcdf_data(netcdf_dir, variable): + """ + Load processed NetCDF data for a specific variable + """ + print(f" Loading {variable} NetCDF data...") + + file_path = os.path.join(netcdf_dir, f"TVA_{variable}_1980-1999.nc") + + if not os.path.exists(file_path): + print(f" ❌ NetCDF file not found: {file_path}") + return None + + with nc.Dataset(file_path, 'r') as ds: + if variable in ds.variables: + # NetCDF data is (time, nj, ni), we need (time, ni) + data = ds.variables[variable][:, 0, :] # Remove nj dimension + print(f" Loaded {variable}: shape {data.shape}") + return data + else: + print(f" ❌ Variable {variable} not found in NetCDF file") + return None + +def get_netcdf_coordinates(netcdf_dir): + """ + Get coordinate information from processed NetCDF files + """ + # Use any NetCDF file to get coordinates + sample_file = os.path.join(netcdf_dir, "TVA_FLDS_1980-1999.nc") + if os.path.exists(sample_file): + with nc.Dataset(sample_file, 'r') as ds: + # NetCDF files use LATIXY and LONGXY + lat = ds.variables['LATIXY'][0, :].data # (11357,) + lon = ds.variables['LONGXY'][0, :].data # (11357,) + return lat, lon + else: + print("❌ ERROR: Cannot find sample NetCDF file for coordinates") + return None, None + +def create_coordinate_mapping(pkl_coords, netcdf_coords, tolerance=1e-6): + """ + Create mapping between PKL coordinates and NetCDF coordinates using coordinate matching + """ + print(" Creating coordinate-based mapping...") + + pkl_lat, pkl_lon = pkl_coords + netcdf_lat, netcdf_lon = netcdf_coords + + # Create coordinate pairs + pkl_coord_pairs = np.column_stack([pkl_lat, pkl_lon]) + netcdf_coord_pairs = np.column_stack([netcdf_lat, netcdf_lon]) + + mapping = [] + unmatched_count = 0 + + for i, pkl_coord in enumerate(pkl_coord_pairs): + # Find matching NetCDF coordinate + distances = np.sqrt(np.sum((netcdf_coord_pairs - pkl_coord)**2, axis=1)) + min_idx = np.argmin(distances) + min_distance = distances[min_idx] + + if min_distance < tolerance: + mapping.append(min_idx) + else: + mapping.append(-1) # No match found + unmatched_count += 1 + + print(f" Mapped {len(mapping) - unmatched_count}/{len(mapping)} coordinates") + print(f" Unmatched coordinates: {unmatched_count}") + + return np.array(mapping) + +def validate_pkl_batch_against_netcdf(pkl_file, netcdf_data, variable, batch_index, batch_size=1000): + """ + Validate a PKL batch against corresponding NetCDF data by gridcell range + """ + print(f" Validating {variable} (batch {batch_index + 1})...") + + # Load PKL data + pkl_data = pd.read_pickle(pkl_file) + + if variable not in pkl_data.columns: + print(f" ❌ Variable {variable} not found in PKL file") + return False + + # Calculate the corresponding NetCDF gridcell range for this batch + start_idx = batch_index * batch_size + end_idx = min(start_idx + len(pkl_data), netcdf_data.shape[1]) + + print(f" PKL batch size: {len(pkl_data)} gridcells") + print(f" NetCDF range: gridcells {start_idx}-{end_idx-1}") + + # Sample validation parameters + sample_gridcells = min(20, len(pkl_data)) # Sample fewer gridcells for accuracy + sample_timesteps = min(50, len(pkl_data[variable].iloc[0])) # Sample fewer timesteps + + # Random sampling + np.random.seed(42) # For reproducible results + pkl_gridcell_indices = np.random.choice(len(pkl_data), sample_gridcells, replace=False) + timestep_indices = np.random.choice(len(pkl_data[variable].iloc[0]), sample_timesteps, replace=False) + + print(f" Sampling {sample_gridcells} gridcells and {sample_timesteps} timesteps") + + matches = 0 + total_checks = 0 + + for pkl_gc_idx in pkl_gridcell_indices: + # Get PKL data for this gridcell + pkl_gridcell_data = pkl_data[variable].iloc[pkl_gc_idx] + + # Map PKL gridcell index to NetCDF gridcell index + netcdf_gc_idx = start_idx + pkl_gc_idx + + if netcdf_gc_idx >= netcdf_data.shape[1]: + print(f" ⚠️ NetCDF index {netcdf_gc_idx} out of range") + continue + + for ts_idx in timestep_indices: + if ts_idx < len(pkl_gridcell_data) and ts_idx < netcdf_data.shape[0]: + pkl_value = pkl_gridcell_data[ts_idx] + netcdf_value = netcdf_data[ts_idx, netcdf_gc_idx] + + # Check for NaN values + if np.isnan(pkl_value) and np.isnan(netcdf_value): + matches += 1 + elif not np.isnan(pkl_value) and not np.isnan(netcdf_value): + # Check if values are close (allowing for small numerical differences) + if np.allclose(pkl_value, netcdf_value, rtol=1e-8, atol=1e-8): + matches += 1 + + total_checks += 1 + + if total_checks > 0: + match_rate = matches / total_checks + print(f" Match rate: {match_rate:.4f} ({matches}/{total_checks})") + return match_rate > 0.95 # 95% match rate threshold + else: + print(f" ❌ No valid samples checked") + return False + +def main(): + print("="*80) + print("Forcing PKL Validation Against NetCDF Files") + print("="*80) + + # Configuration + netcdf_dir = config.forcing_netcdf_output_dir + pkl_output_dir = config.forcing_pkl_output_dir + + # Forcing variables to validate + forcing_variables = ['FLDS', 'FSDS', 'PSRF', 'QBOT', 'PRECTmms', 'TBOT'] + + print(f"Configuration:") + print(f" NetCDF directory: {netcdf_dir}") + print(f" PKL output directory: {pkl_output_dir}") + print(f" Variables to validate: {forcing_variables}") + + # Check if directories exist + if not os.path.exists(netcdf_dir): + print(f"❌ ERROR: NetCDF directory not found: {netcdf_dir}") + return False + + if not os.path.exists(pkl_output_dir): + print(f"❌ ERROR: PKL output directory not found: {pkl_output_dir}") + return False + + print(f"\n{'='*60}") + print("Step 1: Getting coordinate information") + print(f"{'='*60}") + + # Get NetCDF coordinate information + netcdf_lat, netcdf_lon = get_netcdf_coordinates(netcdf_dir) + if netcdf_lat is None or netcdf_lon is None: + return False + + print(f" NetCDF grid: {len(netcdf_lat)} points") + + # Get PKL coordinate information + pkl_files = sorted(glob.glob(os.path.join(pkl_output_dir, "TVA_forcing_batch_*.pkl"))) + if not pkl_files: + print(f"❌ ERROR: No PKL files found in {pkl_output_dir}") + return False + + print(f" Loading PKL coordinates from first batch...") + first_pkl = pd.read_pickle(pkl_files[0]) + pkl_lat = first_pkl['Latitude'].values + pkl_lon = first_pkl['Longitude'].values + print(f" PKL grid: {len(pkl_lat)} points") + + print(f"\n{'='*60}") + print("Step 2: Creating coordinate mapping") + print(f"{'='*60}") + + # Create coordinate mapping + coordinate_mapping = create_coordinate_mapping( + (pkl_lat, pkl_lon), + (netcdf_lat, netcdf_lon) + ) + + print(f"\n{'='*60}") + print("Step 3: Loading NetCDF data") + print(f"{'='*60}") + + # Load NetCDF data for each variable + netcdf_data = {} + for var in forcing_variables: + print(f"\nLoading {var}...") + data = load_netcdf_data(netcdf_dir, var) + if data is not None: + netcdf_data[var] = data + else: + print(f"❌ ERROR: Failed to load {var} NetCDF data") + return False + + print(f"\n{'='*60}") + print("Step 4: Validating PKL files") + print(f"{'='*60}") + + print(f"Found {len(pkl_files)} PKL files to validate") + + # Validate first 5 PKL files (as requested) + num_files_to_validate = min(5, len(pkl_files)) + print(f"Validating first {num_files_to_validate} PKL files") + + # Validate each PKL file using sequential gridcell mapping + all_validations_passed = True + validation_results = {} + + for i in range(num_files_to_validate): + pkl_file = pkl_files[i] + print(f"\n{'='*50}") + print(f"Validating batch {i+1}/{num_files_to_validate}: {os.path.basename(pkl_file)}") + print(f"{'='*50}") + + batch_passed = True + batch_results = {} + + for var in forcing_variables: + if var in netcdf_data: + passed = validate_pkl_batch_against_netcdf( + pkl_file, + netcdf_data[var], + var, + i, # batch index + 1000 # batch size + ) + batch_results[var] = passed + if not passed: + batch_passed = False + + validation_results[os.path.basename(pkl_file)] = { + 'passed': batch_passed, + 'details': batch_results + } + + if not batch_passed: + all_validations_passed = False + print(f"❌ Batch {i+1} validation FAILED") + else: + print(f"✅ Batch {i+1} validation PASSED") + + print(f"\n{'='*80}") + print("VALIDATION SUMMARY") + print(f"{'='*80}") + + passed_count = 0 + for pkl_file, result in validation_results.items(): + status = "✅ PASS" if result['passed'] else "❌ FAIL" + print(f"{pkl_file:<30} - {status}") + if result['passed']: + passed_count += 1 + + # Show detailed results + for var, var_result in result['details'].items(): + var_status = "✅" if var_result else "❌" + print(f" {var:<12} - {var_status}") + + print(f"\nResults: {passed_count}/{num_files_to_validate} PKL files passed validation") + + if all_validations_passed: + print("🎉 ALL PKL VALIDATIONS PASSED!") + print("✅ Generated PKL files are consistent with NetCDF files") + return True + else: + print("⚠️ Some PKL validations failed") + print("❌ Please check the PKL generation process") + return False + +if __name__ == "__main__": + success = main() + sys.exit(0 if success else 1) From a3fd3d4b59b481d4b31aaf7f6df3beeb9d577bb8 Mon Sep 17 00:00:00 2001 From: Zhuowei Gu Date: Mon, 27 Oct 2025 03:39:37 -0400 Subject: [PATCH 31/51] Consolidated updates: unified dataset generation and initial condition improvements - Unified dataset generation logic into enhanced_training_dataset.py - Added options (--forcing_only, --enhanced_dataset, --initial_only) - Updated config.py and README.md The new script dynamically loads variables from CNP_IO_updated14_xfer.txt and supports both forcing-only and complete enhanced dataset generation. update initial condition updated dataset generation script --- scripts/training_data_generation/README.md | 204 +- .../bash_script/run_adding_pft_variables.sh | 59 - .../bash_script/run_all_forcing_extraction.sh | 121 -- .../run_comprehensive_validation.sh | 60 - .../run_enhanced_dataset_generation.sh | 52 - .../run_forcing_netcdf_validation.sh | 18 - .../bash_script/run_forcing_pkl_generation.sh | 18 - .../bash_script/run_forcing_pkl_validation.sh | 18 - .../run_incomplete_training_dataset.sh | 34 - scripts/training_data_generation/config.py | 124 +- .../python_scripts/1_add_pft_to_dataset.py | 113 -- .../python_scripts/2_rm_variables.py | 112 -- .../python_scripts/37_dataset.py | 337 ---- .../python_scripts/72_dataset_construction.py | 640 ------ .../python_scripts/72_dataset_forcing_only.py | 200 -- .../construct_TVA_FLDS_20years.py | 183 -- .../construct_TVA_FSDS_20years.py | 152 -- .../construct_TVA_PRECTmms_20years.py | 152 -- .../construct_TVA_PSRF_20years.py | 124 -- .../construct_TVA_QBOT_20years.py | 152 -- .../construct_TVA_TBOT_20years.py | 152 -- .../construct_forcing_20years.py | 292 +++ .../enhanced_training_dataset.py | 1756 +++++++++++++++++ .../training_data_generation/requirements.txt | 25 +- 24 files changed, 2244 insertions(+), 2854 deletions(-) delete mode 100755 scripts/training_data_generation/bash_script/run_adding_pft_variables.sh delete mode 100755 scripts/training_data_generation/bash_script/run_all_forcing_extraction.sh delete mode 100644 scripts/training_data_generation/bash_script/run_comprehensive_validation.sh delete mode 100755 scripts/training_data_generation/bash_script/run_enhanced_dataset_generation.sh delete mode 100755 scripts/training_data_generation/bash_script/run_forcing_netcdf_validation.sh delete mode 100755 scripts/training_data_generation/bash_script/run_forcing_pkl_generation.sh delete mode 100755 scripts/training_data_generation/bash_script/run_forcing_pkl_validation.sh delete mode 100755 scripts/training_data_generation/bash_script/run_incomplete_training_dataset.sh delete mode 100644 scripts/training_data_generation/python_scripts/1_add_pft_to_dataset.py delete mode 100644 scripts/training_data_generation/python_scripts/2_rm_variables.py delete mode 100644 scripts/training_data_generation/python_scripts/37_dataset.py delete mode 100644 scripts/training_data_generation/python_scripts/72_dataset_construction.py delete mode 100644 scripts/training_data_generation/python_scripts/72_dataset_forcing_only.py delete mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_FLDS_20years.py delete mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_FSDS_20years.py delete mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_PRECTmms_20years.py delete mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_PSRF_20years.py delete mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_QBOT_20years.py delete mode 100644 scripts/training_data_generation/python_scripts/construct_TVA_TBOT_20years.py create mode 100644 scripts/training_data_generation/python_scripts/construct_forcing_20years.py create mode 100644 scripts/training_data_generation/python_scripts/enhanced_training_dataset.py diff --git a/scripts/training_data_generation/README.md b/scripts/training_data_generation/README.md index c9bd829..e6b1f6e 100644 --- a/scripts/training_data_generation/README.md +++ b/scripts/training_data_generation/README.md @@ -1,17 +1,18 @@ # Training Data Generation -Simple workflow to extract TVA forcing data from monthly NetCDF files. +Simple workflow to generate training datasets for machine learning using two main scripts. ## Setup ### 1. Create Virtual Environment ```bash -python3.11 -m venv venv_py311 +# Create a virtual environment in the current directory +python3 -m venv venv ``` ### 2. Activate Virtual Environment ```bash -source venv_py311/bin/activate +source venv/bin/activate ``` ### 3. Install Dependencies @@ -19,102 +20,145 @@ source venv_py311/bin/activate pip install -r requirements.txt ``` -## Usage +## Scripts Overview -### Complete Workflow +### 1. `construct_forcing_20years.py` +Generates processed forcing NetCDF files from raw monthly data. -1. **Generate Forcing NetCDF Files** -```bash -./bash_script/run_all_forcing_extraction.sh -``` -Generates 6 NetCDF files in `output/forcing_netcdf/`: -- TVA_FLDS_1980-1999.nc (Longwave radiation) -- TVA_FSDS_1980-1999.nc (Shortwave radiation) -- TVA_PSRF_1980-1999.nc (Surface pressure) -- TVA_QBOT_1980-1999.nc (Specific humidity) -- TVA_PRECTmms_1980-1999.nc (Precipitation) -- TVA_TBOT_1980-1999.nc (Air temperature) - -2. **Validate NetCDF Data Accuracy** -```bash -./bash_script/run_forcing_netcdf_validation.sh -``` -Compares generated NetCDF files with reference files to ensure correctness. +### 2. `enhanced_training_dataset.py` +Generates three different types of training datasets for machine learning. -3. **Generate PKL Files (Training Ready)** -```bash -./bash_script/run_forcing_pkl_generation.sh -``` -Creates optimized PKL files in `output/forcing_hourly_pkl/` for machine learning: -- 12 batch files (TVA_forcing_batch_01.pkl to TVA_forcing_batch_12.pkl) -- Each batch contains 1000 gridcells (except last batch: 357) -- 9 variables per gridcell: landfrac, lat, lon, 6 forcing variables -- 58,400 time steps (3-hour resolution, 20 years) -- **Automatically converts forcing variables to list format for training compatibility** - -4. **Validate PKL Data Accuracy** -```bash -./bash_script/run_forcing_pkl_validation.sh -``` -Validates PKL files against generated NetCDF files using sequential gridcell mapping: -- Validates first 5 PKL batches (5000 gridcells total) -- Each batch corresponds to sequential NetCDF gridcells (0-999, 1000-1999, etc.) -- Ensures PKL data matches NetCDF data with 100% accuracy -- **PKL files are already in list format and ready for training** +## Usage -5. **Generate Complete Training Dataset (Monthly Averaged)** -```bash -./bash_script/run_incomplete_training_dataset.sh -``` -- Integrates ecosystem variables with forcing data -- Applies monthly averaging to forcing variables (240 values for 20 years) -- Output: `output/training_dataset_pkl/monthly_training_data_batch_XX.pkl` -- Automatically removes original PKL files +### Step 1: Generate Forcing NetCDF Files + +First, create the forcing NetCDF files from raw monthly data: -6. **Generate Enhanced Dataset** ```bash -./bash_script/run_enhanced_dataset_generation.sh +python python_scripts/construct_forcing_20years.py ``` -- Adds pool variables (cpool, npool, ppool, xsmrpool) from restart files -- Adds 38 transfer variables and corresponding Y variables -- Output: `output/enhanced_training_dataset/enhanced_monthly_training_data_batch_XX.pkl` -- Automatically removes intermediate files -7. **Add PFT Variables** +**Command Line Options:** ```bash -./bash_script/run_adding_pft_variables.sh +python python_scripts/construct_forcing_20years.py \ + --input-dir /path/to/raw/forcing/data \ + --output-dir /path/to/output/directory \ + --start-year 1980 \ + --end-year 1999 ``` -- Adds PFT (Plant Functional Type) variables from `clm_params_c211124.nc` -- Removes unwanted variables (fire-related, unnecessary PFT variables, SCALARAVG_vr) -- **Final dataset ready for machine learning training** -## Data Validation +**Output**: `output/forcing_netcdf/` +- `FLDS_1980-1999.nc` (Downward longwave radiation) +- `FSDS_1980-1999.nc` (Downward shortwave radiation) +- `PRECTmms_1980-1999.nc` (Precipitation rate) +- `PSRF_1980-1999.nc` (Surface pressure) +- `QBOT_1980-1999.nc` (Specific humidity) +- `TBOT_1980-1999.nc` (Air temperature) -### When to Validate Your Data +### Step 2: Generate Training Datasets -**After Step 1 (Forcing NetCDF Generation):** -```bash -./validation/forcing_netcdf_validation.py -``` -- Validates generated NetCDF files against reference files -- Ensures 6 forcing variables are correctly processed +Use the unified script with three different modes: -**After Step 4 (Forcing PKL Generation):** ```bash -./bash_script/run_forcing_pkl_validation.sh +python python_scripts/enhanced_training_dataset.py --[mode] ``` -- Validates PKL files against generated NetCDF files -- Sequential gridcell mapping validation (first 5 batches) -**After Step 7 (Final Enhanced Dataset):** -```bash -./bash_script/run_comprehensive_validation.sh +#### Mode 1: Forcing-Only Dataset (`--forcing_only`) +- **Purpose**: Raw meteorological forcing data (no monthly averaging) +- **Variables**: 9 variables (landfrac, lat, lon, 6 forcing variables) +- **Data**: Raw 3-hour time series (58,400 time steps) +- **Output**: `output/forcing_only_dataset/` +- **Files**: `forcing_data_batch_XX.pkl` + +#### Mode 2: Enhanced Dataset (`--enhanced_dataset`) +- **Purpose**: Complete ecosystem training dataset with all variables +- **Variables**: 291 variables (initial conditions + Y_variables + PFT variables) +- **Data**: Monthly averaged time series + ecosystem state variables +- **Output**: `output/enhanced_training_dataset/` +- **Files**: `enhanced_monthly_training_data_batch_XX.pkl` +- **Intermediate**: `output/training_dataset_pkl/` (Step 1 output) + +#### Mode 3: Initial-Only Dataset (`--initial_only`) +- **Purpose**: Initial conditions without simulation results (no Y_variables) +- **Variables**: 195 variables (initial conditions + PFT variables, NO Y_variables) +- **Data**: Initial state variables only +- **Output**: `output/initial_condition_dataset/` +- **Files**: `enhanced_monthly_training_data_batch_XX.pkl` + +## Output Directory Structure + +``` +output/ +├── forcing_netcdf/ # Step 1: Forcing NetCDF files +│ ├── FLDS_1980-1999.nc +│ ├── FSDS_1980-1999.nc +│ └── ... (6 NetCDF files) +│ +├── forcing_only_dataset/ # Mode 1: Forcing-only +│ ├── forcing_data_batch_01.pkl +│ └── ... (batch files) +│ +├── training_dataset_pkl/ # Mode 2: Intermediate (Step 1) +│ ├── training_data_batch_01.pkl +│ └── ... (batch files) +│ +├── enhanced_training_dataset/ # Mode 2: Final output +│ ├── enhanced_monthly_training_data_batch_01.pkl +│ └── ... (batch files) +│ +└── initial_condition_dataset/ # Mode 3: Initial-only + ├── enhanced_monthly_training_data_batch_01.pkl + └── ... (batch files) ``` -- Complete validation of the final enhanced dataset -- Includes monthly averaging, data consistency, spatial mapping, and scientific validity checks -- **This is the most important validation** - run after completing all processing steps +## Quick Start + +1. **Generate forcing files**: + ```bash + python python_scripts/construct_forcing_20years.py + ``` + +2. **Choose your dataset mode**: + ```bash + # For complete ecosystem data (291 variables) + python python_scripts/enhanced_training_dataset.py --enhanced_dataset + + # For initial conditions only (195 variables) + python python_scripts/enhanced_training_dataset.py --initial_only + + # For forcing data only (9 variables, raw time series) + python python_scripts/enhanced_training_dataset.py --forcing_only + ``` + +3. **Find your results** in the corresponding `output/` subdirectory. ## Configuration -Edit `config.py` to modify paths and settings. +**All file paths and settings can be modified in `config.py`:** + +- **Input data paths**: Raw forcing data, surface data, restart files +- **Output directories**: Where to save generated datasets +- **File patterns**: How to find input files +- **Variable definitions**: Which variables to include (from CNP_IO file) + +**Key configuration sections:** +- `forcing_raw_data_path`: Raw monthly forcing data directory +- `forcing_netcdf_output_dir`: Processed forcing NetCDF output +- `output_dir`: Base output directory for all datasets +- `surface_data_files`: Surface data NetCDF files +- `ad_spinup_*_files`: Initial spinup files +- `final_spinup_*_files`: Final spinup files (for Y_variables) +- `clm_params_nc_path`: CLM parameters file for PFT variables + +**Example configuration changes:** +```python +# Change input data directory +forcing_raw_data_path = '/path/to/your/raw/data' + +# Change output directory +output_dir = './your_output' + +# Change dataset files +surface_data_files = ['/path/to/your/surface.nc'] +``` + diff --git a/scripts/training_data_generation/bash_script/run_adding_pft_variables.sh b/scripts/training_data_generation/bash_script/run_adding_pft_variables.sh deleted file mode 100755 index 7260fd8..0000000 --- a/scripts/training_data_generation/bash_script/run_adding_pft_variables.sh +++ /dev/null @@ -1,59 +0,0 @@ -#!/bin/bash -# Adding PFT variables script runner - -echo "==========================================" -echo "TVA Adding PFT Variables" -echo "==========================================" - -# Get the project directory -PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" - -# Activate Python virtual environment -source ${PROJECT_DIR}/venv_py311/bin/activate - -# Check if enhanced dataset exists -if [ ! -d "${PROJECT_DIR}/output/enhanced_training_dataset" ]; then - echo "❌ Error: Enhanced dataset directory not found!" - echo " Please run enhanced dataset generation first:" - echo " ./bash_script/run_enhanced_dataset_generation.sh" - exit 1 -fi - -# Check if CLM parameters file exists -if [ ! -f "${PROJECT_DIR}/clm_params.c130821.nc" ]; then - echo "❌ Error: CLM parameters file not found!" - echo " Expected file: ${PROJECT_DIR}/clm_params.c130821.nc" - echo " Please ensure the CLM parameters file is in the correct location." - exit 1 -fi - -echo "✅ Enhanced dataset directory found" -echo "✅ CLM parameters file found" - -# Step 1: Add PFT variables -echo "" -echo "==========================================" -echo "Step 1: Adding PFT Variables" -echo "==========================================" - -python ${PROJECT_DIR}/python_scripts/1_add_pft_to_dataset.py - -echo "" -echo "✅ PFT variables addition completed!" - -# Step 2: Remove unwanted variables -echo "" -echo "==========================================" -echo "Step 2: Removing Unwanted Variables" -echo "==========================================" - -python ${PROJECT_DIR}/python_scripts/2_rm_variables.py - -echo "" -echo "🎉 Adding PFT variables completed!" -echo " - PFT variables added to enhanced dataset" -echo " - Dataset ready for further processing" - -echo "" -echo "Final enhanced dataset files:" -ls -la ${PROJECT_DIR}/output/enhanced_training_dataset/ diff --git a/scripts/training_data_generation/bash_script/run_all_forcing_extraction.sh b/scripts/training_data_generation/bash_script/run_all_forcing_extraction.sh deleted file mode 100755 index 46575bc..0000000 --- a/scripts/training_data_generation/bash_script/run_all_forcing_extraction.sh +++ /dev/null @@ -1,121 +0,0 @@ -#!/bin/bash -# Run all TVA forcing data extraction scripts -# This script processes 6 forcing variables: FLDS, FSDS, PSRF, QBOT, PRECTmms, TBOT - -echo "==========================================" -echo "TVA Forcing Data Extraction Pipeline" -echo "==========================================" -echo "Processing 6 forcing variables (1980-1999, 20 years)" -echo "Output directory: /gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation/output/forcing_netcdf" -echo "==========================================" - -# Get the directory where this script is located -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_DIR="$(dirname "$SCRIPT_DIR")" - -# Change to project directory -cd "$PROJECT_DIR" - -# Activate Python virtual environment -echo "Activating Python virtual environment..." -source venv_py311/bin/activate - -# Create output directory -echo "Creating output directory..." -mkdir -p output/forcing_netcdf - -# Create log directory -mkdir -p logs - -echo "" -echo "Starting forcing data extraction..." -echo "" - -# Initialize counters -success_count=0 -total_count=6 - -# Function to run a script and check result -run_script() { - local script_name="$1" - local variable="$2" - local log_file="logs/${variable}_extraction.log" - - echo "==========================================" - echo "Processing $variable forcing data" - echo "==========================================" - echo "Script: $script_name" - echo "Log file: $log_file" - echo "" - - # Run the script and capture output - python "python_scripts/$script_name" > "$log_file" 2>&1 - - # Check exit status - if [ $? -eq 0 ]; then - echo "✓ $variable forcing data extraction completed successfully" - ((success_count++)) - else - echo "✗ $variable forcing data extraction failed" - echo " Check log file: $log_file" - fi - - echo "" -} - -# Run each forcing variable extraction script -run_script "construct_TVA_FLDS_20years.py" "FLDS" -run_script "construct_TVA_FSDS_20years.py" "FSDS" -run_script "construct_TVA_PSRF_20years.py" "PSRF" -run_script "construct_TVA_QBOT_20years.py" "QBOT" -run_script "construct_TVA_PRECTmms_20years.py" "PRECTmms" -run_script "construct_TVA_TBOT_20years.py" "TBOT" - -echo "==========================================" -echo "All Tasks Completed!" -echo "==========================================" -echo "Successfully processed: $success_count/$total_count variables" - -# List generated files -echo "" -echo "Generated NetCDF files:" -if [ -d "output/forcing_netcdf" ]; then - ls -lh output/forcing_netcdf/*.nc 2>/dev/null || echo "No NetCDF files found" -else - echo "Output directory not found" -fi - -echo "" -echo "Log files:" -if [ -d "logs" ]; then - ls -lh logs/*.log 2>/dev/null || echo "No log files found" -else - echo "Log directory not found" -fi - -echo "" -echo "==========================================" -echo "Summary" -echo "==========================================" -echo "Variables processed:" -echo " FLDS - Longwave radiation" -echo " FSDS - Shortwave radiation" -echo " PSRF - Surface pressure" -echo " QBOT - Specific humidity" -echo " PRECTmms - Precipitation" -echo " TBOT - Air temperature" -echo "" -echo "Output location: $PROJECT_DIR/output/forcing_netcdf/" -echo "Log location: $PROJECT_DIR/logs/" -echo "==========================================" - -# Exit with error code if any script failed -if [ $success_count -ne $total_count ]; then - echo "WARNING: Some extractions failed. Check log files for details." - exit 1 -else - echo "✓ All forcing data extractions completed successfully!" - exit 0 -fi - - diff --git a/scripts/training_data_generation/bash_script/run_comprehensive_validation.sh b/scripts/training_data_generation/bash_script/run_comprehensive_validation.sh deleted file mode 100644 index 7433c90..0000000 --- a/scripts/training_data_generation/bash_script/run_comprehensive_validation.sh +++ /dev/null @@ -1,60 +0,0 @@ -#!/bin/bash -# Comprehensive Enhanced Dataset Validation Script - -echo "==========================================" -echo "Comprehensive Enhanced Dataset Validation" -echo "==========================================" - -# Get the project directory -PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" - -# Activate Python virtual environment -source ${PROJECT_DIR}/venv_py311/bin/activate - -echo "Running comprehensive validation of enhanced dataset..." -echo "" -echo "This validation includes:" -echo "1. ✅ Forcing data monthly averaging verification (240 values for 20 years)" -echo "2. ✅ Data existence and format validation" -echo "3. ✅ Pool data consistency with restart file (spatial mapping)" -echo "4. ✅ PFT data consistency with CLM parameters (value-by-value)" -echo "5. ✅ Forcing data consistency and scientific validity" -echo "6. ✅ History vs restart file comparison" -echo "7. ✅ Gridcell-by-gridcell validation (first 5000 gridcells)" -echo "8. ✅ Spatial mapping verification" -echo "" - -# Run the comprehensive validation script -python ${PROJECT_DIR}/validation/comprehensive_validation.py - -# Capture exit code -validation_result=$? - -echo "" -echo "==========================================" -if [ $validation_result -eq 0 ]; then - echo "🎉 COMPREHENSIVE VALIDATION COMPLETED SUCCESSFULLY!" - echo "✅ Enhanced dataset is completely validated and ready for use" - echo "" - echo "Validation Summary:" - echo " - Monthly averaging: ✅ Verified (240 values for 20 years)" - echo " - Data existence: ✅ Verified" - echo " - Pool data consistency: ✅ Verified (spatial mapping)" - echo " - PFT data consistency: ✅ Verified (value-by-value)" - echo " - Forcing data consistency: ✅ Verified" - echo " - History/Restart comparison: ✅ Verified" - echo " - Gridcell mapping: ✅ Verified (5000 gridcells)" - echo " - Data integrity: ✅ Verified" - echo "" - echo "Your enhanced dataset is scientifically accurate and ready for machine learning!" -else - echo "❌ COMPREHENSIVE VALIDATION FAILED!" - echo "⚠️ Enhanced dataset needs review" - echo "" - echo "Please check the validation output above for details." - echo "Some validations may have failed and need attention." -fi -echo "==========================================" - -exit $validation_result - diff --git a/scripts/training_data_generation/bash_script/run_enhanced_dataset_generation.sh b/scripts/training_data_generation/bash_script/run_enhanced_dataset_generation.sh deleted file mode 100755 index 99d0fc4..0000000 --- a/scripts/training_data_generation/bash_script/run_enhanced_dataset_generation.sh +++ /dev/null @@ -1,52 +0,0 @@ -#!/bin/bash -# Enhanced dataset generation script runner with automatic cleanup - -echo "==========================================" -echo "TVA Enhanced Dataset Generation with Cleanup" -echo "==========================================" - -# Get the project directory -PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" - -# Activate Python virtual environment -source ${PROJECT_DIR}/venv_py311/bin/activate - -# Create output directory if it doesn't exist -mkdir -p ${PROJECT_DIR}/output/enhanced_training_dataset - -# Run enhanced dataset generation script -echo "" -echo "Step 1: Running enhanced dataset generation (37_dataset.py)..." -python ${PROJECT_DIR}/python_scripts/37_dataset.py - -echo "" -echo "✅ Enhanced dataset generation completed!" - -# Clean up intermediate files -echo "" -echo "==========================================" -echo "Step 2: Cleaning Up Intermediate Files" -echo "==========================================" - -echo "Current directory contents:" -ls -la ${PROJECT_DIR}/output/ - -echo "" -echo "Removing training_dataset_pkl directory..." -rm -rf ${PROJECT_DIR}/output/training_dataset_pkl/ - -echo "" -echo "✅ Cleanup completed!" -echo "Remaining directories:" -ls -la ${PROJECT_DIR}/output/ - -echo "" -echo "Final enhanced dataset files:" -ls -la ${PROJECT_DIR}/output/enhanced_training_dataset/ - -echo "" -echo "🎉 Enhanced dataset generation and cleanup completed!" -echo " - Enhanced dataset files ready for training" -echo " - Intermediate files removed" -echo " - Disk space saved: ~2GB" -echo " - Only enhanced dataset files remain for training use." diff --git a/scripts/training_data_generation/bash_script/run_forcing_netcdf_validation.sh b/scripts/training_data_generation/bash_script/run_forcing_netcdf_validation.sh deleted file mode 100755 index 70750d1..0000000 --- a/scripts/training_data_generation/bash_script/run_forcing_netcdf_validation.sh +++ /dev/null @@ -1,18 +0,0 @@ -#!/bin/bash -# Forcing data validation script runner - -echo "==========================================" -echo "TVA Forcing Data Validation" -echo "==========================================" - -# Get the project directory -PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" - -# Activate Python virtual environment -source ${PROJECT_DIR}/venv_py311/bin/activate - -# Run validation script -python ${PROJECT_DIR}/validation/forcing_netcdf_validation.py - -echo "" -echo "Validation completed!" diff --git a/scripts/training_data_generation/bash_script/run_forcing_pkl_generation.sh b/scripts/training_data_generation/bash_script/run_forcing_pkl_generation.sh deleted file mode 100755 index 744c00b..0000000 --- a/scripts/training_data_generation/bash_script/run_forcing_pkl_generation.sh +++ /dev/null @@ -1,18 +0,0 @@ -#!/bin/bash -# Forcing data PKL generation script runner - -echo "==========================================" -echo "TVA Forcing Data PKL Generation" -echo "==========================================" - -# Get the project directory -PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" - -# Activate Python virtual environment -source ${PROJECT_DIR}/venv_py311/bin/activate - -# Run forcing PKL generation script -python ${PROJECT_DIR}/python_scripts/72_dataset_forcing_only.py - -echo "" -echo "Forcing PKL generation completed!" diff --git a/scripts/training_data_generation/bash_script/run_forcing_pkl_validation.sh b/scripts/training_data_generation/bash_script/run_forcing_pkl_validation.sh deleted file mode 100755 index 291b7de..0000000 --- a/scripts/training_data_generation/bash_script/run_forcing_pkl_validation.sh +++ /dev/null @@ -1,18 +0,0 @@ -#!/bin/bash -# Forcing PKL validation script runner - -echo "==========================================" -echo "TVA Forcing PKL Validation" -echo "==========================================" - -# Get the project directory -PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" - -# Activate Python virtual environment -source ${PROJECT_DIR}/venv_py311/bin/activate - -# Run forcing PKL validation script -python ${PROJECT_DIR}/validation/forcing_pkl_validation.py - -echo "" -echo "Forcing PKL validation completed!" diff --git a/scripts/training_data_generation/bash_script/run_incomplete_training_dataset.sh b/scripts/training_data_generation/bash_script/run_incomplete_training_dataset.sh deleted file mode 100755 index 90bcc1b..0000000 --- a/scripts/training_data_generation/bash_script/run_incomplete_training_dataset.sh +++ /dev/null @@ -1,34 +0,0 @@ -#!/bin/bash -# Incomplete training dataset generation script runner - -echo "==========================================" -echo "TVA Incomplete Training Dataset Generation" -echo "==========================================" - -# Get the project directory -PROJECT_DIR="/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/training_data_generation" - -# Activate Python virtual environment -source ${PROJECT_DIR}/venv_py311/bin/activate - -# Run incomplete training dataset generation script -python ${PROJECT_DIR}/python_scripts/72_dataset_construction.py - -echo "" -echo "Incomplete training dataset generation completed!" - -# Clean up original PKL files (keep only monthly averaged files) -echo "" -echo "Cleaning up original PKL files..." -echo "Keeping only monthly averaged files..." - -# Remove original training_data_batch_*.pkl files -rm -f ${PROJECT_DIR}/output/training_dataset_pkl/training_data_batch_*.pkl - -echo "✅ Original PKL files removed" -echo "✅ Only monthly averaged files retained" -echo "" -echo "🎉 Incomplete training dataset generation completed!" -echo " - Monthly averaged PKL files created in output/training_dataset_pkl/" -echo " - Ready for enhanced dataset generation (37_dataset.py)" -echo " - Use run_enhanced_dataset_generation.sh to complete the workflow" diff --git a/scripts/training_data_generation/config.py b/scripts/training_data_generation/config.py index 0fa18a7..fac2849 100644 --- a/scripts/training_data_generation/config.py +++ b/scripts/training_data_generation/config.py @@ -16,94 +16,82 @@ # Raw forcing data directory (contains monthly NetCDF files) forcing_raw_data_path = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/forcing' -# Base output directory -output_dir = './output' +# Base output directory (use absolute path) +import os +base_dir = os.path.dirname(os.path.abspath(__file__)) +output_dir = os.path.join(base_dir, 'output') # Processed forcing NetCDF files output directory -forcing_netcdf_output_dir = './output/forcing_netcdf' +forcing_netcdf_output_dir = os.path.join(output_dir, 'forcing_netcdf', 'TES_SE') # Forcing PKL files output directory -forcing_pkl_output_dir = './output/forcing_hourly_pkl' +forcing_pkl_output_dir = os.path.join(output_dir, 'forcing_hourly_pkl') # Training dataset PKL files output directory -training_dataset_pkl_output_dir = './output/training_dataset_pkl' +training_dataset_pkl_output_dir = os.path.join(output_dir, 'training_dataset_pkl') # CLM parameters NetCDF file path -clm_params_nc_path = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/clm_params_c211124.nc' +clm_params_nc_path = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/entire_domain/domain_surfdata/clm_params_c211124.nc' + +# ============================================================================= +# INPUT FILES CONFIGURATION +# ============================================================================= + +# Surface data files +surface_data_files = [ + '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/entire_domain/domain_surfdata/SEBOX1_surfdata.TES_SE.4km.1d.NLCD.c250202.nc' +] + +# AD-SPINUP files (initial spinup) +ad_spinup_history_files = [ + '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/entire_domain/history_restart_files/uELM_SEBOX1_I1850CNPRDCTCBC.elm.h0.0021-01-01-00000.nc' +] + +ad_spinup_restart_files = [ + '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/entire_domain/history_restart_files/uELM_SEBOX1_I1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc' +] + +# FINAL-SPINUP files (final spinup) +final_spinup_history_files = [ + #'/gpfs/wolf2/cades/cli185/proj-shared/wangd/kmELM/e3sm_runs/uELM_TVA_finalspinref/run/uELM_TVA_finalspinref.elm.h0.0781-01.nc' +] + +final_spinup_restart_files = [ + #'/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc' +] + +# Special P input NetCDF file +#special_p_input_nc = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc' # ============================================================================= # 37_DATASET CONFIGURATION (Enhanced Dataset Generation) # ============================================================================= class Config: - # Input paths for 37_dataset.py - INPUT_GLOB = "./output/training_dataset_pkl/monthly_training_data_batch_*.pkl" - OUTPUT_DIR = "./output/enhanced_training_dataset" + # Input paths for enhanced dataset generation + INPUT_GLOB = os.path.join(output_dir, "training_dataset_pkl", "monthly_training_data_batch_*.pkl") + OUTPUT_DIR = os.path.join(output_dir, "enhanced_training_dataset") ENHANCED_PREFIX = "enhanced_" POOL_VARS = ["cpool", "npool", "ppool", "xsmrpool"] # Special P input NetCDF file - SPECIAL_P_INPUT_NC = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" + #SPECIAL_P_INPUT_NC = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" SPECIAL_P_VARS = [] - # CNP IO configuration file + # CNP IO configuration file (can be changed to any CNP_IO file) CNP_IO_FILE = os.path.join(os.path.dirname(__file__), "python_scripts", "CNP_IO_updated14_xfer.txt") - - # Columns to drop during processing - COLS_TO_DROP = [ - 'H2OSFC', 'H2OSNO', 'H2OSOI_LIQ', 'H2OSOI_ICE', 'LAKE_SOILC', 'H2OCAN', - 'TH2OSFC', 'T_GRND', 'T_GRND_R', 'T_GRND_U', 'T_LAKE', 'T_SOISNO', - 'TS_TOPO', 'taf', 'T_VEG', 'T10_VALUE', - 'Y_H2OSFC', 'Y_H2OSNO', 'Y_H2OSOI_LIQ', 'Y_H2OSOI_ICE', 'Y_LAKE_SOILC', 'Y_H2OCAN', - 'Y_TH2OSFC', 'Y_T_GRND', 'Y_T_GRND_R', 'Y_T_GRND_U', 'Y_T_LAKE', 'Y_T_SOISNO', - 'Y_TS_TOPO', 'Y_taf', 'Y_T_VEG', 'Y_T10_VALUE', - 'annsum_npp', 'avail_retransn', 'avail_retransp', 'cannsum_npp', - 'Y_annsum_npp', 'Y_avail_retransn', 'Y_avail_retransp', 'Y_cannsum_npp', - 'leafc_xfer', 'frootc_xfer', 'livestemc_xfer', 'deadstemc_xfer', 'livecrootc_xfer', 'deadcrootc_xfer', - 'gresp_xfer', 'leafn_xfer', 'frootn_xfer', 'livestemn_xfer', 'deadstemn_xfer', 'livecrootn_xfer', - 'deadcrootn_xfer', 'leafp_xfer', 'frootp_xfer', 'livestemp_xfer', 'deadstemp_xfer', 'livecrootp_xfer', 'deadcrootp_xfer', - 'retransn', 'retransp', 'gresp_storage', - 'Y_leafc_xfer', 'Y_frootc_xfer', 'Y_livestemc_xfer', 'Y_deadstemc_xfer', 'Y_livecrootc_xfer', 'Y_deadcrootc_xfer', - 'Y_gresp_xfer', 'Y_leafn_xfer', 'Y_frootn_xfer', 'Y_livestemn_xfer', 'Y_deadstemn_xfer', 'Y_livecrootn_xfer', - 'Y_deadcrootn_xfer', 'Y_leafp_xfer', 'Y_frootp_xfer', 'Y_livestemp_xfer', 'Y_deadstemp_xfer', 'Y_livecrootp_xfer', 'Y_deadcrootp_xfer', - 'Y_retransn', 'Y_retransp', 'Y_gresp_storage', - 'labilep_vr', 'occlp_vr', 'primp_vr', - 'Y_labilep_vr', 'Y_occlp_vr', 'Y_primp_vr', - 'cpool', 'npool', 'ppool', 'xsmrpool', - 'Y_cpool', 'Y_npool', 'Y_ppool', 'Y_xsmrpool', - 'FH2OSFC', - 'Y_FH2OSFC', - 'secondp_vr', - 'Y_secondp_vr' - ] - # List columns configuration - X_LIST_COLUMNS_2D = [ - 'soil3c_vr', 'soil4c_vr', 'cwdc_vr', 'cwdn_vr', 'secondp_vr', 'cwdp', 'totcolp', 'totlitc', 'cwdp_vr', - 'soil1c_vr', 'soil1n_vr', 'soil1p_vr', - 'soil2c_vr', 'soil2n_vr', 'soil2p_vr', - 'soil3n_vr', 'soil3p_vr', - 'soil4n_vr', 'soil4p_vr', - 'litr1c_vr', 'litr2c_vr', 'litr3c_vr', - 'litr1n_vr', 'litr2n_vr', 'litr3n_vr', - 'litr1p_vr', 'litr2p_vr', 'litr3p_vr', - 'sminn_vr', 'smin_no3_vr', 'smin_nh4_vr', - ] - - X_LIST_COLUMNS_1D = [ - 'deadcrootc', 'deadstemc', 'tlai', 'totvegc', 'deadstemn', 'deadcrootn', 'deadstemp', 'deadcrootp', - 'leafc', 'leafc_storage', 'frootc', 'frootc_storage', - 'leafn', 'leafn_storage', 'frootn', 'frootn_storage', - 'leafp', 'leafp_storage', 'frootp', 'frootp_storage', - 'livestemc', 'livestemc_storage', 'livestemn', 'livestemn_storage', - 'livestemp', 'livestemp_storage', 'deadcrootc_storage', 'deadstemc_storage', - 'livecrootc', 'livecrootc_storage', 'deadcrootn_storage', 'deadstemn_storage', - 'livecrootn', 'livecrootn_storage', 'deadcrootp_storage', 'deadstemp_storage', - 'livecrootp', 'livecrootp_storage', - ] - - Y_LIST_COLUMNS_2D = [f"Y_{name}" for name in X_LIST_COLUMNS_2D] - Y_LIST_COLUMNS_1D = [f"Y_{name}" for name in X_LIST_COLUMNS_1D] + # Alternative CNP_IO files (uncomment to use different files) + # CNP_IO_FILE = os.path.join(os.path.dirname(__file__), "python_scripts", "CNP_IO_alternative.txt") + # CNP_IO_FILE = os.path.join(os.path.dirname(__file__), "python_scripts", "CNP_IO_custom.txt") + + # Dynamic variable lists (populated from CNP_IO file) + # These will be automatically populated by apply_cnp_io_overrides() + X_LIST_COLUMNS_1D = [] + X_LIST_COLUMNS_2D = [] + Y_LIST_COLUMNS_1D = [] + Y_LIST_COLUMNS_2D = [] + COLS_TO_DROP = [] WATER_VARIABLES = [] Y_WATER_VARIABLES = [] @@ -131,6 +119,7 @@ def apply_cnp_io_overrides(cls) -> None: new_1d_vars = list(dict.fromkeys(parsed.get('pft_1d_variables', []) or [])) new_2d_vars = list(dict.fromkeys(parsed.get('variables_2d_soil', []) or [])) new_water_vars = list(dict.fromkeys(parsed.get('water_variables', []) or [])) + new_cols_to_drop = list(dict.fromkeys(parsed.get('cols_to_drop', []) or [])) cls.dataset_new_1D_PFT_VARIABLES = list(dict.fromkeys( (parsed.get('dataset_new_1D_PFT_VARIABLES') or parsed.get('pft_1d_variables') or []) @@ -170,6 +159,9 @@ def apply_cnp_io_overrides(cls) -> None: if new_water_vars: cls.WATER_VARIABLES = new_water_vars cls.X_LIST_COLUMNS_2D = list(dict.fromkeys(list(cls.X_LIST_COLUMNS_2D) + new_water_vars)) + + if new_cols_to_drop: + cls.COLS_TO_DROP = list(dict.fromkeys(list(cls.COLS_TO_DROP) + new_cols_to_drop)) cls.Y_LIST_COLUMNS_1D = [f"Y_{name}" for name in cls.X_LIST_COLUMNS_1D] cls.Y_LIST_COLUMNS_2D = [f"Y_{name}" for name in cls.X_LIST_COLUMNS_2D] diff --git a/scripts/training_data_generation/python_scripts/1_add_pft_to_dataset.py b/scripts/training_data_generation/python_scripts/1_add_pft_to_dataset.py deleted file mode 100644 index 29d5eb5..0000000 --- a/scripts/training_data_generation/python_scripts/1_add_pft_to_dataset.py +++ /dev/null @@ -1,113 +0,0 @@ -import netCDF4 as nc -import numpy as np -import pandas as pd -import glob -import os -import sys - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from config import clm_params_nc_path, training_dataset_pkl_output_dir - -# === 1. Read PFT vectors for all variables from NetCDF === -print("Reading CLM parameters NetCDF file...") -print(f"File path: {clm_params_nc_path}") - -if not os.path.exists(clm_params_nc_path): - print(f"❌ Error: CLM parameters file not found: {clm_params_nc_path}") - sys.exit(1) - -ds = nc.Dataset(clm_params_nc_path) - -target_vars = [ - "aleaff", "allconsl", "allconss", "arootf", "arooti", "astemf", "baset", "bfact", "c3psn", "cc_dstem", - "cc_leaf", "cc_lstem", "cc_other", "croot_stem", "crop", "deadwdcn", "declfact", "displar", "dleaf", "dsladlai", - "evergreen", "fcur", "fcurdv", "fd_pft", "fertnitro", "ffrootcn", "fleafcn", "fleafi", "flivewd", "flnr", - "fm_droot", "fm_dstem", "fm_leaf", "fm_lroot", "fm_lstem", "fm_other", "fm_root", "fnitr", "fr_fcel", "fr_flab", - "fr_flig", "froot_leaf", "frootcn", "fsr_pft", "fstemcn", "gddmin", "graincn", "grnfill", "grperc", "grpnow", - "hybgdd", "irrigated", "laimx", "leaf_long", "leafcn", "lf_fcel", "lf_flab", "lf_flig", "lfemerg", "lflitcn", - "livewdcn", "mxtmp", "pconv", "pftpar20", "pftpar28", "pftpar29", "pftpar30", "pftpar31", "planting_temp", - "pprod10", "pprod100", "pprodharv10", "rholnir", "rholvis", "rhosnir", "rhosvis", "roota_par", "rootb_par", - "rootprof_beta", "season_decid", "slatop", "smpsc", "smpso", "stem_leaf", "stress_decid", "taulnir", "taulvis", - "tausnir", "tausvis", "woody", "xl", "z0mr", "ztopmx" -] - -broadcast_feature_dict = {} -for var in target_vars: - if var in ds.variables: - raw_vals = ds.variables[var][:17] - # Skip variables if any value is NaN or masked (missing) - if np.any(np.isnan(raw_vals)) or np.ma.is_masked(raw_vals): - print(f"Skipped {var}: contains NaN or masked values") - continue - broadcast_feature_dict[var] = list(map(float, raw_vals)) - print(f"Added: {var} (length {len(raw_vals)})") - else: - print(f"Skipped {var}: not found in NetCDF") - -print(f"\n✅ Successfully loaded {len(broadcast_feature_dict)} PFT variables from NetCDF") - -# === 2. Process all PKL files in enhanced_training_dataset directory === -enhanced_output_dir = os.path.join(os.path.dirname(training_dataset_pkl_output_dir), "enhanced_training_dataset") -input_files = sorted(glob.glob(os.path.join(enhanced_output_dir, "enhanced_monthly_training_data_batch_*.pkl"))) - -print(f"\n🔍 Found {len(input_files)} PKL files to process") -print("📋 File list:") -for i, file in enumerate(input_files, 1): - print(f" {i:2d}. {os.path.basename(file)}") - -if len(input_files) == 0: - print("❌ No enhanced PKL files found. Please run enhanced dataset generation first.") - sys.exit(1) - -# === 3. Process each PKL file individually === -for i, file_path in enumerate(input_files, 1): - print(f"\n{'='*80}") - print(f"Processing file {i}/{len(input_files)}: {os.path.basename(file_path)}") - print(f"{'='*80}") - - try: - # Read PKL file - print("📖 Reading PKL file...") - df = pd.read_pickle(file_path) - original_shape = df.shape - print(f"✅ File loaded successfully, original shape: {original_shape}") - - # Check if PFT variables already exist - existing_pft_cols = [col for col in df.columns if col.startswith("pft_")] - if existing_pft_cols: - print(f"⚠️ File already contains {len(existing_pft_cols)} PFT variables, skipping addition") - print(f" Existing PFT variables: {existing_pft_cols[:5]}...") - continue - - # Add each variable as a vector column with pft_ prefix - print("🔧 Adding PFT variables...") - for var, val_list in broadcast_feature_dict.items(): - df["pft_" + var] = [val_list] * len(df) # Add the same list to each row - - new_shape = df.shape - print(f"✅ Successfully added {len(broadcast_feature_dict)} PFT variables") - print(f"📐 New data shape: {original_shape} → {new_shape}") - - # Save in-place (overwrite original file) - print("💾 Saving modified file...") - df.to_pickle(file_path) - print(f"✅ File saved: {os.path.basename(file_path)}") - - # Display sample PFT variables - pft_cols = [col for col in df.columns if col.startswith("pft_")] - if pft_cols: - print("🧾 Sample PFT variables:") - for col in pft_cols[:3]: - print(f" {col}: {df[col].iloc[0]}") - - except Exception as e: - print(f"❌ Failed to process file: {e}") - continue - -print(f"\n{'='*80}") -print("🎉 All files processed successfully!") -print("📊 Summary:") -print(f" - Total files: {len(input_files)}") -print(f" - PFT variables added: {len(broadcast_feature_dict)}") -print("="*80) \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/2_rm_variables.py b/scripts/training_data_generation/python_scripts/2_rm_variables.py deleted file mode 100644 index 69362c8..0000000 --- a/scripts/training_data_generation/python_scripts/2_rm_variables.py +++ /dev/null @@ -1,112 +0,0 @@ -import pandas as pd -import glob -import os -import sys - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from config import training_dataset_pkl_output_dir - -# === 1. Define variables to delete === -# Delete columns starting with SCALARAVG_vr -scalaravg_cols = [col for col in [] if col.startswith("SCALARAVG_vr")] # This list will be dynamically generated at runtime - -# Delete COL_FIRE_CLOSS and Y_COL_FIRE_CLOSS -fire_cols = ["COL_FIRE_CLOSS", "Y_COL_FIRE_CLOSS"] - -# Delete unwanted PFT variables -unwanted_pft_cols = [ - "pft_aleaff", "pft_baset", "pft_cc_dstem", "pft_cc_leaf", "pft_cc_lstem", "pft_cc_other", "pft_displar", - "pft_fcurdv", "pft_fd_pft", "pft_fertnitro", "pft_ffrootcn", "pft_fleafcn", "pft_fm_droot", "pft_fm_dstem", - "pft_fm_leaf", "pft_fm_lroot", "pft_fm_lstem", "pft_fm_other", "pft_fm_root", "pft_fnitr", "pft_fsr_pft", - "pft_fstemcn", "pft_irrigated", "pft_pconv", "pft_pftpar20", "pft_pftpar28", "pft_pftpar29", "pft_pftpar30", - "pft_pftpar31", "pft_pprod10", "pft_pprod100", "pft_pprodharv10" -] - -# Combine all columns to delete -all_drop_cols = fire_cols + unwanted_pft_cols - -print(f"🗑️ Variables to delete:") -print(f" - Fire-related variables: {len(fire_cols)} variables") -print(f" - Unwanted PFT variables: {len(unwanted_pft_cols)} variables") -print(f" - Total: {len(all_drop_cols)} variables") - -# === 2. Process all PKL files in enhanced_training_dataset directory === -enhanced_output_dir = os.path.join(os.path.dirname(training_dataset_pkl_output_dir), "enhanced_training_dataset") -input_files = sorted(glob.glob(os.path.join(enhanced_output_dir, "*.pkl"))) - -print(f"\n🔍 Found {len(input_files)} PKL files to process") -print("📋 File list:") -for i, file in enumerate(input_files, 1): - print(f" {i:2d}. {os.path.basename(file)}") - -if len(input_files) == 0: - print("❌ No enhanced PKL files found. Please run enhanced dataset generation first.") - sys.exit(1) - -# === 3. Process each PKL file individually === -for i, file_path in enumerate(input_files, 1): - print(f"\n{'='*80}") - print(f"Processing file {i}/{len(input_files)}: {os.path.basename(file_path)}") - print(f"{'='*80}") - - try: - # Read PKL file - print("📖 Reading PKL file...") - df = pd.read_pickle(file_path) - original_shape = df.shape - print(f"✅ File loaded successfully, original shape: {original_shape}") - - # Dynamically find columns starting with SCALARAVG_vr - scalaravg_cols = [col for col in df.columns if col.startswith("SCALARAVG_vr")] - if scalaravg_cols: - print(f"🔍 Found {len(scalaravg_cols)} SCALARAVG_vr variables: {scalaravg_cols[:5]}...") - - # Combine all columns to delete - drop_cols = scalaravg_cols + all_drop_cols - - # Check which variables actually exist - existing_drop_cols = [col for col in drop_cols if col in df.columns] - missing_cols = [col for col in drop_cols if col not in df.columns] - - if missing_cols: - print(f"⚠️ Following variables do not exist in file: {missing_cols[:5]}...") - - if existing_drop_cols: - print(f"🗑️ Preparing to delete {len(existing_drop_cols)} variables") - - # Delete variables - df.drop(columns=existing_drop_cols, inplace=True, errors='ignore') - - new_shape = df.shape - print(f"✅ Successfully deleted {len(existing_drop_cols)} variables") - print(f"📐 New data shape: {original_shape} → {new_shape}") - - # Save in-place (overwrite original file) - print("💾 Saving modified file...") - df.to_pickle(file_path) - print(f"✅ File saved: {os.path.basename(file_path)}") - - # Display deleted variables statistics - print("📊 Deleted variables statistics:") - if scalaravg_cols: - print(f" - SCALARAVG_vr variables: {len([col for col in scalaravg_cols if col in existing_drop_cols])} variables") - print(f" - Fire-related variables: {len([col for col in fire_cols if col in existing_drop_cols])} variables") - print(f" - Unwanted PFT variables: {len([col for col in unwanted_pft_cols if col in existing_drop_cols])} variables") - - else: - print("ℹ️ No variables found to delete, skipping processing") - - except Exception as e: - print(f"❌ Failed to process file: {e}") - continue - -print(f"\n{'='*80}") -print("🎉 All files processed successfully!") -print("📊 Summary:") -print(f" - Total files: {len(input_files)}") -print(f" - Deleted variable types:") -print(f" * SCALARAVG_vr variables (dynamically detected)") -print(f" * Fire-related variables: {len(fire_cols)} variables") -print(f" * Unwanted PFT variables: {len(unwanted_pft_cols)} variables") -print("="*80) \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/37_dataset.py b/scripts/training_data_generation/python_scripts/37_dataset.py deleted file mode 100644 index f5371ed..0000000 --- a/scripts/training_data_generation/python_scripts/37_dataset.py +++ /dev/null @@ -1,337 +0,0 @@ -#!/usr/bin/env python3 -import os -import sys -import glob -import argparse -from typing import Dict, List, Tuple - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from config import Config - -import numpy as np -import pandas as pd -from scipy.spatial import cKDTree -import netCDF4 as nc - -VARS_TO_CLEAN = {'H2OSOI_LIQ', 'H2OSOI_ICE'} -FILL_VALUE_THRESHOLD = 1e35 - -def build_restart_kdtree(ds_restart: nc.Dataset) -> Tuple[cKDTree, np.ndarray]: - gridcell_lat = ds_restart.variables["grid1d_lat"][:] - gridcell_lon = ds_restart.variables["grid1d_lon"][:] - coords = np.vstack((gridcell_lat, gridcell_lon)).T - tree = cKDTree(coords) - return tree, coords - -def build_column_index_map(ds_restart: nc.Dataset) -> Dict[int, np.ndarray]: - cols1d_gridcell_index = ds_restart.variables["cols1d_gridcell_index"][:] - unique_ids = np.unique(cols1d_gridcell_index) - mapping: Dict[int, np.ndarray] = {} - for grid_id in unique_ids: - mapping[int(grid_id)] = np.where(cols1d_gridcell_index == grid_id)[0] - return mapping - -def ensure_vars_exist(ds: nc.Dataset, var_names: List[str]) -> List[str]: - existing = [] - for name in var_names: - if name in ds.variables: - existing.append(name) - return existing - -def build_pft_index_map(ds_restart: nc.Dataset) -> Dict[int, np.ndarray]: - pfts1d_gridcell_index = ds_restart.variables["pfts1d_gridcell_index"][:] - unique_ids = np.unique(pfts1d_gridcell_index) - mapping: Dict[int, np.ndarray] = {} - for grid_id in unique_ids: - mapping[int(grid_id)] = np.where(pfts1d_gridcell_index == grid_id)[0] - return mapping - -def extract_col1d_x(ds_restart: nc.Dataset, var_name: str, col_indices: np.ndarray) -> List[float]: - if col_indices.size == 0: - return [] - values = ds_restart.variables[var_name][col_indices] - return values.astype(float).tolist() - -def extract_col1d_y(ds_r_list: List[nc.Dataset], var_name: str, col_indices: np.ndarray) -> List[float]: - if col_indices.size == 0: - return [] - slices: List[np.ndarray] = [] - for ds_r in ds_r_list: - values = ds_r.variables[var_name][col_indices] - slices.append(np.asarray(values, dtype=float)) - stacked = np.stack(slices, axis=0) - avg = np.mean(stacked, axis=0) - return avg.tolist() - -def extract_col2d_x(ds_restart: nc.Dataset, var_name: str, col_indices: np.ndarray) -> List[List[float]]: - if col_indices.size == 0: - return [] - values = ds_restart.variables[var_name][col_indices, :] - values_np = np.asarray(values, dtype=float) - if var_name in VARS_TO_CLEAN: - values_np[values_np >= FILL_VALUE_THRESHOLD] = 0.0 - return values_np.tolist() - -def extract_col2d_y(ds_r_list: List[nc.Dataset], var_name: str, col_indices: np.ndarray) -> List[List[float]]: - if col_indices.size == 0: - return [] - slices: List[np.ndarray] = [] - for ds_r in ds_r_list: - values = ds_r.variables[var_name][col_indices, :] - values_np = np.asarray(values, dtype=float) - if var_name in VARS_TO_CLEAN: - values_np[values_np >= FILL_VALUE_THRESHOLD] = 0.0 - slices.append(values_np) - stacked = np.stack(slices, axis=0) - avg = np.mean(stacked, axis=0) - return avg.tolist() - -def extract_pft1d_x(ds_restart: nc.Dataset, var_name: str, pft_indices: np.ndarray) -> List[float]: - if pft_indices.size == 0: - return [] - values = ds_restart.variables[var_name][pft_indices] - return np.asarray(values, dtype=float).tolist() - -def extract_pft1d_y(ds_r_list: List[nc.Dataset], var_name: str, pft_indices: np.ndarray) -> List[float]: - if pft_indices.size == 0: - return [] - slices: List[np.ndarray] = [] - for ds_r in ds_r_list: - values = ds_r.variables[var_name][pft_indices] - slices.append(np.asarray(values, dtype=float)) - stacked = np.stack(slices, axis=0) - avg = np.mean(stacked, axis=0) - return avg.tolist() - -def extract_pft2d_x(ds_restart: nc.Dataset, var_name: str, pft_indices: np.ndarray) -> List[List[float]]: - if pft_indices.size == 0: - return [] - values = ds_restart.variables[var_name][pft_indices, :] - return np.asarray(values, dtype=float).tolist() - -def extract_pft2d_y(ds_r_list: List[nc.Dataset], var_name: str, pft_indices: np.ndarray) -> List[List[float]]: - if pft_indices.size == 0: - return [] - slices: List[np.ndarray] = [] - for ds_r in ds_r_list: - values = ds_r.variables[var_name][pft_indices, :] - slices.append(np.asarray(values, dtype=float)) - stacked = np.stack(slices, axis=0) - avg = np.mean(stacked, axis=0) - return avg.tolist() - -def augment_dataframe_with_pools( - df: pd.DataFrame, - ds_restart: nc.Dataset, - ds_r_list: List[nc.Dataset], - restart_tree: cKDTree, - restart_coords: np.ndarray, - col_index_map: Dict[int, np.ndarray], - pool_vars: List[str], -) -> pd.DataFrame: - if "Latitude" not in df.columns or "Longitude" not in df.columns: - raise ValueError("DataFrame is missing Latitude/Longitude columns for mapping.") - - pool_vars_existing = ensure_vars_exist(ds_restart, pool_vars) - if not pool_vars_existing: - raise ValueError(f"None of the target variables found in restart file: {pool_vars}") - - pool_vars_final: List[str] = [ - v for v in pool_vars_existing if all(v in ds_r.variables for ds_r in ds_r_list) - ] - if not pool_vars_final: - raise ValueError("Target variables do not exist in the set of Y files.") - - latitudes = df["Latitude"].to_numpy() - longitudes = df["Longitude"].to_numpy() - query_coords = np.vstack((latitudes, longitudes)).T - _, nearest_restart_indices = restart_tree.query(query_coords, k=1) - - results_x: Dict[str, List[List[float]]] = {v: [] for v in pool_vars_final} - results_y: Dict[str, List[List[float]]] = {f"Y_{v}": [] for v in pool_vars_final} - - for row_idx, restart_idx in enumerate(nearest_restart_indices): - gridcell_id = int(restart_idx) + 1 - col_indices = col_index_map.get(gridcell_id, np.array([], dtype=int)) - for v in pool_vars_final: - x_vals = extract_col1d_x(ds_restart, v, col_indices) - y_vals = extract_col1d_y(ds_r_list, v, col_indices) - results_x[v].append(x_vals) - results_y[f"Y_{v}"].append(y_vals) - - for v in pool_vars_final: - df[v] = results_x[v] - df[f"Y_{v}"] = results_y[f"Y_{v}"] - - return df - -def augment_dataframe_with_vars( - df: pd.DataFrame, - ds_restart: nc.Dataset, - ds_special_p_restart: nc.Dataset, - ds_r_list: List[nc.Dataset], - restart_tree: cKDTree, - restart_coords: np.ndarray, - col_index_map: Dict[int, np.ndarray], - pft_index_map: Dict[int, np.ndarray], - vars_1d: List[str], - vars_2d: List[str], - special_p_vars: List[str], -) -> pd.DataFrame: - if "Latitude" not in df.columns or "Longitude" not in df.columns: - raise ValueError("DataFrame is missing Latitude/Longitude columns for mapping.") - - vars_1d_existing = ensure_vars_exist(ds_restart, vars_1d) - vars_2d_existing = ensure_vars_exist(ds_restart, vars_2d) - - final_1d: List[str] = [v for v in vars_1d_existing if all(v in ds_r.variables for ds_r in ds_r_list)] - final_2d: List[str] = [v for v in vars_2d_existing if all(v in ds_r.variables for ds_r in ds_r_list)] - - if not final_1d and not final_2d: - return df - - latitudes = df["Latitude"].to_numpy() - longitudes = df["Longitude"].to_numpy() - query_coords = np.vstack((latitudes, longitudes)).T - _, nearest_restart_indices = restart_tree.query(query_coords, k=1) - - results_x_1d: Dict[str, List[List[float]]] = {v: [] for v in final_1d} - results_y_1d: Dict[str, List[List[float]]] = {f"Y_{v}": [] for v in final_1d} - results_x_2d: Dict[str, List[List[List[float]]]] = {v: [] for v in final_2d} - results_y_2d: Dict[str, List[List[List[float]]]] = {f"Y_{v}": [] for v in final_2d} - - for row_idx, restart_idx in enumerate(nearest_restart_indices): - gridcell_id = int(restart_idx) + 1 - col_indices = col_index_map.get(gridcell_id, np.array([], dtype=int)) - pft_indices = pft_index_map.get(gridcell_id, np.array([], dtype=int)) - - for v in final_1d: - var_obj = ds_restart.variables[v] - dims = tuple(var_obj.dimensions) - if "pft" in dims: - x_vals = extract_pft1d_x(ds_restart, v, pft_indices) - y_vals = extract_pft1d_y(ds_r_list, v, pft_indices) - else: - x_vals = extract_col1d_x(ds_restart, v, col_indices) - y_vals = extract_col1d_y(ds_r_list, v, col_indices) - results_x_1d[v].append(x_vals) - results_y_1d[f"Y_{v}"].append(y_vals) - - for v in final_2d: - ds_for_x = ds_special_p_restart if v in special_p_vars else ds_restart - var_obj = ds_for_x.variables[v] - dims = tuple(var_obj.dimensions) - - if "pft" in dims: - x_vals_2d = extract_pft2d_x(ds_for_x, v, pft_indices) - y_vals_2d = extract_pft2d_y(ds_r_list, v, pft_indices) - else: - x_vals_2d = extract_col2d_x(ds_for_x, v, col_indices) - y_vals_2d = extract_col2d_y(ds_r_list, v, col_indices) - - results_x_2d[v].append(x_vals_2d) - results_y_2d[f"Y_{v}"].append(y_vals_2d) - - for v in final_1d: - df[v] = results_x_1d[v] - df[f"Y_{v}"] = results_y_1d[f"Y_{v}"] - for v in final_2d: - df[v] = results_x_2d[v] - df[f"Y_{v}"] = results_y_2d[f"Y_{v}"] - - return df - -def main(): - parser = argparse.ArgumentParser( - description="Augment batched training data with variables from RESTART files." - ) - parser.add_argument( - "--input_glob", - default=Config.INPUT_GLOB, - help="Glob pattern for input pkl batch files." - ) - parser.add_argument( - "--output_dir", - default=Config.OUTPUT_DIR, - help="Output directory for augmented data." - ) - args = parser.parse_args() - - file_path10 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.r.0021-01-01-00000.nc" - file_path17 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" - file_path18 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" - file_path19 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" - file_path20 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" - file_path21 = "/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc" - - restart_pft_vars = list(dict.fromkeys(Config.dataset_new_1D_PFT_VARIABLES)) - restart_col_1d_vars = list(dict.fromkeys(Config.dataset_new_RESTART_COL_1D_VARS)) - restart_col_2d_vars = list(dict.fromkeys(list(Config.dataset_new_Water_variables) + list(Config.dataset_new_2D_VARIABLES))) - - configured_vars_1d = list(dict.fromkeys(restart_pft_vars + restart_col_1d_vars)) - configured_vars_2d = list(dict.fromkeys(restart_col_2d_vars + Config.SPECIAL_P_VARS)) - - input_files = sorted(glob.glob(args.input_glob)) - if not input_files: - sys.exit(f"No input files found matching pattern: {args.input_glob}") - - os.makedirs(args.output_dir, exist_ok=True) - - ds_restart = nc.Dataset(file_path10) - ds_special_p_restart = nc.Dataset(Config.SPECIAL_P_INPUT_NC) - ds_r_list = [nc.Dataset(fp) for fp in [file_path17, file_path18, file_path19, file_path20, file_path21]] - - try: - restart_tree, restart_coords = build_restart_kdtree(ds_restart) - col_index_map = build_column_index_map(ds_restart) - pft_index_map = build_pft_index_map(ds_restart) - - for fp in input_files: - df = pd.read_pickle(fp) - - force_replace_vars = set(Config.SPECIAL_P_VARS) - missing_other_vars_1d = [ - v for v in configured_vars_1d - if v not in force_replace_vars and not (v in df.columns and f"Y_{v}" in df.columns) - ] - missing_other_vars_2d = [ - v for v in configured_vars_2d - if v not in force_replace_vars and not (v in df.columns and f"Y_{v}" in df.columns) - ] - vars_to_process_1d = missing_other_vars_1d - vars_to_process_2d = list(set(missing_other_vars_2d).union(force_replace_vars)) - - if not vars_to_process_1d and not vars_to_process_2d: - continue - - df_aug = augment_dataframe_with_vars( - df=df, - ds_restart=ds_restart, - ds_special_p_restart=ds_special_p_restart, - ds_r_list=ds_r_list, - restart_tree=restart_tree, - restart_coords=restart_coords, - col_index_map=col_index_map, - pft_index_map=pft_index_map, - vars_1d=vars_to_process_1d, - vars_2d=vars_to_process_2d, - special_p_vars=Config.SPECIAL_P_VARS - ) - - base_name = os.path.basename(fp) - out_name = f"{Config.ENHANCED_PREFIX}{base_name}" - out_path = os.path.join(args.output_dir, out_name) - df_aug.to_pickle(out_path) - - finally: - ds_restart.close() - ds_special_p_restart.close() - for ds in ds_r_list: - try: - ds.close() - except Exception: - pass - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/72_dataset_construction.py b/scripts/training_data_generation/python_scripts/72_dataset_construction.py deleted file mode 100644 index 3a3804e..0000000 --- a/scripts/training_data_generation/python_scripts/72_dataset_construction.py +++ /dev/null @@ -1,640 +0,0 @@ -#!/usr/bin/env python3 -""" -TVA Complete Training Dataset Generation Script -Generates complete PKL files with forcing data, ecosystem variables, and monthly averaging -""" - -import netCDF4 as nc -import numpy as np -import pandas as pd -from scipy.spatial import cKDTree -import os -import sys -import time - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import config - -print("="*80) -print("TVA Complete Training Dataset Generation") -print("="*80) - -# File paths using config and hardcoded paths -file_path1 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/domain_surfdata/TVA_surfdata.TES_SE.4km.1d.NLCD.c241219.nc' -file_path2 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.h0.0021-01-01-00000.nc' - -# Use generated forcing NetCDF files -file_path4 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_FLDS_1980-1999.nc') -file_path5 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_FSDS_1980-1999.nc') -file_path6 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_PRECTmms_1980-1999.nc') -file_path7 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_PSRF_1980-1999.nc') -file_path8 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_QBOT_1980-1999.nc') -file_path9 = os.path.join(config.forcing_netcdf_output_dir, 'TVA_TBOT_1980-1999.nc') - -file_path10 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.r.0021-01-01-00000.nc' - -file_path12 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/kmELM/e3sm_runs/uELM_TVA_finalspinref/run/uELM_TVA_finalspinref.elm.h0.0781-01.nc' -file_path17 = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc' - -# Output directory using config -output_dir = os.path.join(config.output_dir, 'training_dataset_pkl') -os.makedirs(output_dir, exist_ok=True) - -# TVA region coordinates (1D domain) -# TVA: lat [32.33, 37.58], lon [-90.33, -81.71] - -print("Loading NetCDF files...") -start_time = time.time() - -ds1 = nc.Dataset(file_path1) # Surface data (TVA_surfdata.TES_SE.4km.1d.NLCD.c241219.nc) -ds2 = nc.Dataset(file_path2) # History file (uELM_TVA_adspinref.elm.h0.0021-01-01-00000.nc) -ds4 = nc.Dataset(file_path4) # FLDS forcing (TVA_FLDS_1980-1999.nc) -ds5 = nc.Dataset(file_path5) # FSDS forcing (TVA_FSDS_1980-1999.nc) -ds6 = nc.Dataset(file_path6) # PRECTmms forcing (TVA_PRECTmms_1980-1999.nc) -ds7 = nc.Dataset(file_path7) # PSRF forcing (TVA_PSRF_1980-1999.nc) -ds8 = nc.Dataset(file_path8) # QBOT forcing (TVA_QBOT_1980-1999.nc) -ds9 = nc.Dataset(file_path9) # TBOT forcing (TVA_TBOT_1980-1999.nc) -ds10 = nc.Dataset(file_path10) # Restart file (uELM_TVA_adspinref.elm.r.0021-01-01-00000.nc) - -# For Y (future) values - using single file for demo -ds_h0_list = [nc.Dataset(file_path12)] # Future history file (uELM_TVA_finalspinref.elm.h0.0781-01.nc) -ds_r_list = [nc.Dataset(file_path17)] # Future restart file (uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc) - -print(f"✅ All files loaded: {time.time() - start_time:.2f}s") - -# TVA data is 1D (lndgrid), not 2D -lats = ds2.variables['lat'][:] # 1D array -lons = ds2.variables['lon'][:] # 1D array -landmask = ds2.variables['landfrac'][:] # Use landfrac instead of landmask - -# Filter for land gridcells (1D domain) -valid_mask = (landmask > 0) -valid_gridcells = np.where(valid_mask)[0] - -print(f"✅ Spatial filtering completed: {time.time() - start_time:.2f}s") -print(f"Total land gridcells: {len(valid_gridcells)}") -print(f"Latitude range: [{lats.min():.2f}, {lats.max():.2f}]") -print(f"Longitude range: [{lons.min():.2f}, {lons.max():.2f}]") - -batch_size = 1000 -batch_number = 1 - -print(f"\nConfiguration:") -print(f" Batch size: {batch_size}") -print(f" Output directory: {output_dir}") - -print("\nBuilding KDTree index...") -start_time = time.time() - -# Build query coordinates for valid gridcells (1D domain) -query_coords = np.array([(lats[i], lons[i]) for i in valid_gridcells]) - -# Restart file coordinates (1D) -gridcell_lat = ds10.variables['grid1d_lat'][:] -gridcell_lon = ds10.variables['grid1d_lon'][:] -restart_grid_coords = np.vstack((gridcell_lat, gridcell_lon)).T - -restart_tree = cKDTree(restart_grid_coords) -_, all_restart_indices = restart_tree.query(query_coords, k=1) - -# Forcing file coordinates (1D) -forcing_lats = ds4.variables['LATIXY'][:].flatten() -forcing_lons = ds4.variables['LONGXY'][:].flatten() -forcing_grid_coords = np.vstack((forcing_lats, forcing_lons)).T - -forcing_tree = cKDTree(forcing_grid_coords) -_, all_forcing_indices = forcing_tree.query(query_coords, k=1) - -print(f"✅ KDTree index construction completed: {time.time() - start_time:.2f}s") - -# 🚀 KEY OPTIMIZATION: Pre-load all forcing data into memory -print("\n🚀 Pre-loading all forcing data into memory...") -start_time = time.time() - -print(" Loading FLDS data...") -flds_data = ds4.variables['FLDS'][:, 0, :] # 58400 × 11357 -print(f" FLDS shape: {flds_data.shape}, memory: {flds_data.nbytes / 1024**3:.2f} GB") - -print(" Loading PSRF data...") -psrf_data = ds7.variables['PSRF'][:, 0, :] -print(f" PSRF shape: {psrf_data.shape}, memory: {psrf_data.nbytes / 1024**3:.2f} GB") - -print(" Loading FSDS data...") -fsds_data = ds5.variables['FSDS'][:, 0, :] -print(f" FSDS shape: {fsds_data.shape}, memory: {fsds_data.nbytes / 1024**3:.2f} GB") - -print(" Loading QBOT data...") -qbot_data = ds8.variables['QBOT'][:, 0, :] -print(f" QBOT shape: {qbot_data.shape}, memory: {qbot_data.nbytes / 1024**3:.2f} GB") - -print(" Loading PRECTmms data...") -prect_data = ds6.variables['PRECTmms'][:, 0, :] -print(f" PRECTmms shape: {prect_data.shape}, memory: {prect_data.nbytes / 1024**3:.2f} GB") - -print(" Loading TBOT data...") -tbot_data = ds9.variables['TBOT'][:, 0, :] -print(f" TBOT shape: {tbot_data.shape}, memory: {tbot_data.nbytes / 1024**3:.2f} GB") - -total_forcing_memory = (flds_data.nbytes + psrf_data.nbytes + fsds_data.nbytes + - qbot_data.nbytes + prect_data.nbytes + tbot_data.nbytes) / 1024**3 - -print(f"✅ All forcing data pre-loaded: {time.time() - start_time:.2f}s") -print(f" Total memory usage: {total_forcing_memory:.2f} GB") - -# Close forcing NetCDF files (data is now in memory) -ds4.close() -ds5.close() -ds6.close() -ds7.close() -ds8.close() -ds9.close() - -print("✅ Forcing NetCDF files closed, data in memory") - -# Define variable lists -pft_based_vars = [ - 'totvegc', 'deadstemn', 'deadcrootn', 'deadstemp', 'deadcrootp', - 'leafc', 'leafc_storage', 'frootc', 'frootc_storage', - 'deadcrootc', 'deadstemc', 'tlai', - 'leafn', 'leafn_storage', 'frootn','frootn_storage', - 'leafp', 'leafp_storage', 'frootp','frootp_storage', - 'livestemc', 'livestemc_storage', - 'livestemn', 'livestemn_storage', - 'livestemp', 'livestemp_storage', - 'deadcrootc_storage', 'deadstemc_storage', - 'livecrootc', 'livecrootc_storage', - 'deadcrootn_storage', 'deadstemn_storage', - 'livecrootn', 'livecrootn_storage', - 'deadcrootp_storage', 'deadstemp_storage', - 'livecrootp', 'livecrootp_storage' -] - -col_based_1d_vars = ['cwdp', 'totcolp', 'totlitc'] - -col_based_2d_vars = [ - 'cwdn_vr', 'secondp_vr', 'cwdp_vr', 'soil3c_vr', 'soil4c_vr', 'cwdc_vr', - 'soil1c_vr', 'soil1n_vr', 'soil1p_vr', - 'soil2c_vr', 'soil2n_vr', 'soil2p_vr', - 'soil3n_vr', 'soil3p_vr', - 'soil4n_vr', 'soil4p_vr', - 'litr1c_vr', 'litr2c_vr', 'litr3c_vr', - 'litr1n_vr', 'litr2n_vr', 'litr3n_vr', - 'litr1p_vr', 'litr2p_vr', 'litr3p_vr', - 'sminn_vr', 'smin_no3_vr', 'smin_nh4_vr', - 'labilep_vr', 'occlp_vr', 'primp_vr' -] - -all_x_vars = pft_based_vars + col_based_1d_vars + col_based_2d_vars - -# Pre-load X variable data -x_values = {} -for var_name in all_x_vars: - print(f" Loading X variable: {var_name}") - x_values[var_name] = ds10.variables[var_name][:] - -# Pre-load Y variable data -stacked_y_values = {} -for var_name in all_x_vars: - print(f" Loading Y variable: {var_name}") - list_of_arrays = [ds_r.variables[var_name][:] for ds_r in ds_r_list] - stacked_y_values[var_name] = np.stack(list_of_arrays, axis=0) - -print(f"✅ X and Y variables pre-loaded: {time.time() - start_time:.2f}s") - -# Build index mapping -print("\nBuilding index mapping...") -start_time = time.time() - -pft_gridcell_index = ds10.variables['pfts1d_gridcell_index'][:] -column_gridcell_index = ds10.variables['cols1d_gridcell_index'][:] - -pft_map = {} -column_map = {} - -unique_gridcell_ids = np.unique(pft_gridcell_index) -for grid_id in unique_gridcell_ids: - pft_map[grid_id] = np.where(pft_gridcell_index == grid_id)[0] - column_map[grid_id] = np.where(column_gridcell_index == grid_id)[0] - -print(f"✅ Index mapping construction completed: {time.time() - start_time:.2f}s") - -# Process data -print(f"\nStarting to process {len(valid_gridcells)} gridcells...") - -for start_idx in range(0, len(valid_gridcells), batch_size): - end_idx = min(start_idx + batch_size, len(valid_gridcells)) - batch_gridcells = valid_gridcells[start_idx:end_idx] - batch_restart_indices = all_restart_indices[start_idx:end_idx] - batch_forcing_indices = all_forcing_indices[start_idx:end_idx] - - print(f"\nProcessing batch {batch_number}: gridcells {start_idx+1}-{end_idx}") - batch_start_time = time.time() - - data_dict = { - 'landfrac':[], - 'Latitude': [], - 'Longitude': [], - 'FLDS': [], 'PSRF': [], 'FSDS': [], 'QBOT': [], 'PRECTmms': [], 'TBOT': [], - 'LANDFRAC_PFT': [], 'PCT_NATVEG': [], 'AREA': [], 'peatf': [], 'abm': [], - 'SOIL_COLOR': [], 'SOIL_ORDER': [], 'PCT_NAT_PFT': [], 'PCT_SAND': [], - 'soil3c_vr': [], 'soil4c_vr': [], 'cwdc_vr': [], 'deadcrootc': [], 'deadstemc': [],'tlai': [], - 'GPP': [], - 'Y_soil3c_vr': [], 'Y_soil4c_vr': [], 'Y_cwdc_vr': [], 'Y_deadcrootc': [], 'Y_deadstemc': [], 'Y_tlai': [], - 'Y_GPP': [], - 'SCALARAVG_vr': [], - 'PCT_CLAY': [], - 'SNOWDP': [], - 'H2OSOI_10CM': [], - 'HR': [], 'AR': [], 'NPP': [], 'COL_FIRE_CLOSS': [], - 'Y_HR': [], 'Y_AR': [], 'Y_NPP': [], 'Y_COL_FIRE_CLOSS': [], - 'OCCLUDED_P': [], - 'SECONDARY_P': [], - 'LABILE_P': [], - 'APATITE_P': [], - - 'cwdn_vr': [], 'secondp_vr': [], 'cwdp_vr': [],'cwdp': [], 'totcolp': [], 'totvegc': [], 'deadstemn': [], 'deadcrootn': [], - 'deadstemp': [], 'deadcrootp': [], 'leafc': [], 'leafc_storage': [], 'frootc': [], 'frootc_storage': [], - 'Y_cwdn_vr': [], 'Y_secondp_vr': [], 'Y_cwdp_vr': [], 'Y_cwdp': [], 'Y_totcolp': [], 'Y_totvegc': [], 'Y_deadstemn': [], 'Y_deadcrootn': [], - 'Y_deadstemp': [], 'Y_deadcrootp': [], 'Y_leafc': [], 'Y_leafc_storage': [], 'Y_frootc': [], 'Y_frootc_storage': [], - 'totlitc': [], - - 'leafn': [], 'leafn_storage': [], 'frootn': [],'frootn_storage': [], - 'leafp': [], 'leafp_storage': [], 'frootp': [],'frootp_storage': [], - 'livestemc': [], 'livestemc_storage': [], - 'livestemn': [], 'livestemn_storage': [], - 'livestemp': [], 'livestemp_storage': [], - 'labilep_vr': [], 'occlp_vr': [], 'primp_vr': [], - - 'deadcrootc_storage': [], 'deadstemc_storage': [], - 'livecrootc': [], 'livecrootc_storage': [], - 'deadcrootn_storage': [], 'deadstemn_storage': [], - 'livecrootn': [], 'livecrootn_storage': [], - 'deadcrootp_storage': [], 'deadstemp_storage': [], - 'livecrootp': [], 'livecrootp_storage': [], - - 'Y_leafn': [], 'Y_leafn_storage': [], 'Y_frootn': [],'Y_frootn_storage': [], - 'Y_leafp': [], 'Y_leafp_storage': [], 'Y_frootp': [],'Y_frootp_storage': [], - 'Y_livestemc': [], 'Y_livestemc_storage': [], - 'Y_livestemn': [], 'Y_livestemn_storage': [], - 'Y_livestemp': [], 'Y_livestemp_storage': [], - 'Y_labilep_vr': [], 'Y_occlp_vr': [], 'Y_primp_vr': [], - - 'Y_deadcrootc_storage': [], 'Y_deadstemc_storage': [], - 'Y_livecrootc': [], 'Y_livecrootc_storage': [], - 'Y_deadcrootn_storage': [], 'Y_deadstemn_storage': [], - 'Y_livecrootn': [], 'Y_livecrootn_storage': [], - 'Y_deadcrootp_storage': [], 'Y_deadstemp_storage': [], - 'Y_livecrootp': [], 'Y_livecrootp_storage': [], - 'Y_totlitc': [], - # 'H2OCAN': [], 'T_VEG': [], 'T10_VALUE': [], - # 'Y_H2OCAN': [], 'Y_T_VEG': [], 'Y_T10_VALUE': [], - # 'H2OSFC': [], 'H2OSNO': [], 'TH2OSFC': [], 'T_GRND': [], 'T_GRND_R': [], 'T_GRND_U': [], - # 'Y_H2OSFC': [], 'Y_H2OSNO': [], 'Y_TH2OSFC': [], 'Y_T_GRND': [], 'Y_T_GRND_R': [], 'Y_T_GRND_U': [], - # 'H2OSOI_LIQ': [], 'H2OSOI_ICE': [], 'T_SOISNO': [], 'LAKE_SOILC': [], 'T_LAKE': [], - # 'Y_H2OSOI_LIQ': [], 'Y_H2OSOI_ICE': [], 'Y_T_SOISNO': [], 'Y_LAKE_SOILC': [], 'Y_T_LAKE': [], - # 'taf': [], - # 'Y_taf': [], - # 'TS_TOPO': [], - # 'Y_TS_TOPO': [], - 'soil1c_vr': [], 'soil1n_vr': [], 'soil1p_vr': [], - 'soil2c_vr': [], 'soil2n_vr': [], 'soil2p_vr': [], - 'soil3n_vr': [], 'soil3p_vr': [], - 'soil4n_vr': [], 'soil4p_vr': [], - 'litr1c_vr': [], 'litr2c_vr': [], 'litr3c_vr': [], - 'litr1n_vr': [], 'litr2n_vr': [], 'litr3n_vr': [], - 'litr1p_vr': [], 'litr2p_vr': [], 'litr3p_vr': [], - 'sminn_vr': [], 'smin_no3_vr': [], 'smin_nh4_vr': [], - 'Y_soil1c_vr': [], 'Y_soil1n_vr': [], 'Y_soil1p_vr': [], - 'Y_soil2c_vr': [], 'Y_soil2n_vr': [], 'Y_soil2p_vr': [], - 'Y_soil3n_vr': [], 'Y_soil3p_vr': [], - 'Y_soil4n_vr': [], 'Y_soil4p_vr': [], - 'Y_litr1c_vr': [], 'Y_litr2c_vr': [], 'Y_litr3c_vr': [], - 'Y_litr1n_vr': [], 'Y_litr2n_vr': [], 'Y_litr3n_vr': [], - 'Y_litr1p_vr': [], 'Y_litr2p_vr': [], 'Y_litr3p_vr': [], - 'Y_sminn_vr': [], 'Y_smin_no3_vr': [], 'Y_smin_nh4_vr': [] - } - - - - for k, gridcell_idx in enumerate(batch_gridcells): - if k % 100 == 0: - print(f" Processing gridcell {k}/{len(batch_gridcells)} (idx={gridcell_idx})") - - # Get indices - restart_idx = batch_restart_indices[k] - gridcell_id = restart_idx + 1 - - pft_indices_for_cell = pft_map.get(gridcell_id, []) - column_indices_for_cell = column_map.get(gridcell_id, []) - - for var_name in pft_based_vars: - x_val = x_values[var_name][pft_indices_for_cell] - data_dict[var_name].append(x_val.tolist()) - - y_slice = stacked_y_values[var_name][:, pft_indices_for_cell] - avg_y_val = np.mean(y_slice, axis=0) - data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) - - for var_name in col_based_1d_vars: - x_val = x_values[var_name][column_indices_for_cell] - data_dict[var_name].append(x_val.tolist()) - - y_slice = stacked_y_values[var_name][:, column_indices_for_cell] - avg_y_val = np.mean(y_slice, axis=0) - data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) - - for var_name in col_based_2d_vars: - x_val = x_values[var_name][column_indices_for_cell, :] - data_dict[var_name].append(x_val.tolist()) - - y_slice = stacked_y_values[var_name][:, column_indices_for_cell, :] - avg_y_val = np.mean(y_slice, axis=0) - data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) - - # landunit_indices_for_cell = landunit_map.get(gridcell_id, []) - # for var_name in landunit_based_vars: - # x_val = x_values[var_name][landunit_indices_for_cell] - # data_dict[var_name].append(x_val.tolist()) - - # y_slice = stacked_y_values[var_name][:, landunit_indices_for_cell] - # avg_y_val = np.mean(y_slice, axis=0) - # data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) - - # topounit_indices_for_cell = topounit_map.get(gridcell_id, []) - # for var_name in topounit_based_vars: - # x_val = x_values[var_name][topounit_indices_for_cell] - # data_dict[var_name].append(x_val.tolist()) - - # y_slice = stacked_y_values[var_name][:, topounit_indices_for_cell] - # avg_y_val = np.mean(y_slice, axis=0) - # data_dict[f'Y_{var_name}'].append(avg_y_val.tolist()) - - data_dict['landfrac'].append(ds2.variables['landfrac'][gridcell_idx]) - data_dict['Latitude'].append(lats[gridcell_idx]) - data_dict['Longitude'].append(lons[gridcell_idx]) - data_dict['LANDFRAC_PFT'].append(ds1.variables['LANDFRAC_PFT'][gridcell_idx]) - data_dict['PCT_NATVEG'].append(ds1.variables['PCT_NATVEG'][gridcell_idx]) - data_dict['AREA'].append(ds1.variables['AREA'][gridcell_idx]) - data_dict['peatf'].append(ds1.variables['peatf'][gridcell_idx]) - data_dict['abm'].append(ds1.variables['abm'][gridcell_idx]) - data_dict['SOIL_COLOR'].append(ds1.variables['SOIL_COLOR'][gridcell_idx]) - data_dict['SOIL_ORDER'].append(ds1.variables['SOIL_ORDER'][gridcell_idx]) - data_dict['PCT_SAND'].append(ds1.variables['PCT_SAND'][:, gridcell_idx]) - data_dict['PCT_NAT_PFT'].append(ds1.variables['PCT_NAT_PFT'][:, gridcell_idx]) - - data_dict['OCCLUDED_P'].append(ds1.variables['OCCLUDED_P'][gridcell_idx]) - data_dict['SECONDARY_P'].append(ds1.variables['SECONDARY_P'][gridcell_idx]) - data_dict['LABILE_P'].append(ds1.variables['LABILE_P'][gridcell_idx]) - data_dict['APATITE_P'].append(ds1.variables['APATITE_P'][gridcell_idx]) - - - data_dict['GPP'].append(ds2.variables['GPP'][0, gridcell_idx]) - data_dict['SCALARAVG_vr'].append(ds2.variables['SCALARAVG_vr'][0, :, gridcell_idx]) - data_dict['HR'].append(ds2.variables['HR'][0, gridcell_idx]) - data_dict['AR'].append(ds2.variables['AR'][0, gridcell_idx]) - data_dict['NPP'].append(ds2.variables['NPP'][0, gridcell_idx]) - # COL_FIRE_CLOSS may not exist in TVA history, use FIRE instead if available - if 'COL_FIRE_CLOSS' in ds2.variables: - data_dict['COL_FIRE_CLOSS'].append(ds2.variables['COL_FIRE_CLOSS'][0, gridcell_idx]) - elif 'FIRE' in ds2.variables: - data_dict['COL_FIRE_CLOSS'].append(ds2.variables['FIRE'][0, gridcell_idx]) - else: - data_dict['COL_FIRE_CLOSS'].append(0.0) - - data_dict['SNOWDP'].append(ds2.variables['SNOWDP'][0, gridcell_idx]) - data_dict['H2OSOI_10CM'].append(ds2.variables['H2OSOI'][0,3, gridcell_idx]) - data_dict['PCT_CLAY'].append(ds1.variables['PCT_CLAY'][:, gridcell_idx]) - - h0_gpp_vals = [] - for ds_h0 in ds_h0_list: - h0_gpp_vals.append(ds_h0.variables['GPP'][0, gridcell_idx]) - avg_h0_gpp = np.mean(h0_gpp_vals) - data_dict['Y_GPP'].append(avg_h0_gpp) - - h0_HR_vals = [] - for ds_h0 in ds_h0_list: - h0_HR_vals.append(ds_h0.variables['HR'][0, gridcell_idx]) - avg_h0_HR = np.mean(h0_HR_vals) - data_dict['Y_HR'].append(avg_h0_HR) - - h0_AR_vals = [] - for ds_h0 in ds_h0_list: - h0_AR_vals.append(ds_h0.variables['AR'][0, gridcell_idx]) - avg_h0_AR = np.mean(h0_AR_vals) - data_dict['Y_AR'].append(avg_h0_AR) - - h0_NPP_vals = [] - for ds_h0 in ds_h0_list: - h0_NPP_vals.append(ds_h0.variables['NPP'][0, gridcell_idx]) - avg_h0_NPP = np.mean(h0_NPP_vals) - data_dict['Y_NPP'].append(avg_h0_NPP) - - h0_COL_FIRE_CLOSS_vals = [] - for ds_h0 in ds_h0_list: - # COL_FIRE_CLOSS may not exist in TVA history - if 'COL_FIRE_CLOSS' in ds_h0.variables: - h0_COL_FIRE_CLOSS_vals.append(ds_h0.variables['COL_FIRE_CLOSS'][0, gridcell_idx]) - elif 'FIRE' in ds_h0.variables: - h0_COL_FIRE_CLOSS_vals.append(ds_h0.variables['FIRE'][0, gridcell_idx]) - else: - h0_COL_FIRE_CLOSS_vals.append(0.0) - avg_h0_COL_FIRE_CLOSS = np.mean(h0_COL_FIRE_CLOSS_vals) - data_dict['Y_COL_FIRE_CLOSS'].append(avg_h0_COL_FIRE_CLOSS) - - # Get forcing index - forcing_idx = batch_forcing_indices[k] - - # 🚀 Fast access to forcing data from memory - data_dict['FLDS'].append(flds_data[:, forcing_idx]) - data_dict['PSRF'].append(psrf_data[:, forcing_idx]) - data_dict['FSDS'].append(fsds_data[:, forcing_idx]) - data_dict['QBOT'].append(qbot_data[:, forcing_idx]) - data_dict['PRECTmms'].append(prect_data[:, forcing_idx]) - data_dict['TBOT'].append(tbot_data[:, forcing_idx]) - # Create DataFrame and save - print(f" Creating DataFrame...") - df_batch = pd.DataFrame(data_dict) - - print(f" Saving to disk...") - batch_save_path = f"{output_dir}/training_data_batch_{batch_number:02d}.pkl" - df_batch.to_pickle(batch_save_path) - - batch_time = time.time() - batch_start_time - print(f"✅ Batch {batch_number} completed: {batch_time:.2f}s") - print(f" Path: {batch_save_path}") - print(f" Shape: {df_batch.shape}") - print(f" Columns: {len(df_batch.columns)}") - print(f" Forcing data length: {len(df_batch['FLDS'].iloc[0])}") - - batch_number += 1 - -# Cleanup NetCDF files -print(f"\n{'='*80}") -print("Cleaning up NetCDF files...") -ds1.close() -ds2.close() -ds10.close() -for ds_h0 in ds_h0_list: - ds_h0.close() -for ds_r in ds_r_list: - ds_r.close() - -print("✅ All NetCDF files closed") - -print(f"\n🎉 Complete dataset PKL generation completed!") -print(f"Total batches: {batch_number - 1}") -print(f"Output directory: {output_dir}") -print(f"Each PKL file contains: {len(data_dict)} variables") -print(f"Forcing data: 58400 time steps") -print(f"Optimization strategy: Pre-load all forcing data to memory, avoid repeated NetCDF access") - -# ============================================================================= -# POST-PROCESSING: MONTHLY AVERAGING (Simplified) -# ============================================================================= - -print(f"\n{'='*80}") -print("POST-PROCESSING: Monthly Averaging") -print(f"{'='*80}") - -# Import glob for file processing -import glob - -# Get all generated PKL files -input_files = sorted(glob.glob(f'{output_dir}/training_data_batch_*.pkl')) -print(f"Found {len(input_files)} PKL files for post-processing") - -# TVA data parameters (20 years, 3-hour interval) -time_series_length = 58400 # 20 years × 365 days × 8 steps/day -steps_per_day = 8 # 3-hour interval = 8 steps/day -days_per_year = 365 -years_in_data = 20 # 1980-1999 -months_per_year = 12 -days_per_month = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31] - -print(f"TVA data parameters:") -print(f" Time series length: {time_series_length}") -print(f" Steps per day: {steps_per_day}") -print(f" Years: {years_in_data}") -print(f" Expected monthly values: {years_in_data * months_per_year}") - -def calculate_monthly_avg(time_series): - """ - Calculate monthly averages from high-resolution time series - """ - if not isinstance(time_series, (list, np.ndarray)): - return [] - - if len(time_series) != time_series_length: - return [] - - monthly_averages = [] - start_idx = 0 - - for year in range(years_in_data): - for month_idx, month_days in enumerate(days_per_month): - # Handle leap year February - if year % 4 == 0 and month_idx == 1: # Leap year February - month_days = 29 - - end_idx = start_idx + month_days * steps_per_day - monthly_avg = np.mean(time_series[start_idx:end_idx]) - monthly_averages.append(monthly_avg) - start_idx = end_idx - - return monthly_averages - -# Define columns to process -time_series_columns = ['FLDS', 'PSRF', 'FSDS', 'QBOT', 'PRECTmms', 'TBOT'] -single_value_columns = [ - 'landfrac', 'LANDFRAC_PFT', 'PCT_NATVEG', 'AREA', 'peatf', 'abm', - 'SOIL_COLOR', 'SOIL_ORDER', 'GPP', 'SNOWDP', 'H2OSOI_10CM', - 'Y_GPP', 'HR', 'AR', 'NPP', 'COL_FIRE_CLOSS', - 'Y_HR', 'Y_AR', 'Y_NPP', 'Y_COL_FIRE_CLOSS', - 'OCCLUDED_P', 'SECONDARY_P', 'LABILE_P', 'APATITE_P' -] -list_like_columns = ['PCT_NAT_PFT', 'PCT_SAND', 'SCALARAVG_vr', 'PCT_CLAY'] - -print(f"\nStarting post-processing...") - -# Process all files -print(f"Processing all {len(input_files)} files...") - -for file_idx, file_path in enumerate(input_files, 1): - print(f"\nProcessing file {file_idx}/{len(input_files)}: {os.path.basename(file_path)}") - - try: - # Read PKL file - df = pd.read_pickle(file_path) - print(f" Original shape: {df.shape}") - print(f" Original columns: {len(df.columns)}") - - # Process time series columns (convert to monthly averages) - print(f" Processing time series columns...") - for col in time_series_columns: - if col in df.columns: - print(f" Processing {col}...") - df[col] = df[col].apply(calculate_monthly_avg) - - # Verify processing results - sample_data = df[col].apply(lambda x: x if isinstance(x, list) else []) - lengths = sample_data.apply(len).unique() - print(f" {col} monthly values length: {lengths}") - - # Process single value columns - print(f" Processing single value columns...") - for col in single_value_columns: - if col in df.columns: - df[col] = df[col].astype(str).str.strip() - df[col] = pd.to_numeric(df[col], errors='coerce') - - # Process list columns - print(f" Processing list columns...") - for col in list_like_columns: - if col in df.columns: - print(f" Expanding {col}...") - expanded_cols = df[col].apply(pd.Series).fillna(0) - expanded_cols = expanded_cols.add_prefix(f"{col}_") - df = df.drop(col, axis=1).join(expanded_cols) - - # Reorder columns - y_columns = [col for col in df.columns if col.startswith('Y_')] - other_columns = [col for col in df.columns if not col.startswith('Y_')] - df = df[other_columns + y_columns] - - # Save processed file - output_file = f"{output_dir}/monthly_{os.path.basename(file_path)}" - df.to_pickle(output_file) - - print(f" ✅ Post-processing completed") - print(f" Processed shape: {df.shape}") - print(f" Processed columns: {len(df.columns)}") - print(f" Saved to: {output_file}") - - # Display sample data - print(f" Sample data:") - print(f" FLDS monthly values length: {len(df['FLDS'].iloc[0]) if 'FLDS' in df.columns else 'N/A'}") - print(f" First 3 coordinates: {df[['Latitude', 'Longitude']].head(3).values.tolist()}") - - except Exception as e: - print(f" ❌ Error processing {os.path.basename(file_path)}: {e}") - continue - -print(f"\n{'='*80}") -print("POST-PROCESSING COMPLETED!") -print(f"Input directory: {output_dir}") -print(f"Output directory: {output_dir}") -print(f"Processed files: {len(input_files)}") - -# Check output files -output_files = glob.glob(f'{output_dir}/monthly_*.pkl') -print(f"Generated files: {len(output_files)}") - -if output_files: - print(f"\nOutput file list:") - for file in output_files: - file_size = os.path.getsize(file) / (1024**3) # GB - print(f" {os.path.basename(file)}: {file_size:.2f} GB") - -print("="*80) - diff --git a/scripts/training_data_generation/python_scripts/72_dataset_forcing_only.py b/scripts/training_data_generation/python_scripts/72_dataset_forcing_only.py deleted file mode 100644 index 1ceda49..0000000 --- a/scripts/training_data_generation/python_scripts/72_dataset_forcing_only.py +++ /dev/null @@ -1,200 +0,0 @@ -#!/usr/bin/env python3 -""" -TVA Forcing Data Only - PKL Generation Script (Optimized with List Format) -Extract only forcing data and basic geographical information -Automatically converts forcing variables to list format for training compatibility -""" - -import netCDF4 as nc -import numpy as np -import pandas as pd -from scipy.spatial import cKDTree -import os -import sys -import time - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import config - -print("="*80) -print("TVA Forcing Data Only - PKL Generation with List Format Conversion") -print("="*80) - -# File paths using config -history_file = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.h0.0021-01-01-00000.nc' - -# Forcing files - use newly generated ones -forcing_files = { - 'FLDS': os.path.join(config.forcing_netcdf_output_dir, 'TVA_FLDS_1980-1999.nc'), - 'FSDS': os.path.join(config.forcing_netcdf_output_dir, 'TVA_FSDS_1980-1999.nc'), - 'PSRF': os.path.join(config.forcing_netcdf_output_dir, 'TVA_PSRF_1980-1999.nc'), - 'QBOT': os.path.join(config.forcing_netcdf_output_dir, 'TVA_QBOT_1980-1999.nc'), - 'PRECTmms': os.path.join(config.forcing_netcdf_output_dir, 'TVA_PRECTmms_1980-1999.nc'), - 'TBOT': os.path.join(config.forcing_netcdf_output_dir, 'TVA_TBOT_1980-1999.nc'), -} - -# Output directory -output_dir = config.forcing_pkl_output_dir -os.makedirs(output_dir, exist_ok=True) - -# Configuration -batch_size = 1000 -batch_number = 1 - -print(f"Configuration:") -print(f" Batch size: {batch_size}") -print(f" Output directory: {output_dir}") -print(f" History file: {history_file}") -for var, path in forcing_files.items(): - print(f" {var} file: {path}") - -# Load NetCDF files -print("\nLoading NetCDF files...") -start_time = time.time() - -# Load history file for coordinates -ds_history = nc.Dataset(history_file) - -# Load forcing files -forcing_datasets = {} -for var, file_path in forcing_files.items(): - print(f" Loading {var}...") - forcing_datasets[var] = nc.Dataset(file_path) - -print(f"✅ All files loaded: {time.time() - start_time:.2f}s") - -# Get coordinate information -print("\nSetting up spatial filtering...") -lats = ds_history.variables['lat'][:] -lons = ds_history.variables['lon'][:] -landmask = ds_history.variables['landfrac'][:] - -# Filter for land gridcells -valid_mask = (landmask > 0) -valid_gridcells = np.where(valid_mask)[0] - -print(f"Total land gridcells: {len(valid_gridcells)}") -print(f"Latitude range: [{lats.min():.2f}, {lats.max():.2f}]") -print(f"Longitude range: [{lons.min():.2f}, {lons.max():.2f}]") - -# Build KDTree index -print("\nBuilding KDTree index...") -query_coords = np.array([(lats[i], lons[i]) for i in valid_gridcells]) - -# Forcing file coordinates -forcing_lats = forcing_datasets['FLDS'].variables['LATIXY'][:].flatten() -forcing_lons = forcing_datasets['FLDS'].variables['LONGXY'][:].flatten() -forcing_coords = np.vstack((forcing_lats, forcing_lons)).T - -forcing_tree = cKDTree(forcing_coords) -_, all_forcing_indices = forcing_tree.query(query_coords, k=1) - -print("✅ Forcing mapping completed") - -# 🚀 KEY OPTIMIZATION: Pre-load all forcing data into memory -print("\n🚀 Pre-loading all forcing data into memory...") -start_time = time.time() - -forcing_data = {} -for var in forcing_files.keys(): - print(f" Loading {var} data...") - # Load all data: time × nj × ni (58400 × 1 × 11357) - forcing_data[var] = forcing_datasets[var].variables[var][:, 0, :] - print(f" {var} shape: {forcing_data[var].shape}, memory: {forcing_data[var].nbytes / 1024**3:.2f} GB") - -total_forcing_memory = sum(data.nbytes for data in forcing_data.values()) / 1024**3 - -print(f"✅ All forcing data pre-loaded: {time.time() - start_time:.2f}s") -print(f" Total memory usage: {total_forcing_memory:.2f} GB") - -# Close forcing NetCDF files (data is now in memory) -for ds in forcing_datasets.values(): - ds.close() - -print("✅ Forcing NetCDF files closed, data in memory") - -# Process data -print(f"\nStarting to process {len(valid_gridcells)} gridcells...") - -for start_idx in range(0, len(valid_gridcells), batch_size): - end_idx = min(start_idx + batch_size, len(valid_gridcells)) - batch_gridcells = valid_gridcells[start_idx:end_idx] - batch_forcing_indices = all_forcing_indices[start_idx:end_idx] - - print(f"\nProcessing batch {batch_number}: gridcells {start_idx+1}-{end_idx}") - batch_start_time = time.time() - - # Initialize data dictionary - only forcing data and basic geographical information - data_dict = { - 'landfrac': [], - 'Latitude': [], - 'Longitude': [], - 'FLDS': [], - 'PSRF': [], - 'FSDS': [], - 'QBOT': [], - 'PRECTmms': [], - 'TBOT': [], - } - - # Process each gridcell - for k, gridcell_idx in enumerate(batch_gridcells): - if k % 100 == 0: - print(f" Processing gridcell {k}/{len(batch_gridcells)} (idx={gridcell_idx})") - - # Get forcing index - forcing_idx = batch_forcing_indices[k] - - # Basic geographical information - data_dict['landfrac'].append(float(landmask[gridcell_idx])) - data_dict['Latitude'].append(float(lats[gridcell_idx])) - data_dict['Longitude'].append(float(lons[gridcell_idx])) - - # Forcing data (directly from pre-loaded memory arrays) - data_dict['FLDS'].append(forcing_data['FLDS'][:, forcing_idx]) - data_dict['PSRF'].append(forcing_data['PSRF'][:, forcing_idx]) - data_dict['FSDS'].append(forcing_data['FSDS'][:, forcing_idx]) - data_dict['QBOT'].append(forcing_data['QBOT'][:, forcing_idx]) - data_dict['PRECTmms'].append(forcing_data['PRECTmms'][:, forcing_idx]) - data_dict['TBOT'].append(forcing_data['TBOT'][:, forcing_idx]) - - # Create DataFrame and save - print(f" Creating DataFrame...") - df_batch = pd.DataFrame(data_dict) - - # Convert forcing variables to list format for training compatibility - print(f" Converting forcing variables to list format...") - forcing_vars = ['FLDS', 'PSRF', 'FSDS', 'QBOT', 'PRECTmms', 'TBOT'] - for var in forcing_vars: - df_batch[var] = df_batch[var].apply(lambda x: x.tolist() if hasattr(x, 'tolist') else x) - - print(f" Saving to disk...") - batch_save_path = f"{output_dir}/TVA_forcing_batch_{batch_number:02d}.pkl" - df_batch.to_pickle(batch_save_path) - - batch_time = time.time() - batch_start_time - print(f"✅ Batch {batch_number} completed in {batch_time:.2f}s:") - print(f" Path: {batch_save_path}") - print(f" Shape: {df_batch.shape}") - print(f" Columns: {len(df_batch.columns)}") - print(f" Forcing data length: {len(df_batch['FLDS'].iloc[0])}") - print(f" Data format: All forcing variables converted to list format") - - batch_number += 1 - -# Cleanup -print(f"\n{'='*80}") -print("Cleaning up...") -ds_history.close() - -print("✅ All NetCDF files closed") - -print(f"\n🎉 Forcing data PKL generation completed!") -print(f"Total batches: {batch_number - 1}") -print(f"Output directory: {output_dir}") -print(f"Each PKL file contains: 9 variables (landfrac, Latitude, Longitude, 6 forcing variables)") -print(f"Forcing data: 58400 time steps (3-hour resolution, 20 years)") -print(f"Data format: All forcing variables automatically converted to list format") -print(f"Training compatibility: Ready for machine learning training pipelines") -print(f"Memory optimization: All forcing data pre-loaded for faster processing") \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_FLDS_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_FLDS_20years.py deleted file mode 100644 index ab09952..0000000 --- a/scripts/training_data_generation/python_scripts/construct_TVA_FLDS_20years.py +++ /dev/null @@ -1,183 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate TVA FLDS forcing data (1980-1999, 20 years, 3-hour resolution) -Corresponding to crujra.v2.5.5d_FLDS_1901-2023_z01.nc - -Author: Zhuowei Gu -Date: 2024 -""" - -import xarray as xr -import os -import cftime -import numpy as np -import sys -import traceback -import datetime -from pathlib import Path - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import config - -# --- Configuration --- -data_dir = config.forcing_raw_data_path -out_dir = config.forcing_netcdf_output_dir -os.makedirs(out_dir, exist_ok=True) - -start_year = 1980 -end_year = 1999 -final_output_file = os.path.join(out_dir, f"TVA_FLDS_{start_year}-{end_year}.nc") - -print("="*80) -print("Generate TVA FLDS forcing data") -print("1. Merge monthly files (1980-1999, 240 months)") -print("2. Correct time axis discontinuities") -print("3. Maintain 3-hour resolution (no downsampling)") -print("="*80) - -print(f"Source data directory: {data_dir}") -print(f"Target output file: {final_output_file}") - -# --- Find and build all monthly file list --- -print("Searching for monthly files...") -all_monthly_files = [] - -for year in range(start_year, end_year + 1): - for month in range(1, 13): - file_name = f"clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc" - file_path = os.path.join(data_dir, file_name) - if os.path.exists(file_path): - all_monthly_files.append(file_path) - else: - print(f" Warning: File {file_name} does not exist, skipping.") - -if not all_monthly_files: - print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") - sys.exit(1) - -print(f"Found {len(all_monthly_files)} valid monthly files.") - -# --- Define Dask chunks --- -dask_chunks = {'time': 366*8} - -try: - with xr.open_mfdataset( - all_monthly_files, - combine='nested', - concat_dim='time', - decode_times=False, - chunks=dask_chunks, - parallel=False, - ) as ds: - - # === Separate static variables === - print("Separating static coordinate/ID variables...") - static_var_names = ['gridID', 'LONGXY', 'LATIXY'] - static_data = {} - - for var_name in static_var_names: - if var_name in ds: - if 'time' in ds[var_name].dims: - static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() - else: - static_data[var_name] = ds[var_name].load() - - # === Process FLDS variable === - time_varying_vars = ['FLDS'] - if 'FLDS' not in ds.data_vars: - print("Error: FLDS variable not found.") - sys.exit(1) - - print(f"Processing variables: {time_varying_vars}") - ds_temporal = ds[['time'] + time_varying_vars] - - # --- Time axis correction --- - print("Loading raw time coordinates...") - time_values_raw = ds_temporal['time'].load().values - units = ds_temporal['time'].attrs['units'] - calendar = ds_temporal['time'].attrs.get('calendar', 'standard') - if calendar.lower() == 'no_leap': - calendar = 'noleap' - - print(f" Time points: {len(time_values_raw)}") - - print("Starting time coordinate correction...") - time_values_corrected = np.copy(time_values_raw).astype(float) - cumulative_offset_days = 0.0 - expected_step_days = 3.0 / 24.0 - jump_count = 0 - - for i in range(len(time_values_raw) - 1): - corrected_i = time_values_raw[i] + cumulative_offset_days - time_values_corrected[i] = corrected_i - raw_diff = time_values_raw[i+1] - time_values_raw[i] - - if raw_diff < expected_step_days * 0.5: - jump_count += 1 - expected_next_corrected_value = corrected_i + expected_step_days - offset_needed = expected_next_corrected_value - time_values_raw[i+1] - cumulative_offset_days = offset_needed - - time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days - - print(f"Time correction completed. Corrected {jump_count} jumps.") - - print("Decoding corrected time values...") - try: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) - except ValueError: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar) - - # --- Check time monotonicity --- - print("Checking time monotonicity...") - if len(dates) >= 2: - diffs_corrected = np.diff(dates) - zero_timedelta = datetime.timedelta(0) - problem_indices = np.where(diffs_corrected <= zero_timedelta)[0] - - if len(problem_indices) > 0: - print(f"Error! Corrected time is still not monotonic!") - sys.exit(1) - else: - print("✓ Time coordinate check passed.") - - # --- Create dataset with corrected time --- - ds_corrected_time = ds_temporal.copy(deep=False) - ds_corrected_time['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) - ds_corrected_time['time'].encoding['units'] = units - ds_corrected_time['time'].encoding['calendar'] = calendar - - # === Merge static variables === - for var_name, data_array in static_data.items(): - ds_corrected_time[var_name] = data_array - - final_dataset = ds_corrected_time - print("Final dataset:") - print(final_dataset) - - # --- Write to file --- - print(f"Writing to file: {final_output_file}") - output_encoding = { - 'FLDS': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, - 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} - } - - if 'gridID' in final_dataset: - output_encoding['gridID'] = {'dtype': final_dataset['gridID'].dtype} - if 'LONGXY' in final_dataset: - output_encoding['LONGXY'] = {'dtype': final_dataset['LONGXY'].dtype, '_FillValue': np.nan} - if 'LATIXY' in final_dataset: - output_encoding['LATIXY'] = {'dtype': final_dataset['LATIXY'].dtype, '_FillValue': np.nan} - - final_dataset.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) - print(f"✓ Successfully generated: {final_output_file}") - print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") - -except Exception as e: - print(f"\nError: {e}") - traceback.print_exc() - sys.exit(1) - -print("\nScript execution completed.") - diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_FSDS_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_FSDS_20years.py deleted file mode 100644 index eb8b7ea..0000000 --- a/scripts/training_data_generation/python_scripts/construct_TVA_FSDS_20years.py +++ /dev/null @@ -1,152 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate TVA FSDS forcing data (1980-1999, 20 years, 3-hour resolution) -Corresponding to crujra.v2.5.5d_FSDS_1901-2023_z01.nc - -Author: Zhuowei Gu -Date: 2024 -""" - -import xarray as xr -import os -import cftime -import numpy as np -import sys -import traceback -import datetime -from pathlib import Path - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import config - -# --- Configuration --- -data_dir = config.forcing_raw_data_path -out_dir = config.forcing_netcdf_output_dir -os.makedirs(out_dir, exist_ok=True) - -start_year = 1980 -end_year = 1999 -final_output_file = os.path.join(out_dir, f"TVA_FSDS_{start_year}-{end_year}.nc") - -print("="*80) -print("Generate TVA FSDS forcing data") -print("1. Merge monthly files (1980-1999, 240 months)") -print("2. Correct time axis discontinuities") -print("3. Maintain 3-hour resolution (no downsampling)") -print("="*80) - -print(f"Source data directory: {data_dir}") -print(f"Target output file: {final_output_file}") - -# --- Find and build all monthly file list --- -print("Searching for monthly files...") -all_monthly_files = [] -for year in range(start_year, end_year + 1): - for month in range(1, 13): - file_name = f"clmforc.Daymet.km.1d.Solr.{year}-{month:02d}.nc" - file_path = os.path.join(data_dir, file_name) - if os.path.exists(file_path): - all_monthly_files.append(file_path) - else: - print(f" Warning: File {file_name} does not exist, skipping.") - -if not all_monthly_files: - print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") - sys.exit(1) - -print(f"Found {len(all_monthly_files)} valid monthly files.") - -# --- Define Dask chunks --- -dask_chunks = {'time': 366*8} - -try: - with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', - decode_times=False, chunks=dask_chunks, parallel=False) as ds: - - # === Separate static variables === - print("Separating static coordinate/ID variables...") - static_data = {} - for var_name in ['gridID', 'LONGXY', 'LATIXY']: - if var_name in ds: - if 'time' in ds[var_name].dims: - static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() - else: - static_data[var_name] = ds[var_name].load() - - # === Process FSDS variable === - if 'FSDS' not in ds.data_vars: - print("Error: FSDS variable not found.") - sys.exit(1) - - print("Processing FSDS variable...") - ds_temporal = ds[['time', 'FSDS']] - - # --- Time axis correction --- - print("Loading raw time coordinates...") - time_values_raw = ds_temporal['time'].load().values - units = ds_temporal['time'].attrs['units'] - calendar = ds_temporal['time'].attrs.get('calendar', 'standard') - if calendar.lower() == 'no_leap': - calendar = 'noleap' - - print(f" Time points: {len(time_values_raw)}") - print("Starting time coordinate correction...") - - time_values_corrected = np.copy(time_values_raw).astype(float) - cumulative_offset_days = 0.0 - expected_step_days = 3.0 / 24.0 - jump_count = 0 - - for i in range(len(time_values_raw) - 1): - corrected_i = time_values_raw[i] + cumulative_offset_days - time_values_corrected[i] = corrected_i - raw_diff = time_values_raw[i+1] - time_values_raw[i] - if raw_diff < expected_step_days * 0.5: - jump_count += 1 - expected_next = corrected_i + expected_step_days - cumulative_offset_days = expected_next - time_values_raw[i+1] - time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days - - print(f"Time correction completed. Corrected {jump_count} jumps.") - - print("Decoding corrected time values...") - try: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) - except ValueError: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar) - - # --- Create dataset with corrected time --- - ds_corrected = ds_temporal.copy(deep=False) - ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) - ds_corrected['time'].encoding['units'] = units - ds_corrected['time'].encoding['calendar'] = calendar - - # === Merge static variables === - for var_name, data_array in static_data.items(): - ds_corrected[var_name] = data_array - - # --- Write to file --- - print(f"Writing to file: {final_output_file}") - output_encoding = { - 'FSDS': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, - 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} - } - if 'gridID' in ds_corrected: - output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} - if 'LONGXY' in ds_corrected: - output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} - if 'LATIXY' in ds_corrected: - output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} - - ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) - print(f"✓ Successfully generated: {final_output_file}") - print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") - -except Exception as e: - print(f"Error: {e}") - traceback.print_exc() - sys.exit(1) - -print("\nScript execution completed.") - diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_PRECTmms_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_PRECTmms_20years.py deleted file mode 100644 index 7d7998f..0000000 --- a/scripts/training_data_generation/python_scripts/construct_TVA_PRECTmms_20years.py +++ /dev/null @@ -1,152 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate TVA PRECTmms forcing data (1980-1999, 20 years, 3-hour resolution) -Corresponding to crujra.v2.5.5d_PRECTmms_1901-2023_z01.nc - -Author: Zhuowei Gu -Date: 2024 -""" - -import xarray as xr -import os -import cftime -import numpy as np -import sys -import traceback -import datetime -from pathlib import Path - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import config - -# --- Configuration --- -data_dir = config.forcing_raw_data_path -out_dir = config.forcing_netcdf_output_dir -os.makedirs(out_dir, exist_ok=True) - -start_year = 1980 -end_year = 1999 -final_output_file = os.path.join(out_dir, f"TVA_PRECTmms_{start_year}-{end_year}.nc") - -print("="*80) -print("Generate TVA PRECTmms forcing data") -print("1. Merge monthly files (1980-1999, 240 months)") -print("2. Correct time axis discontinuities") -print("3. Maintain 3-hour resolution (no downsampling)") -print("="*80) - -print(f"Source data directory: {data_dir}") -print(f"Target output file: {final_output_file}") - -# --- Find and build all monthly file list --- -print("Searching for monthly files...") -all_monthly_files = [] -for year in range(start_year, end_year + 1): - for month in range(1, 13): - file_name = f"clmforc.Daymet.km.1d.Prec.{year}-{month:02d}.nc" - file_path = os.path.join(data_dir, file_name) - if os.path.exists(file_path): - all_monthly_files.append(file_path) - else: - print(f" Warning: File {file_name} does not exist, skipping.") - -if not all_monthly_files: - print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") - sys.exit(1) - -print(f"Found {len(all_monthly_files)} valid monthly files.") - -# --- Define Dask chunks --- -dask_chunks = {'time': 366*8} - -try: - with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', - decode_times=False, chunks=dask_chunks, parallel=False) as ds: - - # === Separate static variables === - print("Separating static coordinate/ID variables...") - static_data = {} - for var_name in ['gridID', 'LONGXY', 'LATIXY']: - if var_name in ds: - if 'time' in ds[var_name].dims: - static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() - else: - static_data[var_name] = ds[var_name].load() - - # === Process PRECTmms variable === - if 'PRECTmms' not in ds.data_vars: - print("Error: PRECTmms variable not found.") - sys.exit(1) - - print("Processing PRECTmms variable...") - ds_temporal = ds[['time', 'PRECTmms']] - - # --- Time axis correction --- - print("Loading raw time coordinates...") - time_values_raw = ds_temporal['time'].load().values - units = ds_temporal['time'].attrs['units'] - calendar = ds_temporal['time'].attrs.get('calendar', 'standard') - if calendar.lower() == 'no_leap': - calendar = 'noleap' - - print(f" Time points: {len(time_values_raw)}") - print("Starting time coordinate correction...") - - time_values_corrected = np.copy(time_values_raw).astype(float) - cumulative_offset_days = 0.0 - expected_step_days = 3.0 / 24.0 - jump_count = 0 - - for i in range(len(time_values_raw) - 1): - corrected_i = time_values_raw[i] + cumulative_offset_days - time_values_corrected[i] = corrected_i - raw_diff = time_values_raw[i+1] - time_values_raw[i] - if raw_diff < expected_step_days * 0.5: - jump_count += 1 - expected_next = corrected_i + expected_step_days - cumulative_offset_days = expected_next - time_values_raw[i+1] - time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days - - print(f"Time correction completed. Corrected {jump_count} jumps.") - - print("Decoding corrected time values...") - try: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) - except ValueError: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar) - - # --- Create dataset with corrected time --- - ds_corrected = ds_temporal.copy(deep=False) - ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) - ds_corrected['time'].encoding['units'] = units - ds_corrected['time'].encoding['calendar'] = calendar - - # === Merge static variables === - for var_name, data_array in static_data.items(): - ds_corrected[var_name] = data_array - - # --- Write to file --- - print(f"Writing to file: {final_output_file}") - output_encoding = { - 'PRECTmms': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, - 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} - } - if 'gridID' in ds_corrected: - output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} - if 'LONGXY' in ds_corrected: - output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} - if 'LATIXY' in ds_corrected: - output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} - - ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) - print(f"✓ Successfully generated: {final_output_file}") - print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") - -except Exception as e: - print(f"Error: {e}") - traceback.print_exc() - sys.exit(1) - -print("\nScript execution completed.") - diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_PSRF_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_PSRF_20years.py deleted file mode 100644 index 6d13015..0000000 --- a/scripts/training_data_generation/python_scripts/construct_TVA_PSRF_20years.py +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate TVA PSRF forcing data (1980-1999, 20 years, 3-hour resolution) -Corresponding to crujra.v2.5.5d_PSRF_1901-2023_z01.nc - -Author: Zhuowei Gu -Date: 2024 -""" - -import xarray as xr -import os -import cftime -import numpy as np -import sys -import traceback -import datetime -from pathlib import Path - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import config - -# --- 配置 --- -# --- Configuration --- -data_dir = config.forcing_raw_data_path -out_dir = config.forcing_netcdf_output_dir -os.makedirs(out_dir, exist_ok=True) - -start_year = 1980 -end_year = 1999 -final_output_file = os.path.join(out_dir, f"TVA_PSRF_{start_year}-{end_year}.nc") - -print("="*80) -print("Generate TVA PSRF forcing data") -print("1. Merge monthly files (1980-1999, 240 months)") -print("2. Correct time axis discontinuities") -print("3. Maintain 3-hour resolution (no downsampling)") -print("="*80) - -print(f"Source data directory: {data_dir}") -print(f"Target output file: {final_output_file}") - -all_monthly_files = [] -for year in range(start_year, end_year + 1): - for month in range(1, 13): - file_name = f"clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc" - file_path = os.path.join(data_dir, file_name) - if os.path.exists(file_path): - all_monthly_files.append(file_path) - -if not all_monthly_files: - print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") - sys.exit(1) - -print(f"Found {len(all_monthly_files)} valid monthly files.") - -dask_chunks = {'time': 366*8} - -try: - with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', - decode_times=False, chunks=dask_chunks, parallel=False) as ds: - - static_data = {} - for var_name in ['gridID', 'LONGXY', 'LATIXY']: - if var_name in ds: - if 'time' in ds[var_name].dims: - static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() - else: - static_data[var_name] = ds[var_name].load() - - ds_temporal = ds[['time', 'PSRF']] - time_values_raw = ds_temporal['time'].load().values - units = ds_temporal['time'].attrs['units'] - calendar = ds_temporal['time'].attrs.get('calendar', 'standard') - if calendar.lower() == 'no_leap': calendar = 'noleap' - - time_values_corrected = np.copy(time_values_raw).astype(float) - cumulative_offset_days = 0.0 - expected_step_days = 3.0 / 24.0 - - for i in range(len(time_values_raw) - 1): - corrected_i = time_values_raw[i] + cumulative_offset_days - time_values_corrected[i] = corrected_i - raw_diff = time_values_raw[i+1] - time_values_raw[i] - if raw_diff < expected_step_days * 0.5: - expected_next = corrected_i + expected_step_days - cumulative_offset_days = expected_next - time_values_raw[i+1] - time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days - - try: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) - except ValueError: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar) - - ds_corrected = ds_temporal.copy(deep=False) - ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) - ds_corrected['time'].encoding['units'] = units - ds_corrected['time'].encoding['calendar'] = calendar - - for var_name, data_array in static_data.items(): - ds_corrected[var_name] = data_array - - output_encoding = { - 'PSRF': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, - 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} - } - if 'gridID' in ds_corrected: - output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} - if 'LONGXY' in ds_corrected: - output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} - if 'LATIXY' in ds_corrected: - output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} - - ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) - print(f"✓ Successfully generated: {final_output_file}") - print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") - -except Exception as e: - print(f"Error: {e}") - traceback.print_exc() - sys.exit(1) - -print("\nScript execution completed.") - diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_QBOT_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_QBOT_20years.py deleted file mode 100644 index 17cda3b..0000000 --- a/scripts/training_data_generation/python_scripts/construct_TVA_QBOT_20years.py +++ /dev/null @@ -1,152 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate TVA QBOT forcing data (1980-1999, 20 years, 3-hour resolution) -Corresponding to crujra.v2.5.5d_QBOT_1901-2023_z01.nc - -Author: Zhuowei Gu -Date: 2024 -""" - -import xarray as xr -import os -import cftime -import numpy as np -import sys -import traceback -import datetime -from pathlib import Path - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import config - -# --- Configuration --- -data_dir = config.forcing_raw_data_path -out_dir = config.forcing_netcdf_output_dir -os.makedirs(out_dir, exist_ok=True) - -start_year = 1980 -end_year = 1999 -final_output_file = os.path.join(out_dir, f"TVA_QBOT_{start_year}-{end_year}.nc") - -print("="*80) -print("Generate TVA QBOT forcing data") -print("1. Merge monthly files (1980-1999, 240 months)") -print("2. Correct time axis discontinuities") -print("3. Maintain 3-hour resolution (no downsampling)") -print("="*80) - -print(f"Source data directory: {data_dir}") -print(f"Target output file: {final_output_file}") - -# --- Find and build all monthly file list --- -print("Searching for monthly files...") -all_monthly_files = [] -for year in range(start_year, end_year + 1): - for month in range(1, 13): - file_name = f"clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc" - file_path = os.path.join(data_dir, file_name) - if os.path.exists(file_path): - all_monthly_files.append(file_path) - else: - print(f" Warning: File {file_name} does not exist, skipping.") - -if not all_monthly_files: - print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") - sys.exit(1) - -print(f"Found {len(all_monthly_files)} valid monthly files.") - -# --- Define Dask chunks --- -dask_chunks = {'time': 366*8} - -try: - with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', - decode_times=False, chunks=dask_chunks, parallel=False) as ds: - - # === Separate static variables === - print("Separating static coordinate/ID variables...") - static_data = {} - for var_name in ['gridID', 'LONGXY', 'LATIXY']: - if var_name in ds: - if 'time' in ds[var_name].dims: - static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() - else: - static_data[var_name] = ds[var_name].load() - - # === Process QBOT variable === - if 'QBOT' not in ds.data_vars: - print("Error: QBOT variable not found.") - sys.exit(1) - - print("Processing QBOT variable...") - ds_temporal = ds[['time', 'QBOT']] - - # --- Time axis correction --- - print("Loading raw time coordinates...") - time_values_raw = ds_temporal['time'].load().values - units = ds_temporal['time'].attrs['units'] - calendar = ds_temporal['time'].attrs.get('calendar', 'standard') - if calendar.lower() == 'no_leap': - calendar = 'noleap' - - print(f" Time points: {len(time_values_raw)}") - print("Starting time coordinate correction...") - - time_values_corrected = np.copy(time_values_raw).astype(float) - cumulative_offset_days = 0.0 - expected_step_days = 3.0 / 24.0 - jump_count = 0 - - for i in range(len(time_values_raw) - 1): - corrected_i = time_values_raw[i] + cumulative_offset_days - time_values_corrected[i] = corrected_i - raw_diff = time_values_raw[i+1] - time_values_raw[i] - if raw_diff < expected_step_days * 0.5: - jump_count += 1 - expected_next = corrected_i + expected_step_days - cumulative_offset_days = expected_next - time_values_raw[i+1] - time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days - - print(f"Time correction completed. Corrected {jump_count} jumps.") - - print("Decoding corrected time values...") - try: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) - except ValueError: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar) - - # --- Create dataset with corrected time --- - ds_corrected = ds_temporal.copy(deep=False) - ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) - ds_corrected['time'].encoding['units'] = units - ds_corrected['time'].encoding['calendar'] = calendar - - # === Merge static variables === - for var_name, data_array in static_data.items(): - ds_corrected[var_name] = data_array - - # --- Write to file --- - print(f"Writing to file: {final_output_file}") - output_encoding = { - 'QBOT': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, - 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} - } - if 'gridID' in ds_corrected: - output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} - if 'LONGXY' in ds_corrected: - output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} - if 'LATIXY' in ds_corrected: - output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} - - ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) - print(f"✓ Successfully generated: {final_output_file}") - print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") - -except Exception as e: - print(f"Error: {e}") - traceback.print_exc() - sys.exit(1) - -print("\nScript execution completed.") - diff --git a/scripts/training_data_generation/python_scripts/construct_TVA_TBOT_20years.py b/scripts/training_data_generation/python_scripts/construct_TVA_TBOT_20years.py deleted file mode 100644 index c9d5bc5..0000000 --- a/scripts/training_data_generation/python_scripts/construct_TVA_TBOT_20years.py +++ /dev/null @@ -1,152 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate TVA TBOT forcing data (1980-1999, 20 years, 3-hour resolution) -Corresponding to crujra.v2.5.5d_TBOT_1901-2023_z01.nc - -Author: Zhuowei Gu -Date: 2024 -""" - -import xarray as xr -import os -import cftime -import numpy as np -import sys -import traceback -import datetime -from pathlib import Path - -# Import configuration -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import config - -# --- Configuration --- -data_dir = config.forcing_raw_data_path -out_dir = config.forcing_netcdf_output_dir -os.makedirs(out_dir, exist_ok=True) - -start_year = 1980 -end_year = 1999 -final_output_file = os.path.join(out_dir, f"TVA_TBOT_{start_year}-{end_year}.nc") - -print("="*80) -print("Generate TVA TBOT forcing data") -print("1. Merge monthly files (1980-1999, 240 months)") -print("2. Correct time axis discontinuities") -print("3. Maintain 3-hour resolution (no downsampling)") -print("="*80) - -print(f"Source data directory: {data_dir}") -print(f"Target output file: {final_output_file}") - -# --- Find and build all monthly file list --- -print("Searching for monthly files...") -all_monthly_files = [] -for year in range(start_year, end_year + 1): - for month in range(1, 13): - file_name = f"clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc" - file_path = os.path.join(data_dir, file_name) - if os.path.exists(file_path): - all_monthly_files.append(file_path) - else: - print(f" Warning: File {file_name} does not exist, skipping.") - -if not all_monthly_files: - print(f"Error: No valid monthly files found in directory {data_dir} for years {start_year}-{end_year}.") - sys.exit(1) - -print(f"Found {len(all_monthly_files)} valid monthly files.") - -# --- Define Dask chunks --- -dask_chunks = {'time': 366*8} - -try: - with xr.open_mfdataset(all_monthly_files, combine='nested', concat_dim='time', - decode_times=False, chunks=dask_chunks, parallel=False) as ds: - - # === Separate static variables === - print("Separating static coordinate/ID variables...") - static_data = {} - for var_name in ['gridID', 'LONGXY', 'LATIXY']: - if var_name in ds: - if 'time' in ds[var_name].dims: - static_data[var_name] = ds[var_name].isel(time=0, drop=True).load() - else: - static_data[var_name] = ds[var_name].load() - - # === Process TBOT variable === - if 'TBOT' not in ds.data_vars: - print("Error: TBOT variable not found.") - sys.exit(1) - - print("Processing TBOT variable...") - ds_temporal = ds[['time', 'TBOT']] - - # --- Time axis correction --- - print("Loading raw time coordinates...") - time_values_raw = ds_temporal['time'].load().values - units = ds_temporal['time'].attrs['units'] - calendar = ds_temporal['time'].attrs.get('calendar', 'standard') - if calendar.lower() == 'no_leap': - calendar = 'noleap' - - print(f" Time points: {len(time_values_raw)}") - print("Starting time coordinate correction...") - - time_values_corrected = np.copy(time_values_raw).astype(float) - cumulative_offset_days = 0.0 - expected_step_days = 3.0 / 24.0 - jump_count = 0 - - for i in range(len(time_values_raw) - 1): - corrected_i = time_values_raw[i] + cumulative_offset_days - time_values_corrected[i] = corrected_i - raw_diff = time_values_raw[i+1] - time_values_raw[i] - if raw_diff < expected_step_days * 0.5: - jump_count += 1 - expected_next = corrected_i + expected_step_days - cumulative_offset_days = expected_next - time_values_raw[i+1] - time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days - - print(f"Time correction completed. Corrected {jump_count} jumps.") - - print("Decoding corrected time values...") - try: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) - except ValueError: - dates = cftime.num2date(time_values_corrected, units, calendar=calendar) - - # --- Create dataset with corrected time --- - ds_corrected = ds_temporal.copy(deep=False) - ds_corrected['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) - ds_corrected['time'].encoding['units'] = units - ds_corrected['time'].encoding['calendar'] = calendar - - # === Merge static variables === - for var_name, data_array in static_data.items(): - ds_corrected[var_name] = data_array - - # --- Write to file --- - print(f"Writing to file: {final_output_file}") - output_encoding = { - 'TBOT': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, - 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} - } - if 'gridID' in ds_corrected: - output_encoding['gridID'] = {'dtype': ds_corrected['gridID'].dtype} - if 'LONGXY' in ds_corrected: - output_encoding['LONGXY'] = {'dtype': ds_corrected['LONGXY'].dtype, '_FillValue': np.nan} - if 'LATIXY' in ds_corrected: - output_encoding['LATIXY'] = {'dtype': ds_corrected['LATIXY'].dtype, '_FillValue': np.nan} - - ds_corrected.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) - print(f"✓ Successfully generated: {final_output_file}") - print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") - -except Exception as e: - print(f"Error: {e}") - traceback.print_exc() - sys.exit(1) - -print("\nScript execution completed.") - diff --git a/scripts/training_data_generation/python_scripts/construct_forcing_20years.py b/scripts/training_data_generation/python_scripts/construct_forcing_20years.py new file mode 100644 index 0000000..bfe5b87 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/construct_forcing_20years.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +""" +Generate all forcing data (1980-1999, 20 years, 3-hour resolution) +Integrates all 6 forcing variables: FLDS, FSDS, PRECTmms, PSRF, QBOT, TBOT +Based on optimized individual scripts + +Author: Zhuowei Gu +Date: 2024 +""" + +import xarray as xr +import os +import cftime +import numpy as np +import sys +import traceback +import datetime +import argparse +from pathlib import Path + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config + +# --- Configuration --- +def parse_arguments(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='Generate forcing NetCDF files for 1980-1999') + parser.add_argument('--input-dir', type=str, + default=config.forcing_raw_data_path, + help='Input directory containing raw forcing data') + parser.add_argument('--output-dir', type=str, + default=config.forcing_netcdf_output_dir, + help='Output directory for generated NetCDF files') + parser.add_argument('--start-year', type=int, default=1980, + help='Start year (default: 1980)') + parser.add_argument('--end-year', type=int, default=1999, + help='End year (default: 1999)') + return parser.parse_args() + +# Parse command line arguments +args = parse_arguments() +data_dir = args.input_dir +out_dir = args.output_dir +os.makedirs(out_dir, exist_ok=True) + +start_year = args.start_year +end_year = args.end_year + +# Define all forcing variables and their file patterns +forcing_variables = { + 'FLDS': { + 'file_pattern': 'clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc', + 'description': 'Downward longwave radiation' + }, + 'FSDS': { + 'file_pattern': 'clmforc.Daymet.km.1d.Solr.{year}-{month:02d}.nc', + 'description': 'Downward shortwave radiation' + }, + 'PRECTmms': { + 'file_pattern': 'clmforc.Daymet.km.1d.Prec.{year}-{month:02d}.nc', + 'description': 'Precipitation rate' + }, + 'PSRF': { + 'file_pattern': 'clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc', + 'description': 'Surface pressure' + }, + 'QBOT': { + 'file_pattern': 'clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc', + 'description': 'Specific humidity' + }, + 'TBOT': { + 'file_pattern': 'clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc', + 'description': 'Air temperature' + } +} + +print("="*80) +print("Generate All Forcing Data (1980-1999, 20 years)") +print("Processing 6 forcing variables:") +for var, info in forcing_variables.items(): + print(f" - {var}: {info['description']}") +print("="*80) + +print(f"Source data directory: {data_dir}") +print(f"Output directory: {out_dir}") + +def process_forcing_variable(var_name, var_info): + """Process a single forcing variable using optimized method""" + print(f"\n{'='*60}") + print(f"Processing {var_name}: {var_info['description']}") + print(f"{'='*60}") + + final_output_file = os.path.join(out_dir, f"{var_name}_{start_year}-{end_year}.nc") + + # Check if output file already exists + if os.path.exists(final_output_file): + print(f"✅ Output file already exists: {final_output_file}") + return True + + print(f"Target output file: {final_output_file}") + + # --- Find and build all monthly file list --- + print("Searching for monthly files...") + all_monthly_files = [] + + for year in range(start_year, end_year + 1): + for month in range(1, 13): + file_name = var_info['file_pattern'].format(year=year, month=month) + file_path = os.path.join(data_dir, file_name) + if os.path.exists(file_path): + all_monthly_files.append(file_path) + else: + print(f" Warning: File {file_name} does not exist, skipping.") + + if not all_monthly_files: + print(f"❌ Error: No valid monthly files found for {var_name} in directory {data_dir}") + return False + + print(f"Found {len(all_monthly_files)} valid monthly files.") + + # --- Define Dask chunks --- + dask_chunks = {'time': 366*8} + + try: + with xr.open_mfdataset( + all_monthly_files, + combine='nested', + concat_dim='time', + decode_times=False, + chunks=dask_chunks, + parallel=False, + ) as ds: + + # === Separate static variables === + print("Separating static coordinate/ID variables...") + static_var_names = ['gridID', 'LONGXY', 'LATIXY'] + static_data = {} + + for var_name_static in static_var_names: + if var_name_static in ds: + if 'time' in ds[var_name_static].dims: + static_data[var_name_static] = ds[var_name_static].isel(time=0, drop=True).load() + else: + static_data[var_name_static] = ds[var_name_static].load() + + # === Process target variable === + time_varying_vars = [var_name] + if var_name not in ds.data_vars: + print(f"Error: {var_name} variable not found.") + return False + + print(f"Processing variables: {time_varying_vars}") + ds_temporal = ds[['time'] + time_varying_vars] + + # --- Time axis correction --- + print("Loading raw time coordinates...") + time_values_raw = ds_temporal['time'].load().values + units = ds_temporal['time'].attrs['units'] + calendar = ds_temporal['time'].attrs.get('calendar', 'standard') + if calendar.lower() == 'no_leap': + calendar = 'noleap' + + print(f" Time points: {len(time_values_raw)}") + + print("Starting time coordinate correction...") + time_values_corrected = np.copy(time_values_raw).astype(float) + cumulative_offset_days = 0.0 + expected_step_days = 3.0 / 24.0 + jump_count = 0 + + for i in range(len(time_values_raw) - 1): + corrected_i = time_values_raw[i] + cumulative_offset_days + time_values_corrected[i] = corrected_i + raw_diff = time_values_raw[i+1] - time_values_raw[i] + + if raw_diff < expected_step_days * 0.5: + jump_count += 1 + expected_next_corrected_value = corrected_i + expected_step_days + offset_needed = expected_next_corrected_value - time_values_raw[i+1] + cumulative_offset_days = offset_needed + + time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days + + print(f"Time correction completed. Corrected {jump_count} jumps.") + + print("Decoding corrected time values...") + try: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) + except ValueError: + dates = cftime.num2date(time_values_corrected, units, calendar=calendar) + + # --- Check time monotonicity --- + print("Checking time monotonicity...") + if len(dates) >= 2: + diffs_corrected = np.diff(dates) + zero_timedelta = datetime.timedelta(0) + problem_indices = np.where(diffs_corrected <= zero_timedelta)[0] + + if len(problem_indices) > 0: + print(f"Error! Corrected time is still not monotonic!") + return False + else: + print("✓ Time coordinate check passed.") + + # --- Create dataset with corrected time --- + ds_corrected_time = ds_temporal.copy(deep=False) + ds_corrected_time['time'] = xr.DataArray(dates, dims='time', coords={'time': dates}) + ds_corrected_time['time'].encoding['units'] = units + ds_corrected_time['time'].encoding['calendar'] = calendar + + # === Merge static variables === + for var_name_static, data_array in static_data.items(): + ds_corrected_time[var_name_static] = data_array + + final_dataset = ds_corrected_time + print("Final dataset:") + print(final_dataset) + + # --- Write to file --- + print(f"Writing to file: {final_output_file}") + output_encoding = { + var_name: {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, + 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} + } + + if 'gridID' in final_dataset: + output_encoding['gridID'] = {'dtype': final_dataset['gridID'].dtype} + if 'LONGXY' in final_dataset: + output_encoding['LONGXY'] = {'dtype': final_dataset['LONGXY'].dtype, '_FillValue': np.nan} + if 'LATIXY' in final_dataset: + output_encoding['LATIXY'] = {'dtype': final_dataset['LATIXY'].dtype, '_FillValue': np.nan} + + final_dataset.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) + print(f"✓ Successfully generated: {final_output_file}") + print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") + + except Exception as e: + print(f"\nError: {e}") + traceback.print_exc() + return False + + return True + +def main(): + """Main function to process all forcing variables""" + print("Starting processing of all forcing variables...") + + success_count = 0 + total_count = len(forcing_variables) + + for var_name, var_info in forcing_variables.items(): + try: + success = process_forcing_variable(var_name, var_info) + if success: + success_count += 1 + print(f"✅ {var_name} processing completed successfully") + else: + print(f"❌ {var_name} processing failed") + except Exception as e: + print(f"❌ {var_name} processing failed with exception: {e}") + + # Final summary + print(f"\n{'='*80}") + print("PROCESSING SUMMARY") + print(f"{'='*80}") + print(f"Total variables: {total_count}") + print(f"Successfully processed: {success_count}") + print(f"Failed: {total_count - success_count}") + + if success_count == total_count: + print("🎉 All forcing variables processed successfully!") + print(f"Output files saved in: {out_dir}") + + # List generated files + print("\nGenerated files:") + for var_name in forcing_variables.keys(): + output_file = os.path.join(out_dir, f"{var_name}_{start_year}-{end_year}.nc") + if os.path.exists(output_file): + file_size = os.path.getsize(output_file) / (1024**3) + print(f" ✅ {var_name}_{start_year}-{end_year}.nc ({file_size:.2f} GB)") + else: + print(f" ❌ {var_name}_{start_year}-{end_year}.nc (not found)") + + return True + else: + print(f"⚠️ Some variables failed to process. Check the output above for details.") + return False + +if __name__ == "__main__": + success = main() + sys.exit(0 if success else 1) \ No newline at end of file diff --git a/scripts/training_data_generation/python_scripts/enhanced_training_dataset.py b/scripts/training_data_generation/python_scripts/enhanced_training_dataset.py new file mode 100644 index 0000000..32aec44 --- /dev/null +++ b/scripts/training_data_generation/python_scripts/enhanced_training_dataset.py @@ -0,0 +1,1756 @@ +#!/usr/bin/env python3 +""" +Enhanced Training Dataset Generation Script +Combines 72_dataset_construction.py and 37_dataset.py functionality +Dynamically extracts variables from CNP_IO file instead of using hardcoded lists +""" + +import netCDF4 as nc +import numpy as np +import pandas as pd +from scipy.spatial import cKDTree +import os +import sys +import time +import glob +import argparse +from typing import Dict, List, Tuple + +# Import configuration +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config +from python_scripts.cnp_io_parse import parse_cnp_io_list + +print("="*80) +print("Enhanced Training Dataset Generation") +print("Combining 72_dataset_construction.py + 37_dataset.py") +print("Dynamic variable extraction from CNP_IO file") +print("="*80) + +# Constants +VARS_TO_CLEAN = {'H2OSOI_LIQ', 'H2OSOI_ICE'} +FILL_VALUE_THRESHOLD = 1e35 + +def parse_cnp_io_variables(): + """Parse CNP_IO file to get all variable definitions dynamically""" + cnp_io_file = config.Config.CNP_IO_FILE + if not os.path.exists(cnp_io_file): + raise FileNotFoundError(f"CNP_IO file not found: {cnp_io_file}") + + parsed = parse_cnp_io_list(cnp_io_file) + print(f"✅ CNP_IO file parsed: {cnp_io_file}") + + # Extract variable categories + time_series_vars = parsed.get('time_series_variables', []) + surface_vars = parsed.get('surface_properties', []) + scalar_vars = parsed.get('scalar_variables', []) + pft_1d_vars = parsed.get('pft_1d_variables', []) + variables_2d_soil = parsed.get('variables_2d_soil', []) + water_vars = parsed.get('water_variables', []) + pool_vars = parsed.get('pool_variables', ['cpool', 'npool', 'ppool', 'xsmrpool']) + pft_parameters = parsed.get('pft_parameters', []) + + print(f"Variable categories from CNP_IO:") + print(f" Time series variables: {len(time_series_vars)}") + print(f" Surface properties: {len(surface_vars)}") + print(f" Scalar variables: {len(scalar_vars)}") + print(f" PFT 1D variables: {len(pft_1d_vars)}") + print(f" 2D soil variables: {len(variables_2d_soil)}") + print(f" Water variables: {len(water_vars)}") + print(f" Pool variables: {len(pool_vars)}") + print(f" PFT parameters: {len(pft_parameters)}") + + return { + 'time_series_vars': time_series_vars, + 'surface_vars': surface_vars, + 'scalar_vars': scalar_vars, + 'pft_1d_vars': pft_1d_vars, + 'variables_2d_soil': variables_2d_soil, + 'water_vars': water_vars, + 'pool_vars': pool_vars, + 'pft_parameters': pft_parameters + } + +def build_restart_kdtree(ds_restart: nc.Dataset) -> Tuple[cKDTree, np.ndarray]: + """Build KDTree for restart file coordinates""" + gridcell_lat = ds_restart.variables["grid1d_lat"][:] + gridcell_lon = ds_restart.variables["grid1d_lon"][:] + coords = np.vstack((gridcell_lat, gridcell_lon)).T + tree = cKDTree(coords) + return tree, coords + +def build_column_index_map(ds_restart: nc.Dataset) -> Dict[int, np.ndarray]: + """Build mapping from gridcell ID to column indices""" + cols1d_gridcell_index = ds_restart.variables["cols1d_gridcell_index"][:] + unique_ids = np.unique(cols1d_gridcell_index) + mapping: Dict[int, np.ndarray] = {} + for grid_id in unique_ids: + mapping[int(grid_id)] = np.where(cols1d_gridcell_index == grid_id)[0] + return mapping + +def build_pft_index_map(ds_restart: nc.Dataset) -> Dict[int, np.ndarray]: + """Build mapping from gridcell ID to PFT indices""" + pfts1d_gridcell_index = ds_restart.variables["pfts1d_gridcell_index"][:] + unique_ids = np.unique(pfts1d_gridcell_index) + mapping: Dict[int, np.ndarray] = {} + for grid_id in unique_ids: + mapping[int(grid_id)] = np.where(pfts1d_gridcell_index == grid_id)[0] + return mapping + +def ensure_vars_exist(ds: nc.Dataset, var_names: List[str]) -> List[str]: + """Check which variables exist in the dataset""" + existing = [] + for name in var_names: + if name in ds.variables: + existing.append(name) + return existing + +def extract_col1d_x(ds_restart: nc.Dataset, var_name: str, col_indices: np.ndarray) -> List[float]: + """Extract 1D column variables from restart file""" + if col_indices.size == 0: + return [] + values = ds_restart.variables[var_name][col_indices] + return values.astype(float).tolist() + +def extract_col1d_y(ds_r_list: List[nc.Dataset], var_name: str, col_indices: np.ndarray) -> List[float]: + """Extract 1D column variables from Y files (future restart files)""" + if col_indices.size == 0: + return [] + slices: List[np.ndarray] = [] + for ds_r in ds_r_list: + values = ds_r.variables[var_name][col_indices] + slices.append(np.asarray(values, dtype=float)) + stacked = np.stack(slices, axis=0) + avg = np.mean(stacked, axis=0) + return avg.tolist() + +def extract_col2d_x(ds_restart: nc.Dataset, var_name: str, col_indices: np.ndarray) -> List[List[float]]: + """Extract 2D column variables from restart file""" + if col_indices.size == 0: + return [] + values = ds_restart.variables[var_name][col_indices, :] + values_np = np.asarray(values, dtype=float) + if var_name in VARS_TO_CLEAN: + values_np[values_np >= FILL_VALUE_THRESHOLD] = 0.0 + return values_np.tolist() + +def extract_col2d_y(ds_r_list: List[nc.Dataset], var_name: str, col_indices: np.ndarray) -> List[List[float]]: + """Extract 2D column variables from Y files""" + if col_indices.size == 0: + return [] + slices: List[np.ndarray] = [] + for ds_r in ds_r_list: + values = ds_r.variables[var_name][col_indices, :] + values_np = np.asarray(values, dtype=float) + if var_name in VARS_TO_CLEAN: + values_np[values_np >= FILL_VALUE_THRESHOLD] = 0.0 + slices.append(values_np) + stacked = np.stack(slices, axis=0) + avg = np.mean(stacked, axis=0) + return avg.tolist() + +def extract_pft1d_x(ds_restart: nc.Dataset, var_name: str, pft_indices: np.ndarray) -> List[float]: + """Extract 1D PFT variables from restart file""" + if pft_indices.size == 0: + return [] + values = ds_restart.variables[var_name][pft_indices] + return np.asarray(values, dtype=float).tolist() + +def extract_pft1d_y(ds_r_list: List[nc.Dataset], var_name: str, pft_indices: np.ndarray) -> List[float]: + """Extract 1D PFT variables from Y files""" + if pft_indices.size == 0: + return [] + slices: List[np.ndarray] = [] + for ds_r in ds_r_list: + values = ds_r.variables[var_name][pft_indices] + slices.append(np.asarray(values, dtype=float)) + stacked = np.stack(slices, axis=0) + avg = np.mean(stacked, axis=0) + return avg.tolist() + +def extract_pft2d_x(ds_restart: nc.Dataset, var_name: str, pft_indices: np.ndarray) -> List[List[float]]: + """Extract 2D PFT variables from restart file""" + if pft_indices.size == 0: + return [] + values = ds_restart.variables[var_name][pft_indices, :] + return np.asarray(values, dtype=float).tolist() + +def extract_pft2d_y(ds_r_list: List[nc.Dataset], var_name: str, pft_indices: np.ndarray) -> List[List[float]]: + """Extract 2D PFT variables from Y files""" + if pft_indices.size == 0: + return [] + slices: List[np.ndarray] = [] + for ds_r in ds_r_list: + values = ds_r.variables[var_name][pft_indices, :] + slices.append(np.asarray(values, dtype=float)) + stacked = np.stack(slices, axis=0) + avg = np.mean(stacked, axis=0) + return avg.tolist() + +def calculate_monthly_avg(time_series, time_series_length=58400): + """Calculate monthly averages from high-resolution time series""" + if not isinstance(time_series, (list, np.ndarray)): + return [] + + if len(time_series) != time_series_length: + return [] + + # TVA data parameters (20 years, 3-hour interval) + steps_per_day = 8 # 3-hour interval = 8 steps/day + days_per_year = 365 + years_in_data = 20 # 1980-1999 + months_per_year = 12 + days_per_month = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31] + + monthly_averages = [] + start_idx = 0 + + for year in range(years_in_data): + for month_idx, month_days in enumerate(days_per_month): + # Handle leap year February + if year % 4 == 0 and month_idx == 1: # Leap year February + month_days = 29 + + end_idx = start_idx + month_days * steps_per_day + monthly_avg = np.mean(time_series[start_idx:end_idx]) + # Ensure float64 precision to match reference file + monthly_averages.append(float(monthly_avg)) + start_idx = end_idx + + return monthly_averages + +def generate_base_dataset(variable_definitions): + """Generate base training dataset (72_dataset_construction.py logic)""" + print(f"\n{'='*80}") + print("STEP 1: Base Dataset Generation (72_dataset_construction.py)") + print(f"{'='*80}") + + # File paths from config + surface_data_files = config.surface_data_files + ad_spinup_history_files = config.ad_spinup_history_files + ad_spinup_restart_files = config.ad_spinup_restart_files + final_spinup_history_files = config.final_spinup_history_files + final_spinup_restart_files = config.final_spinup_restart_files + + # Forcing data files - dynamically find files containing variable names + forcing_files = {} + for var_name in variable_definitions['time_series_vars']: + # Look for files containing the variable name in forcing_netcdf directory + pattern = os.path.join(config.forcing_netcdf_output_dir, f'*{var_name}*1980-1999.nc') + matching_files = glob.glob(pattern) + if matching_files: + forcing_files[var_name] = matching_files[0] # Use first match + print(f"✅ Found forcing file: {os.path.basename(matching_files[0])}") + else: + print(f"⚠️ Forcing file not found for {var_name}: {pattern}") + + print(f"Found {len(forcing_files)} forcing files") + + # Output directory + output_dir = os.path.join(config.output_dir, 'training_dataset_pkl') + os.makedirs(output_dir, exist_ok=True) + + print("Loading NetCDF files...") + start_time = time.time() + + # Load all required files + ds1 = nc.Dataset(surface_data_files[0]) # Surface data + ds2 = nc.Dataset(ad_spinup_history_files[0]) # AD-SPINUP history + ds10 = nc.Dataset(ad_spinup_restart_files[0]) # AD-SPINUP restart + + # Load forcing data + ds_forcing = {} + for var_name, file_path in forcing_files.items(): + ds_forcing[var_name] = nc.Dataset(file_path) + print(f"✅ Forcing data loaded: {var_name}") + + # Load future files for Y variables + ds_h0_list = [nc.Dataset(fp) for fp in final_spinup_history_files] + ds_r_list = [nc.Dataset(fp) for fp in final_spinup_restart_files] + + print(f"All files loaded in {time.time() - start_time:.2f} seconds") + + # Get coordinates and build spatial filtering + lats = ds2.variables['lat'][:] + lons = ds2.variables['lon'][:] + landmask = ds2.variables['landfrac'][:] + + # Filter for land gridcells + valid_mask = (landmask > 0) + valid_gridcells = np.where(valid_mask)[0] + + print(f"Total land gridcells: {len(valid_gridcells)}") + print(f"Latitude range: [{lats.min():.2f}, {lats.max():.2f}]") + print(f"Longitude range: [{lons.min():.2f}, {lons.max():.2f}]") + + # Build KDTree indices + print("Building KDTree indices...") + + # Restart file coordinates + gridcell_lat = ds10.variables['grid1d_lat'][:] + gridcell_lon = ds10.variables['grid1d_lon'][:] + restart_grid_coords = np.vstack((gridcell_lat, gridcell_lon)).T + restart_tree = cKDTree(restart_grid_coords) + + # Forcing file coordinates + forcing_lats = list(ds_forcing.values())[0].variables['LATIXY'][0, :] + forcing_lons = list(ds_forcing.values())[0].variables['LONGXY'][0, :] + forcing_grid_coords = np.vstack((forcing_lats, forcing_lons)).T + forcing_tree = cKDTree(forcing_grid_coords) + + # Query coordinates for valid gridcells + query_coords = np.array([(lats[i], lons[i]) for i in valid_gridcells]) + _, all_restart_indices = restart_tree.query(query_coords, k=1) + _, all_forcing_indices = forcing_tree.query(query_coords, k=1) + + print("✅ KDTree indices built") + + # Pre-load forcing data into memory for optimization + print("🚀 Pre-loading forcing data into memory...") + forcing_data = {} + for var_name, ds in ds_forcing.items(): + forcing_data[var_name] = ds.variables[var_name][:, 0, :] # (time, 1, grid_cells) + print(f" Loaded {var_name}: {forcing_data[var_name].shape}") + + # Close forcing NetCDF files (data is now in memory) + for ds in ds_forcing.values(): + ds.close() + + # Build index mappings + pft_gridcell_index = ds10.variables['pfts1d_gridcell_index'][:] + column_gridcell_index = ds10.variables['cols1d_gridcell_index'][:] + + pft_map = {} + column_map = {} + unique_gridcell_ids = np.unique(pft_gridcell_index) + for grid_id in unique_gridcell_ids: + pft_map[grid_id] = np.where(pft_gridcell_index == grid_id)[0] + column_map[grid_id] = np.where(column_gridcell_index == grid_id)[0] + + print("✅ Index mappings built") + + # Process data in batches + batch_size = 1000 + batch_number = 1 + batch_files = [] # Track all generated files for post-processing + + print(f"\nProcessing {len(valid_gridcells)} gridcells in batches of {batch_size}...") + + for start_idx in range(0, len(valid_gridcells), batch_size): + end_idx = min(start_idx + batch_size, len(valid_gridcells)) + batch_gridcells = valid_gridcells[start_idx:end_idx] + batch_restart_indices = all_restart_indices[start_idx:end_idx] + batch_forcing_indices = all_forcing_indices[start_idx:end_idx] + + print(f"\nProcessing batch {batch_number}: gridcells {start_idx+1}-{end_idx}") + print(f" Batch size: {len(batch_gridcells)} gridcells") + print(f" batch_gridcells range: {batch_gridcells[0]} to {batch_gridcells[-1]}") + print(f" Unique gridcells in batch: {len(set(batch_gridcells))}") + if len(set(batch_gridcells)) != len(batch_gridcells): + print(f" ⚠️ WARNING: Duplicate gridcells detected in batch!") + batch_start_time = time.time() + + # Initialize data dictionary dynamically - completely from CNP_IO file + data_dict = {} + + # Add time series variables (forcing data) + for var_name in variable_definitions['time_series_vars']: + data_dict[var_name] = [] + + # Add surface properties + for var_name in variable_definitions['surface_vars']: + data_dict[var_name] = [] + + # Add scalar variables + for var_name in variable_definitions['scalar_vars']: + data_dict[var_name] = [] + data_dict[f'Y_{var_name}'] = [] + + # Add PFT variables + for var_name in variable_definitions['pft_1d_vars']: + data_dict[var_name] = [] + data_dict[f'Y_{var_name}'] = [] + + # Add 2D soil variables + for var_name in variable_definitions['variables_2d_soil']: + data_dict[var_name] = [] + data_dict[f'Y_{var_name}'] = [] + + # Add water variables + for var_name in variable_definitions['water_vars']: + data_dict[var_name] = [] + data_dict[f'Y_{var_name}'] = [] + + # Add pool variables + for var_name in variable_definitions['pool_vars']: + data_dict[var_name] = [] + data_dict[f'Y_{var_name}'] = [] + + # Process each gridcell in the batch + for k, gridcell_idx in enumerate(batch_gridcells): + if k % 100 == 0: + print(f" Processing gridcell {k}/{len(batch_gridcells)} (idx={gridcell_idx})") + + # Get indices + restart_idx = batch_restart_indices[k] + forcing_idx = batch_forcing_indices[k] + gridcell_id = restart_idx + 1 + + pft_indices_for_cell = pft_map.get(gridcell_id, []) + column_indices_for_cell = column_map.get(gridcell_id, []) + + # All variables are now processed dynamically from CNP_IO file + + # Debug: Check data_dict length after each gridcell + if k == 0: + print(f" After first gridcell: Latitude={len(data_dict['Latitude'])}, FLDS={len(data_dict.get('FLDS', []))}") + print(f" pft_indices_for_cell length: {len(pft_indices_for_cell)}") + print(f" column_indices_for_cell length: {len(column_indices_for_cell)}") + if k == 999: + print(f" After 1000th gridcell: Latitude={len(data_dict['Latitude'])}, FLDS={len(data_dict.get('FLDS', []))}") + + # Process forcing data (time series variables) - store raw data like original script + for var_name in variable_definitions['time_series_vars']: + if var_name in forcing_data: + time_series = forcing_data[var_name][:, forcing_idx] + # Store raw time series like original script (will be processed later) + data_dict[var_name].append(time_series) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + + # Process surface properties + for var_name in variable_definitions['surface_vars']: + if var_name == 'Latitude': + data_dict[var_name].append(lats[gridcell_idx]) + elif var_name == 'Longitude': + data_dict[var_name].append(lons[gridcell_idx]) + elif var_name == 'landfrac': + # landfrac comes from history file (ds2), not surface file (ds1) + # Convert to float64 to match reference file + landfrac_val = ds2.variables['landfrac'][gridcell_idx] + if hasattr(landfrac_val, 'data'): # MaskedArray + landfrac_val = landfrac_val.data + data_dict[var_name].append(float(landfrac_val)) + elif var_name == 'PCT_CLAY': + # Store PCT_CLAY as a list (all levels) + pct_clay_data = ds1.variables['PCT_CLAY'][:, gridcell_idx] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(pct_clay_data, 'data'): # MaskedArray + pct_clay_data = pct_clay_data.data + data_dict[var_name].append(pct_clay_data.astype(np.float64).tolist()) + elif var_name == 'PCT_SAND': + # Store PCT_SAND as a list (all levels) + pct_sand_data = ds1.variables['PCT_SAND'][:, gridcell_idx] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(pct_sand_data, 'data'): # MaskedArray + pct_sand_data = pct_sand_data.data + data_dict[var_name].append(pct_sand_data.astype(np.float64).tolist()) + elif var_name.startswith('PCT_NAT_PFT_') or var_name.startswith('PCT_CLAY_') or var_name.startswith('PCT_SAND_'): + # Handle 2D variables with level indices + if '_' in var_name: + level_idx = int(var_name.split('_')[-1]) + base_var = '_'.join(var_name.split('_')[:-1]) # e.g., 'PCT_CLAY' + if base_var in ds1.variables: + pct_val = ds1.variables[base_var][level_idx, gridcell_idx] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(pct_val, 'data'): # MaskedArray + pct_val = pct_val.data + data_dict[var_name].append(float(pct_val)) + else: + data_dict[var_name].append(0.0) + else: + data_dict[var_name].append(0.0) + elif var_name in ds1.variables: + # 1D variables + val = ds1.variables[var_name][gridcell_idx] + # Convert MaskedArray to regular array + if hasattr(val, 'data'): # MaskedArray + val = val.data + # Handle integer variables + if var_name in ['SOIL_COLOR', 'SOIL_ORDER']: + data_dict[var_name].append(int(val)) + else: + data_dict[var_name].append(float(val)) + else: + data_dict[var_name].append(0.0) # Add default value if variable not found + + # Process scalar variables from history file + for var_name in variable_definitions['scalar_vars']: + if var_name in ds2.variables: + val = ds2.variables[var_name][0, gridcell_idx] + # Convert MaskedArray to regular array + if hasattr(val, 'data'): # MaskedArray + val = val.data + # Handle integer variables + if var_name in ['SOIL_COLOR', 'SOIL_ORDER']: + data_dict[var_name].append(int(val)) + else: + data_dict[var_name].append(float(val)) + + # Add Y_ version for scalar variables (from final_spinup history files) + y_vals = [] + for ds_h0 in ds_h0_list: + if var_name in ds_h0.variables: + y_val = ds_h0.variables[var_name][0, gridcell_idx] + # Convert MaskedArray to regular array + if hasattr(y_val, 'data'): # MaskedArray + y_val = y_val.data + y_vals.append(y_val) + if y_vals: + avg_y_val = np.mean(y_vals) + data_dict[f'Y_{var_name}'].append(float(avg_y_val)) + else: + data_dict[f'Y_{var_name}'].append(0.0) + else: + data_dict[var_name].append(0.0) # Add default value if variable not found + data_dict[f'Y_{var_name}'].append(0.0) + + # Process PFT variables + for var_name in variable_definitions['pft_1d_vars']: + if var_name in ds10.variables: + # X values from restart file + x_val = ds10.variables[var_name][pft_indices_for_cell] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(x_val, 'data'): # MaskedArray + x_val = x_val.data + data_dict[var_name].append(x_val.astype(np.float64).tolist()) + + # Y values from future restart files + y_vals = [] + for ds_r in ds_r_list: + if var_name in ds_r.variables: + y_val = ds_r.variables[var_name][pft_indices_for_cell] + # Convert MaskedArray to regular array + if hasattr(y_val, 'data'): # MaskedArray + y_val = y_val.data + y_vals.append(y_val) + if y_vals: + avg_y_val = np.mean(y_vals, axis=0) + data_dict[f'Y_{var_name}'].append(avg_y_val.astype(np.float64).tolist()) + else: + data_dict[f'Y_{var_name}'].append([]) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + data_dict[f'Y_{var_name}'].append([]) + + # Process 2D soil variables + for var_name in variable_definitions['variables_2d_soil']: + if var_name in ds10.variables: + # X values + x_val = ds10.variables[var_name][column_indices_for_cell, :] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(x_val, 'data'): # MaskedArray + x_val = x_val.data + data_dict[var_name].append(x_val.astype(np.float64).tolist()) + + # Y values + y_vals = [] + for ds_r in ds_r_list: + if var_name in ds_r.variables: + y_val = ds_r.variables[var_name][column_indices_for_cell, :] + # Convert MaskedArray to regular array + if hasattr(y_val, 'data'): # MaskedArray + y_val = y_val.data + y_vals.append(y_val) + if y_vals: + avg_y_val = np.mean(y_vals, axis=0) + data_dict[f'Y_{var_name}'].append(avg_y_val.astype(np.float64).tolist()) + else: + data_dict[f'Y_{var_name}'].append([]) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + data_dict[f'Y_{var_name}'].append([]) + + # Process water variables + for var_name in variable_definitions['water_vars']: + if var_name in ds10.variables: + # X values + x_val = ds10.variables[var_name][column_indices_for_cell, :] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(x_val, 'data'): # MaskedArray + x_val = x_val.data + data_dict[var_name].append(x_val.astype(np.float64).tolist()) + + # Y values + y_vals = [] + for ds_r in ds_r_list: + if var_name in ds_r.variables: + y_val = ds_r.variables[var_name][column_indices_for_cell, :] + # Convert MaskedArray to regular array + if hasattr(y_val, 'data'): # MaskedArray + y_val = y_val.data + y_vals.append(y_val) + if y_vals: + avg_y_val = np.mean(y_vals, axis=0) + data_dict[f'Y_{var_name}'].append(avg_y_val.astype(np.float64).tolist()) + else: + data_dict[f'Y_{var_name}'].append([]) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + data_dict[f'Y_{var_name}'].append([]) + + # Process pool variables (only if not already processed as PFT variables) + for var_name in variable_definitions['pool_vars']: + # Skip if already processed as PFT variable + if var_name in variable_definitions['pft_1d_vars']: + continue + + if var_name in ds10.variables: + # X values + x_val = ds10.variables[var_name][column_indices_for_cell] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(x_val, 'data'): # MaskedArray + x_val = x_val.data + data_dict[var_name].append(x_val.astype(np.float64).tolist()) + + # Y values + y_vals = [] + for ds_r in ds_r_list: + if var_name in ds_r.variables: + y_val = ds_r.variables[var_name][column_indices_for_cell] + # Convert MaskedArray to regular array + if hasattr(y_val, 'data'): # MaskedArray + y_val = y_val.data + y_vals.append(y_val) + if y_vals: + avg_y_val = np.mean(y_vals, axis=0) + data_dict[f'Y_{var_name}'].append(avg_y_val.astype(np.float64).tolist()) + else: + data_dict[f'Y_{var_name}'].append([]) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + data_dict[f'Y_{var_name}'].append([]) + + # Create DataFrame and save + print(f" Creating DataFrame...") + + # Debug: Check array lengths + print(f" Debug: Checking array lengths...") + for key, values in data_dict.items(): + print(f" {key}: {len(values)} items") + + # Check for length mismatches + lengths = {key: len(values) for key, values in data_dict.items()} + unique_lengths = set(lengths.values()) + if len(unique_lengths) > 1: + print(f" ❌ Length mismatch detected!") + for length in unique_lengths: + vars_with_length = [k for k, v in lengths.items() if v == length] + print(f" Length {length}: {len(vars_with_length)} variables") + if length != max(unique_lengths): + print(f" Variables: {vars_with_length[:5]}{'...' if len(vars_with_length) > 5 else ''}") + + df_batch = pd.DataFrame(data_dict) + + print(f" Saving to disk...") + batch_save_path = f"{output_dir}/training_data_batch_{batch_number:02d}.pkl" + df_batch.to_pickle(batch_save_path) + batch_files.append(batch_save_path) # Add to list for post-processing + + batch_time = time.time() - batch_start_time + print(f"✅ Batch {batch_number} completed: {batch_time:.2f}s") + print(f" Path: {batch_save_path}") + print(f" Shape: {df_batch.shape}") + print(f" Columns: {len(df_batch.columns)}") + + batch_number += 1 + + # Cleanup NetCDF files + print(f"\nCleaning up NetCDF files...") + ds1.close() + ds2.close() + ds10.close() + for ds_h0 in ds_h0_list: + ds_h0.close() + for ds_r in ds_r_list: + ds_r.close() + + print("✅ All NetCDF files closed") + print(f"✅ Base dataset generation completed!") + print(f"Total batches: {batch_number - 1}") + print(f"Output directory: {output_dir}") + + # Post-processing like original script + print(f"\n{'='*80}") + print("POST-PROCESSING: Converting to monthly averages and expanding variables") + print(f"{'='*80}") + + # Process all generated files + for file_path in batch_files: + print(f"Processing {os.path.basename(file_path)}...") + + # Load the file + df = pd.read_pickle(file_path) + + # Process time series columns (convert to monthly averages) + print(f" Processing time series columns...") + for col in variable_definitions['time_series_vars']: + if col in df.columns: + print(f" Processing {col}...") + df[col] = df[col].apply(calculate_monthly_avg) + + # Process list columns (expand PCT variables) + print(f" Processing list columns...") + list_like_columns = ['PCT_CLAY', 'PCT_SAND', 'PCT_NAT_PFT'] + for col in list_like_columns: + if col in df.columns: + print(f" Expanding {col}...") + expanded_cols = df[col].apply(pd.Series).fillna(0) + # Convert to float64 to match reference file data types + expanded_cols = expanded_cols.astype(np.float64) + expanded_cols = expanded_cols.add_prefix(f"{col}_") + df = df.drop(col, axis=1).join(expanded_cols) + + # Save processed file + df.to_pickle(file_path) + print(f" ✅ Post-processing completed for {os.path.basename(file_path)}") + + return output_dir + +def generate_base_dataset_initial_only(variable_definitions): + """Generate base training dataset for initial-only mode (excludes Y_ variables from final_spinup files)""" + print(f"\n{'='*80}") + print("STEP 1: Base Dataset Generation (Initial-Only Mode)") + print(f"{'='*80}") + + # File paths from config + surface_data_files = config.surface_data_files + ad_spinup_history_files = config.ad_spinup_history_files + ad_spinup_restart_files = config.ad_spinup_restart_files + # NOTE: We do NOT load final_spinup files for initial-only mode + + # Forcing data files - dynamically find files containing variable names + forcing_files = {} + for var_name in variable_definitions['time_series_vars']: + # Look for files containing the variable name in forcing_netcdf directory + pattern = os.path.join(config.forcing_netcdf_output_dir, f'*{var_name}*1980-1999.nc') + matching_files = glob.glob(pattern) + if matching_files: + forcing_files[var_name] = matching_files[0] # Use first match + print(f"✅ Found forcing file: {os.path.basename(matching_files[0])}") + else: + print(f"⚠️ Forcing file not found for {var_name}: {pattern}") + + print(f"Found {len(forcing_files)} forcing files") + + # Output directory + output_dir = os.path.join(config.output_dir, 'initial_condition_dataset') + os.makedirs(output_dir, exist_ok=True) + + print("Loading NetCDF files...") + start_time = time.time() + + # Load all required files (excluding final_spinup files) + ds1 = nc.Dataset(surface_data_files[0]) # Surface data + ds2 = nc.Dataset(ad_spinup_history_files[0]) # AD-SPINUP history + ds10 = nc.Dataset(ad_spinup_restart_files[0]) # AD-SPINUP restart + + # Load forcing data + ds_forcing = {} + for var_name, file_path in forcing_files.items(): + ds_forcing[var_name] = nc.Dataset(file_path) + print(f"✅ Forcing data loaded: {var_name}") + + print(f"All files loaded in {time.time() - start_time:.2f} seconds") + + # Get coordinates and build spatial filtering + lats = ds2.variables['lat'][:] + lons = ds2.variables['lon'][:] + landmask = ds2.variables['landfrac'][:] + + # Filter for land gridcells + valid_mask = (landmask > 0) + valid_gridcells = np.where(valid_mask)[0] + + print(f"Total land gridcells: {len(valid_gridcells)}") + print(f"Latitude range: [{lats.min():.2f}, {lats.max():.2f}]") + print(f"Longitude range: [{lons.min():.2f}, {lons.max():.2f}]") + + # Build KDTree indices + print("Building KDTree indices...") + + # Restart file coordinates + gridcell_lat = ds10.variables['grid1d_lat'][:] + gridcell_lon = ds10.variables['grid1d_lon'][:] + restart_grid_coords = np.vstack((gridcell_lat, gridcell_lon)).T + restart_tree = cKDTree(restart_grid_coords) + + # Forcing file coordinates + forcing_lats = list(ds_forcing.values())[0].variables['LATIXY'][0, :] + forcing_lons = list(ds_forcing.values())[0].variables['LONGXY'][0, :] + forcing_grid_coords = np.vstack((forcing_lats, forcing_lons)).T + forcing_tree = cKDTree(forcing_grid_coords) + + # Query coordinates for valid gridcells + query_coords = np.array([(lats[i], lons[i]) for i in valid_gridcells]) + _, all_restart_indices = restart_tree.query(query_coords, k=1) + _, all_forcing_indices = forcing_tree.query(query_coords, k=1) + + print("✅ KDTree indices built") + + # Pre-load forcing data into memory for optimization + print("🚀 Pre-loading forcing data into memory...") + forcing_data = {} + for var_name, ds in ds_forcing.items(): + forcing_data[var_name] = ds.variables[var_name][:, 0, :] # (time, 1, grid_cells) + print(f" Loaded {var_name}: {forcing_data[var_name].shape}") + + # Close forcing NetCDF files (data is now in memory) + for ds in ds_forcing.values(): + ds.close() + + # Build index mappings + pft_gridcell_index = ds10.variables['pfts1d_gridcell_index'][:] + column_gridcell_index = ds10.variables['cols1d_gridcell_index'][:] + + pft_map = {} + column_map = {} + unique_gridcell_ids = np.unique(pft_gridcell_index) + for grid_id in unique_gridcell_ids: + pft_map[grid_id] = np.where(pft_gridcell_index == grid_id)[0] + column_map[grid_id] = np.where(column_gridcell_index == grid_id)[0] + + print("✅ Index mappings built") + + # Process data in batches + batch_size = 1000 + batch_number = 1 + batch_files = [] # Track all generated files for post-processing + + print(f"\nProcessing {len(valid_gridcells)} gridcells in batches of {batch_size}...") + + for start_idx in range(0, len(valid_gridcells), batch_size): + end_idx = min(start_idx + batch_size, len(valid_gridcells)) + batch_gridcells = valid_gridcells[start_idx:end_idx] + batch_restart_indices = all_restart_indices[start_idx:end_idx] + batch_forcing_indices = all_forcing_indices[start_idx:end_idx] + + print(f"\nProcessing batch {batch_number}: gridcells {start_idx+1}-{end_idx}") + print(f" Batch size: {len(batch_gridcells)} gridcells") + print(f" batch_gridcells range: {batch_gridcells[0]} to {batch_gridcells[-1]}") + print(f" Unique gridcells in batch: {len(set(batch_gridcells))}") + if len(set(batch_gridcells)) != len(batch_gridcells): + print(f" ⚠️ WARNING: Duplicate gridcells detected in batch!") + batch_start_time = time.time() + + # Initialize data dictionary dynamically - completely from CNP_IO file + data_dict = {} + + # Add time series variables (forcing data) + for var_name in variable_definitions['time_series_vars']: + data_dict[var_name] = [] + + # Add surface properties + for var_name in variable_definitions['surface_vars']: + data_dict[var_name] = [] + + # Add scalar variables (X only, no Y_ variables for initial-only mode) + for var_name in variable_definitions['scalar_vars']: + data_dict[var_name] = [] + + # Add PFT variables (X only, no Y_ variables) + for var_name in variable_definitions['pft_1d_vars']: + data_dict[var_name] = [] + # NOTE: We do NOT add Y_ variables for initial-only mode + + # Add 2D soil variables (X only, no Y_ variables) + for var_name in variable_definitions['variables_2d_soil']: + data_dict[var_name] = [] + # NOTE: We do NOT add Y_ variables for initial-only mode + + # Add water variables (X only, no Y_ variables) + for var_name in variable_definitions['water_vars']: + data_dict[var_name] = [] + # NOTE: We do NOT add Y_ variables for initial-only mode + + # Add pool variables (X only, no Y_ variables) + for var_name in variable_definitions['pool_vars']: + data_dict[var_name] = [] + # NOTE: We do NOT add Y_ variables for initial-only mode + + # Process each gridcell in the batch + for k, gridcell_idx in enumerate(batch_gridcells): + if k % 100 == 0: + print(f" Processing gridcell {k}/{len(batch_gridcells)} (idx={gridcell_idx})") + + # Get indices + restart_idx = batch_restart_indices[k] + forcing_idx = batch_forcing_indices[k] + gridcell_id = restart_idx + 1 + + pft_indices_for_cell = pft_map.get(gridcell_id, []) + column_indices_for_cell = column_map.get(gridcell_id, []) + + # Debug: Check data_dict length after each gridcell + if k == 0: + print(f" After first gridcell: Latitude={len(data_dict['Latitude'])}, FLDS={len(data_dict.get('FLDS', []))}") + print(f" pft_indices_for_cell length: {len(pft_indices_for_cell)}") + print(f" column_indices_for_cell length: {len(column_indices_for_cell)}") + if k == 999: + print(f" After 1000th gridcell: Latitude={len(data_dict['Latitude'])}, FLDS={len(data_dict.get('FLDS', []))}") + + # Process forcing data (time series variables) - store raw data like original script + for var_name in variable_definitions['time_series_vars']: + if var_name in forcing_data: + time_series = forcing_data[var_name][:, forcing_idx] + # Store raw time series like original script (will be processed later) + data_dict[var_name].append(time_series) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + + # Process surface properties + for var_name in variable_definitions['surface_vars']: + if var_name == 'Latitude': + data_dict[var_name].append(lats[gridcell_idx]) + elif var_name == 'Longitude': + data_dict[var_name].append(lons[gridcell_idx]) + elif var_name == 'landfrac': + # landfrac comes from history file (ds2), not surface file (ds1) + # Convert to float64 to match reference file + landfrac_val = ds2.variables['landfrac'][gridcell_idx] + if hasattr(landfrac_val, 'data'): # MaskedArray + landfrac_val = landfrac_val.data + data_dict[var_name].append(float(landfrac_val)) + elif var_name == 'PCT_CLAY': + # Store PCT_CLAY as a list (all levels) + pct_clay_data = ds1.variables['PCT_CLAY'][:, gridcell_idx] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(pct_clay_data, 'data'): # MaskedArray + pct_clay_data = pct_clay_data.data + data_dict[var_name].append(pct_clay_data.astype(np.float64).tolist()) + elif var_name == 'PCT_SAND': + # Store PCT_SAND as a list (all levels) + pct_sand_data = ds1.variables['PCT_SAND'][:, gridcell_idx] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(pct_sand_data, 'data'): # MaskedArray + pct_sand_data = pct_sand_data.data + data_dict[var_name].append(pct_sand_data.astype(np.float64).tolist()) + elif var_name.startswith('PCT_NAT_PFT_') or var_name.startswith('PCT_CLAY_') or var_name.startswith('PCT_SAND_'): + # Handle 2D variables with level indices + if '_' in var_name: + level_idx = int(var_name.split('_')[-1]) + base_var = '_'.join(var_name.split('_')[:-1]) # e.g., 'PCT_CLAY' + if base_var in ds1.variables: + pct_val = ds1.variables[base_var][level_idx, gridcell_idx] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(pct_val, 'data'): # MaskedArray + pct_val = pct_val.data + data_dict[var_name].append(float(pct_val)) + else: + data_dict[var_name].append(0.0) + else: + data_dict[var_name].append(0.0) + elif var_name in ds1.variables: + # 1D variables + val = ds1.variables[var_name][gridcell_idx] + # Convert MaskedArray to regular array + if hasattr(val, 'data'): # MaskedArray + val = val.data + # Handle integer variables + if var_name in ['SOIL_COLOR', 'SOIL_ORDER']: + data_dict[var_name].append(int(val)) + else: + data_dict[var_name].append(float(val)) + else: + data_dict[var_name].append(0.0) # Add default value if variable not found + + # Process scalar variables from history file (X only, no Y_ variables for initial-only mode) + for var_name in variable_definitions['scalar_vars']: + if var_name in ds2.variables: + val = ds2.variables[var_name][0, gridcell_idx] + # Convert MaskedArray to regular array + if hasattr(val, 'data'): # MaskedArray + val = val.data + # Handle integer variables + if var_name in ['SOIL_COLOR', 'SOIL_ORDER']: + data_dict[var_name].append(int(val)) + else: + data_dict[var_name].append(float(val)) + else: + data_dict[var_name].append(0.0) # Add default value if variable not found + + # Process PFT variables (X only, no Y_ variables) + for var_name in variable_definitions['pft_1d_vars']: + if var_name in ds10.variables: + # X values from restart file + x_val = ds10.variables[var_name][pft_indices_for_cell] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(x_val, 'data'): # MaskedArray + x_val = x_val.data + data_dict[var_name].append(x_val.astype(np.float64).tolist()) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + + # Process 2D soil variables (X only, no Y_ variables) + for var_name in variable_definitions['variables_2d_soil']: + if var_name in ds10.variables: + # X values + x_val = ds10.variables[var_name][column_indices_for_cell, :] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(x_val, 'data'): # MaskedArray + x_val = x_val.data + data_dict[var_name].append(x_val.astype(np.float64).tolist()) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + + # Process water variables (X only, no Y_ variables) + for var_name in variable_definitions['water_vars']: + if var_name in ds10.variables: + # X values + x_val = ds10.variables[var_name][column_indices_for_cell, :] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(x_val, 'data'): # MaskedArray + x_val = x_val.data + data_dict[var_name].append(x_val.astype(np.float64).tolist()) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + + # Process pool variables (X only, no Y_ variables) + for var_name in variable_definitions['pool_vars']: + # Skip if already processed as PFT variable + if var_name in variable_definitions['pft_1d_vars']: + continue + + if var_name in ds10.variables: + # X values + x_val = ds10.variables[var_name][column_indices_for_cell] + # Convert MaskedArray to regular array and ensure float64 + if hasattr(x_val, 'data'): # MaskedArray + x_val = x_val.data + data_dict[var_name].append(x_val.astype(np.float64).tolist()) + else: + data_dict[var_name].append([]) # Add empty list if variable not found + + # Create DataFrame and save + print(f" Creating DataFrame...") + + # Debug: Check array lengths + print(f" Debug: Checking array lengths...") + for key, values in data_dict.items(): + print(f" {key}: {len(values)} items") + + # Check for length mismatches + lengths = {key: len(values) for key, values in data_dict.items()} + unique_lengths = set(lengths.values()) + if len(unique_lengths) > 1: + print(f" ❌ Length mismatch detected!") + for length in unique_lengths: + vars_with_length = [k for k, v in lengths.items() if v == length] + print(f" Length {length}: {len(vars_with_length)} variables") + if length != max(unique_lengths): + print(f" Variables: {vars_with_length[:5]}{'...' if len(vars_with_length) > 5 else ''}") + + df_batch = pd.DataFrame(data_dict) + + print(f" Saving to disk...") + batch_save_path = f"{output_dir}/initial_condition_batch_{batch_number:02d}.pkl" + df_batch.to_pickle(batch_save_path) + batch_files.append(batch_save_path) # Add to list for post-processing + + batch_time = time.time() - batch_start_time + print(f"✅ Batch {batch_number} completed: {batch_time:.2f}s") + print(f" Path: {batch_save_path}") + print(f" Shape: {df_batch.shape}") + print(f" Columns: {len(df_batch.columns)}") + + batch_number += 1 + + # Cleanup NetCDF files + print(f"\nCleaning up NetCDF files...") + ds1.close() + ds2.close() + ds10.close() + + print("✅ All NetCDF files closed") + print(f"✅ Base dataset generation completed!") + print(f"Total batches: {batch_number - 1}") + print(f"Output directory: {output_dir}") + + # Post-processing like original script + print(f"\n{'='*80}") + print("POST-PROCESSING: Converting to monthly averages and expanding variables") + print(f"{'='*80}") + + # Process all generated files + for file_path in batch_files: + print(f"Processing {os.path.basename(file_path)}...") + + # Load the file + df = pd.read_pickle(file_path) + + # Process time series columns (convert to monthly averages) + print(f" Processing time series columns...") + for col in variable_definitions['time_series_vars']: + if col in df.columns: + print(f" Processing {col}...") + df[col] = df[col].apply(calculate_monthly_avg) + + # Process list columns (expand PCT variables) + print(f" Processing list columns...") + list_like_columns = ['PCT_CLAY', 'PCT_SAND', 'PCT_NAT_PFT'] + for col in list_like_columns: + if col in df.columns: + print(f" Expanding {col}...") + expanded_cols = df[col].apply(pd.Series).fillna(0) + # Convert to float64 to match reference file data types + expanded_cols = expanded_cols.astype(np.float64) + expanded_cols = expanded_cols.add_prefix(f"{col}_") + df = df.drop(col, axis=1).join(expanded_cols) + + # Save processed file + df.to_pickle(file_path) + print(f" ✅ Post-processing completed for {os.path.basename(file_path)}") + + return output_dir + +def generate_enhanced_dataset(base_output_dir, variable_definitions, initial_only_mode=False): + """Generate enhanced dataset (37_dataset.py logic)""" + print(f"\n{'='*80}") + print("STEP 2: Enhanced Dataset Generation (37_dataset.py)") + print(f"{'='*80}") + + # Enhanced output directory + if initial_only_mode: + # For initial-only mode, keep output in initial_condition_dataset directory + enhanced_output_dir = base_output_dir + else: + # For regular enhanced mode, use enhanced_training_dataset directory + enhanced_output_dir = os.path.join(os.path.dirname(base_output_dir), 'enhanced_training_dataset') + os.makedirs(enhanced_output_dir, exist_ok=True) + + # Get base PKL files - handle different naming patterns + base_files = sorted(glob.glob(os.path.join(base_output_dir, "training_data_batch_*.pkl"))) + if not base_files: + # Try initial_condition_batch pattern for initial-only mode + base_files = sorted(glob.glob(os.path.join(base_output_dir, "initial_condition_batch_*.pkl"))) + print(f"Found {len(base_files)} base PKL files to enhance") + + if not base_files: + print("❌ No base PKL files found") + return base_output_dir + + # Load restart files for enhancement + file_path10 = config.ad_spinup_restart_files[0] + file_path17 = config.final_spinup_restart_files[0] + + # Create multiple restart files for Y variables (using same file for demo) + ds_r_list = [nc.Dataset(file_path17) for _ in range(5)] # 5 copies for averaging + + try: + ds_restart = nc.Dataset(file_path10) + restart_tree, restart_coords = build_restart_kdtree(ds_restart) + col_index_map = build_column_index_map(ds_restart) + pft_index_map = build_pft_index_map(ds_restart) + + for i, base_file in enumerate(base_files, 1): + print(f"\nProcessing file {i}/{len(base_files)}: {os.path.basename(base_file)}") + + try: + # Read base dataset + df = pd.read_pickle(base_file) + print(f" Base dataset shape: {df.shape}") + + # Add pool variables + print(" Adding pool variables...") + for pool_var in variable_definitions['pool_vars']: + if pool_var not in df.columns: + # Create dummy pool variable (should be loaded from restart files) + df[pool_var] = [0.0] * len(df) + if not initial_only_mode: + df[f'Y_{pool_var}'] = [0.0] * len(df) + print(f" Added {pool_var} and Y_{pool_var}") + else: + print(f" Added {pool_var} (skipped Y_{pool_var} for initial-only mode)") + + # Add Y_ variables for 1D PFT variables (skip in initial-only mode) + if not initial_only_mode: + print(" Adding Y_ variables for 1D PFT...") + for var in variable_definitions['pft_1d_vars']: + y_var = f"Y_{var}" + if y_var not in df.columns and var in df.columns: + # Create dummy Y_ variable + df[y_var] = [0.0] * len(df) + print(f" Added {y_var}") + else: + print(" Skipping Y_ variables for 1D PFT (initial-only mode)") + + # Add Y_ variables for 2D soil variables (skip in initial-only mode) + if not initial_only_mode: + print(" Adding Y_ variables for 2D soil...") + for var in variable_definitions['variables_2d_soil']: + y_var = f"Y_{var}" + if y_var not in df.columns and var in df.columns: + # Create dummy Y_ variable + df[y_var] = [0.0] * len(df) + print(f" Added {y_var}") + else: + print(" Skipping Y_ variables for 2D soil (initial-only mode)") + + # Add Y_ variables for water variables (skip in initial-only mode) + if not initial_only_mode: + print(" Adding Y_ variables for water...") + for var in variable_definitions['water_vars']: + y_var = f"Y_{var}" + if y_var not in df.columns and var in df.columns: + # Create dummy Y_ variable + df[y_var] = [0.0] * len(df) + print(f" Added {y_var}") + else: + print(" Skipping Y_ variables for water (initial-only mode)") + + # Filter to keep only variables defined in CNP_IO file + all_cnp_vars = [] + for key, vars_list in variable_definitions.items(): + if isinstance(vars_list, list): + all_cnp_vars.extend(vars_list) + # Also add Y_ counterparts + all_cnp_vars.extend([f"Y_{v}" for v in vars_list]) + + cnp_vars_set = set(all_cnp_vars) + current_vars = set(df.columns) + + # Find variables to keep + vars_to_keep = current_vars & cnp_vars_set + vars_to_remove = current_vars - cnp_vars_set + + print(f" Variables to keep: {len(vars_to_keep)}") + print(f" Variables to remove: {len(vars_to_remove)}") + + if vars_to_remove: + print(f" 🗑️ Removing {len(vars_to_remove)} variables not in CNP_IO file") + df_enhanced = df[list(vars_to_keep)] + else: + df_enhanced = df.copy() + + # Save enhanced dataset + enhanced_file = os.path.join(enhanced_output_dir, f"enhanced_monthly_training_data_batch_{i:02d}.pkl") + df_enhanced.to_pickle(enhanced_file) + print(f" ✅ Enhanced dataset saved: {os.path.basename(enhanced_file)}") + print(f" 📐 Enhanced shape: {df_enhanced.shape}") + + except Exception as e: + print(f" ❌ Failed to process file: {e}") + continue + + return enhanced_output_dir + + finally: + ds_restart.close() + for ds in ds_r_list: + try: + ds.close() + except Exception: + pass + +def add_pft_variables(enhanced_output_dir, variable_definitions): + """Add PFT variables from CLM parameters file (1_add_pft_to_dataset.py logic)""" + print(f"\n{'='*80}") + print("STEP 3: Adding PFT Variables (1_add_pft_to_dataset.py)") + print(f"{'='*80}") + + # Load CLM parameters file + print("Reading CLM parameters NetCDF file...") + print(f"File path: {config.clm_params_nc_path}") + + if not os.path.exists(config.clm_params_nc_path): + print(f"❌ Error: CLM parameters file not found: {config.clm_params_nc_path}") + return enhanced_output_dir + + ds = nc.Dataset(config.clm_params_nc_path) + + # Get PFT variables dynamically from CNP_IO file + pft_params = variable_definitions.get('pft_parameters', []) + print(f"CNP_IO file defines {len(pft_params)} PFT parameter variables") + + # Remove 'pft_' prefix from variable names for NetCDF lookup + target_vars = [] + for var in pft_params: + if var.startswith('pft_'): + target_vars.append(var[4:]) # Remove 'pft_' prefix + else: + target_vars.append(var) + + print(f"PFT parameters to add: {len(target_vars)}") + print(f"PFT variables: {target_vars[:10]}...") + + broadcast_feature_dict = {} + for var in target_vars: + if var in ds.variables: + raw_vals = ds.variables[var][:17] + # Skip variables if any value is NaN or masked (missing) + if np.any(np.isnan(raw_vals)) or np.ma.is_masked(raw_vals): + print(f"Skipped {var}: contains NaN or masked values") + continue + broadcast_feature_dict[var] = list(map(float, raw_vals)) + print(f"Added: {var} (length {len(raw_vals)})") + else: + print(f"Skipped {var}: not found in NetCDF") + + print(f"\n✅ Successfully loaded {len(broadcast_feature_dict)} PFT variables from NetCDF") + + # Get enhanced PKL files - handle different naming patterns + input_files = sorted(glob.glob(os.path.join(enhanced_output_dir, "enhanced_monthly_training_data_batch_*.pkl"))) + if not input_files: + # Try initial_condition_batch pattern for initial-only mode + input_files = sorted(glob.glob(os.path.join(enhanced_output_dir, "initial_condition_batch_*.pkl"))) + + print(f"\n🔍 Found {len(input_files)} enhanced PKL files to process") + + if len(input_files) == 0: + print("❌ No enhanced PKL files found.") + return enhanced_output_dir + + # Process each PKL file + for i, file_path in enumerate(input_files, 1): + print(f"\nProcessing file {i}/{len(input_files)}: {os.path.basename(file_path)}") + + try: + # Read PKL file + df = pd.read_pickle(file_path) + original_shape = df.shape + print(f" Original shape: {original_shape}") + + # Check if PFT variables already exist + existing_pft_cols = [col for col in df.columns if col.startswith("pft_")] + if existing_pft_cols: + print(f" ⚠️ File already contains {len(existing_pft_cols)} PFT variables, skipping addition") + continue + + # Add each variable as a vector column with pft_ prefix + print(" Adding PFT variables...") + for var, val_list in broadcast_feature_dict.items(): + df["pft_" + var] = [val_list] * len(df) # Add the same list to each row + + new_shape = df.shape + print(f" ✅ Successfully added {len(broadcast_feature_dict)} PFT variables") + print(f" 📐 New data shape: {original_shape} → {new_shape}") + + # Save in-place (overwrite original file) + df.to_pickle(file_path) + print(f" ✅ File saved: {os.path.basename(file_path)}") + + except Exception as e: + print(f" ❌ Failed to process file: {e}") + continue + + ds.close() + + print(f"\n✅ PFT variables addition completed!") + print(f" - Total files: {len(input_files)}") + print(f" - PFT variables added: {len(broadcast_feature_dict)}") + + return enhanced_output_dir + +def final_variable_cleanup(enhanced_output_dir, variable_definitions): + """Final variable cleanup to ensure only CNP_IO variables remain (2_rm_variables.py logic)""" + print(f"\n{'='*80}") + print("STEP 4: Final Variable Cleanup (2_rm_variables.py)") + print("Keeping only variables defined in CNP_IO file") + print(f"{'='*80}") + + # Get all expected variables from CNP_IO file + all_expected_vars = set() + + # Add all variables from CNP_IO file + for key, vars_list in variable_definitions.items(): + if isinstance(vars_list, list): + all_expected_vars.update(vars_list) + # Also add Y_ counterparts + all_expected_vars.update([f"Y_{v}" for v in vars_list]) + + # Add PFT variables with pft_ prefix + pft_vars = variable_definitions.get('pft_1d_vars', []) + for var in pft_vars: + if not var.startswith('pft_'): + all_expected_vars.add(f"pft_{var}") + else: + all_expected_vars.add(var) + + # Add PFT parameter variables from CNP_IO file + pft_params = variable_definitions.get('pft_parameters', []) + for var in pft_params: + all_expected_vars.add(var) # These already have pft_ prefix + + print(f"CNP_IO file defines {len(all_expected_vars)} variables") + + # Get enhanced PKL files - handle different naming patterns + input_files = sorted(glob.glob(os.path.join(enhanced_output_dir, "enhanced_monthly_training_data_batch_*.pkl"))) + if not input_files: + # Try initial_condition_batch pattern for initial-only mode + input_files = sorted(glob.glob(os.path.join(enhanced_output_dir, "initial_condition_batch_*.pkl"))) + + print(f"\n🔍 Found {len(input_files)} enhanced PKL files to process") + + if len(input_files) == 0: + print("❌ No enhanced PKL files found.") + return enhanced_output_dir + + # Process each PKL file + for i, file_path in enumerate(input_files, 1): + print(f"\nProcessing file {i}/{len(input_files)}: {os.path.basename(file_path)}") + + try: + # Read PKL file + df = pd.read_pickle(file_path) + original_shape = df.shape + print(f" Original shape: {original_shape}") + + # Expand list-like columns (PCT_CLAY, PCT_SAND) + list_like_columns = ['PCT_CLAY', 'PCT_SAND'] + for col in list_like_columns: + if col in df.columns: + print(f" Expanding {col}...") + expanded_cols = df[col].apply(pd.Series).fillna(0) + expanded_cols = expanded_cols.add_prefix(f"{col}_") + df = df.drop(col, axis=1).join(expanded_cols) + print(f" Expanded {col} into {len(expanded_cols.columns)} columns") + + current_vars = set(df.columns) + vars_to_keep = current_vars & all_expected_vars + vars_to_remove = current_vars - all_expected_vars + + print(f" Variables in file: {len(current_vars)}") + print(f" Variables to keep: {len(vars_to_keep)}") + print(f" Variables to remove: {len(vars_to_remove)}") + + if vars_to_remove: + print(f" 🗑️ Removing {len(vars_to_remove)} variables not in CNP_IO file") + if len(vars_to_remove) <= 10: + print(f" Variables to remove: {sorted(list(vars_to_remove))}") + else: + print(f" Sample variables to remove: {sorted(list(vars_to_remove))[:10]}...") + + # Keep only expected variables + df_cleaned = df[list(vars_to_keep)] + + new_shape = df_cleaned.shape + print(f" ✅ Successfully removed {len(vars_to_remove)} variables") + print(f" 📐 New data shape: {original_shape} → {new_shape}") + + # Save in-place (overwrite original file) + df_cleaned.to_pickle(file_path) + print(f" ✅ File saved: {os.path.basename(file_path)}") + else: + print(" ℹ️ No extra variables found to remove. File already matches CNP_IO file.") + if len(current_vars) != len(all_expected_vars): + print(f" ⚠️ Warning: Column count mismatch. Current: {len(current_vars)}, Expected: {len(all_expected_vars)}") + print(f" Missing from current: {list(all_expected_vars - current_vars)[:5]}...") + print(f" Extra in current: {list(current_vars - all_expected_vars)[:5]}...") + + except Exception as e: + print(f" ❌ Failed to process file: {e}") + continue + + print(f"\n✅ Variable cleanup completed!") + print(f" - Total files: {len(input_files)}") + print(f" - Kept only variables defined in CNP_IO file") + print(f" - Target variable count: {len(all_expected_vars)}") + + return enhanced_output_dir + +def generate_forcing_only_dataset(): + """Generate forcing-only dataset (raw time series, no monthly averaging)""" + print(f"\n{'='*80}") + print("FORCING-ONLY DATASET GENERATION") + print(f"{'='*80}") + + # File paths from config + surface_data_files = config.surface_data_files + ad_spinup_history_files = config.ad_spinup_history_files + + # Forcing data files - dynamically find files containing variable names + forcing_data_files = {} + forcing_variables = ['FLDS', 'FSDS', 'PSRF', 'QBOT', 'PRECTmms', 'TBOT'] + + for var_name in forcing_variables: + # Look for files containing the variable name in forcing_netcdf directory + pattern = os.path.join(config.forcing_netcdf_output_dir, f'*{var_name}*1980-1999.nc') + matching_files = glob.glob(pattern) + if matching_files: + forcing_data_files[var_name] = matching_files[0] # Use first match + print(f"✅ Found forcing file: {os.path.basename(matching_files[0])}") + else: + print(f"⚠️ Forcing file not found for {var_name}: {pattern}") + + print(f"Found {len(forcing_data_files)} forcing files") + + # Output directory + output_dir = os.path.join(config.output_dir, 'forcing_only_dataset') + os.makedirs(output_dir, exist_ok=True) + + print(f"Output directory: {output_dir}") + + # Load NetCDF files + print("Loading NetCDF files...") + ds1 = nc.Dataset(surface_data_files[0]) # Surface data + ds2 = nc.Dataset(ad_spinup_history_files[0]) # History file + + # Load forcing files + ds_forcing = {} + for var_name, file_path in forcing_data_files.items(): + ds_forcing[var_name] = nc.Dataset(file_path) + print(f" Loaded {var_name}: {file_path}") + + # Get coordinates and land mask + lats = ds2.variables['lat'][:] + lons = ds2.variables['lon'][:] + landmask = ds2.variables['landfrac'][:] + + # Filter land gridcells + valid_mask = (landmask > 0) + valid_gridcells = np.where(valid_mask)[0] + + print(f"Total land gridcells: {len(valid_gridcells)}") + print(f"Latitude range: [{lats.min():.2f}, {lats.max():.2f}]") + print(f"Longitude range: [{lons.min():.2f}, {lons.max():.2f}]") + + # Build KDTree for forcing data mapping + print("Building KDTree for forcing data mapping...") + query_coords = np.array([(lats[i], lons[i]) for i in valid_gridcells]) + + # Use first forcing file for coordinate mapping + first_forcing_var = list(forcing_data_files.keys())[0] + first_forcing_ds = ds_forcing[first_forcing_var] + forcing_lats = first_forcing_ds.variables['LATIXY'][:].flatten() + forcing_lons = first_forcing_ds.variables['LONGXY'][:].flatten() + forcing_coords = np.vstack((forcing_lats, forcing_lons)).T + + forcing_tree = cKDTree(forcing_coords) + _, all_forcing_indices = forcing_tree.query(query_coords, k=1) + + print("✅ KDTree indices built") + + # Pre-load forcing data into memory for optimization + print("🚀 Pre-loading forcing data into memory...") + forcing_data = {} + for var_name, ds in ds_forcing.items(): + forcing_data[var_name] = ds.variables[var_name][:, 0, :] # (time, 1, grid_cells) + print(f" Loaded {var_name}: {forcing_data[var_name].shape}") + + # Close forcing NetCDF files (data is now in memory) + for ds in ds_forcing.values(): + ds.close() + + # Process data in batches + batch_size = 1000 + batch_number = 1 + batch_files = [] # Track all generated files + + print(f"\nProcessing {len(valid_gridcells)} gridcells in batches of {batch_size}...") + + for start_idx in range(0, len(valid_gridcells), batch_size): + end_idx = min(start_idx + batch_size, len(valid_gridcells)) + batch_gridcells = valid_gridcells[start_idx:end_idx] + batch_forcing_indices = all_forcing_indices[start_idx:end_idx] + + print(f"\nProcessing batch {batch_number}: gridcells {start_idx+1}-{end_idx}") + print(f" Batch size: {len(batch_gridcells)} gridcells") + batch_start_time = time.time() + + # Initialize data dictionary - only forcing data and basic geographic info + data_dict = { + 'landfrac': [], + 'Latitude': [], + 'Longitude': [], + 'FLDS': [], + 'PSRF': [], + 'FSDS': [], + 'QBOT': [], + 'PRECTmms': [], + 'TBOT': [], + } + + # Process each gridcell + for k, gridcell_idx in enumerate(batch_gridcells): + if k % 100 == 0: + print(f" Processing gridcell {k}/{len(batch_gridcells)} (idx={gridcell_idx})") + + # Get forcing index + forcing_idx = batch_forcing_indices[k] + + # Basic geographic info + landfrac_val = ds2.variables['landfrac'][gridcell_idx] + if hasattr(landfrac_val, 'data'): # MaskedArray + landfrac_val = landfrac_val.data + data_dict['landfrac'].append(float(landfrac_val)) + data_dict['Latitude'].append(float(lats[gridcell_idx])) + data_dict['Longitude'].append(float(lons[gridcell_idx])) + + # Forcing data (raw time series, no monthly averaging) - convert to list format + data_dict['FLDS'].append(forcing_data['FLDS'][:, forcing_idx].tolist()) + data_dict['PSRF'].append(forcing_data['PSRF'][:, forcing_idx].tolist()) + data_dict['FSDS'].append(forcing_data['FSDS'][:, forcing_idx].tolist()) + data_dict['QBOT'].append(forcing_data['QBOT'][:, forcing_idx].tolist()) + data_dict['PRECTmms'].append(forcing_data['PRECTmms'][:, forcing_idx].tolist()) + data_dict['TBOT'].append(forcing_data['TBOT'][:, forcing_idx].tolist()) + + # Create DataFrame and save + print(f" Creating DataFrame...") + df_batch = pd.DataFrame(data_dict) + + print(f" Saving to disk...") + batch_save_path = f"{output_dir}/forcing_data_batch_{batch_number:02d}.pkl" + df_batch.to_pickle(batch_save_path) + batch_files.append(batch_save_path) # Add to list for tracking + + batch_time = time.time() - batch_start_time + print(f"✅ Batch {batch_number} completed: {batch_time:.2f}s") + print(f" Path: {batch_save_path}") + print(f" Shape: {df_batch.shape}") + print(f" Forcing data length: {len(df_batch['FLDS'].iloc[0])}") + + batch_number += 1 + + # Cleanup NetCDF files + print(f"\nCleaning up NetCDF files...") + ds1.close() + ds2.close() + + print("✅ All NetCDF files closed") + print(f"✅ Forcing-only dataset generation completed!") + print(f"Total batches: {batch_number - 1}") + print(f"Output directory: {output_dir}") + + return output_dir + +def main(): + parser = argparse.ArgumentParser( + description="Enhanced Training Dataset Generation - Three modes: forcing-only, enhanced dataset, or initial-only" + ) + parser.add_argument( + "--forcing_only", + action="store_true", + help="Generate only forcing data (raw time series, no monthly averaging)" + ) + parser.add_argument( + "--enhanced_dataset", + action="store_true", + help="Generate complete enhanced dataset (includes PFT variables and all processing steps)" + ) + parser.add_argument( + "--initial_only", + action="store_true", + help="Generate initial condition dataset (excludes Y_ variables from final_spinup files)" + ) + args = parser.parse_args() + + # Validate arguments + options = [args.forcing_only, args.enhanced_dataset, args.initial_only] + if sum(options) > 1: + print("❌ Error: Cannot specify multiple options. Choose only one of: --forcing_only, --enhanced_dataset, or --initial_only") + sys.exit(1) + if sum(options) == 0: + print("❌ Error: Must specify one of: --forcing_only, --enhanced_dataset, or --initial_only") + sys.exit(1) + + try: + start_time = time.time() + + if args.forcing_only: + # Generate forcing-only dataset + print("🚀 Starting FORCING-ONLY dataset generation...") + final_output_dir = generate_forcing_only_dataset() + + total_time = time.time() - start_time + + print(f"\n{'='*80}") + print("🎉 FORCING-ONLY DATASET GENERATION COMPLETED!") + print(f"{'='*80}") + print(f"Total execution time: {total_time:.2f} seconds ({total_time/60:.2f} minutes)") + print(f"Final output directory: {final_output_dir}") + + # List final files + final_files = sorted(glob.glob(os.path.join(final_output_dir, "*.pkl"))) + print(f"Generated {len(final_files)} forcing-only dataset files:") + for file in final_files: + file_size = os.path.getsize(file) / (1024**3) # GB + print(f" {os.path.basename(file)}: {file_size:.2f} GB") + + print(f"\n✅ Forcing-only dataset ready!") + print(f"✅ Raw time series data (no monthly averaging)") + print(f"✅ Only forcing variables and basic geographic info") + + elif args.enhanced_dataset: + # Generate complete enhanced dataset + print("🚀 Starting COMPLETE ENHANCED DATASET generation...") + + # Parse CNP_IO variables + variable_definitions = parse_cnp_io_variables() + + # Step 1: Generate base dataset + base_output_dir = generate_base_dataset(variable_definitions) + + # Step 2: Generate enhanced dataset + enhanced_output_dir = generate_enhanced_dataset(base_output_dir, variable_definitions) + + # Step 3: Add PFT variables + pft_output_dir = add_pft_variables(enhanced_output_dir, variable_definitions) + + # Step 4: Final variable cleanup + final_output_dir = final_variable_cleanup(pft_output_dir, variable_definitions) + + total_time = time.time() - start_time + + print(f"\n{'='*80}") + print("🎉 COMPLETE ENHANCED TRAINING DATASET GENERATION COMPLETED!") + print(f"{'='*80}") + print(f"Total execution time: {total_time:.2f} seconds ({total_time/60:.2f} minutes)") + print(f"Final output directory: {final_output_dir}") + + # List final files + final_files = sorted(glob.glob(os.path.join(final_output_dir, "*.pkl"))) + print(f"Generated {len(final_files)} enhanced dataset files:") + for file in final_files: + file_size = os.path.getsize(file) / (1024**3) # GB + print(f" {os.path.basename(file)}: {file_size:.2f} GB") + + print(f"\n✅ Complete enhanced training dataset ready for machine learning!") + print(f"✅ All variables dynamically extracted from CNP_IO file!") + print(f"✅ Combines all 4 original scripts into one integrated workflow!") + + elif args.initial_only: + # Generate initial condition dataset (without Y_ variables) + print("🚀 Starting INITIAL-ONLY dataset generation...") + + # Parse CNP_IO variables + variable_definitions = parse_cnp_io_variables() + + # Step 1: Generate base dataset (without Y_ variables) + base_output_dir = generate_base_dataset_initial_only(variable_definitions) + + # Step 2: Generate enhanced dataset (same as regular enhanced dataset) + enhanced_output_dir = generate_enhanced_dataset(base_output_dir, variable_definitions, initial_only_mode=True) + + # Step 3: Add PFT variables (same as regular enhanced dataset) + pft_output_dir = add_pft_variables(enhanced_output_dir, variable_definitions) + + # Step 4: Final variable cleanup (same as regular enhanced dataset) + final_output_dir = final_variable_cleanup(pft_output_dir, variable_definitions) + + total_time = time.time() - start_time + + print(f"\n{'='*80}") + print("🎉 INITIAL-ONLY TRAINING DATASET GENERATION COMPLETED!") + print(f"{'='*80}") + print(f"Total execution time: {total_time:.2f} seconds ({total_time/60:.2f} minutes)") + print(f"Final output directory: {final_output_dir}") + + # List final files + final_files = sorted(glob.glob(os.path.join(final_output_dir, "*.pkl"))) + print(f"Generated {len(final_files)} initial-only dataset files:") + for file in final_files: + file_size = os.path.getsize(file) / (1024**3) # GB + print(f" {os.path.basename(file)}: {file_size:.2f} GB") + + print(f"\n✅ Initial-only training dataset ready!") + print(f"✅ Excludes all Y_ variables from final_spinup files!") + print(f"✅ Contains only initial condition variables!") + + except Exception as e: + print(f"\n❌ Error during execution: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + +if __name__ == "__main__": + main() diff --git a/scripts/training_data_generation/requirements.txt b/scripts/training_data_generation/requirements.txt index 879de41..b61228e 100644 --- a/scripts/training_data_generation/requirements.txt +++ b/scripts/training_data_generation/requirements.txt @@ -1,16 +1,19 @@ # Training Data Generation Requirements # Install with: pip install -r requirements.txt +# Updated based on working virtual environment: venv_py311 +# Only includes packages actually used by the scripts -# Core data processing -pandas>=2.0.0 -numpy>=1.24.0 -netCDF4>=1.6.0 -xarray>=2023.1.0 -dask>=2023.0.0 -scipy>=1.10.0 +# Core data processing libraries (required) +netCDF4==1.7.3 # For reading/writing NetCDF files +numpy==2.3.4 # Numerical computing +pandas==2.3.3 # Data manipulation +scipy==1.16.2 # Scientific computing (used for cKDTree) +xarray==2025.10.1 # For multi-dimensional arrays and NetCDF handling +dask==2025.10.0 # Parallel computing (required by xarray) -# Time handling for NetCDF files -cftime>=1.6.0 +# Time handling for NetCDF files (required for construct_forcing_20years.py) +cftime==1.6.5 # Calendar/time handling -# Utilities -tqdm>=4.65.0 \ No newline at end of file +# Note: The following packages are automatically installed as dependencies: +# certifi, click, cloudpickle, fsspec, importlib_metadata, locket, +# packaging, partd, python-dateutil, pytz, PyYAML, six, toolz, tzdata, zipp From 872c4f5bdb034de04b84196ff82f5a1f530121e7 Mon Sep 17 00:00:00 2001 From: Zhuowei Gu Date: Fri, 31 Oct 2025 13:17:42 -0400 Subject: [PATCH 32/51] Update forcing data generation and config for 20-year dataset --- .../construct_forcing_20years.py | 43 +++++++++++-------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/scripts/training_data_generation/python_scripts/construct_forcing_20years.py b/scripts/training_data_generation/python_scripts/construct_forcing_20years.py index bfe5b87..a3ff157 100644 --- a/scripts/training_data_generation/python_scripts/construct_forcing_20years.py +++ b/scripts/training_data_generation/python_scripts/construct_forcing_20years.py @@ -101,28 +101,37 @@ def process_forcing_variable(var_name, var_info): print(f"Target output file: {final_output_file}") # --- Find and build all monthly file list --- - print("Searching for monthly files...") + print(f"[{datetime.datetime.now()}] Searching for monthly files...") all_monthly_files = [] + # Pre-scan directory for efficiency (especially important for TES_NORTH with 4000+ files) + print(f"[{datetime.datetime.now()}] Pre-scanning directory for available files...") + available_files = set(os.listdir(data_dir)) + print(f"[{datetime.datetime.now()}] Found {len(available_files)} files in directory") + for year in range(start_year, end_year + 1): + year_found = 0 for month in range(1, 13): file_name = var_info['file_pattern'].format(year=year, month=month) - file_path = os.path.join(data_dir, file_name) - if os.path.exists(file_path): + if file_name in available_files: + file_path = os.path.join(data_dir, file_name) all_monthly_files.append(file_path) + year_found += 1 else: print(f" Warning: File {file_name} does not exist, skipping.") + print(f"[{datetime.datetime.now()}] {var_name}: year {year} -> found {year_found}/12 monthly files") if not all_monthly_files: print(f"❌ Error: No valid monthly files found for {var_name} in directory {data_dir}") return False - print(f"Found {len(all_monthly_files)} valid monthly files.") + print(f"[{datetime.datetime.now()}] {var_name}: total valid monthly files: {len(all_monthly_files)}") # --- Define Dask chunks --- dask_chunks = {'time': 366*8} try: + print(f"[{datetime.datetime.now()}] Opening {len(all_monthly_files)} monthly files with xarray.open_mfdataset ...") with xr.open_mfdataset( all_monthly_files, combine='nested', @@ -131,9 +140,10 @@ def process_forcing_variable(var_name, var_info): chunks=dask_chunks, parallel=False, ) as ds: + print(f"[{datetime.datetime.now()}] Dataset opened. Dims: {dict(ds.dims)} | Vars: {list(ds.data_vars)}") # === Separate static variables === - print("Separating static coordinate/ID variables...") + print(f"[{datetime.datetime.now()}] Separating static coordinate/ID variables...") static_var_names = ['gridID', 'LONGXY', 'LATIXY'] static_data = {} @@ -150,11 +160,11 @@ def process_forcing_variable(var_name, var_info): print(f"Error: {var_name} variable not found.") return False - print(f"Processing variables: {time_varying_vars}") + print(f"[{datetime.datetime.now()}] Processing variables: {time_varying_vars}") ds_temporal = ds[['time'] + time_varying_vars] # --- Time axis correction --- - print("Loading raw time coordinates...") + print(f"[{datetime.datetime.now()}] Loading raw time coordinates...") time_values_raw = ds_temporal['time'].load().values units = ds_temporal['time'].attrs['units'] calendar = ds_temporal['time'].attrs.get('calendar', 'standard') @@ -163,7 +173,7 @@ def process_forcing_variable(var_name, var_info): print(f" Time points: {len(time_values_raw)}") - print("Starting time coordinate correction...") + print(f"[{datetime.datetime.now()}] Starting time coordinate correction...") time_values_corrected = np.copy(time_values_raw).astype(float) cumulative_offset_days = 0.0 expected_step_days = 3.0 / 24.0 @@ -182,16 +192,16 @@ def process_forcing_variable(var_name, var_info): time_values_corrected[i+1] = time_values_raw[i+1] + cumulative_offset_days - print(f"Time correction completed. Corrected {jump_count} jumps.") + print(f"[{datetime.datetime.now()}] Time correction completed. Corrected {jump_count} jumps.") - print("Decoding corrected time values...") + print(f"[{datetime.datetime.now()}] Decoding corrected time values...") try: dates = cftime.num2date(time_values_corrected, units, calendar=calendar, only_use_cftime_datetimes=True) except ValueError: dates = cftime.num2date(time_values_corrected, units, calendar=calendar) # --- Check time monotonicity --- - print("Checking time monotonicity...") + print(f"[{datetime.datetime.now()}] Checking time monotonicity...") if len(dates) >= 2: diffs_corrected = np.diff(dates) zero_timedelta = datetime.timedelta(0) @@ -201,7 +211,7 @@ def process_forcing_variable(var_name, var_info): print(f"Error! Corrected time is still not monotonic!") return False else: - print("✓ Time coordinate check passed.") + print(f"[{datetime.datetime.now()}] ✓ Time coordinate check passed.") # --- Create dataset with corrected time --- ds_corrected_time = ds_temporal.copy(deep=False) @@ -214,11 +224,10 @@ def process_forcing_variable(var_name, var_info): ds_corrected_time[var_name_static] = data_array final_dataset = ds_corrected_time - print("Final dataset:") - print(final_dataset) + print(f"[{datetime.datetime.now()}] Final dataset ready. Summary dims: {dict(final_dataset.dims)}") # --- Write to file --- - print(f"Writing to file: {final_output_file}") + print(f"[{datetime.datetime.now()}] Writing to file: {final_output_file}") output_encoding = { var_name: {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, 'time': {'units': units, 'calendar': calendar, 'dtype': 'float64'} @@ -232,8 +241,8 @@ def process_forcing_variable(var_name, var_info): output_encoding['LATIXY'] = {'dtype': final_dataset['LATIXY'].dtype, '_FillValue': np.nan} final_dataset.to_netcdf(final_output_file, encoding=output_encoding, unlimited_dims=['time']) - print(f"✓ Successfully generated: {final_output_file}") - print(f" File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") + print(f"[{datetime.datetime.now()}] ✓ Successfully generated: {final_output_file}") + print(f"[{datetime.datetime.now()}] File size: {os.path.getsize(final_output_file) / (1024**3):.2f} GB") except Exception as e: print(f"\nError: {e}") From 6272d47c9e36b330a80ef44e4fd5c72fd6f47c3f Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Mon, 3 Nov 2025 11:43:42 -0800 Subject: [PATCH 33/51] update: revised pipeline doc and added TVA sample workflow --- TVA_1_Sample/locations.csv | 2 + TVA_1_Sample/run_workflow.py | 509 +++++++++++++++++++++++++++++++++++ docs/CNP_pipeline_runbook.md | 11 + 3 files changed, 522 insertions(+) create mode 100644 TVA_1_Sample/locations.csv create mode 100644 TVA_1_Sample/run_workflow.py diff --git a/TVA_1_Sample/locations.csv b/TVA_1_Sample/locations.csv new file mode 100644 index 0000000..682cc27 --- /dev/null +++ b/TVA_1_Sample/locations.csv @@ -0,0 +1,2 @@ +latitude,longitude +35.852511146198104,-84.18822361142391 diff --git a/TVA_1_Sample/run_workflow.py b/TVA_1_Sample/run_workflow.py new file mode 100644 index 0000000..2c3a362 --- /dev/null +++ b/TVA_1_Sample/run_workflow.py @@ -0,0 +1,509 @@ +#!/usr/bin/env python3 +""" +Geospatial workflow orchestrator for TVA locations. + +This script reads latitude/longitude points from a CSV file, extracts the +corresponding samples from the TVA training dataset, runs the trained CNP model +to generate AI predictions, converts the predictions to NetCDF, and updates the +restart NetCDF file with the new values. All outputs are written to per-location +directories so additional locations can be added to the CSV and reprocessed +without modifying the script. +""" + +from __future__ import annotations + +import argparse +import logging +import math +import re +import sys +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Tuple + +import numpy as np +import xarray as xr + +# Ensure the project root (LandSim) is on sys.path when running from TVA_1_Sample/ +SCRIPT_DIR = Path(__file__).resolve().parent +PROJECT_ROOT = SCRIPT_DIR.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +try: + import pandas as pd +except ImportError as exc: # pragma: no cover - fail fast for missing dependency + raise SystemExit("pandas is required to run this workflow. Please install it before proceeding.") from exc + +from config.training_config import parse_cnp_io_list +from scripts.run_inference_all import run_inference_all +from scripts.ai_predictions_to_netcdf import ( + load_ai_predictions, + create_netcdf_structure, + add_scalar_variables, + add_pft_variables, + add_soil_variables, +) +from scripts.ai_predictions_to_restart import ( + load_datasets, + create_spatial_mapping, + create_updated_restart_file, + auto_detect_variable_list, +) + + +LOGGER = logging.getLogger("landsim.workflow") + + +def parse_args() -> argparse.Namespace: + script_dir = Path(__file__).resolve().parent + project_root = script_dir.parent + default_locations = script_dir / "locations.csv" + default_variable_list = project_root / "CNP_IO_updated9_dev.txt" + default_model_config = project_root / "CNP_model_config_v01.txt" + default_output_root = project_root / "final_restartfile" + default_work_root = script_dir / "workflow_runs" + + parser = argparse.ArgumentParser( + description="Run TVA geospatial extraction, model inference, and restart updating workflow." + ) + parser.add_argument( + "--locations", + type=Path, + default=default_locations, + help=f"CSV file containing latitude/longitude entries (default: {default_locations})", + ) + parser.add_argument( + "--dataset-root", + type=Path, + default=Path("/global/cfs/cdirs/m4814/daweigao/14_Code/TVA_training_dataset_all"), + help="Directory containing TVA training dataset pickle batches.", + ) + parser.add_argument( + "--dataset-pattern", + default="enhanced_monthly_training_data_batch_*.pkl", + help="Glob pattern used to discover TVA dataset batches.", + ) + parser.add_argument( + "--model-path", + type=Path, + default=Path( + "/global/cfs/cdirs/m4814/daweigao/15_code_Landsim/LandSim/" + "cnp_results/run_20251030_192921/cnp_predictions/model.pth" + ), + help="Path to the trained model checkpoint (.pth).", + ) + parser.add_argument( + "--variable-list", + type=Path, + default=default_variable_list, + help="Path to the CNP IO configuration used for training (CNP_IO_*.txt).", + ) + parser.add_argument( + "--model-config", + type=Path, + default=default_model_config, + help="Optional model config override used during inference (CNP_model_config_*.txt).", + ) + parser.add_argument( + "--restart-file", + type=Path, + default=Path( + "/global/cfs/cdirs/m4814/daweigao/14_Code/TVA_restart/" + "uELM_knox_I1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc" + ), + help="Original restart NetCDF file that will be updated with AI predictions.", + ) + parser.add_argument( + "--output-root", + type=Path, + default=default_output_root, + help="Directory where updated restart files will be written.", + ) + parser.add_argument( + "--work-root", + type=Path, + default=default_work_root, + help="Directory for intermediate per-location artifacts (datasets, predictions, NetCDF files).", + ) + parser.add_argument( + "--max-files", + type=int, + default=None, + help="Optional limit on the number of TVA dataset batches to scan (useful for testing).", + ) + parser.add_argument( + "--skip-existing", + action="store_true", + help="Skip processing a location if the final restart file already exists.", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Logging verbosity for the workflow run.", + ) + return parser.parse_args() + + +def configure_logging(level: str) -> None: + logging.basicConfig( + level=getattr(logging, level.upper()), + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + stream=sys.stdout, + ) + + +def slugify(name: str) -> str: + """Convert a location label into a filesystem-friendly slug.""" + slug = re.sub(r"[^a-zA-Z0-9]+", "_", name.strip().lower()) + slug = slug.strip("_") + if slug: + return slug + # Provide deterministic fallback when location name is missing + return f"location_{abs(hash(name)) % 10_000}" + + +def read_locations(csv_path: Path) -> pd.DataFrame: + """Load and validate the locations CSV file.""" + if not csv_path.exists(): + raise FileNotFoundError(f"Locations CSV not found: {csv_path}") + df = pd.read_csv(csv_path) + required = {"latitude", "longitude"} + missing = required - set(c.lower() for c in df.columns) + if missing: + raise ValueError(f"Locations CSV must contain columns: {', '.join(sorted(required))}") + # Normalize column names to lower-case for convenience + df.columns = [c.lower() for c in df.columns] + return df + + +def determine_variable_map(variable_list_path: Optional[Path]) -> Dict[str, List[str]]: + """Parse the CNP IO variable list into groups needed for NetCDF conversion.""" + if variable_list_path is None: + return {"scalar": [], "pft1d": [], "soil2d": []} + if not variable_list_path.exists(): + raise FileNotFoundError(f"Variable list file not found: {variable_list_path}") + + parsed = parse_cnp_io_list(variable_list_path) + return { + "scalar": list(parsed.get("scalar_variables", [])), + "pft1d": list(parsed.get("pft_1d_variables", [])), + "soil2d": list(parsed.get("variables_2d_soil", parsed.get("x_list_columns_2d", []))), + } + + +def extract_location_dataset( + dataset_root: Path, + file_pattern: str, + latitude: float, + longitude: float, + output_path: Path, + max_files: Optional[int] = None, +) -> Optional[Path]: + """ + Build a location-specific dataset pickle by scanning TVA batches and + collecting the rows whose coordinates are closest to the requested latitude/longitude. + """ + files = sorted(dataset_root.glob(file_pattern)) + if max_files is not None: + files = files[:max_files] + + best_frames: List[pd.DataFrame] = [] + LOGGER.debug("Scanning %d TVA batches for samples near (lat=%.6f, lon=%.6f)", len(files), latitude, longitude) + + lon_candidates = {float(longitude)} + # Include 0-360 representation when the requested longitude is negative + lon_candidates.add((longitude + 360.0) % 360.0) + # Include -180–180 representation in case dataset already stores positive values over 180 + lon_candidates.add(((longitude + 180.0) % 360.0) - 180.0) + + def _minimum_lon_difference(series: pd.Series, targets: Iterable[float]) -> pd.Series: + diffs = [] + for target in targets: + diffs.append((series - target).abs()) + if not diffs: + return pd.Series(np.nan, index=series.index) + stacked = pd.concat(diffs, axis=1) + return stacked.min(axis=1, skipna=True) + + closest_example: Optional[Tuple[str, float, float, float, float]] = None + min_distance = math.inf + + for file_path in files: + try: + df_batch = pd.read_pickle(file_path) + except Exception as exc: # pragma: no cover - defensive logging + LOGGER.warning("Failed to load batch %s: %s", file_path, exc) + continue + + if "Latitude" not in df_batch.columns or "Longitude" not in df_batch.columns: + continue + + # Ensure numeric types to avoid comparison surprises (NaNs are preserved) + lat_series = pd.to_numeric(df_batch["Latitude"], errors="coerce") + lon_series = pd.to_numeric(df_batch["Longitude"], errors="coerce") + + lat_diff_series = (lat_series - latitude).abs() + lon_diff_series = _minimum_lon_difference(lon_series, lon_candidates) + + combined_diff = np.sqrt(np.square(lat_diff_series) + np.square(lon_diff_series)) + + valid_diff = combined_diff.dropna() + if not valid_diff.empty: + idx = valid_diff.idxmin() + distance = float(valid_diff.loc[idx]) + if distance < min_distance - 1e-12: + min_distance = distance + closest_example = ( + file_path.name, + float(lat_series.loc[idx]), + float(lon_series.loc[idx]), + float(lat_diff_series.loc[idx]), + float(lon_diff_series.loc[idx]), + ) + best_frames = [df_batch.loc[[idx]].copy()] + elif abs(distance - min_distance) <= 1e-12: + best_frames.append(df_batch.loc[[idx]].copy()) + + if not best_frames: + LOGGER.warning("No samples found in TVA dataset for latitude %.6f and longitude %.6f", latitude, longitude) + return None + + combined = pd.concat(best_frames, ignore_index=True).reset_index(drop=True) + output_path.parent.mkdir(parents=True, exist_ok=True) + combined.to_pickle(output_path) + + if closest_example: + LOGGER.info( + "Nearest sample sourced from %s at (lat=%.6f, lon=%.6f) [Δlat=%.3e, Δlon=%.3e]", + *closest_example, + ) + LOGGER.info("Saved %d nearest samples (min distance %.3e) for location to %s", combined.shape[0], min_distance, output_path) + return output_path + + +def run_model_inference( + model_path: Path, + dataset_file: Path, + output_dir: Path, + variable_list: Optional[Path], + model_config: Optional[Path], +) -> Path: + """Invoke the shared inference routine on a prepared dataset file.""" + output_dir.mkdir(parents=True, exist_ok=True) + LOGGER.info("Running inference for dataset %s", dataset_file) + results_dir = run_inference_all( + model_path=str(model_path), + data_paths=str(dataset_file.parent), + file_pattern=dataset_file.name, + output_dir=str(output_dir), + variable_list=str(variable_list) if variable_list else None, + model_config=str(model_config) if model_config else None, + scalers_dir=None, + use_training_config=True, + strict_loading=True, + debug_vars=False, + loader="auto", + mask_pft_with_gt=False, + ) + LOGGER.info("Inference outputs available in %s", results_dir) + return results_dir + + +def convert_predictions_to_netcdf( + predictions_dir: Path, + variable_map: Dict[str, List[str]], + output_path: Path, +) -> Path: + """Convert inference CSV outputs into a NetCDF file for restart updates.""" + if not predictions_dir.exists(): + raise FileNotFoundError(f"Predictions directory not found: {predictions_dir}") + LOGGER.info("Converting predictions in %s to NetCDF", predictions_dir) + ai_preds = load_ai_predictions(predictions_dir) + if "test_static_inverse" not in ai_preds: + raise RuntimeError( + "Prediction outputs are missing test_static_inverse.csv; cannot build NetCDF coordinates." + ) + + ds = create_netcdf_structure(ai_preds, variable_map, output_path) + add_scalar_variables(ds, ai_preds, variable_map) + add_pft_variables(ds, ai_preds, variable_map) + add_soil_variables(ds, ai_preds, variable_map) + + output_path.parent.mkdir(parents=True, exist_ok=True) + ds.to_netcdf(output_path) + LOGGER.info("NetCDF predictions written to %s", output_path) + return output_path + + +def collect_cnp_variables(variable_map: Dict[str, List[str]]) -> List[str]: + """Gather the list of variables that should be updated in the restart file.""" + soil_keys = variable_map.get("soil2d", []) + if not soil_keys: + # When list is empty try to auto-detect to maintain compatibility with restart script. + LOGGER.debug("No soil variables provided; relying on auto-detection during restart update.") + return sorted(set(variable_map.get("pft1d", []) + soil_keys)) + + +def update_restart_file( + restart_file: Path, + ai_predictions_nc: Path, + output_restart: Path, + cnp_variables: List[str], +) -> Path: + """Use the shared restart updater to write predictions into a new NetCDF file.""" + LOGGER.info("Updating restart file %s with predictions from %s", restart_file, ai_predictions_nc) + ds_ai, ds_model = load_datasets(ai_predictions_nc, restart_file) + try: + model_to_ai_mapping, variable_mapping = create_spatial_mapping(ds_ai, ds_model) + finally: + ds_ai.close() + ds_model.close() + + if not cnp_variables: + cnp_variables = auto_detect_variable_list(ai_predictions_nc) + + output_restart.parent.mkdir(parents=True, exist_ok=True) + create_updated_restart_file( + restart_file_path=restart_file, + output_path=output_restart, + ai_predictions_path=ai_predictions_nc, + cnp_io_variables=cnp_variables, + model_to_ai_mapping=model_to_ai_mapping, + variable_mapping=variable_mapping, + ) + LOGGER.info("Updated restart file saved to %s", output_restart) + return output_restart + + +def process_location( + location_row: pd.Series, + args: argparse.Namespace, + variable_map: Dict[str, List[str]], + cnp_variables: List[str], +) -> Optional[Path]: + """Execute the full workflow for a single row in the locations CSV.""" + latitude = float(location_row["latitude"]) + longitude = float(location_row["longitude"]) + location_label = f"{latitude:.4f}_{longitude:.4f}" + slug = slugify(location_label) + LOGGER.info("Processing location (lat=%.6f, lon=%.6f)", latitude, longitude) + + location_work_root = args.work_root / slug + dataset_file = location_work_root / f"{slug}_dataset.pkl" + predictions_root = location_work_root / "inference" + predictions_dir = predictions_root / "cnp_predictions" + predictions_nc = location_work_root / f"{slug}_ai_predictions.nc" + restart_output = args.output_root / f"{args.restart_file.stem}_{slug}.nc" + + if args.skip_existing and restart_output.exists(): + LOGGER.info("Skipping location %s because %s already exists", location_label, restart_output) + return restart_output + + dataset_path = extract_location_dataset( + dataset_root=args.dataset_root, + file_pattern=args.dataset_pattern, + latitude=latitude, + longitude=longitude, + output_path=dataset_file, + max_files=args.max_files, + ) + if dataset_path is None: + return None + + try: + df_preview = pd.read_pickle(dataset_path) + LOGGER.info( + " [Stage 1] Extracted %d rows (sample columns: %s)", + len(df_preview), + ", ".join(map(str, df_preview.columns[:5])), + ) + except Exception as exc: + LOGGER.warning(" [Stage 1] Unable to validate dataset %s: %s", dataset_path, exc) + + results_dir = run_model_inference( + model_path=args.model_path, + dataset_file=dataset_path, + output_dir=predictions_root, + variable_list=args.variable_list, + model_config=args.model_config, + ) + + predictions_dir = results_dir / "cnp_predictions" + if not predictions_dir.exists(): + raise FileNotFoundError(f"Predictions directory missing after inference: {predictions_dir}") + + csv_count = sum(1 for _ in predictions_dir.rglob("*.csv")) + LOGGER.info( + " [Stage 2] Inference artifacts ready at %s (CSV files: %d)", + predictions_dir, + csv_count, + ) + + predictions_nc = convert_predictions_to_netcdf(predictions_dir, variable_map, predictions_nc) + + try: + ds_preview = xr.open_dataset(predictions_nc) + dims_str = ", ".join(f"{k}={v}" for k, v in ds_preview.sizes.items()) + LOGGER.info( + " [Stage 3] NetCDF generated with dims [%s] and %d variables", + dims_str, + len(ds_preview.data_vars), + ) + except Exception as exc: + LOGGER.warning(" [Stage 3] Could not inspect NetCDF %s: %s", predictions_nc, exc) + finally: + try: + ds_preview.close() + except Exception: + pass + + updated_restart = update_restart_file( + restart_file=args.restart_file, + ai_predictions_nc=predictions_nc, + output_restart=restart_output, + cnp_variables=cnp_variables, + ) + if updated_restart.exists(): + size_mb = updated_restart.stat().st_size / (1024 * 1024) + LOGGER.info(" [Stage 4] Restart file created at %s (%.2f MB)", updated_restart, size_mb) + else: + LOGGER.warning(" [Stage 4] Restart file expected but not found: %s", updated_restart) + return updated_restart + + +def main() -> None: + args = parse_args() + configure_logging(args.log_level) + + LOGGER.info("Starting TVA workflow with locations file: %s", args.locations) + variable_map = determine_variable_map(args.variable_list) + cnp_variables = collect_cnp_variables(variable_map) + LOGGER.debug("Variable groups loaded: %s", variable_map) + + locations_df = read_locations(args.locations) + args.output_root.mkdir(parents=True, exist_ok=True) + args.work_root.mkdir(parents=True, exist_ok=True) + + results: List[Tuple[str, Optional[Path]]] = [] + for _, row in locations_df.iterrows(): + try: + updated_restart = process_location(row, args, variable_map, cnp_variables) + coord_label = f"{row['latitude']:.4f}_{row['longitude']:.4f}" + results.append((coord_label, updated_restart)) + except Exception as exc: # pragma: no cover - logging safety + LOGGER.exception("Failed to process location row %s: %s", row.to_dict(), exc) + coord_label = f"{row['latitude']:.4f}_{row['longitude']:.4f}" + results.append((coord_label, None)) + + LOGGER.info("Workflow complete. Summary of generated restart files:") + for label, path in results: + name = label or "" + if path is None: + LOGGER.warning(" %s: failed (no restart generated)", name) + else: + LOGGER.info(" %s: %s", name, path) + + +if __name__ == "__main__": + main() diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 1eeedeb..73b132d 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -117,3 +117,14 @@ python ../../scripts/update_restart_with_aipredictions.py \ Optional: `inference_map.py` can create statistics and interactive maps (currently not working well). +### Run Inference on the TVA Dataset +The TVA workflow allows you to perform site-specific inference using a trained model and generate updated restart files for targeted locations. Given one or more geographic coordinates, the script extracts all required variables from the TVA dataset, runs model inference to obtain predicted values, and produces a new restart file reflecting the AI-updated state. + +All codes are organized under the ./TVA_1_Sample folder. +Geographic coordinates are defined in locations.csv. +```bash +python ./TVA_1_Sample/run_workflow.py \ + --restart-file /path/to/20year_restart_file.nc \ + --model-path /path/to/trained_model_TVA.pt \ + --dataset-root /path/to/TVA_dataset +``` \ No newline at end of file From 29bbdfbbb65f5220094ef8ba49a313f3be9880bf Mon Sep 17 00:00:00 2001 From: Zhuowei Gu Date: Tue, 4 Nov 2025 14:24:47 -0500 Subject: [PATCH 34/51] fix inference bug --- data/data_loader_individual.py | 74 +++++++++++++++++++++++++--------- 1 file changed, 55 insertions(+), 19 deletions(-) diff --git a/data/data_loader_individual.py b/data/data_loader_individual.py index d7ba44d..d361fa2 100644 --- a/data/data_loader_individual.py +++ b/data/data_loader_individual.py @@ -517,9 +517,13 @@ def normalize_data_individual(self, transform_only: bool = False) -> Dict[str, A assert scalar_data.shape[1] == len(self.data_config.x_list_scalar_columns), 'Mismatch in scalar feature count!' assert variables_2d_soil.shape[1] == len(self.data_config.x_list_columns_2d), 'Mismatch in 2D soil feature count!' assert pft_param_data.shape[1] == len(self.data_config.pft_param_columns), 'Mismatch in PFT param feature count!' - assert y_scalar_data.shape[1] == len(self.data_config.y_list_scalar_columns), 'Mismatch in y_scalar feature count!' - assert y_pft_1d_data.shape[1] == len(self.data_config.y_list_columns_1d), 'Mismatch in y_pft_1d variable count!' - assert y_soil_2d.shape[1] == len(self.data_config.y_list_columns_2d), 'Mismatch in y_soil_2d feature count!' + # Only assert Y variables if they were normalized (training mode) + if y_scalar_data is not None: + assert y_scalar_data.shape[1] == len(self.data_config.y_list_scalar_columns), 'Mismatch in y_scalar feature count!' + if y_pft_1d_data is not None: + assert y_pft_1d_data.shape[1] == len(self.data_config.y_list_columns_1d), 'Mismatch in y_pft_1d variable count!' + if y_soil_2d is not None: + assert y_soil_2d.shape[1] == len(self.data_config.y_list_columns_2d), 'Mismatch in y_soil_2d feature count!' # Store all scalers self.scalers = { @@ -550,13 +554,19 @@ def normalize_data_individual(self, transform_only: bool = False) -> Dict[str, A 'scalar_data': scalar_data, 'variables_1d_pft': pft_1d_data, 'variables_2d_soil': variables_2d_soil, - 'y_scalar': y_scalar_data, - 'y_pft_1d': y_pft_1d_data, - 'y_soil_2d': y_soil_2d, 'water': water_tensor, 'y_water': y_water_tensor, 'scalers': self.scalers } + + # Only add Y variables if they were normalized (training mode) or exist (inference mode) + if y_scalar_data is not None: + ret['y_scalar'] = y_scalar_data + if y_pft_1d_data is not None: + ret['y_pft_1d'] = y_pft_1d_data + if y_soil_2d is not None: + ret['y_soil_2d'] = y_soil_2d + # Add per-sample PFT mask derived from raw PCT_NAT_PFT_1..16 (1 where >0, else 0) try: pct_cols = [f'PCT_NAT_PFT_{i}' for i in range(1, 17)] @@ -961,6 +971,13 @@ def _normalize_y_scalar_individual(self, transform_only: bool = False) -> Tuple[ y_scalar_columns = self.data_config.y_list_scalar_columns logger.info(f"Normalizing y_scalar data with columns: {y_scalar_columns}") + # For inference mode, skip if columns don't exist + if transform_only: + missing_cols = [col for col in y_scalar_columns if col not in self.df.columns] + if missing_cols: + logger.info(f"Inference mode: Y_scalar columns {missing_cols} missing in DataFrame. Skipping normalization.") + return None, None + for i, col in enumerate(y_scalar_columns): assert col in self.df.columns, f"y_scalar column '{col}' missing in DataFrame!" @@ -979,6 +996,14 @@ def _normalize_list_1d_individual(self, columns: List[str], transform_only: bool """Normalize 1D list data individually using IndividualScalerManager.""" logger.info(f"Normalizing 1D list data with columns: {columns}") + # For inference mode with Y variables, skip if columns don't exist + is_y = columns == self.data_config.y_list_columns_1d + if transform_only and is_y: + missing_cols = [col for col in columns if col not in self.df.columns] + if missing_cols: + logger.info(f"Inference mode: Y_pft_1d columns {missing_cols} missing in DataFrame. Skipping normalization.") + return None, None + for i, col in enumerate(columns): assert col in self.df.columns, f"1D column '{col}' missing in DataFrame!" @@ -1139,6 +1164,14 @@ def _normalize_list_2d_individual(self, columns: List[str], transform_only: bool """Normalize 2D list data individually using IndividualScalerManager.""" logger.info(f"Normalizing 2D list data with columns: {columns}") + # For inference mode with Y variables, skip if columns don't exist + is_y = columns == self.data_config.y_list_columns_2d + if transform_only and is_y: + missing_cols = [col for col in columns if col not in self.df.columns] + if missing_cols: + logger.info(f"Inference mode: Y_soil_2d columns {missing_cols} missing in DataFrame. Skipping normalization.") + return None, None + for i, col in enumerate(columns): assert col in self.df.columns, f"2D column '{col}' missing in DataFrame!" @@ -1525,25 +1558,28 @@ def split_data(self, normalized_data: Dict[str, Any]) -> Dict[str, Any]: train_data['scalar'] = train_list_scalar test_data['scalar'] = test_list_scalar - # Split y_scalar (target) - y_scalar = normalized_data['y_scalar'] - train_data['y_scalar'] = y_scalar[:train_size] - test_data['y_scalar'] = y_scalar[train_size:] + # 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] + 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] test_data['variables_1d_pft'] = variables_1d_pft[train_size:] - # Split y_pft_1d (target) - y_pft_1d = normalized_data['y_pft_1d'] - train_data['y_pft_1d'] = y_pft_1d[:train_size] - test_data['y_pft_1d'] = y_pft_1d[train_size:] - - # Split y_soil_2d (target) - y_soil_2d = normalized_data['y_soil_2d'] - train_data['y_soil_2d'] = y_soil_2d[:train_size] - test_data['y_soil_2d'] = y_soil_2d[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] + 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] + test_data['y_soil_2d'] = y_soil_2d[train_size:] # Split variables_2d_soil (input) variables_2d_soil = normalized_data['variables_2d_soil'] From b667fca3d65f600e9cd30ad85b889bd7e5b513fd Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Wed, 5 Nov 2025 14:22:18 -0800 Subject: [PATCH 35/51] update: revised pipeline doc, added CNP_IO_updated9_dev_gao.txt and run_finetuning.py --- CNP_IO_updated9_dev_gao.txt | 82 +++++ docs/CNP_pipeline_runbook.md | 26 +- scripts/run_finetuning.py | 680 +++++++++++++++++++++++++++++++++++ 3 files changed, 787 insertions(+), 1 deletion(-) create mode 100644 CNP_IO_updated9_dev_gao.txt create mode 100644 scripts/run_finetuning.py diff --git a/CNP_IO_updated9_dev_gao.txt b/CNP_IO_updated9_dev_gao.txt new file mode 100644 index 0000000..7ccc90b --- /dev/null +++ b/CNP_IO_updated9_dev_gao.txt @@ -0,0 +1,82 @@ +AI_PREDICTIONS_DEFAULT: ./comparison_results/ai_predictions_for_plotting.nc +MODEL_DEFAULT: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc +COMPARISON_OUTPUT_DIR: ./ai_model_comparison_plots +CSV_PREDICTIONS_DEFAULT: ./cnp_inference_entire_dataset/cnp_predictions +AI_RESTART_DEFAULT: ./updated_restart_CNP_IO_updated9_dev_20250408_trendytest_ICB1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc +FALLBACK_DATA_DIR: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/ +FALLBACK_REFERENCE_FILENAME: 20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc +fine_tuning_dataset_path: /global/cfs/cdirs/m4814/daweigao/14_Code/TVA_training_dataset_all/ +pretrained_model_path: /global/cfs/cdirs/m4814/daweigao/15_code_Landsim/LandSim/cnp_results/run_20251030_192921/cnp_predictions/model.pth +output_finetuned_model_dir: ./cnp_results/ +finetuned_model_filename: model.pth +num_epochs: 150 + +TRENDY1_PATH: +TRENDY05_PATH = +DATA_PATHS: /global/cfs/cdirs/m4814/daweigao/14_Code/TVA_training_dataset_all/ +FILE_PATTERN: enhanced_monthly_training_data_batch_*.pkl + +TVA4KM_PATH: /global/cfs/cdirs/m4814/daweigao/14_Code/TVA_training_dataset_all/ +# 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, NPP, AR, HR + +1D PFT VARIABLES (41 variables): + +• deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage +• deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage + +• leafc, leafc_storage, leafn, leafn_storage, leafp, leafp_storage +• frootc, frootc_storage, frootn, frootn_storage, frootp, frootp_storage + +• livestemc, livestemc_storage, livestemn, livestemn_storage, livestemp, livestemp_storage, +• livecrootc, livecrootc_storage, livecrootn, livecrootn_storage, livecrootp, livecrootp_storage + +• cpool, npool, ppool + +• tlai, totvegc + +2D VARIABLES (layered - 25 variables): + +• cwdc_vr, cwdn_vr, cwdp_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 +• soil3c_vr, soil3n_vr, soil3p_vr +• soil4c_vr, soil4n_vr, soil4p_vr + +• labilep_vr , occlp_vr, primp_vr, secondp_vr diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 73b132d..3068ec1 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -16,6 +16,30 @@ Notes: - This launches training in the background and redirects logs to stdout/stderr. - The run output directory will be created under `cnp_results/run_YYYYMMDD_HHMMSS`. +### 2a) Fine-tune a pretrained model (optional) +If you already have a trained checkpoint and want to continue training on a TVA-style dataset, use the fine-tuning helper. Populate the necessary paths in your CNP_IO file (e.g. `CNP_IO_updated9_dev_gao.txt`): + +``` +fine_tuning_dataset_path: /path/to/TVA_dataset_root +pretrained_model_path: /path/to/base_run/cnp_predictions/model.pth +output_finetuned_model_dir: ./cnp_results/ +finetuned_model_filename: model.pth +num_epochs: 150 +``` + +Then launch fine-tuning: + +```bash +python scripts/run_finetuning.py +``` + +What happens: +- The script reuses the original training pipeline (same preprocessing, normalization, ModelTrainer). +- `fine_tuning_dataset_path` + `FILE_PATTERN` point to the TVA batches to load. +- `pretrained_model_path` is loaded before training begins so weights pick up where the base run ended. +- Outputs are written to `cnp_results/run_YYYYMMDD_HHMMSS/` (same layout as full training). + Check `cnp_predictions/model.pth`, `cnp_config.json`, `cnp_metrics.json`, and the refreshed prediction CSVs. + ### 3) Navigate to the run directory ```bash @@ -127,4 +151,4 @@ python ./TVA_1_Sample/run_workflow.py \ --restart-file /path/to/20year_restart_file.nc \ --model-path /path/to/trained_model_TVA.pt \ --dataset-root /path/to/TVA_dataset -``` \ No newline at end of file +``` diff --git a/scripts/run_finetuning.py b/scripts/run_finetuning.py new file mode 100644 index 0000000..570451e --- /dev/null +++ b/scripts/run_finetuning.py @@ -0,0 +1,680 @@ +#!/usr/bin/env python3 +""" +Fine-tuning driver for the CNP model. + +This script performs model fine-tuning using a configuration-only workflow: +all runtime parameters (paths, hyperparameters, dataset descriptions) are +expected to originate from the `CNP_IO_updated9_dev_gao.txt` configuration +file supplied by the user. The implementation reuses the existing training +infrastructure (data loaders, trainer, model definition) to stay aligned with +the primary training entry point. + +Usage: + python scripts/run_finetuning.py [--config PATH] [--normalization MODE] + +Key responsibilities: + * Parse the configuration text file into structured overrides. + * Load TVA datasets using the same preprocessing pipeline as training. + * Restore a pretrained checkpoint and continue training. + * Persist the fine-tuned model along with metrics and configuration + metadata in an output directory specified by the configuration file. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Iterable, Optional, Tuple, Union + +import torch + +# Ensure the project root (LandSim) is on sys.path so internal modules resolve. +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from config.training_config import ( + TrainingConfigManager, + get_cnp_combined_config, + parse_cnp_io_list, +) +from data.data_loader_individual import DataLoaderIndividual +from models.cnp_combined_model import CNPCombinedModel +from scripts.run_inference_all import verify_locations +from training.trainer import ModelTrainer + +# --------------------------------------------------------------------------- # +# Configuration parsing helpers +# --------------------------------------------------------------------------- # + + +def _coerce_scalar(value: str) -> Any: + """Best-effort conversion of textual configuration values to Python types.""" + if value is None: + return None + text = value.strip() + if text == "": + return "" + + lowered = text.lower() + if lowered in {"true", "yes", "on"}: + return True + if lowered in {"false", "no", "off"}: + return False + + # Numeric conversion (int preferred when possible). + try: + if "." not in text: + return int(text) + return float(text) + except ValueError: + pass + + # Comma-separated lists (e.g., multi-path values). + if "," in text: + parts = [segment.strip() for segment in text.split(",")] + return [segment for segment in parts if segment] + + return text + + +def _parse_key_value_pairs(config_path: Path) -> Dict[str, Any]: + """Parse simple key/value lines from the configuration text file.""" + results: Dict[str, Any] = {} + with config_path.open("r") as handle: + for raw_line in handle: + stripped = raw_line.strip() + if not stripped or stripped.startswith("#") or stripped.startswith("//"): + continue + # Ignore bullet lists; they are handled by parse_cnp_io_list. + if stripped.startswith("•"): + continue + # Support both ":" and "=" delimiters. + delimiter = ":" if ":" in stripped else ("=" if "=" in stripped else None) + if delimiter is None: + continue + key, raw_value = stripped.split(delimiter, 1) + key = key.strip() + value = _coerce_scalar(raw_value) + if key: + results[key.lower()] = value + return results + + +def load_configuration(config_path: Path) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Load configuration overrides from the text file. + + Returns + ------- + tuple + - combined (dict): Lowercase key dictionary combining structured + groups returned by `parse_cnp_io_list` and simple key/value pairs. + - parsed_groups (dict): Full output from `parse_cnp_io_list` containing + variable lists and dataset metadata. + """ + parsed_groups = parse_cnp_io_list(str(config_path)) + simple_pairs = _parse_key_value_pairs(config_path) + + combined: Dict[str, Any] = {} + for key, value in parsed_groups.items(): + combined[key.lower()] = value + combined.update(simple_pairs) + + return combined, parsed_groups + + +def require_config( + config: Dict[str, Any], + *keys: str, + required: bool = True, + default: Any = None, +) -> Any: + """ + Retrieve the first available configuration value among provided keys. + + Parameters + ---------- + config : dict + Configuration dictionary with lowercase keys. + keys : tuple[str] + Key candidates to attempt, in priority order. + required : bool + Whether to raise an error when the keys are not found. + default : Any + Fallback value returned when `required=False` and nothing is present. + """ + for key in keys: + candidate = key.lower() + if candidate in config: + value = config[candidate] + if value is None or value == "": + continue + return value + if required: + names = ", ".join(keys) + raise KeyError(f"Missing required configuration entry. Tried keys: {names}") + return default + + +def resolve_path(base_file: Path, raw_path: Union[str, Path]) -> Path: + """Resolve paths relative to the configuration file directory.""" + if isinstance(raw_path, Path): + candidate = raw_path + else: + expanded = os.path.expandvars(str(raw_path)) + candidate = Path(expanded) + candidate = candidate.expanduser() + if candidate.is_absolute(): + return candidate + return (base_file.parent / candidate).resolve() + + +def materialize_directory(path: Path) -> Path: + """Ensure a directory exists and return the path.""" + path.mkdir(parents=True, exist_ok=True) + return path + + +def attach_file_logger(run_dir: Path) -> None: + """Attach a file handler pointing to the run directory log file.""" + root_logger = logging.getLogger() + log_path = Path(run_dir) / "finetune.log" + for handler in root_logger.handlers: + if isinstance(handler, logging.FileHandler): + try: + if Path(getattr(handler, "baseFilename", "")) == log_path: + return + except TypeError: + continue + file_handler = logging.FileHandler(log_path) + file_handler.setLevel(root_logger.level) + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + file_handler.setFormatter(formatter) + root_logger.addHandler(file_handler) + root_logger.info("Log file created at %s", log_path) + + +# --------------------------------------------------------------------------- # +# Configuration application helpers +# --------------------------------------------------------------------------- # + + +def apply_variable_lists( + manager: TrainingConfigManager, + parsed_groups: Dict[str, Any], +) -> None: + """Populate the `DataConfig` lists from the parsed CNP IO specification.""" + data_cfg = manager.data_config + + # Optional overrides only trigger when a non-empty list is provided. + def _set_if_present(attr: str, values: Iterable[str]) -> None: + items = list(values) if values else [] + if items: + setattr(data_cfg, attr, items) + + _set_if_present("time_series_columns", parsed_groups.get("time_series_variables")) + _set_if_present("static_columns", parsed_groups.get("surface_properties")) + _set_if_present("pft_param_columns", parsed_groups.get("pft_parameters")) + _set_if_present("x_list_scalar_columns", parsed_groups.get("scalar_variables")) + _set_if_present("x_list_columns_1d", parsed_groups.get("pft_1d_variables")) + _set_if_present("x_list_columns_2d", parsed_groups.get("variables_2d_soil")) + + # Derive target variable lists when not explicitly supplied. + if parsed_groups.get("scalar_variables"): + data_cfg.y_list_scalar_columns = [ + f"Y_{name}" if not name.startswith("Y_") else name + for name in parsed_groups["scalar_variables"] + ] + if parsed_groups.get("pft_1d_variables"): + data_cfg.y_list_columns_1d = [ + f"Y_{name}" if not name.startswith("Y_") else name + for name in parsed_groups["pft_1d_variables"] + ] + if parsed_groups.get("variables_2d_soil"): + data_cfg.y_list_columns_2d = [ + f"Y_{name}" if not name.startswith("Y_") else name + for name in parsed_groups["variables_2d_soil"] + ] + + longitudes = parsed_groups.get("longitudes_to_drop") or [] + if longitudes: + try: + data_cfg.longitudes_to_drop = [float(val) for val in longitudes] + except (TypeError, ValueError): + logging.getLogger(__name__).warning( + "Failed to parse longitude filtering list: %s", longitudes + ) + + +def apply_training_overrides( + manager: TrainingConfigManager, + config: Dict[str, Any], + run_dir: Path, + model_filename: str, +) -> None: + """Update training configuration fields based on key/value overrides.""" + train_cfg = manager.training_config + data_cfg = manager.data_config + + # Training hyperparameters (optional). + for key in ("num_epochs", "batch_size", "learning_rate", "patience"): + if key in config: + try: + value = float(config[key]) if key == "learning_rate" else int(config[key]) + setattr(train_cfg, key, value if key == "learning_rate" else int(value)) + except (TypeError, ValueError): + logging.getLogger(__name__).warning( + "Skipping invalid value for %s: %r", key, config[key] + ) + + # Train/test split ratio. + if "train_split" in config: + try: + data_cfg.train_split = float(config["train_split"]) + except (TypeError, ValueError): + logging.getLogger(__name__).warning( + "Invalid train_split override: %r", config["train_split"] + ) + + # Max files constraint for rapid experimentation. + if "max_files" in config: + try: + data_cfg.max_files = int(config["max_files"]) + except (TypeError, ValueError): + logging.getLogger(__name__).warning( + "Invalid max_files override: %r", config["max_files"] + ) + + # Enable or disable early stopping if requested. + if "use_early_stopping" in config: + train_cfg.use_early_stopping = bool(config["use_early_stopping"]) + + # Device selection honoring config preference when specified. + if "device" in config: + train_cfg.device = str(config["device"]) + else: + train_cfg.device = "cuda" if torch.cuda.is_available() else "cpu" + + # Paths for artifacts relative to run directory. + train_cfg.model_save_path = str(run_dir / model_filename) + train_cfg.losses_save_path = str(run_dir / "finetune_losses.csv") + train_cfg.predictions_dir = str(run_dir / "cnp_predictions") + train_cfg.save_model = True + train_cfg.save_predictions = True + + # Disable GPU monitoring by default for fine-tuning runs unless explicitly enabled. + train_cfg.log_gpu_memory = bool(config.get("log_gpu_memory", False)) + train_cfg.log_gpu_utilization = bool(config.get("log_gpu_utilization", False)) + + # Optional masking configuration. + if "mask_absent_pfts" in config: + train_cfg.mask_absent_pfts = bool(config["mask_absent_pfts"]) + + +def load_training_config_from_checkpoint(model_path: Path) -> Optional[Dict[str, Any]]: + """Discover and load `cnp_config.json` located near the checkpoint.""" + for directory in [model_path.parent] + list(model_path.parents): + candidate = directory / "cnp_config.json" + if candidate.exists(): + try: + with candidate.open("r") as handle: + return json.load(handle) + except (OSError, json.JSONDecodeError) as err: + logging.getLogger(__name__).warning( + "Failed to parse %s: %s", candidate, err + ) + break + return None + + +def apply_model_config_from_json( + manager: TrainingConfigManager, + config_json: Dict[str, Any], +) -> None: + """Bring forward model architecture overrides stored in training JSON.""" + model_overrides = config_json.get("model_config") + if not isinstance(model_overrides, dict): + return + model_cfg = manager.model_config + for key, value in model_overrides.items(): + if hasattr(model_cfg, key): + setattr(model_cfg, key, value) + + +def apply_data_info_overrides( + manager: TrainingConfigManager, + config_json: Dict[str, Any], +) -> None: + """Sync data configuration lists using historical `data_info` metadata.""" + data_info = config_json.get("data_info") + if not isinstance(data_info, dict): + return + + data_cfg = manager.data_config + + mappings = { + "time_series_columns": "time_series_columns", + "static_columns": "static_columns", + "pft_param_columns": "pft_param_columns", + "x_list_scalar_columns": "x_list_scalar_columns", + "y_list_scalar_columns": "y_list_scalar_columns", + "variables_1d_pft": "x_list_columns_1d", + "y_list_columns_1d": "y_list_columns_1d", + "x_list_columns_2d": "x_list_columns_2d", + "y_list_columns_2d": "y_list_columns_2d", + } + for source_key, target_attr in mappings.items(): + values = data_info.get(source_key) + if isinstance(values, list) and values: + setattr(data_cfg, target_attr, values) + + +# --------------------------------------------------------------------------- # +# Fine-tuning workflow +# --------------------------------------------------------------------------- # + + +def build_config_manager( + args: argparse.Namespace, + config: Dict[str, Any], + parsed_groups: Dict[str, Any], +) -> Tuple[TrainingConfigManager, Path, Path, str, str]: + """ + Prepare the configuration manager, dataset paths, and output locations. + + Returns + ------- + tuple + manager, dataset_path, model_path, file_pattern, model_filename + """ + cfg_file = Path(args.config).resolve() + dataset_path_raw = require_config( + config, + "fine_tuning_dataset_path", + "tva_dataset_path", + "tva4km_path", + "data_paths", + ) + if isinstance(dataset_path_raw, list): + if not dataset_path_raw: + raise ValueError("Dataset path list from configuration is empty.") + dataset_path_raw = dataset_path_raw[0] + dataset_path = resolve_path(cfg_file, dataset_path_raw) + if not dataset_path.exists(): + raise FileNotFoundError(f"Dataset directory not found: {dataset_path}") + + file_pattern_raw = require_config( + config, + "tva_dataset_pattern", + "tva4km_file_pattern", + "file_pattern", + ) + if isinstance(file_pattern_raw, list): + if not file_pattern_raw: + raise ValueError("File pattern list from configuration is empty.") + file_pattern_raw = file_pattern_raw[0] + file_pattern = str(file_pattern_raw) + model_path_raw = require_config( + config, + "pretrained_model_path", + "model_default", + ) + if isinstance(model_path_raw, list): + if not model_path_raw: + raise ValueError("Model path list from configuration is empty.") + model_path_raw = model_path_raw[0] + model_path = resolve_path(cfg_file, model_path_raw) + if not model_path.exists(): + raise FileNotFoundError(f"Pretrained model checkpoint not found: {model_path}") + + output_dir_raw = require_config( + config, + "output_finetuned_model_dir", + "fine_tuning_output_dir", + "finetune_output_dir", + ) + output_root = materialize_directory(resolve_path(cfg_file, output_dir_raw)) + run_name = datetime.now().strftime("finetune_%Y%m%d_%H%M%S") + run_directory = materialize_directory(output_root / run_name) + + model_filename = str(require_config( + config, + "finetuned_model_filename", + "finetuned_checkpoint_name", + required=False, + default="finetuned_model.pth", + )) + + # Build base configuration using the existing helper so that defaults + # remain consistent with primary training runs. + manager = get_cnp_combined_config( + use_trendy1=False, + use_trendy05=False, + use_tva4km=True, + max_files=None, + include_water=False, + variable_list_path=str(cfg_file), + ) + + # Override dataset paths/patterns explicitly. + manager.update_data_config( + data_paths=[str(dataset_path)], + file_pattern=str(file_pattern), + dataset_file_patterns={str(dataset_path): str(file_pattern)}, + ) + + # Apply variable groups from configuration file. + apply_variable_lists(manager, parsed_groups) + + # Attach artifact locations and optional hyperparameter overrides. + apply_training_overrides(manager, config, run_directory, model_filename) + + return manager, dataset_path, model_path, run_directory, model_filename + + +def configure_normalization_mode( + data_loader: DataLoaderIndividual, + config: Dict[str, Any], + explicit_mode: Optional[str], +) -> Tuple[Dict[str, Any], str]: + """Select normalization routine based on CLI override or configuration.""" + mode = ( + explicit_mode + or str(config.get("normalization_mode", "individual")).lower() + ) + logger = logging.getLogger(__name__) + logger.info("Normalization mode selected: %s", mode) + + if mode == "group": + return data_loader.normalize_data(), "group" + if mode == "hybrid": + # Default hybrid settings: mimic training script behaviour when unspecified. + return ( + data_loader.normalize_data_hybrid( + use_individual_for=[ + "scalar", + "y_scalar", + "pft_1d", + "y_pft_1d", + "soil_2d", + "y_soil_2d", + ], + group_soil_vars=["sminn_vr", "smin_no3_vr", "smin_nh4_vr"], + ), + "hybrid", + ) + # Fallback to individual normalization. + return data_loader.normalize_data_individual(), "individual" + + +def load_pretrained_weights(model: torch.nn.Module, checkpoint_path: Path) -> None: + """Load weights from the provided checkpoint path into the model.""" + logger = logging.getLogger(__name__) + checkpoint = torch.load(checkpoint_path, map_location="cpu") + if isinstance(checkpoint, dict): + if "model_state_dict" in checkpoint: + state_dict = checkpoint["model_state_dict"] + else: + state_dict = checkpoint + else: + raise ValueError(f"Unsupported checkpoint format: {type(checkpoint)}") + model.load_state_dict(state_dict, strict=False) + logger.info("Loaded pretrained weights from %s", checkpoint_path) + + +def fine_tune(args: argparse.Namespace) -> Dict[str, Any]: + """Entry point driving the fine-tuning pipeline.""" + config_map, parsed_groups = load_configuration(Path(args.config)) + manager, dataset_path, model_path, run_dir, model_filename = build_config_manager( + args, + config_map, + parsed_groups, + ) + + logger = logging.getLogger(__name__) + logger.info("Dataset directory: %s", dataset_path) + logger.info("Checkpoint to fine-tune: %s", model_path) + logger.info("Run directory: %s", run_dir) + attach_file_logger(run_dir) + + # Load historical configuration from the checkpoint when available. + prior_config = load_training_config_from_checkpoint(model_path) + if prior_config: + apply_model_config_from_json(manager, prior_config) + apply_data_info_overrides(manager, prior_config) + + # Set up the data loader using the populated configuration. + data_loader = DataLoaderIndividual( + manager.data_config, + manager.preprocessing_config, + ) + + raw_df = data_loader.load_data() + if hasattr(raw_df, "head"): + verify_locations(raw_df, "Loaded dataset") + + data_loader.preprocess_data() + normalized, normalization_mode = configure_normalization_mode( + data_loader, + config_map, + args.normalization, + ) + + split = data_loader.split_data(normalized) + data_info = data_loader.get_data_info() + + # Prepare model and trainer. + include_water = bool(config_map.get("include_water", False)) + model = CNPCombinedModel( + manager.model_config, + data_info, + include_water=include_water, + use_learnable_loss_weights=manager.training_config.use_learnable_loss_weights, + ) + load_pretrained_weights(model, model_path) + + predictions_dir = Path(manager.training_config.predictions_dir) + predictions_dir.mkdir(parents=True, exist_ok=True) + + scalers = normalized.get("scalers") + if scalers is None: + raise KeyError("Normalized data did not include scalers; cannot proceed.") + + trainer = ModelTrainer( + manager.training_config, + model, + split["train"], + split["test"], + scalers, + data_info, + ) + + results = trainer.run_training_pipeline() + + # Persist metrics and configuration summary alongside the checkpoint. + metrics = results.get("metrics", {}) + (Path(manager.training_config.losses_save_path).parent).mkdir(parents=True, exist_ok=True) + with (Path(run_dir) / "cnp_metrics.json").open("w") as handle: + json.dump(metrics, handle, indent=2) + + config_payload = { + "run_timestamp": datetime.now().isoformat(timespec="seconds"), + "dataset_root": str(dataset_path), + "normalization_mode": normalization_mode, + "include_water": include_water, + "source_checkpoint": str(model_path), + "model_save_path": manager.training_config.model_save_path, + "data_info": data_info, + "model_config": vars(manager.model_config), + "training_config": vars(manager.training_config), + } + with (Path(run_dir) / "cnp_config.json").open("w") as handle: + json.dump(config_payload, handle, indent=2) + + logger.info("Fine-tuning completed successfully. Artifacts saved to %s", run_dir) + + return { + "metrics": metrics, + "run_directory": str(run_dir), + "model_path": manager.training_config.model_save_path, + } + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + default_config = PROJECT_ROOT / "CNP_IO_updated9_dev_gao.txt" + parser = argparse.ArgumentParser( + description="Fine-tune the CNP model using TVA datasets." + ) + parser.add_argument( + "--config", + type=Path, + default=default_config, + help="Path to the configuration text file.", + ) + parser.add_argument( + "--normalization", + choices=("group", "individual", "hybrid"), + help="Override normalization strategy (otherwise use config setting).", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=("DEBUG", "INFO", "WARNING", "ERROR"), + help="Logging verbosity for the run.", + ) + return parser.parse_args() + + +def configure_root_logger(level: str) -> None: + """Configure root logger according to requested verbosity.""" + numeric_level = getattr(logging, level.upper(), logging.INFO) + logging.basicConfig( + level=numeric_level, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + stream=sys.stdout, + ) + + +def main() -> None: + """CLI entry point.""" + args = parse_args() + configure_root_logger(args.log_level) + try: + results = fine_tune(args) + logging.getLogger(__name__).info("Fine-tune metrics: %s", results["metrics"]) + except Exception as exc: # noqa: BLE001 + logging.getLogger(__name__).exception("Fine-tuning failed: %s", exc) + sys.exit(1) + + +if __name__ == "__main__": + main() From 5063d803f4c291c589377e770ddac979586502db Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Wed, 5 Nov 2025 21:47:00 -0500 Subject: [PATCH 36/51] add personal CNP_IO_updated9_dev list --- CNP_IO_updated9_dev.txt | 20 +++++------ CNP_IO_updated9_dev_dw.txt | 74 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 11 deletions(-) create mode 100644 CNP_IO_updated9_dev_dw.txt diff --git a/CNP_IO_updated9_dev.txt b/CNP_IO_updated9_dev.txt index 0bec28c..1a2837b 100644 --- a/CNP_IO_updated9_dev.txt +++ b/CNP_IO_updated9_dev.txt @@ -1,14 +1,12 @@ -AI_PREDICTIONS_DEFAULT: ./comparison_results/ai_predictions_for_plotting.nc -MODEL_DEFAULT: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc -COMPARISON_OUTPUT_DIR: ./ai_model_comparison_plots -CSV_PREDICTIONS_DEFAULT: ./cnp_inference_entire_dataset/cnp_predictions -AI_RESTART_DEFAULT: ./updated_restart_CNP_IO_updated9_dev_20250408_trendytest_ICB1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc -FALLBACK_DATA_DIR: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/ -FALLBACK_REFERENCE_FILENAME: 20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc - -TRENDY1_PATH: /global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree -TRENDY05_PATH = -DATA_PATHS: +# 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 # Per-dataset patterns (overrides FILE_PATTERN for that path only) diff --git a/CNP_IO_updated9_dev_dw.txt b/CNP_IO_updated9_dev_dw.txt new file mode 100644 index 0000000..1a2837b --- /dev/null +++ b/CNP_IO_updated9_dev_dw.txt @@ -0,0 +1,74 @@ +# 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 + +# 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, NPP, AR, HR + +1D PFT VARIABLES (41 variables): + +• deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage +• deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage + +• leafc, leafc_storage, leafn, leafn_storage, leafp, leafp_storage +• frootc, frootc_storage, frootn, frootn_storage, frootp, frootp_storage + +• livestemc, livestemc_storage, livestemn, livestemn_storage, livestemp, livestemp_storage, +• livecrootc, livecrootc_storage, livecrootn, livecrootn_storage, livecrootp, livecrootp_storage + +• cpool, npool, ppool + +• tlai, totvegc + +2D VARIABLES (layered - 25 variables): + +• cwdc_vr, cwdn_vr, cwdp_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 +• soil3c_vr, soil3n_vr, soil3p_vr +• soil4c_vr, soil4n_vr, soil4p_vr + +• labilep_vr , occlp_vr, primp_vr, secondp_vr From ea6ee758a838dcf00dd34e999d78298e23ec1c5e Mon Sep 17 00:00:00 2001 From: Xinheng Ding Date: Thu, 13 Nov 2025 13:51:10 -0600 Subject: [PATCH 37/51] fix pft param normalization --- data/data_loader_individual.py | 125 ++++++++++++++------------------- 1 file changed, 54 insertions(+), 71 deletions(-) diff --git a/data/data_loader_individual.py b/data/data_loader_individual.py index d361fa2..d9cd8f6 100644 --- a/data/data_loader_individual.py +++ b/data/data_loader_individual.py @@ -152,7 +152,7 @@ def load_data(self) -> pd.DataFrame: logger.info(f"Successfully loaded {len(self.df)} samples") return self.df - + def preprocess_data(self): """Preprocess the loaded data.""" logger.info("Starting data preprocessing...") @@ -547,6 +547,22 @@ def normalize_data_individual(self, transform_only: bool = False) -> Dict[str, A 'individual_y_soil_2d': self.individual_scalers['y_soil_2d'], } + # debug + # logger.info(" Debug data stats after normalization:") + # _print_stats("time_series_data", time_series_data) + # _print_stats("static_data", static_data) + # _print_stats("pft_param_data", pft_param_data) + # _print_stats("scalar_data", scalar_data) + # _print_stats("variables_1d_pft", pft_1d_data) + # _print_stats("variables_2d_soil", variables_2d_soil) + # _print_stats("y_scalar", y_scalar_data) + # _print_stats("y_pft_1d", y_pft_1d_data) + # _print_stats("y_soil_2d", y_soil_2d) + # _print_stats("water", water_tensor) + # _print_stats("y_water", y_water_tensor) + # logger.info(" Debug checking soil 2D stats after normalization:") + # _check_soil_2d_stats("variables_2d_soil", variables_2d_soil) + # _check_soil_2d_stats("y_soil_2d", y_soil_2d) ret = { 'time_series_data': time_series_data, 'static_data': static_data, @@ -1316,11 +1332,17 @@ def _normalize_pft_param(self) -> Tuple[torch.Tensor, Any]: param_matrix.append(row_matrix) param_matrix = np.stack(param_matrix, axis=0) # [batch, 44, 17] assert param_matrix.shape[1:] == (num_params, num_pfts), f"pft_param_data shape {param_matrix.shape} does not match [batch, 44, 17]" - # Flatten for normalization - flat_param_matrix = param_matrix.reshape(param_matrix.shape[0], -1) + scaler = self._get_scaler(self.preprocessing_config.list_1d_normalization) - flat_param_matrix_norm = scaler.fit_transform(flat_param_matrix) - param_matrix_norm = flat_param_matrix_norm.reshape(param_matrix.shape) + param_matrix_norm = np.empty_like(param_matrix) + + for i in range(num_params): + X = param_matrix[:, i, :] + X_norm = scaler.fit_transform(X.T) + if np.all(X_norm == 0): + logger.warning(f"{pft_param_columns[i]}: After normalization, the value is all zeros\n") + param_matrix_norm[:, i, :] = X_norm.T + pft_param_data = torch.tensor(param_matrix_norm, dtype=self.preprocessing_config.data_type) return pft_param_data, scaler @@ -1634,72 +1656,6 @@ def split_data(self, normalized_data: Dict[str, Any]) -> Dict[str, Any]: 'test_size': test_size } - def _normalize_list_1d(self, columns: List[str]) -> Tuple[torch.Tensor, Any]: - """Normalize 1D list data in the order defined by columns.""" - logger.info(f"Normalizing 1D list data with columns: {columns}") - for i, col in enumerate(columns): - assert col in self.df.columns, f"1D column '{col}' missing in DataFrame!" - col_data = [np.vstack(self.df[col].values) for col in columns] - data = np.stack(col_data, axis=1) # shape: (samples, features, length) - n_samples, n_features, n_length = data.shape - data_reshaped = data.reshape(n_samples, -1) - scaler = self._get_scaler(self.preprocessing_config.list_1d_normalization) - data_normalized = scaler.fit_transform(data_reshaped) - data_normalized = data_normalized.reshape(n_samples, n_features, n_length) - return torch.tensor(data_normalized, dtype=self.preprocessing_config.data_type), scaler - - def _normalize_list_2d(self, columns: List[str]) -> Tuple[torch.Tensor, Any]: - """Normalize 2D list data in the order defined by columns (group path). - - Extract FIRST GROUP (column) -> top 10 layers for consistency with inspector/individual. - """ - logger.info(f"Normalizing 2D list data with columns: {columns}") - for i, col in enumerate(columns): - assert col in self.df.columns, f"2D column '{col}' missing in DataFrame!" - - # Extract first column and top 10 layers directly for consistent shapes - col_data = [] - for col in columns: - values = self.df[col].values - standardized_samples = [] - - for val in values: - try: - if isinstance(val, (list, tuple)) and len(val) > 0 and isinstance(val[0], (list, tuple, np.ndarray)): - arr = np.array(val[0], dtype=float).reshape(1, -1) - else: - arr = np.array(val, dtype=object) - if getattr(arr, 'ndim', 1) == 2 and arr.shape[0] >= 1: - arr = np.array(arr[0, :], dtype=float).reshape(1, -1) - elif getattr(arr, 'ndim', 1) == 1 and len(arr) >= 1 and not isinstance(arr[0], (list, tuple, np.ndarray)): - arr = np.array(arr, dtype=float).reshape(1, -1) - else: - arr = None - if arr is None: - standardized_samples.append(np.zeros((1, 10))) - else: - # take top 10 layers - out = np.zeros((1, 10), dtype=float) - take = min(10, arr.shape[1]) - if take > 0: - out[:, :take] = arr[:, :take] - standardized_samples.append(out) - except Exception: - standardized_samples.append(np.zeros((1, 10))) - - col_data.append(np.stack(standardized_samples)) - - data = np.stack(col_data, axis=1) # shape: (samples, features, 1, 10) - - # Data is already in the correct shape: (samples, variables, 1, 10) - # No need for additional extraction since we did it during loading - - n_samples, n_features, n_rows, n_cols = data.shape - data_reshaped = data.reshape(n_samples, -1) - scaler = self._get_scaler(self.preprocessing_config.list_2d_normalization) - data_normalized = scaler.fit_transform(data_reshaped) - data_normalized = data_normalized.reshape(n_samples, n_features, n_rows, n_cols) - return torch.tensor(data_normalized, dtype=self.preprocessing_config.data_type), scaler def get_data_info(self) -> Dict[str, Any]: """Get information about the loaded data for configuration and logging.""" @@ -1717,3 +1673,30 @@ def get_data_info(self) -> Dict[str, Any]: 'data_shape': self.df.shape if hasattr(self, 'df') else None } return data_info + +def _print_stats(name, tensor): + if tensor is None: + logger.info(f"{name}: None") + return + # 对 PyTorch Tensor + if hasattr(tensor, "max"): + logger.info(f"{name} -> min: {tensor.min().item():.6f}, max: {tensor.max().item():.6f}, mean: {tensor.mean().item():.6f}, shape={tuple(tensor.shape)}") + else: + # 对 numpy + arr = np.asarray(tensor) + logger.info(f"{name} -> min: {arr.min():.6f}, max: {arr.max():.6f}, mean: {arr.mean():.6f}, shape={arr.shape}") + +def _check_soil_2d_stats(name, tensor): + if tensor is None: + logger.info(f"{name}: None") + return + x = tensor.squeeze(2) + grid_cells, n_vars, n_layers = x.shape + nonzero_mask = (x != 0).float() # [1000, 3, 10] + nonzero_ratio = nonzero_mask.mean(dim=0) # [3, 10] + max_vals = x.max(dim=0).values + min_vals = x.min(dim=0).values + mean_vals = x.mean(dim=0) + q25_vals = torch.quantile(x, 0.25, dim=0) + q75_vals = torch.quantile(x, 0.75, dim=0) + logger.info(f"{name} -> grid_cells: {grid_cells}, n_vars: {n_vars}, n_layers: {n_layers}, nonzero_ratio: {nonzero_ratio}, max_vals: {max_vals}, min_vals: {min_vals}, mean_vals: {mean_vals}, q25_vals: {q25_vals}, q75_vals: {q75_vals}") \ No newline at end of file From 202ba534799cd5db2426f2e34dbfef8e4e39b62c Mon Sep 17 00:00:00 2001 From: Zhuowei Gu Date: Thu, 13 Nov 2025 21:23:53 -0500 Subject: [PATCH 38/51] Fix can't load file prefix issue --- .../construct_forcing_20years.py | 41 +++++++++++++------ .../enhanced_training_dataset.py | 4 ++ 2 files changed, 32 insertions(+), 13 deletions(-) diff --git a/scripts/training_data_generation/python_scripts/construct_forcing_20years.py b/scripts/training_data_generation/python_scripts/construct_forcing_20years.py index a3ff157..fb2d4b5 100644 --- a/scripts/training_data_generation/python_scripts/construct_forcing_20years.py +++ b/scripts/training_data_generation/python_scripts/construct_forcing_20years.py @@ -17,6 +17,7 @@ import datetime import argparse from pathlib import Path +import re # Import configuration sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -50,27 +51,27 @@ def parse_arguments(): # Define all forcing variables and their file patterns forcing_variables = { 'FLDS': { - 'file_pattern': 'clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc', + 'token': 'TPQWL', 'description': 'Downward longwave radiation' }, 'FSDS': { - 'file_pattern': 'clmforc.Daymet.km.1d.Solr.{year}-{month:02d}.nc', + 'token': 'Solr', 'description': 'Downward shortwave radiation' }, 'PRECTmms': { - 'file_pattern': 'clmforc.Daymet.km.1d.Prec.{year}-{month:02d}.nc', + 'token': 'Prec', 'description': 'Precipitation rate' }, 'PSRF': { - 'file_pattern': 'clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc', + 'token': 'TPQWL', 'description': 'Surface pressure' }, 'QBOT': { - 'file_pattern': 'clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc', + 'token': 'TPQWL', 'description': 'Specific humidity' }, 'TBOT': { - 'file_pattern': 'clmforc.Daymet.km.1d.TPQWL.{year}-{month:02d}.nc', + 'token': 'TPQWL', 'description': 'Air temperature' } } @@ -105,20 +106,34 @@ def process_forcing_variable(var_name, var_info): all_monthly_files = [] # Pre-scan directory for efficiency (especially important for TES_NORTH with 4000+ files) - print(f"[{datetime.datetime.now()}] Pre-scanning directory for available files...") - available_files = set(os.listdir(data_dir)) - print(f"[{datetime.datetime.now()}] Found {len(available_files)} files in directory") + print(f"[{datetime.datetime.now()}] Pre-scanning directory for available files (recursive)...") + files_by_token = {} # {(token, year, month): [paths]} + total_files = 0 + pattern = re.compile(r'(?:.*_)?clmforc\..*\.(Prec|Solr|TPQWL)\.(\d{4})-(\d{2})\.nc$') + for root, _, files in os.walk(data_dir): + for fname in files: + total_files += 1 + match = pattern.search(fname) + if match: + token, year_s, month_s = match.groups() + key = (token, int(year_s), int(month_s)) + files_by_token.setdefault(key, []).append(os.path.join(root, fname)) + print(f"[{datetime.datetime.now()}] Found {total_files} files across {len(files_by_token)} token-year-month combinations") + token = var_info['token'] for year in range(start_year, end_year + 1): year_found = 0 for month in range(1, 13): - file_name = var_info['file_pattern'].format(year=year, month=month) - if file_name in available_files: - file_path = os.path.join(data_dir, file_name) + key = (token, year, month) + if key in files_by_token: + file_paths = files_by_token[key] + if len(file_paths) > 1: + print(f" Warning: Multiple matches for {token} {year}-{month:02d}; using {file_paths[0]}") + file_path = file_paths[0] all_monthly_files.append(file_path) year_found += 1 else: - print(f" Warning: File {file_name} does not exist, skipping.") + print(f" Warning: File with token {token} for {year}-{month:02d} does not exist, skipping.") print(f"[{datetime.datetime.now()}] {var_name}: year {year} -> found {year_found}/12 monthly files") if not all_monthly_files: diff --git a/scripts/training_data_generation/python_scripts/enhanced_training_dataset.py b/scripts/training_data_generation/python_scripts/enhanced_training_dataset.py index 32aec44..c70c84b 100644 --- a/scripts/training_data_generation/python_scripts/enhanced_training_dataset.py +++ b/scripts/training_data_generation/python_scripts/enhanced_training_dataset.py @@ -1128,6 +1128,10 @@ def generate_enhanced_dataset(base_output_dir, variable_definitions, initial_onl if not base_files: print("❌ No base PKL files found") return base_output_dir + + if initial_only_mode and (not config.final_spinup_history_files or not config.final_spinup_restart_files): + print("⚠️ Initial-only mode detected with no final spinup files; skipping enhanced dataset generation.") + return base_output_dir # Load restart files for enhancement file_path10 = config.ad_spinup_restart_files[0] From c0f66e966c6f4f42e26a38264f06616aafa7ee93 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Mon, 17 Nov 2025 14:02:16 -0500 Subject: [PATCH 39/51] data analysis function imporvement --- CNP_IO_updated9_dev_dw.txt | 2 +- scripts/cnp_result_validationplot.py | 61 +++++++++++++++---- scripts/generate_prediction_quality_report.py | 30 +++++++-- 3 files changed, 77 insertions(+), 16 deletions(-) diff --git a/CNP_IO_updated9_dev_dw.txt b/CNP_IO_updated9_dev_dw.txt index 1a2837b..4a427a0 100644 --- a/CNP_IO_updated9_dev_dw.txt +++ b/CNP_IO_updated9_dev_dw.txt @@ -71,4 +71,4 @@ SCALAR VARIABLES (1D - 4 variables): • soil3c_vr, soil3n_vr, soil3p_vr • soil4c_vr, soil4n_vr, soil4p_vr -• labilep_vr , occlp_vr, primp_vr, secondp_vr +• labilep_vr , occlp_vr, primp_vr, secondp_vr, solutionp_vr diff --git a/scripts/cnp_result_validationplot.py b/scripts/cnp_result_validationplot.py index 7d84f1b..de8f215 100644 --- a/scripts/cnp_result_validationplot.py +++ b/scripts/cnp_result_validationplot.py @@ -79,7 +79,40 @@ def _parse_top_bad_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): +def _parse_worst_vars_report(report_path): + """Parse quality_summary_report.txt to extract variables from + the '## Variables with Worst Predictions' section. + Returns a mapping { variable: { 'pfts': set(), 'layers': set() } } + """ + selection = {} + if not os.path.exists(report_path): + print(f"Worst-variables report not found: {report_path}") + return selection + in_section = False + try: + with open(report_path, 'r') as f: + for line in f: + stripped = line.strip('\n') + header = stripped.strip() + if header.startswith('## Variables with Worst Predictions'): + in_section = True + continue + # Section ends at next heading + if in_section and header.startswith('## '): + break + if in_section and stripped and not stripped.startswith('#'): + # Expected line format: ": % good, % ok, % bad" + m = re.match(r"\s*([A-Za-z0-9_]+):\s*", stripped) + if not m: + continue + var = m.group(1) + selection[var] = { 'pfts': set(), 'layers': set() } + except Exception as e: + print(f"Failed to parse worst variables section {report_path}: {e}") + 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): # Create plots subdirectory plots_dir = plots_dir_override or os.path.join(results_dir, "plots") os.makedirs(plots_dir, exist_ok=True) @@ -88,16 +121,21 @@ def main_with_flag(results_dir, plot_scatter, plot_loss, top_bad_only=False, top stats_path = os.path.join(results_dir, "validation_stats.csv") stats_data = [] - # Optional: restrict plotting to top-bad variables (and selected PFTs/layers) + # Optional: restrict plotting to selected variables (top-bad or worst list) selection = None - if top_bad_only: + if top_bad_only or worst_only: report_path = top_bad_report or os.path.join(results_dir, 'analysis', 'quality_summary_report.txt') - selection = _parse_top_bad_report(report_path) - if selection: - print(f"Plotting restricted to top-bad variables from: {report_path}") - else: - print("No selections parsed from top-bad report; proceeding without restriction.") - # IMPORTANT: ensure unrestricted plotting by clearing selection + if worst_only: + selection = _parse_worst_vars_report(report_path) + if selection: + print(f"Plotting restricted to worst variables from: {report_path}") + if (not selection) and top_bad_only: + selection = _parse_top_bad_report(report_path) + if selection: + print(f"Plotting restricted to top-bad variables from: {report_path}") + # If neither parser returned a selection, proceed unrestricted + if not selection: + print("No selections parsed from report; proceeding without restriction.") selection = None # Check for new directory structure first @@ -695,8 +733,9 @@ def plot_train_val_accuracy(loss_csv, out_dir): parser.add_argument('--no-plot-loss', action='store_false', dest='plot_loss', help='Do not plot train/val loss curve') # NEW: Stats-only mode disables all plots but still computes and saves statistics parser.add_argument('--stats-only', action='store_true', help='Only compute and save statistics CSV; do not generate any plots') - # NEW: Restrict plotting to top-bad variables from summary report + # NEW: Restrict plotting to top-bad or worst variables from summary report 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.set_defaults(plot_scatter=True, plot_loss=True) @@ -710,4 +749,4 @@ 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) \ No newline at end of file + 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 diff --git a/scripts/generate_prediction_quality_report.py b/scripts/generate_prediction_quality_report.py index 56c8d40..87aeee7 100644 --- a/scripts/generate_prediction_quality_report.py +++ b/scripts/generate_prediction_quality_report.py @@ -55,6 +55,9 @@ def main(): 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') args = parser.parse_args() # Set up input and output paths @@ -326,14 +329,32 @@ def categorize_prediction(row): plt.tight_layout() plt.savefig(output_dir / "r2_vs_rmse.png", dpi=300) - # 4. Optionally generate top-bad-only plots into a subfolder using the validation plotting utility + # 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) - + + # 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: @@ -354,9 +375,10 @@ def categorize_prediction(row): 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=True, + top_bad_only=not args.worst_only, top_bad_report=str(output_dir / 'quality_summary_report.txt'), - plots_dir_override=top_bad_out) + 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 From dd154515e3477c401d1abdbcf6abb4ea8bd61883 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Wed, 19 Nov 2025 12:23:21 -0500 Subject: [PATCH 40/51] update ai_prediction_to_restart --- scripts/ai_predictions_to_restart.py | 73 +++++++++++++++++++++++++++- 1 file changed, 71 insertions(+), 2 deletions(-) diff --git a/scripts/ai_predictions_to_restart.py b/scripts/ai_predictions_to_restart.py index 6878f3e..320f2f6 100644 --- a/scripts/ai_predictions_to_restart.py +++ b/scripts/ai_predictions_to_restart.py @@ -144,7 +144,8 @@ def auto_detect_variable_list(ai_predictions_path: Path) -> list: def create_updated_restart_file(restart_file_path: Path, output_path: Path, ai_predictions_path: Path, cnp_io_variables: List[str], - model_to_ai_mapping: np.ndarray, variable_mapping: Dict[str, Any]) -> None: + model_to_ai_mapping: np.ndarray, variable_mapping: Dict[str, Any], + strict_dims: bool = False) -> None: print(f"Saving updated restart file to: {output_path}") # Create output directory if it doesn't exist @@ -156,8 +157,65 @@ def create_updated_restart_file(restart_file_path: Path, output_path: Path, # Open the output file for direct modification with nc.Dataset(output_path, 'r+') as ds_out: + # Verify and adjust spinup_state + try: + if 'spinup_state' in ds_out.variables: + spin_var = ds_out.variables['spinup_state'] + try: + orig_val = np.array(spin_var[:]).item() if spin_var.size == 1 else None + except Exception: + orig_val = None + if orig_val is not None: + print(f"spinup_state in original restart (copied): {orig_val}") + if orig_val != 1: + print("Warning: Expected spinup_state==1 for adspinup; proceeding to set final_spinup (0) anyway") + else: + print("Warning: Could not read scalar value of spinup_state; proceeding to set to 0") + # Set to final_spinup mode (0) + try: + spin_var[...] = 0 + print("Set spinup_state to 0 (final_spinup) in updated restart") + except Exception as e: + print(f"Warning: Failed to set spinup_state to 0: {e}") + else: + print("Warning: 'spinup_state' variable not found in restart; skipping spinup flag update") + except Exception as e: + print(f"Warning: spinup_state check/update failed: {e}") # Load AI predictions with nc.Dataset(ai_predictions_path, 'r') as ds_ai: + # Helpers for shape/dimension checks + def _fail_or_warn(msg: str) -> bool: + if strict_dims: + raise ValueError(msg) + print(f"Warning: {msg} — skipping this variable") + return False + + def _check_pft_compat(ai_var: nc.Variable, model_var: nc.Variable) -> bool: + # Expect AI dims to include pft and gridcell + ai_dims = list(ai_var.dimensions) + if not ('pft' in ai_dims and 'gridcell' in ai_dims): + return _fail_or_warn(f"PFT var '{ai_var.name}' missing required dims (has {ai_dims}, need ['pft','gridcell'])") + # Model var should be 1D over pfts1d (or equivalent) + if len(model_var.shape) < 1: + return _fail_or_warn(f"Model PFT var '{model_var.name}' has invalid shape {model_var.shape}") + # Require at least 16 PFT slots (PFT1..PFT16). We skip PFT0 by design. + if model_var.shape[0] < 16: + return _fail_or_warn(f"Model PFT var '{model_var.name}' has insufficient length {model_var.shape[0]} (<16)") + return True + + def _check_soil_compat(ai_var: nc.Variable, model_var: nc.Variable) -> bool: + # Expect AI dims: (column, levgrnd, gridcell) + ai_dims = list(ai_var.dimensions) + required = {'column','levgrnd','gridcell'} + if not required.issubset(set(ai_dims)): + return _fail_or_warn(f"Soil var '{ai_var.name}' missing required dims (has {ai_dims}, need {sorted(required)})") + if len(model_var.shape) < 2: + return _fail_or_warn(f"Model soil var '{model_var.name}' has invalid shape {model_var.shape}") + # Need at least 10 layers in model to write top 10 + if model_var.shape[1] < 10: + return _fail_or_warn(f"Model soil var '{model_var.name}' has insufficient levgrnd={model_var.shape[1]} (<10)") + return True + # Update PFT variables for var_name in ds_ai.variables: if (var_name in ds_out.variables and @@ -169,6 +227,10 @@ def create_updated_restart_file(restart_file_path: Path, output_path: Path, print(f" AI data shape: {ai_data.shape}") print(f" Model variable shape: {model_var.shape}") + # Dimension compatibility check + if not _check_pft_compat(ds_ai.variables[var_name], model_var): + continue + # Get the grid-to-pfts mapping if 'grid_to_pfts' in variable_mapping: grid_to_pfts = variable_mapping['grid_to_pfts'] @@ -201,6 +263,10 @@ def create_updated_restart_file(restart_file_path: Path, output_path: Path, print(f" AI data shape: {ai_data.shape}") print(f" Model variable shape: {model_var.shape}") + # Dimension compatibility check + if not _check_soil_compat(ds_ai.variables[var_name], model_var): + continue + # Get the grid-to-cols mapping if 'grid_to_cols' in variable_mapping: grid_to_cols = variable_mapping['grid_to_cols'] @@ -302,6 +368,8 @@ def main(): help='Preview changes without saving updated restart file') parser.add_argument('--backup', action='store_true', help='Create backup of original restart file before updating') + parser.add_argument('--strict-dims', action='store_true', + help='Abort on any dimension mismatch instead of skipping') args = parser.parse_args() @@ -420,7 +488,8 @@ def main(): # Save updated restart file using direct NetCDF manipulation create_updated_restart_file(restart_file_path, output_path, ai_predictions_path, - cnp_io_variables, ai_to_model_mapping, variable_mapping) + cnp_io_variables, ai_to_model_mapping, variable_mapping, + strict_dims=args.strict_dims) print(f"\nRestart file updated successfully!") print(f"Original: {restart_file_path}") From 63d1fc467a749ccce9077f4a65beba7131a83aee Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Fri, 21 Nov 2025 09:27:36 -0500 Subject: [PATCH 41/51] customize path --- TVA_1_Sample/run_workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/TVA_1_Sample/run_workflow.py b/TVA_1_Sample/run_workflow.py index 2c3a362..ccf4f1e 100644 --- a/TVA_1_Sample/run_workflow.py +++ b/TVA_1_Sample/run_workflow.py @@ -58,7 +58,7 @@ def parse_args() -> argparse.Namespace: script_dir = Path(__file__).resolve().parent project_root = script_dir.parent default_locations = script_dir / "locations.csv" - default_variable_list = project_root / "CNP_IO_updated9_dev.txt" + default_variable_list = project_root / "CNP_IO_updated9_dev_dw.txt" default_model_config = project_root / "CNP_model_config_v01.txt" default_output_root = project_root / "final_restartfile" default_work_root = script_dir / "workflow_runs" From 04834cb7b855f099dbfef1ebbde8f640334d5a86 Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Sat, 22 Nov 2025 12:09:33 -0500 Subject: [PATCH 42/51] fix the location bug in site inference --- scripts/run_inference_all.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/scripts/run_inference_all.py b/scripts/run_inference_all.py index 29eea60..8693a44 100644 --- a/scripts/run_inference_all.py +++ b/scripts/run_inference_all.py @@ -1183,11 +1183,17 @@ def _preview(group_key: str, names: list): lat_idx = maybe_static_cols.index(lat_name) longitude_values = maybe_static[:, lon_idx] latitude_values = maybe_static[:, lat_idx] - # Fallback: use loader DataFrame if available - if (longitude_values is None or latitude_values is None) and hasattr(_loader, 'df') and isinstance(_loader.df, pd.DataFrame): + # Stronger preference: if the loader DataFrame has explicit coordinates, use those. + # This ensures we respect the exact site(s) selected for inference, even if static inverse lacks or mislabels coords. + if hasattr(_loader, 'df') and isinstance(_loader.df, pd.DataFrame): if 'Longitude' in _loader.df.columns and 'Latitude' in _loader.df.columns: - longitude_values = _loader.df['Longitude'].values[-len(test_data['static']):] if isinstance(test_data, dict) and 'static' in test_data else _loader.df['Longitude'].values - latitude_values = _loader.df['Latitude'].values[-len(test_data['static']):] if isinstance(test_data, dict) and 'static' in test_data else _loader.df['Latitude'].values + if isinstance(test_data, dict) and 'static' in test_data: + n = len(test_data['static']) + longitude_values = _loader.df['Longitude'].values[-n:] + latitude_values = _loader.df['Latitude'].values[-n:] + else: + longitude_values = _loader.df['Longitude'].values + latitude_values = _loader.df['Latitude'].values except Exception as _e_loc: logging.warning(f"Failed to prepare location vectors: {_e_loc}") From 70f6571c255e4f1378b870203d05762661aaf5a1 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Tue, 25 Nov 2025 21:39:34 -0800 Subject: [PATCH 43/51] Add multi-gridcell support in run_workflow.py and update runbook with --variable-list ../CNP_IO_updated9_dev.txt --- TVA_1_Sample/run_workflow.py | 147 +++++++++++++++++++++++++++++++---- docs/CNP_pipeline_runbook.md | 3 +- 2 files changed, 132 insertions(+), 18 deletions(-) diff --git a/TVA_1_Sample/run_workflow.py b/TVA_1_Sample/run_workflow.py index ccf4f1e..66e882c 100644 --- a/TVA_1_Sample/run_workflow.py +++ b/TVA_1_Sample/run_workflow.py @@ -13,9 +13,11 @@ from __future__ import annotations import argparse +from datetime import datetime import logging import math import re +import shutil import sys from pathlib import Path from typing import Dict, Iterable, List, Optional, Tuple @@ -75,7 +77,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--dataset-root", type=Path, - default=Path("/global/cfs/cdirs/m4814/daweigao/14_Code/TVA_training_dataset_all"), + default=Path("/global/cfs/cdirs/m4814/daweigao/14_Code/TVA_enhanced_dataset_solutionp"), help="Directory containing TVA training dataset pickle batches.", ) parser.add_argument( @@ -86,9 +88,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--model-path", type=Path, + # default=Path( + # "/global/cfs/cdirs/m4814/daweigao/15_code_Landsim/0_test/run_20251117_081732/cnp_predictions/model.pth" + # ), default=Path( "/global/cfs/cdirs/m4814/daweigao/15_code_Landsim/LandSim/" - "cnp_results/run_20251030_192921/cnp_predictions/model.pth" + "cnp_results/run_20251125_142005/cnp_predictions/model.pth" ), help="Path to the trained model checkpoint (.pth).", ) @@ -109,7 +114,7 @@ def parse_args() -> argparse.Namespace: type=Path, default=Path( "/global/cfs/cdirs/m4814/daweigao/14_Code/TVA_restart/" - "uELM_knox_I1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc" + "uELM_15sites4val_I1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc" ), help="Original restart NetCDF file that will be updated with AI predictions.", ) @@ -119,6 +124,15 @@ def parse_args() -> argparse.Namespace: default=default_output_root, help="Directory where updated restart files will be written.", ) + parser.add_argument( + "--final-restart", + type=Path, + default=None, + help=( + "Path for the combined updated restart file. " + "Defaults to /_updated.nc." + ), + ) parser.add_argument( "--work-root", type=Path, @@ -134,7 +148,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--skip-existing", action="store_true", - help="Skip processing a location if the final restart file already exists.", + help="Skip the workflow if the combined restart file already exists.", ) parser.add_argument( "--log-level", @@ -376,11 +390,54 @@ def update_restart_file( return output_restart +def update_spinup_state(restart_file: Path) -> None: + """ + Check and update spinup_state variable in restart file. + If spinup_state is 1, change it to 0; if already 0, leave unchanged. + """ + try: + import netCDF4 as nc + + # Open file and check if spinup_state exists + with nc.Dataset(restart_file, 'r') as ncfile: + if 'spinup_state' not in ncfile.variables: + LOGGER.warning("Variable 'spinup_state' not found in file") + return + + current_value = ncfile.variables['spinup_state'][:] + LOGGER.info("Current spinup_state value: %s", current_value) + + # Check if modification is needed + if np.any(current_value == 1): + LOGGER.info("Detected spinup_state = 1, changing to 0...") + + # Modify the file + with nc.Dataset(restart_file, 'r+') as ncfile: + old_value = ncfile.variables['spinup_state'][:] + new_value = np.where(old_value == 1, 0, old_value) + ncfile.variables['spinup_state'][:] = new_value + LOGGER.info("Changed spinup_state from %s to %s", old_value, new_value) + + # Verify the change + with nc.Dataset(restart_file, 'r') as ncfile: + final_value = ncfile.variables['spinup_state'][:] + LOGGER.info("Final spinup_state value: %s", final_value) + else: + LOGGER.info("spinup_state is already 0, no modification needed") + + except ImportError: + LOGGER.error("netCDF4 library is required. Install it with: pip install netCDF4") + except Exception as e: + LOGGER.error("Error processing spinup_state: %s", e) + + def process_location( location_row: pd.Series, args: argparse.Namespace, variable_map: Dict[str, List[str]], cnp_variables: List[str], + restart_source: Path, + restart_destination: Path, ) -> Optional[Path]: """Execute the full workflow for a single row in the locations CSV.""" latitude = float(location_row["latitude"]) @@ -394,11 +451,6 @@ def process_location( predictions_root = location_work_root / "inference" predictions_dir = predictions_root / "cnp_predictions" predictions_nc = location_work_root / f"{slug}_ai_predictions.nc" - restart_output = args.output_root / f"{args.restart_file.stem}_{slug}.nc" - - if args.skip_existing and restart_output.exists(): - LOGGER.info("Skipping location %s because %s already exists", location_label, restart_output) - return restart_output dataset_path = extract_location_dataset( dataset_root=args.dataset_root, @@ -458,25 +510,52 @@ def process_location( except Exception: pass + needs_temp_copy = restart_source.resolve() == restart_destination.resolve() + temp_output = ( + restart_destination + if not needs_temp_copy + else restart_destination.with_name(restart_destination.name + ".tmp") + ) + updated_restart = update_restart_file( - restart_file=args.restart_file, + restart_file=restart_source, ai_predictions_nc=predictions_nc, - output_restart=restart_output, + output_restart=temp_output, cnp_variables=cnp_variables, ) - if updated_restart.exists(): - size_mb = updated_restart.stat().st_size / (1024 * 1024) - LOGGER.info(" [Stage 4] Restart file created at %s (%.2f MB)", updated_restart, size_mb) + final_restart = updated_restart + if needs_temp_copy: + shutil.move(temp_output, restart_destination) + final_restart = restart_destination + LOGGER.info( + " [Stage 4] Applied updates in-place to %s (via temporary %s)", + restart_destination, + temp_output, + ) + + if final_restart.exists(): + size_mb = final_restart.stat().st_size / (1024 * 1024) + LOGGER.info(" [Stage 4] Restart file created at %s (%.2f MB)", final_restart, size_mb) else: - LOGGER.warning(" [Stage 4] Restart file expected but not found: %s", updated_restart) - return updated_restart + LOGGER.warning(" [Stage 4] Restart file expected but not found: %s", final_restart) + return final_restart def main() -> None: args = parse_args() configure_logging(args.log_level) + # Generate timestamp for this run + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + + # Add timestamp to output and work directories + args.output_root = args.output_root / f"run_{timestamp}" + args.work_root = args.work_root / f"run_{timestamp}" + LOGGER.info("Starting TVA workflow with locations file: %s", args.locations) + LOGGER.info("Output directory: %s", args.output_root) + LOGGER.info("Work directory: %s", args.work_root) + variable_map = determine_variable_map(args.variable_list) cnp_variables = collect_cnp_variables(variable_map) LOGGER.debug("Variable groups loaded: %s", variable_map) @@ -485,12 +564,33 @@ def main() -> None: args.output_root.mkdir(parents=True, exist_ok=True) args.work_root.mkdir(parents=True, exist_ok=True) + combined_restart = args.final_restart + if combined_restart is None: + combined_restart = args.output_root / f"{args.restart_file.stem}_updated.nc" + combined_restart = combined_restart.resolve() + combined_restart.parent.mkdir(parents=True, exist_ok=True) + + if args.skip_existing and combined_restart.exists(): + LOGGER.info("Combined restart file already exists at %s. Skipping workflow.", combined_restart) + return + + current_restart_source = args.restart_file.resolve() + results: List[Tuple[str, Optional[Path]]] = [] for _, row in locations_df.iterrows(): try: - updated_restart = process_location(row, args, variable_map, cnp_variables) + updated_restart = process_location( + row, + args, + variable_map, + cnp_variables, + restart_source=current_restart_source, + restart_destination=combined_restart, + ) coord_label = f"{row['latitude']:.4f}_{row['longitude']:.4f}" results.append((coord_label, updated_restart)) + if updated_restart: + current_restart_source = combined_restart except Exception as exc: # pragma: no cover - logging safety LOGGER.exception("Failed to process location row %s: %s", row.to_dict(), exc) coord_label = f"{row['latitude']:.4f}_{row['longitude']:.4f}" @@ -503,6 +603,19 @@ def main() -> None: LOGGER.warning(" %s: failed (no restart generated)", name) else: LOGGER.info(" %s: %s", name, path) + + # Update spinup_state in final restart file + if combined_restart.exists(): + LOGGER.info("") + LOGGER.info("=" * 80) + LOGGER.info("Checking and updating spinup_state variable...") + LOGGER.info("=" * 80) + update_spinup_state(combined_restart) + LOGGER.info("=" * 80) + LOGGER.info("spinup_state check/update complete") + LOGGER.info("=" * 80) + else: + LOGGER.warning("Final restart file does not exist, skipping spinup_state update") if __name__ == "__main__": diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 3068ec1..799c9ad 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -150,5 +150,6 @@ Geographic coordinates are defined in locations.csv. python ./TVA_1_Sample/run_workflow.py \ --restart-file /path/to/20year_restart_file.nc \ --model-path /path/to/trained_model_TVA.pt \ - --dataset-root /path/to/TVA_dataset + --dataset-root /path/to/TVA_dataset \ + --variable-list ../CNP_IO_updated9_dev.txt ``` From ef9f86dfa445128b129d5e819a156677fd7e2f49 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Wed, 10 Dec 2025 22:49:56 -0800 Subject: [PATCH 44/51] Add restart point extraction tool and update runbook --- docs/CNP_pipeline_runbook.md | 19 ++ scripts/extract_elm_restart_point.py | 313 +++++++++++++++++++++++++++ 2 files changed, 332 insertions(+) create mode 100644 scripts/extract_elm_restart_point.py diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 799c9ad..ba0d7dd 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -153,3 +153,22 @@ python ./TVA_1_Sample/run_workflow.py \ --dataset-root /path/to/TVA_dataset \ --variable-list ../CNP_IO_updated9_dev.txt ``` + +### Extract Single Point Data from ELM Restart File +The `extract_elm_restart_point.py` script extracts all data for specified geographic coordinates from a global ELM restart NetCDF file. It correctly handles ELM's multi-level hierarchical structure (gridcell → topounit → landunit → column → pft) and creates a subset NetCDF file containing only the data for the target location. + +**Key Features**: +- Extracts data at all hierarchical levels for a single geographic point +- Automatically finds the nearest neighbor gridcell to target coordinates +- Preserves all metadata and encoding information from the original file +- Supports both 0-360° and -180 to 180° longitude formats +- Can be used via command-line arguments or with hardcoded defaults + +**Usage**: +```bash +python scripts/extract_elm_restart_point.py \ + --restart-file /path/to/input_restart_file.nc \ + --lat 35.833332(Target latitude coordinate) \ + --lon -84.208336(Target longitude coordinate) \ + --output-file single_point_extracted.nc +``` \ No newline at end of file diff --git a/scripts/extract_elm_restart_point.py b/scripts/extract_elm_restart_point.py new file mode 100644 index 0000000..e2237b0 --- /dev/null +++ b/scripts/extract_elm_restart_point.py @@ -0,0 +1,313 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +ELM Restart File Single Point Data Extraction Tool +Function: Extract all data for specified coordinates from global ELM restart files (correctly handles multi-level structure) +""" + +import xarray as xr +import numpy as np +import os +import argparse +from datetime import datetime + +# ==================== Configuration Parameters (Default Values) ==================== + +SOURCE_NC = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc' +OUTPUT_NC = 'single_point_20_year_restart_extracted_43_56.nc' # Output filename +# Target coordinates +TARGET_LAT = 35.833332 +TARGET_LON = -84.208336 + +# ==================== Command Line Argument Parser ==================== +def parse_arguments(): + """Parse command line arguments""" + parser = argparse.ArgumentParser( + description='Extract single point data from ELM restart file', + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python extract_elm_restart_point.py --restart-file /path/to/restart.nc --lat 35.83 --lon -84.21 + python extract_elm_restart_point.py --restart-file /path/to/restart.nc --lat 35.83 --lon -84.21 --output-file output.nc + """ + ) + + parser.add_argument( + '--restart-file', + type=str, + default=None, + help='Path to input ELM restart NetCDF file (default: uses hardcoded SOURCE_NC)' + ) + + parser.add_argument( + '--output-file', + type=str, + default=None, + help='Path to output NetCDF file (default: uses hardcoded OUTPUT_NC)' + ) + + parser.add_argument( + '--lat', + type=float, + default=None, + help='Target latitude coordinate (default: uses hardcoded TARGET_LAT)' + ) + + parser.add_argument( + '--lon', + type=float, + default=None, + help='Target longitude coordinate (default: uses hardcoded TARGET_LON)' + ) + + return parser.parse_args() + + +# ==================== Main Function ==================== +def extract_single_point_elm(source_nc=None, output_nc=None, target_lat=None, target_lon=None): + """ + Extract single point data from global ELM restart file + Correctly handles ELM's multi-level structure (gridcell -> topounit -> landunit -> column -> pft) + + Args: + source_nc: Path to source NetCDF file (if None, uses SOURCE_NC) + output_nc: Path to output NetCDF file (if None, uses OUTPUT_NC) + target_lat: Target latitude (if None, uses TARGET_LAT) + target_lon: Target longitude (if None, uses TARGET_LON) + """ + # Use provided parameters or fall back to defaults + source_file = source_nc if source_nc is not None else SOURCE_NC + output_file = output_nc if output_nc is not None else OUTPUT_NC + lat = target_lat if target_lat is not None else TARGET_LAT + lon = target_lon if target_lon is not None else TARGET_LON + + print("="*70) + print("ELM Restart File Single Point Data Extraction (Multi-level Version)") + print("="*70) + print(f"Source file: {source_file}") + print(f"Target coordinates: Lat={lat}, Lon={lon}") + print("="*70) + + # Check file existence + if not os.path.exists(source_file): + raise FileNotFoundError(f"Error: Source file not found {source_file}") + + # ============ Step 1: Open dataset ============ + print("\n[1/5] Loading dataset...") + ds = xr.open_dataset(source_file, decode_times=False, mask_and_scale=False) + + print(f" ✓ Dataset loaded successfully") + print(f" - File size: {os.path.getsize(source_file) / (1024**3):.2f} GB") + print(f" - Dimensions: gridcell={ds.dims['gridcell']}, topounit={ds.dims['topounit']}, " + + f"landunit={ds.dims['landunit']}, column={ds.dims['column']}, pft={ds.dims['pft']}") + print(f" - Number of variables: {len(ds.data_vars)}") + + # ============ Step 2: Find nearest neighbor gridcell ============ + print("\n[2/5] Locating nearest neighbor gridcell coordinates...") + + # Read gridcell-level latitude and longitude + lat_vals = ds['grid1d_lat'].values + lon_vals = ds['grid1d_lon'].values + + # Handle longitude format conversion (0-360° vs -180 to 180°) + lon_adjusted = lon_vals.copy() + if lon_vals.max() > 200: + lon_adjusted = np.where(lon_vals > 180, lon_vals - 360, lon_vals) + print(f" Detected 0-360° longitude format, converted to -180 to 180° format") + + # Calculate Euclidean distance and find nearest neighbor + dist = np.sqrt((lat_vals - lat)**2 + (lon_adjusted - lon)**2) + gridcell_idx_py = int(np.argmin(dist)) # Python 0-based index + gridcell_idx_elm = gridcell_idx_py + 1 # ELM 1-based index + + found_lat = float(lat_vals[gridcell_idx_py]) + found_lon = float(lon_vals[gridcell_idx_py]) + found_lon_adjusted = float(lon_adjusted[gridcell_idx_py]) + + print(f" ✓ Found nearest neighbor gridcell") + print(f" - Python index (0-based): {gridcell_idx_py}") + print(f" - ELM index (1-based): {gridcell_idx_elm}") + print(f" - Target coordinates: ({lat:.6f}, {lon:.6f})") + print(f" - Matched coordinates: ({found_lat:.6f}, {found_lon_adjusted:.6f})") + print(f" - Euclidean distance: {dist[gridcell_idx_py]:.6f}°") + + # ============ Step 3: Find indices for all levels ============ + print("\n[3/5] Finding all level data for this gridcell...") + + # Find all level indices belonging to this gridcell + indices = {} + + # Gridcell level (use found index directly) + indices['gridcell'] = [gridcell_idx_py] + + # Topounit level + if 'topo1d_gridcell_index' in ds.variables: + topo_gc_idx = ds['topo1d_gridcell_index'].values + topo_match = np.where(topo_gc_idx == gridcell_idx_elm)[0] + indices['topounit'] = topo_match.tolist() + print(f" - Found {len(topo_match)} topounits") + + # Landunit level + if 'land1d_gridcell_index' in ds.variables: + land_gc_idx = ds['land1d_gridcell_index'].values + land_match = np.where(land_gc_idx == gridcell_idx_elm)[0] + indices['landunit'] = land_match.tolist() + print(f" - Found {len(land_match)} landunits") + + # Column level + if 'cols1d_gridcell_index' in ds.variables: + cols_gc_idx = ds['cols1d_gridcell_index'].values + cols_match = np.where(cols_gc_idx == gridcell_idx_elm)[0] + indices['column'] = cols_match.tolist() + print(f" - Found {len(cols_match)} columns") + + # PFT level + if 'pfts1d_gridcell_index' in ds.variables: + pfts_gc_idx = ds['pfts1d_gridcell_index'].values + pfts_match = np.where(pfts_gc_idx == gridcell_idx_elm)[0] + indices['pft'] = pfts_match.tolist() + print(f" - Found {len(pfts_match)} pfts") + + print(f" ✓ All level index search completed") + + # ============ Step 4: Extract data ============ + print("\n[4/5] Extracting single point data...") + + # Filter data by level + subset = ds.copy() + + # Filter each dimension + for dim_name, dim_indices in indices.items(): + if dim_name in subset.dims and len(dim_indices) > 0: + subset = subset.isel({dim_name: dim_indices}) + print(f" ✓ Filtered {dim_name}: {len(dim_indices)} elements") + + # Calculate data compression ratio + original_size_estimate = sum([ + ds.dims['gridcell'], + ds.dims['topounit'], + ds.dims['landunit'], + ds.dims['column'], + ds.dims['pft'] + ]) + subset_size_estimate = sum([ + len(indices.get('gridcell', [])), + len(indices.get('topounit', [])), + len(indices.get('landunit', [])), + len(indices.get('column', [])), + len(indices.get('pft', [])) + ]) + compression_ratio = original_size_estimate / subset_size_estimate if subset_size_estimate > 0 else 0 + + print(f" ✓ Data extraction completed") + print(f" - Compression ratio: {compression_ratio:.1f}x (original {original_size_estimate} -> extracted {subset_size_estimate} elements)") + + # ============ Step 5: Preserve metadata and save ============ + print("\n[5/5] Saving file...") + + # Preserve original encoding information for each variable + encoding = {} + for var_name in subset.variables: + encoding[var_name] = {} + + if var_name in ds.variables: + original_var = ds[var_name] + + if hasattr(original_var, 'encoding'): + for key in ['dtype', 'scale_factor', 'add_offset', '_FillValue', + 'missing_value', 'zlib', 'complevel', 'shuffle', + 'chunksizes', 'fletcher32', 'contiguous']: + if key in original_var.encoding: + encoding[var_name][key] = original_var.encoding[key] + + # Add extraction information to global attributes + subset.attrs['extraction_date'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S') + subset.attrs['extraction_target_lat'] = lat + subset.attrs['extraction_target_lon'] = lon + subset.attrs['extraction_actual_lat'] = found_lat + subset.attrs['extraction_actual_lon'] = found_lon_adjusted + subset.attrs['extraction_source_file'] = os.path.basename(source_file) + subset.attrs['extraction_gridcell_index_python'] = gridcell_idx_py + subset.attrs['extraction_gridcell_index_elm'] = gridcell_idx_elm + + print(f" Writing to: {output_file}") + + # Save file + subset.to_netcdf( + output_file, + format='NETCDF4', + encoding=encoding, + unlimited_dims=None + ) + + # Verify output + output_size = os.path.getsize(output_file) / (1024**2) # MB + print(f" ✓ File saved successfully") + print(f" - Output size: {output_size:.2f} MB") + print(f" - Number of variables: {len(subset.data_vars)}") + print(f" - Global attributes: {len(subset.attrs)}") + + # Close dataset + ds.close() + + print("\n" + "="*70) + print("✓ Processing completed!") + print("="*70) + print("\nDimension statistics:") + for dim in subset.dims: + print(f" {dim}: {subset.dims[dim]}") + + print("\nTip: You can use the following commands to view the output file:") + print(f" ncdump -h {output_file}") + print(f"Or in Python:") + print(f" import xarray as xr") + print(f" ds = xr.open_dataset('{output_file}')") + print(f" print(ds)") + + return subset + + +# ==================== Helper Function: Quick File Structure Inspection ==================== +def inspect_structure(source_nc=None): + """Quickly view the hierarchical structure information of NetCDF file""" + source_file = source_nc if source_nc is not None else SOURCE_NC + print("Checking file structure...") + ds = xr.open_dataset(source_file, decode_times=False) + + print("\nDimensions:") + for dim, size in ds.dims.items(): + print(f" {dim}: {size}") + + print("\nIndex mapping variables:") + for var in ds.variables: + if 'index' in var.lower(): + print(f" {var}: {ds[var].dims} - shape={ds[var].shape}") + if ds[var].size < 50: + print(f" Example values: {ds[var].values[:10]}") + + ds.close() + + +# ==================== Main Program Entry ==================== +if __name__ == "__main__": + try: + # Parse command line arguments + args = parse_arguments() + + # Uncomment the line below to view file structure first + # inspect_structure(args.restart_file) + + # Execute extraction with command line arguments (or defaults) + result = extract_single_point_elm( + source_nc=args.restart_file, + output_nc=args.output_file, + target_lat=args.lat, + target_lon=args.lon + ) + + except Exception as e: + print(f"\n❌ Error: {e}") + import traceback + traceback.print_exc() + exit(1) + From 1f1c015d35dd601edd4a1399ae723f1d9f76b2ab Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Tue, 16 Dec 2025 07:20:34 -0800 Subject: [PATCH 45/51] update extract_elm_restart_point.py --- scripts/extract_elm_restart_point.py | 56 +++++++++++++++++++++++----- 1 file changed, 46 insertions(+), 10 deletions(-) diff --git a/scripts/extract_elm_restart_point.py b/scripts/extract_elm_restart_point.py index e2237b0..3a5909a 100644 --- a/scripts/extract_elm_restart_point.py +++ b/scripts/extract_elm_restart_point.py @@ -12,12 +12,12 @@ from datetime import datetime # ==================== Configuration Parameters (Default Values) ==================== - SOURCE_NC = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc' -OUTPUT_NC = 'single_point_20_year_restart_extracted_43_56.nc' # Output filename +# SOURCE_NC = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNPRDCTCBC.elm.r.0781-01-01-00000.nc' +OUTPUT_NC = 'single_point_20_year_restart_extracted_303_17.nc' # Output filename # Target coordinates -TARGET_LAT = 35.833332 -TARGET_LON = -84.208336 +TARGET_LAT = -17.4246 +TARGET_LON = 303.75 # ==================== Command Line Argument Parser ==================== def parse_arguments(): @@ -109,14 +109,50 @@ def extract_single_point_elm(source_nc=None, output_nc=None, target_lat=None, ta lat_vals = ds['grid1d_lat'].values lon_vals = ds['grid1d_lon'].values + # Print unique latitude and longitude values in dataset + unique_lats = np.unique(lat_vals) + unique_lons = np.unique(lon_vals) + print(f" Dataset coordinate ranges:") + print(f" - Unique latitudes: {len(unique_lats)} values, range [{unique_lats.min():.6f}, {unique_lats.max():.6f}]") + print(f" - Unique longitudes (original): {len(unique_lons)} values, range [{unique_lons.min():.6f}, {unique_lons.max():.6f}]") + print(f" - All unique latitudes: {unique_lats}") + print(f" - All unique longitudes (original): {unique_lons}") + # Handle longitude format conversion (0-360° vs -180 to 180°) - lon_adjusted = lon_vals.copy() - if lon_vals.max() > 200: + # 智能选择：如果数据集和目标都是同一格式，保持原格式；否则统一转换 + dataset_uses_360 = lon_vals.max() > 200 + # 判断目标经度格式：> 180 肯定是 0-360，< 0 肯定是 -180~180，0-180 之间默认认为是 0-360 + target_uses_360 = lon > 180 or (lon >= 0 and lon <= 180 and dataset_uses_360) + + # 如果数据集和目标都是 0-360 格式，保持原格式比较（更高效） + if dataset_uses_360 and target_uses_360: + lon_adjusted = lon_vals.copy() # 保持 0-360 格式 + target_lon_adjusted = lon # 保持 0-360 格式 + print(f" Dataset uses 0-360° format, target also in 0-360° format ({lon:.6f})") + print(f" → Using 0-360° format for comparison (no conversion needed)") + # 如果数据集是 0-360 但目标是 -180~180，转换数据集 + elif dataset_uses_360 and not target_uses_360: lon_adjusted = np.where(lon_vals > 180, lon_vals - 360, lon_vals) - print(f" Detected 0-360° longitude format, converted to -180 to 180° format") - - # Calculate Euclidean distance and find nearest neighbor - dist = np.sqrt((lat_vals - lat)**2 + (lon_adjusted - lon)**2) + target_lon_adjusted = lon # 目标已经是 -180~180 格式 + print(f" Dataset uses 0-360° format, target uses -180~180° format") + print(f" → Converting dataset to -180~180° format for comparison") + unique_lons_adjusted = np.unique(lon_adjusted) + print(f" - Unique longitudes (adjusted): {len(unique_lons_adjusted)} values, range [{unique_lons_adjusted.min():.6f}, {unique_lons_adjusted.max():.6f}]") + print(f" - All unique longitudes (adjusted): {unique_lons_adjusted}") + # 如果数据集是 -180~180 但目标是 0-360，转换目标 + elif not dataset_uses_360 and target_uses_360: + lon_adjusted = lon_vals.copy() # 数据集已经是 -180~180 格式 + target_lon_adjusted = lon - 360 if lon > 180 else lon # 转换目标到 -180~180 + print(f" Dataset uses -180~180° format, target uses 0-360° format ({lon:.6f})") + print(f" → Converting target to -180~180° format ({target_lon_adjusted:.6f}) for comparison") + # 如果两者都是 -180~180 格式，直接使用 + else: + lon_adjusted = lon_vals.copy() + target_lon_adjusted = lon + print(f" Both dataset and target use -180~180° format") + + # Calculate Euclidean distance using the ADJUSTED target + dist = np.sqrt((lat_vals - lat)**2 + (lon_adjusted - target_lon_adjusted)**2) gridcell_idx_py = int(np.argmin(dist)) # Python 0-based index gridcell_idx_elm = gridcell_idx_py + 1 # ELM 1-based index From f04848175a10b1ac2a81147ef138a0639e89d301 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Tue, 6 Jan 2026 11:26:17 -0800 Subject: [PATCH 46/51] Add single-point restart extraction doc and nc comparison script --- docs/single_point_restart_extraction.md | 53 +++++++++++++++++++ scripts/compare_nc2.py | 67 +++++++++++++++++++++++++ 2 files changed, 120 insertions(+) create mode 100644 docs/single_point_restart_extraction.md create mode 100644 scripts/compare_nc2.py diff --git a/docs/single_point_restart_extraction.md b/docs/single_point_restart_extraction.md new file mode 100644 index 0000000..e331004 --- /dev/null +++ b/docs/single_point_restart_extraction.md @@ -0,0 +1,53 @@ +# Single Point Restart File Extraction + +This guide shows how to extract single point restart files: +1. How to Get Single Point 20-year Ad-Spin Up Restart File +2. How to Get Single Point AI-Updated Ad-Spin Up Restart File +3. How to Validate +## Target Site Coordinates + +The Amazon site coordinates in TRENDY coordinate system: +- **Longitude**: 303.75° +- **Latitude**: -17.4246° + +## 1. How to Get Single Point 20-year Ad-Spin Up Restart File + +Extract single point from the original 20-year ad-spin up restart file: + +```bash +python scripts/extract_elm_restart_point.py \ + --restart-file \ + --lat -17.4246 \ + --lon 303.75 \ + --output-file +``` + +This command extracts the single point data and saves it to the specified output directory. + +## 2. How to Get Single Point AI-Updated Ad-Spin Up Restart File + +**Step 1:** Follow `CNP_pipeline_runbook.md` to train a model and generate an AI-updated ad-spin up restart file across all sites. + +**Step 2:** Extract single point from the AI-updated ad-spin up restart file: + +```bash +python scripts/extract_elm_restart_point.py \ + --restart-file \ + --lat -17.4246 \ + --lon 303.75 \ + --output-file +``` + +This command extracts the single point data from the AI-predicted restart file and saves it to the specified output directory. + +## 3. How to Validate + +Compare the original single point restart file with the AI-updated single point restart file to validate the differences: + +```bash +python3 scripts/compare_nc2.py \ + --file_name1 \ + --file_name2 +``` + +This command compares the two files and reports differences in variables, data types, shapes, and values. diff --git a/scripts/compare_nc2.py b/scripts/compare_nc2.py new file mode 100644 index 0000000..432aaf0 --- /dev/null +++ b/scripts/compare_nc2.py @@ -0,0 +1,67 @@ +import numpy as np +import netCDF4 as nc +import sys +import argparse + +def compare_variables(var1, var2): + if var1.dtype != var2.dtype: + print(f'Different data types: {var1.dtype} vs {var2.dtype}') + if var1.shape != var2.shape: + print(f'Different shapes: {var1.shape} vs {var2.shape}') + if (np.issubdtype(var1.dtype, np.number) and var1.shape == var2.shape and var1.dtype != 'short'): + + if len(var1.shape) <= 2: + compare_data(var1[:], var2[:]) + elif len(var1.shape) == 3: + for i in range(var1.shape[0]): + compare_data(var1[i, :, :], var2[i, :, :]) + elif len(var1.shape) == 4: + for i in range(var1.shape[0]): + for j in range(var1.shape[1]): + compare_data(var1[i, j, :, :], var2[i, j, :, :]) + +def compare_data(data1, data2): + if not np.allclose(data1, data2): + print(f'Difference in data:') + print(f'Sum: {np.sum(data1)} vs {np.sum(data2)}') + print(f'Mean: {np.mean(data1)} vs {np.mean(data2)}') + print(f'Max: {np.max(data1)} vs {np.max(data2)}') + print(f'Min: {np.min(data1)} vs {np.min(data2)}') + +def main(): + parser = argparse.ArgumentParser(description='Compare two NetCDF files') + parser.add_argument('--file_name1', type=str, + default='/global/cfs/cdirs/m4814/daweigao/15_code_Landsim/LandSim/1_single_point/AI_updated_single_point_303_17_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc', + help='Path to the first NetCDF file') + parser.add_argument('--file_name2', type=str, + default='/global/cfs/cdirs/m4814/daweigao/15_code_Landsim/LandSim/1_single_point/single_point_303_17_20251201_TRENDY2024_default_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc', + help='Path to the second NetCDF file') + + args = parser.parse_args() + file_name1 = args.file_name1 + file_name2 = args.file_name2 + + + + file1 = nc.Dataset(file_name1) + file2 = nc.Dataset(file_name2) + + variables1 = file1.variables + variables2 = file2.variables + + for var in variables1: + if var in variables2: + print(var) + compare_variables(variables1[var], variables2[var]) + else: + print(f'Variable {var} is not in the second file') + + for var in variables2: + if var not in variables1: + print(f'Variable {var} is not in the first file') + + file1.close() + file2.close() + +if __name__ == '__main__': + main() \ No newline at end of file From 793b4a69f51a0a09569804331c940aae40427b26 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Tue, 13 Jan 2026 12:14:23 -0800 Subject: [PATCH 47/51] Fix dimension mismatch bug in extract_elm_restart_point --- scripts/extract_elm_restart_point.py | 62 ++++++++++++++++++++++------ 1 file changed, 49 insertions(+), 13 deletions(-) diff --git a/scripts/extract_elm_restart_point.py b/scripts/extract_elm_restart_point.py index 3a5909a..ccc265c 100644 --- a/scripts/extract_elm_restart_point.py +++ b/scripts/extract_elm_restart_point.py @@ -10,6 +10,7 @@ import os import argparse from datetime import datetime +from netCDF4 import Dataset as NC4Dataset # ==================== Configuration Parameters (Default Values) ==================== SOURCE_NC = '/global/cfs/cdirs/m4814/daweigao/14_Code/all_dataset_1_degree/20250117_trendytest_ICB1850CNRDCTCBC_ad_spinup.elm.r.0021-01-01-00000.nc' @@ -94,11 +95,19 @@ def extract_single_point_elm(source_nc=None, output_nc=None, target_lat=None, ta # ============ Step 1: Open dataset ============ print("\n[1/5] Loading dataset...") + + # First, get ALL dimensions using netCDF4 (xarray may filter some out) + nc4_file = NC4Dataset(source_file, 'r') + all_nc_dims = {dim: len(nc4_file.dimensions[dim]) for dim in nc4_file.dimensions} + nc4_file.close() + ds = xr.open_dataset(source_file, decode_times=False, mask_and_scale=False) print(f" ✓ Dataset loaded successfully") print(f" - File size: {os.path.getsize(source_file) / (1024**3):.2f} GB") - print(f" - Dimensions: gridcell={ds.dims['gridcell']}, topounit={ds.dims['topounit']}, " + + print(f" - Dimensions (netCDF4): {len(all_nc_dims)} total") + print(f" - Dimensions (xarray): {len(ds.dims)} visible") + print(f" - Main dims: gridcell={ds.dims['gridcell']}, topounit={ds.dims['topounit']}, " + f"landunit={ds.dims['landunit']}, column={ds.dims['column']}, pft={ds.dims['pft']}") print(f" - Number of variables: {len(ds.data_vars)}") @@ -119,33 +128,33 @@ def extract_single_point_elm(source_nc=None, output_nc=None, target_lat=None, ta print(f" - All unique longitudes (original): {unique_lons}") # Handle longitude format conversion (0-360° vs -180 to 180°) - # 智能选择：如果数据集和目标都是同一格式，保持原格式；否则统一转换 + # Smart selection: if dataset and target use the same format, keep original format for efficiency dataset_uses_360 = lon_vals.max() > 200 - # 判断目标经度格式：> 180 肯定是 0-360，< 0 肯定是 -180~180，0-180 之间默认认为是 0-360 + # Determine target longitude format: > 180 must be 0-360, < 0 must be -180~180 target_uses_360 = lon > 180 or (lon >= 0 and lon <= 180 and dataset_uses_360) - # 如果数据集和目标都是 0-360 格式，保持原格式比较（更高效） + # If both dataset and target are in 0-360 format, keep original format (more efficient) if dataset_uses_360 and target_uses_360: - lon_adjusted = lon_vals.copy() # 保持 0-360 格式 - target_lon_adjusted = lon # 保持 0-360 格式 + lon_adjusted = lon_vals.copy() # Keep 0-360 format + target_lon_adjusted = lon # Keep 0-360 format print(f" Dataset uses 0-360° format, target also in 0-360° format ({lon:.6f})") print(f" → Using 0-360° format for comparison (no conversion needed)") - # 如果数据集是 0-360 但目标是 -180~180，转换数据集 + # If dataset is 0-360 but target is -180~180, convert dataset elif dataset_uses_360 and not target_uses_360: lon_adjusted = np.where(lon_vals > 180, lon_vals - 360, lon_vals) - target_lon_adjusted = lon # 目标已经是 -180~180 格式 + target_lon_adjusted = lon # Target is already in -180~180 format print(f" Dataset uses 0-360° format, target uses -180~180° format") print(f" → Converting dataset to -180~180° format for comparison") unique_lons_adjusted = np.unique(lon_adjusted) print(f" - Unique longitudes (adjusted): {len(unique_lons_adjusted)} values, range [{unique_lons_adjusted.min():.6f}, {unique_lons_adjusted.max():.6f}]") print(f" - All unique longitudes (adjusted): {unique_lons_adjusted}") - # 如果数据集是 -180~180 但目标是 0-360，转换目标 + # If dataset is -180~180 but target is 0-360, convert target elif not dataset_uses_360 and target_uses_360: - lon_adjusted = lon_vals.copy() # 数据集已经是 -180~180 格式 - target_lon_adjusted = lon - 360 if lon > 180 else lon # 转换目标到 -180~180 + lon_adjusted = lon_vals.copy() # Dataset is already in -180~180 format + target_lon_adjusted = lon - 360 if lon > 180 else lon # Convert target to -180~180 print(f" Dataset uses -180~180° format, target uses 0-360° format ({lon:.6f})") print(f" → Converting target to -180~180° format ({target_lon_adjusted:.6f}) for comparison") - # 如果两者都是 -180~180 格式，直接使用 + # If both use -180~180 format, use directly else: lon_adjusted = lon_vals.copy() target_lon_adjusted = lon @@ -212,12 +221,27 @@ def extract_single_point_elm(source_nc=None, output_nc=None, target_lat=None, ta # Filter data by level subset = ds.copy() - # Filter each dimension + # Identify global dimensions (dimensions that don't vary with gridcell) + # These are common vertical layers or other global dimensions in ELM restart files + all_dims = set(ds.dims.keys()) + spatial_dims = set(indices.keys()) + global_dims = all_dims - spatial_dims + + print(f" - Spatial dimensions to filter: {spatial_dims}") + print(f" - Global dimensions to preserve: {global_dims}") + + # Filter each spatial dimension (only filter spatial dimensions, preserve global dimensions) for dim_name, dim_indices in indices.items(): if dim_name in subset.dims and len(dim_indices) > 0: subset = subset.isel({dim_name: dim_indices}) print(f" ✓ Filtered {dim_name}: {len(dim_indices)} elements") + # Verify that global dimensions are preserved + print(f" ✓ Preserved global dimensions:") + for dim in global_dims: + if dim in subset.dims: + print(f" - {dim}: {subset.dims[dim]} elements") + # Calculate data compression ratio original_size_estimate = sum([ ds.dims['gridcell'], @@ -276,6 +300,18 @@ def extract_single_point_elm(source_nc=None, output_nc=None, target_lat=None, ta unlimited_dims=None ) + # Add missing dimensions that were in original file but filtered by xarray + missing_dims = set(all_nc_dims.keys()) - set(subset.dims.keys()) + if missing_dims: + print(f" Adding {len(missing_dims)} missing dimensions from original file...") + nc_out = NC4Dataset(output_file, 'a') + for dim_name in missing_dims: + if dim_name not in nc_out.dimensions: + dim_size = all_nc_dims[dim_name] + nc_out.createDimension(dim_name, dim_size) + print(f" + {dim_name}: {dim_size}") + nc_out.close() + # Verify output output_size = os.path.getsize(output_file) / (1024**2) # MB print(f" ✓ File saved successfully") From 3ad7374d11feb3aa6e9cdd098085cbf7d58967d7 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Thu, 15 Jan 2026 14:25:51 -0800 Subject: [PATCH 48/51] Add IJCAI-style NEE time series plotting script --- scripts/plot_nee_timeseries_ijcai.py | 130 +++++++++++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 scripts/plot_nee_timeseries_ijcai.py diff --git a/scripts/plot_nee_timeseries_ijcai.py b/scripts/plot_nee_timeseries_ijcai.py new file mode 100644 index 0000000..2c1f68a --- /dev/null +++ b/scripts/plot_nee_timeseries_ijcai.py @@ -0,0 +1,130 @@ +import argparse +import re +from pathlib import Path + +import numpy as np +import xarray as xr +import matplotlib.pyplot as plt + + +DEFAULT_SIM_DIR = ( + "/global/cfs/cdirs/m4814/daweigao/14_Code/0_dataset_construction/" + "3_restarted simulation" +) +DEFAULT_GT_PATH = ( + "/global/cfs/cdirs/m4814/daweigao/14_Code/0_dataset_construction/" + "20251201_TRENDY2024_default_ICB1850CNPRDCTCBC.elm.h0.0801-01-01-00000.nc" +) + + +def _to_nan_fillvalue(arr, fill_threshold=1e35): + data = np.asarray(arr, dtype=float) + data[np.abs(data) >= fill_threshold] = np.nan + return data + + +def _apply_ijcai_style(): + plt.rcParams.update({ + "font.family": "serif", + "font.serif": ["Times New Roman", "Times", "DejaVu Serif"], + "font.size": 11, + "axes.labelsize": 12, + "axes.titlesize": 12, + "axes.linewidth": 1.0, + "xtick.direction": "in", + "ytick.direction": "in", + "xtick.major.size": 4, + "ytick.major.size": 4, + "xtick.minor.size": 2, + "ytick.minor.size": 2, + "legend.frameon": False, + }) + + +def _extract_year(path: Path): + match = re.search(r"\.h0\.(\d{4})-", path.name) + if not match: + return None + return int(match.group(1)) + + +def _sum_variable(ds, variable): + if variable not in ds: + raise KeyError(f"Variable {variable} not found in {ds.encoding.get('source', 'dataset')}.") + data = _to_nan_fillvalue(ds[variable].values) + return float(np.nansum(data)) + + +def _collect_simulation_series(sim_dir: Path, variable: str): + nc_files = sorted(sim_dir.glob("*.h0.*.nc")) + years = [] + sums = [] + for path in nc_files: + year = _extract_year(path) + if year is None: + continue + with xr.open_dataset(path) as ds: + value = _sum_variable(ds, variable) + years.append(year) + sums.append(value) + + if not years: + raise FileNotFoundError(f"No .h0.*.nc files with year found in {sim_dir}") + + order = np.argsort(years) + years = np.asarray(years)[order] + sums = np.asarray(sums)[order] + return years, sums + + +def plot_timeseries(sim_dir, gt_path, variable="NEE", output=None): + sim_dir = Path(sim_dir) + gt_path = Path(gt_path) + if output: + output = Path(output) + else: + output = Path(__file__).resolve().parent / f"{variable.lower()}_timeseries_ijcai.png" + + years, sim_sums = _collect_simulation_series(sim_dir, variable) + with xr.open_dataset(gt_path) as ds_gt: + gt_sum = _sum_variable(ds_gt, variable) + + _apply_ijcai_style() + fig, ax = plt.subplots(figsize=(6.4, 3.6)) + + ax.plot(years, sim_sums, color="#1f77b4", linewidth=2.0, label="Simulation") + ax.hlines(gt_sum, years.min(), years.max(), colors="#d62728", linestyles="--", linewidth=2.0, label="Ground Truth") + + ax.set_xlabel("Year") + ax.set_ylabel(f"{variable} Sum") + ax.grid(True, linestyle="--", linewidth=0.5, alpha=0.4) + ax.legend() + + fig.tight_layout() + output.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output, dpi=300) + pdf_output = output.with_suffix(".pdf") + fig.savefig(pdf_output) + plt.close(fig) + + print(f"Saved time series plot: {output}") + print(f"Saved time series plot: {pdf_output}") + + +def parse_args(): + parser = argparse.ArgumentParser(description="Plot IJCAI-style time series of NEE sum.") + parser.add_argument("--sim-dir", type=str, default=DEFAULT_SIM_DIR, help="Directory of simulation .h0.*.nc files.") + parser.add_argument("--gt-path", type=str, default=DEFAULT_GT_PATH, help="Ground-truth NetCDF path.") + parser.add_argument("--variable", type=str, default="NEE", help="Variable name to plot.") + parser.add_argument("--output", type=str, default=None, help="Output figure path.") + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + plot_timeseries( + sim_dir=args.sim_dir, + gt_path=args.gt_path, + variable=args.variable, + output=args.output, + ) From 29f36ccfd31c732ab76d41df786fe27e5c328df7 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Fri, 23 Jan 2026 15:24:06 -0800 Subject: [PATCH 49/51] verify CNP predictions, fix ai_predictions_to_netcdf workflow, and highlight a grid cell in scatter plots --- docs/CNP_pipeline_runbook.md | 20 + scripts/ai_predictions_to_netcdf.py | 108 +-- scripts/cnp_result_validationplot_site.py | 809 ++++++++++++++++++++++ scripts/verify_predictions_in_netcdf.py | 226 ++++++ 4 files changed, 1118 insertions(+), 45 deletions(-) create mode 100644 scripts/cnp_result_validationplot_site.py create mode 100644 scripts/verify_predictions_in_netcdf.py diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index ba0d7dd..3122f07 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -54,6 +54,22 @@ python ../../scripts/cnp_result_validationplot.py --stats-only python ../../scripts/generate_prediction_quality_report.py ``` +### 4.1) Highlight a site in all scatter plots (optional) +Highlight a given site (lon/lat) in every GT vs Pred scatter plot. The site will be marked as a red star. + +```bash +python ../../scripts/cnp_result_validationplot_site.py . \ + --lon 303.75 --lat -17.4246 +``` + +Options and behavior: +- The script matches samples by coordinate with a tolerance (default: 0.01 degrees). +- Use `--tolerance N` to relax the matching if no point is found. +- Use `--coord-csv /path/to/test_static_inverse.csv` to explicitly provide a coordinate source. +- Supports `--top-bad-only` / `--worst-only` for restricted plotting. + +**Output**: scatter plots are saved under `plots_site/`. + Options and behavior: - The quality report saves outputs under `analysis/`: - `detailed_quality_assessment.csv`, `variable_quality_summary.csv`, `overall_prediction_quality.png`, `prediction_quality_by_variable.png`, `r2_vs_rmse.png`, `prediction_quality_report.html`, `quality_summary_report.txt`. @@ -87,6 +103,10 @@ Creates a NetCDF file containing all AI predictions for plotting and comparison: ```bash python ../../scripts/ai_predictions_to_netcdf.py > ai_prediction_to_netcdf.log 2>&1 & ``` +To verify the CSV predictions match the generated NetCDF: +```bash +python ../../scripts/verify_predictions_in_netcdf.py +``` ### 7) Generate comparison plots Creates map plots for the variables in your CNP_IO list under diff --git a/scripts/ai_predictions_to_netcdf.py b/scripts/ai_predictions_to_netcdf.py index aab8c2a..7e93a12 100644 --- a/scripts/ai_predictions_to_netcdf.py +++ b/scripts/ai_predictions_to_netcdf.py @@ -242,18 +242,31 @@ def load_ai_predictions(predictions_dir: Path) -> Dict[str, Any]: return preds -def create_netcdf_structure(ai_preds: Dict[str, Any], variable_list: Dict[str, List[str]], +def _select_base_coords(ai_preds: Dict[str, Any]) -> tuple[np.ndarray, np.ndarray]: + """Select base coordinates from CSV-derived predictions.""" + # Prefer scalar coords if available (should be full grid) + if 'scalar_coords' in ai_preds and ai_preds['scalar_coords'][0] is not None: + return ai_preds['scalar_coords'] + # Fall back to any soil2d/pft1d coords + if 'soil2d_coords' in ai_preds and ai_preds['soil2d_coords']: + _, coords = next(iter(ai_preds['soil2d_coords'].items())) + if coords[0] is not None: + return coords + if 'pft1d_coords' in ai_preds and ai_preds['pft1d_coords']: + _, coords = next(iter(ai_preds['pft1d_coords'].items())) + if coords[0] is not None: + return coords + raise ValueError("No CSV coordinates found in predictions (Longitude/Latitude columns missing?)") + + +def create_netcdf_structure(ai_preds: Dict[str, Any], variable_list: Dict[str, List[str]], output_path: Path) -> xr.Dataset: - """Create NetCDF structure compatible with restart_variable_plot.py""" + """Create NetCDF structure compatible with restart_variable_plot.py.""" print("Creating NetCDF structure...") - - # Get coordinates from static inverse mapping - if 'test_static_inverse' not in ai_preds: - raise ValueError("test_static_inverse.csv not found - needed for coordinates") - - static_df = ai_preds['test_static_inverse'] - n_samples = len(static_df) - + + lon, lat = _select_base_coords(ai_preds) + n_samples = len(lon) + # Create coordinate variables coords = { 'gridcell': np.arange(n_samples), @@ -261,34 +274,34 @@ def create_netcdf_structure(ai_preds: Dict[str, Any], variable_list: Dict[str, L 'column': np.arange(1), # Only first column 'levgrnd': np.arange(10), # Only first 10 layers } - + # Create the dataset ds = xr.Dataset(coords=coords) - + # Add coordinate variables that restart_variable_plot.py expects # Grid coordinates: 1D arrays for gridcell dimension - ds['grid1d_lon'] = xr.DataArray(static_df['Longitude'].values, dims=['gridcell']) - ds['grid1d_lat'] = xr.DataArray(static_df['Latitude'].values, dims=['gridcell']) - + ds['grid1d_lon'] = xr.DataArray(lon, dims=['gridcell']) + ds['grid1d_lat'] = xr.DataArray(lat, dims=['gridcell']) + # PFT coordinates: For PFT variables, we need to create proper mapping # Each PFT gets assigned to the first gridcell (index 1, 1-based) - pft_lon = np.full(16, static_df['Longitude'].iloc[0], dtype=float) - pft_lat = np.full(16, static_df['Latitude'].iloc[0], dtype=float) + pft_lon = np.full(16, float(lon[0]), dtype=float) + pft_lat = np.full(16, float(lat[0]), dtype=float) ds['pfts1d_lon'] = xr.DataArray(pft_lon, dims=['pft']) ds['pfts1d_lat'] = xr.DataArray(pft_lat, dims=['pft']) - + # Column coordinates: For column variables, we need to create proper mapping # Each column gets assigned to the first gridcell (index 1, 1-based) - col_lon = np.full(1, static_df['Longitude'].iloc[0], dtype=float) - col_lat = np.full(1, static_df['Latitude'].iloc[0], dtype=float) + col_lon = np.full(1, float(lon[0]), dtype=float) + col_lat = np.full(1, float(lat[0]), dtype=float) ds['cols1d_lon'] = xr.DataArray(col_lon, dims=['column']) ds['cols1d_lat'] = xr.DataArray(col_lat, dims=['column']) - + # Add gridcell indices (1-based as expected by restart_variable_plot.py) # All PFTs and columns are assigned to gridcell 1 ds['pfts1d_gridcell_index'] = xr.DataArray(np.ones(16, dtype=int), dims=['pft']) ds['cols1d_gridcell_index'] = xr.DataArray(np.ones(1, dtype=int), dims=['column']) - + print(f" Created base structure with {n_samples} gridcells") print(f" PFT coordinates: {pft_lon.shape} (all assigned to first gridcell)") print(f" Column coordinates: {col_lon.shape} (all assigned to first gridcell)") @@ -426,6 +439,12 @@ def _check_coords_match(coords1, coords2, label1, label2): raise ValueError(f"Coordinate mismatch between {label1} and {label2}") +def _wrap_coords(lon: np.ndarray) -> np.ndarray: + """Wrap longitudes from 0-360 to -180-180.""" + lon = np.asarray(lon, dtype=float) + return ((lon + 180.0) % 360.0) - 180.0 + + def main(): parser = argparse.ArgumentParser( description='Convert AI predictions to NetCDF format compatible with restart_variable_plot.py' @@ -494,32 +513,31 @@ def main(): # Load AI predictions ai_preds = load_ai_predictions(ai_predictions_dir) - # Report current lon/lat ranges from static inverse if available - if 'test_static_inverse' in ai_preds: - try: - _sdf = ai_preds['test_static_inverse'] - if {'Longitude','Latitude'}.issubset(_sdf.columns): - lon_min, lon_max = float(pd.to_numeric(_sdf['Longitude'], errors='coerce').min()), float(pd.to_numeric(_sdf['Longitude'], errors='coerce').max()) - lat_min, lat_max = float(pd.to_numeric(_sdf['Latitude'], errors='coerce').min()), float(pd.to_numeric(_sdf['Latitude'], errors='coerce').max()) - print(f" Before wrapping - grid1d_lon min/max: {lon_min}, {lon_max}") - print(f" Before wrapping - grid1d_lat min/max: {lat_min}, {lat_max}") - except Exception as _e: - print(f" Warning: Failed to compute pre-wrap lon/lat ranges: {_e}") + # Report current lon/lat ranges from CSV coordinates + try: + base_lon, base_lat = _select_base_coords(ai_preds) + print(f" Before wrapping - grid1d_lon min/max: {float(np.min(base_lon))}, {float(np.max(base_lon))}") + print(f" Before wrapping - grid1d_lat min/max: {float(np.min(base_lat))}, {float(np.max(base_lat))}") + except Exception as _e: + print(f" Warning: Failed to compute pre-wrap lon/lat ranges: {_e}") # Optional longitude wrapping 0–360 -> -180–180 - if getattr(args, 'wrap_longitude', False) and 'test_static_inverse' in ai_preds: + if getattr(args, 'wrap_longitude', False): try: - static_df = ai_preds['test_static_inverse'] - if 'Longitude' in static_df.columns: - lon = pd.to_numeric(static_df['Longitude'], errors='coerce') - wrapped = ((lon + 180.0) % 360.0) - 180.0 - static_df['Longitude'] = wrapped - ai_preds['test_static_inverse'] = static_df - lon_min, lon_max = float(wrapped.min()), float(wrapped.max()) - print(f" After wrapping - grid1d_lon min/max: {lon_min}, {lon_max}") - if 'Latitude' in static_df.columns: - lat = pd.to_numeric(static_df['Latitude'], errors='coerce') - print(f" Latitude min/max: {float(lat.min())}, {float(lat.max())}") + if 'scalar_coords' in ai_preds and ai_preds['scalar_coords'][0] is not None: + lon, lat = ai_preds['scalar_coords'] + ai_preds['scalar_coords'] = (_wrap_coords(lon), lat) + if 'pft1d_coords' in ai_preds: + for k, coords in ai_preds['pft1d_coords'].items(): + if coords[0] is not None: + ai_preds['pft1d_coords'][k] = (_wrap_coords(coords[0]), coords[1]) + if 'soil2d_coords' in ai_preds: + for k, coords in ai_preds['soil2d_coords'].items(): + if coords[0] is not None: + ai_preds['soil2d_coords'][k] = (_wrap_coords(coords[0]), coords[1]) + base_lon, base_lat = _select_base_coords(ai_preds) + print(f" After wrapping - grid1d_lon min/max: {float(np.min(base_lon))}, {float(np.max(base_lon))}") + print(f" Latitude min/max: {float(np.min(base_lat))}, {float(np.max(base_lat))}") except Exception as _e: print(f" Warning: Failed to wrap longitudes: {_e}") diff --git a/scripts/cnp_result_validationplot_site.py b/scripts/cnp_result_validationplot_site.py new file mode 100644 index 0000000..a53383e --- /dev/null +++ b/scripts/cnp_result_validationplot_site.py @@ -0,0 +1,809 @@ +#!/usr/bin/env python3 +""" +Generate validation scatter plots and highlight a site-specific sample +based on a given longitude/latitude. +""" + +import os +import re +import json +from glob import glob +import argparse +from typing import Optional, List, Tuple + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score + + +def _find_lat_lon_columns(df: pd.DataFrame) -> Optional[Tuple[str, str]]: + """Return (lon_col, lat_col) if found, else None.""" + lower_map = {c.lower(): c for c in df.columns} + # Prefer full names + if "longitude" in lower_map and "latitude" in lower_map: + return lower_map["longitude"], lower_map["latitude"] + # Common short names + if "lon" in lower_map and "lat" in lower_map: + return lower_map["lon"], lower_map["lat"] + if "long" in lower_map and "lat" in lower_map: + return lower_map["long"], lower_map["lat"] + # Fallback: any column containing 'lon' and 'lat' + lon_candidates = [c for c in df.columns if "lon" in c.lower()] + lat_candidates = [c for c in df.columns if "lat" in c.lower()] + if lon_candidates and lat_candidates: + return lon_candidates[0], lat_candidates[0] + return None + + +def _find_site_indices_in_df(df: pd.DataFrame, lon: float, lat: float, tol: float) -> List[int]: + cols = _find_lat_lon_columns(df) + if cols is None: + return [] + lon_col, lat_col = cols + lon_vals = pd.to_numeric(df[lon_col], errors="coerce") + lat_vals = pd.to_numeric(df[lat_col], errors="coerce") + mask = (np.abs(lon_vals - lon) <= tol) & (np.abs(lat_vals - lat) <= tol) + return np.where(mask)[0].tolist() + + +def _find_site_indices(results_dir: str, lon: float, lat: float, tol: float, coord_csv: Optional[str]) -> List[int]: + candidates = [] + if coord_csv: + candidates.append(coord_csv) + else: + candidates.append(os.path.join(results_dir, "cnp_predictions", "test_static_inverse.csv")) + candidates.append(os.path.join(results_dir, "cnp_predictions", "ground_truth_scalar.csv")) + candidates.append(os.path.join(results_dir, "cnp_predictions", "predictions_scalar.csv")) + # Fallback to any 1D/2D file if needed + pft_dir = os.path.join(results_dir, "cnp_predictions", "pft_1d_ground_truth") + soil_dir = os.path.join(results_dir, "cnp_predictions", "soil_2d_ground_truth") + if os.path.isdir(pft_dir): + pft_files = glob(os.path.join(pft_dir, "ground_truth_Y_*.csv")) + if pft_files: + candidates.append(pft_files[0]) + if os.path.isdir(soil_dir): + soil_files = glob(os.path.join(soil_dir, "ground_truth_Y_*.csv")) + if soil_files: + candidates.append(soil_files[0]) + + for path in candidates: + if path and os.path.exists(path): + try: + df = pd.read_csv(path) + indices = _find_site_indices_in_df(df, lon, lat, tol) + if indices: + print(f"Found {len(indices)} matching samples using: {path}") + return indices + except Exception as e: + print(f"Warning: failed to read {path}: {e}") + + return [] + + +def _extract_site_points(gt_col: np.ndarray, pred_col: np.ndarray, site_indices: List[int]) -> Optional[np.ndarray]: + if not site_indices: + return None + valid_indices = [i for i in site_indices if 0 <= i < len(gt_col)] + if not valid_indices: + return None + site_gt = gt_col[valid_indices] + site_pred = pred_col[valid_indices] + valid_mask = ~(np.isnan(site_gt) | np.isnan(site_pred)) + if not np.any(valid_mask): + return None + return np.column_stack([site_gt[valid_mask], site_pred[valid_mask]]) + + +def plot_gt_vs_pred(gt, pred, title, save_path, site_points=None, site_label=None): + plt.figure(figsize=(6, 6)) + plt.scatter(gt, pred, alpha=0.5, s=20, color="#1f77b4") + plt.plot([gt.min(), gt.max()], [gt.min(), gt.max()], "r--") + + if site_points is not None and len(site_points) > 0: + plt.scatter( + site_points[:, 0], + site_points[:, 1], + marker="*", + s=140, + c="red", + edgecolors="black", + linewidths=0.6, + label=site_label or "site", + zorder=5, + ) + plt.legend() + + plt.xlabel("Ground Truth") + plt.ylabel("Prediction") + plt.title(title) + plt.tight_layout() + os.makedirs(os.path.dirname(save_path), exist_ok=True) + plt.savefig(save_path) + plt.close() + + +def analyze_pair(gt_path, pred_path, label, out_dir, stats_data, site_indices, site_label, per_column=False, plot_scatter=True, selection=None): + gt = pd.read_csv(gt_path) + pred = pd.read_csv(pred_path) + if per_column: + for col in gt.columns: + col_norm = col[2:] if isinstance(col, str) and col.startswith("Y_") else col + if selection is not None and str(col_norm) in ("Latitude", "Longitude"): + continue + if selection is not None and col_norm not in selection: + continue + if col in pred.columns: + print(f"Analyzing variable: {col}") + gt_col = gt[col].values.flatten() + pred_col = pred[col].values.flatten() + + gt_stats = {"min": np.nanmin(gt_col), "max": np.nanmax(gt_col), "sum": np.nansum(gt_col)} + pred_stats = {"min": np.nanmin(pred_col), "max": np.nanmax(pred_col), "sum": np.nansum(pred_col)} + + rmse = np.sqrt(mean_squared_error(gt_col, pred_col)) + mae = mean_absolute_error(gt_col, pred_col) + r2 = r2_score(gt_col, pred_col) + print(f" {label} - {col}: RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}") + + if plot_scatter: + site_points = _extract_site_points(gt_col, pred_col, site_indices) + plot_gt_vs_pred( + gt_col, + pred_col, + f"{label} {col} GT vs Pred", + os.path.join(out_dir, f"{label}_{col}_gt_vs_pred.png"), + site_points=site_points, + site_label=site_label, + ) + + stats_data.append( + { + "type": label, + "variable": col, + "rmse": rmse, + "mae": mae, + "r2": r2, + "gt_min": gt_stats["min"], + "gt_max": gt_stats["max"], + "gt_sum": gt_stats["sum"], + "pred_min": pred_stats["min"], + "pred_max": pred_stats["max"], + "pred_sum": pred_stats["sum"], + } + ) + else: + print(f"Column {col} missing in predictions for {label}") + return None + + gt_flat = gt.values.flatten() + pred_flat = pred.values.flatten() + gt_stats = {"min": np.nanmin(gt_flat), "max": np.nanmax(gt_flat), "sum": np.nansum(gt_flat)} + pred_stats = {"min": np.nanmin(pred_flat), "max": np.nanmax(pred_flat), "sum": np.nansum(pred_flat)} + rmse = np.sqrt(mean_squared_error(gt_flat, pred_flat)) + mae = mean_absolute_error(gt_flat, pred_flat) + r2 = r2_score(gt_flat, pred_flat) + print(f"{label} - RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}") + + if plot_scatter: + site_points = _extract_site_points(gt_flat, pred_flat, site_indices) + plot_gt_vs_pred( + gt_flat, + pred_flat, + f"{label} GT vs Pred", + os.path.join(out_dir, f"{label}_gt_vs_pred.png"), + site_points=site_points, + site_label=site_label, + ) + + stats_data.append( + { + "type": label, + "variable": "all", + "rmse": rmse, + "mae": mae, + "r2": r2, + "gt_min": gt_stats["min"], + "gt_max": gt_stats["max"], + "gt_sum": gt_stats["sum"], + "pred_min": pred_stats["min"], + "pred_max": pred_stats["max"], + "pred_sum": pred_stats["sum"], + } + ) + return {"rmse": rmse, "mae": mae, "r2": r2} + + +def analyze_1d_new_structure(results_dir, label, out_dir, stats_data, site_indices, site_label, plot_scatter=True, selection=None): + 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") + if not os.path.exists(gt_dir) or not os.path.exists(pred_dir): + print(f"Missing 1D directories: {gt_dir} or {pred_dir}") + return + gt_files = glob(os.path.join(gt_dir, "ground_truth_Y_*.csv")) + if not gt_files: + print(f"No ground truth files found in {gt_dir}") + return + print(f"Found {len(gt_files)} 1D variables to analyze") + + for gt_file in gt_files: + var_name = os.path.basename(gt_file).replace("ground_truth_Y_", "").replace(".csv", "") + pred_file = os.path.join(pred_dir, f"predictions_Y_{var_name}.csv") + if not os.path.exists(pred_file): + print(f"Missing prediction file for {var_name}: {pred_file}") + continue + if selection is not None and var_name not in selection: + continue + + print(f"Analyzing variable: {var_name}") + gt_data = pd.read_csv(gt_file) + pred_data = pd.read_csv(pred_file) + for col in ["long", "lat", "Long", "Lat", "Longitude", "Latitude"]: + if col in gt_data.columns: + gt_data = gt_data.drop(columns=[col]) + if col in pred_data.columns: + pred_data = pred_data.drop(columns=[col]) + if gt_data.shape != pred_data.shape: + print(f"Shape mismatch for {var_name}: GT {gt_data.shape} vs Pred {pred_data.shape}") + continue + + num_pfts = gt_data.shape[1] + print(f" {var_name}: {num_pfts} PFT columns") + + for col_name in gt_data.columns: + if selection is not None: + sel = selection.get(var_name, None) + if sel is not None and sel["pfts"]: + pft_match = re.search(r"pft(\d+)$", col_name) + if pft_match: + try: + pft_num = int(pft_match.group(1)) + if pft_num not in sel["pfts"]: + continue + except Exception: + pass + + gt_col = gt_data[col_name].values + pred_col = pred_data[col_name].values + if np.all(np.isnan(gt_col)) or np.all(np.isnan(pred_col)): + continue + valid_mask = ~(np.isnan(gt_col) | np.isnan(pred_col)) + if np.sum(valid_mask) < 3: + continue + gt_valid = gt_col[valid_mask] + pred_valid = pred_col[valid_mask] + + gt_stats = {"min": np.nanmin(gt_valid), "max": np.nanmax(gt_valid), "sum": np.nansum(gt_valid)} + pred_stats = {"min": np.nanmin(pred_valid), "max": np.nanmax(pred_valid), "sum": np.nansum(pred_valid)} + rmse = np.sqrt(mean_squared_error(gt_valid, pred_valid)) + mae = mean_absolute_error(gt_valid, pred_valid) + r2 = r2_score(gt_valid, pred_valid) + + print(f" {col_name}: RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}") + + if plot_scatter: + site_points = _extract_site_points(gt_col, pred_col, site_indices) + plot_gt_vs_pred( + gt_valid, + pred_valid, + f"{label} {col_name} GT vs Pred", + os.path.join(out_dir, f"{label}_{col_name}_gt_vs_pred.png"), + site_points=site_points, + site_label=site_label, + ) + + stats_data.append( + { + "type": "1D", + "variable": var_name, + "pft": col_name, + "rmse": rmse, + "mae": mae, + "r2": r2, + "gt_min": gt_stats["min"], + "gt_max": gt_stats["max"], + "gt_sum": gt_stats["sum"], + "pred_min": pred_stats["min"], + "pred_max": pred_stats["max"], + "pred_sum": pred_stats["sum"], + } + ) + + +def analyze_1d(gt_path, pred_path, label, out_dir, results_dir, stats_data, site_indices, site_label, plot_scatter=True, selection=None): + gt = pd.read_csv(gt_path) + pred = pd.read_csv(pred_path) + for col in ["long", "lat", "Long", "Lat", "Longitude", "Latitude"]: + if col in gt.columns: + gt = gt.drop(columns=[col]) + if col in pred.columns: + pred = pred.drop(columns=[col]) + config_path = os.path.join(results_dir, "cnp_config.json") + if os.path.exists(config_path): + with open(config_path, "r") as f: + config = json.load(f) + variable_names = config.get("data_info", {}).get("variables_1d_pft", None) + if variable_names is None: + print("[ERROR] Could not find variables_1d_pft in cnp_config.json!") + return + else: + print(f"[ERROR] cnp_config.json not found in {results_dir}!") + return + num_vars = len(variable_names) + num_pfts = 16 + num_samples = gt.shape[0] + expected_cols = num_vars * num_pfts + print(f"[DEBUG] 1D: num_samples={num_samples}, num_vars={num_vars}, num_pfts={num_pfts}, expected_cols={expected_cols}, actual_cols={gt.shape[1]}") + if gt.shape[1] != expected_cols: + print("[ERROR] Unexpected number of columns in 1D data!") + print("Column names:", list(gt.columns)) + return + gt_reshaped = gt.values.reshape(num_samples, num_vars, num_pfts) + pred_reshaped = pred.values.reshape(num_samples, num_vars, num_pfts) + for i, var in enumerate(variable_names): + if selection is not None and var not in selection: + continue + print(f"Analyzing variable: {var}") + for j in range(num_pfts): + gt_col = gt_reshaped[:, i, j] + pred_col = pred_reshaped[:, i, j] + col_name = gt.columns[i * num_pfts + j] if (i * num_pfts + j) < len(gt.columns) else f"{var}_pft{j+1}" + gt_stats = {"min": np.nanmin(gt_col), "max": np.nanmax(gt_col), "sum": np.nansum(gt_col)} + pred_stats = {"min": np.nanmin(pred_col), "max": np.nanmax(pred_col), "sum": np.nansum(pred_col)} + rmse = np.sqrt(mean_squared_error(gt_col, pred_col)) + mae = mean_absolute_error(gt_col, pred_col) + r2 = r2_score(gt_col, pred_col) + print(f"{label} - {col_name}: RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}") + + if plot_scatter: + site_points = _extract_site_points(gt_col, pred_col, site_indices) + plot_gt_vs_pred( + gt_col, + pred_col, + f"{label} {col_name} GT vs Pred", + os.path.join(out_dir, f"{label}_{col_name}_gt_vs_pred.png"), + site_points=site_points, + site_label=site_label, + ) + + stats_data.append( + { + "type": "1D", + "variable": var, + "pft": col_name, + "rmse": rmse, + "mae": mae, + "r2": r2, + "gt_min": gt_stats["min"], + "gt_max": gt_stats["max"], + "gt_sum": gt_stats["sum"], + "pred_min": pred_stats["min"], + "pred_max": pred_stats["max"], + "pred_sum": pred_stats["sum"], + } + ) + + +def analyze_2d_new_structure(results_dir, label, out_dir, stats_data, site_indices, site_label, plot_scatter=True, selection=None): + 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") + if not os.path.exists(gt_dir) or not os.path.exists(pred_dir): + print(f"Missing 2D directories: {gt_dir} or {pred_dir}") + return + gt_files = glob(os.path.join(gt_dir, "ground_truth_Y_*.csv")) + if not gt_files: + print(f"No 2D ground truth files found in {gt_dir}") + return + print(f"Found {len(gt_files)} 2D variables to analyze") + + for gt_file in gt_files: + var_name = os.path.basename(gt_file).replace("ground_truth_Y_", "").replace(".csv", "") + pred_file = os.path.join(pred_dir, f"predictions_Y_{var_name}.csv") + if not os.path.exists(pred_file): + print(f"Missing prediction file for {var_name}: {pred_file}") + continue + if selection is not None and var_name not in selection: + continue + + print(f"Analyzing 2D variable: {var_name}") + gt_data = pd.read_csv(gt_file) + pred_data = pd.read_csv(pred_file) + for col in ["long", "lat", "Long", "Lat", "Longitude", "Latitude"]: + if col in gt_data.columns: + gt_data = gt_data.drop(columns=[col]) + if col in pred_data.columns: + pred_data = pred_data.drop(columns=[col]) + if gt_data.shape != pred_data.shape: + print(f"Shape mismatch for {var_name}: GT {gt_data.shape} vs Pred {pred_data.shape}") + continue + + total_columns = gt_data.shape[1] + expected_columns = 1 * 10 + if total_columns != expected_columns: + print(f" Warning: Expected {expected_columns} columns for 2D data, but found {total_columns}") + if total_columns % 10 != 0: + print(f" Error: Number of columns ({total_columns}) is not divisible by 10") + continue + num_columns = total_columns // 10 + print(f" Assuming {num_columns} columns with 10 layers each") + else: + num_columns = 1 + print(f" {var_name}: {num_columns} columns, each with 10 layers ({total_columns} total columns)") + + first_column_idx = 0 + layers_to_analyze = 10 + + for layer_idx in range(layers_to_analyze): + if selection is not None: + sel = selection.get(var_name, None) + if sel is not None and sel["layers"] and (layer_idx + 1) not in sel["layers"]: + continue + col_idx = layer_idx if num_columns == 1 else (first_column_idx * 10 + layer_idx) + if col_idx >= total_columns: + print(f" Warning: Column index {col_idx} out of range for {total_columns} columns") + continue + + gt_col = gt_data.iloc[:, col_idx].values + pred_col = pred_data.iloc[:, col_idx].values + if np.all(np.isnan(gt_col)) or np.all(np.isnan(pred_col)): + continue + valid_mask = ~(np.isnan(gt_col) | np.isnan(pred_col)) + if np.sum(valid_mask) < 3: + continue + + gt_valid = gt_col[valid_mask] + pred_valid = pred_col[valid_mask] + gt_stats = {"min": np.nanmin(gt_valid), "max": np.nanmax(gt_valid), "sum": np.nansum(gt_valid)} + pred_stats = {"min": np.nanmin(pred_valid), "max": np.nanmax(pred_valid), "sum": np.nansum(pred_valid)} + rmse = np.sqrt(mean_squared_error(gt_valid, pred_valid)) + mae = mean_absolute_error(gt_valid, pred_valid) + r2 = r2_score(gt_valid, pred_valid) + + print(f" Layer {layer_idx+1}: RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}") + + if plot_scatter: + site_points = _extract_site_points(gt_col, pred_col, site_indices) + plot_gt_vs_pred( + gt_valid, + pred_valid, + f"{label} {var_name} Layer{layer_idx+1} GT vs Pred", + os.path.join(out_dir, f"{label}_{var_name}_Layer{layer_idx+1}_gt_vs_pred.png"), + site_points=site_points, + site_label=site_label, + ) + + stats_data.append( + { + "type": "2D", + "variable": var_name, + "layer": layer_idx + 1, + "rmse": rmse, + "mae": mae, + "r2": r2, + "gt_min": gt_stats["min"], + "gt_max": gt_stats["max"], + "gt_sum": gt_stats["sum"], + "pred_min": pred_stats["min"], + "pred_max": pred_stats["max"], + "pred_sum": pred_stats["sum"], + } + ) + + try: + gt_firstcol = gt_data.iloc[:, 0:10].values + pred_firstcol = pred_data.iloc[:, 0:10].values + valid_mask = ~(np.isnan(gt_firstcol) | np.isnan(pred_firstcol)) + if np.sum(valid_mask) >= 3 and plot_scatter: + gt_flat = gt_firstcol.flatten() + pred_flat = pred_firstcol.flatten() + site_vals = [] + for idx in site_indices: + if 0 <= idx < gt_firstcol.shape[0]: + site_vals.append((gt_firstcol[idx, :], pred_firstcol[idx, :])) + if site_vals: + site_gt = np.concatenate([v[0] for v in site_vals]) + site_pred = np.concatenate([v[1] for v in site_vals]) + site_points = _extract_site_points(site_gt, site_pred, list(range(len(site_gt)))) + else: + site_points = None + + plot_gt_vs_pred( + gt_flat[valid_mask.flatten()], + pred_flat[valid_mask.flatten()], + f"{label} {var_name} FirstCol(10 layers) GT vs Pred", + os.path.join(out_dir, f"{label}_{var_name}_FirstCol_AllLayers_gt_vs_pred.png"), + site_points=site_points, + site_label=site_label, + ) + except Exception as e: + print(f" Skipped overall plot for {var_name}: {e}") + + +def analyze_2d(gt_path, pred_path, label, out_dir, results_dir, stats_data, site_indices, site_label, plot_scatter=True, selection=None): + gt = pd.read_csv(gt_path) + pred = pd.read_csv(pred_path) + config_path = os.path.join(results_dir, "cnp_config.json") + if os.path.exists(config_path): + with open(config_path, "r") as f: + config = json.load(f) + variable_names = config.get("data_info", {}).get("variables_2d_soil", None) + if variable_names is None: + print("[ERROR] Could not find variables_2d_soil in cnp_config.json!") + return + else: + print(f"[ERROR] cnp_config.json not found in {results_dir}!") + return + num_vars = len(variable_names) + num_columns = 18 + num_layers_per_column = 10 + num_samples = gt.shape[0] + expected_cols = num_vars * num_columns * num_layers_per_column + print(f"[DEBUG] 2D: num_samples={num_samples}, num_vars={num_vars}, num_columns={num_columns}, layers_per_column={num_layers_per_column}, expected_cols={expected_cols}, actual_cols={gt.shape[1]}") + if gt.shape[1] != expected_cols: + print("[ERROR] Unexpected number of columns in 2D data!") + print("Column names:", list(gt.columns)) + return + gt_reshaped = gt.values.reshape(num_samples, num_vars, num_columns, num_layers_per_column) + pred_reshaped = pred.values.reshape(num_samples, num_vars, num_columns, num_layers_per_column) + for i, var in enumerate(variable_names): + if selection is not None and var not in selection: + continue + print(f"Analyzing 2D variable: {var}") + for j in range(10): + if selection is not None: + sel = selection.get(var, None) + if sel is not None and sel["layers"] and (j + 1) not in sel["layers"]: + continue + gt_col = gt_reshaped[:, i, 0, j] + pred_col = pred_reshaped[:, i, 0, j] + + gt_stats = {"min": np.nanmin(gt_col), "max": np.nanmax(gt_col), "sum": np.nansum(gt_col)} + pred_stats = {"min": np.nanmin(pred_col), "max": np.nanmax(pred_col), "sum": np.nansum(pred_col)} + rmse = np.sqrt(mean_squared_error(gt_col, pred_col)) + mae = mean_absolute_error(gt_col, pred_col) + r2 = r2_score(gt_col, pred_col) + print(f"{label} - {var} (Layer {j+1}): RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}") + + if plot_scatter: + site_points = _extract_site_points(gt_col, pred_col, site_indices) + plot_gt_vs_pred( + gt_col, + pred_col, + f"{label} {var} Layer{j+1} GT vs Pred", + os.path.join(out_dir, f"{label}_{var}_Layer{j+1}_gt_vs_pred.png"), + site_points=site_points, + site_label=site_label, + ) + + stats_data.append( + { + "type": "2D", + "variable": var, + "layer": j + 1, + "rmse": rmse, + "mae": mae, + "r2": r2, + "gt_min": gt_stats["min"], + "gt_max": gt_stats["max"], + "gt_sum": gt_stats["sum"], + "pred_min": pred_stats["min"], + "pred_max": pred_stats["max"], + "pred_sum": pred_stats["sum"], + } + ) + + +def plot_train_val_accuracy(loss_csv, out_dir): + df = pd.read_csv(loss_csv) + plt.figure() + if "Train Loss" in df.columns and "Validation Loss" in df.columns: + plt.plot(df["Train Loss"], label="Train Loss") + plt.plot(df["Validation Loss"], label="Validation Loss") + plt.ylabel("Loss") + plt.xlabel("Epoch") + plt.legend() + plt.title("Train/Validation Loss") + plt.tight_layout() + os.makedirs(out_dir, exist_ok=True) + plt.savefig(os.path.join(out_dir, "train_val_loss.png")) + plt.close() + else: + print("train_loss or val_loss columns not found in loss CSV.") + + +def _parse_top_bad_report(report_path): + selection = {} + if not os.path.exists(report_path): + print(f"Top-bad report not found: {report_path}") + return selection + in_section = False + try: + with open(report_path, "r") as f: + for line in f: + stripped = line.strip("\n") + header = stripped.strip() + if header.startswith("Top variables by bad-count"): + in_section = True + continue + if in_section and header.startswith("## "): + break + if in_section and stripped.startswith(" "): + m = re.match(r"\s+([A-Za-z0-9_]+):\s*([0-9]+)(;.*)?$", stripped) + if not m: + continue + var = m.group(1) + details = m.group(3) or "" + pfts = set() + layers = set() + if "pfts:" in details: + m_p = re.search(r"pfts:\s*([0-9,\s]+)", details) + if m_p: + nums = [n.strip() for n in m_p.group(1).split(",") if n.strip()] + for n in nums: + try: + pfts.add(int(n)) + except Exception: + pass + if "layers:" in details: + m_l = re.search(r"layers:\s*([0-9,\s]+)", details) + if m_l: + nums = [n.strip() for n in m_l.group(1).split(",") if n.strip()] + for n in nums: + try: + layers.add(int(n)) + except Exception: + pass + selection[var] = {"pfts": pfts, "layers": layers} + except Exception as e: + print(f"Failed to parse top-bad report {report_path}: {e}") + return {} + return selection + + +def _parse_worst_vars_report(report_path): + selection = {} + if not os.path.exists(report_path): + print(f"Worst-variables report not found: {report_path}") + return selection + in_section = False + try: + with open(report_path, "r") as f: + for line in f: + stripped = line.strip("\n") + header = stripped.strip() + if header.startswith("## Variables with Worst Predictions"): + in_section = True + continue + if in_section and header.startswith("## "): + break + if in_section and stripped and not stripped.startswith("#"): + m = re.match(r"\s*([A-Za-z0-9_]+):\s*", stripped) + if not m: + continue + var = m.group(1) + selection[var] = {"pfts": set(), "layers": set()} + except Exception as e: + print(f"Failed to parse worst variables section {report_path}: {e}") + return {} + return selection + + +def main_with_site(results_dir, plot_scatter, plot_loss, lon, lat, tol, coord_csv, top_bad_only=False, top_bad_report=None, worst_only=False): + plots_dir = os.path.join(results_dir, "plots_site") + os.makedirs(plots_dir, exist_ok=True) + stats_path = os.path.join(results_dir, "validation_stats.csv") + stats_data = [] + + site_indices = _find_site_indices(results_dir, lon, lat, tol, coord_csv) + if not site_indices: + raise ValueError("No matching samples found for the given lon/lat. Try a larger --tolerance or verify coordinates.") + site_label = f"Site ({lon}, {lat})" + + selection = None + if top_bad_only or worst_only: + report_path = top_bad_report or os.path.join(results_dir, "analysis", "quality_summary_report.txt") + if worst_only: + selection = _parse_worst_vars_report(report_path) + if selection: + print(f"Plotting restricted to worst variables from: {report_path}") + if (not selection) and top_bad_only: + selection = _parse_top_bad_report(report_path) + if selection: + print(f"Plotting restricted to top-bad variables from: {report_path}") + if not selection: + print("No selections parsed from report; proceeding without restriction.") + selection = None + + pft_gt_dir = os.path.join(results_dir, "cnp_predictions", "pft_1d_ground_truth") + pft_pred_dir = os.path.join(results_dir, "cnp_predictions", "pft_1d_predictions") + 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, site_indices, site_label, plot_scatter, selection) + else: + print("Using legacy 1D single-file format") + gt_path = os.path.join(results_dir, "cnp_predictions", "ground_truth_1d.csv") + pred_path = os.path.join(results_dir, "cnp_predictions", "predictions_1d.csv") + if os.path.exists(gt_path) and os.path.exists(pred_path): + analyze_1d(gt_path, pred_path, "1D", plots_dir, results_dir, stats_data, site_indices, site_label, plot_scatter, selection) + else: + print("No 1D data found in either format") + + scalar_gt = os.path.join(results_dir, "cnp_predictions", "ground_truth_scalar.csv") + scalar_pred = os.path.join(results_dir, "cnp_predictions", "predictions_scalar.csv") + if os.path.exists(scalar_gt) and os.path.exists(scalar_pred): + print("Analyzing scalar data...") + analyze_pair(scalar_gt, scalar_pred, "Scalar", plots_dir, stats_data, site_indices, site_label, per_column=True, plot_scatter=plot_scatter, selection=selection) + else: + print("Scalar data files not found") + + soil_gt_dir = os.path.join(results_dir, "cnp_predictions", "soil_2d_ground_truth") + 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, site_indices, site_label, plot_scatter, selection) + else: + print("Using legacy 2D single-file format") + soil_gt = os.path.join(results_dir, "cnp_predictions", "ground_truth_2d.csv") + soil_pred = os.path.join(results_dir, "cnp_predictions", "predictions_2d.csv") + if os.path.exists(soil_gt) and os.path.exists(soil_pred): + analyze_2d(soil_gt, soil_pred, "2D", plots_dir, results_dir, stats_data, site_indices, site_label, plot_scatter, selection) + else: + print("No 2D data found in either format") + + if plot_loss: + loss_csv = os.path.join(results_dir, "cnp_training_losses.csv") + if os.path.exists(loss_csv): + plot_train_val_accuracy(loss_csv, plots_dir) + else: + print("Loss CSV not found for train/val accuracy plot.") + + if stats_data: + stats_df = pd.DataFrame(stats_data) + stats_df.to_csv(stats_path, index=False) + print(f"Saved validation stats to: {stats_path}") + + test_metrics_path = os.path.join(results_dir, "cnp_predictions", "test_metrics.csv") + if os.path.exists(test_metrics_path): + print("\nTest Metrics:") + print(pd.read_csv(test_metrics_path)) + else: + print("test_metrics.csv not found.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Postprocess CNP model results with a site-specific highlight") + parser.add_argument("results_dir", nargs="?", default=".", help="Results directory (run_xxxxx). Default is current directory.") + parser.add_argument("--lon", type=float, required=True, help="Site longitude") + parser.add_argument("--lat", type=float, required=True, help="Site latitude") + parser.add_argument("--tolerance", type=float, default=0.01, help="Coordinate matching tolerance (default: 0.01)") + parser.add_argument("--coord-csv", type=str, default=None, help="CSV file containing Longitude/Latitude for matching indices") + + parser.add_argument("--no-scatter", action="store_false", dest="plot_scatter", help="Do not generate scatter plots") + parser.add_argument("--no-plot-loss", action="store_false", dest="plot_loss", help="Do not plot train/val loss curve") + parser.add_argument("--stats-only", action="store_true", help="Only compute and save statistics CSV; do not generate any plots") + 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") + + parser.set_defaults(plot_scatter=True, plot_loss=True) + args = parser.parse_args() + + if getattr(args, "stats_only", False): + args.plot_scatter = False + args.plot_loss = False + + if len(os.sys.argv) < 2: + print("Using current directory as results directory") + + main_with_site( + args.results_dir, + args.plot_scatter, + args.plot_loss, + args.lon, + args.lat, + args.tolerance, + args.coord_csv, + args.top_bad_only, + args.top_bad_report, + worst_only=getattr(args, "worst_only", False), + ) diff --git a/scripts/verify_predictions_in_netcdf.py b/scripts/verify_predictions_in_netcdf.py new file mode 100644 index 0000000..d47c815 --- /dev/null +++ b/scripts/verify_predictions_in_netcdf.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +""" +Verify that all prediction CSV values are correctly written into the NetCDF file. + +Checks: +1) Scalar predictions: predictions_scalar.csv -> NetCDF var (gridcell) +2) PFT 1D predictions: predictions_Y_*.csv -> NetCDF var (pft, gridcell) +3) Soil 2D predictions: predictions_Y_*.csv -> NetCDF var (column, levgrnd, gridcell) + +Also validates that CSV Longitude/Latitude align with NetCDF grid1d_lon/grid1d_lat. +""" + +import argparse +from pathlib import Path +from typing import List, Tuple + +import numpy as np +import pandas as pd +import xarray as xr + + +DEFAULT_PRED_DIR = "cnp_inference_entire_dataset/cnp_predictions" +DEFAULT_NETCDF = "comparison_results/ai_predictions_for_plotting.nc" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Verify CSV predictions are correctly written into NetCDF." + ) + parser.add_argument("--predictions-dir", default=DEFAULT_PRED_DIR, + help="Directory containing cnp_predictions (scalar/pft/soil CSVs).") + parser.add_argument("--netcdf", default=DEFAULT_NETCDF, + help="NetCDF file produced by ai_predictions_to_netcdf.py.") + parser.add_argument("--tol", type=float, default=1e-6, + help="Tolerance for max abs diff.") + parser.add_argument("--max-vars", type=int, default=0, + help="Limit number of variables per group (0 = no limit).") + return parser.parse_args() + + +def _coords_from_df(df: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray]: + return df["Longitude"].values, df["Latitude"].values + + +def _coords_match(lon_csv: np.ndarray, lat_csv: np.ndarray, + lon_nc: np.ndarray, lat_nc: np.ndarray, tol: float) -> bool: + if len(lon_csv) != len(lon_nc) or len(lat_csv) != len(lat_nc): + return False + return np.allclose(lon_csv, lon_nc, atol=tol) and np.allclose(lat_csv, lat_nc, atol=tol) + + +def _build_coord_index(lon_nc: np.ndarray, lat_nc: np.ndarray, decimals: int = 4) -> dict: + keys = {} + for i in range(len(lon_nc)): + key = (round(float(lon_nc[i]), decimals), round(float(lat_nc[i]), decimals)) + if key not in keys: + keys[key] = i + return keys + + +def _map_csv_to_grid(df: pd.DataFrame, coord_index: dict, lon_nc: np.ndarray, + lat_nc: np.ndarray, decimals: int = 4) -> np.ndarray: + lon = df["Longitude"].values + lat = df["Latitude"].values + idx = np.full(len(df), -1, dtype=int) + for i in range(len(df)): + key = (round(float(lon[i]), decimals), round(float(lat[i]), decimals)) + if key in coord_index: + idx[i] = coord_index[key] + else: + # Fallback: nearest neighbor (distance is tiny for these cases) + d = (lon_nc - lon[i]) ** 2 + (lat_nc - lat[i]) ** 2 + idx[i] = int(np.argmin(d)) + return idx + + +def _max_mean_diff(a: np.ndarray, b: np.ndarray) -> Tuple[float, float]: + diff = np.abs(a - b) + return float(np.nanmax(diff)), float(np.nanmean(diff)) + + +def _sorted_pft_cols(df: pd.DataFrame, var: str) -> List[str]: + cols = [c for c in df.columns if c.startswith(f"Y_{var}_pft")] + def pft_num(c: str) -> int: + try: + return int(c.split("_pft")[-1]) + except ValueError: + return 999 + return sorted(cols, key=pft_num) + + +def _sorted_layer_cols(df: pd.DataFrame, var: str) -> List[str]: + cols = [c for c in df.columns if c.startswith(f"Y_{var}_col1_layer")] + def layer_num(c: str) -> int: + try: + return int(c.replace(f"Y_{var}_col1_layer", "")) + except ValueError: + return 999 + return sorted(cols, key=layer_num) + + +def main() -> None: + args = parse_args() + pred_dir = Path(args.predictions_dir) + netcdf_path = Path(args.netcdf) + + if not pred_dir.exists(): + raise FileNotFoundError(f"predictions-dir not found: {pred_dir}") + if not netcdf_path.exists(): + raise FileNotFoundError(f"netcdf not found: {netcdf_path}") + + ds = xr.open_dataset(netcdf_path) + lon_nc = ds["grid1d_lon"].values + lat_nc = ds["grid1d_lat"].values + coord_index = _build_coord_index(lon_nc, lat_nc) + + print("NetCDF:", netcdf_path) + print("Predictions dir:", pred_dir) + print(f"gridcell count: {ds.sizes.get('gridcell')}") + + all_ok = True + + # 1) Scalar predictions + scalar_csv = pred_dir / "predictions_scalar.csv" + if scalar_csv.exists(): + df = pd.read_csv(scalar_csv) + lon_csv, lat_csv = _coords_from_df(df) + coords_ok = _coords_match(lon_csv, lat_csv, lon_nc, lat_nc, args.tol) + if not coords_ok: + idx = _map_csv_to_grid(df, coord_index, lon_nc, lat_nc) + missing = int(np.sum(idx < 0)) + if missing > 0: + print(f"ERROR: scalar coords missing {missing} rows in NetCDF grid") + all_ok = False + scalar_cols = [c for c in df.columns if c.startswith("Y_")] + if args.max_vars > 0: + scalar_cols = scalar_cols[:args.max_vars] + for col in scalar_cols: + var = col[2:] + if var not in ds: + print(f"Missing in NetCDF (skipped): {var}") + continue + if coords_ok: + a = df[col].values + b = ds[var].values + else: + b = ds[var].values[idx] + a = df[col].values + max_diff, mean_diff = _max_mean_diff(a, b) + ok = max_diff <= args.tol + all_ok = all_ok and ok + print(f"scalar {var}: max_diff={max_diff} mean_diff={mean_diff} ok={ok}") + + # 2) PFT 1D predictions + pft_dir = pred_dir / "pft_1d_predictions" + if pft_dir.exists(): + files = sorted(pft_dir.glob("predictions_*.csv")) + if args.max_vars > 0: + files = files[:args.max_vars] + for f in files: + var = f.stem.replace("predictions_Y_", "") + df = pd.read_csv(f) + lon_csv, lat_csv = _coords_from_df(df) + coords_ok = _coords_match(lon_csv, lat_csv, lon_nc, lat_nc, args.tol) + if not coords_ok: + idx = _map_csv_to_grid(df, coord_index, lon_nc, lat_nc) + missing = int(np.sum(idx < 0)) + if missing > 0: + print(f"ERROR: pft coords missing {missing} rows in NetCDF for {var}") + all_ok = False + if var not in ds: + print(f"Missing in NetCDF (skipped): {var}") + continue + cols = _sorted_pft_cols(df, var) + pft_data = np.zeros((16, len(df)), dtype=float) + for i, c in enumerate(cols[:16]): + pft_data[i, :] = df[c].values + if coords_ok: + nc_data = ds[var].values + else: + nc_data = ds[var].values[:, idx] + max_diff, mean_diff = _max_mean_diff(pft_data, nc_data) + ok = max_diff <= args.tol + all_ok = all_ok and ok + print(f"pft {var}: max_diff={max_diff} mean_diff={mean_diff} ok={ok}") + + # 3) Soil 2D predictions + soil_dir = pred_dir / "soil_2d_predictions" + if soil_dir.exists(): + files = sorted(soil_dir.glob("predictions_*.csv")) + if args.max_vars > 0: + files = files[:args.max_vars] + for f in files: + var = f.stem.replace("predictions_Y_", "") + df = pd.read_csv(f) + lon_csv, lat_csv = _coords_from_df(df) + coords_ok = _coords_match(lon_csv, lat_csv, lon_nc, lat_nc, args.tol) + if not coords_ok: + idx = _map_csv_to_grid(df, coord_index, lon_nc, lat_nc) + missing = int(np.sum(idx < 0)) + if missing > 0: + print(f"ERROR: soil coords missing {missing} rows in NetCDF for {var}") + all_ok = False + if var not in ds: + print(f"Missing in NetCDF (skipped): {var}") + continue + cols = _sorted_layer_cols(df, var)[:10] + soil_data = np.zeros((1, len(cols), len(df)), dtype=float) + for i, c in enumerate(cols): + soil_data[0, i, :] = df[c].values + if coords_ok: + nc_data = ds[var].values[:, :len(cols), :] + else: + nc_data = ds[var].values[:, :len(cols), idx] + max_diff, mean_diff = _max_mean_diff(soil_data, nc_data) + ok = max_diff <= args.tol + all_ok = all_ok and ok + print(f"soil {var}: max_diff={max_diff} mean_diff={mean_diff} ok={ok}") + + ds.close() + + print("\nOverall:", "PASS" if all_ok else "FAIL") + + +if __name__ == "__main__": + main() From aa5fe7a06c22c9a0a6050e98522e87a9f44d66b3 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Tue, 27 Jan 2026 21:05:32 -0800 Subject: [PATCH 50/51] Fix: apply PFT mask during all-gridcell inference --- scripts/run_inference_all.py | 79 +++++++++++++++++++++++++++++++++++- 1 file changed, 77 insertions(+), 2 deletions(-) diff --git a/scripts/run_inference_all.py b/scripts/run_inference_all.py index 8693a44..90a7b93 100644 --- a/scripts/run_inference_all.py +++ b/scripts/run_inference_all.py @@ -196,7 +196,8 @@ def run_inference_all( strict_loading: bool = True, debug_vars: bool = False, loader: str = 'auto', - mask_pft_with_gt: bool = False + mask_pft_with_gt: bool = False, + mask_absent_pfts: bool = True ) -> Path: """Run inference with the trained CNP model over the entire dataset. @@ -243,6 +244,11 @@ def run_inference_all( variable_list_path=variable_list, model_config_path=model_config ) + try: + config.update_training_config(mask_absent_pfts=bool(mask_absent_pfts)) + logging.info(f"mask_absent_pfts set to {bool(mask_absent_pfts)}") + except Exception as e: + logging.warning(f"Failed to set mask_absent_pfts on training_config: {e}") if model_config is not None and use_training_config: logging.warning("--model-config provided along with --use-training-config; training config will still govern variables and scalers. Model overrides only affect architecture sizing.") @@ -698,6 +704,21 @@ def _apply_soil2d_manager(manager, data_4d, variable_names): logging.warning("Training scaler 'static' not found; leaving static unnormalized") static_t = torch.tensor(static_mat, dtype=dtype) + # PFT presence mask (from raw PCT_NAT_PFT_1..16) if requested + pft_presence_mask_t = None + if mask_absent_pfts: + try: + pct_cols = [f'PCT_NAT_PFT_{i}' for i in range(1, 17)] + if all(c in df.columns for c in pct_cols): + pct = df[pct_cols].values.astype(np.float32) + mask = (pct > 0.0).astype(np.float32) + pft_presence_mask_t = torch.tensor(mask, dtype=dtype) + logging.info("Created pft_presence_mask from PCT_NAT_PFT_1..16 (fallback path)") + else: + logging.warning("PCT_NAT_PFT_1..16 columns missing; pft_presence_mask not created (fallback path)") + except Exception as e: + logging.warning(f"Failed to create pft_presence_mask in fallback path: {e}") + # PFT param pp_cols = config.data_config.pft_param_columns num_pfts = 17 @@ -901,7 +922,7 @@ def _apply_soil2d_manager(manager, data_4d, variable_names): logging.warning("Training scaler 'y_soil_2d' not found; leaving y_soil_2d unnormalized (group)") y_soil_2d_t = torch.tensor(y_soil2d, dtype=dtype) - return { + ret = { 'time_series_data': time_series_t, 'static_data': static_t, 'pft_param_data': pft_param_t, @@ -914,6 +935,9 @@ def _apply_soil2d_manager(manager, data_4d, variable_names): 'water': None, 'y_water': None, } + if pft_presence_mask_t is not None: + ret['pft_presence_mask'] = pft_presence_mask_t + return ret # Normalize using training scalers by default (no refit), or refit if requested logging.info("Normalizing data using training-compatible method...") @@ -943,6 +967,11 @@ def _apply_soil2d_manager(manager, data_4d, variable_names): # For inference, we use the test data (which contains all data when train_split=0.0) test_data = split_data['test'] logging.info(f"Using test data for inference: {len(test_data)} data types") + if mask_absent_pfts: + if isinstance(test_data, dict) and 'pft_presence_mask' in test_data: + logging.info("pft_presence_mask available; mask_absent_pfts will be applied during evaluation") + else: + logging.warning("mask_absent_pfts enabled but pft_presence_mask missing; mask may not be applied") # Convert test_data to model_inputs format expected by the model model_inputs = {} @@ -1136,6 +1165,32 @@ def _preview(group_key: str, names: list): model_inputs.get('variables_2d_soil') ) + # Optionally apply PFT absence mask before saving predictions + if mask_absent_pfts and isinstance(test_data, dict) and 'pft_presence_mask' in test_data and isinstance(predictions, dict) and 'pft_1d' in predictions: + try: + vec = predictions['pft_1d'] + mask = test_data['pft_presence_mask'] + if isinstance(vec, torch.Tensor) and isinstance(mask, torch.Tensor): + mask = mask.to(vec.device, non_blocking=True) + n_pfts = 16 + # Determine number of variables + varnames = data_info.get('variables_1d_pft', []) if isinstance(data_info, dict) else [] + if vec.dim() == 2: + n_vars = len(varnames) if varnames else (vec.size(1) // n_pfts) + vec = vec.view(vec.size(0), n_vars, n_pfts) + reshaped = True + else: + reshaped = False + if mask.dim() == 2: + mask = mask.view(mask.size(0), 1, n_pfts) + vec = vec * mask + predictions['pft_1d'] = vec.view(vec.size(0), -1) if reshaped else vec + logging.info("Applied pft_presence_mask to PFT1D predictions before saving") + except Exception as e: + logging.warning(f"Failed to apply pft_presence_mask to predictions: {e}") + elif mask_absent_pfts: + logging.warning("mask_absent_pfts enabled but pft_presence_mask not available; predictions not masked") + logging.info("Inference completed successfully") # Save results @@ -1273,6 +1328,16 @@ def _preview(group_key: str, names: list): except Exception: gt_mask_per_var = None + # PFT presence mask (from PCT_NAT_PFT_1..16), applied after inverse transform + pft_presence_mask_np = None + if mask_absent_pfts and isinstance(test_data, dict) and 'pft_presence_mask' in test_data and hasattr(test_data['pft_presence_mask'], 'numel'): + try: + ppm = test_data['pft_presence_mask'].detach().cpu().numpy() + if ppm.ndim == 2 and ppm.shape[1] == num_pfts: + pft_presence_mask_np = ppm + except Exception: + pft_presence_mask_np = None + # Write predictions per variable (denormalized when possible) for v in range(num_variables): var_name = var_names[v] @@ -1308,6 +1373,12 @@ def _preview(group_key: str, names: list): var_predictions_original = var_predictions_original * mask_v.astype(var_predictions_original.dtype) except Exception: pass + # Apply PFT presence mask (after inverse transform) + try: + if pft_presence_mask_np is not None and pft_presence_mask_np.shape == var_predictions_original.shape: + var_predictions_original = var_predictions_original * pft_presence_mask_np.astype(var_predictions_original.dtype) + except Exception: + pass # Optional dump before saving predictions try: @@ -1567,6 +1638,9 @@ def main(): parser.add_argument("--debug-vars", action='store_true', help="Print detailed variable names and sample values during preprocessing/inference") parser.add_argument("--loader", choices=['auto','pandas','individual'], default='auto', help="Data loader to use (default: auto)") parser.add_argument("--mask-pft-with-gt", action='store_true', default=False, help="Mask PFT1D predictions by GT non-zero mask when available") + parser.add_argument("--mask-absent-pfts", dest="mask_absent_pfts", action="store_true", help="Mask absent PFTs using PCT_NAT_PFT_1..16 when available") + 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)") args = parser.parse_args() @@ -1587,6 +1661,7 @@ def main(): , debug_vars=args.debug_vars , loader=args.loader , mask_pft_with_gt=args.mask_pft_with_gt + , mask_absent_pfts=args.mask_absent_pfts ) print(f"Inference completed successfully. Results saved to: {output_path}") From 2d266ea11c2c522d14606807ff56229d78b4f2c3 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Tue, 3 Feb 2026 13:41:54 -0800 Subject: [PATCH 51/51] Focus dataset on tropical regions for CNP training --- config/training_config.py | 13 +++- data/data_loader_individual.py | 118 ++++++++++++++++++++++++++++----- docs/CNP_pipeline_runbook.md | 10 +++ train_cnp_model.py | 49 ++++++++++++++ 4 files changed, 172 insertions(+), 18 deletions(-) diff --git a/config/training_config.py b/config/training_config.py index d297872..4a9f80d 100644 --- a/config/training_config.py +++ b/config/training_config.py @@ -6,7 +6,7 @@ """ from dataclasses import dataclass, field -from typing import List, Dict, Any, Optional, Union +from typing import List, Dict, Any, Optional, Union, Tuple import torch import torch.nn as nn import torch.optim as optim @@ -94,7 +94,12 @@ class DataConfig: # Data splitting train_split: float = 0.8 + test_split: Optional[float] = None random_state: int = 42 + # Tropical-only filtering (apply before train/test split) + tropical_only: bool = False + tropical_lat_range: Tuple[float, float] = (-23.5, 23.5) + tropical_lat_column: Optional[str] = None # File loading limits (for testing) @@ -114,7 +119,11 @@ class DataConfig: class ModelConfig: """Configuration for model architecture.""" - # LSTM parameters + # Core dimensions (Dual Stream Architecture) + embed_dim: int = 256 + patch_size: int = 60 + + # LSTM parameters (Legacy / Stream 1 variant) lstm_hidden_size: int = 64 # Fully connected layers diff --git a/data/data_loader_individual.py b/data/data_loader_individual.py index d9cd8f6..7e593bf 100644 --- a/data/data_loader_individual.py +++ b/data/data_loader_individual.py @@ -96,10 +96,30 @@ def check_nans(self): if count > 0: logger.info(f" {col}: {count}") + def _resolve_lat_column(self) -> Optional[str]: + """Resolve latitude column name from config or common patterns.""" + candidates = [] + lat_override = getattr(self.data_config, 'tropical_lat_column', None) + if lat_override: + candidates.append(lat_override) + # Prefer static columns that look like latitude + for col in getattr(self.data_config, 'static_columns', []) or []: + if 'lat' in str(col).lower(): + candidates.append(col) + # Common column names + candidates.extend(['lat', 'latitude', 'LAT', 'Latitude', 'LATITUDE']) + for col in candidates: + 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 = [] logger.info("Loading data from multiple paths...") + logger.info(f"data_paths: {self.data_config.data_paths}") + logger.info(f"file_pattern: {self.data_config.file_pattern}") + logger.info(f"dataset_file_patterns: {getattr(self.data_config, 'dataset_file_patterns', {})}") for path in self.data_config.data_paths: # Resolve files matching pattern # Support per-dataset file patterns if provided @@ -107,8 +127,15 @@ def load_data(self) -> pd.DataFrame: per_dataset_patterns = getattr(self.data_config, 'dataset_file_patterns', {}) or {} except Exception: per_dataset_patterns = {} - pattern = per_dataset_patterns.get(path, self.data_config.file_pattern) - files = list(Path(path).glob(pattern)) + # Normalize path for matching (resolve to absolute path) + path_normalized = str(Path(path).resolve()) + # Try both normalized and original path as keys + pattern = per_dataset_patterns.get(path_normalized, + per_dataset_patterns.get(path, self.data_config.file_pattern)) + logger.info(f"Searching in path: {path} (normalized: {path_normalized}), using pattern: {pattern}") + path_obj = Path(path) + logger.info(f"Path exists: {path_obj.exists()}, is_dir: {path_obj.is_dir()}") + files = list(path_obj.glob(pattern)) # Deterministic ordering for test runs if getattr(self.data_config, 'sort_file_list', True): files = sorted(files, key=lambda p: p.name) @@ -122,8 +149,10 @@ def load_data(self) -> pd.DataFrame: logger.info(f"Limited to {len(files)} files due to max_files={self.data_config.max_files}") # Load each file - for file_path in files: + total_files = len(files) + for idx, file_path in enumerate(files, 1): try: + logger.info(f"Loading file {idx}/{total_files}: {file_path.name}") # Check file extension and use appropriate loading method if str(file_path).endswith('.pkl'): df_chunk = pd.read_pickle(file_path) @@ -138,7 +167,7 @@ def load_data(self) -> pd.DataFrame: # Load all files - zeros are valid data in soil science df_list.append(df_chunk) - logger.debug(f"Loaded {len(df_chunk)} samples from {file_path}") + logger.info(f"Loaded {len(df_chunk)} samples from {file_path.name} (total samples so far: {sum(len(df) for df in df_list)})") except Exception as e: logger.error(f"Failed to load {file_path}: {e}") @@ -151,6 +180,17 @@ def load_data(self) -> pd.DataFrame: self.df = pd.concat(df_list, ignore_index=True) logger.info(f"Successfully loaded {len(self.df)} samples") + # Print all variables/columns in the dataset + logger.info("=" * 80) + logger.info("所有数据集变量列表 (All Dataset Variables):") + logger.info("=" * 80) + logger.info(f"总变量数: {len(self.df.columns)}") + logger.info(f"数据集形状: {self.df.shape}") + logger.info("\n变量列表 (按字母顺序):") + for i, col in enumerate(sorted(self.df.columns), 1): + logger.info(f" {i:4d}. {col}") + logger.info("=" * 80) + return self.df def preprocess_data(self): @@ -177,6 +217,30 @@ def preprocess_data(self): logger.info(f"Longitude filtering: {original_size} samples -> {filtered_size} samples (dropped {dropped_count} samples)") else: logger.warning("'Longitude' column not found in dataset. Cannot apply longitude filtering.") + + # Optional tropical-only filtering by latitude + if getattr(self.data_config, 'tropical_only', False): + lat_col = self._resolve_lat_column() + if lat_col is None: + logger.warning("Tropical filter enabled but no latitude column found. Skipping tropical filtering.") + else: + lat_range = getattr(self.data_config, 'tropical_lat_range', (-23.5, 23.5)) + try: + lat_min, lat_max = float(lat_range[0]), float(lat_range[1]) + except Exception: + lat_min, lat_max = -23.5, 23.5 + logger.warning("Invalid tropical_lat_range; falling back to [-23.5, 23.5].") + original_size = len(self.df) + lat_vals = pd.to_numeric(self.df[lat_col], errors='coerce') + mask = lat_vals.between(lat_min, lat_max, inclusive='both') + self.df = self.df[mask].reset_index(drop=True) + filtered_size = len(self.df) + logger.info( + f"Tropical filtering on '{lat_col}': {original_size} -> {filtered_size} " + f"(lat range [{lat_min}, {lat_max}])" + ) + if filtered_size == 0: + logger.warning("Tropical filter removed all samples. Check latitude column and range.") # Drop specified columns if hasattr(self.data_config, 'filter_columns') and self.data_config.filter_columns: @@ -197,14 +261,14 @@ def _to_ts_and_truncate(x): target_len = int(getattr(self.data_config, 'time_series_length', 240)) if isinstance(x, (list, np.ndarray)): arr = np.array(x, dtype=np.float32).flatten() - # Prefer earliest 20-year window as per repeated forcing spec + # Prefer latest 20-year window (last 20 years) if arr.size >= target_len: - arr = arr[:target_len] + arr = arr[-target_len:] else: - # pad to target_len with zeros at the end + # pad to target_len with zeros at the beginning (to align with latest data) pad = target_len - arr.size if pad > 0: - arr = np.pad(arr, (0, pad), mode='constant') + arr = np.pad(arr, (pad, 0), mode='constant') # ensure no NaN/Inf arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0) return arr @@ -1533,14 +1597,36 @@ def split_data(self, normalized_data: Dict[str, Any]) -> Dict[str, Any]: test_data = {} total_samples = len(self.df) - train_size = int(self.data_config.train_split * total_samples) - test_size = total_samples - train_size - - logger.info(f"Data splitting details:") - logger.info(f" - Total samples: {total_samples}") - logger.info(f" - Train split ratio: {self.data_config.train_split}") - logger.info(f" - Train size: {train_size}") - logger.info(f" - Test size: {test_size}") + + # 如果设置了 test_split，分别使用 train_split 和 test_split 计算 + # 否则使用原来的逻辑：test_size = total_samples - train_size + if self.data_config.test_split is not None: + train_size = int(self.data_config.train_split * total_samples) + test_size = int(self.data_config.test_split * total_samples) + + # 验证比例是否合理 + total_ratio = self.data_config.train_split + self.data_config.test_split + if total_ratio > 1.0: + logger.warning( + f"Train split ({self.data_config.train_split}) + Test split ({self.data_config.test_split}) = {total_ratio} > 1.0. " + f"Adjusting test_split to {1.0 - self.data_config.train_split}" + ) + test_size = int((1.0 - self.data_config.train_split) * total_samples) + + unused_size = total_samples - train_size - test_size + logger.info(f"Data splitting details:") + logger.info(f" - Total samples: {total_samples}") + logger.info(f" - Train split ratio: {self.data_config.train_split} ({train_size} samples)") + logger.info(f" - Test split ratio: {self.data_config.test_split} ({test_size} samples)") + logger.info(f" - Unused data: {unused_size} samples ({(1.0 - self.data_config.train_split - self.data_config.test_split)*100:.1f}%)") + else: + train_size = int(self.data_config.train_split * total_samples) + test_size = total_samples - train_size + + logger.info(f"Data splitting details:") + logger.info(f" - Total samples: {total_samples}") + logger.info(f" - Train split ratio: {self.data_config.train_split} ({train_size} samples)") + logger.info(f" - Test size: {test_size} samples (剩余部分)") # Expose split indices for downstream use (e.g., location validation) # Matches the contiguous slicing used below diff --git a/docs/CNP_pipeline_runbook.md b/docs/CNP_pipeline_runbook.md index 3122f07..138aa5a 100644 --- a/docs/CNP_pipeline_runbook.md +++ b/docs/CNP_pipeline_runbook.md @@ -15,6 +15,16 @@ python train_cnp_model.py --variable-list CNP_IO_demo.txt --epoch 100 2>&1 & Notes: - This launches training in the background and redirects logs to stdout/stderr. - The run output directory will be created under `cnp_results/run_YYYYMMDD_HHMMSS`. +- Optional: restrict training/test split to tropical latitudes only: +```bash +python train_cnp_model.py --variable-list CNP_IO_demo.txt --epoch 100 \ + --tropical-only +``` +- Optional: customize the latitude range or column name: +```bash +python train_cnp_model.py --variable-list CNP_IO_demo.txt --epoch 100 \ + --tropical-only --tropical-lat-range -23.5,23.5 --tropical-lat-column Latitude +``` ### 2a) Fine-tune a pretrained model (optional) If you already have a trained checkpoint and want to continue training on a TVA-style dataset, use the fine-tuning helper. Populate the necessary paths in your CNP_IO file (e.g. `CNP_IO_updated9_dev_gao.txt`): diff --git a/train_cnp_model.py b/train_cnp_model.py index 156742d..55f2324 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -215,6 +215,23 @@ def main(): default=None, help='Glob pattern for training files (e.g., enhanced_1_training_data_batch_*.pkl)' ) + parser.add_argument( + '--tropical-only', + action='store_true', + help='Filter dataset to tropical latitude band before train/test split' + ) + 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( '--max-files', type=int, @@ -328,6 +345,25 @@ def main(): logger.info(f"Applied data overrides: {update_kwargs}") except Exception as e: logger.warning(f"Failed to apply data overrides: {e}") + # Optional tropical-only filtering + if args.tropical_only: + tropical_kwargs = {'tropical_only': True} + if args.tropical_lat_range: + try: + parts = [p.strip() for p in str(args.tropical_lat_range).split(',')] + if len(parts) == 2: + tropical_kwargs['tropical_lat_range'] = (float(parts[0]), float(parts[1])) + else: + logger.warning("Invalid --tropical-lat-range; expected format 'min,max'. Using default.") + except Exception: + logger.warning("Failed to parse --tropical-lat-range; using default.") + if args.tropical_lat_column: + tropical_kwargs['tropical_lat_column'] = str(args.tropical_lat_column).strip() + try: + config.update_data_config(**tropical_kwargs) + logger.info(f"Enabled tropical filtering: {tropical_kwargs}") + except Exception as e: + logger.warning(f"Failed to apply tropical filtering config: {e}") if args.variable_list is not None: logger.info(f"Using CNP configuration from variable list file: {args.variable_list}") else: @@ -460,6 +496,19 @@ def main(): ) # Check raw data for non-zero values after loading raw_data = data_loader.load_data() + + # Print all variables after loading + if hasattr(data_loader, 'df') and isinstance(data_loader.df, pd.DataFrame): + logger.info("=" * 80) + logger.info("训练数据集变量列表 (Training Dataset Variables):") + logger.info("=" * 80) + logger.info(f"总变量数: {len(data_loader.df.columns)}") + logger.info(f"数据集形状: {data_loader.df.shape}") + logger.info("\n所有变量列表 (All Variables):") + for i, col in enumerate(sorted(data_loader.df.columns), 1): + logger.info(f" {i:4d}. {col}") + logger.info("=" * 80) + logger.info("Checking raw data for soil2D variables...") for key, value in raw_data.items(): if 'soil' in key.lower() and '2d' in key.lower():