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