Source code for rebound.plotting

# -*- coding: utf-8 -*-
import math
from .particle import Particle
from itertools import cycle

[docs]def OrbitPlot(sim, figsize=None, fancy=False, slices=0, xlim=None, ylim=None, unitlabel=None, color=False, periastron=False, orbit_type="trail", lw=1., plotparticles=[], primary=None, Narc=128): """ Convenience function for plotting instantaneous orbits. Parameters ---------- figsize : tuple of float, optional Tuple defining the figure size (default: (5,5)) fancy : bool (default: False) Changes various settings to create a fancy looking plot slices : float, optional Default is 0, showing the orbits in the xy plane only. Set to a value between 0 and 1 to create three plots, showing the orbits from different directions. The value corresponds to the size of the additional plots relative to the main plot. xlim : tuple of float, optional Limits for x axes (default: None = automatically determined) ylim : tuple of float, optional Limits for y axes (default: None = automatically determined) unitlabel : str, optional String describing the units, shown on axis labels (default: None) color : bool, str or list, optional By default plots are black and white. If set to True, plots use a color cycle. If a string or list of strings, e.g. ['red', 'cyan'], will cycle between passed colors. periastron : bool, optional Draw a marker at periastron (default: False) orbit_type : str, optional This argument determines the type of orbit show. By default it shows the orbit as a trailing and fading line ("trail"). Other object are: "solid", None. lw : float, optional Linewidth used in plots (default: 1.) plotparticles : list, optional List of particles to plot. Can be a list of any valid keys for accessing sim.particles, i.e., integer indices or hashes (default: plot all particles) primary : rebound.Particle, optional Primary to use for the osculating orbit (default: Jacobi center of mass) Narc : int, optional Number of points used in an orbit. Increase this number for highly eccentric orbits. (default: 128) Returns ------- fig, ax_main, (ax_top, ax_right) The function return the matplotlib figure as well as the axes (three axes if slices>0.) Examples -------- The following example illustrates a typical use case. >>> sim = rebound.Simulation() >>> sim.add(m=1) >>> sim.add(a=1) >>> fig, ax_main = rebound.OrbitPlot(sim) >>> fig.savefig("image.png") # save figure to file >>> fig.show() # show figure on screen """ try: import matplotlib.pyplot as plt from matplotlib import gridspec from mpl_toolkits.axes_grid1 import make_axes_locatable import numpy as np except: raise ImportError("Error importing matplotlib and/or numpy. Plotting functions not available. If running from within a jupyter notebook, try calling '%matplotlib inline' beforehand.") if unitlabel is not None: unitlabel = " " + unitlabel else: unitlabel = "" if figsize is None: if slices>0.: figsize = (8,8) else: figsize = (5,5) fig = plt.figure(figsize=figsize) ax_main = plt.subplot(111,aspect="equal") ax_main.set_xlabel("x"+unitlabel) ax_main.set_ylabel("y"+unitlabel) if slices>0.: divider = make_axes_locatable(ax_main) divider.set_aspect(True) ax_top = divider.append_axes("top", size="%.2f%%"%(100.*slices), sharex=ax_main) ax_top.set_aspect('equal', adjustable='datalim') ax_right = divider.append_axes("right", size="%.2f%%"%(100.*slices), sharey=ax_main) ax_right.set_aspect('equal', adjustable='datalim') plt.setp(ax_top.get_xticklabels(), visible=False) plt.setp(ax_top.get_xticklines(), visible=False) ax_top.set_ylabel("z"+unitlabel) plt.setp(ax_right.get_yticklabels(), visible=False) plt.setp(ax_right.get_yticklines(), visible=False) ax_right.set_xlabel("z"+unitlabel) OrbitPlotOneSlice(sim, ax_main, Narc=Narc, color=color, periastron=periastron, orbit_type=orbit_type, lw=lw, axes="xy",fancy=fancy, plotparticles=plotparticles, primary=primary) if slices>0.: OrbitPlotOneSlice(sim, ax_right, Narc=Narc, color=color, periastron=periastron, orbit_type=orbit_type, lw=lw,fancy=fancy, axes="zy", plotparticles=plotparticles, primary=primary) OrbitPlotOneSlice(sim, ax_top, Narc=Narc, color=color, periastron=periastron, orbit_type=orbit_type, lw=lw,fancy=fancy, axes="xz", plotparticles=plotparticles, primary=primary) if xlim is not None: ax_main.set_xlim(xlim) if ylim is not None: ax_main.set_ylim(ylim) if fancy: ax_main.apply_aspect() if slices>0.: ax_top.apply_aspect() ax_right.apply_aspect() ax_main.autoscale(False) if slices>0.: ax_top.autoscale(False) ax_right.autoscale(False) OrbitPlotAddFancyStars(ax_main,lw) if slices>0.: OrbitPlotAddFancyStars(ax_top,lw,slices) OrbitPlotAddFancyStars(ax_right,lw,slices) if slices>0.: return fig, ax_main, ax_top, ax_right else: return fig, ax_main
def get_color(color): """ Takes a string for a color name defined in matplotlib and returns of a 3-tuple of RGB values. Will simply return passed value if it's a tuple of length three. Parameters ---------- color : str Name of matplotlib color to calculate RGB values for. """ if isinstance(color, tuple) and len(color) == 3: # already a tuple of RGB values return color try: import matplotlib.colors as mplcolors except: raise ImportError("Error importing matplotlib. If running from within a jupyter notebook, try calling '%matplotlib inline' beforehand.") try: hexcolor = mplcolors.cnames[color] except KeyError: raise AttributeError("Color not recognized in matplotlib.") hexcolor = hexcolor.lstrip('#') lv = len(hexcolor) return tuple(int(hexcolor[i:i + lv // 3], 16)/255. for i in range(0, lv, lv // 3)) # tuple of rgb values def fading_line(x, y, color='black', alpha_initial=1., alpha_final=0., glow=False, **kwargs): """ Returns a matplotlib LineCollection connecting the points in the x and y lists, with a single color and alpha varying from alpha_initial to alpha_final along the line. Can pass any kwargs you can pass to LineCollection, like linewidgth. Parameters ---------- x : list or array of floats for the positions on the (plot's) x axis y : list or array of floats for the positions on the (plot's) y axis color : matplotlib color for the line. Can also pass a 3-tuple of RGB values (default: 'black') alpha_initial: Limiting value of alpha to use at the beginning of the arrays. alpha_final: Limiting value of alpha to use at the end of the arrays. """ try: from matplotlib.collections import LineCollection from matplotlib.colors import LinearSegmentedColormap import numpy as np except: raise ImportError("Error importing matplotlib and/or numpy. Plotting functions not available. If running from within a jupyter notebook, try calling '%matplotlib inline' beforehand.") if "lw" not in kwargs: kwargs["lw"] = 1 lw = kwargs["lw"] if glow: kwargs["lw"] = 1*lw fl1 = fading_line(x, y, color, alpha_initial, alpha_final, glow=False, **kwargs) kwargs["lw"] = 2*lw alpha_initial *= 0.5 alpha_final *= 0.5 fl2 = fading_line(x, y, color, alpha_initial, alpha_final, glow=False, **kwargs) kwargs["lw"] = 6*lw alpha_initial *= 0.5 alpha_final *= 0.5 fl3 = fading_line(x, y, color, alpha_initial, alpha_final, glow=False, **kwargs) return [fl3,fl2,fl1] color = get_color(color) cdict = {'red': ((0.,color[0],color[0]),(1.,color[0],color[0])), 'green': ((0.,color[1],color[1]),(1.,color[1],color[1])), 'blue': ((0.,color[2],color[2]),(1.,color[2],color[2])), 'alpha': ((0.,alpha_initial, alpha_initial), (1., alpha_final, alpha_final))} Npts = len(x) if len(y) != Npts: raise AttributeError("x and y must have same dimension.") segments = np.zeros((Npts-1,2,2)) segments[0][0] = [x[0], y[0]] for i in range(1,Npts-1): pt = [x[i], y[i]] segments[i-1][1] = pt segments[i][0] = pt segments[-1][1] = [x[-1], y[-1]] individual_cm = LinearSegmentedColormap('indv1', cdict) lc = LineCollection(segments, cmap=individual_cm, **kwargs) lc.set_array(np.linspace(0.,1.,len(segments))) return lc def OrbitPlotOneSlice(sim, ax, Narc=128, color=False, periastron=False, orbit_type="trial", lw=1., axes="xy", plotparticles=[], primary=None, fancy=False): import matplotlib.pyplot as plt from matplotlib.collections import LineCollection from matplotlib.colors import LinearSegmentedColormap import numpy as np import random #ax.set_aspect("equal") p_orb_pairs = [] if not plotparticles: plotparticles = range(1, sim.N_real) for i in plotparticles: p = sim.particles[i] p_orb_pairs.append((p, p.calculate_orbit(primary=primary))) if color: if color == True: colors = [(1.,0.,0.),(0.,0.75,0.75),(0.75,0.,0.75),(0.75, 0.75, 0,),(0., 0., 0.),(0., 0., 1.),(0., 0.5, 0.)] if isinstance(color, str): colors = [get_color(color)] if isinstance(color, list): colors = [] for c in color: colors.append(get_color(c)) else: if fancy: colors = [(181./206.,66./206.,191./206.)] else: colors = ["black"] coloriterator = cycle(colors) coords = {'x':0, 'y':1, 'z':2} axis0 = coords[axes[0]] axis1 = coords[axes[1]] prim = sim.particles[0] if primary is None else primary if fancy: sun = (256./256.,256./256.,190./256.) opa = 0.035 size = 6000. for i in range(100): ax.scatter(getattr(prim,axes[0]),getattr(prim,axes[1]), alpha=opa, s=size*lw, facecolor=sun, edgecolor=None, zorder=3) size *= 0.95 ax.scatter(getattr(prim,axes[0]),getattr(prim,axes[1]), s=size*lw, facecolor=sun, edgecolor=None, zorder=3) else: ax.scatter(getattr(prim,axes[0]),getattr(prim,axes[1]), marker="*", s=35*lw, facecolor="black", edgecolor=None, zorder=3) proj = {} for p, o in p_orb_pairs: prim = p.jacobi_com if primary is None else primary colori = next(coloriterator) if fancy: ax.scatter(getattr(p,axes[0]), getattr(p,axes[1]), s=25*lw, facecolor=colori, edgecolor=None, zorder=3) else: ax.scatter(getattr(p,axes[0]), getattr(p,axes[1]), s=25*lw, facecolor="black", edgecolor=None, zorder=3) if orbit_type is not None: alpha_final = 0. if orbit_type=="trail" else 1. # fade to 0 with trail hyperbolic = o.a < 0. # Boolean for whether orbit is hyperbolic if hyperbolic is False: pts = np.array(p.sample_orbit(Npts=Narc+1, primary=prim)) proj['x'],proj['y'],proj['z'] = [pts[:,i] for i in range(3)] lc = fading_line(proj[axes[0]], proj[axes[1]], colori, alpha_final=alpha_final, lw=lw, glow=fancy) if type(lc) is list: for l in lc: ax.add_collection(l) else: ax.add_collection(lc) else: pts = np.array(p.sample_orbit(Npts=Narc+1, primary=prim, useTrueAnomaly=False)) # true anomaly stays close to limiting value and switches quickly at pericenter for hyperbolic orbit, so use mean anomaly proj['x'],proj['y'],proj['z'] = [pts[:,i] for i in range(3)] lc = fading_line(proj[axes[0]], proj[axes[1]], colori, alpha_final=alpha_final, lw=lw, glow=fancy) if type(lc) is list: for l in lc: ax.add_collection(l) else: ax.add_collection(lc) alpha = 0.2 if orbit_type=="trail" else 1. pts = np.array(p.sample_orbit(Npts=Narc+1, primary=prim, trailing=False, useTrueAnomaly=False)) proj['x'],proj['y'],proj['z'] = [pts[:,i] for i in range(3)] lc = fading_line(proj[axes[0]], proj[axes[1]], colori, alpha_initial=alpha, alpha_final=alpha, lw=lw, glow=fancy) if type(lc) is list: for l in lc: ax.add_collection(l) else: ax.add_collection(lc) if periastron: newp = Particle(a=o.a, f=0., inc=o.inc, omega=o.omega, Omega=o.Omega, e=o.e, m=p.m, primary=prim, simulation=sim) ax.plot([getattr(prim,axes[0]), getattr(newp,axes[0])], [getattr(prim,axes[1]), getattr(newp,axes[1])], linestyle="dotted", c=colori, zorder=1, lw=lw) ax.scatter([getattr(newp,axes[0])],[getattr(newp,axes[1])], marker="o", s=5.*lw, facecolor="none", edgecolor=colori, zorder=1) def OrbitPlotAddFancyStars(ax,lw,slices=1.): import numpy as np # safe the current random seed to restore later os = np.random.get_state() # always produce the same stars np.random.seed(1) ax.set_facecolor((0.,0.,0.)) for pos in ['top', 'bottom', 'right', 'left']: ax.spines[pos].set_edgecolor((0.3,0.3,0.3)) starcolor = (1.,1.,1.) starsurfacedensity = 0.8 area = np.sqrt(np.sum(np.square(ax.transAxes.transform([1.,1.]) - ax.transAxes.transform([0.,0.]))))*slices nstars = int(starsurfacedensity*area) #small stars xy = np.random.uniform(size=(nstars,2)) ax.scatter(xy[:,0],xy[:,1], transform=ax.transAxes, alpha=0.05, s=8*lw, facecolor=starcolor, edgecolor=None, zorder=3) ax.scatter(xy[:,0],xy[:,1], transform=ax.transAxes, alpha=0.1, s=4*lw, facecolor=starcolor, edgecolor=None, zorder=3) ax.scatter(xy[:,0],xy[:,1], transform=ax.transAxes, alpha=0.2, s=0.5*lw, facecolor=starcolor, edgecolor=None, zorder=3) #large stars xy = np.random.uniform(size=(nstars//4,2)) ax.scatter(xy[:,0],xy[:,1], transform=ax.transAxes, alpha=0.1, s=15*lw, facecolor=starcolor, edgecolor=None, zorder=3) ax.scatter(xy[:,0],xy[:,1], transform=ax.transAxes, alpha=0.1, s=5*lw, facecolor=starcolor, edgecolor=None, zorder=3) ax.scatter(xy[:,0],xy[:,1], transform=ax.transAxes, alpha=0.5, s=2*lw, facecolor=starcolor, edgecolor=None, zorder=3) np.random.set_state(os)