feat(mates): add mate type definitions and geometry references
Some checks failed
CI / lint (push) Has been cancelled
CI / type-check (push) Has been cancelled
CI / test (push) Has been cancelled

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
This commit is contained in:
2026-02-03 12:55:37 -06:00
parent 5d1988b513
commit 9f53fdb154
4 changed files with 625 additions and 0 deletions

17
solver/mates/__init__.py Normal file
View File

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

279
solver/mates/primitives.py Normal file
View File

@@ -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

0
tests/mates/__init__.py Normal file
View File

View File

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