#!/usr/bin/env python
'''General plotting functions
'''
#Python standard import
from itertools import combinations
#Third party import
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.cm as cm
from mpl_toolkits.mplot3d import Axes3D
import pyorb
#Local import
from .. import constants
[docs]def grid_earth(ax, num_lat=25, num_lon=50, alpha=0.1, res = 100, color='black', hide_ax=True):
'''Add a 3d spherical grid to the given axis that represent the Earth.
'''
lons = np.linspace(-180, 180, num_lon+1) * np.pi/180
lons = lons[:-1]
lats = np.linspace(-90, 90, num_lat) * np.pi/180
lonsl = np.linspace(-180, 180, res) * np.pi/180
latsl = np.linspace(-90, 90, res) * np.pi/180
r_e = constants.R_earth
for lat in lats:
x = r_e*np.cos(lonsl)*np.cos(lat)
y = r_e*np.sin(lonsl)*np.cos(lat)
z = r_e*np.ones(np.size(lonsl))*np.sin(lat)
ax.plot(x,y,z,alpha=alpha,linestyle='-', marker='',color=color)
for lon in lons:
x = r_e*np.cos(lon)*np.cos(latsl)
y = r_e*np.sin(lon)*np.cos(latsl)
z = r_e*np.sin(latsl)
ax.plot(x,y,z,alpha=alpha,color=color)
if hide_ax:
# Hide grid lines
ax.grid(False)
# Hide axes ticks
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
ax.set_axis_off()
return ax
[docs]def orbits(o, **options):
"""This function creates several scatter plots of a set of orbital elements based on the
different possible axis planar projections, calculates all possible permutations of plane
intersections based on the number of columns
:param numpy.ndarray o: Rows are distinct orbits and columns are orbital elements in the order a, e, i, omega, Omega
:param options: dictionary containing all the optional settings
#TODO: this needs updating
Currently the options fields are:
:marker [char]: the marker type
:size [int]: the size of the marker
:title [string]: the title of the plot
:title_font_size [int]: the title font size
:axis_labels [list of strings]: labels for each column
:tick_font_size [int]: the axis tick font size
:window [tuple/list]: the size of the plot window in pixels (assuming dpi = 80)
:save [string]: will not display figure and will instead save it to this path
:show [bool]: if False will do draw() instead of show() allowing script to continue
:tight_rect [list of 4 floats]: configuration for the tight_layout function
Example::
import dpt_tools as dpt
import numpy as np
np.random.seed(19680221)
orbs = np.matrix([
11 + 3 *np.random.randn(1000),
0.5 + 0.2*np.random.randn(1000),
60 + 10 *np.random.randn(1000),
120 + 5 *np.random.randn(1000),
33 + 2 *np.random.randn(1000),
]).T
dpt.orbits(orbs,
title = "My orbital element distribution",
size = 10,
)
"""
if not isinstance(o, np.ndarray):
o = np.array(o)
data_axis = options.get('axis', 0)
if data_axis == 1:
dim_axis = 0
else:
dim_axis = 1
scale = options.get('scale', np.ones((o.shape[dim_axis],), dtype=o.dtype))
#turn on TeX interperter
plt.rc('text', usetex=True)
lis = list(range(o.shape[dim_axis]))
axis_plot = list(combinations(lis, 2))
axis_labels = options.setdefault('axis_labels', None)
limits = options.get('limits', None)
if isinstance(axis_labels, str):
if axis_labels == 'earth-orbit':
axis_labels = ["$a$ [$R_E$]","$e$ [1]","$i$ [deg]","$\omega$ [deg]","$\Omega$ [deg]","$M_0$ [deg]" ]
scale = [1/constants.R_earth] + [1]*5
elif axis_labels == 'earth-state':
axis_labels = ["$x$ [$R_E$]","$y$ [$R_E$]","$z$ [$R_E$]","$v_x$ [km/s]","$v_y$ [km/s]","$v_z$ [km/s]" ]
scale = [1/constants.R_earth]*3 + [1e-3]*3
elif axis_labels == 'sol-orbit':
axis_labels = ["$a$ [AU]","$e$ [1]","$i$ [deg]","$\omega$ [deg]","$\Omega$ [deg]","$M_0$ [deg]" ]
scale = [1/pyorb.AU]*3 + [1e-3]*3
elif axis_labels == 'sol-state':
axis_labels = ["$x$ [AU]","$y$ [AU]","$z$ [AU]","$v_x$ [km/s]","$v_y$ [km/s]","$v_z$ [km/s]" ]
scale = [1/pyorb.AU]*3 + [1e-3]*3
else:
axis_labels = ['']*6
elif axis_labels is None:
axis_labels = ['']*6
if o.shape[dim_axis] == 2:
subplot_cnt = (1,2)
subplot_perms = 2
elif o.shape[dim_axis] == 3:
subplot_cnt = (1,3)
subplot_perms = 3
elif o.shape[dim_axis] == 4:
subplot_cnt = (2,3)
subplot_perms = 6
elif o.shape[dim_axis] == 5:
subplot_cnt = (2,5)
subplot_perms = 10
else:
subplot_cnt = (3,5)
subplot_perms = 15
subplot_cnt_ind = 1
if 'window' in options:
size_in = options['window']
size_in = tuple(x/80.0 for x in size_in)
else:
size_in=(19, 10)
fig = plt.figure(figsize=size_in,dpi=80)
fig.suptitle(
options.get('title','Orbital elements distribution'),
fontsize=options.get('title_font_size',24),
)
axes = []
for I in range( subplot_perms ):
ax = fig.add_subplot(subplot_cnt[0],subplot_cnt[1],subplot_cnt_ind)
axes.append(ax)
if dim_axis == 1:
x = o[:,axis_plot[I][0]]*scale[axis_plot[I][0]]
y = o[:,axis_plot[I][1]]*scale[axis_plot[I][1]]
else:
x = o[axis_plot[I][0],:]*scale[axis_plot[I][0]]
y = o[axis_plot[I][1],:]*scale[axis_plot[I][1]]
sc = ax.scatter(
x.flatten(),
y.flatten(),
marker=options.get('marker','.'),
c=options.setdefault('color','b'),
s=options.get('size',2),
)
if isinstance(options['color'],np.ndarray):
plt.colorbar(sc)
# x_ticks = np.linspace(np.min(o[:,axis_plot[I][0]]),np.max(o[:,axis_plot[I][0]]), num=4)
# plt.xticks( [round(x,1) for x in x_ticks] )
ax.set_xlabel(
axis_labels[axis_plot[I][0]],
fontsize=options.setdefault('ax_font_size',22),
)
ax.set_ylabel(
axis_labels[axis_plot[I][1]],
fontsize=options['ax_font_size'],
)
plt.xticks(fontsize=options.setdefault('tick_font_size',17))
plt.yticks(fontsize=options['tick_font_size'])
if limits is not None:
if len(limits) > axis_plot[I][0]:
ax.set_xlim(*limits[axis_plot[I][0]])
if len(limits) > axis_plot[I][1]:
ax.set_ylim(*limits[axis_plot[I][1]])
subplot_cnt_ind += 1
plt.tight_layout(rect=options.setdefault('tight_rect',[0, 0.03, 1, 0.95]))
if 'save' in options:
fig.savefig(options['save'],bbox_inches='tight')
else:
if options.setdefault('show', False):
plt.show()
else:
plt.draw()
return fig, axes