#!/usr/bin/env python
# coding: utf-8

# # Assigning particles unique IDs and removing particles from the simulation
# 
# For some applications, it is useful to keep track of which particle is which, and this can get jumbled up when particles are added or removed from the simulation.  It can thefore be useful for particles to have unique IDs associated with them.
# 
# Let's set up a simple simulation with 10 bodies, and give them IDs in the order we add the particles (if you don't set them explicitly, all particle IDs default to 0):

# In[1]:


import rebound
import numpy as np

def setupSimulation(Nplanets):
    sim = rebound.Simulation()
    sim.integrator = "ias15" # IAS15 is the default integrator, so we don't need this line
    sim.add(m=1.,id=0)
    for i in range(1,Nbodies):
        sim.add(m=1e-5,x=i,vy=i**(-0.5),id=i)
    sim.move_to_com()
    return sim

Nbodies=10
sim = setupSimulation(Nbodies)
print([sim.particles[i].id for i in range(sim.N)])


# Now let's do a simple example where we do a short initial integration to isolate the particles that interest us for a longer simulation:

# In[2]:


Noutputs = 1000
xs = np.zeros((Nbodies, Noutputs))
ys = np.zeros((Nbodies, Noutputs))
times = np.linspace(0.,50*2.*np.pi, Noutputs, endpoint=False)
for i, time in enumerate(times):
    sim.integrate(time)
    xs[:,i] = [sim.particles[j].x for j in range(Nbodies)]
    ys[:,i] = [sim.particles[j].y for j in range(Nbodies)]
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
fig,ax = plt.subplots(figsize=(15,5))
for i in range(Nbodies):
    plt.plot(xs[i,:], ys[i,:])
ax.set_aspect('equal')


# At this stage, we might be interested in particles that remained within some semimajor axis range, particles that were in resonance with a particular planet, etc.  Let's imagine a simple (albeit arbitrary) case where we only want to keep particles that had $x > 0$ at the end of the preliminary integration. Let's first print out the particle ID and x position.

# In[3]:


print("ID\tx")
for i in range(Nbodies):
    print("{0}\t{1}".format(i, xs[i,-1]))


# Next, let's use the `remove()` function to filter out particle. As an argument, we pass the corresponding index in the particles array.

# In[4]:


for i in reversed(range(1,Nbodies)):
    if xs[i,-1] < 0:
        sim.remove(i)
print("Number of particles after cut = {0}".format(sim.N))
print("IDs of remaining particles = {0}".format([p.id for p in sim.particles]))


# By default, the `remove()` function removes the `i`-th particle from the `particles` array, and shifts all particles with higher indices down by 1.  This ensures that the original order in the `particles` array is preserved (e.g., to help with output).
# 
# By running through the planets in reverse order above, we are guaranteed that when a particle with index `i` gets removed, the particle replacing it doesn't need to also be removed (we already checked it).
# 
# If you have many particles and many removals (or you don't care about the ordering), you can save the reshuffling of all particles with higher indices with the flag `keepSorted=0`:

# In[5]:


sim.remove(2, keepSorted=0)
print("Number of particles after cut = {0}".format(sim.N))
print("IDs of remaining particles = {0}".format([p.id for p in sim.particles]))


# We see that the `particles` array is no longer sorted by ID.  Note that the default `keepSorted=1` only *keeps* things sorted (i.e., if they were sorted by ID to start with).  If you custom-assign IDs out of order as you add particles, the default will simply preserve the original order.
# 
# You might also have been surprised that the above `sim.remove(2, keepSorted=0)` succeeded, since there was no `id=2` left in the simulation.  That's because `remove()` takes the index in the particles array, so we removed the 3rd particle (with `id=4`).  If you'd like to remove a particle by `id`, use the `id` keyword, e.g.

# In[6]:


sim.remove(id=9)
print("Number of particles after cut = {0}".format(sim.N))
print("IDs of remaining particles = {0}".format([p.id for p in sim.particles]))