PySDM_examples.Szumowski_et_al_1998.storage

 1import os
 2import tempfile
 3from pathlib import Path
 4
 5import numpy as np
 6
 7
 8class Storage:
 9    class Exception(BaseException):
10        pass
11
12    def __init__(self, dtype=np.float32, path=None):
13        self.temp_dir = None
14        if path is None:
15            self.setup_temporary_directory()
16        else:
17            Path(path).mkdir(parents=True, exist_ok=True)
18            self.dir_path = Path(path).absolute()
19        self.dtype = dtype
20        self.grid = None
21        self._data_range = None
22
23    def __del__(self):
24        self.cleanup()
25
26    def cleanup(self):
27        if self.temp_dir is not None:
28            self.temp_dir.cleanup()
29
30    def setup_temporary_directory(self):
31        self.cleanup()
32        self.temp_dir = tempfile.TemporaryDirectory()
33        self.dir_path = self.temp_dir.name
34
35    def init(self, settings):
36        self.grid = settings.grid
37        self._data_range = {}
38        if self.temp_dir is not None and any(os.scandir(self.temp_dir.name)):
39            self.setup_temporary_directory()
40
41    def _filepath(self, name: str, step: int = None):
42        if step is None:
43            filename = f"{name}.npy"
44        else:
45            filename = f"{name}_{step:06}.npy"
46        path = os.path.join(self.dir_path, filename)
47        return path
48
49    def save(self, data: (float, np.ndarray), step: int, name: str):
50        if isinstance(data, (int, float)):
51            path = self._filepath(name)
52            np.save(
53                path,
54                np.concatenate(
55                    (() if step == 0 else np.load(path), (self.dtype(data),))
56                ),
57            )
58        elif data.shape[0:2] == self.grid:
59            np.save(self._filepath(name, step), data.astype(self.dtype))
60        else:
61            raise NotImplementedError()
62
63        if name not in self._data_range:
64            self._data_range[name] = (np.inf, -np.inf)
65        just_nans = np.isnan(data).all()
66        self._data_range[name] = (
67            min(
68                self._data_range[name][0] if just_nans else np.nanmin(data),
69                self._data_range[name][0],
70            ),
71            max(
72                self._data_range[name][1] if just_nans else np.nanmax(data),
73                self._data_range[name][1],
74            ),
75        )
76
77    def data_range(self, name):
78        return self._data_range[name]
79
80    def load(self, name: str, step: int = None) -> np.ndarray:
81        try:
82            data = np.load(self._filepath(name, step))
83        except FileNotFoundError as err:
84            raise Storage.Exception() from err
85        return data