Add a mini visualization of a Keplerian orbit to a matplotlib axis

miniKeplerViz.py

def keplerPointsReshaped(a, e, Δθ = 0., nPoints = 1_000):
    θ = np.linspace(0., 2. * np.pi, nPoints)
    r = a * (1. - e**2) / (1. + e * np.cos(θ))
    x = r * np.cos(θ + Δθ)
    y = r * np.sin(θ + Δθ)
    return np.array([x, y]).T.reshape(-1, 1, 2)

def addKeplerVizToAx(ax, bounds,
                     a = 1., e = 0., q = 1., Δθ = 0.,
                     nPoints = 1_000,
                     traj_color = 'blue', max_width = 2.,
                     body_color = 'black', body_scale = 1.):

    axins = ax.inset_axes(
        bounds,
        transform=ax.transData,
        xlim=(-2, 2), ylim=(-2, 2),
        xticklabels=[], yticklabels=[])

    lwidths=np.linspace(0., max_width, nPoints)
    lwidths[ lwidths < 0. ] = 0.

    points = 1. / (1. + q) * keplerPointsReshaped(a, e, Δθ, nPoints)

    segments = np.concatenate([points[:-1], points[1:]], axis=1)
    lc = LineCollection(segments, linewidths=lwidths, color=traj_color)

    axins.add_collection(lc)

    disc = patches.Circle(points[0,0], radius = q / (1. + q) * body_scale * 0.05,
                          color=body_color, fill=True, zorder=2.1)
    axins.add_patch(disc)

    points = -q / (1. + q) * keplerPointsReshaped(a, e, Δθ, nPoints)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)
    lc = LineCollection(segments, linewidths=lwidths, color=traj_color)
    axins.add_collection(lc)

    disc = patches.Circle(points[0,0], radius = 1. / (1. + q) * body_scale * 0.05,
                          color=body_color, fill=True, zorder=2.1)
    axins.add_patch(disc)

    axins.tick_params(axis='both', which='both', bottom=False, top=False,
                      left=False, right=False, labelbottom=False, labelleft=False)
    axins.set_frame_on(False)

    axins.set_aspect('equal')

    return axins