From e1dc67d089cb22b9c13209947f3fa7fc316494b6 Mon Sep 17 00:00:00 2001
From: "Edgar P. Burkhart" <git@edgarpierre.fr>
Date: Tue, 15 Mar 2022 11:37:59 +0100
Subject: [PATCH] Switched to multithreading

---
 swash/processing/read_swash.py |  3 +++
 swash/processing/sws_npz.py    | 40 +++++++++++++++-------------------
 2 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/swash/processing/read_swash.py b/swash/processing/read_swash.py
index bacbd09..e5859da 100644
--- a/swash/processing/read_swash.py
+++ b/swash/processing/read_swash.py
@@ -31,6 +31,9 @@ class ReadSwash:
             return self.read_nohead(path).reshape((self._n_t, self._n_x))[0, :]
         return self.read_nohead(path).reshape((self._n_t, self._n_x))
 
+    def read_const(self, path):
+        return self.read_scalar(path, const=True)
+
     def read_vector(self, path):
         return self.read_nohead(path).reshape((self._n_t, 2, self._n_x))
 
diff --git a/swash/processing/sws_npz.py b/swash/processing/sws_npz.py
index 32622d4..7aef1cd 100644
--- a/swash/processing/sws_npz.py
+++ b/swash/processing/sws_npz.py
@@ -1,6 +1,7 @@
 import argparse
 import configparser
 import logging
+from multiprocessing import ThreadPool
 import pathlib
 
 import numpy as np
@@ -27,28 +28,21 @@ rsws = ReadSwash()
 rsws.read_time(sws_out.joinpath("tsec.dat"))
 rsws.read_x(sws_out.joinpath("xp.dat"))
 
-inp.mkdir(exist_ok=True)
-log.info(f"Wrinting output in '{inp}'")
-log.info("Reading 'dep'")
-np.save(inp.joinpath("dep"), rsws.read_scalar(sws_out.joinpath("dep.dat")))
-log.info("Reading 'botl'")
-np.save(
-    inp.joinpath("botl"),
-    rsws.read_scalar(sws_out.joinpath("botl.dat"), const=True),
-)
-log.info("Reading 'watl'")
-np.save(inp.joinpath("watl"), rsws.read_scalar(sws_out.joinpath("watl.dat")))
-log.info("Reading 'vel'")
-np.save(inp.joinpath("vel"), rsws.read_vector(sws_out.joinpath("vel.dat")))
-log.info("Reading 'press'")
-np.save(inp.joinpath("press"), rsws.read_scalar(sws_out.joinpath("press.dat")))
-log.info("Reading 'zk'")
-np.save(inp.joinpath("zk"), rsws.read_scalar_lay(sws_out.joinpath("zk.dat")))
-log.info("Reading 'velk'")
-np.save(
-    inp.joinpath("velk"), rsws.read_vector_lay(sws_out.joinpath("velk.dat"))
-)
+var = {
+    "dep": rsws.read_scalar,
+    "botl": rsws.read_const,
+    "watl": rsws.read_scalar,
+    "vel": rsws.read_vector,
+    "press": rsws.read_scalar,
+    "zk": rsws.read_scalar_lay,
+    "velk": rsws.read_vector_lay,
+}
 
-log.info(f"Writing npz file in '{inp}'")
 inp.mkdir(exist_ok=True)
-np.savez_compressed(inp.joinpath("sws"), t=rsws.t, x=rsws.x, **rsws.data)
+with ThreadPool() as pool:
+    pool.map(
+        lambda name, f: np.save(
+            inp.joinpath(name), f(sws_out.joinpath(name).with_suffix(".dat"))
+        ),
+        var.items(),
+    )