Add example applying forces to shuttle model

- fix shuttle model to recalculate inertia matrix after centering
- refactor 3d plotting functions
- add functions for using quaternions
- add quaternion based model for integrating thrusts
- add exponential map based model for integrating thrusts
- refactor state space models for approximately maintaining constraints
- add euler step function
- add animation function for 3d

controls/stepping.py

@@ -53,4 +53,46 @@ def solve_cts_eqns(cts_state_model, x0, u_fcn, DT, T, cts_sim_model=None):
             lambda t, x: cts_sim_model.step(x, u[:, i]),
             t_span, x[:, i], t_eval=t_eval)
         x[:, [i + 1]] = sol.y
+        x[:, i + 1] = cts_state_model.reset_state(x[:, i + 1])
     return x, u
+
+
+def solve_eqns_euler(cts_state_model, x0, u_fcn, DT, T, cts_sim_model=None, dt_euler=None):
+    """
+    :param cts_state_model: model used for controller
+    :param x0: m, |
+    :param u_fcn: (x_t, t) -> u_t \in n, |
+    :param DT: sampling time
+    :param T: final time, steps made for t < T
+    :param cts_sim_model: true physical model used to advance simulation
+    :param dt_euler: sampling time used in Euler steps
+    :return:
+        x: m, k | ith column is state at t = i*DT
+        u: n, k | ith column is control at t
+    """
+    cts_sim_model = cts_sim_model if cts_sim_model else cts_state_model
+    m, n = cts_state_model.get_mn()
+    t = np.arange(0, T, DT)
+    x = np.zeros((m, t.size))
+    u = np.zeros((n, t.size))
+    x[:, 0] = x0.copy()
+    dt_euler = dt_euler or DT / 10
+    for i in range(0, t.size - 1):
+        operators = cts_state_model.get_process_operators(x[:, i])
+        u[:, i] = u_fcn(x[:, i], i * DT, **operators)
+        t_span = np.arange(i*DT, (i+1)*DT + dt_euler, dt_euler)
+        x[:, i + 1] = euler_step(
+            lambda t, x: cts_sim_model.step(x, u[:, i]),
+            t_span, x[:, i])[:, -1]
+        x[:, i + 1] = cts_state_model.reset_state(x[:, i + 1])
+    return x, u
+
+
+def euler_step(fun, t_span, y0):
+    y = np.zeros((y0.size, t_span.size))
+    y[:, 0] = y0
+    h = t_span[1:] - t_span[:-1]
+    for i in range(h.size):
+        y_dot = fun(t_span[i + 1], y[:, i])
+        y[:, i + 1] = y[:, i] + h[i] * y_dot
+    return y

examples/display_forces_on_shuttle.py

+import numpy as np
+from rigid_body_models.shuttle import Shuttle
+from state_space_models import quaternion_floating_integrator as qfli
+from state_space_models import floating_integrator as fli
+from misc import quaternion as qua
+from controls import time_based
+from controls import stepping
+from vis import animate_3d as ani
+
+
+def main_quaternion_integrator():
+    DT = 0.1
+    model = Shuttle(rho=1e1)
+    cts_state_model = qfli.FreeFloating(model)
+    x0 = np.array([
+        0., 0, 0,
+        0, 0, 0., 0,
+        0, 0, 0,
+        0, 0, 0,
+    ])
+    x0[3:7] = qua.exp_map2quaternion(np.array([0, 0, 1.0*np.pi/2]))
+    print(x0[3:7], np.linalg.norm(x0[3:7]))
+    u_fcn = time_based.make_control_signal(
+        [1, 8],
+        np.array([
+            [0., 500., 0, 500],
+            [0., 0, 0, 0],
+        ]).T
+    )
+    T = 5
+    x, u = stepping.solve_cts_eqns(cts_state_model, x0, u_fcn, DT, T)
+    np.set_printoptions(suppress=True)
+    print(x.T[-20:].round(2))
+
+    scale = .01
+    ani.loop_sim(model, x.T[:, :3], x.T[:, 3:7], u.T * scale, DT, is_quaternion=True)
+
+
+def main_exp_map_integrator():
+    DT = 0.1
+    model = Shuttle(rho=1e1)
+    cts_state_model = fli.FreeFloating(model)
+    x0 = np.array([
+        0., 0, 0,
+        0, 0, 1.0 * np.pi / 2,
+        0, 0, 0,
+        0, 0, 0,
+    ])
+    u_fcn = time_based.make_control_signal(
+        [1, 8],
+        np.array([
+            [0., 0., 1000, 1000],
+            [0., 0, 0, 0],
+        ]).T
+    )
+    T = 5
+    x, u = stepping.solve_cts_eqns(cts_state_model, x0, u_fcn, DT, T)
+    np.set_printoptions(suppress=True)
+    print(x.T[-20:].round(2))
+
+    scale = .01
+    ani.loop_sim(model, x.T[:, :3], x.T[:, 3:6], u.T * scale, DT, is_quaternion=False)
+
+
+if __name__ == '__main__':
+    main_quaternion_integrator()
+    # main_exp_map_integrator()

examples/display_shuttle_model.py

@@ -7,12 +7,13 @@ from vis import triangulated_thruster_model as vttm
 def main_display_mesh(is_thrusters=True):
     model = Shuttle(rho=1e6)
     fig = go.Figure()
-    vtrs.draw_model(fig, model)
+    fig.add_traces([vtrs.make_model_mesh(model)])
     if is_thrusters:
-        vttm.draw_thrust_positions(fig, model)
+        fig.add_traces(vttm.make_thrust_vectors(model))
         title_str = 'Shuttle thrusters'
     else:
-        vtrs.draw_surface_normals(fig, model)
+        fig.add_traces(vtrs.make_surface_normals(
+            model.vertices, model.triangles))
         title_str = 'Shuttle surface normals'
     fig.update_layout(
         title=title_str, autosize=False,
@@ -26,4 +27,4 @@ def main_display_mesh(is_thrusters=True):
 
 
 if __name__ == '__main__':
-    main_display_mesh(is_thrusters=False)
+    main_display_mesh(is_thrusters=True)

misc/matrix_building.py

@@ -33,5 +33,21 @@ def Rt_planar_xyz1(theta, t):
     return mat
 
 
+def Rt_3d(R=(), t=()):
+    """
+    Make homogeneous transformation matrix
+    - identity operation if not specified
+    :param R: 3, 3
+    :param t: 3,
+    :return: 4, 4
+    """
+    if len(R) == 0:
+        R = np.eye(3)
+    if len(t) == 0:
+        t = np.zeros((3,))
+    Rt = np.eye(4)
+    Rt[:3, :3] = R
+    Rt[:3, -1] = t
+    return Rt
 
 

misc/quaternion.py

+import numpy as np
+
+EPS = np.finfo(float).eps
+EPSrt4 = EPS ** (1/4)
+
+
+def dot_mult_matrix(q):
+    """
+    Quaternion dot multiplication has an equivalent matrix multiplication:
+    q \odot [x; 0] = C(q)x  (eqn 2.98 in MC)
+
+    with C(q) =  ( q_4 I_3 + [q_{1:3} \cross]
+                        -q_{1:3}'             )
+    :param q: 4, | quaternion
+    :return: 4, 3 | C(q) matrix
+    """
+    C = np.zeros((4, 3))
+    C[:3, :3] = cross_matrix(q[:3])
+    C[[0, 1, 2], [0, 1, 2]] = q[3]
+    C[-1] = -q[:3]
+    return C
+
+
+def quaternion2rotation_matrix(q):
+    """
+    Rotating a vector has an equivalent matrix multiplication:
+    (q \kron {x ; 0} \kron q*)[:3] = A(q)x
+
+    with A(q) = (q_4^2 - ||q_{1:3}||^2) I_3
+                - 2 q_4 [q_{1:3} \cross]
+                + 2 q_{1:3} q_{1:3}'       (eqn 2.125 in MC)
+    :param q: 4, | quaternion
+    :return: 3, 3 |
+    """
+    n = np.linalg.norm(q[:3])
+    A = (q[3] ** 2 - n ** 2) * np.eye(3) \
+        - 2 * q[3] * cross_matrix(q[:3]) \
+        + 2 * np.outer(q[:3], q[:3])
+    return A
+
+
+def exp_map2rotation_matrix(v):
+    """
+    :param v: 3, | exponential map representation of rotation
+    :return: 3, 3
+    """
+    q = exp_map2quaternion(v)
+    A = quaternion2rotation_matrix(q)
+    return A
+
+
+def exp_map2quaternion(v):
+    """
+    Convert exponential map representation to quaternion,
+    where v \in \reals^3 that encodes a unit axis and the angle
+    that is rotated around the axis.
+
+    ||v||_2 = \theta = the angle [rad]
+    v / ||v||_2 = unit axis
+
+    This is identified by setting \theta = 0 (no rotation) at v = 0,
+    but there is numerical error in dividing by ||v||_2 as v -> 0.
+
+    As proposed in (Grassia '99), achieve loss-less conversion by
+    using the taylor expansion of the sinc function.
+    S. F. Grassia. "Practical parameterization of rotations using the
+    exponential map," Journal of Graphics Tools, vol. 3 no. 3, pp.29-48, 1998.
+
+    :param v: 3, | exponential map representation of rotation
+    :return: q: 4, | quaternion
+    """
+    theta = np.linalg.norm(v)
+    q = np.zeros((4,))
+    if theta > EPSrt4:
+        q[:3] = np.sin(theta / 2) / theta
+    else:
+        q[:3] = 0.5 + theta ** 2 / 48
+    q[:3] *= v
+    q[3] = np.cos(theta / 2)
+    return q
+
+
+def quaternion2exp_map(q):
+    """
+    :param q: 4, | quaternion (not necessarily unit norm)
+    :return: v: 3, | exponential map representation of rotation
+    """
+    q_ = 1 * q
+    n = np.linalg.norm(q)
+    if n > 1:
+        # since arc cos range as [-1, 1]
+        q_ /= n
+    theta = 2 * np.arccos(q_[-1])
+    if np.abs(theta) < EPS:
+        return np.zeros((3,))
+    v = q_[:3] * theta / np.sin(theta / 2)
+    return v
+
+
+def reset_exp_map(v):
+    """
+    Reset proposed in (Grassia '99)
+    :param v: 3, | exponential map representation of rotation
+    :return:
+    """
+    theta = np.linalg.norm(v)
+    if theta < np.pi:
+        return v
+    return v * (1 - 2*np.pi / theta)
+
+
+def cross_matrix(a):
+    """
+    a \cross x = [a \cross] x = C(a)x
+    where
+    C(a) = [a \cross]
+    :param a: 3,
+    :return: 3, 3 | C(a)
+    """
+    C = np.array([
+        [0, -a[2], a[1]],
+        [a[2], 0, -a[0]],
+        [-a[1], a[0], 0],
+    ])
+    return C

rigid_body_models/shuttle.py

@@ -82,3 +82,7 @@ class Shuttle(trs.TriangulatedSurface):
         self.center = 0 * center
         self.vertices -= center
         self.r = (self.r - center[:, np.newaxis])
+        # recalculate inertia matrix after translating to center
+        _, _, im = trs.calculate_mass_properties(
+            self.vertices, self.triangles, rho)
+        self.im = im

rigid_body_models/triangulated_surface.py

@@ -69,15 +69,3 @@ def calculate_mass_properties(vertices, triangles, rho):
         [-s[6], -s[5], s[7] + s[8]],
     ])
     return m, center, im
-
-# 3d triangulation object -> display surface normals
-# 3d thruster object with triangulation -> display surface normals
-
-# shuttle = 3d thruster obj w tri
-
-# floating integrator = B based on thrust, A using exp map for rotations
-
-# structure for computing im from assemblies of components?
-# ~ recursive
-# - number of helper functions?
-# - then m, center, im as inputs to class?

state_space_models/empirical_drag.py

@@ -106,5 +106,8 @@ class EmpiricalDragCts(object):
         dfxdx = A - dfxdx_all_drag
         return dict(A=A, B=B_t, fx=fx, gx=gx, dfxdx=dfxdx)
 
+    def reset_state(self, x_t):
+        return x_t
+
     def get_mn(self):
         return self.m, self.n

state_space_models/floating_integrator.py

+"""
+Object state defined by:
+
+x_t: 12, | (p_t, q_t, v_t, H_t) full state vector
+p_t: 3, | 3d position of center of mass [m]
+q_t: 3, | exponential map representation for 3d orientation
+    where q/|q| is the axis
+    and |q| is the rotation about that axis [rad]
+v_t: 3, | velocity of center of mass [m/s]
+H_t: 3, | angular momentum H_I^c [kg m^2 / s]
+    about center of mass in body frame
+"""
+import numpy as np
+from misc import quaternion as qua
+
+
+class FreeFloating(object):
+    """
+    Model 6dof (translation + rotation) motion imparted by thrust inputs
+    in absence of other forces.
+
+    x_dot = fx + B(x)
+
+    p_dot = v_t + 0
+    q_dot = [0.5(J_B^c)^-1 A(q)H_I^c] \dot q + 0
+        - where q is quaternion representation of q_t
+        - dot is the quaternion dot multiplication
+    v_dot = 0 + A(q) f/m = A(q) b_part
+    H_dot = 0 + \sum_{i=1}^n r_i \cross f_i u_i
+
+    Use exponential map to represent rotations, converting to
+    quaternions to create the rotation matrix when needed.
+    """
+
+    def __init__(self, mounted_thruster_model):
+        self.m = 12
+        self.A = np.zeros((self.m, self.m))
+        self.A[:3, 6:9] = np.eye(3)
+        f = mounted_thruster_model.f
+        r = mounted_thruster_model.r
+        self.n = f.shape[1]
+        self.B = np.zeros((self.m, self.n))
+        self.b_part = f / mounted_thruster_model.m
+        self.B[9:, :] = np.cross(r.T, f.T).T
+        self.J_Bc_inv = np.linalg.inv(mounted_thruster_model.im)
+
+    def step(self, x_t, u_t):
+        operators = self.get_process_operators(x_t)
+        fx, gx = operators['fx'], operators['gx']
+        x_dot = fx + gx.dot(u_t)
+        return x_dot
+
+    def get_process_operators(self, x_t):
+        """
+        Return operators which may be based on current state
+        :param x_t:
+        :return: dict{fx, gx, B(x_t)}
+        """
+        B_t = self.B
+        q = qua.exp_map2quaternion(x_t[3:6])
+        Aq = qua.quaternion2rotation_matrix(q)
+        B_t[6:9, :] = Aq.T.dot(self.b_part)
+        gx = B_t
+        omega = self.J_Bc_inv.dot(x_t[9:])
+        q_quaternion = 0.5 * qua.dot_mult_matrix(q).dot(omega)
+        fx_q_dot = qua.quaternion2exp_map(q_quaternion)
+        fx = self.A.dot(x_t)
+        fx[3:6] = fx_q_dot
+        return dict(fx=fx, gx=gx, A=self.A, B=B_t)
+
+    def reset_state(self, x_t):
+        """
+        Outside of IVP solver, reset exp map here to avoid
+        singularity at 2pi (= identity rotation but magnitude != 0)
+        :param x_t:
+        :return:
+        """
+        x_t[3:6] = qua.reset_exp_map(x_t[3:6])
+        return x_t
+
+    def get_mn(self):
+        return self.m, self.n

state_space_models/quaternion_floating_integrator.py

+"""
+Object state defined by:
+
+x_t: 13, | (p_t, q_t, v_t, H_t) full state vector
+p_t: 3, | 3d position of center of mass [m]
+q_t: 4, | quaternion representation for 3d orientation
+v_t: 3, | velocity of center of mass [m/s]
+H_t: 3, | angular momentum H_I^c [kg m^2 / s]
+    about center of mass in body frame
+"""
+import numpy as np
+from misc import quaternion as qua
+
+
+class FreeFloating(object):
+    """
+    Model 6dof (translation + rotation) motion imparted by thrust inputs
+    in absence of other forces.
+
+    x_dot = fx + B(x)
+
+    p_dot = v_t + 0
+    q_dot = [0.5(J_B^c)^-1 A(q)H_I^c] \dot q + 0
+        - A(q) is rotation matrix for orientation q
+        - dot is the quaternion dot multiplication
+    v_dot = 0 + A(q) f/m = A(q) b_part
+    H_dot = 0 + \sum_{i=1}^n r_i \cross f_i u_i
+
+    Approximately integrates rotation by normalizing the quaternion
+    after each initial value problem (IVP) is solved.
+    - Exact would be enforcing unit length throughout IVP
+    - Smaller step sizes (timestep) => better accuracy
+    """
+
+    def __init__(self, mounted_thruster_model):
+        self.m = 13
+        self.A = np.zeros((self.m, self.m))
+        self.A[:3, 7:10] = np.eye(3)
+        f = mounted_thruster_model.f
+        r = mounted_thruster_model.r
+        self.n = f.shape[1]
+        self.B = np.zeros((self.m, self.n))
+        self.b_part = f / mounted_thruster_model.m
+        self.B[10:, :] = np.cross(r.T, f.T).T
+        self.J_Bc_inv = np.linalg.inv(mounted_thruster_model.im)
+
+    def step(self, x_t, u_t):
+        operators = self.get_process_operators(x_t)
+        fx, gx = operators['fx'], operators['gx']
+        x_dot = fx + gx.dot(u_t)
+        return x_dot
+
+    def get_process_operators(self, x_t):
+        """
+        Return operators which may be based on current state
+        :param x_t:
+        :return: dict{fx, gx, B(x_t)}
+        """
+        B_t = self.B
+        q = x_t[3:7]
+        Aq = qua.quaternion2rotation_matrix(q / np.linalg.norm(q))
+        B_t[7:10, :] = Aq.T.dot(self.b_part)
+        gx = B_t
+        omega = self.J_Bc_inv.dot(x_t[10:])
+        fx_q_dot = 0.5 * qua.dot_mult_matrix(q).dot(omega)
+        fx = self.A.dot(x_t)
+        fx[3:7] = fx_q_dot
+        return dict(fx=fx, gx=gx, A=self.A, B=B_t)
+
+    def reset_state(self, x_t):
+        """
+        Outside of IVP solver, normalize quaternion to unit length
+        :param x_t:
+        :return:
+        """
+        q = x_t[3:7]
+        q /= max(np.linalg.norm(q), 1e-8)
+        return x_t
+
+    def get_mn(self):
+        return self.m, self.n

state_space_models/simple_planar_integrator.py

@@ -63,6 +63,9 @@ class NoDragCts(object):
         B_t[2:4, :] = mb.rot_2d(x_t[4]).dot(self.b_part)
         return dict(A=self.Ac, B=B_t)
 
+    def reset_state(self, x_t):
+        return x_t
+
     def get_mn(self):
         return self.m, self.n
 
@@ -90,6 +93,9 @@ class AccelApproxDragCts(object):
     def get_process_operators(self, x_t):
         return self.no_drag_cts_model.get_process_operators(x_t)
 
+    def reset_state(self, x_t):
+        return x_t
+
     def get_mn(self):
         return self.no_drag_cts_model.get_mn()
 
@@ -135,5 +141,8 @@ class ZeroOrderHoldDiscrete(object):
         self.Bd_t = eM[0:m, m:].dot(Bc_t)
         return dict(A=self.Ad_t, B=self.Bd_t)
 
+    def reset_state(self, x_t):
+        return x_t
+
     def get_mn(self):
         return self.cts_model.get_mn()

vis/animate_3d.py

+import plotly.graph_objects as go
+from misc import quaternion
+from misc.matrix_building import Rt_3d
+from vis import triangulated_surface as vtrs
+from vis import triangulated_thruster_model as vttm
+
+
+def loop_sim(model, x, q, u, dt, scene=None, is_quaternion=False, **kwargs):
+    """
+
+    :param model: triangulated surface
+    :param x: k, 3 | 3d position of center of model
+    :param q: k, 3 | exponential map representation for 3d orientation
+    - iff is_quaternion=True, instead treat as k, 4 quaternions
+    :param u: k, n | control inputs
+    :param dt:
+    :param scene:
+    :param kwargs:
+    :return:
+    """
+    k = x.shape[0]
+    frame_data = []
+    for i in range(k):
+        if is_quaternion:
+            Aq = quaternion.quaternion2rotation_matrix(q[i])
+        else:
+            Aq = quaternion.exp_map2rotation_matrix(q[i])
+        Rt = Rt_3d(Aq.T, x[i])
+        data = [vtrs.make_model_mesh(model, Rt)]
+        data.extend(vttm.make_thrust_vectors(model, scales=u[i], Rt=Rt))
+        frame_data.append(data)
+    drawn_frames = [go.Frame(
+        data=frame_data[i],
+        layout=go.Layout(
+            title_text='t = {:2.2f}s'.format(i * dt),
+            scene=dict(
+                xaxis=dict(range=x[i, 0] + [-10, 10], autorange=False),
+                yaxis=dict(range=x[i, 1] + [-10, 10], autorange=False),
+                zaxis=dict(range=x[i, 2] + [-10, 10], autorange=False),
+                aspectmode='cube',
+            ),
+        )
+    ) for i in range(len(frame_data))]
+
+    fig = go.Figure(
+        data=frame_data[0],
+        layout=go.Layout(
+            title='t = {:2.2f}s'.format(0),
+            autosize=False,
+            width=1200, height=800,
+            scene=dict(
+                xaxis=dict(range=x[0, 0] + [-10, 10], autorange=False),
+                yaxis=dict(range=x[0, 1] + [-10, 10], autorange=False),
+                zaxis=dict(range=x[0, 2] + [-10, 10], autorange=False),
+                aspectmode='cube',
+            ),
+            updatemenus=[dict(
+                type="buttons",
+                buttons=[dict(label='Play',
+                              method='animate',
+                              args=[None, {'frame': {'duration': 200, 'redraw': True}, }])])]
+        ),
+        frames=drawn_frames
+    )
+    fig.update_layout(transition_duration=500)
+    fig.show()

vis/triangulated_surface.py

 import numpy as np
 import plotly.graph_objects as go
 from vis.xyz_utils import make_3d_vector
+from misc.matrix_building import Rt_3d
 
 
-def draw_model(fig, model):
-    xyz = model.vertices
+def make_model_mesh(model, Rt=()):
+    """
+    :param model:
+    :param Rt: 4, 4 | homogeneous coordinate rigid body transform
+        - body frame to given frame
+    :return:
+    """
+    if len(Rt) == 0:
+        Rt = Rt_3d()
+    xyz = Rt[:3, :3].dot(model.vertices.T).T + Rt[:3, -1]
     ijk = model.triangles
     mesh = go.Mesh3d(
         x=xyz[:, 0], y=xyz[:, 1], z=xyz[:, 2],
         i=ijk[:, 0], j=ijk[:, 1], k=ijk[:, 2],
         color='white', flatshading=True
     )
-    fig.add_traces([mesh])
-
-
-def draw_surface_normals(fig, model):
-    fig.add_traces(make_surface_normals(model.vertices, model.triangles))
+    return mesh
 
 
 def make_surface_normals(vertices, triangles):

vis/triangulated_thruster_model.py

 import numpy as np
-import plotly.graph_objects as go
 from vis.xyz_utils import make_3d_vector
+from misc.matrix_building import Rt_3d
 
 
-def draw_thrust_positions(fig, model):
+def make_thrust_vectors(model, scales=(), Rt=()):
     """
     Draw thrust vectors with direction opposite to the force being
     applied, like the direction of exhaust.
     (since force direction may end up *inside* the 3d model)
-    :param fig:
     :param model:
+    :param scales: multipliers for vector line and head length
+    :param Rt: 4, 4 | homogeneous coordinate rigid body transform
+        - body frame to given frame
     :return:
     """
     object_list = []
     n_thrusters = model.r.shape[1]
+    if len(scales) == 0:
+        scales = np.ones((n_thrusters,))
+    if len(Rt) == 0:
+        Rt = Rt_3d()
+    r = Rt[:3, :3].dot(model.r) + Rt[:3, -1][:, np.newaxis]
+    f = Rt[:3, :3].dot(model.f)
     for i in range(n_thrusters):
         name = 'thruster {}'.format(i)
-        arrow_objs = make_3d_vector(model.r[:, i], -model.f[:, i], name=name)
+        arrow_objs = make_3d_vector(
+            r[:, i], -f[:, i], scale=scales[i], name=name)
         object_list.extend(arrow_objs)
-    fig.add_traces(object_list)
+    return object_list

vis/xyz_utils.py

@@ -2,16 +2,16 @@ import numpy as np
 import plotly.graph_objects as go
 
 
-def make_3d_vector(root, vector, head_length=1., name=''):
+def make_3d_vector(root, vector, scale=1., name=''):
     """
     Make a 3d vector for plotting
     :param root: 3,
     :param vector: 3,
-    :param head_length:
+    :param scale: multiplier for vector line and head length
     :param name:
     :return:
     """
-    head_xyz = root + vector
+    head_xyz = root + scale * vector
     xyz = np.vstack((root, head_xyz))
     line_obj = go.Scatter3d(
         x=xyz[:, 0],
@@ -21,7 +21,7 @@ def make_3d_vector(root, vector, head_length=1., name=''):
         line=dict(width=4),
         name=name,
     )
-    n = head_length * vector / np.linalg.norm(vector)
+    n = scale * vector / np.linalg.norm(vector)
     arrow_head_obj = go.Cone(
         x=[head_xyz[0]], y=[head_xyz[1]], z=[head_xyz[2]],
         u=[n[0]], v=[n[1]], w=[n[2]],