| """Nonlinear least-squares geometry constraint solver.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| from collections import defaultdict |
|
|
| import numpy as np |
| from scipy.optimize import least_squares |
|
|
| from statement_to_tikz.ir import ( |
| AngleMeasure, |
| CircleDef, |
| Collinear, |
| EqualAngle, |
| EqualLength, |
| GeometryIR, |
| Intersection, |
| Length, |
| Midpoint, |
| OnCircle, |
| OnLine, |
| Parallel, |
| Perpendicular, |
| SolveMode, |
| SolvedScene, |
| ) |
|
|
| DEFAULT_TOL = 1e-4 |
|
|
|
|
| def _get(coords: dict[str, tuple[float, float]], name: str) -> np.ndarray: |
| return np.asarray(coords[name], dtype=float) |
|
|
|
|
| def _vec(a: np.ndarray, b: np.ndarray) -> np.ndarray: |
| return b - a |
|
|
|
|
| def _cross2(u: np.ndarray, v: np.ndarray) -> float: |
| return float(u[0] * v[1] - u[1] * v[0]) |
|
|
|
|
| def _dot(u: np.ndarray, v: np.ndarray) -> float: |
| return float(np.dot(u, v)) |
|
|
|
|
| def _norm(u: np.ndarray) -> float: |
| return float(np.linalg.norm(u)) |
|
|
|
|
| def _angle_cos(a: np.ndarray, v: np.ndarray, b: np.ndarray) -> float: |
| u = a - v |
| w = b - v |
| nu, nw = _norm(u), _norm(w) |
| if nu < 1e-12 or nw < 1e-12: |
| return 1.0 |
| return _dot(u, w) / (nu * nw) |
|
|
|
|
| def _circle_geometry( |
| ir: GeometryIR, coords: dict[str, tuple[float, float]], circle_id: str |
| ) -> tuple[np.ndarray, float] | None: |
| circ = next((c for c in ir.circles if c.id == circle_id), None) |
| if circ is None: |
| return None |
| return _circle_center_radius_from_def(circ, coords) |
|
|
|
|
| def _circle_center_radius_from_def( |
| circ: CircleDef, coords: dict[str, tuple[float, float]] |
| ) -> tuple[np.ndarray, float] | None: |
| if circ.center is not None and circ.center in coords: |
| c = _get(coords, circ.center) |
| if circ.radius is not None: |
| return c, float(circ.radius) |
| if circ.through: |
| for p in circ.through: |
| if p in coords: |
| return c, _norm(_get(coords, p) - c) |
| return c, 1.0 |
| if circ.through and len(circ.through) >= 3: |
| pts = [p for p in circ.through[:3] if p in coords] |
| if len(pts) < 3: |
| return None |
| a, b, c = (_get(coords, p) for p in pts) |
| ax, ay = a |
| bx, by = b |
| cx, cy = c |
| d = 2 * (ax * (by - cy) + bx * (cy - ay) + cx * (ay - by)) |
| if abs(d) < 1e-12: |
| return None |
| ux = ( |
| (ax**2 + ay**2) * (by - cy) |
| + (bx**2 + by**2) * (cy - ay) |
| + (cx**2 + cy**2) * (ay - by) |
| ) / d |
| uy = ( |
| (ax**2 + ay**2) * (cx - bx) |
| + (bx**2 + by**2) * (ax - cx) |
| + (cx**2 + cy**2) * (bx - ax) |
| ) / d |
| center = np.array([ux, uy]) |
| return center, _norm(a - center) |
| return None |
|
|
|
|
| def constraint_residuals( |
| ir: GeometryIR, coords: dict[str, tuple[float, float]] |
| ) -> list[float]: |
| """Return list of scalar residuals for each constraint (may be multi-valued flattened).""" |
| r: list[float] = [] |
| for c in ir.constraints: |
| if isinstance(c, EqualLength): |
| la = _norm(_get(coords, c.a2) - _get(coords, c.a1)) |
| lb = _norm(_get(coords, c.b2) - _get(coords, c.b1)) |
| r.append(la - lb) |
| elif isinstance(c, Length): |
| la = _norm(_get(coords, c.b) - _get(coords, c.a)) |
| r.append(la - c.value) |
| elif isinstance(c, EqualAngle): |
| cos1 = _angle_cos( |
| _get(coords, c.a1), _get(coords, c.v1), _get(coords, c.b1) |
| ) |
| cos2 = _angle_cos( |
| _get(coords, c.a2), _get(coords, c.v2), _get(coords, c.b2) |
| ) |
| r.append(cos1 - cos2) |
| elif isinstance(c, AngleMeasure): |
| cos_t = math.cos(math.radians(c.degrees)) |
| cos_m = _angle_cos( |
| _get(coords, c.a), _get(coords, c.vertex), _get(coords, c.b) |
| ) |
| r.append(cos_m - cos_t) |
| elif isinstance(c, Perpendicular): |
| u = _vec(_get(coords, c.a1), _get(coords, c.a2)) |
| v = _vec(_get(coords, c.b1), _get(coords, c.b2)) |
| r.append(_dot(u, v)) |
| elif isinstance(c, Parallel): |
| u = _vec(_get(coords, c.a1), _get(coords, c.a2)) |
| v = _vec(_get(coords, c.b1), _get(coords, c.b2)) |
| r.append(_cross2(u, v)) |
| elif isinstance(c, OnLine): |
| a, b, p = _get(coords, c.a), _get(coords, c.b), _get(coords, c.point) |
| r.append(_cross2(b - a, p - a)) |
| elif isinstance(c, OnCircle): |
| geom = _circle_geometry(ir, coords, c.circle) |
| if geom is None: |
| r.append(0.0) |
| else: |
| center, rad = geom |
| r.append(_norm(_get(coords, c.point) - center) - rad) |
| elif isinstance(c, Midpoint): |
| m = _get(coords, c.point) |
| mid = 0.5 * (_get(coords, c.a) + _get(coords, c.b)) |
| r.extend([float(m[0] - mid[0]), float(m[1] - mid[1])]) |
| elif isinstance(c, Collinear): |
| pts = [_get(coords, p) for p in c.points] |
| base = pts[1] - pts[0] |
| for p in pts[2:]: |
| r.append(_cross2(base, p - pts[0])) |
| elif isinstance(c, Intersection): |
| a, b = _get(coords, c.a), _get(coords, c.b) |
| cc, d = _get(coords, c.c), _get(coords, c.d) |
| p = _get(coords, c.point) |
| r.append(_cross2(b - a, p - a)) |
| r.append(_cross2(d - cc, p - cc)) |
| else: |
| raise TypeError(f"unknown constraint {type(c)}") |
|
|
| for circ in ir.circles: |
| if circ.center and circ.through and circ.center in coords: |
| center = _get(coords, circ.center) |
| if circ.radius is not None: |
| rad = circ.radius |
| else: |
| refs = [p for p in circ.through if p in coords] |
| if not refs: |
| continue |
| rad = _norm(_get(coords, refs[0]) - center) |
| for p in circ.through: |
| if p in coords: |
| r.append(_norm(_get(coords, p) - center) - rad) |
| elif circ.through and len(circ.through) >= 3: |
| pts = [p for p in circ.through if p in coords] |
| if len(pts) >= 3: |
| geom = _circle_center_radius_from_def(circ, coords) |
| if geom is not None: |
| center, rad = geom |
| for p in pts: |
| r.append(_norm(_get(coords, p) - center) - rad) |
|
|
| return r |
|
|
|
|
| def evaluate_residuals( |
| ir: GeometryIR, coords: dict[str, tuple[float, float]] |
| ) -> list[float]: |
| return [abs(x) for x in constraint_residuals(ir, coords)] |
|
|
|
|
| def _initial_coords(ir: GeometryIR) -> dict[str, tuple[float, float]]: |
| coords: dict[str, tuple[float, float]] = {} |
| n = len(ir.points) |
| for i, p in enumerate(ir.points): |
| if p.hint is not None: |
| coords[p.name] = (float(p.hint[0]), float(p.hint[1])) |
| else: |
| ang = 2 * math.pi * i / max(n, 1) - math.pi / 2 |
| coords[p.name] = (math.cos(ang), math.sin(ang)) |
| return coords |
|
|
|
|
| def _gauge_point_names(ir: GeometryIR) -> tuple[str | None, str | None]: |
| """Choose two points to fix: prefer a length-constrained segment as base.""" |
| names = ir.point_names() |
| if not names: |
| return None, None |
|
|
| for c in ir.constraints: |
| if isinstance(c, Length): |
| return c.a, c.b |
|
|
| for c in ir.constraints: |
| if isinstance(c, Midpoint): |
| return c.a, c.b |
|
|
| if ir.segments: |
| return ir.segments[0].a, ir.segments[0].b |
|
|
| p0 = names[0] |
| p1 = names[1] if len(names) > 1 else None |
| return p0, p1 |
|
|
|
|
| def _triangle_area_residuals( |
| ir: GeometryIR, coords: dict[str, tuple[float, float]] |
| ) -> list[float]: |
| """Soft non-degeneracy: penalize near-zero area for 3-cycles of segments.""" |
| adj: dict[str, set[str]] = defaultdict(set) |
| for seg in ir.segments: |
| adj[seg.a].add(seg.b) |
| adj[seg.b].add(seg.a) |
|
|
| residuals: list[float] = [] |
| seen: set[tuple[str, str, str]] = set() |
| for a, nbrs in adj.items(): |
| for b in nbrs: |
| for c in nbrs: |
| if b >= c: |
| continue |
| if c not in adj[b]: |
| continue |
| key = tuple(sorted((a, b, c))) |
| if key in seen: |
| continue |
| seen.add(key) |
| if not all(p in coords for p in key): |
| continue |
| pa, pb, pc = (_get(coords, p) for p in key) |
| area2 = abs(_cross2(pb - pa, pc - pa)) |
| residuals.append(max(0.0, 0.5 - area2)) |
| return residuals |
|
|
|
|
| def _remap_init_to_gauge( |
| init: dict[str, tuple[float, float]], |
| p0: str, |
| p1: str, |
| gauged: dict[str, tuple[float, float]], |
| ) -> dict[str, tuple[float, float]]: |
| """Rigidly map hinted positions onto the gauged base segment.""" |
| if p0 not in init or p1 not in init: |
| return gauged |
| o0 = np.asarray(init[p0], dtype=float) |
| o1 = np.asarray(init[p1], dtype=float) |
| n0 = np.asarray(gauged[p0], dtype=float) |
| n1 = np.asarray(gauged[p1], dtype=float) |
| o_len = _norm(o1 - o0) or 1.0 |
| n_len = _norm(n1 - n0) or 1.0 |
| scale = n_len / o_len |
| od = (o1 - o0) / o_len |
| nd = (n1 - n0) / n_len |
| ang = math.atan2(nd[1], nd[0]) - math.atan2(od[1], od[0]) |
| ca, sa = math.cos(ang), math.sin(ang) |
| rot = np.array([[ca, -sa], [sa, ca]]) |
| remapped = dict(gauged) |
| for n, xy in init.items(): |
| if n in (p0, p1): |
| continue |
| local = (np.asarray(xy, dtype=float) - o0) * scale |
| remapped[n] = tuple(n0 + rot @ local) |
| return remapped |
|
|
|
|
| def _degrees_of_freedom(ir: GeometryIR) -> tuple[int, int]: |
| """Rough DOF count: 2*|points| - 3 (gauge) vs number of scalar residuals.""" |
| n_coords = 2 * len(ir.points) |
| gauge = min(3, n_coords) |
| probe = _initial_coords(ir) |
| n_res = len(constraint_residuals(ir, probe)) |
| return n_coords - gauge, n_res |
|
|
|
|
| def solve_geometry( |
| ir: GeometryIR, |
| *, |
| tol: float = DEFAULT_TOL, |
| max_nfev: int = 2000, |
| ) -> SolvedScene: |
| """Solve for point coordinates. Returns exact mode if residuals within tol.""" |
| names = ir.point_names() |
| if not names: |
| return SolvedScene( |
| ir=ir, |
| coordinates={}, |
| mode=SolveMode.failed, |
| max_residual=0.0, |
| message="no points", |
| ) |
|
|
| init = _initial_coords(ir) |
| dof, n_res = _degrees_of_freedom(ir) |
| if n_res == 0 or n_res < max(1, dof - 1): |
| res = evaluate_residuals(ir, init) |
| return SolvedScene( |
| ir=ir, |
| coordinates=init, |
| mode=SolveMode.failed, |
| max_residual=max(res) if res else 0.0, |
| residuals=res, |
| message=f"underconstrained (dof≈{dof}, residuals={n_res})", |
| ) |
|
|
| p0, p1 = _gauge_point_names(ir) |
|
|
| free_names: list[str] = [n for n in names if n != p0] |
|
|
| def apply_gauge(partial: dict[str, tuple[float, float]]) -> dict[str, tuple[float, float]]: |
| full = dict(partial) |
| if p0: |
| full[p0] = (0.0, 0.0) |
| if p1 and p1 in full: |
| x, _y = full[p1] |
| if abs(x) < 1e-9: |
| x = 1.0 |
| full[p1] = (abs(x), 0.0) |
| elif p1: |
| full[p1] = (1.0, 0.0) |
| return full |
|
|
| def pack_free(coords: dict[str, tuple[float, float]]) -> np.ndarray: |
| vals: list[float] = [] |
| for n in free_names: |
| x, y = coords[n] |
| if n == p1: |
| vals.append(x if abs(x) > 1e-9 else 1.0) |
| else: |
| vals.extend([x, y]) |
| return np.asarray(vals, dtype=float) |
|
|
| def unpack_free(v: np.ndarray) -> dict[str, tuple[float, float]]: |
| partial: dict[str, tuple[float, float]] = {} |
| idx = 0 |
| for n in free_names: |
| if n == p1: |
| partial[n] = (float(v[idx]), 0.0) |
| idx += 1 |
| else: |
| partial[n] = (float(v[idx]), float(v[idx + 1])) |
| idx += 2 |
| return apply_gauge(partial) |
|
|
| def fun(v: np.ndarray) -> np.ndarray: |
| coords = unpack_free(v) |
| res = list(constraint_residuals(ir, coords)) |
| res.extend(_triangle_area_residuals(ir, coords)) |
| if not res: |
| extras = [] |
| if p1 and p1 in coords: |
| extras.append(coords[p1][0] - 1.0) |
| return np.asarray(extras or [0.0], dtype=float) |
| return np.asarray(res, dtype=float) |
|
|
| gauged_init = apply_gauge(init) |
| if p0 and p1: |
| gauged_init = apply_gauge(_remap_init_to_gauge(init, p0, p1, gauged_init)) |
|
|
| x0 = pack_free(gauged_init) |
| if x0.size == 0: |
| coords = gauged_init |
| res = evaluate_residuals(ir, coords) |
| max_r = max(res) if res else 0.0 |
| mode = SolveMode.exact if max_r <= tol else SolveMode.failed |
| return SolvedScene( |
| ir=ir, |
| coordinates=coords, |
| mode=mode, |
| max_residual=max_r, |
| residuals=res, |
| message="only gauge points", |
| ) |
|
|
| try: |
| result = least_squares( |
| fun, x0, ftol=1e-12, xtol=1e-12, gtol=1e-12, max_nfev=max_nfev |
| ) |
| coords = unpack_free(result.x) |
| except Exception as exc: |
| coords = gauged_init |
| res = evaluate_residuals(ir, coords) |
| return SolvedScene( |
| ir=ir, |
| coordinates=coords, |
| mode=SolveMode.failed, |
| max_residual=max(res) if res else float("inf"), |
| residuals=res, |
| message=f"solver error: {exc}", |
| ) |
|
|
| res = evaluate_residuals(ir, coords) |
| max_r = max(res) if res else 0.0 |
| if max_r <= tol: |
| mode = SolveMode.exact |
| msg = "constraints satisfied" |
| else: |
| mode = SolveMode.failed |
| msg = f"residuals too large ({max_r:.4g})" |
|
|
| return SolvedScene( |
| ir=ir, |
| coordinates={k: (float(v[0]), float(v[1])) for k, v in coords.items()}, |
| mode=mode, |
| max_residual=float(max_r), |
| residuals=res, |
| message=msg, |
| ) |
|
|
|
|
| def scene_from_coordinates( |
| ir: GeometryIR, |
| coords: dict[str, tuple[float, float]], |
| *, |
| mode: SolveMode, |
| message: str = "", |
| ) -> SolvedScene: |
| res = evaluate_residuals(ir, coords) |
| max_r = max(res) if res else 0.0 |
| return SolvedScene( |
| ir=ir, |
| coordinates=coords, |
| mode=mode, |
| max_residual=float(max_r), |
| residuals=res, |
| message=message, |
| ) |
|
|
