{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.11.13"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"d1032cdc","cell_type":"markdown","source":"# MABe Mouse Behavior Detection - Training Notebook\n\n**XGBoost Model Training with Optuna Hyperparameter Optimization**\n\nThis notebook trains XGBoost models for mouse behavior detection using cross-validation and hyperparameter optimization.\n","metadata":{}},{"id":"786f7a40","cell_type":"markdown","source":"## Kaggle Setup\n\n**Note for Kaggle Users:**\n- This notebook automatically detects Kaggle environment and uses `/kaggle/input/MABe-mouse-behavior-detection` for data\n- Enable GPU in notebook settings for faster training (Settings → Accelerator → GPU)\n- Enable Internet in notebook settings to allow automatic dataset creation (Settings → Internet → ON)\n- Models and thresholds will be saved to `/kaggle/working/` (available for download)\n- To use Weights & Biases logging, add `WANDB_API_KEY` to Kaggle secrets and set `USE_WANDB = True` in configuration\n","metadata":{}},{"id":"d3a5cda2","cell_type":"code","source":"import argparse\nimport json\nimport os\nimport time\nimport warnings\nfrom collections import defaultdict\nfrom pathlib import Path\nimport joblib\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport optuna\nfrom optuna.pruners import MedianPruner\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom sklearn.metrics import f1_score, precision_score, recall_score\nfrom sklearn.base import clone\nfrom sklearn.impute import SimpleImputer\nfrom xgboost import XGBClassifier\nfrom tqdm import tqdm\n# from koolbox import Trainer  # Replaced with custom implementation below\nimport itertools\nimport glob\nimport gc\ntry:\n    from imblearn.over_sampling import SMOTE\n    SMOTE_AVAILABLE = True\nexcept ImportError:\n    SMOTE_AVAILABLE = False\n    print(\"Warning: imbalanced-learn not installed. Install with: pip install imbalanced-learn\")\n\ntry:\n    import wandb\n    WANDB_AVAILABLE = True\nexcept ImportError:\n    WANDB_AVAILABLE = False\n    print(\"Warning: wandb not installed. Install with: pip install wandb\")\n\noptuna.logging.set_verbosity(optuna.logging.WARNING)\nwarnings.filterwarnings('ignore')","metadata":{},"outputs":[],"execution_count":null},{"id":"4dfb9478","cell_type":"code","source":"# Install kagglehub for automatic dataset creation\n# This allows the notebook to automatically create a Kaggle dataset from trained models\ntry:\n    import kagglehub\n    print(\"[OK] kagglehub already installed\")\nexcept ImportError:\n    print(\"Installing kagglehub...\")\n    !pip install kagglehub -q\n    import kagglehub\n    print(\"[OK] kagglehub installed successfully\")\n","metadata":{},"outputs":[],"execution_count":null},{"id":"91f95b41","cell_type":"code","source":"# Custom Trainer class (koolbox replacement for Kaggle compatibility)\n# koolbox is not available on Kaggle, so we provide a custom implementation\n\nclass Trainer:\n    \"\"\"\n    Custom Trainer class to replace koolbox.Trainer\n    Performs cross-validation and stores out-of-fold predictions and trained models.\n    \"\"\"\n    def __init__(self, estimator, cv, cv_args=None, metric=None, task=\"binary\", \n                 verbose=False, save=False, save_path=None):\n        self.estimator = estimator\n        self.cv = cv\n        self.cv_args = cv_args or {}\n        self.metric = metric\n        self.task = task\n        self.verbose = verbose\n        self.save = save\n        self.save_path = save_path\n        self.oof_preds = None\n        self.estimators = []\n    \n    def fit(self, X, y):\n        \"\"\"Perform cross-validation and store OOF predictions and models\"\"\"\n        groups = self.cv_args.get('groups', None)\n        \n        # Initialize OOF predictions\n        if self.task == \"binary\":\n            self.oof_preds = np.zeros(len(X))\n        else:\n            n_classes = len(np.unique(y))\n            self.oof_preds = np.zeros((len(X), n_classes))\n        \n        # Perform cross-validation\n        for fold_idx, (train_idx, val_idx) in enumerate(self.cv.split(X, y, groups=groups)):\n            # Clone estimator for this fold\n            fold_estimator = clone(self.estimator)\n            \n            # Get fold data\n            X_train_fold = X.iloc[train_idx] if hasattr(X, 'iloc') else X[train_idx]\n            X_val_fold = X.iloc[val_idx] if hasattr(X, 'iloc') else X[val_idx]\n            y_train_fold = y[train_idx] if hasattr(y, 'iloc') else y[train_idx]\n            y_val_fold = y[val_idx] if hasattr(y, 'iloc') else y[val_idx]\n            \n            # Train model\n            fold_estimator.fit(X_train_fold, y_train_fold)\n            self.estimators.append(fold_estimator)\n            \n            # Get predictions\n            if hasattr(fold_estimator, 'predict_proba'):\n                if self.task == \"binary\":\n                    y_pred_proba = fold_estimator.predict_proba(X_val_fold)[:, 1]\n                else:\n                    y_pred_proba = fold_estimator.predict_proba(X_val_fold)\n            else:\n                y_pred_proba = fold_estimator.predict(X_val_fold)\n            \n            # Store OOF predictions\n            if self.task == \"binary\":\n                self.oof_preds[val_idx] = y_pred_proba\n            else:\n                self.oof_preds[val_idx] = y_pred_proba\n            \n            # Save model if requested\n            if self.save and self.save_path:\n                os.makedirs(self.save_path, exist_ok=True)\n                model_path = os.path.join(self.save_path, f\"fold_{fold_idx}_trainer_model.pkl\")\n                joblib.dump(fold_estimator, model_path)\n    \n    def predict(self, X):\n        \"\"\"Predict using ensemble of all CV models\"\"\"\n        if not self.estimators:\n            raise ValueError(\"Model must be fitted before prediction\")\n        \n        predictions = []\n        for estimator in self.estimators:\n            if hasattr(estimator, 'predict_proba'):\n                pred = estimator.predict_proba(X)\n            else:\n                pred = estimator.predict(X)\n            predictions.append(pred)\n        \n        # Average predictions across folds\n        return np.mean(predictions, axis=0)\n","metadata":{},"outputs":[],"execution_count":null},{"id":"e11801d2","cell_type":"code","source":"# Configuration\nimport os\nfrom pathlib import Path\n\n# Detect Kaggle environment\nIS_KAGGLE = os.path.exists('/kaggle/input')\n\nif IS_KAGGLE:\n    # Kaggle paths\n    TRAIN_PATH = '/kaggle/input/MABe-mouse-behavior-detection/train.csv'\n    TEST_PATH = '/kaggle/input/MABe-mouse-behavior-detection/test.csv'\n    TRAIN_ANNOTATION_PATH = '/kaggle/input/MABe-mouse-behavior-detection/train_annotation'\n    TRAIN_TRACKING_PATH = '/kaggle/input/MABe-mouse-behavior-detection/train_tracking'\n    TEST_TRACKING_PATH = '/kaggle/input/MABe-mouse-behavior-detection/test_tracking'\n    MODEL_PATH = '/kaggle/input/social-action-recognition-in-mice-xgb-catboost'  # Update with your model dataset\n    MODEL_NAME = 'xgboost'\n    OUTPUT_DIR = '/kaggle/working'\nelse:\n    # Local paths (fallback)\n    TRAIN_PATH = 'train.csv'\n    TEST_PATH = 'test.csv'\n    TRAIN_ANNOTATION_PATH = 'train_annotation'\n    TRAIN_TRACKING_PATH = 'train_tracking'\n    TEST_TRACKING_PATH = 'test_tracking'\n    MODEL_PATH = 'xgboost'\n    MODEL_NAME = 'xgboost'\n    OUTPUT_DIR = '.'\n\n# Model hyperparameters\nN_SPLITS = 3\nN_ESTIMATORS = 250\nLEARNING_RATE = 0.08\nMAX_DEPTH = 6\nMIN_CHILD_WEIGHT = 5\nSUBSAMPLE = 0.8\nCOLSAMPLE_BYTREE = 0.8\n\n# Hyperparameter optimization\nOPTIMIZE_HYPERPARAMS = True\nHYPEROPT_TRIALS = 50\nHYPEROPT_TIMEOUT = None\n\n# Threshold tuning\nN_TRIALS = 200\n\n# Class imbalance handling\nUSE_CLASS_WEIGHTS = True\nUSE_SMOTE = True\nSMOTE_THRESHOLD = 0.01\n\n# Early stopping\nUSE_EARLY_STOPPING = True\nEARLY_STOPPING_ROUNDS = 10\n\n# Adaptive optimization\nADAPTIVE_TRIALS = False\nMIN_TRIALS = 20\nUSE_SECTION_SPECIFIC_RANGES = True\n\n# GPU/CUDA support\nUSE_GPU = os.path.exists('/kaggle/input')  # Auto-detect on Kaggle\n\n# Weights & Biases logging (optional)\nUSE_WANDB = False  # Set to True and add WANDB_API_KEY to Kaggle secrets\nWANDB_PROJECT = 'mabe-mouse-behavior'\nWANDB_ENTITY = None\nWANDB_RUN_NAME = None\nPREVIOUS_RUN_ID = None\n\n# Create CFG-like object for compatibility\nclass CFG:\n    def __init__(self):\n        self.train_path = TRAIN_PATH\n        self.test_path = TEST_PATH\n        self.train_annotation_path = TRAIN_ANNOTATION_PATH\n        self.train_tracking_path = TRAIN_TRACKING_PATH\n        self.test_tracking_path = TEST_TRACKING_PATH\n        self.model_path = MODEL_PATH\n        self.model_name = MODEL_NAME\n        self.mode = 'validate'\n        self.n_splits = N_SPLITS\n        from sklearn.model_selection import StratifiedGroupKFold\n        self.cv = StratifiedGroupKFold(self.n_splits)\n        self.use_gpu = USE_GPU\n        self.n_trials = N_TRIALS\n        self.optimize_hyperparams = OPTIMIZE_HYPERPARAMS\n        self.hyperopt_trials = HYPEROPT_TRIALS\n        self.hyperopt_timeout = HYPEROPT_TIMEOUT\n        self.use_class_weights = USE_CLASS_WEIGHTS\n        self.use_smote = USE_SMOTE\n        self.smote_threshold = SMOTE_THRESHOLD\n        self.use_early_stopping = USE_EARLY_STOPPING\n        self.early_stopping_rounds = EARLY_STOPPING_ROUNDS\n        self.adaptive_trials = ADAPTIVE_TRIALS\n        self.min_trials = MIN_TRIALS\n        self.use_section_specific_ranges = USE_SECTION_SPECIFIC_RANGES\n        self.previous_run_id = PREVIOUS_RUN_ID\n        self.previous_best_scores = {}\n        \n        # Base XGBoost parameters\n        xgb_params = {\n            'verbosity': 0,\n            'random_state': 42,\n            'n_estimators': N_ESTIMATORS,\n            'learning_rate': LEARNING_RATE,\n            'max_depth': MAX_DEPTH,\n            'min_child_weight': MIN_CHILD_WEIGHT,\n            'subsample': SUBSAMPLE,\n            'colsample_bytree': COLSAMPLE_BYTREE\n        }\n        \n        if self.use_gpu:\n            xgb_params['device'] = 'cuda'\n        \n        self.base_xgb_params = xgb_params\n        from xgboost import XGBClassifier\n        self.model = XGBClassifier(**xgb_params)\n        self.optimized_params = {}\n\ncfg = CFG()\n\n# Initialize wandb if enabled\nif USE_WANDB:\n    try:\n        import wandb\n        import os\n        # Try to get API key from Kaggle secrets\n        if IS_KAGGLE:\n            try:\n                from kaggle_secrets import UserSecretsClient\n                user_secrets = UserSecretsClient()\n                wandb_api_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n                os.environ[\"WANDB_API_KEY\"] = wandb_api_key\n            except:\n                print(\"Warning: Could not load WANDB_API_KEY from Kaggle secrets\")\n        \n        wandb.init(\n            project=WANDB_PROJECT,\n            entity=WANDB_ENTITY,\n            name=WANDB_RUN_NAME,\n            config={\n                'mode': cfg.mode,\n                'n_splits': cfg.n_splits,\n                'use_gpu': cfg.use_gpu,\n                'optimize_hyperparams': cfg.optimize_hyperparams,\n            }\n        )\n        print(f\"W&B logging enabled: {wandb.run.url}\")\n    except ImportError:\n        print(\"Warning: wandb not installed. Install with: pip install wandb\")\n","metadata":{},"outputs":[],"execution_count":null},{"id":"c9248deb","cell_type":"markdown","source":"## F-Beta Scoring Functions","metadata":{}},{"id":"d436ec14","cell_type":"code","source":"class HostVisibleError(Exception):\n    pass\n\n\n\ndef single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1) -> float:\n    label_frames: defaultdict[str, set[int]] = defaultdict(set)\n    prediction_frames: defaultdict[str, set[int]] = defaultdict(set)\n\n    for row in lab_solution.to_dicts():\n        label_frames[row['label_key']].update(range(row['start_frame'], row['stop_frame']))\n\n    for video in lab_solution['video_id'].unique():\n        active_labels: str = lab_solution.filter(pl.col('video_id') == video)['behaviors_labeled'].first()\n        active_labels: set[str] = set(json.loads(active_labels))\n        predicted_mouse_pairs: defaultdict[str, set[int]] = defaultdict(set)\n\n        for row in lab_submission.filter(pl.col('video_id') == video).to_dicts():\n            if ','.join([str(row['agent_id']), str(row['target_id']), row['action']]) not in active_labels:\n                continue\n\n            new_frames = set(range(row['start_frame'], row['stop_frame']))\n            new_frames = new_frames.difference(prediction_frames[row['prediction_key']])\n            prediction_pair = ','.join([str(row['agent_id']), str(row['target_id'])])\n            if predicted_mouse_pairs[prediction_pair].intersection(new_frames):\n                raise HostVisibleError('Multiple predictions for the same frame from one agent/target pair')\n            prediction_frames[row['prediction_key']].update(new_frames)\n            predicted_mouse_pairs[prediction_pair].update(new_frames)\n\n    tps = defaultdict(int)\n    fns = defaultdict(int)\n    fps = defaultdict(int)\n    for key, pred_frames in prediction_frames.items():\n        action = key.split('_')[-1]\n        matched_label_frames = label_frames[key]\n        tps[action] += len(pred_frames.intersection(matched_label_frames))\n        fns[action] += len(matched_label_frames.difference(pred_frames))\n        fps[action] += len(pred_frames.difference(matched_label_frames))\n\n    distinct_actions = set()\n    for key, frames in label_frames.items():\n        action = key.split('_')[-1]\n        distinct_actions.add(action)\n        if key not in prediction_frames:\n            fns[action] += len(frames)\n\n    action_f1s = []\n    for action in distinct_actions:\n        if tps[action] + fns[action] + fps[action] == 0:\n            action_f1s.append(0)\n        else:\n            action_f1s.append((1 + beta**2) * tps[action] / ((1 + beta**2) * tps[action] + beta**2 * fns[action] + fps[action]))\n    return sum(action_f1s) / len(action_f1s)\n\n\n\ndef mouse_fbeta(solution: pd.DataFrame, submission: pd.DataFrame, beta: float = 1) -> float:\n    if len(solution) == 0 or len(submission) == 0:\n        raise ValueError('Missing solution or submission data')\n\n    expected_cols = ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\n\n    for col in expected_cols:\n        if col not in solution.columns:\n            raise ValueError(f'Solution is missing column {col}')\n        if col not in submission.columns:\n            raise ValueError(f'Submission is missing column {col}')\n\n    solution: pl.DataFrame = pl.DataFrame(solution)\n    submission: pl.DataFrame = pl.DataFrame(submission)\n    assert (solution['start_frame'] <= solution['stop_frame']).all()\n    assert (submission['start_frame'] <= submission['stop_frame']).all()\n    solution_videos = set(solution['video_id'].unique())\n    submission = submission.filter(pl.col('video_id').is_in(solution_videos))\n\n    solution = solution.with_columns(\n        pl.concat_str(\n            [\n                pl.col('video_id').cast(pl.Utf8),\n                pl.col('agent_id').cast(pl.Utf8),\n                pl.col('target_id').cast(pl.Utf8),\n                pl.col('action'),\n            ],\n            separator='_',\n        ).alias('label_key'),\n    )\n    submission = submission.with_columns(\n        pl.concat_str(\n            [\n                pl.col('video_id').cast(pl.Utf8),\n                pl.col('agent_id').cast(pl.Utf8),\n                pl.col('target_id').cast(pl.Utf8),\n                pl.col('action'),\n            ],\n            separator='_',\n        ).alias('prediction_key'),\n    )\n\n    lab_scores = []\n    for lab in solution['lab_id'].unique():\n        lab_solution = solution.filter(pl.col('lab_id') == lab).clone()\n        lab_videos = set(lab_solution['video_id'].unique())\n        lab_submission = submission.filter(pl.col('video_id').is_in(lab_videos)).clone()\n        lab_scores.append(single_lab_f1(lab_solution, lab_submission, beta=beta))\n\n    return sum(lab_scores) / len(lab_scores)\n\n\n\ndef score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, beta: float = 1) -> float:\n    \"\"\"F1 score for the MABe Challenge\"\"\"\n    solution = solution.drop(row_id_column_name, axis='columns', errors='ignore')\n    submission = submission.drop(row_id_column_name, axis='columns', errors='ignore')\n    return mouse_fbeta(solution, submission, beta=beta)\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"2a0dd48c","cell_type":"markdown","source":"## Hyperparameter Optimization","metadata":{}},{"id":"15659dd8","cell_type":"code","source":"def load_best_scores_from_wandb(run_id, entity=\"bookmyrumah-fisa\", project=\"mabe-mouse-behavior\"):\n    \"\"\"\n    Load best CV F1 scores from a wandb run, organized by section.\n    \n    Args:\n        run_id: W&B run ID (e.g., \"iv8w2k9w\" or full path)\n        entity: W&B entity/team name\n        project: W&B project name\n    \n    Returns:\n        dict: Dictionary mapping section numbers to best CV F1 scores\n    \"\"\"\n    if not WANDB_AVAILABLE:\n        return {}\n    \n    try:\n        api = wandb.Api()\n        \n        # Construct run path\n        if \"/\" in run_id:\n            run_path = run_id  # Full path provided\n        else:\n            run_path = f\"{entity}/{project}/{run_id}\"\n        \n        run = api.run(run_path)\n        \n        # Get history to extract best_cv_f1 per section\n        history = run.history()\n        \n        best_scores = {}\n        if len(history) > 0 and 'best_cv_f1' in history.columns and 'section' in history.columns:\n            for _, row in history.iterrows():\n                section = row.get('section')\n                score = row.get('best_cv_f1')\n                \n                # Filter out NaN values\n                if section is not None and not pd.isna(section) and score is not None and not pd.isna(score):\n                    section_key = int(section)\n                    score_val = float(score)\n                    \n                    # Keep the best score for each section\n                    if section_key not in best_scores or score_val > best_scores[section_key]:\n                        best_scores[section_key] = score_val\n        \n        print(f\"Loaded {len(best_scores)} best scores from previous run {run_id}\")\n        return best_scores\n    except Exception as e:\n        print(f\"Warning: Could not load best scores from wandb run {run_id}: {e}\")\n        return {}\n\n\n\ndef optimize_xgb_hyperparameters(X, y, groups, cv, use_gpu, n_trials=50, timeout=None, section=None, action=None,\n                                  use_class_weights=True, use_early_stopping=False, early_stopping_rounds=10,\n                                  use_section_specific_ranges=False, use_smote=False, smote_threshold=0.01,\n                                  target_best_score=None):\n    \"\"\"\n    Optimize XGBoost hyperparameters using Optuna.\n    \n    Args:\n        X: Feature matrix\n        y: Target labels\n        groups: Group labels for StratifiedGroupKFold\n        cv: Cross-validation object\n        use_gpu: Whether to use GPU\n        n_trials: Number of Optuna trials\n        timeout: Timeout in seconds (None for no timeout)\n        section: Section number for wandb logging (optional)\n        action: Action name for wandb logging (optional)\n        use_class_weights: Whether to use class weights for imbalanced data\n        use_early_stopping: Whether to use early stopping with MedianPruner\n        early_stopping_rounds: Number of rounds without improvement before stopping\n        target_best_score: Target CV F1 score from previous run for comparison (optional)\n    \n    Returns:\n        dict: Best hyperparameters\n    \"\"\"\n    # Adaptive trial reduction for slow sections\n    if section == 4:  # Slow section\n        n_trials = max(20, n_trials // 2)\n        if n_trials < 50:\n            print(f\"  Reduced trials to {n_trials} for faster training (section {section})\")\n    \n    def objective(trial):\n        # Section-specific hyperparameter ranges (only if enabled)\n        if use_section_specific_ranges and section == 3:  # Worst performing section - more aggressive regularization\n            params = {\n                'verbosity': 0,\n                'random_state': 42,\n                'n_estimators': trial.suggest_int('n_estimators', 200, 600, step=50),\n                'learning_rate': trial.suggest_float('learning_rate', 0.005, 0.1, log=True),\n                'max_depth': trial.suggest_int('max_depth', 2, 6),\n                'min_child_weight': trial.suggest_int('min_child_weight', 5, 15),\n                'subsample': trial.suggest_float('subsample', 0.5, 0.9, step=0.1),\n                'colsample_bytree': trial.suggest_float('colsample_bytree', 0.5, 0.9, step=0.1),\n                'gamma': trial.suggest_float('gamma', 1, 10, step=0.5),\n                'reg_alpha': trial.suggest_float('reg_alpha', 5, 20, step=0.5),\n                'reg_lambda': trial.suggest_float('reg_lambda', 5, 20, step=0.5),\n            }\n        elif use_section_specific_ranges and section == 6:  # Low CV F1 section - different strategy\n            params = {\n                'verbosity': 0,\n                'random_state': 42,\n                'n_estimators': trial.suggest_int('n_estimators', 300, 700, step=50),\n                'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.2, log=True),\n                'max_depth': trial.suggest_int('max_depth', 4, 12),\n                'min_child_weight': trial.suggest_int('min_child_weight', 1, 5),\n                'subsample': trial.suggest_float('subsample', 0.7, 1.0, step=0.1),\n                'colsample_bytree': trial.suggest_float('colsample_bytree', 0.7, 1.0, step=0.1),\n                'gamma': trial.suggest_float('gamma', 0, 3, step=0.5),\n                'reg_alpha': trial.suggest_float('reg_alpha', 0, 5, step=0.5),\n                'reg_lambda': trial.suggest_float('reg_lambda', 0, 5, step=0.5),\n            }\n        else:\n            # Default ranges\n            params = {\n                'verbosity': 0,\n                'random_state': 42,\n                'n_estimators': trial.suggest_int('n_estimators', 100, 500, step=50),\n                'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3, log=True),\n                'max_depth': trial.suggest_int('max_depth', 3, 10),\n                'min_child_weight': trial.suggest_int('min_child_weight', 1, 10),\n                'subsample': trial.suggest_float('subsample', 0.6, 1.0, step=0.1),\n                'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 1.0, step=0.1),\n                'gamma': trial.suggest_float('gamma', 0, 5, step=0.5),\n                'reg_alpha': trial.suggest_float('reg_alpha', 0, 10, step=0.5),\n                'reg_lambda': trial.suggest_float('reg_lambda', 0, 10, step=0.5),\n            }\n        \n        if use_gpu:\n            params['device'] = 'cuda'\n        \n        # Calculate class weights for imbalanced data\n        if use_class_weights:\n            n_positive = np.sum(y)\n            n_negative = len(y) - n_positive\n            if n_positive > 0 and n_negative > 0:\n                # Scale weight inversely proportional to class frequency\n                scale_pos_weight = n_negative / n_positive\n                # Cap at reasonable maximum to avoid extreme weights\n                scale_pos_weight = min(scale_pos_weight, 100.0)\n                params['scale_pos_weight'] = scale_pos_weight\n        \n        # Perform cross-validation\n        model = XGBClassifier(**params)\n        cv_scores = []\n        \n        for train_idx, val_idx in cv.split(X, y, groups=groups):\n            X_train_fold, X_val_fold = X.iloc[train_idx], X.iloc[val_idx]\n            y_train_fold, y_val_fold = y[train_idx], y[val_idx]\n            \n            # Recalculate class weights for each fold\n            if use_class_weights:\n                n_pos_fold = np.sum(y_train_fold)\n                n_neg_fold = len(y_train_fold) - n_pos_fold\n                if n_pos_fold > 0 and n_neg_fold > 0:\n                    scale_pos_weight_fold = min(n_neg_fold / n_pos_fold, 100.0)\n                    model.set_params(scale_pos_weight=scale_pos_weight_fold)\n            \n            model.fit(X_train_fold, y_train_fold)\n            y_pred = model.predict_proba(X_val_fold)[:, 1]\n            \n            # Use a default threshold of 0.5 for optimization\n            y_pred_binary = (y_pred >= 0.5).astype(int)\n            score = f1_score(y_val_fold, y_pred_binary, zero_division=0)\n            cv_scores.append(score)\n        \n        return np.mean(cv_scores)\n    \n    # Use MedianPruner for early stopping\n    if use_early_stopping:\n        pruner = MedianPruner(n_startup_trials=5, n_warmup_steps=early_stopping_rounds)\n        study = optuna.create_study(direction='maximize', pruner=pruner)\n    else:\n        study = optuna.create_study(direction='maximize')\n    \n    study.optimize(objective, n_trials=n_trials, timeout=timeout, n_jobs=1)\n    \n    best_params = study.best_params.copy()\n    best_params['verbosity'] = 0\n    best_params['random_state'] = 42\n    if use_gpu:\n        best_params['device'] = 'cuda'\n    \n    print(f\"  Best CV F1: {study.best_value:.4f}\")\n    \n    # Compare with target score if provided\n    if target_best_score is not None:\n        improvement = study.best_value - target_best_score\n        improvement_pct = (improvement / target_best_score * 100) if target_best_score > 0 else 0\n        \n        print(f\"  Previous best: {target_best_score:.4f}\")\n        print(f\"  Current best:  {study.best_value:.4f}\")\n        print(f\"  Improvement:   {improvement:+.4f} ({improvement_pct:+.2f}%)\")\n    \n    print(f\"  Best params: {best_params}\")\n    \n    # Log to wandb if enabled\n    if WANDB_AVAILABLE and wandb.run is not None:\n        log_dict = {\n            'best_cv_f1': study.best_value,\n            **{f'best_param_{k}': v for k, v in best_params.items()}\n        }\n        if section is not None:\n            log_dict['section'] = section\n            log_dict[f'best_cv_f1/section_{section}'] = study.best_value\n        if action is not None:\n            log_dict['action'] = action\n        \n        # Add comparison metrics if target score is provided\n        if target_best_score is not None:\n            improvement = study.best_value - target_best_score\n            improvement_pct = (improvement / target_best_score * 100) if target_best_score > 0 else 0\n            log_dict['previous_best_cv_f1'] = target_best_score\n            log_dict['cv_f1_improvement'] = improvement\n            log_dict['cv_f1_improvement_pct'] = improvement_pct\n        \n        wandb.log(log_dict)\n    \n    return best_params\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"f6bf89dc","cell_type":"markdown","source":"## Data Loading and Preprocessing","metadata":{}},{"id":"8ff28b00","cell_type":"code","source":"drop_body_parts = [\n    'headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright',\n    'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft', 'headpiece_topfrontright',\n    'spine_1', 'spine_2', 'tail_middle_1', 'tail_middle_2', 'tail_midpoint'\n]\n\ndef create_solution_df(dataset, cfg):\n    solution = []\n    for _, row in tqdm(dataset.iterrows(), total=len(dataset), desc=\"Creating solution\"):\n        lab_id = row['lab_id']\n        if lab_id.startswith('MABe22'):\n            continue\n        \n        video_id = row['video_id']\n        path = f\"{cfg.train_annotation_path}/{lab_id}/{video_id}.parquet\"\n        try:\n            annot = pd.read_parquet(path)\n        except FileNotFoundError:\n            continue\n    \n        annot['lab_id'] = lab_id\n        annot['video_id'] = video_id\n        annot['behaviors_labeled'] = row['behaviors_labeled']\n        annot['target_id'] = np.where(annot.target_id != annot.agent_id, annot['target_id'].apply(lambda s: f\"mouse{s}\"), 'self')\n        annot['agent_id'] = annot['agent_id'].apply(lambda s: f\"mouse{s}\")\n        solution.append(annot)\n    \n    solution = pd.concat(solution)\n    return solution\n\n\n\ndef generate_mouse_data(dataset, traintest, cfg, traintest_directory=None, generate_single=True, generate_pair=True):\n    if traintest_directory is None:\n        if traintest == 'train':\n            traintest_directory = cfg.train_tracking_path\n        else:\n            traintest_directory = cfg.test_tracking_path\n        \n    for _, row in dataset.iterrows():\n        lab_id = row.lab_id\n        if lab_id.startswith('MABe22') or type(row.behaviors_labeled) != str:\n            continue\n        \n        video_id = row.video_id\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n        if len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        \n        del vid\n        gc.collect()\n        \n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n        pvid /= row.pix_per_cm_approx\n\n        vid_behaviors = json.loads(row.behaviors_labeled)\n        vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in vid_behaviors}))\n        vid_behaviors = [b.split(',') for b in vid_behaviors]\n        vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action'])\n        \n        if traintest == 'train':\n            try:\n                annot = pd.read_parquet(path.replace('train_tracking', 'train_annotation'))\n            except FileNotFoundError:\n                continue\n\n        if generate_single:\n            vid_behaviors_subset = vid_behaviors.query(\"target == 'self'\")\n            for mouse_id_str in np.unique(vid_behaviors_subset.agent):\n                try:\n                    mouse_id = int(mouse_id_str[-1])\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"agent == @mouse_id_str\").action)\n                    single_mouse = pvid.loc[:, mouse_id]\n                    assert len(single_mouse) == len(pvid)\n                    single_mouse_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': mouse_id_str,\n                        'target_id': 'self',\n                        'video_frame': single_mouse.index\n                    })\n                    if traintest == 'train':\n                        single_mouse_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=single_mouse.index)\n                        annot_subset = annot.query(\"(agent_id == @mouse_id) & (target_id == @mouse_id)\")\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\n                            single_mouse_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'single', single_mouse, single_mouse_meta, single_mouse_label\n                    else:\n                        yield 'single', single_mouse, single_mouse_meta, vid_agent_actions\n                except KeyError:\n                    pass\n\n        if generate_pair:\n            vid_behaviors_subset = vid_behaviors.query(\"target != 'self'\")\n            if len(vid_behaviors_subset) > 0:\n                for agent, target in itertools.permutations(np.unique(pvid.columns.get_level_values('mouse_id')), 2):\n                    agent_str = f\"mouse{agent}\"\n                    target_str = f\"mouse{target}\"\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"(agent == @agent_str) & (target == @target_str)\").action)\n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\n                    assert len(mouse_pair) == len(pvid)\n                    mouse_pair_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': agent_str,\n                        'target_id': target_str,\n                        'video_frame': mouse_pair.index\n                    })\n                    if traintest == 'train':\n                        mouse_pair_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=mouse_pair.index)\n                        annot_subset = annot.query(\"(agent_id == @agent) & (target_id == @target)\")\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\n                            mouse_pair_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'pair', mouse_pair, mouse_pair_meta, mouse_pair_label\n                    else:\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"47ebfffc","cell_type":"markdown","source":"## Feature Engineering","metadata":{}},{"id":"908ed501","cell_type":"code","source":"def safe_rolling(series, window, func, min_periods=None):\n    if min_periods is None:\n        min_periods = max(1, window // 4)\n    return series.rolling(window, min_periods=min_periods, center=True).apply(func, raw=True)\n\n\n\ndef _scale(n_frames_at_30fps, fps, ref=30.0):\n    return max(1, int(round(n_frames_at_30fps * float(fps) / ref)))\n\n\n\ndef _scale_signed(n_frames_at_30fps, fps, ref=30.0):\n    if n_frames_at_30fps == 0:\n        return 0\n    s = 1 if n_frames_at_30fps > 0 else -1\n    mag = max(1, int(round(abs(n_frames_at_30fps) * float(fps) / ref)))\n    return s * mag\n\n\n\ndef _fps_from_meta(meta_df, fallback_lookup, default_fps=30.0):\n    if 'frames_per_second' in meta_df.columns and pd.notnull(meta_df['frames_per_second']).any():\n        return float(meta_df['frames_per_second'].iloc[0])\n    vid = meta_df['video_id'].iloc[0]\n    return float(fallback_lookup.get(vid, default_fps))\n\n\n\ndef add_curvature_features(X, center_x, center_y, fps):\n    vel_x = center_x.diff()\n    vel_y = center_y.diff()\n    acc_x = vel_x.diff()\n    acc_y = vel_y.diff()\n\n    cross_prod = vel_x * acc_y - vel_y * acc_x\n    vel_mag = np.sqrt(vel_x**2 + vel_y**2)\n    curvature = np.abs(cross_prod) / (vel_mag**3 + 1e-6)\n\n    for w in [25, 50, 75]:\n        ws = _scale(w, fps)\n        X[f'curv_mean_{w}'] = curvature.rolling(ws, min_periods=max(1, ws // 5)).mean()\n\n    angle = np.arctan2(vel_y, vel_x)\n    angle_change = np.abs(angle.diff())\n    w = 30\n    ws = _scale(w, fps)\n    X[f'turn_rate_{w}'] = angle_change.rolling(ws, min_periods=max(1, ws // 5)).sum()\n\n    return X\n\n\n\ndef add_multiscale_features(X, center_x, center_y, fps):\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)\n\n    scales = [20, 40, 60, 80]\n    for scale in scales:\n        ws = _scale(scale, fps)\n        if len(speed) >= ws:\n            X[f'sp_m{scale}'] = speed.rolling(ws, min_periods=max(1, ws // 4)).mean()\n            X[f'sp_s{scale}'] = speed.rolling(ws, min_periods=max(1, ws // 4)).std()\n\n    if len(scales) >= 2 and f'sp_m{scales[0]}' in X.columns and f'sp_m{scales[-1]}' in X.columns:\n        X['sp_ratio'] = X[f'sp_m{scales[0]}'] / (X[f'sp_m{scales[-1]}'] + 1e-6)\n\n    return X\n\n\n\ndef add_state_features(X, center_x, center_y, fps):\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)\n    w_ma = _scale(15, fps)\n    speed_ma = speed.rolling(w_ma, min_periods=max(1, w_ma // 3)).mean()\n\n    try:\n        bins = [-np.inf, 0.5 * fps, 2.0 * fps, 5.0 * fps, np.inf]\n        speed_states = pd.cut(speed_ma, bins=bins, labels=[0, 1, 2, 3]).astype(float)\n\n        for window in [20, 40, 60, 80]:\n            ws = _scale(window, fps)\n            if len(speed_states) >= ws:\n                for state in [0, 1, 2, 3]:\n                    X[f's{state}_{window}'] = (\n                        (speed_states == state).astype(float)\n                        .rolling(ws, min_periods=max(1, ws // 5)).mean()\n                    )\n                state_changes = (speed_states != speed_states.shift(1)).astype(float)\n                X[f'trans_{window}'] = state_changes.rolling(ws, min_periods=max(1, ws // 5)).sum()\n    except Exception:\n        pass\n\n    return X\n\n\n\ndef add_longrange_features(X, center_x, center_y, fps):\n    for window in [30, 60, 120]:\n        ws = _scale(window, fps)\n        if len(center_x) >= ws:\n            X[f'x_ml{window}'] = center_x.rolling(ws, min_periods=max(5, ws // 6)).mean()\n            X[f'y_ml{window}'] = center_y.rolling(ws, min_periods=max(5, ws // 6)).mean()\n\n    for span in [30, 60, 120]:\n        s = _scale(span, fps)\n        X[f'x_e{span}'] = center_x.ewm(span=s, min_periods=1).mean()\n        X[f'y_e{span}'] = center_y.ewm(span=s, min_periods=1).mean()\n\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)\n    for window in [30, 60, 120]:\n        ws = _scale(window, fps)\n        if len(speed) >= ws:\n            X[f'sp_pct{window}'] = speed.rolling(ws, min_periods=max(5, ws // 6)).rank(pct=True)\n\n    return X\n\n\n\ndef add_interaction_features(X, mouse_pair, avail_A, avail_B, fps):\n    if 'body_center' not in avail_A or 'body_center' not in avail_B:\n        return X\n\n    rel_x = mouse_pair['A']['body_center']['x'] - mouse_pair['B']['body_center']['x']\n    rel_y = mouse_pair['A']['body_center']['y'] - mouse_pair['B']['body_center']['y']\n    rel_dist = np.sqrt(rel_x**2 + rel_y**2)\n\n    A_vx = mouse_pair['A']['body_center']['x'].diff()\n    A_vy = mouse_pair['A']['body_center']['y'].diff()\n    B_vx = mouse_pair['B']['body_center']['x'].diff()\n    B_vy = mouse_pair['B']['body_center']['y'].diff()\n\n    A_lead = (A_vx * rel_x + A_vy * rel_y) / (np.sqrt(A_vx**2 + A_vy**2) * rel_dist + 1e-6)\n    B_lead = (B_vx * (-rel_x) + B_vy * (-rel_y)) / (np.sqrt(B_vx**2 + B_vy**2) * rel_dist + 1e-6)\n\n    for window in [30, 60]:\n        ws = _scale(window, fps)\n        X[f'A_ld{window}'] = A_lead.rolling(ws, min_periods=max(1, ws // 6)).mean()\n        X[f'B_ld{window}'] = B_lead.rolling(ws, min_periods=max(1, ws // 6)).mean()\n\n    approach = -rel_dist.diff()\n    chase = approach * B_lead\n    w = 30\n    ws = _scale(w, fps)\n    X[f'chase_{w}'] = chase.rolling(ws, min_periods=max(1, ws // 6)).mean()\n\n    for window in [60, 120]:\n        ws = _scale(window, fps)\n        A_sp = np.sqrt(A_vx**2 + A_vy**2)\n        B_sp = np.sqrt(B_vx**2 + B_vy**2)\n        X[f'sp_cor{window}'] = A_sp.rolling(ws, min_periods=max(1, ws // 6)).corr(B_sp)\n\n    return X\n\n\n\ndef transform_single(single_mouse, body_parts_tracked, fps):\n    available_body_parts = single_mouse.columns.get_level_values(0)\n\n    X = pd.DataFrame({\n        f\"{p1}+{p2}\": np.square(single_mouse[p1] - single_mouse[p2]).sum(axis=1, skipna=False)\n        for p1, p2 in itertools.combinations(body_parts_tracked, 2)\n        if p1 in available_body_parts and p2 in available_body_parts\n    })\n    X = X.reindex(columns=[f\"{p1}+{p2}\" for p1, p2 in itertools.combinations(body_parts_tracked, 2)], copy=False)\n\n    if all(p in single_mouse.columns for p in ['ear_left', 'ear_right', 'tail_base']):\n        lag = _scale(10, fps)\n        shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(lag)\n        speeds = pd.DataFrame({\n            'sp_lf': np.square(single_mouse['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False),\n            'sp_rt': np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False),\n            'sp_lf2': np.square(single_mouse['ear_left'] - shifted['tail_base']).sum(axis=1, skipna=False),\n            'sp_rt2': np.square(single_mouse['ear_right'] - shifted['tail_base']).sum(axis=1, skipna=False),\n        })\n        X = pd.concat([X, speeds], axis=1)\n\n    if 'nose+tail_base' in X.columns and 'ear_left+ear_right' in X.columns:\n        X['elong'] = X['nose+tail_base'] / (X['ear_left+ear_right'] + 1e-6)\n\n    if all(p in available_body_parts for p in ['nose', 'body_center', 'tail_base']):\n        v1 = single_mouse['nose'] - single_mouse['body_center']\n        v2 = single_mouse['tail_base'] - single_mouse['body_center']\n        X['body_ang'] = (v1['x'] * v2['x'] + v1['y'] * v2['y']) / (\n            np.sqrt(v1['x']**2 + v1['y']**2) * np.sqrt(v2['x']**2 + v2['y']**2) + 1e-6)\n\n    if 'body_center' in available_body_parts:\n        cx = single_mouse['body_center']['x']\n        cy = single_mouse['body_center']['y']\n\n        for w in [5, 15, 30, 60]:\n            ws = _scale(w, fps)\n            roll = dict(min_periods=1, center=True)\n            X[f'cx_m{w}'] = cx.rolling(ws, **roll).mean()\n            X[f'cy_m{w}'] = cy.rolling(ws, **roll).mean()\n            X[f'cx_s{w}'] = cx.rolling(ws, **roll).std()\n            X[f'cy_s{w}'] = cy.rolling(ws, **roll).std()\n            X[f'x_rng{w}'] = cx.rolling(ws, **roll).max() - cx.rolling(ws, **roll).min()\n            X[f'y_rng{w}'] = cy.rolling(ws, **roll).max() - cy.rolling(ws, **roll).min()\n            X[f'disp{w}'] = np.sqrt(cx.diff().rolling(ws, min_periods=1).sum()**2 +\n                                     cy.diff().rolling(ws, min_periods=1).sum()**2)\n            X[f'act{w}'] = np.sqrt(cx.diff().rolling(ws, min_periods=1).var() +\n                                   cy.diff().rolling(ws, min_periods=1).var())\n\n        X = add_curvature_features(X, cx, cy, fps)\n        X = add_multiscale_features(X, cx, cy, fps)\n        X = add_state_features(X, cx, cy, fps)\n        X = add_longrange_features(X, cx, cy, fps)\n\n    if all(p in available_body_parts for p in ['nose', 'tail_base']):\n        nt_dist = np.sqrt((single_mouse['nose']['x'] - single_mouse['tail_base']['x'])**2 +\n                          (single_mouse['nose']['y'] - single_mouse['tail_base']['y'])**2)\n        for lag in [10, 20, 40]:\n            l = _scale(lag, fps)\n            X[f'nt_lg{lag}'] = nt_dist.shift(l)\n            X[f'nt_df{lag}'] = nt_dist - nt_dist.shift(l)\n\n    if all(p in available_body_parts for p in ['ear_left', 'ear_right']):\n        ear_d = np.sqrt((single_mouse['ear_left']['x'] - single_mouse['ear_right']['x'])**2 +\n                        (single_mouse['ear_left']['y'] - single_mouse['ear_right']['y'])**2)\n        for off in [-30, -20, -10, 10, 20, 30]:\n            o = _scale_signed(off, fps)\n            X[f'ear_o{off}'] = ear_d.shift(-o)\n        w = _scale(30, fps)\n        X['ear_con'] = ear_d.rolling(w, min_periods=1, center=True).std() / \\\n                       (ear_d.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n\n    return X.astype(np.float32, copy=False)\n\n\n\ndef transform_pair(mouse_pair, body_parts_tracked, fps):\n    avail_A = mouse_pair['A'].columns.get_level_values(0)\n    avail_B = mouse_pair['B'].columns.get_level_values(0)\n\n    X = pd.DataFrame({\n        f\"12+{p1}+{p2}\": np.square(mouse_pair['A'][p1] - mouse_pair['B'][p2]).sum(axis=1, skipna=False)\n        for p1, p2 in itertools.product(body_parts_tracked, repeat=2)\n        if p1 in avail_A and p2 in avail_B\n    })\n    X = X.reindex(columns=[f\"12+{p1}+{p2}\" for p1, p2 in itertools.product(body_parts_tracked, repeat=2)], copy=False)\n\n    if ('A', 'ear_left') in mouse_pair.columns and ('B', 'ear_left') in mouse_pair.columns:\n        lag = _scale(10, fps)\n        shA = mouse_pair['A']['ear_left'].shift(lag)\n        shB = mouse_pair['B']['ear_left'].shift(lag)\n        speeds = pd.DataFrame({\n            'sp_A': np.square(mouse_pair['A']['ear_left'] - shA).sum(axis=1, skipna=False),\n            'sp_AB': np.square(mouse_pair['A']['ear_left'] - shB).sum(axis=1, skipna=False),\n            'sp_B': np.square(mouse_pair['B']['ear_left'] - shB).sum(axis=1, skipna=False),\n        })\n        X = pd.concat([X, speeds], axis=1)\n\n    if 'nose+tail_base' in X.columns and 'ear_left+ear_right' in X.columns:\n        X['elong'] = X['nose+tail_base'] / (X['ear_left+ear_right'] + 1e-6)\n\n    if all(p in avail_A for p in ['nose', 'tail_base']) and all(p in avail_B for p in ['nose', 'tail_base']):\n        dir_A = mouse_pair['A']['nose'] - mouse_pair['A']['tail_base']\n        dir_B = mouse_pair['B']['nose'] - mouse_pair['B']['tail_base']\n        X['rel_ori'] = (dir_A['x'] * dir_B['x'] + dir_A['y'] * dir_B['y']) / (\n            np.sqrt(dir_A['x']**2 + dir_A['y']**2) * np.sqrt(dir_B['x']**2 + dir_B['y']**2) + 1e-6)\n\n    if all(p in avail_A for p in ['nose']) and all(p in avail_B for p in ['nose']):\n        cur = np.square(mouse_pair['A']['nose'] - mouse_pair['B']['nose']).sum(axis=1, skipna=False)\n        lag = _scale(10, fps)\n        shA_n = mouse_pair['A']['nose'].shift(lag)\n        shB_n = mouse_pair['B']['nose'].shift(lag)\n        past = np.square(shA_n - shB_n).sum(axis=1, skipna=False)\n        X['appr'] = cur - past\n\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        cd = np.sqrt((mouse_pair['A']['body_center']['x'] - mouse_pair['B']['body_center']['x'])**2 +\n                     (mouse_pair['A']['body_center']['y'] - mouse_pair['B']['body_center']['y'])**2)\n        X['v_cls'] = (cd < 5.0).astype(float)\n        X['cls'] = ((cd >= 5.0) & (cd < 15.0)).astype(float)\n        X['med'] = ((cd >= 15.0) & (cd < 30.0)).astype(float)\n        X['far'] = (cd >= 30.0).astype(float)\n\n    cd_full = None\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        cd_full = np.square(mouse_pair['A']['body_center'] - mouse_pair['B']['body_center']).sum(axis=1, skipna=False)\n\n        for w in [5, 15, 30, 60]:\n            ws = _scale(w, fps)\n            roll = dict(min_periods=1, center=True)\n            X[f'd_m{w}'] = cd_full.rolling(ws, **roll).mean()\n            X[f'd_s{w}'] = cd_full.rolling(ws, **roll).std()\n            X[f'd_mn{w}'] = cd_full.rolling(ws, **roll).min()\n            X[f'd_mx{w}'] = cd_full.rolling(ws, **roll).max()\n\n            d_var = cd_full.rolling(ws, **roll).var()\n            X[f'int{w}'] = 1 / (1 + d_var)\n\n            Axd = mouse_pair['A']['body_center']['x'].diff()\n            Ayd = mouse_pair['A']['body_center']['y'].diff()\n            Bxd = mouse_pair['B']['body_center']['x'].diff()\n            Byd = mouse_pair['B']['body_center']['y'].diff()\n            coord = Axd * Bxd + Ayd * Byd\n            X[f'co_m{w}'] = coord.rolling(ws, **roll).mean()\n            X[f'co_s{w}'] = coord.rolling(ws, **roll).std()\n\n    if 'nose' in avail_A and 'nose' in avail_B:\n        nn = np.sqrt((mouse_pair['A']['nose']['x'] - mouse_pair['B']['nose']['x'])**2 +\n                     (mouse_pair['A']['nose']['y'] - mouse_pair['B']['nose']['y'])**2)\n        for lag in [10, 20, 40]:\n            l = _scale(lag, fps)\n            X[f'nn_lg{lag}'] = nn.shift(l)\n            X[f'nn_ch{lag}'] = nn - nn.shift(l)\n            is_cl = (nn < 10.0).astype(float)\n            X[f'cl_ps{lag}'] = is_cl.rolling(l, min_periods=1).mean()\n\n    if 'body_center' in avail_A and 'body_center' in avail_B and cd_full is not None:\n        Avx = mouse_pair['A']['body_center']['x'].diff()\n        Avy = mouse_pair['A']['body_center']['y'].diff()\n        Bvx = mouse_pair['B']['body_center']['x'].diff()\n        Bvy = mouse_pair['B']['body_center']['y'].diff()\n        val = (Avx * Bvx + Avy * Bvy) / (np.sqrt(Avx**2 + Avy**2) * np.sqrt(Bvx**2 + Bvy**2) + 1e-6)\n\n        for off in [-30, -20, -10, 0, 10, 20, 30]:\n            o = _scale_signed(off, fps)\n            X[f'va_{off}'] = val.shift(-o)\n\n        w = _scale(30, fps)\n        X['int_con'] = cd_full.rolling(w, min_periods=1, center=True).std() / \\\n                       (cd_full.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n\n        X = add_interaction_features(X, mouse_pair, avail_A, avail_B, fps)\n\n    return X.astype(np.float32, copy=False)\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"4bfc840e","cell_type":"markdown","source":"## Training and Validation Functions","metadata":{}},{"id":"5a8a3e9e","cell_type":"code","source":"def robustify(submission, dataset, traintest, cfg, traintest_directory=None):\n    if traintest_directory is None:\n        if traintest == 'train':\n            traintest_directory = cfg.train_tracking_path\n        else:\n            traintest_directory = cfg.test_tracking_path\n\n    old_submission = submission.copy()\n    submission = submission[submission.start_frame < submission.stop_frame]\n    if len(submission) != len(old_submission):\n        print(\"ERROR: Dropped frames with start >= stop\")\n    \n    old_submission = submission.copy()\n    group_list = []\n    for _, group in submission.groupby(['video_id', 'agent_id', 'target_id']):\n        group = group.sort_values('start_frame')\n        mask = np.ones(len(group), dtype=bool)\n        last_stop_frame = 0\n        for i, (_, row) in enumerate(group.iterrows()):\n            if row['start_frame'] < last_stop_frame:\n                mask[i] = False\n            else:\n                last_stop_frame = row['stop_frame']\n        group_list.append(group[mask])\n        \n    submission = pd.concat(group_list)\n    \n    if len(submission) != len(old_submission):\n        print(\"ERROR: Dropped duplicate frames\")\n        \n    s_list = []\n    for idx, row in dataset.iterrows():\n        lab_id = row['lab_id']\n        if lab_id.startswith('MABe22'):\n            continue\n        \n        video_id = row['video_id']\n        if (submission.video_id == video_id).any():\n            continue\n        \n        if type(row.behaviors_labeled) != str:\n            continue\n\n        print(f\"Video {video_id} has no predictions.\")\n        \n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n    \n        vid_behaviors = json.loads(row['behaviors_labeled'])\n        vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in vid_behaviors}))\n        vid_behaviors = [b.split(',') for b in vid_behaviors]\n        vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action'])\n    \n        start_frame = vid.video_frame.min()\n        stop_frame = vid.video_frame.max() + 1\n    \n        for (agent, target), actions in vid_behaviors.groupby(['agent', 'target']):\n            batch_length = int(np.ceil((stop_frame - start_frame) / len(actions)))\n            for i, (_, action_row) in enumerate(actions.iterrows()):\n                batch_start = start_frame + i * batch_length\n                batch_stop = min(batch_start + batch_length, stop_frame)\n                s_list.append((video_id, agent, target, action_row['action'], batch_start, batch_stop))\n\n    if len(s_list) > 0:\n        submission = pd.concat([\n            submission,\n            pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n        ])\n        print(\"ERROR: Filled empty videos\")\n\n    submission = submission.reset_index(drop=True)\n    return submission\n\n\n\ndef predict_multiclass(pred, meta, thresholds):\n    ama = np.argmax(pred.values, axis=1)\n    max_proba = pred.max(axis=1).values\n\n    threshold_array = np.array([thresholds.get(col, 0.27) for col in pred.columns])\n    action_thresholds = threshold_array[ama]\n\n    ama = np.where(max_proba >= action_thresholds, ama, -1)\n    ama = pd.Series(ama, index=meta.video_frame)\n    \n    changes_mask = (ama != ama.shift(1)).values\n    ama_changes = ama[changes_mask]\n    meta_changes = meta[changes_mask]\n    \n    mask = ama_changes.values >= 0\n    mask[-1] = False\n    \n    submission_part = pd.DataFrame({\n        'video_id': meta_changes['video_id'][mask].values,\n        'agent_id': meta_changes['agent_id'][mask].values,\n        'target_id': meta_changes['target_id'][mask].values,\n        'action': pred.columns[ama_changes[mask].values],\n        'start_frame': ama_changes.index[mask],\n        'stop_frame': ama_changes.index[1:][mask[:-1]]\n    })\n    \n    stop_video_id = meta_changes['video_id'][1:][mask[:-1]].values\n    stop_agent_id = meta_changes['agent_id'][1:][mask[:-1]].values\n    stop_target_id = meta_changes['target_id'][1:][mask[:-1]].values\n    for i in range(len(submission_part)):\n        video_id = submission_part.video_id.iloc[i]\n        agent_id = submission_part.agent_id.iloc[i]\n        target_id = submission_part.target_id.iloc[i]\n        if stop_video_id[i] != video_id or stop_agent_id[i] != agent_id or stop_target_id[i] != target_id:\n            new_stop_frame = meta.query(\"(video_id == @video_id)\").video_frame.max() + 1\n            submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n\n    return submission_part\n\n\n\ndef tune_threshold(oof_action, y_action, n_trials=200, positive_ratio=None):\n    \"\"\"\n    Tune threshold with consideration for class imbalance.\n    \n    Args:\n        oof_action: Out-of-fold predictions\n        y_action: True labels\n        n_trials: Number of Optuna trials\n        positive_ratio: Ratio of positive samples (for imbalanced classes)\n    \n    Returns:\n        float: Optimal threshold\n    \"\"\"\n    def objective(trial):\n        threshold = trial.suggest_float(\"threshold\", 0, 1, step=0.01)\n        y_pred = (oof_action >= threshold).astype(int)\n        score = f1_score(y_action, y_pred, zero_division=0)\n        \n        # For very imbalanced classes, also consider precision\n        if positive_ratio is not None and positive_ratio < 0.01:\n            precision = precision_score(y_action, y_pred, zero_division=0)\n            # Weighted combination favoring F1 but considering precision\n            score = 0.7 * score + 0.3 * precision\n        \n        return score\n\n    study = optuna.create_study(direction=\"maximize\")\n    study.optimize(objective, n_trials=n_trials, n_jobs=1)  # Changed from n_jobs=-1 to n_jobs=1 for Windows compatibility\n    return study.best_params[\"threshold\"]\n\n\n\ndef cross_validate_with_smote(X, y, groups, cv, estimator, use_smote, smote_threshold, \n                              use_class_weights, use_early_stopping, early_stopping_rounds,\n                              metric, save_path=None, verbose=False):\n    \"\"\"\n    Perform cross-validation with SMOTE applied within each training fold.\n    \n    Windows-safe version with controlled threading and multiprocessing safety.\n    This ensures group integrity is preserved and prevents data leakage by only\n    applying SMOTE to training data, not validation data.\n    \n    Parameters\n    ----------\n    X : pd.DataFrame\n        Feature matrix\n    y : np.ndarray or pd.Series\n        Target variable\n    groups : np.ndarray or pd.Series\n        Group labels for group-based CV\n    cv : sklearn.model_selection.BaseCrossValidator\n        Cross-validation strategy\n    estimator : sklearn.base.BaseEstimator\n        Base estimator to train\n    use_smote : bool\n        Whether to apply SMOTE\n    smote_threshold : float\n        Positive ratio threshold for applying SMOTE\n    use_class_weights : bool\n        Whether to use class weights\n    use_early_stopping : bool\n        Whether to use early stopping\n    early_stopping_rounds : int\n        Number of rounds for early stopping\n    metric : callable\n        Metric function for evaluation\n    save_path : str, optional\n        Path to save trained models\n    verbose : bool\n        Whether to print progress\n        \n    Returns\n    -------\n    oof_preds : np.ndarray\n        Out-of-fold predictions on original data\n    estimators : list\n        List of trained estimators (one per fold)\n    fold_scores : list\n        List of scores for each fold\n    \"\"\"\n    import os\n    import joblib\n    \n    # --- Multiprocessing safety (Windows)\n    os.environ[\"LOKY_MAX_CPU_COUNT\"] = \"8\"  # limit process detection\n    os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n    os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n    \n    # Convert to numpy arrays for indexing\n    if isinstance(y, pd.Series):\n        y = y.values\n    if isinstance(groups, pd.Series):\n        groups = groups.values\n    \n    oof_preds = np.zeros(len(X))\n    estimators = []\n    fold_scores = []\n    \n    # Prepare imputer for SMOTE (SMOTE doesn't accept NaN)\n    imputer = None\n    if use_smote and X.isna().any().any():\n        imputer = SimpleImputer(strategy='median')\n    \n    # Calculate positive ratio for SMOTE decision\n    positive_ratio = np.sum(y) / len(y) if len(y) > 0 else 0\n    should_use_smote = (use_smote and SMOTE_AVAILABLE and \n                       positive_ratio < smote_threshold and \n                       len(np.unique(y)) == 2)\n    \n    # Get CV splits\n    split = cv.split(X, y, groups=groups)\n    \n    for fold_idx, (train_idx, val_idx) in enumerate(split):\n        # Get fold data\n        X_train_fold = X.iloc[train_idx].copy()\n        X_val_fold = X.iloc[val_idx].copy()\n        y_train_fold = y[train_idx]\n        y_val_fold = y[val_idx]\n        \n        # Apply SMOTE only to training fold if needed\n        if should_use_smote:\n            try:\n                # Check if we have enough positive samples in training fold\n                n_positive_train = np.sum(y_train_fold)\n                if n_positive_train >= 2:  # Need at least 2 positive samples for SMOTE\n                    # Impute NaN values in training data\n                    if imputer is not None:\n                        X_train_imputed = pd.DataFrame(\n                            imputer.fit_transform(X_train_fold),\n                            columns=X_train_fold.columns,\n                            index=X_train_fold.index\n                        )\n                    else:\n                        X_train_imputed = X_train_fold\n                    \n                    # Calculate k_neighbors based on positive samples in training fold\n                    k_neighbors = min(5, n_positive_train - 1)\n                    if k_neighbors >= 1:\n                        # Apply SMOTE to training fold only (n_jobs=1 for Windows safety)\n                        smote = SMOTE(random_state=42 + fold_idx, k_neighbors=k_neighbors, n_jobs=1)\n                        X_train_resampled, y_train_resampled = smote.fit_resample(\n                            X_train_imputed, y_train_fold\n                        )\n                        \n                        # Convert back to DataFrame if needed\n                        if not isinstance(X_train_resampled, pd.DataFrame):\n                            X_train_resampled = pd.DataFrame(\n                                X_train_resampled,\n                                columns=X_train_fold.columns\n                            )\n                    else:\n                        X_train_resampled, y_train_resampled = X_train_imputed, y_train_fold\n                else:\n                    # Not enough positive samples, skip SMOTE for this fold\n                    X_train_resampled, y_train_resampled = X_train_fold, y_train_fold\n            except Exception as e:\n                if verbose:\n                    print(f\"  Fold {fold_idx}: SMOTE failed ({e}), using original training data\")\n                X_train_resampled, y_train_resampled = X_train_fold, y_train_fold\n        else:\n            X_train_resampled, y_train_resampled = X_train_fold, y_train_fold\n        \n        # Clone estimator for this fold\n        fold_estimator = clone(estimator)\n        \n        # Force single-threaded CPU + GPU use (Windows safety)\n        if hasattr(fold_estimator, \"set_params\"):\n            # Set n_jobs=1 to avoid multiprocessing issues on Windows\n            # GPU will still be used if device='cuda' is set\n            try:\n                fold_estimator.set_params(n_jobs=1, verbosity=0)\n            except (ValueError, TypeError):\n                # Some parameters might not be settable, ignore\n                pass\n        \n        # Set class weights if enabled\n        if use_class_weights:\n            n_positive = np.sum(y_train_resampled)\n            n_negative = len(y_train_resampled) - n_positive\n            if n_positive > 0 and n_negative > 0:\n                scale_pos_weight = min(n_negative / n_positive, 100.0)\n                fold_estimator.set_params(scale_pos_weight=scale_pos_weight)\n        \n        # Prepare fit arguments\n        fit_args = {}\n        if use_early_stopping:\n            fit_args[\"eval_set\"] = [(X_val_fold, y_val_fold)]\n            fit_args[\"early_stopping_rounds\"] = early_stopping_rounds\n            fit_args[\"verbose\"] = False\n        \n        # Train on resampled training data\n        fold_estimator.fit(X_train_resampled, y_train_resampled, **fit_args)\n        estimators.append(fold_estimator)\n        \n        # Predict on original validation data (not resampled)\n        if hasattr(fold_estimator, 'predict_proba'):\n            y_pred_proba = fold_estimator.predict_proba(X_val_fold)[:, 1]\n        else:\n            y_pred_proba = fold_estimator.predict(X_val_fold)\n        \n        oof_preds[val_idx] = y_pred_proba\n        \n        # Calculate fold score\n        fold_score = metric(y_val_fold, (y_pred_proba >= 0.5).astype(int), zero_division=0)\n        fold_scores.append(fold_score)\n        \n        # Save model if save_path is provided\n        if save_path:\n            os.makedirs(save_path, exist_ok=True)\n            model_path = os.path.join(save_path, f\"fold_{fold_idx}_model.pkl\")\n            joblib.dump(fold_estimator, model_path)\n        \n        if verbose:\n            print(f\"  Fold {fold_idx}: {metric.__name__} = {fold_score:.4f}\")\n    \n    return oof_preds, estimators, fold_scores\n\n\n\ndef cross_validate_classifier(X, label, meta, body_parts_tracked_str, section, cfg):\n    section_start_time = time.time()\n    os.makedirs(f\"{cfg.model_name}/{section}\", exist_ok=True)\n    \n    oof = pd.DataFrame(index=meta.video_frame)\n    \n    f1_list = []\n    submission_list = []\n    thresholds = {}\n    \n    # Track action index for step tracking\n    action_idx = 0\n    \n    # Optimize hyperparameters if enabled (once per section)\n    section_key = f\"{section}\"\n    if cfg.optimize_hyperparams and section_key not in cfg.optimized_params:\n        print(f\"  Optimizing hyperparameters for section {section}...\")\n        # Use the first action with sufficient positive samples for optimization\n        optimization_done = False\n        for action in label.columns:\n            action_mask = ~ label[action].isna().values\n            y_action = label[action][action_mask].values.astype(int)\n            X_action = X[action_mask]\n            groups_action = meta.video_id[action_mask]\n            \n            if len(np.unique(groups_action)) >= cfg.n_splits and not (y_action == 0).all() and np.sum(y_action) > 10:\n                try:\n                    # Get target score for this section if available\n                    target_score = cfg.previous_best_scores.get(section, None)\n                    \n                    best_params = optimize_xgb_hyperparameters(\n                        X_action, y_action, groups_action, cfg.cv, \n                        cfg.use_gpu, n_trials=cfg.hyperopt_trials, \n                        timeout=cfg.hyperopt_timeout,\n                        section=section, action=action,\n                        use_class_weights=cfg.use_class_weights,\n                        use_early_stopping=cfg.use_early_stopping,\n                        early_stopping_rounds=cfg.early_stopping_rounds,\n                        use_section_specific_ranges=cfg.use_section_specific_ranges,\n                        use_smote=cfg.use_smote,  # Add this\n                        smote_threshold=cfg.smote_threshold,  # Add this\n                        target_best_score=target_score  # Pass target score for comparison\n                    )\n                    cfg.optimized_params[section_key] = best_params\n                    cfg.model = XGBClassifier(**best_params)\n                    optimization_done = True\n                    joblib.dump(best_params, f\"{cfg.model_name}/{section}/best_hyperparams.pkl\")\n                    print(f\"  Hyperparameter optimization complete for section {section}\")\n                    break\n                except Exception as e:\n                    print(f\"  Warning: Hyperparameter optimization failed: {e}\")\n                    print(f\"  Using default hyperparameters for section {section}\")\n                    cfg.optimized_params[section_key] = cfg.base_xgb_params.copy()\n                    cfg.model = XGBClassifier(**cfg.base_xgb_params)\n        \n        if not optimization_done:\n            print(f\"  No suitable action found for optimization, using default hyperparameters\")\n            cfg.optimized_params[section_key] = cfg.base_xgb_params.copy()\n            cfg.model = XGBClassifier(**cfg.base_xgb_params)\n    \n    # Load optimized params if available (for submit mode or if already optimized)\n    elif section_key in cfg.optimized_params:\n        cfg.model = XGBClassifier(**cfg.optimized_params[section_key])\n    elif cfg.mode == 'submit' or (cfg.mode == 'validate' and not cfg.optimize_hyperparams):\n        # Try to load saved hyperparameters (works for both submit and validate modes)\n        # First try in current model directory\n        hyperparams_path = f\"{cfg.model_name}/{section}/best_hyperparams.pkl\"\n        if not os.path.exists(hyperparams_path) and cfg.mode == 'submit':\n            # Fallback to model_path for submit mode\n            hyperparams_path = f\"{cfg.model_path}/{cfg.model_name}/{section}/best_hyperparams.pkl\"\n        \n        if os.path.exists(hyperparams_path):\n            cfg.optimized_params[section_key] = joblib.load(hyperparams_path)\n            # Ensure required parameters are set\n            if 'verbosity' not in cfg.optimized_params[section_key]:\n                cfg.optimized_params[section_key]['verbosity'] = 0\n            if 'random_state' not in cfg.optimized_params[section_key]:\n                cfg.optimized_params[section_key]['random_state'] = 42\n            if cfg.use_gpu and 'device' not in cfg.optimized_params[section_key]:\n                cfg.optimized_params[section_key]['device'] = 'cuda'\n            elif not cfg.use_gpu and 'device' in cfg.optimized_params[section_key]:\n                # Remove device if GPU is not enabled\n                cfg.optimized_params[section_key] = {k: v for k, v in cfg.optimized_params[section_key].items() if k != 'device'}\n            \n            cfg.model = XGBClassifier(**cfg.optimized_params[section_key])\n            print(f\"  Loaded optimized hyperparameters for section {section}\")\n        else:\n            cfg.model = XGBClassifier(**cfg.base_xgb_params)\n    \n    for action in label.columns:\n        action_mask = ~ label[action].isna().values\n        y_action = label[action][action_mask].values.astype(int)\n        X_action = X[action_mask]\n        groups_action = meta.video_id[action_mask]\n        \n        if len(np.unique(groups_action)) < cfg.n_splits:\n            continue\n\n        if not (y_action == 0).all():\n            try:\n                action_start_time = time.time()\n                \n                # Calculate positive ratio for class imbalance handling\n                positive_ratio = np.sum(y_action) / len(y_action) if len(y_action) > 0 else 0\n                \n                with warnings.catch_warnings():\n                    warnings.filterwarnings('ignore', category=RuntimeWarning)\n                    \n                    # Determine if SMOTE should be used\n                    use_smote = (cfg.use_smote and SMOTE_AVAILABLE and \n                                positive_ratio < cfg.smote_threshold and \n                                len(np.unique(y_action)) == 2)\n                    \n                    # Prepare model with class weights (if not using SMOTE, or as fallback)\n                    model_with_weights = clone(cfg.model)\n                    if cfg.use_class_weights and not use_smote:\n                        # Only use class weights if not using SMOTE (SMOTE balances classes)\n                        n_positive = np.sum(y_action)\n                        n_negative = len(y_action) - n_positive\n                        if n_positive > 0 and n_negative > 0:\n                            scale_pos_weight = min(n_negative / n_positive, 100.0)\n                            model_with_weights.set_params(scale_pos_weight=scale_pos_weight)\n                    \n                    # Use custom CV with SMOTE if enabled, otherwise use Trainer\n                    if use_smote:\n                        print(f\"\\t  Using SMOTE within CV folds (positive_ratio: {positive_ratio:.4f})\")\n                        oof_action, trained_estimators, fold_scores = cross_validate_with_smote(\n                            X=X_action,\n                            y=y_action,\n                            groups=groups_action,\n                            cv=cfg.cv,\n                            estimator=model_with_weights,\n                            use_smote=True,\n                            smote_threshold=cfg.smote_threshold,\n                            use_class_weights=cfg.use_class_weights,\n                            use_early_stopping=cfg.use_early_stopping,\n                            early_stopping_rounds=cfg.early_stopping_rounds,\n                            metric=f1_score,\n                            save_path=f\"{cfg.model_name}/{section}/{action}\",\n                            verbose=False\n                        )\n                        \n                        # Create a mock trainer object for compatibility with existing code\n                        class MockTrainer:\n                            def __init__(self, oof_preds, estimators):\n                                self.oof_preds = oof_preds\n                                self.estimators = estimators\n                        \n                        trainer = MockTrainer(oof_action, trained_estimators)\n                    else:\n                        # Use standard Trainer without SMOTE\n                        trainer = Trainer(\n                            estimator=model_with_weights,\n                            cv=cfg.cv,\n                            cv_args={\"groups\": groups_action},\n                            metric=f1_score,\n                            task=\"binary\",\n                            verbose=False,\n                            save=True,\n                            save_path=f\"{cfg.model_name}/{section}/{action}\"\n                        )\n                        trainer.fit(X_action, y_action)\n                        oof_action = trainer.oof_preds\n                    \n                    threshold = tune_threshold(oof_action, y_action, n_trials=cfg.n_trials, \n                                              positive_ratio=positive_ratio)\n                    thresholds[action] = threshold\n            \n                    f1 = f1_score(y_action, (oof_action >= threshold), zero_division=0)\n                    f1_list.append((body_parts_tracked_str, action, f1))\n                    \n                    # Analyze feature importance for underperforming sections (F1 < 0.5) or Section 3\n                    analyze_feature_importance = (f1 < 0.5) or (section == 3)\n                    if analyze_feature_importance:\n                        try:\n                            # Try to get feature importance from trainer or saved models\n                            feature_importances = None\n                            feature_names = X_action.columns.tolist() if hasattr(X_action, 'columns') else None\n                            \n                            # Try to get from trainer.estimators (list of trained models from CV folds)\n                            if hasattr(trainer, 'estimators') and len(trainer.estimators) > 0:\n                                # Average feature importance across all CV folds\n                                all_importances = []\n                                for est in trainer.estimators:\n                                    if hasattr(est, 'feature_importances_'):\n                                        all_importances.append(est.feature_importances_)\n                                if len(all_importances) > 0:\n                                    feature_importances = np.mean(all_importances, axis=0)\n                            # Fallback: try trainer.estimator (base estimator before training)\n                            elif hasattr(trainer, 'estimator') and hasattr(trainer.estimator, 'feature_importances_'):\n                                feature_importances = trainer.estimator.feature_importances_\n                            # Last resort: try to load from saved model files\n                            elif hasattr(trainer, 'save_path'):\n                                saved_models = glob.glob(f\"{trainer.save_path}/*_trainer_*.pkl\")\n                                if len(saved_models) > 0:\n                                    try:\n                                        saved_trainer = joblib.load(saved_models[0])\n                                        if hasattr(saved_trainer, 'estimators') and len(saved_trainer.estimators) > 0:\n                                            all_importances = []\n                                            for est in saved_trainer.estimators:\n                                                if hasattr(est, 'feature_importances_'):\n                                                    all_importances.append(est.feature_importances_)\n                                            if len(all_importances) > 0:\n                                                feature_importances = np.mean(all_importances, axis=0)\n                                    except:\n                                        pass\n                            \n                            # If we have feature importances, analyze them\n                            if feature_importances is not None and len(feature_importances) > 0:\n                                if feature_names is None:\n                                    feature_names = [f'f{i}' for i in range(len(feature_importances))]\n                                \n                                # Create feature importance DataFrame\n                                importance_df = pd.DataFrame({\n                                    'feature': feature_names[:len(feature_importances)],\n                                    'importance': feature_importances\n                                }).sort_values('importance', ascending=False)\n                                \n                                # Log top features to wandb\n                                if WANDB_AVAILABLE and wandb.run is not None:\n                                    top_n = min(20, len(importance_df))\n                                    top_features = importance_df.head(top_n)\n                                    \n                                    # Log as table\n                                    wandb.log({\n                                        f'feature_importance/{section}/{action}': wandb.Table(\n                                            dataframe=top_features,\n                                            columns=['feature', 'importance']\n                                        ),\n                                        f'feature_importance/top1/{section}/{action}': top_features.iloc[0]['importance'] if len(top_features) > 0 else 0,\n                                        f'feature_importance/top5_mean/{section}/{action}': top_features.head(5)['importance'].mean() if len(top_features) >= 5 else 0,\n                                    })\n                                \n                                # Print top features for monitoring\n                                if section == 3 or f1 < 0.3:\n                                    print(f\"\\t  Top 10 features for {action} (F1={f1:.4f}):\")\n                                    for idx, row in importance_df.head(10).iterrows():\n                                        print(f\"\\t    {row['feature']}: {row['importance']:.4f}\")\n                        except Exception as e:\n                            print(f\"\\t  Warning: Could not analyze feature importance: {e}\")\n                    \n                    action_time = time.time() - action_start_time\n                    n_samples = len(y_action)\n                    n_positive = np.sum(y_action)\n                    n_negative = n_samples - n_positive\n                    \n                    # Log to wandb if enabled\n                    if WANDB_AVAILABLE and wandb.run is not None:\n                        wandb.log({\n                            f'f1/{section}/{action}': f1,\n                            f'threshold/{section}/{action}': threshold,\n                            f'training_time/{section}/{action}': action_time,\n                            f'n_samples/{section}/{action}': n_samples,\n                            f'n_positive/{section}/{action}': n_positive,\n                            f'n_negative/{section}/{action}': n_negative,\n                            f'positive_ratio/{section}/{action}': n_positive / n_samples if n_samples > 0 else 0,\n                            'section': section,\n                            'action': action,\n                            'action_idx': action_idx,\n                            'step': action_idx,  # For time-series visualization\n                        })\n                    \n                    action_idx += 1\n                    \n                    joblib.dump(oof_action, f\"{cfg.model_name}/{section}/{action}/oof_pred_probs.pkl\")\n                    joblib.dump(threshold, f\"{cfg.model_name}/{section}/{action}/threshold.pkl\")\n                    \n                    print(f\"\\tF1: {f1:.4f} ({threshold:.2f}) Section: {section} Action: {action}\")\n    \n                    del trainer\n                    gc.collect()\n\n            except Exception as e:\n                oof_action = np.zeros(len(y_action))\n                print(f\"\\tF1: 0.0000 (0.00) Section: {section} Action: {action} - Error: {e}\")\n        \n        else:\n            oof_action = np.zeros(len(y_action))\n            print(f\"\\tF1: 0.0000 (0.00) Section: {section} Action: {action}\")\n        \n        oof_column = np.zeros(len(label))\n        oof_column[action_mask] = oof_action\n        oof[action] = oof_column\n\n        del oof_action, action_mask, X_action, y_action, groups_action\n        gc.collect()\n\n    submission_part = predict_multiclass(oof, meta, thresholds)\n    submission_list.append(submission_part)\n    \n    # Log section-level summary metrics\n    if WANDB_AVAILABLE and wandb.run is not None and len(f1_list) > 0:\n        section_f1s = [f1 for _, _, f1 in f1_list]\n        section_time = time.time() - section_start_time\n        wandb.log({\n            f'section_{section}/mean_f1': np.mean(section_f1s),\n            f'section_{section}/std_f1': np.std(section_f1s),\n            f'section_{section}/min_f1': np.min(section_f1s),\n            f'section_{section}/max_f1': np.max(section_f1s),\n            f'section_{section}/n_actions': len(section_f1s),\n            f'section_{section}/total_time': section_time,\n            'section': section,\n        })\n    \n    return submission_list, f1_list, thresholds\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"ffa5d086","cell_type":"markdown","source":"## Main Training Loop","metadata":{}},{"id":"d010b073","cell_type":"code","source":"# Load data\nprint(\"Loading data...\")\ntrain = pd.read_csv(cfg.train_path)\ntrain['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\ntrain_without_mabe22 = train.query(\"~lab_id.str.startswith('MABe22_')\")\n\ntest = pd.read_csv(cfg.test_path)\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\n# Initialize thresholds\nthresholds = {\n    \"single\": {},\n    \"pair\": {}\n}\nsolution = create_solution_df(train_without_mabe22, cfg)\n\n# Main training/validation loop\nf1_list = []\nsubmission_list = []\ntotal_sections = len(body_parts_tracked_list) - 1\n\n# Log total sections to wandb\nif WANDB_AVAILABLE and wandb.run is not None:\n    wandb.config.update({'total_sections': total_sections})\n\nfor section in range(1, len(body_parts_tracked_list)):\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        print(f\"{section}/{total_sections} Processing videos with: {body_parts_tracked}\\n\")\n        \n        # Special monitoring for Section 3\n        if section == 3:\n            print(\"=\" * 80)\n            print(\"MONITORING SECTION 3 - Close performance tracking enabled\")\n            print(\"=\" * 80)\n        \n        # Log section progress to wandb\n        if WANDB_AVAILABLE and wandb.run is not None:\n            wandb.log({\n                'progress/section': section,\n                'progress/total_sections': total_sections,\n                'progress/section_pct': (section / total_sections) * 100,\n            })\n        \n        if len(body_parts_tracked) > 5:\n            body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n    \n        train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n\n        _fps_lookup = (\n            train_subset[['video_id', 'frames_per_second']]\n            .drop_duplicates('video_id')\n            .set_index('video_id')['frames_per_second']\n            .to_dict()\n        )\n        \n        single_mouse_list = []\n        single_mouse_label_list = []\n        single_mouse_meta_list = []\n        \n        mouse_pair_list = []\n        mouse_pair_label_list = []\n        mouse_pair_meta_list = []\n    \n        for switch, data, meta, label in generate_mouse_data(train_subset, 'train', cfg):\n            if switch == 'single':\n                single_mouse_list.append(data)\n                single_mouse_meta_list.append(meta)\n                single_mouse_label_list.append(label)\n            else:\n                mouse_pair_list.append(data)\n                mouse_pair_meta_list.append(meta)\n                mouse_pair_label_list.append(label)\n            \n            del data, meta, label\n        gc.collect()\n    \n    \n        if len(single_mouse_list) > 0:\n            single_feats_parts = []\n            for data_i, meta_i in zip(single_mouse_list, single_mouse_meta_list):\n                fps_i = _fps_from_meta(meta_i, _fps_lookup, default_fps=30.0)\n                X_i = transform_single(data_i, body_parts_tracked, fps_i).astype(np.float32)\n                single_feats_parts.append(X_i)\n                del X_i, fps_i\n            gc.collect()\n\n            X_tr = pd.concat(single_feats_parts, axis=0, ignore_index=True)\n            single_mouse_label = pd.concat(single_mouse_label_list, axis=0, ignore_index=True)\n            single_mouse_meta = pd.concat(single_mouse_meta_list, axis=0, ignore_index=True)\n            \n            del single_feats_parts, single_mouse_list, single_mouse_label_list, single_mouse_meta_list\n            gc.collect()\n\n            temp_submission_list, temp_f1_list, temp_thresholds = cross_validate_classifier(\n                X_tr, single_mouse_label, single_mouse_meta, body_parts_tracked_str, section, cfg)\n            \n            if f\"{section}\" not in thresholds[\"single\"].keys():\n                thresholds[\"single\"][f\"{section}\"] = {}\n            for k, v in temp_thresholds.items():\n                thresholds[\"single\"][f\"{section}\"][k] = v\n            \n            f1_list.extend(temp_f1_list)\n            submission_list.extend(temp_submission_list)\n            \n            del temp_submission_list, temp_f1_list, temp_thresholds, X_tr\n            gc.collect()\n            \n        if len(mouse_pair_list) > 0:\n            pair_feats_parts = []\n            for data_i, meta_i in zip(mouse_pair_list, mouse_pair_meta_list):\n                fps_i = _fps_from_meta(meta_i, _fps_lookup, default_fps=30.0)\n                X_i = transform_pair(data_i, body_parts_tracked, fps_i).astype(np.float32)\n                pair_feats_parts.append(X_i)\n                del X_i, fps_i\n            gc.collect()\n\n            X_tr = pd.concat(pair_feats_parts, axis=0, ignore_index=True)\n            mouse_pair_label = pd.concat(mouse_pair_label_list, axis=0, ignore_index=True)\n            mouse_pair_meta = pd.concat(mouse_pair_meta_list, axis=0, ignore_index=True)\n            \n            del pair_feats_parts, mouse_pair_list, mouse_pair_label_list, mouse_pair_meta_list\n            gc.collect()\n\n            temp_submission_list, temp_f1_list, temp_thresholds = cross_validate_classifier(\n                X_tr, mouse_pair_label, mouse_pair_meta, body_parts_tracked_str, section, cfg)\n\n            if f\"{section}\" not in thresholds[\"pair\"].keys():\n                thresholds[\"pair\"][f\"{section}\"] = {}\n            for k, v in temp_thresholds.items():\n                thresholds[\"pair\"][f\"{section}\"][k] = v\n            \n            f1_list.extend(temp_f1_list)\n            submission_list.extend(temp_submission_list)\n            \n            del temp_submission_list, temp_f1_list, temp_thresholds, X_tr\n            gc.collect()\n            \n    except Exception as e:\n        print(f\"\\tError in section {section}: {e}\")\n        import traceback\n        traceback.print_exc()\n    print()\n","metadata":{},"outputs":[],"execution_count":null},{"id":"d40d23b8","cell_type":"markdown","source":"## Save Results and Thresholds","metadata":{}},{"id":"cf3f8d1f","cell_type":"code","source":"# Finalize results\nsubmission = pd.concat(submission_list)\nsubmission_robust = robustify(submission, train, 'train', cfg)\ncompetition_score = score(solution, submission_robust, '')\nprint(f\"Competition metric: {competition_score:.4f}\")\n\nf1_df = pd.DataFrame(f1_list, columns=['body_parts_tracked_str', 'action', 'binary F1 score'])\nmean_f1 = f1_df['binary F1 score'].mean()\nprint(f\"Mean F1:            {mean_f1:.4f}\")\n\n# Log final metrics to wandb\nif WANDB_AVAILABLE and wandb.run is not None:\n    wandb.log({\n        'competition_score': competition_score,\n        'mean_f1': mean_f1,\n        'total_actions': len(f1_df),\n    })\n    \n    # Log F1 scores per action as a table\n    wandb.log({\n        'f1_scores_table': wandb.Table(dataframe=f1_df)\n    })\n\nos.makedirs(cfg.model_name, exist_ok=True)\njoblib.dump(thresholds, f\"{cfg.model_name}/thresholds.pkl\")\njoblib.dump(f1_df, f\"{cfg.model_name}/scores.pkl\")\nprint(f\"\\nResults saved to {cfg.model_name}/\")\n\n# Finish wandb run\nif WANDB_AVAILABLE and wandb.run is not None:\n    wandb.finish()\n","metadata":{},"outputs":[],"execution_count":null},{"id":"b70fc34c","cell_type":"markdown","source":"## Create Kaggle Dataset (Optional)","metadata":{}},{"id":"f0e77cba","cell_type":"code","source":"# ============================================================================\n# AUTO-CREATE KAGGLE DATASET FROM TRAINED MODELS\n# ============================================================================\n# This cell automatically creates a Kaggle dataset from your trained models\n# so you can use them in other notebooks (e.g., inference notebook)\n#\n# [WARN] IMPORTANT: Update the configuration below before running!\n# ============================================================================\n\nimport os\nimport json\nfrom pathlib import Path\nfrom datetime import datetime\n\n# ============================================================================\n# CONFIGURATION - UPDATE THESE VALUES\n# ============================================================================\nYOUR_USERNAME = \"bvproperty\"  # [WARN] CHANGE THIS to your Kaggle username (e.g., \"bvproperty\")\nDATASET_SLUG = \"mabe-xgboost-models\"  # [WARN] Change if you want a different dataset name\nAUTO_CREATE_DATASET = True  # Set to False to skip dataset creation\n\n# ============================================================================\n# DATASET CREATION FUNCTION\n# ============================================================================\n\ndef create_kaggle_dataset_from_models(model_dir: str, dataset_handle: str, version_notes: str = None):\n    \"\"\"\n    Create a Kaggle dataset from trained models using kagglehub.\n    \n    Args:\n        model_dir: Path to model directory (e.g., '/kaggle/working/xgboost')\n        dataset_handle: Kaggle dataset handle (format: 'username/dataset-slug')\n        version_notes: Optional version notes\n    \"\"\"\n    try:\n        from kagglehub import dataset_upload\n        \n        model_path = Path(model_dir)\n        if not model_path.exists():\n            print(f\"[ERROR] Model directory not found: {model_dir}\")\n            return False\n        \n        # Check what was saved\n        print(f\"\\n{'='*60}\")\n        print(\"CHECKING TRAINED MODELS\")\n        print(f\"{'='*60}\")\n        \n        thresholds_file = model_path / 'thresholds.pkl'\n        scores_file = model_path / 'scores.pkl'\n        \n        print(f\"Model directory: {model_dir}\")\n        print(f\"Thresholds file: {'[OK] Found' if thresholds_file.exists() else '[ERROR] Missing'}\")\n        print(f\"Scores file: {'[OK] Found' if scores_file.exists() else '[ERROR] Missing'}\")\n        \n        # Count sections\n        sections = [d for d in model_path.iterdir() if d.is_dir() and d.name.isdigit()]\n        print(f\"Sections found: {len(sections)}\")\n        \n        if version_notes is None:\n            timestamp = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n            version_notes = f\"XGBoost models - Training completed {timestamp}\"\n        \n        print(f\"\\n{'='*60}\")\n        print(\"CREATING KAGGLE DATASET\")\n        print(f\"{'='*60}\")\n        print(f\"Dataset: {dataset_handle}\")\n        print(\"Uploading... (this may take several minutes)\")\n        \n        dataset_upload(\n            handle=dataset_handle,\n            local_dataset_dir=str(model_path),\n            version_notes=version_notes,\n            ignore_patterns=[\"*.pyc\", \"__pycache__/\", \".git/\", \"*.ipynb_checkpoints\", \"*.ipynb\"]\n        )\n        \n        print(f\"\\n{'='*60}\")\n        print(\"[OK] SUCCESS!\")\n        print(f\"{'='*60}\")\n        print(f\"Dataset created: {dataset_handle}\")\n        print(f\"View at: https://www.kaggle.com/datasets/{dataset_handle}\")\n        print(f\"\\n[NOTE] Next steps:\")\n        print(f\"   1. Go to your inference notebook\")\n        print(f\"   2. Add this dataset as input: {dataset_handle}\")\n        dataset_slug = dataset_handle.split('/')[-1]\n        print(f\"   3. Update MODEL_PATH to: '/kaggle/input/{dataset_slug}/xgboost'\")\n        return True\n        \n    except ImportError:\n        print(f\"\\n{'='*60}\")\n        print(\"[WARN]  kagglehub not available\")\n        print(f\"{'='*60}\")\n        print(\"To enable automatic dataset creation:\")\n        print(\"   1. Enable Internet in notebook settings (Settings → Internet → ON)\")\n        print(\"   2. The kagglehub installation cell should have run earlier\")\n        print(\"\\nAlternatively, download models manually:\")\n        print(\"   1. Go to notebook Output tab\")\n        print(\"   2. Download the 'xgboost' folder\")\n        print(\"   3. Upload as a Kaggle dataset manually\")\n        return False\n    except Exception as e:\n        print(f\"\\n[ERROR] Error creating dataset: {e}\")\n        import traceback\n        traceback.print_exc()\n        print(\"\\n💡 You can still download models manually from the Output tab\")\n        return False\n\n# ============================================================================\n# RUN DATASET CREATION\n# ============================================================================\n\nif AUTO_CREATE_DATASET:\n    if YOUR_USERNAME == \"bvproperty\":\n        print(\"\\n[WARN]  WARNING: Please update YOUR_USERNAME in the configuration above!\")\n        print(\"   Dataset creation skipped. Update the config and re-run this cell.\")\n    else:\n        model_dir = f'/kaggle/working/{cfg.model_name}'\n        dataset_handle = f\"{YOUR_USERNAME}/{DATASET_SLUG}\"\n        \n        if os.path.exists(model_dir):\n            create_kaggle_dataset_from_models(\n                model_dir=model_dir,\n                dataset_handle=dataset_handle,\n                version_notes=f\"XGBoost models trained on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\"\n            )\n        else:\n            print(f\"\\n[ERROR] Model directory not found: {model_dir}\")\n            print(\"   Training may not have completed successfully.\")\nelse:\n    print(\"\\n[SKIP] Dataset creation disabled (AUTO_CREATE_DATASET = False)\")\n    print(f\"   Models are available in: /kaggle/working/{cfg.model_name}\")\n    print(\"   You can download them manually from the Output tab\")\n","metadata":{},"outputs":[],"execution_count":null}]}