add regime-pooling code

cjoshliu · Sep 16, 2023 · 467e728 · 467e728
1 parent cccb36a
commit 467e728
Showing 1 changed file with 40 additions and 1 deletion.
diff --git a/preprocessing/vid_to_dset.ipynb b/preprocessing/vid_to_dset.ipynb
@@ -23,7 +23,23 @@
     "# directory containing simulation and experiment video subdirectories\n",
     "SRC_DIR = '/path/to/saved/raw_data' # CHANGE THIS!!!\n",
     "# directory to save datasets as subdirectories\n",
-    "DST_DIR = '/path/to/save/datasets' # CHANGE THIS!!!"
+    "DST_DIR = '/path/to/save/datasets' # CHANGE THIS!!!\n",
+    "# datasets to pool in regime_pool\n",
+    "POOLS = {'pool_0': ['c1_-020_c2_050_dT_010_s_0',\n",
+    "                    'c1_-020_c2_070_dT_010_s_1',\n",
+    "                    'c1_-020_c2_090_dT_010_s_0',\n",
+    "                    'c1_-020_c2_110_dT_010_s_0',\n",
+    "                    'c1_-020_c2_130_dT_010_s_0'],\n",
+    "         'pool_1': ['c1_-020_c2_050_dT_010_s_4',\n",
+    "                    'c1_-020_c2_070_dT_010_s_5',\n",
+    "                    'c1_-020_c2_090_dT_010_s_6',\n",
+    "                    'c1_-020_c2_110_dT_010_s_4',\n",
+    "                    'c1_-020_c2_130_dT_010_s_12'],\n",
+    "         'pool_2': ['c1_-020_c2_050_dT_010_s_6',\n",
+    "                    'c1_-020_c2_070_dT_010_s_6',\n",
+    "                    'c1_-020_c2_090_dT_010_s_22',\n",
+    "                    'c1_-020_c2_110_dT_010_s_5',\n",
+    "                    'c1_-020_c2_130_dT_010_s_15']}"
    ]
   },
   {
@@ -189,6 +205,29 @@
     "                 os.path.join(DST_DIR, 'exp_orderp_vary'),\n",
     "                 ruler_path=ruler_path)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# make regime_pool datasets\n",
+    "src_datasets = os.path.join(DST_DIR, 'sim_regime_vary')\n",
+    "dst_datasets = os.path.join(DST_DIR, 'regime_pool')\n",
+    "\n",
+    "for key in POOLS.keys():\n",
+    "    print('Starting {}:'.format(key))\n",
+    "    dst_dataset = os.path.join(dst_datasets, key)\n",
+    "    Path(dst_dataset).mkdir(parents=True)\n",
+    "    for i in tqdm(range(len(POOLS[key]))):\n",
+    "        src_dataset = os.path.join(src_datasets, POOLS[key][i])\n",
+    "        for j in range(2000):\n",
+    "            src_fn = os.path.join(src_dataset, str(j).zfill(5)+'.tif')\n",
+    "            dst_fn = os.path.join(dst_dataset, str(j+i*2000).zfill(5)+'.tif')\n",
+    "            shutil.copyfile(src_fn, dst_fn)\n",
+    "    shutil.copyfile(ruler_path, os.path.join(dst_dataset, '10000.tif'))"
+   ]
   }
  ],
  "metadata": {