import numpy as np
import plotly.graph_objects as go
#
# Load vector field: x y z vx vy vz
# -------------------------
data = np.loadtxt("./vector_field.txt")
x, y, z, vx, vy, vz = data.T
#
# Compute vector magnitudes
magnitude = np.sqrt(vx**2 + vy**2 + vz**2)
#
# Plot using Plotly Cone
# -------------------------
fig = go.Figure(data=go.Cone(
    x=x,
    y=y,
    z=z,
    u=vx,
    v=vy,
    w=vz,
    colorscale='jet',     # color map
    cmin=magnitude.min(),
    cmax=magnitude.max(),
    # color=magnitude,          # color by vector magnitude
    sizemode="absolute",      # size in absolute units
    sizeref=3,              # adjust arrow size globally
    anchor="tail"             # arrow starts at the base point
))
#
# Layout
# -------------------------
fig.update_layout(
    scene=dict(
        xaxis_title='X',
        yaxis_title='Y',
        zaxis_title='Z',
        aspectmode='data'  # maintain correct aspect ratio
    ),
    title="3D Vector Field",
    margin=dict(l=0, r=0, b=0, t=50)
)
#
# Show figure
# -------------------------
fig.show()