1
Fork 0
internship/data/processing/orbitals.py

103 lines
2.8 KiB
Python

import argparse
import configparser
import logging
import pathlib
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
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"))
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(figsize=(5/2.54, 2/3*10/2.54), dpi=200, constrained_layout=True)
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] * 1e-2,
z0[:-1] * 1e-2,
np.diff(x0)[:] * 1e-2,
np.diff(z0)[:] * 1e-2,
color="k",
scale_units="xy",
scale=1,
)
ax2dv.grid(c="k", alpha=.2)
ax2dv.set(aspect="equal", xlabel="x (m)", ylabel="z (m)")
ax2dv.set(ylim=(-10, 10))
ax2dv.yaxis.set_minor_locator(MultipleLocator(1))
fig2dv.savefig(out_root.joinpath("orbitals.pdf"))
fig2dv.savefig(out_root.joinpath("out_orbitals.jpg"))
plt.show()