# `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()