diff --git a/commands.txt b/commands.txt index 4a738ed..b33e87c 100644 --- a/commands.txt +++ b/commands.txt @@ -484,4 +484,39 @@ python scripts/compare_5p_amazon_site_scatter.py \ cnp_inference_entire_dataset/cnp_predictions/soil_2d_predictions_4P_bias_corrected_phase2 \ --lon 292.5 \ --lat -15.549738 \ - --output-dir cnp_results/run_20260228_214757_natveg_improved/analysis/amazon_5p_comparison_bias_correction_site_B \ No newline at end of file + --output-dir cnp_results/run_20260228_214757_natveg_improved/analysis/amazon_5p_comparison_bias_correction_site_B + + + + cd /mnt/proj-shared/AI4BGC_7xw/AI4BGC + +# 1) Phase1 global → Africa_phase1_global_restart.nc +python scripts/extract_elm_restart_point.py \ + --restart-file cnp_results/run_20260311_165538_phase1_global/updated_restart_base.nc \ + --lat 0.0 \ + --lon 28.0 \ + --output-file cnp_results/run_20260311_165538_phase1_global/Africa_phase1_global_restart.nc + +# 2) Phase2 tropical → Africa_phase2_tropical_restart.nc +python scripts/extract_elm_restart_point.py \ + --restart-file cnp_results/run_20260311_174843_phase2_tropical/updated_restart_phase2_tropical_5P_raw.nc \ + --lat 0.0 \ + --lon 28.0 \ + --output-file cnp_results/run_20260311_174843_phase2_tropical/Africa_phase2_tropical_restart.nc + +# 3) Phase3 two-regions → Africa_phase3_tworegions_restart.nc +python scripts/extract_elm_restart_point.py \ + --restart-file cnp_results/run_20260311_204845_phase3_tworegions/updated_restart_phase3_tworegions_5P_bias_corrected_tropical.nc \ + --lat 0.0 \ + --lon 28.0 \ + --output-file cnp_results/run_20260311_204845_phase3_tworegions/Africa_phase3_tworegions_restart.nc + +python scripts/generate_site_5p_restart_comparison.py --lon 303.75 --lat -17.434553 --site-name amazon_303_17S +python scripts/generate_site_5p_restart_comparison.py --lon 292.5 --lat -15.55 --site-name amazon_292_15S + +python scripts/generate_site_5p_restart_comparison.py --lon 28 --lat 0 --site-name africa_28_0 + +python scripts/generate_site_5p_restart_comparison.py --lon 14 --lat 0 --site-name africa + + + diff --git a/docs/WORKFLOW_PHASE1_PHASE2_PHASE3_RESTARTS.md b/docs/WORKFLOW_PHASE1_PHASE2_PHASE3_RESTARTS.md index 4cc625c..81b46e6 100644 --- a/docs/WORKFLOW_PHASE1_PHASE2_PHASE3_RESTARTS.md +++ b/docs/WORKFLOW_PHASE1_PHASE2_PHASE3_RESTARTS.md @@ -1,3 +1,37 @@ +### 5. Site‑based 5P restart comparison (Amazon / Africa or any site) + +Once `phase1_global`, `phase2_tropical`, and `phase3_tworegions` restarts are created, +you can **automatically generate single‑point restarts and 5P comparison plots** for +any longitude/latitude using: + +- `scripts/generate_site_5p_restart_comparison.py` +- Documented in `docs/SITE_5P_RESTART_COMPARISON.md` + +Example (Africa site at lon=28, lat=0): + +```bash +python scripts/generate_site_5p_restart_comparison.py \ + --lon 28.0 \ + --lat 0.0 \ + --site-name africa_28_0 +``` + +Example (Amazon validation site): + +```bash +python scripts/generate_site_5p_restart_comparison.py \ + --lon 303.75 \ + --lat -17.434553 \ + --site-name amazon_303_17S +``` + +The script will: + +- Extract single‑point restarts from the three phase runs at the requested site. +- Load Phase1 **ground‑truth 5P vertical profiles** at that site. +- Produce line plots comparing GT vs `phase1_global`, `phase2_tropical`, + and `phase3_tworegions` for the 5 P variables. + ## Workflow: phase1_global, phase2_tropical, phase3_tworegions restarts This document summarizes how to create three restart files from existing diff --git a/scripts/generate_site_5p_restart_comparison.py b/scripts/generate_site_5p_restart_comparison.py new file mode 100644 index 0000000..6a5324a --- /dev/null +++ b/scripts/generate_site_5p_restart_comparison.py @@ -0,0 +1,318 @@ +import argparse +import os +from typing import Dict + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr + +# This script lives in the same directory as extract_elm_restart_point.py +from extract_elm_restart_point import extract_single_point_elm + + +# Fixed run directories and global restart paths (same as current workflow) +PHASE1_RUN_DIR = ( + "/mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260311_165538_phase1_global" +) +PHASE2_RUN_DIR = ( + "/mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260311_174843_phase2_tropical" +) +PHASE3_RUN_DIR = ( + "/mnt/proj-shared/AI4BGC_7xw/AI4BGC/cnp_results/run_20260311_204845_phase3_tworegions" +) + +PHASE1_GLOBAL_RESTART = os.path.join(PHASE1_RUN_DIR, "updated_restart_base.nc") +PHASE2_GLOBAL_RESTART = os.path.join( + PHASE2_RUN_DIR, "updated_restart_phase2_tropical_5P_raw.nc" +) +PHASE3_GLOBAL_RESTART = os.path.join( + PHASE3_RUN_DIR, "updated_restart_phase3_tworegions_5P_bias_corrected_tropical.nc" +) + +# Ground-truth directory (per-variable CSVs ground_truth_Y_{var}.csv) +GT_PROFILES_DIR = os.path.join( + PHASE1_RUN_DIR, + "cnp_inference_entire_dataset", + "cnp_predictions", + "soil_2d_ground_truth", +) + +FIVE_P = ["labilep_vr", "occlp_vr", "solutionp_vr", "secondp_vr", "primp_vr"] +LEV_DIM = "levgrnd" + + +def _safe_site_name(site_name: str, lon: float, lat: float) -> str: + """Create a filesystem-safe site identifier.""" + if site_name: + base = site_name.strip().replace(" ", "_") + else: + base = "site" + return f"{base}_lon{lon:.2f}_lat{lat:.2f}".replace(".", "p").replace("-", "m") + + +def _find_row_for_site(df: pd.DataFrame, lon: float, lat: float, atol: float = 1e-4) -> int: + """Return row index matching lon/lat; fall back to nearest neighbour if needed.""" + if "Longitude" not in df.columns or "Latitude" not in df.columns: + raise ValueError("DataFrame must contain 'Longitude' and 'Latitude' columns.") + + lon_vals = df["Longitude"].values + lat_vals = df["Latitude"].values + mask = np.isclose(lon_vals, lon, atol=atol) & np.isclose(lat_vals, lat, atol=atol) + + idx = np.where(mask)[0] + if idx.size == 1: + return int(idx[0]) + + # Fallback: nearest neighbour + d2 = (lon_vals - lon) ** 2 + (lat_vals - lat) ** 2 + return int(np.argmin(d2)) + + +def _extract_profile_from_gt(df: pd.DataFrame, var: str, row_idx: int) -> np.ndarray: + """Extract vertical profile for one variable at row_idx from GT CSV.""" + prefix = f"Y_{var}_col1_layer" + layer_cols = sorted( + [c for c in df.columns if c.startswith(prefix)], + key=lambda c: int(c.split("layer")[-1]), + ) + if not layer_cols: + raise ValueError(f"No columns starting with '{prefix}' found.") + values = df.loc[row_idx, layer_cols].astype(float).values + return values + + +def _load_profiles_from_restart(path: str) -> Dict[str, np.ndarray]: + """Load 1D vertical profiles (top 10 layers) for the 5P variables from a single-point restart.""" + ds = xr.open_dataset(path, decode_times=False) + profs: Dict[str, np.ndarray] = {} + for var in FIVE_P: + if var not in ds: + print(f" WARNING: {var} not found in {path}") + continue + da = ds[var] + arr = da.values + if arr.ndim == 1: + prof = arr + elif arr.ndim == 2: + if da.dims[0] == LEV_DIM: + prof = arr[:, 0] + elif da.dims[1] == LEV_DIM: + prof = arr[0, :] + else: + prof = arr.reshape(-1) + else: + prof = arr.reshape(-1) + profs[var] = np.asarray(prof[:10], dtype=float) + return profs + + +def generate_site_restarts(lon: float, lat: float, site_name: str) -> Dict[str, str]: + """Extract phase1/phase2/phase3 single-point restarts for this site.""" + site_id = _safe_site_name(site_name, lon, lat) + + out_paths = { + "phase1_global": os.path.join(PHASE1_RUN_DIR, f"{site_id}_phase1_global_restart.nc"), + "phase2_tropical": os.path.join( + PHASE2_RUN_DIR, f"{site_id}_phase2_tropical_restart.nc" + ), + "phase3_tworegions": os.path.join( + PHASE3_RUN_DIR, f"{site_id}_phase3_tworegions_restart.nc" + ), + } + + print(f"\n=== Extracting single-point restarts for {site_id} ===") + extract_single_point_elm( + source_nc=PHASE1_GLOBAL_RESTART, + output_nc=out_paths["phase1_global"], + target_lat=lat, + target_lon=lon, + ) + extract_single_point_elm( + source_nc=PHASE2_GLOBAL_RESTART, + output_nc=out_paths["phase2_tropical"], + target_lat=lat, + target_lon=lon, + ) + extract_single_point_elm( + source_nc=PHASE3_GLOBAL_RESTART, + output_nc=out_paths["phase3_tworegions"], + target_lat=lat, + target_lon=lon, + ) + + return out_paths + + +def compare_5p_for_site( + lon: float, + lat: float, + site_name: str, + restart_paths: Dict[str, str], + output_dir: str, +) -> None: + """Compute and plot 5P profiles for GT + three restarts at the given site.""" + os.makedirs(output_dir, exist_ok=True) + site_id = _safe_site_name(site_name, lon, lat) + + # Load ground truth profiles + gt_profiles: Dict[str, np.ndarray] = {} + if os.path.isdir(GT_PROFILES_DIR): + print(f"\nLoading ground truth from {GT_PROFILES_DIR}") + for var in FIVE_P: + csv_path = os.path.join(GT_PROFILES_DIR, f"ground_truth_Y_{var}.csv") + if not os.path.exists(csv_path): + print(f" WARNING: GT CSV not found for {var}: {csv_path}") + continue + df_gt = pd.read_csv(csv_path) + row_idx = _find_row_for_site(df_gt, lon=lon, lat=lat, atol=1e-4) + prof = _extract_profile_from_gt(df_gt, var=var, row_idx=row_idx) + gt_profiles[var] = prof[:10] + else: + print( + f"WARNING: GT profiles directory not found at {GT_PROFILES_DIR}; " + "plots will omit ground truth." + ) + + # Load restart profiles + all_profiles: Dict[str, Dict[str, np.ndarray]] = {} + print("\nLoading site restarts:") + for name, path in restart_paths.items(): + print(f" {name}: {path}") + if not os.path.exists(path): + print(" WARNING: file does not exist, skipping.") + continue + all_profiles[name] = _load_profiles_from_restart(path) + + # Pairwise numeric diffs + print("\n=== Pairwise max |diff| between restarts (per 5P variable) ===") + restart_names = list(all_profiles.keys()) + for var in FIVE_P: + print(f"\nVariable: {var}") + for i in range(len(restart_names)): + for j in range(i + 1, len(restart_names)): + n1, n2 = restart_names[i], restart_names[j] + p1 = all_profiles[n1].get(var) + p2 = all_profiles[n2].get(var) + if p1 is None or p2 is None: + continue + nlev = min(len(p1), len(p2)) + diff = float(np.nanmax(np.abs(p1[:nlev] - p2[:nlev]))) + print(f" max |{n1} - {n2}| = {diff}") + + # Line plots + print(f"\nSaving plots to: {output_dir}") + for var in FIVE_P: + plt.figure(figsize=(5, 6)) + has_any = False + + # GT first + if var in gt_profiles: + gt = np.asarray(gt_profiles[var], dtype=float) + nlev_gt = min(10, len(gt)) + layers_gt = np.arange(1, nlev_gt + 1) + plt.plot( + layers_gt, + gt[:nlev_gt], + marker="o", + linestyle="-", + color="black", + label="ground_truth", + linewidth=2.0, + ) + has_any = True + + # Restarts + for name, profs in all_profiles.items(): + if var not in profs: + continue + prof = np.asarray(profs[var], dtype=float) + nlev = min(10, len(prof)) + layers = np.arange(1, nlev + 1) + plt.plot( + layers, + prof[:nlev], + marker="o", + linestyle="-", + label=name, + alpha=0.9, + ) + has_any = True + + if not has_any: + plt.close() + continue + + plt.xlabel("Soil layer index") + plt.ylabel(f"{var} value") + plt.title( + f"{var} vertical profile at {site_id}\n" + f"(lon={lon}, lat={lat})" + ) + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + + out_path = os.path.join(output_dir, f"{site_id}_restart_profile_{var}.png") + plt.savefig(out_path, dpi=200) + plt.close() + print(f" Wrote {out_path}") + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Given a lon/lat, extract single-point restarts for phase1_global, " + "phase2_tropical, phase3_tworegions and generate 5P vertical-profile " + "comparison plots vs ground truth." + ) + ) + parser.add_argument("--lon", type=float, required=True, help="Site longitude (degrees, 0–360).") + parser.add_argument("--lat", type=float, required=True, help="Site latitude (degrees).") + parser.add_argument( + "--site-name", + type=str, + default="site", + help="Optional human-readable site name used in filenames and plot titles.", + ) + parser.add_argument( + "--output-dir", + type=str, + default="", + help=( + "Output directory for plots. Default: " + f"{PHASE3_RUN_DIR}/analysis/_5p_restart_comparison" + ), + ) + args = parser.parse_args() + + lon = args.lon + lat = args.lat + site_name = args.site_name + site_id = _safe_site_name(site_name, lon, lat) + + if args.output_dir: + output_dir = os.path.abspath(args.output_dir) + else: + output_dir = os.path.join( + PHASE3_RUN_DIR, + "analysis", + f"{site_id}_5p_restart_comparison", + ) + + print(f"Site: {site_id} (lon={lon}, lat={lat})") + print(f"Output directory: {output_dir}") + + restart_paths = generate_site_restarts(lon=lon, lat=lat, site_name=site_name) + compare_5p_for_site( + lon=lon, + lat=lat, + site_name=site_name, + restart_paths=restart_paths, + output_dir=output_dir, + ) + + +if __name__ == "__main__": + main() +