#!/usr/bin/env python3 """ rv-to-elements.py ================= The inverse map, one step at a time: (r, v) -> (a, e, i, Omega, omega, f). Part III of the deck is an algorithm, and an algorithm read off a table is a list of formulae. Read off the *picture* it is something else: each line of the table puts one more object into the scene, and the six elements are what is left when the scene is complete. 1-2 the state r, v at one instant, and v_r = (r.v)/r 3 h = r x v fixes the orbital PLANE 4 i angle from Z to h -- no quadrant test needed 5 n = K x h fixes the NODE LINE, where the planes meet 6 Omega angle from X to n, in the reference plane 7 e = (v x h)/mu-rhat fixes the APSE LINE, and with p = h^2/mu the conic 8 omega angle from n to e, in the orbital plane 9 f angle from e to r, in the orbital plane 10 a from the energy alone: a = (2/r - v^2/mu)^-1 Nothing is drawn before the step that earns it. The orbit itself arrives only at step 7, because until e is known the plane holds a conic of unknown shape. The three quadrant tests -- sign(n_J), sign(e_K), sign(r.v) -- are shown as they fire. The default state is Example 4 of the deck, chosen because all three of them fire at once: the raw arccos is wrong for every one of the three angles that carries a test. Usage ----- python codes/rv-to-elements.py # play in a window python codes/rv-to-elements.py --verify # the algorithm, in text python codes/rv-to-elements.py --png img/fig-rv-to-elements.png python codes/rv-to-elements.py --storyboard img/fig-rv-to-elements-steps.png python codes/rv-to-elements.py --save img/rv-to-elements.gif python codes/rv-to-elements.py --example 1 --save img/ex1.gif python codes/rv-to-elements.py --state -6045 -3490 2500 -3.457 6.618 2.533 Units are km and km/s throughout, mu = 398600 km^3/s^2 (Earth). """ from __future__ import annotations import argparse import os import numpy as np MU_EARTH = 398600.0 # -- palette (shared with the other figures in this deck) --------------------- CREAM, INK, DIM = "#fcfbf7", "#2f2a24", "#8b857c" MAROON, BLUE, GREEN, AMBER, PURPLE = ("#6b1f1f", "#2f5ea8", "#2d6a2d", "#b45309", "#7e22ce") TEAL = "#0f766e" # velocity: the one vector with no element of its own AXIS = "#6b655c" # the reference frame is backdrop, not an element RED = "#b91c1c" # a quadrant test that fires # States worked in Part V of the deck. Example 4 is the default: it is the one # case where all three quadrant tests fire, so the picture and the algorithm # disagree with a bare arccos three times over. EXAMPLES = { "1": (np.array([-6045.0, -3490.0, 2500.0]), np.array([-3.457, 6.618, 2.533]), "Example 1 -- a retrograde orbit; only $f$ escapes a flipped quadrant"), "4": (np.array([-1222.492, 9623.627, 6500.568]), np.array([-3.297635, 0.034197, -4.916837]), "Example 4 -- all three quadrant tests fire"), } # -- the algorithm itself ----------------------------------------------------- def decompose(r_vec, v_vec, mu=MU_EARTH): """Steps 1-10 of Part III, keeping every intermediate the picture needs. Angles come back in radians. Both forms of each angle are kept: `*_acos` is what a bare arccos returns, `Om`/`om`/`f` is the value after the quadrant test, so the animation can show the two side by side. """ r_vec = np.asarray(r_vec, dtype=float) v_vec = np.asarray(v_vec, dtype=float) r = np.linalg.norm(r_vec) # step 1 v = np.linalg.norm(v_vec) rdotv = r_vec @ v_vec # step 2 v_r = rdotv / r h_vec = np.cross(r_vec, v_vec) # step 3 h = np.linalg.norm(h_vec) inc = np.arccos(np.clip(h_vec[2] / h, -1.0, 1.0)) # step 4 n_vec = np.cross([0.0, 0.0, 1.0], h_vec) # step 5 n = np.linalg.norm(n_vec) # step 6: arccos first, then the sign of n_J decides the half-plane Om_acos = np.arccos(np.clip(n_vec[0] / n, -1.0, 1.0)) if n > 0 else 0.0 Om = (2 * np.pi - Om_acos) if n_vec[1] < 0.0 else Om_acos e_vec = np.cross(v_vec, h_vec) / mu - r_vec / r # step 7 e = np.linalg.norm(e_vec) om_acos = (np.arccos(np.clip((n_vec @ e_vec) / (n * e), -1.0, 1.0)) if n > 0 and e > 0 else 0.0) # step 8 om = (2 * np.pi - om_acos) if e_vec[2] < 0.0 else om_acos f_acos = np.arccos(np.clip((e_vec @ r_vec) / (e * r), -1.0, 1.0)) f = (2 * np.pi - f_acos) if rdotv < 0.0 else f_acos # step 9 a = 1.0 / (2.0 / r - v**2 / mu) # step 10 p = h**2 / mu return dict(mu=mu, r_vec=r_vec, v_vec=v_vec, r=r, v=v, rdotv=rdotv, v_r=v_r, h_vec=h_vec, h=h, n_vec=n_vec, n=n, e_vec=e_vec, e=e, inc=inc, Om=Om, om=om, f=f, Om_acos=Om_acos, om_acos=om_acos, f_acos=f_acos, a=a, p=p, r_p=p / (1.0 + e), r_a=(p / (1.0 - e) if e < 1.0 else np.inf), T=(2 * np.pi * np.sqrt(a**3 / mu) if a > 0 else np.inf)) def state_from_elements(a, e, inc, Om, om, f, mu=MU_EARTH): """The forward map of Part IV -- used only to check the round trip.""" p = a * (1.0 - e**2) r = p / (1.0 + e * np.cos(f)) r_pqw = r * np.array([np.cos(f), np.sin(f), 0.0]) v_pqw = np.sqrt(mu / p) * np.array([-np.sin(f), e + np.cos(f), 0.0]) cO, sO, ci, si, cw, sw = (np.cos(Om), np.sin(Om), np.cos(inc), np.sin(inc), np.cos(om), np.sin(om)) Q = np.array([[cO * cw - sO * sw * ci, -cO * sw - sO * cw * ci, sO * si], [sO * cw + cO * sw * ci, -sO * sw + cO * cw * ci, -cO * si], [sw * si, cw * si, ci]]) return Q @ r_pqw, Q @ v_pqw # -- geometry helpers --------------------------------------------------------- def rodrigues(axis, ang): """Rotation matrix about a unit `axis` by `ang` (active rotation).""" k = np.asarray(axis, dtype=float) k = k / np.linalg.norm(k) K = np.array([[0.0, -k[2], k[1]], [k[2], 0.0, -k[0]], [-k[1], k[0], 0.0]]) return np.eye(3) + np.sin(ang) * K + (1.0 - np.cos(ang)) * (K @ K) def _basis(nhat, prefer=None): """An orthonormal pair spanning the plane through the origin normal nhat.""" nhat = np.asarray(nhat, dtype=float) u = np.array([1.0, 0.0, 0.0]) if prefer is None else np.asarray(prefer, float) if abs(u @ nhat) > 0.9: u = np.array([0.0, 1.0, 0.0]) u = u - (u @ nhat) * nhat u /= np.linalg.norm(u) return u, np.cross(nhat, u) def _square(nhat, half, prefer=None): """Corners of a square patch of the plane through the origin, normal nhat.""" u, w = _basis(nhat, prefer) return np.array([half * (su * u + sw * w) for su, sw in ((1, 1), (-1, 1), (-1, -1), (1, -1))]) def _grid(nhat, half, n=4, prefer=None): """Line segments of a square grid ruled on that plane (for depth cues).""" u, w = _basis(nhat, prefer) segs = [] for t in np.linspace(-half, half, n + 1): segs.append(np.array([t * u - half * w, t * u + half * w])) segs.append(np.array([-half * u + t * w, half * u + t * w])) return segs def _arc(v_from, axis, ang, radius, n=90): """Points of an arc: `v_from` swept about `axis` through `ang`, scaled.""" v = np.asarray(v_from, dtype=float) v = v / np.linalg.norm(v) ts = np.linspace(0.0, ang, max(int(n), 2)) return np.array([radius * (rodrigues(axis, t) @ v) for t in ts]) def _sector(v_from, axis, ang, radius, n=90): """The arc closed back on the vertex: the piece of plane the angle lives in.""" return np.vstack([np.zeros(3), _arc(v_from, axis, ang, radius, n=n)]) def ease(t): """Cosine ease in/out on [0, 1] -- no jerk at the start or end of a step.""" return 0.5 * (1.0 - np.cos(np.pi * np.clip(t, 0.0, 1.0))) # -- the steps ---------------------------------------------------------------- # key, the caption carried while that step is being drawn STEPS = [ ("rv", r"**1-2.** The state: $\mathbf{r}$, $\mathbf{v}$ at one instant, " r"and $v_r=(\mathbf{r}\cdot\mathbf{v})/r$"), ("h", r"**3.** $\mathbf{h}=\mathbf{r}\times\mathbf{v}$ -- normal to the " r"orbital plane, so it fixes the plane"), ("i", r"**4.** $i=\arccos(h_K/h)$ -- the only angle that needs no " r"quadrant test"), ("n", r"**5.** $\mathbf{n}=\hat{\mathbf{K}}\times\mathbf{h}$ -- the node " r"line, where the two planes meet"), ("Om", r"**6.** $\Omega$, from $\hat{\mathbf{I}}$ to $\mathbf{n}$ in the " r"reference plane -- test: $\mathrm{sign}(n_J)$"), ("e", r"**7.** $\mathbf{e}$ -- the apse line; with $p=h^2/\mu$ the conic " r"is now fixed"), ("om", r"**8.** $\omega$, from $\mathbf{n}$ to $\mathbf{e}$ in the orbital " r"plane -- test: $\mathrm{sign}(e_K)$"), ("f", r"**9.** $f$, from $\mathbf{e}$ to $\mathbf{r}$ -- test: " r"$\mathrm{sign}(\mathbf{r}\cdot\mathbf{v})$"), ("a", r"**10.** $a=(2/r-v^2/\mu)^{-1}$ -- from the energy alone. " r"The state is spent."), ] STEP_KEYS = [k for k, _ in STEPS] def _reveal(key, step, t): """How far along the reveal of `key` we are: 0 before, t during, 1 after.""" j = STEP_KEYS.index(key) if j < step: return 1.0 if j > step: return 0.0 return float(np.clip(t, 0.0, 1.0)) def draw_scene(ax, S, step, t=1.0, zoom=1.68): """Redraw the 3D scene as it stands `t` of the way through `step`.""" from mpl_toolkits.mplot3d.art3d import Poly3DCollection elev, azim = ax.elev, ax.azim ax.clear() e, inc, Om, om, f = S["e"], S["inc"], S["Om"], S["om"], S["f"] r_vec, v_vec, h_vec, n_vec, e_vec = (S["r_vec"], S["v_vec"], S["h_vec"], S["n_vec"], S["e_vec"]) zhat = np.array([0.0, 0.0, 1.0]) hhat = h_vec / S["h"] nhat = n_vec / S["n"] if S["n"] > 0 else np.array([1.0, 0.0, 0.0]) ehat = e_vec / e if e > 0 else nhat # Scene half-size: fixed by the *apoapsis*, so the frame never rescales as # the steps reveal more of the orbit. L = (S["r_a"] if np.isfinite(S["r_a"]) else 2.4 * S["r_p"]) * 1.12 L = max(L, 1.25 * S["r"]) # vectors carrying physical units of their own are drawn to a length that # reads well against L; only their *directions* are meant to be measured v_scale = 0.42 * L / S["v"] h_scale = 0.78 * L / S["h"] n_scale = 0.88 * L / S["n"] if S["n"] > 0 else 0.0 def q(vec, color, lw=2.0, alpha=1.0, ratio=0.10, frac=1.0): if alpha <= 0.02 or frac <= 0.02: return ax.quiver(0, 0, 0, *(frac * np.asarray(vec)), color=color, lw=lw, alpha=alpha, arrow_length_ratio=ratio / max(frac, 0.2)) # -- reference plane, its grid, and the frame axes ------------------------ sq = L * np.array([[1, 1, 0], [-1, 1, 0], [-1, -1, 0], [1, -1, 0]], float) ax.add_collection3d(Poly3DCollection([sq], facecolor=DIM, alpha=0.09, edgecolor=DIM, linewidths=0.8)) for seg in _grid(zhat, L, n=6): ax.plot(*seg.T, color=DIM, lw=0.45, alpha=0.35) ax.text(-0.50 * L, -0.95 * L, 0.0, "reference plane (equator)", color=DIM, fontsize=9.5, ha="center") for vec, lab, ln, ratio in (([1.0, 0, 0], r"$\hat{\mathbf{I}}$", 1.02, 0.10), ([0, 1.0, 0], r"$\hat{\mathbf{J}}$", 0.95, 0.09), ([0, 0, 1.0], r"$\hat{\mathbf{K}}$", 0.80, 0.11)): ax.quiver(0, 0, 0, *(ln * L * np.asarray(vec)), color=AXIS, lw=1.2, arrow_length_ratio=ratio) ax.text(*((ln + 0.08) * L * np.asarray(vec)), lab, color=AXIS, fontsize=11, ha="center") # -- central body --------------------------------------------------------- ax.plot([0], [0], [0], "o", color=AMBER, ms=11, zorder=8) # -- step 1-2: the state -------------------------------------------------- g = _reveal("rv", step, t) if g > 0: gr = np.clip(2.0 * g, 0.0, 1.0) # r first, then v from its tip gv = np.clip(2.0 * g - 1.0, 0.0, 1.0) q(r_vec, MAROON, lw=2.4, frac=gr) ax.plot(*(gr * r_vec).reshape(3, 1), "o", color=MAROON, ms=9, zorder=9) ax.text(*(0.80 * gr * r_vec + 0.085 * L * hhat), r"$\mathbf{r}$", color=MAROON, fontsize=14, ha="center", va="center") # the dotted drop to the reference plane is the only cue for the sign # of z, and z is what every quadrant test is ultimately about ax.plot(*np.vstack([gr * r_vec, [gr * r_vec[0], gr * r_vec[1], 0.0]]).T, color=MAROON, lw=0.9, ls=":", alpha=0.65) if gv > 0: tip = r_vec + gv * v_scale * v_vec ax.quiver(*r_vec, *(gv * v_scale * v_vec), color=TEAL, lw=2.4, arrow_length_ratio=0.12 / max(gv, 0.25)) ax.text(*(tip + 0.07 * L * (v_vec / S["v"])), r"$\mathbf{v}$", color=TEAL, fontsize=14, ha="center", va="center") # -- step 3: h, and the plane it defines ---------------------------------- g = _reveal("h", step, t) if g > 0: # the parallelogram spanned by r and v: its area is |r x v| = h, which # is the one picture that makes h a *number* and not just a direction par = np.array([np.zeros(3), r_vec, r_vec + v_scale * v_vec, v_scale * v_vec]) ax.add_collection3d(Poly3DCollection([par], facecolor=BLUE, alpha=0.10 * g, edgecolor=BLUE, linewidths=0.7)) quad = _square(hhat, 0.86 * L, prefer=nhat) ax.add_collection3d(Poly3DCollection([quad], facecolor=BLUE, alpha=0.13 * g, edgecolor=BLUE, linewidths=0.9)) for seg in _grid(hhat, 0.86 * L, n=2, prefer=nhat): ax.plot(*seg.T, color=BLUE, lw=0.4, alpha=0.30 * g) ax.text(*(quad[1] * 0.99), "orbital plane", color=BLUE, fontsize=9.5, ha="center", alpha=g) q(h_scale * h_vec, BLUE, lw=2.4, alpha=1.0, ratio=0.11, frac=g) ax.text(*(g * h_scale * h_vec + 0.06 * L * hhat), r"$\mathbf{h}$", color=BLUE, fontsize=14, alpha=g, ha="center") # -- step 4: the inclination ---------------------------------------------- g = _reveal("i", step, t) if g > 0: _angle(ax, zhat, nhat, g * inc, 0.42 * L, BLUE, r"$i$", alpha=g, lab_r=1.24) # -- step 5: the node line ------------------------------------------------ g = _reveal("n", step, t) if g > 0: ax.plot(*np.vstack([-g * L * nhat, g * L * nhat]).T, color=GREEN, lw=1.8, ls=(0, (7, 4)), alpha=g) q(n_scale * n_vec, GREEN, lw=2.2, alpha=g, ratio=0.10, frac=g) ax.text(*(1.02 * g * n_scale * n_vec - 0.05 * L * zhat), r"$\mathbf{n}$ (ascending node)", color=GREEN, fontsize=10.5, ha="center", va="top", alpha=g) # -- step 6: Omega, in the reference plane -------------------------------- g = _reveal("Om", step, t) if g > 0: # below the reference plane: nothing else in the scene lives there _angle(ax, np.array([1.0, 0.0, 0.0]), zhat, g * Om, 0.54 * L, AMBER, r"$\Omega$", alpha=g, fill=True, lab_r=1.12, lab_off=-0.34) # -- step 7: the apse line, and with it the conic -------------------------- g = _reveal("e", step, t) if g > 0: q(0.85 * L * ehat, PURPLE, lw=2.2, alpha=g, ratio=0.10, frac=g) ax.text(*(0.93 * L * ehat), r"$\mathbf{e}$", color=PURPLE, fontsize=14, alpha=g, ha="center") # the conic, traced from periapsis outwards in both directions so that # it grows the way the apse line does -- symmetric about e, as it is _orbit(ax, S, hhat, ehat, span=g, L=L) # -- step 8: omega, in the orbital plane ---------------------------------- g = _reveal("om", step, t) if g > 0: _angle(ax, nhat, hhat, g * om, 0.26 * L, PURPLE, r"$\omega$", alpha=g, fill=True, lab_r=1.26, lab_off=+0.40) # -- step 9: the true anomaly --------------------------------------------- g = _reveal("f", step, t) if g > 0: _angle(ax, ehat, hhat, g * f, 0.13 * L, MAROON, r"$f$", alpha=g, lab_r=1.35, lab_off=-0.55) # v_r as the projection of v on rhat: the sign test, drawn rhat = r_vec / S["r"] ax.quiver(*r_vec, *(g * v_scale * S["v_r"] * rhat), color=RED, lw=2.0, arrow_length_ratio=0.16 / max(g, 0.3), alpha=g) qhat = np.cross(hhat, rhat) # in-plane, across from the arcs ax.text(*(r_vec + 0.5 * g * v_scale * S["v_r"] * rhat + 0.085 * L * qhat), r"$v_r$", color=RED, fontsize=12, alpha=g, ha="center", va="center") # -- step 10: the apsides make the size visible ---------------------------- g = _reveal("a", step, t) if g > 0: r_p_vec = S["r_p"] * ehat ax.plot(*r_p_vec.reshape(3, 1), "o", color=PURPLE, ms=7, zorder=9) if np.isfinite(S["r_a"]): r_a_vec = -S["r_a"] * ehat ax.plot(*np.vstack([r_p_vec, r_p_vec + g * (r_a_vec - r_p_vec)]).T, color=INK, lw=1.4, ls=(0, (6, 3)), alpha=0.85 * g) ax.plot(*r_a_vec.reshape(3, 1), "o", mfc=CREAM, mec=INK, mew=1.4, ms=6, zorder=9, alpha=g) ax.text(*(0.72 * r_a_vec + 0.28 * r_p_vec + 0.07 * L * hhat), r"$2a$", color=INK, fontsize=13, alpha=g, ha="center", va="center") # -- cosmetics ------------------------------------------------------------- zlo, zhi = -0.92 * L, 0.96 * L ax.set_xlim(-L, L) ax.set_ylim(-L, L) ax.set_zlim(zlo, zhi) ax.set_box_aspect((1.0, 1.0, (zhi - zlo) / (2.0 * L)), zoom=zoom) ax.set_facecolor(CREAM) ax.axis("off") ax.view_init(elev=elev, azim=azim) def _angle(ax, v_from, axis, ang, radius, color, label, alpha=1.0, fill=False, lab_r=1.20, lab_off=0.0): """Draw an angle where it is actually measured: an arc about `axis`. `lab_off` lifts the label off the arc along `axis`, i.e. along the normal of the plane the angle is swept in. Three of the four angles have arcs that are nearly concentric in projection, so a purely radial offset cannot separate their labels; an offset along the axis can, and it says something true while doing it -- each label floats over its own plane. """ from mpl_toolkits.mplot3d.art3d import Poly3DCollection if abs(ang) < 1e-4: return n = max(8, int(90 * abs(ang) / (2 * np.pi)) + 8) pts = _arc(v_from, axis, ang, radius, n=n) if fill: # the filled sector shows which plane the angle is swept in -- the # whole difficulty of the 3-1-3 sequence is that the three planes differ ax.add_collection3d(Poly3DCollection( [_sector(v_from, axis, ang, radius, n=n)], facecolor=color, alpha=0.16 * alpha, edgecolor="none")) for end in (pts[0], pts[-1]): ax.plot(*np.vstack([np.zeros(3), end]).T, color=color, lw=0.9, alpha=0.75 * alpha) ax.plot(*pts.T, color=color, lw=2.0, alpha=alpha) nhat_ax = np.asarray(axis, dtype=float) nhat_ax = nhat_ax / np.linalg.norm(nhat_ax) ax.text(*(pts[len(pts) // 2] * lab_r + lab_off * radius * nhat_ax), label, color=color, fontsize=14, ha="center", va="center", alpha=alpha) def _orbit(ax, S, hhat, ehat, span=1.0, L=None): """The conic p/(1+e cos f), traced symmetrically outwards from periapsis.""" e, p = S["e"], S["p"] f_max = np.pi if e < 1.0 else 0.985 * np.arccos(-1.0 / e) fs = np.linspace(-span * f_max, span * f_max, 721) rs = p / (1.0 + e * np.cos(fs)) qhat = np.cross(hhat, ehat) pts = rs[:, None] * (np.cos(fs)[:, None] * ehat + np.sin(fs)[:, None] * qhat) if L is not None: # keep a hyperbola inside the frame pts = pts[np.linalg.norm(pts, axis=1) <= 1.45 * L] if len(pts) < 2: return ax.plot(*pts.T, color=INK, lw=2.4) shadow = pts.copy() shadow[:, 2] = 0.0 ax.plot(*shadow.T, color=INK, lw=0.9, ls=(0, (3, 4)), alpha=0.42) # -- the running tally -------------------------------------------------------- def _vec(v, p=3): return "(" + ",".join(f"{x:9.{p}f}" for x in v) + ")" def panel_lines(S, step, t=1.0): """The lines of the side panel, as (text, colour, size, weight) tuples. The panel is the point of the whole animation: it is the state on the left being spent, one line at a time, into the six numbers on the right. """ out = [("GIVEN", DIM, 10.5, "bold")] out += [(f"r = {_vec(S['r_vec'])} km", MAROON, 9.5, "normal"), (f"v = {_vec(S['v_vec'])} km/s", TEAL, 9.5, "normal"), (f"|r| = {S['r']:.3f} km |v| = {S['v']:.6f} km/s", INK, 9.5, "normal")] if _reveal("rv", step, t) >= 1.0 or step > 0: sgn = "inbound" if S["rdotv"] < 0 else "outbound" out += [(f"r.v = {S['rdotv']:.1f} -> v_r = {S['v_r']:+.6f} km/s " f"({sgn})", RED if S["rdotv"] < 0 else INK, 9.5, "normal")] out += [("", INK, 6, "normal"), ("BUILT FROM IT", DIM, 10.5, "bold")] if _reveal("h", step, t) > 0: out += [(f"h = r x v = {_vec(S['h_vec'], 1)}", BLUE, 9.5, "normal"), (f" h = {S['h']:.2f} km^2/s", BLUE, 9.5, "normal")] if _reveal("n", step, t) > 0: out += [(f"n = K x h = {_vec(S['n_vec'], 1)}", GREEN, 9.5, "normal"), (f" n = {S['n']:.2f}", GREEN, 9.5, "normal")] if _reveal("e", step, t) > 0: out += [(f"e = {_vec(S['e_vec'], 6)}", PURPLE, 9.5, "normal")] out += [("", INK, 6, "normal"), ("THE SIX ELEMENTS", DIM, 10.5, "bold")] def ang_row(key, sym, val, acos, test, sgn_txt, color): g = _reveal(key, step, t) if g <= 0: return [] flipped = abs(val - acos) > 1e-9 rows = [(f"{sym} = {np.rad2deg(val):11.3f} deg", color, 13, "bold")] if key != "i": if flipped: rows += [(f" arccos gave {np.rad2deg(acos):.3f}; " f"{test} < 0", RED, 8.5, "normal"), (f" -> 360 - {np.rad2deg(acos):.3f}", RED, 8.5, "normal")] else: rows += [(f" {test} > 0, arccos stands ({sgn_txt})", DIM, 8.5, "normal")] else: rows += [(" no quadrant test: i in [0, 180]", DIM, 8.5, "normal")] return rows if _reveal("a", step, t) > 0: out += [(f"a = {S['a']:11.3f} km", INK, 13, "bold")] if _reveal("e", step, t) > 0: out += [(f"e = {S['e']:11.6f}", PURPLE, 13, "bold")] out += ang_row("i", "i", S["inc"], S["inc"], "", "", BLUE) out += ang_row("Om", "Ω", S["Om"], S["Om_acos"], "n_J", "node in J > 0", AMBER) out += ang_row("om", "ω", S["om"], S["om_acos"], "e_K", "periapsis above", PURPLE) out += ang_row("f", "f", S["f"], S["f_acos"], "r.v", "outbound", MAROON) if _reveal("a", step, t) >= 1.0: out += [("", INK, 6, "normal")] ra = "inf" if not np.isfinite(S["r_a"]) else f"{S['r_a']:.1f}" out += [(f"r_p = {S['r_p']:.1f} km r_a = {ra} km", DIM, 9.5, "normal"), (f"p = {S['p']:.1f} km" + ("" if not np.isfinite(S["T"]) else f" T = {S['T']:.1f} s"), DIM, 9.5, "normal")] return out def draw_panel(ax, S, step, t=1.0): ax.clear() ax.set_xlim(0, 1) ax.set_ylim(0, 1) ax.axis("off") ax.set_facecolor(CREAM) y = 0.985 for text, color, size, weight in panel_lines(S, step, t): if text: ax.text(0.0, y, text, color=color, fontsize=size, weight=weight, family="DejaVu Sans Mono", ha="left", va="top", transform=ax.transAxes) y -= (size + 5.5) / 478.0 # -- frame schedule ----------------------------------------------------------- def schedule(draw=26, hold=16): """(step, t) for every frame: `draw` frames to build it, `hold` to read it.""" frames = [] for j in range(len(STEPS)): for k in range(draw): frames.append((j, ease((k + 1) / draw))) frames += [(j, 1.0)] * hold frames += [(len(STEPS) - 1, 1.0)] * (2 * hold) # a beat on the answer return frames def _figure(figsize=(12.6, 6.8), elev=24, azim=22): import matplotlib.pyplot as plt plt.rcParams.update({"figure.facecolor": CREAM, "savefig.facecolor": CREAM, "axes.facecolor": CREAM, "text.color": INK, "font.size": 11}) fig = plt.figure(figsize=figsize) ax3 = fig.add_axes([-0.020, 0.050, 0.695, 0.905], projection="3d") ax3.view_init(elev=elev, azim=azim) axp = fig.add_axes([0.650, 0.050, 0.345, 0.905]) return fig, ax3, axp def render(path, S, title, fps=20, dpi=100, draw=26, hold=16, elev=24, azim=22): import matplotlib matplotlib.use("Agg") from matplotlib import animation frames = schedule(draw=draw, hold=hold) fig, ax3, axp = _figure(elev=elev, azim=azim) cap = fig.text(0.5, 0.022, "", ha="center", fontsize=13.5, color=INK) fig.text(0.5, 0.972, title, ha="center", fontsize=12, color=DIM) def frame(k): step, t = frames[k] draw_scene(ax3, S, step, t) draw_panel(axp, S, step, t) cap.set_text(_caption(step)) return () ani = animation.FuncAnimation(fig, frame, frames=len(frames), blit=False) os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) writer = "ffmpeg" if path.endswith(".mp4") else "pillow" print(f"writing {path} ({len(frames)} frames) ...") ani.save(path, writer=writer, fps=fps, dpi=dpi, savefig_kwargs=dict(facecolor=CREAM)) print("done:", path, f"{os.path.getsize(path) / 1e6:.2f} MB") def _caption(step): """Captions carry mathtext, so the bold markers have to go.""" return STEPS[step][1].replace("**", "") def save_png(path, S, title, step=None, dpi=160, elev=24, azim=22): import matplotlib matplotlib.use("Agg") step = len(STEPS) - 1 if step is None else step fig, ax3, axp = _figure(elev=elev, azim=azim) fig.text(0.5, 0.972, title, ha="center", fontsize=12, color=DIM) fig.text(0.5, 0.022, _caption(step), ha="center", fontsize=13.5, color=INK) draw_scene(ax3, S, step, 1.0) draw_panel(axp, S, step, 1.0) os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) fig.savefig(path, dpi=dpi, facecolor=CREAM) print("wrote", path) def storyboard(path, S, title, dpi=150, elev=24, azim=22, cols=3): """All nine steps on one sheet -- the animation, for a printed handout.""" import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt plt.rcParams.update({"figure.facecolor": CREAM, "savefig.facecolor": CREAM, "axes.facecolor": CREAM, "text.color": INK}) rows = int(np.ceil(len(STEPS) / cols)) fig = plt.figure(figsize=(5.0 * cols, 3.45 * rows)) fig.text(0.5, 0.996, title, ha="center", va="top", fontsize=13, color=DIM) for j in range(len(STEPS)): ax = fig.add_subplot(rows, cols, j + 1, projection="3d") ax.view_init(elev=elev, azim=azim) draw_scene(ax, S, j, 1.0, zoom=1.46) ax.set_title(_caption(j), color=INK, fontsize=9.0, pad=-14, wrap=True) fig.subplots_adjust(left=0.005, right=0.995, top=0.955, bottom=0.005, wspace=0.01, hspace=0.05) os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) fig.savefig(path, dpi=dpi, facecolor=CREAM, bbox_inches="tight") print("wrote", path) def interactive(S, title, elev=24, azim=22): """Step through it by hand: left/right arrows, or the two buttons.""" import matplotlib.pyplot as plt from matplotlib.widgets import Button fig, ax3, axp = _figure(elev=elev, azim=azim) fig.text(0.5, 0.972, title, ha="center", fontsize=12, color=DIM) cap = fig.text(0.5, 0.022, "", ha="center", fontsize=13.5, color=INK) state = {"step": 0} def show(): draw_scene(ax3, S, state["step"], 1.0) draw_panel(axp, S, state["step"], 1.0) cap.set_text(_caption(state["step"])) fig.canvas.draw_idle() def go(d): def _cb(_=None): state["step"] = int(np.clip(state["step"] + d, 0, len(STEPS) - 1)) show() return _cb bprev = Button(fig.add_axes([0.30, 0.955, 0.06, 0.038]), "< back", color="#eee9dd", hovercolor="#e2dbc9") bnext = Button(fig.add_axes([0.37, 0.955, 0.06, 0.038]), "next >", color="#eee9dd", hovercolor="#e2dbc9") bprev.on_clicked(go(-1)) bnext.on_clicked(go(+1)) fig.canvas.mpl_connect("key_press_event", lambda ev: ( go(+1)() if ev.key in ("right", " ") else go(-1)() if ev.key == "left" else None)) show() plt.show() # -- verification ------------------------------------------------------------- def verify(S): """The same nine steps in text, then the round trip back through Part IV.""" deg = np.rad2deg print("step-by-step decomposition (r, v) -> (a, e, i, Omega, omega, f)\n") print(f" 1 r = {_vec(S['r_vec'])} km, |r| = {S['r']:.6f} km") print(f" v = {_vec(S['v_vec'], 6)} km/s, |v| = {S['v']:.6f} km/s") print(f" 2 r.v = {S['rdotv']:.4f} -> v_r = {S['v_r']:+.6f} km/s " f"({'inbound' if S['rdotv'] < 0 else 'outbound'})") print(f" 3 h = {_vec(S['h_vec'], 4)}, h = {S['h']:.4f} km^2/s") print(f" 4 i = arccos({S['h_vec'][2] / S['h']:+.6f}) = " f"{deg(S['inc']):.4f} deg [no test]") print(f" 5 n = {_vec(S['n_vec'], 4)}, n = {S['n']:.4f}") print(f" 6 Om = {deg(S['Om']):.4f} deg arccos gave " f"{deg(S['Om_acos']):.4f}, n_J = {S['n_vec'][1]:+.2f} " f"[{'FLIP' if abs(S['Om'] - S['Om_acos']) > 1e-9 else 'stands'}]") print(f" 7 e = {_vec(S['e_vec'], 6)}, e = {S['e']:.6f}") print(f" 8 om = {deg(S['om']):.4f} deg arccos gave " f"{deg(S['om_acos']):.4f}, e_K = {S['e_vec'][2]:+.6f} " f"[{'FLIP' if abs(S['om'] - S['om_acos']) > 1e-9 else 'stands'}]") print(f" 9 f = {deg(S['f']):.4f} deg arccos gave " f"{deg(S['f_acos']):.4f}, r.v = {S['rdotv']:+.2f} " f"[{'FLIP' if abs(S['f'] - S['f_acos']) > 1e-9 else 'stands'}]") print(f" 10 a = {S['a']:.4f} km, p = {S['p']:.4f} km, " f"r_p = {S['r_p']:.4f} km, r_a = {S['r_a']:.4f} km") r2, v2 = state_from_elements(S["a"], S["e"], S["inc"], S["Om"], S["om"], S["f"], S["mu"]) print("\nround trip through Part IV:") print(f" r = {_vec(r2, 6)} km |dr| = " f"{np.linalg.norm(r2 - S['r_vec']):.3e} km") print(f" v = {_vec(v2, 6)} km/s |dv| = " f"{np.linalg.norm(v2 - S['v_vec']):.3e} km/s") # what one dropped quadrant test costs -- Example 3 of the deck r_bad, _ = state_from_elements(S["a"], S["e"], S["inc"], S["Om_acos"], S["om"], S["f"], S["mu"]) print(f"\nwith Omega left at the raw arccos ({deg(S['Om_acos']):.3f} deg):") print(f" r = {_vec(r_bad, 2)} km |dr| = " f"{np.linalg.norm(r_bad - S['r_vec']):.1f} km " f"-- and |r| is still {np.linalg.norm(r_bad):.3f} km") def main(): ap = argparse.ArgumentParser(description=__doc__.split("\n")[3]) ap.add_argument("--example", choices=sorted(EXAMPLES), default="4", help="which worked example of the deck to decompose") ap.add_argument("--state", nargs=6, type=float, metavar=("RX", "RY", "RZ", "VX", "VY", "VZ"), help="an arbitrary state, km and km/s, overriding --example") ap.add_argument("--mu", type=float, default=MU_EARTH) ap.add_argument("--verify", action="store_true", help="the ten steps in text, plus the round trip") ap.add_argument("--png", metavar="PATH", help="save the final frame") ap.add_argument("--storyboard", metavar="PATH", help="save all nine steps on one sheet") ap.add_argument("--save", metavar="PATH", help="render the animation to PATH (.gif or .mp4)") ap.add_argument("--step", type=int, default=None, help="which step --png should show (0-8, default: last)") ap.add_argument("--fps", type=int, default=20) ap.add_argument("--dpi", type=int, default=100) ap.add_argument("--draw", type=int, default=26, help="frames per reveal") ap.add_argument("--hold", type=int, default=16, help="frames per pause") ap.add_argument("--elev", type=float, default=24.0) ap.add_argument("--azim", type=float, default=22.0) args = ap.parse_args() if args.state: r_vec = np.array(args.state[:3]) v_vec = np.array(args.state[3:]) title = (r"$\mathbf{r}=$" + _vec(r_vec, 1) + " km, " + r"$\mathbf{v}=$" + _vec(v_vec, 4) + " km/s") else: r_vec, v_vec, title = EXAMPLES[args.example] S = decompose(r_vec, v_vec, mu=args.mu) if args.verify: verify(S) if args.png: save_png(args.png, S, title, step=args.step, elev=args.elev, azim=args.azim) if args.storyboard: storyboard(args.storyboard, S, title, elev=args.elev, azim=args.azim) if args.save: render(args.save, S, title, fps=args.fps, dpi=args.dpi, draw=args.draw, hold=args.hold, elev=args.elev, azim=args.azim) if not (args.verify or args.png or args.storyboard or args.save): interactive(S, title, elev=args.elev, azim=args.azim) if __name__ == "__main__": main()