98 lines
2.5 KiB
Python
98 lines
2.5 KiB
Python
import argparse
|
|
import configparser
|
|
import logging
|
|
import pathlib
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
parser = argparse.ArgumentParser(description="Plot orbitals")
|
|
parser.add_argument("-v", "--verbose", action="count", default=0)
|
|
parser.add_argument("-c", "--config", default="config.ini")
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(level=max((10, 20 - 10 * args.verbose)))
|
|
log = logging.getLogger("bathy")
|
|
|
|
log.info("Starting time-series pre-processing")
|
|
config = configparser.ConfigParser()
|
|
config.read(args.config)
|
|
|
|
inp_root = pathlib.Path(config.get("inp", "root"))
|
|
out_root = pathlib.Path(config.get("out", "root"))
|
|
|
|
out_ts = out_root.joinpath("ts.dat")
|
|
|
|
raw_ts = []
|
|
for tsi in config.get("inp", "raw_ts").split(","):
|
|
raw_ts.append(
|
|
np.loadtxt(
|
|
inp_root.joinpath(tsi),
|
|
dtype=[("state", int), ("z", float), ("y", float), ("x", float)],
|
|
delimiter=",",
|
|
max_rows=2304,
|
|
)
|
|
)
|
|
n = len(raw_ts)
|
|
raw_ts = np.concatenate(raw_ts)
|
|
log.debug(f"{raw_ts=}")
|
|
|
|
if (errs := np.count_nonzero(raw_ts["state"])) != 0:
|
|
log.warning(f"{errs} transmission errors!")
|
|
log.debug(f"{dict(zip(*np.unique(raw_ts['state'], return_counts=True)))}")
|
|
|
|
t = np.linspace(0, 30 * 60 * n, 2304 * n + 1)[:-1]
|
|
log.debug(f"{t=}")
|
|
|
|
flt = (t > 1370) & (t < 1405)
|
|
|
|
figt, axt = plt.subplots(3)
|
|
axt[0].plot(t, raw_ts["x"])
|
|
axt[1].plot(t, raw_ts["y"])
|
|
axt[2].plot(t, raw_ts["z"])
|
|
for ax in axt:
|
|
ax.axvline(t[flt].min(), c="k")
|
|
ax.axvline(t[flt].max(), c="k")
|
|
ax.grid()
|
|
ax.set(xlim=(t.min(), t.max()))
|
|
|
|
ts_flt = raw_ts[flt]
|
|
z0 = ts_flt["z"]
|
|
figtz, axtz = plt.subplots(3)
|
|
axtz[0].plot(t[flt], ts_flt["x"])
|
|
axtz[1].plot(t[flt], ts_flt["y"])
|
|
axtz[2].plot(t[flt], z0)
|
|
for ax in axtz:
|
|
ax.grid()
|
|
ax.set(xlim=(t[flt].min(), t[flt].max()))
|
|
|
|
fig3d = plt.figure()
|
|
ax3d = fig3d.add_subplot(projection="3d")
|
|
ax3d.plot(ts_flt["x"], ts_flt["y"], z0, c="#0066ff")
|
|
ax3d.quiver3D(
|
|
ts_flt["x"][:-1],
|
|
ts_flt["y"][:-1],
|
|
z0[:-1],
|
|
np.diff(ts_flt["x"])[:],
|
|
np.diff(ts_flt["y"])[:],
|
|
np.diff(z0)[:],
|
|
color="#0066ff",
|
|
)
|
|
ax3d.set(xlabel="x (cm)", ylabel="y (cm)", zlabel="z (cm)")
|
|
|
|
theta = np.angle(raw_ts["x"] + 1j * raw_ts["y"]).mean()
|
|
fig2dv, ax2dv = plt.subplots()
|
|
x0 = ts_flt["x"] * np.cos(theta) + ts_flt["y"] * np.sin(theta)
|
|
ax2dv.plot(x0, z0, c="#0066ff", lw=1)
|
|
ax2dv.quiver(
|
|
x0[:-1],
|
|
z0[:-1],
|
|
np.diff(x0)[:],
|
|
np.diff(z0)[:],
|
|
color="#0066ff",
|
|
)
|
|
ax2dv.grid()
|
|
ax2dv.set(aspect="equal")
|
|
fig2dv.savefig("out_orbitals.pdf")
|
|
|
|
plt.show()
|