1def plot_planet(plnt, t0=0, tf=None, N=60, units=1.0, color='k', alpha=1.0, s=40, legend=(False, False), axes=None): 2 """ 3 ax = plot_planet(plnt, t0=0, tf=None, N=60, units=1.0, color='k', alpha=1.0, s=40, legend=(False, False), axes=None): 4 5 - axes: 3D axis object created using fig.gca(projection='3d') 6 - plnt: pykep.planet object we want to plot 7 - t0: a pykep.epoch or float (mjd2000) indicating the first date we want to plot the planet position 8 - tf: a pykep.epoch or float (mjd2000) indicating the final date we want to plot the planet position. 9 if None this is computed automatically from the orbital period (prone to error for non periodic orbits) 10 - units: the length unit to be used in the plot 11 - color: color to use to plot the orbit (passed to matplotlib) 12 - s: planet size (passed to matplotlib) 13 - legend 2-D tuple of bool or string: The first element activates the planet scatter plot, 14 the second to the actual orbit. If a bool value is used, then an automated legend label is generated (if True), if a string is used, the string is the legend. Its also possible but deprecated to use a single bool value. In which case that value is used for both the tuple components. 15 16 Plots the planet position and its orbit. 17 18 Example:: 19 20 import pykep as pk 21 import matplotlib.pyplot as plt 22 23 fig = plt.figure() 24 ax = fig.gca(projection = '3d') 25 pl = pk.planet.jpl_lp('earth') 26 t_plot = pk.epoch(219) 27 ax = pk.orbit_plots.plot_planet(pl, ax = ax, color='b') 28 """ 29 from pykep import MU_SUN, SEC2DAY, epoch, AU, RAD2DEG 30 from pykep.planet import keplerian 31 from math import pi, sqrt 32 import numpy as np 33 import matplotlib.pylab as plt 34 from mpl_toolkits.mplot3d import Axes3D 35 36 if axes is None: 37 fig = plt.figure() 38 ax = fig.gca(projection='3d') 39 else: 40 ax = axes 41 42 if type(t0) is not epoch: 43 t0 = epoch(t0) 44 45 # This is to make the tuple API compatible with the old API 46 if type(legend) is bool: 47 legend = (legend, legend) 48 49 if tf is None: 50 # orbit period at epoch 51 T = plnt.compute_period(t0) * SEC2DAY 52 else: 53 if type(tf) is not epoch: 54 tf = epoch(tf) 55 T = (tf.mjd2000 - t0.mjd2000) 56 if T < 0: 57 raise ValueError("tf should be after t0 when plotting an orbit") 58 59 # points where the orbit will be plotted 60 when = np.linspace(0, T, N) 61 62 # Ephemerides Calculation for the given planet 63 x = np.array([0.0] * N) 64 y = np.array([0.0] * N) 65 z = np.array([0.0] * N) 66 67 for i, day in enumerate(when): 68 r, v = plnt.eph(epoch(t0.mjd2000 + day)) 69 x[i] = r[0] / units 70 y[i] = r[1] / units 71 z[i] = r[2] / units 72 73 # Actual plot commands 74 if (legend[0] is True): 75 label1 = plnt.name + " " + t0.__repr__()[0:11] 76 elif (legend[0] is False): 77 label1 = None 78 elif (legend[0] is None): 79 label1 = None 80 else: 81 label1 = legend[0] 82 83 if (legend[1] is True): 84 label2 = plnt.name + " orbit" 85 elif (legend[1] is False): 86 label2 = None 87 elif (legend[1] is None): 88 label2 = None 89 else: 90 label2 = legend[1] 91 92 ax.plot(x, y, z, label=label2, c=color, alpha=alpha) 93 ax.scatter([x[0]], [y[0]], [z[0]], s=s, marker='o', alpha=0.8, c=[color], label=label1) 94 95 if legend[0] or legend[1]: 96 ax.legend() 97 return ax 98 99 100def plot_lambert(l, N=60, sol=0, units=1.0, color='b', legend=False, axes=None, alpha=1.): 101 """ 102 ax = plot_lambert(l, N=60, sol=0, units='pykep.AU', legend='False', axes=None, alpha=1.) 103 104 - axes: 3D axis object created using fig.gca(projection='3d') 105 - l: pykep.lambert_problem object 106 - N: number of points to be plotted along one arc 107 - sol: solution to the Lambert's problem we want to plot (must be in 0..Nmax*2) 108 where Nmax is the maximum number of revolutions for which there exist a solution. 109 - units: the length unit to be used in the plot 110 - color: matplotlib color to use to plot the line 111 - legend: when True it plots also the legend with info on the Lambert's solution chosen 112 113 Plots a particular solution to a Lambert's problem 114 115 Example:: 116 117 import pykep as pk 118 import matplotlib.pyplot as plt 119 120 fig = plt.figure() 121 ax = fig.gca(projection='3d') 122 123 t1 = pk.epoch(0) 124 t2 = pk.epoch(640) 125 dt = (t2.mjd2000 - t1.mjd2000) * pk.DAY2SEC 126 127 pl = pk.planet.jpl_lp('earth') 128 pk.orbit_plots.plot_planet(pl, t0=t1, axes=ax, color='k') 129 rE,vE = pl.eph(t1) 130 131 pl = pk.planet.jpl_lp('mars') 132 pk.orbit_plots.plot_planet(pl, t0=t2, axes=ax, color='r') 133 rM, vM = pl.eph(t2) 134 135 l = lambert_problem(rE,rM,dt,pk.MU_SUN) 136 pk.orbit_plots.plot_lambert(l, ax=ax, color='b') 137 pk.orbit_plots.plot_lambert(l, sol=1, axes=ax, color='g') 138 pk.orbit_plots.plot_lambert(l, sol=2, axes=ax, color='g', legend = True) 139 140 plt.show() 141 """ 142 from pykep import propagate_lagrangian, AU 143 import numpy as np 144 import matplotlib.pylab as plt 145 from mpl_toolkits.mplot3d import Axes3D 146 147 if axes is None: 148 fig = plt.figure() 149 ax = fig.gca(projection='3d') 150 else: 151 ax = axes 152 153 if sol > l.get_Nmax() * 2: 154 raise ValueError("sol must be in 0 .. NMax*2 \n * Nmax is the maximum number of revolutions for which there exist a solution to the Lambert's problem \n * You can compute Nmax calling the get_Nmax() method of the lambert_problem object") 155 156 # We extract the relevant information from the Lambert's problem 157 r = l.get_r1() 158 v = l.get_v1()[sol] 159 T = l.get_tof() 160 mu = l.get_mu() 161 162 # We define the integration time ... 163 dt = T / (N - 1) 164 165 # ... and allocate the cartesian components for r 166 x = np.array([0.0] * N) 167 y = np.array([0.0] * N) 168 z = np.array([0.0] * N) 169 170 # We calculate the spacecraft position at each dt 171 for i in range(N): 172 x[i] = r[0] / units 173 y[i] = r[1] / units 174 z[i] = r[2] / units 175 r, v = propagate_lagrangian(r, v, dt, mu) 176 177 # And we plot 178 if legend: 179 label = 'Lambert solution (' + str((sol + 1) // 2) + ' revs.)' 180 else: 181 label = None 182 ax.plot(x, y, z, c=color, label=label, alpha=alpha) 183 184 if legend: 185 ax.legend() 186 187 return ax 188 189 190def plot_kepler(r0, v0, tof, mu, N=60, units=1, color='b', label=None, axes=None): 191 """ 192 ax = plot_kepler(r0, v0, tof, mu, N=60, units=1, color='b', label=None, axes=None): 193 194 - axes: 3D axis object created using fig.gca(projection='3d') 195 - r0: initial position (cartesian coordinates) 196 - v0: initial velocity (cartesian coordinates) 197 - tof: propagation time 198 - mu: gravitational parameter 199 - N: number of points to be plotted along one arc 200 - units: the length unit to be used in the plot 201 - color: matplotlib color to use to plot the line 202 - label adds a label to the plotted arc. 203 204 Plots the result of a keplerian propagation 205 206 Example:: 207 208 import pykep as pk 209 pi = 3.14 210 pk.orbit_plots.plot_kepler(r0 = [1,0,0], v0 = [0,1,0], tof = pi/3, mu = 1) 211 """ 212 213 from pykep import propagate_lagrangian 214 import matplotlib.pylab as plt 215 from mpl_toolkits.mplot3d import Axes3D 216 from copy import deepcopy 217 218 if axes is None: 219 fig = plt.figure() 220 ax = fig.gca(projection='3d') 221 else: 222 ax = axes 223 224 # We define the integration time ... 225 dt = tof / (N - 1) 226 227 # ... and calculate the cartesian components for r 228 x = [0.0] * N 229 y = [0.0] * N 230 z = [0.0] * N 231 232 # We calculate the spacecraft position at each dt 233 r = deepcopy(r0) 234 v = deepcopy(v0) 235 for i in range(N): 236 x[i] = r[0] / units 237 y[i] = r[1] / units 238 z[i] = r[2] / units 239 r, v = propagate_lagrangian(r, v, dt, mu) 240 241 # And we plot 242 ax.plot(x, y, z, c=color, label=label) 243 return ax 244 245 246def plot_taylor(r0, v0, m0, thrust, tof, mu, veff, N=60, units=1, color='b', legend=False, axes=None): 247 """ 248 ax = plot_taylor(r0, v0, m0, thrust, tof, mu, veff, N=60, units=1, color='b', legend=False, axes=None): 249 250 - axes: 3D axis object created using fig.gca(projection='3d') 251 - r0: initial position (cartesian coordinates) 252 - v0: initial velocity (cartesian coordinates) 253 - m0: initial mass 254 - u: cartesian components for the constant thrust 255 - tof: propagation time 256 - mu: gravitational parameter 257 - veff: the product Isp * g0 258 - N: number of points to be plotted along one arc 259 - units: the length unit to be used in the plot 260 - color: matplotlib color to use to plot the line 261 - legend: when True it plots also the legend 262 263 Plots the result of a taylor propagation of constant thrust 264 265 Example:: 266 267 import pykep as pk 268 import matplotlib.pyplot as plt 269 pi = 3.14 270 271 fig = plt.figure() 272 ax = fig.gca(projection = '3d') 273 pk.orbit_plots.plot_taylor([1,0,0],[0,1,0],100,[1,1,0],40, 1, 1, N = 1000, axes = ax) 274 plt.show() 275 """ 276 277 from pykep import propagate_taylor 278 import matplotlib.pyplot as plt 279 280 if axes is None: 281 fig = plt.figure() 282 ax = fig.gca(projection='3d') 283 else: 284 ax = axes 285 286 # We define the integration time ... 287 dt = tof / (N - 1) 288 289 # ... and calcuate the cartesian components for r 290 x = [0.0] * N 291 y = [0.0] * N 292 z = [0.0] * N 293 294 # Replace r0, v0, m0 for r, v, m 295 r = r0 296 v = v0 297 m = m0 298 # We calculate the spacecraft position at each dt 299 for i in range(N): 300 x[i] = r[0] / units 301 y[i] = r[1] / units 302 z[i] = r[2] / units 303 r, v, m = propagate_taylor(r, v, m, thrust, dt, mu, veff, -10, -10) 304 305 # And we plot 306 if legend: 307 label = 'constant thrust arc' 308 else: 309 label = None 310 ax.plot(x, y, z, c=color, label=label) 311 312 if legend: 313 ax.legend() 314 315 if axes is None: # show only if axis is not set 316 plt.show() 317 return ax 318 319 320def plot_taylor_disturbance(r0, v0, m0, thrust, disturbance, tof, mu, veff, N=60, units=1, color='b', legend=False, axes=None): 321 """ 322 ax = plot_taylor_disturbance(r, v, m, thrust, disturbance, t, mu, veff, N=60, units=1, color='b', legend=False, axes=None): 323 324 - axes: 3D axis object created using fig.gca(projection='3d') 325 - r0: initial position (cartesian coordinates) 326 - v0: initial velocity (cartesian coordinates) 327 - m0: initial mass 328 - thrust: cartesian components for the constant thrust 329 - disturbance: cartesian components for a constant disturbance (will not affect mass) 330 - tof: propagation time 331 - mu: gravitational parameter 332 - veff: the product Isp * g0 333 - N: number of points to be plotted along one arc 334 - units: the length unit to be used in the plot 335 - color: matplotlib color to use to plot the line 336 - legend: when True it plots also the legend 337 338 Plots the result of a taylor propagation of constant thrust 339 """ 340 341 from pykep import propagate_taylor_disturbance 342 import matplotlib.pyplot as plt 343 344 if axes is None: 345 fig = plt.figure() 346 ax = fig.gca(projection='3d') 347 else: 348 ax = axes 349 350 # We define the integration time ... 351 dt = tof / (N - 1) 352 353 # ... and calcuate the cartesian components for r 354 x = [0.0] * N 355 y = [0.0] * N 356 z = [0.0] * N 357 358 # Replace r0, v0 and m0 359 r = r0 360 v = v0 361 m = m0 362 # We calculate the spacecraft position at each dt 363 for i in range(N): 364 x[i] = r[0] / units 365 y[i] = r[1] / units 366 z[i] = r[2] / units 367 r, v, m = propagate_taylor_disturbance( 368 r, v, m, thrust, disturbance, dt, mu, veff, -10, -10) 369 370 # And we plot 371 if legend: 372 label = 'constant thrust arc' 373 else: 374 label = None 375 ax.plot(x, y, z, c=color, label=label) 376 return ax 377 378 379def plot_sf_leg(leg, N=5, units=1, color='b', legend=False, plot_line=True, plot_segments=False, axes=None): 380 """ 381 ax = plot_sf_leg(leg, N=5, units=1, color='b', legend=False, no_trajectory=False, axes=None): 382 383 - axes: 3D axis object created using fig.gca(projection='3d') 384 - leg: a pykep.sims_flanagan.leg 385 - N: number of points to be plotted along one arc 386 - units: the length unit to be used in the plot 387 - color: matplotlib color to use to plot the trajectory and the grid points 388 - legend when True it plots also the legend 389 - plot_line: when True plots also the trajectory (between mid-points and grid points) 390 391 Plots a Sims-Flanagan leg 392 393 Example:: 394 395 from mpl_toolkits.mplot3d import Axes3D 396 import matplotlib.pyplot as plt 397 398 fig = plt.figure() 399 ax = fig.gca(projection='3d') 400 t1 = epoch(0) 401 pl = planet_ss('earth') 402 rE,vE = pl.eph(t1) 403 plot_planet(pl,t0=t1, units=AU, axes=ax) 404 405 t2 = epoch(440) 406 pl = planet_ss('mars') 407 rM, vM = pl.eph(t2) 408 plot_planet(pl,t0=t2, units=AU, axes=ax) 409 410 sc = sims_flanagan.spacecraft(4500,0.5,2500) 411 x0 = sims_flanagan.sc_state(rE,vE,sc.mass) 412 xe = sims_flanagan.sc_state(rM,vM,sc.mass) 413 l = sims_flanagan.leg(t1,x0,[1,0,0]*5,t2,xe,sc,MU_SUN) 414 415 plot_sf_leg(l, units=AU, axes=ax) 416 """ 417 from pykep import propagate_lagrangian, AU, DAY2SEC, G0, propagate_taylor 418 import numpy as np 419 from scipy.linalg import norm 420 from math import exp 421 import matplotlib.pylab as plt 422 from mpl_toolkits.mplot3d import Axes3D 423 424 if axes is None: 425 fig = plt.figure() 426 ax = fig.gca(projection='3d') 427 else: 428 ax = axes 429 430 # We compute the number of segments for forward and backward propagation 431 n_seg = len(leg.get_throttles()) 432 fwd_seg = (n_seg + 1) // 2 433 back_seg = n_seg // 2 434 435 # We extract information on the spacecraft 436 sc = leg.get_spacecraft() 437 isp = sc.isp 438 max_thrust = sc.thrust 439 440 # And on the leg 441 throttles = leg.get_throttles() 442 mu = leg.get_mu() 443 444 # Forward propagation 445 446 # x,y,z contain the cartesian components of all points (grid+midpoints) 447 x = [0.0] * (fwd_seg * 2 + 1) 448 y = [0.0] * (fwd_seg * 2 + 1) 449 z = [0.0] * (fwd_seg * 2 + 1) 450 451 state = leg.get_xi() 452 453 # Initial conditions 454 r = state.r 455 v = state.v 456 m = state.m 457 x[0] = r[0] / units 458 y[0] = r[1] / units 459 z[0] = r[2] / units 460 461 # We compute all points by propagation 462 for i, t in enumerate(throttles[:fwd_seg]): 463 dt = (t.end.mjd - t.start.mjd) * DAY2SEC 464 alpha = min(norm(t.value), 1.0) 465 # Keplerian propagation and dV application 466 if leg.high_fidelity is False: 467 dV = [max_thrust / m * dt * dumb for dumb in t.value] 468 if plot_line: 469 plot_kepler(r, v, dt / 2, mu, N=N, units=units, 470 color=(alpha, 0, 1 - alpha), axes=ax) 471 r, v = propagate_lagrangian(r, v, dt / 2, mu) 472 x[2 * i + 1] = r[0] / units 473 y[2 * i + 1] = r[1] / units 474 z[2 * i + 1] = r[2] / units 475 # v= v+dV 476 v = [a + b for a, b in zip(v, dV)] 477 if plot_line: 478 plot_kepler(r, v, dt / 2, mu, N=N, units=units, 479 color=(alpha, 0, 1 - alpha), axes=ax) 480 r, v = propagate_lagrangian(r, v, dt / 2, mu) 481 x[2 * i + 2] = r[0] / units 482 y[2 * i + 2] = r[1] / units 483 z[2 * i + 2] = r[2] / units 484 m *= exp(-norm(dV) / isp / G0) 485 # Taylor propagation of constant thrust u 486 else: 487 u = [max_thrust * dumb for dumb in t.value] 488 if plot_line: 489 plot_taylor(r, v, m, u, dt / 2, mu, isp * G0, N=N, 490 units=units, color=(alpha, 0, 1 - alpha), axes=ax) 491 r, v, m = propagate_taylor( 492 r, v, m, u, dt / 2, mu, isp * G0, -12, -12) 493 x[2 * i + 1] = r[0] / units 494 y[2 * i + 1] = r[1] / units 495 z[2 * i + 1] = r[2] / units 496 if plot_line: 497 plot_taylor(r, v, m, u, dt / 2, mu, isp * G0, N=N, 498 units=units, color=(alpha, 0, 1 - alpha), axes=ax) 499 r, v, m = propagate_taylor( 500 r, v, m, u, dt / 2, mu, isp * G0, -12, -12) 501 x[2 * i + 2] = r[0] / units 502 y[2 * i + 2] = r[1] / units 503 z[2 * i + 2] = r[2] / units 504 505 x_grid = x[::2] 506 y_grid = y[::2] 507 z_grid = z[::2] 508 x_midpoint = x[1::2] 509 y_midpoint = y[1::2] 510 z_midpoint = z[1::2] 511 if plot_segments: 512 ax.scatter(x_grid[:-1], y_grid[:-1], z_grid[:-1], 513 label='nodes', marker='o') 514 ax.scatter(x_midpoint, y_midpoint, z_midpoint, 515 label='mid-points', marker='x') 516 ax.scatter(x_grid[-1], y_grid[-1], z_grid[-1], 517 marker='^', c='y', label='mismatch point') 518 519 # Backward propagation 520 521 # x,y,z will contain the cartesian components of 522 x = [0.0] * (back_seg * 2 + 1) 523 y = [0.0] * (back_seg * 2 + 1) 524 z = [0.0] * (back_seg * 2 + 1) 525 526 state = leg.get_xf() 527 528 # Final conditions 529 r = state.r 530 v = state.v 531 m = state.m 532 x[-1] = r[0] / units 533 y[-1] = r[1] / units 534 z[-1] = r[2] / units 535 536 for i, t in enumerate(throttles[-1:-back_seg - 1:-1]): 537 dt = (t.end.mjd - t.start.mjd) * DAY2SEC 538 alpha = min(norm(t.value), 1.0) 539 if leg.high_fidelity is False: 540 dV = [max_thrust / m * dt * dumb for dumb in t.value] 541 if plot_line: 542 plot_kepler(r, v, -dt / 2, mu, N=N, units=units, 543 color=(alpha, 0, 1 - alpha), axes=ax) 544 r, v = propagate_lagrangian(r, v, -dt / 2, mu) 545 x[-2 * i - 2] = r[0] / units 546 y[-2 * i - 2] = r[1] / units 547 z[-2 * i - 2] = r[2] / units 548 # v= v+dV 549 v = [a - b for a, b in zip(v, dV)] 550 if plot_line: 551 plot_kepler(r, v, -dt / 2, mu, N=N, units=units, 552 color=(alpha, 0, 1 - alpha), axes=ax) 553 r, v = propagate_lagrangian(r, v, -dt / 2, mu) 554 x[-2 * i - 3] = r[0] / units 555 y[-2 * i - 3] = r[1] / units 556 z[-2 * i - 3] = r[2] / units 557 m *= exp(norm(dV) / isp / G0) 558 else: 559 u = [max_thrust * dumb for dumb in t.value] 560 if plot_line: 561 plot_taylor(r, v, m, u, -dt / 2, mu, isp * G0, N=N, 562 units=units, color=(alpha, 0, 1 - alpha), axes=ax) 563 r, v, m = propagate_taylor( 564 r, v, m, u, -dt / 2, mu, isp * G0, -12, -12) 565 x[-2 * i - 2] = r[0] / units 566 y[-2 * i - 2] = r[1] / units 567 z[-2 * i - 2] = r[2] / units 568 if plot_line: 569 plot_taylor(r, v, m, u, -dt / 2, mu, isp * G0, N=N, 570 units=units, color=(alpha, 0, 1 - alpha), axes=ax) 571 r, v, m = propagate_taylor( 572 r, v, m, u, -dt / 2, mu, isp * G0, -12, -12) 573 x[-2 * i - 3] = r[0] / units 574 y[-2 * i - 3] = r[1] / units 575 z[-2 * i - 3] = r[2] / units 576 577 x_grid = x[::2] 578 y_grid = y[::2] 579 z_grid = z[::2] 580 x_midpoint = x[1::2] 581 y_midpoint = y[1::2] 582 z_midpoint = z[1::2] 583 584 if plot_segments: 585 ax.scatter(x_grid[1:], y_grid[1:], z_grid[ 586 1:], marker='o', label='nodes') 587 ax.scatter(x_midpoint, y_midpoint, z_midpoint, 588 marker='x', label='mid-points') 589 ax.scatter(x_grid[0], y_grid[0], z_grid[0], 590 marker='^', c='y', label='mismatch point') 591 592 if legend: 593 ax.legend() 594 595 if axes is None: # show only if axis is not set 596 plt.show() 597 return ax 598