# `simulation` module

Definition of the Simulation class and the Galaxy constructor.

Source code

    """Definition of the Simulation class and the Galaxy constructor."""

    import os
    import pickle
    import numpy as np
    import matplotlib.pyplot as plt

    from utils import random_unit_vectors, cascade_round
    from distributions import PLUMMER, HERNQUIST, UNIFORM, EXP, NFW
    import acceleration

    ##############################################################################
    ##############################################################################
    class Simulation:
        """"Main class for the gravitational simulation.

        Attributes:
            r_vec (array): position of the particles in the current timestep.
                Shape: (number of particles, 3)
            rprev_vec (array): position of the particles in the previous timestep.
                Shape: (number of particles, 3)
            v_vec (array): velocity in the current timestep.
                Shape: (number of particles, 3)
            a_vec (array): acceleration in the current timestep.
                Shape: (number of particles, 3)
            mass (array): mass of each particle in the simulation.
                Shape: (number of particles,)
            type (array): non-unique identifier for each particle.
                Shape: (number of particles,)
            tracks (array): list of positions through the simulation for central
                masses. Shape: (tracked particles, n+1, 3).
            CONFIG (array): configuration used to create the simulation.
                It will be saved along the state of the simulation.

            dt (float): timestep of the simulation
            n (int): current timestep. Initialized as n=0.
            soft (float): softening length used by the simulation.
            verbose (boolean): When True progress statements will be printed.
        """

        def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs):
            """Constructor for the Simulation class.

            Arguments:
                dt (float): timestep of the simulation
                n (int): current timestep. Initialized as n=0.
                soft (float): softening length used by the simulation.
                verbose (bool): When True progress statements will be printed.
                CONFIG (dict): configuration file used to create the simulation.
                method (string): Optional. Algorithm to use when computing the
                    gravitational forces. One of 'bruteForce', 'bruteForce_numba',
                    'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'.
            """
            self.n = 0
            self.t = 0
            self.dt = dt
            self.soft = soft
            self.verbose = verbose
            self.CONFIG = CONFIG
            # Initialize empty arrays for all necessary properties
            self.r_vec = np.empty((0, 3))
            self.v_vec = np.empty((0, 3))
            self.a_vec = np.empty((0, 3))
            self.mass = np.empty((0,))
            self.type = np.empty((0, 2))
            algorithms = {
                'bruteForce': acceleration.bruteForce,
                'bruteForceNumba': acceleration.bruteForceNumba,
                'bruteForceNumbaOptimized': acceleration.bruteForceNumbaOptimized,
                'bruteForceCPP': acceleration.bruteForceCPP,
                'barnesHutCPP': acceleration.barnesHutCPP
            }
            try:
                self.acceleration = algorithms[method]
            except: raise Exception("Method '{}' unknown".format(method))

        def add(self, body):
            """Add a body to the simulation. It must expose the public attributes
               body.r_vec, body.v_vec, body.a_vec, body.type, body.mass.

            Arguments:
                body: Object to be added to the simulation (e.g. a Galaxy object)
            """
            # Extend all relevant attributes by concatenating the body
            for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
                simattr, bodyattr = getattr(self, name), getattr(body, name)
                setattr(self, name, np.concatenate([simattr, bodyattr], axis=0))
            # Order based on mass
            order = np.argsort(-self.mass)
            for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
                setattr(self, name, getattr(self, name)[order])

            # Update the list of objects to keep track of
            self.tracks = np.empty((np.sum(self.type[:,0]=='center'), 0, 3))

        def step(self):
            """Perform a single step of the simulation.
               Makes use of a 4th order Verlet integrator.
            """
            # Calculate the acceleration
            self.a_vec = self.acceleration(self.r_vec, self.mass, soft=self.soft)
            # Update the state using the Verlet algorithm
            # (A custom algorithm is written mainly for learning purposes)
            self.r_vec, self.rprev_vec = (2*self.r_vec - self.rprev_vec
                + self.a_vec * self.dt**2, self.r_vec)
            self.n += 1
            # Update tracks
            self.tracks = np.concatenate([self.tracks,
                self.r_vec[self.type[:,0]=='center'][:,np.newaxis]], axis=1)

        def run(self, tmax, saveEvery, outputFolder, **kwargs):
            """Run the galactic simulation.

            Attributes:
                tmax (float): Time to which the simulation will run to.
                    This is measured here since the start of the simulation,
                    not since pericenter.
                saveEvery (int): The state is saved every saveEvery steps.
                outputFolder (string): It will be saved to /data/outputFolder/
            """
            # When the simulation starts, intialize self.rprev_vec
            self.rprev_vec = self.r_vec - self.v_vec * self.dt
            if self.verbose: print('Simulation starting. Bon voyage!')
            while(self.t < tmax):
                self.step()
                if(self.n % saveEvery == 0):
                    self.save('data/{}'.format(outputFolder))

            print('Simulation complete.')

        def save(self, outputFolder):
            """Save the state of the simulation to the outputFolder.
               Two files are saved:
                    sim{self.n}.pickle: serializing the state.
                    sim{self.n}.png: a simplified 2D plot of x, y.
            """
            # Create the output folder if it doesn't exist
            if not os.path.exists(outputFolder): os.makedirs(outputFolder)

            # Compute some useful quantities
            # v_vec is not required by the integrator, but useful
            self.v_vec = (self.r_vec - self.rprev_vec)/self.dt
            self.t = self.n * self.dt # prevents numerical rounding errors

            # Serialize state
            file = open(outputFolder+'/data{}.pickle'.format(self.n), "wb")
            pickle.dump({'r_vec': self.r_vec, 'v_vec': self.v_vec,
                         'type': self.type, 'mass': self.mass,
                         'CONFIG': self.CONFIG, 't': self.t,
                         'tracks': self.tracks}, file)

            # Save simplified plot of the current state.
            # Its main use is to detect ill-behaved situations early on.
            fig = plt.figure()
            plt.xlim(-5, 5); plt.ylim(-5, 5); plt.axis('equal')
            # Dark halo is plotted in red, disk in blue, bulge in green
            PLTCON = [('dark', 'r', 0.3), ('disk', 'b', 1.0), ('bulge', 'g', 0.5)]
            for type_, c, a in PLTCON:
                plt.scatter(self.r_vec[self.type[:,0]==type_][:,0],
                    self.r_vec[self.type[:,0]==type_][:,1], s=0.1, c=c, alpha=a)
            # Central mass as a magenta star
            plt.scatter(self.r_vec[self.type[:,0]=='center'][:,0],
                self.r_vec[self.type[:,0]=='center'][:,1], s=100, marker="*", c='m')
            # Save to png file
            fig.savefig(outputFolder+'/sim{}.png'.format(self.n), dpi=150)
            plt.close(fig)

        def project(self, theta, phi, view=0):
            """Projects the 3D simulation onto a plane as viewed from the
               direction described by the (theta, phi, view). Angles in radians.
               (This is used by the simulated annealing algorithm)

            Parameters:
                theta (float): polar angle.
                phi (float): azimuthal angle.
                view (float): rotation along line of sight.
            """
            M1 = np.array([[np.cos(phi), np.sin(phi), 0],
                           [-np.sin(phi), np.cos(phi), 0],
                           [0, 0, 1]])
            M2 = np.array([[1, 0, 0],
                           [0, np.cos(theta), np.sin(theta)],
                           [0, -np.sin(theta), np.cos(theta)]])
            M3 = np.array([[np.cos(view), np.sin(view), 0],
                           [-np.sin(view), np.cos(view), 0],
                           [0, 0, 1]])

            M = np.matmul(M1, np.matmul(M2, M3)) # combine rotations
            r = np.tensordot(self.r_vec, M, axes=[1, 0])

            return r[:,0:2]

        def setOrbit(self, galaxy1, galaxy2, e, rmin, R0):
            """Sets the two galaxies galaxy1, galaxy2 in an orbit.

            Parameters:
                galaxy1 (Galaxy): 1st galaxy (main)
                galaxy2 (Galaxy): 2nd galaxy (companion)
                e: eccentricity of the orbit
                rmin: distance of closest approach
                R0: initial separation
            """
            m1, m2 = np.sum(galaxy1.mass), np.sum(galaxy2.mass)
            # Relevant formulae:
            # $r_0 = r (1 + e) \cos(\phi)$, where $r_0$ ($\neq R_0$) is the semi-latus rectum
            # $r_0 = r_\textup{min} (1 + e)$
            # $v^2 = G M (2/r - 1/a)$, where a is the semimajor axis

            # Solve the reduced two-body problem with reduced mass $\mu$ (mu)
            M = m1 + m2
            r0 = rmin * (1 + e)
            try:
                phi = np.arccos((r0/R0 - 1) / e) # inverting the orbit equation
                phi = -np.abs(phi) # Choose negative (incoming) angle
                ainv = (1 - e) / rmin # ainv = $1/a$, as a may be infinite
                v = np.sqrt(M * (2/R0 - ainv))
                # Finally, calculate the angle of motion. angle = tan(dy/dx)
                # $dy/dx = ((dr/d\phi) sin(\phi) + r \cos(\phi))/((dr/d\phi) cos(\phi) - r \sin(\phi))$
                dy = R0/r0 * e * np.sin(phi)**2 + np.cos(phi)
                dx = R0/r0 * e * np.sin(phi) * np.cos(phi) - np.sin(phi)
                vangle = np.arctan2(dy, dx)
            except: raise Exception('''The orbital parameters cannot be satisfied.
                For elliptical orbits check that R0 is posible (<rmax).''')

            # We now need the actual motion of each body
            R_vec = np.array([[R0*np.cos(phi), R0*np.sin(phi), 0.]])
            V_vec = np.array([[v*np.cos(vangle), v*np.sin(vangle), 0.]])

            galaxy1.r_vec += m2/M * R_vec
            galaxy1.v_vec += m2/M * V_vec
            galaxy2.r_vec += -m1/M * R_vec
            galaxy2.v_vec += -m1/M * V_vec

            # Explicitely add the galaxies to the simulation
            self.add(galaxy1)
            self.add(galaxy2)

            if self.verbose: print('Galaxies set in orbit.')

    ##############################################################################
    ##############################################################################
    class Galaxy():
        """"Helper class for creating galaxies.

        Attributes:
            r_vec (array): position of the particles in the current timestep.
                Shape: (number of particles, 3)
            v_vec (array): velocity in the current timestep.
                Shape: (number of particles, 3)
            a_vec (array): acceleration in the current timestep.
                Shape: (number of particles, 3)
            mass (array): mass of each particle in the simulation.
                Shape: (number of particles,)
            type (array): non-unique identifier for each particle.
                Shape: (number of particles,)    """
        def __init__(self, orientation, centralMass, bulge, disk, halo, sim):
            """Constructor for the Galaxy class.

               Parameters:
                    orientation (tupple): (inclination i, argument of pericenter w)
                    centralMass (float): mass at the center of the galaxy
                    bulge (dict): passed to the addBulge method.
                    disk (dict): passed to the addDisk method.
                    halo (dict): passed to the addHalo method.
                    sim (Simulation): simulation object
            """
            if sim.verbose: print('Initializing galaxy')
            # Build the central mass
            self.r_vec = np.zeros((1, 3))
            self.v_vec = np.zeros((1, 3))
            self.a_vec = np.zeros((1, 3))
            self.mass = np.array([centralMass])
            self.type = np.array([['center', 0]])
            # Build the other components
            self.addBulge(**bulge)
            if sim.verbose: print('Bulge created.')
            self.addDisk(**disk)
            if sim.verbose: print('Disk created.')
            self.addHalo(**halo)
            if sim.verbose: print('Halo created.')
            # Correct particles velocities to generate circular orbits
            # $a_\textup{centripetal} = v^2/r$
            a_vec = sim.acceleration(self.r_vec, self.mass, soft=sim.soft)
            a = np.linalg.norm(a_vec, axis=-1, keepdims=True)
            r = np.linalg.norm(self.r_vec, axis=-1, keepdims=True)
            v = np.linalg.norm(self.v_vec[1:], axis=-1, keepdims=True)
            direction_unit = self.v_vec[1:]/v
            # Set orbital velocities (except for central mass)
            self.v_vec[1:] = np.sqrt(a[1:] * r[1:]) * direction_unit
            self.a_vec = np.zeros_like(self.r_vec)
            # Rotate the galaxy into its correct orientation
            self.rotate(*(np.array(orientation)/360 * 2*np.pi))
            if sim.verbose: print('Galaxy set in rotation and oriented.')

        def addBulge(self, model, totalMass, particles, l):
            """Adds a bulge to the galaxy.

                Parameters:
                    model (string): parametrization of the bulge.
                        'plummer' and 'hernquist' are supported.
                    totalMass (float): total mass of the bulge
                    particles (int): number of particles in the bulge
                    l (float): characteristic length scale of the model.
            """
            if particles == 0: return None
            # Divide the mass equally among all particles
            mass = np.ones(particles) * totalMass/particles
            self.mass = np.concatenate([self.mass, mass], axis=0)
            # Create particles according to the radial distribution from model
            if model == 'plummer':
                r = PLUMMER.ppf(np.random.rand(particles), scale=l)
            elif model == 'hernquist':
                r = HERNQUIST.ppf(np.random.rand(particles), scale=l)
            else: raise Exception("""Bulge distribution not allowed.
                        'plummer' and 'hernquist' are the supported values""")
            r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
            self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
            # Set them orbitting along random directions normal to r_vec
            v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
            self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
            # Label the particles
            type_ = [['bulge', 0]]*particles
            self.type = np.concatenate([self.type, type_], axis=0)

        def addDisk(self, model, totalMass, particles, l):
            """Adds a disk to the galaxy.

                Parameters:
                    model (string): parametrization of the disk.
                        'rings', 'uniform' and 'exp' are supported.
                    totalMass (float): total mass of the bulge
                    particles (int): number of particles in the bulge
                    l: fot 'exp' and 'uniform' characteristic length of the
                        model. For 'rings' tupple of the form (inner radius,
                        outer radius, number of rings)
            """
            if particles == 0: return None
            # Create particles according to the radial distribution from model
            if model == 'uniform':
                r = UNIFORM.ppf(np.random.rand(particles), scale=l)
                type_ = [['disk', 0]]*particles
                r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
                self.type = np.concatenate([self.type, type_], axis=0)
            elif model == 'rings':
                # l = [inner radius, outter radius, number of rings]
                distances = np.linspace(*l)
                # Aim for roughly constant areal density
                # Cascade rounding preserves the total number of particles
                perRing = cascade_round(particles * distances / np.sum(distances))
                particles = int(np.sum(perRing)) # prevents numerical errors
                r_vec = np.empty((0, 3))
                for d, n, i in zip(distances, perRing, range(l[2])):
                    type_ = [['disk', i+1]]*int(n)
                    self.type = np.concatenate([self.type, type_], axis=0)
                    phi = np.linspace(0, 2 * np.pi, n, endpoint=False)
                    ringr = d * np.array([[np.cos(i), np.sin(i), 0] for i in phi])
                    r_vec = np.concatenate([r_vec, ringr], axis=0)
            elif model == 'exp':
                r = EXP.ppf(np.random.rand(particles), scale=l)
                r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
                type_ = [['disk', 0]]*particles
                self.type = np.concatenate([self.type, type_], axis=0)
            else:
                raise Exception("""Disk distribution not allowed.
                        'uniform', 'rings' and 'exp' are the supported values""")
            self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
            # Divide the mass equally among all particles
            mass = np.ones(particles) * totalMass/particles
            self.mass = np.concatenate([self.mass, mass], axis=0)
            # Set them orbitting along the spin plane
            v_vec = np.cross(r_vec, [0, 0, 1])
            self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)

        def addHalo(self, model, totalMass, particles, rs):
            """Adds a halo to the galaxy.

                Parameters:
                    model (string): parametrization of the halo.
                        Only 'NFW' is supported.
                    totalMass (float): total mass of the halo
                    particles (int): number of particles in the halo
                    rs (float): characteristic length scale of the NFW profile.
            """
            if particles == 0: return None
            # Divide the mass equally among all particles
            mass = np.ones(particles)*totalMass/particles
            self.mass = np.concatenate([self.mass, mass], axis=0)
            # Create particles according to the radial distribution from model
            if model == 'NFW':
                r = NFW.ppf(np.random.rand(particles), scale=rs)
            else: raise Exception("""Bulge distribution not allowed.
                        'plummer' and 'hernquist' are the supported values""")
            r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
            self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
            # Orbit along random directions normal to the radial vector
            v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
            self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
            # Label the particles
            type_ = [['dark'], 0]*particles
            self.type = np.concatenate([self.type, type_], axis=0)

        def rotate(self, theta, phi):
            """Rotates the galaxy so that its spin is along the (theta, phi)
               direction.

            Parameters:
                theta (float): polar angle.
                phi (float): azimuthal angle.
            """
            M1 = np.array([[1, 0, 0],
                           [0, np.cos(theta), np.sin(theta)],
                           [0, -np.sin(theta), np.cos(theta)]])
            M2 = np.array([[np.cos(phi), np.sin(phi), 0],
                           [-np.sin(phi), np.cos(phi), 0],
                           [0, 0, 1]])
            M = np.matmul(M1, M2) # combine rotations
            self.r_vec = np.tensordot(self.r_vec, M, axes=[1, 0])
            self.v_vec = np.tensordot(self.v_vec, M, axes=[1, 0])

## Classes

`class Galaxy`

  "Helper class for creating galaxies.

  ## Attributes

  `r_vec` :`array`
    position of the particles in the current timestep. Shape: (number of particles, 3)
  `v_vec` :`array`
    velocity in the current timestep. Shape: (number of particles, 3)
  `a_vec` :`array`
    acceleration in the current timestep. Shape: (number of particles, 3)
  `mass` :`array`
    mass of each particle in the simulation. Shape: (number of particles,)
  `type` :`array`
    non-unique identifier for each particle. Shape: (number of particles,)
  Source code

      class Galaxy():
          """"Helper class for creating galaxies.

          Attributes:
              r_vec (array): position of the particles in the current timestep.
                  Shape: (number of particles, 3)
              v_vec (array): velocity in the current timestep.
                  Shape: (number of particles, 3)
              a_vec (array): acceleration in the current timestep.
                  Shape: (number of particles, 3)
              mass (array): mass of each particle in the simulation.
                  Shape: (number of particles,)
              type (array): non-unique identifier for each particle.
                  Shape: (number of particles,)    """
          def __init__(self, orientation, centralMass, bulge, disk, halo, sim):
              """Constructor for the Galaxy class.

                 Parameters:
                      orientation (tupple): (inclination i, argument of pericenter w)
                      centralMass (float): mass at the center of the galaxy
                      bulge (dict): passed to the addBulge method.
                      disk (dict): passed to the addDisk method.
                      halo (dict): passed to the addHalo method.
                      sim (Simulation): simulation object
              """
              if sim.verbose: print('Initializing galaxy')
              # Build the central mass
              self.r_vec = np.zeros((1, 3))
              self.v_vec = np.zeros((1, 3))
              self.a_vec = np.zeros((1, 3))
              self.mass = np.array([centralMass])
              self.type = np.array([['center', 0]])
              # Build the other components
              self.addBulge(**bulge)
              if sim.verbose: print('Bulge created.')
              self.addDisk(**disk)
              if sim.verbose: print('Disk created.')
              self.addHalo(**halo)
              if sim.verbose: print('Halo created.')
              # Correct particles velocities to generate circular orbits
              # $a_\textup{centripetal} = v^2/r$
              a_vec = sim.acceleration(self.r_vec, self.mass, soft=sim.soft)
              a = np.linalg.norm(a_vec, axis=-1, keepdims=True)
              r = np.linalg.norm(self.r_vec, axis=-1, keepdims=True)
              v = np.linalg.norm(self.v_vec[1:], axis=-1, keepdims=True)
              direction_unit = self.v_vec[1:]/v
              # Set orbital velocities (except for central mass)
              self.v_vec[1:] = np.sqrt(a[1:] * r[1:]) * direction_unit
              self.a_vec = np.zeros_like(self.r_vec)
              # Rotate the galaxy into its correct orientation
              self.rotate(*(np.array(orientation)/360 * 2*np.pi))
              if sim.verbose: print('Galaxy set in rotation and oriented.')

          def addBulge(self, model, totalMass, particles, l):
              """Adds a bulge to the galaxy.

                  Parameters:
                      model (string): parametrization of the bulge.
                          'plummer' and 'hernquist' are supported.
                      totalMass (float): total mass of the bulge
                      particles (int): number of particles in the bulge
                      l (float): characteristic length scale of the model.
              """
              if particles == 0: return None
              # Divide the mass equally among all particles
              mass = np.ones(particles) * totalMass/particles
              self.mass = np.concatenate([self.mass, mass], axis=0)
              # Create particles according to the radial distribution from model
              if model == 'plummer':
                  r = PLUMMER.ppf(np.random.rand(particles), scale=l)
              elif model == 'hernquist':
                  r = HERNQUIST.ppf(np.random.rand(particles), scale=l)
              else: raise Exception("""Bulge distribution not allowed.
                          'plummer' and 'hernquist' are the supported values""")
              r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
              self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
              # Set them orbitting along random directions normal to r_vec
              v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
              self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
              # Label the particles
              type_ = [['bulge', 0]]*particles
              self.type = np.concatenate([self.type, type_], axis=0)

          def addDisk(self, model, totalMass, particles, l):
              """Adds a disk to the galaxy.

                  Parameters:
                      model (string): parametrization of the disk.
                          'rings', 'uniform' and 'exp' are supported.
                      totalMass (float): total mass of the bulge
                      particles (int): number of particles in the bulge
                      l: fot 'exp' and 'uniform' characteristic length of the
                          model. For 'rings' tupple of the form (inner radius,
                          outer radius, number of rings)
              """
              if particles == 0: return None
              # Create particles according to the radial distribution from model
              if model == 'uniform':
                  r = UNIFORM.ppf(np.random.rand(particles), scale=l)
                  type_ = [['disk', 0]]*particles
                  r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
                  self.type = np.concatenate([self.type, type_], axis=0)
              elif model == 'rings':
                  # l = [inner radius, outter radius, number of rings]
                  distances = np.linspace(*l)
                  # Aim for roughly constant areal density
                  # Cascade rounding preserves the total number of particles
                  perRing = cascade_round(particles * distances / np.sum(distances))
                  particles = int(np.sum(perRing)) # prevents numerical errors
                  r_vec = np.empty((0, 3))
                  for d, n, i in zip(distances, perRing, range(l[2])):
                      type_ = [['disk', i+1]]*int(n)
                      self.type = np.concatenate([self.type, type_], axis=0)
                      phi = np.linspace(0, 2 * np.pi, n, endpoint=False)
                      ringr = d * np.array([[np.cos(i), np.sin(i), 0] for i in phi])
                      r_vec = np.concatenate([r_vec, ringr], axis=0)
              elif model == 'exp':
                  r = EXP.ppf(np.random.rand(particles), scale=l)
                  r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
                  type_ = [['disk', 0]]*particles
                  self.type = np.concatenate([self.type, type_], axis=0)
              else:
                  raise Exception("""Disk distribution not allowed.
                          'uniform', 'rings' and 'exp' are the supported values""")
              self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
              # Divide the mass equally among all particles
              mass = np.ones(particles) * totalMass/particles
              self.mass = np.concatenate([self.mass, mass], axis=0)
              # Set them orbitting along the spin plane
              v_vec = np.cross(r_vec, [0, 0, 1])
              self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)

          def addHalo(self, model, totalMass, particles, rs):
              """Adds a halo to the galaxy.

                  Parameters:
                      model (string): parametrization of the halo.
                          Only 'NFW' is supported.
                      totalMass (float): total mass of the halo
                      particles (int): number of particles in the halo
                      rs (float): characteristic length scale of the NFW profile.
              """
              if particles == 0: return None
              # Divide the mass equally among all particles
              mass = np.ones(particles)*totalMass/particles
              self.mass = np.concatenate([self.mass, mass], axis=0)
              # Create particles according to the radial distribution from model
              if model == 'NFW':
                  r = NFW.ppf(np.random.rand(particles), scale=rs)
              else: raise Exception("""Bulge distribution not allowed.
                          'plummer' and 'hernquist' are the supported values""")
              r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
              self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
              # Orbit along random directions normal to the radial vector
              v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
              self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
              # Label the particles
              type_ = [['dark'], 0]*particles
              self.type = np.concatenate([self.type, type_], axis=0)

          def rotate(self, theta, phi):
              """Rotates the galaxy so that its spin is along the (theta, phi)
                 direction.

              Parameters:
                  theta (float): polar angle.
                  phi (float): azimuthal angle.
              """
              M1 = np.array([[1, 0, 0],
                             [0, np.cos(theta), np.sin(theta)],
                             [0, -np.sin(theta), np.cos(theta)]])
              M2 = np.array([[np.cos(phi), np.sin(phi), 0],
                             [-np.sin(phi), np.cos(phi), 0],
                             [0, 0, 1]])
              M = np.matmul(M1, M2) # combine rotations
              self.r_vec = np.tensordot(self.r_vec, M, axes=[1, 0])
              self.v_vec = np.tensordot(self.v_vec, M, axes=[1, 0])

  ### Methods

  `def __init__(self, orientation, centralMass, bulge, disk, halo, sim)`

    Constructor for the Galaxy class.

    ## Parameters

    `orientation` :`tupple`
      (inclination i, argument of pericenter w)
    `centralMass` :`float`
      mass at the center of the galaxy
    `bulge` :`dict`
      passed to the addBulge method.
    `disk` :`dict`
      passed to the addDisk method.
    `halo` :`dict`
      passed to the addHalo method.
    `sim` :`Simulation`
      simulation object
    Source code

        def __init__(self, orientation, centralMass, bulge, disk, halo, sim):
            """Constructor for the Galaxy class.

               Parameters:
                    orientation (tupple): (inclination i, argument of pericenter w)
                    centralMass (float): mass at the center of the galaxy
                    bulge (dict): passed to the addBulge method.
                    disk (dict): passed to the addDisk method.
                    halo (dict): passed to the addHalo method.
                    sim (Simulation): simulation object
            """
            if sim.verbose: print('Initializing galaxy')
            # Build the central mass
            self.r_vec = np.zeros((1, 3))
            self.v_vec = np.zeros((1, 3))
            self.a_vec = np.zeros((1, 3))
            self.mass = np.array([centralMass])
            self.type = np.array([['center', 0]])
            # Build the other components
            self.addBulge(**bulge)
            if sim.verbose: print('Bulge created.')
            self.addDisk(**disk)
            if sim.verbose: print('Disk created.')
            self.addHalo(**halo)
            if sim.verbose: print('Halo created.')
            # Correct particles velocities to generate circular orbits
            # $a_\textup{centripetal} = v^2/r$
            a_vec = sim.acceleration(self.r_vec, self.mass, soft=sim.soft)
            a = np.linalg.norm(a_vec, axis=-1, keepdims=True)
            r = np.linalg.norm(self.r_vec, axis=-1, keepdims=True)
            v = np.linalg.norm(self.v_vec[1:], axis=-1, keepdims=True)
            direction_unit = self.v_vec[1:]/v
            # Set orbital velocities (except for central mass)
            self.v_vec[1:] = np.sqrt(a[1:] * r[1:]) * direction_unit
            self.a_vec = np.zeros_like(self.r_vec)
            # Rotate the galaxy into its correct orientation
            self.rotate(*(np.array(orientation)/360 * 2*np.pi))
            if sim.verbose: print('Galaxy set in rotation and oriented.')

  `def addBulge(self, model, totalMass, particles, l)`

    Adds a bulge to the galaxy.

    ## Parameters

    `model` :`string`
      parametrization of the bulge. 'plummer' and 'hernquist' are supported.
    `totalMass` :`float`
      total mass of the bulge
    `particles` :`int`
      number of particles in the bulge
    `l` :`float`
      characteristic length scale of the model.
    Source code

        def addBulge(self, model, totalMass, particles, l):
            """Adds a bulge to the galaxy.

                Parameters:
                    model (string): parametrization of the bulge.
                        'plummer' and 'hernquist' are supported.
                    totalMass (float): total mass of the bulge
                    particles (int): number of particles in the bulge
                    l (float): characteristic length scale of the model.
            """
            if particles == 0: return None
            # Divide the mass equally among all particles
            mass = np.ones(particles) * totalMass/particles
            self.mass = np.concatenate([self.mass, mass], axis=0)
            # Create particles according to the radial distribution from model
            if model == 'plummer':
                r = PLUMMER.ppf(np.random.rand(particles), scale=l)
            elif model == 'hernquist':
                r = HERNQUIST.ppf(np.random.rand(particles), scale=l)
            else: raise Exception("""Bulge distribution not allowed.
                        'plummer' and 'hernquist' are the supported values""")
            r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
            self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
            # Set them orbitting along random directions normal to r_vec
            v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
            self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
            # Label the particles
            type_ = [['bulge', 0]]*particles
            self.type = np.concatenate([self.type, type_], axis=0)

  `def addDisk(self, model, totalMass, particles, l)`

    Adds a disk to the galaxy.

    ## Parameters

    `model` :`string`
      parametrization of the disk. 'rings', 'uniform' and 'exp' are supported.
    `totalMass` :`float`
      total mass of the bulge
    `particles` :`int`
      number of particles in the bulge
    `l`
      fot 'exp' and 'uniform' characteristic length of the model. For 'rings' tupple of the form (inner radius, outer radius, number of rings)
    Source code

        def addDisk(self, model, totalMass, particles, l):
            """Adds a disk to the galaxy.

                Parameters:
                    model (string): parametrization of the disk.
                        'rings', 'uniform' and 'exp' are supported.
                    totalMass (float): total mass of the bulge
                    particles (int): number of particles in the bulge
                    l: fot 'exp' and 'uniform' characteristic length of the
                        model. For 'rings' tupple of the form (inner radius,
                        outer radius, number of rings)
            """
            if particles == 0: return None
            # Create particles according to the radial distribution from model
            if model == 'uniform':
                r = UNIFORM.ppf(np.random.rand(particles), scale=l)
                type_ = [['disk', 0]]*particles
                r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
                self.type = np.concatenate([self.type, type_], axis=0)
            elif model == 'rings':
                # l = [inner radius, outter radius, number of rings]
                distances = np.linspace(*l)
                # Aim for roughly constant areal density
                # Cascade rounding preserves the total number of particles
                perRing = cascade_round(particles * distances / np.sum(distances))
                particles = int(np.sum(perRing)) # prevents numerical errors
                r_vec = np.empty((0, 3))
                for d, n, i in zip(distances, perRing, range(l[2])):
                    type_ = [['disk', i+1]]*int(n)
                    self.type = np.concatenate([self.type, type_], axis=0)
                    phi = np.linspace(0, 2 * np.pi, n, endpoint=False)
                    ringr = d * np.array([[np.cos(i), np.sin(i), 0] for i in phi])
                    r_vec = np.concatenate([r_vec, ringr], axis=0)
            elif model == 'exp':
                r = EXP.ppf(np.random.rand(particles), scale=l)
                r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
                type_ = [['disk', 0]]*particles
                self.type = np.concatenate([self.type, type_], axis=0)
            else:
                raise Exception("""Disk distribution not allowed.
                        'uniform', 'rings' and 'exp' are the supported values""")
            self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
            # Divide the mass equally among all particles
            mass = np.ones(particles) * totalMass/particles
            self.mass = np.concatenate([self.mass, mass], axis=0)
            # Set them orbitting along the spin plane
            v_vec = np.cross(r_vec, [0, 0, 1])
            self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)

  `def addHalo(self, model, totalMass, particles, rs)`

    Adds a halo to the galaxy.

    ## Parameters

    `model` :`string`
      parametrization of the halo. Only 'NFW' is supported.
    `totalMass` :`float`
      total mass of the halo
    `particles` :`int`
      number of particles in the halo
    `rs` :`float`
      characteristic length scale of the NFW profile.
    Source code

        def addHalo(self, model, totalMass, particles, rs):
            """Adds a halo to the galaxy.

                Parameters:
                    model (string): parametrization of the halo.
                        Only 'NFW' is supported.
                    totalMass (float): total mass of the halo
                    particles (int): number of particles in the halo
                    rs (float): characteristic length scale of the NFW profile.
            """
            if particles == 0: return None
            # Divide the mass equally among all particles
            mass = np.ones(particles)*totalMass/particles
            self.mass = np.concatenate([self.mass, mass], axis=0)
            # Create particles according to the radial distribution from model
            if model == 'NFW':
                r = NFW.ppf(np.random.rand(particles), scale=rs)
            else: raise Exception("""Bulge distribution not allowed.
                        'plummer' and 'hernquist' are the supported values""")
            r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
            self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
            # Orbit along random directions normal to the radial vector
            v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
            self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
            # Label the particles
            type_ = [['dark'], 0]*particles
            self.type = np.concatenate([self.type, type_], axis=0)

  `def rotate(self, theta, phi)`

    Rotates the galaxy so that its spin is along the (theta, phi) direction.

    ## Parameters

    `theta` :`float`
      polar angle.
    `phi` :`float`
      azimuthal angle.
    Source code

        def rotate(self, theta, phi):
            """Rotates the galaxy so that its spin is along the (theta, phi)
               direction.

            Parameters:
                theta (float): polar angle.
                phi (float): azimuthal angle.
            """
            M1 = np.array([[1, 0, 0],
                           [0, np.cos(theta), np.sin(theta)],
                           [0, -np.sin(theta), np.cos(theta)]])
            M2 = np.array([[np.cos(phi), np.sin(phi), 0],
                           [-np.sin(phi), np.cos(phi), 0],
                           [0, 0, 1]])
            M = np.matmul(M1, M2) # combine rotations
            self.r_vec = np.tensordot(self.r_vec, M, axes=[1, 0])
            self.v_vec = np.tensordot(self.v_vec, M, axes=[1, 0])

`class Simulation`

  "Main class for the gravitational simulation.

  ## Attributes

  `r_vec` :`array`
    position of the particles in the current timestep. Shape: (number of particles, 3)
  `rprev_vec` :`array`
    position of the particles in the previous timestep. Shape: (number of particles, 3)
  `v_vec` :`array`
    velocity in the current timestep. Shape: (number of particles, 3)
  `a_vec` :`array`
    acceleration in the current timestep. Shape: (number of particles, 3)
  `mass` :`array`
    mass of each particle in the simulation. Shape: (number of particles,)
  `type` :`array`
    non-unique identifier for each particle. Shape: (number of particles,)
  `tracks` :`array`
    list of positions through the simulation for central masses. Shape: (tracked particles, n+1, 3).
  `CONFIG` :`array`
    configuration used to create the simulation. It will be saved along the state of the simulation.
  `dt` :`float`
    timestep of the simulation
  `n` :`int`
    current timestep. Initialized as n=0.
  `soft` :`float`
    softening length used by the simulation.
  `verbose` :`boolean`
    When True progress statements will be printed.
  Source code

      class Simulation:
          """"Main class for the gravitational simulation.

          Attributes:
              r_vec (array): position of the particles in the current timestep.
                  Shape: (number of particles, 3)
              rprev_vec (array): position of the particles in the previous timestep.
                  Shape: (number of particles, 3)
              v_vec (array): velocity in the current timestep.
                  Shape: (number of particles, 3)
              a_vec (array): acceleration in the current timestep.
                  Shape: (number of particles, 3)
              mass (array): mass of each particle in the simulation.
                  Shape: (number of particles,)
              type (array): non-unique identifier for each particle.
                  Shape: (number of particles,)
              tracks (array): list of positions through the simulation for central
                  masses. Shape: (tracked particles, n+1, 3).
              CONFIG (array): configuration used to create the simulation.
                  It will be saved along the state of the simulation.

              dt (float): timestep of the simulation
              n (int): current timestep. Initialized as n=0.
              soft (float): softening length used by the simulation.
              verbose (boolean): When True progress statements will be printed.
          """

          def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs):
              """Constructor for the Simulation class.

              Arguments:
                  dt (float): timestep of the simulation
                  n (int): current timestep. Initialized as n=0.
                  soft (float): softening length used by the simulation.
                  verbose (bool): When True progress statements will be printed.
                  CONFIG (dict): configuration file used to create the simulation.
                  method (string): Optional. Algorithm to use when computing the
                      gravitational forces. One of 'bruteForce', 'bruteForce_numba',
                      'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'.
              """
              self.n = 0
              self.t = 0
              self.dt = dt
              self.soft = soft
              self.verbose = verbose
              self.CONFIG = CONFIG
              # Initialize empty arrays for all necessary properties
              self.r_vec = np.empty((0, 3))
              self.v_vec = np.empty((0, 3))
              self.a_vec = np.empty((0, 3))
              self.mass = np.empty((0,))
              self.type = np.empty((0, 2))
              algorithms = {
                  'bruteForce': acceleration.bruteForce,
                  'bruteForceNumba': acceleration.bruteForceNumba,
                  'bruteForceNumbaOptimized': acceleration.bruteForceNumbaOptimized,
                  'bruteForceCPP': acceleration.bruteForceCPP,
                  'barnesHutCPP': acceleration.barnesHutCPP
              }
              try:
                  self.acceleration = algorithms[method]
              except: raise Exception("Method '{}' unknown".format(method))

          def add(self, body):
              """Add a body to the simulation. It must expose the public attributes
                 body.r_vec, body.v_vec, body.a_vec, body.type, body.mass.

              Arguments:
                  body: Object to be added to the simulation (e.g. a Galaxy object)
              """
              # Extend all relevant attributes by concatenating the body
              for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
                  simattr, bodyattr = getattr(self, name), getattr(body, name)
                  setattr(self, name, np.concatenate([simattr, bodyattr], axis=0))
              # Order based on mass
              order = np.argsort(-self.mass)
              for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
                  setattr(self, name, getattr(self, name)[order])

              # Update the list of objects to keep track of
              self.tracks = np.empty((np.sum(self.type[:,0]=='center'), 0, 3))

          def step(self):
              """Perform a single step of the simulation.
                 Makes use of a 4th order Verlet integrator.
              """
              # Calculate the acceleration
              self.a_vec = self.acceleration(self.r_vec, self.mass, soft=self.soft)
              # Update the state using the Verlet algorithm
              # (A custom algorithm is written mainly for learning purposes)
              self.r_vec, self.rprev_vec = (2*self.r_vec - self.rprev_vec
                  + self.a_vec * self.dt**2, self.r_vec)
              self.n += 1
              # Update tracks
              self.tracks = np.concatenate([self.tracks,
                  self.r_vec[self.type[:,0]=='center'][:,np.newaxis]], axis=1)

          def run(self, tmax, saveEvery, outputFolder, **kwargs):
              """Run the galactic simulation.

              Attributes:
                  tmax (float): Time to which the simulation will run to.
                      This is measured here since the start of the simulation,
                      not since pericenter.
                  saveEvery (int): The state is saved every saveEvery steps.
                  outputFolder (string): It will be saved to /data/outputFolder/
              """
              # When the simulation starts, intialize self.rprev_vec
              self.rprev_vec = self.r_vec - self.v_vec * self.dt
              if self.verbose: print('Simulation starting. Bon voyage!')
              while(self.t < tmax):
                  self.step()
                  if(self.n % saveEvery == 0):
                      self.save('data/{}'.format(outputFolder))

              print('Simulation complete.')

          def save(self, outputFolder):
              """Save the state of the simulation to the outputFolder.
                 Two files are saved:
                      sim{self.n}.pickle: serializing the state.
                      sim{self.n}.png: a simplified 2D plot of x, y.
              """
              # Create the output folder if it doesn't exist
              if not os.path.exists(outputFolder): os.makedirs(outputFolder)

              # Compute some useful quantities
              # v_vec is not required by the integrator, but useful
              self.v_vec = (self.r_vec - self.rprev_vec)/self.dt
              self.t = self.n * self.dt # prevents numerical rounding errors

              # Serialize state
              file = open(outputFolder+'/data{}.pickle'.format(self.n), "wb")
              pickle.dump({'r_vec': self.r_vec, 'v_vec': self.v_vec,
                           'type': self.type, 'mass': self.mass,
                           'CONFIG': self.CONFIG, 't': self.t,
                           'tracks': self.tracks}, file)

              # Save simplified plot of the current state.
              # Its main use is to detect ill-behaved situations early on.
              fig = plt.figure()
              plt.xlim(-5, 5); plt.ylim(-5, 5); plt.axis('equal')
              # Dark halo is plotted in red, disk in blue, bulge in green
              PLTCON = [('dark', 'r', 0.3), ('disk', 'b', 1.0), ('bulge', 'g', 0.5)]
              for type_, c, a in PLTCON:
                  plt.scatter(self.r_vec[self.type[:,0]==type_][:,0],
                      self.r_vec[self.type[:,0]==type_][:,1], s=0.1, c=c, alpha=a)
              # Central mass as a magenta star
              plt.scatter(self.r_vec[self.type[:,0]=='center'][:,0],
                  self.r_vec[self.type[:,0]=='center'][:,1], s=100, marker="*", c='m')
              # Save to png file
              fig.savefig(outputFolder+'/sim{}.png'.format(self.n), dpi=150)
              plt.close(fig)

          def project(self, theta, phi, view=0):
              """Projects the 3D simulation onto a plane as viewed from the
                 direction described by the (theta, phi, view). Angles in radians.
                 (This is used by the simulated annealing algorithm)

              Parameters:
                  theta (float): polar angle.
                  phi (float): azimuthal angle.
                  view (float): rotation along line of sight.
              """
              M1 = np.array([[np.cos(phi), np.sin(phi), 0],
                             [-np.sin(phi), np.cos(phi), 0],
                             [0, 0, 1]])
              M2 = np.array([[1, 0, 0],
                             [0, np.cos(theta), np.sin(theta)],
                             [0, -np.sin(theta), np.cos(theta)]])
              M3 = np.array([[np.cos(view), np.sin(view), 0],
                             [-np.sin(view), np.cos(view), 0],
                             [0, 0, 1]])

              M = np.matmul(M1, np.matmul(M2, M3)) # combine rotations
              r = np.tensordot(self.r_vec, M, axes=[1, 0])

              return r[:,0:2]

          def setOrbit(self, galaxy1, galaxy2, e, rmin, R0):
              """Sets the two galaxies galaxy1, galaxy2 in an orbit.

              Parameters:
                  galaxy1 (Galaxy): 1st galaxy (main)
                  galaxy2 (Galaxy): 2nd galaxy (companion)
                  e: eccentricity of the orbit
                  rmin: distance of closest approach
                  R0: initial separation
              """
              m1, m2 = np.sum(galaxy1.mass), np.sum(galaxy2.mass)
              # Relevant formulae:
              # $r_0 = r (1 + e) \cos(\phi)$, where $r_0$ ($\neq R_0$) is the semi-latus rectum
              # $r_0 = r_\textup{min} (1 + e)$
              # $v^2 = G M (2/r - 1/a)$, where a is the semimajor axis

              # Solve the reduced two-body problem with reduced mass $\mu$ (mu)
              M = m1 + m2
              r0 = rmin * (1 + e)
              try:
                  phi = np.arccos((r0/R0 - 1) / e) # inverting the orbit equation
                  phi = -np.abs(phi) # Choose negative (incoming) angle
                  ainv = (1 - e) / rmin # ainv = $1/a$, as a may be infinite
                  v = np.sqrt(M * (2/R0 - ainv))
                  # Finally, calculate the angle of motion. angle = tan(dy/dx)
                  # $dy/dx = ((dr/d\phi) sin(\phi) + r \cos(\phi))/((dr/d\phi) cos(\phi) - r \sin(\phi))$
                  dy = R0/r0 * e * np.sin(phi)**2 + np.cos(phi)
                  dx = R0/r0 * e * np.sin(phi) * np.cos(phi) - np.sin(phi)
                  vangle = np.arctan2(dy, dx)
              except: raise Exception('''The orbital parameters cannot be satisfied.
                  For elliptical orbits check that R0 is posible (<rmax).''')

              # We now need the actual motion of each body
              R_vec = np.array([[R0*np.cos(phi), R0*np.sin(phi), 0.]])
              V_vec = np.array([[v*np.cos(vangle), v*np.sin(vangle), 0.]])

              galaxy1.r_vec += m2/M * R_vec
              galaxy1.v_vec += m2/M * V_vec
              galaxy2.r_vec += -m1/M * R_vec
              galaxy2.v_vec += -m1/M * V_vec

              # Explicitely add the galaxies to the simulation
              self.add(galaxy1)
              self.add(galaxy2)

              if self.verbose: print('Galaxies set in orbit.')

  ### Methods

  `def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs)`

    Constructor for the Simulation class.

    ## Arguments

    `dt` :`float`
      timestep of the simulation
    `n` :`int`
      current timestep. Initialized as n=0.
    `soft` :`float`
      softening length used by the simulation.
    `verbose` :`bool`
      When True progress statements will be printed.
    `CONFIG` :`dict`
      configuration file used to create the simulation.
    `method` :`string`
      Optional. Algorithm to use when computing the gravitational forces. One of 'bruteForce', 'bruteForce_numba', 'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'.
    Source code

        def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs):
            """Constructor for the Simulation class.

            Arguments:
                dt (float): timestep of the simulation
                n (int): current timestep. Initialized as n=0.
                soft (float): softening length used by the simulation.
                verbose (bool): When True progress statements will be printed.
                CONFIG (dict): configuration file used to create the simulation.
                method (string): Optional. Algorithm to use when computing the
                    gravitational forces. One of 'bruteForce', 'bruteForce_numba',
                    'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'.
            """
            self.n = 0
            self.t = 0
            self.dt = dt
            self.soft = soft
            self.verbose = verbose
            self.CONFIG = CONFIG
            # Initialize empty arrays for all necessary properties
            self.r_vec = np.empty((0, 3))
            self.v_vec = np.empty((0, 3))
            self.a_vec = np.empty((0, 3))
            self.mass = np.empty((0,))
            self.type = np.empty((0, 2))
            algorithms = {
                'bruteForce': acceleration.bruteForce,
                'bruteForceNumba': acceleration.bruteForceNumba,
                'bruteForceNumbaOptimized': acceleration.bruteForceNumbaOptimized,
                'bruteForceCPP': acceleration.bruteForceCPP,
                'barnesHutCPP': acceleration.barnesHutCPP
            }
            try:
                self.acceleration = algorithms[method]
            except: raise Exception("Method '{}' unknown".format(method))

  `def add(self, body)`

    Add a body to the simulation. It must expose the public attributes body.r_vec, body.v_vec, body.a_vec, body.type, body.mass.

    ## Arguments

    `body`
      Object to be added to the simulation (e.g. a Galaxy object)
    Source code

        def add(self, body):
            """Add a body to the simulation. It must expose the public attributes
               body.r_vec, body.v_vec, body.a_vec, body.type, body.mass.

            Arguments:
                body: Object to be added to the simulation (e.g. a Galaxy object)
            """
            # Extend all relevant attributes by concatenating the body
            for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
                simattr, bodyattr = getattr(self, name), getattr(body, name)
                setattr(self, name, np.concatenate([simattr, bodyattr], axis=0))
            # Order based on mass
            order = np.argsort(-self.mass)
            for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
                setattr(self, name, getattr(self, name)[order])

            # Update the list of objects to keep track of
            self.tracks = np.empty((np.sum(self.type[:,0]=='center'), 0, 3))

  `def project(self, theta, phi, view=0)`

    Projects the 3D simulation onto a plane as viewed from the direction described by the (theta, phi, view). Angles in radians. (This is used by the simulated annealing algorithm)

    ## Parameters

    `theta` :`float`
      polar angle.
    `phi` :`float`
      azimuthal angle.
    `view` :`float`
      rotation along line of sight.
    Source code

        def project(self, theta, phi, view=0):
            """Projects the 3D simulation onto a plane as viewed from the
               direction described by the (theta, phi, view). Angles in radians.
               (This is used by the simulated annealing algorithm)

            Parameters:
                theta (float): polar angle.
                phi (float): azimuthal angle.
                view (float): rotation along line of sight.
            """
            M1 = np.array([[np.cos(phi), np.sin(phi), 0],
                           [-np.sin(phi), np.cos(phi), 0],
                           [0, 0, 1]])
            M2 = np.array([[1, 0, 0],
                           [0, np.cos(theta), np.sin(theta)],
                           [0, -np.sin(theta), np.cos(theta)]])
            M3 = np.array([[np.cos(view), np.sin(view), 0],
                           [-np.sin(view), np.cos(view), 0],
                           [0, 0, 1]])

            M = np.matmul(M1, np.matmul(M2, M3)) # combine rotations
            r = np.tensordot(self.r_vec, M, axes=[1, 0])

            return r[:,0:2]

  `def run(self, tmax, saveEvery, outputFolder, **kwargs)`

    Run the galactic simulation.

    ## Attributes

    `tmax` :`float`
      Time to which the simulation will run to. This is measured here since the start of the simulation, not since pericenter.
    `saveEvery` :`int`
      The state is saved every saveEvery steps.
    `outputFolder` :`string`
      It will be saved to /data/outputFolder/
    Source code

        def run(self, tmax, saveEvery, outputFolder, **kwargs):
            """Run the galactic simulation.

            Attributes:
                tmax (float): Time to which the simulation will run to.
                    This is measured here since the start of the simulation,
                    not since pericenter.
                saveEvery (int): The state is saved every saveEvery steps.
                outputFolder (string): It will be saved to /data/outputFolder/
            """
            # When the simulation starts, intialize self.rprev_vec
            self.rprev_vec = self.r_vec - self.v_vec * self.dt
            if self.verbose: print('Simulation starting. Bon voyage!')
            while(self.t < tmax):
                self.step()
                if(self.n % saveEvery == 0):
                    self.save('data/{}'.format(outputFolder))

            print('Simulation complete.')

  `def save(self, outputFolder)`

    Save the state of the simulation to the outputFolder. Two files are saved: sim{self.n}.pickle: serializing the state. sim{self.n}.png: a simplified 2D plot of x, y.

    Source code

        def save(self, outputFolder):
            """Save the state of the simulation to the outputFolder.
               Two files are saved:
                    sim{self.n}.pickle: serializing the state.
                    sim{self.n}.png: a simplified 2D plot of x, y.
            """
            # Create the output folder if it doesn't exist
            if not os.path.exists(outputFolder): os.makedirs(outputFolder)

            # Compute some useful quantities
            # v_vec is not required by the integrator, but useful
            self.v_vec = (self.r_vec - self.rprev_vec)/self.dt
            self.t = self.n * self.dt # prevents numerical rounding errors

            # Serialize state
            file = open(outputFolder+'/data{}.pickle'.format(self.n), "wb")
            pickle.dump({'r_vec': self.r_vec, 'v_vec': self.v_vec,
                         'type': self.type, 'mass': self.mass,
                         'CONFIG': self.CONFIG, 't': self.t,
                         'tracks': self.tracks}, file)

            # Save simplified plot of the current state.
            # Its main use is to detect ill-behaved situations early on.
            fig = plt.figure()
            plt.xlim(-5, 5); plt.ylim(-5, 5); plt.axis('equal')
            # Dark halo is plotted in red, disk in blue, bulge in green
            PLTCON = [('dark', 'r', 0.3), ('disk', 'b', 1.0), ('bulge', 'g', 0.5)]
            for type_, c, a in PLTCON:
                plt.scatter(self.r_vec[self.type[:,0]==type_][:,0],
                    self.r_vec[self.type[:,0]==type_][:,1], s=0.1, c=c, alpha=a)
            # Central mass as a magenta star
            plt.scatter(self.r_vec[self.type[:,0]=='center'][:,0],
                self.r_vec[self.type[:,0]=='center'][:,1], s=100, marker="*", c='m')
            # Save to png file
            fig.savefig(outputFolder+'/sim{}.png'.format(self.n), dpi=150)
            plt.close(fig)

  `def setOrbit(self, galaxy1, galaxy2, e, rmin, R0)`

    Sets the two galaxies galaxy1, galaxy2 in an orbit.

    ## Parameters

    `galaxy1` :`Galaxy`
      1st galaxy (main)
    `galaxy2` :`Galaxy`
      2nd galaxy (companion)
    `e`
      eccentricity of the orbit
    `rmin`
      distance of closest approach
    `R0`
      initial separation
    Source code

        def setOrbit(self, galaxy1, galaxy2, e, rmin, R0):
            """Sets the two galaxies galaxy1, galaxy2 in an orbit.

            Parameters:
                galaxy1 (Galaxy): 1st galaxy (main)
                galaxy2 (Galaxy): 2nd galaxy (companion)
                e: eccentricity of the orbit
                rmin: distance of closest approach
                R0: initial separation
            """
            m1, m2 = np.sum(galaxy1.mass), np.sum(galaxy2.mass)
            # Relevant formulae:
            # $r_0 = r (1 + e) \cos(\phi)$, where $r_0$ ($\neq R_0$) is the semi-latus rectum
            # $r_0 = r_\textup{min} (1 + e)$
            # $v^2 = G M (2/r - 1/a)$, where a is the semimajor axis

            # Solve the reduced two-body problem with reduced mass $\mu$ (mu)
            M = m1 + m2
            r0 = rmin * (1 + e)
            try:
                phi = np.arccos((r0/R0 - 1) / e) # inverting the orbit equation
                phi = -np.abs(phi) # Choose negative (incoming) angle
                ainv = (1 - e) / rmin # ainv = $1/a$, as a may be infinite
                v = np.sqrt(M * (2/R0 - ainv))
                # Finally, calculate the angle of motion. angle = tan(dy/dx)
                # $dy/dx = ((dr/d\phi) sin(\phi) + r \cos(\phi))/((dr/d\phi) cos(\phi) - r \sin(\phi))$
                dy = R0/r0 * e * np.sin(phi)**2 + np.cos(phi)
                dx = R0/r0 * e * np.sin(phi) * np.cos(phi) - np.sin(phi)
                vangle = np.arctan2(dy, dx)
            except: raise Exception('''The orbital parameters cannot be satisfied.
                For elliptical orbits check that R0 is posible (<rmax).''')

            # We now need the actual motion of each body
            R_vec = np.array([[R0*np.cos(phi), R0*np.sin(phi), 0.]])
            V_vec = np.array([[v*np.cos(vangle), v*np.sin(vangle), 0.]])

            galaxy1.r_vec += m2/M * R_vec
            galaxy1.v_vec += m2/M * V_vec
            galaxy2.r_vec += -m1/M * R_vec
            galaxy2.v_vec += -m1/M * V_vec

            # Explicitely add the galaxies to the simulation
            self.add(galaxy1)
            self.add(galaxy2)

            if self.verbose: print('Galaxies set in orbit.')

  `def step(self)`

    Perform a single step of the simulation. Makes use of a 4th order Verlet integrator.

    Source code

        def step(self):
            """Perform a single step of the simulation.
               Makes use of a 4th order Verlet integrator.
            """
            # Calculate the acceleration
            self.a_vec = self.acceleration(self.r_vec, self.mass, soft=self.soft)
            # Update the state using the Verlet algorithm
            # (A custom algorithm is written mainly for learning purposes)
            self.r_vec, self.rprev_vec = (2*self.r_vec - self.rprev_vec
                + self.a_vec * self.dt**2, self.r_vec)
            self.n += 1
            # Update tracks
            self.tracks = np.concatenate([self.tracks,
                self.r_vec[self.type[:,0]=='center'][:,np.newaxis]], axis=1)