{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":4535,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":3327,"modelId":986}],"dockerImageVersionId":31234,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport math\nimport warnings\nfrom pathlib import Path\nfrom typing import Dict, List, Tuple, Optional\n\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom sklearn.model_selection import GroupKFold\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.optim as optim\n\nfrom transformers import AutoModel, AutoConfig\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nwarnings.filterwarnings('ignore')","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.638564Z","iopub.execute_input":"2025-12-30T00:30:55.639177Z","iopub.status.idle":"2025-12-30T00:30:55.644193Z","shell.execute_reply.started":"2025-12-30T00:30:55.639147Z","shell.execute_reply":"2025-12-30T00:30:55.643576Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# DINOv2 native resolution\nIMG_SIZE = 518\nPATCH_STRIDE = 400 \nORIG_H, ORIG_W = 1000, 2000\nDINO_PATH = '/kaggle/input/dinov2/pytorch/large/1'\n\n# Hyperparameters\nBATCH_SIZE = 4\nLEARNING_RATE = 1e-4\nWEIGHT_DECAY = 1e-5\nNUM_EPOCHS = 2\nPATIENCE = 10\n\n# Target names and weights (as per competition)\nTARGET_NAMES = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\nTARGET_WEIGHTS = {\n    'Dry_Green_g': 0.1,\n    'Dry_Dead_g': 0.1,\n    'Dry_Clover_g': 0.1,\n    'GDM_g': 0.2,\n    'Dry_Total_g': 0.5\n}\n\n# State encoding (one-hot)\nSTATES = ['NSW', 'Tas', 'Vic', 'WA']\nSTATE_TO_IDX = {s: i for i, s in enumerate(STATES)}\n\n# Dataset paths\nTRAIN_CSV = '/kaggle/input/csiro-biomass/train.csv'\nTRAIN_IMG_DIR = '/kaggle/input/csiro-biomass/'\n\nTEST_CSV = '/kaggle/input/csiro-biomass/test.csv'\nTEST_IMG_DIR = 'test_images/'\nMODEL_PATH = 'afhn_fold0_best.pth' \n\n# Device configuration\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {DEVICE}\")","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.645500Z","iopub.execute_input":"2025-12-30T00:30:55.645730Z","iopub.status.idle":"2025-12-30T00:30:55.664712Z","shell.execute_reply.started":"2025-12-30T00:30:55.645699Z","shell.execute_reply":"2025-12-30T00:30:55.663909Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ZeroInflatedLogNormalLoss(nn.Module):\n    \"\"\"\n    Zero-Inflated LogNormal (ZILN) Loss for semi-continuous targets.\n    \n    Combines:\n    1. Binary classification (zero vs non-zero)\n    2. LogNormal regression for positive values\n    \n    Model outputs 3 values per target:\n    - logit_prob: Logit for P(y > 0)\n    - mu: Mean of log-normal distribution\n    - sigma_raw: Raw scale parameter (transformed via softplus)\n    \"\"\"\n\n    def __init__(self, eps: float = 1e-6):\n        super().__init__()\n        self.eps = eps\n    \n    def forward(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Args:\n            preds: (batch, 3) - [logit_prob, mu, sigma_raw]\n            target: (batch,) - Ground truth values (>= 0)\n        \n        Returns:\n            Scalar loss value\n        \"\"\"\n        logit_prob = preds[:, 0]\n        mu = preds[:, 1]\n        sigma_raw = preds[:, 2]\n        \n        # Ensure sigma > 0\n        sigma = F.softplus(sigma_raw) + self.eps\n\n        # Binary indicator: 1 if y > 0, else 0\n        is_positive = (target > 0).float()\n        \n        # Classification loss (zero vs non-zero)\n        class_loss = F.binary_cross_entropy_with_logits(\n            logit_prob, is_positive, reduction='none'\n        )\n        \n        # Regression loss (LogNormal NLL for positives)\n        safe_target = torch.clamp(target, min=self.eps)\n        log_target = torch.log(safe_target)\n        \n        reg_loss = (\n            log_target +\n            torch.log(sigma) +\n            0.5 * math.log(2 * math.pi) +\n            (log_target - mu).pow(2) / (2 * sigma.pow(2))\n        )\n        \n        # Combined loss: classification always, regression only for positives\n        total_loss = class_loss + (is_positive * reg_loss)\n        \n        return total_loss.mean()\n\nclass TweedieLoss(nn.Module):\n    \"\"\"\n    Tweedie loss for compound Poisson-Gamma distribution.\n    Suitable for strictly positive, heavy-tailed targets like Total Biomass.\n    \"\"\"\n    def __init__(self, p: float = 1.5, epsilon: float = 1e-8):\n        super().__init__()\n        assert 1 < p < 2, \"Tweedie power p must be in (1, 2)\"\n        self.p = p # Power parameter: 1 < p < 2 for Compound Poisson-Gamma\n        self.epsilon = epsilon\n\n    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Args:\n            pred: Predicted mean (batch,)\n            target: Ground truth (batch,)\n        \"\"\"\n        # Ensure positivity\n        pred = torch.clamp(pred, min=self.epsilon) \n        target = torch.clamp(target, min=self.epsilon) \n        \n        # Tweedie deviance\n        term1 = -target * torch.pow(pred, 1 - self.p) / (1 - self.p)\n        term2 = torch.pow(pred, 2 - self.p) / (2 - self.p)\n        \n        loss = term1 + term2\n        return loss.mean()","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.665495Z","iopub.execute_input":"2025-12-30T00:30:55.665665Z","iopub.status.idle":"2025-12-30T00:30:55.683504Z","shell.execute_reply.started":"2025-12-30T00:30:55.665650Z","shell.execute_reply":"2025-12-30T00:30:55.682822Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# MULTIPLE INSTANCE LEARNING (MIL) MODULE\n# ============================================================================\n\nclass GatedAttentionMIL(nn.Module):\n    \"\"\"\n    Gated Attention mechanism for MIL aggregation.\n    \n    Based on: Ilse et al. (2018) \"Attention-based Deep Multiple Instance Learning\"\n    \"\"\"\n    def __init__(self, input_dim: int = 1024, hidden_dim: int = 256):\n        super().__init__()\n        self.attention_V = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.Tanh()\n        )\n        self.attention_U = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.Sigmoid()\n        )\n        self.attention_weights = nn.Linear(hidden_dim, 1)\n    \n    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Args:\n            x: (batch, n_patches, embed_dim)\n        \n        Returns:\n            aggregated: (batch, embed_dim) - Weighted sum of patches\n            attention: (batch, n_patches, 1) - Attention weights per patch\n        \"\"\"\n\n        # Gated attention computation\n        A_V = self.attention_V(x) # (batch, n_patches, hidden_dim)\n        A_U = self.attention_U(x) # (batch, n_patches, hidden_dim)\n\n        # Element-wise gating\n        A = self.attention_weights(A_V * A_U) # (batch, n_patches, 1)\n\n        # Softmax over patches\n        A = torch.softmax(A, dim=1) # (batch, n_patches, 1)\n\n        # Weighted aggregation\n        M = torch.bmm(A.transpose(1, 2), x) # (batch, 1, embed_dim)\n\n        return M.squeeze(1), A","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.684527Z","iopub.execute_input":"2025-12-30T00:30:55.684844Z","iopub.status.idle":"2025-12-30T00:30:55.703681Z","shell.execute_reply.started":"2025-12-30T00:30:55.684825Z","shell.execute_reply":"2025-12-30T00:30:55.703015Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class FiLM(nn.Module):\n    \"\"\"\n    Feature-wise Linear Modulation for metadata injection.\n    \n    Modulates visual features with metadata (state, height, NDVI):\n    \"\"\"\n\n    def __init__(self, meta_dim: int, feat_dim: int):\n        super().__init__()\n        self.scale = nn.Linear(meta_dim, feat_dim)\n        self.shift = nn.Linear(meta_dim, feat_dim)\n    \n    def forward(self, features: torch.Tensor, metadata: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Args:\n            features: (batch, feat_dim)\n            metadata: (batch, meta_dim)\n        \n        Returns:\n            modulated: (batch, feat_dim)\n        \"\"\"\n        gamma = self.scale(metadata) # (batch, feat_dim)\n        beta = self.shift(metadata)  # (batch, feat_dim)\n\n        return features * gamma + beta","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.704964Z","iopub.execute_input":"2025-12-30T00:30:55.705182Z","iopub.status.idle":"2025-12-30T00:30:55.720815Z","shell.execute_reply.started":"2025-12-30T00:30:55.705164Z","shell.execute_reply":"2025-12-30T00:30:55.720171Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class AFHN(nn.Module):\n    \"\"\"\n    Agri-Foundational Hurdle Network for biomass prediction.\n    \n    Architecture:\n    1. DINOv2-Large (frozen) backbone for patch features\n    2. Gated Attention MIL for patch aggregation\n    3. FiLM layers for metadata modulation\n    4. Hierarchical heads with physical constraints\n    \"\"\"\n\n    def __init__(\n        self,\n        num_components: int = 5,\n        meta_dim: int = 10,\n        use_ziln: bool = True \n    ):\n        super().__init__()\n        self.num_components = num_components\n        self.use_ziln = use_ziln\n\n        # 1. Visual backbone: DINOv2-Large (frozen)\n        print(\"Loading DINOv2-Large backbone from local storage...\")\n        # Load Config and Model offline\n        config = AutoConfig.from_pretrained(DINO_PATH)\n        self.backbone = AutoModel.from_pretrained(DINO_PATH, config=config)\n\n        # Freeze all backbone parameters\n        for param in self.backbone.parameters():\n            param.requires_grad = False\n        \n        self.vis_dim = 1024  # DINOv2-Large embedding dimension\n\n        # 2. MIL Aggregator\n        self.mil = GatedAttentionMIL(input_dim=self.vis_dim, hidden_dim=256)\n\n        # 3. Metadata Encoder + FiLM\n        self.meta_encoder = nn.Sequential(\n            nn.Linear(meta_dim, 128),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(128, 64)\n        )\n        self.film = FiLM(meta_dim=64, feat_dim=self.vis_dim)\n\n        # 4. Prediction Heads\n\n        # A. Total Biomass Head (ZILN or Tweedie)\n        if self.use_ziln:\n            self.head_total = nn.Sequential(\n                nn.Linear(self.vis_dim, 256),\n                nn.GELU(),\n                nn.Dropout(0.2),\n                nn.Linear(256, 3)  # [logit_prob, mu, sigma_raw]\n            )\n        else:\n            self.head_total = nn.Sequential(\n                nn.Linear(self.vis_dim, 256),\n                nn.GELU(),\n                nn.Dropout(0.2),\n                nn.Linear(256, 1),\n                nn.Softplus()  # Ensure positive prediction\n            )\n        \n        # B. Component Ratios Head (Softmax for sum-to-one constraint)\n        self.head_ratios = nn.Sequential(\n            nn.Linear(self.vis_dim, 256),\n            nn.GELU(),\n            nn.Dropout(0.2),\n            nn.Linear(256, num_components)\n        )\n\n        # C. Component Gates (Hurdle mechanism)\n        self.head_gate = nn.Sequential(\n            nn.Linear(self.vis_dim, 128),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(128, num_components)\n        )\n    \n    def forward_features(\n        self,\n        x_patches: torch.Tensor,\n        metadata: torch.Tensor\n    ) -> Tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Extract and aggregate patch features.\n        \n        Args:\n            x_patches: (batch, n_patches, 3, H, W)\n            metadata: (batch, meta_dim)\n        \n        Returns:\n            modulated_feat: (batch, vis_dim)\n            attn_weights: (batch, n_patches, 1)\n        \"\"\"\n        batch_size, n_patches, c, h, w = x_patches.shape\n\n        # Flaten patches for batch processing\n        x_flat = x_patches.view(batch_size * n_patches, c, h, w) # (B*n_patches, 3, H, W)\n\n        # Extract patch features (frozen backbone, no grad)\n        with torch.no_grad():\n            # Transformers expects 'pixel_values' argument\n            outputs = self.backbone(pixel_values=x_flat)\n            \n            # Get last_hidden_state: (B*n_patches, sequence_length, hidden_size)\n            last_hidden_state = outputs.last_hidden_state\n            \n            # The CLS token is usually at index 0\n            feat_flat = last_hidden_state[:, 0, :]  # (batch*n_patches, embed_dim)\n\n        # Reshape to (batch, n_patches, embed_dim)\n        feat_seq = feat_flat.view(batch_size, n_patches, -1)\n\n        # MIL aggregation\n        global_feat, attn_weights = self.mil(feat_seq) # (batch, vis_dim), (batch, n_patches, 1)\n\n        # Meta Modulation via FiLM\n        meta_emb = self.meta_encoder(metadata) # (batch, meta_feat_dim)\n        modulated_feat = self.film(global_feat, meta_emb) # (batch, vis_dim)\n\n        return modulated_feat, attn_weights\n\n    def forward(\n        self,\n        x_patches: torch.Tensor,\n        metadata: torch.Tensor\n    ) -> Dict[str, torch.Tensor]:\n        \"\"\"\n        Forward pass returning all predictions.\n        \n        Returns:\n            Dict with keys: 'total', 'ratios', 'gates'\n        \"\"\"\n        feat, _ = self.forward_features(x_patches, metadata)\n        \n        total = self.head_total(feat)\n        raw_ratios = self.head_ratios(feat)\n        ratios = F.softmax(raw_ratios, dim=1)\n        gate_logits = self.head_gate(feat)\n        \n        return {\n            'total': total,\n            'ratios': ratios,\n            'gates': gate_logits\n        }\n    \n    def predict_components(\n        self,\n        x_patches: torch.Tensor,\n        metadata: torch.Tensor\n    ) -> torch.Tensor:\n        \"\"\"\n        Inference: predict all 5 component biomass values.\n        \n        Returns:\n            components: (batch, 5) - Predicted biomass for each component\n        \"\"\"\n        outputs = self.forward(x_patches, metadata)\n\n        # Decode total biomass\n        if self.use_ziln:\n            total_ziln = outputs['total']\n            prob_nonzero = torch.sigmoid(total_ziln[:, 0])\n            mu = total_ziln[:, 1]\n            sigma = F.softplus(total_ziln[:, 2]) + 1e-6\n            # Expected valud of ZILN\n            expected_total = prob_nonzero * torch.exp(mu +0.5 * sigma.pow(2))\n        else:\n            expected_total = outputs['total'].squeeze(-1)\n        \n        # Decode gates (Bernoulli probability)\n        gates = torch.sigmoid(outputs['gates'])\n\n        # Apply gates to ratios and renormalize\n        gated_ratios = outputs['ratios'] * gates\n        sum_gated = gated_ratios.sum(dim=1, keepdim=True) + 1e-6\n        final_ratios = gated_ratios / sum_gated\n\n        # Component predictions (enforces sum constraint)\n        components = expected_total.unsqueeze(1) * final_ratios\n\n        return components","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.795760Z","iopub.execute_input":"2025-12-30T00:30:55.795983Z","iopub.status.idle":"2025-12-30T00:30:55.810002Z","shell.execute_reply.started":"2025-12-30T00:30:55.795964Z","shell.execute_reply":"2025-12-30T00:30:55.809419Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class PatchBiomassDataset(Dataset):\n    \"\"\"\n    Dataset for CSIRO Biomass with on-the-fly patch extraction.\n    \n    Each sample contains:\n    - Image patches (multiple instances)\n    - Metadata (height, NDVI, state, species, date)\n    - Target biomass values (5 components)\n    \"\"\"\n\n    def __init__(\n        self,\n        df: pd.DataFrame,\n        image_dir: str,\n        transform: Optional[A.Compose] = None,\n        patch_size: int = IMG_SIZE,\n        stride: int = PATCH_STRIDE,\n        is_test: bool = False\n    ):\n        # For test set, group by image_path (5 rows per image)\n        if is_test:\n            self.df = df.groupby('image_path').first().reset_index()\n        else:\n            # For train, group by sample_id (each image appears once)\n            self.df = df.groupby('sample_id').first().reset_index()\n        \n        self.image_dir = Path(image_dir)\n        self.transform = transform\n        self.patch_size = patch_size\n        self.stride = stride\n        self.is_test = is_test\n\n    def __len__(self) -> int:\n        return len(self.df)\n\n    def extract_patches(self, image: np.ndarray) -> torch.Tensor:\n        \"\"\"\n        Extract overlapping patches from high-res image.\n        \n        Args:\n            image: (H, W, 3) numpy array\n        \n        Returns:\n            patches: (n_patches, 3, patch_size, patch_size)\n        \"\"\"\n        h, w = image.shape[:2]\n        patches = []\n\n        for y in range(0, h - self.patch_size + 1, self.stride):\n            for x in range(0, w - self.patch_size + 1, self.stride):\n                patch = image[y:y+self.patch_size, x:x+self.patch_size]\n                \n                if self.transform:\n                    augmented = self.transform(image=patch)\n                    patch_tensor = augmented['image']\n                else:\n                    patch_tensor = torch.from_numpy(patch).permute(2, 0, 1).float() / 255.0\n                \n                patches.append(patch_tensor)\n        \n        return torch.stack(patches)\n\n\n    def __getitem__(self, idx: int) -> Dict:\n        row = self.df.iloc[idx]\n\n        # Load image\n        img_path = self.image_dir / row['image_path']\n        image = np.array(Image.open(img_path).convert('RGB'))\n\n        # Extract patches\n        patches = self.extract_patches(image)\n\n        # Encode metadata\n        # State one-hot\n        state_onehot = np.zeros(len(STATES))\n        if row['State'] in STATE_TO_IDX:\n            state_onehot[STATE_TO_IDX[row['State']]] = 1.0\n\n        # Date encoding (cyclical)\n        date = pd.to_datetime(row['Sampling_Date'])\n        month = date.month\n        month_sin = np.sin(2 * np.pi * month / 12)\n        month_cos = np.cos(2 * np.pi * month / 12)\n\n        # Combine metadata: [height, ndvi, state_onehot(4), month_sin, month_cos]\n        metadata = np.concatenate([\n            [row['Height_Ave_cm'], row['Pre_GSHH_NDVI']],\n            state_onehot,\n            [month_sin, month_cos]\n        ]).astype(np.float32)\n\n        metadata = torch.from_numpy(metadata)\n\n        # Targets (if not test)\n        if not self.is_test:\n            targets = torch.tensor([\n                row['Dry_Green_g'],\n                row['Dry_Dead_g'],\n                row['Dry_Clover_g'],\n                row['GDM_g'],\n                row['Dry_Total_g']\n            ], dtype=torch.float32)\n            \n            return {\n                'patches': patches,\n                'metadata': metadata,\n                'targets': targets,\n                'sample_id': row['sample_id']\n            }\n        else:\n            return {\n                'patches': patches,\n                'metadata': metadata,\n                'sample_id': row.name # Index for test\n            }","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.811203Z","iopub.execute_input":"2025-12-30T00:30:55.811563Z","iopub.status.idle":"2025-12-30T00:30:55.831900Z","shell.execute_reply.started":"2025-12-30T00:30:55.811526Z","shell.execute_reply":"2025-12-30T00:30:55.831315Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_transforms(is_train: bool = True) -> A.Compose:\n    \"\"\"\n    Albumentations transforms for image augmentation.\n    \n    DINOv2 expects ImageNet normalization.\n    \"\"\"\n    if is_train:\n        return A.Compose([\n            A.RandomResizedCrop(\n                size=(IMG_SIZE, IMG_SIZE),\n                scale=(0.7, 1.0), \n                ratio=(0.9, 1.1), \n                p=1.0\n            ),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            A.Transpose(p=0.5),\n\n            A.OneOf([\n                A.MotionBlur(p=0.2),\n                A.GaussianBlur(p=0.2),\n            ], p=0.2),\n\n            # Photometric (cautious with hue to avoid confusing green/dead)\n            A.OneOf([\n                A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05, p=0.3),\n                A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),\n                A.RandomGamma(gamma_limit=(80, 120), p=0.3),\n                A.HueSaturationValue(p=0.3),\n            ], p=0.3),\n\n            # Noise and regularization\n            A.GaussNoise(var_limit=(10.0, 50.0), p=0.2),\n            A.CoarseDropout(max_holes=8, max_height=32, max_width=32, fill_value=0, p=0.3),\n\n            # Normalization (DINOv2 uses ImageNet stats)\n            A.Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225]\n            ),\n            ToTensorV2()\n        ])\n    else:\n        return A.Compose([\n            A.Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225]\n            ),\n            ToTensorV2()\n        ])","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.832666Z","iopub.execute_input":"2025-12-30T00:30:55.832897Z","iopub.status.idle":"2025-12-30T00:30:55.852001Z","shell.execute_reply.started":"2025-12-30T00:30:55.832867Z","shell.execute_reply":"2025-12-30T00:30:55.851394Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def weighted_r2_score(\n    y_true: np.ndarray,\n    y_pred: np.ndarray,\n    weights: np.ndarray\n) -> float:\n    \"\"\"\n    Compute weighted R2 score as per competition metric.\n    \n    R2_w = 1 - (SS_res_w / SS_tot_w)\n    \"\"\"\n    # Weighted mean\n    y_mean = np.average(y_true, weights=weights)\n\n    # Weighted residual sum of squares\n    ss_res = np.sum(weights * (y_true - y_pred) ** 2)\n\n    # Weighted total sum of squares\n    ss_tot = np.sum(weights * (y_true - y_mean) ** 2)\n\n    r2 = 1 - (ss_res / ss_tot + 1e-8)\n    return r2\n\ndef compute_metrics(\n    targets: torch.Tensor,\n    predictions: torch.Tensor\n) -> Dict[str, float]:\n    \"\"\"\n    Compute all evaluation metrics.\n    \n    Args:\n        targets: (n_samples, 5)\n        predictions: (n_samples, 5)\n    \n    Returns:\n        Dict of metric names to values\n    \"\"\"\n    targets_np = targets.cpu().numpy()\n    preds_np = predictions.cpu().numpy()\n\n    # Per-component MAE\n    component_mae = {}\n    for i, name in enumerate(TARGET_NAMES):\n        mae = np.abs(targets_np[:, i] - preds_np[:, i]).mean()\n        component_mae[f'{name}_MAE'] = mae\n\n    # Weighted R2\n    # Expand weights to match row-wise structure\n    n_samples = len(targets_np)\n    weights_expanded = np.array([TARGET_WEIGHTS[name] for name in TARGET_NAMES])\n    weights_tiled = np.tile(weights_expanded, n_samples)\n    \n    targets_flat = targets_np.flatten()\n    preds_flat = preds_np.flatten()\n    \n    r2_weighted = weighted_r2_score(targets_flat, preds_flat, weights_tiled)\n    \n    return {\n        **component_mae,\n        'Weighted_R2': r2_weighted\n    }","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.853355Z","iopub.execute_input":"2025-12-30T00:30:55.853621Z","iopub.status.idle":"2025-12-30T00:30:55.872057Z","shell.execute_reply.started":"2025-12-30T00:30:55.853602Z","shell.execute_reply":"2025-12-30T00:30:55.871423Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_epoch(\n    model: AFHN,\n    dataloader: DataLoader,\n    optimizer: optim.Optimizer,\n    criterion_ziln: nn.Module, # This was criterion_components\n    device: torch.device\n) -> float:\n    \"\"\"Train for one epoch.\"\"\"\n    model.train()\n    total_loss = 0\n    \n    pbar = tqdm(dataloader, desc='Training')\n    for batch in pbar:\n        patches = batch['patches'].to(device)\n        metadata = batch['metadata'].to(device)\n        targets = batch['targets'].to(device)\n        \n        optimizer.zero_grad()\n        \n        pred_components = model.predict_components(patches, metadata)\n        loss_components = F.mse_loss(pred_components, targets)\n        \n        outputs = model(patches, metadata)\n        target_total = targets[:, -1] \n        \n        loss_total_biomass = criterion_ziln(outputs['total'], target_total)\n        \n        target_gates = (targets > 0).float()\n        loss_gate = F.binary_cross_entropy_with_logits(\n            outputs['gates'], target_gates\n        )\n        \n        gt_total = target_total.unsqueeze(1) + 1e-6\n        gt_ratios = targets / gt_total # Shape: (B, 5)\n        \n        pred_gates = torch.sigmoid(outputs['gates'])\n        \n        gated_ratios = outputs['ratios'] * pred_gates\n        sum_gated = gated_ratios.sum(dim=1, keepdim=True) + 1e-6\n        pred_effective_ratios = gated_ratios / sum_gated\n        \n        # Use MSE only for the ratios (0-1 range), which is stable\n        loss_ratios = F.mse_loss(pred_effective_ratios, gt_ratios)\n        \n\n        loss = loss_total_biomass + loss_ratios + 0.5 * loss_gate\n        \n        if torch.isnan(loss):\n            print(\"Warning: NaN loss detected!\")\n            continue\n\n        loss.backward()\n        \n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n        optimizer.step()\n        \n        total_loss += loss.item()\n        pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n    \n    return total_loss / len(dataloader)\n\ndef validate(\n    model: AFHN,\n    dataloader: DataLoader,\n    device: torch.device\n) -> Tuple[float, Dict[str, float]]:\n    \"\"\"Validate model.\"\"\"\n    model.eval()\n\n    all_targets = []\n    all_predictions = []\n\n    with torch.no_grad():\n        for batch in tqdm(dataloader, desc='Validation'):\n            patches = batch['patches'].to(device)\n            metadata = batch['metadata'].to(device)\n            targets = batch['targets']\n\n            predictions = model.predict_components(patches, metadata)\n\n            all_targets.append(targets)\n            all_predictions.append(predictions.cpu())\n    \n    all_targets = torch.cat(all_targets)\n    all_predictions = torch.cat(all_predictions)\n\n    metrics = compute_metrics(all_targets, all_predictions)\n\n    return metrics['Weighted_R2'], metrics","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.872830Z","iopub.execute_input":"2025-12-30T00:30:55.873106Z","iopub.status.idle":"2025-12-30T00:30:55.890542Z","shell.execute_reply.started":"2025-12-30T00:30:55.873073Z","shell.execute_reply":"2025-12-30T00:30:55.889970Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def main():\n    \"\"\"Main training function.\"\"\"\n    # Load data\n    print(\"Loading training data...\")\n    train_df = pd.read_csv(TRAIN_CSV)\n    \n    train_df['sample_id'] = train_df['sample_id'].str.split('__').str[0]\n    \n    # Pivot from long to wide format\n    train_df_wide = train_df.pivot_table(\n        index=['sample_id', 'image_path', 'Sampling_Date', 'State',\n               'Species', 'Pre_GSHH_NDVI', 'Height_Ave_cm'],\n        columns='target_name',\n        values='target'\n    ).reset_index()\n\n    if train_df_wide.isnull().values.any():\n        print(f\"Warning: Found {train_df_wide.isnull().sum().sum()} missing values in pivot table.\")\n        print(\"Filling missing target values with 0.0 (assuming missing = no biomass).\")\n        train_df_wide = train_df_wide.fillna(0.0)\n    else:\n        print(\"Success: No missing values found in pivot table!\")\n    \n    # Group K-Fold split (by State to ensure generalization)\n    groups = train_df_wide['State'].values\n    n_groups = len(np.unique(groups))\n    print(f\"Found {n_groups} unique states: {np.unique(groups)}\")\n    \n    # Adjust n_splits to not exceed the number of groups\n    n_splits = 5\n    if n_groups < n_splits:\n        print(f\"Warning: Requested {n_splits} splits but only found {n_groups} groups.\")\n        print(f\"Adjusting n_splits to {n_groups} (Leave-One-Group-Out CV).\")\n        n_splits = n_groups\n\n    # Initialize GroupKFold with the safe number of splits\n    gkf = GroupKFold(n_splits=n_splits)\n    \n    fold = 0\n    for train_idx, val_idx in gkf.split(train_df_wide, groups=groups):\n        print(f\"\\n{'='*60}\")\n        print(f\"Fold {fold + 1}/5\")\n        print(f\"{'='*60}\")\n        \n        train_fold = train_df_wide.iloc[train_idx]\n        val_fold = train_df_wide.iloc[val_idx]\n        \n        # Datasets\n        train_dataset = PatchBiomassDataset(\n            train_fold, TRAIN_IMG_DIR,\n            transform=get_transforms(is_train=True)\n        )\n        val_dataset = PatchBiomassDataset(\n            val_fold, TRAIN_IMG_DIR,\n            transform=get_transforms(is_train=False)\n        )\n        \n        train_loader = DataLoader(\n            train_dataset, batch_size=BATCH_SIZE,\n            shuffle=True, num_workers=4, pin_memory=True\n        )\n        val_loader = DataLoader(\n            val_dataset, batch_size=BATCH_SIZE,\n            shuffle=False, num_workers=4, pin_memory=True\n        )\n        \n        # Model\n        model = AFHN(num_components=5, meta_dim=8, use_ziln=True).to(DEVICE)\n        \n        # Optimizer & Scheduler\n        optimizer = optim.AdamW(\n            model.parameters(),\n            lr=LEARNING_RATE,\n            weight_decay=WEIGHT_DECAY\n        )\n        scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(\n            optimizer, T_0=10, T_mult=2\n        )\n        \n        # Loss functions\n        criterion_ziln = ZeroInflatedLogNormalLoss(eps=1e-6)\n        criterion_total = TweedieLoss(p=1.5)\n        \n        # Training loop\n        best_r2 = -float('inf')\n        patience_counter = 0\n        \n        for epoch in range(NUM_EPOCHS):\n            print(f\"\\nEpoch {epoch + 1}/{NUM_EPOCHS}\")\n            \n            train_loss = train_epoch(\n                model, \n                train_loader, \n                optimizer,\n                criterion_ziln,          \n                DEVICE\n            )\n            \n            val_r2, val_metrics = validate(model, val_loader, DEVICE)\n            \n            scheduler.step()\n            \n            print(f\"Train Loss: {train_loss:.4f}\")\n            print(f\"Val Weighted R2: {val_r2:.4f}\")\n            for k, v in val_metrics.items():\n                if k != 'Weighted_R2':\n                    print(f\"  {k}: {v:.4f}\")\n            \n            # Early stopping\n            if val_r2 > best_r2:\n                best_r2 = val_r2\n                patience_counter = 0\n                torch.save(model.state_dict(), f'afhn_fold{fold}_best.pth')\n                print(f\"Saved best model (R2={best_r2:.4f})\")\n            else:\n                patience_counter += 1\n                if patience_counter >= PATIENCE:\n                    print(f\"Early stopping at epoch {epoch + 1}\")\n                    break\n        \n        print(f\"\\nFold {fold + 1} Best R2: {best_r2:.4f}\")\n        fold += 1\n        \n        # Only train one fold for demo (remove break for full CV)\n        break\n\n\ndef inference_example():\n    \"\"\"Example inference on test set.\"\"\"\n    # Load test data\n    test_df = pd.read_csv(TEST_CSV)\n\n    test_dataset = PatchBiomassDataset(\n        test_df, TEST_IMG_DIR,\n        transform=get_transforms(is_train=False),\n        is_test=True\n    )\n\n    test_loader = DataLoader(\n        test_dataset, batch_size=4,\n        shuffle=False, num_workers=4\n    )\n\n    # Load model\n    model = AFHN(num_components=5, meta_dim=8, use_ziln=True).to(DEVICE)\n    model.load_state_dict(torch.load(MODEL_PATH))\n\n    # Predict\n    predictions = []\n    sample_ids = []\n\n    with torch.no_grad():\n        for batch in tqdm(test_loader, desc='Inference'):\n            patches = batch['patches'].to(DEVICE)\n            metadata = batch['metadata'].to(DEVICE)\n            \n            preds = model.predict_components(patches, metadata)\n            predictions.append(preds.cpu().numpy())\n            sample_ids.extend(batch['sample_id'])\n    \n    predictions = np.vstack(predictions)\n\n    # Create submission\n    submission_rows = []\n\n    for i, img_id in enumerate(sample_ids):\n        for j, target_name in enumerate(TARGET_NAMES):\n            submission_rows.append({\n                'sample_id': f\"{img_id}__{target_name}\",\n                'target': predictions[i, j]\n            })\n    \n    submission_df = pd.DataFrame(submission_rows)\n    submission_df.to_csv('submission.csv', index=False)\n    print(\"Submission saved to submission.csv\")","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.891531Z","iopub.execute_input":"2025-12-30T00:30:55.891780Z","iopub.status.idle":"2025-12-30T00:30:55.911110Z","shell.execute_reply.started":"2025-12-30T00:30:55.891761Z","shell.execute_reply":"2025-12-30T00:30:55.910526Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == '__main__':\n    main()\n    # inference_example()","metadata":{"execution":{"iopub.status.busy":"2025-12-30T00:30:55.952145Z","iopub.execute_input":"2025-12-30T00:30:55.952717Z","iopub.status.idle":"2025-12-30T01:25:43.605255Z","shell.execute_reply.started":"2025-12-30T00:30:55.952697Z","shell.execute_reply":"2025-12-30T01:25:43.604108Z"},"trusted":true},"outputs":[],"execution_count":null}]}