improve data import with method chaining

Author: drbenvincent

docs/notebooks/geolift1.ipynb

@@ -181,11 +181,12 @@
     }
    ],
    "source": [
-    "df = cp.load_data(\"geolift1\")\n",
-    "# convert time column to datetime format\n",
-    "df[\"time\"] = pd.to_datetime(df[\"time\"])\n",
-    "df.set_index(\"time\", inplace=True)\n",
-    "# define the treamtment time\n",
+    "df = (\n",
+    "    cp.load_data(\"geolift1\")\n",
+    "    .assign(time=lambda x: pd.to_datetime(x[\"time\"]))\n",
+    "    .set_index(\"time\")\n",
+    ")\n",
+    "\n",
     "treatment_time = pd.to_datetime(\"2022-01-01\")\n",
     "df.head()"
    ]
@@ -589,7 +590,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.10.8 ('CausalPy')",
+   "display_name": "CausalPy",
    "language": "python",
    "name": "python3"
   },
@@ -603,12 +604,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.8 | packaged by conda-forge | (main, Nov 22 2022, 08:25:29) [Clang 14.0.6 ]"
+   "version": "3.10.6 | packaged by conda-forge | (main, Aug 22 2022, 20:38:29) [Clang 13.0.1 ]"
   },
   "orig_nbformat": 4,
   "vscode": {
    "interpreter": {
-    "hash": "46d31859cc45aa26a1223a391e7cf3023d69984b498bed11e66c690302b7e251"
+    "hash": "02f5385db19eab57520277c5168790c7855381ee953bdbb5c89c321e1f17586e"
    }
   }
  },

docs/notebooks/its_covid.ipynb

@@ -142,9 +142,12 @@
     }
    ],
    "source": [
-    "df = cp.load_data(\"its\")\n",
-    "df[\"date\"] = pd.to_datetime(df[\"date\"])\n",
-    "df.set_index(\"date\", inplace=True)\n",
+    "df = (\n",
+    "    cp.load_data(\"its\")\n",
+    "    .assign(date=lambda x: pd.to_datetime(x[\"date\"]))\n",
+    "    .set_index(\"date\")\n",
+    ")\n",
+    "\n",
     "treatment_time = pd.to_datetime(\"2017-01-01\")\n",
     "df.head()"
    ]
@@ -228,7 +231,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 2 seconds.\n",
+      "Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 3 seconds.\n",
       "Sampling: [beta, sigma, y_hat]\n",
       "Sampling: [y_hat]\n",
       "Sampling: [y_hat]\n",
@@ -498,7 +501,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.10.8 ('CausalPy')",
+   "display_name": "CausalPy",
    "language": "python",
    "name": "python3"
   },
@@ -512,12 +515,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.8"
+   "version": "3.10.6"
   },
   "orig_nbformat": 4,
   "vscode": {
    "interpreter": {
-    "hash": "46d31859cc45aa26a1223a391e7cf3023d69984b498bed11e66c690302b7e251"
+    "hash": "02f5385db19eab57520277c5168790c7855381ee953bdbb5c89c321e1f17586e"
    }
   }
  },

docs/notebooks/sc2_pymc.ipynb

@@ -123,9 +123,12 @@
     }
    ],
    "source": [
-    "df = cp.load_data(\"its\")\n",
-    "df[\"date\"] = pd.to_datetime(df[\"date\"])\n",
-    "df.set_index(\"date\", inplace=True)\n",
+    "df = (\n",
+    "    cp.load_data(\"its\")\n",
+    "    .assign(date=lambda x: pd.to_datetime(x[\"date\"]))\n",
+    "    .set_index(\"date\")\n",
+    ")\n",
+    "\n",
     "treatment_time = pd.to_datetime(\"2017-01-01\")\n",
     "df.head()"
    ]
@@ -188,7 +191,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.10.6 ('CausalPy')",
+   "display_name": "CausalPy",
    "language": "python",
    "name": "python3"
   },
@@ -207,7 +210,7 @@
   "orig_nbformat": 4,
   "vscode": {
    "interpreter": {
-    "hash": "46d31859cc45aa26a1223a391e7cf3023d69984b498bed11e66c690302b7e251"
+    "hash": "02f5385db19eab57520277c5168790c7855381ee953bdbb5c89c321e1f17586e"
    }
   }
  },