Multi-CPU processing and save disk space & downloading time

SeismoSocialDistancing.ipynb

@@ -53,6 +53,7 @@
     "import pandas as pd\n",
     "import tqdm\n",
     "import warnings\n",
+    "import multiprocessing\n",
     "\n",
     "from obspy import UTCDateTime, read\n",
     "from obspy.clients.fdsn import Client\n",
@@ -137,11 +138,15 @@
     "for day in pbar:\n",
     "    datestr = day.strftime(\"%Y-%m-%d\")\n",
     "    fn = \"{}_{}_{}.mseed\".format(dataset, datestr, nslc)\n",
-    "    if day != UTCDateTime().datetime and os.path.isfile(fn):\n",
+    "    mseedid = \"{:}.{:}.{:}.{:}\".format(network, station, location, channel)\n",
+    "    fn_npz = \"{}_{}_{}.npz\".format(dataset, datestr, mseedid)\n",
+    "    if day != UTCDateTime().datetime and (os.path.isfile(fn) or os.path.isfile(fn_npz)):\n",
+    "        pbar.set_description(\"Either mSEED or .npz already exists ... skipping date %s\" % fn)\n",
     "        continue\n",
     "    else:\n",
     "        pbar.set_description(\"Fetching %s\" % fn)\n",
     "        try: \n",
+    "            print(datestr)\n",
     "            st = c.get_waveforms(network, station, location, channel,\n",
     "                                  UTCDateTime(day)-1801, UTCDateTime(day)+86400+1801,\n",
     "                                  attach_response=True)\n",
@@ -150,7 +155,7 @@
     "            continue\n",
     "        st.write(fn)\n",
     "        \n",
-    "resp = c.get_stations(UTCDateTime(day), network=network, station=station, location=location,\n",
+    "resp = c.get_stations(start=UTCDateTime(daylist[0]), end=UTCDateTime(daylist[-1]), network=network, station=station, location=location,\n",
     "                      channel=channel, level=\"response\")\n",
     "print(resp)\n"
    ]
@@ -176,32 +181,39 @@
    "outputs": [],
    "source": [
     "force_reprocess = False\n",
-    "pbar = tqdm.tqdm(datelist)\n",
-    "for day in pbar:\n",
-    "    datestr = day.strftime(\"%Y-%m-%d\")\n",
-    "    fn_in = \"{}_{}_{}.mseed\".format(dataset, datestr, nslc)\n",
-    "    pbar.set_description(\"Processing %s\" % fn_in)\n",
-    "    if not os.path.isfile(fn_in):\n",
-    "        continue\n",
-    "    stall = read(fn_in, headonly=True)\n",
-    "    for mseedid in list(set([tr.id for tr in stall])):\n",
-    "        fn_out = \"{}_{}_{}.npz\".format(dataset, datestr, mseedid)\n",
-    "        if os.path.isfile(fn_out) and not force_reprocess:\n",
+    "def process(datelist_):\n",
+    "    pbar = tqdm.tqdm(datelist_)\n",
+    "    for day in pbar:\n",
+    "        print(day)\n",
+    "        datestr = day.strftime(\"%Y-%m-%d\")\n",
+    "        fn_in = \"{}_{}_{}.mseed\".format(dataset, datestr, nslc)\n",
+    "        pbar.set_description(\"Processing %s\" % fn_in)\n",
+    "        if not os.path.isfile(fn_in):\n",
     "            continue\n",
-    "        st = read(fn_in, sourcename=mseedid)\n",
-    "        st.attach_response(resp)\n",
-    "        ppsd = PPSD(st[0].stats, metadata=resp,\n",
-    "                    ppsd_length=1800, overlap=0.5,\n",
-    "                    period_smoothing_width_octaves=0.025,\n",
-    "                    period_step_octaves=0.0125,\n",
-    "                    period_limits=(0.008, 50),\n",
-    "                    db_bins=(-200, 20, 0.25))\n",
-    "        with warnings.catch_warnings():\n",
-    "            warnings.simplefilter(\"ignore\")\n",
-    "            ppsd.add(st)\n",
-    "        ppsd.save_npz(fn_out[:-4])\n",
-    "        del st, ppsd\n",
-    "    del stall\n"
+    "        stall = read(fn_in, headonly=True)\n",
+    "        for mseedid in list(set([tr.id for tr in stall])):\n",
+    "            fn_out = \"{}_{}_{}.npz\".format(dataset, datestr, mseedid)\n",
+    "            if os.path.isfile(fn_out) and not force_reprocess:\n",
+    "                continue\n",
+    "            st = read(fn_in, sourcename=mseedid)\n",
+    "            st.attach_response(resp)\n",
+    "            ppsd = PPSD(st[0].stats, metadata=resp,\n",
+    "                        ppsd_length=1800, overlap=0.5,\n",
+    "                        period_smoothing_width_octaves=0.025,\n",
+    "                        period_step_octaves=0.0125,\n",
+    "                        period_limits=(0.008, 50),\n",
+    "                        db_bins=(-200, 20, 0.25))\n",
+    "            with warnings.catch_warnings():\n",
+    "                warnings.simplefilter(\"ignore\")\n",
+    "                ppsd.add(st)\n",
+    "            ppsd.save_npz(fn_out[:-4])\n",
+    "            del st, ppsd\n",
+    "        del stall\n",
+    "        \n",
+    "nproc = multiprocessing.cpu_count()\n",
+    "datelist_splt = np.array_split(datelist, nproc)\n",
+    "with multiprocessing.Pool(processes=nproc) as pool:\n",
+    "    pool.map(process, datelist_splt)\n"
    ]
   },
   {
@@ -480,7 +492,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.6"
+   "version": "3.7.3"
   }
  },
  "nbformat": 4,