-
Notifications
You must be signed in to change notification settings - Fork 38
Expand file tree
/
Copy pathbase.py
More file actions
152 lines (126 loc) · 4.17 KB
/
Copy pathbase.py
File metadata and controls
152 lines (126 loc) · 4.17 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
from pathlib import Path
import typing as T
import numpy as np
from datetime import datetime
from .raw import get_simsize as get_simsize_raw
try:
from .hdf import (
get_simsize as get_simsize_h5,
write_grid as write_grid_h5,
write_state as write_state_h5,
write_Efield as write_Efield_h5,
write_precip as write_precip_h5,
)
except ModuleNotFoundError:
get_simsize_h5 = write_grid_h5 = write_state_h5 = None
try:
from .nc4 import (
get_simsize as get_simsize_nc,
write_grid as write_grid_nc,
write_state as write_state_nc,
write_Efield as write_Efield_nc,
write_precip as write_precip_nc,
)
except ModuleNotFoundError:
get_simsize_nc = write_grid_nc = write_state_nc = None
Pathlike = T.Union[str, Path]
def get_simsize(path: Pathlike) -> T.Tuple[int, ...]:
""" get simulation dimenions """
path = Path(path).expanduser().resolve()
if path.is_dir():
for suffix in [".h5", ".nc", ".dat"]:
fn = path / ("simsize" + suffix)
if fn.is_file():
break
else:
fn = path
if not fn.stem == "simsize":
fn = path.parent / ("simsize" + path.suffix)
if not fn.is_file():
raise FileNotFoundError(path)
if fn.suffix == ".h5":
if get_simsize_h5 is None:
raise ModuleNotFoundError("pip install h5py")
return get_simsize_h5(fn)
elif fn.suffix == ".nc":
if get_simsize_nc is None:
raise ModuleNotFoundError("pip install netcdf4")
return get_simsize_nc(fn)
else:
return get_simsize_raw(fn)
def write_grid(p: T.Dict[str, T.Any], xg: T.Dict[str, T.Any]):
""" writes grid to disk
Parameters
----------
p: dict
simulation parameters
xg: dict
grid values
"""
if p["format"] in ("hdf5", "h5"):
if write_grid_h5 is None:
raise ImportError("pip install h5py")
write_grid_h5(p, xg)
elif p["format"] in ("netcdf", "nc"):
if write_grid_nc is None:
raise ImportError("pip install netcdf4")
write_grid_nc(p, xg)
else:
raise ValueError(f'unknown file format {p["format"]}')
def write_Efield(p: T.Dict[str, T.Any], xg: T.Dict[str, T.Any]):
""" writes grid to disk
Parameters
----------
p: dict
simulation parameters
xg: dict
grid values
"""
if p["format"] in ("hdf5", "h5"):
if write_Efield_h5 is None:
raise ImportError("pip install h5py")
write_Efield_h5(p, xg)
elif p["format"] in ("netcdf", "nc"):
if write_Efield_nc is None:
raise ImportError("pip install netcdf4")
write_Efield_nc(p, xg)
else:
raise ValueError(f'unknown file format {p["format"]}')
def write_precip(p: T.Dict[str, T.Any], xg: T.Dict[str, T.Any]):
""" writes grid to disk
Parameters
----------
p: dict
simulation parameters
xg: dict
grid values
"""
if p["format"] in ("hdf5", "h5"):
if write_precip_h5 is None:
raise ImportError("pip install h5py")
write_precip_h5(xg)
elif p["format"] in ("netcdf", "nc"):
if write_precip_nc is None:
raise ImportError("pip install netcdf4")
write_precip_nc(xg)
else:
raise ValueError(f'unknown file format {p["format"]}')
def write_state(time: datetime, ns: np.ndarray, vs: np.ndarray, Ts: np.ndarray, out_dir: Path, file_format: str):
"""
WRITE STATE VARIABLE DATA.
NOTE THAT WE don't write ANY OF THE ELECTRODYNAMIC
VARIABLES SINCE THEY ARE NOT NEEDED TO START THINGS
UP IN THE FORTRAN CODE.
INPUT ARRAYS SHOULD BE TRIMMED TO THE CORRECT SIZE
I.E. THEY SHOULD NOT INCLUDE GHOST CELLS
"""
if file_format in ("hdf5", "h5"):
if write_grid_h5 is None:
raise ImportError("pip install h5py")
write_state_h5(time, ns, vs, Ts, out_dir)
elif file_format in ("netcdf", "nc"):
if write_grid_nc is None:
raise ImportError("pip install netcdf4")
write_state_nc(time, ns, vs, Ts, out_dir)
else:
raise ValueError(f"unknown file format {file_format}")