def main():
# Initialize the GUI
interface = pt.ShotViewer()
# Create a system state
table = pt.PocketTable(model_name='7_foot')
balls = pt.get_nine_ball_rack(table, ordered=True)
cue = pt.Cue(cueing_ball=balls['cue'])
# Set up a shot
shot = pt.System(cue=cue, table=table, balls=balls)
shot.cue.aim_at_ball(shot.balls['1'])
shot.cue.strike(V0=8)
# Now instead of simulating, save the system as a pickle file
filepath = pt.utils.get_temp_file_path()
shot.save(filepath)
# Ok now make a new system and attach the old state
shot2 = pt.System()
shot2.load(filepath)
# Now simulate the first
shot.simulate(continuize=True)
interface.show(shot, title='Original system state')
# Now simulate the second
shot2.simulate(continuize=True)
interface.show(shot2, title='Pickled system state')
# Now make a copy of the second. This is a 'deep' copy, and
# uses the same underlying methods that created shot2
shot3 = shot2.copy()
interface.show(shot3, title='Copied system state')
def main(args):
# Run once to compile all numba functions. By doing this,
# compilation times will be excluded in the timing.
system = pt.System(path=path)
system.simulate(continuize=False, quiet=True)
system = pt.System(path=path)
if args.type == 'time':
with pt.terminal.TimeCode():
system.simulate(continuize=False, quiet=False)
if args.type == 'profile':
with pt.utils.PProfile(args.path):
system.simulate(continuize=False)
def worker(output_queue):
while True:
# setup table, cue, and cue ball
table = pt.PocketTable(model_name='7_foot')
balls = pt.get_nine_ball_rack(table,
spacing_factor=spacing_factor,
ordered=True)
balls['cue'].rvw[0] = get_cue_pos(balls['cue'], table)
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)
try:
shot.simulate(continuize=False, quiet=True)
except Exception as e:
continue
if len(
shot.events.filter_ball(
balls['cue']).filter_type('ball-pocket')):
# Cue ball was potted. Illegal shot
output_queue.put(shot.as_dict())
continue
output_queue.put(shot.as_dict())
# Code never reaches here because worker is terminated by main thread
return
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 main(args):
if not args.no_viz:
interface = pt.ShotViewer()
while True:
# setup table, cue, and cue ball
table = pt.PocketTable(l=4, w=2)
balls = {}
balls['cue'] = place_ball('cue', balls, table)
for i in range(args.N):
balls[str(i)] = place_ball(str(i), balls, table)
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=40)
# Evolve the shot
shot = pt.System(cue=cue, table=table, balls=balls)
try:
shot.simulate(continuize=False, quiet=False)
except KeyboardInterrupt:
shot.progress.end()
break
except:
shot.progress.end()
shot.run.info("Shot calculation failed", ":(")
continue
if not args.no_viz:
interface.show(shot)
def main(args):
if not args.no_viz:
interface = pt.ShotViewer()
if args.seed:
import numpy as np
np.random.seed(args.seed)
table = pt.PocketTable(model_name='7_foot')
balls = pt.get_nine_ball_rack(table,
ordered=True,
spacing_factor=args.spacing_factor)
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=args.V0)
# Evolve the shot
shot = pt.System(cue=cue, table=table, balls=balls)
shot.simulate()
if not args.no_viz:
interface.show(shot)
if args.save:
shot.save(args.save)
def load_prev_data():
if best_break_path.exists():
shot = pt.System(path=best_break_path)
else:
shot = None
if best_break_stats.exists():
stats = pt.utils.load_pickle(best_break_stats)
else:
stats = Counter({0: 0})
return shot, stats
def main(args):
shot = pt.System(path=args.path)
shot.simulate(quiet=True)
shot.reset_balls()
with pt.terminal.TimeCode(success_msg='Trajectories simulated in: '):
shot.simulate(quiet=True)
with pt.terminal.TimeCode(success_msg='Trajectories continuized in: '):
shot.continuize(dt=1 / 60 * 2)
class Interface(pt.ShotViewer):
def __init__(self, *args, **kwargs):
pt.ShotViewer.__init__(self, *args, **kwargs)
def change_mode(self, *args, **kwargs):
with pt.terminal.TimeCode(
success_msg='Animation sequence rendered in: '):
super().change_mode(*args, **kwargs)
interface = Interface()
interface.show(shot)
def process_shots(shots, stats, break_count, session_best, best_break,
interface):
for shot_dict in shots:
shot = pt.System(d=shot_dict)
break_count += 1
if len(
shot.events.filter_ball(
shot.balls['cue']).filter_type('ball-pocket')):
# Cue ball was potted. Illegal shot
stats['scratch'] += 1
else:
# Count how many balls were potted, ignoring cue ball
numbered_balls = [
ball for ball in shot.balls.values() if ball.id != 'cue'
]
balls_potted = len(
shot.events.filter_type('ball-pocket').filter_ball(
numbered_balls))
stats[balls_potted] += 1
if balls_potted > session_best:
session_best = balls_potted
if balls_potted > best_break:
shot.continuize(dt=0.003)
shot.save(Path(__file__).parent / 'best_break.pkl')
best_break = balls_potted
if not args.no_viz:
interface.show(shot, f"{best_break} balls potted")
pt.utils.save_pickle(stats, best_break_stats)
return stats, break_count, session_best, best_break
def main(args):
if args.clear:
clear()
interface = pt.ShotViewer() if not args.no_viz else None
shot, stats = load_prev_data()
session_best = 0
best_break = max([x for x in list(stats.keys()) if x != 'scratch'])
break_count = sum(list(stats.values())) + 1
print_stats(stats, run)
if not args.no_viz and shot is not None:
interface.show(pt.System(path=best_break_path),
f"The best break so far ({best_break} balls)")
shots = []
buffer_size = 200
queue_size = args.threads * 5
manager = multiprocessing.Manager()
output_queue = manager.Queue(queue_size)
processes = []
for _ in range(args.threads):
processes.append(
multiprocessing.Process(target=worker, args=(output_queue, )))
for proc in processes:
proc.start()
while True:
try:
shot = output_queue.get()
shots.append(shot)
if buffer_size > 0 and len(shots) % buffer_size == 0:
stats, break_count, session_best, best_break = process_shots(
shots, stats, break_count, session_best, best_break,
interface)
print_stats(stats, run)
shots = []
except KeyboardInterrupt:
run.info_single('Cancelling upon user request...',
nl_before=1,
nl_after=1)
break
except Exception as worker_error:
for proc in processes:
proc.terminate()
run.info_single('Thread interrupted. Ending...',
nl_before=1,
nl_after=1)
break
for proc in processes:
proc.terminate()
shots = []
def setup_and_run():
system = pt.System(path='benchmark_short.pkl')
system.simulate(continuize=False, quiet=True)
#! /usr/bin/env python
"""For some reason, `numba_cache` in pooltool/constants.py must be set to False prior to running this script"""
import pooltool as pt
import IPython
# Run once to compile all numba functions. By doing this,
# compilation times will be excluded in the timing.
system = pt.System(path='benchmark_short.pkl')
system.simulate(continuize=False, quiet=True)
def setup_and_run():
system = pt.System(path='benchmark_short.pkl')
system.simulate(continuize=False, quiet=True)
ipython = IPython.get_ipython()
ipython.magic("timeit setup_and_run()")
#! /usr/bin/env python
import numpy as np
import pooltool as pt
# Setup a shot
table = pt.PocketTable(model_name='7_foot')
balls = {
'cue': pt.Ball('cue', xyz=(table.w / 2, table.l / 3, pt.R)),
'1': pt.Ball('1', xyz=(table.w / 4, table.l * 0.2, pt.R)),
}
cue = pt.Cue(cueing_ball=balls['cue'])
cue.set_state(phi=225, V0=2)
system = pt.System(cue=cue, table=table, balls=balls)
collection = pt.SystemCollection()
for x in np.linspace(0, 0.7, 20):
shot = system.copy()
shot.cue.set_state(b=-x)
shot.cue.strike()
shot.simulate()
collection.append(shot)
interface = pt.ShotViewer()
interface.show(collection)