# apps170608_collision.binary5_two_peaksΒΆ

import matplotlib.pyplot as plt
from matplotlib.patches import Circle
import numpy as np
import scipy.stats as st
import numpy.random as rnd
rnd.seed(100)    # Set the seed

from irfpy.util import collision as coll

lim = 200
n_particle = 500
r = 0.5

def main():
# Only 2D.

poslist = rnd.uniform(0, lim, size=[n_particle, 3])

poslist[:, 2] = 0
vellist = rnd.uniform(0.1, 0.4, size=[n_particle, 3])
vellist[:int(n_particle * 0.2)] = rnd.uniform(-0.4, -0.38, size=[int(n_particle * 0.2), 3])
vellist[:, 2] = 0

## Time series of mean, std, skew, and kutosis
vxmean = [np.mean(vellist[:, 0])]   # bulk velocity
vxstd = [np.std(vellist[:, 0])]    # thermal velocity
vxskew = [st.skew(vellist[:, 0])]
vxkurt = [st.kurtosis(vellist[:, 0])]
ncoll = [0]

snapshot(poslist, vellist, 0, ())

for t in range(1, 1000):
print('t={}'.format(t))
poslist, vellist, cpair = progress(poslist, vellist)

vxmean.append(np.mean(vellist[:, 0]))
vxstd.append(np.std(vellist[:, 0]))
vxskew.append(st.skew(vellist[:, 0]))
vxkurt.append(st.kurtosis(vellist[:, 0]))
ncoll.append(ncoll[-1] + len(cpair))

print(vxmean[-1], vxstd[-1], vxskew[-1], vxkurt[-1])

if t % 10 == 0 or t <= 100:
#        if t % 60 == 0:
snapshot(poslist, vellist, t, cpair)
timeser(vxmean, vxstd, vxskew, vxkurt, ncoll)

def timeser(mean, std, skew, kurt, ncoll):
fig = plt.figure()
ax.plot(mean)
ax.axhline(0, color='r', ls=':')
ax.set_ylabel('Mean')
ax.plot(std)
ax.set_ylabel('StD')
ax.plot(skew)
ax.axhline(0, color='r', ls=':')
ax.set_ylabel('Skewness')
ax.plot(kurt)
ax.axhline(0, color='r', ls=':')
ax.set_ylabel('Kurtsis')

ax.plot(ncoll)
ax.set_ylabel('N-coll')

fig.savefig('binary5_timeseries')
plt.close(fig)

def progress(poslist, vellist):
### Very crude way checking the collision.
coll_l = r * 2
collision_pair = []

n = n_particle
for i in range(n):
for j in range(i+1, n):
if np.linalg.norm(poslist[i, :] - poslist[j, :]) < coll_l:    #  Collision
print('  Collision between {} and {}'.format(i, j))
v1, v2 = coll.binary_collision(1, poslist[i, :], vellist[i, :], 1, poslist[j, :], vellist[j, :], coll_l)
vellist[i, :] = v1
vellist[j, :] = v2
collision_pair.append((i, j))

poslist = poslist + vellist

while poslist.max() >= lim:
poslist = np.where(poslist >= lim, poslist - lim, poslist)

while poslist.min() < 0:
poslist = np.where(poslist < 0, poslist + lim, poslist)

return poslist, vellist, collision_pair

def snapshot(poslist, vellist, t, cpair):

fig = plt.figure(figsize=(8, 14))

for i in range(poslist.shape[0]):
c = Circle([poslist[i, 0], poslist[i, 1]], r, color='b', edgecolor='none')
ax.plot([poslist[i, 0], poslist[i, 0] + vellist[i, 0]], [poslist[i, 1], poslist[i, 1] + vellist[i, 1]], 'k')

for i, j in cpair:
ax.plot([poslist[i, 0], poslist[j, 0]], [poslist[i, 1], poslist[j, 1]], 'r:')

ax.set_aspect(1)
ax.set_xlim(0, lim)
ax.set_ylim(0, lim)
ax.set_title('Position of atoms')

import scipy.stats

xvalue = np.linspace(-0.5, 0.5, 100)
normed = True

ax2.hist(vellist[:, 0], bins=xvalue, normed=normed)
m1 = np.mean(vellist[:, 0])
m2 = np.std(vellist[:, 0])
pdf = scipy.stats.norm.pdf(xvalue, loc=m1, scale=m2)
ax2.plot(xvalue, pdf)
ax2.set_xlim(-0.5, 0.5)
ax2.set_ylim(0, 10)
ax2.set_title('Normalized probability dstrbtn (vx)')

ax2.hist(vellist[:, 1], bins=xvalue, normed=normed)
m1 = np.mean(vellist[:, 1])
m2 = np.std(vellist[:, 1])
pdf = scipy.stats.norm.pdf(xvalue, loc=m1, scale=m2)
ax2.plot(xvalue, pdf)
ax2.set_xlim(-0.5, 0.5)
ax2.set_ylim(0, 10)
ax2.set_title('Normalized probability dstrbtn (vy)')

fig.savefig('binary5_t{:03d}.png'.format(t))
plt.close()

if __name__ == '__main__':
main()