def __init__(self, path=None, cue=None, table=None, balls=None, d=None):
SystemHistory.__init__(self)
SystemRender.__init__(self)
EvolveShotEventBased.__init__(self)
if path and (cue or table or balls):
raise ConfigError(
"System :: if path provided, cue, table, and balls must be None"
)
if d and (cue or table or balls):
raise ConfigError(
"System :: if d provided, cue, table, and balls must be None")
if d and path:
raise ConfigError(
"System :: Preload a system with either `d` or `path`, not both"
)
if path:
path = Path(path)
self.load(path)
elif d:
self.load_from_dict(d)
else:
self.cue = cue
self.table = table
self.balls = balls
self.t = None
self.meta = None
def aim_at_ball(self, ball, cut=None):
"""Set phi to aim directly at a ball
Parameters
==========
ball : pooltool.objects.ball.Ball
A ball
cut : float, None
The cut angle in degrees, within [-89, 89]
"""
self.aim_at_pos(ball.rvw[0])
if cut is None:
return
if cut > 89 or cut < -89:
raise ConfigError(
"Cue.aim_at_ball :: cut must be less than 89 and more than -89"
)
# Ok a cut angle has been requested. Unfortunately, there exists no analytical function
# phi(cut), at least as far as I have been able to calculate. Instead, it is a nasty
# transcendental equation that must be solved. The gaol is to make its value 0. To do this,
# I sweep from 0 to the max possible angle with 100 values and find where the equation flips
# from positive to negative. The dphi that makes the equation lies somewhere between those
# two values, so then I do a new parameter sweep between the value that was positive and the
# value that was negative. Then I rinse and repeat this a total of 5 times.
left = True if cut < 0 else False
cut = np.abs(cut) * np.pi / 180
R = ball.R
d = np.linalg.norm(ball.rvw[0] - self.cueing_ball.rvw[0])
lower_bound = 0
upper_bound = (np.pi / 2 - np.arccos((2 * R) / d))
for _ in range(5):
dphis = np.linspace(lower_bound, upper_bound, 100)
transcendental = np.arctan(
2 * R * np.sin(cut - dphis) /
(d - 2 * R * np.cos(cut - dphis))) - dphis
for i in range(len(transcendental)):
if transcendental[i] < 0:
lower_bound = dphis[i - 1] if i > 0 else 0
upper_bound = dphis[i]
dphi = dphis[i]
break
else:
raise ConfigError(
"Cue.aim_at_ball :: Wow this should never happen. The algorithm "
"that finds the cut angle needs to be looked at again, because "
"the transcendental equation could not be solved.")
self.phi = (self.phi + 180 / np.pi * (dphi if left else -dphi)) % 360
def main(args):
if not args.force:
raise ConfigError(
"Many of the unit tests are generated automatically by parsing the output of this script. "
"That means the output serves as a ground truth. By running this script, you are deciding "
"that a new ground truth should be issued, which is clearly no joke. Provide the flag --force "
"to proceed. The trajectories of the balls in this simmulation will be taken as true and used "
"to compare the identicality of future versions of the code.")
table = pt.PocketTable(model_name='7_foot')
balls = pt.get_nine_ball_rack(table, ordered=True)
cue = pt.Cue(cueing_ball=balls['cue'])
# Aim at the head ball then strike the cue ball
cue.aim_at_ball(balls['1'])
cue.strike(V0=8)
# Evolve the shot
shot = pt.System(cue=cue, table=table, balls=balls)
shot.simulate(continuize=True, dt=0.01)
# Visualize the shot
interface = pt.ShotViewer()
interface.show(shot, 'This is the new benchmark.')
# Save the shot
output_dir = Path(pt.__file__).parent / 'tests' / 'data'
output_dir.mkdir(exist_ok=True)
shot.save(output_dir / 'benchmark.pkl')
def __init__(self,
ball_id,
m=None,
R=None,
u_s=None,
u_r=None,
u_sp=None,
g=None,
e_c=None,
f_c=None,
rel_model_path=None,
xyz=None,
initial_orientation=None):
"""Initialize a ball
Parameters
==========
rel_model_path : str
path should be relative to pooltool/models/balls directory
"""
self.id = ball_id
if not (isinstance(self.id, int) or isinstance(self.id, str)):
raise ConfigError("ball_id must be integer or string")
# physical properties
self.m = m or c.m
self.R = R or c.R
self.I = 2 / 5 * self.m * self.R**2
self.g = g or c.g
# felt properties
self.u_s = u_s or c.u_s
self.u_r = u_r or c.u_r
self.u_sp = u_sp or c.u_sp
# restitution properties
self.e_c = e_c or c.e_c
self.f_c = f_c or c.f_c
self.t = 0
self.s = c.stationary
x, y, z = xyz if xyz is not None else (np.nan, np.nan, np.nan)
self.rvw = np.array([
[x, y, z], # positions (r)
[0, 0, 0], # velocities (v)
[0, 0, 0]
]) # angular velocities (w)
self.update_next_transition_event()
self.history = BallHistory()
self.history_cts = BallHistory()
self.events = Events()
if initial_orientation is None:
self.initial_orientation = self.get_random_orientation()
self.rel_model_path = rel_model_path
BallRender.__init__(self, rel_model_path=self.rel_model_path)
def init_cloth(self):
if not self.has_model or not ani.settings['graphics']['table']:
node = render.find('scene').attachNewNode('cloth')
path = ani.model_dir / 'table' / 'custom' / 'custom.glb'
model = loader.loadModel(panda_path(path))
model.reparentTo(node)
model.setScale(self.w, self.l, 1)
else:
path_dir = ani.model_dir / 'table' / self.name
pbr_path = path_dir / (self.name + '_pbr.glb')
standard_path = path_dir / (self.name + '.glb')
if ani.settings['graphics']['physical_based_rendering']:
path = pbr_path
if not path.exists():
path = standard_path
else:
path = standard_path
if not path.exists():
raise ConfigError(f"Couldn't find table model at {standard_path} or {pbr_path}")
node = loader.loadModel(panda_path(path))
node.reparentTo(render.find('scene'))
node.setName('cloth')
self.nodes['cloth'] = node
self.collision_nodes = {}
def set_animation(self):
if self.parallel:
self.shot_animation = Parallel()
# `max_dur` is the shot duration of the longest shot in the collection. All shots beside
# this one will have a buffer appended where the balls stay in their final state until
# the last shot finishes.
max_dur = max([shot.events[-1].time for shot in self])
# FIXME `leading_buffer` should be utilized here to sync up all shots that have cue
# trajectories such that the ball animations all start at the moment of the stick-ball
# collision
pass
for shot in self:
shot_dur = shot.events[-1].time
shot.init_shot_animation(
trailing_buffer=max_dur - shot_dur,
leading_buffer=0,
)
self.shot_animation.append(shot.shot_animation)
else:
if not self.active:
raise ConfigError(
"SystemCollectionRender.set_animation :: self.active not set"
)
self.active.init_shot_animation()
self.shot_animation = self.active.shot_animation
def append(self, system):
if len(self):
# In order to append a system, the table must be damn-near identical to existing systems
# in this collection. Otherwise we raise an error
if system.table.as_dict() != self[0].table.as_dict():
raise ConfigError(
f"Cannot append System '{system}', which has a different table than "
f"the rest of the SystemCollection")
utils.ListLike.append(self, system)
def init_collision(self, cue):
if not cue.rendered:
raise ConfigError(
"BallRender.init_collision :: `cue` must be rendered")
collision_node = self.nodes['ball'].attachNewNode(
CollisionNode(f"ball_csphere_{self.id}"))
collision_node.node().addSolid(
CollisionCapsule(0, 0, -self.R, 0, 0, self.R,
cue.tip_radius + self.R))
if ani.settings['graphics']['debug']:
collision_node.show()
self.nodes[f"ball_csphere_{self.id}"] = collision_node
def set_meta(self, meta):
"""Define any meta data for the shot
This method provides the opportunity to associate information to the system. If the
system is saved or copied, this information will be retained under the attribute `meta`.
Parameters
==========
meta : pickleable object
Any information can be stored, so long as it is pickleable.
"""
if not utils.is_pickleable(meta):
raise ConfigError(
"System.set_meta :: Cannot set unpickleable object")
self.meta = meta
def __init__(self, cushion_id, p1, p2, direction=2):
"""A linear cushion segment defined by the line between points p1 and p2
Parameters
==========
p1 : tuple
A length-3 tuple that defines a 3D point in space where the cushion segment starts
p2 : tuple
A length-3 tuple that defines a 3D point in space where the cushion segment ends
direction: 0 or 1, default 2
For most table geometries, the playing surface only exists on one side of the cushion,
so collisions only need to be checked for one direction. This direction can be specified
with either 0 or 1. To determine whether 0 or 1 should be used, please experiment
(FIXME: determine the rule). By default, both collision directions are checked, which
can be specified explicitly by passing 2, however this makes collision checks twice as
slow for event-based shot evolution algorithms.
"""
self.id = cushion_id
self.p1 = np.array(p1, dtype=np.float64)
self.p2 = np.array(p2, dtype=np.float64)
p1x, p1y, p1z = self.p1
p2x, p2y, p2z = self.p2
if p1z != p2z:
raise ValueError(f"LinearCushionSegment with id '{self.id}' has points p1 and p2 with different cushion heights (h)")
self.height = p1z
if (p2x - p1x) == 0:
self.lx = 1
self.ly = 0
self.l0 = -p1x
else:
self.lx = - (p2y - p1y) / (p2x - p1x)
self.ly = 1
self.l0 = (p2y - p1y) / (p2x - p1x) * p1x - p1y
self.normal = utils.unit_vector_fast(np.array([self.lx, self.ly, 0]))
if direction not in {0, 1, 2}:
raise ConfigError("LinearCushionSegment :: `direction` must be 0, 1, or 2.")
self.direction = direction
def is_hit(shot, clean=False):
# Which ball is in motion?
for ball in shot.balls.values():
if ball.history.s[0] in (c.rolling, c.sliding):
cue = ball.id
break
else:
raise ConfigError("three_cushion.is_hit :: no ball is in motion")
get_other_agent = lambda event: event.agents[0].id if event.agents[
0].id != cue else event.agents[1].id
get_agent_ids = lambda event: [agent.id for agent in event.agents]
first_hit = False
second_hit = False
cushion_count = 0
for event in shot.events:
if event.event_type == 'ball-cushion' and event.agents[0].id == cue:
cushion_count += 1
if not first_hit and event.event_type == 'ball-ball':
first_hit_agent = get_other_agent(event)
first_hit = True
continue
if not second_hit and event.event_type == 'ball-ball':
agents = get_agent_ids(event)
if cue not in agents:
if clean:
return False
elif get_other_agent(event) == first_hit_agent:
if clean:
return False
else:
second_hit_agent = get_other_agent(event)
second_hit = True
break
else:
return False
if cushion_count < 3:
return False
return True
def get_shot_components(shot):
# Which ball is in motion?
for ball in shot.balls.values():
if ball.history.s[0] in (c.rolling, c.sliding):
cue = ball.id
break
else:
raise ConfigError(
"three_cushion.get_shot_components :: no ball is in motion")
get_other_agent = lambda event: event.agents[0].id if event.agents[
0].id != cue else event.agents[1].id
get_agent_ids = lambda event: [agent.id for agent in event.agents]
first_hit, second_hit = False, False
shot_components = []
for event in shot.events:
if event.event_type == 'ball-cushion':
ball, cushion = event.agents
if ball.id != cue:
continue
shot_components.append(cushion.id)
if not first_hit and event.event_type == 'ball-ball':
first_hit_agent = get_other_agent(event)
shot_components.append(first_hit_agent)
first_hit = True
continue
if not second_hit and event.event_type == 'ball-ball':
agents = get_agent_ids(event)
if cue not in agents:
continue
elif get_other_agent(event) == first_hit_agent:
continue
else:
shot_components.append(get_other_agent(event))
second_hit = True
break
return tuple(shot_components)
def which_hit_first(shot):
# Which ball is in motion?
for ball in shot.balls.values():
if ball.history.s[0] in (c.rolling, c.sliding):
cue = ball.id
break
else:
raise ConfigError("three_cushion.is_hit :: no ball is in motion")
get_other_agent = lambda event: event.agents[0].id if event.agents[
0].id != cue else event.agents[1].id
get_agent_ids = lambda event: [agent.id for agent in event.agents]
for event in shot.events:
if event.event_type == 'ball-ball':
first_hit_agent = get_other_agent(event)
break
else:
return False
return first_hit_agent
def show(self, shot_or_shots=None, title=''):
if shot_or_shots is None:
# No passed shots. This is ok if self.shots has already been defined, but will complain
# otherwise
if not len(self.shots):
raise ConfigError("ShotViewer.show :: No shots passed and no shots set.")
else:
# Create a new SystemCollection based on type of shot_or_shots
if issubclass(type(shot_or_shots), System):
self.shots = SystemCollection()
self.shots.append(shot_or_shots)
elif issubclass(type(shot_or_shots), SystemCollection):
self.shots = shot_or_shots
if self.shots.active is None:
self.shots.set_active(0)
self.standby_screen.hide()
self.instructions.show()
self.create_title(title)
self.title_node.show()
self.init_help_page()
self.help_hint.hide()
self.mouse = Mouse()
self.init_system_nodes()
self.init_hud()
params = dict(
init_animations = True,
single_instance = True,
)
self.change_mode('shot', enter_kwargs=params)
self.player_cam.load_state('last_scene', ok_if_not_exists=True)
self.taskMgr.run()
def init_collision_handling(self, collision_handler):
if not ani.settings['gameplay']['cue_collision']:
return
if not self.rendered:
raise ConfigError(
"Cue.init_collision_handling :: Cue has not been rendered, "
"so collision handling cannot be initialized.")
bounds = self.get_node('cue_stick').get_tight_bounds()
x = 0
X = bounds[1][0] - bounds[0][0]
cnode = CollisionNode(f"cue_cseg")
cnode.set_into_collide_mask(0)
collision_node = self.get_node('cue_stick_model').attachNewNode(cnode)
collision_node.node().addSolid(CollisionSegment(x, 0, 0, X, 0, 0))
self.nodes['cue_cseg'] = collision_node
base.cTrav.addCollider(collision_node, collision_handler)
if ani.settings['graphics']['debug']:
collision_node.show()
def is_partial(self):
if self.partial:
raise ConfigError(f"Cannot call `{self.__class__.__name__}.resolve` when event is partial. Add agent objects.")
import pooltool as pt
from pooltool.error import ConfigError
from pooltool.system import System
from pathlib import Path
import pytest
import numpy as np
data_dir = Path(__file__).parent / 'data'
benchmark_path = data_dir / 'benchmark.pkl'
if not benchmark_path.exists():
raise ConfigError(
'benchmark.pkl missing. Run pooltool/tests/get_dataset.py to generate.'
)
shot_ref = System()
shot_ref.load(benchmark_path)
@pytest.fixture
def ref():
return shot_ref.copy()
shot_trial = shot_ref.copy()
shot_trial.simulate(continuize=True, dt=0.01)
shot_trial.reset_balls()