location and site inference bug fix

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

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"

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}")

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: "<var>: <good>% good, <ok>% ok, <bad>% 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)
+    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))

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

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}")