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(), + )