def simulation_step(dt, mass, radius, loc, vel, domain):
""" returns positions and velocities of particles after one step"""
# TODO: should we have different raidus?
r12_x = squareform(pdist(loc, lambda r1, r2: r1[0] - r2[0]))
r12_y = squareform(pdist(loc, lambda r1, r2: r1[1] - r2[1]))
r12 = np.stack((r12_x, r12_y), axis=-1)
dist = np.apply_along_axis(np.linalg.norm, 2, r12)
# finding particles that are colliding
ind1, ind2 = np.where(dist < 2 * radius)
unique = (ind1 < ind2)
ind1 = ind1[unique]
ind2 = ind2[unique]
# collisions
for id1, id2 in zip(ind1, ind2):
vel[id1], vel[id2] = ec.collision_2d(vel[id1], vel[id2], r12[id1, id2],
mass[id1], mass[id2])
# advection
loc[:] = loc[:] + vel[:] * dt
# find outside domain points ...
out_left = (loc[:, 0] < domain[0][0] + radius)
out_right = (loc[:, 0] > domain[0][1] - radius)
out_top = (loc[:, 1] < domain[1][0] + radius)
out_bottom = (loc[:, 1] > domain[1][1] - radius)
# ... change their velocities
vel[out_left | out_right, 0] *= -1
vel[out_top | out_bottom, 1] *= -1
return loc, vel
def simulation_step(dt, mass, radius, loc, vel, domain):
""" calculating positions and velocities of particles after one step
Arguments:
dt: time steps
mass: masses
radius: radius for all balls
loc: locations of balls
vel: 2d velocities of balls
domain: domain for the box
Returns:
positions and velocities of particles after one step
"""
# transport
loc = transport(loc, vel, dt)
r12_x = squareform(pdist(loc, lambda r1, r2: r1[0] - r2[0]))
r12_y = squareform(pdist(loc, lambda r1, r2: r1[1] - r2[1]))
r12 = np.stack((r12_x, r12_y), axis=-1)
dist = np.apply_along_axis(np.linalg.norm, 2, r12)
# finding particles that are colliding
ind1, ind2 = np.where(dist < 2 * radius)
unique = (ind1 < ind2)
ind1 = ind1[unique]
ind2 = ind2[unique]
# collisions
for id1, id2 in zip(ind1, ind2):
vel[id1], vel[id2] = ec.collision_2d(vel[id1], vel[id2], r12[id1, id2],
mass[id1], mass[id2])
# find outside domain points ...
out_left = (loc[:, 0] < domain[0][0] + radius)
out_right = (loc[:, 0] > domain[0][1] - radius)
out_top = (loc[:, 1] < domain[1][0] + radius)
out_bottom = (loc[:, 1] > domain[1][1] - radius)
# ... change their velocities
vel[out_left | out_right, 0] *= -1
vel[out_top | out_bottom, 1] *= -1
return loc, vel
def test_collision_2d_1():
v1_f, v2_f = ec.collision_2d(v1=[1, 0], v2=[-2, 0], r12=[1, 0], m1=1, m2=1)
assert (v1_f == [-2, 0]).all()
assert (v2_f == [1, 0]).all()
def test_collision_2d_mom():
v1_i, v2_i = [3, 1], [1, -2]
m1, m2 = 1, 2
v1_f, v2_f = ec.collision_2d(v1=v1_i, v2=v2_i, r12=[0, 1], m1=m1, m2=m2)
assert (momentum([v1_i,v2_i], [m1, m2]) == momentum([v1_f, v2_f], [m1, m2])).all()
def test_collision_2d_en():
v1_i, v2_i = [3, 1], [1, -2]
m1, m2 = 1, 2
v1_f, v2_f = ec.collision_2d(v1=v1_i, v2=v2_i, r12=[0, 1], m1=m1, m2=m2)
assert E_kin([v1_i,v2_i], [m1, m2]) == E_kin([v1_f,v2_f], [m1, m2])
def test_collision_2d_2():
v1_f, v2_f = ec.collision_2d(v1=[0, 1], v2=[0, -2], r12=[0, 1], m1=1, m2=1)
assert (v1_f == [0, -2]).all()
assert (v2_f == [0, 1]).all()