internship/data/processing/orbitals.py

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()