cortex-solver 3.0.0__py3-none-any.whl

cortex/core/scene_solver.py

@@ -0,0 +1,586 @@
+"""Scene solver for scene-level constraints.
+
+The solver initializes all objects at the origin and applies constraints
+in the order they appear in the recipe. It produces a SolveResult with
+placements for each object (and any mirrored objects).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+import random
+from typing import Any, Iterable
+
+from cortex.types import (
+    Dimensions,
+    PartSpec,
+    SceneConstraint,
+    SceneConstraintType,
+    SceneRecipe,
+    SolveResult,
+)
+
+
+Vector3 = list[float]
+Placement = dict[str, list[float]]
+Placements = dict[str, Placement]
+
+
+@dataclass(frozen=True)
+class ZoneBounds:
+    min: Vector3
+    max: Vector3
+
+
+def solve_scene(recipe: SceneRecipe) -> SolveResult:
+    """Solve all scene constraints for a recipe.
+
+    Constraints are processed in order. Object placements are always
+    stored as "location", "rotation", and "scale" vectors.
+    """
+
+    placements = _init_placements(recipe.objects)
+    specs = dict(recipe.objects)
+    warnings: list[str] = []
+
+    for constraint in recipe.constraints:
+        if constraint.type == SceneConstraintType.GRID:
+            _apply_grid(constraint, placements, warnings)
+        elif constraint.type == SceneConstraintType.ALONG_PATH:
+            _apply_along_path(constraint, placements, warnings)
+        elif constraint.type == SceneConstraintType.RADIAL:
+            _apply_radial(constraint, placements, warnings)
+        elif constraint.type == SceneConstraintType.FACING:
+            _apply_facing(constraint, placements, warnings)
+        elif constraint.type == SceneConstraintType.DISTANCE:
+            _apply_distance(constraint, placements, warnings)
+        elif constraint.type == SceneConstraintType.AGAINST_EDGE:
+            _apply_against_edge(constraint, placements, recipe, warnings)
+        elif constraint.type == SceneConstraintType.RANDOM_SCATTER:
+            _apply_random_scatter(constraint, placements, recipe, warnings)
+        elif constraint.type == SceneConstraintType.STACK_VERTICAL:
+            _apply_stack_vertical(constraint, placements, specs, warnings)
+        elif constraint.type == SceneConstraintType.MIRROR_SCENE:
+            _apply_mirror_scene(constraint, placements, specs, warnings)
+        else:
+            warnings.append(f"Unknown scene constraint type: {constraint.type}")
+
+    return SolveResult(success=True, positions=placements, warnings=warnings)
+
+
+# ---------------------------------------------------------------------------
+# Initialization helpers
+# ---------------------------------------------------------------------------
+
+
+def _init_placements(objects: dict[str, PartSpec]) -> Placements:
+    placements: Placements = {}
+    for name in objects:
+        placements[name] = _default_pose()
+    return placements
+
+
+def _default_pose() -> Placement:
+    return {
+        "location": [0.0, 0.0, 0.0],
+        "rotation": [0.0, 0.0, 0.0],
+        "scale": [1.0, 1.0, 1.0],
+    }
+
+
+def _ensure_object_list(value: str | list[str]) -> list[str]:
+    if isinstance(value, list):
+        return value
+    return [value]
+
+
+# ---------------------------------------------------------------------------
+# Vector helpers
+# ---------------------------------------------------------------------------
+
+
+def _vec_add(a: Vector3, b: Vector3) -> Vector3:
+    return [a[0] + b[0], a[1] + b[1], a[2] + b[2]]
+
+
+def _vec_sub(a: Vector3, b: Vector3) -> Vector3:
+    return [a[0] - b[0], a[1] - b[1], a[2] - b[2]]
+
+
+def _vec_mul(a: Vector3, scalar: float) -> Vector3:
+    return [a[0] * scalar, a[1] * scalar, a[2] * scalar]
+
+
+def _vec_length(a: Vector3) -> float:
+    return math.sqrt(a[0] ** 2 + a[1] ** 2 + a[2] ** 2)
+
+
+def _normalize(a: Vector3) -> Vector3:
+    length = _vec_length(a)
+    if length == 0:
+        return [1.0, 0.0, 0.0]
+    return [a[0] / length, a[1] / length, a[2] / length]
+
+
+def _perpendicular_xy(a: Vector3) -> Vector3:
+    return [-a[1], a[0], 0.0]
+
+
+def _lerp(a: float, b: float, t: float) -> float:
+    return a + (b - a) * t
+
+
+def _distance(a: Vector3, b: Vector3) -> float:
+    return _vec_length(_vec_sub(a, b))
+
+
+def _degrees_atan2(dy: float, dx: float) -> float:
+    return math.degrees(math.atan2(dy, dx))
+
+
+# ---------------------------------------------------------------------------
+# Constraint implementations
+# ---------------------------------------------------------------------------
+
+
+def _apply_grid(
+    constraint: SceneConstraint, placements: Placements, warnings: list[str]
+) -> None:
+    objects = _ensure_object_list(constraint.object)
+    origin = _as_vec3(constraint.params.get("origin"), default=[0.0, 0.0, 0.0])
+    rows = int(constraint.params.get("rows", 0))
+    cols = int(constraint.params.get("cols", 0))
+    spacing_x = float(constraint.params.get("spacing_x", 0.0))
+    spacing_y = float(constraint.params.get("spacing_y", 0.0))
+
+    capacity = rows * cols
+    if capacity and len(objects) > capacity:
+        warnings.append(
+            f"GRID capacity {capacity} smaller than object count {len(objects)}"
+        )
+
+    for index, name in enumerate(objects):
+        row = 0 if cols == 0 else index // cols
+        col = index if cols == 0 else index % cols
+        position = [
+            origin[0] + col * spacing_x,
+            origin[1] + row * spacing_y,
+            origin[2],
+        ]
+        _set_location(placements, name, position, warnings)
+
+
+def _apply_along_path(
+    constraint: SceneConstraint,
+    placements: Placements,
+    warnings: list[str],
+) -> None:
+    objects = _ensure_object_list(constraint.object)
+    start = _as_vec3(constraint.params.get("start"), default=[0.0, 0.0, 0.0])
+    end = _as_vec3(constraint.params.get("end"), default=[0.0, 0.0, 0.0])
+    spacing = float(constraint.params.get("spacing", 0.0))
+    offset_lateral = float(constraint.params.get("offset_lateral", 0.0))
+
+    direction = _vec_sub(end, start)
+    if _vec_length(direction) == 0:
+        warnings.append(
+            "ALONG_PATH start and end are identical; using default direction"
+        )
+    direction = _normalize(direction)
+    perpendicular = _perpendicular_xy(direction)
+
+    for index, name in enumerate(objects):
+        along = _vec_mul(direction, spacing * index)
+        lateral = _vec_mul(perpendicular, offset_lateral)
+        position = _vec_add(start, _vec_add(along, lateral))
+        _set_location(placements, name, position, warnings)
+
+
+def _apply_radial(
+    constraint: SceneConstraint,
+    placements: Placements,
+    warnings: list[str],
+) -> None:
+    objects = _ensure_object_list(constraint.object)
+    center = _as_vec3(constraint.params.get("center"), default=[0.0, 0.0, 0.0])
+    radius = float(constraint.params.get("radius", 0.0))
+    count = int(constraint.params.get("count", len(objects)))
+    start_angle = float(constraint.params.get("start_angle", 0.0))
+    face_center = bool(constraint.params.get("face_center", False))
+
+    if count <= 0:
+        warnings.append("RADIAL count must be positive; using 1")
+        count = 1
+
+    step = 360.0 / count
+
+    for index, name in enumerate(objects):
+        angle_deg = start_angle + index * step
+        angle_rad = math.radians(angle_deg)
+        position = [
+            center[0] + radius * math.cos(angle_rad),
+            center[1] + radius * math.sin(angle_rad),
+            center[2],
+        ]
+        _set_location(placements, name, position, warnings)
+        if face_center:
+            rotation = _get_rotation(placements, name, warnings)
+            rotation[2] = _degrees_atan2(
+                center[1] - position[1], center[0] - position[0]
+            )
+            _set_rotation(placements, name, rotation, warnings)
+
+
+def _apply_facing(
+    constraint: SceneConstraint,
+    placements: Placements,
+    warnings: list[str],
+) -> None:
+    objects = _ensure_object_list(constraint.object)
+    target = constraint.params.get("target")
+    if not isinstance(target, str):
+        warnings.append("FACING constraint requires 'target' object name")
+        return
+
+    target_location = _get_location(placements, target, warnings)
+    if target_location is None:
+        return
+
+    for name in objects:
+        location = _get_location(placements, name, warnings)
+        if location is None:
+            continue
+        rotation = _get_rotation(placements, name, warnings)
+        rotation[2] = _degrees_atan2(
+            target_location[1] - location[1],
+            target_location[0] - location[0],
+        )
+        _set_rotation(placements, name, rotation, warnings)
+
+
+def _apply_distance(
+    constraint: SceneConstraint,
+    placements: Placements,
+    warnings: list[str],
+) -> None:
+    objects = _ensure_object_list(constraint.object)
+    min_distance = float(constraint.params.get("min_distance", 0.0))
+
+    for idx, name_a in enumerate(objects):
+        loc_a = _get_location(placements, name_a, warnings)
+        if loc_a is None:
+            continue
+        for name_b in objects[idx + 1 :]:
+            loc_b = _get_location(placements, name_b, warnings)
+            if loc_b is None:
+                continue
+            distance = _distance(loc_a, loc_b)
+            if distance < min_distance:
+                warnings.append(
+                    f"DISTANCE constraint violated by {name_a} and {name_b}"
+                )
+
+
+def _apply_against_edge(
+    constraint: SceneConstraint,
+    placements: Placements,
+    recipe: SceneRecipe,
+    warnings: list[str],
+) -> None:
+    objects = _ensure_object_list(constraint.object)
+    if len(objects) != 1:
+        warnings.append("AGAINST_EDGE constraint expects a single object")
+        return
+
+    edge = str(constraint.params.get("edge", "")).lower()
+    zone_name = constraint.params.get("zone")
+    position_along = float(constraint.params.get("position_along", 0.0))
+
+    bounds = _get_zone_bounds(recipe, zone_name, warnings)
+    if bounds is None:
+        return
+
+    x = bounds.min[0]
+    y = bounds.min[1]
+    if edge == "north":
+        x = _lerp(bounds.min[0], bounds.max[0], position_along)
+        y = bounds.max[1]
+    elif edge == "south":
+        x = _lerp(bounds.min[0], bounds.max[0], position_along)
+        y = bounds.min[1]
+    elif edge == "east":
+        x = bounds.max[0]
+        y = _lerp(bounds.min[1], bounds.max[1], position_along)
+    elif edge == "west":
+        x = bounds.min[0]
+        y = _lerp(bounds.min[1], bounds.max[1], position_along)
+    else:
+        warnings.append(f"AGAINST_EDGE unknown edge: {edge}")
+        return
+
+    location = _get_location(placements, objects[0], warnings)
+    if location is None:
+        return
+    location[0] = x
+    location[1] = y
+    _set_location(placements, objects[0], location, warnings)
+
+
+def _apply_random_scatter(
+    constraint: SceneConstraint,
+    placements: Placements,
+    recipe: SceneRecipe,
+    warnings: list[str],
+) -> None:
+    objects = _ensure_object_list(constraint.object)
+    zone_name = constraint.params.get("zone")
+    min_spacing = float(constraint.params.get("min_spacing", 0.0))
+    seed = constraint.params.get("seed")
+    rng = random.Random(seed)
+
+    bounds = _get_zone_bounds(recipe, zone_name, warnings)
+    if bounds is None:
+        return
+
+    for name in objects:
+        placed = False
+        for _ in range(1000):
+            position = [
+                rng.uniform(bounds.min[0], bounds.max[0]),
+                rng.uniform(bounds.min[1], bounds.max[1]),
+                rng.uniform(bounds.min[2], bounds.max[2]),
+            ]
+            if _is_far_enough(position, placements, min_spacing):
+                _set_location(placements, name, position, warnings)
+                placed = True
+                break
+        if not placed:
+            warnings.append(
+                f"RANDOM_SCATTER failed to place {name} in zone {zone_name}"
+            )
+
+
+def _apply_stack_vertical(
+    constraint: SceneConstraint,
+    placements: Placements,
+    specs: dict[str, PartSpec],
+    warnings: list[str],
+) -> None:
+    objects = _ensure_object_list(constraint.object)
+    base_position = _as_vec3(
+        constraint.params.get("base_position"), default=[0.0, 0.0, 0.0]
+    )
+    gap = float(constraint.params.get("gap", 0.0))
+
+    prev_pos = None
+    prev_height = 0.0
+    for index, name in enumerate(objects):
+        dimensions = _get_dimensions(specs, name, warnings)
+        curr_height = dimensions.height if dimensions else 0.0
+        if index == 0:
+            position = base_position
+        else:
+            prev_top = 0.0 if prev_pos is None else prev_pos[2] + prev_height / 2.0
+            position = [
+                base_position[0],
+                base_position[1],
+                prev_top + gap + curr_height / 2.0,
+            ]
+        _set_location(placements, name, position, warnings)
+        prev_pos = position
+        prev_height = curr_height
+
+
+def _apply_mirror_scene(
+    constraint: SceneConstraint,
+    placements: Placements,
+    specs: dict[str, PartSpec],
+    warnings: list[str],
+) -> None:
+    targets = _ensure_object_list(constraint.object)
+    axis = str(constraint.params.get("axis", "")).lower()
+    center = float(constraint.params.get("center", 0.0))
+
+    axis_idx = _axis_to_index(axis)
+    if axis_idx is None:
+        warnings.append(f"MIRROR_SCENE unknown axis: {axis}")
+        return
+
+    for name in targets:
+        location = _get_location(placements, name, warnings)
+        if location is None:
+            continue
+        mirrored_name = f"{name}_mirrored"
+        mirrored_location = list(location)
+        mirrored_location[axis_idx] = 2 * center - location[axis_idx]
+
+        placements[mirrored_name] = _default_pose()
+        _set_location(placements, mirrored_name, mirrored_location, warnings)
+        _set_rotation(
+            placements,
+            mirrored_name,
+            _get_rotation(placements, name, warnings),
+            warnings,
+        )
+        _set_scale(
+            placements, mirrored_name, _get_scale(placements, name, warnings), warnings
+        )
+
+        if name in specs:
+            specs[mirrored_name] = specs[name]
+        else:
+            warnings.append(f"MIRROR_SCENE missing spec for {name}")
+
+
+# ---------------------------------------------------------------------------
+# Constraint helpers
+# ---------------------------------------------------------------------------
+
+
+def _as_vec3(value: Any, default: Vector3) -> Vector3:
+    if isinstance(value, (list, tuple)) and len(value) >= 3:
+        return [float(value[0]), float(value[1]), float(value[2])]
+    return list(default)
+
+
+def _axis_to_index(axis: str) -> int | None:
+    if axis == "x":
+        return 0
+    if axis == "y":
+        return 1
+    if axis == "z":
+        return 2
+    return None
+
+
+def _get_zone_bounds(
+    recipe: SceneRecipe,
+    zone_name: Any,
+    warnings: list[str],
+) -> ZoneBounds | None:
+    if not isinstance(zone_name, str):
+        warnings.append("Zone name must be a string")
+        return None
+
+    zone = recipe.zones.get(zone_name)
+    if zone is None:
+        warnings.append(f"Unknown zone: {zone_name}")
+        return None
+
+    bounds = zone.get("bounds") if isinstance(zone, dict) else None
+    if not isinstance(bounds, dict):
+        warnings.append(f"Zone {zone_name} missing bounds")
+        return None
+
+    min_vec = _as_vec3(bounds.get("min"), default=[0.0, 0.0, 0.0])
+    max_vec = _as_vec3(bounds.get("max"), default=[0.0, 0.0, 0.0])
+    return ZoneBounds(min=min_vec, max=max_vec)
+
+
+def _get_dimensions(
+    specs: dict[str, PartSpec],
+    name: str,
+    warnings: list[str],
+) -> Dimensions | None:
+    spec = specs.get(name)
+    if spec is None:
+        warnings.append(f"Missing dimensions for {name}")
+        return None
+    return spec.dimensions
+
+
+def _set_location(
+    placements: Placements,
+    name: str,
+    location: Vector3,
+    warnings: list[str],
+) -> None:
+    pose = placements.get(name)
+    if pose is None:
+        warnings.append(f"Unknown object: {name}")
+        return
+    pose["location"] = [float(location[0]), float(location[1]), float(location[2])]
+
+
+def _set_rotation(
+    placements: Placements,
+    name: str,
+    rotation: Vector3,
+    warnings: list[str],
+) -> None:
+    pose = placements.get(name)
+    if pose is None:
+        warnings.append(f"Unknown object: {name}")
+        return
+    pose["rotation"] = [float(rotation[0]), float(rotation[1]), float(rotation[2])]
+
+
+def _set_scale(
+    placements: Placements,
+    name: str,
+    scale: Vector3,
+    warnings: list[str],
+) -> None:
+    pose = placements.get(name)
+    if pose is None:
+        warnings.append(f"Unknown object: {name}")
+        return
+    pose["scale"] = [float(scale[0]), float(scale[1]), float(scale[2])]
+
+
+def _get_location(
+    placements: Placements,
+    name: str,
+    warnings: list[str],
+) -> Vector3 | None:
+    pose = placements.get(name)
+    if pose is None:
+        warnings.append(f"Unknown object: {name}")
+        return None
+    return list(pose["location"])
+
+
+def _get_rotation(
+    placements: Placements,
+    name: str,
+    warnings: list[str],
+) -> Vector3:
+    pose = placements.get(name)
+    if pose is None:
+        warnings.append(f"Unknown object: {name}")
+        return [0.0, 0.0, 0.0]
+    return list(pose["rotation"])
+
+
+def _get_scale(
+    placements: Placements,
+    name: str,
+    warnings: list[str],
+) -> Vector3:
+    pose = placements.get(name)
+    if pose is None:
+        warnings.append(f"Unknown object: {name}")
+        return [1.0, 1.0, 1.0]
+    return list(pose["scale"])
+
+
+def _is_far_enough(
+    position: Vector3,
+    placements: Placements,
+    min_spacing: float,
+) -> bool:
+    if min_spacing <= 0:
+        return True
+    for pose in placements.values():
+        distance = _distance(position, pose["location"])
+        if distance < min_spacing:
+            return False
+    return True
+
+
+def _iter_locations(names: Iterable[str], placements: Placements) -> Iterable[Vector3]:
+    for name in names:
+        pose = placements.get(name)
+        if pose is not None:
+            yield list(pose["location"])