From aa5fe7a06c22c9a0a6050e98522e87a9f44d66b3 Mon Sep 17 00:00:00 2001 From: Daewi Gao Date: Tue, 27 Jan 2026 21:05:32 -0800 Subject: [PATCH 1/2] 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 2/2] 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():