How to use poliastro - 10 common examples

def test_earth_satellite_orbit():
    r = [3_539.08827417, 5_310.19903462, 3_066.31301457] * u.km
    v = [-6.49780849, 3.24910291, 1.87521413] * u.km / u.s
    ss = Orbit.from_vectors(Earth, r, v)
    C_D = 2.2 * u.one  # dimentionless (any value would do)
    A = ((np.pi / 4.0) * (u.m ** 2)).to(u.km ** 2)
    m = 100 * u.kg
    spacecraft = Spacecraft(A, C_D, m)
    earth_satellite = EarthSatellite(ss, spacecraft)
    assert isinstance(earth_satellite.orbit, Orbit)

def test_set_frame_plots_same_colors():
    # TODO: Review
    op = StaticOrbitPlotter()
    op.plot_body_orbit(Jupiter, J2000)
    colors1 = [orb[2] for orb in op.trajectories]
    op.set_body_frame(Jupiter)
    colors2 = [orb[2] for orb in op.trajectories]
    assert colors1 == colors2

def test_plot_trajectory_sets_label():
    expected_label = "67P"

    op = StaticOrbitPlotter()
    trajectory = churi.sample()
    op.plot_body_orbit(Mars, J2000, label="Mars")

    op.plot_trajectory(trajectory, label=expected_label)

    legend = plt.gca().get_legend()
    assert legend.get_texts()[1].get_text() == expected_label

def test_redraw_keeps_trajectories():
    # See https://github.com/poliastro/poliastro/issues/518
    op = StaticOrbitPlotter()
    trajectory = churi.sample()
    op.plot_body_orbit(Mars, J2000, label="Mars")
    op.plot_trajectory(trajectory, label="67P")

    assert len(op.trajectories) == 2

    op.set_body_frame(Mars)

    assert len(op.trajectories) == 2

def test_plot_trajectory_sets_label():
    expected_label = "67P"

    op = StaticOrbitPlotter()
    trajectory = churi.sample()
    op.plot_body_orbit(Mars, J2000, label="Mars")

    op.plot_trajectory(trajectory, label=expected_label)

    legend = plt.gca().get_legend()
    assert legend.get_texts()[1].get_text() == expected_label

def near_parabolic():
    r = [8.0e3, 1.0e3, 0.0] * u.km
    v = [-0.5, -0.5, 0.0] * u.km / u.s

    return Orbit.from_vectors(Earth, r, v)

    "attractor,location_str", [(None, "@ssb"), (Earth, "500@399"), (Venus, "500@299")]
)
@pytest.mark.parametrize(
    "plane,refplane_str",
    [(Planes.EARTH_EQUATOR, "earth"), (Planes.EARTH_ECLIPTIC, "ecliptic")],
)
def test_ephem_from_horizons_calls_horizons_with_correct_parameters(
    horizons_mock, attractor, location_str, plane, refplane_str
):
    unused_name = "Strange Object"
    unused_id_type = "id_type"
    epochs = Time(["2020-03-01 12:00:00"], scale="tdb")

    horizons_mock().vectors.return_value = {
        "x": [1] * u.au,
        "y": [0] * u.au,
        "z": [0] * u.au,

def test_propagate_instance():
    tof = 1.0 * u.min
    ss0 = Orbit.from_classical(
        Earth,
        1000 * u.km,
        0.75 * u.one,
        63.4 * u.deg,
        0 * u.deg,
        270 * u.deg,
        80 * u.deg,
    )
    C_D = 2.2 * u.one  # dimentionless (any value would do)
    A = ((np.pi / 4.0) * (u.m ** 2)).to(u.km ** 2)
    m = 100 * u.kg
    spacecraft = Spacecraft(A, C_D, m)
    earth_satellite = EarthSatellite(ss0, spacecraft)
    orbit_with_j2 = earth_satellite.propagate(tof=tof, gravity=EarthGravity.J2)
    orbit_without_perturbation = earth_satellite.propagate(tof)
    orbit_with_atmosphere_and_j2 = earth_satellite.propagate(
        tof=tof, gravity=EarthGravity.J2, atmosphere=COESA76()

def test_propagating_to_certain_nu_is_correct():
    # take an elliptic orbit
    a = 1.0 * u.AU
    ecc = 1.0 / 3.0 * u.one
    _a = 0.0 * u.rad
    nu = 10 * u.deg
    elliptic = Orbit.from_classical(Sun, a, ecc, _a, _a, _a, nu)

    elliptic_at_perihelion = elliptic.propagate_to_anomaly(0.0 * u.rad)
    r_per, _ = elliptic_at_perihelion.rv()

    elliptic_at_aphelion = elliptic.propagate_to_anomaly(np.pi * u.rad)
    r_ap, _ = elliptic_at_aphelion.rv()

    assert_quantity_allclose(norm(r_per), a * (1.0 - ecc))
    assert_quantity_allclose(norm(r_ap), a * (1.0 + ecc))

    # TODO: Test specific values
    assert elliptic_at_perihelion.epoch > elliptic.epoch
    assert elliptic_at_aphelion.epoch > elliptic.epoch

    # test 10 random true anomaly values
    # TODO: Rework this test
    for nu in np.random.uniform(low=-np.pi, high=np.pi, size=10):
        elliptic = elliptic.propagate_to_anomaly(nu * u.rad)
        r, _ = elliptic.rv()
        assert_quantity_allclose(norm(r), a * (1.0 - ecc ** 2) / (1 + ecc * np.cos(nu)))

def test_propagation_hyperbolic():
    # Data from Curtis, example 3.5
    r0 = [Earth.R.to(u.km).value + 300, 0, 0] * u.km
    v0 = [0, 15, 0] * u.km / u.s
    expected_r_norm = 163180 * u.km
    expected_v_norm = 10.51 * u.km / u.s

    ss0 = Orbit.from_vectors(Earth, r0, v0)
    tof = 14941 * u.s
    ss1 = ss0.propagate(tof)
    r, v = ss1.rv()

    assert_quantity_allclose(norm(r), expected_r_norm, rtol=1e-4)
    assert_quantity_allclose(norm(v), expected_v_norm, rtol=1e-3)