Skip to content
Permalink
Browse files

Keep trajectories on redraw, fix poliastro#518

  • Loading branch information...
Juanlu001 committed Jan 7, 2019
1 parent cf14bc2 commit cbe42dba09e403bfd0fc1ce6d157e62d7c1204cb
@@ -1,5 +1,7 @@
import os.path
from collections import namedtuple
from itertools import cycle
from typing import List

import numpy as np
import plotly.colors
@@ -12,6 +14,10 @@
from poliastro.util import norm


class Trajectory(namedtuple("Trajectory", ["trajectory", "state", "label", "color"])):
pass


class BaseOrbitPlotter:
"""
Parent Class for the 2D and 3D OrbitPlotter Classes based on Plotly.
@@ -21,7 +27,7 @@ def __init__(self, figure=None):
self._figure = figure or FigureWidget()
self._layout = None

self._trajectories = []
self._trajectories = [] # type: List[Trajectory]

self._attractor = None
self._attractor_radius = np.inf * u.km
@@ -1,3 +1,5 @@
from typing import List

import matplotlib as mpl
import numpy as np
from astropy import units as u
@@ -8,6 +10,8 @@
from poliastro.twobody.propagation import mean_motion
from poliastro.util import norm

from ._base import Trajectory


class StaticOrbitPlotter:
"""StaticOrbitPlotter class.
@@ -43,11 +47,11 @@ def __init__(self, ax=None, num_points=150, dark=False):
self._frame = None
self._attractor = None
self._attractor_radius = np.inf * u.km
self._orbits = list(tuple()) # type: List[Tuple[Orbit, str, str]]
self._trajectories = [] # type: List[Trajectory]

@property
def orbits(self):
return self._orbits
def trajectories(self):
return self._trajectories

def set_frame(self, p_vec, q_vec, w_vec):
"""Sets perifocal frame.
@@ -64,40 +68,51 @@ def set_frame(self, p_vec, q_vec, w_vec):
else:
self._frame = p_vec, q_vec, w_vec

if self._orbits:
if self._trajectories:
self._redraw()

def _redraw(self):
for artist in self.ax.lines + self.ax.collections:
artist.remove()
self._attractor = None
for orbit, label, color in self._orbits:
self.plot(orbit, label, color)

for trajectory, state, label, color in self._trajectories:
self._plot(trajectory, state, label, color)

self.ax.relim()
self.ax.autoscale()

def _plot_trajectory(self, trajectory, color=None):
rr = trajectory.represent_as(CartesianRepresentation).xyz.transpose()
x, y = self._project(rr)
lines = self.ax.plot(x.to(u.km).value, y.to(u.km).value, "--", color=color)

return lines

def plot_trajectory(self, trajectory, *, label=None, color=None):
"""Plots a precomputed trajectory.
Parameters
----------
trajectory : ~astropy.coordinates.BaseRepresentation, ~astropy.coordinates.BaseCoordinateFrame
Trajectory to plot.
label : str, optional
Label.
color : str, optional
Color string.
"""
lines = []
rr = trajectory.represent_as(CartesianRepresentation).xyz.transpose()
x, y = self._project(rr)
a, = self.ax.plot(
x.to(u.km).value, y.to(u.km).value, "--", color=color, label=label
)
lines.append(a)
lines = self._plot_trajectory(trajectory, color)

if label:
a.set_label(label)
lines[0].set_label(label)
self.ax.legend(
loc="upper left", bbox_to_anchor=(1.05, 1.015), title="Names and epochs"
)

self._trajectories.append(
Trajectory(trajectory, None, label, lines[0].get_color())
)

return lines

def set_attractor(self, attractor):
@@ -137,34 +152,30 @@ def _redraw_attractor(self, min_radius=0 * u.km):
mpl.patches.Circle((0, 0), self._attractor_radius.value, lw=0, color=color)
)

def plot(self, orbit, label=None, color=None, method=mean_motion):
"""Plots state and osculating orbit in their plane.
"""
if not self._frame:
self.set_frame(*orbit.pqw())

self.set_attractor(orbit.attractor)
self._redraw_attractor(orbit.r_p * 0.15) # Arbitrary threshold
positions = orbit.sample(self.num_points, method)
def _plot(self, trajectory, state=None, label=None, color=None):
lines = self._plot_trajectory(trajectory, color)

x0, y0 = self._project(orbit.r[None])
# Plot current position
l, = self.ax.plot(x0.to(u.km).value, y0.to(u.km).value, "o", mew=0, color=color)
if state is not None:
x0, y0 = self._project(state[None])

if (orbit, label, l.get_color()) not in self._orbits:
self._orbits.append((orbit, label, l.get_color()))

lines = self.plot_trajectory(trajectory=positions, color=l.get_color())
lines.append(l)
# Plot current position
l, = self.ax.plot(
x0.to(u.km).value,
y0.to(u.km).value,
"o",
mew=0,
color=lines[0].get_color(),
)
lines.append(l)

if label:
# This will apply the label to either the point or the osculating
# orbit depending on the last plotted line, as they share variable
if not self.ax.get_legend():
size = self.ax.figure.get_size_inches() + [8, 0]
self.ax.figure.set_size_inches(size)
label = generate_label(orbit, label)
l.set_label(label)

# This will apply the label to either the point or the osculating
# orbit depending on the last plotted line
lines[-1].set_label(label)
self.ax.legend(
loc="upper left", bbox_to_anchor=(1.05, 1.015), title="Names and epochs"
)
@@ -174,3 +185,24 @@ def plot(self, orbit, label=None, color=None, method=mean_motion):
self.ax.set_aspect(1)

return lines

def plot(self, orbit, label=None, color=None, method=mean_motion):
"""Plots state and osculating orbit in their plane.
"""
if not self._frame:
self.set_frame(*orbit.pqw())

self.set_attractor(orbit.attractor)
self._redraw_attractor(orbit.r_p * 0.15) # Arbitrary Threshhold
positions = orbit.sample(self.num_points, method)

if label:
label = generate_label(orbit, label)

lines = self._plot(positions, orbit.r, label, color)

self._trajectories.append(
Trajectory(positions, orbit.r, label, lines[0].get_color())
)

return lines
@@ -100,3 +100,19 @@ def test_set_frame_plots_same_colors():
op.set_frame(*jupiter.pqw())
colors2 = [orb[2] for orb in op.trajectories]
assert colors1 == colors2


def test_redraw_keeps_trajectories():
# See https://github.com/poliastro/poliastro/issues/518
op = StaticOrbitPlotter()
earth = Orbit.from_body_ephem(Earth)
mars = Orbit.from_body_ephem(Mars)
trajectory = earth.sample()
op.plot(mars, label="Mars")
op.plot_trajectory(trajectory, label="Earth")

assert len(op.trajectories) == 2

op.set_frame(*mars.pqw())

assert len(op.trajectories) == 2

0 comments on commit cbe42db

Please sign in to comment.
You can’t perform that action at this time.