From 9f53fdb15428e922b61b108fc6df896e4955d64b Mon Sep 17 00:00:00 2001 From: forbes-0023 Date: Tue, 3 Feb 2026 12:55:37 -0600 Subject: [PATCH] feat(mates): add mate type definitions and geometry references MateType enum (8 types), GeometryType enum (5 types), GeometryRef and Mate dataclasses with validation, serialization, and context-dependent DOF removal via dof_removed(). Closes #11 --- solver/mates/__init__.py | 17 ++ solver/mates/primitives.py | 279 ++++++++++++++++++++++++++++ tests/mates/__init__.py | 0 tests/mates/test_primitives.py | 329 +++++++++++++++++++++++++++++++++ 4 files changed, 625 insertions(+) create mode 100644 solver/mates/__init__.py create mode 100644 solver/mates/primitives.py create mode 100644 tests/mates/__init__.py create mode 100644 tests/mates/test_primitives.py diff --git a/solver/mates/__init__.py b/solver/mates/__init__.py new file mode 100644 index 0000000..994e51d --- /dev/null +++ b/solver/mates/__init__.py @@ -0,0 +1,17 @@ +"""Mate-level constraint types for assembly analysis.""" + +from solver.mates.primitives import ( + GeometryRef, + GeometryType, + Mate, + MateType, + dof_removed, +) + +__all__ = [ + "GeometryRef", + "GeometryType", + "Mate", + "MateType", + "dof_removed", +] diff --git a/solver/mates/primitives.py b/solver/mates/primitives.py new file mode 100644 index 0000000..9a3b0d2 --- /dev/null +++ b/solver/mates/primitives.py @@ -0,0 +1,279 @@ +"""Mate type definitions and geometry references for assembly constraints. + +Mates are the user-facing constraint primitives in CAD (e.g. SolidWorks-style +Coincident, Concentric, Parallel). Each mate references geometry on two bodies +and removes a context-dependent number of degrees of freedom. +""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from typing import Any + +__all__ = [ + "GeometryRef", + "GeometryType", + "Mate", + "MateType", + "dof_removed", +] + + +# --------------------------------------------------------------------------- +# Enums +# --------------------------------------------------------------------------- + + +class MateType(enum.Enum): + """CAD mate types with default DOF-removal counts. + + Values are ``(ordinal, default_dof)`` tuples so that mate types + sharing the same DOF count remain distinct enum members. Use the + :attr:`default_dof` property to get the scalar constraint count. + + The actual DOF removed can be context-dependent (e.g. COINCIDENT + removes 3 DOF for face-face but only 1 for face-point). Use + :func:`dof_removed` for the context-aware count. + """ + + COINCIDENT = (0, 3) + CONCENTRIC = (1, 2) + PARALLEL = (2, 2) + PERPENDICULAR = (3, 1) + TANGENT = (4, 1) + DISTANCE = (5, 1) + ANGLE = (6, 1) + LOCK = (7, 6) + + @property + def default_dof(self) -> int: + """Default number of DOF removed by this mate type.""" + return self.value[1] + + +class GeometryType(enum.Enum): + """Types of geometric references used by mates.""" + + FACE = "face" + EDGE = "edge" + POINT = "point" + AXIS = "axis" + PLANE = "plane" + + +# Geometry types that require a direction vector. +_DIRECTIONAL_TYPES = frozenset( + { + GeometryType.FACE, + GeometryType.AXIS, + GeometryType.PLANE, + } +) + + +# --------------------------------------------------------------------------- +# Dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class GeometryRef: + """A reference to a specific geometric entity on a body. + + Attributes: + body_id: Index of the body this geometry belongs to. + geometry_type: What kind of geometry (face, edge, etc.). + geometry_id: CAD identifier string (e.g. ``"Face001"``). + origin: 3D position of the geometry reference point. + direction: Unit direction vector. Required for FACE, AXIS, PLANE; + ``None`` for POINT. + """ + + body_id: int + geometry_type: GeometryType + geometry_id: str + origin: np.ndarray = field(default_factory=lambda: np.zeros(3)) + direction: np.ndarray | None = None + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable dict.""" + return { + "body_id": self.body_id, + "geometry_type": self.geometry_type.value, + "geometry_id": self.geometry_id, + "origin": self.origin.tolist(), + "direction": self.direction.tolist() if self.direction is not None else None, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> GeometryRef: + """Construct from a dict produced by :meth:`to_dict`.""" + direction_raw = data.get("direction") + return cls( + body_id=data["body_id"], + geometry_type=GeometryType(data["geometry_type"]), + geometry_id=data["geometry_id"], + origin=np.asarray(data["origin"], dtype=np.float64), + direction=( + np.asarray(direction_raw, dtype=np.float64) if direction_raw is not None else None + ), + ) + + +@dataclass +class Mate: + """A mate constraint between geometry on two bodies. + + Attributes: + mate_id: Unique identifier for this mate. + mate_type: The type of constraint (Coincident, Concentric, etc.). + ref_a: Geometry reference on the first body. + ref_b: Geometry reference on the second body. + value: Scalar parameter for DISTANCE and ANGLE mates (0 otherwise). + tolerance: Numeric tolerance for constraint satisfaction. + """ + + mate_id: int + mate_type: MateType + ref_a: GeometryRef + ref_b: GeometryRef + value: float = 0.0 + tolerance: float = 1e-6 + + def validate(self) -> None: + """Raise ``ValueError`` if this mate has incompatible geometry. + + Checks: + - Self-mate (both refs on same body) + - CONCENTRIC requires AXIS geometry on both refs + - PARALLEL requires directional geometry (not POINT) + - TANGENT requires surface geometry (FACE or EDGE) + - Directional geometry types must have a direction vector + """ + if self.ref_a.body_id == self.ref_b.body_id: + msg = f"Self-mate: ref_a and ref_b both reference body {self.ref_a.body_id}" + raise ValueError(msg) + + for label, ref in [("ref_a", self.ref_a), ("ref_b", self.ref_b)]: + if ref.geometry_type in _DIRECTIONAL_TYPES and ref.direction is None: + msg = ( + f"{label}: geometry type {ref.geometry_type.value} requires a direction vector" + ) + raise ValueError(msg) + + if self.mate_type is MateType.CONCENTRIC: + for label, ref in [("ref_a", self.ref_a), ("ref_b", self.ref_b)]: + if ref.geometry_type is not GeometryType.AXIS: + msg = ( + f"CONCENTRIC mate requires AXIS geometry, " + f"got {ref.geometry_type.value} on {label}" + ) + raise ValueError(msg) + + if self.mate_type is MateType.PARALLEL: + for label, ref in [("ref_a", self.ref_a), ("ref_b", self.ref_b)]: + if ref.geometry_type is GeometryType.POINT: + msg = f"PARALLEL mate requires directional geometry, got POINT on {label}" + raise ValueError(msg) + + if self.mate_type is MateType.TANGENT: + _surface = frozenset({GeometryType.FACE, GeometryType.EDGE}) + for label, ref in [("ref_a", self.ref_a), ("ref_b", self.ref_b)]: + if ref.geometry_type not in _surface: + msg = ( + f"TANGENT mate requires surface geometry " + f"(FACE or EDGE), got {ref.geometry_type.value} " + f"on {label}" + ) + raise ValueError(msg) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable dict.""" + return { + "mate_id": self.mate_id, + "mate_type": self.mate_type.name, + "ref_a": self.ref_a.to_dict(), + "ref_b": self.ref_b.to_dict(), + "value": self.value, + "tolerance": self.tolerance, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Mate: + """Construct from a dict produced by :meth:`to_dict`.""" + return cls( + mate_id=data["mate_id"], + mate_type=MateType[data["mate_type"]], + ref_a=GeometryRef.from_dict(data["ref_a"]), + ref_b=GeometryRef.from_dict(data["ref_b"]), + value=data.get("value", 0.0), + tolerance=data.get("tolerance", 1e-6), + ) + + +# --------------------------------------------------------------------------- +# Context-dependent DOF removal +# --------------------------------------------------------------------------- + +# Lookup table: (MateType, ref_a GeometryType, ref_b GeometryType) -> DOF removed. +# Entries with None match any geometry type for that position. +_DOF_TABLE: dict[tuple[MateType, GeometryType | None, GeometryType | None], int] = { + # COINCIDENT — context-dependent + (MateType.COINCIDENT, GeometryType.FACE, GeometryType.FACE): 3, + (MateType.COINCIDENT, GeometryType.POINT, GeometryType.POINT): 3, + (MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE): 3, + (MateType.COINCIDENT, GeometryType.EDGE, GeometryType.EDGE): 2, + (MateType.COINCIDENT, GeometryType.FACE, GeometryType.POINT): 1, + (MateType.COINCIDENT, GeometryType.POINT, GeometryType.FACE): 1, + # CONCENTRIC + (MateType.CONCENTRIC, GeometryType.AXIS, GeometryType.AXIS): 2, + # PARALLEL + (MateType.PARALLEL, GeometryType.AXIS, GeometryType.AXIS): 2, + (MateType.PARALLEL, GeometryType.FACE, GeometryType.FACE): 2, + (MateType.PARALLEL, GeometryType.PLANE, GeometryType.PLANE): 2, + # TANGENT + (MateType.TANGENT, GeometryType.FACE, GeometryType.FACE): 1, + (MateType.TANGENT, GeometryType.FACE, GeometryType.EDGE): 1, + (MateType.TANGENT, GeometryType.EDGE, GeometryType.FACE): 1, + # Types where DOF is always the same regardless of geometry + (MateType.PERPENDICULAR, None, None): 1, + (MateType.DISTANCE, None, None): 1, + (MateType.ANGLE, None, None): 1, + (MateType.LOCK, None, None): 6, +} + + +def dof_removed( + mate_type: MateType, + ref_a: GeometryRef, + ref_b: GeometryRef, +) -> int: + """Return the number of DOF removed by a mate given its geometry context. + + Looks up the exact ``(mate_type, ref_a.geometry_type, ref_b.geometry_type)`` + combination first, then falls back to a wildcard ``(mate_type, None, None)`` + entry, and finally to :attr:`MateType.default_dof`. + + Args: + mate_type: The mate constraint type. + ref_a: Geometry reference on the first body. + ref_b: Geometry reference on the second body. + + Returns: + Number of scalar DOF removed by this mate. + """ + key = (mate_type, ref_a.geometry_type, ref_b.geometry_type) + if key in _DOF_TABLE: + return _DOF_TABLE[key] + + wildcard = (mate_type, None, None) + if wildcard in _DOF_TABLE: + return _DOF_TABLE[wildcard] + + return mate_type.default_dof diff --git a/tests/mates/__init__.py b/tests/mates/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/mates/test_primitives.py b/tests/mates/test_primitives.py new file mode 100644 index 0000000..c28fff8 --- /dev/null +++ b/tests/mates/test_primitives.py @@ -0,0 +1,329 @@ +"""Tests for solver.mates.primitives -- mate type definitions.""" + +from __future__ import annotations + +from typing import ClassVar + +import numpy as np +import pytest + +from solver.mates.primitives import ( + GeometryRef, + GeometryType, + Mate, + MateType, + dof_removed, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_ref( + body_id: int, + geom_type: GeometryType, + *, + geometry_id: str = "Geom001", + origin: np.ndarray | None = None, + direction: np.ndarray | None = None, +) -> GeometryRef: + """Factory for GeometryRef with sensible defaults.""" + if origin is None: + origin = np.zeros(3) + if direction is None and geom_type in { + GeometryType.FACE, + GeometryType.AXIS, + GeometryType.PLANE, + }: + direction = np.array([0.0, 0.0, 1.0]) + return GeometryRef( + body_id=body_id, + geometry_type=geom_type, + geometry_id=geometry_id, + origin=origin, + direction=direction, + ) + + +# --------------------------------------------------------------------------- +# MateType +# --------------------------------------------------------------------------- + + +class TestMateType: + """MateType enum construction and DOF values.""" + + EXPECTED_DOF: ClassVar[dict[str, int]] = { + "COINCIDENT": 3, + "CONCENTRIC": 2, + "PARALLEL": 2, + "PERPENDICULAR": 1, + "TANGENT": 1, + "DISTANCE": 1, + "ANGLE": 1, + "LOCK": 6, + } + + def test_member_count(self) -> None: + assert len(MateType) == 8 + + @pytest.mark.parametrize("name,dof", EXPECTED_DOF.items()) + def test_default_dof_values(self, name: str, dof: int) -> None: + assert MateType[name].default_dof == dof + + def test_value_is_tuple(self) -> None: + assert MateType.COINCIDENT.value == (0, 3) + assert MateType.COINCIDENT.default_dof == 3 + + def test_access_by_name(self) -> None: + assert MateType["LOCK"] is MateType.LOCK + + def test_no_alias_collision(self) -> None: + ordinals = [m.value[0] for m in MateType] + assert len(ordinals) == len(set(ordinals)) + + +# --------------------------------------------------------------------------- +# GeometryType +# --------------------------------------------------------------------------- + + +class TestGeometryType: + """GeometryType enum.""" + + def test_member_count(self) -> None: + assert len(GeometryType) == 5 + + def test_string_values(self) -> None: + for gt in GeometryType: + assert isinstance(gt.value, str) + assert gt.value == gt.name.lower() + + def test_access_by_name(self) -> None: + assert GeometryType["FACE"] is GeometryType.FACE + + +# --------------------------------------------------------------------------- +# GeometryRef +# --------------------------------------------------------------------------- + + +class TestGeometryRef: + """GeometryRef dataclass.""" + + def test_construction(self) -> None: + ref = _make_ref(0, GeometryType.AXIS, geometry_id="Axis001") + assert ref.body_id == 0 + assert ref.geometry_type is GeometryType.AXIS + assert ref.geometry_id == "Axis001" + np.testing.assert_array_equal(ref.origin, np.zeros(3)) + assert ref.direction is not None + + def test_default_direction_none(self) -> None: + ref = GeometryRef( + body_id=0, + geometry_type=GeometryType.POINT, + geometry_id="Point001", + ) + assert ref.direction is None + + def test_to_dict_round_trip(self) -> None: + ref = _make_ref( + 1, + GeometryType.FACE, + origin=np.array([1.0, 2.0, 3.0]), + direction=np.array([0.0, 1.0, 0.0]), + ) + d = ref.to_dict() + restored = GeometryRef.from_dict(d) + assert restored.body_id == ref.body_id + assert restored.geometry_type is ref.geometry_type + assert restored.geometry_id == ref.geometry_id + np.testing.assert_array_almost_equal(restored.origin, ref.origin) + assert restored.direction is not None + np.testing.assert_array_almost_equal(restored.direction, ref.direction) + + def test_to_dict_with_none_direction(self) -> None: + ref = GeometryRef( + body_id=2, + geometry_type=GeometryType.POINT, + geometry_id="Point002", + origin=np.array([5.0, 6.0, 7.0]), + ) + d = ref.to_dict() + assert d["direction"] is None + restored = GeometryRef.from_dict(d) + assert restored.direction is None + + +# --------------------------------------------------------------------------- +# Mate +# --------------------------------------------------------------------------- + + +class TestMate: + """Mate dataclass.""" + + def test_construction(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(1, GeometryType.FACE) + m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b) + assert m.mate_id == 0 + assert m.mate_type is MateType.COINCIDENT + + def test_value_default_zero(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(1, GeometryType.FACE) + m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b) + assert m.value == 0.0 + + def test_tolerance_default(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(1, GeometryType.FACE) + m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b) + assert m.tolerance == 1e-6 + + def test_to_dict_round_trip(self) -> None: + ref_a = _make_ref(0, GeometryType.AXIS, origin=np.array([1.0, 0.0, 0.0])) + ref_b = _make_ref(1, GeometryType.AXIS, origin=np.array([2.0, 0.0, 0.0])) + m = Mate( + mate_id=5, + mate_type=MateType.CONCENTRIC, + ref_a=ref_a, + ref_b=ref_b, + value=0.0, + tolerance=1e-8, + ) + d = m.to_dict() + restored = Mate.from_dict(d) + assert restored.mate_id == m.mate_id + assert restored.mate_type is m.mate_type + assert restored.ref_a.body_id == m.ref_a.body_id + assert restored.ref_b.body_id == m.ref_b.body_id + assert restored.value == m.value + assert restored.tolerance == m.tolerance + + def test_from_dict_missing_optional(self) -> None: + d = { + "mate_id": 1, + "mate_type": "DISTANCE", + "ref_a": _make_ref(0, GeometryType.POINT).to_dict(), + "ref_b": _make_ref(1, GeometryType.POINT).to_dict(), + } + m = Mate.from_dict(d) + assert m.value == 0.0 + assert m.tolerance == 1e-6 + + +# --------------------------------------------------------------------------- +# dof_removed +# --------------------------------------------------------------------------- + + +class TestDofRemoved: + """Context-dependent DOF removal counts.""" + + def test_coincident_face_face(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(1, GeometryType.FACE) + assert dof_removed(MateType.COINCIDENT, ref_a, ref_b) == 3 + + def test_coincident_point_point(self) -> None: + ref_a = _make_ref(0, GeometryType.POINT) + ref_b = _make_ref(1, GeometryType.POINT) + assert dof_removed(MateType.COINCIDENT, ref_a, ref_b) == 3 + + def test_coincident_edge_edge(self) -> None: + ref_a = _make_ref(0, GeometryType.EDGE) + ref_b = _make_ref(1, GeometryType.EDGE) + assert dof_removed(MateType.COINCIDENT, ref_a, ref_b) == 2 + + def test_coincident_face_point(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(1, GeometryType.POINT) + assert dof_removed(MateType.COINCIDENT, ref_a, ref_b) == 1 + + def test_concentric_axis_axis(self) -> None: + ref_a = _make_ref(0, GeometryType.AXIS) + ref_b = _make_ref(1, GeometryType.AXIS) + assert dof_removed(MateType.CONCENTRIC, ref_a, ref_b) == 2 + + def test_lock_any(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(1, GeometryType.POINT) + assert dof_removed(MateType.LOCK, ref_a, ref_b) == 6 + + def test_distance_any(self) -> None: + ref_a = _make_ref(0, GeometryType.POINT) + ref_b = _make_ref(1, GeometryType.EDGE) + assert dof_removed(MateType.DISTANCE, ref_a, ref_b) == 1 + + def test_unknown_combo_uses_default(self) -> None: + """Unlisted geometry combos fall back to default_dof.""" + ref_a = _make_ref(0, GeometryType.EDGE) + ref_b = _make_ref(1, GeometryType.POINT) + result = dof_removed(MateType.COINCIDENT, ref_a, ref_b) + assert result == MateType.COINCIDENT.default_dof + + +# --------------------------------------------------------------------------- +# Mate.validate +# --------------------------------------------------------------------------- + + +class TestMateValidation: + """Mate.validate() compatibility checks.""" + + def test_valid_concentric(self) -> None: + ref_a = _make_ref(0, GeometryType.AXIS) + ref_b = _make_ref(1, GeometryType.AXIS) + m = Mate(mate_id=0, mate_type=MateType.CONCENTRIC, ref_a=ref_a, ref_b=ref_b) + m.validate() # should not raise + + def test_invalid_concentric_face(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(1, GeometryType.AXIS) + m = Mate(mate_id=0, mate_type=MateType.CONCENTRIC, ref_a=ref_a, ref_b=ref_b) + with pytest.raises(ValueError, match="CONCENTRIC"): + m.validate() + + def test_valid_coincident_face_face(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(1, GeometryType.FACE) + m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b) + m.validate() # should not raise + + def test_invalid_self_mate(self) -> None: + ref_a = _make_ref(0, GeometryType.FACE) + ref_b = _make_ref(0, GeometryType.FACE, geometry_id="Face002") + m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b) + with pytest.raises(ValueError, match="Self-mate"): + m.validate() + + def test_invalid_parallel_point(self) -> None: + ref_a = _make_ref(0, GeometryType.POINT) + ref_b = _make_ref(1, GeometryType.AXIS) + m = Mate(mate_id=0, mate_type=MateType.PARALLEL, ref_a=ref_a, ref_b=ref_b) + with pytest.raises(ValueError, match="PARALLEL"): + m.validate() + + def test_invalid_tangent_axis(self) -> None: + ref_a = _make_ref(0, GeometryType.AXIS) + ref_b = _make_ref(1, GeometryType.FACE) + m = Mate(mate_id=0, mate_type=MateType.TANGENT, ref_a=ref_a, ref_b=ref_b) + with pytest.raises(ValueError, match="TANGENT"): + m.validate() + + def test_missing_direction_for_axis(self) -> None: + ref_a = GeometryRef( + body_id=0, + geometry_type=GeometryType.AXIS, + geometry_id="Axis001", + origin=np.zeros(3), + direction=None, # missing! + ) + ref_b = _make_ref(1, GeometryType.AXIS) + m = Mate(mate_id=0, mate_type=MateType.CONCENTRIC, ref_a=ref_a, ref_b=ref_b) + with pytest.raises(ValueError, match="direction"): + m.validate()