{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13877892,"sourceType":"datasetVersion","datasetId":8835951},{"sourceId":14060843,"sourceType":"datasetVersion","datasetId":8624642}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false},"papermill":{"default_parameters":{},"duration":1494.813501,"end_time":"2025-10-25T19:51:35.870094","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-10-25T19:26:41.056593","version":"2.6.0"},"widgets":{"application/vnd.jupyter.widget-state+json":{"state":{"00962fe9010e42b7a5b123af50d2a026":{"model_module":"@jupyter-widgets/base","model_module_version":"2.0.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"2.0.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border_bottom":null,"border_left":null,"border_right":null,"border_top":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"0b098fa23a3b4031bf88b32b20044978":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HTMLStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HTMLStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"StyleView","background":null,"description_width":"","font_size":null,"text_color":null}},"108c3e5db1154f47a5aca68386ed4f24":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"ProgressStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"38c36db3be3843b19da108d4cd1f9421":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HBoxModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"2.0.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_3e9806911d1d489d963f76955d8468c4","IPY_MODEL_e6ad60bc760e4b51885b15335c048a65","IPY_MODEL_6f4723180bf14101b60deeac9fd59e50"],"layout":"IPY_MODEL_00962fe9010e42b7a5b123af50d2a026","tabbable":null,"tooltip":null}},"3e9806911d1d489d963f76955d8468c4":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"2.0.0","_view_name":"HTMLView","description":"","description_allow_html":false,"layout":"IPY_MODEL_aeb2544f9ad9413b8b71b5f5643b1019","placeholder":"​","style":"IPY_MODEL_ec012e9ac53d477daee0988a35117f0b","tabbable":null,"tooltip":null,"value":"100%"}},"3edd91e1eafb4f6bbf55275ea60b6c97":{"model_module":"@jupyter-widgets/base","model_module_version":"2.0.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"2.0.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border_bottom":null,"border_left":null,"border_right":null,"border_top":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"6f4723180bf14101b60deeac9fd59e50":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"2.0.0","_view_name":"HTMLView","description":"","description_allow_html":false,"layout":"IPY_MODEL_e8f95c8096294ce6b6fb2b65c08af536","placeholder":"​","style":"IPY_MODEL_0b098fa23a3b4031bf88b32b20044978","tabbable":null,"tooltip":null,"value":" 863/863 [00:08<00:00, 96.46it/s]"}},"aeb2544f9ad9413b8b71b5f5643b1019":{"model_module":"@jupyter-widgets/base","model_module_version":"2.0.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"2.0.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border_bottom":null,"border_left":null,"border_right":null,"border_top":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"e6ad60bc760e4b51885b15335c048a65":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"FloatProgressModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"2.0.0","_view_name":"ProgressView","bar_style":"success","description":"","description_allow_html":false,"layout":"IPY_MODEL_3edd91e1eafb4f6bbf55275ea60b6c97","max":863,"min":0,"orientation":"horizontal","style":"IPY_MODEL_108c3e5db1154f47a5aca68386ed4f24","tabbable":null,"tooltip":null,"value":863}},"e8f95c8096294ce6b6fb2b65c08af536":{"model_module":"@jupyter-widgets/base","model_module_version":"2.0.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"2.0.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border_bottom":null,"border_left":null,"border_right":null,"border_top":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"ec012e9ac53d477daee0988a35117f0b":{"model_module":"@jupyter-widgets/controls","model_module_version":"2.0.0","model_name":"HTMLStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"2.0.0","_model_name":"HTMLStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"2.0.0","_view_name":"StyleView","background":null,"description_width":"","font_size":null,"text_color":null}}},"version_major":2,"version_minor":0}}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Imports and configs","metadata":{"papermill":{"duration":0.007171,"end_time":"2025-10-25T19:26:46.426493","exception":false,"start_time":"2025-10-25T19:26:46.419322","status":"completed"},"tags":[]}},{"cell_type":"code","source":"\"\"\"F Beta customized for the data format of the MABe challenge.\"\"\"\n\nimport json\n\nfrom collections import defaultdict\n\nimport pandas as pd\nimport polars as pl\n\n\nclass HostVisibleError(Exception):\n pass\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() # ty: ignore\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 # Since the labels are sparse, we can't evaluate prediction keys not in the active labels.\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 # Ignore truly redundant predictions.\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 # A single agent can have multiple targets per frame (ex: evading all other mice) but only one action per target per frame.\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\ndef mouse_fbeta(solution: pd.DataFrame, submission: pd.DataFrame, beta: float = 1) -> float:\n \"\"\"\n Doctests:\n >>> solution = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10},\n ... ])\n >>> mouse_fbeta(solution, submission)\n 1.0\n\n >>> solution = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'mount', 'start_frame': 0, 'stop_frame': 10}, # Wrong action\n ... ])\n >>> mouse_fbeta(solution, submission)\n 0.0\n\n >>> solution = pd.DataFrame([\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'mount', 'start_frame': 15, 'stop_frame': 24, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9},\n ... ])\n >>> \"%.12f\" % mouse_fbeta(solution, submission)\n '0.500000000000'\n\n >>> solution = pd.DataFrame([\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'mount', 'start_frame': 15, 'stop_frame': 24, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 345, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9, 'lab_id': 2, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 345, 'agent_id': 1, 'target_id': 2, 'action': 'mount', 'start_frame': 15, 'stop_frame': 24, 'lab_id': 2, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9},\n ... ])\n >>> \"%.12f\" % mouse_fbeta(solution, submission)\n '0.250000000000'\n\n >>> # Overlapping solution events, one prediction matching both.\n >>> solution = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 10, 'stop_frame': 20, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 20},\n ... ])\n >>> mouse_fbeta(solution, submission)\n 1.0\n\n >>> solution = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 30, 'stop_frame': 40, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 40},\n ... ])\n >>> mouse_fbeta(solution, submission)\n 0.6666666666666666\n \"\"\"\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 # Need to align based on video IDs as we can't rely on the row IDs for handling public/private splits.\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\ndef score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, beta: float = 1) -> float:\n \"\"\"\n F1 score for the MABe Challenge\n \"\"\"\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)","metadata":{"_kg_hide-input":true,"jupyter":{"source_hidden":true},"papermill":{"duration":2.975797,"end_time":"2025-10-25T19:26:49.408873","exception":false,"start_time":"2025-10-25T19:26:46.433076","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import StratifiedGroupKFold, StratifiedShuffleSplit, train_test_split\nfrom sklearn.metrics import f1_score, fbeta_score\nfrom sklearn.base import clone\nfrom xgboost import XGBClassifier\nimport xgboost as xgb\nfrom catboost import CatBoostClassifier\nfrom tqdm.notebook import tqdm\nfrom typing import Dict, Optional, Tuple, Any\nfrom collections import defaultdict\nimport polars as pl\nfrom scipy import signal, stats\nimport cupy as cp\n\nimport numpy as np\nimport itertools, traceback\nimport warnings\nimport optuna\nimport joblib\nimport glob\nimport gc, re, os, math, time\n\noptuna.logging.set_verbosity(optuna.logging.WARNING)\nwarnings.filterwarnings('ignore')\n\nfrom xgboost.core import XGBoostError\nclass XGBoostError(Exception):\n pass\n","metadata":{"papermill":{"duration":8.578306,"end_time":"2025-10-25T19:26:57.993688","exception":false,"start_time":"2025-10-25T19:26:49.415382","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CFG:\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 seed = 42\n\n use_gpu = True\n use_cudf = False\n use_holdout = True\n\n mode = \"submit\"\n\n model_path = (\n \"/kaggle/input/mabe-xgb-cluster/XGB_v9.6\"#\"/kaggle/input/mabe-xgb-cluster\"\n )\n\n n_splits = 3\n cv = StratifiedGroupKFold(n_splits)\n holdout_frac = 0.2\n\n\n xgb_model = XGBClassifier(**({\n \"verbosity\": 0, \"random_state\": 42,\n \"n_estimators\": 250, \"learning_rate\": 0.08,\n \"max_depth\": 6, \"min_child_weight\": 5,\n \"subsample\": 0.8, \"colsample_bytree\": 0.8,\n \"tree_method\": \"gpu_hist\", \"gpu_id\": 0\n }))\n\n\n # ----------------- Ensemble registry -----------------\n models = [xgb_model]\n model_tags = [\"xgb_baseline\"]\n\n# -------- sanitize / expand all paths once --------\ndef _expand_path(p: str) -> str:\n return os.path.abspath(os.path.expanduser(p))\n\nfor _attr in [\n \"train_path\",\n \"test_path\",\n \"train_annotation_path\",\n \"train_tracking_path\",\n \"test_tracking_path\",\n \"model_path\",\n]:\n val = getattr(CFG, _attr, None)\n if isinstance(val, str):\n setattr(CFG, _attr, _expand_path(val))\n\nprint(f\"Using GPU? {CFG.use_gpu}\")\nprint(f\"Model path: {CFG.model_path}\")\n","metadata":{"papermill":{"duration":0.014723,"end_time":"2025-10-25T19:26:58.043155","exception":false,"start_time":"2025-10-25T19:26:58.028432","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if CFG.use_cudf:\n try:\n import pandas as pd_cpu\n import cudf\n import cudf.pandas\n\n # Install the pandas accelerator: patches pandas to use cuDF under the hood\n cudf.pandas.install()\n\n import pandas as pd\n print(\"✅ Using cuDF.pandas accelerator (pandas API on GPU)\")\n except Exception as e:\n print(f\"⚠️ Failed to enable cuDF.pandas accelerator, falling back to CPU pandas: {e}\")\n import pandas as pd\n cudf = None\nelse:\n import pandas as pd_cpu\n cudf = None\n print(\"Using plain CPU pandas\")\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"labs_dict = {\n \"MABe22\": {\n \"labs\": [\"MABe22_keypoints\", \"MABe22_movies\"],\n \"frames\": 14266800,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"forepaw_left\", \"forepaw_right\",\n \"hindpaw_left\", \"hindpaw_right\", \"neck\", \"nose\", \"tail_base\",\n \"body_center\"\n ],\n \"actions\": [],\n \"remove_video_id\": [],\n },\n \"CalMS21_supplemental\": {\n \"labs\": [\"CalMS21_supplemental\"],\n \"frames\": 5544597,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\n \"approach\", \"attack\", \"attemptmount\", \"dominancemount\", \"intromit\",\n \"mount\", \"sniff\", \"sniffbody\", \"sniffface\", \"sniffgenital\"\n ],\n \"remove_video_id\": [],\n },\n \"BoisterousParrot\": {\n \"labs\": [\"BoisterousParrot\"],\n \"frames\": 4752204,\n \"body_parts\": [\"ear_left\", \"ear_right\", \"nose\", \"tail_base\", \"body_center\"],\n \"actions\": [\"shepherd\"],\n \"remove_video_id\": [],\n },\n \"CalMS21_task2\": {\n \"labs\": [\"CalMS21_task2\"],\n \"frames\": 1491318,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"attack\", \"mount\", \"sniff\"],\n \"remove_video_id\": [],\n },\n \"TranquilPanther\": {\n \"labs\": [\"TranquilPanther\"],\n \"frames\": 1234590,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"intromit\", \"mount\", \"rear\", \"selfgroom\", \"sniff\", \"sniffgenital\"],\n \"remove_video_id\": [],\n },\n \"CalMS21_task1\": {\n \"labs\": [\"CalMS21_task1\"],\n \"frames\": 981129,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\n \"approach\", \"attack\", \"genitalgroom\", \"intromit\", \"mount\",\n \"sniff\", \"sniffbody\", \"sniffface\", \"sniffgenital\"\n ],\n \"remove_video_id\": [],\n },\n \"PleasantMeerkat\": {\n \"labs\": [\"PleasantMeerkat\"],\n \"frames\": 972000,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"lateral_left\", \"lateral_right\",\n \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"attack\", \"chase\", \"escape\", \"follow\"],\n \"remove_video_id\": [\"1375833299\"],\n },\n \"ElegantMink\": {\n \"labs\": [\"ElegantMink\"],\n \"frames\": 910341,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\n \"allogroom\", \"attack\", \"attemptmount\", \"intromit\", \"mount\",\n \"rear\", \"selfgroom\", \"sniff\", \"sniffgenital\"\n ],\n \"remove_video_id\": [],\n },\n \"LyricalHare\": {\n \"labs\": [\"LyricalHare\"],\n \"frames\": 887304,\n \"body_parts\": [\"ear_left\", \"ear_right\", \"nose\", \"tail_base\", \"body_center\"],\n \"actions\": [\"approach\", \"attack\", \"defend\", \"escape\", \"freeze\", \"rear\", \"sniff\"],\n \"remove_video_id\": [],\n },\n \"AdaptableSnail\": {\n \"labs\": [\"AdaptableSnail\"],\n \"frames\": 733220,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"lateral_left\", \"lateral_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"approach\", \"attack\", \"avoid\", \"chase\", \"chaseattack\", \"rear\", \"submit\"],\n \"remove_video_id\": [\"1212811043\"],\n },\n \"SparklingTapir\": {\n \"labs\": [\"SparklingTapir\"],\n \"frames\": 704056,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"lateral_left\", \"lateral_right\",\n \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"attack\", \"defend\", \"escape\", \"mount\", \"sniffgenital\"],\n \"remove_video_id\": [],\n },\n \"InvincibleJellyfish\": {\n \"labs\": [\"InvincibleJellyfish\"],\n \"frames\": 565521,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\n \"allogroom\", \"attack\", \"dig\", \"dominancegroom\", \"escape\",\n \"selfgroom\", \"sniff\", \"sniffgenital\"\n ],\n \"remove_video_id\": [],\n },\n \"JovialSwallow\": {\n \"labs\": [\"JovialSwallow\"],\n \"frames\": 468000,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"attack\", \"chase\", \"sniff\"],\n \"remove_video_id\": [],\n },\n \"GroovyShrew\": {\n \"labs\": [\"GroovyShrew\"],\n \"frames\": 450756,\n \"body_parts\": [\"ear_left\", \"ear_right\", \"head\", \"tail_base\", \"body_center\"],\n \"actions\": [\n \"approach\", \"attemptmount\", \"climb\", \"defend\", \"dig\", \"escape\",\n \"intromit\", \"mount\", \"rear\", \"rest\", \"run\", \"selfgroom\", \"sniff\", \"sniffgenital\"\n ],\n \"remove_video_id\": [],\n },\n \"NiftyGoldfinch\": {\n \"labs\": [\"NiftyGoldfinch\"],\n \"frames\": 354249,\n \"body_parts\": [\"ear_left\", \"ear_right\", \"nose\", \"tail_base\", \"body_center\"],\n \"actions\": [\n \"approach\", \"attack\", \"biteobject\", \"chase\", \"climb\", \"defend\",\n \"dig\", \"escape\", \"exploreobject\", \"flinch\", \"follow\", \"rear\",\n \"selfgroom\", \"sniff\", \"sniffface\", \"sniffgenital\", \"tussle\"\n ],\n \"remove_video_id\": [],\n },\n \"DeliriousFly\": {\n \"labs\": [\"DeliriousFly\"],\n \"frames\": 324027,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"lateral_left\", \"lateral_right\",\n \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"attack\", \"dominance\", \"sniff\"],\n \"remove_video_id\": [],\n },\n \"UppityFerret\": {\n \"labs\": [\"UppityFerret\"],\n \"frames\": 314354,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"lateral_left\", \"lateral_right\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\n \"huddle\", \"intromit\", \"mount\", \"rear\", \"reciprocalsniff\",\n \"selfgroom\", \"sniff\", \"sniffgenital\"\n ],\n \"remove_video_id\": [],\n },\n \"ReflectiveManatee\": {\n \"labs\": [\"ReflectiveManatee\"],\n \"frames\": 303023,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"lateral_left\", \"lateral_right\",\n \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"attack\", \"sniff\"],\n \"remove_video_id\": [],\n },\n \"CRIM13\": {\n \"labs\": [\"CRIM13\"],\n \"frames\": 208605,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"approach\", \"attack\", \"disengage\", \"mount\", \"rear\", \"selfgroom\", \"sniff\"],\n \"remove_video_id\": [],\n },\n \"CautiousGiraffe\": {\n \"labs\": [\"CautiousGiraffe\"],\n \"frames\": 150000,\n \"body_parts\": [\n \"ear_left\", \"ear_right\", \"hip_left\", \"hip_right\",\n \"neck\", \"nose\", \"tail_base\", \"body_center\"\n ],\n \"actions\": [\"approach\", \"chase\", \"escape\", \"reciprocalsniff\", \"sniff\", \"sniffbody\", \"sniffgenital\"],\n \"remove_video_id\": [],\n },\n}\n","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data loading and preprocessing","metadata":{"papermill":{"duration":0.00762,"end_time":"2025-10-25T19:26:58.057724","exception":false,"start_time":"2025-10-25T19:26:58.050104","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# ==================== DATA LOADING ====================\nprune_df = pd.read_csv('/kaggle/input/mabe-pruning/prunable_by_action.csv')\n\ntrain = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\n\n# drop likely-sleeping MABe22 clips: condition == \"lights on\"\ntrain = train.loc[~(train['lab_id'].astype(str).str.contains('MABe22', na=False) &\n train['mouse1_condition'].astype(str).str.lower().eq('lights on'))].copy()\n\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('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\ntest['sleeping'] = (\n test['lab_id'].astype(str).str.contains('MABe22', na=False) &\n test['mouse1_condition'].astype(str).str.lower().eq('lights on')\n)\ntest['n_mice'] = 4 - test[['mouse1_strain','mouse2_strain','mouse3_strain','mouse4_strain']].isna().sum(axis=1)\n\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\ndrop_body_parts = [\n 'headpiece_bottombackleft', 'headpiece_bottombackright',\n 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright',\n 'headpiece_topbackleft', 'headpiece_topbackright',\n 'headpiece_topfrontleft', 'headpiece_topfrontright',\n 'spine_1', 'spine_2', 'tail_middle_1', 'tail_middle_2', 'tail_midpoint'\n]\n\n# ---- context-only allowed bodyparts (renamed) ----\nALLOWED_BPARTS_CTX = {\n 'lateral_left', 'lateral_right', 'neck', 'nose',\n 'ear_left', 'ear_right', 'tail_base', 'hip_left', 'hip_right'\n}\n\n_sex_cols = [f'mouse{i}_sex' for i in range(1,5)]\n_train_sex_lut = (train[['video_id'] + _sex_cols].drop_duplicates('video_id')\n .set_index('video_id').to_dict('index'))\n_test_sex_lut = (test[['video_id'] + _sex_cols].drop_duplicates('video_id')\n .set_index('video_id').to_dict('index'))\n_FEATURE_TEMPLATES = {}\n\n# ==================== GENERATOR (main data pruned by drop_body_parts; context pruned by ALLOWED_BPARTS_CTX) ====================\ndef generate_mouse_data(dataset, traintest, traintest_directory=None,\n generate_single=True, generate_pair=True):\n \"\"\"\n Yields a 5-tuple:\n (switch, data_df, context_df, meta_df, y_or_actions)\n\n - switch: 'single' or 'pair'\n - data_df:\n * single: focal mouse XY (columns (bodypart, coord)), pruned by drop_body_parts\n * pair: two mice under top level ['A','B'], pruned by drop_body_parts\n - context_df: XY for ALL other mice (not in focal), pruned by ALLOWED_BPARTS_CTX.\n Columns MultiIndex (mouse_id, bodypart, coord) where coord in {'x','y'}.\n - meta_df: per-row metadata\n - y_or_actions:\n * train: DataFrame of booleans for actions\n * test: list of actions\n \"\"\"\n assert traintest in ['train', 'test']\n if traintest_directory is None:\n traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n\n def _to_num(x):\n if isinstance(x, (int, np.integer)): return int(x)\n m = re.search(r'(\\d+)$', str(x))\n return int(m.group(1)) if m else None\n\n for _, row in dataset.iterrows():\n lab_id = row.lab_id\n lab_id_s = str(lab_id)\n video_id = row.video_id\n tracker = row.tracking_method\n fps = float(row.frames_per_second)\n n_mice = int(row.n_mice)\n arena_w = float(row.get('arena_width_cm', np.nan))\n arena_h = float(row.get('arena_height_cm', np.nan))\n sleeping = bool(getattr(row, 'sleeping', False))\n arena_shape = row.get('arena_shape', 'rectangular')\n\n if not isinstance(row.behaviors_labeled, str):\n continue\n\n # ---- tracking ----\n path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n vid = pd.read_parquet(path)\n\n # MAIN DATA pruning: mirror baseline behavior with drop_body_parts (only when enough bodyparts present)\n if len(np.unique(vid.bodypart)) > 5:\n vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n\n # pivot to columns = (mouse_id, bodypart, coord) with coord in {'x','y'}\n pvid = vid.pivot(columns=['mouse_id','bodypart'], index='video_frame', values=['x','y'])\n del vid\n pvid = pvid.reorder_levels([1,2,0], axis=1).T.sort_index().T\n pvid.columns = pvid.columns.set_names(['mouse_id','bodypart','coord'])\n # scale to cm; cast\n pvid = (pvid / float(row.pix_per_cm_approx)).astype('float32', copy=False)\n\n # available mouse_id labels in tracking (ints or strings)\n avail = list(pvid.columns.get_level_values('mouse_id').unique())\n avail_set = set(avail) | set(map(str, avail)) | {f\"mouse{_to_num(a)}\" for a in avail if _to_num(a) is not None}\n\n def _resolve(agent_str):\n m = re.search(r'(\\d+)$', str(agent_str))\n cand = [agent_str]\n if m:\n n = int(m.group(1))\n cand = [n, n-1, str(n), f\"mouse{n}\", agent_str]\n for c in cand:\n if c in avail_set:\n if c in set(avail): return c\n for a in avail:\n if str(a) == str(c) or f\"mouse{_to_num(a)}\" == str(c):\n return a\n return None\n\n # ---- behaviors ----\n vb = json.loads(row.behaviors_labeled)\n vb = sorted(list({b.replace(\"'\", \"\") for b in vb}))\n vb = pd.DataFrame([b.split(',') for b in vb], columns=['agent','target','action'])\n vb['agent'] = vb['agent'].astype(str)\n vb['target'] = vb['target'].astype(str)\n vb['action'] = vb['action'].astype(str).str.lower()\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 def _mk_meta(index, agent_id, target_id):\n m = pd.DataFrame({\n 'lab_id': lab_id,\n 'video_id': video_id,\n 'agent_id': agent_id,\n 'target_id': target_id,\n 'video_frame': index.astype('int32', copy=False),\n 'frames_per_second': np.float32(fps),\n 'sleeping': sleeping,\n 'tracker': tracker,\n 'arena_shape': arena_shape,\n 'arena_width_cm': np.float32(arena_w),\n 'arena_height_cm': np.float32(arena_h),\n 'n_mice': np.int8(n_mice),\n })\n for c in ('lab_id','video_id','agent_id','target_id','arena_shape'):\n m[c] = m[c].astype('category')\n return m\n\n # helper to build context: all mice except exclude_ids, then prune by ALLOWED_BPARTS_CTX\n def _context_df(exclude_ids):\n if not isinstance(exclude_ids, (list, tuple, set)):\n exclude_ids = [exclude_ids]\n exclude_ids = set(exclude_ids)\n mask_other = ~pvid.columns.get_level_values('mouse_id').isin(exclude_ids)\n ctx = pvid.loc[:, mask_other]\n if ctx.shape[1] == 0:\n # Keep proper MultiIndex even if empty\n empty_cols = pd.MultiIndex.from_tuples([], names=['mouse_id','bodypart','coord'])\n return pd.DataFrame(index=pvid.index, columns=empty_cols, dtype=np.float32)\n # Context-only pruning\n bmask_ctx = ctx.columns.get_level_values('bodypart').isin(ALLOWED_BPARTS_CTX)\n ctx = ctx.loc[:, bmask_ctx]\n return ctx\n\n # ---------- SINGLE ----------\n if generate_single:\n vb_single = vb.query(\"target == 'self'\")\n for agent_str in pd.unique(vb_single['agent']):\n col_lab = _resolve(agent_str)\n if col_lab is None:\n continue\n actions = sorted(vb_single.loc[vb_single['agent'].eq(agent_str), 'action'].unique().tolist())\n if not actions:\n continue\n\n single = pvid.loc[:, col_lab] # (bodypart, coord); already pruned by drop_body_parts upstream\n meta_df = _mk_meta(single.index, agent_str, 'self')\n\n # context: all other mice (context-only pruning applied)\n ctx = _context_df(exclude_ids=[col_lab])\n\n if traintest == 'train':\n a_num = _to_num(col_lab)\n y = pd.DataFrame(False, index=single.index.astype('int32', copy=False), columns=actions)\n a_sub = annot.query(\"(agent_id == @a_num) & (target_id == @a_num)\")\n for i in range(len(a_sub)):\n ar = a_sub.iloc[i]\n a = str(ar.action).lower()\n if a in y.columns:\n y.loc[int(ar['start_frame']):int(ar['stop_frame']), a] = True\n yield 'single', single, ctx, meta_df, y\n else:\n yield 'single', single, ctx, meta_df, list(actions)\n\n # ---------- PAIR (ONLY LABELED PAIRS) ----------\n if generate_pair:\n vb_pair = vb.query(\"target != 'self'\")\n if len(vb_pair) > 0:\n allowed_pairs = set(map(tuple, vb_pair[['agent','target']].itertuples(index=False, name=None)))\n\n for agent_num, target_num in itertools.permutations(\n np.unique(pvid.columns.get_level_values('mouse_id')), 2):\n agent_str = f\"mouse{_to_num(agent_num)}\"\n target_str = f\"mouse{_to_num(target_num)}\"\n if (agent_str, target_str) not in allowed_pairs:\n continue\n\n a_col = _resolve(agent_str)\n b_col = _resolve(target_str)\n if a_col is None or b_col is None:\n continue\n\n actions = sorted(\n vb_pair.query(\"(agent == @agent_str) & (target == @target_str)\")['action'].unique().tolist()\n )\n if not actions:\n continue\n\n # focal pair (top level A/B); main data already pruned by drop_body_parts upstream\n pair_xy = pd.concat([pvid[a_col], pvid[b_col]], axis=1, keys=['A','B'])\n meta_df = _mk_meta(pair_xy.index, agent_str, target_str)\n\n # context: mice other than A and B (context-only pruning applied)\n ctx = _context_df(exclude_ids=[a_col, b_col])\n\n if traintest == 'train':\n a_num = _to_num(a_col); b_num = _to_num(b_col)\n y = pd.DataFrame(False, index=pair_xy.index.astype('int32', copy=False), columns=actions)\n a_sub = annot.query(\"(agent_id == @a_num) & (target_id == @b_num)\")\n for i in range(len(a_sub)):\n ar = a_sub.iloc[i]\n a = str(ar.action).lower()\n if a in y.columns:\n y.loc[int(ar['start_frame']):int(ar['stop_frame']), a] = True\n yield 'pair', pair_xy, ctx, meta_df, y\n else:\n yield 'pair', pair_xy, ctx, meta_df, list(actions)\n","metadata":{"jupyter":{"source_hidden":true},"papermill":{"duration":0.028134,"end_time":"2025-10-25T19:26:58.378526","exception":false,"start_time":"2025-10-25T19:26:58.350392","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def create_solution_df(dataset):\n solution = []\n for _, row in tqdm(dataset.iterrows(), total=len(dataset)):\n \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 \n return solution\n\nif CFG.mode == 'validate':\n solution = create_solution_df(train_without_mabe22)","metadata":{"jupyter":{"source_hidden":true},"papermill":{"duration":0.018377,"end_time":"2025-10-25T19:26:58.330166","exception":false,"start_time":"2025-10-25T19:26:58.311789","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Transforming coordinates","metadata":{"papermill":{"duration":0.006144,"end_time":"2025-10-25T19:26:58.391332","exception":false,"start_time":"2025-10-25T19:26:58.385188","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# ==================== ADVANCED FEATURE ENGINEERING (FPS-AWARE) ====================\ndef _safe_add(X, fn, *args, **kwargs):\n \"\"\"Call feature-adder `fn` safely; if a KeyError (missing body part) occurs, skip it.\"\"\"\n try:\n out = fn(X, *args, **kwargs)\n return out if out is not None else X\n except KeyError as e:\n if globals().get(\"verbose\", False):\n print(f\"[transform_single] Skipping {getattr(fn, '__name__', 'fn')} — missing part {e}\")\n return X\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# ---------- small utils ----------\ndef _roll_future_mean(s: pd.Series, w: int, min_p: int = 1) -> pd.Series:\n # mean over [t, t+w-1]\n return s.iloc[::-1].rolling(w, min_periods=min_p).mean().iloc[::-1]\n\ndef _roll_future_var(s: pd.Series, w: int, min_p: int = 2) -> pd.Series:\n # var over [t, t+w-1]\n return s.iloc[::-1].rolling(w, min_periods=min_p).var().iloc[::-1]\n\ndef _speed(cx: pd.Series, cy: pd.Series, fps: float) -> pd.Series:\n return np.hypot(cx.diff(), cy.diff()).fillna(0.0) * float(fps)\n\ndef safe_rolling(series, window, func, min_periods=None):\n \"\"\"Safe rolling operation with NaN handling\"\"\"\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\ndef _scale(n_frames_at_30fps, fps, ref=30.0):\n \"\"\"Scale a frame count defined at 30 fps to the current video's fps.\"\"\"\n return max(1, int(round(n_frames_at_30fps * float(fps) / ref)))\n\ndef _scale_signed(n_frames_at_30fps, fps, ref=30.0):\n \"\"\"Signed version of _scale for forward/backward shifts (keeps at least 1 frame when |n|>=1).\"\"\"\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\ndef add_curvature_features(X, center_x, center_y, fps):\n \"\"\"Trajectory curvature (window lengths scaled by 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) # invariant to time scaling\n\n for w in [30, 60]:\n ws = _scale(w, fps)\n X[f'curv_mean_{w}'] = curvature.rolling(ws, min_periods=max(1, ws // 6)).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 // 6)).sum()\n\n return X\n\n\n\n# ---- arena helpers (NO barrier handling) ------------------------------------------\ndef _norm_arena_shape(s: object) -> str:\n s = \"\" if s is None else str(s).strip().lower()\n if 'circle' in s: return 'circular'\n if 'square' in s: return 'square'\n if 'rect' in s: return 'rectangular' # includes \"split rectangular\"\n return 'rectangular'\n\ndef _extract_arena_meta(meta_df):\n \"\"\"Return (shape, width_cm, height_cm) from meta (row-aligned DataFrame).\"\"\"\n if meta_df is None or len(meta_df) == 0:\n return 'rectangular', np.nan, np.nan\n shape = _norm_arena_shape(meta_df.get('arena_shape', pd.Series(['rectangular'])).iloc[0])\n w = float(meta_df.get('arena_width_cm', pd.Series([np.nan])).iloc[0])\n h = float(meta_df.get('arena_height_cm', pd.Series([np.nan])).iloc[0])\n return shape, w, h\n\ndef _norm_xy_series(x_cm: pd.Series, y_cm: pd.Series, w_cm: float, h_cm: float):\n \"\"\"Normalize cm coords to [0,1] using arena size; min-max fallback if dims missing.\"\"\"\n if not (np.isfinite(w_cm) and w_cm > 0 and np.isfinite(h_cm) and h_cm > 0):\n xn = (x_cm - x_cm.min()) / (x_cm.max() - x_cm.min() + 1e-9)\n yn = (y_cm - y_cm.min()) / (y_cm.max() - y_cm.min() + 1e-9)\n else:\n xn = x_cm / float(w_cm)\n yn = y_cm / float(h_cm)\n return xn.clip(0, 1), yn.clip(0, 1)\n\ndef distance_to_wall_norm(xn, yn, arena_shape: str):\n \"\"\"\n Normalized distance to nearest boundary in [0, 0.5].\n - rectangular/square: min distance to 4 walls in normalized coords\n - circular: distance to circle edge (radius=0.5, center=(0.5,0.5))\n \"\"\"\n shp = (arena_shape or '').strip().lower()\n x = np.asarray(xn, dtype=float)\n y = np.asarray(yn, dtype=float)\n\n if shp == 'circular':\n rho = np.sqrt((x - 0.5)**2 + (y - 0.5)**2)\n return np.maximum(0.0, 0.5 - rho)\n\n # rectangular family\n return np.minimum.reduce([x, 1.0 - x, y, 1.0 - y])\n\ndef _heading_unit(df, have_parts) -> np.ndarray:\n \"\"\"Unit heading vector from tail_base -> nose; NaNs handled.\"\"\"\n if all(p in have_parts for p in ['nose', 'tail_base']):\n v = df['nose'] - df['tail_base'] # columns ['x','y']\n vx = np.asarray(v['x'], dtype=float)\n vy = np.asarray(v['y'], dtype=float)\n n = np.sqrt(vx*vx + vy*vy) + 1e-9\n return np.c_[vx/n, vy/n]\n # fallback: zero heading\n z = np.zeros((len(df), 2), dtype=float)\n return z\n\ndef _wall_normal_and_tangent(xn, yn, arena_shape: str):\n \"\"\"\n Inward unit normal and unit tangent at nearest boundary (per frame).\n - rectangular family: choose the nearest side\n - circular: normal points to center; tangent is perpendicular\n Returns (nvec, tvec) with shape (N,2) each.\n \"\"\"\n shp = (arena_shape or '').strip().lower()\n x = np.asarray(xn, dtype=float)\n y = np.asarray(yn, dtype=float)\n N = len(x)\n nvec = np.zeros((N,2), dtype=float)\n tvec = np.zeros((N,2), dtype=float)\n\n if shp == 'circular':\n nx = 0.5 - x\n ny = 0.5 - y\n mag = np.sqrt(nx*nx + ny*ny) + 1e-9\n nvec[:,0] = nx / mag; nvec[:,1] = ny / mag\n tvec[:,0] = -nvec[:,1]; tvec[:,1] = nvec[:,0]\n return nvec, tvec\n\n # rectangular family\n dists = np.stack([x, 1.0-x, y, 1.0-y], axis=1) # left, right, bottom, top\n normals = np.array([[+1,0], [-1,0], [0,+1], [0,-1]], float) # inward normals\n idx = np.argmin(dists, axis=1)\n nvec = normals[idx] # (N,2)\n tvec[:,0] = -nvec[:,1]; tvec[:,1] = nvec[:,0]\n return nvec, tvec\n\n\n\ndef add_wall_geometry_single(X: pd.DataFrame, single_mouse: pd.DataFrame, meta_df: pd.DataFrame | None):\n if 'body_center' not in single_mouse.columns.get_level_values(0):\n return X\n shape, w, h = _extract_arena_meta(meta_df)\n\n cx = single_mouse['body_center']['x']; cy = single_mouse['body_center']['y']\n xn, yn = _norm_xy_series(cx, cy, w, h)\n xna = np.asarray(xn, dtype=float); yna = np.asarray(yn, dtype=float)\n\n X['wall_dist'] = distance_to_wall_norm(xna, yna, shape)\n\n # orientation vs wall (parallel/toward)\n parts = single_mouse.columns.get_level_values(0)\n head = _heading_unit(single_mouse, parts) # (N,2)\n nvec, tvec = _wall_normal_and_tangent(xna, yna, shape)\n X['toward_wall'] = (head * nvec).sum(axis=1) # +1 = heading into wall; -1 = away\n X['along_wall'] = np.abs((head * tvec).sum(axis=1)) # 1 = parallel to wall\n return X\n\ndef add_wall_geometry_pair(X: pd.DataFrame, mouse_pair: pd.DataFrame, meta_df: pd.DataFrame | None):\n if 'body_center' not in mouse_pair['A'].columns.get_level_values(0) or \\\n 'body_center' not in mouse_pair['B'].columns.get_level_values(0):\n return X\n shape, w, h = _extract_arena_meta(meta_df)\n\n Ax = mouse_pair['A']['body_center']['x']; Ay = mouse_pair['A']['body_center']['y']\n Bx = mouse_pair['B']['body_center']['x']; By = mouse_pair['B']['body_center']['y']\n Axn, Ayn = _norm_xy_series(Ax, Ay, w, h)\n Bxn, Byn = _norm_xy_series(Bx, By, w, h)\n Axna = np.asarray(Axn, dtype=float); Ayna = np.asarray(Ayn, dtype=float)\n Bxna = np.asarray(Bxn, dtype=float); Byna = np.asarray(Byn, dtype=float)\n\n X['wall_A'] = distance_to_wall_norm(Axna, Ayna, shape)\n X['wall_B'] = distance_to_wall_norm(Bxna, Byna, shape)\n\n partsA = mouse_pair['A'].columns.get_level_values(0)\n partsB = mouse_pair['B'].columns.get_level_values(0)\n hA = _heading_unit(mouse_pair['A'], partsA)\n hB = _heading_unit(mouse_pair['B'], partsB)\n nA, tA = _wall_normal_and_tangent(Axna, Ayna, shape)\n nB, tB = _wall_normal_and_tangent(Bxna, Byna, shape)\n X['toward_wall_A'] = (hA * nA).sum(axis=1)\n X['toward_wall_B'] = (hB * nB).sum(axis=1)\n X['along_wall_A'] = np.abs((hA * tA).sum(axis=1))\n X['along_wall_B'] = np.abs((hB * tB).sum(axis=1))\n return X\n\ndef _u_wall_rect(xn, yn):\n x = np.asarray(xn, dtype=float)\n y = np.asarray(yn, dtype=float)\n dL, dR, dB, dT = x, (1.0 - x), y, (1.0 - y)\n mL = (dL <= dR) & (dL <= dB) & (dL <= dT)\n mR = (dR < dL) & (dR <= dB) & (dR <= dT)\n mB = (dB < dL) & (dB < dR) & (dB <= dT)\n mT = ~(mL | mR | mB)\n ux = np.zeros_like(dL); uy = np.zeros_like(dL)\n ux[mL], uy[mL] = -1.0, 0.0\n ux[mR], uy[mR] = 1.0, 0.0\n ux[mB], uy[mB] = 0.0, -1.0\n ux[mT], uy[mT] = 0.0, 1.0\n return ux, uy\n\ndef _u_wall_circ(xn, yn):\n x = np.asarray(xn, dtype=float)\n y = np.asarray(yn, dtype=float)\n rx = (x - 0.5)\n ry = (y - 0.5)\n rn = np.sqrt(rx*rx + ry*ry) + 1e-9\n return rx / rn, ry / rn # radial (outward) = wall normal\n\ndef distance_to_corner_norm(xn, yn, arena_shape):\n \"\"\"\n Nearest-corner distance in normalized coords.\n - For rectangular/square arenas (x,y in [0,1]): min distance to any corner.\n - For circular/other shapes: returns NaN (let your SimpleImputer handle it).\n \"\"\"\n s = str(arena_shape).strip().lower()\n x = np.asarray(xn, dtype=float)\n y = np.asarray(yn, dtype=float)\n if 'circ' in s:\n return np.full_like(x, np.nan, dtype=float)\n d00 = np.hypot(x - 0.0, y - 0.0)\n d10 = np.hypot(x - 1.0, y - 0.0)\n d01 = np.hypot(x - 0.0, y - 1.0)\n d11 = np.hypot(x - 1.0, y - 1.0)\n return np.minimum.reduce([d00, d10, d01, d11])\n\n\n\n\ndef add_multiscale_features(X, center_x, center_y, fps):\n \"\"\"Multi-scale temporal features (speed in cm/s; windows scaled by fps).\"\"\"\n speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)\n scales = [10, 40, 160]\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 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 return X\n\ndef add_state_features(X, center_x, center_y, fps):\n \"\"\"Behavioral state transitions; bins adjusted so semantics are fps-invariant.\"\"\"\n speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps) # cm/s\n w_ma = _scale(15, fps)\n speed_ma = speed.rolling(w_ma, min_periods=max(1, w_ma // 3)).mean()\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 for window in [60, 120]:\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 // 6)).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 // 6)).sum()\n except Exception:\n pass\n return X\n\ndef add_state_features(X, center_x, center_y, fps):\n \"\"\"\n Behavioral state transitions; bins adjusted so semantics are fps-invariant.\n\n Slimmed version, keeping only:\n - s0_60: fraction of time in very-low-speed state over ~2s\n - s1_60: fraction of time in low/medium-speed state over ~2s\n - trans_60: number of state changes over ~2s\n\n Rare high-speed bins (states 2 & 3) and longer 120-frame horizon are dropped\n as they showed negligible importance.\n \"\"\"\n speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps) # cm/s\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 # same binning as before, just fewer outputs\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 window = 60 # ~2s @30fps\n ws = _scale(window, fps)\n if len(speed_states) >= ws:\n # keep only low-speed states 0 and 1\n for state in [0, 1]:\n X[f's{state}_{window}'] = (\n (speed_states == state)\n .astype(float)\n .rolling(ws, min_periods=max(1, ws // 6))\n .mean()\n )\n\n # transitions over same window\n state_changes = (speed_states != speed_states.shift(1)).astype(float)\n X[f'trans_{window}'] = (\n state_changes.rolling(ws, min_periods=max(1, ws // 6)).sum()\n )\n\n except Exception:\n # fail-safe: just return X unchanged if anything goes sideways\n pass\n\n return X\n\ndef add_longrange_features(X, center_x, center_y, fps):\n \"\"\"Long-range temporal features (windows & spans scaled by fps).\"\"\"\n for window in [120, 240]:\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 for span in [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 speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps) # cm/s\n for window in [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 return X\n\ndef add_contact_bins(X, pair, fps, prox_cm=(3.0, 5.0, 10.0)):\n def _part(df, name): return (df[name]['x'], df[name]['y']) if name in df.columns.get_level_values(0) else (None,None)\n Ax, Ay = _part(pair['A'], 'nose'); Bx, By = _part(pair['B'], 'nose')\n if Ax is None or Bx is None: return X\n d = np.sqrt((Ax-Bx)**2 + (Ay-By)**2)\n if all(n in pair['A'].columns.get_level_values(0) for n in ['nose','tail_base']) and \\\n all(n in pair['B'].columns.get_level_values(0) for n in ['nose','tail_base']):\n da = pair['A']['nose']-pair['A']['tail_base']; db = pair['B']['nose']-pair['B']['tail_base']\n cos_face = ((da['x']*db['x'] + da['y']*db['y']) /\n (np.sqrt(da['x']**2+da['y']**2)*np.sqrt(db['x']**2+db['y']**2) + 1e-6))\n else:\n cos_face = pd.Series(0.0, index=d.index)\n for thr in prox_cm:\n near = (d < thr) & (cos_face > 0) # facing each other\n for w in [30, 60]:\n ws = _scale(w, fps)\n X[f'near{int(thr)}_face{w}'] = near.rolling(ws, min_periods=max(1,ws//6)).mean()\n return X\n\ndef add_xcorr_lag(X, pair, fps, max_sec=1.0):\n if 'body_center' not in pair['A'].columns.get_level_values(0) or 'body_center' not in pair['B'].columns.get_level_values(0):\n return X\n Av = pair['A']['body_center'].diff()\n Bv = pair['B']['body_center'].diff()\n Avn = np.sqrt(Av['x']**2 + Av['y']**2)\n Bvn = np.sqrt(Bv['x']**2 + Bv['y']**2)\n maxlag = _scale(int(max_sec*30), fps)\n if maxlag <= 0:\n X['lead_lag_s'] = 0.0\n return X\n def _bestlag(a, b):\n lags = range(-maxlag, maxlag + 1)\n vals = np.array([pd.Series(a).corr(pd.Series(b).shift(l)) for l in lags], dtype=float)\n if np.all(np.isnan(vals)):\n return 0\n return list(lags)[int(np.nanargmax(vals))]\n bl = _bestlag(Avn.fillna(0), Bvn.fillna(0))\n X['lead_lag_s'] = bl / float(fps) # >0 => B lags A\n return X\n\ndef add_interaction_features(X, mouse_pair, avail_A, avail_B, fps):\n \"\"\"Social interaction features (windows scaled by fps).\"\"\"\n if 'body_center' not in avail_A or 'body_center' not in avail_B:\n return X\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 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 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 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 approach = -rel_dist.diff() # decreasing distance => positive approach\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 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 return X\n\nfrom scipy import signal\n\ndef _sec_to_frames(sec, fps):\n # at least 1 frame, keep as int\n return max(1, int(round(float(sec) * float(fps))))\n\ndef _rolling_safe(x, window_frames, min_frac=0.2):\n return x.rolling(window_frames, min_periods=max(1, int(round(window_frames*min_frac))))\n\n# ========================================================\n# FPS-AWARE FEATURES\n# ========================================================\n\ndef add_frequency_features(X, center_x, center_y, fps: float = 30.0,\n seg_sec: float = 4.0, # Welch window length (s)\n step_sec: float = 2.0): # step ~ 50% overlap (s)\n \"\"\"\n Frequency-domain features from speed signal. Uses Welch with fs=fps so\n 'dominant_freq_hz' is in Hz and robust to variable frame rates.\n Adds:\n - dominant_freq_hz (scalar column, broadcast)\n - spectral_entropy (scalar column, broadcast)\n - rolling_psd_mean (series, length = len(X))\n \"\"\"\n # speed in pixel/second\n dx = center_x.diff()\n dy = center_y.diff()\n speed_ps = np.sqrt(dx*dx + dy*dy).fillna(0) * float(fps)\n\n n = len(speed_ps)\n if n >= 16:\n nperseg = min(_sec_to_frames(seg_sec, fps), n)\n noverlap = min(nperseg//2, nperseg-1) # avoid equal to nperseg\n f, psd = signal.welch(speed_ps.values, fs=float(fps),\n nperseg=nperseg, noverlap=noverlap)\n if len(psd) and np.isfinite(psd).any():\n dom_hz = float(f[int(np.nanargmax(psd))])\n p = np.maximum(psd, 0)\n p = p / (p.sum() + 1e-12)\n spec_ent = float(-np.sum(p * np.log2(p + 1e-12)) / np.log2(len(p)))\n else:\n dom_hz, spec_ent = 0.0, 0.0\n else:\n dom_hz, spec_ent = 0.0, 0.0\n\n X['dominant_freq_hz'] = dom_hz\n X['spectral_entropy'] = spec_ent\n\n # Rolling PSD mean (very light-weight proxy) over sliding windows in seconds\n if n >= 2 * _sec_to_frames(seg_sec, fps):\n step = _sec_to_frames(step_sec, fps)\n win = _sec_to_frames(seg_sec, fps)\n means, idxs = [], []\n arr = speed_ps.values\n for start in range(0, n - win + 1, step):\n seg = arr[start:start+win]\n f_i, psd_i = signal.welch(seg, fs=float(fps),\n nperseg=min(win, len(seg)),\n noverlap=min(win//2, max(0, len(seg)-1)))\n means.append(float(np.nanmean(psd_i)) if len(psd_i) else 0.0)\n idxs.append(start + win//2)\n # interpolate centers back to full length\n X['rolling_psd_mean'] = pd.Series(\n np.interp(np.arange(n), idxs, means, left=means[0], right=means[-1]),\n index=X.index\n )\n else:\n X['rolling_psd_mean'] = 0.0\n\n return X\n\n\ndef add_statistical_moments(X, center_x, center_y, fps: float = 30.0,\n windows_sec=(1.0, 2.0)):\n \"\"\"\n Higher-order moments of speed (pixel/s) in time windows measured in seconds.\n Adds per window:\n - sp_skew_{ws}\n - sp_kurt_{ws}\n \"\"\"\n speed_ps = np.sqrt(center_x.diff()**2 + center_y.diff()**2).fillna(0) * float(fps)\n\n for ws in windows_sec:\n w = _sec_to_frames(ws, fps)\n X[f'sp_skew_{int(ws)}s'] = _rolling_safe(speed_ps, w).skew()\n X[f'sp_kurt_{int(ws)}s'] = _rolling_safe(speed_ps, w).kurt()\n\n return X.fillna(0)\n\n\ndef add_acceleration_features(X, center_x, center_y, fps: float = 30.0,\n windows_sec=(0.5, 1.0, 2.0)):\n \"\"\"\n Velocity/acceleration in physical units:\n vel = Δpos * fps -> pixel/s\n acc = Δvel * fps -> pixel/s^2\n jerk = Δacc -> pixel/s^2 per frame (proxy)\n Adds per window:\n - acc_mean_{ws}s / acc_std_{ws}s / acc_max_{ws}s\n Also:\n - jerk_mean_1s\n \"\"\"\n fps = float(fps)\n vx = center_x.diff() * fps\n vy = center_y.diff() * fps\n ax = vx.diff() * fps\n ay = vy.diff() * fps\n acc = np.sqrt(ax**2 + ay**2).fillna(0)\n\n for ws in windows_sec:\n w = _sec_to_frames(ws, fps)\n roll = _rolling_safe(acc, w, min_frac=0.25)\n X[f'acc_mean_{int(ws)}s'] = roll.mean()\n X[f'acc_std_{int(ws)}s'] = roll.std()\n X[f'acc_max_{int(ws)}s'] = roll.max()\n\n # jerk over ~1s\n wj = _sec_to_frames(1.0, fps)\n jerk = acc.diff().abs()\n X['jerk_mean_1s'] = _rolling_safe(jerk, wj, min_frac=0.25).mean()\n\n return X.fillna(0)\n\n\ndef add_relative_acceleration_features(X, mouse_pair, avail_A, avail_B,\n fps: float = 30.0,\n windows_sec=(1.0, 2.0)):\n \"\"\"\n Relative acceleration magnitude between A and B (pixel/s^2), with\n windows measured in seconds.\n Adds per window:\n - rel_acc_mean_{ws}s / rel_acc_std_{ws}s\n \"\"\"\n if 'body_center' not in avail_A or 'body_center' not in avail_B:\n return X\n\n fps = float(fps)\n A_vx = mouse_pair['A']['body_center']['x'].diff() * fps\n A_vy = mouse_pair['A']['body_center']['y'].diff() * fps\n B_vx = mouse_pair['B']['body_center']['x'].diff() * fps\n B_vy = mouse_pair['B']['body_center']['y'].diff() * fps\n\n A_ax = A_vx.diff() * fps\n A_ay = A_vy.diff() * fps\n B_ax = B_vx.diff() * fps\n B_ay = B_vy.diff() * fps\n\n rel_acc_mag = np.sqrt((A_ax - B_ax)**2 + (A_ay - B_ay)**2).fillna(0)\n\n for ws in windows_sec:\n w = _sec_to_frames(ws, fps)\n roll = _rolling_safe(rel_acc_mag, w, min_frac=0.25)\n X[f'rel_acc_mean_{int(ws)}s'] = roll.mean()\n X[f'rel_acc_std_{int(ws)}s'] = roll.std()\n\n return X.fillna(0)\n\ndef add_single_stillness_and_headjitter(X, df, fps):\n parts = df.columns.get_level_values(0)\n if 'body_center' not in parts:\n return X\n\n cx = df['body_center']['x']; cy = df['body_center']['y']\n speed = (np.sqrt(cx.diff()**2 + cy.diff()**2) * float(fps)).fillna(0)\n\n # 1) stillness ratio in short/med windows (fps-aware)\n thr_cm_s = 0.3 * (fps / 30.0)\n for w in [15, 45]: # ~0.5s, 1.5s @30fps\n ws = _scale(w, fps)\n still = (speed < thr_cm_s).astype(float)\n X[f'still_frac_{w}'] = still.rolling(ws, min_periods=max(1, ws//3)).mean()\n\n # 2) head micro-jitter when center is still\n if 'nose' in parts:\n nx = df['nose']['x']; ny = df['nose']['y']\n nspd = (np.sqrt(nx.diff()**2 + ny.diff()**2) * float(fps)).fillna(0)\n win = _scale(30, fps)\n head_var = nspd.rolling(max(3, win//3), min_periods=1).var()\n still_mask = (speed < thr_cm_s).astype(float).rolling(win, min_periods=1).mean()\n X['head_jitter_ps'] = (head_var * still_mask).fillna(0.0)\n\n # 3) posture compactness / elongation (nose–tail)\n if 'nose' in parts and 'tail_base' in parts:\n d = np.sqrt((df['nose']['x']-df['tail_base']['x'])**2 + (df['nose']['y']-df['tail_base']['y'])**2)\n d_med = float(np.nanmedian(d)) + 1e-9\n X['body_elong_norm'] = (d / d_med).fillna(1.0)\n\n return X\n\n\ndef add_groom_microfeatures(X, df, fps):\n parts = df.columns.get_level_values(0)\n if 'body_center' not in parts or 'nose' not in parts:\n return X\n\n cx = df['body_center']['x']; cy = df['body_center']['y']\n nx = df['nose']['x']; ny = df['nose']['y']\n\n cs = (np.sqrt(cx.diff()**2 + cy.diff()**2) * float(fps)).fillna(0)\n ns = (np.sqrt(nx.diff()**2 + ny.diff()**2) * float(fps)).fillna(0)\n\n w30 = _scale(30, fps)\n X['head_body_decouple'] = (ns / (cs + 1e-3)).clip(0, 10).rolling(w30, min_periods=max(1, w30//3)).median()\n\n r = np.sqrt((nx - cx)**2 + (ny - cy)**2)\n X['nose_rad_std'] = r.rolling(w30, min_periods=max(1, w30//3)).std().fillna(0)\n\n if 'tail_base' in parts:\n ang = np.arctan2(df['nose']['y']-df['tail_base']['y'], df['nose']['x']-df['tail_base']['x'])\n dang = np.abs(ang.diff()).fillna(0)\n X['head_orient_jitter'] = dang.rolling(w30, min_periods=max(1, w30//3)).mean()\n\n return X\n\n\ndef add_verticality_proxies_single(X, df, meta_df, fps):\n \"\"\"Rear/Climb proxies using arena geometry from meta_df when available.\"\"\"\n if 'body_center' not in df.columns.get_level_values(0):\n return X\n\n shape, w, h = _extract_arena_meta(meta_df)\n cx = df['body_center']['x']; cy = df['body_center']['y']\n xn, yn = _norm_xy_series(cx, cy, w, h)\n xna = np.asarray(xn, dtype=float); yna = np.asarray(yn, dtype=float)\n\n nvec, tvec = _wall_normal_and_tangent(xna, yna, str(shape))\n wall_d = distance_to_wall_norm(xna, yna, str(shape))\n near = (wall_d <= 0.05).astype(float) # within 5% of arena (normalized)\n\n fpsf = float(fps)\n vx = (cx.diff() * fpsf).fillna(0); vy = (cy.diff() * fpsf).fillna(0)\n sp = np.sqrt(vx**2 + vy**2).fillna(0)\n\n parts = df.columns.get_level_values(0)\n px = df['nose']['x'] if 'nose' in parts else cx\n py = df['nose']['y'] if 'nose' in parts else cy\n nvx = (px.diff() * fpsf).fillna(0); nvy = (py.diff() * fpsf).fillna(0)\n\n head_norm = nvx.values * nvec[:,0] + nvy.values * nvec[:,1]\n head_tangent = np.abs(nvx.values * tvec[:,0] + nvy.values * tvec[:,1])\n\n w30 = _scale(30, fps)\n X['near_wall_5p'] = pd.Series(near, index=cx.index).rolling(w30, min_periods=max(1, w30//3)).mean()\n\n # Rear: near wall, slow COM, positive head motion into wall\n rear_score = near * (sp.values < 2.0 * (fps/30.0)) * (\n pd.Series(head_norm).rolling(w30, min_periods=1).mean().clip(lower=0)\n )\n X['rear_proxy'] = pd.Series(rear_score, index=cx.index).fillna(0)\n\n # Climb: near wall, strong tangent oscillation\n climb_score = near * (pd.Series(head_tangent).rolling(w30, min_periods=1).mean())\n X['climb_proxy'] = pd.Series(climb_score, index=cx.index).fillna(0)\n\n return X\n\n\ndef add_dig_proxies_single(X, df, meta_df, fps):\n \"\"\"Digging proxy: corner proximity + normal-axis oscillation energy.\"\"\"\n if 'body_center' not in df.columns.get_level_values(0):\n return X\n\n shape, w, h = _extract_arena_meta(meta_df)\n cx = df['body_center']['x']; cy = df['body_center']['y']\n xn, yn = _norm_xy_series(cx, cy, w, h)\n xna = np.asarray(xn, dtype=float); yna = np.asarray(yn, dtype=float)\n\n cdist = distance_to_corner_norm(xna, yna, str(shape))\n X['near_corner_10p'] = pd.Series((cdist <= 0.10).astype(float), index=cx.index)\n\n parts = df.columns.get_level_values(0)\n px = df['nose']['x'] if 'nose' in parts else cx\n py = df['nose']['y'] if 'nose' in parts else cy\n nvx = (px.diff() * float(fps)).fillna(0); nvy = (py.diff() * float(fps)).fillna(0)\n\n nvec, _ = _wall_normal_and_tangent(xna, yna, str(shape))\n vnorm = nvx.values * nvec[:,0] + nvy.values * nvec[:,1]\n vnorm = pd.Series(vnorm, index=cx.index).fillna(0)\n\n w_short = _scale(20, fps)\n X['dig_norm_energy'] = vnorm.rolling(w_short, min_periods=max(1, w_short//3)) \\\n .apply(lambda a: float(np.mean(a*a)), raw=True).fillna(0)\n return X\n\n\ndef add_directional_entropy(X, center_x, center_y, fps):\n \"\"\"1s directional entropy of motion (low for straight, high for jittery).\"\"\"\n ang = np.arctan2(center_y.diff(), center_x.diff())\n dang = np.abs((ang - ang.shift(1)).fillna(0))\n w = _scale(30, fps)\n bins = pd.cut(dang, bins=[-1e-9, 0.02, 0.08, 0.2, np.inf], labels=[0,1,2,3]).astype('float').fillna(0)\n\n def _ent(vals):\n v = vals.astype(int)\n p = np.bincount(v, minlength=4).astype(float)\n p = p / (p.sum() + 1e-12)\n return float(-(p * (np.log2(p + 1e-12))).sum() / np.log2(4))\n\n X['dir_entropy_1s'] = bins.rolling(w, min_periods=max(1, w//3)).apply(_ent, raw=False).fillna(0)\n return X\n\ndef add_head_orientation_tree_feats_single(\n X: pd.DataFrame,\n single_mouse: pd.DataFrame,\n fps: float,\n cx: Optional[pd.Series] = None,\n cy: Optional[pd.Series] = None,\n) -> pd.DataFrame:\n \"\"\"\n Add tree-friendly head orientation & micromotion features for single-mouse.\n - Uses cos/sin of angle (no wrap discontinuity).\n - Δθ computed wrap-safely via vector formulation.\n - Produces: head_ang_cos, head_ang_sin, head_ang_vel, head_ang_jitter.\n\n Args\n ----\n X : DataFrame being augmented.\n single_mouse : MultiIndex columns [body_part, (x|y)].\n fps : frames-per-second (float).\n cx, cy : Optional center coordinates (pd.Series). If not provided,\n will use single_mouse['body_center'] if available.\n\n Returns\n -------\n X : The same DataFrame with new columns added (if inputs available).\n \"\"\"\n parts = set(single_mouse.columns.get_level_values(0))\n\n # Need nose + center (either provided or from 'body_center')\n if 'nose' not in parts:\n return X\n if cx is None or cy is None:\n if 'body_center' not in parts:\n return X\n cx = single_mouse['body_center']['x']\n cy = single_mouse['body_center']['y']\n\n # Vector from center -> nose\n hx = single_mouse['nose']['x'] - cx\n hy = single_mouse['nose']['y'] - cy\n\n # Smooth representation of angle\n head_ang = np.arctan2(hy, hx)\n c = np.cos(head_ang)\n s = np.sin(head_ang)\n X['head_ang_cos'] = c\n X['head_ang_sin'] = s\n\n # Wrap-safe delta theta using trig identities\n c_prev = c.shift(1)\n s_prev = s.shift(1)\n sin_d = s * c_prev - c * s_prev\n cos_d = c * c_prev + s * s_prev\n dtheta = np.arctan2(sin_d, cos_d).fillna(0.0) # already wrapped to (-π, π]\n\n # Velocity and short-term jitter\n X['head_ang_vel'] = (dtheta * float(fps)).abs()\n wj = max(3, _scale(30, fps)) # keep consistent with your windowing helper\n X['head_ang_jitter'] = (\n dtheta.rolling(wj, min_periods=max(1, wj // 3)).std().fillna(0.0)\n )\n\n return X\n\ndef add_tortuosity_and_dir_persist_single(\n X: pd.DataFrame,\n cx: pd.Series,\n cy: pd.Series,\n fps: float,\n window_30: int = 30, # 30 @30fps ≈ 1s; scaled via _scale\n) -> pd.DataFrame:\n \"\"\"\n Add path tortuosity (~1s) and directional persistence (~1s) from body_center (cx, cy).\n - Tortuosity = path length / net displacement (bounded).\n - Directional persistence = ||_rolling|| (∈[0,1]).\n \"\"\"\n # window size in frames (scaled for actual fps)\n w = max(3, _scale(window_30, fps))\n\n # --- tortuosity ---\n dx = cx.diff()\n dy = cy.diff()\n step = np.hypot(dx, dy) # per-frame step\n path_len = step.rolling(w, min_periods=1, center=True).sum()\n net_disp = np.hypot(cx - cx.shift(w), cy - cy.shift(w))\n X['tortuosity_1s'] = (path_len / (net_disp + 1e-6)).clip(0, 50)\n\n # --- directional persistence (use cos/sin to avoid angle wrap) ---\n heading = np.arctan2(dy, dx).fillna(method='ffill').fillna(0.0)\n c = np.cos(heading); s = np.sin(heading)\n c_m = c.rolling(w, min_periods=1, center=True).mean()\n s_m = s.rolling(w, min_periods=1, center=True).mean()\n X['dir_persist_1s'] = np.sqrt(c_m * c_m + s_m * s_m).clip(0, 1)\n\n # hygiene\n X.replace([np.inf, -np.inf], np.nan, inplace=True)\n return X\n\ndef add_dir_persist_single(\n X: pd.DataFrame,\n cx: pd.Series,\n cy: pd.Series,\n fps: float,\n window_30: int = 30, # ~1s at 30fps; scaled via _scale\n) -> pd.DataFrame:\n \"\"\"\n Add directional persistence (~1s) from body_center (cx, cy).\n Uses unit displacement vectors (ux, uy) -> rolling mean -> norm in [0,1].\n Avoids angle discontinuities and is tree-friendly.\n \"\"\"\n w = max(3, _scale(window_30, fps))\n dx = cx.diff()\n dy = cy.diff()\n step = np.hypot(dx, dy)\n eps = 1e-6\n ux = (dx / (step + eps)).fillna(0.0)\n uy = (dy / (step + eps)).fillna(0.0)\n\n ux_m = ux.rolling(w, min_periods=1, center=True).mean()\n uy_m = uy.rolling(w, min_periods=1, center=True).mean()\n X['dir_persist_1s'] = np.sqrt(ux_m * ux_m + uy_m * uy_m).clip(0, 1)\n\n return X\n\ndef add_posture_area_features_single(\n X: pd.DataFrame,\n single_mouse: pd.DataFrame,\n available_body_parts,\n clip_max: float | None = None,\n ) -> pd.DataFrame:\n \"\"\"\n Add posture area features using (nose, body_center, tail_base).\n - pose_area: triangle area formed by these three keypoints\n - pose_area_d: absolute first difference of pose_area (motion proxy)\n\n If any required body part is missing, returns X unchanged.\n \"\"\"\n req = {'nose', 'body_center', 'tail_base'}\n if not req.issubset(set(available_body_parts)):\n return X\n\n x1, y1 = single_mouse['nose']['x'], single_mouse['nose']['y']\n x2, y2 = single_mouse['body_center']['x'], single_mouse['body_center']['y']\n x3, y3 = single_mouse['tail_base']['x'], single_mouse['tail_base']['y']\n\n area = 0.5 * np.abs(x1 * (y2 - y3) + x2 * (y3 - y1) + x3 * (y1 - y2))\n if clip_max is not None:\n area = area.clip(0, clip_max)\n\n X['pose_area'] = area\n X['pose_area_d'] = area.diff().abs().fillna(0.0)\n return X\n\ndef add_head_micromove_spectrum_single(\n X: pd.DataFrame,\n single_mouse: pd.DataFrame,\n cx: pd.Series,\n cy: pd.Series,\n fps: float,\n available_body_parts,\n win_seconds: float = 4.0, # Welch window ~4s by default\n min_len: int = 16, # need at least this many samples\n col_name: str = \"head_spec_entropy\",\n ) -> pd.DataFrame:\n \"\"\"\n Add 'head_spec_entropy' — spectral entropy of head micromovements\n (nose relative to body_center). Returns X unchanged if requirements\n are not met or SciPy isn't available.\n \"\"\"\n if 'nose' not in set(available_body_parts):\n return X\n\n nx = single_mouse['nose']['x']\n ny = single_mouse['nose']['y']\n\n # speed of nose relative to center (per-frame), in px/s\n hx_sp = (np.hypot((nx - cx).diff(), (ny - cy).diff()) * float(fps)).fillna(0.0)\n n = len(hx_sp)\n if n < min_len:\n X[col_name] = 0.0\n return X\n\n try:\n nperseg = max(8, min(int(round(win_seconds * float(fps))), n))\n noverlap = min(nperseg // 2, nperseg - 1)\n\n # Welch PSD\n _, psd = signal.welch(\n hx_sp.values, fs=float(fps),\n nperseg=nperseg, noverlap=noverlap\n )\n\n if psd.size and np.isfinite(psd).any():\n p = np.maximum(psd, 0.0)\n p = p / (p.sum() + 1e-12)\n # spectral entropy in [0, 1]\n ent = -np.sum(p * np.log2(p + 1e-12)) / np.log2(len(p))\n X[col_name] = float(np.clip(ent, 0.0, 1.0))\n else:\n X[col_name] = 0.0\n except Exception:\n # SciPy missing or numerical issue — fail safe\n X[col_name] = 0.0\n\n return X\n\ndef add_vertical_offsets_single(\n X: pd.DataFrame,\n single_mouse: pd.DataFrame,\n available_body_parts,\n col_nose_over_ear: str = \"nose_over_ear_y\",\n col_ear_over_tail: str = \"ear_over_tail_y\",\n ) -> pd.DataFrame:\n \"\"\"\n Add vertical offsets using ear midpoint as reference:\n - nose_over_ear_y = nose_y - mean(ear_left_y, ear_right_y)\n - ear_over_tail_y = mean(ear_left_y, ear_right_y) - tail_base_y\n\n No-ops if the required parts are missing.\n \"\"\"\n bset = set(available_body_parts)\n have_ears = {\"ear_left\", \"ear_right\"}.issubset(bset)\n if not have_ears:\n return X\n\n ear_y = 0.5 * (single_mouse[\"ear_left\"][\"y\"] + single_mouse[\"ear_right\"][\"y\"])\n\n if \"nose\" in bset:\n X[col_nose_over_ear] = single_mouse[\"nose\"][\"y\"] - ear_y\n\n if \"tail_base\" in bset:\n X[col_ear_over_tail] = ear_y - single_mouse[\"tail_base\"][\"y\"]\n\n return X\n\ndef fill_or_create_body_center(df: pd.DataFrame) -> pd.DataFrame:\n \"\"\"\n Compute/fill ('body_center','x'/'y') for a pose DataFrame with MultiIndex columns:\n level 0 = body part (e.g., 'neck','lateral_left', ...),\n level 1 = {'x','y'}.\n If ('body_center','x'/'y') are missing, they are created. Existing non-NaN values are preserved.\n\n Heuristic (uses only parts with >=105 annotated videos):\n 1) If BOTH flanks present: body_center := midpoint(lateral_left, lateral_right)\n 2) Else:\n front := neck -> nose -> (ear_left,ear_right) midpoint\n hind := mean(tail_base, hip_mid) if both; else hip_mid; else tail_base\n body_center := midpoint(front, hind)\n\n Any NaNs in a source part are handled frame-wise by falling back to the next rule.\n \"\"\"\n if not isinstance(df.columns, pd.MultiIndex) or df.columns.nlevels < 2:\n raise ValueError(\"Expected MultiIndex columns with levels [bodypart, coord ('x'/'y')].\")\n\n out = df.copy()\n idx = out.index\n\n # Ensure target columns exist (create as NaN if missing)\n if ('body_center','x') not in out.columns:\n out[('body_center','x')] = np.nan\n if ('body_center','y') not in out.columns:\n out[('body_center','y')] = np.nan\n\n # Convenience handles (work off a copy; we’ll fillna back into out at the end)\n bc_x = out[('body_center','x')].copy()\n bc_y = out[('body_center','y')].copy()\n\n # Helper to fetch x/y (returns NaN series if missing)\n def get_xy(part: str):\n x = out[(part,'x')] if (part,'x') in out.columns else pd.Series(np.nan, index=idx)\n y = out[(part,'y')] if (part,'y') in out.columns else pd.Series(np.nan, index=idx)\n return x, y\n\n # Frequently-annotated parts\n llx, lly = get_xy('lateral_left')\n lrx, lry = get_xy('lateral_right')\n nx, ny = get_xy('neck')\n nozx,nozy= get_xy('nose')\n elx, ely = get_xy('ear_left')\n erx, ery = get_xy('ear_right')\n tbx, tby = get_xy('tail_base')\n hlx, hly = get_xy('hip_left')\n hrx, hry = get_xy('hip_right')\n\n # Frames still needing fill (if either coord is missing)\n need_bc = bc_x.isna() | bc_y.isna()\n\n # ---- Rule 1: Flank midpoint ----\n flank_ok = llx.notna() & lly.notna() & lrx.notna() & lry.notna()\n m = need_bc & flank_ok\n if m.any():\n bc_x.loc[m & bc_x.isna()] = ((llx + lrx) / 2.0).loc[m & bc_x.isna()]\n bc_y.loc[m & bc_y.isna()] = ((lly + lry) / 2.0).loc[m & bc_y.isna()]\n need_bc = bc_x.isna() | bc_y.isna()\n\n # ---- If still missing, anchor-based ----\n if need_bc.any():\n # Front: neck -> nose -> ear midpoint\n front_x = pd.Series(np.nan, index=idx)\n front_y = pd.Series(np.nan, index=idx)\n\n neck_ok = nx.notna() & ny.notna()\n use = need_bc & neck_ok\n if use.any():\n front_x.loc[use] = nx.loc[use]; front_y.loc[use] = ny.loc[use]\n\n front_missing = need_bc & (front_x.isna() | front_y.isna())\n nose_ok = nozx.notna() & nozy.notna()\n use = front_missing & nose_ok\n if use.any():\n front_x.loc[use] = nozx.loc[use]; front_y.loc[use] = nozy.loc[use]\n\n front_missing = need_bc & (front_x.isna() | front_y.isna())\n ears_ok = elx.notna() & erx.notna() & ely.notna() & ery.notna()\n use = front_missing & ears_ok\n if use.any():\n front_x.loc[use] = ((elx + erx) / 2.0).loc[use]\n front_y.loc[use] = ((ely + ery) / 2.0).loc[use]\n\n front_ok = front_x.notna() & front_y.notna()\n\n # Hind: mean(tail_base, hip_mid) -> hip_mid -> tail_base\n hip_mid_x = pd.Series(np.nan, index=idx)\n hip_mid_y = pd.Series(np.nan, index=idx)\n hips_ok = hlx.notna() & hly.notna() & hrx.notna() & hry.notna()\n if hips_ok.any():\n hip_mid_x.loc[hips_ok] = (hlx.loc[hips_ok] + hrx.loc[hips_ok]) / 2.0\n hip_mid_y.loc[hips_ok] = (hly.loc[hips_ok] + hry.loc[hips_ok]) / 2.0\n\n hind_x = pd.Series(np.nan, index=idx)\n hind_y = pd.Series(np.nan, index=idx)\n\n both_hind_avail = hips_ok & tbx.notna() & tby.notna()\n use = need_bc & both_hind_avail\n if use.any():\n hind_x.loc[use] = ((hip_mid_x + tbx) / 2.0).loc[use]\n hind_y.loc[use] = ((hip_mid_y + tby) / 2.0).loc[use]\n\n hind_missing = need_bc & (hind_x.isna() | hind_y.isna())\n use = hind_missing & hips_ok\n if use.any():\n hind_x.loc[use] = hip_mid_x.loc[use]; hind_y.loc[use] = hip_mid_y.loc[use]\n\n hind_missing = need_bc & (hind_x.isna() | hind_y.isna())\n tail_only = tbx.notna() & tby.notna()\n use = hind_missing & tail_only\n if use.any():\n hind_x.loc[use] = tbx.loc[use]; hind_y.loc[use] = tby.loc[use]\n\n hind_ok = hind_x.notna() & hind_y.notna()\n\n # Midpoint(front, hind)\n m = need_bc & front_ok & hind_ok\n if m.any():\n cand_x = (front_x + hind_x) / 2.0\n cand_y = (front_y + hind_y) / 2.0\n bc_x.loc[m & bc_x.isna()] = cand_x.loc[m & bc_x.isna()]\n bc_y.loc[m & bc_y.isna()] = cand_y.loc[m & bc_y.isna()]\n\n # ---- Fill existing columns (preserve non-NaNs) ----\n out[('body_center','x')] = out[('body_center','x')].fillna(bc_x)\n out[('body_center','y')] = out[('body_center','y')].fillna(bc_y)\n\n # Optional: keep columns tidy\n out = out.reindex(columns=sorted(out.columns, key=lambda c: (str(c[0]), c[1])))\n return out\n\ndef add_cumulative_distance_single(X, cx, cy, fps, horizon_frames_base: int = 180, colname: str = \"path_cum180\"):\n L = max(1, _scale(horizon_frames_base, fps)) # frames\n # step length (cm per frame since coords are cm)\n step = np.hypot(cx.diff(), cy.diff())\n # centered rolling sum over ~2L+1 frames (acausal)\n path = step.rolling(2*L + 1, min_periods=max(5, L//6), center=True).sum()\n X[colname] = path.fillna(0.0).astype(np.float32)\n return X\n\ndef add_long_distance_activity_single(\n X,\n cx,\n cy,\n fps,\n horizon_frames_base: int = 180, # interpreted at 30fps; scaled via _scale\n):\n \"\"\"\n Acausal (centered) long-horizon features:\n - longdist_dispL : displacement between t-L and t+L [cm]\n - longdist_speedL: mean speed over a centered window [cm/s]\n\n Assumes cx, cy are already in centimeters.\n \"\"\"\n import numpy as np\n\n L = int(max(1, _scale(horizon_frames_base, fps)))\n\n # centered long-distance displacement: ||(x,y)_{t+L} - (x,y)_{t-L}||\n disp = np.hypot(cx.shift(-L) - cx.shift(L), cy.shift(-L) - cy.shift(L))\n\n # centered mean speed over ~2L+1 frames\n inst_speed = np.hypot(cx.diff(), cy.diff()).mul(float(fps))\n mean_speed = inst_speed.rolling(\n 2 * L + 1, min_periods=max(5, L // 6), center=True\n ).mean()\n\n X[\"longdist_dispL\"] = disp.fillna(0.0).clip(lower=0).astype(np.float32)\n X[\"longdist_speedL\"] = mean_speed.fillna(0.0).clip(lower=0).astype(np.float32)\n return X\n\n# ===============================================================\n# 1) Past–vs–Future speed asymmetry (acausal, continuous)\n# Δv = mean_future(speed) - mean_past(speed)\n# ===============================================================\ndef add_speed_asymmetry_future_past_single(\n X: pd.DataFrame, cx: pd.Series, cy: pd.Series, fps: float,\n horizon_base: int = 30, agg: str = \"mean\"\n) -> pd.DataFrame:\n w = max(3, _scale(horizon_base, fps))\n v = _speed(cx, cy, fps)\n if agg == \"median\":\n v_past = v.rolling(w, min_periods=max(3, w//4), center=False).median()\n v_fut = v.iloc[::-1].rolling(w, min_periods=max(3, w//4)).median().iloc[::-1]\n else:\n v_past = v.rolling(w, min_periods=max(3, w//4), center=False).mean()\n v_fut = _roll_future_mean(v, w, min_p=max(3, w//4))\n X[\"spd_asym_1s\"] = (v_fut - v_past).fillna(0.0)\n return X\n\n# ===============================================================\n# 2) Time-reversal asymmetry (3rd-order increment, smoothed)\n# A_h = [(Δ+ x)^3 + (Δ+ y)^3] - [(Δ- x)^3 + (Δ- y)^3]\n# ===============================================================\ndef add_time_reversal_asymmetry3_single(\n X: pd.DataFrame, cx: pd.Series, cy: pd.Series, fps: float,\n h_base: int = 2, smooth_base: int = 11\n) -> pd.DataFrame:\n h = max(1, _scale(h_base, fps))\n dx_f = cx.shift(-h) - cx; dy_f = cy.shift(-h) - cy\n dx_p = cx - cx.shift(h); dy_p = cy - cy.shift(h)\n a = (dx_f**3 + dy_f**3) - (dx_p**3 + dy_p**3)\n ws = max(3, _scale(smooth_base, fps))\n X[\"trev_asym3\"] = a.rolling(ws, min_periods=max(3, ws//4), center=True).mean().fillna(0.0)\n return X\n\n# ===============================================================\n# 3) Distribution shift (future vs past) via symmetric KL of\n# Gaussian fits on speed (cheap, parameter-free)\n# ===============================================================\ndef add_gauss_shift_speed_future_past_single(\n X: pd.DataFrame, cx: pd.Series, cy: pd.Series, fps: float,\n window_base: int = 30, eps: float = 1e-6\n) -> pd.DataFrame:\n w = max(5, _scale(window_base, fps))\n v = _speed(cx, cy, fps)\n\n mu_p = v.rolling(w, min_periods=max(3, w//4)).mean()\n va_p = v.rolling(w, min_periods=max(3, w//4)).var().clip(lower=eps)\n\n mu_f = _roll_future_mean(v, w, min_p=max(3, w//4))\n va_f = _roll_future_var(v, w, min_p=max(3, w//4)).clip(lower=eps)\n\n # KL(Np||Nf) + KL(Nf||Np)\n kl_pf = 0.5 * ((va_p/va_f) + ((mu_f - mu_p)**2)/va_f - 1.0 + np.log(va_f/va_p))\n kl_fp = 0.5 * ((va_f/va_p) + ((mu_p - mu_f)**2)/va_p - 1.0 + np.log(va_p/va_f))\n X[\"spd_symkl_1s\"] = (kl_pf + kl_fp).replace([np.inf, -np.inf], np.nan).fillna(0.0)\n return X\n\n# ===============================================================\n# 4) Future-weighted spectral flux (past vs future FFT of speed)\n# flux = sum(ReLU(|F_future|-|F_past|)) / N\n# ===============================================================\ndef add_spectral_flux_future_single(\n X: pd.DataFrame, cx: pd.Series, cy: pd.Series, fps: float,\n window_base: int = 64\n) -> pd.DataFrame:\n w = max(16, _scale(window_base, fps))\n v = _speed(cx, cy, fps).astype(float).fillna(0.0)\n # construct past/future windows via rolling means of magnitude spectra\n # do simple Hann window + next pow2 FFT\n nfft = 1 << (int(math.ceil(math.log2(max(8, w)))))\n hann = 0.5 - 0.5*np.cos(2*np.pi*np.arange(w)/max(1, w-1))\n\n def mag_fft(seg: np.ndarray) -> np.ndarray:\n # expects length w; zero-pad to nfft\n z = np.zeros(nfft, dtype=float)\n ln = min(w, len(seg))\n z[:ln] = seg[:ln] * hann[:ln]\n return np.abs(np.fft.rfft(z))\n\n # build arrays by shifting once (vectorized enough for w ~ 64)\n past = v.rolling(w, min_periods=w).apply(\n lambda arr: float(np.maximum(0.0, mag_fft(arr)).sum()), raw=True\n )\n fut = v.iloc[::-1].rolling(w, min_periods=w).apply(\n lambda arr: float(np.maximum(0.0, mag_fft(arr)).sum()), raw=True\n ).iloc[::-1]\n X[\"spec_flux_fut\"] = (fut - past).fillna(0.0) / float(nfft)\n return X\n\n# ===============================================================\n# 5) Future-only Katz fractal dimension of speed (cheap FD proxy)\n# ===============================================================\ndef add_katz_fd_future_speed_single(\n X: pd.DataFrame, cx: pd.Series, cy: pd.Series, fps: float,\n horizon_base: int = 30\n) -> pd.DataFrame:\n w = max(8, _scale(horizon_base, fps))\n v = _speed(cx, cy, fps).fillna(0.0)\n\n def katz_fd(seg: np.ndarray) -> float:\n n = len(seg)\n if n < 3:\n return 1.0\n L = np.sum(np.abs(np.diff(seg))) + 1e-12\n d = np.max(np.abs(seg - seg[0])) + 1e-12\n return float(np.log10(n) / (np.log10(n) + np.log10(L/d)))\n\n X[\"katz_fd_fut\"] = (\n v.iloc[::-1].rolling(w, min_periods=w).apply(lambda a: katz_fd(a[::-1]), raw=True).iloc[::-1]\n ).fillna(1.0)\n return X\n\n# ===============================================================\n# 6) Future-only Sample Entropy of speed (m=2, r=0.2*std)\n# Keep window small (<= ~1 s) for cost.\n# ===============================================================\ndef add_sampen_future_speed_single(\n X: pd.DataFrame, cx: pd.Series, cy: pd.Series, fps: float,\n horizon_base: int = 30, m: int = 2, r_frac: float = 0.2\n) -> pd.DataFrame:\n w = max(m+3, _scale(horizon_base, fps))\n v = _speed(cx, cy, fps).astype(float).fillna(0.0)\n\n def sampen(seg: np.ndarray) -> float:\n n = len(seg)\n if n <= m + 3:\n return 0.0\n sd = np.std(seg)\n r = r_frac * (sd + 1e-12)\n # build (n-m+1) template vectors\n Xm = np.lib.stride_tricks.sliding_window_view(seg, m)\n Xm1 = np.lib.stride_tricks.sliding_window_view(seg, m+1)\n # Chebyshev distances via broadcasting (careful with size)\n # use a cap on window to keep cost bounded\n if Xm.shape[0] > 64: # cap pairs\n Xm = Xm[-64:]\n Xm1 = Xm1[-64:]\n Dm = np.max(np.abs(Xm[:, None, :] - Xm[None, :, :]), axis=2)\n Dm1 = np.max(np.abs(Xm1[:, None, :] - Xm1[None, :, :]), axis=2)\n # exclude self-matches on diagonal\n Cm = (Dm <= r).sum() - Xm.shape[0]\n Cm1 = (Dm1 <= r).sum() - Xm1.shape[0]\n Cm = max(Cm, 1)\n Cm1 = max(Cm1, 1)\n return float(-np.log(Cm1 / Cm))\n\n X[\"sampen_fut\"] = (\n v.iloc[::-1].rolling(w, min_periods=w).apply(lambda a: sampen(a[::-1]), raw=False).iloc[::-1]\n ).fillna(0.0)\n return X\n\n# ===============================================================\n# 7) Savitzky–Golay central derivatives (acausal) on speed\n# Produces SG acceleration & jerk at t (tree-friendly, smooth)\n# ===============================================================\ndef add_savgol_center_derivs_single(\n X: pd.DataFrame, cx: pd.Series, cy: pd.Series, fps: float,\n win_base: int = 11, poly: int = 2\n) -> pd.DataFrame:\n w = int(max(5, _scale(win_base, fps)))\n if w % 2 == 0: w += 1\n v = _speed(cx, cy, fps).astype(float)\n\n try:\n from scipy.signal import savgol_filter\n dt = 1.0 / float(fps)\n v_sg = savgol_filter(v.to_numpy(), window_length=w, polyorder=poly, mode=\"nearest\")\n a_sg = savgol_filter(v.to_numpy(), window_length=w, polyorder=poly, deriv=1, delta=dt, mode=\"nearest\")\n j_sg = savgol_filter(v.to_numpy(), window_length=w, polyorder=poly, deriv=2, delta=dt, mode=\"nearest\")\n X[\"acc_sg\"] = pd.Series(a_sg, index=v.index)\n X[\"jerk_sg\"] = pd.Series(j_sg, index=v.index)\n except Exception:\n # fallback: central finite differences + centered smoothing\n dt = 1.0 / float(fps)\n a = (v.shift(-1) - v.shift(1)) / (2*dt)\n j = (v.shift(-1) - 2*v + v.shift(1)) / (dt*dt)\n X[\"acc_sg\"] = a.rolling(w, center=True, min_periods=max(3, w//4)).mean().fillna(0.0)\n X[\"jerk_sg\"] = j.rolling(w, center=True, min_periods=max(3, w//4)).mean().fillna(0.0)\n return X\n\ndef add_selfgroom_features_single(X: pd.DataFrame, single_df: pd.DataFrame, fps: float) -> pd.DataFrame:\n \"\"\"\n Append self-groom-focused, low-cost features (float32) to X.\n `single_df` columns: MultiIndex (bodypart, coord) with coord in {'x','y'}.\n \"\"\"\n eps = 1e-6\n w1 = max(2, int(round(fps))) # ~1 s window (for EWMs)\n whalf = max(2, int(round(0.5*fps))) # ~0.5 s\n\n # ---------- small helpers ----------\n def _has(part): \n return (part, 'x') in single_df.columns and (part, 'y') in single_df.columns\n\n def _xy(part):\n return (single_df[(part, 'x')].to_numpy(dtype=np.float32, copy=False),\n single_df[(part, 'y')].to_numpy(dtype=np.float32, copy=False))\n\n def _speed(x, y): # cm/s\n dx = np.diff(x, prepend=x[:1])\n dy = np.diff(y, prepend=y[:1])\n return np.sqrt(dx*dx + dy*dy) * np.float32(fps)\n\n def _angvel(ax, ay, bx, by): # rad/s for vector a->b\n ang = np.arctan2(by - ay, bx - ax)\n dang = np.diff(np.unwrap(ang), prepend=ang[:1]) * np.float32(fps)\n return np.abs(dang).astype(np.float32, copy=False)\n\n # ---------- choose robust COM proxy ----------\n if _has('hip_left') and _has('hip_right'):\n hx, hy = _xy('hip_left'); kx, ky = _xy('hip_right')\n comx = 0.5*(hx + kx); comy = 0.5*(hy + ky)\n elif _has('lateral_left') and _has('lateral_right'):\n lx, ly = _xy('lateral_left'); rx, ry = _xy('lateral_right')\n comx = 0.5*(lx + rx); comy = 0.5*(ly + ry)\n elif _has('nose') and _has('tail_base'):\n nx, ny = _xy('nose'); tx, ty = _xy('tail_base')\n comx = 0.5*(nx + tx); comy = 0.5*(ny + ty)\n elif _has('neck'):\n comx, comy = _xy('neck')\n else:\n # last resort: first available keypoint\n bp0 = sorted(set([bp for bp, _c in single_df.columns]))[0]\n comx, comy = _xy(bp0)\n\n com_speed = _speed(comx, comy)\n\n # ---------- head & spine signals ----------\n if _has('nose'):\n hx, hy = _xy('nose')\n elif _has('neck'):\n hx, hy = _xy('neck')\n elif _has('ear_left') and _has('ear_right'):\n elx, ely = _xy('ear_left'); erx, ery = _xy('ear_right')\n hx, hy = 0.5*(elx+erx), 0.5*(ely+ery)\n else:\n # fallback to COM if no head landmarks\n hx, hy = comx, comy\n\n head_speed = _speed(hx, hy)\n\n if _has('nose') and _has('tail_base'):\n nx, ny = _xy('nose'); tx, ty = _xy('tail_base')\n spine_angvel = _angvel(nx, ny, tx, ty)\n else:\n spine_angvel = np.zeros_like(head_speed, dtype=np.float32)\n\n # head vs COM motion (high during grooming, low locomotion)\n head_com_ratio = (head_speed / (com_speed + eps)).astype(np.float32, copy=False)\n\n # short-term variability of head motion (rhythmic strokes → higher variability)\n hs = pd.Series(head_speed)\n hs_ewm_std = hs.ewm(span=whalf, adjust=False).std().astype(np.float32)\n hs_ewm_mean = (hs.ewm(span=whalf, adjust=False).mean() + eps).astype(np.float32)\n head_speed_cv = (hs_ewm_std.values / hs_ewm_mean.values).astype(np.float32, copy=False)\n\n # low locomotion index (~ inverse of smoothed COM speed)\n cs = pd.Series(com_speed)\n com_speed_smooth = cs.ewm(span=w1, adjust=False).mean().astype(np.float32).values\n low_loco = (1.0 / (1.0 + com_speed_smooth)).astype(np.float32, copy=False)\n\n # optional: paw ↔ head proximity (when forepaws exist; NaN otherwise)\n paw_head_min = np.full(len(hx), np.nan, dtype=np.float32)\n have_paws = _has('forepaw_left') and _has('forepaw_right')\n if have_paws:\n plx, ply = _xy('forepaw_left'); prx, pry = _xy('forepaw_right')\n dL = np.sqrt((plx - hx)**2 + (ply - hy)**2, dtype=np.float32)\n dR = np.sqrt((prx - hx)**2 + (pry - hy)**2, dtype=np.float32)\n paw_head_min = np.minimum(dL, dR).astype(np.float32, copy=False)\n\n # ---------- append (no copies of X) ----------\n X['sg_com_speed'] = com_speed.astype(np.float32, copy=False)\n X['sg_head_speed'] = head_speed.astype(np.float32, copy=False)\n X['sg_head_com_ratio'] = head_com_ratio\n X['sg_spine_angvel'] = spine_angvel\n X['sg_head_speed_cv_0p5s'] = head_speed_cv\n X['sg_low_locomotion_idx'] = low_loco\n X['sg_paw_head_min'] = paw_head_min # may be NaN if paws absent\n return X\n\ndef _unit_vec(dx: pd.Series, dy: pd.Series, eps=1e-9):\n den = np.sqrt(dx*dx + dy*dy) + eps\n return dx/den, dy/den\n\ndef _angle(x: pd.Series, y: pd.Series):\n # radians in [-pi, pi]\n return np.arctan2(y.to_numpy(dtype=float), x.to_numpy(dtype=float))\n\ndef add_pursuit_features(X, mouse_pair, fps):\n # ---- positions & headings (body_center; nose->tail heading) -------------------\n Ax, Ay = mouse_pair['A']['body_center']['x'], mouse_pair['A']['body_center']['y']\n Bx, By = mouse_pair['B']['body_center']['x'], mouse_pair['B']['body_center']['y']\n Avx, Avy = Ax.diff(), Ay.diff()\n Bvx, Bvy = Bx.diff(), By.diff()\n\n # heading vectors (nose - tail_base) if available; else use velocity as fallback\n def _heading(side):\n cols = mouse_pair[side].columns.get_level_values(0)\n if 'nose' in cols and 'tail_base' in cols:\n hx = mouse_pair[side]['nose']['x'] - mouse_pair[side]['tail_base']['x']\n hy = mouse_pair[side]['nose']['y'] - mouse_pair[side]['tail_base']['y']\n else:\n hx, hy = mouse_pair[side]['body_center']['x'].diff(), mouse_pair[side]['body_center']['y'].diff()\n return _unit_vec(hx, hy)\n\n hAx, hAy = _heading('A')\n hBx, hBy = _heading('B')\n\n # LOS from A->B\n rx, ry = (Bx - Ax), (By - Ay)\n uLx, uLy = _unit_vec(rx, ry) # LOS unit\n\n # ---- 1) LOS geometry & closing ------------------------------------------------\n # distance & its derivative (closing_speed), LOS angle & its rate\n dist = np.sqrt(rx*rx + ry*ry)\n vrel_x, vrel_y = (Bvx - Avx), (Bvy - Avy) # relative velocity\n closing_speed = (vrel_x * uLx + vrel_y * uLy) # <0 = closing\n X['closing_speed'] = closing_speed.astype(np.float32)\n\n # TTC (only when closing); inverse/logized versions (bounded)\n ttc = dist / np.maximum(1e-6, -closing_speed)\n ttc = ttc.where(closing_speed < 0, np.inf)\n X['ttc_inv'] = (1.0 / (ttc.replace([np.inf, -np.inf], np.nan))).fillna(0.0).astype(np.float32)\n X['ttc_log'] = np.log1p(1.0 + X['ttc_inv']).astype(np.float32)\n\n # LOS angle and centered rate (non-causal)\n los_ang = _angle(rx, ry)\n w = _scale(5, fps) # ~0.17s @30fps\n # smooth & differentiate\n los_s = pd.Series(los_ang).rolling(w, center=True, min_periods=1).mean()\n X['los_rate'] = pd.Series(los_s).diff().astype(np.float32)\n\n # Bearing consistency (lower std ⇒ steadier pursuit LOS)\n for W in (10, 20, 30):\n wW = _scale(W, fps)\n X[f'bearing_consistency_w{W}'] = (1.0 - pd.Series(los_ang).rolling(wW, center=True, min_periods=1).std()).astype(np.float32)\n\n # ---- 2) Facing & sector cues --------------------------------------------------\n # A facing B; B facing A (cosines)\n faceAB = hAx * uLx + hAy * uLy\n faceBA = hBx * (-uLx) + hBy * (-uLy)\n X['faceAB'] = faceAB.astype(np.float32)\n X['faceBA'] = faceBA.astype(np.float32)\n\n # Rear-follow sector: A is behind B within ~30deg and within ~2 body lengths\n # rear if B's heading is aligned with vector B->A\n uBAx, uBAy = -uLx, -uLy\n rear_align = (hBx * uBAx + hBy * uBAy)\n lateral = (uBAx * (-hBy) + uBAy * (hBx)) # signed lateral offset wrt B heading\n # X['rear_followAB'] = ((rear_align > np.cos(np.deg2rad(30))) & (dist < dist.quantile(0.7))).astype(float)\n X['lat_offset_AB'] = np.abs(lateral).astype(np.float32)\n\n # ---- 3) Speed dominance & decomposition --------------------------------------\n vA_mag = np.sqrt(Avx*Avx + Avy*Avy)\n vB_mag = np.sqrt(Bvx*Bvx + Bvy*Bvy)\n X['vA_mag'] = vA_mag.astype(np.float32)\n X['vB_mag'] = vB_mag.astype(np.float32)\n X['speed_dom'] = (vA_mag - vB_mag).astype(np.float32)\n # X['speed_ratio'] = (vA_mag / (vB_mag + 1e-6)).clip(0, 10).astype(np.float32)\n\n # relative velocity components wrt LOS\n X['vrel_par'] = (vrel_x * uLx + vrel_y * uLy).astype(np.float32)\n # X['vrel_perp'] = (vrel_x * (-uLy) + vrel_y * (uLx)).astype(np.float32)\n\n # ---- 4) Heading rate & curvature (centered windows) --------------------------\n # heading angles\n hA_ang = _angle(hAx, hAy)\n hB_ang = _angle(hBx, hBy)\n X['head_rate_A'] = pd.Series(hA_ang).diff().astype(np.float32)\n X['head_rate_B'] = pd.Series(hB_ang).diff().astype(np.float32)\n\n # curvature κ = |x'y'' - y'x''| / ( (x'^2 + y'^2)^(3/2) )\n def _curvature(x, y, W):\n x1, y1 = x.diff(), y.diff()\n x2, y2 = x1.diff(), y1.diff()\n num = (x1 * y2 - y1 * x2).abs()\n den = (x1*x1 + y1*y1)**1.5 + 1e-9\n k = (num / den).rolling(W, center=True, min_periods=1).median()\n return k\n\n for W in (10, 20, 30):\n wW = _scale(W, fps)\n X[f'curvA_w{W}'] = _curvature(Ax, Ay, wW).astype(np.float32)\n X[f'curvB_w{W}'] = _curvature(Bx, By, wW).astype(np.float32)\n X[f'curv_diff_w{W}'] = (X[f'curvB_w{W}'] - X[f'curvA_w{W}']).astype(np.float32)\n\n # ---- 5) Speed cross-corr lag (leader–follower) --------------------------------\n # small bounded scan ±20 frames; record argmax lag (sign tells who leads)\n L = _scale(20, fps)\n # normalize to zero-mean per video slice to stabilize\n sA = (vA_mag - vA_mag.rolling(_scale(60, fps), center=True, min_periods=1).mean()).fillna(0.0)\n sB = (vB_mag - vB_mag.rolling(_scale(60, fps), center=True, min_periods=1).mean()).fillna(0.0)\n lags = range(-L, L+1, max(1, L//10)) # coarse scan\n cc = {lag: (sA * sB.shift(lag)).rolling(_scale(30, fps), center=True, min_periods=1).mean()\n for lag in lags}\n cc_df = pd.DataFrame(cc)\n X['lag_cc_speed'] = cc_df.idxmax(axis=1).astype(np.float32)\n\n # ---- 6) Contact-imminent & lunge ------------------------------------------------\n # rolling mins (centered) for impending contact, and closing acceleration\n d_center = dist\n # X['closing_accel'] = pd.Series(closing_speed).diff().astype(np.float32)\n for W in (5, 10):\n wW = _scale(W, fps)\n X[f'dmin_center_w{W}'] = d_center.rolling(wW, center=True, min_periods=1).min().astype(np.float32)\n\n # Optional: if noses exist, add nose–nose min too (cheap)\n colsA = mouse_pair['A'].columns.get_level_values(0)\n colsB = mouse_pair['B'].columns.get_level_values(0)\n if 'nose' in colsA and 'nose' in colsB:\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 W in (5, 10):\n wW = _scale(W, fps)\n X[f'dmin_nose_w{W}'] = nn.rolling(wW, center=True, min_periods=1).min().astype(np.float32)\n\n # ---- 7) Wall-aware LOS rate (optional; needs your wall normals) ----------------\n # If you already computed u_wall_tangent unit vectors externally, you can project los_rate on it.\n return X\n\ndef _add_pair_meta(X: pd.DataFrame, meta_df: Optional[pd.DataFrame]) -> pd.DataFrame:\n try:\n if meta_df is not None and len(meta_df):\n r0 = meta_df.iloc[0]\n vid, agent, target = r0.get('video_id'), r0.get('agent_id'), r0.get('target_id')\n is_test = bool(r0.get('is_test', False))\n mss = float(_pair_same_sex(vid, agent, target, is_test=is_test))\n else:\n mss = 0.0\n except Exception:\n mss = 0.0\n X = X.copy()\n X['meta_same_sex'] = np.float32(mss)\n return X\n\n\ndef _add_speed_ear(X: pd.DataFrame, mp: pd.DataFrame, fps: float) -> pd.DataFrame:\n if ('A','ear_left') not in mp.columns or ('B','ear_left') not in mp.columns:\n return X\n lag = _scale(10, fps)\n Ael, Bel = mp['A']['ear_left'], mp['B']['ear_left']\n shA, shB = Ael.shift(lag), Bel.shift(lag)\n F = pd.DataFrame({\n 'sp_A': np.square(Ael - shA).sum(axis=1, skipna=False),\n 'sp_AB': np.square(Ael - shB).sum(axis=1, skipna=False),\n 'sp_B': np.square(Bel - shB).sum(axis=1, skipna=False),\n }, index=X.index)\n return pd.concat([X, F], axis=1)\n\n\ndef _add_elongation_proxy(X: pd.DataFrame) -> pd.DataFrame:\n # works whether your cols are \"nose+tail_base\" or \"12+nose+tail_base\"\n a = 'nose+tail_base' if 'nose+tail_base' in X.columns else '12+nose+tail_base'\n b = 'ear_left+ear_right' if 'ear_left+ear_right' in X.columns else '12+ear_left+ear_right'\n if a in X.columns and b in X.columns:\n Xt = X.copy()\n Xt['elong'] = Xt[a] / (Xt[b] + 1e-6)\n return Xt\n return X\n\n\ndef _add_rel_ori_and_contact(X: pd.DataFrame, mp: pd.DataFrame, fps: float) -> pd.DataFrame:\n colsA = mp['A'].columns.get_level_values(0)\n colsB = mp['B'].columns.get_level_values(0)\n if all(p in colsA for p in ('nose','tail_base')) and all(p in colsB for p in ('nose','tail_base')):\n dir_A = mp['A']['nose'] - mp['A']['tail_base']\n dir_B = mp['B']['nose'] - mp['B']['tail_base']\n num = dir_A['x']*dir_B['x'] + dir_A['y']*dir_B['y']\n den = (np.sqrt(dir_A['x']**2 + dir_A['y']**2) * np.sqrt(dir_B['x']**2 + dir_B['y']**2) + 1e-6)\n F = pd.DataFrame({'rel_ori': num/den}, index=X.index)\n X = pd.concat([X, F], axis=1)\n X = add_contact_bins(X, mp, fps) # your existing helper\n return X\n\n\ndef _add_approach_rate(X: pd.DataFrame, mp: pd.DataFrame, fps: float) -> pd.DataFrame:\n if ('A','nose') not in mp.columns or ('B','nose') not in mp.columns:\n return X\n lag = _scale(10, fps)\n cur = np.square(mp['A']['nose'] - mp['B']['nose']).sum(axis=1, skipna=False)\n past = np.square(mp['A']['nose'].shift(lag) - mp['B']['nose'].shift(lag)).sum(axis=1, skipna=False)\n Xt = X.copy()\n Xt['appr'] = cur - past\n return Xt\n\n\ndef _add_distance_bins(X: pd.DataFrame, mp: pd.DataFrame) -> pd.DataFrame:\n if ('A','body_center') not in mp.columns or ('B','body_center') not in mp.columns:\n return X\n dx = mp['A']['body_center']['x'] - mp['B']['body_center']['x']\n dy = mp['A']['body_center']['y'] - mp['B']['body_center']['y']\n cd = np.sqrt(dx*dx + dy*dy)\n F = pd.DataFrame({\n 'v_cls': (cd < 5.0).astype(float),\n 'cls' : ((cd >= 5.0) & (cd < 15.0)).astype(float),\n 'med' : ((cd >= 15.0) & (cd < 30.0)).astype(float),\n 'far' : (cd >= 30.0).astype(float),\n }, index=X.index)\n return pd.concat([X, F], axis=1)\n\n\ndef _add_temporal_interactions(X: pd.DataFrame, mp: pd.DataFrame, fps: float) -> pd.DataFrame:\n if ('A','body_center') not in mp.columns or ('B','body_center') not in mp.columns:\n return X\n cd_full = np.square(mp['A']['body_center'] - mp['B']['body_center']).sum(axis=1, skipna=False)\n Axd, Ayd = mp['A']['body_center']['x'].diff(), mp['A']['body_center']['y'].diff()\n Bxd, Byd = mp['B']['body_center']['x'].diff(), mp['B']['body_center']['y'].diff()\n coord = Axd*Bxd + Ayd*Byd\n\n def _ws(w): return int(max(1, round(w * (fps/30.0))))\n roll = lambda w: dict(window=_ws(w), min_periods=1, center=True)\n\n F = pd.DataFrame(index=X.index)\n for w in (5, 15, 30, 60):\n r = roll(w)\n F[f'd_m{w}'] = cd_full.rolling(**r).mean()\n F[f'd_s{w}'] = cd_full.rolling(**r).std()\n F[f'd_mn{w}'] = cd_full.rolling(**r).min()\n F[f'd_mx{w}'] = cd_full.rolling(**r).max()\n F[f'int{w}'] = 1 / (1 + cd_full.rolling(**r).var())\n F[f'co_m{w}'] = coord.rolling(**r).mean()\n F[f'co_s{w}'] = coord.rolling(**r).std()\n return pd.concat([X, F], axis=1)\n\n\ndef _add_nose_nose(X: pd.DataFrame, mp: pd.DataFrame, fps: float) -> pd.DataFrame:\n if ('A','nose') not in mp.columns or ('B','nose') not in mp.columns:\n return X\n nn = np.sqrt((mp['A']['nose']['x'] - mp['B']['nose']['x'])**2 +\n (mp['A']['nose']['y'] - mp['B']['nose']['y'])**2)\n F = pd.DataFrame(index=X.index)\n for lag in (10, 20, 40):\n l = _scale(lag, fps)\n F[f'nn_lg{lag}'] = nn.shift(l)\n F[f'nn_ch{lag}'] = nn - nn.shift(l)\n F[f'cl_ps{lag}'] = (nn < 10.0).rolling(l, min_periods=1).mean()\n return pd.concat([X, F], axis=1)\n\n\ndef _add_val_and_intcon(X: pd.DataFrame, mp: pd.DataFrame, fps: float) -> pd.DataFrame:\n if ('A','body_center') not in mp.columns or ('B','body_center') not in mp.columns:\n return X\n Avx = mp['A']['body_center']['x'].diff()\n Avy = mp['A']['body_center']['y'].diff()\n Bvx = mp['B']['body_center']['x'].diff()\n Bvy = mp['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 F = pd.DataFrame({'va_0': val}, index=X.index)\n for off in (-20, 0, 20):\n o = _scale_signed(off, fps)\n F[f'va_{off}'] = val.shift(-o)\n\n cd_full = np.square(mp['A']['body_center'] - mp['B']['body_center']).sum(axis=1, skipna=False)\n w = int(max(1, round(30 * (fps/30.0))))\n F['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 # plug your advanced extras\n X2 = pd.concat([X, F], axis=1)\n return add_interaction_features(X2, mp, mp['A'].columns.get_level_values(0), mp['B'].columns.get_level_values(0), fps)\n\n\ndef _add_arena(X: pd.DataFrame, mp: pd.DataFrame, shape: str, w_cm: float, h_cm: float) -> pd.DataFrame:\n if ('A','body_center') not in mp.columns or ('B','body_center') not in mp.columns:\n return X\n\n Ax, Ay = mp['A']['body_center']['x'], mp['A']['body_center']['y']\n Bx, By = mp['B']['body_center']['x'], mp['B']['body_center']['y']\n\n def _norm01(v: pd.Series) -> pd.Series:\n q0, q1 = np.nanquantile(v.to_numpy(), [0.005, 0.995])\n rng = max(q1 - q0, 1e-6)\n return ((v - q0) / rng).clip(0.0, 1.0)\n\n if np.isfinite(w_cm) and w_cm > 0 and np.isfinite(h_cm) and h_cm > 0:\n Axn, Ayn = (Ax / w_cm).clip(0,1), (Ay / h_cm).clip(0,1)\n Bxn, Byn = (Bx / w_cm).clip(0,1), (By / h_cm).clip(0,1)\n else:\n Axn, Ayn = _norm01(Ax), _norm01(Ay)\n Bxn, Byn = _norm01(Bx), _norm01(By)\n\n Axn, Ayn = np.asarray(Axn, float), np.asarray(Ayn, float)\n Bxn, Byn = np.asarray(Bxn, float), np.asarray(Byn, float)\n\n wA = distance_to_wall_norm(Axn, Ayn, arena_shape=shape)\n wB = distance_to_wall_norm(Bxn, Byn, arena_shape=shape)\n cA = np.sqrt((Axn - 0.5)**2 + (Ayn - 0.5)**2)\n cB = np.sqrt((Bxn - 0.5)**2 + (Byn - 0.5)**2)\n\n n = len(Axn)\n def _ori_unit(side):\n cols = mp[side].columns.get_level_values(0)\n if 'nose' in cols and 'tail_base' in cols:\n ox = (mp[side]['nose']['x'] - mp[side]['tail_base']['x']).to_numpy(dtype=float)\n oy = (mp[side]['nose']['y'] - mp[side]['tail_base']['y']).to_numpy(dtype=float)\n d = np.sqrt(ox*ox + oy*oy) + 1e-9\n return ox/d, oy/d\n return np.zeros(n, float), np.zeros(n, float)\n\n oAx, oAy = _ori_unit('A')\n oBx, oBy = _ori_unit('B')\n\n if 'circ' in shape:\n uAx, uAy = _u_wall_circ(Axn, Ayn)\n uBx, uBy = _u_wall_circ(Bxn, Byn)\n else:\n uAx, uAy = _u_wall_rect(Axn, Ayn)\n uBx, uBy = _u_wall_rect(Bxn, Byn)\n\n tAx, tAy = -uAy, uAx\n tBx, tBy = -uBy, uBx\n\n F = pd.DataFrame({\n 'wall_A': wA, 'wall_B': wB, 'wall_min': np.minimum(wA, wB),\n 'center_A': cA, 'center_B': cB,\n 'oriA_wall_norm': oAx*uAx + oAy*uAy,\n 'oriB_wall_norm': oBx*uBx + oBy*uBy,\n 'oriA_wall_tan': np.abs(oAx*tAx + oAy*tAy),\n 'oriB_wall_tan': np.abs(oBx*tBx + oBy*tBy),\n 'corner_A': distance_to_corner_norm(Axn, Ayn, arena_shape=shape),\n 'corner_B': distance_to_corner_norm(Bxn, Byn, arena_shape=shape),\n 'corner_min': np.minimum(\n distance_to_corner_norm(Axn, Ayn, arena_shape=shape),\n distance_to_corner_norm(Bxn, Byn, arena_shape=shape)),\n }, index=X.index)\n return pd.concat([X, F], axis=1)\n\ndef add_shepherd_proxies_pair(\n X: pd.DataFrame,\n mouse_pair: pd.DataFrame,\n meta_df: Optional[pd.DataFrame],\n fps: float,\n) -> pd.DataFrame:\n \"\"\"\n Shepherd-focused pairwise features.\n\n Intuition:\n - shepherd episodes typically happen with:\n * target pressed near a wall,\n * agent more central than target,\n * agent behind target (rear_followAB == 1),\n * target running along the wall for a bit.\n - we turn that into a couple of cheap, non-redundant shep_* features.\n\n Requirements:\n - ('A','body_center'), ('B','body_center') present.\n - Uses (and if needed creates) 'wall_A', 'wall_B', 'along_wall_B'.\n - Reuses 'rear_followAB' from add_pursuit_features if present.\n \"\"\"\n # avoid double-adding\n if 'shep_wall_trap' in X.columns:\n return X\n\n # need body_center\n colsA = set(mouse_pair['A'].columns.get_level_values(0))\n colsB = set(mouse_pair['B'].columns.get_level_values(0))\n if 'body_center' not in colsA or 'body_center' not in colsB:\n return X\n\n X = X.copy()\n\n # --- ensure we have basic wall geometry (no redundancy with other blocks) ----\n # If you already call add_wall_geometry_pair() elsewhere, this will be a no-op\n # because column names won't collide.\n need_wall = not ({'wall_A', 'wall_B'} <= set(X.columns))\n if need_wall or ('along_wall_B' not in X.columns):\n # This will add: wall_A, wall_B, toward_wall_A/B, along_wall_A/B\n X = add_wall_geometry_pair(X, mouse_pair, meta_df)\n\n # fallbacks if for some reason those columns are still missing\n if not ({'wall_A', 'wall_B'} <= set(X.columns)):\n # last-resort: compute only wall distances (no tangents) so we can still run\n shape, w_cm, h_cm = _extract_arena_meta(meta_df)\n Ax = mouse_pair['A']['body_center']['x']; Ay = mouse_pair['A']['body_center']['y']\n Bx = mouse_pair['B']['body_center']['x']; By = mouse_pair['B']['body_center']['y']\n Axn, Ayn = _norm_xy_series(Ax, Ay, w_cm, h_cm)\n Bxn, Byn = _norm_xy_series(Bx, By, w_cm, h_cm)\n X['wall_A'] = distance_to_wall_norm(Axn, Ayn, shape)\n X['wall_B'] = distance_to_wall_norm(Bxn, Byn, shape)\n\n # rename to short handles\n wallA = pd.Series(X['wall_A'], index=X.index).astype(float)\n wallB = pd.Series(X['wall_B'], index=X.index).astype(float)\n alongB = pd.Series(X.get('along_wall_B', 0.0), index=X.index).astype(float)\n\n # reuse pursuit feature if available; else zero\n rear = pd.Series(X.get('rear_followAB', 0.0), index=X.index).astype(float)\n\n # --- 1) Instantaneous \"wall trap\" proxy --------------------------------------\n # B is close to wall, A more central, A behind B.\n # distance_to_wall_norm is in [0, ~0.5]; \"near wall\" ~ lowest 15%.\n near_wall_B = (wallB <= 0.15).astype(float)\n A_outside = (wallA > wallB).astype(float) # A is more central than B\n # strong motion along wall (parallel) – already encoded by along_wall_B\n corridor = (alongB > 0.7).astype(float)\n\n shep_wall_trap = near_wall_B * A_outside * rear * corridor\n\n # --- 2) Smoothed trap score over ~2 s window --------------------------------\n w_long = _scale(60, fps) # ~2 s at 30 fps\n shep_wall_trap_smooth = (\n pd.Series(shep_wall_trap, index=X.index)\n .rolling(w_long, min_periods=max(3, w_long // 4))\n .mean()\n .fillna(0.0)\n )\n\n # --- 3) \"Driven into wall\" dynamics -----------------------------------------\n # B’s wall distance decreasing over ~1 s while A is behind.\n w_h = _scale(30, fps) # ~1 s\n wallB_series = pd.Series(wallB, index=X.index)\n\n wallB_past = wallB_series.rolling(w_h, min_periods=max(3, w_h // 4)).mean()\n wallB_future = _roll_future_mean(wallB_series, w_h, min_p=max(3, w_h // 4))\n\n drive = wallB_future - wallB_past # < 0 when moving *toward* wall\n shep_wall_drive = (-drive).clip(lower=0.0) * rear\n\n # --- 4) gate to frames where the pair is actually interacting ---------------\n Ax = mouse_pair['A']['body_center']['x']; Ay = mouse_pair['A']['body_center']['y']\n Bx = mouse_pair['B']['body_center']['x']; By = mouse_pair['B']['body_center']['y']\n dist = np.sqrt((Ax - Bx)**2 + (Ay - By)**2)\n # use 80% quantile as \"typical interaction distance\"\n d_thr = float(np.nanquantile(dist.to_numpy(), 0.8)) if len(dist) else 0.0\n near_pair = (dist < d_thr).astype(float)\n\n shep_wall_trap = (shep_wall_trap * near_pair).astype(np.float32)\n shep_wall_trap_smooth = (shep_wall_trap_smooth * near_pair).astype(np.float32)\n shep_wall_drive = (shep_wall_drive * near_pair).astype(np.float32)\n\n X['shep_wall_trap'] = shep_wall_trap\n X['shep_wall_trap_smooth'] = shep_wall_trap_smooth\n X['shep_wall_drive'] = shep_wall_drive\n\n return X\n","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"######### CONTEXT FEATS (LITE) #########################\n\n# -------------------- small helpers --------------------\n\ndef _quick_stats_series(s: pd.Series):\n if s.size == 0:\n return dict(mean=np.nan, std=np.nan, p5=np.nan, p50=np.nan, p95=np.nan)\n q = s.quantile([0.05, 0.50, 0.95])\n return dict(\n mean=float(s.mean()),\n std=float(s.std(ddof=0)),\n p5=float(q.loc[0.05]),\n p50=float(q.loc[0.50]),\n p95=float(q.loc[0.95]),\n )\n\ndef _frames(seconds: float, fps: float) -> int:\n return max(1, int(round(float(seconds) * float(fps))))\n\ndef _unit_heading_from_pose(df_like: pd.DataFrame, fps: float) -> Tuple[pd.Series, pd.Series]:\n \"\"\"\n Prefer anatomical heading (nose - body_center). If unavailable or tiny, fallback to velocity.\n Returns unit (hx, hy). If still tiny, defaults to (1,0).\n \"\"\"\n eps = 1e-6\n cols = df_like.columns\n has_nose = ('nose', 'x') in cols and ('nose', 'y') in cols\n has_bc = ('body_center', 'x') in cols and ('body_center', 'y') in cols\n if has_nose and has_bc:\n hx = (df_like[('nose','x')] - df_like[('body_center','x')]).astype('float32')\n hy = (df_like[('nose','y')] - df_like[('body_center','y')]).astype('float32')\n elif has_bc:\n vx = df_like[('body_center','x')].diff() * float(fps)\n vy = df_like[('body_center','y')].diff() * float(fps)\n hx, hy = vx, vy\n else:\n idx = df_like.index\n hx = pd.Series(1.0, index=idx)\n hy = pd.Series(0.0, index=idx)\n\n norm = np.sqrt(hx**2 + hy**2).clip(eps)\n hx_u = (hx / norm).astype('float32')\n hy_u = (hy / norm).astype('float32')\n hx_u = hx_u.fillna(1.0)\n hy_u = hy_u.fillna(0.0)\n return hx_u, hy_u\n\ndef _iter_context_mice_simple(ctx: pd.DataFrame):\n \"\"\"\n Yields per-mouse sub-DFs shaped (bodypart, coord) from a context_df with\n columns MultiIndex (mouse_id, bodypart, coord). If already 2-level, yields once.\n \"\"\"\n if ctx is None or ctx.shape[1] == 0:\n return\n if not isinstance(ctx.columns, pd.MultiIndex) or ctx.columns.nlevels < 3:\n yield ctx\n return\n for mid in ctx.columns.get_level_values(0).unique():\n yield ctx.xs(mid, level=0, axis=1)\n\n# -------------------- LITE: SINGLE --------------------\n\ndef add_context_single(\n X: pd.DataFrame,\n focal_df: pd.DataFrame,\n context_df: Optional[pd.DataFrame] = None,\n meta_df: Optional[pd.DataFrame] = None, # unused in LITE but kept for API compatibility\n fps: float = 30.0,\n large_fill: float = 1e2,\n verbose: bool = False,\n) -> pd.DataFrame:\n \"\"\"\n LITE social-context features for SINGLE (updated):\n\n Kept / expanded core signals:\n • ctx_nn_dist : nearest-neighbour distance (cm)\n • ctx_nn_dist_inv : 1 / max(ctx_nn_dist, 1 cm)\n • ctx_nn_approach : frame-to-frame decrease of NN distance * fps (cm/s, ≥0)\n • ctx_nn_bearing_cos : cos(angle between focal heading and NN direction)\n • ctx_nn_bearing_abs : |ctx_nn_bearing_cos| (alignment magnitude)\n • ctx_pressure : ∑_others 1 / max(dist, 1)^2\n • ctx_pressure_front : pressure from others in front (cos ≥ 0.3)\n • ctx_pressure_rear : pressure from behind (cos ≤ -0.3)\n • ctx_pressure_left : pressure to the left (positive cross)\n • ctx_pressure_right : pressure to the right (negative cross)\n\n Removed (low SHAP importance):\n • All *_r10 count features\n • ctx_nn_ttc\n \"\"\"\n if 'body_center' not in focal_df.columns.get_level_values(0):\n raise ValueError(\"focal_df must already contain ('body_center','x'|'y').\")\n\n # Focal geometry\n cx_f = focal_df[('body_center','x')].astype('float32')\n cy_f = focal_df[('body_center','y')].astype('float32')\n vx_f = cx_f.diff() * float(fps)\n vy_f = cy_f.diff() * float(fps)\n hx_u, hy_u = _unit_heading_from_pose(focal_df, fps=fps)\n\n # Collect context body centers (others)\n others: List[Tuple[pd.Series, pd.Series]] = []\n if context_df is not None and context_df.shape[1] > 0:\n if not focal_df.index.equals(context_df.index):\n context_df = context_df.reindex(focal_df.index)\n for ctx_one in _iter_context_mice_simple(context_df):\n ctx_one = fill_or_create_body_center_any(ctx_one)\n if ('body_center','x') in ctx_one.columns and ('body_center','y') in ctx_one.columns:\n ox = ctx_one[('body_center','x')].astype('float32')\n oy = ctx_one[('body_center','y')].astype('float32')\n others.append((ox, oy))\n\n n = len(others)\n m = len(X.index)\n idx = X.index\n\n # Initialize outputs (defaults if no others)\n X['ctx_nn_dist'] = np.float32(large_fill)\n X['ctx_nn_dist_inv'] = np.float32(0.0)\n X['ctx_nn_approach'] = np.float32(0.0)\n X['ctx_nn_bearing_cos'] = np.float32(0.0)\n X['ctx_nn_bearing_abs'] = np.float32(0.0)\n X['ctx_pressure'] = np.float32(0.0)\n X['ctx_pressure_front'] = np.float32(0.0)\n X['ctx_pressure_rear'] = np.float32(0.0)\n X['ctx_pressure_left'] = np.float32(0.0)\n X['ctx_pressure_right'] = np.float32(0.0)\n\n if n == 0:\n if verbose:\n print(\"[add_context_single:LITE] No context mice → only default constants.\")\n return X\n\n # Build per-other stacks without creating per-slot columns\n D = np.full((m, n), fill_value=float(large_fill), dtype=np.float32) # distances\n Ux = np.zeros((m, n), dtype=np.float32) # unit vec x from focal → other\n Uy = np.zeros((m, n), dtype=np.float32)\n\n for j, (ox, oy) in enumerate(others):\n dx = (ox - cx_f).astype('float32')\n dy = (oy - cy_f).astype('float32')\n dist = np.sqrt(dx**2 + dy**2).astype(np.float32)\n miss = (ox.isna() | oy.isna())\n dist = dist.mask(miss, other=np.float32(large_fill)).fillna(np.float32(large_fill))\n D[:, j] = dist.values\n\n denom = dist.replace(0, np.nan)\n ux = (dx / denom).fillna(0.0).astype('float32')\n uy = (dy / denom).fillna(0.0).astype('float32')\n Ux[:, j] = ux.values\n Uy[:, j] = uy.values\n\n # Aggregates based on nearest neighbour and pressure-style crowding\n d_nn = D.min(axis=1) # numpy array\n s_nn = pd.Series(d_nn, index=idx)\n\n X['ctx_nn_dist'] = s_nn.astype('float32')\n\n # Inverse NN distance (cap denominator at 1 cm to avoid extreme values)\n denom_nn = np.maximum(d_nn, 1.0)\n X['ctx_nn_dist_inv'] = pd.Series(1.0 / denom_nn, index=idx).astype('float32')\n\n # Approach speed (only closing, ≥ 0)\n app = (s_nn.shift(1) - s_nn) * float(fps)\n app = app.clip(lower=0.0).fillna(0.0).astype('float32')\n X['ctx_nn_approach'] = app\n\n # Bearing to NN (cosine via dot with heading)\n row_idx = np.arange(m, dtype=int)\n argmin_idx = D.argmin(axis=1)\n hx = hx_u.values.astype(np.float32)\n hy = hy_u.values.astype(np.float32)\n nn_cos = hx * Ux[row_idx, argmin_idx] + hy * Uy[row_idx, argmin_idx]\n nn_cos = np.clip(nn_cos, -1.0, 1.0)\n X['ctx_nn_bearing_cos'] = pd.Series(nn_cos, index=idx).astype('float32')\n X['ctx_nn_bearing_abs'] = pd.Series(np.abs(nn_cos), index=idx).astype('float32')\n\n # Distance-weighted crowding / \"pressure\"\n # inv2 = 1 / max(dist, 1)^2, so far-away mice contribute ~0\n denom = np.maximum(1.0, D)\n inv2 = (1.0 / (denom**2)).astype(np.float32)\n pressure = inv2.sum(axis=1)\n X['ctx_pressure'] = pd.Series(pressure, index=idx).astype('float32')\n\n # Sector pressures (front/rear via cos; left/right via cross sign)\n cos_mat = (hx[:, None] * Ux + hy[:, None] * Uy).astype(np.float32)\n cross = (hx[:, None] * Uy - hy[:, None] * Ux).astype(np.float32)\n\n front_mask = (cos_mat >= 0.3).astype(np.float32)\n rear_mask = (cos_mat <= -0.3).astype(np.float32)\n left_mask = (cross > 0.0).astype(np.float32)\n right_mask = (cross < 0.0).astype(np.float32)\n\n X['ctx_pressure_front'] = pd.Series((inv2 * front_mask).sum(axis=1), index=idx).astype('float32')\n X['ctx_pressure_rear'] = pd.Series((inv2 * rear_mask ).sum(axis=1), index=idx).astype('float32')\n X['ctx_pressure_left'] = pd.Series((inv2 * left_mask ).sum(axis=1), index=idx).astype('float32')\n X['ctx_pressure_right'] = pd.Series((inv2 * right_mask).sum(axis=1), index=idx).astype('float32')\n\n if verbose:\n agg = [\n 'ctx_nn_dist','ctx_nn_dist_inv','ctx_nn_approach',\n 'ctx_nn_bearing_cos','ctx_nn_bearing_abs',\n 'ctx_pressure','ctx_pressure_front','ctx_pressure_rear',\n 'ctx_pressure_left','ctx_pressure_right',\n ]\n stats = {c: _quick_stats_series(X[c]) for c in agg}\n print(\"[add_context_single:LITE] Aggregates:\",\n {k: {kk: round(vv, 4) for kk, vv in d.items()} for k, d in stats.items()})\n\n return X\n\n# -------------------- LITE: PAIR --------------------\n\ndef add_context_pair(\n X: pd.DataFrame,\n pair_df: pd.DataFrame,\n context_df: Optional[pd.DataFrame] = None,\n meta_df: Optional[pd.DataFrame] = None, # unused in LITE but kept for API compatibility\n fps: float = 30.0,\n large_fill: float = 1e2,\n verbose: bool = False,\n) -> pd.DataFrame:\n \"\"\"\n LITE social-context features for PAIR (A and B) (updated):\n\n For A and B separately:\n • {prefix}_nn_dist : nearest-neighbour distance\n • {prefix}_nn_dist_inv : 1 / max(nn_dist, 1)\n • {prefix}_nn_approach : closing speed toward NN (cm/s, ≥0)\n • {prefix}_nn_bearing_cos : cos(angle between heading and NN)\n • {prefix}_nn_bearing_abs : |bearing_cos|\n • {prefix}_pressure : ∑_others 1 / max(dist,1)^2\n • {prefix}_pressure_front : pressure from front (cos ≥ 0.3)\n • {prefix}_pressure_rear : pressure from behind (cos ≤ -0.3)\n • {prefix}_pressure_left : pressure to left (cross > 0)\n • {prefix}_pressure_right : pressure to right (cross < 0)\n\n AB corridor occlusion (3rd+ mice between A and B):\n • ctx_between_AB_min_dist : min distance of any other to segment AB\n • ctx_between_AB_pressure : ∑_others 1 / max(dist_to_AB,1)^2\n\n Removed (low SHAP importance):\n • All *_r10 count features\n • ctxA_nn_ttc, ctxB_nn_ttc\n • ctx_between_AB_r10\n \"\"\"\n need = [('A','body_center','x'), ('A','body_center','y'),\n ('B','body_center','x'), ('B','body_center','y')]\n if any(c not in pair_df.columns for c in need):\n raise ValueError(\"pair_df must contain ('A'|'B','body_center','x'|'y').\")\n\n def _xy(df_like, who, part):\n return df_like[(who, part, 'x')].astype('float32'), df_like[(who, part, 'y')].astype('float32')\n\n cx_A, cy_A = _xy(pair_df, 'A', 'body_center')\n cx_B, cy_B = _xy(pair_df, 'B', 'body_center')\n hxA_u, hyA_u = _unit_heading_from_pose(pair_df.xs('A', level=0, axis=1), fps=fps)\n hxB_u, hyB_u = _unit_heading_from_pose(pair_df.xs('B', level=0, axis=1), fps=fps)\n\n # Gather others (3rd+ mice)\n others: List[Tuple[pd.Series, pd.Series]] = []\n if context_df is not None and context_df.shape[1] > 0:\n if not pair_df.index.equals(context_df.index):\n context_df = context_df.reindex(pair_df.index)\n for ctx_one in _iter_context_mice_simple(context_df):\n ctx_one = fill_or_create_body_center_any(ctx_one)\n if ('body_center','x') in ctx_one.columns and ('body_center','y') in ctx_one.columns:\n ox = ctx_one[('body_center','x')].astype('float32')\n oy = ctx_one[('body_center','y')].astype('float32')\n others.append((ox, oy))\n\n # Initialize outputs\n def _init(prefix: str):\n X[f'{prefix}_nn_dist'] = np.float32(large_fill)\n X[f'{prefix}_nn_dist_inv'] = np.float32(0.0)\n X[f'{prefix}_nn_approach'] = np.float32(0.0)\n X[f'{prefix}_nn_bearing_cos'] = np.float32(0.0)\n X[f'{prefix}_nn_bearing_abs'] = np.float32(0.0)\n X[f'{prefix}_pressure'] = np.float32(0.0)\n X[f'{prefix}_pressure_front'] = np.float32(0.0)\n X[f'{prefix}_pressure_rear'] = np.float32(0.0)\n X[f'{prefix}_pressure_left'] = np.float32(0.0)\n X[f'{prefix}_pressure_right'] = np.float32(0.0)\n\n _init('ctxA'); _init('ctxB')\n\n idx = X.index\n m = len(idx)\n n = len(others)\n\n # Fast exit if no others\n if n == 0:\n X['ctx_between_AB_min_dist'] = np.float32(large_fill)\n X['ctx_between_AB_pressure'] = np.float32(0.0)\n if verbose:\n print(\"[add_context_pair:LITE] No context mice → only corridor defaults.\")\n return X\n\n # Build stacks relative to A and B (no per-slot columns)\n def _stacks(cx, cy):\n D = np.full((m, n), fill_value=float(large_fill), dtype=np.float32)\n Ux = np.zeros((m, n), dtype=np.float32)\n Uy = np.zeros((m, n), dtype=np.float32)\n for j, (ox, oy) in enumerate(others):\n dx = (ox - cx).astype('float32')\n dy = (oy - cy).astype('float32')\n dist = np.sqrt(dx**2 + dy**2).astype(np.float32)\n miss = (ox.isna() | oy.isna())\n dist = dist.mask(miss, other=np.float32(large_fill)).fillna(np.float32(large_fill))\n D[:, j] = dist.values\n denom = dist.replace(0, np.nan)\n Ux[:, j] = (dx / denom).fillna(0.0).astype('float32').values\n Uy[:, j] = (dy / denom).fillna(0.0).astype('float32').values\n return D, Ux, Uy\n\n DA, UxA, UyA = _stacks(cx_A, cy_A)\n DB, UxB, UyB = _stacks(cx_B, cy_B)\n\n def _per_focal(prefix: str, D, Ux, Uy, hx_u, hy_u):\n d_nn = D.min(axis=1)\n s_nn = pd.Series(d_nn, index=idx)\n\n # NN distance + inverse\n X[f'{prefix}_nn_dist'] = s_nn.astype('float32')\n denom_nn = np.maximum(d_nn, 1.0)\n X[f'{prefix}_nn_dist_inv'] = pd.Series(1.0 / denom_nn, index=idx).astype('float32')\n\n # Approach speed\n app = (s_nn.shift(1) - s_nn) * float(fps)\n app = app.clip(lower=0.0).fillna(0.0).astype('float32')\n X[f'{prefix}_nn_approach'] = app\n\n # Bearing cosine to NN\n hx = hx_u.values.astype(np.float32)\n hy = hy_u.values.astype(np.float32)\n nn_cos = hx * Ux[np.arange(m), D.argmin(axis=1)] + hy * Uy[np.arange(m), D.argmin(axis=1)]\n nn_cos = np.clip(nn_cos, -1.0, 1.0)\n X[f'{prefix}_nn_bearing_cos'] = pd.Series(nn_cos, index=idx).astype('float32')\n X[f'{prefix}_nn_bearing_abs'] = pd.Series(np.abs(nn_cos), index=idx).astype('float32')\n\n # Pressure-style crowding\n denom = np.maximum(1.0, D)\n inv2 = (1.0 / (denom**2)).astype(np.float32)\n X[f'{prefix}_pressure'] = pd.Series(inv2.sum(axis=1), index=idx).astype('float32')\n\n cos_mat = (hx[:, None] * Ux + hy[:, None] * Uy).astype(np.float32)\n cross = (hx[:, None] * Uy - hy[:, None] * Ux).astype(np.float32)\n\n front_mask = (cos_mat >= 0.3).astype(np.float32)\n rear_mask = (cos_mat <= -0.3).astype(np.float32)\n left_mask = (cross > 0.0).astype(np.float32)\n right_mask = (cross < 0.0).astype(np.float32)\n\n X[f'{prefix}_pressure_front'] = pd.Series((inv2 * front_mask).sum(axis=1), index=idx).astype('float32')\n X[f'{prefix}_pressure_rear'] = pd.Series((inv2 * rear_mask ).sum(axis=1), index=idx).astype('float32')\n X[f'{prefix}_pressure_left'] = pd.Series((inv2 * left_mask ).sum(axis=1), index=idx).astype('float32')\n X[f'{prefix}_pressure_right'] = pd.Series((inv2 * right_mask).sum(axis=1), index=idx).astype('float32')\n\n _per_focal('ctxA', DA, UxA, UyA, hxA_u, hyA_u)\n _per_focal('ctxB', DB, UxB, UyB, hxB_u, hyB_u)\n\n # Corridor occlusion between A and B\n def _dist_point_to_segment(ax, ay, bx, by, ox, oy):\n ABx, ABy = (bx - ax), (by - ay)\n AOx, AOy = (ox - ax), (oy - ay)\n AB2 = ABx**2 + ABy**2 + 1e-9\n t = ((AOx * ABx + AOy * ABy) / AB2).clip(0.0, 1.0)\n projx = ax + t * ABx\n projy = ay + t * ABy\n return np.sqrt((ox - projx)**2 + (oy - projy)**2)\n\n between_pressure = np.zeros(m, dtype=np.float32)\n min_between = np.full(m, fill_value=float(large_fill), dtype=np.float32)\n ax, ay = cx_A, cy_A\n bx, by = cx_B, cy_B\n for (ox, oy) in others:\n dseg = _dist_point_to_segment(ax, ay, bx, by, ox, oy).astype('float32')\n # coverage along corridor as distance-weighted \"pressure\"\n denom_seg = np.maximum(dseg, 1.0)\n between_pressure += (1.0 / (denom_seg**2)).fillna(0.0).values\n min_between = np.minimum(min_between, dseg.fillna(np.float32(large_fill)).values)\n\n X['ctx_between_AB_min_dist'] = pd.Series(min_between, index=idx).astype('float32')\n X['ctx_between_AB_pressure'] = pd.Series(between_pressure, index=idx).astype('float32')\n\n if verbose:\n aggA = [c for c in X.columns if c.startswith('ctxA_')]\n aggB = [c for c in X.columns if c.startswith('ctxB_')]\n both = aggA + aggB + ['ctx_between_AB_min_dist','ctx_between_AB_pressure']\n stats = {c: _quick_stats_series(X[c]) for c in both}\n print(\"[add_context_pair:LITE] Aggregates:\",\n {k: {kk: round(vv, 4) for kk, vv in d.items()} for k, d in stats.items()})\n\n return X\n","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#################################### body parts filler\n\n# ---------- geometry helpers ----------\ndef _norm2(x, y):\n return np.sqrt(np.square(x) + np.square(y))\n\ndef _safe_unit(dx, dy, eps=1e-6):\n n = _norm2(dx, dy)\n n = n.where(n > eps, np.nan)\n return dx / n, dy / n\n\ndef _rot90(dx, dy):\n # rotate +90 degrees (x, y) -> (-y, x)\n return -dy, dx\n\n\n# ---------- CPU implementations (use pd_cpu ONLY) ----------\n\ndef _fill_or_create_body_center_single_cpu(df: pd_cpu.DataFrame) -> pd_cpu.DataFrame:\n \"\"\"\n Single-mouse version (2-level MultiIndex: level0=bodypart, level1 in {'x','y'}).\n Preserves existing non-NaNs in ('body_center', 'x'/'y'); creates/fills if missing.\n \"\"\"\n if not isinstance(df.columns, pd_cpu.MultiIndex) or df.columns.nlevels < 2:\n raise ValueError(\"Expected 2-level MultiIndex columns: [bodypart, coord 'x'/'y'].\")\n\n out = df.copy()\n idx = out.index\n\n # Ensure target columns exist (create as NaN if missing)\n if ('body_center', 'x') not in out.columns:\n out[('body_center', 'x')] = np.nan\n if ('body_center', 'y') not in out.columns:\n out[('body_center', 'y')] = np.nan\n\n bc_x = out[('body_center', 'x')].copy()\n bc_y = out[('body_center', 'y')].copy()\n\n def get_xy(part: str):\n x = out[(part, 'x')] if (part, 'x') in out.columns else pd_cpu.Series(np.nan, index=idx)\n y = out[(part, 'y')] if (part, 'y') in out.columns else pd_cpu.Series(np.nan, index=idx)\n return x, y\n\n # Frequently-annotated parts\n llx, lly = get_xy('lateral_left')\n lrx, lry = get_xy('lateral_right')\n nx, ny = get_xy('neck')\n nozx, nozy = get_xy('nose')\n elx, ely = get_xy('ear_left')\n erx, ery = get_xy('ear_right')\n tbx, tby = get_xy('tail_base')\n hlx, hly = get_xy('hip_left')\n hrx, hry = get_xy('hip_right')\n\n need_bc = bc_x.isna() | bc_y.isna()\n\n # Rule 1: flank midpoint\n flank_ok = llx.notna() & lly.notna() & lrx.notna() & lry.notna()\n m = need_bc & flank_ok\n if m.any():\n midx = (llx + lrx) / 2.0\n midy = (lly + lry) / 2.0\n bc_x.loc[m & bc_x.isna()] = midx.loc[m & bc_x.isna()]\n bc_y.loc[m & bc_y.isna()] = midy.loc[m & bc_y.isna()]\n need_bc = bc_x.isna() | bc_y.isna()\n\n # Anchor-based fallback\n if need_bc.any():\n # Front: neck -> nose -> ear midpoint\n front_x = pd_cpu.Series(np.nan, index=idx)\n front_y = pd_cpu.Series(np.nan, index=idx)\n\n neck_ok = nx.notna() & ny.notna()\n use = need_bc & neck_ok\n if use.any():\n front_x.loc[use] = nx.loc[use]\n front_y.loc[use] = ny.loc[use]\n\n front_missing = need_bc & (front_x.isna() | front_y.isna())\n nose_ok = nozx.notna() & nozy.notna()\n use = front_missing & nose_ok\n if use.any():\n front_x.loc[use] = nozx.loc[use]\n front_y.loc[use] = nozy.loc[use]\n\n front_missing = need_bc & (front_x.isna() | front_y.isna())\n ears_ok = elx.notna() & erx.notna() & ely.notna() & ery.notna()\n use = front_missing & ears_ok\n if use.any():\n front_x.loc[use] = ((elx + erx) / 2.0).loc[use]\n front_y.loc[use] = ((ely + ery) / 2.0).loc[use]\n\n front_ok = front_x.notna() & front_y.notna()\n\n # Hind: mean(tail_base, hip_mid) -> hip_mid -> tail_base\n hip_mid_x = pd_cpu.Series(np.nan, index=idx)\n hip_mid_y = pd_cpu.Series(np.nan, index=idx)\n hips_ok = hlx.notna() & hly.notna() & hrx.notna() & hry.notna()\n if hips_ok.any():\n hip_mid_x.loc[hips_ok] = (hlx.loc[hips_ok] + hrx.loc[hips_ok]) / 2.0\n hip_mid_y.loc[hips_ok] = (hly.loc[hips_ok] + hry.loc[hips_ok]) / 2.0\n\n hind_x = pd_cpu.Series(np.nan, index=idx)\n hind_y = pd_cpu.Series(np.nan, index=idx)\n\n both_hind_avail = hips_ok & tbx.notna() & tby.notna()\n use = need_bc & both_hind_avail\n if use.any():\n hind_x.loc[use] = ((hip_mid_x + tbx) / 2.0).loc[use]\n hind_y.loc[use] = ((hip_mid_y + tby) / 2.0).loc[use]\n\n hind_missing = need_bc & (hind_x.isna() | hind_y.isna())\n use = hind_missing & hips_ok\n if use.any():\n hind_x.loc[use] = hip_mid_x.loc[use]\n hind_y.loc[use] = hip_mid_y.loc[use]\n\n hind_missing = need_bc & (hind_x.isna() | hind_y.isna())\n tail_only = tbx.notna() & tby.notna()\n use = hind_missing & tail_only\n if use.any():\n hind_x.loc[use] = tbx.loc[use]\n hind_y.loc[use] = tby.loc[use]\n\n hind_ok = hind_x.notna() & hind_y.notna()\n\n m = need_bc & front_ok & hind_ok\n if m.any():\n cand_x = (front_x + hind_x) / 2.0\n cand_y = (front_y + hind_y) / 2.0\n bc_x.loc[m & bc_x.isna()] = cand_x.loc[m & bc_x.isna()]\n bc_y.loc[m & bc_y.isna()] = cand_y.loc[m & bc_y.isna()]\n\n # Fill existing columns (preserve non-NaNs)\n out[('body_center', 'x')] = out[('body_center', 'x')].where(\n out[('body_center', 'x')].notna(), bc_x\n )\n out[('body_center', 'y')] = out[('body_center', 'y')].where(\n out[('body_center', 'y')].notna(), bc_y\n )\n\n # Keep original column order\n return out\n\n\ndef _fill_or_create_body_center_any_cpu(df: pd_cpu.DataFrame) -> pd_cpu.DataFrame:\n if not isinstance(df.columns, pd_cpu.MultiIndex):\n raise ValueError(\"Expected MultiIndex columns.\")\n\n if df.columns.nlevels == 2:\n return _fill_or_create_body_center_single_cpu(df)\n\n if df.columns.nlevels >= 3:\n # keep original top-level order of mice\n mice_order = pd_cpu.Index(\n df.columns.get_level_values(0),\n name=df.columns.names[0]\n ).unique()\n\n pieces = []\n for m in mice_order:\n sub = df.xs(m, axis=1, level=0, drop_level=True)\n sub_filled = _fill_or_create_body_center_single_cpu(sub)\n pieces.append(pd_cpu.concat({m: sub_filled}, axis=1))\n\n out = pd_cpu.concat(pieces, axis=1)\n # ONLY reorder by top level; keep all (including new) subcolumns\n out = out.reindex(columns=mice_order, level=0)\n out.columns.names = df.columns.names\n return out\n\n raise ValueError(\"Unhandled column structure.\")\n\n\ndef _fill_or_create_nose_single_cpu(\n df: pd_cpu.DataFrame,\n alpha_default: float = 0.91,\n alpha_clip: tuple[float, float] = (0.30, 1.50),\n) -> pd_cpu.DataFrame:\n \"\"\"\n Fill/create ('nose','x'/'y') from ears (+ neck or tail_base if available).\n - Uses per-DF robust alpha if some nose labels exist: median(||nose-ear_mid|| / ear_dist).\n - Otherwise uses alpha_default.\n Preserves existing non-NaNs in ('nose','x'/'y').\n \"\"\"\n if not isinstance(df.columns, pd_cpu.MultiIndex) or df.columns.nlevels < 2:\n raise ValueError(\"Expected 2-level MultiIndex columns: [bodypart, coord 'x'/'y'].\")\n\n out = df.copy()\n idx = out.index\n\n # Ensure columns exist\n for col in [('nose', 'x'), ('nose', 'y')]:\n if col not in out.columns:\n out[col] = np.nan\n\n # Pull coordinates (Series or all-NaN Series if missing)\n def get_xy(part: str):\n x = out[(part, 'x')] if (part, 'x') in out.columns else pd_cpu.Series(np.nan, index=idx)\n y = out[(part, 'y')] if (part, 'y') in out.columns else pd_cpu.Series(np.nan, index=idx)\n return x, y\n\n elx, ely = get_xy('ear_left')\n erx, ery = get_xy('ear_right')\n nx, ny = get_xy('neck')\n tbx, tby = get_xy('tail_base')\n\n nose_x = out[('nose', 'x')].copy()\n nose_y = out[('nose', 'y')].copy()\n\n # Masks\n have_ears = elx.notna() & ely.notna() & erx.notna() & ery.notna()\n need_nose = nose_x.isna() | nose_y.isna()\n work = have_ears & need_nose\n if not work.any():\n return out # nothing to do\n\n # Ear midpoint & distance\n emx = (elx + erx) / 2.0\n emy = (ely + ery) / 2.0\n edx = erx - elx\n edy = ery - ely\n ear_dist = _norm2(edx, edy)\n\n # If some nose exists, estimate alpha robustly\n has_true_nose = have_ears & nose_x.notna() & nose_y.notna()\n if has_true_nose.any():\n true_dx = nose_x - emx\n true_dy = nose_y - emy\n step = _norm2(true_dx, true_dy)\n with np.errstate(invalid='ignore', divide='ignore'):\n alpha_series = (step / ear_dist).clip(lower=alpha_clip[0], upper=alpha_clip[1])\n alpha_val = float(np.nanmedian(alpha_series.values))\n if not np.isfinite(alpha_val):\n alpha_val = alpha_default\n else:\n alpha_val = alpha_default\n\n # Candidate heading vectors:\n # 1) neck -> ear_mid (forward)\n nx_ok = nx.notna() & ny.notna()\n v1x = emx - nx\n v1y = emy - ny\n v1x, v1y = _safe_unit(v1x, v1y)\n\n # 2) tail_base -> ear_mid (headward)\n tb_ok = tbx.notna() & tby.notna()\n v2x = emx - tbx\n v2y = emy - tby\n v2x, v2y = _safe_unit(v2x, v2y)\n\n # 3) perpendicular to ear axis (choose sign using tail_base if available; else default)\n ux, uy = _safe_unit(edx, edy)\n ppx, ppy = _rot90(ux, uy) # +90°\n mmx = emx - tbx # headward proxy for sign if tail exists\n mmy = emy - tby\n # sign decision: use dot(pp, headward) ≥ 0 else flip\n dot = ppx * mmx + ppy * mmy\n sign = np.where(tb_ok, np.sign(dot).replace(0, 1), 1) # if dot==0 or no tail, use +1\n ppx = ppx * sign\n ppy = ppy * sign\n\n # Choose best available direction per-frame\n dir_x = pd_cpu.Series(np.nan, index=idx)\n dir_y = pd_cpu.Series(np.nan, index=idx)\n\n use1 = work & nx_ok\n dir_x.loc[use1] = v1x.loc[use1]\n dir_y.loc[use1] = v1y.loc[use1]\n\n use2 = work & dir_x.isna() & tb_ok\n dir_x.loc[use2] = v2x.loc[use2]\n dir_y.loc[use2] = v2y.loc[use2]\n\n use3 = work & dir_x.isna()\n dir_x.loc[use3] = ppx.loc[use3]\n dir_y.loc[use3] = ppy.loc[use3]\n\n # Step length = alpha * ear_dist\n step = (alpha_val * ear_dist).where(work, np.nan)\n\n cand_x = emx + step * dir_x\n cand_y = emy + step * dir_y\n\n # Fill only where missing, preserve existing non-NaNs\n nose_x = nose_x.where(nose_x.notna(), cand_x)\n nose_y = nose_y.where(nose_y.notna(), cand_y)\n\n out[('nose', 'x')] = nose_x\n out[('nose', 'y')] = nose_y\n return out\n\n\ndef _fill_or_create_nose_any_cpu(\n df: pd_cpu.DataFrame,\n **kwargs,\n) -> pd_cpu.DataFrame:\n \"\"\"\n Works like _fill_or_create_nose_single_cpu but supports multi-mouse data\n (top level = mouse id). Leaves column order intact.\n \"\"\"\n if not isinstance(df.columns, pd_cpu.MultiIndex):\n raise ValueError(\"Expected MultiIndex columns.\")\n\n if df.columns.nlevels == 2:\n return _fill_or_create_nose_single_cpu(df, **kwargs)\n\n # nlevels >= 3: split by first level (mouse), process, then concat\n mice_order = pd_cpu.Index(\n df.columns.get_level_values(0),\n name=df.columns.names[0]\n ).unique()\n pieces = []\n for m in mice_order:\n sub = df.xs(m, axis=1, level=0, drop_level=True)\n sub_filled = _fill_or_create_nose_single_cpu(sub, **kwargs)\n pieces.append(pd_cpu.concat({m: sub_filled}, axis=1))\n\n out = pd_cpu.concat(pieces, axis=1)\n out = out.reindex(columns=mice_order, level=0) # preserve top-level order\n out.columns.names = df.columns.names\n return out\n\n\n# ---------- Public wrappers (work with cudf.pandas OR plain pandas) ----------\n\ndef fill_or_create_body_center_single(df: pd.DataFrame) -> pd.DataFrame:\n \"\"\"\n Wrapper that runs the logic on CPU pandas (pd_cpu) and returns\n a DataFrame in the currently-active pandas namespace (pd).\n \"\"\"\n df_cpu = pd_cpu.DataFrame(df)\n out_cpu = _fill_or_create_body_center_single_cpu(df_cpu)\n return pd.DataFrame(out_cpu)\n\n\ndef fill_or_create_body_center_any(df: pd.DataFrame) -> pd.DataFrame:\n df_cpu = pd_cpu.DataFrame(df)\n out_cpu = _fill_or_create_body_center_any_cpu(df_cpu)\n return pd.DataFrame(out_cpu)\n\n\ndef fill_or_create_nose_single(\n df: pd.DataFrame,\n alpha_default: float = 0.91,\n alpha_clip: tuple[float, float] = (0.30, 1.50),\n) -> pd.DataFrame:\n df_cpu = pd_cpu.DataFrame(df)\n out_cpu = _fill_or_create_nose_single_cpu(\n df_cpu,\n alpha_default=alpha_default,\n alpha_clip=alpha_clip,\n )\n return pd.DataFrame(out_cpu)\n\n\ndef fill_or_create_nose_any(df: pd.DataFrame, **kwargs) -> pd.DataFrame:\n df_cpu = pd_cpu.DataFrame(df)\n out_cpu = _fill_or_create_nose_any_cpu(df_cpu, **kwargs)\n return pd.DataFrame(out_cpu)\n\nimport numpy as np\nimport pandas as pd\n\n\ndef interpolate_ears_single(\n df: pd.DataFrame,\n max_linear_gap: int = 15,\n min_template_frames: int = 20,\n) -> pd.DataFrame:\n \"\"\"\n Fill short NaN gaps in ear_left/ear_right via linear interpolation and\n long gaps via a simple head-geometry triangulation using nose/neck/body_center.\n\n Expects 2-level MultiIndex columns: (bodypart, 'x'/'y').\n Keeps all other columns unchanged.\n \"\"\"\n if not isinstance(df.columns, pd.MultiIndex) or df.columns.nlevels < 2:\n raise ValueError(\"Expected 2-level MultiIndex columns [bodypart, coord].\")\n\n out = df.copy()\n idx = out.index\n\n def has_part(part: str) -> bool:\n return (part, \"x\") in out.columns and (part, \"y\") in out.columns\n\n # nothing to do if ears missing from schema\n if not (has_part(\"ear_left\") and has_part(\"ear_right\")):\n return out\n\n # ---------- 1) Short-gap interpolation (per ear, per coord) ----------\n for part in (\"ear_left\", \"ear_right\"):\n for coord in (\"x\", \"y\"):\n col = (part, coord)\n s = out[col].astype(float)\n # fills only gaps up to max_linear_gap, leaves longer gaps as NaN\n s_lin = s.interpolate(\n method=\"linear\",\n limit=max_linear_gap,\n limit_direction=\"both\",\n )\n out[col] = s_lin\n\n # Track which frames had *real* ear labels for template learning\n orig = {}\n for part in (\"ear_left\", \"ear_right\"):\n orig[(part, \"x\")] = df[(part, \"x\")].notna()\n orig[(part, \"y\")] = df[(part, \"y\")].notna()\n\n # ---------- 2) Build head center & local frame (u = forward, p = lateral) ----------\n def get_xy(part: str):\n if has_part(part):\n return out[(part, \"x\")].astype(float), out[(part, \"y\")].astype(float)\n return pd.Series(np.nan, index=idx), pd.Series(np.nan, index=idx)\n\n nx, ny = get_xy(\"nose\")\n kx, ky = get_xy(\"neck\")\n cx, cy = get_xy(\"body_center\")\n\n Hx = pd.Series(np.nan, index=idx) # head center\n Hy = pd.Series(np.nan, index=idx)\n ux = pd.Series(np.nan, index=idx) # forward unit vector\n uy = pd.Series(np.nan, index=idx)\n\n nose_ok = nx.notna() & ny.notna()\n neck_ok = kx.notna() & ky.notna()\n bc_ok = cx.notna() & cy.notna()\n\n # case 1: nose + neck -> head midpoint, axis neck->nose\n m = nose_ok & neck_ok\n if m.any():\n Hx.loc[m] = 0.5 * (nx[m] + kx[m])\n Hy.loc[m] = 0.5 * (ny[m] + ky[m])\n vx = nx[m] - kx[m]\n vy = ny[m] - ky[m]\n nrm = np.sqrt(vx**2 + vy**2) + 1e-6\n ux.loc[m] = vx / nrm\n uy.loc[m] = vy / nrm\n\n # case 2: nose + body_center\n m = Hx.isna() & nose_ok & bc_ok\n if m.any():\n Hx.loc[m] = 0.5 * (nx[m] + cx[m])\n Hy.loc[m] = 0.5 * (ny[m] + cy[m])\n vx = nx[m] - cx[m]\n vy = ny[m] - cy[m]\n nrm = np.sqrt(vx**2 + vy**2) + 1e-6\n ux.loc[m] = vx / nrm\n uy.loc[m] = vy / nrm\n\n # case 3: neck + body_center\n m = Hx.isna() & neck_ok & bc_ok\n if m.any():\n Hx.loc[m] = 0.5 * (kx[m] + cx[m])\n Hy.loc[m] = 0.5 * (ky[m] + cy[m])\n vx = kx[m] - cx[m]\n vy = ky[m] - cy[m]\n nrm = np.sqrt(vx**2 + vy**2) + 1e-6\n ux.loc[m] = vx / nrm\n uy.loc[m] = vy / nrm\n\n # crude fallbacks: use nose or body_center directly if nothing else\n m = Hx.isna() & nose_ok\n if m.any():\n Hx.loc[m] = nx[m]\n Hy.loc[m] = ny[m]\n ux.loc[m] = 1.0\n uy.loc[m] = 0.0\n\n m = Hx.isna() & bc_ok\n if m.any():\n Hx.loc[m] = cx[m]\n Hy.loc[m] = cy[m]\n ux.loc[m] = 1.0\n uy.loc[m] = 0.0\n\n # lateral axis = +90° rotation of forward\n px = -uy\n py = ux\n\n # ---------- 3) Learn ear templates in head frame ----------\n # For each ear, we express E - H in [u, p] basis and take median (fwd, lat).\n templates: dict[str, tuple[float, float] | None] = {}\n for part in (\"ear_left\", \"ear_right\"):\n ex = out[(part, \"x\")].astype(float)\n ey = out[(part, \"y\")].astype(float)\n\n valid = (\n orig[(part, \"x\")]\n & orig[(part, \"y\")]\n & Hx.notna()\n & Hy.notna()\n & ux.notna()\n & uy.notna()\n )\n if valid.sum() < min_template_frames:\n # too few good frames -> don't trust triangulation for this ear\n templates[part] = None\n continue\n\n vx = ex[valid] - Hx[valid]\n vy = ey[valid] - Hy[valid]\n\n fwd = vx * ux[valid] + vy * uy[valid] # along head axis\n lat = vx * px[valid] + vy * py[valid] # lateral (left/right)\n\n f_med = float(np.nanmedian(fwd.values))\n l_med = float(np.nanmedian(lat.values))\n templates[part] = (f_med, l_med)\n\n # Optional symmetry fix: enforce opposite lateral offsets if both sides exist\n if templates.get(\"ear_left\") is not None and templates.get(\"ear_right\") is not None:\n fL, sL = templates[\"ear_left\"]\n fR, sR = templates[\"ear_right\"]\n if sL * sR > 0: # both same sign -> flip right\n sR = -sL\n templates[\"ear_right\"] = (fR, sR)\n\n # ---------- 4) Fill long gaps via triangulation ----------\n for part in (\"ear_left\", \"ear_right\"):\n tpl = templates.get(part)\n if tpl is None:\n continue\n\n f_med, l_med = tpl\n colx = (part, \"x\")\n coly = (part, \"y\")\n\n missing = out[colx].isna() | out[coly].isna()\n mask = (\n missing\n & Hx.notna()\n & Hy.notna()\n & ux.notna()\n & uy.notna()\n )\n if not mask.any():\n continue\n\n ex_hat = Hx + f_med * ux + l_med * px\n ey_hat = Hy + f_med * uy + l_med * py\n\n ex = out[colx].copy()\n ey = out[coly].copy()\n ex.loc[mask] = ex_hat.loc[mask]\n ey.loc[mask] = ey_hat.loc[mask]\n out[colx] = ex\n out[coly] = ey\n\n # ---------- 5) Optional final clean-up: nearest neighbor fill ----------\n # (if some frames still NaN because anchors were totally missing)\n for part in (\"ear_left\", \"ear_right\"):\n for coord in (\"x\", \"y\"):\n col = (part, coord)\n s = out[col].astype(float)\n if s.isna().any():\n out[col] = s.interpolate(method=\"nearest\", limit_direction=\"both\")\n\n return out\n\n\ndef interpolate_ears_any(\n df: pd.DataFrame,\n max_linear_gap: int = 15,\n min_template_frames: int = 20,\n) -> pd.DataFrame:\n \"\"\"\n Apply interpolate_ears_single to:\n - a single-mouse pose DF with 2-level MultiIndex columns, or\n - a multi-mouse DF with ≥3 levels (top level = mouse id).\n\n Keeps original top-level order intact.\n \"\"\"\n if not isinstance(df.columns, pd.MultiIndex):\n raise ValueError(\"Expected MultiIndex columns.\")\n\n # Single mouse: just run the core helper\n if df.columns.nlevels == 2:\n return interpolate_ears_single(\n df,\n max_linear_gap=max_linear_gap,\n min_template_frames=min_template_frames,\n )\n\n # Multi-mouse: split by first level, process, then concat\n mice_order = pd.Index(\n df.columns.get_level_values(0),\n name=df.columns.names[0],\n ).unique()\n\n pieces = []\n for m in mice_order:\n sub = df.xs(m, axis=1, level=0, drop_level=True)\n sub_interp = interpolate_ears_single(\n sub,\n max_linear_gap=max_linear_gap,\n min_template_frames=min_template_frames,\n )\n pieces.append(pd.concat({m: sub_interp}, axis=1))\n\n out = pd.concat(pieces, axis=1)\n out = out.reindex(columns=mice_order, level=0)\n out.columns.names = df.columns.names\n return out\n\nimport numpy as np\nimport pandas as pd\n\n# ---------------------------------------------------------\n# Small helpers (re-use your existing _norm2 / _rot90)\n# ---------------------------------------------------------\n\ndef _safe_unit_pd(dx: pd.Series, dy: pd.Series, eps: float = 1e-6):\n \"\"\"Return unit vector (ux, uy) with NaNs when norm < eps.\"\"\"\n n = np.sqrt(dx * dx + dy * dy)\n m = n > eps\n ux = pd.Series(np.nan, index=dx.index, dtype=float)\n uy = pd.Series(np.nan, index=dy.index, dtype=float)\n ux[m] = dx[m] / n[m]\n uy[m] = dy[m] / n[m]\n return ux, uy\n\n\ndef _interp_short_gaps_1d(s: pd.Series, max_gap: int = 15) -> pd.Series:\n \"\"\"\n Linearly interpolate only gaps of length <= max_gap.\n Longer gaps remain NaN for now.\n \"\"\"\n s = s.copy()\n isnan = s.isna().to_numpy()\n if not isnan.any():\n return s\n\n idx = np.arange(len(s))\n # find start/end indices of each NaN run\n diff = np.diff(isnan.astype(int), prepend=0, append=0)\n starts = np.where(diff == 1)[0]\n ends = np.where(diff == -1)[0] - 1 # inclusive\n\n for st, en in zip(starts, ends):\n gap_len = en - st + 1\n if gap_len <= max_gap:\n # local indices for interpolation\n left = st - 1\n right = en + 1\n if left < 0 or right >= len(s) or pd.isna(s.iloc[left]) or pd.isna(s.iloc[right]):\n # can't safely interp this small gap, skip\n continue\n x = [left, right]\n y = [s.iloc[left], s.iloc[right]]\n xi = np.arange(st, en + 1)\n s.iloc[st:en+1] = np.interp(xi, x, y)\n return s\n\n\n# ---------------------------------------------------------\n# Build body frame: origin + (forward, side) unit vectors\n# ---------------------------------------------------------\n\ndef _build_body_frame_single(df: pd.DataFrame):\n \"\"\"\n Build a body-centric frame per frame:\n - origin (ox, oy)\n - forward unit (fx, fy) ~ head direction\n - side unit (sx, sy) ~ left-right\n\n Uses as many BPs as available, with a chain of fallbacks:\n - origin: body_center -> hip_mid -> lateral_mid -> ear_mid -> neck -> nose\n - head: nose -> neck -> ear_mid -> body_center\n - tail: tail_base if present\n - hips: hip_mid if both hips present\n - forward: head - hip_mid; else head - origin; else COM velocity\n Returns:\n ox, oy, fx, fy, sx, sy, valid_mask\n \"\"\"\n idx = df.index\n\n def _get_xy(part: str):\n if (part, \"x\") in df.columns and (part, \"y\") in df.columns:\n return df[(part, \"x\")], df[(part, \"y\")]\n return pd.Series(np.nan, index=idx, dtype=float), pd.Series(np.nan, index=idx, dtype=float)\n\n # key parts\n bcx, bcy = _get_xy(\"body_center\")\n elx, ely = _get_xy(\"ear_left\")\n erx, ery = _get_xy(\"ear_right\")\n nx, ny = _get_xy(\"neck\")\n nozx,nozy= _get_xy(\"nose\")\n tbx, tby = _get_xy(\"tail_base\")\n hlx, hly = _get_xy(\"hip_left\")\n hrx, hry = _get_xy(\"hip_right\")\n llx, lly = _get_xy(\"lateral_left\")\n lrx, lry = _get_xy(\"lateral_right\")\n\n # ---- origin: body_center -> hip_mid -> lateral_mid -> ear_mid -> neck -> nose ----\n ox = pd.Series(np.nan, index=idx, dtype=float)\n oy = pd.Series(np.nan, index=idx, dtype=float)\n\n # body_center\n bc_ok = bcx.notna() & bcy.notna()\n ox[bc_ok] = bcx[bc_ok]\n oy[bc_ok] = bcy[bc_ok]\n\n need_o = ox.isna() | oy.isna()\n\n # hip_mid\n hips_ok = hlx.notna() & hly.notna() & hrx.notna() & hry.notna()\n use = need_o & hips_ok\n if use.any():\n hipx = (hlx + hrx) / 2.0\n hipy = (hly + hry) / 2.0\n ox[use] = hipx[use]\n oy[use] = hipy[use]\n need_o = ox.isna() | oy.isna()\n\n # lateral_mid\n lat_ok = llx.notna() & lly.notna() & lrx.notna() & lry.notna()\n use = need_o & lat_ok\n if use.any():\n latx = (llx + lrx) / 2.0\n laty = (lly + lry) / 2.0\n ox[use] = latx[use]\n oy[use] = laty[use]\n need_o = ox.isna() | oy.isna()\n\n # ear_mid\n ears_ok = elx.notna() & ely.notna() & erx.notna() & ery.notna()\n use = need_o & ears_ok\n if use.any():\n emx = (elx + erx) / 2.0\n emy = (ely + ery) / 2.0\n ox[use] = emx[use]\n oy[use] = emy[use]\n need_o = ox.isna() | oy.isna()\n\n # neck\n neck_ok = nx.notna() & ny.notna()\n use = need_o & neck_ok\n if use.any():\n ox[use] = nx[use]\n oy[use] = ny[use]\n need_o = ox.isna() | oy.isna()\n\n # nose\n nose_ok = nozx.notna() & nozy.notna()\n use = need_o & nose_ok\n if use.any():\n ox[use] = nozx[use]\n oy[use] = nozy[use]\n\n # ---- head proxy: nose -> neck -> ear_mid -> body_center ----\n hx = pd.Series(np.nan, index=idx, dtype=float)\n hy = pd.Series(np.nan, index=idx, dtype=float)\n\n use = nose_ok\n hx[use] = nozx[use]; hy[use] = nozy[use]\n\n need_h = hx.isna() | hy.isna()\n use = need_h & neck_ok\n if use.any():\n hx[use] = nx[use]; hy[use] = ny[use]\n need_h = hx.isna() | hy.isna()\n\n use = need_h & ears_ok\n if use.any():\n emx = (elx + erx) / 2.0\n emy = (ely + ery) / 2.0\n hx[use] = emx[use]; hy[use] = emy[use]\n need_h = hx.isna() | hy.isna()\n\n use = need_h & bc_ok\n if use.any():\n hx[use] = bcx[use]; hy[use] = bcy[use]\n\n # ---- hind proxy for forward direction ----\n hipx = pd.Series(np.nan, index=idx, dtype=float)\n hipy = pd.Series(np.nan, index=idx, dtype=float)\n if hips_ok.any():\n hipx[hips_ok] = (hlx[hips_ok] + hrx[hips_ok]) / 2.0\n hipy[hips_ok] = (hly[hips_ok] + hry[hips_ok]) / 2.0\n\n # ---- forward axis: head - hip_mid -> head - origin -> COM velocity ----\n fx = pd.Series(np.nan, index=idx, dtype=float)\n fy = pd.Series(np.nan, index=idx, dtype=float)\n\n # 1) head - hip_mid\n hh_ok = hx.notna() & hy.notna() & hipx.notna() & hipy.notna()\n use = hh_ok\n if use.any():\n dx = hx[use] - hipx[use]\n dy = hy[use] - hipy[use]\n ux, uy = _safe_unit_pd(dx, dy)\n fx[use] = ux; fy[use] = uy\n\n need_f = fx.isna() | fy.isna()\n\n # 2) head - origin\n use = need_f & hx.notna() & hy.notna() & ox.notna() & oy.notna()\n if use.any():\n dx = hx[use] - ox[use]\n dy = hy[use] - oy[use]\n ux, uy = _safe_unit_pd(dx, dy)\n fx[use] = ux; fy[use] = uy\n need_f = fx.isna() | fy.isna()\n\n # 3) origin velocity as last resort\n if need_f.any():\n vx = ox.diff()\n vy = oy.diff()\n use = need_f & vx.notna() & vy.notna()\n ux, uy = _safe_unit_pd(vx, vy)\n fx[use] = ux[use]\n fy[use] = uy[use]\n\n # ---- side axis: perpendicular to forward ----\n sx = pd.Series(np.nan, index=idx, dtype=float)\n sy = pd.Series(np.nan, index=idx, dtype=float)\n ok_f = fx.notna() & fy.notna()\n if ok_f.any():\n # rotate +90°: (x,y) -> (-y,x)\n sx[ok_f] = -fy[ok_f]\n sy[ok_f] = fx[ok_f]\n\n valid = ox.notna() & oy.notna() & fx.notna() & fy.notna() & sx.notna() & sy.notna()\n return ox, oy, fx, fy, sx, sy, valid\n\n\n# ---------------------------------------------------------\n# Generic filler for one body part (tail_base / neck)\n# ---------------------------------------------------------\n\ndef _fill_part_gaps_single(\n df: pd.DataFrame,\n part: str,\n max_short_gap: int = 15,\n) -> pd.DataFrame:\n \"\"\"\n Fill gaps for a single body part using:\n 1) linear interpolation for gaps <= max_short_gap\n 2) body-frame median offsets for long gaps (when calibratable)\n 3) final ffill+bfill for any residual NaNs.\n\n If the part is not present at all, returns df unchanged.\n \"\"\"\n if not isinstance(df.columns, pd.MultiIndex):\n raise ValueError(\"Expected MultiIndex columns with levels [bodypart, coord].\")\n\n if (part, \"x\") not in df.columns or (part, \"y\") not in df.columns:\n # part not present, nothing to do\n return df\n\n out = df.copy()\n idx = out.index\n\n px = out[(part, \"x\")].astype(float)\n py = out[(part, \"y\")].astype(float)\n\n # 1) short-gap interpolation\n px_interp = _interp_short_gaps_1d(px, max_gap=max_short_gap)\n py_interp = _interp_short_gaps_1d(py, max_gap=max_short_gap)\n\n # 2) long gaps -> geometry, but only where we can calibrate\n ox, oy, fx, fy, sx, sy, valid_frame = _build_body_frame_single(out)\n\n # frames with good label & good body frame\n good_label = px.notna() & py.notna() & valid_frame\n if good_label.sum() >= 5:\n # project known part positions into body frame\n dx = px[good_label] - ox[good_label]\n dy = py[good_label] - oy[good_label]\n u = dx * fx[good_label] + dy * fy[good_label]\n v = dx * sx[good_label] + dy * sy[good_label]\n\n u_med = float(np.nanmedian(u.values))\n v_med = float(np.nanmedian(v.values))\n\n # frames with long gaps: missing in original AND still missing after short interp\n missing_long = px.isna() & px_interp.isna() & valid_frame\n\n if missing_long.any():\n # reconstruct in world coords using median offset\n ox_m = ox[missing_long]\n oy_m = oy[missing_long]\n fx_m = fx[missing_long]\n fy_m = fy[missing_long]\n sx_m = sx[missing_long]\n sy_m = sy[missing_long]\n\n rec_x = ox_m + u_med * fx_m + v_med * sx_m\n rec_y = oy_m + u_med * fy_m + v_med * sy_m\n\n # fill in reconstructed only where still NaN after short interpolation\n px_interp[missing_long] = rec_x\n py_interp[missing_long] = rec_y\n\n # 3) final nearest-neighbor fallback for any residual NaNs\n px_final = px_interp.fillna(method=\"ffill\").fillna(method=\"bfill\")\n py_final = py_interp.fillna(method=\"ffill\").fillna(method=\"bfill\")\n\n out[(part, \"x\")] = px_final\n out[(part, \"y\")] = py_final\n return out\n\n\n# ---------------------------------------------------------\n# Public helper: fill gaps for tail_base and neck\n# ---------------------------------------------------------\n\ndef fill_gaps_tail_base_and_neck_single(\n df: pd.DataFrame,\n max_short_gap: int = 15,\n) -> pd.DataFrame:\n \"\"\"\n Gap-filler for 'tail_base' and 'neck' on a single-mouse pose DataFrame.\n\n - Uses all available body parts to build a body frame:\n body_center, hips, laterals, ears, neck, nose, tail_base (if present).\n - Short gaps (≤ max_short_gap): linear interpolation.\n - Long gaps: body-frame median offset reconstruction when possible.\n - If a part is not present at all, it's skipped.\n \"\"\"\n out = df.copy()\n for part in (\"tail_base\", \"neck\"):\n out = _fill_part_gaps_single(out, part=part, max_short_gap=max_short_gap)\n return out\n\ndef fill_gaps_tail_base_and_neck_any(\n df: pd.DataFrame,\n max_short_gap: int = 15,\n) -> pd.DataFrame:\n \"\"\"\n Apply `fill_gaps_tail_base_and_neck_single` to:\n - single-mouse (2-level MultiIndex: [bodypart, coord])\n - multi-mouse / pair (>=3 levels, top level = mouse id like 'A','B').\n\n Preserves column order. If a mouse doesn't have tail_base/neck,\n the single helper just no-ops for that side.\n \"\"\"\n if not isinstance(df.columns, pd.MultiIndex):\n raise ValueError(\"Expected MultiIndex columns.\")\n\n # --- simple case: single-mouse pose (2-level MI) -------------------------\n if df.columns.nlevels == 2:\n return fill_gaps_tail_base_and_neck_single(\n df,\n max_short_gap=max_short_gap,\n )\n\n # --- multi-mouse case: split by first level (e.g. 'A','B'), process each -\n top0 = df.columns.get_level_values(0)\n mice_order = pd.Index(top0, name=df.columns.names[0]).unique()\n\n pieces = []\n for m in mice_order:\n # slice this mouse's columns → drop top level so we have [bodypart, coord]\n sub = df.xs(m, axis=1, level=0, drop_level=True)\n\n sub_filled = fill_gaps_tail_base_and_neck_single(\n sub,\n max_short_gap=max_short_gap,\n )\n\n # reattach top-level mouse label\n pieces.append(pd.concat({m: sub_filled}, axis=1))\n\n out = pd.concat(pieces, axis=1)\n\n # restore original mouse order on top level, keep sublevel order as-is\n out = out.reindex(columns=mice_order, level=0)\n out.columns.names = df.columns.names\n return out\n","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def drop_least_important_for_action(\n X: pd.DataFrame,\n prunable_by_action: pd.DataFrame,\n action: str,\n n_drop: float,\n verbose: bool = True,\n) -> tuple[pd.DataFrame, list[str]]:\n \"\"\"\n Drop a *fraction* (or count) of the least important features for a given action,\n and always drop constant / nearly constant features.\n\n Parameters\n ----------\n X : pd.DataFrame\n Feature matrix for this action.\n prunable_by_action : pd.DataFrame\n Output of build_prunable_feats_by_actionz(fi_action_agg).\n Must contain columns: ['action', 'feature', 'action_importance']\n and optionally 'useless_rank_action'.\n action : str\n Action name to prune for (must match `prunable_by_action[\"action\"]` values).\n n_drop : float\n If 0 < n_drop < 1.0 → interpreted as a *fraction* of X.columns to drop.\n If n_drop >= 1 → interpreted as an absolute maximum number of cols to drop.\n If n_drop <= 0 → only constant / nearly-constant columns are dropped.\n verbose : bool, default True\n If True, prints which columns were dropped.\n\n Returns\n -------\n X_pruned : pd.DataFrame\n Copy of X with selected columns dropped.\n dropped_feats : list[str]\n List of feature names actually dropped (subset of X.columns).\n \"\"\"\n n_cols = X.shape[1]\n if n_cols == 0:\n if verbose:\n print(f\"[PRUNE] X has 0 columns for action='{action}', nothing to drop.\")\n return X, []\n\n # ===== 1) Detect constant and nearly-constant columns (always dropped) =====\n # constant: nunique(dropna=False) <= 1\n # nearly-constant: most frequent value covers >= NEAR_CONST_THRESH of non-null values\n NEAR_CONST_THRESH = 0.995\n\n const_feats = []\n near_const_feats = []\n\n for col in X.columns:\n s = X[col]\n nunique = s.nunique(dropna=False)\n\n # all values identical (or all NaN)\n if nunique <= 1:\n const_feats.append(col)\n continue\n\n # near-constant check\n non_null = s.dropna()\n if non_null.empty:\n # all NaN; treat as constant\n const_feats.append(col)\n continue\n\n vc = non_null.value_counts()\n top_freq = vc.iloc[0]\n frac_top = top_freq / float(len(non_null))\n\n if frac_top >= NEAR_CONST_THRESH:\n near_const_feats.append(col)\n\n always_drop = list(dict.fromkeys(const_feats + near_const_feats)) # unique, preserve order\n\n # ===== 2) Decide how many *total* columns to drop (including const/near) =====\n if n_drop > 0:\n if 0.0 < n_drop < 1.0:\n # interpret as fraction of total columns\n total_to_drop = int(np.floor(n_cols * n_drop))\n else:\n total_to_drop = int(n_drop)\n total_to_drop = max(0, min(total_to_drop, n_cols))\n else:\n # only const / near-const will be dropped\n total_to_drop = 0\n\n # budget for importance-based pruning (don't count const/near-const against the fraction)\n budget_prunable = max(0, total_to_drop - len(always_drop))\n\n # If no budget for importance-based pruning, we still might drop const/near-const\n if budget_prunable <= 0:\n to_drop_prunable = []\n else:\n # ===== 3) Use importance table to drop least important features up to budget =====\n df_act = prunable_by_action[prunable_by_action[\"action\"] == action].copy()\n if df_act.empty:\n if verbose:\n print(\n f\"[PRUNE] No importance info for action='{action}', \"\n f\"only dropping const/nearly-const cols.\"\n )\n to_drop_prunable = []\n else:\n # sort by \"useless\" within this action:\n if \"useless_rank_action\" in df_act.columns:\n df_act = df_act.sort_values(\"useless_rank_action\", ascending=True)\n else:\n # fallback: lowest importance first\n df_act = df_act.sort_values(\"action_importance\", ascending=True)\n\n cand_feats = df_act[\"feature\"].tolist()\n\n cols_in_X = set(X.columns)\n already_dropped = set(always_drop)\n\n # candidate features in order of least → most important,\n # that are actually in X and not already in const/near-const list\n usable_cands = [\n f for f in cand_feats\n if (f in cols_in_X) and (f not in already_dropped)\n ]\n\n to_drop_prunable = usable_cands[:budget_prunable]\n\n # ===== 4) Final union of drops =====\n final_to_drop = list(dict.fromkeys(always_drop + to_drop_prunable)) # unique, keep order\n\n if not final_to_drop:\n if verbose:\n print(\n f\"[PRUNE] action='{action}': nothing to drop \"\n f\"(n_drop={n_drop}, const={len(const_feats)}, near={len(near_const_feats)}).\"\n )\n return X, []\n\n X_pruned = X.drop(columns=final_to_drop, errors=\"ignore\")\n\n if verbose:\n print(\n f\"[PRUNE] action='{action}': n_cols={n_cols}, \"\n f\"requested n_drop={n_drop}, target_total={total_to_drop}, \"\n f\"const={len(const_feats)}, near_const={len(near_const_feats)}, \"\n f\"importance_based={len(to_drop_prunable)}, \"\n f\"dropped_total={len(final_to_drop)}.\"\n )\n print(\" dropped:\", \", \".join(final_to_drop[:20]),\n \"...\" if len(final_to_drop) > 20 else \"\")\n\n return X_pruned, final_to_drop\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _smooth(pred, fps: float, kind: str = \"mean\", smooth_sec: float = 5/30):\n \"\"\"\n Smoothing helper with optional GPU acceleration (for mean smoothing).\n\n If CFG.use_gpu is True and CuPy is available, the 'mean' case is\n computed on GPU with semantics matching:\n\n df.rolling(window=win, min_periods=1, center=True).mean()\n\n For 'median' and 'savgol' (or if GPU is not available), we fall back\n to the original Pandas / SciPy implementations.\n \"\"\"\n def _as_df(x):\n if isinstance(x, pd.DataFrame):\n return x\n if isinstance(x, pd.Series):\n return x.to_frame()\n return pd.DataFrame(np.asarray(x))\n\n # window size (guard against weird fps/smooth_sec)\n win = max(1, int(round((smooth_sec or 0.0) * (fps or 0.0))))\n\n # --- GPU fast path for mean smoothing ----------------------------------\n use_gpu = getattr(CFG, \"use_gpu\", False)\n if kind == \"mean\" and use_gpu and win > 1:\n try:\n\n # Make a 1D float32 array on GPU\n arr = np.asarray(pred, dtype=np.float32).ravel()\n if arr.size == 0:\n return arr # nothing to do\n\n x = cp.asarray(arr) # (N,)\n n = x.shape[0]\n\n # For center=True:\n # window at i uses indices [i - left, i + right] intersect [0, n-1]\n # where:\n # left = floor((win - 1) / 2)\n # right = win - 1 - left (so total window size = win)\n left = (win - 1) // 2\n right = win - 1 - left\n\n idx = cp.arange(n, dtype=cp.int64)\n l = cp.maximum(0, idx - left)\n r = cp.minimum(n - 1, idx + right)\n\n # prefix sums with one-element padding\n cumsum = cp.cumsum(x)\n cumsum_padded = cp.concatenate(\n (cp.zeros(1, dtype=x.dtype), cumsum)\n ) # shape (n+1,)\n\n # sum over [l, r] = cumsum[r] - cumsum[l-1]\n # but with padding we use cumsum_padded[r+1] - cumsum_padded[l]\n window_sum = cumsum_padded[r + 1] - cumsum_padded[l]\n counts = (r - l + 1).astype(x.dtype) # number of elements per window\n\n out = (window_sum / counts).astype(cp.float32)\n\n # back to numpy 1D\n return cp.asnumpy(out).ravel()\n\n except Exception:\n # If anything goes wrong (no CuPy, etc.), fall back to CPU\n pass\n\n # ------------------- CPU paths below -----------------------------------\n\n if kind == \"median\":\n df = _as_df(pred)\n return (\n df.rolling(window=win, min_periods=1, center=True)\n .median()\n .to_numpy()\n .ravel()\n )\n\n if kind == \"savgol\":\n try:\n from scipy.signal import savgol_filter\n\n arr = _as_df(pred).to_numpy()\n n = arr.shape[0]\n # odd window, close to 'win', bounded by length, >=3\n w = max(3, (win // 2) * 2 + 1)\n w = min(w, n if n % 2 == 1 else max(3, n - 1))\n if w < 3 or n < 3:\n # fallback if too short for savgol\n df = _as_df(pred)\n return (\n df.rolling(window=win, min_periods=1, center=True)\n .median()\n .to_numpy()\n .ravel()\n )\n return savgol_filter(\n arr,\n window_length=w,\n polyorder=2,\n axis=0,\n mode=\"interp\"\n ).ravel()\n except Exception:\n df = _as_df(pred)\n return (\n df.rolling(window=win, min_periods=1, center=True)\n .median()\n .to_numpy()\n .ravel()\n )\n\n # default: mean on CPU\n df = _as_df(pred)\n return (\n df.rolling(window=win, min_periods=1, center=True)\n .mean()\n .to_numpy()\n .ravel()\n )\n","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def transform_pair(mouse_pair,\n body_parts_tracked,\n fps,\n meta_df=None,\n context_df=None,\n lab_id=None\n ):\n \"\"\"\n Lean pair transform: builds base inter-mouse distances then composes\n optional blocks via X = fn(X, ...) if the required body parts are present.\n \"\"\"\n # arena meta (for optional arena block)\n shape, w_cm, h_cm = _extract_arena_meta(meta_df)\n shape = str(shape).strip().lower()\n\n # ensure centers exist\n mouse_pair = fill_or_create_body_center_any(mouse_pair)\n mouse_pair = fill_or_create_nose_any(mouse_pair)\n avail_A = mouse_pair['A'].columns.get_level_values(0)\n avail_B = mouse_pair['B'].columns.get_level_values(0)\n\n # base inter-mouse squared distances across all requested part pairs\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 }, index=mouse_pair.index)\n X = X.reindex(columns=[f\"12+{p1}+{p2}\" for p1, p2 in itertools.product(body_parts_tracked, repeat=2)],\n copy=False)\n\n if not context_df.empty:\n X = add_context_pair(\n X,\n pair_df=mouse_pair, # expects ('A'/'B','body_center','x'|'y') present\n context_df=context_df, # other mice (none of 'A'/'B' inside)\n meta_df=meta_df,\n fps=fps,\n large_fill=5e2\n )\n\n # composed, gated helpers — each is no-op if inputs missing\n X = add_pursuit_features(X, mouse_pair, fps) # your pursuit block\n X = _add_pair_meta(X, meta_df)\n X = _add_speed_ear(X, mouse_pair, fps)\n X = _add_elongation_proxy(X)\n X = _add_rel_ori_and_contact(X, mouse_pair, fps)\n X = _add_approach_rate(X, mouse_pair, fps)\n # X = _add_distance_bins(X, mouse_pair)\n X = _add_temporal_interactions(X, mouse_pair, fps)\n X = _add_nose_nose(X, mouse_pair, fps)\n X = _add_val_and_intcon(X, mouse_pair, fps)\n \n if lab_id == \"BoisterousParrot\": \n X = add_shepherd_proxies_pair(X, mouse_pair, meta_df, fps=fps)\n\n\n return X.astype(np.float32, copy=False)","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def transform_single(single_mouse,\n body_parts_tracked,\n fps,\n meta_df=None,\n context_df=None,\n lab_id=None\n ):\n \"\"\"Single-mouse transform (fps-aware, arena-aware via meta_df).\"\"\"\n single_mouse = fill_or_create_body_center_any(single_mouse)\n single_mouse = fill_or_create_nose_any(single_mouse)\n available_body_parts = single_mouse.columns.get_level_values(0)\n\n # --- base squared distances across parts ---\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 not context_df.empty:\n X = add_context_single(\n X,\n focal_df=single_mouse, # expects ('A'/'B','body_center','x'|'y') present\n context_df=context_df, # other mice (none of 'A'/'B' inside)\n meta_df=meta_df,\n fps=fps,\n large_fill=5e2\n )\n\n # --- speed-like via lagged displacements ---\n if all(p in single_mouse.columns for p in ['ear_left', 'ear_right', 'tail_base']):\n lag = _scale(10, fps)\n sh = single_mouse[['ear_left','ear_right','tail_base']].shift(lag)\n X = pd.concat([X, pd.DataFrame({\n 'sp_lf': np.square(single_mouse['ear_left'] - sh['ear_left']).sum(axis=1, skipna=False),\n 'sp_rt': np.square(single_mouse['ear_right'] - sh['ear_right']).sum(axis=1, skipna=False),\n 'sp_lf2': np.square(single_mouse['ear_left'] - sh['tail_base']).sum(axis=1, skipna=False),\n 'sp_rt2': np.square(single_mouse['ear_right'] - sh['tail_base']).sum(axis=1, skipna=False),\n })], axis=1)\n\n # elongation proxy\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 # body angle\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 den = (np.sqrt(v1['x']**2 + v1['y']**2) * np.sqrt(v2['x']**2 + v2['y']**2) + 1e-6)\n X['body_ang'] = (v1['x'] * v2['x'] + v1['y'] * v2['y']) / den\n\n # core temporal (fps-scaled)\n if 'body_center' in available_body_parts:\n cx = single_mouse['body_center']['x']; cy = single_mouse['body_center']['y']\n for w in [5, 15, 30, 60]:\n ws = _scale(w, fps); 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, 1).sum()**2 + cy.diff().rolling(ws, 1).sum()**2)\n X[f'act{w}'] = np.sqrt(cx.diff().rolling(ws, 1).var() + cy.diff().rolling(ws, 1).var())\n\n # your existing advanced center features\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 # arena-aware behavioral proxies you already had\n X = add_single_stillness_and_headjitter(X, single_mouse, fps)\n X = add_groom_microfeatures(X, single_mouse, fps)\n X = add_verticality_proxies_single(X, single_mouse, meta_df=meta_df, fps=fps)\n if lab_id in [\"InvincibleJellyfish\", \"GroovyShrew\",\"NiftyGoldfinch\"]:\n X = add_dig_proxies_single(X, single_mouse, meta_df=meta_df, fps=fps)\n X = add_tortuosity_and_dir_persist_single(X, cx, cy, fps)\n X = add_dir_persist_single(X, cx, cy, fps)\n\n # # # # --- acausal future-aware features (keep windows fps-scaled and modest) ---\n X = add_speed_asymmetry_future_past_single(X, cx, cy, fps, horizon_base=30) # Δv future - past (~1s)\n # X = add_time_reversal_asymmetry3_single(X, cx, cy, fps, h_base=2, smooth_base=11) # 3rd-order asym (tiny h)\n \n # # # # # distribution shift of speed (Gaussian symmetric-KL; cheap stand-in for MMD/ED)\n X = add_gauss_shift_speed_future_past_single(X, cx, cy, fps, window_base=30)\n \n # # # # # spectral & complexity cues (use small windows to keep cost low)\n # X = add_spectral_flux_future_single(X, cx, cy, fps, window_base=64) # FFT-flux future vs past\n # X = add_katz_fd_future_speed_single(X, cx, cy, fps, horizon_base=30) # future-only FD proxy\n # X = add_sampen_future_speed_single(X, cx, cy, fps, horizon_base=30, m=2, r_frac=0.2)# future-only SampEn\n \n # # # # # smooth, acausal kinematic derivatives at center\n # X = add_savgol_center_derivs_single(X, cx, cy, fps, win_base=11, poly=2)\n\n # 1) head orientation & micromotion\n if 'nose' in available_body_parts:\n X = add_head_orientation_tree_feats_single(X, single_mouse, fps, cx=cx, cy=cy)\n X = add_head_micromove_spectrum_single(X, single_mouse, cx, cy, fps, available_body_parts)\n\n # nose-tail dynamics\n if all(p in available_body_parts for p in ['nose', 'tail_base']):\n nt = 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 lg in [10, 20, 40]:\n L = _scale(lg, fps)\n X[f'nt_lg{lg}'] = nt.shift(L)\n X[f'nt_df{lg}'] = nt - nt.shift(L)\n\n # ears\n if all(p in available_body_parts for p in ['ear_left', 'ear_right']):\n ed = 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 [-20, -10, 10, 20]:\n o = _scale_signed(off, fps)\n X[f'ear_o{off}'] = ed.shift(-o)\n w = _scale(30, fps)\n X['ear_con'] = ed.rolling(w, min_periods=1, center=True).std() / (ed.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n\n # hygiene\n X.replace([np.inf, -np.inf], np.nan, inplace=True)\n if 'tortuosity_1s' in X:\n X['tortuosity_1s'] = X['tortuosity_1s'].clip(0, 50)\n if 'dir_persist_1s' in X:\n X['dir_persist_1s'] = X['dir_persist_1s'].clip(0, 1)\n\n return X.astype(np.float32, copy=False)\n","metadata":{"jupyter":{"source_hidden":true},"papermill":{"duration":0.051971,"end_time":"2025-10-25T19:26:58.497403","exception":false,"start_time":"2025-10-25T19:26:58.445432","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def limit_rows_stratified(X, y, max_rows=None, seed=42, verbose=True):\n if (max_rows is None) or (len(y) <= max_rows):\n return X, y\n y_arr = np.asarray(y); n = len(y_arr)\n\n if len(np.unique(y_arr)) < 2:\n rng = np.random.RandomState(seed)\n idx = rng.choice(n, max_rows, replace=False)\n else:\n sss = StratifiedShuffleSplit(n_splits=1, train_size=max_rows, random_state=seed)\n idx, _ = next(sss.split(np.zeros(n), y_arr))\n\n idx = np.sort(idx) # <- preserve original temporal order\n\n X_sub = X.iloc[idx] if hasattr(X, \"iloc\") else X[idx]\n y_sub = y.iloc[idx].to_numpy() if hasattr(y, \"iloc\") else y_arr[idx]\n if verbose:\n print(f\"\\t[limit] {n}→{len(idx)} rows (pos {100*y_arr.mean():.1f}%→{100*y_sub.mean():.1f}%)\")\n return X_sub, y_sub\n\ndef decide_max_rows_for_xgb(X, model, ram_gb=30, util=0.7):\n p = getattr(model, \"get_params\", model.get_xgb_params)() if hasattr(model, \"get_xgb_params\") else model.get_params()\n tm = str(p.get(\"tree_method\", \"hist\"))\n max_depth = int(p.get(\"max_depth\", 6))\n n_cols = int(X.shape[1])\n\n # rough bytes/row (float32) for DMatrix + grads/hess + work buffers\n if tm == \"gpu_hist\":\n per_row = (8 * n_cols) # ~ data(4) + overhead\n fixed = int(0.5 * 1024**3) # ~0.5GB misc host-side\n else:\n per_row = ((12 + 2*max_depth) * n_cols) # deeper trees → more workspace\n fixed = int(1.0 * 1024**3) # ~1GB misc\n\n budget = int(ram_gb * 1024**3 * util) - fixed\n if budget <= 0:\n return min(len(X), 500_000)\n\n cap = budget // per_row\n return None if cap >= len(X) else int(cap)\n","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training, validation and submission","metadata":{"papermill":{"duration":0.006521,"end_time":"2025-10-25T19:26:58.511707","exception":false,"start_time":"2025-10-25T19:26:58.505186","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# ========================= GATING HELPERS ===========================\n\nfrom collections import defaultdict\n\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import OneHotEncoder, StandardScaler\nfrom sklearn.compose import ColumnTransformer\nfrom sklearn.pipeline import Pipeline\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.impute import SimpleImputer\nfrom sklearn.ensemble import RandomForestClassifier, ExtraTreesClassifier\nfrom sklearn.base import clone\n\n_XGB_GPU_KWARGS = {\"tree_method\": \"hist\", \"device\": \"cuda\"}\n\n# ------------------------ GATE CONSTANTS -------------------------\n# \"Almost never happens\" threshold, in FRACTION of frames.\n# 1e-3 ≈ 0.1% of frames.\nGATE_MIN_POS_RATE = 5e-4\n\n# Require a minimum number of clips in each class\nGATE_MIN_NEG_COUNT = 1\nGATE_MIN_POS_COUNT = 1\n\n# Only keep gate models with clearly better-than-random AUC\nGATE_MIN_AUC = 0.7\n# Gating behaviour:\n# - only DOWN-WEIGHT clips with gate_prob < GATE_NEG_CLIP_MAX_PROB\n# - others are treated as neutral (gate_prob ~ 0.5, no effect)\nGATE_NEG_CLIP_MAX_PROB = 0.20 # only g < 0.20 are candidates for gating\n\n# Within that \"negative\" region, we distinguish:\n# gate_prob < GATE_STRONG_PROB -> strong gating\n# GATE_STRONG_PROB <= gate_prob < GATE_NEG_CLIP_MAX_PROB -> soft gating\nGATE_STRONG_PROB = 0.10 # very unlikely region\n\n# Base logit-space strength λ. The actual per-frame strength is adapted\n# based on gate_prob (soft vs strong) around this base value.\nGATE_STRENGTH = 2.0 # logit-space base strength λ\n\n\n# Clip-level feature names (used in multiple places)\nCLIP_FEAT_COLS = [\n \"clip_speed_mean_norm\",\n \"clip_speed_p90_norm\",\n \"clip_speed_std_norm\",\n \"clip_speed_iqr_norm\",\n \"clip_speed_high_frac\",\n \"clip_speed_regime_shift_norm\",\n \"clip_speed_entropy\",\n \"clip_freq_dom_hz\",\n \"clip_freq_power_low_frac\",\n \"clip_freq_power_high_frac\",\n \"clip_closeness_med_norm\",\n \"clip_closeness_entropy\",\n]\n\n# Cache for per-(lab, switch) video-level tables across ALL actions\n# key: (lab_id, switch) -> dict[action_name -> video_tab DataFrame]\nVIDEO_TAB_CACHE: dict[tuple[str, str], dict[str, pd.DataFrame]] = {}\n\n\n# ------------------------------------------------------------\n# Helper: which actions are actually labeled in this clip?\n# ------------------------------------------------------------\ndef _get_labeled_actions_from_meta(meta_df: pd.DataFrame, y_cols) -> set[str]:\n \"\"\"\n Return the set of actions that are *actually labeled* for this clip/video.\n\n If `behaviors_labeled` is missing or empty, we fall back to treating all\n y_cols as labeled (to avoid silently dropping entire labs).\n \"\"\"\n if meta_df is None or meta_df.empty or \"behaviors_labeled\" not in meta_df.columns:\n return set(y_cols)\n\n raw = meta_df[\"behaviors_labeled\"].iloc[0]\n\n # String case: e.g. \"attack sniff mount\" or \"attack,sniff|mount\"\n if isinstance(raw, str):\n raw = raw.strip()\n if not raw:\n return set(y_cols)\n for ch in [\",\", \";\", \"|\"]:\n raw = raw.replace(ch, \" \")\n tokens = raw.split()\n return set(t for t in tokens if t)\n\n # Iterable case: list/tuple/set\n if isinstance(raw, (list, tuple, set)):\n return set(map(str, raw))\n\n # Anything else (NaN, weird type) -> treat as \"all are labeled\"\n return set(y_cols)\n\n\n# ------------------------------------------------------------\n# 0) Clip-level activity features from data_df / ctx_df\n# ------------------------------------------------------------\ndef compute_clip_activity_features(\n data_df: pd.DataFrame | None,\n ctx_df: pd.DataFrame | None = None,\n fps: float = 30.0,\n arena_width_cm: float | None = None,\n arena_height_cm: float | None = None,\n) -> dict:\n \"\"\"\n Compute a compact set of clip-level features (time/freq/closeness).\n \"\"\"\n\n feat_names = CLIP_FEAT_COLS\n feats = {k: np.nan for k in feat_names}\n\n # ---- arena scale (diagonal) for normalisation ----\n w = float(arena_width_cm) if arena_width_cm is not None else np.nan\n h = float(arena_height_cm) if arena_height_cm is not None else np.nan\n if np.isfinite(w) and np.isfinite(h) and (w > 0) and (h > 0):\n arena_scale = float(np.sqrt(w * w + h * h))\n else:\n arena_scale = 1.0 # fallback; effectively \"no normalisation\"\n\n fps = float(fps) if np.isfinite(fps) and fps > 0 else 30.0\n\n # ------------------------------------------------------------------\n # 0.1) Time-domain motion: build a 1D \"speed_norm\" from data_df\n # ------------------------------------------------------------------\n if data_df is None or data_df.empty:\n speed_norm = None\n else:\n num_data = data_df.select_dtypes(include=[np.number])\n if num_data.shape[0] < 2 or num_data.shape[1] == 0:\n speed_norm = None\n else:\n arr = num_data.to_numpy(dtype=np.float32, copy=False)\n diff = np.diff(arr, axis=0) # (T-1, F)\n frame_motion = np.mean(np.abs(diff), axis=1) # (T-1,)\n\n if frame_motion.size == 0 or not np.isfinite(frame_motion).any():\n speed_norm = None\n else:\n frame_speed = frame_motion * fps\n speed_norm = frame_speed / max(arena_scale, 1e-3)\n\n if speed_norm is not None:\n sn = np.asarray(speed_norm, dtype=np.float32)\n if np.isfinite(sn).any():\n m = float(np.nanmean(sn))\n p90 = float(np.nanquantile(sn, 0.90))\n s = float(np.nanstd(sn))\n\n feats[\"clip_speed_mean_norm\"] = m\n feats[\"clip_speed_p90_norm\"] = p90\n feats[\"clip_speed_std_norm\"] = s\n\n # robust spread\n with np.errstate(all=\"ignore\"):\n q25 = float(np.nanquantile(sn, 0.25))\n q75 = float(np.nanquantile(sn, 0.75))\n feats[\"clip_speed_iqr_norm\"] = q75 - q25\n\n hi_thr = m + s\n if np.isfinite(hi_thr) and hi_thr > 0:\n feats[\"clip_speed_high_frac\"] = float(np.mean(sn > hi_thr))\n else:\n feats[\"clip_speed_high_frac\"] = 0.0\n\n # -------------- Regime-shift feature --------------\n n = sn.size\n if n >= 4 and np.isfinite(s) and s > 1e-6:\n idx1 = max(1, n // 3)\n idx2 = max(idx1 + 1, (2 * n) // 3)\n early = sn[:idx1]\n late = sn[idx2:]\n if early.size > 0 and late.size > 0:\n m_early = float(np.nanmean(early))\n m_late = float(np.nanmean(late))\n feats[\"clip_speed_regime_shift_norm\"] = abs(m_late - m_early) / s\n\n # -------------- Speed entropy --------------\n valid_sn = sn[np.isfinite(sn)]\n if valid_sn.size > 0:\n hist, _ = np.histogram(valid_sn, bins=10)\n total = hist.sum()\n if total > 0:\n p = hist.astype(np.float64) / float(total)\n p = p[p > 0.0]\n if p.size > 0:\n feats[\"clip_speed_entropy\"] = float(\n -(p * np.log(p + 1e-12)).sum()\n )\n\n # ----------------- Frequency-domain on speed_norm -----------------\n if n >= 4:\n sig = sn - np.nanmean(sn)\n sig = np.nan_to_num(sig, nan=0.0, posinf=0.0, neginf=0.0)\n\n spec = np.fft.rfft(sig)\n psd = (spec * np.conj(spec)).real\n freqs = np.fft.rfftfreq(n, d=1.0 / fps)\n\n total_power = float(psd.sum())\n if total_power > 0.0:\n # dominant frequency (skip DC if possible)\n if psd.size > 1:\n dom_idx = int(np.argmax(psd[1:])) + 1\n else:\n dom_idx = 0\n feats[\"clip_freq_dom_hz\"] = float(freqs[dom_idx])\n\n low_mask = freqs <= 1.0\n high_mask = freqs > 3.0\n\n if np.any(low_mask):\n feats[\"clip_freq_power_low_frac\"] = float(\n psd[low_mask].sum() / total_power\n )\n if np.any(high_mask):\n feats[\"clip_freq_power_high_frac\"] = float(\n psd[high_mask].sum() / total_power\n )\n\n # ------------------------------------------------------------------\n # 0.2) Closeness: agent-target distance normalised by arena size\n # ------------------------------------------------------------------\n def _extract_agent_target_xy(df: pd.DataFrame | None):\n if df is None or df.empty:\n return None\n\n cols = list(df.columns)\n cols_str = [str(c).lower() for c in cols]\n\n numeric_mask = [\n pd.api.types.is_numeric_dtype(df[c]) for c in cols\n ]\n num_idxs = [i for i, is_num in enumerate(numeric_mask) if is_num]\n if not num_idxs:\n return None\n\n def _find_xy(role_token: str):\n role_idxs = [\n i for i in num_idxs if role_token in cols_str[i]\n ]\n if not role_idxs:\n return None, None\n\n x_cands = [cols[i] for i in role_idxs if \"x\" in cols_str[i]]\n y_cands = [cols[i] for i in role_idxs if \"y\" in cols_str[i]]\n\n if not x_cands or not y_cands:\n return None, None\n return x_cands[0], y_cands[0]\n\n ax_col, ay_col = _find_xy(\"agent\")\n tx_col, ty_col = _find_xy(\"target\")\n\n if None in (ax_col, ay_col, tx_col, ty_col):\n return None\n\n ax = pd.to_numeric(df[ax_col], errors=\"coerce\").to_numpy(dtype=np.float32)\n ay = pd.to_numeric(df[ay_col], errors=\"coerce\").to_numpy(dtype=np.float32)\n tx = pd.to_numeric(df[tx_col], errors=\"coerce\").to_numpy(dtype=np.float32)\n ty = pd.to_numeric(df[ty_col], errors=\"coerce\").to_numpy(dtype=np.float32)\n\n n_min = min(ax.shape[0], ay.shape[0], tx.shape[0], ty.shape[0])\n if n_min < 1:\n return None\n\n ax, ay, tx, ty = ax[:n_min], ay[:n_min], tx[:n_min], ty[:n_min]\n return ax, ay, tx, ty\n\n xy = _extract_agent_target_xy(ctx_df)\n if xy is None:\n xy = _extract_agent_target_xy(data_df)\n\n if xy is not None:\n ax, ay, tx, ty = xy\n dist = np.sqrt((ax - tx) ** 2 + (ay - ty) ** 2)\n if dist.size > 0 and np.isfinite(dist).any():\n if arena_scale > 0:\n dist_norm = dist / arena_scale\n else:\n dist_norm = dist\n dist_norm = np.clip(dist_norm, 0.0, 1e3)\n\n dist_clip = np.clip(dist_norm, 0.0, 1.0)\n closeness = 1.0 - dist_clip\n\n med = float(np.nanmedian(closeness))\n feats[\"clip_closeness_med_norm\"] = med\n\n valid_cl = closeness[np.isfinite(closeness)]\n if valid_cl.size > 0:\n hist, _ = np.histogram(valid_cl, bins=10, range=(0.0, 1.0))\n total = hist.sum()\n if total > 0:\n p = hist.astype(np.float64) / float(total)\n p = p[p > 0.0]\n if p.size > 0:\n feats[\"clip_closeness_entropy\"] = float(\n -(p * np.log(p + 1e-12)).sum()\n )\n\n return feats\n\n\n# ------------------------------------------------------------\n# 1) Build video-level tables for ALL actions (single pass)\n# ------------------------------------------------------------\ndef _build_video_level_tables_for_lab_switch(\n lab_id: str,\n switch: str, # \"single\" or \"pair\"\n train_df: pd.DataFrame,\n train_tracking_dir: str = CFG.train_tracking_path,\n) -> dict[str, pd.DataFrame]:\n \"\"\"\n Single heavy pass over tracking for a given (lab_id, switch).\n\n Returns:\n dict[action_name -> video_tab_for_that_action]\n\n IMPORTANT:\n A clip contributes to action `A` *only if* `A` is present\n in its `behaviors_labeled` for that video.\n \"\"\"\n assert switch in {\"single\", \"pair\"}\n\n dataset = train_df[train_df.lab_id == lab_id].copy()\n rows_in_lab = len(dataset)\n if dataset.empty:\n raise ValueError(f\"No rows in train for lab_id={lab_id}\")\n\n # video-level meta lookup for fps if needed\n if \"frames_per_second\" in train_df.columns:\n _fps_lookup = (\n train_df[[\"video_id\", \"frames_per_second\"]]\n .drop_duplicates(\"video_id\")\n .set_index(\"video_id\")[\"frames_per_second\"]\n .to_dict()\n )\n else:\n _fps_lookup = {}\n\n gen = generate_mouse_data(\n dataset,\n \"train\",\n traintest_directory=train_tracking_dir,\n generate_single=(switch == \"single\"),\n generate_pair=(switch == \"pair\"),\n )\n\n clip_rows: dict[tuple, dict] = {}\n pos_counts: dict[tuple, dict[str, int]] = defaultdict(lambda: defaultdict(int))\n all_actions: set[str] = set()\n\n # NEW: which actions are actually labeled per (video_id, agent_id, target_id)\n labeled_actions_per_key: dict[tuple, set[str]] = defaultdict(set)\n\n for sw, data_df, ctx_df, meta_df, y_df in gen:\n if sw != switch:\n del data_df, ctx_df, meta_df, y_df\n gc.collect()\n continue\n\n y_cols = [c for c in y_df.columns if pd.api.types.is_numeric_dtype(y_df[c])]\n if not y_cols:\n del data_df, ctx_df, meta_df, y_df\n gc.collect()\n continue\n all_actions.update(y_cols)\n\n video_id = meta_df[\"video_id\"].iloc[0]\n agent_id = meta_df[\"agent_id\"].iloc[0]\n target_id = meta_df[\"target_id\"].iloc[0]\n key = (video_id, agent_id, target_id)\n\n # record which actions are labeled in this clip\n labeled_here = _get_labeled_actions_from_meta(meta_df, y_cols)\n if labeled_here:\n labeled_actions_per_key[key].update(labeled_here)\n else:\n # ultra-defensive: fall back to all y_cols\n labeled_actions_per_key[key].update(y_cols)\n\n fps = 30.0\n try:\n if \"frames_per_second\" in meta_df.columns:\n fps_val = meta_df[\"frames_per_second\"].iloc[0]\n else:\n fps_val = _fps_lookup.get(video_id, 30.0)\n fps_val = pd.to_numeric(fps_val, errors=\"coerce\")\n fps = float(fps_val) if np.isfinite(fps_val) and fps_val > 0 else 30.0\n except Exception:\n fps = 30.0\n\n def _safe_scalar(col, default=np.nan):\n if col not in meta_df.columns:\n return default\n v = pd.to_numeric(meta_df[col].iloc[0], errors=\"coerce\")\n return float(v) if np.isfinite(v) else default\n\n arena_w = _safe_scalar(\"arena_width_cm\")\n arena_h = _safe_scalar(\"arena_height_cm\")\n\n if key not in clip_rows:\n clip_feats = compute_clip_activity_features(\n data_df=data_df,\n ctx_df=ctx_df,\n fps=fps,\n arena_width_cm=arena_w,\n arena_height_cm=arena_h,\n )\n row = {\n \"video_id\": video_id,\n \"agent_id\": agent_id,\n \"target_id\": target_id,\n \"n_frames\": len(meta_df),\n }\n row.update(clip_feats)\n clip_rows[key] = row\n else:\n clip_rows[key][\"n_frames\"] += len(meta_df)\n\n y_int = y_df[y_cols].astype(int)\n counts = y_int.sum(axis=0)\n for act_name, n_pos_chunk in counts.items():\n pos_counts[key][act_name] += int(n_pos_chunk)\n\n del data_df, ctx_df, meta_df, y_df\n gc.collect()\n\n if not clip_rows:\n print(\n f\"[build_video_level_tables_for_all_actions] lab={lab_id}, switch={switch} | \"\n \"no frames where any action is even defined.\"\n )\n return {}\n\n base_df = pd.DataFrame(list(clip_rows.values()))\n\n for c in CLIP_FEAT_COLS:\n if c not in base_df.columns:\n base_df[c] = np.nan\n\n base_df = base_df[\n [\"video_id\", \"agent_id\", \"target_id\", \"n_frames\"] + CLIP_FEAT_COLS\n ]\n\n print(\n f\"[build_video_level_tables_for_all_actions] lab={lab_id}, switch={switch} \"\n f\"| rows_in_lab={rows_in_lab}, n_clips={len(base_df)}, n_actions_seen={len(all_actions)}\"\n )\n\n video_tabs: dict[str, pd.DataFrame] = {}\n keys_list = list(clip_rows.keys())\n\n for act in sorted(all_actions):\n # keep only clips where this action is *actually labeled*\n mask = np.array(\n [act in labeled_actions_per_key.get(k, set()) for k in keys_list],\n dtype=bool,\n )\n if not mask.any():\n # action is never labeled in this lab/switch\n continue\n\n vt = base_df.loc[mask].copy()\n\n n_pos_all = [pos_counts[k].get(act, 0) for k in keys_list]\n vt[\"n_pos\"] = np.asarray(n_pos_all, dtype=np.int64)[mask]\n vt[\"n_frames\"] = vt[\"n_frames\"].clip(lower=1)\n vt[\"pos_rate\"] = vt[\"n_pos\"] / vt[\"n_frames\"]\n\n vt = vt[\n [\"video_id\", \"agent_id\", \"target_id\", \"n_frames\", \"n_pos\", \"pos_rate\"]\n + CLIP_FEAT_COLS\n ]\n video_tabs[act] = vt\n\n return video_tabs\n\n\n# ------------------------------------------------------------\n# 1) Build video-level table for one (lab_id, action, switch)\n# (now backed by cached ALL-actions tables)\n# ------------------------------------------------------------\ndef build_video_level_table_for_action(\n lab_id: str,\n action: str,\n switch: str, # \"single\" or \"pair\"\n train_df: pd.DataFrame,\n train_tracking_dir: str = CFG.train_tracking_path,\n) -> pd.DataFrame:\n \"\"\"\n For a given lab_id + action + switch:\n\n - Uses a cached, single-pass video-level table over ALL actions\n for this (lab_id, switch).\n - Returns the per-action video_tab with:\n n_frames, n_pos, pos_rate\n + clip-level activity/closeness features.\n \"\"\"\n assert switch in {\"single\", \"pair\"}\n\n cache_key = (str(lab_id), str(switch))\n if cache_key not in VIDEO_TAB_CACHE:\n VIDEO_TAB_CACHE[cache_key] = _build_video_level_tables_for_lab_switch(\n lab_id=lab_id,\n switch=switch,\n train_df=train_df,\n train_tracking_dir=train_tracking_dir,\n )\n\n video_tabs_all = VIDEO_TAB_CACHE[cache_key]\n\n if action not in video_tabs_all:\n print(\n f\"[build_video_level_table_for_action] lab={lab_id}, action={action}, switch={switch} | \"\n \"no frames where this action is even defined.\"\n )\n return pd.DataFrame(\n columns=[\n \"video_id\",\n \"agent_id\",\n \"target_id\",\n \"n_frames\",\n \"n_pos\",\n \"pos_rate\",\n ] + CLIP_FEAT_COLS\n )\n\n return video_tabs_all[action].copy()\n\n\n# ------------------------------------------------------------\n# 2) Join train metadata onto that video-level table\n# ------------------------------------------------------------\ndef add_video_metadata(video_tab: pd.DataFrame, train_df: pd.DataFrame) -> pd.DataFrame:\n \"\"\"\n Adds per-video metadata (sex, strain, condition, arena, etc.)\n to the video-level action table and computes arena_area_cm2.\n \"\"\"\n if video_tab.empty:\n return video_tab\n\n meta_cols = [\n \"video_id\",\n \"video_duration_sec\", \"frames_per_second\",\n \"arena_type\", \"arena_shape\",\n \"mouse1_strain\", \"mouse2_strain\", \"mouse3_strain\", \"mouse4_strain\",\n \"mouse1_condition\", \"mouse2_condition\",\n \"mouse1_sex\", \"mouse2_sex\",\n \"arena_width_cm\", \"arena_height_cm\",\n \"n_mice\",\n \"tracking_method\",\n ]\n\n train_meta = (\n train_df[meta_cols]\n .drop_duplicates(\"video_id\")\n )\n\n merged = video_tab.merge(train_meta, on=\"video_id\", how=\"left\")\n\n if \"arena_width_cm\" in merged.columns and \"arena_height_cm\" in merged.columns:\n with np.errstate(all=\"ignore\"):\n merged[\"arena_area_cm2\"] = (\n pd.to_numeric(merged[\"arena_width_cm\"], errors=\"coerce\") *\n pd.to_numeric(merged[\"arena_height_cm\"], errors=\"coerce\")\n )\n else:\n merged[\"arena_area_cm2\"] = np.nan\n\n return merged\n\n\n# ------------------------------------------------------------\n# 3) Train a video-level gate: \"high pos_rate\" vs \"never / tiny\"\n# ------------------------------------------------------------\ndef train_action_presence_model(\n video_tab: pd.DataFrame,\n lab_id: str,\n action: str,\n switch: str,\n min_pos_rate: float = GATE_MIN_POS_RATE,\n min_auc: float = GATE_MIN_AUC,\n):\n \"\"\"\n Trains a simple video-level model to discriminate clips\n that have a *meaningful* amount of `action` vs clips that\n effectively never have it.\n\n Label definition (soft):\n has_action = 1 if pos_rate >= min_pos_rate\n = 0 otherwise\n\n Safety nets:\n - requires both classes present\n - requires each class to have >= 5% of examples (5% rule)\n - requires AUC >= min_auc to be considered usable\n \"\"\"\n if video_tab.empty:\n print(f\"[Gate model] No video rows for lab={lab_id}, action={action}, switch={switch}.\")\n return None, None, None, None\n\n vt = video_tab.copy()\n\n if \"arena_area_cm2\" not in vt.columns:\n if \"arena_width_cm\" in vt.columns and \"arena_height_cm\" in vt.columns:\n with np.errstate(all=\"ignore\"):\n vt[\"arena_area_cm2\"] = (\n pd.to_numeric(vt[\"arena_width_cm\"], errors=\"coerce\") *\n pd.to_numeric(vt[\"arena_height_cm\"], errors=\"coerce\")\n )\n else:\n vt[\"arena_area_cm2\"] = np.nan\n\n # --- define has_action via GATE_MIN_POS_RATE, NOT pure 0-pos ---\n vt[\"has_action\"] = (vt[\"pos_rate\"] >= float(min_pos_rate)).astype(int)\n y = vt[\"has_action\"].values\n\n base_num_cols = [\n \"n_frames\",\n \"n_mice\",\n \"arena_area_cm2\",\n \"video_duration_sec\",\n \"frames_per_second\",\n ]\n activity_cols = CLIP_FEAT_COLS\n num_cols = base_num_cols + activity_cols\n\n cat_cols = [\n \"arena_type\", \"arena_shape\",\n \"mouse1_strain\", \"mouse2_strain\", \"mouse3_strain\", \"mouse4_strain\",\n \"mouse1_condition\", \"mouse2_condition\",\n \"mouse1_sex\", \"mouse2_sex\",\n \"tracking_method\",\n ]\n\n X = vt[num_cols + cat_cols].copy()\n\n for c in num_cols:\n X[c] = pd.to_numeric(X[c], errors=\"coerce\")\n\n for c in cat_cols:\n X[c] = X[c].astype(\"category\")\n X[c] = X[c].cat.add_categories([\"\"]).fillna(\"\")\n\n uniq, counts = np.unique(y, return_counts=True)\n class_dist = dict(zip(uniq, counts))\n n_neg = class_dist.get(0, 0)\n n_pos = class_dist.get(1, 0)\n total = int(n_neg + n_pos) if (n_neg is not None and n_pos is not None) else 0\n frac_pos = (n_pos / total) if total > 0 else 0.0\n frac_neg = (n_neg / total) if total > 0 else 0.0\n\n # -------- safety net 1: both classes present --------\n if len(uniq) < 2:\n print(\n f\"[Gate model] lab={lab_id}, action={action}, switch={switch} | \"\n f\"only one class; skip. class_dist={class_dist}\"\n )\n return None, None, None, None\n\n # -------- safety net 2: 5% rule on has_action label --------\n MIN_FRAC = 0.75\n if (frac_pos < MIN_FRAC) or (frac_neg < MIN_FRAC):\n print(\n f\"[Gate model] lab={lab_id}, action={action}, switch={switch} | \"\n f\"class fractions too extreme (frac_pos={frac_pos:.3f}, frac_neg={frac_neg:.3f}); \"\n f\"skip gate.\"\n )\n return None, None, None, {\n \"class_dist\": class_dist,\n \"frac_pos\": frac_pos,\n \"frac_neg\": frac_neg,\n }\n\n # -------- safety net 3: tiny absolute class count --------\n if min(counts) < 2:\n print(\n f\"[Gate model] lab={lab_id}, action={action}, switch={switch} | \"\n f\"min class count {min(counts)} < 2; skip gate.\"\n )\n return None, None, None, {\"class_dist\": class_dist}\n\n # ----------------- rest of the function unchanged -----------------\n X_train, X_val, y_train, y_val = train_test_split(\n X, y, test_size=0.3, stratify=y, random_state=42\n )\n\n num_pipe = Pipeline(\n steps=[\n (\"impute\", SimpleImputer(strategy=\"median\")),\n (\"scale\", StandardScaler()),\n ]\n )\n cat_pipe = OneHotEncoder(handle_unknown=\"ignore\")\n\n pre_template = ColumnTransformer(\n transformers=[\n (\"num\", num_pipe, num_cols),\n (\"cat\", cat_pipe, cat_cols),\n ]\n )\n\n candidate_clfs = {\n \"logreg\": LogisticRegression(\n max_iter=300,\n class_weight=\"balanced\",\n ),\n \"logreg_C0_5\": LogisticRegression(\n max_iter=300,\n class_weight=\"balanced\",\n C=0.5,\n ),\n \"rf\": RandomForestClassifier(\n n_estimators=200,\n max_depth=4,\n min_samples_leaf=2,\n class_weight=\"balanced\",\n random_state=42,\n n_jobs=-1,\n ),\n \"rf_deep\": RandomForestClassifier(\n n_estimators=300,\n max_depth=None,\n min_samples_leaf=2,\n class_weight=\"balanced_subsample\",\n random_state=42,\n n_jobs=-1,\n ),\n \"et\": ExtraTreesClassifier(\n n_estimators=300,\n max_depth=None,\n min_samples_leaf=2,\n class_weight=\"balanced_subsample\",\n random_state=42,\n n_jobs=-1,\n ),\n }\n\n candidate_clfs[\"xgb_gpu\"] = xgb.XGBClassifier(\n n_estimators=400,\n max_depth=4,\n learning_rate=0.05,\n subsample=0.8,\n colsample_bytree=0.8,\n objective=\"binary:logistic\",\n eval_metric=\"auc\",\n **_XGB_GPU_KWARGS,\n )\n\n best_auc = -np.inf\n best_name = None\n best_pipe = None\n\n for name, clf in candidate_clfs.items():\n pipe = Pipeline(\n steps=[\n (\"pre\", clone(pre_template)),\n (\"clf\", clf),\n ]\n )\n try:\n pipe.fit(X_train, y_train)\n y_val_proba = pipe.predict_proba(X_val)[:, 1]\n auc = roc_auc_score(y_val, y_val_proba)\n except Exception as e:\n print(\n f\"[Gate model] lab={lab_id}, action={action}, switch={switch} | \"\n f\"model={name} failed with error: {e!r}; skipping this model.\"\n )\n continue\n\n if auc > best_auc:\n best_auc = auc\n best_name = name\n best_pipe = pipe\n\n if best_pipe is None or best_auc < float(min_auc):\n print(\n f\"[Gate model] lab={lab_id}, action={action}, switch={switch} | \"\n f\"best AUC={best_auc:.3f} < {min_auc:.3f}; gate skipped. class_dist={class_dist}\"\n )\n return None, None, None, {\n \"auc\": best_auc,\n \"class_dist\": class_dist,\n \"frac_pos\": frac_pos,\n \"frac_neg\": frac_neg,\n }\n\n best_pipe.fit(X, y)\n\n metrics = {\n \"auc\": float(best_auc),\n \"class_dist\": class_dist,\n \"frac_pos\": float(frac_pos),\n \"min_pos_rate\": float(min_pos_rate),\n \"model_type\": best_name,\n }\n\n print(\n f\"[Gate model] lab={lab_id}, action={action}, switch={switch} | \"\n f\"model={best_name}, AUC={best_auc:.3f}, class_dist={class_dist}, frac_pos={frac_pos:.3f}\"\n )\n\n return best_pipe, num_cols, cat_cols, metrics\n\n# ------------------------------------------------------------\n# 4) Save/load gate model\n# ------------------------------------------------------------\ndef maybe_train_and_save_gate_model(\n lab_id: str,\n action: str,\n switch: str,\n train_df: pd.DataFrame,\n out_root: str,\n train_tracking_dir: str = CFG.train_tracking_path,\n) -> None:\n \"\"\"\n Convenience wrapper:\n - builds video_tab (with clip activity; now cached per lab/switch)\n - joins metadata (incl. arena_area_cm2)\n - trains gate\n - saves it if useful (AUC high enough, enough examples)\n \"\"\"\n video_tab = build_video_level_table_for_action(\n lab_id=lab_id,\n action=action,\n switch=switch,\n train_df=train_df,\n train_tracking_dir=train_tracking_dir,\n )\n\n if video_tab.empty:\n return\n\n video_tab = add_video_metadata(video_tab, train_df)\n\n gate_model, num_cols, cat_cols, metrics = train_action_presence_model(\n video_tab,\n lab_id=lab_id,\n action=action,\n switch=switch,\n min_pos_rate=GATE_MIN_POS_RATE,\n min_auc=GATE_MIN_AUC,\n )\n\n if gate_model is None:\n return\n\n out_dir = os.path.join(out_root, str(lab_id), str(switch))\n os.makedirs(out_dir, exist_ok=True)\n\n obj = {\n \"model\": gate_model,\n \"num_cols\": num_cols,\n \"cat_cols\": cat_cols,\n \"metrics\": metrics,\n }\n path = os.path.join(out_dir, f\"gate_{action}.pkl\")\n joblib.dump(obj, path, compress=3)\n\n\ndef load_gate_model(lab_id: str, action: str, switch: str, out_root: str):\n \"\"\"\n Returns saved gate model dict or None.\n \"\"\"\n path = os.path.join(out_root, str(lab_id), str(switch), f\"gate_{action}.pkl\")\n if not os.path.exists(path):\n return None\n return joblib.load(path)\n\n\n# ------------------------------------------------------------\n# 5) Soft logit gating for frame-level preds (DOWN-WEIGHT only)\n# ------------------------------------------------------------\ndef apply_gate_to_preds_soft_logit(\n pred: pd.DataFrame,\n gate_prob: pd.Series | np.ndarray,\n actions: list[str] | None = None,\n strength: float = GATE_STRENGTH,\n eps: float = 1e-6,\n only_downweight: bool = True,\n neg_clip: float = GATE_NEG_CLIP_MAX_PROB,\n) -> pd.DataFrame:\n \"\"\"\n Softly gate frame-level probabilities using a clip-level gate probability\n in logit space.\n \"\"\"\n if actions is None:\n actions = list(pred.columns)\n\n gate_series = pd.Series(gate_prob, index=pred.index)\n gate_series = pd.to_numeric(gate_series, errors=\"coerce\")\n gate_series = gate_series.fillna(0.5).clip(eps, 1.0 - eps)\n\n def _logit(p):\n return np.log(p) - np.log(1.0 - p)\n\n def _inv_logit(z):\n return 1.0 / (1.0 + np.exp(-z))\n\n g_eff = pd.Series(0.5, index=gate_series.index, dtype=np.float32)\n\n strength_vec = np.zeros(len(gate_series), dtype=np.float32)\n base_strength = float(strength)\n neg_clip = float(neg_clip)\n strong_cut = float(min(GATE_STRONG_PROB, neg_clip))\n\n if only_downweight:\n g_vals = gate_series.to_numpy(dtype=np.float32)\n\n neutral_mask = g_vals >= neg_clip\n strong_mask = g_vals < strong_cut\n soft_mask = (~neutral_mask) & (~strong_mask)\n\n g_eff.iloc[soft_mask] = g_vals[soft_mask]\n g_eff.iloc[strong_mask] = g_vals[strong_mask]\n\n strength_vec[soft_mask] = base_strength * 0.7\n strength_vec[strong_mask] = base_strength * 1.3\n else:\n g_eff = gate_series.astype(np.float32)\n strength_vec[:] = base_strength\n\n g_vals_eff = g_eff.to_numpy(dtype=np.float64)\n g_vals_eff = np.clip(g_vals_eff, eps, 1.0 - eps)\n\n logit_gate = _logit(g_vals_eff) * strength_vec\n\n pred_gated = pred.copy()\n\n for act in actions:\n if act not in pred_gated.columns:\n continue\n p = pred_gated[act].astype(float).clip(eps, 1.0 - eps)\n p_vals = p.to_numpy(dtype=np.float64)\n z = _logit(p_vals) + logit_gate\n pred_gated[act] = _inv_logit(z).astype(np.float32)\n\n return pred_gated\n\n\n# ------------------------------------------------------------\n# Shared gating helpers for test / holdout\n# ------------------------------------------------------------\n\nGATE_META_COLS = [\n \"video_id\",\n \"video_duration_sec\", \"frames_per_second\",\n \"arena_type\", \"arena_shape\",\n \"mouse1_strain\", \"mouse2_strain\", \"mouse3_strain\", \"mouse4_strain\",\n \"mouse1_condition\", \"mouse2_condition\", \"mouse3_condition\", \"mouse4_condition\",\n \"mouse1_sex\", \"mouse2_sex\",\n \"arena_width_cm\", \"arena_height_cm\",\n \"n_mice\",\n \"tracking_method\",\n]\n\n\ndef build_video_meta_lookup_for_gating(df: pd.DataFrame) -> pd.DataFrame:\n available_meta_cols = [c for c in GATE_META_COLS if c in df.columns]\n if \"video_id\" not in available_meta_cols:\n available_meta_cols = [\"video_id\"] + [\n c for c in available_meta_cols if c != \"video_id\"\n ]\n\n video_meta_lookup = (\n df[available_meta_cols]\n .drop_duplicates(\"video_id\")\n .set_index(\"video_id\")\n )\n return video_meta_lookup\n\n\ndef build_gate_feature_row_from_chunk(\n meta_df: pd.DataFrame,\n data_df: pd.DataFrame | None,\n ctx_df: pd.DataFrame | None,\n fps_i: float,\n video_meta_lookup: pd.DataFrame,\n) -> dict | None:\n if meta_df is None or meta_df.empty:\n return None\n\n vid = meta_df[\"video_id\"].iloc[0]\n if vid in video_meta_lookup.index:\n base_meta = video_meta_lookup.loc[vid].to_dict()\n else:\n base_meta = {}\n\n def _mget(key, default=np.nan):\n return base_meta.get(key, default)\n\n arena_w = _mget(\"arena_width_cm\", np.nan)\n arena_h = _mget(\"arena_height_cm\", np.nan)\n\n clip_feats = compute_clip_activity_features(\n data_df=data_df,\n ctx_df=ctx_df,\n fps=fps_i,\n arena_width_cm=arena_w,\n arena_height_cm=arena_h,\n )\n\n with np.errstate(all=\"ignore\"):\n arena_area_cm2 = (\n float(arena_w) * float(arena_h)\n if (np.isfinite(arena_w) and np.isfinite(arena_h))\n else np.nan\n )\n\n row = {\n \"n_frames\": len(meta_df),\n \"n_mice\": _mget(\"n_mice\", np.nan),\n \"arena_width_cm\": arena_w,\n \"arena_height_cm\": arena_h,\n \"arena_area_cm2\": arena_area_cm2,\n \"video_duration_sec\": _mget(\"video_duration_sec\", np.nan),\n \"frames_per_second\": _mget(\"frames_per_second\", fps_i),\n \"arena_type\": _mget(\"arena_type\", \"\"),\n \"arena_shape\": _mget(\"arena_shape\", \"\"),\n \"mouse1_strain\": _mget(\"mouse1_strain\", \"\"),\n \"mouse2_strain\": _mget(\"mouse2_strain\", \"\"),\n \"mouse3_strain\": _mget(\"mouse3_strain\", \"\"),\n \"mouse4_strain\": _mget(\"mouse4_strain\", \"\"),\n \"mouse1_condition\": _mget(\"mouse1_condition\", \"\"),\n \"mouse2_condition\": _mget(\"mouse2_condition\", \"\"),\n \"mouse3_condition\": _mget(\"mouse3_condition\", \"\"),\n \"mouse4_condition\": _mget(\"mouse4_condition\", \"\"),\n \"mouse1_sex\": _mget(\"mouse1_sex\", \"\"),\n \"mouse2_sex\": _mget(\"mouse2_sex\", \"\"),\n \"tracking_method\": _mget(\"tracking_method\", \"\"),\n }\n row.update(clip_feats)\n return row\n\n\ndef get_or_train_gate_model_for_action(\n lab_id: str,\n action_name: str,\n switch_tr: str,\n cfg: dict,\n gate_models: dict[str, dict],\n gate_skipped: set[str],\n) -> dict | None:\n \"\"\"\n Train & cache gate model for one (lab_id, action, switch) on TRAIN only.\n Mutates gate_models / gate_skipped in-place, returns gate_info or None.\n \"\"\"\n if action_name in gate_skipped:\n return None\n\n if action_name in gate_models:\n return gate_models[action_name]\n\n try:\n train_subset_gate = train[train.lab_id.isin(cfg[\"labs\"])]\n\n video_tab_gate = build_video_level_table_for_action(\n lab_id=lab_id,\n action=action_name,\n switch=switch_tr,\n train_df=train_subset_gate,\n train_tracking_dir=CFG.train_tracking_path,\n )\n except Exception as e:\n print(\n f\"[get_or_train_gate_model_for_action] GATE build failed for \"\n f\"lab={lab_id}, action={action_name}, switch={switch_tr}: {e}\"\n )\n gate_skipped.add(action_name)\n return None\n\n if video_tab_gate.empty:\n gate_skipped.add(action_name)\n return None\n\n video_tab_gate = add_video_metadata(video_tab_gate, train_subset_gate)\n\n gate_model, num_cols, cat_cols, metrics = train_action_presence_model(\n video_tab_gate,\n lab_id=lab_id,\n action=action_name,\n switch=switch_tr,\n min_pos_rate=GATE_MIN_POS_RATE,\n min_auc=GATE_MIN_AUC,\n )\n\n if gate_model is None:\n gate_skipped.add(action_name)\n return None\n\n info = {\n \"model\": gate_model,\n \"num_cols\": num_cols,\n \"cat_cols\": cat_cols,\n }\n gate_models[action_name] = info\n return info\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def robustify(submission, dataset, traintest, traintest_directory=None):\n if traintest_directory is None:\n traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\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 \n return submission","metadata":{"_kg_hide-input":false,"jupyter":{"source_hidden":true},"papermill":{"duration":0.023773,"end_time":"2025-10-25T19:26:58.542199","exception":false,"start_time":"2025-10-25T19:26:58.518426","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_multiclass(pred, meta, thresholds):\n \"\"\"\n Choose actions by how much each probability exceeds its per-action threshold\n in absolute terms: margin = p - thr.\n\n For each frame:\n - Compute margins = p - thr for all actions.\n - Pick the action with the largest margin.\n - If max margin < 0, emit -1 (no action).\n \"\"\"\n probs = pred.values # shape: (n_frames, n_actions)\n\n # Build per-action thresholds aligned with pred.columns\n thr_vec = np.array([thresholds.get(col, 0.27) for col in pred.columns], dtype=float)\n\n # Absolute margin above threshold: >0 means above; <0 means below\n margins = probs - thr_vec # broadcast over columns\n\n # Best action per frame according to this margin\n best_idx = np.argmax(margins, axis=1)\n best_margin = margins[np.arange(len(best_idx)), best_idx]\n\n # If best_margin < 0 -> no action\n ama = np.where(best_margin >= 0.0, best_idx, -1)\n ama = pd.Series(ama, index=meta.video_frame)\n\n # --- unchanged segment-building logic below ---\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 # last change only defines the stop of the previous segment\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 # Clamp segments that would cross (video, agent, target) boundaries\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\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\n stop_agent_id[i] != agent_id or\n 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","metadata":{"jupyter":{"source_hidden":true},"papermill":{"duration":0.017663,"end_time":"2025-10-25T19:51:34.056443","exception":false,"start_time":"2025-10-25T19:51:34.03878","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def tune_threshold(\n oof_action: np.ndarray,\n y_action: np.ndarray,\n thresholds: np.ndarray | None = None,\n) -> float:\n \"\"\"\n 1D threshold tuning for a single action, maximizing F1 over a grid.\n \"\"\"\n oof_action = np.asarray(oof_action, dtype=np.float32)\n y_action = np.asarray(y_action, dtype=np.int8)\n\n if thresholds is None:\n thresholds = np.linspace(0.0, 1.0, 101, dtype=np.float32)\n else:\n thresholds = np.asarray(thresholds, dtype=np.float32)\n\n # Broadcast: preds.shape = (n_samples, n_thresholds)\n preds = (oof_action[:, None] >= thresholds[None, :])\n\n # True / false positives / negatives per threshold\n y_pos = (y_action == 1)\n y_neg = ~y_pos\n\n tp = (preds & y_pos[:, None]).sum(axis=0).astype(np.float32)\n fp = (preds & y_neg[:, None]).sum(axis=0).astype(np.float32)\n fn = ((~preds) & y_pos[:, None]).sum(axis=0).astype(np.float32)\n\n # Precision, recall, F1 (avoid division by zero)\n denom_prec = tp + fp\n denom_rec = tp + fn\n\n prec = np.divide(tp, denom_prec, out=np.zeros_like(tp), where=denom_prec > 0)\n rec = np.divide(tp, denom_rec, out=np.zeros_like(tp), where=denom_rec > 0)\n\n denom_f1 = (prec + rec)\n f1 = np.divide(2 * prec * rec, denom_f1, out=np.zeros_like(denom_f1), where=denom_f1 > 0)\n\n best_idx = int(f1.argmax())\n best_thr = float(thresholds[best_idx])\n return best_thr\n\ndef tune_ensemble(\n preds_by_model: dict[str, pd.DataFrame],\n label: pd.DataFrame,\n actions: list[str] | None = None,\n thresholds_by_model: dict[str, dict[str, float]] | None = None,\n n_trials: int = 200,\n thr_bounds: tuple[float, float] = (0.01, 0.99),\n beta: float = 1.0,\n seed: int = 42,\n fps: float = 30.0,\n base_smooth_sec: float = 5.0 / 30.0,\n model_tags: list[str] | None = None, # order of models in the ensemble\n optuna_n_jobs: int = 1, # how many parallel Optuna threads\n) -> tuple[dict[str, float], float, dict[str, float], dict[str, dict[str, float]]]:\n \"\"\"\n Jointly tune:\n - ensemble model weights (per action; special-cased for 1 or 2 models),\n - per-action thresholds,\n - per-action smoothing windows,\n using multi-class decision rule (argmax margin p - thr).\n\n Parameters\n ----------\n preds_by_model : dict[model_tag -> DataFrame]\n Each DataFrame: shape (N, A) with action columns.\n label : DataFrame\n Ground-truth labels (columns = actions).\n actions : list[str] | None\n Actions to consider; defaults to label.columns.\n thresholds_by_model : dict[model_tag -> dict[action -> thr]] | None\n Per-model single thresholds, used to initialize the search\n (we take per-action mean across models).\n model_tags : list[str] | None\n Order of models in the ensemble. If None, uses preds_by_model.keys()\n in insertion order.\n optuna_n_jobs : int\n Passed to study.optimize(n_jobs=...). 1 = sequential; -1 = all cores.\n\n Returns\n -------\n best_thresholds : dict[str, float] # action -> threshold\n best_score : float # best F-beta\n best_smooth : dict[str, float] # action -> smoothing seconds\n best_weights : dict[str, dict[str, float]]\n Nested dict: best_weights[action][model_tag] = weight (sum ~ 1 per action)\n \"\"\"\n\n if not preds_by_model:\n raise ValueError(\"tune_ensemble: preds_by_model is empty.\")\n\n # --- model order ---\n if model_tags is None:\n # keep dict insertion order\n model_tags = list(preds_by_model.keys())\n else:\n # keep only models that actually exist in preds_by_model\n model_tags = [t for t in model_tags if t in preds_by_model]\n\n if len(model_tags) == 0:\n raise ValueError(\"tune_ensemble: model_tags resolved to empty list.\")\n\n M = len(model_tags)\n single_model = (M == 1)\n two_models = (M == 2)\n\n # --- actions / columns ---\n if actions is None:\n actions = list(label.columns)\n actions = list(actions)\n\n # Align label columns and reset index to avoid duplicate-index reindex issues\n label = label[actions].copy().reset_index(drop=True)\n N = len(label)\n\n # Convert preds to aligned numpy arrays per model (align by position)\n P_raw_models: dict[str, np.ndarray] = {}\n for tag in model_tags:\n df = preds_by_model[tag][actions].copy().reset_index(drop=True)\n\n if len(df) != N:\n raise ValueError(\n f\"tune_ensemble: length mismatch for model '{tag}': \"\n f\"preds={len(df)}, label={N}\"\n )\n\n P_raw_models[tag] = df.to_numpy(dtype=np.float32, copy=False)\n\n # Common shape\n N, A = next(iter(P_raw_models.values())).shape\n\n # ---- Precompute masks / y per action (skip NaNs) ----\n masks: list[np.ndarray] = [] # boolean masks per action\n ys: list[np.ndarray] = [] # y_true (0/1) per action\n actions_with_pos: list[int] = [] # indices into `actions`\n\n for j, act in enumerate(actions):\n y = label[act].to_numpy(dtype=float)\n mask = ~np.isnan(y)\n y_masked = y[mask].astype(np.int8)\n\n if y_masked.sum() > 0:\n masks.append(mask)\n ys.append(y_masked)\n actions_with_pos.append(j)\n\n low_global, high_global = thr_bounds\n default_thr = (low_global + high_global) / 2.0\n\n if not actions_with_pos:\n # No positives for any action in this lab; nothing meaningful to tune\n best_thresholds = {a: default_thr for a in actions}\n best_smooth = {a: 0.0 for a in actions}\n best_weights = {\n a: {tag: 1.0 / M for tag in model_tags}\n for a in actions\n }\n return best_thresholds, 0.0, best_smooth, best_weights\n\n # ---- Base thresholds = mean over models (if provided) ----\n base_thr_vec = np.empty(A, dtype=np.float32)\n for j, act in enumerate(actions):\n if thresholds_by_model:\n vals = []\n for tag in model_tags:\n td = thresholds_by_model.get(tag, {})\n if act in td:\n vals.append(float(td[act]))\n if vals:\n base_thr_vec[j] = float(np.mean(vals))\n else:\n base_thr_vec[j] = default_thr\n else:\n base_thr_vec[j] = default_thr\n\n base_thr_vec = np.clip(base_thr_vec, low_global, high_global)\n\n # ---- Cache for smoothed columns: key = (action_idx, smooth_sec) ----\n smooth_cache: dict[tuple[int, float], np.ndarray] = {}\n\n sampler = optuna.samplers.TPESampler(seed=seed, constant_liar=True)\n study = optuna.create_study(direction=\"maximize\", sampler=sampler)\n\n # ---- Enqueue an initial trial ----\n init_params: dict[str, float] = {}\n tag0 = model_tags[0]\n tag1 = model_tags[1] if two_models else None\n\n for j, act in enumerate(actions):\n init_params[f\"thr_{act}\"] = float(base_thr_vec[j])\n init_params[f\"smooth_sec_{act}\"] = float(base_smooth_sec)\n\n if single_model:\n # no weight params\n continue\n elif two_models:\n # single weight per action: w_{act}_{tag0}{tag1}\n init_params[f\"w_{act}_{tag0}{tag1}\"] = 0.5\n else:\n # M > 2: per-action, per-model weights\n for tag in model_tags:\n init_params[f\"w_{act}_{tag}\"] = 1.0 / M\n\n study.enqueue_trial(init_params)\n\n smooth_step = base_smooth_sec / 3.0 if base_smooth_sec > 0 else 0.01\n smooth_low = 0.0\n smooth_high = base_smooth_sec * 3.0 if base_smooth_sec > 0 else 0.5\n\n def objective(trial: optuna.Trial) -> float:\n # ---- sample weights ----\n # For M = 1: implicit weights = 1.0, no params\n # For M = 2: per-action scalar w2, w1 = 1 - w2\n # For M > 2: per-action, per-model weights (normalized)\n if not single_model and not two_models:\n # w_mat[j, m] = weight for action j, model m\n w_mat = np.empty((A, M), dtype=np.float32)\n for j, act in enumerate(actions):\n w_raw = np.empty(M, dtype=np.float32)\n for m, tag in enumerate(model_tags):\n w_raw[m] = trial.suggest_float(\n f\"w_{act}_{tag}\",\n 0.0,\n 1.0,\n step=0.05,\n )\n s_raw_sum = float(w_raw.sum())\n if s_raw_sum <= 0:\n w_mat[j, :] = 1.0 / M\n else:\n w_mat[j, :] = w_raw / s_raw_sum\n else:\n w_mat = None # not used for M <= 2\n\n if two_models:\n # per-action scalar for second model; first gets 1 - w2\n w2_vec = np.empty(A, dtype=np.float32)\n for j, act in enumerate(actions):\n w2 = trial.suggest_float(\n f\"w_{act}_{tag0}{tag1}\",\n 0.0,\n 1.0,\n step=0.05,\n )\n w2_vec[j] = w2\n else:\n w2_vec = None\n\n # ---- sample thresholds & smoothing per action ----\n thr_vec = np.empty(A, dtype=np.float32)\n smooth_secs = np.empty(A, dtype=np.float32)\n\n for j, act in enumerate(actions):\n center = base_thr_vec[j]\n low_j = max(low_global, center * 0.5)\n high_j = max(min(high_global, center * 1.5), 0.25)\n\n thr = trial.suggest_float(f\"thr_{act}\", low_j, high_j, step=0.005)\n thr_vec[j] = thr\n\n s = trial.suggest_float(\n f\"smooth_sec_{act}\",\n smooth_low,\n smooth_high,\n step=smooth_step,\n )\n smooth_secs[j] = s\n\n thr_vec_safe = np.clip(thr_vec, 0.0, None)\n\n # ---- build ensemble preds ----\n if single_model:\n # just copy that model's predictions\n P_raw_ens = P_raw_models[tag0].copy()\n else:\n P_raw_ens = np.zeros((N, A), dtype=np.float32)\n if two_models:\n # two-model, per-action scalar weights\n for j in range(A):\n w2 = float(w2_vec[j])\n w1 = 1.0 - w2\n col = (\n w1 * P_raw_models[tag0][:, j]\n + w2 * P_raw_models[tag1][:, j]\n )\n P_raw_ens[:, j] = col\n else:\n # M > 2: per-action, per-model weights\n for j in range(A):\n col = np.zeros(N, dtype=np.float32)\n for m, tag in enumerate(model_tags):\n col += w_mat[j, m] * P_raw_models[tag][:, j]\n P_raw_ens[:, j] = col\n\n # ---- per-action smoothing ----\n P_eff = P_raw_ens.copy()\n for j in range(A):\n s = smooth_secs[j]\n if s <= 0.0:\n continue\n key = (j, float(np.round(s, 6)))\n if key not in smooth_cache:\n col = P_raw_ens[:, j]\n smooth_col = _smooth(col, fps=fps, kind=\"mean\", smooth_sec=s)\n smooth_cache[key] = np.asarray(smooth_col, dtype=np.float32).reshape(-1)\n P_eff[:, j] = smooth_cache[key]\n\n # ---- multiclass decision: argmax margin (p - thr) ----\n margins = P_eff - thr_vec_safe # (N, A)\n best_idx = np.argmax(margins, axis=1)\n best_margin = margins[np.arange(N), best_idx]\n pred_class = np.where(best_margin >= 0.0, best_idx, -1)\n\n # ---- compute mean F-beta over actions-with-positives ----\n scores = []\n for k, j in enumerate(actions_with_pos):\n mask = masks[k]\n y_true = ys[k]\n y_pred = (pred_class[mask] == j).astype(np.int8)\n if y_true.sum() == 0:\n continue\n f1 = fbeta_score(y_true, y_pred, beta=beta)\n scores.append(f1)\n\n if not scores:\n return 0.0\n\n return float(np.mean(scores))\n\n study.optimize(\n objective,\n n_trials=n_trials,\n show_progress_bar=False,\n n_jobs=optuna_n_jobs,\n )\n\n best_trial = study.best_trial\n\n # ---- extract best thresholds / smoothing ----\n best_thresholds: dict[str, float] = {}\n best_smooth_dict: dict[str, float] = {}\n\n for j, act in enumerate(actions):\n best_thresholds[act] = float(best_trial.params[f\"thr_{act}\"])\n s = float(best_trial.params.get(f\"smooth_sec_{act}\", 0.0))\n best_smooth_dict[act] = 0.0 if s <= 0.0 else s\n\n # ---- extract best per-action ensemble weights ----\n best_weights: dict[str, dict[str, float]] = {}\n\n for j, act in enumerate(actions):\n if single_model:\n best_weights[act] = {tag0: 1.0}\n elif two_models:\n # read per-action scalar w2; clamp to [0,1]\n pname = f\"w_{act}_{tag0}{tag1}\"\n w2 = float(best_trial.params.get(pname, 0.5))\n w2 = min(max(w2, 0.0), 1.0)\n w1 = 1.0 - w2\n best_weights[act] = {tag0: w1, tag1: w2}\n else:\n # M > 2: normalize per-action, per-model\n w_raw_act: dict[str, float] = {}\n total_w = 0.0\n for tag in model_tags:\n param_name = f\"w_{act}_{tag}\"\n w_val = float(best_trial.params.get(param_name, 0.0))\n w_raw_act[tag] = w_val\n total_w += w_val\n\n if total_w <= 0:\n best_weights[act] = {tag: 1.0 / M for tag in model_tags}\n else:\n best_weights[act] = {\n tag: (w_raw_act[tag] / total_w) for tag in model_tags\n }\n\n best_score = float(study.best_value)\n\n # ---- print median weight per model across actions ----\n median_weights: dict[str, float] = {}\n for tag in model_tags:\n vals = [best_weights[act][tag] for act in actions if tag in best_weights[act]]\n median_weights[tag] = float(np.median(vals)) if vals else 0.0\n\n weights_str = \" \".join(\n f\"{k}={v:.3f}\" for k, v in sorted(median_weights.items())\n )\n print(f\"Ensemble median weights: {weights_str}\")\n\n return best_thresholds, best_score, best_smooth_dict, best_weights\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def compute_lab_f1_from_thresholds(\n pred: pd.DataFrame,\n label: pd.DataFrame,\n thresholds: dict,\n actions: list[str] | None = None,\n beta: float = 1.0,\n) -> float:\n \"\"\"\n Compute mean lab F-beta over actions with positives, using the same\n multiclass decision rule based on the p - thr margin:\n\n margin_j = P_j - thr_j\n j* = argmax_j margin_j\n if margin_{j*} < 0: predict \"no action\"\n else: predict action j*\n\n NOTE: this assumes `pred` already has whatever smoothing you want applied.\n \"\"\"\n\n if actions is None:\n actions = list(pred.columns)\n\n # Align columns\n pred = pred[actions]\n label = label[actions]\n\n P = pred.to_numpy(dtype=np.float32) # (N, A)\n N, A = P.shape\n\n # Per-action thresholds aligned with columns; default if missing\n thr_vec = np.array(\n [thresholds.get(act, 0.27) for act in actions],\n dtype=np.float32\n )\n thr_vec_safe = np.clip(thr_vec, 0.0, None)\n\n # Precompute masks and y for actions with positives\n masks = [] # list[ndarray[bool]]\n ys = [] # list[ndarray[int]]\n actions_with_pos = [] # list[int] indices into `actions`\n\n for j, act in enumerate(actions):\n y = label[act].to_numpy(dtype=float)\n mask = ~np.isnan(y)\n y_masked = y[mask].astype(np.int8)\n\n if y_masked.sum() > 0:\n masks.append(mask)\n ys.append(y_masked)\n actions_with_pos.append(j)\n\n if not actions_with_pos:\n return 0.0\n\n # Decision rule: argmax(P - thr), -1 if max < 0\n margins = P - thr_vec_safe # (N, A)\n best_idx = np.argmax(margins, axis=1)\n best_margin = margins[np.arange(N), best_idx]\n pred_class = np.where(best_margin >= 0.0, best_idx, -1)\n\n f1s = []\n for k, j in enumerate(actions_with_pos):\n mask = masks[k]\n y_true = ys[k]\n y_pred = (pred_class[mask] == j).astype(np.int8)\n\n if y_true.sum() == 0:\n continue\n\n f1 = fbeta_score(y_true, y_pred, beta=beta)\n f1s.append(f1)\n\n if not f1s:\n return 0.0\n\n return float(np.mean(f1s))\n\ndef score_per_lab(lab_id, tmp_submission_list, lab_solution, train_subset):\n \"\"\"\n Compute the competition metric for a single lab_id (one or more lab_ids),\n using the same robustify + score pipeline as the global metric.\n\n Parameters\n ----------\n lab_id : str\n Key from labs_dict (your cid).\n tmp_submission_list : list[pd.DataFrame]\n Per-lab submissions (both single & pair) accumulated for this cid.\n lab_solution : pd.DataFrame\n Subset of the global solution with only the lab_id(s) of this cid.\n train_subset : pd.DataFrame\n Same subset of the train metadata passed to robustify (has behaviors_labeled, etc).\n\n Returns\n -------\n float\n Competition metric for this cid (mouse_fbeta averaged over its lab_id(s)).\n \"\"\"\n if not tmp_submission_list:\n print(f\"[per-lab] {lab_id}: no submissions, skipping\")\n return float(\"nan\")\n\n # Combine single+pair submissions for this cid\n lab_submission = pd.concat(tmp_submission_list, ignore_index=True)\n\n # Apply your robustify exactly as in the global scoring\n lab_submission_robust = robustify(lab_submission, train_subset, \"train\")\n\n # Score using your official metric, but only on this cid's labs\n lab_score = score(lab_solution, lab_submission_robust, row_id_column_name=\"\")\n\n print(f\"[per-lab] {lab_id}: competition metric = {lab_score:.4f}\")\n return lab_score","metadata":{"jupyter":{"source_hidden":true},"papermill":{"duration":0.017663,"end_time":"2025-10-25T19:51:34.056443","exception":false,"start_time":"2025-10-25T19:51:34.03878","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_lab_holdout_videos(\n train_subset: pd.DataFrame,\n lab_solution: pd.DataFrame,\n holdout_frac_videos: float = 0.15,\n min_pos_videos_for_strat: int = 2,\n random_state: int = 42,\n) -> list:\n \"\"\"\n Choose ~holdout_frac_videos of videos in train_subset as holdout,\n grouped by video_id and roughly stratified over actions.\n\n We use `lab_solution` (segment-level GT) to know which actions exist\n in which videos. Edge cases:\n - <=1 video: no holdout.\n - We first try to ensure each action with enough positive videos\n gets at least one positive video into holdout, then fill up\n randomly to target count.\n \"\"\"\n rng = np.random.default_rng(random_state)\n\n vids = train_subset[\"video_id\"].unique()\n n_videos = len(vids)\n if n_videos <= 1:\n return []\n\n # action -> set of videos with at least one positive segment\n act_to_vids = defaultdict(set)\n # NOTE: lab_solution may contain actions not in this cluster; that's fine\n for vid, act in lab_solution[[\"video_id\", \"action\"]].drop_duplicates().itertuples(index=False):\n act_to_vids[act].add(vid)\n\n target_n = max(1, int(round(holdout_frac_videos * n_videos)))\n selected = set()\n\n # Stratified pass: try to give each sufficiently-represented action a holdout video\n actions_sorted = sorted(act_to_vids.keys(), key=lambda a: len(act_to_vids[a]), reverse=True)\n for act in actions_sorted:\n vids_pos = np.array(list(act_to_vids[act]))\n if len(vids_pos) < min_pos_videos_for_strat:\n continue # too rare, don't force it\n\n candidates = [v for v in vids_pos if v not in selected]\n if not candidates:\n continue\n\n v = rng.choice(candidates)\n selected.add(v)\n\n if len(selected) >= target_n:\n break\n\n # Fill up to target_n with random remaining videos\n if len(selected) < target_n:\n remaining = [v for v in vids if v not in selected]\n if remaining:\n n_extra = min(target_n - len(selected), len(remaining))\n extra = rng.choice(remaining, size=n_extra, replace=False)\n selected.update(extra)\n\n return sorted(selected)\n\nfrom sklearn.model_selection import GroupKFold, StratifiedKFold\n\ndef make_group_stratified_splits(\n groups,\n y,\n n_splits: int,\n random_state: int = 42,\n):\n \"\"\"\n Build CV splits that:\n - keep all rows from the same group (video_id) together\n - are stratified (when possible) on per-group positive rate for this action.\n\n Edge cases handled:\n - 0 or 1 unique group -> return [] (no CV possible)\n - fewer groups than n_splits -> clamp to n_splits_eff = min(G, n_splits)\n - all labels same -> GroupKFold (no stratification possible)\n - pathological binning (StratifiedKFold raises) -> GroupKFold fallback\n - accidental empty train/val folds -> dropped, fallback to GroupKFold if all dropped\n \"\"\"\n\n groups = np.asarray(groups)\n y = np.asarray(y).astype(np.int8)\n assert groups.shape[0] == y.shape[0], \"groups and y must have same length\"\n\n uniq_groups, group_idx = np.unique(groups, return_inverse=True)\n G = len(uniq_groups)\n\n # ---- 0 or 1 group: cannot do real CV ----\n if G <= 1:\n # signal \"no CV\", caller will keep oof=0's and skip training effectively\n return []\n\n # ---- clamp requested splits to #groups ----\n n_splits_eff = min(int(n_splits), G)\n if n_splits_eff < 2:\n # should not really happen with G >= 2, but be defensive\n return []\n\n # ---- trivial label structure: all 0 or all 1 -> GroupKFold only ----\n if y.min() == y.max():\n gkf = GroupKFold(n_splits=n_splits_eff)\n return list(gkf.split(np.zeros_like(y), y, groups=groups))\n\n # ---- compute per-group positive rate ----\n group_pos = np.zeros(G, dtype=float)\n group_cnt = np.zeros(G, dtype=int)\n np.add.at(group_pos, group_idx, y)\n np.add.at(group_cnt, group_idx, 1)\n group_pos /= np.maximum(group_cnt, 1)\n\n # ---- convert per-group pos rate into bins for stratification ----\n # up to 5 bins, but not more than G or n_splits_eff\n n_bins = min(5, n_splits_eff, G)\n if n_bins < 2:\n # can't meaningfully stratify -> GroupKFold\n gkf = GroupKFold(n_splits=n_splits_eff)\n return list(gkf.split(np.zeros_like(y), y, groups=groups))\n\n order = np.argsort(group_pos)\n group_bins = np.empty(G, dtype=int)\n for rank, g_idx in enumerate(order):\n # roughly equal number of groups per bin\n group_bins[g_idx] = int(rank * n_bins / G)\n\n skf = StratifiedKFold(\n n_splits=n_splits_eff,\n shuffle=True,\n random_state=random_state,\n )\n group_indices = np.arange(G)\n\n try:\n folds_on_groups = list(skf.split(group_indices, group_bins))\n except ValueError:\n # e.g. some bins have < n_splits_eff groups -> fallback to GroupKFold\n gkf = GroupKFold(n_splits=n_splits_eff)\n return list(gkf.split(np.zeros_like(y), y, groups=groups))\n\n # ---- map group-level folds -> row-level folds ----\n splits: list[tuple[np.ndarray, np.ndarray]] = []\n row_indices = np.arange(len(y))\n\n for train_g_idx, val_g_idx in folds_on_groups:\n train_groups = uniq_groups[train_g_idx]\n val_groups = uniq_groups[val_g_idx]\n\n train_mask = np.isin(groups, train_groups)\n val_mask = np.isin(groups, val_groups)\n\n tr_idx = row_indices[train_mask]\n val_idx = row_indices[val_mask]\n\n # be safe: drop degenerate folds\n if len(tr_idx) == 0 or len(val_idx) == 0:\n continue\n\n splits.append((tr_idx, val_idx))\n\n # if everything degenerated (very rare), fall back once more to GroupKFold\n if not splits:\n gkf = GroupKFold(n_splits=n_splits_eff)\n splits = list(gkf.split(np.zeros_like(y), y, groups=groups))\n\n return splits\n","metadata":{"jupyter":{"source_hidden":true},"papermill":{"duration":0.017663,"end_time":"2025-10-25T19:51:34.056443","exception":false,"start_time":"2025-10-25T19:51:34.03878","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def submit_lab(\n lab_id: str,\n switch_tr: str,\n thresholds_for_lab: dict,\n smooth_cfg_for_lab: dict | None = None,\n model_weights_for_lab: dict | None = None,\n model_tags: list[str] | None = None,\n) -> list[pd.DataFrame]:\n \"\"\"\n Run inference for one lab + switch ('single' or 'pair') using an ensemble\n of tree models saved under CFG.model_path / model_tag / lab_id / switch / action.\n\n This version also applies clip-level gating in logit space, using gate\n models trained on the TRAIN set for this lab + switch via:\n - build_video_level_table_for_action\n - add_video_metadata\n - train_action_presence_model\n - apply_gate_to_preds_soft_logit\n\n Parameters\n ----------\n lab_id : str\n Cluster / lab key.\n switch_tr : {\"single\", \"pair\"}\n Which switch to run.\n thresholds_for_lab : dict[str, float]\n Dict {action -> threshold} for THIS lab & switch (ensemble thresholds).\n smooth_cfg_for_lab : dict[str, float] or None\n Dict {action -> smooth_sec} for THIS lab & switch.\n smooth_sec == 0.0 means \"no smoothing\".\n model_weights_for_lab : dict[str, dict[str, float]] or None\n Per-action ensemble weights:\n { action -> { model_tag -> weight } }\n If None or empty, defaults to uniform over available models.\n model_tags : list[str] or None\n List of model_tags to use. If None, falls back to CFG.model_tags\n or, if that is missing, to tags found in model_weights_for_lab.\n\n Returns\n -------\n submission_list : list[pd.DataFrame]\n List of segment DataFrames for this lab & switch.\n \"\"\"\n\n cfg = labs_dict[lab_id]\n body_parts_tracked = list(cfg[\"body_parts\"])\n\n # restrict test to this lab's underlying dataset labs\n test_subset = test[test.lab_id.isin(cfg[\"labs\"])]\n\n if test_subset.empty:\n print(\n f\"[submit_lab] No test rows found for lab {lab_id} \"\n f\"(switch={switch_tr}); skipping inference.\"\n )\n return []\n\n # Ensure weights dict is always a dict\n model_weights_for_lab = model_weights_for_lab or {}\n\n # Determine which model tags to use\n if model_tags is not None:\n model_tags_eff = list(model_tags)\n else:\n # derive from weights if present, else CFG.model_tags\n if model_weights_for_lab:\n tags = set()\n for w_per_action in model_weights_for_lab.values():\n tags.update(w_per_action.keys())\n model_tags_eff = sorted(tags)\n else:\n model_tags_eff = list(getattr(CFG, \"model_tags\", []))\n\n if not model_tags_eff:\n print(f\"[submit_lab] No model_tags provided for lab {lab_id}; skipping.\")\n return []\n\n generator = generate_mouse_data(\n test_subset,\n \"test\",\n generate_single=(switch_tr == \"single\"),\n generate_pair=(switch_tr == \"pair\"),\n )\n\n fps_lookup = (\n test_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 # per-lab smoothing cfg (already lab-specific)\n lab_smooth_cfg = smooth_cfg_for_lab or {}\n\n submission_list: list[pd.DataFrame] = []\n base_dir = os.path.expanduser(str(CFG.model_path))\n\n # ------------------------------------------------------------------\n # GATE MODEL CACHE (per lab_id + switch)\n # ------------------------------------------------------------------\n gate_models: dict[str, dict] = {} # action -> {\"model\": gate_model, \"num_cols\": [...], \"cat_cols\": [...]}\n gate_skipped: set[str] = set() # actions for which gate is not usable / not trainable\n\n # ------------------------------------------------------------------\n # Per-video metadata lookup for gating from TEST META\n # ------------------------------------------------------------------\n video_meta_lookup = build_video_meta_lookup_for_gating(test_subset)\n\n # =================== MAIN TEST LOOP ======================\n for switch_te, data_te, ctx_te, meta_te, actions_te in generator:\n assert switch_te == switch_tr\n try:\n fps_i = _fps_from_meta(meta_te, fps_lookup, default_fps=30.0)\n\n # Build gate feature row (for this clip), shared across actions\n gate_feature_row = build_gate_feature_row_from_chunk(\n meta_df=meta_te,\n data_df=data_te,\n ctx_df=ctx_te,\n fps_i=fps_i,\n video_meta_lookup=video_meta_lookup,\n )\n\n # Build features ONCE, keep as immutable base\n if switch_te == \"single\":\n X_base = transform_single(\n data_te, body_parts_tracked, fps_i, meta_te, ctx_te, lab_id\n )\n else:\n X_base = transform_pair(\n data_te, body_parts_tracked, fps_i, meta_te, ctx_te, lab_id\n )\n\n # raw tracking arrays no longer needed\n del data_te, ctx_te\n gc.collect()\n\n # Raw per-model per-action predictions (no smoothing yet)\n # model_tag -> { action -> np.ndarray (N,) }\n model_action_preds: dict[str, dict[str, np.ndarray]] = {\n tag: {} for tag in model_tags_eff\n }\n\n # -------- per (model, action) predictions ----------\n for action in actions_te:\n # require a threshold entry for this action in this lab\n if action not in thresholds_for_lab:\n continue\n\n for model_tag in model_tags_eff:\n # Model directory:\n # CFG.model_path / model_tag / lab_id / switch / action\n adir = os.path.join(\n base_dir,\n str(model_tag),\n str(lab_id),\n str(switch_te),\n str(action),\n )\n cols_path = os.path.join(adir, \"used_cols.pkl\")\n\n if not os.path.exists(cols_path):\n # model not trained for this (lab, switch, action)\n continue\n\n cols = joblib.load(cols_path)\n # Never mutate X_base; per-action view\n X_act = X_base.reindex(columns=cols, fill_value=0)\n\n # -------- prefer *_final_trainer_*; fallback to *_trainer_* --------\n json_final = sorted(\n glob.glob(os.path.join(adir, \"*_final_trainer.json\"))\n )\n cbm_final = sorted(\n glob.glob(os.path.join(adir, \"*_final_trainer.cbm\"))\n )\n pkl_final = sorted(\n glob.glob(os.path.join(adir, \"*_final_trainer.pkl\"))\n )\n\n if json_final or cbm_final or pkl_final:\n json_files = json_final\n cbm_files = cbm_final\n pkl_files = pkl_final\n else:\n json_files = sorted(\n glob.glob(os.path.join(adir, \"*_trainer_*.json\"))\n )\n cbm_files = sorted(\n glob.glob(os.path.join(adir, \"*_trainer_*.cbm\"))\n )\n pkl_files = sorted(\n glob.glob(os.path.join(adir, \"*_trainer_*.pkl\"))\n )\n # ------------------------------------------------------------------\n\n fold_preds: list[np.ndarray] = []\n\n # XGBoost JSON models\n if json_files:\n if xgb is None:\n raise ImportError(\n \"xgboost is required to load JSON models but is not installed.\"\n )\n X_mat = X_act.to_numpy(dtype=np.float32, copy=False)\n for path in json_files:\n booster = xgb.Booster()\n booster.load_model(path)\n yhat = booster.inplace_predict(X_mat)\n fold_preds.append(np.asarray(yhat, dtype=np.float32))\n del booster\n gc.collect()\n\n # CatBoost CBM models\n if cbm_files:\n if CatBoostClassifier is None:\n raise ImportError(\n \"catboost is required to load .cbm models but is not installed.\"\n )\n for path in cbm_files:\n cb_clf = CatBoostClassifier()\n cb_clf.load_model(path)\n yhat = cb_clf.predict_proba(X_act)[:, 1]\n fold_preds.append(np.asarray(yhat, dtype=np.float32))\n del cb_clf\n gc.collect()\n\n # Generic pickle models (if no format-specific ones were used)\n if pkl_files and not (json_files or cbm_files):\n for path in pkl_files:\n tr = joblib.load(path)\n if hasattr(tr, \"predict_proba\"):\n yhat = tr.predict_proba(X_act)[:, 1]\n else:\n yhat = tr.predict(X_act)\n fold_preds.append(np.asarray(yhat, dtype=np.float32))\n del tr\n gc.collect()\n\n if not fold_preds:\n del X_act\n gc.collect()\n continue\n\n if len(fold_preds) == 1:\n raw_pred = fold_preds[0]\n else:\n raw_pred = np.mean(\n np.column_stack(fold_preds), axis=1\n ).astype(np.float32)\n\n model_action_preds[model_tag][action] = raw_pred\n\n del X_act, fold_preds, raw_pred\n gc.collect()\n\n del X_base\n gc.collect()\n\n # -------- ensemble per action (gating + weights + smoothing) ----------\n pred_ens = pd.DataFrame(index=meta_te.video_frame)\n\n for action in actions_te:\n if action not in thresholds_for_lab:\n continue\n\n preds_list = []\n weights_list = []\n\n # per-action weight dict: {model_tag -> weight}\n w_dict = model_weights_for_lab.get(action, {})\n\n for model_tag in model_tags_eff:\n yhat = model_action_preds.get(model_tag, {}).get(action)\n if yhat is None:\n continue\n w = float(w_dict.get(model_tag, 1.0))\n preds_list.append(yhat)\n weights_list.append(w)\n\n if not preds_list:\n # no models had this action\n continue\n\n w_arr = np.asarray(weights_list, dtype=np.float32)\n sum_w = float(w_arr.sum())\n if sum_w <= 0.0:\n w_norm = np.full_like(w_arr, 1.0 / len(w_arr), dtype=np.float32)\n else:\n w_norm = w_arr / sum_w\n\n stacked = np.column_stack(preds_list) # (N, M_eff)\n raw_ens = (stacked * w_norm).sum(axis=1)\n\n # ------------------ GATING FOR THIS ACTION -------------------\n gate_info = get_or_train_gate_model_for_action(\n lab_id=lab_id,\n action_name=action,\n switch_tr=switch_tr,\n cfg=cfg,\n gate_models=gate_models,\n gate_skipped=gate_skipped,\n )\n\n if gate_info is not None and gate_feature_row is not None:\n num_cols = gate_info[\"num_cols\"]\n cat_cols = gate_info[\"cat_cols\"]\n full_row = dict(gate_feature_row)\n\n # ensure all expected columns exist\n for c in num_cols:\n full_row.setdefault(c, np.nan)\n for c in cat_cols:\n full_row.setdefault(c, \"\")\n\n X_gate = pd.DataFrame([full_row])\n\n # numeric: coerce\n for c in num_cols:\n X_gate[c] = pd.to_numeric(X_gate[c], errors=\"coerce\")\n\n # categorical: category + \"\" token (safe add_categories)\n for c in cat_cols:\n X_gate[c] = X_gate[c].astype(\"category\")\n if \"\" not in X_gate[c].cat.categories:\n X_gate[c] = X_gate[c].cat.add_categories([\"\"])\n X_gate[c] = X_gate[c].fillna(\"\")\n\n p_gate_clip = float(\n gate_info[\"model\"]\n .predict_proba(X_gate[num_cols + cat_cols])[:, 1][0]\n )\n\n gate_series = pd.Series(\n p_gate_clip,\n index=pred_ens.index, # same as meta_te.video_frame\n )\n\n # apply soft logit gating (down-weight only) on ensemble probas\n tmp_df = pd.DataFrame(\n {action: raw_ens},\n index=pred_ens.index,\n )\n tmp_gated = apply_gate_to_preds_soft_logit(\n pred=tmp_df,\n gate_prob=gate_series,\n actions=[action],\n strength=GATE_STRENGTH,\n eps=1e-6,\n only_downweight=True,\n neg_clip=GATE_NEG_CLIP_MAX_PROB,\n )\n raw_ens = tmp_gated[action].to_numpy(dtype=np.float32)\n\n del X_gate, gate_series, tmp_df, tmp_gated\n gc.collect()\n\n # ------------------ SMOOTH AFTER GATING ----------------------\n smooth_sec = float(lab_smooth_cfg.get(action, 0.0))\n if smooth_sec > 0.0:\n pred_ens[action] = _smooth(\n raw_ens, fps_i, kind=\"mean\", smooth_sec=smooth_sec\n )\n else:\n pred_ens[action] = raw_ens\n\n # -------- convert to segments ----------\n if pred_ens.shape[1] != 0:\n submission_part = predict_multiclass(\n pred_ens,\n meta_te,\n thresholds_for_lab,\n )\n submission_list.append(submission_part)\n\n del model_action_preds, pred_ens\n gc.collect()\n\n except KeyError:\n try:\n del data_te\n except Exception:\n pass\n gc.collect()\n\n return submission_list\n","metadata":{"jupyter":{"source_hidden":true},"papermill":{"duration":0.017663,"end_time":"2025-10-25T19:51:34.056443","exception":false,"start_time":"2025-10-25T19:51:34.03878","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_holdout(\n lab_id: str,\n switch_tr: str,\n thresholds_for_lab: dict,\n smooth_cfg_for_switch: dict,\n subset_df: pd.DataFrame,\n fps_lookup: dict,\n body_parts: tuple,\n model_weights_for_lab: dict | None = None,\n model_tags: list[str] | None = None,\n save_shap_data: bool = True,\n gate_models_for_lab: dict[str, dict] | None = None, # <- NEW\n) -> list[pd.DataFrame]:\n \"\"\"\n Submit-style inference for a *subset* of the TRAIN data\n (our holdout videos), using already trained CV models, ensemble\n thresholds and per-action smoothing.\n\n (docstring unchanged, omitted here for brevity)\n \"\"\"\n\n submission_list: list[pd.DataFrame] = []\n\n if subset_df is None or subset_df.empty:\n return submission_list\n\n gate_models_for_lab = gate_models_for_lab or {}\n\n # container to accumulate SHAP-ready data per action\n shap_store = defaultdict(list) if save_shap_data else None\n\n # per-lab smoothing cfg (already for this lab + switch)\n lab_smooth_cfg = smooth_cfg_for_switch or {}\n\n # Ensure weights dict is always a dict\n model_weights_for_lab = model_weights_for_lab or {}\n\n # Determine which model tags to use\n if model_tags is not None:\n model_tags_eff = list(model_tags)\n else:\n if model_weights_for_lab:\n # union of all model_tag keys across actions\n tags = set()\n for w_per_action in model_weights_for_lab.values():\n tags.update(w_per_action.keys())\n model_tags_eff = sorted(tags)\n else:\n model_tags_eff = list(getattr(CFG, \"model_tags\", []))\n\n # If still empty, nothing to ensemble => fallback to no predictions\n if not model_tags_eff:\n return submission_list\n\n base_dir = os.path.expanduser(str(CFG.model_path))\n\n # ------------------------------------------------------------\n # Build per-video metadata lookup for gating\n # ------------------------------------------------------------\n video_meta_lookup = build_video_meta_lookup_for_gating(subset_df)\n\n # only generate the requested switch\n generator = generate_mouse_data(\n subset_df,\n \"train\",\n generate_single=(switch_tr == \"single\"),\n generate_pair=(switch_tr == \"pair\"),\n )\n\n for switch_te, data_te, ctx_te, meta_te, label_te in generator:\n assert switch_te == switch_tr\n try:\n fps_i = _fps_from_meta(meta_te, fps_lookup, default_fps=30.0)\n\n # Build gate feature row (for this clip), if any gates exist\n gate_feature_row = (\n build_gate_feature_row_from_chunk(meta_te, data_te, ctx_te, fps_i, video_meta_lookup)\n if gate_models_for_lab\n else None\n )\n\n # Build features ONCE, keep as immutable base\n if switch_te == \"single\":\n X_base = transform_single(\n data_te, list(body_parts), fps_i, meta_te, ctx_te, lab_id\n )\n else:\n X_base = transform_pair(\n data_te, list(body_parts), fps_i, meta_te, ctx_te, lab_id\n )\n\n # We don't need raw tracking arrays after this\n del data_te, ctx_te\n gc.collect()\n\n actions_te = list(label_te.columns)\n\n # Raw per-model per-action predictions (no smoothing yet)\n # model_tag -> { action -> np.ndarray (N,) }\n model_action_preds: dict[str, dict[str, np.ndarray]] = {\n tag: {} for tag in model_tags_eff\n }\n\n # ---- Compute predictions per (model, action) ----\n for action in actions_te:\n # Only care about actions that have thresholds\n if action not in thresholds_for_lab:\n continue\n\n for model_tag in model_tags_eff:\n # Models are saved under:\n # CFG.model_path / model_tag / lab_id / switch / action\n adir = os.path.join(\n base_dir,\n str(model_tag),\n str(lab_id),\n str(switch_te),\n str(action),\n )\n\n cols_path = os.path.join(adir, \"used_cols.pkl\")\n if not os.path.exists(cols_path):\n # This model wasn't trained for this action\n continue\n\n cols = joblib.load(cols_path)\n X_act = X_base.reindex(columns=cols, fill_value=0)\n\n # Collect fold models\n json_files = sorted(glob.glob(os.path.join(adir, \"*_trainer_*.json\")))\n cbm_files = sorted(glob.glob(os.path.join(adir, \"*_trainer_*.cbm\")))\n pkl_files = sorted(glob.glob(os.path.join(adir, \"*_trainer_*.pkl\")))\n\n fold_preds = []\n\n # XGBoost JSON models\n if json_files:\n if xgb is None:\n raise ImportError(\n \"xgboost is required to load JSON models but is not installed.\"\n )\n X_mat = X_act.to_numpy(dtype=np.float32, copy=False)\n for path in json_files:\n booster = xgb.Booster()\n booster.load_model(path)\n yhat = booster.inplace_predict(X_mat)\n fold_preds.append(np.asarray(yhat, dtype=np.float32))\n del booster\n gc.collect()\n\n # CatBoost CBM models\n if cbm_files:\n if CatBoostClassifier is None:\n raise ImportError(\n \"catboost is required to load .cbm models but is not installed.\"\n )\n for path in cbm_files:\n cb_clf = CatBoostClassifier()\n cb_clf.load_model(path)\n yhat = cb_clf.predict_proba(X_act)[:, 1]\n fold_preds.append(np.asarray(yhat, dtype=np.float32))\n del cb_clf\n gc.collect()\n\n # Generic pickle models\n if pkl_files and not (json_files or cbm_files):\n for path in pkl_files:\n tr = joblib.load(path)\n if hasattr(tr, \"predict_proba\"):\n yhat = tr.predict_proba(X_act)[:, 1]\n else:\n yhat = tr.predict(X_act)\n fold_preds.append(np.asarray(yhat, dtype=np.float32))\n del tr\n gc.collect()\n\n if not fold_preds:\n del X_act\n gc.collect()\n continue\n\n if len(fold_preds) == 1:\n raw_pred = fold_preds[0]\n else:\n raw_pred = np.mean(\n np.column_stack(fold_preds), axis=1\n ).astype(np.float32)\n\n model_action_preds[model_tag][action] = raw_pred\n\n del X_act, fold_preds, raw_pred\n gc.collect()\n\n # We won't need X_base for model loading anymore;\n # keep it around if we need SHAP data later\n X_base_for_shap = X_base\n del X_base\n gc.collect()\n\n # ---- Build ensemble predictions per action ----\n pred_ens = pd.DataFrame(index=meta_te.video_frame)\n\n for action in actions_te:\n if action not in thresholds_for_lab:\n continue\n\n preds_list = []\n weights_list = []\n\n # per-action weight dict: {model_tag -> weight}\n w_dict = model_weights_for_lab.get(action, {})\n\n for model_tag in model_tags_eff:\n yhat = model_action_preds.get(model_tag, {}).get(action)\n if yhat is None:\n continue\n w = float(w_dict.get(model_tag, 1.0))\n preds_list.append(yhat)\n weights_list.append(w)\n\n if not preds_list:\n continue\n\n w_arr = np.asarray(weights_list, dtype=np.float32)\n sum_w = float(w_arr.sum())\n if sum_w <= 0.0:\n w_norm = np.full_like(w_arr, 1.0 / len(w_arr), dtype=np.float32)\n else:\n w_norm = w_arr / sum_w\n\n stacked = np.column_stack(preds_list) # (N, M_eff)\n raw_ens = (stacked * w_norm).sum(axis=1)\n\n # --------------- APPLY GATE (if available) ----------------\n gate_info = gate_models_for_lab.get(action) if gate_models_for_lab else None\n if gate_info is not None and gate_feature_row is not None:\n num_cols = gate_info[\"num_cols\"]\n cat_cols = gate_info[\"cat_cols\"]\n full_row = dict(gate_feature_row)\n\n for c in num_cols:\n full_row.setdefault(c, np.nan)\n for c in cat_cols:\n full_row.setdefault(c, \"\")\n\n X_gate = pd.DataFrame([full_row])\n\n for c in num_cols:\n X_gate[c] = pd.to_numeric(X_gate[c], errors=\"coerce\")\n for c in cat_cols:\n X_gate[c] = X_gate[c].astype(\"category\")\n if \"\" not in X_gate[c].cat.categories:\n X_gate[c] = X_gate[c].cat.add_categories([\"\"])\n X_gate[c] = X_gate[c].fillna(\"\")\n\n p_gate_clip = float(\n gate_info[\"model\"]\n .predict_proba(X_gate[num_cols + cat_cols])[:, 1][0]\n )\n\n gate_series = pd.Series(\n p_gate_clip,\n index=pred_ens.index, # same as meta_te.video_frame\n )\n\n tmp_df = pd.DataFrame(\n {action: raw_ens},\n index=pred_ens.index,\n )\n tmp_gated = apply_gate_to_preds_soft_logit(\n pred=tmp_df,\n gate_prob=gate_series,\n actions=[action],\n strength=GATE_STRENGTH,\n eps=1e-6,\n only_downweight=True,\n neg_clip=GATE_NEG_CLIP_MAX_PROB,\n )\n raw_ens = tmp_gated[action].to_numpy(dtype=np.float32)\n\n del X_gate, gate_series, tmp_df, tmp_gated\n gc.collect()\n\n # --------------- SMOOTH AFTER GATING ----------------------\n smooth_sec = float(lab_smooth_cfg.get(action, 0.0))\n if smooth_sec > 0.0:\n pred_ens[action] = _smooth(\n raw_ens, fps=fps_i, kind=\"mean\", smooth_sec=smooth_sec\n )\n else:\n pred_ens[action] = raw_ens\n\n # ---- SHAP-ready data ----\n if save_shap_data and pred_ens.shape[1] != 0:\n meta_small = meta_te[[\"video_id\", \"agent_id\", \"target_id\", \"video_frame\"]].reset_index(drop=True)\n X_feat = X_base_for_shap.reset_index(drop=True)\n\n for action in actions_te:\n if action not in thresholds_for_lab:\n continue\n if action not in pred_ens.columns:\n continue\n\n y_true = label_te[action].astype(int).to_numpy()\n y_pred = pred_ens[action].to_numpy()\n\n shap_chunk = pd.concat(\n [\n meta_small,\n pd.DataFrame(\n {\n \"lab_id\": lab_id,\n \"switch\": switch_tr,\n \"fps\": float(fps_i),\n \"action\": action,\n \"y_true\": y_true,\n \"y_pred\": y_pred,\n }\n ),\n X_feat,\n ],\n axis=1,\n )\n shap_store[action].append(shap_chunk)\n\n del X_base_for_shap\n gc.collect()\n\n if pred_ens.shape[1] != 0:\n submission_part = predict_multiclass(\n pred_ens,\n meta_te,\n thresholds_for_lab,\n )\n submission_list.append(submission_part)\n\n del model_action_preds, pred_ens\n gc.collect()\n\n except KeyError:\n try:\n del data_te\n except Exception:\n pass\n gc.collect()\n\n # ---- Save SHAP-ready holdout data (if requested) ----\n if save_shap_data and shap_store:\n shap_base = os.path.join(base_dir, \"shap\", str(lab_id), str(switch_tr))\n os.makedirs(shap_base, exist_ok=True)\n\n for action, chunks in shap_store.items():\n if not chunks:\n continue\n shap_df_all = pd.concat(chunks, axis=0, ignore_index=True)\n out_path = os.path.join(shap_base, f\"{action}_holdout_data.pkl\")\n joblib.dump(shap_df_all, out_path, compress=(\"xz\", 3))\n # print(\n # f\"[predict_holdout] Saved SHAP-ready holdout data for \"\n # f\"lab={lab_id}, switch={switch_tr}, action={action} -> {out_path} \"\n # f\"| shape={shap_df_all.shape}\"\n # )\n\n return submission_list\n","metadata":{"papermill":{"duration":0.017663,"end_time":"2025-10-25T19:51:34.056443","exception":false,"start_time":"2025-10-25T19:51:34.03878","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def cross_validate_classifier_xgb(\n X,\n label,\n meta,\n lab_id,\n fps=30.0,\n switch=\"\",\n model_tag=\"xgb\", # model identifier for paths/logging\n with_estop=False,\n es_rounds=50, # <--- patience for ES (rounds with no improvement)\n es_min_delta=5e-4, # <--- required improvement in logloss\n with_focal=False, # <--- toggle custom focal loss\n focal_gamma=2.0, # <--- focal loss gamma\n focal_alpha=0.5, # <--- will be overridden per action based on pos rate\n):\n start_total = time.time()\n\n def _has_both_classes(y):\n y = np.asarray(y)\n return (y == 0).any() and (y == 1).any()\n\n def _sample_eval_from_train(train_idx, y, frac=0.15, seed=42):\n \"\"\"Sample an eval subset *from train_idx* for early stopping.\"\"\"\n if frac <= 0.0 or len(train_idx) < 2:\n return None\n rng = np.random.RandomState(seed)\n train_idx = np.asarray(train_idx)\n n_eval = max(1, int(round(len(train_idx) * frac)))\n if n_eval >= len(train_idx):\n return None\n eval_idx = rng.choice(train_idx, size=n_eval, replace=False)\n return np.sort(eval_idx)\n\n # OOF will still be a DataFrame, but we’ll store float32 columns\n oof = pd.DataFrame(index=meta.video_frame)\n f1_list, thresholds = [], {}\n tot_f1 = 0.0\n\n n_rows = len(label)\n\n for action in label.columns:\n action_t0 = time.time()\n\n # --- build mask once, as numpy ---\n action_mask = label[action].notna().values\n\n # if no labeled rows at all, just fill zeros and move on\n if not action_mask.any():\n oof[action] = np.zeros(n_rows, dtype=np.float32)\n del action_mask\n gc.collect()\n continue\n\n # base series / features / groups on labeled rows\n y_series = label[action][action_mask].astype(np.int8) # Series\n X_action = X[action_mask] # DataFrame\n groups_action = meta.video_id[action_mask] # Series\n\n # ------------------------------------------------------------\n # 0-POS VIDEO FILTER: drop 0-pos videos ONLY if they are >5% of videos\n # ------------------------------------------------------------\n pos_by_video = y_series.groupby(groups_action).sum()\n videos_keep = pos_by_video[pos_by_video > 0].index\n videos_drop = pos_by_video[pos_by_video == 0].index\n\n n_total_vids = len(pos_by_video)\n n_keep = len(videos_keep)\n n_drop = len(videos_drop)\n frac_zero = (n_drop / n_total_vids) if n_total_vids > 0 else 0.0\n\n # apply 0-pos filter only if more than 5% of videos are 0-pos\n apply_zero_filter = frac_zero > 0.75\n\n if apply_zero_filter and n_keep == 0:\n # all videos have 0 positives -> nothing to learn; OOF stays 0\n oof[action] = np.zeros(n_rows, dtype=np.float32)\n train_time = time.time() - action_t0\n print(\n f\"\\t[{model_tag}] F1=0.000 P=0.000 R=0.000 thr=0.00 \"\n f\"pos=0.0% time={train_time:.1f}s X=(0, {X.shape[1]}) \"\n f\"Lab:{lab_id} Act:{action} (all videos have 0 positives) \"\n f\"| vids: {n_keep}/{n_total_vids}\"\n )\n del action_mask, X_action, y_series, groups_action, pos_by_video, videos_keep, videos_drop\n gc.collect()\n continue\n\n # keep only frames from videos that actually have at least one positive\n # IF we decided to apply the 0-pos filter; otherwise keep everything\n if apply_zero_filter:\n keep_rows = groups_action.isin(videos_keep).values\n y_action = y_series.values[keep_rows].astype(np.int8)\n X_action = X_action[keep_rows]\n groups_action = groups_action[keep_rows]\n else:\n # no filtering: use all labeled frames\n keep_rows = np.ones(len(y_series), dtype=bool)\n y_action = y_series.values.astype(np.int8)\n # X_action and groups_action already refer to all labeled rows\n\n # prune features (this returns a new DataFrame)\n X_action, _ = drop_least_important_for_action(\n X_action,\n prunable_by_action=prune_df,\n action=action,\n n_drop=0.0, # adjust if you want pruning; 0 == no importance-based drop\n verbose=False,\n )\n\n # ensure float32 to halve feature memory & xgboost usage\n X_action = X_action.astype(np.float32, copy=False)\n used_cols = list(X_action.columns)\n\n # pos_rate in % and fraction\n pos_rate = 100.0 * (y_action.mean() if len(y_action) else 0.0)\n pos_frac = float(y_action.mean() if len(y_action) else 0.0)\n\n # === focal loss config (optional, alpha as function of pos_frac) ===\n use_focal = False\n focal_obj = None\n # map pos_frac -> alpha in [0.25, 0.85], rarer => closer to 0.85\n focal_alpha_eff = focal_alpha\n if with_focal and pos_frac > 0.0 and \"make_focal_loss_obj\" in globals():\n # linear schedule up to pos_frac_hi: 0 -> 0.85, pos_frac_hi+ -> 0.25\n pos_frac_hi = 0.10 # 10% positives treated as \"common\"\n alpha_min, alpha_max = 0.25, 0.85\n rarity = 1.0 - min(pos_frac, pos_frac_hi) / pos_frac_hi\n rarity = float(np.clip(rarity, 0.0, 1.0))\n focal_alpha_eff = alpha_min + (alpha_max - alpha_min) * rarity\n\n focal_obj = make_focal_loss_obj(alpha=focal_alpha_eff, gamma=focal_gamma)\n use_focal = True\n\n # defaults so we can always write back\n oof_action = None\n sel_idx = X_action.index\n\n # === common CV config ===\n n_splits = getattr(CFG, \"n_splits\", 5)\n seed = getattr(CFG, \"seed\", 42)\n val_frac = float(getattr(CFG, \"val_nfrac\", 0.15))\n\n if not (y_action == 0).all():\n # ---------- legacy behaviour for too-few-groups ----------\n if len(np.unique(groups_action.values)) < n_splits:\n # Not enough groups for StratifiedGroupKFold -> skip training\n oof_action = np.zeros(len(y_action), dtype=np.float32)\n sel_idx = X_action.index\n train_time = time.time() - action_t0\n print(\n f\"\\t[{model_tag}] F1=0.000 P=0.000 R=0.000 thr=0.00 \"\n f\"pos={pos_rate:.1f}% time={train_time:.1f}s X={X_action.shape} \"\n f\"Lab:{lab_id} Act:{action} (too few groups for CV) \"\n f\"| vids: {n_keep}/{n_total_vids}\"\n )\n # ------------------------------------------------------------------\n else:\n try:\n with warnings.catch_warnings():\n warnings.filterwarnings(\"ignore\", category=RuntimeWarning)\n\n # ---- row limiting (actually use max_rows) ----\n max_rows = None # hook if you ever want to cap rows\n if max_rows is not None:\n X_action, y_action = limit_rows_stratified(\n X_action,\n y_action,\n max_rows=max_rows,\n )\n # keep groups aligned with the reduced X_action\n groups_action = groups_action.loc[X_action.index]\n\n sel_idx = X_action.index # final kept rows\n\n # ==== grouped + stratified CV (seeded) ====\n skf = StratifiedGroupKFold(\n n_splits=n_splits,\n shuffle=True,\n random_state=seed,\n )\n splits = list(\n skf.split(X_action, y_action, groups_action.values)\n )\n # ==========================================\n\n # base dir for saving models\n base_dir = os.path.expanduser(CFG.model_path)\n save_dir = os.path.join(base_dir, model_tag, lab_id, switch, action)\n os.makedirs(save_dir, exist_ok=True)\n\n # OOF container (for this action only)\n oof_action = np.zeros(len(X_action), dtype=np.float32)\n\n # collect per-fold best rounds when using ES\n best_rounds_action = []\n\n if len(splits) == 0:\n # no valid CV possible: leave oof_action as zeros\n print(\n f\"\\t[WARN] [{model_tag}] No valid CV splits for \"\n f\"Lab:{lab_id} Act:{action} \"\n f\"(groups={len(np.unique(groups_action.values))}); \"\n f\"using all-zero OOF. | vids: {n_keep}/{n_total_vids}\"\n )\n else:\n for fold_idx, (tr_idx, val_idx) in enumerate(splits):\n est = clone(CFG.model)\n\n # If requested and this is an XGBoost model, plug in focal loss\n if (\n use_focal\n and xgb is not None\n and isinstance(est, xgb.XGBModel)\n and focal_obj is not None\n ):\n est.set_params(\n objective=focal_obj,\n )\n\n # --- default: full train fold, no ES ---\n X_tr = X_action.iloc[tr_idx]\n y_tr = y_action[tr_idx]\n X_val = X_action.iloc[val_idx]\n fit_kwargs = {}\n\n # --- optional early stopping with eval carved from train ---\n if with_estop:\n eval_idx = _sample_eval_from_train(\n tr_idx,\n y_action,\n frac=val_frac,\n seed=seed + fold_idx,\n )\n if eval_idx is not None and len(eval_idx) > 1:\n # carve eval out of train_idx\n # eval_idx are global positions in X_action / y_action\n mask_eval = np.isin(tr_idx, eval_idx)\n real_tr_idx = tr_idx[~mask_eval]\n\n # if real_tr_idx gets too small, skip ES for this fold\n if len(real_tr_idx) > 1:\n X_tr = X_action.iloc[real_tr_idx]\n y_tr = y_action[real_tr_idx]\n X_eval = X_action.iloc[eval_idx]\n y_eval = y_action[eval_idx]\n\n fit_kwargs = dict(\n eval_set=[(X_eval, y_eval)],\n eval_metric=\"logloss\", # stay on logloss\n verbose=False,\n )\n\n # only attach callback if this is actually an XGB model\n if xgb is not None and isinstance(\n est, xgb.XGBModel\n ):\n early_cb = xgb.callback.EarlyStopping(\n rounds=es_rounds,\n min_delta=es_min_delta,\n save_best=True,\n )\n fit_kwargs[\"callbacks\"] = [early_cb]\n\n est.fit(X_tr, y_tr, **fit_kwargs)\n\n if (\n with_estop\n and hasattr(est, \"best_iteration\")\n and est.best_iteration is not None\n ):\n best_rounds_action.append(int(est.best_iteration) + 1)\n\n # probability if available, else raw prediction\n if hasattr(est, \"predict_proba\") and not use_focal:\n # normal path: trust XGBClassifier's predict_proba\n yhat = est.predict_proba(X_val)[:, 1]\n else:\n # for focal (custom objective), be explicit:\n # use raw margin + sigmoid\n if xgb is not None and isinstance(est, xgb.XGBModel):\n dval = xgb.DMatrix(\n X_val.values,\n feature_names=list(X_val.columns),\n )\n margin = est.get_booster().predict(\n dval,\n output_margin=True,\n )\n yhat = 1.0 / (1.0 + np.exp(-margin))\n else:\n # fallback for non-XGB models\n if hasattr(est, \"predict_proba\"):\n yhat = est.predict_proba(X_val)[:, 1]\n else:\n yhat = est.predict(X_val)\n\n oof_action[val_idx] = yhat.astype(np.float32, copy=False)\n\n model_base = f\"{action}_trainer_fold{fold_idx}\"\n model_json = os.path.join(save_dir, model_base + \".json\")\n model_pkl = os.path.join(save_dir, model_base + \".pkl\")\n\n # Detect XGBoost sklearn wrapper\n is_xgb = (\n (xgb is not None)\n and isinstance(\n getattr(est, \"estimator_\", est), xgb.XGBModel\n )\n if hasattr(est, \"estimator_\")\n else (xgb is not None and isinstance(est, xgb.XGBModel))\n )\n\n if is_xgb:\n # Prefer native XGB JSON format\n try:\n if hasattr(est, \"save_model\"):\n est.save_model(model_json)\n elif hasattr(est, \"get_booster\"):\n est.get_booster().save_model(model_json)\n else:\n joblib.dump(est, model_pkl)\n except Exception as e:\n print(\n f\"\\t[WARN] [{model_tag}] XGBoost save_model \"\n f\"failed ({e}); falling back to joblib.\",\n flush=True,\n )\n joblib.dump(est, model_pkl)\n else:\n # Non-XGB models: normal joblib\n joblib.dump(est, model_pkl)\n\n del est, X_tr, y_tr, X_val, yhat\n gc.collect()\n\n # --------- ONE-LINE SMOOTHING FOR THR/F1 ----------\n oof_for_thr = _smooth(\n oof_action,\n fps=fps,\n kind=\"mean\",\n smooth_sec=5.0 / 30.0,\n )\n # --------------------------------------------------\n\n # independent per-action threshold (on smoothed OOF)\n threshold = tune_threshold(oof_for_thr, y_action)\n thresholds[action] = float(threshold)\n\n preds = (oof_for_thr >= threshold)\n f1 = f1_score(y_action, preds, zero_division=0)\n tp = int(((preds == 1) & (y_action == 1)).sum())\n fp = int(((preds == 1) & (y_action == 0)).sum())\n fn = int(((preds == 0) & (y_action == 1)).sum())\n prec = (tp / (tp + fp)) if (tp + fp) > 0 else 0.0\n rec = (tp / (tp + fn)) if (tp + fn) > 0 else 0.0\n\n f1_list.append((lab_id, action, f1))\n tot_f1 += f1\n\n # store RAW OOF to disk (ensemble can re-smooth later)\n joblib.dump(\n oof_action.astype(np.float16),\n os.path.join(save_dir, \"oof_pred_probs.pkl\"),\n )\n joblib.dump(\n used_cols,\n os.path.join(save_dir, \"used_cols.pkl\"),\n )\n\n train_time = time.time() - action_t0\n\n # mean best round over folds (if ES was used)\n if with_estop and best_rounds_action:\n br_mean = int(np.round(np.mean(best_rounds_action)))\n br_str = f\" br={br_mean}\"\n else:\n br_str = \"\"\n\n focal_str = (\n f\" | focal_alpha={focal_alpha_eff:.3f}\"\n if use_focal\n else \"\"\n )\n\n print(\n f\"\\t[{model_tag}] F1={f1:.3f} P={prec:.3f} R={rec:.3f} \"\n f\"thr={threshold:.2f} pos={pos_rate:.1f}% \"\n f\"time={train_time:.1f}s X={X_action.shape} \"\n f\"Act:{action}{br_str} | vids: {n_keep}/{n_total_vids}\"\n f\"{focal_str}\"\n )\n\n except Exception as e:\n print(e)\n # keep consistent shapes; all zeros for this action\n oof_action = np.zeros(len(y_action), dtype=np.float32)\n sel_idx = X_action.index\n train_time = time.time() - action_t0\n print(\n f\"\\t[{model_tag}] F1=0.000 P=0.000 R=0.000 thr=0.00 \"\n f\"pos={pos_rate:.1f}% time={train_time:.1f}s X={X_action.shape} \"\n f\"Lab:{lab_id} Act:{action} | vids: {n_keep}/{n_total_vids}\"\n )\n else:\n # no positives -> always-zero model (after dropping 0-pos videos this is rare)\n oof_action = np.zeros(len(y_action), dtype=np.float32)\n sel_idx = X_action.index\n train_time = time.time() - action_t0\n print(\n f\"\\t[{model_tag}] F1=0.000 P=0.000 R=0.000 thr=0.00 \"\n f\"pos={pos_rate:.1f}% time={train_time:.1f}s X={X_action.shape} \"\n f\"Lab:{lab_id} Act:{action} | vids: {n_keep}/{n_total_vids}\"\n )\n\n # ---- write OOF back using kept indices (float32) ----\n oof_column = np.zeros(n_rows, dtype=np.float32)\n oof_column[sel_idx] = oof_action\n oof[action] = oof_column\n\n # memory hygiene: drop per-action big stuff ASAP\n del (\n oof_action,\n oof_column,\n action_mask,\n X_action,\n y_action,\n groups_action,\n used_cols,\n sel_idx,\n y_series,\n pos_by_video,\n videos_keep,\n videos_drop,\n keep_rows,\n )\n gc.collect()\n\n total_time = time.time() - start_total\n print(\n f\"Avg F1 of {lab_id} ({switch}, {model_tag}): \"\n f\"{tot_f1/len(label.columns):.4f} | total={total_time:.1f}s\"\n )\n\n return oof, f1_list, thresholds\n","metadata":{"papermill":{"duration":0.017663,"end_time":"2025-10-25T19:51:34.056443","exception":false,"start_time":"2025-10-25T19:51:34.03878","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def cross_validate_classifier_cb(\n X,\n label,\n meta,\n lab_id,\n fps=30.0,\n switch=\"\",\n model_tag=\"cat\", # model identifier for paths/logging\n with_estop=False,\n):\n start_total = time.time()\n\n def _has_both_classes(y):\n y = np.asarray(y)\n return (y == 0).any() and (y == 1).any()\n\n def _sample_eval_from_train(train_idx, y, frac=0.15, seed=42):\n \"\"\"Sample an eval subset *from train_idx* for early stopping.\"\"\"\n if frac <= 0.0 or len(train_idx) < 2:\n return None\n rng = np.random.RandomState(seed)\n train_idx = np.asarray(train_idx)\n n_eval = max(1, int(round(len(train_idx) * frac)))\n if n_eval >= len(train_idx):\n return None\n eval_idx = rng.choice(train_idx, size=n_eval, replace=False)\n return np.sort(eval_idx)\n\n # OOF will still be a DataFrame, but we’ll store float32 columns\n oof = pd.DataFrame(index=meta.video_frame)\n f1_list, thresholds = [], {}\n tot_f1 = 0.0\n\n n_rows = len(label)\n\n for action in label.columns:\n action_t0 = time.time()\n\n # --- build mask once, as numpy ---\n action_mask = label[action].notna().values\n\n # if no labeled rows at all, just fill zeros and move on\n if not action_mask.any():\n oof[action] = np.zeros(n_rows, dtype=np.float32)\n del action_mask\n gc.collect()\n continue\n\n # ---- base series / features / groups on labeled rows ----\n y_series = label[action][action_mask].astype(np.int8) # Series\n X_action = X[action_mask] # DataFrame\n groups_action = meta.video_id[action_mask] # Series\n\n # ------------------------------------------------------------\n # 0-POS VIDEO FILTER: drop 0-pos videos ONLY if they are >5% of videos\n # ------------------------------------------------------------\n pos_by_video = y_series.groupby(groups_action).sum()\n n_total_vids = len(pos_by_video)\n\n videos_pos = pos_by_video[pos_by_video > 0].index\n videos_zero = pos_by_video[pos_by_video == 0].index\n n_keep = len(videos_pos)\n n_drop = len(videos_zero)\n\n frac_zero = (n_drop / n_total_vids) if n_total_vids > 0 else 0.0\n apply_zero_filter = frac_zero > 0.75\n\n # if we do NOT apply the filter, all videos are effectively \"kept\"\n if not apply_zero_filter:\n n_keep = n_total_vids # every video is kept for training\n\n if apply_zero_filter and n_keep == 0:\n # all videos have 0 positives and filter is active -> nothing to learn\n oof[action] = np.zeros(n_rows, dtype=np.float32)\n train_time = time.time() - action_t0\n print(\n f\"\\t[{model_tag}] F1=0.000 P=0.000 R=0.000 thr=0.00 \"\n f\"pos=0.0% time={train_time:.1f}s X=(0, {X.shape[1]}) \"\n f\"Lab:{lab_id} Act:{action} | vids:{n_keep}/{n_total_vids}\"\n )\n del action_mask, X_action, y_series, groups_action, pos_by_video, videos_pos, videos_zero\n gc.collect()\n continue\n\n # keep only frames from videos that actually have at least one positive\n # IF we decided to apply the 0-pos filter; otherwise keep everything\n if apply_zero_filter:\n keep_rows = groups_action.isin(videos_pos).values\n else:\n keep_rows = np.ones(len(groups_action), dtype=bool)\n\n y_action = y_series.values[keep_rows].astype(np.int8)\n X_action = X_action[keep_rows]\n groups_action = groups_action[keep_rows]\n\n # prune features (this returns a new DataFrame)\n X_action, _ = drop_least_important_for_action(\n X_action,\n prunable_by_action=prune_df,\n action=action,\n n_drop=0.0, # 0 == no importance-based drop\n verbose=False,\n )\n\n # ensure float32 to halve feature memory & CatBoost usage\n X_action = X_action.astype(np.float32, copy=False)\n used_cols = list(X_action.columns)\n\n pos_rate = 100.0 * (y_action.mean() if len(y_action) else 0.0)\n\n # groups for grouped CV; Series with same index as X_action (already aligned)\n oof_action = None\n sel_idx = X_action.index\n\n # === common CV config ===\n n_splits = getattr(CFG, \"n_splits\", 5)\n seed = getattr(CFG, \"seed\", 42)\n es_rounds = int(getattr(CFG, \"es_rounds\", 50))\n val_frac = float(getattr(CFG, \"val_nfrac\", 0.15))\n\n if not (y_action == 0).all():\n # ---------- guard: too few groups ----------\n if len(np.unique(groups_action.values)) < n_splits:\n oof_action = np.zeros(len(y_action), dtype=np.float32)\n sel_idx = X_action.index\n train_time = time.time() - action_t0\n print(\n f\"\\t[{model_tag}] F1=0.000 P=0.000 R=0.000 thr=0.00 \"\n f\"pos={pos_rate:.1f}% time={train_time:.1f}s X={X_action.shape} \"\n f\"Lab:{lab_id} Act:{action} | vids:{n_keep}/{n_total_vids}\"\n )\n # -------------------------------------------------------\n else:\n try:\n with warnings.catch_warnings():\n warnings.filterwarnings(\"ignore\", category=RuntimeWarning)\n\n # ---- row limiting hook ----\n max_rows = None # e.g. 5_000_000 if you ever want to cap rows\n if max_rows is not None:\n X_action, y_action = limit_rows_stratified(\n X_action,\n y_action,\n max_rows=max_rows,\n )\n # keep groups aligned with the reduced X_action\n groups_action = groups_action.loc[X_action.index]\n\n sel_idx = X_action.index # final kept rows\n\n # ==== StratifiedGroupKFold, seeded ====\n skf = StratifiedGroupKFold(\n n_splits=n_splits,\n shuffle=True,\n random_state=seed,\n )\n splits = list(\n skf.split(X_action, y_action, groups_action.values)\n )\n # =======================================\n\n # ---------- path: CFG.model_path + model_tag ----------\n base_dir = os.path.expanduser(CFG.model_path)\n save_dir = os.path.join(base_dir, model_tag, lab_id, switch, action)\n os.makedirs(save_dir, exist_ok=True)\n # ------------------------------------------------------\n\n # OOF container (for this action only)\n oof_action = np.zeros(len(X_action), dtype=np.float32)\n\n # collect per-fold best rounds when using ES\n best_rounds_action = []\n\n if len(splits) == 0:\n # no valid CV possible: leave oof_action as zeros\n print(\n f\"\\t[WARN] [{model_tag}] No valid CV splits for \"\n f\"Lab:{lab_id} Act:{action} \"\n f\"(groups={len(np.unique(groups_action.values))}); \"\n f\"using all-zero OOF.\"\n )\n else:\n for fold_idx, (tr_idx, val_idx) in enumerate(splits):\n est = clone(CFG.cb_model)\n\n # default: full train fold, no ES\n X_tr = X_action.iloc[tr_idx]\n y_tr = y_action[tr_idx]\n X_val = X_action.iloc[val_idx]\n\n fit_kwargs = {}\n\n if with_estop:\n eval_idx = _sample_eval_from_train(\n tr_idx,\n y_action,\n frac=val_frac,\n seed=seed + fold_idx,\n )\n if eval_idx is not None and len(eval_idx) > 1:\n # carve eval out of train_idx (NO leakage)\n mask_eval = np.isin(tr_idx, eval_idx)\n real_tr_idx = tr_idx[~mask_eval]\n\n if len(real_tr_idx) > 1:\n X_tr = X_action.iloc[real_tr_idx]\n y_tr = y_action[real_tr_idx]\n X_eval = X_action.iloc[eval_idx]\n y_eval = y_action[eval_idx]\n\n # set ES params (overrides defaults, but only in ES mode)\n est.set_params(\n use_best_model=True,\n od_type=\"Iter\",\n od_wait=es_rounds,\n random_seed=seed + fold_idx,\n # metric_period can be large to reduce spam; verbose=False anyway\n metric_period=50,\n )\n\n # single eval_set (X, y) for CatBoost\n fit_kwargs = dict(\n eval_set=(X_eval, y_eval),\n verbose=False,\n )\n\n est.fit(X_tr, y_tr, **fit_kwargs)\n\n if with_estop:\n best_iter = None\n # best_iteration attribute or get_best_iteration method\n if hasattr(est, \"get_best_iteration\"):\n try:\n best_iter = est.get_best_iteration()\n except Exception:\n best_iter = None\n elif hasattr(est, \"best_iteration_\"):\n best_iter = getattr(est, \"best_iteration_\", None)\n\n if best_iter is not None and best_iter >= 0:\n # convert 0-based iteration index to \"round count\"\n best_rounds_action.append(int(best_iter) + 1)\n\n # probability if available, else raw prediction\n if hasattr(est, \"predict_proba\"):\n yhat = est.predict_proba(X_val)[:, 1]\n else:\n yhat = est.predict(X_val)\n\n oof_action[val_idx] = yhat.astype(np.float32, copy=False)\n\n model_base = f\"{action}_trainer_fold{fold_idx}\"\n model_cbm = os.path.join(save_dir, model_base + \".cbm\")\n model_pkl = os.path.join(save_dir, model_base + \".pkl\")\n\n # Prefer native CatBoost format\n try:\n est.save_model(model_cbm)\n except Exception as e:\n print(\n f\"\\t[WARN] [{model_tag}] CatBoost save_model failed \"\n f\"({e}); falling back to joblib.\",\n flush=True,\n )\n joblib.dump(est, model_pkl)\n\n del est, X_tr, y_tr, X_val, yhat\n gc.collect()\n\n # --------- ONE-LINE SMOOTHING FOR THR/F1 ----------\n oof_for_thr = _smooth(\n oof_action,\n fps=fps,\n kind=\"mean\",\n smooth_sec=5.0 / 30.0,\n )\n # --------------------------------------------------\n\n # independent per-action threshold (on smoothed OOF)\n threshold = tune_threshold(oof_for_thr, y_action)\n thresholds[action] = float(threshold)\n\n preds = (oof_for_thr >= threshold)\n f1 = f1_score(y_action, preds, zero_division=0)\n tp = int(((preds == 1) & (y_action == 1)).sum())\n fp = int(((preds == 1) & (y_action == 0)).sum())\n fn = int(((preds == 0) & (y_action == 1)).sum())\n prec = (tp / (tp + fp)) if (tp + fp) > 0 else 0.0\n rec = (tp / (tp + fn)) if (tp + fn) > 0 else 0.0\n\n f1_list.append((lab_id, action, f1))\n tot_f1 += f1\n\n # store RAW OOF to disk (ensemble can re-smooth later)\n joblib.dump(\n oof_action.astype(np.float16),\n os.path.join(save_dir, \"oof_pred_probs.pkl\"),\n )\n joblib.dump(\n used_cols,\n os.path.join(save_dir, \"used_cols.pkl\"),\n )\n\n train_time = time.time() - action_t0\n\n # mean best round over folds (if ES was used)\n if with_estop and best_rounds_action:\n br_mean = int(np.round(np.mean(best_rounds_action)))\n br_str = f\" br={br_mean}\"\n else:\n br_str = \"\"\n\n print(\n f\"\\t[{model_tag}] F1={f1:.3f} P={prec:.3f} R={rec:.3f} \"\n f\"thr={threshold:.2f} pos={pos_rate:.1f}% \"\n f\"time={train_time:.1f}s X={X_action.shape} \"\n f\"Act:{action}{br_str} | vids:{n_keep}/{n_total_vids}\"\n )\n\n except Exception as e:\n print(e)\n # keep consistent shapes; all zeros for this action\n oof_action = np.zeros(len(y_action), dtype=np.float32)\n sel_idx = X_action.index\n train_time = time.time() - action_t0\n print(\n f\"\\t[{model_tag}] F1=0.000 P=0.000 R=0.000 thr=0.00 \"\n f\"pos={pos_rate:.1f}% time={train_time:.1f}s X={X_action.shape} \"\n f\"Lab:{lab_id} Act:{action} | vids:{n_keep}/{n_total_vids}\"\n )\n else:\n # no positives -> always-zero model (after dropping 0-pos videos this is rare)\n oof_action = np.zeros(len(y_action), dtype=np.float32)\n sel_idx = X_action.index\n train_time = time.time() - action_t0\n print(\n f\"\\t[{model_tag}] F1=0.000 P=0.000 R=0.000 thr=0.00 \"\n f\"pos={pos_rate:.1f}% time={train_time:.1f}s X={X_action.shape} \"\n f\"Lab:{lab_id} Act:{action} | vids:{n_keep}/{n_total_vids}\"\n )\n\n # ---- write OOF back using kept indices (float32) ----\n oof_column = np.zeros(n_rows, dtype=np.float32)\n oof_column[sel_idx] = oof_action\n oof[action] = oof_column\n\n # memory hygiene: drop per-action big stuff ASAP\n del (\n oof_action,\n oof_column,\n action_mask,\n X_action,\n y_action,\n groups_action,\n used_cols,\n sel_idx,\n y_series,\n pos_by_video,\n videos_pos,\n videos_zero,\n keep_rows,\n )\n gc.collect()\n\n total_time = time.time() - start_total\n print(\n f\"Avg F1 of {lab_id} ({switch}, {model_tag}): \"\n f\"{tot_f1/len(label.columns):.4f} | total={total_time:.1f}s\"\n )\n\n return oof, f1_list, thresholds\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _infer_estop_rounds_for_action(\n base_dir: str,\n model_tag: str,\n lab_id: str,\n switch_tr: str,\n action: str,\n) -> int | None:\n \"\"\"\n Inspect CV fold models for (model_tag, lab_id, switch_tr, action) and\n infer an average best_round / num_iterations.\n\n Returns an integer >=1 or None if nothing usable is found.\n \"\"\"\n adir = os.path.join(\n base_dir,\n str(model_tag),\n str(lab_id),\n str(switch_tr),\n str(action),\n )\n\n rounds: list[int] = []\n\n # --- XGBoost JSON boosters ---\n json_files = glob.glob(os.path.join(adir, \"*trainer_fold*.json\"))\n if xgb is not None:\n for path in json_files:\n try:\n booster = xgb.Booster()\n booster.load_model(path)\n try:\n r = int(booster.num_boosted_rounds())\n except Exception:\n # fallback: count dumped trees\n r = int(len(booster.get_dump()))\n if r > 0:\n rounds.append(r)\n except Exception:\n continue\n\n # --- CatBoost CBM models ---\n cbm_files = glob.glob(os.path.join(adir, \"*trainer_fold*.cbm\"))\n if CatBoostClassifier is not None:\n for path in cbm_files:\n try:\n cb = CatBoostClassifier()\n cb.load_model(path)\n try:\n r = int(cb.get_tree_count())\n except Exception:\n r = int(getattr(cb, \"tree_count_\", 0))\n if r > 0:\n rounds.append(r)\n except Exception:\n continue\n\n # --- Generic pickles (XGB / Cat) ---\n pkl_files = glob.glob(os.path.join(adir, \"*trainer_fold*.pkl\"))\n for path in pkl_files:\n try:\n est_cv = joblib.load(path)\n except Exception:\n continue\n\n try:\n r = 0\n if xgb is not None and isinstance(est_cv, xgb.XGBModel):\n try:\n booster = est_cv.get_booster()\n try:\n r = int(booster.num_boosted_rounds())\n except Exception:\n r = int(len(booster.get_dump()))\n except Exception:\n r = int(getattr(est_cv, \"n_estimators\", 0))\n elif CatBoostClassifier is not None and isinstance(est_cv, CatBoostClassifier):\n try:\n r = int(est_cv.get_tree_count())\n except Exception:\n r = int(getattr(est_cv, \"tree_count_\", 0))\n\n if r > 0:\n rounds.append(r)\n finally:\n del est_cv\n\n if not rounds:\n return None\n\n avg_r = int(round(float(np.mean(rounds))))\n return max(1, avg_r)\n\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_lab(\n lab_id: str,\n switch_tr: str,\n model_tags: list[str] | None = None,\n):\n \"\"\"\n Retrain final models on the FULL train set for one lab + switch\n (no CV, no early stopping, no predictions).\n\n For each action and each model_tag, we:\n - build X/y over all labeled frames for that lab+switch\n - optionally restrict to used_cols.pkl (if present from CV)\n - fit a fresh clone of CFG.models[model_tag]\n * if 'estop' in model_tag:\n - load all CV fold models for this (lab, switch, action, model_tag)\n - infer avg best_round (trees/iterations) and\n use it as a fixed num_rounds for the full-data fit\n - save a single final model as:\n {action}_final_trainer.(json|cbm|pkl)\n under:\n CFG.model_path / model_tag / lab_id / switch / action\n \"\"\"\n\n cfg = labs_dict[lab_id]\n body_parts_tracked = list(cfg[\"body_parts\"])\n\n # restrict train to this lab's underlying dataset labs\n train_subset = train[train.lab_id.isin(cfg[\"labs\"])]\n\n if train_subset.empty:\n print(\n f\"[train_lab] No train rows found for lab {lab_id} \"\n f\"(switch={switch_tr}); skipping training.\"\n )\n return\n\n # -------- resolve model tags & mapping to estimators --------\n all_tags = list(getattr(CFG, \"model_tags\", []))\n all_models = list(getattr(CFG, \"models\", []))\n tag_to_model = {tag: m for tag, m in zip(all_tags, all_models)}\n\n if model_tags is not None:\n model_tags_eff = [t for t in model_tags if t in tag_to_model]\n else:\n model_tags_eff = all_tags\n\n if not model_tags_eff:\n print(f\"[train_lab] No valid model_tags for lab {lab_id}; skipping.\")\n return\n\n base_dir = os.path.expanduser(str(CFG.model_path))\n\n # FPS lookup\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 # --------- build full features & labels for this lab+switch ---------\n generator = generate_mouse_data(\n train_subset,\n \"train\",\n generate_single=(switch_tr == \"single\"),\n generate_pair=(switch_tr == \"pair\"),\n )\n\n feats_parts: list[pd.DataFrame] = []\n label_parts: list[pd.DataFrame] = []\n fps_last = 30.0\n\n for switch_cur, data_cur, ctx_cur, meta_cur, label_cur in generator:\n if switch_cur != switch_tr:\n del data_cur, ctx_cur, meta_cur, label_cur\n gc.collect()\n continue\n\n fps_i = _fps_from_meta(meta_cur, fps_lookup, default_fps=30.0)\n fps_last = fps_i\n\n if switch_cur == \"single\":\n X_i = transform_single(\n data_cur, body_parts_tracked, fps_i, meta_cur, ctx_cur, lab_id\n ).astype(np.float32)\n else:\n X_i = transform_pair(\n data_cur, body_parts_tracked, fps_i, meta_cur, ctx_cur, lab_id\n ).astype(np.float32)\n\n feats_parts.append(X_i)\n label_parts.append(label_cur)\n\n del data_cur, ctx_cur, meta_cur, label_cur, X_i\n gc.collect()\n\n if not feats_parts:\n print(f\"[train_lab] No feature chunks for lab {lab_id} (switch={switch_tr}); skipping.\")\n return\n\n # full features & labels (aligned RangeIndex)\n X_tr = pd.concat(feats_parts, axis=0, ignore_index=True)\n y_df = pd.concat(label_parts, axis=0, ignore_index=True)\n\n del feats_parts, label_parts\n gc.collect()\n\n actions = list(y_df.columns)\n\n print(\n f\"[train_lab] Lab {lab_id}, switch={switch_tr}: \"\n f\"X_tr shape={X_tr.shape}, actions={len(actions)}, models={model_tags_eff}\"\n )\n\n # --------- train per action, per model_tag on full data ---------\n for act in actions:\n y_raw = y_df[act].to_numpy(dtype=float)\n mask = ~np.isnan(y_raw)\n\n if not mask.any():\n print(f\"[train_lab] Lab {lab_id} {switch_tr} action '{act}': no labels, skip.\")\n continue\n\n y_action = y_raw[mask].astype(int)\n X_action_all = X_tr.loc[mask]\n\n # optional: skip if no positives\n if y_action.sum() == 0:\n print(\n f\"[train_lab] Lab {lab_id} {switch_tr} action '{act}': \"\n f\"no positive frames, skip.\"\n )\n continue\n\n for model_tag in model_tags_eff:\n base_est = tag_to_model.get(model_tag)\n if base_est is None:\n continue\n\n est = clone(base_est)\n\n # action directory: same structure as CV / submit\n adir = os.path.join(\n base_dir,\n str(model_tag),\n str(lab_id),\n str(switch_tr),\n str(act),\n )\n os.makedirs(adir, exist_ok=True)\n\n cols_path = os.path.join(adir, \"used_cols.pkl\")\n\n # Use existing used_cols from CV if available; else use all cols and save.\n if os.path.exists(cols_path):\n used_cols = joblib.load(cols_path)\n X_action = X_action_all.reindex(columns=used_cols, fill_value=0)\n else:\n X_action = X_action_all\n used_cols = list(X_action.columns)\n joblib.dump(used_cols, cols_path)\n\n # --------- estop-aware num_rounds config ---------\n use_estop = (\"estop\" in str(model_tag).lower())\n if use_estop:\n n_rounds = _infer_estop_rounds_for_action(\n base_dir=base_dir,\n model_tag=str(model_tag),\n lab_id=str(lab_id),\n switch_tr=str(switch_tr),\n action=str(act),\n )\n if n_rounds is not None:\n if xgb is not None and isinstance(est, xgb.XGBModel):\n # fixed number of trees, no ES\n est.set_params(n_estimators=n_rounds)\n print(\n f\"[train_lab] Using estop n_estimators={n_rounds} \"\n f\"for XGB model_tag={model_tag}, action={act}\"\n )\n elif CatBoostClassifier is not None and isinstance(est, CatBoostClassifier):\n est.set_params(iterations=n_rounds)\n print(\n f\"[train_lab] Using estop iterations={n_rounds} \"\n f\"for CatBoost model_tag={model_tag}, action={act}\"\n )\n else:\n print(\n f\"[train_lab] WARN: estop requested for model_tag={model_tag}, \"\n f\"action={act} but no CV fold models found; using default \"\n f\"params (no estop rounding).\"\n )\n\n # --------- fit on FULL data (no CV, no ES, no eval_set) ---------\n t0 = time.time()\n est.fit(X_action, y_action)\n train_time = time.time() - t0\n\n # --------- save a single FINAL model with *_final_trainer_* ---------\n model_base = f\"{act}_final_trainer\"\n saved = False\n\n # XGBoost native format\n if xgb is not None and isinstance(est, xgb.XGBModel):\n model_json = os.path.join(adir, model_base + \".json\")\n try:\n est.save_model(model_json)\n saved = True\n except Exception as e:\n print(\n f\"[train_lab] WARN: XGB save_model failed for \"\n f\"{lab_id}/{switch_tr}/{act}/{model_tag} ({e}); \"\n f\"falling back to joblib.\"\n )\n\n # CatBoost native format\n if (not saved) and (CatBoostClassifier is not None) and isinstance(est, CatBoostClassifier):\n model_cbm = os.path.join(adir, model_base + \".cbm\")\n try:\n est.save_model(model_cbm)\n saved = True\n except Exception as e:\n print(\n f\"[train_lab] WARN: CatBoost save_model failed for \"\n f\"{lab_id}/{switch_tr}/{act}/{model_tag} ({e}); \"\n f\"falling back to joblib.\"\n )\n\n # Generic / fallback\n if not saved:\n model_pkl = os.path.join(adir, model_base + \".pkl\")\n joblib.dump(est, model_pkl)\n\n print(\n f\"[train_lab] Trained {model_tag} for lab={lab_id}, \"\n f\"switch={switch_tr}, action={act} | \"\n f\"X_action shape={X_action.shape} | \"\n f\"frames={len(y_action)} (pos={y_action.mean()*100:.1f}%) | \"\n f\"time={train_time:.1f}s\"\n )\n\n del est, X_action\n gc.collect()\n\n # cleanup big matrices\n del X_tr, y_df\n gc.collect()\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Main training / validate loop with HOLDOUT + ENSEMBLE\n# =========================\n\nif CFG.mode != \"shap_analysis\":\n if CFG.mode in [\"train\", \"validate\"]:\n # Fresh containers for this ensemble run\n lab_thresholds = {\"single\": {}, \"pair\": {}}\n lab_smooth_cfg = {\"single\": {}, \"pair\": {}}\n # per-switch, per-lab, per-action -> {model_tag: weight}\n lab_model_weights = {\"single\": {}, \"pair\": {}}\n else:\n base_dir = os.path.expanduser(CFG.model_path)\n thr_path = os.path.join(base_dir, \"thresholds.pkl\")\n sm_path = os.path.join(base_dir, \"smooth_cfg.pkl\")\n mw_path = os.path.join(base_dir, \"model_weights.pkl\")\n\n # These are ensemble-level dicts, not per-model\n lab_thresholds = joblib.load(thr_path)\n lab_smooth_cfg = joblib.load(sm_path)\n\n # Model weights are also a single dict for the whole ensemble\n if os.path.exists(mw_path):\n lab_model_weights = joblib.load(mw_path)\n else:\n # fallback for old runs without model_weights.pkl\n lab_model_weights = {\"single\": {}, \"pair\": {}}\n\n f1_list, submission_list = [], []\n lab_comp_scores = {}\n\n for lab_id in labs_dict.keys():\n cfg = labs_dict[lab_id]\n lab_key = str(lab_id)\n\n try:\n bp = tuple(cfg[\"body_parts\"]) # tiny mem win vs list copy\n\n print(\n f\"{lab_id}: bps:{len(bp)} | frames:{cfg['frames']:,} | \"\n f\"labs:{len(cfg['labs'])} | actions:{len(cfg['actions'])} \\n\"\n )\n\n train_subset = train[train.lab_id.isin(cfg[\"labs\"])]\n test_subset = test[test.lab_id.isin(cfg[\"labs\"])]\n if CFG.mode == \"submit\" and test_subset.empty:\n print(f\"{lab_id}: no test data for this lab; skipping\\n\")\n continue\n\n # --- pure TRAIN mode: fit final models and skip rest ---\n if CFG.mode == \"train\":\n train_lab(\n lab_id=lab_id,\n switch_tr=\"single\",\n model_tags=CFG.model_tags,\n )\n train_lab(\n lab_id=lab_id,\n switch_tr=\"pair\",\n model_tags=CFG.model_tags,\n )\n del train_subset, test_subset\n gc.collect()\n continue\n\n # --- carve out holdout videos in validate mode (guarded by CFG.use_holdout) ---\n lab_solution = None\n holdout_videos = []\n train_subset_train = train_subset\n train_subset_holdout = None\n\n if CFG.mode == \"validate\" and CFG.use_holdout:\n lab_solution = solution[solution.lab_id.isin(cfg[\"labs\"])].copy()\n holdout_videos = make_lab_holdout_videos(\n train_subset_train,\n lab_solution,\n holdout_frac_videos=CFG.holdout_frac,\n min_pos_videos_for_strat=2,\n random_state=getattr(CFG, \"seed\", 42),\n )\n if holdout_videos:\n train_subset_holdout = train_subset_train[\n train_subset_train.video_id.isin(holdout_videos)\n ]\n train_subset_train = train_subset_train[\n ~train_subset_train.video_id.isin(holdout_videos)\n ]\n print(\n f\"{lab_id}: holdout videos={len(holdout_videos)} \"\n f\"({len(np.unique(train_subset_holdout.video_id))}/\"\n f\"{len(np.unique(train_subset.video_id))} videos)\"\n )\n else:\n print(f\"{lab_id}: no holdout videos selected\")\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, single_mouse_label_list, single_mouse_meta_list, single_mouse_ctx_list = [], [], [], []\n mouse_pair_list, mouse_pair_label_list, mouse_pair_meta_list, mouse_pair_ctx_list = [], [], [], []\n\n # will collect all submissions (single + pair) for this lab\n cluster_submission_list = []\n\n # --------- SPLIT SINGLE / PAIR CHUNKS ----------\n for switch, data, ctx_i, meta_df_i, label_i in generate_mouse_data(\n train_subset_train, \"train\" # <-- train only on non-holdout videos\n ):\n if switch == \"single\":\n single_mouse_list.append(data)\n single_mouse_meta_list.append(meta_df_i)\n single_mouse_label_list.append(label_i)\n single_mouse_ctx_list.append(ctx_i)\n else:\n mouse_pair_list.append(data)\n mouse_pair_meta_list.append(meta_df_i)\n mouse_pair_label_list.append(label_i)\n mouse_pair_ctx_list.append(ctx_i)\n\n # free asap\n del data, ctx_i, meta_df_i, label_i\n gc.collect()\n\n # ---------- SINGLE ----------\n if len(single_mouse_list) > 0:\n single_feats_parts = []\n fps_single = 30.0\n for data_i, meta_i, ctx_i in zip(\n single_mouse_list, single_mouse_meta_list, single_mouse_ctx_list\n ):\n fps_i = _fps_from_meta(meta_i, _fps_lookup, default_fps=30.0)\n fps_single = fps_i # last seen; ok since fps only used for smoothing\n X_i = transform_single(\n data_i, bp, fps_i, meta_i, ctx_i, lab_id\n ).astype(np.float32)\n single_feats_parts.append(X_i)\n del X_i\n gc.collect()\n\n X_tr_single = pd.concat(single_feats_parts, axis=0, ignore_index=True)\n single_mouse_label = pd.concat(\n single_mouse_label_list, axis=0, ignore_index=True\n )\n single_mouse_meta = pd.concat(\n single_mouse_meta_list, axis=0, ignore_index=True\n )\n\n del (\n single_feats_parts,\n single_mouse_list,\n single_mouse_label_list,\n single_mouse_meta_list,\n single_mouse_ctx_list,\n )\n gc.collect()\n\n if CFG.mode == \"validate\":\n actions_single = list(single_mouse_label.columns)\n\n # ---- run CV for all models and collect OOF / thresholds ----\n single_oof_by_model: dict[str, pd.DataFrame] = {}\n single_thr_by_model: dict[str, dict[str, float]] = {}\n\n for model, model_tag in zip(CFG.models, CFG.model_tags):\n use_estop = True if \"estop\" in model_tag else False\n use_focal = True if \"focal\" in model_tag else False\n if isinstance(model, XGBClassifier):\n CFG.model = model\n oof_single_m, f1_single_m, thr_single_m = cross_validate_classifier_xgb(\n X_tr_single,\n single_mouse_label,\n single_mouse_meta,\n lab_id,\n fps=fps_single,\n switch=\"single\",\n model_tag=model_tag,\n with_estop=use_estop,\n with_focal=use_focal,\n )\n else:\n CFG.cb_model = model\n oof_single_m, f1_single_m, thr_single_m = cross_validate_classifier_cb(\n X_tr_single,\n single_mouse_label,\n single_mouse_meta,\n lab_id,\n fps=fps_single,\n switch=\"single\",\n model_tag=model_tag,\n with_estop=use_estop,\n )\n\n # collect global F1 logs\n f1_list.extend(f1_single_m)\n\n # store OOF & thresholds for ensemble tuning\n single_oof_by_model[model_tag] = oof_single_m[actions_single]\n single_thr_by_model[model_tag] = thr_single_m\n\n del oof_single_m, f1_single_m, thr_single_m\n gc.collect()\n\n # ----- per-switch ensemble tuning (single) -----\n tune_t0 = time.time()\n joint_thr_single, joint_f1_single, best_smooth_single, model_weights_single = tune_ensemble(\n preds_by_model=single_oof_by_model,\n label=single_mouse_label[actions_single],\n actions=actions_single,\n thresholds_by_model=single_thr_by_model,\n model_tags=CFG.model_tags,\n seed=getattr(CFG, \"seed\", 42),\n fps=fps_single,\n base_smooth_sec=5 / 30,\n )\n tune_time = time.time() - tune_t0\n\n print(\n f\"Lab {lab_id} (single) ensemble tuning: \"\n f\"F1_joint={joint_f1_single:.4f} \"\n f\"| tuning={tune_time:.1f}s\"\n f\"| weights={model_weights_single}\"\n )\n\n # store joint thresholds for this lab / switch\n lab_thresholds[\"single\"].setdefault(lab_key, {})\n lab_thresholds[\"single\"][lab_key].update(joint_thr_single)\n\n # store best smoothing for this lab / switch\n lab_smooth_cfg[\"single\"].setdefault(lab_key, {})\n lab_smooth_cfg[\"single\"][lab_key].update(best_smooth_single)\n\n # per-action model weights: {action -> {model_tag -> w}}\n lab_model_weights[\"single\"].setdefault(lab_key, {})\n lab_model_weights[\"single\"][lab_key].update(model_weights_single)\n\n # ---- Train GATE models on non-holdout portion only (leak-free) ----\n gate_models_single: dict[str, dict] = {}\n if (\n CFG.use_holdout\n and train_subset_holdout is not None\n and holdout_videos\n ):\n for act in actions_single:\n try:\n video_tab_gate = build_video_level_table_for_action(\n lab_id=lab_id,\n action=act,\n switch=\"single\",\n train_df=train_subset_train, # <- NO holdout here\n train_tracking_dir=CFG.train_tracking_path,\n )\n except Exception as e:\n print(\n f\"[Gate model] Failed to build video_tab \"\n f\"for lab={lab_id}, action={act}, switch=single: {e}\"\n )\n continue\n\n if video_tab_gate.empty:\n continue\n\n # metadata ONLY from non-holdout subset\n video_tab_gate = add_video_metadata(video_tab_gate, train_subset_train)\n\n gate_model, num_cols, cat_cols, _ = train_action_presence_model(\n video_tab_gate,\n lab_id=lab_id,\n action=act,\n switch=\"single\",\n min_pos_rate=GATE_MIN_POS_RATE,\n min_auc=GATE_MIN_AUC,\n )\n\n if gate_model is not None:\n gate_models_single[act] = {\n \"model\": gate_model,\n \"num_cols\": num_cols,\n \"cat_cols\": cat_cols,\n }\n\n # === HOLDOUT submission for SINGLE switch (ensemble-aware + GATING) ===\n if (\n CFG.use_holdout\n and train_subset_holdout is not None\n and holdout_videos\n ):\n holdout_single_list = predict_holdout(\n lab_id=lab_id,\n switch_tr=\"single\",\n thresholds_for_lab=lab_thresholds[\"single\"].get(lab_key, {}),\n smooth_cfg_for_switch=lab_smooth_cfg[\"single\"].get(lab_key, {}),\n subset_df=train_subset_holdout,\n fps_lookup=_fps_lookup,\n body_parts=bp,\n model_weights_for_lab=lab_model_weights[\"single\"].get(lab_key, {}),\n model_tags=CFG.model_tags,\n save_shap_data=True,\n gate_models_for_lab=gate_models_single, # <- pass gates\n )\n submission_list.extend(holdout_single_list)\n cluster_submission_list.extend(holdout_single_list)\n\n del (\n single_oof_by_model,\n single_thr_by_model,\n joint_thr_single,\n best_smooth_single,\n model_weights_single,\n )\n gc.collect()\n\n else:\n # submit mode: run inference with saved thresholds + smoothing + weights\n temp_submission_list = submit_lab(\n lab_id=lab_id,\n switch_tr=\"single\",\n thresholds_for_lab=lab_thresholds[\"single\"].get(lab_key, {}),\n smooth_cfg_for_lab=lab_smooth_cfg[\"single\"].get(lab_key, {}),\n model_weights_for_lab=lab_model_weights[\"single\"].get(lab_key, {}),\n model_tags=CFG.model_tags,\n )\n submission_list.extend(temp_submission_list)\n cluster_submission_list.extend(temp_submission_list)\n del temp_submission_list\n gc.collect()\n\n del X_tr_single, single_mouse_label, single_mouse_meta\n gc.collect()\n\n # ---------- PAIR ----------\n if len(mouse_pair_list) > 0:\n pair_feats_parts = []\n fps_pair = 30.0\n for data_i, meta_i, ctx_i in zip(\n mouse_pair_list, mouse_pair_meta_list, mouse_pair_ctx_list\n ):\n fps_i = _fps_from_meta(meta_i, _fps_lookup, default_fps=30.0)\n fps_pair = fps_i\n X_i = transform_pair(\n data_i, bp, fps_i, meta_i, ctx_i, lab_id\n ).astype(np.float32)\n pair_feats_parts.append(X_i)\n del X_i\n gc.collect()\n\n X_tr_pair = pd.concat(pair_feats_parts, axis=0, ignore_index=True)\n mouse_pair_label = pd.concat(\n mouse_pair_label_list, axis=0, ignore_index=True\n )\n mouse_pair_meta = pd.concat(\n mouse_pair_meta_list, axis=0, ignore_index=True\n )\n\n del (\n pair_feats_parts,\n mouse_pair_list,\n mouse_pair_label_list,\n mouse_pair_meta_list,\n mouse_pair_ctx_list,\n )\n gc.collect()\n\n if CFG.mode == \"validate\":\n actions_pair = list(mouse_pair_label.columns)\n\n # ---- run CV for all models and collect OOF / thresholds ----\n pair_oof_by_model: dict[str, pd.DataFrame] = {}\n pair_thr_by_model: dict[str, dict[str, float]] = {}\n\n for model, model_tag in zip(CFG.models, CFG.model_tags):\n use_estop = True if \"estop\" in model_tag else False\n use_focal = True if \"focal\" in model_tag else False\n if isinstance(model, XGBClassifier):\n CFG.model = model\n oof_pair_m, f1_pair_m, thr_pair_m = cross_validate_classifier_xgb(\n X_tr_pair,\n mouse_pair_label,\n mouse_pair_meta,\n lab_id,\n fps=fps_pair,\n switch=\"pair\",\n model_tag=model_tag,\n with_estop=use_estop,\n with_focal=use_focal,\n )\n else:\n CFG.cb_model = model\n oof_pair_m, f1_pair_m, thr_pair_m = cross_validate_classifier_cb(\n X_tr_pair,\n mouse_pair_label,\n mouse_pair_meta,\n lab_id,\n fps=fps_pair,\n switch=\"pair\",\n model_tag=model_tag,\n with_estop=use_estop,\n )\n\n f1_list.extend(f1_pair_m)\n\n pair_oof_by_model[model_tag] = oof_pair_m[actions_pair]\n pair_thr_by_model[model_tag] = thr_pair_m\n\n del oof_pair_m, f1_pair_m, thr_pair_m\n gc.collect()\n\n # ----- per-switch ensemble tuning (pair) -----\n tune_t0 = time.time()\n joint_thr_pair, joint_f1_pair, best_smooth_pair, model_weights_pair = tune_ensemble(\n preds_by_model=pair_oof_by_model,\n label=mouse_pair_label[actions_pair],\n actions=actions_pair,\n thresholds_by_model=pair_thr_by_model,\n model_tags=CFG.model_tags,\n seed=getattr(CFG, \"seed\", 42),\n fps=fps_pair,\n base_smooth_sec=5 / 30,\n )\n tune_time = time.time() - tune_t0\n\n print(\n f\"Lab {lab_id} (pair) ensemble tuning: \"\n f\"F1_joint={joint_f1_pair:.4f} \"\n f\"| tuning={tune_time:.1f}s\"\n f\"| weights={model_weights_pair}\"\n )\n\n lab_thresholds[\"pair\"].setdefault(lab_key, {})\n lab_thresholds[\"pair\"][lab_key].update(joint_thr_pair)\n\n lab_smooth_cfg[\"pair\"].setdefault(lab_key, {})\n lab_smooth_cfg[\"pair\"][lab_key].update(best_smooth_pair)\n\n # per-action ensemble weights\n lab_model_weights[\"pair\"].setdefault(lab_key, {})\n lab_model_weights[\"pair\"][lab_key].update(model_weights_pair)\n\n # ---- Train GATE models on non-holdout portion only (leak-free) ----\n gate_models_pair: dict[str, dict] = {}\n if (\n CFG.use_holdout\n and train_subset_holdout is not None\n and holdout_videos\n ):\n for act in actions_pair:\n try:\n video_tab_gate = build_video_level_table_for_action(\n lab_id=lab_id,\n action=act,\n switch=\"pair\",\n train_df=train_subset_train, # <- NO holdout here\n train_tracking_dir=CFG.train_tracking_path,\n )\n except Exception as e:\n print(\n f\"[Gate model] Failed to build video_tab \"\n f\"for lab={lab_id}, action={act}, switch=pair: {e}\"\n )\n continue\n\n if video_tab_gate.empty:\n continue\n\n video_tab_gate = add_video_metadata(video_tab_gate, train_subset_train)\n\n gate_model, num_cols, cat_cols, _ = train_action_presence_model(\n video_tab_gate,\n lab_id=lab_id,\n action=act,\n switch=\"pair\",\n min_pos_rate=GATE_MIN_POS_RATE,\n min_auc=GATE_MIN_AUC,\n )\n\n if gate_model is not None:\n gate_models_pair[act] = {\n \"model\": gate_model,\n \"num_cols\": num_cols,\n \"cat_cols\": cat_cols,\n }\n\n # === HOLDOUT submission for PAIR switch (ensemble-aware + GATING) ===\n if (\n CFG.use_holdout\n and train_subset_holdout is not None\n and holdout_videos\n ):\n holdout_pair_list = predict_holdout(\n lab_id=lab_id,\n switch_tr=\"pair\",\n thresholds_for_lab=lab_thresholds[\"pair\"].get(lab_key, {}),\n smooth_cfg_for_switch=lab_smooth_cfg[\"pair\"].get(lab_key, {}),\n subset_df=train_subset_holdout,\n fps_lookup=_fps_lookup,\n body_parts=bp,\n model_weights_for_lab=lab_model_weights[\"pair\"].get(lab_key, {}),\n model_tags=CFG.model_tags,\n save_shap_data=True,\n gate_models_for_lab=gate_models_pair, # <- pass gates\n )\n submission_list.extend(holdout_pair_list)\n cluster_submission_list.extend(holdout_pair_list)\n\n del (\n pair_oof_by_model,\n pair_thr_by_model,\n joint_thr_pair,\n best_smooth_pair,\n model_weights_pair,\n )\n gc.collect()\n\n else:\n temp_submission_list = submit_lab(\n lab_id=lab_id,\n switch_tr=\"pair\",\n thresholds_for_lab=lab_thresholds[\"pair\"].get(lab_key, {}),\n smooth_cfg_for_lab=lab_smooth_cfg[\"pair\"].get(lab_key, {}),\n model_weights_for_lab=lab_model_weights[\"pair\"].get(lab_key, {}),\n model_tags=CFG.model_tags,\n )\n submission_list.extend(temp_submission_list)\n cluster_submission_list.extend(temp_submission_list)\n del temp_submission_list\n gc.collect()\n\n del X_tr_pair, mouse_pair_label, mouse_pair_meta\n gc.collect()\n\n # ---------- per-lab competition metric (NOW ON HOLDOUT ONLY) ----------\n if CFG.mode == \"validate\" and CFG.use_holdout and cluster_submission_list:\n if train_subset_holdout is not None and holdout_videos:\n lab_solution_holdout = lab_solution[\n lab_solution.video_id.isin(holdout_videos)\n ].copy()\n lab_score = score_per_lab(\n lab_id,\n cluster_submission_list,\n lab_solution_holdout,\n train_subset_holdout,\n )\n else:\n # fallback: no holdout -> old behavior (still using cluster_submission_list)\n lab_solution_full = solution[solution.lab_id.isin(cfg[\"labs\"])].copy()\n lab_score = score_per_lab(\n lab_id,\n cluster_submission_list,\n lab_solution_full,\n train_subset,\n )\n lab_comp_scores[lab_id] = lab_score\n\n # drop per-cluster intermediates\n del train_subset, train_subset_train, train_subset_holdout, _fps_lookup\n gc.collect()\n\n except Exception as e:\n print(f\"\\t{e}\")\n print(traceback.format_exc(), flush=True)\n print()\n\n # After the loop: save tuned thresholds / smoothing / ensemble weights for submit mode\n if CFG.mode == \"validate\":\n base_dir = os.path.expanduser(CFG.model_path)\n os.makedirs(base_dir, exist_ok=True)\n\n joblib.dump(lab_thresholds, os.path.join(base_dir, \"thresholds.pkl\"))\n joblib.dump(lab_smooth_cfg, os.path.join(base_dir, \"smooth_cfg.pkl\"))\n joblib.dump(lab_model_weights, os.path.join(base_dir, \"model_weights.pkl\"))\n\n # ---- aggregate per-lab competition scores ----\n if lab_comp_scores:\n print(\"\\nPer-lab competition metrics:\")\n for lab_key, score in sorted(lab_comp_scores.items()):\n print(f\" {lab_key}: {score:.4f}\")\n\n competition_score = float(np.nanmean(list(lab_comp_scores.values())))\n print(f\"\\nOverall competition metric (mean over labs): {competition_score:.4f}\")\n else:\n competition_score = float(\"nan\")\n print(\"\\nNo per-lab competition scores were computed; overall metric is NaN.\")\n\n # bare average F1 over all (lab, action, model) entries\n if f1_list:\n avg_bare_f1 = float(np.mean([x[2] for x in f1_list]))\n print(f\"Average bare F1: {avg_bare_f1:.4f}\")\n else:\n print(\"Average bare F1: n/a (no F1 entries)\")\n","metadata":{"papermill":{"duration":0.017663,"end_time":"2025-10-25T19:51:34.056443","exception":false,"start_time":"2025-10-25T19:51:34.03878","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission","metadata":{"papermill":{"duration":0.009079,"end_time":"2025-10-25T19:51:34.072884","exception":false,"start_time":"2025-10-25T19:51:34.063805","status":"completed"},"tags":[]}},{"cell_type":"code","source":"if CFG.mode == 'submit':\n if len(submission_list) > 0:\n submission = pd.concat(submission_list)\n else:\n submission = pd.DataFrame(\n dict(\n video_id=438887472,\n agent_id='mouse1',\n target_id='self',\n action='rear',\n start_frame='278',\n stop_frame='500'\n ), index=[44])\n \n submission_robust = robustify(submission, test, 'test')\n submission_robust.index.name = 'row_id'\n submission_robust.to_csv('submission.csv')\n submission.head()","metadata":{"papermill":{"duration":0.154772,"end_time":"2025-10-25T19:51:34.236441","exception":false,"start_time":"2025-10-25T19:51:34.081669","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# SHAP","metadata":{}},{"cell_type":"code","source":"if CFG.mode == \"shap_analysis\":\n import os\n import glob\n import gc\n import traceback\n from typing import Dict, Any, Optional, Union, List\n\n import numpy as np\n import pandas as pd\n import joblib\n import shap\n import matplotlib.pyplot as plt\n import seaborn as sns\n\n from sklearn.metrics import (\n f1_score, precision_score, recall_score,\n confusion_matrix, brier_score_loss,\n precision_recall_curve,\n )\n from sklearn.preprocessing import StandardScaler\n from sklearn.decomposition import PCA\n from sklearn.cluster import KMeans\n\n import xgboost as xgb\n try:\n from catboost import CatBoostClassifier\n _HAS_CATBOOST = True\n except ImportError:\n _HAS_CATBOOST = False\n\n # -------------------------------------------------------------\n # Small helper: expand '~' in paths\n # -------------------------------------------------------------\n def _expand_path(p: str) -> str:\n return os.path.abspath(os.path.expanduser(str(p)))\n\n\n # -------------------------------------------------------------\n # Split DataFrame-style holdout artefact (current format)\n # -------------------------------------------------------------\n def _split_holdout_dataframe(df: pd.DataFrame):\n \"\"\"\n Accepts a DataFrame like the ensemble holdout artefacts and tries to\n infer:\n - y (binary label)\n - probs (predicted probability from ensemble)\n - groups (video_id / clip_id / group)\n - X (feature columns = everything else)\n \"\"\"\n\n df = df.copy()\n\n label_candidates = [\"y_true\", \"y\", \"label\", \"target\"]\n proba_candidates = [\n \"oof_pred\", \"proba\", \"prob\", \"pred\", \"p\", \"p_hat\",\n \"y_pred\", \"y_hat\", \"score\"\n ]\n group_candidates = [\"video_id\", \"clip_id\", \"group\", \"groups\"]\n\n label_col = next((c for c in label_candidates if c in df.columns), None)\n proba_col = next((c for c in proba_candidates if c in df.columns), None)\n group_col = next((c for c in group_candidates if c in df.columns), None)\n\n if label_col is None:\n raise KeyError(\n \"Could not find label column in holdout DataFrame; \"\n f\"tried {label_candidates}, got columns={list(df.columns)}\"\n )\n if proba_col is None:\n raise KeyError(\n \"Could not find probability column in holdout DataFrame; \"\n f\"tried {proba_candidates}, got columns={list(df.columns)}\"\n )\n\n y = df[label_col].to_numpy(int)\n probs = df[proba_col].to_numpy(float)\n\n if group_col is not None:\n groups = df[group_col]\n else:\n groups = pd.Series(\n np.zeros(len(df), dtype=int),\n index=df.index,\n name=\"video_id\",\n )\n\n drop_cols = {label_col, proba_col}\n if group_col is not None:\n drop_cols.add(group_col)\n\n feature_cols = [c for c in df.columns if c not in drop_cols]\n X = df[feature_cols]\n\n return X, y, probs, groups\n\n\n # -------------------------------------------------------------\n # Load holdout artefact (dict OR DataFrame)\n # -------------------------------------------------------------\n def _load_holdout_action(path: str):\n \"\"\"\n Load a *_holdout_data.pkl for one (lab, switch, action).\n\n CASE 1: dict artefact (future-proof)\n Keys:\n 'X', 'y_true'/'y', 'oof_pred'/'probs'/'pred',\n optional: 'threshold', 'video_id', 'model'\n\n CASE 2: DataFrame artefact (current ensemble format)\n Columns:\n 'y_true', 'y_pred', 'video_id', etc + numeric features.\n\n We:\n - split into X / y_true / probs / groups\n - recompute a good F1 threshold on this holdout\n \"\"\"\n obj = joblib.load(path)\n\n # ---- Case 1: dict-style artefact --------------------------------\n if isinstance(obj, dict):\n if \"X\" in obj:\n X = obj[\"X\"]\n elif \"X_action\" in obj:\n X = obj[\"X_action\"]\n elif \"X_holdout\" in obj:\n X = obj[\"X_holdout\"]\n else:\n raise KeyError(f\"Holdout {path} missing 'X' key\")\n\n if not isinstance(X, pd.DataFrame):\n X = pd.DataFrame(X)\n\n if \"y_true\" in obj:\n y_true = np.asarray(obj[\"y_true\"], dtype=int)\n elif \"y\" in obj:\n y_true = np.asarray(obj[\"y\"], dtype=int)\n else:\n raise KeyError(f\"Holdout {path} missing 'y_true'/'y' key\")\n\n if \"oof_pred\" in obj:\n probs = np.asarray(obj[\"oof_pred\"], dtype=float)\n elif \"probs\" in obj:\n probs = np.asarray(obj[\"probs\"], dtype=float)\n elif \"pred\" in obj:\n probs = np.asarray(obj[\"pred\"], dtype=float)\n else:\n raise KeyError(\n f\"Holdout {path} missing 'oof_pred'/'probs'/'pred' key\"\n )\n\n thr = float(obj.get(\"threshold\", 0.5))\n\n vid = obj.get(\"video_id\", None)\n if vid is None:\n groups = pd.Series(\n np.zeros_like(y_true, dtype=int),\n name=\"video_id\",\n index=X.index if len(X) == len(y_true) else None,\n )\n else:\n groups = pd.Series(vid, name=\"video_id\")\n if len(groups) != len(y_true):\n groups = groups.iloc[:len(y_true)]\n\n model = obj.get(\"model\", None)\n\n if len(X) != len(y_true):\n X = X.iloc[:len(y_true)]\n\n return X, y_true, probs, thr, groups, model\n\n # ---- Case 2: DataFrame-style artefact ---------------------------\n elif isinstance(obj, pd.DataFrame):\n X, y_true, probs, groups = _split_holdout_dataframe(obj)\n\n if len(y_true) == 0:\n thr = 0.5\n else:\n precision, recall, thresholds = precision_recall_curve(y_true, probs)\n with np.errstate(divide=\"ignore\", invalid=\"ignore\"):\n f1_all = 2 * precision[:-1] * recall[:-1] / (\n precision[:-1] + recall[:-1] + 1e-12\n )\n if len(f1_all) > 0:\n thr = float(thresholds[np.nanargmax(f1_all)])\n else:\n thr = 0.5\n\n model = None # we will load true models separately\n\n return X, y_true, probs, thr, groups, model\n\n else:\n raise TypeError(\n f\"Holdout artefact at {path} must be dict or DataFrame, \"\n f\"got {type(obj)}\"\n )\n\n\n # -------------------------------------------------------------\n # Used-cols handling and feature alignment\n # -------------------------------------------------------------\n def _get_used_cols_for_action(\n model_tag: str,\n lab_id: str,\n switch: str,\n action: str,\n debug: bool = False,\n ) -> Optional[List[str]]:\n \"\"\"\n Load used_cols.pkl for a given (model_tag, lab_id, switch, action).\n Returns list-of-feature-names or None if not found.\n \"\"\"\n root = _expand_path(CFG.model_path)\n model_dir = os.path.join(root, str(model_tag), str(lab_id), str(switch), str(action))\n path = os.path.join(model_dir, \"used_cols.pkl\")\n\n if not os.path.exists(path):\n if debug:\n print(f\"[DEBUG _get_used_cols_for_action] No used_cols.pkl at {path}\")\n return None\n\n cols = joblib.load(path)\n if not isinstance(cols, (list, tuple)):\n if debug:\n print(f\"[DEBUG _get_used_cols_for_action] used_cols.pkl at {path} \"\n f\"has unexpected type {type(cols)}; ignoring.\")\n return None\n\n if debug:\n print(f\"[DEBUG _get_used_cols_for_action] Loaded {len(cols)} used_cols \"\n f\"from {path}\")\n return list(cols)\n\n\n def _align_X_to_model_features(\n X: pd.DataFrame,\n model_tag: str,\n lab_id: str,\n switch: str,\n action: str,\n debug: bool = False,\n ) -> pd.DataFrame:\n \"\"\"\n Align a holdout DataFrame X to the feature space used by a particular\n model (based on used_cols.pkl if available).\n\n - Always drops non-numeric columns.\n - If used_cols.pkl exists: reorders to exactly those features, adding\n missing columns as zeros and dropping extras.\n - If used_cols.pkl missing: returns numeric subset as-is.\n \"\"\"\n # start from numeric features only\n X_num = X.select_dtypes(include=[\"number\", \"bool\"]).copy()\n\n used_cols = _get_used_cols_for_action(\n model_tag, lab_id, switch, action, debug=debug\n )\n if used_cols is None:\n if debug:\n print(\"[DEBUG _align_X_to_model_features] No used_cols.pkl; \"\n \"using numeric columns only.\")\n return X_num\n\n # Add any missing columns as zeros\n missing_in_X = [c for c in used_cols if c not in X_num.columns]\n for c in missing_in_X:\n X_num[c] = 0.0\n\n extra_in_X = [c for c in X_num.columns if c not in used_cols]\n\n if debug:\n print(f\"[DEBUG _align_X_to_model_features] Aligning X to \"\n f\"{len(used_cols)} features from used_cols.pkl\")\n print(f\"[DEBUG _align_X_to_model_features] missing_in_X={missing_in_X} \"\n f\"(total {len(missing_in_X)}), extra_in_X={extra_in_X} \"\n f\"(total {len(extra_in_X)})\")\n\n # Return only used_cols, in correct order\n X_aligned = X_num[used_cols]\n if debug:\n print(f\"[DEBUG _align_X_to_model_features] X_aligned.shape=\"\n f\"{X_aligned.shape}\")\n return X_aligned\n\n\n # -------------------------------------------------------------\n # Load base model(s) for a given tag / lab / switch / action\n # -------------------------------------------------------------\n def _load_model_for_action(\n model_tag: str,\n lab_id: str,\n switch: str,\n action: str,\n debug: bool = False,\n ) -> Optional[Union[Any, List[Any]]]:\n \"\"\"\n Try to load the trained model(s) for a given base model type.\n\n Looks under:\n /////\n\n Supports:\n - *_trainer_fold*.pkl / .json / .cbm (CV models)\n - any *.json / *.cbm / *.pkl as fallback\n\n Returns:\n - single model object, or\n - list of models (e.g. folds), or\n - None if nothing found.\n \"\"\"\n root = _expand_path(CFG.model_path)\n model_dir = os.path.join(root, str(model_tag), str(lab_id), str(switch), str(action))\n\n if debug:\n print(f\"[DEBUG _load_model_for_action] model_tag={model_tag}\")\n print(f\"[DEBUG _load_model_for_action] model_dir={model_dir} \"\n f\"(exists={os.path.isdir(model_dir)})\")\n\n if not os.path.isdir(model_dir):\n print(f\"[LOAD_MODEL] Missing dir for {model_tag}/{lab_id}/{switch}/{action}: \"\n f\"{model_dir}\")\n return None\n\n if debug:\n files_here = sorted(os.listdir(model_dir))\n print(\"[DEBUG _load_model_for_action] Files in model_dir:\")\n for f in files_here:\n print(f\" {f}\")\n\n def _load_one(path: str):\n # NOTE: this function should raise if the object is *not* an estimator\n if path.endswith(\".pkl\"):\n obj = joblib.load(path)\n # Heuristic: skip obvious non-estimators\n if isinstance(obj, (np.ndarray, list, dict, pd.DataFrame)):\n if debug:\n print(f\"[DEBUG _load_model_for_action] Skipping non-estimator \"\n f\"PKL {path}; type={type(obj)}, repr={repr(obj)[:120]}\")\n raise TypeError(\"Not an estimator\")\n return obj\n elif path.endswith(\".json\"):\n # XGBoost JSON\n m = xgb.XGBClassifier()\n m.load_model(path)\n return m\n elif path.endswith(\".cbm\"):\n if not _HAS_CATBOOST:\n raise ImportError(\"CatBoost not installed, cannot load .cbm\")\n m = CatBoostClassifier()\n m.load_model(path)\n return m\n else:\n raise ValueError(f\"Unsupported model file extension: {path}\")\n\n # ------------------------------------------------------------------\n # 1) Prefer explicit *_trainer_fold* (CV models)\n # ------------------------------------------------------------------\n patterns = [\n \"*_trainer_fold*.pkl\",\n \"*_trainer_fold*.json\",\n \"*_trainer_fold*.cbm\",\n ]\n for pat in patterns:\n files = sorted(glob.glob(os.path.join(model_dir, pat)))\n if not files:\n continue\n if debug:\n print(f\"[DEBUG _load_model_for_action] Found fold files with pattern {pat}:\")\n for f in files:\n print(f\" {os.path.basename(f)}\")\n models: List[Any] = []\n for f in files:\n try:\n m = _load_one(f)\n models.append(m)\n except Exception as e:\n print(f\"[LOAD_MODEL] Failed to load fold model {f}: {e}\")\n if not models:\n # we found fold-like files, but none loaded -> try fallback later\n break\n if len(models) == 1:\n if debug:\n print(f\"[DEBUG _load_model_for_action] Loaded single fold model from \"\n f\"{files[0]}\")\n return models[0]\n if debug:\n print(f\"[DEBUG _load_model_for_action] Loaded {len(models)} fold models\")\n return models\n\n # ------------------------------------------------------------------\n # 2) Fallback: any *.json / *.cbm / *.pkl in that order\n # (but *skip* obvious non-estimator PKLs like oof_pred_probs.pkl)\n # ------------------------------------------------------------------\n fallback_patterns = [\"*.json\", \"*.cbm\", \"*.pkl\"]\n for pat in fallback_patterns:\n files = sorted(glob.glob(os.path.join(model_dir, pat)))\n if not files:\n continue\n if debug:\n print(f\"[DEBUG _load_model_for_action] Trying fallback pattern {pat}:\")\n for f in files:\n print(f\" {os.path.basename(f)}\")\n\n for f in files:\n try:\n m = _load_one(f)\n if debug:\n print(f\"[DEBUG _load_model_for_action] Successfully loaded fallback \"\n f\"model from {f}\")\n return m\n except Exception as e:\n # non-estimator pkl, or some loading error -> try next file/pattern\n print(f\"[LOAD_MODEL] Failed to load {f}: {e}\")\n # this pattern had files, but none worked: move on to next pattern\n\n print(f\"[LOAD_MODEL] No usable model files found for \"\n f\"{model_tag}/{lab_id}/{switch}/{action}\")\n return None\n\n\n # -------------------------------------------------------------\n # Predict proba for binary classification (supports lists)\n # -------------------------------------------------------------\n def _predict_proba_binary(model, X: pd.DataFrame) -> np.ndarray:\n \"\"\"\n Returns P(y=1) for:\n - single XGBoost/CatBoost/sklearn-like model\n - list of such models (averages probs)\n \"\"\"\n if isinstance(model, (list, tuple)):\n probs_list = [_predict_proba_binary(m, X) for m in model]\n return np.mean(np.vstack(probs_list), axis=0)\n\n # single model\n if hasattr(model, \"predict_proba\"):\n p = model.predict_proba(X)\n p = np.asarray(p)\n if p.ndim == 2:\n if p.shape[1] == 2:\n return p[:, 1]\n else: # weird shape, just take last column\n return p[:, -1]\n else: # (n,)\n return p.astype(float)\n\n # CatBoost without predict_proba\n if _HAS_CATBOOST and isinstance(model, CatBoostClassifier):\n p = model.predict(X, prediction_type=\"Probability\")\n p = np.asarray(p)\n if p.ndim == 2 and p.shape[1] == 2:\n return p[:, 1]\n elif p.ndim == 1:\n return p.astype(float)\n else:\n return p.ravel().astype(float)\n\n # Fallback: treat raw predictions as logits\n raw = model.predict(X)\n raw = np.asarray(raw, dtype=float).ravel()\n return 1.0 / (1.0 + np.exp(-raw))\n\n\n # -------------------------------------------------------------\n # SHAP core: single-model + multi-fold helper (GPU-aware)\n # -------------------------------------------------------------\n def _compute_shap_for_one(m, X_sub: pd.DataFrame, debug: bool = False) -> np.ndarray:\n \"\"\"\n Defensive SHAP wrapper for a *single* model, GPU-preferred:\n\n - Try GPUTreeExplainer for tree models (XGBoost Booster, etc)\n - Fall back to TreeExplainer\n - If that fails, fall back to shap.Explainer\n\n Always returns a 2D float array: (n_samples, n_features).\n \"\"\"\n if debug:\n print(\"\\n[DEBUG] _compute_shap_for_one()\")\n print(f\" model type: {type(m)}\")\n print(f\" X_sub.shape={X_sub.shape}\")\n print(f\" X_sub dtypes summary:\\n{X_sub.dtypes.head(20)}\")\n\n vals = None\n\n # ---- GPU/tree path ----\n try:\n # Try to get GPU tree explainer if available\n try:\n from shap.explainers._tree import GPUTreeExplainer as _GPUTreeExplainer\n has_gpu_explainer = True\n except Exception:\n _GPUTreeExplainer = None\n has_gpu_explainer = False\n\n # Use underlying Booster for XGBoost models if possible\n model_for_shap = m\n try:\n import xgboost as _xgb # local import to avoid hard dep\n if isinstance(m, _xgb.XGBModel) and hasattr(m, \"get_booster\"):\n model_for_shap = m.get_booster()\n elif isinstance(m, _xgb.Booster):\n model_for_shap = m\n except Exception:\n model_for_shap = m\n\n if has_gpu_explainer and _GPUTreeExplainer is not None:\n if debug:\n print(\"[DEBUG] Using GPUTreeExplainer\")\n # Use X_sub as background; fine for diagnostics\n explainer = _GPUTreeExplainer(\n model_for_shap,\n X_sub,\n feature_perturbation=\"tree_path_dependent\",\n )\n else:\n if debug:\n print(\"[DEBUG] Using TreeExplainer (CPU)\")\n explainer = shap.TreeExplainer(\n model_for_shap,\n feature_perturbation=\"tree_path_dependent\",\n )\n\n vals = explainer.shap_values(X_sub)\n\n except Exception:\n # ---- Generic fallback: shap.Explainer ----\n if debug:\n print(\"[DEBUG] Tree-based explainer path failed; \"\n \"falling back to generic shap.Explainer\")\n traceback.print_exc()\n explainer = shap.Explainer(m, X_sub)\n vals_obj = explainer(X_sub)\n if hasattr(vals_obj, \"values\"):\n vals = vals_obj.values\n else:\n vals = np.asarray(vals_obj)\n\n # SHAP can return list-of-arrays for multiclass; take last class\n if isinstance(vals, list):\n vals = vals[-1]\n\n vals = np.asarray(vals, dtype=float)\n if vals.ndim == 1:\n vals = vals.reshape(-1, 1)\n\n if debug:\n print(f\"[DEBUG] Final vals.shape={vals.shape}, dtype={vals.dtype}\")\n return vals\n\n\n def _compute_shap_values_for_model(\n model,\n X_sub: pd.DataFrame,\n debug: bool = False,\n ) -> np.ndarray:\n \"\"\"\n SHAP values for:\n - single model: (n_samples, n_features)\n - list/tuple of models (folds): averages SHAP over folds\n\n Uses GPU-accelerated tree explainer when possible.\n \"\"\"\n if isinstance(model, (list, tuple)):\n if debug:\n print(f\"[DEBUG] Computing SHAP for {len(model)} sub-models and averaging.\")\n shap_list = []\n for i, m in enumerate(model):\n if debug:\n print(f\"[DEBUG] sub-model {i}: type={type(m)}\")\n vals_m = _compute_shap_for_one(m, X_sub, debug=debug)\n shap_list.append(vals_m)\n shap_vals = np.mean(np.stack(shap_list, axis=0), axis=0)\n if debug:\n print(f\"[DEBUG] Averaged SHAP over {len(shap_list)} models; \"\n f\"shap_vals.shape={shap_vals.shape}\")\n return shap_vals\n else:\n if debug:\n print(\"[DEBUG] Computing SHAP for single model\")\n shap_vals = _compute_shap_for_one(model, X_sub, debug=debug)\n if debug:\n print(f\"[DEBUG] shap_vals.shape={shap_vals.shape}\")\n return shap_vals\n\n\n # ---------------------------------------------------------------------\n # SMALL HELPER: stratified SHAP importance on a subsample (per model)\n # ---------------------------------------------------------------------\n def _compute_shap_importance_xgb_gpu_strat(\n model,\n X_val: pd.DataFrame,\n y_val: np.ndarray,\n max_samples: int = 5000,\n min_pos_fraction: float = 0.3,\n random_state: int = 42,\n ) -> pd.DataFrame:\n \"\"\"\n Compute SHAP feature importance for a trained model using\n stratified subsampling on y_val (binary).\n\n Uses GPU-accelerated SHAP when available.\n\n If model is a list of tree models, averages |SHAP| over them.\n Returns ['feature', 'importance'] with importance = mean(|SHAP|).\n \"\"\"\n if X_val.empty:\n return pd.DataFrame(columns=[\"feature\", \"importance\"])\n\n y_val = np.asarray(y_val).astype(int)\n idx_all = np.arange(len(X_val))\n\n pos_idx = idx_all[y_val == 1]\n neg_idx = idx_all[y_val == 0]\n\n # --- stratified subsample to keep some positives ---\n if len(pos_idx) == 0:\n if len(X_val) > max_samples:\n X_sub = X_val.sample(max_samples, random_state=random_state)\n else:\n X_sub = X_val\n else:\n n_total = min(len(X_val), max_samples)\n n_pos_target = int(np.ceil(n_total * min_pos_fraction))\n n_pos = min(len(pos_idx), n_pos_target)\n n_neg = max(0, n_total - n_pos)\n\n rng = np.random.RandomState(random_state)\n\n if len(pos_idx) <= n_pos:\n pos_sample = pos_idx\n else:\n pos_sample = rng.choice(pos_idx, size=n_pos, replace=False)\n\n if len(neg_idx) <= n_neg:\n neg_sample = neg_idx\n else:\n neg_sample = rng.choice(neg_idx, size=n_neg, replace=False)\n\n keep_idx = np.concatenate([pos_sample, neg_sample])\n rng.shuffle(keep_idx)\n X_sub = X_val.iloc[keep_idx]\n\n # --- SHAP per-model, GPU-aware via _compute_shap_for_one ---\n if isinstance(model, (list, tuple)):\n acc = None\n for m in model:\n vals_m = _compute_shap_for_one(m, X_sub, debug=False)\n imp_m = np.abs(vals_m).mean(axis=0)\n if acc is None:\n acc = imp_m\n else:\n acc += imp_m\n shap_abs_mean = acc / len(model)\n else:\n vals = _compute_shap_for_one(model, X_sub, debug=False)\n shap_abs_mean = np.abs(vals).mean(axis=0)\n\n return pd.DataFrame({\n \"feature\": X_sub.columns,\n \"importance\": shap_abs_mean,\n })\n\n\n # -------------------------------------------------------------\n # Model weights handling (for ensemble SHAP)\n # -------------------------------------------------------------\n def _load_model_weights() -> Optional[dict]:\n root = _expand_path(CFG.model_path)\n path = os.path.join(root, \"model_weights.pkl\")\n if not os.path.exists(path):\n return None\n try:\n return joblib.load(path)\n except Exception as e:\n print(f\"[WEIGHTS] Could not load model_weights.pkl: {e}\")\n return None\n\n def _resolve_model_weight(\n weights_obj: dict,\n model_tag: str,\n lab_id: Optional[str] = None,\n switch: Optional[str] = None,\n action: Optional[str] = None,\n ) -> Optional[float]:\n \"\"\"\n Try a few plausible nesting patterns:\n\n 1) weights[model_tag] -> scalar\n 2) weights[lab_id][switch][action][model_tag]\n 3) weights[lab_id][action][model_tag]\n 4) weights[action][model_tag]\n\n If nothing matches, return None (=> will fall back to uniform weight\n in that context).\n \"\"\"\n if not isinstance(weights_obj, dict):\n return None\n\n # Simple global mapping\n if model_tag in weights_obj and isinstance(weights_obj[model_tag], (int, float)):\n return float(weights_obj[model_tag])\n\n # Helper: walk dict by a list of keys\n def _walk_with_keys(root, keys):\n cur = root\n for k in keys:\n if k is None:\n continue\n if isinstance(cur, dict) and k in cur:\n cur = cur[k]\n else:\n return None\n return cur\n\n # Try nested patterns\n for key_chain in [\n [lab_id, switch, action],\n [lab_id, action],\n [action],\n ]:\n cur = _walk_with_keys(weights_obj, key_chain)\n if isinstance(cur, dict) and model_tag in cur and isinstance(cur[model_tag], (int, float)):\n return float(cur[model_tag])\n\n return None\n\n\n # =====================================================================\n # 1) GLOBAL / AGGREGATED ANALYSIS (per model + optional ensemble)\n # =====================================================================\n def run_aggregated_shap_analysis(\n lab_comp_scores: Dict[str, float],\n switches=(\"single\", \"pair\"),\n shap_max_samples_per_action: int = 5000,\n shap_min_pos_fraction: float = 0.3,\n random_state: int = 42,\n top_k_global: int = 40,\n top_k_diff: int = 20,\n top_k_heatmap: int = 35,\n n_clusters: int = 4,\n save_csv_path: Optional[str] = \"feature_importance_shap_with_f1.csv\",\n debug: bool = False,\n ) -> pd.DataFrame:\n \"\"\"\n Aggregated SHAP + F1 analysis across all labs/switches/actions/models.\n\n Holdout artefacts:\n shap_root / lab_id / switch / {action}_holdout_data.pkl\n\n For each base model_tag in CFG.model_tags:\n - load model(s) from /////\n - align X to that model's used_cols.pkl\n - compute SHAP FI on the holdout X\n - compute F1 for that model (using the holdout threshold)\n - store per-model row with feature importances\n\n If model_weights.pkl is present and interpretable, we also build\n weighted-average \"ensemble\" SHAP rows per (lab, switch, action).\n \"\"\"\n root = _expand_path(CFG.model_path)\n cand1 = os.path.join(root, \"shap\")\n cand2 = os.path.join(root, \"shap\")\n shap_root = cand1 if os.path.isdir(cand1) else cand2\n\n print(f\"[AGG] Scanning SHAP holdout artefacts under: {shap_root}\")\n\n if not os.path.isdir(shap_root):\n print(f\"[AGG] SHAP root directory does not exist: {shap_root}\")\n return pd.DataFrame()\n\n # Load model weights (optional)\n weights_obj = _load_model_weights()\n\n rows = []\n model_tags = getattr(CFG, \"model_tags\", [])\n\n for lab_id in sorted(os.listdir(shap_root)):\n lab_dir = os.path.join(shap_root, lab_id)\n if not os.path.isdir(lab_dir):\n continue\n\n for switch in switches:\n switch_dir = os.path.join(lab_dir, str(switch))\n if not os.path.isdir(switch_dir):\n continue\n\n holdout_files = sorted(\n glob.glob(os.path.join(switch_dir, \"*_holdout_data.pkl\"))\n )\n if not holdout_files:\n continue\n\n for path in holdout_files:\n action = os.path.basename(path).replace(\"_holdout_data.pkl\", \"\")\n try:\n X_raw, y_true, probs_ens, threshold, _, _ = _load_holdout_action(path)\n except Exception as e:\n print(\n f\"[WARN] Failed to load holdout data for \"\n f\"{lab_id}/{switch}/{action}: {e}\"\n )\n continue\n\n if X_raw.empty or len(y_true) == 0:\n continue\n\n # Ensemble F1 on this holdout (same for all model_tags)\n preds_ens = (probs_ens >= float(threshold)).astype(int)\n f1_ens = f1_score(y_true, preds_ens, zero_division=0)\n\n # For each base model, compute SHAP FI + its own F1\n for model_tag in model_tags:\n model = _load_model_for_action(\n model_tag, lab_id, switch, action, debug=debug\n )\n if model is None:\n continue\n\n # Align X to the feature space of this model_tag\n X_action = _align_X_to_model_features(\n X_raw, model_tag, lab_id, switch, action, debug=debug\n )\n\n if X_action.empty:\n continue\n\n try:\n probs_model = _predict_proba_binary(model, X_action)\n preds_model = (probs_model >= float(threshold)).astype(int)\n f1_model = f1_score(y_true, preds_model, zero_division=0)\n\n fi_df = _compute_shap_importance_xgb_gpu_strat(\n model,\n X_action,\n y_true,\n max_samples=shap_max_samples_per_action,\n min_pos_fraction=shap_min_pos_fraction,\n random_state=random_state,\n )\n if fi_df.empty:\n continue\n\n row = {\n \"lab_id\": lab_id,\n \"switch\": switch,\n \"action\": action,\n \"model_tag\": model_tag,\n \"lab_score\": lab_comp_scores.get(lab_id, np.nan),\n \"binary F1 score\": f1_model, # F1 of THIS model\n \"ensemble_F1\": f1_ens, # F1 of final ensemble\n }\n for feat, imp in zip(fi_df[\"feature\"], fi_df[\"importance\"]):\n row[f\"{feat}_imp\"] = float(imp)\n\n rows.append(row)\n\n except Exception as e:\n print(\n f\"[WARN] SHAP importance failed for \"\n f\"{lab_id}/{switch}/{action}/{model_tag}: {e}\"\n )\n finally:\n gc.collect()\n\n if not rows:\n print(\"[INFO] No SHAP rows collected; nothing to analyze.\")\n return pd.DataFrame()\n\n fi_all_with_f1 = pd.DataFrame(rows)\n\n # Optional: build weighted \"ensemble\" SHAP rows per (lab,switch,action)\n if weights_obj is not None and not fi_all_with_f1.empty:\n feat_cols = [c for c in fi_all_with_f1.columns if c.endswith(\"_imp\")]\n ensemble_rows = []\n grouped = fi_all_with_f1.groupby([\"lab_id\", \"switch\", \"action\"])\n\n for (lab_id, switch, action), sub in grouped:\n tags = sub[\"model_tag\"].tolist()\n weights = []\n for tag in tags:\n w = _resolve_model_weight(\n weights_obj,\n model_tag=tag,\n lab_id=lab_id,\n switch=switch,\n action=action,\n )\n weights.append(w)\n\n w_arr = np.array([np.nan if w is None else float(w) for w in weights],\n dtype=float)\n if not np.isfinite(w_arr).any():\n # no usable weights for this triple\n continue\n w_arr = np.nan_to_num(w_arr, nan=0.0)\n if w_arr.sum() <= 0:\n continue\n w_arr = w_arr / w_arr.sum()\n\n sub_sorted = sub.copy()\n # align order of weights with sub rows\n # (sub is already in correct order since tags is sub[\"model_tag\"] list)\n agg_row = {\n \"lab_id\": lab_id,\n \"switch\": switch,\n \"action\": action,\n \"model_tag\": \"ensemble\",\n \"lab_score\": sub_sorted[\"lab_score\"].iloc[0],\n \"binary F1 score\": sub_sorted[\"ensemble_F1\"].iloc[0],\n \"ensemble_F1\": sub_sorted[\"ensemble_F1\"].iloc[0],\n }\n\n for col in feat_cols:\n vals = sub_sorted[col].to_numpy(float)\n agg_row[col] = float(np.dot(w_arr, vals))\n\n ensemble_rows.append(agg_row)\n\n if ensemble_rows:\n fi_all_with_f1 = pd.concat(\n [fi_all_with_f1, pd.DataFrame(ensemble_rows)],\n axis=0,\n ignore_index=True,\n )\n print(\n f\"[AGG] Added {len(ensemble_rows)} ensemble SHAP rows \"\n f\"(weighted by model_weights.pkl).\"\n )\n\n feat_cols = [c for c in fi_all_with_f1.columns if c.endswith(\"_imp\")]\n print(\"Combined FI + F1 shape:\", fi_all_with_f1.shape)\n print(\"Number of feature importance columns:\", len(feat_cols))\n\n if save_csv_path is not None:\n fi_all_with_f1.to_csv(save_csv_path, index=False)\n print(f\"[AGG] Saved fi_all_with_f1 to: {save_csv_path}\")\n\n # Use ensemble rows for global plots if available; else all rows\n fi_plot = fi_all_with_f1\n if \"model_tag\" in fi_all_with_f1.columns:\n ens_mask = fi_all_with_f1[\"model_tag\"] == \"ensemble\"\n if ens_mask.any():\n fi_plot = fi_all_with_f1[ens_mask].copy()\n\n # ---------- PLOTS ----------\n def _plot_global_feature_importance(fi, feat_cols, top_k: int = 30):\n if not feat_cols or fi.empty:\n print(\"[PLOT] No feature columns for global FI.\")\n return\n imp_mean = fi[feat_cols].mean().sort_values(ascending=False).head(top_k)\n plt.figure(figsize=(10, max(4, 0.25 * len(imp_mean))))\n sns.barplot(\n x=imp_mean.values,\n y=imp_mean.index,\n orient=\"h\",\n )\n plt.title(f\"Global top-{top_k} SHAP feature importances (mean |SHAP|)\")\n plt.xlabel(\"Mean importance across labs/actions/switches\")\n plt.ylabel(\"Feature\")\n plt.tight_layout()\n plt.show()\n\n def _plot_good_vs_bad_lab_diff(\n fi, feat_cols, q_low: float = 0.25, q_high: float = 0.75, top_k: int = 25\n ):\n if \"lab_score\" not in fi.columns or not feat_cols or fi.empty:\n print(\"[PLOT] Missing lab_score or feature columns for good/bad diff.\")\n return\n\n lab_scores = fi.groupby(\"lab_id\")[\"lab_score\"].mean()\n X_fi_lab = fi.groupby(\"lab_id\")[feat_cols].mean()\n\n low_thr = lab_scores.quantile(q_low)\n high_thr = lab_scores.quantile(q_high)\n\n bad_labs = lab_scores[lab_scores <= low_thr].index\n good_labs = lab_scores[lab_scores >= high_thr].index\n\n print(f\"Bad labs (<= q{int(q_low * 100)}):\", list(bad_labs))\n print(f\"Good labs (>= q{int(q_high * 100)}):\", list(good_labs))\n\n if len(bad_labs) == 0 or len(good_labs) == 0:\n print(\"[WARN] Not enough labs in good/bad groups; skipping diff plot.\")\n return\n\n mean_bad = X_fi_lab.loc[bad_labs].mean(axis=0)\n mean_good = X_fi_lab.loc[good_labs].mean(axis=0)\n\n delta = (mean_bad - mean_good).sort_values(ascending=False)\n\n top_pos = delta.head(top_k)\n top_neg = delta.tail(top_k)\n\n df_diff = pd.concat([top_pos, top_neg])\n df_diff = df_diff.sort_values(ascending=True)\n\n plt.figure(figsize=(10, max(4, 0.25 * len(df_diff))))\n sns.barplot(\n x=df_diff.values,\n y=df_diff.index,\n orient=\"h\",\n palette=[\"tab:red\" if v > 0 else \"tab:blue\" for v in df_diff.values],\n )\n plt.axvline(0.0, color=\"black\", linewidth=1)\n plt.title(\n f\"Feature importance difference (bad labs - good labs)\\n\"\n f\"(lab_score quantiles: low={low_thr:.3f}, high={high_thr:.3f})\"\n )\n plt.xlabel(\"Mean importance (bad - good)\")\n plt.ylabel(\"Feature\")\n plt.tight_layout()\n plt.show()\n\n def _plot_lab_feature_heatmap(\n fi, feat_cols, top_k_features: int = 30, normalize_rows: bool = True\n ):\n if not feat_cols or fi.empty:\n print(\"[PLOT] No feature columns for heatmap.\")\n return\n\n X_fi_lab = fi.groupby(\"lab_id\")[feat_cols].mean()\n\n feat_global_mean = X_fi_lab.mean().sort_values(ascending=False)\n top_feats = feat_global_mean.head(top_k_features).index.tolist()\n mat = X_fi_lab[top_feats].copy()\n\n if normalize_rows:\n row_sums = mat.sum(axis=1).replace(0, np.nan)\n mat = mat.div(row_sums, axis=0)\n\n plt.figure(figsize=(1.2 * len(top_feats), 0.4 * len(mat)))\n sns.heatmap(\n mat,\n cmap=\"viridis\",\n cbar_kws={\n \"label\": \"Relative importance\"\n if normalize_rows\n else \"Mean importance\"\n },\n )\n plt.title(\"Labs × top features (SHAP mean importance)\")\n plt.xlabel(\"Feature\")\n plt.ylabel(\"Lab\")\n plt.tight_layout()\n plt.show()\n\n def _cluster_labs_in_fi_space(\n fi, feat_cols, n_clusters: int = 3, random_state: int = 42\n ):\n if not feat_cols or fi.empty:\n print(\"[PLOT] No feature columns for clustering.\")\n return\n\n from IPython.display import display # type: ignore\n\n lab_scores = fi.groupby(\"lab_id\")[[\"lab_score\", \"binary F1 score\"]].mean()\n X_fi_lab = fi.groupby(\"lab_id\")[feat_cols].mean().fillna(0.0)\n\n scaler = StandardScaler()\n X_scaled = scaler.fit_transform(X_fi_lab.values)\n\n km = KMeans(n_clusters=n_clusters, random_state=random_state, n_init=\"auto\")\n clusters = km.fit_predict(X_scaled)\n\n pca = PCA(n_components=2, random_state=random_state)\n X_pca = pca.fit_transform(X_scaled)\n\n lab_emb = pd.DataFrame(\n X_pca,\n columns=[\"PC1\", \"PC2\"],\n index=X_fi_lab.index,\n )\n lab_emb[\"cluster\"] = clusters\n lab_emb[\"lab_score\"] = lab_scores[\"lab_score\"]\n lab_emb[\"binary_F1_mean\"] = lab_scores[\"binary F1 score\"]\n\n print(\"Explained variance ratio (PCA):\", pca.explained_variance_ratio_)\n\n plt.figure(figsize=(8, 6))\n sns.scatterplot(\n data=lab_emb,\n x=\"PC1\",\n y=\"PC2\",\n hue=\"cluster\",\n size=\"lab_score\",\n sizes=(50, 300),\n palette=\"tab10\",\n )\n for lab, row in lab_emb.iterrows():\n plt.text(\n row[\"PC1\"],\n row[\"PC2\"],\n lab,\n fontsize=8,\n ha=\"center\",\n va=\"center\",\n )\n plt.title(\"Labs clustered in SHAP importance space (PCA)\")\n plt.xlabel(\"PC1\")\n plt.ylabel(\"PC2\")\n plt.tight_layout()\n plt.show()\n\n display(\n lab_emb[[\"cluster\", \"lab_score\", \"binary_F1_mean\"]]\n .sort_values([\"cluster\", \"lab_score\"], ascending=[True, False])\n )\n\n # ---- actually call plots (using ensemble rows if present) ----\n feat_cols_plot = [c for c in fi_plot.columns if c.endswith(\"_imp\")]\n _plot_global_feature_importance(fi_plot, feat_cols_plot, top_k=top_k_global)\n _plot_good_vs_bad_lab_diff(fi_plot, feat_cols_plot, top_k=top_k_diff)\n _plot_lab_feature_heatmap(fi_plot, feat_cols_plot, top_k_features=top_k_heatmap)\n _cluster_labs_in_fi_space(fi_plot, feat_cols_plot, n_clusters=n_clusters)\n\n return fi_all_with_f1\n\n\n # =====================================================================\n # 2) PER-ACTION DIAGNOSTICS USING HOLDOUT DATA (ENSEMBLE SHAP)\n # =====================================================================\n def diagnose_action_failure(\n lab_id: str,\n action: str,\n switch: str = \"\",\n model_tag: Optional[str] = None,\n max_shap_samples_per_group: int = 200,\n random_state: int = 42,\n print_top_k: int = 15,\n debug: bool = False,\n ) -> Dict[str, Any]:\n \"\"\"\n Debug helper for a given (lab_id, action, switch).\n\n Holdout artefact:\n shap_root / lab_id / switch / {action}_holdout_data.pkl\n\n This simplified version:\n - computes metrics using the *ensemble* probabilities in the holdout\n - samples up to max_shap_samples_per_group examples from the holdout\n (uniformly; no TP/FP/FN breakdown)\n - computes SHAP on that subset:\n * either for a single base model_tag (if model_tag != 'ensemble')\n * or as a weighted ensemble over CFG.model_tags (default)\n\n Returns:\n - metrics\n - per_video breakdown\n - SHAP values for the chosen subset (single group)\n \"\"\"\n rng = np.random.RandomState(random_state)\n\n # Default: ensemble mode\n if model_tag is None:\n model_tag = \"ensemble\"\n\n # -------- resolve SHAP root ----------\n root = _expand_path(CFG.model_path)\n cand1 = os.path.join(root, \"shap\")\n cand2 = os.path.join(root, \"shap\")\n shap_root = cand1 if os.path.isdir(cand1) else cand2\n\n if debug:\n print(f\"[DEBUG] model_path resolved to: {root}\")\n print(f\"[DEBUG] shap_root candidates: {cand1} (exists={os.path.isdir(cand1)}), \"\n f\"{cand2} (exists={os.path.isdir(cand2)})\")\n print(f\"[DEBUG] Using shap_root={shap_root}\")\n\n base_dir = os.path.join(shap_root, str(lab_id), str(switch))\n holdout_path = os.path.join(base_dir, f\"{action}_holdout_data.pkl\")\n\n if debug:\n print(f\"[DEBUG] lab_id={lab_id}, action={action}, switch={switch}\")\n print(f\"[DEBUG] base_dir={base_dir}\")\n print(f\"[DEBUG] holdout_path={holdout_path}\")\n\n if not os.path.exists(holdout_path):\n raise FileNotFoundError(f\"Missing holdout_data.pkl at {holdout_path}\")\n\n # ---- load holdout artefact (ensemble preds) ----\n X_raw, y_action, probs_ens, threshold, groups_action, _ = _load_holdout_action(\n holdout_path\n )\n\n if debug:\n print(f\"[DEBUG] Loaded holdout artefact:\")\n print(f\" X_raw.shape={X_raw.shape}\")\n print(f\" y_action.shape={y_action.shape}, dtype={y_action.dtype}\")\n print(f\" probs_ens.shape={probs_ens.shape}, dtype={probs_ens.dtype}\")\n print(f\" threshold={threshold}\")\n print(f\" groups_action.shape={groups_action.shape}, \"\n f\"unique videos={groups_action.nunique()}\")\n\n if len(X_raw) != len(y_action):\n raise ValueError(\n f\"Length mismatch: X_raw={len(X_raw)}, \"\n f\"y_action={len(y_action)}. Check holdout artefact at {holdout_path}.\"\n )\n\n # ---- basic metrics (ensemble) ----\n preds = (probs_ens >= float(threshold)).astype(int)\n\n f1 = f1_score(y_action, preds, zero_division=0)\n prec = precision_score(y_action, preds, zero_division=0)\n rec = recall_score(y_action, preds, zero_division=0)\n cm = confusion_matrix(y_action, preds, labels=[0, 1])\n brier = brier_score_loss(y_action, probs_ens)\n\n tn, fp, fn, tp = cm.ravel()\n pos_rate = y_action.mean()\n\n print(f\"\\n=== Diagnostics for lab={lab_id}, action={action}, switch={switch} ===\")\n print(\n f\"Pos rate: {pos_rate*100:.2f}% \"\n f\"(n_pos={y_action.sum()}, n_neg={(y_action==0).sum()})\"\n )\n print(f\"F1 (ensemble)={f1:.3f} Precision={prec:.3f} Recall={rec:.3f}\")\n print(f\"Confusion matrix [0,1]:\\n{cm}\")\n print(f\"Brier score (calibration, ensemble): {brier:.4f}\")\n print(f\"Threshold used (ensemble): {threshold:.3f}\")\n\n # ---- per-video breakdown (by ensemble outcome) ----\n vid = groups_action.values\n df_eval = pd.DataFrame(\n {\n \"video_id\": vid,\n \"y\": y_action,\n \"pred\": preds,\n \"proba\": probs_ens,\n }\n )\n\n per_vid = (\n df_eval.groupby(\"video_id\")\n .apply(\n lambda g: pd.Series(\n {\n \"n\": len(g),\n \"pos_rate\": g[\"y\"].mean(),\n \"precision\": precision_score(\n g[\"y\"], g[\"pred\"], zero_division=0\n ),\n \"recall\": recall_score(g[\"y\"], g[\"pred\"], zero_division=0),\n \"f1\": f1_score(g[\"y\"], g[\"pred\"], zero_division=0),\n }\n )\n )\n .sort_values(\"f1\")\n )\n\n print(\"\\n--- Worst videos by F1 (ensemble) ---\")\n print(per_vid.head(5))\n print(\"\\n--- Best videos by F1 (ensemble) ---\")\n print(per_vid.tail(5))\n\n # ================== SHAP SUBSET SELECTION ==================\n n_total = len(y_action)\n idx_all = np.arange(n_total)\n\n if max_shap_samples_per_group is None or max_shap_samples_per_group >= n_total:\n idx_interest = idx_all\n else:\n idx_interest = rng.choice(\n idx_all, size=max_shap_samples_per_group, replace=False\n )\n\n # ===========================================================\n # CASE 1: Single-model SHAP (model_tag != 'ensemble')\n # ===========================================================\n if model_tag is not None and model_tag != \"ensemble\":\n model = _load_model_for_action(model_tag, lab_id, switch, action, debug=debug)\n print(f\"[DEBUG] Loaded model for model_tag={model_tag}: type={type(model)}\")\n\n if model is None:\n print(\n f\"[INFO] No model found for model_tag={model_tag}; \"\n \"skipping SHAP error analysis.\"\n )\n return {\n \"f1_ensemble\": f1,\n \"prec_ensemble\": prec,\n \"rec_ensemble\": rec,\n \"cm_ensemble\": cm,\n \"per_video\": per_vid,\n \"oof_ensemble\": probs_ens,\n \"y\": y_action,\n }\n\n X_full = _align_X_to_model_features(\n X_raw, model_tag, lab_id, switch, action, debug=debug\n )\n X_shap = X_full.iloc[idx_interest]\n\n try:\n shap_vals = _compute_shap_values_for_model(\n model, X_shap, debug=debug\n )\n except Exception as e:\n print(f\"[WARN] SHAP failed for model_tag={model_tag}: {e}\")\n if debug:\n print(\"[DEBUG] Full traceback for SHAP failure:\")\n traceback.print_exc()\n return {\n \"f1_ensemble\": f1,\n \"prec_ensemble\": prec,\n \"rec_ensemble\": rec,\n \"cm_ensemble\": cm,\n \"per_video\": per_vid,\n \"oof_ensemble\": probs_ens,\n \"y\": y_action,\n }\n\n feat_names = list(X_shap.columns)\n\n # === Global mean |SHAP| over selected subset ===\n mean_abs = np.abs(shap_vals).mean(axis=0)\n df_all = pd.DataFrame(\n {\n \"feature\": feat_names,\n \"mean_abs_shap\": mean_abs,\n }\n ).sort_values(\"mean_abs_shap\", ascending=False)\n\n print(f\"\\n[SHAP {model_tag}] Top {print_top_k} features (mean |SHAP|):\")\n print(df_all.head(print_top_k).to_string(index=False))\n\n # Summary plots (optional)\n try:\n print(\"\\n[SHAP] Global violin summary (selected subset)\")\n shap.summary_plot(\n shap_vals,\n X_shap,\n feature_names=feat_names,\n plot_type=\"violin\",\n max_display=20,\n show=True,\n )\n\n print(\"\\n[SHAP] Global bar summary (selected subset)\")\n shap.summary_plot(\n shap_vals,\n X_shap,\n feature_names=feat_names,\n plot_type=\"bar\",\n max_display=20,\n show=True,\n )\n except Exception as e:\n print(f\"[WARN] Could not generate SHAP plots: {e}\")\n if debug:\n traceback.print_exc()\n\n return {\n \"f1_ensemble\": f1,\n \"prec_ensemble\": prec,\n \"rec_ensemble\": rec,\n \"cm_ensemble\": cm,\n \"per_video\": per_vid,\n \"oof_ensemble\": probs_ens,\n \"y\": y_action,\n \"X_action\": X_raw,\n \"shap_values\": shap_vals,\n \"shap_index\": idx_interest,\n \"feature_names\": feat_names,\n }\n\n # ===========================================================\n # CASE 2: ENSEMBLE SHAP (default, model_tag == 'ensemble')\n # ===========================================================\n weights_obj = _load_model_weights()\n base_tags = getattr(CFG, \"model_tags\", [])\n\n models_info = [] # list of dicts {tag, model, X_full, weight_raw}\n for tag in base_tags:\n model = _load_model_for_action(tag, lab_id, switch, action, debug=debug)\n if model is None:\n continue\n X_full_tag = _align_X_to_model_features(\n X_raw, tag, lab_id, switch, action, debug=debug\n )\n if X_full_tag.empty:\n continue\n\n w_raw = None\n if weights_obj is not None:\n w_raw = _resolve_model_weight(\n weights_obj,\n model_tag=tag,\n lab_id=lab_id,\n switch=switch,\n action=action,\n )\n\n models_info.append(\n {\n \"tag\": tag,\n \"model\": model,\n \"X_full\": X_full_tag,\n \"weight_raw\": w_raw,\n }\n )\n\n if not models_info:\n print(\"[INFO] No base models loaded; cannot compute ensemble SHAP.\")\n return {\n \"f1_ensemble\": f1,\n \"prec_ensemble\": prec,\n \"rec_ensemble\": rec,\n \"cm_ensemble\": cm,\n \"per_video\": per_vid,\n \"oof_ensemble\": probs_ens,\n \"y\": y_action,\n }\n\n # Resolve / normalize weights\n w_list = []\n for info in models_info:\n if info[\"weight_raw\"] is None:\n w_list.append(np.nan)\n else:\n w_list.append(float(info[\"weight_raw\"]))\n w_arr = np.array(w_list, dtype=float)\n\n if not np.isfinite(w_arr).any():\n # Fallback: uniform weights\n w_arr = np.ones(len(models_info), dtype=float)\n else:\n w_arr = np.nan_to_num(w_arr, nan=0.0)\n if w_arr.sum() <= 0:\n w_arr = np.ones(len(models_info), dtype=float)\n w_arr = w_arr / w_arr.sum()\n\n # Build canonical feature space across all models\n feat_union = set()\n for info in models_info:\n feat_union.update(info[\"X_full\"].columns)\n feat_names = sorted(feat_union)\n n_feats = len(feat_names)\n feat_idx_global = {f: i for i, f in enumerate(feat_names)}\n\n # Build X_shap_ensemble for plotting on selected subset\n X_all_num = X_raw.select_dtypes(include=[\"number\", \"bool\"])\n for f_name in feat_names:\n if f_name not in X_all_num.columns:\n X_all_num[f_name] = 0.0\n X_all_num = X_all_num[feat_names]\n X_shap_ensemble = X_all_num.iloc[idx_interest].copy()\n\n # Compute ensemble SHAP: weighted combination of per-model SHAP\n shap_ensemble = np.zeros((len(idx_interest), n_feats), dtype=float)\n\n for (info, w) in zip(models_info, w_arr):\n X_full_tag = info[\"X_full\"]\n X_shap_tag = X_full_tag.iloc[idx_interest]\n try:\n vals_tag = _compute_shap_values_for_model(\n info[\"model\"], X_shap_tag, debug=debug\n ) # (n_samples, n_feat_tag)\n except Exception as e:\n print(f\"[WARN] SHAP failed for model_tag={info['tag']} in ensemble: {e}\")\n if debug:\n traceback.print_exc()\n continue\n\n feat_tag = list(X_full_tag.columns)\n if vals_tag.shape[1] != len(feat_tag):\n print(\n f\"[WARN] SHAP values shape mismatch for {info['tag']} \"\n f\"(got {vals_tag.shape[1]} cols, expected {len(feat_tag)}); skipping.\"\n )\n continue\n\n idx_global = [feat_idx_global[f] for f in feat_tag]\n contrib = np.zeros_like(shap_ensemble)\n contrib[:, idx_global] = vals_tag\n shap_ensemble += w * contrib\n\n # === Global mean |SHAP| over selected subset ===\n mean_abs_all = np.abs(shap_ensemble).mean(axis=0)\n df_all = pd.DataFrame(\n {\n \"feature\": feat_names,\n \"mean_abs_shap\": mean_abs_all,\n }\n ).sort_values(\"mean_abs_shap\", ascending=False)\n\n print(f\"\\n[SHAP ensemble] Top {print_top_k} features over selected subset:\")\n print(df_all.head(print_top_k).to_string(index=False))\n\n # ------------------------------------------------------------------\n # Extra SHAP plots (ensemble): violin + bar + optional heatmap\n # ------------------------------------------------------------------\n try:\n # 1) Global violin summary\n print(\"\\n[SHAP ensemble] Global violin summary (selected subset)\")\n shap.summary_plot(\n shap_ensemble,\n X_shap_ensemble,\n feature_names=feat_names,\n plot_type=\"violin\",\n max_display=20,\n show=True,\n )\n\n # # 2) Global bar summary (mean |SHAP|)\n # print(\"\\n[SHAP ensemble] Global bar summary (selected subset)\")\n # shap.summary_plot(\n # shap_ensemble,\n # X_shap_ensemble,\n # feature_names=feat_names,\n # plot_type=\"bar\",\n # max_display=20,\n # show=True,\n # )\n\n # 3) Optional: heatmap for one random video (using selected subset)\n try:\n vid_interest = groups_action.iloc[idx_interest].values\n unique_vids = np.unique(vid_interest)\n n_vid = 1\n if n_vid > 0 and len(unique_vids) > 0:\n rng_vid = np.random.RandomState(random_state)\n chosen_vids = rng_vid.choice(\n unique_vids, size=min(n_vid, len(unique_vids)), replace=False\n )\n for v in chosen_vids:\n mask_vid = (vid_interest == v)\n if mask_vid.sum() == 0:\n continue\n\n print(\n f\"\\n[SHAP ensemble] Heatmap for video_id={v} \"\n f\"(n_samples={mask_vid.sum()})\"\n )\n sv = shap_ensemble[mask_vid]\n Xv = X_shap_ensemble.iloc[mask_vid]\n\n try:\n exp_v = shap.Explanation(\n values=sv,\n base_values=np.zeros(sv.shape[0], dtype=float),\n data=Xv.values,\n feature_names=feat_names,\n )\n shap.plots.heatmap(exp_v, max_display=20)\n except Exception:\n shap.plots.heatmap(\n sv,\n feature_names=feat_names,\n max_display=20,\n )\n except Exception as e_inner:\n print(f\"[WARN] Could not create SHAP heatmap plots: {e_inner}\")\n if debug:\n traceback.print_exc()\n\n except Exception as e:\n print(f\"[WARN] Could not generate some SHAP ensemble plots: {e}\")\n if debug:\n traceback.print_exc()\n\n return {\n \"f1_ensemble\": f1,\n \"prec_ensemble\": prec,\n \"rec_ensemble\": rec,\n \"cm_ensemble\": cm,\n \"per_video\": per_vid,\n \"oof_ensemble\": probs_ens,\n \"y\": y_action,\n \"X_action\": X_raw,\n \"shap_values_ensemble\": shap_ensemble,\n \"shap_index\": idx_interest,\n \"feature_names\": feat_names,\n }\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if CFG.mode == \"shap_analysis\":\n all_diag = {} # (lab_id, switch, action) -> diag object\n\n shap_root = os.path.join(str(CFG.model_path), \"shap\")\n \n for switch in (\"single\", \"pair\"):\n for lab_id in labs_dict.keys():\n shap_dir = os.path.join(shap_root, str(lab_id), str(switch))\n if not os.path.isdir(shap_dir):\n # no SHAP data for this lab/switch\n continue\n\n # infer actions from saved *_holdout_data.pkl files\n shap_paths = sorted(\n glob.glob(os.path.join(shap_dir, \"*_holdout_data.pkl\"))\n )\n if not shap_paths:\n continue\n\n for path in shap_paths:\n # filename is \"_holdout_data.pkl\"\n fname = os.path.basename(path)\n if not fname.endswith(\"_holdout_data.pkl\"):\n continue\n action = fname.replace(\"_holdout_data.pkl\", \"\")\n\n key = (lab_id, switch, action)\n print(f\"[shap_analysis] lab={lab_id}, switch={switch}, action={action}\")\n\n try:\n diag = diagnose_action_failure(\n lab_id=lab_id,\n action=action,\n switch=switch,\n max_shap_samples_per_group=300,\n random_state=42,\n print_top_k=15,\n debug=False,\n )\n all_diag[key] = diag\n except Exception as e:\n print(\n f\"[shap_analysis] FAILED for lab={lab_id}, switch={switch}, \"\n f\"action={action}: {e}\"\n )\n continue\n\n # optional: keep all_diag for later inspection\n # e.g. joblib.dump(all_diag, os.path.join(shap_root, \"all_diag.pkl\"), compress=3)\n","metadata":{},"outputs":[],"execution_count":null}]}