549 lines
22 KiB
Python
549 lines
22 KiB
Python
"""
|
|
Self-contained utility module for two-step harvest date prediction and Excel export.
|
|
Includes model architecture, feature engineering, and core prediction logic.
|
|
"""
|
|
|
|
import sys
|
|
import pandas as pd
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import pickle
|
|
import yaml
|
|
from pathlib import Path
|
|
from typing import Tuple, Dict, Any, List
|
|
from sklearn.preprocessing import StandardScaler
|
|
|
|
# ============================================================================
|
|
# TORCH MODELS (from src/models.py, inlined for self-containment)
|
|
# ============================================================================
|
|
|
|
class HarvestDetectionLSTM(nn.Module):
|
|
"""Unidirectional LSTM for harvest detection with dual outputs."""
|
|
def __init__(self, input_size: int, hidden_size: int = 128,
|
|
num_layers: int = 1, dropout: float = 0.5):
|
|
super(HarvestDetectionLSTM, self).__init__()
|
|
self.input_size = input_size
|
|
self.hidden_size = hidden_size
|
|
self.num_layers = num_layers
|
|
|
|
self.lstm = nn.LSTM(
|
|
input_size=input_size,
|
|
hidden_size=hidden_size,
|
|
num_layers=num_layers,
|
|
dropout=dropout if num_layers > 1 else 0,
|
|
bidirectional=False,
|
|
batch_first=True
|
|
)
|
|
|
|
self.imminent_head = nn.Sequential(
|
|
nn.Linear(hidden_size, 16),
|
|
nn.ReLU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(16, 1),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
self.detected_head = nn.Sequential(
|
|
nn.Linear(hidden_size, 16),
|
|
nn.ReLU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(16, 1),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
lstm_out, _ = self.lstm(x)
|
|
batch_size, seq_len, hidden_size = lstm_out.shape
|
|
lstm_flat = lstm_out.reshape(-1, hidden_size)
|
|
|
|
imminent_flat = self.imminent_head(lstm_flat).reshape(batch_size, seq_len)
|
|
detected_flat = self.detected_head(lstm_flat).reshape(batch_size, seq_len)
|
|
|
|
return imminent_flat, detected_flat
|
|
|
|
|
|
class HarvestDetectionGRU(nn.Module):
|
|
"""Unidirectional GRU for harvest detection with dual outputs."""
|
|
def __init__(self, input_size: int, hidden_size: int = 128,
|
|
num_layers: int = 1, dropout: float = 0.5):
|
|
super(HarvestDetectionGRU, self).__init__()
|
|
self.input_size = input_size
|
|
self.hidden_size = hidden_size
|
|
self.num_layers = num_layers
|
|
|
|
self.gru = nn.GRU(
|
|
input_size=input_size,
|
|
hidden_size=hidden_size,
|
|
num_layers=num_layers,
|
|
dropout=dropout if num_layers > 1 else 0,
|
|
bidirectional=False,
|
|
batch_first=True
|
|
)
|
|
|
|
self.imminent_head = nn.Sequential(
|
|
nn.Linear(hidden_size, 16),
|
|
nn.ReLU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(16, 1),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
self.detected_head = nn.Sequential(
|
|
nn.Linear(hidden_size, 16),
|
|
nn.ReLU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(16, 1),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
gru_out, _ = self.gru(x)
|
|
batch_size, seq_len, hidden_size = gru_out.shape
|
|
gru_flat = gru_out.reshape(-1, hidden_size)
|
|
|
|
imminent_flat = self.imminent_head(gru_flat).reshape(batch_size, seq_len)
|
|
detected_flat = self.detected_head(gru_flat).reshape(batch_size, seq_len)
|
|
|
|
return imminent_flat, detected_flat
|
|
|
|
|
|
def create_model(model_type: str, input_size: int, hidden_size: int = 128,
|
|
num_layers: int = 1, dropout: float = 0.5, device = None) -> nn.Module:
|
|
"""Create a model from registry."""
|
|
registry = {'LSTM': HarvestDetectionLSTM, 'GRU': HarvestDetectionGRU}
|
|
if model_type not in registry:
|
|
raise ValueError(f"Unknown model type: {model_type}")
|
|
|
|
model = registry[model_type](
|
|
input_size=input_size,
|
|
hidden_size=hidden_size,
|
|
num_layers=num_layers,
|
|
dropout=dropout
|
|
)
|
|
|
|
if device:
|
|
model = model.to(device)
|
|
|
|
# Print model info
|
|
total_params = sum(p.numel() for p in model.parameters())
|
|
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
print(f"Model: {model_type}")
|
|
print(f" Input size: {input_size}")
|
|
print(f" Hidden size: {hidden_size}")
|
|
print(f" Num layers: {num_layers}")
|
|
print(f" Dropout: {dropout}")
|
|
print(f" Total parameters: {total_params:,}")
|
|
print(f" Trainable parameters: {trainable_params:,}")
|
|
print(f" Device: {device}")
|
|
|
|
return model
|
|
|
|
|
|
# ============================================================================
|
|
# FEATURE ENGINEERING (from src/feature_engineering.py, simplified for inline)
|
|
# ============================================================================
|
|
|
|
def compute_ci_features(ci_series: pd.Series, dah_series: pd.Series = None) -> pd.DataFrame:
|
|
"""Compute all CI-based features (state, velocity, acceleration, min/max/range/std/CV)."""
|
|
features = pd.DataFrame(index=ci_series.index)
|
|
|
|
# State (moving averages)
|
|
features['CI_raw'] = ci_series
|
|
features['7d_MA'] = ci_series.rolling(window=7, min_periods=1).mean()
|
|
features['14d_MA'] = ci_series.rolling(window=14, min_periods=1).mean()
|
|
features['21d_MA'] = ci_series.rolling(window=21, min_periods=1).mean()
|
|
|
|
# Velocity (gradient of MA)
|
|
for window in [7, 14, 21]:
|
|
ma = ci_series.rolling(window=window, min_periods=1).mean()
|
|
features[f'{window}d_velocity'] = ma.diff() / 1.0 # Simplified gradient
|
|
|
|
# Acceleration (gradient of velocity)
|
|
for window in [7, 14, 21]:
|
|
ma = ci_series.rolling(window=window, min_periods=1).mean()
|
|
vel = ma.diff()
|
|
features[f'{window}d_acceleration'] = vel.diff()
|
|
|
|
# Min, Max, Range
|
|
for window in [7, 14, 21]:
|
|
features[f'{window}d_min'] = ci_series.rolling(window=window, min_periods=1).min()
|
|
features[f'{window}d_max'] = ci_series.rolling(window=window, min_periods=1).max()
|
|
features[f'{window}d_range'] = features[f'{window}d_max'] - features[f'{window}d_min']
|
|
|
|
# Std and CV
|
|
for window in [7, 14, 21]:
|
|
features[f'{window}d_std'] = ci_series.rolling(window=window, min_periods=1).std()
|
|
ma = ci_series.rolling(window=window, min_periods=1).mean()
|
|
features[f'{window}d_CV'] = features[f'{window}d_std'] / (ma + 1e-6)
|
|
|
|
# DAH normalized (Days After Harvest)
|
|
if dah_series is not None:
|
|
features['DAH_normalized'] = dah_series / 450.0
|
|
|
|
return features.fillna(0)
|
|
|
|
|
|
def extract_features(data_df: pd.DataFrame, feature_names: List[str], ci_column: str = 'FitData',
|
|
season_anchor_day: int = None, lookback_start: int = 0) -> np.ndarray:
|
|
"""
|
|
Extract and return specified features as numpy array.
|
|
|
|
Args:
|
|
data_df: DataFrame with Date and CI data (may be a window after a harvest)
|
|
feature_names: List of feature names to extract
|
|
ci_column: Name of CI column
|
|
season_anchor_day: Day in FULL sequence where this season started (for DAH reset)
|
|
DAH will be recalculated as: 1, 2, 3, ... from this point
|
|
lookback_start: Starting index in original full data (for season reset calculation)
|
|
|
|
Returns:
|
|
NumPy array of shape (len(data_df), len(feature_names))
|
|
"""
|
|
# Compute all CI features
|
|
ci_series = data_df[ci_column].astype(float)
|
|
|
|
# Compute DAH (age/days since season start) - NOT day-of-year!
|
|
# DAH is a continuous counter: 1, 2, 3, ..., 475 (doesn't cycle at 365)
|
|
# It only resets to 1 after a harvest is detected (new season)
|
|
dah_series = None
|
|
if 'DAH_normalized' in feature_names:
|
|
if season_anchor_day is not None and lookback_start >= season_anchor_day:
|
|
# Season was reset after harvest. Recalculate DAH as simple counter from 1
|
|
# This is a window starting at or after harvest, so DAH should be: 1, 2, 3, ...
|
|
dah_series = pd.Series(np.arange(1, len(data_df) + 1), index=data_df.index)
|
|
elif 'DAH' in data_df.columns:
|
|
# Use DAH directly from CSV - already calculated as continuous age counter
|
|
dah_series = pd.Series(data_df['DAH'].astype(float).values, index=data_df.index)
|
|
else:
|
|
# Fallback: create continuous age counter (1, 2, 3, ...)
|
|
dah_series = pd.Series(np.arange(1, len(data_df) + 1), index=data_df.index)
|
|
|
|
all_features = compute_ci_features(ci_series, dah_series)
|
|
|
|
# Select requested features
|
|
requested = [f for f in feature_names if f in all_features.columns]
|
|
if not requested:
|
|
raise ValueError(f"No valid features found. Requested: {feature_names}, Available: {all_features.columns.tolist()}")
|
|
|
|
return all_features[requested].values
|
|
|
|
|
|
# ============================================================================
|
|
# MAIN UTILITY FUNCTIONS
|
|
# ============================================================================
|
|
|
|
def load_model_and_config(model_dir: Path):
|
|
"""Load model, config, and scalers from a given directory."""
|
|
cwd = Path.cwd()
|
|
|
|
# Try different naming conventions
|
|
candidates = [
|
|
# Standard names
|
|
(model_dir / "config.json", model_dir / "model.pt", model_dir / "scalers.pkl"),
|
|
# Model 307 specific names
|
|
(model_dir / "model_config.json", model_dir / "model_307.pt", model_dir / "model_scalers.pkl"),
|
|
# CWD standard names
|
|
(cwd / "config.json", cwd / "model.pt", cwd / "scalers.pkl"),
|
|
# CWD Model 307 names
|
|
(cwd / "model_config.json", cwd / "model_307.pt", cwd / "model_scalers.pkl"),
|
|
]
|
|
|
|
config_file = model_file = scalers_file = None
|
|
for cfg, mdl, scl in candidates:
|
|
if cfg.exists() and mdl.exists() and scl.exists():
|
|
config_file, model_file, scalers_file = cfg, mdl, scl
|
|
print(f"Found model files in: {cfg.parent}")
|
|
break
|
|
|
|
if not (config_file and model_file and scalers_file):
|
|
missing = []
|
|
for cfg, mdl, scl in candidates:
|
|
if not cfg.exists():
|
|
missing.append(str(cfg))
|
|
if not mdl.exists():
|
|
missing.append(str(mdl))
|
|
if not scl.exists():
|
|
missing.append(str(scl))
|
|
raise FileNotFoundError(
|
|
f"Missing model files. Checked multiple locations. Missing: {missing}"
|
|
)
|
|
|
|
with open(config_file) as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = create_model(
|
|
model_type=config['model']['type'],
|
|
input_size=len(config['features']),
|
|
hidden_size=config['model']['hidden_size'],
|
|
num_layers=config['model']['num_layers'],
|
|
dropout=config['model']['dropout'],
|
|
device=device
|
|
)
|
|
|
|
print(f"Loading weights from: {model_file}")
|
|
model.load_state_dict(torch.load(model_file, map_location=device, weights_only=False))
|
|
model.eval()
|
|
|
|
with open(scalers_file, 'rb') as f:
|
|
scalers = pickle.load(f)
|
|
|
|
return model, config, scalers
|
|
|
|
|
|
def load_harvest_data(data_file: Path) -> pd.DataFrame:
|
|
"""Load harvest data CSV."""
|
|
print(f"Loading data from: {data_file}")
|
|
df = pd.read_csv(data_file)
|
|
print(f"Loaded {len(df)} rows")
|
|
return df
|
|
|
|
|
|
def run_phase1_growing_window(field_data, model, config, scalers, ci_column, device,
|
|
threshold=0.45, consecutive_days=2):
|
|
"""
|
|
Phase 1: Growing window detection with DOY season reset.
|
|
|
|
For each detected harvest, reset DOY counter for the next season.
|
|
This allows the model to detect multiple consecutive harvests in multi-year data.
|
|
"""
|
|
harvest_dates = []
|
|
season_anchor_day = 0
|
|
current_pos = 0
|
|
|
|
while current_pos < len(field_data):
|
|
consecutive_above_threshold = 0
|
|
min_window_size = 120
|
|
|
|
for window_end in range(current_pos + 1, len(field_data) + 1):
|
|
window_data = field_data.iloc[current_pos:window_end].copy().reset_index(drop=True)
|
|
|
|
if len(window_data) < min_window_size:
|
|
continue
|
|
|
|
try:
|
|
reset_doy = current_pos > season_anchor_day
|
|
|
|
features = extract_features(
|
|
window_data,
|
|
config['features'],
|
|
ci_column=ci_column,
|
|
season_anchor_day=season_anchor_day if reset_doy else None,
|
|
lookback_start=current_pos
|
|
)
|
|
|
|
features_scaled = features.copy().astype(float)
|
|
for fi, scaler in enumerate(scalers):
|
|
try:
|
|
features_scaled[:, fi] = scaler.transform(features[:, fi].reshape(-1, 1)).flatten()
|
|
except Exception:
|
|
pass
|
|
|
|
with torch.no_grad():
|
|
x_tensor = torch.tensor(features_scaled, dtype=torch.float32).unsqueeze(0).to(device)
|
|
imminent_probs, detected_probs = model(x_tensor)
|
|
|
|
last_prob = detected_probs[0, -1].item()
|
|
|
|
if last_prob > threshold:
|
|
consecutive_above_threshold += 1
|
|
else:
|
|
consecutive_above_threshold = 0
|
|
|
|
if consecutive_above_threshold >= consecutive_days:
|
|
harvest_idx = current_pos + window_end - consecutive_days
|
|
harvest_date = field_data.iloc[harvest_idx]['Date']
|
|
|
|
harvest_dates.append((harvest_date, harvest_idx))
|
|
season_anchor_day = harvest_idx + 1
|
|
current_pos = harvest_idx + 1
|
|
break
|
|
|
|
except Exception as e:
|
|
continue
|
|
else:
|
|
break
|
|
|
|
return harvest_dates
|
|
|
|
|
|
def run_phase2_refinement(field_data, phase1_harvests, model, config, scalers, ci_column, device):
|
|
"""
|
|
Phase 2: Refinement with ±40 day window.
|
|
For each Phase 1 harvest, extract window and refine with argmax.
|
|
Returns list of (harvest_date, harvest_idx) tuples.
|
|
"""
|
|
refined_harvests = []
|
|
field_data = field_data.sort_values('Date').reset_index(drop=True)
|
|
|
|
for i, (phase1_harvest_date, phase1_idx) in enumerate(phase1_harvests):
|
|
try:
|
|
# Determine season start
|
|
if i == 0:
|
|
season_start_date = field_data.iloc[0]['Date']
|
|
else:
|
|
prev_harvest_idx = phase1_harvests[i-1][1]
|
|
season_start_idx = prev_harvest_idx + 1
|
|
if season_start_idx >= len(field_data):
|
|
break
|
|
season_start_date = field_data.iloc[season_start_idx]['Date']
|
|
|
|
# Extract ±40 day window
|
|
window_start_date = season_start_date - pd.Timedelta(days=40)
|
|
window_end_date = phase1_harvest_date + pd.Timedelta(days=40)
|
|
|
|
# FIXED: Use proper index selection
|
|
mask_start = field_data['Date'] >= window_start_date
|
|
mask_end = field_data['Date'] <= window_end_date
|
|
|
|
if mask_start.any():
|
|
window_start_idx = mask_start.idxmax() # First True index
|
|
else:
|
|
window_start_idx = 0
|
|
|
|
if mask_end.any():
|
|
# Last True index: find where condition becomes False from the right
|
|
true_indices = np.where(mask_end)[0]
|
|
window_end_idx = true_indices[-1] + 1 # +1 for slicing (exclusive end)
|
|
else:
|
|
window_end_idx = len(field_data)
|
|
|
|
if window_end_idx <= window_start_idx:
|
|
refined_harvests.append((phase1_harvest_date, phase1_idx))
|
|
continue
|
|
|
|
window_data = field_data.iloc[window_start_idx:window_end_idx].copy().reset_index(drop=True)
|
|
|
|
# Extract features for full window
|
|
features = extract_features(window_data, config['features'], ci_column=ci_column)
|
|
|
|
# Apply scalers
|
|
for fi, scaler in enumerate(scalers):
|
|
try:
|
|
features[:, fi] = scaler.transform(features[:, fi].reshape(-1, 1)).flatten()
|
|
except Exception:
|
|
pass
|
|
|
|
# Run model once on full window
|
|
with torch.no_grad():
|
|
x_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(device)
|
|
imminent_probs, detected_probs = model(x_tensor)
|
|
detected_probs = detected_probs.squeeze(0).cpu().numpy()
|
|
|
|
# Find refined harvest (argmax in window)
|
|
refined_idx_in_window = int(np.argmax(detected_probs))
|
|
refined_idx_global = window_start_idx + refined_idx_in_window
|
|
refined_harvest_date = field_data.iloc[refined_idx_global]['Date']
|
|
|
|
refined_harvests.append((refined_harvest_date, refined_idx_global))
|
|
|
|
except Exception:
|
|
refined_harvests.append((phase1_harvest_date, phase1_idx))
|
|
|
|
return refined_harvests
|
|
|
|
|
|
def run_two_step_refinement(df: pd.DataFrame, model, config, scalers, device=None,
|
|
phase1_threshold=0.45, phase1_consecutive=2):
|
|
"""
|
|
Two-step harvest detection for each field:
|
|
1. Phase 1: Growing window with threshold confirmation
|
|
2. Phase 2: ±40 day refinement with argmax
|
|
|
|
Args:
|
|
phase1_threshold (float): Probability threshold for Phase 1 (default 0.45, tuned for Model 307)
|
|
phase1_consecutive (int): Consecutive days required (default 2, reduced from 3 for robustness)
|
|
|
|
Returns list of dicts with field, season_start_date, season_end_date, etc.
|
|
"""
|
|
if device is None:
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
results = []
|
|
# CI column is 'FitData' (interpolated CI) - NOT 'value' (raw with NAs)
|
|
# 'FitData' is already gap-filled by R stage 03, ready for ML
|
|
ci_column = 'FitData'
|
|
|
|
# Group by field and count total fields for progress
|
|
field_groups = list(df.groupby('field'))
|
|
total_fields = len(field_groups)
|
|
harvests_found = 0
|
|
|
|
print(f" Processing {total_fields} fields...")
|
|
print(f" Phase 1 parameters: threshold={phase1_threshold}, consecutive_days={phase1_consecutive}")
|
|
|
|
for idx, (field, field_data) in enumerate(field_groups, 1):
|
|
# Simple progress indicator
|
|
pct = int((idx / total_fields) * 100)
|
|
bar_length = 40
|
|
filled = int((idx / total_fields) * bar_length)
|
|
bar = "█" * filled + "░" * (bar_length - filled)
|
|
print(f" [{bar}] {pct:3d}% ({idx}/{total_fields} fields)", end='\r')
|
|
|
|
field_data = field_data.sort_values('Date').reset_index(drop=True)
|
|
|
|
# Phase 1: Growing window detection
|
|
phase1_harvests = run_phase1_growing_window(field_data, model, config, scalers, ci_column, device,
|
|
threshold=phase1_threshold, consecutive_days=phase1_consecutive)
|
|
|
|
if not phase1_harvests:
|
|
continue
|
|
|
|
# Phase 2: Refinement
|
|
phase2_harvests = run_phase2_refinement(field_data, phase1_harvests, model, config, scalers, ci_column, device)
|
|
|
|
# Store results
|
|
for i, (harvest_date, harvest_idx) in enumerate(phase2_harvests):
|
|
if i == 0:
|
|
season_start_date = field_data.iloc[0]['Date']
|
|
else:
|
|
prev_harvest_idx = phase2_harvests[i-1][1]
|
|
season_start_idx = prev_harvest_idx + 1
|
|
if season_start_idx >= len(field_data):
|
|
break
|
|
season_start_date = field_data.iloc[season_start_idx]['Date']
|
|
|
|
season_end_date = harvest_date
|
|
|
|
result = {
|
|
'field': field,
|
|
'season': i + 1,
|
|
'season_start_date': season_start_date,
|
|
'season_end_date': season_end_date,
|
|
'phase2_harvest_date': harvest_date,
|
|
}
|
|
results.append(result)
|
|
harvests_found += 1
|
|
|
|
print() # New line after progress bar
|
|
print(f" ✓ Complete: Found {harvests_found} harvest events across {total_fields} fields")
|
|
if results:
|
|
print(f" Sample harvest dates: {results[0]['phase2_harvest_date']}")
|
|
if len(results) > 1:
|
|
print(f" {results[-1]['phase2_harvest_date']}")
|
|
|
|
return results
|
|
|
|
|
|
def build_production_harvest_table(refined_results: List[Dict]) -> pd.DataFrame:
|
|
"""
|
|
Build a DataFrame from refined results with columns for production pipeline.
|
|
One row per field/season with season start and end dates (formatted as YYYY-MM-DD).
|
|
"""
|
|
if not refined_results:
|
|
print("WARNING: No refined results to build table")
|
|
return pd.DataFrame(columns=['field', 'season', 'season_start_date', 'season_end_date'])
|
|
|
|
# Build DataFrame
|
|
df = pd.DataFrame(refined_results)
|
|
|
|
# Ensure date columns are datetime
|
|
df['season_start_date'] = pd.to_datetime(df['season_start_date']).dt.strftime('%Y-%m-%d')
|
|
df['season_end_date'] = pd.to_datetime(df['season_end_date']).dt.strftime('%Y-%m-%d')
|
|
df['phase2_harvest_date'] = pd.to_datetime(df['phase2_harvest_date']).dt.strftime('%Y-%m-%d')
|
|
|
|
print(f"Built production table with {len(df)} field/season combinations")
|
|
|
|
return df
|