{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":102335,"databundleVersionId":12518947,"sourceType":"competition"},{"sourceId":12373710,"sourceType":"datasetVersion","datasetId":7801972},{"sourceId":12418209,"sourceType":"datasetVersion","datasetId":7799425}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport sys\nimport pandas as pd\nimport numpy as np\nsys.path.append('/kaggle/input/cmi-modules-tom')\nfrom feature_engineering import *\nfrom scipy.spatial.transform import Rotation as R, Slerp\nimport pyro\nfrom pyro.distributions import Normal, Uniform\nimport pyro.distributions as dist\nfrom pyro.infer import EmpiricalMarginal, Importance, MCMC, NUTS\nimport argparse\nfrom configs_IMUonly_exp1 import *\nimport torch\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nimport seaborn as sns\nfrom mpl_toolkits.mplot3d.art3d import Poly3DCollection\nfrom scipy.linalg import eigh","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:15:53.165711Z","iopub.execute_input":"2025-07-09T16:15:53.166362Z","iopub.status.idle":"2025-07-09T16:15:53.172171Z","shell.execute_reply.started":"2025-07-09T16:15:53.166338Z","shell.execute_reply":"2025-07-09T16:15:53.171462Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## MCMC CMI Data","metadata":{}},{"cell_type":"markdown","source":"* acc_measured = acc_true + noise\n\n* vel[t+1] = vel[t] + acc_true[t] * dt\n\n* pos[t+1] = pos[t] + vel[t] * dt","metadata":{}},{"cell_type":"code","source":"g = 9.81  # gravity\ndt = 0.005\nnum_samples = 300","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:20:12.169124Z","iopub.execute_input":"2025-07-09T16:20:12.169419Z","iopub.status.idle":"2025-07-09T16:20:12.173153Z","shell.execute_reply.started":"2025-07-09T16:20:12.169399Z","shell.execute_reply":"2025-07-09T16:20:12.172333Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def quaternion_to_rotation_matrix(quat):\n    \"\"\"Convert (T, 4) quaternion to rotation matrix (T, 3, 3).\"\"\"\n    rotations = R.from_quat(quat).as_matrix()\n    return torch.tensor(rotations, dtype=torch.float32)\n\ndef rotate_acceleration(acc_local, quat):\n    \"\"\"Rotate local acceleration to global frame using quaternions.\"\"\"\n    R_t = quaternion_to_rotation_matrix(quat)  # (T, 3, 3)\n    acc_global = torch.einsum(\"tij,tj->ti\", R_t, acc_local)\n    return acc_global\n\ndef kinematic_integration(acc_global, dt):\n    \"\"\"Integrate acceleration to get velocity and position.\"\"\"\n    v = torch.cumsum(acc_global * dt, dim=0)\n    p = torch.cumsum(v * dt, dim=0)\n    return v, p\n\n# Pyro model\ndef imu_model(acc_obs, quat_obs, angvel_obs, dt=0.005):\n    T = acc_obs.shape[0]\n\n    # Priors\n    bias = pyro.sample(\"acc_bias\", dist.Normal(torch.zeros(3), 0.1 * torch.ones(3)))\n    noise_scale = pyro.sample(\"noise_scale\", dist.HalfNormal(0.1))\n\n    # Rotate acceleration to global frame\n    acc_corrected = acc_obs - bias\n    acc_global = rotate_acceleration(acc_corrected, quat_obs)\n\n    # Gravity correction\n    acc_global = acc_global - torch.tensor([0, 0, g])\n    v, p = kinematic_integration(acc_global, dt)\n\n    # Observation likelihood (we observe positions here as proxy)\n    for t in range(T):\n        pyro.sample(f\"obs_{t}\", dist.Normal(p[t], noise_scale).to_event(1), obs=None)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:19:29.085039Z","iopub.execute_input":"2025-07-09T16:19:29.085334Z","iopub.status.idle":"2025-07-09T16:19:29.091976Z","shell.execute_reply.started":"2025-07-09T16:19:29.085313Z","shell.execute_reply":"2025-07-09T16:19:29.091249Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df = pd.read_csv('/kaggle/input/cmi-detect-behavior-with-sensor-data/train.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:19:48.932499Z","iopub.execute_input":"2025-07-09T16:19:48.932774Z","iopub.status.idle":"2025-07-09T16:20:10.252584Z","shell.execute_reply.started":"2025-07-09T16:19:48.932754Z","shell.execute_reply":"2025-07-09T16:20:10.251919Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"seq_id = df.sequence_id.unique()[5]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:20:19.399608Z","iopub.execute_input":"2025-07-09T16:20:19.400334Z","iopub.status.idle":"2025-07-09T16:20:19.429225Z","shell.execute_reply.started":"2025-07-09T16:20:19.400311Z","shell.execute_reply":"2025-07-09T16:20:19.428241Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df[df.sequence_id == seq_id].shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:20:19.568703Z","iopub.execute_input":"2025-07-09T16:20:19.568958Z","iopub.status.idle":"2025-07-09T16:20:19.613124Z","shell.execute_reply.started":"2025-07-09T16:20:19.568937Z","shell.execute_reply":"2025-07-09T16:20:19.612556Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_seq = df[df.sequence_id == seq_id].copy().groupby('sequence_id', group_keys=False).apply(lambda x: compute_combined_features(x, CFG))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:20:20.983881Z","iopub.execute_input":"2025-07-09T16:20:20.984510Z","iopub.status.idle":"2025-07-09T16:20:21.105719Z","shell.execute_reply.started":"2025-07-09T16:20:20.984484Z","shell.execute_reply":"2025-07-09T16:20:21.104751Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"acc_obs = torch.tensor(df_seq[['acc_x', 'acc_y', 'acc_z']].values, dtype=torch.float32)\nquat_obs = torch.tensor(df_seq[['rot_w', 'rot_x', 'rot_y', 'rot_z']].values)\nangvel_obs = torch.tensor(df_seq[['angular_vel_x', 'angular_vel_y', 'angular_vel_z']].values)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:20:23.654004Z","iopub.execute_input":"2025-07-09T16:20:23.654611Z","iopub.status.idle":"2025-07-09T16:20:23.661127Z","shell.execute_reply.started":"2025-07-09T16:20:23.654586Z","shell.execute_reply":"2025-07-09T16:20:23.660377Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.manual_seed(42)\nnuts_kernel = NUTS(imu_model)\nmcmc = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=100)\nmcmc.run(acc_obs, quat_obs, angvel_obs, dt)\nposterior_samples = mcmc.get_samples()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:22:02.362048Z","iopub.execute_input":"2025-07-09T16:22:02.362334Z","iopub.status.idle":"2025-07-09T16:22:20.054997Z","shell.execute_reply.started":"2025-07-09T16:22:02.362313Z","shell.execute_reply":"2025-07-09T16:22:20.054261Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"num_samples = posterior_samples[\"obs_0\"].shape[0]\nT = len([k for k in posterior_samples if k.startswith(\"obs_\")])\npositions = torch.stack(\n    [posterior_samples[f\"obs_{t}\"] for t in range(T)], dim=1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:22:35.159476Z","iopub.execute_input":"2025-07-09T16:22:35.160014Z","iopub.status.idle":"2025-07-09T16:22:35.164512Z","shell.execute_reply.started":"2025-07-09T16:22:35.159988Z","shell.execute_reply":"2025-07-09T16:22:35.163796Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"positions = torch.stack([posterior_samples[f\"obs_{t}\"] for t in range(T)], dim=1)\npositions_np = positions.numpy()\nmean_pos = positions_np.mean(axis=0)  # shape (T, 3)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:22:36.504329Z","iopub.execute_input":"2025-07-09T16:22:36.505103Z","iopub.status.idle":"2025-07-09T16:22:36.509129Z","shell.execute_reply.started":"2025-07-09T16:22:36.505077Z","shell.execute_reply":"2025-07-09T16:22:36.508568Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"num_samples, T, _ = positions_np.shape\n\n# Prepare long-form DataFrame for ridge plots\ndef make_df_for_dim(dim, name):\n    return pd.DataFrame({\n        \"position\": positions_np[:, :, dim].flatten(),\n        \"time\": np.repeat(np.arange(T), num_samples),\n        \"axis\": name\n    })\n\ndf_x = make_df_for_dim(0, \"x\")\ndf_y = make_df_for_dim(1, \"y\")\ndf_z = make_df_for_dim(2, \"z\")\ndf_all = pd.concat([df_x, df_y, df_z], ignore_index=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:22:36.814985Z","iopub.execute_input":"2025-07-09T16:22:36.815309Z","iopub.status.idle":"2025-07-09T16:22:36.822736Z","shell.execute_reply.started":"2025-07-09T16:22:36.815288Z","shell.execute_reply":"2025-07-09T16:22:36.822028Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig = plt.figure(figsize=(14, 10))\n\n# --- Main 3D trajectory plot ---\nax3d = fig.add_subplot(2, 2, 1, projection='3d')\nmean_pos = positions.mean(0).numpy()\n\nax3d.plot(mean_pos[:, 0], mean_pos[:, 1], mean_pos[:, 2], color='blue', lw=2)\nax3d.set_title(\"Mean 3D Trajectory\")\nax3d.set_xlabel(\"X\")\nax3d.set_ylabel(\"Y\")\nax3d.set_zlabel(\"Z\")\n\n# --- Ridge plot helper ---\ndef plot_ridges(axis_name, subplot_idx):\n    df_axis = df_all[df_all[\"axis\"] == axis_name]\n    ax = fig.add_subplot(2, 2, subplot_idx)\n    sns.set(style=\"white\")\n    g = sns.FacetGrid(df_axis, row=\"time\", aspect=15, height=0.25)\n    g.map(sns.kdeplot, \"position\", fill=True, alpha=1)\n    g.map(plt.axhline, y=0, lw=2, clip_on=False)\n    g.set_titles(\"\")\n    g.set(yticks=[], ylabel=\"\", xlabel=f\"{axis_name.upper()} Position\")\n    g.despine(bottom=True, left=True)\n    plt.close(g.fig)  # Avoid duplicated output\n    ax.imshow(g.fig.canvas.buffer_rgba())\n    ax.axis(\"off\")\n\n# NOTE: FacetGrid + subplot combination is tricky in one figure.\n# Better alternative:\n# --- Violin plots for x, y, z per time ---\nfor i, axis_name in enumerate([\"x\", \"y\", \"z\"]):\n    ax = fig.add_subplot(2, 2, i + 2)\n    df_axis = df_all[df_all[\"axis\"] == axis_name]\n    sns.violinplot(x=\"time\", y=\"position\", data=df_axis, ax=ax, scale=\"width\", inner=\"quartile\")\n    ax.set_title(f\"Posterior {axis_name.upper()} Distribution\")\n    ax.set_xlabel(\"Time Step\")\n    ax.set_ylabel(f\"{axis_name.upper()} Position\")\n    ax.set_xticks([])\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:22:38.868277Z","iopub.execute_input":"2025-07-09T16:22:38.868550Z","iopub.status.idle":"2025-07-09T16:22:42.849203Z","shell.execute_reply.started":"2025-07-09T16:22:38.868530Z","shell.execute_reply":"2025-07-09T16:22:42.848468Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_cov_ellipsoid(ax, mean, cov, n_std=2.0, resolution=20, alpha=0.2, color='red'):\n    \"\"\"Plot a 3D ellipsoid representing covariance.\"\"\"\n    # Eigenvalues (radii) and eigenvectors (rotation)\n    eigvals, eigvecs = eigh(cov)\n    radii = n_std * np.sqrt(eigvals)\n    \n    # Create a unit sphere\n    u = np.linspace(0, 2 * np.pi, resolution)\n    v = np.linspace(0, np.pi, resolution)\n    x = np.outer(np.cos(u), np.sin(v))\n    y = np.outer(np.sin(u), np.sin(v))\n    z = np.outer(np.ones_like(u), np.cos(v))\n    sphere = np.stack((x, y, z), axis=-1)  # shape (res, res, 3)\n    \n    # Transform sphere to ellipsoid\n    ellipsoid = sphere @ np.diag(radii) @ eigvecs.T + mean\n    \n    # Convert to Poly3D and add to plot\n    ax.plot_surface(\n        ellipsoid[:, :, 0], ellipsoid[:, :, 1], ellipsoid[:, :, 2],\n        rstride=1, cstride=1, color=color, alpha=alpha, linewidth=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:22:50.591770Z","iopub.execute_input":"2025-07-09T16:22:50.592268Z","iopub.status.idle":"2025-07-09T16:22:50.597919Z","shell.execute_reply.started":"2025-07-09T16:22:50.592243Z","shell.execute_reply":"2025-07-09T16:22:50.597222Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig = plt.figure(figsize=(12, 8))\nax = fig.add_subplot(111, projection='3d')\n\n# Plot mean trajectory\nax.plot(mean_pos[:, 0], mean_pos[:, 1], mean_pos[:, 2], lw=2, color='blue', label=\"Mean Trajectory\")\n\n# Add ellipsoids at selected time steps\nskip = max(1, T // 30)\nfor t in range(0, T, skip):\n    cloud = positions_np[:, t, :]  # shape (N, 3)\n    mean_t = cloud.mean(axis=0)\n    cov_t = np.cov(cloud.T)\n\n    plot_cov_ellipsoid(ax, mean_t, cov_t, n_std=2.0, alpha=0.01, color='red')\n\nax.set_xlabel(\"X\")\nax.set_ylabel(\"Y\")\nax.set_zlabel(\"Z\")\nax.set_title(\"3D Trajectory with Posterior Covariance Ellipsoids\")\nax.legend()\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-09T16:22:52.370014Z","iopub.execute_input":"2025-07-09T16:22:52.370326Z","iopub.status.idle":"2025-07-09T16:22:52.871695Z","shell.execute_reply.started":"2025-07-09T16:22:52.370303Z","shell.execute_reply":"2025-07-09T16:22:52.870986Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}