diff --git a/pooltool/ani/animate.py b/pooltool/ani/animate.py index c64b7297..dd973291 100755 --- a/pooltool/ani/animate.py +++ b/pooltool/ani/animate.py @@ -522,7 +522,8 @@ def _create_system(self): ballset=None, spacing_factor=1e-3, ) - cue = Cue(cue_ball_id=game.shot_constraints.cueball(balls)) + cue = Cue.from_game_type(game_type) + cue.cue_ball_id = game.shot_constraints.cueball(balls) shot = System(table=table, balls=balls, cue=cue) self.attach_system(shot) diff --git a/pooltool/models/cue/cue.blend b/pooltool/models/cue/cue/cue.blend similarity index 100% rename from pooltool/models/cue/cue.blend rename to pooltool/models/cue/cue/cue.blend diff --git a/pooltool/models/cue/cue.glb b/pooltool/models/cue/cue/cue.glb similarity index 100% rename from pooltool/models/cue/cue.glb rename to pooltool/models/cue/cue/cue.glb diff --git a/pooltool/models/cue/cue_texture.png b/pooltool/models/cue/cue/cue_texture.png similarity index 100% rename from pooltool/models/cue/cue_texture.png rename to pooltool/models/cue/cue/cue_texture.png diff --git a/pooltool/models/cue/cue_texture.svg b/pooltool/models/cue/cue/cue_texture.svg similarity index 100% rename from pooltool/models/cue/cue_texture.svg rename to pooltool/models/cue/cue/cue_texture.svg diff --git a/pooltool/models/cue/uv_layout.png b/pooltool/models/cue/cue/uv_layout.png similarity index 100% rename from pooltool/models/cue/uv_layout.png rename to pooltool/models/cue/cue/uv_layout.png diff --git a/pooltool/models/cue/cue_snooker/ash_wood.jpg b/pooltool/models/cue/cue_snooker/ash_wood.jpg new file mode 100644 index 00000000..5a6a02ee Binary files /dev/null and b/pooltool/models/cue/cue_snooker/ash_wood.jpg differ diff --git a/pooltool/models/cue/cue_snooker/cue.glb b/pooltool/models/cue/cue_snooker/cue.glb new file mode 100644 index 00000000..accd3d13 Binary files /dev/null and b/pooltool/models/cue/cue_snooker/cue.glb differ diff --git a/pooltool/models/cue/cue_snooker/snooker_cue.blend b/pooltool/models/cue/cue_snooker/snooker_cue.blend new file mode 100644 index 00000000..c56a55ba Binary files /dev/null and b/pooltool/models/cue/cue_snooker/snooker_cue.blend differ diff --git a/pooltool/objects/__init__.py b/pooltool/objects/__init__.py index 5fcff7f0..8526ddae 100644 --- a/pooltool/objects/__init__.py +++ b/pooltool/objects/__init__.py @@ -15,7 +15,7 @@ ) from pooltool.objects.ball.params import BallParams, PrebuiltBallParams from pooltool.objects.ball.sets import BallSet, get_ballset, get_ballset_names -from pooltool.objects.cue.datatypes import Cue, CueSpecs +from pooltool.objects.cue.datatypes import Cue, CueSpecs, PrebuiltCueSpecs from pooltool.objects.table.collection import TableName from pooltool.objects.table.components import ( CircularCushionSegment, @@ -42,6 +42,7 @@ "BallHistory", "BallOrientation", "CueSpecs", + "PrebuiltCueSpecs", "Cue", "Pocket", "LinearCushionSegment", diff --git a/pooltool/objects/cue/datatypes.py b/pooltool/objects/cue/datatypes.py index 00aa6808..e3997f5f 100755 --- a/pooltool/objects/cue/datatypes.py +++ b/pooltool/objects/cue/datatypes.py @@ -4,6 +4,9 @@ from attrs import define, evolve, field, fields_dict +from pooltool.game.datatypes import GameType +from pooltool.utils.strenum import StrEnum, auto + @define(frozen=True) class CueSpecs: @@ -27,28 +30,107 @@ class CueSpecs: end_mass: The mass of the of the cue's end. This controls the amount of deflection (squirt) that occurs when using sidespin. Lower means less deflection. It is - defined here: - https://billiards.colostate.edu/technical_proofs/new/TP_A-31.pdf. + defined here: https://drdavepoolinfo.com/technical_proofs/new/TP_A-31.pdf. """ - brand: str = field(default="Predator") - M: float = field(default=0.567) - length: float = field(default=1.4732) - tip_radius: float = field(default=0.0106045) # nickel radius - shaft_radius_at_tip: float = field( - default=0.0065 - ) # 13 mm shaft diameter at the tip - shaft_radius_at_butt: float = field(default=0.02) - end_mass: float = field(default=0.170097 / 30) + brand: str = field() + M: float = field() + length: float = field() + tip_radius: float = field() + shaft_radius_at_tip: float = field() + shaft_radius_at_butt: float = field() + end_mass: float = field() + + @classmethod + def default(cls, game_type: GameType = GameType.EIGHTBALL) -> CueSpecs: + """Return prebuilt cue specs based on game type. + + Args: + game_type: + What type of game is being played? + + Returns: + The prebuilt cue specs associated with the passed game type. + """ + return _get_default_cue_specs(game_type) - @staticmethod - def default() -> CueSpecs: - """Construct a default cue spec""" - return CueSpecs() + @classmethod + def prebuilt(cls, name: PrebuiltCueSpecs) -> CueSpecs: + """Return prebuilt cue specs based on name. - @staticmethod - def snooker() -> CueSpecs: - raise NotImplementedError() + Args: + name: + A :class:`PrebuiltCueSpecs` member. + """ + return _prebuilt_cue_specs(name) + + +class PrebuiltCueSpecs(StrEnum): + """An Enum specifying prebuilt cue specs. + + Attributes: + POOL_GENERIC: + SNOOKER_GENERIC: + BILLIARD_GENERIC: + """ + + POOL_GENERIC = auto() + SNOOKER_GENERIC = auto() + BILLIARD_GENERIC = auto() + + +CUE_MODELS: dict[PrebuiltCueSpecs, str] = { + PrebuiltCueSpecs.POOL_GENERIC: "cue", + PrebuiltCueSpecs.SNOOKER_GENERIC: "cue_snooker", + PrebuiltCueSpecs.BILLIARD_GENERIC: "cue", +} + +CUE_SPECS: dict[PrebuiltCueSpecs, CueSpecs] = { + PrebuiltCueSpecs.POOL_GENERIC: CueSpecs( + brand="Pooltool", + M=0.567, + length=1.4732, + tip_radius=0.0106045, + shaft_radius_at_tip=0.0065, + shaft_radius_at_butt=0.02, + end_mass=0.170097 / 30, + ), + PrebuiltCueSpecs.SNOOKER_GENERIC: CueSpecs( + brand="Pooltool", + M=0.478, + length=1.475, + tip_radius=0.0106045, + shaft_radius_at_tip=0.0049, + shaft_radius_at_butt=0.0124, + end_mass=0.140 / 30, + ), + # TODO: These are just copied from the pool cue specs + PrebuiltCueSpecs.BILLIARD_GENERIC: CueSpecs( + brand="Pooltool", + M=0.567, + length=1.4732, + tip_radius=0.0106045, + shaft_radius_at_tip=0.0065, + shaft_radius_at_butt=0.02, + end_mass=0.210 / 30, + ), +} + +_default_map: dict[GameType, PrebuiltCueSpecs] = { + GameType.EIGHTBALL: PrebuiltCueSpecs.POOL_GENERIC, + GameType.NINEBALL: PrebuiltCueSpecs.POOL_GENERIC, + GameType.THREECUSHION: PrebuiltCueSpecs.BILLIARD_GENERIC, + GameType.SNOOKER: PrebuiltCueSpecs.SNOOKER_GENERIC, + GameType.SUMTOTHREE: PrebuiltCueSpecs.BILLIARD_GENERIC, +} + + +def _get_default_cue_specs(game_type: GameType) -> CueSpecs: + return _prebuilt_cue_specs(_default_map[game_type]) + + +def _prebuilt_cue_specs(name: PrebuiltCueSpecs) -> CueSpecs: + return CUE_SPECS[name] @define @@ -104,6 +186,10 @@ class Cue: The ball ID of the ball being cued. specs: The cue specs. + model_name: + The name of the cue model directory under ``pooltool/models/cue/``. + + Important if rendering the cue in a scene. """ id: str = field(default="cue_stick") @@ -114,6 +200,7 @@ class Cue: b: float = field(default=0.25) cue_ball_id: str = field(default="cue") specs: CueSpecs = field(factory=CueSpecs.default) + model_name: str | None = field(default=None) def __repr__(self): lines = [ @@ -180,6 +267,23 @@ def set_state( if cue_ball_id is not None: self.cue_ball_id = cue_ball_id + @classmethod + def from_game_type(cls, game_type: GameType, id: str | None = None) -> Cue: + if game_type not in _default_map: + raise NotImplementedError( + f"There is no cue stick associated with '{game_type}'" + ) + + if id is None: + id = fields_dict(cls)["id"].default + assert id is not None + + prebuilt = _default_map[game_type] + cue = cls(id=id) + cue.specs = CueSpecs.prebuilt(prebuilt) + cue.model_name = CUE_MODELS[prebuilt] + return cue + @classmethod def default(cls) -> Cue: """Construct a cue with defaults""" diff --git a/pooltool/objects/cue/render.py b/pooltool/objects/cue/render.py index 2122b06e..ea9a1b90 100644 --- a/pooltool/objects/cue/render.py +++ b/pooltool/objects/cue/render.py @@ -57,7 +57,8 @@ def set_render_state_as_object_state(self): cue_stick.setZ(tip_offset_b * self.follow._ball.params.R) # b def init_model(self): - path = utils.panda_path(model_dir / "cue" / "cue.glb") + name = self._cue.model_name or "cue" + path = utils.panda_path(model_dir / "cue" / name / "cue.glb") cue_stick_model = Global.loader.loadModel(path) cue_stick_model.setName("cue_stick_model") diff --git a/tests/events/example_system.msgpack b/tests/events/example_system.msgpack index 36ad1082..8856fd95 100644 Binary files a/tests/events/example_system.msgpack and b/tests/events/example_system.msgpack differ diff --git a/tests/evolution/event_based/test_data/case1.msgpack b/tests/evolution/event_based/test_data/case1.msgpack index b493d40a..ef8f1e28 100644 Binary files a/tests/evolution/event_based/test_data/case1.msgpack and b/tests/evolution/event_based/test_data/case1.msgpack differ diff --git a/tests/evolution/event_based/test_data/case2.msgpack b/tests/evolution/event_based/test_data/case2.msgpack index 65e0e17d..9612ad01 100644 Binary files a/tests/evolution/event_based/test_data/case2.msgpack and b/tests/evolution/event_based/test_data/case2.msgpack differ diff --git a/tests/evolution/event_based/test_data/case3.msgpack b/tests/evolution/event_based/test_data/case3.msgpack index 0a0c61f5..806ee0f3 100644 Binary files a/tests/evolution/event_based/test_data/case3.msgpack and b/tests/evolution/event_based/test_data/case3.msgpack differ diff --git a/tests/evolution/event_based/test_data/case4.msgpack b/tests/evolution/event_based/test_data/case4.msgpack index 43c4eed9..285568bc 100644 Binary files a/tests/evolution/event_based/test_data/case4.msgpack and b/tests/evolution/event_based/test_data/case4.msgpack differ diff --git a/tests/objects/cue/test_datatypes.py b/tests/objects/cue/test_datatypes.py index a52eebe4..28b0cfea 100644 --- a/tests/objects/cue/test_datatypes.py +++ b/tests/objects/cue/test_datatypes.py @@ -1,7 +1,8 @@ import pytest from attrs.exceptions import FrozenInstanceError -from pooltool.objects.cue.datatypes import Cue +from pooltool.game.datatypes import GameType +from pooltool.objects.cue.datatypes import Cue, CueSpecs, PrebuiltCueSpecs def test_cue_copy(): @@ -20,3 +21,30 @@ def test_cue_copy(): cue.phi += 1 assert cue != copy assert cue.phi != copy.phi + + +def test_cue_specs_construction(): + # Can't instantiate without setting parameters + with pytest.raises(TypeError): + CueSpecs() # type: ignore + + # All prebuilt/default methods construct properly + CueSpecs.default() + CueSpecs.default(GameType.SNOOKER) + CueSpecs.prebuilt(PrebuiltCueSpecs.POOL_GENERIC) + + +def test_cue_from_game_type(): + assert Cue.from_game_type(GameType.EIGHTBALL).specs == CueSpecs.default() + assert Cue.from_game_type(GameType.SNOOKER).specs == CueSpecs.default( + GameType.SNOOKER + ) + + # ID is passed through + assert Cue.from_game_type(GameType.SNOOKER).id == "cue_stick" + assert Cue.from_game_type(GameType.SNOOKER, id="other").id == "other" + + +def test_cue_from_game_type_unknown(): + with pytest.raises(NotImplementedError): + Cue.from_game_type("unknown_game") # type: ignore diff --git a/tests/ruleset/test_shots/01_test_shot_no_point.msgpack b/tests/ruleset/test_shots/01_test_shot_no_point.msgpack index 915cad11..e7a49835 100644 Binary files a/tests/ruleset/test_shots/01_test_shot_no_point.msgpack and b/tests/ruleset/test_shots/01_test_shot_no_point.msgpack differ diff --git a/tests/ruleset/test_shots/01a_test_shot_no_point.msgpack b/tests/ruleset/test_shots/01a_test_shot_no_point.msgpack index 0c4aab28..942d7306 100644 Binary files a/tests/ruleset/test_shots/01a_test_shot_no_point.msgpack and b/tests/ruleset/test_shots/01a_test_shot_no_point.msgpack differ diff --git a/tests/ruleset/test_shots/02_test_shot_ispoint.msgpack b/tests/ruleset/test_shots/02_test_shot_ispoint.msgpack index 0e7ac238..ddac4a41 100644 Binary files a/tests/ruleset/test_shots/02_test_shot_ispoint.msgpack and b/tests/ruleset/test_shots/02_test_shot_ispoint.msgpack differ diff --git a/tests/ruleset/test_shots/02a_test_shot_ispoint.msgpack b/tests/ruleset/test_shots/02a_test_shot_ispoint.msgpack index 7f5c3cef..50561aee 100644 Binary files a/tests/ruleset/test_shots/02a_test_shot_ispoint.msgpack and b/tests/ruleset/test_shots/02a_test_shot_ispoint.msgpack differ diff --git a/tests/ruleset/test_shots/03_test_shot_ispoint.msgpack b/tests/ruleset/test_shots/03_test_shot_ispoint.msgpack index 2b005135..c67f5b90 100644 Binary files a/tests/ruleset/test_shots/03_test_shot_ispoint.msgpack and b/tests/ruleset/test_shots/03_test_shot_ispoint.msgpack differ diff --git a/tests/ruleset/test_shots/03a_test_shot_ispoint.msgpack b/tests/ruleset/test_shots/03a_test_shot_ispoint.msgpack index b1cbec3e..984fd9b0 100644 Binary files a/tests/ruleset/test_shots/03a_test_shot_ispoint.msgpack and b/tests/ruleset/test_shots/03a_test_shot_ispoint.msgpack differ diff --git a/tests/ruleset/test_shots/04_test_shot_no_point.msgpack b/tests/ruleset/test_shots/04_test_shot_no_point.msgpack index cdc7e03b..9b086cf3 100644 Binary files a/tests/ruleset/test_shots/04_test_shot_no_point.msgpack and b/tests/ruleset/test_shots/04_test_shot_no_point.msgpack differ diff --git a/tests/ruleset/test_shots/04a_test_shot_no_point.msgpack b/tests/ruleset/test_shots/04a_test_shot_no_point.msgpack index 74475770..c344bc84 100644 Binary files a/tests/ruleset/test_shots/04a_test_shot_no_point.msgpack and b/tests/ruleset/test_shots/04a_test_shot_no_point.msgpack differ diff --git a/tests/ruleset/test_shots/05_test_shot_ispoint.msgpack b/tests/ruleset/test_shots/05_test_shot_ispoint.msgpack index ac886a69..d3267dd1 100644 Binary files a/tests/ruleset/test_shots/05_test_shot_ispoint.msgpack and b/tests/ruleset/test_shots/05_test_shot_ispoint.msgpack differ diff --git a/tests/ruleset/test_shots/05a_test_shot_ispoint.msgpack b/tests/ruleset/test_shots/05a_test_shot_ispoint.msgpack index 7e65d65f..19600a5d 100644 Binary files a/tests/ruleset/test_shots/05a_test_shot_ispoint.msgpack and b/tests/ruleset/test_shots/05a_test_shot_ispoint.msgpack differ