import argparse
import numpy as np
import plotly.graph_objs as go
import polars as pl
def build_sphere(size, num_points=500, color=(86, 180, 233), opacity=1.0):
theta = np.linspace(0, 2 * np.pi, num_points)
phi = np.linspace(0, np.pi, num_points)
x0 = size * np.outer(np.cos(theta), np.sin(phi))
y0 = size * np.outer(np.sin(theta), np.sin(phi))
z0 = size * np.outer(np.ones(num_points), np.cos(phi))
trace = go.Surface(
x=x0,
y=y0,
z=z0,
colorscale=[
[0, f"rgba({int(color[0])}, {int(color[1])}, {int(color[2])}, {opacity})"],
[1, f"rgba({int(color[0])}, {int(color[1])}, {int(color[2])}, {opacity})"],
],
hoverinfo="none",
)
trace.update(showscale=False)
return trace
def plot_traj(
df: pl.DataFrame,
colored_by="prop_mass (kg)",
color_descr="prop mass (kg)",
scale=1.0,
radius=6378.1363, color=(86, 180, 233),
):
df = df.with_columns(
pl.col("Epoch (UTC)").str.to_datetime("%Y-%m-%dT%H:%M:%S%.f")
).sort("Epoch (UTC)", descending=False)
fig = go.Figure()
traces = [build_sphere(radius, color, opacity=0.8)]
if isinstance(df, list):
all_traj = df
else:
all_traj = [df]
for traj_df in all_traj:
traces += [
go.Scatter3d(
x=traj_df["x (km)"],
y=traj_df["y (km)"],
z=traj_df["z (km)"],
mode="lines",
line=dict(
color=traj_df[colored_by] * scale,
colorscale="Viridis",
colorbar=dict(title=color_descr),
),
name="GEO",
)
]
fig.add_traces(traces)
fig.update_layout(scene_aspectmode="data")
fig.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="3D traj plotter")
parser.add_argument("pq", type=str, help="Path to the parquet file")
args = parser.parse_args()
df = pl.read_parquet(args.pq)
plot_traj(df)
plot_traj(
df,
colored_by="penumbra event light-source: Sun J2000, shadows casted by: Earth J2000, Moon J2000",
color_descr="Illumination %",
scale=100.0,
)