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

@@ -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",
+]

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

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()