deepcad-canonical-equivalence

#!/usr/bin/env python3
"""DeepCAD canonical-equivalence solver.

For each pair of compact CAD command programs, decide whether both programs
describe the same canonical solid.

Approach: normalize each program into a canonical *signature* that captures only
what affects the solid, then compare the two signatures for equality.

The solid is determined by:
  * the multiset of *extrude* features (each canonicalized: operation, extent,
    direction, depth, and the multiset of consumed profile geometries), grouped
    into body/channel partitions; and
  * the multiset of *opaque* solid-affecting payload commands compared literally.

Everything else (parameter names, ids, sketch ids, command order, construction
geometry, annotation commands) is normalized away or dropped.
"""

from __future__ import annotations

import ast
import json
import sys
from pathlib import Path


# Absolute tolerance for numeric comparison is 1e-6; rounding to 6 decimals
# yields a grid finer than that tolerance.
ROUND_DIGITS = 6

SUPPORTED_KINDS = {"param", "sketch", "rect", "circle", "slot", "extrude"}
PROFILE_KINDS = {"rect", "circle", "slot"}
ANNOTATION_KINDS = {"constraint", "dimension", "note", "metadata", "view", "comment"}

# Fields on a profile command that are labels/wiring, not geometry.
PROFILE_NON_GEOM = {"cmd", "kind", "type", "id", "sketch", "construction", "plane"}


# --------------------------------------------------------------------------- #
# Expression evaluation
# --------------------------------------------------------------------------- #


class ExprError(Exception):
    pass


def _eval_ast(node, resolver):
    if isinstance(node, ast.Expression):
        return _eval_ast(node.body, resolver)
    if isinstance(node, ast.Constant):
        if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
            raise ExprError("non-numeric constant")
        return float(node.value)
    if isinstance(node, ast.Num):  # pragma: no cover (py<3.8 compat)
        return float(node.n)
    if isinstance(node, ast.BinOp):
        left = _eval_ast(node.left, resolver)
        right = _eval_ast(node.right, resolver)
        if isinstance(node.op, ast.Add):
            return left + right
        if isinstance(node.op, ast.Sub):
            return left - right
        if isinstance(node.op, ast.Mult):
            return left * right
        if isinstance(node.op, ast.Div):
            return left / right
        raise ExprError("unsupported binary operator")
    if isinstance(node, ast.UnaryOp):
        operand = _eval_ast(node.operand, resolver)
        if isinstance(node.op, ast.UAdd):
            return +operand
        if isinstance(node.op, ast.USub):
            return -operand
        raise ExprError("unsupported unary operator")
    if isinstance(node, ast.Name):
        return resolver(node.id)
    raise ExprError("unsupported expression node")


def eval_expr(expr, resolver):
    tree = ast.parse(expr, mode="eval")
    return _eval_ast(tree, resolver)


class ParamTable:
    """Resolves parameter names to numeric values, allowing references between
    parameters (in any order) and detecting cycles."""

    def __init__(self, exprs):
        self._exprs = exprs
        self._cache = {}

    def resolve(self, name, stack=None):
        if name in self._cache:
            return self._cache[name]
        if name not in self._exprs:
            raise ExprError(f"unknown parameter: {name}")
        if stack is None:
            stack = ()
        if name in stack:
            raise ExprError(f"cyclic parameter: {name}")
        raw = self._exprs[name]
        value = self._eval_value(raw, stack + (name,))
        self._cache[name] = value
        return value

    def _eval_value(self, raw, stack):
        if isinstance(raw, bool):
            raise ExprError("boolean parameter")
        if isinstance(raw, (int, float)):
            return float(raw)
        if isinstance(raw, str):
            tree = ast.parse(raw, mode="eval")
            return _eval_ast(tree, lambda n: self.resolve(n, stack))
        raise ExprError("unsupported parameter expression")

    def resolver(self):
        return lambda name: self.resolve(name)


def evaluate_value(value, params):
    """Evaluate a JSON value (number, expression string, or nested list) to a
    canonical numeric form. If a string cannot be evaluated as an arithmetic
    expression, it is preserved literally."""
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return round(float(value), ROUND_DIGITS)
    if isinstance(value, list):
        return [evaluate_value(v, params) for v in value]
    if isinstance(value, str):
        try:
            return round(eval_expr(value, params.resolver()), ROUND_DIGITS)
        except (ExprError, SyntaxError, ZeroDivisionError, ValueError, TypeError):
            return value
    return value


# --------------------------------------------------------------------------- #
# Command helpers
# --------------------------------------------------------------------------- #


def cmd_kind(command):
    for key in ("cmd", "kind", "type"):
        if key in command and isinstance(command[key], str):
            return command[key].lower()
    return None


def is_construction(command):
    return command.get("construction") is True


def normalize_plane(plane):
    if isinstance(plane, str):
        return plane.upper()
    return plane


def canonical_json(obj):
    return json.dumps(obj, sort_keys=True, separators=(",", ":"))


# --------------------------------------------------------------------------- #
# Signature construction
# --------------------------------------------------------------------------- #


def build_signature(program):
    if not isinstance(program, list):
        program = []

    # 1. Parameters (register regardless of construction flag).
    param_exprs = {}
    for command in program:
        if cmd_kind(command) == "param" and "name" in command:
            param_exprs[command["name"]] = command.get("expr")
    params = ParamTable(param_exprs)

    # 2. Sketch planes.
    sketch_planes = {}
    for command in program:
        if cmd_kind(command) == "sketch" and "id" in command:
            sketch_planes[command["id"]] = normalize_plane(command.get("plane"))

    # 3. Profiles (skip construction profiles).
    profiles = {}
    for command in program:
        kind = cmd_kind(command)
        if kind in PROFILE_KINDS and not is_construction(command) and "id" in command:
            profiles[command["id"]] = canonical_profile(command, kind, sketch_planes, params)

    # 4. Extrude features and 5. opaque payloads.
    features = []  # list of (partition_key, feature_sig)
    opaque = []
    for command in program:
        if is_construction(command):
            continue
        kind = cmd_kind(command)
        if kind == "extrude":
            features.append(canonical_extrude(command, profiles, params))
        elif kind in SUPPORTED_KINDS or kind in ANNOTATION_KINDS:
            continue  # param/sketch/profile already handled; annotations ignored
        else:
            opaque.append(canonical_opaque(command))

    # 6. Partition grouping.
    default_partition = []
    labeled = {}
    for partition_key, feature_sig in features:
        if partition_key is None:
            default_partition.append(feature_sig)
        else:
            labeled.setdefault(partition_key, []).append(feature_sig)

    default_partition.sort()
    labeled_partitions = sorted(canonical_json(sorted(group)) for group in labeled.values())
    opaque.sort()

    return {
        "default_partition": default_partition,
        "labeled_partitions": labeled_partitions,
        "opaque": opaque,
    }


def canonical_profile(command, kind, sketch_planes, params):
    plane = command.get("plane")
    if plane is None:
        plane = sketch_planes.get(command.get("sketch"))
    plane = normalize_plane(plane)

    fields = {}
    for key, value in command.items():
        if key in PROFILE_NON_GEOM:
            continue
        evaluated = evaluate_value(value, params)
        if key == "angle" and isinstance(evaluated, (int, float)) and not isinstance(evaluated, bool):
            evaluated = round(((evaluated % 180.0) + 180.0) % 180.0, ROUND_DIGITS)
        fields[key] = evaluated

    return canonical_json({"type": kind, "plane": plane, "fields": fields})


def canonical_extrude(command, profiles, params):
    operation = command.get("operation", "new")
    extent = command.get("extent", "one_side")
    direction = command.get("direction", "normal")

    operation = operation.lower() if isinstance(operation, str) else operation
    extent = extent.lower() if isinstance(extent, str) else extent
    direction = direction.lower() if isinstance(direction, str) else direction

    depth = command.get("depth")
    depth = evaluate_value(depth, params) if depth is not None else None

    # Resolve consumed profiles (order not semantic).
    ref_ids = []
    if "profiles" in command and isinstance(command["profiles"], list):
        ref_ids.extend(command["profiles"])
    if "profile" in command and command["profile"] is not None:
        ref_ids.append(command["profile"])

    profile_sigs = []
    for pid in ref_ids:
        if pid in profiles:
            profile_sigs.append(profiles[pid])
    profile_sigs.sort()

    feature = {
        "operation": operation,
        "extent": extent,
        "direction": direction,
        "depth": depth,
        "profiles": profile_sigs,
    }
    feature_sig = canonical_json(feature)

    # Partition key: body or channel value (same namespace); None = default.
    partition_key = None
    if command.get("body") is not None:
        partition_key = ("v", command["body"])
    elif command.get("channel") is not None:
        partition_key = ("v", command["channel"])

    return partition_key, feature_sig


def canonical_opaque(command):
    payload = {k: v for k, v in command.items() if k != "id"}
    return canonical_json(payload)


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #


def decide(pair):
    sig_a = build_signature(pair.get("program_a", []))
    sig_b = build_signature(pair.get("program_b", []))
    equivalent = sig_a == sig_b
    return equivalent, sig_a, sig_b


def main() -> int:
    input_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/workspace/data/public_pairs.json")
    output_path = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("/workspace/predictions.json")

    payload = json.loads(input_path.read_text(encoding="utf-8"))

    predictions = []
    for pair in payload.get("pairs", []):
        equivalent, sig_a, sig_b = decide(pair)
        predictions.append(
            {
                "pair_id": pair["pair_id"],
                "equivalent": bool(equivalent),
                "signature_a": canonical_json(sig_a),
                "signature_b": canonical_json(sig_b),
            }
        )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps({"predictions": predictions}, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env python3
"""DeepCAD canonical-equivalence solver.

For each pair of compact CAD command programs, decide whether both programs
describe the same canonical solid.

Approach: normalize each program into a canonical *signature* that captures only
what affects the solid, then compare the two signatures for equality.

The solid is determined by:
  * the multiset of *extrude* features (each canonicalized: operation, extent,
    direction, depth, and the multiset of consumed profile geometries), grouped
    into body/channel partitions; and
  * the multiset of *opaque* solid-affecting payload commands compared literally.

Everything else (parameter names, ids, sketch ids, command order, construction
geometry, annotation commands) is normalized away or dropped.
"""

from __future__ import annotations

import ast
import json
import sys
from pathlib import Path


# Absolute tolerance for numeric comparison is 1e-6; rounding to 6 decimals
# yields a grid finer than that tolerance.
ROUND_DIGITS = 6

SUPPORTED_KINDS = {"param", "sketch", "rect", "circle", "slot", "extrude"}
PROFILE_KINDS = {"rect", "circle", "slot"}
ANNOTATION_KINDS = {"constraint", "dimension", "note", "metadata", "view", "comment"}

# Fields on a profile command that are labels/wiring, not geometry.
PROFILE_NON_GEOM = {"cmd", "kind", "type", "id", "sketch", "construction", "plane"}


# --------------------------------------------------------------------------- #
# Expression evaluation
# --------------------------------------------------------------------------- #


class ExprError(Exception):
    pass


def _eval_ast(node, resolver):
    if isinstance(node, ast.Expression):
        return _eval_ast(node.body, resolver)
    if isinstance(node, ast.Constant):
        if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
            raise ExprError("non-numeric constant")
        return float(node.value)
    if isinstance(node, ast.Num):  # pragma: no cover (py<3.8 compat)
        return float(node.n)
    if isinstance(node, ast.BinOp):
        left = _eval_ast(node.left, resolver)
        right = _eval_ast(node.right, resolver)
        if isinstance(node.op, ast.Add):
            return left + right
        if isinstance(node.op, ast.Sub):
            return left - right
        if isinstance(node.op, ast.Mult):
            return left * right
        if isinstance(node.op, ast.Div):
            return left / right
        raise ExprError("unsupported binary operator")
    if isinstance(node, ast.UnaryOp):
        operand = _eval_ast(node.operand, resolver)
        if isinstance(node.op, ast.UAdd):
            return +operand
        if isinstance(node.op, ast.USub):
            return -operand
        raise ExprError("unsupported unary operator")
    if isinstance(node, ast.Name):
        return resolver(node.id)
    raise ExprError("unsupported expression node")


def eval_expr(expr, resolver):
    tree = ast.parse(expr, mode="eval")
    return _eval_ast(tree, resolver)


class ParamTable:
    """Resolves parameter names to numeric values, allowing references between
    parameters (in any order) and detecting cycles."""

    def __init__(self, exprs):
        self._exprs = exprs
        self._cache = {}

    def resolve(self, name, stack=None):
        if name in self._cache:
            return self._cache[name]
        if name not in self._exprs:
            raise ExprError(f"unknown parameter: {name}")
        if stack is None:
            stack = ()
        if name in stack:
            raise ExprError(f"cyclic parameter: {name}")
        raw = self._exprs[name]
        value = self._eval_value(raw, stack + (name,))
        self._cache[name] = value
        return value

    def _eval_value(self, raw, stack):
        if isinstance(raw, bool):
            raise ExprError("boolean parameter")
        if isinstance(raw, (int, float)):
            return float(raw)
        if isinstance(raw, str):
            tree = ast.parse(raw, mode="eval")
            return _eval_ast(tree, lambda n: self.resolve(n, stack))
        raise ExprError("unsupported parameter expression")

    def resolver(self):
        return lambda name: self.resolve(name)


def evaluate_value(value, params):
    """Evaluate a JSON value (number, expression string, or nested list) to a
    canonical numeric form. If a string cannot be evaluated as an arithmetic
    expression, it is preserved literally."""
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return round(float(value), ROUND_DIGITS)
    if isinstance(value, list):
        return [evaluate_value(v, params) for v in value]
    if isinstance(value, str):
        try:
            return round(eval_expr(value, params.resolver()), ROUND_DIGITS)
        except (ExprError, SyntaxError, ZeroDivisionError, ValueError, TypeError):
            return value
    return value


# --------------------------------------------------------------------------- #
# Command helpers
# --------------------------------------------------------------------------- #


def cmd_kind(command):
    for key in ("cmd", "kind", "type"):
        if key in command and isinstance(command[key], str):
            return command[key].lower()
    return None


def is_construction(command):
    return command.get("construction") is True


def normalize_plane(plane):
    if isinstance(plane, str):
        return plane.upper()
    return plane


def canonical_json(obj):
    return json.dumps(obj, sort_keys=True, separators=(",", ":"))


# --------------------------------------------------------------------------- #
# Signature construction
# --------------------------------------------------------------------------- #


def build_signature(program):
    if not isinstance(program, list):
        program = []

    # 1. Parameters (register regardless of construction flag).
    param_exprs = {}
    for command in program:
        if isinstance(command, dict) and cmd_kind(command) == "param" and "name" in command:
            param_exprs[command["name"]] = command.get("expr")
    params = ParamTable(param_exprs)

    # 2. Sketch planes.
    sketch_planes = {}
    for command in program:
        if isinstance(command, dict) and cmd_kind(command) == "sketch" and "id" in command:
            sketch_planes[command["id"]] = normalize_plane(command.get("plane"))

    # 3. Profiles (skip construction profiles).
    profiles = {}
    for command in program:
        if not isinstance(command, dict):
            continue
        kind = cmd_kind(command)
        if kind in PROFILE_KINDS and not is_construction(command) and "id" in command:
            profiles[command["id"]] = canonical_profile(command, kind, sketch_planes, params)

    # 4. Extrude features and 5. opaque payloads.
    features = []  # list of (partition_key, feature_sig)
    opaque = []
    for command in program:
        if not isinstance(command, dict):
            continue
        if is_construction(command):
            continue
        kind = cmd_kind(command)
        if kind == "extrude":
            features.append(canonical_extrude(command, profiles, params))
        elif kind in SUPPORTED_KINDS or kind in ANNOTATION_KINDS:
            continue  # param/sketch/profile already handled; annotations ignored
        else:
            opaque.append(canonical_opaque(command))

    # 6. Partition grouping.
    default_partition = []
    labeled = {}
    for partition_key, feature_sig in features:
        if partition_key is None:
            default_partition.append(feature_sig)
        else:
            labeled.setdefault(partition_key, []).append(feature_sig)

    default_partition.sort()
    labeled_partitions = sorted(canonical_json(sorted(group)) for group in labeled.values())
    opaque.sort()

    return {
        "default_partition": default_partition,
        "labeled_partitions": labeled_partitions,
        "opaque": opaque,
    }


def canonical_profile(command, kind, sketch_planes, params):
    plane = command.get("plane")
    if plane is None:
        plane = sketch_planes.get(command.get("sketch"))
    plane = normalize_plane(plane)

    fields = {}
    for key, value in command.items():
        if key in PROFILE_NON_GEOM:
            continue
        evaluated = evaluate_value(value, params)
        if key == "angle" and isinstance(evaluated, (int, float)) and not isinstance(evaluated, bool):
            evaluated = round(((evaluated % 180.0) + 180.0) % 180.0, ROUND_DIGITS)
        fields[key] = evaluated

    return canonical_json({"type": kind, "plane": plane, "fields": fields})


def canonical_extrude(command, profiles, params):
    operation = command.get("operation", "new")
    extent = command.get("extent", "one_side")
    direction = command.get("direction", "normal")

    operation = operation.lower() if isinstance(operation, str) else operation
    extent = extent.lower() if isinstance(extent, str) else extent
    direction = direction.lower() if isinstance(direction, str) else direction

    depth = command.get("depth")
    depth = evaluate_value(depth, params) if depth is not None else None

    # Resolve consumed profiles (order not semantic).
    ref_ids = []
    if "profiles" in command and isinstance(command["profiles"], list):
        ref_ids.extend(command["profiles"])
    if "profile" in command and command["profile"] is not None:
        ref_ids.append(command["profile"])

    profile_sigs = []
    for pid in ref_ids:
        if pid in profiles:
            profile_sigs.append(profiles[pid])
    profile_sigs.sort()

    feature = {
        "operation": operation,
        "extent": extent,
        "direction": direction,
        "depth": depth,
        "profiles": profile_sigs,
    }
    feature_sig = canonical_json(feature)

    # Partition key: body or channel value (same namespace); None = default.
    partition_key = None
    if command.get("body") is not None:
        partition_key = ("v", command["body"])
    elif command.get("channel") is not None:
        partition_key = ("v", command["channel"])

    return partition_key, feature_sig


def canonical_opaque(command):
    payload = {k: v for k, v in command.items() if k != "id"}
    return canonical_json(payload)


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #


def decide(pair):
    sig_a = build_signature(pair.get("program_a", []))
    sig_b = build_signature(pair.get("program_b", []))
    equivalent = sig_a == sig_b
    return equivalent, sig_a, sig_b


def main() -> int:
    input_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/workspace/data/public_pairs.json")
    output_path = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("/workspace/predictions.json")

    payload = json.loads(input_path.read_text(encoding="utf-8"))

    predictions = []
    for pair in payload.get("pairs", []):
        equivalent, sig_a, sig_b = decide(pair)
        predictions.append(
            {
                "pair_id": pair["pair_id"],
                "equivalent": bool(equivalent),
                "signature_a": canonical_json(sig_a),
                "signature_b": canonical_json(sig_b),
            }
        )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps({"predictions": predictions}, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())