{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "725806fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 4 files from '/Users/mpitell/Downloads/Drag Tracking Output' | 3 particle(s) detected:\n",
      "  0.1 mm/s: 154 frames\n",
      "  0.15 mm/s: 135 frames\n",
      "  0.2 mm/s: 133 frames\n",
      "  0.25 mm/s: 140 frames\n",
      "   frame     speed       p0_cx       p1_cx      p2_cx\n",
      "0      0  0.1 mm/s  137.377744  112.645468  85.141857\n",
      "1      1  0.1 mm/s  137.276239  112.657680  85.057428\n",
      "2      2  0.1 mm/s  137.164824  112.674773  85.139626\n",
      "3      3  0.1 mm/s  137.595084  112.609946  84.987524\n",
      "4      4  0.1 mm/s  137.128077  112.740710  84.912927\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import glob\n",
    "import os\n",
    "\n",
    "folder = '/Users/mpitell/Downloads/Drag Tracking Output'\n",
    "dfs = []\n",
    "\n",
    "# Column names for up to 3 particles (9 columns)\n",
    "ALL_COL_NAMES = [\n",
    "    \"p0_cx\", \"p0_cy\", \"p0_r\",\n",
    "    \"p1_cx\", \"p1_cy\", \"p1_r\",\n",
    "    \"p2_cx\", \"p2_cy\", \"p2_r\",\n",
    "]\n",
    "\n",
    "for fpath in sorted(glob.glob(os.path.join(folder, \"*.csv\"))):\n",
    "    fname = os.path.splitext(os.path.basename(fpath))[0]  # e.g. \"0.4 mms\"\n",
    "    speed = fname.replace(\"mms\", \"mm/s\")\n",
    "\n",
    "    # Detect number of columns in file\n",
    "    ncols = len(pd.read_csv(fpath, header=None, nrows=1).columns)\n",
    "    n_particles = ncols // 3\n",
    "    col_names = ALL_COL_NAMES[:ncols]\n",
    "\n",
    "    df = pd.read_csv(fpath, header=None, usecols=list(range(ncols)), names=col_names)\n",
    "    df[\"frame\"] = df.index\n",
    "    df[\"speed\"] = speed\n",
    "    df[\"n_particles\"] = n_particles\n",
    "    dfs.append(df)\n",
    "\n",
    "data = pd.concat(dfs, ignore_index=True)\n",
    "\n",
    "# Global particle count (max across all files)\n",
    "N_PARTICLES = int(data[\"n_particles\"].max())\n",
    "PARTICLE_COLS = [f\"p{i}_cx\" for i in range(N_PARTICLES)]\n",
    "\n",
    "print(f\"Loaded {len(dfs)} files from '{folder}' | {N_PARTICLES} particle(s) detected:\")\n",
    "for d in dfs:\n",
    "    print(f\"  {d['speed'].iloc[0]}: {len(d)} frames\")\n",
    "print(data[[\"frame\", \"speed\"] + PARTICLE_COLS].head())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b9ef82",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "import matplotlib.pyplot as plt\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "\n",
    "TRIM = 15\n",
    "\n",
    "phase_colors = {\n",
    "    \"Stationary\": \"#0072B2\",\n",
    "    \"Movement1\":  \"#D55E00\",\n",
    "    \"Movement2\":  \"#009E73\",\n",
    "}\n",
    "\n",
    "# One distinct color per particle for the signal traces\n",
    "particle_trace_colors = [\"#333333\", \"#CC79A7\", \"#F0E442\"]\n",
    "\n",
    "all_ranges   = {}\n",
    "speed_groups = {s: g.reset_index(drop=True) for s, g in data.groupby(\"speed\")}\n",
    "speeds_list  = sorted(speed_groups.keys())\n",
    "\n",
    "confirm_all_btn = widgets.Button(description=\"Confirm All & Continue\",\n",
    "                                  button_style=\"success\",\n",
    "                                  layout=widgets.Layout(width=\"220px\"))\n",
    "global_status = widgets.Label(\"\")\n",
    "display(widgets.HBox([confirm_all_btn, global_status]))\n",
    "\n",
    "slider_refs = {}\n",
    "\n",
    "for speed in speeds_list:\n",
    "    group  = speed_groups[speed]\n",
    "    frames = group[\"frame\"].to_numpy()\n",
    "    n      = len(frames)\n",
    "\n",
    "    s1 = int(n * 0.00); e1 = int(n * 0.20)\n",
    "    s2 = int(n * 0.20); e2 = int(n * 0.50)\n",
    "    s3 = int(n * 0.50); e3 = int(n * 0.80)\n",
    "\n",
    "    sliders = {\n",
    "        \"Stationary\": widgets.IntRangeSlider(value=[s1, e1], min=0, max=n-1, step=1,\n",
    "                          description=\"Stationary:\", continuous_update=True,\n",
    "                          layout=widgets.Layout(width=\"600px\"),\n",
    "                          style={\"description_width\": \"100px\"}),\n",
    "        \"Movement1\":  widgets.IntRangeSlider(value=[s2, e2], min=0, max=n-1, step=1,\n",
    "                          description=\"Movement 1:\", continuous_update=True,\n",
    "                          layout=widgets.Layout(width=\"600px\"),\n",
    "                          style={\"description_width\": \"100px\"}),\n",
    "        \"Movement2\":  widgets.IntRangeSlider(value=[s3, e3], min=0, max=n-1, step=1,\n",
    "                          description=\"Movement 2:\", continuous_update=True,\n",
    "                          layout=widgets.Layout(width=\"600px\"),\n",
    "                          style={\"description_width\": \"100px\"}),\n",
    "    }\n",
    "    slider_refs[speed] = sliders\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(5, 3))\n",
    "\n",
    "    # Plot each particle's x signal\n",
    "    n_p = int(group[\"n_particles\"].iloc[0])\n",
    "    for i in range(n_p):\n",
    "        col = f\"p{i}_cx\"\n",
    "        if col in group.columns:\n",
    "            ax.plot(frames, group[col].to_numpy(),\n",
    "                    color=particle_trace_colors[i % len(particle_trace_colors)],\n",
    "                    lw=0.8, label=f\"Particle {i}\")\n",
    "\n",
    "    ax.set_xlabel(\"Frame\"); ax.set_ylabel(\"X (px)\")\n",
    "    ax.set_title(f\"Speed: {speed}\")\n",
    "\n",
    "    # Legend proxy patches for phases (static)\n",
    "    for phase, color in phase_colors.items():\n",
    "        ax.axvspan(0, 0, alpha=0.4, color=color, label=phase)\n",
    "    ax.legend(fontsize=7, loc=\"upper right\")\n",
    "    fig.tight_layout()\n",
    "\n",
    "    active_spans = []\n",
    "\n",
    "    def redraw(change, ax=ax, frames=frames, sliders=sliders,\n",
    "               active_spans=active_spans, fig=fig):\n",
    "        for p in active_spans:\n",
    "            p.remove()\n",
    "        active_spans.clear()\n",
    "\n",
    "        for phase, sl in sliders.items():\n",
    "            color = phase_colors[phase]\n",
    "            lo, hi = sl.value\n",
    "            active_spans.append(\n",
    "                ax.axvspan(frames[lo], frames[hi], alpha=0.15, color=color))\n",
    "            t_lo = min(lo + TRIM, hi)\n",
    "            t_hi = max(hi - TRIM, t_lo)\n",
    "            active_spans.append(\n",
    "                ax.axvspan(frames[t_lo], frames[t_hi], alpha=0.40, color=color))\n",
    "\n",
    "        fig.canvas.draw_idle()\n",
    "\n",
    "    for sl in sliders.values():\n",
    "        sl.observe(redraw, names=\"value\")\n",
    "\n",
    "    plt.show()\n",
    "    display(widgets.VBox(list(sliders.values())))\n",
    "    redraw(None)\n",
    "\n",
    "def on_confirm_all(_):\n",
    "    for speed, sliders in slider_refs.items():\n",
    "        all_ranges[speed] = {\n",
    "            phase: tuple(sl.value) for phase, sl in sliders.items()\n",
    "        }\n",
    "    summary = \"\\n\".join(\n",
    "        f\"  {s}: \" + \", \".join(f\"{ph}={v}\" for ph, v in r.items())\n",
    "        for s, r in all_ranges.items()\n",
    "    )\n",
    "    global_status.value = \"✓ Ranges saved — run the analysis cell.\"\n",
    "    print(\"Confirmed ranges:\\n\" + summary)\n",
    "\n",
    "confirm_all_btn.on_click(on_confirm_all)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "0bf72919",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Particle radius: 1.000 +/- 0.0000 um\n",
      "    speed  particle     phase  F_drag_pN  F_se_pN    dx_um  dx_se_um\n",
      " 0.1 mm/s         0 Movement1   1.884956      0.0 0.081575  0.002619\n",
      " 0.1 mm/s         0 Movement2   1.884956      0.0 0.117347  0.003291\n",
      " 0.1 mm/s         1 Movement1   1.884956      0.0 0.068395  0.002828\n",
      " 0.1 mm/s         1 Movement2   1.884956      0.0 0.106229  0.002717\n",
      " 0.1 mm/s         2 Movement1   1.884956      0.0 0.052097  0.002083\n",
      " 0.1 mm/s         2 Movement2   1.884956      0.0 0.084052  0.002118\n",
      "0.15 mm/s         0 Movement1   2.827433      0.0 0.133638  0.004520\n",
      "0.15 mm/s         0 Movement2   2.827433      0.0 0.175811  0.003540\n",
      "0.15 mm/s         1 Movement1   2.827433      0.0 0.114427  0.004062\n",
      "0.15 mm/s         1 Movement2   2.827433      0.0 0.144567  0.002992\n",
      "0.15 mm/s         2 Movement1   2.827433      0.0 0.085323  0.004092\n",
      "0.15 mm/s         2 Movement2   2.827433      0.0 0.135507  0.003782\n",
      " 0.2 mm/s         0 Movement1   3.769911      0.0 0.179852  0.003646\n",
      " 0.2 mm/s         0 Movement2   3.769911      0.0 0.218980  0.002660\n",
      " 0.2 mm/s         1 Movement1   3.769911      0.0 0.156521  0.003255\n",
      " 0.2 mm/s         1 Movement2   3.769911      0.0 0.191989  0.003276\n",
      " 0.2 mm/s         2 Movement1   3.769911      0.0 0.127746  0.002778\n",
      " 0.2 mm/s         2 Movement2   3.769911      0.0 0.167028  0.003381\n",
      "0.25 mm/s         0 Movement1   4.712389      0.0 0.222488  0.006108\n",
      "0.25 mm/s         0 Movement2   4.712389      0.0 0.252513  0.006417\n",
      "0.25 mm/s         1 Movement1   4.712389      0.0 0.201752  0.006412\n",
      "0.25 mm/s         1 Movement2   4.712389      0.0 0.220653  0.007192\n",
      "0.25 mm/s         2 Movement1   4.712389      0.0 0.163626  0.005710\n",
      "0.25 mm/s         2 Movement2   4.712389      0.0 0.187163  0.006309\n",
      "Particle 0 Movement1: k = 21.2304 +/- 0.2854 pN/um  (95% CI: +/- 0.9084)\n",
      "Particle 0 Movement2: k = 17.4799 +/- 0.6202 pN/um  (95% CI: +/- 1.9737)\n",
      "Particle 1 Movement1: k = 24.0227 +/- 0.5768 pN/um  (95% CI: +/- 1.8356)\n",
      "Particle 1 Movement2: k = 20.1524 +/- 0.6587 pN/um  (95% CI: +/- 2.0963)\n",
      "Particle 2 Movement1: k = 29.9907 +/- 1.1696 pN/um  (95% CI: +/- 3.7222)\n",
      "Particle 2 Movement2: k = 23.2384 +/- 0.9756 pN/um  (95% CI: +/- 3.1047)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "98d1c19f5430431cb03a1bb253960aa2",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 127288 characters]",
      "text/html": [
       "\n",
       "            <div style=\"display: inline-block;\">\n",
       "                <div class=\"jupyter-widgets widget-label\" style=\"text-align: center;\">\n",
       "                    Figure\n",
       "                </div>\n",
       "                <img src='data:image/png;base64,[encoded data, 127288 characters]' width=700.0/>\n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from scipy import stats\n",
    "\n",
    "PIXELS_PER_MICRON = 18.0\n",
    "ETA = 0.001  # Pa·s, viscosity of water\n",
    "# TRIM is defined in the slider cell above\n",
    "\n",
    "# Particle radius (m) — used for Stokes drag force\n",
    "R    = 1e-6   # 1 µm\n",
    "R_se = 0\n",
    "print(f\"Particle radius: {R*1e6:.3f} +/- {R_se*1e6:.4f} um\")\n",
    "\n",
    "def trimmed_vals(series, start, end, trim=TRIM):\n",
    "    s = start + trim\n",
    "    e = end - trim\n",
    "    if s >= e:\n",
    "        s, e = start, end\n",
    "    return series.iloc[s:e + 1]\n",
    "\n",
    "results = []\n",
    "\n",
    "for speed_str, group in data.groupby(\"speed\"):\n",
    "    ranges    = all_ranges[speed_str]\n",
    "    v         = float(speed_str.split()[0]) * 1e-3  # mm/s -> m/s\n",
    "    F_drag_pN = 6 * np.pi * ETA * R * v * 1e12\n",
    "    F_se_pN   = F_drag_pN * (R_se / R) if R > 0 else 0\n",
    "\n",
    "    n_p = int(group[\"n_particles\"].iloc[0])\n",
    "    s_start, s_end = ranges[\"Stationary\"]\n",
    "\n",
    "    for i in range(n_p):\n",
    "        col = f\"p{i}_cx\"\n",
    "        if col not in group.columns:\n",
    "            continue\n",
    "\n",
    "        x_stat_vals = trimmed_vals(group[col], s_start, s_end)\n",
    "        x_stat      = x_stat_vals.mean()\n",
    "        x_stat_se   = x_stat_vals.std() / np.sqrt(len(x_stat_vals))\n",
    "\n",
    "        for key in [\"Movement1\", \"Movement2\"]:\n",
    "            if key not in ranges:\n",
    "                continue\n",
    "            m_start, m_end = ranges[key]\n",
    "            vals      = trimmed_vals(group[col], m_start, m_end)\n",
    "            x_move    = vals.mean()\n",
    "            x_move_se = vals.std() / np.sqrt(len(vals))\n",
    "\n",
    "            dx_um    = abs(x_move - x_stat) / PIXELS_PER_MICRON\n",
    "            dx_se_um = np.sqrt(x_stat_se**2 + x_move_se**2) / PIXELS_PER_MICRON\n",
    "\n",
    "            results.append({\n",
    "                \"speed\":      speed_str,\n",
    "                \"particle\":   i,\n",
    "                \"phase\":      key,\n",
    "                \"F_drag_pN\":  F_drag_pN,\n",
    "                \"F_se_pN\":    F_se_pN,\n",
    "                \"dx_um\":      dx_um,\n",
    "                \"dx_se_um\":   dx_se_um,\n",
    "            })\n",
    "\n",
    "results_df = pd.DataFrame(results)\n",
    "print(results_df.to_string(index=False))\n",
    "\n",
    "# ---- Color = particle, filled = Movement1, open = Movement2 ----\n",
    "particle_colors = [\"#0072B2\", \"#D55E00\", \"#009E73\", \"#CC79A7\", \"#E69F00\"]\n",
    "phase_style = {\n",
    "    \"Movement1\": dict(filled=True,  linestyle=\"-\",  label_suffix=\"Dir 1\"),\n",
    "    \"Movement2\": dict(filled=False, linestyle=\":\",  label_suffix=\"Dir 2\"),\n",
    "}\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 5))\n",
    "\n",
    "for i in range(N_PARTICLES):\n",
    "    color   = particle_colors[i % len(particle_colors)]\n",
    "    df_part = results_df[results_df[\"particle\"] == i]\n",
    "    if df_part.empty:\n",
    "        continue\n",
    "\n",
    "    for phase, pstyle in phase_style.items():\n",
    "        df_ph = df_part[df_part[\"phase\"] == phase]\n",
    "        if df_ph.empty:\n",
    "            continue\n",
    "\n",
    "        mfc = color if pstyle[\"filled\"] else \"none\"\n",
    "\n",
    "        for _, row in df_ph.iterrows():\n",
    "            ax.errorbar(\n",
    "                row[\"dx_um\"], row[\"F_drag_pN\"],\n",
    "                xerr=row[\"dx_se_um\"], yerr=row[\"F_se_pN\"],\n",
    "                fmt=\"o\",\n",
    "                markersize=6,\n",
    "                markerfacecolor=mfc,\n",
    "                markeredgecolor=color,\n",
    "                ecolor=color,\n",
    "                capsize=3,\n",
    "                linestyle=\"none\",\n",
    "                zorder=3,\n",
    "                label=\"_nolegend_\",\n",
    "            )\n",
    "\n",
    "        # ---- origin-constrained fit per particle per direction ----\n",
    "        x_p = df_ph[\"dx_um\"].to_numpy(dtype=float)\n",
    "        y_p = df_ph[\"F_drag_pN\"].to_numpy(dtype=float)\n",
    "        if len(x_p) < 2:\n",
    "            continue\n",
    "\n",
    "        k       = np.sum(x_p * y_p) / np.sum(x_p**2)\n",
    "        resid   = y_p - k * x_p\n",
    "        sigma_k = np.sqrt(np.sum(resid**2) / max((len(x_p) - 1) * np.sum(x_p**2), 1e-30))\n",
    "\n",
    "        t_val = stats.t.ppf(0.975, df=max(len(x_p) - 1, 1))\n",
    "        ci_k  = t_val * sigma_k\n",
    "        x_fit = np.linspace(0, x_p.max() * 1.05, 300)\n",
    "\n",
    "        ax.plot(\n",
    "            x_fit, k * x_fit,\n",
    "            linestyle=pstyle[\"linestyle\"],\n",
    "            color=color,\n",
    "            linewidth=1.8,\n",
    "            zorder=2,\n",
    "            label=fr\"P{i} {pstyle['label_suffix']}: $k={k:.3g}\\pm{sigma_k:.2g}$ pN/$\\mu$m\",\n",
    "        )\n",
    "        ax.fill_between(x_fit, (k - ci_k) * x_fit, (k + ci_k) * x_fit,\n",
    "                        color=color, alpha=0.10, zorder=1)\n",
    "\n",
    "        print(f\"Particle {i} {phase}: k = {k:.4f} +/- {sigma_k:.4f} pN/um  (95% CI: +/- {ci_k:.4f})\")\n",
    "\n",
    "ax.set_xlabel(r\"$\\Delta X$ ($\\mu$m)\")\n",
    "ax.set_ylabel(\"Drag Force (pN)\")\n",
    "ax.set_title(\"Drag Force vs Displacement\")\n",
    "ax.legend(fontsize=7, frameon=False)\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
