diff --git a/controls/stepping.py b/controls/stepping.py
index a34a946879ff3351d97932d1f76398947b43eb2b..4b3d72ee34965b8a7fbc49222284a7a1e892c3fe 100644
--- a/controls/stepping.py
+++ b/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
diff --git a/examples/display_forces_on_shuttle.py b/examples/display_forces_on_shuttle.py
new file mode 100644
index 0000000000000000000000000000000000000000..e93ca26ab79c9b40039082dac5d65357c4127064
--- /dev/null
+++ b/examples/display_forces_on_shuttle.py
@@ -0,0 +1,67 @@
+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()
diff --git a/examples/display_shuttle_model.py b/examples/display_shuttle_model.py
index f70228894e35a6fcbb572f54db0e89a6d57e0c48..829955ea339d8e6f550eb9bcdb08baaf023ec4fc 100644
--- a/examples/display_shuttle_model.py
+++ b/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)
diff --git a/misc/matrix_building.py b/misc/matrix_building.py
index e8f267be51b6951fa7a43c41f0320a1ed19e61e5..5451c618bbc20bad7b0595cc914bbd8d68613304 100644
--- a/misc/matrix_building.py
+++ b/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
diff --git a/misc/quaternion.py b/misc/quaternion.py
new file mode 100644
index 0000000000000000000000000000000000000000..9ae058a41ab5abcfd368c6e806ff0239c5f951e9
--- /dev/null
+++ b/misc/quaternion.py
@@ -0,0 +1,125 @@
+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
diff --git a/rigid_body_models/shuttle.py b/rigid_body_models/shuttle.py
index ba62dbae604c5d968adb35cacd4c1fe11d5c3d7b..24e8a3b49436381f6bb4e2a3e83afa5cd26c527d 100644
--- a/rigid_body_models/shuttle.py
+++ b/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
diff --git a/rigid_body_models/triangulated_surface.py b/rigid_body_models/triangulated_surface.py
index e87bdcf1cc93f5b7c9358788971df48ba18b6a62..3a9a0e773ae4643329c4eb52953f32947d1ea2ce 100644
--- a/rigid_body_models/triangulated_surface.py
+++ b/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?
diff --git a/state_space_models/empirical_drag.py b/state_space_models/empirical_drag.py
index 49b8888c0ebb030293c6936cc9b5c31f837bd6a7..27f2b1e2135a38190d406d4d9761e2d9b2b4929f 100644
--- a/state_space_models/empirical_drag.py
+++ b/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
diff --git a/state_space_models/floating_integrator.py b/state_space_models/floating_integrator.py
new file mode 100644
index 0000000000000000000000000000000000000000..2a78777306826f2c60e7085c97aca9395f5d826b
--- /dev/null
+++ b/state_space_models/floating_integrator.py
@@ -0,0 +1,82 @@
+"""
+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
diff --git a/state_space_models/quaternion_floating_integrator.py b/state_space_models/quaternion_floating_integrator.py
new file mode 100644
index 0000000000000000000000000000000000000000..07a18a43c670c787c1193abf8d7fd1d0528c2607
--- /dev/null
+++ b/state_space_models/quaternion_floating_integrator.py
@@ -0,0 +1,81 @@
+"""
+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
diff --git a/state_space_models/simple_planar_integrator.py b/state_space_models/simple_planar_integrator.py
index 7af129e21bc9e2fc5d76f3fd5016aa6f01eaad95..963591ca2e1e896c60bc719e28bede923f45ef1e 100644
--- a/state_space_models/simple_planar_integrator.py
+++ b/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()
diff --git a/vis/animate_3d.py b/vis/animate_3d.py
new file mode 100644
index 0000000000000000000000000000000000000000..00960fb82b05e53d3e2fd88d6e76dabece4a1194
--- /dev/null
+++ b/vis/animate_3d.py
@@ -0,0 +1,66 @@
+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()
diff --git a/vis/triangulated_surface.py b/vis/triangulated_surface.py
index 4dc48d0436837de9723e73749a0f7932f0417536..3f48e81f071b1821ef7e699efeb8048d219e937a 100644
--- a/vis/triangulated_surface.py
+++ b/vis/triangulated_surface.py
@@ -1,21 +1,26 @@
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):
diff --git a/vis/triangulated_thruster_model.py b/vis/triangulated_thruster_model.py
index 4ca925faba2adbb3f28832b92de56185cd71dc81..2227ca08ec81b7660e32da48c421adaa6b0400f3 100644
--- a/vis/triangulated_thruster_model.py
+++ b/vis/triangulated_thruster_model.py
@@ -1,21 +1,30 @@
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
diff --git a/vis/xyz_utils.py b/vis/xyz_utils.py
index 6c2e9f18e583538b2745b022f30f9fb071369320..cdc9b9718cc9e5ce0755bf92b9bf5ce910fed5f6 100644
--- a/vis/xyz_utils.py
+++ b/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]],