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