<commit_before>"""
This package implements an event-based pool physics simulator based on the paper
(available at http://web.stanford.edu/group/billiards/AnEventBasedPoolPhysicsSimulator.pdf): ::

  AN EVENT-BASED POOL PHYSICS SIMULATOR
  <NAME>, <NAME>
  DOI: 10.1007/11922155_19 · Source: DBLP
  Conference: Advances in Computer Games, 11th International Conference,
  Taipei, Taiwan, September 6-9, 2005.

"""
import logging
_logger = logging.getLogger(__name__)
from bisect import bisect
from itertools import chain
from time import perf_counter

import numpy as np


from ..table import PoolTable
from .events import (CueStrikeEvent,
                     BallEvent,
                     BallStationaryEvent,
                     BallRestEvent,
                     BallMotionEvent,
                     BallCollisionEvent,
                     MarlowBallCollisionEvent,
                     SimpleBallCollisionEvent,
                     RailCollisionEvent)


PIx2 = np.pi*2
RAD2DEG = 180/np.pi
INCH2METER = 0.0254


class PoolPhysics(object):
    _ZERO_TOLERANCE = 1e-8
    _ZERO_TOLERANCE_SQRD = _ZERO_TOLERANCE**2
    _IMAG_TOLERANCE = 1e-8
    _IMAG_TOLERANCE_SQRD = _IMAG_TOLERANCE**2
    _BALL_MASS = 0.17
    _BALL_RADIUS = 1.125*INCH2METER
    _FR_COEFF_ROLLING = 0.016
    _FR_COEFF_SLIDING = 0.2
    _FR_COEFF_SPINING = 0.044
    _FR_COEFF_BALL2BALL = 0.06
    _GRAV_ACCEL = 9.81
    def __init__(self,
                 num_balls=16,
                 ball_mass=_BALL_MASS,
                 ball_radius=_BALL_RADIUS,
                 mu_r=_FR_COEFF_ROLLING,
                 mu_sp=_FR_COEFF_SPINING,
                 mu_s=_FR_COEFF_SLIDING,
                 mu_b=_FR_COEFF_BALL2BALL,
                 g=_GRAV_ACCEL,
                 balls_on_table=None,
                 ball_positions=None,
                 ball_collision_model="simple",
                 table=None,
                 enable_sanity_check=True,
                 enable_occlusion=True,
                 realtime=False,
                 collision_search_time_limit=0.2/90,
                 collision_search_time_forward=0.2,
                 **kwargs):
        r"""
        Pool physics simulator

        :param mu_r:  :math:`\mu_r`,    rolling friction coefficient
        :param mu_sp: :math:`\mu_{sp}`, spinning friction coefficient
        :param mu_s:  :math:`\mu_s`,    sliding friction coefficient
        :param mu_b:  :math:`\mu_b`,    ball-to-ball collision friction coefficient
        :param c_b:   :math:`c_b`,      ball material's speed of sound
        :param E_Y_b: :math:`{E_Y}_b`,  ball material's Young's modulus
        :param g:     :math:`g`,        downward acceleration due to gravity
        """
        if ball_collision_model == 'simple':
            self._ball_collision_event_class = SimpleBallCollisionEvent
        elif ball_collision_model == 'marlow':
            self._ball_collision_event_class = MarlowBallCollisionEvent
        else:
            raise Exception('dont know that collision model!')
        self.num_balls = num_balls
        # allocate for lower-level memory management:
        self._BALL_MOTION_EVENTS = [BallMotionEvent(0.0, i, float('inf'),
                                                    a=np.zeros((3,3), dtype=np.float64),
                                                    b=np.zeros((2,3), dtype=np.float64))
                                    for i in range(self.num_balls)]
        self._BALL_REST_EVENTS = [BallRestEvent(0.0, i, r=np.zeros(3, dtype=np.float64))
                                  for i in range(self.num_balls)]
        if table is None:
            table = PoolTable(num_balls=num_balls, ball_radius=ball_radius)
        self.table = table
        if balls_on_table is None:
            balls_on_table = range(self.num_balls)
        self.ball_mass = ball_mass
        self.ball_radius = ball_radius
        self.ball_diameter = 2*ball_radius
        self.ball_I = 2/5 * ball_mass * ball_radius**2 # moment of inertia
        self.mu_r = mu_r
        self.mu_sp = mu_sp
        self.mu_s = mu_s
        self.mu_b = mu_b
        self.g = g
        self.t = 0.0
        self._balls_on_table = balls_on_table
        self._balls_at_rest = set(balls_on_table)
        self._on_table = np.array(self.num_balls * [False])
        self._on_table[np.array(balls_on_table, dtype=np.int64)] = True
        self._realtime = realtime
        self._collision_search_time_limit = collision_search_time_limit
        self._collision_search_time_forward = collision_search_time_forward
        self._enable_occlusion = enable_occlusion
        self._enable_sanity_check = enable_sanity_check
        self._p = np.empty(5, dtype=np.float64)
        self._mask = np.array(4*[True])
        self._a_ij = np.zeros((self.num_balls, 3), dtype=np.float64)
        self._a_ij_mag = np.zeros((self.num_balls, self.num_balls), dtype=np.float64)
        self._r_ij = np.zeros((self.num_balls, self.num_balls, 3), dtype=np.float64)
        self._r_ij_mag = np.zeros((self.num_balls, self.num_balls), dtype=np.float64)
        self._theta_ij = np.zeros((self.num_balls, self.num_balls), dtype=np.float64)
        self._psi_ij = np.zeros((self.num_balls, self.num_balls), dtype=np.float64)
        self._occ_ij = np.array(self.num_balls*[self.num_balls*[False]])
        self._velocity_meshes = None
        self._angular_velocity_meshes = None
        self.reset(ball_positions=ball_positions, balls_on_table=balls_on_table)

    def reset(self, ball_positions=None, balls_on_table=None):
        """
        Reset the state of the balls to at rest, at the specified positions.
        """
        self._ball_motion_events = {}
        if balls_on_table is None:
            balls_on_table = range(self.num_balls)
        self.balls_on_table = balls_on_table
        if ball_positions is None:
            ball_positions = self.table.calc_racked_positions()
        else:
            ball_positions = ball_positions[self.balls_on_table]
        self.t = 0.0
        for e in self._BALL_MOTION_EVENTS:
            e._a[:] = 0
            e._b[:] = 0
            e.t = self.t
            e.T = 0.0
        for e, r in zip(self._BALL_REST_EVENTS, ball_positions):
            e._r[:] = r
            e.t = self.t
            e.T = float('inf')
        self.ball_events = {i: [self._BALL_REST_EVENTS[i]]
                            for i in self.balls_on_table}
        self.events = list(chain.from_iterable(self.ball_events.values()))
        self._a_ij[:] = 0
        self._a_ij_mag[:] = 0
        # update occlusion buffers:
        self._occ_ij[:] = False
        self._occ_ij[range(self.num_balls), range(self.num_balls)] = True
        self._r_ij[:] = 0
        self._r_ij_mag[:] = 0
        self._psi_ij[:] = 0<commit_msg>fix balls_on_table<commit_after>"""
This package implements an event-based pool physics simulator based on the paper
(available at http://web.stanford.edu/group/billiards/AnEventBasedPoolPhysicsSimulator.pdf): ::

  AN EVENT-BASED POOL PHYSICS SIMULATOR
  <NAME>, <NAME>
  DOI: 10.1007/11922155_19 · Source: DBLP
  Conference: Advances in Computer Games, 11th International Conference,
  Taipei, Taiwan, September 6-9, 2005.

"""
import logging
_logger = logging.getLogger(__name__)
from bisect import bisect
from itertools import chain
from time import perf_counter

import numpy as np


from ..table import PoolTable
from .events import (CueStrikeEvent,
                     BallEvent,
                     BallStationaryEvent,
                     BallRestEvent,
                     BallMotionEvent,
                     BallCollisionEvent,
                     MarlowBallCollisionEvent,
                     SimpleBallCollisionEvent,
                     RailCollisionEvent)


PIx2 = np.pi*2
RAD2DEG = 180/np.pi
INCH2METER = 0.0254


class PoolPhysics(object):
    _ZERO_TOLERANCE = 1e-8
    _ZERO_TOLERANCE_SQRD = _ZERO_TOLERANCE**2
    _IMAG_TOLERANCE = 1e-8
    _IMAG_TOLERANCE_SQRD = _IMAG_TOLERANCE**2
    _BALL_MASS = 0.17
    _BALL_RADIUS = 1.125*INCH2METER
    _FR_COEFF_ROLLING = 0.016
    _FR_COEFF_SLIDING = 0.2
    _FR_COEFF_SPINING = 0.044
    _FR_COEFF_BALL2BALL = 0.06
    _GRAV_ACCEL = 9.81
    def __init__(self,
                 num_balls=16,
                 ball_mass=_BALL_MASS,
                 ball_radius=_BALL_RADIUS,
                 mu_r=_FR_COEFF_ROLLING,
                 mu_sp=_FR_COEFF_SPINING,
                 mu_s=_FR_COEFF_SLIDING,
                 mu_b=_FR_COEFF_BALL2BALL,
                 g=_GRAV_ACCEL,
                 balls_on_table=None,
                 ball_positions=None,
                 ball_collision_model="simple",
                 table=None,
                 enable_sanity_check=True,
                 enable_occlusion=True,
                 realtime=False,
                 collision_search_time_limit=0.2/90,
                 collision_search_time_forward=0.2,
                 **kwargs):
        r"""
        Pool physics simulator

        :param mu_r:  :math:`\mu_r`,    rolling friction coefficient
        :param mu_sp: :math:`\mu_{sp}`, spinning friction coefficient
        :param mu_s:  :math:`\mu_s`,    sliding friction coefficient
        :param mu_b:  :math:`\mu_b`,    ball-to-ball collision friction coefficient
        :param c_b:   :math:`c_b`,      ball material's speed of sound
        :param E_Y_b: :math:`{E_Y}_b`,  ball material's Young's modulus
        :param g:     :math:`g`,        downward acceleration due to gravity
        """
        if ball_collision_model == 'simple':
            self._ball_collision_event_class = SimpleBallCollisionEvent
        elif ball_collision_model == 'marlow':
            self._ball_collision_event_class = MarlowBallCollisionEvent
        else:
            raise Exception('dont know that collision model!')
        self.num_balls = num_balls
        # allocate for lower-level memory management:
        self._BALL_MOTION_EVENTS = [BallMotionEvent(0.0, i, float('inf'),
                                                    a=np.zeros((3,3), dtype=np.float64),
                                                    b=np.zeros((2,3), dtype=np.float64))
                                    for i in range(self.num_balls)]
        self._BALL_REST_EVENTS = [BallRestEvent(0.0, i, r=np.zeros(3, dtype=np.float64))
                                  for i in range(self.num_balls)]
        if table is None:
            table = PoolTable(num_balls=num_balls, ball_radius=ball_radius)
        self.table = table
        if balls_on_table is None:
            balls_on_table = range(self.num_balls)
        self.ball_mass = ball_mass
        self.ball_radius = ball_radius
        self.ball_diameter = 2*ball_radius
        self.ball_I = 2/5 * ball_mass * ball_radius**2 # moment of inertia
        self.mu_r = mu_r
        self.mu_sp = mu_sp
        self.mu_s = mu_s
        self.mu_b = mu_b
        self.g = g
        self.t = 0.0
        self._balls_on_table = balls_on_table
        self._balls_at_rest = set(balls_on_table)
        self._on_table = np.array(self.num_balls * [False])
        self._on_table[np.array(balls_on_table, dtype=np.int64)] = True
        self._realtime = realtime
        self._collision_search_time_limit = collision_search_time_limit
        self._collision_search_time_forward = collision_search_time_forward
        self._enable_occlusion = enable_occlusion
        self._enable_sanity_check = enable_sanity_check
        self._p = np.empty(5, dtype=np.float64)
        self._mask = np.array(4*[True])
        self._a_ij = np.zeros((self.num_balls, 3), dtype=np.float64)
        self._a_ij_mag = np.zeros((self.num_balls, self.num_balls), dtype=np.float64)
        self._r_ij = np.zeros((self.num_balls, self.num_balls, 3), dtype=np.float64)
        self._r_ij_mag = np.zeros((self.num_balls, self.num_balls), dtype=np.float64)
        self._theta_ij = np.zeros((self.num_balls, self.num_balls), dtype=np.float64)
        self._psi_ij = np.zeros((self.num_balls, self.num_balls), dtype=np.float64)
        self._occ_ij = np.array(self.num_balls*[self.num_balls*[False]])
        self._velocity_meshes = None
        self._angular_velocity_meshes = None
        self.reset(ball_positions=ball_positions, balls_on_table=balls_on_table)

    def reset(self, ball_positions=None, balls_on_table=None):
        """
        Reset the state of the balls to at rest, at the specified positions.
        """
        self._ball_motion_events = {}
        if balls_on_table is None:
            balls_on_table = range(self.num_balls)
        self.balls_on_table = balls_on_table
        if ball_positions is None:
            ball_positions = self.table.calc_racked_positions()[self.balls_on_table]
        else:
            ball_positions = ball_positions[self.balls_on_table]
        self.t = 0.0
        for e in self._BALL_MOTION_EVENTS:
            e._a[:] = 0
            e._b[:] = 0
            e.t = self.t
            e.T = 0.0
        for ii, i in enumerate(self.balls_on_table):
            e = self._BALL_REST_EVENTS[i]
            e._r[:] = ball_positions[ii]
            e.t = self.t
            e.T = float('inf')
        for e, r in zip(self._BALL_REST_EVENTS, ball_positions):
            e._r[:] = r
        self.ball_events = {i: [self._BALL_REST_EVENTS[i]]
                            for i in self.balls_on_table}
        self.events = list(chain.from_iterable(self.ball_events.values()))
        self._a_ij[:] = 0
        self._a_ij_mag[:] = 0
        # update occlusion buffers:
        self._occ_ij[:] = False
        self._occ_ij[range(self.num_balls), range(self.num_balls)] = True
        self._r_ij[:] = 0
        self._r_ij_mag[:] = 0
        self._psi_ij[:] = 0
{'fraction_non_alphanumeric': 0.0931336742147553, 'fraction_numerical': 0.018699780861943023, 'mean_word_length': 2.0842532101824736, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2103020', 'n_tokens_mistral': 4716, 'n_tokens_neox': 4426, 'n_words': 829}