From 239e45c7f90077d40b74e49d0848b6a7338c3a2e Mon Sep 17 00:00:00 2001 From: forbes-0023 Date: Tue, 3 Feb 2026 13:05:58 -0600 Subject: [PATCH] feat(mates): add mate-based synthetic assembly generator SyntheticMateGenerator wraps existing joint generator with reverse mapping (joint->mates) and configurable noise injection (redundant, missing, incompatible mates). Batch generation via generate_mate_training_batch(). Closes #14 --- solver/mates/__init__.py | 6 + solver/mates/generator.py | 315 ++++++++++++++++++++++++++++++++++ tests/mates/test_generator.py | 155 +++++++++++++++++ 3 files changed, 476 insertions(+) create mode 100644 solver/mates/generator.py create mode 100644 tests/mates/test_generator.py diff --git a/solver/mates/__init__.py b/solver/mates/__init__.py index fa8a5d5..0496bf4 100644 --- a/solver/mates/__init__.py +++ b/solver/mates/__init__.py @@ -5,6 +5,10 @@ from solver.mates.conversion import ( analyze_mate_assembly, convert_mates_to_joints, ) +from solver.mates.generator import ( + SyntheticMateGenerator, + generate_mate_training_batch, +) from solver.mates.patterns import ( JointPattern, PatternMatch, @@ -26,8 +30,10 @@ __all__ = [ "MateAnalysisResult", "MateType", "PatternMatch", + "SyntheticMateGenerator", "analyze_mate_assembly", "convert_mates_to_joints", "dof_removed", + "generate_mate_training_batch", "recognize_patterns", ] diff --git a/solver/mates/generator.py b/solver/mates/generator.py new file mode 100644 index 0000000..26d3834 --- /dev/null +++ b/solver/mates/generator.py @@ -0,0 +1,315 @@ +"""Mate-based synthetic assembly generator. + +Wraps SyntheticAssemblyGenerator to produce mate-level training data. +Generates joint-based assemblies via the existing generator, then +reverse-maps joints to plausible mate combinations. Supports noise +injection (redundant, missing, incompatible mates) for robust training. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +from solver.datagen.generator import SyntheticAssemblyGenerator +from solver.datagen.types import Joint, JointType, RigidBody +from solver.mates.conversion import MateAnalysisResult, analyze_mate_assembly +from solver.mates.primitives import GeometryRef, GeometryType, Mate, MateType + +if TYPE_CHECKING: + from typing import Any + +__all__ = [ + "SyntheticMateGenerator", + "generate_mate_training_batch", +] + + +# --------------------------------------------------------------------------- +# Reverse mapping: JointType -> list of (MateType, geom_a, geom_b) combos +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _MateSpec: + """Specification for a mate to generate from a joint.""" + + mate_type: MateType + geom_a: GeometryType + geom_b: GeometryType + + +_JOINT_TO_MATES: dict[JointType, list[_MateSpec]] = { + JointType.REVOLUTE: [ + _MateSpec(MateType.CONCENTRIC, GeometryType.AXIS, GeometryType.AXIS), + _MateSpec(MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE), + ], + JointType.CYLINDRICAL: [ + _MateSpec(MateType.CONCENTRIC, GeometryType.AXIS, GeometryType.AXIS), + ], + JointType.BALL: [ + _MateSpec(MateType.COINCIDENT, GeometryType.POINT, GeometryType.POINT), + ], + JointType.FIXED: [ + _MateSpec(MateType.LOCK, GeometryType.FACE, GeometryType.FACE), + ], + JointType.SLIDER: [ + _MateSpec(MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE), + _MateSpec(MateType.PARALLEL, GeometryType.AXIS, GeometryType.AXIS), + ], + JointType.PLANAR: [ + _MateSpec(MateType.COINCIDENT, GeometryType.FACE, GeometryType.FACE), + ], +} + + +# --------------------------------------------------------------------------- +# Generator +# --------------------------------------------------------------------------- + + +class SyntheticMateGenerator: + """Generates mate-based assemblies for training data. + + Wraps SyntheticAssemblyGenerator to produce joint-based assemblies, + then reverse-maps each joint to a plausible set of mate constraints. + + Args: + seed: Random seed for reproducibility. + redundant_prob: Probability of injecting a redundant mate per joint. + missing_prob: Probability of dropping a mate from a multi-mate pattern. + incompatible_prob: Probability of injecting a mate with wrong geometry. + """ + + def __init__( + self, + seed: int = 42, + *, + redundant_prob: float = 0.0, + missing_prob: float = 0.0, + incompatible_prob: float = 0.0, + ) -> None: + self._joint_gen = SyntheticAssemblyGenerator(seed=seed) + self._rng = np.random.default_rng(seed) + self.redundant_prob = redundant_prob + self.missing_prob = missing_prob + self.incompatible_prob = incompatible_prob + + def _make_geometry_ref( + self, + body_id: int, + geom_type: GeometryType, + joint: Joint, + *, + is_ref_a: bool = True, + ) -> GeometryRef: + """Create a GeometryRef from joint geometry. + + Uses joint anchor, axis, and body_id to produce a ref + with realistic geometry for the given type. + """ + origin = joint.anchor_a if is_ref_a else joint.anchor_b + + direction: np.ndarray | None = None + if geom_type in {GeometryType.AXIS, GeometryType.PLANE, GeometryType.FACE}: + direction = joint.axis.copy() + + geom_id = f"{geom_type.value.capitalize()}001" + + return GeometryRef( + body_id=body_id, + geometry_type=geom_type, + geometry_id=geom_id, + origin=origin.copy(), + direction=direction, + ) + + def _reverse_map_joint( + self, + joint: Joint, + next_mate_id: int, + ) -> list[Mate]: + """Convert a joint to its mate representation.""" + specs = _JOINT_TO_MATES.get(joint.joint_type, []) + if not specs: + # Fallback: emit a single DISTANCE mate + specs = [_MateSpec(MateType.DISTANCE, GeometryType.POINT, GeometryType.POINT)] + + mates: list[Mate] = [] + for spec in specs: + ref_a = self._make_geometry_ref(joint.body_a, spec.geom_a, joint, is_ref_a=True) + ref_b = self._make_geometry_ref(joint.body_b, spec.geom_b, joint, is_ref_a=False) + mates.append( + Mate( + mate_id=next_mate_id + len(mates), + mate_type=spec.mate_type, + ref_a=ref_a, + ref_b=ref_b, + ) + ) + return mates + + def _inject_noise( + self, + mates: list[Mate], + next_mate_id: int, + ) -> list[Mate]: + """Apply noise injection to the mate list. + + Modifies the list in-place and may add new mates. + Returns the (possibly extended) list. + """ + result = list(mates) + extra: list[Mate] = [] + + for mate in mates: + # Redundant: duplicate a mate + if self._rng.random() < self.redundant_prob: + dup = Mate( + mate_id=next_mate_id + len(extra), + mate_type=mate.mate_type, + ref_a=mate.ref_a, + ref_b=mate.ref_b, + value=mate.value, + tolerance=mate.tolerance, + ) + extra.append(dup) + + # Incompatible: wrong geometry type + if self._rng.random() < self.incompatible_prob: + bad_geom = GeometryType.POINT + bad_ref = GeometryRef( + body_id=mate.ref_a.body_id, + geometry_type=bad_geom, + geometry_id="BadGeom001", + origin=mate.ref_a.origin.copy(), + direction=None, + ) + extra.append( + Mate( + mate_id=next_mate_id + len(extra), + mate_type=MateType.CONCENTRIC, + ref_a=bad_ref, + ref_b=mate.ref_b, + ) + ) + + result.extend(extra) + + # Missing: drop mates from multi-mate patterns (only if > 1 mate + # for same body pair) + if self.missing_prob > 0: + filtered: list[Mate] = [] + for mate in result: + if self._rng.random() < self.missing_prob: + continue + filtered.append(mate) + # Ensure at least one mate remains + if not filtered and result: + filtered = [result[0]] + result = filtered + + return result + + def generate( + self, + n_bodies: int = 4, + *, + grounded: bool = False, + ) -> tuple[list[RigidBody], list[Mate], MateAnalysisResult]: + """Generate a mate-based assembly. + + Args: + n_bodies: Number of rigid bodies. + grounded: Whether to ground the first body. + + Returns: + (bodies, mates, analysis_result) tuple. + """ + bodies, joints, _analysis = self._joint_gen.generate_chain_assembly( + n_bodies, + joint_type=JointType.REVOLUTE, + grounded=grounded, + ) + + mates: list[Mate] = [] + next_id = 0 + for joint in joints: + joint_mates = self._reverse_map_joint(joint, next_id) + mates.extend(joint_mates) + next_id += len(joint_mates) + + # Apply noise + mates = self._inject_noise(mates, next_id) + + ground_body = bodies[0].body_id if grounded else None + result = analyze_mate_assembly(bodies, mates, ground_body) + + return bodies, mates, result + + +# --------------------------------------------------------------------------- +# Batch generation +# --------------------------------------------------------------------------- + + +def generate_mate_training_batch( + batch_size: int = 100, + n_bodies_range: tuple[int, int] = (3, 8), + seed: int = 42, + *, + redundant_prob: float = 0.0, + missing_prob: float = 0.0, + incompatible_prob: float = 0.0, + grounded_ratio: float = 1.0, +) -> list[dict[str, Any]]: + """Produce a batch of mate-level training examples. + + Args: + batch_size: Number of assemblies to generate. + n_bodies_range: (min, max_exclusive) body count. + seed: Random seed. + redundant_prob: Probability of redundant mate injection. + missing_prob: Probability of missing mate injection. + incompatible_prob: Probability of incompatible mate injection. + grounded_ratio: Fraction of assemblies that are grounded. + + Returns: + List of dicts with bodies, mates, patterns, and labels. + """ + rng = np.random.default_rng(seed) + examples: list[dict[str, Any]] = [] + + for i in range(batch_size): + gen = SyntheticMateGenerator( + seed=seed + i, + redundant_prob=redundant_prob, + missing_prob=missing_prob, + incompatible_prob=incompatible_prob, + ) + n = int(rng.integers(*n_bodies_range)) + grounded = bool(rng.random() < grounded_ratio) + + bodies, mates, result = gen.generate(n, grounded=grounded) + + examples.append( + { + "bodies": [ + { + "body_id": b.body_id, + "position": b.position.tolist(), + } + for b in bodies + ], + "mates": [m.to_dict() for m in mates], + "patterns": [p.to_dict() for p in result.patterns], + "labels": result.labels.to_dict() if result.labels else None, + "n_bodies": len(bodies), + "n_mates": len(mates), + "n_joints": len(result.joints), + } + ) + + return examples diff --git a/tests/mates/test_generator.py b/tests/mates/test_generator.py new file mode 100644 index 0000000..bb8dd71 --- /dev/null +++ b/tests/mates/test_generator.py @@ -0,0 +1,155 @@ +"""Tests for solver.mates.generator -- synthetic mate generator.""" + +from __future__ import annotations + +from solver.mates.generator import SyntheticMateGenerator, generate_mate_training_batch +from solver.mates.primitives import MateType + +# --------------------------------------------------------------------------- +# SyntheticMateGenerator +# --------------------------------------------------------------------------- + + +class TestSyntheticMateGenerator: + """SyntheticMateGenerator core functionality.""" + + def test_generate_basic(self) -> None: + """Generate a simple assembly with mates.""" + gen = SyntheticMateGenerator(seed=42) + bodies, mates, result = gen.generate(3) + assert len(bodies) == 3 + assert len(mates) > 0 + assert result.analysis is not None + + def test_deterministic_with_seed(self) -> None: + """Same seed produces same output.""" + gen1 = SyntheticMateGenerator(seed=123) + _, mates1, _ = gen1.generate(3) + + gen2 = SyntheticMateGenerator(seed=123) + _, mates2, _ = gen2.generate(3) + + assert len(mates1) == len(mates2) + for m1, m2 in zip(mates1, mates2, strict=True): + assert m1.mate_type == m2.mate_type + assert m1.ref_a.body_id == m2.ref_a.body_id + + def test_grounded(self) -> None: + """Grounded assembly should work.""" + gen = SyntheticMateGenerator(seed=42) + bodies, _mates, result = gen.generate(3, grounded=True) + assert len(bodies) == 3 + assert result.analysis is not None + + def test_revolute_produces_two_mates(self) -> None: + """A revolute joint should reverse-map to 2 mates.""" + gen = SyntheticMateGenerator(seed=42) + _bodies, mates, _result = gen.generate(2) + # 2 bodies -> 1 revolute joint -> 2 mates (concentric + coincident) + assert len(mates) == 2 + mate_types = {m.mate_type for m in mates} + assert MateType.CONCENTRIC in mate_types + assert MateType.COINCIDENT in mate_types + + +class TestReverseMapping: + """Reverse mapping from joints to mates.""" + + def test_revolute_mapping(self) -> None: + """REVOLUTE -> Concentric + Coincident.""" + gen = SyntheticMateGenerator(seed=42) + _bodies, mates, _result = gen.generate(2) + types = [m.mate_type for m in mates] + assert MateType.CONCENTRIC in types + assert MateType.COINCIDENT in types + + def test_round_trip_analysis(self) -> None: + """Generated mates round-trip through analysis successfully.""" + gen = SyntheticMateGenerator(seed=42) + _bodies, _mates, result = gen.generate(4) + assert result.analysis is not None + assert result.labels is not None + # Should produce joints from the mates + assert len(result.joints) > 0 + + +class TestNoiseInjection: + """Noise injection mechanisms.""" + + def test_redundant_injection(self) -> None: + """Redundant prob > 0 produces more mates than clean version.""" + gen_clean = SyntheticMateGenerator(seed=42, redundant_prob=0.0) + _, mates_clean, _ = gen_clean.generate(4) + + gen_noisy = SyntheticMateGenerator(seed=42, redundant_prob=1.0) + _, mates_noisy, _ = gen_noisy.generate(4) + + assert len(mates_noisy) > len(mates_clean) + + def test_missing_injection(self) -> None: + """Missing prob > 0 produces fewer mates than clean version.""" + gen_clean = SyntheticMateGenerator(seed=42, missing_prob=0.0) + _, mates_clean, _ = gen_clean.generate(4) + + gen_noisy = SyntheticMateGenerator(seed=42, missing_prob=0.5) + _, mates_noisy, _ = gen_noisy.generate(4) + + # With 50% drop rate on 6 mates, very likely to drop at least one + assert len(mates_noisy) <= len(mates_clean) + + def test_incompatible_injection(self) -> None: + """Incompatible prob > 0 adds mates with wrong geometry.""" + gen = SyntheticMateGenerator(seed=42, incompatible_prob=1.0) + _, mates, _ = gen.generate(3) + # Should have extra mates beyond the clean count + gen_clean = SyntheticMateGenerator(seed=42) + _, mates_clean, _ = gen_clean.generate(3) + assert len(mates) > len(mates_clean) + + +# --------------------------------------------------------------------------- +# generate_mate_training_batch +# --------------------------------------------------------------------------- + + +class TestGenerateMateTrainingBatch: + """Batch generation function.""" + + def test_batch_structure(self) -> None: + """Each example has required keys.""" + examples = generate_mate_training_batch(batch_size=3, seed=42) + assert len(examples) == 3 + for ex in examples: + assert "bodies" in ex + assert "mates" in ex + assert "patterns" in ex + assert "labels" in ex + assert "n_bodies" in ex + assert "n_mates" in ex + assert "n_joints" in ex + + def test_batch_deterministic(self) -> None: + """Same seed produces same batch.""" + batch1 = generate_mate_training_batch(batch_size=5, seed=99) + batch2 = generate_mate_training_batch(batch_size=5, seed=99) + for ex1, ex2 in zip(batch1, batch2, strict=True): + assert ex1["n_bodies"] == ex2["n_bodies"] + assert ex1["n_mates"] == ex2["n_mates"] + + def test_batch_grounded_ratio(self) -> None: + """Batch respects grounded_ratio parameter.""" + # All grounded + examples = generate_mate_training_batch(batch_size=5, seed=42, grounded_ratio=1.0) + assert len(examples) == 5 + + def test_batch_with_noise(self) -> None: + """Batch with noise injection runs without error.""" + examples = generate_mate_training_batch( + batch_size=3, + seed=42, + redundant_prob=0.3, + missing_prob=0.1, + ) + assert len(examples) == 3 + for ex in examples: + assert ex["n_mates"] >= 0