analysis.utils module

Source code
import pickle
import numpy as np

def loadData(fileName, timestep, absolute=False):
    """Load the data in data/fileName/ at the given timestep.

    Parameters:
        fileName (string): data is loaded from data/fileName/
        timestep (integer): timestep at which to load the data
        absolute (bool): true to load relative to parent folder.
            Used for interactive jupyter notebook.

    Returns:
        data instance that was saved by Simulation.save
    """
    prepend = './..' if absolute else './'
    return pickle.load(open(prepend+'/data/{}/data{}.pickle'.format(fileName, timestep), "rb" ))

def calculateEccentricity(M, ref_r, ref_v, r, v):
    """Calculates the orbital eccentricity of a list of particles.

    Parameters:
        M (string): mass of the central object
        ref_r (array): reference position of the central object
        ref_v (array): reference velocity of the central object
        r (array): list of positions of all particles (n particles, 3)
        v (array): list of velocities of all particles (n particles, 3)

    Returns:
        Array of shape (n particles, 3) with the eccentricities of the orbits
    """
    v1 = v - ref_v
    r1 = r - ref_r
    h_vec = np.cross(r1, v1)
    h = np.linalg.norm(h_vec, axis=-1)
    # Make use of $\vec{e} = (v \times (r \times v)) / M$, and e = |\vec{e}|
    e_vec = np.cross(v1, h_vec)/M - r1/np.linalg.norm(r1, axis=-1, keepdims=True)
    e = np.linalg.norm(e_vec, axis=-1)
    return e

#############################  PLOTTING UTILS ########################
# Matplotlib can be tedious at times. These short functions make the
# repetitive parts simpler and ensure consistency.

def plotCOM(ax):
    """Plots cross at (0,0)"""
    ax.scatter([0],[0],marker='+', c='black', s=500, 
               alpha=.5, linewidth=1, zorder=-1)

def plotCenterMasses(ax, data):
    """Plots stars at the position of the central masses"""
    ax.scatter(data['r_vec'][data['type'][:,0]=='center'][:,0], 
               data['r_vec'][data['type'][:,0]=='center'][:,1], 
               s=100, marker="*", c='black', alpha=.7)

def plotTracks(ax, tracks):
    """Plots the tracks of the central masses"""
    for track in tracks:
        ax.plot(track[:,0], track[:,1], 
                c='black', alpha=1.0, linewidth=1)

def setSize(ax, x=None, y=None, mode=None):
    """Sets the size of the plot. If mode='square', the x and y axis
    will have the same scale."""
    if mode=='square': ax.axis('square')
    if x is not None: ax.set_xlim(x)
    if y is not None: ax.set_ylim(y)

def setAxes(ax, x=None, y=None, xcoords=None, ycoords=None, mode=None):
    """"Sets the axis labels (x, y) and their position (None). The mode
    keyword can be used to hide all the axes ('hide'), only plot the bottom
    axis ('bottom') or only plot the bottom and left axes ('bottomleft')"""
    if x is not None: ax.set_xlabel(x)
    if y is not None: ax.set_ylabel(y)
    if mode=='hide':
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.set_xticklabels([])
        ax.set_xticks([])
        ax.set_yticklabels([])
        ax.set_yticks([])
    elif mode=='bottomleft':
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
    elif mode=='bottom':
        ax.get_yaxis().set_ticks([])
        ax.set_ylabel('')   
        ax.spines['left'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
    if xcoords is not None: ax.xaxis.set_label_coords(*xcoords)
    if ycoords is not None: ax.yaxis.set_label_coords(*ycoords)

def stylizePlot(axs):
    """Adds ticks (a la root style) for prettier plots"""
    for ax in axs:
        ax.tick_params(axis='both', which='both', direction='in', top=True, right=True)
        ax.minorticks_on()

Functions

def calculateEccentricity(M, ref_r, ref_v, r, v)

Calculates the orbital eccentricity of a list of particles.

Parameters

M : string
mass of the central object
ref_r : array
reference position of the central object
ref_v : array
reference velocity of the central object
r : array
list of positions of all particles (n particles, 3)
v : array
list of velocities of all particles (n particles, 3)

Returns

Array of shape (n particles, 3) with the eccentricities of the orbits

Source code
def calculateEccentricity(M, ref_r, ref_v, r, v):
    """Calculates the orbital eccentricity of a list of particles.

    Parameters:
        M (string): mass of the central object
        ref_r (array): reference position of the central object
        ref_v (array): reference velocity of the central object
        r (array): list of positions of all particles (n particles, 3)
        v (array): list of velocities of all particles (n particles, 3)

    Returns:
        Array of shape (n particles, 3) with the eccentricities of the orbits
    """
    v1 = v - ref_v
    r1 = r - ref_r
    h_vec = np.cross(r1, v1)
    h = np.linalg.norm(h_vec, axis=-1)
    # Make use of $\vec{e} = (v \times (r \times v)) / M$, and e = |\vec{e}|
    e_vec = np.cross(v1, h_vec)/M - r1/np.linalg.norm(r1, axis=-1, keepdims=True)
    e = np.linalg.norm(e_vec, axis=-1)
    return e
def loadData(fileName, timestep, absolute=False)

Load the data in data/fileName/ at the given timestep.

Parameters

fileName : string
data is loaded from data/fileName/
timestep : integer
timestep at which to load the data
absolute : bool
true to load relative to parent folder. Used for interactive jupyter notebook.

Returns

data instance that was saved by Simulation.save

Source code
def loadData(fileName, timestep, absolute=False):
    """Load the data in data/fileName/ at the given timestep.

    Parameters:
        fileName (string): data is loaded from data/fileName/
        timestep (integer): timestep at which to load the data
        absolute (bool): true to load relative to parent folder.
            Used for interactive jupyter notebook.

    Returns:
        data instance that was saved by Simulation.save
    """
    prepend = './..' if absolute else './'
    return pickle.load(open(prepend+'/data/{}/data{}.pickle'.format(fileName, timestep), "rb" ))
def plotCOM(ax)

Plots cross at (0,0)

Source code
def plotCOM(ax):
    """Plots cross at (0,0)"""
    ax.scatter([0],[0],marker='+', c='black', s=500, 
               alpha=.5, linewidth=1, zorder=-1)
def plotCenterMasses(ax, data)

Plots stars at the position of the central masses

Source code
def plotCenterMasses(ax, data):
    """Plots stars at the position of the central masses"""
    ax.scatter(data['r_vec'][data['type'][:,0]=='center'][:,0], 
               data['r_vec'][data['type'][:,0]=='center'][:,1], 
               s=100, marker="*", c='black', alpha=.7)
def plotTracks(ax, tracks)

Plots the tracks of the central masses

Source code
def plotTracks(ax, tracks):
    """Plots the tracks of the central masses"""
    for track in tracks:
        ax.plot(track[:,0], track[:,1], 
                c='black', alpha=1.0, linewidth=1)
def setAxes(ax, x=None, y=None, xcoords=None, ycoords=None, mode=None)

"Sets the axis labels (x, y) and their position (None). The mode keyword can be used to hide all the axes ('hide'), only plot the bottom axis ('bottom') or only plot the bottom and left axes ('bottomleft')

Source code
def setAxes(ax, x=None, y=None, xcoords=None, ycoords=None, mode=None):
    """"Sets the axis labels (x, y) and their position (None). The mode
    keyword can be used to hide all the axes ('hide'), only plot the bottom
    axis ('bottom') or only plot the bottom and left axes ('bottomleft')"""
    if x is not None: ax.set_xlabel(x)
    if y is not None: ax.set_ylabel(y)
    if mode=='hide':
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.set_xticklabels([])
        ax.set_xticks([])
        ax.set_yticklabels([])
        ax.set_yticks([])
    elif mode=='bottomleft':
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
    elif mode=='bottom':
        ax.get_yaxis().set_ticks([])
        ax.set_ylabel('')   
        ax.spines['left'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
    if xcoords is not None: ax.xaxis.set_label_coords(*xcoords)
    if ycoords is not None: ax.yaxis.set_label_coords(*ycoords)
def setSize(ax, x=None, y=None, mode=None)

Sets the size of the plot. If mode='square', the x and y axis will have the same scale.

Source code
def setSize(ax, x=None, y=None, mode=None):
    """Sets the size of the plot. If mode='square', the x and y axis
    will have the same scale."""
    if mode=='square': ax.axis('square')
    if x is not None: ax.set_xlim(x)
    if y is not None: ax.set_ylim(y)
def stylizePlot(axs)

Adds ticks (a la root style) for prettier plots

Source code
def stylizePlot(axs):
    """Adds ticks (a la root style) for prettier plots"""
    for ax in axs:
        ax.tick_params(axis='both', which='both', direction='in', top=True, right=True)
        ax.minorticks_on()