-
Notifications
You must be signed in to change notification settings - Fork 1
/
fluxes.py
52 lines (42 loc) · 1.45 KB
/
fluxes.py
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
#!/usr/bin/env python3
import numpy
from matplotlib import pyplot
import h5py
# Need one extra so that each node has upper and lower
# edges in the mesh plot. Avoids dropping the highest row
# Will do similar when updating the days
HEIGHTS = numpy.linspace(0, 360, 101)
UPPER = 311
def plot_flux(g: h5py.Group, ax: pyplot.Axes):
dayvec = numpy.empty(g["days"].size + 1, dtype=g["days"].dtype)
g["days"].read_direct(dayvec, dest_sel=numpy.s_[: dayvec.size - 1])
dayvec[-1] = UPPER
flux = numpy.divide(g["flux"], numpy.max(g["flux"]))
ax.pcolormesh(dayvec, HEIGHTS, flux.T, rasterized=True)
def main() -> pyplot.Figure:
fig, axes = pyplot.subplots(2, 2, sharex=True, sharey=True, figsize=(8, 4))
with h5py.File("./data.h5", "r") as data:
for col, title, fmt in zip(
(axes[:, 0], axes[:, 1]),
("Hybrid", "Serpent predictor"),
("hybrid/{}", "serpent/ce/{}"),
):
for step, ax in zip((5, 10), col):
plot_flux(data[fmt.format(step)], ax)
ax.set_title(f"{title} - {step} day step")
axes[0, 0].set(
ylabel="Axial height [cm]",
)
axes[1, 0].set(
ylabel=axes[0, 0].get_ylabel(),
xlabel="Time [d]",
)
axes[1, 1].set_xlabel(axes[1, 0].get_xlabel())
return fig
if __name__ == "__main__":
import sys
fig = main()
if len(sys.argv) == 1:
pyplot.show()
else:
fig.savefig(sys.argv[1])