def get_ball_ball_collision_coeffs_fast(rvw1, rvw2, s1, s2, mu1, mu2, m1, m2, g1, g2, R):
"""Get the quartic coeffs required to determine the ball-ball collision time (just-in-time compiled)
Notes
=====
- Speed comparison in pooltool/tests/speed/get_ball_ball_collision_coeffs.py
"""
c1x, c1y = rvw1[0, 0], rvw1[0, 1]
c2x, c2y = rvw2[0, 0], rvw2[0, 1]
if s1 == const.spinning or s1 == const.pocketed or s1 == const.stationary:
a1x, a1y, b1x, b1y = 0, 0, 0, 0
else:
phi1 = utils.angle_fast(rvw1[1])
v1 = np.linalg.norm(rvw1[1])
u1 = (np.array([1,0,0], dtype=np.float64)
if s1 == const.rolling
else utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw1, R)), -phi1))
K1 = -0.5*mu1*g1
cos_phi1 = np.cos(phi1)
sin_phi1 = np.sin(phi1)
a1x = K1*(u1[0]*cos_phi1 - u1[1]*sin_phi1)
a1y = K1*(u1[0]*sin_phi1 + u1[1]*cos_phi1)
b1x = v1*cos_phi1
b1y = v1*sin_phi1
if s2 == const.spinning or s2 == const.pocketed or s2 == const.stationary:
a2x, a2y, b2x, b2y = 0., 0., 0., 0.
else:
phi2 = utils.angle_fast(rvw2[1])
v2 = np.linalg.norm(rvw2[1])
u2 = (np.array([1,0,0], dtype=np.float64)
if s2 == const.rolling
else utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw2, R)), -phi2))
K2 = -0.5*mu2*g2
cos_phi2 = np.cos(phi2)
sin_phi2 = np.sin(phi2)
a2x = K2*(u2[0]*cos_phi2 - u2[1]*sin_phi2)
a2y = K2*(u2[0]*sin_phi2 + u2[1]*cos_phi2)
b2x = v2*cos_phi2
b2y = v2*sin_phi2
Ax, Ay = a2x-a1x, a2y-a1y
Bx, By = b2x-b1x, b2y-b1y
Cx, Cy = c2x-c1x, c2y-c1y
a = Ax**2 + Ay**2
b = 2*Ax*Bx + 2*Ay*By
c = Bx**2 + 2*Ax*Cx + 2*Ay*Cy + By**2
d = 2*Bx*Cx + 2*By*Cy
e = Cx**2 + Cy**2 - 4*R**2
return a, b, c, d, e
def get_ball_ball_collision_coeffs(rvw1, rvw2, s1, s2, mu1, mu2, m1, m2, g1, g2, R):
"""Get the quartic coeffs required to determine the ball-ball collision time"""
c1x, c1y = rvw1[0, 0], rvw1[0, 1]
c2x, c2y = rvw2[0, 0], rvw2[0, 1]
if s1 in const.nontranslating:
a1x, a1y, b1x, b1y = 0, 0, 0, 0
else:
phi1 = utils.angle_fast(rvw1[1])
v1 = np.linalg.norm(rvw1[1])
u1 = (np.array([1,0,0])
if s1 == const.rolling
else utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw1, R)), -phi1))
K1 = -0.5*mu1*g1
cos_phi1 = np.cos(phi1)
sin_phi1 = np.sin(phi1)
a1x = K1*(u1[0]*cos_phi1 - u1[1]*sin_phi1)
a1y = K1*(u1[0]*sin_phi1 + u1[1]*cos_phi1)
b1x = v1*cos_phi1
b1y = v1*sin_phi1
if s2 in const.nontranslating:
a2x, a2y, b2x, b2y = 0, 0, 0, 0
else:
phi2 = utils.angle_fast(rvw2[1])
v2 = np.linalg.norm(rvw2[1])
u2 = (np.array([1,0,0])
if s2 == const.rolling
else utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw2, R)), -phi2))
K2 = -0.5*mu2*g2
cos_phi2 = np.cos(phi2)
sin_phi2 = np.sin(phi2)
a2x = K2*(u2[0]*cos_phi2 - u2[1]*sin_phi2)
a2y = K2*(u2[0]*sin_phi2 + u2[1]*cos_phi2)
b2x = v2*cos_phi2
b2y = v2*sin_phi2
Ax, Ay = a2x-a1x, a2y-a1y
Bx, By = b2x-b1x, b2y-b1y
Cx, Cy = c2x-c1x, c2y-c1y
a = Ax**2 + Ay**2
b = 2*Ax*Bx + 2*Ay*By
c = Bx**2 + 2*Ax*Cx + 2*Ay*Cy + By**2
d = 2*Bx*Cx + 2*By*Cy
e = Cx**2 + Cy**2 - 4*R**2
return a, b, c, d, e
def evolve_slide_state(rvw, R, m, u_s, u_sp, g, t):
if t == 0:
return rvw
# Angle of initial velocity in table frame
phi = utils.angle_fast(rvw[1])
rvw_B0 = utils.coordinate_rotation_fast(rvw.T, -phi).T
# Relative velocity unit vector in ball frame
u_0 = utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw, R)), -phi)
# Calculate quantities according to the ball frame. NOTE w_B in this code block
# is only accurate of the x and y evolution of angular velocity. z evolution of
# angular velocity is done in the next block
rvw_B = np.empty((3,3), dtype=np.float64)
rvw_B[0,0] = rvw_B0[1,0]*t - 0.5*u_s*g*t**2 * u_0[0]
rvw_B[0,1] = -0.5*u_s*g*t**2 * u_0[1]
rvw_B[0,2] = 0
rvw_B[1,:] = rvw_B0[1] - u_s*g*t*u_0
rvw_B[2,:] = rvw_B0[2] - 5/2/R*u_s*g*t * utils.cross_fast(u_0, np.array([0,0,1], dtype=np.float64))
# This transformation governs the z evolution of angular velocity
rvw_B[2, 2] = rvw_B0[2, 2]
rvw_B = evolve_perpendicular_spin_state(rvw_B, R, u_sp, g, t)
# Rotate to table reference
rvw_T = utils.coordinate_rotation_fast(rvw_B.T, phi).T
rvw_T[0] += rvw[0] # Add initial ball position
return rvw_T
def get_ball_pocket_collision_coeffs_fast(rvw, s, a, b, r, mu, m, g, R):
"""Get the quartic coeffs required to determine the ball-pocket collision time (just-in-time compiled)
Notes
=====
- Speed comparison in pooltool/tests/speed/get_ball_circular_cushion_collision_coeffs.py
"""
if s == const.spinning or s == const.pocketed or s == const.stationary:
return np.inf, np.inf, np.inf, np.inf, np.inf
phi = utils.angle_fast(rvw[1])
v = np.linalg.norm(rvw[1])
u = (np.array([1,0,0], dtype=np.float64)
if s == const.rolling
else utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw, R)), -phi))
K = -0.5*mu*g
cos_phi = np.cos(phi)
sin_phi = np.sin(phi)
ax = K*(u[0]*cos_phi - u[1]*sin_phi)
ay = K*(u[0]*sin_phi + u[1]*cos_phi)
bx, by = v*cos_phi, v*sin_phi
cx, cy = rvw[0, 0], rvw[0, 1]
A = 0.5 * (ax**2 + ay**2)
B = ax*bx + ay*by
C = ax*(cx-a) + ay*(cy-b) + 0.5*(bx**2 + by**2)
D = bx*(cx-a) + by*(cy-b)
E = 0.5*(a**2 + b**2 + cx**2 + cy**2 - r**2) - (cx*a + cy*b)
return A, B, C, D, E
def get_ball_linear_cushion_collision_time(rvw, s, lx, ly, l0, p1, p2, mu, m, g, R):
"""Get the time until collision between ball and linear cushion segment"""
if s in const.nontranslating:
return np.inf
phi = utils.angle_fast(rvw[1])
v = np.linalg.norm(rvw[1])
u = (np.array([1,0,0]
if s == const.rolling
else utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw, R)), -phi)))
cos_phi = np.cos(phi)
ax = -0.5*mu*g*(u[0]*np.cos(phi) - u[1]*np.sin(phi))
ay = -0.5*mu*g*(u[0]*np.sin(phi) + u[1]*np.cos(phi))
bx, by = v*np.cos(phi), v*np.sin(phi)
cx, cy = rvw[0, 0], rvw[0, 1]
A = lx*ax + ly*ay
B = lx*bx + ly*by
C1 = l0 + lx*cx + ly*cy + R*np.sqrt(lx**2 + ly**2)
C2 = l0 + lx*cx + ly*cy - R*np.sqrt(lx**2 + ly**2)
root1, root2 = utils.quadratic_fast(A,B,C1)
root3, root4 = utils.quadratic_fast(A,B,C2)
roots = np.array([
root1,
root2,
root3,
root4,
])
roots = roots[
(abs(roots.imag) <= const.tol) & \
(roots.real > const.tol)
].real
# All roots beyond this point are real and positive
for i, root in enumerate(roots):
rvw_dtau, _ = evolve_state_motion(s, rvw, R, m, mu, 1, mu, g, root)
s_score = - np.dot(p1 - rvw_dtau[0], p2 - p1) / np.dot(p2 - p1, p2 - p1)
if not (0 <= s_score <= 1):
roots[i] = np.inf
return roots.min() if len(roots) else np.inf
def resolve(self):
self.is_partial()
cue_stick, ball = self.agents
self.agent1_state_initial = (np.copy(ball.rvw), ball.s)
v, w = physics.cue_strike(ball.m, cue_stick.M, ball.R, cue_stick.V0, cue_stick.phi, cue_stick.theta, cue_stick.a, cue_stick.b)
rvw = np.array([ball.rvw[0], v, w])
s = (c.rolling
if abs(np.sum(utils.get_rel_velocity_fast(rvw, ball.R))) <= c.tol
else c.sliding)
ball.set(rvw, s)
ball.update_next_transition_event()
self.agent1_state_final = (np.copy(ball.rvw), ball.s)
def get_ball_pocket_collision_coeffs(rvw, s, a, b, r, mu, m, g, R):
"""Get the quartic coeffs required to determine the ball-pocket collision time
Parameters
==========
a : float
The x-coordinate of the pocket's center
b : float
The y-coordinate of the pocket's center
r : float
The radius of the pocket's center
mu : float
The rolling or sliding coefficient of friction. Should match the value of s
"""
if s in const.nontranslating:
return np.inf, np.inf, np.inf, np.inf, np.inf
phi = utils.angle_fast(rvw[1])
v = np.linalg.norm(rvw[1])
u = (np.array([1,0,0])
if s == const.rolling
else utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw, R)), -phi))
K = -0.5*mu*g
cos_phi = np.cos(phi)
sin_phi = np.sin(phi)
ax = K*(u[0]*cos_phi - u[1]*sin_phi)
ay = K*(u[0]*sin_phi + u[1]*cos_phi)
bx, by = v*cos_phi, v*sin_phi
cx, cy = rvw[0, 0], rvw[0, 1]
A = 0.5 * (ax**2 + ay**2)
B = ax*bx + ay*by
C = ax*(cx-a) + ay*(cy-b) + 0.5*(bx**2 + by**2)
D = bx*(cx-a) + by*(cy-b)
E = 0.5*(a**2 + b**2 + cx**2 + cy**2 - r**2) - (cx*a + cy*b)
return A, B, C, D, E
def get_slide_time_fast(rvw, R, u_s, g):
return 2*np.linalg.norm(utils.get_rel_velocity_fast(rvw, R)) / (7*u_s*g)
def get_ball_linear_cushion_collision_time_fast(rvw, s, lx, ly, l0, p1, p2, direction, mu, m, g, R):
"""Get the time until collision between ball and linear cushion segment (just-in-time compiled)
Notes
=====
- Speed comparison in pooltool/tests/speed/get_ball_circular_cushion_collision_coeffs.py
"""
if s == const.spinning or s == const.pocketed or s == const.stationary:
return np.inf
phi = utils.angle_fast(rvw[1])
v = np.linalg.norm(rvw[1])
u = (np.array([1,0,0], dtype=np.float64)
if s == const.rolling
else utils.coordinate_rotation_fast(utils.unit_vector_fast(utils.get_rel_velocity_fast(rvw, R)), -phi))
K = -0.5*mu*g
cos_phi = np.cos(phi)
sin_phi = np.sin(phi)
ax = K*(u[0]*cos_phi - u[1]*sin_phi)
ay = K*(u[0]*sin_phi + u[1]*cos_phi)
bx, by = v*cos_phi, v*sin_phi
cx, cy = rvw[0, 0], rvw[0, 1]
A = lx*ax + ly*ay
B = lx*bx + ly*by
if direction == 0:
C = l0 + lx * cx + ly * cy + R * np.sqrt(lx ** 2 + ly ** 2)
root1, root2 = utils.quadratic_fast(A,B,C)
roots = [root1, root2]
elif direction == 1:
C = l0 + lx * cx + ly * cy - R * np.sqrt(lx ** 2 + ly ** 2)
root1, root2 = utils.quadratic_fast(A,B,C)
roots = [root1, root2]
else:
C1 = l0 + lx * cx + ly * cy + R * np.sqrt(lx ** 2 + ly ** 2)
C2 = l0 + lx * cx + ly * cy - R * np.sqrt(lx ** 2 + ly ** 2)
root1, root2 = utils.quadratic_fast(A,B,C1)
root3, root4 = utils.quadratic_fast(A,B,C2)
roots = [root1, root2, root3, root4]
min_time = np.inf
for root in roots:
if np.abs(root.imag) > const.tol:
continue
if root.real <= const.tol:
continue
rvw_dtau, _ = evolve_state_motion(s, rvw, R, m, mu, 1, mu, g, root)
s_score = - np.dot(p1 - rvw_dtau[0], p2 - p1) / np.dot(p2 - p1, p2 - p1)
if not (0 <= s_score <= 1):
continue
if root.real < min_time:
min_time = root.real
return min_time